from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import torch
from torch.utils.data import Dataset


def validate_array_shape(X: np.ndarray) -> np.ndarray:
    X = np.asarray(X, dtype=np.float32)
    if X.ndim != 3:
        raise ValueError(f"Expected X with 3 dimensions [n_samples, n_steps, n_features], got {X.shape}.")
    return X


def validate_timestamp_shape(timestamps: np.ndarray, X: np.ndarray) -> np.ndarray:
    timestamps = np.asarray(timestamps, dtype=np.float32)
    if timestamps.ndim == 2:
        if timestamps.shape != X.shape[:2]:
            raise ValueError(
                "Expected timestamps with shape [n_samples, n_steps] when using 2D timestamps, "
                f"got {timestamps.shape} for X shape {X.shape}."
            )
        timestamps = np.repeat(timestamps[:, :, None], X.shape[2], axis=2)
    elif timestamps.ndim == 3:
        if timestamps.shape[0] != X.shape[0] or timestamps.shape[1] != X.shape[1]:
            raise ValueError(
                "Expected timestamps to match X on sample and step dimensions, "
                f"got timestamps shape {timestamps.shape} for X shape {X.shape}."
            )
        if timestamps.shape[2] == 1:
            timestamps = np.repeat(timestamps, X.shape[2], axis=2)
        elif timestamps.shape[2] != X.shape[2]:
            raise ValueError(
                "Expected timestamps feature dimension to be 1 or match X, "
                f"got timestamps shape {timestamps.shape} for X shape {X.shape}."
            )
    else:
        raise ValueError(
            "Expected timestamps with 2 or 3 dimensions, "
            f"got {timestamps.shape}."
        )

    if np.isnan(timestamps).any():
        raise ValueError("timestamps must not contain NaN.")
    if np.any(np.diff(timestamps, axis=1) < 0):
        raise ValueError("timestamps must be non-decreasing along the time dimension.")
    return timestamps.astype(np.float32)


def fill_and_get_mask_np(X: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    missing_mask = (~np.isnan(X)).astype(np.float32)
    filled = np.nan_to_num(X, nan=0.0).astype(np.float32)
    return filled, missing_mask


def parse_delta_np(missing_mask: np.ndarray, timestamps: np.ndarray | None = None) -> np.ndarray:
    deltas = np.zeros_like(missing_mask, dtype=np.float32)
    for t in range(1, missing_mask.shape[1]):
        if timestamps is None:
            step_gap = 1.0
        else:
            step_gap = np.abs(timestamps[:, t, :] - timestamps[:, t - 1, :])
        deltas[:, t, :] = step_gap + (1.0 - missing_mask[:, t - 1, :]) * deltas[:, t - 1, :]
    return deltas


def make_holdout_validation(
    X: np.ndarray,
    holdout_rate: float = 0.1,
    random_state: int | None = None,
    timestamps: np.ndarray | None = None,
) -> dict[str, np.ndarray]:
    X = validate_array_shape(X)
    if not 0.0 <= holdout_rate < 1.0:
        raise ValueError("holdout_rate must be in [0, 1).")
    if timestamps is not None:
        timestamps = validate_timestamp_shape(timestamps, X)

    rng = np.random.default_rng(random_state)
    observed = ~np.isnan(X)
    holdout = (rng.random(X.shape) < holdout_rate) & observed

    X_masked = X.copy()
    X_masked[holdout] = np.nan

    result = {
        "X": X_masked.astype(np.float32),
        "X_ori": X.astype(np.float32),
        "indicating_mask": holdout.astype(np.float32),
    }
    if timestamps is not None:
        result["timestamps"] = timestamps.astype(np.float32)
    return result


@dataclass
class SampleTensors:
    X: torch.Tensor
    missing_mask: torch.Tensor
    timestamps: torch.Tensor
    deltas: torch.Tensor
    back_X: torch.Tensor
    back_missing_mask: torch.Tensor
    back_timestamps: torch.Tensor
    back_deltas: torch.Tensor
    X_ori: torch.Tensor | None = None
    indicating_mask: torch.Tensor | None = None


class CurEDataset(Dataset):
    def __init__(
        self,
        X: np.ndarray,
        timestamps: np.ndarray | None = None,
        X_ori: np.ndarray | None = None,
        indicating_mask: np.ndarray | None = None,
    ) -> None:
        X = validate_array_shape(X)
        self.X_raw = X
        if timestamps is None:
            base_timestamps = np.broadcast_to(
                np.arange(X.shape[1], dtype=np.float32)[None, :, None],
                X.shape,
            ).copy()
        else:
            base_timestamps = validate_timestamp_shape(timestamps, X)
        self.timestamps = base_timestamps
        self.X_ori = None if X_ori is None else validate_array_shape(X_ori)
        self.indicating_mask = None if indicating_mask is None else validate_array_shape(indicating_mask)

        X_filled, missing_mask = fill_and_get_mask_np(X)
        deltas = parse_delta_np(missing_mask, self.timestamps)
        back_X = np.flip(X_filled, axis=1).copy()
        back_missing_mask = np.flip(missing_mask, axis=1).copy()
        back_timestamps = np.flip(self.timestamps, axis=1).copy()
        back_deltas = parse_delta_np(back_missing_mask, back_timestamps)

        self.X = torch.from_numpy(X_filled)
        self.missing_mask = torch.from_numpy(missing_mask)
        self.timestamps = torch.from_numpy(self.timestamps)
        self.deltas = torch.from_numpy(deltas)
        self.back_X = torch.from_numpy(back_X)
        self.back_missing_mask = torch.from_numpy(back_missing_mask)
        self.back_timestamps = torch.from_numpy(back_timestamps)
        self.back_deltas = torch.from_numpy(back_deltas)
        if self.X_ori is not None:
            self.X_ori = torch.from_numpy(self.X_ori.astype(np.float32))
        if self.indicating_mask is not None:
            self.indicating_mask = torch.from_numpy(self.indicating_mask.astype(np.float32))

    def __len__(self) -> int:
        return self.X.shape[0]

    def __getitem__(self, idx: int):
        sample = [
            torch.tensor(idx, dtype=torch.long),
            self.X[idx],
            self.missing_mask[idx],
            self.timestamps[idx],
            self.deltas[idx],
            self.back_X[idx],
            self.back_missing_mask[idx],
            self.back_timestamps[idx],
            self.back_deltas[idx],
        ]
        if self.X_ori is not None and self.indicating_mask is not None:
            sample.extend([self.X_ori[idx], self.indicating_mask[idx]])
        return sample
