
import os
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')

BASE_DIR = "logs"
REWARD_THRESHOLD = 0.03
PARTIAL_LIFT_THRESHOLD = 0.01

MODELS = {
    "CAP + Oracle": "rum/oracle",
    "CAP + Gemini": "rum/gemini",
    "CAP + Molmo": "rum/molmo",
    "CAP + Moondream": "rum/moondream",
}


def compute_outcomes_from_csv(csv_path):
    if not os.path.exists(csv_path):
        return None, None
    df = pd.read_csv(csv_path, sep="\t")
    total_episodes = len(df)
    
    # Count successes
    successes = (df["max_reward"] > REWARD_THRESHOLD).sum()
    
    def get_failure_mode(row):
        if row["max_reward"] > REWARD_THRESHOLD:
            return "Success"
        
        bodies_contacted = str(row.get("grasped_bodies", ""))
        object_name = str(row.get("object_name", ""))
        grasping_object = row.get("grasping_object", False)
        is_grasping = row.get("is_grasping", False)  # Final gripper state only
        
        # Extract body names from grasped_bodies list
        has_target_contact = object_name in bodies_contacted
        has_gripper_contact = "left" in bodies_contacted or "right" in bodies_contacted
        has_any_object_contact = "object" in bodies_contacted
        has_wrong_object_contact = has_any_object_contact and not has_target_contact
        
        # Decision tree (most specific to least specific)
        # Note: is_grasping is only final state, grasping_object tracks if target was ever grasped
        
        # 1. Successfully grasped target but didn't lift high enough
        if grasping_object and row["max_reward"] >= PARTIAL_LIFT_THRESHOLD:
            return "Did not lift enough"
        
        # 2. Grasped target but barely lifted (or dropped immediately)
        if grasping_object:
            return "Object touched but not grasped"
        
        # 3. Made contact with target but never achieved grasp
        if has_target_contact:
            return "Object touched but not grasped"
        
        # 4. Grasped wrong object with significant lift
        if has_wrong_object_contact and row["max_reward"] >= PARTIAL_LIFT_THRESHOLD:
            return "Picked wrong object"
        
        # 5. Contacted wrong object (with or without final grasp state)
        if has_wrong_object_contact:
            return "Picked wrong object"
        
        # 6. Gripper closed at end or had gripper contact but no object identification
        if is_grasping or has_gripper_contact:
            return "Empty Grasp"
        
        # 7. Never made meaningful contact with anything
        return "Did not grasp"
    
    df["outcome"] = df.apply(get_failure_mode, axis=1)
    
    outcomes = {
        "Success": 0,
        "Did not lift enough": 0,
        "Object touched but not grasped": 0,
        "Picked wrong object": 0,
        "Empty Grasp": 0,
        "Did not grasp": 0,
    }
    
    for _, row in df.iterrows():
        mode = row["outcome"]
        if mode in outcomes:
            outcomes[mode] += 1
    
    return outcomes, total_episodes


def plot_failure_modes():
    outcome_data = {}
    episode_totals = {}
    
    # Get data for 5 objects (4 distractors) for each model
    for model_name, model_folder in MODELS.items():
        # Try multiple folder name variations and nested evaluation folders
        possible_paths = [
            os.path.join(BASE_DIR, model_folder, "5_objects", "log.csv"),
            os.path.join(BASE_DIR, model_folder, "5_object", "log.csv"),
        ]
        
        csv_path = None
        for path in possible_paths:
            if os.path.exists(path):
                csv_path = path
                break
        
        # Check for nested evaluation folders
        if csv_path is None:
            for folder_name in ["5_objects", "5_object"]:
                folder_path = os.path.join(BASE_DIR, model_folder, folder_name)
                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"Warning: Could not find CSV data for {model_name}")
            continue
            
        outcomes, total_episodes = compute_outcomes_from_csv(csv_path)
        if outcomes is None:
            print(f"Warning: Could not find CSV data for {model_name}")
            continue
        outcome_data[model_name] = outcomes
        episode_totals[model_name] = total_episodes
    
    if not outcome_data:
        print("No data found!")
        return
    
    model_names = list(outcome_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:
        total = episode_totals.get(model, 0)
        print(f"\n{model} (n={total}):")
        print("-" * 60)
        for outcome, count in outcome_data[model].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:
        total = episode_totals.get(model, 0) or 1
        count = outcome_data[model].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='.1f')
    )
    
    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()
