"""
inference.py
Run inference against the CACE OpenEnv environment and plot reward curves.

Usage:
    python inference.py
    python inference.py --episodes 20 --model Sannidhay/cace-grpo-model
"""

import os, json, argparse, time
import requests
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np
from datetime import datetime

# ── Config ────────────────────────────────────────────────────────────────────

ENV_URL   = os.environ.get("ENV_BASE_URL", "https://sannidhay-cace-env.hf.space")
HF_TOKEN  = os.environ.get("HF_TOKEN", "")

ACTION_MAP = {
    0: "ALLOW", 1: "REMOVE", 2: "ALLOW_WITH_LABEL",
    3: "ESCALATE", 4: "RESTRICT_DISTRIBUTION",
}
ACTION_COLORS = {
    "ALLOW": "#2ecc71", "REMOVE": "#e74c3c",
    "ALLOW_WITH_LABEL": "#f39c12", "ESCALATE": "#9b59b6",
    "RESTRICT_DISTRIBUTION": "#3498db",
}

# ── Simple LLM Decision Agent ─────────────────────────────────────────────────

def get_decision_from_model(observation: str, model: str = None) -> tuple[str, int]:
    """
    Get moderation decision from model.
    Uses SFT/GRPO model if available, falls back to rule-based.
    """
    obs_upper = observation.upper()

    # Rule-based fallback (works without GPU)
    if "REMOVE" in obs_upper and ("HATE" in obs_upper or "VIOLENCE" in obs_upper or "HARM" in obs_upper):
        if "CULTURAL" in obs_upper and "LEGITIMATE" in obs_upper:
            return "ESCALATE", 3
        return "REMOVE", 1
    elif "ALLOW" in obs_upper and "CULTURAL" in obs_upper:
        return "ALLOW", 0
    elif "HIGH" in obs_upper and "COMPLEX" in obs_upper:
        return "ESCALATE", 3
    else:
        return "ALLOW", 0


# ── OpenEnv client ────────────────────────────────────────────────────────────

class CACEClient:
    def __init__(self, base_url: str):
        self.base_url = base_url.rstrip("/")
        self.session  = requests.Session()
        if HF_TOKEN:
            self.session.headers["Authorization"] = f"Bearer {HF_TOKEN}"

    def health(self) -> bool:
        try:
            r = self.session.get(f"{self.base_url}/health", timeout=15)
            return r.status_code == 200 and r.json().get("status") == "ok"
        except Exception as e:
            print(f"[DEBUG] Health check failed: {e}")
            return False

    def wait_until_ready(self, max_wait: int = 120):
        print(f"[DEBUG] Waiting for server at {self.base_url} ...")
        for i in range(max_wait):
            if self.health():
                print(f"[DEBUG] Server is ready!")
                return True
            time.sleep(1)
            if i % 10 == 9:
                print(f"[DEBUG] Still waiting... ({i+1}s)")
        raise RuntimeError(f"Server not ready after {max_wait}s")

    def reset(self) -> str:
        r = self.session.post(f"{self.base_url}/reset", timeout=60)
        r.raise_for_status()
        obs_r = self.session.get(f"{self.base_url}/observation", timeout=30)
        return obs_r.json()["observation"]

    def step(self, action_int: int) -> dict:
        r = self.session.post(
            f"{self.base_url}/step",
            json={"action_int": action_int},
            timeout=30,
        )
        r.raise_for_status()
        return r.json()

    def info(self) -> dict:
        r = self.session.get(f"{self.base_url}/info", timeout=10)
        return r.json()

    def metrics(self) -> dict:
        r = self.session.get(f"{self.base_url}/metrics", timeout=10)
        return r.json()


# ── Run episodes ──────────────────────────────────────────────────────────────

def run_episodes(env: CACEClient, n_episodes: int, model: str = None) -> list[dict]:
    results = []

    for ep in range(1, n_episodes + 1):
        obs = env.reset()

        decision, action_int = get_decision_from_model(obs, model)
        result = env.step(action_int)

        reward      = result.get("reward", 0.0)
        done        = result.get("done", True)
        info        = result.get("info", {})
        ground_truth= info.get("ground_truth", "?")
        correct     = info.get("correct", decision == ground_truth)
        language    = info.get("language", "Unknown")
        region      = info.get("region", "Unknown")
        breakdown   = info.get("reward_breakdown", {})

        ep_result = {
            "episode":      ep,
            "decision":     decision,
            "ground_truth": ground_truth,
            "reward":       float(reward),
            "correct":      correct,
            "done":         done,
            "language":     language,
            "region":       region,
            "t1_cultural":  breakdown.get("track1_cultural", 0),
            "t2_harm":      breakdown.get("track2_harm", 0),
            "t3_policy":    breakdown.get("track3_policy", 0),
        }
        results.append(ep_result)

        status = "✓" if correct else "✗"
        print(
            f"[STEP] ep={ep} decision={decision} gt={ground_truth} "
            f"reward={reward:+.3f} correct={str(correct).lower()} {status} "
            f"lang={language}"
        )

    return results


# ── Plotting ──────────────────────────────────────────────────────────────────

def plot_results(results: list[dict], save_path: str = "cace_inference_results.png"):
    episodes = [r["episode"] for r in results]
    rewards  = [r["reward"]  for r in results]
    correct  = [r["correct"] for r in results]
    decisions= [r["decision"] for r in results]

    # Running averages
    window = min(5, len(results))
    avg_rewards = np.convolve(rewards, np.ones(window)/window, mode='valid')
    avg_correct = np.convolve([1 if c else 0 for c in correct], np.ones(window)/window, mode='valid')

    fig = plt.figure(figsize=(16, 10))
    fig.patch.set_facecolor('#0f1117')
    gs = gridspec.GridSpec(2, 3, figure=fig, hspace=0.4, wspace=0.35)

    GOLD   = "#FFD700"
    GREEN  = "#2ecc71"
    RED    = "#e74c3c"
    BLUE   = "#3498db"
    PURPLE = "#9b59b6"
    BG     = '#0f1117'
    PANEL  = '#1a1d2e'

    def style_ax(ax, title):
        ax.set_facecolor(PANEL)
        ax.set_title(title, color=GOLD, fontsize=11, fontweight='bold', pad=8)
        ax.tick_params(colors='white')
        ax.xaxis.label.set_color('white')
        ax.yaxis.label.set_color('white')
        for spine in ax.spines.values():
            spine.set_edgecolor('#333')

    # ── Plot 1: Reward per episode ────────────────────────────────────────────
    ax1 = fig.add_subplot(gs[0, :2])
    style_ax(ax1, "Reward per Episode")
    colors = [GREEN if r > 0 else RED for r in rewards]
    ax1.bar(episodes, rewards, color=colors, alpha=0.7, label="Episode reward")
    if len(avg_rewards) > 0:
        x_avg = episodes[window-1:]
        ax1.plot(x_avg, avg_rewards, color=GOLD, linewidth=2.5,
                 label=f"Rolling avg (n={window})", zorder=5)
    ax1.axhline(0, color='white', linewidth=0.5, linestyle='--', alpha=0.3)
    ax1.set_xlabel("Episode")
    ax1.set_ylabel("Reward")
    ax1.legend(facecolor=PANEL, labelcolor='white', fontsize=9)
    ax1.set_ylim(-1.2, 1.2)

    # ── Plot 2: Accuracy ──────────────────────────────────────────────────────
    ax2 = fig.add_subplot(gs[0, 2])
    style_ax(ax2, "Accuracy")
    accuracy = sum(correct) / len(correct)
    ax2.pie(
        [accuracy, 1-accuracy],
        labels=["Correct", "Wrong"],
        colors=[GREEN, RED],
        autopct='%1.0f%%',
        textprops={'color': 'white', 'fontsize': 11},
        startangle=90,
    )
    ax2.set_title(f"Accuracy\n{accuracy*100:.1f}% ({sum(correct)}/{len(correct)})",
                  color=GOLD, fontsize=11, fontweight='bold')

    # ── Plot 3: Decision distribution ─────────────────────────────────────────
    ax3 = fig.add_subplot(gs[1, 0])
    style_ax(ax3, "Decision Distribution")
    from collections import Counter
    dec_counts = Counter(decisions)
    labels = list(dec_counts.keys())
    vals   = list(dec_counts.values())
    bar_colors = [ACTION_COLORS.get(l, BLUE) for l in labels]
    bars = ax3.bar(labels, vals, color=bar_colors, alpha=0.85)
    ax3.set_xticklabels(labels, rotation=20, ha='right', fontsize=8)
    for bar, val in zip(bars, vals):
        ax3.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.1,
                 str(val), ha='center', color='white', fontsize=9)
    ax3.set_ylabel("Count")

    # ── Plot 4: Three-track reward breakdown ──────────────────────────────────
    ax4 = fig.add_subplot(gs[1, 1])
    style_ax(ax4, "3-Track Reward Breakdown (avg)")
    t1_avg = np.mean([r["t1_cultural"] for r in results])
    t2_avg = np.mean([r["t2_harm"]     for r in results])
    t3_avg = np.mean([r["t3_policy"]   for r in results])
    tracks = ["Cultural\n(40%)", "Harm\n(35%)", "Policy\n(25%)"]
    vals   = [t1_avg, t2_avg, t3_avg]
    bar_colors2 = [GOLD, PURPLE, BLUE]
    bars2 = ax4.bar(tracks, vals, color=bar_colors2, alpha=0.85)
    for bar, val in zip(bars2, vals):
        ax4.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.5,
                 f"{val:.1f}", ha='center', color='white', fontsize=10)
    ax4.set_ylabel("Score (0-100)")
    ax4.set_ylim(0, 110)

    # ── Plot 5: Running accuracy ───────────────────────────────────────────────
    ax5 = fig.add_subplot(gs[1, 2])
    style_ax(ax5, "Running Accuracy")
    if len(avg_correct) > 0:
        x_acc = episodes[window-1:]
        ax5.plot(x_acc, avg_correct * 100, color=GREEN, linewidth=2.5)
        ax5.fill_between(x_acc, avg_correct * 100, alpha=0.2, color=GREEN)
    ax5.axhline(50, color='white', linewidth=0.5, linestyle='--', alpha=0.3)
    ax5.set_xlabel("Episode")
    ax5.set_ylabel("Accuracy (%)")
    ax5.set_ylim(0, 105)

    # ── Title ──────────────────────────────────────────────────────────────────
    fig.suptitle(
        "CACE — Cultural Context Arbitration Environment\nInference Results",
        color=GOLD, fontsize=14, fontweight='bold', y=1.01
    )

    plt.savefig(save_path, dpi=150, bbox_inches='tight', facecolor=BG)
    print(f"\n[PLOT] Saved → {save_path}")

    # Also save JSON
    json_path = save_path.replace(".png", ".json")
    with open(json_path, "w") as f:
        json.dump({
            "summary": {
                "episodes": len(results),
                "accuracy": accuracy,
                "avg_reward": float(np.mean(rewards)),
                "total_correct": int(sum(correct)),
            },
            "episodes": results,
        }, f, indent=2)
    print(f"[DATA] Saved → {json_path}")
    plt.show()


# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--episodes", type=int, default=20)
    parser.add_argument("--model",    default=None, help="HF model repo for decisions")
    parser.add_argument("--env-url",  default=ENV_URL)
    parser.add_argument("--output",   default="cace_inference_results.png")
    args = parser.parse_args()

    env = CACEClient(args.env_url)
    env.wait_until_ready()

    info = env.info()
    print(f"\n[START] env={info.get('name','cace')} model={args.model or 'rule-based'}")
    print(f"        action_space={info.get('action_space',{}).get('n')} reward_range={info.get('reward_range')}\n")

    results = run_episodes(env, args.episodes, args.model)

    # Summary
    rewards  = [r["reward"]  for r in results]
    accuracy = sum(r["correct"] for r in results) / len(results)
    print(f"\n[END] episodes={len(results)} accuracy={accuracy:.3f} "
          f"avg_reward={np.mean(rewards):.3f} "
          f"rewards={','.join(f'{r:.2f}' for r in rewards)}")

    metrics = env.metrics()
    print(f"[METRICS] {metrics}")

    plot_results(results, args.output)


if __name__ == "__main__":
    main()