from __future__ import annotations
import os

import numpy as np
import IMP.core
import IMP.atom

import scikit_fluorescence.io.zipped

from scikit_fluorescence import typing


pdb_keys_formats = [
    ('i', 'i4'),
    ('chain', '|U1'),
    ('res_id', 'i4'),
    ('res_name', '|U5'),
    ('atom_id', 'i4'),
    ('atom_name', '|U5'),
    ('element', '|U1'),
    ('xyz', '3f8'),
    ('charge', 'f8'),
    ('radius', 'f8'),
    ('bfactor', 'f8'),
    ('mass', 'f8')
]
keys, formats = list(zip(*pdb_keys_formats))


def read_coordinates(
    filename: str
) -> np.ndarray:
    """

    Examples
    --------
    >>> import scikit_fluorescence as skf
    >>> import scikit_fluorescence.io
    >>> atoms = skf.io.structure.read_coordinates('./test/data/1fat.cif')

    :param filename:
    :return:
    """
    model = IMP.Model()
    if not os.path.isfile(filename):
        raise FileNotFoundError("The file %s could not be found." % filename)
    if filename.upper().endswith('.PDB'):
        mp = IMP.atom.read_pdb(filename, model, IMP.atom.NonWaterPDBSelector())
        return convert_atoms(
            IMP.atom.get_by_type(mp, IMP.atom.ATOM_TYPE)
        )
    if filename.upper().endswith('.CIF'):
        print("Opening ci")
        mp = IMP.atom.read_mmcif(filename, model, IMP.atom.NonWaterPDBSelector())
        return convert_atoms(
            IMP.atom.get_by_type(mp, IMP.atom.ATOM_TYPE)
        )


def find_atom_index(
        atoms: np.array,
        chain_identifier: str,
        residue_seq_number: int,
        atom_name: str,
        residue_name: str,
        verbose: bool = False,
        ignore_multiple_selections: bool = True
):
    """
    Find the index of an atom by a set of identifiers

    :param atoms:
    :param chain_identifier:
    :param residue_seq_number:
    :param atom_name:
    :param residue_name:
    :param ignore_multiple_selections:
    :param verbose:
    :return:
    """
    # Determine Labeling position
    if residue_seq_number is None or atom_name is None:
        raise ValueError(
            "Either attachment_atom number or residue number and atom_name need to be provided."
        )
    if verbose:
        print("find_atom_index:")
        print("-- Chain ID: %s" % chain_identifier)
        print("-- Residue seq. number: %s" % residue_seq_number)
        print("-- Residue name: %s" % residue_name)
        print("-- Atom name: %s" % atom_name)
    if chain_identifier is None or chain_identifier == '':
        attachment_atom_index = np.where(
            (atoms['res_id'] == residue_seq_number) &
            (atoms['atom_name'] == atom_name)
        )[0]
        if verbose:
            print(
                "-- WARNING no chain specified. Possible attachment atoms: % s"
                % attachment_atom_index
            )
    else:
        attachment_atom_index = np.where(
            (atoms['res_id'] == residue_seq_number) &
            (atoms['atom_name'] == atom_name) &
            (atoms['chain'] == chain_identifier)
        )[0]
    if len(attachment_atom_index) != 1 and not ignore_multiple_selections:
        print("Search atom index:")
        print("-- Chain ID: %s" % chain_identifier)
        print("-- Residue seq. number: %s" % residue_seq_number)
        print("-- Residue name: %s" % residue_name)
        print("-- Atom name: %s" % atom_name)
        raise ValueError("Invalid selection of attachment atom")
    else:
        attachment_atom_index = attachment_atom_index[0]
    if verbose:
        print("Atom index: %s" % attachment_atom_index)
    return attachment_atom_index


def write_xyz(
        filename: str,
        points: np.array,
        verbose: bool = False
):
    """
    Writes the points as xyz-format file. The xyz-format file can be opened and displayed for instance
    in PyMol

    :param filename: string
    :param points: array
    :param verbose: bool

    """
    if verbose:
        print("write_xyz\n")
        print("Filename: %s\n" % filename)
    with scikit_fluorescence.io.zipped.open_maybe_zipped(
            filename=filename,
            mode='w'
    ) as fp:
        npoints = len(points)
        fp.write('%i\n' % npoints)
        fp.write('Name\n')
        for p in points:
            fp.write('D %.3f %.3f %.3f\n' % (p[0], p[1], p[2]))


def read_xyz(
        filename: str
) -> np.array:
    t = np.loadtxt(
        filename,
        skiprows=2,
        usecols=(1, 2, 3),
        delimiter=" "
    )
    return t


def write_pdb(
        filename: str,
        atoms=None,
        append_model: bool = False,
        append_coordinates: bool = False,
        verbose: bool = False
):
    """ Writes a structured numpy array containing the PDB-info to a PDB-file

    If append_model and append_coordinates are False the file is overwritten. Otherwise the atomic-coordinates
    are appended to the existing file.


    :param filename: target-filename
    :param atoms: structured numpy array
    :param append_model: bool
        If True the atoms are appended as a new models
    :param append_coordinates:
        If True the coordinates are appended to the file

    """
    mode = 'a+' if append_model or append_coordinates else 'w'
    if verbose:
        print("Writing to file: ", filename)
    with scikit_fluorescence.io.zipped.open_maybe_zipped(
            filename=filename,
            mode=mode
    ) as fp:
        # http://cupnet.net/pdb_format/
        al = [
            "%-6s%5d %4s%1s%3s %1s%4d%1s   %8.3f%8.3f%8.3f%6.2f%6.2f          %2s%2s\n" %
            (
                "ATOM ", at['atom_id'], at['atom_name'], " ", at['res_name'],
                at['chain'], at['res_id'], " ",
                at['xyz'][0], at['xyz'][1], at['xyz'][2], 0.0, at['bfactor'],
                at['element'], "  "
            )
            for at in atoms
        ]
        if append_model:
            fp.write('MODEL')
        fp.write("".join(al))
        if append_model:
            fp.write('ENDMDL')


def write_points(
        filename: str,
        points: np.ndarray,
        verbose: bool = False,
        mode: str = 'xyz',
        density: np.ndarray = None
):
    """

    :param filename:
    :param points:
    :param verbose:
    :param mode:
    :param density:
    :return:
    """
    if mode == 'pdb':
        atoms = np.empty(
            len(points),
            dtype={
                'names': keys,
                'formats': formats
            }
        )
        atoms['xyz'] = points
        if density is not None:
            atoms['bfactor'] = density
        write_pdb(
            filename,
            atoms,
            verbose=verbose
        )
    else:
        write_xyz(
            filename,
            points,
            verbose=verbose
        )


def convert_atoms(
        ps: typing.List[IMP.atom.Hierarchy],
        radius_no_interaction: bool = True
) -> np.ndarray:
    """Converts a list of IMP.atom.Hierarchy to a numpy record array

    Parameters
    ----------
    ps: list of IMP.atom.Hierarchy

    radius_no_interaction: bool
        If set to True the returned radii are the radii where the potential
        energy is zero. Otherwise, the radii are correspond to the minimal
        distance Rmin in a 6-12 LJ potential
        :math:`E = eij ((Rmin/rij)**12 - 2*(Rmin/rij)**6))`

    Returns
    -------
    atoms: numpy array containing the atom information

    """
    atoms = np.zeros(
        len(ps),
        dtype={
            'names': keys,
            'formats': formats
        }
    )
    radius_scaleling = 1.0
    if radius_no_interaction:
        # IMP uses the CHARM force field radius Rmin
        # E = eij ((Rmin/rij)**12 - 2*(Rmin/rij)**6))
        # The LJ radius is the distance where E = 0
        radius_scaleling = 2**(-1./6.)
    t = IMP.atom.get_element_table()
    for i, atom in enumerate(ps):
        a = IMP.atom.Atom(atom)
        r = IMP.atom.Residue(a.get_parent())
        c = IMP.atom.Chain(r.get_parent())
        atoms[i]['i'] = i
        atoms[i]['chain'] = c.get_id()
        atoms[i]['res_id'] = r.get_index()
        atoms[i]['res_name'] = r.get_name()
        atoms[i]['atom_id'] = a.get_input_index()
        atoms[i]['atom_name'] = str(a.get_atom_type())[1:-1]
        atoms[i]['element'] = t.get_name(a.get_element())
        atoms[i]['xyz'] = IMP.core.XYZR(atom).get_coordinates()
        atoms[i]['radius'] = IMP.core.XYZR(atom).get_radius() * radius_scaleling
        atoms[i]['bfactor'] = a.get_temperature_factor()
        atoms[i]['mass'] = IMP.atom.Mass(atom).get_mass()
    return atoms

