
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/5_objects"
REWARD_THRESHOLD = 0.03
PARTIAL_LIFT_THRESHOLD = 0.005


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():
    checkpoint_order = ["checkpoint_10", "checkpoint_50", "checkpoint_64", "checkpoint_80"]
    checkpoint_labels = {
        "checkpoint_10": "Checkpoint 24%",
        "checkpoint_50": "Checkpoint 39%",
        "checkpoint_64": "Checkpoint 68%",
        "checkpoint_80": "Checkpoint 83%"
    }
    outcome_data = {}
    episode_totals = {}
    
    for checkpoint in checkpoint_order:
        csv_path = os.path.join(BASE_DIR, checkpoint, "log.csv")
        outcomes, total_episodes = compute_outcomes_from_csv(csv_path)
        if outcomes is None:
            print(f"Warning: Could not find CSV data for {checkpoint}")
            continue
        outcome_data[checkpoint] = outcomes
        episode_totals[checkpoint] = total_episodes
    
    if not outcome_data:
        print("No data found!")
        return
    
    checkpoint_names = list(outcome_data.keys())
    
    # Print results to terminal
    print("\n" + "="*80)
    print("FAILURE MODE ANALYSIS BY CHECKPOINT")
    print("="*80)
    for checkpoint in checkpoint_names:
        checkpoint_label = checkpoint_labels.get(checkpoint, checkpoint)
        total = episode_totals.get(checkpoint, 0)
        print(f"\n{checkpoint_label} (n={total}):")
        print("-" * 60)
        for outcome, count in outcome_data[checkpoint].items():
            percentage = (count / total * 100) if total > 0 else 0
            print(f"  {outcome:40s}: {count:3d} ({percentage:5.1f}%)")
    print("\n" + "="*80 + "\n")
    
    outcome_order = [
        "Success",
        "Did not lift enough",
        "Object touched but not grasped",
        "Picked wrong object",
        "Empty Grasp",
        "Did not grasp",
    ]
    
    colors = {
        "Success": "#388038",  # Green
        "Did not lift enough": "#F7D45B",  # Yellow
        "Object touched but not grasped": "#66ACF7",  # Blue
        "Picked wrong object": "#F0529C",  # Pink
        "Empty Grasp": "#9B66BB",  # Purple
        "Did not grasp": "#870927",  # Dark red
    }
    
    # Calculate average percentage for each outcome to sort by size
    outcome_totals = {}
    for checkpoint in checkpoint_names:
        total = episode_totals.get(checkpoint, 0) or 1
        for outcome in outcome_order:
            count = outcome_data[checkpoint].get(outcome, 0)
            percentage = count / total * 100
            if outcome not in outcome_totals:
                outcome_totals[outcome] = 0
            outcome_totals[outcome] += percentage
    
    # Sort outcomes by total percentage (largest first), excluding Success
    failure_modes = [o for o in outcome_order if o != "Success"]
    sorted_failures = sorted(failure_modes, key=lambda x: outcome_totals.get(x, 0), reverse=True)
    
    # Build stacking order: Success at bottom, then failures from largest to smallest going up
    stacking_order = ["Success"] + sorted_failures
    
    chart_data = []
    for checkpoint in checkpoint_names:
        checkpoint_label = checkpoint_labels.get(checkpoint, checkpoint)
        
        total = episode_totals.get(checkpoint, 0) or 1
        for outcome in stacking_order:
            count = outcome_data[checkpoint].get(outcome, 0)
            percentage = count / total * 100
            if percentage > 0:  # Only include non-zero values
                chart_data.append({
                    'Checkpoint': checkpoint_label,
                    'Outcome': outcome,
                    'Percentage': percentage,
                    'Color': colors.get(outcome, "#999999")
                })
    
    df = pd.DataFrame(chart_data)
    
    # Get only outcomes that actually appear in the data
    outcomes_in_data = df['Outcome'].unique().tolist()
    
    # Filter stacking_order to only include outcomes present in data
    filtered_stacking_order = [o for o in stacking_order if o in outcomes_in_data]
    
    # Add sort index to control stacking order
    outcome_to_index = {outcome: i for i, outcome in enumerate(filtered_stacking_order)}
    df['sort_index'] = df['Outcome'].map(outcome_to_index)
    
    # Create color scale only for outcomes in data
    color_scale = alt.Scale(
        domain=filtered_stacking_order,
        range=[colors[o] for o in filtered_stacking_order]
    )
    
    # Create stacked bar chart
    chart = alt.Chart(df).mark_bar(
        stroke='white',
        strokeWidth=1
    ).encode(
        x=alt.X('Checkpoint:N', title=None, axis=alt.Axis(labelFontSize=18, labelAngle=0)),
        y=alt.Y('Percentage:Q', title='Share of Episodes (%)', axis=alt.Axis(labelFontSize=18, titleFontSize=20)).scale(domain=[0, 100]),
        color=alt.Color('Outcome:N', scale=color_scale, sort=filtered_stacking_order, legend=alt.Legend(
            title=None,
            labelFontSize=16,
            symbolSize=200,
            orient="bottom",
            direction="horizontal",
            labelLimit=0,
            columns=3
        )),
        order=alt.Order('sort_index:Q'),
        tooltip=['Checkpoint', 'Outcome', alt.Tooltip('Percentage:Q', format='.1f')]
    ).properties(
        width=600,
        height=400,
        title={
            'text': '       EgoGym-Pick Failure Modes by Checkpoint',
            'fontSize': 22,
            'anchor': 'start',
            'dx': 60,
            'dy': -20
        },
        padding={"left": 5, "right": 5, "top": 20, "bottom": 40}
    ).configure_view(
        strokeWidth=0
    )
    
    # Save chart
    chart.save("failure_modes_by_checkpoint.html")
    chart.save("failure_modes_by_checkpoint.pdf", scale_factor=3)
    chart.save("failure_modes_by_checkpoint.png", scale_factor=3)
    print("\nPlot saved to: failure_modes_by_checkpoint.html and failure_modes_by_checkpoint.png")


if __name__ == "__main__":
    plot_failure_modes()
