from __future__ import absolute_import
# #START_LICENSE###########################################################
#
#
# This file is part of the Environment for Tree Exploration program
# (ETE).  http://etetoolkit.org
#
# ETE is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# ETE is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
# License for more details.
#
# You should have received a copy of the GNU General Public License
# along with ETE.  If not, see <http://www.gnu.org/licenses/>.
#
#
#                     ABOUT THE ETE PACKAGE
#                     =====================
#
# ETE is distributed under the GPL copyleft license (2008-2015).
#
# If you make use of ETE in published work, please cite:
#
# Jaime Huerta-Cepas, Joaquin Dopazo and Toni Gabaldon.
# ETE: a python Environment for Tree Exploration. Jaime BMC
# Bioinformatics 2010,:24doi:10.1186/1471-2105-11-24
#
# Note that extra references to the specific methods implemented in
# the toolkit may be available in the documentation.
#
# More info at http://etetoolkit.org. Contact: huerta@embl.de
#
#
# #END_LICENSE#############################################################

from math import log, exp

from six.moves import range
from numpy import floor, pi as PI, sin

from .. import Tree

def get_rooting(tol, seed_species, agename = False):
    '''
    returns dict of species age for a given TOL and a given seed

    **Example:**

    ::

      tol  = "((((((((Drosophila melanogaster,(Drosophila simulans,Drosophila secchellia)),(Drosophila yakuba,Drosophila erecta))[&&NHX:name=melanogaster subgroup],Drosophila ananassae)[&&NHX:name=melanogaster group],(Drosophila pseudoobscura,Drosophila persimilis)[&&NHX:name=obscura group])[&&NHX:name=Sophophora Old World],Drosophila willistoni)[&&NHX:name=subgenus Sophophora],(Drosophila grimshawi,(Drosophila virilis,Drosophila mojavensis))[&&NHX:name=subgenus Drosophila])[&&NHX:name=genus Drosophila],(Anopheles gambiae,Aedes aegypti)[&&NHX:name=Culicidae])[&&NHX:name=Arthropoda],Caenorhabditis elegans)[&&NHX:name=Animalia];"
      seed = "Drosophila melanogaster"
      ROOTING, age2name = get_rooting (tol, seed, True)

      ROOTING == {"Aedes aegypti"           : 7,
                  "Anopheles gambiae"       : 7,
                  "Caenorhabditis elegans"  : 8,
                  "Drosophila ananassae"    : 3,
                  "Drosophila erecta"       : 2,
                  "Drosophila grimshawi"    : 6,
                  "Drosophila melanogaster" : 1,
                  "Drosophila mojavensis"   : 6,
                  "Drosophila persimilis"   : 4,
                  "Drosophila pseudoobscura": 4,
                  "Drosophila secchellia"   : 1,
                  "Drosophila simulans"     : 1,
                  "Drosophila virilis"      : 6,
                  "Drosophila willistoni"   : 5,
                  "Drosophila yakuba"       : 2}

      age2name == {1: "Drosophila melanogaster. Drosophila simulans. Drosophila secchellia",
                   2: "melanogaster subgroup",
                   3: "melanogaster group",
                   4: "Sophophora Old World",
                   5: "subgenus Sophophora",
                   6: "genus Drosophila",
                   7: "Arthropoda",
                   8: "Animalia"}

    :argument seed_species: species name
    :argument False agename: if True, also returns the inverse dictionary

    :returns: ROOTING dictionary with age of each species

    '''

    tol = Tree (tol)
    try:
        node = tol.search_nodes (name=seed_species)[0]
    except IndexError:
        exit ('ERROR: Seed species not found in tree\n')
    age = 1
    ROOTING = {}
    if agename:
        age2name = {}
    while not node.is_root():
        node = node.up
        for leaf in node.get_leaf_names():
            if agename:
                if node.name == 'NoName':
                    nam = '.'.join (node.get_leaf_names())
                else:
                    nam = node.name
                age2name.setdefault (age, nam)
            ROOTING.setdefault (leaf, age)
        age += 1
    if agename:
        return ROOTING, age2name
    return ROOTING

def translate(sequence):
    '''
    little function to translate DNA to protein...
    from: http://python.genedrift.org/
    TODO : inseqgroup functions?

    :argument sequence: string

    :returns: translated sequence
    '''
    #dictionary with the genetic code
    gencode = {
        'ATA':'I', 'ATC':'I', 'ATT':'I', 'ATG':'M',
        'ACA':'T', 'ACC':'T', 'ACG':'T', 'ACT':'T',
        'AAC':'N', 'AAT':'N', 'AAA':'K', 'AAG':'K',
        'AGC':'S', 'AGT':'S', 'AGA':'R', 'AGG':'R',
        'CTA':'L', 'CTC':'L', 'CTG':'L', 'CTT':'L',
        'CCA':'P', 'CCC':'P', 'CCG':'P', 'CCT':'P',
        'CAC':'H', 'CAT':'H', 'CAA':'Q', 'CAG':'Q',
        'CGA':'R', 'CGC':'R', 'CGG':'R', 'CGT':'R',
        'GTA':'V', 'GTC':'V', 'GTG':'V', 'GTT':'V',
        'GCA':'A', 'GCC':'A', 'GCG':'A', 'GCT':'A',
        'GAC':'D', 'GAT':'D', 'GAA':'E', 'GAG':'E',
        'GGA':'G', 'GGC':'G', 'GGG':'G', 'GGT':'G',
        'TCA':'S', 'TCC':'S', 'TCG':'S', 'TCT':'S',
        'TTC':'F', 'TTT':'F', 'TTA':'L', 'TTG':'L',
        'TAC':'Y', 'TAT':'Y', 'TAA':'.', 'TAG':'.',
        'TGC':'C', 'TGT':'C', 'TGA':'.', 'TGG':'W',
        '---':'-', 'nnn':'x', 'NNN':'X'
    }
    ambig = {'Y':['A', 'G'], 'R':['C', 'T'], 'M':['G', 'T'], 'K':['A', 'C'], \
             'S':['G', 'C'],'W':['A', 'T'], 'V':['C', 'G', 'T'], \
             'H':['A', 'G', 'T'], 'D':['A', 'C', 'T'], 'B':['A', 'C', 'G'], \
             'N':['A', 'C', 'G', 'T']}
    proteinseq = ''
    #loop to read DNA sequence in codons, 3 nucleotides at a time
    sequence = sequence.upper()
    for n in range(0, len(sequence), 3):
        #checking to see if the dictionary has the key
        try:
            proteinseq += gencode[sequence[n:n+3]]
        except KeyError:
            newcod = []
            for nt in sequence[n:n+3]:
                if nt in ambig:
                    newcod.append(ambig[nt])
                else :
                    newcod.append(list (nt))
            aa = ''
            for nt1 in newcod[0]:
                for nt2 in newcod[1]:
                    for nt3 in newcod[2]:
                        try:
                            if aa == '':
                                aa  = gencode[nt1+nt2+nt3]
                            elif gencode[nt1+nt2+nt3] != aa:
                                aa = 'X'
                                break
                        except KeyError:
                            aa = 'X'
                            break
            proteinseq += aa
    return proteinseq

# reused from pycogent
ROUND_ERROR = 1e-14
MAXLOG      = 7.09782712893383996843E2
MAXLGM      = 2.556348e305
big         = 4.503599627370496e15
biginv      = 2.22044604925031308085e-16
MACHEP      = 1.11022302462515654042E-16
LS2PI       =  0.91893853320467274178
LOGPI       = 1.14472988584940017414

def chi_high(x, df):
    """Returns right-hand tail of chi-square distribution (x to infinity).

    df, the degrees of freedom, ranges from 1 to infinity (assume integers).
    Typically, df is (r-1)*(c-1) for a r by c table.

    Result ranges from 0 to 1.

    See Cephes docs for details.
    """
    x = fix_rounding_error(x)

    if x < 0:
        raise ValueError("chi_high: x must be >= 0 (got %s)." % x)
    if df < 1:
        raise ValueError("chi_high: df must be >= 1 (got %s)." % df)
    return igamc(float(df)/2, x/2)

def fix_rounding_error(x):
    """If x is almost in the range 0-1, fixes it.

    Specifically, if x is between -ROUND_ERROR and 0, returns 0.
    If x is between 1 and 1+ROUND_ERROR, returns 1.
    """
    if -ROUND_ERROR < x < 0:
        return 0
    elif 1 < x < 1+ROUND_ERROR:
        return 1
    return x

def igamc(a,x):
    """Complemented incomplete Gamma integral: see Cephes docs."""
    if x <= 0 or a <= 0:
        return 1
    if x < 1 or x < a:
        return 1 - igam(a, x)
    ax = a * log(x) - x - lgam(a)
    if ax < -MAXLOG:    #underflow
        return 0
    ax = exp(ax)
    #continued fraction
    y = 1 - a
    z = x + y + 1
    c = 0
    pkm2 = 1
    qkm2 = x
    pkm1 = x + 1
    qkm1 = z * x
    ans = pkm1/qkm1

    while 1:
        c += 1
        y += 1
        z += 2
        yc = y * c
        pk = pkm1 * z - pkm2 * yc
        qk = qkm1 * z - qkm2 * yc
        if qk != 0:
            r = pk/qk
            t = abs((ans-r)/r)
            ans = r
        else:
            t = 1
        pkm2 = pkm1
        pkm1 = pk
        qkm2 = qkm1
        qkm1 = qk
        if abs(pk) > big:
            pkm2 *= biginv
            pkm1 *= biginv
            qkm2 *= biginv
            qkm1 *= biginv
        if t <= MACHEP:
            break
    return ans * ax

def lgam(x):
    """Natural log of the gamma fuction: see Cephes docs for details"""
    if x < -34:
        q = -x
        w = lgam(q)
        p = floor(q)
        if p == q:
            raise OverflowError("lgam returned infinity.")
        z = q - p
        if z > 0.5:
            p += 1
            z = p - q
        z = q * sin(PI * z)
        if z == 0:
            raise OverflowError("lgam returned infinity.")
        z = LOGPI - log(z) - w
        return z
    if x < 13:
        z = 1
        p = 0
        u = x
        while u >= 3:
            p -= 1
            u = x + p
            z *= u
        while u < 2:
            if u == 0:
                raise OverflowError("lgam returned infinity.")
            z /= u
            p += 1
            u = x + p
        if z < 0:
            z = -z
        if u == 2:
            return log(z)
        p -= 2
        x = x + p
        p = x * polevl(x, GB)/polevl(x,GC)
        return log(z) + p
    if x > MAXLGM:
        raise OverflowError("Too large a value of x in lgam.")
    q = (x - 0.5) * log(x) - x + LS2PI
    if x > 1.0e8:
        return q
    p = 1/(x*x)
    if x >= 1000:
        q += ((  7.9365079365079365079365e-4 * p
                 -2.7777777777777777777778e-3) *p
              + 0.0833333333333333333333) / x
    else:
        q += polevl(p, GA)/x
    return q

def polevl(x, coef):
    """evaluates a polynomial y = C_0 + C_1x + C_2x^2 + ... + C_Nx^N

    Coefficients are stored in reverse order, i.e. coef[0] = C_N
    """
    result = 0
    for c in coef:
        result = result * x + c
    return result

def igam(a, x):
    """Left tail of incomplete gamma function: see Cephes docs for details"""
    if x <= 0 or a <= 0:
        return 0
    if x > 1 and x > a:
        return 1 - igamc(a,x)

    #Compute x**a * exp(x) / Gamma(a)

    ax = a * log(x) - x - lgam(a)
    if ax < -MAXLOG:    #underflow
        return 0.0
    ax = exp(ax)

    #power series
    r = a
    c = 1
    ans = 1
    while 1:
        r += 1
        c *= x/r
        ans += c
        if c/ans <= MACHEP:
            break

    return ans * ax / a

#Coefficients for Gamma follow:
GA = [
    8.11614167470508450300E-4,
    -5.95061904284301438324E-4,
    7.93650340457716943945E-4,
    -2.77777777730099687205E-3,
    8.33333333333331927722E-2,
]

GB = [
    -1.37825152569120859100E3,
    -3.88016315134637840924E4,
    -3.31612992738871184744E5,
    -1.16237097492762307383E6,
    -1.72173700820839662146E6,
    -8.53555664245765465627E5,
]

GC = [
    1.00000000000000000000E0,
    -3.51815701436523470549E2,
    -1.70642106651881159223E4,
    -2.20528590553854454839E5,
    -1.13933444367982507207E6,
    -2.53252307177582951285E6,
    -2.01889141433532773231E6,
]

GP = [
    1.60119522476751861407E-4,
    1.19135147006586384913E-3,
    1.04213797561761569935E-2,
    4.76367800457137231464E-2,
    2.07448227648435975150E-1,
    4.94214826801497100753E-1,
    9.99999999999999996796E-1,
]

GQ = [
    -2.31581873324120129819E-5,
    5.39605580493303397842E-4,
    -4.45641913851797240494E-3,
    1.18139785222060435552E-2,
    3.58236398605498653373E-2,
    -2.34591795718243348568E-1,
    7.14304917030273074085E-2,
    1.00000000000000000320E0,
]


# -*- coding: utf-8 -*-
# @Brief: 模型loss函数

from tensorflow.keras import backend, losses
import tensorflow as tf
import numpy as np

def rpn_cls_loss(ratio=3):
    def cls_loss(y_true, y_pred):
        """
        计算rpn 是否有物体的loss，可以直接用交叉熵计算损失，但这里计算
        :param y_true: 真实值 [batch_size, num_anchor, 1]
        :param y_pred: 预测值 [batch_size, num_anchor, 1]
        :return: rpn cls_loss
        """

        label_true = y_true[:, :, -1]       # 取出真实值的label rpn的label 1是物体，0是背景，-1是需要忽略的
        label_pred = y_pred

        # 找出存在目标的先验框 有目标是1
        indices_for_object = tf.where(backend.equal(label_true, 1))     # 如果x,y都为None，就返回condition的坐标
        labels_for_object = tf.gather_nd(y_true, indices_for_object)    # 根据indices选取真实值标签
        classification_for_object = tf.gather_nd(label_pred, indices_for_object)    # 选取预测值标签

        cls_loss_for_object = backend.binary_crossentropy(labels_for_object, classification_for_object)

        # 找出实际上为背景的先验框 没有目标是0
        indices_for_back = tf.where(backend.equal(label_true, 0))
        labels_for_back = tf.gather_nd(y_true, indices_for_back)
        classification_for_back = tf.gather_nd(label_pred, indices_for_back)

        # 计算每一个先验框应该有的权重
        cls_loss_for_back = backend.binary_crossentropy(labels_for_back, classification_for_back)

        # 标准化，计算是正样本的数量
        normalizer_pos = tf.where(backend.equal(label_true, 1))
        normalizer_pos = backend.cast(backend.shape(normalizer_pos)[0], 'float32')
        normalizer_pos = backend.maximum(backend.cast_to_floatx(1.0), normalizer_pos)

        # 计算负样本的数量
        normalizer_neg = tf.where(backend.equal(label_true, 0))
        normalizer_neg = backend.cast(backend.shape(normalizer_neg)[0], 'float32')
        normalizer_neg = backend.maximum(backend.cast_to_floatx(1.0), normalizer_neg)

        # 将所获得的loss除上样本的数量
        cls_loss_for_object = backend.sum(cls_loss_for_object) / normalizer_pos         # 物体的loss
        cls_loss_for_back = ratio * backend.sum(cls_loss_for_back) / normalizer_neg     # 背景的loss

        # 总的loss
        loss = cls_loss_for_object + cls_loss_for_back

        return loss

    return cls_loss

def rpn_regr_loss(sigma=1.0):
    sigma_squared = sigma ** 2

    def smooth_l1(y_true, y_pred):
        """
        计算rpn 建议框坐标的loss
        使用smooth l1 loss
        f(x) =  0.5 * (sigma * x)^2          if |x| < 1 / sigma / sigma
                |x| - 0.5 / sigma^2          otherwise
        :param sigma: 是平滑参数，控制平滑区域
        :param y_true: 真实值 [batch_size, num_anchor, 4+1]
        :param y_pred: 预测值 [batch_size, num_anchor, 4]
        :return: rpn regr_loss
        """
        regression_pred = y_pred
        regression_true = y_true[:, :, :-1]   # 取rpn上的坐标
        label_true = y_true[:, :, -1]         # 取框内是否有物体的预测值

        # 找到只有物体的框，不要背景
        indices = tf.where(backend.equal(label_true, 1))                    # 如果x,y都为None，就返回condition的坐标
        regression_pred = tf.gather_nd(regression_pred, indices)            # 根据有物体的索引，取出预测框的相关坐标
        regression_true = tf.gather_nd(regression_true, indices)            # 取出真实框的坐标

        # 计算 smooth L1 loss
        regression_diff = backend.abs(regression_pred - regression_true)
        regression_loss = tf.where(  # tf.where用做判断条件
            backend.less(regression_diff, 1.0 / sigma_squared),             # 绝对值是否小于1
            0.5 * sigma_squared * backend.pow(regression_diff, 2),          # 如果是
            regression_diff - 0.5 / sigma_squared                           # 如果不是
        )

        # 除于N_cls
        normalizer = backend.maximum(1, backend.shape(indices)[0])
        normalizer = backend.cast(normalizer, dtype='float32')
        loss = backend.sum(regression_loss) / normalizer

        return loss

    return smooth_l1

def class_loss_regr(num_classes):
    epsilon = 1e-4

    def class_loss_regr_fixed_num(y_true, y_pred):
        """
        计算classifier的回归损失
        :param y_true: 真实值 [batch_size, num_rois, num_classes * 8]
        :param y_pred: 预测值 [batch_size, num_rois, num_classes * 4]
        :return: classifier regr_loss
        """
        regr_loss = 0
        batch_size = len(y_true)
        for i in range(batch_size):
            x = y_true[i, :, 4 * num_classes:] - y_pred[i, :, :]                    # 取出y_true后一半的数据，与y_pred做差值
            x_abs = backend.abs(x)                                                  # 计算绝对值
            x_bool = backend.cast(backend.less_equal(x_abs, 1.0), 'float32')        # 小于1的值

            # 1、差值绝对值小于1时0.5 * X^2，大于1的绝对值减0.5然后相加
            # 2、在乘上是否要计算这个loss
            # 3、求和在除以个数，得均值
            loss = 4 * backend.sum(
                y_true[i, :, :4 * num_classes] * (x_bool * (0.5 * x * x) + (1 - x_bool) * (x_abs - 0.5))) / backend.sum(
                epsilon + y_true[i, :, :4 * num_classes])
            regr_loss += loss

        return regr_loss / backend.constant(batch_size)

    return class_loss_regr_fixed_num

def class_loss_cls(y_true, y_pred):
    """
    计算具体的分类loss
    :param y_true: 真实值 [batch_size, num_rois, 4+1]
    :param y_pred: 预测值 [batch_size, num_rois, 4+1]
    :return: classifier class_loss
    """
    return backend.mean(losses.categorical_crossentropy(y_true, y_pred))


"""Module for representation and analysis of MS measurements"""

import re
import numpy as np
import json  # FIXME: This is for MSCalibration.export, but shouldn't have to be here.
import warnings

from ..measurements import Measurement, Calibration
from ..spectra import Spectrum
from ..plotters.ms_plotter import MSPlotter, STANDARD_COLORS
from ..exceptions import QuantificationError
from ..constants import (
    AVOGADROS_CONSTANT,
    BOLTZMAN_CONSTANT,
    STANDARD_TEMPERATURE,
    STANDARD_PRESSURE,
    DYNAMIC_VISCOSITIES,
    MOLECULAR_DIAMETERS,
    MOLAR_MASSES,
)
from ..data_series import ValueSeries
from ..db import Saveable
from ..tools import deprecate

class MSMeasurement(Measurement):
    """Class implementing raw MS functionality"""

    extra_column_attrs = {"ms_measurement": ("tspan_bg",)}
    default_plotter = MSPlotter

    def __init__(self, name, **kwargs):
        tspan_bg = kwargs.pop("tspan_bg", None)
        super().__init__(name, **kwargs)
        self.tspan_bg = tspan_bg

    @property
    def ms_calibration(self):
        ms_cal_list = []
        tspan_bg = None
        signal_bgs = {}
        for cal in self.calibration_list:
            ms_cal_list = ms_cal_list + getattr(cal, "ms_cal_list", [])
            for mass, bg in getattr(cal, "signal_bgs", {}).items():
                if mass not in signal_bgs:
                    signal_bgs[mass] = bg
            tspan_bg = tspan_bg or getattr(cal, "tspan_bg", None)
        return MSCalibration(ms_cal_results=ms_cal_list, signal_bgs=signal_bgs)

    @property
    def signal_bgs(self):
        return self.ms_calibration.signal_bgs

    def set_bg(self, tspan_bg=None, mass_list=None):
        """Set background values for mass_list to the average signal during tspan_bg."""
        mass_list = mass_list or self.mass_list
        tspan_bg = tspan_bg or self.tspan_bg
        signal_bgs = {}
        for mass in mass_list:
            t, v = self.grab(mass, tspan_bg)
            signal_bgs[mass] = np.mean(v)
        self.add_calibration(MSCalibration(signal_bgs=signal_bgs))

    def reset_bg(self, mass_list=None):
        """Reset background values for the masses in mass_list"""
        mass_list = mass_list or self.mass_list
        for mass in mass_list:
            if mass in self.signal_bgs:
                del self.signal_bgs[mass]

    def grab(
        self,
        item,
        tspan=None,
        tspan_bg=None,
        include_endpoints=False,
        remove_background=False,
    ):
        """Returns t, S where S is raw signal in [A] for a given signal name (ie mass)

        Args:
            item (str): Name of the signal.
            tspan (list): Timespan for which the signal is returned.
            tspan_bg (list): Timespan that corresponds to the background signal.
                If not given, no background is subtracted.
            remove_background (bool): Whether to remove a pre-set background if
                available. This is special to MSMeasurement.
                Defaults to False, but in grab_flux it defaults to True.
            include_endpoints (bool): Whether to ensure tspan[0] and tspan[-1] are in t
        """
        time, value = super().grab(
            item, tspan=tspan, include_endpoints=include_endpoints
        )

        if tspan_bg:
            _, bg = self.grab(item, tspan=tspan_bg)
            return time, value - np.average(bg)
        elif remove_background:
            if item in self.signal_bgs:
                return time, value - self.signal_bgs[item]
            elif self.tspan_bg:
                _, bg = self.grab(item, tspan=self.tspan_bg)
                return time, value - np.average(bg)
        return time, value

    def grab_for_t(self, item, t, tspan_bg=None, remove_background=False):
        """Return a numpy array with the value of item interpolated to time t

        Args:
            item (str): The name of the value to grab
            t (np array): The time vector to grab the value for
            tspan_bg (iterable): Optional. A timespan defining when `item` is at its
                baseline level. The average value of `item` in this interval will be
                subtracted from what is returned.
            remove_background (bool): Whether to remove a pre-set background if
                available. This is special to MSMeasurement.
                Defaults to False, but in grab_flux it defaults to True.
        """
        t_0, v_0 = self.grab(
            item, tspan_bg=tspan_bg, remove_background=remove_background
        )
        v = np.interp(t, t_0, v_0)
        return v

    def grab_signal(self, *args, **kwargs):
        """Alias for grab()"""
        return self.grab(*args, **kwargs)

    @deprecate(
        "0.1", "Use `remove_background` instead.", "0.3", kwarg_name="removebackground"
    )
    def grab_flux(
        self,
        mol,
        tspan=None,
        tspan_bg=None,
        remove_background=True,
        removebackground=None,
        include_endpoints=False,
    ):
        """Return the flux of mol (calibrated signal) in [mol/s]

        Note:
        `grab_flux(mol, ...)` is identical to `grab(f"n_dot_{mol}", ...)` with
        remove_background=True by default. An MSCalibration does the maths.

        Args:
            mol (str or MSCalResult): Name of the molecule or a ms_calibration thereof
            tspan (list): Timespan for which the signal is returned.
            tspan_bg (list): Timespan that corresponds to the background signal.
                If not given, no background is subtracted.
            remove_background (bool): Whether to remove a pre-set background if available
                Defaults to True.
            removebackground (bool): DEPRECATED. Use `remove_background`.
            include_endpoints (bool): Whether to interpolate for tspan[0] and tspan[-1]
        """
        if removebackground is not None:
            remove_background = removebackground
        if isinstance(mol, MSCalResult):
            t, signal = self.grab(
                mol.mass,
                tspan=tspan,
                tspan_bg=tspan_bg,
                remove_background=remove_background,
                include_endpoints=include_endpoints,
            )
            return t, signal / mol.F
        return self.grab(
            # grab() invokes __getitem__, which invokes the `Calibration`. Specifically,
            # `MSCalibration.calibrate_series()` interprets item names starting with
            # "n_" as molecule fluxes, and checks itself for a sensitivity factor.
            f"n_dot_{mol}",
            tspan=tspan,
            tspan_bg=tspan_bg,
            remove_background=remove_background,
            include_endpoints=include_endpoints,
        )

    @deprecate(
        "0.1", "Use `remove_background` instead.", "0.3", kwarg_name="removebackground"
    )
    def grab_flux_for_t(
        self,
        mol,
        t,
        tspan_bg=None,
        remove_background=False,
        removebackground=None,
    ):
        """Return the flux of mol (calibrated signal) in [mol/s] for a given time vec

        Args:
            mol (str): Name of the molecule.
            t (np.array): The time vector along which to give the flux
            tspan_bg (tspan): Timespan that corresponds to the background signal.
                If not given, no background is subtracted.
            remove_background (bool): Whether to remove a pre-set background if available
            removebackground (bool): DEPRECATED. Use `remove_background`.
        """
        if removebackground is not None:
            remove_background = removebackground
        t_0, y_0 = self.grab_flux(
            mol,
            tspan_bg=tspan_bg,
            remove_background=remove_background,
        )
        y = np.interp(t, t_0, y_0)
        return y

    def get_flux_series(self, mol):
        """Return a ValueSeries with the calibrated flux of mol"""
        return self[f"n_dot_{mol}"]

    def integrate_signal(self, mass, tspan, tspan_bg, ax=None):
        """Integrate a ms signal with background subtraction and evt. plotting

        TODO: Should this, like grab_signal does now, have the option of using a
            background saved in the object rather than calculating a new one?

        Args:
            mass (str): The mass for which to integrate the signal
            tspan (tspan): The timespan over which to integrate
            tspan_bg (tspan): Timespan at which the signal is at its background value
            ax (Axis): axis to plot on. Defaults to None
        """
        t, S = self.grab_signal(mass, tspan=tspan, include_endpoints=True)
        if tspan_bg:
            t_bg, S_bg_0 = self.grab_signal(mass, tspan=tspan_bg, include_endpoints=True)
            S_bg = np.mean(S_bg_0) * np.ones(t.shape)
        else:
            S_bg = np.zeros(t.shape)
        if ax:
            if ax == "new":
                fig, ax = self.plotter.new_ax()
            ax.fill_between(t, S_bg, S, color=STANDARD_COLORS[mass], alpha=0.2)
        return np.trapz(S - S_bg, t)

    @property
    def mass_list(self):
        """List of the masses for which ValueSeries are contained in the measurement"""
        return [self.as_mass(col) for col in self.series_names if self.is_mass(col)]

    def is_mass(self, item):
        if re.search("^M[0-9]+$", item):
            return True
        if item in self.reverse_aliases and self.is_mass(self.reverse_aliases[item][0]):
            return True
        return False

    def as_mass(self, item):
        if re.search("^M[0-9]+$", item):
            return item
        new_item = self.reverse_aliases[item][0]
        if self.is_mass(new_item):
            return self.as_mass(new_item)
        raise TypeError(f"{self} does not recognize '{item}' as a mass.")

class MSCalResult(Saveable):
    """A class for a mass spec ms_calibration result.

    FIXME: I think that something inheriting directly from Saveable does not belong in
        a technique module.
    """

    table_name = "ms_cal_results"
    column_attrs = {"name", "mol", "mass", "cal_type", "F"}

    def __init__(
        self,
        name=None,
        mol=None,
        mass=None,
        cal_type=None,
        F=None,
    ):
        super().__init__()
        self.name = name or f"{mol}@{mass}"
        self.mol = mol
        self.mass = mass
        self.cal_type = cal_type
        self.F = F

    def __repr__(self):
        return (
            f"{self.__class__.__name__}(name={self.name}, mol={self.mol}, "
            f"mass={self.mass}, F={self.F})"
        )

    @property
    def color(self):
        return STANDARD_COLORS[self.mass]

class MSCalibration(Calibration):
    """Class for mass spec calibrations. TODO: replace with powerful external package"""

    extra_linkers = {"ms_calibration_results": ("ms_cal_results", "ms_cal_result_ids")}
    # FIXME: signal_bgs are not saved at present. Should they be a separate table
    #   of Saveable objects like ms_cal_results or should they be a single json value?
    child_attrs = [
        "ms_cal_results",
    ]

    def __init__(
        self,
        name=None,
        date=None,
        tstamp=None,  # FIXME: No need to have both a date and a tstamp?
        setup=None,
        ms_cal_results=None,
        signal_bgs=None,
        technique="MS",
        measurement=None,
    ):
        """
        Args:
            name (str): Name of the ms_calibration
            date (str): Date of the ms_calibration
            setup (str): Name of the setup where the ms_calibration is made
            ms_cal_results (list of MSCalResult): The mass spec calibrations
            measurement (MSMeasurement): The measurement
        """
        super().__init__(
            name=name or f"EC-MS ms_calibration for {setup} on {date}",
            technique=technique,
            tstamp=tstamp,
            measurement=measurement,
        )
        self.date = date
        self.setup = setup
        self.ms_cal_results = ms_cal_results or []
        self.signal_bgs = signal_bgs or {}

    @property
    def ms_cal_result_ids(self):
        return [cal.id for cal in self.ms_cal_results]

    @property
    def mol_list(self):
        return list({cal.mol for cal in self.ms_cal_results})

    @property
    def mass_list(self):
        return list({cal.mass for cal in self.ms_cal_results})

    @property
    def name_list(self):
        return list({cal.name for cal in self.ms_cal_results})

    def __contains__(self, mol):
        return mol in self.mol_list or mol in self.name_list

    def __iter__(self):
        yield from self.ms_cal_results

    def calibrate_series(self, key, measurement=None):
        """Return a calibrated series for `key` if possible.

        If key starts with "n_", it is interpreted as a molecule flux. This method then
        searches the calibration for a sensitivity factor for that molecule uses it to
        divide the relevant mass signal from the measurement. Example acceptable keys:
        "n_H2", "n_dot_H2".
        If the key does not start with "n_", or the calibration can't find a relevant
        sensitivity factor and mass signal, this method returns None.
        """
        measurement = measurement or self.measurement
        if key.startswith("n_"):  # it's a flux!
            mol = key.split("_")[-1]
            try:
                mass, F = self.get_mass_and_F(mol)
            except QuantificationError:
                # Calibrations just return None when they can't get what's requested.
                return
            signal_series = measurement[mass]
            y = signal_series.data
            if mass in measurement.signal_bgs:
                # FIXME: How to make this optional to user of MSMeasuremt.grab()?
                y = y - measurement.signal_bgs[mass]
            n_dot = y / F
            return ValueSeries(
                name=f"n_dot_{mol}",
                unit_name="mol/s",
                data=n_dot,
                tseries=signal_series.tseries,
            )

    def get_mass_and_F(self, mol):
        """Return the mass and sensitivity factor to use for simple quant. of mol"""
        cal_list_for_mol = [cal for cal in self if cal.mol == mol or cal.name == mol]
        Fs = [cal.F for cal in cal_list_for_mol]
        if not Fs:
            raise QuantificationError(f"{self} has no sensitivity factor for {mol}")
        index = np.argmax(np.array(Fs))

        the_good_cal = cal_list_for_mol[index]
        return the_good_cal.mass, the_good_cal.F

    def get_F(self, mol, mass):
        """Return the sensitivity factor for mol at mass"""
        cal_list_for_mol_at_mass = [
            cal
            for cal in self
            if (cal.mol == mol or cal.name == mol) and cal.mass == mass
        ]
        F_list = [cal.F for cal in cal_list_for_mol_at_mass]
        if not F_list:
            raise QuantificationError(
                f"{self} has no sensitivity factor for {mol} at {mass}"
            )
        return np.mean(np.array(F_list))

    def scaled_to(self, ms_cal_result):
        """Return a new ms_calibration w scaled sensitivity factors to match one given"""
        F_0 = self.get_F(ms_cal_result.mol, ms_cal_result.mass)
        scale_factor = ms_cal_result.F / F_0
        calibration_as_dict = self.as_dict()
        new_cal_list = []
        for cal in self.ms_cal_results:
            cal = MSCalResult(
                name=cal.name,
                mass=cal.mass,
                mol=cal.mol,
                F=cal.F * scale_factor,
                cal_type=cal.cal_type + " scaled",
            )
            new_cal_list.append(cal)
        calibration_as_dict["ms_cal_results"] = new_cal_list
        del calibration_as_dict["ms_cal_result_ids"]
        # ^ FIXME: ms_cal_result_ids via MemoryBackend
        calibration_as_dict["name"] = calibration_as_dict["name"] + " scaled"
        return self.__class__.from_dict(calibration_as_dict)

    @classmethod
    def read(cls, path_to_file):
        """Read an MSCalibration from a json-formatted text file"""
        with open(path_to_file) as f:
            obj_as_dict = json.load(f)
        # put the MSCalResults (exported as dicts) into objects:
        obj_as_dict["ms_cal_results"] = [
            MSCalResult.from_dict(ms_cal_as_dict)
            for ms_cal_as_dict in obj_as_dict["ms_cal_results"]
        ]
        return cls.from_dict(obj_as_dict)

    def export(self, path_to_file=None):
        """Export an ECMSCalibration as a json-formatted text file"""
        path_to_file = path_to_file or (self.name + ".ix")
        self_as_dict = self.as_dict()
        # replace the ms_cal_result ids with the dictionaries of the results themselves:
        del self_as_dict["ms_cal_result_ids"]
        self_as_dict["ms_cal_results"] = [cal.as_dict() for cal in self.ms_cal_results]
        with open(path_to_file, "w") as f:
            json.dump(self_as_dict, f, indent=4)

class MSInlet:
    """A class for describing the inlet to the mass spec

    Every MSInlet describes the rate and composition of the gas entering a mass
    spectrometer. The default is a Spectro Inlets EC-MS chip.
    TODO: Replace with powerful external package.
    """

    def __init__(
        self,
        *,
        l_cap=1e-3,
        w_cap=6e-6,
        h_cap=6e-6,
        gas="He",
        T=STANDARD_TEMPERATURE,
        p=STANDARD_PRESSURE,
        verbose=True,
    ):
        """Create an MSInlet object given its properties.

        Args:
            l_cap (float): capillary length [m]. Defaults to design parameter.
            w_cap (float): capillary width [m]. Defaults to design parameter.
            h_cap (float): capillary height [m]. Defaults to design parameter.
            p (float): system pressure in [Pa] (if to change from that in medium)
            T (float): system temperature in [K] (if to change from that in medium)
            gas (str): the gas at the start of the inlet.
            verbose (bool): whether to print stuff to the terminal
        """
        self.verbose = verbose
        self.l_cap = l_cap
        self.l_cap_eff = {}
        self.w_cap = w_cap
        self.h_cap = h_cap
        self.p = p
        self.T = T
        self.gas = gas  # TODO: Gas mixture class. This must be a pure gas now.

    def calc_l_cap_eff(
        self, n_dot_measured, gas=None, w_cap=None, h_cap=None, T=None, p=None
    ):
        """Calculate gas specific effective length of the capillary in [m]
        and add {gas:value} to l_cap_eff (dict)

        Args:
            w_cap (float): Capillary width [m], defaults to self.w_cap
            h_cap (float): Capillary height [m], defaults to self.h_cap
            n_dot_measured (float): Measured flux of gas [mol/s]
            gas (dict or str): The gas in the chip, defaults to self.gas
            T (float): Temperature [K], if to be updated
            p (float): Pressure [Pa], if to be updated
        Returns:
            float: Gas specific effective length in [m]
        """

        n_dot_predicted = self.calc_n_dot_0(gas=gas, w_cap=w_cap, h_cap=h_cap, T=T, p=p)

        l_cap_gas_specific_eff = self.l_cap * n_dot_predicted / n_dot_measured
        self.l_cap_eff[
            gas
        ] = l_cap_gas_specific_eff  # add effective l_cap for specific gas

        return l_cap_gas_specific_eff

    def update_l_cap(self, gases=[]):
        """Update self.l_cap from average of values in dict l_cap_eff

        Args:
            gases (list): List of gases to average l_cap, default all
        Returns:
            float: Averaged effective capilllary length in [m]
        """
        if self.l_cap_eff and not gases:
            self.l_cap = np.mean(list(self.l_cap_eff.values()))
        elif self.l_cap_eff and gases:
            _l_cap = 0
            for gas in gases:
                _l_cap += self.l_cap_eff[gas]
            self.l_cap = _l_cap / len(gases)

        return self.l_cap

    def calc_n_dot_0(self, gas=None, w_cap=None, h_cap=None, l_cap=None, T=None, p=None):
        """Calculate the total molecular flux through the capillary in [s^-1]

        Uses Equation 4.10 of Trimarco, 2017. "Real-time detection of sub-monolayer
        desorption phenomena during electrochemical reactions: Instrument development
        and applications." PhD Thesis, Technical University of Denmark.

        Args:
            w_cap (float): Capillary width [m], defaults to self.w_cap
            h_cap (float): Capillary height [m], defaults to self.h_cap
            l_cap (float): Capillary length [m], defaults to self.l_cap
            gas (dict or str): The gas in the chip, defaults to self.gas
            T (float): Temperature [K], if to be updated
            p (float): Pressure [Pa], if to be updated
        Returns:
            float: The total molecular flux in [s^-1] through the capillary
        """

        if w_cap is None:
            w_cap = self.w_cap  # capillary width in [m]
        if h_cap is None:
            h_cap = self.h_cap  # capillary height in [m]
        if l_cap is None:
            l_cap = self.l_cap  # effective capillary length in [m]
        if T is None:
            T = self.T
        if p is None:
            p = self.p
        pi = np.pi

        # TODO: make it so that DYNAMIC_VISCOSITIES[gas] can just be a float if someone
        #   enters it without having access to the temperature-dependent values.
        if T < DYNAMIC_VISCOSITIES[gas][0, 0] or T > DYNAMIC_VISCOSITIES[gas][-1, 0]:
            warnings.warn(
                "Insufficient data in constants.py to appropriately estimate "
                f"the dynamic viscosity for {gas} at temperature: {T}K",
                stacklevel=2,
            )

        _eta_v = DYNAMIC_VISCOSITIES[gas][:, 1]  # list of known eta(T) for 'gas'
        _eta_T = DYNAMIC_VISCOSITIES[gas][:, 0]  # list of paired Ts for eta(T)

        eta = np.interp(T, _eta_T, _eta_v)  # dynamic viscosity of gas at T in [Pa*s]

        s = MOLECULAR_DIAMETERS[gas]  # molecule diameter in [m]
        m = MOLAR_MASSES[gas] * 1e-3 / AVOGADROS_CONSTANT  # molecule mass in [kg]

        d = ((w_cap * h_cap) / pi) ** 0.5 * 2
        # d = 4.4e-6  #used in Henriksen2009
        a = d / 2
        p_1 = p
        lambda_ = d  # defining the transitional pressure
        # ...from setting mean free path equal to capillary d
        p_t = BOLTZMAN_CONSTANT * T / (2**0.5 * pi * s**2 * lambda_)
        p_2 = 0
        p_m = (p_1 + p_t) / 2  # average pressure in the transitional flow region
        v_m = (8 * BOLTZMAN_CONSTANT * T / (pi * m)) ** 0.5
        # a reciprocal velocity used for short-hand:
        nu = (m / (BOLTZMAN_CONSTANT * T)) ** 0.5

        # ... and now, we're ready for the capillary equation.
        #   (need to turn of black and flake8 for tolerable format)
        # fmt: off
        #   Equation 4.10 of Daniel Trimarco's PhD Thesis:
        N_dot = (                                                               # noqa
            1 / (BOLTZMAN_CONSTANT * T) * 1 / l_cap * (                         # noqa
                (p_t - p_2) * a**3 * 2 * pi / 3 * v_m + (p_1 - p_t) * (         # noqa
                    a**4 * pi / (8 * eta) * p_m  + a**3 * 2 * pi / 3 * v_m * (  # noqa
                        (1 + 2 * a * nu * p_m / eta) / (                        # noqa
                        1 + 2.48 * a * nu * p_m / eta                           # noqa
                        )                                                       # noqa
                    )                                                           # noqa
                )                                                               # noqa
            )                                                                   # noqa
        )                                                                       # noqa
        # fmt: on
        n_dot = N_dot / AVOGADROS_CONSTANT
        return n_dot

    def gas_flux_calibration(
        self,
        measurement,
        mol,
        mass,
        tspan=None,
        tspan_bg=None,
        ax=None,
        carrier_mol=None,
        mol_conc_ppm=None,
    ):
        """
        Args:
            measurement (MSMeasurement): The measurement with the ms_calibration data
            mol (str): The name of the molecule to calibrate
            mass (str): The mass to calibrate at
            tspan (iter): The timespan to average the signal over. Defaults to all
            tspan_bg (iter): Optional timespan at which the signal is at its background.
            ax (matplotlib axis): The axis on which to indicate what signal is used
                with a thicker line. Defaults to none
            carrier_mol (str): The name of the molecule of the carrier gas if
                a dilute analyte is used. Calibration assumes total flux of the
                capillary is the same as the flux of pure carrier gas. Defaults
                to None.
            mol_conc_ppm (float): Concentration of the dilute analyte in the carrier gas
                in ppm. Defaults to None.

        Returns MSCalResult: a ms_calibration result containing the sensitivity factor
            for mol at mass
        """
        t, S = measurement.grab_signal(mass, tspan=tspan, tspan_bg=tspan_bg)
        if ax:
            ax.plot(t, S, color=STANDARD_COLORS[mass], linewidth=5)
        if carrier_mol:
            if mol_conc_ppm:
                cal_type = "carrier_gas_flux_calibration"
            else:
                raise QuantificationError(
                    "Cannot use carrier gas calibration without analyte"
                    " concentration. mol_conc_ppm is missing."
                )
        elif mol_conc_ppm:
            raise QuantificationError(
                "Cannot use carrier gas calibration without carrier"
                " gas definition. carrier_mol is missing."
            )
        else:
            cal_type = "gas_flux_calibration"
            mol_conc_ppm = 10**6
            carrier_mol = mol

        n_dot = self.calc_n_dot_0(gas=carrier_mol) * mol_conc_ppm / 10**6
        F = np.mean(S) / n_dot
        return MSCalResult(
            name=f"{mol}@{mass}",
            mol=mol,
            mass=mass,
            cal_type=cal_type,
            F=F,
        )

class MSSpectrum(Spectrum):
    """Nothing to add to normal spectrum yet.
    TODO: Methods for co-plotting ref spectra from a database
    """

    pass


import glob
import os
import pysiaf

import astropy.coordinates as crd
import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np
from astroquery.irsa import Irsa
from matplotlib import cm
from scipy.io import readsav
from astropy.io import fits
from exoctk.utils import get_env_variables
from pysiaf.utils import rotations

EXOCTK_DATA = os.environ.get('EXOCTK_DATA')
if not EXOCTK_DATA:
    print('WARNING: The $EXOCTK_DATA environment variable is not set. '
          'Contamination overlap will not work. Please set the '
          'value of this variable to point to the location of the exoctk_data '
          'download folder.  Users may retreive this folder by clicking the '
          '"ExoCTK Data Download" button on the ExoCTK website, or by using '
          'the exoctk.utils.download_exoctk_data() function.'
          )
    TRACES_PATH = None
else:
    TRACES_PATH = os.path.join(EXOCTK_DATA, 'exoctk_contam', 'traces')

def sossFieldSim(ra, dec, binComp='', dimX=256):
    """ Produce a SOSS field simulation for a target.

    Parameters
    ----------
    ra: float
        The RA of the target.
    dec: float
        The Dec of the target.
    binComp: sequence
        The parameters of a binary companion.
    dimX: int
        The subarray size.

    Returns
    -------
    simuCub : np.ndarray
        The simulated data cube.
    """

    # STEP 1
    # Pulling stars from IRSA point-source catalog
    targetcrd = crd.SkyCoord(ra=ra, dec=dec, unit=(u.hour, u.deg))
    targetRA = targetcrd.ra.value
    targetDEC = targetcrd.dec.value
    info = Irsa.query_region(targetcrd,
                             catalog='fp_psc',
                             spatial='Cone',
                             radius=2.5 * u.arcmin)

    # Coordinates of all stars in FOV, including target
    allRA = info['ra'].data.data
    allDEC = info['dec'].data.data
    Jmag = info['j_m'].data.data
    Hmag = info['h_m'].data.data
    Kmag = info['k_m'].data.data

    # J-H band, H-K band. This will be used to derive the stellar Temps later
    J_Hobs = Jmag - Hmag
    H_Kobs = Hmag - Kmag

    # Determining target index by calculating the relative distance between
    # each source and the target. The target will have the smallest distance
    # from itself (oof) so whatever that index is will be the targetIndex
    aa = ((targetRA - allRA) * np.cos(targetDEC))
    distance = np.sqrt(aa**2 + (targetDEC - allDEC)**2)
    targetIndex = np.argmin(distance)

    # Add any missing companion
    if binComp != '':
        deg2rad = np.pi / 180
        bb = binComp[0] / 3600 / np.cos(allDEC[targetIndex] * deg2rad)
        allRA = np.append(allRA, (allRA[targetIndex] + bb))
        allDEC = np.append(allDEC, (allDEC[targetIndex] + binComp[1] / 3600))
        Jmag = np.append(Jmag, binComp[2])
        Hmag = np.append(Kmag, binComp[3])
        Kmag = np.append(Kmag, binComp[4])
        J_Hobs = Jmag - Hmag
        H_Kobs = Hmag - Kmag

    # Number of stars
    nStars = allRA.size

    # Restoring model parameters
    modelParam = readsav(os.path.join(TRACES_PATH, 'NIRISS', 'modelsInfo.sav'),
                         verbose=False)
    models = modelParam['models']
    modelPadX = modelParam['modelpadx']
    modelPadY = modelParam['modelpady']
    dimXmod = modelParam['dimxmod']
    dimYmod = modelParam['dimymod']
    jhMod = modelParam['jhmod']
    hkMod = modelParam['hkmod']
    teffMod = modelParam['teffmod']

    # Find/assign Teff of each star
    starsT = np.empty(nStars)
    for j in range(nStars):
        color_separation = (J_Hobs[j] - jhMod)**2 + (H_Kobs[j] - hkMod)**2
        min_separation_ind = np.argmin(color_separation)
        starsT[j] = teffMod[min_separation_ind]

    sweetSpot = dict(x=856, y=107, RA=allRA[targetIndex],
                     DEC=allDEC[targetIndex], jmag=Jmag[targetIndex])

    radeg = 180 / np.pi
    niriss_pixel_scale = 0.065  # arcsec
    # offset between all stars and target
    dRA = (allRA - sweetSpot['RA']) * np.cos(sweetSpot['DEC'] / radeg) * 3600
    dDEC = (allDEC - sweetSpot['DEC']) * 3600

    # Put field stars positions and magnitudes in structured array
    _ = dict(RA=allRA, DEC=allDEC, dRA=dRA, dDEC=dDEC, jmag=Jmag, T=starsT,
             x=np.empty(nStars), y=np.empty(nStars), dx=np.empty(nStars),
             dy=np.empty(nStars))
    stars = np.empty(nStars,
                     dtype=[(key, val.dtype) for key, val in _.items()])
    for key, val in _.items():
        stars[key] = val

    # Initialize final fits cube that contains the modelled traces
    # with contamination
    PAmin = 0  # instrument PA, degrees
    PAmax = 360
    dPA = 1  # degrees

    # Set of IPA values to cover
    PAtab = np.arange(PAmin, PAmax, dPA)    # degrees
    nPA = len(PAtab)

    dimY = 2048
    # cube of trace simulation at every degree of field rotation,
    # +target at O1 and O2
    simuCube = np.zeros([nPA + 2, dimY, dimX])

    saveFiles = glob.glob(
        os.path.join(
            TRACES_PATH,
            'NIRISS',
            '*modelOrder12*.sav'))

    # Big loop to generate a simulation at each instrument PA

    for kPA in range(PAtab.size):
        APA = PAtab[kPA]
        print('Generating field at APA : {}'.format(str(APA)))

        V3PA = APA + 0.57  # from APT

        sindx = np.sin((np.pi / 2) + APA / radeg) * stars['dDEC']
        cosdx = np.cos((np.pi / 2) + APA / radeg) * stars['dDEC']
        nps = niriss_pixel_scale
        stars['dx'] = (np.cos((np.pi / 2) + APA / radeg)
                       * stars['dRA'] - sindx) / nps
        stars['dy'] = (np.sin((np.pi / 2) + APA / radeg)
                       * stars['dRA'] + cosdx) / nps
        stars['x'] = stars['dx'] + sweetSpot['x']
        stars['y'] = stars['dy'] + sweetSpot['y']

        # Retain stars that are within the Direct Image NIRISS POM FOV
        ind, = np.where((stars['x'] >= -162) & (stars['x'] <= 2047 + 185) &
                        (stars['y'] >= -154) & (stars['y'] <= 2047 + 174))
        starsInFOV = stars[ind]

        for i in range(len(ind)):
            intx = round(starsInFOV['dx'][i])
            inty = round(starsInFOV['dy'][i])

            k = np.where(teffMod == starsInFOV['T'][i])[0][0]

            fluxscale = 10.0**(-0.4 *
                               (starsInFOV['jmag'][i] - sweetSpot['jmag']))

            # deal with subection sizes.
            # these variables will determine where the
            # trace will land on the array based on the
            # neighbor's position relative to the target's position
            mx0 = int(modelPadX - intx)
            mx1 = int(modelPadX - intx + dimX)
            my0 = int(modelPadY - inty)
            my1 = int(modelPadY - inty + dimY)

            if (mx0 > dimXmod) or (my0 > dimYmod):
                continue
            if (mx1 < 0) or (my1 < 0):
                continue

            x0 = (mx0 < 0) * (-mx0)
            y0 = (my0 < 0) * (-my0)
            mx0 *= (mx0 >= 0)
            mx1 = dimXmod if mx1 > dimXmod else mx1
            my0 *= (my0 >= 0)
            my1 = dimYmod if my1 > dimYmod else my1

            # if target and first kPA, add target traces of order 1 and 2
            # in output cube
            if (intx == 0) & (inty == 0) & (kPA == 0):
                fNameModO12 = saveFiles[k]

                modelO12 = readsav(fNameModO12, verbose=False)['modelo12']
                ord1 = modelO12[0, my0:my1, mx0:mx1] * fluxscale
                ord2 = modelO12[1, my0:my1, mx0:mx1] * fluxscale
                simuCube[0, y0:y0 + my1 - my0, x0:x0 + mx1 - mx0] = ord1
                simuCube[1, y0:y0 + my1 - my0, x0:x0 + mx1 - mx0] = ord2

            if (intx != 0) or (inty != 0):
                mod = models[k, my0:my1, mx0:mx1]
                simuCube[kPA + 2, y0:y0 + my1 - my0,
                         x0:x0 + mx1 - mx0] += mod * fluxscale
    return simuCube

def gtsFieldSim(ra, dec, filter, binComp=''):
    """ Produce a Grism Time Series field simulation for a target.
    Parameters
    ----------
    ra : float
        The RA of the target.
    dec : float
        The Dec of the target.
    filter : str
        The NIRCam filter being used. Can either be:
        'F444W' or 'F322W2' (case-sensitive)
    binComp : sequence
        The parameters of a binary companion.

    Returns
    -------
    simuCube : np.ndarray
        The simulated data cube. Index 0 and 1 (axis=0) show the trace of
        the target for orders 1 and 2 (respectively). Index 2-362 show the trace
        of the target at every position angle (PA) of the instrument.
    """
    # Instantiate a pySIAF object
    siaf = pysiaf.Siaf('NIRCam')

    full = siaf.apertures['NRCA5_FULL']
    if filter == 'F444W':
        aper = siaf.apertures['NRCA5_GRISM256_F444W']
    elif filter == 'F322W2':
        aper = siaf.apertures['NRCA5_GRISM256_F322W2']

    # Calling the variables
    deg2rad = np.pi / 180
    subX, subY = aper.XSciSize, aper.YSciSize
    rad = 2.5  # arcmins
    pixel_scale = 0.063  # arsec/pixel
    V3PAs = np.arange(0, 360, 1)
    nPA = len(V3PAs)
    # Generate cube of field simulation at every degree of APA rotation
    simuCube = np.zeros([nPA + 1, subY, subX])
    xSweet, ySweet = aper.reference_point('det')
    add_to_v3pa = aper.V3IdlYAngle
    # NIRCam Full Frame dimensions
    rows, cols = full.corners('det')
    minrow, maxrow = rows.min(), rows.max()
    mincol, maxcol = cols.min(), cols.max()

    #############################STEP 1#####################################
    ########################################################################
    # Converting to degrees
    targetcrd = crd.SkyCoord(ra=ra, dec=dec, unit=(u.hour, u.deg))
    targetRA = targetcrd.ra.value
    targetDEC = targetcrd.dec.value

    # Querying for neighbors with 2MASS IRSA's fp_psc (point-source catalog)
    info = Irsa.query_region(targetcrd, catalog='fp_psc', spatial='Cone',
                             radius=rad * u.arcmin)

    # Coordinates of all the stars in FOV, including target
    allRA = info['ra'].data.data
    allDEC = info['dec'].data.data

    # Initiating a dictionary to hold all relevant star information
    stars = {}
    stars['RA'], stars['DEC'] = allRA, allDEC

    #############################STEP 2#####################################
    ########################################################################
    sindRA = (targetRA - stars['RA']) * np.cos(targetDEC)
    cosdRA = targetDEC - stars['DEC']
    distance = np.sqrt(sindRA**2 + cosdRA**2)
    if np.min(distance) > 1.0*(10**-4):
        coords = crd.SkyCoord(ra=ra, dec=dec, unit=(u.hour, u.deg)).to_string('decimal')
        ra, dec = coords.split(' ')[0], coords.split(' ')[1]
        raise Exception('Unable to detect a source with coordinates [RA: {}, DEC: {}] within IRSA`s 2MASS Point-Source Catalog. Please enter different coordinates or contact the JWST help desk.'.format(str(ra), str(dec)))

    targetIndex = np.argmin(distance)

    # Restoring model parameters
    modelParam = readsav(os.path.join(TRACES_PATH, 'NIRISS', 'modelsInfo.sav'),
                         verbose=False)
    models = modelParam['models']
    modelPadX = modelParam['modelpadx']
    modelPadY = modelParam['modelpady']
    dimXmod = modelParam['dimxmod']
    dimYmod = modelParam['dimymod']
    jhMod = modelParam['jhmod']
    hkMod = modelParam['hkmod']
    teffMod = modelParam['teffmod']

    #############################STEP 3#####################################
    ########################################################################
    # JHK bands of all stars in FOV, including target
    Jmag = info['j_m'].data.data
    Hmag = info['h_m'].data.data
    Kmag = info['k_m'].data.data
    # J-H band, H-K band. This will be used to derive the Teff
    J_Hobs = Jmag - Hmag
    H_Kobs = Hmag - Kmag

    # Add any missing companion
    if binComp != '':
        bb = binComp[0] / 3600 / np.cos(allDEC[targetIndex] * deg2rad)
        allRA = np.append(allRA, (allRA[targetIndex] + bb))
        allDEC = np.append(allDEC, (allDEC[targetIndex] + binComp[1] / 3600))
        Jmag = np.append(Jmag, binComp[2])
        Hmag = np.append(Kmag, binComp[3])
        Kmag = np.append(Kmag, binComp[4])
        J_Hobs = Jmag - Hmag
        H_Kobs = Hmag - Kmag

    # Number of stars
    nStars = stars['RA'].size

    # Find/assign Teff of each star
    starsT = np.empty(nStars)
    for j in range(nStars):
        color_separation = (J_Hobs[j] - jhMod)**2 + (H_Kobs[j] - hkMod)**2
        min_separation_ind = np.argmin(color_separation)
        starsT[j] = teffMod[min_separation_ind]

    # Record keeping
    stars['Temp'] = starsT
    stars['Jmag'] = Jmag

    #############################STEP 4#####################################
    ########################################################################
    # Calculate corresponding V2/V3 (TEL) coordinates for Sweetspot
    v2targ, v3targ = aper.det_to_tel(xSweet, ySweet)

    for V3PA in range(0, nPA, 1):
        # Get APA from V3PA
        APA = V3PA + add_to_v3pa
        if APA > 360:
            APA = APA-360
        elif APA < 0:
            APA = APA+360

        print('Generating field at APA : {}'.format(str(APA)))

        # Get target's attitude matrix for each Position Angle
        attitude = rotations.attitude_matrix(v2targ, v3targ,
                                             targetRA, targetDEC,
                                             APA)

        xdet, ydet = [], []
        xsci, ysci = [], []
        for starRA, starDEC in zip(stars['RA'], stars['DEC']):
            # Get the TEL coordinates of each star w attitude matrix
            V2, V3 = rotations.sky_to_tel(attitude, starRA, starDEC)
            # Convert to arcsec and turn to a float
            V2, V3 = V2.to(u.arcsec).value, V3.to(u.arcsec).value

            XDET, YDET = aper.tel_to_det(V2, V3)
            XSCI, YSCI = aper.det_to_sci(XDET, YDET)

            xdet.append(XDET)
            ydet.append(YDET)
            xsci.append(XSCI)
            ysci.append(YSCI)

        # Record keeping
        stars['xdet'], stars['ydet'] = np.array(xdet), np.array(ydet)
        stars['xsci'], stars['ysci'] = np.array(xsci), np.array(ysci)

        sci_targx, sci_targy = stars['xsci'][targetIndex],\
            stars['ysci'][targetIndex]

    #############################STEP 5#####################################
    ########################################################################
        inFOV = []
        for star in range(0, nStars):

            x, y = stars['xdet'][star], stars['ydet'][star]
            if (mincol < x) & (x < maxcol) & (minrow < y) & (y < maxrow):
                inFOV.append(star)

        inFOV = np.array(inFOV)

    #############################STEP 6#####################################
    ########################################################################
        nircam_path = 'NIRCam_F444W' if filter == 'F444W' else 'NIRCam_F322W2'
        fitsFiles = glob.glob(
            os.path.join(
                TRACES_PATH,
                nircam_path,
                'rot*.fits'))
        fitsFiles = np.sort(fitsFiles)

        for idx in inFOV:

            sci_dx = round(sci_targx - stars['xsci'][idx])
            sci_dy = round(sci_targy - stars['ysci'][idx])
            temp = stars['Temp'][idx]

            for file in fitsFiles:
                if str(temp) in file:
                    trace = fits.getdata(file, 1)[0]

            fluxscale = 10.0**(-0.4 * \
                               (stars['Jmag'][idx] - stars['Jmag'][targetIndex]))

            # Padding array
            pad_trace = np.pad(trace, pad_width=5000, mode='constant',
                               constant_values=0)

            # Determine the highest pixel value of trace
            maxY, maxX = np.where(pad_trace == pad_trace.max())
            peakY, peakX = maxY[0], maxX[0]

            # Use relative distances (sci_dx, sci_dy) to find target
            xTarg = peakX + sci_dx
            yTarg = peakY + sci_dy

            # Use the (xTarg, yTarg) coordinates to slice out subarray
            # remember X is columns, Y is rows
            dimX0, dimX1 = xTarg - sci_targx, xTarg + subX - sci_targx
            dimY0, dimY1 = yTarg - sci_targy, yTarg + subY - sci_targy

            if dimX0 < 0:
                dimX0 = 0
                dimX1 = subX
            if dimY0 < 0:
                dimY0 = 0
                dimY1 = subY

            traceX, traceY = np.shape(pad_trace)[1], np.shape(pad_trace)[0]
            if dimX1 > traceX:
                dimX1 = traceX
                dimX0 = traceX - subX
            if dimY1 > traceY:
                dimY1 = traceY
                dimY0 = traceY - subY

            if (dimX1 < 0) or (dimY1 < 0):
                continue

            # -1 because pySIAF is 1-indexed
            mx0, mx1 = int(dimX0) - 1, int(dimX1) - 1
            my0, my1 = int(dimY0) - 1, int(dimY1) - 1

            # Fleshing out index 0 of the simulation cube (trace of target)
            if (sci_dx == 0) & (sci_dy == 0):  # this is the target

                tr = pad_trace[my0:my1, mx0:mx1] * fluxscale
                trX, trY = np.shape(tr)[1], np.shape(tr)[0]

                simuCube[0, 0:trY, 0:trX] = tr

            # Fleshing out indexes 1-361 of the simulation cube
            # (trace of neighboring stars at every position angle)
            else:

                tr = pad_trace[my0:my1, mx0:mx1] * fluxscale
                trX, trY = np.shape(tr)[1], np.shape(tr)[0]
                simuCube[V3PA + 1, 0:trY, 0:trX] += tr

    return simuCube

def lrsFieldSim(ra, dec, binComp=''):
    """ Produce a Grism Time Series field simulation for a target.
    Parameters
    ----------
    ra : float
        The RA of the target.
    dec : float
        The Dec of the target.
    binComp : sequence
        The parameters of a binary companion.

    Returns
    -------
    simuCube : np.ndarray
        The simulated data cube. Index 0 and 1 (axis=0) show the trace of
        the target for orders 1 and 2 (respectively). Index 2-362 show the trace
        of the target at every position angle (PA) of the instrument.
    """
    #############################INSTRUMENT PARAMETERS######################
    ########################################################################
    # Instantiate a pySIAF object
    siaf = pysiaf.Siaf('MIRI')
    aper = siaf.apertures['MIRIM_SLITLESSPRISM']
    full = siaf.apertures['MIRIM_FULL']

    # Calling the variables
    deg2rad = np.pi / 180
    subX, subY = aper.XSciSize, aper.YSciSize
    rad = 2.0  # arcmins
    pixel_scale = 0.11  # arsec/pixel
    V3PAs = np.arange(0, 360, 1)
    nPA = len(V3PAs)
    # Generate cube of field simulation at every degree of APA rotation
    simuCube = np.zeros([nPA + 1, subY, subX])
    xSweet, ySweet = aper.reference_point('det')
    add_to_v3pa = aper.V3IdlYAngle
    # MIRI Full Frame dimensions
    rows, cols = full.corners('det')
    minrow, maxrow = rows.min(), rows.max()
    mincol, maxcol = cols.min(), cols.max()

    #############################STEP 1#####################################
    ########################################################################
    # Converting to degrees
    targetcrd = crd.SkyCoord(ra=ra, dec=dec, unit=(u.hour, u.deg))
    targetRA = targetcrd.ra.value
    targetDEC = targetcrd.dec.value

    # Querying for neighbors with 2MASS IRSA's fp_psc (point-source catalog)
    info = Irsa.query_region(targetcrd, catalog='fp_psc', spatial='Cone',
                             radius=rad * u.arcmin)

    # Coordinates of all the stars in FOV, including target
    allRA = info['ra'].data.data
    allDEC = info['dec'].data.data

    # Initiating a dictionary to hold all relevant star information
    stars = {}
    stars['RA'], stars['DEC'] = allRA, allDEC

    #############################STEP 2#####################################
    ########################################################################
    sindRA = (targetRA - stars['RA']) * np.cos(targetDEC)
    cosdRA = targetDEC - stars['DEC']
    distance = np.sqrt(sindRA**2 + cosdRA**2)
    if np.min(distance) > 1.0*(10**-4):
        coords = crd.SkyCoord(ra=ra, dec=dec, unit=(u.hour, u.deg)).to_string('decimal')
        ra, dec = coords.split(' ')[0], coords.split(' ')[1]
        raise Exception('Unable to detect a source with coordinates [RA: {}, DEC: {}] within IRSA`s 2MASS Point-Source Catalog. Please enter different coordinates or contact the JWST help desk.'.format(str(ra), str(dec)))

    targetIndex = np.argmin(distance)

    # Restoring model parameters
    modelParam = readsav(os.path.join(TRACES_PATH, 'NIRISS', 'modelsInfo.sav'),
                         verbose=False)
    models = modelParam['models']
    modelPadX = modelParam['modelpadx']
    modelPadY = modelParam['modelpady']
    dimXmod = modelParam['dimxmod']
    dimYmod = modelParam['dimymod']
    jhMod = modelParam['jhmod']
    hkMod = modelParam['hkmod']
    teffMod = modelParam['teffmod']

    #############################STEP 3#####################################
    ########################################################################
    # JHK bands of all stars in FOV, including target
    Jmag = info['j_m'].data.data
    Hmag = info['h_m'].data.data
    Kmag = info['k_m'].data.data
    # J-H band, H-K band. This will be used to derive the Teff
    J_Hobs = Jmag - Hmag
    H_Kobs = Hmag - Kmag

    # Add any missing companion
    if binComp != '':
        bb = binComp[0] / 3600 / np.cos(allDEC[targetIndex] * deg2rad)
        allRA = np.append(allRA, (allRA[targetIndex] + bb))
        allDEC = np.append(allDEC, (allDEC[targetIndex] + binComp[1] / 3600))
        Jmag = np.append(Jmag, binComp[2])
        Hmag = np.append(Kmag, binComp[3])
        Kmag = np.append(Kmag, binComp[4])
        J_Hobs = Jmag - Hmag
        H_Kobs = Hmag - Kmag

    # Number of stars
    nStars = stars['RA'].size

    # Find/assign Teff of each star
    starsT = np.empty(nStars)
    for j in range(nStars):
        color_separation = (J_Hobs[j] - jhMod)**2 + (H_Kobs[j] - hkMod)**2
        min_separation_ind = np.argmin(color_separation)
        starsT[j] = teffMod[min_separation_ind]

    # Record keeping
    stars['Temp'] = starsT
    stars['Jmag'] = Jmag

    #############################STEP 4#####################################
    ########################################################################
    # Calculate corresponding V2/V3 (TEL) coordinates for Sweetspot
    v2targ, v3targ = aper.det_to_tel(xSweet, ySweet)

    for V3PA in range(0, nPA, 1):
        # Get APA from V3PA
        APA = V3PA + add_to_v3pa
        if APA > 360:
            APA = APA-360
        elif APA < 0:
            APA = APA+360

        print('Generating field at APA : {}'.format(str(APA)))

        # Get target's attitude matrix for each Position Angle
        attitude = rotations.attitude_matrix(v2targ, v3targ,
                                             targetRA, targetDEC,
                                             APA)

        xdet, ydet = [], []
        xsci, ysci = [], []
        for starRA, starDEC in zip(stars['RA'], stars['DEC']):
            # Get the TEL coordinates of each star w attitude matrix
            V2, V3 = rotations.sky_to_tel(attitude, starRA, starDEC)
            # Convert to arcsec and turn to a float
            V2, V3 = V2.to(u.arcsec).value, V3.to(u.arcsec).value

            XDET, YDET = aper.tel_to_det(V2, V3)
            XSCI, YSCI = aper.det_to_sci(XDET, YDET)

            xdet.append(XDET)
            ydet.append(YDET)
            xsci.append(XSCI)
            ysci.append(YSCI)

        # Record keeping
        stars['xdet'], stars['ydet'] = np.array(xdet), np.array(ydet)
        stars['xsci'], stars['ysci'] = np.array(xsci), np.array(ysci)

        sci_targx, sci_targy = stars['xsci'][targetIndex],\
            stars['ysci'][targetIndex]
    #############################STEP 5#####################################
    ########################################################################
        inFOV = []
        for star in range(0, nStars):

            x, y = stars['xdet'][star], stars['ydet'][star]
            if (mincol < x) & (x < maxcol) & (minrow < y) & (y < maxrow):
                inFOV.append(star)

        inFOV = np.array(inFOV)

    #############################STEP 6#####################################
    ########################################################################
        fitsFiles = glob.glob(os.path.join(TRACES_PATH, 'MIRI', 'LOW*.fits'))
        fitsFiles = np.sort(fitsFiles)

        for idx in inFOV:

            sci_dx = round(sci_targx - stars['xsci'][idx])
            sci_dy = round(sci_targy - stars['ysci'][idx])

            temp = stars['Temp'][idx]

            for file in fitsFiles:
                if str(temp) in file:
                    trace = fits.getdata(file)[0]

            fluxscale = 10.0**(-0.4 * \
                               (stars['Jmag'][idx] - stars['Jmag'][targetIndex]))

            # Padding array
            pad_trace = np.pad(trace, pad_width=5000, mode='constant',
                               constant_values=0)

            # Determine the highest pixel value of trace
            maxY, maxX = np.where(pad_trace == pad_trace.max())
            peakY, peakX = maxY[0], maxX[0]

            # Use relative distances (sci_dx, sci_dy) to find target
            # xTarg,yTarg are essentially the "sweetspot" in the PADDED arr
            xTarg = peakX + sci_dx
            yTarg = peakY + sci_dy

            # Use the (xTarg, yTarg) coordinates to slice out subarray
            # remember X is columns, Y is rows
            dimX0, dimX1 = xTarg - sci_targx, xTarg + subX - sci_targx
            dimY0, dimY1 = yTarg - sci_targy, yTarg + subY - sci_targy

            if dimX0 < 0:
                dimX0 = 0
                dimX1 = subX
            if dimY0 < 0:
                dimY0 = 0
                dimY1 = subY

            traceX, traceY = np.shape(pad_trace)[1], np.shape(pad_trace)[0]
            if dimX1 > traceX:
                dimX1 = traceX
                dimX0 = traceX - subX
            if dimY1 > traceY:
                dimY1 = traceY
                dimY0 = traceY - subY

            if (dimX1 < 0) or (dimY1 < 0):
                continue

            # -1 because pySIAF is 1-indexed
            mx0, mx1 = int(dimX0) - 1, int(dimX1) - 1
            my0, my1 = int(dimY0) - 1, int(dimY1) - 1

            # Fleshing out index 0 of the simulation cube (trace of target)
            if (sci_dx == 0) & (sci_dy == 0):  # this is the target

                tr = pad_trace[my0:my1, mx0:mx1] * fluxscale
                trX, trY = np.shape(tr)[1], np.shape(tr)[0]

                simuCube[0, 0:trY, 0:trX] = tr

            # Fleshing out indexes 1-361 of the simulation cube
            # (trace of neighboring stars at every position angle)
            else:

                tr = pad_trace[my0:my1, mx0:mx1] * fluxscale
                trX, trY = np.shape(tr)[1], np.shape(tr)[0]
                simuCube[V3PA + 1, 0:trY, 0:trX] += tr

    return simuCube

def fieldSim(ra, dec, instrument, binComp='', testing=False):
    """ Wraps ``sossFieldSim``, ``gtsFieldSim``, and ``lrsFieldSim`` together.
    Produces a field simulation for a target using any instrument (NIRISS,
    NIRCam, or MIRI).

    Parameters
    ----------
    ra : float
        The RA of the target.
    dec : float
        The Dec of the target.
    instrument : str
        The instrument the contamination is being calculated for.
        Can either be (case-sensitive):
        'NIRISS', 'NIRCam F322W2', 'NIRCam F444W', 'MIRI'
    binComp : sequence
        The parameters of a binary companion.
    testing : bool
        Shoud be ``True`` if running fieldSim for testing / troubleshooting
        purposes. This will generate a matplotlib figure showing the target
        FOV. The neighboring stars in this FOV will be included in the
        contamination calculation (contamFig.py).

    Returns
    -------
    simuCube : np.ndarray
        The simulated data cube. Index 0 and 1 (axis=0) show the trace of
        the target for orders 1 and 2 (respectively). Index 2-362 show the trace
        of the target at every position angle (PA) of the instrument.
    plt.plot() : matplotlib object
        A plot. Only if `testing` parameter is set to True.
    """

    # Calling the variables which depend on what instrument you use
    if instrument == 'NIRISS':
        simuCube = sossFieldSim(ra, dec, binComp)

    elif instrument == 'NIRCam F444W':
        simuCube = gtsFieldSim(ra, dec, 'F444W', binComp)

    elif instrument == 'NIRCam F322W2':
        simuCube = gtsFieldSim(ra, dec, 'F322W2', binComp)

    elif instrument == 'MIRI':
        simuCube = lrsFieldSim(ra, dec, binComp)

    return simuCube

if __name__ == '__main__':
    ra, dec = "04 25 29.0162", "-30 36 01.603"  # Wasp 79
    #sossFieldSim(ra, dec)
    if EXOCTK_DATA:
        fieldSim(ra, dec, instrument='NIRISS')

import numpy,from scipy,import astropy,from astropy


#!/usr/bin/env python

import itertools as itt
import logging
import os
from datetime import datetime
from getpass import getuser

import numpy as np
import pandas as pd

from .generate import get_percentile_diff, get_inducible_pairs
from .. import hgnc, mi, up, snp as rs
from ..struct.hetnet import HetNet, encode_color_path

log = logging.getLogger()

CP1 = ('p', 'm')
CP1_TERMINAL = ('p', ('m', ()))
CP1_STR = encode_color_path(CP1_TERMINAL, color_path_type='terminal')
CP2 = ('p', 'p', ('g', (('regulated', 'T'),)))
CP2_STR = encode_color_path(CP2, color_path_type='terminal')

def convert_simple_to_terminal(cp):
    *head, tail = cp
    tail = (tail, ())

    if isinstance(cp, list):
        return list(head) + [tail]

    return tuple(head) + (tail,)

def generate_toy(seed=None):
    np.random.seed(seed)
    h = HetNet({
        "g": {
            "regulated": ["T", "F"]  # significant regulation
        },
        "p": {
        },
        "m": {
        },
        "s": {
        }
    })

    n_genes = 426
    n_proteins = 275
    n_mirnas = 16
    n_snps = 14

    n_mirna_encoding_genes = 8
    n_protein_encoding_genes = 275

    n_ppis = 292
    n_mtis = 51

    # ??? not sure about numbers
    n_regulated = 25
    n_coexprs = 400

    regulated = np.random.choice(np.arange(n_genes), size=n_regulated, replace=False)
    genes = {}
    for i in range(n_genes):
        gene = hgnc(i)
        genes[i] = gene
        h.add_node(gene, dict(color='g', annotations={'regulated': ("T" if i in regulated else "F")}))

    protein_encoding_genes = sorted(genes.values())[:n_protein_encoding_genes]
    proteins = {}
    for i, gene in zip(range(n_proteins), protein_encoding_genes):
        protein = up(i)
        proteins[i] = protein
        h.add_node(protein, dict(color='p', annotations={}))
        h.add_edge(gene, protein)

    mirna_encoding_genes = list(
        np.random.choice(list(set(genes.values()) - set(protein_encoding_genes)), size=n_mirna_encoding_genes,
                         replace=False))
    mirnas = {}
    for i, gene in zip(range(n_mirnas), 2 * mirna_encoding_genes):
        mirna = mi(i)
        mirnas[i] = mirna
        h.add_node(mirna, dict(color='m', annotations={}))
        h.add_edge(gene, mirna)

    h.mirna_encoding = mirna_encoding_genes

    mutations = np.random.choice(list(set(genes.values()) - set(protein_encoding_genes) - set(mirna_encoding_genes)),
                                 size=n_snps, replace=False)
    snps = {}
    for i, gene in zip(range(n_snps), mutations):
        snp = rs(i)
        snps[i] = snp
        h.add_node(snp, dict(color='s', annotations={}))
        h.add_edge(gene, snp)

    pp = list(itt.combinations(proteins, 2))
    for i in np.random.choice(len(pp), size=n_ppis, replace=False):
        a, b = pp[i]
        h.add_edge(proteins[a], proteins[b])

    gg = list(itt.combinations(genes, 2))
    for i in np.random.choice(len(gg), size=n_coexprs, replace=False):
        a, b = gg[i]
        h.add_edge(genes[a], genes[b])

    mp = list(itt.product(proteins, mirnas))
    for i in np.random.choice(len(mp), size=n_mtis, replace=False):
        a, b = mp[i]
        h.add_edge(proteins[a], mirnas[b])

    h.graph['generation_manifest'] = {
        'user': getuser(),
        'generation_time': str(datetime.now()),
        'np_random_seed': seed,
        'protein_encoding': protein_encoding_genes
    }

    return h

def induce_toy(graph, target_nodes, upper=100, lower=90, loc_coef=1.1, scale_coef=1.2, seed=None):
    """
    Generate network pertaining to actual biology, and induce 2 features over protein-coding genes:
        - protein-mirna
        - protein-protein-regulated_gene

    :param graph: the network to induce
    :param target_nodes: the nodes to induce in the network
    :param upper: the upper percentile to calculate for the color paths
    :param lower: the lower percentile to calculate for the color paths
    :param loc_coef: the factor to multiply the upper percentile for induction for the mean of random gaussian sampling
    :param scale_coef: the factor to multiply the upper-lower percentile difference for the standard deviation of random
                        gaussian sampling
    :param seed: seed for numpy random number generator
    :return:
    """

    np.random.seed(seed)
    pairs = []

    # TODO factor out induction parameters
    p100a, p98a, pt2a = get_percentile_diff(graph, CP1, upper, lower, color_path_type='simple')
    log.debug(
        'CP {} -> {}%: {}, {}%: {}, D{}%: {}'.format(encode_color_path(CP1, color_path_type="simple"), upper, p100a,
                                                     lower, p98a, upper - lower, pt2a))

    p100b, p98b, pt2b = get_percentile_diff(graph, CP2, upper, lower, color_path_type='terminal')
    log.debug(
        'CP {} -> {}%: {}, {}%: {}, D{}%: {}'.format(encode_color_path(CP2, color_path_type="terminal"), upper, p100b,
                                                     lower, p98b, upper - lower, pt2b))

    for node in target_nodes:
        for edge in get_inducible_pairs(graph, node, CP1,
                                        int(np.random.normal(loc=(loc_coef * p100a), scale=(scale_coef * pt2a))),
                                        color_path_type='simple'):
            pairs.append(edge)
        for edge in get_inducible_pairs(graph, node, CP2,
                                        int(np.random.normal(loc=(loc_coef * p100b), scale=(scale_coef * pt2b))),
                                        color_path_type='terminal'):
            pairs.append(edge)

    """
    params = [
        (CP1, upper, lower, 1.1, 1.2, 'simple'),
        (CP2, upper, lower, 1.1, 1.2, 'terminal')
    ]
    pairs = []
    for cp, up, lo, lc, sc, cpt in params:
        p100, p98, pt2 = get_percentile_diff(graph, cp, up, lo, color_path_type=cpt)
        for node in target_nodes:
            ne = int(np.random.normal(loc=(lc * p100), scale=(sc * pt2)))
            pairs.extend(get_inducible_pairs(graph, node, cp, ne, color_path_type=cpt))
    """

    for a, b in set(pairs):
        graph.add_edge(a, b)

    graph.graph['induction_manifest'] = {
        'user': getuser(),
        'induction_time': str(datetime.now()),
        'induced': sorted(target_nodes),
        'upper': upper,
        'lower': lower,
        'loc_coef': loc_coef,
        'scale_coef': scale_coef,
        'np_random_seed': seed
    }

    return graph

def main(directory, percent=0.8, seed=None):
    """
    :param directory: output directory
    :param percent: if given, outputs a training and test manifest
    :param seed: seed for numpy random number generator
    """

    np.random.seed(seed)

    h = generate_toy()
    n_induce = int(0.5 + 7 / percent)
    target_nodes = np.random.choice(h.graph['generation_manifest']['protein_encoding'], size=n_induce, replace=False)
    hn = induce_toy(h, target_nodes)
    hn.to_resource(directory)

    induced = hn.graph['induction_manifest']['induced']
    nodes = sorted(hn.graph['generation_manifest']['protein_encoding'])

    full_induction_manifest = pd.DataFrame([node in induced for node in nodes], index=nodes, columns=['induced'])
    full_induction_manifest.to_csv(os.path.join(directory, 'full_induce_manifest.csv'))

    not_induced = list(set(nodes) - set(induced))
    n_induced = len(induced)
    n_not_induced = len(not_induced)

    np.random.shuffle(induced)
    np.random.shuffle(not_induced)

    head_induced = induced[:int(percent * n_induced)]
    tail_induced = induced[int(percent * n_induced):]

    head_not_induced = not_induced[:int(percent * n_not_induced)]
    tail_not_induced = not_induced[int(percent * n_not_induced):]

    head = sorted(head_induced) + sorted(head_not_induced)
    tail = sorted(tail_induced) + sorted(tail_not_induced)

    full_induction_manifest.loc[head].to_csv(os.path.join(directory, 'training_induce_manifest.csv'))
    full_induction_manifest.loc[tail].to_csv(os.path.join(directory, 'test_induce_manifest.csv'))