"""Regression analysis plotting functions."""

from typing import cast

import numpy as np
import pandas as pd
import plotly.graph_objects as go

from entropice.dashboard.utils.colors import get_palette


def plot_regression_scatter(
    y_true: np.ndarray | pd.Series,
    y_pred: np.ndarray | pd.Series,
    title: str = "True vs Predicted",
) -> go.Figure:
    """Create scatter plot of true vs predicted values for regression.

    Args:
        y_true: True target values.
        y_pred: Predicted target values.
        title: Title for the plot.

    Returns:
        Plotly figure with regression scatter plot.

    """
    # Convert to numpy arrays if needed
    y_true_np = cast(np.ndarray, y_true.to_numpy()) if isinstance(y_true, pd.Series) else y_true
    y_pred_np = cast(np.ndarray, y_pred.to_numpy()) if isinstance(y_pred, pd.Series) else y_pred

    # Calculate metrics
    mse = np.mean((y_true_np - y_pred_np) ** 2)
    mae = np.mean(np.abs(y_true_np - y_pred_np))
    r2 = 1 - (np.sum((y_true_np - y_pred_np) ** 2) / np.sum((y_true_np - np.mean(y_true_np)) ** 2))

    # Get colormap
    hex_colors = get_palette("r2", n_colors=256)

    # Calculate point density for coloring
    from scipy.stats import gaussian_kde

    try:
        # Create KDE for density estimation
        xy = np.vstack([y_true_np, y_pred_np])
        kde = gaussian_kde(xy)
        density = kde(xy)
    except (np.linalg.LinAlgError, ValueError):
        # Fallback if KDE fails (e.g., all points identical)
        density = np.ones(len(y_true_np))

    # Create figure
    fig = go.Figure()

    # Add scatter plot
    fig.add_trace(
        go.Scatter(
            x=y_true_np,
            y=y_pred_np,
            mode="markers",
            marker={
                "size": 6,
                "color": density,
                "colorscale": [[i / 255, c] for i, c in enumerate(hex_colors)],
                "showscale": False,
                "opacity": 0.6,
            },
            text=[f"True: {true:.3f}<br>Pred: {pred:.3f}" for true, pred in zip(y_true_np, y_pred_np)],
            hovertemplate="%{text}<extra></extra>",
            name="Data",
        )
    )

    # Add diagonal line (perfect prediction)
    min_val = min(y_true_np.min(), y_pred_np.min())
    max_val = max(y_true_np.max(), y_pred_np.max())
    fig.add_trace(
        go.Scatter(
            x=[min_val, max_val],
            y=[min_val, max_val],
            mode="lines",
            line={"color": "red", "dash": "dash", "width": 2},
            name="Perfect Prediction",
            hovertemplate="y = x<extra></extra>",
        )
    )

    # Add metrics as annotation
    metrics_text = f"R² = {r2:.4f}<br>MSE = {mse:.4f}<br>MAE = {mae:.4f}"

    fig.add_annotation(
        x=0.02,
        y=0.98,
        xref="paper",
        yref="paper",
        text=metrics_text,
        showarrow=False,
        bgcolor="white",
        bordercolor="black",
        borderwidth=1,
        xanchor="left",
        yanchor="top",
        font={"size": 12},
    )

    fig.update_layout(
        title=title,
        xaxis_title="True Values",
        yaxis_title="Predicted Values",
        height=500,
        showlegend=True,
        legend={"x": 0.98, "y": 0.02, "xanchor": "right", "yanchor": "bottom"},
    )

    # Make axes equal
    fig.update_xaxes(scaleanchor="y", scaleratio=1)

    return fig


def plot_residuals(
    y_true: np.ndarray | pd.Series,
    y_pred: np.ndarray | pd.Series,
    title: str = "Residual Plot",
) -> go.Figure:
    """Create residual plot for regression diagnostics.

    Args:
        y_true: True target values.
        y_pred: Predicted target values.
        title: Title for the plot.

    Returns:
        Plotly figure with residual plot.

    """
    # Convert to numpy arrays if needed
    y_true_np = cast(np.ndarray, y_true.to_numpy()) if isinstance(y_true, pd.Series) else y_true
    y_pred_np = cast(np.ndarray, y_pred.to_numpy()) if isinstance(y_pred, pd.Series) else y_pred

    # Calculate residuals
    residuals = y_true_np - y_pred_np

    # Get colormap
    hex_colors = get_palette("r2", n_colors=256)

    # Create figure
    fig = go.Figure()

    # Add scatter plot
    fig.add_trace(
        go.Scatter(
            x=y_pred,
            y=residuals,
            mode="markers",
            marker={
                "size": 6,
                "color": np.abs(residuals),
                "colorscale": [[i / 255, c] for i, c in enumerate(hex_colors)],
                "showscale": True,
                "colorbar": {"title": "Abs Residual"},
                "opacity": 0.6,
            },
            text=[f"Pred: {pred:.3f}<br>Residual: {res:.3f}" for pred, res in zip(y_pred, residuals)],
            hovertemplate="%{text}<extra></extra>",
        )
    )

    # Add zero line
    fig.add_hline(y=0, line_dash="dash", line_color="red", line_width=2)

    fig.update_layout(
        title=title,
        xaxis_title="Predicted Values",
        yaxis_title="Residuals (True - Predicted)",
        height=400,
        showlegend=False,
    )

    return fig
