from __future__ import annotations

import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler

from .config import CHAMBER_VERSION, active_roots, load_config, registered_paths, sha256_path
from .corpus import canonical_ladder, load_prepared_profiles, read_primary_ids, valid_radial_rows
from .evidence import branch_bundle, gate_b_evidence_report, repeat_bundle, write_gate_b_provenance_amendment

PRIMARY_MODES = [
    "LEAVE_ONE_RADIUS_OUT",
    "THINNING_N10_DETERMINISTIC",
    "COVERAGE",
    "ML_SENSITIVITY",
]
MODE_TO_BRANCH = {
    "LEAVE_ONE_RADIUS_OUT": "L",
    "THINNING_N10_DETERMINISTIC": "T",
    "COVERAGE": "C",
    "ML_SENSITIVITY": "M",
}


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def _load_r1() -> pd.DataFrame:
    frame = pd.read_csv(registered_paths()["r1_profiles"])
    return frame.rename(columns={"galaxy_id": "source_galaxy_id"})[
        ["source_galaxy_id", "kappa", "rho", "Ak", "nu_V"]
    ]


def _repeat_frame(repeat: str) -> pd.DataFrame:
    payload = repeat_bundle(repeat).result_json()
    rows: list[dict[str, Any]] = []
    for entry in payload.get("profiles", []):
        gid = entry["name"]
        for point in entry["curve"]:
            rows.append({
                "source_galaxy_id": gid,
                "kappa": float(point["kappa"]),
                "rho": float(point["rho"]),
                "Ak": float(point["Ak"]),
                "nu_V": int(point.get("nu", point.get("nu_V"))),
            })
    return pd.DataFrame(rows)


def _validate_repeat(frame: pd.DataFrame, ids: list[str], label: str) -> pd.DataFrame:
    if len(frame) != len(ids) * 40:
        raise ValueError(f"{label}: expected {len(ids)*40} curve rows, found {len(frame)}")
    counts = frame.groupby("source_galaxy_id").size()
    if set(counts.index) != set(ids) or not (counts == 40).all():
        raise ValueError(f"{label}: incomplete or unknown galaxy profiles")
    return frame.sort_values(["source_galaxy_id", "kappa"], kind="mergesort").reset_index(drop=True)


def _matrix(frame: pd.DataFrame, ids: list[str], field: str) -> np.ndarray:
    return np.vstack([
        frame.loc[frame.source_galaxy_id == gid].sort_values("kappa")[field].to_numpy(float)
        for gid in ids
    ])


def load_reference() -> tuple[list[str], dict[str, np.ndarray], dict[str, str]]:
    ids = read_primary_ids()
    frames = {"R1": _validate_repeat(_load_r1(), ids, "R1")}
    sources = {"R1": str(registered_paths()["r1_profiles"].resolve())}
    hashes = {"R1": sha256_path(registered_paths()["r1_profiles"])}
    for repeat in ("R2", "R3", "R4", "R5"):
        bundle = repeat_bundle(repeat)
        frames[repeat] = _validate_repeat(_repeat_frame(repeat), ids, repeat)
        sources[repeat] = f"{bundle.source_path.resolve()}::{bundle.result_name}"
        hashes[repeat] = bundle.result_sha256
    reference = {
        field: np.mean(np.stack([_matrix(frames[r], ids, field) for r in frames]), axis=0)
        for field in ("rho", "Ak", "nu_V")
    }
    reference["kappa"] = np.sort(frames["R1"].kappa.unique())
    return ids, reference, {"source_files": sources, "source_sha256": hashes}


def _parse_number(value: Any, kind: type, default: Any = None) -> Any:
    if value in (None, "", "None"):
        return default
    return kind(value)


def _branch_frame(mode: str) -> tuple[pd.DataFrame, list[dict[str, str]], dict[str, Any]]:
    bundle = branch_bundle(mode)
    payload = bundle.result_json()
    manifest_rows = bundle.run_manifest_rows()
    meta = {row.get("ladder_name", ""): row for row in manifest_rows}
    rows: list[dict[str, Any]] = []
    for entry in payload.get("profiles", []):
        name = entry["name"]
        item = meta.get(name)
        if item is None:
            raise ValueError(f"{mode}: no run-manifest row for {name}")
        source = item.get("source_galaxy_id") or item.get("galaxy_id")
        realization = item.get("realization_id") or name
        for point in entry["curve"]:
            rows.append({
                "mode": mode,
                "branch": MODE_TO_BRANCH[mode],
                "ladder_name": name,
                "source_galaxy_id": source,
                "realization_id": realization,
                "item_id": item.get("item_id", ""),
                "n_source": _parse_number(item.get("n_source"), int, 0),
                "n_realization": _parse_number(item.get("n_realization") or item.get("n"), int, 0),
                "generation_seed": _parse_number(item.get("generation_seed"), int, None),
                "radial_coverage_min": _parse_number(item.get("radial_coverage_min"), float, None),
                "radial_coverage_max": _parse_number(item.get("radial_coverage_max"), float, None),
                "upsilon_disk": _parse_number(item.get("upsilon_disk"), float, 0.5),
                "upsilon_bulge": _parse_number(item.get("upsilon_bulge"), float, 0.7),
                "kappa": float(point["kappa"]),
                "rho": float(point["rho"]),
                "Ak": float(point["Ak"]),
                "nu_V": int(point.get("nu", point.get("nu_V"))),
            })
    frame = pd.DataFrame(rows)
    source = {
        "source": str(bundle.source_path.resolve()),
        "result_member": bundle.result_name,
        "result_sha256": bundle.result_sha256,
        "profile_count": len(payload.get("profiles", [])),
    }
    return frame, manifest_rows, source


def gate_c_readiness() -> dict[str, Any]:
    gate_b = gate_b_evidence_report()
    expected = load_config()["gate_c_expected_counts"]
    branches: dict[str, Any] = {}
    ready = bool(gate_b.get("gate_b_pass"))
    for mode in PRIMARY_MODES:
        try:
            frame, manifest, source = _branch_frame(mode)
            profile_counts = frame.groupby("ladder_name").size() if not frame.empty else pd.Series(dtype=int)
            count = len(profile_counts)
            valid = count == expected[mode] and bool((profile_counts == 40).all()) and len(manifest) == expected[mode]
            branches[mode] = {**source, "expected": expected[mode], "found": count, "valid": bool(valid)}
            ready &= bool(valid)
        except Exception as exc:
            branches[mode] = {"expected": expected[mode], "found": 0, "valid": False, "error": str(exc)}
            ready = False
    return {"ready": bool(ready), "gate_b": gate_b, "branches": branches}


def _ordinary_control_space(ids: list[str]) -> tuple[np.ndarray, dict[str, list[int]], list[str]]:
    prepared = load_prepared_profiles()
    rows: list[list[float]] = []
    names = [
        "log_ladder_size", "profile_range", "median_gap", "gap_iqr", "gap_variance",
        "log_radial_span", "median_normalized_radial_gap", "normalized_radial_gap_variance",
    ]
    for gid in ids:
        ladder = canonical_ladder(gid)
        gaps = np.diff(ladder)
        radial = valid_radial_rows(prepared, gid)
        span = float(radial.radius_kpc.max() - radial.radius_kpc.min())
        nr = np.sort(radial.radius_normalized.to_numpy(float))
        rgaps = np.diff(nr)
        rows.append([
            np.log(len(ladder)), float(np.ptp(ladder)), float(np.median(gaps)),
            float(np.quantile(gaps, 0.75) - np.quantile(gaps, 0.25)), float(np.var(gaps)),
            np.log(max(span, 1e-12)), float(np.median(rgaps)), float(np.var(rgaps)),
        ])
    matrix = np.asarray(rows, dtype=float)
    z = StandardScaler().fit_transform(matrix)
    impostors: dict[str, list[int]] = {}
    for i, gid in enumerate(ids):
        dist = np.sqrt(np.sum((z - z[i]) ** 2, axis=1))
        order = [int(j) for j in np.argsort(dist, kind="mergesort") if int(j) != i]
        impostors[gid] = order[:10]
    return matrix, impostors, names


def _reference_neighbours(reference_rho: np.ndarray) -> list[set[int]]:
    scaler = StandardScaler().fit(reference_rho)
    z = scaler.transform(reference_rho)
    out: list[set[int]] = []
    for i in range(len(z)):
        dist = np.sqrt(np.mean((z - z[i]) ** 2, axis=1))
        order = [int(j) for j in np.argsort(dist, kind="mergesort") if int(j) != i]
        out.append(set(order[:5]))
    return out


def _distance_products(
    reference: dict[str, np.ndarray], query: dict[str, np.ndarray],
    ref_n: np.ndarray, query_n: np.ndarray,
) -> dict[str, np.ndarray]:
    rho_scaler = StandardScaler().fit(reference["rho"])
    ref_rho = rho_scaler.transform(reference["rho"])
    qry_rho = rho_scaler.transform(query["rho"])
    rho = np.sqrt(np.mean((qry_rho[:, None, :] - ref_rho[None, :, :]) ** 2, axis=2))

    ref_denom = np.maximum(np.ceil((ref_n - 1.0) / 2.0), 1.0)[:, None]
    qry_denom = np.maximum(np.ceil((query_n - 1.0) / 2.0), 1.0)[:, None]
    ref_channels = [reference["rho"], 1.0 - reference["Ak"], reference["nu_V"] / ref_denom]
    qry_channels = [query["rho"], 1.0 - query["Ak"], query["nu_V"] / qry_denom]
    channel_mse = []
    for ref_ch, qry_ch in zip(ref_channels, qry_channels):
        scaler = StandardScaler().fit(ref_ch)
        rz = scaler.transform(ref_ch)
        qz = scaler.transform(qry_ch)
        channel_mse.append(np.mean((qz[:, None, :] - rz[None, :, :]) ** 2, axis=2))
    combined = np.sqrt(np.mean(np.stack(channel_mse, axis=0), axis=0))

    pca = PCA(n_components=3).fit(ref_rho)
    ref_pc = pca.transform(ref_rho)
    qry_pc = pca.transform(qry_rho)
    pca3 = np.sqrt(np.sum((qry_pc[:, None, :] - ref_pc[None, :, :]) ** 2, axis=2))
    return {"rho_rms": rho, "combined_rms": combined, "rho_pca3": pca3}


def _query_matrices(frame: pd.DataFrame) -> tuple[pd.DataFrame, dict[str, np.ndarray], np.ndarray]:
    metadata_cols = [
        "mode", "branch", "ladder_name", "source_galaxy_id", "realization_id", "item_id",
        "n_source", "n_realization", "generation_seed", "radial_coverage_min", "radial_coverage_max",
        "upsilon_disk", "upsilon_bulge",
    ]
    metadata = frame[metadata_cols].drop_duplicates("ladder_name").sort_values("ladder_name", kind="mergesort").reset_index(drop=True)
    matrices = {}
    for field in ("rho", "Ak", "nu_V"):
        matrices[field] = np.vstack([
            frame.loc[frame.ladder_name == name].sort_values("kappa")[field].to_numpy(float)
            for name in metadata.ladder_name
        ])
    return metadata, matrices, metadata.n_realization.to_numpy(float)


def _retrieval_rows(
    metadata: pd.DataFrame, distances: dict[str, np.ndarray], ids: list[str],
    impostors: dict[str, list[int]], ref_neighbours: list[set[int]],
) -> pd.DataFrame:
    id_to_index = {gid: i for i, gid in enumerate(ids)}
    rows: list[dict[str, Any]] = []
    for metric, matrix in distances.items():
        for qidx, meta in metadata.iterrows():
            source = meta.source_galaxy_id
            own_idx = id_to_index[source]
            dist = matrix[qidx]
            own = float(dist[own_idx])
            wrong = np.delete(dist, own_idx)
            nearest_wrong = float(np.min(wrong))
            median_wrong = float(np.median(wrong))
            order = np.argsort(dist, kind="mergesort")
            tol = 1e-12
            rank = int(1 + np.sum(dist < own - tol))
            tie_count = int(np.sum(np.isclose(dist, own, rtol=0, atol=tol)))
            strict_top1 = bool(rank == 1 and tie_count == 1)
            top5 = bool(own_idx in order[:5])
            row = {
                **meta.to_dict(),
                "metric": metric,
                "own_source_distance": own,
                "nearest_wrong_source_distance": nearest_wrong,
                "median_wrong_source_distance": median_wrong,
                "own_to_nearest_wrong_ratio": own / nearest_wrong if nearest_wrong > 0 else np.inf,
                "within_to_between_median_ratio": own / median_wrong if median_wrong > 0 else np.inf,
                "distance_margin_to_nearest_wrong": nearest_wrong - own,
                "source_rank": rank,
                "tie_count_at_own_distance": tie_count,
                "strict_top1": strict_top1,
                "top5": top5,
            }
            if metric == "rho_rms":
                candidate_indices = [own_idx, *impostors[source]]
                candidate_dist = dist[candidate_indices]
                candidate_order = np.argsort(candidate_dist, kind="mergesort")
                own_local = 0
                local_own = float(candidate_dist[own_local])
                local_rank = int(1 + np.sum(candidate_dist < local_own - tol))
                local_ties = int(np.sum(np.isclose(candidate_dist, local_own, rtol=0, atol=tol)))
                query_wrong_order = [int(j) for j in order if int(j) != own_idx][:5]
                qset = set(query_wrong_order)
                rset = ref_neighbours[own_idx]
                row.update({
                    "geometry_matched_source_rank": local_rank,
                    "geometry_matched_strict_top1": bool(local_rank == 1 and local_ties == 1),
                    "geometry_matched_top5": bool(own_local in candidate_order[:5]),
                    "top5_neighbor_jaccard": len(qset & rset) / len(qset | rset),
                })
            else:
                row.update({
                    "geometry_matched_source_rank": np.nan,
                    "geometry_matched_strict_top1": np.nan,
                    "geometry_matched_top5": np.nan,
                    "top5_neighbor_jaccard": np.nan,
                })
            rows.append(row)
    return pd.DataFrame(rows)


def _variant_group(mode: str, realization_id: str) -> str:
    if mode == "LEAVE_ONE_RADIUS_OUT":
        return "L_ALL_DELETIONS"
    return realization_id


def _summary_table(primary: pd.DataFrame, thresholds: dict[str, float]) -> pd.DataFrame:
    rows = []
    for mode in PRIMARY_MODES:
        group = primary.loc[primary["mode"] == mode]
        observed = {
            "strict_top1_accuracy": float(group.strict_top1.mean()),
            "top5_accuracy": float(group.top5.mean()),
            "median_own_to_nearest_wrong_ratio": float(group.own_to_nearest_wrong_ratio.median()),
            "geometry_matched_top1_accuracy": float(group.geometry_matched_strict_top1.mean()),
            "geometry_matched_top5_accuracy": float(group.geometry_matched_top5.mean()),
            "median_top5_neighbor_jaccard": float(group.top5_neighbor_jaccard.median()),
            "realization_count": int(len(group)),
            "galaxy_count": int(group.source_galaxy_id.nunique()),
        }
        checks = {
            "strict_top1_pass": observed["strict_top1_accuracy"] >= thresholds["strict_top1_accuracy_min"],
            "top5_pass": observed["top5_accuracy"] >= thresholds["top5_accuracy_min"],
            "ratio_pass": observed["median_own_to_nearest_wrong_ratio"] <= thresholds["median_own_to_nearest_wrong_ratio_max"],
            "geometry_top1_pass": observed["geometry_matched_top1_accuracy"] >= thresholds["geometry_matched_top1_accuracy_min"],
        }
        rows.append({"mode": mode, "branch": MODE_TO_BRANCH[mode], **observed, **checks, "branch_pass": bool(all(checks.values()))})
    return pd.DataFrame(rows)


def _variant_summary(primary: pd.DataFrame) -> pd.DataFrame:
    work = primary.copy()
    work["variant_group"] = [
        _variant_group(mode, rid) for mode, rid in zip(work["mode"], work["realization_id"])
    ]
    return work.groupby(["mode", "branch", "variant_group"], as_index=False).agg(
        realization_count=("ladder_name", "size"),
        galaxy_count=("source_galaxy_id", "nunique"),
        strict_top1_accuracy=("strict_top1", "mean"),
        top5_accuracy=("top5", "mean"),
        median_own_to_nearest_wrong_ratio=("own_to_nearest_wrong_ratio", "median"),
        geometry_matched_top1_accuracy=("geometry_matched_strict_top1", "mean"),
        median_top5_neighbor_jaccard=("top5_neighbor_jaccard", "median"),
    )


def analyze_gate_c() -> dict[str, Any]:
    readiness = gate_c_readiness()
    if not readiness["ready"]:
        raise ValueError("Gate C is not ready: Gate B evidence or one or more primary branch exports are incomplete")
    ids, reference, reference_sources = load_reference()
    manifest = pd.read_csv(registered_paths()["ladder_manifest"]).set_index("galaxy_id")
    ref_n = np.asarray([float(manifest.loc[gid, "unique_ladder_values"]) for gid in ids])
    _, impostors, control_names = _ordinary_control_space(ids)
    ref_neighbours = _reference_neighbours(reference["rho"])

    retrieval_parts = []
    branch_sources: dict[str, Any] = {}
    for mode in PRIMARY_MODES:
        frame, _, source = _branch_frame(mode)
        counts = frame.groupby("ladder_name").size()
        if not (counts == 40).all():
            raise ValueError(f"{mode}: one or more profiles do not have 40 κ points")
        metadata, query, query_n = _query_matrices(frame)
        products = _distance_products(reference, query, ref_n, query_n)
        retrieval_parts.append(_retrieval_rows(metadata, products, ids, impostors, ref_neighbours))
        branch_sources[mode] = source
    retrieval = pd.concat(retrieval_parts, ignore_index=True)
    primary = retrieval.loc[retrieval.metric == "rho_rms"].copy()
    thresholds = load_config()["gate_c_thresholds"]
    branch_summary = _summary_table(primary, thresholds)
    variant_summary = _variant_summary(primary)
    passed_count = int(branch_summary.branch_pass.sum())
    if passed_count == 4:
        verdict = "GATE_C_ALL_PRIMARY_BRANCHES_PASS"
    elif passed_count > 0:
        verdict = "GATE_C_PARTIAL_PRIMARY_BRANCH_PASS"
    else:
        verdict = "GATE_C_NO_PRIMARY_BRANCH_PASSED_STOP"

    root = Path(active_roots()["analysis_root"])
    tables = root / "tables"
    reports = root / "reports"
    tables.mkdir(parents=True, exist_ok=True)
    reports.mkdir(parents=True, exist_ok=True)
    retrieval_path = tables / "gate_c_retrieval_results.csv"
    summary_path = tables / "gate_c_branch_summary.csv"
    variant_path = tables / "gate_c_variant_summary.csv"
    centroids_path = tables / "gate_c_reference_centroids.csv"
    controls_path = tables / "gate_c_geometry_matched_impostors.csv"

    retrieval.to_csv(retrieval_path, index=False)
    branch_summary.to_csv(summary_path, index=False)
    variant_summary.to_csv(variant_path, index=False)
    centroid_rows = []
    for i, gid in enumerate(ids):
        for j, kappa in enumerate(reference["kappa"]):
            centroid_rows.append({
                "galaxy_id": gid, "kappa": kappa,
                "rho": reference["rho"][i, j], "Ak": reference["Ak"][i, j], "nu_V": reference["nu_V"][i, j],
            })
    pd.DataFrame(centroid_rows).to_csv(centroids_path, index=False)
    control_rows = []
    for gid, indices in impostors.items():
        for rank, idx in enumerate(indices, start=1):
            control_rows.append({"source_galaxy_id": gid, "impostor_rank": rank, "impostor_galaxy_id": ids[idx]})
    pd.DataFrame(control_rows).to_csv(controls_path, index=False)

    amendment_path = write_gate_b_provenance_amendment()
    record = {
        "chamber": "GAL-FINGERPRINT-I",
        "controller_version": CHAMBER_VERSION,
        "runtime_package_version": CHAMBER_VERSION,
        "measurement_kernel": "STRUC-I",
        "measurement_kernel_version": "1.0.4",
        "protocol_version": "0.1",
        "gate": "GATE_C",
        "timestamp": utc_now(),
        "verdict": verdict,
        "passed_branch_count": passed_count,
        "all_primary_branches_passed": passed_count == 4,
        "gate_d_permitted_by_stop_rule": passed_count > 0,
        "final_scientific_verdict_assigned": False,
        "qualifier": "OBSERVATIONAL_UNCERTAINTY_NOT_TESTED",
        "gate_b_evidence": readiness["gate_b"],
        "reference_sources": reference_sources,
        "branch_sources": branch_sources,
        "primary_metric": "rho_rms",
        "sensitivity_metrics": ["combined_rms", "rho_pca3"],
        "geometry_control_fields": control_names,
        "thresholds": thresholds,
        "branch_summary": branch_summary.to_dict(orient="records"),
        "outputs": {
            "retrieval_results": str(retrieval_path.resolve()),
            "branch_summary": str(summary_path.resolve()),
            "variant_summary": str(variant_path.resolve()),
            "reference_centroids": str(centroids_path.resolve()),
            "geometry_matched_impostors": str(controls_path.resolve()),
            "gate_b_provenance_amendment": str(amendment_path.resolve()),
        },
    }
    json_path = reports / "GAL_FINGERPRINT_I_GATE_C_RESULT_RECORD.json"
    md_path = reports / "GAL_FINGERPRINT_I_GATE_C_RESULT_RECORD.md"
    report_path = reports / "GAL_FINGERPRINT_I_GATE_C_REPORT.md"
    json_path.write_text(json.dumps(record, indent=2, ensure_ascii=False), encoding="utf-8")

    table_lines = [
        "| Branch | Realizations | Top-1 | Top-5 | Median own/wrong | Geometry top-1 | Pass |",
        "|---|---:|---:|---:|---:|---:|:---:|",
    ]
    for row in record["branch_summary"]:
        table_lines.append(
            f"| {row['branch']} — {row['mode']} | {row['realization_count']} | {row['strict_top1_accuracy']:.6f} | "
            f"{row['top5_accuracy']:.6f} | {row['median_own_to_nearest_wrong_ratio']:.6f} | "
            f"{row['geometry_matched_top1_accuracy']:.6f} | {'Yes' if row['branch_pass'] else 'No'} |"
        )
    table = "\n".join(table_lines)
    md = f"""# GAL-FINGERPRINT-I Gate C result record

```text
verdict                              {verdict}
passed primary branches              {passed_count} / 4
gate D permitted by stop rule        {'yes' if passed_count > 0 else 'no'}
final scientific verdict assigned    no
qualifier                            OBSERVATIONAL_UNCERTAINTY_NOT_TESTED
controller version                   {CHAMBER_VERSION}
measurement kernel                   STRUC-I v1.0.4
```

## Registered primary branches

{table}

Gate C tests sampling and baryonic-preparation robustness. It does not include random thinning or chamber-null models, and it does not assign the final GAL-FINGERPRINT-I verdict.
"""
    md_path.write_text(md, encoding="utf-8")
    report_path.write_text(md.replace("result record", "registered analysis report", 1), encoding="utf-8")
    record["outputs"].update({
        "result_json": str(json_path.resolve()),
        "result_markdown": str(md_path.resolve()),
        "report_markdown": str(report_path.resolve()),
    })
    json_path.write_text(json.dumps(record, indent=2, ensure_ascii=False), encoding="utf-8")
    return record
