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

"""Functions for handling CIF files"""

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

logger = logging.getLogger(__name__)

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


def cos(value):
    return math.cos(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}')
    lines.append(f'_chemical_formula_structural   {empirical_formula}')
    lines.append(f"_chemical_formula_sum   '{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(' _atom_site_type_symbol')
    lines.append(' _atom_site_label')
    lines.append(' _atom_site_symmetry_multiplicity')
    lines.append(' _atom_site_Cartn_x')
    lines.append(' _atom_site_Cartn_y')
    lines.append(' _atom_site_Cartn_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)

    for element, name, xyz in zip(
        atoms['elements'], names, atoms['coordinates']
    ):
        x, y, z = xyz
        lines.append(f'{element} {name}  1  {x} {y} {z}  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'_chemical_formula_structural   {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)

    for element, name, xyz in zip(
        atoms['elements'], names, atoms['coordinates']
    ):
        x, y, z = xyz
        lines.append(f'MOL1 {name} {element} {x} {y} {z} {x} {y} {z} 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)


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

    result = to_mmcif(system)
    print(result)
