# Copyright (c) 2009-2022 The Regents of the University of Michigan.
# Part of HOOMD-blue, released under the BSD 3-Clause License.

"""Compute properties of hard particle configurations.

The HPMC compute classes analyze the system configuration and provide results
as loggable quantities for use with `hoomd.logging.Logger` or by direct access
via the Python API. `FreeVolume` computes the free volume available to small
particles, such as depletants, and `SDF` computes the pressure in system of
convex particles with a fixed box size.
"""

from __future__ import print_function

from hoomd import _hoomd
from hoomd.operation import Compute
from hoomd.hpmc import _hpmc
from hoomd.hpmc import integrate
from hoomd.data.parameterdicts import ParameterDict
from hoomd.logging import log
import hoomd
import numpy

class FreeVolume(Compute):
    r"""Compute the free volume available to a test particle.

    Args:
        test_particle_type (str): Test particle type.
        num_samples (int): Number of samples to evaluate.

    `FreeVolume` computes the free volume in the simulation state available to a
    given test particle shape using Monte Carlo integration. Use it in
    combination with `hoomd.hpmc.integrate.HPMCIntegrator`, which defines the
    particle shape parameters. Particles of ``test_particle_type`` may or may
    not be present in the simulation state.

    `FreeVolume` generates `num_samples` (:math:`n_\mathrm{samples}`) trial
    particle configurations with positions :math:`\vec{r}^t_j` uniformly
    distributed in the simulation box, and orientations :math:`\mathbf{q}^t_j`
    uniformly distributed among rotations matching the box dimensionality.
    `FreeVolume` counts the number of successful samples that do not overlap
    particles in the simulation state:

    .. math::

        n_\mathrm{success} = \sum_{j=1}^{n_\mathrm{samples}}
            \prod_{i=0}^{N_\mathrm{particles}-1}
            \prod_{\vec{A} \in B_\mathrm{images}}
            \left[
            \mathrm{overlap}\left(
            S_i(\mathbf{q}_i),
            S_t(\mathbf{q}^t_j, \vec{r}^t_j - (\vec{r}_i + \vec{A}))
            \right) = \emptyset
            \right]

    where :math:`\mathrm{overlap}` is the shape overlap function defined in
    `hoomd.hpmc.integrate`, :math:`S_i` is the shape of particle :math:`i`,
    :math:`S_t` is the shape of the test particle, :math:`\vec{A} = h\vec{a}_1 +
    k\vec{a}_2 + l\vec{a}_3` is a vector that translates by periodic box images,
    the set of box images includes all image vectors necessary to find overlaps
    between particles in the primary image with particles in periodic images,
    and the square brackets denote the Iverson bracket.

    The free volume :math:`V_\mathrm{free}` is given by:

    .. math::
        V_\mathrm{free} = \frac{n_\mathrm{success}}
                               {n_\mathrm{samples}} V_\mathrm{box}

    where :math:`V_\mathrm{box}` is the volume of the simulation box (or area in
    2D).

    Note:

        `FreeVolume` respects the HPMC integrator's ``interaction_matrix``.

    .. rubric:: Mixed precision

    `FreeVolume` uses reduced precision floating point arithmetic when checking
    for particle overlaps in the local particle reference frame.

    .. rubric:: Box images

    On CPU devices, `FreeVolume` does not apply the minimum image convention. It
    supports small boxes where particles may overlap with non-primary images of
    other particles, including self overlap. On GPU devices, `FreeVolume`
    applies the minimum image convention.

    Examples::

        fv = hoomd.hpmc.compute.FreeVolume(test_particle_type='B',
                                           num_samples=1000)

    Attributes:
        test_particle_type (str): Test particle type.

        num_samples (int): Number of samples to evaluate.

    """

    def __init__(self, test_particle_type, num_samples):
        # store metadata
        param_dict = ParameterDict(test_particle_type=str, num_samples=int)
        param_dict.update(
            dict(test_particle_type=test_particle_type,
                 num_samples=num_samples))
        self._param_dict.update(param_dict)

    def _attach(self):
        integrator = self._simulation.operations.integrator
        if not isinstance(integrator, integrate.HPMCIntegrator):
            raise RuntimeError("The integrator must be an HPMC integrator.")

        # Extract 'Shape' from '<hoomd.hpmc.integrate.Shape object>'
        integrator_name = integrator.__class__.__name__
        try:
            if isinstance(self._simulation.device, hoomd.device.CPU):
                cpp_cls = getattr(_hpmc, 'ComputeFreeVolume' + integrator_name)
            else:
                cpp_cls = getattr(_hpmc,
                                  'ComputeFreeVolume' + integrator_name + 'GPU')
        except AttributeError:
            raise RuntimeError("Unsupported integrator.")

        cl = _hoomd.CellList(self._simulation.state._cpp_sys_def)
        self._cpp_obj = cpp_cls(self._simulation.state._cpp_sys_def,
                                integrator._cpp_obj, cl)

        super()._attach()

    @log(requires_run=True)
    def free_volume(self):
        """Free volume available to the test particle \
        :math:`[\\mathrm{length}^{2}]` in 2D and \
        :math:`[\\mathrm{length}^{3}]` in 3D."""
        self._cpp_obj.compute(self._simulation.timestep)
        return self._cpp_obj.free_volume

class SDF(Compute):
    r"""Compute the scale distribution function.

    Args:
        xmax (float): Maximum *x* value at the right hand side of the rightmost
            bin :math:`[\mathrm{length}]`.
        dx (float): Bin width :math:`[\mathrm{length}]`.

    `SDF` computes the proability distribution :math:`s(x)` of particles
    overlapping as a function of separation. It estimates :math:`s(x)`
    numerically by computing a histogram with
    :math:`\lfloor x_\mathrm{max}/ \delta x \rfloor` bins of width `dx`
    (:math:`\delta x`).

    See Also:
         `Anderson 2016 <https://dx.doi.org/10.1016/j.cpc.2016.02.024>`_
         describes the theory relating `SDF` to the system pressure.

    .. rubric:: Implementation

    For each pair of particles :math:`i` and :math:`j` `SDF` scales the particle
    separation vector by the factor :math:`(1-x)` and finds the smallest
    positive value of :math:`x` leading to an overlap of the particle shapes:

    .. math::

        x_{ij}(\vec{A}) = \min \{ x \in \mathbb{R}_{> 0} :
            \mathrm{overlap}\left(
                S_i(\mathbf{q}_i),
                S_j(\mathbf{q}_j, (1-x)(\vec{r}^t_j - (\vec{r}_i + \vec{A})))
            \right) \ne \emptyset \}

    where :math:`\mathrm{overlap}` is the shape overlap function defined in
    `hoomd.hpmc.integrate`, :math:`S_i` is the shape of particle
    :math:`i`, and :math:`\vec{A} = h\vec{a}_1 + k\vec{a}_2 + l\vec{a}_3` is a
    vector that translates by periodic box images.

    :math:`x_i` is the minimum value of :math:`x_{ij}` for a single particle:

    .. math::

        x_i = \min \{ x_{ij} : \vec{A} \in B_\mathrm{images},
                     j \in [0,N_\mathrm{particles}) \}

    where the set of box images includes all image vectors necessary to find
    overlaps between particles in the primary image with particles in periodic
    images.

    `SDF` adds a single count to the histogram for each particle :math:`i`:

    .. math::

        s(x + \delta x/2) = \frac{1}{N_\mathrm{particles} \cdot \delta x}
            \sum_{i=0}^{N_\mathrm{particles}-1}
            [x \le x_i < x + \delta x]

    where the square brackets denote the Iverson bracket, and :math:`s(x +
    \delta x/2)` is evaluated for :math:`\{ x \in \mathbb{R}, 0 \le x <
    x_\mathrm{max}, x = k \cdot \delta x, k \in \mathbb{Z}^* \}`.

    .. rubric:: Pressure

    The extrapolation of :math:`s(x)` to :math:`x = 0`, :math:`s(0+)` is related
    to the pressure :math:`P`:

    .. math::
        \beta P = \rho \left(1 + \frac{s(0+)}{2d} \right)

    where :math:`d` is the dimensionality of the system, :math:`\rho` is the
    number density, and :math:`\beta = \frac{1}{kT}`. This measurement of the
    pressure is inherently noisy due to the nature of the sampling. Average
    `betaP` over many timesteps to obtain accurate results.

    Assuming particle diameters are ~1, these paramater values typically
    achieve good results:

      * ``xmax = 0.02``
      * ``dx = 1e-4``

    In systems near densest packings, ``dx=1e-5`` may be needed along with
    smaller ``xmax``. Check that :math:`\sum_k s(x_k) \cdot dx \approx 0.5`.

    Warning:
        `SDF` does not compute correct pressures for simulations with
        concave particles or enthalpic interactions.

    Note:
        `SDF` always runs on the CPU.

    .. rubric:: Mixed precision

    `SDF` uses reduced precision floating point arithmetic when checking
    for particle overlaps in the local particle reference frame.

    .. rubric:: Box images

    `SDF` does not apply the minimum image convention. It supports small boxes
    where particles may overlap with non-primary images of other particles,
    including self overlap.

    Attributes:
        xmax (float): Maximum *x* value at the right hand side of the rightmost
            bin :math:`[\mathrm{length}]`.

        dx (float): Bin width :math:`[\mathrm{length}]`.
    """

    def __init__(self, xmax, dx):
        # store metadata
        param_dict = ParameterDict(xmax=float(xmax), dx=float(dx))
        self._param_dict.update(param_dict)

    def _attach(self):
        integrator = self._simulation.operations.integrator
        if not isinstance(integrator, integrate.HPMCIntegrator):
            raise RuntimeError("The integrator must be an HPMC integrator.")

        # Extract 'Shape' from '<hoomd.hpmc.integrate.Shape object>'
        integrator_name = integrator.__class__.__name__

        cpp_cls = getattr(_hpmc, 'ComputeSDF' + integrator_name)

        self._cpp_obj = cpp_cls(self._simulation.state._cpp_sys_def,
                                integrator._cpp_obj, self.xmax, self.dx)

        super()._attach()

    @log(category='sequence', requires_run=True)
    def sdf(self):
        """(*N_bins*,) `numpy.ndarray` of `float`): :math:`s[k]` - The scale \
        distribution function :math:`[\\mathrm{probability\\ density}]`.

        The :math:`x` at the center of bin :math:`k` is:
        :math:`x = k \\cdot \\delta x + \\delta x/2`.

        Attention:
            In MPI parallel execution, the array is available on rank 0 only.
            `sdf` is `None` on ranks >= 1.
        """
        self._cpp_obj.compute(self._simulation.timestep)
        return self._cpp_obj.sdf

    @log(requires_run=True)
    def betaP(self):  # noqa: N802 - allow function name
        """float: Beta times pressure in NVT simulations \
        :math:`\\left[ \\mathrm{length}^{-d} \\right]`.

        Uses a polynomial curve fit of degree 5 to estimate :math:`s(0+)` and
        computes the pressure via:

        .. math::
            \\beta P = \\rho \\left(1 + \\frac{s(0+)}{2d} \\right)

        where :math:`d` is the dimensionality of the system, :math:`\\rho` is
        the number density, and :math:`\\beta = \\frac{1}{kT}`.

        Attention:
            In MPI parallel execution, `betaP` is available on rank 0 only.
            `betaP` is `None` on ranks >= 1.
        """
        if not numpy.isnan(self.sdf).all():
            # get the values to fit
            n_fit = int(numpy.ceil(self.xmax / self.dx))
            sdf_fit = self.sdf[0:n_fit]
            # construct the x coordinates
            x_fit = numpy.arange(0, self.xmax, self.dx)
            x_fit += self.dx / 2
            # perform the fit and extrapolation
            p = numpy.polyfit(x_fit, sdf_fit, 5)

            box = self._simulation.state.box
            N = self._simulation.state.N_particles
            rho = N / box.volume
            return rho * (1 + numpy.polyval(p, 0.0) / (2 * box.dimensions))
        else:
            return None


import numpy as np

from bayes_implicit_solvent.continuous_parameter_experiments.elemental_types_mh import log_prior, mols, ll, data_path, \
    smiles

smiles_list = smiles
from bayes_implicit_solvent.typers import RADIUS_UNIT

from bayes_implicit_solvent.freesolv import smiles_list
from bayes_implicit_solvent.typers import AtomSpecificationProposal

np.random.seed(0)

from bayes_implicit_solvent.gb_models.obc2_parameters import mbondi_model

initial_tree = mbondi_model
initial_tree.remove_node('[#14]') # otherwise everything is -inf, because this type will be empty
initial_tree.proposal_sigmas['radius'] = 1e-2 * RADIUS_UNIT
initial_tree.proposal_sigmas['scale_factor'] = 1e-2

# add one more parameter per element appearing in FreeSolv but not specified in obc2 parameter set to initial tree
for i in [17, 35, 53]:
    smirks = '[#{}]'.format(i)
    initial_tree.add_child(smirks, '*')
    initial_tree.un_delete_able_types.add(smirks)

#specifiers = ['X1', 'X2', 'X3', 'X4', 'a', 'A']
#atom_specification_proposal = AtomSpecificationProposal(atomic_specifiers=specifiers)
#smirks_elaboration_proposal = atom_specification_proposal

all_bond_specifiers = ['@', '-', '#', '=', ':']
from bayes_implicit_solvent.smarts import atomic_number_dict

all_bondable_types = list(atomic_number_dict.keys())

# atomic_decorators list:
ring_specifiers = ['r0', 'r3', 'r4', 'r5', 'r6', 'r7', 'a', 'A']
charge_specifiers = ['-1', '+0', '+1', '+2']
hydrogen_count_specifiers = ['H0', 'H1', 'H2', 'H3', 'H4']
connectivity_specifiers = ['X1', 'X2', 'X3', 'X4']

all_specifier_lists = [
    ring_specifiers,
    charge_specifiers,
    hydrogen_count_specifiers,
    connectivity_specifiers,
]

from itertools import chain

all_atomic_specifiers = list(chain(*all_specifier_lists))
#all_bondable_types += ['[{}]'.format(s) for s in all_atomic_specifiers]
#all_decorators = all_bondable_types + all_atomic_specifiers + all_bond_specifiers

from bayes_implicit_solvent.typers import BondProposal, BondSpecificationProposal, AtomSpecificationProposal, SMIRKSElaborationProposal
#bond_proposal = BondProposal(bondable_types=all_bondable_types)
atom_specification_proposal = AtomSpecificationProposal(atomic_specifiers=all_atomic_specifiers)
#bond_specification_proposal = BondSpecificationProposal(bond_specifiers=all_bond_specifiers)

smirks_elaborators = [
    #bond_proposal,
    atom_specification_proposal,
    #bond_specification_proposal,
]
smirks_elaboration_proposal = SMIRKSElaborationProposal(smirks_elaborators=smirks_elaborators)

print('initial tree:')
print(initial_tree)

n_configuration_samples = 5

import os

name = 'tree_rjmc_n_config={}_{}_ll'.format(n_configuration_samples, ll)
smiles_subset_fname = os.path.join(data_path,
                                   'smiles_subset_{}.txt'.format(name))
with open(smiles_subset_fname, 'w') as f:
    f.writelines(['{}\n'.format(s) for s in smiles_list])

from bayes_implicit_solvent.prior_checking import check_no_empty_types

error_y_trees = []

for mol in mols:
    thinning = int(len(mol.vacuum_traj) / n_configuration_samples)
    mol.vacuum_traj = mol.vacuum_traj[::thinning]

def log_prob(tree):
    log_prior_value = check_no_empty_types(tree)

    theta = np.hstack([tree.get_radii(), tree.get_scale_factors()])

    log_prior_value += log_prior(theta)

    if log_prior_value > -np.inf:
        try:
            # TODO: Parallelize. Note that multiprocessing.Pool won't work here because it doesn't play nice with SwigPy objects
            # TODO: update to allow scale factors to be variable also
            log_likelihood_value = 0
            for mol in mols:
                radii = tree.assign_radii(mol.mol) / RADIUS_UNIT
                scale_factors = tree.assign_scale_factors(mol.mol)

                log_likelihood_value += mol.log_prob(radii, scale_factors)
        except:
            global error_y_trees
            error_y_trees.append(tree)
            print('Warning! Encountered un-anticipated exception!')
            return - np.inf
        return log_prior_value + log_likelihood_value
    else:
        return log_prior_value

from bayes_implicit_solvent.samplers import tree_rjmc
from pickle import dump

n_iterations = 10000

result = tree_rjmc(initial_tree, log_prob, smirks_elaboration_proposal, n_iterations=n_iterations,
                   fraction_cross_model_proposals=0.1)
with open('elaborate_tree_rjmc_run_n_compounds={}_n_iter={}_gaussian_ll.pkl'.format(len(mols), n_iterations),
          'wb') as f:
    dump(result, f)

with open('error_y_trees.pkl', 'wb') as f:
    dump(error_y_trees, f)


#!/usr/bin/env python
# coding: utf-8
'''
--------------------------------------------------------------------------------
Vaccines Manufacturing Model
--------------------------------------------------------------------------------
Bryden Wood
License: MIT, see full license in LICENSE.txt
--------------------------------------------------------------------------------
Date: 2020-11-20
Authors:
    Jiabin Li
    Wynne Lim
    Stamatios Stamatiadis
    David Reader
--------------------------------------------------------------------------------
The purpose of the manufacturing model is to identify plausible scenarios
for how long it will take to manufacture vaccines to treat the populations
in need, and is split into two parts: Manufacturing Preparation, predicting
when primary (drug substance) and secondary (drug product) manufacturing
could start; and Manufacturing Capacity, predicting when enough doses are
made to meet the needs of the target populations. The model uses output from
the Research and Development model, expert input and interviews, literature
and manufacturing capacity survey data as inputs. The model uses Monte Carlo
techniques to derive the dates at which the vaccine production targets will
be met and should therefore be run many times to smooth statistical
fluctuations.
--------------------------------------------------------------------------------
Developed for compatibility with Python 3.7.

This model functions as an importable module to R&D Vaccine Predictions model
('R&D model').

The model uses the following libraries from the default Python installation:
    random
    datetime
    json
    copy

The model uses the following additional libraries:
    numpy
    scipy
    pandas
--------------------------------------------------------------------------------
'''

# Module imports
import pandas as pd
import numpy as np
import random
import datetime
import json
import copy

# --------------------------------------------------------------------------------

def initialise(params, vaccines, manufacturing):
    '''
    This function creates and populates a set of global dataframes which can
    then be accessed from other functions. External data is loaded from JSON
    files.

    Parameters
    ----------
    None

    Returns
    ----------
    None
    '''

    # manufacturing model initialisation.

    global primary_cap
    global ratio_pri_v
    global ratio_pri_d
    global plants_categories

    global secondary_cap
    global ratio_sec
    global sec_plants_available

    global df_Categories
    global df_product

    global default
    global m_params

    global funding
    global demand
    global targetDoses

    global output_summary
    global cumulative_summary
    global df_schedule

    global ramp_up

    # load offline default data and online parameters
    default = jread(manufacturing)
    m_params = params
    ramp_up = params['ramp_up']

    primary_cap, ratio_pri_v, ratio_pri_d = getPrimaryInput()

    # initialise the product info
    df_Categories = default['Capacity Category']
    df_product = getProduct(vaccines)

    # get the number of primary plants available for each category
    plants_categories = [6, 9, 6, 0, 6]

    # initialise the schedule input
    df_schedule, funding = getScheduleInput()

    # initialise the secondary input
    secondary_cap, ratio_sec, sec_plants_available = getSecondaryInput()

    # initialise the demand data
    demand = m_params['Doses needed']
    targets = ["Group 1", "Group 2", "Group 3", "Group 4"]
    wastage = demand['Percentage of drug wastage']
    dose_perVx = demand['Number of doses per vaccine']
    # set the target numbers
    targetDoses = [(demand[x] / (1 - wastage / 100)) * dose_perVx for x in targets]
    targetDoses = np.cumsum(targetDoses)

    # define global variables
    output_summary = pd.DataFrame()
    cumulative_summary = pd.DataFrame()

# --------------------------------------------------------------------------------

def getSecondaryInput():
    '''
    This function extracts the relevant information pertaining to 'secondary'
    from the dictionaries: 'default' and 'm_params' and
    stores them as a separate dictionary 'secDefaultCap' and 'parameter' respectively

    Parameters
    ----------
    None

    Returns
    ----------
    secDefaultCap: a dictionary that stores the default secondary capacity of each country.

    ratio_sec: an array of ratio for low, most likely, high.

    sec_plants_available: a value for the number of secondary plants available for each category
    '''

    # secDefaultCap refers to the offline input table with the default (highest available) capacity for each country
    secDefaultCap = default['Sec Capacity']
    parameter = m_params['Sec Input']

    # get the ratio for low, most likely, high
    parameter_default = 4268.328666666668  # this is the value that is the original value we calculated

    cols = [name + ' Available Capacity (million doses/month)' for name in ['Lowest', 'Most Likely', 'Highest']]
    ratio_sec = [parameter[col] / parameter_default for col in cols]

    # get the number of secondary plants available for each category
    sec_plants_available = 0
    for key, value in secDefaultCap.items():
        if (value > 0): sec_plants_available += 1

    return secDefaultCap, ratio_sec, sec_plants_available

# --------------------------------------------------------------------------------

def getIteration(df_rnd):
    '''
    Takes the output from the R&D model and creates the iteration data required
    for the manufacturing functions. This includes joining with vaccine data in
    the database and restricting the number of successful vaccines in each
    category based on the model assumptions.

    Parameters
    ----------
    df_rnd: a dataframe of successful vaccines from the R&D model for the
        current iteration

    Returns
    ----------
    df_iteration: a dataframe of the successful vaccines with data processed
        for use in the manufacturing model
    '''
    global mfg_at_risk

    # generate iteration table
    tryID = df_rnd['try']
    v = df_rnd['vaccines']
    df = pd.DataFrame(v, columns=['Vaccine', 'phase1_start', 'phase1_end', 'phase2_start', 'phase2_end', 'Phase III',
                                  'phase3_end', 'approval_start', 'Approval (month)'])
    df['try'] = tryID

    # filter the vaccines being approved
    df_rnd = df[df['Approval (month)'] != '']
    df_rnd.reset_index(drop=True, inplace=True)

    # find platform and category
    df_iteration = df_rnd
    df_iteration = pd.merge(df_iteration, df_product, on='Vaccine', how='left')

    # get product manufacturing at risk info
    df_iteration_copy = df_iteration.copy()
    df_iteration_copy['Mfg before approval'] = df_iteration_copy['Funding'].apply(
        lambda x: funding[x]['Manufacturing start before approval?'])
    mfg_at_risk = df_iteration_copy[['Phase III', 'Mfg before approval', 'Approval (month)']].to_dict('records')
    mfg_at_risk = dict(zip(df_rnd['Vaccine'], mfg_at_risk))

    df_iteration.drop(
        ['phase1_start', 'phase1_end', 'phase2_start', 'phase2_end', 'Phase III', 'phase3_end', 'approval_start'],
        axis=1, inplace=True)

    platforms = np.array(df_iteration['Platform'])
    categories = []

    # get capacity category info
    for i in df_iteration.index.values:
        categories.append(df_Categories[platforms[i]])

    df_iteration['Category'] = categories

    # Keep the first 3 of vaccine in the same category
    df_iteration = df_iteration.sort_values(by=['Approval (month)'])
    df_iteration = df_iteration.groupby('Category').head(3)
    df_iteration.reset_index(drop=True, inplace=True)

    # rearrange the table
    # df_iteration = df_iteration[['try',]]

    return df_iteration

# --------------------------------------------------------------------------------

def getProduct(vaccines):
    '''
    This function loads the vaccine database, selects the columns that are
    required (number, name and platform) and stores them in a dataframe.

    Parameters
    ----------
    None

    Returns
    ----------
    df_product: filtered vaccine table containing only essential columns
    '''

    # filter the vaccine data
    data = vaccines
    df_product = pd.DataFrame(columns=['Vaccine', 'Platform', 'Funding'])

    vaccines = []
    platforms = []
    funding = []

    for i in range(len(data)):
        vaccines.append(data[i]['number'])
        platforms.append(data[i]['platform'])
        funding.append(data[i]['funding_category'])

    df_product['Vaccine'] = vaccines
    df_product['Vaccine'] = df_product['Vaccine'].astype(int)
    df_product['Platform'] = platforms
    df_product['Funding'] = funding

    return df_product

# --------------------------------------------------------------------------------

def getManufacturingStartTime(df_iteration):
    '''
    Gets the manufacturing start time based on the current iteration parameters
    and the required manufacturing preparation timeline using the getSchedule
    function.

    Parameters
    ----------
    df_iteration: a dataframe of the successful vaccines with data processed
        for use in the manufacturing model

    Returns
    ----------
    None
    '''

    platforms = np.array(df_iteration['Platform'])
    approval_months = np.array(df_iteration['Approval (month)'])
    funding = np.array(df_iteration['Funding'])
    pri_starts = np.zeros(len(df_iteration))
    sec_starts = np.zeros(len(df_iteration))

    # get primary and secondary start time from scheduling (getSchedule) function
    for i in df_iteration.index.values:
        t = getSchedule(platforms[i], approval_months[i], funding[i])
        pri_starts[i] = t[0]
        sec_starts[i] = t[1]

    df_iteration['Primary Start'] = pri_starts
    df_iteration['Secondary Start'] = sec_starts
    df_iteration['Primary available'] = 0
    df_iteration['Primary assigned'] = 0

    df_iteration.drop(['Funding'], axis=1, inplace=True)

# --------------------------------------------------------------------------------
def getPrimaryInput():
    '''
    This function extracts the relevant information pertaining to 'primary'
    from the dictionaries: 'default' and 'm_params' and
    stores them as a separate dictionary 'primary_Cap' and 'parameter' respectively

    Parameters
    ----------
    None

    Returns
    ----------
    primary_cap: a dictionary that stores the default primary capacity of each country
                for each of the platform

    ratio_pri_v: a dictionary containing the ratios for the lowest, most likely and highest volume
                that is used to convert default value to online value

    ratio_pri_d: a dictionary containing the ratio for the number of doses (million doses per month)
                that is used to convert default value to online value

    '''

    # primary_cap refers to the offline input table with the default primary capacity for each country
    primary_cap = default['Pri Capacity']

    # parameter is the input table for primary capacity - from online platform
    parameter = m_params['Pri Input']

    # dfPrimary['1'], dfPrimary['2'], dfPrimary['3'], dfPrimary['4'], dfPrimary['5'] = 0, 0, 0, 0, 0

    platforms = ['DNA', 'Inactivated', 'Live-attenuated', 'Non-replicating viral vector', 'Replicating viral vector',
                 'Protein subunit', 'Other', 'RNA']
    platform_cat = ['DNA', 'Protein subunit', 'Inactivated', 'RNA']

    # dv -> default volume  dd -> default doses
    parameter_dv = [268, 1113, 779, 167728]
    parameter_dd = [112, 6488, 6488, 11596, 11596, 5798, 6355, 280]

    # ratio -> used to convert default value to online value
    ratio_pri_v = {}
    ratio_pri_d = {}
    cols = [name + ' Volume V (m3)' for name in ['Lowest', 'Most Likely', 'Highest']]

    for p in platform_cat:
        ratio_pri_v[p] = [parameter[p][col] / parameter_dv[platform_cat.index(p)] for col in cols]

    for p in platforms:
        ratio_pri_d[p] = parameter[p]['Doses for most likely volume (Million Doses per month) N_v'] / parameter_dd[
            platforms.index(p)]

    return primary_cap, ratio_pri_v, ratio_pri_d

# --------------------------------------------------------------------------------
def primary(df_iteration):
    '''
    Takes the current iteration data and allocates primary production capacity
    to each vaccine candidate based on the model assumptions.

    Parameters
    ----------
    df_iteration: a dataframe of the successful vaccines with data processed
        for use in the manufacturing model

    Returns
    ----------
    df_priAllocation: a dataframe of primary capacity allocations for the
        successful vaccines
    '''

    # sort the iteration table based on the primary start time
    df_iteration = df_iteration.sort_values(by=['Primary Start']).reset_index()

    # get capacity available for each category
    platform_cat = ['DNA', 'Protein subunit', 'Inactivated', 'RNA']
    category_list = [1, 2, 3, 5]

    platforms_bycat = {
        '1': ['DNA'],
        '2': ['Non-replicating viral vector', 'Replicating viral vector', 'Protein subunit', 'Other'],
        '3': ['Inactivated', 'Live-attenuated'],
        '5': ['RNA']
    }
    sorted_cap = {}

    primary_cap_copy = copy.deepcopy(primary_cap)

    for i in category_list:
        category_cap = primary_cap_copy[str(i)]
        example_p = platform_cat[category_list.index(i)]
        for p in platforms_bycat[str(i)]:
            ratio_cat = random.triangular(ratio_pri_v[example_p][0], ratio_pri_v[example_p][2],
                                          ratio_pri_v[example_p][1])
            for key, value in category_cap.items():
                category_cap[key][p] = category_cap[key][p] * ratio_cat

        platform_r = platforms_bycat[str(i)][0]

        sorted_temp = sorted(category_cap.items(), key=lambda x: x[1][platform_r], reverse=True)

        sorted_cap[str(i)] = sorted_temp

    # primary info initialisation

    # calculate the number of plants available per vaccine
    vx_categories = [df_iteration[df_iteration['Category'] == i].index.size for i in range(1, 6)]

    # assign no. of plants available to each vaccine product
    pri_plants_available = plants_categories[:]

    # Initialise arrays
    iter_categories = np.array(df_iteration['Category'])
    iter_pri_avail = np.array(df_iteration['Primary available'])
    iter_try = np.array(df_iteration['try'])
    iter_vaccines = np.array(df_iteration['Vaccine'])
    iter_platforms = np.array(df_iteration['Platform'])
    iter_pri_starts = np.array(df_iteration['Primary Start'])
    iter_pri_assigned = np.array(df_iteration['Primary assigned'])

    for i in range(len(df_iteration)):
        category = iter_categories[i] - 1

        if vx_categories[category] > 1:
            p = round(pri_plants_available[category] / vx_categories[category])
            iter_pri_avail[i] = p if p <= 3 else 3  # one vaccine can be assigned to maximum 3 plants
        else:
            p = pri_plants_available[category]
            iter_pri_avail[i] = p if p <= 3 else 3
        pri_plants_available[category] -= p if p <= 3 else 3
        vx_categories[category] -= 1

    df_iteration['Primary available'] = iter_pri_avail

    ### Primary product allocation
    df_priAllocation = pd.DataFrame(columns=['try', 'Country', 'Vaccine', 'Platform', 'Primary Start', 'Capacity'])

    prialloc_try = []
    prialloc_countries = []
    prialloc_vaccines = []
    prialloc_platforms = []
    prialloc_pri_starts = []
    prialloc_capacities = []

    for i in range(len(df_iteration)):
        while iter_pri_assigned[i] < iter_pri_avail[i]:
            iter_pri_assigned[i] += 1
            category = str(iter_categories[i])
            platform = iter_platforms[i]
            sorted_primary = sorted_cap[category]
            for country in sorted_primary:
                if country[1][platform] > 0:
                    prialloc_try.append(iter_try[i])
                    prialloc_countries.append(country[0])
                    prialloc_vaccines.append(iter_vaccines[i])
                    prialloc_platforms.append(iter_platforms[i])
                    prialloc_pri_starts.append(iter_pri_starts[i])
                    prialloc_capacities.append(country[1][platform] / 1000 * ratio_pri_d[platform])
                    sorted_primary.remove(country)
                    break

    df_priAllocation['try'] = prialloc_try
    df_priAllocation['Country'] = prialloc_countries
    df_priAllocation['Vaccine'] = prialloc_vaccines
    df_priAllocation['Platform'] = prialloc_platforms
    df_priAllocation['Primary Start'] = prialloc_pri_starts
    df_priAllocation['Capacity'] = prialloc_capacities

    return df_priAllocation

# --------------------------------------------------------------------------------

def secondary(df_iteration, df_priAllocation):
    '''
    Takes the current iteration data and allocates secondary production capacity
    to each vaccine candidate based on the model assumptions.

    Parameters
    ----------
    df_iteration: a dataframe of the successful vaccines with data processed
        for use in the manufacturing model
    df_priAllocation: a dataframe of primary capacity allocations for the
        successful vaccines

    Returns
    ----------
    df_secCumProduction: a dataframe of secondary capacity allocations for the
        successful vaccines, and cumulative dose count each month
    '''

    # sort the iteration table based on the secondary start time
    df_iteration = df_iteration.sort_values(by=['Secondary Start']).reset_index(drop=True)

    # get the number of plants available
    sec_plantsC = sec_plants_available
    secondary_copy = copy.deepcopy(secondary_cap)

    # get random secondary capacity for each country based on a triangular distribution based on low,most likely,high value
    ratios = random.triangular(ratio_sec[0], ratio_sec[2], ratio_sec[1])
    for key, value in secondary_copy.items():
        secondary_copy[key] = secondary_copy[key] * ratios

    # sort the secondary plants by capacity
    sort_secondary = sorted(secondary_copy.items(), key=lambda x: x[1], reverse=True)

    # get the no. of plants available for each vaccine
    v_count = df_iteration.index.size
    df_iteration['Secondary available'] = 0
    df_iteration['Secondary assigned'] = 0

    # Initialise arrays
    iter_sec_avail = np.zeros(len(df_iteration))
    iter_sec_assigned = np.zeros(len(df_iteration))
    iter_try = np.array(df_iteration['try'])
    iter_vaccines = np.array(df_iteration['Vaccine'])
    iter_platforms = np.array(df_iteration['Platform'])
    iter_sec_starts = np.array(df_iteration['Secondary Start'])

    for i in range(len(df_iteration)):
        if v_count > 1:
            p = round(sec_plantsC / v_count)
            iter_sec_avail[i] = p if p <= 3 else 3  # one vaccine can be assigned to maximum 3 plants
        else:
            p = sec_plantsC
            iter_sec_avail[i] = p if p <= 3 else 3  # one vaccine can be assigned to maximum 3 plants
        v_count -= 1
        sec_plantsC -= p

    df_iteration['Secondary available'] = iter_sec_avail

    ### secondary product allocation

    # initialisation
    df_secAllocation = pd.DataFrame(columns=['try', 'Country', 'Vaccine', 'Platform', 'Secondary Start', 'Capacity'])
    df_iteration['Secondary assigned'] = 0

    iter_sec_assigned = np.array(df_iteration['Secondary assigned'])

    secalloc_try = []
    secalloc_countries = []
    secalloc_vaccines = []
    secalloc_platforms = []
    secalloc_sec_starts = []
    secalloc_capacities = []

    # assign plants to vaccines
    for i in range(len(df_iteration)):
        while iter_sec_assigned[i] < iter_sec_avail[i]:
            iter_sec_assigned[i] += 1
            for country in sort_secondary:
                if country[1] > 0:
                    secalloc_try.append(iter_try[i])
                    secalloc_countries.append(country[0])
                    secalloc_vaccines.append(iter_vaccines[i])
                    secalloc_platforms.append(iter_platforms[i])
                    secalloc_sec_starts.append(iter_sec_starts[i])
                    secalloc_capacities.append(country[1] / 1000)
                    sort_secondary.remove(country)
                    break

    df_secAllocation['try'] = secalloc_try
    df_secAllocation['Country'] = secalloc_countries
    df_secAllocation['Vaccine'] = secalloc_vaccines
    df_secAllocation['Platform'] = secalloc_platforms
    df_secAllocation['Secondary Start'] = secalloc_sec_starts
    df_secAllocation['Capacity'] = secalloc_capacities

    # check whether primary or secondary is the bottle neck
    df_priThroughput = df_priAllocation.groupby('Vaccine')['Capacity'].sum()
    df_secThroughput = df_secAllocation.groupby('Vaccine')['Capacity'].sum()

    for i in range(len(df_secAllocation)):
        vx_ID = secalloc_vaccines[i]
        if df_secThroughput[vx_ID] > df_priThroughput[vx_ID]:
            df_secAllocation.loc[i, 'Capacity'] *= (df_priThroughput[vx_ID] / df_secThroughput[vx_ID])

    # get secondary throughput table
    df_secThroughput = df_secAllocation.copy()
    capacity_arr = np.array(df_secThroughput['Capacity'])
    start_arr = np.array(df_secThroughput['Secondary Start'])

    duration = ramp_up['duration']
    pre_approval = ramp_up['pre_approval']

    cum_arr = np.cumsum(np.linspace(0, 1, duration + 1))

    monthly_throughput = {}
    for j in range(1, 101):
        col_month = []
        for i in range(len(df_secThroughput)):
            vx_ID = secalloc_vaccines[i]
            phase3 = mfg_at_risk[vx_ID]['Phase III']
            mfg_before_approval = mfg_at_risk[vx_ID]['Mfg before approval']
            approv_month = mfg_at_risk[vx_ID]['Approval (month)']

            if phase3 > start_arr[i]:  # secondary -> phase 3 -> approval
                # calculate the cumulative initial dose
                if mfg_before_approval == 1 and j > phase3:
                    if j >= approv_month:
                        initial_doses = pre_approval * (approv_month - phase3)
                    else:
                        initial_doses = 0
                else:
                    initial_doses = 0

                # calculate the cumulative ramp up doses
                if j > approv_month:  # after ramp up
                    if j - approv_month >= duration:
                        cum_factor = (j - (approv_month + duration - 1)) + cum_arr[duration - 1] + initial_doses
                    else:  # in the ramp up period
                        cum_factor = cum_arr[int(j - approv_month)] + initial_doses
                    capacity_month = cum_factor * capacity_arr[i]

                else:
                    capacity_month = initial_doses * capacity_arr[i]

            else:
                if approv_month > start_arr[i]:  # phase 3 - > secondary - > approval
                    # calculate the cumulative initial dose
                    if mfg_before_approval == 1 and j > start_arr[i]:
                        if j >= approv_month:
                            initial_doses = pre_approval * (approv_month - start_arr[i])
                        else:
                            initial_doses = 0
                    else:
                        initial_doses = 0

                    # calculate the cumulative ramp up doses
                    if j > approv_month:  # after ramp up
                        if j - approv_month >= duration:
                            cum_factor = (j - (approv_month + duration - 1)) + cum_arr[duration - 1] + initial_doses
                        else:  # in the ramp up period
                            cum_factor = cum_arr[int(j - approv_month)] + initial_doses
                        capacity_month = cum_factor * capacity_arr[i]

                    else:
                        capacity_month = initial_doses * capacity_arr[i]

                else:  # phase 3 - > approval - > secondary
                    # calculate the cumulative initial dose
                    initial_doses = 0

                    # calculate the cumulative ramp up doses
                    if j > start_arr[i]:  # after ramp up
                        if j - start_arr[i] >= duration:
                            cum_factor = (j - (start_arr[i] + duration - 1)) + cum_arr[duration - 1] + initial_doses
                        else:  # in the ramp up period
                            cum_factor = cum_arr[int(j - start_arr[i])] + initial_doses
                        capacity_month = cum_factor * capacity_arr[i]

                    else:
                        capacity_month = initial_doses * capacity_arr[i]

            col_month.append(capacity_month)

        monthly_throughput[j] = col_month

    df_monthly = pd.DataFrame(monthly_throughput)
    df_secThroughput = pd.concat([df_secThroughput, df_monthly], axis=1)

    # get cumulative production for secondary
    df_secCumProduction = df_secThroughput.copy()

    return df_secCumProduction

# --------------------------------------------------------------------------------

def getSchedule(platform, approval_month, funding_cat):
    '''
    Calculates the critical path for a vaccine's manufacturing preparation
    schedule, and returns the months that primary and secondary manufacturing
    can start.

    Parameters
    ----------
    platform: the vaccine's platform
    approval_month: the month that the vaccine reached approval from R&D

    Returns
    ----------
    m_month: the months that primary and secondary manufacturing can start
    '''

    # get the gantt table
    df_gantt = df_schedule[platform].copy()

    funding_ratio = funding[funding_cat]['Gantt duration factor*']
    # update the dependencies based on the funding category

    if funding[funding_cat]['Simultaneous tech transfer?'] == 1:  # Simultaneous Tech Transfer
        df_gantt.loc['12', 'Predecessor'] = '0'
        df_gantt.loc['22', 'Predecessor'] = '21'
    else:
        pass

    if funding[funding_cat]['Manufacturing start before approval?'] == 1:  # Manufacturing before approval
        df_gantt.loc['17', 'Predecessor'] = '17'
        df_gantt.loc['27', 'Predecessor'] = '25,27'
    else:
        pass

    # Remove rows not in use
    df_gantt['Predecessor'] = df_gantt['Predecessor'].astype(str)

    # get time for each activity
    df_gantt['Time (days)'] = df_gantt[['Type', 'Value', 'Low', 'Most Likely', 'High']].apply(getValue, axis=1)

    # apply gantt duration factor
    df_gantt['Time (days)'] = df_gantt['Time (days)'] * funding_ratio

    # Initialise columns
    df_gantt['end_date'] = datetime.date(2020, 3, 1)
    df_gantt['start_date'] = datetime.date(2020, 3, 1)

    # Calculate end time

    predecessors = np.array(df_gantt['Predecessor'])
    start_dates = np.array(df_gantt['start_date'])
    end_dates = np.array(df_gantt['end_date'])
    time_days = np.array(df_gantt['Time (days)'])

    df_gantt = df_gantt[df_gantt['Predecessor'] != '-1']

    for i in df_gantt.index.values.astype(int):
        # assign start date
        if predecessors[i] != '0':

            if str(predecessors[i]).find(',') == -1:

                pTask_index = int(df_gantt.index[df_gantt['Task ID'] == int(predecessors[i])][0])
                start_dates[i] = end_dates[pTask_index]

            else:
                temp_list = []
                for p in predecessors[i].split(','):
                    pTask_index = int(df_gantt.index[df_gantt['Task ID'] == int(p)][0])
                    temp_list.append(end_dates[pTask_index])
                start_dates[i] = max(temp_list)

            # assign end date
            if i == 5:  ## assign vaccine approval time

                date_today = datetime.date.today()
                mon = (date_today.month + approval_month) % 12

                # Assign the year and month
                if mon != 0:  # not the month of Dec
                    year = date_today.year + int((date_today.month + approval_month) / 12)
                    newmon = mon
                    nodaysinapprmonth = (datetime.date(year, newmon + 1, 1) - datetime.date(year, newmon, 1)).days

                else:  # mon=0 means that the approved month is 12
                    year = (date_today.year + int((date_today.month + approval_month) / 12)) - 1
                    newmon = 12
                    nodaysinapprmonth = (datetime.date(year + 1, 1, 1) - datetime.date(year, 12, 1)).days

                # Assign the date
                if date_today.day <= nodaysinapprmonth:
                    target_date = datetime.date(year, newmon, date_today.day)

                else:
                    target_date = datetime.date(year, newmon, nodaysinapprmonth)

                end_dates[i] = target_date
            else:
                end_dates[i] = start_dates[i] + datetime.timedelta(days=time_days[i])

        else:
            end_dates[i] = start_dates[i] + datetime.timedelta(days=time_days[i])

    m_month = []

    date_today = datetime.date.today()

    comm_prod_index = int(df_gantt[df_gantt['Activities'] == 'Start commercial production'].index.values[0])
    dp_prod_index = int(df_gantt[df_gantt['Activities'] == 'Start DP production'].index.values[0])

    m_time = [end_dates[i] for i in [comm_prod_index, dp_prod_index]]

    for t in m_time:
        m_month.append((t.year - date_today.year) * 12 + t.month - date_today.month)

    # apply a minimum 3 months' gap
    if m_month[0] + 3 > m_month[1]: m_month[1] = m_month[0] + 3

    return m_month

# --------------------------------------------------------------------------------

# --------------------------------------------------------------------------------

def jread(filename):
    '''
    Reads JSON files and returns the values.

    Parameters
    ----------
    filename: filename of JSON file to load

    Returns
    ----------
    values: data from the JSON file
    '''

    f = open(filename)
    values = json.load(f)
    f.close()

    return values

# --------------------------------------------------------------------------------

def getValue(values):
    '''
    Returns a single time value for an activity of the schedule based on the
    type of distribution specified.

    Parameters
    ----------
    df_schedule: dataframe of the schedule activity data
    i: row index
    x: value or distribution type

    Returns
    ----------
    computed value in days
    '''
    #### Get distribution
    return {
        'Static': values[1],
        'Triangular': random.triangular(values[2], values[4], values[3])
    }[values[0]]

# --------------------------------------------------------------------------------

def getTarget(df_iteration, df_secCumProduction):
    '''
    Finds the months at which the vaccine dose targets are hit.

    Parameters
    ----------
    df_iteration: a dataframe of the successful vaccines with data processed
        for use in the manufacturing model
    df_secCumProduction: text

    Returns
    ----------
    targetMonth: array of months when each dose target was hit
    df: vaccines for the current iteration, months the target was hit,
    '''
    targetDoses_copy = targetDoses.copy()
    targetMonth = [0, 0, 0, 0]
    # find the target month
    for i in range(1, 101):
        total = df_secCumProduction[i].sum()
        for t in range(4):
            if total > targetDoses_copy[t] and targetMonth[t] == 0:
                targetMonth[t] = i
                targetDoses_copy[t] = total
    df = df_iteration.iloc[:, :7]
    df['Target1 (month)'] = targetMonth[0]
    df['Target2 (month)'] = targetMonth[1]
    df['Target3 (month)'] = targetMonth[2]
    df['Target4 (month)'] = targetMonth[3]

    vaccine_list = df['Vaccine']
    target_dose = {}
    for k in range(4):
        t = []
        for j in range(len(df)):
            if targetMonth[k] > 0:
                t.append(
                    df_secCumProduction.loc[df_secCumProduction['Vaccine'] == vaccine_list[j], targetMonth[k]].sum())
            else:
                t.append(0)
        target_dose['Target' + str(k + 1) + ' (bn doses)'] = t

    df_target = pd.DataFrame(target_dose)
    df_target = pd.concat([df, df_target], axis=1)

    return targetMonth, df_target

# --------------------------------------------------------------------------------

def timeline(df):
    '''
    Creates the data required for the timeline bar chart.

    Parameters
    ----------
    df: dataframe of output summary data from the manufacturing model

    Returns
    ----------
    df: dataframe containing data for the timeline chart
    '''
    if df.empty:
        df = pd.DataFrame(index=['None'],
                          columns=['Approval (month)', 'Platform', 'Primary start time', 'Secondary start time'],
                          data=0)
    else:
        df = df.copy()
        df = df.iloc[:, 0:7]  # pull out columns A to F from dataframe
        df['Primary start time'] = df['Primary Start'].sub(df['Approval (month)']).astype(int)
        df['Secondary start time'] = df['Secondary Start'].sub(df['Primary Start']).astype(int)
        df['Approval (month)'] = df['Approval (month)'].astype(int)

        df = df.drop(columns=['try', 'Vaccine', 'Category', 'Primary Start', 'Secondary Start'])

        df = df.groupby(['Platform']).mean().apply(np.ceil).astype(int)

    return df

# --------------------------------------------------------------------------------

def doseBreakdown(df):
    '''
    Creates the data for the dose breakdown pie charts.

    Parameters
    ----------
    df: dataframe of output summary data from the manufacturing model

    Returns
    ----------
    target1_pie: table of data to create the pie chart for dose target 1
    target2_pie: table of data to create the pie chart for dose target 2
    target3_pie: table of data to create the pie chart for dose target 3
    target4_pie: table of data to create the pie chart for dose target 4
    '''
    #### Get dose breakdown
    if df.empty:
        target1_pie = pd.DataFrame(index=['None'], columns=['Target1 (bn doses)', 'Target 1 %'], data=0)
        target2_pie = pd.DataFrame(index=['None'], columns=['Target2 (bn doses)', 'Target 2 %'], data=0)
        target3_pie = pd.DataFrame(index=['None'], columns=['Target3 (bn doses)', 'Target 3 %'], data=0)
        target4_pie = pd.DataFrame(index=['None'], columns=['Target4 (bn doses)', 'Target 4 %'], data=0)
    else:
        df_copy = df.copy()

        # Create a dataframe for each Target where 0 values have been filtered out
        df_doses = [df_copy[df_copy['Target' + str(i) + ' (bn doses)'] != 0][
                        ['try', 'Platform', 'Target' + str(i) + ' (bn doses)']].reset_index(drop=True) for i in
                    range(1, 5)]

        # Calculate number of tries for each Target above
        tries = [i['try'].nunique() for i in df_doses]

        # Calculate average
        df_grouped = [(df_doses[k].groupby(['Platform'])['Target' + str(k + 1) + ' (bn doses)'].sum()) / tries[k] for k
                      in range(len(df_doses))]

        target1, target2, target3, target4 = [i for i in df_grouped]

        # Calculate the % of each Target
        target1_percent, target2_percent, target3_percent, target4_percent = [df_grouped[i] / df_grouped[i].sum() * 100
                                                                              for i in range(len(df_grouped))]

        df = df.groupby(['Platform'])[
                 ['Target1 (bn doses)', 'Target2 (bn doses)', 'Target3 (bn doses)', 'Target4 (bn doses)']].sum() / df[
                 'try'].nunique()

        target1_pie = pd.DataFrame(data={'Target1 (bn doses)': target1, 'Target 1 %': target1_percent})
        target2_pie = pd.DataFrame(data={'Target2 (bn doses)': target2, 'Target 2 %': target2_percent})
        target3_pie = pd.DataFrame(data={'Target3 (bn doses)': target3, 'Target 3 %': target3_percent})
        target4_pie = pd.DataFrame(data={'Target4 (bn doses)': target4, 'Target 4 %': target4_percent})

    return target1_pie, target2_pie, target3_pie, target4_pie

# --------------------------------------------------------------------------------

def getHistogram(target_no, df):
    '''
    Creates the data for the histogram chart.

    Parameters
    ----------
    target_no: vaccine dose target to generate the histogram for
    df: dataframe of output summary data from the manufacturing model

    Returns
    ----------
    return1: text
    '''
    #### Get histogram
    if df.empty:
        target = pd.DataFrame(np.array([[0, 0, 0.9, 0.75, 0.5]]), index=['None'],
                              columns=['Number of Runs', 'Cumulative %', '90%', '75%', '50%'])
    else:
        df = df.copy()
        # count the number of tries reach target n on a particular month
        target = pd.DataFrame(df.groupby('Target' + str(target_no) + ' (month)')['try'].nunique())

        # remove failed tries (month 0)
        if target.index[0] == 0:
            target.drop(0, inplace=True)

        # calculate the cumulative sum
        target['Cumulative %'] = np.cumsum(target.loc[:, 'try'] / sum(target['try']))

        # add referening line
        target['90%'], target['75%'], target['50%'] = (0.9, 0.75, 0.5)

        # rename the column
        target.rename(index=str, columns={"try": "Number of Runs"}, inplace=True)

    return target

# --------------------------------------------------------------------------------

def cumulativeProduction(target4_hist, cumulative_summary):
    '''
    Creates the data required for the cumulative production chart.

    Parameters
    ----------
    target4_hist: histogram data for the fourth vaccine dose target
    cumulative_summary: cumulative summary of the dose production

    Returns
    ----------
    df_cumulative: dataframe containing data for the cumulative chart
    '''
    #### Get cumulative production
    if cumulative_summary.empty:
        df_cumulative = pd.DataFrame(index=['None'], columns=['Month', '10%', '25%', '50%', '75%', '90%', 'Target 4'],
                                     data=0)
        df_cumulative['Target 4'] = targetDoses[3]
    else:
        df_hist = target4_hist.copy()

        percentage = (0.1, 0.25, 0.5, 0.75, 0.9)
        month = []

        for i in percentage:
            df = df_hist[df_hist['Cumulative %'] > i]
            month.append(int(df.index[0]))

        target4 = targetDoses[3]
        df = cumulative_summary.copy()

        # Define list with column names to sum()
        lst = [i for i in range(1, 101)]

        df = df.groupby(['try'])[lst].sum()

        df_cumulative = pd.DataFrame()  # reset the dataframe

        for i in month:
            if i > 0:
                df1 = df[(df[i - 1] < target4) & (df[i] > target4)].head(1).T
            else:
                df1 = df[df[100] == 0].head(1).T
            df_cumulative = pd.concat([df_cumulative, df1], axis=1)

        df_cumulative['Target 4'] = targetDoses[3]
        df_cumulative.reset_index(inplace=True)

        df_cumulative.columns = ['Month', '10%', '25%', '50%', '75%', '90%', 'Target 4']

    return df_cumulative

# --------------------------------------------------------------------------------

def getOutput():
    '''
    Processes the output data from the manufacturing model by calling
    functions for each individual charts required.

    Parameters
    ----------
    None

    Returns:
    ----------
    timeline_bar: table of data to create the timeline bar chart

    target1_pie: table of data to create the pie chart for dose target 1
    target2_pie: table of data to create the pie chart for dose target 2
    target3_pie: table of data to create the pie chart for dose target 3
    target4_pie: table of data to create the pie chart for dose target 4

    target1_hist: table of data to create the histogram for dose target 1
    target2_hist: table of data to create the histogram for dose target 2
    target3_hist: table of data to create the histogram for dose target 3
    target4_hist: table of data to create the histogram for dose target 4

    cum_line: table of data to create the cumulative % line chart.
    '''

    # get the timeline table
    timeline_bar = timeline(output_summary)

    # get the pie chart tables
    target1_pie, target2_pie, target3_pie, target4_pie = doseBreakdown(output_summary)

    # get the histogram tables
    target1_hist = getHistogram(1, output_summary)
    target2_hist = getHistogram(2, output_summary)
    target3_hist = getHistogram(3, output_summary)
    target4_hist = getHistogram(4, output_summary)

    # get the trendline table and highlights if targets have been met
    if target1_hist.index[0] == 'None':
        cum_line = {'Month': {'None': 0}, '10%': {'None': 0}, '25%': {'None': 0}, '50%': {'None': 0},
                    '75%': {'None': 0}, '90%': {'None': 0}}
        highlights = [0, 0, 0, 0]
    else:
        cum_line = cumulativeProduction(target4_hist, cumulative_summary).to_dict()
        targets_list = [target1_hist, target2_hist, target3_hist, target4_hist]
        highlights = [int(i.index[i['Cumulative %'] >= 0.5][0]) for i in targets_list]

    return timeline_bar.to_dict(), target1_pie.to_dict(), target2_pie.to_dict(), target3_pie.to_dict(), target4_pie.to_dict() \
        , target1_hist.to_dict(), target2_hist.to_dict(), target3_hist.to_dict(), target4_hist.to_dict(), cum_line, highlights

# --------------------------------------------------------------------------------

def runTrial(trialData):
    '''
    This function takes the trial output from the R&D model and runs the main
    functions of the manufacturing model.

    Parameters
    ----------
    trialData: iteration trial data from the R&D model

    Returns
    ----------
    None
    '''

    global output_summary
    global cumulative_summary

    if bool(trialData):
        ## get the iteration table for this try
        df_iteration = getIteration(trialData)
        ## append the primary and secondary manufacturing start time to the iteration table
        getManufacturingStartTime(df_iteration)

        ## get the primary allocation and throughput
        df_priAllocation = primary(df_iteration)

        ## get the secondary allocation and throughput, as well as cumulative total
        df_secCumProduction = secondary(df_iteration, df_priAllocation, )

        ## get the output for this try
        targetMonth, output_mfg = getTarget(df_iteration, df_secCumProduction)

        # merge the output table
        output_summary = output_summary.append(output_mfg, ignore_index=True)
        cumulative_summary = cumulative_summary.append(df_secCumProduction, ignore_index=True)

# --------------------------------------------------------------------------------
def getScheduleInput():
    '''
    Modifies the default Gantt chart to create a Gantt chart for each platform.

    Parameters
    ----------
    None

    Returns:
    ----------
    df_schedule:

    funding: a dictionary containing relevant information for each funding type
    '''

    # the tables below are used for value updating (static, low, most likely, high)
    gantt = pd.DataFrame(default.get('Generic Gantt')).transpose()
    parameter_v = default.get('Gantt Timelines (default)')
    parameter = m_params.get('Gantt Timelines')

    # the dictionary below is based on funding category
    funding = m_params.get('Timelines by funding criteria')

    ##############################################################################
    # Update the values of other activities to the default gantt, once only
    ranges = ['Low', 'Most Likely', 'High']

    # apply changes to all platforms
    low_d, most_likely_d, high_d = [parameter_v['All Platform'][x] for x in ranges]
    low, most_likely, high = [parameter['All Platform'][x] for x in ranges]

    low_r = low / low_d
    most_likely_r = most_likely / most_likely_d
    high_r = high / high_d

    # Convert numerical values into integers
    gantt[['Value', 'Low', 'Most Likely', 'High']] = gantt[['Value', 'Low', 'Most Likely', 'High']].astype(int)

    mask = (gantt['Activities'] != 'Scale up and and process development') & (
            gantt['Activities'] != 'Technology transfer') & (
                   gantt['Activities'] != 'DP technology transfer')

    gantt.loc[mask, 'Value'] = round(most_likely_r * gantt.loc[mask, 'Value'], 0).astype(int)
    gantt.loc[mask, 'Low'] = round(low_r * gantt.loc[mask, 'Low'], 0).astype(int)
    gantt.loc[mask, 'Most Likely'] = round(most_likely_r * gantt.loc[mask, 'Most Likely'], 0).astype(int)
    gantt.loc[mask, 'High'] = round(high_r * gantt.loc[mask, 'High'], 0).astype(int)

    gantt_updated = gantt

    del parameter['All Platform']
    ##############################################################################

    # get the iteration table for different platforms
    platforms = parameter.keys()

    # gerenate the gantt for different platform
    df_schedule = {}

    for i in platforms:
        gantt_updated_copy = gantt_updated.copy()
        low_d, most_likely_d, high_d = [parameter_v[i][x] for x in ranges]
        low, most_likely, high = [parameter[i][x] for x in ranges]

        low_r = low / low_d
        most_likely_r = most_likely / most_likely_d
        high_r = high / high_d

        mask = (gantt_updated_copy['Activities'] == 'Scale up and and process development') | (
                gantt_updated_copy['Activities'] == 'Technology transfer') | (
                       gantt_updated_copy['Activities'] == 'DP technology transfer')

        gantt_updated_copy.loc[mask, 'Low'] = round(low_r * gantt_updated_copy.loc[mask, 'Low'], 0).astype(int)
        gantt_updated_copy.loc[mask, 'Most Likely'] = round(most_likely_r * gantt_updated_copy.loc[mask, 'Most Likely'],
                                                            0).astype(int)
        gantt_updated_copy.loc[mask, 'High'] = round(high_r * gantt_updated_copy.loc[mask, 'High'], 0).astype(int)

        df_schedule[i] = gantt_updated_copy

    return df_schedule, funding

##############################################################################


# Copyright (C) 2021 Sarah Roggendorf and Jørgen S. Dokken
#
# SPDX-License-Identifier:    MIT

from typing import Dict, Tuple

import basix
import dolfinx_cuas
import numpy as np
import ufl
from dolfinx import common as _common
from dolfinx import fem as _fem
from dolfinx import log as _log
from dolfinx import mesh as dmesh
from dolfinx.graph import create_adjacencylist
from petsc4py import PETSc as _PETSc
from dolfinx.cpp.mesh import MeshTags_int32
import dolfinx_contact
import dolfinx_contact.cpp
from dolfinx_contact.helpers import (epsilon, lame_parameters,
                                     rigid_motions_nullspace, sigma_func)

__all__ = ["nitsche_custom"]

def nitsche_custom(mesh: dmesh.Mesh, mesh_data: Tuple[MeshTags_int32, int, int],
                   physical_parameters: dict = {}, nitsche_parameters: Dict[str, float] = {},
                   plane_loc: float = 0.0, vertical_displacement: float = -0.1,
                   nitsche_bc: bool = True, quadrature_degree: int = 5, form_compiler_params: Dict = {},
                   jit_params: Dict = {}, petsc_options: Dict = {}, newton_options: Dict = {}) -> _fem.Function:
    """
    Use custom kernel to compute the one sided contact problem with a mesh coming into contact
    with a rigid surface (not meshed).

    Parameters
    ==========
    mesh
        The input mesh
    mesh_data
        A triplet with a mesh tag for facets and values v0, v1. v0 should be the value in the mesh tags
        for facets to apply a Dirichlet condition on. v1 is the value for facets which should have applied
        a contact condition on
    physical_parameters
        Optional dictionary with information about the linear elasticity problem.
        Valid (key, value) tuples are: ('E': float), ('nu', float), ('strain', bool)
    nitsche_parameters
        Optional dictionary with information about the Nitsche configuration.
        Valid (keu, value) tuples are: ('gamma', float), ('theta', float) where theta can be -1, 0 or 1 for
        skew-symmetric, penalty like or symmetric enforcement of Nitsche conditions
    plane_loc
        The location of the plane in y-coordinate (2D) and z-coordinate (3D)
    vertical_displacement
        The amount of verticial displacment enforced on Dirichlet boundary
    nitsche_bc
        Use Nitche's method to enforce Dirichlet boundary conditions
    quadrature_degree
        The quadrature degree to use for the custom contact kernels
    form_compiler_params
        Parameters used in FFCX compilation of this form. Run `ffcx --help` at
        the commandline to see all available options. Takes priority over all
        other parameter values, except for `scalar_type` which is determined by
        DOLFINX.
    jit_params
        Parameters used in CFFI JIT compilation of C code generated by FFCX.
        See https://github.com/FEniCS/dolfinx/blob/main/python/dolfinx/jit.py
        for all available parameters. Takes priority over all other parameter values.
    petsc_options
        Parameters that is passed to the linear algebra backend
        PETSc. For available choices for the 'petsc_options' kwarg,
        see the `PETSc-documentation
        <https://petsc4py.readthedocs.io/en/stable/manual/ksp/>`
    newton_options
        Dictionary with Newton-solver options. Valid (key, item) tuples are:
        ("atol", float), ("rtol", float), ("convergence_criterion", "str"),
        ("max_it", int), ("error_on_nonconvergence", bool), ("relaxation_parameter", float)
    """
    # Compute lame parameters
    plane_strain = physical_parameters.get("strain", False)
    E = physical_parameters.get("E", 1e3)
    nu = physical_parameters.get("nu", 0.1)
    mu_func, lambda_func = lame_parameters(plane_strain)
    mu = mu_func(E, nu)
    lmbda = lambda_func(E, nu)
    sigma = sigma_func(mu, lmbda)

    # Nitche parameters and variables
    theta = nitsche_parameters.get("theta", 1)
    gamma = nitsche_parameters.get("gamma", 1)

    # Unpack mesh data
    (facet_marker, dirichlet_value, contact_value) = mesh_data
    assert(facet_marker.dim == mesh.topology.dim - 1)

    # Outward unit normal of plane
    n_vec = np.zeros(mesh.geometry.dim)
    n_vec[mesh.geometry.dim - 1] = 1

    # Setup function space and functions used in Jacobian and residual formulation
    V = _fem.VectorFunctionSpace(mesh, ("CG", 1))
    u = _fem.Function(V)
    v = ufl.TestFunction(V)
    du = ufl.TrialFunction(V)
    u = _fem.Function(V)
    v = ufl.TestFunction(V)

    # Compute classical (volume) contributions of the equations of linear elasticity
    dx = ufl.Measure("dx", domain=mesh)
    J = ufl.inner(sigma(du), epsilon(v)) * dx
    F = ufl.inner(sigma(u), epsilon(v)) * dx

    # Nitsche for Dirichlet
    # https://doi.org/10.1016/j.cma.2018.05.024
    if nitsche_bc:
        ds = ufl.Measure("ds", domain=mesh, subdomain_data=facet_marker)
        h = ufl.Circumradius(mesh)
        n = ufl.FacetNormal(mesh)
        disp_vec = np.zeros(mesh.geometry.dim)
        disp_vec[mesh.geometry.dim - 1] = vertical_displacement
        u_D = ufl.as_vector(disp_vec)
        F += - ufl.inner(sigma(u) * n, v) * ds(dirichlet_value)\
             - theta * ufl.inner(sigma(v) * n, u - u_D) * \
            ds(dirichlet_value) + E * gamma / h * ufl.inner(u - u_D, v) * ds(dirichlet_value)
        J += - ufl.inner(sigma(du) * n, v) * ds(dirichlet_value)\
            - theta * ufl.inner(sigma(v) * n, du) * \
            ds(dirichlet_value) + E * gamma / h * ufl.inner(du, v) * ds(dirichlet_value)
    else:
        raise RuntimeError("Dirichlet bc not implemented in custom assemblers yet.")

    # Custom assembly of contact boundary condition
    q_rule = dolfinx_contact.QuadratureRule(mesh.topology.cell_type, quadrature_degree,
                                            mesh.topology.dim - 1, basix.QuadratureType.Default)
    consts = np.array([E * gamma, theta])
    consts = np.hstack((consts, n_vec))

    # Compute coefficients for mu and lambda as DG-0 functions
    V2 = _fem.FunctionSpace(mesh, ("DG", 0))
    lmbda2 = _fem.Function(V2)
    lmbda2.interpolate(lambda x: np.full((1, x.shape[1]), lmbda))
    mu2 = _fem.Function(V2)
    mu2.interpolate(lambda x: np.full((1, x.shape[1]), mu))

    # Compute integral entities on exterior facets (cell_index, local_index)
    bottom_facets = facet_marker.indices[facet_marker.values == contact_value]
    integral = _fem.IntegralType.exterior_facet
    integral_entities = dolfinx_contact.compute_active_entities(mesh, bottom_facets, integral)
    # Pack mu and lambda on facets
    coeffs = dolfinx_cuas.pack_coefficients([mu2, lmbda2], integral_entities)
    # Pack circumradius of facets
    h_facets = dolfinx_contact.pack_circumradius(mesh, integral_entities)

    # Create contact class
    data = np.array([contact_value, dirichlet_value], dtype=np.int32)
    offsets = np.array([0, 2], dtype=np.int32)
    surfaces = create_adjacencylist(data, offsets)
    contact = dolfinx_contact.cpp.Contact([facet_marker], surfaces, [(0, 1)],
                                          V._cpp_object, quadrature_degree=quadrature_degree)
    # Compute gap from contact boundary
    g_vec = contact.pack_gap_plane(0, -plane_loc)

    # Concatenate coefficients
    coeffs = np.hstack([coeffs, h_facets, g_vec])

    # Create RHS kernels
    L_custom = _fem.form(F, jit_params=jit_params, form_compiler_params=form_compiler_params)
    kernel_rhs = dolfinx_contact.cpp.generate_contact_kernel(V._cpp_object, dolfinx_contact.Kernel.Rhs, q_rule,
                                                             [u._cpp_object, mu2._cpp_object, lmbda2._cpp_object])

    def create_b():
        return _fem.petsc.create_vector(L_custom)

    def assemble_residual(x, b):
        u.vector[:] = x.array
        u_packed = dolfinx_cuas.pack_coefficients([u._cpp_object], integral_entities)
        c = np.hstack([u_packed, coeffs])
        contact.assemble_vector(b, 0, kernel_rhs, c, consts)
        _fem.petsc.assemble_vector(b, L_custom)

    # Create Jacobian kernels
    a_custom = _fem.form(J, jit_params=jit_params, form_compiler_params=form_compiler_params)
    kernel_J = dolfinx_contact.cpp.generate_contact_kernel(
        V._cpp_object, dolfinx_contact.Kernel.Jac, q_rule, [u._cpp_object, mu2._cpp_object, lmbda2._cpp_object])

    def create_A():
        return _fem.petsc.create_matrix(a_custom)

    def assemble_jacobian(x, A):
        u.vector[:] = x.array
        u_packed = dolfinx_cuas.pack_coefficients([u._cpp_object], integral_entities)
        c = np.hstack([u_packed, coeffs])
        contact.assemble_matrix(A, [], 0, kernel_J, c, consts)
        _fem.petsc.assemble_matrix(A, a_custom)

    # Setup non-linear problem and Newton-solver
    problem = dolfinx_cuas.NonlinearProblemCUAS(assemble_residual, assemble_jacobian, create_b, create_A)
    solver = dolfinx_cuas.NewtonSolver(mesh.comm, problem)

    # Create rigid motion null-space
    null_space = rigid_motions_nullspace(V)
    solver.A.setNearNullSpace(null_space)

    # Set Newton solver options
    solver.atol = newton_options.get("atol", 1e-9)
    solver.rtol = newton_options.get("rtol", 1e-9)
    solver.convergence_criterion = newton_options.get("convergence_criterion", "incremental")
    solver.max_it = newton_options.get("max_it", 50)
    solver.error_on_nonconvergence = newton_options.get("error_on_nonconvergence", True)
    solver.relaxation_parameter = newton_options.get("relaxation_parameter", 1.0)

    def _u_initial(x):
        values = np.zeros((mesh.geometry.dim, x.shape[1]))
        values[-1] = -0.01 - plane_loc
        return values

    # Set initial_condition:
    u.interpolate(_u_initial)

    # Define solver and options
    ksp = solver.krylov_solver
    option_prefix = ksp.getOptionsPrefix()

    # Set PETSc options
    opts = _PETSc.Options()
    opts.prefixPush(option_prefix)
    for k, v in petsc_options.items():
        opts[k] = v
    opts.prefixPop()
    ksp.setFromOptions()

    dofs_global = V.dofmap.index_map_bs * V.dofmap.index_map.size_global
    _log.set_log_level(_log.LogLevel.INFO)

    # Solve non-linear problem
    with _common.Timer(f"{dofs_global} Solve Nitsche"):
        n, converged = solver.solve(u)
    u.x.scatter_forward()

    if solver.error_on_nonconvergence:
        assert(converged)
    print(f"{dofs_global}, Number of interations: {n:d}")

    return u


"""
The classes in this file do preprocessing on data and monte carlo to be used
to do a point source analysis.
"""

__author__ = 'John Evans'
__copyright__ = 'Copyright 2020 John Evans'
__credits__ = ['John Evans', 'Jason Fan', 'Michael Larson']
__license__ = 'Apache License 2.0'
__version__ = '0.0.1'
__maintainer__ = 'John Evans'
__email__ = 'john.evans@icecube.wisc.edu'
__status__ = 'Development'

from typing import Optional, Tuple, Union
import numpy as np
import numpy.lib.recfunctions as rf
from scipy.interpolate import UnivariateSpline as Spline

from dataclasses import dataclass
from dataclasses import field
from dataclasses import InitVar

from .. import sources
from .. import _models
from . import spectral

@dataclass
class _ThreeMLEventModelBase(_models.EventModelBase):
    """Docstring"""
    _sin_dec_bins: np.array = field(init=False)
    _log_energy_bins: np.array = field(init=False)
    _edge_point: Tuple[float, float] = field(init=False)
    _background_sob_map: np.ndarray = field(init=False)
    _ratio: np.ndarray = field(init=False)
    _reduced_sim_reconstructed: np.ndarray = field(init=False)

@dataclass
class _ThreeMLEventModelDefaultsBase(_models.TdEventModelDefaultsBase):
    """Docstring"""
    signal_sin_dec_bins: InitVar[Union[np.array, int]] = field(default=50)
    log_energy_bins: InitVar[Union[np.array, int]] = field(default=50)
    _spectrum: spectral.BaseSpectrum = field(
        default=spectral.PowerLaw(1e3, 1e-14, -2))

@dataclass
class ThreeMLEventModel(
    _models.TdEventModel,
    _ThreeMLEventModelDefaultsBase,
    _ThreeMLEventModelBase,
):
    """Docstring"""
    def __post_init__(
        self,
        source: sources.Source,
        data: np.ndarray,
        sim: np.ndarray,
        grl: np.ndarray,
        gamma: float,
        sampling_width: Optional[float],
        background_sin_dec_bins: Union[np.array, int],
        background_window: float,
        withinwindow: bool,
        signal_sin_dec_bins: Union[np.array, int],
        log_energy_bins: Union[np.array, int],
    ) -> None:
        """
        Args:
            source:
            grl:
            background_sin_dec_bins: If an int, then the number of bins
                spanning -1 -> 1, otherwise, a numpy array of bin edges.
            background_window:
            withinwindow:
        """
        super().__post_init__(
            source,
            data,
            sim,
            grl,
            gamma,
            sampling_width,
            background_sin_dec_bins,
            background_window,
            withinwindow,
        )
        if isinstance(signal_sin_dec_bins, int):
            signal_sin_dec_bins = np.linspace(-1, 1, 1 + signal_sin_dec_bins)
        self._sin_dec_bins = signal_sin_dec_bins

        if isinstance(log_energy_bins, int):
            log_energy_bins = np.linspace(1, 8, 1 + log_energy_bins)

        self._log_energy_bins = log_energy_bins
        self._background_sob_map = self._init_background_sob_map()
        self._init_reduced_sim_reconstructed(source)
        self._ratio = self._init_sob_ratio()

    def _init_background_sob_map(self) -> None:
        """Create the backgroub SOB map
        """
        # background
        bins = np.array([self._sin_dec_bins, self._log_energy_bins])
        bg_h, _, _ = np.histogram2d(self._data['sindec'], self._data['logE'],
                                    bins=bins, density=True)
        with np.errstate(divide='ignore', invalid='ignore'):
            bg_h /= np.sum(bg_h, axis=1)[:, None]
        return bg_h

    def _init_sob_ratio(self, *args, **kwargs) -> None:
        """Create the SOB map with a spectrum
        """
        bins = np.array([self._sin_dec_bins, self._log_energy_bins])
        bin_centers = bins[1, :-1] + np.diff(bins[1]) / 2
        sig_w = self._reduced_sim_reconstructed['ow'] * self._spectrum(
            self._reduced_sim_reconstructed['trueE'])
        sig_h, _, _ = np.histogram2d(self._reduced_sim_reconstructed['sindec'],
                                     self._reduced_sim_reconstructed['logE'],
                                     bins=bins, weights=sig_w, density=True)

        # Normalize histograms by dec band
        with np.errstate(divide='ignore', invalid='ignore'):  # divide warnings
            sig_h /= np.sum(sig_h, axis=1)[:, None]

        if 'k' not in kwargs:
            kwargs['k'] = 1
        if 's' not in kwargs:
            kwargs['s'] = 0
        if 'ext' not in kwargs:
            kwargs['ext'] = 3

        with np.errstate(divide='ignore', invalid='ignore'):  # divide warnings
            ratio = sig_h / self._background_sob_map

        with np.errstate(divide='ignore', invalid='ignore'):  # divide warnings
            for i in range(ratio.shape[0]):
                # Pick out the values we want to use.
                # We explicitly want to avoid NaNs and infinities
                values = ratio[i]
                good = np.isfinite(values) & (values > 0)
                x_good, y_good = bin_centers[good], values[good]

                # Do a linear interpolation across the energy range
                if len(x_good) > 1:
                    spline = Spline(x_good, y_good, *args, **kwargs)
                    ratio[i] = spline(bin_centers)
                elif len(x_good) == 1:
                    ratio[i] = y_good
                else:
                    ratio[i] = 0
        return ratio

    def _init_reduced_sim_reconstructed(self, source: sources.Source) -> None:
        """Gets a small simulation dataset to use for injecting signal.

        Prunes the simulation set to only events close to a given source and
        calculate the weight for each event. Adds the weights as a new column
        to the simulation set.

        Args:
            source:

        Returns:
            A reweighted simulation set around the source declination.
        """
        if self._sampling_width is not None:
            self._cut_sim_reconstructed(source)
        else:
            self._reduced_sim_reconstructed = self._sim

    def _cut_sim_reconstructed(self, source: sources.Source) -> np.ndarray:
        """Select simulation events in a reconstruction dec band

        Args:
            source:
        """
        if self._sampling_width is not None:
            self._edge_point = (np.searchsorted(
                                self._sin_dec_bins,
                                np.sin(source.dec - self._sampling_width)) - 1,
                                np.searchsorted(
                                self._sin_dec_bins,
                                np.sin(source.dec + self._sampling_width)) - 1)
        else:
            self._edge_point = (self._sin_dec_bins[0], self._sin_dec_bins[-1])
        sindec_dist = np.abs(source.dec - self._sim['dec'])
        close = sindec_dist < self._sampling_width
        self._reduced_sim_reconstructed = self._sim[close].copy()

    def _weight_reduced_sim(self, reduced_sim: np.ndarray) -> np.ndarray:
        """Docstring"""
        try:
            reduced_sim = rf.append_fields(reduced_sim, 'weight',
                                           np.zeros(len(reduced_sim)),
                                           dtypes=np.float32)

        except ValueError:  # weight already exist
            pass

        # Assign the weights using the newly defined "time profile"
        # classes above. If you want to make this a more complicated
        # shape, talk to me and we can work it out.
        reduced_sim['weight'] = reduced_sim['ow'] * self._spectrum(
            reduced_sim['trueE'])
        return reduced_sim

    def reweight_reduced_sim(self, spectrum: spectral.BaseSpectrum):
        """Docstring"""
        self._reduced_sim['weight'] = self._reduced_sim['ow'] * spectrum(
            self._reduced_sim['trueE'])

    def prepro_index(self, events: np.ndarray) -> np.ndarray:
        """Find the sindec index and energy index for events

        More function info...

        Args:
            events: An array of events including their positional data.

        Returns:
            A list of index
        """
        # Get the bin that each event belongs to
        try:
            sin_dec_idx = np.searchsorted(self._sin_dec_bins[:-1],
                                          events['sindec']) - 1
        except ValueError:
            sin_dec_idx = np.searchsorted(self._sin_dec_bins[:-1],
                                          np.sin(events['dec'])) - 1

        log_energy_idx = np.searchsorted(self._log_energy_bins[:-1],
                                         events['logE']) - 1

        sin_dec_idx[sin_dec_idx < self._edge_point[0]] = self._edge_point[0]
        # If events fall outside the sampling width, just gonna approxiamte the
        # weight using the nearest non-zero sinDec bin.
        sin_dec_idx[sin_dec_idx > self._edge_point[1]] = self._edge_point[1]
        return sin_dec_idx, log_energy_idx

    def _energy_sob(
        self,
        sin_dec_idx: np.ndarray,
        log_energy_idx: np.ndarray
    ) -> np.ndarray:
        """Gets the sob vs. gamma required for each event and specific .

        More function info...

        Args:
            sin_dec_idx: An array of sin dec index of events
            log_energy_idx: An array of log energy index of events

        Returns:
            signal-over-background for each event.
        """
        return self._ratio[sin_dec_idx, log_energy_idx]

    def get_ns(self, livetime: float) -> float:
        """Gets expected number of neutrino

        More function info...

        Args:
            livetime: livetime

        Returns:
            expected number of neutrino
        """
        ns = (
            self._spectrum(self._reduced_sim['trueE'])
            * self._reduced_sim['ow']
            * livetime
        ).sum()
        return ns

    def get_sob_energy(
        self,
        sin_dec_idx: np.ndarray,
        log_energy_idx: np.ndarray,
    ) -> np.ndarray:
        """Docstring"""
        return self._energy_sob(sin_dec_idx, log_energy_idx)

    @property
    def edge_point(self) -> Tuple[float, float]:
        """Docstring"""
        return self._edge_point

    @property
    def spectrum(self) -> spectral.BaseSpectrum:
        """Docstring"""
        return self._spectrum

    @spectrum.setter
    def spectrum(self, spectrum: spectral.BaseSpectrum):
        """Docstring"""
        self._spectrum = spectrum


#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. 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.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
CISD analytical nuclear gradients
'''

import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.ci import cisd
from pyscf.grad import rhf as rhf_grad
from pyscf.grad import ccsd as ccsd_grad

def grad_elec(cigrad, civec=None, eris=None, atmlst=None, verbose=logger.INFO):
    myci = cigrad.base
    if civec is None: civec = myci.ci
    assert(not isinstance(civec, (list, tuple)))
    nocc = myci.nocc
    nmo = myci.nmo
    d1 = cisd._gamma1_intermediates(myci, civec, nmo, nocc)
    fd2intermediate = lib.H5TmpFile()
    d2 = cisd._gamma2_outcore(myci, civec, nmo, nocc, fd2intermediate, True)
    t1 = t2 = l1 = l2 = civec
    return ccsd_grad.grad_elec(cigrad, t1, t2, l1, l2, eris, atmlst, d1, d2, verbose)

def as_scanner(grad_ci, state=0):
    '''Generating a nuclear gradients scanner/solver (for geometry optimizer).

    The returned solver is a function. This function requires one argument
    "mol" as input and returns total CISD energy.

    The solver will automatically use the results of last calculation as the
    initial guess of the new calculation.  All parameters assigned in the
    CISD and the underlying SCF objects (conv_tol, max_memory etc) are
    automatically applied in the solver.

    Note scanner has side effects.  It may change many underlying objects
    (_scf, with_df, with_x2c, ...) during calculation.

    Examples::

    >>> from pyscf import gto, scf, ci
    >>> mol = gto.M(atom='H 0 0 0; F 0 0 1')
    >>> ci_scanner = ci.CISD(scf.RHF(mol)).nuc_grad_method().as_scanner()
    >>> e_tot, grad = ci_scanner(gto.M(atom='H 0 0 0; F 0 0 1.1'))
    >>> e_tot, grad = ci_scanner(gto.M(atom='H 0 0 0; F 0 0 1.5'))
    '''
    from pyscf import gto
    if isinstance(grad_ci, lib.GradScanner):
        return grad_ci

    logger.info(grad_ci, 'Create scanner for %s', grad_ci.__class__)

    class CISD_GradScanner(grad_ci.__class__, lib.GradScanner):
        def __init__(self, g):
            lib.GradScanner.__init__(self, g)

        def __call__(self, mol_or_geom, state=state, **kwargs):
            if isinstance(mol_or_geom, gto.Mole):
                mol = mol_or_geom
            else:
                mol = self.mol.set_geom_(mol_or_geom, inplace=False)

            ci_scanner = self.base
            if ci_scanner.nroots > 1 and state >= ci_scanner.nroots:
                raise ValueError('State ID greater than the number of CISD roots')

            mf_scanner = ci_scanner._scf
            mf_scanner(mol)
            ci_scanner.mo_coeff = mf_scanner.mo_coeff
            ci_scanner.mo_occ = mf_scanner.mo_occ

            if getattr(ci_scanner.ci, 'size', 0) != ci_scanner.vector_size():
                ci_scanner.ci = None
            eris = ci_scanner.ao2mo(ci_scanner.mo_coeff)
            ci_scanner.kernel(ci0=ci_scanner.ci, eris=eris)

# TODO: Check root flip
            if ci_scanner.nroots > 1:
                e_tot = ci_scanner.e_tot[state]
                civec = ci_scanner.ci[state]
            else:
                e_tot = ci_scanner.e_tot
                civec = ci_scanner.ci

            self.mol = mol
            de = self.kernel(civec, eris=eris, **kwargs)
            return e_tot, de
        @property
        def converged(self):
            ci_scanner = self.base
            if ci_scanner.nroots > 1:
                ci_conv = ci_scanner.converged[state]
            else:
                ci_conv = ci_scanner.converged
            return all((ci_scanner._scf.converged, ci_conv))

    # cache eris object in CCSD base class. eris object is used many times
    # when calculating gradients
    g_ao2mo = grad_ci.base.__class__.ao2mo
    def _save_eris(self, *args, **kwargs):
        self._eris = g_ao2mo(self, *args, **kwargs)
        return self._eris
    grad_ci.base.__class__.ao2mo = _save_eris

    return CISD_GradScanner(grad_ci)

class Gradients(rhf_grad.GradientsBasics):
    def __init__(self, myci):
        self.state = 0  # of which the gradients to be computed.
        rhf_grad.GradientsBasics.__init__(self, myci)

    def dump_flags(self, verbose=None):
        log = logger.new_logger(self, verbose)
        log.info('\n')
        if not self.base.converged:
            log.warn('Ground state %s not converged',
                     self.base.__class__.__name__)
        log.info('******** %s for %s ********',
                 self.__class__, self.base.__class__)
        if self.state != 0 and self.base.nroots > 1:
            log.info('State ID = %d', self.state)
        return self

    grad_elec = grad_elec

    def kernel(self, civec=None, eris=None, atmlst=None, state=None,
               verbose=None):
        log = logger.new_logger(self, verbose)
        myci = self.base
        if civec is None: civec = myci.ci
        if civec is None: civec = myci.kernel(eris=eris)
        if (isinstance(civec, (list, tuple)) or
            (isinstance(civec, numpy.ndarray) and civec.ndim > 1)):
            if state is None:
                state = self.state
            else:
                self.state = state

            civec = civec[state]
            logger.info(self, 'Multiple roots are found in CISD solver. '
                        'Nuclear gradients of root %d are computed.', state)

        if atmlst is None:
            atmlst = self.atmlst
        else:
            self.atmlst = atmlst

        if self.verbose >= logger.WARN:
            self.check_sanity()
        if self.verbose >= logger.INFO:
            self.dump_flags()

        de = self.grad_elec(civec, eris, atmlst, verbose=log)
        self.de = de + self.grad_nuc(atmlst=atmlst)
        if self.mol.symmetry:
            self.de = self.symmetrize(self.de, atmlst)
        self._finalize()
        return self.de

    # Calling the underlying SCF nuclear gradients because it may be modified
    # by external modules (e.g. QM/MM, solvent)
    def grad_nuc(self, mol=None, atmlst=None):
        mf_grad = self.base._scf.nuc_grad_method()
        return mf_grad.grad_nuc(mol, atmlst)

    def _finalize(self):
        if self.verbose >= logger.NOTE:
            logger.note(self, '--------- %s gradients for state %d ----------',
                        self.base.__class__.__name__, self.state)
            self._write(self.mol, self.de, self.atmlst)
            logger.note(self, '----------------------------------------------')

    as_scanner = as_scanner

Grad = Gradients

cisd.CISD.Gradients = lib.class_as_method(Gradients)

if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf

    mol = gto.M(
        atom = [
            ["O" , (0. , 0.     , 0.)],
            [1   , (0. ,-0.757  , 0.587)],
            [1   , (0. , 0.757  , 0.587)]],
        basis = '631g'
    )
    mf = scf.RHF(mol)
    ehf = mf.scf()

    myci = cisd.CISD(mf)
    myci.kernel()
    g1 = myci.Gradients().kernel()
# O     0.0000000000    -0.0000000000     0.0065498854
# H    -0.0000000000     0.0208760610    -0.0032749427
# H    -0.0000000000    -0.0208760610    -0.0032749427
    print(lib.finger(g1) - -0.032562200777204092)

    mcs = myci.as_scanner()
    mol.set_geom_([
            ["O" , (0. , 0.     , 0.001)],
            [1   , (0. ,-0.757  , 0.587)],
            [1   , (0. , 0.757  , 0.587)]])
    e1 = mcs(mol)
    mol.set_geom_([
            ["O" , (0. , 0.     ,-0.001)],
            [1   , (0. ,-0.757  , 0.587)],
            [1   , (0. , 0.757  , 0.587)]])
    e2 = mcs(mol)
    print(g1[0,2] - (e1-e2)/0.002*lib.param.BOHR)

    print('-----------------------------------')
    mol = gto.M(
        atom = [
            ["O" , (0. , 0.     , 0.)],
            [1   , (0. ,-0.757  , 0.587)],
            [1   , (0. , 0.757  , 0.587)]],
        basis = '631g'
    )
    mf = scf.RHF(mol)
    ehf = mf.scf()

    myci = cisd.CISD(mf)
    myci.frozen = [0,1,10,11,12]
    myci.max_memory = 1
    myci.kernel()
    g1 = Gradients(myci).kernel()
# O    -0.0000000000     0.0000000000     0.0106763547
# H     0.0000000000    -0.0763194988    -0.0053381773
# H     0.0000000000     0.0763194988    -0.0053381773
    print(lib.finger(g1) - 0.1022427304650084)

    mcs = myci.as_scanner()
    mol.set_geom_([
            ["O" , (0. , 0.     , 0.001)],
            [1   , (0. ,-0.757  , 0.587)],
            [1   , (0. , 0.757  , 0.587)]])
    e1 = mcs(mol)
    mol.set_geom_([
            ["O" , (0. , 0.     ,-0.001)],
            [1   , (0. ,-0.757  , 0.587)],
            [1   , (0. , 0.757  , 0.587)]])
    e2 = mcs(mol)
    print(g1[0,2] - (e1-e2)/0.002*lib.param.BOHR)

    mol = gto.M(
        atom = 'H 0 0 0; H 0 0 1.76',
        basis = '631g',
        unit='Bohr')
    mf = scf.RHF(mol).run(conv_tol=1e-14)
    myci = cisd.CISD(mf)
    myci.conv_tol = 1e-10
    myci.kernel()
    g1 = Gradients(myci).kernel()
#[[ 0.          0.         -0.07080036]
# [ 0.          0.          0.07080036]]