#!/usr/bin/env python3
"""Unrelaxed organic adsorption probes on alpha-corundum (0001).

Cases: methanol on bare Al-terminated corundum, methanol on a hydroxylated
surface, and phenol on bare corundum. This is an execution/scale probe, not
a fitted or independently validated adsorption-energy dataset. Finite
clusters retain exposed edges. Periodic slabs have true xy periodicity;
GFN2 is used only on finite clusters, MACE can also evaluate the slabs.

Generate or evaluate one component at a time in an external output folder:
    python 39_corundum_adsorption.py --output /path/to/results --prepare-only
    python 39_corundum_adsorption.py --output /path/to/results --case bare-methanol
    python 39_corundum_adsorption.py --output /path/to/results --case bare-phenol \
        --repeat 2 --method mace --boundary slab

The bulk motif matches examples/periodic/_systems.py. The hydroxylated
motif adds one dissociated H2O per top Al: Al-OH plus a proton on a nearby
substrate O. This follows the chemistry of Ahmad & Martsinovich (2023),
Sec. 2, doi:10.1098/rsta.2022.0254; coordinates here are independently
constructed starting geometries, not that paper's optimized structures.
Bare methanol's O-to-Al coordination is motivated by Borck & Schroder,
arXiv:cond-mat/0307039. No literature adsorption energy is used as a target.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
import time

import numpy as np
from ase import Atom as ASEAtom, Atoms
from ase.build import molecule, surface
from ase.spacegroup import crystal

CASES = ('bare-methanol', 'hydroxylated-methanol', 'bare-phenol')
BOHR_ANGSTROM = 0.529177210903


def make_surface(repeat=1, hydroxylated=False):
    """Stoichiometric 30-atom conventional-cell slab, then xy replication."""
    if type(repeat) is not int or repeat < 1:
        raise ValueError('repeat must be a positive integer')
    bulk = crystal(['Al', 'O'], basis=[(0, 0, .35216), (.30624, 0, .25)],
                   spacegroup=167, cellpar=[4.7589, 4.7589, 12.991, 90, 90, 120])
    slab = surface(bulk, (0, 0, 1), 1, vacuum=12.0) * (repeat, repeat, 1)
    if hydroxylated:
        top = slab.positions[:, 2].max()
        sites = [i for i, a in enumerate(slab) if a.symbol == 'Al' and abs(a.z-top) < 1e-6]
        oxygen = [i for i, a in enumerate(slab) if a.symbol == 'O' and top-a.z < 1.0]
        used = set()
        for i in sites:
            # Distinct nearest subsurface oxygen for each top Al. Minimum
            # image distances define the periodic motif before the finite cut.
            candidates = [j for j in oxygen if j not in used]
            j = min(candidates, key=lambda j: slab.get_distance(i, j, mic=True))
            used.add(j)
            p = slab.positions[i]
            slab += ASEAtom('O', p + [0, 0, 1.8])
            slab += ASEAtom('H', p + [0, 0, 2.77])
            slab += ASEAtom('H', slab.positions[j] + [0, 0, .97])
    return slab


def make_adsorbate(name):
    """Return a molecule with its O anchor at zero and carbon above it."""
    if name == 'methanol':
        ads = molecule('CH3OH')
        oxygen = next(i for i, a in enumerate(ads) if a.symbol == 'O')
        carbon = next(i for i, a in enumerate(ads) if a.symbol == 'C')
        ads.translate(-ads.positions[oxygen].copy())
        ads.rotate(ads.positions[carbon], [0, 0, 1], center=(0, 0, 0))
        return ads, oxygen
    if name != 'phenol':
        raise ValueError('unknown adsorbate')
    # Ideal planar starting geometry: regular 1.40 A aromatic ring,
    # C-O 1.36 A, C-H 1.09 A, O-H .97 A; not an optimized conformer.
    angles = np.arange(6) * np.pi / 3
    ring = np.column_stack((1.40*np.cos(angles), 1.40*np.sin(angles), np.zeros(6)))
    hydrogen = np.column_stack((2.49*np.cos(angles[1:]), 2.49*np.sin(angles[1:]), np.zeros(5)))
    oxygen = np.array([2.76, 0., 0.])
    hydroxyl_h = oxygen + .97*np.array([.5, np.sqrt(3)/2, 0.])
    ads = Atoms('C6H5OH', positions=np.vstack((ring, hydrogen, oxygen, hydroxyl_h)))
    ads.translate(-oxygen)
    ads.rotate([-1, 0, 0], [0, 0, 1], center=(0, 0, 0))
    return ads, 11


def build_case(case='bare-methanol', repeat=1):
    """Return combined/surface/molecule with identical fragment geometries."""
    if case not in CASES:
        raise ValueError(f'case must be one of {CASES}')
    hydroxylated = case == 'hydroxylated-methanol'
    slab = make_surface(repeat, hydroxylated)
    organic, oxygen = make_adsorbate('phenol' if case == 'bare-phenol' else 'methanol')
    al = [i for i, a in enumerate(slab) if a.symbol == 'Al']
    top_z = max(slab.positions[i, 2] for i in al)
    sites = [i for i in al if abs(slab.positions[i, 2]-top_z) < 1e-6]
    center = .5*(slab.cell[0]+slab.cell[1])
    site = min(sites, key=lambda i: np.linalg.norm(slab.positions[i, :2]-center[:2]))
    # Bare O-Al contact 2.1 A. On the hydroxylated surface, O accepts
    # a hydrogen bond at 1.8 A above the outward surface OH hydrogen.
    height = 1.8 + .97 + 1.8 if hydroxylated else 2.1
    organic.translate(slab.positions[site] + [0, 0, height])
    combined = slab + organic
    combined.set_cell(slab.cell); combined.set_pbc([True, True, False])
    organic.set_cell(slab.cell); organic.set_pbc(False)
    for atoms in (slab, organic, combined):
        distances = atoms.get_all_distances(mic=bool(atoms.pbc.any()))
        np.fill_diagonal(distances, np.inf)
        if distances.min() < .65:
            raise ValueError('constructed geometry has an atomic contact below 0.65 Angstrom')
    return {'combined': combined, 'surface': slab, 'molecule': organic}, {
        'case': case, 'repeat_xy': repeat, 'surface_atoms': len(slab),
        'adsorbate_atoms': len(organic), 'anchor_index': len(slab)+oxygen,
        'surface_site_index': site, 'relaxed': False,
        'hydroxylation': 'one dissociated H2O per top Al' if hydroxylated else 'bare',
        'finite_edge_passivation': 'none', 'charge': 0, 'multiplicity': 1,
    }


def to_molecule(atoms):
    from vibeqc import Atom, Molecule
    return Molecule([Atom(int(z), p/BOHR_ANGSTROM) for z, p in
                     zip(atoms.numbers, atoms.positions)], 0, 1)


def to_periodic_system(atoms):
    """Construct this example's xy-periodic slab without the 3D ASE bridge."""
    from vibeqc import PeriodicSystem
    if atoms.pbc.tolist() != [True, True, False]:
        raise ValueError('Corundum slab requires xy periodicity')
    return PeriodicSystem(2, atoms.cell.array.T/BOHR_ANGSTROM,
                          to_molecule(atoms).atoms, charge=0, multiplicity=1)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--case', choices=CASES, default='bare-methanol')
    parser.add_argument('--repeat', type=int, default=1)
    parser.add_argument('--component', choices=['combined', 'surface', 'molecule'], default='combined')
    parser.add_argument('--boundary', choices=['finite', 'slab'], default='finite')
    parser.add_argument('--method', choices=['gfn2_xtb', 'mace'], default='gfn2_xtb')
    parser.add_argument('--weight-sha256', help='Expected MACE-MPA-0 checkpoint digest')
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--prepare-only', action='store_true')
    args = parser.parse_args()
    if args.boundary == 'slab' and args.method != 'mace' and not args.prepare_only:
        parser.error('this example uses GFN2 only for finite clusters; use mace for slabs')
    if args.weight_sha256 and args.method != 'mace':
        parser.error('--weight-sha256 requires --method mace')
    from vibeqc.output import write
    from vibeqc.output.formats.xyz import write_xyz
    components, metadata = build_case(args.case, args.repeat)
    args.output.mkdir(parents=True, exist_ok=True)
    for name, atoms in components.items():
        pbc = 'T T F' if args.boundary == 'slab' and name != 'molecule' else 'F F F'
        lattice = ' '.join(format(x, '.15g') for x in atoms.cell.array.ravel())
        write_xyz(args.output/f'{args.case}-{args.repeat}-{name}', to_molecule(atoms),
                  comment=f'Lattice="{lattice}" Properties=species:S:1:pos:R:3 pbc="{pbc}" unrelaxed=true')
    write(json.dumps(metadata, sort_keys=True)+'\n')
    if args.prepare_only:
        return
    atoms = components[args.component]; start = time.monotonic()
    # Keep the prepared extended XYZ inputs intact when the job writer emits
    # its own final-geometry files beside the calculation result.
    result_stem = args.output/f'{args.case}-{args.repeat}-{args.component}-result'
    options = {}
    if args.method == 'mace':
        from vibeqc.mlip import MLIPOptions
        options['mlip_options'] = MLIPOptions(
            model='medium-mpa-0', device='cpu', dtype='float64',
            weight_sha256=args.weight_sha256,
        )
    if args.boundary == 'slab' and args.component != 'molecule':
        from vibeqc import run_periodic_job
        result = run_periodic_job(
            to_periodic_system(atoms), None, method='mace',
            output=result_stem,
            output_qvf=True, record_hostname=False, **options,
        )
        identity = getattr(result, 'weight_identity', None)
    else:
        from vibeqc import run_job
        result = run_job(to_molecule(atoms), method=args.method,
                         output=result_stem,
                         record_hostname=False, **options)
        identity = getattr(result, 'weight_identity' if args.method == 'mace'
                           else 'parameter_provenance', None)
    gradient = np.asarray(result.gradient())
    write(json.dumps({'energy_ha':float(result.energy), 'gradient_ha_per_bohr':gradient.tolist(),
                      'max_force_ha_per_bohr':float(np.linalg.norm(gradient, axis=1).max()),
                      'seconds':time.monotonic()-start, 'identity':identity,
                      'relaxed':False, 'accuracy_qualified':False}, sort_keys=True)+'\n')


if __name__ == '__main__':
    main()