"""
Activation Functions Playground — interactive parameter exploration.
"""

import numpy as np

META = {
    "title": "Activation Functions",
    "icon": "⚡",
    "description": (
        "Compare different activation functions side by side. "
        "Adjust the zoom and alpha parameter (for Leaky ReLU / ELU) to "
        "see how each function transforms its input."
    ),
}

PARAMS = [
    {
        "name": "x_range",
        "label": "X Range (±)",
        "type": "slider",
        "min": 2,
        "max": 10,
        "step": 1,
        "default": 5,
    },
    {
        "name": "alpha",
        "label": "Alpha (Leaky ReLU / ELU)",
        "type": "slider",
        "min": 0.01,
        "max": 1.0,
        "step": 0.01,
        "default": 0.1,
    },
    {
        "name": "functions",
        "label": "Highlight Function",
        "type": "select",
        "options": ["All", "ReLU", "Sigmoid", "Tanh", "Leaky ReLU", "ELU", "Swish"],
        "default": "All",
    },
]


def compute(params: dict) -> dict:
    x_range = float(params.get("x_range", 5))
    alpha = float(params.get("alpha", 0.1))
    highlight = params.get("functions", "All")

    x = np.linspace(-x_range, x_range, 500)

    # Define activation functions
    activations = {
        "ReLU": np.maximum(0, x),
        "Sigmoid": 1 / (1 + np.exp(-x)),
        "Tanh": np.tanh(x),
        "Leaky ReLU": np.where(x > 0, x, alpha * x),
        "ELU": np.where(x > 0, x, alpha * (np.exp(np.clip(x, -10, 0)) - 1)),
        "Swish": x * (1 / (1 + np.exp(-x))),
    }

    colors = {
        "ReLU": "#6C63FF",
        "Sigmoid": "#FF6584",
        "Tanh": "#34d399",
        "Leaky ReLU": "#fbbf24",
        "ELU": "#a78bfa",
        "Swish": "#38bdf8",
    }

    data = []
    for name, y in activations.items():
        is_highlighted = (highlight == "All" or highlight == name)
        data.append({
            "x": x.tolist(),
            "y": y.tolist(),
            "mode": "lines",
            "type": "scatter",
            "name": name,
            "line": {
                "color": colors[name],
                "width": 3 if is_highlighted else 1.5,
            },
            "opacity": 1.0 if is_highlighted else 0.25,
        })

    # Add zero lines
    data.append({
        "x": [-x_range, x_range], "y": [0, 0],
        "mode": "lines", "type": "scatter", "name": "y=0",
        "line": {"color": "rgba(255,255,255,0.15)", "width": 1, "dash": "dash"},
        "showlegend": False,
    })
    data.append({
        "x": [0, 0], "y": [-2, x_range],
        "mode": "lines", "type": "scatter", "name": "x=0",
        "line": {"color": "rgba(255,255,255,0.15)", "width": 1, "dash": "dash"},
        "showlegend": False,
    })

    title_text = f"Activation Functions" if highlight == "All" else f"{highlight} Activation"

    layout = {
        "template": "plotly_dark",
        "paper_bgcolor": "rgba(0,0,0,0)",
        "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": title_text, "font": {"size": 20}},
        "xaxis": {"title": "Input (z)", "gridcolor": "rgba(255,255,255,0.08)", "zeroline": False},
        "yaxis": {"title": "Output f(z)", "gridcolor": "rgba(255,255,255,0.08)",
                  "range": [-2, x_range], "zeroline": False},
        "legend": {"x": 0.02, "y": 0.98, "bgcolor": "rgba(0,0,0,0.3)"},
        "margin": {"l": 60, "r": 30, "t": 60, "b": 50},
    }
    return {"data": data, "layout": layout}
