"""
Collection of process functions for AiiDA, used for structure transformation
"""

import re
from typing import Any, List, Tuple

import numpy as np
from aiida import orm
from aiida.engine import calcfunction
from aiida.orm import StructureData
from aiida.tools.data.array.kpoints import get_kpoints_path
from aiida.tools.data.structure import spglib_tuple_to_structure, structure_to_spglib_tuple
from ase import Atoms
from ase.build import niggli_reduce as niggli_reduce_
from ase.build import sort
from ase.build.supercells import make_supercell as ase_supercell
from ase.neb import NEB
from spglib import niggli_reduce as niggli_reduce_spg
from spglib import refine_cell


@calcfunction
def magnetic_structure_decorate(structure: orm.StructureData, magmom: orm.List) -> dict[str, Any]:
    """
    Create Quantum Espresso style decorated structure with
    given magnetic moments.
    """

    magmom = magmom.get_list()
    assert len(magmom) == len(structure.sites), (
        f'Mismatch between the magmom ({len(magmom)}) and the nubmer of sites ({len(structure.sites)}).'
    )
    old_species = [structure.get_kind(site.kind_name).symbol for site in structure.sites]
    new_species, magmom_mapping = create_additional_species(old_species, magmom)
    new_structure = StructureData()
    new_structure.set_cell(structure.cell)
    new_structure.set_pbc(structure.pbc)
    for site, name in zip(structure.sites, new_species):
        this_symbol = structure.get_kind(site.kind_name).symbol
        new_structure.append_atom(position=site.position, symbols=this_symbol, name=name)

    # Keep the label
    new_structure.label = structure.label
    return {'structure': new_structure, 'mapping': orm.Dict(dict=magmom_mapping)}


@calcfunction
def magnetic_structure_dedecorate(structure: orm.StructureData, mapping: orm.Dict) -> dict[str, Any]:
    """
    Remove decorations of a structure with multiple names for the same specie
    given that the decoration was previously created to give different species
    name for different initialisation of magnetic moments.
    """

    mapping = mapping.get_dict()
    # Get a list of decroated names
    old_species = [structure.get_kind(site.kind_name).name for site in structure.sites]
    new_species, magmom = convert_to_plain_list(old_species, mapping)

    new_structure = StructureData()
    new_structure.set_cell(structure.cell)
    new_structure.set_pbc(structure.pbc)

    for site, name in zip(structure.sites, new_species):
        this_symbol = structure.get_kind(site.kind_name).symbol
        new_structure.append_atom(position=site.position, symbols=this_symbol, name=name)
    new_structure.label = structure.label
    return {'structure': new_structure, 'magmom': orm.List(list=magmom)}


@calcfunction
def rattle(structure: orm.StructureData, amp: orm.Float) -> orm.StructureData:
    """
    Rattle the structure by a certain amplitude
    """
    native_keys = ['cell', 'pbc1', 'pbc2', 'pbc3', 'kinds', 'sites', 'mp_id']
    # Keep the foreign keys as it is
    foreign_attrs = {key: value for key, value in structure.attributes.items() if key not in native_keys}
    atoms = structure.get_ase()
    atoms.rattle(amp.value)
    # Clean any tags etc
    atoms.set_tags(None)
    atoms.set_masses(None)
    # Convert it back
    out = StructureData(ase=atoms)
    out.base.attributes.set_many(foreign_attrs)
    out.label = structure.label + ' RATTLED'
    return out


@calcfunction
def get_primitive(structure: orm.StructureData) -> orm.StructureData:
    """Create primitive structure use pymatgen interface"""
    pstruct = structure.get_pymatgen()
    ps = pstruct.get_primitive_structure()
    out = StructureData(pymatgen=ps)
    out.label = structure.label + ' PRIMITIVE'
    return out


@calcfunction
def get_standard_primitive(structure: orm.StructureData, **kwargs: Any) -> orm.StructureData:
    """Create the standard primitive structure via seekpath"""

    parameters = kwargs.get('parameters', {'symprec': 1e-5})
    if isinstance(parameters, orm.Dict):
        parameters = parameters.get_dict()

    out = get_kpoints_path(structure, **parameters)['primitive_structure']
    out.label = structure.label + ' PRIMITIVE'
    return out


@calcfunction
def spglib_refine_cell(structure: orm.StructureData, symprec: Any) -> orm.StructureData:
    """Create the standard primitive structure via seekpath"""

    structure_tuple, kind_info, kinds = structure_to_spglib_tuple(structure)

    lattice, positions, types = refine_cell(structure_tuple, symprec.value)

    refined = spglib_tuple_to_structure((lattice, positions, types), kind_info, kinds)

    return refined


@calcfunction
def get_standard_conventional(structure: orm.StructureData) -> orm.StructureData:
    """Create the standard primitive structure via seekpath"""
    out = get_kpoints_path(structure)['conv_structure']
    out.label = structure.label + ' PRIMITIVE'
    return out


@calcfunction
def get_refined_structure(structure: orm.StructureData, symprec: orm.Float, angle_tolerance: Any) -> orm.StructureData:
    """Create refined structure use pymatgen's interface"""
    from pymatgen.symmetry.analyzer import SpacegroupAnalyzer  # noqa: PLC0415

    pstruct = structure.get_pymatgen()
    ana = SpacegroupAnalyzer(pstruct, symprec=symprec.value, angle_tolerance=angle_tolerance.value)
    ps = ana.get_refined_structure()
    out = StructureData(pymatgen=ps)
    out.label = structure.label + ' REFINED'
    return out


@calcfunction
def get_conventional_standard_structure(structure: orm.StructureData, symprec: orm.Float, angle_tolerance: Any) -> Any:
    """Create conventional standard structure use pymatgen's interface"""
    from pymatgen.symmetry.analyzer import SpacegroupAnalyzer  # noqa: PLC0415

    pstruct = structure.get_pymatgen()
    ana = SpacegroupAnalyzer(pstruct, symprec=symprec.value, angle_tolerance=angle_tolerance.value)
    ps = ana.get_conventional_standard_structure()
    out = StructureData(pymatgen=ps)
    out.label = structure.label + ' CONVENTIONAL STANDARD'
    return out


@calcfunction
def make_supercell(structure: Any, supercell: Any, **kwargs: Any) -> dict[str, Any]:
    """Make supercell structure, keep the tags in order"""

    if 'tags' in kwargs:
        tags = kwargs['tags']
    else:
        tags = None

    atoms = structure.get_ase()
    atoms.set_tags(tags)

    slist = supercell.get_list()
    if isinstance(slist[0], int):
        satoms = atoms.repeat(slist)
    else:
        satoms = ase_supercell(atoms, np.array(slist))
    if 'no_sort' not in kwargs:
        satoms = sort(satoms)

    if tags:
        stags = satoms.get_tags().tolist()
    satoms.set_tags(None)

    out = StructureData(ase=satoms)
    out.label = structure.label + ' SUPER {} {} {}'.format(*slist)

    if tags:
        return {'structure': out, 'tags': orm.List(list=stags)}
    else:
        return {'structure': out}


@calcfunction
def niggli_reduce(structure: orm.StructureData) -> orm.StructureData:
    """Peroform niggli reduction using ase as the backend - this will rotate the structure into the standard setting"""
    atoms = structure.get_ase()
    niggli_reduce_(atoms)
    new_structure = StructureData(ase=atoms)
    new_structure.label = structure.label + ' NIGGLI'
    return new_structure


@calcfunction
def niggli_reduce_spglib(structure: orm.StructureData) -> orm.StructureData:
    """Peroform niggli reduction using spglib as backend - this does not rotate the structure"""
    atoms = structure.get_ase()
    reduced_cell = niggli_reduce_spg(atoms.cell)
    atoms.set_cell(reduced_cell)
    atoms.wrap()
    new_structure = StructureData(ase=atoms)
    new_structure.label = structure.label + ' NIGGLI'
    return new_structure


@calcfunction
def neb_interpolate(
    init_structure: orm.StructureData, final_strucrture: orm.StructureData, nimages: orm.Int
) -> dict[str, orm.StructureData]:
    """
    Interpolate NEB frames using the starting and the final structures

    Get around the PBC wrapping problem by calculating the MIC displacements
    from the initial to the final structure

    The initial structure is not changed, while the final structure is
    modified to be consistent with the initial structure in terms of absolute
    displacements, i.e. the final structure is *unwrapped*.
    """
    ainit = init_structure.get_ase()
    afinal = final_strucrture.get_ase()
    disps = []

    # Find distances
    acombined = ainit.copy()
    acombined.extend(afinal)
    # Get piece-wise MIC distances
    for i in range(len(ainit)):
        dist = acombined.get_distance(i, i + len(ainit), vector=True, mic=True)
        disps.append(dist.tolist())
    disps = np.asarray(disps)
    afinal = ainit.copy()

    # Displace the atoms according to MIC distances
    afinal.positions += disps
    neb = NEB([ainit.copy() for i in range(int(nimages) + 1)] + [afinal.copy()])
    neb.interpolate()
    out_init = StructureData(ase=neb.images[0])
    out_init.label = init_structure.label + ' INIT'
    out_final = StructureData(ase=neb.images[-1])
    out_final.label = init_structure.label + ' FINAL'

    outputs = {'image_init': out_init}
    for i, out in enumerate(neb.images[1:-1]):
        outputs[f'image_{i + 1:02d}'] = StructureData(ase=out)
        outputs[f'image_{i + 1:02d}'].label = init_structure.label + f' FRAME {i + 1:02d}'
    outputs['image_final'] = out_final
    return outputs


@calcfunction
def fix_atom_order(reference: orm.StructureData, to_fix: orm.StructureData) -> orm.StructureData:
    """
    Fix atom order by finding NN distances bet ween two frames. This resolves
    the issue where two closely matching structures having diffferent atomic orders.
    Note that the two frames must be close enough for this to work
    """

    aref = reference.get_ase()
    afix = to_fix.get_ase()

    # Index of the reference atom in the second structure
    new_indices = np.zeros(len(aref), dtype=int)

    # Find distances
    acombined = aref.copy()
    acombined.extend(afix)
    # Get piece-wise MIC distances
    for i in range(len(aref)):
        dists = []
        for j in range(len(aref)):
            dist = acombined.get_distance(i, j + len(aref), mic=True)
            dists.append(dist)
        min_idx = np.argmin(dists)
        min_dist = min(dists)
        if min_dist > 0.5:
            print(f'Large displacement found - moving atom {j} to {i} - please check if this is correct!')
        new_indices[i] = min_idx

    afixed = afix[new_indices]
    fixed_structure = StructureData(ase=afixed)
    fixed_structure.label = to_fix.label + ' UPDATED ORDER'
    return fixed_structure


def match_atomic_order_(atoms: Atoms, atoms_ref: Atoms) -> Tuple[Atoms, List[int]]:
    """
    Reorder the atoms to that of the reference.

    Only works for identical or nearly identical structures that are ordered differently.
    Returns a new `Atoms` object with order similar to that of `atoms_ref` as well as the sorting indices.
    """

    # Find distances
    acombined = atoms_ref.copy()
    acombined.extend(atoms)
    new_index = []
    # Get piece-wise MIC distances
    jidx = list(range(len(atoms), len(atoms) * 2))
    for i in range(len(atoms)):
        dists = acombined.get_distances(i, jidx, mic=True)
        # Find the index of the atom with the smallest distance
        min_idx = np.where(dists == dists.min())[0][0]
        new_index.append(min_idx)
    assert len(set(new_index)) == len(atoms), 'The detected mapping is not unique!'
    return atoms[new_index], new_index


def create_additional_species(species: list[str], magmoms: list[float]) -> tuple[list[str], dict[str, float]]:
    """
    Create additional species depending on magnetic moments.
    For example, create Fe1 and Fe2 if there are Fe with different
    magnetisations.

    :return: a tuples of (newspecies, magmom_mapping)
    """

    unique_species: set[str] = set(species)
    new_species: list[str] = []
    current_species_mapping: dict[str, dict[str, float]] = {sym: {} for sym in unique_species}
    for symbol, magmom in zip(species, magmoms):
        current_symbol: str = symbol
        # Mappings for this original symbol
        mapping: dict[str, float] = current_species_mapping[symbol]
        # First check if this magmom has been treated
        not_seen = True
        for sym_, mag_ in mapping.items():
            if mag_ == magmom:
                current_symbol = sym_
                not_seen = False
        # This symbol has not been seen yet
        if not_seen:
            if current_symbol in mapping:
                # The other species having the same symbol has been assigned
                counter = len(mapping) + 1
                current_symbol = f'{symbol}{counter}'
            mapping[current_symbol] = magmom
        new_species.append(current_symbol)

    # Rename symbols that has more than one species, so A becomes A1
    for symbol, mapping in current_species_mapping.items():
        if len(mapping) > 1:
            mapping[f'{symbol}1'] = mapping[symbol]
            mapping.pop(symbol)
            # Refresh the new_species list
            new_species = [f'{sym}1' if sym == symbol else sym for sym in new_species]

    all_mapping: dict[str, float] = {}
    for value in current_species_mapping.values():
        all_mapping.update(value)

    return new_species, all_mapping


def convert_to_plain_list(species: list[str], magmom_mapping: dict[str, float]) -> tuple[list[str], list[float]]:
    """
    Covert from a decorated species list to a plain list of symbols
    and magnetic moments.

    :return: A tuple of (symbols, magmoms)
    """
    magmoms = []
    symbols = []
    for symbol in species:
        magmoms.append(magmom_mapping[symbol])
        match = re.match(r'(\w+)\d+', symbol)
        if match:
            symbols.append(match.group(1))
        else:
            symbols.append(symbol)
    return symbols, magmoms
