from numpy import exp

from . helpers import lbm_dubois, ffm_alsallami
from . model import Model

"""
opentiva.propofol
=================

This module contains the classes for propofol models.

All Classes have the same following parameters and attributes.

Parameters
----------
sex
    0 for male or 1 for female
age
    in years
weight
    in kg
height
   in cm

Additional
~~~~~~~~~~
opiates_coadministered
   bool; true if opiates co-administered, Eleveld model only
bolus_data
   bool; true for bolus data, Schuttler model only
venous_data
   bool; true for venous data, Schuttler model only

Attributes
----------
compartments : int
    number of compartments to model; 1, 2 or 3
concentration_unit : str
    drug concentration unit
target_unit : str
    target concentration unit
age_lower : float
    lower age limit of model; -1 if no limit
age_upper : float
    upper age limit of model; -1 if no limit
weight_lower : float
    lower weight limit of model; -1 if no limit
weight_upper : float
    upper weight limit of model; -1 if no limit
pmid : str
    Pubmed ID of model's reference
doi : str
    Digital Object Identifier (DOI) of model's reference
warning : str
    Warnings relating to non-validated anthropometric values
v1 : float
    volume of central compartment
k10 : float
    equilibrium rate constant from compartment 1 to 0
k12 : float
    equilibrium rate constant from compartment 1 to 2
k13 : float
    equilibrium rate constant from compartment 1 to 3
k21 : float
    equilibrium rate constant from compartment 2 to 1
k31 : float
    equilibrium rate constant from compartment 3 to 1
ke0 : float
    effect compartment equilibrium rate constant
"""

class MarshDiprifusor(Model):
    """MarshDiprifusor class holds pharmacokinetic parameters for the
       Diprifusor Marsh propofol model with Keo 0.26.

       Reference: PMID: 1859758 DOI: 10.1093/bja/67.1.41
    """

    def __init__(self, sex: int, age: float, weight: float, height: float):
        super().__init__(sex, age, weight, height)

        self.compartments = 3
        self.concentration_unit = "mg/ml"
        self.target_unit = "ug/ml"
        self.age_lower = 16
        self.age_upper = -1
        self.weight_lower = -1
        self.weight_upper = 150
        self.pmid = "1859758"
        self.doi = "10.1093/bja/67.1.41"
        self.validate_anthropometric_values()

        self.v1 = 0.228 * weight
        self.v2 = 0.463 * weight
        self.v3 = 2.893 * weight

        self.k10 = 0.119
        self.k12 = 0.112
        self.k13 = 0.0419
        self.k21 = 0.055
        self.k31 = 0.0033
        self.ke0 = 0.26

class MarshModified(Model):
    """MarshModified class holds pharmacokinetic parameters for the Modified
       Marsh propofol model with  Keo 1.2.

       Reference: PMID: 1859758 DOI: 10.1093/bja/67.1.41
    """

    def __init__(self, sex: int, age: float, weight: float, height: float):
        super().__init__(sex, age, weight, height)

        self.compartments = 3
        self.concentration_unit = "mg/ml"
        self.target_unit = "ug/ml"
        self.age_lower = 16
        self.age_upper = -1
        self.weight_lower = -1
        self.weight_upper = 150
        self.pmid = "1859758"
        self.doi = "10.1093/bja/67.1.41"
        self.validate_anthropometric_values()

        self.v1 = 0.228 * weight
        self.v2 = 0.463 * weight
        self.v3 = 2.893 * weight

        self.k10 = 0.119
        self.k12 = 0.112
        self.k13 = 0.0419
        self.k21 = 0.055
        self.k31 = 0.0033
        self.ke0 = 1.2

class Schnider(Model):
    """Schnider class holds pharmacokinetic parameters for the Schnider propofol
    model.

    Reference: PMID: 9605675 DOI: 10.1097/00000542-199805000-00006
    """

    def __init__(self, sex: int, age: float, weight: float, height: float):
        super().__init__(sex, age, weight, height)

        self.compartments = 3
        self.concentration_unit = "mg/ml"
        self.target_unit = "ug/ml"
        self.age_lower = -1
        self.age_upper = -1
        self.weight_lower = -1
        self.weight_upper = -1
        if sex == 0:
            self.bmi_upper = 42
        elif sex == 1:
            self.bmi_upper = 35
        self.pmid = "9605675"
        self.doi = "10.1097/00000542-199805000-00006"
        self.validate_anthropometric_values()

        lbm = lbm_dubois(sex, weight, height)

        self.v1 = 4.27
        self.v2 = 18.9 - 0.391 * (age - 52)
        self.v3 = 238

        self.cl1 = 1.89 + 0.0456 * (weight - 77) - 0.0681 * (lbm - 59) \
                   + 0.0264 * (height - 177)
        self.cl2 = 1.29 - 0.024 * (age - 53)
        self.cl3 = 0.836

        self.k10 = self.cl1 / self.v1
        self.k12 = self.cl2 / self.v1
        self.k13 = self.cl3 / self.v1
        self.k21 = self.cl2 / self.v2
        self.k31 = self.cl3 / self.v3
        self.ke0 = 0.456  # TTPE 1.6 minutes is used in original model

class Paedfusor(Model):
    """Paedfusor class holds pharmacokinetic parameters for the Paedfusor propofol
    model.

    Reference: PMID: 15941735 DOI: 10.1093/bja/aei567
    """

    def __init__(self, sex: int, age: float, weight: float, height: float):
        super().__init__(sex, age, weight, height)

        self.compartments = 3
        self.concentration_unit = "mg/ml"
        self.target_unit = "ug/ml"
        self.age_lower = 1
        self.age_upper = 16
        self.weight_lower = 5
        self.weight_upper = 61
        self.pmid = "15941735"
        self.doi = "10.1093/bja/aei567"
        self.validate_anthropometric_values()

        self.k12 = 0.114
        self.k13 = 0.0419
        self.k21 = 0.055
        self.k31 = 0.0033
        self.ke0 = 0.26

        if age <= 12:
            self.v1 = 458.4 * weight / 1000
            self.k10 = 0.1527 * (weight ** -0.3)
        elif age >= 13:
            self.v1 = 400 * weight / 1000
            self.k10 = 0.0678
        elif age >= 14:
            self.v1 = 342 * weight / 1000
            self.k10 = 0.0792
        elif age >= 15:
            self.v1 = 284 * weight / 1000
            self.k10 = 0.0954
        elif age >= 16:
            self.v1 = 228.57 * weight / 1000
            self.k10 = 0.119

        self.v2 = self.v1 * self.k12 / self.k21 / 1000
        self.v3 = self.v1 * self.k13 / self.k31 / 1000

class Kataria(Model):
    """Kataria class holds pharmacokinetic parameters for the Kataria propofol
    model.

    Reference: PMID: 8291699 DOI: 10.1097/00000542-199401000-00018
    """

    def __init__(self, sex: int, age: float, weight: float, height: float):
        super().__init__(sex, age, weight, height)

        self.compartments = 3
        self.concentration_unit = "mg/ml"
        self.target_unit = "ug/ml"
        self.age_lower = 3
        self.age_upper = 11
        self.weight_lower = 15
        self.weight_upper = 61
        self.pmid = "8291699"
        self.doi = "10.1097/00000542-199401000-00018"
        self.validate_anthropometric_values()

        self.v1 = weight * 0.41
        self.v2 = weight * 0.78 + 3.1 * age - 16
        self.v3 = weight * 6.9
        self.cl1 = weight * 0.035
        self.cl2 = weight * 0.077
        self.cl3 = weight * 0.026

        self.k10 = self.cl1 / self.v1
        self.k12 = self.cl2 / self.v1
        self.k13 = self.cl3 / self.v1
        self.k21 = self.cl2 / self.v2
        self.k31 = self.cl3 / self.v3
        self.ke0 = 0.41

class Eleveld(Model):
    """Eleveld class holds pharmacokinetic parameters for the Eleveld propofol
    model.

    Reference: PMID: 29661412 DOI: 10.1016/j.bja.2018.01.018
    """

    def __init__(self, sex: int, age: float, weight: float, height: float,
                 opiates_coadministered: bool):
        super().__init__(sex, age, weight, height)

        self.compartments = 3
        self.concentration_unit = "mg/ml"
        self.target_unit = "ug/ml"
        self.age_lower = -1
        self.age_upper = -1
        self.weight_lower = -1
        self.weight_upper = -1
        self.pmid = "29661412"
        self.doi = "10.1016/j.bja.2018.01.018"
        self.validate_anthropometric_values()

        # Reference
        sex_ref = 0
        age_ref = 35
        weight_ref = 70
        height_ref = 170

        # Theta constants
        theta_1 = 6.28  # V1ref L
        theta_2 = 25.5  # V2ref L
        theta_3 = 273  # V3ref L
        theta_4 = 1.79  # CLref (male) L/min
        theta_5 = 1.83  # Q2ref L/min
        theta_6 = 1.11  # Q3ref L/min
        theta_7 = 0.191  # Typical residual error
        theta_8 = 42.3  # CL maturation E50 weeks
        theta_9 = 9.06  # CL maturation slope
        theta_10 = -0.0156  # Smaller V2 with age
        theta_11 = -0.00286  # Lower CL with age
        theta_12 = 33.6  # Weight for 50% of maximal V1 Kg
        theta_13 = -0.0138  # Smaller V3 with age
        theta_14 = 68.3  # Maturation of Q3 weeks
        theta_15 = 2.1  # CLref (female) L$min1
        theta_16 = 1.3  # Higher Q2 for maturation of Q3
        theta_17 = 1.42  # V1 venous samples (children)
        theta_18 = 0.68  # Higher Q2 venous samples

        # Post menstrual age
        pma = age * 52 + 40
        pma_ref = age_ref * 52 + 40

        def ageing(x, age):
            return exp(x * (age - age_ref))

        def sigmoid(x, e50, y):
            return (x ** y) / ((x ** y) + (e50 ** y))

        def central(x):
            return sigmoid(x, theta_12, 1)

        def opiates(x, present):
            if present:
                return exp(x * age)
            else:
                return 1

        # cl1 maturation
        cl1_mat = sigmoid(pma, theta_8, theta_9)
        cl1_mat_ref = sigmoid(pma_ref, theta_8, theta_9)

        # cl3 maturation
        cl3_mat = sigmoid(pma, theta_14, 1)
        cl3_mat_ref = sigmoid(pma_ref, theta_14, 1)

        # fat free mass
        ffm = ffm_alsallami(sex, age, weight, height)
        ffm_ref = ffm_alsallami(sex_ref, age_ref, weight_ref, height_ref)

        self.v1 = theta_1 * (central(weight) / central(weight_ref))
        self.v2 = theta_2 * (weight / weight_ref) * ageing(theta_10, age)
        self.v3 = theta_3 * (ffm / ffm_ref) * opiates(theta_13,
                                                      opiates_coadministered)

        if sex == 0:
            self.cl1 = theta_4 * ((weight / weight_ref) ** 0.75) * \
                       (cl1_mat / cl1_mat_ref) * \
                       opiates(theta_11, opiates_coadministered)
        elif sex == 1:
            self.cl1 = theta_15 * (weight / weight_ref) ** 0.75 * \
                       (cl1_mat / cl1_mat_ref) * \
                       opiates(theta_11, opiates_coadministered)

        self.cl2 = theta_5 * (self.v2 / theta_2) ** 0.75 * \
                    (1 + theta_16 * (1 - cl3_mat))

        self.cl3 = theta_6 * (self.v3 / theta_3) ** 0.75 * \
                   (cl3_mat / cl3_mat_ref)

        self.k10 = self.cl1 / self.v1
        self.k12 = self.cl2 / self.v1
        self.k13 = self.cl3 / self.v1
        self.k21 = self.cl2 / self.v2
        self.k31 = self.cl3 / self.v3
        self.ke0 = 0.146 * ((weight / weight_ref) ** -0.25)

        self.ce50 = 3.08 * ageing(-0.00635, age)

class Short(Model):
    """Short class holds pharmacokinetic parameters for the Short
    propofol model; uses Keo from Eleveld

    Reference: PMID: 8130049 DOI: 10.1093/bja/72.3.302
    """

    def __init__(self, sex: int, age: float, weight: float, height: float):
        super().__init__(sex, age, weight, height)

        self.compartments = 3
        self.concentration_unit = "mg/ml"
        self.target_unit = "ug/ml"
        self.age_lower = 4
        self.age_upper = 7
        self.weight_lower = 15
        self.weight_upper = 22
        self.pmid = "8130049"
        self.doi = "10.1093/bja/72.3.302"
        self.validate_anthropometric_values()

        self.v1 = 0.432 * weight

        self.k10 = 0.0967
        self.k12 = 0.1413
        self.k13 = 0.1092
        self.k21 = 0.0392
        self.k31 = 0.0049
        self.ke0 = 0.146 * ((weight / 70) ** -0.25)

class Schuttler(Model):
    """Schuttler class holds pharmacokinetic parameters for the Schuttler
       propofol model; uses Keo from Eleveld

    Reference: PMID:  DOI: 10.1097/00000542-200003000-00017
    """

    def __init__(self, sex: int, age: float, weight: float, height: float,
                 bolus_data: bool = False, venous_data: bool = False):
        super().__init__(sex, age, weight, height)

        self.compartments = 3
        self.concentration_unit = "mg/ml"
        self.target_unit = "ug/ml"
        self.age_lower = 2
        self.age_upper = 88
        self.weight_lower = 2
        self.weight_upper = 88
        self.pmid = "10719952"
        self.doi = "10.1097/00000542-200003000-00017"
        self.validate_anthropometric_values()

        # Theta constants
        theta_1 = 1.44
        theta_2 = 9.3
        theta_3 = 2.25
        theta_4 = 44.2
        theta_5 = 0.92
        theta_6 = 266
        theta_7 = 0.75
        theta_8 = 0.62
        theta_9 = 0.61
        theta_10 = 0.045
        theta_11 = 0.55
        theta_12 = 0.71
        theta_13 = 20.39
        theta_14 = 20.40
        theta_15 = 1.61
        theta_16 = 2.02
        theta_17 = 0.73
        theta_18 = 20.48

        if bolus_data:
            bol = 1
        else:
            bol = 0

        if venous_data:
            ven = 1
        else:
            ven = 0

        self.v1 = theta_2 * (weight / 70) ** theta_12 * \
                  (age / 30) ** theta_13 * (1 + bol * theta_15)
        self.v2 = theta_4 * (weight / 70) ** theta_9 * \
                  (1 + bol * theta_17)
        self.v3 = theta_6

        if age <= 60:
            self.cl1 = theta_1 * (weight / 70) ** theta_7
        else:
            self.cl1 = theta_1 * (weight / 70) ** theta_7 - \
                       (age - 60) * theta_10
        self.cl2 = theta_3 * (weight / 70) ** theta_8 * \
                   (1 + ven * theta_14) * (1 + bol * theta_16)
        self.cl3 = theta_5 * (weight / 70) ** theta_11 * \
                   (1 + bol * theta_18)

        self.k10 = self.cl1 / self.v1
        self.k12 = self.cl2 / self.v1
        self.k13 = self.cl3 / self.v1
        self.k21 = self.cl2 / self.v2
        self.k31 = self.cl3 / self.v3
        self.ke0 = 0.146 * ((weight / 70) ** -0.25)


"""
对transformer encoder的paddle实现。
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from functools import partial
import numpy as np
import paddle.fluid as fluid
import paddle.fluid.layers as layers

# transformer_encoder结构的实现
# transformer_encoder的主题为依次连接的多层encoder层
def transformer_encoder(inputs,
                        attention_bias,
                        num_attention_layers,
                        num_attention_heads,
                        hidden_state_size,
                        key_size,
                        value_size,
                        inner_hidden_size,
                        attention_dropout,
                        activate_dropout,
                        post_and_pre_process_dropout,
                        hidden_act,
                        preprocess_cmd="n",
                        postprocess_cmd="da",
                        param_initializer=None,
                        name='',
                        is_test=False):

    enc_input = inputs
    enc_output_list = []
    # 多层encoder依次相连
    for i in range(num_attention_layers):
        enc_output = encoder_layer(
            enc_input,
            attention_bias,
            num_attention_heads,
            hidden_state_size,
            key_size,
            value_size,
            inner_hidden_size,
            attention_dropout,
            activate_dropout,
            post_and_pre_process_dropout,
            hidden_act,
            preprocess_cmd,
            postprocess_cmd,
            param_initializer=param_initializer,
            name=name + '_layer_' + str(i),
            is_test=is_test)
        enc_output_list.append(enc_output)
        enc_input = enc_output

    enc_output = pre_process_layer(
        enc_output, preprocess_cmd, post_and_pre_process_dropout, name="post_encoder", is_test=is_test)

    return enc_output, enc_output_list

# encoder层的实现
# encoder层的主体为multi_head_attention和前馈神经网络
def encoder_layer(inputs,
                  attention_bias,
                  num_attention_heads,
                  hidden_state_size,
                  key_size,
                  value_size,
                  inner_hidden_size,
                  attention_dropout,
                  activate_dropout,
                  post_and_pre_process_dropout,
                  hidden_act,
                  preprocess_cmd="n",
                  postprocess_cmd="da",
                  param_initializer=None,
                  name='',
                  is_test=False
                  ):
    # multi_head_attention
    attn_output = multi_head_attention(
        pre_process_layer(
            inputs,
            preprocess_cmd,
            post_and_pre_process_dropout,
            name=name + '_pre_att', is_test=is_test),
        None,
        None,
        attention_bias,
        key_size,
        value_size,
        hidden_state_size,
        num_attention_heads,
        attention_dropout,
        param_initializer=param_initializer,
        name=name + '_multi_head_att', is_test=is_test)
    attn_output = post_process_layer(
        inputs,
        attn_output,
        postprocess_cmd,
        post_and_pre_process_dropout,
        name=name + '_post_att', is_test=is_test)

    # 前馈神经网络层
    ffd_output = feed_forward_layer(
        pre_process_layer(
            attn_output,
            preprocess_cmd,
            post_and_pre_process_dropout,
            name=name + '_pre_ffn',
            is_test=is_test),
        hidden_state_size,
        inner_hidden_size,
        activate_dropout,
        hidden_act,
        param_initializer=param_initializer,
        name=name + '_ffn',
        is_test=is_test)

    return post_process_layer(
        attn_output,
        ffd_output,
        postprocess_cmd,
        post_and_pre_process_dropout,
        name=name + '_post_ffn',
        is_test=is_test)

# 对encoder_layer中前馈层的实现
# 其主体为两个全连接层，一层实现hidden_size到inner_hidden_size的转化，另一层实现inner_hidden_size到hidden_size的转化
def feed_forward_layer(x,
                       hidden_size,
                       inner_hidden_size,
                       dropout_rate,
                       hidden_act,
                       param_initializer=None,
                       name='ffn',
                       is_test=False):

    # 将维数为hidden_size的向量转化为inner_hidden_size
    hidden = layers.fc(input=x,
                       size=inner_hidden_size,
                       num_flatten_dims=2,
                       act=hidden_act,
                       param_attr=fluid.ParamAttr(
                           name=name + '_fc_0.w_0',
                           initializer=param_initializer),
                       bias_attr=name + '_fc_0.b_0')
    if dropout_rate:
        hidden = layers.dropout(
            hidden,
            dropout_prob=dropout_rate,
            dropout_implementation="upscale_in_train",
            is_test=is_test)
    # 将维数为inner_hidden_size的向量转化为hidden_size
    out = layers.fc(input=hidden,
                    size=hidden_size,
                    num_flatten_dims=2,
                    param_attr=fluid.ParamAttr(
                        name=name + '_fc_1.w_0', initializer=param_initializer),
                    bias_attr=name + '_fc_1.b_0')
    return out

# 实现dropout、输入相加、normalization的前处理、后处理层
# 一般接在其他层之前或之后，对输入输出进行处理
def pre_post_process_layer(prev_out, out, process_cmd, dropout_rate=0.,
                           name='', is_test=False):

    for cmd in process_cmd:
        if cmd == "a":  # 两个输入相加
            out = out + prev_out if prev_out else out
        elif cmd == "n":  # 进行normalization
            out_type = out.dtype
            if out_type == fluid.core.VarDesc.VarType.FP16:
                out = layers.cast(x=out, dtype="float32")
            out = layers.layer_norm(
                out,
                begin_norm_axis=len(out.shape) - 1,
                param_attr=fluid.ParamAttr(
                    name=name + '_layer_norm_scale',
                    initializer=fluid.initializer.Constant(1.)),
                bias_attr=fluid.ParamAttr(
                    name=name + '_layer_norm_bias',
                    initializer=fluid.initializer.Constant(0.)))
            if out_type == fluid.core.VarDesc.VarType.FP16:
                out = layers.cast(x=out, dtype="float16")
        elif cmd == "d":  # 进行dropout
            if dropout_rate:
                out = layers.dropout(
                    out,
                    dropout_prob=dropout_rate,
                    dropout_implementation="upscale_in_train",
                    is_test=is_test)
    return out

# 对多头attention计算的实现
# 通过queries和keys计算attention值，并依此对values进行加权求和
# attention_bias可实现对某些位置的mask作用
def multi_head_attention(queries,
                         keys,
                         values,
                         attention_bias,
                         key_size,
                         value_size,
                         hidden_size,
                         n_head=1,
                         dropout_rate=0.,
                         cache=None,
                         param_initializer=None,
                         name='multi_head_att',
                         is_test=False):

    keys = queries if keys is None else keys
    values = keys if values is None else values

    if not (len(queries.shape) == len(keys.shape) == len(values.shape) == 3):
        raise ValueError(
            "Inputs: quries, keys and values should all be 3-D tensors.")

    # 以下为函数定义部分
    # 定义了一系列计算attention所需的函数

    # 定义q、k、v三个矩阵，将输入的queries、keys、values与三个矩阵相乘得到用于计算的q、k、v
    def __compute_qkv(queries, keys, values, n_head, key_size, value_size):
        q = layers.fc(input=queries,
                      size=key_size * n_head,
                      num_flatten_dims=2,
                      param_attr=fluid.ParamAttr(
                          name=name + '_query_fc.w_0',
                          initializer=param_initializer),
                      bias_attr=name + '_query_fc.b_0')
        k = layers.fc(input=keys,
                      size=key_size * n_head,
                      num_flatten_dims=2,
                      param_attr=fluid.ParamAttr(
                          name=name + '_key_fc.w_0',
                          initializer=param_initializer),
                      bias_attr=name + '_key_fc.b_0')
        v = layers.fc(input=values,
                      size=value_size * n_head,
                      num_flatten_dims=2,
                      param_attr=fluid.ParamAttr(
                          name=name + '_value_fc.w_0',
                          initializer=param_initializer),
                      bias_attr=name + '_value_fc.b_0')
        return q, k, v

    # 将输入的[batch_size, max_sequence_length, n_head * hidden_dim]维度的向量转换为[batch_size, n_head, max_sequence_length,
    # hidden_dim]维度，以便后续的多头attention计算
    def __split_heads(x, num_head):
        hidden_size = x.shape[-1]
        reshaped = layers.reshape(
            x=x, shape=[0, 0, num_head, hidden_size // num_head], inplace=True)
        return layers.transpose(x=reshaped, perm=[0, 2, 1, 3])

    # 对_split_heads的逆操作，将[batch_size, n_head, max_sequence_length,hidden_dim]维的向量转化回[batch_size,
    # max_sequence_length, n_head * hidden_dim]维度
    def __combine_heads(x):

        if len(x.shape) == 3: return x
        if len(x.shape) != 4:
            raise ValueError("Input(x) should be a 4-D Tensor.")

        trans_x = layers.transpose(x, perm=[0, 2, 1, 3])
        return layers.reshape(
            x=trans_x,
            shape=[0, 0, trans_x.shape[2] * trans_x.shape[3]],
            inplace=True)

    # 计算q与k的点积，并凭此对v加权求和
    def scaled_dot_product_attention(q, k, v, attn_bias, d_key, dropout_rate, is_test=False):
        scaled_q = layers.scale(x=q, scale=d_key**-0.5)
        product = layers.matmul(x=scaled_q, y=k, transpose_y=True)
        if attn_bias:
            product += attn_bias
        weights = layers.softmax(product)
        if dropout_rate:
            weights = layers.dropout(
                weights,
                dropout_prob=dropout_rate,
                dropout_implementation="upscale_in_train",
                is_test=is_test)
        out = layers.matmul(weights, v)
        return out

    # 函数定义部分结束
    # 开始计算attention
    q, k, v = __compute_qkv(queries, keys, values, n_head, key_size, value_size)

    if cache is not None:  # use cache and concat time steps
        # Since the inplace reshape in __split_heads changes the shape of k and
        # v, which is the cache input for next time step, reshape the cache
        # input from the previous time step first.
        k = cache["k"] = layers.concat(
            [layers.reshape(
                cache["k"], shape=[0, 0, hidden_size]), k], axis=1)
        v = cache["v"] = layers.concat(
            [layers.reshape(
                cache["v"], shape=[0, 0, hidden_size]), v], axis=1)

    # 转化向量维度
    q = __split_heads(q, n_head)
    k = __split_heads(k, n_head)
    v = __split_heads(v, n_head)

    ctx_multiheads = scaled_dot_product_attention(q, k, v, attention_bias, key_size,
                                                  dropout_rate, is_test=is_test)

    out = __combine_heads(ctx_multiheads)

    # 投影回模型所需的hidden_size，本模型中out的维度与hidden_size相同
    proj_out = layers.fc(input=out,
                         size=hidden_size,
                         num_flatten_dims=2,
                         param_attr=fluid.ParamAttr(
                             name=name + '_output_fc.w_0',
                             initializer=param_initializer),
                         bias_attr=name + '_output_fc.b_0')
    return proj_out

# 定义用于预处理和后处理的层
pre_process_layer = partial(pre_post_process_layer, None)
post_process_layer = pre_post_process_layer


# -*- coding: utf-8 -*-
"""
"""
from __future__ import division, print_function, unicode_literals
import numpy as np
import declarative

#import phasor.numerics.dispatched as dmath
#import sympy

from ..utilities.future_from_2 import super
from ..base.autograft import invalidate_auto
from ..base.bases import Element

from ..base.multi_unit_args import (
    unitless_refval_attribute,
)

from .utils import (
    TargetLeft,
    TargetRight,
    TargetIdx,
    matrix_focus,
    np_check_sorted,
    str_m,
)

from .substrates import (
    substrate_environment,
)

from . import standard_attrs as attrs

class MatrixAtsBase(Element):
    #TODO report in ctree

    @declarative.dproperty
    def plotname(self, arg = declarative.NOARG):
        if arg is declarative.NOARG:
            arg = self.name
        return arg

    @declarative.dproperty
    def reversed(self, arg = declarative.NOARG):
        elname = "reversed"
        if arg is declarative.NOARG:
            arg = False

        ooa = self.ctree
        if self.inst_prototype_t in ["full"]:
            #TODO make this do the correct thing
            arg = getattr(self.inst_prototype, elname)
        else:
            ooa = self.ctree.useidx('immediate')

        ooa[elname] = arg
        #arg = ooa.setdefault(elname, arg)
        return ooa[elname]

    @declarative.mproperty
    def env_reversed(self):
        #TODO put this into the environment_query
        #print("PREV: ", self.parent.env_reversed, " ME: ", self.reversed)
        p_env_reversed = self.parent.environment_query((MatrixAtsBase, "reversed"))
        arg = bool(p_env_reversed) ^ bool(self.reversed)
        #arg = self.ctree.setdefault("env_reversed", arg)
        self.ctree.env_reversed = arg
        return self.ctree.env_reversed

    @declarative.mproperty(simple_delete = True)
    @invalidate_auto
    def matrix_inv(self):
        #print(self.__class__)
        #print(self.matrix)
        return self.matrix**(-1)

    def matrix_between(self, tidx1, tidx2):
        if tidx1 == TargetLeft:
            if tidx2 == TargetLeft:
                return np.eye(2)
            elif tidx2 == TargetRight:
                return self.matrix
            else:
                raise RuntimeError("Unknown Target {0}".format(tidx1))
        elif tidx1 == TargetRight:
            if tidx2 == TargetLeft:
                return self.matrix_inv
            elif tidx2 == TargetRight:
                return np.eye(2)
            else:
                raise RuntimeError("Unknown Target {0}".format(tidx1))
        else:
            print(self.__class__)
            raise RuntimeError("Unknown Target {0}".format(tidx1))

    def matrix_target_to_z_single(self, tidx1, z_m, invert = False):
        raise NotImplementedError()

    def matrix_target_to_z(self, tidx1, z_m, fill, invert = False):
        z_m = np.asarray(z_m)
        if len(z_m.shape) > 0:
            if len(z_m.shape) == 1 and np_check_sorted(z_m):
                return self.matrix_target_to_z_linsorted(tidx1, z_m, fill, invert = invert)

            fill_reshape = fill.reshape(2, 2, -1)
            idx_remux = np.indices(z_m.shape)
            idx_remux = idx_remux.reshape(len(z_m.shape), -1)
            z_m = z_m.reshape(-1)
            sidx = np.argsort(z_m)
            z_m = z_m[sidx]
            idx_remux = idx_remux[:, sidx]
            self.matrix_target_to_z_linsorted(tidx1, z_m, fill_reshape, invert = invert)
            fill = fill_reshape[(slice(None), slice(None)) + tuple(idx_remux[i] for i in range(idx_remux.shape[0]))]
            return
        else:
            return self.matrix_target_to_z_single(z_m)

    def matrix_target_to_z_linsorted(self, tidx1, z_m, fill, invert = False):
        it = np.nditer(z_m, flags = ['multi_index'])
        while not it.finished:
            fill[(slice(None), slice(None)) + it.multi_index] = self.matrix_target_to_z_single(tidx1, it.value)
            it.iternext()
        return fill

    @declarative.mproperty
    def constraints(self):
        return []

    def as_target(self, direction = 'left'):
        sub_target = self.parent._target_to_child(self)
        if direction == 'left':
            return TargetIdx(TargetLeft + sub_target)
        elif direction == 'right':
            return TargetIdx(TargetRight + sub_target)
        else:
            return None

    def environment_query_local(self, query):
        if query == (MatrixAtsBase, "reversed"):
            return self.env_reversed
        return super().environment_query_local(query)

    @declarative.dproperty
    def root(self):
        return self.environment_query((MatrixAtsBase, "root"))

class MatrixAtsCompositeBase(MatrixAtsBase):
    pass

class ThinBase(MatrixAtsBase, declarative.OverridableObject):
    width_m      = 0

    _loc_default = ('loc_m', None)
    loc_m = attrs.generate_loc_m()

    def matrix_target_to_z_single(self, tidx1, z_m, invert = False):
        if z_m != 0:
            raise RuntimeError("Only located at 0")

        #invert if viewing from the right
        if tidx1 == TargetRight:
            invert = not invert

        if not invert:
            return self.matrix
        else:
            return self.matrix_inv

    def matrix_target_to_z(self, tidx1, z_m, fill, invert = False):
        return self.matrix_target_to_z_linsorted(self, z_m, fill, invert = invert)

    def matrix_target_to_z_linsorted(self, tidx1, z_m, fill, invert = False):
        if not all(z_m == 0):
            raise RuntimeError("Only located at 0")

        #invert if viewing from the right
        if tidx1 == TargetRight:
            invert = not invert

        if not invert:
            fill[:, :, ...] = self.matrix
        else:
            fill[:, :, ...] = self.matrix_inv
        return

    def target_pos(self, target):
        return 0

    def system_data_targets(self, typename):
        dmap = {}
        return dmap

class NoP(ThinBase):
    @declarative.mproperty
    def matrix(self):
        return np.matrix([[1, 0], [0, 1]])

    @declarative.mproperty(simple_delete = True)
    def matrix_inv(self):
        return np.matrix([[1, 0], [0, 1]])

class ThinLens(ThinBase):

    f_m = attrs.generate_f_m()

    @declarative.mproperty
    def matrix(self):
        mat = matrix_focus(f_m = self.f_m.val)
        return mat

    @declarative.mproperty
    def matrix_inv(self):
        mat = matrix_focus(f_m = -self.f_m.val)
        return mat

    def lens_description(self, z, from_target):
        return declarative.Bunch(
            f_m = self.f_m.val,
            z = z,
            type = 'lens',
            name = self.plotname,
            str = 'thin lens f_m = {f_m}'.format(f_m = str_m(self.f_m.val)),
        )

    def detune_description(self, z, q_left):
        q_right = q_left.propagate_matrix(self.matrix)
        cplg02 = q_right.cplg02 - q_left.cplg02
        return declarative.Bunch(
            cplg02   = cplg02,
            type    = 'lens',
            q       = q_left,
            obj     = self,
        )

    def system_data_targets(self, typename):
        dmap = {}
        if typename == 'lens_description':
            dmap[TargetIdx()] = self.lens_description
        elif typename == 'detune_description':
            dmap[TargetIdx()] = self.detune_description
        return dmap

class LensInterface(ThinBase):
    substrate_from = substrate_environment
    substrate_to   = substrate_environment

    R_m = attrs.generate_R_m()

    @declarative.mproperty
    def matrix(self):
        n_from = self.substrate_from.n(self)
        n_to   = self.substrate_to.n(self)
        if self.R_m.val is not None:
            if not self.env_reversed:
                mat = np.matrix([
                    [1, 0],
                    [(n_from/n_to - 1)/self.R_m.val, n_from / n_to],
                ])
            else:
                mat = np.matrix([
                    [1, 0],
                    [(n_to/n_from - 1)/-self.R_m.val, n_to / n_from],
                ])
        else:
            if not self.env_reversed:
                mat = np.matrix([
                    [1 , 0],
                    [0 , n_from / n_to],
                ])
            else:
                mat = np.matrix([
                    [1 , 0],
                    [0 , n_to / n_from],
                ])
        return mat

    @declarative.mproperty
    def matrix_inv(self):
        n_to = self.substrate_from.n(self)
        n_from = self.substrate_to.n(self)
        if self.R_m.val is not None:
            if not self.env_reversed:
                mat = np.matrix([
                    [1, 0],
                    [(n_from/n_to - 1)/self.R_m.val, n_from / n_to],
                ])
            else:
                mat = np.matrix([
                    [1, 0],
                    [(n_to/n_from - 1)/-self.R_m.val, n_to / n_from],
                ])
        else:
            if not self.env_reversed:
                mat = np.matrix([
                    [1 , 0],
                    [0 , n_from / n_to],
                ])
            else:
                mat = np.matrix([
                    [1 , 0],
                    [0 , n_to / n_from],
                ])
        return mat

class Mirror(ThinBase):
    R_m = attrs.generate_R_m()

    @declarative.mproperty
    def matrix(self):
        if self.R_m.val is not None:
            mat = np.matrix([
                [1,      0],
                [-2/self.R_m.val, 1],
            ])
        else:
            mat = np.matrix([
                [1 , 0],
                [0 , 1],
            ])
        return mat

    def mirror_description(self, z, from_target):
        if self.R_m.val is not None:
            f_m = -1/self.matrix[1, 0]
            return declarative.Bunch(
                R_m = self.R_m.val,
                f_m = f_m,
                z = z,
                name = self.plotname,
                type = 'mirror',
                str = 'Mirror, R_m = {R_m}, f_m = {f_m}'.format(R_m = str_m(self.R_m.val), f_m = str_m(f_m)),
            )
        else:
            return declarative.Bunch(
                R_m = self.R_m.val,
                z = z,
                type = 'mirror',
                name = self.plotname,
                str = 'Mirror, flat',
            )

    def detune_description(self, z, q_left):
        q_right = q_left.propagate_matrix(self.matrix)
        cplg02 = q_right.cplg02 + q_left.cplg02
        return declarative.Bunch(
            cplg02   = cplg02,
            type    = 'mirror',
            q       = q_left,
            obj     = self,
        )

    def system_data_targets(self, typename):
        dmap = {}
        if typename == 'mirror_description':
            dmap[TargetIdx()] = self.mirror_description
        elif typename == 'detune_description':
            dmap[TargetIdx()] = self.detune_description
        return dmap

class ABCDGeneric(ThinBase):
    _A_default = 1
    @declarative.dproperty_adv
    def A(desc):
        return unitless_refval_attribute(
            desc,
            prototypes    = ['full', 'base'],
            default_attr  = '_A_default',
            allow_fitting = True,
        )

    _B_default = 0
    @declarative.dproperty_adv
    def B(desc):
        return unitless_refval_attribute(
            desc,
            prototypes    = ['full', 'base'],
            default_attr  = '_B_default',
            allow_fitting = True,
        )

    _C_default = 0
    @declarative.dproperty_adv
    def C(desc):
        return unitless_refval_attribute(
            desc,
            prototypes    = ['full', 'base'],
            default_attr  = '_C_default',
            allow_fitting = True,
        )

    @declarative.dproperty
    def D(self):
        return (self.C.val * self.B.val + 1)/self.A.val

    @declarative.mproperty
    def matrix(self):
        return np.matrix([[self.A.val, self.B.val], [self.C.val, self.D]])


import numpy as np
import warnings
import nestpy

print(f'Using nestpy version {nestpy.__version__}')

@np.vectorize
def quanta_from_NEST(en, model, e_field, A, Z, create_s2, **kwargs):
    """
    Function which uses NEST to yield photons and electrons
    for a given set of parameters.

    Note:
        In case the energy deposit is outside of the range of NEST a -1
        is returned.

    Args:
        en (numpy.array): Energy deposit of the interaction [keV]
        model (numpy.array): Nest Id for qunata generation (integers)
        e_field (numpy.array): Field value in the interaction site [V/cm]
        A (numpy.array): Atomic mass number
        Z (numpy.array): Atomic number
        create_s2 (bool): Specifies if S2 can be produced by interaction,
            in this case electrons are generated.
        kwargs: Additional keyword arguments which can be taken by
            GetYields e.g. density.

    Returns:
        photons (numpy.array): Number of generated photons
        electrons (numpy.array): Number of generated electrons
        excitons (numpy.array): Number of generated excitons
    """
    nc = nestpy.NESTcalc(nestpy.VDetector())
    density = 2.862  # g/cm^3

    # Fix for Kr83m events.
    # Energies have to be very close to 32.1 keV or 9.4 keV
    # See: https://github.com/NESTCollaboration/nest/blob/master/src/NEST.cpp#L567
    # and: https://github.com/NESTCollaboration/nest/blob/master/src/NEST.cpp#L585
    max_allowed_energy_difference = 1  # keV
    if model == 11:
        if abs(en - 32.1) > max_allowed_energy_difference:
            en = 32.1
        if abs(en - 9.4) > max_allowed_energy_difference:
            en = 9.4

    # Some addition taken from
    # https://github.com/NESTCollaboration/nestpy/blob/e82c71f864d7362fee87989ed642cd875845ae3e/src/nestpy/helpers.py#L94-L100
    if model == 0 and en > 2e2:
        warnings.warn(f"Energy deposition of {en} keV beyond NEST validity for NR model of 200 keV - Remove Interaction")
        return -1, -1, -1
    if model == 7 and en > 3e3:
        warnings.warn(f"Energy deposition of {en} keV beyond NEST validity for gamma model of 3 MeV - Remove Interaction")
        return -1, -1, -1
    if model == 8 and en > 3e3:
        warnings.warn(f"Energy deposition of {en} keV beyond NEST validity for beta model of 3 MeV - Remove Interaction")
        return -1, -1, -1

    y = nc.GetYields(interaction=nestpy.INTERACTION_TYPE(model),
                     energy=en,
                     drift_field=e_field,
                     A=A,
                     Z=Z,
                     **kwargs
                     )

    event_quanta = nc.GetQuanta(y)  # Density argument is not use in function...

    photons = event_quanta.photons
    excitons = event_quanta.excitons
    electrons = 0
    if create_s2:
        electrons = event_quanta.electrons

    return photons, electrons, excitons


#!/usr/bin/env python
# coding: utf-8

# load a bunch of stuff
from __future__ import division
# load
import cantera as ct
import numpy as np
import scipy
import pylab
import matplotlib
import matplotlib.pyplot  as plt
import matplotlib.gridspec as gridspec
from matplotlib.pyplot import cm
from matplotlib.ticker import NullFormatter, MaxNLocator, LogLocator
plt.switch_backend('agg')  # needed for saving figures
import csv
from pydas.dassl import DASSL
import os
import rmgpy
import rmg
import re
import operator
import pandas as pd
import pylab
from cycler import cycler
import seaborn as sns
import os
import multiprocessing

# this chemkin file is from the cti generated by rmg
gas = ct.Solution('./rh/cantera/chem_annotated.cti', 'gas')
surf = ct.Interface('./rh/cantera/chem_annotated.cti', 'surface1', [gas])
# gas = ct.Solution('chem_vlachos_bidentate.cti', 'gas')
# surf = ct.Interface('chem_vlachos_bidentate.cti', 'surface1', [gas])

print("This mechanism contains {} gas reactions and {} surface reactions".format(gas.n_reactions, surf.n_reactions))

i_ar = gas.species_index('Ar')
i_ch4 = gas.species_index('CH4(2)')
i_o2 = gas.species_index('O2(3)')
i_co2 = gas.species_index('CO2(4)')
i_h2o = gas.species_index('H2O(5)')
i_h2 = gas.species_index('H2(6)')
i_co = gas.species_index('CO(7)')

# unit conversion factors to SI
mm = 0.001
cm = 0.01
ms = mm
minute = 60.0

#######################################################################
# Input Parameters
#######################################################################
t_in = 700  # K - in the paper, it was ~698.15K at the start of the cat surface and ~373.15 for the gas inlet temp
t_cat = t_in
length = 70 * mm  # Reactor length- m
diam = 16.5*mm  # Reactor diameter - in m
area = (diam/2.0)**2*np.pi  # Reactor cross section area (area of tube) in m^2
porosity = 0.81  # Monolith channel porosity, from Horn ref 17 sec 2.2.2
cat_area_per_vol = 16000  # I made this up, in m-1. 4500 is lowest that "work" for all base
flow_rate = 4.7  # slpm
flow_rate = flow_rate*.001/60  # m^3/s
tot_flow = 0.208  # from Horn 2007, constant inlet flow rate in mol/min, equivalent to 4.7 slpm
velocity = flow_rate/area  # m/s
# The PFR will be simulated by a chain of 'NReactors' stirred reactors.
NReactors = 7001

on_catalyst = 1000  # catalyst length 10mm, but it doesn't say where.  let's guess at 1 cm?
off_catalyst = 2000
dt = 1.0

reactor_len = length/(NReactors-1)
rvol = area * reactor_len * porosity
# catalyst area in one reactor
cat_area = cat_area_per_vol * rvol

def plotZoom(a):
    gas_out, surf_out, gas_names, surf_names, dist_array, T_array = a

    fig, axs = plt.subplots(1, 2)
    axs[0].set_prop_cycle(cycler('color', ['m', 'g', 'b', 'y', 'c', 'r', 'k', 'g']))

    for i in range(len(gas_out[0, :])):
        if i != i_ar:
            if gas_out[:, i].max() > 5.e-3:
                #             print(gas_names[i])
                axs[0].plot(dist_array, gas_out[:, i], label=gas_names[i])
                species_name = gas_names[i]
                if species_name.endswith(')'):
                    if species_name[-3] == '(':
                        species_name = species_name[0:-3]
                    else:
                        species_name = species_name[0:-4]
                if species_name == "O2":
                    axs[0].annotate("O$_2$", fontsize=18, color='y',
                                    xy=(dist_array[1100], gas_out[:, i][1100] + gas_out[:, i][1100] / 100.0),
                                    va='bottom', ha='center')
                elif species_name == "CO2":
                    axs[0].annotate("CO$_2$", fontsize=18, color='c',
                                    xy=(dist_array[2400], gas_out[:, i][2400] + gas_out[:, i][2400] / 10.0), va='bottom',
                                    ha='center')
                elif species_name == "CO":
                    axs[0].annotate("CO", fontsize=18, color='g', xy=(dist_array[2100], gas_out[:, i][2100] + 0.001),
                                    va='bottom', ha='center')
                elif species_name == "H2":
                    axs[0].annotate("H$_2$", fontsize=18, color='k', xy=(dist_array[2200], gas_out[:, i][2200] - 0.001),
                                    va='top', ha='center')
                elif species_name == "CH4":
                    axs[0].annotate("CH$_4$", fontsize=18, color='b',
                                    xy=(dist_array[1100], gas_out[:, i][1100] + gas_out[:, i][1100] / 100.0),
                                    va='bottom', ha='center')
                elif species_name == "H2O":
                    axs[0].annotate("H$_2$O", fontsize=18, color='r',
                                    xy=(dist_array[2100], gas_out[:, i][2100] + gas_out[:, i][2100] / 40.0 + 0.001), va='bottom',
                                    ha='center')
                else:
                    axs[0].annotate(species_name, fontsize=18,
                                    xy=(dist_array[-1], gas_out[:, i][-1] + gas_out[:, i][-1] / 10.0), va='top',
                                    ha='center')
            else:
                axs[0].plot(0, 0)

    axs[1].set_prop_cycle(cycler('color', ['m', 'g', 'b', 'y', 'c', 'r', 'k', 'g']))
    ax2 = axs[0].twinx()
    ax2.plot(dist_array, T_array, label='temperature', color='r', linestyle=':')
    axs[0].set_prop_cycle(cycler('color', ['m', 'g', 'b', 'y', 'c', 'r', 'k', 'g']))

    axs[0].plot([dist_array[on_catalyst], dist_array[on_catalyst]], [0, 0.2], linestyle='--', color='xkcd:grey')
    axs[0].plot([dist_array[off_catalyst], dist_array[off_catalyst]], [0, 0.2], linestyle='--', color='xkcd:grey')
    axs[0].annotate("catalyst", fontsize=18, xy=(dist_array[on_catalyst], 0.175), va='bottom', ha='left')
    axs[1].plot([dist_array[on_catalyst], dist_array[on_catalyst]], [600.0, 2000], linestyle='--', color='xkcd:grey')
    axs[1].plot([dist_array[off_catalyst], dist_array[off_catalyst]], [600.0, 2000], linestyle='--', color='xkcd:grey')
    axs[1].annotate("catalyst", fontsize=18, xy=(dist_array[on_catalyst], 1800), va='bottom', ha='left')

    for item in (
            axs[0].get_xticklabels() + axs[0].get_yticklabels() + ax2.get_xticklabels() + ax2.get_yticklabels()):
        item.set_fontsize(18)

    axs[0].legend(loc='upper center', bbox_to_anchor=(0.5, -0.2), fancybox=False, shadow=False, ncol=4)
    axs[0].set_ylim(0., 0.1)
    axs[1].set_ylim(600.0, 2000)
    axs[0].set_xlim(8, 25)
    axs[1].set_xlim(8, 25)
    axs[0].set_xlabel('Distance (mm)', fontsize=22)
    axs[1].set_xlabel('Distance (mm)', fontsize=22)  # axs[0,1].set_xlabel('time (s)'); axs[1,1].set_xlabel('time (s)')
    axs[0].set_ylabel('flow/ mol/min', fontsize=22)
    ax2.set_ylabel('Temperature (K)', fontsize=22)
    ax2.set_ylim(600, 2000)
    ax2.set_xlim(8, 25)
    fig.delaxes(axs[1])  # THIS DELETES THE EXTRA SUBPLOT!
    fig.set_figheight(6)
    fig.set_figwidth(24)

    for n in range(len(gas_names)):
        if gas_names[n] == 'CH4(2)':
            c_in = gas_out[0][n]
        if gas_names[n] == 'O2(3)':
            o_in = gas_out[0][n]
    ratio = c_in / (o_in * 2)
    ratio = round(ratio, 1)

    out_dir = 'figures'
    os.path.exists(out_dir) or os.makedirs(out_dir)

def plotSurf(a):
    gas_out, surf_out, gas_names, surf_names, dist_array, T_array = a

    fig, axs = plt.subplots(1, 2)
    axs[0].set_prop_cycle(cycler('color', ['m', 'g', 'b', 'y', 'c', 'r', 'k', 'g']))

    for i in range(len(gas_out[0, :])):
        if i != i_ar:
            if gas_out[:, i].max() > 5.e-3:
                axs[0].plot(dist_array, gas_out[:, i], label=gas_names[i])
                species_name = gas_names[i]
                if species_name.endswith(')'):
                    if species_name[-3] == '(':
                        species_name = species_name[0:-3]
                    else:
                        species_name = species_name[0:-4]
                if species_name == "O2":
                    axs[0].annotate("O$_2$", fontsize=18,
                                    xy=(dist_array[2200], gas_out[:, i][2200] + gas_out[:, i][2200] / 100.0),
                                    va='bottom', ha='center')
                elif species_name == "CO2":
                    axs[0].annotate("CO$_2$", fontsize=18,
                                    xy=(dist_array[2200], gas_out[:, i][2200] + gas_out[:, i][2200] / 10.0), va='top',
                                    ha='center')
                elif species_name == "CO":
                    axs[0].annotate("CO", fontsize=18, xy=(dist_array[2200], gas_out[:, i][2200] + 0.001),
                                    va='bottom', ha='center')
                elif species_name == "CH2O":
                    axs[0].annotate("CH$_2$O", fontsize=18, xy=(dist_array[2200], gas_out[:, i][2200] + 0.001),
                                    va='bottom', ha='center')
                elif species_name == "CH4":
                    axs[0].annotate("CH$_4$", fontsize=18,
                                    xy=(dist_array[2200], gas_out[:, i][2200] + gas_out[:, i][2200] / 100.0),
                                    va='bottom', ha='center')
                elif species_name == "H2O":
                    axs[0].annotate("H$_2$O", fontsize=18,
                                    xy=(dist_array[2200], gas_out[:, i][2200] + gas_out[:, i][2200] / 40.0), va='top',
                                    ha='center')
                else:
                    axs[0].annotate(species_name, fontsize=18,
                                    xy=(dist_array[-1], gas_out[:, i][-1] + gas_out[:, i][-1] / 10.0), va='top',
                                    ha='center')
            else:
                axs[0].plot(0, 0)

    axs[1].set_prop_cycle(cycler('color', ['m', 'g', 'b', 'y', 'c', 'r', 'k', 'g']))
    # Plot two temperatures (of gas-phase and surface vs only surface.)
    for i in range(len(surf_out[0, :])):
        if surf_out[:, i].max() > 5.e-3:
            axs[1].semilogy(dist_array, surf_out[:, i], label=surf_names[i])
    axs[0].set_prop_cycle(cycler('color', ['m', 'g', 'b', 'y', 'c', 'r', 'k', 'g']))

    axs[0].plot([dist_array[on_catalyst], dist_array[on_catalyst]], [0, 0.2], linestyle='--', color='xkcd:grey')
    axs[0].plot([dist_array[off_catalyst], dist_array[off_catalyst]], [0, 0.2], linestyle='--', color='xkcd:grey')
    axs[0].annotate("catalyst", fontsize=18, xy=(dist_array[on_catalyst], 0.175), va='bottom', ha='left')
    axs[1].plot([dist_array[on_catalyst], dist_array[on_catalyst]], [0, 1.2], linestyle='--', color='xkcd:grey')
    axs[1].plot([dist_array[off_catalyst], dist_array[off_catalyst]], [0, 1.2], linestyle='--', color='xkcd:grey')
    axs[1].annotate("catalyst", fontsize=18, xy=(dist_array[on_catalyst], 1.1), va='bottom', ha='left')

    for item in (
            axs[0].get_xticklabels() + axs[0].get_yticklabels() + axs[1].get_xticklabels() + axs[1].get_yticklabels()):
        item.set_fontsize(18)

    axs[1].legend(loc='upper center', bbox_to_anchor=(0.5, -0.2), fancybox=False, shadow=False, ncol=2)
    axs[0].legend(loc='upper center', bbox_to_anchor=(0.5, -0.2), fancybox=False, shadow=False, ncol=4)
    axs[0].set_ylim(0., 0.1)
    axs[1].set_ylim(1e-10, 1.2)
    axs[0].set_xlim(5, 25)
    axs[1].set_xlim(9, 21)
    axs[0].set_xlabel('Distance (mm)', fontsize=22)
    axs[1].set_xlabel('Distance (mm)', fontsize=22)
    axs[0].set_ylabel('flow/ mol/min', fontsize=22)
    axs[1].set_ylabel('Site fraction', fontsize=22)
    fig.delaxes(axs[0])  # THIS DELETES THE EXTRA SUBPLOT!

    fig.set_figheight(6)
    fig.set_figwidth(18)

    for n in range(len(gas_names)):
        if gas_names[n] == 'CH4(2)':
            c_in = gas_out[0][n]
        if gas_names[n] == 'O2(3)':
            o_in = gas_out[0][n]
    ratio = c_in / (o_in * 2)
    ratio = round(ratio, 1)

    out_dir = 'figures'
    os.path.exists(out_dir) or os.makedirs(out_dir)
    fig.clf()

def monolithFull(gas, surf, temp, mol_in, verbose=False, sens=False):
    """
    Verbose prints out values as you go along
    Sens is for sensitivity, in the form [perturbation, reaction #]
    """
    ch4, o2, ar = mol_in
    ratio = ch4/(2*o2)
    ratio = round(ratio, 1)
    ch4 = str(ch4)
    o2 = str(o2)
    ar = str(ar)
    X = str('CH4(2):' + ch4 + ', O2(3):' + o2 + ', Ar:' + ar)
    gas.TPX = 273.15, ct.one_atm, X  # need to initialize mass flow rate at STP
    mass_flow_rate = flow_rate * gas.density_mass
    gas.TPX = temp, ct.one_atm, X
    temp_cat = temp
    surf.TP = temp_cat, ct.one_atm
    surf.coverages = 'X(1):1.0'
    gas.set_multiplier(1.0)

    TDY = gas.TDY
    cov = surf.coverages

    if verbose is True:
        print('  distance(mm)   X_CH4        X_O2        X_H2       X_CO       X_H2O       X_CO2')

    # create a new reactor
    gas.TDY = TDY
    r = ct.IdealGasReactor(gas)
    r.volume = rvol

    # create a reservoir to represent the reactor immediately upstream. Note
    # that the gas object is set already to the state of the upstream reactor
    upstream = ct.Reservoir(gas, name='upstream')

    # create a reservoir for the reactor to exhaust into. The composition of
    # this reservoir is irrelevant.
    downstream = ct.Reservoir(gas, name='downstream')

    # Add the reacting surface to the reactor. The area is set to the desired
    # catalyst area in the reactor.
    rsurf = ct.ReactorSurface(surf, r, A=cat_area)

    # The mass flow rate into the reactor will be fixed by using a
    # MassFlowController object.
    # mass_flow_rate = velocity * gas.density_mass * area  # kg/s
    # mass_flow_rate = flow_rate * gas.density_mass
    m = ct.MassFlowController(upstream, r, mdot=mass_flow_rate)

    # We need an outlet to the downstream reservoir. This will determine the
    # pressure in the reactor. The value of K will only affect the transient
    # pressure difference.
    v = ct.PressureController(r, downstream, master=m, K=1e-5)

    sim = ct.ReactorNet([r])
    sim.max_err_test_fails = 12

    # set relative and absolute tolerances on the simulation
    sim.rtol = 1.0e-10
    sim.atol = 1.0e-20

    gas_names = gas.species_names
    surf_names = surf.species_names
    gas_out = []
    surf_out = []
    dist_array = []
    T_array = []

    surf.set_multiplier(0.0)  # no surface reactions until the gauze
    for n in range(NReactors):
        # Set the state of the reservoir to match that of the previous reactor
        gas.TDY = r.thermo.TDY
        upstream.syncState()
        if n == on_catalyst:
            surf.set_multiplier(1.0)
            if sens is not False:
                surf.set_multiplier(1.0 + sens[0], sens[1])
        if n == off_catalyst:
            surf.set_multiplier(0.0)
        sim.reinitialize()
        sim.advance_to_steady_state()
        dist = n * reactor_len * 1.0e3  # distance in mm
        dist_array.append(dist)
        T_array.append(surf.T)
        # print "mass_flow_rate", mass_flow_rate,  v.mdot(sim.time), "kg/s"
        kmole_flow_rate = mass_flow_rate / gas.mean_molecular_weight  # kmol/s
        gas_out.append(1000 * 60 * kmole_flow_rate * gas.X.copy())  # molar flow rate in moles/minute
        surf_out.append(surf.X.copy())

        # make reaction diagrams
        out_dir = 'rxnpath'
        os.path.exists(out_dir) or os.makedirs(out_dir)
        elements = ['H', 'O']
        locations_of_interest = [1000, 1150, 1160, 1183, 1196, 1999]
        if sens is False:
            for l in locations_of_interest:
                if n == l:
                    location = str(n / 100)

                    diagram = ct.ReactionPathDiagram(surf, 'X')
                    diagram.title = 'rxn path'
                    diagram.label_threshold = 1e-9
                    dot_file = out_dir + '/rxnpath-' + str(ratio) + '-x-' + location + 'mm.dot'
                    img_file = out_dir + '/rxnpath-' + str(ratio) + '-x-' + location + 'mm.png'
                    img_path = os.path.join(out_dir, img_file)
                    diagram.write_dot(dot_file)
                    os.system('dot {0} -Tpng -o{1} -Gdpi=200'.format(dot_file, img_file))

                    for element in elements:
                        diagram = ct.ReactionPathDiagram(surf, element)
                        diagram.title = element + 'rxn path'
                        diagram.label_threshold = 1e-9
                        dot_file = out_dir + '/rxnpath-' + str(ratio) + '-surf-' + location + 'mm-' + element + '.dot'
                        img_file = out_dir + '/rxnpath-' + str(ratio) + '-surf-' + location + 'mm-' + element + '.png'
                        img_path = os.path.join(out_dir, img_file)
                        diagram.write_dot(dot_file)
                        os.system('dot {0} -Tpng -o{1} -Gdpi=200'.format(dot_file, img_file))
        else:
            pass

        if verbose is True:
            if not n % 100:
                print('  {0:10f}  {1:10f}  {2:10f}  {3:10f} {4:10f} {5:10f} {6:10f}'.format(dist, *gas[
                    'CH4(2)', 'O2(3)', 'H2(6)', 'CO(7)', 'H2O(5)', 'CO2(4)'].X * 1000 * 60 * kmole_flow_rate))

    gas_out = np.array(gas_out)
    surf_out = np.array(surf_out)
    gas_names = np.array(gas_names)
    surf_names = np.array(surf_names)
    data_out = gas_out, surf_out, gas_names, surf_names, dist_array, T_array
    return data_out

ratio = 1.0
fo2 = 1 / (2. * ratio + 1 + 79 / 21)
fch4 = 2 * fo2 * ratio
far = 79 * fo2 / 21
ratio_in = [fch4, fo2, far]

a = monolithFull(gas, surf, t_in, ratio_in)
gas_out, surf_out, gas_names, surf_names, dist_array, T_array = a

def deriv(p):
    deriv = []
    for x in range(7000):
        deriv.append((p[x+1] - p[x])/.01)
    deriv.append(0.)
    return deriv

methane = gas_out[:,3]
oxygen = gas_out[:,4]
hydrogen = gas_out[:,7]
co2 = gas_out[:,5]
co = gas_out[:,8]
h2o = gas_out[:,6]

d_ch4 = deriv(methane)
d_o2 = deriv(oxygen)
d_h2 = deriv(hydrogen)
d_co2 = deriv(co2)
d_co = deriv(co)
d_h2o = deriv(h2o)

plt.figure(figsize=(6,5))
plt.plot(dist_array,d_ch4,color='green')
plt.plot(dist_array,d_o2,color='orange')
plt.plot(dist_array,d_h2,color='k')
plt.plot(dist_array,d_co2,color='blue')
plt.plot(dist_array,d_co,color='limegreen')
plt.plot(dist_array,d_h2o,color='dodgerblue')

plt.annotate("CH$_4$", fontsize=18, color='g',
            xy=(dist_array[1140], -.15),
            va='bottom', ha='center')
plt.annotate("O$_2$", fontsize=18, color='orange',
            xy=(dist_array[1113], -.17),
            va='bottom', ha='center')
plt.annotate("H$_2$", fontsize=18, color='k',
            xy=(dist_array[1144], .15),
            va='bottom', ha='center')
plt.annotate("CO", fontsize=18, color='limegreen',
            xy=(dist_array[1137], .03),
            va='bottom', ha='center')
plt.annotate("CO$_2$", fontsize=18, color='blue',
            xy=(dist_array[1125], -0.03),
            va='bottom', ha='center')
plt.annotate("H$_2$O", fontsize=18, color='dodgerblue',
            xy=(dist_array[1105], .1),
            va='bottom', ha='center')

plt.xlabel('Position (mm)', fontsize=22)
plt.ylabel('d Flow/dx', fontsize=22)
plt.xticks(fontsize=18)
plt.yticks(fontsize=18)
plt.xticks(np.arange(10.5,12.5,0.5))
plt.yticks(np.arange(-.2,.25, .1))
plt.xlim(10.5,12.)
plt.savefig('./paperplots/deriv.pdf',bbox_inches='tight',dpi=300)
plt.clf()

temp_profile = np.array([
        [0.,373.15],
        [1.,373.15],
        [2.,373.15],
        [3.,373.15],
        [4.,373.15],
        [5.,373.15],
        [6.,413.15],
        [7.,448.15],
        [8.,473.15],
        [9.,548.15],
        [10.,648.15],
        [11.,963.15],
        [12.,1133.15],
        [13.,1173.15],
        [14.,1163.15],
        [15.,1123.15],
        [16.,1098.15],
        [17.,1073.15],
        [18.,1068.15],
        [19.,1063.15],
        [20.,1048.15],
        [21.,1033.15],
        [22.,1028.15],
        [23.,1023.15],
        [24.,1023.15],
        [25.,1023.15],])

ch4_in = 21/71*.208
ch4_profile = np.array([
        [0.,ch4_in],
        [1.,ch4_in],
        [2.,ch4_in],
        [3.,ch4_in],
        [4.,ch4_in],
        [5.,ch4_in],
        [6.,ch4_in],
        [7.,ch4_in],
        [8.,ch4_in],
        [9.,ch4_in],
        [10.,ch4_in],
        [10.25,0.059],
        [10.5,0.048],
        [10.75,0.041],
        [11.,0.036],
        [11.5,0.028],
        [12.,0.0225],
        [12.25,0.022],
        [12.5,0.021],
        [13.,0.021],
        [13.5,0.022],
        [14.,0.022],
        [14.5,0.021],
        [15.,0.019],
        [15.5,0.017],
        [16.,0.016],
        [17.,0.016],
        [18.,0.015],
        [19.,0.015],
        [20.,0.015],
        [21.,0.014],
        [22.,0.014],
        [23.,0.014],
        [24.,0.014],
        [25.,0.014],])

o2_in = (1 / (2. + 1 + 79 / 21)) * .208
o2_profile = np.array([
        [0.,o2_in],
        [1.,o2_in],
        [2.,o2_in],
        [3.,o2_in],
        [4.,o2_in],
        [5.,o2_in],
        [6.,o2_in],
        [7.,o2_in],
        [8.,o2_in],
        [9.,o2_in],
        [9.5,o2_in],
        [9.75,o2_in],
        [10.,o2_in],
        [10.1,0.028],
        [10.25,0.022],
        [10.5,0.013],
        [10.75,0.008],
        [11.,0.004],
        [11.25,0.001],
        [11.5,0.],
        [12.,0.],
        [13.,0.],
        [14.,0.],
        [15.,0.],
        [16.,0.],
        [17.,0.],
        [18.,0.],
        [19.,0.],
        [20.,0.],
        [21.,0.],
        [22.,0.],
        [23.,0.],
        [24.,0.],
        [25.,0.],])

h2_profile = np.array([
        [0.,0],
        [1.,0],
        [2.,0],
        [3.,0],
        [4.,0],
        [5.,0],
        [6.,0],
        [7.,0],
        [8.,0],
        [9.,0],
        [10.,0],
        [10.5,0.015],
        [11.,0.034],
        [12.,0.051],
        [13.,0.065],
        [14.,0.071],
        [15.,0.083],
        [16.,0.089],
        [17.,0.0895],
        [18.,0.09],
        [19.,0.0905],
        [20.,0.091],
        [21.,0.092],
        [22.,0.092],
        [23.,0.092],
        [24.,0.092],
        [25.,0.092],])

co_profile = np.array([
        [0.,0],
        [1.,0],
        [2.,0],
        [3.,0],
        [4.,0],
        [5.,0],
        [6.,0],
        [7.,0],
        [8.,0],
        [9.,0],
        [10.,0],
        [10.5,0.015],
        [11.,0.026],
        [12.,0.035],
        [13.,0.04],
        [14.,0.036],
        [15.,0.04],
        [16.,0.044],
        [17.,0.045],
        [18.,0.045],
        [19.,0.045],
        [20.,0.045],
        [21.,0.045],
        [22.,0.045],
        [23.,0.045],
        [24.,0.045],
        [25.,0.045],])

h2o_profile = np.array([
        [0.,0],
        [1.,0],
        [2.,0],
        [3.,0],
        [4.,0],
        [5.,0],
        [6.,0],
        [7.,0],
        [8.,0],
        [9.,0],
        [10.,0],
        [10.5,0.015],
        [11.,0.025],
        [11.5,0.026],
        [12.,0.025],
        [12.5,0.022],
        [13.,0.019],
        [14.,0.011],
        [15.,0.009],
        [16.,0.009],
        [17.,0.009],
        [18.,0.008],
        [19.,0.008],
        [20.,0.008],
        [21.,0.008],
        [22.,0.008],
        [23.,0.008],
        [24.,0.008],
        [25.,0.008],])

co2_profile = np.array([
        [0.,0],
        [1.,0],
        [2.,0],
        [3.,0],
        [4.,0],
        [5.,0],
        [6.,0],
        [7.,0],
        [8.,0],
        [9.,0],
        [10.,0],
        [11.,0.002],
        [12.,0.002],
        [12.5,0.002],
        [13.,0.002],
        [14.,0.002],
        [15.,0.002],
        [16.,0.002],
        [17.,0.002],
        [18.,0.002],
        [19.,0.002],
        [20.,0.002],
        [21.,0.002],
        [22.,0.002],
        [23.,0.002],
        [24.,0.002],
        [25.,0.002],])

from scipy.interpolate import interp1d

dist = np.linspace(0,25,2501)
ch4 = interp1d(ch4_profile[:,0],ch4_profile[:,1],kind='cubic')
o2 = interp1d(o2_profile[:,0],o2_profile[:,1],kind='cubic')
h2 = interp1d(h2_profile[:,0],h2_profile[:,1],kind='cubic')
co = interp1d(co_profile[:,0],co_profile[:,1],kind='cubic')
h2o = interp1d(h2o_profile[:,0],h2o_profile[:,1],kind='cubic')
co2 = interp1d(co2_profile[:,0],co2_profile[:,1],kind='cubic')
temp = interp1d(temp_profile[:,0],temp_profile[:,1],kind='cubic')

plt.plot(dist,ch4(dist),color='g')
plt.annotate("CH$_4$", fontsize=18, color='g',
            xy=(dist_array[700], .062),
            va='bottom', ha='center')
plt.plot(dist,o2(dist),color='orange')
plt.annotate("O$_2$", fontsize=18, color='orange',
            xy=(dist_array[700], .032),
            va='bottom', ha='center')
plt.plot(dist,h2(dist),color='k')
plt.annotate("H$_2$", fontsize=18, color='k',
            xy=(dist_array[2300], .072),
            va='bottom', ha='center')
plt.plot(dist,co(dist),color='limegreen')
plt.annotate("CO", fontsize=18, color='limegreen',
            xy=(dist_array[2300], .03),
            va='bottom', ha='center')
plt.plot(dist,co2(dist),color='blue')
plt.annotate("CO$_2$", fontsize=18, color='blue',
            xy=(dist_array[2200], .002),
            va='bottom', ha='center')
plt.plot(dist,h2o(dist),color='dodgerblue')
plt.annotate("H$_2$O", fontsize=18, color='dodgerblue',
            xy=(dist_array[1600], .01),
            va='bottom', ha='center')
plt.plot([dist_array[on_catalyst], dist_array[on_catalyst]], [0, 0.1], linestyle='--', color='xkcd:grey')
plt.plot([dist_array[off_catalyst], dist_array[off_catalyst]], [0, 0.1], linestyle='--', color='xkcd:grey')
plt.annotate("catalyst", fontsize=18, xy=(dist_array[1500], 0.095), va='center', ha='center')

plt.ylabel('Flow (mol/min)', fontsize=22)
plt.xlabel('Position (mm)', fontsize=22)
plt.ylim((0,0.1))
plt.xlim((5,25))
plt.xticks(fontsize=18)
plt.yticks(fontsize=18)
plt.xticks(np.arange(6,26,2))

ax2 = plt.twinx()
ax2.plot(dist,temp(dist), ':', color='r')
ax2.set_ylim(300, 2000)
ax2.annotate("T",fontsize=18, color='r',
             xy=(dist_array[2300], 1310),
             va='top', ha='center')
ax2.set_ylabel('Temperature (K)', fontsize=22)
for item in (ax2.get_yticklabels()):
             item.set_fontsize(18)
plt.savefig('./paperplots/horn_data.pdf',bbox_inches='tight',dpi=300)
plt.clf()

# plotting ch4
plt.plot(dist_array,gas_out[:,3], label=gas_names[3],color='g')
plt.annotate("CH$_4$", fontsize=18, color='g',
            xy=(dist_array[700], gas_out[:, 3][500] + gas_out[:, 3][500] / 100.0),
            va='bottom', ha='center')
plt.plot(dist,ch4(dist),'--',color='g')
# plotting o2
plt.plot(dist_array,gas_out[:,4], label=gas_names[4],color='orange')
plt.annotate("O$_2$", fontsize=18, color='orange',
            xy=(dist_array[700], gas_out[:, 4][500] + gas_out[:, 4][500] / 100.0),
            va='bottom', ha='center')
plt.plot(dist,o2(dist),'--',color='orange')

plt.plot([dist_array[on_catalyst], dist_array[on_catalyst]], [0, 0.1], linestyle=':', color='xkcd:grey')
plt.plot([dist_array[off_catalyst], dist_array[off_catalyst]], [0, 0.1], linestyle=':', color='xkcd:grey')
plt.annotate("catalyst", fontsize=18, xy=(dist_array[1500], 0.095), va='center', ha='center')

plt.ylabel('Flow (mol/min)', fontsize=22)
plt.xlabel('Position (mm)', fontsize=22)
plt.ylim((0,0.1))
plt.xticks(fontsize=18)
plt.yticks(fontsize=18)
plt.xticks(np.arange(6,26,2))

ax2 = plt.twinx()
ax2.plot(dist,temp(dist), '--', color='r')
ax2.plot(dist_array,T_array,color='r')
ax2.annotate("T",fontsize=18, color='r',
             xy=(dist_array[2300], T_array[2300] -150),
             va='top', ha='center')
ax2.set_ylim(300, 2000)
ax2.set_ylabel('Temperature (K)', fontsize=22)
for item in (ax2.get_yticklabels()):
             item.set_fontsize(18)

plt.xlim((5,25))  # cutting off where the horn paper cuts off the plots
plt.savefig('./paperplots/compare_ch4_o2_temp.pdf',bbox_inches='tight',dpi=300)
plt.clf()

# plotting h2
plt.plot(dist_array,gas_out[:,7], label=gas_names[7],color='k')
plt.annotate("H$_2$", fontsize=18, color='k',
            xy=(dist_array[2200], gas_out[:, 7][2200] + gas_out[:, 7][2200] / 100.0),
            va='bottom', ha='center')
plt.plot(dist,h2(dist),'--',color='k')
# plotting co
plt.plot(dist_array,gas_out[:,8], label=gas_names[8],color='limegreen')
plt.annotate("CO", fontsize=18, color='limegreen',
            xy=(dist_array[1800], gas_out[:, 8][1800] + gas_out[:, 8][1800] / 100.0),
            va='bottom', ha='center')
plt.plot(dist,co(dist),'--',color='limegreen')

plt.plot([dist_array[on_catalyst], dist_array[on_catalyst]], [0, 0.1], linestyle=':', color='xkcd:grey')
plt.plot([dist_array[off_catalyst], dist_array[off_catalyst]], [0, 0.1], linestyle=':', color='xkcd:grey')
plt.annotate("catalyst", fontsize=18, xy=(dist_array[1500], 0.095), va='center', ha='center')

plt.ylabel('Flow (mol/min)', fontsize=22)
plt.xlabel('Position (mm)', fontsize=22)
plt.ylim((0,0.1))
plt.xticks(fontsize=18)
plt.yticks(fontsize=18)
plt.xticks(np.arange(6,26,2))

plt.xlim((5,25))  # cutting off where the horn paper cuts off the plots
plt.savefig('./paperplots/compare_h2_co.pdf',bbox_inches='tight',dpi=300)
plt.clf()

# plotting h2o
plt.plot(dist_array,gas_out[:,6], label=gas_names[6],color='dodgerblue')
plt.annotate("H$_2$O", fontsize=18, color='dodgerblue',
            xy=(dist_array[2200], gas_out[:, 6][2200] -.01),
            va='bottom', ha='center')
plt.plot(dist,h2o(dist),'--',color='dodgerblue')
# plotting co2
plt.plot(dist_array,gas_out[:,5], label=gas_names[5],color='blue')
plt.annotate("CO$_2$", fontsize=18, color='blue',
            xy=(dist_array[1400], .009),
            va='bottom', ha='center')
plt.plot(dist,co2(dist),'--',color='blue')

plt.plot([dist_array[on_catalyst], dist_array[on_catalyst]], [0, 0.1], linestyle=':', color='xkcd:grey')
plt.plot([dist_array[off_catalyst], dist_array[off_catalyst]], [0, 0.1], linestyle=':', color='xkcd:grey')
plt.annotate("catalyst", fontsize=18, xy=(dist_array[1500], 0.095), va='center', ha='center')

plt.ylabel('Flow (mol/min)', fontsize=22)
plt.xlabel('Position (mm)', fontsize=22)
plt.ylim((0,0.1))
plt.xticks(fontsize=18)
plt.yticks(fontsize=18)
plt.xticks(np.arange(6,26,2))

plt.xlim((5,25))  # cutting off where the horn paper cuts off the plots
plt.savefig('./paperplots/compare_h2o_co2.pdf',bbox_inches='tight',dpi=300)
plt.clf()

# plotting ch4
plt.plot(dist_array,gas_out[:,3], label=gas_names[3],color='g')
plt.annotate("CH$_4$", fontsize=18, color='g',
            xy=(dist_array[700], gas_out[:, 3][500] + gas_out[:, 3][500] / 100.0),
            va='bottom', ha='center')
plt.plot(dist,ch4(dist),'--',color='g')
# plotting o2
plt.plot(dist_array,gas_out[:,4], label=gas_names[4],color='orange')
plt.annotate("O$_2$", fontsize=18, color='orange',
            xy=(dist_array[700], gas_out[:, 4][500] + gas_out[:, 4][500] / 100.0),
            va='bottom', ha='center')
plt.plot(dist,o2(dist),'--',color='orange')
# plotting h2
plt.plot(dist_array,gas_out[:,7], label=gas_names[7],color='k')
plt.annotate("H$_2$", fontsize=18, color='k',
            xy=(dist_array[2200], gas_out[:, 7][2200] + gas_out[:, 7][2200] / 100.0),
            va='bottom', ha='center')
plt.plot(dist,h2(dist),'--',color='k')
# plotting co
plt.plot(dist_array,gas_out[:,8], label=gas_names[8],color='limegreen')
plt.annotate("CO", fontsize=18, color='limegreen',
            xy=(dist_array[1800], gas_out[:, 8][1800] + gas_out[:, 8][1800] / 100.0),
            va='bottom', ha='center')
plt.plot(dist,co(dist),'--',color='limegreen')
# plotting h2o
plt.plot(dist_array,gas_out[:,6], label=gas_names[6],color='dodgerblue')
plt.annotate("H$_2$O", fontsize=18, color='dodgerblue',
            xy=(dist_array[2200], gas_out[:, 6][2200] -.01),
            va='bottom', ha='center')
plt.plot(dist,h2o(dist),'--',color='dodgerblue')
# plotting co2
plt.plot(dist_array,gas_out[:,5], label=gas_names[5],color='blue')
plt.annotate("CO$_2$", fontsize=18, color='blue',
            xy=(dist_array[1400], .009),
            va='bottom', ha='center')
plt.plot(dist,co2(dist),'--',color='blue')

plt.plot([dist_array[on_catalyst], dist_array[on_catalyst]], [0, 0.1], linestyle=':', color='xkcd:grey')
plt.plot([dist_array[off_catalyst], dist_array[off_catalyst]], [0, 0.1], linestyle=':', color='xkcd:grey')
plt.annotate("catalyst", fontsize=18, xy=(dist_array[1500], 0.095), va='center', ha='center')

plt.ylabel('Flow (mol/min)', fontsize=22)
plt.xlabel('Position (mm)', fontsize=22)
plt.ylim((0,0.1))
plt.xticks(fontsize=18)
plt.yticks(fontsize=18)
plt.xticks(np.arange(6,26,2))

plt.xlim((5,25))  # cutting off where the horn paper cuts off the plots
plt.savefig('./paperplots/all_compare.pdf',bbox_inches='tight',dpi=300)

plt.clf()

fig, axs = plt.subplots(3, 1)

# plotting ch4
axs[0].plot(dist_array,gas_out[:,3], label=gas_names[3],color='g')
axs[0].annotate("CH$_4$", fontsize=18, color='g',
            xy=(dist_array[700], gas_out[:, 3][500] + gas_out[:, 3][500] / 100.0),
            va='bottom', ha='center')
axs[0].plot(dist,ch4(dist),'--',color='g')
# plotting o2
axs[0].plot(dist_array,gas_out[:,4], label=gas_names[4],color='orange')
axs[0].annotate("O$_2$", fontsize=18, color='orange',
            xy=(dist_array[700], gas_out[:, 4][500] + gas_out[:, 4][500] / 100.0),
            va='bottom', ha='center')
axs[0].plot(dist,o2(dist),'--',color='orange')

axs[0].plot([dist_array[on_catalyst], dist_array[on_catalyst]], [0, 0.1], linestyle=':', color='xkcd:grey')
axs[0].plot([dist_array[off_catalyst], dist_array[off_catalyst]], [0, 0.1], linestyle=':', color='xkcd:grey')
axs[0].annotate("catalyst", fontsize=18, xy=(dist_array[1500], 0.095), va='center', ha='center')

axs[0].set_ylabel('Flow (mol/min)', fontsize=22)
axs[0].set_ylim((0,0.1))
for item in (axs[0].get_xticklabels()):
             item.set_fontsize(18)
for item in (axs[0].get_yticklabels()):
             item.set_fontsize(18)
axs[0].set_xticks(np.arange(6,26,2))

ax2 = axs[0].twinx()
ax2.plot(dist,temp(dist), '--', color='r')
ax2.plot(dist_array,T_array,color='r')
ax2.annotate("T (K)",fontsize=18, color='r',
             xy=(dist_array[2300], T_array[2300] -150),
             va='top', ha='center')
ax2.set_ylim(300, 2000)
for item in (ax2.get_yticklabels()):
             item.set_fontsize(18)
axs[0].set_xlim((5,25))

# plotting h2
axs[1].plot(dist_array,gas_out[:,7], label=gas_names[7],color='k')
axs[1].annotate("H$_2$", fontsize=18, color='k',
            xy=(dist_array[2200], 0.0575),
            va='bottom', ha='center')
axs[1].plot(dist,h2(dist),'--',color='k')
# plotting co
axs[1].plot(dist_array,gas_out[:,8], label=gas_names[8],color='limegreen')
axs[1].annotate("CO", fontsize=18, color='limegreen',
            xy=(dist_array[1800], gas_out[:, 8][1800] + gas_out[:, 8][1800] / 100.0 + 0.005),
            va='bottom', ha='center')
axs[1].plot(dist,co(dist),'--',color='limegreen')

axs[1].plot([dist_array[on_catalyst], dist_array[on_catalyst]], [0, 0.1], linestyle=':', color='xkcd:grey')
axs[1].plot([dist_array[off_catalyst], dist_array[off_catalyst]], [0, 0.1], linestyle=':', color='xkcd:grey')

axs[1].set_ylabel('Flow (mol/min)', fontsize=22)
axs[1].set_ylim((0,0.1))
for item in (axs[1].get_xticklabels()):
             item.set_fontsize(18)
for item in (axs[1].get_yticklabels()):
             item.set_fontsize(18)
axs[1].set_xticks(np.arange(6,26,2))

axs[1].set_xlim((5,25))  # cutting off where the horn paper cuts off the plots

# plotting h2o
axs[2].plot(dist_array,gas_out[:,6], label=gas_names[6],color='dodgerblue')
axs[2].annotate("H$_2$O", fontsize=18, color='dodgerblue',
            xy=(dist_array[2200], 0.02),
            va='bottom', ha='center')
axs[2].plot(dist,h2o(dist),'--',color='dodgerblue')
# plotting co2
axs[2].plot(dist_array,gas_out[:,5], label=gas_names[5],color='blue')
axs[2].annotate("CO$_2$", fontsize=18, color='blue',
            xy=(dist_array[1225], .005),
            va='bottom', ha='center')
axs[2].plot(dist,co2(dist),'--',color='blue')

axs[2].plot([dist_array[on_catalyst], dist_array[on_catalyst]], [0, 0.1], linestyle=':', color='xkcd:grey')
axs[2].plot([dist_array[off_catalyst], dist_array[off_catalyst]], [0, 0.1], linestyle=':', color='xkcd:grey')

axs[2].set_ylabel('Flow (mol/min)', fontsize=22)
axs[2].set_ylim((0,0.1))
for item in (axs[2].get_xticklabels()):
             item.set_fontsize(18)
for item in (axs[2].get_yticklabels()):
             item.set_fontsize(18)
axs[2].set_xticks(np.arange(6,26,2))

axs[2].set_xlim((5,25))  # cutting off where the horn paper cuts off the plots

axs[2].set_xlabel('Position (mm)', fontsize=22)
fig.set_figheight(10)
fig.set_figwidth(7)
plt.savefig('./paperplots/all_compare_subplots.pdf',bbox_inches='tight',dpi=300)
plt.clf()

def simulationWorker(ratio):
    fo2 = 1 / (2. * ratio + 1 + 79 / 21)
    fch4 = 2 * fo2 * ratio
    far = 79 * fo2 / 21
    ratio_in = [fch4, fo2, far]  # mol fractions

    a = monolithFull(gas, surf, t_in, ratio_in)
    print("Finished simulation at a C/O ratio of {:.1f}".format(ratio))
    gas_out, surf_out, gas_names, surf_names, dist_array, T_array = a
    return [ratio, [gas_out, gas_names, dist_array, T_array]]

ratios = [.6, .7, .8, .9, 1., 1.1, 1.2, 1.3, 1.4, 1.6, 1.8, 2., 2.2, 2.4, 2.6]  # 15 items
data = []
num_threads = len(ratios)
pool = multiprocessing.Pool(processes=num_threads)
data = pool.map(simulationWorker, ratios, 1)
pool.close()
pool.join()

end_temp = []
max_temp = []
dist_max_temp = []
ch4_in = []
ch4_out = []
ch4_conv = []
o2_conv = []
co_sel = []
co_out = []
h2_sel = []
h2_out = []
h2o_sel = []
h2o_out = []
co2_sel = []
co2_out = []
ratios_real = []
for r in data:
    for x in range(len(r[1][1])):
        if r[1][1][x] == 'CH4(2)':
            ch4_i = r[1][0][0][x]
            ch4_in.append(ch4_i)
            ch4_o = r[1][0][-1][x]
            ch4_out.append(ch4_o)
            ch4_depletion = ch4_i - ch4_o
            ch4_conv.append(ch4_depletion / ch4_i)
        if r[1][1][x] == 'O2(3)':
            o2_in = r[1][0][0][x]
            o2_out = r[1][0][-1][x]
            o2_conv.append((o2_in - o2_out) / o2_in)
    ratios_real.append(ch4_i / (2 * o2_in))
    end_temp.append(r[1][3][-1])
    max_temp.append(max(r[1][3]))
    dist_max_temp.append(r[1][2][r[1][3].index(max(r[1][3]))])

    for x in range(len(r[1][1])):
        if r[1][1][x] == 'Ar':
            ar = r[1][0][-1][x]
        if r[1][1][x] == 'CO(7)':
            co_o = r[1][0][-1][x]
            co_out.append(co_o)
            co_sel.append(co_o / ch4_depletion)
        if r[1][1][x] == 'H2O(5)':
            h2o_o = r[1][0][-1][x]
            h2o_out.append(h2o_o)
            h2o_sel.append(h2o_o / (ch4_depletion * 2))
        if r[1][1][x] == 'H2(6)':
            h2_o = r[1][0][-1][x]
            h2_out.append(h2_o)
            h2_sel.append(h2_o / (ch4_depletion * 2))
        if r[1][1][x] == 'CO2(4)':
            co2_o = r[1][0][-1][x]
            co2_out.append(co2_o)
            co2_sel.append(co2_o / ch4_depletion)

pch4_conv = [x *100 for x in ch4_conv]
po2_conv = [x * 100 for x in o2_conv]
# #horn data
o2_horn = [100] * 15
ch4_horn = [100,99,95,88,78,69,62,57,50,44,37,32,27,24,21]
t_horn_c = [1175,1000,820,790,780,775,770,768,765,760,755,750,740,730,725] #celcius
t_horn = [x + 275.13 for x in t_horn_c]

plt.figure(figsize=(6,5))
plt.plot(ratios_real, pch4_conv, 'bo-', label='CH4', color='g')
plt.annotate("CH$_4$", fontsize=18, color='g',xy=(1.6,60), va='top', ha='center')
plt.plot(ratios_real,ch4_horn, '--', color='g')
plt.plot(ratios_real, po2_conv, 'bo-', label='O2', color='orange')
plt.annotate("O$_2$", fontsize=18, color='orange',xy=(1.8,95), va='top', ha='center')
plt.plot(ratios_real,o2_horn, '--', color='orange')
plt.yticks(np.arange(0,106,10))
plt.ylabel('Exit conversion (%)', fontsize=22)
plt.xlabel('C/O Ratio', fontsize=22)
plt.xticks(np.arange(.6,2.7,0.4))
plt.xticks(fontsize=18)
plt.yticks(fontsize=18)

ax2 = plt.twinx()
ax2.plot(ratios_real, end_temp, 'bo-', color='r')
ax2.set_ylim(900, 1900)
ax2.set_yticks(np.arange(900,2200,100))
for item in (ax2.get_yticklabels()):
             item.set_fontsize(13)
ax2.annotate("T",fontsize=18, color='r',
             xy=(1.3, 1150),
             va='top', ha='center')
ax2.set_ylabel('Temperature (K)', fontsize=22)
for item in (ax2.get_yticklabels()):
             item.set_fontsize(18)
ax2.plot(ratios_real,t_horn,'--',color='r')
plt.tight_layout()
plt.savefig('./paperplots/conversion.pdf', bbox_inches='tight', dpi=300)
plt.clf()

pco_sel = [x *100 for x in co_sel]
ph2_sel = [x * 100 for x in h2_sel]
pco2_sel = [x * 100 for x in co2_sel]
ph2o_sel = [x * 100 for x in h2o_sel]

# horn data
horn_co = [88,90,93,94,95,95,94,94,93,92,91,90,89,88,87]
horn_h2 = [79,88,95,94,93,92,91,90,89,88,85,84,80,78,76]
horn_h2o = [21,12,5,6,7,8,9,10,11,12,15,16,20,22,24]
horn_co2 = [12,10,7,6,5,5,6,6,7,8,9,10,11,12,13]

plt.figure(figsize=(6,5))
plt.plot(ratios_real, pco_sel, 'bo-', label='CO', color='limegreen')
plt.annotate("CO", fontsize=18, color='limegreen',xy=(1.8,98), va='top', ha='center')
plt.plot(ratios_real,horn_co,'--',color='limegreen')
plt.plot(ratios_real, ph2_sel, 'bo-', label='H2', color='k')
plt.annotate("H$_2$", fontsize=18, color='k',xy=(0.8,75), va='top', ha='center')
plt.plot(ratios_real, pco2_sel, 'bo-', label='CO2', color='blue')
plt.plot(ratios_real,horn_h2,'--',color='k')
plt.annotate("CO$_2$", fontsize=18, color='blue',xy=(1.8,8), va='top', ha='center')
plt.plot(ratios_real,horn_co2,'--',color='blue')
plt.plot(ratios_real, ph2o_sel, 'bo-', label='H2O', color='dodgerblue')
plt.annotate("H$_2$O", fontsize=18, color='dodgerblue',xy=(1.,25), va='top', ha='center')
plt.ylabel('Exit Selectivity (%)', fontsize=18)
plt.plot(ratios_real,horn_h2o,'--',color='dodgerblue')
plt.xlabel('C/O Ratio', fontsize=22)
plt.xticks(np.arange(.6,2.7,0.4))
plt.xticks(fontsize=18)
plt.yticks(fontsize=18)
plt.yticks(np.arange(0,110,10))
plt.tight_layout()
plt.savefig('./paperplots/selectivity_comparison.pdf', bbox_inches='tight', dpi=300)
plt.clf()

plt.figure(figsize=(6,5))
plt.plot(ratios_real, pco_sel, 'bo-', label='CO', color='limegreen')
plt.annotate("CO", fontsize=18, color='limegreen',xy=(1.8,98), va='top', ha='center')
plt.plot(ratios_real,horn_co,'--',color='limegreen')
plt.plot(ratios_real, ph2_sel, 'bo-', label='H2', color='k')
plt.annotate("H$_2$", fontsize=18, color='k',xy=(0.8,75), va='top', ha='center')
plt.plot(ratios_real,horn_h2,'--',color='k')
plt.ylabel('Exit Selectivity (%)', fontsize=22)
plt.xlabel('C/O Ratio', fontsize=22)
plt.xticks(np.arange(.6,2.7,0.2))
plt.xticks(fontsize=18)
plt.yticks(fontsize=18)
plt.yticks(np.arange(10,100,10))
plt.tight_layout()
plt.savefig('./paperplots/selectivity_syngas.pdf', bbox_inches='tight', dpi=300)
plt.clf()

plt.figure(figsize=(6,5))
plt.plot(ratios_real, pco2_sel, 'bo-', label='CO2', color='blue')
plt.annotate("CO$_2$", fontsize=18, color='blue',xy=(1.8,8), va='top', ha='center')
plt.plot(ratios_real,horn_co2,'--',color='blue')
plt.plot(ratios_real, ph2o_sel, 'bo-', label='H2O', color='dodgerblue')
plt.annotate("H$_2$O", fontsize=18, color='dodgerblue',xy=(1.,25), va='top', ha='center')
plt.plot(ratios_real,horn_h2o,'--',color='dodgerblue')
plt.ylabel('Exit Selectivity (%)', fontsize=22)
plt.xlabel('C/O Ratio', fontsize=22)
plt.xticks(np.arange(.6,2.7,0.2))
plt.xticks(fontsize=18)
plt.yticks(fontsize=18)
plt.yticks(np.arange(10,100,10))
plt.tight_layout()
plt.savefig('./paperplots/selectivity_full.pdf', bbox_inches='tight', dpi=300)
plt.clf()


"""
This file provides simple functions to calculate wavelength dependent effects.

The functions can also be used to estimate the Weak Lensing Channel ghosts as a function
of spectral type.

:requires: NumPy
:requires: matplotlib
:requires: pysynphot

:version: 0.1

:author: Sami-Matias Niemi
:contact: s.niemi@ucl.ac.uk
"""
import matplotlib
matplotlib.rc('text', usetex=True)
matplotlib.rcParams['font.size'] = 17
matplotlib.rc('xtick', labelsize=14)
matplotlib.rc('axes', linewidth=1.1)
matplotlib.rcParams['legend.fontsize'] = 11
matplotlib.rcParams['legend.handlelength'] = 3
matplotlib.rcParams['xtick.major.size'] = 5
matplotlib.rcParams['ytick.major.size'] = 5
import numpy as np
import matplotlib.pyplot as plt
import pysynphot as S

def _VISbandpass(scale=0.8, area=10000.):
    """
    Returns Weak Lensing Channel bandpass objects.

    Sets the primary mirror collecting area to 1m**2. This affects the count rate calcuations.

    :param scale: scale the throughput to EoL situation [default=0.8]
    :type scale: float
    :param area: collecting area of the primary mirror [default=10000.]
    :type area: float

    :return: BoL and EoL bandpass objects
    """
    data = np.loadtxt('/Users/sammy/EUCLID/throughputs/VIS.txt')

    bp = S.ArrayBandpass(wave=data[:, 0], throughput=data[:, 1], waveunits='angstrom', name='VIS')

    bpEoL = S.ArrayBandpass(wave=data[:, 0].copy(), throughput=data[:, 1].copy()*scale,
                            waveunits='angstrom', name='VIS EoL')

    #set the primary mirror collecting area to 1m**2, effects the count rate estimates
    bp.primary_area = area
    bpEoL.primary_area = area

    return bp, bpEoL

def _VISbandpassGhost(scale=0.8, area=10000.):
    """
    Returns Weak Lensing Channel ghost objects.

    Sets the primary mirror collecting area to 1m**2. This affects the count rate calcuations.

    :param scale: scale the throughput to EoL situation [default=0.8]
    :type scale: float
    :param area: collecting area of the primary mirror [default=10000.]
    :type area: float

    :return: BoL and EoL bandpass objects
    """
    data = np.loadtxt('/Users/sammy/EUCLID/throughputs/ghost.txt')

    bp = S.ArrayBandpass(wave=data[:, 0], throughput=data[:, 1], waveunits='angstrom', name='VIS')

    bpEoL = S.ArrayBandpass(wave=data[:, 0].copy(), throughput=data[:, 1].copy()*scale,
                            waveunits='angstrom', name='VIS EoL')

    #set the primary mirror collecting area to 1m**2, effects the count rate estimates
    bp.primary_area = area
    bpEoL.primary_area = area

    return bp, bpEoL

def throughputs(output='throughputs.pdf'):
    """
    Plot throughputs, compares to HST WFC3 UVIS F600LP

    :param output: name of the output file
    :type output: str

    :return: None
    """
    #comparison
    bp1 = S.ObsBandpass('wfc3,uvis2,f600lp')

    #VIS
    bp, bpEoL = _VISbandpass()

    #ghost
    bpG, bpEoLG = _VISbandpassGhost()

    #plot
    plt.semilogy(bp1.wave/10., bp1.throughput, 'r-', label='WFC3 F600LP')
    plt.semilogy(bp.wave/10., bp.throughput, 'b-', label='VIS Best Estimate')
    plt.semilogy(bpEoL.wave/10., bpEoL.throughput, 'g--', label='VIS EoL Req.')
    plt.semilogy(bpG.wave/10., bpG.throughput, 'm-', label='VIS Ghost')
    plt.semilogy(bpEoLG.wave/10., bpEoLG.throughput, 'y-.', label='VIS Ghost EoL')
    plt.xlim(230, 1100)
    plt.xlabel(r'Wavelength [nm]')
    plt.ylabel(r'Total System Throughput')
    plt.legend(shadow=True, fancybox=True, loc='best')
    plt.savefig(output)
    plt.close()

def testFlatSpectrum(mag=18):
    """
    Test the pysynphot flat spectra, how to make flat in lambda and nu.

    :param mag:
    :return:
    """
    unitflux1 = S.FlatSpectrum(mag, fluxunits='abmag')                 #F_lam
    unitflux2 = S.FlatSpectrum(mag, fluxunits='abmag', waveunits='Hz') #F_nu

    #unitflux2.convert(S.units.Angstrom)

    plt.plot(unitflux1.wave, unitflux1.flux, 'r-')
    plt.plot(unitflux2.waveunits.ToAngstrom(unitflux2.wave), unitflux2.flux, 'b--')
    #plt.plot(unitflux2.wave, unitflux2.flux, 'b--')

    plt.xlim(3000, 11000)
    #plt.ylim(17.5, 18.5)
    plt.savefig('flatSpectra.pdf')
    plt.close()

def flatSpectrum(mag=18):
    unitflux = S.FlatSpectrum(mag, fluxunits='abmag')

    #observing bandpass
    bp1 = S.ObsBandpass('wfc3,uvis2,f600lp')
    #VIS
    bp, bpEoL = _VISbandpass()

    #observations
    obs1 = S.Observation(unitflux, bp1)
    obs2 = S.Observation(unitflux, bp)
    obsEoL = S.Observation(unitflux, bpEoL)

    #converts
    obs1.convert('counts')
    obs2.convert('counts')
    obsEoL.convert('counts')

    print 'Count rates in e/s (WFC3 vs VIS):'
    print obs1.countrate(range=[3500, 11000]) #source countrate e/s (all counts, no aperture assumed)
    print obs2.countrate(range=[3500, 11000]) #source countrate e/s (all counts, no aperture assumed)
    print obsEoL.countrate(range=[3500, 11000]) #source countrate e/s (all counts, no aperture assumed)

    print 'Count rate VIS in-band / total:'
    print obs2.countrate(range=[5500, 9000]) / obs2.countrate()

    #print 'integrated ?'
    #print obs1.integrate() #
    #print obs2.integrate() #

    #bin widths
    bw1 = np.diff(obs1.binwave)
    bw2 = np.diff(obs2.binwave)

    #plot
    plt.title(r'Flat Spectrum $F_{\lambda}$ 18 mag$_{AB}$ [photons cm$^{-2}$ s$^{-1}$ \AA$^{-1}$]')
    plt.plot(obs1.binwave[1:], obs1.binflux[1:]/bw1, 'r-', label='WFC3 F600LP')
    plt.plot(obs2.binwave[1:], obs2.binflux[1:]/bw2, 'b-', label='VIS Best Estimate')
    plt.plot(obsEoL.binwave[1:], obsEoL.binflux[1:]/bw2, 'g--', label='VIS EoL Req.')
    plt.xlim(3000, 11000)
    plt.xlabel(r'Wavelength [\AA]')
    plt.ylabel(r'Counts per Wavelength Unit')
    plt.legend(shadow=True, fancybox=True)
    plt.savefig('comparisonFlatspectrum.pdf')
    plt.close()

def G2star():
    #object flux
    G2 = S.FileSpectrum('/Users/sammy/synphot/pickles/dat_uvk/pickles_uk_26.fits')

    #observing bandpass
    bp1 = S.ObsBandpass('wfc3,uvis2,f600lp')
    #VIS
    bp, bpEoL = _VISbandpass()

    #observations
    obs1 = S.Observation(G2, bp1)
    obs2 = S.Observation(G2, bp)
    obsEoL = S.Observation(G2, bpEoL)
    obs1.convert('counts')
    obs2.convert('counts')
    obsEoL.convert('counts')

    print 'effective wavelength [AA]'
    print obs2.efflam()  #effective wavelength

    print 'Countrate in e/s:'
    #source countrate e/s (all counts, no aperture assumed)
    print obs2.countrate(range=[3500, 11000])
    print obsEoL.countrate(range=[3500, 11000])
    print obs2.countrate(range=[5500, 9000])
    print obsEoL.countrate(range=[5500, 9000])

    print 'Count rate VIS in-band / total:'
    print obs2.countrate(range=[5500, 9000]) / obs2.countrate()

    #bin widths
    bw1 = np.diff(obs1.binwave)
    bw2 = np.diff(obs2.binwave)

    #plot
    plt.title('G2 Star (Pickles\_uk\_26)')
    plt.plot(obs1.binwave[1:], obs1.binflux[1:]/bw1, 'r-', label='WFC3 F600LP')
    plt.plot(obs2.binwave[1:], obs2.binflux[1:]/bw2, 'b-', label='VIS Best Estimate')
    plt.plot(obsEoL.binwave[1:], obsEoL.binflux[1:]/bw2, 'g--', label='VIS EoL Req.')
    plt.xlim(3000, 11000)
    plt.xlabel(r'Wavelength [\AA]')
    plt.ylabel(r'Counts per Wavelength Unit')
    plt.legend(shadow=True, fancybox=True)
    plt.savefig('comparisonG2.pdf')

def ghostCalculations(sourceSpectrum, title, output):
    """
    Weak Lensing Channel Ghost calculations.

    :param sourceSpectrum: pysynphot spectrum object
    :type sourceSpectrum: object
    :param title: title of the plot
    :type title: str
    :param output: name of the output file
    :type output: str

    :return: pysynphot observation object for VIS and Ghost
    :rtype: list of objects
    """
    bp, bpEoL = _VISbandpass()
    bpG, bpEoLG = _VISbandpassGhost()

    #BoL, the best estimate
    obs = S.Observation(sourceSpectrum, bp)
    obsG = S.Observation(sourceSpectrum, bpG)

    #convert to counts and derive the count rate in e/s
    obs.convert('counts')
    obsG.convert('counts')
    c = obs.countrate(range=[3200, 11500])
    cG = obsG.countrate(range=[3200, 11500])
    #c = obs.countrate(binned=False)
    #cG = obsG.countrate(binned=False)

    print 'effective stimulation in magnitude (AB)'
    print obs.effstim('abmag'), obsG.effstim('abmag')

    print 'effective wavelength'
    print obs.efflam(), obs.efflam()

    print 'source vs. ghost count rates [e/s]'
    print c, cG

    print 'ghost count rate / total count rate = %e' % (cG / c)
    print 'ghost count rate / total count rate [1750 dilution] = %e' % (cG / c / 1750.)

    #binflux is wavelength bin dependent...
    bw1 = np.diff(obs.binwave)
    bw2 = np.diff(obsG.binwave)

    #for a, b in zip(obsG.binwave[1:], bw2):
    #    print a, b

    y1 = obs.binflux[1:]/bw1
    y2 = obsG.binflux[1:]/bw2

    scale = np.max(y1)

    #make a plot
    plt.title(title)
    plt.semilogy(obs.binwave[1:], y1/scale, 'b-', label='Source')
    plt.semilogy(obsG.binwave[1:], y2/scale, 'r--', label='Ghost')
    plt.xlim(3000, 11000)
    plt.ylim(1e-7, 2.)
    plt.xlabel(r'Wavelength [\AA]')
    plt.ylabel(r'Normalised Counts per Wavelength Unit')
    plt.legend(shadow=True, fancybox=True, loc='best')
    plt.savefig(output)
    plt.close()

    return obs, obsG

def ghostResults():
    """
    Calculate the VIS channel ghost contribution.

    Synphot table for the Pickles stellar library:
    http://www.stsci.edu/hst/HST_overview/documents/synphot/AppA_Catalogs5.html

    :return:
    """
    #G2V
    print '\n\n\nG2V:'
    G2 = S.FileSpectrum('/Users/sammy/synphot/pickles/dat_uvk/pickles_uk_26.fits')
    obsBoL2, obsBoLG2 = ghostCalculations(G2, 'G2V Star (Pickles\_uk\_26)', 'G2Ghost.pdf')

    #Flat
    print '\n\n\nFlat F_lam:'
    unitflux = S.FlatSpectrum(18.)
    obsBoL, obsBoLG = ghostCalculations(unitflux,
                                        r'Flat Spectrum $F_{\lambda}$ [photons cm$^{-2}$ s$^{-1}$ \AA$^{-1}$]',
                                        'FlatGhost.pdf')

    #O5V
    print '\n\n\nO5V:'
    O5V = S.FileSpectrum('/Users/sammy/synphot/pickles/dat_uvk/pickles_uk_1.fits')
    obsBoL2, obsBoLG2 = ghostCalculations(O5V, 'O5V Star (Pickles\_uk\_1)', 'O5Ghost.pdf')

    #G2IV
    print '\n\n\nG2IV:'
    sp = S.FileSpectrum('/Users/sammy/synphot/pickles/dat_uvk/pickles_uk_54.fits')
    obsBoL2, obsBoLG2 = ghostCalculations(sp, 'G2IV Star (Pickles\_uk\_54)', 'G2IVGhost.pdf')

    #G5III
    print '\n\n\nG5III:'
    sp = S.FileSpectrum('/Users/sammy/synphot/pickles/dat_uvk/pickles_uk_73.fits')
    obsBoL2, obsBoLG2 = ghostCalculations(sp, 'G5III Star (Pickles\_uk\_73)', 'G5IIIGhost.pdf')

    #K2I
    print '\n\n\nK2I:'
    sp = S.FileSpectrum('/Users/sammy/synphot/pickles/dat_uvk/pickles_uk_128.fits')
    obsBoL2, obsBoLG2 = ghostCalculations(sp, 'K2I Star (Pickles\_uk\_128)', 'K2IGhost.pdf')

if __name__ == '__main__':
    #testFlatSpectrum()
    #throughputs()
    #flatSpectrum()
    #G2star()

    ghostResults()