# -*- coding: utf-8 -*-

"""Functions for handling CIF files"""

from collections import Counter
from functools import reduce
import logging
import math
import numpy

logger = logging.getLogger(__name__)

bond_order = {'single': 'SING', 'double': 'DOUB', 'triple': 'TRIP'}


def cos(value):
    return math.cos(math.radians(value))


def sin(value):
    return math.sin(math.radians(value))


def to_cif(system):
    """Transform a SEAMM structure to CIF file"""
    atoms = system['atoms']
    elements = atoms['elements']
    periodicity = system['periodicity']

    # Get the chemical formula
    counts = Counter(elements)

    # Order the elements ... Merck CH then alphabetical,
    # or if no C or H, then just alphabetically
    formula_list = []
    if 'C' in counts and 'H' in counts:
        formula_list.append(('C', counts.pop('C')))
        formula_list.append(('H', counts.pop('H')))

    for element in sorted(counts.keys()):
        formula_list.append((element, counts[element]))

    counts = []
    for _, count in formula_list:
        counts.append(count)

    formula = ''
    for element, count in formula_list:
        if formula != '':
            formula += ' '
        if count > 1:
            formula += element + str(count)
        else:
            formula += element

    # And the empirical formula
    divisor = reduce(math.gcd, counts)
    empirical_formula_list = []
    for element, count in formula_list:
        empirical_formula_list.append((element, int(count / divisor)))

    empirical_formula = ''
    for element, count in empirical_formula_list:
        if count > 1:
            empirical_formula += element + str(count)
        else:
            empirical_formula += element

    # And created the file, line-by-line
    lines = []
    lines.append('# Generated by MolSSI SEAMM')
    lines.append(f'data_{empirical_formula}')

    # Cell information
    if periodicity == 3:
        a, b, c, alpha, beta, gamma = system['cell']
        volume = (
            a * b * c * (1 - cos(alpha)**2 - cos(beta)**2 - cos(gamma)**2) +
            2 * math.sqrt(cos(alpha) * cos(beta) * cos(gamma))
        )
        lines.append("_symmetry_space_group_name_H-M   'P 1'")
        lines.append(f'_cell_length_a   {a}')
        lines.append(f'_cell_length_b   {b}')
        lines.append(f'_cell_length_c   {c}')
        lines.append(f'_cell_angle_alpha   {alpha}')
        lines.append(f'_cell_angle_beta    {beta}')
        lines.append(f'_cell_angle_gamma   {gamma}')
        lines.append('_symmetry_Int_Tables_number   1')
        lines.append(f'_cell_volume   {volume}')
        lines.append(f'_cell_formula_units_Z   {divisor}')
        lines.append('loop_')
        lines.append(' _symmetry_equiv_pos_site_id')
        lines.append(' _symmetry_equiv_pos_as_xyz')
        lines.append("  1  'x, y, z'")

    lines.append(f'_chemical_formula_structural   {empirical_formula}')
    lines.append(f"_chemical_formula_sum   '{formula}'")

    # The atoms
    lines.append('loop_')
    lines.append(' _atom_site_type_symbol')
    lines.append(' _atom_site_label')
    lines.append(' _atom_site_symmetry_multiplicity')
    lines.append(' _atom_site_fract_x')
    lines.append(' _atom_site_fract_y')
    lines.append(' _atom_site_fract_z')
    lines.append(' _atom_site_occupancy')

    # Need unique names
    if 'names' in atoms:
        original_names = atoms['names']
    else:
        original_names = atoms['elements']

    names = []
    tmp = {}
    for name in original_names:
        if name in tmp:
            tmp[name] += 1
            names.append(name + str(tmp[name]))
        else:
            tmp[name] = 1
            names.append(name)

    XYZ = into_cell(system)
    UVW = to_frac_explicit(system['cell'], XYZ)

    for element, name, uvw in zip(atoms['elements'], names, UVW):
        u, w, v = uvw
        lines.append(f'{element} {name}  1  {u:.3f} {v:.3f} {w:.3f}  1')

    # And that is it!
    return '\n'.join(lines)


def to_mmcif(system):
    """Transform a SEAMM structure to a MMCIF file"""
    atoms = system['atoms']
    elements = atoms['elements']
    bonds = system['bonds']
    periodicity = system['periodicity']

    # Get the chemical formula
    counts = Counter(elements)

    # Order the elements ... Merck CH then alphabetical,
    # or if no C or H, then just alphabetically
    formula_list = []
    if 'C' in counts and 'H' in counts:
        formula_list.append(('C', counts.pop('C')))
        formula_list.append(('H', counts.pop('H')))

    for element in sorted(counts.keys()):
        formula_list.append((element, counts[element]))

    counts = []
    for _, count in formula_list:
        counts.append(count)

    formula = ''
    for element, count in formula_list:
        if formula != '':
            formula += ' '
        if count > 1:
            formula += element + str(count)
        else:
            formula += element

    # And the empirical formula
    divisor = reduce(math.gcd, counts)
    empirical_formula_list = []
    for element, count in formula_list:
        empirical_formula_list.append((element, int(count / divisor)))

    empirical_formula = ''
    for element, count in empirical_formula_list:
        if count > 1:
            empirical_formula += element + str(count)
        else:
            empirical_formula += element

    # And created the file, line-by-line
    lines = []
    lines.append('# Generated by MolSSI SEAMM')
    lines.append(f'data_{empirical_formula}')
    lines.append(f"_chem_comp.name '{formula}'")
    lines.append(f"_chem_comp.id '{empirical_formula}'")
    lines.append(f"_chem_comp.formula   '{formula}'")

    # Cell information
    if periodicity == 3:
        a, b, c, alpha, beta, gamma = system['cell']
        volume = (
            a * b * c * (1 - cos(alpha)**2 - cos(beta)**2 - cos(gamma)**2) +
            2 * math.sqrt(cos(alpha) * cos(beta) * cos(gamma))
        )
        lines.append("_symmetry_space_group_name_H-M   'P 1'")
        lines.append(f'_cell_length_a   {a}')
        lines.append(f'_cell_length_b   {b}')
        lines.append(f'_cell_length_c   {c}')
        lines.append(f'_cell_angle_alpha   {alpha}')
        lines.append(f'_cell_angle_beta    {beta}')
        lines.append(f'_cell_angle_gamma   {gamma}')
        lines.append('_symmetry_Int_Tables_number   1')
        lines.append(f'_cell_volume   {volume}')
        lines.append(f'_cell_formula_units_Z   {divisor}')
        lines.append('loop_')
        lines.append(' _symmetry_equiv_pos_site_id')
        lines.append(' _symmetry_equiv_pos_as_xyz')
        lines.append("  1  'x, y, z'")

    # The atoms
    lines.append('loop_')
    lines.append(' _chem_comp_atom.comp_id')
    lines.append(' _chem_comp_atom.atom_id')
    lines.append(' _chem_comp_atom.type_symbol')
    lines.append(' _chem_comp_atom.model_Cartn_x')
    lines.append(' _chem_comp_atom.model_Cartn_y')
    lines.append(' _chem_comp_atom.model_Cartn_z')
    lines.append(' _chem_comp_atom.pdbx_model_Cartn_x_ideal')
    lines.append(' _chem_comp_atom.pdbx_model_Cartn_y_ideal')
    lines.append(' _chem_comp_atom.pdbx_model_Cartn_z_ideal')
    lines.append(' _chem_comp_atom.pdbx_component_comp_id')
    lines.append(' _chem_comp_atom.pdbx_residue_numbering')

    # Need unique names
    if 'names' in atoms:
        original_names = atoms['names']
    else:
        original_names = atoms['elements']

    names = []
    tmp = {}
    for name in original_names:
        if name in tmp:
            tmp[name] += 1
            names.append(name + str(tmp[name]))
        else:
            tmp[name] = 1
            names.append(name)

    if periodicity == 3:
        XYZa = into_cell(system)
    else:
        XYZa = atoms['coordinates']

    for element, name, xyza, xyz in zip(
        atoms['elements'], names, XYZa, atoms['coordinates']
    ):
        xa, ya, za = xyza
        x, y, z = xyz
        lines.append(
            f'MOL1 {name} {element} {xa:.3f} {ya:3f} {za:.3f} '
            f'{x:.3f} {y:.3f} {z:.3f} HET 1'
        )

    # The bonds
    lines.append('# ')
    lines.append('loop_')
    lines.append(' _chem_comp_bond.comp_id')
    lines.append(' _chem_comp_bond.atom_id_1')
    lines.append(' _chem_comp_bond.atom_id_2')
    lines.append(' _chem_comp_bond.value_order ')
    for i, j, order in bonds:
        nm1 = names[i - 1]
        nm2 = names[j - 1]
        value_order = bond_order[order]
        lines.append(f'MOL1 {nm1} {nm2} {value_order}')

    # And that is it!
    return '\n'.join(lines)


def to_cart(system):
    """Convert fractional coordinates to Cartesian.

    see https://en.wikipedia.org/wiki/Fractional_coordinates for a description.

    Parameters
    ----------
    system : System
        The incoming structure

    Returns
    -------
    xyz : [N][float*3]
        The Cartesian coordinates.
    """
    a, b, c, alpha, beta, gamma = system['cell']

    ca = cos(alpha)
    cb = cos(beta)
    cg = cos(gamma)
    sg = sin(gamma)

    V = a * b * c * math.sqrt(1 - ca**2 - cb**2 - cg**2 + 2 * ca * cb * cg)
    T = numpy.matrix(
        [
            [a, b * cg, c * cb],
            [0, b * sg, c * (ca - cb * cg) / sg],
            [0, 0, V / (a * b * sg)]
        ]
    )  # yapf: disable
    UVW = numpy.matrix(system['atoms']['coordinates'])
    XYZ = UVW * T.T

    return XYZ.tolist()


def to_cart_explicit(cell, uvw):
    """Convert fractional coordinates to Cartesian.

    see https://en.wikipedia.org/wiki/Fractional_coordinates for a description.

    Parameters
    ----------
    cell : Cell
        The cell parameters
    uvw : [[float]*3]*n_atoms
        The fractional coordinates

    Returns
    -------
    xyz : [N][float*3]
        The Cartesian coordinates.
    """
    a, b, c, alpha, beta, gamma = cell

    ca = cos(alpha)
    cb = cos(beta)
    cg = cos(gamma)
    sg = sin(gamma)

    V = a * b * c * math.sqrt(1 - ca**2 - cb**2 - cg**2 + 2 * ca * cb * cg)
    T = numpy.matrix(
        [
            [a, b * cg, c * cb],
            [0, b * sg, c * (ca - cb * cg) / sg],
            [0, 0, V / (a * b * sg)]
        ]
    )  # yapf: disable
    UVW = numpy.matrix(uvw)
    XYZ = UVW * T.T

    return XYZ.tolist()


def to_frac(system):
    """Convert Cartesian coordinates to fractional.

    see https://en.wikipedia.org/wiki/Fractional_coordinates for a description.

    Parameters
    ----------
    system : System
        The (molecular) system to convert

    Returns
    -------
    uvw : [N][float*3]
        The Fractional coordinates.
    """
    a, b, c, alpha, beta, gamma = system['cell']

    ca = cos(alpha)
    cb = cos(beta)
    cg = cos(gamma)
    sg = sin(gamma)

    V = a * b * c * math.sqrt(1 - ca**2 - cb**2 - cg**2 + 2 * ca * cb * cg)
    T = numpy.matrix(
        [
            [1 / a, -cg / (a * sg), b * c * (ca * cg - cb) / (V * sg)],
            [0, 1 / (b * sg), a * c * (cb * cg - ca) / (V * sg)],
            [0, 0, a * b * sg / V]
        ]
    )  # yapf: disable
    XYZ = numpy.matrix(system['atoms']['coordinates'])
    UVW = XYZ * T.T

    return UVW.tolist()


def to_frac_explicit(cell, xyz):
    """Convert Cartesian coordinates to fractional.

    see https://en.wikipedia.org/wiki/Fractional_coordinates for a description.

    Parameters
    ----------
    cell : Cell
        The cell parameters
    xyz : [[float]*3]*n_atoms
        The Cartesians coordinates

    Returns
    -------
    uvw : [N][float*3]
        The Fractional coordinates.
    """
    a, b, c, alpha, beta, gamma = cell

    ca = cos(alpha)
    cb = cos(beta)
    cg = cos(gamma)
    sg = sin(gamma)

    V = a * b * c * math.sqrt(1 - ca**2 - cb**2 - cg**2 + 2 * ca * cb * cg)
    T = numpy.matrix(
        [
            [1 / a, -cg / (a * sg), b * c * (ca * cg - cb) / (V * sg)],
            [0, 1 / (b * sg), a * c * (cb * cg - ca) / (V * sg)],
            [0, 0, a * b * sg / V]
        ]
    )  # yapf: disable
    XYZ = numpy.matrix(xyz)
    UVW = XYZ * T.T

    return UVW.tolist()


def into_cell(system, as_molecules=True):
    """Return the coordinates translated back into the principal unit cell.

    NB. One and only one set of coordinates must be given.

    Parameters
    ----------
    system : System

    Returns
    -------
    abc : [N][float*3]
        The transformed coordinates, either Cartesian or fractional,
        depending on which was input.
    """
    cell = system['cell']
    xyz = system['atoms']['coordinates']

    if 'bonds' not in system or len(system['bonds']) == 0:
        as_molecules = False

    use_xyz = xyz is not None
    if use_xyz:
        uvw = numpy.array(to_frac(system))

    if as_molecules:
        molecules = find_molecules(system)

        for indices in molecules:
            indices = numpy.array([i - 1 for i in indices])
            uvw_mol = numpy.take(uvw, indices, axis=0)
            center = numpy.average(uvw_mol, axis=0)
            delta = numpy.floor(center)
            uvw_mol -= delta
            numpy.put_along_axis(
                uvw, numpy.expand_dims(indices, axis=1), uvw_mol, axis=0
            )
    else:
        delta = numpy.floor(uvw)
        uvw -= delta

    if use_xyz:
        return to_cart_explicit(cell, uvw)
    else:
        return uvw.tolist()


def find_molecules(system):
    """Find the separate molecules in a system.

    Parameters
    ----------
    system : System
        The system definition

    Returns
    -------
    [][int] : A list of lists of indices for the molecules
    """

    if 'atoms' not in system:
        return []

    # Create a list of bonded atoms for each atom
    n_atoms = len(system['atoms']['elements'])
    if n_atoms == 0:
        return []

    if 'bonds' not in system:
        # No bonds, so just atoms....
        neighbors = [None]
        for i in range(n_atoms):
            neighbors.append(i + 1)  # 1-based at the moment

        return neighbors

    visited = [False] * (n_atoms + 1)
    visited[0] = True  # 1 offset stuff :-)
    neighbors = {i: [] for i in range(1, n_atoms + 1)}
    molecules = []

    for tmp in system['bonds']:
        i, j, _ = tmp
        neighbors[i].append(j)
        neighbors[j].append(i)

    while True:
        # Find first atom not yet visited
        try:
            i = visited.index(False)
        except ValueError:
            break
        atoms = [i]
        visited[i] = True
        next_atoms = neighbors[i]
        while len(next_atoms) > 0:
            tmp = []
            for i in next_atoms:
                if not visited[i]:
                    atoms.append(i)
                    visited[i] = True
                    tmp.extend(neighbors[i])
            next_atoms = tmp
        molecules.append(sorted(atoms))
    return molecules


if __name__ == "__main__":  # pragma: no cover
    system = {
        'atoms':
            {
                'coordinates':
                    [
                        (5.587, 2.245, 2.963), (5.888, 0.873, 3.131),
                        (6.399, 0.594, 2.352), (5.012, 2.591, 3.825),
                        (4.995, 2.381, 2.054), (6.516, 2.817, 2.895),
                        (3.245, 0.521, 1.0), (4.589, 0.126, 1.199),
                        (5.131, 0.622, 0.562), (2.608, -0.019, 1.705),
                        (2.94, 0.281, -0.021), (3.15, 1.595, 1.178),
                        (5.772, 5.176, 5.155), (4.563, 4.735, 4.567),
                        (3.965, 4.489, 5.293), (6.469, 5.462, 4.363),
                        (6.208, 4.368, 5.746), (5.574, 6.042, 5.792),
                        (3.334, 1.734, 5.872), (4.583, 1.858, 6.526),
                        (5.106, 2.506, 6.023), (2.719, 1.012, 6.415),
                        (3.489, 1.383, 4.849), (2.829, 2.704, 5.863),
                        (4.169, 5.213, 1.378), (3.523, 4.055, 0.884),
                        (3.209, 4.262, -0.012), (4.536, 5.013, 2.388),
                        (3.458, 6.043, 1.409), (5.012, 5.465, 0.731),
                        (0.468, 4.848, 1.229), (0.078, 3.672, 1.912),
                        (0.498, 2.923, 1.457), (0.012, 5.712, 1.718),
                        (0.129, 4.798, 0.191), (1.556, 4.947, 1.264),
                        (1.338, 5.396, 5.109), (2.246, 5.806, 4.104),
                        (1.792, 6.474, 3.563), (1.82, 4.643, 5.738),
                        (0.448, 4.964, 4.644), (1.061, 6.256, 5.724),
                        (0.897, 2.508, 4.153), (1.601, 3.524, 3.463),
                        (2.463, 3.623, 3.902), (-0.08, 2.371, 3.683),
                        (0.759, 2.802, 5.196), (1.459, 1.573, 4.098)
                    ],
                'elements':
                    [
                        'C', 'O', 'H', 'H', 'H', 'H', 'C', 'O', 'H', 'H', 'H',
                        'H', 'C', 'O', 'H', 'H', 'H', 'H', 'C', 'O', 'H', 'H',
                        'H', 'H', 'C', 'O', 'H', 'H', 'H', 'H', 'C', 'O', 'H',
                        'H', 'H', 'H', 'C', 'O', 'H', 'H', 'H', 'H', 'C', 'O',
                        'H', 'H', 'H', 'H'
                    ],
                'names':
                    [
                        'C', 'O', 'H', 'H', 'H', 'H', 'C', 'O', 'H', 'H', 'H',
                        'H', 'C', 'O', 'H', 'H', 'H', 'H', 'C', 'O', 'H', 'H',
                        'H', 'H', 'C', 'O', 'H', 'H', 'H', 'H', 'C', 'O', 'H',
                        'H', 'H', 'H', 'C', 'O', 'H', 'H', 'H', 'H', 'C', 'O',
                        'H', 'H', 'H', 'H'
                    ]
            },
        'bonds':
            [
                (1, 2, 'single'), (1, 4, 'single'), (1, 5, 'single'),
                (1, 6, 'single'), (2, 3, 'single'), (7, 8, 'single'),
                (7, 10, 'single'), (7, 11, 'single'), (7, 12, 'single'),
                (8, 9, 'single'), (13, 14, 'single'), (13, 16, 'single'),
                (13, 17, 'single'), (13, 18, 'single'), (14, 15, 'single'),
                (19, 20, 'single'), (19, 22, 'single'), (19, 23, 'single'),
                (19, 24, 'single'), (20, 21, 'single'), (25, 26, 'single'),
                (25, 28, 'single'), (25, 29, 'single'), (25, 30, 'single'),
                (26, 27, 'single'), (31, 32, 'single'), (31, 34, 'single'),
                (31, 35, 'single'), (31, 36, 'single'), (32, 33, 'single'),
                (37, 38, 'single'), (37, 40, 'single'), (37, 41, 'single'),
                (37, 42, 'single'), (38, 39, 'single'), (43, 44, 'single'),
                (43, 46, 'single'), (43, 47, 'single'), (43, 48, 'single'),
                (44, 45, 'single')
            ],
        'cell':
            [
                8.47201833303022, 8.47201833303022, 8.47201833303022, 90.0,
                90.0, 90.0
            ],
        'periodicity': 3,
        'units':
            {
                'coordinates': 'angstrom'
            }
    }  # yapf: disable

    # import pprint

    a, b, c, alpha, beta, gamma = system['cell']
    XYZ = system['atoms']['coordinates']
    tmp = []
    for xyz in XYZ[0:6]:
        x, y, z = xyz
        x += a
        tmp.append((x, y, z))
    for xyz in XYZ[6:12]:
        x, y, z = xyz
        x -= a
        tmp.append((x, y, z))
    for xyz in XYZ[12:18]:
        x, y, z = xyz
        x -= 2 * a
        y += b
        tmp.append((x, y, z))
    for xyz in XYZ[18:24]:
        x, y, z = xyz
        x -= a
        y -= b
        z += 3 * c
        tmp.append((x, y, z))
    for xyz in XYZ[24:]:
        tmp.append(xyz)

    system['atoms']['coordinates'] = tmp

    result = to_mmcif(system)
    with open('final.mmcif', 'w') as fd:
        print(result, file=fd)
    print(result)

    # pprint.pprint(to_frac(system))
    # print()
    # pprint.pprint(find_molecules(system))
    # print()
    # pprint.pprint(into_cell(system))
