"""
libcode_gulp.py — Interface between the evolutionary algorithm and the GULP
force-field calculator.

Workflow:
  1. `make_ainput_gulp`     — write one .gin file per structure.
  2. `run_gulp_in_folder`   — run GULP sequentially on every .gin in a folder
                               using subprocess (no intermediate .sh scripts).
  3. `calculator_all_gulp`  — orchestrate parallel GULP runs across folders.
  4. `get_geometry_gulp`    — parse a .got output file into an ASE Atoms object.
  5. `structure_recolector` — gather results from all folders after runs finish.

Design note — why folders are still needed:
    GULP writes temporary files (restart files, scratch data) whose names are
    fixed by GULP itself and cannot be parameterised.  Running multiple GULP
    processes in the same directory would cause those files to collide and
    corrupt each other.  Each parallel worker therefore operates in its own
    subdirectory.  Python's subprocess 'cwd' parameter replaces the 'cd'
    command that the old bash scripts used.
"""
import os
import sys
import time
import subprocess
import numpy as np
from ase import Atoms
from multiprocessing import Process
from aegon.libutils import sort_by_energy, prepare_folders, split_poscarlist

#------------------------------------------------------------------------------------------
def make_ainput_gulp(headgulp, singleposcar, folder):
    """Write one GULP input (.gin) file for a single structure.

    Placeholders in headgulp:
        'LATTICEVECTORS' — replaced by the 3×3 cell matrix (Angstroms).
        'COORDINATES'    — replaced by fractional coordinates (symbol x y z).

    in:
        headgulp     (list of str): Template lines from the ---GULP--- block.
        singleposcar (Atoms)      : Structure to write; must have info['i'] (label).
        folder       (str)        : Destination directory (trailing '/' included).
    out:
        None  — writes <folder><label>.gin; skips if .got or .gin already exists.
    """
    iposname = singleposcar.info['i']
    nameinp  = iposname + '.gin'
    nameout  = iposname + '.got'
    if os.path.isfile(folder + nameout):
        print("%s ... skipped (output %s already exists)" % (nameinp, folder + nameout))
    elif os.path.isfile(folder + nameinp):
        print("%s ... skipped (input already exists)" % nameinp)
    else:
        with open(folder + nameinp, 'w') as fh:
            for iline in headgulp:
                if iline == 'LATTICEVECTORS':
                    matrix = singleposcar.cell
                    print("%12.9f %16.9f %16.9f" % (matrix[0,0], matrix[0,1], matrix[0,2]), file=fh)
                    print("%12.9f %16.9f %16.9f" % (matrix[1,0], matrix[1,1], matrix[1,2]), file=fh)
                    print("%12.9f %16.9f %16.9f" % (matrix[2,0], matrix[2,1], matrix[2,2]), file=fh)
                elif iline == 'COORDINATES':
                    symbols            = singleposcar.get_chemical_symbols()
                    direct_coordinates = singleposcar.get_scaled_positions()
                    for symbol, coord in zip(symbols, direct_coordinates):
                        print("%-2s %16.9f %16.9f %16.9f" % (symbol, coord[0], coord[1], coord[2]), file=fh)
                else:
                    print(iline, file=fh)

#------------------------------------------------------------------------------------------
def run_gulp_in_folder(gulp_path, folder, poscarlist):
    """Run GULP sequentially on every structure in a folder using subprocess.

    Replaces the old approach of writing a bash launcher script (.sh) and
    executing it with os.system.  Benefits:
      - No intermediate .sh file is written, read, or left behind on failure.
      - subprocess handles stdin/stdout redirection natively (no shell required).
      - 'cwd=folder' keeps GULP's temporary files inside the designated folder,
        preventing collisions between parallel workers.
      - The GULP exit code is available for future error-handling extensions.

    Structures whose .got file already exists are skipped to support restarts.

    in:
        gulp_path  (str)        : Absolute path to the GULP executable.
        folder     (str)        : Working directory for this batch of structures.
        poscarlist (list Atoms) : Structures to optimize; each has info['i'] (label).
    out:
        None  — results are written to <folder>/<label>.got files.
    """
    for iposcar in poscarlist:
        name     = iposcar.info['i']
        gin_path = os.path.join(folder, name + '.gin')
        got_path = os.path.join(folder, name + '.got')
        if os.path.isfile(got_path):
            continue  # Already computed; supports restart without recomputing.
        try:
            with open(gin_path, 'r') as fin, open(got_path, 'w') as fout:
                subprocess.run(
                    [gulp_path],
                    stdin=fin,
                    stdout=fout,
                    stderr=subprocess.DEVNULL,  # GULP's stderr is not informative.
                    cwd=folder,                 # Run inside the folder so GULP's
                )                               # temp files do not collide across
        except Exception as e:                  # parallel workers.
            # Log the error but continue so the remaining structures are attempted.
            print("GULP subprocess error for %s: %s" % (name, e))

#------------------------------------------------------------------------------------------
def get_termination_gulp(pathfilename):
    """Check whether a GULP output file reports successful optimization.

    in:
        pathfilename (str): Full path to the .got file.
    out:
        1     — '**** Optimisation achieved ****' found (converged).
        0     — file exists but optimization did not converge.
        False — file does not exist.
    """
    if not os.path.isfile(pathfilename):
        return False
    normal = 0  # Default: file exists but convergence not achieved.
    with open(pathfilename, 'r') as f:
        for line in f:
            if '**** Optimisation achieved ****' in line:
                normal = 1
                break
    return normal

#------------------------------------------------------------------------------------------
def get_energy_gulp(pathfilename):
    """Read the final energy (or enthalpy for constant-pressure runs) from a .got file.

    'Final enthalpy' takes precedence over 'Final energy' when both appear;
    this matches GULP output conventions for constant-pressure optimizations.

    in:
        pathfilename (str): Full path to the GULP output file.
    out:
        eneineV (float): Final energy / enthalpy in eV.
    raises:
        ValueError if no energy line is found.
    """
    eneineV = None
    with open(pathfilename, 'r') as f:
        for line in f:
            if 'Final energy =' in line:
                eneineV = float(line.split()[3])
            if 'Final enthalpy =' in line:
                # Enthalpy overrides energy for constant-pressure runs.
                eneineV = float(line.split()[3])
    if eneineV is None:
        raise ValueError("No 'Final energy' or 'Final enthalpy' found in '%s'" % pathfilename)
    return eneineV

#------------------------------------------------------------------------------------------
def get_geometry_gulp(pathfilename):
    """Parse the final geometry from a GULP output file and return an ASE Atoms object.

    Reads the file once into memory, then makes two in-memory passes:
      - Pass 1: locate the last Cartesian lattice vector block.
      - Pass 2: locate the last fractional coordinate block.
    Keeping only the last occurrence of each block ensures we capture the
    final (converged) geometry rather than any intermediate snapshot.

    in:
        pathfilename (str): Full path to the .got file.
    out:
        Atoms with info['i'] (label), info['e'] (eV), info['c'] (convergence flag).
        False if the file is missing, did not converge, or geometry parsing failed.
    """
    nt = get_termination_gulp(pathfilename)
    if nt is False:
        return False  # Output file does not exist.
    if nt == 0:
        return False  # Did not converge; partial geometry is unreliable.

    namein = os.path.basename(pathfilename).split('.')[0]

    # Single disk read; all subsequent operations work on the in-memory list.
    with open(pathfilename, 'r') as f:
        all_lines = f.readlines()

    ans1 = 'Final Cartesian lattice vectors (Angstroms) :'
    ans2 = 'Cartesian lattice vectors (Angstroms) :'

    # --- Pass 1: find the LAST lattice vector block ---
    matrix = None
    for i, line in enumerate(all_lines):
        if (ans1 in line) or (ans2 in line):
            try:
                # File layout after header: 1 blank line, then 3 vector lines.
                a1 = list(map(float, all_lines[i + 2].split()))
                a2 = list(map(float, all_lines[i + 3].split()))
                a3 = list(map(float, all_lines[i + 4].split()))
                matrix = np.array([a1, a2, a3])
            except (ValueError, IndexError):
                pass  # Malformed block; keep searching for a later occurrence.

    if matrix is None:
        return False

    # --- Pass 2: find the LAST fractional coordinate block ---
    symbols            = []
    direct_coordinates = []
    for i, line in enumerate(all_lines):
        if 'Final fractional coordinates of atoms :' in line:
            syms_tmp, coords_tmp = [], []
            # Skip: blank + separator + 2 column-header lines + separator = 5 lines,
            # then the first data line (7 tokens: No. symbol c/s x y z charge).
            j = i + 6
            while j < len(all_lines):
                ls = all_lines[j].split()
                if len(ls) == 7:
                    syms_tmp.append(str(ls[1]))
                    coords_tmp.append([float(ls[3]), float(ls[4]), float(ls[5])])
                    j += 1
                else:
                    break  # Trailing separator line.
            # Overwrite with this (later) occurrence to keep the final geometry.
            symbols, direct_coordinates = syms_tmp, coords_tmp

    if not symbols:
        return False

    poscarx = Atoms(symbols=symbols, scaled_positions=direct_coordinates,
                    cell=matrix, pbc=True)
    poscarx.info['c'] = nt
    poscarx.info['i'] = namein
    poscarx.info['e'] = get_energy_gulp(pathfilename)
    return poscarx

#------------------------------------------------------------------------------------------
def get_all_xt_geometry_gulp(poscarlist, path):
    """Collect successfully parsed geometries from a list of GULP output files.

    in:
        poscarlist (list Atoms): Input structures; each must have info['i'] (label).
        path       (str)       : Directory containing .got files (trailing '/' included).
    out:
        List of parsed Atoms objects, or False if none were parsed successfully.
    """
    poscarout = []
    for ipos in poscarlist:
        filename = ipos.info['i'] + '.got'
        pos01 = get_geometry_gulp(path + filename)
        if pos01 is not False:
            poscarout.append(pos01)
    return poscarout if poscarout else False

#------------------------------------------------------------------------------------------
def structure_recolector(folders, poscars_split):
    """Gather optimized structures from all parallel GULP folders and print statistics.

    Classifies each calculation as:
        N.T. — Normal Termination (converged).
        A.T. — Abnormal Termination (ran but did not converge).
        N.E. — No output file (GULP did not start or crashed before writing output).

    in:
        folders       (list str)      : Folder names used for the parallel runs.
        poscars_split (list of lists) : Structures assigned to each folder.
    out:
        moleculeout (list Atoms): All converged structures.
    raises:
        SystemExit(1) if no structures converged at all.
    """
    list_ne, list_at, list_nt = [], [], []
    ntot       = 0
    moleculeout = []
    for ifolder, iposcars in zip(folders, poscars_split):
        ntot += len(iposcars)
        molx = get_all_xt_geometry_gulp(iposcars, ifolder + '/')
        if molx is not False:
            moleculeout = moleculeout + molx
        for iposcar in iposcars:
            termination = get_termination_gulp(ifolder + '/' + iposcar.info['i'] + '.got')
            if termination is False:
                list_ne.append(iposcar.info['i'])
            elif termination == 0:
                list_at.append(iposcar.info['i'])
            elif termination == 1:
                list_nt.append(iposcar.info['i'])
    nnt, nat = len(list_nt), len(list_at)
    pnt = int(float(nnt) * 100.0 / float(ntot))
    pat = int(float(nat) * 100.0 / float(ntot))
    print("Calculations with N.T. = %d (%d percent)" % (nnt, pnt))
    print("Calculations with A.T. = %d (%d percent)" % (nat, pat))
    if not moleculeout:
        print("ZERO calculations finished satisfactorily — aborting.")
        sys.exit(1)
    return moleculeout

#------------------------------------------------------------------------------------------
def calculator_all_gulp(poscar_list_in, block_gulp, gulp_path, nproc, base_name):
    """Run GULP geometry optimizations in parallel using Python multiprocessing.

    Distributes structures across up to `nproc` folders, writes .gin input files,
    launches one subprocess-based worker per folder in parallel, then collects
    results.  No intermediate bash (.sh) scripts are written.

    in:
        poscar_list_in (list Atoms): Structures to optimize.
        block_gulp     (list str)  : GULP template lines (from ---GULP--- block).
        gulp_path      (str)       : Absolute path to the GULP executable.
        nproc          (int)       : Maximum number of parallel worker processes.
        base_name      (str)       : Prefix for folder names (e.g., 'stage').
    out:
        poscar_list_out (list Atoms): Converged, optimized structures.
    """
    start     = time.time()
    # Avoid creating empty folders when there are fewer structures than workers.
    n_workers    = min(nproc, len(poscar_list_in))
    folder_names  = prepare_folders(poscar_list_in, n_workers, base_name)
    split_poscars = split_poscarlist(poscar_list_in, n_workers)

    # Write all .gin input files before starting any worker.
    for poscars, ifolder in zip(split_poscars, folder_names):
        for iposcar in poscars:
            make_ainput_gulp(block_gulp, iposcar, ifolder + '/')

    # Launch one worker process per folder.  Each worker runs GULP sequentially
    # on its own structures via subprocess; parallelism is at the folder level.
    procs = []
    print("Local submit: %sproc01 to %sproc%s" % (base_name, base_name, str(n_workers).zfill(2)), end=" ")
    for poscars, ifolder in zip(split_poscars, folder_names):
        proc = Process(target=run_gulp_in_folder, args=(gulp_path, ifolder, poscars))
        procs.append(proc)
        proc.start()
    for proc in procs:
        proc.join()

    end = time.time()
    print('TOTAL TIME: %5.2f s' % (end - start))
    poscar_list_out = structure_recolector(folder_names, split_poscars)
    return poscar_list_out

#------------------------------------------------------------------------------------------
def display_info(moleculein, stage_string, dicc_term):
    """Print a ranked energy summary table for a list of structures.

    in:
        moleculein   (list Atoms): Structures; each must have info['e'] and info['i'].
        stage_string (str)       : Label shown in the summary header.
        dicc_term    (dict)      : Maps structure label → termination code (0 or 1).
    out:
        None
    """
    print("-------------------------------------------------------------------")
    print("------------------------- SUMMARY %s -------------------------" % stage_string)
    print("Number File--------Name   Energy (ev)   Delta----E T")
    molzz = sort_by_energy(moleculein, 1)
    emin  = molzz[0].info['e']
    for ii, imol in enumerate(molzz):
        ei     = imol.info['e']
        id_    = imol.info['i']
        nt     = dicc_term[id_]
        deltae = ei - emin
        kk     = str(ii + 1).zfill(6)
        print("%s %s %13.8f %12.8f %d" % (kk, id_, ei, deltae, nt))
#------------------------------------------------------------------------------------------
