#!/usr/bin/env python3
# -----------------------------------------------------------------------------
# Script developed by A. Saraiva-Souza, research scientist
# Description: Converts a CIF file to a simplified atoms.xyz and generates a
#              RESCU input file with preconfigured settings.
# -----------------------------------------------------------------------------

import argparse
import os
import numpy as np
from ase.io import read, write

# ------------------------------ Utils ---------------------------------
def sort_by_z(atoms):
    """Return a copy of the Atoms object sorted by the z coordinate."""
    positions = atoms.get_positions()
    order = np.argsort(positions[:, 2])
    return atoms[order]

def fmt_row(row):
    """Format a lattice-vector row in the desired style."""
    return f"[{row[0]:12.6f} {row[1]:12.6f} {row[2]:12.6f}        ];"

def build_latvec_block(cell):
    rows = [fmt_row(cell[0]), fmt_row(cell[1]), fmt_row(cell[2])]
    return "domain.latvec = [\n  " + "\n  ".join(rows) + "\n]"

def write_xyz(atoms, xyz_path):
    write(xyz_path, atoms, format="xyz")

def generate_rescu_input(
    base_name,
    atoms,
    xyz_name,
    pseudo_dir="./pseudos",
    pseudo_suffix="ONCV_PBE_DZP",
    kgrid=(9, 9, 1),
    sigma=0.05,
    lowres=0.3,
    boundary=(1, 1, 2),
):
    # ---------------- lattice ----------------
    cell = atoms.get_cell()
    latvec_block = build_latvec_block(cell)

    # ---------------- elements ----------------
    elements = sorted(set(atoms.get_chemical_symbols()))
    elem_lines = []
    for i, el in enumerate(elements, 1):
        matname = f"{el}_{pseudo_suffix}.mat"
        elem_lines.append(f"element({i}).species   = '{el}'")
        elem_lines.append(
            "element({i}).path      = '{path}'".format(
                i=i,
                path=matname if pseudo_dir in ['', '.', './']
                else os.path.join(pseudo_dir, matname)
            )
        )
    element_block = "\n".join(elem_lines)

    # ---------------- names/paths ----------------
    savepath = f"results/{base_name}_scf"
    outfile  = f"{base_name}_scf.out"

    # ---------------- assembled input ----------------
    txt = f"""## -------------------------------------------------------------
## MPI Configuration
## -------------------------------------------------------------
#mpi.status = true
#mpi.parak  = true
#smi.status = true

# -------------------------------------------------------------
# Calculation Options
# -------------------------------------------------------------
option.timeLimit       = 2
option.maxSCFiteration = 400
option.precision       = 'real-med'

# -------------------------------------------------------------
# Unit Settings
# -------------------------------------------------------------
units.length = 'Angstrom'
units.energy = 'eV'

# -------------------------------------------------------------
# Domain Settings (Lattice Vectors and Boundaries)
# -------------------------------------------------------------
{latvec_block}
domain.boundary = [{boundary[0]},{boundary[1]},{boundary[2]}]
domain.lowres   = {lowres}

# -------------------------------------------------------------
# Functional and LibXC Information
# -------------------------------------------------------------
functional.libxc = true % required for vdW
functional.list  = {{'XC_GGA_X_OPTPBE_VDW','XC_GGA_C_OP_PBE'}}

# -------------------------------------------------------------
# Atomic and Elemental Information
# -------------------------------------------------------------
{element_block}
atom.xyz = '{xyz_name}'

# ==============================================================
# ===================== SCF  ===================================
# ==============================================================

# -------------------------------------------------------------
# LCAO (Linear Combination of Atomic Orbitals) Configuration
# -------------------------------------------------------------
LCAO.status   = true
LCAO.mulliken = true

# -------------------------------------------------------------
# General Calculation Information
# -------------------------------------------------------------
info.calculationType = 'self-consistent'
info.savepath        = '{savepath}'
info.outfile         = '{outfile}'

# -------------------------------------------------------------
# K-Point Grid and Smearing
# -------------------------------------------------------------
kpoint.type  = 'MonkhorstPack'
kpoint.gridn = [{kgrid[0]},{kgrid[1]},{kgrid[2]}]
kpoint.sigma = {sigma}

# -------------------------------------------------------------
# Mixing Method for SCF Convergence
# -------------------------------------------------------------
mixing.method = 'broyden'
mixing.beta   = 0.05
mixing.type   = 'density'
mixing.tol    = [1e-5 1e-5]

# -------------------------------------------------------------
# Eigensolver Settings
# -------------------------------------------------------------
eigensolver.algo      = 'cfsi'
eigensolver.tol       = 5e-2
eigensolver.emptyBand = 36

# -------------------------------------------------------------
# Force Calculation Settings
# -------------------------------------------------------------
force.status = true
"""
    return txt

# ------------------------------ Main ----------------------------------
def main():
    parser = argparse.ArgumentParser(
        description="Convert a .cif into <basename>.xyz and <basename>_scf.input (RESCU)."
    )
    parser.add_argument("-i", "--input", required=True, help="Input .cif file")
    parser.add_argument("--pseudo-dir", default="./pseudos",
                        help="Directory containing pseudopotentials")
    parser.add_argument("--pseudo-suffix", default="ONCV_PBE_DZP",
                        help="Suffix used in pseudo files (e.g., ONCV_PBE_DZP)")
    parser.add_argument("--k", nargs=3, type=int, default=(9, 9, 1),
                        metavar=("KX", "KY", "KZ"), help="k-point grid")
    parser.add_argument("--sigma", type=float, default=0.05, help="Smearing (eV)")
    parser.add_argument("--lowres", type=float, default=0.3, help="domain.lowres")
    parser.add_argument("--boundary", nargs=3, type=int, default=(1, 1, 2),
                        metavar=("BX", "BY", "BZ"), help="domain.boundary")
    args = parser.parse_args()

    cif_path = args.input
    base_name = os.path.splitext(os.path.basename(cif_path))[0]

    # Read CIF
    atoms = read(cif_path)

    # Sort by z
    atoms = sort_by_z(atoms)

    # Save XYZ as <basename>.xyz
    xyz_name = f"{base_name}.xyz"
    write_xyz(atoms, xyz_name)
    print(f"✅ Created: {xyz_name} (z-sorted).")

    # Ensure results/ exists (many RESCU workflows expect it)
    os.makedirs("results", exist_ok=True)

    # Generate organized RESCU input
    input_txt = generate_rescu_input(
        base_name=base_name,
        atoms=atoms,
        xyz_name=xyz_name,
        pseudo_dir=args.pseudo_dir,
        pseudo_suffix=args.pseudo_suffix,
        kgrid=tuple(args.k),
        sigma=args.sigma,
        lowres=args.lowres,
        boundary=tuple(args.boundary),
    )

    input_name = f"{base_name}_scf.input"
    with open(input_name, "w") as f:
        f.write(input_txt)
    print(f"✅ Created: {input_name}")

if __name__ == "__main__":
    main()