#!/usr/bin/env python3
"""
Example: train a BatteryMHM state-of-health predictor on your own cells.

Replace `load_my_cells()` with your real data: for each cell you need an
early-cycle capacity curve (a 1-D array of per-cycle discharge capacities)
and a target SOH label (e.g. capacity retention at end of life).

Run:  python examples/predict_soh.py
"""

import numpy as np

from batterymhm import MHMEnsemble, compute_metrics, mhm_full_features, seq_to_harmonics


def load_my_cells():
    """
    Stand-in data loader — returns (curves, soh_labels).
    Swap this out for your measured cells (e.g. parsed from a cycler export).
    """
    rng = np.random.default_rng(7)
    curves, labels = [], []
    for _ in range(120):
        rate = rng.uniform(0.001, 0.004)
        cap = 1.0 - rate * np.sqrt(np.arange(1, 61)) + rng.normal(0, 0.0015, 60)
        curves.append(np.clip(cap, 0, 1.05))
        labels.append(float(cap[-1]))
    return curves, np.asarray(labels)


def featurize(curves):
    """Quantise each curve to HINs and build the MHM descriptor matrix."""
    dicts = [mhm_full_features(seq_to_harmonics(list(c), bins=9)) for c in curves]
    keys = sorted(dicts[0])                       # stable column order
    X = np.array([[d[k] for k in keys] for d in dicts], dtype=float)
    return X, keys


def main():
    curves, y = load_my_cells()
    X, keys = featurize(curves)

    n = len(y)
    cut = int(0.75 * n)
    model = MHMEnsemble().fit(X[:cut], y[:cut], feature_names=keys)
    pred = model.predict(X[cut:])

    print("Held-out performance:", compute_metrics(y[cut:], pred))
    print("\nMost informative MHM features:")
    for name, imp in model.top_features(8):
        print(f"  {imp:6.4f}  {name}")


if __name__ == "__main__":
    main()