import time
import warnings

from ase.units import Ang, fs
from ase.utils import reader, writer


v_unit = Ang / (1000.0 * fs)


@reader
def read_aims(fd, apply_constraints=True):
    """Import FHI-aims geometry type files.

    Reads unitcell, atom positions and constraints from
    a geometry.in file.

    If geometric constraint (symmetry parameters) are in the file
    include that information in atoms.info["symmetry_block"]
    """

    lines = fd.readlines()
    return parse_geometry_lines(lines, apply_constraints=apply_constraints)


def parse_geometry_lines(lines, apply_constraints=True):

    from ase import Atoms
    from ase.constraints import (
        FixAtoms,
        FixCartesian,
        FixScaledParametricRelations,
        FixCartesianParametricRelations,
    )
    import numpy as np

    atoms = Atoms()

    positions = []
    cell = []
    symbols = []
    velocities = []
    magmoms = []
    symmetry_block = []
    charges = []
    fix = []
    fix_cart = []
    xyz = np.array([0, 0, 0])
    i = -1
    n_periodic = -1
    periodic = np.array([False, False, False])
    cart_positions, scaled_positions = False, False
    for line in lines:
        inp = line.split()
        if inp == []:
            continue
        if inp[0] in ["atom", "atom_frac"]:

            if inp[0] == "atom":
                cart_positions = True
            else:
                scaled_positions = True

            if xyz.all():
                fix.append(i)
            elif xyz.any():
                fix_cart.append(FixCartesian(i, xyz))
            floatvect = float(inp[1]), float(inp[2]), float(inp[3])
            positions.append(floatvect)
            symbols.append(inp[4])
            magmoms.append(0.0)
            charges.append(0.0)
            xyz = np.array([0, 0, 0])
            i += 1

        elif inp[0] == "lattice_vector":
            floatvect = float(inp[1]), float(inp[2]), float(inp[3])
            cell.append(floatvect)
            n_periodic = n_periodic + 1
            periodic[n_periodic] = True

        elif inp[0] == "initial_moment":
            magmoms[-1] = float(inp[1])

        elif inp[0] == "initial_charge":
            charges[-1] = float(inp[1])

        elif inp[0] == "constrain_relaxation":
            if inp[1] == ".true.":
                fix.append(i)
            elif inp[1] == "x":
                xyz[0] = 1
            elif inp[1] == "y":
                xyz[1] = 1
            elif inp[1] == "z":
                xyz[2] = 1

        elif inp[0] == "velocity":
            floatvect = [v_unit * float(l) for l in inp[1:4]]
            velocities.append(floatvect)

        elif inp[0] in [
            "symmetry_n_params",
            "symmetry_params",
            "symmetry_lv",
            "symmetry_frac",
        ]:
            symmetry_block.append(" ".join(inp))

    if xyz.all():
        fix.append(i)
    elif xyz.any():
        fix_cart.append(FixCartesian(i, xyz))

    if cart_positions and scaled_positions:
        raise Exception(
            "Can't specify atom positions with mixture of "
            "Cartesian and fractional coordinates"
        )
    elif scaled_positions and periodic.any():
        atoms = Atoms(
            symbols, scaled_positions=positions, cell=cell, pbc=periodic
        )
    else:
        atoms = Atoms(symbols, positions)

    if len(velocities) > 0:
        if len(velocities) != len(positions):
            raise Exception(
                "Number of positions and velocities have to coincide."
            )
        atoms.set_velocities(velocities)

    fix_params = []

    if len(symmetry_block) > 5:
        params = symmetry_block[1].split()[1:]

        lattice_expressions = []
        lattice_params = []

        atomic_expressions = []
        atomic_params = []

        n_lat_param = int(symmetry_block[0].split(" ")[2])

        lattice_params = params[:n_lat_param]
        atomic_params = params[n_lat_param:]

        for ll, line in enumerate(symmetry_block[2:]):
            expression = " ".join(line.split(" ")[1:])
            if ll < 3:
                lattice_expressions += expression.split(",")
            else:
                atomic_expressions += expression.split(",")

        fix_params.append(
            FixCartesianParametricRelations.from_expressions(
                list(range(3)),
                lattice_params,
                lattice_expressions,
                use_cell=True,
            )
        )

        fix_params.append(
            FixScaledParametricRelations.from_expressions(
                list(range(len(atoms))), atomic_params, atomic_expressions
            )
        )

    if any(magmoms):
        atoms.set_initial_magnetic_moments(magmoms)
    if any(charges):
        atoms.set_initial_charges(charges)

    if periodic.any():
        atoms.set_cell(cell)
        atoms.set_pbc(periodic)
    if len(fix):
        atoms.set_constraint([FixAtoms(indices=fix)] + fix_cart + fix_params)
    else:
        atoms.set_constraint(fix_cart + fix_params)

    if fix_params and apply_constraints:
        atoms.set_positions(atoms.get_positions())
    return atoms


@writer
def write_aims(
    fd,
    atoms,
    scaled=False,
    geo_constrain=False,
    velocities=False,
    ghosts=None,
    info_str=None,
    wrap=False,
):
    """Method to write FHI-aims geometry files.

    Writes the atoms positions and constraints (only FixAtoms is
    supported at the moment).

    Args:
        fd: file object
            File to output structure to
        atoms: ase.atoms.Atoms
            structure to output to the file
        scaled: bool
            If True use fractional coordinates instead of Cartesian coordinates
        symmetry_block: list of str
            List of geometric constraints as defined in:
            https://arxiv.org/abs/1908.01610
        velocities: bool
            If True add the atomic velocity vectors to the file
        ghosts: list of Atoms
            A list of ghost atoms for the system
        info_str: str
            A string to be added to the header of the file
        wrap: bool
            Wrap atom positions to cell before writing
    """

    from ase.constraints import FixAtoms, FixCartesian

    import numpy as np

    if geo_constrain:
        if not scaled:
            warnings.warn(
                "Setting scaled to True because a symmetry_block is detected."
            )
            scaled = True

    fd.write("#=======================================================\n")
    if hasattr(fd, 'name'):
        fd.write("# FHI-aims file: " + fd.name + "\n")
    fd.write("# Created using the Atomic Simulation Environment (ASE)\n")
    fd.write("# " + time.asctime() + "\n")

    # If writing additional information is requested via info_str:
    if info_str is not None:
        fd.write("\n# Additional information:\n")
        if isinstance(info_str, list):
            fd.write("\n".join(["#  {}".format(s) for s in info_str]))
        else:
            fd.write("# {}".format(info_str))
        fd.write("\n")

    fd.write("#=======================================================\n")

    i = 0
    if atoms.get_pbc().any():
        for n, vector in enumerate(atoms.get_cell()):
            fd.write("lattice_vector ")
            for i in range(3):
                fd.write("%16.16f " % vector[i])
            fd.write("\n")
    fix_cart = np.zeros([len(atoms), 3])

    # else aims crashes anyways
    # better be more explicit
    # write_magmoms = np.any([a.magmom for a in atoms])

    if atoms.constraints:
        for constr in atoms.constraints:
            if isinstance(constr, FixAtoms):
                fix_cart[constr.index] = [1, 1, 1]
            elif isinstance(constr, FixCartesian):
                fix_cart[constr.a] = -constr.mask + 1

    if ghosts is None:
        ghosts = np.zeros(len(atoms))
    else:
        assert len(ghosts) == len(atoms)

    if geo_constrain:
        wrap = False
    scaled_positions = atoms.get_scaled_positions(wrap=wrap)

    for i, atom in enumerate(atoms):
        if ghosts[i] == 1:
            atomstring = "empty "
        elif scaled:
            atomstring = "atom_frac "
        else:
            atomstring = "atom "
        fd.write(atomstring)
        if scaled:
            for pos in scaled_positions[i]:
                fd.write("%16.16f " % pos)
        else:
            for pos in atom.position:
                fd.write("%16.16f " % pos)
        fd.write(atom.symbol)
        fd.write("\n")
        # (1) all coords are constrained:
        if fix_cart[i].all():
            fd.write("    constrain_relaxation .true.\n")
        # (2) some coords are constrained:
        elif fix_cart[i].any():
            xyz = fix_cart[i]
            for n in range(3):
                if xyz[n]:
                    fd.write("    constrain_relaxation %s\n" % "xyz"[n])
        if atom.charge:
            fd.write("    initial_charge %16.6f\n" % atom.charge)
        if atom.magmom:
            fd.write("    initial_moment %16.6f\n" % atom.magmom)

        # Write velocities if this is wanted
        if velocities and atoms.get_velocities() is not None:
            fd.write(
                "    velocity {:.16f} {:.16f} {:.16f}\n".format(
                    *atoms.get_velocities()[i] / v_unit
                )
            )

    if geo_constrain:
        for line in get_sym_block(atoms):
            fd.write(line)


def get_sym_block(atoms):
    """Get symmetry block for Parametric constraints in atoms.constraints"""
    import numpy as np

    from ase.constraints import (
        FixScaledParametricRelations,
        FixCartesianParametricRelations,
    )

    # Initialize param/expressions lists
    atomic_sym_params = []
    lv_sym_params = []
    atomic_param_constr = np.zeros((len(atoms),), dtype="<U100")
    lv_param_constr = np.zeros((3,), dtype="<U100")

    # Populate param/expressions list
    for constr in atoms.constraints:
        if isinstance(constr, FixScaledParametricRelations):
            atomic_sym_params += constr.params

            if np.any(atomic_param_constr[constr.indices] != ""):
                warnings.warn(
                    "multiple parametric constraints defined for the same "
                    "atom, using the last one defined"
                )

            atomic_param_constr[constr.indices] = [
                ", ".join(expression) for expression in constr.expressions
            ]
        elif isinstance(constr, FixCartesianParametricRelations):
            lv_sym_params += constr.params

            if np.any(lv_param_constr[constr.indices] != ""):
                warnings.warn(
                    "multiple parametric constraints defined for the same "
                    "lattice vector, using the last one defined"
                )

            lv_param_constr[constr.indices] = [
                ", ".join(expression) for expression in constr.expressions
            ]

    if np.all(atomic_param_constr == "") and np.all(lv_param_constr == ""):
        return []

    # Check Constraint Parameters
    if len(atomic_sym_params) != len(np.unique(atomic_sym_params)):
        warnings.warn(
            "Some parameters were used across constraints, they will be "
            "combined in the aims calculations"
        )
        atomic_sym_params = np.unique(atomic_sym_params)

    if len(lv_sym_params) != len(np.unique(lv_sym_params)):
        warnings.warn(
            "Some parameters were used across constraints, they will be "
            "combined in the aims calculations"
        )
        lv_sym_params = np.unique(lv_sym_params)

    if np.any(atomic_param_constr == ""):
        raise IOError(
            "FHI-aims input files require all atoms have defined parametric "
            "constraints"
        )

    cell_inds = np.where(lv_param_constr == "")[0]
    for ind in cell_inds:
        lv_param_constr[ind] = "{:.16f}, {:.16f}, {:.16f}".format(
            *atoms.cell[ind]
        )

    n_atomic_params = len(atomic_sym_params)
    n_lv_params = len(lv_sym_params)
    n_total_params = n_atomic_params + n_lv_params

    sym_block = []
    if n_total_params > 0:
        sym_block.append("#" + "="*55 + "\n")
        sym_block.append("# Parametric constraints\n")
        sym_block.append("#" + "="*55 + "\n")
        sym_block.append(
            "symmetry_n_params {:d} {:d} {:d}\n".format(
                n_total_params, n_lv_params, n_atomic_params
            )
        )
        sym_block.append(
            "symmetry_params %s\n"
            % " ".join(lv_sym_params + atomic_sym_params)
        )

        for constr in lv_param_constr:
            sym_block.append("symmetry_lv {:s}\n".format(constr))

        for constr in atomic_param_constr:
            sym_block.append("symmetry_frac {:s}\n".format(constr))
    return sym_block


def _parse_atoms(fd, n_atoms, molecular_dynamics=False):
    """parse structure information from aims output to Atoms object"""
    from ase import Atoms, Atom

    next(fd)
    atoms = Atoms()
    for i in range(n_atoms):
        inp = next(fd).split()
        if "lattice_vector" in inp[0]:
            cell = []
            for i in range(3):
                cell += [[float(inp[1]), float(inp[2]), float(inp[3])]]
                inp = next(fd).split()
            atoms.set_cell(cell)
            inp = next(fd).split()
        atoms.append(Atom(inp[4], (inp[1], inp[2], inp[3])))
        if molecular_dynamics:
            inp = next(fd).split()

    return atoms


@reader
def read_aims_output(fd, index=-1):
    """Import FHI-aims output files with all data available, i.e.
    relaxations, MD information, force information etc etc etc."""
    from ase import Atoms, Atom
    from ase.calculators.singlepoint import SinglePointCalculator
    from ase.constraints import FixAtoms, FixCartesian

    molecular_dynamics = False
    cell = []
    images = []
    fix = []
    fix_cart = []
    f = None
    pbc = False
    found_aims_calculator = False
    stress = None
    for line in fd:
        # if "List of parameters used to initialize the calculator:" in line:
        #     next(fd)
        #     calc = read_aims_calculator(fd)
        #     calc.out = filename
        #     found_aims_calculator = True
        if "| Number of atoms                   :" in line:
            inp = line.split()
            n_atoms = int(inp[5])
        if "| Unit cell:" in line:
            if not pbc:
                pbc = True
                for i in range(3):
                    inp = next(fd).split()
                    cell.append([inp[1], inp[2], inp[3]])
        if "Found relaxation constraint for atom" in line:
            xyz = [0, 0, 0]
            ind = int(line.split()[5][:-1]) - 1
            if "All coordinates fixed" in line:
                if ind not in fix:
                    fix.append(ind)
            if "coordinate fixed" in line:
                coord = line.split()[6]
                if coord == "x":
                    xyz[0] = 1
                elif coord == "y":
                    xyz[1] = 1
                elif coord == "z":
                    xyz[2] = 1
                keep = True
                for n, c in enumerate(fix_cart):
                    if ind == c.a:
                        keep = False
                if keep:
                    fix_cart.append(FixCartesian(ind, xyz))
                else:
                    fix_cart[n].mask[xyz.index(1)] = 0

        if "Atomic structure:" in line and not molecular_dynamics:
            next(fd)
            atoms = Atoms()
            for _ in range(n_atoms):
                inp = next(fd).split()
                atoms.append(Atom(inp[3], (inp[4], inp[5], inp[6])))

        if "Complete information for previous time-step:" in line:
            molecular_dynamics = True

        if "Updated atomic structure:" in line and not molecular_dynamics:
            atoms = _parse_atoms(fd, n_atoms=n_atoms)
        elif "Atomic structure (and velocities)" in line:
            next(fd)
            atoms = Atoms()
            velocities = []
            for i in range(n_atoms):
                inp = next(fd).split()
                atoms.append(Atom(inp[4], (inp[1], inp[2], inp[3])))
                inp = next(fd).split()
                floatvect = [v_unit * float(l) for l in inp[1:4]]
                velocities.append(floatvect)
            atoms.set_velocities(velocities)
            if len(fix):
                atoms.set_constraint([FixAtoms(indices=fix)] + fix_cart)
            else:
                atoms.set_constraint(fix_cart)
            images.append(atoms)

        # if we enter here, the SocketIO/PIMD Wrapper was used
        elif (
            "Atomic structure that "
            "was used in the preceding time step of the wrapper"
        ) in line:
            # parse atoms and add calculator information, i.e., the energies
            # and forces that were already collected
            atoms = _parse_atoms(fd, n_atoms=n_atoms)
            results = images[-1].calc.results
            atoms.calc = SinglePointCalculator(atoms, **results)

            # replace last image with updated atoms
            images[-1] = atoms

            # make sure `atoms` does not point to `images[-1` later on
            atoms = atoms.copy()

        # FlK: add analytical stress and replace stress=None
        if "Analytical stress tensor - Symmetrized" in line:
            # scroll to significant lines
            for _ in range(4):
                next(fd)
            stress = []
            for _ in range(3):
                inp = next(fd)
                stress.append([float(i) for i in inp.split()[2:5]])

        if "Total atomic forces" in line:
            f = []
            for i in range(n_atoms):
                inp = next(fd).split()
                # FlK: use inp[-3:] instead of inp[1:4] to make sure this works
                # when atom number is not preceded by a space.
                f.append([float(i) for i in inp[-3:]])
            if not found_aims_calculator:
                e = images[-1].get_potential_energy()
                # FlK: Add the stress if it has been computed
                if stress is None:
                    calc = SinglePointCalculator(atoms, energy=e, forces=f)
                else:
                    calc = SinglePointCalculator(
                        atoms, energy=e, forces=f, stress=stress
                    )
                images[-1].calc = calc
            e = None
            f = None

        if "Total energy corrected" in line:
            e = float(line.split()[5])
            if pbc:
                atoms.set_cell(cell)
                atoms.pbc = True
            if not found_aims_calculator:
                atoms.calc = SinglePointCalculator(atoms, energy=e)
            if not molecular_dynamics:
                if len(fix):
                    atoms.set_constraint([FixAtoms(indices=fix)] + fix_cart)
                else:
                    atoms.set_constraint(fix_cart)
                images.append(atoms)
            e = None
            # if found_aims_calculator:
            # calc.set_results(images[-1])
            # images[-1].calc = calc

        # FlK: add stress per atom
        if "Per atom stress (eV) used for heat flux calculation" in line:
            # scroll to boundary
            next(l for l in fd if "-------------" in l)

            stresses = []
            for l in [next(fd) for _ in range(n_atoms)]:
                # Read stresses
                xx, yy, zz, xy, xz, yz = [float(d) for d in l.split()[2:8]]
                stresses.append([[xx, xy, xz], [xy, yy, yz], [xz, yz, zz]])

            if not found_aims_calculator:
                e = images[-1].get_potential_energy()
                f = images[-1].get_forces()
                stress = images[-1].get_stress(voigt=False)

                calc = SinglePointCalculator(
                    atoms, energy=e, forces=f, stress=stress, stresses=stresses
                )
                images[-1].calc = calc

    fd.close()
    if molecular_dynamics:
        images = images[1:]

    # return requested images, code borrowed from ase/io/trajectory.py
    if isinstance(index, int):
        return images[index]
    else:
        step = index.step or 1
        if step > 0:
            start = index.start or 0
            if start < 0:
                start += len(images)
            stop = index.stop or len(images)
            if stop < 0:
                stop += len(images)
        else:
            if index.start is None:
                start = len(images) - 1
            else:
                start = index.start
                if start < 0:
                    start += len(images)
            if index.stop is None:
                stop = -1
            else:
                stop = index.stop
                if stop < 0:
                    stop += len(images)
        return [images[i] for i in range(start, stop, step)]