# Copyright (c) Stanford University, The Regents of the University of
#               California, and others.
#
# All Rights Reserved.
#
# See Copyright-SimVascular.txt for additional details.
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject
# to the following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
# IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
# TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER
# OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
import os
from sv_ml.base.model import AbstractModel
import numpy as np
import tensorflow as tf
import sv_ml.modules.layers as tf_util

# import matplotlib
# matplotlib.use('Agg')
# import matplotlib.pyplot as plt

import sv_ml.modules.vessel_regression as vr

def get_batch(X,Y, batch_size=16):
    ids = np.random.choice(X.shape[0], size=batch_size)

    x   = np.array([X[i] for i in ids])
    y   = np.array([Y[i] for i in ids])

    return x,y

class Model(AbstractModel):
    def setup(self):
        self.start_iter = 0
        self.build_model()
        self.configure_trainer()
        self.finalize()

        self.losses = []
        self.iters = []

    def train_step(self,x,y):
        self.global_step = self.global_step+1

        if np.sum(np.isnan(x)) > 0: return
        if np.sum(np.isnan(y)) > 0: return

        self.sess.run(self.train_op,{self.x:x,self.y:y})

    def save(self, model_path=None):
        if model_path == None:
            model_path = os.path.join(self.config['MODEL_DIR'], self.config['MODEL_NAME'])
        else:
            model_path = os.path.join(model_path, self.config['MODEL_NAME'])
        self.saver.save(self.sess,model_path)

    def load(self, model_path=None):
        if model_path == None:
            model_path = os.path.join(self.config['MODEL_DIR'],self.config['MODEL_NAME'])
        else:
            model_path = os.path.join(model_path,self.config['MODEL_NAME'])
        self.saver.restore(self.sess, model_path)

    def predict(self,x):
        S = list(x.shape)
        if len(S) == 3:
            x_ = x.reshape([1]+S)
            return self._predict(x_)[0]
        else:
            out = []
            for i in range(S[0]):
                x_ = x[i].reshape([1]+S[1:4])
                y = self._predict(x_)[0].copy()
                out.append(y)
            return np.array(out)

    def calculate_loss(self,x,y):
        return self.sess.run(self.loss,{self.x:x,self.y:y})

    def build_model(self):
        raise RuntimeError("Abstract not implemented")

    def build_loss(self):
        self.loss = tf.reduce_mean(tf.square(self.y-self.yhat))

    def configure_trainer(self):
        LEARNING_RATE = self.config["LEARNING_RATE"]
        self.global_step = tf.Variable(0, trainable=False)
        boundaries = [2000,
                      5000,
                      10000,
                      15000,
                      150000]

        values = [LEARNING_RATE,
                  LEARNING_RATE/3,
                  LEARNING_RATE/10,
                  LEARNING_RATE/100,
                  LEARNING_RATE/1000,
                  LEARNING_RATE/10000,]

        learning_rate = tf.train.piecewise_constant(self.global_step, boundaries, values)

        self.opt = tf.train.AdamOptimizer(learning_rate)

        #self.opt = tf.train.MomentumOptimizer(learning_rate, momentum=0.9)
        self.train_op = self.opt.minimize(self.loss)

    def train(self, X,Y):
        for i in range(self.start_iter,
            self.config['TRAIN_STEPS']+self.start_iter):
            x,y = get_batch(X,Y, self.config['BATCH_SIZE'])

            self.train_step(x,y)

            if i % self.config['LOG_STEP'] == 0:
                l = self.calculate_loss(x,y)
                self.losses.append(l)
                self.iters.append(i)

                self.log(i,x,y)
                self.log(i,X[:4],Y[:4])
                self.save()

    def _predict(self,x):
        return self.sess.run(self.yhat,{self.x:x})

    def finalize(self):
        self.sess = tf.Session()
        self.sess.run(tf.global_variables_initializer())

    def log(self,i,x,y):
        pass
        # f = open(self.config['ITER_FILE'], 'w')
        # f.write(str(i))
        # f.close()
        #
        # l = self.calculate_loss(x,y)
        # yhat = self.predict(x)[0]
        #
        # print("{}: loss={}\n".format(i,l))
        # print("yhat = {}".format(yhat))
        #
        # f = open(self.config["LOG_FILE"],"a+")
        # f.write("{}: loss={}\n".format(i,l))
        # f.write("{}: yhat={}\n".format(i,yhat))
        # f.close()
        #
        # self.save()
        #
        # x_ = x[0,:,:,0]
        # y_ = y[0]
        # if self.config['DATASET'] == 'axial2d_point':
        #     ctrue = vr.point_pred_to_contour(y_)
        #     cpred = vr.point_pred_to_contour(yhat)
        #
        # else:
        #     ctrue = vr.pred_to_contour(y_)
        #     cpred = vr.pred_to_contour(yhat)
        #
        # plt.figure()
        # plt.imshow(x_,cmap='gray',extent=[-1, 1, 1, -1])
        # plt.colorbar()
        # plt.scatter(cpred[:,0], cpred[:,1], color='r', label='predicted',s=4)
        # plt.scatter(ctrue[:,0], ctrue[:,1], color='y', label='true', s=4)
        # plt.show()
        # plt.close()
        #
        # W = x_.shape[0]
        # s = int(0.25*W)
        # e = int(0.75*W)
        #
        # plt.figure()
        # plt.imshow(x_[s:e,s:e],cmap='gray',extent=[-0.5, 0.5, 0.5, -0.5])
        # plt.colorbar()
        # plt.scatter(cpred[:,0], cpred[:,1], color='r', label='predicted',s=4)
        # plt.scatter(ctrue[:,0], ctrue[:,1], color='y', label='true', s=4)
        # plt.show()
        # plt.close()
        #
        #
        # c = self.config
        # log_dir = c['RESULTS_DIR']+'/'+c['NAME']+'/log/'
        #
        # plt.figure()
        # plt.plot(self.iters, self.losses)
        # plt.savefig(log_dir+'loss.png',dpi=300)
        # plt.close()
        #
        # plt.figure()
        # plt.plot(self.iters, self.losses)
        # plt.ylim(0,1)
        # plt.savefig(log_dir+'loss_1.png',dpi=300)
        # plt.close()
        #
        # plt.figure()
        # plt.plot(self.iters, self.losses)
        # plt.ylim(0,0.1)
        # plt.savefig(log_dir+'loss_2.png',dpi=300)
        # plt.close()
        #
        # plt.figure()
        # plt.plot(self.iters, self.losses)
        # plt.ylim(0,0.01)
        # plt.savefig(log_dir+'loss_3.png',dpi=300)
        # plt.close()

class FcNet(Model):
    def build_model(self):
        CROP_DIMS   = self.config['CROP_DIMS']
        LEAK        = self.config['LEAK']
        C           = self.config['NUM_CHANNELS']
        LAMBDA      = self.config['L2_REG']
        INIT        = self.config['INIT']
        NUM_POINTS  = self.config['NUM_CONTOUR_POINTS']

        leaky_relu = tf.contrib.keras.layers.LeakyReLU(LEAK)

        self.x = tf.placeholder(shape=[None,CROP_DIMS,CROP_DIMS,C],dtype=tf.float32)
        self.y = tf.placeholder(shape=[None,NUM_POINTS],dtype=tf.float32)

        o = self.x
        if "INPUT_POOL" in self.config:
            d = self.config['INPUT_POOL']

            o = tf.nn.pool(o, [d,d], "MAX", "VALID", strides=[d,d])

        s = o.get_shape().as_list()

        o_vec = tf.reshape(o,shape=[-1,s[1]*s[2]*s[3]])

        for i,h in enumerate(self.config['HIDDEN_SIZES']):

            o_vec = tf_util.fullyConnected(o_vec, h,
                leaky_relu, std=INIT, scope='fc_'+str(i))

            if "DROPOUT" in self.config:
                o_vec = tf.nn.dropout(o_vec, self.config['DROPOUT'])

        self.yhat = tf_util.fullyConnected(o_vec, NUM_POINTS,
            tf.sigmoid, std=INIT, scope='fc_final')

        self.build_loss()

        self.saver = tf.train.Saver()

class I2INetReg(Model):
    def build_model(self):
        CROP_DIMS   = self.config['CROP_DIMS']
        C           = self.config['NUM_CHANNELS']
        LEAK        = self.config['LEAK']
        NUM_FILTERS = self.config['NUM_FILTERS']
        LAMBDA      = self.config['L2_REG']
        INIT        = self.config['INIT']
        NUM_POINTS  = self.config['NUM_CONTOUR_POINTS']

        leaky_relu = tf.contrib.keras.layers.LeakyReLU(LEAK)

        self.x = tf.placeholder(shape=[None,CROP_DIMS,CROP_DIMS,C],dtype=tf.float32)
        self.y = tf.placeholder(shape=[None,NUM_POINTS],dtype=tf.float32)

        self.yclass,self.yhat,_,_ = tf_util.I2INet(self.x,nfilters=NUM_FILTERS,
            activation=leaky_relu,init=INIT)

        o = leaky_relu(self.yhat)

        s = o.get_shape().as_list()

        o_vec = tf.reshape(o,shape=[-1,s[1]*s[2]*s[3]])

        for i in range(self.config['FC_LAYERS']-1):
            if "HIDDEN_SIZES" in self.config:
                h = self.config['HIDDEN_SIZES'][i]
            else:
                h = self.config['HIDDEN_SIZE']

            o_vec = tf_util.fullyConnected(o_vec, h,
                leaky_relu, std=INIT, scope='fc_'+str(i))

        self.yhat = tf_util.fullyConnected(o_vec, NUM_POINTS,
            tf.identity, std=INIT, scope='fc_final')

        self.build_loss()

        self.saver = tf.train.Saver()

class ResNetReg(Model):
    def build_model(self):
        CROP_DIMS   = self.config['CROP_DIMS']
        C           = self.config['NUM_CHANNELS']
        LEAK        = self.config['LEAK']
        LAMBDA      = self.config['L2_REG']
        INIT        = self.config['INIT']

        NLAYERS     = int(self.config['NLAYERS']/2)
        NFILTERS_SMALL = self.config['NFILTERS_SMALL']
        NFILTERS_LARGE = self.config['NFILTERS_LARGE']

        NUM_POINTS  = self.config['NUM_CONTOUR_POINTS']

        leaky_relu = tf.contrib.keras.layers.LeakyReLU(LEAK)

        self.x = tf.placeholder(shape=[None,CROP_DIMS,CROP_DIMS,C],dtype=tf.float32)
        self.y = tf.placeholder(shape=[None,NUM_POINTS],dtype=tf.float32)

        self.yclass,self.yhat,_,_ = tf_util.resNet(self.x,
            nlayers_before=NLAYERS, nlayers_after=NLAYERS,
            nfilters=NFILTERS_SMALL, nfilters_large=NFILTERS_LARGE,
            output_filters=NFILTERS_LARGE, activation=leaky_relu, init=INIT)

        o = leaky_relu(self.yhat)

        d = self.config['POOL']

        o = tf.nn.pool(o, [d,d], "MAX", "VALID", strides=[d,d])

        s = o.get_shape().as_list()

        o_vec = tf.reshape(o,shape=[-1,s[1]*s[2]*s[3]])

        for i in range(self.config['FC_LAYERS']):
            if "HIDDEN_SIZES" in self.config:
                h = self.config['HIDDEN_SIZES'][i]
            else:
                h = self.config['HIDDEN_SIZE']

            o_vec = tf_util.fullyConnected(o_vec, h,
                leaky_relu, std=INIT, scope='fc_'+str(i))

        self.yhat = tf_util.fullyConnected(o_vec, NUM_POINTS,
            tf.sigmoid, std=INIT, scope='fc_final')

        self.build_loss()

        self.saver = tf.train.Saver()

class ResNetRegMultiscale(Model):
    def build_model(self):
        CROP_DIMS   = self.config['CROP_DIMS']
        C           = self.config['NUM_CHANNELS']
        LEAK        = self.config['LEAK']
        LAMBDA      = self.config['L2_REG']
        INIT        = self.config['INIT']

        NLAYERS     = int(self.config['NLAYERS']/2)
        NFILTERS_SMALL = self.config['NFILTERS_SMALL']
        NFILTERS_LARGE = self.config['NFILTERS_LARGE']

        NUM_POINTS  = self.config['NUM_CONTOUR_POINTS']

        leaky_relu = tf.contrib.keras.layers.LeakyReLU(LEAK)

        self.x = tf.placeholder(shape=[None,CROP_DIMS,CROP_DIMS,C],dtype=tf.float32)

        if self.config['MULTI_TYPE'] == "POOL":
            self.x_1 = tf.nn.pool(self.x, [2,2], "MAX", "VALID", strides=[2,2])
            self.x_2 = tf.nn.pool(self.x_1, [2,2], "MAX", "VALID", strides=[2,2])
        elif self.config['MULTI_TYPE'] == "CROP":
            self.x_1 = tf.image.central_crop(self.x, central_fraction=0.5)
            self.x_2 = tf.image.central_crop(self.x_1, central_fraction=0.5)
        else:
            raise RuntimeError("Unrecognized multi type")

        self.y = tf.placeholder(shape=[None,NUM_POINTS],dtype=tf.float32)

        self.yclass,self.yhat,_,_ = tf_util.resNet(self.x,
            nlayers_before=NLAYERS, nlayers_after=NLAYERS,
            nfilters=NFILTERS_SMALL, nfilters_large=NFILTERS_LARGE,
            output_filters=NFILTERS_LARGE, activation=leaky_relu, init=INIT)

        self.yclass_1,self.yhat_1,_,_ = tf_util.resNet(self.x_1,
            nlayers_before=NLAYERS, nlayers_after=NLAYERS,
            nfilters=NFILTERS_SMALL, nfilters_large=NFILTERS_LARGE,
            output_filters=NFILTERS_LARGE, activation=leaky_relu, init=INIT,
            scope="resnet_1")

        self.yclass_2,self.yhat_2,_,_ = tf_util.resNet(self.x_2,
            nlayers_before=NLAYERS, nlayers_after=NLAYERS,
            nfilters=NFILTERS_SMALL, nfilters_large=NFILTERS_LARGE,
            output_filters=NFILTERS_LARGE, activation=leaky_relu, init=INIT,
            scope="resnet_2")

        o   = leaky_relu(self.yhat)
        o_1 = leaky_relu(self.yhat_1)
        o_2 = leaky_relu(self.yhat_2)

        s   = o.get_shape().as_list()
        s_1 = o_1.get_shape().as_list()
        s_2 = o_2.get_shape().as_list()

        o_vec   = tf.reshape(o,shape=[-1,s[1]*s[2]*s[3]])
        o_vec_1 = tf.reshape(o_1,shape=[-1,s_1[1]*s_1[2]*s_1[3]])
        o_vec_2 = tf.reshape(o_2,shape=[-1,s_2[1]*s_2[2]*s_2[3]])

        o = tf.concat([o_vec, o_vec_1, o_vec_2], axis=1)

        print(o)

        for i in range(self.config['FC_LAYERS']-1):
            if "HIDDEN_SIZES" in self.config:
                h = self.config['HIDDEN_SIZES'][i]
            else:
                h = self.config['HIDDEN_SIZE']

            o = tf_util.fullyConnected(o, h,
                leaky_relu, std=INIT, scope='fc_'+str(i))

        self.yhat = tf_util.fullyConnected(o, NUM_POINTS,
            tf.identity, std=INIT, scope='fc_final')

        self.build_loss()

        self.saver = tf.train.Saver()

class ConvNet(Model):
    def build_model(self):
        CROP_DIMS   = self.config['CROP_DIMS']
        C           = self.config['NUM_CHANNELS']
        LEAK        = self.config['LEAK']
        LAMBDA      = self.config['L2_REG']
        INIT        = self.config['INIT']

        NLAYERS     = self.config['NLAYERS']
        NFILTERS    = self.config['NFILTERS']

        NUM_POINTS  = self.config['NUM_CONTOUR_POINTS']
        DIMS = [self.config['CONV_DIMS']]*2

        leaky_relu = tf.contrib.keras.layers.LeakyReLU(LEAK)

        self.x = tf.placeholder(shape=[None,CROP_DIMS,CROP_DIMS,C],dtype=tf.float32)
        self.y = tf.placeholder(shape=[None,NUM_POINTS],dtype=tf.float32)

        o = self.x

        if "INPUT_POOL" in self.config:
            d = self.config['INPUT_POOL']

            o = tf.nn.pool(o, [d,d], "MAX", "VALID", strides=[d,d])

        for i in range(NLAYERS):
            o = tf_util.conv2D(o,dims=DIMS,nfilters=NFILTERS,
                               init=INIT,
                          activation=leaky_relu,scope="conv_{}".format(i))

        s   = o.get_shape().as_list()
        o = tf.reshape(o,shape=[-1,s[1]*s[2]*s[3]])

        for i in range(self.config['FC_LAYERS']):
            if "HIDDEN_SIZES" in self.config:
                h = self.config['HIDDEN_SIZES'][i]
            else:
                h = self.config['HIDDEN_SIZE']

            o = tf_util.fullyConnected(o, h,
                leaky_relu, std=INIT, scope='fc_'+str(i))
            print(o)

            if "DROPOUT" in self.config:
                o = tf.nn.dropout(o, self.config['DROPOUT'])

        self.yhat = tf_util.fullyConnected(o, NUM_POINTS,
            tf.sigmoid, std=INIT, scope='fc_final')

        self.build_loss()

        self.saver = tf.train.Saver()

class ConvNetMulti(Model):
    def build_model(self):
        CROP_DIMS   = self.config['CROP_DIMS']
        C           = self.config['NUM_CHANNELS']
        LEAK        = self.config['LEAK']
        LAMBDA      = self.config['L2_REG']
        INIT        = self.config['INIT']

        NLAYERS     = int(self.config['NLAYERS']/2)
        NFILTERS    = self.config['NFILTERS']

        NUM_POINTS  = self.config['NUM_CONTOUR_POINTS']
        DIMS = [self.config['CONV_DIMS']]*2

        leaky_relu = tf.contrib.keras.layers.LeakyReLU(LEAK)

        self.x = tf.placeholder(shape=[None,CROP_DIMS,CROP_DIMS,C],dtype=tf.float32)
        self.y = tf.placeholder(shape=[None,NUM_POINTS],dtype=tf.float32)

        self.x_1 = tf.nn.pool(self.x, [2,2], "MAX", "VALID", strides=[2,2])
        self.x_2 = tf.nn.pool(self.x_1, [2,2], "MAX", "VALID", strides=[2,2])

        o = self.x

        for i in range(NLAYERS):
            o = tf_util.conv2D(o,dims=DIMS, nfilters=NFILTERS,init=INIT,activation=leaky_relu,
                scope="conv_{}".format(i))

        o_1 = self.x_1

        for i in range(NLAYERS):
            o_1 = tf_util.conv2D(o_1,dims=DIMS, nfilters=NFILTERS,init=INIT,activation=leaky_relu, scope="conv_1_{}".format(i))

        o_2 = self.x_2

        for i in range(NLAYERS):
            o_2 = tf_util.conv2D(o_2,dims=DIMS, nfilters=NFILTERS,init=INIT,activation=leaky_relu, scope="conv_2_{}".format(i))

        s = o.get_shape().as_list()
        s_1 = o_1.get_shape().as_list()
        s_2 = o_2.get_shape().as_list()

        o_vec   = tf.reshape(o,shape=[-1,s[1]*s[2]*s[3]])
        o_vec_1 = tf.reshape(o_1,shape=[-1,s_1[1]*s_1[2]*s_1[3]])
        o_vec_2 = tf.reshape(o_2,shape=[-1,s_2[1]*s_2[2]*s_2[3]])

        o = tf.concat([o_vec, o_vec_1, o_vec_2], axis=1)

        for i in range(self.config['FC_LAYERS']-1):
            if "HIDDEN_SIZES" in self.config:
                h = self.config['HIDDEN_SIZES'][i]
            else:
                h = self.config['HIDDEN_SIZE']

            o = tf_util.fullyConnected(o, h,
                leaky_relu, std=INIT, scope='fc_'+str(i))

        self.yhat = tf_util.fullyConnected(o, NUM_POINTS,
            tf.identity, std=INIT, scope='fc_final')

        self.build_loss()

        self.saver = tf.train.Saver()

class GoogleNet(Model):
    def build_model(self):
        CROP_DIMS   = self.config['CROP_DIMS']
        C           = self.config['NUM_CHANNELS']
        LEAK        = self.config['LEAK']
        LAMBDA      = self.config['L2_REG']
        INIT        = self.config['INIT']
        DROPOUT     = self.config['DROPOUT']
        NUM_POINTS  = self.config['NUM_CONTOUR_POINTS']

        leaky_relu = tf.contrib.keras.layers.LeakyReLU(LEAK)

        self.x = tf.placeholder(shape=[None,CROP_DIMS,CROP_DIMS,C],dtype=tf.float32)
        self.y = tf.placeholder(shape=[None,NUM_POINTS],dtype=tf.float32)

        o,o_side = tf_util.GoogleNet(self.x, activation=leaky_relu, init=INIT,
            scope='googlenet', output_size=NUM_POINTS, dropout=DROPOUT)

        print(o)
        print(o_side)

        self.yhat = tf.nn.sigmoid(o)
        self.yhat_side = tf.nn.sigmoid(o_side)

        self.build_loss()

        self.saver = tf.train.Saver()

        self.dropout_mask_op = op = tf.get_default_graph().get_tensor_by_name(
            "googlenet/dropout_1/random_uniform:0")

        self.dropout_mask = None
        self.dropout_fixed = False

    def build_loss(self):
        self.loss = tf.reduce_mean(tf.square(self.y-self.yhat))
        self.loss += 0.3*tf.reduce_mean(tf.square(self.y-self.yhat_side))

    def sample(self):
        self.dropout_mask = (np.random.uniform(size=(1,9216))<=0.6).astype(int)
        self.dropout_fixed = True

    def _predict(self,x):
        if not self.dropout_fixed:
            return self.sess.run(self.yhat,{self.x:x})
        else:
            return self.sess.run(self.yhat,{self.x:x,
                self.dropout_mask_op:self.dropout_mask})


# -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2014-2018 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake 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 Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.

"""
Module exports :class:`YuEtAl2013`, :class:`YuEtAl2013Tibet`,
:class:`YuEtAl2013Eastern`, :class:`YuEtAl2013Stable`

"""
import numpy as np
from scipy.constants import g

from openquake.hazardlib.gsim.base import GMPE, CoeffsTable
from openquake.hazardlib import const
from openquake.hazardlib.imt import PGA, PGV, SA

def gc(coeff, mag):
    """
    Returns the set of coefficients to be used for the calculation of GM
    as a function of earthquake magnitude

    :param coeff:
        A dictionary of parameters for the selected IMT
    :param mag:
        Magnitude value
    :returns:
        The set of coefficients
    """
    if mag > 6.5:
        a1ca = coeff['ua']
        a1cb = coeff['ub']
        a1cc = coeff['uc']
        a1cd = coeff['ud']
        a1ce = coeff['ue']
        a2ca = coeff['ia']
        a2cb = coeff['ib']
        a2cc = coeff['ic']
        a2cd = coeff['id']
        a2ce = coeff['ie']
    else:
        a1ca = coeff['a']
        a1cb = coeff['b']
        a1cc = coeff['c']
        a1cd = coeff['d']
        a1ce = coeff['e']
        a2ca = coeff['ma']
        a2cb = coeff['mb']
        a2cc = coeff['mc']
        a2cd = coeff['md']
        a2ce = coeff['me']
    return a1ca, a1cb, a1cc, a1cd, a1ce, a2ca, a2cb, a2cc, a2cd, a2ce

def rbf(ra, coeff, mag):
    """
    Calculate the median ground motion for a given magnitude and distance

    :param ra:
        Distance value [km]
    :param coeff:
        The set of coefficients
    :param mag:
        Magnitude value
    :returns:

    """
    a1ca, a1cb, a1cc, a1cd, a1ce, a2ca, a2cb, a2cc, a2cd, a2ce = gc(coeff, mag)
    term1 = a1ca + a1cb * mag + a1cc * np.log(ra + a1cd*np.exp(a1ce*mag))
    term2 = a2ca + a2cb * mag
    term3 = a2cd*np.exp(a2ce*mag)
    return np.exp((term1 - term2) / a2cc) - term3

def fnc(ra, *args):
    """
    Function used in the minimisation problem.

    :param ra:
        Semi-axis of the ellipses used in the Yu et al.
    :returns:
        The absolute difference between the epicentral distance and the
        adjusted distance
    """
    #
    # epicentral distance
    repi = args[0]
    #
    # azimuth
    theta = args[1]
    #
    # magnitude
    mag = args[2]
    #
    # coefficients
    coeff = args[3]
    #
    # compute the difference between epicentral distances
    rb = rbf(ra, coeff, mag)
    t1 = ra**2 * (np.sin(np.radians(theta)))**2
    t2 = rb**2 * (np.cos(np.radians(theta)))**2
    xx = ra * rb / (t1+t2)**0.5
    return xx-repi

def get_ras(repi, theta, mag, coeff):
    """
    Computes equivalent distance

    :param repi:
        Epicentral distance
    :param theta:
        Azimuth value
    :param mag:
        Magnitude
    :param coeff:
        GMPE coefficients
    """
    rx = 150.
    ras = 300.
    dff = 1.e0
    while abs(dff) > 1e-5:
        #
        # calculate the difference between epicentral distances
        dff = fnc(ras, repi, theta, mag, coeff)
        #
        # update the value of distance computed
        ras -= np.sign(dff) * rx
        rx = rx / 2.
        if rx < 1e-3:
            break
    return ras

class YuEtAl2013Ms(GMPE):
    """
    Implements the Yu et al. (2013) GMPE used for the calculation of the 2015
    version of the national seismic hazard maps for China. Note that magnitude
    supported is Ms.
    """

    #: Supported tectonic region type is active shallow crust
    DEFINED_FOR_TECTONIC_REGION_TYPE = const.TRT.ACTIVE_SHALLOW_CRUST

    #: Supported intensity measure types are peak ground velocity and
    #: peak ground acceleration
    DEFINED_FOR_INTENSITY_MEASURE_TYPES = set([
        PGA,
        PGV,
        SA
    ])

    #: Supported intensity measure component is geometric mean (supposed)
    DEFINED_FOR_INTENSITY_MEASURE_COMPONENT = const.IMC.AVERAGE_HORIZONTAL

    #: Supported standard deviation types is total
    DEFINED_FOR_STANDARD_DEVIATION_TYPES = set([
        const.StdDev.TOTAL
    ])

    #: No site parameters required
    REQUIRES_SITES_PARAMETERS = set(())

    #: Required rupture parameter is magnitude
    REQUIRES_RUPTURE_PARAMETERS = set(('mag',))

    #: Required distance measures are epicentral distance and azimuth
    REQUIRES_DISTANCES = set(('repi', 'azimuth'))

    def get_mean_and_stddevs(self, sites, rup, dists, imt, stddev_types):
        """
        See :meth:`superclass method
        <.base.GroundShakingIntensityModel.get_mean_and_stddevs>`
        for spec of input and result values.
        """
        # Check that the requested standard deviation type is available
        assert all(stddev_type in self.DEFINED_FOR_STANDARD_DEVIATION_TYPES
                    for stddev_type in stddev_types)
        #
        # Set parameters
        mag = rup.mag
        epi = dists.repi
        theta = dists.azimuth
        #
        # Set coefficients
        coeff = self.COEFFS[imt]
        a1ca, a1cb, a1cc, a1cd, a1ce, a2ca, a2cb, a2cc, a2cd, a2ce = \
            gc(coeff, mag)
        #
        # Get correction coefficients. Here for each site we find the
        # the geometry of the ellipses
        ras = []
        for epi, theta in zip(dists.repi, dists.azimuth):
            res = get_ras(epi, theta, mag, coeff)
            ras.append(res)
        ras = np.array(ras)
        rbs = rbf(ras, coeff, mag)
        #
        # Compute values of ground motion for the two cases. The value of
        # 225 is hardcoded under the assumption that the hypocentral depth
        # corresponds to 15 km (i.e. 15**2)
        mean1 = (a1ca + a1cb * mag +
                    a1cc * np.log((ras**2+225)**0.5 +
                                a1cd * np.exp(a1ce * mag)))
        mean2 = (a2ca + a2cb * mag +
                    a2cc * np.log((rbs**2+225)**0.5 +
                                a2cd * np.exp(a2ce * mag)))
        #
        # Get distances
        x = (mean1 * np.sin(np.radians(dists.azimuth)))**2
        y = (mean2 * np.cos(np.radians(dists.azimuth)))**2
        mean = mean1 * mean2 / np.sqrt(x+y)
        if isinstance(imt, (PGA)):
            mean = np.exp(mean)/g/100
        elif isinstance(imt, (PGV)):
            mean = np.exp(mean)
        else:
            raise ValueError('Unsupported IMT')
        #
        # Get the standard deviation
        stddevs = self._compute_std(coeff, stddev_types, len(dists.repi))
        #
        # Return results
        return np.log(mean), stddevs

    def _compute_std(self, C, stddev_types, num_sites):
        return [np.ones(num_sites)*C['sigma']]

    #: Coefficient table
    COEFFS = CoeffsTable(sa_damping=5, table="""\
IMT a b c d e ua ub uc ud ue ma mb mc md me ia ib ic id ie sigma
PGA 4.1193 1.656 -2.389 1.772 0.424 7.8269 1.0856 -2.389 1.772 0.424 2.2609 1.6399 -2.118 0.825 0.465 6.003 1.0649 -2.118 0.825 0.465 0.5428
PGV -1.2581 1.932 -2.181 1.772 0.424 3.013 1.2742 -2.181 1.772 0.424 -3.1073 1.9389 -1.945 0.825 0.465 1.3087 1.2627 -1.945 0.825 0.465 0.6233
        """)

class YuEtAl2013MsTibet(YuEtAl2013Ms):
    #: Supported tectonic region type is Tibetan plateau
    DEFINED_FOR_TECTONIC_REGION_TYPE = const.TRT.ACTIVE_SHALLOW_CRUST
    #: Coefficient table
    COEFFS = CoeffsTable(sa_damping=5, table="""\
IMT a b c d e ua ub uc ud ue ma mb mc md me ia ib ic id ie sigma
PGA 5.4901 1.4835 -2.416 2.647 0.366 8.7561 0.9453 -2.416 2.647 0.366 2.3069 1.4007 -1.854 0.612 0.457 5.6511 0.8924 -1.854 0.612 0.457 0.5428
PGV -0.1472 1.7618 -2.205 2.647 0.366 3.9422 1.1293 -2.205 2.647 0.366 -2.9923 1.7043 -1.696 0.612 0.457 1.0189 1.0902 -1.696 0.612 0.457 0.6233
     """)

class YuEtAl2013MsEastern(YuEtAl2013Ms):
    #: Supported tectonic region type is eastern part of China
    DEFINED_FOR_TECTONIC_REGION_TYPE = const.TRT.STABLE_CONTINENTAL
    #: Coefficient table
    COEFFS = CoeffsTable(sa_damping=5, table="""\
IMT a b c d e ua ub uc ud ue ma mb mc md me ia ib ic id ie sigma
PGA 4.5517 1.5433 -2.315 2.088 0.399 8.1259 0.9936 -2.315 2.088 0.399 2.7048 1.518 -2.004 0.944 0.447 6.3319 0.9614 -2.004 0.944 0.447 0.5428
PGV -0.8349 1.8193 -2.103 2.088 0.399 3.3051 1.1799 -2.103 2.088 0.399 -2.6381 1.8124 -1.825 0.944 0.447 1.6376 1.1546 -1.825 0.944 0.447 0.6233
     """)

class YuEtAl2013MsStable(YuEtAl2013Ms):
    #: Supported tectonic region type is stable part of China
    DEFINED_FOR_TECTONIC_REGION_TYPE = const.TRT.STABLE_CONTINENTAL
    #: Coefficient table
    COEFFS = CoeffsTable(sa_damping=5, table="""\
IMT a b c d e ua ub uc ud ue ma mb mc md me ia ib ic id ie sigma
PGA 5.5591 1.1454 -2.079 2.802 0.295 8.5238 0.6854 -2.079 2.802 0.295 3.9445 1.0833 -1.723 1.295 0.331 6.187 0.7383 -1.723 1.295 0.331 0.5428
PGV 0.2139 1.4283 -1.889 2.802 0.295 3.772 0.8786 -1.889 2.802 0.295 -1.3547 1.3823 -1.559 1.295 0.331 1.5433 0.9361 -1.559 1.295 0.331 0.6233
     """)

class YuEtAl2013Mw(YuEtAl2013Ms):
    """
    This is a modified version of the original Yu et al. (2013) that supports
    the use of Mw rather than Ms. The Mw to Ms conversion equation used is the
    one proposed by Cheng et al. (2017). Note that this version does not
    propagate the uncertainty related to the magnitude conversion process.
    """

    def get_mean_and_stddevs(self, sites, rup, dists, imt, stddev_types):
        """
        See :meth:`superclass method
        <.base.GroundShakingIntensityModel.get_mean_and_stddevs>`
        for spec of input and result values.
        """
        # Check that the requested standard deviation type is available
        assert all(stddev_type in self.DEFINED_FOR_STANDARD_DEVIATION_TYPES
                   for stddev_type in stddev_types)
        #
        # Set parameters
        magn = rup.mag
        epi = dists.repi
        theta = dists.azimuth
        #
        # Convert Mw into Ms
        if magn < 6.58:
            mag = (magn - 0.59) / 0.86
        else:
            mag = (magn + 2.42) / 1.28
        #
        # Set coefficients
        coeff = self.COEFFS[imt]
        a1ca, a1cb, a1cc, a1cd, a1ce, a2ca, a2cb, a2cc, a2cd, a2ce = \
            gc(coeff, mag)
        #
        # Get correction coefficients. Here for each site we find the
        # the geometry of the ellipses
        ras = []
        for epi, theta in zip(dists.repi, dists.azimuth):
            res = get_ras(epi, theta, mag, coeff)
            ras.append(res)
        ras = np.array(ras)
        rbs = rbf(ras, coeff, mag)
        #
        # Compute values of ground motion for the two cases. The value of
        # 225 is hardcoded under the assumption that the hypocentral depth
        # corresponds to 15 km (i.e. 15**2)
        mean1 = (a1ca + a1cb * mag +
                 a1cc * np.log((ras**2+225)**0.5 +
                               a1cd * np.exp(a1ce * mag)))
        mean2 = (a2ca + a2cb * mag +
                 a2cc * np.log((rbs**2+225)**0.5 +
                               a2cd * np.exp(a2ce * mag)))
        #
        # Get distances
        x = (mean1 * np.sin(np.radians(dists.azimuth)))**2
        y = (mean2 * np.cos(np.radians(dists.azimuth)))**2
        mean = mean1 * mean2 / np.sqrt(x+y)
        if isinstance(imt, (PGA)):
            mean = np.exp(mean)/g/100
        elif isinstance(imt, (PGV)):
            mean = np.exp(mean)
        else:
            raise ValueError('Unsupported IMT')
        #
        # Get the standard deviation
        stddevs = self._compute_std(coeff, stddev_types, len(dists.repi))
        #
        # Return results
        return np.log(mean), stddevs

class YuEtAl2013MwTibet(YuEtAl2013Mw):
    #: Supported tectonic region type is Tibetan plateau
    DEFINED_FOR_TECTONIC_REGION_TYPE = const.TRT.ACTIVE_SHALLOW_CRUST
    #: Coefficient table
    COEFFS = CoeffsTable(sa_damping=5, table="""\
IMT a b c d e ua ub uc ud ue ma mb mc md me ia ib ic id ie sigma
PGA 5.4901 1.4835 -2.416 2.647 0.366 8.7561 0.9453 -2.416 2.647 0.366 2.3069 1.4007 -1.854 0.612 0.457 5.6511 0.8924 -1.854 0.612 0.457 0.5428
PGV -0.1472 1.7618 -2.205 2.647 0.366 3.9422 1.1293 -2.205 2.647 0.366 -2.9923 1.7043 -1.696 0.612 0.457 1.0189 1.0902 -1.696 0.612 0.457 0.6233
     """)

class YuEtAl2013MwEastern(YuEtAl2013Mw):
    #: Supported tectonic region type is eastern part of China
    DEFINED_FOR_TECTONIC_REGION_TYPE = const.TRT.STABLE_CONTINENTAL
    #: Coefficient table
    COEFFS = CoeffsTable(sa_damping=5, table="""\
IMT a b c d e ua ub uc ud ue ma mb mc md me ia ib ic id ie sigma
PGA 4.5517 1.5433 -2.315 2.088 0.399 8.1259 0.9936 -2.315 2.088 0.399 2.7048 1.518 -2.004 0.944 0.447 6.3319 0.9614 -2.004 0.944 0.447 0.5428
PGV -0.8349 1.8193 -2.103 2.088 0.399 3.3051 1.1799 -2.103 2.088 0.399 -2.6381 1.8124 -1.825 0.944 0.447 1.6376 1.1546 -1.825 0.944 0.447 0.6233
     """)

class YuEtAl2013MwStable(YuEtAl2013Mw):
    #: Supported tectonic region type is stable part of China
    DEFINED_FOR_TECTONIC_REGION_TYPE = const.TRT.STABLE_CONTINENTAL
    #: Coefficient table
    COEFFS = CoeffsTable(sa_damping=5, table="""\
IMT a b c d e ua ub uc ud ue ma mb mc md me ia ib ic id ie sigma
PGA 5.5591 1.1454 -2.079 2.802 0.295 8.5238 0.6854 -2.079 2.802 0.295 3.9445 1.0833 -1.723 1.295 0.331 6.187 0.7383 -1.723 1.295 0.331 0.5428
PGV 0.2139 1.4283 -1.889 2.802 0.295 3.772 0.8786 -1.889 2.802 0.295 -1.3547 1.3823 -1.559 1.295 0.331 1.5433 0.9361 -1.559 1.295 0.331 0.6233
     """)


"""
Customised implementation of FermiDos. Will Move back to pymatgen at some point.
"""
import logging
from typing import Dict, Optional, Tuple, Union

import numpy as np
from monty.json import MSONable
from pymatgen import Spin, Structure
from pymatgen.electronic_structure.dos import Dos

from amset.constants import boltzmann_au, ev_to_hartree, hartree_to_ev
from amset.electronic_structure.fd import fd

__author__ = "Alex Ganose"
__maintainer__ = "Alex Ganose"
__email__ = "aganose@lbl.gov"

logger = logging.getLogger(__name__)

class FermiDos(Dos, MSONable):
    """
    This wrapper class helps relate the density of states, doping levels
    (i.e. carrier concentrations) and corresponding fermi levels. A negative
    doping concentration indicates the majority carriers are electrons
    (n-type doping); a positive doping concentration indicates holes are the
    majority carriers (p-type doping).

    Args:
        efermi: The Fermi level energy in Hartree.
        energies: A sequences of energies in Hartree.
        densities ({Spin: np.array}): representing the density of states
            for each Spin.
        structure: A structure. If not provided, the structure
            of the dos object will be used. If the dos does not have an
            associated structure object, an error will be thrown.
        dos_weight: The weighting for the dos. Defaults to 2 for non-spin
            polarized calculations and 1 for spin-polarized calculations.
        atomic_units: Whether energies are given in eV or Hartree.
        num_electrons: The number of electrons in the system. If None, this will be
            calculated by integrating up to the intrinsic Fermi level.
    """

    def __init__(
        self,
        efermi: float,
        energies: np.ndarray,
        densities: Dict[Spin, np.ndarray],
        structure: Structure,
        dos_weight: Optional[float] = None,
        atomic_units: bool = True,
        num_electrons: Optional[float] = None,
    ):
        # structure should be atomic structure
        super().__init__(efermi, energies, densities)
        self.structure = structure
        self.atomic_units = atomic_units

        if not dos_weight:
            dos_weight = 2 if len(self.densities) == 1 else 1

        self.dos_weight = dos_weight
        self.tdos = np.array(self.get_densities()) * self.dos_weight
        self.de = self.energies[1] - self.energies[0]
        self._num_electrons = num_electrons  # this is just for msonability

        if num_electrons is None:
            # integrate up to Fermi level to get number of electrons
            self.nelect = self.tdos[self.energies <= self.efermi].sum() * self.de
        else:
            self.nelect = num_electrons

        logger.info(
            "Intrinsic DOS Fermi level: {:.4f} eV".format(
                self.efermi * hartree_to_ev if atomic_units else self.efermi
            )
        )
        logger.info("DOS contains {:.3f} electrons".format(self.nelect))

    def get_doping(
        self,
        fermi_level: float,
        temperature: float,
        return_electron_hole_conc: bool = False,
    ) -> Union[float, Tuple[float, float, float]]:
        """
        Calculate the doping (majority carrier concentration) at a given
        fermi level  and temperature. A simple Left Riemann sum is used for
        integrating the density of states over energy & equilibrium Fermi-Dirac
        distribution.

        Args:
            fermi_level: The fermi_level level in Hartree.
            temperature: The temperature in Kelvin.
            return_electron_hole_conc: Whether to also return the separate
                electron and hole concentrations at the doping level.

        Returns:
            If return_electron_hole_conc is False: the doping concentration in
            units of 1/Bohr^3. Negative values indicate that the majority carriers
            are electrons (n-type doping) whereas positive values indicates the
            majority carriers are holes (p-type doping).

            If return_electron_hole_conc is True: the doping concentration,
            electron concentration and hole concentration as a tuple.
        """
        wdos = _get_weighted_dos(
            self.energies,
            self.tdos,
            fermi_level,
            temperature,
            atomic_units=self.atomic_units,
        )

        num_electrons = wdos.sum() * self.de
        conc = (self.nelect - num_electrons) / self.structure.volume

        if return_electron_hole_conc:
            cb_conc = wdos[self.energies > self.efermi].sum() * self.de
            vb_conc = wdos[self.energies <= self.efermi].sum() * self.de
            cb_conc = cb_conc / self.structure.volume
            vb_conc = (self.nelect - vb_conc) / self.structure.volume
            return conc, cb_conc, vb_conc

        else:
            return conc

    def get_num_electrons(self, fermi_level: float, temperature: float) -> float:
        """
        Calculate the number of electrons at a given fermi level and temperature.
        A simple Left Riemann sum is used for integrating the density of states over
        energy & equilibrium Fermi-Dirac distribution.

        Args:
            fermi_level: The fermi_level level in Hartree.
            temperature: The temperature in Kelvin.

        Returns:
            The number of electrons.
        """
        wdos = _get_weighted_dos(
            self.energies,
            self.tdos,
            fermi_level,
            temperature,
            atomic_units=self.atomic_units,
        )

        num_electrons = wdos.sum() * self.de
        return num_electrons

    def get_fermi_from_num_electrons(
        self,
        num_electrons: float,
        temperature: float,
        tol: float = 0.01,
        nstep: int = 50,
        step: float = 0.1,
        precision: int = 10,
    ):
        # this is finding the Fermi level of metals
        fermi = self.efermi  # initialize target fermi
        relative_error = float("inf")
        for _ in range(precision):
            frange = np.arange(-nstep, nstep + 1) * step + fermi
            calc_nelectrons = [self.get_num_electrons(f, temperature) for f in frange]
            relative_error = abs(np.array(calc_nelectrons) / num_electrons - 1.0)
            fermi = frange[np.argmin(relative_error)]
            step /= 10.0

        if min(relative_error) > tol:
            raise ValueError(
                "Could not find fermi within {}% of num electrons={}".format(
                    tol * 100, num_electrons
                )
            )

        return fermi

    def get_fermi(
        self,
        concentration: float,
        temperature: float,
        tol: float = 0.01,
        nstep: int = 50,
        step: float = 0.1,
        precision: int = 10,
        return_electron_hole_conc=False,
    ):
        """
        Finds the fermi level at which the doping concentration at the given
        temperature (T) is equal to concentration. A greedy algorithm is used
        where the relative error is minimized by calculating the doping at a
        grid which continually becomes finer.

        Args:
            concentration: The doping concentration in 1/Bohr^3. Negative values
                represent n-type doping and positive values represent p-type
                doping.
            temperature: The temperature in Kelvin.
            return_electron_hole_conc: Whether to also return the separate
                electron and hole concentrations at the doping level.

        Returns:
            If return_electron_hole_conc is False: The Fermi level in eV. Note
            that this is different from the default dos.efermi.

            If return_electron_hole_conc is True: the Fermi level, electron
            concentration and hole concentration at the Fermi level as a tuple.
            The electron and hole concentrations are in Bohr^-3.
        """
        fermi = self.efermi  # initialize target fermi
        relative_error = float("inf")
        for _ in range(precision):
            frange = np.arange(-nstep, nstep + 1) * step + fermi
            calc_doping = np.array([self.get_doping(f, temperature) for f in frange])
            relative_error = abs(calc_doping / concentration - 1.0)
            fermi = frange[np.argmin(relative_error)]
            step /= 10.0

        if min(relative_error) > tol:
            raise ValueError(
                "Could not find fermi within {}% of concentration={}".format(
                    tol * 100, concentration
                )
            )

        if return_electron_hole_conc:
            _, n_elec, n_hole = self.get_doping(
                fermi, temperature, return_electron_hole_conc=True
            )
            return fermi, n_elec, n_hole
        else:
            return fermi

def _get_weighted_dos(energies, dos, fermi_level, temperature, atomic_units=True):
    if temperature == 0.0:
        occ = np.where(energies < fermi_level, 1.0, 0.0)
        occ[energies == fermi_level] = 0.5
    else:
        kbt = temperature * boltzmann_au
        if atomic_units:
            occ = fd(energies, fermi_level, kbt)
        else:
            occ = fd(energies * ev_to_hartree, fermi_level * ev_to_hartree, kbt)

    wdos = dos * occ
    return wdos


import nest
import os
import subprocess
import sys
import itertools
from scipy.optimize import minimize
import scipy.special as sp_spec
import numpy as np
# from nest_elephant_tvb.simulation.file_tvb.Zerlaut import ZerlautAdaptationSecondOrder as model
from nest_elephant_tvb.Tvb import Matteo_2 as model
from nest_elephant_tvb.Tvb import excitatory

# Matteo function
# # excitatory
# excitatory={
#         'C_m':200.0,
#         't_ref':5.0,
#         'V_reset':-64.5,
#         'E_L':-64.5,
#         'g_L':10.0,
#         'I_e':0.0,
#         'a':0.0,
#         'b':0.0,
#         'Delta_T':2.0,
#         'tau_w':500.0,
#         'V_th':-50.0,
#         'E_ex':0.0,
#         'tau_syn_ex':5.0,
#         'E_in':-80.0,
#         'tau_syn_in':5.0,
#     'V_peak': 10.0,
#     'N_tot':10**4,
#     'p_connect':0.05,
#     'g':0.2,
#     'Q_e':1.0,
#     'Q_i':2.5,
# }
# #inhibitory
# inhibitory={
#         'C_m':200.0,
#         't_ref':5.0,
#         'V_reset':-65.0,
#         'E_L':-65.,
#         'g_L':10.0,
#         'I_e':0.0,
#         'a':0.0,
#         'b':0.0,
#         'Delta_T':0.5,
#         'tau_w':1.0,
#         'V_th':-50.0,
#         'E_ex':0.0,
#         'tau_syn_ex':5.0,
#         'E_in':-80.0,
#         'tau_syn_in':5.0,
#     'V_peak': 10.0,
#     'N_tot':10**4,
#     'p_connect':0.05,
#     'g':0.2,
#     'Q_e':1.0,
#     'Q_i':2.5
# }

def compute_rate(data,begin,end,nb):
    """
    Compute the firing rate
    :param data: the spike of all neurons between end and begin
    :param begin: the time of the first spike
    :param end: the time of the last spike
    :return: the mean and the standard deviation of firing rate, the maximum and minimum of firing rate
    """
    #get data
    n_fil = data[:, 0]
    n_fil = n_fil.astype(int)
    #count the number of the same id
    count_of_n = np.bincount(n_fil)
    #compute the rate
    rate_each_n_incomplet = count_of_n / (end - begin)
    #fill the table with the neurons which are not firing
    rate_each_n = np.concatenate(
        (rate_each_n_incomplet, np.zeros(-np.shape(rate_each_n_incomplet)[0] + nb +1)))
    #save the value

    return rate_each_n[1:]

def load_event(events):
    """
    Get the id of the neurons which create the spike and time
    :param path: the path to the file
    :return: The spike of all neurons
    """
    data_concatenated =  np.concatenate(([events['senders']],[events['times']]))
    if data_concatenated.size < 5:
        print('empty file')
        return None
    data_raw = data_concatenated[np.argsort(data_concatenated[:, 1])]
    return np.swapaxes(data_raw,0,1)

def load_spike(path):
    """
    Get the id of the neurons which create the spike and time
    :param path: the path to the file
    :return: The spike of all neurons
    """
    if not os.path.exists(path + "/spike_detector.gdf"):
        print('no file')
        return None
    data_concatenated = np.loadtxt(path + "/spike_detector.gdf")
    if data_concatenated.size < 5:
        print('empty file')
        return None
    data_raw = data_concatenated[np.argsort(data_concatenated[:, 1])]
    return data_raw

def create_transfer_function(parameter,excitatory):
    model_test = model()
    model_test.g_L = np.array(parameter['g_L'])
    model_test.E_L_e =  np.array(parameter['E_L'])
    model_test.E_L_i = np.array(parameter['E_L'])
    model_test.C_m = np.array(parameter['C_m'])
    model_test.b_e = np.array(parameter['b'])
    model_test.a_e = np.array(parameter['a'])
    model_test.b_i = np.array(parameter['b'])
    model_test.a_i = np.array(parameter['a'])
    model_test.tau_w_e = np.array(parameter['tau_w'])
    model_test.tau_w_i = np.array(parameter['tau_w'])
    model_test.E_e = np.array(parameter['E_ex'])
    model_test.E_i = np.array(parameter['E_in'])
    model_test.Q_e = np.array(parameter['Q_e'])
    model_test.Q_i = np.array(parameter['Q_i'])
    model_test.tau_e = np.array(parameter['tau_syn_ex'])
    model_test.tau_i = np.array(parameter['tau_syn_in'])
    model_test.N_tot = np.array(parameter['N_tot'])
    model_test.p_connect = np.array(parameter['p_connect'])
    model_test.g = np.array(parameter['g'])
    model_test.T = np.array(parameter['t_ref'])
    model_test.external_input_in_in = np.array(0.0)
    model_test.external_input_in_ex = np.array(0.0)
    model_test.external_input_ex_in = np.array(0.0)
    model_test.external_input_ex_ex = np.array(0.0)
    model_test.K_ext_e=np.array(1)
    model_test.K_ext_i=np.array(0)
    if excitatory:
        def TF(fe,fi,p,f_ext_e=0.0,f_ext_i=0.0,w=0.0):
            model_test.P_e=p
            return model_test.TF_excitatory(fe,fi,f_ext_e,f_ext_i,w)
    else:
       def TF(fe,fi,p,f_ext_e=0.0,f_ext_i=0.0,w=0.0):
            model_test.P_i=p
            return model_test.TF_inhibitory(fe,fi,f_ext_e,f_ext_i,w)
    return TF

def effective_Vthre(Y, muV, sV, Tv):
    Vthre_eff = muV+np.sqrt(2)*sV*sp_spec.erfcinv(Y*2.*Tv) # effective threshold
    return Vthre_eff

def engin(parameters,excitatory,max_frequency=40.0,precision=0.5,frequency=None,rescale=None):
    name_file ='/home/kusch/Documents/project/co_simulation/co-simulation_mouse/test_nest/test_file/fitting/'
    for name,value in parameters.items():
        name_file += name+'_'+str(value)+'/'
    if frequency is None:
        frequency = np.arange(0.0,max_frequency,precision)
    frequencies = np.array(list(itertools.product(frequency,frequency)))

    if os.path.exists(name_file+'/P.npy'):
        return np.load(name_file+'/P.npy')
    elif os.path.exists(name_file+'/rate.npy'):
        rate = np.load(name_file+'/rate.npy')
    else:
        if os.path.exists(name_file+'/spike_detector.gdf'):
            print('analysis')
            data = load_spike(name_file)

        else:
            if not os.path.exists(name_file):
                os.makedirs(name_file)
            #initialisation of the parameter
            params = {      'g_L':parameters['g_L'],
                            'E_L':parameters['E_L'],
                            'V_reset':parameters['V_reset'],
                            'I_e':parameters['I_e'],
                            'C_m':parameters['C_m'],
                            'V_th':parameters['V_th'],
                            't_ref':parameters['t_ref'],
                            'tau_w':parameters['tau_w'],
                            'Delta_T':parameters['Delta_T'],
                            'b':parameters['b'],
                            'a':parameters['a'],
                            'V_peak':parameters['V_peak'],
                            'E_ex':parameters['E_ex'],
                            'E_in':parameters['E_in'],
                            'tau_syn_ex':parameters['tau_syn_ex'],
                            'tau_syn_in':parameters['tau_syn_in'],
                            'gsl_error_tol':1e-8
                            }
            Number_connexion_ex = parameters['N_tot']*parameters['p_connect']*(1-parameters['g'])
            Number_connexion_in = parameters['N_tot']*parameters['p_connect']*parameters['g']
            simtime=100000.0
            master_seed = 5
            local_num_threads = 8
            # simulation
            simulation = False
            error = 1.0e-6

            while error > 1.0e-20 and not simulation:
                params['gsl_error_tol'] = error
                # initialisation of nest
                nest.ResetKernel()
                nest.SetKernelStatus({
                    # Resolution of the simulation (in ms).
                    "resolution": 0.05,
                    # Print the time progress, this should only be used when the simulation
                    # is run on a local machine.
                    "print_time": True,
                    # If True, data will be overwritten,
                    # If False, a NESTError is raised if the files already exist.
                    "overwrite_files": True,
                    # Number of threads per MPI process.
                    'local_num_threads': local_num_threads,
                    # Path to save the output data
                    'data_path':  name_file,
                    # Masterseed for NEST and NumPy
                    'grng_seed': master_seed + local_num_threads,
                    # Seeds for the individual processes
                    'rng_seeds': range(master_seed + 1 + local_num_threads, master_seed + 1 + (2 * local_num_threads)),
                    })

                #create the network
                nest.SetDefaults('aeif_cond_exp', params)
                neurons = nest.Create('aeif_cond_exp', frequency.shape[0]**2)
                poisson_generator_ex = nest.Create('poisson_generator', frequency.shape[0])
                poisson_generator_in = nest.Create('poisson_generator', frequency.shape[0])
                nest.SetStatus(poisson_generator_ex,'rate',frequency*Number_connexion_ex)
                nest.SetStatus(poisson_generator_in,'rate',frequency*Number_connexion_in)
                nest.CopyModel("static_synapse", "excitatory",
                               {"weight": parameters['Q_e'], "delay": 1.0})
                nest.CopyModel("static_synapse", "inhibitory",
                               {"weight": -parameters['Q_i'], "delay": 1.0})
                for inh in range(len(frequency)):
                    for ex in range(len(frequency)):
                        nest.Connect(poisson_generator_ex[ex],neurons[ex+inh*len(frequency)],syn_spec="excitatory")
                        nest.Connect(poisson_generator_in[inh],neurons[ex+inh*len(frequency)],syn_spec="inhibitory")

                #create spike detector
                spikes_dec = nest.Create("spike_detector")
                nest.SetStatus(spikes_dec, [{"label": "spike",
                                          # "withtime": True,
                                          # "withgid": True,
                                          # "to_file": True,
                                            "record_to":"ascii",
                                           }])
                nest.Connect(neurons,spikes_dec)
                try :
                    nest.Simulate(simtime)
                    simulation = True
                    print('end')
                except nest.NESTError as exception:
                    template = "An exception of type {0} occurred. Arguments:\n{1!r}"
                    message = template.format(type(exception).__name__, exception.args)
                    print(message)
                    error = error/10.0
            print('analysis')

            # Concatenate the different spike files
            if subprocess.call([os.path.join(os.path.dirname(__file__),'script.sh'),name_file]) == 1:
                sys.stderr.write('ERROR bad concatenation of spikes file\n')
                exit(1)

            #Compute rate
            # data = load_event(nest.GetStatus(spikes_dec)[0]['events'] )
            data = load_spike(name_file)

        if data is None:
            print('compute rate')
            rate = np.zeros_like(frequencies)
            return [0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.]
        else:
            rate = compute_rate(data,0.0,100000.0,len(frequencies))
        np.save(name_file+'/rate.npy', rate)
        del data

    nb_freq = len (frequency)
    if rescale is not None:
        index = np.arange(0,nb_freq,1)
        index = index[np.where(np.logical_not(np.isin(index,rescale)))]
        rate = rate.reshape(nb_freq,nb_freq)
        rate = np.ravel(rate[index,:][:,index])
        frequencies =  np.array(list(itertools.product(frequency[index],frequency[index])))
        nb_freq -=len(rescale)

    # data = np.load('/home/kusch/Documents/project/Zerlaut/travail/Zerlaut/mean_field_for_multi_input_integration/transfer_functions/data/FS-cell_CONFIG1.npy')
    # data = np.load('/home/kusch/Documents/project/Zerlaut/travail/Zerlaut/mean_field_for_multi_input_integration/transfer_functions/data/RS-cell_CONFIG1.npy')
    # rate = np.ravel(data[0])*1e-3
    # frequencies = np.empty((len(rate),2))
    # frequencies[:,1] = np.ravel(data[2])
    # frequencies[:,0] = np.ravel(np.repeat([data[3]],data[2].shape[1],axis=0))

    #Compute mean of value for the model
    # rate+=1e-10
    muV, sV, Tv =model.get_fluct_regime_vars(frequencies[:,1]*1e-3,frequencies[:,0]*1e-3,0.00,0.0,0.0,parameters['Q_e'],parameters['tau_syn_ex'],parameters['E_ex'],parameters['Q_i'],parameters['tau_syn_in'],parameters['E_in'],
                          parameters['g_L'],parameters['C_m'],parameters['E_L'],parameters['N_tot'],parameters['p_connect'],parameters['g'],0.0,0.0)
    Tv+= parameters['g_L']/parameters['C_m']
    i_non_zeros = np.where((rate>=1/100000.0)  &(rate*Tv<1.0))
    # i_non_zeros=np.arange(0,len(rate),1)

    Vthre_eff = effective_Vthre(rate[i_non_zeros], muV[i_non_zeros], sV[i_non_zeros], Tv[i_non_zeros])*1e-3
    TvN = Tv[i_non_zeros]*parameters['g_L']/parameters['C_m']

    import matplotlib.pylab as plt
    # all = frequency[index]
    all =frequency
    index = i_non_zeros[0]; not_index = np.where(np.logical_not(np.isin(np.arange(0,nb_freq*nb_freq,1),i_non_zeros)))[0];
    fig = plt.figure();ax = fig.add_subplot(111, projection='3d');ax.scatter(muV, sV, Tv*parameters['g_L']/parameters['C_m'], marker='x',s=0.1);ax.set_xlabel('Vm');ax.set_ylabel('sV');ax.set_zlabel('Tv')
    fig = plt.figure();ax = fig.add_subplot(111, projection='3d');ax.scatter(muV[index], sV[index], Tv[index]*parameters['g_L']/parameters['C_m'], marker='x',s=0.1);ax.set_xlabel('Vm');ax.set_ylabel('sV');ax.set_zlabel('Tv');ax.scatter(muV[not_index], sV[not_index], Tv[not_index]*parameters['g_L']/parameters['C_m'], marker='o',s=0.1);
    plt.figure(); plt.plot(all,muV.reshape(nb_freq,nb_freq),'x',markersize=0.5);
    plt.figure(); plt.plot(all,sV.reshape(nb_freq,nb_freq),'x',markersize=0.5);
    plt.figure(); plt.plot(all,Tv.reshape(nb_freq,nb_freq)*parameters['g_L']/parameters['C_m'],'x',markersize=0.5)
    plt.figure(); plt.plot(index,TvN,'bx',markersize=0.5); plt.plot(not_index,Tv[not_index]*parameters['g_L']/parameters['C_m'],'rx',markersize=0.5)
    plt.figure(); plt.plot(index,muV[index],'bx',markersize=0.5); plt.plot(not_index,muV[not_index],'rx',markersize=0.5)
    plt.figure(); plt.plot(index,sV[index],'bx',markersize=0.5); plt.plot(not_index,sV[not_index],'rx',markersize=0.5)
    plt.figure();plt.plot(Vthre_eff,rate[index],'x',markersize=0.5);
    plt.figure();plt.plot(Vthre_eff,muV[index],'x',markersize=0.5);
    plt.figure();plt.plot(Vthre_eff,sV[index],'x',markersize=0.5);
    plt.figure();plt.plot(Vthre_eff,TvN,'x',markersize=0.5);
    plt.figure(); plt.plot(frequencies[:,1].reshape(nb_freq,nb_freq).transpose(),rate.reshape(nb_freq,nb_freq).transpose()*1e3)
    plt.show()

    TF = create_transfer_function(parameters,excitatory=excitatory)
    P = np.zeros(20)
    P[:5] = Vthre_eff.mean(), 1e-3, 1e-3, 1e-3, 1e-3

    def Res(p):
        pp=p
        vthre = model.threshold_func(muV[i_non_zeros], sV[i_non_zeros], TvN, *pp)
        return np.mean((Vthre_eff-vthre)**2)
        # return np.mean(np.abs(Vthre_eff-vthre)*1e3)
        # return np.mean(np.abs((Vthre_eff - vthre) / Vthre_eff))
    plsq = minimize(Res, P, method='SLSQP',options={'ftol': 1e-15, 'disp': True, 'maxiter':40000})
    # plsq = minimize(Res, P, method='SLSQP',tol=1e-10,\
    #                 options={'ftol': 1e-12, 'eps':1e-8,'disp': True, 'maxiter':100000})

    P = plsq.x
    def Res_1(p):
        return np.mean((rate[i_non_zeros] - TF(frequencies[i_non_zeros,1]*1e-3,frequencies[i_non_zeros,0]*1e-3,p)) ** 2)
        # return np.mean(np.abs((rate[i_non_zeros] - TF(frequencies[i_non_zeros,1]*1e-3,frequencies[i_non_zeros,0]*1e-3,p)))*1e3)
        # return np.mean(np.abs( (rate[i_non_zeros] - TF(frequencies[i_non_zeros, 1] * 1e-3, frequencies[i_non_zeros, 0] * 1e-3, p)) / rate[i_non_zeros]))
    # plsq = minimize(Res_1, P, method='nelder-mead',options={'xtol': 1e-5, 'disp': True, 'maxiter': 50000})
    # plsq = minimize(Res_1, P, method='nelder-mead',tol=1e-12, \
    #                 options={'xtol': 1., 'disp': True, 'maxiter': 100000, 'maxfev':100000})
    # p = plsq.x
    # index = np.argsort(np.abs(((rate - TF(frequencies[:, 1] * 1e-3, frequencies[:, 0] * 1e-3, p)) * 1e3)))[-10:]
    # print('first part ')
    # print(P)
    # print("frequency", frequencies[index])
    # print("expected : ", rate[index] * 1e3)
    # print("got : ", TF(frequencies[:, 1] * 1e-3, frequencies[:, 0] * 1e-3, p)[index] * 1e3)
    # print("error : ", np.abs(((rate - TF(frequencies[:, 1] * 1e-3, frequencies[:, 0] * 1e-3, p)) * 1e3))[index])
    # print("max error ", np.abs(((rate - TF(frequencies[:, 1] * 1e-3, frequencies[:, 0] * 1e-3, p)) * 1e3))[index[-1]])

    P = plsq.x
    def Res_2(p):
        return np.mean((rate - TF(frequencies[:,1]*1e-3,frequencies[:,0]*1e-3,p)) ** 2)
        # return np.mean(np.abs((rate - TF(frequencies[:,1]*1e-3,frequencies[:,0]*1e-3,p)))*1e3)
        # return np.mean(np.abs((rate - TF(frequencies[:,1]*1e-3,frequencies[:,0]*1e-3,p))/rate))
    plsq = minimize(Res_2, P, method='nelder-mead',options={'xtol': 1e-5, 'disp': True, 'maxiter': 50000})
    # plsq = minimize(Res_2, P, method='nelder-mead',tol=1e-12, \
    #                 options={'xtol': 1.0, 'disp': True, 'maxiter': 100000, 'maxfev':100000})

    # network = np.empty((651,9))
    # range_rate = np.arange(0.0,105.0,5.0)
    # range_b =  np.arange(0.,3.1,0.1)
    # range_test = np.array(list(itertools.product(range_b,range_rate)))
    # for i in range(651):
    #     network[i,:7]=np.load("/home/kusch/Documents/project/co_simulation/co-simulation_mouse/test_nest/test_file/fitting/network_"+str(i)+".npy")
    # network[:,7:]=range_test
    # network = network[np.where(network[:,6]==0)[0],:]
    # index_network = 2 if excitatory else 0
    # P = plsq.x
    # def Res_3(p):
        # return np.mean((( network[:,index_network]*1e-3- TF(network[:,2]*1e-3,network[:,0]*1e-3,p,network[:,8]*1e-3,0.0,network[:,4]))*1e3) ** 2)
        # return np.mean(np.abs( network[:,index_network]- TF(network[:,2]*1e-3,network[:,0]*1e-3,p,network[:,8]*1e-3,0.0,network[:,4])*1e3))
        # return np.mean(np.abs( (network[:,index_network] - TF(network[:,2]*1e-3,network[:,0]*1e-3,p,network[:,8]*1e-3,0.0,network[:,4])*1e3))/network[:,index_network])
    # plsq = minimize(Res_3, P, method='nelder-mead',tol=1e-12, \
    #             options={'xtol': 1.0, 'disp': True, 'maxiter': 100000, 'maxfev':100000})

    p = plsq.x
    index = np.argsort(np.abs(((rate - TF(frequencies[:, 1] * 1e-3, frequencies[:, 0] * 1e-3, p)) * 1e3)))[-10:]
    # print("frequency", frequencies[index])
    # print("expected : ", rate[index] * 1e3)
    # print("got : ", TF(frequencies[:, 1] * 1e-3, frequencies[:, 0] * 1e-3, p)[index] * 1e3)
    # print("error : ", np.abs(((rate - TF(frequencies[:, 1] * 1e-3, frequencies[:, 0] * 1e-3, p)) * 1e3))[index])
    # print("max error ", np.abs(((rate - TF(frequencies[:, 1] * 1e-3, frequencies[:, 0] * 1e-3, p)) * 1e3))[index[-1]])
    # np.save(name_file+'/P.npy', plsq.x)

    TF(frequencies[:,1]*1e-3,frequencies[:,0]*1e-3,p)
    plt.figure(); plt.plot(frequencies[:,1].reshape(nb_freq,nb_freq).transpose(),TF(frequencies[:,1]*1e-3,frequencies[:,0]*1e-3,p).reshape(nb_freq,nb_freq).transpose()*1e3)
    plt.figure(); plt.plot(frequencies[:,1].reshape(nb_freq,nb_freq).transpose(),rate.reshape(nb_freq,nb_freq).transpose()*1e3)
    plt.show()
    return plsq.x

# frequency = np.array([0.0,0.1,0.2,0.3,0.4,0.5,0.6,0.61,0.62,0.63,0.64,0.65,0.66,0.67,0.68,0.69,0.7,0.71,0.72,0.73,0.74,0.75,0.76,0.77,0.78,0.79,0.8,0.81,0.82,0.83,0.84,0.85,0.86,0.87,0.88,0.89,0.9,0.92,0.95,0.97,1.0,1.1,1.2,1.3,1.4,1.5,1.6,1.7,1.8,1.9,2.0,2.2,2.5,2.7,3.0,3.2,3.5,3.7,4.0,4.2,4.5,4.7,5.0,5.2,5.7,6.0,6.5,7.0,7.5,8.0,8.5,9.0,
#              10.0,11.0,12.0,13.0,14.0,15.0,16.0,17.0,18.0,19.0,20.0,21.0,22.0,23.0,24.0,25.0,26.0,27.0,28.0,29.0,30.0,31.0,32.0,33.0,34.0,35.0,36.0,37.0,38.0,39.0,40.0,
#                       50.0,60.0])
#                       41.0,42.0,43.0,44.0,45.0,46.0,47.0,48.0,49.0,50.0,51.0,52.0,53.0,54.0,55.0,56.0,57.0,58.0,59.0,60.0])
                      # 41.0, 42.0, 43.0, 44.0, 45.0, 46.0, 47.0, 48.0, 49.0, 50.0, 51.0, 52.0, 53.0, 54.0, 55.0, 56.0,
                      # 57.0, 58.0, 59.0, 60.0,61.0,62.0,63.0,64.0,65.0,66.0,67.0,68.0,69.0,70.0,71.0,72.0,73.0,74.0,75.0,76.0,77.0,78.0,79.0,80.0,
                      # 85.0,90.0,100.0,110.0,120.0,130.0,140.0,150.0,160.0,170.0,180.0,200.0])

print("EXCITATORY")
print("'P_e':",np.array2string(engin(parameters=excitatory,excitatory=True,max_frequency=40.0,precision=0.5), separator=', '),",",sep='')
# print("'P_e':",np.array2string(engin(parameters=excitatory,excitatory=True,max_frequency=40.0,precision=0.5,rescale=[32,33,56,57,59,60,62,63]), separator=', '),",",sep='')
print("INHIBITORY")
# print("'P_i':",np.array2string(engin(parameters=inhibitory,excitatory=False,max_frequency=40.0,precision=0.5), separator=', '),",",sep='')
# print("'P_i':",np.array2string(engin(parameters=inhibitory,excitatory=False,max_frequency=40.0,precision=0.5,rescale=[5,6]), separator=', '),",",sep='')
# print(engin(parameters=default,frequency=frequency))
# print(engin(parameters=default,max_frequency=60.0,precision=0.5))
# print(engin(parameters=default,max_frequency=100.0,precision=0.5))
# np.set_printoptions(linewidth=500)
# int_first = 200
print("EXCITATORY")
# print("'P_e':",np.array2string(engin(parameters=excitatory,excitatory=True,max_frequency=100.0,precision=0.5,rescale=[10,12,123,124,174]), separator=', '),",",sep='')
# rescale = np.arange(int_first,200)
# print("'P_e':",np.array2string(engin(parameters=excitatory,excitatory=True,max_frequency=100.0,precision=0.5,rescale=rescale), separator=', '),",",sep='')
# print("'P_e':",np.array2string(engin(parameters=excitatory,max_frequency=40.0,precision=0.5), separator=', '),",",sep='')
print("INHIBITORY")
# print("'P_i':",np.array2string(engin(parameters=inhibitory,excitatory=False,max_frequency=100.0,precision=0.5,rescale=[3,4,44,45,59,60,81,82,105,106,109,110,111,112,113,114,194,195]), separator=', '),",",sep='')
# print("'P_i':",np.array2string(engin(parameters=inhibitory,excitatory=False,max_frequency=100.0,precision=0.5,rescale=[44,45,59,60,81,82,105,106,109,110,111,112,194,195]), separator=', '),",",sep='')
# rescale = np.arange(int_first,200)
# print("'P_i':",np.array2string(engin(parameters=inhibitory,excitatory=False,max_frequency=100.0,precision=0.5,rescale=rescale), separator=', '),",",sep='')
# print("'P_i':",np.array2string(engin(parameters=inhibitory,max_frequency=40.0,precision=0.5), separator=', '),",",sep='')

def model_end (g_L,E_L,V_reset,I_e,C_m,V_th,t_ref,tau_w,Delta_T,b,a,E_ex,E_in,tau_syn_ex,tau_syn_in,Q_e,Q_i,N_tot,p_connect,g):
    parameters= {   'g_L':g_L,
                    'E_L':E_L,
                    'V_reset':V_reset,
                    'I_e':I_e,
                    'C_m':C_m,
                    'V_th':V_th,
                    't_ref':t_ref,
                    'tau_w':tau_w,
                    'Delta_T':Delta_T,
                    'b':b,
                    'a':a,
                    'V_peak':-10.0,
                    'E_ex':E_ex,
                    'E_in':E_in,
                    'tau_syn_ex':tau_syn_ex,
                    'tau_syn_in':tau_syn_in,
                    'Q_e': Q_e,
                    'Q_i' : Q_i,
                    'N_tot' : N_tot,
                    'p_connect' : p_connect,
                    'g': g
                    }

    return None, engin(parameters=parameters,max_frequency=40.0)


""" scan info for specific reaction classes
"""

import math
import numpy
import more_itertools as mit
from phydat import phycon, bnd
from automol.graph import ts
from automol.par import ReactionClass
import automol.zmat
from automol.util import dict_
# from automol.util import numpy_to_float
from automol.reac._util import hydrogen_migration_atom_keys
from automol.reac._util import ring_forming_scission_chain
from automol.reac._util import insertion_forming_bond_keys
from automol.reac._util import elimination_breaking_bond_keys

# Wrapper function to obtain all of the scan data for a reaction
def build_scan_info(zrxn, zma, var=False):
    """ Build all of the scan information
    """

    # Obtain the reactions scan and constraint coordinates
    scan_names = scan_coordinate(zrxn, zma)
    const_names = constraint_coordinates(zrxn, zma)

    constraint_dct = automol.zmat.constraint_dct(zma, const_names)

    # Build the grid
    grids = scan_grid(zrxn, zma, var=var)

    # Set the update guess
    update_guess = scan_update_guess(zrxn, var=var)

    return scan_names, constraint_dct, grids, update_guess

# SCAN AND CONSTRAINT COORDINATES #
# Unimolecular reactions
# 1. Hydrogen migrations
def hydrogen_migration_scan_coordinate(rxn, zma):
    """ Obtain the scan coordinate for a hydrogen migration.

    :param rxn: a Reaction object
    :returns: the name of the scan coordinate in the z-matrix
    :rtype: str
    """
    frm_bnd_key, = ts.forming_bond_keys(rxn.forward_ts_graph)
    scan_name = automol.zmat.distance_coordinate_name(zma, *frm_bnd_key)
    return (scan_name,)

def hydrogen_migration_constraint_coordinates(rxn, zma):
    """ Obtain the constraint coordinates for a hydrogen migration

    :param rxn: a Reaction object
    :returns: the names of the constraint coordinates in the z-matrix
    :rtype: tuple[str]
    """
    att_key, _, _, ngb_key = hydrogen_migration_atom_keys(rxn)
    dist_name = automol.zmat.distance_coordinate_name(zma, att_key, ngb_key)
    return (dist_name,)

def hydrogen_migration_grid(zrxn, zma, npoints=(18,)):
    """ Build forward 1D grid  for addition reaction
    """

    # Obtain the reactions scan and constraint coordinates
    scan_name, = hydrogen_migration_scan_coordinate(zrxn, zma)

    # Build the scan grid
    npoints1 = npoints[0]
    interval = 0.3 * phycon.ANG2BOHR

    frm_bnd_len = _ts_bnd_len(zma, scan_name)
    rmin1 = 2.0 * phycon.ANG2BOHR
    rmin2 = frm_bnd_len + (0.05 * phycon.ANG2BOHR)
    rmax = frm_bnd_len

    if rmax > rmin1:
        npoints = math.ceil((rmax-rmin1)/interval)
        if npoints < 1:
            grid1 = []
        else:
            grid1 = numpy.linspace(rmax, rmin1, npoints)
    else:
        grid1 = []

    grid2 = numpy.linspace(rmin1, rmin2, npoints1)
    grid = numpy.concatenate((grid1, grid2), axis=None)

    return (grid,)

# 2. Beta scissions
def beta_scission_scan_coordinate(rxn, zma):
    """ Obtain the scan coordinate for a beta scission

    :param rxn: a Reaction object
    :returns: the name of the scan coordinate in the z-matrix
    :rtype: str
    """
    brk_bnd_key, = ts.breaking_bond_keys(rxn.forward_ts_graph)
    scan_name = automol.zmat.distance_coordinate_name(zma, *brk_bnd_key)
    return (scan_name,)

def beta_scission_grid(zrxn, zma, npoints=(14,)):
    """ Build forward 1D grid for a beta scission reaction
    """

    # Obtain the reactions scan and constraint coordinates
    scan_name, = beta_scission_scan_coordinate(zrxn, zma)

    # Build the scan grid
    npoints1 = npoints[0]

    frm_bnd_len = _ts_bnd_len(zma, scan_name)
    if frm_bnd_len is not None:
        rmin = frm_bnd_len + (0.1 * phycon.ANG2BOHR)
        rmax = frm_bnd_len + (0.8 * phycon.ANG2BOHR)
    else:
        rmin = 1.4 * phycon.ANG2BOHR
        rmax = 2.0 * phycon.ANG2BOHR
    grid = numpy.linspace(rmin, rmax, npoints1)

    return (grid,)

# 3. Ring-forming scissions
def ring_forming_scission_scan_coordinate(rxn, zma):
    """ Obtain the scan coordinate for a ring-forming scission

    :param rxn: a Reaction object
    :returns: the name of the scan coordinate in the z-matrix
    :rtype: str
    """
    brk_bnd_key, = ts.breaking_bond_keys(rxn.forward_ts_graph)
    scan_name = automol.zmat.distance_coordinate_name(zma, *brk_bnd_key)
    return (scan_name,)

def ring_forming_scission_constraint_coordinates(rxn, zma):
    """ Obtain the constraint coordinates for a ring-forming scission

    :param rxn: a Reaction object
    :returns: the names of the constraint coordinates in the z-matrix
    :rtype: str
    """
    chain_keys = ring_forming_scission_chain(rxn)
    ang_keys_lst = sorted(mit.windowed(chain_keys[1:], 3))
    dih_keys_lst = sorted(mit.windowed(chain_keys, 4))
    ang_names = [automol.zmat.central_angle_coordinate_name(zma, *ks)
                 for ks in ang_keys_lst]
    dih_names = [automol.zmat.dihedral_angle_coordinate_name(zma, *ks)
                 for ks in dih_keys_lst]
    const_names = tuple(ang_names + dih_names)
    return const_names

def ring_forming_scission_grid(zrxn, zma, npoints=(7,)):
    """ Build forward WD grid for a ring forming scission reaction
        # the following allows for a 2-d grid search in the initial ts_search
        # for now try 1-d grid and see if it is effective
    """

    # Obtain the scan coordinate
    scan_name, = ring_forming_scission_scan_coordinate(zrxn, zma)

    # Build the grid
    npoints1 = npoints[0]

    brk_bnd_len = _ts_bnd_len(zma, scan_name)
    if brk_bnd_len is not None:
        r1min = brk_bnd_len + (0.1 * phycon.ANG2BOHR)
        r1max = brk_bnd_len + (0.7 * phycon.ANG2BOHR)
    else:
        r1min = (1.54 + 0.1) * phycon.ANG2BOHR
        r1max = (1.54 + 0.7) * phycon.ANG2BOHR

    grid1 = numpy.linspace(r1min, r1max, npoints1)
    grid = tuple(val.item() for val in grid1)

    return (grid,)

# 4. Eliminations
def elimination_scan_coordinate(rxn, zma):
    """ Obtain the scan coordinate for an elimination

    :param rxn: a Reaction object
    :returns: the name of the scan coordinate in the z-matrix
    :rtype: str
    """
    frm_bnd_key, = ts.forming_bond_keys(rxn.forward_ts_graph)

    brk_bnd_key1, _ = elimination_breaking_bond_keys(rxn)

    frm_name = automol.zmat.distance_coordinate_name(zma, *frm_bnd_key)
    brk_name = automol.zmat.distance_coordinate_name(zma, *brk_bnd_key1)

    return (frm_name, brk_name)

def elimination_grid(zrxn, zma, npoints=(7, 5)):
    """ Build forward 2D grid for elimination reaction
    """

    # Obtain the scan coordinate
    frm_name, brk_name = elimination_scan_coordinate(zrxn, zma)

    # Build the grid
    npoints1, npoints2 = npoints

    frm_bnd_len = _ts_bnd_len(zma, frm_name)
    brk_bnd_len = _ts_bnd_len(zma, brk_name)
    if frm_bnd_len is not None:
        r1min = frm_bnd_len + (0.1 * phycon.ANG2BOHR)
        r1max = frm_bnd_len + (0.6 * phycon.ANG2BOHR)
    else:
        r1min = (0.85 + 0.1) * phycon.ANG2BOHR
        r1max = (0.85 + 0.8) * phycon.ANG2BOHR
    if brk_bnd_len is not None:
        r2min = brk_bnd_len + (0.3 * phycon.ANG2BOHR)
        r2max = brk_bnd_len + (1.2 * phycon.ANG2BOHR)
    else:
        r2min = (1.50 + 0.3) * phycon.ANG2BOHR
        r2max = (1.50 + 1.2) * phycon.ANG2BOHR

    grid1 = numpy.linspace(r1min, r1max, npoints1) * phycon.ANG2BOHR
    grid2 = numpy.linspace(r2min, r2max, npoints2) * phycon.ANG2BOHR

    return (grid1, grid2)

# Bimolecular reactions
# 1. Hydrogen abstractions
def hydrogen_abstraction_scan_coordinate(rxn, zma):
    """ Obtain the scan coordinate for a hydrogen abstraction

    :param rxn: a Reaction object
    :returns: the name of the scan coordinate in the z-matrix
    :rtype: str
    """
    frm_bnd_key, = ts.forming_bond_keys(rxn.forward_ts_graph)
    scan_name = automol.zmat.distance_coordinate_name(zma, *frm_bnd_key)
    return (scan_name,)

def hydrogen_abstraction_grid(zrxn, zma, npoints=(8,)):
    """ Build forward 1D grid for hydrogen abstraction reaction
    """

    # Obtain the scan coordinate
    scan_name, = hydrogen_abstraction_scan_coordinate(zrxn, zma)

    # Build the grid
    npoints1 = npoints[0]

    frm_bnd_len = _ts_bnd_len(zma, scan_name)
    if frm_bnd_len is not None:
        rmin = frm_bnd_len + (0.1 * phycon.ANG2BOHR)
        rmax = frm_bnd_len + (1.0 * phycon.ANG2BOHR)
    else:
        rmin = 0.7 * phycon.ANG2BOHR
        rmax = 2.2 * phycon.ANG2BOHR
    grid = numpy.linspace(rmin, rmax, npoints1)

    return (grid,)

def radrad_hydrogen_abstraction_grid(zrxn, zma, npoints=(8, 4)):
    """ Build forward 1D grid for elimination reaction
    """

    # Obtain the scan coordinate
    scan_name, = hydrogen_abstraction_scan_coordinate(zrxn, zma)

    # Build the grid
    npoints1, npoints2 = npoints

    # Get the first grid from close to the mid
    frm_bnd_len = _ts_bnd_len(zma, scan_name)
    if frm_bnd_len is not None:
        rmin = frm_bnd_len + (0.1 * phycon.ANG2BOHR)
        rmax = frm_bnd_len + (1.0 * phycon.ANG2BOHR)
    else:
        rmin = 0.7 * phycon.ANG2BOHR
        rmax = 2.2 * phycon.ANG2BOHR

    grid1 = numpy.linspace(rmin, rmax, npoints1)
    grid1 = numpy.flip(grid1)

    # Get the outer grid from mid to long-distance
    rend2 = 4.0 * phycon.ANG2BOHR
    grid2 = numpy.linspace(rmax, rend2, npoints2+1)
    # grid2 = numpy.delete(grid2, 0)

    return (grid1, grid2)

# 2. Additions
def addition_scan_coordinate(rxn, zma):
    """ Obtain the scan coordinate for an addition

    :param rxn: a Reaction object
    :returns: the name of the scan coordinate in the z-matrix
    :rtype: str
    """
    frm_bnd_key, = ts.forming_bond_keys(rxn.forward_ts_graph)
    scan_name = automol.zmat.distance_coordinate_name(zma, *frm_bnd_key)
    return (scan_name,)

def addition_grid(zrxn, zma, npoints=(14,)):
    """ Build forward 1D grid for addition reaction
    """

    # Obtain the scan coordinate
    scan_name, = addition_scan_coordinate(zrxn, zma)

    # Build the grid
    npoints1 = npoints[0]

    frm_bnd_len = _ts_bnd_len(zma, scan_name)
    if frm_bnd_len is not None:
        rmin = frm_bnd_len + (0.1 * phycon.ANG2BOHR)
        rmax = frm_bnd_len + (1.2 * phycon.ANG2BOHR)
    else:
        rmin = 1.6 * phycon.ANG2BOHR
        rmax = 2.8 * phycon.ANG2BOHR

    grid = _geometric_progression(
        rmin, rmax, npoints1, gfact=1.1, rstp=0.05)

    return (grid,)

def radrad_addition_grid(zrxn, zma, npoints=(8, 4)):
    """ Build forward 1D grid for a beta scission reaction
    """

    # Obtain the scan coordinate
    scan_name, = addition_scan_coordinate(zrxn, zma)

    # Build the grid
    npoints1, npoints2 = npoints

    # Get the first grid from close to the mid
    frm_bnd_len = _ts_bnd_len(zma, scan_name)
    if frm_bnd_len is not None:
        rmin = frm_bnd_len + (0.2 * phycon.ANG2BOHR)
        rmax = frm_bnd_len + (1.2 * phycon.ANG2BOHR)
    else:
        rmin = 1.5 * phycon.ANG2BOHR
        rmax = 2.8 * phycon.ANG2BOHR

    grid1 = numpy.linspace(rmin, rmax, npoints1)
    grid1 = numpy.flip(grid1)
    # grid1 = _geometric_progression(
    #     rmin, rmax, npoints1, gfact=1.1, rstp=0.05)

    # Get the outer grid from mid to long-distance
    # Add extra point since initial will be dropped
    rend2 = 4.00 * phycon.ANG2BOHR
    grid2 = numpy.linspace(rmax, rend2, npoints2+1)
    # grid2 = numpy.delete(grid2, 0)

    return (grid1, grid2)

# 3. Insertions
def insertion_scan_coordinate(rxn, zma):
    """ Obtain the scan coordinate for an insertion

    :param rxn: a Reaction object
    :returns: the name of the scan coordinate in the z-matrix
    :rtype: str
    """
    frm_bnd_key, _ = insertion_forming_bond_keys(rxn)
    scan_name = automol.zmat.distance_coordinate_name(zma, *frm_bnd_key)
    return (scan_name,)

def insertion_grid(zrxn, zma, npoints=(16,)):
    """ Build forward 1D grid for insertion reaction
    """

    # Obtain the scan coordinate
    scan_name, = insertion_scan_coordinate(zrxn, zma)

    # Build the grid
    npoints1 = npoints[0]

    frm_bnd_len = _ts_bnd_len(zma, scan_name)
    if frm_bnd_len is not None:
        rmin = frm_bnd_len
        rmax = frm_bnd_len + (1.4 * phycon.ANG2BOHR)
    else:
        rmin = 1.4 * phycon.ANG2BOHR
        rmax = 2.4 * phycon.ANG2BOHR

    grid = numpy.linspace(rmin, rmax, npoints1)

    return (grid,)

# 4. Substitution
def substitution_scan_coordinate(rxn, zma):
    """ Obtain the scan coordinate for a substitution

    :param rxn: a Reaction object
    :returns: the name of the scan coordinate in the z-matrix
    :rtype: str
    """
    frm_bnd_key, = ts.forming_bond_keys(rxn.forward_ts_graph)
    scan_name = automol.zmat.distance_coordinate_name(zma, *frm_bnd_key)
    return (scan_name,)

def substitution_grid(zrxn, zma, npoints=(14,)):
    """ Build forward 1D grid for substitution reaction
    """

    # Obtain the scan coordinate
    scan_name, = substitution_scan_coordinate(zrxn, zma)

    # Build the grid
    npoints1 = npoints[0]

    frm_bnd_len = _ts_bnd_len(zma, scan_name)
    if frm_bnd_len is not None:
        rmin = frm_bnd_len
        rmax = frm_bnd_len + (1.4 * phycon.ANG2BOHR)
    else:
        rmin = 0.7 * phycon.ANG2BOHR
        rmax = 2.4 * phycon.ANG2BOHR

    grid = numpy.linspace(rmin, rmax, npoints1)

    return (grid,)

# Aux function to return empty tuple when info not require for class
def _return_empty_tuple(*_):
    """ Return empty tuple
    """
    return ()

# Wrapper functions to handle rxn obj and zma for any reaction class
SCAN_COORD_DCT = {
    # unimolecular
    ReactionClass.Typ.HYDROGEN_MIGRATION:
    hydrogen_migration_scan_coordinate,
    ReactionClass.Typ.BETA_SCISSION: beta_scission_scan_coordinate,
    ReactionClass.Typ.RING_FORM_SCISSION:
    ring_forming_scission_scan_coordinate,
    ReactionClass.Typ.ELIMINATION: elimination_scan_coordinate,
    # bimolecular
    ReactionClass.Typ.HYDROGEN_ABSTRACTION:
    hydrogen_abstraction_scan_coordinate,
    ReactionClass.Typ.ADDITION: addition_scan_coordinate,
    ReactionClass.Typ.INSERTION: insertion_scan_coordinate,
    ReactionClass.Typ.SUBSTITUTION: substitution_scan_coordinate,
}

def scan_coordinate(rxn, zma):
    """ Obtain the scan coordinates

    :param rxn: a hydrogen migration Reaction object
    """
    return SCAN_COORD_DCT[rxn.class_](rxn, zma)

CONSTRAINT_COORD_DCT = {
    # unimolecular
    ReactionClass.Typ.HYDROGEN_MIGRATION:
    hydrogen_migration_constraint_coordinates,
    ReactionClass.Typ.BETA_SCISSION: _return_empty_tuple,
    ReactionClass.Typ.RING_FORM_SCISSION:
    ring_forming_scission_constraint_coordinates,
    ReactionClass.Typ.ELIMINATION: _return_empty_tuple,
    # bimolecular
    ReactionClass.Typ.HYDROGEN_ABSTRACTION:
    _return_empty_tuple,
    ReactionClass.Typ.ADDITION: _return_empty_tuple,
    ReactionClass.Typ.INSERTION: _return_empty_tuple,
    ReactionClass.Typ.SUBSTITUTION: _return_empty_tuple,
}

def constraint_coordinates(rxn, zma):
    """ Obtain the constraint coordinates

    :param rxn: a hydrogen migration Reaction object
    """
    return CONSTRAINT_COORD_DCT[rxn.class_](rxn, zma)

TIGHT_TS_GRID_DCT = {
    ReactionClass.Typ.BETA_SCISSION: beta_scission_grid,
    ReactionClass.Typ.ADDITION: addition_grid,
    ReactionClass.Typ.HYDROGEN_MIGRATION: hydrogen_migration_grid,
    ReactionClass.Typ.ELIMINATION: elimination_grid,
    ReactionClass.Typ.RING_FORM_SCISSION: ring_forming_scission_grid,
    ReactionClass.Typ.HYDROGEN_ABSTRACTION: hydrogen_abstraction_grid,
    ReactionClass.Typ.SUBSTITUTION: substitution_grid,
    ReactionClass.Typ.INSERTION: insertion_grid
}
VAR_TS_GRID_DCT = {
    ReactionClass.Typ.ADDITION: radrad_addition_grid,
    ReactionClass.Typ.HYDROGEN_ABSTRACTION: radrad_hydrogen_abstraction_grid
}

def scan_grid(zrxn, zma, var=False):
    """ Set the grid for a transition state search
    """

    if not var:
        grid = TIGHT_TS_GRID_DCT[zrxn.class_](zrxn, zma)
    else:
        grid = VAR_TS_GRID_DCT[zrxn.class_](zrxn, zma)

    return grid

# UPDATE GUESS DICTIONARY #
TIGHT_TS_UPDATE_GUESS_DCT = {
    ReactionClass.Typ.BETA_SCISSION: False,
    ReactionClass.Typ.ADDITION: False,
    ReactionClass.Typ.HYDROGEN_MIGRATION: True,
    ReactionClass.Typ.ELIMINATION: False,
    ReactionClass.Typ.RING_FORM_SCISSION: False,
    ReactionClass.Typ.HYDROGEN_ABSTRACTION: False,
    ReactionClass.Typ.SUBSTITUTION: False,
    ReactionClass.Typ.INSERTION: False
}
VAR_TS_UPDATE_GUESS_DCT = {
    ReactionClass.Typ.ADDITION: True,
    ReactionClass.Typ.HYDROGEN_ABSTRACTION: True
}

def scan_update_guess(zrxn, var=False):
    """ Set boolean to control whether the initial guess structure for
        optimization updates along scan, i.e., uses optimized geometry
        from previous grid point
    """

    if not var:
        _update = TIGHT_TS_UPDATE_GUESS_DCT[zrxn.class_]
    else:
        _update = VAR_TS_UPDATE_GUESS_DCT[zrxn.class_]

    return _update

# Helper functions
def _ts_bnd_len(zma, scan_coord):
    """ Obtain the current value of the bond defined by the scam coordinate
    """

    symbs = automol.zmat.symbols(zma)
    dist_coo, = automol.zmat.coordinates(zma)[scan_coord]
    ts_bnd_symbs = tuple(sorted(map(symbs.__getitem__, dist_coo)))
    ts_bnd_len = dict_.values_by_unordered_tuple(bnd.LEN_DCT, ts_bnd_symbs)

    return ts_bnd_len

def _geometric_progression(rmin, rmax, npoints, gfact=1.1, rstp=0.05):
    """ Build a grid using a geometric progresion
    """
    grid = [rmin]
    rgrid = rmin
    for _ in range(npoints):
        rgrid += rstp
        if rgrid == rmax:
            break
        grid.append(rgrid)
        rstp = rstp * gfact
    grid = numpy.array(grid)

    return grid