from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple

import numpy as np
import torch
from torch.nn.utils.rnn import pad_sequence

from dataset import (
    DISEASE_HISTORY_MODE_TIMED,
    AllFutureHealthDataset,
    HealthDataset,
    normalize_disease_history_mode,
)
from model_architectures import resolve_model_architecture
from models import DeepHealth
from targets import PAD_IDX


def load_json_config(path: str | Path | None) -> Dict[str, Any]:
    if path is None:
        return {}
    config_path = Path(path)
    if not config_path.exists():
        return {}
    with config_path.open("r", encoding="utf-8") as file:
        return json.load(file)


def cfg_get(
    args: argparse.Namespace | Dict[str, Any] | None,
    cfg: Dict[str, Any],
    name: str,
    default: Any,
) -> Any:
    if args is not None:
        value = (
            args.get(name)
            if isinstance(args, dict)
            else getattr(args, name, None)
        )
        if value is not None:
            return value
    return cfg.get(name, default)


def resolve_eval_device(device_arg: Optional[str]) -> torch.device:
    device_name = device_arg or ("cuda" if torch.cuda.is_available() else "cpu")
    device = torch.device(device_name)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError(
            f"Requested device {device_name!r}, but CUDA is not available."
        )
    return device


def validate_training_mode_config(cfg: Dict[str, Any]) -> None:
    model_target_mode = str(
        cfg.get("model_target_mode", "next_token")
    ).lower()
    if model_target_mode not in {"next_token", "all_future"}:
        raise ValueError(
            "model_target_mode must be next_token or all_future, got "
            f"{model_target_mode!r}"
        )
    disease_history_mode = normalize_disease_history_mode(
        cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED)
    )
    if disease_history_mode != DISEASE_HISTORY_MODE_TIMED:
        expected = {
            "model_target_mode": "all_future",
            "time_mode": "relative",
            "dist_mode": "weibull",
            "model_architecture": "traj_mixer_v5",
        }
        actual = {
            "model_target_mode": model_target_mode,
            "time_mode": str(cfg.get("time_mode", "")).lower(),
            "dist_mode": str(cfg.get("dist_mode", "")).lower(),
            "model_architecture": str(
                cfg.get("model_architecture", "")
            ).lower(),
        }
        mismatches = [
            f"{name}={actual[name]!r} (expected {value!r})"
            for name, value in expected.items()
            if actual[name] != value
        ]
        extra_info_types = cfg.get("extra_info_types", None)
        if extra_info_types != []:
            mismatches.append("extra_info_types must be []")
        if mismatches:
            raise ValueError(
                f"disease_history_mode={disease_history_mode!r} is only valid "
                "for the no-extra TrajMixer + all_future + relative + Weibull "
                "ablation; " + "; ".join(mismatches)
            )

    if model_target_mode != "next_token":
        return

    time_mode = str(cfg.get("time_mode", "")).lower()
    target_mode = str(cfg.get("target_mode", "")).lower()
    if time_mode != "absolute" or target_mode != "delphi2m":
        raise ValueError(
            "next_token is reserved for Delphi2M reproduction and requires "
            "time_mode='absolute' and target_mode='delphi2m'; got "
            f"time_mode={time_mode!r}, target_mode={target_mode!r}"
        )


def split_indices(
    n: int,
    train_ratio: float,
    val_ratio: float,
    test_ratio: float,
    seed: int,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    total = float(train_ratio) + float(val_ratio) + float(test_ratio)
    if not np.isclose(total, 1.0, atol=1e-6):
        raise ValueError(
            f"train/val/test ratios must sum to 1.0, got {total}"
        )
    indices = np.random.RandomState(seed).permutation(n)
    n_train = int(n * train_ratio)
    n_val = int(n * val_ratio)
    return (
        indices[:n_train],
        indices[n_train:n_train + n_val],
        indices[n_train + n_val:],
    )


def select_indices_by_eid_file(
    dataset: Any,
    eid_file: str | Path,
) -> Tuple[np.ndarray, Path]:
    """Return dataset indices whose patient EIDs occur in ``eid_file``."""
    from train_util import load_eid_file

    path = Path(eid_file)
    if not path.is_absolute():
        direct = Path.cwd() / path
        path = direct if direct.is_file() else Path(__file__).resolve().parent / path
    if not path.is_file():
        raise FileNotFoundError(f"EID split file not found: {path}")

    samples = getattr(dataset, "samples", None)
    if samples is None:
        raise TypeError("EID-based evaluation requires dataset.samples")
    selected_eids = load_eid_file(path)
    indices = np.asarray(
        [
            index
            for index, sample in enumerate(samples)
            if int(sample["eid"]) in selected_eids
        ],
        dtype=np.int64,
    )
    if indices.size == 0:
        raise ValueError(
            f"No dataset patients matched the EID split file: {path}"
        )
    return indices, path.resolve()


def build_model_from_dataset(
    args: argparse.Namespace,
    cfg: Dict[str, Any],
    dataset: HealthDataset,
    state_dict: Optional[Dict[str, Any]] = None,
) -> DeepHealth:
    validate_training_mode_config(cfg)
    model_target_mode = str(
        cfg_get(args, cfg, "model_target_mode", "next_token")
    ).lower()
    if model_target_mode not in {"next_token", "all_future"}:
        raise ValueError(
            "model_target_mode must be next_token or all_future, got "
            f"{model_target_mode!r}"
        )
    risk_head_bias = bool(cfg_get(args, cfg, "risk_head_bias", False))
    if state_dict is not None:
        # The checkpoint schema is authoritative. This keeps all older
        # bias-free checkpoints loadable while restoring the new baseline bias.
        risk_head_bias = "risk_head.bias" in state_dict
    model_architecture = resolve_model_architecture(cfg, state_dict)
    continuous_value_center = None
    continuous_value_scale = None
    if dataset.n_cont_types > 0:
        scaling = str(cfg.get("continuous_value_scaling", "")).lower()
        if scaling != "robust":
            raise RuntimeError(
                "Continuous-variable checkpoints must declare "
                "continuous_value_scaling='robust'; unscaled checkpoints are "
                "not supported"
            )
        if state_dict is None:
            raise RuntimeError(
                "A checkpoint state_dict is required to restore RobustScale buffers"
            )
        center_key = "tokenizer.continuous_value_center"
        scale_key = "tokenizer.continuous_value_scale"
        missing = [
            key for key in (center_key, scale_key)
            if key not in state_dict
        ]
        if missing:
            raise RuntimeError(
                "Checkpoint is missing required RobustScale buffers: "
                + ", ".join(missing)
            )
        continuous_value_center = state_dict[center_key]
        continuous_value_scale = state_dict[scale_key]
    return DeepHealth(
        vocab_size=dataset.vocab_size,
        n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
        n_head=int(cfg_get(args, cfg, "n_head", 10)),
        n_layer=int(cfg["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=int(cfg_get(args, cfg, "n_bins", 16)),
        continuous_value_center=continuous_value_center,
        continuous_value_scale=continuous_value_scale,
        extra_pool_reduce=str(
            cfg_get(args, cfg, "extra_pool_reduce", "mean")
        ),
        target_mode=model_target_mode,
        time_mode=str(cfg_get(args, cfg, "time_mode", "absolute")),
        dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
        dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
        model_architecture=model_architecture,
        risk_head_bias=risk_head_bias,
    )


def validate_dataset_metadata(
    dataset: HealthDataset,
    cfg: Dict[str, Any],
) -> None:
    metadata = cfg.get("dataset_metadata")
    if not isinstance(metadata, dict):
        return

    actual: Dict[str, Any] = {
        "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],
    }
    mismatches = [
        f"{key}: train_config={metadata.get(key)!r}, current_dataset={value!r}"
        for key, value in actual.items()
        if key in metadata and metadata.get(key) != value
    ]
    if mismatches:
        raise RuntimeError(
            "Current dataset metadata does not match train_config.json. "
            "Use the same prepared data and extra_info_types as training. "
            + "; ".join(mismatches)
        )


def build_first_occurrence_map(
    dataset: HealthDataset,
    subset_indices: np.ndarray,
) -> Dict[int, Tuple[np.ndarray, np.ndarray]]:
    first_lists: Dict[int, List[Tuple[int, float]]] = {}
    for patient_id, dataset_index in enumerate(subset_indices.tolist()):
        sample = dataset.samples[int(dataset_index)]
        sequence_events = np.asarray(sample["event_seq"], dtype=np.int64)
        sequence_times = np.asarray(sample["time_seq"], dtype=np.float32)
        target_events = np.asarray(
            sample["target_event_seq"], dtype=np.int64
        )
        target_times = np.asarray(
            sample["target_time_seq"], dtype=np.float32
        )
        if sequence_events.size == 0 or target_events.size == 0:
            continue

        full_events = np.concatenate(
            [sequence_events, target_events[-1:]]
        )
        full_times = np.concatenate([sequence_times, target_times[-1:]])
        unique_tokens, first_indices = np.unique(
            full_events, return_index=True
        )
        for token, first_index in zip(
            unique_tokens.tolist(), first_indices.tolist()
        ):
            first_lists.setdefault(int(token), []).append(
                (patient_id, float(full_times[int(first_index)]))
            )

    return {
        int(token): (
            np.asarray([patient for patient, _ in pairs], dtype=np.int32),
            np.asarray([time for _, time in pairs], dtype=np.float32),
        )
        for token, pairs in first_lists.items()
        if pairs
    }


class AllFutureSequenceEvalDataset:
    """
    Eval-only sequence view for all-future checkpoints.

    All-future training uses the observed history without reusing the
    next-step view that contains imputed <NO_EVENT> gap tokens. Legacy label-1
    assessment events are removed by the shared base dataset for every
    extra-info selection.
    """

    def __init__(
        self,
        data_prefix: str,
        labels_file: str,
        min_history_events: int = 1,
        min_future_events: int = 1,
        extra_info_types: Iterable[int] | None = None,
        disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED,
    ) -> None:
        base = AllFutureHealthDataset(
            data_prefix=data_prefix,
            labels_file=labels_file,
            split="train",
            min_history_events=min_history_events,
            min_future_events=min_future_events,
            extra_info_types=extra_info_types,
            disease_history_mode=disease_history_mode,
        )

        self.base = base
        self.label_code_to_id = base.label_code_to_id
        self.label_id_to_code = base.label_id_to_code
        self.vocab_size = base.vocab_size
        self.n_types = base.n_types
        self.n_cont_types = base.n_cont_types
        self.n_categories = base.n_categories
        self.cont_type_ids = base.cont_type_ids
        self.extra_info_types = base.extra_info_types
        self.disease_history_mode = base.disease_history_mode

        self.samples: List[Dict[str, Any]] = []
        for patient in base.patients:
            labels = np.asarray(patient["labels"], dtype=np.int64)
            times = np.asarray(patient["times"], dtype=np.float32)
            if labels.size < 2:
                continue
            self.samples.append(
                {
                    "eid": int(patient["eid"]),
                    "event_seq": labels[:-1],
                    "time_seq": times[:-1],
                    "target_event_seq": labels[1:],
                    "target_time_seq": times[1:],
                    "sex": int(patient["sex"]),
                    "other_type": np.asarray(patient["other_type"], dtype=np.int64),
                    "other_value": np.asarray(patient["other_value"], dtype=np.float32),
                    "other_value_kind": np.asarray(patient["other_value_kind"], dtype=np.int64),
                    "other_time": np.asarray(patient["other_time"], dtype=np.float32),
                }
            )

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

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        s = self.samples[idx]
        return {
            "event_seq": torch.from_numpy(s["event_seq"]).long(),
            "time_seq": torch.from_numpy(s["time_seq"]).float(),
            "target_event_seq": torch.from_numpy(s["target_event_seq"]).long(),
            "target_time_seq": torch.from_numpy(s["target_time_seq"]).float(),
            "sex": torch.tensor(s["sex"], dtype=torch.long),
            "other_type": torch.from_numpy(s["other_type"]).long(),
            "other_value": torch.from_numpy(s["other_value"]).float(),
            "other_value_kind": torch.from_numpy(s["other_value_kind"]).long(),
            "other_time": torch.from_numpy(s["other_time"]).float(),
        }


def load_sequence_eval_dataset(
    *,
    model_target_mode: str,
    data_prefix: str,
    labels_file: str,
    no_event_interval_years: float,
    min_history_events: int,
    min_future_events: int,
    extra_info_types: Iterable[int] | None,
    disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED,
):
    mode = str(model_target_mode).lower()
    if mode == "next_token":
        return HealthDataset(
            data_prefix=data_prefix,
            labels_file=labels_file,
            no_event_interval_years=no_event_interval_years,
            extra_info_types=extra_info_types,
        )
    if mode == "all_future":
        return AllFutureSequenceEvalDataset(
            data_prefix=data_prefix,
            labels_file=labels_file,
            min_history_events=min_history_events,
            min_future_events=min_future_events,
            extra_info_types=extra_info_types,
            disease_history_mode=disease_history_mode,
        )
    raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")


def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
    event_seq = pad_sequence(
        [s["event_seq"] for s in batch], batch_first=True, padding_value=PAD_IDX
    )
    time_seq = pad_sequence(
        [s["time_seq"] for s in batch], batch_first=True, padding_value=0.0
    )
    target_event_seq = pad_sequence(
        [s["target_event_seq"] for s in batch], batch_first=True, padding_value=PAD_IDX
    )
    target_time_seq = pad_sequence(
        [s["target_time_seq"] for s in batch], batch_first=True, padding_value=0.0
    )
    other_type = pad_sequence(
        [s["other_type"] for s in batch], batch_first=True, padding_value=0
    )
    other_value = pad_sequence(
        [s["other_value"] for s in batch], batch_first=True, padding_value=0.0
    )
    other_value_kind = pad_sequence(
        [s["other_value_kind"] for s in batch], batch_first=True, padding_value=0
    )
    other_time = pad_sequence(
        [s["other_time"] for s in batch], batch_first=True, padding_value=0.0
    )

    return {
        "event_seq": event_seq,
        "time_seq": time_seq,
        "padding_mask": event_seq > PAD_IDX,
        "target_event_seq": target_event_seq,
        "target_time_seq": target_time_seq,
        "readout_mask": event_seq > PAD_IDX,
        "sex": torch.stack([s["sex"] for s in batch]),
        "other_type": other_type,
        "other_value": other_value,
        "other_value_kind": other_value_kind,
        "other_time": other_time,
    }