"""
The :py:mod:`h2_mobility` module contains a class to read the required data and
a class to evaluate the power-to-mobility system.
"""

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import pvlib

class ReadData:
    """

    This class enables to read data from the data files.

    Parameters
    ----------
    filename_climate : str
        The directory of the file with information on the
        solar irradiance.

    """

    def __init__(self, filename_climate):
        self.filename_climate = filename_climate
        self.path = os.path.dirname(os.path.abspath(__file__))

    def load_climate(self):
        """

        This method loads the hourly solar irradiance data
        and ambient temperature data,
        situated in the 'sol_irr' and 'T_amb' columns of the
        climate data file.

        Returns
        -------
        sol_irr : ndarray
            The hourly solar irradiance data for a Typical
            Meteorological Year.
        t_amb : ndarray
            The hourly ambient temperature data for a Typical
            Meteorological Year.

        """
        data = pd.read_csv(self.filename_climate)
        sol_irr = data['sol_irr'].to_numpy()
        t_amb = data['T_amb'].to_numpy()

        return sol_irr, t_amb

    def load_parameters(self):
        """

        This method loads the deterministic values of the model
        parameters, defined in the design_space file. This is
        useful when the deterministic performance of a specific
        design needs to be evaluated.

        Returns
        -------
        param_dict : dict
            Dictionary with the names of the model parameters
            and the corresponding deterministic values.

        """
        param_dict = {}
        design_space = os.path.join(self.path, 'design_space')

        # read the deterministic values for the parameters in `design_space`
        with open(design_space, 'r') as file:
            for line in file:
                tmp = line.split()
                if tmp[1] == 'par':
                    param_dict[tmp[0]] = float(tmp[2])

        return param_dict

class Evaluation:
    """

    This class evaluates the power-to-mobility system.
    For a given design, the solar irradiance, ambient temperature
    and the characterization of the model parameters,
    the levelized cost of driving, carbon intensity and the annual
    grid consumption are quantified.

    Parameters
    ----------
    sol_irr : ndarray
        The hourly solar irradiance for the evaluated year.
    t_amb : ndarray
        The hourly ambient temperature for the evaluated year.
    parameters : dict
        Dictionary with the model parameters and design variables values.

    """

    def __init__(self, sol_irr, t_amb, par):
        self.par = par

        # the solar irradiance and ambient temperature are scaled with the
        # corresponding uncertainty
        self.sol_irr = sol_irr * self.par['u_sol_irr']
        self.t_amb = t_amb + self.par['u_t_amb']

        self.length = len(self.sol_irr)

        # the result dictionary
        self.res = {}

        # the system lifetime
        self.par['life_sys'] = 20.

        # initialize the operating hours of the electrolyzer array
        self.res['running_hours_pemel'] = 0.

        # initialize the storage tank size and its starting status
        self.m_h2_max = self.tank()
        self.m_h2_min = 0.05 * self.m_h2_max
        self.m_h2 = self.m_h2_min

        # instantiate the profiles for grid electricity price
        self.demand_profiles()

        # the number of PEM electrolyzer cells, corresponding to the
        # nominal capacity of the considered PEM cell and the provided
        # PEM capacity
        self.n_pemel_array = self.par['n_pemel'] / 0.4

        # the fitted polynomials on the electrolyzer and compressor
        self.polyfit_pemel()
        self.polyfit_pemel_compr()

    def demand_profiles(self):
        """
        Set the grid electricity price for buying and selling electricity.
        A contract with fixed electricity price is considered, for which the
        price for buying electricity consists of three segments: the energy
        price itself (i.e. 'elec cost'), the profit made on this price by the
        electricity provider (i.e. 'elec_cost_profit') and the fraction of the
        energy price to the final retail price (i.e. 'elec_cost_ratio', e.g.
        when this value equal 0.3, then the energy price corresponds to 30% of
        the final bill, while 70% corresponds to transportation cost,
        taxes,...). The price for selling electricity back to the grid
        corresponds to the energy price.

        In addition, the demand profiles from the hydrogen buses and diesel
        buses is determined, based on the European daily refueling profile [1].

        [1] T. A. Gunawan, I.Williamson, D. Raine, and R. F. Monaghan,
        “Decarbonising city bus networks in ireland with renewable hydrogen,”
        International Journal of Hydrogen Energy, 2020.

        """

        # electricity cost profile [euro/Wh]
        self.elec_profile = np.ones(self.length) * (
            (self.par['elec_cost'] +
             self.par['elec_cost_profit']) /
            self.par['elec_cost_ratio']) / 1e6

        # electricity selling profile [euro/Wh]
        self.elec_profile_sale = np.ones(
            self.length) * self.par['elec_cost'] / 1e6

        self.diesel_profile = np.ones(self.length) * self.par['diesel_cost']

        # number of km driven per day per bus
        self.par['n_km_bus'] = 250.

        # number of buses in the fleet
        self.par['n_bus'] = 40.

        # energy consumed by the diesel buses and hydrogen buses per day
        # [kWh/day]
        energy_h2 = (self.par['cons_h2_bus'] * self.par['n_km_bus'] *
                     self.par['n_h2_bus'])
        energy_diesel = (self.par['cons_diesel_bus'] * self.par['n_km_bus'] *
                         (self.par['n_bus'] - self.par['n_h2_bus']))

        h2_required = energy_h2 / 33.33  # kg
        diesel_required = energy_diesel / 10.  # litre

        # the daily refueling profile of the buses.
        fill_profile = np.array([0.09, 0.015, 0.005, 0.04, 0.04, 0., 0.01,
                                 0.01, 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                                 0.08, 0.08, 0.13, 0.13, 0.13, 0.13, 0.11])

        # daily refueling profile for the hydrogen buses and diesel buses
        day_h2 = np.ones(24) * fill_profile * h2_required
        day_diesel = np.ones(24) * fill_profile * diesel_required

        # annual refueling profiles
        self.load_h2 = list(day_h2) * int(365 * self.length / 8760)
        self.load_diesel = list(day_diesel) * int(365 * self.length / 8760)

        # dispenser capacity such that the hourly hydrogen demand can be
        # complied with
        dispenser_mass_flow_rate = 33.333
        self.par['n_disp'] = max(self.load_h2) / dispenser_mass_flow_rate

    #############################
    # photovoltaic array module #
    #############################

    def quantify_mpp(self, sol_irr, t_amb, pv_system):
        """

        Quantify the maximum power of the photovoltaic array
        for a given solar irradiance and ambient temperature.

        Parameters
        ----------
        sol_irr : float
            The solar irradiance [W/m2].
        t_amb : float
            The ambient temperature [C].
        pv_system : pandas.core.series.Series
            The pv system characteristics
        Returns
        -------
        pmp : float
            The maximum power.

        """

        # quantify the parameters for the pv system using De Soto method
        pv_inputs = pvlib.pvsystem.calcparams_desoto(sol_irr,
                                                     t_amb,
                                                     pv_system['alpha_sc'],
                                                     pv_system['a_ref'],
                                                     pv_system['I_L_ref'],
                                                     pv_system['I_o_ref'],
                                                     pv_system['R_sh_ref'],
                                                     pv_system['R_s'],
                                                     EgRef=1.121,
                                                     dEgdT=-0.0002677,
                                                     irrad_ref=1000.,
                                                     temp_ref=25.)

        # determine the maximum power for the given pv system
        pmp = pvlib.pvsystem.max_power_point(pv_inputs[0],
                                             pv_inputs[1],
                                             pv_inputs[2],
                                             pv_inputs[3],
                                             pv_inputs[4],
                                             method='newton')['p_mp']

        return pmp

    def photovoltaic(self):
        """

        The hourly photovoltaic power is quantified via the PVlib package.
        Using this package, first the characteristics for a typical
        photovoltaic panel are defined. Based on these characteristics,
        the maximum power point is quantified for each hour, based on the
        corresponding solar irradiance and ambient temperature. Finally, the
        hourly power production is scaled by the considered photovoltaic array
        capacity.

        """

        p_pv = np.zeros(len(self.sol_irr))

        # get the specific photovoltaic panel characteristics
        pv_database = pvlib.pvsystem.retrieve_sam('CECmod')
        pv_system = pv_database.SunPower_SPR_X19_240_BLK

        # determine the maximum power point at reference conditions
        p_mpp_ref = self.quantify_mpp(1000., 25., pv_system)  # W

        # maximum power point determination for each hour in the timeframe
        for i, irr in enumerate(self.sol_irr):
            if irr > 0.:
                p_mpp = self.quantify_mpp(irr, self.t_amb[i], pv_system)
                p_pv[i] = p_mpp / p_mpp_ref * self.par['n_pv'] * 1e3  # W
            else:
                p_pv[i] = 0.

        # store the hourly pv power in the result dictionary
        self.res['p_pv'] = p_pv

        # the dc-dc converter capacity in kW
        self.res['n_dcdc_pv'] = max(p_pv) / 1e3

    #############################
    # electrolyzer array module #
    #############################

    def pemel(self, i_pemel):
        """
        The electrolyzer model, based on the work of Saeed et al. [2]. For a
        given current, the model determines the operating voltage by
        considering the activation, concentration and ohmic overpotentials.
        The model quantifies the operating voltage, power, efficiency and
        hydrogen production.

        [2] Saeed, E. W., & Warkozek, E. G. (2015). Modeling and Analysis of
            Renewable PEM Fuel Cell System. Energy Procedia, 74, 87–101.
            https://doi.org/10.1016/j.egypro.2015.07.527

        Parameters
        ----------
        i_pemel : float
            The electrolyzer input current [A].

        Returns
        -------
        res : dict
            Dictionary with the operating conditions of the electrolyzer for a
            given current. It contains items on the operating voltage, power,
            efficiency and hydrogen mass flow rate.

        """
        par_pemel = {'T': 353.,
                     'a': 1.,
                     'p_o2': 1.,
                     'p_h2': 1.,
                     'p_h2o': 1.,
                     'i_L': 2.,
                     'A': 100.,
                     'i_0': 1e-4,
                     'n': 2.,
                     't_mem': 50e-4,
                     'alpha': 0.3,
                     'R': 8.3143,
                     'F': 96485.,
                     'HHV': 141.7e6,
                     }

        res = {}
        i = i_pemel / par_pemel['A']

        # minimum operating voltage of electrolyzer
        e_0 = (1.48 - 0.85e-3 * (par_pemel['T'] - 298.15) + 4.3085e-5 *
               par_pemel['T'] * np.log(par_pemel['p_h2'] *
                                       np.sqrt(par_pemel['p_o2']) /
                                       par_pemel['p_h2o']))

        # activation overpotential
        v_act = (np.log(i / par_pemel['i_0']) /
                 (par_pemel['alpha'] * par_pemel['n'] * par_pemel['F']) *
                 par_pemel['R'] * par_pemel['T'])

        # ohmic overpotential
        lambda_mem = (0.043 + 17.81 * par_pemel['a'] -
                      39.85 * par_pemel['a']**2. +
                      36. * par_pemel['a']**3.)
        sigma_mem = ((0.005139 * lambda_mem - 0.00326) *
                     np.exp(1268 * (1. / 303. - 1. / par_pemel['T'])))
        v_ohm = i * par_pemel['t_mem'] / sigma_mem

        # the concentration overpotential
        v_con = - (par_pemel['R'] * par_pemel['T'] /
                   (par_pemel['n'] * par_pemel['F']) *
                   np.log(1. - i / par_pemel['i_L']))

        # model outputs
        res['v_pemel'] = (e_0 + v_act + v_ohm + v_con) * self.n_pemel_array
        res['m_pemel'] = self.current_to_mh2(i_pemel) * self.n_pemel_array
        res['p_pemel'] = i_pemel * res['v_pemel']
        res['eff_pemel'] = (res['m_pemel'] * par_pemel['HHV'] /
                            (res['p_pemel'] * 3600.))
        return res

    def current_to_mh2(self, current):
        """
        When current is provided, this function determines the
        corresponding hydrogen mass flow rate per hour.

        Parameters
        ----------
        current : float
            The electrolyzer input current [A].

        Returns
        -------
        m_h2 : float
            The produced hydrogen mass flow rate [kg/h].

        """
        far_cons = 96485.
        m_h2 = current / (2. * far_cons) * 2.02e-3 * 3600.

        return m_h2

    def mh2_to_power(self, m_h2):
        """
        When the hydrogen mass flow rate is provided, this function determines
        the corresponding required power per hour.

        Parameters
        ----------
        m_h2 : float
            The produced hydrogen mass flow rate [kg/h].

        Returns
        -------
        power : float
            The required power to produce the hydrogen [W].

        """

        far_cons = 96485.
        current = m_h2 * (2. * far_cons) / (2.02e-3 * 3600. *
                                            self.n_pemel_array)
        power = self.pemel(current)['p_pemel']

        return power

    def polyfit_pemel(self):
        """
        The electrolyzer stack is evaluated over a range of input currents.
        Following these evaluations, a polynomial is fitted on the
        power - current relation of the electrolyzer. This polynomial enables
        to rapidly determine the input current when a certain amount of power
        is available. Since this relation is fairly linear, the polynomial
        should reach good agreement with the actual power - current relation,
        while maintaining the level of fidelity of the actual model.

        """

        # evaluate the electrolyzer stack for a set of currents
        i_list = np.arange(start=3, stop=200, step=4)
        p_pemel = np.zeros(len(i_list))
        for index, i in enumerate(i_list):
            res = self.pemel(i)
            p_pemel[index] = res['p_pemel']

        # generate a polynomial fitted on the power - current points
        self.p_to_i_pemel = polyfit_func(p_pemel, i_list)

    #####################
    # compressor module #
    #####################

    def compressor(self, m_h2):
        """
        The compressor module defined the required compression power to
        compress the hydrogen mass flow rate [3].

        [3] Zhao, L., Brouwer, J., & Samuelsen, S. (2014). Dynamic analysis of
        a self-sustainable renewable hydrogen fueling station. ASME 2014 12th
        International Conference on Fuel Cell Science, Engineering and
        Technology, FUELCELL 2014 Collocated with the ASME 2014 8th
        International Conference on Energy Sustainability.
        https://doi.org/10.1115/FuelCell2014-6330

        Parameters
        ----------
        m_h2 : float
            Hydrogen mass flow rate [kg/h].

        Returns
        -------
        power : float
            The required compression power [W].

        """

        # convert the flow rate into kg/s
        m_h2 *= 1. / 3600.

        par_c = {
            'T_in': 353.,
            'p_in': 1.,
            'p_out': 350.,
            'eta_c': 0.85,
            'R': 4.124,
            'n': 1.609,
        }

        power = (m_h2 *
                 par_c['n'] *
                 par_c['R'] *
                 par_c['T_in'] *
                 ((par_c['p_out'] /
                   par_c['p_in'])**((par_c['n'] -
                                     1.) /
                                    par_c['n']) -
                     1.) *
                 1000. /
                 (par_c['eta_c'] *
                     (par_c['n'] -
                      1.)))

        return power

    def polyfit_pemel_compr(self):
        """
        The power consumption by the electrolyzer stack and compressor are
        evaluated over a range of hydrogen mass flow rates. Following these
        evaluations, a polynomial is fitted on the mass flow rate - power
        relation. This polynomial enables to rapidly determine the input
        mass flow rate when a certain amount of power is available.
        Since this relation is fairly linear, the polynomial should reach good
        agreement with the actual mass flow rate - power relation, while
        maintaining the level of fidelity of the actual model.

        """

        # the electrolyzer array operating limits
        pemel_lower_lim = self.par['n_pemel'] * 10.
        pemel_upper_lim = self.par['n_pemel'] * 1e3

        # the operating current at these limits
        current_ub = self.p_to_i_pemel(pemel_upper_lim)
        current_lb = self.p_to_i_pemel(pemel_lower_lim)

        # the characteristics of the electrolyzer at the limits
        pemel_ub = self.pemel(current_ub)
        pemel_lb = self.pemel(current_lb)

        # the compression power at the hydrogen production limits
        p_compr_ub = self.compressor(pemel_ub['m_pemel'])
        p_compr_lb = self.compressor(pemel_lb['m_pemel'])

        # the compression and electrolyzer power at the limits
        pemel_compr_ub = pemel_ub['p_pemel'] + p_compr_ub
        pemel_compr_lb = pemel_lb['p_pemel'] + p_compr_lb

        # the operating bounds for the compressor and electrolyzer
        self.bounds = {'pemel_lb': pemel_lb,
                       'pemel_ub': pemel_ub,
                       'pemel_compr_lb': pemel_compr_lb,
                       'pemel_compr_ub': pemel_compr_ub
                       }

        # evaluate the electrolyzer and compressor for a set of mass flow rates
        step = pemel_ub['m_pemel'] / 50.
        m_h2_list = np.arange(start=step, stop=pemel_ub['m_pemel'] - step,
                              step=step)
        p_pemel_compr = np.zeros(len(m_h2_list))
        for index, m_h2 in enumerate(m_h2_list):
            p_pemel = self.mh2_to_power(m_h2)
            p_compr = self.compressor(m_h2)
            p_pemel_compr[index] = p_pemel + p_compr

        # generate a polynomial fitted on the power - mass flow rate points
        self.p_to_m_pemel_comp = polyfit_func(p_pemel_compr, m_h2_list)

    #############################
    # tank and dispenser module #
    #############################

    def tank(self):
        """

        The maximum storage capacity of the hydrogen tank.

        Returns
        -------
        m_max : float
            The hydrogen storage capacity. [kg]

        """

        # conversion from energy (kWh) into mass (kg)
        m_max = self.par['n_tank'] / 33.33

        return m_max

    def dispenser(self, m_h2):
        """
        The cooling power is determined before the dispensation of the hydrogen
        into the bus fuel tank [4].

        [4] A. Elgowainy, K. Reddi, D. Y. Lee, N. Rustagi, and E. Gupta,
        “Techno-economic and thermodynamic analysis of pre-cooling systems at
        gaseous hydrogen refueling stations,” Int. J. Hydrogen Energy, vol. 42,
        no. 49, pp. 29067–29079, Dec. 2017.

        Parameters
        ----------
        m_h2 : float
            Hydrogen mass flow rate [kg/h].

        Returns
        -------
        power : float
            The required cooling power for dispensation [W].

        """

        demand_day_h2 = sum(self.load_h2[:24])

        if m_h2 < 1e-6 or demand_day_h2 < 1e-4:
            power = 0.
        else:
            t_amb = 10.4
            m_h2 *= 1. / 3600.
            ei_pcu = ((0.3 / (1.6 * np.exp(-0.018 * t_amb))) +
                      (25. * np.log(t_amb) - 21.) / demand_day_h2)  # kWhe/kgH2
            power = ei_pcu * m_h2 * 3.6e6  # W

        return power

    ##############################
    # management strategy module #
    ##############################

    def p_for_inst_demand(self, m_h2):
        """
        When a hydrogen demand is provided, this method determines the power
        to generate the hydrogen in the electrolyzer array and the power to
        compress the hydrogen in the compressor. The method indicates when the
        electrolyzer array capacity is not sufficient to produce the desired
        hydrogen mass flow rate.

        Parameters
        ----------
        m_h2 : float
            Hydrogen mass flow rate [kg/h].

        Returns
        -------
        p_pemel : float
            The required electrolyzer array power [W].
        p_compr : float
            The required compressor power [W].
        bool
            True when the electrolyzer capacity is insufficient to produce the
            hydrogen.

        """

        # check if the desired hydrogen can be produced by the PEM
        if m_h2 > self.bounds['pemel_ub']['m_pemel']:
            return 0., 0., True

        # determine the power needed from electrolyzer and compressor to
        # deliver the required hydrogen mass flow rate
        p_pemel = self.mh2_to_power(m_h2)
        p_compr = self.compressor(m_h2)
        self.res['running_hours_pemel'] += 1.
        return p_pemel, p_compr, False

    def prod_mh2(self, p_in):
        """
        When there is power available to produce hydrogen, this method
        distributes this power over the electrolyzer array and compressor,
        such that the hydrogen is produced and compressed with this given
        power. If the given power will lead to an excess of hydrogen (i.e
        the storage tank is full before the energy is fully consumed), the
        power supplied to the electrolyzer array and compressor is recalculated
        such that no excess hydrogen is produced.

        Parameters
        ----------
        p_in : float
            The power available to produce hydrogen [W].

        Returns
        -------
        p_pemel : float
            The power consumed by the electrolyzer array [W].
        p_compr : float
            The power consumed by the compressor [W].

        """

        no_run = False

        # produce H2 only if tank is not yet full
        # nothing happens when the desired power is under the lower limit
        if self.m_h2 >= self.m_h2_max or p_in < self.bounds['pemel_compr_lb']:
            p_pemel = 0.
            p_compr = 0.
            no_run = True  # the electrolyzer did not run

        # if the power is higher than the upper limit, work at the upper limit
        elif p_in > self.bounds['pemel_compr_ub']:

            # this is what can be produced at the upper limit for H2
            m_h2 = self.bounds['pemel_ub']['m_pemel']

            # produced hydrogen is added to the tank
            self.m_h2 += m_h2

        else:

            # quantify the hydrogen created with the available power
            m_h2 = self.p_to_m_pemel_comp(p_in)

            # produced hydrogen is added to the tank
            self.m_h2 += m_h2

        # if the new hydrogen capacity exceeds the maximum storage capacity
        if self.m_h2 > self.m_h2_max:

            # the hydrogen that is still allowed in the tank
            m_h2 -= (self.m_h2 - self.m_h2_max)

            # the power for the PEM to generate this hydrogen
            p_rev = self.mh2_to_power(m_h2)

            # check if this power is higher than the lower limit for the PEM
            if p_rev > self.bounds['pemel_lb']['p_pemel']:
                self.m_h2 = self.m_h2_max

            # the space left in the tank is too small
            else:
                self.m_h2 -= (m_h2 + self.m_h2 - self.m_h2_max)
                p_pemel = 0.
                p_compr = 0.
                no_run = True

        # if power is applied to the PEM and compessor
        if not no_run:
            p_compr = self.compressor(m_h2)
            p_pemel = self.mh2_to_power(m_h2)
            self.res['running_hours_pemel'] += 1.

        return p_pemel, p_compr

    def extract_h2_from_tank(self, demand):
        """
        Extract the hydrogen demand from the storage tank. If more hydrogen is
        demanded than available in the tank, extract only the available amount
        of hydrogen.

        Parameters
        ----------
        demand : float
            The hydrogen demand. [kg]

        Returns
        -------
        demand_left : float
            The amount of hydrogen demand that is not covered by the tank. [kg]

        """

        # extract the hydrogen from the tank
        self.m_h2 -= demand

        if self.m_h2 < self.m_h2_min:

            # if the demand was too high, set the tank to its minimum level and
            # define the demand that was not covered by the tank
            demand_left = self.m_h2_min - self.m_h2
            self.m_h2 = self.m_h2_min
        else:

            # if the tank covers the demand, there is no hydrogen demand left
            demand_left = 0.

        return demand_left

    ##########################################
    # model evaluation
    ##########################################

    def evaluation(self):
        """

        This is the main method of the Evaluation class.
        For each hour, the power management strategy is applied.
        The hydrogen demand is extracted from the hydrogen tank, when
        sufficient hydrogen is available in the tank. When the hydrogen in the
        tank does not comply with the hydrogen demand, the power to run the
        electrolyzer array and compressor is calculated to generate and
        compress the remaining hydrogen. To generate this power and the
        dispensation power, the photovoltaic power is called upon first. If
        necessary, the remaining power is covered by the grid.
        When the hydrogen tank does comply with the hydrogen demand, the
        photovoltaic power is used to cover the required dispensation power.
        When excess photovoltaic power is present, this power is used to
        generate and compress hydrogen in the electrolyzer array and
        compressor, respectively. Finally, the lifetime, cost and
        CO2-emission of the system are quantified.

        Returns
        -------
        bool
            True when the capacity of the electrolyzer array is sufficient to
            cover the instantaneous hydrogen demand during the year.

        """

        n_compr = np.zeros(self.length)
        n_disp = np.zeros(self.length)
        n_dcdc_pem = np.zeros(self.length)
        n_dcac = np.zeros(self.length)

        self.res['m_h2_array'] = np.ones(self.length)
        self.res['grid_e_buy'] = np.ones(self.length)
        self.res['grid_e_sold'] = np.ones(self.length)
        self.res['grid_co2'] = 0.

        # get the hourly photovoltaic array power
        self.photovoltaic()

        for t in range(self.length):
            e_grid_buy = 0.
            e_grid_sold = 0.

            # define the power for dispensation
            p_disp = self.dispenser(self.load_h2[t])
            n_disp[t] = p_disp

            # define if there is any H2 demand left after assessing the tank
            demand_left = self.extract_h2_from_tank(self.load_h2[t])
            if demand_left > 0.:

                # power needed by the PEM and compressor to generate the demand
                p_pemel, p_compr, check = self.p_for_inst_demand(demand_left)
                n_dcdc_pem[t] = p_pemel
                n_compr[t] = p_compr

                if check:
                    return False

                # positive if the PV energy can comply with the PEM power
                net_p_1 = self.res['p_pv'][t] - p_pemel

                if net_p_1 > 0.:
                    # if yes, check if remaining PV power can comply with the
                    # compressor and dispensation demand
                    net_p_2 = net_p_1 - (p_compr + p_disp)

                    if net_p_2 > 0.:
                        # when still excess power available, sell to the grid
                        e_grid_sold = net_p_2
                        n_dcac[t] = e_grid_sold

                    else:
                        # if the compressor and dispensation demand cannot be
                        # covered, buy remaining demand from the grid
                        # (compressor and dispensation are AC-powered)
                        e_grid_buy = abs(net_p_2)

                else:
                    # if PEM, compressor and dispensation cannot be covered by
                    # PV energy, buy required electricity from the grid
                    # only power for PEM has to pass through AC-DC conversion
                    p_to_buy = abs(net_p_1) + p_compr + p_disp
                    e_grid_buy = p_to_buy
                    n_dcac[t] = abs(net_p_1)

            else:  # excess PV energy available to generate hydrogen

                # use solar energy to power the dispensation
                net_p = self.res['p_pv'][t] - p_disp

                if net_p <= 0.:
                    # remaining power covered by the grid
                    e_grid_buy = abs(net_p)

                    # the solar power is all used for dispensation and thus
                    # sent through the DC-AC converter
                    n_dcac[t] = self.res['p_pv'][t]

                else:

                    # quantify how much of the remaining PV energy can be used
                    # by the PEM and compressor to generate hydrogen
                    p_pemel, p_compr = self.prod_mh2(net_p)
                    n_compr[t] = p_compr

                    # the remaining excess PV energy can be sold
                    e_grid_sold = net_p - p_pemel - p_compr
                    n_dcac[t] = p_disp + e_grid_sold + p_compr
                    n_dcdc_pem[t] = p_pemel

            # store evolution of storage tank, electricity bought/sold
            # and the CO2 amount from buying grid electricity
            self.res['m_h2_array'][t] = ((self.m_h2 - self.m_h2_min) /
                                         (self.m_h2_max - self.m_h2_min))
            self.res['grid_e_buy'][t] = e_grid_buy
            self.res['grid_e_sold'][t] = e_grid_sold
            self.res['grid_co2'] += e_grid_buy * self.par['co2_elec']

        # define the capacity of the converters, compression and cooling
        self.res['n_compr'] = max(n_compr) / 1e3
        self.res['n_cooling'] = max(n_disp) / 1e3
        self.res['n_dcdc_pemel'] = max(n_dcdc_pem) / 1e3
        self.res['n_dcac'] = max(n_dcac) / 1e3

        # the cost of annual diesel consumption
        self.res['diesel_cost'] = self.load_diesel * self.diesel_profile

        # determine the lifetime of the electrolyzer
        self.lifetime()

        # determine the system cost
        self.cost()

        # determine the system annual CO2 emission
        self.lca()

        return True

    def lifetime(self):
        """

        The lifetime method determines the lifetime of
        the electrolyzer array, based on the number of
        operating hours during the evaluated year.

        """

        # lifetime of the electrolyzer array
        if self.res['running_hours_pemel'] == 0.:
            self.res['life_pemel'] = 1e8
        else:
            self.res['life_pemel'] = (self.par['life_pemel'] /
                                      self.res['running_hours_pemel'])

    def cost(self):
        """

        Based on the capital recovery factor, the CAPEX,
        OPEX and replacement cost of the system components,
        the levelized cost of mobility [euro/km] is determined. The formula
        for the annualized system cost is adopted from Coppitters et al. [5].

        [5] Coppitters, D., De Paepe, W., & Contino, F. (2020). Robust design
            optimization and stochastic performance analysis of a
            grid-connected photovoltaic system with battery storage and
            hydrogen storage. Energy, 213, 118798.
            https://doi.org/10.1016/j.energy.2020.118798
        """

        # the capital recovery factor
        inv_rate = ((self.par['int_rate'] - self.par['infl_rate']) /
                    (1. + self.par['infl_rate']))
        crf = (((1. + inv_rate)**self.par['life_sys'] - 1.) /
               (inv_rate * (1. + inv_rate)**self.par['life_sys']))**(-1)

        # annual cost of photovoltaic array and DC-DC converter
        pv_cost = self.par['n_pv'] * (crf * self.par['capex_pv'] +
                                      self.par['opex_pv'])
        pv_dcdc_cost = self.res['n_dcdc_pv'] * (self.par['capex_dcdc'] *
                                                (crf + self.par['opex_dcdc']))
        components_cost = pv_cost + pv_dcdc_cost

        # annual cost of electrolyzer array and DC-DC converter
        pemel_cost = self.par['n_pemel'] * (self.par['capex_pemel'] *
                                            (crf + self.par['opex_pemel']))
        pemel_dcdc_cost = (self.res['n_dcdc_pemel'] *
                           (self.par['capex_dcdc'] *
                            (crf + self.par['opex_dcdc'])))
        components_cost += pemel_cost + pemel_dcdc_cost

        # annual cost of hydrogen storage tank
        tank_cost = self.par['n_tank'] * (self.par['capex_tank'] *
                                          (crf + self.par['opex_tank']))
        components_cost += tank_cost

        # annual cost of compressor
        compressor_cost = (self.par['capex_compr'] *
                           (self.res['n_compr']**0.5861) *
                           (crf + self.par['opex_compr']))
        components_cost += compressor_cost

        # annual cost of dispenser
        dispenser_cost = self.par['n_disp'] * (self.par['capex_disp'] *
                                               (crf + self.par['opex_disp']))
        components_cost += dispenser_cost

        # annual cost of dispensation cooling
        cooling_cost = self.res['n_cooling'] * (self.par['capex_cool'] *
                                                (crf + self.par['opex_cool']))
        components_cost += cooling_cost

        # annual cost of DC-AC inverter
        dcac_cost = self.res['n_dcac'] * (self.par['capex_dcac'] *
                                          (crf + self.par['opex_dcac']))
        components_cost += dcac_cost

        # annual cost of buses
        h2_bus_cost = (self.par['capex_h2_bus'] * crf + self.par['n_km_bus'] *
                       365. * self.par['opex_h2_bus']) * self.par['n_h2_bus']
        diesel_bus_cost = (self.par['capex_diesel_bus'] * crf +
                           self.par['opex_diesel_bus'] * 365. *
                           self.par['n_km_bus']) * (self.par['n_bus'] -
                                                    self.par['n_h2_bus'])
        components_cost += h2_bus_cost + diesel_bus_cost

        # annual replacement cost of the electrolyzer and buses
        arc = crf * sum([(1. + inv_rate)**(-(i + 1.) *
                                           self.res['life_pemel']) *
                         self.par['n_pemel'] *
                         self.par['repl_pemel'] *
                         self.par['capex_pemel'] for
                         i in range(int(self.par['life_sys'] /
                                        self.res['life_pemel']))])
        arc += crf * sum([(1. + inv_rate)**(-(i + 1.) * 10.) *
                          self.par['n_h2_bus'] * self.par['repl_h2_bus']
                          for i in range(int(self.par['life_sys'] / 10.))])
        arc += crf * sum([(1. + inv_rate)**(-(i + 1.) * 10.) *
                          (self.par['n_bus'] - self.par['n_h2_bus']) *
                          self.par['repl_diesel_bus'] for i in
                          range(int(self.par['life_sys'] / 10.))])

        grid_e_cost = sum(self.res['grid_e_buy'] * self.elec_profile)
        grid_e_gain = sum(self.res['grid_e_sold'] * self.elec_profile_sale)

        # total annual cost
        cost = (arc + components_cost + grid_e_cost - grid_e_gain +
                sum(self.res['diesel_cost']))

        # levelized cost of mobility in euro/km
        self.res['lcom'] = cost / (self.par['n_km_bus'] * self.par['n_bus'] *
                                   365.)

    def lca(self):
        """
        The life cycle assessment is performed based on the CO2 emissions from
        constructing the system components and the emissions related to
        consuming grid electricity and diesel. The annual CO2-equivalent
        emissions of the system is divided by the annual distance covered by
        the bus fleet, resulting in the levelized cost of mobility.

        """

        # annual CO2 emission of photovoltaic array production
        pv_co2 = self.par['n_pv'] * self.par['co2_pv']
        pv_dcdc_co2 = self.res['n_dcdc_pv'] * self.par['co2_dcdc']
        comp_co2 = pv_co2 + pv_dcdc_co2

        # annual CO2 emission of electrolyzer array production
        pemel_co2 = (self.par['n_pemel'] * self.par['co2_pemel'] *
                     (1 + int(self.par['life_sys'] / self.res['life_pemel'])))
        pemel_dcdc_co2 = self.res['n_dcdc_pemel'] * self.par['co2_dcdc']
        comp_co2 += pemel_co2 + pemel_dcdc_co2

        # annual CO2 emission of hydrogen storage tank production
        tank_co2 = self.par['n_tank'] / 33.33 * 16. * self.par['co2_tank']
        comp_co2 += tank_co2

        # annual CO2 emission of compressor production
        compressor_co2 = self.res['n_compr'] * self.par['co2_compr']
        comp_co2 += compressor_co2

        # annual CO2 emission of dispenser production
        dispenser_co2 = self.par['n_disp'] * self.par['co2_disp']
        comp_co2 += dispenser_co2

        # annual CO2 emission of cooling unit production
        cooling_co2 = self.res['n_cooling'] * self.par['co2_cool']
        comp_co2 += cooling_co2

        # annual CO2 emission of DC-AC inverter production
        dcac_co2 = self.res['n_dcac'] * self.par['co2_dcac']
        comp_co2 += dcac_co2

        # annual CO2 emission of diesel and hydrogen engine production
        diesel_engine_co2 = (self.par['co2_diesel_engine'] * 200. *  # 200 kW
                             (self.par['n_bus'] - self.par['n_h2_bus']))
        h2_engine_co2 = self.par['co2_fc_engine'] * 200. * self.par['n_h2_bus']
        comp_co2 += ((diesel_engine_co2 + h2_engine_co2) *
                     (1 + int(self.par['life_sys'] / 10.)))  # 10y lifetime

        # annual CO2 emission of diesel consumption
        diesel_co2 = sum(self.load_diesel) * self.par['co2_diesel']

        # annual CO2 emission of system
        co2 = (comp_co2 / self.par['life_sys'] + self.res['grid_co2'] +
               diesel_co2)

        # CO2 emitted per km driven by the fleet
        self.res['lco2'] = co2 / (self.par['n_km_bus'] * self.par['n_bus'] *
                                  365. * self.length / 8760.)

    def print_results(self, succes=True):
        """

        This method prints the levelized cost of electricity,
        the self-sufficiency ratio and the annual energy produced
        by the photovoltaic array.

        """
        if not succes:
            print("""Evaluation failed: the electrolyzer array capacity
                     of %f kW was not sufficient to cover the instantaneous
                     hydrogen demand.""" % self.par['n_pemel'])
        else:
            print('outputs:')
            print('LCOE:'.ljust(30) + '%.5f euro/km' % self.res['lcom'])
            print('LCO2:'.ljust(30) + '%.5f kg co2-eq/km' % self.res['lco2'])
            print(
                'PV electricity generated:'.ljust(30) +
                '%.5f MWh' %
                (sum(self.res['p_pv']) / 1e6))
            print('grid energy bought:'.ljust(30) + '%.5f MWh' %
                  (sum(self.res['grid_e_buy']) / 1e6))
            print('grid energy sold:'.ljust(30) + '%.5f MWh' %
                  (sum(self.res['grid_e_sold']) / 1e6))
            print(
                'compressor capacity:'.ljust(30) +
                '%.5f kW' %
                self.res['n_compr'])
            print(
                'cooling capacity:'.ljust(30) +
                '%.5f kW' %
                self.res['n_cooling'])
            print('life electrolyzer:'.ljust(30) + '%.5f year' %
                  self.res['life_pemel'])

            plt.plot(self.res['m_h2_array'])
            plt.show(block=False)

def polyfit_func(x_in, y_in, threshold=0.99999999):
    """
    The function fits a polynomial to the points of x_in and y_in. The
    polynomial starts with order 1. To evaluate its performance, the R-squared
    performance indicator is quantified. If the value for R-squared does
    not reach the defined threshold, the polynomial order is increased and
    the polynomial is fitted again on the points, until the threshold is
    satisfied. Once satisfied, the function returns the polynomial.

    Parameters
    ----------
    x_in : ndarray
        The x-coordinates for the sample points.
    y_in : ndarray
        The y-coordinates for the sample points.
    threshold : float, optional
        The threshold for the R-squared parameter. The default is 0.99999.

    Returns
    -------
    poly_func : numpy.poly1d
        A one-dimensional polynomial.

    """
    order = 0
    r_squared = 0.
    while r_squared < threshold:
        order += 1

        # the polynomial
        poly_coeff = np.polyfit(x_in, y_in, order)
        poly_func = np.poly1d(poly_coeff)

        # r-squared
        yhat = poly_func(x_in)
        ybar = np.sum(y_in) / len(y_in)
        ssreg = np.sum((yhat - ybar)**2.)
        sstot = np.sum((y_in - ybar)**2.)
        r_squared = ssreg / sstot

    return poly_func


"""
This module contains the class RacosClassification, which provides a classifier for all Racos algorithms.

Author:
    Yu-Ren Liu

Updated by:
    Ze-Wen Li
"""

from zoopt.dimension import Dimension, Dimension2
from zoopt.utils.tool_function import ToolFunction
import copy
import numpy as np
from zoopt.utils.zoo_global import gl

class RacosClassification:
    """
    This class implements a classifier used by all Racos algorithms.
    """

    def __init__(self, dim, positive, negative, ub=1):
        """
        Initialization

        :param dim: a Dimension object
        :param positive: positive population
        :param negative: negative population
        :param ub: uncertain bits, which is a parameter for Racos
        """
        self.__solution_space = dim
        self.__sample_region = []
        self.__label_index = []
        # Solution
        self.__positive_solution = positive
        self.__negative_solution = negative
        self.__x_positive = None
        self.__uncertain_bit = ub

        regions = dim.get_regions()
        for i in range(dim.get_size()):
            temp = [regions[i][0], regions[i][1]]
            self.__sample_region.append(temp)
        return

    def reset_classifier(self):
        """
        Reset this classifier.

        :return: no return value
        """
        regions = self.__solution_space.get_regions()
        for i in range(self.__solution_space.get_size()):
            self.__sample_region[i][0] = regions[i][0]
            self.__sample_region[i][1] = regions[i][1]
            self.__label_index = []
        self.__x_positive = None
        return

    # This algos always works, whether discrete or continuous, we always use this function.
    def mixed_classification(self):
        """
        The process to train this classifier, which can handle mixed search space(continuous and discrete).

        :return: no return value
        """
        if type(self.__solution_space) == Dimension:
            self.__x_positive = self.__positive_solution[np.random.randint(
                0, len(self.__positive_solution))]
            len_negative = len(self.__negative_solution)
            index_set = list(range(self.__solution_space.get_size()))
            remain_index_set = list(range(self.__solution_space.get_size()))
            types = self.__solution_space.get_types()
            order = self.__solution_space.get_order()
            while len_negative > 0:
                if len(remain_index_set) == 0:
                    ToolFunction.log('ERROR: sampled two same solutions, please raise issues on github')
                k = remain_index_set[np.random.randint(0, len(remain_index_set))]
                x_pos_k = self.__x_positive.get_x_index(k)
                # continuous
                if types[k] is True:
                    x_negative = self.__negative_solution[
                        np.random.randint(0, len_negative)]
                    x_neg_k = x_negative.get_x_index(k)
                    if x_pos_k < x_neg_k:
                        r = np.random.uniform(x_pos_k, x_neg_k)
                        if r < self.__sample_region[k][1]:
                            self.__sample_region[k][1] = r
                            i = 0
                            while i < len_negative:
                                if self.__negative_solution[i].get_x_index(k) >= r:
                                    len_negative -= 1
                                    itemp = self.__negative_solution[i]
                                    self.__negative_solution[i] = self.__negative_solution[len_negative]
                                    self.__negative_solution[len_negative] = itemp
                                else:
                                    i += 1
                    else:
                        r = np.random.uniform(x_neg_k, x_pos_k)
                        if r > self.__sample_region[k][0]:
                            self.__sample_region[k][0] = r
                            i = 0
                            while i < len_negative:
                                if self.__negative_solution[i].get_x_index(k) <= r:
                                    len_negative -= 1
                                    itemp = self.__negative_solution[i]
                                    self.__negative_solution[i] = self.__negative_solution[len_negative]
                                    self.__negative_solution[len_negative] = itemp
                                else:
                                    i += 1
                # discrete
                else:
                    if order[k] is True:
                        x_negative = self.__negative_solution[
                            np.random.randint(0, len_negative)]
                        x_neg_k = x_negative.get_x_index(k)
                        if x_pos_k < x_neg_k:
                            # different from continuous version
                            r = np.random.randint(x_pos_k, x_neg_k)
                            if r < self.__sample_region[k][1]:
                                self.__sample_region[k][1] = r
                                i = 0
                                while i < len_negative:
                                    if self.__negative_solution[i].get_x_index(k) >= r:
                                        len_negative -= 1
                                        itemp = self.__negative_solution[i]
                                        self.__negative_solution[i] = self.__negative_solution[len_negative]
                                        self.__negative_solution[len_negative] = itemp
                                    else:
                                        i += 1
                        else:
                            r = np.random.randint(x_neg_k, x_pos_k + 1)
                            if r > self.__sample_region[k][0]:
                                self.__sample_region[k][0] = r
                                i = 0
                                while i < len_negative:
                                    if self.__negative_solution[i].get_x_index(k) <= r:
                                        len_negative -= 1
                                        itemp = self.__negative_solution[i]
                                        self.__negative_solution[i] = self.__negative_solution[len_negative]
                                        self.__negative_solution[len_negative] = itemp
                                    else:
                                        i += 1
                    else:
                        delete = 0
                        i = 0
                        while i < len_negative:
                            if self.__negative_solution[i].get_x_index(k) != x_pos_k:
                                len_negative -= 1
                                delete += 1
                                itemp = self.__negative_solution[i]
                                self.__negative_solution[i] = self.__negative_solution[len_negative]
                                self.__negative_solution[len_negative] = itemp
                            else:
                                i += 1
                        remain_index_set.remove(k)
                        if delete != 0:
                            index_set.remove(k)
                        if len(index_set) == 0:
                            index_set.append(k)
            self.set_uncertain_bit(index_set)
            return
        elif type(self.__solution_space) == Dimension2:
            self.__x_positive = self.__positive_solution[np.random.randint(0, len(self.__positive_solution))]
            len_negative = len(self.__negative_solution)
            index_set = list(range(self.__solution_space.get_size()))
            remain_index_set = list(range(self.__solution_space.get_size()))
            types = self.__solution_space.get_types()
            order_or_precision = self.__solution_space.get_order_or_precision()
            while len_negative > 0:
                if len(remain_index_set) == 0:
                    ToolFunction.log('ERROR: sampled two same solutions, please raise issues on github')
                k = remain_index_set[np.random.randint(0, len(remain_index_set))]
                x_pos_k = self.__x_positive.get_x_index(k)

                # continuous
                if types[k]:
                    x_negative = self.__negative_solution[np.random.randint(0, len_negative)]
                    x_neg_k = x_negative.get_x_index(k)

                    if x_pos_k < x_neg_k:
                        r = round(np.random.uniform(x_pos_k, x_neg_k), gl.float_precisions[k])
                        if r < self.__sample_region[k][1]:
                            self.__sample_region[k][1] = r
                            i = 0
                            while i < len_negative:
                                if self.__negative_solution[i].get_x_index(k) >= r:
                                    len_negative -= 1
                                    itemp = self.__negative_solution[i]
                                    self.__negative_solution[i] = self.__negative_solution[len_negative]
                                    self.__negative_solution[len_negative] = itemp
                                else:
                                    i += 1
                    else:
                        r = round(np.random.uniform(x_neg_k, x_pos_k), gl.float_precisions[k])
                        if r > self.__sample_region[k][0]:
                            self.__sample_region[k][0] = r
                            i = 0
                            while i < len_negative:
                                if self.__negative_solution[i].get_x_index(k) <= r:
                                    len_negative -= 1
                                    itemp = self.__negative_solution[i]
                                    self.__negative_solution[i] = self.__negative_solution[len_negative]
                                    self.__negative_solution[len_negative] = itemp
                                else:
                                    i += 1
                # discrete
                else:
                    if order_or_precision[k] is True:
                        x_negative = self.__negative_solution[np.random.randint(0, len_negative)]
                        x_neg_k = x_negative.get_x_index(k)
                        if x_pos_k < x_neg_k:
                            r = np.random.randint(x_pos_k, x_neg_k)
                            if r < self.__sample_region[k][1]:
                                self.__sample_region[k][1] = r
                                i = 0
                                while i < len_negative:
                                    if self.__negative_solution[i].get_x_index(k) >= r:
                                        len_negative -= 1
                                        itemp = self.__negative_solution[i]
                                        self.__negative_solution[i] = self.__negative_solution[len_negative]
                                        self.__negative_solution[len_negative] = itemp
                                    else:
                                        i += 1
                        else:
                            r = np.random.randint(x_neg_k, x_pos_k + 1)
                            if r > self.__sample_region[k][0]:
                                self.__sample_region[k][0] = r
                                i = 0
                                while i < len_negative:
                                    if self.__negative_solution[i].get_x_index(k) <= r:
                                        len_negative -= 1
                                        itemp = self.__negative_solution[i]
                                        self.__negative_solution[i] = self.__negative_solution[len_negative]
                                        self.__negative_solution[len_negative] = itemp
                                    else:
                                        i += 1
                    else:
                        delete = 0
                        i = 0
                        while i < len_negative:
                            if self.__negative_solution[i].get_x_index(k) != x_pos_k:
                                len_negative -= 1
                                delete += 1
                                itemp = self.__negative_solution[i]
                                self.__negative_solution[i] = self.__negative_solution[len_negative]
                                self.__negative_solution[len_negative] = itemp
                            else:
                                i += 1
                        remain_index_set.remove(k)
                        if delete != 0:
                            index_set.remove(k)
                        if len(index_set) == 0:
                            index_set.append(k)
            self.set_uncertain_bit(index_set)
            return

    def set_uncertain_bit(self, index_set):
        """
        Choose uncertain bits from iset.

        :param iset: index set
        :return: no return value
        """
        ub = min(self.__uncertain_bit, len(index_set))
        self.__label_index = np.random.choice(index_set, ub, replace=False)
        return

    def rand_sample(self):
        """
        Random sample from self.__solution_space.get_dim().

        :return: sampled x
        """
        if type(self.__solution_space) == Dimension:
            x = copy.deepcopy(self.__x_positive.get_x())
            for index in self.__label_index:
                if self.__solution_space.get_type(index) is True:
                    x[index] = np.random.uniform(self.__sample_region[index][0], self.__sample_region[index][1])
                else:
                    x[index] = np.random.randint(self.__sample_region[index][0], self.__sample_region[index][1] + 1)
            return x
        elif type(self.__solution_space) == Dimension2:
            x = copy.deepcopy(self.__x_positive.get_x())
            for index in self.__label_index:
                # continuous
                if self.__solution_space.get_type(index):
                    x[index] = round(np.random.uniform(self.__sample_region[index][0], self.__sample_region[index][1]),
                                     gl.float_precisions[index])
                # discrete
                else:
                    x[index] = np.random.randint(self.__sample_region[index][0], self.__sample_region[index][1] + 1)
            return x

    def get_sample_region(self):
        return self.__sample_region

    def get_sample_space(self):
        if type(self.__solution_space) == Dimension:
            size = self.__solution_space.get_size()
            regions = self.__sample_region
            types = self.__solution_space.get_types()
            return Dimension(size, regions, types)
        elif type(self.__solution_space) == Dimension2:
            types = self.__solution_space.get_types()
            regions = self.__sample_region
            order_or_precision = self.__solution_space.get_order_or_precision()
            dim_li = []
            for i in range(len(types)):
                dim_li.append((types[i], regions[i], order_or_precision[i]))
            return Dimension2(dim_li)
        else:
            ToolFunction.log('get sample space wrong')

    def get_positive_solution(self):
        return self.__positive_solution

    def get_negative_solution(self):
        return self.__negative_solution

    def get_x_positive(self):
        return self.__x_positive

    def get_label_index(self):
        return self.__label_index

    # for debugging
    def print_neg(self):
        """
        Print negative population.

        :return: no return value
        """
        ToolFunction.log('------print neg------')
        for x in self.__negative_solution:
            x.print_solution()

    def print_pos(self):
        """
        Print positive population.

        :return: no return value
        """

        ToolFunction.log('------print pos------')
        for x in self.__positive_solution:
            x.print_solution()

    def print_sample_region(self):
        """
        Print sample region.

        :return: no return value
        """
        ToolFunction.log('------print sample region------')
        ToolFunction.log(self.__sample_region)


from pdfminer.psparser import LIT, PSLiteral, PSStackParser, PSKeyword, PSEOF, keyword_name
from pdfminer.pdftypes import PDFObjRef, resolve1, dict_value, stream_value, list_value, PDFStream

from PIL import ImageCms
from io import BytesIO
import numpy as np
from itertools import product

class colorSpaces:
    @property
    def defaults(self):
        default_values = [
            (GrayColorSpace, LIT('DeviceGray'), LIT('G')),
            (RGBColorSpace, LIT('DeviceRGB'), LIT('RGB')),
            (CMYKColorSpace, LIT('DeviceCMYK'), LIT('CMYK')),
            (CalGrayColorSpace, LIT('CalGray')),
            (CalRGBColorSpace, LIT('CalRGB')),
            (LabColorSpace, LIT('Lab')),
            (ICCBasedColorSpace, LIT('ICCBased')),
            (IndexedColorSpace, LIT('Indexed')),
            (SeparationColorSpace, LIT('Separation')),
            # (DeviceNColorSpace, LIT('DeviceN')),
            (PatternColorSpace, LIT('Pattern')),
            (NColorSpace, LIT('DeviceN')),
        ]
        refs = {}
        for tpl in default_values:
            for i, x in enumerate(tpl):
                if i > 0:
                    refs[x] = tpl[0]
        return refs

    def parse(self, obj, args=[]):
        if isinstance(obj, PDFObjRef):
            obj = resolve1(obj)

        if isinstance(obj, PSLiteral):
            cs = self.defaults.get(obj)
            if not cs:
                return None
                # raise TypeError('unknown color space: %s' % obj.name)
            return cs(*args)

        if isinstance(obj, list):
            return self.parse(obj[0], args=obj[1:])

class ColorSpace:
    overprintMask = 0x0f
    pipe = lambda *val: val
    getGray = pipe
    getRGB = pipe
    getCMYK = pipe
    mapGray = pipe
    mapRGB = pipe
    mapCMYK = pipe

class GrayColorSpace(ColorSpace):
    mode = 'L'
    ncomps = 1

    def getRGB(self, gray):
        # [gray] · [1, 1, 1]
        r = g = b = gray
        return r, g, b

    def getCMYK(self, gray):
        # [gray] · [0, 0, 0, 1]
        c = m = y = 0
        k = gray
        return c, m, y, k

class CalGrayColorSpace(GrayColorSpace):
    whiteX = whiteY = whiteZ = 1
    blackX = blackY = blackZ = 0
    gamma = 1

    def __init__(self, obj):
        obj = resolve1(obj)
        params = dict_value(obj)
        self.whiteX, self.whiteY, self.whiteZ = params['WhitePoint']
        self.blackX, self.blackY, self.blackZ = params['BlackPoint']
        self.gamma = params['Gamma']

class RGBColorSpace(ColorSpace):
    mode = 'RGB'
    ncomps = 3

    def getGray(self, r, g, b):
        return 0.299 * r + 0.587 * g + 0.114 * b

    def getCMYK(self, r, g, b):
        c = 1 - r
        m = 1 - g
        y = 1 - b
        k = min(c, m, y)
        return c - k, m - k, y - k, k

    def mapGray(self, arr):
        return self.getGray(arr[..., 0], arr[..., 1], arr[..., 2])

    def mapCMYK(self, arr):
        k = arr.max(-1)
        out = np.empty_like(arr)
        out[..., 0] = k - arr[..., 0]
        out[..., 1] = k - arr[..., 1]
        out[..., 2] = k - arr[..., 2]
        k = k[..., np.newaxis]
        return np.concatenate((out, 255 - k), axis=-1)

class CalRGBColorSpace(RGBColorSpace):
    matrix = [
        1,	0,	0,
        0,	1,	0,
        0,	0,	1
    ]

    def __init__(self, obj):
        obj = resolve1(obj)
        params = dict_value(obj)
        self.whiteX, self.whiteY, self.whiteZ = params.get(
            'WhitePoint', (1, 1, 1))
        self.blackX, self.blackY, self.blackZ = params.get(
            'BlackPoint', (0, 0, 0))
        self.gammaR, self.gammaG, self.gammaB = params.get('Gamma', (1, 1, 1))
        self.matrix = params.get('Matrix', self.matrix)

class CMYKColorSpace(ColorSpace):
    mode = 'CMYK'
    ncomps = 4
    factors = [
        [1,		 1,	     1],
        [0.1373, 0.1216, 0.1255],
        [1,		 0.9490, 0],
        [0.1098, 0.1020, 0],
        [0.9255, 0,	     0.5490],
        [0.1412, 0,	     0],
        [0.9294, 0.1098, 0.1412],
        [0.1333, 0,	     0],
        [0,		 0.6784, 0.9373],
        [0,		 0.0588, 0.1412],
        [0,		 0.6510, 0.3137],
        [0,		 0.0745, 0],
        [0.1804, 0.1922, 0.5725],
        [0,		 0,		 0.0078],
        [0.2118, 0.2119, 0.2235],
        [0,		 0,	     0]
    ]

    def getGray(self, c, m, y, k):
        return 1 - k - 0.3 * c - 0.59 * m - 0.11 * y

    def getRGB(self, c, m, y, k, r=0, g=0, b=0):
        c1, m1, y1, k1 = 1-c, 1-m, 1-y, 1-k
        for i, (b0, b1, b2, b3) in enumerate(product([c1, c], [m1, m], [y1, y], [k1, k])):
            x = b0 * b1 * b2 * b3
            r += self.factors[i][0] * x
            g += self.factors[i][1] * x
            b += self.factors[i][2] * x
        return r, g, b

    def mapGray(self, arr):
        return 255 - arr[..., 3] - 0.3 * arr[..., 0] - 0.59 * arr[..., 1] - 0.11 * arr[..., 2]

    def mapRGB(self, arr):
        arr = arr.astype('float') / 255
        out = np.empty_like(arr[..., :-1], dtype='float')
        self.getRGB(*(arr[..., i] for i in range(4)),
                    *(out[..., i] for i in range(3)))
        return (out * 255).astype('uint8')

xyzrgb = [
    [3.240449,	-1.537136,	-0.498531],
    [-0.969265,	1.876011,	0.041556],
    [0.055643,	-0.204026,	1.057229]
]

class LabColorSpace(ColorSpace):
    mode = 'LAB'
    ncomps = 3

    def __init__(self, obj):
        obj = resolve1(obj)
        params = dict_value(obj)
        self.whiteX, self.whiteY, self.whiteZ = params.get(
            'WhitePoint', (1, 1, 1))
        self.blackX, self.blackY, self.blackZ = params.get(
            'BlackPoint', (0, 0, 0))
        self.aMin, self.bMin, self.aMax, self.bMax = params.get(
            'Range', (-100, -100, 100, 100))
        self.kr = 1 / (
            xyzrgb[0][0] * self.whiteX +
            xyzrgb[0][1] * self.whiteY +
            xyzrgb[0][2] * self.whiteZ
        )
        self.kg = 1 / (
            xyzrgb[1][0] * self.whiteX +
            xyzrgb[1][1] * self.whiteY +
            xyzrgb[1][2] * self.whiteZ
        )
        self.kb = 1 / (
            xyzrgb[2][0] * self.whiteX +
            xyzrgb[2][1] * self.whiteY +
            xyzrgb[2][2] * self.whiteZ
        )

    def getGray(self, l, a, b):
        r, g, b = self.getRGB(l, a, b)
        return 0.299 * r + 0.587 * g + 0.114 * b + 0.5

    def getRGB(self, l, a, b):
        def lab2xyz(t): return t ** 3 if (t >= 6 /
                                          29) else (108 / 841 * (t - 4 / 29))

        # convert L*a*b* to CIE 1931 XYZ color space
        t1 = (l + 16) / 116
        t2 = t1 + a / 500
        X = lab2xyz(t2)
        X *= self.whiteX
        Y = lab2xyz(t1)
        Y *= self.whiteY
        t2 = t1 - b / 200
        Z = lab2xyz(t2)
        Z *= self.whiteZ

        # convert XYZ to RGB, including gamut mapping and gamma correction
        r = xyzrgb[0][0] * X + xyzrgb[0][1] * Y + xyzrgb[0][2] * Z
        g = xyzrgb[1][0] * X + xyzrgb[1][1] * Y + xyzrgb[1][2] * Z
        b = xyzrgb[2][0] * X + xyzrgb[2][1] * Y + xyzrgb[2][2] * Z

        return r ** 0.5, g ** 0.5, b ** 0.5

    def getCMYK(self, l, a, b):
        r, g, b = self.getRGB(l, a, b)
        c = 1 - r
        m = 1 - g
        y = 1 - b
        k = min(c, m, y)
        return c - k, m - k, y - k, k

class ICCBasedColorSpace(ColorSpace):
    @property
    def defaults(self):
        return {
            'L': GrayColorSpace,
            'RGB': RGBColorSpace,
            'CMYK': CMYKColorSpace,
            'LAB': LabColorSpace
        }

    mode = 'RGB'

    def __init__(self, obj):
        obj = resolve1(obj)
        fp = BytesIO(obj.get_data())
        self.profile = ImageCms.ImageCmsProfile(fp)
        fp.close()
        self.mode = self.profile.profile.color_space
        if self.mode == 'LAB':
            alt = resolve1(obj['Alternate'])
            if isinstance(alt, list):
                alt = alt[1]
            self.base = self.defaults[self.mode](alt)
        else:
            self.base = self.defaults[self.mode]()
        self.ncomps = len(self.mode)

    def getGray(self, *val):
        return self.base.getGray(*val)

    def getRGB(self, *val):
        return self.base.getRGB(*val)

    def getCMYK(self, *val):
        return self.base.getCMYK(*val)

class IndexedColorSpace(ColorSpace):
    mode = 'P'
    basemode = 'RGB'
    palette = list(map(lambda i: (i, i, i), range(256)))
    ncomps = 1

    def __init__(self, base, hival, obj):
        cs = colorSpaces()
        self.base = cs.parse(resolve1(base))
        self.hival = int(resolve1(hival))

        obj = resolve1(obj)
        data = b''
        if isinstance(obj, bytes):
            data = obj
        elif isinstance(obj, PDFStream):
            data = obj.get_data()
        if data:
            n = self.base.ncomps
            self.palette = [[data[i * n + j] for j in range(n)] for i in range(len(data) // n)]

    def lookup(self, index):
        i = max(0, min(index, len(self.palette) - 1))
        return self.palette[i]

    def getGray(self, index):
        return self.base.getGray(*self.lookup(index))

    def getRGB(self, index):
        return self.base.getRGB(*self.lookup(index))

    def getCMYK(self, index):
        return self.base.getCMYK(*self.lookup(index))

    def mapPixels(self, arr):
        palette = np.array(self.palette, dtype='uint8')
        return palette[arr]

    def mapGray(self, arr):
        return self.base.mapGray(arr)

    def mapRGB(self, arr):
        return self.base.mapRGB(arr)

    def mapCMYK(self, arr):
        return self.base.mapCMYK(arr)

class functionParser:
    def _min(self, x, num):
        if isinstance(x, (int, float)):
            return min(x, num)
        x[x >= num] = num
        return x

    def _max(self, x, num):
        if isinstance(x, (int, float)):
            return max(x, num)
        x[x < num] = num
        return x

class SampledFunctionParser(functionParser):
    def __init__(self, spec, domain):
        self.domain = domain
        self.frange = list_value(spec['Range'])
        self.nins = len(self.domain) >> 1
        self.nouts = len(self.frange) >> 1

        self.sizes = list_value(spec['Size'])[:self.nins]
        self.bits = int(spec['BitsPerSample'])

        if 'Encode' in spec:
            self.encode = list_value(spec['Encode'])
        else:
            self.encode = [0] * (self.nins << 1)
            self.encode[1::2] = [size-1 for size in self.sizes]

        self.decode = list_value(
            spec['Decode']) if 'Decode' in spec else self.frange[:]

        # domain = [0 1]
        # range = [0 1 0 1 0 1 0 1]
        # bits = 8
        # sizes = [1024]
        # encode = [0 1023]
        # decode = [0 1 0 1 0 1 0 1]

    def interpolate(self, x, xmin, xmax, ymin, ymax):
        return (ymax - ymin) / (xmax-xmin) * (x-xmin) + ymin

    def parse(self, *args):
        e = []
        for i in range(self.nins):
            x = self._min(
                self._max(args[i], self.domain[i*2]), self.domain[i*2+1])
            x = self.interpolate(
                x, self.domain[i*2], self.domain[i*2+1], self.encode[i*2], self.encode[i*2+1])
            e.append(self._min(self._max(x, 0), self.sizes[i]-1))
        return e

def SampledFunction(spec, domain):
    parser = SampledFunctionParser(spec, domain)
    return parser.parse

class ExponentialFunctionParser(functionParser):
    def __init__(self, spec, domain):
        self.c0, self.c1 = [0], [1]
        if spec.get('C0'):
            self.c0 = [float(x) for x in list_value(spec['C0'])]
        if spec.get('C1'):
            self.c1 = [float(x) for x in list_value(spec['C1'])]
        self.n = spec['N']
        self.frange = None
        if spec.get('Range'):
            self.frange = list_value(spec.get('Range'))
        self.domain = domain

    def parse(self, ipt):
        ipt /= 255
        ipt = self._min(self._max(ipt, self.domain[0]), self.domain[1])
        opt = []
        for i in range(len(self.c0)):
            x = self.c0[i] + pow(ipt, self.n) * (self.c1[i] - self.c0[i])
            if self.frange:
                x = self._min(self._max(x, self.frange[0]), self.frange[1])
            opt.append(x * 255)
        return opt

def ExponentialFunction(spec, domain):
    parser = ExponentialFunctionParser(spec, domain)
    return parser.parse

def StitchingFunction(spec, domain):
    pass

class PSFunctionParser(PSStackParser):
    def __init__(self, fp):
        super().__init__(fp)
        self.run()

    def run(self):
        try:
            self.nextobject()
        except PSEOF:
            pass
        _, self.argstack = self.curstack.pop()
        self.reset()

    def parse(self, *args):
        argstack = list(args) + self.argstack
        self.curstack = []
        while argstack:
            obj = argstack.pop(0)
            if isinstance(obj, PSKeyword):
                name = keyword_name(obj)
                if not isinstance(name, str):
                    name = name.decode()
                result = getattr(self, 'do_'+name)()
                if result is not None:
                    if isinstance(result, (list, tuple)):
                        self.curstack += list(result)
                    else:
                        self.curstack.append(result)
            else:
                self.curstack.append(obj)
        return self.curstack

    def do_keyword(self, pos, token):
        self.push((pos, token))

    def do_roll(self):
        n, j = self.pop(2)
        vals = self.pop(n)
        j %= n
        if not j:
            return vals
        return (vals*2)[n-j:n*2-j]

    def do_dup(self):
        x = self.pop(1)
        return x + x

    def do_exch(self):
        a, b = self.pop(2)
        return b, a

    def do_sub(self):
        a, b = self.pop(2)
        if isinstance(b, (int, float)):
            return b - a
        b[b < a] = 0
        b[b >= a] -= a
        return b

    def do_pop(self):
        self.pop(1)

    def do_index(self):
        i = self.pop(1)[0]
        return self.curstack[-i-1]

    def do_cvr(self):
        num = self.pop(1)[0]
        return float(num)

    def do_mul(self):
        a, b = self.pop(2)
        return a * b

def PostScriptFunction(spec, domain):
    parser = PSFunctionParser(BytesIO(spec.get_data()))
    return parser.parse

def func_parse(spec):
    func_type = int(spec.get('FunctionType'))
    domain = list_value(spec.get('Domain'))

    func_refs = {
        0: SampledFunction,
        2: ExponentialFunction,
        3: StitchingFunction,
        4: PostScriptFunction
    }

    func_builder = func_refs[func_type]
    return func_builder(spec, domain)

class SeparationColorSpace(ColorSpace):
    mode = 'P'

    def __init__(self, alt, base, func, *args):
        cs = colorSpaces()
        self.base = cs.parse(resolve1(base))
        spec = resolve1(func)
        self.ncomps = len(spec['Domain']) >> 1
        self.func = func_parse(spec)

    def transform(self, *val):
        transformed = self.func(*val)
        new_val = []
        for i in range(self.base.ncomps):
            new_val.append(transformed[i])
        return new_val

    def mapPixels(self, arr):
        if not self.func:
            return arr
        if len(arr.shape) == 2:
            arr = arr[..., np.newaxis]
        w, h, d = arr.shape
        arr = arr.astype('float')
        transformed = self.transform(*[arr[..., i] for i in range(d)])
        result = None
        for layer in transformed:
            if isinstance(layer, (int, float)):
                layer = np.ones((w, h), dtype='float') * layer
            layer = layer.astype('uint8')
            if result is None:
                result = layer
            else:
                result = np.dstack([result, layer])
        return result

    def getGray(self, *val):
        val = self.transform(*val)
        return self.base.getGray(*val)

    def getRGB(self, *val):
        val = self.transform(*val)
        return self.base.getRGB(*val)

    def getCMYK(self, *val):
        val = self.transform(*val)
        return self.base.getCMYK(*val)

    def mapGray(self, arr):
        return self.base.mapGray(arr)

    def mapRGB(self, arr):
        return self.base.mapRGB(arr)

    def mapCMYK(self, arr):
        return self.base.mapCMYK(arr)

class NColorSpace(SeparationColorSpace):
    mode = 'P'

    def __init__(self, names, alt, func, *attrs):
        self.names = list_value(names)
        self.base = colorSpaces().parse(resolve1(alt))
        spec = resolve1(func)
        self.ncomps = len(spec['Domain']) >> 1
        self.func = func_parse(spec)

class PatternColorSpace(ColorSpace):
    under = None
    mode = 'P'
    ncomps = 1

    def __init__(self, *args):
        if args:
            cs = colorSpaces()
            self.under = cs.parse(resolve1(args[0]))

defaults = colorSpaces().defaults
parse = colorSpaces().parse


#!/usr/bin/env python3
from __future__ import print_function
import sys
sys.path.append('../lib/')
import os
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import pints
import pints.io
import pints.plot

# Model etc.
import model as m
import parameters
from parameters import simvc_fix, get_qc
from parameters import simvc_typical_values
import parametertransform
from priors import IKrWithoutConductanceLogPrior
from priors import VoltageOffsetWithConductanceLogPrior
from protocols import leak_staircase as protocol_def

"""
Run fit for mutliple cell experiment data with same kinetics:

Optimise a likelihood
    L(\theta) = \prod_i L_{kinetics, i}(\theta | \phi^*_i)
where
    i, measurement index
    \theta, the parameters of the kinetics
    L_{kinetics, i}, the likelihood of the kinetic parameters for measurement i
and
    \phi^*_i = \argmax_{\phi_i} L_{individual, i}(\phi_i | \theta)
where
    \phi_i, the parameters for individual measurement, such as g_{Kr} and
        voltage artefact parameters
    L_{individual, i}, the likelihood of \phi_i given \theta.

This returns us \theta* (= \argmax_\theta L(\theta)) and {\phi^*_i}.
"""

try:
    file_name = sys.argv[1]
    run_id = int(sys.argv[2])
except:
    print('Usage: python %s [str:file_name]' % os.path.basename(__file__)
            + '[int:run_id]')
    sys.exit()

common_param = parameters.ikr[1:]
independent_param = ['ikr.g',
    #'voltageclamp.rseries',
    'voltageclamp.voffset_eff',
    'voltageclamp.gLeak']
n_common_param = len(common_param)
n_independent_param = len(independent_param)

def get_param(common, independent):
    model_p = np.append(independent[0], common)
    vclamp_p = independent[1:]
    return np.append(model_p, vclamp_p)

def get_fix_param(var, val):
    """
        var: variable name.
        val: variable value to fix.
    """
    out = {}
    for i, j in zip(var, val):
        out[i] = j
    return out

def update_logposterior_fix_param(l, p):
    """
        l: log-posterior.
        p: dict, fix params.

        NOTE: Only works for specific implementation.
    """
    l._log_likelihood._problem._model.set_fix_parameters(p)

file_dir = '../data'
file_list = [
        'herg25oc1',
        'herg27oc1',
        'herg27oc2',
        'herg30oc1',
        'herg30oc2',
        'herg33oc1',
        'herg33oc2',
        'herg37oc3',
        ]
temperatures = [25.0, 27.0, 27.0, 30.0, 30.0, 33.0, 33.0, 37.0]
useFilterCap = False

if file_name not in file_list:
    raise ValueError('Input `file_name` must be in `file_list`')
temperature = temperatures[file_list.index(file_name)]

savedir = './out/' + file_name + '-scheme3-simvclinleak'
if not os.path.isdir(savedir):
    os.makedirs(savedir)

print('Temperature: ', temperature)

# Control fitting seed --> OR DONT
np.random.seed(run_id)
fit_seed = np.random.randint(0, 2**30)
# fit_seed = 542811797
print('Fit seed: ', fit_seed)
np.random.seed(fit_seed)

# Leak param
leakbeforeparam = np.loadtxt('../qc/' + file_name + '-staircaseramp-leak_before.txt')
cell_id_file = '../qc/%s-staircaseramp-cell_id.txt' % file_name
cell_ids = []
with open(cell_id_file, 'r') as f:
    for l in f:
        if not l.startswith('#'):
            cell_ids.append(l.split()[0])

# Get cells
selectedfile = '../manualselection/manualv2selected-%s.txt' % (file_name)
selectedwell = []
with open(selectedfile, 'r') as f:
    for l in f:
        if not l.startswith('#'):
            selectedwell.append(l.split()[0])
selectedwell = selectedwell[:]
n_cells = len(selectedwell)
print(file_name + ' selected ' + str(n_cells) + ' cells')

# Set parameter transformation
transform_to_ikr = parametertransform.log_transform_to_ikr
transform_from_ikr = parametertransform.log_transform_from_ikr
transform_to_vc = parametertransform.log_transform_to_vc
transform_from_vc = parametertransform.log_transform_from_vc
transform_to_leak = parametertransform.donothing
transform_from_leak = parametertransform.donothing

transform_to_common_param = transform_to_ikr
transform_from_common_param = transform_from_ikr

n_gvc_parameters = VoltageOffsetWithConductanceLogPrior(None,
        None).n_parameters()
transform_to_gvc = parametertransform.ComposeTransformation(
        np.exp, transform_to_vc, 1)
transform_from_gvc = parametertransform.ComposeTransformation(
        np.log, transform_from_vc, 1)
transform_to_independent_param = parametertransform.ComposeTransformation(
        transform_to_gvc, transform_to_leak, n_gvc_parameters)
transform_from_independent_param = parametertransform.ComposeTransformation(
        transform_from_gvc, transform_from_leak, n_gvc_parameters)

# For individual bit in MwG
# Model
model_common = m.Model('../mmt-model-files/simplified-voltage-clamp-ikr-linleak.mmt',
                protocol_def=protocol_def,
                temperature=273.15 + temperature,  # K
                transform=transform_to_common_param,
                useFilterCap=useFilterCap)  # ignore capacitive spike
model_common.set_parameters(common_param)

model_independent = m.Model('../mmt-model-files/simplified-voltage-clamp-ikr-linleak.mmt',
                protocol_def=protocol_def,
                temperature=273.15 + temperature,  # K
                transform=transform_to_independent_param,
                useFilterCap=useFilterCap)  # ignore capacitive spike
model_independent.set_parameters(independent_param)

# Load all data
log_posteriors_common = []
log_posteriors_independent = []
data_all = []
times_all = []
noise_sigma_all = []
x0_common_all = []
x0_independent_all = []
transform_x0_common_all = []
transform_x0_independent_all = []
for cell in selectedwell:
    # Load data
    data_file_name = file_name + '-staircaseramp-' + cell + '.csv'
    time_file_name = file_name + '-staircaseramp-times.csv'
    print('Fitting to ', data_file_name)

    data = np.loadtxt(file_dir + '/' + data_file_name,
                      delimiter=',', skiprows=1) # headers
    times = np.loadtxt(file_dir + '/' + time_file_name,
                       delimiter=',', skiprows=1) # headers
    times = times * 1e3  # s -> ms
    noise_sigma = np.std(data[:500])
    print('Estimated noise level: ', noise_sigma)

    if useFilterCap:
        # Apply capacitance filter to data
        data = data * model_common.cap_filter(times)

    data_all.append(data)
    times_all.append(times)
    noise_sigma_all.append(noise_sigma)

    # Get independent CMA-ES results as x0
    cma_file = './fit-results/%s/%s-staircaseramp-%s-solution-542811797.txt' \
            % (file_name, file_name, cell)
    x0i_param = np.loadtxt(cma_file)
    # Change unit... (note, conductance unit different to Beattie's ones)
    x0i_param = x0i_param * 1e-3  # in mV, ms
    x0_common = x0i_param[1:n_common_param + 1]  # remove conductance
    # Get 'independent parameters'
    cma_ip_file = './out/%s-fixkinetics-simvclinleak/%s-staircaseramp-%s-solution-542811797.txt' \
            % (file_name, file_name, cell)
    x0_independent = np.loadtxt(cma_ip_file)
    # x0_independent = np.array([-1.5, 1e-2])
    tx0_common = transform_from_common_param(x0_common)
    # tx0_independent = transform_from_vc(x0i_param[n_common_param + 1:])
    tx0_independent = transform_from_independent_param(x0_independent)

    x0_common_all.append(x0_common)
    x0_independent_all.append(x0_independent)
    transform_x0_common_all.append(tx0_common)
    transform_x0_independent_all.append(tx0_independent)

    # Get fix parameters' value
    rseal, cm, rseries = get_qc('../qc', file_name, cell)
    print('Est. Rseal, Cm, Rseries:', rseal, cm, rseries, '(GOhm, pF, GOhm)')
    alpha = 0.8  # rseries %compensation
    cell_idx = cell_ids.index(cell)
    ga, Ea = leakbeforeparam[cell_idx]
    simvc_fix_values = [cm, rseries * alpha, rseries] + [ga, Ea]
    fix_p = get_fix_param(simvc_fix + ['voltageclamp.rseries', 'voltageclamp.gLeak_est', 'voltageclamp.ELeak_est'],
            simvc_fix_values)

    #
    # For individual bit in MwG
    #
    model_common.set_fix_parameters(fix_p)
    model_independent.set_fix_parameters(fix_p)

    problem_common = pints.SingleOutputProblem(model_common, times, data)
    loglikelihood_common = pints.GaussianKnownSigmaLogLikelihood(
            problem_common, noise_sigma)
    logprior_common = IKrWithoutConductanceLogPrior(transform_to_common_param,
            transform_from_common_param)
    log_posterior_common = pints.LogPosterior(loglikelihood_common,
            logprior_common)
    log_posteriors_common.append(log_posterior_common)

    problem_independent = pints.SingleOutputProblem(model_independent, times,
            data)
    loglikelihood_independent = pints.GaussianKnownSigmaLogLikelihood(
            problem_independent, noise_sigma)
    vc_prior = VoltageOffsetWithConductanceLogPrior(
            transform_to_gvc, transform_from_gvc)
    ileak_prior = pints.UniformLogPrior([-1e3], [1e3])
    prior_independent = pints.ComposedLogPrior(vc_prior, ileak_prior)
    log_posterior_independent = pints.LogPosterior(loglikelihood_independent,
            prior_independent)
    log_posteriors_independent.append(log_posterior_independent)

    print('Score at default parameters: ',
            log_posterior_common(transform_x0_common_all[-1]))

    fix_p = get_fix_param(independent_param, x0_independent_all[-1])
    update_logposterior_fix_param(log_posterior_common, fix_p)
    print('Score at updated parameters: ',
            log_posterior_common(transform_x0_common_all[-1]))

    print('Score at default parameters: ',
            log_posterior_independent(transform_x0_independent_all[-1]))

    fix_p = get_fix_param(common_param, x0_common_all[-1])
    update_logposterior_fix_param(log_posterior_independent, fix_p)
    print('Score at updated parameters: ',
            log_posterior_independent(transform_x0_independent_all[-1]))

    for _ in range(5):
        assert(log_posterior_common(transform_x0_common_all[-1]) ==\
                log_posterior_common(transform_x0_common_all[-1]))
        assert(log_posterior_independent(transform_x0_independent_all[-1]) ==\
                log_posterior_independent(transform_x0_independent_all[-1]))

data_all = np.asarray(data_all).T
times_all = np.asarray(times_all).T
times = times_all[:, 0]

transform_x0_common_all_mean = np.mean(transform_x0_common_all, axis=0)
x0_common_all_mean = np.copy(transform_to_common_param(transform_x0_common_all_mean))
total_log_posteriors_common = pints.SumOfIndependentLogPDFs(
        log_posteriors_common)

# quick sanity checks
assert(n_cells == len(log_posteriors_independent))
assert(n_cells == len(log_posteriors_common))

#
# Define the big loglikelihood
#
tip0 = np.copy(transform_x0_independent_all)
tcp0 = np.copy(transform_x0_common_all_mean)
n_opt_iterations_independent = 200

final_posterior = m.ParallelMultiLevelLogLikelihood(
            total_common_logposterior=total_log_posteriors_common,
            list_of_independent_logposteriors=log_posteriors_independent,
            common_param=common_param,
            independent_param=independent_param,
            transform_to_common_param=transform_to_common_param,
            transform_to_independent_param=transform_to_independent_param,
            transformed_independent_parameters0=tip0,
            n_opt=n_opt_iterations_independent,
            restart_fit=True,
            onlyNelderMead=True,
            n_workers=None)

print('Test total_log_posteriors_common:', total_log_posteriors_common(tcp0))
# This is important to run once...
# To set what's the prior lp within final_posterior object
print('Test final_posterior: ', final_posterior(tcp0))

#
# Run optimisation
#
N = 1

logposteriors = []
params_c = []
params_i = []
n_iter = 1000
epsilon = 0.5  # when to switch from CMA-ES to Nelder-Mead

for _ in range(N):

    x0 = tcp0

    if epsilon > 0:
        # Create optimiser part 1
        opt = pints.OptimisationController(
                final_posterior, x0,
                method=pints.CMAES)
        opt.set_max_iterations(int(epsilon * n_iter))
        opt.set_max_unchanged_iterations(iterations=50, threshold=1e4)
        opt.set_parallel(False)

        # Run optimisation
        try:
            with np.errstate(all='ignore'):
                # Tell numpy not to issue warnings
                p, s = opt.run()
        except ValueError:
            import traceback
            traceback.print_exc()

        del(opt)
        x0 = np.copy(p)

    if epsilon < 1:
        # Create optimiser part 2
        opt = pints.OptimisationController(
                final_posterior, x0,
                method=pints.NelderMead)
        #opt.set_max_iterations(int((1 - epsilon) * n_iter))
        opt.set_max_unchanged_iterations(iterations=50, threshold=1e3)
        opt.set_parallel(False)

        # Run optimisation
        try:
            with np.errstate(all='ignore'):
                # Tell numpy not to issue warnings
                p, s = opt.run()
        except ValueError:
            import traceback
            traceback.print_exc()

        del(opt)

    # Found parameters
    found_tcp = np.copy(p)
    found_cp = transform_to_common_param(p)
    _ = final_posterior(found_tcp)
    found_ip = final_posterior.get_current_independent_param()

    logposteriors.append(s)
    params_c.append(found_cp)
    params_i.append(found_ip)

#
# Done
#
obtained_logposterior0 = logposteriors[0]
obtained_parameters0 = params_c[0]
obtained_parameters_i0 = params_i[0]

# Show results
saveas = file_name
print('Found solution:          Prior parameters:' )
# Store output
with open('%s/%s-solution-%s.txt' % (savedir, saveas, fit_seed), 'w') as f:
    for k, x in enumerate(obtained_parameters0):
        print(pints.strfloat(x) + '    ' + pints.strfloat(x0_common_all_mean[k]))
        f.write(pints.strfloat(x) + '\n')
np.savetxt('%s/%s-solution_i-%s.txt' % (savedir, saveas, fit_seed), obtained_parameters_i0)

with open('%s/%s-cells-%s.txt' % (savedir, saveas, fit_seed), 'w') as f:
    for c in selectedwell:
        f.write(c + '\n')

# Plot results
model_ideal = m.Model('../mmt-model-files/ideal-ikr.mmt',
                protocol_def=protocol_def,
                temperature=273.15 + temperature,  # K
                transform=parametertransform.donothing,
                useFilterCap=useFilterCap)  # ignore capacitive spike
voltage = model_ideal.voltage(times)
output0 = final_posterior.problem_evaluate(
        transform_from_common_param(obtained_parameters0),
        obtained_parameters_i0)
if not os.path.isdir(savedir + '/plot-solution-%s' % fit_seed):
    os.makedirs(savedir + '/plot-solution-%s' % fit_seed)
for i in range(n_cells):
    fig, axes = plt.subplots(2, 1, sharex=True, figsize=(8, 6))
    axes[0].plot(times, voltage, c='#7f7f7f')
    axes[0].set_ylabel('Voltage (mV)')
    axes[1].plot(times, data_all[:, i], alpha=0.5, label='data')
    ideal = model_ideal.simulate(
            np.append(obtained_parameters_i0[i][0], obtained_parameters0),
            times)
    axes[1].plot(times, ideal, label='same kinetics')
    axes[1].plot(times, output0[:, i], label='found solution 1')
    axes[1].legend()
    axes[1].set_ylabel('Current (pA)')
    axes[1].set_xlabel('Time (ms)')
    plt.subplots_adjust(hspace=0)
    plt.savefig(savedir + '/plot-solution-%s/%s' % (fit_seed, selectedwell[i]))
    plt.close()

print('Done.')

## eof


# Copyright 2021 The NetKet Authors - All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from functools import partial
from typing import Callable
import warnings

import jax
from jax import numpy as jnp

from netket import jax as nkjax
from netket.stats import Stats
from netket.utils.types import PyTree
from netket.utils.dispatch import dispatch

from netket.operator import (
    AbstractOperator,
    DiscreteOperator,
    ContinuousOperator,
    Squared,
)

from netket.vqs import expect

from netket.vqs.mc import (
    kernels,
    get_local_kernel,
    get_local_kernel_arguments,
)

from .state import MCState

# Dispatches to select what expect-kernel to use
@dispatch
def get_local_kernel(vstate: MCState, Ô: Squared, chunk_size: int):  # noqa: F811
    return kernels.local_value_squared_kernel_chunked

@dispatch
def get_local_kernel(  # noqa: F811
    vstate: MCState, Ô: DiscreteOperator, chunk_size: int
):
    return kernels.local_value_kernel_chunked

def _local_continuous_kernel(kernel, logpsi, pars, σ, args, *, chunk_size=None):
    def _kernel(σ):
        return kernel(logpsi, pars, σ, args)

    return nkjax.vmap_chunked(_kernel, in_axes=0, chunk_size=chunk_size)(σ)

@dispatch
def get_local_kernel(  # noqa: F811
    vstate: MCState, Ô: ContinuousOperator, chunk_size: int
):
    return nkjax.HashablePartial(_local_continuous_kernel, Ô._expect_kernel)

# If batch_size is None, ignore it and remove it from signature so that we fall back
# to already implemented methods
@expect.dispatch
def expect_nochunking(vstate: MCState, operator: AbstractOperator, chunk_size: None):
    return expect(vstate, operator)

# if no implementation exists for batched, fall back to unbatched methods.
@expect.dispatch
def expect_fallback(
    vstate: MCState, operator: AbstractOperator, chunk_size
):  # noqa: F811
    warnings.warn(
        f"Ignoring chunk_size={chunk_size} for expect_and_grad method with signature "
        f"({type(vstate)}, {type(operator)}) because no implementation supporting "
        f"chunking for this signature exists."
    )

    return expect(vstate, operator)

@expect.dispatch
def expect_mcstate_operator_chunked(
    vstate: MCState, Ô: AbstractOperator, chunk_size: int
) -> Stats:  # noqa: F811
    σ, args = get_local_kernel_arguments(vstate, Ô)

    local_estimator_fun = get_local_kernel(vstate, Ô, chunk_size)

    return _expect_chunking(
        chunk_size,
        local_estimator_fun,
        vstate._apply_fun,
        vstate.sampler.machine_pow,
        vstate.parameters,
        vstate.model_state,
        σ,
        args,
    )

@partial(jax.jit, static_argnums=(0, 1, 2))
def _expect_chunking(
    chunk_size: int,
    local_value_kernel: Callable,
    model_apply_fun: Callable,
    machine_pow: int,
    parameters: PyTree,
    model_state: PyTree,
    σ: jnp.ndarray,
    args: PyTree,
) -> Stats:
    σ_shape = σ.shape

    if jnp.ndim(σ) != 2:
        σ = σ.reshape((-1, σ_shape[-1]))

    def logpsi(w, σ):
        return model_apply_fun({"params": w, **model_state}, σ)

    def log_pdf(w, σ):
        return machine_pow * model_apply_fun({"params": w, **model_state}, σ).real

    _, Ō_stats = nkjax.expect(
        log_pdf,
        partial(local_value_kernel, logpsi, chunk_size=chunk_size),
        parameters,
        σ,
        args,
        n_chains=σ_shape[0],
    )

    return Ō_stats