"""NON-CANONICAL exploration for her-model (ocx25).

The canonical prepare_targets.py / evaluate.py from the ocx25-data Skill could not be
read, so this script reconstructs an approximate target: the per-sample `avg` HER
full-cell voltage at 50 mA/cm2 from experiments.csv (NOT voltage_she_v). Pearson r^2 and
Spearman rho are invariant to an affine full-cell->SHE conversion; MAE/r2 would be too only
if the conversion is a pure offset.

Features use only composition (target composition string) and a small hard-coded table of
textbook elemental properties (Pauling electronegativity, d-band filling group number,
atomic radius pm), fixed a priori. LOCO folds = sorted element set of the target composition.
Usage: python3 -I explore_loco.py <inputs_dir> <out_dir>
"""
import sys, json
import numpy as np, pandas as pd
from scipy.stats import pearsonr, spearmanr
from sklearn.linear_model import RidgeCV
from sklearn.ensemble import RandomForestRegressor
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

inp, out = sys.argv[1], sys.argv[2]
# Pauling EN, group number, metallic radius (pm) -- textbook values
PROPS = {
 'Ag':(1.93,11,144),'Au':(2.54,11,144),'Cd':(1.69,12,151),'Cu':(1.90,11,128),'Ga':(1.81,13,135),
 'Hg':(2.00,12,151),'In':(1.78,13,167),'Mn':(1.55,7,127),'Mo':(2.16,6,139),'Nb':(1.60,5,146),
 'Ni':(1.91,10,124),'Pd':(2.20,10,137),'Pt':(2.28,10,139),'Rh':(2.28,9,134),'Ru':(2.20,8,134),
 'Ir':(2.20,9,136),'Sn':(1.96,14,151),'Tl':(1.62,13,170),'W':(2.36,6,139),'Zn':(1.65,12,134),
 'Co':(1.88,9,125),'Fe':(1.83,8,126),'Bi':(2.02,15,156),'Pb':(2.33,14,175),'Sb':(2.05,15,145),
 'Cr':(1.66,6,128),'Ti':(1.54,4,147),'V':(1.63,5,134),'Zr':(1.33,4,160),'Ta':(1.50,5,146),
 'Re':(1.90,7,137),'Os':(2.20,8,135),'Al':(1.61,13,143),'Ge':(2.01,14,122),'Se':(2.55,16,120),
 'Te':(2.10,16,140),'Y':(1.22,3,180),'Sc':(1.36,3,162),'Hf':(1.30,4,159),'La':(1.10,3,187),
}
# Known good HER metals (Pt-group) indicator, fixed a priori from textbook volcano plots
PGM = {'Pt','Pd','Rh','Ir','Ru','Os'}

def parse(c):
    t = c.split('-'); return {t[i]: float(t[i+1]) for i in range(0, len(t), 2)}

e = pd.read_csv(f'{inp}/experiments.csv')
h = e[(e.reaction == 'HER') & (e.rep == 'avg') & (e['current density'] == 50)].copy()
h['comp'] = h['composition'].map(parse)
h['elements'] = h['comp'].map(lambda d: '-'.join(sorted(k for k, v in d.items() if v > 0)))
els = sorted({k for d in h.comp for k in d})
missing = [x for x in els if x not in PROPS]
assert not missing, missing

def feats(d):
    tot = sum(d.values()); f = {k: v / tot for k, v in d.items()}
    P = np.array([PROPS[k] for k in f]); w = np.array(list(f.values()))
    mean = w @ P; var = w @ (P - mean) ** 2
    return list(mean) + list(np.sqrt(var)) + [sum(v for k, v in f.items() if k in PGM),
            max(f.values()), len(f)]
FN = ['en_mean','grp_mean','rad_mean','en_sd','grp_sd','rad_sd','pgm_frac','max_frac','n_el']
X = np.array([feats(d) for d in h.comp]); y = h.voltage.values; g = h.elements.values
frac = np.array([[d.get(k, 0) / sum(d.values()) for k in els] for d in h.comp])

models = {
 'mean_only': None,
 'ridge_props': lambda: make_pipeline(StandardScaler(), RidgeCV(alphas=np.logspace(-3, 3, 13))),
 'rf_props': lambda: RandomForestRegressor(n_estimators=500, min_samples_leaf=3, random_state=0),
}
res = {}; preds = {}
for name, mk in models.items():
    p = np.zeros_like(y); fold = np.empty(len(y), dtype=object)
    for sysn in np.unique(g):
        te = g == sysn; tr = ~te
        if mk is None: p[te] = y[tr].mean()
        else:
            m = mk(); m.fit(X[tr], y[tr]); p[te] = m.predict(X[te])
        fold[te] = sysn
    r = pearsonr(y, p)[0] if np.std(p) > 0 else float('nan')
    res[name] = dict(pearson_r=float(r), pearson_r2=float(r**2) if r == r else None,
                     r2=float(1 - ((y-p)**2).sum() / ((y-y.mean())**2).sum()),
                     mae=float(np.abs(y-p).mean()), spearman_rho=float(spearmanr(y, p)[0]) if np.std(p)>0 else None,
                     n=int(len(y)), n_folds=int(len(np.unique(g))))
    preds[name] = p
    pd.DataFrame({'sample_id': h['sample id'], 'prediction': p, 'fold': fold}).to_csv(
        f'{out}/pred_{name}.csv', index=False)

# system-level bootstrap of rf - ridge pearson_r2, seed 0
rng = np.random.default_rng(0); systems = np.unique(g); diffs = []
idx_by = {s: np.where(g == s)[0] for s in systems}
for _ in range(2000):
    idx = np.concatenate([idx_by[s] for s in rng.choice(systems, len(systems))])
    a = pearsonr(y[idx], preds['rf_props'][idx])[0]**2; b = pearsonr(y[idx], preds['ridge_props'][idx])[0]**2
    diffs.append(a - b)
res['bootstrap_rf_minus_ridge_pearson_r2'] = dict(method='resample whole element systems, 2000x', seed=0,
    median=float(np.median(diffs)), ci95=[float(np.percentile(diffs, 2.5)), float(np.percentile(diffs, 97.5))])
# worst systems for rf
err = pd.DataFrame({'sys': g, 'ae': np.abs(y - preds['rf_props'])}).groupby('sys').ae.mean().sort_values(ascending=False)
res['rf_worst_systems_mae'] = {k: round(float(v), 3) for k, v in err.head(8).items()}
res['n_rows'] = int(len(y)); res['n_systems'] = int(len(systems)); res['features'] = FN
print(json.dumps(res, indent=1))
json.dump(res, open(f'{out}/explore_results.json', 'w'), indent=1)