# Standard library
import os
import pickle
from glob import glob

# Third-party
import numpy as np
from astropy import units as u
from astropy.table import Table

# Project
from . import data_dir

__all__ = ['phot_system_list',
           'FilterSystem',
           'get_filter_names',
           'get_filter_properties',
           'phot_system_lookup',
           'load_zero_point_converter']

# list of photometric systems with pre-calculated filter properties
phot_system_list = [
    'HST_WFC3', 'HST_ACSWF', 'SDSSugriz', 'CFHTugriz', 'DECam', 'HSC',
    'JWST', 'LSST', 'UBVRIplus', 'UKIDSS', 'WFIRST'
]

class FilterSystem(object):
    """
    Class for calculating filter parameters from throughput curves.
    The parameter definitions are taken from
    `Fukugita et al. (1996) AJ 111, 1748.
    <https://ui.adsabs.harvard.edu/abs/1996AJ....111.1748F/abstract>`_

    Parameters
    ----------
    filter_curve_files : list
        List of the file names of the filter curves.
    filter_names : str
        Names of filters. Must be the same length as ``filter_curve_files``.
    **kwargs
        Optional argumnets for `~numpy.loadtxt`, which is used to load the
        filter curve files.
    """

    def __init__(self, filter_curve_files, filter_names, **kwargs):

        self.filter_names = filter_names
        for fn, name in zip(filter_curve_files, filter_names):
            data = np.loadtxt(fn, **kwargs)
            table = Table(data=data, names=['wave', 'trans'])
            setattr(self, name, table)
            self.filter_names.append(name)

    def _get_trans(self, bandpass):
        """
        Get the filter throughput curve for the given bandpass.
        """
        lam = getattr(self, bandpass)['wave']
        trans = getattr(self, bandpass)['trans']
        return lam, trans

    def effective_throughput(self, bandpass):
        """
        Calculate the effective throughput for the given bandpass.

        Parameters
        ----------
        bandpass : str
            The name of the filter.

        Returns
        -------
        theff : float
            Effective effective throughput.
        """
        lam, trans = self._get_trans(bandpass)
        theff = np.trapz(trans, np.log(lam))
        return theff

    def lam_eff(self, bandpass):
        """
        Calculate effective wavelength for the given bandpass.

        Parameters
        ----------
        bandpass : str
            The name of the filter.

        Returns
        -------
            leff : float
                The effective wavelength.
        """
        lam, trans = self._get_trans(bandpass)
        log_leff = np.trapz(np.log(lam) * trans, np.log(lam))
        log_leff /= np.trapz(trans, np.log(lam))
        leff =  np.exp(log_leff) * u.angstrom
        return leff

    def lam_pivot(self, bandpass):
        """
        Calculate pivot wavelength for the given bandpass.

        Parameters
        ----------
        bandpass : str
            The name of the filter.

        Returns
        -------
            lpivot : float
                The pivot wavelength.
        """
        lam, trans = self._get_trans(bandpass)
        lpivot = np.trapz(lam * trans, lam)
        lpivot /= np.trapz(trans, np.log(lam))
        lpivot = np.sqrt(lpivot) * u.angstrom
        return lpivot

    def dlam(self, bandpass):
        """
        Calculate the bandpass width.

        Parameters
        ----------
        bandpass : str
            The name of the filter.

        Returns
        -------
            width : float
                The bandpass width.
        """
        lam, trans = self._get_trans(bandpass)
        norm = np.trapz(trans, lam)
        width = (norm / trans.max()) * u.angstrom
        return width

class ZeroPointConverter(object):
    """
    Help class for converting between AB, ST, and Vega magnitudes & colors.

    Parameters
    ----------
    zpt_table : `~astropy.table.Table`
        MIST zero point table, which can be `downloaded here
        <http://waps.cfa.harvard.edu/MIST/BC_tables/zeropoints.txt>`_.
    """

    def __init__(self, zpt_table):
        for f, system, v_to_st, v_to_ab in zpt_table:
            setattr(self, f, [system, v_to_st, v_to_ab])
        self.zpt_table = zpt_table

    def to_vega(self, bandpass):
        """
        Convert to Vega magnitudes.

        Parameters
        ----------
        bandpass : str
            The name of the filter.

        Returns
        -------
        zpt_convert : float
            Zero point conversion magnitude.
        """
        system, v_to_st, v_to_ab = getattr(self, bandpass)
        if system == 'Vega':
            zpt_convert = 0.
        elif system == 'AB':
            zpt_convert = -v_to_ab
        return zpt_convert

    def to_ab(self, bandpass):
        """
        Convert to AB magnitudes.

        Parameters
        ----------
        bandpass : str
            The name of the filter.

        Returns
        -------
        zpt_convert : float
            Zero point conversion magnitude.
        """
        system, v_to_st, v_to_ab = getattr(self, bandpass)
        if system == 'AB':
            zpt_convert = 0.
        elif system == 'Vega':
            zpt_convert = v_to_ab
        return zpt_convert

    def to_st(self, bandpass):
        """
        Convert to ST magnitudes.

        Parameters
        ----------
        bandpass : str
            The name of the filter.

        Returns
        -------
        zpt_convert : float
            Zero point conversion magnitude.
        """
        system, v_to_st, v_to_ab = getattr(self, bandpass)
        if system == 'AB':
            zpt_convert = v_to_st - v_to_ab
        elif system == 'Vega':
            zpt_convert = v_to_st
        return zpt_convert

    def color_to_vega(self, blue, red):
        """
        Convert to Vega colors.

        Parameters
        ----------
        blue : str
            The name of the blue filter.
        red : str
            The name of the red filter.

        Returns
        -------
        zpt_convert : float
            Zero point conversion magnitude.
        """
        blue_convert = self.to_vega(blue)
        red_convert = self.to_vega(red)
        return blue_convert - red_convert

    def color_to_ab(self, blue, red):
        """
        Convert to AB colors.

        Parameters
        ----------
        blue : str
            The name of the blue filter.
        red : str
            The name of the red filter.

        Returns
        -------
        zpt_convert : float
            Zero point conversion magnitude.
        """
        blue_convert = self.to_ab(blue)
        red_convert = self.to_ab(red)
        return blue_convert - red_convert

    def color_to_st(self, blue, red):
        """
        Convert to ST colors.

        Parameters
        ----------
        blue : str
            The name of the blue filter.
        red : str
            The name of the red filter.

        Returns
        -------
        zpt_convert : float
            Zero point conversion magnitude.
        """
        blue_convert = self.to_st(blue)
        red_convert = self.to_st(red)
        return blue_convert - red_convert

def get_filter_names(phot_system=None):
    """
    Get MIST photometric systems and filter names.

    Parameters
    ----------
    phot_system : str or None
        The desired photometric system.

    Returns
    -------
    mist_filter_names : list or dict
        If ``phot_system`` given, then a list of filter names is returned. If
        ``phot_system`` is ``None``, then a dict of all photometric systems
        and filters is returned.
    """
    fn = os.path.join(data_dir, 'mist_filter_names.pkl')
    pickle_in = open(fn, 'rb')
    mist_filter_names = pickle.load(pickle_in)
    pickle_in.close()
    if phot_system is not None:
        if type(phot_system) == str:
            phot_system = [phot_system]
        names = []
        for p in phot_system:
            names.extend(mist_filter_names[p])
        mist_filter_names = names
    return mist_filter_names

def get_filter_properties():
    """

    """
    pkl_fn = os.path.join(data_dir, 'filter_properties.pkl')
    with open(pkl_fn, 'rb') as pkl_file:
        data = pickle.load(pkl_file)
        props = Table(data)
        props['dlam'] *= u.angstrom
        props['lam_eff'] *= u.angstrom
    return props

def phot_system_lookup(filter_name=None):
    """
    Lookup the photometric system name associated with a given filter name.

    Parameters
    ----------
    filter_name : str, optional
        Filter name to lookup.

    Returns
    -------
    lookup : dict or str
        If ``filter_name`` is ``None``, a dictionary with the filter names as
        keywords and the photometric systems as values. If ``filter_name`` is
        not ``None``, its photometric system is returned.
    """
    fn = os.path.join(data_dir, 'phot_system_lookup.pkl')
    pickle_in = open(fn, 'rb')
    lookup = pickle.load(pickle_in)
    pickle_in.close()
    if filter_name is not None:
        lookup = lookup[filter_name]
    return lookup

def load_zero_point_converter():
    """
    Create and return a `~artpop.filters.ZeroPointConverter` object.
    """
    from astropy.io import ascii
    fn = os.path.join(data_dir, 'zeropoints.txt')
    table= ascii.read(fn)
    return ZeroPointConverter(table)


# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License

import re
import numpy as np
import warnings
import collections
import os
from monty.io import zopen
import fnmatch
from collections import defaultdict
from pymatgen.electronic_structure.core import Spin, Orbital
from pymatgen.io.vasp.outputs import Vasprun
from pymatgen.electronic_structure.dos import Dos, LobsterCompleteDos
from pymatgen.core.structure import Structure
from pymatgen.electronic_structure.bandstructure import Kpoint
from pymatgen.io.vasp.inputs import Kpoints
from collections import OrderedDict
from pymatgen.electronic_structure.bandstructure import LobsterBandStructureSymmLine

"""
Module for reading Lobster output files. For more information
on LOBSTER see www.cohp.de.
"""

__author__ = "Marco Esters, Janine George"
__copyright__ = "Copyright 2017, The Materials Project"
__version__ = "0.2"
__maintainer__ = "Marco Esters, Janine George"
__email__ = "esters@uoregon.edu, janine.george@uclouvain.be"
__date__ = "Dec 13, 2017"

class Cohpcar:
    """
    Class to read COHPCAR/COOPCAR files generated by LOBSTER.

    Args:
        are_coops: Determines if the file is a list of COHPs or COOPs.
          Default is False for COHPs.

        filename: Name of the COHPCAR file. If it is None, the default
          file name will be chosen, depending on the value of are_coops.

    .. attribute: cohp_data

         Dict that contains the COHP data of the form:
           {bond: {"COHP": {Spin.up: cohps, Spin.down:cohps},
                   "ICOHP": {Spin.up: icohps, Spin.down: icohps},
                   "length": bond length,
                   "sites": sites corresponding to the bond}
         Also contains an entry for the average, which does not have
         a "length" key.

    .. attribute: efermi

         The Fermi energy in eV.

    .. attribute: energies

         Sequence of energies in eV. Note that LOBSTER shifts the energies
         so that the Fermi energy is at zero.

    .. attribute: is_spin_polarized

         Boolean to indicate if the calculation is spin polarized.

    .. attribute: orb_res_cohp

        orb_cohp[label] = {bond_data["orb_label"]: {"COHP": {Spin.up: cohps, Spin.down:cohps},
                                                     "ICOHP": {Spin.up: icohps, Spin.down: icohps},
                                                     "orbitals": orbitals,
                                                     "length": bond lengths,
                                                     "sites": sites corresponding to the bond}}

    """

    def __init__(self, are_coops=False, filename=None):
        self.are_coops = are_coops
        if filename is None:
            filename = "COOPCAR.lobster" if are_coops \
                else "COHPCAR.lobster"

        with zopen(filename, "rt") as f:
            contents = f.read().split("\n")

        # The parameters line is the second line in a COHPCAR file. It
        # contains all parameters that are needed to map the file.
        parameters = contents[1].split()
        # Subtract 1 to skip the average
        num_bonds = int(parameters[0]) - 1
        self.efermi = float(parameters[-1])
        if int(parameters[1]) == 2:
            spins = [Spin.up, Spin.down]
            self.is_spin_polarized = True
        else:
            spins = [Spin.up]
            self.is_spin_polarized = False

        # The COHP data start in row num_bonds + 3
        data = np.array([np.array(row.split(), dtype=float)
                         for row in contents[num_bonds + 3:]]).transpose()
        self.energies = data[0]

        cohp_data = {"average": {"COHP": {spin: data[1 + 2 * s * (num_bonds + 1)]
                                          for s, spin in enumerate(spins)},
                                 "ICOHP": {spin: data[2 + 2 * s * (num_bonds + 1)]
                                           for s, spin in enumerate(spins)}}}
        orb_cohp = {}

        # present for Lobster versions older than Lobster 2.2.0
        veryold = False
        # the labeling had to be changed: there are more than one COHP for each atom combination
        # this is done to make the labeling consistent with ICOHPLIST.lobster
        bondnumber = 0
        for bond in range(num_bonds):
            bond_data = self._get_bond_data(contents[3 + bond])

            label = str(bondnumber)

            orbs = bond_data["orbitals"]
            cohp = {spin: data[2 * (bond + s * (num_bonds + 1)) + 3]
                    for s, spin in enumerate(spins)}

            icohp = {spin: data[2 * (bond + s * (num_bonds + 1)) + 4]
                     for s, spin in enumerate(spins)}
            if orbs is None:
                bondnumber = bondnumber + 1
                label = str(bondnumber)
                cohp_data[label] = {"COHP": cohp, "ICOHP": icohp,
                                    "length": bond_data["length"],
                                    "sites": bond_data["sites"]}

            elif label in orb_cohp:
                orb_cohp[label].update({bond_data["orb_label"]:
                                            {"COHP": cohp,
                                             "ICOHP": icohp,
                                             "orbitals": orbs,
                                             "length": bond_data["length"],
                                             "sites": bond_data["sites"]}})
            else:
                # present for Lobster versions older than Lobster 2.2.0
                if bondnumber == 0:
                    veryold = True
                if veryold:
                    bondnumber += 1
                    label = str(bondnumber)

                orb_cohp[label] = {bond_data["orb_label"]: {"COHP": cohp,
                                                            "ICOHP": icohp,
                                                            "orbitals": orbs,
                                                            "length": bond_data["length"],
                                                            "sites": bond_data["sites"]}}

        # present for lobster older than 2.2.0
        if veryold:
            for bond in orb_cohp:
                cohp_data[bond] = {"COHP": None, "ICOHP": None,
                                   "length": bond_data["length"],
                                   "sites": bond_data["sites"]}

        self.orb_res_cohp = orb_cohp if orb_cohp else None
        self.cohp_data = cohp_data

    @staticmethod
    def _get_bond_data(line):
        """
        Subroutine to extract bond label, site indices, and length from
        a LOBSTER header line. The site indices are zero-based, so they
        can be easily used with a Structure object.

        Example header line: No.4:Fe1->Fe9(2.4524893531900283)
        Example header line for orbtial-resolved COHP:
            No.1:Fe1[3p_x]->Fe2[3d_x^2-y^2](2.456180552772262)

        Args:
            line: line in the COHPCAR header describing the bond.

        Returns:
            Dict with the bond label, the bond length, a tuple of the site
            indices, a tuple containing the orbitals (if orbital-resolved),
            and a label for the orbitals (if orbital-resolved).
        """

        orb_labs = ["s", "p_y", "p_z", "p_x", "d_xy", "d_yz", "d_z^2",
                    "d_xz", "d_x^2-y^2", "f_y(3x^2-y^2)", "f_xyz",
                    "f_yz^2", "f_z^3", "f_xz^2", "f_z(x^2-y^2)", "f_x(x^2-3y^2)"]

        line = line.rsplit("(", 1)
        # bondnumber = line[0].replace("->", ":").replace(".", ":").split(':')[1]
        length = float(line[-1][:-1])

        sites = line[0].replace("->", ":").split(":")[1:3]
        site_indices = tuple(int(re.split(r"\D+", site)[1]) - 1
                             for site in sites)

        # species = tuple(re.split(r"\d+", site)[0] for site in sites)
        if "[" in sites[0]:
            orbs = [re.findall(r"\[(.*)\]", site)[0] for site in sites]
            orbitals = [tuple((int(orb[0]), Orbital(orb_labs.index(orb[1:]))))
                        for orb in orbs]
            orb_label = "%d%s-%d%s" % (orbitals[0][0], orbitals[0][1].name,
                                       orbitals[1][0], orbitals[1][1].name)

        else:
            orbitals = None
            orb_label = None

        # a label based on the species alone is not feasible, there can be more than one bond for each atom combination
        # label = "%s" % (bondnumber)

        bond_data = {"length": length, "sites": site_indices,
                     "orbitals": orbitals, "orb_label": orb_label}
        return bond_data

class Icohplist:
    """
    Class to read ICOHPLIST/ICOOPLIST files generated by LOBSTER.

    Args:
        are_coops: Determines if the file is a list of ICOHPs or ICOOPs.
          Defaults to False for ICOHPs.

        filename: Name of the ICOHPLIST file. If it is None, the default
          file name will be chosen, depending on the value of are_coops.

    .. attribute: are_coops
         Boolean to indicate if the populations are COOPs or COHPs.

    .. attribute: is_spin_polarized
         Boolean to indicate if the calculation is spin polarized.

    .. attribute: Icohplist
        Dict containing the listfile data of the form:
           {bond: "length": bond length,
                  "number_of_bonds": number of bonds
                  "icohp": {Spin.up: ICOHP(Ef) spin up, Spin.down: ...}}

    .. attribute: IcohpCollection
        IcohpCollection Object

    """

    def __init__(self, are_coops=False, filename=None):

        self.are_coops = are_coops
        if filename is None:
            filename = "ICOOPLIST.lobster" if are_coops \
                else "ICOHPLIST.lobster"

        # LOBSTER list files have an extra trailing blank line
        # and we don't need the header.
        with zopen(filename) as f:
            data = f.read().split("\n")[1:-1]
        if len(data) == 0:
            raise IOError("ICOHPLIST file contains no data.")

        # Which Lobster version?
        if len(data[0].split()) == 8:
            version = '3.1.1'
        elif len(data[0].split()) == 6:
            version = '2.2.1'
            warnings.warn('Please consider using the new Lobster version. See www.cohp.de.')
        else:
            raise ValueError

        # If the calculation is spin polarized, the line in the middle
        # of the file will be another header line.
        if "distance" in data[len(data) // 2]:
            num_bonds = len(data) // 2
            if num_bonds == 0:
                raise IOError("ICOHPLIST file contains no data.")
            self.is_spin_polarized = True
        else:
            num_bonds = len(data)
            self.is_spin_polarized = False

        list_labels = []
        list_atom1 = []
        list_atom2 = []
        list_length = []
        list_translation = []
        list_num = []
        list_icohp = []
        for bond in range(num_bonds):
            line = data[bond].split()
            icohp = {}
            if version == '2.2.1':
                label = "%s" % (line[0])
                atom1 = str(line[1])
                atom2 = str(line[2])
                length = float(line[3])
                icohp[Spin.up] = float(line[4])
                num = int(line[5])
                translation = [0, 0, 0]
                if self.is_spin_polarized:
                    icohp[Spin.down] = float(data[bond + num_bonds + 1].split()[4])

            elif version == '3.1.1':
                label = "%s" % (line[0])
                atom1 = str(line[1])
                atom2 = str(line[2])
                length = float(line[3])
                translation = [int(line[4]), int(line[5]), int(line[6])]
                icohp[Spin.up] = float(line[7])
                num = int(1)

                if self.is_spin_polarized:
                    icohp[Spin.down] = float(data[bond + num_bonds + 1].split()[7])

            list_labels.append(label)
            list_atom1.append(atom1)
            list_atom2.append(atom2)
            list_length.append(length)
            list_translation.append(translation)
            list_num.append(num)
            list_icohp.append(icohp)

        # to avoid circular dependencies
        from pymatgen.electronic_structure.cohp import IcohpCollection
        self._icohpcollection = IcohpCollection(are_coops=are_coops, list_labels=list_labels, list_atom1=list_atom1,
                                                list_atom2=list_atom2, list_length=list_length,
                                                list_translation=list_translation, list_num=list_num,
                                                list_icohp=list_icohp, is_spin_polarized=self.is_spin_polarized)

    @property
    def icohplist(self):
        """
        Returns: icohplist compatible with older version of this class
        """
        icohplist_new = {}
        for key, value in self._icohpcollection._icohplist.items():
            icohplist_new[key] = {"length": value._length, "number_of_bonds": value._num,
                                  "icohp": value._icohp, "translation": value._translation}
        return icohplist_new

    @property
    def icohpcollection(self):
        """
        Returns: IcohpCollection object
        """
        return self._icohpcollection

class Doscar:
    """
    Class to deal with Lobster's projected DOS and local projected DOS.
    The beforehand quantum-chemical calculation was performed with VASP

    Args:
        doscar: DOSCAR filename, typically "DOSCAR.lobster"
        vasprun: vasprun filename, typically "vasprun.xml"
        dftprogram: so far only "vasp" is implemented

    .. attribute:: completedos

        LobsterCompleteDos Object

    .. attribute:: pdos
        List of Dict including numpy arrays with pdos. Access as pdos[atomindex]['orbitalstring']['Spin.up/Spin.down']

    .. attribute:: tdos
        Dos Object of the total density of states

    .. attribute:: energies
        numpy array of the energies at which the DOS was calculated (in eV, relative to Efermi)

    .. attribute:: tdensities
        tdensities[Spin.up]: numpy array of the total density of states for the Spin.up contribution at each of the energies
        tdensities[Spin.down]: numpy array of the total density of states for the Spin.down contribution at each of the energies

        if is_spin_polarized=False:
        tdensities[Spin.up]: numpy array of the total density of states

    .. attribute:: is_spin_polarized
        Boolean. Tells if the system is spin polarized

    """

    def __init__(self, doscar="DOSCAR.lobster", vasprun="vasprun.xml", dftprogram="Vasp"):

        self._doscar = doscar
        if dftprogram == "Vasp":
            self._vasprun = vasprun
            self._VASPRUN = Vasprun(filename=self._vasprun, ionic_step_skip=None,
                                    ionic_step_offset=0, parse_dos=False,
                                    parse_eigen=False, parse_projected_eigen=False,
                                    parse_potcar_file=False, occu_tol=1e-8,
                                    exception_on_bad_xml=True)
            self._final_structure = self._VASPRUN.final_structure
            self._is_spin_polarized = self._VASPRUN.is_spin

        self._parse_doscar()

    def _parse_doscar(self):
        doscar = self._doscar

        tdensities = {}
        f = open(doscar)
        natoms = int(f.readline().split()[0])
        efermi = float([f.readline() for nn in range(4)][3].split()[17])
        dos = []
        orbitals = []
        for atom in range(natoms + 1):
            line = f.readline()
            ndos = int(line.split()[2])
            orbitals.append(line.split(';')[-1].split())
            line = f.readline().split()
            cdos = np.zeros((ndos, len(line)))
            cdos[0] = np.array(line)
            for nd in range(1, ndos):
                line = f.readline().split()
                cdos[nd] = np.array(line)
            dos.append(cdos)
        f.close()

        doshere = np.array(dos[0])
        energies = doshere[:, 0]
        if not self._is_spin_polarized:
            tdensities[Spin.up] = doshere[:, 1]
            pdoss = []
            spin = Spin.up
            for atom in range(natoms):
                pdos = defaultdict(dict)
                data = dos[atom + 1]
                _, ncol = data.shape
                orbnumber = 0
                for j in range(1, ncol):
                    orb = orbitals[atom + 1][orbnumber]
                    pdos[orb][spin] = data[:, j]
                    orbnumber = orbnumber + 1
                pdoss.append(pdos)
        else:
            tdensities[Spin.up] = doshere[:, 1]
            tdensities[Spin.down] = doshere[:, 2]
            pdoss = []
            for atom in range(natoms):
                pdos = defaultdict(dict)
                data = dos[atom + 1]
                _, ncol = data.shape
                orbnumber = 0
                for j in range(1, ncol):
                    if j % 2 == 0:
                        spin = Spin.down
                    else:
                        spin = Spin.up
                    orb = orbitals[atom + 1][orbnumber]
                    pdos[orb][spin] = data[:, j]
                    if j % 2 == 0:
                        orbnumber = orbnumber + 1
                pdoss.append(pdos)
        self._efermi = efermi
        self._pdos = pdoss
        self._tdos = Dos(efermi, energies, tdensities)
        self._energies = energies
        self._tdensities = tdensities
        final_struct = self._final_structure

        pdossneu = {final_struct[i]: pdos for i, pdos in enumerate(self._pdos)}

        self._completedos = LobsterCompleteDos(final_struct, self._tdos, pdossneu)

    @property
    def completedos(self):
        return self._completedos

    @property
    def pdos(self):
        return self._pdos

    @property
    def tdos(self):
        return self._tdos

    @property
    def energies(self):
        return self._energies

    @property
    def tdensities(self):
        return self._tdensities

    @property
    def is_spin_polarized(self):
        return self._is_spin_polarized

class Charge:
    """Class to read CHARGE files generated by LOBSTER
        Args:
            filename: filename for the CHARGE file, typically "CHARGE.lobster"

        .. attribute: atomlist
            List of atoms in CHARGE.lobster
        .. attribute: types
            List of types of atoms in CHARGE.lobster
        .. attribute: Mulliken
            List of Mulliken charges of atoms in CHARGE.lobster
        .. attribute: Loewdin
            List of Loewdin charges of atoms in CHARGE.Loewdin
        .. attribute: num_atoms
            Number of atoms in CHARGE.lobster

    """

    def __init__(self, filename="CHARGE.lobster"):
        with zopen(filename) as f:
            data = f.read().split("\n")[3:-3]
        if len(data) == 0:
            raise IOError("CHARGES file contains no data.")

        self.num_atoms = len(data)
        self.atomlist = []
        self.types = []
        self.Mulliken = []
        self.Loewdin = []
        for atom in range(0, self.num_atoms):
            line = data[atom].split()
            self.atomlist.append(line[1] + line[0])
            self.types.append(line[1])
            self.Mulliken.append(float(line[2]))
            self.Loewdin.append(float(line[3]))

    def get_structure_with_charges(self, structure_filename):
        """
        get a Structure with Mulliken and Loewdin charges as site properties
        Args:
            structure_filename: filename of POSCAR
        Returns:
            Structure Object with Mulliken and Loewdin charges as site properties
        """

        struct = Structure.from_file(structure_filename)
        Mulliken = self.Mulliken
        Loewdin = self.Loewdin
        site_properties = {"Mulliken Charges": Mulliken, "Loewdin Charges": Loewdin}
        new_struct = struct.copy(site_properties=site_properties)
        return new_struct

class Lobsterout:
    """
    Class to read in the lobsterout and evaluate the spilling, save the basis, save warnings, save infos
    Args:
      filename: filename of lobsterout

      .. attribute: basis_functions
        list of basis functions that were used in lobster run as strings

      .. attribute: basis_type
         list of basis type that were used in lobster run as strings

      .. attribute: chargespilling
         list of charge spilling (first entry: result for spin 1, second entry: result for spin 2 or not present)

      .. attribute: dftprogram
         string representing the dft program used for the calculation of the wave function

      .. attribute: elements
         list of strings of elements that were present in lobster calculation

      .. attribute: has_CHARGE
         Boolean, indicates that CHARGE.lobster is present

      .. attribute: has_COHPCAR
        Boolean, indicates that COHPCAR.lobster and ICOHPLIST.lobster are present

      .. attribute: has_COOPCAR
        Boolean, indicates that COOPCAR.lobster and ICOOPLIST.lobster are present

      .. attribute: has_DOSCAR
        Boolean, indicates that DOSCAR.lobster is present

      .. attribute: has_Projection
        Boolean, indcates that projectionData.lobster is present

      .. attribute: has_bandoverlaps
        Boolean, indcates that bandOverlaps.lobster is present

      .. attribute: has_density_of_energies
        Boolean, indicates that DensityOfEnergy.lobster is present

      .. attribute: has_fatbands
        Boolean, indicates that fatband calculation was performed

      .. attribute: has_grosspopulation
        Boolean, indicates that GROSSPOP.lobster is present

      .. attribute: info_lines
        string with additional infos on the run

      .. attribute: info_orthonormalization
        string with infos on orthonormalization

      .. attribute: is_restart_from_projection
        Boolean that indicates that calculation was restartet from existing projection file

      .. attribute: lobster_version
        string that indicates Lobster version

      .. attribute: number_of_spins
        Integer indicating the number of spins

      .. attribute: number_of_threads
       integer that indicates how many threads were used

      .. attribute: timing
        dict with infos on timing

      .. attribute: totalspilling
        list of values indicating the total spilling for spin channel 1 (and spin channel 2)

      .. attribute: warninglines
        string with all warnings

    """

    def __init__(self, filename="lobsterout"):

        warnings.warn("Make sure the lobsterout is read in correctly. This is a brand new class.")
        # read in file
        with zopen(filename) as f:
            data = f.read().split("\n")  # [3:-3]
        if len(data) == 0:
            raise IOError("lobsterout does not contain any data")

        LobsterDict = {}
        # check if Lobster starts from a projection
        self.is_restart_from_projection = self._starts_from_projection(data=data)

        self.lobster_version = self._get_lobster_version(data=data)

        self.number_of_threads = int(self._get_threads(data=data))
        self.dftprogram = self._get_dft_program(data=data)

        self.number_of_spins = self._get_number_of_spins(data=data)
        chargespilling, totalspilling = self._get_spillings(data=data, number_of_spins=self.number_of_spins)
        self.chargespilling = chargespilling
        self.totalspilling = totalspilling

        elements, basistype, basisfunctions = self._get_elements_basistype_basisfunctions(data=data)
        self.elements = elements
        self.basis_type = basistype
        self.basis_functions = basisfunctions

        wall_time, user_time, sys_time = self._get_timing(data=data)
        timing = {}
        timing['walltime'] = wall_time
        timing['usertime'] = user_time
        timing['sys_time'] = sys_time
        self.timing = timing

        warninglines = self._get_all_warning_lines(data=data)
        self.warninglines = warninglines

        orthowarning = self._get_warning_orthonormalization(data=data)
        self.info_orthonormalization = orthowarning
        # print(orthowarning)

        infos = self._get_all_info_lines(data=data)
        self.info_lines = infos

        self.has_DOSCAR = self._has_DOSCAR(data=data)
        self.has_COHPCAR = self._has_COOPCAR(data=data)
        self.has_COOPCAR = self._has_COHPCAR(data=data)
        self.has_CHARGE = self._has_CHARGE(data=data)
        self.has_Projection = self._has_projection(data=data)
        self.has_bandoverlaps = self._has_bandoverlaps(data=data)
        self.has_fatbands = self._has_fatband(data=data)
        self.has_grosspopulation = self._has_grosspopulation(data=data)
        self.has_density_of_energies = self._has_density_of_energies(data=data)

    def get_doc(self):
        """

        Returns: LobsterDict with all the information stored in lobsterout
        """

        LobsterDict = {}
        # check if Lobster starts from a projection
        LobsterDict['restart_from_projection'] = self.is_restart_from_projection
        LobsterDict['lobster_version'] = self.lobster_version
        LobsterDict['threads'] = self.number_of_threads
        LobsterDict['Dftprogram'] = self.dftprogram

        number_of_spins = self.number_of_spins
        LobsterDict['chargespilling'] = self.chargespilling
        LobsterDict['totalspilling'] = self.totalspilling

        LobsterDict['elements'] = self.elements
        LobsterDict['basistype'] = self.basis_type
        LobsterDict['basisfunctions'] = self.basis_functions

        LobsterDict['timing'] = self.timing

        LobsterDict['warnings'] = self.warninglines

        LobsterDict['orthonormalization'] = self.info_orthonormalization

        LobsterDict['infos'] = self.info_lines

        LobsterDict['hasDOSCAR'] = self.has_DOSCAR
        LobsterDict['hasCOHPCAR'] = self.has_COHPCAR
        LobsterDict['hasCOOPCAR'] = self.has_COOPCAR
        LobsterDict['hasCHARGE'] = self.has_CHARGE
        LobsterDict['hasProjection'] = self.has_Projection
        LobsterDict['hasbandoverlaps'] = self.has_bandoverlaps
        LobsterDict['hasfatband'] = self.has_fatbands
        LobsterDict['hasGrossPopuliation'] = self.has_grosspopulation
        LobsterDict['hasDensityOfEnergies'] = self.has_density_of_energies

        return LobsterDict

    def _get_lobster_version(self, data):
        for row in data:
            splitrow = row.split()
            if len(splitrow) > 1:
                if splitrow[0] == "LOBSTER":
                    return splitrow[1]

    def _has_bandoverlaps(self, data):
        if 'WARNING: I dumped the band overlap matrices to the file bandOverlaps.lobster.' in data:
            return True
        else:
            return False

    def _starts_from_projection(self, data):
        if 'loading projection from projectionData.lobster...' in data:
            return True
        else:
            return False

    def _has_DOSCAR(self, data):
        if 'writing DOSCAR.lobster...' in data and not 'SKIPPING writing DOSCAR.lobster...' in data:
            return True
        else:
            return False

    def _has_COOPCAR(self, data):
        if 'writing COOPCAR.lobster and ICOOPLIST.lobster...' in data and not 'SKIPPING writing COOPCAR.lobster and ICOOPLIST.lobster...' in data:
            return True
        else:
            return False

    def _has_COHPCAR(self, data):
        if 'writing COHPCAR.lobster and ICOHPLIST.lobster...' in data and not 'SKIPPING writing COHPCAR.lobster and ICOHPLIST.lobster...' in data:
            return True
        else:
            return False

    def _has_CHARGE(self, data):
        # weitere optionen testen -> auch hier kann uebersprungen werden
        if not 'SKIPPING writing CHARGE.lobster...' in data:
            return True
        else:
            return False

    def _has_grosspopulation(self, data):
        if 'writing CHARGE.lobster and GROSSPOP.lobster...' in data:
            return True
        else:
            return False

    def _has_projection(self, data):
        if 'saving projection to projectionData.lobster...' in data:
            return True
        else:
            return False

    def _has_fatband(self, data):
        for row in data:
            splitrow = row.split()
            if len(splitrow) > 1:
                if splitrow[1] == 'FatBand':
                    return True
        return False

    def _has_density_of_energies(self, data):
        if "writing DensityOfEnergy.lobster..." in data:
            return True
        else:
            return False

    def _get_dft_program(self, data):
        for row in data:
            splitrow = row.split()
            if len(splitrow) > 4:
                if splitrow[3] == "program...":
                    return splitrow[4]

    def _get_number_of_spins(self, data):
        if "spillings for spin channel 2" in data:
            # print('two spin channels')
            return 2
        else:
            return 1

    def _get_threads(self, data):
        for row in data:
            splitrow = row.split()
            if len(splitrow) > 11:
                if (splitrow[11]) == "threads":
                    return splitrow[10]

    def _get_spillings(self, data, number_of_spins):
        charge_spilling = []
        total_spilling = []

        for row in data:
            splitrow = row.split()

            if len(splitrow) > 2:
                if splitrow[2] == 'spilling:':
                    # print(splitrow)
                    if splitrow[1] == 'charge':
                        charge_spilling.append(np.float(splitrow[3].replace('%', '')) / 100.0)
                    if splitrow[1] == 'total':
                        total_spilling.append(np.float(splitrow[3].replace('%', '')) / 100.0)

            if len(charge_spilling) == number_of_spins and len(total_spilling) == number_of_spins:
                break;

        return charge_spilling, total_spilling

    def _get_elements_basistype_basisfunctions(self, data):
        begin = False
        end = False
        elements = []
        basistype = []
        basisfunctions = []
        for row in data:

            if begin and not end:
                # print(row)
                splitrow = row.split()
                if splitrow[0] not in ['INFO:', 'WARNING:', 'setting', 'calculating', 'post-processing', 'saving',
                                       'spillings', 'writing']:

                    elements.append(splitrow[0])
                    basistype.append(splitrow[1].replace('(', '').replace(')', ''))
                    # last sign is a ''
                    basisfunctions.append(splitrow[2:])
                else:
                    end = True
            if "setting up local basis functions..." in row:
                begin = True
                # print(row)
        return elements, basistype, basisfunctions

    def _get_timing(self, data):
        # will give back wall, user and sys time
        begin = False
        # end=False
        # time=[]

        for row in data:
            splitrow = row.split()
            if 'finished' in splitrow:
                begin = True
            if begin:
                if 'wall' in splitrow:
                    wall_time = (splitrow[2:10])
                if 'user' in splitrow:
                    user_time = (splitrow[0:8])
                if 'sys' in splitrow:
                    sys_time = (splitrow[0:8])

        wall_time_dict = {"h": wall_time[0], "min": wall_time[2], "s": wall_time[4], "ms": wall_time[6]}
        user_time_dict = {"h": user_time[0], "min": user_time[2], "s": user_time[4], "ms": user_time[6]}
        sys_time_dict = {"h": sys_time[0], "min": sys_time[2], "s": sys_time[4], "ms": sys_time[6]}

        return wall_time_dict, user_time_dict, sys_time_dict

    def _get_warning_orthonormalization(self, data):
        orthowarning = []
        for row in data:
            splitrow = row.split()
            if 'orthonormalized' in splitrow:
                orthowarning.append(" ".join(splitrow[1:]))
        return orthowarning

    def _get_all_warning_lines(self, data):
        warnings = []
        for row in data:
            splitrow = row.split()
            if len(splitrow) > 0:
                if splitrow[0] == 'WARNING:':
                    warnings.append(" ".join(splitrow[1:]))

        return warnings
        # print(warnings)

    def _get_all_info_lines(self, data):
        infos = []
        for row in data:
            splitrow = row.split()
            if len(splitrow) > 0:
                if splitrow[0] == 'INFO:':
                    infos.append(" ".join(splitrow[1:]))

        return infos

class Fatband:
    """
    reads in FATBAND_x_y.lobster files
    Args:
      filenames (list or string): can be a list of file names or a path to a folder folder from which all "FATBAND_*" files will be read
      vasprun: corresponding vasprun file
      Kpointsfile: KPOINTS file for bandstructure calculation, typically "KPOINTS"

      .. attribute: efermi
      efermi that was read in from vasprun.xml

      .. attribute: eigenvals
      {Spin.up:[][],Spin.down:[][]}, the first index of the array
            [][] refers to the band and the second to the index of the
            kpoint. The kpoints are ordered according to the order of the
            kpoints array. If the band structure is not spin polarized, we
            only store one data set under Spin.up.

      .. attribute: is_spinpolarized
      Boolean that tells you whether this was a spin-polarized calculation

      .. attribute: kpoints_array
        list of kpoint as numpy arrays, in frac_coords of the
            given lattice by default

      .. attribute: label_dict
     (dict) of {} this link a kpoint (in frac coords or
            cartesian coordinates depending on the coords).

      .. attribute: lattice
      lattice object of reciprocal lattice as read in from vasprun.xml

      .. attribute: nbands
      number of bands used in the calculation

      .. attribute: p_eigenvals
      dict of orbital projections as {spin: array of dict}.
      The indices of the array are [band_index, kpoint_index].
      The dict is then built the following way:
      {"string of element": "string of orbital as read in from FATBAND file"}
      If the band structure is not spin polarized, we only store one data set under Spin.up.

      .. attribute: structure
      structure read in from vasprun.xml

    """

    def __init__(self, filenames=".", vasprun='vasprun.xml', Kpointsfile='KPOINTS'):

        warnings.warn('Make sure all relevant FATBAND files were generated and read in!')
        warnings.warn('Use Lobster 3.2.0 or newer for fatband calculations!')

        VASPRUN = Vasprun(filename=vasprun, ionic_step_skip=None,
                          ionic_step_offset=0, parse_dos=True,
                          parse_eigen=False, parse_projected_eigen=False,
                          parse_potcar_file=False, occu_tol=1e-8,
                          exception_on_bad_xml=True)
        self.structure = VASPRUN.final_structure
        self.lattice = self.structure.lattice.reciprocal_lattice
        self.efermi = VASPRUN.efermi
        kpoints_object = Kpoints.from_file(Kpointsfile)

        atomtype = []
        atomnames = []
        orbital_names = []

        if type(filenames) is not type([]) or filenames is None:
            filenames_new = []
            if filenames is None:
                filenames = '.'
            for file in os.listdir(filenames):
                if fnmatch.fnmatch(file, 'FATBAND_*.lobster'):
                    filenames_new.append(os.path.join(filenames, file))
            filenames = filenames_new
        if (len(filenames)) == 0:
            raise ValueError("No FATBAND files in folder or given")
        for ifilename, filename in enumerate(filenames):
            with zopen(filename, "rt") as f:
                contents = f.read().split("\n")

            # TODO: could be replaced for future versions of Lobster, get atomname from filename
            atomnames.append(os.path.split(filename)[1].split('_')[1].capitalize())
            parameters = contents[0].split()
            atomtype.append(re.split(r"[0-9]+", parameters[3])[0].capitalize())
            orbital_names.append(parameters[4])

        # get atomtype orbital dict
        atom_orbital_dict = {}
        for iatom, atom in enumerate(atomnames):
            if atom not in atom_orbital_dict:
                atom_orbital_dict[atom] = []
            atom_orbital_dict[atom].append(orbital_names[iatom])
        # test if there are the same orbitals twice or if two different formats were used or if all necessary orbitals are there
        for key, items in atom_orbital_dict.items():
            if len(set(items)) != len(items):
                raise (ValueError("The are two FATBAND files for the same atom and orbital. The program will stop."))
            split = []
            for item in items:
                split.append(item.split("_")[0])
            for orb, number in collections.Counter(split).items():
                if number != 1 and number != 3 and number != 5 and number != 7:
                    raise (ValueError(
                        "Make sure all relevant orbitals were generated and that no duplicates (2p and 2p_x) are present"))

        kpoints_array = []
        for ifilename, filename in enumerate(filenames):
            with zopen(filename, "rt") as f:
                contents = f.read().split("\n")

            if ifilename == 0:
                self.nbands = int(parameters[6])
                self.number_kpts = kpoints_object.num_kpts - int(contents[1].split()[2]) + 1

            if len(contents[1:]) == self.nbands + 2:
                self.is_spinpolarized = False
            elif len(contents[1:]) == self.nbands * 2 + 2:
                self.is_spinpolarized = True
            else:
                linenumbers = []
                for iline, line in enumerate(contents[1:self.nbands * 2 + 4]):
                    # print(line)
                    # if line in ['\n', '\r\n']:
                    #    linenumbers.append(iline)
                    if line.split()[0] == '#':
                        linenumbers.append(iline)

                if ifilename == 0:
                    if len(linenumbers) == 2:
                        self.is_spinpolarized = True
                    else:
                        self.is_spinpolarized = False

            if ifilename == 0:
                eigenvals = {}
                eigenvals[Spin.up] = [[collections.defaultdict(float)
                                       for i in range(self.number_kpts)]
                                      for j in range(self.nbands)]
                if self.is_spinpolarized:
                    eigenvals[Spin.down] = [[collections.defaultdict(float)
                                             for i in range(self.number_kpts)]
                                            for j in range(self.nbands)]

                p_eigenvals = {}
                p_eigenvals[Spin.up] = [
                    [{str(e): {str(orb): collections.defaultdict(float) for orb in atom_orbital_dict[e]}
                      for e in atomnames}
                     for i in range(self.number_kpts)]
                    for j in range(self.nbands)]

                if self.is_spinpolarized:
                    p_eigenvals[Spin.down] = [
                        [{str(e): {str(orb): collections.defaultdict(float) for orb in atom_orbital_dict[e]}
                          for e in atomnames}
                         for i in range(self.number_kpts)]
                        for j in range(self.nbands)]

            ikpoint = -1
            for iline, line in enumerate(contents[1:-1]):
                if line.split()[0] == '#':
                    Kpointnumber = int(line.split()[2])
                    KPOINT = np.array([float(line.split()[4]), float(line.split()[5]), float(line.split()[6])])
                    if ifilename == 0:
                        kpoints_array.append(KPOINT)

                    linenumber = 0
                    iband = 0
                    ikpoint += 1
                if linenumber == self.nbands:
                    iband = 0
                if line.split()[0] != '#':

                    if linenumber < self.nbands:
                        if ifilename == 0:
                            eigenvals[Spin.up][iband][ikpoint] = float(line.split()[1]) + self.efermi

                        p_eigenvals[Spin.up][iband][ikpoint][atomnames[ifilename]][orbital_names[ifilename]] = float(
                            line.split()[2])
                    if linenumber >= self.nbands and self.is_spinpolarized:
                        # print(line.split())
                        # print(iband)
                        if ifilename == 0:
                            eigenvals[Spin.down][iband][ikpoint] = float(line.split()[1]) + self.efermi
                        p_eigenvals[Spin.down][iband][ikpoint][atomnames[ifilename]][
                            orbital_names[ifilename]] = float(line.split()[2])

                    linenumber += 1
                    iband += 1

        self.kpoints_array = kpoints_array
        self.eigenvals = eigenvals
        self.p_eigenvals = p_eigenvals

        label_dict = {}
        for ilabel, label in enumerate(kpoints_object.labels[-self.number_kpts:], start=0):

            if label is not None:
                label_dict[label] = kpoints_array[ilabel]

        self.label_dict = label_dict

    def get_bandstructure(self):
        """
        returns a LobsterBandStructureSymmLine object which can be plotted with a normal BSPlotter
        """

        return LobsterBandStructureSymmLine(kpoints=self.kpoints_array, eigenvals=self.eigenvals, lattice=self.lattice,
                                            efermi=self.efermi, labels_dict=self.label_dict,
                                            structure=self.structure,
                                            projections=self.p_eigenvals)


"""Contains code for generating single polynucleotides."""

from collections import OrderedDict

import numpy
from tools.geometry import (angle_between_vectors, unit_vector, find_foot,
                            Quaternion, dihedral, find_transformations,
                            cylindrical_to_cartesian, Axis)
from ampal.base_ampal import Atom
from ampal.nucleic_acid import Polynucleotide, Nucleotide

_helix_parameters = {
    # nucleotides_per_turn, rise_per_nucleotide
    'a_dna': (11.0, 2.6),
    'b_dna': (10.2, 3.4),
    'z_dna': (12.3, 3.7)
}

_helix_handedness = {
    'a_dna': 'r',
    'b_dna': 'r',
    'z_dna': 'r'
}

_backbone_properties = {
    # helix_type : {atom_type: [radius, angular_offset (in radians), z-shift]}
    'a_dna': {
        'P': [8.92, 0, 0],  # P
    },
    'b_dna': {
        'mol_code_format': 'D{}',
        'labels': ["P", "OP1", "OP2", "O5'", "C5'", "C4'",
                   "O4'", "C3'", "O3'", "C2'", "C1'"],
        'atoms': [
            (8.92, 0.0, 0.0),
            (10.23, 0.07, 0.21),
            (9.0, -0.14, 0.77),
            (7.72, -0.5, -3.11),
            (7.68, -0.36, -4.05),
            (7.57, -0.17, -3.33),
            (6.22, -0.11, -3.17),
            (8.19, -0.19, -1.93),
            (8.74, -0.04, -1.55),
            (7.02, -0.25, -1.06),
            (5.79, -0.14, -1.83),
            (4.58, -0.30, -1.83)
        ]
    },
    'z_dna': {
        'P': [8.92, 0, 0],  # P
    }
}

_bases = {
    'G': {
        'labels': ["O4'", "C1'", "N9", "C8", "N7", "C5",
                   "C6", "O6", "N1", "C2", "N2", "N3", "C4"],
        'ref_atom': "N9",
        'rot_adj': 15.0,
        'atoms': [
            (33.930, 30.325, 17.814),
            (34.631, 29.403, 17.056),
            (35.908, 29.190, 17.720),
            (36.117, 28.636, 18.964),
            (37.376, 28.606, 19.309),
            (38.040, 29.196, 18.238),
            (39.419, 29.435, 18.041),
            (40.348, 29.171, 18.808),
            (39.672, 30.068, 16.810),
            (38.706, 30.393, 15.889),
            (39.131, 30.975, 14.756),
            (37.406, 30.170, 16.069),
            (37.150, 29.567, 17.255)
        ]
    },
    'C': {
        'labels': ["O4'", "C1'", "N1", "C2", "O2", "N3",
                   "C4", "N4", "C5", "C6"],
        'ref_atom': "N1",
        'rot_adj': 47.0,
        'atoms': [
            (35.422, 27.992, 13.920),
            (36.655, 27.383, 13.630),
            (37.240, 26.863, 14.903),
            (38.622, 26.960, 15.116),
            (39.337, 27.469, 14.237),
            (39.137, 26.506, 16.285),
            (38.336, 25.972, 17.205),
            (38.892, 25.545, 18.338),
            (36.925, 25.864, 17.008),
            (36.426, 26.317, 15.853)
        ]
    },
    'A': {
        'labels': ["O4'", "C1'", "N9", "C8", "N7", "C5",
                   "C6", "N6", "N1", "C2", "N3", "C4"],
        'ref_atom': "N9",
        'rot_adj': 5.0,
        'atoms': [
            (39.099, 26.091,  9.763),
            (39.785, 24.859,  9.897),
            (39.558, 24.371, 11.250),
            (38.407, 23.814, 11.737),
            (38.478, 23.462, 12.998),
            (39.769, 23.803, 13.357),
            (40.461, 23.687, 14.561),
            (39.913, 23.168, 15.662),
            (41.745, 24.130, 14.598),
            (42.280, 24.652, 13.481),
            (41.715, 24.823, 12.282),
            (40.447, 24.374, 12.291)
        ]
    },
    'T': {
        'labels': ["O4'", "C1'", "N1", "C2", "O2", "N3",
                   "C4", "O4", "C5", "C7", "C6"],
        'ref_atom': "N1",
        'rot_adj': 30.0,
        'atoms': [
            (43.43,  23.953,  9.349),
            (44.116, 22.801,  9.767),
            (43.234, 22.076, 10.744),
            (43.702, 21.817, 12.016),
            (44.824, 22.102, 12.389),
            (42.814, 21.185, 12.840),
            (41.527, 20.804, 12.542),
            (40.808, 20.241, 13.364),
            (41.084, 21.115, 11.196),
            (39.697, 20.746, 10.751),
            (41.948, 21.737, 10.374)
        ]
    }
}

class NucleicAcidStrand(Polynucleotide):
    """Generates a `Polynucleotide` from a sequence of bases.

    Parameters
    ----------
    sequence: str
        The nucleotide sequence of the nucleic acid.
    helix_type: str, optional
        The type of nucleic acid helix to generate.
    phos_3_prime: bool, optional
        If false the 5' and the 3' phosphor will be omitted.

    Attributes
    ----------
    base_sequence : str
        Nucleotide sequence for the nucleic acid.
    num_monomers : int
        Number of bases in the sequence.
    helix_type : str
        The type of nucleic acid helix. Currently only b-DNA is
        implemented.
    phos_3_prime : bool
        If `true`, the 3' end of the strand will have a phosphate.
    nucleotides_per_turn : float
        Number of nucleotides per turn of the nucleic acid super
        helix.
    rise_per_nucleotide : float
        Rise along the nucleic acid super helix per base.
    handedness : str
        Handedness of the super helix.
    helix_start : 3D Vector (tuple or list or numpy.array)
        Initial coordinate of the backbone primitive.
    helix_end : 3D Vector (tuple or list or numpy.array)
        Last coordinate of the backbone primitive.

    Raises
    ------
    ValueError
        Raised if a sequence is not provided.
    """

    def __init__(self, sequence='GATC', helix_type='b_dna', phos_3_prime=False):
        super().__init__()
        if not sequence:
            raise ValueError(
                'A sequence must be provided to build a region of DNA.')
        self.base_sequence = sequence
        self.num_monomers = len(sequence)
        self.helix_type = helix_type
        self.phos_3_prime = phos_3_prime
        self.nucleotides_per_turn, self.rise_per_nucleotide = _helix_parameters[
            self.helix_type]
        self.handedness = _helix_handedness[self.helix_type]
        self.helix_start = numpy.array([0.0, 0.0, 0.0])
        self.helix_end = self.helix_length * numpy.array([0.0, 0.0, 1.0])
        self.build()

    @classmethod
    def from_start_and_end(cls, start, end, sequence, helix_type='b_dna',
                           phos_3_prime=False):
        """Generates a helical `Polynucleotide` that is built along an axis.

        Parameters
        ----------
        start: [float, float, float]
            Start of the build axis.
        end: [float, float, float]
            End of build axis.
        sequence: str
            The nucleotide sequence of the nucleic acid.
        helix_type: str
            The type of nucleic acid helix to generate.
        phos_3_prime: bool
            If false the 5' and the 3' phosphor will be omitted.
        """
        start = numpy.array(start)
        end = numpy.array(end)
        instance = cls(sequence, helix_type=helix_type,
                       phos_3_prime=phos_3_prime)
        instance.move_to(start=start, end=end)
        return instance

    @property
    def axis(self):
        """The super-helical axis."""
        return Axis(start=self.helix_start, end=self.helix_end)

    @property
    def ax_unit(self):
        """The unit tangent of the super-helical axis."""
        return self.axis.unit_tangent

    @property
    def rad_unit(self):
        """The unit normal of the super-helical axis."""
        return self.axis.unit_normal

    @property
    def tan_unit(self):
        """The unit binormal of the super-helical axis."""
        return self.axis.unit_binormal

    @property
    def helix_length(self):
        """Length of the helix in Ångstroms."""
        return self.num_monomers * self.rise_per_nucleotide

    def translate(self, vector):
        super(Polynucleotide, self).translate(vector=vector)
        self.helix_start += vector
        self.helix_end += vector
        return

    def rotate(self, angle, axis, point=None, radians=False):
        super(Polynucleotide, self).rotate(
            angle=angle, axis=axis, point=point, radians=radians)
        # modify helix_start and helix_end accordingly.
        q = Quaternion.angle_and_axis(angle=angle, axis=axis, radians=radians)
        self.helix_start = q.rotate_vector(v=self.helix_start, point=point)
        self.helix_end = q.rotate_vector(v=self.helix_end, point=point)
        return

    def build(self):
        """Build single DNA strand along z-axis, starting with P on x-axis"""
        ang_per_res = (2 * numpy.pi) / self.nucleotides_per_turn
        atom_offset_coords = _backbone_properties[self.helix_type]['atoms']
        if self.handedness == 'l':
            handedness = -1
        else:
            handedness = 1
        base_atom_labels = _backbone_properties[self.helix_type]['labels']
        monomers = []
        mol_code_format = _backbone_properties[self.helix_type]['mol_code_format']
        for i, b in enumerate(self.base_sequence):
            nucleotide = Nucleotide(
                mol_code=mol_code_format.format(b), ampal_parent=self)
            atoms_dict = OrderedDict()
            if (i == (len(self.base_sequence) - 1)) and not self.phos_3_prime:
                # Do not include phosphate on last nucleotide
                atom_labels = base_atom_labels[3:] + [_bases[b]['labels'][2]]
                atom_offsets = {k: v for k, v in zip(
                    atom_labels, atom_offset_coords[3:])}
            else:
                atom_labels = base_atom_labels + [_bases[b]['labels'][2]]
                atom_offsets = {k: v for k, v in zip(
                    atom_labels, atom_offset_coords)}
            for atom_label in atom_labels:
                r, zeta, z_shift = atom_offsets[atom_label]
                rot_ang = ((i * ang_per_res) + zeta) * handedness
                z = (self.rise_per_nucleotide * i) + z_shift
                coords = cylindrical_to_cartesian(
                    radius=r, azimuth=rot_ang, z=z, radians=True)
                atom = Atom(
                    coordinates=coords, element=atom_label[0],
                    ampal_parent=nucleotide, res_label=atom_label)
                atoms_dict[atom_label] = atom
            base_ref = _bases[b]['ref_atom']
            rot_adj = _bases[b]['rot_adj']
            base_dict = OrderedDict(
                zip(_bases[b]['labels'], _bases[b]['atoms']))
            translation, angle, axis, point = find_transformations(
                base_dict[base_ref], base_dict["C1'"],
                atoms_dict[base_ref]._vector, atoms_dict["C1'"]._vector)
            q1 = Quaternion.angle_and_axis(angle, axis)
            # Align N9 C1'
            for k, v in base_dict.items():
                base_dict[k] = q1.rotate_vector(v, point) + translation
            # Rotate to align O4'
            axis = numpy.array(base_dict["C1'"]) - base_dict[base_ref]
            angle = dihedral(base_dict["O4'"], base_dict[base_ref],
                             base_dict["C1'"], atoms_dict["O4'"]) - rot_adj
            q2 = Quaternion.angle_and_axis(angle, axis)
            for k, v in list(base_dict.items()):
                if k not in atoms_dict:
                    atom = Atom(q2.rotate_vector(v, base_dict[base_ref]),
                                element=k[0], ampal_parent=nucleotide,
                                res_label=k)
                    atoms_dict[k] = atom
            nucleotide.atoms = atoms_dict
            monomers.append(nucleotide)
        self._monomers = monomers
        self.relabel_monomers()
        self.relabel_atoms()
        return

    def move_to(self, start, end):
        """Moves the `Polynucleotide` to lie on the `start` and `end` vector.

        Parameters
        ----------
        start : 3D Vector (tuple or list or numpy.array)
            The coordinate of the start of the helix primitive.
        end : 3D Vector (tuple or list or numpy.array)
            The coordinate of the end of the helix primitive.

        Raises
        ------
        ValueError
            Raised if `start` and `end` are very close together.
        """
        start = numpy.array(start)
        end = numpy.array(end)
        if numpy.allclose(start, end):
            raise ValueError('start and end must NOT be identical')
        translation, angle, axis, point = find_transformations(
            self.helix_start, self.helix_end, start, end)
        if not numpy.isclose(angle, 0.0):
            self.rotate(angle=angle, axis=axis, point=point, radians=False)
        self.translate(vector=translation)
        return

__author__ = 'Jack W. Heal, Christopher W. Wood'


"""
The main logic for building side chain with Monte Carlo Side Chain Ensemble algorithm,
 as defined in Bhowmick, Asmit, and Teresa Head-Gordon.
  "A Monte Carlo method for generating side chain structural ensembles."
  Structure 23.1 (2015): 44-55.

Coded by Jie Li
Date created: Jul 29, 2021
"""

import numpy as np
import numba as nb
from numba import jit, njit

from mcsce.core.rotamer_library import DunbrakRotamerLibrary
from mcsce.core.definitions import aa3to1
from mcsce.core.build_definitions import sidechain_templates
from mcsce.libs.libcalc import calc_torsion_angles, place_sidechain_template
from mcsce.libs.libscbuild import rotate_sidechain
from tqdm import tqdm
from copy import deepcopy
from functools import partial

from mcsce.libs.libstructure import Structure

_rotamer_library = DunbrakRotamerLibrary()

sidechain_placeholders = []
energy_calculators = []

@jit(nb.int32(nb.float64[:]), nopython=True, nogil=True)
def choose_random(array):
    cumulative = [np.sum(array[:i + 1]) for i in range(len(array))]
    pointer = np.random.random() * cumulative[-1]
    for idx, num in enumerate(cumulative):
        if num > pointer:
            return idx
    return len(array) - 1

def initialize_func_calc(efunc_creator, aa_seq=None, structure=None):
    """
    Helper function for initializing the energy function calculators according to the specified
    amino acid sequence or a structure object.

    Parameters
    ----------
    efunc_creator: partial function
        A partial function for creating energy functions for evaluating the generated conformations
        It should accept atom_label, res_num and res_label as inputs

    aa_seq: list
        List of 3-letter codes for the amino acid sequence

    structure: Structure object
        The structure with backbone atoms only
    """
    # declare global variabales
    global sidechain_placeholders
    global energy_calculators

    sidechain_placeholders = []
    energy_calculators = []
    print("Start preparing energy calculators at different sidechain completion levels")
    if structure is None:
        # create structure object according to the amino acid sequence
        fasta = "".join(aa3to1.get(res, 'X') for res in aa_seq)
        structure = Structure(fasta=fasta)
        structure.build()
    elif aa_seq is None:
        # extract amino acid sequence from structure object
        aa_seq = structure.residue_types
    structure = deepcopy(structure)
    sidechain_placeholders.append(deepcopy(structure))
    energy_calculators.append(efunc_creator(structure.atom_labels,
                                            structure.res_nums,
                                            structure.res_labels))
    for idx, resname in tqdm(enumerate(aa_seq), total=len(aa_seq)):
        template = sidechain_templates[resname]
        structure.add_side_chain(idx + 1, template)
        sidechain_placeholders.append(deepcopy(structure))
        if resname not in ["GLY", "ALA"]:
            n_sidechain_atoms = len(template[1])
            all_indices = np.arange(len(structure.atom_labels))
            energy_func = efunc_creator(structure.atom_labels,
                                        structure.res_nums,
                                        structure.res_labels,
                                        partial_indices=[all_indices[-n_sidechain_atoms:],
                                                         all_indices[:-n_sidechain_atoms]])
            energy_calculators.append(energy_func)
        else:
            energy_calculators.append(None)
    print("Finished preparing all energy functions. Now start conformer generation")

def create_side_chain_structure(inputs):
    """
    A function that takes a backbone-only structure and generates a conformation using the Monte Carlo approach

    Parameters
    ----------
    (parameters are packaged into a list to allow multiprocessing)
    backbone_coords: np.array with shape (N, 3)
        Backbone atom coordinates that defines the conformation for growing side chains
        Order of atoms are assumed to be ["N", "CA", "C", "O", "H1", "H2", "H3"] for N-terminal residue,
        ["N", "CA", "C", "O", "H"] for middle residues, and
        then ["N", "CA", "C", "O", "H", "OXT"] for C-terminal residue

    beta: float
        The beta value used for Boltzmann weighting

    save_addr: str or None
        The path for saving the generated structure. When None, the structure is not saved locally

    Returns
    ----------
    all_atom_structure: Structure object
        The structure including side chain atoms

    succeeded: bool
        Whether the side chain growth completed successfully

    energy:
        The energy for the generated conformation
    """
    backbone_coords, beta, save_addr = inputs
    assert len(sidechain_placeholders) > 0, "Energy functions have not yet initialized!"
    structure = deepcopy(sidechain_placeholders[0])
    structure.coords = backbone_coords
    N_CA_C_coords = structure.get_sorted_minimal_backbone_coords()
    all_backbone_dihedrals = calc_torsion_angles(N_CA_C_coords)
    all_phi = np.concatenate([[np.nan], all_backbone_dihedrals[2::3]]) * 180 / np.pi
    all_psi = np.concatenate([all_backbone_dihedrals[::3], [np.nan]]) * 180 / np.pi
    structure_coords = backbone_coords
    accumulated_energy = energy_calculators[0](structure_coords[None], structure_coords[None, :0])[0] # energies of backbone only
    for idx, resname in enumerate(structure.residue_types):
        # copy coordinates from the previous growing step to the current placeholder
        previous_coords = structure_coords
        # structure = deepcopy(sidechain_placeholders[idx])
        template_struc, sidechain_atom_idx = sidechain_templates[resname]
        n_sidechain_atoms = len(sidechain_atom_idx)
        new_coords = sidechain_placeholders[idx + 1].coords
        new_coords[:-n_sidechain_atoms] = previous_coords
        structure_coords = new_coords
        # coords = structure.coords
        residue_bb_coords = N_CA_C_coords[idx * 3: (idx + 1) * 3]
        if resname in ["GLY", "ALA"]:
            # For glycine and alanine, no degrees of freedom for bond rotations, so just move side chain to appropriate position
            sc_conformation = place_sidechain_template(residue_bb_coords, template_struc.coords)
            structure_coords[-n_sidechain_atoms:] = sc_conformation[sidechain_atom_idx]
            continue
        energy_func = energy_calculators[idx + 1]
        # get all candidate conformations (rotamers) for this side chain
        candidiate_conformations, candidate_probs = _rotamer_library.retrieve_torsion_and_prob(resname, all_phi[idx], all_psi[idx])
        # perturb chi angles of the side chains by ~0.5 degrees
        candidiate_conformations += np.random.normal(scale=0.5, size=candidiate_conformations.shape)

        energies = []

        all_coords = np.tile(structure_coords[None], (len(candidiate_conformations), 1, 1))
        # for each rotamer, decide the resulting conformation and calculate its energy
        for tor_idx, tors in enumerate(candidiate_conformations):
            sidechain_with_specific_chi = rotate_sidechain(resname, tors)  # THIS STEP MIGHT BE SLOW
            sc_conformation = place_sidechain_template(residue_bb_coords, sidechain_with_specific_chi[0])
            all_coords[tor_idx, -n_sidechain_atoms:] = sc_conformation[sidechain_atom_idx]
        energies = energy_func(all_coords[:, -n_sidechain_atoms:], all_coords[:, : -n_sidechain_atoms])
        minimum_energy = min(energies)  # Keep track of the minimum energy so that the renormalized energies can be converted back

        # print(idx, resname, len(candidate_probs), (~np.isinf(energies)).sum())
        # If all energies are inf, end this growth
        if np.isinf(energies).all():
            return None, False, None, None

        # renormalize energies to avoid numerical issues
        energies -= minimum_energy
        adjusted_weights = candidate_probs * np.exp(-beta * energies)
        selected_idx = choose_random(adjusted_weights)

        # set the coordinates of the actual structure that has side chain for this residue grown
        structure_coords = all_coords[selected_idx]
        accumulated_energy += energies[selected_idx] + minimum_energy # The raw energy for this step

    # all side chains have been created, then reorder all atoms and return the final structure
    structure = deepcopy(sidechain_placeholders[-1])
    structure.coords = structure_coords
    structure.reorder_with_resnum()
    if save_addr is not None:
        structure.write_PDB(save_addr)
    return structure, True, accumulated_energy, save_addr

def create_side_chain(structure, n_trials, temperature, parallel_worker=16, return_first_valid=False):
    """
    Using the MCSCE workflow to add sidechains to a backbone-only PDB structure. The building process will be repeated for n_trial times, but only the lowest energy conformation will be returned

    Parameters
    ----------
    structure: Structure object
        The structure with backbone atoms only

    n_trials: int
        The total number of trials for the generation procedure
        if n_trials <=0, then sequentially generate structures until the first valid structure is generated

    efunc_creator: partial function
        A partial function for creating energy functions for evaluating the generated conformations
        It should accept atom_label, res_num and res_label as inputs

    temperature: float
        The temperature value used for Boltzmann weighting

    parallel_worker: int
        Number of workers for parallel execution

    return_first_valid: bool
        Controls the behavior of whether execute parallel building and return the first valid structure (no clashes), instead of generating a collection of structures and return the lowest energy one

    Returns
    ----------
    lowest_energy_conformation: Structure object
        The generated lowest-energy conformation structure with sidechains, or None when every trial of the conformation generation failed
    """
    # convert temperature to beta
    beta = 1 / (temperature * 0.008314462) # k unit: kJ/mol/K

    conformations = []
    energies = []
    if return_first_valid:
        # Sequential execution with maximal n_trial times, but return the first valid structure
        for _ in range(n_trials):
            conf, succeeded, energy, _ = create_side_chain_structure([structure.coords, beta, None])
            if succeeded:
                return conf
        return None
    else:
        # Emsemble building with either sequential execution or parallelization
        if parallel_worker == 1:
            # sequential execution
            for idx in tqdm(range(n_trials)):
                conf, succeeded, energy, _ = create_side_chain_structure([structure.coords, beta, None])
                if succeeded:
                    conformations.append(conf)
                    energies.append(energy)
        else:
            import multiprocessing
            pool = multiprocessing.Pool(parallel_worker)

            with tqdm(total=n_trials) as pbar:
                for result in pool.imap_unordered(
                create_side_chain_structure, \
                [[structure.coords, beta, None]] * n_trials):
                    conf, succeeded, energy, _ = result
                    if succeeded:
                        conformations.append(conf)
                        energies.append(energy)
                    pbar.update()

            pool.close()
            pool.join()

        # define the lowest energy conformation and return
        if len(energies) == 0:
            return None
        else:
            lowest_energy_idx = np.argmin(energies)
            return conformations[lowest_energy_idx]

def create_side_chain_ensemble(structure, n_conformations, temperature, save_path, parallel_worker=16):
    """
    Create a given number of conformation ensemble for the backbone-only structure of a protein

    Parameters
    ----------
    structure: Structure object
        The structure with backbone atoms only

    n_conformations: int
        The total number of conformations to be generated

    temperature: float
        The temperature value used for Boltzmann weighting

    save_path: str
        The folder path for saving all succeessfully generated PDB files

    Returns
    ----------
    conformations: list of Structure object
        The generated conformation structures

    all_success_count: int
        Total number of the generated structures that are not early-stopped due to unresolvable clashes
    """

    # convert temperature to beta
    beta = 1 / (temperature * 0.008314462) # k unit: kJ/mol/K

    conformations = []
    success_indicator = []
    energies = {}

    if parallel_worker == 1:
        for idx in tqdm(range(n_conformations)):
            conf, succeeded, energy, save_dir = create_side_chain_structure([structure.coords, beta, save_path + f"/{idx}.pdb"])
            conformations.append(conf)
            success_indicator.append(succeeded)
            if succeeded:
                energies[save_dir] = energy
    else:
        # import pathos
        import multiprocessing
        # pool = pathos.multiprocessing.ProcessPool(parallel_worker)
        pool = multiprocessing.Pool(parallel_worker)
        # result_iterator = pool.starmap(create_side_chain_structure,
        # [[beta, save_path + f"/{n}.pdb"] for n in range(n_conformations)])
        # result_iterator = pool.uimap(create_side_chain_structure,
        #                     [beta] * n_conformations, [save_path + f"/{n}.pdb" for n in range(n_conformations)])

        with tqdm(total=n_conformations) as pbar:
            for result in pool.imap_unordered(create_side_chain_structure,\
                [[structure.coords, beta, save_path + f"/{n}.pdb"] for n in range(n_conformations)]):
                conformations.append(result[0])
                success_indicator.append(result[1])
                if result[1]:
                    # A succeeded case
                    energies[result[3]] = result[2]
                pbar.update()

        pool.close()
        pool.join()
    with open(save_path + "/energies.csv", "w") as f:
        f.write("File name,Energy(kJ/mol)\n")
        for item in energies:
            f.write("%s,%f\n" % (item, energies[item]))
    return conformations, success_indicator