#!/usr/bin/env python
# make_beo_stacking_xyz.py

import argparse
import math
import sys
from pathlib import Path

import numpy as np


"""
python make_beo_stacking_xyz.py --cif BeO.cif --nx 3 --ny 3 --stacking ABABABCBCBC --be-z-offset 2
python make_beo_stacking_xyz.py --stacking ABCABCABC --outfile beo_zb111.xyz
python make_beo_stacking_xyz.py --stacking ABABABCBCB --outfile beo_fault.xyz
"""

def try_read_beo_params_from_cif(cif_path):
    try:
        from pymatgen.core import Structure
    except Exception:
        print("ERROR: pymatgen is required for --cif")
        print("pip install pymatgen")
        input("\nPress ENTER to terminate>>\n")
        sys.exit(1)

    st = Structure.from_file(cif_path)
    a = float(st.lattice.a)
    c = float(st.lattice.c)

    be_indices = [i for i, s in enumerate(st) if s.specie.symbol.lower() == "be"]
    o_indices = [i for i, s in enumerate(st) if s.specie.symbol.lower() == "o"]

    if not be_indices or not o_indices:
        raise ValueError("CIF must contain Be and O atoms.")

    # nearest Be-O distance
    dmin = 1.0e30
    for i in be_indices:
        for j in o_indices:
            d = st.get_distance(i, j)
            dmin = min(dmin, d)

    # O close-packed layer spacing along c
    oz = sorted([st[i].frac_coords[2] % 1.0 for i in o_indices])
    oz_unique = []
    for z in oz:
        if not oz_unique or abs(z - oz_unique[-1]) > 1.0e-4:
            oz_unique.append(z)

    if len(oz_unique) >= 2:
        diffs = []
        for i in range(len(oz_unique)):
            z1 = oz_unique[i]
            z2 = oz_unique[(i + 1) % len(oz_unique)]
            dzf = (z2 - z1) % 1.0
            if dzf > 1.0e-5:
                diffs.append(dzf)
        dz = min(diffs) * c
    else:
        dz = 0.5 * c

    return a, c, dmin, dz


def make_triangular_layer(a, nx, ny):
    """
    2D triangular lattice patch.
    """
    a1 = np.array([a, 0.0])
    a2 = np.array([0.5 * a, 0.5 * math.sqrt(3.0) * a])

    xy = []
    for i in range(nx):
        for j in range(ny):
            r = i * a1 + j * a2
            xy.append(r)

    return np.array(xy), a1, a2


def stacking_shift(label, a1, a2):
    """
    A/B/C close-packed layer registry shifts.
    """
    label = label.upper()
    if label == "A":
        return np.array([0.0, 0.0])
    elif label == "B":
        return (a1 + a2) / 3.0
    elif label == "C":
        return 2.0 * (a1 + a2) / 3.0
    else:
        raise ValueError(f"Unknown stacking label: {label}")


def write_xyz(path, atoms, comment):
    with open(path, "w", encoding="utf-8") as f:
        f.write(f"{len(atoms)}\n")
        f.write(comment + "\n")
        for elem, x, y, z in atoms:
            f.write(f"{elem:2s} {x:15.8f} {y:15.8f} {z:15.8f}\n")


def main():
    parser = argparse.ArgumentParser(
        description="Generate BeO-like WZ/ZB/stacking-fault XYZ model."
    )

    parser.add_argument("--cif", type=str, default=None,
                        help="BeO CIF file. If given, a, c, d, dz are estimated.")
    parser.add_argument("--a", type=float, default=2.70,
                        help="In-plane triangular lattice constant in Angstrom.")
    parser.add_argument("--c", type=float, default=4.38,
                        help="Wurtzite c lattice constant in Angstrom.")
    parser.add_argument("--d", type=float, default=None,
                        help="Be-O distance. Used mainly for comment/output.")
    parser.add_argument("--dz", type=float, default=2.0,
                        help="A-B close-packed layer spacing in Angstrom. Default: c/2.")
    parser.add_argument("--be-z-offset", type=float, default=0.55,
                        help="Be z offset from O layer in Angstrom.")

    parser.add_argument("--stacking", type=str, default="ABABAB",
                        help='Stacking sequence, e.g. "ABABAB", "ABCABC", "ABABABCBCB".')
    parser.add_argument("--nx", type=int, default=8,
                        help="Number of in-plane cells along a1.")
    parser.add_argument("--ny", type=int, default=8,
                        help="Number of in-plane cells along a2.")
    parser.add_argument("--center", type=int, default=1, choices=[0, 1],
                        help="Center xy coordinates: 1=yes, 0=no.")
    parser.add_argument("--outfile", type=str, default="beo_stacking.xyz",
                        help="Output XYZ file.")

    args = parser.parse_args()

    if args.cif is not None:
        a_cif, c_cif, d_cif, dz_cif = try_read_beo_params_from_cif(args.cif)
        args.a = a_cif
        args.c = c_cif
        if args.d is None:
            args.d = d_cif
        if args.dz is None:
            args.dz = dz_cif

    if args.d is None:
        args.d = 1.65

    if args.dz is None:
        args.dz = 0.5 * args.c

    base_xy, a1, a2 = make_triangular_layer(args.a, args.nx, args.ny)

    if args.center == 1:
        base_xy = base_xy - base_xy.mean(axis=0)

    atoms = []

    for k, label in enumerate(args.stacking.upper()):
        shift = stacking_shift(label, a1, a2)
        z_o = k * args.dz
        z_be = z_o + args.be_z_offset

        for xy in base_xy:
            x, y = xy + shift
            atoms.append(("O", x, y, z_o))
            atoms.append(("Be", x, y, z_be))

    comment = (
        f"BeO stacking={args.stacking}; "
        f"a={args.a:.6g} A, c={args.c:.6g} A, "
        f"d(Be-O)~{args.d:.6g} A, dz={args.dz:.6g} A, "
        f"be_z_offset={args.be_z_offset:.6g} A"
    )

    write_xyz(args.outfile, atoms, comment)

    print("Output:", args.outfile)
    print("Stacking:", args.stacking)
    print(f"a  = {args.a:.6f} A")
    print(f"c  = {args.c:.6f} A")
    print(f"d  = {args.d:.6f} A")
    print(f"dz = {args.dz:.6f} A")
    print(f"Be z offset = {args.be_z_offset:.6f} A")
    print(f"Atoms = {len(atoms)}")


if __name__ == "__main__":
    main()