#!/usr/bin/env python3
"""
Example: build the MHM harmonic descriptor for a crystal composition.

Given a list of element symbols (the atoms in a formula unit), each element is
folded to a Harmonic Identity Number (HIN) by its atomic number, and the
neighbour pairs are scored through the 8×8 matter Chi matrix plus the HIN-9
energy channel. The resulting descriptor can be concatenated with any other
composition features and fed to a regressor for property prediction.

Run:  python examples/materials_descriptor.py
"""

from batterymhm import chi_score, element_hin, mhm_matter8_neighbor_histograms

CATHODES = {
    "LiFePO4": ["Li", "Fe", "P", "O", "O", "O", "O"],
    "LiCoO2":  ["Li", "Co", "O", "O"],
    "LiNi0.8Mn0.1Co0.1O2 (NMC811)": ["Li", "Ni", "Mn", "Co", "O", "O"],
    "NaFePO4": ["Na", "Fe", "P", "O", "O", "O", "O"],
}


def describe(name, elements):
    hins = [element_hin(e) for e in elements]
    # Treat the formula's atoms as a neighbour list paired with itself.
    feats = mhm_matter8_neighbor_histograms(hins, hins)
    active = sum(1 for v in feats.values() if abs(v) > 1e-12)
    # A couple of human-readable Chi compatibilities between the cation and anion.
    cation, anion = elements[0], "O" if "O" in elements else elements[-1]
    chi = chi_score(element_hin(cation), element_hin(anion))
    print(f"{name}")
    print(f"   element HINs        : {dict(zip(elements, hins))}")
    print(f"   descriptor size     : {len(feats)} features ({active} active)")
    print(f"   Chi({cation}–{anion}) compatibility : {chi:.3f}\n")


def main():
    print("MHM harmonic descriptors for common cathode chemistries:\n")
    for name, elements in CATHODES.items():
        describe(name, elements)


if __name__ == "__main__":
    main()