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

# 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')

BASE_DIR = "logs"
REWARD_THRESHOLD = 0.03

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_pi0_bars():
    model_folder = "pi0"
    num_objects = [1, 2, 3, 4, 5]
    
    rows = []
    
    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: π-0.5, {n_obj} objects")
            continue
        
        s, t = compute_success_from_csv(csv_path)
        
        # Beta posterior
        a, b = 1 + s, 1 + (t - s)
        mean = 100 * a / (a + b)
        lo = 100 * beta.ppf(0.025, a, b)
        hi = 100 * beta.ppf(0.975, a, b)
        
        print(f"π-0.5 | {n_obj} objects: {s}/{t} = {mean:.1f}%")
        
        rows.append({
            "num_objects": n_obj,
            "mean": mean,
            "lo": lo,
            "hi": hi,
            "successes": s,
            "total": t
        })
    
    df = pd.DataFrame(rows)
    
    # Create bar chart
    bars = alt.Chart(df).mark_bar(
        color='#E0BE16',
        width=60
    ).encode(
        x=alt.X(
            'num_objects:O',
            title='Number of Objects',
            axis=alt.Axis(
                labelFontSize=20,
                titleFontSize=24,
                labelAngle=0,
                titlePadding=15
            )
        ),
        y=alt.Y(
            'mean:Q',
            title='Success Rate (%)',
            scale=alt.Scale(domain=[0, 100]),
            axis=alt.Axis(
                labelFontSize=20,
                titleFontSize=24,
                titlePadding=14,
                grid=True,
                gridOpacity=0.3
            )
        ),
        tooltip=[
            alt.Tooltip('num_objects:O', title='Objects'),
            alt.Tooltip('mean:Q', title='Success Rate (%)', format='.1f'),
            alt.Tooltip('successes:Q', title='Successes'),
            alt.Tooltip('total:Q', title='Total')
        ]
    )
    
    # Create error bars for 95% CI
    error_bars = alt.Chart(df).mark_errorbar(
        ticks=True,
        thickness=2
    ).encode(
        x=alt.X('num_objects:O'),
        y=alt.Y('lo:Q', title=''),
        y2=alt.Y2('hi:Q')
    )
    
    # Add text labels on top of bars
    text = alt.Chart(df).mark_text(
        dy=-10,
        fontSize=16,
        fontWeight='bold'
    ).encode(
        x=alt.X('num_objects:O'),
        y=alt.Y('mean:Q'),
        text=alt.Text('mean:Q', format='.1f')
    )
    
    # Combine layers
    chart = (bars + error_bars + text).properties(
        width=500,
        height=400,
        title={
            'text': '           π-0.5 Success Rate by Number of Objects',
            'fontSize': 24,
            'anchor': 'start',
            'dx': 40,
            'dy': -10
        },
        padding={'left': 10, 'right': 10, 'top': 40, 'bottom': 40}
    ).configure_view(
        strokeWidth=0
    )
    
    return chart


if __name__ == "__main__":
    chart = plot_pi0_bars()
    chart.save("pi0_bars.html")
    chart.save("pi0_bars.png", scale_factor=3)
    chart.save("pi0_bars.pdf", scale_factor=3)
    print("\nPlot saved to: pi0_bars.html, pi0_bars.png, and pi0_bars.pdf")
