#!/usr/bin/env python3
"""Train PPO/SAC trên PositionPlayEnv (curriculum stage 2a — position play).

Chạy từ gốc repo (venv đã kích hoạt):
    python scripts/train_position.py                       # PPO, 300k (smoke test)
    python scripts/train_position.py --total-steps 1000000 # run dài + EvalCallback
    python scripts/train_position.py --pos-coef 0          # ablation: stage-1-trên-env-mới
    python scripts/train_position.py --pos-coef 1.0        # ablation POS_COEF cao

Smoke test 300k trả lời đúng 2 câu hỏi (design doc §7):
    1. pot% có giữ ~15-20% không?  (sập → giảm --pos-coef 0.5 → 0.25)
    2. Q có tăng không?            (không → kiểm tra tần suất gate mở)

Output:
    models/<run>/final_model.zip     — model cuối
    models/<run>/best_model.zip      — model TỐT NHẤT theo EvalCallback (dùng cái này!)
    logs/<run>/monitor.csv           — episode log (nguồn learning curve)
    logs/<run>/learning_curve.png    — 3 panel: reward / rates / Q + spin
    + đánh giá cuối: pot / scratch / Q / makeable% / spin usage vs random baseline
"""

from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

sys.stdout.reconfigure(encoding="utf-8")
sys.stderr.reconfigure(encoding="utf-8")

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

INFO_KEYS = ("potted", "scratch", "contact", "aim_cos", "tgt_disp",
             "pos_q", "b2_potted", "abs_side", "abs_vert")


def make_env(scratch_penalty: float | None = None, pos_coef: float | None = None,
             aim_mode: str | None = None):
    # import trong hàm để subprocess (Windows spawn) tự import lại được
    from poolcoach_rl.envs import PositionPlayEnv

    kwargs = {}
    if scratch_penalty is not None:
        kwargs["scratch_penalty"] = scratch_penalty
    if pos_coef is not None:
        kwargs["pos_coef"] = pos_coef
    if aim_mode is not None:
        kwargs["aim_mode"] = aim_mode
    return PositionPlayEnv(**kwargs)


def plot_learning_curve(monitor_csv: Path, out_png: Path, window: int = 500):
    """3 panel: reward / pot-scratch-contact-aim / Q + spin usage."""
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import numpy as np

    # dòng 1 là metadata JSON, dòng 2 là header
    data = np.genfromtxt(monitor_csv, delimiter=",", names=True, skip_header=1)
    rewards = np.atleast_1d(data["r"])
    if len(rewards) < 2 * window:
        window = max(10, len(rewards) // 10)
    window = max(1, min(window, len(rewards)))

    def roll(x: np.ndarray) -> np.ndarray:
        return np.convolve(x, np.ones(window) / window, mode="valid")

    def col(name: str) -> np.ndarray:
        return (np.atleast_1d(data[name]) if name in data.dtype.names
                else np.zeros_like(rewards))

    episodes = np.arange(window, len(rewards) + 1)
    pot, scratch = roll(col("potted")), roll(col("scratch"))
    contact, aim = roll(col("contact")), roll(col("aim_cos"))
    pos_q, b2_pot = roll(col("pos_q")), roll(col("b2_potted"))
    a_side, a_vert = roll(col("abs_side")), roll(col("abs_vert"))

    # mean Q CHỈ trên các cú pot (Q=0 khi không pot làm loãng đường cong):
    # rolling(pos_q) / rolling(potted) — cẩn thận chia 0 giai đoạn đầu
    with np.errstate(divide="ignore", invalid="ignore"):
        q_on_pot = np.where(pot > 1e-6, pos_q / pot, np.nan)

    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(9, 10), sharex=True)
    ax1.plot(episodes, roll(rewards), lw=1.5)
    ax1.set_ylabel(f"Reward (rolling mean, w={window})")
    ax1.set_title("PoolCoach stage 2a — PositionPlayEnv learning curve")
    ax1.grid(alpha=0.3)

    ax2.plot(episodes, aim, lw=1.2, c="tab:purple", ls="--", alpha=0.8,
             label="Aim-ghost cos (kỳ vọng → 1.0)")
    ax2.plot(episodes, contact, lw=1.2, c="tab:green", alpha=0.7, label="Contact rate")
    ax2.plot(episodes, pot, lw=1.5, label="Pot rate (kỳ vọng GIỮ ~15-20%)")
    ax2.plot(episodes, scratch, lw=1.5, label="Scratch rate")
    ax2.axhline(0.04, ls="--", c="gray", lw=1, label="Random pot (~2-4%)")
    ax2.set_ylabel("Tỉ lệ")
    ax2.legend(loc="upper right")  # loc cố định: tránh warning "best" chậm
    ax2.grid(alpha=0.3)

    ax3.plot(episodes, pos_q, lw=1.5, c="tab:red",
             label="Q mean (mọi ep; 0 khi không pot)")
    ax3.plot(episodes, q_on_pot, lw=1.2, c="tab:orange", alpha=0.8,
             label="Q | pot (mean trên cú pot — câu hỏi chính)")
    ax3.plot(episodes, a_side, lw=1.0, c="tab:blue", ls=":", alpha=0.8,
             label="|side| usage (0-1)")
    ax3.plot(episodes, a_vert, lw=1.0, c="tab:cyan", ls=":", alpha=0.8,
             label="|vert| usage (0-1)")
    ax3.plot(episodes, b2_pot, lw=1.0, c="gray", alpha=0.6,
             label="B2 lucky-pot rate")
    ax3.set_xlabel("Episode (= số cú đánh)")
    ax3.set_ylabel("Q / spin")
    ax3.legend(loc="upper right")
    ax3.grid(alpha=0.3)

    fig.tight_layout()
    fig.savefig(out_png, dpi=150)
    print(f"Learning curve -> {out_png}")


def evaluate(model, n_episodes: int = 200, scratch_penalty: float | None = None,
             pos_coef: float | None = None, obs_slice: int | None = None,
             aim_mode: str | None = None) -> dict:
    """Đánh giá policy deterministic trên PositionPlayEnv.

    obs_slice: nếu set (vd 4) → cắt obs còn N chiều đầu trước khi predict —
    dùng cho model stage 1 (position-blind baseline, design doc §6).
    aim_mode: khớp env lúc train để aim_cos/reward mean so sánh được.
    """
    import numpy as np

    env = make_env(scratch_penalty, pos_coef, aim_mode)
    potted = scratched = contacted = b2_potted = 0
    rewards, aims, disps, qs_on_pot, sides, verts = [], [], [], [], [], []
    makeable = 0
    for _ in range(n_episodes):
        obs, _ = env.reset()
        if obs_slice is not None:
            obs = obs[:obs_slice]
        action, _ = model.predict(obs, deterministic=True)
        _, r, _, _, info = env.step(action)
        rewards.append(r)
        potted += info.get("potted", 0)
        scratched += info.get("scratch", 0)
        contacted += info.get("contact", 0)
        b2_potted += info.get("b2_potted", 0)
        aims.append(info.get("aim_cos", 0.0))
        disps.append(info.get("tgt_disp", 0.0))
        sides.append(info.get("abs_side", 0.0))
        verts.append(info.get("abs_vert", 0.0))
        if info.get("potted", 0) and not info.get("scratch", 0):
            q = info.get("pos_q", 0.0)
            qs_on_pot.append(q)
            makeable += q > 0.5

    n_pot = len(qs_on_pot)
    return {
        "pot_rate": potted / n_episodes,
        "scratch_rate": scratched / n_episodes,
        "contact_rate": contacted / n_episodes,
        "b2_pot_rate": b2_potted / n_episodes,
        "q_mean_on_pot": float(np.mean(qs_on_pot)) if n_pot else 0.0,
        "makeable_rate": makeable / n_pot if n_pot else 0.0,
        "aim_cos_mean": float(np.mean(aims)),
        "tgt_disp_mean": float(np.mean(disps)),
        "abs_side_mean": float(np.mean(sides)),
        "abs_vert_mean": float(np.mean(verts)),
        "reward_mean": float(np.mean(rewards)),
    }


def print_stats(stats: dict):
    print(f"  Pot rate      : {stats['pot_rate']:.1%}   (mục tiêu: GIỮ ~15-20%)")
    print(f"  Scratch rate  : {stats['scratch_rate']:.1%}   (stage 1: ~20%; kỳ vọng giảm dần)")
    print(f"  Contact rate  : {stats['contact_rate']:.1%}")
    print(f"  Q | pot       : {stats['q_mean_on_pot']:.3f}   (câu hỏi chính: có tăng không?)")
    print(f"  Makeable Q>.5 : {stats['makeable_rate']:.1%}   (trên các cú pot)")
    print(f"  B2 lucky pot  : {stats['b2_pot_rate']:.1%}")
    print(f"  Aim cos mean  : {stats['aim_cos_mean']:+.3f}")
    print(f"  Tgt disp mean : {stats['tgt_disp_mean']:.3f} m")
    print(f"  |side| / |vert|: {stats['abs_side_mean']:.2f} / {stats['abs_vert_mean']:.2f} (0-1)")
    print(f"  Reward mean   : {stats['reward_mean']:+.4f}")


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--algo", choices=["ppo", "sac"], default="ppo")
    p.add_argument("--total-steps", type=int, default=300_000)
    p.add_argument("--n-envs", type=int, default=8)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--ent-coef", type=float, default=0.01,
                   help="PPO entropy coef")
    p.add_argument("--scratch-penalty", type=float, default=None,
                   help="ghi đè SCRATCH_PENALTY; 0.5 hay -0.5 đều hiểu -0.5")
    p.add_argument("--pos-coef", type=float, default=None,
                   help="ghi đè POS_COEF (mặc định 0.5 trong env); "
                        "0 = position-blind (ablation), thử 0.25 nếu pot sập")
    p.add_argument("--eval-freq", type=int, default=25_000,
                   help="EvalCallback: eval mỗi N bước/env (bài học mất đỉnh 23.7%)")
    p.add_argument("--init-from", default=None, metavar="MODEL_ZIP",
                   help="warm-start: load weights từ model cùng env (vd fine-tune "
                        "model 1M với --pos-coef khác, khỏi trả lại 200k bước ramp)")
    p.add_argument("--aim-mode", choices=["best_cut", "any"], default=None,
                   help="'any': aim reward max trên mọi lỗ khả thi — mở khoá "
                        "chọn lỗ cho position play (mặc định env: best_cut)")
    p.add_argument("--run-name", default=None)
    p.add_argument("--plot-only", metavar="MONITOR_CSV",
                   help="chỉ vẽ lại curve từ monitor.csv có sẵn rồi thoát")
    args = p.parse_args()

    if args.plot_only:
        csv = Path(args.plot_only)
        plot_learning_curve(csv, csv.parent / "learning_curve.png")
        return

    from functools import partial

    from stable_baselines3 import PPO, SAC
    from stable_baselines3.common.callbacks import EvalCallback
    from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv, VecMonitor

    if args.scratch_penalty is not None:
        args.scratch_penalty = -abs(args.scratch_penalty)
        print(f"[config] scratch_penalty override: {args.scratch_penalty}")
    if args.pos_coef is not None:
        print(f"[config] pos_coef override: {args.pos_coef}")
    if args.aim_mode is not None:
        print(f"[config] aim_mode: {args.aim_mode}")

    run = args.run_name or f"{args.algo}_pos_{time.strftime('%Y%m%d_%H%M%S')}"
    log_dir = ROOT / "logs" / run
    model_dir = ROOT / "models" / run
    log_dir.mkdir(parents=True, exist_ok=True)
    model_dir.mkdir(parents=True, exist_ok=True)

    env_fn = partial(make_env, args.scratch_penalty, args.pos_coef, args.aim_mode)
    venv = SubprocVecEnv([env_fn for _ in range(args.n_envs)])
    venv = VecMonitor(venv, filename=str(log_dir / "monitor.csv"),
                      info_keywords=INFO_KEYS)

    # EvalCallback: env riêng (DummyVecEnv 1 env là đủ — episode 1 bước),
    # deterministic, lưu best model (bài học run 1M: final != best)
    eval_env = VecMonitor(DummyVecEnv([env_fn]))
    eval_cb = EvalCallback(
        eval_env,
        best_model_save_path=str(model_dir),
        log_path=str(log_dir / "eval"),
        eval_freq=max(args.eval_freq // args.n_envs, 1),
        n_eval_episodes=100,
        deterministic=True,
        verbose=1,
    )

    common = dict(env=venv, verbose=1, seed=args.seed,
                  tensorboard_log=str(log_dir))
    if args.init_from:
        # warm-start: giữ weights, ghi đè hyperparams truyền qua kwargs
        cls = PPO if args.algo == "ppo" else SAC
        model = cls.load(args.init_from, ent_coef=args.ent_coef, **common) \
            if args.algo == "ppo" else cls.load(args.init_from, **common)
        print(f"[config] warm-start từ {args.init_from}")
    elif args.algo == "ppo":
        model = PPO("MlpPolicy", n_steps=128, batch_size=256,
                    ent_coef=args.ent_coef, **common)
    else:
        # gradient_steps=8 cân với train_freq=8 (bài học 12/07: mặc định
        # gradient_steps=1 làm SAC update thiếu 8 lần so với chuẩn)
        model = SAC("MlpPolicy", buffer_size=200_000, learning_starts=2_000,
                    train_freq=(8, "step"), gradient_steps=8, **common)

    try:  # progress bar cần tqdm + rich; thiếu thì train không bar
        import tqdm  # noqa: F401
        import rich  # noqa: F401

        progress = True
    except ImportError:
        progress = False

    t0 = time.time()
    model.learn(total_timesteps=args.total_steps, progress_bar=progress,
                callback=eval_cb)
    dt = time.time() - t0
    print(f"\nTrain {args.total_steps:,} bước trong {dt/60:.1f} phút "
          f"({args.total_steps/dt:.0f} steps/s)")

    model.save(model_dir / "final_model")
    print(f"Final model -> {model_dir / 'final_model.zip'}")
    print(f"Best model  -> {model_dir / 'best_model.zip'}  (theo EvalCallback — DÙNG CÁI NÀY)")
    venv.close()
    eval_env.close()

    plot_learning_curve(log_dir / "monitor.csv",
                        log_dir / "learning_curve.png")

    print("\n== Đánh giá FINAL model (deterministic, 200 cú) ==")
    stats = evaluate(model, scratch_penalty=args.scratch_penalty,
                     pos_coef=args.pos_coef, aim_mode=args.aim_mode)
    print_stats(stats)

    best_path = model_dir / "best_model.zip"
    if best_path.exists():
        print("\n== Đánh giá BEST model (deterministic, 200 cú) ==")
        cls = PPO if args.algo == "ppo" else SAC
        best = cls.load(best_path)
        stats = evaluate(best, scratch_penalty=args.scratch_penalty,
                         pos_coef=args.pos_coef, aim_mode=args.aim_mode)
        print_stats(stats)


if __name__ == "__main__":
    main()