"""Reproduce Delphi2M with absolute-time next-token supervision."""
from __future__ import annotations

import argparse
import json
import logging
import math
import time
from pathlib import Path
from typing import Any, Dict

import numpy as np
import torch
from torch.nn.utils import clip_grad_norm_
from torch.optim import AdamW
from torch.utils.data import DataLoader, RandomSampler
from tqdm.auto import tqdm

from dataset import HealthDataset, collate_fn
from losses import build_loss
from model_architectures import (
    DEFAULT_MODEL_ARCHITECTURE,
    SUPPORTED_MODEL_ARCHITECTURES,
)
from models import DeepHealth, DeepHealthOutput
from targets import PAD_IDX, RESERVED_IDX
from train_util import (
    ContinuousRobustScalerStats,
    configure_torch_for_training,
    create_unique_run_dir,
    fit_continuous_robust_scaler,
    format_extra_info_types,
    get_lr,
    get_model_parameter_counts,
    load_extra_info_types_file,
    move_batch_to_device,
    resolve_device,
    save_checkpoint,
    save_config,
    set_optimizer_lr,
    set_seed,
    setup_logging,
    split_dataset,
    split_dataset_by_eid_files,
)


MODEL_INPUT_KEYS = (
    "event_seq",
    "time_seq",
    "sex",
    "padding_mask",
    "other_type",
    "other_value",
    "other_value_kind",
    "other_time",
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Train DeepHealth with next-token/point supervision")

    parser.add_argument("--data_prefix", type=str, default="ukb")
    parser.add_argument("--labels_file", type=str, default="labels.csv")
    parser.add_argument("--runs_root", type=str, default="runs")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--extra_info_types_file", type=str, default=None)
    parser.add_argument("--no_event_interval_years", type=float, default=5.0)

    parser.add_argument("--train_ratio", type=float, default=0.7)
    parser.add_argument("--val_ratio", type=float, default=0.15)
    parser.add_argument("--test_ratio", type=float, default=0.15)
    parser.add_argument("--train_eid_file", type=str, default="ukb_train_eid.csv")
    parser.add_argument("--val_eid_file", type=str, default="ukb_val_eid.csv")
    parser.add_argument("--test_eid_file", type=str, default="ukb_test_eid.csv")

    parser.add_argument("--n_embd", type=int, default=120)
    parser.add_argument("--n_head", type=int, default=10)
    parser.add_argument("--n_layer", type=int, default=12)
    parser.add_argument("--n_bins", type=int, default=16)
    parser.add_argument("--extra_pool_reduce", type=str, default="mean",
                        choices=["mean", "sum"])
    parser.add_argument("--dropout", type=float, default=0.0)
    parser.add_argument(
        "--model_architecture",
        type=str,
        default=DEFAULT_MODEL_ARCHITECTURE,
        choices=SUPPORTED_MODEL_ARCHITECTURES,
    )

    parser.add_argument("--t_min", type=float, default=0.0027378507871321013)
    parser.add_argument("--max_exp_input", type=float, default=60.0)
    parser.add_argument("--ce_weight", type=float, default=1.0)
    parser.add_argument("--time_weight", type=float, default=1.0)

    parser.add_argument("--batch_size", type=int, default=128)
    parser.add_argument("--base_lr", type=float, default=3e-4)
    parser.add_argument("--weight_decay", type=float, default=0.1)
    parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99))
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--max_epochs", type=int, default=200)
    parser.add_argument("--warmup_epochs", type=int, default=10)
    parser.add_argument("--patience", type=int, default=15)
    parser.add_argument("--min_lr_ratio", type=float, default=0.1)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--device", type=str, default="cuda")
    parser.add_argument("--progress_interval", type=int, default=20)

    args = parser.parse_args()
    use_eid_split = all(
        getattr(args, name)
        for name in ("train_eid_file", "val_eid_file", "test_eid_file")
    )
    if not use_eid_split and not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
        raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
    args.extra_info_types = (
        load_extra_info_types_file(args.extra_info_types_file)
        if args.extra_info_types_file is not None
        else None
    )
    return args


def build_model(
    args: argparse.Namespace,
    dataset: HealthDataset,
    scaler_stats: ContinuousRobustScalerStats,
) -> DeepHealth:
    if tuple(int(x) for x in dataset.cont_type_ids) != scaler_stats.cont_type_ids:
        raise ValueError(
            "RobustScale statistics are not aligned with dataset.cont_type_ids"
        )
    center = scaler_stats.center if dataset.n_cont_types > 0 else None
    scale = scaler_stats.scale if dataset.n_cont_types > 0 else None
    return DeepHealth(
        vocab_size=dataset.vocab_size,
        n_embd=args.n_embd,
        n_head=args.n_head,
        n_layer=args.n_layer,
        n_types=dataset.n_types,
        n_cont_types=dataset.n_cont_types,
        n_categories=dataset.n_categories,
        cont_type_ids=dataset.cont_type_ids,
        n_bins=args.n_bins,
        continuous_value_center=center,
        continuous_value_scale=scale,
        extra_pool_reduce=args.extra_pool_reduce,
        target_mode="next_token",
        time_mode="absolute",
        dist_mode="exponential",
        dropout=args.dropout,
        model_architecture=args.model_architecture,
    )


def build_next_step_loss(args: argparse.Namespace):
    return build_loss(
        "delphi2m",
        ignored_tokens={PAD_IDX, RESERVED_IDX},
        t_min=args.t_min,
        max_exp_input=args.max_exp_input,
        ce_weight=args.ce_weight,
        time_weight=args.time_weight,
    )


def build_augmented_next_step_targets(
    batch_cpu: Dict[str, torch.Tensor],
    model_out: DeepHealthOutput,
) -> Dict[str, torch.Tensor]:
    hidden_len = model_out.hidden.size(1)
    event_len = int(model_out.event_len)
    extra_len = hidden_len - event_len
    device = model_out.hidden.device
    non_blocking = device.type == "cuda"
    if extra_len <= 0:
        return {
            "target_event_seq": batch_cpu["target_event_seq"].to(device, non_blocking=non_blocking),
            "target_time_seq": batch_cpu["target_time_seq"].to(device, non_blocking=non_blocking),
        }

    bsz = batch_cpu["target_event_seq"].size(0)
    other_valid = batch_cpu["other_type"] > 0
    extra_time = batch_cpu["other_time"].new_zeros(bsz, extra_len)
    extra_mask = torch.zeros(bsz, extra_len, dtype=torch.bool)
    for b in range(bsz):
        unique_time = torch.unique(batch_cpu["other_time"][b, other_valid[b]], sorted=True)
        n_time = min(int(unique_time.numel()), extra_len)
        if n_time > 0:
            extra_time[b, :n_time] = unique_time[:n_time]
            extra_mask[b, :n_time] = True

    target_event_seq = torch.cat(
        [
            batch_cpu["target_event_seq"],
            torch.full(
                (bsz, extra_len),
                PAD_IDX,
                dtype=batch_cpu["target_event_seq"].dtype,
            ),
        ],
        dim=1,
    )
    target_time_seq = torch.cat(
        [
            batch_cpu["target_time_seq"],
            torch.zeros(
                bsz,
                extra_len,
                dtype=batch_cpu["target_time_seq"].dtype,
            ),
        ],
        dim=1,
    )
    for b in range(bsz):
        valid_event = batch_cpu["padding_mask"][b].bool()
        if not valid_event.any():
            continue
        n_event = int(valid_event.sum().item())
        events = torch.cat(
            [
                batch_cpu["event_seq"][b, :n_event],
                batch_cpu["target_event_seq"][b, n_event - 1:n_event],
            ]
        )
        times = torch.cat(
            [
                batch_cpu["time_seq"][b, :n_event],
                batch_cpu["target_time_seq"][b, n_event - 1:n_event],
            ]
        )
        valid_full = events > PAD_IDX
        events = events[valid_full]
        times = times[valid_full]
        if events.numel() == 0:
            continue

        for j in range(extra_len):
            if not bool(extra_mask[b, j]):
                continue
            pos = event_len + j
            t = extra_time[b, j]
            future = times > t
            if not future.any():
                continue

            first_idx = int(torch.nonzero(future, as_tuple=False)[0].item())
            next_time = times[first_idx]
            next_event = events[first_idx]
            target_event_seq[b, pos] = next_event
            target_time_seq[b, pos] = next_time

    return {
        "target_event_seq": target_event_seq.to(device, non_blocking=non_blocking),
        "target_time_seq": target_time_seq.to(device, non_blocking=non_blocking),
    }


def compute_next_step_loss(
    model: DeepHealth,
    criterion,
    batch: Dict[str, torch.Tensor],
    device: torch.device,
) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
    batch_cpu = batch
    batch = move_batch_to_device(
        {key: batch_cpu[key] for key in MODEL_INPUT_KEYS},
        device,
    )
    model_out = model(
        event_seq=batch["event_seq"],
        time_seq=batch["time_seq"],
        sex=batch["sex"],
        padding_mask=batch["padding_mask"],
        other_type=batch["other_type"],
        other_value=batch["other_value"],
        other_value_kind=batch["other_value_kind"],
        other_time=batch["other_time"],
        return_output=True,
    )
    if not isinstance(model_out, DeepHealthOutput):
        raise TypeError("DeepHealth return_output=True must return DeepHealthOutput")
    targets = build_augmented_next_step_targets(
        batch_cpu=batch_cpu,
        model_out=model_out,
    )
    logits = model.calc_risk(model_out.hidden)

    loss, parts = criterion(
        logits=logits,
        target_events=targets["target_event_seq"],
        target_times=targets["target_time_seq"],
        current_times=model_out.time_seq,
        padding_mask=model_out.padding_mask,
        return_components=True,
    )
    if not torch.isfinite(loss):
        raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}")
    return loss, parts


def run_epoch(
    logger: logging.Logger,
    args: argparse.Namespace,
    model: DeepHealth,
    criterion,
    loader: DataLoader,
    optimizer: AdamW | None,
    device: torch.device,
    is_train: bool,
) -> float:
    model.train(is_train)
    total = torch.zeros((), device=device)
    n_batches = 0
    skipped = 0
    parts_sum: Dict[str, torch.Tensor] = {}
    desc = "train" if is_train else "val"
    progress_interval = max(1, int(args.progress_interval))

    progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)
    for batch_idx, batch in enumerate(progress):
        try:
            loss, parts = compute_next_step_loss(
                model, criterion, batch, device
            )
            if is_train:
                if optimizer is None:
                    raise ValueError("optimizer is required for training")
                optimizer.zero_grad(set_to_none=True)
                loss.backward()
                if args.grad_clip > 0:
                    clip_grad_norm_(model.parameters(), args.grad_clip)
                optimizer.step()

            total = total + loss.detach()
            n_batches += 1
            for name, value in parts.items():
                parts_sum[name] = parts_sum.get(name, torch.zeros((), device=device)) + value.detach()
            if (batch_idx + 1) % progress_interval == 0:
                avg = total / max(1, n_batches)
                postfix = {
                    "loss": f"{float(loss.detach().cpu()):.4f}",
                    "avg": f"{float(avg.detach().cpu()):.4f}",
                    "skipped": skipped,
                }
                for name, value in parts_sum.items():
                    postfix[name] = f"{float((value / max(1, n_batches)).detach().cpu()):.4f}"
                progress.set_postfix(postfix)
        except RuntimeError as exc:
            if "Loss is not finite" not in str(exc):
                raise
            skipped += 1
            logger.warning(f"Batch {batch_idx} skipped: {str(exc)[:120]}")

    if skipped:
        logger.info(f"Skipped {skipped} batches due to non-finite loss")
    return float((total / max(1, n_batches)).detach().cpu()) if n_batches else float("inf")


def build_metadata(
    args: argparse.Namespace,
    dataset: HealthDataset,
    run_name: str,
    train_subset,
    val_subset,
    test_subset,
    scaler_stats: ContinuousRobustScalerStats,
) -> Dict[str, Any]:
    return {
        "run_name": run_name,
        "dataset_class": "NextStepHealthDataset",
        "collate_fn": "next_step_collate_fn",
        "model_class": "DeepHealth",
        "model_architecture": args.model_architecture,
        "model_target_mode": "next_token",
        "target_mode": "delphi2m",
        "event_stream_version": "disease_death_only_v1",
        "uses_assessment_event_token": False,
        "time_mode": "absolute",
        "dist_mode": "exponential",
        "extra_info_types_file": (
            Path(args.extra_info_types_file).name
            if args.extra_info_types_file is not None
            else None
        ),
        "extra_info_types": [int(x) for x in dataset.extra_info_types],
        "continuous_value_scaling": "robust",
        "continuous_value_scaler": scaler_stats.as_metadata(),
        "dataset_metadata": {
            "vocab_size": int(dataset.vocab_size),
            "n_types": int(dataset.n_types),
            "n_cont_types": int(dataset.n_cont_types),
            "n_categories": int(dataset.n_categories),
            "cont_type_ids": [int(x) for x in dataset.cont_type_ids],
            "extra_info_types": [int(x) for x in dataset.extra_info_types],
            "event_stream_version": "disease_death_only_v1",
            "uses_assessment_event_token": False,
        },
        "split_sizes": {
            "train": int(len(train_subset)),
            "val": int(len(val_subset)),
            "test": int(len(test_subset)),
        },
        "resolved_readout_name": "token",
        "resolved_loss_name": "delphi2m",
    }


def main() -> None:
    args = parse_args()
    set_seed(args.seed)
    device = resolve_device(args.device)
    configure_torch_for_training(device)

    run_dir, run_name = create_unique_run_dir(
        lambda timestamp: (
            "absolute_exponential_next_token_delphi2m_"
            f"gap_{args.no_event_interval_years:g}y_{timestamp}"
        ),
        runs_root=Path(args.runs_root) / args.model_architecture,
    )
    logger = setup_logging(run_dir)

    logger.info(f"Starting next-step training run: {run_name}")
    logger.info(f"Device: {device}")
    logger.info(f"Model architecture: {args.model_architecture}")
    logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
    logger.info("Continuous value scaling: RobustScale (required)")
    logger.info("time_mode=absolute, readout=token, target_mode=delphi2m")

    dataset = HealthDataset(
        data_prefix=args.data_prefix,
        labels_file=args.labels_file,
        no_event_interval_years=args.no_event_interval_years,
        extra_info_types=args.extra_info_types,
    )
    if args.train_eid_file and args.val_eid_file and args.test_eid_file:
        train_subset, val_subset, test_subset = split_dataset_by_eid_files(
            dataset=dataset,
            train_eid_file=args.train_eid_file,
            val_eid_file=args.val_eid_file,
            test_eid_file=args.test_eid_file,
        )
        logger.info(
            "Using eid split files: "
            f"train={args.train_eid_file}, val={args.val_eid_file}, test={args.test_eid_file}"
        )
    else:
        train_subset, val_subset, test_subset = split_dataset(
            dataset=dataset,
            train_ratio=args.train_ratio,
            val_ratio=args.val_ratio,
            test_ratio=args.test_ratio,
            seed=args.seed,
        )
        logger.info(
            f"Using random ratio split: train={args.train_ratio}, "
            f"val={args.val_ratio}, test={args.test_ratio}, seed={args.seed}"
        )
    logger.info(
        f"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
    )

    if dataset.n_cont_types > 0:
        logger.info(
            "Fitting continuous RobustScaler on the complete training subset: "
            f"patients={len(train_subset):,}, features={dataset.n_cont_types}"
        )
    scaler_stats = fit_continuous_robust_scaler(dataset, train_subset)
    if dataset.n_cont_types > 0:
        logger.info(
            "Continuous RobustScaler fitted: "
            f"observations={int(scaler_stats.observation_count.sum()):,}, "
            f"min_per_feature={int(scaler_stats.observation_count.min()):,}, "
            f"max_per_feature={int(scaler_stats.observation_count.max()):,}"
        )

    train_loader = DataLoader(
        train_subset,
        batch_size=args.batch_size,
        sampler=RandomSampler(train_subset, generator=torch.Generator().manual_seed(args.seed)),
        collate_fn=collate_fn,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
        persistent_workers=args.num_workers > 0,
        prefetch_factor=2 if args.num_workers > 0 else None,
    )
    val_loader = DataLoader(
        val_subset,
        batch_size=args.batch_size,
        shuffle=False,
        collate_fn=collate_fn,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
        persistent_workers=args.num_workers > 0,
        prefetch_factor=2 if args.num_workers > 0 else None,
    )
    test_loader = DataLoader(
        test_subset,
        batch_size=args.batch_size,
        shuffle=False,
        collate_fn=collate_fn,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
        persistent_workers=args.num_workers > 0,
        prefetch_factor=2 if args.num_workers > 0 else None,
    )

    model = build_model(args, dataset, scaler_stats).to(device)
    parameter_counts = get_model_parameter_counts(model)
    logger.info(
        "Model parameters: "
        f"total={parameter_counts['model_parameter_count']:,}, "
        f"trainable={parameter_counts['trainable_parameter_count']:,}"
    )
    criterion = build_next_step_loss(args)
    optimizer = AdamW(
        model.parameters(),
        lr=args.base_lr,
        betas=tuple(args.betas),
        weight_decay=args.weight_decay,
    )
    adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)

    train_metadata = build_metadata(
        args,
        dataset,
        run_name,
        train_subset,
        val_subset,
        test_subset,
        scaler_stats,
    )
    train_metadata.update(parameter_counts)
    save_config(
        args,
        run_dir / "train_config.json",
        extra=train_metadata,
    )

    best_val = float("inf")
    patience = 0
    history = []
    best_model_path = run_dir / "best_model.pt"
    start = time.time()

    for epoch in range(args.max_epochs):
        lr = get_lr(epoch, args, adaptive_lr)
        set_optimizer_lr(optimizer, lr)

        train_loss = run_epoch(
            logger, args, model, criterion, train_loader, optimizer, device, True
        )
        with torch.no_grad():
            val_loss = run_epoch(
                logger, args, model, criterion, val_loader, None, device, False
            )

        is_best = val_loss < best_val
        if is_best:
            best_val = val_loss
            patience = 0
            save_checkpoint(model, best_model_path)
        else:
            patience += 1

        logger.info(
            f"Epoch {epoch + 1}/{args.max_epochs} | lr={lr:.6f} | "
            f"train_loss={train_loss:.6f} | val_loss={val_loss:.6f} | "
            f"best_val_loss={best_val:.6f} | patience={patience}/{args.patience} | "
            f"elapsed={time.time() - start:.1f}s"
        )
        history.append({
            "epoch": epoch + 1,
            "lr": lr,
            "train_loss": train_loss,
            "val_loss": val_loss,
            "best_val_loss": best_val,
            "is_best": int(is_best),
        })
        if patience >= args.patience:
            logger.info(f"Early stopping triggered at epoch {epoch + 1}")
            break

    with (run_dir / "history.json").open("w", encoding="utf-8") as f:
        json.dump(history, f, indent=2)

    logger.info("Evaluating best model on next-step test split...")
    model.load_state_dict(torch.load(best_model_path, map_location=device))
    with torch.no_grad():
        test_loss = run_epoch(
            logger, args, model, criterion, test_loader, None, device, False
        )
    logger.info(f"Test loss: {test_loss:.6f}")
    logger.info(f"Best checkpoint: {best_model_path}")


if __name__ == "__main__":
    main()
