import argparse
import json
from pathlib import Path


def main():
    parser = argparse.ArgumentParser(description="Plot OversightForge metrics from rollout artifact")
    parser.add_argument("--input", default="artifacts/training_rollout.json")
    parser.add_argument("--eval-input", default="artifacts/level4_eval.json")
    parser.add_argument("--output-dir", default="artifacts/plots")
    args = parser.parse_args()

    try:
        import matplotlib.pyplot as plt
    except Exception as exc:
        raise RuntimeError(
            "matplotlib is required for plotting. Install with: pip install matplotlib"
        ) from exc

    data = json.loads(Path(args.input).read_text(encoding="utf-8"))
    results = data.get("results", [])

    rewards = [float(ep.get("reward", 0.0)) for ep in results]
    discovered = [1 if (ep.get("reward_breakdown") or {}).get("fraud_discovered") else 0 for ep in results]
    code_ok = [1 if (ep.get("reward_breakdown") or {}).get("correct_violation_code") else 0 for ep in results]
    syco_ok = [1 if (ep.get("reward_breakdown") or {}).get("sycophancy_resisted") else 0 for ep in results]
    calls = [int(ep.get("steps_taken", 0)) for ep in results]

    def rolling_pct(values, window=10):
        out = []
        for idx in range(len(values)):
            start = max(0, idx - window + 1)
            section = values[start : idx + 1]
            out.append((sum(section) / max(1, len(section))) * 100.0)
        return out

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    # 1. reward curve
    plt.figure(figsize=(9, 4))
    plt.plot(rewards, color="#1f6f78")
    plt.title("Reward Curve")
    plt.xlabel("Episode")
    plt.ylabel("Reward")
    plt.tight_layout()
    plt.savefig(out_dir / "reward_curve.png", dpi=150)
    plt.close()

    # 2. fraud discovery rate (rolling)
    plt.figure(figsize=(9, 4))
    plt.plot(rolling_pct(discovered), color="#2a7f4f")
    plt.title("Fraud Discovery Rate (Rolling 10)")
    plt.xlabel("Episode")
    plt.ylabel("%")
    plt.ylim(0, 100)
    plt.tight_layout()
    plt.savefig(out_dir / "fraud_discovery_rate.png", dpi=150)
    plt.close()

    # 3. AML code accuracy (rolling)
    plt.figure(figsize=(9, 4))
    plt.plot(rolling_pct(code_ok), color="#c85c2d")
    plt.title("AML Code Accuracy (Rolling 10)")
    plt.xlabel("Episode")
    plt.ylabel("%")
    plt.ylim(0, 100)
    plt.tight_layout()
    plt.savefig(out_dir / "aml_code_accuracy.png", dpi=150)
    plt.close()

    # 4. sycophancy resistance (rolling)
    plt.figure(figsize=(9, 4))
    plt.plot(rolling_pct(syco_ok), color="#5a3e36")
    plt.title("Sycophancy Resistance (Rolling 10)")
    plt.xlabel("Episode")
    plt.ylabel("%")
    plt.ylim(0, 100)
    plt.tight_layout()
    plt.savefig(out_dir / "sycophancy_resistance.png", dpi=150)
    plt.close()

    # 5. detection calls
    plt.figure(figsize=(9, 4))
    plt.plot(calls, color="#444")
    plt.title("Detection Calls Per Episode")
    plt.xlabel("Episode")
    plt.ylabel("Calls")
    plt.tight_layout()
    plt.savefig(out_dir / "detection_calls.png", dpi=150)
    plt.close()

    eval_path = Path(args.eval_input)
    if eval_path.exists():
        eval_data = json.loads(eval_path.read_text(encoding="utf-8"))
        eval_results = eval_data.get("results", [])
        eval_discovery = [1 if (ep.get("reward_breakdown") or {}).get("fraud_discovered") else 0 for ep in eval_results]
        train_rate = (sum(discovered) / max(1, len(discovered))) * 100.0
        eval_rate = (sum(eval_discovery) / max(1, len(eval_discovery))) * 100.0

        plt.figure(figsize=(7, 4))
        plt.bar(["Train", "Level4 Held-Out"], [train_rate, eval_rate], color=["#1f6f78", "#c85c2d"])
        plt.title("Generalization: Train vs Held-Out Discovery Rate")
        plt.ylabel("%")
        plt.ylim(0, 100)
        plt.tight_layout()
        plt.savefig(out_dir / "generalization_train_vs_heldout.png", dpi=150)
        plt.close()

    print(f"Saved plots to: {out_dir}")


if __name__ == "__main__":
    main()