#!/usr/bin/env python3
"""Generate professional visualizations for co-founder presentation.

EB1 Router - Keiro Lab
Design principles (Dean/Hassabis style):
- Clean, minimal aesthetic
- Strategic use of color: blues for our results, greys for baselines
- Clear story at a glance
- No chartjunk
"""

import altair as alt
import pandas as pd
import numpy as np
from pathlib import Path

# Color palette - Keiro brand blues and professional greys
KEIRO_BLUE = "#2563EB"       # Primary - EB1 results
KEIRO_BLUE_LIGHT = "#60A5FA" # Secondary
KEIRO_BLUE_DARK = "#1E40AF"  # Accent
GREY_DARK = "#4B5563"        # Baselines
GREY_MEDIUM = "#9CA3AF"
GREY_LIGHT = "#E5E7EB"
SUCCESS_GREEN = "#10B981"
WARNING_AMBER = "#F59E0B"

# Configure Altair
alt.data_transformers.disable_max_rows()


def save_chart(chart, name, output_dir="charts"):
    """Save chart as HTML and PNG."""
    Path(output_dir).mkdir(exist_ok=True)
    chart.save(f"{output_dir}/{name}.html")
    print(f"Saved: {output_dir}/{name}.html")


# =============================================================================
# Chart 1: Router Performance vs Baselines
# =============================================================================
def chart_router_vs_baselines():
    """frac_oracle comparison: Router vs Random, Best Single, Oracle."""
    data = pd.DataFrame({
        "Strategy": ["Random\nBaseline", "Best Single\n(Opus 4.5)", "EB1 Router", "Oracle\n(Upper Bound)"],
        "frac_oracle": [0.50, 0.68, 0.68, 1.00],
        "is_ours": [False, False, True, False]
    })

    # Add improvement annotation
    data["improvement"] = ["", "", "+36% vs random", ""]

    bars = alt.Chart(data).mark_bar(
        cornerRadiusTopLeft=4,
        cornerRadiusTopRight=4,
    ).encode(
        x=alt.X("Strategy:N",
                axis=alt.Axis(labelAngle=0, title=None, labelFontSize=12),
                sort=["Random\nBaseline", "Best Single\n(Opus 4.5)", "EB1 Router", "Oracle\n(Upper Bound)"]),
        y=alt.Y("frac_oracle:Q",
                scale=alt.Scale(domain=[0, 1.05]),
                axis=alt.Axis(title="frac_oracle", format=".0%", titleFontSize=14)),
        color=alt.condition(
            alt.datum.is_ours,
            alt.value(KEIRO_BLUE),
            alt.value(GREY_MEDIUM)
        )
    )

    labels = alt.Chart(data).mark_text(
        dy=-15, fontSize=14, fontWeight="bold"
    ).encode(
        x=alt.X("Strategy:N", sort=["Random\nBaseline", "Best Single\n(Opus 4.5)", "EB1 Router", "Oracle\n(Upper Bound)"]),
        y="frac_oracle:Q",
        text=alt.Text("frac_oracle:Q", format=".0%"),
        color=alt.condition(
            alt.datum.is_ours,
            alt.value(KEIRO_BLUE_DARK),
            alt.value(GREY_DARK)
        )
    )

    chart = (bars + labels).properties(
        width=500,
        height=350,
        title=alt.TitleParams(
            text="Router Performance vs Baselines",
            subtitle="EB1 Router achieves 68% of oracle performance",
            fontSize=18,
            subtitleFontSize=12,
            subtitleColor=GREY_DARK
        )
    ).configure_axis(
        grid=False
    ).configure_view(
        strokeWidth=0
    )

    return chart


# =============================================================================
# Chart 2: Cost Optimization Impact
# =============================================================================
def chart_cost_optimization():
    """Cost vs Accuracy with routing strategy."""
    data = pd.DataFrame({
        "Strategy": ["Always\nPremium", "Variance-Based\nRouting", "EB1 Router\n(Entropy-Based)"],
        "Accuracy": [0.642, 0.647, 0.650],
        "Cost": [30.0, 13.0, 12.0],
        "is_ours": [False, False, True]
    })

    # Create dual-axis effect with layered charts
    accuracy_bars = alt.Chart(data).mark_bar(
        cornerRadiusTopLeft=4,
        cornerRadiusTopRight=4,
        width=40
    ).encode(
        x=alt.X("Strategy:N",
                axis=alt.Axis(labelAngle=0, title=None),
                sort=["Always\nPremium", "Variance-Based\nRouting", "EB1 Router\n(Entropy-Based)"]),
        y=alt.Y("Accuracy:Q",
                scale=alt.Scale(domain=[0.60, 0.68]),
                axis=alt.Axis(title="Accuracy", format=".1%")),
        color=alt.condition(
            alt.datum.is_ours,
            alt.value(KEIRO_BLUE),
            alt.value(GREY_MEDIUM)
        ),
        xOffset=alt.value(-22)
    )

    cost_bars = alt.Chart(data).mark_bar(
        cornerRadiusTopLeft=4,
        cornerRadiusTopRight=4,
        width=40
    ).encode(
        x=alt.X("Strategy:N", sort=["Always\nPremium", "Variance-Based\nRouting", "EB1 Router\n(Entropy-Based)"]),
        y=alt.Y("Cost:Q",
                scale=alt.Scale(domain=[0, 35]),
                axis=alt.Axis(title="Cost ($/1M tokens)", orient="right")),
        color=alt.condition(
            alt.datum.is_ours,
            alt.value(KEIRO_BLUE_LIGHT),
            alt.value(GREY_LIGHT)
        ),
        xOffset=alt.value(22)
    )

    chart = alt.layer(accuracy_bars, cost_bars).resolve_scale(
        y='independent'
    ).properties(
        width=450,
        height=350,
        title=alt.TitleParams(
            text="Cost Optimization: Same Accuracy, 60% Less Cost",
            subtitle="Blue = Accuracy, Light Blue = Cost",
            fontSize=18,
            subtitleFontSize=12
        )
    )

    return chart


# =============================================================================
# Chart 3: Oracle Selection Diversity
# =============================================================================
def chart_oracle_diversity():
    """Which models does the oracle select?"""
    data = pd.DataFrame({
        "Model": ["GPT-5.1", "Gemini Flash\nLite", "GPT-5.1\nCodex", "Claude\nHaiku",
                  "Claude\nOpus", "GPT-5\nNano", "Gemini 3\nPro"],
        "Selection": [29.6, 26.7, 19.5, 12.3, 4.3, 3.9, 3.7],
        "tier": ["premium", "cheap", "premium", "mid", "premium", "cheap", "premium"]
    })

    # Sort by selection rate
    data = data.sort_values("Selection", ascending=True)

    color_scale = alt.Scale(
        domain=["cheap", "mid", "premium"],
        range=[SUCCESS_GREEN, WARNING_AMBER, KEIRO_BLUE]
    )

    bars = alt.Chart(data).mark_bar(
        cornerRadiusTopRight=4,
        cornerRadiusBottomRight=4,
    ).encode(
        y=alt.Y("Model:N",
                sort=alt.EncodingSortField(field="Selection", order="ascending"),
                axis=alt.Axis(title=None, labelFontSize=11)),
        x=alt.X("Selection:Q",
                scale=alt.Scale(domain=[0, 35]),
                axis=alt.Axis(title="Oracle Selection Rate (%)", titleFontSize=12)),
        color=alt.Color("tier:N",
                       scale=color_scale,
                       legend=alt.Legend(title="Cost Tier", orient="bottom"))
    )

    labels = alt.Chart(data).mark_text(
        dx=5, fontSize=12, fontWeight="bold", align="left"
    ).encode(
        y=alt.Y("Model:N", sort=alt.EncodingSortField(field="Selection", order="ascending")),
        x="Selection:Q",
        text=alt.Text("Selection:Q", format=".1f"),
        color=alt.value(GREY_DARK)
    )

    chart = (bars + labels).properties(
        width=450,
        height=300,
        title=alt.TitleParams(
            text="Oracle Model Selection Distribution",
            subtitle="Cheap models (green) win 30%+ of the time",
            fontSize=18,
            subtitleFontSize=12
        )
    ).configure_axis(
        grid=False
    ).configure_view(
        strokeWidth=0
    )

    return chart


# =============================================================================
# Chart 4: Uncertainty Calibration
# =============================================================================
def chart_uncertainty_calibration():
    """Accuracy by entropy quartile - shows model is well-calibrated."""
    data = pd.DataFrame({
        "Entropy": ["Low\n(Confident)", "Medium", "High\n(Uncertain)"],
        "Accuracy": [17.9, 7.0, 7.0],
        "AvgModelsCorrect": [2.7, 3.4, 3.8],
        "is_confident": [True, False, False]
    })

    bars = alt.Chart(data).mark_bar(
        cornerRadiusTopLeft=4,
        cornerRadiusTopRight=4,
    ).encode(
        x=alt.X("Entropy:N",
                axis=alt.Axis(labelAngle=0, title="Prediction Entropy", titleFontSize=12),
                sort=["Low\n(Confident)", "Medium", "High\n(Uncertain)"]),
        y=alt.Y("Accuracy:Q",
                scale=alt.Scale(domain=[0, 22]),
                axis=alt.Axis(title="Router Accuracy (%)", titleFontSize=12)),
        color=alt.condition(
            alt.datum.is_confident,
            alt.value(KEIRO_BLUE),
            alt.value(GREY_MEDIUM)
        )
    )

    labels = alt.Chart(data).mark_text(
        dy=-12, fontSize=14, fontWeight="bold"
    ).encode(
        x=alt.X("Entropy:N", sort=["Low\n(Confident)", "Medium", "High\n(Uncertain)"]),
        y="Accuracy:Q",
        text=alt.Text("Accuracy:Q", format=".1f"),
        color=alt.condition(
            alt.datum.is_confident,
            alt.value(KEIRO_BLUE_DARK),
            alt.value(GREY_DARK)
        )
    )

    # Add subtitle about what this means
    chart = (bars + labels).properties(
        width=400,
        height=300,
        title=alt.TitleParams(
            text="Uncertainty Calibration: Confident = Accurate",
            subtitle="Low entropy predictions are 2.5x more accurate",
            fontSize=16,
            subtitleFontSize=12,
            subtitleColor=GREY_DARK
        )
    ).configure_axis(
        grid=False
    ).configure_view(
        strokeWidth=0
    )

    return chart


# =============================================================================
# Chart 5: Embedding Comparison
# =============================================================================
def chart_embedding_comparison():
    """Performance across different embedding methods."""
    data = pd.DataFrame({
        "Embedding": ["OpenAI\n512d", "OpenAI\n256d", "Gemini\n768d", "Gemini\n256d", "MiniLM\n384d"],
        "frac_oracle": [68.0, 65.2, 64.8, 62.1, 58.3],
        "is_best": [True, False, False, False, False]
    })

    bars = alt.Chart(data).mark_bar(
        cornerRadiusTopLeft=4,
        cornerRadiusTopRight=4,
    ).encode(
        x=alt.X("Embedding:N",
                axis=alt.Axis(labelAngle=0, title=None),
                sort=["OpenAI\n512d", "OpenAI\n256d", "Gemini\n768d", "Gemini\n256d", "MiniLM\n384d"]),
        y=alt.Y("frac_oracle:Q",
                scale=alt.Scale(domain=[50, 72]),
                axis=alt.Axis(title="frac_oracle (%)", titleFontSize=12)),
        color=alt.condition(
            alt.datum.is_best,
            alt.value(KEIRO_BLUE),
            alt.value(GREY_MEDIUM)
        )
    )

    labels = alt.Chart(data).mark_text(
        dy=-12, fontSize=13, fontWeight="bold"
    ).encode(
        x=alt.X("Embedding:N", sort=["OpenAI\n512d", "OpenAI\n256d", "Gemini\n768d", "Gemini\n256d", "MiniLM\n384d"]),
        y="frac_oracle:Q",
        text=alt.Text("frac_oracle:Q", format=".1f"),
        color=alt.condition(
            alt.datum.is_best,
            alt.value(KEIRO_BLUE_DARK),
            alt.value(GREY_DARK)
        )
    )

    chart = (bars + labels).properties(
        width=450,
        height=300,
        title=alt.TitleParams(
            text="Embedding Comparison: Matryoshka 512d Wins",
            subtitle="OpenAI text-embedding-3-large with dimension truncation",
            fontSize=16,
            subtitleFontSize=12
        )
    ).configure_axis(
        grid=False
    ).configure_view(
        strokeWidth=0
    )

    return chart


# =============================================================================
# Chart 6: Key Metrics Dashboard
# =============================================================================
def chart_key_metrics():
    """Single-view dashboard of key metrics."""
    data = pd.DataFrame({
        "Metric": ["frac_oracle", "NDCG@3", "ECE", "Cost Savings"],
        "Value": [68.0, 70.8, 88.7, 57.0],  # ECE inverted for "better = higher"
        "Target": [50.0, 50.0, 85.0, 0.0],
        "color": [KEIRO_BLUE, KEIRO_BLUE, SUCCESS_GREEN, SUCCESS_GREEN]
    })

    bars = alt.Chart(data).mark_bar(
        cornerRadiusTopLeft=4,
        cornerRadiusTopRight=4,
    ).encode(
        x=alt.X("Metric:N", axis=alt.Axis(labelAngle=0, title=None)),
        y=alt.Y("Value:Q",
                scale=alt.Scale(domain=[0, 100]),
                axis=alt.Axis(title="Score (%)", titleFontSize=12)),
        color=alt.Color("color:N", scale=None)
    )

    # Target line for baselines
    target = alt.Chart(data).mark_tick(
        color=GREY_DARK,
        thickness=2,
        size=35
    ).encode(
        x="Metric:N",
        y="Target:Q"
    )

    labels = alt.Chart(data).mark_text(
        dy=-12, fontSize=14, fontWeight="bold", color=GREY_DARK
    ).encode(
        x="Metric:N",
        y="Value:Q",
        text=alt.Text("Value:Q", format=".0f")
    )

    chart = (bars + target + labels).properties(
        width=400,
        height=300,
        title=alt.TitleParams(
            text="EB1 Router: Key Metrics",
            subtitle="Grey tick = baseline/target",
            fontSize=16,
            subtitleFontSize=12
        )
    ).configure_axis(
        grid=False
    ).configure_view(
        strokeWidth=0
    )

    return chart


# =============================================================================
# Main: Generate all charts
# =============================================================================
def main():
    output_dir = Path(__file__).parent.parent / "charts"
    output_dir.mkdir(exist_ok=True)

    print("Generating presentation charts...")
    print(f"Output directory: {output_dir}")
    print()

    charts = [
        ("01_router_vs_baselines", chart_router_vs_baselines()),
        ("02_cost_optimization", chart_cost_optimization()),
        ("03_oracle_diversity", chart_oracle_diversity()),
        ("04_uncertainty_calibration", chart_uncertainty_calibration()),
        ("05_embedding_comparison", chart_embedding_comparison()),
        ("06_key_metrics", chart_key_metrics()),
    ]

    for name, chart in charts:
        filepath = output_dir / f"{name}.html"
        chart.save(str(filepath))
        print(f"  Saved: {filepath}")

    print()
    print("All charts generated. Open the HTML files in a browser to view.")
    print("Tip: Use browser print (Cmd+P) to save as PDF for presentation.")


if __name__ == "__main__":
    main()
