"""Plots for cross-dataset correlation and similarity analysis."""

from typing import Literal

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import seaborn as sns
from plotly.subplots import make_subplots
from scipy.cluster.hierarchy import dendrogram, linkage
from scipy.spatial.distance import squareform
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler


def select_top_variable_features(
    data_dict: dict[str, pd.Series],
    n_features: int = 500,
    method: Literal["variance", "iqr", "cv"] = "variance",
) -> dict[str, pd.Series]:
    """Select the top N most variable features from a dataset.

    Args:
        data_dict: Dictionary mapping variable names to pandas Series
        n_features: Number of features to select
        method: Method to measure variability ('variance', 'iqr', or 'cv')

    Returns:
        Filtered dictionary with only the most variable features

    """
    if len(data_dict) <= n_features:
        return data_dict

    # Calculate variability metric for each feature
    variability_scores = {}
    for var_name, values in data_dict.items():
        clean_values = values.dropna()
        if len(clean_values) == 0:
            variability_scores[var_name] = 0
            continue

        if method == "variance":
            variability_scores[var_name] = clean_values.var()
        elif method == "iqr":
            variability_scores[var_name] = clean_values.quantile(0.75) - clean_values.quantile(0.25)
        elif method == "cv":
            mean_val = clean_values.mean()
            if abs(mean_val) > 1e-10:
                variability_scores[var_name] = clean_values.std() / abs(mean_val)
            else:
                variability_scores[var_name] = 0

    # Sort by variability and select top N
    sorted_vars = sorted(variability_scores.items(), key=lambda x: x[1], reverse=True)
    top_vars = [var_name for var_name, _ in sorted_vars[:n_features]]

    return {k: v for k, v in data_dict.items() if k in top_vars}


def create_matplotlib_correlation_heatmap(
    data_dict: dict[str, pd.Series],
    method: Literal["pearson", "kendall", "spearman"] = "pearson",
    cluster: bool = False,
    max_labels: int = 100,
) -> plt.Figure:
    """Create a correlation heatmap using matplotlib for large datasets.

    Args:
        data_dict: Dictionary mapping variable names to pandas Series with cell_ids as index
        method: Correlation method ('pearson', 'spearman', or 'kendall')
        cluster: Whether to reorder variables by hierarchical clustering
        max_labels: Maximum number of labels to show on axes

    Returns:
        Matplotlib Figure with correlation heatmap

    """
    if len(data_dict) < 2:
        fig, ax = plt.subplots(figsize=(8, 6))
        ax.text(0.5, 0.5, "Need at least 2 variables for correlation analysis", ha="center", va="center", fontsize=12)
        ax.axis("off")
        return fig

    # Create DataFrame from all variables
    df = pd.DataFrame(data_dict)

    # Drop rows with any NaN to get valid correlations
    df_clean = df.dropna()

    if len(df_clean) == 0:
        fig, ax = plt.subplots(figsize=(8, 6))
        ax.text(0.5, 0.5, "No overlapping non-null data between variables", ha="center", va="center", fontsize=12)
        ax.axis("off")
        return fig

    # Subsample if too many rows for performance
    if len(df_clean) > 50000:
        df_clean = df_clean.sample(n=50000, random_state=42)

    # Calculate correlation matrix
    corr = df_clean.corr(method=method)

    # Apply hierarchical clustering if requested
    if cluster and len(corr) > 2:
        # Use correlation distance for clustering
        corr_dist = 1 - np.abs(corr.values)
        np.fill_diagonal(corr_dist, 0)
        condensed_dist = squareform(corr_dist, checks=False)
        linkage_matrix = linkage(condensed_dist, method="average")
        dendro = dendrogram(linkage_matrix, no_plot=True)
        cluster_order = dendro["leaves"]

        # Reorder correlation matrix
        corr = corr.iloc[cluster_order, cluster_order]

    # Determine figure size based on number of variables
    n_vars = len(corr)
    fig_size = max(10, min(n_vars * 0.3, 50))

    # Create figure
    fig, ax = plt.subplots(figsize=(fig_size, fig_size))

    # Create heatmap using seaborn
    sns.heatmap(
        corr,
        cmap="RdBu_r",
        center=0,
        vmin=-1,
        vmax=1,
        square=True,
        linewidths=0.5 if n_vars < 50 else 0,
        cbar_kws={"shrink": 0.8, "label": f"{method.title()} Correlation"},
        ax=ax,
        xticklabels=n_vars <= max_labels,
        yticklabels=n_vars <= max_labels,
    )

    # Set title
    title = f"Correlation Matrix ({method.title()}, {n_vars} variables)"
    if cluster:
        title += " - Hierarchically Clustered"
    ax.set_title(title, fontsize=14, pad=20)

    # Rotate labels if shown
    if n_vars <= max_labels:
        plt.setp(ax.get_xticklabels(), rotation=45, ha="right", fontsize=8)
        plt.setp(ax.get_yticklabels(), rotation=0, fontsize=8)

    plt.tight_layout()
    return fig


def create_full_correlation_heatmap(
    data_dict: dict[str, pd.Series],
    method: Literal["pearson", "kendall", "spearman"] = "pearson",
    cluster: bool = False,
) -> go.Figure:
    """Create a correlation heatmap for all variables across all datasets.

    Args:
        data_dict: Dictionary mapping variable names to pandas Series with cell_ids as index
        method: Correlation method ('pearson', 'spearman', or 'kendall')
        cluster: Whether to reorder variables by hierarchical clustering

    Returns:
        Plotly Figure with correlation heatmap

    """
    if len(data_dict) < 2:
        return go.Figure().add_annotation(
            text="Need at least 2 variables for correlation analysis",
            xref="paper",
            yref="paper",
            x=0.5,
            y=0.5,
            showarrow=False,
        )

    # Create DataFrame from all variables
    df = pd.DataFrame(data_dict)

    # Drop rows with any NaN to get valid correlations
    df_clean = df.dropna()

    if len(df_clean) == 0:
        return go.Figure().add_annotation(
            text="No overlapping non-null data between variables",
            xref="paper",
            yref="paper",
            x=0.5,
            y=0.5,
            showarrow=False,
        )

    # Subsample if too many rows for performance
    if len(df_clean) > 50000:
        df_clean = df_clean.sample(n=50000, random_state=42)

    # Calculate correlation matrix
    corr = df_clean.corr(method=method)

    # Apply hierarchical clustering if requested
    if cluster and len(corr) > 2:
        # Use correlation distance for clustering
        corr_dist = 1 - np.abs(corr.values)
        np.fill_diagonal(corr_dist, 0)  # Ensure diagonal is 0
        condensed_dist = squareform(corr_dist, checks=False)
        linkage_matrix = linkage(condensed_dist, method="average")
        dendro = dendrogram(linkage_matrix, no_plot=True)
        cluster_order = dendro["leaves"]

        # Reorder correlation matrix
        corr = corr.iloc[cluster_order, cluster_order]

    # Create heatmap
    fig = go.Figure(
        data=go.Heatmap(
            z=corr.values,
            x=corr.columns,
            y=corr.index,
            colorscale="RdBu_r",
            zmid=0,
            zmin=-1,
            zmax=1,
            text=np.round(corr.values, 2),
            texttemplate="%{text}",
            textfont={"size": 8},
            colorbar={"title": f"{method.title()}<br>Correlation"},
            hovertemplate=f"<b>%{{y}}</b> vs <b>%{{x}}</b><br>{method.title()} correlation: %{{z:.3f}}<extra></extra>",
        )
    )

    title = f"Cross-Dataset Variable Correlations ({method.title()})"
    if cluster:
        title += " - Hierarchically Clustered"

    fig.update_layout(
        title=title,
        height=max(600, len(corr) * 15),
        width=max(800, len(corr) * 15),
        xaxis={"side": "bottom", "tickangle": 45},
        yaxis={"tickangle": 0},
    )

    return fig


def create_pca_biplot(
    data_dict: dict[str, pd.Series],
    n_components: int = 2,
    show_loadings: bool = True,
) -> go.Figure:
    """Create a PCA biplot showing feature relationships in principal component space.

    Args:
        data_dict: Dictionary mapping variable names to pandas Series with cell_ids as index
        n_components: Number of principal components to compute
        show_loadings: Whether to show loading vectors

    Returns:
        Plotly Figure with PCA biplot

    """
    if len(data_dict) < 2:
        return go.Figure().add_annotation(
            text="Need at least 2 variables for PCA analysis",
            xref="paper",
            yref="paper",
            x=0.5,
            y=0.5,
            showarrow=False,
        )

    # Create DataFrame and clean
    df = pd.DataFrame(data_dict).dropna()

    if len(df) < 10:
        return go.Figure().add_annotation(
            text="Not enough overlapping data for PCA analysis",
            xref="paper",
            yref="paper",
            x=0.5,
            y=0.5,
            showarrow=False,
        )

    # Subsample if needed
    if len(df) > 50000:
        df = df.sample(n=50000, random_state=42)

    # Standardize features
    scaler = StandardScaler()
    scaled_data = scaler.fit_transform(df)

    # Perform PCA
    n_components = min(n_components, len(data_dict), len(df))
    pca = PCA(n_components=n_components)
    pca_result = pca.fit_transform(scaled_data)

    # Create figure
    fig = go.Figure()

    if n_components >= 2:
        # Scatter plot of observations (subsampled for visibility)
        sample_size = min(5000, len(pca_result))
        rng = np.random.default_rng(42)
        indices = rng.choice(len(pca_result), sample_size, replace=False)

        fig.add_trace(
            go.Scatter(
                x=pca_result[indices, 0],
                y=pca_result[indices, 1],
                mode="markers",
                marker={"size": 3, "opacity": 0.3, "color": "lightblue"},
                name="Observations",
                showlegend=True,
            )
        )

        # Add loading vectors if requested
        if show_loadings:
            loadings = pca.components_.T * np.sqrt(pca.explained_variance_)

            # Scale loadings for visibility
            max_loading = np.abs(loadings[:, :2]).max()
            max_data = max(np.abs(pca_result[:, :2]).max(axis=0))
            scale = max_data / max_loading * 0.8

            for i, var_name in enumerate(df.columns):
                fig.add_trace(
                    go.Scatter(
                        x=[0, loadings[i, 0] * scale],
                        y=[0, loadings[i, 1] * scale],
                        mode="lines+text",
                        line={"color": "red", "width": 2},
                        text=["", var_name],
                        textposition="top center",
                        textfont={"size": 10, "color": "darkred"},
                        name=var_name,
                        showlegend=False,
                        hovertemplate=f"<b>{var_name}</b><br>"
                        + f"PC1 loading: {loadings[i, 0]:.3f}<br>"
                        + f"PC2 loading: {loadings[i, 1]:.3f}<extra></extra>",
                    )
                )

        var_exp = pca.explained_variance_ratio_
        fig.update_xaxes(title=f"PC1 ({var_exp[0] * 100:.1f}% variance)")
        fig.update_yaxes(title=f"PC2 ({var_exp[1] * 100:.1f}% variance)")

    fig.update_layout(
        title="PCA Biplot - Feature Relationships in Principal Component Space",
        height=700,
        showlegend=show_loadings,
        hovermode="closest",
    )

    return fig


def create_dendrogram_plot(
    data_dict: dict[str, pd.Series],
    method: Literal["pearson", "kendall", "spearman"] = "pearson",
) -> go.Figure:
    """Create a dendrogram showing hierarchical clustering of variables based on correlation.

    Args:
        data_dict: Dictionary mapping variable names to pandas Series with cell_ids as index
        method: Correlation method to use for distance calculation

    Returns:
        Plotly Figure with dendrogram

    """
    if len(data_dict) < 2:
        return go.Figure().add_annotation(
            text="Need at least 2 variables for clustering",
            xref="paper",
            yref="paper",
            x=0.5,
            y=0.5,
            showarrow=False,
        )

    # Create DataFrame and clean
    df = pd.DataFrame(data_dict).dropna()

    if len(df) < 10:
        return go.Figure().add_annotation(
            text="Not enough overlapping data for clustering",
            xref="paper",
            yref="paper",
            x=0.5,
            y=0.5,
            showarrow=False,
        )

    # Subsample if needed
    if len(df) > 50000:
        df = df.sample(n=50000, random_state=42)

    # Calculate correlation and convert to distance
    corr = df.corr(method=method)
    corr_dist = 1 - np.abs(corr.values)
    np.fill_diagonal(corr_dist, 0)

    # Perform hierarchical clustering
    condensed_dist = squareform(corr_dist, checks=False)
    linkage_matrix = linkage(condensed_dist, method="average")

    # Create dendrogram
    dendro = dendrogram(linkage_matrix, labels=list(df.columns), no_plot=True)

    # Extract dendrogram data
    icoord = np.array(dendro["icoord"])
    dcoord = np.array(dendro["dcoord"])

    # Create plotly figure
    fig = go.Figure()

    # Add dendrogram lines
    for i in range(len(icoord)):
        fig.add_trace(
            go.Scatter(
                x=icoord[i],
                y=dcoord[i],
                mode="lines",
                line={"color": "black", "width": 1.5},
                hoverinfo="skip",
                showlegend=False,
            )
        )

    # Add labels
    labels = dendro["ivl"]
    label_pos = np.arange(5, len(labels) * 10 + 5, 10)

    fig.update_layout(
        title=f"Hierarchical Clustering of Variables (based on {method.title()} correlation)",
        xaxis={
            "title": "Variables",
            "tickmode": "array",
            "tickvals": label_pos,
            "ticktext": labels,
            "tickangle": 45,
        },
        yaxis={"title": "Distance (1 - |correlation|)"},
        height=600,
        showlegend=False,
    )

    return fig


def create_mutual_information_matrix(data_dict: dict[str, pd.Series]) -> go.Figure:
    """Create a mutual information matrix to capture non-linear relationships.

    Args:
        data_dict: Dictionary mapping variable names to pandas Series with cell_ids as index

    Returns:
        Plotly Figure with mutual information heatmap

    """
    from sklearn.feature_selection import mutual_info_regression

    if len(data_dict) < 2:
        return go.Figure().add_annotation(
            text="Need at least 2 variables for mutual information analysis",
            xref="paper",
            yref="paper",
            x=0.5,
            y=0.5,
            showarrow=False,
        )

    # Create DataFrame and clean
    df = pd.DataFrame(data_dict).dropna()

    if len(df) < 100:
        return go.Figure().add_annotation(
            text="Not enough overlapping data for mutual information analysis",
            xref="paper",
            yref="paper",
            x=0.5,
            y=0.5,
            showarrow=False,
        )

    # Subsample if needed
    if len(df) > 10000:  # MI is computationally expensive
        df = df.sample(n=10000, random_state=42)

    # Calculate pairwise mutual information
    n_vars = len(df.columns)
    mi_matrix = np.zeros((n_vars, n_vars))

    for i, col_i in enumerate(df.columns):
        for j, col_j in enumerate(df.columns):
            if i == j:
                mi_matrix[i, j] = 1.0  # Self-information (normalized)
            elif i < j:
                # Calculate MI
                X = df[col_i].to_numpy().reshape(-1, 1)  # noqa: N806
                y = df[col_j].to_numpy()
                mi = mutual_info_regression(X, y, random_state=42)[0]

                # Normalize by entropy (approximate)
                mi_norm = mi / (np.std(y) * np.std(X.flatten()) + 1e-10)
                mi_matrix[i, j] = mi_norm
                mi_matrix[j, i] = mi_norm

    # Create heatmap
    fig = go.Figure(
        data=go.Heatmap(
            z=mi_matrix,
            x=list(df.columns),
            y=list(df.columns),
            colorscale="YlOrRd",
            zmin=0,
            text=np.round(mi_matrix, 3),
            texttemplate="%{text}",
            textfont={"size": 8},
            colorbar={"title": "Normalized<br>Mutual Info"},
            hovertemplate="<b>%{y}</b> vs <b>%{x}</b><br>Mutual Information: %{z:.3f}<extra></extra>",
        )
    )

    fig.update_layout(
        title="Mutual Information Matrix (captures non-linear relationships)",
        height=max(600, n_vars * 15),
        width=max(800, n_vars * 15),
        xaxis={"side": "bottom", "tickangle": 45},
        yaxis={"tickangle": 0},
    )

    return fig


def create_feature_variance_plot(data_dict: dict[str, pd.Series]) -> go.Figure:
    """Create a bar plot showing variance/spread of each feature.

    Args:
        data_dict: Dictionary mapping variable names to pandas Series with cell_ids as index

    Returns:
        Plotly Figure with variance bar plot

    """
    if len(data_dict) == 0:
        return go.Figure().add_annotation(
            text="No variables provided",
            xref="paper",
            yref="paper",
            x=0.5,
            y=0.5,
            showarrow=False,
        )

    # Calculate statistics for each variable
    stats_data = []
    for var_name, values in data_dict.items():
        clean_values = values.dropna()
        if len(clean_values) > 0:
            stats_data.append(
                {
                    "Variable": var_name,
                    "Std Dev": clean_values.std(),
                    "Variance": clean_values.var(),
                    "IQR": clean_values.quantile(0.75) - clean_values.quantile(0.25),
                    "Range": clean_values.max() - clean_values.min(),
                    "CV": clean_values.std() / (abs(clean_values.mean()) + 1e-10),
                }
            )

    stats_df = pd.DataFrame(stats_data).sort_values("Variance", ascending=False)

    # Create subplots for different variance metrics
    fig = make_subplots(
        rows=2,
        cols=2,
        subplot_titles=["Standard Deviation", "Interquartile Range (IQR)", "Range", "Coefficient of Variation"],
        vertical_spacing=0.12,
        horizontal_spacing=0.10,
    )

    metrics = [
        ("Std Dev", 1, 1),
        ("IQR", 1, 2),
        ("Range", 2, 1),
        ("CV", 2, 2),
    ]

    for metric, row, col in metrics:
        fig.add_trace(
            go.Bar(
                x=stats_df["Variable"],
                y=stats_df[metric],
                name=metric,
                marker={"color": stats_df[metric], "colorscale": "Viridis"},
                showlegend=False,
                hovertemplate=f"<b>%{{x}}</b><br>{metric}: %{{y:.3f}}<extra></extra>",
            ),
            row=row,
            col=col,
        )

        fig.update_xaxes(tickangle=45, row=row, col=col)
        fig.update_yaxes(title_text=metric, row=row, col=col)

    fig.update_layout(
        title="Feature Variance and Spread Metrics",
        height=800,
        showlegend=False,
    )

    return fig
