# src/plotting.py
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import pandas as pd
import numpy as np
import json
from collections import defaultdict
from typing import Dict, List, Optional, Union
from config import (
    LANGUAGE_NAMES,
    ALL_UG40_LANGUAGES,
    GOOGLE_SUPPORTED_LANGUAGES,
    METRICS_CONFIG,
    EVALUATION_TRACKS,
    MODEL_CATEGORIES,
    CHART_CONFIG,
)


def create_leaderboard_plot(
    df: pd.DataFrame, track: str, metric: str = "quality", top_n: int = 15
) -> go.Figure:
    """Create leaderboard plot with confidence intervals."""
    
    if df.empty:
        fig = go.Figure()
        fig.add_annotation(
            text="No models available for this track",
            xref="paper", yref="paper",
            x=0.5, y=0.5, showarrow=False,
            font=dict(size=16)
        )
        fig.update_layout(
            title=f"No Data Available - {track.title()} Track",
            paper_bgcolor="rgba(0,0,0,0)",
            plot_bgcolor="rgba(0,0,0,0)"
        )
        return fig
    
    # Get top N models for this track
    metric_col = f"{track}_{metric}"
    ci_lower_col = f"{track}_ci_lower"
    ci_upper_col = f"{track}_ci_upper"
    
    if metric_col not in df.columns:
        fig = go.Figure()
        fig.add_annotation(
            text=f"Metric {metric} not available for {track} track",
            xref="paper", yref="paper",
            x=0.5, y=0.5, showarrow=False,
        )
        return fig
    
    # Filter and sort
    valid_models = df[(df[metric_col] > 0)].head(top_n)
    
    if valid_models.empty:
        fig = go.Figure()
        fig.add_annotation(text="No valid models found", x=0.5, y=0.5, showarrow=False)
        return fig
    
    # Create color mapping by category
    colors = [MODEL_CATEGORIES.get(cat, {}).get("color", "#808080") for cat in valid_models["model_category"]]
    
    # Main bar plot
    fig = go.Figure()
    
    # Add bars with error bars if confidence intervals available
    error_x = None
    if ci_lower_col in valid_models.columns and ci_upper_col in valid_models.columns:
        error_x = dict(
            type="data",
            array=valid_models[ci_upper_col] - valid_models[metric_col],
            arrayminus=valid_models[metric_col] - valid_models[ci_lower_col],
            visible=True,
            thickness=2,
            width=4,
        )
    
    fig.add_trace(go.Bar(
        y=valid_models["model_name"],
        x=valid_models[metric_col],
        orientation="h",
        marker=dict(color=colors, line=dict(color="black", width=0.5)),
        error_x=error_x,
        text=[f"{score:.3f}" for score in valid_models[metric_col]],
        textposition="auto",
        hovertemplate=(
            "<b>%{y}</b><br>" +
            f"{metric.title()}: %{{x:.4f}}<br>" +
            "Category: %{customdata[0]}<br>" +
            "Author: %{customdata[1]}<br>" +
            "Samples: %{customdata[2]}<br>" +
            "<extra></extra>"
        ),
        customdata=list(zip(
            valid_models["model_category"],
            valid_models["author"],
            valid_models.get(f"{track}_samples", [0] * len(valid_models))
        )),
    ))
    
    # Customize layout
    track_info = EVALUATION_TRACKS[track]
    fig.update_layout(
        title=f"🏆 {track_info['name']} - {metric.title()} Score",
        xaxis_title=f"{metric.title()} Score (with 95% CI)",
        yaxis_title="Models",
        height=max(400, len(valid_models) * 35 + 100),
        margin=dict(l=20, r=20, t=60, b=20),
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        font=dict(size=12),
    )
    
    # Reverse y-axis to show best model at top
    fig.update_yaxes(autorange="reversed")
    
    return fig


def create_language_pair_heatmap(
    model_results: Dict, track: str, metric: str = "quality_score"
) -> go.Figure:
    """Create language pair heatmap for a model."""
    
    if not model_results or "tracks" not in model_results:
        fig = go.Figure()
        fig.add_annotation(text="No model results available", x=0.5, y=0.5, showarrow=False)
        return fig
    
    track_data = model_results["tracks"].get(track, {})
    if track_data.get("error") or "pair_metrics" not in track_data:
        fig = go.Figure()
        fig.add_annotation(text=f"No data available for {track} track", x=0.5, y=0.5, showarrow=False)
        return fig
    
    pair_metrics = track_data["pair_metrics"]
    track_languages = EVALUATION_TRACKS[track]["languages"]
    
    # Create matrix for heatmap
    n_langs = len(track_languages)
    matrix = np.full((n_langs, n_langs), np.nan)
    
    for i, src_lang in enumerate(track_languages):
        for j, tgt_lang in enumerate(track_languages):
            if src_lang != tgt_lang:
                pair_key = f"{src_lang}_to_{tgt_lang}"
                if pair_key in pair_metrics and metric in pair_metrics[pair_key]:
                    matrix[i, j] = pair_metrics[pair_key][metric]["mean"]
    
    # Create language labels
    lang_labels = [LANGUAGE_NAMES.get(lang, lang.upper()) for lang in track_languages]
    
    # Create heatmap
    fig = go.Figure(data=go.Heatmap(
        z=matrix,
        x=lang_labels,
        y=lang_labels,
        colorscale="Viridis",
        showscale=True,
        colorbar=dict(
            title=f"{metric.replace('_', ' ').title()}",
            titleside="right",
            len=0.8,
        ),
        hovertemplate=(
            "Source: %{y}<br>" +
            "Target: %{x}<br>" +
            f"{metric.replace('_', ' ').title()}: %{{z:.3f}}<br>" +
            "<extra></extra>"
        ),
        zmin=0,
        zmax=1 if metric == "quality_score" else None,
    ))
    
    # Customize layout
    track_info = EVALUATION_TRACKS[track]
    fig.update_layout(
        title=f"🗺️ {track_info['name']} - {metric.replace('_', ' ').title()} by Language Pair",
        xaxis_title="Target Language",
        yaxis_title="Source Language",
        height=600,
        width=700,
        font=dict(size=12),
        xaxis=dict(side="bottom"),
        yaxis=dict(autorange="reversed"),
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
    )
    
    return fig


def create_performance_comparison_plot(df: pd.DataFrame, track: str) -> go.Figure:
    """Create performance comparison plot showing confidence intervals."""
    
    if df.empty:
        fig = go.Figure()
        fig.add_annotation(text="No data available", x=0.5, y=0.5, showarrow=False)
        return fig
    
    metric_col = f"{track}_quality"
    ci_lower_col = f"{track}_ci_lower"
    ci_upper_col = f"{track}_ci_upper"
    
    # Filter to models with data for this track
    valid_models = df[
        (df[metric_col] > 0) & 
        (df[ci_lower_col].notna()) & 
        (df[ci_upper_col].notna())
    ].head(10)
    
    if valid_models.empty:
        fig = go.Figure()
        fig.add_annotation(text="No models with confidence intervals", x=0.5, y=0.5, showarrow=False)
        return fig
    
    fig = go.Figure()
    
    # Add confidence intervals as error bars
    for i, (_, model) in enumerate(valid_models.iterrows()):
        category = model["model_category"]
        color = MODEL_CATEGORIES.get(category, {}).get("color", "#808080")
        
        # Main point
        fig.add_trace(go.Scatter(
            x=[model[metric_col]],
            y=[i],
            mode="markers",
            marker=dict(
                size=12,
                color=color,
                line=dict(color="black", width=1),
            ),
            name=model["model_name"],
            showlegend=False,
            hovertemplate=(
                f"<b>{model['model_name']}</b><br>" +
                f"Quality: {model[metric_col]:.4f}<br>" +
                f"95% CI: [{model[ci_lower_col]:.4f}, {model[ci_upper_col]:.4f}]<br>" +
                f"Category: {category}<br>" +
                "<extra></extra>"
            ),
        ))
        
        # Confidence interval line
        fig.add_trace(go.Scatter(
            x=[model[ci_lower_col], model[ci_upper_col]],
            y=[i, i],
            mode="lines",
            line=dict(color=color, width=3),
            showlegend=False,
            hoverinfo="skip",
        ))
    
    # Customize layout
    track_info = EVALUATION_TRACKS[track]
    fig.update_layout(
        title=f"📊 {track_info['name']} - Performance Comparison",
        xaxis_title="Quality Score",
        yaxis_title="Models",
        height=max(400, len(valid_models) * 40 + 100),
        yaxis=dict(
            tickmode="array",
            tickvals=list(range(len(valid_models))),
            ticktext=valid_models["model_name"].tolist(),
            autorange="reversed",
        ),
        showlegend=False,
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
    )
    
    return fig


def create_language_pair_comparison_plot(pairs_df: pd.DataFrame, track: str) -> go.Figure:
    """Create language pair comparison plot showing all models across all pairs."""
    
    if pairs_df.empty:
        fig = go.Figure()
        fig.add_annotation(
            text="No language pair data available", 
            x=0.5, y=0.5, showarrow=False
        )
        return fig
    
    # Get unique language pairs and models
    language_pairs = sorted(pairs_df['Language Pair'].unique())
    models = sorted(pairs_df['Model'].unique())
    
    if len(language_pairs) == 0 or len(models) == 0:
        fig = go.Figure()
        fig.add_annotation(
            text="Insufficient data for comparison", 
            x=0.5, y=0.5, showarrow=False
        )
        return fig
    
    # Create subplot for each metric
    fig = make_subplots(
        rows=2, cols=1,
        subplot_titles=('Quality Score by Language Pair', 'BLEU Score by Language Pair'),
        vertical_spacing=0.1,
        shared_xaxes=True
    )
    
    # Quality Score comparison
    for model in models:
        model_data = pairs_df[pairs_df['Model'] == model]
        category = model_data['Category'].iloc[0] if not model_data.empty else 'community'
        color = MODEL_CATEGORIES.get(category, {}).get('color', '#808080')
        
        fig.add_trace(
            go.Bar(
                name=model,
                x=model_data['Language Pair'],
                y=model_data['Quality Score'],
                marker_color=color,
                opacity=0.8,
                legendgroup=model,
                showlegend=True,
                hovertemplate=(
                    f"<b>{model}</b><br>" +
                    "Language Pair: %{x}<br>" +
                    "Quality Score: %{y:.4f}<br>" +
                    f"Category: {category}<br>" +
                    "<extra></extra>"
                )
            ),
            row=1, col=1
        )
        
        # BLEU Score comparison
        fig.add_trace(
            go.Bar(
                name=model,
                x=model_data['Language Pair'],
                y=model_data['BLEU'],
                marker_color=color,
                opacity=0.8,
                legendgroup=model,
                showlegend=False,
                hovertemplate=(
                    f"<b>{model}</b><br>" +
                    "Language Pair: %{x}<br>" +
                    "BLEU: %{y:.2f}<br>" +
                    f"Category: {category}<br>" +
                    "<extra></extra>"
                )
            ),
            row=2, col=1
        )
    
    # Update layout
    track_info = EVALUATION_TRACKS[track]
    fig.update_layout(
        title=f"📊 {track_info['name']} - Language Pair Performance Comparison",
        height=800,
        barmode='group',
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        )
    )
    
    # Rotate x-axis labels for better readability
    fig.update_xaxes(tickangle=45, row=2, col=1)
    fig.update_yaxes(title_text="Quality Score", row=1, col=1)
    fig.update_yaxes(title_text="BLEU Score", row=2, col=1)
    
    return fig


def create_category_comparison_plot(df: pd.DataFrame, track: str) -> go.Figure:
    """Create category-wise comparison plot."""
    
    if df.empty:
        fig = go.Figure()
        fig.add_annotation(text="No data available", x=0.5, y=0.5, showarrow=False)
        return fig
    
    metric_col = f"{track}_quality"
    
    # Filter to models with data
    valid_models = df[df[metric_col] > 0]
    
    if valid_models.empty:
        fig = go.Figure()
        fig.add_annotation(text="No valid models found", x=0.5, y=0.5, showarrow=False)
        return fig
    
    fig = go.Figure()
    
    # Create box plot for each category
    for category, info in MODEL_CATEGORIES.items():
        category_models = valid_models[valid_models["model_category"] == category]
        
        if len(category_models) > 0:
            fig.add_trace(go.Box(
                y=category_models[metric_col],
                name=info["name"],
                marker_color=info["color"],
                boxpoints="all",  # Show all points
                jitter=0.3,
                pointpos=-1.8,
                hovertemplate=(
                    f"<b>{info['name']}</b><br>" +
                    "Quality: %{y:.4f}<br>" +
                    "Model: %{customdata}<br>" +
                    "<extra></extra>"
                ),
                customdata=category_models["model_name"],
            ))
    
    # Customize layout
    track_info = EVALUATION_TRACKS[track]
    fig.update_layout(
        title=f"📈 {track_info['name']} - Performance by Category",
        xaxis_title="Model Category",
        yaxis_title="Quality Score",
        height=500,
        showlegend=False,
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
    )
    
    return fig