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


def plot_bayesian_correlation_hardcoded():
    # Set random seed for reproducibility
    np.random.seed(42)
    
    # Register and enable custom font theme
    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 values
    checkpoint_data = {
        "checkpoint_10": {
            "sim_success": 1001,
            "sim_total": 5000,
            "real_success": 60,
            "real_total": 250
        },
        "checkpoint_50": {
            "sim_success": 1847,
            "sim_total": 5000,
            "real_success": 96,
            "real_total": 250
        },
        "checkpoint_64": {
            "sim_success": 3137,
            "sim_total": 5000,
            "real_success": 168,
            "real_total": 250
        },
        "checkpoint_80": {
            "sim_success": 3994,
            "sim_total": 5000,
            "real_success": 208,
            "real_total": 250
        }
    }
    
    stats = {}
    checkpoint_order = ["checkpoint_10", "checkpoint_50", "checkpoint_64", "checkpoint_80"]
    
    for checkpoint in checkpoint_order:
        data = checkpoint_data[checkpoint]
        
        # Simulation stats
        a_sim = 1 + data["sim_success"]
        b_sim = 1 + (data["sim_total"] - data["sim_success"])
        sim_mean = 100 * a_sim / (a_sim + b_sim)
        
        # Real stats
        a_real = 1 + data["real_success"]
        b_real = 1 + (data["real_total"] - data["real_success"])
        real_mean = 100 * a_real / (a_real + b_real)
        
        print(f"{checkpoint}: sim={data['sim_success']}/{data['sim_total']} ({sim_mean:.1f}%), real={real_mean:.1f}% ({data['real_success']}/{data['real_total']})")
        
        stats[checkpoint] = {
            "sim_mean": sim_mean,
            "sim_alpha": a_sim,
            "sim_beta": b_sim,
            "real_mean": real_mean,
            "real_alpha": a_real,
            "real_beta": b_real,
        }
    
    if not stats:
        print("No data found!")
        return
    
    checkpoint_names = list(stats.keys())
    sim_rates = np.array([stats[c]["sim_mean"] for c in checkpoint_names])
    real_rates = np.array([stats[c]["real_mean"] for c in checkpoint_names])
    
    r, p_value = pearsonr(sim_rates, real_rates)
    
    # Compute confidence intervals and violin plot data
    violin_data = []
    point_data = []
    error_data = []
    
    for i, checkpoint in enumerate(checkpoint_names):
        sim_lo = beta.ppf(0.025, stats[checkpoint]["sim_alpha"], stats[checkpoint]["sim_beta"]) * 100
        sim_hi = beta.ppf(0.975, stats[checkpoint]["sim_alpha"], stats[checkpoint]["sim_beta"]) * 100
        
        # Generate violin data
        y_samples = beta.rvs(stats[checkpoint]["real_alpha"], stats[checkpoint]["real_beta"], size=10000) * 100
        
        xc = sim_rates[i]
        yc = real_rates[i]
        
        # Store point data
        point_data.append({
            'checkpoint': checkpoint,
            'sim_rate': xc,
            'real_rate': yc,
            'sim_lo': sim_lo,
            'sim_hi': sim_hi
        })
        
        # Store violin data for density transform
        for y_val in y_samples:
            violin_data.append({
                'checkpoint': checkpoint,
                'sim_rate': xc,
                'real_rate': y_val
            })
    
    # Compute bootstrap correlation confidence interval
    n_samples = 1000
    r_samples = []
    for _ in range(n_samples):
        sim_sample = [beta.rvs(stats[c]["sim_alpha"], stats[c]["sim_beta"]) * 100 for c in checkpoint_names]
        real_sample = [beta.rvs(stats[c]["real_alpha"], stats[c]["real_beta"]) * 100 for c in checkpoint_names]
        r_sample, _ = pearsonr(sim_sample, real_sample)
        r_samples.append(r_sample)
    r_samples = np.array(r_samples)
    r_lo = np.percentile(r_samples, 2.5)
    r_hi = np.percentile(r_samples, 97.5)
    
    # Create DataFrames
    point_df = pd.DataFrame(point_data)
    
    # For violins, we need much less data - just sample points
    violin_sample_data = []
    for checkpoint in checkpoint_names:
        # Use only 500 samples per checkpoint to avoid memory issues
        y_samples = beta.rvs(stats[checkpoint]["real_alpha"], stats[checkpoint]["real_beta"], size=500) * 100
        xc = point_df[point_df['checkpoint'] == checkpoint]['sim_rate'].values[0]
        for y_val in y_samples:
            violin_sample_data.append({
                'checkpoint': checkpoint,
                'sim_rate': xc,
                'real_rate': y_val
            })
    
    violin_df = pd.DataFrame(violin_sample_data)
    
    # Compute regression line
    z = np.polyfit(sim_rates, real_rates, 1)
    p = np.poly1d(z)
    xs = np.linspace(sim_rates.min() - 5, sim_rates.max() + 5, 200)
    regression_df = pd.DataFrame({'sim_rate': xs, 'real_rate': p(xs)})
    
    # Manually create violin shapes positioned at sim_rate coordinates
    violin_width_scale = 2.5
    violin_polygon_data = []
    
    # Color mapping from lowest to highest checkpoint
    checkpoint_colors = {
        "checkpoint_10": "#F8F0FA",  # very light purple
        "checkpoint_50": "#D6BAE2",  # medium purple
        "checkpoint_64": "#9B66BB",  # medium-dark purple
        "checkpoint_80": "#4B136D"   # deep purple
    }
    
    for checkpoint in checkpoint_names:
        # Get samples and compute KDE
        y_samples = beta.rvs(stats[checkpoint]["real_alpha"], stats[checkpoint]["real_beta"], size=2000) * 100
        kde = gaussian_kde(y_samples)
        
        # Create density curve points
        y_points = np.linspace(y_samples.min(), y_samples.max(), 100)
        densities = kde(y_points)
        densities = densities / densities.max() * violin_width_scale
        
        xc = point_df[point_df['checkpoint'] == checkpoint]['sim_rate'].values[0]
        color = checkpoint_colors[checkpoint]
        
        # Create closed polygon: left side up, then right side down
        for i, (y, d) in enumerate(zip(y_points, densities)):
            violin_polygon_data.append({
                'checkpoint': checkpoint,
                'x': xc - d,
                'y': y,
                'order': i,
                'color': color
            })
        # Right side going back down
        for i, (y, d) in enumerate(zip(y_points[::-1], densities[::-1])):
            violin_polygon_data.append({
                'checkpoint': checkpoint,
                'x': xc + d,
                'y': y,
                'order': len(y_points) + i,
                'color': color
            })
    
    violin_polygon_df = pd.DataFrame(violin_polygon_data)
    
    # Create violin shapes with color encoding
    violins = alt.Chart(violin_polygon_df).mark_line(
        fillOpacity=0.6,
        stroke='#8B4789',
        strokeWidth=1.1,
        interpolate='linear',
        filled=True
    ).encode(
        x=alt.X('x:Q', title='EgoGym Performance (%)').scale(zero=False),
        y=alt.Y('y:Q', title='Real Performance (%)').scale(zero=False),
        order='order:Q',
        detail='checkpoint:N',
        fill=alt.Fill('color:N', scale=None, legend=None)
    )
    
    
    # Create regression line
    regression_line = alt.Chart(regression_df).mark_line(
        strokeDash=[5, 5],
        color='black',
        opacity=0.5,
        size=1.5
    ).encode(
        x=alt.X('sim_rate:Q').scale(zero=False),
        y=alt.Y('real_rate:Q').scale(zero=False)
    )
    
    # Create error bars for simulation
    error_bars = alt.Chart(point_df).mark_errorbar(ticks=True, thickness=1.5).encode(
        x=alt.X('sim_lo:Q', title=''),
        x2=alt.X2('sim_hi:Q'),
        y='real_rate:Q'
    )
    
    # Create horizontal lines at mean (for each violin)
    mean_line_data = []
    for checkpoint in checkpoint_names:
        xc = point_df[point_df['checkpoint'] == checkpoint]['sim_rate'].values[0]
        yc = point_df[point_df['checkpoint'] == checkpoint]['real_rate'].values[0]
        mean_line_data.append({
            'checkpoint': checkpoint,
            'x_left': xc - 1.0,
            'x_right': xc + 1.0,
            'y': yc
        })
    mean_line_df = pd.DataFrame(mean_line_data)
    
    mean_lines = alt.Chart(mean_line_df).mark_rule(color='black', size=1.2).encode(
        x=alt.X('x_left:Q'),
        x2=alt.X2('x_right:Q'),
        y='y:Q'
    )
    
    # Create central points
    points = alt.Chart(point_df).mark_point(
        filled=True,
        size=100,
        color='#A894DB',
        stroke='#8B4789',
        strokeWidth=1.5,
        shape='diamond'
    ).encode(
        x=alt.X('sim_rate:Q').scale(zero=False),
        y=alt.Y('real_rate:Q').scale(zero=False),
        tooltip=['checkpoint', 'sim_rate', 'real_rate']
    )
    
    # Combine layers (violins, regression line, error bars, mean lines and points on top)
    chart = (violins + regression_line + mean_lines + error_bars + points).properties(
        title={
            'text': '            Blind EgoGym-Pick Sim-to-Real Correlation',
            'fontSize': 24,
            'anchor': 'start',
            'dx': 15,
            'dy': -8
        },
        width=500,
        height=400
    ).configure_axis(
        labelFontSize=16,
        titleFontSize=18,
        titleFontStyle='normal',
        grid=True,
        gridOpacity=0.3,
        tickCount=6
    ).configure_view(
        strokeWidth=0
    )
    
    # Add background box for correlation text (wider and thinner)
    correlation_box = alt.Chart(pd.DataFrame([{
        'x': point_df['sim_rate'].max() - 18,
        'y': point_df['real_rate'].min() - 10,
        'x2': point_df['sim_rate'].max() + 4,
        'y2': point_df['real_rate'].min() - 4
    }])).mark_rect(
        fill='#F7D45B',
        stroke='gray',
        strokeWidth=1.5,
        opacity=0.9
    ).encode(
        x=alt.X('x:Q'),
        y=alt.Y('y:Q'),
        x2='x2:Q',
        y2='y2:Q'
    )
    
    # Add correlation text annotation in box
    correlation_text = alt.Chart(pd.DataFrame([{
        'x': point_df['sim_rate'].max() - 7,
        'y': point_df['real_rate'].min() - 7,
        'text': f'95% CI r: {r_lo:.3f}, {r_hi:.3f}'
    }])).mark_text(
        align='center',
        baseline='middle',
        fontSize=16,
        font='Produkt'
    ).encode(
        x=alt.X('x:Q'),
        y=alt.Y('y:Q'),
        text='text:N'
    )
    
    # Combine with correlation box and text
    chart = chart + correlation_box + correlation_text
    
    # Save chart
    chart.save("checkpoint_correlation.html")
    chart.save("checkpoint_correlation.png", scale_factor=2.0)
    chart.save("checkpoint_correlation.pdf", scale_factor=2.0)
    print(f"\nCorrelation: r = {r:.3f}, p = {p_value:.4f}")
    print("\nPlot saved to: checkpoint_correlation.html and checkpoint_correlation.png")


if __name__ == "__main__":
    plot_bayesian_correlation_hardcoded()
