"""Pretrain a TimesNet + ConvNeXtV2 autoencoder on training-set exposure."""
from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import numpy as np
import torch
from torch.optim import AdamW
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

from backbones import TimesNetExposureAutoencoder
from dataset import ExposureCache
from train_util import (
    configure_torch_for_training,
    create_unique_run_dir,
    load_eid_file,
    resolve_device,
    set_seed,
    setup_logging,
)


class ExposureWindowDataset(Dataset):
    def __init__(self, cache: ExposureCache, row_indices: np.ndarray):
        self.cache = cache
        self.row_indices = np.asarray(row_indices, dtype=np.int64)

    def __len__(self) -> int:
        return len(self.row_indices)

    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
        row = int(self.row_indices[index])
        return {
            "daily": torch.from_numpy(
                np.array(self.cache.daily[row], dtype=np.float32, copy=True)
            ),
            "monthly": torch.from_numpy(
                np.array(self.cache.monthly[row], dtype=np.float32, copy=True)
            ),
        }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Pretrain a TimesNet + ConvNeXtV2 exposure autoencoder"
    )
    parser.add_argument("--exposure_cache_dir", required=True)
    parser.add_argument("--train_eid_file", default="ukb_train_eid.csv")
    parser.add_argument("--val_eid_file", default="ukb_val_eid.csv")
    parser.add_argument("--runs_root", default="runs")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--n_embd", type=int, default=120)
    parser.add_argument("--d_model", type=int, default=None)
    parser.add_argument("--n_layers", type=int, default=2)
    parser.add_argument("--top_k", type=int, default=3)
    parser.add_argument("--n_convnext_blocks", type=int, default=2)
    parser.add_argument("--conv_kernel_size", type=int, default=7)
    parser.add_argument("--mlp_ratio", type=float, default=4.0)
    parser.add_argument("--dropout", type=float, default=0.0)
    parser.add_argument("--mask_ratio", type=float, default=0.25)
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--base_lr", type=float, default=3e-4)
    parser.add_argument("--weight_decay", type=float, default=0.05)
    parser.add_argument("--max_epochs", type=int, default=100)
    parser.add_argument("--warmup_epochs", type=int, default=5)
    parser.add_argument("--patience", type=int, default=12)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument(
        "--data_parallel",
        action="store_true",
        help="Use torch.nn.DataParallel across multiple CUDA devices.",
    )
    parser.add_argument(
        "--gpu_ids",
        default=None,
        help="Comma-separated CUDA device ids for --data_parallel, e.g. 0,1,2,3.",
    )
    args = parser.parse_args()
    if not 0.0 <= args.mask_ratio < 1.0:
        parser.error("--mask_ratio must be in [0, 1)")
    if args.gpu_ids:
        try:
            args.gpu_ids = [
                int(part.strip())
                for part in args.gpu_ids.split(",")
                if part.strip()
            ]
        except ValueError as exc:
            parser.error("--gpu_ids must be a comma-separated list of integers")
        if not args.gpu_ids:
            parser.error("--gpu_ids did not contain any valid CUDA device ids")
        args.data_parallel = True
    return args


def select_rows(cache: ExposureCache, eids: set[int], split: str) -> np.ndarray:
    valid_row = np.asarray(cache.row_index, dtype=np.int64) >= 0
    selected_events = valid_row & np.isin(cache.eids, np.fromiter(eids, np.int64))
    rows = np.unique(
        np.asarray(cache.row_index[selected_events], dtype=np.int64)
    )
    if len(rows) == 0:
        raise ValueError(f"{split} exposure rows are empty after EID filtering")
    return rows


def maybe_wrap_data_parallel(
    model: TimesNetExposureAutoencoder,
    args: argparse.Namespace,
    device: torch.device,
    logger,
):
    if not args.data_parallel:
        return model
    if device.type != "cuda":
        raise ValueError("--data_parallel requires --device cuda or cuda:<id>")
    if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
        raise ValueError("--data_parallel requires at least two CUDA devices")
    primary = (
        int(device.index)
        if device.index is not None
        else int(torch.cuda.current_device())
    )
    device_ids = (
        args.gpu_ids
        if args.gpu_ids
        else list(range(torch.cuda.device_count()))
    )
    device_ids = [primary, *[idx for idx in device_ids if idx != primary]]
    if len(device_ids) < 2:
        raise ValueError("--data_parallel needs at least two device ids")
    if any(idx < 0 or idx >= torch.cuda.device_count() for idx in device_ids):
        raise ValueError(f"CUDA device id is out of range: {device_ids}")
    logger.info(f"Using DataParallel on CUDA devices: {device_ids}")
    return torch.nn.DataParallel(
        model, device_ids=device_ids, output_device=primary
    )


def unwrap_model(model) -> TimesNetExposureAutoencoder:
    return model.module if isinstance(model, torch.nn.DataParallel) else model


def channel_stats(
    cache: ExposureCache, rows: np.ndarray, chunk_size: int = 256
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    results = []
    for source in (cache.daily, cache.monthly):
        sums = np.zeros(source.shape[-1], dtype=np.float64)
        squares = np.zeros_like(sums)
        counts = np.zeros_like(sums)
        for start in tqdm(range(0, len(rows), chunk_size), desc="Channel statistics"):
            values = np.asarray(source[rows[start:start + chunk_size]], dtype=np.float64)
            finite = np.isfinite(values)
            clean = np.where(finite, values, 0.0)
            sums += clean.sum(axis=(0, 1))
            squares += np.square(clean).sum(axis=(0, 1))
            counts += finite.sum(axis=(0, 1))
        mean = sums / np.maximum(counts, 1.0)
        variance = squares / np.maximum(counts, 1.0) - np.square(mean)
        std = np.sqrt(np.maximum(variance, 1e-12))
        results.extend([mean.astype(np.float32), std.astype(np.float32)])
    return tuple(results)


def masked_mse(
    prediction: torch.Tensor, target: torch.Tensor, mask: torch.Tensor
) -> torch.Tensor:
    error = (prediction - target).square() * mask
    return error.sum() / mask.sum().clamp_min(1.0)


def run_epoch(
    model,
    loader: DataLoader,
    device: torch.device,
    stats: tuple[torch.Tensor, ...],
    mask_ratio: float,
    optimizer: AdamW | None,
    scaler: torch.amp.GradScaler,
    grad_clip: float,
    amp_enabled: bool,
) -> float:
    training = optimizer is not None
    model.train(training)
    total_loss = 0.0
    total_samples = 0
    daily_mean, daily_std, monthly_mean, monthly_std = stats
    context = torch.enable_grad if training else torch.no_grad
    with context():
        for batch in tqdm(loader, desc="train" if training else "val", leave=False):
            daily = batch["daily"].to(device, non_blocking=True)
            monthly = batch["monthly"].to(device, non_blocking=True)
            daily_observed = torch.isfinite(daily)
            monthly_observed = torch.isfinite(monthly)
            daily = (torch.nan_to_num(daily) - daily_mean) / daily_std
            monthly = (torch.nan_to_num(monthly) - monthly_mean) / monthly_std
            daily = daily * daily_observed
            monthly = monthly * monthly_observed

            if training and mask_ratio > 0:
                daily_input_mask = daily_observed & (
                    torch.rand_like(daily) >= mask_ratio
                )
                monthly_input_mask = monthly_observed & (
                    torch.rand_like(monthly) >= mask_ratio
                )
            else:
                daily_input_mask = daily_observed
                monthly_input_mask = monthly_observed
            daily_input = daily * daily_input_mask
            monthly_input = monthly * monthly_input_mask

            if training:
                optimizer.zero_grad(set_to_none=True)
            with torch.autocast(
                device_type=device.type, dtype=torch.float16,
                enabled=amp_enabled,
            ):
                daily_hat, monthly_hat, _ = model(
                    daily_input, monthly_input,
                    daily_input_mask, monthly_input_mask,
                )
                loss = (
                    masked_mse(daily_hat, daily, daily_observed)
                    + masked_mse(monthly_hat, monthly, monthly_observed)
                )
            if training:
                scaler.scale(loss).backward()
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
                scaler.step(optimizer)
                scaler.update()
            batch_size = daily.size(0)
            total_loss += float(loss.detach()) * batch_size
            total_samples += batch_size
    return total_loss / max(total_samples, 1)


def learning_rate(epoch: int, args: argparse.Namespace) -> float:
    if epoch < args.warmup_epochs:
        return args.base_lr * (epoch + 1) / max(args.warmup_epochs, 1)
    progress = (epoch - args.warmup_epochs) / max(
        args.max_epochs - args.warmup_epochs - 1, 1
    )
    return args.base_lr * 0.5 * (1.0 + math.cos(math.pi * progress))


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 stamp: f"exposure_ae_{stamp}", Path(args.runs_root)
    )
    logger = setup_logging(run_dir)
    cache = ExposureCache(args.exposure_cache_dir)
    train_eids = load_eid_file(args.train_eid_file)
    val_eids = load_eid_file(args.val_eid_file)
    if train_eids & val_eids:
        raise ValueError("train and validation EID files must be disjoint")
    train_rows = select_rows(cache, train_eids, "Training")
    val_rows = select_rows(cache, val_eids, "Validation")
    raw_stats = channel_stats(cache, train_rows)
    stats = tuple(
        torch.as_tensor(value, device=device).view(1, 1, -1)
        for value in raw_stats
    )
    loader_kwargs = dict(
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
        persistent_workers=args.num_workers > 0,
    )
    train_loader = DataLoader(
        ExposureWindowDataset(cache, train_rows), shuffle=True, **loader_kwargs
    )
    val_loader = DataLoader(
        ExposureWindowDataset(cache, val_rows), shuffle=False, **loader_kwargs
    )
    model = TimesNetExposureAutoencoder(
        n_embd=args.n_embd, d_model=args.d_model, n_layers=args.n_layers,
        top_k=args.top_k, n_convnext_blocks=args.n_convnext_blocks,
        conv_kernel_size=args.conv_kernel_size, mlp_ratio=args.mlp_ratio,
        dropout=args.dropout,
    ).to(device)
    model = maybe_wrap_data_parallel(model, args, device, logger)
    optimizer = AdamW(
        model.parameters(), lr=args.base_lr,
        weight_decay=args.weight_decay, betas=(0.9, 0.95),
    )
    amp_enabled = bool(args.amp and device.type == "cuda")
    scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled)
    logger.info(
        f"Run {run_name}: device={device}, train_rows={len(train_rows):,}, "
        f"val_rows={len(val_rows):,}"
    )
    config = vars(args) | {
        "train_rows": len(train_rows),
        "val_rows": len(val_rows),
        "daily_mean": raw_stats[0].tolist(),
        "daily_std": raw_stats[1].tolist(),
        "monthly_mean": raw_stats[2].tolist(),
        "monthly_std": raw_stats[3].tolist(),
    }
    (run_dir / "train_config.json").write_text(
        json.dumps(config, indent=2), encoding="utf-8"
    )

    best_loss = float("inf")
    stale_epochs = 0
    history = []
    for epoch in range(args.max_epochs):
        lr = learning_rate(epoch, args)
        for group in optimizer.param_groups:
            group["lr"] = lr
        train_loss = run_epoch(
            model, train_loader, device, stats, args.mask_ratio, optimizer,
            scaler, args.grad_clip, amp_enabled,
        )
        val_loss = run_epoch(
            model, val_loader, device, stats, 0.0, None,
            scaler, args.grad_clip, amp_enabled,
        )
        logger.info(
            f"Epoch {epoch + 1:03d} | lr={lr:.3e} | "
            f"train={train_loss:.6f} | val={val_loss:.6f}"
        )
        history.append(
            {"epoch": epoch + 1, "lr": lr,
             "train_loss": train_loss, "val_loss": val_loss}
        )
        if val_loss < best_loss:
            best_loss = val_loss
            stale_epochs = 0
            checkpoint_model = unwrap_model(model)
            torch.save(
                {
                    "model_state_dict": checkpoint_model.state_dict(),
                    "encoder_state_dict": checkpoint_model.encoder.state_dict(),
                    "model_config": {
                        key: config[key] for key in (
                            "n_embd", "d_model", "n_layers", "top_k",
                            "n_convnext_blocks", "conv_kernel_size",
                            "mlp_ratio", "dropout",
                        )
                    },
                    "normalization": {
                        "daily_mean": raw_stats[0],
                        "daily_std": raw_stats[1],
                        "monthly_mean": raw_stats[2],
                        "monthly_std": raw_stats[3],
                    },
                    "epoch": epoch + 1,
                    "val_loss": val_loss,
                },
                run_dir / "best.pt",
            )
        else:
            stale_epochs += 1
        (run_dir / "history.json").write_text(
            json.dumps(history, indent=2), encoding="utf-8"
        )
        if stale_epochs >= args.patience:
            logger.info(f"Early stopping after {epoch + 1} epochs")
            break
    logger.info(f"Best validation loss: {best_loss:.6f}")
    logger.info(f"Checkpoint: {run_dir / 'best.pt'}")


if __name__ == "__main__":
    main()
