# -*- coding: utf-8 -*-
"""
Conversion between ``crystals`` data structures and other modules.

These functions are not expected to be used on their own; see the associated
 `Crystal` methods instead, like `Crystal.to_cif`.
"""
from contextlib import redirect_stdout
from io import StringIO
from itertools import groupby
from os import PathLike
from typing import Iterable, Optional

import numpy as np
from CifFile import CifBlock, CifFile

from . import __version__
from .atom import Atom

try:
    import ase
except ImportError:
    WITH_ASE = False
else:
    WITH_ASE = True

CIF_HEADER = f"""
#\\#CIF_2.0
#
# File generated by `crystals` {__version__}, documented at https://crystals.rtfd.io
# Reference:  L. P. René de Cotret et al, An open-source software ecosystem for the interactive exploration
#             of ultrafast electron scattering data, Advanced Structural and Chemical Imaging 4:11 (2018) DOI: 10.1186/s40679-018-0060-y.
#
# For more information on this type of file, please visit https://www.iucr.org/resources/cif
"""


# TODO: test against known XYZ file
def write_xyz(crystal: Iterable[Atom], fname: PathLike, comment: Optional[str] = None):
    """
    Generate an atomic coordinates .xyz file from a crystal structure.

    Parameters
    ----------
    crystal : crystals.Crystal
        Crystal to be converted.
    fname : path-like
        The XYZ file will be written to this file. If the file already exists,
        it will be overwritten.
    comment : str or None, optional
        Comment to include at the second line of ``fname``.
    """
    # Format is specified here:
    #   http://openbabel.org/wiki/XYZ_%28format%29
    comment = comment or ""

    with open(fname, "wt", encoding="ascii") as file:
        # First two lines are:
        #   1. Number of atoms described in the file
        #   2. Optional comment
        file.write(str(len(crystal)) + "\n")
        file.write(comment + "\n")

        # Write atomic data row-by-row
        # For easier human readability, atoms are sorted
        # by element
        for atom in sorted(crystal):
            x, y, z = atom.coords_cartesian
            row = f"  {atom.element:<2}       {x:10.5f}       {y:10.5f}       {z:10.5f}"
            file.write(row + "\n")


# TODO: write the asymmetric cell + symmetry operatrs
def write_cif(crystal: "Crystal", fname: PathLike):
    """
    Generate an atomic coordinates .cif file from a crystal structure.

    .. versionadded:: 1.2.0

    Parameters
    ----------
    crystal : crystals.Crystal
        Crystal to be converted.
    fname : path-like
        The CIF file will be written to this file. If the file already exists,
        it will be overwritten.
    comment : str or None, optional
        Comment to include at the second line of ``fname``.
    """
    cf = CifFile(strict=False)
    a, b, c, alpha, beta, gamma = crystal.lattice_parameters
    lattice_items = {
        "_cell_length_a": a,
        "_cell_length_b": b,
        "_cell_length_c": c,
        "_cell_angle_alpha": alpha,
        "_cell_angle_beta": beta,
        "_cell_angle_gamma": gamma,
    }

    sym = crystal.symmetry()
    symmetry_items = {
        "_symmetry_Int_Tables_number": sym["international_number"],
        "_symmetry_space_group_name_Hall": sym["hall_symbol"],
    }

    block = CifBlock()
    for key, val in lattice_items.items():
        block[key] = val

    for key, val in symmetry_items.items():
        block[key] = val

    # Note that we are using all atoms in the unit-cell,
    # and not the asymmetric unit cell + symmetry operators
    # This is valid CIF! And it is much simpler to implement
    # TODO: how to determine asymmetric cell + symmetry operations?
    atoms = list(crystal.unitcell)
    symbols = [atm.symbol for atm in atoms]
    xf = [atm.coords_fractional[0] for atm in atoms]
    yf = [atm.coords_fractional[1] for atm in atoms]
    zf = [atm.coords_fractional[2] for atm in atoms]

    block.CreateLoop(
        datanames=[
            "_atom_site_type_symbol",
            "_atom_site_fract_x",
            "_atom_site_fract_y",
            "_atom_site_fract_z",
        ],
        length_check=False,
    )
    block["_atom_site_type_symbol"] = symbols
    block["_atom_site_fract_x"] = xf
    block["_atom_site_fract_y"] = yf
    block["_atom_site_fract_z"] = zf

    # Name of the block cannot be empty!
    block_name = crystal.chemical_formula.replace(" ", "_")
    cf[block_name] = block

    # Converting to string writes to stdout for some reason
    with redirect_stdout(StringIO()):
        lines = str(cf).splitlines()

    with open(fname, "w", encoding="utf-8") as f:
        f.write(CIF_HEADER)
        f.write("\n".join(lines[13::]))  # Skip the fixed header


def ase_atoms(crystal: "Crystal", **kwargs):
    """
    Convert a :class:`crystals.Crystal` object into an :class:`ase.Atoms` object.
    Keyword arguments are passed to :class:`ase.Atoms` constructor.

    Parameters
    ----------
    crystal : :class:`Crystal`
        Crystal to be converted.

    Returns
    -------
    atoms : ase.Atoms
        Group of atoms ready for ASE's routines.

    Raises
    ------
    ImportError : If ASE is not installed
    """
    if not WITH_ASE:
        raise ImportError("ASE is not installed/importable.")

    return ase.Atoms(
        symbols=[
            ase.Atom(
                symbol=atom.element,
                position=atom.coords_cartesian,
                magmom=atom.magmom,
                mass=atom.mass,
            )
            for atom in crystal
        ],
        cell=np.array(crystal.lattice_vectors),
        **kwargs,
    )


def write_poscar(
    crystal: "Crystal",
    fname: PathLike,
    comment: Optional[str] = None,
    scaling_factor: float = 1.0,
    cartesian: bool = False,
):
    """
    Generate a POSCAR file from a crystal structure.

    Parameters
    ----------
    crystal : crystals.Crystal
        Crystal to be converted.
    fname : path-like
        The POSCAR file will be written to this file. If the file already exists,
        it will be overwritten.
    comment : str or None, optional
        Comment to include as the first line of ``fname``.
    scaling_factor: float, optional
        Scaling factor with which to scale lattice vectors
    cartesian: bool, optional
        If True, the resulting file will be expressed in cartesian coordinates.
        Default is to use fractional coordinates.
    """
    # Format is specified here:
    #   https://www.vasp.at/wiki/index.php/POSCAR
    comment = comment or "POSCAR"

    with open(fname, "wt", encoding="ascii") as file:
        file.write(comment + "\n")

        file.write(str(scaling_factor))
        file.write("\n")
        for vec in crystal.lattice_vectors:
            x, y, z = vec / scaling_factor
            file.write(f"{x:10.6f}  {y:10.6f}  {z:10.6f}")
            file.write("\n")

        grouped = [
            (elem, list(atoms))
            for elem, atoms in groupby(sorted(crystal), key=lambda atom: atom.element)
        ]

        file.write(" ".join(["%5s" % elem for elem, atoms in grouped]))
        file.write("\n")
        file.write(" ".join(["%5d" % len(atoms) for elem, atoms in grouped]))
        file.write("\n")

        if cartesian:
            file.write("Cartesian\n")
            for elem, atoms in grouped:
                for atom in atoms:
                    x, y, z = atom.coords_cartesian / scaling_factor
                    row = f"{x:10.5f}  {y:10.5f}  {z:10.5f}"
                    file.write(row + "\n")
        else:
            file.write("Direct\n")
            for elem, atoms in grouped:
                for atom in atoms:
                    x, y, z = atom.coords_fractional
                    row = f"{x:10.5f}  {y:10.5f}  {z:10.5f}"
                    file.write(row + "\n")
