
import numpy as np
import pandas as pd
import altair as alt

# Register custom font for 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')

# Hardcoded data from experiments
MODEL_DATA = {
    "CAP + Oracle": {
        "total": 5000,
        "outcomes": {
            "Success": 3994,
            "Did not lift enough": 89,
            "Object touched but not grasped": 412,
            "Picked wrong object": 59,
            "Empty Grasp": 446,
            "Did not grasp": 0,
        }
    },
    "CAP + Gemini": {
        "total": 5000,
        "outcomes": {
            "Success": 3441,
            "Did not lift enough": 88,
            "Object touched but not grasped": 389,
            "Picked wrong object": 448,
            "Empty Grasp": 634,
            "Did not grasp": 0,
        }
    },
    "CAP + Molmo": {
        "total": 5000,
        "outcomes": {
            "Success": 3008,
            "Did not lift enough": 94,
            "Object touched but not grasped": 467,
            "Picked wrong object": 877,
            "Empty Grasp": 554,
            "Did not grasp": 0,
        }
    },
    "CAP + Moondream": {
        "total": 5000,
        "outcomes": {
            "Success": 3033,
            "Did not lift enough": 83,
            "Object touched but not grasped": 342,
            "Picked wrong object": 746,
            "Empty Grasp": 796,
            "Did not grasp": 0,
        }
    }
}


def plot_failure_modes():
    model_names = list(MODEL_DATA.keys())
    
    # Print results to terminal
    print("\n" + "="*80)
    print("WRONG OBJECT SELECTION RATE BY VLM (4 Distractors)")
    print("="*80)
    for model in model_names:
        data = MODEL_DATA[model]
        total = data["total"]
        print(f"\n{model} (n={total}):")
        print("-" * 60)
        for outcome, count in data["outcomes"].items():
            percentage = (count / total * 100) if total > 0 else 0
            print(f"  {outcome:40s}: {count:3d} ({percentage:5.1f}%)")
    print("\n" + "="*80 + "\n")
    
    # Only show "Picked wrong object" percentage
    chart_data = []
    for model in model_names:
        data = MODEL_DATA[model]
        total = data["total"] or 1
        count = data["outcomes"].get("Picked wrong object", 0)
        proportion = count / total
        percentage = proportion * 100
        
        # Calculate standard error for binomial proportion
        se = np.sqrt(proportion * (1 - proportion) / total) * 100
        
        chart_data.append({
            'Model': model,
            'Percentage': percentage,
            'SE': se,
        })
    
    df = pd.DataFrame(chart_data)
    
    # Sort by percentage (lowest to highest)
    df = df.sort_values('Percentage')
    
    # Create ordered list for x-axis
    model_order = df['Model'].tolist()
    
    # Assign colors from purple family (lightest for best, darkest for worst)
    purple_shades = ["#D6BAE2", "#9B66BB", "#7B4FA2", "#4B136D"]  # Lightest to darkest
    color_mapping = {model: purple_shades[i] for i, model in enumerate(model_order)}
    df['Color'] = df['Model'].map(color_mapping)
    
    # Create simple bar chart with gradient colors
    bars = alt.Chart(df).mark_bar(
        stroke='white',
        strokeWidth=1,
        cornerRadiusTopLeft=8,
        cornerRadiusTopRight=8
    ).encode(
        x=alt.X('Model:N', title=None, sort=model_order, axis=alt.Axis(labelFontSize=18, labelAngle=0)),
        y=alt.Y('Percentage:Q', title=None, axis=alt.Axis(labelFontSize=18)),
        color=alt.Color('Model:N', scale=alt.Scale(domain=model_order, range=[color_mapping[m] for m in model_order]), legend=None),
        tooltip=['Model', alt.Tooltip('Percentage:Q', format='.1f')]
    )
    
    # Add error bars
    error_bars = alt.Chart(df).mark_errorbar(
        color='black',
        thickness=2,
        ticks=alt.MarkConfig(width=10)
    ).encode(
        x=alt.X('Model:N', sort=model_order),
        y=alt.Y('Percentage:Q'),
        yError=alt.YError('SE:Q')
    )
    
    # Add percentage labels on top of bars
    text = alt.Chart(df).mark_text(
        align='center',
        baseline='bottom',
        dy=-15,
        fontSize=16,
        fontWeight='bold'
    ).encode(
        x=alt.X('Model:N', sort=model_order),
        y=alt.Y('Percentage:Q'),
        text=alt.Text('Percentage:Q', format='.0f')
    ).transform_calculate(
        label='datum.Percentage + "%"'
    )
    
    chart = (bars + error_bars + text).properties(
        width=600,
        height=400,
        title={
            'text': '       Wrong Object Selection Rate by VLM (4 Distractors)',
            'fontSize': 22,
            'anchor': 'start',
            'dx': 60,
            'dy': -20
        },
        padding={"left": 5, "right": 5, "top": 20, "bottom": 0}
    ).configure_view(
        strokeWidth=0
    )
    
    # Save chart
    chart.save("failure_modes_by_model.html")
    chart.save("failure_modes_by_model.png", scale_factor=3)
    chart.save("failure_modes_by_model.pdf")
    print("Saved: failure_modes_by_model.html, .png, .pdf")


if __name__ == "__main__":
    plot_failure_modes()
