import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LogisticRegression

# Funkcja do tworzenia jednego zestawu danych
def generuj_zbior(n, seed, k_true, x0_true):
    rng = np.random.default_rng(seed)
    czas = rng.uniform(0, 60, n)
    prawd = 1 / (1 + np.exp(-k_true * (czas - x0_true)))
    wynik = rng.binomial(1, prawd)
    return czas, wynik

# Definicje 4 zestawów
zbiory = [
    {"n": 100, "seed": 1, "k_true": 0.2, "x0_true": 30, "etykieta": "Zbiór 1"},
    {"n": 100, "seed": 2, "k_true": 0.15, "x0_true": 25, "etykieta": "Zbiór 2"},
    {"n": 100, "seed": 3, "k_true": 0.25, "x0_true": 35, "etykieta": "Zbiór 3"},
    {"n": 100, "seed": 4, "k_true": 0.1, "x0_true": 28, "etykieta": "Zbiór 4"},
]

# Siatka do rysowania krzywej
siatka = np.linspace(0, 60, 500).reshape(-1, 1)

# ===== 4 oddzielne wykresy =====
fig, axs = plt.subplots(2, 2, figsize=(12, 10))
axs = axs.ravel()

krzywe = []  # tu zapiszemy krzywe do 5. wykresu

for i, zbior in enumerate(zbiory):
    # Dane
    czas, wynik = generuj_zbior(zbior["n"], zbior["seed"], zbior["k_true"], zbior["x0_true"])
    X = czas.reshape(-1, 1)

    # Model
    model = LogisticRegression()
    model.fit(X, wynik)

    # Przewidywania
    przewidziane = model.predict_proba(siatka)[:, 1] * 100
    krzywe.append((przewidziane, zbior["etykieta"]))  # zapisujemy do 5. wykresu

    # Punkty uczące
    axs[i].scatter(czas[wynik == 0], wynik[wynik == 0] * 100,
                   c="red", alpha=0.6, label="Porażka (0)")
    axs[i].scatter(czas[wynik == 1], wynik[wynik == 1] * 100,
                   c="green", alpha=0.6, label="Sukces (1)")

    # Krzywa
    axs[i].plot(siatka, przewidziane, "b-", linewidth=2, label="Regresja logistyczna")

    # Opisy
    axs[i].set_title(zbior["etykieta"])
    axs[i].set_xlabel("Czas treningu [min]")
    axs[i].set_ylabel("Sukces [%] (0/100)")
    axs[i].grid(True, alpha=0.3)
    axs[i].legend()

plt.suptitle("Regresja logistyczna: 4 różne zestawy danych", fontsize=14)
plt.tight_layout(rect=[0, 0, 1, 0.97])
plt.show()

# ===== 5. wykres: wszystkie krzywe razem =====
plt.figure(figsize=(10, 6))

for przewidziane, etykieta in krzywe:
    plt.plot(siatka, przewidziane, linewidth=2, label=etykieta)

plt.xlabel("Czas treningu [min]")
plt.ylabel("Prawdopodobieństwo sukcesu [%]")
plt.title("Porównanie krzywych regresji logistycznej dla 4 zestawów")
plt.grid(True, alpha=0.3)
plt.legend()
plt.tight_layout()
plt.show()