"""
Regularization (Ridge & Lasso) — comprehensive topic content with interactive graphs.
"""

import json
import numpy as np


def _coefficient_paths():
    """Show how coefficients shrink with increasing regularization."""
    alphas = np.logspace(-2, 3, 100)
    np.random.seed(42)
    n_features = 5
    true_coefs = [3.5, -2.1, 1.5, 0.3, -0.1]

    ridge_paths = []
    lasso_paths = []
    for coef in true_coefs:
        ridge = [coef / (1 + a * 0.5) for a in alphas]
        lasso = [max(0, abs(coef) - a * 0.02) * np.sign(coef) for a in alphas]
        ridge_paths.append(ridge)
        lasso_paths.append(lasso)

    colors = ["#6C63FF", "#FF6584", "#34d399", "#fbbf24", "#a78bfa"]
    data = []
    for i, (r, l) in enumerate(zip(ridge_paths, lasso_paths)):
        data.append({"x": alphas.tolist(), "y": r, "mode": "lines", "type": "scatter",
                     "name": f"Ridge β{i+1}", "line": {"color": colors[i], "width": 2.5},
                     "legendgroup": f"f{i}"})
        data.append({"x": alphas.tolist(), "y": l, "mode": "lines", "type": "scatter",
                     "name": f"Lasso β{i+1}", "line": {"color": colors[i], "width": 2.5, "dash": "dash"},
                     "legendgroup": f"f{i}"})

    layout = {
        "template": "plotly_dark", "paper_bgcolor": "rgba(0,0,0,0)", "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": "Coefficient Paths: Ridge (solid) vs Lasso (dashed)", "font": {"size": 18}},
        "xaxis": {"title": "Regularization Strength (α)", "type": "log", "gridcolor": "rgba(255,255,255,0.08)"},
        "yaxis": {"title": "Coefficient Value", "gridcolor": "rgba(255,255,255,0.08)"},
        "margin": {"l": 60, "r": 30, "t": 50, "b": 50},
    }
    return json.dumps({"data": data, "layout": layout})


def _l1_vs_l2_geometry():
    """Show L1 (diamond) vs L2 (circle) constraint regions."""
    t = np.linspace(0, 2 * np.pi, 200)
    # L2 constraint (circle)
    l2_x = np.cos(t).tolist()
    l2_y = np.sin(t).tolist()
    # L1 constraint (diamond)
    l1_x = (np.abs(np.cos(t)) * np.cos(t) / np.abs(np.cos(t) + 1e-10)).tolist()
    n = 50
    l1_x = [np.cos(a) * 1 / (abs(np.cos(a)) + abs(np.sin(a))) for a in np.linspace(0, 2*np.pi, n)]
    l1_y = [np.sin(a) * 1 / (abs(np.cos(a)) + abs(np.sin(a))) for a in np.linspace(0, 2*np.pi, n)]

    # Contour ellipses of loss function
    loss_t = np.linspace(0, 2 * np.pi, 100)
    for r in [0.5, 1.0, 1.5]:
        pass

    data = [
        {"x": l2_x, "y": l2_y, "mode": "lines", "type": "scatter", "name": "L2 Constraint (Ridge)",
         "line": {"color": "#6C63FF", "width": 3}, "fill": "toself", "fillcolor": "rgba(108,99,255,0.1)"},
        {"x": l1_x, "y": l1_y, "mode": "lines", "type": "scatter", "name": "L1 Constraint (Lasso)",
         "line": {"color": "#FF6584", "width": 3}, "fill": "toself", "fillcolor": "rgba(255,101,132,0.1)"},
        {"x": [0, 0.7], "y": [1, 0.5], "mode": "markers+text", "type": "scatter", "name": "Optimal Points",
         "text": ["L1 hits axis (sparse!)", "L2 (non-zero)"], "textposition": "top right",
         "textfont": {"color": "#e0e0e0", "size": 11},
         "marker": {"size": 12, "color": "#fbbf24", "symbol": "diamond"}},
    ]
    layout = {
        "template": "plotly_dark", "paper_bgcolor": "rgba(0,0,0,0)", "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": "L1 vs L2 Constraint Geometry", "font": {"size": 18}},
        "xaxis": {"title": "β₁", "gridcolor": "rgba(255,255,255,0.08)", "range": [-1.5, 1.5], "scaleanchor": "y"},
        "yaxis": {"title": "β₂", "gridcolor": "rgba(255,255,255,0.08)", "range": [-1.5, 1.5]},
        "margin": {"l": 60, "r": 30, "t": 50, "b": 50},
    }
    return json.dumps({"data": data, "layout": layout})


def _mse_vs_alpha():
    """Show training vs validation MSE as alpha changes."""
    alphas = np.logspace(-3, 3, 60)
    train_mse = [0.5 + 0.3 * (1 - np.exp(-0.001 * a)) + 0.02 * np.log10(a + 1) for a in alphas]
    val_mse = [0.8 * np.exp(-0.01 * a) + 0.05 + 0.03 * np.log10(a + 1) + 0.15 * (a > 50) * np.log10(a) for a in alphas]

    data = [
        {"x": alphas.tolist(), "y": train_mse, "mode": "lines", "type": "scatter", "name": "Training MSE",
         "line": {"color": "#34d399", "width": 3}},
        {"x": alphas.tolist(), "y": val_mse, "mode": "lines", "type": "scatter", "name": "Validation MSE",
         "line": {"color": "#FF6584", "width": 3}},
    ]
    layout = {
        "template": "plotly_dark", "paper_bgcolor": "rgba(0,0,0,0)", "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": "MSE vs Regularization Strength", "font": {"size": 18}},
        "xaxis": {"title": "Alpha (α)", "type": "log", "gridcolor": "rgba(255,255,255,0.08)"},
        "yaxis": {"title": "MSE", "gridcolor": "rgba(255,255,255,0.08)"},
        "margin": {"l": 60, "r": 30, "t": 50, "b": 50},
    }
    return json.dumps({"data": data, "layout": layout})


def get_content() -> dict:
    return {
        "title": "Regularization",
        "icon": "🛡️",
        "subtitle": "Preventing overfitting by penalizing model complexity",
        "sections": [
            {"id": "introduction", "heading": "What is Regularization?", "type": "text",
             "content": ("Regularization adds a <strong>penalty term</strong> to the loss function that discourages "
                         "large coefficient values. This constrains the model from becoming too complex and "
                         "<strong>prevents overfitting</strong>.\n\n"
                         "The two most common forms are <strong>L1 (Lasso)</strong> and <strong>L2 (Ridge)</strong> regularization.")},
            {"id": "math", "heading": "The Mathematics", "type": "math",
             "content": [
                 {"label": "Ridge Regression (L2)", "formula": "J(\\theta) = \\text{MSE} + \\alpha \\sum_{j=1}^{p} \\theta_j^2",
                  "explanation": "Adds the sum of squared coefficients. Shrinks all coefficients toward zero but never exactly to zero."},
                 {"label": "Lasso Regression (L1)", "formula": "J(\\theta) = \\text{MSE} + \\alpha \\sum_{j=1}^{p} |\\theta_j|",
                  "explanation": "Adds the sum of absolute coefficients. Can shrink coefficients to exactly zero — performing feature selection."},
                 {"label": "Elastic Net", "formula": "J(\\theta) = \\text{MSE} + \\alpha \\left[ \\rho \\sum |\\theta_j| + \\frac{1-\\rho}{2} \\sum \\theta_j^2 \\right]",
                  "explanation": "Combines L1 and L2. The mixing ratio ρ controls the balance. Best when features are correlated."},
             ]},
            {"id": "geometry", "heading": "Why Lasso Produces Sparse Models", "type": "graph",
             "description": "The L1 constraint is a diamond shape — the loss contours are more likely to hit at a corner (axis), forcing coefficients to zero. The L2 circle doesn't have corners, so coefficients stay non-zero.",
             "graph_json": _l1_vs_l2_geometry()},
            {"id": "paths", "heading": "Coefficient Shrinkage Paths", "type": "graph",
             "description": "As regularization strength increases, Ridge (solid) shrinks all coefficients smoothly. Lasso (dashed) drives weaker features to exactly zero — automatic feature selection!",
             "graph_json": _coefficient_paths()},
            {"id": "tuning", "heading": "Choosing Alpha", "type": "graph",
             "description": "Too little regularization (small α) → overfitting. Too much (large α) → underfitting. Use cross-validation (RidgeCV / LassoCV) to find the optimal α.",
             "graph_json": _mse_vs_alpha()},
            {"id": "comparison", "heading": "When to Use What", "type": "list",
             "entries": [
                 {"title": "Ridge (L2)", "detail": "When all features are potentially relevant and you want to reduce their impact. Good for multicollinearity."},
                 {"title": "Lasso (L1)", "detail": "When you suspect many features are irrelevant and want automatic feature selection. Produces sparse models."},
                 {"title": "Elastic Net", "detail": "When features are correlated AND you want feature selection. Combines the best of both Ridge and Lasso."},
             ]},
            {"id": "implementation", "heading": "Python Implementation", "type": "code", "language": "python",
             "content": ("from sklearn.linear_model import Ridge, Lasso, ElasticNet, RidgeCV, LassoCV\n"
                         "from sklearn.model_selection import train_test_split\nfrom sklearn.datasets import make_regression\n"
                         "from sklearn.metrics import mean_squared_error\nimport numpy as np\n\n"
                         "# Generate data with some irrelevant features\nX, y = make_regression(n_samples=200, n_features=20,\n"
                         "                        n_informative=5, noise=10, random_state=42)\n"
                         "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)\n\n"
                         "# Ridge with cross-validation\nridge = RidgeCV(alphas=np.logspace(-3, 3, 50))\nridge.fit(X_train, y_train)\n"
                         "print(f'Ridge: alpha={ridge.alpha_:.3f}, MSE={mean_squared_error(y_test, ridge.predict(X_test)):.2f}')\n\n"
                         "# Lasso with cross-validation\nlasso = LassoCV(alphas=np.logspace(-3, 3, 50), cv=5)\nlasso.fit(X_train, y_train)\n"
                         "print(f'Lasso: alpha={lasso.alpha_:.3f}, MSE={mean_squared_error(y_test, lasso.predict(X_test)):.2f}')\n"
                         "print(f'  Non-zero coefficients: {np.sum(lasso.coef_ != 0)} / {len(lasso.coef_)}')\n")},
            {"id": "tips", "heading": "Key Takeaways", "type": "list",
             "entries": [
                 {"title": "Always Standardize", "detail": "Regularization penalizes coefficient magnitude. If features aren't scaled, it penalizes them unfairly."},
                 {"title": "Use CV Variants", "detail": "RidgeCV, LassoCV, and ElasticNetCV automatically find the best alpha via cross-validation."},
                 {"title": "Don't Regularize the Intercept", "detail": "sklearn excludes the intercept from regularization by default. This is correct — the intercept shouldn't be penalized."},
                 {"title": "Regularization is Everywhere", "detail": "Dropout in neural networks, max_depth in trees, C in SVM — these are all forms of regularization."},
             ]},
        ],
    }
