import os
import numpy as np
import pandas as pd
import altair as alt
from scipy.stats import beta, gaussian_kde

# Register custom font for PNG export
alt.themes.register('custom_theme', lambda: {
    'config': {
        'title': {'font': 'Produkt'},
        'axis': {'labelFont': 'Produkt', 'titleFont': 'Produkt'},
        'legend': {'labelFont': 'Produkt', 'titleFont': 'Produkt'},
        'mark': {'font': 'Produkt'},
        'text': {'font': 'Produkt'},
    }
})
alt.themes.enable('custom_theme')

alt.renderers.enable("colab")

BASE_DIR = "logs"
REWARD_THRESHOLD = 0.03
SHOW_LABELS = 0  # Set to True to show percentage labels above points

MODELS = {
    "CAP + Gemini-ER": "rum/gemini",
    "π-0.5": "pi0",
    "CAP + Oracle": "rum/oracle",
    "CAP + Molmo": "rum/molmo",
    "CAP + Moondream": "rum/moondream",
}

NUM_OBJECTS = [1, 2, 3, 4, 5]


def compute_success_from_csv(csv_path):
    df = pd.read_csv(csv_path, sep="\t")
    successes = (df["max_reward"] > REWARD_THRESHOLD).sum()
    total = len(df)
    return successes, total


def plot_bayesian_success_by_objects():
    rows = []

    for model_name, model_folder in MODELS.items():
        for n_obj in NUM_OBJECTS:
            possible_folders = [
                f"{n_obj}_objects",
                f"{n_obj}_object",
                f"{n_obj}-objects",
                f"{n_obj}-object",
            ]

            csv_path = None
            for folder in possible_folders:
                # First try direct path
                candidate = os.path.join(
                    BASE_DIR, model_folder, folder, "log.csv"
                )
                if os.path.exists(candidate):
                    csv_path = candidate
                    break
                
                # Try nested evaluation folder structure
                folder_path = os.path.join(BASE_DIR, model_folder, folder)
                if os.path.exists(folder_path) and os.path.isdir(folder_path):
                    for subdir in os.listdir(folder_path):
                        subdir_path = os.path.join(folder_path, subdir)
                        if os.path.isdir(subdir_path):
                            candidate = os.path.join(subdir_path, "log.csv")
                            if os.path.exists(candidate):
                                csv_path = candidate
                                break
                    if csv_path:
                        break

            if csv_path is None:
                print(f"Missing data: {model_name}, {n_obj} objects")
                continue

            s, t = compute_success_from_csv(csv_path)

            # Beta posterior
            a, b = 1 + s, 1 + (t - s)
            mean = a / (a + b)
            lo = beta.ppf(0.025, a, b)
            hi = beta.ppf(0.975, a, b)

            print(f"{model_name} | {n_obj} objects: {s}/{t} = {mean:.2%}")

            rows.append({
                "model": model_name,
                "num_objects": n_obj,
                "num_distractors": n_obj - 1,
                "mean": mean,
                "lo": lo,
                "hi": hi,
                "err": hi - mean,
                "alpha": a,
                "beta": b,
            })

    df = pd.DataFrame(rows)
    
    # Normalize each model's success rates relative to 0 distractors (1 object)
    normalized_rows = []
    for model_name in df["model"].unique():
        model_df = df[df["model"] == model_name].copy()
        # Get the baseline (0 distractors = 1 object)
        baseline_row = model_df[model_df["num_distractors"] == 0]
        if len(baseline_row) == 0:
            print(f"Warning: No 0 distractor data for {model_name}")
            continue
        baseline_mean = baseline_row.iloc[0]["mean"]
        if baseline_mean == 0:
            print(f"Warning: Baseline is 0 for {model_name}, skipping normalization")
            continue
        
        # Store baseline for violin normalization
        model_df["baseline_mean"] = baseline_mean
        
        # Normalize all metrics
        model_df["mean"] = model_df["mean"] / baseline_mean
        model_df["lo"] = model_df["lo"] / baseline_mean
        model_df["hi"] = model_df["hi"] / baseline_mean
        model_df["err"] = model_df["hi"] - model_df["mean"]
        normalized_rows.append(model_df)
    
    df = pd.concat(normalized_rows, ignore_index=True)
    
    # Calculate average success rate per model and sort by descending order
    model_means = df.groupby("model")["mean"].mean().sort_values(ascending=False)
    sorted_models = model_means.index.tolist()
    
    # Create color mapping maintaining original colors
    color_map = {
        "CAP + Gemini-ER": "#6095e1",
        "π-0.5": "#E0BE16",
        "CAP + Oracle": "#A918C0",
        "CAP + Molmo": "#F0529c",
        "CAP + Moondream": "#000000",
        "CAP + 4o": "#FF6B35"
    }
    
    # Ensure all sorted models have colors
    sorted_colors = []
    for model in sorted_models:
        sorted_colors.append(color_map[model])

    print(f"Sorted models: {sorted_models}")
    print(f"Sorted colors: {sorted_colors}")

    base = alt.Chart(df).encode(
        x=alt.X(
            "num_distractors:Q",
            title="Number of Distractor Objects",
            axis=alt.Axis(labelFontSize=20, titleFontSize=24, labelAngle=0, titlePadding=15, tickMinStep=1),
            scale=alt.Scale(domain=[0, 4.5])
        ),
        color=alt.Color(
            "model:N",
            scale=alt.Scale(
                domain=sorted_models,
                range=sorted_colors
            ),
            legend=alt.Legend(
                title=None,
                labelFontSize=18,
                symbolSize=200,
                orient="bottom",
                direction="horizontal",
                symbolType="square",
                labelLimit=0
            )
        )
    )

    lines = base.mark_line(
        strokeWidth=3,
        point=alt.OverlayMarkDef(size=100)
    ).encode(
        y=alt.Y(
            "mean:Q",
            title="Relative Success Rate (%)",
            scale=alt.Scale(domain=[0.6, 1.0]),
            axis=alt.Axis(
                values=[0.6, 0.7, 0.8, 0.9, 1.0],
                labelExpr="datum.value * 100",
                labelFontSize=20,
                titleFontSize=24,
                titlePadding=14,
                grid=False
            )
        )
    )

    # Create violin plots only for non-zero distractors
    df_with_violins = df[df["num_distractors"] != 0]
    
    # Set random seed for reproducibility
    np.random.seed(42)
    
    # Create violin polygons
    violin_width_scale = 0.08  # Width of violins
    violin_polygon_data = []
    
    for _, row in df_with_violins.iterrows():
        # Generate samples from beta distribution (these are in 0-1 range)
        samples = beta.rvs(row["alpha"], row["beta"], size=2000)
        # Normalize samples using the stored baseline_mean
        baseline_mean = row["baseline_mean"]
        if baseline_mean > 0:
            samples = samples / baseline_mean
        
        # Compute KDE
        kde = gaussian_kde(samples)
        
        # Create density curve points
        y_points = np.linspace(samples.min(), samples.max(), 100)
        densities = kde(y_points)
        densities = densities / densities.max() * violin_width_scale
        
        xc = row["num_distractors"]
        model = row["model"]
        color = sorted_colors[sorted_models.index(model)]
        
        # Create closed polygon: left side up, then right side down
        for i, (y, d) in enumerate(zip(y_points, densities)):
            violin_polygon_data.append({
                "model": model,
                "x": xc - d,
                "y": y,
                "order": i,
                "color": color,
                "group": f"{model}_{xc}"
            })
        # Right side going back down
        for i, (y, d) in enumerate(zip(y_points[::-1], densities[::-1])):
            violin_polygon_data.append({
                "model": model,
                "x": xc + d,
                "y": y,
                "order": len(y_points) + i,
                "color": color,
                "group": f"{model}_{xc}"
            })
    
    violin_polygon_df = pd.DataFrame(violin_polygon_data)
    
    # Create violin shapes - now uses same quantitative scale as base
    violins = alt.Chart(violin_polygon_df).mark_line(
        fillOpacity=0.4,
        strokeWidth=0.5,
        interpolate='linear',
        filled=True
    ).encode(
        x=alt.X('x:Q', scale=alt.Scale(domain=[0, 4.5]), title=''),
        y=alt.Y('y:Q', scale=alt.Scale(domain=[0.6, 1.0]), title=''),
        order='order:Q',
        detail='group:N',
        fill=alt.Fill('color:N', scale=None, legend=None),
        stroke=alt.Stroke('color:N', scale=None, legend=None)
    )


    labels = alt.Chart(df).mark_text(
        dy=-8,
        fontSize=14,
        color="black"
    ).encode(
        x=alt.X("num_distractors:O"),
        y="hi:Q",
        text=alt.Text("mean:Q", format=".0%")
    )

    # Combine layers based on SHOW_LABELS setting
    if SHOW_LABELS:
        chart = violins + lines + labels
    else:
        chart = violins + lines

    return (
        chart
        .properties(
            width=720,
            height=420,
            title={
                "text": "           EgoGym Pick Relative Success Rate vs Number of Distractors",
                "fontSize": 22,
                "anchor": "start",
                "dx": 65,
                "offset": 20
            },
            padding={"left": 10, "right": 10, "top": 40, "bottom": 40}
        )
        .configure_view(stroke=None)
    )


chart = plot_bayesian_success_by_objects()
chart.save("success_by_objects.png", scale_factor=3)
chart.save("success_by_objects.pdf", scale_factor=3)