from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd

from sepsis_mcp.constants import STATIC_COLUMNS
from sepsis_mcp.dataset import build_labeled_patient_frame


def _build_window_matrices(
    patient_frame: pd.DataFrame,
    sample_index: int,
    dynamic_columns: list[str],
    lookback_hours: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    values = np.zeros((lookback_hours, len(dynamic_columns)), dtype=np.float32)
    masks = np.zeros((lookback_hours, len(dynamic_columns)), dtype=np.float32)
    deltas = np.zeros((lookback_hours, len(dynamic_columns)), dtype=np.float32)

    window_start = max(0, sample_index - lookback_hours + 1)
    actual_window = patient_frame.iloc[window_start : sample_index + 1]
    pad_length = lookback_hours - len(actual_window)

    for feature_offset, column in enumerate(dynamic_columns):
        series = actual_window[column].tolist()
        for time_offset, value in enumerate(series, start=pad_length):
            if pd.notna(value):
                values[time_offset, feature_offset] = float(value)
                masks[time_offset, feature_offset] = 1.0

            if time_offset == 0:
                deltas[time_offset, feature_offset] = 0.0
            elif masks[time_offset, feature_offset] == 1.0:
                deltas[time_offset, feature_offset] = 0.0
            else:
                previous_delta = deltas[time_offset - 1, feature_offset]
                deltas[time_offset, feature_offset] = previous_delta + 1.0

    return values, masks, deltas


def _global_missing_rate(masks: np.ndarray) -> float:
    return float(1.0 - masks.mean())


def build_patient_grud_samples(
    patient_frame: pd.DataFrame,
    patient_id: str,
    dynamic_columns: list[str],
    lookback_hours: int = 12,
    horizon_hours: int = 6,
) -> list[dict[str, Any]]:
    labeled = build_labeled_patient_frame(
        patient_frame,
        patient_id=patient_id,
        horizon_hours=horizon_hours,
    )

    samples: list[dict[str, Any]] = []
    for _, row in labeled.iterrows():
        sample_index = int(row["sample_index"])
        values, masks, deltas = _build_window_matrices(
            patient_frame,
            sample_index=sample_index,
            dynamic_columns=dynamic_columns,
            lookback_hours=lookback_hours,
        )
        current_row = patient_frame.iloc[sample_index]
        static = np.array(
            [
                float(current_row[column]) if pd.notna(current_row[column]) else 0.0
                for column in STATIC_COLUMNS
            ],
            dtype=np.float32,
        )
        samples.append(
            {
                "patient_id": patient_id,
                "sample_index": sample_index,
                "values": values,
                "masks": masks,
                "deltas": deltas,
                "static": static,
                "label": int(row["target_in_6h"]),
                "global_missing_rate": _global_missing_rate(masks),
            }
        )

    return samples


def stack_grud_samples(samples: list[dict[str, Any]]) -> dict[str, Any]:
    return {
        "values": np.stack([sample["values"] for sample in samples]).astype(np.float32),
        "masks": np.stack([sample["masks"] for sample in samples]).astype(np.float32),
        "deltas": np.stack([sample["deltas"] for sample in samples]).astype(np.float32),
        "static": np.stack([sample["static"] for sample in samples]).astype(np.float32),
        "labels": np.array([sample["label"] for sample in samples], dtype=np.float32),
        "global_missing_rates": np.array(
            [sample["global_missing_rate"] for sample in samples],
            dtype=np.float32,
        ),
        "patient_ids": [sample["patient_id"] for sample in samples],
        "sample_indices": [sample["sample_index"] for sample in samples],
    }


def fit_grud_scaler(stacked: dict[str, Any]) -> dict[str, np.ndarray]:
    values = stacked["values"]
    masks = stacked["masks"]
    static = stacked["static"]

    dynamic_mean = np.zeros(values.shape[-1], dtype=np.float32)
    dynamic_std = np.ones(values.shape[-1], dtype=np.float32)
    for feature_index in range(values.shape[-1]):
        observed = values[:, :, feature_index][masks[:, :, feature_index] == 1]
        if observed.size:
            dynamic_mean[feature_index] = np.float32(observed.mean())
            feature_std = np.float32(observed.std())
            dynamic_std[feature_index] = feature_std if feature_std > 0 else 1.0

    static_mean = np.nanmean(static, axis=0).astype(np.float32)
    static_std = np.nanstd(static, axis=0).astype(np.float32)
    static_std[static_std == 0] = 1.0

    return {
        "dynamic_mean": dynamic_mean,
        "dynamic_std": dynamic_std,
        "static_mean": static_mean,
        "static_std": static_std,
    }


def transform_grud_stacked(
    stacked: dict[str, Any],
    scaler: dict[str, np.ndarray],
) -> dict[str, Any]:
    transformed = dict(stacked)
    values = stacked["values"].copy()
    masks = stacked["masks"]

    for feature_index in range(values.shape[-1]):
        observed = masks[:, :, feature_index] == 1
        values[:, :, feature_index][observed] = (
            values[:, :, feature_index][observed] - scaler["dynamic_mean"][feature_index]
        ) / scaler["dynamic_std"][feature_index]
        values[:, :, feature_index][~observed] = 0.0

    static = stacked["static"].copy()
    static = np.where(np.isfinite(static), static, scaler["static_mean"])
    static = (static - scaler["static_mean"]) / scaler["static_std"]

    transformed["values"] = values.astype(np.float32)
    transformed["static"] = static.astype(np.float32)
    return transformed