#!/usr/bin/env python
# coding: utf-8

# ## Battery Cathodes: Intercalation Potential and Volume Change
# 
# Cathode materials are often judged by more than their voltage alone. A useful battery cathode should store lithium at a favorable potential, but it should also remain structurally stable during charge and discharge. Large changes in unit-cell volume can lead to mechanical stress, cracking, and capacity fade over repeated cycling.
# 
# In this notebook we use the inorganic-materials UMA model, UMA-S-1.2-OMat, to optimize olivine LiFePO4 and its delithiated FePO4 counterpart. From these calculations we estimate two useful screening descriptors: the average lithium intercalation potential and the relative volume change upon delithiation.

# ## Overview
# 
# The workflow follows the reaction
# 
# $$\mathrm{LiFePO_4 \rightarrow FePO_4 + Li}$$
# 
# and evaluates it with a machine-learning potential. We will:
# 
# - optimize bulk Li to obtain a consistent Li metal reference energy
# - optimize lithiated LiFePO4 and delithiated FePO4
# - compare the optimized structures visually and through their cell volumes
# - compute the average intercalation potential from the total energies
# - compute the relative volume change per formula unit and compare with experiment

from math import gcd

from ase.build import bulk
from scm.plams import AMSJob, Settings, Units, fromASE, view, plot_image_grid


# ## Settings and Visualization Helpers
# 
# We use the UMA-S-1.2-OMat machine-learning potential with a geometry optimization and lattice relaxation so that both atomic positions and unit-cell parameters can respond to lithiation state. The D4 dispersion add-on is also consistently included in all calculations.
# 
# This section also defines a few small helper functions for structure handling, formula counting, and visualization.

def umat_settings(nproc=1):
    s = Settings()
    s.input.ams.Task = "GeometryOptimization"
    s.input.ams.GeometryOptimization.OptimizeLattice = "Yes"
    s.input.ams.EngineAddons.D4Dispersion.Enabled = "Yes"
    s.input.mlpotential.Model = "UMA-S-1.2-OMat"
    s.runscript.nproc = nproc
    s.runscript.preamble_lines = ["export OMP_NUM_THREADS=1"]
    return s


def formula_units(system):
    counts = [count for count in system.get_formula(as_dict=True).values() if count > 0]
    nfu = counts[0]
    for count in counts[1:]:
        nfu = gcd(nfu, count)
    return nfu


def render(system, supercell=(3, 3, 1), direction="along_y", width=400, height=400, padding=-1):
    m = system.copy().supercell(*supercell)
    return view(
        m,
        direction=direction,
        width=width,
        height=height,
        padding=padding,
        fixed_atom_size=False,
        show_atom_labels=False,
    )


def make_lifepo4():
    mol = AMSJob.from_input(
        """
System
    Atoms
        Li 0.0 0.0 0.0 
        Li 2.327458595 2.98537755 5.118098025 
        Li 2.327458595 0.0 5.118098025 
        Li 0.0 2.98537755 0.0 
        Fe 2.466481095 1.492688775 7.9960175444 
        Fe 0.1390225 4.478066325 7.3582765306 
        Fe 4.51589469 1.492688775 2.8779195194 
        Fe 2.188436095 4.478066325 2.2401785056 
        P 1.9486533972 1.492688775 0.9608338493 
        P 4.2761119922 4.478066325 4.1572641757 
        P 0.3788051978 1.492688775 6.0789318743 
        P 2.7062637928 4.478066325 9.2753622007 
        O 3.466919761 1.492688775 0.964563612 
        O 3.3223789033 3.25739339 8.5385692027 
        O 3.3223789033 5.69873926 8.5385692027 
        O 1.187997429 4.478066325 9.271632437999999 
        O 3.3043602309 4.478066325 0.4535511069 
        O 1.139461166 4.478066325 4.153534413 
        O 1.3325382867 0.27201584 1.6976268473 
        O 0.9949203083 2.71336171 6.8157248723 
        O 3.6599968817 3.25739339 3.4204711777 
        O 3.6599968817 5.69873926 3.4204711777 
        O 1.3505569591 1.492688775 9.782644943099999 
        O 0.9769016359 1.492688775 4.6645469181 
        O 3.6780155541 4.478066325 5.5716491319 
        O 3.515456024 1.492688775 6.082661637 
        O 0.9949203083 0.27201584 6.8157248723 
        O 1.3325382867 2.71336171 1.6976268473 
    End
    Lattice
        4.65491719 0.0 0.0
        0.0 5.9707551 0.0
        0.0 0.0 10.23619605
    End
End
        """
    ).molecule 
    
    if isinstance(mol, dict) and len(mol) == 1 and "" in mol:
        return mol[""]
    else:
        return mol


# ## Lithium Metal Reference
# 
# To convert total-energy differences into an intercalation voltage, we need a reference for metallic lithium. We therefore optimize bulk bcc Li and extract the energy per Li atom.

li_bulk = fromASE(bulk("Li", "bcc", a=3.49, cubic=True))
li_job = AMSJob(name="Li_bcc_opt", molecule=li_bulk, settings=umat_settings())
li_result = li_job.run()

li_bulk_opt = li_result.get_main_molecule()
n_li = li_bulk_opt.get_formula(as_dict=True)["Li"]
e_li_metal = li_result.get_energy(unit="hartree") / n_li
v_li_bulk = li_bulk_opt.unit_cell_volume()

print(f"Li metal reference energy: {e_li_metal:.8f} Ha/atom")
print(f"Li bcc cell volume: {v_li_bulk:.5f} A^3")


plot_image_grid(
    {"optimized Li metal": render(li_bulk_opt, direction="along_c", width=420, height=420, padding=0)}, 
    figsize=(4, 4)
, save_path="picture1.png");


# ## Lithiated LiFePO4 Optimization
# 
# We define the olivine LiFePO4 unit cell directly in the notebook and optimize the fully lithiated cathode. The resulting total energy is one of the three quantities needed for the voltage calculation, and the relaxed cell volume provides the reference volume for the charged-state comparison.

lfp = make_lifepo4()
print(lfp.get_formula(False))


lfp_job = AMSJob(name="LiFePO4_opt", molecule=lfp, settings=umat_settings())
lfp_result = lfp_job.run()

lfp_opt = lfp_result.get_main_molecule()
e_lfp = lfp_result.get_energy(unit="hartree")
v_lfp = lfp_opt.unit_cell_volume()

print(f"LiFePO4 total energy: {e_lfp:.8f} Ha")
print(f"LiFePO4 cell volume: {v_lfp:.5f} A^3")
print(lfp_opt.get_formula(False))


plot_image_grid({
    "initial LiFePO4": render(lfp),
    "optimized LiFePO4": render(lfp_opt),
}, save_path="picture2.png");


# ## Delithiated FePO4
# 
# To represent delithiation, we remove all Li atoms from the optimized LiFePO4 structure and use the resulting FePO4 framework as the starting point for a second optimization.
# 
# This mirrors the structural change associated with charge extraction and lets us compare the two end members on the same footing.

fp = lfp_opt.copy()
for atom in list(fp):
    if atom.symbol == "Li":
        fp.delete_atom(atom)

print(fp.get_formula(False))


fp_job = AMSJob(name="FePO4_opt", molecule=fp, settings=umat_settings())
fp_result = fp_job.run()

fp_opt = fp_result.get_main_molecule()
e_fp = fp_result.get_energy(unit="hartree")
v_fp = fp_opt.unit_cell_volume()

print(f"FePO4 total energy: {e_fp:.8f} Ha")
print(f"FePO4 cell volume: {v_fp:.5f} A^3")
print(fp_opt.get_formula(False))


plot_image_grid({
    "optimized LiFePO4": render(lfp_opt),
    "optimized FePO4": render(fp_opt),
}, save_path="picture3.png");


# ## Analysis
# 
# With optimized energies and cell volumes for Li, LiFePO4, and FePO4 in hand, we can evaluate two useful cathode descriptors.
# 
# The average intercalation potential is obtained from the reaction energy, while the relative volume change measures how strongly the crystal expands or contracts on delithiation. For meaningful comparison across cells, both quantities are reduced to a per-formula-unit basis where needed.

n_formula_units = formula_units(lfp_opt)

e_lfp_fu = e_lfp / n_formula_units
e_fp_fu = e_fp / n_formula_units
v_lfp_fu = v_lfp / n_formula_units
v_fp_fu = v_fp / n_formula_units

voltage = Units.convert(e_fp_fu + e_li_metal - e_lfp_fu, "hartree", "eV")
delta_v_percent = 100.0 * (v_fp_fu - v_lfp_fu) / v_lfp_fu

print(f"LiFePO4 energy per formula unit: {e_lfp_fu:.8f} Ha")
print(f"FePO4 energy per formula unit: {e_fp_fu:.8f} Ha")
print(f"LiFePO4 volume per formula unit: {v_lfp_fu:.5f} A^3")
print(f"FePO4 volume per formula unit: {v_fp_fu:.5f} A^3")
print(f"Average intercalation potential: {voltage:.3f} V")
print(f"Relative volume change: {delta_v_percent:.3f} %")


# ## Comparison to Experiment
# 
# A quick comparison to representative literature values helps place the calculation in context. Agreement at this level is useful as a sanity check before extending the same workflow to doped olivines, sodium analogues, or larger screening studies.

import pandas as pd

experimental = {
    "LiFePO4 volume (A^3/f.u.)": 72.85,
    "FePO4 volume (A^3/f.u.)": 68.09,
    "Volume change (%)": -6.534,
    "Intercalation potential (V)": 3.5,
}

calculated = {
    "LiFePO4 volume (A^3/f.u.)": v_lfp_fu,
    "FePO4 volume (A^3/f.u.)": v_fp_fu,
    "Volume change (%)": delta_v_percent,
    "Intercalation potential (V)": voltage,
}

df = pd.DataFrame([experimental, calculated], index=["Experimental", "Calculated"])
df.loc["Delta"] = df.loc["Calculated"] - df.loc["Experimental"]
df = df.round(3)


print(df)


# ## Extensions
# 
# The same workflow can be generalized in several directions:
# 
# - replace Fe with a dopant and examine the trade-off between voltage and structural change
# - adapt the workflow to a sodium analogue and compare Na insertion with Li insertion
# - loop over multiple cathode materials and collect a benchmark table automatically