package_code_imports/v5_scripts/train_v5_cloud.py

from __future__ import annotations

import argparse
import json
import sys
import time
from dataclasses import asdict
from pathlib import Path
from types import SimpleNamespace
from typing import Any

import numpy as np
import pandas as pd
import torch


SCRIPT_DIR = Path(__file__).resolve().parent
ROOT_DIR = SCRIPT_DIR.parents[1]
V4P4_SCRIPT_DIR = ROOT_DIR / "v4p4_world_model" / "scripts"
V3P5_SCRIPT_DIR = ROOT_DIR / "v3p5_static" / "scripts"
for path in (SCRIPT_DIR, V4P4_SCRIPT_DIR, V3P5_SCRIPT_DIR):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

from action_ontology_v5 import build_action_ontology, load_json  # noqa: E402
from build_v5_action_tensors import (  # noqa: E402
    build_action_arrays,
    build_target_trial_labels,
    medication_flags_from_stage0,
    resolve_split_path,
    validate_medication_flags,
)
from config_v5_train import DEFAULT_STAGE0_DIR, DEFAULT_TENSOR_DIR, V5_LOSS_WEIGHTS  # noqa: E402
from loss_v5 import sctm_v5_loss  # noqa: E402
from model_v5 import SCTMv5, SCTMv5Config  # noqa: E402
from smoke_v3p4_architecture import build_field_value_mask, load_service_prior  # noqa: E402
from train_v4p4_cloud import TrainIndexSampler, autocast_context, batch_from_indices, load_npz_to_memory, metadata_config  # noqa: E402


def tensorize_actions(action_arrays: dict[str, np.ndarray], idx: np.ndarray, device: torch.device) -> dict[str, torch.Tensor]:
    return {
        "action_type_ids": torch.as_tensor(action_arrays["action_type_ids"][idx], dtype=torch.long, device=device),
        "action_value_ids": torch.as_tensor(action_arrays["action_value_ids"][idx], dtype=torch.long, device=device),
        "action_mask": torch.as_tensor(action_arrays["action_mask"][idx], dtype=torch.bool, device=device),
        "action_available_at": torch.as_tensor(action_arrays["action_available_at"][idx], dtype=torch.long, device=device),
    }


def validate_action_arrays(action_arrays: dict[str, np.ndarray], arrays: dict[str, np.ndarray], ontology: dict[str, Any], *, source: Path | str) -> None:
    required = ("action_type_ids", "action_value_ids", "action_mask", "action_available_at")
    missing = [key for key in required if key not in action_arrays]
    if missing:
        raise ValueError(f"Action tensor source {source} is missing required arrays: {missing}")
    n_windows, seq_len = arrays["valid_mask"].shape
    n_slots = len(ontology["slots"])
    for key in required:
        shape = tuple(action_arrays[key].shape)
        expected = (n_windows, seq_len, n_slots)
        if shape != expected:
            raise ValueError(f"Action tensor source {source} has {key} shape {shape}, expected {expected}")
    if not np.array_equal(action_arrays["action_mask"].astype(bool), action_arrays["action_mask"]):
        raise ValueError(f"Action tensor source {source} has non-boolean-compatible action_mask values")


def metrics_to_float(metrics: dict[str, torch.Tensor]) -> dict[str, float]:
    return {key: float(value.detach().cpu()) for key, value in metrics.items()}


def aggregate_eval_metrics(rows: list[dict[str, float]]) -> dict[str, float]:
    if not rows:
        return {}
    keys = sorted(set().union(*(row.keys() for row in rows)))
    weights = np.asarray([max(0.0, float(row.get("valid_next_positions", 0.0))) for row in rows], dtype=np.float64)
    weight_sum = float(weights.sum())
    out: dict[str, float] = {}
    for key in keys:
        values = np.asarray([float(row[key]) for row in rows if key in row], dtype=np.float64)
        if values.size == 0:
            continue
        if key in {"valid_next_positions", "next_contact_positions"}:
            out[key] = float(values.sum())
            continue
        if key.startswith("loss_") and weight_sum > 0 and values.size == len(rows):
            out[key] = float(np.sum(values * weights) / weight_sum)
        else:
            out[key] = float(values.mean())
    return out


def load_split_bundle(
    tensor_dir: Path,
    stage0_dir: Path,
    split: str,
    max_windows: int = 0,
    action_dir: Path | None = None,
) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray], dict[str, Any], dict[str, Any]]:
    meta = load_json(tensor_dir / "tensor_metadata.json")
    vocab = load_json(tensor_dir / "cat_value_vocab.json")
    ordinal_direction = load_json(stage0_dir / "ordinal_direction_table.json")
    shard = resolve_split_path(tensor_dir, split)
    arrays = load_npz_to_memory(shard)
    if max_windows > 0:
        arrays = {k: v[:max_windows] for k, v in arrays.items()}
    ontology = build_action_ontology(meta, vocab)
    medication_flags = medication_flags_from_stage0(arrays, stage0_dir)
    action_path = action_dir / f"v5_{split}_action_tensors.npz" if action_dir is not None else None
    if action_path is not None and action_path.exists():
        with np.load(action_path, allow_pickle=False) as z:
            action_keys = {"action_type_ids", "action_value_ids", "action_mask", "action_available_at", "lai_present", "oral_present"}
            action_arrays = {k: z[k][:max_windows] if max_windows > 0 else z[k] for k in z.files if k in action_keys}
        validate_action_arrays(action_arrays, arrays, ontology, source=action_path)
        validate_medication_flags(action_arrays, medication_flags, source=action_path)
    else:
        action_arrays = build_action_arrays(arrays, meta, ordinal_direction, vocab, ontology, medication_flags=medication_flags)
        _ = build_target_trial_labels(arrays, (30.0, 60.0, 90.0), medication_flags=medication_flags, stage0_dir=stage0_dir)
        validate_action_arrays(action_arrays, arrays, ontology, source="built_in_memory")
        validate_medication_flags(action_arrays, medication_flags, source="built_in_memory")
    return arrays, action_arrays, meta, ontology


def make_adamw(
    parameters: Any,
    *,
    lr: float,
    weight_decay: float,
    fused: bool,
    device: torch.device,
) -> torch.optim.Optimizer:
    kwargs: dict[str, Any] = {"lr": lr, "weight_decay": weight_decay}
    if fused and device.type == "cuda":
        try:
            return torch.optim.AdamW(parameters, fused=True, **kwargs)
        except TypeError:
            pass
    return torch.optim.AdamW(parameters, **kwargs)


def _extract_model_state_dict(checkpoint: Any) -> tuple[dict[str, torch.Tensor], str]:
    if isinstance(checkpoint, dict):
        for key in ("model", "model_state_dict", "state_dict"):
            state = checkpoint.get(key)
            if isinstance(state, dict):
                return state, key
    if isinstance(checkpoint, dict) and all(isinstance(v, torch.Tensor) for v in checkpoint.values()):
        return checkpoint, "root"
    raise ValueError("Checkpoint does not contain a model state dict under model/model_state_dict/state_dict or root tensor keys.")


def _normalize_state_dict_keys(state: dict[str, torch.Tensor], model_keys: set[str]) -> tuple[dict[str, torch.Tensor], str]:
    if any(key in model_keys for key in state):
        return state, "as_is"
    for prefix in ("module.", "_orig_mod."):
        stripped = {key.removeprefix(prefix): value for key, value in state.items() if key.startswith(prefix)}
        if stripped and any(key in model_keys for key in stripped):
            return stripped, f"strip_{prefix}"
    return state, "as_is_no_direct_match"


def load_checkpoint_weights(model: torch.nn.Module, checkpoint_path: Path, *, strict: bool) -> dict[str, Any]:
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    raw_state, source_key = _extract_model_state_dict(checkpoint)
    model_keys = set(model.state_dict().keys())
    state, key_transform = _normalize_state_dict_keys(raw_state, model_keys)
    loadable_keys = sorted(key for key in state if key in model_keys and tuple(state[key].shape) == tuple(model.state_dict()[key].shape))
    if not strict:
        state = {key: value for key, value in state.items() if key in loadable_keys}
        if not state:
            raise RuntimeError(f"Checkpoint {checkpoint_path} did not match any model parameters; refusing to continue.")
    result = model.load_state_dict(state, strict=strict)
    return {
        "path": str(checkpoint_path),
        "strict": bool(strict),
        "source_key": source_key,
        "key_transform": key_transform,
        "checkpoint_parameter_keys": int(len(raw_state)),
        "loadable_parameter_keys": int(len(loadable_keys)),
        "missing_keys": list(result.missing_keys),
        "unexpected_keys": list(result.unexpected_keys),
    }


def set_base_trainable(model: torch.nn.Module, trainable: bool) -> dict[str, int]:
    action_prefixes = (
        "action_encoder.",
        "action_pwe_delta_head.",
        "action_active_delta_head.",
        "action_missing_delta_head.",
        "action_field_delta_head.",
        "action_numeric_delta_head.",
        "action_ordinal_delta_head.",
        "action_event_delta_head.",
        "behavior_policy_head.",
    )
    frozen = 0
    trainable_count = 0
    for name, param in model.named_parameters():
        if trainable or name.startswith(action_prefixes):
            param.requires_grad_(True)
            trainable_count += param.numel()
        else:
            param.requires_grad_(False)
            frozen += param.numel()
    return {"trainable_parameters": int(trainable_count), "frozen_parameters": int(frozen)}


def checkpoint_state(
    model: torch.nn.Module,
    config: SCTMv5Config,
    ontology: dict[str, Any],
    args: argparse.Namespace,
    prior_audit: dict[str, Any],
    *,
    step: int | None = None,
    best_step: int | None = None,
    best_val_loss: float | None = None,
    current_val_loss: float | None = None,
) -> dict[str, Any]:
    raw_model = getattr(model, "_orig_mod", model)
    payload = {
        "model": raw_model.state_dict(),
        "config": asdict(config),
        "ontology": ontology,
        "args": vars(args),
        "service_prior_audit": prior_audit,
    }
    if step is not None:
        payload["step"] = int(step)
    if best_step is not None:
        payload["best_step"] = int(best_step)
    if best_val_loss is not None:
        payload["best_val_loss"] = float(best_val_loss)
    if current_val_loss is not None:
        payload["current_val_loss"] = float(current_val_loss)
    return payload


@torch.inference_mode()
def evaluate(model: SCTMv5, arrays: dict[str, np.ndarray], action_arrays: dict[str, np.ndarray], config: SCTMv5Config, args: argparse.Namespace, device: torch.device) -> dict[str, float]:
    model.eval()
    n = arrays["valid_mask"].shape[0]
    n_eval = min(n, args.eval_windows) if args.eval_windows > 0 else n
    losses = []
    for start in range(0, n_eval, args.eval_batch_size):
        idx = np.arange(start, min(start + args.eval_batch_size, n_eval))
        batch = batch_from_indices(arrays, idx, device)
        batch.update(tensorize_actions(action_arrays, idx, device))
        with autocast_context(device, args.precision):
            out = model(batch, rollout_steps=1, compute_pwe_diagnostics=False)
            _, metrics = sctm_v5_loss(out, batch, config, return_metrics=True)
        losses.append(metrics_to_float(metrics))
    return aggregate_eval_metrics(losses)


def main() -> None:
    parser = argparse.ArgumentParser(description="Train SCTM-v5 action-conditioned heads on observed-action likelihood.")
    parser.add_argument("--tensor-dir", type=Path, default=DEFAULT_TENSOR_DIR)
    parser.add_argument("--stage0-dir", type=Path, default=DEFAULT_STAGE0_DIR)
    parser.add_argument("--out-dir", type=Path, required=True)
    parser.add_argument("--train-split", default="train")
    parser.add_argument("--val-split", default="val")
    parser.add_argument("--action-dir", type=Path, default=None, help="Optional directory containing precomputed v5_{split}_action_tensors.npz files.")
    parser.add_argument("--max-windows", type=int, default=0)
    parser.add_argument("--max-steps", type=int, default=100)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--eval-batch-size", type=int, default=16)
    parser.add_argument("--log-every", type=int, default=25)
    parser.add_argument("--eval-every", type=int, default=500)
    parser.add_argument("--eval-windows", type=int, default=64)
    parser.add_argument("--save-every", type=int, default=1000)
    parser.add_argument("--save-eval-checkpoints", action="store_true", help="Save a checkpoint at every validation point for post-hoc model selection audits.")
    parser.add_argument("--top-k-checkpoints", type=int, default=0, help="If >0, keep only the top-k eval checkpoints by validation loss in eval_checkpoints/.")
    parser.add_argument("--early-stopping-patience", type=int, default=0, help="If >0, stop after this many consecutive validation checks without improvement.")
    parser.add_argument("--early-stopping-min-delta", type=float, default=0.0, help="Minimum val loss improvement required to reset early-stopping patience.")
    parser.add_argument("--min-steps-before-stopping", type=int, default=0, help="Do not trigger early stopping before this step.")
    parser.add_argument("--learning-rate", type=float, default=1.0e-4)
    parser.add_argument("--weight-decay", type=float, default=1.0e-4)
    parser.add_argument("--grad-clip", type=float, default=1.0)
    parser.add_argument("--precision", choices=("fp32", "bf16", "fp16"), default="fp32")
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--fused-adamw", action="store_true")
    parser.add_argument("--compile", action="store_true")
    parser.add_argument("--matmul-precision", choices=("highest", "high", "medium"), default=None)
    parser.add_argument("--d-model", type=int, default=64)
    parser.add_argument("--n-heads", type=int, default=4)
    parser.add_argument("--n-layers", type=int, default=2)
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--seed", type=int, default=20260525)
    parser.add_argument("--init-checkpoint", type=Path, default=None, help="Optional v4/v5 checkpoint used to initialize matching model weights with strict=False.")
    parser.add_argument("--resume-checkpoint", type=Path, default=None, help="Optional v5 checkpoint used to resume with strict=True.")
    parser.add_argument("--freeze-base-steps", type=int, default=0, help="If >0, train only v5 action/behavior heads for this many initial steps, then unfreeze all parameters.")
    for key, value in V5_LOSS_WEIGHTS.items():
        parser.add_argument(f"--{key.replace('_', '-')}", type=float, default=float(value))
    args = parser.parse_args()

    args.out_dir.mkdir(parents=True, exist_ok=True)
    device = torch.device(args.device)
    if args.matmul_precision:
        torch.set_float32_matmul_precision(args.matmul_precision)
    rng = np.random.default_rng(args.seed)
    torch.manual_seed(args.seed)
    train_arrays, train_actions, meta, ontology = load_split_bundle(args.tensor_dir, args.stage0_dir, args.train_split, args.max_windows, args.action_dir)
    val_arrays, val_actions, _, _ = load_split_bundle(args.tensor_dir, args.stage0_dir, args.val_split, args.max_windows, args.action_dir)
    vocab = load_json(args.tensor_dir / "cat_value_vocab.json")
    base_cfg = metadata_config(meta, vocab, SimpleNamespace(d_model=args.d_model, n_heads=args.n_heads, n_layers=args.n_layers, dropout=args.dropout, n_missing=5))
    cfg_dict = asdict(base_cfg)
    cfg_dict.update(
        {
            "n_action_slots": len(ontology["slots"]),
            "action_value_vocab_size": len(ontology["action_value_vocab"]),
            "n_action_availability": len(ontology["action_availability_vocab"]),
        }
    )
    config = SCTMv5Config(**cfg_dict)
    field_value_mask, _ = build_field_value_mask(meta, vocab, {k: train_arrays[k] for k in ["cat_value_ids", "missing_ids"]})
    prior, prior_audit = load_service_prior(args.stage0_dir, int(meta.get("n_service_states", 8)), "service_state_transitions_train.json")
    model = SCTMv5(config, field_value_mask=field_value_mask, service_prior_bias=prior).to(device)
    checkpoint_audit: dict[str, Any] | None = None
    if args.resume_checkpoint is not None:
        checkpoint_audit = load_checkpoint_weights(model, args.resume_checkpoint, strict=True)
    elif args.init_checkpoint is not None:
        checkpoint_audit = load_checkpoint_weights(model, args.init_checkpoint, strict=False)
    freeze_audit = set_base_trainable(model, trainable=args.freeze_base_steps <= 0)
    if args.compile:
        model = torch.compile(model)
    optimizer = make_adamw(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay, fused=args.fused_adamw, device=device)
    sampler = TrainIndexSampler(rng, train_arrays["valid_mask"].shape[0], sample_with_replacement=False)
    loss_kwargs = {key: getattr(args, key) for key in V5_LOSS_WEIGHTS}
    log_rows: list[dict[str, Any]] = []
    best_val = float("inf")
    best_step = 0
    completed_steps = 0
    no_improve_evals = 0
    eval_checkpoint_records: list[tuple[float, int, Path]] = []
    start_time = time.time()
    for step in range(1, args.max_steps + 1):
        completed_steps = step
        if args.freeze_base_steps > 0 and step == args.freeze_base_steps + 1:
            raw_model = getattr(model, "_orig_mod", model)
            freeze_audit = set_base_trainable(raw_model, trainable=True)
            optimizer = make_adamw(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay, fused=args.fused_adamw, device=device)
        model.train()
        idx = sampler.next(args.batch_size)
        batch = batch_from_indices(train_arrays, idx, device)
        batch.update(tensorize_actions(train_actions, idx, device))
        optimizer.zero_grad(set_to_none=True)
        need_log = step == 1 or step % args.log_every == 0 or step % args.eval_every == 0 or step == args.max_steps
        need_eval = step == 1 or step % args.eval_every == 0 or step == args.max_steps
        with autocast_context(device, args.precision):
            out = model(batch, rollout_steps=1, compute_pwe_diagnostics=False)
            loss, metrics = sctm_v5_loss(out, batch, config, **loss_kwargs, return_metrics=need_log)
        loss.backward()
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
        optimizer.step()
        row = {
            "step": step,
            "elapsed_sec": round(time.time() - start_time, 3),
            "loss": float(loss.detach().cpu()) if need_log else float("nan"),
            "grad_norm": float(grad_norm.detach().cpu()) if need_log else float("nan"),
        }
        if need_log:
            row.update({f"train_{k}": v for k, v in metrics_to_float(metrics).items()})
        if need_eval:
            val = evaluate(model, val_arrays, val_actions, config, args, device)
            row.update({f"val_{k}": v for k, v in val.items()})
            row.update({f"freeze_{k}": v for k, v in freeze_audit.items()})
            val_loss = float(val.get("loss_total", float("inf")))
            improved = val_loss < (best_val - float(args.early_stopping_min_delta))
            if improved:
                best_val = val_loss
                best_step = step
                no_improve_evals = 0
                torch.save(
                    checkpoint_state(
                        model,
                        config,
                        ontology,
                        args,
                        prior_audit,
                        step=step,
                        best_step=best_step,
                        best_val_loss=best_val,
                        current_val_loss=val_loss,
                    ),
                    args.out_dir / "best.pt",
                )
            else:
                no_improve_evals += 1
            if args.save_eval_checkpoints or args.top_k_checkpoints > 0:
                eval_dir = args.out_dir / "eval_checkpoints"
                eval_dir.mkdir(parents=True, exist_ok=True)
                eval_path = eval_dir / f"step_{step:06d}_val_{val_loss:.6f}.pt"
                torch.save(
                    checkpoint_state(
                        model,
                        config,
                        ontology,
                        args,
                        prior_audit,
                        step=step,
                        best_step=best_step,
                        best_val_loss=best_val,
                        current_val_loss=val_loss,
                    ),
                    eval_path,
                )
                eval_checkpoint_records.append((val_loss, step, eval_path))
                if args.top_k_checkpoints > 0:
                    eval_checkpoint_records = sorted(eval_checkpoint_records, key=lambda x: (x[0], x[1]))
                    for _, _, stale_path in eval_checkpoint_records[args.top_k_checkpoints :]:
                        stale_path.unlink(missing_ok=True)
                    eval_checkpoint_records = eval_checkpoint_records[: args.top_k_checkpoints]
            row["best_step"] = best_step
            row["best_val_loss"] = best_val
            row["early_stopping_no_improve_evals"] = no_improve_evals
            should_stop = (
                args.early_stopping_patience > 0
                and step >= args.min_steps_before_stopping
                and no_improve_evals >= args.early_stopping_patience
            )
            row["early_stopped"] = bool(should_stop)
            print(json.dumps(row, ensure_ascii=False), flush=True)
        log_rows.append(row)
        if need_eval and row.get("early_stopped"):
            break
        if args.save_every > 0 and step % args.save_every == 0:
            torch.save(
                checkpoint_state(
                    model,
                    config,
                    ontology,
                    args,
                    prior_audit,
                    step=step,
                    best_step=best_step,
                    best_val_loss=best_val,
                ),
                args.out_dir / "latest.pt",
            )
    pd.DataFrame(log_rows).to_csv(args.out_dir / "train_log.csv", index=False)
    torch.save(
        checkpoint_state(
            model,
            config,
            ontology,
            args,
            prior_audit,
            step=completed_steps,
            best_step=best_step,
            best_val_loss=best_val,
        ),
        args.out_dir / "latest.pt",
    )
    summary = {
        "status": "complete",
        "out_dir": str(args.out_dir),
        "steps": completed_steps,
        "planned_steps": args.max_steps,
        "train_rows": len(log_rows),
        "best_step": best_step,
        "best_val_loss": best_val,
        "early_stopped": bool(completed_steps < args.max_steps),
        "checkpoint_audit": checkpoint_audit,
        "freeze_audit": freeze_audit,
        "eval_checkpoint_count": len(eval_checkpoint_records),
    }
    (args.out_dir / "train_result.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(summary, ensure_ascii=False), flush=True)


if __name__ == "__main__":
    main()