# ======================================================================
# Copyright TOTAL / CERFACS / LIRMM (02/2020)
# Contributor: Adrien Suau (<adrien.suau@cerfacs.fr>
#                           <adrien.suau@lirmm.fr>)
#              Siyuan Niu  (<siyuan.niu@lirmm.fr>)
# This software is governed by the CeCILL-B license under French law and
# abiding  by the  rules of  distribution of free software. You can use,
# modify  and/or  redistribute  the  software  under  the  terms  of the
# CeCILL-B license as circulated by CEA, CNRS and INRIA at the following
# URL "http://www.cecill.info".
#
# As a counterpart to the access to  the source code and rights to copy,
# modify and  redistribute granted  by the  license, users  are provided
# only with a limited warranty and  the software's author, the holder of
# the economic rights,  and the  successive licensors  have only limited
# liability.
#
# In this respect, the user's attention is drawn to the risks associated
# with loading,  using, modifying and/or  developing or reproducing  the
# software by the user in light of its specific status of free software,
# that  may mean  that it  is complicated  to manipulate,  and that also
# therefore  means that  it is reserved for  developers and  experienced
# professionals having in-depth  computer knowledge. Users are therefore
# encouraged  to load and  test  the software's  suitability as  regards
# their  requirements  in  conditions  enabling  the  security  of their
# systems  and/or  data to be  ensured and,  more generally,  to use and
# operate it in the same conditions as regards security.
#
# The fact that you  are presently reading this  means that you have had
# knowledge of the CeCILL-B license and that you accept its terms.
# ======================================================================

import logging
import typing as ty

import numpy
from qiskit.circuit.quantumregister import Qubit
from qiskit.dagcircuit.dagcircuit import DAGNode

from hamap.gates import TwoQubitGate, SwapTwoQubitGate
from hamap.hardware.IBMQHardwareArchitecture import IBMQHardwareArchitecture
from hamap.layer import QuantumLayer, update_layer

logger = logging.getLogger("hamap.heuristics")

def _gate_op_cost(
    op: DAGNode,
    distance_matrix: numpy.ndarray,
    mapping: ty.Dict[Qubit, int],
    hardware: IBMQHardwareArchitecture,
) -> float:
    if hardware.is_ignored_operation(op):
        return 0
    if len(op.qargs) == 1:
        # SABRE ignores 1-qubit gates
        return 0
    elif len(op.qargs) == 2:
        # This is a CNOT
        source, sink = op.qargs
        return distance_matrix[mapping[source], mapping[sink]]
    else:
        logger.warning(
            f"Found a quantum operation applied on '{len(op.qargs)}' qubits. This "
            f"operation will be excluded from the cost computation."
        )
        return 0

def sabre_heuristic(
    hardware: IBMQHardwareArchitecture,
    front_layer: QuantumLayer,
    topological_nodes: ty.List[DAGNode],
    current_node_index: int,
    current_mapping: ty.Dict[Qubit, int],
    initial_mapping: ty.Dict[Qubit, int],
    trans_mapping: ty.Dict[Qubit, int],
    distance_matrix: numpy.ndarray,
    tentative_gate: TwoQubitGate,
    look_ahead_depth: int = 20,
    look_ahead_weight: float = 0.5,
) -> float:
    """The heuristic cost function used in the SABRE optimiser.

    :param hardware: the SABRE optimiser does not take into account the hardware data to
        compute the heuristic cost, only to generate the possible SWAPs to evaluate
        with this heuristic. The SABRE heuristic only uses the distance matrix.
        Nevertheless, this implementation uses the hardware data to check if some
        gates are ignored (such as barriers for example).
    :param front_layer: the current front layer. Used to compute an "immediate" cost,
        i.e. a quantity that will tell us if the SWAP/Bridge is useful to execute
        gates in the front layer.
    :param topological_nodes: the list of all the DAGNodes of the quantum circuit,
        sorted in topological order.
    :param current_node_index: index of the first non-processed node.
    :param current_mapping: the mapping *before* applying the given SWAP.
    :param distance_matrix: the pre-computed distance matrix between each qubits.
    :param tentative_gate: the SWAP we want to estimate the usefulness of.
    :param look_ahead_depth: the depth of the look-ahead. The procedure will consider
        gates that will be executed in the future (i.e. not in the front layer) up to
        the given depth. Note that 1-qubit gates are not ignored, which means that a
        depth of 3 will not guarantee that there is at least 3 CNOTs in the
        look-ahead set.
    :param look_ahead_weight: weight of the look-ahead. The actual gates (i.e. the
        gates in the front layer) have a weight of 1.
    :return: the heuristic cost of the given SWAP/Bridge according to the current
        state of the algorithm.
    """
    # First, compute the proposed new mapping
    new_mapping = tentative_gate.update_mapping(current_mapping)
    # Compute H_basic, the cost associated with the distance.
    H_basic = 0.0
    H_tentative = 0.0
    H_tentative_gate_number = 0
    H_basic_gate_number = 0
    for op in front_layer.ops:
        # Only add the gate to the cost if the gate is not already implemented by the
        # SWAP/Bridge
        if not tentative_gate.implements_operation(op, initial_mapping, trans_mapping):
        #if isinstance(tentative_gate, SwapTwoQubitGate):
            H_basic += _gate_op_cost(op, distance_matrix, new_mapping, hardware)
            H_basic_gate_number += 1
    # Compute H, the cost cost that encourage parallelism and adds some look-ahead
    # ability.

    if isinstance(tentative_gate, SwapTwoQubitGate):
        H_tentative += tentative_gate.cost(hardware, current_mapping, distance_matrix)
        H_tentative_gate_number += 3
    else:
        H_tentative += tentative_gate.cost(hardware, initial_mapping, distance_matrix)
        H_tentative_gate_number += 4

    #H = 0.0
    future_nodes_layer = QuantumLayer(max_depth=look_ahead_depth)
    # We do not use the return of update_layer because we do not care about the
    # number of gates that were added. Still, we add the firsts look_ahead_depth layers
    # of our future gates in this set to have this look-ahead ability.
    _ = update_layer(future_nodes_layer, topological_nodes, current_node_index)
    # The decay is not implemented in the code the authors gave us and not
    # sufficiently explained in the paper to implement it without guessing. Not
    # implementing it for the moment...
    #H += (H_basic / H_basic_gate_number) if H_basic_gate_number != 0 else 0
    H = (H_basic + H_tentative) / (H_basic_gate_number + H_tentative_gate_number)
    #print(type(tentative_gate), tentative_gate.left, tentative_gate.right, H, H_basic, H_basic_gate_number)
    #print(front_layer.ops[0].qargs)
    H_extended = 0.0
    if future_nodes_layer:
        # Only add this cost if there are nodes in the future_node_layer
        H_extended += (
            look_ahead_weight
            * sum(
                _gate_op_cost(op, distance_matrix, new_mapping, hardware)
                for op in future_nodes_layer.ops
            )
            / len(future_nodes_layer)
        )

    H += H_extended
    #print(f"H extended {H_extended} and final H is {H}, gate number {len(future_nodes_layer)}")
    return H

def sabre_heuristic_with_effect(
    hardware: IBMQHardwareArchitecture,
    front_layer: QuantumLayer,
    topological_nodes: ty.List[DAGNode],
    current_node_index: int,
    current_mapping: ty.Dict[Qubit, int],
    initial_mapping: ty.Dict[Qubit, int],
    trans_mapping: ty.Dict[Qubit, int],
    distance_matrix: numpy.ndarray,
    tentative_gate: SwapTwoQubitGate,
    look_ahead_depth: int = 20,
    look_ahead_weight: float = 0.5,
) -> ty.Tuple[float, float]:
    """The heuristic cost function used by SABRE, modified to return the effect.

    The effect of the SWAP is a float number that is negative if the SWAP gate has a
    bad effect on the following gates, else positive.

    :param hardware: the SABRE optimiser does not take into account the hardware data to
        compute the heuristic cost, only to generate the possible SWAPs to evaluate
        with this heuristic. The SABRE heuristic only uses the distance matrix.
        Nevertheless, this implementation uses the hardware data to check if some
        gates are ignored (such as barriers for example).
    :param front_layer: the current front layer. Used to compute an "immediate" cost,
        i.e. a quantity that will tell us if the SWAP/Bridge is useful to execute
        gates in the front layer.
    :param topological_nodes: the list of all the DAGNodes of the quantum circuit,
        sorted in topological order.
    :param current_node_index: index of the first non-processed node.
    :param current_mapping: the mapping *before* applying the given SWAP.
    :param distance_matrix: the pre-computed distance matrix between each qubits.
    :param tentative_gate: the SWAP we want to estimate the usefulness of.
    :param look_ahead_depth: the depth of the look-ahead. The procedure will consider
        gates that will be executed in the future (i.e. not in the front layer) up to
        the given depth. Note that 1-qubit gates are not ignored, which means that a
        depth of 3 will not guarantee that there is at least 3 CNOTs in the
        look-ahead set.
    :param look_ahead_weight: weight of the look-ahead. The actual gates (i.e. the
        gates in the front layer) have a weight of 1.
    :return: the heuristic cost of the given SWAP according to the current
        state of the algorithm along with the effect of the SWAP on the non-executed
        gates.
    """
    # First, compute the proposed new mapping
    new_mapping = tentative_gate.update_mapping(current_mapping)
    # Compute H_basic, the cost associated with the distance.
    H_basic = 0.0
    H_basic_gate_number = 0
    for op in front_layer.ops:
        # Only add the gate to the cost if the gate is not already implemented by the
        # SWAP/Bridge
        if not tentative_gate.implements_operation(op, initial_mapping, trans_mapping):
            H_basic += _gate_op_cost(op, distance_matrix, new_mapping, hardware)
            H_basic_gate_number += 1
    # Compute H, the cost cost that encourage parallelism and adds some look-ahead
    # ability.
    #H = tentative_gate.cost(hardware, current_mapping)
    H = 0.0
    future_nodes_layer = QuantumLayer(max_depth=look_ahead_depth)
    # We do not use the return of update_layer because we do not care about the
    # number of gates that were added. Still, we add the firsts look_ahead_depth layers
    # of our future gates in this set to have this look-ahead ability.
    _ = update_layer(future_nodes_layer, topological_nodes, current_node_index)
    # The decay is not implemented in the code the authors gave us and not
    # sufficiently explained in the paper to implement it without guessing. Not
    # implementing it for the moment...
    swap_effect = 0.0
    H += (H_basic / H_basic_gate_number) if H_basic_gate_number != 0 else 0
    if future_nodes_layer:
        # Only add this cost if there are nodes in the future_node_layer
        H += (
            look_ahead_weight
            * sum(
                _gate_op_cost(op, distance_matrix, new_mapping, hardware)
                for op in future_nodes_layer.ops
            )
            / len(future_nodes_layer)
        )
        swap_effect += sum(
            _gate_op_cost(op, distance_matrix, current_mapping, hardware)
            - _gate_op_cost(op, distance_matrix, new_mapping, hardware)
            for op in future_nodes_layer.ops
        )
    return H, swap_effect


# The Start-Up
from pyrosetta import *  ## for newer versions make all start-up commands with pyrosetta
init()
from math import exp, log, pi, sqrt
from random import random as rnd
from random import randint
import sys
import os
import numpy as np
from pyrosetta.rosetta.core.scoring import *
from pyrosetta.rosetta.core.scoring.methods import *
from pyrosetta.rosetta.core.scoring.methods import EnergyMethodOptions
from pyrosetta.rosetta.protocols.grafting import *
from pyrosetta.rosetta.protocols.simple_moves import *
from pyrosetta.rosetta.protocols.moves import *
from pyrosetta.rosetta.core.fragment import *
from pyrosetta.rosetta.protocols.minimization_packing import *
import scipy.optimize as op
from scipy.interpolate import interp1d
import argparse

# Argument Parsing
parser = argparse.ArgumentParser(description='Program')
parser.add_argument('-in', '--Input_FASTA_File', action='store', type=str, required=True,
	help='Name of the text file containing the FASTA sequence of the protein of interest. Carot should not be in same line as sequence, UniProt format preferred.')
parser.add_argument('-abnstruct', '--Number_of_AbInitio_Structures', action='store', type=str, required=False,
	help='Number of structures to sample during Centroid AbInitio portion. Default = 1000')
parser.add_argument('-diso', '--Input_DisoPred_File', action='store', type=str, required=True,
	help='Name of the file containing the per residue disordered probability prediction generated by RaptorX')
parser.add_argument('-t_frag', '--Three_Mer_Frag_Library', action='store', type=str, required=True,
	help='Name of the file containing the three-mer fragment library generated by disorder corrected method')
parser.add_argument('-n_frag', '--Nine_Mer_Frag_Library', action='store', type=str, required=True,
	help='Name of the file containing the nine-mer fragment library generated by disorder corrected method')
parser.add_argument('-rsd_wt_helix', '--Residue_Weight_Helix', action='store', type=float, required=False,
	help='Reweights the env, pair and cbeta scores for helical residues by specificied factor. Default: 0.5')
parser.add_argument('-rsd_wt_loop', '--Residue_Weight_Loop', action='store', type=float, required=False,
	help='Reweights the env, pair and cbeta scores for loop residues by specificied factor. Default: 0.5')
parser.add_argument('-rsd_wt_sheet', '--Residue_Weight_Sheet', action='store', type=float, required=False,
	help='Reweights the env, pair and cbeta scores for sheet residues by specificied factor. Default: 1.0')
parser.add_argument('-rg_weight', '--Rg_Weight', action='store', type=float, required=False,
	help='Reweights the weight of rg in each step by specificied factor. Default: 0.5')
parser.add_argument('-cycles', '--AbInitio_Cycles', action='store', type=str, required=False,
	help='Number of sampling cycles within each stage of AbInitio, identical to increase_cycles flag in C++ ClassicAbInitio and AbRelax. Default 10')
parser.add_argument('-abinitiovo', '--AbInitioVO', action='store_true', required=False,
	help='Boolean for running protocol with AbInitioVO score function, runs with flag')
parser.add_argument('-refinesubset', '--Refine_Subset', action='store', type=int, required=False,
	help='Only subjects the lowest X% of structures to Relax refinement where X is specified as input following flag. Default 100')
parser.add_argument('-relnstruct', '--Number_of_Relax_Structures', action='store', type=str, required=False,
	help='Number of independent full-atom Relax sampling trajectories from a single AbInitio structure. Default 50')
args = parser.parse_args()

## Imports from Parser and Defaults
if args.Number_of_AbInitio_Structures:
	abnstruct = int(args.Number_of_AbInitio_Structures)
else:
	abnstruct = 20
if args.AbInitio_Cycles:
	cycles = int(args.AbInitio_Cycles)
else:
	cycles = 10
if args.Number_of_Relax_Structures:
	relnstruct = int(args.Number_of_Relax_Structures)
else:
	relnstruct = 50
if args.Refine_Subset:
	refine_number = int((args.Refine_Subset/100)*int(abnstruct))
else:
	refine_number = int(abnstruct)

# Importing sequence from FASTA
fasta_file = open(args.Input_FASTA_File, 'r')
fasta_lines = fasta_file.readlines()
fasta_counter = 0
fasta_sequence = ' '
for fasta_line in fasta_lines:
	if '>' not in fasta_line:
		if fasta_counter == 0:
			if '\n' in fasta_line:
				fasta_sequence = fasta_line.split('\n')[0]
			else:
				fasta_sequence = fasta_line
			fasta_counter = 1
		else:
			if '\n' in fasta_line:
				fasta_sequence = fasta_sequence + fasta_line.split('\n')[0]
			else:
				fasta_sequence = fasta_sequence + fasta_line

# The Poses
#p=Pose()
#make_pose_from_sequence(p, fasta_sequence, "centroid")
p = pose_from_sequence(fasta_sequence, "centroid")
starting_p = Pose()
starting_p.assign(p)
fa_p = Pose()
fa_p.assign(p)
pvdw = Pose()
pvdwc = Pose()
pcen = Pose()
maplow = MoveMap()
maplow.set_bb(True)
maplow.set_chi(True)

# Score Functions
## Preparing the Rg Score Term
#### Determining the Per Residue Disorder Probability and Number of Segments
disorder_dtypes = [('res_num', np.float_), ('AA', np.unicode_, 2), ('ast', np.unicode_, 1), ('disprob', np.float_)]
disorder_dat = np.genfromtxt(args.Input_DisoPred_File, dtype=disorder_dtypes, delimiter=' ', skip_header=3)
diso_dat = np.empty((len(disorder_dat),1))
diso_segments = []
diso_cutoff = 0.5
seg_cutoff = 10
seg_res_counter = 0
per_residue_disorder = 0
diso_current = disorder_dat[0][3]
segment_break_list = []
segment_switch = False
for seg_residue in range(len(disorder_dat)):
	if disorder_dat[seg_residue][3] >= diso_cutoff and diso_current >= diso_cutoff:
		seg_res_counter = 0
	elif disorder_dat[seg_residue][3] < diso_cutoff and diso_current < diso_cutoff:
		seg_res_counter = 0
	elif disorder_dat[seg_residue][3] < diso_cutoff and diso_current >= diso_cutoff:
		seg_res_counter = seg_res_counter + 1
	elif disorder_dat[seg_residue][3] >= diso_cutoff and diso_current < diso_cutoff:
		seg_res_counter = seg_res_counter + 1
	if seg_res_counter > 9:
		seg_res_counter = 0
		diso_current = disorder_dat[seg_residue][3]
		if diso_current >= diso_cutoff:
			transition_id = 'order' # meaning all residues prior to end are ordered
		else:
			transition_id = 'disorder' # meaning all residues prior to end are disordered
		segment_break_list.append((seg_residue-9, transition_id))
	per_residue_disorder = per_residue_disorder + (disorder_dat[seg_residue][3]/float(len(disorder_dat)))
if disorder_dat[len(disorder_dat)-1][3] >= diso_cutoff:
	segment_break_list.append((len(disorder_dat), 'disorder'))
else:
	segment_break_list.append((len(disorder_dat), 'order'))

#### Computing the Mean Hydrophobicity per Residue
scaling_info_holder = np.zeros([len(segment_break_list),5]) # Need spots for segH,segQ,frac_pos,frac_neg,seq_len -> scalvQ,scalvH,scalv,rad_gyr,seq_len => scalvQ -> sf_rg_term_potential

##### Kyte, J. Doolitte, R.F. J. Mol. Biol. (1982) 157, 105-132
##### Normalization proposed in Uversky, V.N. et. al. Proteins: Struc. Func. Gen. (2000), 41, 415-427
aa_hydro_idx = {'I':4.5, 'V':4.2, 'L':3.8, 'F':2.8, 'C':2.5, 'M':1.9, 'A':1.8, 'G':-0.4, 'T':-0.7, 'W':-0.9, 'S':-0.8, 'Y':-1.3, 'P':-1.6, 'H':-3.2, 'E':-3.5, 'Q':-3.5, 'D':-3.5, 'N':-3.5, 'K':-3.9, 'R':-4.5}
for diso_segment_idx, diso_segment_item in enumerate(segment_break_list):
	start_res = 1
	if diso_segment_idx > 0:
		start_res = segment_break_list[diso_segment_idx-1][0] + 1
	end_res = diso_segment_item[0] + 1
	seq_len = end_res-start_res
	pro_len = p.total_residue()
	scaling_info_holder[diso_segment_idx][4] = seq_len
	sequence_avg_H = 0
	for res_idx in range(start_res-1, end_res-1, 1):
		window_start = 0
		window_end = 0
		window_H_avg = 0
		if res_idx < 5:
			window_start = 0
			window_end = res_idx + 1
		elif pro_len - res_idx < 5:
			window_start = res_idx
			window_end = pro_len
		else:
			window_start = res_idx - 2
			window_end = res_idx + 3
		window = p.sequence()[window_start:window_end]
		for window_res_idx, window_res_item in enumerate(window):
			window_H_avg = window_H_avg + ((aa_hydro_idx[str(window_res_item)]+4.5)/9.0)
		window_H_avg = window_H_avg/(len(window))
		sequence_avg_H = sequence_avg_H + window_H_avg/scaling_info_holder[diso_segment_idx][4]
		scaling_info_holder[diso_segment_idx][0] = sequence_avg_H

#### Computing the Mean Net Charge per Residue
for diso_segment_idx, diso_segment_item in enumerate(segment_break_list):
	start_res = 1
	if diso_segment_idx > 0:
		start_res = segment_break_list[diso_segment_idx-1][0] + 1
	end_res = diso_segment_item[0] + 1
	seq_len = end_res-start_res
	scaling_info_holder[diso_segment_idx][4] = seq_len
	sequence_Q = 0
	frac_pos = 0 # fraction of positively charged residues Arg Lys
	frac_neg = 0 # fraction of negatively charged residues Asp Glu
	for res_idx in range(start_res-1, end_res-1, 1):
		res_id = p.sequence()[res_idx]
		if res_id == 'E' or res_id == 'D':
			frac_neg = frac_neg + 1/seq_len
			sequence_Q = sequence_Q - 1/seq_len
		elif res_id == 'K' or res_id == 'R':
			frac_pos = frac_pos + 1/seq_len
			sequence_Q = sequence_Q + 1/seq_len
		else:
			continue
	sequence_Q = np.sqrt(sequence_Q**2)
	scaling_info_holder[diso_segment_idx][1] = sequence_Q
	scaling_info_holder[diso_segment_idx][2] = frac_pos
	scaling_info_holder[diso_segment_idx][3] = frac_neg

#### Calculations from Hofmann, H. et al. PNAS (2012) 109, 40, 16155-16160
##### Compute charge and hydrophobicity scaling options
###### Constants
scalv = 0
con_a = 0.394
con_z = 0.09
con_x0 = 0.114
con_c = 1.72
con_d = 0.9

###### Calculations
###### Folded vs Unfolded
###### From Uversky, if sequence_Q < 2.785*sequence_avg_H - 1.151 -> Folded, but fails for synuclein
for QH_set_idx in range(len(scaling_info_holder)):
	if 'ordered' in segment_break_list[QH_set_idx]:
		scaling_info_holder[QH_set_idx][0] = 0.33
		scaling_info_holder[QH_set_idx][1] = 0.33
	else:
		scalv_Q = (1.0/3.0) + con_a*((1+np.exp(con_x0-scaling_info_holder[QH_set_idx][1]/con_z))**(-1))
		scaling_info_holder[QH_set_idx][1] = scalv_Q
		scalv_H = (1.0/3.0) + con_a*((1+np.exp((con_x0+con_c*scaling_info_holder[QH_set_idx][0]-con_d)/con_z))**(-1))
		scaling_info_holder[QH_set_idx][0] = scalv_H

###### Determining the appropriate scaling value
####### Constants
elemchar = 1.602*10**(-19)
ss_eo = 8.854*10**(-12)
ss_er = 78.7 #permittivity of water at 298 K
ss_kb = 1.38*10**(-23) # boltzmann constant
ss_T = 298 # in Kelvin
ss_I = 0.15 # in Molar
ss_kappa = 1/(0.304*np.sqrt(ss_I)) # in nanometers
ss_lB = (elemchar**2)/(4*np.pi*ss_eo*ss_er*ss_kb*ss_T)

####### Calculation
for QH_set_idx in range(len(scaling_info_holder)):
	scalvstar = ((4*np.pi*ss_lB*((scaling_info_holder[QH_set_idx][2]-scaling_info_holder[QH_set_idx][3])**2))/(ss_kappa**2)) - ((np.pi*ss_lB*((scaling_info_holder[QH_set_idx][2]+scaling_info_holder[QH_set_idx][3])**2))/(ss_kappa))
	if scaling_info_holder[QH_set_idx][2] == 0:
		if scaling_info_holder[QH_set_idx][3] == 0:
			scalv = scaling_info_holder[QH_set_idx][0]
		else:
			scalv = scaling_info_holder[QH_set_idx][1]
	elif scaling_info_holder[QH_set_idx][3] == 0:
		if scaling_info_holder[QH_set_idx][2] == 0:
			scalv = scaling_info_holder[QH_set_idx][0]
		else:
			scalv = scaling_info_holder[QH_set_idx][1]
	elif scalvstar < 0:
		scalv = scaling_info_holder[QH_set_idx][1]
	else:
		scalv = scaling_info_holder[QH_set_idx][0]
	scaling_info_holder[QH_set_idx][2] = scalv

##### Computing the Expected Radius of Gyration
rg_b = 0.38 # in nanometers
rg_lp = 0.53 # in nanometers
for QH_set_idx in range(len(scaling_info_holder)):
	rad_gyr = (np.sqrt((2*rg_lp*rg_b)/((2*scaling_info_holder[QH_set_idx][2]+1)*(2*scaling_info_holder[QH_set_idx][2]+2)))*(scaling_info_holder[QH_set_idx][4]**scaling_info_holder[QH_set_idx][2]))*10 # in Angstroms
	scaling_info_holder[QH_set_idx][3] = rad_gyr

## Making the Actual Potential Energy Function
## Constants for Probability
gamma = 1.1615

## Concocting the Potential
def pr_saw(var_a, input_vars):
	r, rg, g, delta = input_vars
	return (var_a[0]*4*np.pi/rg)*((r/rg)**(2+g))*np.exp(-var_a[1]*((r/rg)**delta))

def solve_pr_saw(var_a, *input_vars):
	r, rg, g, delta = input_vars
	return (np.sum((var_a[0]*4*np.pi/rg)*((r/rg)**(2+g))*np.exp(-var_a[1]*((r/rg)**delta)))-1, np.sum((var_a[0]*4*np.pi/rg)*((r/rg)**(2+g))*np.exp(-var_a[1]*((r/rg)**delta))*(r**2))-rg**2)

rg_sf_term_potential_list = []
for QH_set_idx in range(len(scaling_info_holder)):
	var_g = (gamma-1)/scaling_info_holder[QH_set_idx][2]
	var_delta = 1/(1-scaling_info_holder[QH_set_idx][2])
	r_set=np.arange(0.0,5*scaling_info_holder[QH_set_idx][3],0.01)
	saw_inputs = (r_set, scaling_info_holder[QH_set_idx][3], var_g, var_delta)
	a0 = [1.0, 1.0]
	a1 = op.fsolve(solve_pr_saw,a0,args=saw_inputs)
	rg_sf_term_potential = interp1d(r_set, (1-(pr_saw(a1, saw_inputs)/np.max(pr_saw(a1, saw_inputs)))))
	rg_sf_term_potential_list.append(rg_sf_term_potential)

## Making the Actual Score Term
from pyrosetta.rosetta.core.scoring.methods import ContextIndependentOneBodyEnergy ## newer versions make this pyrosetta.rosetta
@pyrosetta.EnergyMethod() ## for newer versions make pyrosetta.EnergyMethod()
class SeqCorrRgMethod(WholeStructureEnergy):
	"""A scoring method that using a predicted radius of gyration from the
	primary sequence to construct a polymer-scaled potential

	"""
	def __init__(self):
		"""Construct LengthScoreMethod."""
		WholeStructureEnergy.__init__(self, self.creator())

	def finalize_total_energy(self, pose, sfxn, emap):
		"""Calculate energy of res of pose and emap"""
		pose = pose ## for newer versions remove line
		e_val = 0
		for segment_idx in range(len(segment_break_list)):
			r_xyz = np.zeros([int(scaling_info_holder[segment_idx][4]), 3])
			rg_sq = np.zeros([int(scaling_info_holder[segment_idx][4]), 1])
			start_res = 1
			if segment_idx > 0:
				start_res = segment_break_list[segment_idx-1][0] + 1
			end_res = segment_break_list[segment_idx][0] + 1
			for res_num in range(start_res, end_res, 1):
				for res_xyz in range(len(r_xyz[0])):
					r_xyz[res_num-start_res][res_xyz] = pose.residue(res_num).nbr_atom_xyz()[res_xyz]
			r_cen_mass = np.average(r_xyz, axis=0)
			for res_num in range(len(r_xyz)):
				rg_sq[res_num] = (np.linalg.norm(r_xyz[res_num] - r_cen_mass))**2
			rg_val = np.sqrt(np.average(rg_sq))
			rg_sf_term_potential = rg_sf_term_potential_list[segment_idx]
			e_val = e_val + float(rg_sf_term_potential(rg_val))
		emap.set(self.scoreType, e_val) ## for newer versions remove .get()

new_rg_score = SeqCorrRgMethod.scoreType

## DSSP-based Reweighting Score Terms
dssp = rosetta.protocols.moves.DsspMover()
dssp_E_weight = 0.0
dssp_L_weight = 0.0
dssp_H_weight = 0.0

if args.Residue_Weight_Sheet:
	dssp_E_weight = 1.0 - float(args.Residue_Weight_Sheet)
else:
	dssp_E_weight = 0.0

if args.Residue_Weight_Loop:
	dssp_L_weight = 1.0 - float(args.Residue_Weight_Loop)
else:
	dssp_L_weight = 0.5

if args.Residue_Weight_Helix:
	dssp_H_weight = 1.0 - float(args.Residue_Weight_Helix)
else:
	dssp_H_weight = 0.5

sec_struct_weight = {'L':float(dssp_L_weight), 'H':float(dssp_H_weight), 'E':float(dssp_E_weight)}

### Preparing Unweighted Score Functions
sf_env = ScoreFunction()
sf_env.set_weight(env, 1.0)

sf_pair = ScoreFunction()
sf_pair.set_weight(pair, 1.0)

sf_cbeta = ScoreFunction()
sf_cbeta.set_weight(cbeta, 1.0)

### Centroid DSSP Weighted Pair Term
@pyrosetta.EnergyMethod()
class SecondaryStructurePenalty(ContextIndependentTwoBodyEnergy):
	"""A scoring method that assigns different score weights to residues
	depending on the secondary structure.

	"""
	def __init__(self):
		"""Construct LengthScoreMethod."""
		ContextIndependentTwoBodyEnergy.__init__(self, self.creator())

	def setup_for_scoring(self, pose, sf):
		pose = pose
		dssp.apply(pose)

	def defines_intrares_energy(self, weights):
		"""Return True if intra-residue energy is Defined."""
		return True

	def eval_intrares_energy(self, res, pose, sf, emap):
		"""Calculate intra-residue energy if defined."""
		pose = pose
		emv = EMapVector()
		sf_env.eval_cd_1b(res,pose,emv)
		weighted_score_env = emv[env]*sec_struct_weight[pose.secstruct()[res.seqpos()-1]]
		sf_cbeta.eval_cd_1b(res,pose,emv)
		weighted_score_cbeta = emv[cbeta]*sec_struct_weight[pose.secstruct()[res.seqpos()-1]]
		weighted_score = -1*(weighted_score_env + weighted_score_cbeta)
		emap.set(self.scoreType, weighted_score)

	def atomic_interaction_cutoff(self):
		"""Get the cutoff."""
		return 0.0

	def residue_pair_energy(self, res1, res2, pose, sf, emap):
		"""Calculate energy of res of pose and set emap"""
		pose = pose
		emv = EMapVector()
		sf_pair.eval_ci_2b(res1,res2,pose,emv)
		weighted_score = -1*(emv[pair]*sec_struct_weight[pose.secstruct()[res1.seqpos()-1]]*sec_struct_weight[pose.secstruct()[res2.seqpos()-1]])
		emap.set(self.scoreType, weighted_score)

secstrucpenalty_score = SecondaryStructurePenalty.scoreType

## Preparing Score Functions
if args.Rg_Weight:
	rg_reweight_factor = args.Rg_Weight
else:
	rg_reweight_factor = 0.5

sf0 = create_score_function('score0')
sf_stage_1 = create_score_function('score0')

sf_stage_2 = create_score_function('score1')
if args.AbInitioVO == True:
	sf_stage_2.set_weight(new_rg_score, ((400/24)*5))
sf_stage_2.set_weight(secstrucpenalty_score, 1.0)

sf_stage_3a = create_score_function('score2')
if args.AbInitioVO == True:
	sf_stage_3a.set_weight(new_rg_score, ((400/24)*5))
sf_stage_3a.set_weight(secstrucpenalty_score, 1.0)

sf_stage_3b = create_score_function('score5')
if args.AbInitioVO == True:
	sf_stage_3b.set_weight(new_rg_score, ((400/24)*5))
sf_stage_3b.set_weight(secstrucpenalty_score, 1.0)

sf_stage_4 = create_score_function('score3')
if args.AbInitioVO == True:
	sf_stage_4.set_weight(rg, 0.0)
	sf_stage_4.set_weight(new_rg_score, ((400/24)*5))
else:
	sf_stage_4.set_weight(rg, sf_stage_4.get_weight(rg)*rg_reweight_factor)
sf_stage_4.set_weight(hbond_sr_bb, 1.0)
sf_stage_4.set_weight(hbond_lr_bb, 1.0)
sf_stage_4.set_weight(rama, 1.0)
sf_stage_4.set_weight(secstrucpenalty_score, 1.0)

## Setting Up Stage 4
sf_stage_4_term_list = [vdw, hbond_sr_bb, hbond_lr_bb, rama, new_rg_score]
cen_score_holder_dtypes = [('out_name', np.unicode_, 100)]
for score_4_term in range(len(sf_stage_4_term_list)):
	cen_score_holder_dtypes.append((str(sf_stage_4_term_list[score_4_term]), float))
cen_score_holder = np.zeros([int(abnstruct)], dtype=cen_score_holder_dtypes)

## Full Atom Score Functions
sfbeta = create_score_function('ref2015')
sfrelax = create_score_function('ref2015_cart')

## Radius of Gyration Reference Score Function
sfrg = ScoreFunction()
sfrg.set_weight(rg, 1.0)

# The Movers
## Fragment Movers
# Importing the fragment files
fragset9 = ConstantLengthFragSet(9)
fragset9.read_fragment_file(args.Nine_Mer_Frag_Library)
fragset3 = ConstantLengthFragSet(3)
fragset3.read_fragment_file(args.Three_Mer_Frag_Library)

# Constructing the Fragment Mover
fragmover9 = ClassicFragmentMover(fragset9, maplow)
fragmover3 = ClassicFragmentMover(fragset3, maplow)
gunncost = GunnCost()
smoothfragmover3 = SmoothFragmentMover(fragset3, maplow, gunncost)

## Phi-Psi Movers
smMovercen = SmallMover(maplow, 0.8, 1)
shMovercen = ShearMover(maplow, 0.8, 1)
smMovercen.angle_max("H", 2.0)
smMovercen.angle_max("E", 2.0)
smMovercen.angle_max("L", 5.0)
shMovercen.angle_max("H", 2.0)
shMovercen.angle_max("E", 2.0)
shMovercen.angle_max("L", 5.0)

smMoverfa = SmallMover(maplow, 0.8, 1)
shMoverfa = ShearMover(maplow, 0.8, 1)
smMoverfa.angle_max("H", 2.0)
smMoverfa.angle_max("E", 2.0)
smMoverfa.angle_max("L", 5.0)
shMoverfa.angle_max("H", 2.0)
shMoverfa.angle_max("E", 2.0)
shMoverfa.angle_max("L", 5.0)

##Random Movers
random_stage_1 = RandomMover()
random_stage_1.add_mover(fragmover9)

random_stage_2 = RandomMover()
random_stage_2.add_mover(fragmover9)

random_stage_3 = RandomMover()
random_stage_3.add_mover(fragmover9)

random_stage_4a = RandomMover()
random_stage_4a.add_mover(fragmover3)

random_stage_4b = RandomMover()
random_stage_4b.add_mover(smoothfragmover3)

# Relax Mover
relax = rosetta.protocols.relax.FastRelax()
relax.min_type('lbfgs_armijo_nonmonotone')
relax.dualspace(True)
relax.set_scorefxn(sfrelax)
relax.max_iter(200)

# The Monte Carlo
## PyRosetta Monte Carlo Objects
mc_stage_1 = MonteCarlo(p, sf_stage_1, 2.0)
mc_stage_2 = MonteCarlo(p, sf_stage_2, 2.0)
mc_stage_3a = MonteCarlo(p, sf_stage_3a, 2.0)
mc_stage_3b = MonteCarlo(p, sf_stage_3b, 2.0)
mc_stage_4 = MonteCarlo(p, sf_stage_4, 2.0)

## Setting up Trial Movers
trial_stage_1 = TrialMover(random_stage_1, mc_stage_1)
trial_stage_2 = TrialMover(random_stage_2, mc_stage_2)
trial_stage_3a = TrialMover(random_stage_3, mc_stage_3a)
trial_stage_3b = TrialMover(random_stage_3, mc_stage_3b)
trial_stage_4a = TrialMover(random_stage_4a, mc_stage_4)
trial_stage_4b = TrialMover(random_stage_4b, mc_stage_4)

## Setting up Repeat Movers
stage1 = RepeatMover(trial_stage_1, 2000)
stage2 = RepeatMover(trial_stage_2, 2000)
stage3a = RepeatMover(trial_stage_3a, 4000)
stage3b = RepeatMover(trial_stage_3b, 4000)
stage4a = RepeatMover(trial_stage_4a, 4000)
stage4b = RepeatMover(trial_stage_4b, 4000)

# Converting the Pose
switch = SwitchResidueTypeSetMover('fa_standard')

# Side-Chain Movers
switch.apply(fa_p)
task = standard_packer_task(fa_p)
task.restrict_to_repacking()
pack_mover = PackRotamersMover(sfbeta, task)
rot_mover = RotamerTrialsMover(sfbeta, task)

# The Simulation and Output
for i in range(int(abnstruct)):
	# Setting up the Input Structure
	p.assign(starting_p)
	for pose_res_num in range(p.total_residue()):
		p.set_phi(pose_res_num+1,-150)
		p.set_psi(pose_res_num+1,150)
		p.set_omega(pose_res_num+1,180)
	mc_stage_1.reset(p)
	mc_stage_2.reset(p)
	mc_stage_3a.reset(p)
	mc_stage_3b.reset(p)
	mc_stage_4.reset(p)
	for j in range(cycles):
		print('Performing Stage 1 Sampling: Phase: ' + str(j))
		sf_env.set_weight(env, sf_stage_1.get_weight(env))
		sf_pair.set_weight(pair, sf_stage_1.get_weight(pair))
		sf_cbeta.set_weight(cbeta, sf_stage_1.get_weight(cbeta))
		stage1.apply(p)
		mc_stage_1.recover_low(p)
		sf_stage_1.show(p)
		mc_stage_1.reset(p)
	for k in range(cycles):
		print('Performing Stage 2 Sampling: Phase: ' + str(k))
		sf_env.set_weight(env, sf_stage_2.get_weight(env))
		sf_pair.set_weight(pair, sf_stage_2.get_weight(pair))
		sf_cbeta.set_weight(cbeta, sf_stage_2.get_weight(cbeta))
		stage2.apply(p)
		mc_stage_2.recover_low(p)
		sf_stage_2.show(p)
		mc_stage_2.reset(p)
	for l in range(int(cycles/2)):
		print('Performing Stage 3 Sampling: Phase: ' + str(l))
		if l % 2 == 0:
			sf_env.set_weight(env, sf_stage_3a.get_weight(env))
			sf_pair.set_weight(pair, sf_stage_3a.get_weight(pair))
			sf_cbeta.set_weight(cbeta, sf_stage_3a.get_weight(cbeta))
			stage3a.apply(p)
			mc_stage_3a.recover_low(p)
			mc_stage_3a.reset(p)
		else:
			sf_env.set_weight(env, sf_stage_3b.get_weight(env))
			sf_pair.set_weight(pair, sf_stage_3b.get_weight(pair))
			sf_cbeta.set_weight(cbeta, sf_stage_3b.get_weight(cbeta))
			stage3b.apply(p)
			mc_stage_3b.recover_low(p)
			mc_stage_3b.reset(p)
		sf_stage_3b.show(p)
	for m in range(cycles):
		print('Performing Stage 4 Sampling: Phase: ' + str(m))
		sf_env.set_weight(env, sf_stage_4.get_weight(env))
		sf_pair.set_weight(pair, sf_stage_4.get_weight(pair))
		sf_cbeta.set_weight(cbeta, sf_stage_4.get_weight(cbeta))
		if m < 2:
			stage4a.apply(p)
			mc_stage_4.recover_low(p)
		else:
			stage4b.apply(p)
			mc_stage_4.recover_low(p)
		sf_stage_4.show(p)
		mc_stage_4.reset(p)
	for record_idx in range(len(sf_stage_4_term_list)):
		cen_score_holder[i][record_idx] = p.energies().total_energies()[sf_stage_4_term_list[record_idx]]
	pcen.assign(p)
	if args.AbInitioVO == True:
		outf = open("AbInitioVO_Centroid.sc", 'a')
		pdb_out = "AbInitioVO_Centroid_out_%i.pdb" %i
	else:
		outf = open("AbInitio_Centroid.sc", 'a')
		pdb_out = "AbInitio_Centroid_out_%i.pdb" %i
	pcen.dump_pdb(pdb_out)
	switch.apply(p)
	sfbeta.show(p)
	outf.write("%s\t%.3f\t%.4f\t%.3f\n" % (pdb_out, sfbeta(p), sf_stage_4(pcen), sfrg(p)))
	outf.close()
np.savetxt('Centroid_Score_Breakdown.txt', cen_score_holder, fmt='%s', delimiter=' ', newline='\n')

## Deciding who to minimize
dtype_list = [('out_name','S50'),('full_sc',float),('cen_score',float),('out_rg',float)]
if args.AbInitioVO == True:
	cen_out_data = np.genfromtxt('AbInitioVO_Centroid.sc', dtype=dtype_list)
else:
	cen_out_data = np.genfromtxt('AbInitio_Centroid.sc', dtype=dtype_list)
cen_out_sort = np.sort(cen_out_data, order='cen_score')
for cen_out_struct_idx in range(refine_number):
	cen_out_struct_item = cen_out_data[cen_out_struct_idx]['out_name']
	relax_p_in = pose_from_pdb(str(cen_out_struct_item.decode('UTF-8')))
	relax_p = Pose()
	print('Relaxing Output ' + str(cen_out_struct_idx+1) + ' of ' + str(refine_number))
	for relnstruct_idx in range(relnstruct):
		relax_p.assign(relax_p_in)
		relax.apply(relax_p)
		sfrelax.show(relax_p)
		if args.AbInitioVO == True:
			outf = open("AbInitioVO_FullAtom.sc", 'a')
			pdb_out = "AbInitioVO_FullAtom_out_" + str(cen_out_struct_idx) + "_" + str(relnstruct_idx) + ".pdb"
		else:
			outf = open("AbInitio_FullAtom.sc", 'a')
			pdb_out = "AbInitio_FullAtom_out_" + str(cen_out_struct_idx) + "_" + str(relnstruct_idx) + ".pdb"
		outf.write("%s\t%s\t%.4f\t%.4f\n" % (pdb_out, str(cen_out_struct_item), sfrelax(relax_p), sfrg(relax_p)))
		relax_p.dump_pdb(pdb_out)
		outf.close()


"""
implementation of DDPG

DPG + DQN: deterministic PG + deep Q network
model-free, off-policy, actor-critic
DQN - discrete space; DDPG - continuous space

Adapted from https://github.com/sweetice/Deep-reinforcement-learning-with-pytorch/blob/master/Char05%20DDPG/DDPG.py

Lily Xu, 2021
"""

import os, sys
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

class Critic(nn.Module):
    def __init__(self, input_size, output_size):
        """
        input: state, action
        output: value of each action
        """

        hidden_size1 = 16
        hidden_size2 = 32
        super(Critic, self).__init__()
        self.linear1 = nn.Linear(input_size, hidden_size1)
        self.linear2 = nn.Linear(hidden_size1, hidden_size2)
        self.linear3 = nn.Linear(hidden_size2, output_size)

    def forward(self, state, action):
        """
        params
        state and action - both torch Tensors
        """
        x = torch.cat([state, action], 1)
        x = self.linear1(x)
        x = F.relu(x)
        x = self.linear2(x)
        x = F.relu(x)
        x = self.linear3(x)
        return x

class Actor(nn.Module):
    def __init__(self, input_size, output_size, learning_rate=3e-4):
        """
        input: state
        output: action to take
        """
        hidden_size1 = 16
        hidden_size2 = 32
        super(Actor, self).__init__()

        self.linear1 = nn.Linear(input_size, hidden_size1)
        self.linear2 = nn.Linear(hidden_size1, hidden_size2)
        self.linear3 = nn.Linear(hidden_size2, output_size)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, state):
        x = self.linear1(state)
        x = F.relu(x)
        x = self.linear2(x)
        x = F.relu(x)
        x = self.linear3(x)
        x = self.softmax(x)
        return x

########### replay buffer

import random
from collections import deque

class ReplayBuffer:
    """ experience replay to overcome challenge of data not being independently distributed
    during on-policy training """
    def __init__(self, max_size):
        self.buffer = deque(maxlen=max_size)

    def push(self, s, a, r, ss, done):
        """ state, action, reward, next_state """
        experience = (s, a, r, ss, done)
        self.buffer.append(experience)

    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)

        s_batch    = [experience[0] for experience in batch]
        a_batch    = [experience[1] for experience in batch]
        r_batch    = [experience[2] for experience in batch]
        ss_batch   = [experience[3] for experience in batch]
        done_batch = [experience[4] for experience in batch]

        return s_batch, a_batch, r_batch, ss_batch, done_batch

    def __len__(self):
        return len(self.buffer)

# https://github.com/openai/gym/blob/master/gym/core.py
class NormalizedEnv:
    """ Wrap action """

    def _action(self, action):
        act_k = (self.action_space.high - self.action_space.low)/ 2.
        act_b = (self.action_space.high + self.action_space.low)/ 2.
        return act_k * action + act_b

    def _reverse_action(self, action):
        act_k_inv = 2./(self.action_space.high - self.action_space.low)
        act_b = (self.action_space.high + self.action_space.low)/ 2.
        return act_k_inv * (action - act_b)

class DDPG:
    def __init__(self, n_targets, actor_learning_rate=1e-4,
        critic_learning_rate=1e-3, gamma=0.99, tau=1e-2,
        memory_max_size=50000):

        # params
        self.states_dim = 2 * n_targets + 1
        self.actions_dim = n_targets
        self.gamma = gamma
        self.tau = tau

        ##### networks
        # randomly initialize critic and actor network
        self.actor         = Actor(self.states_dim, self.actions_dim)
        self.critic        = Critic(self.states_dim + self.actions_dim, 1)

        # initialize critic and actor target network
        self.actor_target  = Actor(self.states_dim, self.actions_dim)
        self.critic_target = Critic(self.states_dim + self.actions_dim, 1)

        for target_param, param in zip(self.actor_target.parameters(), self.actor.parameters()):
            target_param.data.copy_(param.data)

        for target_param, param in zip(self.critic_target.parameters(), self.critic.parameters()):
            target_param.data.copy_(param.data)

        self.actor_optimizer  = torch.optim.Adam(self.actor.parameters(), lr=actor_learning_rate)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=critic_learning_rate)

        self.memory = ReplayBuffer(memory_max_size)
        self.loss = nn.MSELoss()

    def select_action(self, state):
        if not torch.is_tensor(state):
            state = torch.from_numpy(state).float().unsqueeze(0)
        else:
            state = state.float().unsqueeze(0)
        action = self.actor.forward(state)
        action = action.detach().numpy()[0]
        return action

    def update(self, batch_size):
        states, actions, rewards, next_states, dones = self.memory.sample(batch_size)
        states      = torch.FloatTensor(states)
        actions     = torch.FloatTensor(actions)
        rewards     = torch.FloatTensor(rewards)
        next_states = torch.FloatTensor(next_states)

        # update critic by minimizing loss
        Q_vals = self.critic.forward(states, actions)
        next_actions = self.actor_target.forward(next_states)
        next_Q = self.critic_target.forward(next_states, next_actions.detach())
        Q_prime = rewards + self.gamma * next_Q

        critic_loss = self.loss(Q_vals, Q_prime)
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        # update actor policy using sampled policy gradient
        policy_loss = -self.critic.forward(states, self.actor.forward(states)).mean()
        self.actor_optimizer.zero_grad()
        policy_loss.backward()
        self.actor_optimizer.step()

        # update target networks (slowly using soft updates)
        for target_param, param in zip(self.actor_target.parameters(), self.actor.parameters()):
            target_param.data.copy_(self.tau * param.data + (1. - self.tau) * target_param.data)

        for target_param, param in zip(self.critic_target.parameters(), self.critic.parameters()):
            target_param.data.copy_(self.tau * param.data + (1. - self.tau) * target_param.data)


import tensorflow as tf
import scipy.io as sio
import numpy as np
from PIL import Image

img_H = 64
img_W = 64
img_C = 3
GAN_type = "SNGAN"  # DCGAN, WGAN, WGAN-GP, SNGAN, LSGAN, RSGAN, RaSGAN
batchsize = 128
epsilon = 1e-14#if epsilon is too big, training of DCGAN is failure.

def deconv(inputs, shape, strides, out_num, is_sn=False):
    filters = tf.get_variable("kernel", shape=shape, initializer=tf.random_normal_initializer(stddev=0.02))
    bias = tf.get_variable("bias", shape=[shape[-2]], initializer=tf.constant_initializer([0]))
    if is_sn:
        return tf.nn.conv2d_transpose(inputs, spectral_norm("sn", filters), out_num, strides) + bias
    else:
        return tf.nn.conv2d_transpose(inputs, filters, out_num, strides) + bias

def conv(inputs, shape, strides, is_sn=False):
    filters = tf.get_variable("kernel", shape=shape, initializer=tf.random_normal_initializer(stddev=0.02))
    bias = tf.get_variable("bias", shape=[shape[-1]], initializer=tf.constant_initializer([0]))
    if is_sn:
        return tf.nn.conv2d(inputs, spectral_norm("sn", filters), strides, "SAME") + bias
    else:
        return tf.nn.conv2d(inputs, filters, strides, "SAME") + bias

def fully_connected(inputs, num_out, is_sn=False):
    W = tf.get_variable("W", [inputs.shape[-1], num_out], initializer=tf.random_normal_initializer(stddev=0.02))
    b = tf.get_variable("b", [num_out], initializer=tf.constant_initializer([0]))
    if is_sn:
        return tf.matmul(inputs, spectral_norm("sn", W)) + b
    else:
        return tf.matmul(inputs, W) + b

def leaky_relu(inputs, slope=0.2):
    return tf.maximum(slope*inputs, inputs)

def spectral_norm(name, w, iteration=1):
    #Spectral normalization which was published on ICLR2018,please refer to "https://www.researchgate.net/publication/318572189_Spectral_Normalization_for_Generative_Adversarial_Networks"
    #This function spectral_norm is forked from "https://github.com/taki0112/Spectral_Normalization-Tensorflow"
    w_shape = w.shape.as_list()
    w = tf.reshape(w, [-1, w_shape[-1]])
    with tf.variable_scope(name, reuse=False):
        u = tf.get_variable("u", [1, w_shape[-1]], initializer=tf.truncated_normal_initializer(), trainable=False)
    u_hat = u
    v_hat = None

    def l2_norm(v, eps=1e-12):
        return v / (tf.reduce_sum(v ** 2) ** 0.5 + eps)

    for i in range(iteration):
        v_ = tf.matmul(u_hat, tf.transpose(w))
        v_hat = l2_norm(v_)
        u_ = tf.matmul(v_hat, w)
        u_hat = l2_norm(u_)
    sigma = tf.matmul(tf.matmul(v_hat, w), tf.transpose(u_hat))
    w_norm = w / sigma
    with tf.control_dependencies([u.assign(u_hat)]):
        w_norm = tf.reshape(w_norm, w_shape)
    return w_norm

def mapping(x):
    max = np.max(x)
    min = np.min(x)
    return (x - min) * 255.0 / (max - min + epsilon)

def instanceNorm(inputs):
    mean, var = tf.nn.moments(inputs, axes=[1, 2], keep_dims=True)
    scale = tf.get_variable("scale", shape=mean.shape[-1], initializer=tf.constant_initializer([1.0]))
    shift = tf.get_variable("shift", shape=mean.shape[-1], initializer=tf.constant_initializer([0.0]))
    return (inputs - mean) * scale / (tf.sqrt(var + epsilon)) + shift

class Generator:
    def __init__(self, name):
        self.name = name

    def __call__(self, Z):
        with tf.variable_scope(name_or_scope=self.name, reuse=False):
            with tf.variable_scope(name_or_scope="linear"):
                inputs = tf.reshape(tf.nn.relu((fully_connected(Z, 4*4*512))), [batchsize, 4, 4, 512])
            with tf.variable_scope(name_or_scope="deconv1"):
                inputs = tf.nn.relu(instanceNorm(deconv(inputs, [5, 5, 256, 512], [1, 2, 2, 1], [batchsize, 8, 8, 256])))
            with tf.variable_scope(name_or_scope="deconv2"):
                inputs = tf.nn.relu(instanceNorm(deconv(inputs, [5, 5, 128, 256], [1, 2, 2, 1], [batchsize, 16, 16, 128])))
            with tf.variable_scope(name_or_scope="deconv3"):
                inputs = tf.nn.relu(instanceNorm(deconv(inputs, [5, 5, 64, 128], [1, 2, 2, 1], [batchsize, 32, 32, 64])))
            if img_H == 32:
                with tf.variable_scope(name_or_scope="deconv4"):
                    inputs = tf.nn.tanh(deconv(inputs, [5, 5, img_C, 64], [1, 1, 1, 1], [batchsize, img_H, img_W, img_C]))
                return inputs
            with tf.variable_scope(name_or_scope="deconv4"):
                 inputs = tf.nn.tanh(deconv(inputs, [5, 5, img_C, 64], [1, 2, 2, 1], [batchsize, img_H, img_W, img_C]))
            if img_H == 64:
                return inputs

    @property
    def var(self):
        return tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, self.name)

class Discriminator:
    def __init__(self, name):
        self.name = name

    def __call__(self, inputs, reuse=False, is_sn=False):
        with tf.variable_scope(name_or_scope=self.name, reuse=reuse):
            with tf.variable_scope("conv1"):
                inputs = leaky_relu(conv(inputs, [5, 5, img_C, 64], [1, 2, 2, 1], is_sn))
            with tf.variable_scope("conv2"):
                inputs = leaky_relu(instanceNorm(conv(inputs, [5, 5, 64, 128], [1, 2, 2, 1], is_sn)))
            with tf.variable_scope("conv3"):
                inputs = leaky_relu(instanceNorm(conv(inputs, [5, 5, 128, 256], [1, 2, 2, 1], is_sn)))
            with tf.variable_scope("conv4"):
                inputs = leaky_relu(instanceNorm(conv(inputs, [5, 5, 256, 512], [1, 2, 2, 1], is_sn)))
            inputs = tf.layers.flatten(inputs)
            return fully_connected(inputs, 1, is_sn)

    @property
    def var(self):
        return tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=self.name)

class GAN:
    #Architecture of generator and discriminator just like DCGAN.
    def __init__(self):
        self.Z = tf.placeholder("float", [batchsize, 100])
        self.img = tf.placeholder("float", [batchsize, img_H, img_W, img_C])
        D = Discriminator("discriminator")
        G = Generator("generator")
        self.fake_img = G(self.Z)
        if GAN_type == "DCGAN":
            #DCGAN, paper: UNSUPERVISED REPRESENTATION LEARNING WITH DEEP CONVOLUTIONAL GENERATIVE ADVERSARIAL NETWORKS
            self.fake_logit = tf.nn.sigmoid(D(self.fake_img))
            self.real_logit = tf.nn.sigmoid(D(self.img, reuse=True))
            self.d_loss = - (tf.reduce_mean(tf.log(self.real_logit + epsilon)) + tf.reduce_mean(tf.log(1 - self.fake_logit + epsilon)))
            self.g_loss = - tf.reduce_mean(tf.log(self.fake_logit + epsilon))
            self.opt_D = tf.train.AdamOptimizer(2e-4, beta1=0.5).minimize(self.d_loss, var_list=D.var)
            self.opt_G = tf.train.AdamOptimizer(2e-4, beta1=0.5).minimize(self.g_loss, var_list=G.var)
        elif GAN_type == "WGAN":
            #WGAN, paper: Wasserstein GAN
            self.fake_logit = D(self.fake_img)
            self.real_logit = D(self.img, reuse=True)
            self.d_loss = -tf.reduce_mean(self.real_logit) + tf.reduce_mean(self.fake_logit)
            self.g_loss = -tf.reduce_mean(self.fake_logit)
            self.clip = []
            for _, var in enumerate(D.var):
                self.clip.append(var.assign(tf.clip_by_value(var, -0.01, 0.01)))
            self.opt_D = tf.train.RMSPropOptimizer(5e-5).minimize(self.d_loss, var_list=D.var)
            self.opt_G = tf.train.RMSPropOptimizer(5e-5).minimize(self.g_loss, var_list=G.var)
        elif GAN_type == "WGAN-GP":
            #WGAN-GP, paper: Improved Training of Wasserstein GANs
            self.fake_logit = D(self.fake_img)
            self.real_logit = D(self.img, reuse=True)
            e = tf.random_uniform([batchsize, 1, 1, 1], 0, 1)
            x_hat = e * self.img + (1 - e) * self.fake_img
            grad = tf.gradients(D(x_hat, reuse=True), x_hat)[0]
            self.d_loss = tf.reduce_mean(self.fake_logit - self.real_logit) + 10 * tf.reduce_mean(tf.square(tf.sqrt(tf.reduce_sum(tf.square(grad), axis=[1, 2, 3])) - 1))
            self.g_loss = tf.reduce_mean(-self.fake_logit)
            self.opt_D = tf.train.AdamOptimizer(1e-4, beta1=0., beta2=0.9).minimize(self.d_loss, var_list=D.var)
            self.opt_G = tf.train.AdamOptimizer(1e-4, beta1=0., beta2=0.9).minimize(self.g_loss, var_list=G.var)
        elif GAN_type == "LSGAN":
            #LSGAN, paper: Least Squares Generative Adversarial Networks
            self.fake_logit = D(self.fake_img)
            self.real_logit = D(self.img, reuse=True)
            self.d_loss = tf.reduce_mean(0.5 * tf.square(self.real_logit - 1) + 0.5 * tf.square(self.fake_logit))
            self.g_loss = tf.reduce_mean(0.5 * tf.square(self.fake_logit - 1))
            self.opt_D = tf.train.AdamOptimizer(2e-4, beta1=0.5).minimize(self.d_loss, var_list=D.var)
            self.opt_G = tf.train.AdamOptimizer(2e-4, beta1=0.5).minimize(self.g_loss, var_list=G.var)
        elif GAN_type == "SNGAN":
            #SNGAN, paper: SPECTRAL NORMALIZATION FOR GENERATIVE ADVERSARIAL NETWORKS
            self.fake_logit = tf.nn.sigmoid(D(self.fake_img, is_sn=True))
            self.real_logit = tf.nn.sigmoid(D(self.img, reuse=True, is_sn=True))
            self.d_loss = - (tf.reduce_mean(tf.log(self.real_logit + epsilon) + tf.log(1 - self.fake_logit + epsilon)))
            self.g_loss = - tf.reduce_mean(tf.log(self.fake_logit + epsilon))
            self.opt_D = tf.train.AdamOptimizer(2e-4, beta1=0.5).minimize(self.d_loss, var_list=D.var)
            self.opt_G = tf.train.AdamOptimizer(2e-4, beta1=0.5).minimize(self.g_loss, var_list=G.var)
        elif GAN_type == "RSGAN":
            #RSGAN, paper: The relativistic discriminator: a key element missing from standard GAN
            self.fake_logit = D(self.fake_img)
            self.real_logit = D(self.img, reuse=True)
            self.d_loss = - tf.reduce_mean(tf.log(tf.nn.sigmoid(self.real_logit - self.fake_logit) + epsilon))
            self.g_loss = - tf.reduce_mean(tf.log(tf.nn.sigmoid(self.fake_logit - self.real_logit) + epsilon))
            self.opt_D = tf.train.AdamOptimizer(2e-4, beta1=0.5).minimize(self.d_loss, var_list=D.var)
            self.opt_G = tf.train.AdamOptimizer(2e-4, beta1=0.5).minimize(self.g_loss, var_list=G.var)
        elif GAN_type == "RaSGAN":
            #RaSGAN, paper: The relativistic discriminator: a key element missing from standard GAN
            self.fake_logit = D(self.fake_img)
            self.real_logit = D(self.img, reuse=True)
            self.avg_fake_logit = tf.reduce_mean(self.fake_logit)
            self.avg_real_logit = tf.reduce_mean(self.real_logit)
            self.D_r_tilde = tf.nn.sigmoid(self.real_logit - self.avg_fake_logit)
            self.D_f_tilde = tf.nn.sigmoid(self.fake_logit - self.avg_real_logit)
            self.d_loss = - tf.reduce_mean(tf.log(self.D_r_tilde + epsilon)) - tf.reduce_mean(tf.log(1 - self.D_f_tilde + epsilon))
            self.g_loss = - tf.reduce_mean(tf.log(self.D_f_tilde + epsilon)) - tf.reduce_mean(tf.log(1 - self.D_r_tilde + epsilon))
            self.opt_D = tf.train.AdamOptimizer(2e-4, beta1=0.5).minimize(self.d_loss, var_list=D.var)
            self.opt_G = tf.train.AdamOptimizer(2e-4, beta1=0.5).minimize(self.g_loss, var_list=G.var)
        self.sess = tf.Session()
        self.sess.run(tf.global_variables_initializer())

    def __call__(self):
        saver = tf.train.Saver()
        epoch_nums = 200
        facedata = sio.loadmat("./TrainingSet//facedata.mat")["data"]
        #For face data, i random select about 10,000 images from CelebA and resize them to 64x64 by Matlab.
        for epoch in range(epoch_nums):
            for i in range(facedata.__len__()//batchsize-1):
                batch = facedata[i*batchsize:i*batchsize+batchsize, :, :, :] / 255.0
                z = np.random.standard_normal([batchsize, 100])
                d_loss = self.sess.run(self.d_loss, feed_dict={self.img: batch, self.Z: z})
                g_loss = self.sess.run(self.g_loss, feed_dict={self.img: batch, self.Z: z})
                self.sess.run(self.opt_D, feed_dict={self.img: batch, self.Z: z})
                if GAN_type == "WGAN":
                    self.sess.run(self.clip)#WGAN weight clipping
                self.sess.run(self.opt_G, feed_dict={self.img: batch, self.Z: z})
                if i % 10 == 0:
                    print("epoch: %d, step: %d, d_loss: %g, g_loss: %g"%(epoch, i,  d_loss, g_loss))
                    z = np.random.standard_normal([batchsize, 100])
                    imgs = self.sess.run(self.fake_img, feed_dict={self.img: batch, self.Z: z})
                    for j in range(batchsize):
                        if img_C == 1:
                            Image.fromarray(np.reshape(np.uint8(mapping(imgs[j, :, :, :])), [img_H, img_W])).save(
                                "./result//" + str(epoch) + "_" + str(j) + ".jpg")
                        else:
                            Image.fromarray(np.uint8(mapping(imgs[j, :, :, :]))).save(
                                "./result//" + str(epoch) + "_" + str(j) + ".jpg")
            saver.save(self.sess, "./para//model.ckpt")

if __name__ == "__main__":
    gan = GAN()
    gan()


import os
import psi4
import itertools
import subprocess
import numpy as np

class sapt(object):

    def __init__(self, geometries, data, method, basis, outfile):
        self.geometries = geometries
        self.method = method
        self.basis = basis
        self.outfile = outfile
        self.frag_A = data.frag_A
        self.frag_B = data.frag_B
        self.do_fsapt = data.do_fsapt
        self.keywords = data.keywords
        self.step_size = data.irc_stepsize
        self.coordinates = data.coordinates
        self.force_min = data.force_min
        self.nthreads = data.nthreads
        self.set_memory = data.set_memory
        self.memory_allocation = data.memory_allocation
        # Specifics for F-SAPT Calculations
        if(data.do_fsapt):
            self.monomer_A_frags = data.monomer_A_frags
            self.monomer_B_frags = data.monomer_B_frags
            self.monomer_A_labels = data.monomer_A_labels
            self.monomer_B_labels = data.monomer_B_labels

    def psi4_sapt(self):
        self.int_  = []
        self.elst_ = []
        self.exch_ = []
        self.ind_ =  []
        self.disp_ = []
        all_atoms = self.frag_A + self.frag_B
        au_to_kcal = 627.51
        count = 0
        output = open(self.outfile, "a")
        output.write('\n\n--SAPT Decomposition--\n')
        output.write('\n-------------------------------------------------------------------------------------------------')
        output.write('\n{:>15} {:>15} {:>15} {:>15} {:>15} {:>15}\n'.format('IRC Point', 'E_int(kcal)', 'E_elst(kcal)','E_exch(kcal)', 'E_ind(kcal)', 'E_disp(kcal)'))
        output.write('-------------------------------------------------------------------------------------------------\n')
        output.close()
        scale_sapt = False
        # See if the user is requesting a scaled SAPT0 calculation
        if(self.keywords['ssapt0_scale'] or self.keywords['SSAPT0_SCALE']):
            scale_sapt = True
        for geometry in self.geometries:
            output = open(self.outfile, "a")
            psi4.core.set_output_file("psi4_output/irc_%d_sapt.out" %count, False)
            geometry += "symmetry c1"
            psi4.geometry(geometry)
            # Adds fsapt capabilities
            if(self.do_fsapt):
                try:
                    os.mkdir('fsapt_output')
                except FileExistsError:
                    pass
                try:
                    os.mkdir('s-fsapt_output')
                except FileExistsError:
                    pass
                self.keywords['FISAPT_FSAPT_FILEPATH'] = 'fsapt_output/fsapt%d' %count
                if(scale_sapt):
                    self.keywords['FISAPT_FSSAPT_FILEPATH'] = 's-fsapt_output/fsapt%d' %count
            psi4.set_options(self.keywords)
            #print("pyREX:SAPT Calculation on IRC Point %d" %(count))
            psi4.set_num_threads(self.nthreads)
            if(self.set_memory):
                psi4.set_memory(self.memory_allocation)
            e_int = psi4.energy(self.method)
            #TODO: Actually have this create the SAPT files necessary for FSAPT partitioning. Make user options as well
            #      for fragment definitions. Something in the sapt block.
            if(self.do_fsapt):
                fsaptA_outfile = open('fsapt_output/fsapt%d/fA.dat' %count, 'w+')
                if(scale_sapt):
                    s_fsaptA_outfile = open('s-fsapt_output/fsapt%d/fA.dat' %count, 'w+')
                for i in range(len(self.monomer_A_labels)):
                    fsaptA_outfile.write("%s" %self.monomer_A_labels[i])
                    if(scale_sapt):
                        s_fsaptA_outfile.write("%s" %self.monomer_A_labels[i])
                    for j in range(len(self.monomer_A_frags[i])):
                        if(j==len(self.monomer_A_frags[i])-1):
                            fsaptA_outfile.write(" %d \n" %(all_atoms.index(self.monomer_A_frags[i][j])+1))
                            if(scale_sapt):
                                s_fsaptA_outfile.write(" %d \n" %(all_atoms.index(self.monomer_A_frags[i][j])+1))
                        else:
                            fsaptA_outfile.write(" %d " %(all_atoms.index(self.monomer_A_frags[i][j])+1))
                            if(scale_sapt):
                                s_fsaptA_outfile.write(" %d " %(all_atoms.index(self.monomer_A_frags[i][j])+1))
                fsaptA_outfile.close()
                s_fsaptA_outfile.close()
                fsaptB_outfile = open('fsapt_output/fsapt%d/fB.dat' %count, 'w+')
                if(scale_sapt):
                    s_fsaptB_outfile = open('s-fsapt_output/fsapt%d/fB.dat' %count, 'w+')
                for i in range(len(self.monomer_B_labels)):
                    fsaptB_outfile.write("%s" %self.monomer_B_labels[i])
                    if(scale_sapt):
                        s_fsaptB_outfile.write("%s" %self.monomer_B_labels[i])
                    for j in range(len(self.monomer_B_frags[i])):
                        if(j==len(self.monomer_B_frags[i])-1):
                            fsaptB_outfile.write(" %d \n" %(all_atoms.index(self.monomer_B_frags[i][j])+1))
                            if(scale_sapt):
                                s_fsaptB_outfile.write(" %d \n" %(all_atoms.index(self.monomer_B_frags[i][j])+1))
                        else:
                            fsaptB_outfile.write(" %d " %(all_atoms.index(self.monomer_B_frags[i][j])+1))
                            if(scale_sapt):
                                s_fsaptB_outfile.write(" %d " %(all_atoms.index(self.monomer_B_frags[i][j])+1))
                fsaptB_outfile.close()
                s_fsaptB_outfile.close()
            e_int =  psi4.core.variable("SAPT TOTAL ENERGY")
            e_elst = psi4.core.variable("SAPT ELST ENERGY")
            e_exch = psi4.core.variable("SAPT EXCH ENERGY")
            e_ind  = psi4.core.variable("SAPT IND ENERGY")
            e_disp = psi4.core.variable("SAPT DISP ENERGY")

            self.int_.append(e_int)
            self.elst_.append(e_elst)
            self.exch_.append(e_exch)
            self.ind_.append(e_ind)
            self.disp_.append(e_disp)
            output.write('\n{:>15} {:>15.4f} {:>15.4f} {:>15.4f} {:>15.4f} {:>15.4f}\n'.format(count, e_int*au_to_kcal, e_elst*au_to_kcal, e_exch*au_to_kcal, e_ind*au_to_kcal, e_disp*au_to_kcal))
            count = count+1
        output.write('-------------------------------------------------------------------------------------\n')
        output.close()
        return self.int_, self.elst_, self.exch_, self.ind_, self.disp_

    def psi4_super(self):
        self.int_  = []
        all_atoms = self.frag_A + self.frag_B
        au_to_kcal = 627.51
        count = 0
        output = open(self.outfile, "a")
        output.write('\n\n--Supermolecular Interaction Energy--\n')
        output.write('\n-------------------------------------------------------------------------------------------------')
        output.write('\n{:>15} {:>15}\n'.format('IRC Point', 'E_int(kcal)'))
        output.write('-------------------------------------------------------------------------------------------------\n')
        output.close()
        for geometry in self.geometries:
            output = open(self.outfile, "a")
            psi4.core.set_output_file("psi4_output/irc_%d_sapt.out" %count, False)
            geometry += "symmetry c1"
            psi4.geometry(geometry)
            psi4.set_options({'reference': 'rhf', 'basis' : self.basis})
            psi4.set_num_threads(self.nthreads)
            if(self.set_memory):
                psi4.set_memory(self.memory_allocation)
            #print("pyREX:SAPT Calculation on IRC Point %d" %(count))
            e_int = psi4.energy("%s/%s" %(self.method,self.basis), bsse_type='cp')
            output.write('\n{:>15} {:>15.4f}\n'.format(count, e_int*au_to_kcal))
            output.close()
            self.int_.append(e_int)
            count = count + 1
        return self.int_

    def fsapt_post_analysis(self):
        fsapt_script_dir = os.environ['PSIPATH'][:-1] + '/share/psi4/fsapt/fsapt.py'
        link_type = "50-50"
        step_size = self.step_size
        coordinates = self.coordinates
        force_min = self.force_min
        self.monomer_A_labels.append("All")
        self.monomer_B_labels.append("All")
        unique_pairs_list = list(itertools.product(self.monomer_A_labels,self.monomer_B_labels))
        unique_pairs = []
        for i in range(len(unique_pairs_list)):
            unique_pairs.append("%s-%s" %(unique_pairs_list[i][0],unique_pairs_list[i][1]))
        print(unique_pairs)
        # Initialize empty dictionaries to hold pair energy, force(f), and work(w) data
        pair_elst, pair_exch, pair_indab, pair_indba, pair_disp, pair_total = ({} for i in range(6))
        f_elst, f_exch, f_indab, f_indba, f_disp, f_total = ({} for i in range(6))
        w_elst_1, w_exch_1, w_indab_1, w_indba_1, w_disp_1, w_total_1 = ({} for i in range(6))
        w_elst_2, w_exch_2, w_indab_2, w_indba_2, w_disp_2, w_total_2 = ({} for i in range(6))

        # Initialize empty array for every pair property
        for pair in unique_pairs:
            pair_elst[pair], pair_exch[pair], pair_indab[pair], pair_indba[pair], pair_disp[pair], pair_total[pair] = ([] for i in range(6))
        os.chdir("fsapt_output")
        dirs_list = [name for name in os.listdir() if os.path.isdir(name)]
        num_geoms = 0
        for directory in dirs_list:
            if("fsapt" in directory):
                num_geoms = num_geoms + 1
        print(num_geoms)
        for i in range(num_geoms):
            fsapt_dir = "fsapt%d" %i
            os.chdir(fsapt_dir)
            print("running in %s" %fsapt_dir)
            subprocess.call(['python', fsapt_script_dir])
            if(link_type=="By Charge"):
                do_store = True
            else:
                do_store = False
            fsapt_file = open('fsapt.dat', 'r')
            for line in fsapt_file:
                if("Reduced Analysis" in line):
                    if(do_store):
                        for j in range(3):
                            line = next(fsapt_file)
                        for j in range(len(unique_pairs)):
                            pair_data = line.split()
                            pair = "%s-%s" %(pair_data[0],pair_data[1])
                            elst = pair_data[2]
                            exch = pair_data[3]
                            indab = pair_data[4]
                            indba = pair_data[5]
                            disp = pair_data[6]
                            total = pair_data[7]
                            pair_elst[pair].append(float(elst))
                            pair_exch[pair].append(float(exch))
                            pair_indab[pair].append(float(indab))
                            pair_indba[pair].append(float(indba))
                            pair_disp[pair].append(float(disp))
                            pair_total[pair].append(float(total))
                            #print(pair_elst)
                            line = next(fsapt_file)
                    if(link_type=="By Charge"):
                        do_store = False
                    else:
                        do_store = True
            os.chdir("..")
        pair_dicts = [pair_elst, pair_exch, pair_indab, pair_indba, pair_disp, pair_total]
        f_dicts = [f_elst, f_exch, f_indab, f_indba, f_disp, f_total]
        w_dicts_1 = [w_elst_1, w_exch_1, w_indab_1, w_indba_1, w_disp_1, w_total_1]
        w_dicts_2 = [w_elst_2, w_exch_2, w_indab_2, w_indba_2, w_disp_2, w_total_2]

        for dict_ in pair_dicts:
            index = pair_dicts.index(dict_)
            print(index)
            for pair in unique_pairs:
                  f_dicts[index][pair] = -1.0*np.gradient(dict_[pair],step_size)

        # Store energy data in .csv files
        elst_csv = open("fsapt_elst.csv", "w+")
        exch_csv = open("fsapt_exch.csv", "w+")
        indab_csv = open("fsapt_indab.csv", "w+")
        indba_csv = open("fsapt_indba.csv", "w+")
        disp_csv = open("fsapt_disp.csv", "w+")
        total_csv = open("fsapt_total.csv", "w+")

        csv_files = [elst_csv,exch_csv,indab_csv,indba_csv,disp_csv,total_csv]

        for csv_f in csv_files:
            index = csv_files.index(csv_f)
            prop_dict = pair_dicts[index]
            csv_f.write("Coordinate,")
            for pair in unique_pairs:
                if(pair==unique_pairs[-1]):
                    csv_f.write("%s,\n" %pair)
                else:
                    csv_f.write("%s," %pair)
            for i in range(num_geoms):
                csv_f.write("%f," %coordinates[i])
                for pair in unique_pairs:
                    if(pair==unique_pairs[-1]):
                        csv_f.write("%.3f,\n" %prop_dict[pair][i])
                    else:
                        csv_f.write("%.3f," %prop_dict[pair][i])
        elst_csv.close(), exch_csv.close(), indab_csv.close(), indba_csv.close(), disp_csv.close(), total_csv.close()

        # Store force data in .csv files
        f_elst_csv = open("force_elst.csv", "w+")
        f_exch_csv = open("force_exch.csv", "w+")
        f_indab_csv = open("force_indab.csv", "w+")
        f_indba_csv = open("force_indba.csv", "w+")
        f_disp_csv = open("force_disp.csv", "w+")
        f_total_csv = open("force_total.csv", "w+")

        f_csv_files = [f_elst_csv,f_exch_csv,f_indab_csv,f_indba_csv,f_disp_csv,f_total_csv]

        for csv_f in f_csv_files:
            index = f_csv_files.index(csv_f)
            prop_dict = f_dicts[index]
            csv_f.write("Coordinate,")
            for pair in unique_pairs:
                if(pair==unique_pairs[-1]):
                    csv_f.write("%s,\n" %pair)
                else:
                    csv_f.write("%s," %pair)
            for i in range(num_geoms):
                csv_f.write("%f," %coordinates[i])
                for pair in unique_pairs:
                    if(pair==unique_pairs[-1]):
                        csv_f.write("%.3f,\n" %prop_dict[pair][i])
                    else:
                        csv_f.write("%.3f," %prop_dict[pair][i])
        f_elst_csv.close(), f_exch_csv.close(), f_indab_csv.close(), f_indba_csv.close(), f_disp_csv.close(), f_total_csv.close()

        # Calculate Work Contributions for each Pair
        index_min = coordinates.index(force_min)
        f_dict_count = 0
        for dict_ in f_dicts:
            for pair in unique_pairs:
                  w_dicts_1[f_dict_count][pair] = -1.0*np.trapz(dict_[pair][:index_min],dx=step_size)
            f_dict_count = f_dict_count + 1
        f_dict_count = 0
        for dict_ in f_dicts:
            for pair in unique_pairs:
                  w_dicts_2[f_dict_count][pair] = -1.0*np.trapz(dict_[pair][index_min-1:],dx=step_size)
            f_dict_count = f_dict_count + 1

        work_values = open("work_values.dat", "w+")
        work_values.write('\n\n--Reaction Work Decomposition (Region 1)--\n')
        work_values.write('\n-----------------------------------------------------------------------------------------------------------------')
        work_values.write('\n{:>15} {:>15} {:>15} {:>15} {:>15} {:>15} {:>15}\n'.format('Pair(A-B)','W_elst','W_exch',
        'W_indAB', 'W_indBA', 'W_disp', 'W_total'))
        work_values.write('-----------------------------------------------------------------------------------------------------------------\n')
        for pair in unique_pairs:
            work_values.write('\n{:>15s} {:>15.5f} {:>15.5f} {:>15.5f} {:>15.5f} {:>15.5f} {:>15.5f}\n'.format(pair,w_dicts_1[0][pair],w_dicts_1[1][pair], w_dicts_1[2][pair], w_dicts_1[3][pair], w_dicts_1[4][pair], w_dicts_1[5][pair]))
        work_values.write('------------------------------------------------------------------------------------------------------------------\n')
        work_values.write('\n\n--Reaction Work Decomposition (Region 2)--\n')
        work_values.write('\n-----------------------------------------------------------------------------------------------------------------')
        work_values.write('\n{:>15} {:>15} {:>15} {:>15} {:>15} {:>15} {:>15}\n'.format('Pair(A-B)','W_elst','W_exch', 'W_indAB', 'W_indBA', 'W_disp', 'W_total'))
        work_values.write('-----------------------------------------------------------------------------------------------------------------\n')
            #print(w_dicts[0])
        for pair in unique_pairs:
            work_values.write('\n{:>15s} {:>15.5f} {:>15.5f} {:>15.5f} {:>15.5f} {:>15.5f} {:>15.5f}\n'.format(pair,w_dicts_2[0][pair],w_dicts_2[1][pair], w_dicts_2[2][pair], w_dicts_2[3][pair], w_dicts_2[4][pair], w_dicts_2[5][pair]))
        work_values.write('------------------------------------------------------------------------------------------------------------------\n')
        work_values.close()