"""Dataset loading, stratified splitting and per-channel normalisation.

Every cycle sample is an n x 4 matrix [signal_value, second_value, 1, 1].
Samples are zero-padded to a fixed sequence length (divisible by patch size)
and normalised per channel using statistics computed on the training split
only. The train / validation / evaluation split is stratified per part
(8 : 1.5 : 0.5) so every part contributes samples to all three groups.
"""

import csv
import random
from pathlib import Path

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

ROOT = Path(__file__).resolve().parents[1]
INPUT_DIR = ROOT / "inputdatas"
SPLIT_PATH = ROOT / "split.csv"

PATCH_SIZE = 8
N_CHANNELS = 4
# 1552 = 194 patches of 8, covers the observed 1540..1546 cycle lengths.
SEQ_LEN = 1552

SPLIT_SEED = 42
TRAIN_RATIO = 0.80
VAL_RATIO = 0.15
EVAL_RATIO = 0.05


def discover_samples() -> list[dict]:
    samples: list[dict] = []
    for path in sorted(INPUT_DIR.glob("**/*.npy")):
        samples.append(
            {
                "path": str(path),
                "part": path.parent.name,
                "name": path.name,
            }
        )
    return samples


def ensure_split(samples: list[dict]) -> list[dict]:
    """Attach a split label to every sample, writing split.csv on first run."""
    if SPLIT_PATH.exists():
        split_map: dict[str, str] = {}
        with SPLIT_PATH.open(encoding="utf-8") as handle:
            for row in csv.DictReader(handle):
                split_map[row["path"]] = row["split"]
        for item in samples:
            item["split"] = split_map.get(item["path"], "train")
        return samples

    rng = random.Random(SPLIT_SEED)
    parts: dict[str, list[dict]] = {}
    for item in samples:
        parts.setdefault(item["part"], []).append(item)

    for part, items in parts.items():
        order = sorted(items, key=lambda item: item["name"])
        rng.shuffle(order)
        n = len(order)
        n_eval = max(1, round(n * EVAL_RATIO))
        n_val = max(1, round(n * VAL_RATIO))
        n_train = n - n_val - n_eval
        for index, item in enumerate(order):
            if index < n_train:
                item["split"] = "train"
            elif index < n_train + n_val:
                item["split"] = "val"
            else:
                item["split"] = "eval"

    with SPLIT_PATH.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=["path", "part", "name", "split"])
        writer.writeheader()
        for item in sorted(samples, key=lambda item: item["path"]):
            writer.writerow(
                {"path": item["path"], "part": item["part"], "name": item["name"], "split": item["split"]}
            )
    return samples


def _load_raw(path: str) -> np.ndarray:
    return np.load(path).astype(np.float32)


def compute_stats(samples: list[dict]) -> tuple[np.ndarray, np.ndarray]:
    """Per-channel mean/std from the training split only (std==0 -> scale 1)."""
    train = [item for item in samples if item["split"] == "train"]
    sums = np.zeros(N_CHANNELS, dtype=np.float64)
    squares = np.zeros(N_CHANNELS, dtype=np.float64)
    total = 0
    for item in tqdm(train, desc="统计归一化参数"):
        values = _load_raw(item["path"])
        sums += values.sum(axis=0)
        squares += (values.astype(np.float64) ** 2).sum(axis=0)
        total += len(values)
    mean = (sums / total).astype(np.float32)
    variance = squares / total - (mean.astype(np.float64) ** 2)
    std = np.sqrt(np.clip(variance, 0, None)).astype(np.float32)
    std[std == 0] = 1.0
    return mean, std


class CycleDataset(Dataset):
    def __init__(self, samples: list[dict], mean: np.ndarray, std: np.ndarray):
        self.mean = mean
        self.std = std
        self.items: list[tuple[dict, torch.Tensor, int]] = []
        for item in tqdm(samples, desc="加载样本到内存"):
            raw = _load_raw(item["path"])
            length = raw.shape[0]
            normalised = (raw - mean) / std
            padded = np.zeros((SEQ_LEN, N_CHANNELS), dtype=np.float32)
            padded[:length] = normalised
            self.items.append((item, torch.from_numpy(padded), length))

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

    def __getitem__(self, index: int) -> tuple[torch.Tensor, int]:
        _, tensor, length = self.items[index]
        return tensor, length


def build_split_datasets() -> tuple[list[dict], np.ndarray, np.ndarray]:
    samples = ensure_split(discover_samples())
    mean, std = compute_stats(samples)
    return samples, mean, std