"""
Regularization Playground — interactive parameter exploration.
"""

import numpy as np

META = {
    "title": "Regularization",
    "icon": "🛡️",
    "description": (
        "See how L1 (Lasso) and L2 (Ridge) regularization prevent overfitting. "
        "With a high polynomial degree the model overfits — increase alpha to "
        "watch the curve smooth out as large coefficients are penalized."
    ),
}

PARAMS = [
    {
        "name": "alpha",
        "label": "Alpha (Regularization Strength)",
        "type": "slider",
        "min": 0.0,
        "max": 10.0,
        "step": 0.1,
        "default": 0.1,
    },
    {
        "name": "reg_type",
        "label": "Regularization Type",
        "type": "select",
        "options": ["Ridge (L2)", "Lasso (L1)"],
        "default": "Ridge (L2)",
    },
    {
        "name": "degree",
        "label": "Polynomial Degree",
        "type": "slider",
        "min": 1,
        "max": 15,
        "step": 1,
        "default": 8,
    },
    {
        "name": "noise",
        "label": "Noise Level",
        "type": "slider",
        "min": 0.0,
        "max": 5.0,
        "step": 0.25,
        "default": 1.5,
    },
]


def compute(params: dict) -> dict:
    alpha = float(params.get("alpha", 0.1))
    reg_type = params.get("reg_type", "Ridge (L2)")
    degree = int(params.get("degree", 8))
    noise = float(params.get("noise", 1.5))

    np.random.seed(42)
    n = 50
    X = np.linspace(0, 1, n)
    y_true = np.sin(2 * np.pi * X)
    y = y_true + np.random.normal(0, noise * 0.3, size=n)

    # Build polynomial features
    from sklearn.preprocessing import PolynomialFeatures
    from sklearn.linear_model import Ridge, Lasso
    from sklearn.pipeline import make_pipeline

    RegModel = Ridge if "Ridge" in reg_type else Lasso
    model = make_pipeline(PolynomialFeatures(degree), RegModel(alpha=max(alpha, 1e-6)))
    model.fit(X.reshape(-1, 1), y)

    X_smooth = np.linspace(0, 1, 300)
    y_fit = model.predict(X_smooth.reshape(-1, 1))

    # Also fit unregularized for comparison
    model_unreg = make_pipeline(PolynomialFeatures(degree), Ridge(alpha=1e-10))
    model_unreg.fit(X.reshape(-1, 1), y)
    y_unreg = model_unreg.predict(X_smooth.reshape(-1, 1))

    # Compute R²
    y_pred_train = model.predict(X.reshape(-1, 1))
    ss_res = np.sum((y - y_pred_train) ** 2)
    ss_tot = np.sum((y - np.mean(y)) ** 2)
    r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0

    reg_label = "L2 (Ridge)" if "Ridge" in reg_type else "L1 (Lasso)"

    data = [
        {
            "x": X.tolist(), "y": y.tolist(),
            "mode": "markers", "type": "scatter", "name": "Data Points",
            "marker": {"size": 8, "color": "#6C63FF", "opacity": 0.8, "line": {"width": 1, "color": "#fff"}},
        },
        {
            "x": X_smooth.tolist(), "y": np.clip(y_unreg, -5, 5).tolist(),
            "mode": "lines", "type": "scatter", "name": "No Regularization",
            "line": {"color": "rgba(251,191,36,0.4)", "width": 2, "dash": "dot"},
        },
        {
            "x": X_smooth.tolist(), "y": np.clip(y_fit, -5, 5).tolist(),
            "mode": "lines", "type": "scatter",
            "name": f"{reg_label} α={alpha} (R²={r2:.3f})",
            "line": {"color": "#FF6584", "width": 3},
        },
        {
            "x": X_smooth.tolist(), "y": np.sin(2 * np.pi * X_smooth).tolist(),
            "mode": "lines", "type": "scatter", "name": "True Function",
            "line": {"color": "#34d399", "width": 2, "dash": "dash"},
        },
    ]

    layout = {
        "template": "plotly_dark",
        "paper_bgcolor": "rgba(0,0,0,0)", "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": f"Regularization — {reg_label}, α={alpha}, degree={degree}", "font": {"size": 20}},
        "xaxis": {"title": "X", "gridcolor": "rgba(255,255,255,0.08)"},
        "yaxis": {"title": "y", "gridcolor": "rgba(255,255,255,0.08)", "range": [-3, 3]},
        "legend": {"x": 0.02, "y": 0.98},
        "margin": {"l": 60, "r": 30, "t": 60, "b": 50},
    }
    return {"data": data, "layout": layout}
