#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Author: Tomas Cassanelli
import os
import yaml
import numpy as np
from astropy.io import fits
from astropy.time import Time
from astropy.table import Table, QTable
from astropy import units as apu
from astropy.constants import c as light_speed
from .math_functions import rms
from .aperture import e_rs, phase

__all__ = [
    'extract_data_pyoof', 'extract_data_effelsberg', 'str2LaTeX',
    'store_data_csv', 'uv_ratio', 'store_data_ascii', 'table_pyoof_out'
    ]

def extract_data_pyoof(pathfits):
    """
    Extracts data from the `~pyoof` default FITS file OOF holography
    observations, ready to use for the least squares minimization (see
    `~pyoof.fit_zpoly`). The FITS file has to have the following keys on its
    PrimaryHDU header: ``'FREQ'``, ``'WAVEL'``, ``'MEANEL'``, ``'OBJECT'`` and
    ``'DATE_OBS'``. Besides this three BinTableHDU are required for the data
    itself; ``MINUS OOF``, ``ZERO OOF`` and ``PLUS OOF``. The BinTableHDU
    header has to have the ``'DZ'`` key which includes the radial offset,
    :math:`d_z`. Finally the BinTableHDU has the data files ``'U'``, ``'V'``
    and ``'BEAM'``, which is the :math:`x`- and :math:`y`-axis position in
    radians and the ``'BEAM'`` in a flat array, in mJy.

    Parameters
    ----------
    pathfits : `str`
        Path to the FITS file that contains the three beam maps pre-calibrated,
        using the correct PrimaryHDU and the three BinTableHDU (``MINUS OOF``,
        ``ZERO OOF`` and ``PLUS OOF``).

    Returns
    -------
    data_info : `list`
        It contains all extra data besides the beam map. The output
        corresponds to a list,
        ``[name, pthto, obs_object, obs_date, freq, wavel, d_z, meanel]``.
        These are, name of the FITS file, paht of the FITS file, observed
        object, observation date, frequency, wavelength, radial offset and
        mean elevation, respectively.
    data_obs : `list`
        It contains beam maps and :math:`x`-, and :math:`y`-axis
        (:math:`uv`-plane in Fourier space) data for the least squares
        minimization (see `~pyoof.fit_zpoly`). The list has the following order
        ``[beam_data, u_data, v_data]``. ``beam_data`` is the three beam
        observations, minus, zero and plus out-of-focus, in a flat array.
        ``u_data`` and ``v_data`` are the beam axes in a flat array.

    Raises
    ------
    `ValueError`
        If the input file ``pathfits`` is not a FITS file, or if one or more
        of the header keys are missing in the FITS file:
        ``['FREQ', 'WAVEL', 'MEANEL', 'OBJECT', 'DATE_OBS']``.
    """

    if os.path.splitext(pathfits)[1] != '.fits':
        raise ValueError('File must be a FITS file.')

    hdulist = fits.open(pathfits)  # open FITS file, pyoof format
    # path or directory where the FITS file is located
    pthto = os.path.split(pathfits)[0]
    # name of the fit file to fit
    name = os.path.split(pathfits)[1][:-5]

    if not all(
            k in hdulist[0].header
            for k in ['FREQ', 'WAVEL', 'MEANEL', 'OBJECT', 'DATE_OBS']
            ):
        raise ValueError('Not all needed keys found in FITS header.')

    freq = hdulist[0].header['FREQ'] * apu.Hz
    wavel = hdulist[0].header['WAVEL'] * apu.m
    meanel = hdulist[0].header['MEANEL'] * apu.deg
    obs_object = hdulist[0].header['OBJECT']
    obs_date = hdulist[0].header['DATE_OBS']

    beam_data = np.array([hdulist[i].data['BEAM'] for i in range(1, 4)])
    u_data = np.array([hdulist[i].data['U'] for i in range(1, 4)]) * apu.rad
    v_data = np.array([hdulist[i].data['V'] for i in range(1, 4)]) * apu.rad
    d_z = np.array([hdulist[i].header['DZ'] for i in range(1, 4)]) * apu.m

    data_file = [name, pthto]
    data_info = data_file + [obs_object, obs_date, freq, wavel, d_z, meanel]
    data_obs = [beam_data, u_data, v_data]

    return data_info, data_obs

def extract_data_effelsberg(pathfits):
    """
    Extracts data from the Effelsberg OOF holography observations, ready to
    use for the least squares minimization. This function will only work for
    the Effelsberg telescope beam maps.

    Parameters
    ----------
    pathfits : `str`
        Path to the FITS file that contains the three beam maps pre-calibrated,
        from the Effelsberg telescope.

    Returns
    -------
    data_info : `list`
        It contains all extra data besides the beam map. The output
        corresponds to a list,
        ``[name, pthto, obs_object, obs_date, freq, wavel, d_z, meanel]``.
        These are, name of the FITS file, paht of the FITS file, observed
        object, observation date, frequency, wavelength, radial offset and
        mean elevation, respectively.
    data_obs : `list`
        It contains beam maps and :math:`x`-, and :math:`y`-axis
        (:math:`uv`-plane in Fourier space) data for the least squares
        minimization (see `~pyoof.fit_zpoly`). The list has the following order
        ``[beam_data, u_data, v_data]``. ``beam_data`` is the three beam
        observations, minus, zero and plus out-of-focus, in a flat array.
        ``u_data`` and ``v_data`` are the beam axes in a flat array.

    Raises
    ------
    `ValueError`
        If the input file ``pathfits`` is not a FITS file.
    """

    if os.path.splitext(pathfits)[1] != '.fits':
        raise ValueError('File must be a FITS file.')

    pos = [3, 1, 2]  # Positions for OOF holography observations at Effelsberg
    hdulist = fits.open(pathfits)  # main FITS file OOF holography format

    # Observation frequency
    freq = hdulist[0].header['FREQ'] * apu.Hz
    wavel = light_speed / freq

    # Mean elevation
    meanel = hdulist[0].header['MEANEL'] * apu.deg
    obs_object = hdulist[0].header['OBJECT']        # observed object
    obs_date = hdulist[0].header['DATE_OBS']        # observation date
    d_z = np.array([hdulist[i].header['DZ'] for i in pos]) * apu.m

    beam_data = np.array([hdulist[i].data['fnu'] for i in pos])
    u_data = np.array([hdulist[i].data['DX'] for i in pos]) * apu.rad
    v_data = np.array([hdulist[i].data['DY'] for i in pos]) * apu.rad

    # path or directory where the FITS file is located
    pthto = os.path.split(pathfits)[0]
    # name of the fit file to fit
    name = os.path.split(pathfits)[1][:-5]

    data_info = [name, pthto, obs_object, obs_date, freq, wavel, d_z, meanel]
    data_obs = [beam_data, u_data, v_data]

    return data_info, data_obs

def str2LaTeX(python_string):
    """
    Function that solves the underscore problem in a python string to
    :math:`\\LaTeX` string.

    Parameters
    ----------
    python_string : `str`
        String that needs to be changed.

    Returns
    -------
    LaTeX_string : `str`
        String with the new underscore symbol.
    """

    string_list = list(python_string)
    for idx, string in enumerate(string_list):
        if string_list[idx] == '_':
            string_list[idx] = '\\_'

    LaTeX_string = ''.join(string_list)

    return LaTeX_string

def store_data_csv(name, name_dir, order, save_to_csv):
    """
    Stores all important information in a CSV file after the least squares
    minimization has finished, `~pyoof.fit_zpoly`. All data will be stored in
    the ``pyoof_out/name-number`` directory, with ``name`` the name of the
    FITS file, and ``number`` the `~pyoof.fit_zpoly` number of code execution,
    i.e. the output data is never overwritten.

    Parameters
    ----------
    name : `str`
        File name of the FITS file to be optimized.
    name_dir : `str`
        Path to store all the CSV files. The files will depend on the order of
        the Zernike circle polynomial.
    order : `int`
        Order used for the Zernike circle polynomial, :math:`n`.
    save_to_csv : `list`
        It contains all data that will be stored. The list must have the
        following order, ``[beam_data, u_data, v_data, res_optim, jac_optim,
        grad_optim, phase, cov_ptrue, corr_ptrue]``.
    """

    headers = [
        'Normalized beam', 'u vector radians', 'v vector radians', 'Residual',
        'Jacobian', 'Gradient', 'Phase-error radians',
        'Variance-Covariance matrix (first row fitted parameters idx)',
        'Correlation matrix (first row fitted parameters idx)'
        ]

    fnames = [
        f'beam_data.csv', f'u_data.csv', f'v_data.csv',
        f'res_n{order}.csv', f'jac_n{order}.csv',
        f'grad_n{order}.csv', f'phase_n{order}.csv',
        f'cov_n{order}.csv', f'corr_n{order}.csv'
        ]

    if order != 1:
        headers = headers[3:]
        fnames = fnames[3:]
        save_to_csv = save_to_csv[3:]

    for fname, header, file in zip(fnames, headers, save_to_csv):
        np.savetxt(
            fname=os.path.join(name_dir, fname),
            X=file,
            header=' '.join((header, name))
            )

def store_data_ascii(name, name_dir, order, params_solution, params_init):
    """
    Stores in an ``~astropy.table.table.Table`` format the parameters found by
    the least squares minimization (see `~pyoof.fit_zpoly`).

    Parameters
    ----------
    name : `str`
        File name of the FITS file to be optimized.
    name_dir : `str`
        Path to store all the csv files. The files will depend on the order of
        the Zernike circle polynomial (radial) order.
    order : `int`
        Order used for the Zernike circle polynomial, :math:`n`.
    params_solution : `~numpy.ndarray`
        Contains the best fitted parameters, the illumination function
        coefficients, ``I_coeff`` and the Zernike circle polynomial
        coefficients, ``K_coeff`` in one array.
    params_init : `~numpt.ndarray`
        Contains the initial parameters used in the least squares minimization
        to start finding the best fitted combination of them.
    """

    n = order
    N_K_coeff = (n + 1) * (n + 2) // 2

    # Making nice table :)
    ln = [(j, i) for i in range(0, n + 1) for j in range(-i, i + 1, 2)]
    L = np.array(ln)[:, 0]
    N = np.array(ln)[:, 1]

    params_names = ['i_amp', 'c_dB', 'q', 'x_0', 'y_0']
    for i in range(N_K_coeff):
        params_names.append(f'K({N[i]}, {L[i]})')

    # To store fit information and found parameters in ascii file
    tab = Table(
        data=[params_names, params_solution, params_init],
        names=['parname', 'parfit', 'parinit'],
        )

    tab.write(
        os.path.join(name_dir, f'fitpar_n{n}.csv'),
        overwrite=True
        )

def uv_ratio(u, v):
    """
    Calculates the aspect ratio for the 3 power pattern plots, plus some
    corrections for the text on it. Used in the `function` `~pyoof.plot_beam`
    and `~pyoof.plot_data`

    Parameters
    ----------
    u : `~astropy.units.quantity.Quantity`
        Spatial frequencies from the power pattern in angle units.
    v : `~astropy.units.quantity.Quantity`
        Spatial frequencies from the power pattern in angle units.

    Returns
    -------
    width : `float`
        Width for the power pattern figure.
    height : `float`
        Height for the power pattern figure.
    """

    ratio = (u.max() - u.min()) / (v.max() - v.min())

    if type(ratio) == apu.Quantity:
        ratio = ratio.decompose().value

    height = 5
    width = (ratio * 2.25 * height)

    return width, height

def table_pyoof_out(path_pyoof_out, order):
    """
    Auxiliary function to tabulate all data from a series of observations
    gathered in a common ``pyoof_out/`` directory.

    Note: Piston and tilt are not used in error calculations, the phase
    calculations will only be included if these are manually changed in
    ``core.py``.

    Parameters
    ----------
    path_pyoof_out : `list`
        set of paths to the directory ``pyoof_out/`` or where the output from
        the `~pyoof` package is located.
    order : `int`
        Order used for the Zernike circle polynomial, :math:`n`.

    Returns
    -------
    qt : `~astropy.table.table.QTable`
        `~astropy.table.table.QTable` with units of the most important
        quantities from the `~pyoof` package.
    """

    qt = QTable(
        names=[
            'name', 'tel_name', 'obs-object', 'obs-date', 'meanel',
            'i_amp', 'c_dB', 'q', 'phase-rms', 'e_rs',
            'beam-snr-out-l', 'beam-snr-in', 'beam-snr-out-r'
            ],
        dtype=[np.string_] * 4 + [np.float] * 9
        )

    for p, pyoof_out in enumerate(path_pyoof_out):
        with open(os.path.join(pyoof_out, 'pyoof_info.yml'), 'r') as inputfile:
            pyoof_info = yaml.load(inputfile, Loader=yaml.Loader)

        _phase = np.genfromtxt(
            os.path.join(pyoof_out, f'phase_n{order}.csv')
            ) * apu.rad
        phase_rms = rms(_phase, circ=True)
        phase_e_rs = e_rs(_phase, circ=True)

        # random-surface-error efficiency error
        # cov = np.genfromtxt(os.path.join(pyoof_out, f'cov_n{order}.csv'))
        # idx = np.argwhere(cov[0, :].astype(int) > 5)

        params = Table.read(
            os.path.join(pyoof_out, f'fitpar_n{order}.csv'),
            format='ascii'
            )
        I_coeff = params['parfit'][:5]

        qt.add_row([
            pyoof_info['name'], pyoof_info['tel_name'],
            pyoof_info['obs_object'], pyoof_info['obs_date'],
            pyoof_info['meanel'], I_coeff[0], I_coeff[1], I_coeff[2],
            phase_rms, phase_e_rs
            ] + pyoof_info['snr'])

    # updating units
    qt['phase-rms'] *= apu.rad
    qt['meanel'] *= apu.deg
    qt['obs-date'] = Time(qt['obs-date'], format='isot', scale='utc')
    qt['c_dB'] *= apu.dB

    qt.meta = {'order': order}

    return qt


"""CAPRI module."""
import os
import shlex
import shutil
import subprocess
import tempfile
from functools import partial
from math import sqrt
from pathlib import Path

import numpy as np
from Bio import Align
from Bio.Align import substitution_matrices
from Bio.Seq import Seq
from fccpy import read_pdb
from fccpy.contacts import get_intermolecular_contacts
from pdbtools import pdb_segxchain

from haddock import log
from haddock.libs.libontology import PDBFile
from haddock.libs.libpdb import split_by_chain

RES_TO_BE_IGNORED = ["SHA"]

PROT_RES = [
    "ALA",
    "ARG",
    "ASN",
    "ASP",
    "CYS",
    "GLN",
    "GLU",
    "GLY",
    "HIS",
    "ILE",
    "LEU",
    "LYS",
    "MET",
    "PHE",
    "PRO",
    "SER",
    "THR",
    "TRP",
    "TYR",
    "VAL",
    ]

DNA_RES = ["DA", "DC", "DT", "DG"]
# Backbone
PROT_ATOMS = ["C", "N", "CA", "O"]
# Bases
DNA_ATOMS = [
    "C5",
    "N9",
    "N2",
    "C8",
    "O2",
    "N4",
    "N7",
    "C7",
    "N1",
    "N6",
    "C2",
    "O4",
    "C6",
    "N3",
    "C4",
    "O6",
    ]

class CAPRI:
    """CAPRI class."""

    def __init__(
            self,
            reference,
            model_list,
            receptor_chain,
            ligand_chain,
            aln_method,
            path,
            **params,
            ):
        self.reference = reference
        self.model_list = model_list
        self.irmsd_dic = {}
        self.lrmsd_dic = {}
        self.ilrmsd_dic = {}
        self.fnat_dic = {}
        self.dockq_dic = {}
        self.atoms = get_atoms(model_list + [reference])
        self.r_chain = receptor_chain
        self.l_chain = ligand_chain
        self.path = path

        # TODO: For scoring we might need to get one alignment per model
        reference = str(reference)
        model = model_list[0].rel_path
        align_func = get_align(aln_method, **params)
        self.numbering_dic = align_func(reference, model, path)
        if not self.numbering_dic:
            raise CAPRIError("Could not align reference and model")

        # Load the models in the class
        for struct in model_list:
            _ = self.add_chain_from_segid(struct.rel_path)

    def irmsd(self, cutoff=5.0):
        """Calculate the I-RMSD."""
        # Identify reference interface
        ref_interface_resdic = self.identify_interface(self.reference, cutoff)

        # Load interface coordinates
        ref_coord_dic, _ = self.load_coords(
            self.reference, ref_interface_resdic, match=False
            )

        for model in self.model_list:

            mod_coord_dic, _ = self.load_coords(
                model, ref_interface_resdic, match=True
                )

            # Here _coord_dic keys are matched
            #  and formatted as (chain, resnum, atom)
            #  we will use atoms that are present in both
            P = []
            Q = []
            for k in ref_coord_dic.keys() & mod_coord_dic.keys():
                ref_xyz = ref_coord_dic[k]
                mod_xyz = mod_coord_dic[k]

                Q.append(ref_xyz)
                P.append(mod_xyz)

            Q = np.asarray(Q)
            P = np.asarray(P)
            # write_coords("model.pdb", P)
            # write_coords("ref.pdb", Q)

            Q = Q - self.centroid(Q)
            P = P - self.centroid(P)
            U = self.kabsch(P, Q)
            P = np.dot(P, U)
            i_rmsd = self.calc_rmsd(P, Q)
            # write_coords("model_aln.pdb", P)
            # write_coords("ref_aln.pdb", Q)

            self.irmsd_dic[model] = i_rmsd

        return self.irmsd_dic

    def lrmsd(self):
        """Calculate the L-RMSD."""
        ref_coord_dic, _ = self.load_coords(self.reference)

        for model in self.model_list:

            mod_coord_dic, _ = self.load_coords(model, match=True)

            Q = []
            P = []
            # Note: this MUST be sorted since we will use the indexes to
            #  separate between receptor and ligand coordinates
            instersection = sorted(ref_coord_dic.keys() & mod_coord_dic.keys())

            chain_ranges = {}
            for i, segment in enumerate(instersection):
                chain, _, _ = segment
                if chain not in chain_ranges:
                    chain_ranges[chain] = []
                chain_ranges[chain].append(i)

            chain_ranges = make_range(chain_ranges)
            r_start, r_end = chain_ranges[self.r_chain]
            l_start, l_end = chain_ranges[self.l_chain]

            for k in instersection:
                ref_xyz = ref_coord_dic[k]
                mod_xyz = mod_coord_dic[k]

                Q.append(ref_xyz)
                P.append(mod_xyz)

            Q = np.asarray(Q)
            P = np.asarray(P)

            # write_coord_dic("ref.pdb", ref_coord_dic)
            # write_coord_dic("model.pdb", mod_coord_dic)

            # write_coords("ref.pdb", Q)
            # write_coords("model.pdb", P)

            # # move to the origin
            Q = Q - self.centroid(Q)
            P = P - self.centroid(P)

            # get receptor coordinates
            Q_r = Q[r_start: r_end - 1]
            P_r = P[r_start: r_end - 1]

            # Center receptors and get rotation matrix
            # Q_r = Q_r - self.centroid(Q_r)
            # P_r = P_r - self.centroid(P_r)

            U_r = self.kabsch(P_r, Q_r)

            # Center complexes at receptor centroids
            Q = Q - self.centroid(Q_r)
            P = P - self.centroid(P_r)

            # Apply rotation to complex
            #  - complex are now aligned by the receptor
            P = np.dot(P, U_r)

            # write_coords("ref.pdb", Q)
            # write_coords("model.pdb", P)

            # Identify the ligand coordinates
            Q_l = Q[l_start: l_end - 1]
            P_l = P[l_start: l_end - 1]

            # write_coords("ref_l.pdb", Q_l)
            # write_coords("model_l.pdb", P_l)

            # Calculate the RMSD of the ligands
            l_rmsd = self.calc_rmsd(P_l, Q_l)

            # write_coords("ref.pdb", Q)
            # write_coords("model.pdb", P)

            self.lrmsd_dic[model] = l_rmsd

        return self.lrmsd_dic

    def ilrmsd(self, cutoff=10.0):
        """Calculate the Interface Ligand RMSD."""
        # Identify interface
        ref_interface_resdic = self.identify_interface(self.reference, cutoff)

        # Load interface coordinates
        ref_coord_dic, _ = self.load_coords(self.reference, match=False)

        ref_int_coord_dic, _ = self.load_coords(
            self.reference, ref_interface_resdic, match=False
            )

        for model in self.model_list:

            mod_coord_dic, _ = self.load_coords(model, match=True)

            mod_int_coord_dic, _ = self.load_coords(
                model, ref_interface_resdic, match=True
                )

            # write_coord_dic("ref.pdb", ref_int_coord_dic)
            # write_coord_dic("model.pdb", mod_int_coord_dic)

            # find atoms present in both interfaces
            Q_int = []
            P_int = []
            common_keys = ref_int_coord_dic.keys() & mod_int_coord_dic.keys()
            for k in sorted(common_keys):
                ref_xyz = ref_int_coord_dic[k]
                mod_xyz = mod_int_coord_dic[k]

                Q_int.append(ref_xyz)
                P_int.append(mod_xyz)

            Q_int = np.asarray(Q_int)
            P_int = np.asarray(P_int)

            # write_coords("ref.pdb", Q_int)
            # write_coords("model.pdb", P_int)

            # find atoms present in both molecules
            Q = []
            P = []
            intersection = sorted(ref_coord_dic.keys() & mod_coord_dic.keys())
            chain_ranges = {}
            for i, segment in enumerate(intersection):
                chain, _, _ = segment
                if chain not in chain_ranges:
                    chain_ranges[chain] = []
                chain_ranges[chain].append(i)

            chain_ranges = make_range(chain_ranges)
            l_start, l_end = chain_ranges[self.l_chain]

            for k in sorted(ref_coord_dic.keys() & mod_coord_dic.keys()):
                ref_xyz = ref_coord_dic[k]
                mod_xyz = mod_coord_dic[k]

                Q.append(ref_xyz)
                P.append(mod_xyz)

            Q = np.asarray(Q)
            P = np.asarray(P)

            # write_coords("ref.pdb", Q)
            # write_coords("model.pdb", P)

            # put system at origin
            Q_int = Q_int - self.centroid(Q_int)
            P_int = P_int - self.centroid(P_int)

            # # put interfaces at the origin
            # Q_int = Q_int - self.centroid(Q_int)
            # P_int = P_int - self.centroid(P_int)

            # find the rotation that minimizes the interface rmsd
            U_int = self.kabsch(P_int, Q_int)
            P_int = np.dot(P_int, U_int)

            # write_coords("ref.pdb", Q_int)
            # write_coords("model.pdb", P_int)

            # # move the system to the centroid of the interfaces
            Q = Q - self.centroid(Q)
            P = P - self.centroid(P)

            # write_coords("ref_1.pdb", Q)
            # write_coords("model_1.pdb", P)

            Q = Q - self.centroid(Q_int)
            P = P - self.centroid(P_int)

            # write_coords("ref_2.pdb", Q)
            # write_coords("model_2.pdb", P)

            # apply this rotation to the model
            #  - complexes are now aligned by the interfaces
            P = np.dot(P, U_int)

            # write_coords("ref_i.pdb", Q)
            # write_coords("model_i.pdb", P)

            # Calculate the rmsd of the ligand
            Q_l = Q[l_start: l_end - 1]
            P_l = P[l_start: l_end - 1]

            # write_coords("ref_l.pdb", Q_l)
            # write_coords("model_l.pdb", P_l)

            # this will be the interface-ligand-rmsd
            i_l_rmsd = self.calc_rmsd(P_l, Q_l)
            self.ilrmsd_dic[model] = i_l_rmsd

        return self.ilrmsd_dic

    def fnat(self, cutoff=5.0):
        """Calculate the frequency of native contacts."""
        ref_contacts = self.load_contacts(self.reference, cutoff)
        for model in self.model_list:
            model_contacts = self.load_contacts(model, cutoff)
            intersection = ref_contacts & model_contacts
            fnat = len(intersection) / float(len(ref_contacts))
            self.fnat_dic[model] = fnat
        return self.fnat_dic

    def dockq(self):
        """Calculate the DockQ metric."""
        for model in self.model_list:
            irmsd = self.irmsd_dic[model]
            fnat = self.fnat_dic[model]
            lrmsd = self.lrmsd_dic[model]
            dockq = (
                float(fnat)
                + 1 / (1 + (irmsd / 1.5) * (irmsd / 1.5))
                + 1 / (1 + (lrmsd / 8.5) * (lrmsd / 8.5))
                ) / 3
            self.dockq_dic[model] = dockq

        return self.dockq_dic

    def output(
            self,
            clt_threshold,
            sortby_key,
            sort_ascending,
            ):
        """Output the CAPRI results to a .tsv file."""
        self._output_ss(sortby_key, sort_ascending)
        self._output_clt(clt_threshold, sortby_key, sort_ascending,)

    def _output_ss(self, sortby_key, sort_ascending):
        output_l = []
        for model in self.model_list:
            data = {}
            # keep always "model" the first key
            data["model"] = model
            # create the empty rank here so that it will appear
            #  as the second column
            data["caprieval_rank"] = None
            data["score"] = model.score
            if model in self.irmsd_dic:
                data["irmsd"] = self.irmsd_dic[model]
            if model in self.fnat_dic:
                data["fnat"] = self.fnat_dic[model]
            if model in self.lrmsd_dic:
                data["lrmsd"] = self.lrmsd_dic[model]
            if model in self.ilrmsd_dic:
                data["ilrmsd"] = self.ilrmsd_dic[model]
            if model in self.dockq_dic:
                data["dockq"] = self.dockq_dic[model]
            # add cluster data
            data["cluster-id"] = model.clt_id
            data["cluster-ranking"] = model.clt_rank
            data["model-cluster-ranking"] = model.clt_model_rank
            # list of dictionaries
            output_l.append(data)

        output_fname = Path(self.path, "capri_ss.tsv")
        self._dump_file(
            output_l,
            output_fname,
            sortby_key,
            sort_ascending,
            )

    def _output_clt(
            self,
            clt_threshold,
            sortby_key,
            sort_ascending,
            ):
        """Output cluster-based results."""
        has_cluster_info = any(m.clt_id for m in self.model_list)
        if not has_cluster_info:
            return

        # get the cluster data
        clt_data = dict(((m.clt_rank, m.clt_id), []) for m in self.model_list)

        # add models to each cluster
        for model in self.model_list:
            clt_data[(model.clt_rank, model.clt_id)].append(model)

        output_l = []
        for element in clt_data:
            data = {}
            number_of_models_in_cluster = len(clt_data[element])
            # TODO: Refactor these ugly try/excepts
            try:
                score_array = [v.score
                               for v in clt_data[element][: clt_threshold]]
                score_mean, score_stdev = self._calc_stats(
                    score_array, clt_threshold)
            except KeyError:
                score_mean = float("nan")
                score_stdev = float("nan")

            try:
                irmsd_array = [
                    self.irmsd_dic[v]
                    for v in clt_data[element]
                    [: clt_threshold]]
                irmsd_mean, irmsd_stdev = self._calc_stats(
                    irmsd_array, clt_threshold)
            except KeyError:
                irmsd_mean = float("nan")
                irmsd_stdev = float("nan")

            try:
                fnat_array = [self.fnat_dic[v]
                              for v in clt_data[element][:clt_threshold]]
                fnat_mean, fnat_stdev = self._calc_stats(
                    fnat_array, clt_threshold)
            except KeyError:
                fnat_mean = float("nan")
                fnat_stdev = float("nan")

            try:
                lrmsd_array = [self.lrmsd_dic[v]
                               for v in clt_data[element]
                               [: clt_threshold]]
                lrmsd_mean, lrmsd_stdev = self._calc_stats(
                    lrmsd_array, clt_threshold)
            except KeyError:
                lrmsd_mean = float("nan")
                lrmsd_stdev = float("nan")
            try:
                dockq_array = [self.dockq_dic[v]
                               for v in clt_data[element]
                               [: clt_threshold]]
                dockq_mean, dockq_stdev = self._calc_stats(
                    dockq_array, clt_threshold)
            except KeyError:
                dockq_mean = float("nan")
                dockq_stdev = float("nan")

            data["cluster_rank"] = element[0]
            data["cluster_id"] = element[1]
            data["n"] = number_of_models_in_cluster
            if number_of_models_in_cluster < clt_threshold:
                # under-evaluated, the mean was divided by a value
                #  larger than the total number of models in the cluster
                data["under_eval"] = "yes"
            else:
                data["under_eval"] = "-"

            data["score"] = score_mean
            data["score_std"] = score_stdev
            data["irmsd"] = irmsd_mean
            data["irmsd_std"] = irmsd_stdev
            data["fnat"] = fnat_mean
            data["fnat_std"] = fnat_stdev
            data["lrmsd"] = lrmsd_mean
            data["lrmsd_std"] = lrmsd_stdev
            data["dockqn"] = dockq_mean
            data["dockq_std"] = dockq_stdev

            output_l.append(data)

        output_fname = Path(self.path, "capri_clt.tsv")
        info_header = "#" * 40 + os.linesep
        info_header += "# `caprieval` cluster-based analysis" + os.linesep
        info_header += "#" + os.linesep
        info_header += f"# > sortby_key={sortby_key}" + os.linesep
        info_header += f"# > sort_ascending={sort_ascending}" + os.linesep
        info_header += f"# > clt_threshold={clt_threshold}" + os.linesep
        info_header += "#" + os.linesep
        info_header += (
            "# NOTE: if under_eval=yes, it means that there were less models in"
            " a cluster than" + os.linesep
            )
        info_header += (
            "#    clt_threshold, thus these values were under evaluated."
            + os.linesep
            )
        info_header += (
            "#   You might need to tweak the value of clt_threshold or change"
            " some parameters" + os.linesep
            )
        info_header += (
            "#    in `clustfcc` depending on your analysis." + os.linesep
            )
        info_header += "#" + os.linesep
        info_header += "#" * 40

        self._dump_file(
            output_l,
            output_fname,
            sortby_key,
            sort_ascending,
            info_header=info_header,
            )

    def _dump_file(
            self,
            container,
            output_fname,
            sortby_key,
            sort_ascending,
            info_header="",
            ):

        # rank
        ranked_output_l = self._rank(
            container, key='score', ascending=True
            )

        # sort
        sorted_keys = self._sort(
            ranked_output_l, key=sortby_key, ascending=sort_ascending
            )

        header = "\t".join(ranked_output_l[0].keys())

        if info_header:
            header = info_header + os.linesep + header

        with open(output_fname, "w") as out_fh:
            out_fh.write(header + os.linesep)
            for idx, _ in sorted_keys:
                row_l = []
                for value in ranked_output_l[idx].values():
                    if isinstance(value, Path):
                        row_l.append(str(value))
                    elif isinstance(value, PDBFile):
                        row_l.append(str(value.rel_path))
                    elif isinstance(value, int):
                        row_l.append(f"{value}")
                    elif isinstance(value, str):
                        row_l.append(f"{value}")
                    elif value is None:
                        row_l.append("-")
                    else:
                        row_l.append(f"{value:.3f}")
                out_fh.write("\t".join(row_l) + os.linesep)

    @staticmethod
    def _sort(container, key, ascending):
        # Sort the column
        key_values = [(i, k[key]) for i, k in enumerate(container)]
        key_values.sort(key=lambda x: x[1], reverse=not ascending)
        return key_values

    @staticmethod
    def _rank(container, key, ascending):
        rankkey_values = [(i, k[key]) for i, k in enumerate(container)]
        rankkey_values.sort(key=lambda x: x[1], reverse=not ascending)
        for i, k in enumerate(rankkey_values, start=1):
            idx, _ = k
            container[idx]["caprieval_rank"] = i
        return container

    @staticmethod
    def _calc_stats(data, n):
        """Calculate the mean and stdev."""
        mean = sum(data) / n
        var = sum((x - mean) ** 2 for x in data) / n
        stdev = sqrt(var)
        return mean, stdev

    @staticmethod
    def identify_interface(pdb_f, cutoff=5.0):
        """Identify the interface."""
        if isinstance(pdb_f, PDBFile):
            pdb_f = pdb_f.rel_path
        pdb = read_pdb(pdb_f)

        interface_resdic = {}
        for atom_i, atom_j in get_intermolecular_contacts(pdb, cutoff):

            if atom_i.chain not in interface_resdic:
                interface_resdic[atom_i.chain] = []
            if atom_j.chain not in interface_resdic:
                interface_resdic[atom_j.chain] = []

            if atom_i.resid not in interface_resdic[atom_i.chain]:
                interface_resdic[atom_i.chain].append(atom_i.resid)
            if atom_j.resid not in interface_resdic[atom_j.chain]:
                interface_resdic[atom_j.chain].append(atom_j.resid)

        return interface_resdic

    @staticmethod
    def load_contacts(pdb_f, cutoff=5.0):
        """Load residue-based contacts."""
        con_list = []
        if isinstance(pdb_f, PDBFile):
            pdb_f = pdb_f.rel_path
        structure = read_pdb(pdb_f)
        for atom_i, atom_j in get_intermolecular_contacts(structure, cutoff):
            con = (atom_i.chain, atom_i.resid, atom_j.chain, atom_j.resid)
            con_list.append(con)
        return set(con_list)

    @staticmethod
    def calc_rmsd(V, W):
        """Calculate the RMSD from two vectors."""
        diff = np.array(V) - np.array(W)
        N = len(V)
        return np.sqrt((diff * diff).sum() / N)

    @staticmethod
    def kabsch(P, Q):
        """Find the rotation matrix using Kabsch algorithm."""
        # Covariance matrix
        P = np.array(P)
        Q = np.array(Q)
        C = np.dot(np.transpose(P), Q)
        # use SVD
        V, S, W = np.linalg.svd(C)
        d = (np.linalg.det(V) * np.linalg.det(W)) < 0.0
        if d:
            S[-1] = -S[-1]
            V[:, -1] = -V[:, -1]
        # Create Rotation matrix U
        U = np.dot(V, W)
        return U

    @staticmethod
    def centroid(X):
        """Get the centroid."""
        X = np.array(X)
        return X.mean(axis=0)

    @staticmethod
    def add_chain_from_segid(pdb_path):
        """Replace the chainID with the segID."""
        temp_f = tempfile.NamedTemporaryFile(delete=False, mode="w+t")
        with open(pdb_path) as fh:
            for line in list(pdb_segxchain.run(fh)):
                temp_f.writelines(line)
        temp_f.close()
        # REPLACE!
        new_pdb_path = shutil.move(temp_f.name, pdb_path)
        return new_pdb_path

    def load_coords(self, pdb_f, filter_resdic=None, match=False):
        """Load coordinates from PDB."""
        coord_dic = {}
        chain_dic = {}
        idx = 0
        if isinstance(pdb_f, PDBFile):
            pdb_f = pdb_f.rel_path
        with open(pdb_f, "r") as fh:
            for line in fh.readlines():
                if line.startswith("ATOM"):

                    atom_name = line[12:16].strip()
                    resname = line[17:20].strip()
                    chain = line[21]
                    resnum = int(line[22:26])

                    x = float(line[30:38])
                    y = float(line[38:46])
                    z = float(line[46:54])
                    coords = np.asarray([x, y, z])

                    if match:
                        try:
                            resnum = self.numbering_dic[chain][resnum]
                        except KeyError:
                            # this residue is not matched, and so it should
                            #  not be considered
                            # self.log(
                            #     f"WARNING: {chain}.{resnum}.{atom_name}"
                            #     " was not matched!"
                            #     )
                            continue

                    # identifier = f"{chain}.{resnum}.{atom_name}"
                    identifier = (chain, resnum, atom_name)

                    if atom_name not in self.atoms[resname]:
                        continue

                    if chain not in chain_dic:
                        chain_dic[chain] = []

                    if filter_resdic:
                        # Only retrieve coordinates from the filter_resdic
                        if (
                                chain in filter_resdic
                                and resnum in filter_resdic[chain]
                                ):
                            coord_dic[identifier] = coords
                            chain_dic[chain].append(idx)
                            idx += 1

                    else:
                        # retrieve everything
                        coord_dic[identifier] = coords
                        chain_dic[chain].append(idx)
                        idx += 1

        chain_ranges = {}
        for chain in chain_dic:
            min_idx = min(chain_dic[chain])
            max_idx = max(chain_dic[chain])
            chain_ranges[chain] = (min_idx, max_idx)

        return coord_dic, chain_ranges

class CAPRIError(Exception):
    """Raised when something goes wrong with the CAPRI class."""

    def __init__(self, msg=""):
        self.msg = msg
        super().__init__(self.msg)

def get_atoms(pdb_list):
    """Identify what is the molecule type of each PDB."""
    atom_dic = {}
    atom_dic.update(dict((r, PROT_ATOMS) for r in PROT_RES))
    atom_dic.update(dict((r, DNA_ATOMS) for r in DNA_RES))
    for pdb in pdb_list:
        if isinstance(pdb, PDBFile):
            pdb = pdb.rel_path
        with open(pdb) as fh:
            for line in fh.readlines():
                if line.startswith(("ATOM", "HETATM")):
                    resname = line[17:20].strip()
                    atom_name = line[12:16].strip()
                    element = line[76:78].strip()
                    if (
                            resname not in PROT_RES
                            and resname not in DNA_RES
                            and resname not in RES_TO_BE_IGNORED
                            ):
                        # its neither DNA nor protein, use the heavy atoms
                        # WARNING: Atoms that belong to unknown residues must
                        #  be bound to a residue name;
                        #   For example: residue NEP, also contains
                        #  CB and CG atoms, if we do not bind it to the
                        #  residue name, the next functions will include
                        #  CG and CG atoms in the calculations for all
                        #  other residue names
                        if element != "H":
                            if resname not in atom_dic:
                                atom_dic[resname] = []
                            if atom_name not in atom_dic[resname]:
                                atom_dic[resname].append(atom_name)
    return atom_dic

def pdb2fastadic(pdb_f):
    """Write the sequence as a fasta."""
    res_codes = dict(
        [
            ("CYS", "C"),
            ("ASP", "D"),
            ("SER", "S"),
            ("GLN", "Q"),
            ("LYS", "K"),
            ("ILE", "I"),
            ("PRO", "P"),
            ("THR", "T"),
            ("PHE", "F"),
            ("ASN", "N"),
            ("GLY", "G"),
            ("HIS", "H"),
            ("LEU", "L"),
            ("ARG", "R"),
            ("TRP", "W"),
            ("ALA", "A"),
            ("VAL", "V"),
            ("GLU", "E"),
            ("TYR", "Y"),
            ("MET", "M"),
            ("DA", "A"),
            ("DG", "G"),
            ("DC", "C"),
            ("DT", "T"),
            ]
        )
    seq_dic = {}
    with open(pdb_f) as fh:
        for line in fh.readlines():
            if line.startswith("ATOM"):
                res_num = int(line[22:26])
                res_name = line[17:20].strip()
                chain = line[21]
                if res_name in RES_TO_BE_IGNORED:
                    continue
                try:
                    one_letter = res_codes[res_name]
                except KeyError:
                    one_letter = "X"
                if chain not in seq_dic:
                    seq_dic[chain] = {}
                seq_dic[chain][res_num] = one_letter
    return seq_dic

def get_align(method, **kwargs):
    """Get the alignment function."""
    log.info(f"Using {method} alignment")
    if method == "structure":
        return partial(align_strct, lovoalign_exec=kwargs["lovoalign_exec"])
    elif method == "sequence":
        return partial(align_seq)
    else:
        available_alns = ("sequence", "structure")
        raise ValueError(
            f"Alignment method {method!r} not recognized. "
            f"Available options are {', '.join(available_alns)}"
            )

def align_strct(reference, model, output_path, lovoalign_exec=None):
    """Structuraly align and get numbering relationship."""
    if lovoalign_exec is None:
        log.error(
            "Structural alignment needs LovoAlign "
            "get it at github.com/m3g/lovoalign"
            )
        raise CAPRIError("Path to LovoAlign executable required.")

    if not lovoalign_exec:
        raise CAPRIError("lovoalign_exec parameter not defined ")

    if not os.access(lovoalign_exec, os.X_OK):
        raise CAPRIError(f"{lovoalign_exec!r} for LovoAlign is not executable")

    numbering_dic = {}
    protein_a_dic = dict(
        (str(e.stem).split("_")[-1], e) for e in split_by_chain(reference)
        )
    protein_b_dic = dict(
        (str(e.stem).split("_")[-1], e) for e in split_by_chain(model)
        )

    # check if chain ids match
    if protein_a_dic.keys() != protein_b_dic.keys():
        # TODO: Make this a clearer raise
        return numbering_dic

    for chain in protein_a_dic.keys():
        pa_seqdic = pdb2fastadic(protein_a_dic[chain])
        pb_seqdic = pdb2fastadic(protein_b_dic[chain])
        # logging.debug(f"Structurally aligning chain {chain}")
        numbering_dic[chain] = {}
        cmd = (
            f"{lovoalign_exec} -p1 {protein_a_dic[chain]} "
            f"-p2 {protein_b_dic[chain]} "
            f"-c1 {chain} -c2 {chain}"
            )

        # logging.debug(f"Command is: {cmd}")
        p = subprocess.run(shlex.split(cmd), capture_output=True, text=True)
        lovoalign_out = p.stdout.split(os.linesep)

        # we don"t need the splitted proteins anymore
        protein_a_dic[chain].unlink()
        protein_b_dic[chain].unlink()

        # find out where the alignment starts and ends
        alignment_pass = True
        for i, line in enumerate(lovoalign_out):
            if "SEQUENCE ALIGNMENT" in line:
                # there are 2 extra white lines after this header
                alignment_start_index = i + 2
            elif "FINAL" in line:
                # there are 2 extra white lines after this header
                alignment_end_index = i - 2
            elif "ERROR" in line:
                failed_pdb = line.split()[-1]
                _msg = (
                    f"LovoAlign could not read {failed_pdb} " "is it a ligand?"
                    )
                log.warning(_msg)
                alignment_pass = False

                for elem in [k for k in pa_seqdic[chain]]:
                    numbering_dic[chain][elem] = elem

        if not alignment_pass:
            # This alignment failed, move on to the next
            log.warning(
                f"Skipping alignment of chain {chain}, "
                "used sequential matching"
                )
            continue

        aln_l = lovoalign_out[alignment_start_index:alignment_end_index]

        # dump this alignment to a file
        aln_fname = Path(output_path, f"lovoalign_{chain}.aln")
        log.debug(f"Writing alignment to {aln_fname.name}")
        with open(aln_fname, "w") as fh:
            fh.write(os.linesep.join(aln_l))

        # remove the line between the alignment segments
        alignment = [aln_l[i: i + 3][:2] for i in range(0, len(aln_l), 3)]
        # 100% (5 identical nucleotides / min(length(A),length(B))).
        len_seq_a = len(pa_seqdic[chain])
        len_seq_b = len(pb_seqdic[chain])
        identity = (
            (len_seq_a - sum([e[0].count("-") for e in alignment]))
            / min(len_seq_a, len_seq_b)
            * 100
            )

        if identity <= 40.0:
            log.warning(
                f"\"Structural\" identity of chain {chain} is {identity:.2f}%,"
                " please check the results carefully"
                )
        else:
            log.info(
                f"\"Structural\" identity of chain {chain} is {identity:.2f}%"
                )

        # logging.debug("Reading alignment and matching numbering")
        for element in alignment:
            line_a, line_b = element

            resnum_a, seq_a, _ = line_a.split()
            resnum_b, seq_b, _ = line_b.split()

            resnum_a = int(resnum_a) - 1
            resnum_b = int(resnum_b) - 1

            for resname_a, resname_b in zip(seq_a, seq_b):
                if resname_a != "-":
                    resnum_a += 1

                if resname_b != "-":
                    resnum_b += 1

                if resname_a != "-" and resname_b != "-":
                    numbering_dic[chain][resnum_b] = resnum_a

    izone_fname = Path(output_path, "lovoalign.izone")
    log.debug(f"Saving .izone to {izone_fname.name}")
    dump_as_izone(izone_fname, numbering_dic)

    return numbering_dic

def align_seq(reference, model, output_path):
    """Sequence align and get the numbering relationship."""
    seqdic_a = pdb2fastadic(reference)
    seqdic_b = pdb2fastadic(model)

    if seqdic_a.keys() != seqdic_b.keys():
        # TODO: Implement chain-matching here
        return False

    align_dic = {}
    for a, b in zip(seqdic_a, seqdic_b):

        align_dic[a] = {}

        seq_a = Seq("".join(seqdic_a[a].values()))
        seq_b = Seq("".join(seqdic_b[b].values()))

        aligner = Align.PairwiseAligner()
        aligner.substitution_matrix = substitution_matrices.load("BLOSUM62")
        alns = aligner.align(seq_a, seq_b)
        top_aln = alns[0]

        aln_fname = Path(output_path, f"blosum62_{a}.aln")
        log.debug(f"Writing alignment to {aln_fname.name}")
        with open(aln_fname, "w") as fh:
            fh.write(str(top_aln))
        aligned_seg_a, aligned_seg_b = top_aln.aligned

        # this should always be true
        assert len(aligned_seg_a) == len(aligned_seg_b)

        identity = (
            str(top_aln).count("|") / float(min(len(seq_a), len(seq_b)))
            ) * 100

        if not any(e for e in top_aln.aligned):
            # No alignment!
            log.warning(
                f"No alignment for chain {a} is it protein/dna? "
                "Matching sequentially"
                )
            if all("X" in s for s in seq_a) and all("X" in s for s in seq_b):
                # this sequence contains only ligands, do it manually
                if len(seq_a) != len(seq_b):
                    # we cannot handle this
                    raise f"Cannot align chain {b}"
                for res_a, res_b in zip(seqdic_a[a], seqdic_b[b]):
                    align_dic[a].update({res_a: res_b})
        else:
            if identity <= 40.0:
                # Identity is very low
                log.warning(
                    f"Sequence identity of chain {a} is {identity:.2f}%,"
                    " please check the results carefully"
                    )
                log.warning(
                    "Please use alignment_method = \"structure\" instead")
            else:
                log.info(f"Sequence identity of chain {a} is {identity:.2f}%")
            for seg_a, seg_b in zip(aligned_seg_a, aligned_seg_b):
                start_a, end_a = seg_a
                start_b, end_b = seg_b

                reslist_a = list(seqdic_a[a].keys())[start_a:end_a]
                reslist_b = list(seqdic_b[b].keys())[start_b:end_b]

                align_dic[a].update(
                    dict((i, j) for i, j in zip(reslist_a, reslist_b))
                    )
    izone_fname = Path(output_path, "blosum62.izone")
    log.debug(f"Saving .izone to {izone_fname.name}")
    dump_as_izone(izone_fname, align_dic)

    return align_dic

def make_range(chain_range_dic):
    """Expand a chain dictionary into ranges."""
    chain_ranges = {}
    for chain in chain_range_dic:
        min_idx = min(chain_range_dic[chain])
        max_idx = max(chain_range_dic[chain])
        chain_ranges[chain] = (min_idx, max_idx)
    return chain_ranges

def dump_as_izone(fname, numbering_dic):
    """Dump the numbering dictionary as .izone."""
    # FIXME: Collapse the izones so its faster to load in profit
    with open(fname, "w") as fh:
        for chain in numbering_dic:
            for bound_res in numbering_dic[chain]:
                unbound_res = numbering_dic[chain][bound_res]
                #
                izone_str = (
                    "ZONE "
                    f"{chain}{bound_res}:{chain}{unbound_res}"
                    f"{os.linesep}"
                    )
                fh.write(izone_str)

# # debug only
# def write_coord_dic(output_name, coord_dic):
#     """Add a dummy atom to a PDB file according to a list of coordinates."""
#     with open(output_name, "w") as fh:
#         for i, k in enumerate(coord_dic):
#             atom_num = f"{i+1}".rjust(4, " ")
#             chain, resnum, atom = k
#             resnum = int(resnum)
#             resnum = f"{resnum}".rjust(3, " ")
#             atom_name = f"{atom}".rjust(3, " ")
#             x, y, z = coord_dic[k]
#             dum_x = f"{x:.3f}".rjust(7, " ")
#             dum_y = f"{y:.3f}".rjust(7, " ")
#             dum_z = f"{z:.3f}".rjust(7, " ")
#             dummy_line = (
#                 f"ATOM   {atom_num} {atom_name}  DUM {chain} {resnum}   "
#                 f"  {dum_x} {dum_y} {dum_z}  1.00  1.00   "
#                 "        H  " + os.linesep
#                 )
#             fh.write(dummy_line)

# # debug only
# def write_coords(output_name, coor_list):
#     """Add a dummy atom to a PDB file according to a list of coordinates."""
#     with open(output_name, "w") as fh:
#         for i, dummy_coord in enumerate(coor_list):
#             atom_num = f"{i}".rjust(4, " ")
#             resnum = f"{i}".rjust(3, " ")
#             dum_x = f"{dummy_coord[0]:.3f}".rjust(7, " ")
#             dum_y = f"{dummy_coord[1]:.3f}".rjust(7, " ")
#             dum_z = f"{dummy_coord[2]:.3f}".rjust(7, " ")
#             dummy_line = (
#                 f"ATOM   {atom_num}  H   DUM X {resnum}   "
#                 f"  {dum_x} {dum_y} {dum_z}  1.00  1.00   "
#                 "        H  " + os.linesep
#                 )
#             fh.write(dummy_line)

# # debug only
# def write_pymol_viz(resdic):
#     """Write PyMol vizualitation."""
#     for k in resdic:
#         reslist = "+".join(map(str, resdic[k]))
#         cmd = f"sele {k}, chain {k} and resid {reslist}"
#         print(cmd)


import os
import numpy as np
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
from matplotlib import rc
rc('font', **{'family': 'serif'})
from matplotlib import rcParams
rcParams['xtick.direction'] = 'out'
rcParams['ytick.direction'] = 'out'
rcParams['xtick.labelsize'] = 18
rcParams['ytick.labelsize'] = 18
rcParams['lines.linewidth'] = 1.85
rcParams['axes.labelsize'] = 20
rcParams.update({'figure.autolayout': True})
np.set_printoptions(threshold=10000)
import pickle

class Grouping(object):
    def __init__(self, G, geoOption, contig, name, bounds, minCutoff, ratioCutoff, groupCutoff, sig_t, E=None):
        self.name = name
        self.G = G
        self.geoOption = geoOption
        self.contig = contig
        self.minCutoff = minCutoff
        self.ratioCutoff = ratioCutoff
        self.groupCutoff = groupCutoff
        self.structure = bounds
        self.mask = np.array([i for c in range(max(bounds) + 1) for i, g in enumerate(bounds) if g == c])
        self.counts = dict(zip(*np.unique(bounds, return_counts=True)))
        self.maxOrder = np.max(list(self.counts.values()))
        self.fname = '{}_{}_{}g_m{}_r{}_g{}'.format(name, geoOption, G, round(minCutoff, 2), round(ratioCutoff, 2), groupCutoff)
        self.sig_t = sig_t
        self.E = E

    def __str__(self):
        s = 'Name: {}'.format(self.name)
        s += '\nGeometry: {}'.format(self.geoOption)
        s += '\nMaxOrder: {}'.format(self.maxOrder)
        #s += '\nStructure: {}'.format(' '.join(self.structure.astype(str)))
        struct = np.concatenate(([0], np.cumsum([self.counts[i] for i in range(max(self.structure) + 1)])))
        s += '\nStructure:'
        for i in range(len(struct) - 1):
            s += '\n\tGroup {}: '.format(i)
            s += '{}-{}'.format(struct[i] + 1, struct[i+1]) if struct[i] + 1 != struct[i+1] else '{}'.format(struct[i] + 1)
        #s += '\nMask: {}'.format(' '.join(self.mask.astype(str)))
        s += '\nCounts: {}'.format(', '.join(['CG-{}={}'.format(key, item) for key, item in self.counts.items()]))
        s += '\nFileName: {}'.format(self.fname)

        return s

    def writeFile(self):
        # Get the coarse group bounds given the total cross sections
        s = '{}\n{}'.format(self.name, self.structure.tolist())

        with open('XS/structure_{}'.format(self.fname), 'w') as f:
            f.write(s)

def getMats(geoOption):
    if geoOption == 'inf':
        return [0]
    elif geoOption == 'full':
        return [0, 3, 9]
    elif geoOption == 'core1':
        return [0, 1, 2, 9]
    elif geoOption == 'core2':
        return [0, 1, 3, 9]
    elif 'c5g7' in geoOption:
        return [0, 3, 9]
        if 'uo2' in geoOption:
            return [5, 9]
        elif 'mox' in geoOption:
            if 'low' in geoOption:
                return [6, 9]
            elif 'mid' in geoOption:
                return [7, 9]
            elif 'high' in geoOption:
                return [8, 9]
            else:
                return [6, 7, 8, 9]
        else:
            return [5, 6, 7, 8, 9]
    else:
        raise NotImplementedError('{} has not been implented'.format(geoOption))

def getXS(fname):
    '''
    Read the input file and extract the total cross sections

    input
        fname: name of the file to read

    output
        sig_t: numpy array containing the total cross sections for each material
    '''

    # Initialize the total cross section list
    sig_t = []
    with open(fname) as f:
        # Get the number of materials and number of groups
        nMat, nGroups = [int(d) for d in f.readline().split()[:2]]
        # Skip 2 lines to pass the energy bounds and velocity
        E = np.array(f.readline().split()).astype(float)[::-1]
        f.readline()

        # Loop over the materials
        for m in range(nMat):
            # Initialize temp list for total XS for material m
            t = []
            # Skip the material name
            f.readline()
            # Get the number of legendre moments
            nLegendre = int(f.readline().split()[0])
            # Loop over groups and extract the total XS
            for g in range(nGroups):
                t.append(float(f.readline().split()[0]))
            # Skip the scattering moments
            for g in range(nGroups * nLegendre):
                f.readline()
            # Append the temp XS list to the main list
            sig_t.append(t)

    return np.array(sig_t), E

def findContiguousBounds(sig_t, minCutoff, ratioCutoff, groupCutoff):
    '''
    Find the coarse group structure given a set of cross sections

    Inputs:
        sig_t: array of the total XS indexed by material [mat, G]
        minCutoff: XS below which the ratio cutoff is ignored
        ratioCutoff: Ratio of largest to smallest XS in a coarse group
        groupCutoff: Maximum number of fine groups in a coarse group
    Outputs:
        bounds: Starting fine group indices for each coarse group
    '''

    def reset(xs):
        '''Reset the min and max cutoffs to the input cross section'''
        minXS = np.maximum(xs, minCutoff)
        maxXS = xs
        return minXS, maxXS

    sig_t = np.array(sig_t)

    # Get the number of groups
    nGroup = len(sig_t[0])

    # Initialize the boundary at the lowest energy group
    bounds = [nGroup + 1]

    # Initialize the cutoff bounds
    minXS, maxXS = reset(sig_t[:,-1])

    # Loop through the cross sections from greatest to least
    for i, xs in enumerate(sig_t.T[::-1]):
        group = nGroup - i

        # Check if the xs is below the min or above the max
        minXS = np.minimum(xs, minXS)
        maxXS = np.maximum(xs, maxXS)
        ratio = maxXS - minXS

        # Check for a ratio that is too large
        if (max(ratio) > ratioCutoff and max(maxXS) > minCutoff) or bounds[-1] - group > groupCutoff:
            bounds.append(group + 1)
            # Reset the cutoff bounds
            minXS, maxXS = reset(xs)

    # Add the highest energy group bound
    bounds.append(1)
    # Reverse to the natural order of structures
    bounds = np.array(bounds[::-1])

    return np.array([i for i, b in enumerate(bounds[1:] - bounds[:-1]) for j in range(b)])

def findNoncontiguousBounds(sig_t, minCutoff=1.0, ratioCutoff=2.0, groupCutoff=60):
    '''
    Find the coarse group structure given a set of cross sections

    Inputs:
        sig_t: array of the total XS indexed by material [mat, G]
        minCutoff: XS below which the ratio cutoff is ignored
        ratioCutoff: Ratio of largest to smallest XS in a coarse group
        groupCutoff: Maximum number of fine groups in a coarse group
    Outputs:
        structure: mapping of fine group to coarse group
    '''

    G = len(sig_t[0])

    # Get the minimum and maximum of all materials
    minXS = np.min(sig_t, axis=0)
    minmap = np.argsort(minXS)[::-1]
    maxXS = np.max(sig_t, axis=0)
    maxmap = np.argsort(maxXS)[::-1]

    unsorted_groups = list(range(G))

    # Initialize a dictionary
    structure = np.zeros(G) - 1
    coarse_group_index = 0
    while len(unsorted_groups) > 0:
        # Get the largest unsorted cross section
        maximum = maxXS[maxmap[0]]
        # If above the minimum cross section cutoff
        if maximum > minCutoff:
            # Find the ratio for the current group
            #cutoff = max(ratioCutoff, maximum / minXS[maxmap[0]])
            cutoff = ratioCutoff
            # Select all groups with a smaller ratio than the cutoff
            cutmap = minmap[maximum - minXS[minmap] <= cutoff]
            if len(cutmap) == 0:
                cutmap = [maxmap[0]]
        else:
            # Select all remaining groups
            cutmap = minmap
        # Only allow groupCutoff groups into the current group
        cutmap = cutmap[:groupCutoff]
        # Assign selected groups to the coarse_group_index
        structure[cutmap] = coarse_group_index
        # Remove the assigned groups from minmap, maxmap, and unsorted_groups
        minmap = np.array([g for g in minmap if g not in cutmap])
        maxmap = np.array([g for g in maxmap if g not in cutmap])
        unsorted_groups = [g for g in unsorted_groups if g not in cutmap]
        # Increment the coarse_group_index
        coarse_group_index += 1

    return (structure * -1 + max(structure)).astype(int)

def computeBounds(G, geoOption, contig, minCutoff, ratioCutoff, groupCutoff):
    '''
    This function determines which fine groups should belong to each coarse group

    The boundaries will not necessarily be contiguous

    Inputs:
        fname: name of the cross section file in anlxs format
    '''

    # Contig options
    names = {0: 'contiguous_min',
             1: 'contiguous_max',
             2: 'contiguous_mean',
             3: 'continuous_difference',
             4: 'noncontiguous_min',
             5: 'noncontiguous_max',
             6: 'noncontiguous_mean',
             7: 'noncontiguous_difference'}

    # Select the material types for the geoOption
    mats = getMats(geoOption)

    # Load the total cross section for group G
    sig_t, E = getXS('XS/{}gXS.anlxs'.format(G))

    # Slice out only the cross sections used for the current geo
    mask = np.array([r in mats for r in range(len(sig_t))]).astype(bool)
    sig_t = sig_t[mask]

    # Number the energy groups
    groups = np.arange(G)

    # Get the minimum XS across all materials
    minXS = np.min(sig_t, axis=0)
    # Get the maximum XS across all materials
    maxXS = np.max(sig_t, axis=0)
    # Get the mean XS across all materials
    meanXS = np.mean(sig_t, axis=0)

    # Select which XS upon which to compute the bounds
    if contig == 0:  # contiguous min
        mask = findContiguousBounds([minXS], minCutoff, ratioCutoff, groupCutoff)
    elif contig == 1:  # contiguous max
        mask = findContiguousBounds([maxXS], minCutoff, ratioCutoff, groupCutoff)
    elif contig == 2:  # contiguous mean
        mask = findContiguousBounds([meanXS], minCutoff, ratioCutoff, groupCutoff)
    elif contig == 3:  # continuous difference
        mask = findContiguousBounds(sig_t, minCutoff, ratioCutoff, groupCutoff)
    elif contig == 4:  # non-contiguous min
        mask = findNoncontiguousBounds([minXS], minCutoff, ratioCutoff, groupCutoff)
    elif contig == 5:  # non-contiguous max
        mask = findNoncontiguousBounds([maxXS], minCutoff, ratioCutoff, groupCutoff)
    elif contig == 6:  # non-contiguous mean
        mask = findNoncontiguousBounds([meanXS], minCutoff, ratioCutoff, groupCutoff)
    elif contig == 7:  # non-contiguous difference
        mask = findNoncontiguousBounds(sig_t, minCutoff, ratioCutoff, groupCutoff)

    structure = Grouping(G, geoOption, contig, names[contig], mask, minCutoff, ratioCutoff, groupCutoff, sig_t, E)

    makePlot(structure, minXS, maxXS)

    return structure

def makePlot(structure, minXS, maxXS):
    # Plot the cross sections
    colors = ['b', 'g', 'r', 'c', 'm', 'orange', 'purple', 'grey']
    ls = '-' if structure.contig < 4 else '--'

    groups = np.arange(structure.G)
    mask = structure.structure

    for i, m in enumerate(mask):
        submask = mask == m
        plt.fill_between(groups[submask] + 1, minXS[submask], maxXS[submask], label='Group {}'.format(m), linestyle=ls, color=colors[m%len(colors)])
        #plt.semilogy(groups[i], minXS[i], label='Group {}'.format(m), ls=ls, marker='o', c=colors[m%len(colors)])
        #plt.semilogy(groups[i], maxXS[i], label='Group {}'.format(m), ls=ls, marker='s', c=colors[m%len(colors)])
    plt.yscale('log')
    plt.xlabel('Fine group number')
    plt.ylabel('Total cross section [cm$^{-1}$]')
    plt.xlim([0, structure.G+1])
    plt.grid(True)
    #plt.title(structure.name + ' {}'.format(len(structure.counts.keys())))
    plt.savefig('XS/structure_plot_{}.png'.format(structure.fname), transparent=True)
    plt.clf()

    E = structure.E[::-1]
    E = 0.5 * (E[1:] + E[:-1])

    for i, m in enumerate(mask):
        submask = mask == m
        plt.fill_between(E[submask], minXS[submask], maxXS[submask], label='Group {}'.format(m), linestyle=ls, color=colors[m%len(colors)])
        #plt.semilogy(groups[i], minXS[i], label='Group {}'.format(m), ls=ls, marker='o', c=colors[m%len(colors)])
        #plt.semilogy(groups[i], maxXS[i], label='Group {}'.format(m), ls=ls, marker='s', c=colors[m%len(colors)])
    plt.yscale('log')
    plt.xscale('log')
    plt.xlabel('Energy [MeV]')
    plt.ylabel('Total cross section [cm$^{-1}$]')
    plt.grid(True)
    #plt.title(structure.name + ' {}'.format(len(structure.counts.keys())))
    plt.savefig('XS/structure_plot_{}_E.png'.format(structure.fname), transparent=True)
    plt.clf()

def getInfo(task, kill=False):
    Gs = [44, 238, 1968]
    geos = ['inf', 'full', 'core1', 'core2', 'c5g7-small']
    contigs = [1]
    mins = [0]
    ratios = [1.3]
    groups = [60]

    item = 0
    groupBounds = []

    for geo in geos:
        for G in Gs:
            for contig in contigs:
                for min_cutoff in mins:
                    for ratio_cutoff in ratios:
                        for group_cutoff in groups:
                            if item == task and not kill:
                                return G, geo, contig, min_cutoff, ratio_cutoff, group_cutoff
                            else:
                                item += 1
        groupBounds.append(item)
    if kill:
        groupBounds = [0] + groupBounds
        for i, G in enumerate(Gs):
            print('Group {:4d} cutoffs: {:>5d}-{:<5d}'.format(G, groupBounds[i] + 1, groupBounds[i+1]))
    exit()

def f(s):
    '''Convert a structure into a proper string'''
    try:
        return '{}'.format(s[0] + 1) if s[1] - 1 == s[0] else '{}-{}'.format(s[0] + 1, s[1])
    except IndexError:
        return ''

if __name__ == '__main__':
    # getInfo(1, True)

    #task = os.environ['SLURM_ARRAY_TASK_ID']
    #task = int(task) - 1

    for geo in ['full', 'core1', 'core2']:
        print(geo)
        s = '& 44-group & \multicolumn{2}{c|}{238-group} & \multicolumn{4}{c}{1968-group}\\\\\n'
        s += '\hline\n'
        s += 'add\'l. \# & (+0) & (+0) & (+20) & (+0) & (+20) & (+40) & (+60)\\\\\n'
        s += '\hline\n'

        # Build the coarse-group structure
        s44 = computeBounds(44, geo, 1, 0.0, 1.3, 60)
        c44 = np.concatenate(([0], np.cumsum([s44.counts[i] for i in range(max(s44.structure)+1)])))
        s238 = computeBounds(238, geo, 1, 0.0, 1.3, 60)
        c238 = np.concatenate(([0], np.cumsum([s238.counts[i] for i in range(max(s238.structure)+1)])))
        s1968 = computeBounds(1968, geo, 1, 0.0, 1.3, 60)
        c1968 = np.concatenate(([0], np.cumsum([s1968.counts[i] for i in range(max(s1968.structure)+1)])))

        for i in range(20):
            s += 'CG {:2d} & {:5s} & {:7s} & {:7s} & {:9s} & {:9s} & {:9s} & {:9s} \\\\\n'.format(i+1, f(c44[i:i+2]), f(c238[i:i+2]), f(c238[i+20:i+22]), f(c1968[i+0:i+2]), f(c1968[i+20:i+22]), f(c1968[i+40:i+42]), f(c1968[i+60:i+62]))
        print(s)

        #pickle.dump(structure, open('XS/structure{}_{}_{}g_m{}_r{}_g{}.p'.format(contig, geo, G, min_cutoff, ratio_cutoff, group_cutoff), 'wb'))


#!/usr/bin/env python -u
'''
pDMET: Density Matrix Embedding theory for Periodic Systems
Copyright (C) 2018 Hung Q. Pham. All Rights Reserved.
A few functions in pDMET are modifed from QC-DMET Copyright (C) 2015 Sebastian Wouters

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Email: Hung Q. Pham <pqh3.14@gmail.com>
'''

import numpy as np
import scipy.optimize as optimize

'''DFT XC library'''
GGA_X = {
'GGA_X_GAM'                    : 32 , # H. S. Yu, W. Zhang, P. Verma, X. He, and D. G. Truhlar, Phys. Chem. Chem. Phys. 17, 12146 (2015)
'GGA_X_HCTH_A'                 : 34 , # F. A. Hamprecht, A. J. Cohen, D. J. Tozer, and N. C. Handy, J. Chem. Phys. 109, 6264 (1998)
'GGA_X_EV93'                   : 35 , # E. Engel and S. H. Vosko, Phys. Rev. B 47, 13164 (1993)
'GGA_X_BCGP'                   : 38 , # K. Burke, A. Cancio, T. Gould, and S. Pittalis, ArXiv e-prints (2014), arXiv:1409.4834 [cond-mat.mtrl-sci]
'GGA_X_LAMBDA_OC2_N'           : 40 , # M. M. Odashima, K. Capelle, and S. B. Trickey, J. Chem. Theory Comput. 5, 798 (2009)
'GGA_X_B86_R'                  : 41 , # I. Hamada, Phys. Rev. B 89, 121103 (2014)
'GGA_X_LAMBDA_CH_N'            : 44 , # M. M. Odashima, K. Capelle, and S. B. Trickey, J. Chem. Theory Comput. 5, 798 (2009)
'GGA_X_LAMBDA_LO_N'            : 45 , # M. M. Odashima, K. Capelle, and S. B. Trickey, J. Chem. Theory Comput. 5, 798 (2009)
'GGA_X_HJS_B88_V2'             : 46 , # E. Weintraub, T. M. Henderson, and G. E. Scuseria, J. Chem. Theory Comput. 5, 754 (2009)
'GGA_X_Q2D'                    : 48 , # L. Chiodo, L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. Lett. 108, 126402 (2012)
'GGA_X_PBE_MOL'                : 49 , # J. M. del Campo, J. L. G\'azquez, S. B. Trickey, and A. Vela, J. Chem. Phys. 136, 104108 (2012)
'GGA_X_AK13'                   : 56 , # R. Armiento and S. Kummel, Phys. Rev. Lett. 111, 036402 (2013)
'GGA_X_LV_RPW86'               : 58 , # K. Berland and P. Hyldgaard, Phys. Rev. B 89, 035412 (2014)
'GGA_X_PBE_TCA'                : 59 , # V. Tognetti, P. Cortona, and C. Adamo, Chem. Phys. Lett. 460, 536 (2008)
'GGA_X_PBEINT'                 : 60 , # E. Fabiano, L. A. Constantin, and F. Della Sala, Phys. Rev. B 82, 113104 (2010)
'GGA_X_VMT84_GE'               : 68 , # A. Vela, J. C. Pacheco-Kato, J. L. G\'azquez, J. M. del Campo, and S. B. Trickey, J. Chem. Phys. 136, 144115 (2012)
'GGA_X_VMT84_PBE'              : 69 , # A. Vela, J. C. Pacheco-Kato, J. L. G\'azquez, J. M. del Campo, and S. B. Trickey, J. Chem. Phys. 136, 144115 (2012)
'GGA_X_VMT_GE'                 : 70 , # A. Vela, V. Medel, and S. B. Trickey, J. Chem. Phys. 130, 244103 (2009)
'GGA_X_VMT_PBE'                : 71 , # A. Vela, V. Medel, and S. B. Trickey, J. Chem. Phys. 130, 244103 (2009)
'GGA_X_N12'                    : 82 , # R. Peverati and D. G. Truhlar, J. Chem. Theory Comput. 8, 2310 (2012)
'GGA_X_SSB_SW'                 : 90 , # M. Swart, M. Sol\'a, and F. M. Bickelhaupt, J. Comput. Methods Sci. Eng. 9, 69 (2009)
'GGA_X_SSB'                    : 91 , # M. Swart, M. Sol\'a, and F. M. Bickelhaupt, J. Chem. Phys. 131, 094103 (2009)
'GGA_X_SSB_D'                  : 92 , # M. Swart, M. Sol\'a, and F. M. Bickelhaupt, J. Chem. Phys. 131, 094103 (2009)
'GGA_X_BPCCAC'                 : 98 , # E. Br\'emond, D. Pilard, I. Ciofini, H. Chermette, C. Adamo, and P. Cortona, Theor. Chem. Acc. 131, 1184 (2012)
'GGA_X_PBE'                    : 101, # J. P. Perdew, K. Burke, and M. Ernzerhof, Phys. Rev. Lett. 77, 3865 (1996)
'GGA_X_PBE_R'                  : 102, # Y. Zhang and W. Yang, Phys. Rev. Lett. 80, 890 (1998)
'GGA_X_B86'                    : 103, # A. D. Becke, J. Chem. Phys. 84, 4524 (1986)
'GGA_X_HERMAN'                 : 104, # F. Herman, J. P. V. Dyke, and I. B. Ortenburger, Phys. Rev. Lett. 22, 807 (1969)
'GGA_X_B86_MGC'                : 105, # A. D. Becke, J. Chem. Phys. 84, 4524 (1986)
'GGA_X_B88'                    : 106, # A. D. Becke, Phys. Rev. A 38, 3098 (1988)
'GGA_X_G96'                    : 107, # P. M. W. Gill, Mol. Phys. 89, 433 (1996)
'GGA_X_PW86'                   : 108, # J. P. Perdew and W. Yue, Phys. Rev. B 33, 8800 (1986)
'GGA_X_PW91'                   : 109, # J. P. Perdew, in Proceedings of the 75. WE-Heraeus-Seminar and 21st Annual International Symposium on Electronic Structure of Solids, edited by P. Ziesche and H. Eschrig (Akademie Verlag, Berlin, 1991) p. 11
'GGA_X_OPTX'                   : 110, # N. C. Handy and A. J. Cohen, Mol. Phys. 99, 403 (2001)
'GGA_X_DK87_R1'                : 111, # A. E. DePristo and J. D. Kress, J. Chem. Phys. 86, 1425 (1987)
'GGA_X_DK87_R2'                : 112, # A. E. DePristo and J. D. Kress, J. Chem. Phys. 86, 1425 (1987)
'GGA_X_LG93'                   : 113, # D. J. Lacks and R. G. Gordon, Phys. Rev. A 47, 4681 (1993)
'GGA_X_FT97_A'                 : 114, # M. Filatov and W. Thiel, Mol. Phys. 91, 847 (1997)
'GGA_X_FT97_B'                 : 115, # M. Filatov and W. Thiel, Mol. Phys. 91, 847 (1997)
'GGA_X_PBE_SOL'                : 116, # J. P. Perdew, A. Ruzsinszky, G. I. Csonka, O. A. Vydrov, G. E. Scuseria, L. A. Constantin, X. Zhou, and K. Burke, Phys. Rev. Lett. 100, 136406 (2008)
'GGA_X_RPBE'                   : 117, # B. Hammer, L. B. Hansen, and J. K. Norskov, Phys. Rev. B 59, 7413 (1999)
'GGA_X_WC'                     : 118, # Z. Wu and R. E. Cohen, Phys. Rev. B 73, 235116 (2006)
'GGA_X_MPW91'                  : 119, # C. Adamo and V. Barone, J. Chem. Phys. 108, 664 (1998)
'GGA_X_AM05'                   : 120, # R. Armiento and A. E. Mattsson, Phys. Rev. B 72, 085108 (2005)
'GGA_X_PBEA'                   : 121, # G. K. H. Madsen, Phys. Rev. B 75, 195108 (2007)
'GGA_X_MPBE'                   : 122, # C. Adamo and V. Barone, J. Chem. Phys. 116, 5933 (2002)
'GGA_X_XPBE'                   : 123, # X. Xu and W. A. Goddard, J. Chem. Phys. 121, 4068 (2004)
'GGA_X_2D_B86_MGC'             : 124, # S. Pittalis, E. Rasanen, J. G. Vilhena, and M. A. L. Marques, Phys. Rev. A 79, 012503 (2009)
'GGA_X_BAYESIAN'               : 125, # J. J. Mortensen, K. Kaasbjerg, S. L. Frederiksen, J. K. Norskov, J. P. Sethna, and K. W. Jacobsen, Phys. Rev. Lett. 95, 216401 (2005)
'GGA_X_PBE_JSJR'               : 126, # L. S. Pedroza, A. J. R. da Silva, and K. Capelle, Phys. Rev. B 79, 201106 (2009)
'GGA_X_2D_B88'                 : 127, # J. G. Vilhena and M. A. L. Marques, unpublished (2014)
'GGA_X_2D_B86'                 : 128, # J. G. Vilhena and M. A. L. Marques, unpublished (2014)
'GGA_X_2D_PBE'                 : 129, # J. G. Vilhena and M. A. L. Marques, unpublished (2014)
'GGA_X_OPTB88_VDW'             : 139, # J. Klime\v{s}, D. R. Bowler, and A. Michaelides, J. Phys.: Condens. Matter 22, 022201 (2010)
'GGA_X_PBEK1_VDW'              : 140, # J. Klime\v{s}, D. R. Bowler, and A. Michaelides, J. Phys.: Condens. Matter 22, 022201 (2010)
'GGA_X_OPTPBE_VDW'             : 141, # J. Klime\v{s}, D. R. Bowler, and A. Michaelides, J. Phys.: Condens. Matter 22, 022201 (2010)
'GGA_X_RGE2'                   : 142, # A. Ruzsinszky, G. I. Csonka, and G. E. Scuseria, J. Chem. Theory Comput. 5, 763 (2009)
'GGA_X_RPW86'                  : 144, # E. D. Murray, K. Lee, and D. C. Langreth, J. Chem. Theory Comput. 5, 2754 (2009)
'GGA_X_KT1'                    : 145, # T. W. Keal and D. J. Tozer, J. Chem. Phys. 119, 3015 (2003)
'GGA_X_MB88'                   : 149, # V. Tognetti and C. Adamo, J. Phys. Chem. A 113, 14415 (2009)
'GGA_X_SOGGA'                  : 150, # Y. Zhao and D. G. Truhlar, J. Chem. Phys. 128, 184109 (2008)
'GGA_X_SOGGA11'                : 151, # R. Peverati, Y. Zhao, and D. G. Truhlar, J. Phys. Chem. Lett. 2, 1991 (2011)
'GGA_X_C09X'                   : 158, # V. R. Cooper, Phys. Rev. B 81, 161104 (2010)
'GGA_X_OL2'                    : 183, # P. Fuentealba and O. Reyes, Chem. Phys. Lett. 232, 31 (1995)
'GGA_X_APBE'                   : 184, # L. A. Constantin, E. Fabiano, S. Laricchia, and F. Della Sala, Phys. Rev. Lett. 106, 186406 (2011)
'GGA_X_HTBS'                   : 191, # P. Haas, F. Tran, P. Blaha, and K. Schwarz, Phys. Rev. B 83, 205117 (2011)
'GGA_X_AIRY'                   : 192, # L. A. Constantin, A. Ruzsinszky, and J. P. Perdew, Phys. Rev. B 80, 035125 (2009)
'GGA_X_LAG'                    : 193, # L. Vitos, B. Johansson, J. Koll\'ar, and H. L. Skriver, Phys. Rev. B 62, 10046 (2000)
'GGA_X_PBEFE'                  : 265, # R. Sarmiento-P\'erez, S. Botti, and M. A. L. Marques, J. Chem. Theory Comput. 11, 3844 (2015)
'GGA_X_CAP'                    : 270, # J. Carmona-Esp\'indola, J. L. G\'azquez, A. Vela, and S. B. Trickey, J. Chem. Phys. 142, 054105 (2015), 10.1063/1.4906606
'GGA_X_EB88'                   : 271, # P. Elliott and K. Burke, Can. J. Chem. 87, 1485 (2009)
'GGA_X_BEEFVDW'                : 285, # J. Wellendorff, K. T. Lundgaard, A. M\o{gelh\o{j}, V. Petzold, D. D. Landis, J. K. N\o{rskov}, T. Bligaard, and K. W. Jacobsen, }Phys. Rev. B 85, 235149 (2012)
'GGA_X_PBETRANS'               : 291, # Eric Bremond, I. Ciofini, and C. Adamo, Mol. Phys. 114, 1059 (2016)
'GGA_X_CHACHIYO'               : 298, # T. {Chachiyo and H. {Chachiyo}, }ArXiv e-prints (2017), arXiv:1706.01343 [cond-mat.mtrl-sci]
'GGA_X_WPBEH'                  : 524, # J. Heyd, G. E. Scuseria, and M. Ernzerhof, J. Chem. Phys. 118, 8207 (2003)
'GGA_X_HJS_PBE'                : 525, # T. M. Henderson, B. G. Janesko, and G. E. Scuseria, J. Chem. Phys. 128, 194105 (2008)
'GGA_X_HJS_PBE_SOL'            : 526, # T. M. Henderson, B. G. Janesko, and G. E. Scuseria, J. Chem. Phys. 128, 194105 (2008)
'GGA_X_HJS_B97X'               : 528, # T. M. Henderson, B. G. Janesko, and G. E. Scuseria, J. Chem. Phys. 128, 194105 (2008)
'GGA_X_ITYH'                   : 529, # H. Iikura, T. Tsuneda, T. Yanai, and K. Hirao, J. Chem. Phys. 115, 3540 (2001)
'GGA_X_SFAT'                   : 530, # A. Savin and H.-J. Flad, Int. J. Quantum Chem. 56, 327 (1995)
'GGA_X_SG4'                    : 533, # L. A. Constantin, A. Terentjevs, F. Della Sala, P. Cortona, and E. Fabiano, Phys. Rev. B 93, 045126 (2016)
'GGA_X_GG99'                   : 535, # A. T. Gilbert and P. M. Gill, Chem. Phys. Lett. 312, 511 (1999)
'GGA_X_PBEPOW'                 : 539, # Eric Bremond, J. Chem. Phys. 145, 244102 (2016)
'GGA_X_KGG99'                  : 544, # A. T. Gilbert and P. M. Gill, Chem. Phys. Lett. 312, 511 (1999)
'GGA_X_B88M'                   : 570, # E. Proynov, H. Chermette, and D. R. Salahub, J. Chem. Phys. 113, 10013 (2000)
}

GGA_C = {
'GGA_C_GAM'                    : 33 , # H. S. Yu, W. Zhang, P. Verma, X. He, and D. G. Truhlar, Phys. Chem. Chem. Phys. 17, 12146 (2015)
'GGA_C_BCGP'                   : 39 , # K. Burke, A. Cancio, T. Gould, and S. Pittalis, ArXiv e-prints (2014), arXiv:1409.4834 [cond-mat.mtrl-sci]
'GGA_C_Q2D'                    : 47 , # L. Chiodo, L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. Lett. 108, 126402 (2012)
'GGA_C_ZPBEINT'                : 61 , # L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. B 84, 233103 (2011)
'GGA_C_PBEINT'                 : 62 , # E. Fabiano, L. A. Constantin, and F. Della Sala, Phys. Rev. B 82, 113104 (2010)
'GGA_C_ZPBESOL'                : 63 , # L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. B 84, 233103 (2011)
'GGA_C_N12_SX'                 : 79 , # R. Peverati and D. G. Truhlar, Phys. Chem. Chem. Phys. 14, 16187 (2012)
'GGA_C_N12'                    : 80 , # R. Peverati and D. G. Truhlar, J. Chem. Theory Comput. 8, 2310 (2012)
'GGA_C_REGTPSS'                : 83 , # J. P. Perdew, A. Ruzsinszky, G. I. Csonka, L. A. Constantin, and J. Sun, Phys. Rev. Lett. 103, 026403 (2009)
'GGA_C_OP_XALPHA'              : 84 , # T. Tsuneda, T. Suzumura, and K. Hirao, J. Chem. Phys. 110, 10664 (1999)
'GGA_C_OP_G96'                 : 85 , # T. Tsuneda, T. Suzumura, and K. Hirao, J. Chem. Phys. 110, 10664 (1999)
'GGA_C_OP_PBE'                 : 86 , # T. Tsuneda, T. Suzumura, and K. Hirao, J. Chem. Phys. 110, 10664 (1999)
'GGA_C_OP_B88'                 : 87 , # T. Tsuneda, T. Suzumura, and K. Hirao, J. Chem. Phys. 110, 10664 (1999)
'GGA_C_FT97'                   : 88 , # M. Filatov and W. Thiel, Int. J. Quantum Chem. 62, 603 (1997)
'GGA_C_HCTH_A'                 : 97 , # F. A. Hamprecht, A. J. Cohen, D. J. Tozer, and N. C. Handy, J. Chem. Phys. 109, 6264 (1998)
'GGA_C_REVTCA'                 : 99 , # V. Tognetti, P. Cortona, and C. Adamo, Chem. Phys. Lett. 460, 536 (2008)
'GGA_C_TCA'                    : 100, # V. Tognetti, P. Cortona, and C. Adamo, J. Chem. Phys. 128, 034101 (2008)
'GGA_C_PBE'                    : 130, # J. P. Perdew, K. Burke, and M. Ernzerhof, Phys. Rev. Lett. 77, 3865 (1996)
'GGA_C_LYP'                    : 131, # C. Lee, W. Yang, and R. G. Parr, Phys. Rev. B 37, 785 (1988)
'GGA_C_P86'                    : 132, # J. P. Perdew, Phys. Rev. B 33, 8822 (1986)
'GGA_C_PBE_SOL'                : 133, # J. P. Perdew, A. Ruzsinszky, G. I. Csonka, O. A. Vydrov, G. E. Scuseria, L. A. Constantin, X. Zhou, and K. Burke, Phys. Rev. Lett. 100, 136406 (2008)
'GGA_C_PW91'                   : 134, # J. P. Perdew, in Proceedings of the 75. WE-Heraeus-Seminar and 21st Annual International Symposium on Electronic Structure of Solids, edited by P. Ziesche and H. Eschrig (Akademie Verlag, Berlin, 1991) p. 11
'GGA_C_AM05'                   : 135, # R. Armiento and A. E. Mattsson, Phys. Rev. B 72, 085108 (2005)
'GGA_C_XPBE'                   : 136, # X. Xu and W. A. Goddard, J. Chem. Phys. 121, 4068 (2004)
'GGA_C_LM'                     : 137, # D. C. Langreth and M. J. Mehl, Phys. Rev. Lett. 47, 446 (1981)
'GGA_C_PBE_JRGX'               : 138, # L. S. Pedroza, A. J. R. da Silva, and K. Capelle, Phys. Rev. B 79, 201106 (2009)
'GGA_C_RGE2'                   : 143, # A. Ruzsinszky, G. I. Csonka, and G. E. Scuseria, J. Chem. Theory Comput. 5, 763 (2009)
'GGA_C_WL'                     : 147, # L. C. Wilson and M. Levy, Phys. Rev. B 41, 12930 (1990)
'GGA_C_WI'                     : 148, # L. C. Wilson and S. Ivanov, Int. J. Quantum Chem. 69, 523 (1998)
'GGA_C_SOGGA11'                : 152, # R. Peverati, Y. Zhao, and D. G. Truhlar, J. Phys. Chem. Lett. 2, 1991 (2011)
'GGA_C_WI0'                    : 153, # L. C. Wilson and S. Ivanov, Int. J. Quantum Chem. 69, 523 (1998)
'GGA_C_SOGGA11_X'              : 159, # R. Peverati and D. G. Truhlar, J. Chem. Phys. 135, 191102 (2011)
'GGA_C_APBE'                   : 186, # L. A. Constantin, E. Fabiano, S. Laricchia, and F. Della Sala, Phys. Rev. Lett. 106, 186406 (2011)
'GGA_C_OPTC'                   : 200, # A. J. Cohen and N. C. Handy, Mol. Phys. 99, 607 (2001)
'GGA_C_PBELOC'                 : 246, # L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. B 86, 035130 (2012)
'GGA_C_PBEFE'                  : 258, # R. Sarmiento-P\'erez, S. Botti, and M. A. L. Marques, J. Chem. Theory Comput. 11, 3844 (2015)
'GGA_C_OP_PW91'                : 262, # T. Tsuneda, T. Suzumura, and K. Hirao, J. Chem. Phys. 110, 10664 (1999)
'GGA_C_PBE_MOL'                : 272, # J. M. del Campo, J. L. G\'azquez, S. B. Trickey, and A. Vela, J. Chem. Phys. 136, 104108 (2012)
'GGA_C_BMK'                    : 280, # A. D. Boese and J. M. L. Martin, J. Chem. Phys. 121, 3405 (2004)
'GGA_C_TAU_HCTH'               : 281, # A. D. Boese and N. C. Handy, J. Chem. Phys. 116, 9559 (2002)
'GGA_C_HYB_TAU_HCTH'           : 283, # A. D. Boese and N. C. Handy, J. Chem. Phys. 116, 9559 (2002)
'GGA_C_SG4'                    : 534, # L. A. Constantin, A. Terentjevs, F. Della Sala, P. Cortona, and E. Fabiano, Phys. Rev. B 93, 045126 (2016)
'GGA_C_SCAN_E0'                : 553, # J. Sun, A. Ruzsinszky, and J. P. Perdew, Phys. Rev. Lett. 115, 036402 (2015)
'GGA_C_GAPC'                   : 555, # E. Fabiano, P. E. Trevisanutto, A. Terentjevs, and L. A. Constantin, J. Chem. Theory Comput. 10, 2016 (2014), pMID: 26580528
'GGA_C_GAPLOC'                 : 556, # E. Fabiano, P. E. Trevisanutto, A. Terentjevs, and L. A. Constantin, J. Chem. Theory Comput. 10, 2016 (2014), pMID: 26580528
'GGA_C_ZVPBEINT'               : 557, # L. A. Constantin, E. Fabiano, and F. D. Sala, J. Chem. Phys. 137, 194105 (2012)
'GGA_C_ZVPBESOL'               : 558, # L. A. Constantin, E. Fabiano, and F. D. Sala, J. Chem. Phys. 137, 194105 (2012)
'GGA_C_TM_LYP'                 : 559, # A. J. Thakkar and S. P. McCarthy, J. Chem. Phys. 131, 134109 (2009)
'GGA_C_TM_PBE'                 : 560, # A. J. Thakkar and S. P. McCarthy, J. Chem. Phys. 131, 134109 (2009)
'GGA_C_W94'                    : 561, # L. C. Wilson, Chemical Physics 181, 337 (1994)
'GGA_C_CS1'                    : 565, # N. C. Handy a
}

'''Smaller set of XC to test'''
GGA_X = {

'GGA_X_PBE'                    : 101, # J. P. Perdew, K. Burke, and M. Ernzerhof, Phys. Rev. Lett. 77, 3865 (1996)
'GGA_X_B88'                    : 106, # A. D. Becke, Phys. Rev. A 38, 3098 (1988)
'GGA_X_PW91'                   : 109, # J. P. Perdew, in Proceedings of the 75. WE-Heraeus-Seminar and 21st Annual International Symposium on Electronic Structure of Solids, edited by P. Ziesche and H. Eschrig (Akademie Verlag, Berlin, 1991) p. 11
'GGA_X_WC'                     : 118, # Z. Wu and R. E. Cohen, Phys. Rev. B 73, 235116 (2006)
'GGA_X_AM05'                   : 120, # R. Armiento and A. E. Mattsson, Phys. Rev. B 72, 085108 (2005)
}

GGA_C = {
'GGA_C_PBE'                    : 130, # J. P. Perdew, K. Burke, and M. Ernzerhof, Phys. Rev. Lett. 77, 3865 (1996)
'GGA_C_LYP'                    : 131, # C. Lee, W. Yang, and R. G. Parr, Phys. Rev. B 37, 785 (1988)
'GGA_C_P86'                    : 132, # J. P. Perdew, Phys. Rev. B 33, 8822 (1986)
'GGA_C_AM05'                   : 135, # R. Armiento and A. E. Mattsson, Phys. Rev. B 72, 085108 (2005)
'GGA_C_SCAN_E0'                : 553, # J. Sun, A. Ruzsinszky, and J. P. Perdew, Phys. Rev. Lett. 115, 036402 (2015)
}

def get_init_uvec(xc_type='PBE0', dft_HF=None):
    '''Initialize uvec corresponding to each type of the DF-like cost function'''
    if dft_HF is not None:
        HF_percent = dft_HF
    else:
        HF_percent = 1.0

    if xc_type == 'PBE0':
        uvec = [HF_percent]  + [0.] * 2
    elif xc_type == 'RSH-PBE0':
        uvec = [1.] + [1.]  + [0.] * 2
    elif xc_type == 'B3LYP':
        uvec = [HF_percent]  + [0.] * 4
    elif xc_type == 'CAMB3LYP':
        uvec = [HF_percent , .0000001, 1. , 0., 0., 0., 0.]
    elif xc_type == 'manyGGA':
        uvec = [HF_percent]  + [0.] * (len(GGA_X) + len(GGA_C))

    return uvec

def get_bounds(xc_type='PBE0', constraint=1, dft_HF=None):
    '''Prodive the bounds to optimize the DF-like cost function'''

    if constraint == 1:
        if xc_type == 'PBE0':
            lower = [-1.0] + [.0] * 2
            upper = [ 1.0] + [1.0] * 2
            bounds = optimize.Bounds(lower, upper)
        elif xc_type == 'RSH-PBE0':
            lower = [ -1.] * 2 + [.0] * 2
            upper = [ 1.0] * 4
            bounds = optimize.Bounds(lower, upper)
        elif xc_type == 'B3LYP':
            lower = [-1.0] + [.0] * 4
            upper = [ 1.0] + [1.0] * 4
            bounds = optimize.Bounds(lower, upper)
        elif xc_type == 'CAMB3LYP':
            bounds = optimize.Bounds([0.0,.0000001,.0,.0,.0,.0,.0], [1.,1.0,1.,1.,1.,1.,1.])
        elif xc_type == 'manyGGA':
            lower = [-1.0] + [.0] * (len(GGA_X) + len(GGA_C))
            upper = [ 1.0] + [.1] * (len(GGA_X) + len(GGA_C))
            bounds = optimize.Bounds(lower, upper)
    elif constraint == 2:
        if xc_type == 'PBE0':
            lower = [-1.0] * 3
            upper = [ 1.0] * 3
            bounds = optimize.Bounds(lower, upper)
        elif xc_type == 'RSH-PBE0':
            lower = [-1] * 4
            upper = [ 1.0]* 4
        elif xc_type == 'B3LYP':
            lower = [-1.0] * 5
            upper = [ 1.0] * 5
            bounds = optimize.Bounds(lower, upper)
        elif xc_type == 'CAMB3LYP':
            bounds = optimize.Bounds([0.0,.0000001,.0,.0,.0,.0,.0], [1.,1.0,1.,1.,1.,1.,1.])
        elif xc_type == 'manyGGA':
            lower = [-1.0] + [-1.0] * (len(GGA_X) + len(GGA_C))
            upper = [ 1.0] + [ 1.0] * (len(GGA_X) + len(GGA_C))
            bounds = optimize.Bounds(lower, upper)
    elif constraint == 3:
        if xc_type == 'manyGGA':
            lower = [ 0.0] * (len(GGA_X) + len(GGA_C))
            upper = [ 1.0] * (len(GGA_X) + len(GGA_C))
            bounds = optimize.Bounds(lower, upper)
    return bounds

def get_OEH_kpts(local, umat, xc_type='PBE0', dft_HF=None):
    '''Construct the mf Hamiltonian'''

    if dft_HF is not None and xc_type is not 'RSH-PBE0':
        umat[0] = np.float64(dft_HF)

    if xc_type == 'PBE0':
        xc =  "{:.12f}*HF + {:.12f}*PBE, {:.12f}*PBE".format(*umat)
        n, exc, vxc = local.kks._numint.nr_rks(local.cell, local.kks.grids, xc, local.dm_kpts, 0, local.kpts, None)
        veff = vxc + local.vj - umat[0] * 0.5 * local.vk
        OEH_kpts = local.h_core + veff
        if local._is_KROHF:
            OEH_kpts = 0.5 * (OEH_kpts[0] + OEH_kpts[1])
        OEH_kpts = local.ao_2_loc(OEH_kpts, local.ao2lo)

    elif xc_type == 'B3LYP':
        xc =  "{:.12f}*HF + {:.12f}*LDA + {:.12f}*B88, {:.12f}*LYP + {:.12f}*VWN".format(*umat)
        n, exc, vxc = local.kks._numint.nr_rks(local.cell, local.kks.grids, xc, local.dm_kpts, 0, local.kpts, None)
        veff = vxc + local.vj - umat[0] * 0.5 * local.vk
        OEH_kpts = local.h_core + veff
        if local._is_KROHF:
            OEH_kpts = 0.5 * (OEH_kpts[0] + OEH_kpts[1])
        OEH_kpts = local.ao_2_loc(OEH_kpts, local.ao2lo)

    elif xc_type == 'RSH-PBE0':
        # TODO: right now the range-seperation parameter is fixed, it is super expensive to update it every cycle
        # because the vk long-range needed to be generated
        umat = [local.xc_omega] + np.asarray(umat).tolist()
        xc =  "{1:.12f}*SR_HF({0:.2f})+ {2:.12f}*LR_HF({0:.2f}) + {3:.12f}*PBE, {4:.12f}*PBE".format(*umat)
        n, exc, vxc = local.kks._numint.nr_rks(local.cell, local.kks.grids, xc, local.dm_kpts, 0, local.kpts, None)
        veff = vxc + local.vj - 0.5 * (umat[1]*local.vksr + umat[2]*local.vklr)
        OEH_kpts = local.h_core + veff
        OEH_kpts = local.ao_2_loc(OEH_kpts, local.ao2lo)

    elif xc_type == 'CAMB3LYP':
        # TODO: need to debug
        xc =  '{:.12f}*SR_HF({:.12f}) + {:.12f}*LR_HF({:.12f}) + {:.12f}*ITYH + {:.12f}*B88, {:.12f}*VWN5 + {:.12f}*LYP'.format(umat)
        n, exc, vxc = local.kks._numint.nr_rks(local.cell, local.kks.grids, xc, local.dm_kpts, 0, local.kpts, None)
        veff = vxc + local.vj - umat[0] * 0.5 * local.vk
        OEH_kpts = local.h_core + veff
        OEH_kpts = local.ao_2_loc(OEH_kpts, local.ao2lo)

    elif xc_type == 'manyGGA':
        x = '{:.12f}*HF + {:.12f}*' + '+ {:.12f}*'.join(GGA_X.keys())
        c = ', {:.12f}*' + '+ {:.12f}*'.join(GGA_C.keys())
        xc = (x + c).format(*umat)
        n, exc, vxc = local.kks._numint.nr_rks(local.cell, local.kks.grids, xc, local.dm_kpts, 0, local.kpts, None)
        veff = vxc + local.vj - umat[0] * 0.5 * local.vk
        OEH_kpts = local.h_core + veff
        OEH_kpts = local.ao_2_loc(OEH_kpts, local.ao2lo)

    return OEH_kpts