"""Hyperparameter space plotting functions."""

import matplotlib.colors as mcolors
import pandas as pd
import plotly.graph_objects as go

from entropice.dashboard.utils.colors import get_cmap, get_palette


def plot_performance_summary(results: pd.DataFrame, refit_metric: str) -> tuple[pd.DataFrame, pd.DataFrame, dict]:
    """Compute performance summary statistics.

    Args:
        results: DataFrame with CV results.
        refit_metric: The metric used for refit (e.g., 'f1', 'f1_weighted').

    Returns:
        Tuple of (best_scores_df, score_stats_df, best_params_dict).

    """
    # Get all test score columns
    score_cols = [col for col in results.columns if col.startswith("mean_test_")]

    if not score_cols:
        return pd.DataFrame(), pd.DataFrame(), {}

    # Calculate best scores
    best_scores = []
    for col in score_cols:
        metric_name = col.replace("mean_test_", "").replace("_", " ").title()
        best_score = results[col].max()
        best_scores.append({"Metric": metric_name, "Best Score": f"{best_score:.4f}"})

    # Calculate score statistics
    score_stats = []
    for col in score_cols:
        metric_name = col.replace("mean_test_", "").replace("_", " ").title()
        mean_score = results[col].mean()
        std_score = results[col].std()
        score_stats.append(
            {
                "Metric": metric_name,
                "Mean ± Std": f"{mean_score:.4f} ± {std_score:.4f}",
            }
        )

    # Get best parameter combination
    refit_col = f"mean_test_{refit_metric}"
    if refit_col not in results.columns and score_cols:
        refit_col = score_cols[0]

    best_idx = results[refit_col].idxmax()
    best_row = results.loc[best_idx]

    # Extract parameter columns
    param_cols = [col for col in results.columns if col.startswith("param_") and col != "params"]
    best_params = {col.replace("param_", ""): best_row[col] for col in param_cols}

    return pd.DataFrame(best_scores), pd.DataFrame(score_stats), best_params


def plot_parameter_distributions(results: pd.DataFrame, param_grid: dict | None = None) -> dict[str, go.Figure]:
    """Create histogram charts for parameter distributions.

    Args:
        results: DataFrame with CV results.
        param_grid: Optional parameter grid with distribution information.

    Returns:
        Dictionary mapping parameter names to Plotly figures.

    """
    # Get parameter columns
    param_cols = [col for col in results.columns if col.startswith("param_") and col != "params"]

    if not param_cols:
        return {}

    cmap = get_cmap("parameter_distribution")
    bar_color = mcolors.rgb2hex(cmap(0.5))

    charts = {}
    for param_col in param_cols:
        param_name = param_col.replace("param_", "")
        param_values = results[param_col].dropna()

        if len(param_values) == 0:
            continue

        # Determine if parameter is numeric
        if pd.api.types.is_numeric_dtype(param_values):
            # Create histogram for numeric parameters
            fig = go.Figure()
            fig.add_trace(
                go.Histogram(
                    x=param_values,
                    nbinsx=30,
                    marker_color=bar_color,
                    name=param_name,
                )
            )
            fig.update_layout(
                title=f"Distribution of {param_name}",
                xaxis_title=param_name,
                yaxis_title="Count",
                height=400,
                showlegend=False,
            )
        else:
            # Create bar chart for categorical parameters
            value_counts = param_values.value_counts().reset_index()
            value_counts.columns = [param_name, "count"]
            fig = go.Figure()
            fig.add_trace(
                go.Bar(
                    x=value_counts[param_name],
                    y=value_counts["count"],
                    marker_color=bar_color,
                    name=param_name,
                )
            )
            fig.update_layout(
                title=f"Distribution of {param_name}",
                xaxis_title=param_name,
                yaxis_title="Count",
                height=400,
                showlegend=False,
            )

        charts[param_name] = fig

    return charts


def plot_score_vs_parameters(
    results: pd.DataFrame, metric: str, param_grid: dict | None = None
) -> dict[str, go.Figure]:
    """Create scatter plots of score vs each parameter.

    Args:
        results: DataFrame with CV results.
        metric: The metric to plot (e.g., 'f1', 'accuracy').
        param_grid: Optional parameter grid with distribution information.

    Returns:
        Dictionary mapping parameter names to Plotly figures.

    """
    score_col = f"mean_test_{metric}"
    if score_col not in results.columns:
        return {}

    # Get parameter columns
    param_cols = [col for col in results.columns if col.startswith("param_") and col != "params"]

    if not param_cols:
        return {}

    # Get colormap
    hex_colors = get_palette(metric, n_colors=256)

    charts = {}
    for param_col in param_cols:
        param_name = param_col.replace("param_", "")
        param_values = results[param_col].dropna()

        if len(param_values) == 0:
            continue

        # Check if this parameter uses log scale
        use_log = False
        if param_grid and param_name in param_grid:
            param_config = param_grid[param_name]
            if isinstance(param_config, dict) and param_config.get("distribution") == "loguniform":
                use_log = True

        # Create scatter plot
        fig = go.Figure()
        fig.add_trace(
            go.Scatter(
                x=results[param_col],
                y=results[score_col],
                mode="markers",
                marker={
                    "size": 8,
                    "color": results[score_col],
                    "colorscale": [[i / 255, c] for i, c in enumerate(hex_colors)],
                    "showscale": False,
                    "opacity": 0.6,
                },
                text=[
                    f"{param_name}: {val}<br>Score: {score:.4f}"
                    for val, score in zip(results[param_col], results[score_col])
                ],
                hovertemplate="%{text}<extra></extra>",
            )
        )
        fig.update_layout(
            title=f"{metric.replace('_', ' ').title()} vs {param_name}",
            xaxis_title=param_name,
            xaxis_type="log" if use_log else "linear",
            yaxis_title=metric.replace("_", " ").title(),
            height=400,
            showlegend=False,
        )

        charts[param_name] = fig

    return charts


def plot_parameter_correlations(results: pd.DataFrame, metric: str) -> go.Figure | None:
    """Create correlation bar chart between parameters and score.

    Args:
        results: DataFrame with CV results.
        metric: The metric to analyze (e.g., 'f1', 'accuracy').

    Returns:
        Plotly figure or None if no numeric parameters found.

    """
    score_col = f"mean_test_{metric}"
    if score_col not in results.columns:
        return None

    # Get numeric parameter columns
    param_cols = [col for col in results.columns if col.startswith("param_") and col != "params"]
    numeric_params = [col for col in param_cols if pd.api.types.is_numeric_dtype(results[col])]

    if not numeric_params:
        return None

    # Calculate correlations
    correlations = []
    for param_col in numeric_params:
        param_name = param_col.replace("param_", "")
        corr = results[[param_col, score_col]].corr().iloc[0, 1]
        correlations.append({"Parameter": param_name, "Correlation": corr})

    corr_df = pd.DataFrame(correlations).sort_values("Correlation", ascending=False)

    # Get colormap (use diverging colormap for correlation)
    hex_colors = get_palette("correlation", n_colors=256)

    # Create bar chart
    fig = go.Figure()
    fig.add_trace(
        go.Bar(
            x=corr_df["Correlation"],
            y=corr_df["Parameter"],
            orientation="h",
            marker={
                "color": corr_df["Correlation"],
                "colorscale": [[i / 255, c] for i, c in enumerate(hex_colors)],
                "cmin": -1,
                "cmax": 1,
                "showscale": False,
            },
            text=[f"{c:.3f}" for c in corr_df["Correlation"]],
            hovertemplate="%{y}<br>Correlation: %{x:.3f}<extra></extra>",
        )
    )
    fig.update_layout(
        xaxis_title="Correlation with Score",
        yaxis_title="Parameter",
        height=max(300, len(correlations) * 30),
        showlegend=False,
    )

    return fig


def plot_parameter_interactions(results: pd.DataFrame, metric: str, param_grid: dict | None = None) -> list[go.Figure]:
    """Create scatter plots showing parameter interactions.

    Args:
        results: DataFrame with CV results.
        metric: The metric to visualize (e.g., 'f1', 'accuracy').
        param_grid: Optional parameter grid with distribution information.

    Returns:
        List of Plotly figures showing parameter interactions.

    """
    score_col = f"mean_test_{metric}"
    if score_col not in results.columns:
        return []

    # Get numeric parameter columns
    param_cols = [col for col in results.columns if col.startswith("param_") and col != "params"]
    numeric_params = [col for col in param_cols if pd.api.types.is_numeric_dtype(results[col])]

    if len(numeric_params) < 2:
        return []

    # Get colormap
    hex_colors = get_palette(metric, n_colors=256)

    # Create scatter plots for parameter pairs
    charts = []
    param_names = [col.replace("param_", "") for col in numeric_params]

    for i, x_param in enumerate(param_names[:-1]):
        for y_param in param_names[i + 1 :]:
            x_col = f"param_{x_param}"
            y_col = f"param_{y_param}"

            # Check if parameters use log scale
            x_use_log = False
            y_use_log = False
            if param_grid:
                if x_param in param_grid:
                    x_config = param_grid[x_param]
                    if isinstance(x_config, dict) and x_config.get("distribution") == "loguniform":
                        x_use_log = True
                if y_param in param_grid:
                    y_config = param_grid[y_param]
                    if isinstance(y_config, dict) and y_config.get("distribution") == "loguniform":
                        y_use_log = True

            fig = go.Figure()
            fig.add_trace(
                go.Scatter(
                    x=results[x_col],
                    y=results[y_col],
                    mode="markers",
                    marker={
                        "size": 8,
                        "color": results[score_col],
                        "colorscale": [[i / 255, c] for i, c in enumerate(hex_colors)],
                        "showscale": True,
                        "colorbar": {"title": metric.replace("_", " ").title()},
                        "opacity": 0.7,
                    },
                    text=[
                        f"{x_param}: {x_val}<br>{y_param}: {y_val}<br>Score: {score:.4f}"
                        for x_val, y_val, score in zip(results[x_col], results[y_col], results[score_col])
                    ],
                    hovertemplate="%{text}<extra></extra>",
                )
            )
            fig.update_layout(
                title=f"{metric.replace('_', ' ').title()} by {x_param} and {y_param}",
                xaxis_title=x_param,
                xaxis_type="log" if x_use_log else "linear",
                yaxis_title=y_param,
                yaxis_type="log" if y_use_log else "linear",
                height=500,
                width=500,
            )

            charts.append(fig)

    return charts


def plot_score_evolution(results: pd.DataFrame, metric: str) -> go.Figure | None:
    """Create line chart showing score evolution over iterations.

    Args:
        results: DataFrame with CV results.
        metric: The metric to visualize (e.g., 'f1', 'accuracy').

    Returns:
        Plotly figure or None if metric not found.

    """
    score_col = f"mean_test_{metric}"
    if score_col not in results.columns:
        return None

    # Add iteration number
    iterations = list(range(len(results)))
    scores = results[score_col].to_numpy()
    best_so_far = results[score_col].cummax().to_numpy()

    # Get colormap
    cmap = get_cmap("score_evolution")
    score_color = mcolors.rgb2hex(cmap(0.3))
    best_color = mcolors.rgb2hex(cmap(0.7))

    # Create line chart
    fig = go.Figure()

    fig.add_trace(
        go.Scatter(
            x=iterations,
            y=scores,
            mode="lines",
            name="Score",
            line={"color": score_color, "width": 1},
            opacity=0.6,
            hovertemplate="Iteration: %{x}<br>Score: %{y:.4f}<extra></extra>",
        )
    )

    fig.add_trace(
        go.Scatter(
            x=iterations,
            y=best_so_far,
            mode="lines",
            name="Best So Far",
            line={"color": best_color, "width": 2},
            hovertemplate="Iteration: %{x}<br>Best So Far: %{y:.4f}<extra></extra>",
        )
    )

    fig.update_layout(
        title=f"{metric.replace('_', ' ').title()} Evolution",
        xaxis_title="Iteration",
        yaxis_title=metric.replace("_", " ").title(),
        height=300,
        hovermode="x unified",
    )

    return fig
