"""Pair manifests: which trajectories of ``trajectories_fine.npz`` a training run reads, and at which snapshots.

A manifest directory holds ``slice_<name>.json`` and ``master_order.npz`` (the same rows; the loaders re-hash it).
Rows are ``[group_row, member_col, gid]`` into ``vecs_train`` (segment pairs) or ``[block, group_row, member_col, gid]``
with block 0 = ``vecs_train`` (endpoint pairs). ``bifurcation.train`` writes the manifest of a run before training.
"""
from __future__ import annotations

import hashlib
import json
import os
from datetime import datetime, timezone
from pathlib import Path

import numpy as np

from .fine import ENDPOINT_SCHEMA as _fine_endpoint_schema
from .fine import FINE_FILE, _npz_member_memmap, _sha256_file

SEGMENT_SCHEMA = "segment_pairs_v1"
ENDPOINT_SCHEMA = _fine_endpoint_schema


def block_rows_segment(gids_train: np.ndarray, n_ic: int, members: int) -> np.ndarray:
    """Training conditions 0..n_ic-1 x realizations 0..members-1; rows [group_row, member_col, gid]."""
    rows = np.empty((n_ic * members, 3), dtype=np.int64)
    i = 0
    for gi in range(n_ic):
        for mj in range(members):
            rows[i] = (gi, mj, int(gids_train[gi]))
            i += 1
    return rows


def block_rows_endpoint(gids_train: np.ndarray, n_ic: int, members: int) -> np.ndarray:
    """The same enumeration for endpoint pairs; rows [block, group_row, member_col, gid] with block 0 = vecs_train."""
    rows = np.empty((n_ic * members, 4), dtype=np.int64)
    i = 0
    for gi in range(n_ic):
        for mj in range(members):
            rows[i] = (0, gi, mj, int(gids_train[gi]))
            i += 1
    return rows


def segment_geometry(T: float, S: int, t_values: np.ndarray) -> dict:
    """Uniform segmentation of [0, T] on the stored snapshot grid: n_T = index of T (an exact snapshot), S must
    divide n_T; boundary_idx = [0, stride, ..., n_T]; the segment times are read from the stored grid."""
    t_values = np.asarray(t_values, dtype=np.float64)
    hits = np.nonzero(np.isclose(t_values, T, atol=1e-9))[0]
    if len(hits) != 1:
        raise ValueError(f"T={T} is not a stored snapshot time (grid {t_values[:3]}...)")
    n_T = int(hits[0])
    if n_T % S != 0:
        raise ValueError(f"S={S} does not divide n_T={n_T} (T={T}); admissible S: "
                         f"{[s for s in range(1, n_T + 1) if n_T % s == 0]}")
    stride = n_T // S
    boundary_idx = [q * stride for q in range(S + 1)]
    return {"T": float(T), "S": int(S), "n_T": n_T, "stride": stride, "delta": float(T) / S,
            "boundary_idx": boundary_idx, "t_values": [float(t_values[b]) for b in boundary_idx]}


def _now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


def _write_master(slice_dir: Path, rows: np.ndarray, note: str) -> str:
    slice_dir.mkdir(parents=True, exist_ok=True)
    path = slice_dir / "master_order.npz"
    if rows.shape[1] == 3:
        payload = dict(group_row=rows[:, 0].astype(np.int32), member_col=rows[:, 1].astype(np.int32),
                       gid=rows[:, 2].astype(np.int32))
    else:
        payload = dict(block=rows[:, 0].astype(np.int32), group_row=rows[:, 1].astype(np.int32),
                       member_col=rows[:, 2].astype(np.int32), gid=rows[:, 3].astype(np.int32))
    payload["algorithm"] = np.bytes_(note.encode())
    np.savez(path, **payload)
    return _sha256_file(str(path))


def dataset_provenance(packaged_dir: str | os.PathLike) -> dict:
    packaged_dir = Path(packaged_dir)
    return {"dataset_path": str(packaged_dir / FINE_FILE),
            "split_sha256": _sha256_file(str(packaged_dir / "split.json")),
            "normalization_sha256": _sha256_file(str(packaged_dir / "normalization.npz"))}


def write_segment_manifest(slice_dir: Path, rows3: np.ndarray, geometry: dict, name: str,
                           packaged_dir: str | os.PathLike, extra: dict | None = None) -> Path:
    S = int(geometry["S"])
    master_sha = _write_master(slice_dir, rows3, f"complete block ({name}): canonical enumeration")
    man = {"schema": SEGMENT_SCHEMA, "arm": name, "cell": name, **geometry,
           "N_traj": int(rows3.shape[0]), "B": int(rows3.shape[0]) * S,
           "trajectories": rows3.tolist(), "master_order_sha256": master_sha,
           **dataset_provenance(packaged_dir), "built_at": _now(), **(extra or {})}
    out = slice_dir / f"slice_{name}.json"
    out.write_text(json.dumps(man))
    return out


def write_endpoint_manifest(slice_dir: Path, rows4: np.ndarray, idx_T: int, t_values: np.ndarray, name: str,
                            packaged_dir: str | os.PathLike, extra: dict | None = None) -> Path:
    master_sha = _write_master(slice_dir, rows4, f"complete block ({name}): canonical enumeration, block 0 = vecs_train")
    man = {"schema": ENDPOINT_SCHEMA, "arm": name, "S": 1,
           "T": float(t_values[idx_T]), "n_T": int(idx_T),
           "N_traj": int(rows4.shape[0]), "B": int(rows4.shape[0]),
           "boundary_idx": [0, int(idx_T)], "t_values": [0.0, float(t_values[idx_T])],
           "blocks": {"0": "vecs_train"},
           "trajectories": rows4.tolist(), "master_order_sha256": master_sha,
           **dataset_provenance(packaged_dir), "built_at": _now(), **(extra or {})}
    out = slice_dir / f"slice_{name}.json"
    out.write_text(json.dumps(man))
    return out


def fine_grid(packaged_dir: str | os.PathLike) -> dict:
    """Key arrays of the fine file (memmapped; cheap)."""
    fine = str(Path(packaged_dir) / FINE_FILE)
    out = {"t_values": np.asarray(_npz_member_memmap(fine, "t_values")),
           "gids_train": np.asarray(_npz_member_memmap(fine, "gids_train")),
           "gids_eval": np.asarray(_npz_member_memmap(fine, "gids_eval"))}
    for k in ("vecs_train", "vecs_eval"):
        out[k + "_shape"] = tuple(_npz_member_memmap(fine, k).shape)
    return out