"""
Bias-Variance Playground — interactive parameter exploration.
"""

import numpy as np

META = {
    "title": "Bias-Variance Tradeoff",
    "icon": "⚖️",
    "description": (
        "The fundamental tradeoff in ML: simple models underfit (high bias), "
        "complex models overfit (high variance). Adjust model complexity and noise "
        "to see how train vs test error diverge."
    ),
}

PARAMS = [
    {
        "name": "max_degree",
        "label": "Max Polynomial Degree",
        "type": "slider",
        "min": 1,
        "max": 15,
        "step": 1,
        "default": 12,
    },
    {
        "name": "noise",
        "label": "Noise Level",
        "type": "slider",
        "min": 0.1,
        "max": 3.0,
        "step": 0.1,
        "default": 1.0,
    },
    {
        "name": "n_points",
        "label": "Data Points",
        "type": "slider",
        "min": 15,
        "max": 200,
        "step": 5,
        "default": 40,
    },
]


def compute(params: dict) -> dict:
    max_deg = int(params.get("max_degree", 12))
    noise = float(params.get("noise", 1.0))
    n = int(params.get("n_points", 40))

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

    # Split into train/test
    split = int(0.7 * n)
    X_train, X_test = X[:split], X[split:]
    y_train, y_test = y[:split], y[split:]

    degrees = list(range(1, max_deg + 1))
    train_errors = []
    test_errors = []

    for d in degrees:
        coeffs = np.polyfit(X_train, y_train, d)
        y_pred_train = np.polyval(coeffs, X_train)
        y_pred_test = np.polyval(coeffs, X_test)

        train_mse = np.mean((y_train - y_pred_train) ** 2)
        test_mse = np.mean((y_test - y_pred_test) ** 2)
        train_errors.append(float(train_mse))
        test_errors.append(float(min(test_mse, 10)))  # clip extreme values

    # Find best degree
    best_idx = int(np.argmin(test_errors))
    best_degree = degrees[best_idx]

    data = [
        {
            "x": degrees, "y": train_errors,
            "mode": "lines+markers", "type": "scatter", "name": "Train Error (Bias²)",
            "line": {"color": "#6C63FF", "width": 3},
            "marker": {"size": 7, "color": "#6C63FF", "line": {"width": 1, "color": "#fff"}},
        },
        {
            "x": degrees, "y": test_errors,
            "mode": "lines+markers", "type": "scatter", "name": "Test Error (Bias² + Variance)",
            "line": {"color": "#FF6584", "width": 3},
            "marker": {"size": 7, "color": "#FF6584", "line": {"width": 1, "color": "#fff"}},
        },
        {
            "x": [best_degree], "y": [test_errors[best_idx]],
            "mode": "markers", "type": "scatter", "name": f"Sweet Spot (degree={best_degree})",
            "marker": {"size": 16, "color": "#34d399", "symbol": "star",
                        "line": {"width": 2, "color": "#fff"}},
        },
    ]

    # Add annotations for regions
    layout = {
        "template": "plotly_dark",
        "paper_bgcolor": "rgba(0,0,0,0)", "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": f"Bias-Variance Tradeoff — Best at degree={best_degree}", "font": {"size": 20}},
        "xaxis": {"title": "Model Complexity (Polynomial Degree)", "dtick": 1, "gridcolor": "rgba(255,255,255,0.08)"},
        "yaxis": {"title": "Mean Squared Error", "gridcolor": "rgba(255,255,255,0.08)"},
        "legend": {"x": 0.02, "y": 0.98},
        "margin": {"l": 60, "r": 30, "t": 60, "b": 50},
        "annotations": [
            {"x": 2, "y": max(test_errors[0:3]) * 0.9 if test_errors else 1,
             "text": "← Underfitting", "showarrow": False,
             "font": {"size": 14, "color": "#fbbf24"}, "xanchor": "left"},
            {"x": max_deg - 1, "y": max(test_errors[-3:]) * 0.9 if test_errors else 1,
             "text": "Overfitting →", "showarrow": False,
             "font": {"size": 14, "color": "#fbbf24"}, "xanchor": "right"},
        ],
    }
    return {"data": data, "layout": layout}
