#!/usr/bin/env python
"""
Generate paper figures from experimental results.
"""

import json
import matplotlib

matplotlib.use("Agg")  # Non-interactive backend
import matplotlib.pyplot as plt
import numpy as np


def load_statistics():
    """Load statistics from JSON."""
    with open("results/statistics.json", "r") as f:
        return json.load(f)


def load_all_results():
    """Load all raw results."""
    algorithms = ["gradient", "flooding", "random"]
    all_data = {}

    for algo in algorithms:
        with open(f"results/{algo}/all_seeds.json", "r") as f:
            all_data[algo] = json.load(f)

    return all_data


def plot_pdr_vs_airtime(all_data, stats):
    """Generate PDR vs Airtime plot."""
    fig, ax = plt.subplots(figsize=(8, 6))

    colors = {"gradient": "#2E86AB", "flooding": "#E94F37", "random": "#7D8491"}
    markers = {"gradient": "o", "flooding": "s", "random": "^"}

    for algo in ["gradient", "flooding", "random"]:
        algo_stats = stats[algo]

        # Get mean and CI
        pdr_mean = algo_stats["pdr"]["mean"]
        pdr_ci = algo_stats["pdr"]["ci_95"]
        air_mean = algo_stats["airtime_usage_percent"]["mean"]
        air_ci = algo_stats["airtime_usage_percent"]["ci_95"]

        # Plot with error bars
        ax.errorbar(
            air_mean,
            pdr_mean,
            xerr=air_ci,
            yerr=pdr_ci,
            fmt=markers[algo],
            color=colors[algo],
            markersize=12,
            capsize=5,
            capthick=2,
            linewidth=2,
            label=f"{algo.capitalize()} ({pdr_mean:.1f}%, {air_mean:.1f}%)",
        )

        # Add individual points (semi-transparent)
        for r in all_data[algo]:
            if "error" not in r:
                ax.scatter(
                    r["airtime_usage_percent"],
                    r["pdr"],
                    alpha=0.15,
                    color=colors[algo],
                    s=20,
                )

    ax.set_xlabel("Airtime Usage (%)", fontsize=12)
    ax.set_ylabel("Packet Delivery Ratio (%)", fontsize=12)
    ax.set_title(
        "PDR vs Channel Airtime Usage\n(50 seeds per algorithm, error bars = 95% CI)",
        fontsize=14,
    )
    ax.legend(loc="lower right", fontsize=10)
    ax.grid(True, alpha=0.3)
    ax.set_xlim(0, 110)
    ax.set_ylim(0, 35)

    # Add annotation for efficiency frontier
    ax.annotate(
        "Efficiency\nFrontier",
        xy=(40, 12),
        fontsize=10,
        style="italic",
        bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5),
    )

    plt.tight_layout()
    plt.savefig("figures/pdr_vs_airtime.pdf", dpi=300)
    plt.savefig("figures/pdr_vs_airtime.png", dpi=150)
    print("Saved: figures/pdr_vs_airtime.pdf")
    plt.close()


def plot_pdr_vs_tx_cost(all_data, stats):
    """Generate PDR vs TX Cost plot."""
    fig, ax = plt.subplots(figsize=(8, 6))

    colors = {"gradient": "#2E86AB", "flooding": "#E94F37", "random": "#7D8491"}
    markers = {"gradient": "o", "flooding": "s", "random": "^"}

    for algo in ["gradient", "flooding", "random"]:
        algo_stats = stats[algo]

        pdr_mean = algo_stats["pdr"]["mean"]
        pdr_ci = algo_stats["pdr"]["ci_95"]
        tx_mean = algo_stats["tx_per_success"]["mean"]
        tx_ci = algo_stats["tx_per_success"]["ci_95"]

        ax.errorbar(
            tx_mean,
            pdr_mean,
            xerr=tx_ci,
            yerr=pdr_ci,
            fmt=markers[algo],
            color=colors[algo],
            markersize=12,
            capsize=5,
            capthick=2,
            linewidth=2,
            label=f"{algo.capitalize()} (PDR: {pdr_mean:.1f}%)",
        )

        # Add individual points
        for r in all_data[algo]:
            if "error" not in r and r["tx_per_success"] > 0:
                ax.scatter(
                    r["tx_per_success"], r["pdr"], alpha=0.15, color=colors[algo], s=20
                )

    ax.set_xlabel("Transmission Cost (TX per Successful Delivery)", fontsize=12)
    ax.set_ylabel("Packet Delivery Ratio (%)", fontsize=12)
    ax.set_title(
        "PDR vs Transmission Cost\n(50 seeds per algorithm, error bars = 95% CI)",
        fontsize=14,
    )
    ax.legend(loc="lower right", fontsize=10)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig("figures/pdr_vs_tx_cost.pdf", dpi=300)
    plt.savefig("figures/pdr_vs_tx_cost.png", dpi=150)
    print("Saved: figures/pdr_vs_tx_cost.pdf")
    plt.close()


def plot_comparison_bar(stats):
    """Generate comparison bar chart."""
    fig, axes = plt.subplots(1, 3, figsize=(12, 4))

    algorithms = ["gradient", "flooding", "random"]
    colors = ["#2E86AB", "#E94F37", "#7D8491"]

    metrics = [
        ("pdr", "PDR (%)"),
        ("airtime_usage_percent", "Airtime (%)"),
        ("tx_per_success", "TX/Success"),
    ]

    for idx, (metric, label) in enumerate(metrics):
        ax = axes[idx]
        values = [stats[a][metric]["mean"] for a in algorithms]
        errors = [stats[a][metric]["ci_95"] for a in algorithms]

        bars = ax.bar(
            algorithms, values, yerr=errors, color=colors, capsize=5, alpha=0.8
        )
        ax.set_ylabel(label, fontsize=11)
        ax.set_title(label, fontsize=12)
        ax.grid(True, alpha=0.3, axis="y")

        # Add value labels
        for bar, val in zip(bars, values):
            ax.text(
                bar.get_x() + bar.get_width() / 2,
                bar.get_height() + errors[algorithms.index(bar.get_x())],
                f"{val:.1f}",
                ha="center",
                va="bottom",
                fontsize=10,
            )

    plt.tight_layout()
    plt.savefig("figures/comparison_bar.pdf", dpi=300)
    plt.savefig("figures/comparison_bar.png", dpi=150)
    print("Saved: figures/comparison_bar.pdf")
    plt.close()


def main():
    print("Loading results...")
    stats = load_statistics()
    all_data = load_all_results()

    print("Generating figures...")
    plot_pdr_vs_airtime(all_data, stats)
    plot_pdr_vs_tx_cost(all_data, stats)
    plot_comparison_bar(stats)

    print("\nAll figures saved to figures/")


if __name__ == "__main__":
    main()