"""
inference.py — Inference utilities for SleepMaMi pretrained / finetuned checkpoints.

This module provides a thin, dependency-light wrapper around the two
classification backbones used in the paper and exposes three entry points:

  1. extract_micro_features(...)
        -> ClassificationMultiConv1d (epoch-level Micro-Encoder).
           Returns the fused micro feature for each 30-s epoch.

  2. extract_macro_micro_features(...)
        -> ClassificationMultiConv1dMamba (full-night model).
           Returns both the full-night Macro feature (sequence encoder output)
           and the per-epoch Micro feature.

  3. predict_sleep_stages(...)
        -> ClassificationMultiConv1dMamba + classification head.
           Returns the predicted sleep stage for every valid epoch of the night.

Inputs are read from per-modality `.npy` files holding the full night as a
sequence of 30-s epochs, shape (N, sig_len) where sig_len = sig_len_target *
TARGET_RESAMPLE (3000 for the default 30 s @ 100 Hz).

Example
-------
    # --checkpoint defaults to checkpoints/<task>.safetensors
    python inference.py --task sleep_staging --npy_dir /path/to/night

    python inference.py --task micro --npy_dir /path/to/night --num_epochs 20

    python inference.py --task macro_micro \
        --checkpoint checkpoints/macro_micro.safetensors \
        --npy_dir /path/to/night
"""

import argparse
import contextlib
import os

import numpy as np
import torch
from safetensors.torch import load_file as load_safetensors

import models
from configs.global_cfg import TARGET_RESAMPLE
from models.utils import flip_padded_sequences


# Default 6-modality setup used throughout the paper.
DEFAULT_MODALITIES = ["EEG", "EOG", "ECG", "RESP", "EMG", "SaO2"]


# --------------------------------------------------------------------------- #
# Checkpoint loading
# --------------------------------------------------------------------------- #
def _load_state_dict(ckpt_path, map_location="cpu"):
    """Return the raw parameter dict from a released `.safetensors` file
    (weights only, keys without the `model.` prefix) or a Lightning / plain
    `.ckpt` (e.g. one produced by `finetune.py`)."""
    if ckpt_path.endswith(".safetensors"):
        return load_safetensors(ckpt_path, device=str(map_location))
    # Lightning checkpoints pickle non-tensor objects (hparams, numpy arrays),
    # so full unpickling is required — only load .ckpt files you trust.
    ckpt = torch.load(ckpt_path, map_location=map_location, weights_only=False)
    if isinstance(ckpt, dict) and "state_dict" in ckpt:
        return ckpt["state_dict"]
    return ckpt


def load_weights(model, ckpt_path, device="cpu", verbose=True):
    """
    Load weights into `model`, handling the `model.` prefix that PyTorch
    Lightning adds (the LightningModule stores the backbone as `self.model`).

    This mirrors `ConvMAE1d_Finetune.load_pretrained` in `finetune.py`:
    position embeddings are skipped (they are buffers regenerated at build
    time) and matching is attempted both with and without the `model.` prefix.
    """
    src = _load_state_dict(ckpt_path, map_location=device)
    tgt_keys = model.state_dict().keys()

    new_state_dict = {}
    for k in tgt_keys:
        if "pos_embed" in k:
            continue
        if k in src:
            new_state_dict[k] = src[k]
        elif f"model.{k}" in src:
            new_state_dict[k] = src[f"model.{k}"]

    msg = model.load_state_dict(new_state_dict, strict=False)

    if verbose:
        missing = [k for k in msg.missing_keys if "pos_embed" not in k]
        print(f"[load_weights] Loaded {len(new_state_dict)} tensors from {ckpt_path}")
        if missing:
            print(f"[load_weights] Missing (left at init): {missing}")
        if msg.unexpected_keys:
            print(f"[load_weights] Unexpected (ignored): {msg.unexpected_keys}")
    return model


# --------------------------------------------------------------------------- #
# Model construction
# --------------------------------------------------------------------------- #
def build_micro_model(modalities=DEFAULT_MODALITIES, sig_len_target=30,
                      patch_size=50, num_experts=4, num_experts_per=2,
                      num_classes=5, use_revin=True):
    """ClassificationMultiConv1d — epoch-level Micro-Encoder (`classification_multi`)."""
    return models.classification_multi(
        modalities=modalities,
        sig_len=sig_len_target * TARGET_RESAMPLE,
        patch_size=patch_size,
        num_experts=num_experts,
        num_experts_per=num_experts_per,
        num_classes=num_classes,
        use_shared=True,
        use_moe=True,
        use_revin=use_revin,
        ft_task="sleep_staging",
    )


def build_full_night_model(modalities=DEFAULT_MODALITIES, sig_len_target=30,
                           patch_size=50, num_experts=4, num_experts_per=2,
                           num_classes=5, use_revin=True,
                           use_micro=False, use_lstm=False, use_transformer=False,
                           bidi=True):
    """ClassificationMultiConv1dMamba — full-night model (`classification_multi_full_night`)."""
    return models.classification_multi_full_night(
        modalities=modalities,
        sig_len=sig_len_target * TARGET_RESAMPLE,
        patch_size=patch_size,
        num_experts=num_experts,
        num_experts_per=num_experts_per,
        num_classes=num_classes,
        use_shared=True,
        use_moe=True,
        use_revin=use_revin,
        ft_task="sleep_staging",
        use_micro=use_micro,
        use_lstm=use_lstm,
        use_transformer=use_transformer,
        bidi=False if use_transformer else bidi,
    )


# --------------------------------------------------------------------------- #
# Input preparation (reuses the training/eval data pipeline)
# --------------------------------------------------------------------------- #
def _to_device(inputs, device, dtype):
    return {m: v.to(device=device, dtype=dtype) for m, v in inputs.items()}


def _autocast(model, dtype=torch.bfloat16):
    """Autocast context for the forward pass.

    FlashAttention (used in the encoder/fuser) only supports fp16/bf16, so a
    pure float32 forward raises an error. Running under CUDA autocast keeps the
    weights in their loaded dtype while computing attention in `dtype`, exactly
    reproducing the `bf16-mixed` precision the checkpoints were trained with.
    On CPU (no FlashAttention) autocast is a no-op.
    """
    device = next(model.parameters()).device
    if device.type == "cuda":
        return torch.autocast(device_type="cuda", dtype=dtype)
    return contextlib.nullcontext()


# --------------------------------------------------------------------------- #
# Input preparation from .npy files
# --------------------------------------------------------------------------- #
# Each modality is stored as a single .npy array holding the full night as a
# sequence of 30-s epochs, shape (N, sig_len) where sig_len = sig_len_target *
# TARGET_RESAMPLE (3000 for the default 30 s @ 100 Hz). A flat (N*sig_len,)
# array is also accepted and reshaped automatically.
def _read_npy_arrays(npy_paths, modalities, sig_len):
    """Load and shape per-modality arrays to (N, sig_len). Validates epoch counts."""
    arrays, n_epochs = {}, None
    for m in modalities:
        if m not in npy_paths:
            raise ValueError(f"No .npy path provided for modality '{m}'")
        arr = np.load(npy_paths[m]).astype(np.float32)
        if arr.ndim == 1 or arr.shape[-1] != sig_len:
            # if m == 'SaO2':
            #     arr = arr.reshape(-1, 30)
            # else:
            #     arr = arr.reshape(-1, sig_len)
            arr = arr.reshape(-1, sig_len)
        if n_epochs is None:
            n_epochs = arr.shape[0]
        elif arr.shape[0] != n_epochs:
            raise ValueError(
                f"Modality '{m}' has {arr.shape[0]} epochs, expected {n_epochs} "
                f"(all modalities must share the same number of epochs).")
        arrays[m] = arr
    return arrays, n_epochs


def _normalize_epoch(arr, modality):
    """Per-modality normalization matching `data/data_loader.py`.

    - SaO2: scaled to [0, 1] via (clip(x, 0.7, 1.0) - 0.7) / 0.3. Percent-valued
      inputs (max > 1.5) are divided by 100 first. The model downsamples SaO2
      internally, so it must be provided at the signal rate (sig_len samples).
    - others: per-epoch z-normalization. With `use_revin=True` (default) the
      model re-applies this internally, so the step is a harmless no-op there;
      it matters when RevIN is disabled.
    """
    if modality == "SaO2":
        x = arr.copy()
        if np.nanmax(x) > 1.5:        # percent -> fraction
            x = x / 100.0
        return (np.clip(x, 0.7, 1.0) - 0.7) / 0.3
    mu = arr.mean(axis=1, keepdims=True)
    std = arr.std(axis=1, keepdims=True)
    std = np.where(std < 1e-8, 1.0, std)
    return (arr - mu) / std


def load_npy_epoch_inputs(npy_paths, modalities=DEFAULT_MODALITIES, sig_len_target=30,
                          normalize=True, num_epochs=None, device="cpu",
                          dtype=torch.float32):
    """Epoch-level inputs for the Micro-Encoder, built from .npy files.

    Returns dict[str, Tensor] each of shape (N, sig_len, 1).
    """
    sig_len = sig_len_target * TARGET_RESAMPLE
    arrays, n = _read_npy_arrays(npy_paths, modalities, sig_len)
    if num_epochs is not None:
        n = min(num_epochs, n)
    inputs = {}
    for m in modalities:
        a = arrays[m][:n]
        if normalize:
            a = _normalize_epoch(a, m)
        inputs[m] = torch.from_numpy(np.ascontiguousarray(a)).unsqueeze(-1)
    return _to_device(inputs, device, dtype)


def load_npy_full_night_inputs(npy_paths, modalities=DEFAULT_MODALITIES, sig_len_target=30,
                               normalize=True, device="cpu", dtype=torch.float32):
    """Full-night inputs for the Mamba model, built from .npy files.

    Returns
    -------
    inputs : dict[str, Tensor]  each of shape (1, N*sig_len, 1)
    mask   : BoolTensor (1, N)  all True — a single .npy night has no padding.
    """
    sig_len = sig_len_target * TARGET_RESAMPLE
    arrays, n = _read_npy_arrays(npy_paths, modalities, sig_len)
    inputs = {}
    for m in modalities:
        a = arrays[m]
        if normalize:
            a = _normalize_epoch(a, m)
        inputs[m] = torch.from_numpy(np.ascontiguousarray(a).reshape(1, n * sig_len, 1))
    mask = torch.ones(1, n, dtype=torch.bool, device=device)
    return _to_device(inputs, device, dtype), mask


# --------------------------------------------------------------------------- #
# Inference entry points
# --------------------------------------------------------------------------- #
@torch.inference_mode()
def extract_micro_features(model, inputs, amp_dtype=torch.bfloat16):
    """
    (1) Micro feature from ClassificationMultiConv1d.

    Returns a tensor of shape (E, embed_dim[-1] * num_modalities) — one fused
    micro embedding per 30-s epoch.
    """
    model.eval()
    with _autocast(model, amp_dtype):
        feats = model.forward_features(inputs)
    return feats


@torch.inference_mode()
def extract_macro_micro_features(model, inputs, mask, amp_dtype=torch.bfloat16):
    """
    (2) Macro + Micro features from ClassificationMultiConv1dMamba.

    Returns
    -------
    dict with
        'macro' : (1, num_epochs, macro_dim)   full-night sequence feature
                  (for bidirectional models the forward and flipped-backward
                  streams are concatenated, exactly as in finetuning).
        'micro' : (1, num_epochs, micro_dim)   per-epoch micro feature.
    """
    model.eval()
    with _autocast(model, amp_dtype):
        macro, micro = model.forward_features(inputs, mask=mask)

    if not model.use_micro and getattr(model, "bidi", False):
        xf, xr = macro
        xr = flip_padded_sequences(xr, mask)
        macro = torch.cat([xf, xr], dim=-1)

    return {"macro": macro, "micro": micro}


@torch.inference_mode()
def predict_sleep_stages(model, inputs, mask, return_logits=False, amp_dtype=torch.bfloat16):
    """
    (3) Sleep-staging inference (Macro + Micro features + classification head).

    Returns
    -------
    pred   : LongTensor (num_valid_epochs,)   predicted stage per valid epoch.
    logits : (optional) FloatTensor (num_valid_epochs, num_classes).
    """
    model.eval()
    with _autocast(model, amp_dtype):
        logits = model(inputs, mask=mask)          # (1, num_epochs, num_classes)
    valid = mask[0].bool()
    logits = logits[0][valid]                      # (num_valid_epochs, num_classes)
    pred = logits.argmax(dim=-1)
    if return_logits:
        return pred, logits
    return pred


# --------------------------------------------------------------------------- #
# CLI demo
# --------------------------------------------------------------------------- #
def _parse_args():
    p = argparse.ArgumentParser(description="SleepMaMi inference demo")
    p.add_argument("--task", required=True,
                   choices=["micro", "macro_micro", "sleep_staging"],
                   help="Which inference to run.")
    p.add_argument("--checkpoint", default=None,
                   help="Path to a .safetensors (released weights) or Lightning .ckpt file. "
                        "Default: checkpoints/<task>.safetensors")
    # Per-modality .npy files (full night, shape (N, sig_len)).
    p.add_argument("--npy_dir", default=None,
                   help="Directory containing one '<MODALITY>.npy' per modality.")
    p.add_argument("--npy", nargs="+", default=None, metavar="MOD=PATH",
                   help="Explicit per-modality .npy paths, e.g. EEG=eeg.npy EOG=eog.npy")
    p.add_argument("--no_normalize", dest="normalize", action="store_false",
                   help="Skip built-in normalization (use if .npy is already preprocessed).")
    p.set_defaults(normalize=True)
    p.add_argument("--modalities", nargs="+", default=DEFAULT_MODALITIES)
    p.add_argument("--sig_len_target", type=int, default=30)
    p.add_argument("--patch_size", type=int, default=50)
    p.add_argument("--num_classes", type=int, default=5)
    p.add_argument("--num_experts", type=int, default=4)
    p.add_argument("--num_experts_per", type=int, default=2)
    p.add_argument("--no_revin", action="store_true", help="Disable RevIN normalization.")
    p.add_argument("--num_epochs", type=int, default=None,
                   help="(micro only) number of epochs to embed; default = all.")
    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    p.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
    return p.parse_args()


def _resolve_npy_paths(args):
    """Build a {modality: path} dict from --npy / --npy_dir."""
    if args.npy:
        paths = {}
        for item in args.npy:
            if "=" not in item:
                raise ValueError(f"--npy entries must be MOD=PATH, got '{item}'")
            k, v = item.split("=", 1)
            paths[k] = v
        return paths
    if args.npy_dir:
        return {m: os.path.join(args.npy_dir, f"{m}.npy") for m in args.modalities}
    raise SystemExit("Provide an input source: --npy_dir or --npy MOD=PATH ...")


def main():
    args = _parse_args()
    if args.checkpoint is None:
        args.checkpoint = os.path.join("checkpoints", f"{args.task}.safetensors")
    device = args.device
    dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16,
             "float16": torch.float16}[args.dtype]

    npy_paths = _resolve_npy_paths(args)
    common = dict(modalities=args.modalities, sig_len_target=args.sig_len_target,
                  patch_size=args.patch_size, num_experts=args.num_experts,
                  num_experts_per=args.num_experts_per, num_classes=args.num_classes,
                  use_revin=not args.no_revin)

    if args.task == "micro":
        model = build_micro_model(**common).to(device=device, dtype=dtype)
        load_weights(model, args.checkpoint, device=device)
        inputs = load_npy_epoch_inputs(
            npy_paths, modalities=args.modalities, sig_len_target=args.sig_len_target,
            normalize=args.normalize, num_epochs=args.num_epochs,
            device=device, dtype=dtype)
        feats = extract_micro_features(model, inputs)
        print(f"Micro features: {tuple(feats.shape)}  (epochs x feature_dim)")
        print(feats.float().cpu().numpy())

    else:  # macro_micro or sleep_staging — both use the full-night model
        model = build_full_night_model(**common).to(device=device, dtype=dtype)
        load_weights(model, args.checkpoint, device=device)
        inputs, mask = load_npy_full_night_inputs(
            npy_paths, modalities=args.modalities, sig_len_target=args.sig_len_target,
            normalize=args.normalize, device=device, dtype=dtype)

        if args.task == "macro_micro":
            out = extract_macro_micro_features(model, inputs, mask)
            print(f"Macro feature: {tuple(out['macro'].shape)}")
            print(f"Micro feature: {tuple(out['micro'].shape)}")
        else:
            pred = predict_sleep_stages(model, inputs, mask).cpu().numpy()
            print(f"Predicted stages ({len(pred)} epochs): {pred}")


if __name__ == "__main__":
    main()