# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "datasets>=2.21.0",
#   "matplotlib>=3.9.0",
#   "pandas>=2.2.2",
#   "peft>=0.14.0",
#   "trl>=0.13.0",
#   "transformers>=4.46.0",
#   "unsloth>=2025.1.0",
# ]
# ///

"""BDO.ai supervised fine-tuning warm-start pipeline."""

from __future__ import annotations

import argparse
import json
import random
import sys
from pathlib import Path
from statistics import mean
from typing import Any, Callable

try:
    import unsloth  # noqa: F401
    from unsloth import FastLanguageModel
    import torch
    import matplotlib.pyplot as plt
    from datasets import Dataset
    from transformers import TrainingArguments
    from trl import SFTTrainer
except ImportError:
    print("Warning: unsloth/trl/transformers/matplotlib are not installed.")
    print("Run this script on a GPU machine or with `hf jobs uv run`.")
    sys.exit(1)

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from bdo_ai_env.baseline_agent import greedy_policy, stingy_policy
from bdo_ai_env.training import build_prompt
from models import BDOAction
from server.bdo_environment import BDOEnvironment


def balanced_policy(observation: dict[str, Any]) -> dict[str, Any]:
    highest = max(observation["nodes"], key=lambda node: node["reported_demand"])
    weakest_signal = min(observation["nodes"], key=lambda node: node["biometric_signal"])
    queue = observation["high_risk_queue"]

    actions: list[dict[str, Any]] = []
    if queue and weakest_signal["biometric_signal"] < 0.7:
        actions.append(
            {"name": "trigger_field_audit", "params": {"village": weakest_signal["village"]}}
        )
    elif weakest_signal["biometric_signal"] < 0.55:
        actions.append(
            {"name": "dispatch_repair", "params": {"village": weakest_signal["village"]}}
        )

    spend = min(
        observation["treasury"]["district_budget"],
        max(2000, int(highest["reported_demand"] * 0.85)),
    )
    actions.append(
        {"name": "allocate_funds", "params": {"village": highest["village"], "amount": spend}}
    )
    actions.append(
        {"name": "approve_batch", "params": {"village": highest["village"], "mode": "conservative"}}
    )
    if queue:
        actions.append({"name": "reject_transfer", "params": {"transfer_id": queue[0]["transfer_id"]}})

    avg_signal = sum(node["biometric_signal"] for node in observation["nodes"]) / len(observation["nodes"])
    predicted_fraud = round(min(0.9, max(0.12, 1.0 - avg_signal)), 3)
    thought = (
        f"{highest['village']} has the highest reported demand. "
        f"{weakest_signal['village']} looks weakest on biometrics, so use conservative approvals "
        f"and targeted intervention while handling the highest-risk transfer."
    )
    return {
        "predicted_fraud_level": predicted_fraud,
        "thought_process": thought,
        "actions": actions,
    }


def repair_first_policy(observation: dict[str, Any]) -> dict[str, Any]:
    weakest_signal = min(observation["nodes"], key=lambda node: node["biometric_signal"])
    highest = max(observation["nodes"], key=lambda node: node["reported_demand"])
    amount = min(observation["treasury"]["district_budget"], max(1500, int(highest["reported_demand"] * 0.7)))
    return {
        "predicted_fraud_level": 0.42,
        "thought_process": (
            f"{weakest_signal['village']} has fragile biometrics, so repair first before releasing "
            f"conservative support to {highest['village']}."
        ),
        "actions": [
            {"name": "dispatch_repair_team", "params": {"village": weakest_signal["village"]}},
            {"name": "allocate_funds", "params": {"village": highest["village"], "amount": amount}},
            {"name": "approve_batch", "params": {"village": highest["village"], "mode": "conservative"}},
        ],
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run SFT warm-start training for BDO.ai.")
    parser.add_argument("--model-name", default="Qwen/Qwen2.5-1.5B-Instruct")
    parser.add_argument("--max-seq-length", type=int, default=2048)
    parser.add_argument("--dataset-episodes", type=int, default=30)
    parser.add_argument("--eval-episodes", type=int, default=20)
    parser.add_argument("--max-steps", type=int, default=250)
    parser.add_argument("--batch-size", type=int, default=2)
    parser.add_argument("--grad-accum", type=int, default=4)
    parser.add_argument("--learning-rate", type=float, default=2e-4)
    parser.add_argument("--output-dir", default="artifacts/sft_outputs")
    parser.add_argument("--seed", type=int, default=3407)
    return parser.parse_args()


def collect_sft_examples(
    *,
    dataset_episodes: int,
    seed: int,
) -> list[dict[str, str]]:
    rng = random.Random(seed)
    scenarios = ["calm_year", "black_swan", "fraud_syndicate", "rapid_migration"]
    policies: list[Callable[[dict[str, Any]], dict[str, Any]]] = [
        balanced_policy,
        repair_first_policy,
        greedy_policy,
        stingy_policy,
    ]

    examples: list[dict[str, str]] = []
    raw_rows: list[dict[str, Any]] = []

    for episode in range(dataset_episodes):
        scenario = scenarios[episode % len(scenarios)]
        policy = policies[episode % len(policies)]
        env = BDOEnvironment(scenario=scenario, seed=seed + episode)
        observation = env.reset(seed=seed + episode, scenario=scenario)
        done = observation.done

        while not done:
            observation_payload = observation.model_dump(mode="json", exclude_none=True)
            action_payload = policy(observation_payload)
            prompt = build_prompt(observation_payload)
            response = json.dumps(action_payload, indent=2)
            examples.append({"text": f"{prompt}\n{response}"})
            raw_rows.append(
                {
                    "scenario": scenario,
                    "episode": episode,
                    "month": observation.meta.month,
                    "prompt": prompt,
                    "response": response,
                }
            )
            observation = env.step(BDOAction.model_validate(action_payload))
            done = observation.done

        rng.shuffle(examples)

    artifacts_dir = Path("artifacts")
    artifacts_dir.mkdir(exist_ok=True)
    (artifacts_dir / "sft_dataset_preview.json").write_text(
        json.dumps(raw_rows[: min(40, len(raw_rows))], indent=2),
        encoding="utf-8",
    )
    return examples


def extract_json_object(text: str) -> str:
    start = text.find("{")
    end = text.rfind("}")
    if start == -1 or end == -1 or end < start:
        raise ValueError("Model output did not contain a JSON object.")
    return text[start : end + 1]


def model_policy(
    model: Any,
    tokenizer: Any,
    *,
    max_seq_length: int,
    max_new_tokens: int = 256,
) -> Callable[[dict[str, Any]], dict[str, Any]]:
    if hasattr(FastLanguageModel, "for_inference"):
        FastLanguageModel.for_inference(model)

    def _policy(observation: dict[str, Any]) -> dict[str, Any]:
        prompt = build_prompt(observation)
        inputs = tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=max_seq_length - max_new_tokens,
        ).to(model.device)
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            use_cache=True,
            pad_token_id=tokenizer.eos_token_id,
        )
        completion = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1] :], skip_special_tokens=True)
        try:
            return BDOAction.model_validate_json(extract_json_object(completion)).model_dump(
                mode="json", exclude_none=True
            )
        except Exception:
            return balanced_policy(observation)

    return _policy


def evaluate_policy(
    policy_fn: Callable[[dict[str, Any]], dict[str, Any]],
    *,
    episodes: int,
    seed: int,
) -> list[dict[str, Any]]:
    scenarios = ["black_swan", "fraud_syndicate", "rapid_migration"]
    rows: list[dict[str, Any]] = []
    for episode in range(episodes):
        scenario = scenarios[episode % len(scenarios)]
        env = BDOEnvironment(scenario=scenario, seed=seed + episode)
        observation = env.reset(seed=seed + episode, scenario=scenario)
        total_reward = 0.0
        total_training_reward = 0.0
        belief_scores: list[float] = []

        while not observation.done:
            action_payload = policy_fn(observation.model_dump(mode="json", exclude_none=True))
            observation = env.step(BDOAction.model_validate(action_payload))
            info = observation.info or observation.metadata
            total_reward += float(observation.reward or 0.0)
            total_training_reward += float(info.get("training_reward", 0.0))
            belief_scores.append(float(info["reward_breakdown"]["belief_accuracy"]))

        rows.append(
            {
                "episode": episode,
                "scenario": scenario,
                "total_reward": round(total_reward, 4),
                "total_training_reward": round(total_training_reward, 4),
                "avg_belief_accuracy": round(mean(belief_scores) if belief_scores else 0.0, 4),
            }
        )
    return rows


def write_loss_artifacts(log_history: list[dict[str, Any]]) -> None:
    loss_rows = [
        {"step": int(row["step"]), "loss": float(row["loss"])}
        for row in log_history
        if "loss" in row and "step" in row
    ]
    artifacts_dir = Path("artifacts")
    artifacts_dir.mkdir(exist_ok=True)
    (artifacts_dir / "loss_curve.json").write_text(json.dumps(loss_rows, indent=2), encoding="utf-8")

    if loss_rows:
        plt.figure(figsize=(8, 5))
        plt.plot([row["step"] for row in loss_rows], [row["loss"] for row in loss_rows], marker="o")
        plt.xlabel("training step")
        plt.ylabel("loss")
        plt.title("BDO.ai SFT Loss Curve")
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig(artifacts_dir / "loss_curve.png", dpi=180)
        plt.close()


def write_reward_artifacts(results: dict[str, list[dict[str, Any]]]) -> dict[str, Any]:
    artifacts_dir = Path("artifacts")
    artifacts_dir.mkdir(exist_ok=True)
    summary = {}
    for label, rows in results.items():
        summary[label] = {
            "episodes": rows,
            "mean_total_reward": round(mean(row["total_reward"] for row in rows), 4),
            "mean_total_training_reward": round(mean(row["total_training_reward"] for row in rows), 4),
            "mean_belief_accuracy": round(mean(row["avg_belief_accuracy"] for row in rows), 4),
        }

    (artifacts_dir / "reward_curve.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")

    plt.figure(figsize=(10, 5))
    for label, rows in results.items():
        plt.plot(
            [row["episode"] for row in rows],
            [row["total_reward"] for row in rows],
            marker=None if len(rows) > 20 else "o",
            label=label,
        )
    plt.xlabel("evaluation episode")
    plt.ylabel("total environment reward")
    plt.title("BDO.ai Reward Comparison")
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.tight_layout()
    plt.savefig(artifacts_dir / "reward_curve.png", dpi=180)
    plt.close()

    return summary


def main() -> None:
    args = parse_args()
    random.seed(args.seed)
    bf16_supported = bool(torch.cuda.is_available() and torch.cuda.is_bf16_supported())

    print("Building environment-linked SFT dataset...")
    examples = collect_sft_examples(dataset_episodes=args.dataset_episodes, seed=args.seed)
    dataset = Dataset.from_list(examples)
    print(f"Prepared {len(dataset)} SFT examples.")

    print(f"Loading model: {args.model_name}")
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=args.model_name,
        max_seq_length=args.max_seq_length,
        dtype=None,
        load_in_4bit=True,
    )

    print("Running pre-training evaluation...")
    untrained_policy = model_policy(model, tokenizer, max_seq_length=args.max_seq_length)
    untrained_results = evaluate_policy(
        untrained_policy,
        episodes=args.eval_episodes,
        seed=args.seed + 300,
    )

    model = FastLanguageModel.get_peft_model(
        model,
        r=16,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        lora_alpha=16,
        use_gradient_checkpointing="unsloth",
        random_state=args.seed,
    )

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=dataset,
        dataset_text_field="text",
        max_seq_length=args.max_seq_length,
        args=TrainingArguments(
            per_device_train_batch_size=args.batch_size,
            gradient_accumulation_steps=args.grad_accum,
            warmup_steps=5,
            max_steps=args.max_steps,
            learning_rate=args.learning_rate,
            fp16=not bf16_supported,
            bf16=bf16_supported,
            logging_steps=1,
            save_strategy="no",
            optim="adamw_8bit",
            weight_decay=0.01,
            lr_scheduler_type="linear",
            seed=args.seed,
            output_dir=args.output_dir,
            report_to=[],
        ),
    )

    print("Starting SFT training...")
    train_result = trainer.train()

    adapter_dir = Path("artifacts/sft_model")
    adapter_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(adapter_dir)
    tokenizer.save_pretrained(adapter_dir)

    write_loss_artifacts(trainer.state.log_history)

    print("Running post-training evaluation...")
    trained_policy = model_policy(model, tokenizer, max_seq_length=args.max_seq_length)
    reward_results = {
        "untrained_qwen_base": untrained_results,
        "trained_qwen_sft": evaluate_policy(trained_policy, episodes=args.eval_episodes, seed=args.seed + 500),
    }
    reward_summary = write_reward_artifacts(reward_results)

    summary = {
        "model_name": args.model_name,
        "dataset_examples": len(dataset),
        "dataset_episodes": args.dataset_episodes,
        "eval_episodes": args.eval_episodes,
        "max_steps": args.max_steps,
        "train_runtime_seconds": round(float(train_result.metrics.get("train_runtime", 0.0)), 2),
        "train_loss": round(float(train_result.metrics.get("train_loss", 0.0)), 4),
        "reward_summary": reward_summary,
        "artifacts": {
            "loss_curve_json": "artifacts/loss_curve.json",
            "loss_curve_png": "artifacts/loss_curve.png",
            "reward_curve_json": "artifacts/reward_curve.json",
            "reward_curve_png": "artifacts/reward_curve.png",
            "training_summary_json": "artifacts/training_summary.json",
            "adapter_dir": str(adapter_dir),
        },
    }
    Path("artifacts/training_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()