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

"""
This module defines classes to represent crystal orbital Hamilton
populations (COHP) and integrated COHP (ICOHP), but can also be used
for crystal orbital overlap populations (COOP).
"""

import warnings
import re
import sys

import numpy as np

from monty.json import MSONable

from pymatgen.electronic_structure.core import Spin, Orbital
from pymatgen.core.sites import PeriodicSite
from pymatgen.core.structure import Structure
from pymatgen.io.lmto import LMTOCopl
from pymatgen.io.lobster import Cohpcar
from pymatgen.util.num import round_to_sigfigs
from pymatgen.util.coord import get_linear_interpolated_value

__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 Cohp(MSONable):
    """
    Basic COHP object.
    """

    def __init__(self, efermi, energies, cohp, are_coops=False, icohp=None):
        """
        Args:
            are_coops: Indicates whether this object describes COHPs or COOPs.
            efermi: Fermi energy.
            energies: A sequence of energies.
            cohp ({Spin: np.array}): representing the COHP for each spin.
            icohp ({Spin: np.array}): representing the ICOHP for each spin.
        """
        self.are_coops = are_coops
        self.efermi = efermi
        self.energies = np.array(energies)
        self.cohp = cohp
        self.icohp = icohp

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        """
        Returns a string that can be easily plotted (e.g. using gnuplot).
        """
        cohpstring = "COOP" if self.are_coops else "COHP"
        header = ["Energy", cohpstring + "Up"]
        data = [self.energies, self.cohp[Spin.up]]
        if Spin.down in self.cohp:
            header.append(cohpstring + "Down")
            data.append(self.cohp[Spin.down])
        if self.icohp:
            header.append("I" + cohpstring + "Up")
            data.append(self.icohp[Spin.up])
            if Spin.down in self.cohp:
                header.append("I" + cohpstring + "Down")
                data.append(self.icohp[Spin.down])
        formatheader = "#" + " ".join(["{:15s}" for __ in header])
        formatdata = " ".join(["{:.5f}" for __ in header])
        stringarray = [formatheader.format(*header)]
        for i, __ in enumerate(self.energies):
            stringarray.append(formatdata.format(*[d[i] for d in data]))
        return "\n".join(stringarray)

    def as_dict(self):
        """
        Json-serializable dict representation of COHP.
        """
        d = {"@module": self.__class__.__module__,
             "@class": self.__class__.__name__,
             "are_coops": self.are_coops,
             "efermi": self.efermi,
             "energies": self.energies.tolist(),
             "COHP": {str(spin): pops.tolist()
                      for spin, pops in self.cohp.items()}}
        if self.icohp:
            d["ICOHP"] = {str(spin): pops.tolist()
                          for spin, pops in self.icohp.items()}
        return d

    def get_cohp(self, spin=None, integrated=False):
        """
        Returns the COHP or ICOHP for a particular spin.

        Args:
            spin: Spin. Can be parsed as spin object, integer (-1/1)
                or str ("up"/"down")
            integrated: Return COHP (False) or ICOHP (True)

        Returns:
            Returns the CHOP or ICOHP for the input spin. If Spin is
            None and both spins are present, both spins will be returned
            as a dictionary.
        """
        if not integrated:
            populations = self.cohp
        else:
            populations = self.icohp

        if populations is None:
            return None
        if spin is None:
            return populations
        if isinstance(spin, int):
            spin = Spin(spin)
        elif isinstance(spin, str):
            s = {"up": 1, "down": -1}[spin.lower()]
            spin = Spin(s)
        return {spin: populations[spin]}

    def get_icohp(self, spin=None):
        """
        Convenient alternative to get the ICOHP for a particular spin.
        """
        return self.get_cohp(spin=spin, integrated=True)

    def get_interpolated_value(self, energy, integrated=False):
        """
        Returns the COHP for a particular energy.

        Args:
            energy: Energy to return the COHP value for.
        """
        inter = {}
        for spin in self.cohp:
            if not integrated:
                inter[spin] = get_linear_interpolated_value(self.energies,
                                                            self.cohp[spin],
                                                            energy)
            elif self.icohp is not None:
                inter[spin] = get_linear_interpolated_value(self.energies,
                                                            self.icohp[spin],
                                                            energy)
            else:
                raise ValueError("ICOHP is empty.")
        return inter

    def has_antibnd_states_below_efermi(self, spin=None, limit=0.01):
        """
        Returns dict indicating if there are antibonding states below the Fermi level depending on the spin
            spin: Spin
            limit: -COHP smaller -limit will be considered.

        """
        warnings.warn("This method has not been tested on many examples. Check the parameter limit, pls!")

        populations = self.cohp
        number_energies_below_efermi = len([x for x in self.energies if x <= self.efermi])

        if populations is None:
            return None
        if spin is None:
            dict_to_return = {}
            for sp, cohpvalues in populations.items():
                if (max(cohpvalues[0:number_energies_below_efermi])) > limit:
                    dict_to_return[sp] = True
                else:
                    dict_to_return[sp] = False
        else:
            dict_to_return = {}
            if isinstance(spin, int):
                spin = Spin(spin)
            elif isinstance(spin, str):
                s = {"up": 1, "down": -1}[spin.lower()]
                spin = Spin(s)
            if (max(populations[spin][0:number_energies_below_efermi])) > limit:
                dict_to_return[spin] = True
            else:
                dict_to_return[spin] = False

        return dict_to_return

    @classmethod
    def from_dict(cls, d):
        """
        Returns a COHP object from a dict representation of the COHP.

        """
        if "ICOHP" in d:
            icohp = {Spin(int(key)): np.array(val)
                     for key, val in d["ICOHP"].items()}
        else:
            icohp = None
        return Cohp(d["efermi"], d["energies"],
                    {Spin(int(key)): np.array(val)
                     for key, val in d["COHP"].items()},
                    icohp=icohp, are_coops=d["are_coops"])

class CompleteCohp(Cohp):
    """
    A wrapper class that defines an average COHP, and individual COHPs.

    .. attribute: are_coops

         Indicates whether the object is of COOPs or COHPs.

    .. attribute: efermi

         Fermi energy

    .. attribute: energies

         Sequence of energies

    .. attribute: structure

         Structure associated with the COHPs.

    .. attribute: cohp, icohp

         The average COHP/ICOHP.

    .. attribute: all_cohps

         A dict of COHPs for individual bonds of the form {label: COHP}

    .. attribute: orb_res_cohp

        Orbital-resolved COHPs.
    """

    def __init__(self, structure, avg_cohp, cohp_dict, bonds=None,
                 are_coops=False, orb_res_cohp=None):
        """
        Args:
            structure: Structure assosciated with this COHP.
            avg_cohp: The average cohp as a COHP object.
            cohps: A dict of COHP objects for individual bonds of the form
                {label: COHP}
            bonds: A dict containing information on the bonds of the form
                {label: {key: val}}. The key-val pair can be any information
                the user wants to put in, but typically contains the sites,
                the bond length, and the number of bonds. If nothing is
                supplied, it will default to an empty dict.
            are_coops: indicates whether the Cohp objects are COHPs or COOPs.
                Defauls to False for COHPs.
            orb_res_cohp: Orbital-resolved COHPs.
        """
        super().__init__(avg_cohp.efermi, avg_cohp.energies, avg_cohp.cohp,
                         are_coops=are_coops, icohp=avg_cohp.icohp)
        self.structure = structure
        self.are_coops = are_coops
        self.all_cohps = cohp_dict
        self.orb_res_cohp = orb_res_cohp
        if bonds is None:
            self.bonds = {label: {} for label in self.all_cohps.keys()}
        else:
            self.bonds = bonds

    def __str__(self):
        if self.are_coops:
            return "Complete COOPs for " + str(self.structure)
        return "Complete COHPs for " + str(self.structure)

    def as_dict(self):
        """
        Json-serializable dict representation of CompleteCohp.
        """
        d = {"@module": self.__class__.__module__,
             "@class": self.__class__.__name__,
             "are_coops": self.are_coops,
             "efermi": self.efermi,
             "structure": self.structure.as_dict(),
             "energies": self.energies.tolist(),
             "COHP": {"average": {str(spin): pops.tolist()
                                  for spin, pops in
                                  self.cohp.items()}}}

        if self.icohp is not None:
            d["ICOHP"] = {"average": {str(spin): pops.tolist()
                                      for spin, pops in
                                      self.icohp.items()}}

        for label in self.all_cohps.keys():
            d["COHP"].update({label: {str(spin): pops.tolist()
                                      for spin, pops in
                                      self.all_cohps[label].cohp.items()}})
            if self.all_cohps[label].icohp is not None:
                if "ICOHP" not in d.keys():
                    d["ICOHP"] = {label: {str(spin): pops.tolist()
                                          for spin, pops in
                                          self.all_cohps[label].icohp.items()}}
                else:
                    d["ICOHP"].update({label: {str(spin): pops.tolist()
                                               for spin, pops in
                                               self.all_cohps[label].icohp.items()}})
        if False in [bond_dict == {} for bond_dict in self.bonds.values()]:
            d["bonds"] = {bond: {"length": self.bonds[bond]["length"],
                                 "sites": [site.as_dict() for site
                                           in self.bonds[bond]["sites"]]}
                          for bond in self.bonds}
        if self.orb_res_cohp:
            orb_dict = {}
            for label in self.orb_res_cohp:
                orb_dict[label] = {}
                for orbs in self.orb_res_cohp[label]:
                    cohp = {str(spin): pops.tolist() for spin, pops in
                            self.orb_res_cohp[label][orbs]["COHP"].items()}
                    orb_dict[label][orbs] = {"COHP": cohp}
                    icohp = {str(spin): pops.tolist() for spin, pops in
                             self.orb_res_cohp[label][orbs]["ICOHP"].items()}
                    orb_dict[label][orbs]["ICOHP"] = icohp
                    orbitals = [[orb[0], orb[1].name] for orb in
                                self.orb_res_cohp[label][orbs]["orbitals"]]
                    orb_dict[label][orbs]["orbitals"] = orbitals
            d["orb_res_cohp"] = orb_dict

        return d

    def get_cohp_by_label(self, label):
        """
        Get specific COHP object.

        Args:
            label: string (for newer Lobster versions: a number)

        Returns:
            Returns the COHP object to simplify plotting
        """
        if label.lower() == "average":
            return Cohp(efermi=self.efermi, energies=self.energies,
                        cohp=self.cohp, are_coops=self.are_coops, icohp=self.icohp)
        return Cohp(efermi=self.efermi, energies=self.energies,
                    cohp=self.all_cohps[label].get_cohp(spin=None, integrated=False),
                    are_coops=self.are_coops,
                    icohp=self.all_cohps[label].get_icohp(spin=None))

    def get_summed_cohp_by_label_list(self, label_list, divisor=1):
        """
        Returns a COHP object that includes a summed COHP divided by divisor

        Args:
            label_list: list of labels for the COHP that should be included in the summed cohp
            divisor: float/int, the summed cohp will be divided by this divisor
        Returns:
            Returns a COHP object including a summed COHP
        """
        # check if cohps are spinpolarized or not
        first_cohpobject = self.get_cohp_by_label(label_list[0])
        summed_cohp = first_cohpobject.cohp.copy()
        summed_icohp = first_cohpobject.icohp.copy()
        for label in label_list[1:]:
            cohp_here = self.get_cohp_by_label(label)
            summed_cohp[Spin.up] = np.sum([summed_cohp[Spin.up], cohp_here.cohp[Spin.up]], axis=0)
            if Spin.down in summed_cohp:
                summed_cohp[Spin.down] = np.sum([summed_cohp[Spin.down], cohp_here.cohp[Spin.down]], axis=0)
            summed_icohp[Spin.up] = np.sum([summed_icohp[Spin.up], cohp_here.icohp[Spin.up]], axis=0)
            if Spin.down in summed_icohp:
                summed_icohp[Spin.down] = np.sum([summed_icohp[Spin.down], cohp_here.icohp[Spin.down]], axis=0)

        divided_cohp = {}
        divided_icohp = {}
        divided_cohp[Spin.up] = np.divide(summed_cohp[Spin.up], divisor)
        divided_icohp[Spin.up] = np.divide(summed_icohp[Spin.up], divisor)
        if Spin.down in summed_cohp:
            divided_cohp[Spin.down] = np.divide(summed_cohp[Spin.down], divisor)
            divided_icohp[Spin.down] = np.divide(summed_icohp[Spin.down], divisor)

        return Cohp(efermi=first_cohpobject.efermi, energies=first_cohpobject.energies, cohp=divided_cohp,
                    are_coops=first_cohpobject.are_coops,
                    icohp=divided_icohp)

    def get_summed_cohp_by_label_and_orbital_list(self, label_list, orbital_list, divisor=1):
        """
        Returns a COHP object that includes a summed COHP divided by divisor

        Args:
            label_list: list of labels for the COHP that should be included in the summed cohp
            orbital_list: list of orbitals for the COHPs that should be included in the summed cohp (same order as
                label_list)
            divisor: float/int, the summed cohp will be divided by this divisor
        Returns:
            Returns a COHP object including a summed COHP
        """
        # check length of label_list and orbital_list:
        if not len(label_list) == len(orbital_list):
            raise ValueError("label_list and orbital_list don't have the same length!")
        # check if cohps are spinpolarized or not
        first_cohpobject = self.get_orbital_resolved_cohp(label_list[0], orbital_list[0])
        summed_cohp = first_cohpobject.cohp.copy()
        summed_icohp = first_cohpobject.icohp.copy()
        for ilabel, label in enumerate(label_list[1:], 1):
            cohp_here = self.get_orbital_resolved_cohp(label, orbital_list[ilabel])
            summed_cohp[Spin.up] = np.sum([summed_cohp[Spin.up], cohp_here.cohp.copy()[Spin.up]], axis=0)
            if Spin.down in summed_cohp:
                summed_cohp[Spin.down] = np.sum([summed_cohp[Spin.down], cohp_here.cohp.copy()[Spin.down]], axis=0)
            summed_icohp[Spin.up] = np.sum([summed_icohp[Spin.up], cohp_here.icohp.copy()[Spin.up]], axis=0)
            if Spin.down in summed_icohp:
                summed_icohp[Spin.down] = np.sum([summed_icohp[Spin.down], cohp_here.icohp.copy()[Spin.down]], axis=0)

        divided_cohp = {}
        divided_icohp = {}
        divided_cohp[Spin.up] = np.divide(summed_cohp[Spin.up], divisor)
        divided_icohp[Spin.up] = np.divide(summed_icohp[Spin.up], divisor)
        if Spin.down in summed_cohp:
            divided_cohp[Spin.down] = np.divide(summed_cohp[Spin.down], divisor)
            divided_icohp[Spin.down] = np.divide(summed_icohp[Spin.down], divisor)

        return Cohp(efermi=first_cohpobject.efermi, energies=first_cohpobject.energies, cohp=divided_cohp,
                    are_coops=first_cohpobject.are_coops,
                    icohp=divided_icohp)

    def get_orbital_resolved_cohp(self, label, orbitals):
        """
        Get orbital-resolved COHP.

        Args:
            label: bond label (Lobster: labels as in ICOHPLIST/ICOOPLIST.lobster).

            orbitals: The orbitals as a label, or list or tuple of the form
                [(n1, orbital1), (n2, orbital2)]. Orbitals can either be str,
                int, or Orbital.

        Returns:
            A Cohp object if CompleteCohp contains orbital-resolved cohp,
            or None if it doesn't.

        Note: It currently assumes that orbitals are str if they aren't the
            other valid types. This is not ideal, but the easiest way to
            avoid unicode issues between python 2 and python 3.
        """
        if self.orb_res_cohp is None:
            return None
        if isinstance(orbitals, (list, tuple)):
            cohp_orbs = [d["orbitals"] for d in
                         self.orb_res_cohp[label].values()]
            orbs = []
            for orbital in orbitals:
                if isinstance(orbital[1], int):
                    orbs.append(tuple((orbital[0], Orbital(orbital[1]))))
                elif isinstance(orbital[1], Orbital):
                    orbs.append(tuple((orbital[0], orbital[1])))
                elif isinstance(orbital[1], str):
                    orbs.append(tuple((orbital[0], Orbital[orbital[1]])))
                else:
                    raise TypeError("Orbital must be str, int, or Orbital.")
            orb_index = cohp_orbs.index(orbs)
            orb_label = list(self.orb_res_cohp[label].keys())[orb_index]
        elif isinstance(orbitals, str):
            orb_label = orbitals
        else:
            raise TypeError("Orbitals must be str, list, or tuple.")
        try:
            icohp = self.orb_res_cohp[label][orb_label]["ICOHP"]
        except KeyError:
            icohp = None
        return Cohp(self.efermi, self.energies,
                    self.orb_res_cohp[label][orb_label]["COHP"],
                    icohp=icohp, are_coops=self.are_coops)

    @classmethod
    def from_dict(cls, d):
        """
        Returns CompleteCohp object from dict representation.
        """
        cohp_dict = {}
        efermi = d["efermi"]
        energies = d["energies"]
        structure = Structure.from_dict(d["structure"])
        if "bonds" in d.keys():
            bonds = {bond: {"length": d["bonds"][bond]["length"],
                            "sites": tuple(PeriodicSite.from_dict(site)
                                           for site in d["bonds"][bond]["sites"])}
                     for bond in d["bonds"]}
        else:
            bonds = None
        for label in d["COHP"]:
            cohp = {Spin(int(spin)): np.array(d["COHP"][label][spin])
                    for spin in d["COHP"][label]}
            try:
                icohp = {Spin(int(spin)): np.array(d["ICOHP"][label][spin])
                         for spin in d["ICOHP"][label]}
            except KeyError:
                icohp = None
            if label == "average":
                avg_cohp = Cohp(efermi, energies, cohp, icohp=icohp)
            else:
                cohp_dict[label] = Cohp(efermi, energies, cohp, icohp=icohp)

        if "orb_res_cohp" in d.keys():
            orb_cohp = {}
            for label in d["orb_res_cohp"]:
                orb_cohp[label] = {}
                for orb in d["orb_res_cohp"][label]:
                    cohp = {Spin(int(s)): np.array(d["orb_res_cohp"][label][orb]["COHP"][s], dtype=float)
                            for s in d["orb_res_cohp"][label][orb]["COHP"]}
                    try:
                        icohp = {Spin(int(s)): np.array(d["orb_res_cohp"][label][orb]["ICOHP"][s], dtype=float)
                                 for s in d["orb_res_cohp"][label][orb]["ICOHP"]}
                    except KeyError:
                        icohp = None
                    orbitals = [tuple((int(o[0]), Orbital[o[1]])) for o in
                                d["orb_res_cohp"][label][orb]["orbitals"]]
                    orb_cohp[label][orb] = {"COHP": cohp, "ICOHP": icohp,
                                            "orbitals": orbitals}
                # If no total COHPs are present, calculate the total
                # COHPs from the single-orbital populations. Total COHPs
                # may not be present when the cohpgenerator keyword is used
                # in LOBSTER versions 2.2.0 and earlier.
                if label not in d["COHP"] or d["COHP"][label] is None:
                    cohp = {Spin.up: np.sum(np.array(
                        [orb_cohp[label][orb]["COHP"][Spin.up]
                         for orb in orb_cohp[label]]), axis=0)}
                    try:
                        cohp[Spin.down] = np.sum(np.array([orb_cohp[label][orb]["COHP"][Spin.down]
                                                           for orb in orb_cohp[label]]), axis=0)
                    except KeyError:
                        pass

                orb_res_icohp = None in [orb_cohp[label][orb]["ICOHP"] for orb in orb_cohp[label]]
                if (label not in d["ICOHP"] or d["ICOHP"][label] is None) and orb_res_icohp:
                    icohp = {Spin.up: np.sum(np.array([orb_cohp[label][orb]["ICOHP"][Spin.up]
                                                       for orb in orb_cohp[label]]), axis=0)}
                    try:
                        icohp[Spin.down] = np.sum(np.array([orb_cohp[label][orb]["ICOHP"][Spin.down]
                                                            for orb in orb_cohp[label]]), axis=0)
                    except KeyError:
                        pass
        else:
            orb_cohp = None

        if "average" not in d["COHP"].keys():
            # calculate average
            cohp = np.array([np.array(c)
                             for c in d["COHP"].values()]).mean(axis=0)
            try:
                icohp = np.array([np.array(c)
                                  for c in d["ICOHP"].values()]).mean(axis=0)
            except KeyError:
                icohp = None
            avg_cohp = Cohp(efermi, energies, cohp, icohp=icohp)

        return CompleteCohp(structure, avg_cohp, cohp_dict, bonds=bonds,
                            are_coops=d["are_coops"], orb_res_cohp=orb_cohp)

    @classmethod
    def from_file(cls, fmt, filename=None,
                  structure_file=None, are_coops=False):
        """
        Creates a CompleteCohp object from an output file of a COHP
        calculation. Valid formats are either LMTO (for the Stuttgart
        LMTO-ASA code) or LOBSTER (for the LOBSTER code).

        Args:
            cohp_file: Name of the COHP output file. Defaults to COPL
                for LMTO and COHPCAR.lobster/COOPCAR.lobster for LOBSTER.

            are_coops: Indicates whether the populations are COOPs or
                COHPs. Defaults to False for COHPs.

            fmt: A string for the code that was used to calculate
                the COHPs so that the output file can be handled
                correctly. Can take the values "LMTO" or "LOBSTER".

            structure_file: Name of the file containing the structure.
                If no file name is given, use CTRL for LMTO and POSCAR
                for LOBSTER.

        Returns:
            A CompleteCohp object.
        """
        fmt = fmt.upper()
        if fmt == "LMTO":
            # LMTO COOPs and orbital-resolved COHP cannot be handled yet.
            are_coops = False
            orb_res_cohp = None
            if structure_file is None:
                structure_file = "CTRL"
            if filename is None:
                filename = "COPL"
            cohp_file = LMTOCopl(filename=filename, to_eV=True)
        elif fmt == "LOBSTER":
            if structure_file is None:
                structure_file = "POSCAR"
            if filename is None:
                filename = "COOPCAR.lobster" if are_coops \
                    else "COHPCAR.lobster"
            warnings.warn(
                "The bond labels are currently consistent with ICOHPLIST.lobster/ICOOPLIST.lobster, not with "
                "COHPCAR.lobster/COOPCAR.lobster. Please be aware!")
            cohp_file = Cohpcar(filename=filename, are_coops=are_coops)
            orb_res_cohp = cohp_file.orb_res_cohp
        else:
            raise ValueError("Unknown format %s. Valid formats are LMTO "
                             "and LOBSTER." % fmt)

        structure = Structure.from_file(structure_file)
        efermi = cohp_file.efermi
        cohp_data = cohp_file.cohp_data
        energies = cohp_file.energies

        # Lobster shifts the energies so that the Fermi energy is at zero.
        # Shifting should be done by the plotter object though.

        spins = [Spin.up, Spin.down] if cohp_file.is_spin_polarized \
            else [Spin.up]
        if fmt == "LOBSTER":
            energies += efermi

        if orb_res_cohp is not None:
            # If no total COHPs are present, calculate the total
            # COHPs from the single-orbital populations. Total COHPs
            # may not be present when the cohpgenerator keyword is used
            # in LOBSTER versions 2.2.0 and earlier.
            # TODO: Test this more extensively
            # pylint: disable=E1133,E1136
            for label in orb_res_cohp:
                if cohp_file.cohp_data[label]["COHP"] is None:
                    # print(label)
                    cohp_data[label]["COHP"] = {
                        sp: np.sum([orb_res_cohp[label][orbs]["COHP"][sp] for orbs in orb_res_cohp[label]], axis=0)
                        for sp in spins}
                if cohp_file.cohp_data[label]["ICOHP"] is None:
                    cohp_data[label]["ICOHP"] = \
                        {sp: np.sum([orb_res_cohp[label][orbs]["ICOHP"][sp]
                                     for orbs in orb_res_cohp[label]],
                                    axis=0) for sp in spins}

        if fmt == "LMTO":
            # Calculate the average COHP for the LMTO file to be
            # consistent with LOBSTER output.
            avg_data = {"COHP": {}, "ICOHP": {}}
            for i in avg_data:
                for spin in spins:
                    rows = np.array([cohp_data[label][i][spin]
                                     for label in cohp_data])
                    avg = np.average(rows, axis=0)
                    # LMTO COHPs have 5 significant figures
                    avg_data[i].update({spin: np.array([round_to_sigfigs(a, 5)
                                                        for a in avg], dtype=float)})
            avg_cohp = Cohp(efermi, energies,
                            avg_data["COHP"],
                            icohp=avg_data["ICOHP"])
        else:
            avg_cohp = Cohp(efermi, energies,
                            cohp_data["average"]["COHP"],
                            icohp=cohp_data["average"]["COHP"],
                            are_coops=are_coops)
            del cohp_data["average"]

        cohp_dict = {label: Cohp(efermi, energies,
                                 cohp_data[label]["COHP"],
                                 icohp=cohp_data[label]["ICOHP"],
                                 are_coops=are_coops)
                     for label in cohp_data}

        bond_dict = {label: {"length": cohp_data[label]["length"],
                             "sites": [structure.sites[site]
                                       for site in cohp_data[label]["sites"]]}
                     for label in cohp_data}

        return CompleteCohp(structure, avg_cohp, cohp_dict, bonds=bond_dict,
                            are_coops=are_coops, orb_res_cohp=orb_res_cohp)

class IcohpValue(MSONable):
    """
    Class to store information on an ICOHP or ICOOP value

    .. attribute:: num_bonds
            number of bonds used for the average cohp (relevant for Lobster versions <3.0) (int)

    .. attribute:: are_coops
            Boolean to indicate whether ICOOP or not

    .. attribute:: icohp
            dict={Spin.up: icohpvalue for spin.up, Spin.down: icohpvalue for spin.down}

    .. attribute:: summed_icohp:
            sum of icohp/icoop of both spin channels

    """

    def __init__(self, label, atom1, atom2, length, translation, num, icohp, are_coops=False):
        """
        Args:
            label: label for the icohp
            atom1: str of atom that is contributing to the bond
            atom2: str of second atom that is contributing to the bond
            length: float of bond lengths
            translation: translation list, e.g. [0,0,0]
            num: integer describing how often the bond exists
            icohp: dict={Spin.up: icohpvalue for spin.up, Spin.down: icohpvalue for spin.down}
        """
        self._are_coops = are_coops
        self._label = label
        self._atom1 = atom1
        self._atom2 = atom2
        self._length = length
        self._translation = translation
        self._num = num
        self._icohp = icohp
        if Spin.down in self._icohp:
            self._is_spin_polarized = True
        else:
            self._is_spin_polarized = False

    def __str__(self):

        if not self._are_coops:
            if self._is_spin_polarized:
                return ("ICOHP " + str(self._label) + " between " + str(self._atom1) + " and " + str(self._atom2) +
                        " (" + str(self._translation) + "): " + str(self._icohp[Spin.up]) + " eV (Spin up) and " +
                        str(self._icohp[Spin.down]) + " eV (Spin down)")
            return ("ICOHP " + str(self._label) + " between " + str(self._atom1) + " and " + str(self._atom2) +
                    " (" + str(self._translation) + "): " + str(self._icohp[Spin.up]) + " eV (Spin up)")
        if self._is_spin_polarized:
            return ("ICOOP " + str(self._label) + " between " + str(self._atom1) + " and " + str(self._atom2) +
                    " (" + str(self._translation) + "): " + str(self._icohp[Spin.up]) + " (Spin up) and " +
                    str(self._icohp[Spin.down]) + " (Spin down)")
        return ("ICOOP " + str(self._label) + " between " + str(self._atom1) + " and " + str(self._atom2) +
                " (" + str(self._translation) + "): " + str(self._icohp[Spin.up]) + " (Spin up)")

    @property
    def num_bonds(self):
        """
        tells the number of bonds for which the ICOHP value is an average
        Returns:
            Int
        """
        return self._num

    @property
    def are_coops(self):
        """
        tells if ICOOPs or not
        Returns:
            Boolean
        """
        return self._are_coops

    @property
    def is_spin_polarized(self):
        """
        tells if spin polarized calculation or not
        Returns:
            Boolean

        """
        return self._is_spin_polarized

    def icohpvalue(self, spin=Spin.up):
        """
        Args:
            spin: Spin.up or Spin.down
        Returns:
            icohpvalue (float) corresponding to chosen spin
        """
        if not self.is_spin_polarized and spin == Spin.down:
            raise ValueError("The calculation was not performed with spin polarization")

        return self._icohp[spin]

    @property
    def icohp(self):
        """
        dict with icohps for spinup and spindown
        Return:
            dict={Spin.up: icohpvalue for spin.up, Spin.down: icohpvalue for spin.down}
        """
        return self._icohp

    @property
    def summed_icohp(self):
        """
        Adds ICOHPs of both spin channels for spin polarized compounds
        Returns:
             icohp value in eV
        """
        if self._is_spin_polarized:
            sum_icohp = self._icohp[Spin.down] + self._icohp[Spin.up]
        else:
            sum_icohp = self._icohp[Spin.up]
        return sum_icohp

class IcohpCollection(MSONable):
    """
    Class to store IcohpValues

    .. attribute:: are_coops
        Boolean to indicate whether ICOHPs or ICOOPs are stored

    .. attribute:: is_spin_polarized
        Boolean to indicate if the Lobster calculation was done spin polarized or not

    """

    def __init__(self, list_labels, list_atom1, list_atom2, list_length,
                 list_translation, list_num, list_icohp, is_spin_polarized, are_coops=False):
        """
        Args:
            is_spin_polarized: Boolean to indicate if the Lobster calculation was done spin polarized or not Boolean to
                indicate if the Lobster calculation was done spin polarized or not
            are_coops: Boolean to indicate whether ICOHPs or ICOOPs are stored
            list_labels: list of labels for ICOHP/ICOOP values
            list_atom1: list of str of atomnames e.g. "O1"
            list_atom2: list of str of atomnames e.g. "O1"
            list_length: list of lengths of corresponding bonds in Angstrom
            list_translation: list of translation list, e.g. [0,0,0]
            list_num: list of equivalent bonds, usually 1 starting from Lobster 3.0.0
            list_icohp: list of dict={Spin.up: icohpvalue for spin.up, Spin.down: icohpvalue for spin.down}
        """
        self._are_coops = are_coops
        self._icohplist = {}
        self._is_spin_polarized = is_spin_polarized
        self._list_labels = list_labels
        self._list_atom1 = list_atom1
        self._list_atom2 = list_atom2
        self._list_length = list_length
        self._list_translation = list_translation
        self._list_num = list_num
        self._list_icohp = list_icohp

        for ilist, listel in enumerate(list_labels):
            self._icohplist[listel] = IcohpValue(listel, list_atom1[ilist], list_atom2[ilist], list_length[ilist],
                                                 list_translation[ilist], list_num[ilist], list_icohp[ilist])

    def __str__(self):
        joinstr = []
        for value in self._icohplist.values():
            joinstr.append(str(value))
        return "\n".join(joinstr)

    def get_icohp_by_label(self, label, summed_spin_channels=True, spin=Spin.up):
        """
        get an icohp value for a certain bond as indicated by the label (bond labels starting by "1" as in
        ICOHPLIST/ICOOPLIST)

        Args:
            label: label in str format (usually the bond number in Icohplist.lobster/Icooplist.lobster
            summed_spin_channels: Boolean to indicate whether the ICOHPs/ICOOPs of both spin channels should be summed
            spin: if summed_spin_channels is equal to False, this spin indicates which spin channel should be returned

        Returns:
            float describing ICOHP/ICOOP value
        """

        icohp_here = self._icohplist[label]
        if icohp_here._is_spin_polarized:
            if summed_spin_channels:
                return icohp_here.summed_icohp
            return icohp_here.icohpvalue(spin)
        return icohp_here.icohpvalue(spin)

    def get_summed_icohp_by_label_list(self, label_list, divisor=1.0, summed_spin_channels=True, spin=Spin.up):
        """
        get the sum of several ICOHP values that are indicated by a list of labels (labels of the bonds are the same as
        in ICOHPLIST/ICOOPLIST)

        Args:
            label_list: list of labels of the ICOHPs/ICOOPs that should be summed
            divisor: is used to divide the sum
            summed_spin_channels: Boolean to indicate whether the ICOHPs/ICOOPs of both spin channels should be summed
            spin: if summed_spin_channels is equal to False, this spin indicates which spin channel should be returned

        Returns:
             float that is a sum of all ICOHPs/ICOOPs as indicated with label_list
        """
        sum_icohp = 0
        for label in label_list:
            icohp_here = self._icohplist[label]
            if icohp_here.num_bonds != 1:
                warnings.warn("One of the ICOHP values is an average over bonds. This is currently not considered.")
            # prints warning if num_bonds is not equal to 1
            if icohp_here._is_spin_polarized:
                if summed_spin_channels:
                    sum_icohp = sum_icohp + icohp_here.summed_icohp
                else:
                    sum_icohp = sum_icohp + icohp_here.icohpvalue(spin)
            else:
                sum_icohp = sum_icohp + icohp_here.icohpvalue(spin)
        return sum_icohp / divisor

    def get_icohp_dict_by_bondlengths(self, minbondlength=0.0, maxbondlength=8.0):
        """
        get a dict of IcohpValues corresponding to certaind bond lengths
        Args:
            minbondlength: defines the minimum of the bond lengths of the bonds
            maxbondlength: defines the maximum of the bond lengths of the bonds
        Returns:
             dict of IcohpValues, the keys correspond to the values from the initial list_labels
        """
        newicohp_dict = {}
        for value in self._icohplist.values():
            if value._length >= minbondlength and value._length <= maxbondlength:
                newicohp_dict[value._label] = value
        return newicohp_dict

    def get_icohp_dict_of_site(self, site, minsummedicohp=None, maxsummedicohp=None, minbondlength=0.0,
                               maxbondlength=8.0, only_bonds_to=None):
        """
        get a dict of IcohpValue for a certain site (indicated by integer)
        Args:
            site: integer describing the site of interest, order as in Icohplist.lobster/Icooplist.lobster, starts at 0
            minsummedicohp: float, minimal icohp/icoop of the bonds that are considered. It is the summed ICOHP value
                from both spin channels for spin polarized cases
            maxsummedicohp: float, maximal icohp/icoop of the bonds that are considered. It is the summed ICOHP value
                from both spin channels for spin polarized cases
            minbondlength: float, defines the minimum of the bond lengths of the bonds
            maxbondlength: float, defines the maximum of the bond lengths of the bonds
            only_bonds_to: list of strings describing the bonding partners that are allowed, e.g. ['O']
        Returns:
             dict of IcohpValues, the keys correspond to the values from the initial list_labels
        """

        newicohp_dict = {}
        for key, value in self._icohplist.items():
            atomnumber1 = int(re.split(r'(\d+)', value._atom1)[1]) - 1
            atomnumber2 = int(re.split(r'(\d+)', value._atom2)[1]) - 1
            if site in (atomnumber1, atomnumber2):
                # manipulate order of atoms so that searched one is always atom1
                if site == atomnumber2:
                    save = value._atom1
                    value._atom1 = value._atom2
                    value._atom2 = save

                if only_bonds_to is None:
                    second_test = True
                else:
                    second_test = (re.split(r'(\d+)', value._atom2)[0] in only_bonds_to)
                if value._length >= minbondlength and value._length <= maxbondlength and second_test:
                    if minsummedicohp is not None:
                        if value.summed_icohp >= minsummedicohp:
                            if maxsummedicohp is not None:
                                if value.summed_icohp <= maxsummedicohp:
                                    newicohp_dict[key] = value
                            else:
                                newicohp_dict[key] = value
                    else:
                        if maxsummedicohp is not None:
                            if value.summed_icohp <= maxsummedicohp:
                                newicohp_dict[key] = value
                        else:
                            newicohp_dict[key] = value

        return newicohp_dict

    def extremum_icohpvalue(self, summed_spin_channels=True, spin=Spin.up):
        """
        get ICOHP/ICOOP of strongest bond
        Args:
            summed_spin_channels: Boolean to indicate whether the ICOHPs/ICOOPs of both spin channels should be summed

            spin: if summed_spin_channels is equal to False, this spin indicates which spin channel should be returned
        Returns:
            lowest ICOHP/largest ICOOP value (i.e. ICOHP/ICOOP value of strongest bond)
        """
        if not self._are_coops:
            extremum = sys.float_info.max
        else:
            extremum = -sys.float_info.max

        if not self._is_spin_polarized:
            if spin == Spin.down:
                warnings.warn("This spin channel does not exist. I am switching to Spin.up")
            spin = Spin.up

        for value in self._icohplist.values():
            if not value.is_spin_polarized or not summed_spin_channels:
                if not self._are_coops:
                    if value.icohpvalue(spin) < extremum:
                        extremum = value.icohpvalue(spin)
                        # print(extremum)
                else:
                    if value.icohpvalue(spin) > extremum:
                        extremum = value.icohpvalue(spin)
                        # print(extremum)
            else:
                if not self._are_coops:
                    if value.summed_icohp < extremum:
                        extremum = value.summed_icohp
                        # print(extremum)
                else:
                    if value.summed_icohp > extremum:
                        extremum = value.summed_icohp
                        # print(extremum)
        return extremum

    @property
    def is_spin_polarized(self):
        """
        :return: Whether it is spin polarized.
        """
        return self._is_spin_polarized

    @property
    def are_coops(self):
        """
        :return: Whether this is coops.
        """
        return self._are_coops


#!/usr/bin/env python

import os
import time
import json
import tensorflow as tf
import numpy as np
from utils import *
from data_loader import *
from attention import attention
from ffindex import FFindexDB, read_index, read_data

eps  = 1e-9 # small number

N_AA = 20 # regular aa
N_AA_MSA = 21 # regular aa + gap
WMIN = 0.8

N_PRINT_LEVEL = 50

TRAIN_LOG = "Train [%03d/%03d] counter: %5d  time: %10.1f lr: %.6f | loss: %7.4f | %.4f %.4f %.4f %.4f %.4f %.4f %.4f | %.4f %.4f %.4f %.4f %.4f %.4f %.4f"
VALID_LOG = "Valid [%03d/%03d] counter: %5d  time: %10.1f lr: %.6f | loss: %7.4f | %.4f %.4f %.4f %.4f %.4f %.4f %.4f | %.4f %.4f %.4f %.4f %.4f %.4f %.4f"

# ResNet model definition
class ResNet_model(object):
    def __init__(self, sess, n_1d_layer=2, dilation=[1],
                 p_dropout=0.2, l2_coeff=0.001, kernel_size=3, n_hidden_rnn=64, attention_size=50,
                 n_feat_1d=64, n_feat_2d = 64, n_bottle_1d=32, n_bottle_2d=32, use_cpu=False, use_templ=False):
        self.sess = sess # tensorflow session
        #
        self.n_1d_layer = n_1d_layer
        self.dilation = dilation
        self.kernel = kernel_size
        #
        self.n_seq_1d = N_AA + N_AA_MSA + 1
        self.n_str_feat = 8 # 1 dist_map + 6 ori_map + 1 seqsep
        self.SS_dim     = 9
        self.phi_dim    = 36
        self.psi_dim    = 36
        self.omg_dim    = 2
        #
        self.n_hidden_rnn = n_hidden_rnn
        self.attention_size = attention_size
        #
        # hidden layer (1d)
        self.n_feat_1d = n_feat_1d
        self.n_feat_2d = n_feat_2d
        self.n_bottle_1d = n_bottle_1d
        self.n_bottle_2d = n_bottle_2d
        #
        self.p_dropout = p_dropout
        self.rnn_p_dropout = p_dropout
        self.l2_coeff = l2_coeff
        #
        with tf.variable_scope("SStorPred"):
            self.build_model(use_cpu=use_cpu, use_templ=use_templ)

    def build_model(self, use_cpu=False, use_templ=False):
        #
        # Receive inputs
        with tf.variable_scope("input"):
            self.seq      = tf.placeholder(tf.float32, [None, N_AA], name="seq") # n_res, N_AA (blosum)
            self.msa      = tf.placeholder(tf.uint8, [None, None], name="msa") # n_seq, n_res
            self.str_2d   = tf.placeholder(tf.float32, [None, None, None, self.n_str_feat], name="str_2d") # 1, n_res, n_res, n_feat
            self.pth_1d   = tf.placeholder(tf.float32, [None, None, 4])
            self.pth_2d   = tf.placeholder(tf.float32, [None, None, None, 6])
            self.SS       = tf.placeholder(tf.float32, [None, None, self.SS_dim]) # SS label
            self.phi      = tf.placeholder(tf.float32, [None, None, self.phi_dim]) # ref. phi distrib
            self.psi      = tf.placeholder(tf.float32, [None, None, self.psi_dim]) # ref. psi distrib
            self.omg      = tf.placeholder(tf.float32, [None, None, self.omg_dim]) # ref. omg distrib
            self.is_train = tf.placeholder(tf.bool)
            self.n_batch = tf.shape(self.str_2d)[0]
            self.n_res = tf.shape(self.seq)[0]
        #
        # 3-state answer
        self.SS3 = tf.stack([tf.reduce_sum(self.SS[:,:, :3], axis=-1),
                             tf.reduce_sum(self.SS[:,:,3:6], axis=-1),
                             tf.reduce_sum(self.SS[:,:,6: ], axis=-1)], axis=-1)
        #
        # 8-state answer
        self.SS8 = tf.stack([self.SS[:,:,0],
                             tf.reduce_sum(self.SS[:,:,1:3], axis=-1),
                             self.SS[:,:,3],
                             self.SS[:,:,4],
                             self.SS[:,:,5],
                             self.SS[:,:,6],
                             self.SS[:,:,7],
                             self.SS[:,:,8]], axis=-1)
        #================================
        # sequence features
        #================================
        # get pssm features from MSA
        msa1hot = tf.one_hot(self.msa, N_AA_MSA, dtype=tf.float32)
        w_seq = reweight_seq(msa1hot, WMIN)
        pssm = msa2pssm(msa1hot, w_seq)
        #
        seq_1d = tf.concat([self.seq, pssm], axis=-1) # sequence based features
        seq_1d = tf.expand_dims(seq_1d, 0)
        #
        # projection to n_feat_1d
        feat = tf.layers.conv1d(seq_1d, self.n_feat_1d, 1, padding='same')
        #
        #=================================
        # 1D ResNet with combined features
        #=================================
        # Stacking 1-dim residual blocks
        for i in range(self.n_1d_layer):
            d = self.dilation[i%len(self.dilation)]
            feat = self.ResNet_block_1d(feat, self.is_train, step=i, dilation=d)
        feat = tf.nn.elu(inst_norm(feat))
        #
        seq_1d = tf.tile(tf.reshape(feat, [-1]), [self.n_batch*self.n_res])
        seq_1d = tf.reshape(seq_1d, [self.n_batch, self.n_res, self.n_res, self.n_feat_1d])

        if use_templ:
            #=================================
            # Process template info
            #=================================
            t1d = tf.concat([tf.sin(self.pth_1d[:,:,0])[:,:,None],
                            tf.cos(self.pth_1d[:,:,0])[:,:,None],
                            tf.sin(self.pth_1d[:,:,1])[:,:,None],
                            tf.cos(self.pth_1d[:,:,1])[:,:,None],
                            self.pth_1d[:,:,2:]], axis=-1)
            t2d = tf.concat([tf.one_hot(tf.cast(self.pth_2d[:,:,:,0],dtype=tf.uint8),19,dtype=tf.float32),
                            tf.sin(self.pth_2d[:,:,:,1])[:,:,:,None],
                            tf.cos(self.pth_2d[:,:,:,1])[:,:,:,None],
                            tf.sin(self.pth_2d[:,:,:,2])[:,:,:,None],
                            tf.cos(self.pth_2d[:,:,:,2])[:,:,:,None],
                            tf.sin(self.pth_2d[:,:,:,3])[:,:,:,None],
                            tf.cos(self.pth_2d[:,:,:,3])[:,:,:,None],
                            tf.tile(t1d[:,:,None,:],[1,1,self.n_res,1]),
                            tf.tile(t1d[:,None,:,:],[1,self.n_res,1,1])], axis=-1)
            sgnl = tf.nn.elu(tf.layers.conv2d(t2d, self.n_feat_2d, 3, padding='SAME'))
            prob = tf.nn.softmax(tf.layers.conv2d(t2d, 1, 3, padding='SAME'), axis=0)
            t2d = tf.reduce_sum(sgnl*prob,axis=0)
        #
        # combine with 2D str features
        str_2d = tf.layers.conv2d(self.str_2d, self.n_feat_2d, 1, padding='same')
        str_2d = tf.nn.elu(inst_norm(str_2d))

        if use_templ:
            feat = tf.concat((str_2d, t2d[None, :,:,:], seq_1d, tf.transpose(seq_1d, (0,2,1,3))), axis=-1)
        else:
            feat = tf.concat((str_2d, seq_1d, tf.transpose(seq_1d, (0,2,1,3))), axis=-1)
        #
        # projection to n_feat_2d
        feat = tf.layers.conv2d(feat, self.n_feat_2d, 1, padding='same', use_bias=False)
        #
        #=================================
        # 2D ResNet with combined features
        #=================================
        # Stacking 2-dim residual blocks (receptive field size: 61)
        for i in range(8):
            d = self.dilation[i%len(self.dilation)]
            feat = self.ResNet_block_2d(feat, self.is_train, step=i, dilation=d)
        feat = tf.nn.elu(inst_norm(feat))
        #
        #=================================
        # LSTM to extract 1-dimensional features from input features
        #=================================
        with tf.variable_scope("AttBiLSTM") as scope:
            # convert to 1D using LSTM
            weights = {'out': tf.Variable(tf.random_normal([self.n_hidden_rnn*2, self.n_feat_1d]), name='kernel')}
            biases  = {'out': tf.Variable(tf.random_normal([self.n_feat_1d]), name='bias')}
            #
            if use_cpu:
                feat, self.alphas = self.BiLSTM_w_attention_cpu(feat, weights, biases)
            else:
                feat, self.alphas = self.BiLSTM_w_attention_gpu(feat, weights, biases)
            feat = tf.reshape(feat, [self.n_batch, self.n_res, self.n_feat_1d])
        #
        # Stacking 1-dim residual blocks
        with tf.variable_scope("additional_ResNet"):
            for i in range(4):
                d = self.dilation[i%len(self.dilation)]
                feat = self.ResNet_block_1d(feat, self.is_train, step=i, dilation=d)
        #
        # Final branching
        with tf.variable_scope("final_SS") as scope:
            for i in range(4):
                d = self.dilation[i%len(self.dilation)]
                feat = self.ResNet_block_1d(feat, self.is_train, step=i, dilation=d)
            feat = tf.nn.elu(feat)
            SS_logit = tf.layers.conv1d(feat, self.SS_dim, 1, padding='same')
        #
        with tf.variable_scope("final_phi") as scope:
            for i in range(4):
                d = self.dilation[i%len(self.dilation)]
                feat = self.ResNet_block_1d(feat, self.is_train, step=i, dilation=d)
            feat = tf.nn.elu(feat)
            phi_logit = tf.layers.conv1d(feat, self.phi_dim, 1, padding='same')
        #
        with tf.variable_scope("final_psi") as scope:
            for i in range(4):
                d = self.dilation[i%len(self.dilation)]
                feat = self.ResNet_block_1d(feat, self.is_train, step=i, dilation=d)
            feat = tf.nn.elu(feat)
            psi_logit = tf.layers.conv1d(feat, self.psi_dim, 1, padding='same')
        #
        with tf.variable_scope("final_omg") as scope:
            for i in range(4):
                d = self.dilation[i%len(self.dilation)]
                feat = self.ResNet_block_1d(feat, self.is_train, step=i, dilation=d)
            feat = tf.nn.elu(feat)
            omg_logit = tf.layers.conv1d(feat, self.omg_dim, 1, padding='same')
        #
        # calculate probability
        SS_prob = tf.nn.softmax(SS_logit)
        SS3_prob = tf.stack([tf.reduce_sum(SS_prob[:,:, :3], axis=-1),
                             tf.reduce_sum(SS_prob[:,:,3:6], axis=-1),
                             tf.reduce_sum(SS_prob[:,:,6: ], axis=-1)], axis=-1)
        SS8_prob = tf.stack([SS_prob[:,:,0],
                             tf.reduce_sum(SS_prob[:,:,1:3], axis=-1),
                             SS_prob[:,:,3],
                             SS_prob[:,:,4],
                             SS_prob[:,:,5],
                             SS_prob[:,:,6],
                             SS_prob[:,:,7],
                             SS_prob[:,:,8]], axis=-1)
        phi_prob = tf.nn.softmax(phi_logit)
        psi_prob = tf.nn.softmax(psi_logit)
        omg_prob = tf.nn.softmax(omg_logit)
        self.prob_s = [SS_prob, SS8_prob, SS3_prob, phi_prob, psi_prob, omg_prob]

        # calculate loss function (softmax cross-entropy)
        # For SS & omega, it is same as categorical cross entropy
        # For phi psi angles, reference is defined with von Mises distrib.
        # It should be noted that minimizing softmax cross-entropy is same as minimizing KL divergence

        SS_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(self.SS, SS_logit))

        SS3_loss = -tf.reduce_sum(self.SS3 * tf.log(SS3_prob+eps), axis=-1)
        SS3_loss = tf.reduce_mean(SS3_loss)

        SS8_loss = -tf.reduce_sum(self.SS8 * tf.log(SS8_prob+eps), axis=-1)
        SS8_loss = tf.reduce_mean(SS8_loss)

        phi_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(self.phi, phi_logit))
        psi_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(self.psi, psi_logit))
        omg_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(self.omg, omg_logit))

        self.t_vars = tf.trainable_variables()
        #
        # L2-regularization to avoid overfitting
        l2_loss = tf.add_n([tf.nn.l2_loss(var)
                            for var in self.t_vars if 'kernel' in var.name]) * self.l2_coeff
        #
        # total losses
        self.loss_s = [SS_loss, SS8_loss, SS3_loss, phi_loss, psi_loss, omg_loss, l2_loss]
        self.tot_loss = tf.add_n(self.loss_s)

        # calculate accuracy
        equal = tf.cast(tf.equal(tf.argmax(SS_prob, axis=-1), tf.argmax(self.SS, axis=-1)), tf.float32)
        SS_acc = tf.reduce_mean(tf.cast(equal, tf.float32))
        #
        equal = tf.cast(tf.equal(tf.argmax(SS3_prob, axis=-1), tf.argmax(self.SS3, axis=-1)), tf.float32)
        SS3_acc = tf.reduce_mean(tf.cast(equal, tf.float32))
        #
        equal = tf.cast(tf.equal(tf.argmax(SS8_prob, axis=-1), tf.argmax(self.SS8, axis=-1)), tf.float32)
        SS8_acc = tf.reduce_mean(tf.cast(equal, tf.float32))
        #
        phi_equal = tf.equal(tf.argmax(phi_prob, axis=-1), tf.argmax(self.phi, axis=-1))
        phi_acc = tf.reduce_mean(tf.cast(phi_equal, tf.float32))
        #
        psi_equal = tf.equal(tf.argmax(psi_prob, axis=-1), tf.argmax(self.psi, axis=-1))
        psi_acc = tf.reduce_mean(tf.cast(psi_equal, tf.float32))
        #
        omg_equal = tf.equal(tf.argmax(omg_prob, axis=-1), tf.argmax(self.omg, axis=-1))
        omg_acc = tf.reduce_mean(tf.cast(omg_equal, tf.float32))
        #
        equal = tf.stack([phi_equal, psi_equal, omg_equal], axis=-1)
        equal = tf.reduce_all(equal, axis=-1)
        tot_acc = tf.reduce_mean(tf.cast(equal, tf.float32))
        self.acc_s = [SS_acc, SS8_acc, SS3_acc, phi_acc, psi_acc, omg_acc, tot_acc]
        #
        # define trained model saver
        self.saver = tf.train.Saver()

    def BiLSTM_w_attention_gpu(self, x, weights, biases, return_alphas=True):
        # prepare input data shape to match rnn function requirements
        # input x: (batch_size, n_res, n_res, n_ch)
        # Required shape: timesteps tensors list of shape (batch_size', n_input')
        #   - howto? batch_size' = batch_size*n_res , n_input' = n_ch
        #            # of timesteps = n_res

        # reshape input x to (batch_size*n_res, n_res, n_ch), CAUTION: batch_size=1 here
        # first n_res: time-series & time-series should go first in cudnn version
        x = tf.transpose(x, (1,0,2,3))
        x = tf.reshape(x, [self.n_res, self.n_batch*self.n_res, self.n_feat_2d])
        #x = tf.reshape(x, [self.n_res, self.n_batch*self.n_res, self.n_hidden_rnn])

        # define a lstm cell
        lstm = tf.contrib.cudnn_rnn.CudnnLSTM(3, self.n_hidden_rnn, direction='bidirectional', dtype=tf.float32)

        # Get BiLSTM cell output
        outputs, states = lstm(x)
        outputs = tf.concat(outputs, axis=-1)

        # apply attention
        outputs, alphas = attention(outputs, self.attention_size, return_alphas=return_alphas, time_major=True)

        # apply linear activation
        return tf.matmul(outputs, weights['out']) + biases['out'], alphas

    def BiLSTM_w_attention_cpu(self, x, weights, biases, return_alphas=True):
        # prepare input data shape to match rnn function requirements
        # input x: (batch_size, n_res, n_res, n_ch)
        # Required shape: timesteps tensors list of shape (batch_size', n_input')
        #   - howto? batch_size' = batch_size*n_res , n_input' = n_ch
        #            # of timesteps = n_res

        # reshape input x to (batch_size*n_res, n_res, n_ch), CAUTION: batch_size=1 here
        # first n_res: time-series & time-series should go first in cudnn version
        x = tf.transpose(x, (1,0,2,3))
        x = tf.reshape(x, [self.n_res, self.n_batch*self.n_res, self.n_feat_2d])
        #x = tf.reshape(x, [self.n_res, self.n_batch*self.n_res, self.n_hidden_rnn])

        # define a lstm cell
        with tf.variable_scope("cudnn_lstm"):
            single_cell = lambda: tf.contrib.cudnn_rnn.CudnnCompatibleLSTMCell(self.n_hidden_rnn)
            cells_fw = [single_cell() for _ in range(3)]
            cells_bw = [single_cell() for _ in range(3)]
            #
            # Get BiLSTM cell output
            outputs, output_state_fw, output_state_bw = tf.contrib.rnn.stack_bidirectional_dynamic_rnn(cells_fw, cells_bw, x, time_major=True, dtype=tf.float32)
        outputs = tf.concat(outputs, axis=-1)

        # apply attention
        outputs, alphas = attention(outputs, self.attention_size, return_alphas=return_alphas, time_major=True)

        # apply linear activation
        return tf.matmul(outputs, weights['out']) + biases['out'], alphas

    def ResNet_block_2d(self, x, is_train, step=0, dilation=1): # bottleneck block w/ pre-activation
        with tf.variable_scope("ResNet_2d_{}".format(step)) as scope:
            shortcut = x
            # bottleneck layer (kernel: 1, n_feat_2d => n_bottle_2d)
            x = tf.nn.elu(inst_norm(x))
            x = tf.layers.conv2d(x, self.n_bottle_2d, 1, padding='same')
            x = tf.nn.elu(inst_norm(x))
            # convolution
            x = tf.layers.conv2d(x, self.n_bottle_2d, self.kernel, dilation_rate=dilation,
                                 padding='same')
            x = tf.nn.elu(inst_norm(x))
            x = tf.layers.dropout(x, rate=self.p_dropout, training=is_train)
            # project up (kernel: 1, n_bottle_1d => n_feat_1d)
            x = tf.layers.conv2d(x, self.n_feat_2d, 1, padding='same')
            # add
            x += shortcut
        return x

    def ResNet_block_1d(self, x, is_train, step=0, dilation=1): # bottleneck block w/ pre-activation
        with tf.variable_scope("ResNet_1d_{}".format(step)) as scope:
        #with tf.variable_scope("ResNet_1d_{}".format(step), custom_getter=float32_variable_storage_getter) as scope:
            shortcut = x
            # bottleneck layer (kernel: 1, n_feat_1d => n_bottle_1d)
            x = tf.nn.elu(inst_norm(x))
            x = tf.layers.conv1d(x, self.n_bottle_1d, 1, padding='same')
            x = tf.nn.elu(inst_norm(x))
            # convolution
            x = tf.layers.conv1d(x, self.n_bottle_1d, self.kernel, dilation_rate=dilation,
                                 padding='same')
            x = tf.nn.elu(inst_norm(x))
            x = tf.layers.dropout(x, rate=self.p_dropout, training=is_train)
            # project up (kernel: 1, n_bottle_1d => n_feat_1d)
            x = tf.layers.conv1d(x, self.n_feat_1d, 1, padding='same')
            # add
            x += shortcut
        return x

    def save(self, folder, prefix):
        if not os.path.exists(folder):
            os.mkdir(folder)
        self.saver.save(self.sess, folder+"/%s.ckpt"%prefix)

    def load(self, folder, prefix):
        model_fn = os.path.join(folder, "%s.ckpt.index"%prefix)
        if os.path.exists(model_fn):
            self.saver.restore(self.sess, folder+"/%s.ckpt"%prefix)
            return True
        return False

    def close(self):
        self.sess.close()

    def train(self, config_file):
        # read config_file
        with open(config_file) as json_file:
            config = json.load(json_file, object_hook=Json_param)
        #
        train_pdbs = [line.split()[0] for line in open(config.train_list)]
        self.n_train = len(train_pdbs)
        #
        global_step = tf.Variable(0, trainable=False)

        # define optimizer
        if config.lr_schedule == "CosineDecay":
            lr = tf.train.cosine_decay_restarts(learning_rate=config.lr,
                                                global_step=global_step,
                                                first_decay_steps=config.f_decay*self.n_train,
                                                t_mul=config.t_mul,
                                                m_mul=config.m_mul,
                                                alpha=0.1)
        elif config.lr_schedule == 'ExpDecay':
            lr = tf.train.exponential_decay(config.lr,
                                            global_step=global_step,
                                            decay_steps=self.n_train,
                                            decay_rate=0.99,
                                            staircase=True)
        else: # flat lr
            lr = tf.Variable(config.lr, trainable=False)

        if config.optim == 'adam':
            optimizer = tf.train.AdamOptimizer(lr)
        elif config.optim == 'momentum':
            optimizer = tf.train.MomentumOptimizer(lr, 0.5, use_nesterov=True)
        else:
            optimizer = tf.train.GradientDescentOptimizer(lr)

        optim = optimizer.minimize(self.tot_loss, global_step=global_step, var_list=self.t_vars)
        ops_to_run = [optim, lr, self.tot_loss, self.loss_s, self.acc_s]
        #
        # initialize all variables
        init_op = tf.group(tf.global_variables_initializer(),
                           tf.local_variables_initializer())
        self.sess.run(init_op)
        #
        counter = 0
        self.start_time = time.time()
        #
        n_batch = len(train_pdbs)
        min_val_loss = config.best_loss
        #
        # Try to load pre-trained model if exists
        could_load = self.load("model_%d_%d_%d"%(self.n_hidden_rnn, self.n_1d_layer, self.attention_size), 'last_epoch')
        if not could_load:
            could_load = self.load("model_%d_%d_%d"%(self.n_hidden_rnn, self.n_1d_layer, self.attention_size), 'model')
        #
        for epoch in range(config.n_epoch):
            np.random.shuffle(train_pdbs)
            tot_loss_value = 0.0
            tot_loss_s = np.zeros(7, dtype=np.float32)
            tot_acc_s = np.zeros(7, dtype=np.float32)
            n_tot = 0.0
            #
            for pdb in train_pdbs:
                seq, msa, str_2d, pth_1d, pth_2d, SS_labels, phi_labels, psi_labels, omg_labels = load_train_data(pdb, \
                                                is_train=True, mask_diag=config.mask_diag)
                #
                _, decayed_lr, loss_value, loss_s, acc_s = self.sess.run(ops_to_run,
                                         feed_dict={
                                             self.seq: seq,
                                             self.msa: msa,
                                             self.str_2d: str_2d[np.newaxis,:,:,:],
                                             self.pth_1d: pth_1d,
                                             self.pth_2d: pth_2d,
                                             self.SS: SS_labels[np.newaxis,:,:],
                                             self.phi: phi_labels[np.newaxis,:,:],
                                             self.psi: psi_labels[np.newaxis,:,:],
                                             self.omg: omg_labels[np.newaxis,:],
                                             self.is_train: True})
                tot_loss_value += loss_value
                tot_loss_s += np.array(loss_s)
                tot_acc_s += np.array(acc_s)
                n_tot += 1.0
                #
                counter += 1
                if counter % N_PRINT_LEVEL == 0:
                    loss_value = tot_loss_value/n_tot
                    loss_s = tot_loss_s/n_tot
                    acc_s = tot_acc_s/n_tot
                    tot_loss_value = 0.0
                    tot_loss_s = np.zeros(7, dtype=np.float32)
                    tot_acc_s = np.zeros(7, dtype=np.float32)
                    n_tot = 0.0
                    log_list = [epoch, config.n_epoch, counter, time.time()-self.start_time, decayed_lr*100.0, loss_value-loss_s[-1]]
                    log_list.extend(loss_s)
                    log_list.extend(acc_s)
                    print (TRAIN_LOG%tuple(log_list))
            #
            val_loss = self.validation(config, epoch, counter, decayed_lr)
            if val_loss < min_val_loss:
                self.save("model_%d_%d_%d"%(self.n_hidden_rnn, self.n_1d_layer, self.attention_size), 'model')
                min_val_loss = val_loss
            self.save("model_%d_%d_%d"%(self.n_hidden_rnn, self.n_1d_layer, self.attention_size), 'last_epoch')

    def validation(self, config, epoch, counter, decayed_lr):
        ops_to_run = [self.tot_loss, self.loss_s, self.acc_s]
        valid_pdbs = [line.split()[0] for line in open(config.valid_list) if line[0] != "#"]
        #
        tot_loss_value = 0.0
        tot_loss_s = np.zeros(7, dtype=np.float32)
        tot_acc_s = np.zeros(7, dtype=np.float32)
        n_tot = 0.0
        for pdb in valid_pdbs:
            seq, msa, str_2d, pth_1d, pth_2d, SS_labels, phi_labels, psi_labels, omg_labels\
                     = load_train_data(pdb, \
                                            is_train=False, mask_diag=config.mask_diag)
            #
            loss_value, loss_s, acc_s = self.sess.run(ops_to_run,
                                             feed_dict={
                                                 self.seq: seq,
                                                 self.msa: msa,
                                                 self.str_2d: str_2d[np.newaxis,:,:,:],
                                                 self.pth_1d: pth_1d,
                                                 self.pth_2d: pth_2d,
                                                 self.SS: SS_labels[np.newaxis,:,:],
                                                 self.phi: phi_labels[np.newaxis,:,:],
                                                 self.psi: psi_labels[np.newaxis,:,:],
                                                 self.omg: omg_labels[np.newaxis,:],
                                                 self.is_train:False})
            tot_loss_value += loss_value
            tot_loss_s += np.array(loss_s)
            tot_acc_s += np.array(acc_s)
            n_tot += 1.0
        #
        loss_value = tot_loss_value/n_tot
        loss_s = tot_loss_s/n_tot
        acc_s = tot_acc_s/n_tot
        #
        log_list = [epoch, config.n_epoch, counter, time.time()-self.start_time, decayed_lr*100.0, loss_value-loss_s[-1]]
        log_list.extend(loss_s)
        log_list.extend(acc_s)
        print (VALID_LOG%tuple(log_list))
        return loss_value-loss_s[-1]

    def predict(self, config):
        self.load(config.model_dir, 'model')
        #
        ffdb = None
        if config.templ_fn != None:
            ffdb = FFindexDB(read_index(config.TMPDB+'_pdb.ffindex'),
                             read_data(config.TMPDB+"_pdb.ffdata"))
            hhr_fn = config.templ_fn
        #
        out_fn = "%s.npz"%config.outprefix
        a3m_fn = config.a3m_fn
        pdb_fn = config.pdb_fn
        hhr_fn = config.templ_fn
        #
        seq, msa, str_2d, pth_1d, pth_2d = make_input_features(a3m_fn, pdb_fn, hhr_fn, ffdb)
        #
        if hhr_fn != None:
            prob_s, alpha = self.sess.run([self.prob_s, self.alphas], feed_dict={
                                                     self.seq: seq,
                                                     self.msa: msa,
                                                     self.str_2d: str_2d[np.newaxis,:,:,:],
                                                     self.pth_1d: pth_1d,
                                                     self.pth_2d: pth_2d,
                                                     self.is_train:False})
        else:
            prob_s, alpha = self.sess.run([self.prob_s, self.alphas], feed_dict={
                                                     self.seq: seq,
                                                     self.msa: msa,
                                                     self.str_2d: str_2d[np.newaxis,:,:,:],
                                                     self.is_train:False})

        SS_prob = prob_s[0].reshape(-1, self.SS_dim)
        phi_prob = prob_s[3].reshape(-1, self.phi_dim)
        psi_prob = prob_s[4].reshape(-1, self.psi_dim)
        omg_prob = prob_s[5].reshape(-1, self.omg_dim)
        tor_prob = np.concatenate((phi_prob, psi_prob, omg_prob), axis=-1)
        np.savez_compressed(out_fn, ss9=SS_prob.astype(np.float16), tor=tor_prob.astype(np.float16))

    def predict_multi(self, config, a3m_s, pdb_s, hhr_s, out_s):
        self.load(config.model_dir, 'model')
        #
        ffdb = None
        if hhr_s[0] != None:
            ffdb = FFindexDB(read_index(config.TMPDB+'_pdb.ffindex'),
                             read_data(config.TMPDB+"_pdb.ffdata"))
        #
        for i, a3m_fn in enumerate(a3m_s):
            out_fn = out_s[i]
            pdb_fn = pdb_s[i]
            hhr_fn = hhr_s[i]
            if os.path.exists(out_fn):
                continue
            print ("Running..., %s"%out_fn)
            #
            seq, msa, str_2d, pth_1d, pth_2d = make_input_features(a3m_fn, pdb_fn, hhr_fn, ffdb)
            #
            if hhr_fn != None:
                prob_s, alpha = self.sess.run([self.prob_s, self.alphas], feed_dict={
                                                         self.seq: seq,
                                                         self.msa: msa,
                                                         self.str_2d: str_2d[np.newaxis,:,:,:],
                                                         self.pth_1d: pth_1d,
                                                         self.pth_2d: pth_2d,
                                                         self.is_train:False})
            else:
                prob_s, alpha = self.sess.run([self.prob_s, self.alphas], feed_dict={
                                                         self.seq: seq,
                                                         self.msa: msa,
                                                         self.str_2d: str_2d[np.newaxis,:,:,:],
                                                         self.is_train:False})

            SS_prob = prob_s[0].reshape(-1, self.SS_dim)
            phi_prob = prob_s[3].reshape(-1, self.phi_dim)
            psi_prob = prob_s[4].reshape(-1, self.psi_dim)
            omg_prob = prob_s[5].reshape(-1, self.omg_dim)
            tor_prob = np.concatenate((phi_prob, psi_prob, omg_prob), axis=-1)
            np.savez_compressed(out_fn, ss9=SS_prob.astype(np.float16), tor=tor_prob.astype(np.float16))