import numpy as np
import tensorflow as tf

from tensorflow.keras import backend as K
from tensorflow.keras.layers import Layer

class ScheduledDropout(Layer):
    """Applies Scheduled Dropout to the input.
        The Dropout layer randomly sets input units to 0 with a frequency of `rate`
        scheduled by network layer's depth and training step at each step, which
        helps prevent overfitting.
        Inputs not set to 0 are scaled up by 1/(1 - rate) such that the sum over
        all inputs is unchanged.
        Note that the Dropout layer only applies when `training` is set to True
        such that no values are dropped during inference. When using `model.fit`,
        `training` will be appropriately set to True automatically, and in other
        contexts, you can set the kwarg explicitly to True when calling the layer.
        (This is in contrast to setting `trainable=False` for a Dropout layer.
        `trainable` does not affect the layer's behavior, as Dropout does
        not have any variables/weights that can be frozen during training.)
        Arguments:
            drop_rate: Float between 0 and 1. Fraction of the input units to drop.
            cell_num: Cell number in the network
            total_num_cells: Number of cells in the network
            total_training_steps: Number of total steps performed during training
            seed: A Python integer to use as random seed.
        Call arguments:
            inputs: Input tensor (of any rank).
            training: Python boolean indicating whether the layer should behave in
                training mode (adding dropout) or in inference mode (doing nothing).
    """

    def __init__(self, drop_rate, cell_num, total_num_cells, total_training_steps, seed=None, **kwargs):
        super(ScheduledDropout, self).__init__(**kwargs)
        self.drop_rate = drop_rate
        self._cell_num = cell_num
        self._total_num_cells = total_num_cells
        self._total_training_steps = total_training_steps
        self.seed = seed

    def call(self, inputs, training=None):
        if training is None:
            training = K.learning_phase()
        scheduled_drop_rate = self._compute_scheduled_dropout_rate()

        def dropped_inputs():
            noise_shape = tf.shape(inputs)
            random_tensor = 1 - scheduled_drop_rate
            random_tensor += tf.random.uniform(noise_shape, dtype=tf.float32)
            binary_tensor = tf.cast(tf.floor(random_tensor), inputs.dtype)
            keep_prob_inv = tf.cast(1.0 / (1-scheduled_drop_rate), inputs.dtype)
            outputs = inputs * keep_prob_inv * binary_tensor
            return outputs

        if training:
            output = dropped_inputs()
        else:
            output = tf.identity(inputs)
        return output

    def _compute_scheduled_dropout_rate(self):
        drop_rate = self.drop_rate
        if self._total_num_cells is not None:
            # Scale keep prob by layer number
            assert self._cell_num != -1
            # The added 2 is for the reduction cells
            num_cells = self._total_num_cells
            layer_ratio = (self._cell_num + 1) / float(num_cells)
            drop_rate = layer_ratio * drop_rate
        if self._total_training_steps is not None:
            # Decrease the keep probability over time
            current_step = tf.convert_to_tensor(tf.compat.v1.train.get_or_create_global_step())
            tf.compat.v1.get_variable_scope().reuse_variables()
            current_step = tf.cast(current_step, tf.float32)
            drop_path_burn_in_steps = self._total_training_steps
            current_ratio = current_step / drop_path_burn_in_steps
            current_ratio = tf.minimum(1.0, current_ratio)
            drop_rate = current_ratio * drop_rate
        return drop_rate

    def compute_output_shape(self, input_shape):
        return input_shape

    def get_config(self):
        config = {
            'drop_rate': self.drop_rate,
            'cell_num': self._cell_num,
            'total_num_cells': self._total_num_cells,
            'total_training_steps': self._total_training_steps,
            'seed': self.seed
        }
        base_config = super(ScheduledDropout, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class ScheduledDroppath(Layer):
    """Applies Scheduled Droppath to the input.
        The Droppath layer randomly sets whole input path inside to 0 with a
        frequency of `rate` scheduled by network layer's depth and training
        step at each step, which helps prevent overfitting.
        Inputs not set to 0 are scaled up by 1/(1 - rate) such that the sum over
        all inputs is unchanged.
        Note that the Scheduled Droppath layer only applies when `training` is set to True.
        When using `model.fit`, `training` will be appropriately set to True
        automatically, and in other contexts, you can set the kwarg explicitly
        to True when calling the layer. (This is in contrast to setting
        `trainable=False` for a Droppath layer. `trainable` does not affect the
        layer's behavior, as Droppath does not have any variables/weights that
        can be frozen during training.)
        Arguments:
            drop_rate: Float between 0 and 1. Fraction of the inputs to drop.
            cell_num: Cell number in the network
            total_num_cells: Number of cells in the network
            total_training_steps: Number of total steps performed during training
            seed: A Python integer to use as random seed.
        Call arguments:
            inputs: Input tensor (of any rank).
            training: Python boolean indicating whether the layer should behave in
                training mode (adding dropout) or in inference mode (doing nothing).
    """

    def __init__(self, drop_rate, cell_num, total_num_cells, total_training_steps, seed=None, **kwargs):
        super(ScheduledDroppath, self).__init__(**kwargs)
        self.drop_rate = drop_rate
        self._cell_num = cell_num
        self._total_num_cells = total_num_cells
        self._total_training_steps = total_training_steps
        self.seed = seed

    def call(self, inputs, training=None):
        if training is None:
            training = K.learning_phase()
        scheduled_drop_rate = self._compute_scheduled_drop_rate()

        def dropped_inputs():
            noise_shape = [tf.shape(input=inputs)[0], 1, 1, 1]
            random_tensor = 1 - scheduled_drop_rate
            random_tensor += tf.random.uniform(noise_shape, dtype=tf.float32)
            binary_tensor = tf.cast(tf.floor(random_tensor), inputs.dtype)
            keep_prob_inv = tf.cast(1.0 / (1-scheduled_drop_rate), inputs.dtype)
            tf.print('binary tesnor:', binary_tensor)
            outputs = inputs * keep_prob_inv * binary_tensor
            return outputs

        if training:
            output = dropped_inputs()
        else:
            output = tf.identity(inputs)
        return output

    def _compute_scheduled_drop_rate(self):
        drop_rate = self.drop_rate
        if self._total_num_cells is not None:
            # Scale keep prob by layer number
            assert self._cell_num != -1
            # The added 2 is for the reduction cells
            num_cells = self._total_num_cells
            layer_ratio = (self._cell_num + 1) / float(num_cells)
            drop_rate = layer_ratio * drop_rate
        if self._total_training_steps is not None:
            # Decrease the keep probability over time
            current_step = tf.convert_to_tensor(tf.compat.v1.train.get_or_create_global_step())
            tf.compat.v1.get_variable_scope().reuse_variables()
            current_step = tf.cast(current_step, tf.float32)
            drop_path_burn_in_steps = self._total_training_steps
            current_ratio = current_step / drop_path_burn_in_steps
            current_ratio = tf.minimum(1.0, current_ratio)
            drop_rate = current_ratio * drop_rate
        return drop_rate

    def compute_output_shape(self, input_shape):
        return input_shape

    def get_config(self):
        config = {
            'drop_rate': self.drop_rate,
            'cell_num': self._cell_num,
            'total_num_cells': self._total_num_cells,
            'total_training_steps': self._total_training_steps,
            'seed': self.seed
        }
        base_config = super(ScheduledDroppath, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class ConcreteDropout(Layer):
    """Applies Concrete Dropout to the input.
        The Concrete Droppath layer randomly sets input path to 0 with a
        frequency considered as a weight of the layer optimized during training
        time, which helps prevent overfitting.
        Inputs not set to 0 are scaled up by 1/(1 - rate) such that the sum over
        all inputs is unchanged.
        Note that the Concrete Dropout layer only applies when `training` is set
        to True. When using `model.fit`, `training` will be appropriately set to
        True automatically, and in other contexts, you can set the kwarg explicitly
        to True when calling the layer. (This is in contrast to setting
        `trainable=False` for a Concrete Dropout layer. `trainable` does not affect
        the layer's behavior, as Dropout does not have any variables/weights that
        can be frozen during training.)
        Arguments:
            dropout_regularizer: A positive number which satisfies
                $dropout_regularizer = 2 / (\tau * N)$ with model precision
                $\tau$ (inverse observation noise) and N the number of
                instances in the dataset.
            init_min: dropout probability initializer min
            init_max: dropout probability initializer max
            seed: A Python integer to use as random seed.
        Call arguments:
            inputs: Input tensor (of any rank).
            training: Python boolean indicating whether the layer should behave in
                training mode (adding dropout) or in inference mode (doing nothing).
    """

    def __init__(self, dropout_regularizer=1e-5, init_min=0.1, init_max=0.1,
                 seed=None, **kwargs):
        super(ConcreteDropout, self).__init__(**kwargs)
        self.dropout_regularizer = dropout_regularizer
        self.init_min = tf.math.log(init_min) - tf.math.log(1. - init_min)
        self.init_max = tf.math.log(init_max) - tf.math.log(1. - init_max)
        self.p_logit = None
        self.seed = seed

    @tf.function
    def get_p(self):
        return tf.nn.sigmoid(self.p_logit[0])

    def build(self, input_shape=None):
        super(ConcreteDropout, self).build(input_shape)
        # initialise p
        self.p_logit = self.add_weight(name='p_logit',
                                       shape=(1,),
                                       initializer=tf.initializers.RandomUniform(self.init_min, self.init_max),
                                       trainable=True)
        # initialise regulariser / prior KL term
        input_dim = np.prod(input_shape[1:])
        dropout_regularizer = self.get_p() * K.log(self.get_p())
        dropout_regularizer += (1. - self.get_p()) * K.log(1. - self.get_p())
        dropout_regularizer *= self.dropout_regularizer * input_dim
        regularizer = dropout_regularizer
        self.add_loss(regularizer)

    def call(self, inputs, training=None):
        if training:
            return self.concrete_dropout(inputs)
        else:
            tf.identity(inputs)

    def concrete_dropout(self, x):
        '''
        Concrete dropout - used at training time and testing time (gradients can be propagated)
        :param x: input
        :return:  approx. dropped out input
        '''

        eps = K.cast_to_floatx(K.epsilon())
        temp = 0.1
        unif_noise = K.random_uniform(K.shape(x))
        drop_prob = (
            K.log(self.get_p() + eps)
            - K.log(1. - self.get_p() + eps)
            + K.log(unif_noise + eps)
            - K.log(1. - unif_noise + eps)
        )
        drop_prob = K.sigmoid(drop_prob / temp)
        random_tensor = 1. - drop_prob
        retain_prob = 1. - self.get_p()
        x *= random_tensor
        x /= retain_prob
        return x

    def compute_output_shape(self, input_shape):
        return input_shape

    def get_config(self):
        config = {
            'dropout_regularizer': self.dropout_regularizer,
            'init_min': self.init_min.numpy(),
            'init_max': self.init_max.numpy(),
            'p_logit': self.p_logit.numpy(),
            'seed': self.seed
        }
        base_config = super(ConcreteDropout, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class ConcreteDroppath(Layer):
    """Applies Concrete Droppath to the input.
        The Concrete Droppath layer randomly sets input path to 0 with a
        frequency considered as a weight of the layer optimized during training
        time, which helps prevent overfitting.
        Inputs not set to 0 are scaled up by 1/(1 - rate) such that the sum over
        all inputs is unchanged.
        Note that the Concrete Droppath layer only applies when `training` is set
        to True. When using `model.fit`, `training` will be appropriately set to
        True automatically, and in other contexts, you can set the kwarg explicitly
        to True when calling the layer. (This is in contrast to setting
        `trainable=False` for a Concrete Droppath layer. `trainable` does not affect
        the layer's behavior, as Dropout does not have any variables/weights that
        can be frozen during training.)
        Arguments:
            dropout_regularizer: A positive number which satisfies
                $dropout_regularizer = 2 / (\tau * N)$ with model precision
                $\tau$ (inverse observation noise) and N the number of
                instances in the dataset.
            init_min: dropout probability initializer min
            init_max: dropout probability initializer max
            seed: A Python integer to use as random seed.
        Call arguments:
            inputs: Input tensor (of any rank).
            training: Python boolean indicating whether the layer should behave in
                training mode (adding dropout) or in inference mode (doing nothing).
    """

    def __init__(self, dropout_regularizer=1e-5, init_min=0.1, init_max=0.1,
                 seed=None, **kwargs):
        super(ConcreteDroppath, self).__init__(**kwargs)
        self.dropout_regularizer = dropout_regularizer
        self.init_min = tf.math.log(init_min) - tf.math.log(1. - init_min)
        self.init_max = tf.math.log(init_max) - tf.math.log(1. - init_max)
        self.p_logit = None
        self.seed = seed

    @tf.function
    def get_p(self):
        return tf.nn.sigmoid(self.p_logit[0])

    def build(self, input_shape=None):
        super(ConcreteDroppath, self).build(input_shape)
        # initialise p
        self.p_logit = self.add_weight(name='p_logit',
                                       shape=(1,),
                                       initializer=tf.initializers.RandomUniform(self.init_min, self.init_max),
                                       trainable=True)
        # initialise regulariser / prior KL term
        input_dim = 1
        dropout_regularizer = self.get_p() * K.log(self.get_p())
        dropout_regularizer += (1. - self.get_p()) * K.log(1. - self.get_p())
        dropout_regularizer *= self.dropout_regularizer * input_dim
        regularizer = dropout_regularizer
        self.add_loss(regularizer)

    def call(self, inputs, training=None):
        if training:
            return self.concrete_droppath(inputs)
        else:
            tf.identity(inputs)

    def concrete_droppath(self, x):
        """
        Concrete droppath - used at training and testing time (gradients can be propagated)
        :param x: input
        :return:  approx. dropped out input
        """

        eps = K.cast_to_floatx(K.epsilon())
        temp = 0.1
        unif_noise = tf.random.uniform(shape=[K.shape(x)[0], 1, 1, 1])
        drop_prob = (
            K.log(self.get_p() + eps)
            - K.log(1. - self.get_p() + eps)
            + K.log(unif_noise + eps)
            - K.log(1. - unif_noise + eps)
        )
        drop_prob = K.sigmoid(drop_prob / temp)
        random_tensor = 1. - drop_prob
        retain_prob = 1. - self.get_p()
        x *= random_tensor
        x /= retain_prob
        return x

    def compute_output_shape(self, input_shape):
        return input_shape

    def get_config(self):
        config = {
            'dropout_regularizer': self.dropout_regularizer,
            'init_min': self.init_min.numpy(),
            'init_max': self.init_max.numpy(),
            'p_logit': self.p_logit.numpy(),
            'seed': self.seed
        }
        base_config = super(ConcreteDroppath, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


"""
author : JAWAD FAYAZ (email: jfayaz@uci.edu) (website: https://jfayaz.github.io)

------------------------------ Instructions -------------------------------------
This code develops the RotD50 Sa and RotD100 Sa Spectra of the Bi-Directional
Ground Motion records as '.AT2' files provided in the current directory

The two directions of the ground motion record must be named as 'GM1i' and 'GM2i',
where 'i' is the ground motion number which goes from 1 to 'n', 'n' being the total
number of ground motions for which the Spectra needs to be generated. The extension
of the files must be '.AT2'

For example: If the Spectra of two ground motion records are required, 4 files with
the following names must be provided in the given 'GM' folder:
    'GM11.AT2' - Ground Motion 1 in direction 1 (direction 1 can be either one of the bi-directional GM as we are rotating the ground motions it does not matter)
    'GM21.AT2' - Ground Motion 1 in direction 2 (direction 2 is the other direction of the bi-directional GM)
    'GM12.AT2' - Ground Motion 2 in direction 1 (direction 1 can be either one of the bi-directional GM as we are rotating the ground motions it does not matter)
    'GM22.AT2' - Ground Motion 2 in direction 2 (direction 2 is the other direction of the bi-directional GM)

The Ground Motion file must be a vector file with 4 header lines.The first 3 lines can have
any content, however, the 4th header line must be written exactly as per the following example:
    'NPTS=  15864, DT= 0.0050'
The 'ReadGMFile.py' can be edited accordingly  for any other format

You may run this code in python IDE: 'Spyder' or any other similar IDE

Make sure you have the following python libraries installed:
    os
    sys
    pathlib
    fnmatch
    shutil
    IPython
    pandas
    numpy
    matplotlib.pyplot

INPUT:
This codes provides the option to have 3 different regions of developing the Spectra of ground motions with different period intervals (discretizations)
The following inputs within the code are required:
    'Path_to_openpyfiles'--> Path where the library files 'opensees.pyd' and 'LICENSE.rst' of OpenSeesPy are included (for further details go to https://openseespydoc.readthedocs.io/en/latest/windows.html)
    'Int_T_Reg_1'        --> Period Interval for the first region of the Spectrum
    'End_T_Reg_1'        --> Last Period of the first region of the Spectrum (where to end the first region)
    'Int_T_Reg_2'        --> Period Interval for the second region of the Spectrum
    'End_T_Reg_2'        --> Last Period of the second region of the Spectrum (where to end the second region)
    'Int_T_Reg_3'        --> Period Interval for the third region of the Spectrum
    'End_T_Reg_3'        --> Last Period of the third region of the Spectrum (where to end the third region)
    'Plot_Spectra'       --> whether to plot the generated Spectra of the ground motions (options: 'Yes', 'No')

OUTPUT:
The output will be provided in a saperate 'GMi_Spectra.txt' file for each ground motion record, where 'i' denotes the number of ground motion in the same of
provided 'GM1i.AT2' and 'GM2i.AT2' files. The output files will be generated in a saperate folder 'Spectra' which will be created in the current folder
The 'GMi_Spectra.txt' file will consist of space-saperated file with:
    'Periods (secs)' 'RotD50 Sa (g)' 'RotD100 Sa (g)'

%%%%% ========================================================================================================================================================================= %%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

"""

##### ================== INPUTS  ================== #####

# Path where the library files 'opensees.pyd' and 'LICENSE.rst' are included (for further details go to https://openseespydoc.readthedocs.io/en/latest/windows.html)
Path_to_openpyfiles = 'C:\Tcl'

# For periods 0 to 'End_T_Reg_1' in an interval of 'Int_T_Reg_1'
Int_T_Reg_1       = 0.1
End_T_Reg_1       = 1

# For periods ['End_T_Reg_1'+'Int_T_Reg_2'] to 'End_T_Reg_2' in an interval of 'Int_T_Reg_2'
Int_T_Reg_2       = 0.2
End_T_Reg_2       = 2

# For periods ['End_T_Reg_2'+'Int_T_Reg_3'] to 'End_T_Reg_3' in an interval of 'Int_T_Reg_3'
Int_T_Reg_3       = 0.5
End_T_Reg_3       = 5

# Plot Spectra  (options: 'Yes' or 'No')
Plot_Spectra      = 'Yes'

##### =============== CODE BEGINS ================ #######
## Importing Libraries
import os, sys, pathlib, fnmatch
import shutil as st
from IPython import get_ipython

from openseespy.opensees import *

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import warnings
import matplotlib.cbook
warnings.filterwarnings("ignore",category=matplotlib.cbook.mplDeprecation)
wipe()

# Getting Number of Ground Motions from the GM folder
GMdir = os.getcwd()
No_of_GMs = int(len(fnmatch.filter(os.listdir(GMdir),'*.AT2'))/2)
print('\nGenerating Spectra for {} provided GMs \n\n'.format(np.round(No_of_GMs,0)))

# Initializations
DISPLACEMENTS = pd.DataFrame(columns=['uX','uY'])
GM_SPECTRA = pd.DataFrame(columns=['Period(s)','RotD50Sa(g)', 'RotD100Sa(g)'])
SDOF_RESPONSE = [[]]
GM_RESPONSE = [[]]

# Spectra Generation
for iEQ in range(1,No_of_GMs+1):
    print('Generating Spectra for GM: {} ...\n'.format(np.round(iEQ,0)))
    Periods = np.concatenate((list(np.arange(Int_T_Reg_1,End_T_Reg_1+Int_T_Reg_1,Int_T_Reg_1)),list(np.arange(End_T_Reg_1+Int_T_Reg_2,End_T_Reg_2+Int_T_Reg_2,Int_T_Reg_2)),list(np.arange(End_T_Reg_2+Int_T_Reg_3,End_T_Reg_3+Int_T_Reg_3,Int_T_Reg_3))),axis=0)
    ii = 0

    for T in Periods:
        ii = ii+1
        GMinter = 0

        # Storing Periods
        GM_SPECTRA.loc[ii-1,'Period(s)'] = T

        # Setting modelbuilder
        model('basic', '-ndm', 3, '-ndf', 6)

        # Setting SODF Variables
        g = 386.1                   # value of g
        L = 1.0                     # Length
        d = 2                       # Diameter
        r = d/2                     # Radius
        A = np.pi*(r**2)            # Area
        E = 1.0                     # Elastic Modulus
        G = 1.0                     # Shear Modulus
        I3 = np.pi*(r**4)/4         # Moment of Inertia (zz)
        J = np.pi*(r**4)/2          # Polar Moment of Inertia
        I2 = np.pi*(r**4)/4         # Moment of Inertia (yy)
        K = 3*E*I3/(L**3)           # Stiffness
        M = K*(T**2)/4/(np.pi**2)   # Mass
        omega = np.sqrt(K/M)        # Natural Frequency
        Tn = 2*np.pi/omega          # Natural Period

        # Creating nodes
        node(1, 0.0, 0.0, 0.0)
        node(2, 0.0, 0.0, L)

        # Transformation
        transfTag = 1
        geomTransf('Linear',transfTag,0.0,1.0,0.0)

        # Setting boundary condition
        fix(1, 1, 1, 1, 1, 1, 1)

        # Defining materials
        uniaxialMaterial("Elastic", 11, E)

        # Defining elements
        element("elasticBeamColumn",12,1,2,A,E,G,J,I2,I3,1)

        # Defining mass
        mass(2,M,M,0.0,0.0,0.0,0.0)

        # Eigen Value Analysis (Verifying Period)
        numEigen = 1
        eigenValues = eigen(numEigen)
        omega = np.sqrt(eigenValues)
        T = 2*np.pi/omega
        print('   Calculating Spectral Ordinate for Period = {} secs'.format(np.round(T,3)))

        ## Reading GM Files
        exec(open("ReadGMFile.py").read())	            # read in procedure Multinition
        iGMinput = 'GM1'+str(iEQ)+' GM2'+str(iEQ) ;
        GMinput  = iGMinput.split(' ');
        gmXY     = {}
        for i in range(0,2):
            inFile   = GMdir + '\\'+ GMinput[i]+'.AT2';
            dt, NumPts , gmXY = ReadGMFile()

        # Storing GM Histories
        gmX = gmXY[1]
        gmY = gmXY[2]
        gmXY_mat = np.column_stack((gmX,gmX,gmY,gmY))

        # Bidirectional Uniform Earthquake ground motion (uniform acceleration input at all support nodes)
        iGMfile      = 'GM1'+str(iEQ)+' GM2'+str(iEQ) ;
        GMfile       = iGMfile.split(' ')
        GMdirection  = [1,1,2,2];
        GMfact	     = [np.cos(GMinter*np.pi/180),np.sin(-GMinter*np.pi/180), np.sin(GMinter*np.pi/180), np.cos(GMinter*np.pi/180)];
        IDTag        = 2
        loop         = [1,2,3,4]

        for i in loop:
            # Setting time series to be passed to uniform excitation
            timeSeries('Path',IDTag +i, '-dt', dt, '-values', *list(gmXY_mat[:,i-1]), '-factor', GMfact[i-1]*g)
            # Creating UniformExcitation load pattern
            pattern('UniformExcitation',  IDTag+i,   GMdirection[i-1],  '-accel', IDTag+i)

        # Defining Damping
        # Applying Rayleigh Damping from $xDamp
        # D=$alphaM*M + $betaKcurr*Kcurrent + $betaKcomm*KlastCommit + $beatKinit*$Kinitial
        xDamp 		= 0.05;								# 5% damping ratio
        alphaM 		= 0.;								# M-prop. damping; D = alphaM*M
        betaKcurr 	= 0.;         						# K-proportional damping;      +beatKcurr*KCurrent
        betaKcomm 	= 2.*xDamp/omega;   				# K-prop. damping parameter;   +betaKcomm*KlastCommitt
        betaKinit 	= 0.;         						# initial-stiffness proportional damping      +beatKinit*Kini
        rayleigh(alphaM,betaKcurr,betaKinit,betaKcomm); # RAYLEIGH damping

        # Creating the analysis
        wipeAnalysis()			            # clear previously-define analysis parameters
        constraints("Penalty",1e18, 1e18)   # how to handle boundary conditions
        numberer("RCM")                     # renumber dof's to minimize band-width (optimization), if you want to
        system('SparseGeneral')             # how to store and solve the system of equations in the analysis
        algorithm('Linear')	                # use Linear algorithm for linear analysis
        integrator("TRBDF2")                # determine the next time step for an analysis
        algorithm("NewtonLineSearch")       # define type of analysis: time-dependent
        test('EnergyIncr',1.0e-6, 100, 0)
        analysis("Transient")

        # Variables (Can alter the speed of analysis)
        dtAnalysis    = dt
        TmaxAnanlysis = dt*NumPts
        tFinal        = int(TmaxAnanlysis/dtAnalysis)
        tCurrent      = getTime()
        ok            = 0
        time          = [tCurrent]

        # Initializations of response
        u1            = [0.0]
        u2            = [0.0]

        # Performing the transient analysis (Performance is slow in this loop, can be altered by changing the parameters)
        while ok == 0 and tCurrent < tFinal:
            ok = analyze(1, dtAnalysis)
            # if the analysis fails try initial tangent iteration
            if ok != 0:
                print("Iteration failed .. lets try an initial stiffness for this step")
                test('NormDispIncr', 1.0e-12,  100, 0)
                algorithm('ModifiedNewton', '-initial')
                ok =analyze( 1, .001)

                if ok == 0:
                    print("that worked .. back to regular newton")
                    test('NormDispIncr', 1.0e-12,  10 )
                    algorithm('Newton')

            tCurrent = getTime()
            time.append(tCurrent)
            u1.append(nodeDisp(2,1))
            u2.append(nodeDisp(2,2))

        # Storing responses
        DISPLACEMENTS.loc[ii-1,'uX'] = np.array(u1)
        DISPLACEMENTS.loc[ii-1,'uY'] = np.array(u2)
        DISP_X_Y = np.column_stack((np.array(u1),np.array(u2)))

        # Rotating the Spectra (Projections)
        Rot_Matrix = np.zeros((2,2))
        Rot_Disp = np.zeros((180,1))
        for theta in range (0,180,1):
            Rot_Matrix [0,0] = np.cos(np.deg2rad(theta))
            Rot_Matrix [0,1] = np.sin(np.deg2rad(-theta))
            Rot_Matrix [1,0] = np.sin(np.deg2rad(theta))
            Rot_Matrix [1,1] = np.cos(np.deg2rad(theta))
            Rot_Disp[theta,0] = np.max(np.matmul(DISP_X_Y,Rot_Matrix)[:,0])

        # Storing Spectra
        Rot_Acc = np.dot(Rot_Disp,(omega**2)/g)
        GM_SPECTRA.loc[ii-1,'RotD50Sa(g)'] = np.median(Rot_Acc)
        GM_SPECTRA.loc[ii-1,'RotD100Sa(g)']= np.max(Rot_Acc)
        wipe()

    # Writing Spectra to Files
    if not os.path.exists('Spectra'):
        os.makedirs('Spectra')
    GM_SPECTRA.to_csv('Spectra//GM'+str(iEQ)+'_Spectra.txt', sep=' ',header=True,index=False)

    # Plotting Spectra
    if Plot_Spectra == 'Yes':

        def plot_spectra(PlotTitle,SpectraType,iGM):
            axes = fig.add_subplot(1, 1, 1)
            axes.plot(GM_SPECTRA['Period(s)'] , GM_SPECTRA[SpectraType] , '.-',lw=7,markersize=20, label='GM'+str(iGM))
            axes.set_xlabel('Period (sec)',fontsize=30,fontweight='bold')
            axes.set_ylabel(SpectraType,fontsize=30,fontweight='bold')
            axes.set_title(PlotTitle,fontsize=40,fontweight='bold')
            axes.tick_params(labelsize= 25)
            axes.grid(True)
            axes.set_xlim(0, np.ceil(max(GM_SPECTRA['Period(s)'])))
            axes.set_ylim(0, np.ceil(max(GM_SPECTRA[SpectraType])))
            axes.axhline(linewidth=10,color='black')
            axes.axvline(linewidth=10,color='black')
            axes.hold(True)
            axes.legend(fontsize =30)

        fig = plt.figure(1,figsize=(18,12))
        plot_spectra('RotD50 Spectra','RotD50Sa(g)',iEQ)

        fig = plt.figure(2,figsize=(18,12))
        plot_spectra('RotD100 Spectra','RotD100Sa(g)',iEQ)

    SDOF_RESPONSE.insert(iEQ-1,DISPLACEMENTS)
    GM_RESPONSE.insert(iEQ-1,GM_SPECTRA)

    print('\nGenerated Spectra for GM: {}\n\n'.format(np.round(iEQ,0)))


"""
Copyright 2013 Steven Diamond

This file is part of CVXPY.

CVXPY is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

CVXPY is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with CVXPY.  If not, see <http://www.gnu.org/licenses/>.

THIS FILE IS DEPRECATED AND MAY BE REMOVED WITHOUT WARNING!
DO NOT CALL THESE FUNCTIONS IN YOUR CODE!
WE ARE MOVING ALL SOLVER INTERFACES TO THE REDUCTIONS FOLDER.
"""

import cvxpy.interface as intf
import cvxpy.settings as s
from cvxpy.problems.problem_data.compr_matrix import compress_matrix
from cvxpy.problems.solvers.solver import Solver
from cvxpy.problems.kktsolver import get_kktsolver
import scipy.sparse as sp
import scipy
import numpy as np
import copy

class CVXOPT(Solver):
    """An interface for the CVXOPT solver.
    """
    NL_DUAL = 'nl_dual'

    # Solver capabilities.
    LP_CAPABLE = True
    SOCP_CAPABLE = True
    PSD_CAPABLE = True
    EXP_CAPABLE = True
    MIP_CAPABLE = False

    # Map of CVXOPT status to CVXPY status.
    STATUS_MAP = {'optimal': s.OPTIMAL,
                  'primal infeasible': s.INFEASIBLE,
                  'dual infeasible': s.UNBOUNDED,
                  'unknown': s.SOLVER_ERROR}

    def name(self):
        """The name of the solver.
        """
        return s.CVXOPT

    def import_solver(self):
        """Imports the solver.
        """
        import cvxopt
        cvxopt  # For flake8

    def matrix_intf(self):
        """The interface for matrices passed to the solver.
        """
        return intf.DEFAULT_SPARSE_INTF

    def vec_intf(self):
        """The interface for vectors passed to the solver.
        """
        return intf.DEFAULT_NP_INTF

    def split_constr(self, constr_map):
        """Extracts the equality, inequality, and nonlinear constraints.

        Parameters
        ----------
        constr_map : dict
            A dict of the canonicalized constraints.

        Returns
        -------
        tuple
            (eq_constr, ineq_constr, nonlin_constr)
        """
        return (constr_map[s.EQ], constr_map[s.LEQ], constr_map[s.EXP])

    def get_problem_data(self, objective, constraints, cached_data):
        """Returns the argument for the call to the solver.

        Parameters
        ----------
        objective : LinOp
            The canonicalized objective.
        constraints : list
            The list of canonicalized cosntraints.
        cached_data : dict
            A map of solver name to cached problem data.

        Returns
        -------
        dict
            The arguments needed for the solver.
        """
        # Returns CVXOPT matrices so can be used by raw CVXOPT solver.
        data = super(CVXOPT, self).get_problem_data(objective, constraints,
                                                    cached_data)
        # Convert A, b, G, h, c to CVXOPT matrices.
        data[s.A] = intf.sparse2cvxopt(data[s.A])
        data[s.G] = intf.sparse2cvxopt(data[s.G])
        data[s.B] = intf.dense2cvxopt(data[s.B])
        data[s.H] = intf.dense2cvxopt(data[s.H])
        data[s.C] = intf.dense2cvxopt(data[s.C])
        return data

    def solve(self, objective, constraints, cached_data,
              warm_start, verbose, solver_opts):
        """Returns the result of the call to the solver.

        Parameters
        ----------
        objective : LinOp
            The canonicalized objective.
        constraints : list
            The list of canonicalized cosntraints.
        cached_data : dict
            A map of solver name to cached problem data.
        warm_start : bool
            Not used.
        verbose : bool
            Should the solver print output?
        solver_opts : dict
            Additional arguments for the solver.

        Returns
        -------
        tuple
            (status, optimal value, primal, equality dual, inequality dual)
        """
        import cvxopt
        import cvxopt.solvers
        data = super(CVXOPT, self).get_problem_data(objective, constraints,
                                                    cached_data)
        # Save old data in case need to use robust solver.
        data[s.DIMS] = copy.deepcopy(data[s.DIMS])
        # Convert all longs to ints.
        for key, val in data[s.DIMS].items():
            if isinstance(val, list):
                data[s.DIMS][key] = [int(v) for v in val]
            else:
                data[s.DIMS][key] = int(val)
        # User chosen KKT solver option.
        kktsolver = self.get_kktsolver_opt(solver_opts)
        # Cannot have redundant rows unless using robust LDL kktsolver.
        if kktsolver != s.ROBUST_KKTSOLVER:
            # Will detect infeasibility.
            if self.remove_redundant_rows(data) == s.INFEASIBLE:
                return {s.STATUS: s.INFEASIBLE}
        # Convert A, b, G, h, c to CVXOPT matrices.
        data[s.A] = intf.sparse2cvxopt(data[s.A])
        data[s.G] = intf.sparse2cvxopt(data[s.G])
        data[s.B] = intf.dense2cvxopt(data[s.B])
        data[s.H] = intf.dense2cvxopt(data[s.H])
        data[s.C] = intf.dense2cvxopt(data[s.C])
        # Save original cvxopt solver options.
        old_options = cvxopt.solvers.options.copy()
        # Silence cvxopt if verbose is False.
        cvxopt.solvers.options["show_progress"] = verbose

        # Apply any user-specific options.
        # Rename max_iters to maxiters.
        if "max_iters" in solver_opts:
            solver_opts["maxiters"] = solver_opts["max_iters"]
        for key, value in solver_opts.items():
            cvxopt.solvers.options[key] = value

        # Always do 1 step of iterative refinement after solving KKT system.
        if "refinement" not in cvxopt.solvers.options:
            cvxopt.solvers.options["refinement"] = 1

        try:
            # Target cvxopt clp if nonlinear constraints exist.
            if data[s.DIMS][s.EXP_DIM]:
                results_dict = self.cpl_solve(data, kktsolver)
            else:
                results_dict = self.conelp_solve(data, kktsolver)
        # Catch exceptions in CVXOPT and convert them to solver errors.
        except ValueError:
            results_dict = {"status": "unknown"}

        # Restore original cvxopt solver options.
        self._restore_solver_options(old_options)
        return self.format_results(results_dict, data, cached_data)

    def cpl_solve(self, data, kktsolver):
        """Solve using the cpl solver.

        Parameters
        ----------
        data : dict
            All the problem data.
        kktsolver : The kktsolver to use.
        robust : Use the robust kktsolver?

        Returns
        -------
        dict
            The solver output.

        Raises
        ------
        ValueError
            If CVXOPT fails.
        """
        import cvxopt.solvers
        if kktsolver == s.ROBUST_KKTSOLVER:
            # Get custom kktsolver.
            kktsolver = get_kktsolver(data[s.G],
                                      data[s.DIMS],
                                      data[s.A],
                                      data[s.F])
        return cvxopt.solvers.cpl(data[s.C],
                                  data[s.F],
                                  data[s.G],
                                  data[s.H],
                                  data[s.DIMS],
                                  data[s.A],
                                  data[s.B],
                                  kktsolver=kktsolver)

    def conelp_solve(self, data, kktsolver):
        """Solve using the conelp solver.

        Parameters
        ----------
        data : dict
            All the problem data.
        kktsolver : The kktsolver to use.
        robust : Use the robust kktsolver?

        Returns
        -------
        dict
            The solver output.

        Raises
        ------
        ValueError
            If CVXOPT fails.
        """
        import cvxopt.solvers
        if kktsolver == s.ROBUST_KKTSOLVER:
            # Get custom kktsolver.
            kktsolver = get_kktsolver(data[s.G],
                                      data[s.DIMS],
                                      data[s.A])
        return cvxopt.solvers.conelp(data[s.C],
                                     data[s.G],
                                     data[s.H],
                                     data[s.DIMS],
                                     data[s.A],
                                     data[s.B],
                                     kktsolver=kktsolver)

    @staticmethod
    def remove_redundant_rows(data):
        """Remove redundant constraints from A and G.

        Parameters
        ----------
        data : dict
            All the problem data.

        Returns
        -------
        str
            A status indicating if infeasibility was detected.
        """
        # Extract data.
        dims = data[s.DIMS]
        A = data[s.A]
        G = data[s.G]
        b = data[s.B]
        h = data[s.H]
        # Remove redundant rows in A.
        if A.shape[0] > 0:
            # The pivoting improves robustness.
            Q, R, P = scipy.linalg.qr(A.todense(), pivoting=True)
            rows_to_keep = []
            for i in range(R.shape[0]):
                if np.linalg.norm(R[i, :]) > 1e-10:
                    rows_to_keep.append(i)
            R = R[rows_to_keep, :]
            Q = Q[:, rows_to_keep]
            # Invert P from col -> var to var -> col.
            Pinv = np.zeros(P.size, dtype='int')
            for i in range(P.size):
                Pinv[P[i]] = i
            # Rearrage R.
            R = R[:, Pinv]
            A = R
            b_old = b
            b = Q.T.dot(b)
            # If b is not in the range of Q,
            # the problem is infeasible.
            if not np.allclose(b_old, Q.dot(b)):
                return s.INFEASIBLE
            dims[s.EQ_DIM] = int(b.shape[0])
            data["Q"] = intf.dense2cvxopt(Q)
        # Remove obviously redundant rows in G's <= constraints.
        if dims[s.LEQ_DIM] > 0:
            G = G.tocsr()
            G_leq = G[:dims[s.LEQ_DIM], :]
            h_leq = h[:dims[s.LEQ_DIM]].ravel()
            G_other = G[dims[s.LEQ_DIM]:, :]
            h_other = h[dims[s.LEQ_DIM]:].ravel()
            G_leq, h_leq, P_leq = compress_matrix(G_leq, h_leq)
            dims[s.LEQ_DIM] = int(h_leq.shape[0])
            data["P_leq"] = intf.sparse2cvxopt(P_leq)
            G = sp.vstack([G_leq, G_other])
            h = np.hstack([h_leq, h_other])
        # Convert A, b, G, h to CVXOPT matrices.
        data[s.A] = A
        data[s.G] = G
        data[s.B] = b
        data[s.H] = h
        return s.OPTIMAL

    @staticmethod
    def _restore_solver_options(old_options):
        import cvxopt.solvers
        for key, value in list(cvxopt.solvers.options.items()):
            if key in old_options:
                cvxopt.solvers.options[key] = old_options[key]
            else:
                del cvxopt.solvers.options[key]

    def nonlin_constr(self):
        """Returns whether nonlinear constraints are needed.
        """
        return True

    @staticmethod
    def get_kktsolver_opt(solver_opts):
        """Returns the KKT solver selected by the user.

        Removes the KKT solver from solver_opts.

        Parameters
        ----------
        solver_opts : dict
            Additional arguments for the solver.

        Returns
        -------
        str or None
            The KKT solver chosen by the user.
        """
        if "kktsolver" in solver_opts:
            kktsolver = solver_opts["kktsolver"]
            del solver_opts["kktsolver"]
        else:
            kktsolver = 'chol'
        return kktsolver

    def format_results(self, results_dict, data, cached_data):
        """Converts the solver output into standard form.

        Parameters
        ----------
        results_dict : dict
            The solver output.
        data : dict
            Information about the problem.
        cached_data : dict
            A map of solver name to cached problem data.

        Returns
        -------
        dict
            The solver output in standard form.
        """
        import cvxopt
        new_results = {}
        status = self.STATUS_MAP[results_dict['status']]
        new_results[s.STATUS] = status
        if new_results[s.STATUS] in s.SOLUTION_PRESENT:
            primal_val = results_dict['primal objective']
            new_results[s.VALUE] = primal_val + data[s.OFFSET]
            new_results[s.PRIMAL] = results_dict['x']
            new_results[s.EQ_DUAL] = results_dict['y']
            if data[s.DIMS][s.EXP_DIM]:
                new_results[s.INEQ_DUAL] = results_dict['zl']
                new_results[self.NL_DUAL] = results_dict['znl']
            else:
                new_results[s.INEQ_DUAL] = results_dict['z']
            # Need to multiply duals by Q and P_leq.
            if "Q" in data:
                y = results_dict['y']
                # Test if all constraints eliminated.
                if y.size[0] == 0:
                    dual_len = data["Q"].size[0]
                    new_results[s.EQ_DUAL] = cvxopt.matrix(0., (dual_len, 1))
                else:
                    new_results[s.EQ_DUAL] = data["Q"]*y
            if "P_leq" in data:
                leq_len = data[s.DIMS][s.LEQ_DIM]
                P_rows = data["P_leq"].size[1]
                new_len = P_rows + new_results[s.INEQ_DUAL].size[0] - leq_len
                new_dual = cvxopt.matrix(0., (new_len, 1))
                z = new_results[s.INEQ_DUAL][:leq_len]
                # Test if all constraints eliminated.
                if z.size[0] == 0:
                    new_dual[:P_rows] = 0
                else:
                    new_dual[:P_rows] = data["P_leq"].T*z
                new_dual[P_rows:] = new_results[s.INEQ_DUAL][leq_len:]
                new_results[s.INEQ_DUAL] = new_dual

            for key in [s.PRIMAL, s.EQ_DUAL, s.INEQ_DUAL]:
                new_results[key] = intf.cvxopt2dense(new_results[key])
            if data[s.DIMS][s.EXP_DIM]:
                nl_dual = intf.cvxopt2dense(new_results[self.NL_DUAL])
                new_results[s.INEQ_DUAL] = np.vstack([new_results[s.INEQ_DUAL], nl_dual])

        return new_results

import numpy,import scipy,import cvxpy,from cvxpy


"""This module defines the fundamental interfaces and functions related to causal graphs in graphical causal models.

Classes and functions in this module should be considered experimental, meaning there might be breaking API changes in
the future.
"""

from abc import abstractmethod, ABC
from typing import Any, List

import networkx as nx
import numpy as np
from networkx.algorithms.dag import has_cycle
from typing_extensions import Protocol

# This constant is used as key when storing/accessing models as causal mechanisms in graph node attributes
CAUSAL_MECHANISM = 'causal_mechanism'

# This constant is used as key when storing the parents of a node during fitting. It's used for validation purposes
# afterwards.
PARENTS_DURING_FIT = 'parents_during_fit'

class HasNodes(Protocol):
    """This protocol defines a trait for classes having nodes."""
    @property
    @abstractmethod
    def nodes(self):
        """:returns Dict[Any, Dict[Any, Any]]"""
        raise NotImplementedError

class HasEdges(Protocol):
    """This protocol defines a trait for classes having edges."""
    @property
    @abstractmethod
    def edges(self):
        """:returns a Dict[Tuple[Any, Any], Dict[Any, Any]]"""
        raise NotImplementedError

class DirectedGraph(HasNodes, HasEdges, Protocol):
    """A protocol representing a directed graph as needed by graphical causal models.

    This protocol specifically defines a subset of the networkx.DiGraph class, which make that class automatically
    compatible with DirectedGraph. While in most cases a networkx.DiGraph is the class of choice when constructing
    a causal graph, anyone can choose to provide their own implementation of the DirectGraph interface.
    """
    @abstractmethod
    def predecessors(self, node):
        raise NotImplementedError

class StochasticModel(ABC):
    """A stochastic model represents a model used for causal mechanisms for root nodes in a graphical causal model."""

    @abstractmethod
    def fit(self, X: np.ndarray) -> None:
        """Fits the model according to the data."""
        raise NotImplementedError

    @abstractmethod
    def draw_samples(self, num_samples: int) -> np.ndarray:
        """Draws samples for the fitted model."""
        raise NotImplementedError

    @abstractmethod
    def clone(self):
        raise NotImplementedError

class ConditionalStochasticModel(ABC):
    """A conditional stochastic model represents a model used for causal mechanisms for non-root nodes in a graphical
    causal model."""

    @abstractmethod
    def fit(self, X: np.ndarray, Y: np.ndarray) -> None:
        """Fits the model according to the data."""
        raise NotImplementedError

    @abstractmethod
    def draw_samples(self, parent_samples: np.ndarray) -> np.ndarray:
        """Draws samples for the fitted model."""
        raise NotImplementedError

    @abstractmethod
    def clone(self):
        raise NotImplementedError

class FunctionalCausalModel(ConditionalStochasticModel):
    """Represents a Functional Causal Model (FCM), a specific type of conditional stochastic model, that is defined
    as:
        Y := f(X, N), N: Noise
    """

    def draw_samples(self, parent_samples: np.ndarray) -> np.ndarray:
        return self.evaluate(parent_samples, self.draw_noise_samples(parent_samples.shape[0]))

    @abstractmethod
    def draw_noise_samples(self, num_samples: int) -> np.ndarray:
        raise NotImplementedError

    @abstractmethod
    def evaluate(self, parent_samples: np.ndarray, noise_samples: np.ndarray) -> np.ndarray:
        raise NotImplementedError

class InvertibleFunctionalCausalModel(FunctionalCausalModel, ABC):
    @abstractmethod
    def estimate_noise(self, target_samples: np.ndarray, parent_samples: np.ndarray) -> np.ndarray:
        raise NotImplementedError

def is_root_node(causal_graph: DirectedGraph, node: Any) -> bool:
    return list(causal_graph.predecessors(node)) == []

def get_ordered_predecessors(causal_graph: DirectedGraph, node: Any) -> List[Any]:
    """This function returns predecessors of a node in a well-defined order.

    This is necessary, because we select subsets of columns in Dataframes by using a node's parents, and these parents
    might not be returned in a reliable order.
    """
    return sorted(causal_graph.predecessors(node))

def node_connected_subgraph_view(g: DirectedGraph, node: Any) -> Any:
    """Returns a view of the provided graph g that contains only nodes connected to the node passed in"""
    # can't use nx.node_connected_component, because it doesn't work with DiGraphs.
    # Hence a manual loop:
    return nx.induced_subgraph(g, [n for n in g.nodes if nx.has_path(g, n, node)])

def clone_causal_models(source: HasNodes, destination: HasNodes):
    for node in destination.nodes:
        if CAUSAL_MECHANISM in source.nodes[node]:
            destination.nodes[node][CAUSAL_MECHANISM] = source.nodes[node][CAUSAL_MECHANISM].clone()

def validate_acyclic(causal_graph: DirectedGraph) -> None:
    if has_cycle(causal_graph):
        raise RuntimeError('The graph contains a cycle, but an acyclic graph is expected!')

def validate_causal_dag(causal_graph: DirectedGraph) -> None:
    validate_acyclic(causal_graph)
    validate_causal_graph(causal_graph)

def validate_causal_graph(causal_graph: DirectedGraph) -> None:
    for node in causal_graph.nodes:
        validate_node(causal_graph, node)

def validate_node(causal_graph: DirectedGraph, node: Any) -> None:
    validate_causal_model_assignment(causal_graph, node)
    validate_local_structure(causal_graph, node)

def validate_causal_model_assignment(causal_graph: DirectedGraph, target_node: Any) -> None:
    validate_node_has_causal_model(causal_graph, target_node)

    causal_model = causal_graph.nodes[target_node][CAUSAL_MECHANISM]

    if is_root_node(causal_graph, target_node):
        if not isinstance(causal_model, StochasticModel):
            raise RuntimeError('Node %s is a root node and, thus, requires a StochasticModel, '
                               'but a %s was found!' % (target_node, causal_model))
    elif not isinstance(causal_model, ConditionalStochasticModel):
        raise RuntimeError('Node %s has parents and, thus, requires a ConditionalStochasticModel, '
                           'but a %s was found!' % (target_node, causal_model))

def validate_local_structure(causal_graph: DirectedGraph, node: Any) -> None:
    if PARENTS_DURING_FIT not in causal_graph.nodes[node] \
            or causal_graph.nodes[node][PARENTS_DURING_FIT] \
            != get_ordered_predecessors(causal_graph, node):
        raise RuntimeError('The causal mechanism of node %s is not fitted to the graphical structure! Fit all'
                           'causal models in the graph first. If the mechanism is already fitted based on the causal'
                           'parents, consider to update the persisted parents for that node manually.' % node)

def validate_node_has_causal_model(causal_graph: HasNodes, node: Any) -> None:
    validate_node_in_graph(causal_graph, node)

    if CAUSAL_MECHANISM not in causal_graph.nodes[node]:
        raise ValueError("Node %s has no assigned causal mechanism!" % node)

def validate_node_in_graph(causal_graph: HasNodes, node: Any) -> None:
    if node not in causal_graph.nodes:
        raise ValueError("Node %s can not be found in the given graph!" % node)

import numpy,import networkx,from networkx


import torch
import numpy as np
import redner
import pyredner
import time
import skimage.io

# There is a bias-variance trade off in the backward pass.
# If the forward pass and the backward pass are correlated
# the gradients are biased for L2 loss.
# (E[d/dx(f(x) - y)^2] = E[(f(x) - y) d/dx f(x)])
#                      = E[f(x) - y] E[d/dx f(x)]
# The last equation only holds when f(x) and d/dx f(x) are independent.
# It is usually better to use the unbiased one, but we left it as an option here
use_correlated_random_number = False
def set_use_correlated_random_number(v):
    global use_correlated_random_number
    use_correlated_random_number = v

def get_use_correlated_random_number():
    global use_correlated_random_number
    return use_correlated_random_number

print_timing = True

class RenderFunction(torch.autograd.Function):
    """
        The PyTorch interface of Redner.
    """

    @staticmethod
    def serialize_scene(scene,
                        num_samples,
                        max_bounces,
                        channels = [redner.channels.radiance],
                        sampler_type = redner.SamplerType.independent,
                        use_primary_edge_sampling = True,
                        use_secondary_edge_sampling = True):
        """
            Given a PyRedner scene & rendering options, convert them to a linear list of argument,
            so that we can use it in PyTorch.

            Keyword arguments:
            scene -- A pyredner.Scene
            num_samples -- Number of samples per pixel for forward and backward passes,
                           can be an integer or a tuple of 2 integers.
            max_bounces -- Number of bounces for global illumination, 1 means direct lighting only.
            channels -- A list of channels that should present in the output image.
                        Following channels are supported:
                            redner.channels.radiance,
                            redner.channels.alpha,
                            redner.channels.depth,
                            redner.channels.position,
                            redner.channels.geometry_normal,
                            redner.channels.shading_normal,
                            redner.channels.uv,
                            redner.channels.diffuse_reflectance,
                            redner.channels.specular_reflectance,
                            redner.channels.roughness,
                            redner.channels.shape_id,
                            redner.channels.material_id
                        All channels, except for shape id and material id, are differentiable.
            sampler_type -- Which sampling pattern to use.
                            See Chapter 7 of the PBRT book for an explanation of the difference between
                            different samplers.
                            http://www.pbr-book.org/3ed-2018/Sampling_and_Reconstruction.html
                            Following samplers are supported:
                                redner.SamplerType.independent
                                redner.SamplerType.sobol
            use_primary_edge_sampling -- A boolean
            use_secondary_edge_sampling -- A boolean
        """
        cam = scene.camera
        num_shapes = len(scene.shapes)
        num_materials = len(scene.materials)
        num_lights = len(scene.area_lights)
        for light_id, light in enumerate(scene.area_lights):
            scene.shapes[light.shape_id].light_id = light_id
        args = []
        args.append(num_shapes)
        args.append(num_materials)
        args.append(num_lights)
        assert(torch.isfinite(cam.position).all())
        assert(torch.isfinite(cam.look_at).all())
        assert(torch.isfinite(cam.up).all())
        assert(torch.isfinite(cam.ndc_to_cam).all())
        assert(torch.isfinite(cam.cam_to_ndc).all())
        args.append(cam.position)
        args.append(cam.look_at)
        args.append(cam.up)
        args.append(cam.ndc_to_cam)
        args.append(cam.cam_to_ndc)
        args.append(cam.clip_near)
        args.append(cam.resolution)
        args.append(cam.camera_type)
        for shape in scene.shapes:
            assert(torch.isfinite(shape.vertices).all())
            if (shape.uvs is not None):
                assert(torch.isfinite(shape.uvs).all())
            if (shape.normals is not None):
                assert(torch.isfinite(shape.normals).all())
            args.append(shape.vertices)
            args.append(shape.indices)
            args.append(shape.uvs)
            args.append(shape.normals)
            args.append(shape.uv_indices)
            args.append(shape.normal_indices)
            args.append(shape.material_id)
            args.append(shape.light_id)
        for material in scene.materials:
            assert(torch.isfinite(material.diffuse_reflectance.mipmap).all())
            assert(torch.isfinite(material.diffuse_reflectance.uv_scale).all())
            assert(torch.isfinite(material.specular_reflectance.mipmap).all())
            assert(torch.isfinite(material.specular_reflectance.uv_scale).all())
            assert(torch.isfinite(material.roughness.mipmap).all())
            assert(torch.isfinite(material.roughness.uv_scale).all())
            args.append(material.diffuse_reflectance.mipmap)
            args.append(material.diffuse_reflectance.uv_scale)
            args.append(material.specular_reflectance.mipmap)
            args.append(material.specular_reflectance.uv_scale)
            args.append(material.roughness.mipmap)
            args.append(material.roughness.uv_scale)
            if material.normal_map is not None:
                assert(torch.isfinite(material.normal_map.mipmap).all())
                assert(torch.isfinite(material.normal_map.uv_scale).all())
                args.append(material.normal_map.mipmap)
                args.append(material.normal_map.uv_scale)
            else:
                args.append(None)
                args.append(None)
            args.append(material.two_sided)
        for light in scene.area_lights:
            args.append(light.shape_id)
            args.append(light.intensity)
            args.append(light.two_sided)
        if scene.envmap is not None:
            assert(torch.isfinite(scene.envmap.values.mipmap).all())
            assert(torch.isfinite(scene.envmap.values.uv_scale).all())
            assert(torch.isfinite(scene.envmap.env_to_world).all())
            assert(torch.isfinite(scene.envmap.world_to_env).all())
            assert(torch.isfinite(scene.envmap.sample_cdf_ys).all())
            assert(torch.isfinite(scene.envmap.sample_cdf_xs).all())
            args.append(scene.envmap.values.mipmap)
            args.append(scene.envmap.values.uv_scale)
            args.append(scene.envmap.env_to_world)
            args.append(scene.envmap.world_to_env)
            args.append(scene.envmap.sample_cdf_ys)
            args.append(scene.envmap.sample_cdf_xs)
            args.append(scene.envmap.pdf_norm)
        else:
            args.append(None)
            args.append(None)
            args.append(None)
            args.append(None)
            args.append(None)
            args.append(None)
            args.append(None)
        args.append(num_samples)
        args.append(max_bounces)
        args.append(channels)
        args.append(sampler_type)
        args.append(use_primary_edge_sampling)
        args.append(use_secondary_edge_sampling)

        return args

    @staticmethod
    def forward(ctx,
                seed,
                *args):
        """
            Forward rendering pass: given a scene and output an image.
        """
        # Unpack arguments
        current_index = 0
        num_shapes = args[current_index]
        current_index += 1
        num_materials = args[current_index]
        current_index += 1
        num_lights = args[current_index]
        current_index += 1
        cam_position = args[current_index]
        current_index += 1
        cam_look_at = args[current_index]
        current_index += 1
        cam_up = args[current_index]
        current_index += 1
        ndc_to_cam = args[current_index]
        current_index += 1
        cam_to_ndc = args[current_index]
        current_index += 1
        clip_near = args[current_index]
        current_index += 1
        resolution = args[current_index]
        current_index += 1
        camera_type = args[current_index]
        current_index += 1
        camera = redner.Camera(resolution[1],
                               resolution[0],
                               redner.float_ptr(cam_position.data_ptr()),
                               redner.float_ptr(cam_look_at.data_ptr()),
                               redner.float_ptr(cam_up.data_ptr()),
                               redner.float_ptr(ndc_to_cam.data_ptr()),
                               redner.float_ptr(cam_to_ndc.data_ptr()),
                               clip_near,
                               camera_type)
        shapes = []
        for i in range(num_shapes):
            vertices = args[current_index]
            current_index += 1
            indices = args[current_index]
            current_index += 1
            uvs = args[current_index]
            current_index += 1
            normals = args[current_index]
            current_index += 1
            uv_indices = args[current_index]
            current_index += 1
            normal_indices = args[current_index]
            current_index += 1
            material_id = args[current_index]
            current_index += 1
            light_id = args[current_index]
            current_index += 1
            assert(vertices.is_contiguous())
            assert(indices.is_contiguous())
            if uvs is not None:
                assert(uvs.is_contiguous())
            if normals is not None:
                assert(normals.is_contiguous())
            if uv_indices is not None:
                assert(uv_indices.is_contiguous())
            if normal_indices is not None:
                assert(normal_indices.is_contiguous())
            shapes.append(redner.Shape(\
                redner.float_ptr(vertices.data_ptr()),
                redner.int_ptr(indices.data_ptr()),
                redner.float_ptr(uvs.data_ptr() if uvs is not None else 0),
                redner.float_ptr(normals.data_ptr() if normals is not None else 0),
                redner.int_ptr(uv_indices.data_ptr() if uv_indices is not None else 0),
                redner.int_ptr(normal_indices.data_ptr() if normal_indices is not None else 0),
                int(vertices.shape[0]),
                int(uvs.shape[0]) if uvs is not None else 0,
                int(normals.shape[0]) if normals is not None else 0,
                int(indices.shape[0]),
                material_id,
                light_id))
        materials = []
        for i in range(num_materials):
            diffuse_reflectance = args[current_index]
            current_index += 1
            diffuse_uv_scale = args[current_index]
            current_index += 1
            specular_reflectance = args[current_index]
            current_index += 1
            specular_uv_scale = args[current_index]
            current_index += 1
            roughness = args[current_index]
            current_index += 1
            roughness_uv_scale = args[current_index]
            current_index += 1
            normal_map = args[current_index]
            current_index += 1
            normal_map_uv_scale = args[current_index]
            current_index += 1
            two_sided = args[current_index]
            current_index += 1
            assert(diffuse_reflectance.is_contiguous())
            if diffuse_reflectance.dim() == 1:
                diffuse_reflectance = redner.Texture3(\
                    redner.float_ptr(diffuse_reflectance.data_ptr()), 0, 0, 0,
                    redner.float_ptr(diffuse_uv_scale.data_ptr()))
            else:
                diffuse_reflectance = redner.Texture3(\
                    redner.float_ptr(diffuse_reflectance.data_ptr()),
                    int(diffuse_reflectance.shape[2]), # width
                    int(diffuse_reflectance.shape[1]), # height
                    int(diffuse_reflectance.shape[0]), # num levels
                    redner.float_ptr(diffuse_uv_scale.data_ptr()))
            assert(specular_reflectance.is_contiguous())
            if specular_reflectance.dim() == 1:
                specular_reflectance = redner.Texture3(\
                    redner.float_ptr(specular_reflectance.data_ptr()), 0, 0, 0,
                    redner.float_ptr(specular_uv_scale.data_ptr()))
            else:
                specular_reflectance = redner.Texture3(\
                    redner.float_ptr(specular_reflectance.data_ptr()),
                    int(specular_reflectance.shape[2]), # width
                    int(specular_reflectance.shape[1]), # height
                    int(specular_reflectance.shape[0]), # num levels
                    redner.float_ptr(specular_uv_scale.data_ptr()))
            assert(roughness.is_contiguous())
            if roughness.dim() == 1:
                roughness = redner.Texture1(\
                    redner.float_ptr(roughness.data_ptr()), 0, 0, 0,
                    redner.float_ptr(roughness_uv_scale.data_ptr()))
            else:
                assert(roughness.dim() == 4)
                roughness = redner.Texture1(\
                    redner.float_ptr(roughness.data_ptr()),
                    int(roughness.shape[2]), # width
                    int(roughness.shape[1]), # height
                    int(roughness.shape[0]), # num levels
                    redner.float_ptr(roughness_uv_scale.data_ptr()))
            if normal_map is not None:
                assert(normal_map.dim() == 4)
                normal_map = redner.Texture3(\
                    redner.float_ptr(normal_map.data_ptr()),
                    int(normal_map.shape[2]), # width
                    int(normal_map.shape[1]), # height
                    int(normal_map.shape[0]), # num levels
                    redner.float_ptr(normal_map_uv_scale.data_ptr()))
            else:
                normal_map = redner.Texture3(\
                    redner.float_ptr(0), 0, 0, 0, redner.float_ptr(0))
            materials.append(redner.Material(\
                diffuse_reflectance,
                specular_reflectance,
                roughness,
                normal_map,
                two_sided))

        area_lights = []
        for i in range(num_lights):
            shape_id = args[current_index]
            current_index += 1
            intensity = args[current_index]
            current_index += 1
            two_sided = args[current_index]
            current_index += 1

            area_lights.append(redner.AreaLight(\
                shape_id,
                redner.float_ptr(intensity.data_ptr()),
                two_sided))

        envmap = None
        if args[current_index] is not None:
            values = args[current_index]
            current_index += 1
            envmap_uv_scale = args[current_index]
            current_index += 1
            env_to_world = args[current_index]
            current_index += 1
            world_to_env = args[current_index]
            current_index += 1
            sample_cdf_ys = args[current_index]
            current_index += 1
            sample_cdf_xs = args[current_index]
            current_index += 1
            pdf_norm = args[current_index]
            current_index += 1
            values = redner.Texture3(\
                redner.float_ptr(values.data_ptr()),
                int(values.shape[2]), # width
                int(values.shape[1]), # height
                int(values.shape[0]), # num levels
                redner.float_ptr(envmap_uv_scale.data_ptr()))
            envmap = redner.EnvironmentMap(\
                values,
                redner.float_ptr(env_to_world.data_ptr()),
                redner.float_ptr(world_to_env.data_ptr()),
                redner.float_ptr(sample_cdf_ys.data_ptr()),
                redner.float_ptr(sample_cdf_xs.data_ptr()),
                pdf_norm)
        else:
            current_index += 7

        # Options
        num_samples = args[current_index]
        current_index += 1
        max_bounces = args[current_index]
        current_index += 1
        channels = args[current_index]
        current_index += 1
        sampler_type = args[current_index]
        current_index += 1
        use_primary_edge_sampling = args[current_index]
        current_index += 1
        use_secondary_edge_sampling = args[current_index]
        current_index += 1

        start = time.time()
        scene = redner.Scene(camera,
                             shapes,
                             materials,
                             area_lights,
                             envmap,
                             pyredner.get_use_gpu(),
                             pyredner.get_device().index if pyredner.get_device().index is not None else -1,
                             use_primary_edge_sampling,
                             use_secondary_edge_sampling)
        time_elapsed = time.time() - start
        if print_timing:
            print('Scene construction, time: %.5f s' % time_elapsed)

        # check that num_samples is a tuple
        if isinstance(num_samples, int):
            num_samples = (num_samples, num_samples)

        options = redner.RenderOptions(seed, num_samples[0], max_bounces, channels, sampler_type)
        num_channels = redner.compute_num_channels(channels)
        rendered_image = torch.zeros(resolution[0], resolution[1], num_channels,
            device = pyredner.get_device())
        start = time.time()
        redner.render(scene,
                      options,
                      redner.float_ptr(rendered_image.data_ptr()),
                      redner.float_ptr(0),
                      None,
                      redner.float_ptr(0))
        time_elapsed = time.time() - start
        if print_timing:
            print('Forward pass, time: %.5f s' % time_elapsed)

        # # For debugging
        # debug_img = torch.zeros(256, 256, 3)
        # redner.render(scene,
        #               options,
        #               redner.float_ptr(rendered_image.data_ptr()),
        #               redner.float_ptr(0),
        #               None,
        #               redner.float_ptr(debug_img.data_ptr()))
        # pyredner.imwrite(debug_img, 'debug.exr')
        # exit()

        ctx.shapes = shapes
        ctx.materials = materials
        ctx.area_lights = area_lights
        ctx.envmap = envmap
        ctx.scene = scene
        ctx.options = options
        ctx.num_samples = num_samples
        return rendered_image

    @staticmethod
    def backward(ctx,
                 grad_img):
        if not grad_img.is_contiguous():
            grad_img = grad_img.contiguous()
        scene = ctx.scene
        options = ctx.options

        d_cam_position = torch.zeros(3, device = pyredner.get_device())
        d_cam_look = torch.zeros(3, device = pyredner.get_device())
        d_cam_up = torch.zeros(3, device = pyredner.get_device())
        d_ndc_to_cam = torch.zeros(3, 3, device = pyredner.get_device())
        d_cam_to_ndc = torch.zeros(3, 3, device = pyredner.get_device())
        d_camera = redner.DCamera(redner.float_ptr(d_cam_position.data_ptr()),
                                  redner.float_ptr(d_cam_look.data_ptr()),
                                  redner.float_ptr(d_cam_up.data_ptr()),
                                  redner.float_ptr(d_ndc_to_cam.data_ptr()),
                                  redner.float_ptr(d_cam_to_ndc.data_ptr()))
        d_vertices_list = []
        d_uvs_list = []
        d_normals_list = []
        d_shapes = []
        for shape in ctx.shapes:
            num_vertices = shape.num_vertices
            num_uv_vertices = shape.num_uv_vertices
            d_vertices = torch.zeros(num_vertices, 3,
                device = pyredner.get_device())
            d_uvs = torch.zeros(num_uv_vertices, 2,
                device = pyredner.get_device()) if shape.has_uvs() else None
            d_normals = torch.zeros(num_vertices, 3,
                device = pyredner.get_device()) if shape.has_normals() else None
            d_vertices_list.append(d_vertices)
            d_uvs_list.append(d_uvs)
            d_normals_list.append(d_normals)
            d_shapes.append(redner.DShape(\
                redner.float_ptr(d_vertices.data_ptr()),
                redner.float_ptr(d_uvs.data_ptr() if d_uvs is not None else 0),
                redner.float_ptr(d_normals.data_ptr() if d_normals is not None else 0)))

        d_diffuse_list = []
        d_diffuse_uv_scale_list = []
        d_specular_list = []
        d_specular_uv_scale_list = []
        d_roughness_list = []
        d_roughness_uv_scale_list = []
        d_normal_map_list = []
        d_normal_map_uv_scale_list = []
        d_materials = []
        for material in ctx.materials:
            diffuse_size = material.get_diffuse_size()
            specular_size = material.get_specular_size()
            roughness_size = material.get_roughness_size()
            normal_map_size = material.get_normal_map_size()
            if diffuse_size[0] == 0:
                d_diffuse = torch.zeros(3, device = pyredner.get_device())
            else:
                d_diffuse = torch.zeros(diffuse_size[2],
                                        diffuse_size[1],
                                        diffuse_size[0],
                                        3, device = pyredner.get_device())
            if specular_size[0] == 0:
                d_specular = torch.zeros(3, device = pyredner.get_device())
            else:
                d_specular = torch.zeros(specular_size[2],
                                         specular_size[1],
                                         specular_size[0],
                                         3, device = pyredner.get_device())
            if roughness_size[0] == 0:
                d_roughness = torch.zeros(1, device = pyredner.get_device())
            else:
                d_roughness = torch.zeros(roughness_size[2],
                                          roughness_size[1],
                                          roughness_size[0],
                                          1, device = pyredner.get_device())
            if normal_map_size[0] == 0:
                d_normal_map = None
            else:
                d_normal_map = torch.zeros(normal_map_size[2],
                                           normal_map_size[1],
                                           normal_map_size[0],
                                           3, device = pyredner.get_device())
            d_diffuse_list.append(d_diffuse)
            d_specular_list.append(d_specular)
            d_roughness_list.append(d_roughness)
            d_normal_map_list.append(d_normal_map)
            d_diffuse_uv_scale = torch.zeros(2, device = pyredner.get_device())
            d_specular_uv_scale = torch.zeros(2, device = pyredner.get_device())
            d_roughness_uv_scale = torch.zeros(2, device = pyredner.get_device())
            d_diffuse_uv_scale_list.append(d_diffuse_uv_scale)
            d_specular_uv_scale_list.append(d_specular_uv_scale)
            d_roughness_uv_scale_list.append(d_roughness_uv_scale)
            if d_normal_map is None:
                d_normal_map_uv_scale = None
            else:
                d_normal_map_uv_scale = torch.zeros(2, device = pyredner.get_device())
            d_normal_map_uv_scale_list.append(d_normal_map_uv_scale)
            d_diffuse_tex = redner.Texture3(\
                redner.float_ptr(d_diffuse.data_ptr()),
                diffuse_size[0], diffuse_size[1], diffuse_size[2],
                redner.float_ptr(d_diffuse_uv_scale.data_ptr()))
            d_specular_tex = redner.Texture3(\
                redner.float_ptr(d_specular.data_ptr()),
                specular_size[0], specular_size[1], specular_size[2],
                redner.float_ptr(d_specular_uv_scale.data_ptr()))
            d_roughness_tex = redner.Texture1(\
                redner.float_ptr(d_roughness.data_ptr()),
                roughness_size[0], roughness_size[1], roughness_size[2],
                redner.float_ptr(d_roughness_uv_scale.data_ptr()))
            if d_normal_map is None:
                d_normal_map = redner.Texture3(\
                    redner.float_ptr(0), 0, 0, 0, redner.float_ptr(0))
            else:
                d_normal_map = redner.Texture3(\
                    redner.float_ptr(d_normal_map.data_ptr()),
                    normal_map_size[0], normal_map_size[1], normal_map_size[2],
                    redner.float_ptr(d_normal_map_uv_scale.data_ptr()))
            d_materials.append(redner.DMaterial(\
                d_diffuse_tex, d_specular_tex, d_roughness_tex, d_normal_map))

        d_intensity_list = []
        d_area_lights = []
        for light in ctx.area_lights:
            d_intensity = torch.zeros(3, device = pyredner.get_device())
            d_intensity_list.append(d_intensity)
            d_area_lights.append(\
                redner.DAreaLight(redner.float_ptr(d_intensity.data_ptr())))

        d_envmap = None
        if ctx.envmap is not None:
            envmap = ctx.envmap
            size = envmap.get_size()
            d_envmap_values = \
                torch.zeros(size[2],
                            size[1],
                            size[0],
                            3,
                            device = pyredner.get_device())
            d_envmap_uv_scale = torch.zeros(2, device = pyredner.get_device())
            d_envmap_tex = redner.Texture3(\
                redner.float_ptr(d_envmap_values.data_ptr()),
                size[0], size[1], size[2],
                redner.float_ptr(d_envmap_uv_scale.data_ptr()))
            d_world_to_env = torch.zeros(4, 4, device = pyredner.get_device())
            d_envmap = redner.DEnvironmentMap(\
                d_envmap_tex,
                redner.float_ptr(d_world_to_env.data_ptr()))

        d_scene = redner.DScene(d_camera,
                                d_shapes,
                                d_materials,
                                d_area_lights,
                                d_envmap,
                                pyredner.get_use_gpu(),
                                pyredner.get_device().index if pyredner.get_device().index is not None else -1)
        if not get_use_correlated_random_number():
            # Decouple the forward/backward random numbers by adding a big prime number
            options.seed += 1000003

        options.num_samples = ctx.num_samples[1]
        start = time.time()
        redner.render(scene, options,
                      redner.float_ptr(0),
                      redner.float_ptr(grad_img.data_ptr()),
                      d_scene,
                      redner.float_ptr(0))
        time_elapsed = time.time() - start
        if print_timing:
            print('Backward pass, time: %.5f s' % time_elapsed)

        # For debugging
        # pyredner.imwrite(grad_img, 'grad_img.exr')
        # grad_img = torch.ones(256, 256, 3, device = pyredner.get_device())
        # debug_img = torch.zeros(256, 256, 3)
        # start = time.time()
        # redner.render(scene, options,
        #               redner.float_ptr(0),
        #               redner.float_ptr(grad_img.data_ptr()),
        #               d_scene,
        #               redner.float_ptr(debug_img.data_ptr()))
        # time_elapsed = time.time() - start
        # if print_timing:
        #     print('Backward pass, time: %.5f s' % time_elapsed)
        # debug_img = debug_img[:, :, 0]
        # pyredner.imwrite(debug_img, 'debug.exr')
        # pyredner.imwrite(-debug_img, 'debug_.exr')
        # debug_img = debug_img.numpy()
        # print(np.max(debug_img))
        # print(np.unravel_index(np.argmax(debug_img), debug_img.shape))
        # print(np.min(debug_img))
        # print(np.unravel_index(np.argmin(debug_img), debug_img.shape))
        # print(np.sum(debug_img) / 3)
        # debug_max = 0.5
        # debug_min = -0.5
        # debug_img = np.clip((debug_img - debug_min) / (debug_max - debug_min), 0, 1)
        # # debug_img = debug_img[:, :, 0]
        # import matplotlib.cm as cm
        # debug_img = cm.viridis(debug_img)
        # skimage.io.imsave('debug.png', np.power(debug_img, 1/2.2))
        # exit()

        ret_list = []
        ret_list.append(None) # seed
        ret_list.append(None) # num_shapes
        ret_list.append(None) # num_materials
        ret_list.append(None) # num_lights
        ret_list.append(d_cam_position.cpu())
        ret_list.append(d_cam_look.cpu())
        ret_list.append(d_cam_up.cpu())
        ret_list.append(d_ndc_to_cam.cpu())
        ret_list.append(d_cam_to_ndc.cpu())
        ret_list.append(None) # clip near
        ret_list.append(None) # resolution
        ret_list.append(None) # camera_type

        num_shapes = len(ctx.shapes)
        for i in range(num_shapes):
            ret_list.append(d_vertices_list[i])
            ret_list.append(None) # indices
            ret_list.append(d_uvs_list[i])
            ret_list.append(d_normals_list[i])
            ret_list.append(None) # uv_indices
            ret_list.append(None) # normal_indices
            ret_list.append(None) # material id
            ret_list.append(None) # light id

        num_materials = len(ctx.materials)
        for i in range(num_materials):
            ret_list.append(d_diffuse_list[i])
            ret_list.append(d_diffuse_uv_scale_list[i])
            ret_list.append(d_specular_list[i])
            ret_list.append(d_specular_uv_scale_list[i])
            ret_list.append(d_roughness_list[i])
            ret_list.append(d_roughness_uv_scale_list[i])
            ret_list.append(d_normal_map_list[i])
            ret_list.append(d_normal_map_uv_scale_list[i])
            ret_list.append(None) # two sided

        num_area_lights = len(ctx.area_lights)
        for i in range(num_area_lights):
            ret_list.append(None) # shape id
            ret_list.append(d_intensity_list[i].cpu())
            ret_list.append(None) # two sided

        if ctx.envmap is not None:
            ret_list.append(d_envmap_values)
            ret_list.append(d_envmap_uv_scale)
            ret_list.append(None) # env_to_world
            ret_list.append(d_world_to_env.cpu())
            ret_list.append(None) # sample_cdf_ys
            ret_list.append(None) # sample_cdf_xs
            ret_list.append(None) # pdf_norm
        else:
            ret_list.append(None)
            ret_list.append(None)
            ret_list.append(None)
            ret_list.append(None)
            ret_list.append(None)
            ret_list.append(None)
            ret_list.append(None)

        ret_list.append(None) # num samples
        ret_list.append(None) # num bounces
        ret_list.append(None) # channels
        ret_list.append(None) # sampler type
        ret_list.append(None) # use_primary_edge_sampling
        ret_list.append(None) # use_secondary_edge_sampling

        return tuple(ret_list)


meld_NMR_script = '''#!/usr/bin/env python
# encoding: utf-8

import numpy as np
from meld.remd import ladder, adaptor, leader
from meld import comm, vault
from meld import system
from meld import parse
import meld.system.montecarlo as mc
from meld.system.restraints import LinearRamp,ConstantRamp
from collections import namedtuple
import glob as glob

N_REPLICAS = 30
N_STEPS =10000
BLOCK_SIZE = 100

hydrophobes = 'AILMFPWV'
hydrophobes_res = ['ALA','ILE','LEU','MET','PHE','PRO','TRP','VAL']

def gen_state_templates(index, templates):
    n_templates = len(templates)
    print((index,n_templates,index%n_templates))
    a = system.ProteinMoleculeFromPdbFile(templates[index%n_templates])
    #Note that it does not matter which forcefield we use here to build
    #as that information is not passed on, it is used for all the same as
    #in the setup part of the script
    b = system.SystemBuilder(forcefield="ff14sbside")
    c = b.build_system_from_molecules([a])
    pos = c._coordinates
    c._box_vectors=np.array([0.,0.,0.])
    vel = np.zeros_like(pos)
    alpha = index / (N_REPLICAS - 1.0)
    energy = 0
    return system.SystemState(pos, vel, alpha, energy,c._box_vectors)

def get_dist_restraints(filename, s, scaler):
    dists = []
    rest_group = []
    lines = open(filename).read().splitlines()
    lines = [line.strip() for line in lines]
    for line in lines:
        if not line:
            dists.append(s.restraints.create_restraint_group(rest_group, 1))
            rest_group = []
        else:
            cols = line.split()
            i = int(cols[0])
            name_i = cols[1]
            j = int(cols[2])
            name_j = cols[3]
            dist = float(cols[4]) / 10.

            rest = s.restraints.create_restraint('distance', scaler,LinearRamp(0,100,0,1),
                                                  r1=0.0, r2=0.0, r3=dist, r4=dist+0.2, k=350,
                                                  atom_1_res_index=i, atom_2_res_index=j,
                                                  atom_1_name=name_i, atom_2_name=name_j)
            rest_group.append(rest)
    return dists

def get_torsion_restraints(filename, s, scaler):
    torsion_rests = []
    rotamer_group = []
    for line in open(filename,'r'):
        if not line:
            torsion_rests.append(s.restraints.create_restraint_group(rotamer_group, 1))
            rotamer_group = []
        else:
            (res1, at1, res2, at2, res3, at3, res4, at4, rotamer_min, rotamer_max) = line.split()
            rotamer_max = float(rotamer_max)
            rotamer_min = float(rotamer_min)
            rotamer_avg = (rotamer_max+rotamer_min)/2.
            rotamer_sd = rotamer_max - rotamer_min
            rotamer_rest = s.restraints.create_restraint('torsion', scaler, LinearRamp(0,100,0,1),
                                                 phi=rotamer_avg, delta_phi=rotamer_sd, k=2.5,
                                                 atom_1_res_index=int(res1), atom_1_name=at1,
                                                 atom_2_res_index=int(res2), atom_2_name=at2,
                                                 atom_3_res_index=int(res3), atom_3_name=at3,
                                                 atom_4_res_index=int(res4), atom_4_name=at4)
            rotamer_group.append(rotamer_rest)
    return torsion_rests

def setup_system():
    #
    # Start system from minimized structure(s) deposited in TEMPLATES directory
    #

    templates = glob.glob('TEMPLATES/*.pdb')

    #
    # build the system with force field
    #

    p = system.ProteinMoleculeFromPdbFile(templates[0])
    b = system.SystemBuilder(forcefield="ff14sbside")
    s = b.build_system_from_molecules([p])
    n_res = s.residue_numbers[-1]

    #
    # Temperature ladder
    #
    s.temperature_scaler = system.GeometricTemperatureScaler(0, 0.4, 300., 500.)

    #
    # Scalers
    #
    distance_scaler = s.restraints.create_scaler('nonlinear', alpha_min=0.4, alpha_max=1.0, factor=4.0)
    torsion_scaler = s.restraints.create_scaler('constant')

    #
    # Distance Restraints
    #
    for noe in glob.glob('noe_*.dat'):
        NOESY = get_dist_restraints(noe,s,scaler=distance_scaler)
        s.restraints.add_selectively_active_collection(NOESY, int( len(NOESY)*0.95 ) )

    #
    # Torsion Restraints
    #
    for rotamer in glob.glob('rotamer_*.dat'):
        TALOS = get_torsion_restraints(rotamer, s, torsion_scaler)
        s.restraints.add_selectively_active_collection(TALOS, int( len(TALOS) * 0.95) )

    #
    # setup mc minimizer moves
    #
    movers= []
    n_atoms = s.n_atoms
    for i in range(1, n_res +1):
        n=s.index_of_atom(i, 'N') -1
        ca = s.index_of_atom(i, 'CA') -1
        c= s.index_of_atom(i, 'C') -1
        mover= mc.DoubleTorsionMover(n, ca, list(range(ca,n_atoms)),
                                     ca, c, list(range(c, n_atoms)))

        movers.append((mover,1))
    sched= mc.MonteCarloScheduler(movers, n_res *60)

    #
    # create the simulation options
    #
    options = system.RunOptions()
    options.implicit_solvent_model = 'gbNeck2'
    options.use_big_timestep = False
    options.use_bigger_timestep = True
    options.cutoff = 1.8

    options.use_amap = False
    options.amap_alpha_bias = 1.0
    options.amap_beta_bias = 1.0
    options.timesteps = 11111
    options.minimize_steps = 20000

    #
    # MC minimizer?
    #
    options.min_mc = None
    options.run_mc = None
    #options.min_mc = sched

    #
    # Handle how to store the data
    #
    store = vault.DataStore(s.n_atoms, N_REPLICAS, s.get_pdb_writer(), block_size=BLOCK_SIZE)
    store.initialize(mode='w')
    store.save_system(s)
    store.save_run_options(options)

    #
    # Adaptation policy and REMD protocol
    #
    l = ladder.NearestNeighborLadder(n_trials=100)
    policy = adaptor.AdaptationPolicy(2.0, 50, 50)
    a = adaptor.EqualAcceptanceAdaptor(n_replicas=N_REPLICAS, adaptation_policy=policy)

    #
    # create and store the remd_runner
    #
    remd_runner = leader.LeaderReplicaExchangeRunner(N_REPLICAS, max_steps=N_STEPS, ladder=l, adaptor=a)
    store.save_remd_runner(remd_runner)

    #
    # create and store the communicator
    #
    c = comm.MPICommunicator(s.n_atoms, N_REPLICAS)
    store.save_communicator(c)

    #
    # create and save the initial states
    #
    states = [gen_state_templates(i,templates) for i in range(N_REPLICAS)]
    store.save_states(states, 0)

    #
    # save data_store
    #
    store.save_data_store()

    return s.n_atoms

setup_system()
'''