#!/urs/bin/env python
# -*- coding: utf-8 -*-
"""/stations.py
Author: Hunter Mellema
Summary: Defines measuring stations for simulation for CR3BP navigation simulations
Project: Advanced State Estimation Final Project
Author: Hunter Mellema
Date: April 2020
"""
# === Begin Imports ===
# third party
import numpy as np
from numpy.linalg import norm

# local imports
from .constants import (
    THETA_0,
    R_E,
    W_E_ND,
    R_E_ND,
    EARTH_ANG_VEL_ND,
    NON_DIM_DIST_TO_DIM,
    MU_EARTH_MOON,
    E_M_OFFSET_ANGLE,
    OMEGA_CR3BP,
    CR3BP_ANG_VEL
)
from .measurements import R3Msr, Range

# std library
import json

# === End Imports ===

# Measurement Generation Functions
class CR3BPEarthStn(object):
    """ Represents a ground station at some point on the earth's surface that is taking measurements
    of a spacecraft.
    Args:
        stn_id (str, int): a unique string or integer id for the given station
        longitude (float): longitude of the station position in degrees
        latitude (float): latitude of the station in degrees
        el_mask (float): elevation below which the station will not take measurements (in degrees)
        covariance (np.ndarray): covariance for measurements take by this station. Will be added to
            every measurement taken by this station

    """

    def __init__(
        self,
        stn_name,
        stn_id,
        pos_ecef_nd,
        el_mask,
        cov,
        msrs,
        mu,
        lat=None,
        longitude=None,
        el=None,
    ):
        self.stn_name = stn_name
        self.stn_id = stn_id
        self.pos = pos_ecef_nd
        self.el_mask = el_mask
        self.cov = cov
        self.allowed_msrs = msrs
        self.lat = lat
        self.long = longitude
        self.el = el
        self.mu = mu

    @classmethod
    def from_db_object(cls, station_data):
        """ Constructs a stn from a row in the database

        Args:
            station_data (dict): dicitonary representing a row in the stations table

        Return:
            cls

        """
        pos_nd = lat_long_to_ecef(
            station_data["latitude_deg"],
            station_data["longitude_deg"],
            station_data["elevation_km"],
        ) * (1 / NON_DIM_DIST_TO_DIM)

        return cls(
            station_data["stn_name"],
            station_data["stn_id"],
            pos_nd,
            station_data["elevation_mask_deg"],
            json.loads(station_data["covariance"]),
            json.loads(station_data["measurement_types"]),
            station_data["mu"],
            station_data["latitude_deg"],
            station_data["longitude_deg"],
            station_data["elevation_km"],
        )

    def __repr__(self):
        string = """
        ==================================================
        ++++++++++++++++++++++++++++++++++++++++++++++++++
        STN: {}     | STN ID: {}
        ++++++++++++++++++++++++++++++++++++++++++++++++++
        Position:
            {} (Nondimensional CR3BP)
        Location (lat, long, el):
            {} deg N, {} deg E, {} Km,
        Elevation Mask: {} deg
        Measurement Covariance: {}
        ==================================================
        """.format(
            self.stn_name,
            self.stn_id,
            self.pos,
            self.lat,
            self.long,
            self.el,
            self.el_mask,
            self.cov,
        )

        return string

    def gen_msr(self, sc_state, time, msr_type="R3Msr"):
        """Generates a measurement if the spacecraft is visible

        Args:
            sc_state (np.ndarray): spacecraft state vector, pos, vel must be first 6 terms
            time (float): time since reference epoch
            msr_type (str): name of measurement to use (default = R3 (range and range rate))

        Returns:
            tuple(bool, float, np.ndarray)

        """
        valid, el, stn_state = self._check_elevation(sc_state, time)
        if valid:
            return globals()[msr_type].from_stn(
                time, sc_state, stn_state, self.stn_id, self.cov
            )
        else:
            return None

    def _check_elevation(self, sc_state, time):
        """Checks to see if the spacecraft is visible

        Args:
            sc_state (np.ndarray): spacecraft state vector, pos, vel must be first 6 terms
            time (float): time since reference epoch

        Returns:
            tuple(bool, float, np.ndarray)

        """
        stn_state_ns = self.state(time, include_shift=False)
        stn_state_shift = stn_state_ns + np.array([self.mu, 0.0, 0.0, 0.0, 0.0, 0.0])
        line_o_sight = sc_state[0:3] - stn_state_shift[0:3]
        num = np.dot(stn_state_ns[0:3], line_o_sight)
        denom = np.linalg.norm(stn_state_ns[0:3]) * np.linalg.norm(line_o_sight)
        zenel = np.arccos(num / denom)
        el = np.pi / 2 - zenel

        if el > np.deg2rad(self.el_mask):
            flag = True
        else:
            flag = False

        return (flag, el, stn_state_shift)

    def state(self, time_nd, include_shift=True):
        """Finds the station state in ECI at a given nondimensional CR3BP time

        Args:
            time_nd (float): non-dimensional CR3BP time past reference epoch

        Returns:
            np.ndarray([1x6])

        """
        rot_mat_ecef_eci = ecef_to_eci_nd(time_nd)
        rot_mat_eci_cr3bp = eci_to_cr3bp(time_nd)
        pos_eci = rot_mat_ecef_eci @ self.pos
        pos_cr3bp = rot_mat_eci_cr3bp @ pos_eci

        if include_shift:
            pos_cr3bp = pos_cr3bp + np.array([self.mu, 0.0, 0.0])

        vel_ecef = np.cross(EARTH_ANG_VEL_ND.T, self.pos.T).T
        vel_eci = rot_mat_ecef_eci @ vel_ecef
        vel_cr3bp_no_cr3bp_and_vel = rot_mat_eci_cr3bp @ vel_eci
        coriolis = np.cross(CR3BP_ANG_VEL.T, pos_cr3bp.T).T
        vel_cr3bp = vel_cr3bp_no_cr3bp_and_vel + coriolis

        state = np.concatenate((pos_cr3bp.T[0], vel_cr3bp.T[0]))

        return state

def lat_long_to_ecef(latitude, longitude, elevation):
    """Converts from latituted-longitude to cartesian position in ECEF
    Note: Assumes location is on the surface of the earth and uses a spherical
        earth model

    Args:
        latitude (float): latitude to convert (in degrees)
        longitude (float): longitude to convert (in degrees)
        elevation (float): elevation above spherical model (in km)

    Returns:
       np.ndarray([3x3])

    """
    phi = np.deg2rad(latitude)
    lam = np.deg2rad(longitude)

    pos_ecef = (R_E + elevation) * np.array(
        [[np.cos(phi) * np.cos(lam)], [np.cos(phi) * np.sin(lam)], [np.sin(phi)]]
    )

    return pos_ecef

def ecef_to_eci_nd(time_nd, theta_0=THETA_0):
    """Calculates rotation matrix to ECI at a given time
    Args:
       time_nd (float): nondimensional time since reference epoch
       theta_0 (float): initial rotation of the earth at the reference epoch
            [default=filtering.THETA_0]
    Returns:
       np.ndarray([3x3])
    """
    alpha = theta_0 + time_nd * W_E_ND
    rot_mat = np.array(
        [
            [np.cos(alpha), -np.sin(alpha), 0],
            [np.sin(alpha), np.cos(alpha), 0],
            [0, 0, 1],
        ]
    )

    return rot_mat

def eci_to_cr3bp(time_nd):
    """Calculates rotation matrix from ECI to CR3BP at a given nondimensional time

    Args:
       time_nd (float): nondimensional time since reference epoch

    Returns:
       np.ndarray([3x3])

    """
    axis_north_pole = np.array([
        np.sin(-E_M_OFFSET_ANGLE) * np.cos(time_nd * OMEGA_CR3BP),
        np.sin(-E_M_OFFSET_ANGLE) * np.sin(time_nd * OMEGA_CR3BP),
        np.cos(-E_M_OFFSET_ANGLE)
    ])
    z_axis = np.array([0.0, 0.0, 0.1])
    rot = rotate_plane(z_axis, axis_north_pole)

    return rot

def unitcross(a, b):
    c = np.cross(a, b)
    return c / norm(c)

def costheta(a, b):
    return np.dot(a, b) / (norm(a) * norm(b))

def rotate_plane(M, N):
    (x, y, z) = unitcross(M, N)
    c = costheta(M, N)
    s = np.sqrt(1-c*c)
    C = 1-c
    rmat = np.array(
        [
            [ x * x * C + c, x * y * C - z * s, x * z * C + y * s],
            [ y * x * C + z * s, y * y * C + c, y * z * C - x * s],
            [ z * x * C - y * s, z * y * C + x * s, z * z * C + c]
        ]
    )

    return rmat


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

################################################## ##############################
#########################  List of implemented limits  #########################
######################### ######################### #############################

### Define a class to store all the relevant data about the external limits
### as well as the way of recasting them to the fermion portal case. Then loads all the implemented exeternal limits

### L. Darme, S. Ellis, T. You, 29/12/2019

############################ Messy library import part #########################
import numpy as np

# Importing the other sub-Modules

import UsefulFunctions as uf
import Production as br
import Detection as de

verbose=True

class Limit:
    """
    Contain the informations to recast a given experimental limit

    exp:     The experiment associated to the limit

    channel: The type of limit (e.g MissingE, Decay ...)

    descr:   Any relevant additional details on the limit

    ref:     The Inspire reference for the limit

    mx_ini:  The range of Mx value on which the limit is relevant (by convention
             this is typically the mass of the lower state X1 in case there is
             two dark sector states, such as for decay limits)

    lim_ini: The actual limit (typically from a dark photon model)

    lim_inilist: Dictionnary with the initial limits stored in (mx, lim) list for different operator in case they exist ( "V" and "AV")
    """

    def __init__(self, name, exp, channel, Delini,interptype="log"):
        self.exp = exp
        self.channel = channel
        self.delini = Delini
        self.descr = "No description for this limit"
        self.ref = "No reference for this limit"
        self.combthr=1.1 # Standard value used when relevant to combine the upper and lower limit
        try:
            self.name = name
            self.mx_ini, self.lim_ini = uf.LoadData('LimData/'+name+'.dat',interptype)
#             if verbose: print("Loading: ",  self.name,self.lim_ini   )
        except:
            self.name = "NotDefined"
            self.mx_ini = np.logspace(-3., 1., 30.)
            self.lim_ini = np.zeros(np.shape(self.mx_ini ))

        try:
            self.model = model
        except:
            self.model = "No model for this limit"

        self.lim_inifull = {"V":(self.mx_ini,self.lim_ini),"AV":(self.mx_ini,self.lim_ini)}

    def UpdateLimIni(self,mx,lim,optype="V"):
#         print(optype,(mx,lim))
        self.lim_inifull.update({optype:(mx,lim)})
        if optype=="V":
            self.lim_ini = lim
            self.mx_ini = mx

    def recast(self,delrec, geff, optype,useHeavyMeson=False):
        # we define the recast based on the expected signal

        if self.channel == "decay":
            if self.model == "EFT": # For naive NoE based limits
                xi,EffLimtmp = de.GetNaiveDecayLimits( delrec,  self.exp ,10,geff,optype)
                xi,EffLim,EffLim_low = de.RecastDecayLimit(xi, EffLimtmp, delrec,delrec,self.exp,geff,optype)
                xi_full, Lim_full=uf.CombineUpDown(xi,EffLim_low,EffLim,self.combthr)
            else: # Standard recasting case
                xi_full, Lim_full=de.FastDecayLimit(self.exp,self.mx_ini, self.lim_ini , self.delini, delrec, geff, optype,self.combthr,useHeavyMeson)
        elif self.channel == "heavymesondecay":
            if self.model == "EFT":
                xi,EffLimtmp = de.GetNaiveDecayLimits( delrec, self.exp ,10,geff,optype,True)
                xi_heavy,EffLim_heavy,EffLim_heavy_low = de.RecastDecayLimit(xi,EffLimtmp , delrec,delrec,self.exp,geff,optype,True)
                xi_full, Lim_full=uf.CombineUpDown(xi_heavy,EffLim_heavy_low,EffLim_heavy,self.combthr)
            else:
                xi_full, Lim_full=de.FastDecayLimit(self.exp,self.mx_ini, self.lim_ini , self.delini, delrec, geff, optype,self.combthr,useHeavyMeson)
        elif self.channel=="scattering":
            xi_full, Lim_full=de.FastScatLimit(self.exp,self.mx_ini, self.lim_ini , self.delini,delrec, geff, optype)
        elif self.channel=="monogam":
            xi_full, Lim_full=de.FastMonoPhoton(self.exp,self.mx_ini, self.lim_ini , self.delini,delrec, geff, optype)
        elif self.channel=="invisibledecayBmtoKm":
            if self.exp == "babar":
                xi_full, Lim_full=  de.FastInvMesDecay("babar_BmtoKmnunu",delrec, geff, optype)
            elif self.exp == "belle2":
                xi_full, Lim_full=  de.FastInvMesDecay("belle2_BmtoKmnunu",delrec, geff, optype)
        elif self.channel=="invisibledecayBmtoPim":
            xi_full, Lim_full=  de.FastInvMesDecay("babar_BmtoPimnunu",delrec, geff, optype)
        elif self.channel == "invisibledecayB0toPi0":
            xi_full, Lim_full=  de.FastInvMesDecay("belle_B0toPi0nunu",delrec, geff, optype)
        elif self.channel == "invisibledecayB0toK0":
            xi_full, Lim_full=  de.FastInvMesDecay("belle_B0toK0nunu",delrec, geff, optype)
        elif self.channel=="monojet_down":
            xi_full, Lim_full = de.FastMonoJet(self.exp,self.mx_ini, self.lim_ini , self.delini,delrec, geff, optype)
        elif self.channel=="monojet_up":
            xi_full, Lim_full = de.FastMonoJet(self.exp,self.mx_ini, self.lim_ini , self.delini,delrec, geff, optype)
        elif self.channel == "invisibledecayKL0toPi0":
            if self.exp == "e391a":
                xi_full, Lim_full=  de.FastInvMesDecay("e391a_KL0toPi0nunu",delrec, geff, optype)
            if self.exp == "na62":
                xi_full, Lim_full=  de.FastInvMesDecay("na62_KL0toPi0nunu",delrec, geff, optype)
        elif self.channel == "invisibledecayPi0":
            xi_full, Lim_full=  de.FastInvMesDecay("na62_pi0toinvisible",delrec, geff, optype)
        elif self.channel == "invisibledecayJPsi":
            xi_full, Lim_full=  de.FastInvMesDecay("bes_JPsitoinvisible",delrec, geff, optype)
        elif self.channel == "invisibledecayUpsilon":
            xi_full, Lim_full=  de.FastInvMesDecay("babar_Upsilontoinvisible",delrec, geff, optype)
        elif self.channel == "invisibledecayKptoPip":
            if self.exp == "na62":
                xi_full, Lim_full=  de.FastInvMesDecay("na62_KptoPipa",delrec, geff, optype)
            elif self.exp == "e949":
                xi_full, Lim_full=  de.FastInvMesDecay("e949_KptoPipa",delrec, geff, optype)
        elif self.channel == "cosmicrays":
            xi_full, Lim_low_full, Lim_high_full = de.FastCRLimit("t2k", delrec, geff,optype)
            xi_full, Lim_full = uf.CombineUpDown(xi_full, Lim_low_full, Lim_high_full)
        elif self.channel == "low_cooling":
            xi_full, Lim_full= de.FastSN1987Limit(self.lim_inifull ,delrec, geff, optype,False)
        elif self.channel == "high_cooling":
            xi_full, Lim_full= de.FastSN1987Limit( self.lim_inifull ,delrec, geff, optype,True)
        else:
            print("Channel selected: ", self.channel, " is not currently implemented. Possible choices: \n",
            "'decay' : faser, mathusla, ship, charm, seaquest, seaquest_Phase2, lsnd \n",
            "heavymesondecay : ship"
            "'scattering' :  nova, miniboone, sbnd \n",
            "'missingE' : babar, belleII, atlas, lep \n",
             "'cosmicrays' : t2k (decay from cosmic ray showers into t2k), ",
             "cooling : for sn1987_low (pessimistic limit from SN1987 cooling), sn1987_high (optimistic limit from SN1987 cooling) ",
             "Invisible light meson decays: na62 (pi0decay and invisibledecayKptoPip) e949 (pi0decay and invisibledecayKptoPip), e391a (invisibledecayKL0toPi0)",
             "Invisible heavy meson decay: belle (invisibledecayB0toPi0 ,invisibledecayB0toK0) and belleII (invisibledecayBmtoKm)")
            xi_full= np.logspace(-3., 1., 30.);
            Lim_full =np.zeros(np.shape(xi_full))
        return xi_full,Lim_full

AllLimits={} # Dictionnary for all the limits

def UpdateLimitList(limit,name="noname"):

    if name=="noname":
        name=limit.exp+ "_"+limit.channel # We fill up the name automatically if not given

    if name in AllLimits: # Making a new limits if one already existing for the same
#         print("Test:", name, AllLimits)
        i = 1
        name_new=name+"_"+str(i)
        while (name_new in AllLimits):
            i += 1
            name_new=name+"_"+str(i)
        name=name_new
    AllLimits.update({name:limit})

## This function runs over a list of experiments and search channels,
## compute the recasted bound, print it to file and return everything in a dictionnary with a list of label
## Note that it can also just read an external file and use it directly as the limit

def GetLimits(LimList,Del,geff,optype="V",PrintToFile = False, ReadFromFile=False, filename=''):

    Res = {}
    LabelLimit = []

    # Making sure the  input for the effective couplings is okay
    if type(geff) is dict:
        ge=geff
    else:
        if np.ndim(geff[0]) >0:
            geffdiag=(geff[0][0][0],geff[1][0][0],geff[2][0][0])
            gMesDecay=(geff[1][2][1],geff[1][2][1],geff[1][2][0],geff[1][2][0],geff[1][1][0],geff[1][1][0]) # gd32, gd32,gd31,gd31,gd21,gd21
        else:
            geffdiag=  geff
            gMesDecay=(0,0,0,0,0,0)
        ge={"gu11":geffdiag[0],"gd11":geffdiag[1],"gl11":geffdiag[2],"gl22":geffdiag[2] \
        ,"gd32":gMesDecay[0],"gd31":gMesDecay[2],"gd21":gMesDecay[4]}#We make the couping a dictionnary to have more freedom

    resultsDict = {} #e.g. {"exp1_channel1":(mxlist, limlist), "exp2_channel2":{mxlist, limlist}, ... }
    expchannelList = [] #e.g. ["exp1_channel1", "exp2_channel2", etc.]

    for expchannel in LimList:
        if not expchannel in AllLimits:
            print( "Selected limit " + expchannel +" not available, choose from: ", AllLimits.keys())
#             raise Exception()
        else:
            if ReadFromFile:
                if filename == '':
                    filestr = 'Output/Lim_'+expchannel+".dat"
                else:
                    filestr = filename
                filedat = np.loadtxt(filestr)
                Mx, LamLim = np.transpose(filedat)
            else:
                #print(exp, channel)

                Mx, LamLim = AllLimits[expchannel].recast(Del,ge,optype)

            expchannelList.append(expchannel)
            resultsDict[expchannel] = (Mx,LamLim)

            if PrintToFile:
                if filename == '':
                    filestr = 'Output/Lim_'+expchannel+".dat"
                else:
                    filestr = filename
                np.savetxt(filestr, np.transpose(resultsDict[expchannel]))

    return resultsDict, expchannelList

################################################## #########################
#########################                           #########################
#########################  Loading external limits  #########################
#########################                           #########################
######################### ######################### #########################

geffem={"gu11":2/3.,"gd11":-1/3.,"gl11":-1.}

print("Loading limits:")

#########################  Scattering #########################
if verbose: print("Scattering ...")

miniboone_scattering=Limit("miniboone_1807.06137","miniboone","scattering",0.)
miniboone_scattering.descr="""
Limits from MiniBooNE collaboration for light dark matter scattering, produced at the beam dump.

Extracted from Figure 24.a, the data is given as epsilon^2, for alpha_D = 0.5,
we rescale it by 5 ^1/4. since the scattering limit scale as eps^4 alpha_D
"""
miniboone_scattering.ref="inspirehep.net/record/1682906"
miniboone_scattering.UpdateLimIni(miniboone_scattering.mx_ini,np.sqrt(miniboone_scattering.lim_ini)*np.power(5,1/4.) )
UpdateLimitList(miniboone_scattering)

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

sbnd_scattering=Limit("sbnd_1609.01770","sbnd","scattering",0.)
sbnd_scattering.descr="""
Projective Limits from SBND collaboration for light dark matter scattering
produced by 1609.01770, for 2. * 10^20  PoT, based on 10 events reach

Extracted from Figure 9.b, the data is given as epsilon^2, for alpha_D = 0.5
"""
sbnd_scattering.ref="inspirehep.net/record/1485563"
sbnd_scattering.UpdateLimIni(sbnd_scattering.mx_ini,np.sqrt(sbnd_scattering.lim_ini)*np.power(5,1/4.) )
UpdateLimitList(sbnd_scattering)

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

ship_scattering=Limit("ship_1609.01770","ship","scattering",0.)
ship_scattering.descr="""
Projective limits from SHIP collaboration for light dark matter scattering,
produced by  1609.01770, for 2. * 10^20  PoT, based on 10 events reach

Extracted from Figure 24.a, the data is given as epsilon^2, for alpha_D = 0.5
"""
ship_scattering.ref="inspirehep.net/record/1485563"
ship_scattering.UpdateLimIni(ship_scattering.mx_ini,np.sqrt(ship_scattering.lim_ini)*np.power(5,1/4.) )
UpdateLimitList(ship_scattering)

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

nova_scattering=Limit("nova_1807.06501","nova","scattering",0.)
nova_scattering.descr="""
Limit produced by 1807.06501 based on NOvA collaboration neutrino-electron scattering
in 1710.03428, recasted for light dark matter scattering for 2.97 * 10^20  PoT.
This is not a projection, limit for 58 events, although the authors of 1807.06501 merely
call for a full analysis by the collaboration

Extracted from Figure 2, with alpha_D = 0.05, the data is given as y = eps^2 * alpha_D *(mx/ mv)^4= 0.00062 eps^2
"""
nova_scattering.ref="inspirehep.net/record/1682772"
# print("nova",nova_scattering.mx_ini,nova_scattering.lim_ini)
nova_scattering.UpdateLimIni(nova_scattering.mx_ini,np.sqrt(nova_scattering.lim_ini)/np.sqrt(0.00062)/np.power(2,1/4.) )
UpdateLimitList(nova_scattering)
# print("nova2",nova_scattering.mx_ini,nova_scattering.lim_ini)

#########################  Long-lived particles #########################
if verbose: print("Long-lived dark sector ...")

############ Based on: Kling, 2018 MATHUSLA

mathusla_decay=Limit("mathusla_1810.01879","mathusla","decay",0.1)
mathusla_decay.descr="""
Projective limits for MATHUSLA from 1810.01879, with 3 ab-1 of luminosity from HL-LHC

Energy cut > 0.6 GeV
"""
mathusla_decay.ref="inspirehep.net/record/1696950"
mathusla_decay.combthr=1.2
UpdateLimitList(mathusla_decay)

############ Based on: Kling, 2018 FASER

faser_decay=Limit("faser_1810.01879","faser","decay",0.1)
faser_decay.descr="""
Projective limits for FASER from 1810.01879,  with 3 ab-1 of luminosity from HL-LHC

Energy cut > 100 GeV
"""
faser_decay.ref="inspirehep.net/record/1696950"
faser_decay.combthr=1.1
UpdateLimitList(faser_decay)

############ Based on: 2018 SeaQuest phase 1 and 2

seaquest_phase2_decay=Limit("seaquest_phase2_1804.00661","seaquest_phase2","decay",0.1)
seaquest_phase2_decay.descr="""
Projective limits 1804.00661 Figure 12, 10 events limits with 5m−6m fiducial decay region Phase 2 with 10^20 PoT

"""
seaquest_phase2_decay.ref="inspirehep.net/record/1665691"
UpdateLimitList(seaquest_phase2_decay)

seaquest_phase1_decay=Limit("","seaquest_phase1","decay",0.1)
seaquest_phase1_decay.descr="""
Projective limit, 1804.00661, rescaled from Figure 12 given the reduced number of PoT: 1.44×10^18 PoT
, 10 events limits with 5m−6m fiducial decay regions
"""
seaquest_phase1_decay.ref="inspirehep.net/record/1665691"
seaquest_phase1_decay.UpdateLimIni(seaquest_phase1_decay.mx_ini,np.sqrt(seaquest_phase1_decay.lim_ini)/np.power(0.014,1/4.))
UpdateLimitList(seaquest_phase1_decay)

############ Based on: Izaguiire, 2017 LSND (on-shell pi0 only?)

lsnd_decay_2=Limit("lsnd_1703.06881","lsnd","decay",0.1)
lsnd_decay_2.descr="""
Limits for LSND collaboration, based on  original neutrino scattering search from hep-ex/0104049,
recasted for iDM in 1703.06881 Fig 6, with splitting 10%
"""
lsnd_decay_2.ref="inspirehep.net/record/1682906"
UpdateLimitList(lsnd_decay_2,"lsnd_decay_2")

############  LSND ---- based Darme 2018

lsnd_decay=Limit("lsnd_1807.10314","lsnd","decay",0.15)
lsnd_decay.descr="""
Limits for LSND collaboration, based on  original neutrino scattering search from hep-ex/0104049,
recasted for iDM in 1807.10314 Fig5a, with splitting 15%
"""
lsnd_decay.ref="inspirehep.net/record/1684267"
UpdateLimitList(lsnd_decay)

############  CHARM ---- based Tsai 2019

charm_decay=Limit("charm_1908.07525","charm","decay",0.1)
charm_decay.descr="""
Limits for CHARM collaboration, based on  original neutrino scattering search from Phys.Lett. 128B (1983) 361,
 recasted for iDM in 1908.07525 Fig1.c

Ecut > 3 GeV
"""
charm_decay.ref="inspirehep.net/record/1682906"
UpdateLimitList(charm_decay)

############  SHIP decay prediction

ship_decay=Limit("","ship","decay",0.1)
ship_decay.descr="""
Naive limit prediction for SHIP, based on the routine for production and decay implemented in this code.
Does not include the experimental efficiency. The output is the 10 events line.
"""
ship_decay.ref=""
ship_decay.model="EFT"
UpdateLimitList(ship_decay)

ship_heavymesondecay=Limit("","ship","heavymesondecay",0.1)
ship_heavymesondecay.descr="""
Naive limit prediction for SHIP, based on the routine for production and decay implemented in this code.
Does not include the experimental efficiency. The output is the 10 events line.
Production is generated by heavy meson decay
"""
ship_heavymesondecay.ref=""
ship_heavymesondecay.model="EFT" # The generated limits are already in the EFT
UpdateLimitList(ship_heavymesondecay)

#########################  Missing energy searches  #########################
if verbose: print("Mono-X searches ...")

############  BaBAr  # Based on Essig 2013

babar_monogam=Limit("babar_1309.5084","babar","monogam",0.0)
babar_monogam.descr="""
Limits for BaBar collaboration, based on the recasting from 1309.5084 of the upsilon decay
mono-photon search from 0808.0017.

Notice that the absence of "bump" in the reconstructed dark sector invariant mass distribution
significantly weakens the reach. Better control of the background could lead to improvement
"""
babar_monogam.ref="inspirehep.net/record/1254859,inspirehep.net/record/792059"
UpdateLimitList(babar_monogam)

############  Belle II  # Based on Essig 2013

belle2_monogam=Limit("belle2_1309.5084","belle2","monogam",0.0)
belle2_monogam.descr="""
Limits for Belle II collaboration, based on the proejction from 1309.5084

Notice that the absence of "bump" in the reconstructed dark sector invariant mass distribution
significantly weakens the reach. Better control of the background could lead to improvement
"""
belle2_monogam.ref="inspirehep.net/record/1254859"
UpdateLimitList(belle2_monogam)

#########################  MonoJet searches at ATLAS 35.9 fb-1  #########################
data =  np.loadtxt('Data/LimData/atlas_1711.03301.dat')
xi_ATLAS = np.abs(data[:,0])
Lim_ATLAS_Down = data[:,1]
Lim_ATLAS_Up= data[:,2]

atlas_monoXlow=Limit("","atlas","monojet_down",0.0)
atlas_monoXlow.descr="""
Limits for ATLAS collaboration, 1711.03301 based on the upper recasted limit from 1807.03817
and our own recast.
"""
atlas_monoXlow.ref="inspirehep.net/record/1635274"
atlas_monoXlow.UpdateLimIni(xi_ATLAS,Lim_ATLAS_Down)
UpdateLimitList(atlas_monoXlow)

atlas_monoXhigh=Limit("","atlas","monojet_up",0.0)
atlas_monoXhigh.descr="""
Limits for ATLAS collaboration, 1711.03301 based on the upper recasted limit from 1807.03817
and our own recast.
"""
atlas_monoXhigh.ref="inspirehep.net/record/1635274"
atlas_monoXhigh.UpdateLimIni(xi_ATLAS,Lim_ATLAS_Up)
UpdateLimitList(atlas_monoXhigh)

#########################  MonoPhoton searches at DELPHI from 1103.0240  #########################

xi_LimLEP_V,LimLEP_V  = uf.LoadData('LimData/lep_1103.0240_V.dat',"lin")
xi_LimLEP_AV,LimLEP_AV  = uf.LoadData('LimData/lep_1103.0240_AV.dat',"lin")

lep_monoX=Limit("lep_1103.0240","lep","monogam",0.0)
lep_monoX.descr="""
Limits for LEP collaboration, based on the recasting from 1103.0240 for the upper limit.
The lower limit is given by the breakdown of the EFT at the LEP CoM energy.
"""
lep_monoX.ref="inspirehep.net/record/890992"
lep_monoX.UpdateLimIni(xi_LimLEP_V,LimLEP_V)
lep_monoX.UpdateLimIni(xi_LimLEP_AV,LimLEP_AV,"AV")
UpdateLimitList(lep_monoX)

######################### SN1987 cooling rate constraints #####
if verbose: print("SN1987 cooling ...")

sn1987low_cooling=Limit("sn1987low","sn1987","low_cooling",0.0,"lin")
sn1987low_cooling.descr="""
Lower limit from SN1987 cooling constraints, recasted from ***
"""
sn1987low_cooling.ref="inspirehep.net/record/1682906"
newmx=np.power(10.,sn1987low_cooling.mx_ini/1.)
newlim=np.power(10.,sn1987low_cooling.lim_ini/1.)
sn1987low_cooling.UpdateLimIni(newmx,newlim,"V")
sn1987low_cooling.UpdateLimIni(newmx,newlim,"AV")
UpdateLimitList(sn1987low_cooling)
# print(sn1987low_cooling.mx_ini,sn1987low_cooling.lim_ini)

sn1987high_cooling=Limit("sn1987high","sn1987","high_cooling",0.0,"lin")
sn1987high_cooling.descr="""
Upper limit from SN1987 cooling constraints, recasted from ***

In the case of AV coupling, we use instead the lower limit on the pi0->invisible BR
derived in *** from sn1987 cooling
"""
sn1987high_cooling.ref="inspirehep.net/record/1682906"
xi_LimSN_AV_up,LamlimSN_AV_up= uf.LoadData('LimData/sn1987_pi0decay.dat',"log")
sn1987high_cooling.UpdateLimIni(np.power(10.,sn1987high_cooling.mx_ini/1.),np.power(10.,sn1987high_cooling.lim_ini/1.))
sn1987high_cooling.UpdateLimIni(xi_LimSN_AV_up,LamlimSN_AV_up,"AV")
UpdateLimitList(sn1987high_cooling)

# print(sn1987high_cooling.mx_ini,sn1987high_cooling.lim_ini)

#########################  Invisible decay of pi0  #################

if verbose: print("Invisible meson decay ...")

na62_invisibledecayPi0=Limit("na62_talkKaon2019","na62","invisibledecayPi0",0.0)
na62_invisibledecayPi0.descr="""
Limits for invisible decay branching ratio of pi0 meson as constrained by the NA62 collaboration

Currently value only from https://indico.cern.ch/event/769729/sessions/318725/ Ruggiero's talk
BR < 4.4 10^-9 at 90% CL
"""
na62_invisibledecayPi0.ref="indico.cern.ch/event/769729/contributions/3510938/attachments/1905346/3146619/kaon2019_ruggiero_final.pdf"
UpdateLimitList(na62_invisibledecayPi0)

#########################  Invisible decay of heavy mesons  #################
if verbose: print("Invisible heavy meson decay ...")

babar_invisibledecayBmtoKm=Limit("babar_invisibledecayBmtoKm","babar","invisibledecayBmtoKm",0.0)
belle2_invisibledecayBmtoKm=Limit("belle2_invisibledecayBmtoKm","belle2","invisibledecayBmtoKm",0.0)
belle_invisibledecayB0toPi0=Limit("belle_invisibledecayB0toPi0","belle","invisibledecayB0toPi0",0.0)
belle_invisibledecayB0toK0=Limit("belle_invisibledecayB0toK0","belle","invisibledecayB0toK0",0.0)
babar_invisibledecayBmtoPim=Limit("babar_invisibledecayBmtoPim","babar","invisibledecayBmtoPim",0.0)
e391a_invisibledecayKL0toPi0=Limit("e391a_invisibledecayKL0toPi0","e391a","invisibledecayKL0toPi0",0.0)
na62_invisibledecayKL0toPi0=Limit("na62_invisibledecayKL0toPi0","na62","invisibledecayKL0toPi0",0.0)
e949_invisibledecayKptoPip=Limit("e949_invisibledecayKptoPip","e949","invisibledecayKptoPip",0.0)
na62_invisibledecayKptoPip=Limit("na62_invisibledecayKptoPip","na62","invisibledecayKptoPip",0.0)

bes_invisibledecayJPsi=Limit("bes_invisibledecayJPsi","bes","invisibledecayJPsi",0.0)
bes_invisibledecayJPsi.descr="Limit BR JPSi-> inv < 7.2 * 10^-4 from BES collaboration  0710.0039"

babar_invisibledecayUpsilon=Limit("babar_invisibledecayUpsilon","babar","invisibledecayUpsilon",0.0)
babar_invisibledecayUpsilon.descr="Limit BR Upsilon-> inv < 3.0 * 10^-4 from BABAR collaboration  0908.2840"

UpdateLimitList(babar_invisibledecayBmtoKm);UpdateLimitList(belle2_invisibledecayBmtoKm);UpdateLimitList(belle_invisibledecayB0toPi0);
UpdateLimitList(belle_invisibledecayB0toK0);UpdateLimitList(babar_invisibledecayBmtoPim);
UpdateLimitList(e391a_invisibledecayKL0toPi0);UpdateLimitList(na62_invisibledecayKL0toPi0);
UpdateLimitList(e949_invisibledecayKptoPip);UpdateLimitList(na62_invisibledecayKptoPip)
UpdateLimitList(bes_invisibledecayJPsi);UpdateLimitList(babar_invisibledecayUpsilon)

######################### CR limits at T2K ################# Currently not used
# if verbose: print("Decay from CR production limits ...")
#
# t2k_cosmicrays=Limit("t2k_cosmicrays","t2k","cosmicrays",0.0)
# t2k_cosmicrays.descr="""
# Projections from T2K limits based on cosmic ray production of dark sector states
# limits for 500 number of signal events per year, neutrino background for superK is ~150
#
# """
# UpdateLimitList(t2k_cosmicrays)


import numpy as np
from astropy.cosmology import FlatwCDM
from astropy.cosmology import w0waCDM
from astropy.cosmology import FlatLambdaCDM
from astropy.cosmology import LambdaCDM
import matplotlib.pyplot as plt
from astropy import cosmology as cosmo
import wmom as wmom
import matplotlib
import linef
import matplotlib
import plotsetup
from matplotlib import gridspec
import matplotlib.ticker as ticker
plotsetup.halfpaperfig()

gs1 = gridspec.GridSpec(1,4)
gs1.update(bottom=0.6, top=0.95, hspace=0.0)
ax1= plt.subplot(gs1[0])
ax2= plt.subplot(gs1[1])
ax3= plt.subplot(gs1[2])
ax4= plt.subplot(gs1[3])

gs1 = gridspec.GridSpec(1,4)
gs1.update(bottom=0.34, top=0.57, hspace=0.0)
ax5= plt.subplot(gs1[0])
ax6= plt.subplot(gs1[1])
ax7= plt.subplot(gs1[2])
ax8= plt.subplot(gs1[3])

gs1 = gridspec.GridSpec(1,4)
gs1.update(bottom=0.1, top=0.31, hspace=0.0)
ax9= plt.subplot(gs1[0])
ax10= plt.subplot(gs1[1])
ax11= plt.subplot(gs1[2])
ax12= plt.subplot(gs1[3])

#ax13= plt.subplot(gs1[1,0])
#ax14= plt.subplot(gs1[1,1])
#ax15= plt.subplot(gs1[1,2])
#ax16= plt.subplot(gs1[1,3])

ax=[ax1,ax2,ax3,ax4,ax5,ax6,ax7,ax8,ax9,ax10,ax11,ax12]
ran=[0.04,0.04,0.04,0.04,0.02,0.02,0.02,0.02,0.02,0.02,0.02,0.02]
headn=linef.linef('../DATA/SALT2mu_SNLS+SDSS+LOWZ+PS1_Scolnic2+HST/DS17/SALT2mu_FITOPT'+str('000')+'_MUOPT'+str('000')+'.M0DIF','VARN')
z1,mures, murese = np.loadtxt('../DATA/SALT2mu_SNLS+SDSS+LOWZ+PS1_Scolnic2+HST/DS17/SALT2mu_FITOPT'+str('000')+'_MUOPT'+str('000')+'.M0DIF', usecols=(4,5,6), unpack=True, dtype='string', skiprows=headn+1)
z1 = z1.astype(float)
mures1 = mures.astype(float)
murese1 = murese.astype(float)

#FITOPT: color #GLOBAL GROUP 1
#FITOPT: massstep #GLOBAL GROUP 1
#FITOPT: massevol #GLOBAL GROUP 1
#FITOPT: colorevol #GLOBAL GROUP 1

#p1=plt.errorbar(z1, mures, yerr=murese, fmt='ko', ecolor='black', color='black',label='D15 Data')
#xp=[0,0,0,0,1,1,1,1,2,2,2,2,3,3,3,3]
#yp=[0,1,2,3,0,1,2,3,0,1,2,3,0,1,2,3]
#syslabel=['SALT2 Cal','Beta Evol.','Intr Scat.','Selection',
# 'MW Ext.','Mass Split','Mass Evol.','Pec. Vel.',
#          'HST Cal', 'SDSS g','PS1 r','SNLS i','Supercal', 'SDSS r','PS1 i','SNLS z']
#r'$\beta$'+' Evol. & '
#$\gamma$'+' Evol. & '
#$m_{\rm step}$'+'Shift

syslabel=['SALT2 Cal',r'$\mathbf{\beta}$'+' Evol. ','Intr Scat.','Selection',
 'MW Ext.',r'$\mathbf{ m_{\rm step}}$'+' Shift',r'$\mathbf{\gamma}$'+' Evol. ','Pec. Vel.',
          'HST Cal', 'Supercal','PS1 r','SNLS i'];#,'Supercal', 'SDSS r','PS1 i','SNLS z']

#sysnum1=['001','012','000','000','011','000','000','013','012','015','037','046','084','016','038','047']
#sysnum2=['000','000','001','002','000','006','011','000','000','000','000','000','000','000','000','000']
#weights=[1,1,0.5,1,
#         1,1,1,1,
#         1,.5,.3,.5,.5,.5,.5,.5]
sysnum1=['001','012','000','000','011','000','000','013','012','084','037','046']
sysnum2=['000','000','001','002','000','006','011','000','000','000','000','000']
weights=[1,0.5,0.5,1,
                  1,1,1,1,
                  1,.5,.3,.5]

#sysnum=[1,84,85,86,
# 11,83,13,11,
#        28,15,37,50,29,16,38,51]
#sysnum=[1,11,12,86,83,84,85,15,28,37,51]
#syslabel=['SALT2 Cal','MW Ext','HST Cal','Int. Scat','Mass Split','Mass Evol','Color Evol','SDSS g','CSP B','PS1 r', 'SNLS z']
print len(sysnum1)
#stop
for x in range(0,len(sysnum1)):
    z2,mures2, murese2 = np.loadtxt('../DATA/SALT2mu_SNLS+SDSS+LOWZ+PS1_Scolnic2+HST/DS17/SALT2mu_FITOPT'+str(sysnum1[x])+'_MUOPT'+str(sysnum2[x])+'.M0DIF', usecols=(4,5,6), unpack=True, dtype='string', skiprows=headn+1)
    z2 = z2.astype(float)
    mures2 = mures2.astype(float)
    murese2 = murese2.astype(float)

    ax[x].errorbar(z1, (mures1-mures2)*weights[x], yerr=murese2*0, fmt='go', ecolor='g', alpha=.95)
    ax[x].set_ylim(-ran[x]*.99,ran[x]*.99)
    ax[x].set_xlim(0.005,1.7)
    ax[x].set_xscale('log')
    line, = ax[x].plot(np.arange(0,2,.1), np.arange(0,2,.1)*0.0, lw=2)
    if (x<8): ax[x].text(.01,-.88*ran[x], syslabel[x],fontdict={'fontsize':10}, color='black')
    if (x>7): ax[x].text(.01,-.88*ran[x], syslabel[x],fontdict={'fontsize':10}, color='black')

    ax[x].set_xticks([0.01,0.1,1.0])
    if (x<4): ax[x].yaxis.set_major_locator(ticker.MultipleLocator(0.02))
    if (x>3): ax[x].yaxis.set_major_locator(ticker.MultipleLocator(0.01))
    if (x>7): ax[x].yaxis.set_major_locator(ticker.MultipleLocator(0.01))

    if ((x!=0)&(x!=4)&(x!=8)&(x!=12)):
                     ax[x].set_yticklabels(['','','','','','',''])
    if (x<9):
         ax[x].set_xticklabels(['','','','',''])
    if (x>7):
        ax[x].set_xlabel('z',labelpad=-1)
        ax[x].set_xticklabels([' 0.01','0.1','1.0  ',''])
        #ax[x].set_xticklabels(['0.0','','0.5','','1.0'])
    #if ((x==8)|(x==12)):
    #    ax[x].set_yticklabels(['','-0.01','0.0','0.01',''])
    ax[x].tick_params('both', length=2, width=2, which='major')
    ax[x].tick_params('both', length=0, width=2, which='minor')

#plt.text(.2,-.13,r'$Omega_M=0.30,\Omega_{\Lambda}=.70,w=-1.0$',fontdict={'fontsize':20}, color='black')
#plt.text(.2,.1,r'$\Omega_M=0.25,\Omega_{\Lambda}=.75,w=-1.0$',fontdict={'fontsize':20}, color='orange')
#plt.text(.2,.13,r'$\Omega_M=0.30,\Omega_{\Lambda}=.70,w=-1.1$',fontdict={'fontsize':20}, color='red')
#plt.text(.2,.16,r'$\Omega_M=0.30,\Omega_{\Lambda}=.70,w=-1.0, w_a=-0.1$',fontdict={'fontsize':20}, color='blue')

#ax[1,0].text(-0.3,0.05, 'd(m-M) [Mag]',fontdict={'fontsize':18}, color='black',rotation=90)

ax[4].set_ylabel(r'$\Delta \mu_{\rm Sys}  - \Delta \mu_{\rm Baseline} $'+' [mag]',labelpad=-1)

#plt.ylabel('d(m-M) [Mag]', size=24)

plt.show()
plt.savefig('sysgrid.png')
#line, = plt.plot(range(1,150)/100.0,5.0*np.log(c1/c2), lw=2,color='red')
#plt.show()
#plt.savefig('wdemo1.png')


#!/usr/bin/env python
# -*- coding: utf-8 -*-
r"""Provide the wrench limits constraint.

This provides bounds/limits on the wrenches:

.. math::  F_{lb} \leq F \leq F_{ub}

where :math:`(F_{lb}, F_{ub})` are the lower and upper bound on the wrenches, and
:math:`F` is the wrench vector being optimized.

This formulation can be rewritten as the inequality constraint math:`lb \leq x \leq ub` in QP, with
:math:`lb = F_{lb}`, :math:`ub = F_{ub}`, and :math:`x = F`. This can also be rewritten as :math:`Gx \leq h`,
with :math:`G = [-I, I]^\top` and :math:`h = [-F_{lb}^\top, F_{ub}^\top]^\top` where :math:`I` is the square
identity matrix.

The implementation of this class is inspired by [1] (which is licensed under the LGPLv2).

References:
    - [1] "OpenSoT: A whole-body control library for the compliant humanoid robot COMAN", Rocchi et al., 2015
"""

import numpy as np

from pyrobolearn.priorities.constraints.constraint import BoundConstraint, ForceConstraint

__author__ = "Brian Delhaisse"
__copyright__ = "Copyright 2019, PyRoboLearn"
__credits__ = ["Enrico Mingo Hoffman (C++)", "Brian Delhaisse (Python + doc)"]
__license__ = "GNU GPLv3"
__version__ = "1.0.0"
__maintainer__ = "Brian Delhaisse"
__email__ = "briandelhaisse@gmail.com"
__status__ = "Development"

class WrenchLimitsConstraint(BoundConstraint, ForceConstraint):
    r"""Wrench Limits constraint.

     This provides bounds/limits on the wrenches:

    .. math::  F_{lb} \leq F \leq F_{ub}

    where :math:`(F_{lb}, F_{ub})` are the lower and upper bound on the wrenches, and
    :math:`F` is the wrench vector being optimized.

    This formulation can be rewritten as the inequality constraint math:`lb \leq x \leq ub` in QP, with
    :math:`lb = F_{lb}`, :math:`ub = F_{ub}`, and :math:`x = F`. This can also be rewritten as :math:`Gx \leq h`,
    with :math:`G = [-I, I]^\top` and :math:`h = [-F_{lb}^\top, F_{ub}^\top]^\top` where :math:`I` is the square
    identity matrix.

    The implementation of this class is inspired by [1] (which is licensed under the LGPLv2).

    References:
        - [1] "OpenSoT: A whole-body control library for the compliant humanoid robot COMAN", Rocchi et al., 2015
    """

    def __init__(self, model, bounds):
        r"""
        Initialize the constraint.

        Args:
            model (ModelInterface): model interface.
            bounds (tuple[2 * np.array[float[M]]], np.array[float[M]]): wrench limits, where `M` is 3 (vector of
              forces) or 6 (vector of forces and torques). If tuple, it is the lower and upper bounds on the wrenches.
              If np.array, then the lower and upper bound will be set to (-bounds, bounds).
        """
        super(WrenchLimitsConstraint, self).__init__(model)

        # set variables
        self.bounds = bounds

        # first update
        self.update()

    ##############
    # Properties #
    ##############

    @property
    def bounds(self):
        """Get the wrench bounds."""
        return self._bounds

    @bounds.setter
    def bounds(self, bounds):
        """Set the wrench bounds."""
        if bounds is not None:
            if isinstance(bounds, tuple):
                if len(bounds) != 2:
                    raise ValueError("Expecting the bounds to be a tuple of length 2, but got a length of "
                                     "{}".format(len(bounds)))
                for i, bound in enumerate(bounds):
                    if isinstance(bound, (int, float)):
                        bound = np.ones(3) * bound
                        bounds[i] = bound
                    elif isinstance(bound, np.ndarray) and len(bound) != 3:
                        raise ValueError("Expecting the given bound to be of length 3, but instead got a length of "
                                         "{}".format(len(bound)))
                    else:
                        raise TypeError("Expecting the given bound to be a np.array, but got instead: "
                                        "{}".format(type(bound)))
            elif isinstance(bounds, np.ndarray):
                bounds = bounds.reshape(-1)
                if len(bounds) == 3:
                    bounds = (-bounds, bounds)
                elif len(bounds) == 6:
                    bounds = (bounds[:3], bounds[3:])
                elif len(bounds) == 12:
                    bounds = (bounds[:6], bounds[6:])
                else:
                    raise ValueError("Expecting the given bounds to be of length 3 or 6 but got instead a length of: "
                                     "{}".format(len(bounds)))
            elif isinstance(bounds, (int, float)):
                bounds = (-np.ones(3) * bounds, np.ones(3) * bounds)
            else:
                raise TypeError("Expecting the given bounds to be a tuple of np.array, a np.array, or None, but "
                                "instead got: {}".format(type(bounds)))
        self._bounds = bounds

    ###########
    # Methods #
    ###########

    def _update(self):
        """Update the lower and upper bounds."""
        self._b_lower_bound, self._b_upper_bound = self._bounds


__docformat__='reStructuredText'
__author__ = 'Anand Patil, anand.prabhakar.patil@gmail.com'
__all__ = ['extend_children', 'extend_parents', 'ParentDict', 'Stochastic', 'Deterministic', 'Potential']

from copy import copy
from numpy import array, ndarray, reshape, Inf, asarray, dot, sum, float, isnan, size, NaN, asanyarray
import numpy as np
from Node import Node, ZeroProbability, Variable, PotentialBase, StochasticBase, DeterministicBase
import Container
from Container import DictContainer, ContainerBase, file_items, ArrayContainer
import sys
import pdb

d_neg_inf = float(-1.7976931348623157e+308)

# from PyrexLazyFunction import LazyFunction
from LazyFunction import LazyFunction, Counter

def extend_children(children):
    """
    extend_children(children)

    Returns a set containing
    nearest conditionally stochastic (Stochastic, not Deterministic) descendants.
    """
    new_children = copy(children)
    need_recursion = False
    dtrm_children = set()

    for child in children:
        if isinstance(child,Deterministic):
            new_children |= child.children
            dtrm_children.add(child)
            need_recursion = True

    new_children -= dtrm_children

    if need_recursion:
        new_children = extend_children(new_children)

    return new_children

def extend_parents(parents):
    """
    extend_parents(parents)

    Returns a set containing
    nearest conditionally stochastic (Stochastic, not Deterministic) ancestors.
    """
    new_parents = set()

    for parent in parents:

        new_parents.add(parent)

        if isinstance(parent, DeterministicBase):
            new_parents.remove(parent)
            new_parents |= parent.extended_parents

        elif isinstance(parent, ContainerBase):
            for contained_parent in parent.stochastics:
                new_parents.add(contained_parent)
            for contained_parent in parent.deterministics:
                new_parents |= contained_parent.extended_parents

    return new_parents

class ParentDict(DictContainer):
    """
    A special subclass of DictContainer which makes it safe to change
    varibales' parents. When __setitem__ is called, a ParentDict instance
    removes its owner from the old parent's children set (if appropriate)
    and adds its owner to the new parent's children set. It then asks
    its owner to generate a new LazyFunction instance using its new
    parents.

    Also manages the extended_parents attribute of owner.

    NB: StepMethod and Model are expecting variables'
    children to be static. If you want to change indedependence structure
    over the course of an MCMC loop, please do so with indicator variables.

    :SeeAlso: DictContainer
    """
    def __init__(self, regular_dict, owner):
        DictContainer.__init__(self, dict(regular_dict))
        self.owner = owner
        self.owner.extended_parents = extend_parents(self.variables)
        if isinstance(self.owner, StochasticBase) or isinstance(self.owner, PotentialBase):
            self.has_logp = True
        else:
            self.has_logp = False

    def detach_parents(self):
        for parent in self.itervalues():
            if isinstance(parent, Variable):
                parent.children.discard(self.owner)
            elif isinstance(parent, ContainerBase):
                for variable in parent.variables:
                    variable.chidren.discard(self.owner)

        if self.has_logp:
            self.detach_extended_parents()

    def detach_extended_parents(self):
        for e_parent in self.owner.extended_parents:
            if isinstance(e_parent, StochasticBase):
                e_parent.extended_children.discard(self.owner)

    def attach_parents(self):
        for parent in self.itervalues():
            if isinstance(parent, Variable):
                parent.children.add(self.owner)
            elif isinstance(parent, ContainerBase):
                for variable in parent.variables:
                    variable.children.add(self.owner)

        if self.has_logp:
            self.attach_extended_parents()

    def attach_extended_parents(self):
        for e_parent in self.owner.extended_parents:
            if isinstance(e_parent, StochasticBase):
                e_parent.extended_children.add(self.owner)

    def __setitem__(self, key, new_parent):
        old_parent = self[key]

        # Possibly remove owner from old parent's children set.
        if isinstance(old_parent, Variable) or isinstance(old_parent, ContainerBase):

            # Tell all extended parents to forget about owner
            if self.has_logp:
                self.detach_extended_parents()

            self.val_keys.remove(key)
            self.nonval_keys.append(key)

            if isinstance(old_parent, Variable):
                # See if owner only claims the old parent via this key.
                if sum([parent is old_parent for parent in self.itervalues()]) == 1:
                    old_parent.children.remove(self.owner)

            if isinstance(old_parent, ContainerBase):
                for variable in old_parent.variables:
                    if sum([parent is variable for parent in self.itervalues()]) == 1:
                        variable.children.remove(self.owner)

        # If the new parent is a variable, add owner to its children set.
        if isinstance(new_parent, Variable) or isinstance(new_parent, ContainerBase):

            self.val_keys.append(key)
            self.nonval_keys.remove(key)

            if isinstance(new_parent, Variable):
                new_parent.children.add(self.owner)

            elif isinstance(new_parent, ContainerBase):
                for variable in new_parent.variables:
                    new_parent.children.add(self.owner)

        # Totally recompute extended parents
        self.owner.extended_parents = extend_parents(self.variables)
        if self.has_logp:
            self.attach_extended_parents()

        dict.__setitem__(self, key, new_parent)

        file_items(self, self)

        # Tell my owner it needs a new lazy function.
        self.owner.gen_lazy_function()

class Potential(PotentialBase):
    """
    Not a variable; just an arbitrary log-probability term to multiply into the
    joint distribution. Useful for expressing models that aren't directed, such as
    Markov random fields.

    Decorator instantiation:

    @potential(trace = True)
    def A(x = B, y = C):
        return -.5 * (x-y)**2 / 3.

    Direct instantiation:

    :Parameters:

        -logp: function
              The function that computes the potential's value from the values
              of its parents.

        -doc: string
              The docstring for this potential.

        -name: string
              The name of this potential.

        -parents: dictionary
              A dictionary containing the parents of this potential.

        -cache_depth (optional): integer
              An integer indicating how many of this potential's value computations
              should be 'memoized'.

        - plot (optional) : boolean
            A flag indicating whether this variable is to be plotted.

        - verbose (optional) : integer
              Level of output verbosity: 0=none, 1=low, 2=medium, 3=high

    Externally-accessible attribute:

        -logp: float
              Returns the potential's log-probability given its parents' values. Skips
              computation if possible.

    No methods.

    :SeeAlso: Stochastic, Node, LazyFunction, stoch, dtrm, data, Model, Container
    """
    def __init__(self, logp,  doc, name, parents, cache_depth=2, plot=None, verbose=None):

        self.ParentDict = ParentDict

        # This function gets used to evaluate self's value.
        self._logp_fun = logp

        self.errmsg = "Potential %s forbids its parents' current values"%name

        Node.__init__(  self,
                        doc=doc,
                        name=name,
                        parents=parents,
                        cache_depth = cache_depth,
                        verbose=verbose)

        self._plot = plot

        # self._logp.force_compute()

        # Check initial value
        if not isinstance(self.logp, float):
            raise ValueError, "Potential " + self.__name__ + "'s initial log-probability is %s, should be a float." %self.logp.__repr__()

    def gen_lazy_function(self):

        self._logp = LazyFunction(fun = self._logp_fun,
                                    arguments = self.parents,
                                    ultimate_args = self.extended_parents,
                                    cache_depth = self._cache_depth)
        self._logp.force_compute()

    def get_logp(self):
        if self.verbose > 1:
            print '\t' + self.__name__ + ': log-probability accessed.'
        logp = self._logp.get()
        if self.verbose > 1:
            print '\t' + self.__name__ + ': Returning log-probability ', logp

        try:
            logp = float(logp)
        except:
            raise TypeError, self.__name__ + ': computed log-probability ' + str(logp) + ' cannot be cast to float'

        if logp != logp:
            raise ValueError, self.__name__ + ': computed log-probability is NaN'

        # Check if the value is smaller than a double precision infinity:
        if logp <= d_neg_inf:
            if self.verbose > 0:
                raise ZeroProbability, self.errmsg + ": %s" %self._parents.value
            else:
                raise ZeroProbability, self.errmsg

        return logp

    def set_logp(self,value):
        raise AttributeError, 'Potential '+self.__name__+'\'s log-probability cannot be set.'

    logp = property(fget = get_logp, fset=set_logp, doc="Self's log-probability value conditional on parents.")

class Deterministic(DeterministicBase):
    """
    A variable whose value is determined by the values of its parents.

    Decorator instantiation:

    @dtrm(trace=True)
    def A(x = B, y = C):
        return sqrt(x ** 2 + y ** 2)

    :Parameters:
      eval : function
        The function that computes the variable's value from the values
        of its parents.
      doc : string
        The docstring for this variable.
      name: string
        The name of this variable.
      parents: dictionary
        A dictionary containing the parents of this variable.
      trace (optional): boolean
        A boolean indicating whether this variable's value
        should be traced (in MCMC).
      cache_depth (optional): integer
        An integer indicating how many of this variable's
        value computations should be 'memoized'.
      plot (optional) : boolean
        A flag indicating whether this variable is to be plotted.
      verbose (optional) : integer
        Level of output verbosity: 0=none, 1=low, 2=medium, 3=high

    :Attributes:
      value : any object
        Returns the variable's value given its parents' values. Skips
        computation if possible.

    :SeeAlso:
      Stochastic, Potential, deterministic, MCMC, Lambda,
      LinearCombination, Index
    """
    def __init__(self, eval,  doc, name, parents, dtype=None, trace=True, cache_depth=2, plot=None, verbose=None):
        self.ParentDict = ParentDict

        # This function gets used to evaluate self's value.
        self._eval_fun = eval

        Variable.__init__(  self,
                        doc=doc,
                        name=name,
                        parents=parents,
                        cache_depth = cache_depth,
                        dtype=dtype,
                        trace=trace,
                        plot=plot,
                        verbose=verbose)

        # self._value.force_compute()

    def gen_lazy_function(self):

        self._value = LazyFunction(fun = self._eval_fun,
                                    arguments = self.parents,
                                    ultimate_args = self.extended_parents,
                                    cache_depth = self._cache_depth)

        self._value.force_compute()

    def get_value(self):
        if self.verbose > 1:
            print '\t' + self.__name__ + ': value accessed.'
        _value = self._value.get()
        if isinstance(_value, ndarray):
            _value.flags['W'] = False
        if self.verbose > 1:
            print '\t' + self.__name__ + ': Returning value ',_value
        return _value

    def set_value(self,value):
        raise AttributeError, 'Deterministic '+self.__name__+'\'s value cannot be set.'

    value = property(fget = get_value, fset=set_value, doc="Self's value computed from current values of parents.")

class Stochastic(StochasticBase):

    """
    A variable whose value is not determined by the values of its parents.

    Decorator instantiation:

    @stoch(trace=True)
    def X(value = 0., mu = B, tau = C):
        return Normal_like(value, mu, tau)

    @stoch(trace=True)
    def X(value=0., mu=B, tau=C):

        def logp(value, mu, tau):
            return Normal_like(value, mu, tau)

        def random(mu, tau):
            return Normal_r(mu, tau)

        rseed = 1.

    Direct instantiation:

    - logp : function
            The function that computes the variable's log-probability from
            its value and the values of its parents.

    - doc : string
            The docstring for this variable.

    - name : string
            The name of this variable.

    - parents: dict
            A dictionary containing the parents of this variable.

    - random (optional) : function
            A function that draws a new value for this
            variable given its parents' values.

    - trace (optional) : boolean
            A boolean indicating whether this variable's value
            should be traced (in MCMC).

    - value (optional) : number or array
            An initial value for this variable

    - dtype (optional) : type
            A type for this variable.

    - rseed (optional) : integer or rseed
            A seed for this variable's rng. Either value or rseed must
            be given.

    - observed (optional) :  boolean
            A flag indicating whether this variable is data; whether
            its value is known.

    - cache_depth (optional) : integer
            An integer indicating how many of this variable's
            log-probability computations should be 'memoized'.

    - plot (optional) : boolean
            A flag indicating whether this variable is to be plotted.

    - verbose (optional) : integer
            Level of output verbosity: 0=none, 1=low, 2=medium, 3=high

    Externally-accessible attribute:

    - value: any class
          Returns this variable's current value.

    - logp: float
          Returns the variable's log-probability given its value and its
          parents' values. Skips computation if possible.

    last_value: any class
          Returns this variable's last value. Useful for rejecting
          Metropolis-Hastings jumps. See touch() and the warning below.

    Externally-accessible methods:

    random():   Draws a new value for this variable from its distribution and
                returns it.

    :SeeAlso: Deterministic, Node, LazyFunction, stoch, dtrm, data, Model, Container
    """

    def __init__(   self,
                    logp,
                    doc,
                    name,
                    parents,
                    random = None,
                    trace=True,
                    value=None,
                    dtype=None,
                    rseed=False,
                    observed=False,
                    cache_depth=2,
                    plot=None,
                    verbose = None,
                    isdata=None,
                    check_logp=True):

        self.counter = Counter()
        self.ParentDict = ParentDict

        # Support legacy 'isdata' for a while
        if isdata is not None:
            print "Deprecation Warning: the 'isdata' flag has been replaced by 'observed'. Please update your model accordingly."
            self.observed = isdata

        # A flag indicating whether self's value has been observed.
        self._observed = observed
        if observed and value is None:
            raise ValueError, 'Stochastic %s must be given an initial value if observed=True.'%name

        # This function will be used to evaluate self's log probability.
        self._logp_fun = logp

        # This function will be used to draw values for self conditional on self's parents.
        self._random = random

        # A seed for self's rng. If provided, the initial value will be drawn. Otherwise it's
        # taken from the constructor.
        self.rseed = rseed

        self.errmsg = "Stochastic %s's value is outside its support,\n or it forbids its parents' current values."%name

        dtype = np.dtype(dtype)

        # Initialize value, either from value provided or from random function.
        try:
            if dtype.kind != 'O' and value is not None:
                self._value = asanyarray(value, dtype=dtype)
                self._value.flags['W']=False
            else:
                self._value = value
        except:
            cls, inst, tb = sys.exc_info()
            new_inst = cls('Stochastic %s: Failed to cast initial value to required dtype.\n\nOriginal error message:\n'%name + inst.message)
            raise cls, new_inst, tb

        # Store the shape of the stochastic value
        self._shape = np.shape(self._value)

        Variable.__init__(  self,
                        doc=doc,
                        name=name,
                        parents=parents,
                        cache_depth=cache_depth,
                        trace=trace,
                        dtype=dtype,
                        plot=plot,
                        verbose=verbose)

        # self._logp.force_compute()

        # Store the shape of the stochastic value
        self._shape = np.shape(self._value)

        if isinstance(self._value, ndarray):
            self._value.flags['W'] = False

        if check_logp:
            # Check initial value
            if not isinstance(self.logp, float):
                raise ValueError, "Stochastic " + self.__name__ + "'s initial log-probability is %s, should be a float." %self.logp.__repr__()

    def gen_lazy_function(self):
        """
        Will be called by Node at instantiation.
        """

        # If value argument to __init__ was None, draw value from random method.
        if self._value is None:

            # Use random function if provided
            if self._random is not None:
                self.value = self._random(**self._parents.value)

            # Otherwise leave initial value at None and warn.
            else:
                raise ValueError, 'Stochastic ' + self.__name__ + "'s value initialized to None; no initial value or random method provided."

        arguments = {}
        arguments.update(self.parents)
        arguments['value'] = self
        arguments = DictContainer(arguments)

        self._logp = LazyFunction(fun = self._logp_fun,
                                    arguments = arguments,
                                    ultimate_args = self.extended_parents | set([self]),
                                    cache_depth = self._cache_depth)
        self._logp.force_compute()

    def get_value(self):
        # Define value attribute
        if self.verbose > 1:
            print '\t' + self.__name__ + ': value accessed.'
        return self._value

    def set_value(self, value, force=False):
        # Record new value and increment counter

        # Value can't be updated if observed=True
        if self.observed and not force:
            raise AttributeError, 'Stochastic '+self.__name__+'\'s value cannot be updated if observed flag is set'

        if self.verbose > 0:
            print '\t' + self.__name__ + ': value set to ', value

        # Save current value as last_value
        # Don't copy because caching depends on the object's reference.
        self.last_value = self._value

        if self.dtype.kind != 'O':
            self._value = asanyarray(value, dtype=self.dtype)
            self._value.flags['W']=False
        else:
            self._value = value

        self.counter.click()

    value = property(fget=get_value, fset=set_value, doc="Self's current value.")

    def shape():
        doc = "The shape of the value of self."
        def fget(self):
            if self.verbose > 1:
                print '\t' + self.__name__ + ': shape accessed.'
            return self._shape
        return locals()
    shape = property(**shape())

    def revert(self):
        """
        Sets self's value to self's last value. Bypasses the data cleaning in
        the set_value method.
        """
        self.counter.unclick()
        self._value = self.last_value

    def get_logp(self):

        if self.verbose > 0:
            print '\t' + self.__name__ + ': logp accessed.'
        logp = self._logp.get()

        try:
            logp = float(logp)
        except:
            raise TypeError, self.__name__ + ': computed log-probability ' + str(logp) + ' cannot be cast to float'

        if logp != logp:
            raise ValueError, self.__name__ + ': computed log-probability is NaN'

        if self.verbose > 0:
            print '\t' + self.__name__ + ': Returning log-probability ', logp

        # Check if the value is smaller than a double precision infinity:
        if logp <= d_neg_inf:
            if self.verbose > 0:
                raise ZeroProbability, self.errmsg + "\nValue: %s\nParents' values:%s" % (self._value, self._parents.value)
            else:
                raise ZeroProbability, self.errmsg

        return logp

    def set_logp(self, new_logp):
        raise AttributeError, 'Stochastic '+self.__name__+'\'s logp attribute cannot be set'

    logp = property(fget = get_logp, fset=set_logp, doc="Log-probability or log-density of self's current value\n given values of parents.")

    # Sample self's value conditional on parents.
    def random(self):
        """
        Draws a new value for a stoch conditional on its parents
        and returns it.

        Raises an error if no 'random' argument was passed to __init__.
        """

        if self._random:
            # Get current values of parents for use as arguments for _random()
            r = self._random(**self.parents.value)
        else:
            raise AttributeError, 'Stochastic '+self.__name__+' does not know how to draw its value, see documentation'

        if self.shape:
            r = np.reshape(r, self.shape)

        # Set Stochastic's value to drawn value
        if not self.observed:
            self.value = r

        return r

    # Shortcut alias to random
    rand = random

    def _get_isdata(self):
        import warnings
        warnings.warn('"isdata" is deprecated, please use "observed" instead.')
        return self._observed
    def _set_isdata(self, isdata):
        raise ValueError, 'Stochastic %s: "observed" flag cannot be changed.'%self.__name__
    isdata = property(_get_isdata, _set_isdata)

    def _get_observed(self):
        return self._observed
    def _set_observed(self, observed):
        raise ValueError, 'Stochastic %s: "observed" flag cannot be changed.'%self.__name__
    observed = property(_get_observed, _set_observed)

    def _get_coparents(self):
        coparents = set()
        for child in self.extended_children:
            coparents |= child.extended_parents
        coparents.add(self)
        return coparents
    coparents = property(_get_coparents, doc="All the variables whose extended children intersect with self's.")

    def _get_moral_neighbors(self):
        moral_neighbors = self.coparents | self.extended_parents | self.extended_children
        for neighbor in copy(moral_neighbors):
            if isinstance(neighbor, PotentialBase):
                moral_neighbors.remove(neighbor)
        return moral_neighbors
    moral_neighbors = property(_get_moral_neighbors, doc="Self's neighbors in the moral graph: self's Markov blanket with self removed.")

    def _get_markov_blanket(self):
        return self.moral_neighbors | set([self])
    markov_blanket = property(_get_markov_blanket, doc="Self's coparents, self's extended parents, self's children and self.")