"""
Atlas Phase 4 Stage 3 trainer: joint value + policy heads.

Reads (state_tensor, search_policy, legal_mask, outcome) rows from a Parquet
trajectory file produced by the C# `generate` CLI subcommand and trains an
MLP with two heads:

  - **value head** (4-dim softmax): per-seat win probability. Trained against
    the one-hot outcome via cross-entropy on *all* rows.
  - **policy head** (681-dim softmax over the canonical action space): trained
    against the MCTS visit distribution via cross-entropy, **only on rows
    where the search_policy is non-zero** (i.e., MCTS was actually run; rows
    with single-legal-action shortcuts have an all-zero policy and are
    excluded from the policy loss).

Legal-action masking: the policy head's pre-softmax logits get -inf at
illegal indices before softmax, so the trained model only ever places mass
on legal actions.

Exports an ONNX model with two outputs:
  - "value"  : float[batch, 4]
  - "policy" : float[batch, 681]   (already softmaxed and legal-masked is NOT
                                    baked in — the C# side applies the mask
                                    at inference time using the legal moves)

Usage:
  python train.py --data data/2p.parquet --gametype 2p --out models/2p.onnx
"""

import argparse
import time
from pathlib import Path

import numpy as np
import onnx
import pyarrow.parquet as pq
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset

ACTION_SPACE_SIZE = 681  # must match ActionEncoder.ActionSpaceSize

# Schema versions (must match C# StateTensorVersion enum).
SCHEMA_V1_CATAN = 80
SCHEMA_V1_LANDOVER = 128
SCHEMA_V2_CATAN = 516
SCHEMA_V2_LANDOVER = 564
KNOWN_TENSOR_SIZES = {SCHEMA_V1_CATAN, SCHEMA_V1_LANDOVER, SCHEMA_V2_CATAN, SCHEMA_V2_LANDOVER}


def load_trajectories(path: Path):
    """Read Parquet; auto-detect tensor size; return numpy arrays + metadata.

    Returns (state, policy, mask, outcome, has_policy, tensor_size, schema_version).
    """
    table = pq.read_table(path)
    state = np.array(table.column("state_tensor").to_pylist(), dtype=np.float32)
    policy = np.array(table.column("search_policy").to_pylist(), dtype=np.float32)
    mask = np.array(table.column("legal_mask").to_pylist(), dtype=np.float32)
    outcome = np.array(table.column("outcome").to_pylist(), dtype=np.float32)

    tensor_size = state.shape[1]
    if tensor_size not in KNOWN_TENSOR_SIZES:
        raise ValueError(
            f"Unknown tensor size {tensor_size}. "
            f"Recognized: v1 Catan={SCHEMA_V1_CATAN}, v1 Landover={SCHEMA_V1_LANDOVER}, "
            f"v2 Catan={SCHEMA_V2_CATAN}, v2 Landover={SCHEMA_V2_LANDOVER}."
        )
    # Schema version is stamped into element [0] by the C# encoder.
    schema_version = int(state[0, 0])

    if policy.shape[1] != ACTION_SPACE_SIZE:
        raise ValueError(
            f"Policy size mismatch: expected {ACTION_SPACE_SIZE}, got {policy.shape[1]}."
        )

    print(f"  loaded {state.shape[0]} rows  "
          f"tensor_size={tensor_size}  schema=v{schema_version}")

    policy_sums = policy.sum(axis=1)
    has_policy = policy_sums > 1e-6
    print(f"  rows with MCTS policy: {has_policy.sum()} ({has_policy.mean():.1%})")
    return state, policy, mask, outcome, has_policy, tensor_size, schema_version


class AtlasNet(nn.Module):
    """Two-head MLP: value (per-seat win prob) + policy (action distribution).

    Architectures:
      - "small": 2 hidden layers (Stage 4 baseline). ~310K params on v1.
      - "deep":  3 hidden layers + dropout (recommended for v2 rich features).
                 ~1.2M params at hidden=512 on v2.
    """

    def __init__(self, input_size: int, hidden_size: int = 128,
                 value_size: int = 4, action_size: int = ACTION_SPACE_SIZE,
                 arch: str = "small", dropout: float = 0.0):
        super().__init__()
        self.arch = arch
        self.dropout = nn.Dropout(dropout) if dropout > 0 else None

        if arch == "small":
            self.layers = nn.ModuleList([
                nn.Linear(input_size, hidden_size),
                nn.Linear(hidden_size, hidden_size),
            ])
        elif arch == "deep":
            self.layers = nn.ModuleList([
                nn.Linear(input_size, hidden_size),
                nn.Linear(hidden_size, hidden_size),
                nn.Linear(hidden_size, hidden_size),
            ])
        else:
            raise ValueError(f"Unknown arch '{arch}' (try 'small' or 'deep')")

        self.value_head = nn.Linear(hidden_size, value_size)
        self.policy_head = nn.Linear(hidden_size, action_size)

    def forward(self, x):
        h = x
        for layer in self.layers:
            h = F.relu(layer(h))
            if self.dropout is not None:
                h = self.dropout(h)
        value = F.softmax(self.value_head(h), dim=-1)
        # Policy: return softmaxed over all action indices. The C# side applies
        # the legal-action mask before sampling at inference time.
        policy = F.softmax(self.policy_head(h), dim=-1)
        return value, policy


def joint_loss(pred_value, pred_policy, target_outcome, target_policy, has_policy_mask, policy_weight=1.0):
    """Cross-entropy on value (all rows) + cross-entropy on policy (MCTS rows only)."""
    # Value loss: standard CE against one-hot outcome.
    value_loss = -(target_outcome * torch.log(pred_value.clamp(min=1e-8))).sum(dim=1).mean()

    # Policy loss: CE against MCTS visits, masked to rows where MCTS ran.
    # has_policy_mask shape [B], target_policy shape [B, A], pred_policy shape [B, A].
    policy_ce_per_row = -(target_policy * torch.log(pred_policy.clamp(min=1e-8))).sum(dim=1)
    if has_policy_mask.sum() > 0:
        policy_loss = (policy_ce_per_row * has_policy_mask).sum() / has_policy_mask.sum()
    else:
        policy_loss = torch.tensor(0.0, device=pred_value.device)

    return value_loss + policy_weight * policy_loss, value_loss, policy_loss


def train(model, state, policy, outcome, has_policy, *,
          epochs=30, batch_size=256, lr=1e-3, val_frac=0.1, policy_weight=1.0, device="cpu"):
    n = state.shape[0]
    perm = np.random.RandomState(42).permutation(n)
    val_n = max(1, int(n * val_frac))
    train_idx, val_idx = perm[val_n:], perm[:val_n]

    def to_tensors(idx):
        return (
            torch.from_numpy(state[idx]).to(device),
            torch.from_numpy(policy[idx]).to(device),
            torch.from_numpy(outcome[idx]).to(device),
            torch.from_numpy(has_policy[idx].astype(np.float32)).to(device),
        )

    s_train, p_train, o_train, h_train = to_tensors(train_idx)
    s_val, p_val, o_val, h_val = to_tensors(val_idx)

    loader = DataLoader(TensorDataset(s_train, p_train, o_train, h_train),
                        batch_size=batch_size, shuffle=True)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    for epoch in range(1, epochs + 1):
        model.train()
        epoch_v_loss, epoch_p_loss, n_seen = 0.0, 0.0, 0
        for sb, pb, ob, hb in loader:
            optimizer.zero_grad()
            pv, pp = model(sb)
            loss, v_loss, p_loss = joint_loss(pv, pp, ob, pb, hb, policy_weight)
            loss.backward()
            optimizer.step()
            epoch_v_loss += v_loss.item() * sb.shape[0]
            epoch_p_loss += p_loss.item() * sb.shape[0]
            n_seen += sb.shape[0]

        with torch.no_grad():
            model.eval()
            pv, pp = model(s_val)
            _, v_loss, p_loss = joint_loss(pv, pp, o_val, p_val, h_val, policy_weight)
            value_acc = (pv.argmax(dim=1) == o_val.argmax(dim=1)).float().mean().item()
            if h_val.sum() > 0:
                policy_acc = ((pp[h_val.bool()].argmax(dim=1) ==
                               p_val[h_val.bool()].argmax(dim=1)).float().mean().item())
            else:
                policy_acc = 0.0
        print(
            f"  epoch {epoch:3d}  "
            f"train(v={epoch_v_loss / n_seen:.4f}, p={epoch_p_loss / n_seen:.4f})  "
            f"val(v={v_loss.item():.4f}, p={p_loss.item():.4f}, "
            f"v_acc={value_acc:.3f}, p_acc={policy_acc:.3f})"
        )


def export_onnx(model, input_size: int, out_path: Path):
    """Export with two outputs ('value', 'policy'). The C# OnnxPolicyEvaluator
    expects exactly these names."""
    model.eval()
    dummy = torch.zeros(1, input_size, dtype=torch.float32)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    torch.onnx.export(
        model,
        dummy,
        str(out_path),
        input_names=["state"],
        output_names=["value", "policy"],
        dynamic_axes={
            "state": {0: "batch"},
            "value": {0: "batch"},
            "policy": {0: "batch"},
        },
        opset_version=17,
    )
    # PyTorch's exporter splits weights into a sidecar .onnx.data file when
    # they exceed an internal threshold. That makes distributing the model
    # awkward (workers would need to download both). Re-save with all weights
    # inlined into the single .onnx so it's a standalone artefact.
    inlined = onnx.load(str(out_path), load_external_data=True)
    sidecar = out_path.with_suffix(out_path.suffix + ".data")
    onnx.save(inlined, str(out_path))
    if sidecar.exists():
        sidecar.unlink()
    onnx.checker.check_model(str(out_path))
    size_kb = out_path.stat().st_size / 1024
    print(f"  exported {out_path}  ({size_kb:.1f} KB, weights inlined)")


def pick_device():
    """Auto-detect best available accelerator: CUDA > MPS (Apple Silicon) > CPU."""
    if torch.cuda.is_available():
        return torch.device("cuda")
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def upload_model(onnx_path: Path, api_url: str, admin_key: str, version: str):
    """POST the trained ONNX to /api/model/upload. Picks up after the next worker heartbeat."""
    import urllib.request
    with open(onnx_path, "rb") as f:
        body = f.read()
    url = f"{api_url.rstrip('/')}/api/model/upload?version={version}"
    req = urllib.request.Request(url, data=body, method="POST", headers={
        "Authorization": f"Bearer {admin_key}",
        "Content-Type": "application/octet-stream",
    })
    with urllib.request.urlopen(req, timeout=60) as resp:
        result = resp.read().decode()
    print(f"  uploaded → {result}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", type=Path, required=True)
    parser.add_argument("--gametype", choices=["2p", "4p", "landover"], required=True,
                        help="Kept for compatibility with loop.py; tensor size is auto-detected.")
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--epochs", type=int, default=20)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--hidden", type=int, default=128)
    parser.add_argument("--policy-weight", type=float, default=1.0)
    parser.add_argument("--arch", choices=["small", "deep"], default="small",
                        help="Network architecture. 'small' = 2 layers (Stage 4). "
                             "'deep' = 3 layers + dropout (recommended for v2 schema).")
    parser.add_argument("--dropout", type=float, default=0.0,
                        help="Dropout rate (0 = none).")
    parser.add_argument("--device", default="auto",
                        help="auto | cuda | mps | cpu. Default 'auto' picks the best available.")
    parser.add_argument("--upload-to-api",
                        help="If set, POST the trained ONNX to this API URL (e.g. https://atlas-engine.landover.io). "
                             "Requires --upload-version and ATLAS_ADMIN_KEY env var.")
    parser.add_argument("--upload-version",
                        help="Version string for the uploaded model (e.g. 'v3-2026-05-26T03-00Z'). "
                             "If omitted with --upload-to-api, derived from current UTC timestamp.")
    args = parser.parse_args()

    device = pick_device() if args.device == "auto" else torch.device(args.device)
    print(f"Atlas trainer  ({args.gametype})  device={device}")
    print(f"  data: {args.data}")
    print(f"  out:  {args.out}")
    print()

    print("[load]")
    t0 = time.time()
    state, policy, mask, outcome, has_policy, tensor_size, schema = load_trajectories(args.data)
    print(f"  loaded in {time.time() - t0:.2f}s\n")

    print("[train]")
    model = AtlasNet(tensor_size, hidden_size=args.hidden,
                     arch=args.arch, dropout=args.dropout).to(device)
    nparams = sum(p.numel() for p in model.parameters())
    print(f"  model: AtlasNet(input={tensor_size}, hidden={args.hidden}, "
          f"arch={args.arch}, dropout={args.dropout}, "
          f"value=4, policy={ACTION_SPACE_SIZE})  params={nparams}")
    t0 = time.time()
    train(model, state, policy, outcome, has_policy,
          epochs=args.epochs, batch_size=args.batch_size, lr=args.lr,
          policy_weight=args.policy_weight, device=str(device))
    print(f"  trained in {time.time() - t0:.2f}s\n")

    # Export needs CPU because torch.onnx.export's dummy input is on CPU.
    model.to("cpu")
    print("[export]")
    export_onnx(model, tensor_size, args.out)

    if args.upload_to_api:
        import os
        admin_key = os.environ.get("ATLAS_ADMIN_KEY")
        if not admin_key:
            raise SystemExit("--upload-to-api requires ATLAS_ADMIN_KEY env var")
        import datetime
        version = args.upload_version or f"v-{datetime.datetime.utcnow().strftime('%Y%m%dT%H%M%SZ')}"
        print(f"\n[upload] version={version}")
        upload_model(args.out, args.upload_to_api, admin_key, version)


if __name__ == "__main__":
    main()
