"""
train_kaggle.py
================
Standalone training script for WeatherForecastEnv, meant to run on Kaggle
(free GPU, no session-length hyperparameter-search restrictions like Colab).
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
import time
from pathlib import Path
from typing import Optional, Tuple

logger = logging.getLogger("train_kaggle")
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
)


def _add_file_logging(out_dir: Path) -> None:
    formatter = logging.Formatter("%(asctime)s | %(levelname)s | %(message)s")
    file_handler = logging.FileHandler(str(out_dir / "training.log"))
    file_handler.setFormatter(formatter)
    logging.getLogger().addHandler(file_handler)
    logger.info("File logging enabled: %s", out_dir / "training.log")


BEST_HYPERPARAMETERS = dict(
    learning_rate=6.916624987609979e-05,
    ent_coef=0.08779238696445962,
    hidden_size=128,
    spatial_size=12,
    n_steps=4096,
)


def resolve_max_steps(
    n_zones: int,
    budget_mode: str,
    max_steps_arg: Optional[int],
) -> Tuple[int, str]:
    n = max(1, int(n_zones))
    full_ceiling = n + 1

    mode = (budget_mode or "triage").strip().lower()
    if mode not in ("full", "scarce", "triage"):
        raise ValueError(
            f"budget_mode must be one of full|scarce|triage, got {budget_mode!r}"
        )

    if mode == "full":
        derived = full_ceiling
    elif mode == "scarce":
        derived = n
    else:  # triage
        derived = max(1, n - 1)

    if max_steps_arg is not None and int(max_steps_arg) > 0:
        ms = int(max_steps_arg)
        note = f"explicit --max-steps={ms} (budget-mode={mode} would have been {derived})"
        if ms >= full_ceiling:
            logger.warning(
                "BUDGET WARNING: max_steps=%d >= n_zones+1=%d. Full tour is "
                "feasible and reward-optimal (unvisited_zone_penalty → 0). "
                "The policy is NOT forced to differentiate which zone to "
                "inspect. For allocation skill use --budget-mode triage "
                "(or scarce) without overriding --max-steps, or set "
                "--max-steps < %d.",
                ms, full_ceiling, full_ceiling,
            )
        return ms, note

    return derived, f"budget-mode={mode} → max_steps={derived} (full ceiling={full_ceiling})"


def _add_dataset_to_path(dataset_dir: Optional[str]) -> None:
    if dataset_dir:
        p = str(Path(dataset_dir).resolve())
        if p not in sys.path:
            sys.path.insert(0, p)
            logger.info("Added to sys.path: %s", p)
    here = str(Path(__file__).resolve().parent)
    if here not in sys.path:
        sys.path.insert(0, here)


def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Single-config MaskablePPO trainer for WeatherForecastEnv (Kaggle)."
    )
    p.add_argument("--dataset-dir", default=None,
                    help="Path to the uploaded Kaggle dataset directory containing the .py files.")
    p.add_argument("--out", default="./run", help="Output directory for checkpoints/logs/final model.")
    p.add_argument("--resume-from", default=None, help="Path to a checkpoint .zip to resume from.")

    p.add_argument("--steps", type=int, default=50_000,
                    help="Total training timesteps. Start small (e.g. 100_000-300_000) to validate, "
                         "then scale up for a real run.")
    p.add_argument("--n-zones", type=int, default=2,
                    help="Number of zones. Default 2 (matches real-eval design window).")
    p.add_argument(
        "--budget-mode",
        choices=("full", "scarce", "triage"),
        default="triage",
        help="How tight the inspection budget is relative to n_zones. "
             "full=n_zones+1, scarce=n_zones, triage=max(1,n_zones-1). "
             "Default triage forces leaving ≥1 zone unvisited.",
    )
    p.add_argument(
        "--max-steps",
        type=int,
        default=None,
        help="Explicit episode length cap. If omitted, derived from --budget-mode. "
             "Setting this >= n_zones+1 disables triage pressure (WARNING logged).",
    )
    p.add_argument("--seed", type=int, default=42)

    p.add_argument("--lr", type=float, default=BEST_HYPERPARAMETERS["learning_rate"])
    p.add_argument("--ent-coef", type=float, default=BEST_HYPERPARAMETERS["ent_coef"])
    p.add_argument("--hidden-size", type=int, default=BEST_HYPERPARAMETERS["hidden_size"])
    p.add_argument("--spatial-size", type=int, default=BEST_HYPERPARAMETERS["spatial_size"])
    p.add_argument("--n-steps", type=int, default=BEST_HYPERPARAMETERS["n_steps"],
                    help="PPO rollout buffer size. If you change max_steps a lot, consider "
                         "resizing this to roughly 15-25x max_steps.")
    p.add_argument("--batch-size", type=int, default=None,
                    help="Defaults to max(32, n_steps // 8) if not given.")

    p.add_argument("--eval-freq", type=int, default=10_000)
    p.add_argument("--eval-episodes", type=int, default=20)
    p.add_argument("--checkpoint-freq", type=int, default=25_000)
    p.add_argument("--regression-check-freq", type=int, default=5_000)

    p.add_argument("--device", default="auto", help="'cpu', 'cuda', or 'auto'.")
    p.add_argument(
        "--clean-episode-ratio",
        type=float,
        default=0.90,
        help="Fraction of synthetic episodes with no regional hazard (match eval).",
    )
    p.add_argument(
        "--event-spatial-correlation",
        type=float,
        default=0.85,
        help="P(zone dirty | regional event). High → El Niño-style joint risk.",
    )
    p.add_argument(
        "--precip-scale",
        type=float,
        default=40.0,
        help="Fixed (non-learned) divisor applied to forecast_precip before "
             "it reaches the GRU extractor's forecast_proj/zone_encoder. "
             "forecast_precip runs roughly [0, 80] while zone_belief runs "
             "roughly [0, 0.3] with no normalization layer between them; "
             "left unscaled, precip's raw magnitude can suppress the "
             "smaller-but-more-reliable belief signal during optimization, "
             "independent of which feature is actually more informative. "
             "40.0 is the value that produced the validated single-dirty "
             "selection-accuracy results (see model card) -- confirmed "
             "working, not confirmed optimal. Recalibrate against your own "
             "separability probe's dirty-zone precip levels if your event "
             "injection magnitudes differ from the defaults.",
    )
    return p.parse_args()


def build_envs(args: argparse.Namespace):
    from weather_forecast_env import make_weather_env
    from zone_observation import ForecastConfig
    from stable_baselines3.common.monitor import Monitor

    train_config = ForecastConfig(
        n_zones=args.n_zones,
        max_steps=args.max_steps,
        seed=args.seed,
        clean_episode_ratio=args.clean_episode_ratio,
        event_spatial_correlation=args.event_spatial_correlation,
    )
    eval_config = ForecastConfig(
        n_zones=args.n_zones,
        max_steps=args.max_steps,
        seed=args.seed + 10_000,
        clean_episode_ratio=args.clean_episode_ratio,
        event_spatial_correlation=args.event_spatial_correlation,
    )
    train_env = Monitor(make_weather_env(train_config))
    eval_env = Monitor(make_weather_env(eval_config))
    return train_env, eval_env


def build_model(args: argparse.Namespace, train_env):
    from sb3_contrib import MaskablePPO
    from gru_weather_policy import (
        ZoneEquivariantMaskablePolicy,
        create_gru_weather_policy_kwargs,
    )

    if args.resume_from:
        logger.info("Resuming from checkpoint: %s", args.resume_from)
        return MaskablePPO.load(args.resume_from, env=train_env, device=args.device)

    policy_kwargs = create_gru_weather_policy_kwargs(
        hidden_size=args.hidden_size,
        spatial_output_size=args.spatial_size,
        features_dim=args.hidden_size * 2,
        basin_context_hidden=12,
        precip_scale=args.precip_scale,
    )
    batch_size = args.batch_size or max(32, args.n_steps // 8)

    try:
        import tensorboard  # noqa: F401
        tb_log = str(Path(args.out) / "tensorboard")
    except ImportError:
        logger.warning(
            "tensorboard not installed -- continuing without TensorBoard logs "
            "(install with `pip install tensorboard` if you want them)."
        )
        tb_log = None

    return MaskablePPO(
        ZoneEquivariantMaskablePolicy,
        train_env,
        learning_rate=args.lr,
        ent_coef=args.ent_coef,
        policy_kwargs=policy_kwargs,
        n_steps=args.n_steps,
        batch_size=batch_size,
        gamma=0.98,
        gae_lambda=0.95,
        clip_range=0.2,
        device=args.device,
        verbose=1,
        tensorboard_log=tb_log,
    )


class RegressionWatchCallback:

    def __init__(
        self,
        total_timesteps: int,
        check_freq: int = 5_000,
        after_frac: float = 0.2,
        ep_len_threshold: float = 1.5,
    ):
        from stable_baselines3.common.callbacks import BaseCallback
        import numpy as np

        self._np = np
        self._BaseCallback = BaseCallback
        self.total_timesteps = total_timesteps
        self.check_freq = check_freq
        self.after_frac = after_frac
        self.ep_len_threshold = ep_len_threshold
        self._instance = self._build_instance()

    def _build_instance(self):
        np = self._np
        outer = self

        class _Impl(self._BaseCallback):
            def __init__(self):
                super().__init__()
                self._last_check = 0
                self.history = []  # (timestep, ep_len_mean, entropy_loss)

            def _on_step(self) -> bool:
                if self.num_timesteps - self._last_check < outer.check_freq:
                    return True
                self._last_check = self.num_timesteps

                ep_lens = (
                    [ep["l"] for ep in self.model.ep_info_buffer]
                    if self.model.ep_info_buffer
                    else []
                )
                ep_len_mean = float(np.mean(ep_lens)) if ep_lens else float("nan")

                entropy = None
                if self.model.logger is not None:
                    entropy = self.model.logger.name_to_value.get("train/entropy_loss")

                self.history.append((self.num_timesteps, ep_len_mean, entropy))

                frac = self.num_timesteps / max(1, outer.total_timesteps)
                if frac >= outer.after_frac and ep_lens and ep_len_mean < outer.ep_len_threshold:
                    logger.warning(
                        "REGRESSION WARNING at step %d: ep_len_mean=%.2f after %.0f%% "
                        "of training. This matches the original 'terminate immediately' "
                        "collapse pattern -- worth stopping to check config/reward before "
                        "trusting the rest of this run.",
                        self.num_timesteps, ep_len_mean, frac * 100,
                    )
                return True

        return _Impl()

    @property
    def instance(self):
        return self._instance


def run_training(args: argparse.Namespace) -> None:
    resolved_ms, budget_note = resolve_max_steps(
        args.n_zones, args.budget_mode, args.max_steps
    )
    args.max_steps = resolved_ms  # mutate so build_envs / summary see the real value

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

    logger.info(
        "Config: n_zones=%d max_steps=%d (%s) steps=%d clean=%.3f rho=%.3f "
        "lr=%.3g ent_coef=%.3g hidden_size=%d spatial_size=%d n_steps=%d "
        "precip_scale=%.3g device=%s",
        args.n_zones, args.max_steps, budget_note, args.steps,
        args.clean_episode_ratio, args.event_spatial_correlation,
        args.lr, args.ent_coef,
        args.hidden_size, args.spatial_size, args.n_steps,
        args.precip_scale,
        args.device,
    )
    logger.info(
        "Budget pressure: full_ceiling=%d  resolved_max_steps=%d  "
        "must_skip_zones=%s",
        args.n_zones + 1,
        args.max_steps,
        "yes" if args.max_steps < args.n_zones + 1 else "no (full tour allowed)",
    )

    train_env, eval_env = build_envs(args)
    model = build_model(args, train_env)

    from sb3_contrib.common.maskable.callbacks import MaskableEvalCallback
    try:
        from train_curriculum import CheckpointCallback as _ProjectCheckpointCallback
        checkpoint_cb = _ProjectCheckpointCallback(out_dir, save_freq=args.checkpoint_freq)
    except Exception as e:
        logger.warning(
            "Could not import CheckpointCallback from train_curriculum.py (%s); "
            "continuing without periodic checkpoints -- only the final model will "
            "be saved.", e
        )
        checkpoint_cb = None

    eval_cb = MaskableEvalCallback(
        eval_env,
        n_eval_episodes=args.eval_episodes,
        eval_freq=args.eval_freq,
        deterministic=True,
        best_model_save_path=str(out_dir / "best_model"),
        verbose=1,
    )

    regression_watch = RegressionWatchCallback(
        total_timesteps=args.steps, check_freq=args.regression_check_freq,
    )

    callbacks = [c for c in [checkpoint_cb, eval_cb, regression_watch.instance] if c is not None]

    t0 = time.time()
    model.learn(total_timesteps=args.steps, callback=callbacks, progress_bar=False)
    elapsed = time.time() - t0
    logger.info("Training finished in %.1f minutes.", elapsed / 60.0)

    final_path = out_dir / "final_model.zip"
    model.save(str(final_path))
    logger.info("Saved final model: %s", final_path)

    summary = {
        "args": vars(args),
        "budget_note": budget_note,
        "full_tour_ceiling": args.n_zones + 1,
        "must_skip_zones": args.max_steps < args.n_zones + 1,
        "elapsed_minutes": elapsed / 60.0,
        "best_mean_reward": eval_cb.best_mean_reward,
        "regression_check_history": regression_watch.instance.history,
    }
    summary_path = out_dir / "run_summary.json"
    with open(summary_path, "w") as f:
        json.dump(summary, f, indent=2, default=str)
    logger.info("Saved run summary: %s", summary_path)
    logger.info("Best mean eval reward this run: %s", eval_cb.best_mean_reward)


def main() -> None:
    args = _parse_args()
    _add_dataset_to_path(args.dataset_dir)
    run_training(args)


if __name__ == "__main__":
    main()