import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from sklearn import datasets
from sklearn.utils import shuffle

EXPLAIN_MD = """
### Wat testen we hier?
We bekijken of er een **lineair verband** is tussen **BMI** en de **diabetes-progressiescore** in een bekende (openbare) dataset.
Dat doen we met *supervised learning*: het model ziet voorbeelden `(BMI → score)` en leert een lijn \(y = w x + b\) die dit verband benadert.

**Hoe meten we of dat gelukt is?**
- We splitsen de data in **train (80%)** en **test (20%)**.
- We **trainen** het model alleen op de **trainset**.
- We **toetsen** het resultaat op de **testset** die het model niet gezien heeft.
- We rapporteren **MSE** (gemiddelde kwadratische fout) en **R²** (uitlegvariantie) op de testset.

> Let op: in deze sklearn-dataset is BMI **genormaliseerd** (geschaald). De helling `w` geeft wel de **richting en sterkte** aan (positief = hogere BMI hangt samen met hogere score).
"""

STORY_MD = r"""
### Waarom kijken we naar BMI en diabetes?
Stel je voor dat je arts wilt begrijpen of **het gewicht van mensen** (uitgedrukt als *Body Mass Index*, BMI)
iets zegt over hun **gezondheid**. Een van de dingen die onderzocht wordt is het verband tussen BMI en de
**ernst van diabetes**.

We gebruiken hier **echte gegevens** uit een medische dataset (dus **geen foto’s**, maar gemeten waarden van mensen
die in een onderzoek hebben meegedaan). Elke deelnemer heeft:

- een **BMI-waarde** (hoe zwaar of licht iemand is ten opzichte van zijn lengte),
- en een **score** die aangeeft hoe ernstig de diabetes bij die persoon verloopt.

Met lineaire regressie testen we: *kunnen we een lijn tekenen die laat zien of een hogere BMI vaak samenvalt met een
hogere (of juist lagere) score?*

**Waarom is dat belangrijk?**
- Als er wél een duidelijk verband is, kan dit helpen om **risico’s eerder te signaleren**.
- Als er géén verband is, leren we dat BMI misschien niet de juiste voorspeller is en moet er verder gekeken worden
  naar andere factoren.

Kortom: dit experiment laat je zien hoe data ons kan helpen om **patronen in gezondheid** te ontdekken — en dat doen
we hier stap voor stap, live op je scherm.
"""

CONCLUSION_MD = r"""
# **Conclusie**

Mensen met een hogere **BMI** hebben in dit onderzoek gemiddeld vaker een ernstiger verloop van **diabetes**.  
Maar **BMI is niet de enige factor** — leeftijd, erfelijkheid, leefstijl en andere medische waarden spelen ook mee.
"""

def load_bmi_diabetes():
    d = datasets.load_diabetes()
    X = d.data[:, 2]  # BMI feature (genormaliseerd)
    y = d.target      # Progressiescore
    return X.astype(np.float64), y.astype(np.float64), "Diabetes: BMI vs. score"

def train_test_split_1d(x, y, test_size=0.2, seed=42):
    rng = np.random.RandomState(seed)
    idx = np.arange(x.shape[0])
    rng.shuffle(idx)
    n_test = int(len(idx) * test_size)
    test_idx = idx[:n_test]
    train_idx = idx[n_test:]
    return x[train_idx], y[train_idx], x[test_idx], y[test_idx]

def sgd_train_generator(lr, epochs, batch_size, seed, split_seed):
    # Data & split
    x, y, label = load_bmi_diabetes()
    x_tr, y_tr, x_te, y_te = train_test_split_1d(x, y, test_size=0.2, seed=int(split_seed))

    n = x_tr.shape[0]
    w, b = 0.0, 0.0

    x_min, x_max = float(np.min(x)), float(np.max(x))
    train_losses, test_losses = [], []

    rng = np.random.RandomState(int(seed))

    for epoch in range(1, int(epochs) + 1):
        # shuffle train set
        x_tr, y_tr = shuffle(x_tr, y_tr, random_state=rng)

        # SGD over mini-batches
        for start in range(0, n, int(batch_size)):
            end = min(start + int(batch_size), n)
            xb, yb = x_tr[start:end], y_tr[start:end]
            yhat = w * xb + b
            err = yb - yhat
            dw = -(2.0 / xb.size) * np.sum(xb * err)
            db = -(2.0 / xb.size) * np.sum(err)
            w -= lr * dw
            b -= lr * db

        # Metrics on train and test
        y_tr_pred = w * x_tr + b
        y_te_pred = w * x_te + b
        mse_tr = float(np.mean((y_tr - y_tr_pred)**2))
        mse_te = float(np.mean((y_te - y_te_pred)**2))

        # R^2 on test
        ss_res = float(np.sum((y_te - y_te_pred)**2))
        ss_tot = float(np.sum((y_te - np.mean(y_te))**2))
        r2_te = 1.0 - ss_res / ss_tot if ss_tot > 0 else float("nan")

        train_losses.append(mse_tr)
        test_losses.append(mse_te)

        # Plot 1: data (train vs test) + regressielijn
        fig_main = plt.figure(figsize=(7, 4))
        ax1 = fig_main.add_subplot(111)
        ax1.scatter(x_tr, y_tr, alpha=0.6, s=18, label="train")
        ax1.scatter(x_te, y_te, alpha=0.8, s=22, marker="x", label="test")
        xs = np.linspace(x_min, x_max, 200)
        ax1.plot(xs, w * xs + b, linewidth=2, label="model")
        ax1.set_title(f"{label} — Epoch {epoch}/{epochs}")
        ax1.set_xlabel("BMI (genormaliseerd)")
        ax1.set_ylabel("Progressiescore")
        ax1.legend()
        ax1.grid(True, linestyle=":", linewidth=0.6)
        plt.tight_layout()

        # Plot 2: loss-curve (train & test)
        fig_loss = plt.figure(figsize=(7, 3.5))
        ax2 = fig_loss.add_subplot(111)
        ax2.plot(range(1, len(train_losses)+1), train_losses, marker="o", label="Train MSE")
        ax2.plot(range(1, len(test_losses)+1), test_losses, marker="o", linestyle="--", label="Test MSE")
        ax2.set_title("Loss-curve (MSE per epoch) — lager is beter")
        ax2.set_xlabel("Epoch")
        ax2.set_ylabel("MSE")
        ax2.legend()
        ax2.grid(True, linestyle=":", linewidth=0.6)
        plt.tight_layout()

        # Resultaten + opvallende conclusie
        verdict = "positief" if w >= 0 else "negatief"
        summary = (
            f"**Wat levert dit op?**\n"
            f"- Huidige regressielijn: `y = {w:.4f} * x + {b:.4f}`\n"
            f"- Train MSE: `{mse_tr:.2f}` — Test MSE: `{mse_te:.2f}` — Test R²: `{r2_te:.3f}`\n"
            f"- Interpretatie: het verband tussen BMI en progressiescore is **{verdict}** in deze dataset "
            f"(hogere BMI hangt samen met hogere score als `w > 0`).\n\n"
            f"{CONCLUSION_MD}"
        )

        yield fig_main, fig_loss, summary

with gr.Blocks(title="Diabetes: BMI → Progressiescore (Live Regressie)") as demo:
    gr.Markdown("# Diabetes: BMI → Progressiescore (Live Lineaire Regressie)")
    gr.Markdown(EXPLAIN_MD)

    with gr.Row():
        with gr.Column(scale=1):
            lr = gr.Slider(1e-4, 1e-0, value=5e-3, step=1e-4, label="Learning rate")
            epochs = gr.Slider(5, 200, value=60, step=1, label="Epochs")
            batch = gr.Slider(8, 256, value=64, step=1, label="Batchgrootte")
            seed = gr.Slider(0, 9999, value=42, step=1, label="Training seed")
            split_seed = gr.Slider(0, 9999, value=7, step=1, label="Train/test split seed")
            train_btn = gr.Button("Train live")
            # Story direct onder de knop
            gr.Markdown(STORY_MD)
        with gr.Column(scale=2):
            plot_main = gr.Plot(label="Data (train/test) & regressielijn (live)")
            plot_loss = gr.Plot(label="Loss-curve (MSE per epoch) — train vs test")
            results = gr.Markdown()

    # Training starten via knop
    train_btn.click(
        fn=sgd_train_generator,
        inputs=[lr, epochs, batch, seed, split_seed],
        outputs=[plot_main, plot_loss, results]
    )

    # Auto-train bij laden met default-waarden
    demo.load(
        fn=sgd_train_generator,
        inputs=[lr, epochs, batch, seed, split_seed],
        outputs=[plot_main, plot_loss, results]
    )

if __name__ == "__main__":
    demo.launch()