"""Restricted readers and split preparation for the official Q2 inputs."""
from __future__ import annotations

import pickle
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np

ROOT = Path(__file__).resolve().parents[2]
ALIGNED_PATH = ROOT / "E题数据" / "附件2-数据集特征文件" / "aligned_50.pkl"
ATTACHMENT3_ALIGNED = ROOT / "E题数据" / "附件3-模态缺失特征样本" / "对齐版本"
MODALITIES = ("text", "audio", "vision")
EXPECTED_DIMS = {"text": 768, "audio": 74, "vision": 35}


class RestrictedUnpickler(pickle.Unpickler):
    """Allow only primitive containers and NumPy reconstruction primitives."""

    _allowed = {
        ("builtins", name): getattr(__import__("builtins"), name)
        for name in ("set", "frozenset", "slice", "complex", "bytearray")
    }
    _allowed.update({
        ("collections", "OrderedDict"): __import__("collections").OrderedDict,
        ("numpy", "ndarray"): np.ndarray,
        ("numpy", "dtype"): np.dtype,
        ("numpy", "asarray"): np.asarray,
        ("numpy.core.multiarray", "_reconstruct"): np.core.multiarray._reconstruct,
        ("numpy.core.multiarray", "scalar"): np.core.multiarray.scalar,
        ("numpy._core.multiarray", "_reconstruct"): np.core.multiarray._reconstruct,
        ("numpy._core.multiarray", "scalar"): np.core.multiarray.scalar,
    })

    def find_class(self, module: str, name: str) -> Any:
        try:
            return self._allowed[(module, name)]
        except KeyError as exc:
            raise pickle.UnpicklingError(f"blocked pickle global: {module}.{name}") from exc


def restricted_load(path: Path) -> Any:
    with path.open("rb") as stream:
        return RestrictedUnpickler(stream).load()


def _decode(value: Any) -> str:
    if isinstance(value, bytes):
        return value.decode("utf-8", errors="replace")
    if isinstance(value, np.bytes_):
        return bytes(value).decode("utf-8", errors="replace")
    if isinstance(value, np.ndarray) and value.shape == ():
        return _decode(value.item())
    return str(value)


def _one_dim(value: Any, dtype: Any | None = None) -> np.ndarray:
    out = np.asarray(value)
    if out.ndim > 1 and out.shape[-1] == 1:
        out = out.reshape(-1)
    elif out.ndim > 1 and out.shape[0] == 1:
        out = out.reshape(-1)
    else:
        out = out.reshape(-1)
    return out.astype(dtype) if dtype is not None else out


@dataclass
class SplitData:
    name: str
    x: dict[str, np.ndarray]
    mask: np.ndarray
    class_y: np.ndarray | None
    regression_y: np.ndarray | None
    ids: list[str]
    groups: np.ndarray
    alignment_audit: dict[str, Any] | None = None

    @property
    def n(self) -> int:
        return len(self.ids)


def _extract_split(
    name: str, obj: dict[str, Any], with_labels: bool,
    mask_override: np.ndarray | None = None,
    alignment_audit: dict[str, Any] | None = None,
) -> SplitData:
    raw: dict[str, np.ndarray] = {}
    masks = []
    for modality in MODALITIES:
        arr = np.asarray(obj[modality])
        if arr.ndim != 3 or arr.shape[1] != 50 or arr.shape[2] != EXPECTED_DIMS[modality]:
            raise ValueError(f"{name}.{modality}: unexpected feature shape {arr.shape}")
        arr = arr.astype(np.float32)
        if not np.isfinite(arr).all():
            raise ValueError(f"{name}.{modality}: non-finite feature values; refusing to reinterpret them as missing")
        # The dataset documentation defines all-zero aligned rows as missing.
        observed = np.any(arr != 0.0, axis=-1)
        raw[modality] = arr
        masks.append(observed)
    mask = np.stack(masks, axis=-1)
    if mask_override is not None:
        override = np.asarray(mask_override, bool)
        if override.shape != mask.shape:
            raise ValueError(f"{name}: projected mask shape {override.shape} differs from {mask.shape}")
        mask = override
    ids = [_decode(v) for v in _one_dim(obj["id"])]
    if len(ids) != len(mask):
        raise ValueError(f"{name}: id count differs from feature count")
    if len(set(ids)) != len(ids):
        raise ValueError(f"{name}: duplicate video$_$clip primary keys")
    malformed = [sample_id for sample_id in ids if "$_$" not in sample_id or not all(sample_id.split("$_$", 1))]
    if malformed:
        raise ValueError(f"{name}: malformed video$_$clip keys: {malformed[:5]}")
    groups = np.asarray([sample_group(v) for v in ids], dtype=str)
    if with_labels:
        class_y = _one_dim(obj["classification_labels"], np.int64)
        regression_y = _one_dim(obj["regression_labels"], np.float32)
        if len(class_y) != len(ids) or len(regression_y) != len(ids):
            raise ValueError(f"{name}: label count differs from feature count")
        if not np.isfinite(regression_y).all() or np.any(np.abs(regression_y) > 3.0):
            raise ValueError(f"{name}: regression labels must be finite and within [-3,3]")
        if not np.isin(class_y, [0, 1, 2]).all():
            raise ValueError(f"{name}: expected class labels in 0,1,2")
        expected_class = np.where(regression_y < 0.0, 0, np.where(regression_y == 0.0, 1, 2))
        mismatch = np.flatnonzero(class_y != expected_class)
        if len(mismatch):
            examples = [(ids[int(i)], int(class_y[i]), float(regression_y[i])) for i in mismatch[:5]]
            raise ValueError(f"{name}: polarity/regression label mismatch (sample, class, score): {examples}")
    else:
        class_y = regression_y = None
    return SplitData(name, raw, mask, class_y, regression_y, ids, groups, alignment_audit)


def sample_group(sample_id: str) -> str:
    """Official ids are video$_$clip; group on the source video only."""
    return sample_id.split("$_$", 1)[0]


def load_official_splits(path: Path = ALIGNED_PATH, *, version: str = "aligned_50") -> dict[str, SplitData]:
    if version not in {"aligned_50", "unaligned_50"}:
        raise ValueError(f"unsupported feature version: {version}")
    obj = restricted_load(path)
    required = {"train", "valid", "test"}
    if not isinstance(obj, dict) or not required.issubset(obj):
        raise ValueError(f"{path.name} must contain train, valid, and test dictionaries")
    if version == "unaligned_50":
        repo_dir = Path(__file__).resolve().parents[2]
        if str(repo_dir) not in sys.path:
            sys.path.insert(0, str(repo_dir))
        from final.adapter import adapt_official_split

        splits = {}
        for name in ("train", "valid", "test"):
            projected, mask, audit = adapt_official_split(obj[name])
            fields = {**obj[name], **projected}
            splits[name] = _extract_split(name, fields, with_labels=True,
                                          mask_override=mask, alignment_audit=audit)
    else:
        splits = {name: _extract_split(name, obj[name], with_labels=True) for name in ("train", "valid", "test")}
    del obj
    return splits


def load_attachment3_case(path: Path) -> dict[str, np.ndarray]:
    obj = restricted_load(path)
    case = obj.get("test", obj)
    text_bert = np.asarray(case["text_bert"])
    audio = np.asarray(case["audio"])
    vision = np.asarray(case["vision"])
    if text_bert.ndim == 3 and text_bert.shape[0] == 1:
        text_bert = text_bert[0]
    if text_bert.shape != (3, 50):
        raise ValueError(f"{path.name}: expected text_bert (1,3,50), got {np.asarray(case['text_bert']).shape}")
    result = {"input_ids": text_bert[0].astype(np.int64), "attention_mask": text_bert[1].astype(bool), "token_type_ids": text_bert[2].astype(np.int64)}
    for name, arr, dim in (("audio", audio, 74), ("vision", vision, 35)):
        if arr.ndim == 3 and arr.shape[0] == 1:
            arr = arr[0]
        if arr.shape != (50, dim):
            raise ValueError(f"{path.name}: expected {name} (1,50,{dim}), got {np.asarray(case[name]).shape}")
        arr = arr.astype(np.float32)
        if not np.isfinite(arr).all():
            raise ValueError(f"{path.name}: {name} contains non-finite features")
        result[name] = arr
    return result


def fit_preprocessor(train: SplitData) -> dict[str, dict[str, np.ndarray]]:
    """Fit per-dimension mean/std on observed training rows only."""
    fitted: dict[str, dict[str, np.ndarray]] = {}
    for j, name in enumerate(MODALITIES):
        rows = train.x[name][train.mask[:, :, j]]
        mean = rows.mean(axis=0, dtype=np.float64).astype(np.float32)
        std = rows.std(axis=0, dtype=np.float64).astype(np.float32)
        std[std < 1e-5] = 1.0
        fitted[name] = {"mean": mean, "std": std}
    return fitted


def transform_split(split: SplitData, fitted: dict[str, dict[str, np.ndarray]]) -> dict[str, np.ndarray]:
    output = {}
    for j, name in enumerate(MODALITIES):
        arr = (split.x[name] - fitted[name]["mean"]) / fitted[name]["std"]
        arr = np.clip(arr, -10.0, 10.0)
        arr[~split.mask[:, :, j]] = 0.0
        output[name] = arr.astype(np.float32)
    return output