#!/usr/bin/env python3
"""Convert one structure file (CIF/POSCAR/xyz/...) into the inputs of
VASP, Gaussian, CP2K and OpenMX with ASE (+ spglib for --primitive).

  python cif_to_all.py Si.cif                 # keep the cell as read
  python cif_to_all.py Si.cif --primitive     # reduce to the primitive cell
  python cif_to_all.py Si.cif -o out --valence Si=4

Files written to the output directory (default: converted/):
  POSCAR                VASP 5 format, fractional ("Direct") coordinates
  <name>.xyz            plain xyz (NO cell information!)
  <name>.extxyz         extended xyz (keeps Lattice and pbc)
  <name>.gjf            Gaussian input; periodic cells get TV lines (PBC)
  <name>_cp2k.inc       CP2K &CELL + &COORD blocks for @INCLUDE in &SUBSYS
  <name>_openmx.dat     OpenMX Atoms.* blocks (fractional coordinates)

All files are ASCII with Unix (LF) line endings, also when run on Windows.
"""
import argparse
import os

import numpy as np
from ase import Atoms
from ase.io import read, write

# Valence electrons of a few DFT_DATA19 VPS files (used for the initial
# up/down occupation columns of OpenMX). Extend with --valence El=N.
OPENMX_VALENCE = {"H": 1.0, "C": 4.0, "N": 5.0, "O": 6.0, "Si": 4.0,
                  "S": 6.0}
# "Standard" PAO choices of the DFT_DATA19 table for a few elements;
# look up other elements in the OpenMX manual / DFT_DATA19 web page.
OPENMX_PAO = {"H": "H6.0-s2p1", "C": "C6.0-s2p2d1", "O": "O6.0-s2p2d1",
              "Si": "Si7.0-s2p2d1"}


def to_primitive(atoms, symprec):
    """Return the spglib standardized primitive cell as an ASE Atoms."""
    import spglib
    cell = (atoms.cell[:], atoms.get_scaled_positions(), atoms.numbers)
    lat, pos, num = spglib.standardize_cell(cell, to_primitive=True,
                                            no_idealize=False,
                                            symprec=symprec)
    return Atoms(numbers=num, scaled_positions=pos, cell=lat, pbc=True)


def lf_write(path, atoms, **kwargs):
    """ase.io.write with Unix (LF) line endings, also on Windows."""
    with open(path, "w", newline="\n") as f:
        write(f, atoms, **kwargs)


def clean(atoms):
    """Drop CIF bookkeeping arrays (occupancy, spacegroup_kinds, ...)."""
    return Atoms(atoms.get_chemical_symbols(), positions=atoms.positions,
                 cell=atoms.cell, pbc=atoms.pbc)


def write_cp2k(atoms, path):
    """CP2K: cell vectors in Angstrom and Cartesian coordinates."""
    with open(path, "w", newline="\n") as f:
        f.write("# CP2K subsystem generated by cif_to_all.py\n")
        f.write("# Use inside &FORCE_EVAL/&SUBSYS with:  @INCLUDE '%s'\n"
                % os.path.basename(path))
        f.write("# Remember to add one &KIND block (BASIS_SET, POTENTIAL) "
                "per element.\n")
        f.write("&CELL\n")
        for name, vec in zip("ABC", atoms.cell):
            f.write("  %s [angstrom] %14.8f %14.8f %14.8f\n" % (name, *vec))
        periodic = "XYZ" if all(atoms.pbc) else "NONE"
        f.write("  PERIODIC %s\n" % periodic)
        f.write("&END CELL\n")
        f.write("&COORD\n")
        for s, (x, y, z) in zip(atoms.get_chemical_symbols(),
                                atoms.positions):
            f.write("  %-3s %14.8f %14.8f %14.8f\n" % (s, x, y, z))
        f.write("&END COORD\n")


def write_openmx(atoms, path, valence):
    """OpenMX: Atoms.Number, unit vectors and fractional coordinates.

    Columns of Atoms.SpeciesAndCoordinates:
      index  species  x  y  z  n_up  n_down   (n_up + n_down = valence)
    The species names written here are element symbols; they must match
    the names defined in <Definition.of.Atomic.Species.
    """
    symbols = atoms.get_chemical_symbols()
    frac = atoms.get_scaled_positions(wrap=True)
    species = sorted(set(symbols), key=symbols.index)
    with open(path, "w", newline="\n") as f:
        f.write("# OpenMX structure blocks generated by cif_to_all.py\n")
        f.write("# Example species definition for DFT_DATA19 "
                "(check the recommended PAO/VPS table):\n")
        f.write("#   Species.Number  %d\n" % len(species))
        f.write("#   <Definition.of.Atomic.Species\n")
        for s in species:
            pao = OPENMX_PAO.get(s, "<PAO-of-%s>" % s)
            f.write("#    %-3s  %-14s %s_PBE19\n" % (s, pao, s))
        f.write("#   Definition.of.Atomic.Species>\n\n")
        f.write("Atoms.Number  %d\n" % len(atoms))
        f.write("Atoms.SpeciesAndCoordinates.Unit  FRAC\n")
        f.write("<Atoms.SpeciesAndCoordinates\n")
        for i, (s, (x, y, z)) in enumerate(zip(symbols, frac), start=1):
            nval = valence.get(s)
            if nval is None:
                up = dn = "  ?.?"   # forces the user to fill it in
            else:
                up = dn = "%5.2f" % (nval / 2.0)
            f.write("  %3d  %-3s %13.9f %13.9f %13.9f  %s  %s\n"
                    % (i, s, x, y, z, up, dn))
        f.write("Atoms.SpeciesAndCoordinates>\n")
        f.write("Atoms.UnitVectors.Unit  Ang\n")
        f.write("<Atoms.UnitVectors\n")
        for vec in atoms.cell:
            f.write("  %14.9f %14.9f %14.9f\n" % tuple(vec))
        f.write("Atoms.UnitVectors>\n")
    missing = sorted(set(symbols) - set(valence))
    if missing:
        print("  OpenMX: valence unknown for %s -> '?.?' written; "
              "use --valence El=N" % ", ".join(missing))


def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("infile")
    p.add_argument("-o", "--outdir", default="converted")
    p.add_argument("--primitive", action="store_true",
                   help="reduce to the spglib standardized primitive cell")
    p.add_argument("--symprec", type=float, default=1e-3)
    p.add_argument("--valence", nargs="*", default=[],
                   help="OpenMX valence electrons, e.g. Ti=12 O=6")
    args = p.parse_args()

    atoms = clean(read(args.infile))
    print("Read %s: %s, %d atoms, pbc=%s"
          % (args.infile, atoms.get_chemical_formula(), len(atoms),
             atoms.pbc.tolist()))
    if args.primitive:
        atoms = to_primitive(atoms, args.symprec)
        print("Primitive cell: %d atoms" % len(atoms))

    valence = dict(OPENMX_VALENCE)
    for item in args.valence:
        el, n = item.split("=")
        valence[el] = float(n)

    os.makedirs(args.outdir, exist_ok=True)
    name = os.path.splitext(os.path.basename(args.infile))[0]
    if args.primitive:
        name += "_prim"

    def out(fn):
        return os.path.join(args.outdir, fn)

    periodic = atoms.cell.rank == 3
    if periodic:
        lf_write(out("POSCAR"), atoms, format="vasp", direct=True)
    lf_write(out(name + ".xyz"), atoms, format="xyz")
    lf_write(out(name + ".extxyz"), atoms, format="extxyz")
    lf_write(out(name + ".gjf"), atoms, format="gaussian-in",
             method="PBEPBE", basis="6-31G(d)", charge=0, mult=1)
    with open(out(name + ".gjf"), "a", newline="\n") as f:
        f.write("\n")           # Gaussian needs a trailing blank line
    if periodic:
        write_cp2k(atoms, out(name + "_cp2k.inc"))
        write_openmx(atoms, out(name + "_openmx.dat"), valence)

    vol = atoms.get_volume() if periodic else float("nan")
    print("Cell lengths (A): %s" % np.round(atoms.cell.lengths(), 4).tolist())
    print("Volume (A^3)    : %.3f" % vol)
    print("Written to %s/: %s" % (args.outdir,
                                  " ".join(sorted(os.listdir(args.outdir)))))


if __name__ == "__main__":
    main()