#!/usr/bin/env python3
"""
E4-H034 — Elastic dipoles as predictors of MLIP error.

TWO SEPARABLE PARTS, and they must not be confused:

  P1  THEORY (analytic, no data): what error law does a finite-cutoff model imply?
      A point elastic dipole in an infinite isotropic medium has strain
          eps(r) ~ P / (kappa r^3)
      so the elastic energy stored OUTSIDE a cutoff R_c is
          E_out ~ P^2 / (kappa R_c^3)          (integrate eps^2 over r > R_c)
      A model with cutoff R_c therefore *cannot* represent E_out, and the
      candidate error law is  error ~ c * P^2 / (kappa R_c^3).
      This is a derivation with a stated assumption (point dipole, isotropic,
      infinite medium). It is NOT evidence about any real potential.

  P2  PROTOCOL (synthetic, labelled): demonstrate the evaluation that would be
      needed to test whether elastic-dipole descriptors add information beyond
      simple descriptors -- and show that a random train/test split gives a
      spuriously positive answer.  NO claim is made about real MLIPs here.

Run: python3 h034_theory_and_protocol.py
Out: ../results/*.csv
"""
import numpy as np
import csv, os

OUT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "results"))
rng = np.random.default_rng(20261101)


# ---------------------------------------------------------------- P1 theory
def elastic_energy_outside(P, kappa, Rc):
    """Energy of a point dipole outside radius Rc:  E ~ P^2 / (kappa Rc^3).
    P in eV, kappa in eV/A^3 (bulk modulus scale), Rc in A.  Returns eV."""
    return 0.18 * P ** 2 / (kappa * Rc ** 3)


def p1_table():
    print("P1  derived error law  error ~ c * P^2 / (kappa * Rc^3)")
    print(f"    {'P (eV)':>7s} {'kappa (eV/A^3)':>16s} {'Rc (A)':>7s} "
          f"{'E_out (meV)':>12s} {'E_out/E_out(4A)':>16s}")
    rows = []
    base = elastic_energy_outside(3.0, 1.0, 4.0)
    for P in (1.0, 2.0, 3.0, 6.0):
        for Rc in (4.0, 6.0, 8.0):
            E = elastic_energy_outside(P, 1.0, Rc)
            print(f"    {P:7.1f} {1.0:16.2f} {Rc:7.1f} {E*1e3:12.3f} "
                  f"{E/base:16.3f}")
            rows.append(dict(part="P1", P_eV=P, kappa=1.0, Rc_A=Rc,
                             E_out_meV=E * 1e3, ratio_to_base=E / base))
    print("\n    scalings implied by the law: quadratic in |P|, inverse-cube in Rc,")
    print("    inverse-linear in stiffness.  A model with the SAME Rc but a different")
    print("    architecture may differ in c only; the exponents are architecture-blind.")
    return rows


# ---------------------------------------------------------------- P2 protocol
def synth_dataset(n_materials=12, n_defects=40, rng=rng):
    """Synthetic defect data with a MATERIAL-LEVEL latent variable.

    Two generative regimes are created so the protocol can be checked:
      - 'signal' : the error depends on P^2/(kappa Rc^3) AND on simple descriptors
      - 'null'   : the error depends only on simple descriptors; P adds nothing
    Every configuration carries a material id, so grouped splits are possible.
    """
    data = []
    for mid in range(n_materials):
        kappa = rng.uniform(0.5, 2.0)          # material stiffness
        Rc = rng.choice([4.0, 5.0, 6.0])       # model cutoff (architecture proxy)
        mat_bias = rng.normal(0, 1)            # material-level nuisance
        for _ in range(n_defects):
            P = abs(rng.normal(3.0, 1.5))      # elastic dipole magnitude
            size = P + rng.normal(0, 0.5)      # simple descriptor: defect size
            coord = rng.normal(0, 1)           # simple descriptor: chemistry proxy
            e_el = elastic_energy_outside(P, kappa, Rc)
            # true error: simple terms + optional elastic term + material bias
            simple = 0.8 * size + 0.5 * coord + 0.9 * mat_bias
            data.append(dict(material=mid, P=P, size=size, coord=coord,
                             kappa=kappa, Rc=Rc, e_el=e_el, simple=simple))
    return data


def make_target(data, regime):
    y = []
    for d in data:
        if regime == "signal":
            y.append(d["simple"] + 40.0 * d["e_el"] + rng.normal(0, 0.3))
        else:
            y.append(d["simple"] + rng.normal(0, 0.3))
    return np.array(y)


def fit_ols(X, y):
    X1 = np.column_stack([np.ones(len(X)), X])
    beta, *_ = np.linalg.lstsq(X1, y, rcond=None)
    return beta


def predict(beta, X):
    return np.column_stack([np.ones(len(X)), X]) @ beta


def group_cv(data, y, feats, groups, n_folds=6):
    """Leave-materials-out cross-validation; returns RMSE."""
    errs = []
    uniq = np.unique(groups)
    for k in range(n_folds):
        test_m = uniq[k::n_folds]
        te = np.isin(groups, test_m)
        tr = ~te
        Xtr = feats[tr]
        Xte = feats[te]
        beta = fit_ols(Xtr, y[tr])
        errs.append(np.sqrt(np.mean((predict(beta, Xte) - y[te]) ** 2)))
    return float(np.mean(errs))


def p2_protocol():
    print("\nP2  PROTOCOL DEMONSTRATION (synthetic data -- NOT evidence about real MLIPs)")
    data = synth_dataset()
    groups = np.array([d["material"] for d in data])
    x_simple = np.array([[d["size"], d["coord"]] for d in data])
    x_el = np.array([[d["e_el"]] for d in data])
    x_both = np.hstack([x_simple, x_el])

    rows = []
    print(f"    {'regime':>8s} {'split':>10s} {'simple':>10s} {'+dipole':>10s} "
          f"{'improvement':>12s}")
    for regime in ("signal", "null"):
        y = make_target(data, regime)
        # (a) random split -- leaks material identity
        idx = rng.permutation(len(y))
        tr, te = idx[: int(0.7 * len(y))], idx[int(0.7 * len(y)):]
        rmse = {}
        for name, X in (("simple", x_simple), ("both", x_both)):
            beta = fit_ols(X[tr], y[tr])
            rmse[name] = float(np.sqrt(np.mean((predict(beta, X[te]) - y[te]) ** 2)))
        print(f"    {regime:>8s} {'random':>10s} {rmse['simple']:10.3f} "
              f"{rmse['both']:10.3f} {rmse['simple']-rmse['both']:12.3f}")
        rows.append(dict(part="P2", regime=regime, split="random",
                         rmse_simple=rmse["simple"], rmse_both=rmse["both"],
                         improvement=rmse["simple"] - rmse["both"]))
        # (b) grouped split -- material held out
        g1 = group_cv(data, y, x_simple, groups)
        g2 = group_cv(data, y, x_both, groups)
        print(f"    {regime:>8s} {'grouped':>10s} {g1:10.3f} {g2:10.3f} {g1-g2:12.3f}")
        rows.append(dict(part="P2", regime=regime, split="grouped",
                         rmse_simple=g1, rmse_both=g2, improvement=g1 - g2))

    print("\n    reading the table:")
    print("      signal/grouped improvement >> 0  -> the dipole adds information")
    print("      null/grouped   improvement ~ 0   -> it does not")
    print("      any regime/random improvement > 0 -> material leakage, not signal")
    return rows


if __name__ == "__main__":
    os.makedirs(OUT, exist_ok=True)
    rows = p1_table() + p2_protocol()
    path = os.path.join(OUT, "h034_results.csv")
    keys = sorted({k for r in rows for k in r})
    with open(path, "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=keys)
        w.writeheader()
        for r in rows:
            w.writerow(r)
    print(f"\nwrote {path} ({len(rows)} rows)")