#!/usr/bin/env python3
"""Derive the canonical OCx25 modelling targets from the bound reference inputs.

Standard library only. Run it on the exact bytes TfP serves for a Problem:

    python3 prepare_targets.py --experiments experiments.csv --xrf xrf.csv --out-dir targets/

It writes:

    targets/targets_her.csv      one row per (XRF composition, source): HER full-cell
                                 voltage at 50 mA/cm2 and its half-cell conversion
    targets/targets_co2rr.csv    one row per (XRF composition, source): CO2RR
                                 Faradaic efficiencies and production rates at a
                                 full-cell voltage of 3.3 V (log-linear interpolation)

The steps are a port of the processing in the OCx24 reference code
(fairchem, src/fairchem/applications/ocx/core/data_handling.py and voltage.py,
MIT licence), with one deliberate simplification so that the derivation needs
no 1.5 GB computational file: replicates are pooled by synthesis target
(the sample id without its replicate suffix, per source), and the pooled
row carries the element-wise mean of its replicates' XRF compositions,
rather than pooling by the nearest composition in the computational
screening set. Everything else (dropping the avg/std rows,
XRF lookup by source, abs(voltage), mean over replicates, the electron
counts, the log-linear interpolation to 3.3 V without extrapolation, the
half-cell conversion, the element-system label used for leave-one-
composition-out folds) follows the reference code. See reference.md.
"""

from __future__ import annotations

import argparse
import csv
import math
import os
import re
from collections import OrderedDict, defaultdict

FARADAY = 96485.0  # (s A)/mol
PRODUCTS = ["fe_h2", "fe_co", "fe_ch4", "fe_c2h4", "fe_liquid"]
# Electron counts exactly as the reference implementation uses them
# (get_all_production_rates). Note that it uses 2 for C2H4; the physical
# value is 12. Keep the reference convention so numbers stay comparable.
ELECTRONS = {"fe_h2": 2, "fe_co": 2, "fe_ch4": 8, "fe_liquid": 2, "fe_c2h4": 2}
RATE_NAME = {
    "fe_h2": "h2_pr",
    "fe_co": "co_pr",
    "fe_ch4": "ch4_pr",
    "fe_liquid": "liquid_pr",
    "fe_c2h4": "c2h4_pr",
}
INTERPOLATE_V = 3.3


def read_csv(path: str) -> list[dict[str, str]]:
    with open(path, newline="", encoding="utf-8") as f:
        return list(csv.DictReader(f))


def strip_rep(sample_id: str) -> str:
    return re.sub(r"_rep\d+", "", sample_id)


def xrf_mapping(xrf_rows: list[dict[str, str]]) -> dict[str, str]:
    """vsp samples use the XRF measured on the electrode (gde); uoft samples use
    the wafer measurement keyed by the sample id without its replicate suffix."""
    mapping: dict[str, str] = {}
    for row in xrf_rows:
        source = row["source"]
        sid = row["sample id"]
        if source == "vsp":
            value = row.get("xrf-gde_formula", "")
            if value:
                mapping[sid] = value
        elif source == "uoft":
            value = row.get("xrf-wafer_formula", "")
            if value:
                mapping[strip_rep(sid)] = value
    return mapping


def lookup_xrf(row: dict[str, str], mapping: dict[str, str]) -> str | None:
    if row["source"] == "vsp":
        return mapping.get(row["sample id"])
    if row["source"] == "uoft":
        return mapping.get(strip_rep(row["sample id"]))
    return None


def elements_of(comp: str) -> str:
    """'Ag-0.5-Au-0.5' -> 'AgAu' (the leave-one-composition-out group label)."""
    return "".join(comp.split("-")[::2])


def to_float(value: str) -> float:
    value = (value or "").strip()
    if value == "":
        return math.nan
    return float(value)


def mean(values: list[float]) -> float:
    values = [v for v in values if not math.isnan(v)]
    if not values:
        return math.nan
    return sum(values) / len(values)


# --- half-cell conversion (port of fairchem voltage.py) --------------------

def _she2rhe(ushe: float, ph: float, ref_pot: float) -> float:
    return ushe + ref_pot + 0.059 * ph


def _rhe2she(urhe: float, ph: float, ref_pot: float) -> float:
    return urhe - ref_pot - 0.059 * ph


def _correct_potential(pot: float, r: float, ph: float, j: float, area: float, ref_pot: float) -> float:
    if pot < 0:
        return _she2rhe(pot + j / 1000 * area * r, ph, ref_pot)
    return _she2rhe(pot - j / 1000 * area * r, ph, ref_pot)


def _linear_fit(xs: list[float], ys: list[float]) -> tuple[float, float]:
    n = len(xs)
    mx = sum(xs) / n
    my = sum(ys) / n
    sxx = sum((x - mx) ** 2 for x in xs)
    sxy = sum((x - mx) * (y - my) for x, y in zip(xs, ys))
    a = sxy / sxx
    return a, my - a * mx


def get_she(vcell: float) -> float:
    """Full-cell voltage -> cathode potential vs SHE, using the calibration the
    reference code embeds (neutral CO2RR, 4 cm2 cell, sputtered Cu, 0.1 M
    bicarbonate, Ag/AgCl 3 M KCl at 0.23 V vs SHE). Affine in vcell."""
    ref_pot = 0.23
    cathode_ph = 12.5
    area = 4.0
    j = [50.0, 100.0, 200.0]
    cathode_pot = [-1.62, -2.0, -2.3]
    cathode_r = [0.48, 0.34, 0.3]
    fullcell_pot = [3.0, 3.4, 3.7]
    cathode_pot_corr = [
        _correct_potential(cathode_pot[i], cathode_r[i], cathode_ph, j[i], area, ref_pot) for i in range(3)
    ]
    a, b = _linear_fit(fullcell_pot, cathode_pot_corr)
    urhe = a * vcell + b
    return _rhe2she(urhe, cathode_ph, ref_pot)


# --- processing ------------------------------------------------------------

def base_rows(exp_rows: list[dict[str, str]], mapping: dict[str, str], reaction: str) -> list[dict]:
    out = []
    for row in exp_rows:
        if row["reaction"] != reaction:
            continue
        if row["rep"] in ("avg", "std"):
            continue
        xrf = lookup_xrf(row, mapping)
        if not xrf:
            continue
        out.append(
            {
                "sample id": row["sample id"],
                "source": row["source"],
                "composition": row["composition"],
                "xrf comp": xrf,
                "current density": float(row["current density"]),
                "voltage": abs(to_float(row["voltage"])),
                **{p: to_float(row.get(p, "")) for p in PRODUCTS},
            }
        )
    return out


def parse_comp(comp: str) -> dict[str, float]:
    parts = comp.split("-")
    return {parts[i]: float(parts[i + 1]) for i in range(0, len(parts) - 1, 2)}


def format_comp(fractions: dict[str, float]) -> str:
    return "-".join(f"{el}-{fractions[el]:.3f}" for el in sorted(fractions))


def mean_comp(comps: list[str]) -> str:
    """Element-wise mean of replicate XRF compositions, renormalised to 1."""
    totals: dict[str, float] = defaultdict(float)
    for c in comps:
        for el, x in parse_comp(c).items():
            totals[el] += x
    n = len(comps)
    fractions = {el: v / n for el, v in totals.items()}
    s = sum(fractions.values()) or 1.0
    return format_comp({el: v / s for el, v in fractions.items()})


def aggregate(rows: list[dict]) -> "OrderedDict[tuple, dict]":
    """Mean over replicates sharing (synthesis target, source, current density).

    The synthesis target is the sample id without its `_repN` suffix: the same
    composition printed or reduced in the same batch. Pooling by it needs no
    external file; the reference pipeline pools by the nearest composition in
    the computational set instead, which is close but not identical."""
    groups: "OrderedDict[tuple, list[dict]]" = OrderedDict()
    for r in rows:
        groups.setdefault((strip_rep(r["sample id"]), r["source"], r["current density"]), []).append(r)
    agg: "OrderedDict[tuple, dict]" = OrderedDict()
    for key, members in groups.items():
        first = members[0]
        entry = {
            "sample id": key[0],
            "sample ids": sorted({m["sample id"] for m in members}),
            "source": first["source"],
            "composition": first["composition"],
            "xrf comp": mean_comp([m["xrf comp"] for m in members]),
            "current density": first["current density"],
            "n_rows": len(members),
            "voltage": mean([m["voltage"] for m in members]),
        }
        for p in PRODUCTS:
            entry[p] = mean([m[p] for m in members])
        agg[key] = entry
    return agg


def production_rate(fe: float, nel: int, current_density: float) -> float:
    if math.isnan(fe):
        return math.nan
    return ((fe / 100.0) * (current_density / 1000.0)) / (nel * FARADAY)


def log_interp(v: float, points: list[tuple[float, float]]) -> float:
    """Log-linear interpolation of y at voltage v from (voltage, y) points, no
    extrapolation. log(0) is -inf, so a zero neighbour gives 0 as in the
    reference implementation."""
    pts = sorted((x, y) for x, y in points if not math.isnan(x) and not math.isnan(y))
    if len(pts) < 2 or v < pts[0][0] or v > pts[-1][0]:
        return math.nan
    for (x0, y0), (x1, y1) in zip(pts, pts[1:]):
        if x0 <= v <= x1:
            if x1 == x0:
                return y0
            if y0 <= 0 or y1 <= 0:
                return 0.0
            t = (v - x0) / (x1 - x0)
            return math.exp(math.log(y0) * (1 - t) + math.log(y1) * t)
    return math.nan


def her_targets(agg: "OrderedDict[tuple, dict]") -> list[dict]:
    out = []
    for (sid, source, j), e in agg.items():
        if j != 50.0 or math.isnan(e["voltage"]):
            continue
        xrf = e["xrf comp"]
        out.append(
            OrderedDict(
                sample_id=sid,
                source=source,
                xrf_comp=xrf,
                elements=elements_of(xrf),
                target_comp=e["composition"],
                n_rows=e["n_rows"],
                replicate_sample_ids=";".join(e["sample ids"]),
                voltage_cell_v=f"{e['voltage']:.6f}",
                voltage_she_v=f"{get_she(e['voltage']):.6f}",
            )
        )
    return out


def co2rr_targets(agg: "OrderedDict[tuple, dict]") -> list[dict]:
    series: "OrderedDict[tuple, list[dict]]" = OrderedDict()
    for (sid, source, _j), e in agg.items():
        series.setdefault((sid, source), []).append(e)
    out = []
    for (sid, source), entries in series.items():
        volts = [e["voltage"] for e in entries if not math.isnan(e["voltage"])]
        if len(volts) < 2:
            continue
        xrf = entries[0]["xrf comp"]
        row = OrderedDict(
            sample_id=sid,
            source=source,
            xrf_comp=xrf,
            elements=elements_of(xrf),
            target_comp=entries[0]["composition"],
            n_points=len(volts),
            v_min=f"{min(volts):.4f}",
            v_max=f"{max(volts):.4f}",
            voltage_v=f"{INTERPOLATE_V:.1f}",
        )
        ok = True
        for p in PRODUCTS:
            fe = log_interp(INTERPOLATE_V, [(e["voltage"], e[p]) for e in entries])
            pr = log_interp(
                INTERPOLATE_V,
                [(e["voltage"], production_rate(e[p], ELECTRONS[p], e["current density"])) for e in entries],
            )
            if math.isnan(fe):
                ok = False
            row[p] = "" if math.isnan(fe) else f"{fe:.6f}"
            row[RATE_NAME[p]] = "" if math.isnan(pr) else f"{pr:.6e}"
        if not ok:
            # 3.3 V lies outside the measured range (no extrapolation) or a
            # product column is missing for this composition: not a target.
            continue
        out.append(row)
    return out


def write_csv(path: str, rows: list[dict]) -> None:
    if not rows:
        raise SystemExit(f"nothing to write for {path}")
    with open(path, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        w.writeheader()
        w.writerows(rows)


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--experiments", required=True, help="the bound experiments.csv")
    ap.add_argument("--xrf", required=True, help="the bound xrf.csv")
    ap.add_argument("--out-dir", required=True)
    args = ap.parse_args()

    exp_rows = read_csv(args.experiments)
    mapping = xrf_mapping(read_csv(args.xrf))
    os.makedirs(args.out_dir, exist_ok=True)

    her = her_targets(aggregate(base_rows(exp_rows, mapping, "HER")))
    co2 = co2rr_targets(aggregate(base_rows(exp_rows, mapping, "CO2R")))
    write_csv(os.path.join(args.out_dir, "targets_her.csv"), her)
    write_csv(os.path.join(args.out_dir, "targets_co2rr.csv"), co2)

    def systems(rows: list[dict]) -> int:
        return len({r["elements"] for r in rows})

    print(f"HER targets:   {len(her)} compositions, {systems(her)} element systems (LOCO folds)")
    print(f"CO2RR targets: {len(co2)} compositions, {systems(co2)} element systems (LOCO folds)")


if __name__ == "__main__":
    main()