## Copyright (c) 2019 Richard Offer. All right reserved.
#
# see LICENSE.md for license details
#
## Wrapper around random routines to make it easier to re-implement

import logging

import numpy.random as NPRND

import censere.utils as UTILS

# Life table for the total population: United States, 2006
cdc_2006 = [
    0,
    1,
    2,
    3,
    4,
    5,
    6,
    7,
    8,
    9,
    10,
    11,
    12,
    13,
    14,
    15,
    16,
    17,
    18,
    19,
    20,
    21,
    22,
    23,
    24,
    25,
    26,
    27,
    28,
    29,
    30,
    31,
    32,
    33,
    34,
    35,
    36,
    37,
    38,
    39,
    40,
    41,
    42,
    43,
    44,
    45,
    46,
    47,
    48,
    49,
    50,
    51,
    52,
    53,
    54,
    55,
    56,
    57,
    58,
    59,
    60,
    61,
    62,
    63,
    64,
    65,
    66,
    67,
    68,
    69,
    70,
    71,
    72,
    73,
    74,
    75,
    76,
    77,
    78,
    79,
    80,
    81,
    82,
    83,
    84,
    85,
    86,
    87,
    88,
    89,
    90,
    91,
    92,
    93,
    94,
    95,
    96,
    97,
    98,
    99,
    100,
]

# 4.696378 is the sum of all probablilities
#
cdc_2006_weights = [
    0.006713 / 4.696378 ,
    0.000444 / 4.696378,
    0.000300 / 4.696378,
    0.000216 / 4.696378,
    0.000179 / 4.696378,
    0.000168 / 4.696378,
    0.000156 / 4.696378,
    0.000143 / 4.696378,
    0.000125 / 4.696378,
    0.000103 / 4.696378,
    0.000086 / 4.696378,
    0.000088 / 4.696378,
    0.000125 / 4.696378,
    0.000206 / 4.696378,
    0.000317 / 4.696378,
    0.000438 / 4.696378,
    0.000552 / 4.696378,
    0.000657 / 4.696378,
    0.000747 / 4.696378,
    0.000825 / 4.696378,
    0.000905 / 4.696378,
    0.000983 / 4.696378,
    0.001033 / 4.696378,
    0.001049 / 4.696378,
    0.001038 / 4.696378,
    0.001019 / 4.696378,
    0.001006 / 4.696378,
    0.000998 / 4.696378,
    0.001002 / 4.696378,
    0.001018 / 4.696378,
    0.001042 / 4.696378,
    0.001072 / 4.696378,
    0.001113 / 4.696378,
    0.001156 / 4.696378,
    0.001212 / 4.696378,
    0.001276 / 4.696378,
    0.001355 / 4.696378,
    0.001456 / 4.696378,
    0.001585 / 4.696378,
    0.001739 / 4.696378,
    0.001903 / 4.696378,
    0.002077 / 4.696378,
    0.002268 / 4.696378,
    0.002479 / 4.696378,
    0.002706 / 4.696378,
    0.002943 / 4.696378,
    0.003190 / 4.696378,
    0.003453 / 4.696378,
    0.003741 / 4.696378,
    0.004057 / 4.696378,
    0.004405 / 4.696378,
    0.004778 / 4.696378,
    0.005166 / 4.696378,
    0.005554 / 4.696378,
    0.005939 / 4.696378,
    0.006335 / 4.696378,
    0.006760 / 4.696378,
    0.007234 / 4.696378,
    0.007796 / 4.696378,
    0.008470 / 4.696378,
    0.009282 / 4.696378,
    0.010204 / 4.696378,
    0.011178 / 4.696378,
    0.012118 / 4.696378,
    0.013024 / 4.696378,
    0.013999 / 4.696378,
    0.014995 / 4.696378,
    0.016161 / 4.696378,
    0.017527 / 4.696378,
    0.019109 / 4.696378,
    0.020890 / 4.696378,
    0.022925 / 4.696378,
    0.025280 / 4.696378,
    0.027972 / 4.696378,
    0.030997 / 4.696378,
    0.034386 / 4.696378,
    0.038027 / 4.696378,
    0.042036 / 4.696378,
    0.046447 / 4.696378,
    0.051297 / 4.696378,
    0.056623 / 4.696378,
    0.062465 / 4.696378,
    0.068867 / 4.696378,
    0.075871 / 4.696378,
    0.083524 / 4.696378,
    0.091872 / 4.696378,
    0.100962 / 4.696378,
    0.110842 / 4.696378,
    0.121558 / 4.696378,
    0.133155 / 4.696378,
    0.145675 / 4.696378,
    0.159156 / 4.696378,
    0.173631 / 4.696378,
    0.189127 / 4.696378,
    0.205661 / 4.696378,
    0.223242 / 4.696378,
    0.241869 / 4.696378,
    0.261527 / 4.696378,
    0.282188 / 4.696378,
    0.303810 / 4.696378,
    1.0000 / 4.6963780
]

def seed( seed ):

    return NPRND.seed( seed )

def set_state( st ):

    return NPRND.set_state( st )

def get_state( ):

    return NPRND.get_state( )

# return a string of random bytes
# We use this rather than UUID so that
# we can get the same IDs if the same seed is used
#
# uuid.uuid() uses the os.urandom() and can't be replayed.
#
def id():

    return NPRND.bytes(16).hex()

def random():

    return NPRND.random()

## return a random number between start and stop
#
# Return a random integer N such that a <= N < b.
#
def randrange( start, stop ):

    return NPRND.randint( start, stop )

##
#
# Return a random integer N such that a <= N <= b. Alias for randrange(a, b+1)
#
def randint( start, stop ):

    if start == stop:
        return start

    return randrange( start, stop + 1  )

def gauss( mean, sigma):

    return mean + ( NPRND.randn() * sigma )

def triangle( minimum, peak, maximum):

    return NPRND.triangular( minimum, peak, maximum )

def choice( lst):

    return NPRND.choice( lst )

def choices( lst, weights=None ):

    return NPRND.choice( lst, None, p=[ (i/100.0) for i in weights ] )

def life_expectancy():

    return NPRND.choice( cdc_2006, None, p=cdc_2006_weights )

##
# tables: prefix is only valid for life-expectancy values
#
# Difference between randint and randrange is that randint only returns int
# and handles the case where MIN == MAX, randrange will error in that case
# TODO
#   * handle sexes life expectancy independently
#
def parse_random_value( key, default_value=None, key_in_earth_years=False ):

    val = key.split(":")

    value = default_value

    if val[0] == "cdc":
        value = life_expectancy()

    elif val[0] == "triangular" or val[0] == "triangle" :

        values = [ float(i) for i in val[1].split(",") ]
        value = triangle( values[0], values[1], values[2] )

    elif val[0] == "gauss":
        values = [ float(i) for i in val[1].split(",") ]
        value = gauss( values[0], values[1] ),

    elif val[0] == "randint":

        values = [ int(i) for i in val[1].split(",") ]
        value = randint( values[0], values[1] )

    elif val[0] == "randrange":

        values =  [float(i) for i in val[1].split(",") ]
        value = randrange( values[0], values[1] )

    elif val[0] == "half":

        values = [ float(i) for i in val[1].split(",") ]

        r = random()

        if r < 0.5:
            value = ( values[0] ) + ( values[1] * random() )
        elif r < 0.75:
            value = ( values[0] + ( values[1] ) ) + ( values[1] * random() )
        elif r < 0.875:
            value = ( values[0] + ( values[1] * 2 ) ) + ( values[1] * random() )
        elif r < 0.925:
            value = ( values[0] + ( values[1] * 3 ) ) + ( values[1] * random() )
        elif r < 0.9625:
            value = ( values[0] + ( values[1] * 4 ) ) + ( values[1] * random() )
        elif r < 0.98125:
            value = ( values[0] + ( values[1] * 5 ) ) + ( values[1] * random() )
        elif r < 0.990625:
            value = ( values[0] + ( values[1] * 6 ) ) + ( values[1] * random() )
        elif r < 0.9953125:
            value = ( values[0] + ( values[1] * 7 ) ) + ( values[1] * random() )
        else:
            value = ( values[0] + ( values[1] * 8 ) ) + ( values[1] * random() )

    else:

        logging.fatal( 'Invalid value %s', key )

    if key_in_earth_years:

        return UTILS.years_to_sols( value )

    return value


#!/usr/bin/env python2

import sys
sys.path.append('../fml/')

import os
import numpy as np
import fml

from fml.math.distance import get_l2_distance_arad
from fml.kernels import get_atomic_kernels_arad

PTP = {\
         1  :[1,1] ,2:  [1,8]#Row1

        ,3  :[2,1] ,4:  [2,2]#Row2\
        ,5  :[2,3] ,6:  [2,4] ,7  :[2,5] ,8  :[2,6] ,9  :[2,7] ,10 :[2,8]\

        ,11 :[3,1] ,12: [3,2]#Row3\
        ,13 :[3,3] ,14: [3,4] ,15 :[3,5] ,16 :[3,6] ,17 :[3,7] ,18 :[3,8]\

        ,19 :[4,1] ,20: [4,2]#Row4\
        ,31 :[4,3] ,32: [4,4] ,33 :[4,5] ,34 :[4,6] ,35 :[4,7] ,36 :[4,8]\
        ,21 :[4,9] ,22: [4,10],23 :[4,11],24 :[4,12],25 :[4,13],26 :[4,14],27 :[4,15],28 :[4,16],29 :[4,17],30 :[4,18]\

        ,37 :[5,1] ,38: [5,2]#Row5\
        ,49 :[5,3] ,50: [5,4] ,51 :[5,5] ,52 :[5,6] ,53 :[5,7] ,54 :[5,8]\
        ,39 :[5,9] ,40: [5,10],41 :[5,11],42 :[5,12],43 :[5,13],44 :[5,14],45 :[5,15],46 :[5,16],47 :[5,17],48 :[5,18]\

        ,55 :[6,1] ,56: [6,2]#Row6\
        ,81 :[6,3] ,82: [6,4] ,83 :[6,5] ,84 :[6,6] ,85 :[6,7] ,86 :[6,8]
               ,72: [6,10],73 :[6,11],74 :[6,12],75 :[6,13],76 :[6,14],77 :[6,15],78 :[6,16],79 :[6,17],80 :[6,18]\
        ,57 :[6,19],58: [6,20],59 :[6,21],60 :[6,22],61 :[6,23],62 :[6,24],63 :[6,25],64 :[6,26],65 :[6,27],66 :[6,28],67 :[6,29],68 :[6,30],69 :[6,31],70 :[6,32],71 :[6,33]\

        ,87 :[7,1] ,88: [7,2]#Row7\
        ,113:[7,3] ,114:[7,4] ,115:[7,5] ,116:[7,6] ,117:[7,7] ,118:[7,8]\
               ,104:[7,10],105:[7,11],106:[7,12],107:[7,13],108:[7,14],109:[7,15],110:[7,16],111:[7,17],112:[7,18]\
        ,89 :[7,19],90: [7,20],91 :[7,21],92 :[7,22],93 :[7,23],94 :[7,24],95 :[7,25],96 :[7,26],97 :[7,27],98 :[7,28],99 :[7,29],100:[7,30],101:[7,31],101:[7,32],102:[7,14],103:[7,33]}
#@jit#(nopython = True, nogil = True)

#@jit#(nopython = True, nogil = True)

def condense(X1,X2,Z1,Z2,width = 0.2, cutDist = 6.,RWidth = 1 ,CWidth = 0.5,num_cores = 1, printTime = False,molecular = False):#RWidth =1./sqrt(2),Cwidth = 1./sqrt(2)):#,RWidth =1.,Cwidth = 1.):
    def createZINP(Z,l):
        Z_inp = -1*np.ones((len(Z),l,2)).astype(int)
        for i in range(len(Z)):
            Z_inp[i,:len(Z[i])] = np.asarray([PTP[z] for z in Z[i]])

        return Z_inp

    RWidth = float(RWidth)
    Cwidth = float(CWidth)
    Z1_inp = createZINP(Z1,len(X1[0]))
    Z2_inp = createZINP(Z2,len(X2[0]))

    from time import time

    start = time()
    DistMatrix = run(X1,X2,Z1_inp,Z2_inp,width,cutDist,RWidth,Cwidth,num_cores,molecular = molecular)
    if printTime:
        print 'Time Paralell: ', time() - start
    #print DistMatrix[0,0]
    #print DistMatrix.shape
    return DistMatrix

class KeyboardInterruptError(Exception): pass

def run(X1,X2,Z1,Z2,width,cutDist,RWidth,Cwidth,Ncores = 2,NIters = 20, molecular = True ):
    NIters = min(len(X2),NIters)
    from multiprocessing import Pool
    try:

        p = Pool(Ncores)

        w_list  = np.repeat(width,len(X1))
        cD_list = np.repeat(cutDist,len(X1))
        RW_list = np.repeat(RWidth,len(X1))
        CW_list = np.repeat(Cwidth,len(X1))
        const_list =np.repeat(np.NaN,len(X1))

        w_list_2  = np.repeat(width,len(X2))
        cD_list_2 = np.repeat(cutDist,len(X2))
        RW_list_2 = np.repeat(RWidth,len(X2))
        CW_list_2 = np.repeat(Cwidth,len(X2))

        args1 = [X1,Z1,w_list,cD_list,RW_list,CW_list]
        args2 = [X2,Z2,w_list_2,cD_list_2,RW_list_2,CW_list_2]

        if molecular:
            pass_nonmNoDist
            D1 = np.asarray(p.map(pass_nonmNoDist,zip(*args1)))
            D2 = np.asarray(p.map(pass_nonmNoDist,zip(*args2)))
            D1 = unravelKerMatrix(D1)
            D2 = unravelKerMatrix(D2)

        else:
            D1 = np.asarray(p.map(pass_calcNorm,zip(*args1)))
            D2 = np.asarray(p.map(pass_calcNorm,zip(*args2)))

        i_2 = 0
        x2 = np.repeat(np.asarray([X2[i_2:min(len(X2),i_2 + NIters)]]),len(X1),axis = 0)
        z2 = np.repeat(np.asarray([Z2[i_2:min(len(X2),i_2 + NIters)]]),len(X1),axis = 0)
        d2 = np.repeat(np.asarray([D2[i_2:min(len(X2),i_2 + NIters)]]),len(X1),axis = 0)
        d2 = np.repeat(np.asarray([D2[i_2:min(len(X2),i_2 + NIters)]]),len(X1),axis = 0)
        if molecular:
            args = [X1,x2,Z1,z2,w_list,cD_list,RW_list,CW_list]
            results = np.asarray(p.map(pass_calcProdnoDist,zip(*args)))
        else:
            args = [X1,x2,Z1,z2,D1,d2,w_list,cD_list,RW_list,CW_list,const_list]
            results = np.asarray(p.map(pass_calcProd,zip(*args)))

        for i_2 in range(NIters,len(X2),NIters):
            x2 = np.repeat(np.asarray([X2[i_2:min(len(X2),i_2 + NIters)]]),len(X1),axis = 0)
            z2 = np.repeat(np.asarray([Z2[i_2:min(len(X2),i_2 + NIters)]]),len(X1),axis = 0)
            d2 = np.repeat(np.asarray([D2[i_2:min(len(X2),i_2 + NIters)]]),len(X1),axis = 0)

            if molecular:
                args = [X1,x2,Z1,z2,w_list,cD_list,RW_list,CW_list]
                results = append(results,
                    np.asarray(p.map(pass_calcProdnoDist,zip(*args))),
                    axis = 1)
            else:
                args = [X1,x2,Z1,z2,D1,d2,w_list,cD_list,RW_list,CW_list,const_list]
                results = append(results,
                    np.asarray(p.map(pass_calcProd,zip(*args))),
                    axis = 1)
        p.close()
        p.join()
            #print results.shape

        if molecular:
            results = unravelKerMatrix(results)
            #print results[0:5,0:5]
            #print D1[0:5]
            #print D2[0:5]
            #print results.shape,D1[:,newaxis].shape,D2[newaxis].shape

            D = D1[:,newaxis] +  D2[newaxis]
            results =  D - 2 * results
            #print results[0:5,0:5]

        return results
    except KeyboardInterrupt:
        print 'got ^C'
        p.terminate()
        p.join()
        raise KeyboardInterrupt

    except Exception, inst:
        p.terminate()
        p.join()
        print 'got exception: %r' % (inst,)
        print type(inst)     # the exception instance
        print inst.args      # arguments stored in .args
        print inst           # __str__ allows args to be printed directly
        raise inst

def pass_calcNorm(args):
    try:
        return calcNorm(*args)
    except KeyboardInterrupt:
        raise KeyboardInterruptError()

def calcNorm(X,Z,width,cutDist,RWidth,Cwidth):

    DistMatrix  = np.zeros(len(X))
    for i in  range(len(X)):
        if Z[i,0] == -1:
            break
        dd = _MDist(X[i],X[i],width,cutDist,RWidth,Cwidth)
        DistMatrix[i] = dd

    return DistMatrix

#@jit
def pass_calcProd(args):
    try:
        return calcProd(*args)
    except KeyboardInterrupt:
        raise KeyboardInterruptError()
#@jit#(nopython = True, nogil = True)
def calcProd(X1,X2,Z1,Z2,D1,D2,width,cutDist,RWidth,Cwidth,const):
    #maxRDist = 8*RWidth
    #maxCDist = 8*Cwidth

    DistMatrix  = const * np.ones((len(X2), len(X1), len(X2[0])))
    for i in range(len(X2)):
        for j_1 in  range(len(X1)):

            if Z1[j_1,0] == -1:
                break

            for j_2 in range(len(X2[0])):

                if Z2[i,j_2,0] == -1:
                    break

                RDist = abs(Z1[j_1,0] - Z2[i,j_2,0])
                CDist = abs(Z1[j_1,1] - Z2[i,j_2,1])

                #if RDist < maxRDist and CDist < maxCDist:
                dd = _MDist(X1[j_1],X2[i,j_2],width,cutDist,RWidth,Cwidth)
                dd = dd *_StochDist(RDist,CDist,RWidth,Cwidth)

                dd = D1[j_1] + D2[i,j_2] - 2 * dd

                if dd < 0:
                    assert dd > -1E-13, 'Something is wrong in calcProd: ' + str(dd)
                    DistMatrix[i,j_1,j_2] = 0

                DistMatrix[i,j_1,j_2] = dd

    return DistMatrix

def pass_calcProdnoDist(args):
    try:
        return calcProd_noDist(*args)
    except KeyboardInterrupt:
        raise KeyboardInterruptError()

def calcProd_noDist(X1,X2,Z1,Z2,width,cutDist,RWidth,Cwidth):

    DistMatrix  = np.zeros((len(X2), len(X1), len(X2[0])))
    for i in range(len(X2)):
        for j_1 in  range(len(X1)):

            if Z1[j_1,0] == -1:
                break

            for j_2 in range(len(X2[0])):

                if Z2[i,j_2,0] == -1:
                    break

                RDist = abs(Z1[j_1,0] - Z2[i,j_2,0])
                CDist = abs(Z1[j_1,1] - Z2[i,j_2,1])

                dd = _MDist(X1[j_1],X2[i,j_2],width,cutDist,RWidth,Cwidth)
                dd = dd *_StochDist(RDist,CDist,RWidth,Cwidth)
                DistMatrix[i,j_1,j_2] = dd

    return DistMatrix

def pass_nonmNoDist(args):
    try:
        return nonmNoDist(*args)
    except KeyboardInterrupt:
        raise KeyboardInterruptError()

def nonmNoDist(X1,Z1,width,cutDist,RWidth,Cwidth):

    DistMatrix  = np.zeros((len(X1), len(X1)))
    for j_1 in  range(len(X1)):
        if Z1[j_1,0] == -1:
            break
        for j_2 in range(len(X1)):
            if Z1[j_2,0] == -1:
                break
            RDist = abs(Z1[j_1,0] - Z1[j_2,0])
            CDist = abs(Z1[j_1,1] - Z1[j_2,1])
            dd = _MDist(X1[j_1],X1[j_2],width,cutDist,RWidth,Cwidth)
            dd = dd *_StochDist(RDist,CDist,RWidth,Cwidth)
            DistMatrix[j_1,j_2] = dd

    return DistMatrix

def _MDist(X1,X2,width,cutDist,RWidth,Cwidth):
    maxGausDist = 8*width
    #maxRDist = 8*RWidth
    #maxCDist = 8*Cwidth
    AAdist = 0

    for m_1 in range(len(X1[0])):

        if X1[0,m_1] > cutDist:
            break

        for m_2 in range(len(X2[0])):
            if X2[0,m_2] > cutDist:
                break
            if  abs(X2[0,m_2] - X1[0,m_1]) < maxGausDist:
                RDist = abs(X1[1,m_1] - X2[1,m_2])
                CDist = abs(X1[2,m_1] - X2[2,m_2])
                #if RDist < maxRDist and CDist < maxCDist:
                d = _dist(X1[0,m_1], X2[0,m_2],width,cutDist)
                d = d *_StochDist(RDist,CDist,RWidth,Cwidth)
                AAdist += d * (1 + X1[3,m_1]*X2[3,m_2] + X1[4,m_1]*X2[4,m_2])
    return AAdist

def _dist(x1,x2,w1,w2):
    return np.exp(-((x1-x2)**2)/(4*w1**2))*(1 - np.sin(np.pi * x1/(2 * w2)))*(1 - np.sin(np.pi * x2/(2 * w2))) #*(math.cos(pi * x1/(2 * w2))*math.cos(pi * x2/(2 * w2)))

def _StochDist(D1,D2,W1,W2):
    rdist =  W1**2/(W1**2 + D1**2) * W2**2/(W2**2 + D2**2)
    # print "rdist", rdist
    return rdist
    # return W1**2/(W1**2 + D1**2) * W2**2/(W2**2 + D2**2)
    #math.exp(-(D1**2)/(4*W1**2) -(D2**2)/(4*W2**2))

def unravelKerMatrix(KM):

    KM_unravled = np.nansum(KM,axis = (-2,-1))
    return KM_unravled

if __name__ == "__main__":

    mols = []
    path = "tests/xyz/"
    filenames = os.listdir(path)

    print "Generating ARAD descriptors from FML interface ..."
    for filename in filenames:

        mol = fml.Molecule()
        mol.read_xyz(path + filename)
        mol.generate_arad_descriptor(size=30)
        mols.append(mol)

    arad = fml.ARAD()

    x1 = []
    z1 = []

    print "Generating ARAD descriptors from reference implementation ..."
    for mol in mols[:5]:

        z1.append(np.array(mol.nuclear_charges))
        x1.append(arad.describe(np.array(mol.coordinates),np.array(mol.nuclear_charges)))

    x2 = []
    z2 = []

    for mol in mols[5:]:

        z2.append(np.array(mol.nuclear_charges))
        x2.append(arad.describe(np.array(mol.coordinates),np.array(mol.nuclear_charges)))

    print "Calculating ARAD Distance matrix from reference code ..."
    Da = condense(x1,x2,z1,z2, molecular=False)

    np.set_printoptions(linewidth=10000000)

    print "Calculating ARAD Distance matrix from FML code ... "
    Db = get_l2_distance_arad(np.array(x1),np.array(x2),z1,z2)

    print "Example distance matrix (reference)"
    print Da[3,2][:10,:10]
    print Da.shape
    print "Example distance matrix (FML ARAD)"
    print Db[3,2][:10,:10]
    print Db.shape

    print "Max abs element distance difference:", np.nanmax(np.abs(Da - Db))
    print "Max abs element distance difference index:", np.nanargmax(np.abs(Da - Db))

    sigmas = [0.1, 1.0, 10.0, 100.0]

    Ka = np.zeros((len(sigmas),len(x1),len(x2)))

    print "Calculating kernel from reference distance matrix"
    for s, sigma in enumerate(sigmas):

        for i in range(len(x1)):
            for j in range(len(x2)):

                Ka[s,i,j] = np.nansum(np.exp(Da[i,j] * (-1.0 / sigma**2)))

    print "Calculating kernel from FML distance matrix"
    Kb = get_atomic_kernels_arad(np.array(x1), np.array(x2), z1, z2, sigmas)

    print "Example kernel matrix (reference)"
    print Ka[-2]

    print "Example kernel matrix (ARAD)"
    print Kb[-2]

    print "Max abs element kernel difference:", np.nanmax(np.abs(Ka - Kb))
    print "Max abs element kernel difference index:", np.nanargmax(np.abs(Ka - Kb))


from contextlib import AbstractContextManager, suppress
from enum import IntEnum, unique
from time import sleep

from scipy.constants import speed_of_light as c_vacuum

from .base import InstrumentException
from .utils import is_valid_IP, timeout
from .XPS_C8_drivers import XPS

air_refractive_index = 1.0003
c_air = c_vacuum / air_refractive_index  # meters per second

@unique
class XPSC8Errors(IntEnum):
    """
    Enumeration of possible XPS-C8 errors.
    """

    # TODO: Error code and error description
    #       See example:
    #       https://docs.python.org/3/library/enum.html#planet

    NoError = 0
    SocketConnectionError = -1
    WrongObjectTypeForCmd = -8
    ParameterOutOfRangeOrIncorrect = -17
    PositionerNameDoesNotExistOrUnknownCmd = -18
    GroupNameDoesNotExistOrUnknownCmd = -19
    NotAllowedAction = -22
    FollowingError = -25
    EmergencySignal = -26
    MoveAborted = -27
    HomeSearchTimeout = -28
    MotionDoneTimeout = -33
    PositionOutsideTravelLimits = -35
    SlaveErrorDisablingMaster = -44
    InconsistentMechanicalZero = -49
    MotorInitiError = -50
    BothEndRunsActivated = -113
    WarningErrorDuringMove = -120
    NotExpectedPositionAfterMotion = -221

def _errcheck(returned):
    """
    This function checks that any error code is 0 (success)
    Otherwise, raise an InstrumentException with the correct
    error code. The ``returned`` value is either an error code
    or a list (in which case the error code is the first value)
    """
    if returned is None:
        raise ValueError("None has been returned")
    with suppress(ValueError):  # unknown error code.
        if isinstance(returned, list):
            errcode = returned[0]
        else:
            errcode = returned
        error = XPSC8Errors(int(errcode))
        if error != XPSC8Errors.NoError:
            raise InstrumentException(error)
    return returned

class DelayStage(AbstractContextManager):
    """
    Abstract interface to one delay-stage
    connected to a Newport XPS C8.

    Parameters
    ----------
    address : str, optional
        IP address of the XPS, e.g. '192.168.33.101'.

    Raises
    ------
    ValueError : if ``address`` is an invalid IPv4 address.
    InstrumentException : if any connection error occurs.
    """

    _driver = XPS()

    # group and positioner must be overriden in subclasses
    group = ""
    positioner = group + ""

    def __init__(self, address, **kwargs):
        # self.socket_id = None

        # According to TCP_ConnectToServer documentation,
        # port is always 5001
        if not is_valid_IP(address):
            raise ValueError("{} is an invalid IPv4 address".format(address))

        self.socket_id = _errcheck(
            self._driver.TCP_ConnectToServer(IP=address, port=5001, timeOut=10)
        )

        # Reset state by killing the group, and initializing again
        # Note: GroupKill returns [errcode, string] for some reason
        # even though documentation doesn't say that
        _errcheck(self._driver.KillAll(self.socket_id))
        _errcheck(self._driver.GroupInitialize(self.socket_id, self.group))
        _errcheck(self._driver.GroupHomeSearch(self.socket_id, self.group))

        # Get position limits
        errcode, self.min_limit, self.max_limit = _errcheck(
            self._driver.PositionerUserTravelLimitsGet(self.socket_id, self.positioner)
        )

    def __exit__(self, *args, **kwargs):
        self.disconnect()
        super().__exit__(*args, **kwargs)

    @staticmethod
    def delay_to_distance(delay):
        """ Calculate the distance by which to move [mm] for light round-trip
        of ``delay`` picoseconds """
        # Distance to move is half because of back-and-forth motion
        # along the stage

        # Increasing distance means the laser pulse arrives later;
        # since the electron pulse arrives at fixed times
        # this means that increasing distance -> earlier probing
        move_meters = (delay / 1e12) * (c_air / 2)
        return -1 * move_meters * 1e3

    @staticmethod
    def distance_to_delay(dist):
        """ Calculate the extra time [ps] it takes for light to make a round-trip
        if the stage moves by ``dist`` millimeters. """
        # Increasing distance means the laser pulse arrives later;
        # since the electron pulse arrives at fixed times
        # this means that increasing distance -> earlier probing
        extra_path = 2 * float(dist) / 1e3  # extra path [meters]
        return -1 * (extra_path / c_air) * 1e12  # [picoseconds]

    def disconnect(self):
        """ Disconnect from the XPS """
        self._driver.TCP_CloseSocket(self.socket_id)

    def _wait_end_of_move(self, tout=10, tol=5e-3):
        """
        Wait for end of move, i.e. when the current position is close
        enough to the target.

        Parameters
        ----------
        tout : float, optional
            Time-out time in seconds
        tol : float, optional
            Position tolerance.
        """
        with timeout(tout, InstrumentException, exc_message="Movement timeout"):
            while abs(self.current_position() - self.target_position()) > tol:
                sleep(0.1)

    def target_position(self):
        """
        Get the current absolute position setpoint

        Returns
        -------
        pos : float
        """
        errcode, position = _errcheck(
            self._driver.GroupPositionTargetGet(
                self.socket_id, self.positioner, nbElement=1
            )
        )
        return float(position)

    def current_position(self):
        """
        Get current absolute position

        Returns
        -------
        pos : float
        """
        errcode, position = _errcheck(
            self._driver.GroupPositionCurrentGet(
                self.socket_id, self.positioner, nbElement=1
            )
        )
        return float(position)

    def relative_move(self, move):
        """
        Move the delay stage relatively to current position, by distance. This
        function returns when move is completed.

        Parameters
        ----------
        move : float
            Relative move [mm]
        """
        move = float(move)
        # For some reason, the targetDisplacement parameter
        # to GroupMoveRelative should be an iterable...
        _errcheck(
            self._driver.GroupMoveRelative(self.socket_id, self.positioner, [move])
        )
        return self._wait_end_of_move()

    def absolute_move(self, move):
        """
        Move the delay stage to a new absolute position, by distance. This
        function returns when move is completed.

        Parameters
        ----------
        move : float
            Absolute move [mm]
        """
        move = float(move)
        # For some reason, the targetDisplacement parameter
        # to GroupMoveAbsolute should be an iterable...
        _errcheck(
            self._driver.GroupMoveAbsolute(self.socket_id, self.positioner, [move])
        )
        return self._wait_end_of_move()

    def relative_time_shift(self, shift):
        """
        Move the delay stage to achieve a certain relative time-shift. This
        function returns when move is completed.

        Parameters
        ----------
        shift : float
            Time-shift in picoseconds
        """
        shift = float(shift)
        return self.relative_move(self.delay_to_distance(shift))

    def absolute_time(self, time, tzero_position=0.0):
        """
        Move the delay stage to achieve a time-shift with respect
        to a position for time-zero. This function returns when move is completed.

        Parameters
        ----------
        time : float
            Desired absolute time-shift in picoseconds
        tzero_position : float, optional
            Time-zero position in millimeters.
        """
        time, tzero_position = float(time), float(tzero_position)
        return self.absolute_move(self.delay_to_distance(time) + tzero_position)

class ILS250PP(DelayStage):
    """
    Interface to an ILS250PP delay-stage connected
    to a Newport XPS C8 positioner.
    """

    group = "GROUP5"
    positioner = group + ".POSITIONER"

    # Internet address : 132.206.175.95
    # local address    : 192.168.254.254
    def __init__(self, address="192.168.254.254", **kwargs):
        super().__init__(address, **kwargs)


from typing import Tuple, List, Dict, Any
from enum import Enum
import numpy as np
import tensorflow as tf
from tensorflow.python import pywrap_tensorflow
from tensorflow import Tensor
from copy import copy

from decompose.distributions.distribution import DrawType, UpdateType
from decompose.distributions.distribution import Distribution
from decompose.distributions.uniform import Uniform
from decompose.distributions.nnUniform import NnUniform
from decompose.likelihoods.likelihood import Likelihood
from decompose.likelihoods.specificNormal2dLikelihood import SpecificNormal2dLikelihood
from decompose.likelihoods.allSpecificNormal2dLikelihood import AllSpecificNormal2dLikelihood
from decompose.likelihoods.normal2dLikelihood import Normal2dLikelihood
from decompose.likelihoods.normalNdLikelihood import NormalNdLikelihood
from decompose.likelihoods.cvNormal2dLikelihood import CVNormal2dLikelihood
from decompose.likelihoods.cvNormalNdLikelihood import CVNormalNdLikelihood
from decompose.postU.postU import PostU
from decompose.stopCriterions.llhImprovementThreshold import LlhImprovementThreshold
from decompose.stopCriterions.llhStall import LlhStall
from decompose.cv.cv import CV

EstimatorSpec = tf.estimator.EstimatorSpec

class NoiseUniformity(Enum):
    HOMOGENEOUS = 0
    HETEROGENEOUS = 1
    LAST_FACTOR_HETEROGENOUS = 2

HOMOGENEOUS = NoiseUniformity.HOMOGENEOUS
HETEROGENEOUS = NoiseUniformity.HETEROGENEOUS
LAST_FACTOR_HETEROGENOUS = NoiseUniformity.LAST_FACTOR_HETEROGENOUS

class parameterProperty(object):
    """Decorator for descriptors that update tf variables during set.

    This decorator is the same as the python property decorator except
    that its setter method accepts a name which can updates a
    tensorflow variable depending on
    """

    def __init__(self, fget=None, fset=None, fdel=None, doc=None, name=None):
        self.name = name
        self.fget = fget
        self.fset = fset
        self.fdel = fdel
        if doc is None and fget is not None:
            doc = fget.__doc__
        self.__doc__ = doc

    def __set__(self, obj, values):
        if self.fset is None:
            raise AttributeError("can't set attribute")
        newValues = []
        for f, value in enumerate(values):
            if obj.transform and (f == 0):
                name = f"{f}tr"
            else:
                name = f"{f}"
            with tf.variable_scope("U", reuse=tf.AUTO_REUSE):
                var = tf.get_variable(name, dtype=obj.dtype)
            value = tf.assign(var, value)
            newValues.append(value)
        value = tuple(newValues)
        self.fset(obj, value)

    def __get__(self, obj, objtype=None):
        if obj is None:
            return self
        if self.fget is None:
            raise AttributeError("unreadable attribute")
        return self.fget(obj)

    def __delete__(self, obj):
        if self.fdel is None:
            raise AttributeError("can't delete attribute")
        self.fdel(obj)

    def getter(self, fget):
        return type(self)(fget, self.fset, self.fdel, self.__doc__, self.name)

    def setter(self, name):
        if name is None:
            return type(self)(self.fget, None, self.fdel, self.__doc__, None)
        if type(name) is not str:
            raise ValueError("setter takes a name argument as a string")

        def noop(fset):
            return type(self)(self.fget, fset, self.fdel, self.__doc__, name)
        return(noop)

    def deleter(self, fdel):
        return type(self)(self.fget, self.fset, fdel, self.__doc__, self.name)

class Phase(Enum):
    INIT = 1
    EM = 2
    BCD = 3

class TensorFactorisation(object):

    def __init__(self,
                 U: List[Tensor],
                 priorU: List[Distribution],
                 likelihood: Likelihood,
                 dtype: tf.DType,
                 stopCriterion,
                 phase: Phase,
                 noiseUniformity: NoiseUniformity,
                 transform: bool = False) -> None:

        # setup the model
        self.dtype = dtype
        self.__transform = transform
        self.__noiseUniformity = noiseUniformity
        self.likelihood = likelihood
        self.stopCriterion = stopCriterion
        self.postU = []  # type: List[PostU]
        for f, priorUf in enumerate(priorU):
            postUf = PostU(likelihood, priorUf, f)
            self.postU.append(postUf)

        # create or reuse the variables for the filter banks
        for f, Uf in enumerate(copy(U)):
            if transform and (f == 0):
                paramName = "{}tr".format(f)
            else:
                paramName = "{}".format(f)
            with tf.variable_scope("U", reuse=tf.AUTO_REUSE):
                UfVar = tf.get_variable(paramName,
                                        dtype=dtype,
                                        initializer=Uf)
            U[f] = UfVar
        self.__U = tuple(U)
        if phase == Phase.EM or phase == Phase.INIT:
            self.__setEm()
        elif phase == Phase.BCD:
            self.__setBcd()
        else:
            raise ValueError

    @classmethod
    def random(cls,
               priorU: List[Distribution],
               likelihood: Likelihood,
               M: Tuple[int, ...],
               K: int,
               dtype: tf.DType,
               phase: Phase,
               stopCriterion,
               noiseUniformity: NoiseUniformity = HOMOGENEOUS,
               transform: bool = False) -> "TensorFactorisation":

        # initialize U
        dtype = tf.as_dtype(dtype)
        zero = tf.constant(0., dtype=dtype)
        one = tf.constant(1., dtype=dtype)
        normal = tf.distributions.Normal(loc=zero, scale=one)
        F = len(M)
        U = []
        for f in range(F):
            if priorU[f].nonNegative:
                UfInit = tf.abs(normal.sample(sample_shape=(K, M[f])))
            else:
                UfInit = normal.sample(sample_shape=(K, M[f]))
            U.append(UfInit)

        # instantiate
        tefa = TensorFactorisation(U=U,
                                   priorU=priorU,
                                   likelihood=likelihood,
                                   dtype=dtype,
                                   phase=phase,
                                   transform=transform,
                                   noiseUniformity=noiseUniformity,
                                   stopCriterion=stopCriterion)
        return(tefa)

    @property
    def transform(self) -> bool:
        return(self.__transform)

    @property
    def noiseUniformity(self) -> NoiseUniformity:
        return(self.__noiseUniformity)

    @parameterProperty
    def U(self) -> Tuple[tf.Tensor, ...]:
        return(self.__U)

    @U.setter(name="U")
    def U(self, U: Tuple[tf.Tensor, ...]):
        self.__U = U

    def update(self, X: Tensor) -> Tuple[Tensor, ...]:
        # update stopping criterion
        stopCritDeps = self.stopCriterion.update(self, X)

        # perform updated depending on this is train of transformation
        with tf.control_dependencies(stopCritDeps):
            if self.transform:
                self.updateTransform(X)
            else:
                self.updateTrain(X)

        # return the updated tensors
        return(self.U)

    def updateTrain(self, X: Tensor) -> None:
        # store filterbanks in a list
        U = list(self.U)  # type: List[Tensor]

        # update the parameters of the likelihood
        self.likelihood.update(U=U, X=X)

        # update the filters in reversed order
        for f, postUf in reversed(list(enumerate(self.postU))):
            U = self.rescale(U=U, fNonUnit=f)
            U[f] = postUf.update(U=U, X=X, transform=False)

        # update the filter banks
        self.U = tuple(U)

    def updateTransform(self, X: Tensor) -> None:
        # store filterbanks in a list
        U = list(self.U)  # type: List[Tensor]

        # calculate updates of the first filterbank
        U[0] = self.postU[0].update(U=U, X=X, transform=True)

        # update the filter banks
        self.U = tuple(U)

    def rescale(self, U: List[Tensor], fNonUnit: int) -> List[Tensor]:
        """Puts all variance in the factor `fUpdate`-th factor.

        The method assumes that the norm of all filters is larger than 0."""
        F = len(U)

        # calculathe the scale for each source
        scaleOfSources = tf.ones_like(U[0][..., 0])
        for f in range(F):
            scaleOfSources = scaleOfSources*tf.norm(U[f], axis=-1)

        for f in range(F):
            # determine rescaling constant depending on the factor number
            Uf = U[f]
            normUf = tf.norm(Uf, axis=-1)
            if f == fNonUnit:
                # put all variance in the filters of the fUpdate-th factor
                rescaleConstant = scaleOfSources/normUf
            else:
                # normalize the filters all other factors
                rescaleConstant = 1./normUf

            # rescaled filters
            Uf = Uf*rescaleConstant[..., None]
            U[f] = Uf

        return(U)

    def __setEm(self) -> None:
        """Set prior and noise distributions to perform EM updates."""
        for postUf in self.postU:
            postUf.prior.drawType = DrawType.SAMPLE
            postUf.prior.updateType = UpdateType.ALL
        self.likelihood.noiseDistribution.drawType = DrawType.SAMPLE
        self.likelihood.noiseDistribution.updateType = UpdateType.ALL

    def __setBcd(self) -> None:
        """Set prior and noise distributions to perform BCD updates."""
        for postUf in self.postU:
            postUf.prior.drawType = DrawType.MODE
            postUf.prior.updateType = UpdateType.ONLYLATENTS
        self.likelihood.noiseDistribution.drawType = DrawType.MODE
        self.likelihood.noiseDistribution.updateType = UpdateType.ONLYLATENTS

    def loss(self, X: Tensor) -> Tensor:
        """Loss of the data `X` given the parameters."""
        loss = self.likelihood.loss(self.U, X)
        loss = tf.cast(loss, tf.float64)
        return(loss)

    def llh(self, X: Tensor) -> Tensor:
        """Log likelihood of the parameters given data `X`."""

        # log likelihood of the noise
        llh = self.likelihood.llh(self.U, X)

        # log likelihood of the factors
        U = list(self.U)
        for f, postUf in enumerate(self.postU):
            U = self.rescale(U=U, fNonUnit=f)
            UfT = tf.transpose(U[f])
            llhUf = tf.reduce_sum(postUf.prior.llh(UfT))
            llh = llh + llhUf
        llh = tf.cast(llh, tf.float64)
        return(llh)

    def llhIndividual(self, X: Tensor) -> Tensor:
        """Log likelihood of the parameters given data `X`."""

        # log likelihood of the noise
        llhRes = self.likelihood.llh(self.U, X)
        llh = llhRes

        # log likelihood of the factors
        llhU = []
        llhUfk = []
        U = list(self.U)
        for f, postUf in enumerate(self.postU):
            U = self.rescale(U=U, fNonUnit=f)
            UfT = tf.transpose(U[f])
            llhUfk.append(tf.reduce_sum(postUf.prior.llh(UfT), axis=0))
            llhUf = tf.reduce_sum(postUf.prior.llh(UfT))
            llh = llh + llhUf
            llhU.append(llhUf)
        llh = tf.cast(llh, tf.float64)
        return(llh, llhRes, llhU, llhUfk)

    @staticmethod
    def type():
        return(TensorFactorisation)

    def id(self) -> str:
        """Generate a string representation of the model configuration"""
        strId = ""
        for f, postUf in enumerate(self.postU):
            strId += "U{}".format(f) + postUf.prior.id()
        strId += "_" + self.likelihood.id
        return(strId)

    @classmethod
    def __model(cls, data: Tensor, priorTypes: List[Distribution],
                M: Tuple[int, ...],
                K: int, stopCriterion, phase: Phase, dtype: tf.DType,
                reuse=False,
                isFullyObserved: bool = True,
                cv: CV = None,
                transform: bool = False,
                noiseUniformity: NoiseUniformity = HOMOGENEOUS,
                suffix: str = "") -> "TensorFactorisation":
        varscope = "stopCriterion" + phase.name
        stopCriterion.init(ns=varscope)
        F = len(priorTypes)

        # selecting the apropriate likelihood
        useNormal2dLikelihood = (
            F == 2
            and cv is None
            and isFullyObserved
            and (noiseUniformity == HOMOGENEOUS
                 or phase == Phase.INIT))
        useAllSpecificNormal2dLikelihood = (
            F == 2
            and cv is None
            and isFullyObserved
            and noiseUniformity == HETEROGENEOUS
            and phase != Phase.INIT)
        useSpecificNormal2dLikelihood = (
            F == 2
            and cv is None
            and isFullyObserved
            and noiseUniformity == LAST_FACTOR_HETEROGENOUS
            and phase != Phase.INIT)
        useCVNormal2dLikelihood = (
            F == 2
            and (cv is not None
                 or not isFullyObserved)
            and noiseUniformity == HOMOGENEOUS)
        useNormalNdLikelihood = (
            F > 2
            and cv is None
            and isFullyObserved
            and noiseUniformity == HOMOGENEOUS)
        useCVNormalNdLikelihood = (
            F > 2
            and (cv is not None
                 or not isFullyObserved)
            and noiseUniformity == HOMOGENEOUS)

        # instantiate the likelihood
        with tf.variable_scope(f"{suffix}", reuse=reuse):
            if useNormal2dLikelihood:
                likelihood = Normal2dLikelihood(
                    M=M, K=K, dtype=dtype)  # type: Likelihood
            elif useAllSpecificNormal2dLikelihood:
                likelihood = AllSpecificNormal2dLikelihood(
                    M=M, K=K, dtype=dtype)
            elif useSpecificNormal2dLikelihood:
                likelihood = SpecificNormal2dLikelihood(
                    M=M, K=K, dtype=dtype)
            elif useCVNormal2dLikelihood:
                likelihood = CVNormal2dLikelihood(
                    M=M, K=K, dtype=dtype, cv=cv)
            elif useNormalNdLikelihood:
                likelihood = NormalNdLikelihood(
                    M=M, K=K, dtype=dtype)
            elif useCVNormalNdLikelihood:
                likelihood = CVNormalNdLikelihood(
                    M=M, K=K, dtype=dtype, cv=cv)
            else:
                raise NotImplementedError()
            likelihood.init(data)

            # instantiate the priors
            priors = []
            for f, priorType in enumerate(priorTypes):
                prior = priorType.random(shape=(K,), latentShape=(M[f],),
                                         name=f"prior{suffix}{f}", dtype=dtype)
                priors.append(prior)

        # instantiate the model
        tefa = cls.random(priorU=priors, likelihood=likelihood, M=M, K=K,
                          phase=phase, stopCriterion=stopCriterion,
                          dtype=dtype, noiseUniformity=noiseUniformity,
                          transform=transform)
        return(tefa)

    @classmethod
    def __estimatorSpec(cls, mode, features, device: str,
                        isFullyObserved: bool,
                        priors: List[Distribution],
                        K: int, stopCriterionInit, stopCriterionEM,
                        stopCriterionBCD,
                        cv: CV, path: str,
                        noiseUniformity: NoiseUniformity,
                        transform: bool, dtype: tf.DType) -> EstimatorSpec:
        # PREDICT and EVAL are not supported
        if mode != tf.estimator.ModeKeys.TRAIN:
            raise ValueError

        # TRAIN
        with tf.device(device):
            # check the input data
            labels = list(features.keys())
            assert len(labels) == 1
            data = features[labels[0]]
            dataShape = tuple(data.get_shape().as_list())
            assert len(dataShape) == len(priors)

            # shape of the data
            M = data.get_shape().as_list()

            # create llh variable
            inf = np.float64(np.inf)
            with tf.variable_scope("llh"):
                llhVar = tf.get_variable("llh", dtype=tf.float64,
                                         initializer=-inf)

            # create loss variable
            with tf.variable_scope("loss"):
                lossVar = tf.get_variable("loss", dtype=tf.float64,
                                          initializer=inf)

            # create global stopping variable
            with tf.variable_scope("stopCriterion"):
                stopVar = tf.get_variable("stop", dtype=tf.bool,
                                          initializer=False)

            # INIT model
            initPriors = []  # type: List[Distribution]
            for prior in priors:
                if prior.nonNegative:
                    initPriors.append(NnUniform())
                else:
                    initPriors.append(Uniform())
            tefaInit = cls.__model(data=data, priorTypes=initPriors, K=K, M=M,
                                   isFullyObserved=isFullyObserved,
                                   stopCriterion=stopCriterionInit,
                                   dtype=dtype, reuse=False,
                                   transform=transform, cv=cv,
                                   phase=Phase.INIT,
                                   noiseUniformity=noiseUniformity,
                                   suffix="init")

            # EM model
            tefaEM = cls.__model(data=data, priorTypes=priors, K=K, M=M,
                                 isFullyObserved=isFullyObserved,
                                 stopCriterion=stopCriterionEM,
                                 dtype=dtype, phase=Phase.EM,
                                 transform=transform, cv=cv,
                                 noiseUniformity=noiseUniformity,
                                 reuse=tf.AUTO_REUSE)

            # BCD model
            tefaBCD = cls.__model(data=data, priorTypes=priors, K=K, M=M,
                                  isFullyObserved=isFullyObserved,
                                  stopCriterion=stopCriterionBCD,
                                  dtype=dtype, phase=Phase.BCD,
                                  transform=transform, cv=cv,
                                  noiseUniformity=noiseUniformity,
                                  reuse=tf.AUTO_REUSE)

            # replace nan with zeros
            data = tf.where(tf.is_nan(data), tf.zeros_like(data), data)

            # conduct an update depending on the current phase
            stopVarInit = tefaInit.stopCriterion.stopVar
            stopVarEm = tefaEM.stopCriterion.stopVar
            loss = tf.cond(tf.logical_not(stopVarInit),
                           lambda: tefaInit.loss(X=data),
                           lambda: tf.cond(tf.logical_not(stopVarEm),
                                           lambda: tefaEM.loss(X=data),
                                           lambda: tefaBCD.loss(X=data)))

            # conduct an update depending on the current phase
            deps = tf.cond(tf.logical_not(stopVarInit),
                           lambda: tefaInit.update(X=data),
                           lambda: tf.cond(tf.logical_not(stopVarEm),
                                           lambda: tefaEM.update(X=data),
                                           lambda: tefaBCD.update(X=data)))

            # update the global stop variable
            stopVarBcd = tefaBCD.stopCriterion.stopVar
            stop = tf.logical_and(stopVarInit,
                                  tf.logical_and(stopVarEm,
                                                 stopVarBcd))
            with tf.control_dependencies(deps):
                updatedStopVar = tf.assign(stopVar, stop)

            # if stopping criterion is reached store the llh
            updates = tf.cond(stop,
                              lambda: (tf.assign(llhVar, tefaBCD.llh(data)),
                                       tf.assign(lossVar, tefaBCD.loss(data))),
                              lambda: (llhVar, lossVar))

            # increment global step variable
            with tf.control_dependencies([updatedStopVar, *updates]):
                step = tf.train.get_or_create_global_step()
                trainOp = tf.assign(step, step + 1)

            # log summaries
            tf.summary.scalar("loss", loss)
            llh = tf.cond(tf.logical_not(stopVarInit),
                          lambda: tefaInit.llhIndividual(X=data),
                          lambda: tf.cond(tf.logical_not(stopVarEm),
                                          lambda: tefaEM.llhIndividual(X=data),
                                          lambda: tefaBCD.llhIndividual(X=data)))
            llh, llhRes, llhU, llhUfk = llh
            tf.summary.scalar("llh", llh)
            tf.summary.scalar("llhResiduals", llhRes)
            for f, (llhUf, llhUfk) in enumerate(zip(llhU, llhUfk)):
                tf.summary.scalar(f"llhU{f}", llhUf)

            SAVE_EVERY_N_STEPS = 1  # TODO: make configurable
            summary_hook = tf.train.SummarySaverHook(
                SAVE_EVERY_N_STEPS,
                output_dir=path,
                summary_op=tf.summary.merge_all())

        return EstimatorSpec(mode, loss=loss, train_op=trainOp,
                             training_hooks=[summary_hook])

    @classmethod
    def getEstimator(cls, priors: Tuple[Distribution, ...], K: int,
                     dtype: tf.DType = tf.float32,
                     isFullyObserved: bool = True,
                     noiseUniformity: NoiseUniformity = HOMOGENEOUS,
                     stopCriterionInit=LlhStall(10),
                     stopCriterionEM=LlhStall(100),
                     stopCriterionBCD=LlhImprovementThreshold(1e-2),
                     path: str = "/tmp", device: str = "/cpu:0",
                     cv: CV = None):

        def model_fn(features, labels, mode):
            es = cls.__estimatorSpec(mode=mode, features=features,
                                     isFullyObserved=isFullyObserved,
                                     device=device, priors=priors,
                                     noiseUniformity=noiseUniformity,
                                     stopCriterionInit=stopCriterionInit,
                                     stopCriterionEM=stopCriterionEM,
                                     stopCriterionBCD=stopCriterionBCD,
                                     cv=cv, path=path, K=K,
                                     transform=False, dtype=dtype)
            return(es)

        est = tf.estimator.Estimator(model_fn=model_fn,
                                     model_dir=path)
        return(est)

    @classmethod
    def getTransformEstimator(cls, priors: Tuple[Distribution, ...], K: int,
                              chptFile: str, dtype: tf.DType = tf.float32,
                              noiseUniformity: NoiseUniformity = HOMOGENEOUS,
                              stopCriterionInit=LlhStall(10),
                              stopCriterionEM=LlhStall(100),
                              stopCriterionBCD=LlhImprovementThreshold(1e-2),
                              path: str = "/tmp", device: str = "/cpu:0"):
        # configuring warm start settings
        reader = pywrap_tensorflow.NewCheckpointReader(chptFile)
        varList = [v for v in reader.get_variable_to_shape_map().keys()
                   if (v != "U/0" and
                       v != "global_step" and
                       v != "stop" and
                       not v.startswith(f"stopCriterion{Phase.INIT.name}/") and
                       not v.startswith(f"stopCriterion{Phase.EM.name}/") and
                       not v.startswith(f"stopCriterion{Phase.BCD.name}/"))]
        wsVars = "|".join(varList)
        ws = tf.estimator.WarmStartSettings(ckpt_to_initialize_from=chptFile,
                                            vars_to_warm_start=wsVars)

        def model_fn(features, labels, mode):
            es = cls.__estimatorSpec(mode=mode, features=features,
                                     isFullyObserved=True,
                                     device=device, priors=priors,
                                     noiseUniformity=noiseUniformity,
                                     stopCriterionInit=stopCriterionInit,
                                     stopCriterionEM=stopCriterionEM,
                                     stopCriterionBCD=stopCriterionBCD,
                                     K=K, path=path, cv=None,
                                     transform=True, dtype=dtype)
            return(es)

        est = tf.estimator.Estimator(model_fn=model_fn,
                                     model_dir=path,
                                     warm_start_from=ws)
        return(est)


import os
from json import load
from copy import deepcopy

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import binned_statistic

from astropy.io import fits
import astropy.units as u

__all__ = ['Planet', 'Filter', 'PhaseCurve']

planets_path = os.path.join(os.path.dirname(__file__), 'data', 'planets.json')
filters_path = os.path.join(os.path.dirname(__file__), 'data', 'filters.json')
pc_path = os.path.join(os.path.dirname(__file__), 'data', 'lightcurves.fits')

class Planet(object):
    """
    Transiting planet parameters.

    This is meant to be a duck-type drop-in for the ``batman`` package's
    transiting exoplanet parameters ``TransitParams`` object.
    """
    with open(planets_path, 'r') as _f:
        _planets = load(_f)

    def __init__(self, per=None, t0=None, inc=None, rp=None, ecc=None, w=None,
                 a=None, u=None, fp=None, t_secondary=None, T_s=None, rp_a=None,
                 limb_dark='quadratic', name=None):
        """
        Parameters
        ----------
        per : float
            Orbital period [days]
        t0 : float
            Mid-transit time
        inc : float
            Orbital inclination [deg]
        rp : float
            Ratio of planet to star radius
        ecc : float
            Eccentricity
        w : float
            Argument of periastron [deg]
        a : float
            Semimajor axis normalized by the stellar radius
        u : list
            (i.e.) Quadratic limb-darkening parameters
        fp : float
            Planetary flux out of eclipse
        t_secondary : float
            Time of secondary eclipse
        T_s : float
            Temperature of the host star [K]
        rp_a : float
            Radius of the planet over the semimajor axis
        limb_dark : str
            Limb darkening law to use
        name : str
            Name metadata for the planet
        """
        self.per = per
        self.t0 = t0
        self.inc = inc
        self.rp = rp
        self.ecc = ecc
        self.w = w
        self.a = a
        self.u = u
        self.limb_dark = limb_dark
        self.fp = fp
        self.t_secondary = t_secondary
        self.T_s = T_s
        self.rp_a = rp_a
        self.name = name

    @classmethod
    def from_name(cls, name):
        """
        Initialize a Planet instance from the target name.

        There's a small (but growing?) database of planets pre-defined in the
        ``kelp/data/planets.json`` file. If your favorite planet is missing,
        pull requests are welcome!

        Parameters
        ----------
        name : str (i.e.: "Kepler-7" or "KELT-9")
             Name of the planet
        """

        return cls(name=name, **cls._planets[name])

    def eclipse_model(self, xi):
        r"""
        Compute eclipse model at orbital phases ``xi``.

        Parameters
        ----------
        xi : `~numpy.ndarray`
            Orbital phase angle :math:`\xi`

        Returns
        -------
        eclipse : `~numpy.ndarray`
            Eclipse model normalized such that flux is zero in eclipse.
        """
        from batman import TransitModel

        xi_over_pi = xi / np.pi
        eclipse = TransitModel(self, xi_over_pi, transittype='secondary',
                               exp_time=xi_over_pi[1] - xi_over_pi[0],
                               supersample_factor=3,
                               ).light_curve(self)
        eclipse -= eclipse.min()
        return eclipse

class Filter(object):
    """
    Astronomical filter object.
    """
    with open(filters_path, 'r') as _f:
        _filters = load(_f)

    def __init__(self, wavelength, transmittance, name=None):
        """
        Parameters
        ----------
        wavelength : `~numpy.ndarray`
            Wavelength array
        transmittance : `~numpy.ndarray`
            Transmittance array
        """
        self.wavelength = wavelength
        self.transmittance = transmittance
        self.name = name

    @classmethod
    def from_name(cls, name):
        """
        Initialize a Filter instance from the filter name.

        Parameters
        ----------
        name : str
             Name of the filter. Examples include "IRAC 1", "IRAC 2", "Kepler",
             "TESS", and "CHEOPS".
        """
        return cls(np.array(cls._filters[name]['wavelength']) * u.um,
                   np.array(cls._filters[name]['transmittance']),
                   name)

    def plot(self, ax=None, **kwargs):
        """
        Plot the filter transmittance curve.

        Parameters
        ----------
        ax : `~matplotlib.axes.Axes`
            Matplotlib axis object
        kwargs : dict
            Dictionary passed to the `~matplotlib.pyplot.plot` command

        Returns
        -------
        ax : `~matplotlib.axes.Axes`
            Updated axis object
        """
        if ax is None:
            ax = plt.gca()
        ax.set(title=self.name)
        ax.plot(self.wavelength, self.transmittance, **kwargs)

        return ax

    def bin_down(self, bins=10):
        """
        Bin down the filter bandpass wavelengths and transmittances (shortcut
        for faster integration over the bandpass).

        Parameters
        ----------
        bins : int
            Number of bins in the binned transmittance curve.
        """
        bs = binned_statistic(self.wavelength.value, self.transmittance,
                              bins=bins, statistic='median')
        bincenters = 0.5 * (bs.bin_edges[1:] + bs.bin_edges[:-1])

        self.wavelength = bincenters * self.wavelength.unit
        self.transmittance = bs.statistic

class PhaseCurve(object):
    """
    Thermal phase curve.
    """
    fits_file = None
    available = []

    def __init__(self, xi, flux, name=None, channel=None, year=None,
                 renormalize=False):
        """
        Parameters
        ----------
        xi : `~numpy.ndarray`
            Times
        flux : `~numpy.ndarray`
            Flux measurements
        name : str
            Name of the host star
        channel : str
            Name of the Spitzer channel
        year : int
            Year of the observations (for disambiguating)
        renormalize : bool
            Re-normalize the phase curve such that it is represented as
            :math:`F_p/F_s`, in units of ppm
        """
        self.xi = xi[np.argsort(xi)]

        if renormalize:
            in_eclipse = np.abs(xi) < 0.1
            flux_in_eclipse = np.nanmedian(flux[in_eclipse])
            flux = 1e6 * (flux - flux_in_eclipse)

        self.flux = flux[np.argsort(xi)]
        self.name = name
        self.channel = channel
        self.year = year

    @classmethod
    def from_name(cls, name, channel, year=None):
        """
        Initialize a Filter instance from the filter name.

        Parameters
        ----------
        name : str (i.e.: "WASP-18", "KELT-9")
            Name of the host star
        channel : str (i.e.: "1" or "2")
            Name of the filter (IRAC channel number)
        year : int
            Year of the observations (when
            multiple observations are available)
        """
        if cls.fits_file is None:
            with fits.open(pc_path) as fitsfile:
                cls.fits_file = deepcopy(fitsfile)
            for hdu in cls.fits_file[1:]:
                cls.available.append("{0} (Ch {1}; year {2})"
                                     .format(hdu.header['NAME'],
                                             hdu.header['CHANNEL'],
                                             hdu.header['YEAR']))

        recarray = None
        for hdu in cls.fits_file[1:]:
            if (hdu.header['NAME'] == name and
                    hdu.header['CHANNEL'] == channel and
                    hdu.header['YEAR'] == year):
                recarray = hdu.data

        if recarray is not None:
            return cls(recarray['xi'], recarray['flux'],
                       name=name, channel=channel, year=year,
                       renormalize=False)
        else:
            raise KeyError(('Target {1} (Ch {0}, {2}) not ' +
                            'found in FITS registry, ' +
                            'which contains: {3}'
                            ).format(channel, name, year,
                                     ', '.join(sorted(cls.available))))

    def plot(self, ax=None, mask=None, **kwargs):
        """
        Plot the phase curve.

        Parameters
        ----------
        ax : `~matplotlib.axes.Axes`
            Matplotlib axis object
        kwargs : dict
            Dictionary passed to the `~matplotlib.pyplot.plot` command

        Returns
        -------
        ax : `~matplotlib.axes.Axes`
            Updated axis object
        """
        if ax is None:
            ax = plt.gca()

        if mask is None:
            mask = np.ones_like(self.xi).astype(bool)

        ax.plot(self.xi[mask], self.flux[mask], **kwargs)

        return ax

    def _add_to_fits(self, fitsfile):
        """
        Add this phase curve to a FITS archive ``fitsfile``.

        Parameters
        ----------
        fitsfile : FITS file stream
            Open FITS file stream
        """
        ra = np.recarray(len(self.xi), names=["xi", "flux"],
                         formats=['f8', 'f8'])
        ra['xi'] = self.xi
        ra['flux'] = self.flux
        header = fits.Header(dict(YEAR=self.year,
                                  CHANNEL=self.channel,
                                  NAME=self.name))
        fitsfile.append(fits.BinTableHDU(ra, header))

import numpy,from scipy,import astropy,from astropy


r"""
HPF Spectrum
---------------

A container for an HPF spectrum of :math:`M=28` total total orders :math:`m`, each with vectors for wavelength flux and uncertainty, e.g. :math:`F_m(\lambda)`.  HPF additionally has a sky fiber and optionally a Laser Frequency Comb fiber.  Our experimental API currently ignores the LFC fiber.  The sky fiber can be accessed by passing the `sky=True` kwarg when retrieving the

HPFSpectrum
##############
"""

import warnings
import logging
from muler.echelle import EchelleSpectrum, EchelleSpectrumList
import numpy as np
import astropy
from astropy.io import fits
from astropy import units as u
from astropy.wcs import WCS, FITSFixedWarning
from astropy.nddata import StdDevUncertainty
from scipy.interpolate import InterpolatedUnivariateSpline
from astropy.constants import R_jup, R_sun, G, M_jup, R_earth, c
from astropy.time import Time
import copy

log = logging.getLogger(__name__)

for category in [
    astropy.utils.exceptions.AstropyDeprecationWarning,
    FITSFixedWarning,
    RuntimeWarning,
]:
    warnings.filterwarnings("ignore", category=category)

# Convert FITS running index number to echelle order m
grating_order_offsets = {"Goldilocks": 0, "HPF": 0}  # Not implemented yet

class HPFSpectrum(EchelleSpectrum):
    r"""
    A container for HPF spectra

    Args:
        file (str): A path to a reduced HPF spectrum from Goldilocks *or* the HPF instrument team
        order (int): which spectral order to read
        cached_hdus (list) :
            A pre-loaded HDU to reduce file I/O for multiorder access.
            If provided, must give both HDUs.  Optional, default is None.
    """

    def __init__(self, *args, file=None, order=19, cached_hdus=None, **kwargs):

        self.site_name = "mcdonald"
        self.ancillary_spectra = ["sky", "lfc"]
        self.noisy_edges = (3, 2045)
        self.instrumental_resolution = 55_000.0

        if file is not None:
            if "Goldilocks" in file:
                pipeline = "Goldilocks"
            elif "Slope" in file:
                pipeline = "HPF"
            else:
                raise NameError("Cannot identify file as an HPF spectrum")
            grating_order = grating_order_offsets[pipeline] + order

            if cached_hdus is not None:
                hdus = cached_hdus[0]
            else:
                hdus = fits.open(str(file))
            hdr = hdus[0].header

            ## Target Spectrum
            lamb = hdus[7].data[order].astype(np.float64) * u.AA
            flux = hdus[1].data[order].astype(np.float64) * u.ct
            unc = hdus[4].data[order].astype(np.float64) * u.ct

            meta_dict = {
                "x_values": np.arange(0, 2048, 1, dtype=np.int),
                "pipeline": pipeline,
                "m": grating_order,
                "header": hdr,
            }

            uncertainty = StdDevUncertainty(unc)
            mask = (
                np.isnan(flux) | np.isnan(uncertainty.array) | (uncertainty.array <= 0)
            )

            super().__init__(
                spectral_axis=lamb,
                flux=flux,
                mask=mask,
                wcs=None,
                uncertainty=uncertainty,
                meta=meta_dict,
                **kwargs,
            )

            ## Sky Spectrum
            lamb = hdus[8].data[order].astype(np.float64) * u.AA
            flux = hdus[2].data[order].astype(np.float64) * u.ct
            unc = hdus[5].data[order].astype(np.float64) * u.ct
            uncertainty = StdDevUncertainty(unc)
            mask = (
                np.isnan(flux) | np.isnan(uncertainty.array) | (uncertainty.array <= 0)
            )
            sky_spectrum = HPFSpectrum(
                spectral_axis=lamb,
                flux=flux,
                mask=mask,
                wcs=None,
                uncertainty=uncertainty,
                meta=meta_dict.copy(),
                **kwargs,
            )

            ## LFC Spectrum
            lamb = hdus[9].data[order].astype(np.float64) * u.AA
            flux = hdus[3].data[order].astype(np.float64) * u.ct
            unc = hdus[6].data[order].astype(np.float64) * u.ct
            uncertainty = StdDevUncertainty(unc)
            mask = (
                np.isnan(flux) | np.isnan(uncertainty.array) | (uncertainty.array <= 0)
            )
            lfc_spectrum = HPFSpectrum(
                spectral_axis=lamb,
                flux=flux,
                mask=mask,
                wcs=None,
                uncertainty=uncertainty,
                meta=meta_dict.copy(),
                **kwargs,
            )

            ## We could optionally enable lfc and sky metadata for these referece spectra
            ## That's slightly redundant, it enables antipatterns like:
            # `spectrum.sky.lfc` rather than simply `spectrum.lfc`

            # sky_spectrum.meta["lfc"] = lfc_spectrum
            # lfc_spectrum.meta["sky"] = sky_spectrum

            sky_spectrum.meta["provenance"] = "Sky fiber"
            lfc_spectrum.meta["provenance"] = "Laser Frequency Comb"
            self.meta["provenance"] = "Target fiber"

            self.meta["sky"] = sky_spectrum
            self.meta["lfc"] = lfc_spectrum

        else:
            super().__init__(*args, **kwargs)

    @property
    def provenance(self):
        """What is the provenance of each spectrum?"""
        return self.meta["provenance"]

    @property
    def pipeline(self):
        """Which pipeline does this spectrum originate from?"""
        return self.meta["pipeline"]

    @property
    def spectrographname(self):
        """What's the name of the spectrograph?"""
        return "HPF"

    @property
    def sky(self):
        """Sky fiber spectrum stored as its own HPFSpectrum object"""
        return self.meta["sky"]

    @property
    def lfc(self):
        """Sky fiber spectrum stored as its own HPFSpectrum object"""
        return self.meta["lfc"]

    @property
    def RA(self):
        """The right ascension from header files"""
        return self.meta["header"]["RA"] * u.hourangle

    @property
    def DEC(self):
        """The declination from header files"""
        return self.meta["header"]["DEC"] * u.deg

    @property
    def astropy_time(self):
        """The astropy time based on the header"""
        mjd = self.meta["header"]["DATE-OBS"]
        return Time(mjd, format="isot", scale="utc")

    def sky_subtract(self):
        """Subtract science spectrum from sky spectrum

        Note: This operation does not wavelength shift or scale the sky spectrum

        Returns
        -------
        sky_subtractedSpec : (HPFSpectrum)
            Sky subtracted Spectrum
        """
        return self.subtract(self.sky, handle_meta="first_found")

    def blaze_divide_flats(self, flat, order=19):
        """Remove blaze function from spectrum by subtracting by flat spectrum

        Returns
        -------
        blaze corrrected spectrum using flat fields : (HPFSpectrum)

        """
        log.warning("This method is experimental and subject to change")
        new_flux = self.normalize()

        flat_wv = flat[0]
        flat_flux = flat[1]
        if len(flat) == 3:
            flat_err = flat[2]

        master_flat = flat_flux[order] / np.nanmedian(flat_flux[order])

        flat_spline = InterpolatedUnivariateSpline(
            flat_wv[order], np.nan_to_num(master_flat), k=5
        )
        interp_flat = flat_spline(self.wavelength)

        no_flat = new_flux / interp_flat

        return HPFSpectrum(
            spectral_axis=self.wavelength,
            flux=no_flat.flux,
            meta=self.meta,
            mask=self.mask,
        )

class HPFSpectrumList(EchelleSpectrumList):
    r"""
    An enhanced container for a list of HPF spectral orders

    """

    def __init__(self, *args, **kwargs):
        self.normalization_order_index = 14
        super().__init__(*args, **kwargs)

    @staticmethod
    def read(file, precache_hdus=True):
        """Read in a SpectrumList from a file

        Parameters
        ----------
        file : (str)
            A path to a reduced HPF spectrum from plp
        """
        assert ".spectra.fits" in file

        hdus = fits.open(file, memmap=False)
        cached_hdus = [hdus]

        n_orders, n_pix = hdus[7].data.shape

        list_out = []
        for i in range(n_orders):
            spec = HPFSpectrum(file=file, order=i, cached_hdus=cached_hdus)
            list_out.append(spec)
        return HPFSpectrumList(list_out)

    # def sky_subtract(self):
    #     """Sky subtract all orders
    #     """
    #     flux = copy.deepcopy(self.flux)
    #     sky = copy.deepcopy(self.sky)
    #     for i in range(len(self)):
    #         self[i] = flux[i] - sky[i]

    #     return self