"""PyTorch datasets and collator for variable-length NDVI sequences."""

from __future__ import annotations

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


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

def _build_time_feature(start_idx: int, length: int, encoding: str, max_dap: int) -> np.ndarray:
    """Build time feature vector: raw DAP indices or normalized progress."""
    t = np.arange(start_idx, start_idx + length, dtype=np.float32)
    if encoding == "raw":
        return t
    return t / max(1.0, float(max_dap))


def _sample_window(
    total: int,
    min_in: int, max_in: int,
    min_out: int, max_out: int,
    late_bias_p: float,
    late_window_jitter: int,
    rng: np.random.Generator | None = None,
    target_len_choices: list[int] | None = None,
) -> tuple[int, int, int]:
    """Sample (start_idx, input_len, target_len) for a sequence of given total length."""
    randint = rng.integers if rng is not None else np.random.randint

    in_upper = min(max_in, total - min_out)
    input_len = int(randint(min_in, in_upper + 1))

    if target_len_choices:
        valid = [t for t in target_len_choices if min_out <= t <= max_out and (input_len + t) <= total]
        if valid:
            target_len = int(rng.choice(valid) if rng else np.random.choice(valid))
        else:
            t_upper = min(max_out, total - input_len)
            target_len = int(randint(min_out, max(min_out, t_upper) + 1))
    else:
        t_upper = min(max_out, total - input_len)
        target_len = int(randint(min_out, max(min_out, t_upper) + 1))

    target_len = max(min_out, min(target_len, total - input_len))

    start_upper = total - input_len - target_len
    if start_upper < 0:
        target_len = max(min_out, total - input_len)
        start_upper = max(0, total - input_len - target_len)

    use_rand = (rng.random() if rng else np.random.rand()) < late_bias_p
    if use_rand:
        lo = max(0, start_upper - late_window_jitter)
        start_idx = int(randint(lo, start_upper + 1)) if start_upper >= 0 else 0
    else:
        start_idx = int(randint(0, max(1, start_upper + 1)))

    return start_idx, input_len, target_len


# ---------------------------------------------------------------------------
# Training dataset
# ---------------------------------------------------------------------------

class VarSeqNDVIDataset(Dataset):
    """Training dataset with variable-length windows and late-season bias.

    Returns per sample: (input_seq, target_seq, group_stats)
      - input_seq: (L_in, 2) with [ndvi, time_feature]
      - target_seq: (L_out,)
      - group_stats: (2,)
    """

    def __init__(
        self,
        ndvi_array: np.ndarray,
        group_stats_array: np.ndarray,
        min_input_len: int = 65,
        max_input_len: int = 87,
        min_target_len: int = 15,
        max_target_len: int = 30,
        k_samples_per_curve: int = 1,
        *,
        time_encoding: str = "raw",
        max_dap: int | None = None,
        late_bias_p: float = 0.5,
        late_window_jitter: int = 8,
        target_len_choices: list[int] | None = None,
        clamp01: bool = True,
    ):
        self.ndvi = ndvi_array.astype(np.float32)
        self.group_stats = group_stats_array.astype(np.float32)
        self.min_in = int(min_input_len)
        self.max_in = int(max_input_len)
        self.min_out = int(min_target_len)
        self.max_out = int(max_target_len)
        self.seq_len = int(ndvi_array.shape[1])
        self.k = int(k_samples_per_curve)
        self.encoding = str(time_encoding)
        self.max_dap = int(max_dap if max_dap is not None else self.seq_len - 1)
        self.late_bias_p = float(late_bias_p)
        self.late_jitter = int(max(0, late_window_jitter))
        self.target_choices = [int(x) for x in target_len_choices] if target_len_choices else None
        self.clamp01 = bool(clamp01)

    def __len__(self):
        return len(self.ndvi) * self.k

    def __getitem__(self, idx):
        base = idx // self.k
        full_seq = self.ndvi[base]
        gs = self.group_stats[base]

        start, in_len, out_len = _sample_window(
            self.seq_len, self.min_in, self.max_in, self.min_out, self.max_out,
            self.late_bias_p, self.late_jitter, target_len_choices=self.target_choices,
        )

        inp = full_seq[start : start + in_len]
        tgt = full_seq[start + in_len : start + in_len + out_len]

        if self.clamp01:
            inp = np.clip(inp, 0.0, 1.0)
            tgt = np.clip(tgt, 0.0, 1.0)

        time_feat = _build_time_feature(start, in_len, self.encoding, self.max_dap)
        input_seq = torch.tensor(np.stack([inp, time_feat], axis=1), dtype=torch.float32)
        target_seq = torch.tensor(tgt, dtype=torch.float32)
        group_stats = torch.tensor(gs, dtype=torch.float32)
        return input_seq, target_seq, group_stats


# ---------------------------------------------------------------------------
# Validation dataset (frozen windows)
# ---------------------------------------------------------------------------

class FrozenVarSeqNDVIDataset(Dataset):
    """Validation dataset with deterministically sampled windows (frozen at init)."""

    def __init__(
        self,
        ndvi_array: np.ndarray,
        group_stats_array: np.ndarray,
        min_input_len: int,
        max_input_len: int,
        min_target_len: int,
        max_target_len: int,
        samples_per_curve: int = 4,
        seed: int = 123,
        *,
        time_encoding: str = "progress",
        max_dap: int | None = None,
        clamp01: bool = True,
        late_bias_p: float = 0.5,
        late_window_jitter: int = 8,
    ):
        self.ndvi = ndvi_array.astype(np.float32)
        self.group_stats = group_stats_array.astype(np.float32)
        self.seq_len = int(ndvi_array.shape[1])
        self.encoding = str(time_encoding)
        self.max_dap = int(max_dap if max_dap is not None else self.seq_len - 1)
        self.clamp01 = bool(clamp01)

        rng = np.random.default_rng(seed)
        specs = []
        for i in range(len(self.ndvi)):
            made, tries = 0, 0
            while made < int(samples_per_curve) and tries < 50 * samples_per_curve:
                tries += 1
                try:
                    start, in_len, out_len = _sample_window(
                        self.seq_len, min_input_len, max_input_len,
                        min_target_len, max_target_len,
                        late_bias_p, late_window_jitter, rng=rng,
                    )
                    specs.append((i, start, in_len, out_len))
                    made += 1
                except ValueError:
                    continue
        self.specs = specs

    def __len__(self):
        return len(self.specs)

    def __getitem__(self, j):
        i, start, in_len, out_len = self.specs[j]
        row = self.ndvi[i]
        gs = self.group_stats[i]

        inp = row[start : start + in_len]
        tgt = row[start + in_len : start + in_len + out_len]

        if self.clamp01:
            inp = np.clip(inp, 0.0, 1.0).astype(np.float32)
            tgt = np.clip(tgt, 0.0, 1.0).astype(np.float32)

        time_feat = _build_time_feature(start, in_len, self.encoding, self.max_dap)
        input_seq = torch.tensor(np.stack([inp, time_feat], axis=1), dtype=torch.float32)
        target_seq = torch.tensor(tgt, dtype=torch.float32)
        group_stats = torch.tensor(gs, dtype=torch.float32)
        return input_seq, target_seq, group_stats


# ---------------------------------------------------------------------------
# Collator
# ---------------------------------------------------------------------------

class BatchCollator:
    """Pads variable-length NDVI sequences into batches.

    Returns: (padded_inputs, padded_targets, input_lengths, target_lengths, group_stats)
    """

    def __init__(self, pad_value: float = 0.0):
        self.pad_value = float(pad_value)

    def __call__(self, batch):
        inputs, targets, group_stats = zip(*batch)

        input_lengths = torch.as_tensor([len(x) for x in inputs], dtype=torch.long)
        target_lengths = torch.as_tensor([len(x) for x in targets], dtype=torch.long)

        padded_inputs = pad_sequence(inputs, batch_first=True, padding_value=self.pad_value)
        padded_targets = pad_sequence(targets, batch_first=True, padding_value=self.pad_value)
        group_stats = torch.stack(group_stats, dim=0)

        return padded_inputs, padded_targets, input_lengths, target_lengths, group_stats
