"""
File: pylinex/nonlinear/loglikelihood/Loglikelihood.py
Author: Keith Tauscher
Date: 25 Feb 2018

Description: File containing a base class representing a likelihood that can be
             evaluated using a data vector and a Model object (and possibly
             other things, depending on the subclass).
"""
import numpy as np
import numpy.linalg as la
from distpy import TransformList, GaussianDistribution, WindowedDistribution,\
    DistributionSet
from ..util import Savable, Loadable

cannot_instantiate_loglikelihood_error = NotImplementedError("The " +\
    "Loglikelihood class cannot be instantiated directly!")

class Loglikelihood(Savable, Loadable):
    """
    Abstract class representing a likelihood which is Gaussian in the data.
    """
    def __init__(self, *args, **kwargs):
        """
        Initializer throws error because Loglikelihood is supposed to be an
        abstract class that is not directly instantiated.
        """
        raise cannot_instantiate_loglikelihood_error

    @property
    def parameters(self):
        """
        Property storing the names of the parameters of the model defined by
        this likelihood.
        """
        raise cannot_instantiate_loglikelihood_error

    @property
    def num_parameters(self):
        """
        Property storing the number of parameters needed by the Model at the
        heart of this Loglikelihood.
        """
        if not hasattr(self, '_num_parameters'):
            self._num_parameters = len(self.parameters)
        return self._num_parameters

    def fill_hdf5_group(self, group, *args, **kwargs):
        """
        Fills the given hdf5 group with information about this Loglikelihood.

        group: the group to fill with information about this Loglikelihood
        """
        raise cannot_instantiate_loglikelihood_error

    @staticmethod
    def load_from_hdf5_group(group):
        """
        Loads a Loglikelihood object from an hdf5 file group in which it was
        previously saved.

        group: the hdf5 file group from which to load a Loglikelihood object

        returns: the Loglikelihood object loaded from the given hdf5 file group
        """
        raise cannot_instantiate_loglikelihood_error

    def check_parameter_dimension(self, pars):
        """
        Checks to ensure that the given array is 1D and has one element for
        each parameter. The only thing this function does is throw an error if
        the array is the wrong shape.

        pars: array to check
        """
        if pars.shape != (self.num_parameters,):
            raise ValueError("The array of parameters given to this " +\
                "Loglikelihood object was not of the correct size.")

    def __call__(self, pars, return_negative=False):
        """
        Gets the value of this Loglikelihood at the given parameters.
        """
        raise cannot_instantiate_loglikelihood_error

    @property
    def gradient_computable(self):
        """
        Property storing whether the gradient of this Loglikelihood can be
        computed.
        """
        raise cannot_instantiate_loglikelihood_error

    def gradient(self, pars, return_negative=False):
        """
        Computes the gradient of this Loglikelihood for minimization purposes.

        pars: value of the parameters at which to evaluate the gradient
        return_negative: if true, the negative of the gradient of the
                         loglikelihood is returned (this is useful for times
                         when the loglikelihood must be maximized since scipy
                         optimization functions only deal with minimization

        returns: 1D numpy.ndarray of length num_parameters containing gradient
                 of loglikelihood value
        """
        if self.gradient_computable:
            return self.auto_gradient(pars, return_negative=return_negative)
        else:
            raise NotImplementedError("gradient is not computable in an " +\
                "exact form. Use auto_gradient instead to do a numerical " +\
                "approximation.")

    def auto_gradient(self, pars, return_negative=False, differences=1e-6,\
        transform_list=None):
        """
        Computes the gradient of this Loglikelihood for minimization purposes.

        pars: value of the parameters at which to evaluate the gradient
        return_negative: if true, the negative of the gradient of the
                         loglikelihood is returned (this is useful for times
                         when the loglikelihood must be maximized since scipy
                         optimization functions only deal with minimization
        differences: either single number or 1D array of numbers to use as the
                     numerical difference in parameter. Default: 10^(-6)
        transform_list: TransformList object (or something which can be cast to
                        one) defining the transforms to apply to the parameters
                        before computing the gradient. Default: None, parameter
                        space is not transformed

        returns: 1D numpy.ndarray of length num_parameters containing gradient
                 of loglikelihood value
        """
        raise cannot_instantiate_loglikelihood_error

    @property
    def hessian_computable(self):
        """
        Property storing whether the hessian of this Loglikelihood can be
        computed. The hessian of this Loglikelihood is computable as long as
        the model's gradient and hessian are computable.
        """
        raise cannot_instantiate_loglikelihood_error

    def hessian(self, pars, return_negative=False):
        """
        Computes the hessian of this Loglikelihood for minimization purposes.

        pars: value of the parameters at which to evaluate the hessian
        return_negative: if true, the negative of the hessian of the
                         loglikelihood is returned (this is useful for times
                         when the loglikelihood must be maximized since scipy
                         optimization functions only deal with minimization

        returns: square 2D numpy.ndarray of side length num_parameters
                 containing hessian of loglikelihood value
        """
        if self.hessian_computable:
            return self.auto_hessian(pars, return_negative=return_negative)
        else:
            raise NotImplementedError("hessian is not computable in an " +\
                "exact form. Use auto_hessian instead to do a numerical " +\
                "approximation.")

    def auto_hessian(self, pars, return_negative=False,\
        larger_differences=1e-5, smaller_differences=1e-6,\
        transform_list=None):
        """
        Computes the hessian of this Loglikelihood for minimization purposes.

        pars: value of the parameters at which to evaluate the hessian
        return_negative: if true, the negative of the hessian of the
                         loglikelihood is returned (this is useful for times
                         when the loglikelihood must be maximized since scipy
                         optimization functions only deal with minimization
        larger_differences: either single number or 1D array of numbers to use
                            as the numerical difference in parameters.
                            Default: 10^(-5). This is the amount by which the
                            parameters are shifted between evaluations of the
                            gradient. Only used if gradient is not explicitly
                            computable.
        smaller_differences: either single_number or 1D array of numbers to use
                             as the numerical difference in parameters.
                             Default: 10^(-6). This is the amount by which the
                             parameters are shifted during each approximation
                             of the gradient. Only used if hessian is not
                             explicitly computable
        transform_list: TransformList object (or something which can be cast to
                        one) defining the transforms to apply to the parameters
                        before computing the gradient. Default: None, parameter
                        space is not transformed

        returns: square 2D numpy.ndarray of side length num_parameters
                 containing hessian of loglikelihood value
        """
        raise cannot_instantiate_loglikelihood_error

    def __eq__(self, other):
        """
        Checks if self is equal to other.

        other: a Loglikelihood object to check for equality

        returns: True if other and self have the same properties
        """
        raise NotImplementedError("The __eq__ magic method must be defined " +\
            "by each subclass of Loglikelihood individually. The class " +\
            "being used does not have the method defined.")

    def __ne__(self, other):
        """
        Checks if self is equal to other.

        other: a Loglikelihood object to check for equality

        returns: True if other and self do not have the same properties
        """
        return (not self.__eq__(other))

    def fisher_information(self, maximum_likelihood_parameters,\
        larger_differences=1e-5, smaller_differences=1e-6,\
        transform_list=None):
        """
        Calculates the Fisher information matrix of this likelihood assuming
        that the argument associated with the maximum of this likelihood is
        reasonably approximated by the given parameters.

        maximum_likelihood_parameters: the maximum likelihood  parameter vector
                                       (or some approximation of it)
        larger_differences: either single number or 1D array of numbers to use
                            as the numerical difference in parameters.
                            Default: 10^(-5). This is the amount by which the
                            parameters are shifted between evaluations of the
                            gradient. Only used if gradient is not explicitly
                            computable.
        smaller_differences: either single_number or 1D array of numbers to use
                             as the numerical difference in parameters.
                             Default: 10^(-6). This is the amount by which the
                             parameters are shifted during each approximation
                             of the gradient. Only used if hessian is not
                             explicitly computable
        transform_list: TransformList object (or something which can be cast to
                        one) defining the transforms to apply to the parameters
                        before computing the gradient. Default: None, parameter
                        space is not transformed. No matter what,
                        maximum_likelihood_parameters should be the parameters
                        that maximize the likelihood when plugged into the
                        model of this likelihood untransformed.

        returns: numpy.ndarray of shape (num_parameters, num_parameters)
                 containing the Fisher information matrix
        """
        return self.auto_hessian(maximum_likelihood_parameters,\
            larger_differences=larger_differences,\
            smaller_differences=smaller_differences,\
            transform_list=transform_list, return_negative=True)

    def parameter_covariance_fisher_formalism(self,\
        maximum_likelihood_parameters, transform_list=None,\
        max_standard_deviations=np.inf, larger_differences=1e-5,\
        smaller_differences=1e-6):
        """
        Finds the parameter covariance assuming maximum_likelihood_parameters
        contains a reasonable approximation of the true maximum likelihood
        parameter vector.

        maximum_likelihood_parameters: the maximum likelihood parameter vector
                                       (or some approximation of it), given in
                                       untransformed space, no matter the value
                                       of the transform_list argument
        transform_list: TransformList object (or something which can be cast to
                        one) defining the transforms to apply to the parameters
                        before computing the gradient. Default: None, parameter
                        space is not transformed. No matter what,
                        maximum_likelihood_parameters should be the parameters
                        that maximize the likelihood when plugged into the
                        model of this likelihood untransformed.
        max_standard_deviations: single value or array of values containing the
                                 maximum allowable standard deviations of each
                                 parameter. This will stop the covariance from
                                 producing extremely wide results in the case
                                 of an unconstrained parameter. The default
                                 value is numpy.inf, which causes this
                                 correction to be unimportant in all cases.
        larger_differences: either single number or 1D array of numbers to use
                            as the numerical difference in parameters.
                            Default: 10^(-5). This is the amount by which the
                            parameters are shifted between evaluations of the
                            gradient. Only used if gradient is not explicitly
                            computable.
        smaller_differences: either single_number or 1D array of numbers to use
                             as the numerical difference in parameters.
                             Default: 10^(-6). This is the amount by which the
                             parameters are shifted during each approximation
                             of the gradient. Only used if hessian is not
                             explicitly computable

        returns: numpy.ndarray of shape (num_parameters, num_parameters)
                 containing the inverse Fisher information matrix
        """
        inverse_covariance = self.fisher_information(\
            maximum_likelihood_parameters, transform_list=transform_list,\
            larger_differences=larger_differences,\
            smaller_differences=smaller_differences)
        if np.any(max_standard_deviations == 0):
            raise ValueError("At least one of the max_standard deviations " +\
                "was set to 0, which implies the existence of at least one " +\
                "element in the null space of the covariance matrix, which " +\
                "does not make sense.")
        max_standard_deviations =\
            max_standard_deviations * np.ones(self.num_parameters)
        inverse_covariance = inverse_covariance +\
            np.diag(np.power(max_standard_deviations, -2))
        return la.inv(inverse_covariance)

    def parameter_distribution_fisher_formalism(self,\
        maximum_likelihood_parameters, transform_list=None,\
        max_standard_deviations=np.inf,\
        prior_to_impose_in_transformed_space=None,\
        larger_differences=1e-5, smaller_differences=1e-6,\
        covariance_reduction_factor=1):
        """
        Finds the parameter distribution assuming maximum_likelihood_parameters
        contains a reasonable approximation of the true maximum likelihood
        parameter vector.

        maximum_likelihood_parameters: the maximum likelihood  parameter vector
                                       (or some approximation of it)
        transform_list: TransformList object (or something which can be cast to
                        one) defining the transforms to apply to the parameters
                        before computing the gradient. Default: None, parameter
                        space is not transformed. No matter what,
                        maximum_likelihood_parameters should be the parameters
                        that maximize the likelihood when plugged into the
                        model of this likelihood untransformed.
        max_standard_deviations: single value or array of values containing the
                                 maximum allowable standard deviations of each
                                 parameter. This will stop the covariance from
                                 producing extremely wide results in the case
                                 of an unconstrained parameter. The default
                                 value is numpy.inf, which causes this
                                 correction to be unimportant in all cases.
        prior_to_impose_in_transformed_space: if None (default), no prior is
                                                                 imposed and a
                                                                 Gaussian is
                                                                 returned
                                                                 through the
                                                                 Fisher matrix
                                                                 formalism
                                              otherwise, prior_to_impose should
                                                         be a Distribution
                                                         object whose log_value
                                                         function returns
                                                         -np.inf in disallowed
                                                         regions. The prior has
                                                         no effect inside the
                                                         region in which it is
                                                         finite.
        larger_differences: either single number or 1D array of numbers to use
                            as the numerical difference in parameters.
                            Default: 10^(-5). This is the amount by which the
                            parameters are shifted between evaluations of the
                            gradient. Only used if gradient is not explicitly
                            computable.
        smaller_differences: either single_number or 1D array of numbers to use
                             as the numerical difference in parameters.
                             Default: 10^(-6). This is the amount by which the
                             parameters are shifted during each approximation
                             of the gradient. Only used if hessian is not
                             explicitly computable
        covariance_reduction_factor: factor by which to decrease the magnitude
                                     of the covariance distribution: default 1

        returns: DistributionSet object containing GaussianDistribution object
                 approximating distribution in transformed space
        """
        transform_list = TransformList.cast(transform_list,\
            num_transforms=self.num_parameters)
        mean = transform_list(maximum_likelihood_parameters)
        covariance = self.parameter_covariance_fisher_formalism(\
            maximum_likelihood_parameters, transform_list=transform_list,\
            max_standard_deviations=max_standard_deviations,\
            larger_differences=larger_differences,\
            smaller_differences=smaller_differences) /\
            covariance_reduction_factor
        distribution = GaussianDistribution(mean, covariance)
        if type(prior_to_impose_in_transformed_space) is not type(None):
            distribution = WindowedDistribution(distribution,\
                prior_to_impose_in_transformed_space)
        return\
            DistributionSet([(distribution, self.parameters, transform_list)])


""" Build the resource task network for steel plants.

Notations
    Resource Category
    EAF     -   equipment unit in the EAF stage (first stage)
    EAF1    -   (if multiple units in this stage) one equipment unit in the EAF stage
    H_A_S3  -   (RTN1) intermediate product after stage 3
    H_B_S4  -   (RTN1) intermediate product before stage 4
    H_A_LF2 -   (RTN2) intermediate product before unit LF2
    EN      -   electric energy
    Task Category
    EAF     -   process task by unit EAF, i.e. melting
    TR_S3   -   (RTN1) transportation task between stage 3 and stage 4
    TR_LF2_CC1    -   (RTN2) transportation task between units LF2 and CC1

__author__ = xxxzhang
"""

from collections import OrderedDict
import math

import numpy as np
import json
import logging

log = logging.getLogger('steel')

class PlantRtnBuilder():
    """ resources, tasks, and their interaction networks for steel plant scheduling
    """

    def __init__(self, plant):
        self.rtn_t0 = plant['rtn_t0']  # time grid interval for RTN
        self.price_energy = []
        for price in plant['energy_price']:
            self.price_energy += [price] * (60 / self.rtn_t0)
        self.num_t = 24 * 60 / self.rtn_t0
        self.num_stage = len(plant['stage2units'].keys())
        self.num_groups = len(plant['group2heats'].keys())
        self.num_heats = 0
        for group, heats in plant['group2heats'].items():
            self.num_heats += len(heats)

        # basic information
        self.group2heats = plant['group2heats']
        self.stage2units = plant['stage2units']
        self.time_trans_max = plant['trans_time_max']
        self.time_trans = plant['trans_time']
        self.time_setup = plant['setup_time']
        self.time_process = plant['equip2process_time']

        # resources and tasks
        # todo instead using arrays, using dicts in tasks/res
        [self.resources, self.num_resources] = self.build_resources(plant)
        [self.tasks, self.num_tasks] = self.build_tasks(plant)
        [self.task_duration, self.task_cleanup_duration] = self.cal_task_duration(plant)
        [self.rtn_profile, self.heat_consume_time_in_group, self.heat_generate_time_in_group] \
            = self.build_resource_task_profile(plant)

        # redundant information
        self.equip2num = {}
        for unit2num in self.stage2units.values():
            for unit, num in unit2num.items():
                self.equip2num[unit] = num
        self.maxPower = sum([float(plant['equip2mw'][unit])*num for unit, num in self.equip2num.items()])
        self.casters = [unit for unit, num in self.stage2units['4'].items()]
        self.heat_sequence = ['EAF', 'TR_EA', 'AOD', 'TR_AL', 'LF', 'TR_LC']
        self.main_process = ['EAF', 'AOD', 'LF', 'CC1', 'CC2']

        log.debug('TaskCategory : (Task, Duration)')
        for task_type, task_list in self.tasks.items():
            log.debug('%s : %s' % (task_type, ','.join('(%s,%d)' % (x, self.task_duration[x]) for x in task_list)))
        log.debug('ResourceCategory : Resource')
        for res_cat, res_list in self.resources.items():
            log.debug('%s : %s' % (res_cat, ','.join(str(x) for x in res_list)))
        # log.info('ResourceTaskInteraction')
        # log.info(json.dumps(self.rtn_profile, indent=2))

    def build_resources(self, plant):
        # todo, do you need the reverse mapping
        res_cat2idx = OrderedDict()
        r_idx = 1
        # equipment resources
        for stage in range(1, self.num_stage + 1):
            for unit in plant['stage2units'][str(stage)].keys():
                res_cat2idx[unit] = [r_idx]
                r_idx += 1
        # intermediate products and final products, A - after, B - before
        for stage in range(1, self.num_stage+1):
            res_cat2idx['H_A_S%s' % stage] = range(r_idx, r_idx + self.num_heats)
            r_idx += self.num_heats
            if int(stage) == 1:
                continue
            res_cat2idx['H_B_S%s' % stage] = range(r_idx, r_idx + self.num_heats)
            r_idx += self.num_heats
        # energy resource
        res_cat2idx['EN'] = [r_idx]
        res_num = r_idx
        return [res_cat2idx, res_num]

    def build_tasks(self, plant):
        tasks = OrderedDict()
        i_idx = 1
        for stage in range(1, self.num_stage + 1):
            for unit in plant['stage2units'][str(stage)].keys():
                if int(stage) == 4:
                    tasks[unit] = range(i_idx, i_idx + self.num_groups)
                    i_idx += self.num_groups
                else:
                    tasks[unit] = range(i_idx, i_idx + self.num_heats)
                    i_idx += self.num_heats
                    tasks['TR_S%s' % stage] = range(i_idx, i_idx + self.num_heats)
                    i_idx += self.num_heats
        task_num = i_idx - 1
        return [tasks, task_num]

    def cal_task_duration(self, plant):
        """calculate time slots duration of tasks"""
        task_duration = dict()
        task_cleanup_duration = dict()
        for heat in range(1, self.num_heats + 1):
            for task_type in ['EAF', 'AOD', 'LF']:
                task_duration[self.tasks[task_type][heat - 1]] = \
                    int(math.ceil(float(plant['equip2process_time'][task_type][str(heat)])/self.rtn_t0))
            for task_type in ['TR_S1', 'TR_S2', 'TR_S3']:
                task_duration[self.tasks[task_type][heat - 1]] = \
                    int(math.ceil(float(plant['trans_time'][task_type])/self.rtn_t0))
        for group in range(1, self.num_groups + 1):
            for task_type in ['CC1', 'CC2']:
                cast_time = [plant['equip2process_time'][task_type][str(heat)] for heat in self.group2heats[str(group)]]
                total_time = float(sum(cast_time))
                task_duration[self.tasks[task_type][group - 1]] = int(math.ceil(total_time/self.rtn_t0))
                total_time = float(sum(cast_time) + self.time_setup[task_type])
                task_cleanup_duration[self.tasks[task_type][group - 1]] = int(math.ceil(total_time/self.rtn_t0))
        return task_duration, task_cleanup_duration

    def build_resource_task_profile(self, plant):
        """calculate resource task profile, i.e. the interactions between resource and task"""
        rtn_profile = dict()
        for res_category, res_idxes in self.resources.items():
            for res_idx in res_idxes:
                rtn_profile[res_idx] = dict()
        # equipment usage
        for stage in range(1, self.num_stage + 1):
            for unit in plant['stage2units'][str(stage)].keys():
                res = self.resources[unit][0]
                for task in self.tasks[unit]:
                    rtn_profile[res][task] = [-1] + [0] * (self.task_duration[task] - 1) + [1]
                    if stage == 4:
                        rtn_profile[res][task] = [-1] + [0] * (self.task_cleanup_duration[task] - 1) + [1]
        # heat consumption and generation for the first three stages
        for heat_idx in range(0, self.num_heats):
            # process task generate intermediate heat H_A_S (after stage)
            for task_cat, res_cat in [('EAF', 'H_A_S1'), ('AOD', 'H_A_S2'), ('LF', 'H_A_S3')]:
                task = self.tasks[task_cat][heat_idx]
                resource = self.resources[res_cat][heat_idx]
                rtn_profile[resource][task] = [0] * self.task_duration[task] + [1]
            # process task consume intermediate heat H_B_S (before stage)
            for task_cat, res_cat in [('AOD', 'H_B_S2'), ('LF', 'H_B_S3')]:
                task = self.tasks[task_cat][heat_idx]
                resource = self.resources[res_cat][heat_idx]
                rtn_profile[resource][task] = [-1] + [0] * self.task_duration[task]
            # transfer task transports heat
            for task_cat, res_cat in [('TR_S1', ['H_A_S1', 'H_B_S2']),
                                      ('TR_S2', ['H_A_S2', 'H_B_S3']),
                                      ('TR_S3', ['H_A_S3', 'H_B_S4'])]:
                task = self.tasks[task_cat][heat_idx]
                resource1 = self.resources[res_cat[0]][heat_idx]
                rtn_profile[resource1][task] = [-1] + [0] * self.task_duration[task]
                resource2 = self.resources[res_cat[1]][heat_idx]
                rtn_profile[resource2][task] = [0] * self.task_duration[task] + [1]
        # group-heat consumption and generation
        heat_consume_time_in_group = dict()
        heat_generate_time_in_group = dict()
        for unit in plant['stage2units']['4']:
            heat_consume_time_in_group[unit] = dict()
            heat_generate_time_in_group[unit] = dict()
            for group, heats in self.group2heats.items():
                # todo because the task numbers are stored in an array, can we do a better mapping?
                task = self.tasks[unit][int(group)-1]
                duration = 0
                for heat in heats:
                    consume_time = int(math.floor(duration/self.rtn_t0))
                    resource = self.resources['H_B_S4'][heat - 1]
                    rtn_profile[resource][task] = [0] * (self.task_duration[task] + 1)
                    rtn_profile[resource][task][consume_time] = -1
                    duration += plant['equip2process_time'][unit][str(heat)]
                    generate_time = int(math.ceil(duration/self.rtn_t0))
                    resource = self.resources['H_A_S4'][heat - 1]
                    rtn_profile[resource][task] = [0] * (self.task_duration[task] + 1)
                    rtn_profile[resource][task][generate_time] = 1
                    heat_consume_time_in_group[unit][heat] = consume_time
                    heat_generate_time_in_group[unit][heat] = generate_time
        # energy usage
        for stage, units in plant['stage2units'].items():
            stage = int(stage)
            for unit in units.keys():
                norm_mw = float(plant['equip2mw'][unit])
                for task in self.tasks[unit]:
                    total_energy = 0
                    if stage < 4:
                        heat = task - self.tasks[unit][0] + 1
                        total_energy = norm_mw * plant['equip2process_time'][unit][str(heat)] / 60
                    elif stage == 4:
                        group = task - self.tasks[unit][0] + 1
                        for heat in self.group2heats[str(group)]:
                            total_energy += norm_mw * plant['equip2process_time'][unit][str(heat)] / 60
                    profile = [norm_mw * self.rtn_t0 / 60] * (self.task_duration[task] - 1)
                    profile = profile + [total_energy - sum(profile)] + [0]
                    rtn_profile[self.resources['EN'][0]][task] = profile
        return rtn_profile, heat_consume_time_in_group, heat_generate_time_in_group

    def cal_initial_resource(self, res):
        # todo, here needs reverse mapping
        # only equipment initial value not zero
        for stage, stage_units in self.stage2units.items():
            for unit in stage_units.keys():
                if res in self.resources[unit]:
                    return stage_units[unit]
        return 0

    # todo delete this
    def get_unit2num(self):
        unit2num = dict()
        for stage, units in self.stage2units.items():
            for key, value in units.items():
                unit2num[key] = value
        return unit2num

    def get_same_heats(self):
        """find heats which share same processing times"""
        process_time = np.zeros((self.num_heats, len(self.main_process)))
        for i in range(len(self.main_process)):
            for heat_ in range(self.num_heats):
                process_time[heat_, i] = self.time_process[self.main_process[i]][heat_ + 1]
        visited = []
        same_heat = dict()
        for row in range(self.num_heats):
            if row not in visited:
                visited.append(row)
                same_heat[row + 1] = [row + 1]
                for row2 in range(row + 1, self.num_heats):
                    if np.array_equal(process_time[row2], process_time[row]):
                        same_heat[row + 1].append(row2 + 1)
                        visited.append(row2)
        return same_heat

class PlantRtn2Builder(PlantRtnBuilder):
    """Build RTN2 model."""

    def build_resources(self, plant):
        res_cat2idx = OrderedDict()
        r_idx = 1
        # equipment resources
        for stage in range(1, self.num_stage + 1):
            for unit in plant['stage2units'][str(stage)].keys():
                res_cat2idx[unit] = [r_idx]
                r_idx += 1
        # intermediate products, A - after, B - before
        for stage in range(1, self.num_stage+1):
            for unit in plant['stage2units'][str(stage)].keys():
                if int(stage) != 1:
                    res_cat2idx['H_B_%s' % unit] = range(r_idx, r_idx + self.num_heats)
                    r_idx += self.num_heats
                if int(stage) != 4:
                    res_cat2idx['H_A_%s' % unit] = range(r_idx, r_idx + self.num_heats)
                    r_idx += self.num_heats
        # final products
        res_cat2idx['H_FINAL'] = range(r_idx, r_idx + self.num_heats)
        r_idx += self.num_heats
        # energy resource
        res_cat2idx['EN'] = [r_idx]
        res_num = r_idx
        return [res_cat2idx, res_num]

    def build_tasks(self, plant):
        tasks = OrderedDict()
        i_idx = 1
        for stage in range(1, self.num_stage + 1):
            for unit in plant['stage2units'][str(stage)].keys():
                if int(stage) == 4:
                    # casting task
                    tasks[unit] = range(i_idx, i_idx + self.num_groups)
                    i_idx += self.num_groups
                else:
                    # process task for first 3 stages
                    tasks[unit] = range(i_idx, i_idx + self.num_heats)
                    i_idx += self.num_heats
                    # transfer task
                    for unit_2 in plant['stage2units'][str(stage+1)].keys():
                        tasks['TR_%s_%s' % (unit, unit_2)] = range(i_idx, i_idx + self.num_heats)
                        i_idx += self.num_heats
        task_num = i_idx - 1
        return [tasks, task_num]

    def cal_task_duration(self, plant):
        """calculate time slots duration of tasks"""
        task_duration = dict()
        task_cleanup_duration = dict()
        for heat in range(1, self.num_heats + 1):
            for stage in range(1, self.num_stage):
                for unit in plant['stage2units'][str(stage)].keys():
                    # process task
                    task_duration[self.tasks[unit][heat - 1]] = \
                        int(math.ceil(float(plant['equip2process_time'][unit][str(heat)])/self.rtn_t0))
                    if stage == 4:
                        continue
                    # transfer task
                    for unit_2 in plant['stage2units'][str(stage+1)].keys():
                        transfer_task = 'TR_%s_%s' % (unit, unit_2)
                        transfer = '%s-%s' % (unit, unit_2)
                        task_duration[self.tasks[transfer_task][heat - 1]] = \
                            int(math.ceil(float(plant['trans_time'][transfer])/self.rtn_t0))
        for group in range(1, self.num_groups + 1):
            for task_type in ['CC1', 'CC2']:
                cast_time = [plant['equip2process_time'][task_type][str(heat)] for heat in self.group2heats[str(group)]]
                total_time = float(sum(cast_time))
                task_duration[self.tasks[task_type][group - 1]] = int(math.ceil(total_time/self.rtn_t0))
                total_time = float(sum(cast_time) + self.time_setup[task_type])
                task_cleanup_duration[self.tasks[task_type][group - 1]] = int(math.ceil(total_time/self.rtn_t0))
        return task_duration, task_cleanup_duration

    def build_resource_task_profile(self, plant):
        """calculate resource task profile, i.e. the interactions between resource and task"""
        rtn_profile = dict()
        for res_category, res_indices in self.resources.items():
            for res_idx in res_indices:
                rtn_profile[res_idx] = dict()
        # equipment usage
        for stage in range(1, self.num_stage + 1):
            for unit in plant['stage2units'][str(stage)].keys():
                res = self.resources[unit][0]
                for task in self.tasks[unit]:
                    rtn_profile[res][task] = [-1] + [0] * (self.task_duration[task] - 1) + [1]
                    if stage == 4:
                        rtn_profile[res][task] = [-1] + [0] * (self.task_cleanup_duration[task] - 1) + [1]
        # intermediate product with processes and transfers in the first 3 stages
        for heat_idx in range(0, self.num_heats):
            for stage in range(1, self.num_stage + 1):
                if stage == 4:
                    continue
                for unit in plant['stage2units'][str(stage)].keys():
                    # process task generate intermediate heat
                    task_cat = unit
                    res_cat = 'H_A_%s' % unit
                    task = self.tasks[task_cat][heat_idx]
                    resource = self.resources[res_cat][heat_idx]
                    rtn_profile[resource][task] = [0] * self.task_duration[task] + [1]
                    # process task consume intermediate heat
                    if stage != 1:
                        task_cat = unit
                        res_cat = 'H_B_%s' % unit
                        task = self.tasks[task_cat][heat_idx]
                        resource = self.resources[res_cat][heat_idx]
                        rtn_profile[resource][task] = [-1] + [0] * self.task_duration[task]
                    # transfer task
                    for unit_2 in plant['stage2units'][str(stage+1)].keys():
                        task_cat = 'TR_%s_%s' % (unit, unit_2)
                        res_from = 'H_A_%s' % unit
                        res_to = 'H_B_%s' % unit_2
                        task = self.tasks[task_cat][heat_idx]
                        resource1 = self.resources[res_from][heat_idx]
                        rtn_profile[resource1][task] = [-1] + [0] * self.task_duration[task]
                        resource2 = self.resources[res_to][heat_idx]
                        rtn_profile[resource2][task] = [0] * self.task_duration[task] + [1]
        # group-heat consumption and generation
        heat_consume_time_in_group = dict()
        heat_generate_time_in_group = dict()
        for unit in plant['stage2units']['4']:
            heat_consume_time_in_group[unit] = dict()
            heat_generate_time_in_group[unit] = dict()
            for group, heats in self.group2heats.items():
                task = self.tasks[unit][int(group)-1]
                duration = 0
                res_cat_before = 'H_B_%s' % unit
                res_cat_after = 'H_FINAL'
                for heat in heats:
                    consume_time = int(math.floor(duration/self.rtn_t0))
                    resource = self.resources[res_cat_before][heat - 1]
                    rtn_profile[resource][task] = [0] * (self.task_duration[task] + 1)
                    rtn_profile[resource][task][consume_time] = -1
                    duration += plant['equip2process_time'][unit][str(heat)]
                    generate_time = int(math.ceil(duration/self.rtn_t0))
                    resource = self.resources[res_cat_after][heat - 1]
                    rtn_profile[resource][task] = [0] * (self.task_duration[task] + 1)
                    rtn_profile[resource][task][generate_time] = 1
                    heat_consume_time_in_group[unit][heat] = consume_time
                    heat_generate_time_in_group[unit][heat] = generate_time
        # energy usage
        for stage, units in plant['stage2units'].items():
            stage = int(stage)
            for unit in units.keys():
                norm_mw = float(plant['equip2mw'][unit])
                for task in self.tasks[unit]:
                    total_energy = 0
                    if stage < 4:
                        heat = task - self.tasks[unit][0] + 1
                        total_energy = norm_mw * plant['equip2process_time'][unit][str(heat)] / 60
                    elif stage == 4:
                        group = task - self.tasks[unit][0] + 1
                        for heat in self.group2heats[str(group)]:
                            total_energy += norm_mw * plant['equip2process_time'][unit][str(heat)] / 60
                    profile = [norm_mw * self.rtn_t0 / 60] * (self.task_duration[task] - 1)
                    profile = profile + [total_energy - sum(profile)] + [0]
                    rtn_profile[self.resources['EN'][0]][task] = profile
        return rtn_profile, heat_consume_time_in_group, heat_generate_time_in_group


import os
import sys

import gym
import eplus_env

# Assign mpc_path to be the file path where mpc.torch is located.
mpc_path = os.path.abspath(os.path.join(__file__, '..', '..' ))
sys.path.insert(0, mpc_path)

import argparse
from numpy import genfromtxt
import numpy as np
import pickle
import pandas as pd

from diff_mpc import mpc
from diff_mpc.mpc import QuadCost, LinDx

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

from utils import make_dict, R_func

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DEVICE

parser = argparse.ArgumentParser(description='GruRL-Imitation Learning')
parser.add_argument('--seed', type=int, default=42, metavar='N',
                    help='random seed (default: 42)')
parser.add_argument('--lr', type=float, default=5e-4, metavar='G',
                    help='Learning Rate')
parser.add_argument('--T', type=int, default=12, metavar='N',
                    help='Planning Horizon (default: 12)')
parser.add_argument('--step', type=int, default=900, metavar='N',
                    help='Time Step in Simulation, Unit in Seconds (default: 900)') # 15 Minutes Now!
parser.add_argument('--eta', type=int, default=5,
                    help='Hyper Parameter for Balancing Comfort and Energy')
parser.add_argument('--batch_size', type=int, default=256,
                    help='Size of Mini-batch')
parser.add_argument('--save_name', type=str, default='rl',
                    help='save name')
args = parser.parse_args()

torch.manual_seed(args.seed)

# Modify here: Outputs from EnergyPlus; Match the variables.cfg file.
obs_name = ["Outdoor Temp.", "Outdoor RH", "Wind Speed", "Wind Direction", "Diff. Solar Rad.", "Direct Solar Rad.", "Htg SP", "Clg SP", "Indoor Temp.", "Indoor Temp. Setpoint", "PPD", "Occupancy Flag", "Coil Power", "HVAC Power", "Sys In Temp.", "Sys In Mdot", "OA Temp.", "OA Mdot", "MA Temp.", "MA Mdot", "Sys Out Temp.", "Sys Out Mdot"]

# Modify here: Change based on the specific control problem
state_name = ["Indoor Temp."]
dist_name = ["Outdoor Temp.", "Outdoor RH", "Wind Speed", "Wind Direction", "Diff. Solar Rad.", "Direct Solar Rad.", "Occupancy Flag"]
# Caveat: The RL agent controls the difference between Supply Air Temp. and Mixed Air Temp., i.e. the amount of heating from the heating coil. But, the E+ expects Supply Air Temp. Setpoint.
ctrl_name = ["Delta T"]
target_name = ["Indoor Temp. Setpoint"]

n_state = len(state_name)
n_ctrl = len(ctrl_name)
n_dist = len(dist_name)

eta = [0.1, args.eta] # eta: Weight for comfort during unoccupied and occupied mode
step = args.step # step: Timestep; Unit in seconds
T = args.T # T: Number of timesteps in the planning horizon
tol_eps = 90 # tol_eps: Total number of episodes; Each episode is a natural day

# Read Historical Data
dataset = pd.read_pickle("results/Sim-TMY2.pkl")
target = dataset[target_name]
disturbance = dataset[dist_name]
# Min-Max Normalization
disturbance = (disturbance-disturbance.min())/(disturbance.max()-disturbance.min())

dataset["Delta T"] = dataset["Sys Out Temp."]-dataset["MA Temp."]

# Train-Test Split
n_samples = len(dataset)
n_train = int(0.7*n_samples)
n_test = n_samples - n_train
train_set = dataset[:n_train]
test_set = dataset[n_train:]

class Learner():
    def __init__(self, n_state, n_ctrl, n_dist, disturbance, target, u_upper, u_lower):
        self.n_state = n_state
        self.n_ctrl = n_ctrl
        self.n_dist = n_dist
        self.disturbance = disturbance
        self.target = target

        # My Initial Guess
        self.F_hat = torch.ones((self.n_state, self.n_state+self.n_ctrl))
        self.F_hat[0, 0] = 0.9
        self.F_hat[0, 1] = 0.3
        self.F_hat = self.F_hat.double().requires_grad_()

        self.Bd_hat = np.random.rand(self.n_state, self.n_dist)
        self.Bd_hat = torch.tensor(self.Bd_hat).requires_grad_()

        self.optimizer = optim.Adam([self.F_hat, self.Bd_hat], lr=args.lr)

        self.u_lower = u_lower * torch.ones(T, 1, n_ctrl).double()
        self.u_upper = u_upper * torch.ones(T, 1, n_ctrl).double()

    def Cost_function(self, cur_time):
        diag = torch.zeros(T, self.n_state + self.n_ctrl)
        occupied = self.disturbance["Occupancy Flag"][cur_time:cur_time + pd.Timedelta(seconds = (T-1) * step)]
        occupied = np.array(occupied)
        if len(occupied)<T:
            occupied = np.pad(occupied, ((0, T-len(occupied)), ), 'edge')
        eta_w_flag = torch.tensor([eta[int(flag)] for flag in occupied]).unsqueeze(1).double() # Tx1
        diag[:, :n_state] = eta_w_flag
        diag[:, n_state:] = 0.001

        C = []
        for i in range(T):
            C.append(torch.diag(diag[i]))
        C = torch.stack(C).unsqueeze(1) # T x 1 x (m+n) x (m+n)

        x_target = self.target[cur_time : cur_time + pd.Timedelta(seconds = (T-1) * step)] # in pd.Series
        x_target = np.array(x_target)
        if len(x_target)<T:
            x_target = np.pad(x_target, ((0, T-len(x_target)), (0, 0)), 'edge')
        x_target = torch.tensor(x_target)

        c = torch.zeros(T, self.n_state+self.n_ctrl) # T x (m+n)
        c[:, :n_state] = -eta_w_flag*x_target
        c[:, n_state:] = 1 # L1-norm now! Check
        c = c.unsqueeze(1) # T x 1 x (m+n)
        return C, c

    def forward(self, x_init, C, c, cur_time):
        dt = np.array(self.disturbance[cur_time : cur_time + pd.Timedelta(seconds = (T-2) * step)]) # T-1 x n_dist
        if len(dt)<T-1:
            dt = np.pad(dt, ((0, T-1-len(dt)), (0, 0)), 'edge')
        dt = torch.tensor(dt).transpose(0, 1) # n_dist x T-1

        ft = torch.mm(self.Bd_hat, dt).transpose(0, 1) # T-1 x n_state
        ft = ft.unsqueeze(1) # T-1 x 1 x n_state

        x_pred, u_pred, _ = mpc.MPC(n_state=self.n_state,
                                    n_ctrl=self.n_ctrl,
                                    T=T,
                                    u_lower = self.u_lower,
                                    u_upper = self.u_upper,
                                    lqr_iter=20,
                                    verbose=0,
                                    exit_unconverged=False,
                                    )(x_init, QuadCost(C.double(), c.double()),
                                      LinDx(self.F_hat.repeat(T-1, 1, 1, 1),  ft))

        return x_pred[1, 0, :], u_pred[0, 0, :] # Dim.

    def predict(self, x_init, action, cur_time):
        dt = np.array(self.disturbance.loc[cur_time]) # n_dist
        dt = torch.tensor(dt).unsqueeze(1) # n_dist x 1
        ft = torch.mm(self.Bd_hat, dt) # n_state x 1
        tau = torch.stack([x_init, action]) # (n_state + n_ctrl) x 1
        next_state  = torch.mm(self.F_hat, tau) + ft # n_state x 1
        return next_state

    def update_parameters(self, x_true, u_true, x_pred, u_pred):
        # Every thing in T x Dim.
        state_loss = torch.mean((x_true.double() - x_pred)**2)
        action_loss = torch.mean((u_true.double() - u_pred)**2)

        # Note: args.eta balances the importance between predicting states and predicting actions
        traj_loss = args.eta*state_loss + action_loss
        print("From state {}, From action {}".format(state_loss, action_loss))
        self.optimizer.zero_grad()
        traj_loss.backward()
        self.optimizer.step()
        print(self.F_hat)
        print(self.Bd_hat)
        return state_loss.detach(), action_loss.detach()

def evaluate_performance(x_true, u_true, x_pred, u_pred):
    state_loss = torch.mean((x_true.double() - x_pred)**2)
    action_loss = torch.mean((u_true.double() - u_pred)**2)
    return state_loss, action_loss

def main():
    dir = 'results'
    if not os.path.exists(dir):
        os.mkdir(dir)

    perf = []
    n_step = 96 # n_step: Number of Steps per Day
    numOfEpoches = 20

    timeStamp = []
    record_name =["Learner nState", "Expert nState", "Learner action", "Expert action"]
    losses = []
    losses_name = ["train_state_loss", "train_action_loss", "val_state_loss", "val_action_loss"]

    # Initialize the learner
    u_upper = 5
    u_lower = 0
    learner = Learner(n_state, n_ctrl, n_dist, disturbance, target, u_upper, u_lower)

    for epoch in range(numOfEpoches):
        x_true = []
        u_true = []
        x_pred = []
        u_pred = []

        train_state_loss = []
        train_action_loss = []
        for i in range(n_train): # By number of entries in the historical data
            idx = np.random.randint(n_train)
            cur_time = train_set.index[idx]

            expert_moves = train_set[cur_time:cur_time+pd.Timedelta(seconds = step)]
            if len(expert_moves)<2:
                print(cur_time)
                continue

            expert_state = torch.tensor(expert_moves[state_name].values).reshape(-1, n_state) # 2 x n_state
            expert_action = torch.tensor(expert_moves[ctrl_name].values).reshape(-1, n_ctrl) # 2 x n_ctrl
            x_true.append(expert_state[-1])
            u_true.append(expert_action[0])

            obs = train_set.loc[cur_time]
            x_init = torch.tensor(np.array([obs[name] for name in state_name])).unsqueeze(0) # n_batch x n_state, i.e. 1 x n_state
            C, c = learner.Cost_function(cur_time)
            learner_state, learner_action = learner.forward(x_init, C, c, cur_time)

            # Predict next state based on expert's action
            next_state = learner.predict(x_init.squeeze(0), expert_action[0], cur_time)
            x_pred.append(next_state)
            u_pred.append(learner_action)

            if (i % args.batch_size == 0) & (i>0):
                x_true = torch.stack(x_true).reshape(-1, n_state)
                u_true = torch.stack(u_true).reshape(-1, n_ctrl)
                x_pred = torch.stack(x_pred).reshape(-1, n_state)
                u_pred = torch.stack(u_pred).reshape(-1, n_ctrl)
                b_state_loss, b_action_loss = learner.update_parameters(x_true, u_true, x_pred, u_pred)
                train_state_loss.append(b_state_loss)
                train_action_loss.append(b_action_loss)
                x_true = []
                u_true = []
                x_pred = []
                u_pred = []

        # Evaluate performance at the end of each epoch
        x_true = []
        u_true = []
        x_pred = []
        u_pred = []
        timeStamp = []
        for idx in range(n_test):
            cur_time = test_set.index[idx]
            expert_moves = test_set[cur_time:cur_time+pd.Timedelta(seconds = step)]
            if len(expert_moves)<2:
                print(cur_time)
                continue
            expert_state = torch.tensor(expert_moves[state_name].values).reshape(-1, n_state) # 2 x n_state
            expert_action = torch.tensor(expert_moves[ctrl_name].values).reshape(-1, n_ctrl) # 2 x n_ctrl
            x_true.append(expert_state[-1])
            u_true.append(expert_action[0])

            timeStamp.append(cur_time+pd.Timedelta(seconds = step))

            obs = test_set.loc[cur_time]
            x_init = torch.tensor(np.array([obs[name] for name in state_name])).unsqueeze(0) # 1 x n_state
            C, c = learner.Cost_function(cur_time)
            learner_state, learner_action = learner.forward(x_init, C, c, cur_time)
            next_state = learner.predict(x_init.squeeze(0), expert_action[0], cur_time)
            x_pred.append(next_state.detach())
            u_pred.append(learner_action.detach())

        x_true = torch.stack(x_true).reshape(-1, n_state)
        u_true = torch.stack(u_true).reshape(-1, n_ctrl)
        x_pred = torch.stack(x_pred).reshape(-1, n_state)
        u_pred = torch.stack(u_pred).reshape(-1, n_ctrl)
        val_state_loss, val_action_loss = evaluate_performance(x_true, u_true, x_pred, u_pred)
        print("At Epoch {0}, the loss from the state is {1} and from the action is {2}".format(epoch, val_state_loss, val_action_loss))
        losses.append((np.mean(train_state_loss), np.mean(train_action_loss), val_state_loss, val_action_loss))

        record = pd.DataFrame(torch.cat((x_pred, x_true, u_pred, u_true), dim = 1).numpy(), index = np.array(timeStamp), columns = record_name)
        record_df = pd.DataFrame(np.array(record), index = np.array(timeStamp), columns = record_name)
        record_df.to_pickle("results/Imit_{}_{}.pkl".format(args.save_name, epoch))

        # Save weights
        F_hat = learner.F_hat.detach().numpy()
        Bd_hat = learner.Bd_hat.detach().numpy()
        np.save("results/weights/F-{}.npy".format(epoch), F_hat)
        np.save("results/weights/Bd-{}.npy".format(epoch), Bd_hat)

    # Save losses at each epoch
    losses_df = pd.DataFrame(np.array(losses), index = np.arange(numOfEpoches), columns = losses_name)
    losses_df.to_pickle("results/Imit_loss_"+args.save_name+".pkl")

if __name__ == '__main__':
    main()


# Copyright © 2017 Ondrej Martinsky, All rights reserved
# http://github.com/omartinsky/pybor
#
# 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.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from yc_convention import *
import scipy.interpolate
import pylab, re, collections, matplotlib

class PlottingHelper:
    @staticmethod
    def set_tenors_on_axis(axis, start_date):
        tenors = "6M,1Y,2Y,3Y,4Y,5Y,7Y,10Y,15Y,20Y,30Y,40Y,50Y,60Y,70Y".split(",")
        tenordates = [date_step(int(start_date), Tenor(t)) for t in tenors]
        axis.xaxis.set_ticks(tenordates)
        axis.xaxis.set_major_formatter(matplotlib.ticker.FuncFormatter(lambda x, pos: tenors[pos]))

class CurveMap:
    def __init__(self, *arg, **kw):
        super(CurveMap, self).__init__(*arg, **kw)
        self.curves_ = collections.OrderedDict()

    def add_curve(self, c):
        assert_type(c, Curve)
        self.curves_[c.get_id()] = c

    def get_all_dofs(self, curves_for_stage):
        dofs = list()
        for k, v in self.curves_.items():
            if k in curves_for_stage:
                dofs.extend(v.get_all_dofs())
        return dofs

    def set_all_dofs(self, curves_for_stage, dofs):
        i = 0
        for k, v in self.curves_.items():
            if k in curves_for_stage:
                j = i + v.get_dofs_count()
                v.set_all_dofs(dofs[i:j])
                i = j

    def __getitem__(self, item):
        return self.curves_[item]

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

    def keys(self):
        return self.curves_.keys()

    def plot(self, reg=".*", *arg, **kwargs):
        for name, curve in sorted(self.curves_.items()):
            if re.match(reg, name):
                curve.plot(*arg, **kwargs)

class InterpolationMode(enum.Enum):
    LINEAR_LOGDF = 0
    LINEAR_CCZR = 1
    CUBIC_LOGDF = 2

class PlotMode(enum.Enum):
    DF = 0
    ZR = 1
    FWD = 2

class PlotDate(enum.Enum):
    YMD = 0
    EXCEL = 1
    TENOR = 2

class ExponentialInterpolator:
    def __init__(self, interp):
        self.interp = interp
    def value(self, t):
        return exp(self.interp(t))

class ZeroRateInterpolator:
    def __init__(self, interp, t_eval):
        self.interp = interp
        self.t_eval = t_eval
    def value(self, t):
        return exp(self.interp(t) * (t - self.t_eval))

class Curve:
    def __init__(self, curve_id, eval_date, times, dfs, interpolation_mode):
        try:
            times, dfs = array(times), array(dfs)
            assert_type(interpolation_mode, InterpolationMode)
            assert len(times) > 0, "Vector of times is empty"
            assert times[0] != eval_date, "DF at eval date cannot be provided externally. It is assumed to be 1.0 always."
            self.id_ = curve_id
            self.times_ = append(eval_date, times)
            self.dfs_ = append([1.], dfs)
            self.set_interpolator(interpolation_mode)
        except BaseException as ex:
            raise BaseException("Unable to create curve %s" % curve_id) from ex

    def add_another_curve(self, another_curve):
        assert isinstance(another_curve, Curve)
        assert all(self.times_ == another_curve.times_)
        self.dfs_ = another_curve.dfs_ * self.dfs_
        self.set_interpolator(self.interpolation_mode_)

    def set_interpolator(self, interpolation_mode=None):
        if interpolation_mode is not None:
            self.interpolation_mode_ = interpolation_mode
        if self.interpolation_mode_ in [InterpolationMode.LINEAR_LOGDF, InterpolationMode.LINEAR_CCZR]:
            kind = 'linear'
        elif self.interpolation_mode_ in [InterpolationMode.CUBIC_LOGDF]:
            kind = 'cubic'
        else:
            raise BaseException("Invalid interpolation mode. Allowed modes are %s" % enum_values_as_string(InterpolationMode))
        #
        assert len(self.times_) == len(self.dfs_), (len(self.times_), len(self.dfs_))
        #
        if self.interpolation_mode_ in [InterpolationMode.LINEAR_LOGDF, InterpolationMode.CUBIC_LOGDF]:
            logdf = log(self.dfs_)
            interp = scipy.interpolate.interp1d(self.times_, logdf, kind=kind)
            self.interpolator_ = ExponentialInterpolator(interp)
        elif self.interpolation_mode_ in [InterpolationMode.LINEAR_CCZR]:
            t_eval = self.times_[0]
            t_rel = self.times_-t_eval
            cczr1 = log(self.dfs_[1:]) / t_rel[1:]
            cczr = insert(cczr1, 0, cczr1[0]) # ZZCR at t0 is undefined, take it from t1 instead
            interp = scipy.interpolate.interp1d(self.times_, cczr, kind=kind)
            self.interpolator_ = ZeroRateInterpolator(interp, t_eval)
        else:
            raise BaseException("Invalid interpolation mode")

    def __str__(self):
        return self.id_

    def get_id(self):
        return self.id_

    def get_df(self, t):
        try:
            return self.interpolator_.value(t)
        except BaseException as ex:
            raise BaseException("Unable to get discount factor for dates [%i..%i] from curve with dates range [%i..%i]" % (t[0],t[-1],self.times_[0], self.times_[-1])) from ex

    def get_zero_rate(self, t, freq, dcc):
        dfs = self.get_df(t)
        dcf = calculate_dcf(self.times_[0], t, dcc)
        if freq == CouponFreq.ZERO:
            return (1. / dfs - 1.) / dcf
        if freq == CouponFreq.CONTINUOUS:
            return -log(dfs) / dcf

    def get_fwd_rate(self, t_start, t_end, freq, dcc):
        dfs_start = self.get_df(t_start)
        dfs_end = self.get_df(t_end)
        dcf = calculate_dcf(t_start, t_end, dcc)
        if freq == CouponFreq.ZERO:
            return (dfs_start / dfs_end - 1) / dcf
        if freq == CouponFreq.CONTINUOUS:
            return log(dfs_start / dfs_end) / dcf

    def get_fwd_rate_aligned(self, t, freq, dcc):
        # Slightly faster version which relies on the fact that calculation periods are aligned (no overlaps, no gaps)
        dfs = self.get_df(t)
        t1 = t[:-1]
        t2 = t[1:]
        df1 = dfs[:-1]
        df2 = dfs[1:]
        dcf = calculate_dcf(t1, t2, dcc)
        if freq == CouponFreq.ZERO:
            return (df1 / df2 - 1) / dcf
        if freq == CouponFreq.CONTINUOUS:
            return log(df1 / df2) / dcf

    def set_all_dofs(self, dofs):
        self.dfs_ = append([1], dofs)
        self.set_interpolator()

    def get_all_dofs(self):
        return self.dfs_[1:]

    def get_dofs_count(self):
        return len(self.dfs_) - 1

    def plot(self, date_style=PlotDate.YMD, mode=PlotMode.FWD, samples=1000, label=None, convention=None):
        timesample = linspace(self.times_[0], self.times_[-1], samples)
        X = timesample
        if date_style==PlotDate.YMD:
            X = [exceldate_to_pydate(int(x)) for x in X]
        elif date_style==PlotDate.TENOR:
            ax = matplotlib.pyplot.subplot()
            PlottingHelper.set_tenors_on_axis(ax, self.times_[0])
        elif date_style==PlotDate.EXCEL:
            pass
        else:
            raise BaseException("Unknown PlottingDateStyle")
        ###
        if mode==PlotMode.FWD:
            convention = global_conventions.get(self.id_) if convention is None else convention
            Y = self.get_fwd_rate_aligned(timesample, CouponFreq.ZERO, convention.dcc)
            pylab.plot(X[:-1], Y, label=self.id_ if label is None else label)
        elif mode==PlotMode.ZR:
            convention = global_conventions.get(self.id_) if convention is None else convention
            Y = self.get_zero_rate(timesample[1:], CouponFreq.ZERO, convention.dcc)
            pylab.plot(X[:-1], Y, label=self.id_ if label is None else label)
        elif mode==PlotMode.DF:
            Y = self.get_df(timesample)
            pylab.plot(X, Y, label=self.id_ if label is None else label)
        else:
            raise BaseException("Unknown PlottingMode")

class CurveConstructor:
    @staticmethod
    def FromShortRateModel(curve_id, times, r0, speed, mean, sigma, interpolation):
        import random
        times = array(times)
        assert_type(r0, float)
        assert_type(speed, float)
        assert_type(mean, float)
        assert_type(sigma, float)
        assert_type(interpolation, InterpolationMode)
        r = r0
        rates = []
        dts = times[1:] - times[:-1]
        dts = dts / 365.
        for dt in dts:
            rates.append(r)
            dr = speed * (mean - r) * dt + sigma * random.gauss(0, 1) * dt ** .5
            r += dr
        rates = array(rates)
        dfs_fwd = exp(-rates * dts)
        dfs = cumprod(dfs_fwd)
        return Curve(curve_id, times[0], times[1:], dfs, interpolation)


import numpy as np

"""
It is created for using MTC in Mujoco. The dynamics in this model is not continuous. The integration error will be
accumulated overtime. And the system might get unstable if the timestep is too large. It is recommended to set the
timestamp lower than 5e-4 to get decent results.

The model is created based on Song's and Geyer's 2015 paper:
Song, S. and Geyer, H., 2015. A neural circuitry that emphasizes spinal feedback generates diverse behaviours of human
locomotion. The Journal of physiology, 593(16), pp.3493-3511.

V0.1
Passed basic tests. There're slightly difference compared to the simmechanics model.

V0.2
1. Verified with the simmechanics model. Difference in most of the cases can be ignored.
2. Changed the integration method from forward Euler to trapezoid.
3. Muscle force vce etc might vibrate/jitter if in some cases if the timestep is not low enough.
   Need to improve this in the next version.

"""

class MuscleTendonComplex:
    def __init__(self, paraMuscle, stateMuscle, paraMusAttach, offsetCorr, timestep, nameMuscle, angJoi):
        self.frcmax, self.vmax, self.eref, self.lslack, self.lopt, self.tau, self.w, self.c, self.N, self.K = paraMuscle
        self.stim, self.act, self.lmtc, self.lce, self.vce, self.frcmtc = stateMuscle
        self.timestep = timestep
        self.nameMuscle = nameMuscle
        self.angJoi = angJoi
        self.offsetCorr = offsetCorr
        self.r, self.phiref, self.phimaxref, self.rho, self.dirAng, self.phiScale = paraMusAttach
        self.MR =  0.01
        self.typeMuscle = self.angJoi.size
        nJoi = self.typeMuscle
        self.levelArm = np.zeros(nJoi)

        tmpL = np.zeros(nJoi)
        for i in range(0, nJoi):
            if self.offsetCorr[i] == 0:
                tmpL[i] = self.dirAng[i] * (self.angJoi[i] - self.phiref[i]) * self.r[i] * self.rho[i]
                self.levelArm[i] = self.r[i]
            elif self.offsetCorr[i] == 1:
                tmp1 = np.sin((self.phiref[i] - self.phimaxref[i]) * self.phiScale[i])
                tmp2 = np.sin((self.angJoi[i] - self.phimaxref[i]) * self.phiScale[i])
                tmpL[i] = self.dirAng[i] * (tmp2 - tmp1) * self.r[i] * self.rho[i] / self.phiScale[i]
                self.levelArm[i] = np.cos((self.angJoi[i] - self.phimaxref[i]) * self.phiScale[i]) * self.r[i]
            else:
                raise ValueError('Invalid muscle level arm offset correction type. ')
        self.lmtc = self.lslack + self.lopt + np.sum(tmpL)

        self.lce = self.lmtc - self.lslack
        self.lse = self.lmtc - self.lce
        # unitless parameters
        self.Lse = self.lse / self.lslack
        self.Lce = self.lce / self.lopt

        self.actsubstep = (self.stim - self.act) * self.timestep / 2.0 / self.tau + self.act
        self.lcesubstep = self.vce * self.timestep / 2.0 + self.lce

        # test
        self.lce_avg = self.lce
        self.vce_avg = self.vce
        self.frcmtc_avg = 0
        self.act_avg = self.act
        self.frame = 0
        # self.Fse = 0.0
        # self.Fbe = 0.0
        # self.Fpe = 0.0
        # self.Fce = 0.0

    def stepUpdateState(self, angJoi):
        """
        Muscle Tendon Complex Dynamics
        update muscle states based on the muscle dynamics
        Muscle state stim has to be updated outside before this function is called
        """
        # update lmtc and level arm based on the geometry
        self.angJoi = angJoi
        nJoi = self.typeMuscle
        tmpL = np.zeros(nJoi)
        for i in range(0, nJoi):
            if self.offsetCorr[i] == 0:
                tmpL[i] = self.dirAng[i] * (self.angJoi[i] - self.phiref[i]) * self.r[i] * self.rho[i]
                self.levelArm[i] = self.r[i]
            elif self.offsetCorr[i] == 1:
                tmp1 = np.sin((self.phiref[i] - self.phimaxref[i]) * self.phiScale[i])
                tmp2 = np.sin((self.angJoi[i] - self.phimaxref[i]) * self.phiScale[i])
                tmpL[i] = self.dirAng[i] * (tmp2 - tmp1) * self.r[i] * self.rho[i] / self.phiScale[i]
                self.levelArm[i] = np.cos((self.angJoi[i] - self.phimaxref[i]) * self.phiScale[i]) * self.r[i]
            else:
                raise ValueError('Invalid muscle level arm offset correction type. ')
        self.lmtc = self.lslack + self.lopt + np.sum(tmpL)

        # update muscle activation
        # integration, forward-Euler method
        # self.act = (self.stim - self.act) * self.timestep / self.tau + self.act
        # integration, trapezoidal method, 2-step
        self.act = (self.stim - self.actsubstep) * self.timestep / 2.0 / self.tau + self.actsubstep
        self.actsubstep = (self.stim - self.act) * self.timestep / 2.0 / self.tau + self.act

        # update lce and lse based on the lmtc
        # integration, forward-Euler method
        # self.lce = self.vce * self.timestep + self.lce
        # integration, trapezoidal method, 2-step
        self.lce = self.vce * self.timestep / 2.0 + self.lcesubstep
        self.lcesubstep = self.vce * self.timestep / 2.0 + self.lce

        self.lse = self.lmtc - self.lce
        self.Lse = self.lse / self.lslack
        self.Lce = self.lce / self.lopt

        # Serial Elastic element (tendon) force-length relationship
        if self.Lse > 1.0:
            Fse = np.power((self.Lse - 1.0) / self.eref, 2)
        else:
            Fse = 0.0

        # Parallel Elasticity PE
        if self.Lce > 1.0:
            Fpe = np.power((self.Lce - 1.0) / self.w, 2)
        else:
            Fpe = 0.0

        # update frcmtc
        self.frcmtc = Fse * self.frcmax
        #self.frcmtc =  np.clip(self.frcmtc, 0, self.frcmax)

        # Buffer Elasticity BE
        if (self.Lce - (1.0 - self.w)) < 0:
            Fbe = np.power((self.Lce - (1.0 - self.w)) / (self.w / 2), 2)
        else:
            Fbe = 0.0

        # Contractile Element force-length relationship
        tmp = np.power(np.absolute(self.Lce - 1.0) / self.w, 3)
        Fce = np.exp(tmp * np.log(self.c))

        #Fv = (Fse + Fbe) / (Fpe + Fce * self.act)
        if (Fpe + Fce * self.act) < 1e-10:  # avoid numerical error
            if (Fse + Fbe) < 1e-10:
                Fv = 1.0
            else:
                Fv = (Fse + Fbe) / 1e-10
        else:
            Fv = (Fse + Fbe) / (Fpe + Fce * self.act)

        # Contractile Element inverse force-velocity relationship
        if Fv <= 1.0:
            # Concentric
            v = (Fv - 1) / (Fv * self.K + 1.0)
        elif Fv <= self.N:
            # excentric
            tmp = (Fv - self.N) / (self.N - 1.0)
            v = (tmp + 1.0) / (1.0 - tmp * 7.56 * self.K)
        else:
            # excentric overshoot
            v = ((Fv - self.N) * 0.01 + 1)

        self.vce = v * self.lopt * self.vmax
        v_frac = self.vce /  self.vmax
        mr_scale =  self.act * np.absolute(self.frcmax*self.vmax) *self.timestep
        if self.vce <= 1:
            self.MR =  0.01 - 0.11*(v_frac) + 0.06*np.exp(-8*v_frac)
        else:
            self.MR =  0.23 - 0.16*np.exp(-8*v_frac)
        self.MR *= mr_scale
        self.frame += 1
        self.lce_avg = (self.lce_avg*(self.frame - 1) +  self.lce) / self.frame
        self.vce_avg = (self.vce_avg*(self.frame - 1) +  self.vce) / self.frame
        self.frcmtc_avg = (self.frcmtc_avg*(self.frame - 1) +  self.frcmtc) / self.frame
        self.act_avg = (self.act_avg*(self.frame - 1) +  self.act) / self.frame
        #self.MR = np.exp(-self.MR)
        # print(self.MR, np.exp(-self.MR))
        # self.Fv = Fv
        # self.Fse = Fse
        # self.Fbe = Fbe
        # self.Fpe = Fpe
        # self.Fce = Fce

    def reset_state(self):
        self.frame = 0
        self.lce_avg = 0
        self.frcmtc_avg = 0
        self.act_avg = 0
        self.vce_avg = 0