#  Copyright (C) 2019 The authors of Py-ChemShell
#  This file was written by Thomas John, Rowan Hanson and Andrew Logsdail.
#  We are particularly grateful to support from CCP5 and the Nuffield Foundation
#
#  This file is part of Py-ChemShell.
#
#  Py-ChemShell is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation, either version 3 of the
#  License, or (at your option) any later version.
#
#  Py-ChemShell is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with Py-ChemShell.  If not, see
#  <http://www.gnu.org/licenses/>.

# originally created by Rowan Hanson of Cardiff University
# modified by AJL and TWK
# YL 01/02/2021 refactored

# YL 30/01/2021: changed to keyword to dim to match the Py-ChemShell style
#                added an arguement for PBC
# YL TODO 30/01/2021: move to a formal interface to ASE and rename to such as ase2chemsh() or ase.atom2frag()
# TD 14/03/2024: changed dim to ndimensions
def convert_atoms_to_frag(atoms, connect_mode="covalent",  ndimensions=3, pbc='', minus_ve=True, **kwargs):
    """
    Function to convert an ASE atoms object in to a ChemShell Fragment
    Parameters:
        atoms (Ase.Atoms):
            The atoms object containing coordinates, labels, cell vectors
        connect_mode (String):
            The connectivity for the new Fragment object
        ndimensions (Iteger):
            Dimensionality of the periodic system
            e.g., 2 if the atoms object is a model of a 2D surface
        pbc (String):
            Impose a periodic boundary condition

    """

    from ..objects import fragment
    from ..data    import unitconvert
    from ..base    import periodic

    # Get the atomic numbers
    atomZNums = atoms.get_atomic_numbers()

    # Get the atom positions
    atomPositions = atoms.positions*unitconvert.Angstrom2Bohr

    # YL 01/02/2021: now we also use a 3*3 matrix for all cell vectors 
    cellVectors = atoms.cell*unitconvert.Angstrom2Bohr

    # issue a warning if the old keyword is used
    if kwargs.get('twoDimensional', False):
        print(" >>> ChemShell/ASE: For function convert_atoms_to_frag() please use")
        print("                    keyword argument ndimensions=2 instead. twoDimensional")
        print("                    works for now but will raise an error in the future")
        ndimensions = 2

    pbc = periodic.verifyRequestedPBC(pbc, ndimensions)

    # Set the cell vectors for the fragment
    if ndimensions == 2:
        # We also need to shift the entire slab in a -ve direction, so the coordinates are below zero. (ChemShell prefered)
        # To do this, we find the maximum value z-axis value and subtract this from all other values
        # YL 01/02/2021: refactored
        # YL 01/02/2021: -ve is default for 2D but let's add a switch here
        if minus_ve:
            atomPositions[:,2] -= atomPositions[:,2].max()
        cellVectors[2,:] = 0.0

    # Create the fragment to be returned with the necessary arguments
    chemshFragment = fragment.Fragment(znums=atomZNums[:], coords=atomPositions[:])
    chemshFragment.cell.natoms  = len(atomZNums)
    chemshFragment.cell.vectors = cellVectors[:]
    chemshFragment.cell.pbc     = pbc
    chemshFragment.connect_mode = connect_mode

    # Return the fragment to the user
    return chemshFragment


def convert_frag_to_atoms(frag, minus_ve=True):
    """
    Function to convert a ChemShell Fragment object in to an ASE atoms object
    Parameters:
        frag (Chemsh.Fragment):
            The Fragment object containing coordinates, labels, cells vectors, etc.

    """    

    from chemsh.data import unitconvert
    # Attempt to import the Atoms() object
    # If failed, return an error
    try:
        from ase.atoms import Atoms
    except ModuleNotFoundError:
        print(">>> Py-ChemShell (ERROR): Please check ASE is installed correctly. [ModuleNotFoundError]")

    # Get the atomic numbers
    fragZNums = frag.znums

    # Get the atom positions, including unit conversion
    fragPositions = frag.coords*unitconvert.Bohr2Angstrom

    # Get vectors, if appropriate
    cellVectors = frag.cell.vectors*unitconvert.Bohr2Angstrom
    if frag.cell.ndimensions == 2:
        # We also need to manipulate the z-axis coordinates and give cell vectors
        # Traditionally, all z-axis values are +ve in an ASE atoms object, so here lets make sure that is the case as Fragments are -ve
        # This effectively normalises everything to zero for the lowest z-axis value
        # YL 01/02/2021: refactored
        if minus_ve:
            fragPositions[:,2] -= fragPositions[:,2].min()

        # This is a 2D model, needs a 3rd cell vector. Use the largest z-coordinate for this, assuming coordinates are +ve, and double for vacuum
        # YL 01/02/2021: we use a 3*3 matrix for all cell vectors 
        cellVectors[2,:] = [0.0, 0.0, 2.0*fragPositions[:,2].max()]

    # Create the Atoms() object with the necessary arguments
    aseAtoms = Atoms(fragZNums, positions=fragPositions, cell=cellVectors[:])

    # Return the object to the user
    return aseAtoms

