from __future__ import annotations

import csv
import hashlib
import io
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Iterable

import numpy as np
import pandas as pd

from .config import registered_paths, sha256_path

KPC_M = 3.0856775814913673e19
ACC_FACTOR = 1.0e6 / KPC_M
BASE_DISK = 0.5
BASE_BULGE = 0.7


@dataclass(frozen=True)
class Realization:
    item_id: str
    mode: str
    branch: str
    realization_id: str
    repeat_id: str
    galaxy_id: str
    source_galaxy_id: str
    ladder_name: str
    values: tuple[float, ...]
    input_sha256: str
    n_source: int
    n_realization: int
    generation_rule: str
    generation_seed: int | None
    radial_coverage_min: float | None
    radial_coverage_max: float | None
    upsilon_disk: float
    upsilon_bulge: float
    variant_details: str

    def to_public_dict(self, include_values: bool = False) -> dict:
        data = asdict(self)
        data["values"] = list(self.values) if include_values else None
        return data


def read_primary_ids() -> list[str]:
    values = [x.strip() for x in registered_paths()["primary"].read_text(encoding="utf-8").splitlines() if x.strip()]
    if len(values) != 121 or len(set(values)) != 121:
        raise ValueError("Frozen primary sample must contain 121 unique galaxy IDs")
    return values


def read_ladder_csv(path: Path) -> np.ndarray:
    frame = pd.read_csv(path)
    numeric = None
    for col in frame.columns:
        vals = pd.to_numeric(frame[col], errors="coerce")
        if vals.notna().sum() >= 2:
            numeric = vals.dropna().to_numpy(float)
            break
    if numeric is None:
        raise ValueError(f"No numeric ladder column in {path}")
    values = np.unique(np.sort(numeric))
    if len(values) < 2 or not np.isfinite(values).all() or np.any(np.diff(values) <= 0):
        raise ValueError(f"Invalid canonical ladder in {path}")
    return values


def canonical_ladder(galaxy_id: str) -> np.ndarray:
    path = registered_paths()["ladders"] / f"{galaxy_id}.csv"
    if not path.is_file():
        raise FileNotFoundError(path)
    return read_ladder_csv(path)


def values_sha256(values: Iterable[float]) -> str:
    payload = "\n".join(format(float(v), ".17g") for v in values) + "\n"
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


def load_prepared_profiles() -> pd.DataFrame:
    frame = pd.read_csv(registered_paths()["profiles"])
    required = {
        "galaxy_id", "radius_kpc", "radius_normalized", "v_gas_kms",
        "v_disk_kms", "v_bulge_kms", "valid_point",
        "log10_baryonic_acceleration",
    }
    missing = sorted(required.difference(frame.columns))
    if missing:
        raise ValueError(f"Prepared profile table lacks columns: {missing}")
    if frame["valid_point"].dtype != bool:
        frame["valid_point"] = frame["valid_point"].astype(str).str.lower().eq("true")
    return frame


def valid_radial_rows(frame: pd.DataFrame, galaxy_id: str) -> pd.DataFrame:
    rows = frame.loc[(frame["galaxy_id"] == galaxy_id) & frame["valid_point"]].copy()
    rows = rows.sort_values(["radius_kpc", "log10_baryonic_acceleration"], kind="mergesort")
    if len(rows) < 2:
        raise ValueError(f"{galaxy_id}: fewer than two valid radial rows")
    return rows.reset_index(drop=True)


def values_from_rows(rows: pd.DataFrame, disk: float = BASE_DISK, bulge: float = BASE_BULGE) -> np.ndarray:
    gas = rows["v_gas_kms"].to_numpy(float)
    g2 = (
        gas * np.abs(gas)
        + float(disk) * rows["v_disk_kms"].to_numpy(float) ** 2
        + float(bulge) * rows["v_bulge_kms"].to_numpy(float) ** 2
    )
    radius = rows["radius_kpc"].to_numpy(float)
    acc = (g2 / radius) * ACC_FACTOR
    keep = np.isfinite(acc) & np.isfinite(radius) & (radius > 0) & (g2 > 0) & (acc > 0)
    values = np.unique(np.sort(np.log10(acc[keep])))
    if values.size < 2 or np.any(np.diff(values) <= 0):
        raise ValueError("Generated ladder has fewer than two positive unique values")
    return values


def evenly_spaced_indices(n: int, k: int) -> np.ndarray:
    if n < k:
        raise ValueError(f"Cannot select {k} ranks from n={n}")
    indices = np.rint(np.linspace(0, n - 1, k)).astype(int)
    if len(np.unique(indices)) != k:
        raise ValueError(f"Rounded rank rule produced duplicates for n={n}, k={k}")
    return indices


def deletion_indices(n: int) -> list[int]:
    if n <= 12:
        return list(range(n))
    return sorted({int(round(q * (n - 1))) for q in (0.0, 0.25, 0.50, 0.75, 1.0)})


def stable_rng(seed: int, key: str) -> np.random.Generator:
    key_int = int(hashlib.sha256(key.encode("utf-8")).hexdigest()[:16], 16)
    seq = np.random.SeedSequence([int(seed), key_int & 0xFFFFFFFF, key_int >> 32])
    return np.random.Generator(np.random.PCG64(seq))


def random_stratified_indices(n: int, seed: int, galaxy_id: str) -> np.ndarray:
    if n < 10:
        raise ValueError(f"Random n=10 thinning requires n>=10, got {n}")
    if n == 10:
        return np.arange(10, dtype=int)
    interior = np.arange(1, n - 1, dtype=int)
    strata = [part for part in np.array_split(interior, 8) if len(part)]
    if len(strata) != 8:
        raise ValueError(f"Expected eight non-empty interior strata for n={n}")
    rng = stable_rng(seed, galaxy_id)
    selected = np.array(sorted([0] + [int(rng.choice(part)) for part in strata] + [n - 1]), dtype=int)
    if len(np.unique(selected)) != 10:
        raise ValueError(f"Random selection not ten unique ranks for {galaxy_id}")
    return selected


def gap_order_null(values: np.ndarray, seed: int, galaxy_id: str) -> np.ndarray:
    gaps = np.diff(values)
    if np.any(gaps <= 0):
        raise ValueError("Gap-order null requires strictly increasing ladder")
    permuted = stable_rng(seed, galaxy_id).permutation(gaps)
    return np.concatenate(([values[0]], values[0] + np.cumsum(permuted)))


def uniform_range_null(values: np.ndarray, seed: int, galaxy_id: str) -> np.ndarray:
    if len(values) == 2:
        return values.copy()
    interior = np.sort(stable_rng(seed, galaxy_id).uniform(values[0], values[-1], len(values) - 2))
    return np.concatenate(([values[0]], interior, [values[-1]]))


def _safe_name(text: str) -> str:
    return text.replace("-", "_").replace(".", "p").replace("/", "_")


def _make(
    *, mode: str, branch: str, realization_id: str, repeat_id: str, galaxy_id: str,
    values: np.ndarray, n_source: int, generation_rule: str, generation_seed: int | None = None,
    radial_min: float | None = None, radial_max: float | None = None,
    disk: float = BASE_DISK, bulge: float = BASE_BULGE, details: str = "",
) -> Realization:
    values = np.asarray(values, dtype=float)
    if len(values) < 2 or not np.isfinite(values).all() or np.any(np.diff(values) <= 0):
        raise ValueError(f"Invalid generated ladder: {galaxy_id} {realization_id}")
    ladder_name = galaxy_id if mode == "REPEAT" else f"GF1_{galaxy_id}_{_safe_name(realization_id)}"
    item_id = f"{mode}:{repeat_id or '-'}:{galaxy_id}:{realization_id}"
    return Realization(
        item_id=item_id, mode=mode, branch=branch, realization_id=realization_id,
        repeat_id=repeat_id, galaxy_id=galaxy_id, source_galaxy_id=galaxy_id,
        ladder_name=ladder_name, values=tuple(float(v) for v in values),
        input_sha256=values_sha256(values), n_source=int(n_source), n_realization=int(len(values)),
        generation_rule=generation_rule, generation_seed=generation_seed,
        radial_coverage_min=radial_min, radial_coverage_max=radial_max,
        upsilon_disk=float(disk), upsilon_bulge=float(bulge), variant_details=details,
    )


def build_realizations(mode: str, repeat_id: str = "", seeds: list[int] | None = None) -> list[Realization]:
    mode = mode.upper()
    primary = read_primary_ids()
    if mode == "REPEAT":
        if repeat_id not in {"R2", "R3", "R4", "R5"}:
            raise ValueError("REPEAT requires repeat_id R2, R3, R4, or R5")
        out = []
        for gid in primary:
            vals = canonical_ladder(gid)
            out.append(_make(
                mode=mode, branch="R", realization_id=repeat_id, repeat_id=repeat_id,
                galaxy_id=gid, values=vals, n_source=len(vals),
                generation_rule="unchanged frozen primary ladder",
                details="independent whole-chamber repetition using native browser RNG",
            ))
        return out

    frame = load_prepared_profiles()
    out: list[Realization] = []
    for gid in primary:
        rows = valid_radial_rows(frame, gid)
        n = len(rows)
        baseline = values_from_rows(rows)
        if mode == "LEAVE_ONE_RADIUS_OUT":
            for idx in deletion_indices(n):
                kept = rows.drop(index=idx).reset_index(drop=True)
                out.append(_make(
                    mode=mode, branch="L", realization_id=f"L_DEL_R{idx+1:03d}", repeat_id="",
                    galaxy_id=gid, values=values_from_rows(kept), n_source=n,
                    generation_rule="registered single-radius deletion",
                    radial_min=float(kept.radius_normalized.min()), radial_max=float(kept.radius_normalized.max()),
                    details=f"deleted 1-based radial rank {idx+1} of {n}",
                ))
        elif mode == "THINNING_N10_DETERMINISTIC":
            idx = evenly_spaced_indices(n, 10)
            kept = rows.iloc[idx].reset_index(drop=True)
            out.append(_make(
                mode=mode, branch="T", realization_id="T_DET10", repeat_id="", galaxy_id=gid,
                values=values_from_rows(kept), n_source=n,
                generation_rule="registered deterministic n=10 radial-rank thinning",
                radial_min=float(kept.radius_normalized.min()), radial_max=float(kept.radius_normalized.max()),
                details="1+round(j(n-1)/9), j=0..9",
            ))
        elif mode == "THINNING_N10_RANDOM":
            for seed in (seeds or [91001, 91002, 91003]):
                idx = random_stratified_indices(n, seed, gid)
                kept = rows.iloc[idx].reset_index(drop=True)
                out.append(_make(
                    mode=mode, branch="T_RANDOM", realization_id=f"T_RAND10_S{seed}", repeat_id="", galaxy_id=gid,
                    values=values_from_rows(kept), n_source=n,
                    generation_rule="locked random stratified n=10 sensitivity thinning", generation_seed=seed,
                    radial_min=float(kept.radius_normalized.min()), radial_max=float(kept.radius_normalized.max()),
                    details="endpoints plus one draw from each of eight contiguous interior strata",
                ))
        elif mode == "COVERAGE":
            full_rows = rows.iloc[evenly_spaced_indices(n, 8)].reset_index(drop=True)
            out.append(_make(
                mode=mode, branch="C", realization_id="C_FULL8", repeat_id="", galaxy_id=gid,
                values=values_from_rows(full_rows), n_source=n,
                generation_rule="registered full-coverage n=8 radial-rank selection",
                radial_min=float(full_rows.radius_normalized.min()), radial_max=float(full_rows.radius_normalized.max()),
                details="eight evenly spaced ranks over full observed range",
            ))
            inner = rows.loc[rows.radius_normalized <= 0.85].reset_index(drop=True)
            if len(inner) < 8:
                raise ValueError(f"{gid}: fewer than eight valid points within R/Rmax<=0.85")
            inner_rows = inner.iloc[evenly_spaced_indices(len(inner), 8)].reset_index(drop=True)
            out.append(_make(
                mode=mode, branch="C", realization_id="C_INNER85_8", repeat_id="", galaxy_id=gid,
                values=values_from_rows(inner_rows), n_source=n,
                generation_rule="registered inner-85%-coverage n=8 radial-rank selection",
                radial_min=float(inner_rows.radius_normalized.min()), radial_max=float(inner_rows.radius_normalized.max()),
                details="eight evenly spaced ranks among R/Rmax<=0.85",
            ))
        elif mode == "ML_SENSITIVITY":
            for rid, disk, bulge in [
                ("M_D_LOW", 0.4, 0.7), ("M_D_HIGH", 0.6, 0.7),
            ]:
                out.append(_make(
                    mode=mode, branch="M", realization_id=rid, repeat_id="", galaxy_id=gid,
                    values=values_from_rows(rows, disk, bulge), n_source=n,
                    generation_rule="registered one-factor baryonic preparation sensitivity",
                    radial_min=float(rows.radius_normalized.min()), radial_max=float(rows.radius_normalized.max()),
                    disk=disk, bulge=bulge, details=f"Upsilon_disk={disk}; Upsilon_bulge={bulge}",
                ))
            has_bulge = bool(np.any(np.abs(rows.v_bulge_kms.to_numpy(float)) > 0))
            if has_bulge:
                for rid, disk, bulge in [
                    ("M_B_LOW", 0.5, 0.6), ("M_B_HIGH", 0.5, 0.8),
                ]:
                    out.append(_make(
                        mode=mode, branch="M", realization_id=rid, repeat_id="", galaxy_id=gid,
                        values=values_from_rows(rows, disk, bulge), n_source=n,
                        generation_rule="registered one-factor baryonic preparation sensitivity",
                        radial_min=float(rows.radius_normalized.min()), radial_max=float(rows.radius_normalized.max()),
                        disk=disk, bulge=bulge, details=f"Upsilon_disk={disk}; Upsilon_bulge={bulge}",
                    ))
        elif mode == "NULL_GAP_ORDER":
            for seed in (seeds or [92001, 92002, 92003, 92004, 92005]):
                out.append(_make(
                    mode=mode, branch="N_GAP", realization_id=f"N_GAP_S{seed}", repeat_id="", galaxy_id=gid,
                    values=gap_order_null(baseline, seed, gid), n_source=n,
                    generation_rule="registered gap-order permutation null", generation_seed=seed,
                    radial_min=float(rows.radius_normalized.min()), radial_max=float(rows.radius_normalized.max()),
                    details="preserves n, min, max, total range, and exact adjacent-gap multiset",
                ))
        elif mode == "NULL_UNIFORM_RANGE":
            for seed in (seeds or [93001, 93002, 93003, 93004, 93005]):
                out.append(_make(
                    mode=mode, branch="N_UNIFORM", realization_id=f"N_UNIFORM_S{seed}", repeat_id="", galaxy_id=gid,
                    values=uniform_range_null(baseline, seed, gid), n_source=n,
                    generation_rule="registered uniform-range null", generation_seed=seed,
                    radial_min=float(rows.radius_normalized.min()), radial_max=float(rows.radius_normalized.max()),
                    details="preserves n, minimum, and maximum; uniform sorted interior",
                ))
        else:
            raise ValueError(f"Unknown mode: {mode}")
    expected = {
        "LEAVE_ONE_RADIUS_OUT": 782,
        "THINNING_N10_DETERMINISTIC": 121,
        "THINNING_N10_RANDOM": 363,
        "COVERAGE": 242,
        "ML_SENSITIVITY": 302,
        "NULL_GAP_ORDER": 605,
        "NULL_UNIFORM_RANGE": 605,
    }.get(mode)
    if expected is not None and len(out) != expected:
        raise AssertionError(f"{mode} generated {len(out)} realizations; expected {expected}")
    return out


def registered_asset_report() -> dict:
    paths = registered_paths()
    primary = read_primary_ids()
    manifest = pd.read_csv(paths["ladder_manifest"])
    missing = [gid for gid in primary if not (paths["ladders"] / f"{gid}.csv").is_file()]
    ladder_hashes_ok = 0
    manifest_map = manifest.set_index("galaxy_id")
    for gid in primary:
        if gid in missing or gid not in manifest_map.index:
            continue
        actual = sha256_path(paths["ladders"] / f"{gid}.csv")
        expected = str(manifest_map.loc[gid, "struc_i_file_sha256"])
        ladder_hashes_ok += int(actual == expected)
    return {
        "primary_count": len(primary),
        "missing_ladders": missing,
        "ladder_file_hashes_matching_manifest": ladder_hashes_ok,
        "primary_sha256": sha256_path(paths["primary"]),
        "prepared_profiles_sha256": sha256_path(paths["profiles"]),
        "r1_profiles_sha256": sha256_path(paths["r1_profiles"]),
    }
