#!/usr/bin/env python3
"""Build measured Poucher substantivity regression targets and gate baselines."""
from __future__ import annotations

import json
import math
from collections import Counter
from pathlib import Path

import numpy as np

DATA = Path("data")
ARTIFACTS = Path("artifacts")

FEATURE_NAMES = [
    "log10_vapor_pressure_pa",
    "boiling_point_k",
    "molecular_weight",
    "logp",
]


def load_dataset_cas() -> set[str]:
    cas_set: set[str] = set()
    with open(DATA / "empirical_dataset_v8.jsonl") as f:
        for line in f:
            record = json.loads(line)
            if record.get("is_control"):
                continue
            for comp in record.get("formula", []):
                cas = comp.get("cas")
                if cas:
                    cas_set.add(cas)
    return cas_set


def load_poucher_targets() -> dict[str, dict]:
    rows = {}
    with open(DATA / "poucher_substantivity.jsonl") as f:
        for line in f:
            row = json.loads(line)
            rows[row["cas"]] = row
    return rows


def load_aroma_features() -> dict[str, dict]:
    features = {}
    with open(DATA / "aroma_chemicals.jsonl") as f:
        for line in f:
            row = json.loads(line)
            cas = row.get("cas")
            if cas:
                features[cas] = row
    return features


def feature_row(raw: dict | None) -> dict:
    if raw is None:
        return {name: None for name in FEATURE_NAMES}
    vp = raw.get("vapor_pressure_pa")
    return {
        "log10_vapor_pressure_pa": math.log10(vp) if isinstance(vp, (int, float)) and vp > 0 else None,
        "boiling_point_k": raw.get("boiling_point_k"),
        "molecular_weight": raw.get("molecular_weight"),
        "logp": raw.get("logp"),
    }


def write_targets() -> dict:
    dataset_cas = load_dataset_cas()
    poucher = load_poucher_targets()
    aroma = load_aroma_features()

    all_rows = []
    dataset_rows = []
    for cas, target in sorted(poucher.items()):
        raw_features = aroma.get(cas)
        features = feature_row(raw_features)
        coeff = float(target["poucher_coefficient"])
        row = {
            "cas": cas,
            "name": target["name"],
            "target": {
                "raw_poucher_coefficient": coeff,
                "log10_poucher_coefficient": math.log10(coeff),
                "source": target["source"],
                "source_page": target["source_page"],
                "target_type": target["target_type"],
            },
            "features": features,
            "feature_role": "molecular_features_only_not_target_derivation",
            "in_empirical_dataset_v8": cas in dataset_cas,
            "feature_source": "data/aroma_chemicals.jsonl",
            "cas_match_source": target.get("cas_match_source"),
            "coefficient_conflict": target.get("coefficient_conflict", False),
        }
        all_rows.append(row)
        if row["in_empirical_dataset_v8"]:
            dataset_rows.append(row)

    (DATA / "substantivity_targets_poucher.jsonl").write_text(
        "\n".join(json.dumps(r, sort_keys=True) for r in all_rows) + "\n"
    )
    (DATA / "substantivity_targets_dataset.jsonl").write_text(
        "\n".join(json.dumps(r, sort_keys=True) for r in dataset_rows) + "\n"
    )

    values = np.array([r["target"]["raw_poucher_coefficient"] for r in all_rows], dtype=float)
    dataset_values = np.array([r["target"]["raw_poucher_coefficient"] for r in dataset_rows], dtype=float)

    transform = {
        "target_source": "data/poucher_substantivity.jsonl",
        "raw_target": "Poucher measured duration-of-evaporation coefficient, 1-100",
        "model_target": "log10_poucher_coefficient",
        "transform": "log10(raw_poucher_coefficient)",
        "inverse_transform": "10 ** model_prediction",
        "reason": "positive bounded coefficient with long upper tail; features remain separate and are never used to derive the target",
        "tier_band_boundaries_for_presentation_only": {"top": [1, 14], "mid": [15, 60], "base": [61, 100]},
    }
    ARTIFACTS.mkdir(exist_ok=True)
    (ARTIFACTS / "substantivity_transform.json").write_text(json.dumps(transform, indent=2))

    hist_counts, hist_edges = np.histogram(values, bins=[1, 15, 31, 46, 61, 81, 101])
    summary = {
        "poucher_rows": len(all_rows),
        "empirical_dataset_unique_cas": len(dataset_cas),
        "empirical_dataset_measured_coverage": len(dataset_rows),
        "empirical_dataset_measured_coverage_pct": len(dataset_rows) / len(dataset_cas),
        "raw_summary": summarize(values),
        "dataset_raw_summary": summarize(dataset_values) if len(dataset_values) else None,
        "raw_histogram": {
            "bins": ["1-14", "15-30", "31-45", "46-60", "61-80", "81-100"],
            "counts": hist_counts.tolist(),
        },
        "tier_band_distribution_all": dict(Counter(r["target"]["raw_poucher_coefficient"] <= 14 and "top" or r["target"]["raw_poucher_coefficient"] <= 60 and "mid" or "base" for r in all_rows)),
        "feature_completeness_all": feature_completeness(all_rows),
        "feature_completeness_dataset": feature_completeness(dataset_rows),
    }
    (ARTIFACTS / "substantivity_target_summary.json").write_text(json.dumps(summary, indent=2))
    return summary


def summarize(values: np.ndarray) -> dict:
    return {
        "n": int(len(values)),
        "min": float(np.min(values)),
        "p25": float(np.percentile(values, 25)),
        "median": float(np.median(values)),
        "mean": float(np.mean(values)),
        "p75": float(np.percentile(values, 75)),
        "max": float(np.max(values)),
    }


def feature_completeness(rows: list[dict]) -> dict:
    return {
        name: sum(r["features"].get(name) is not None for r in rows)
        for name in FEATURE_NAMES
    }


def matrix_from_rows(rows: list[dict]) -> tuple[np.ndarray, np.ndarray]:
    y = np.array([r["target"]["log10_poucher_coefficient"] for r in rows], dtype=float)
    raw = []
    missing = []
    for r in rows:
        vals = [r["features"].get(name) for name in FEATURE_NAMES]
        raw.append([np.nan if v is None else float(v) for v in vals])
        missing.append([1.0 if v is None else 0.0 for v in vals])
    x = np.array(raw, dtype=float)
    miss = np.array(missing, dtype=float)
    return np.concatenate([x, miss], axis=1), y


def impute_standardize(train_x: np.ndarray, test_x: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    med = np.nanmedian(train_x, axis=0)
    med = np.where(np.isnan(med), 0.0, med)
    train = np.where(np.isnan(train_x), med, train_x)
    test = np.where(np.isnan(test_x), med, test_x)
    mean = train.mean(axis=0)
    std = train.std(axis=0)
    std = np.where(std == 0, 1.0, std)
    return (train - mean) / std, (test - mean) / std


def r2_score(y: np.ndarray, pred: np.ndarray) -> float:
    ss_res = float(np.sum((y - pred) ** 2))
    ss_tot = float(np.sum((y - y.mean()) ** 2))
    return 1.0 - ss_res / ss_tot if ss_tot else float("nan")


def folds(n: int, k: int = 5, seed: int = 20260710) -> list[np.ndarray]:
    rng = np.random.default_rng(seed)
    idx = np.arange(n)
    rng.shuffle(idx)
    return np.array_split(idx, k)


def ridge_cv(x: np.ndarray, y: np.ndarray) -> tuple[float, float]:
    parts = folds(len(y))
    pred = np.zeros_like(y)
    alpha = 1.0
    for test_idx in parts:
        train_idx = np.setdiff1d(np.arange(len(y)), test_idx)
        xt, xv = impute_standardize(x[train_idx], x[test_idx])
        yt = y[train_idx]
        xt1 = np.c_[np.ones(len(xt)), xt]
        xv1 = np.c_[np.ones(len(xv)), xv]
        reg = np.eye(xt1.shape[1]) * alpha
        reg[0, 0] = 0.0
        beta = np.linalg.pinv(xt1.T @ xt1 + reg) @ xt1.T @ yt
        pred[test_idx] = xv1 @ beta
    return r2_score(y, pred), float(np.corrcoef(y, pred)[0, 1])


def stump_gbm_cv(x: np.ndarray, y: np.ndarray) -> tuple[float, float]:
    parts = folds(len(y))
    pred = np.zeros_like(y)
    for test_idx in parts:
        train_idx = np.setdiff1d(np.arange(len(y)), test_idx)
        xt, xv = impute_standardize(x[train_idx], x[test_idx])
        yt = y[train_idx]
        train_pred = np.full(len(yt), yt.mean())
        test_pred = np.full(len(test_idx), yt.mean())
        lr = 0.05
        for _ in range(160):
            residual = yt - train_pred
            best = None
            for j in range(xt.shape[1]):
                thresholds = np.unique(np.quantile(xt[:, j], np.linspace(0.1, 0.9, 9)))
                for threshold in thresholds:
                    left = xt[:, j] <= threshold
                    if left.sum() == 0 or (~left).sum() == 0:
                        continue
                    lv = residual[left].mean()
                    rv = residual[~left].mean()
                    update = np.where(left, lv, rv)
                    sse = float(np.sum((residual - update) ** 2))
                    if best is None or sse < best[0]:
                        best = (sse, j, threshold, lv, rv)
            if best is None:
                break
            _, j, threshold, lv, rv = best
            train_pred += lr * np.where(xt[:, j] <= threshold, lv, rv)
            test_pred += lr * np.where(xv[:, j] <= threshold, lv, rv)
        pred[test_idx] = test_pred
    return r2_score(y, pred), float(np.corrcoef(y, pred)[0, 1])


def run_gate() -> dict:
    rows = [json.loads(line) for line in open(DATA / "substantivity_targets_poucher.jsonl")]
    feature_ready = [
        r for r in rows
        if any(r["features"].get(name) is not None for name in FEATURE_NAMES)
    ]
    x, y = matrix_from_rows(feature_ready)
    ridge_r2, ridge_corr = ridge_cv(x, y)
    gbm_r2, gbm_corr = stump_gbm_cv(x, y)
    result = {
        "target": "log10_poucher_coefficient",
        "n_total_poucher_targets": len(rows),
        "n_with_any_molecular_feature": len(feature_ready),
        "features": FEATURE_NAMES,
        "missing_value_handling": "train-fold median imputation plus missing indicators",
        "ridge_linear_5fold_r2": ridge_r2,
        "ridge_linear_5fold_pearson": ridge_corr,
        "stump_gbm_5fold_r2": gbm_r2,
        "stump_gbm_5fold_pearson": gbm_corr,
        "gate_interpretation": interpret(max(ridge_r2, gbm_r2)),
    }
    (ARTIFACTS / "substantivity_non_circularity_gate.json").write_text(json.dumps(result, indent=2))
    return result


def interpret(best_r2: float) -> str:
    if best_r2 >= 0.95:
        return "STOP: target is essentially reconstructed by simple physicochemical features"
    if best_r2 < 0.2:
        return "STOP_AND_ESCALATE: simple features have very low signal for measured target"
    return "PROCEED: moderate signal, not a closed-form feature formula"


def main() -> None:
    summary = write_targets()
    gate = run_gate()
    print(json.dumps({"summary": summary, "gate": gate}, indent=2))


if __name__ == "__main__":
    main()