#!/usr/bin/env python3
"""
BatteryMHM — self-contained demonstration that the Miller Harmonic Method
(MHM) genuinely carries predictive signal, with NO external data downloads
and NO pre-trained weights.

It does two things:

  1. CELL-HEALTH DEMO (synthetic).
     Generates synthetic lithium-ion capacity-fade curves, each driven by a
     hidden degradation rate. From only an EARLY observation window of each
     curve, it quantises the capacity trace into harmonic identity numbers
     (HINs), extracts the full MHM descriptor, trains the MHM ensemble, and
     predicts each cell's EVENTUAL capacity retention. The MHM model is then
     compared against a mean-predictor baseline — if MHM carries signal, it
     beats the baseline by a wide margin.

  2. MATERIALS DESCRIPTOR DEMO (real element symbols).
     Builds the MHM neighbour/Chi descriptor for real cathode compositions,
     showing the materials feature path runs end-to-end.

This demo does NOT reproduce the published benchmark numbers (those are on
the real public Severson and Matbench datasets — see the README). It exists
so anyone can verify, in seconds and offline, that the open method works.

Run:  python demo.py
"""

import numpy as np

from batterymhm import (
    MHMEnsemble,
    compute_metrics,
    element_hin,
    mhm_full_features,
    mhm_matter8_neighbor_histograms,
    seq_to_harmonics,
)

# ─────────────────────────────────────────────────────────────────────────────
# 1. Synthetic cell-health demonstration
# ─────────────────────────────────────────────────────────────────────────────

def make_synthetic_cells(n_cells=200, total_cycles=200, obs_frac=0.30, seed=42):
    """
    Build synthetic capacity-fade curves with a hidden per-cell degradation
    rate. Returns (observed_curves, eventual_retention).

      observed_curves[i] : capacity vs cycle over the first obs_frac of life
      eventual_retention[i] : capacity / initial at total_cycles (the target)
    """
    rng = np.random.default_rng(seed)
    obs_len = max(8, int(total_cycles * obs_frac))
    cycles = np.arange(1, total_cycles + 1)

    observed, target = [], []
    for _ in range(n_cells):
        # Hidden degradation drivers (the "manufacturing + protocol" signature)
        rate = rng.uniform(0.0008, 0.0040)      # sqrt-fade rate
        knee = rng.uniform(0.0, 1.2e-5)         # late-life acceleration
        c0 = rng.uniform(0.99, 1.01)            # initial capacity spread

        cap = c0 - rate * np.sqrt(cycles) - knee * cycles ** 2
        cap = cap + rng.normal(0.0, 0.0015, size=cap.shape)   # measurement noise
        cap = np.clip(cap, 0.0, 1.05)

        observed.append(cap[:obs_len])
        target.append(float(cap[-1] / c0))      # eventual retention (SOH)

    return observed, np.asarray(target)


def featurize_curve(curve):
    """Quantise a capacity curve to HINs and compute the MHM descriptor."""
    hins = seq_to_harmonics(list(curve), bins=9)
    feats = mhm_full_features(hins)
    return feats


def run_cell_demo():
    print("=" * 70)
    print("1. CELL-HEALTH DEMO  —  predict eventual retention from early cycles")
    print("=" * 70)

    curves, y = make_synthetic_cells()
    feat_dicts = [featurize_curve(c) for c in curves]

    keys = sorted(feat_dicts[0].keys())                 # stable column order
    X = np.array([[d[k] for k in keys] for d in feat_dicts], dtype=float)
    print(f"   cells              : {X.shape[0]}")
    print(f"   MHM features/cell  : {X.shape[1]}")

    # 75 / 25 split
    rng = np.random.default_rng(0)
    idx = rng.permutation(len(y))
    cut = int(0.75 * len(y))
    tr, te = idx[:cut], idx[cut:]

    model = MHMEnsemble(use_stacking=False)             # ET (+XGB if installed)
    model.fit(X[tr], y[tr], feature_names=keys)
    pred = model.predict(X[te])

    m = compute_metrics(y[te], pred)
    baseline_mae = float(np.mean(np.abs(y[te] - y[tr].mean())))

    print(f"\n   MHM ensemble       : {m}")
    print(f"   mean-predictor MAE : {baseline_mae:.4f}")
    factor = baseline_mae / max(m.mae, 1e-9)
    print(f"   → MHM is {factor:.1f}× better than predicting the mean.")
    print("   Top MHM features   :")
    for name, imp in model.top_features(6):
        print(f"       {imp:6.4f}  {name}")
    return m.mae < baseline_mae * 0.6                   # demo passes if clearly better


# ─────────────────────────────────────────────────────────────────────────────
# 2. Materials descriptor demonstration
# ─────────────────────────────────────────────────────────────────────────────

def run_materials_demo():
    print("\n" + "=" * 70)
    print("2. MATERIALS DESCRIPTOR DEMO  —  MHM Chi/HIN features on real cathodes")
    print("=" * 70)

    # (formula, element list) for a few well-known cathode chemistries.
    cathodes = {
        "LiFePO4":  ["Li", "Fe", "P", "O", "O", "O", "O"],
        "LiCoO2":   ["Li", "Co", "O", "O"],
        "LiNiMnCoO2 (NMC)": ["Li", "Ni", "Mn", "Co", "O", "O"],
        "NaFePO4":  ["Na", "Fe", "P", "O", "O", "O", "O"],
    }

    for name, elements in cathodes.items():
        hins = [element_hin(e) for e in elements]
        # Treat the site list as a neighbour list paired with itself (demo).
        feats = mhm_matter8_neighbor_histograms(hins, hins)
        nonzero = sum(1 for v in feats.values() if abs(v) > 1e-12)
        print(f"   {name:20s}  HINs={hins}  "
              f"descriptor={len(feats)} features ({nonzero} active)")

    print("   → Materials feature path runs end-to-end on real compositions.")
    return True


if __name__ == "__main__":
    ok1 = run_cell_demo()
    ok2 = run_materials_demo()
    print("\n" + "=" * 70)
    print("RESULT:", "PASS — the open method runs and carries signal."
          if (ok1 and ok2) else "CHECK — see output above.")
    print("=" * 70)