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

"""
策略值网络实现
依赖TensorFlow 1.10.0版本

@author: CoderJie
@email: coderjie@outlook.com
"""

import numpy as np
import tensorflow as tf

class PolicyValueNet:
    """
    策略值网络
    """
    def __init__(self, board_width, board_height, model_file=None):
        """
        构造函数
        :param board_width: 棋盘宽度
        :param board_height: 棋盘高度
        """

        self.__board_width = board_width
        self.__board_height = board_height

        # 定义模型输入占位符, 张量形状为[?, 4, height, width]
        # None表示张量的第一维为任意长度, 第一维即代表样本数量, 就是表示模型接收任意的样本数,
        # 并且每个样本包含4个 board_height*board_width 的表格
        self.__input_data = tf.placeholder(tf.float32,
                                           shape=[None, 4, board_height, board_width])

        # 定义训练所用的状态值占位符
        self.__train_value = tf.placeholder(tf.float32,
                                            shape=[None, 1])

        # 定义训练所用的动作概率值占位符
        self.__train_prob = tf.placeholder(tf.float32,
                                           shape=[None, board_height * board_width])

        # 定义学习速度占位符
        self.__learning_rate = tf.placeholder(tf.float32)

        # 组建神经网络
        self.__build_net()

        # 创建一个会话, 并进行初始化动作
        self.__session = tf.Session()
        self.__session.run(tf.global_variables_initializer())

        # 保存器, 用于保存或者恢复模型
        self.__saver = tf.train.Saver()

        # 从文件中恢复模型
        if model_file is not None:
            self.restore_model(model_file)

    def __del__(self):
        """
        析构函数
        """
        self.__session.close()

    def train(self, batch_states, batch_probs, batch_values, lr):
        """
        训练模型
        :param batch_states: 批量状态(形状为[?, 4, height, width])
        :param batch_probs: 批量概率(形状为[?, height*width])
        :param batch_values: 批量状态值(形状为[?, 1])
        :param lr: 学习速度(浮点数)
        :return: 损失值, 熵
        """
        loss, entropy, _ = self.__session.run(
                [self.__total_loss, self.__entropy, self.__optimizer],
                feed_dict={self.__input_data: batch_states,
                           self.__train_prob: batch_probs,
                           self.__train_value: batch_values,
                           self.__learning_rate: lr})
        return loss, entropy

    def policy_value(self, board):
        """
        获取策略(动作概率), 状态值
        :param board: 棋盘对象
        :return: 策略(动作概率), 状态值
        """
        state = board.state()
        input_data = state.reshape(-1, 4, self.__board_height, self.__board_width)
        acts_prob, state_value = self.__policy_value(input_data)

        # 只获取有效动作的概率
        avl_acts = board.avl_actions
        acts_prob = list(zip(avl_acts, acts_prob[0][avl_acts]))
        return acts_prob, state_value[0][0]

    def save_model(self, path):
        """
        保存模型
        :param path: 模型路径
        """
        self.__saver.save(self.__session, path)

    def restore_model(self, path):
        """
        恢复模型
        :param path: 模型路径
        """
        self.__saver.restore(self.__session, path)

    def __build_net(self):
        """
        组建卷积神经网络
        :return:
        """
        # 转置输入数据, 输出张量形状[?, height, width, 4]
        # [0, 2, 3, 1] 0 代表原始数据第一维, 1 代表原始数据第二维, 2 代表原始数据第三维, 3 代表原始数据第四维
        # 转置后可以把棋盘数据看作是通道数为4, 大小为height*width的图片
        input_trans = tf.transpose(self.__input_data, [0, 2, 3, 1])

        # 第一层卷积层, 32个卷积核, 输出张量形状[?, height, width, 32]
        # 卷积核数量32, 卷积核大小 3*3, 卷积核的通道数就是输入数据的通道数4, 卷积步长默认为1
        # 填充方式为"same" 表示边界填充0
        # channels_last表示通道为最后一维
        # 默认加上偏置, 偏置初始值默认为0
        # 激活函数选择斜坡函数
        conv1 = tf.layers.conv2d(inputs=input_trans,
                                 filters=32,
                                 kernel_size=[3, 3],
                                 padding="same",
                                 data_format="channels_last",
                                 activation=tf.nn.relu)

        # 第二层卷积层, 64个卷积核, 输出张量形状[?, height, width, 64]
        conv2 = tf.layers.conv2d(inputs=conv1,
                                 filters=64,
                                 kernel_size=[3, 3],
                                 padding="same",
                                 data_format="channels_last",
                                 activation=tf.nn.relu)

        # 第三层卷积层, 128个卷积核, 输出张量形状[?, height, width, 128]
        conv3 = tf.layers.conv2d(inputs=conv2,
                                 filters=128,
                                 kernel_size=[3, 3],
                                 padding="same",
                                 data_format="channels_last",
                                 activation=tf.nn.relu)

        # 策略端降维, 输出张量形状[?, height, width, 4]
        policy_conv = tf.layers.conv2d(inputs=conv3,
                                       filters=4,
                                       kernel_size=[1, 1],
                                       padding="same",
                                       data_format="channels_last",
                                       activation=tf.nn.relu)

        # 策略端修改数据维度, 输出张量形状[?, 4 * height * width]
        policy_flat = tf.reshape(policy_conv,
                                 [-1, 4 * self.__board_height * self.__board_width])

        # 策略端全连接层, 输出张量形状[?, height * width]
        # units代表输出大小
        # 默认加上偏置, 偏置初始值默认为0
        # 使用log_softmax激活函数
        self.__policy_fc = tf.layers.dense(inputs=policy_flat,
                                           units=self.__board_height * self.__board_width,
                                           activation=tf.nn.log_softmax)

        # 价值端降维, 输出张量形状[?, height, width, 2]
        value_conv = tf.layers.conv2d(inputs=conv3,
                                      filters=2,
                                      kernel_size=[1, 1],
                                      padding="same",
                                      data_format="channels_last",
                                      activation=tf.nn.relu)

        # 价值端修改数据维度, 输出张量形状[?, 2 * height * width]
        value_flat = tf.reshape(value_conv,
                                [-1, 2 * self.__board_height * self.__board_width])

        # 价值端64个神经元的全连接层, 输出张量形状[?, 64]
        # 使用斜坡激活函数
        value_fc1 = tf.layers.dense(inputs=value_flat,
                                    units=64,
                                    activation=tf.nn.relu)

        # 价值端1个神经元的全连接层, 输出张量形状[?, 1]
        # 使用双曲正切激活函数, 双曲正切函数的值域是(-1, 1)
        self.__value_fc2 = tf.layers.dense(inputs=value_fc1,
                                           units=1,
                                           activation=tf.nn.tanh)

        # 定义状态值损失函数(平方差)
        value_loss = tf.losses.mean_squared_error(self.__train_value, self.__value_fc2)

        # 定义策略损失函数
        policy_loss = tf.negative(tf.reduce_mean(tf.reduce_sum(
            tf.multiply(self.__train_prob, self.__policy_fc), 1)))

        # l2正则化系数
        l2_beta = 1e-4
        # 获取所有可训练变量
        vs = tf.trainable_variables()
        # 定义正则项
        l2_penalty = l2_beta * tf.add_n(
            [tf.nn.l2_loss(v) for v in vs if 'bias' not in v.name.lower()])

        # 定义总损失函数
        self.__total_loss = value_loss + policy_loss + l2_penalty

        # 定义优化器
        self.__optimizer = tf.train.AdamOptimizer(
            learning_rate=self.__learning_rate).minimize(self.__total_loss)

        # 计算策略的熵值, 用于监控网络
        self.__entropy = tf.negative(tf.reduce_mean(
                tf.reduce_sum(tf.exp(self.__policy_fc) * self.__policy_fc, 1)))

    def __policy_value(self, states):
        """
        获取策略(动作概率), 状态值
        :param states: 批量状态二值数据, 数据形状应该为[?, 4, height, width]
        :return: 策略(形状: [?, height*width]), 状态值(形状: [?, 1])
        """
        log_probs, values = self.__session.run([self.__policy_fc, self.__value_fc2],
                                               feed_dict={self.__input_data: states})
        acts_probs = np.exp(log_probs)
        return acts_probs, values


"""
Code for parsing Spice Double-Array Files (DAF). This is a superset of CK and SPK files. Note that while this
parses Spice files, it only makes minimal use of the spice libraries (date conversion only)
"""

import struct
from spiceypy import etcal  as cspice_etcal
from spiceypy import scdecd as cspice_scdecd
from spiceypy import sct2e  as cspice_sct2e

class daf_summary:
    @staticmethod
    def make(sr,daf,buf,name):
        if daf.subtype=="SPK":
            return daf_SPKsummary(sr,daf,buf,name)
        elif daf.subtype=="CK":
            return daf_CKsummary(sr,daf,buf,name)
        else:
            raise ValueError("Unrecognized subtype %s"%daf.subtype)
    def __init__(self,sr,daf,buf,name):
        """

        :param sr: Summary record which this summary is part of
        :param daf: Double-array file which holds the header (ND, NI) and can read the file
        :param buf: Data for this summary
        """
        self.sr=sr
        self.daf=daf
        self.D=[0.0]*self.daf.ND
        self.I=[0]*self.daf.NI
        self.name=name
        for i in range(self.daf.ND):
            self.D[i]=struct.unpack(self.daf.dformat,buf[i*8:(i+1)*8])[0]
        for i in range(self.daf.NI):
            self.I[i]=struct.unpack(self.daf.iformat,buf[self.daf.ND*8+i*4:self.daf.ND*8+(i+1)*4])[0]
    def __str__(self):
        result="%s\nD:\n" % self.name
        for i,D in enumerate(self.D):
            if i!=0:
                result+=","
            result+="%20.15e"%D
        result+="\nI:\n"
        for i,I in enumerate(self.I):
            if i!=0:
                result+=","
            result+="%d"%I
        return result

class daf_SPKsummary(daf_summary):
    def __init__(self,sr,daf,buf,name):
        super().__init__(sr,daf,buf,name)
    def getet0(self): return self.D[0]
    def getet1(self): return self.D[1]
    def gettarget(self): return self.I[0]
    def getcenter(self): return self.I[1]
    def getframe (self): return self.I[2]
    def gettype  (self): return self.I[3]
    def getaddr0 (self): return self.I[4]
    def getaddr1 (self): return self.I[5]
    et0   =property(fget=getet0   ,doc="Ephemeris time of beginning of segment")
    et1   =property(fget=getet1   ,doc="Ephemeris time of end of segment")
    target=property(fget=gettarget,doc="Target NAIF code")
    center=property(fget=getcenter,doc="Center NAIF code")
    frame =property(fget=getframe ,doc="Reference frame NAIF code")
    type  =property(fget=gettype  ,doc="SPK segment type")
    addr0 =property(fget=getaddr0 ,doc="Address of first element of segment data")
    addr1 =property(fget=getaddr1 ,doc="Address of last element of segment data")
    def __str__(self):
        result=(("%s\n"+
                 "ET0:    %30.14f (%s)\n"+
                 "ET1:    %30.14f (%s)\n"+
                 "Target: %d\n"+
                 "Center: %d\n"+
                 "Frame:  %d\n"+
                 "Type:   %d\n"+
                 "Addr0:  %d\n"+
                 "Addr1:  %d") %
                (self.name,self.et0,cspice_etcal(self.et0),self.et1,cspice_etcal(self.et1),self.target,self.center,self.frame,self.type,self.addr0,self.addr1))
        return result
    def segment(self):
        return daf_SPKSegment(self,self.daf)

class daf_CKsummary(daf_summary):
    def __init__(self,sr,daf,buf,name):
        super().__init__(sr,daf,buf,name)
    def sclk(self):return self.target // 1000 + (0 if self.target % 1000 == 0 else 1)
    def getsclk0(self): return self.D[0]
    def getsclk1(self): return self.D[1]
    def getet0(self): return cspice_sct2e(self.sclk(),self.sclk0)
    def getet1(self): return cspice_sct2e(self.sclk(),self.sclk1)
    def gettarget(self): return self.I[0]
    def getframe (self): return self.I[1]
    def gettype  (self): return self.I[2]
    def getrates (self): return self.I[3]
    def getaddr0 (self): return self.I[4]
    def getaddr1 (self): return self.I[5]
    sclk0 =property(fget=getsclk0 ,doc="DP encoded sclk of beginning of segment")
    sclk1 =property(fget=getsclk1 ,doc="DP encoded sclk of end of segment")
    et  =property(fget=getet0 ,doc="ephemeris time of beginning of segment")
    et0 =property(fget=getet0 ,doc="ephemeris time of beginning of segment")
    et1 =property(fget=getet1 ,doc="ephemeris time of end of segment")
    target=property(fget=gettarget,doc="Target NAIF code")
    frame =property(fget=getframe ,doc="Reference frame NAIF code")
    type  =property(fget=gettype  ,doc="CK segment type")
    rates =property(fget=getrates ,doc="Angular Rates Flag")
    addr0 =property(fget=getaddr0 ,doc="Address of first element of segment data")
    addr1 =property(fget=getaddr1 ,doc="Address of last element of segment data")
    def __str__(self):
        result ="%s\n"%self.name
        result+="SCLK0:    %20.15e" %self.sclk0
        try:
            et=cspice_sct2e(self.sclk(),self.sclk0)
            scdecd=cspice_scdecd(self.sclk(),self.sclk0)
            result+=" (%s ET=%16.6f %s)"%(scdecd,et,cspice_etcal(et))
        except:
            pass #If there is an error, it means that the correct kernels aren't available.
        result+="\n"
        result+="SCLK1:    %20.15e" %self.sclk1
        try:
            et=cspice_sct2e(self.sclk(),self.sclk1)
            scdecd=cspice_scdecd(self.sclk(),self.sclk1)
            result+=" (%s ET=%16.6f %s)"%(scdecd,et,cspice_etcal(et))
        except:
            pass #If there is an error, it means that the correct kernels aren't available.
        result+="\n"
        result+=(("Inst:  %7d\n"+
                  "Frame: %7d\n"+
                  "Type:  %7d\n"+
                  "Rates? %7d\n"+
                  "Addr0: %7d\n"+
                  "Addr1: %7d") %
                (self.target,self.frame,self.type,self.rates,self.addr0,self.addr1))
        return result
    def segment(self):
        return daf_CKSegment(self,self.daf)

class daf_SPKSegment:
    def __init__(self,summary,daf):
        self.summary=summary
        self.daf=daf
        daf.inf.seek(8*(summary.addr0-1))
        segment_size=summary.addr1-summary.addr0+1
        format="%s%d%s"%(daf.dformat[0],segment_size,daf.dformat[1])
        self.buf=struct.unpack(format,daf.inf.read(segment_size*8))
        self.N=int(self.buf[-1]) if self.summary.type not in (15,17) else 1
        result=self.line(0)
        self.csvheader = result.csvheader
    def line(self,i):
        if self.summary.type==1:
            return daf_SPK01line(self.N,i,self.buf)
        elif self.summary.type==2:
            return daf_SPK02line(self.N,i,self.buf)
        elif self.summary.type==3:
            return daf_SPK03line(self.N,i,self.buf)
        elif self.summary.type==9:
            return daf_SPK09line(self.N,i,self.buf)
        elif self.summary.type == 13:
            return daf_SPK13line(self.N, i, self.buf)
        elif self.summary.type == 15:
            return daf_SPK15line(self.N, i, self.buf)
        elif self.summary.type == 17:
            return daf_SPK17line(self.N, i, self.buf)
        else:
            raise ValueError("Unhandled SPK type %d"%self.summary.type)
    def __iter__(self):
        for i in range(self.N):
            result = self.line(i)
            yield result
    def __str__(self):
        result=(("N: %d") %
                (self.N))
        return result

class daf_CKSegment:
    def __init__(self,summary,daf):
        self.summary=summary
        self.daf=daf
        daf.inf.seek(8*(summary.addr0-1))
        segment_size=summary.addr1-summary.addr0+1
        format="%s%d%s"%(daf.dformat[0],segment_size,daf.dformat[1])
        self.buf=struct.unpack(format,daf.inf.read(segment_size*8))
        result = self.line(0)
        self.csvheader = result.csvheader
    def getN(self):
        if self.summary.type==1:
            return int(self.buf[-1])
        elif self.summary.type==2:
            return int(1+len(self.buf)/10.01)
        elif self.summary.type==3:
            return int(self.buf[-1])
        else:
            raise ValueError("Unknown CK type %d"%self.summary.type)
    N=property(fget=getN ,doc="Number of pointing intervals in file")
    def line(self,i):
        sclk=self.summary.target//1000+(0 if self.summary.target%1000==0 else 1)
        if self.summary.type==1:
            return daf_CK01line(sclk,self.N,i,self.buf,self.summary.rates)
        elif self.summary.type==2:
            return daf_CK02line(sclk,self.N,i,self.buf)
        elif self.summary.type == 3:
            return daf_CK03line(sclk,self.N,i,self.buf,self.summary.rates)
        else:
            raise ValueError("Unknown CK type %d" % self.summary.type)
    def __iter__(self):
        for i in range(self.N):
            result = self.line(i)
            yield result
    def __str__(self):
        result=(("N: %d") %
                (self.N))
        return result

class daf_SPKStateline:
    def __init__(self,N,i,buf):
        self.x   =buf[i*6+0]
        self.y   =buf[i*6+1]
        self.z   =buf[i*6+2]
        self.dxdt=buf[i*6+3]
        self.dydt=buf[i*6+4]
        self.dzdt=buf[i*6+5]
        self.et=buf[N*6+i]
        self.csvheader="ET,ETCAL*,x,y,z,dxdt,dydt,dzdt"
    def __str__(self):
        result="%30.14f,%s,%25.15e,%25.15e,%25.15e,%25.15e,%25.15e,%25.15e"%(self.et,cspice_etcal(self.et),self.x,self.y,self.z,self.dxdt,self.dydt,self.dzdt)
        return result

class daf_SPK09line(daf_SPKStateline):
    def __init__(self,N,i,buf):
        super().__init__(N,i,buf)

class daf_SPK13line(daf_SPKStateline):
    def __init__(self,N,i,buf):
        super().__init__(N,i,buf)

class daf_SPK02line:
    def __init__(self,N,i,buf):
        self.RSIZE=int(buf[-2])
        self.asize=(self.RSIZE-2)//3 #Number of elements in each component array
        self.INTLEN=buf[-3]
        self.INIT=buf[-4]
        self.mid    =buf[i*self.RSIZE+0]
        self.radius=buf[i*self.RSIZE+1]
        self.et=self.mid-self.radius
        self.x     =buf[i*self.RSIZE+2+self.asize*0:i*self.RSIZE+2+self.asize*1]
        self.y     =buf[i*self.RSIZE+2+self.asize*1:i*self.RSIZE+2+self.asize*2]
        self.z     =buf[i*self.RSIZE+2+self.asize*2:i*self.RSIZE+2+self.asize*3]
        self.csvheader="ET*,ETCAL*,ETEND*,ETCALEND*,MID,MIDCAL*,RADIUS"
        for i in range(self.asize):
            self.csvheader+=",x[%02d]"%i
        for i in range(self.asize):
            self.csvheader+=",y[%02d]" % i
        for i in range(self.asize):
            self.csvheader+=",z[%02d]" % i
    def __str__(self):
        result="%30.14f,%s,%30.14f,%s,%30.14f,%s,%30.14f"%(self.et,cspice_etcal(self.et),self.mid+self.radius,cspice_etcal(self.mid+self.radius),self.mid,cspice_etcal(self.mid),self.radius)
        for x in self.x:
            result+=",%25.15e"%x
        for y in self.y:
            result+=",%25.15e"%y
        for z in self.z:
            result+=",%25.15e"%z
        return result
    def eval(self,ET):
        pass

class daf_SPK03line:
    def __init__(self,N,i,buf):
        self.RSIZE=int(buf[-2])
        self.asize=(self.RSIZE-2)//6 #Number of elements in each component array
        self.INTLEN=buf[-3]
        self.INIT=buf[-4]
        self.mid    =buf[i*self.RSIZE+0]
        self.radius=buf[i*self.RSIZE+1]
        self.et=self.mid-self.radius
        self.x      =buf[i*self.RSIZE+2+self.asize*0:i*self.RSIZE+2+self.asize*1]
        self.y      =buf[i*self.RSIZE+2+self.asize*1:i*self.RSIZE+2+self.asize*2]
        self.z      =buf[i*self.RSIZE+2+self.asize*2:i*self.RSIZE+2+self.asize*3]
        self.dx     =buf[i*self.RSIZE+2+self.asize*3:i*self.RSIZE+2+self.asize*4]
        self.dy     =buf[i*self.RSIZE+2+self.asize*4:i*self.RSIZE+2+self.asize*5]
        self.dz     =buf[i*self.RSIZE+2+self.asize*5:i*self.RSIZE+2+self.asize*6]
        self.csvheader="ET*,ETCAL*,ETEND*,ETCALEND*,MID,MIDCAL*,RADIUS"
        for i in range(self.asize):
            self.csvheader+=",x[%02d]"%i
        for i in range(self.asize):
            self.csvheader+=",y[%02d]" % i
        for i in range(self.asize):
            self.csvheader+=",z[%02d]" % i
        for i in range(self.asize):
            self.csvheader+=",dx[%02d]"%i
        for i in range(self.asize):
            self.csvheader+=",dy[%02d]" % i
        for i in range(self.asize):
            self.csvheader+=",dz[%02d]" % i
    def __str__(self):
        result="%30.14f,%s,%30.14f,%s,%30.14f,%s,%30.14f"%(self.et,cspice_etcal(self.et),self.mid+self.radius,cspice_etcal(self.mid+self.radius),self.mid,cspice_etcal(self.mid),self.radius)
        for x in self.x:
            result+=",%25.15e"%x
        for y in self.y:
            result+=",%25.15e"%y
        for z in self.z:
            result+=",%25.15e"%z
        for dx in self.dx:
            result+=",%25.15e"%dx
        for dy in self.dy:
            result+=",%25.15e"%dy
        for dz in self.dz:
            result+=",%25.15e"%dz
        return result
    def eval(self,ET):
        pass

class daf_SPK15line:
    def __init__(self,N,i,buf):
        self.et   =buf[0]
        self.polex    =buf[1]
        self.poley    =buf[2]
        self.polez    =buf[3]
        self.perix    =buf[4]
        self.periy    =buf[5]
        self.periz    =buf[6]
        self.l=buf[7]
        self.e =buf[8]
        self.j2flag=buf[9]
        self.centx=buf[10]
        self.centy=buf[11]
        self.centz=buf[12]
        self.centgm=buf[13]
        self.centj2=buf[14]
        self.centre=buf[15]
        self.csvheader="ETPERI,ETPERICAL*,polex,poley,polez,perix,periy,periz,l,e,j2flag,centx,centy,centz,centgm,centj2,centre"
    def __str__(self):
        result="%30.14f,%s,%.15f,%.15f,%.15f,%.15f,%.15f,%.15f,%.15e,%.15e,%1.0f,%17.14f,%17.14f,%17.14f,%.15e,%.15e,%.15e"%(
            self.et,cspice_etcal(self.et),
            self.polex,self.poley,self.polez,self.perix,self.periy,self.periz,
            self.l,self.e,self.j2flag,
            self.centx, self.centy, self.centz,
            self.centgm,self.centj2,self.centre
        )
        return result
    def eval(self,ET):
        pass

class daf_SPK17line:
    def __init__(self,N,i,buf):
        self.et   =buf[0]
        self.a    =buf[1]
        self.h    =buf[2]
        self.k    =buf[3]
        self.L    =buf[4]
        self.p    =buf[5]
        self.q    =buf[6]
        self.omegabardot=buf[7]
        self.Ldot =buf[8]
        self.Omegadot=buf[9]
        self.polera=buf[10]
        self.poledec=buf[11]
        self.csvheader="ETPERI,ETPERICAL*,a,h,k,L,p,q,omegabardot,Ldot,Omegadot,polera,poledec"
    def __str__(self):
        result="%30.14f,%s,%30.14f,%30.14f,%30.14f,%30.14f,%30.14f,%30.14f,%30.14f,%30.14f,%30.14f,%30.14f,%30.14f"%(
            self.et,cspice_etcal(self.et),
            self.a,self.h,self.k,self.L,self.p,self.q,
            self.omegabardot,self.Ldot,self.Omegadot,
            self.polera,self.poledec
        )
        return result
    def eval(self,ET):
        pass

class daf_SPK01line:
    n_elements=71
    def __init__(self,N,i,buf):
        self.e   =buf[i*self.n_elements:(i+1)*self.n_elements]
        self.et=buf[N*self.n_elements+i]
        import numpy as np
        self.FD=np.zeros(15)*float('nan')

        # initialize arrays for spke01
        self.G = np.zeros(15)
        self.REFPOS = np.zeros(3)
        self.REFVEL = np.zeros(3)
        self.KQ = np.array([0, 0, 0])
        self.FC = np.zeros(15)
        self.FC[0] = 1.0
        self.WC = np.zeros(13)
        self.W = np.zeros(17)

        self.spke01(self.e[0],self.e)
        self.csvheader= "ET,ETCAL*,TL"
        for i in range(15):
            self.csvheader += ",G[%02d]" % (i)
        self.csvheader += ",X,DXDT,Y,DYDT,Z,DZDT"
        for i in range(3):
            for j in range(15):
                self.csvheader += ",DT[%02d%02d]" % (i, j)
                self.csvheader += ",KQMAX1,KQ[0],KQ[1],KQ[2]"

    #        self.eval(self.TL)
    def __str__(self):
        result="%30.14f,%s" %(self.et,cspice_etcal(self.et))
        for i in range(self.n_elements):
            result+=",%25.15e"%self.e[i]
        return result
    """
    def getTL(self): return self.e[0]
    def getG (self): return self.e[1:15+1]
    def getREFPOS(self): return self.e[16:22:2]
    def getREFVEL(self): return self.e[17:23:2]
    def getDT    (self): return self.e[22:22+45]
    def getKQMAX1(self): return int(self.e[67])
    def getKQ    (self): return self.e[68:71]
    TL    =property(fget=getTL    ,doc="Ephemeris time of end of segment")
    G     =property(fget=getG     ,doc="Stepsize function vector")
    REFPOS=property(fget=getREFPOS,doc="Reference position vector")
    REFVEL=property(fget=getREFVEL,doc="Reference velocity vector")
    DT    =property(fget=getDT    ,doc="Modified divided difference arrays")
    KQMAX1=property(fget=getKQMAX1,doc="Maximum integration order plus 1")
    KQ    =property(fget=getKQ    ,doc="Integration order array")
    """
    def getFC(self,J):
        print("Accessing element %d, value %f"%(J,self.FD[J]))
        return self.FD[J]
    def setFC(self,J,val):
        print("Writing element %d, was value %f, will be value %f"%(J,self.FD[J],val))
        self.FD[J]=val
    def eval(self,ET):
        import numpy as np
        STATE=np.zeros(6)
        DELTA = ET - self.TL
        TP = DELTA
        MQ2 = self.KQMAX1 - 2
        KS = self.KQMAX1 - 1
        self.setFC(0,1.0)
        WC = np.zeros(13)
        W  = np.zeros(17)

        #     This is clearly collecting some kind of coefficients.
        #     The problem is that we have no idea what they are...
        #
        #     The G coefficients are supposed to be some kind of step size
        #     vector.
        #
        #     TP starts out as the delta t between the request time
        #     and the time for which we last had a state in the MDL file.
        #     We then change it from DELTA  by the components of the stepsize
        #     vector G.
        for J in range(1, MQ2 + 1):
            self.setFC(J,TP / self.G[J-1])
            WC[J-1] = DELTA / self.G[J-1]
            TP = DELTA + self.G[J-1]
        #     Collect KQMAX1 reciprocals.
        for J in range(1, self.KQMAX1 + 1):
            W[J-1] = 1.0 / float(J)

        #     Compute the W(K) terms needed for the position interpolation
        #     (Note,  it is assumed throughout this routine that KS, which
        #     starts out as KQMAX1-1 (the ``maximum integration'')
        #     is at least 2.
        JX = 0
        KS1 = KS - 1
        while KS >= 2:
            JX = JX + 1
            for J in range(1, JX + 1):
                W[J + KS - 1] = self.getFC(J) * W[J + KS1 - 1] - WC[J - 1] * W[J + KS - 1]
            KS = KS1
            KS1 = KS1 - 1

        #     Perform position interpolation: (Note that KS = 1 right now.
        #     We don't know much more than that.)
        for I in range(3):
            KQQ = int(self.KQ[I])
            SUM = 0.0
            for J in range(KQQ, 0, -1):
                SUM = SUM + self.DT[J - 1+I*15] * W[J + KS - 1]
            STATE[I] = self.REFPOS[I] + DELTA * (self.REFVEL[I] + DELTA * SUM)

        #     Again we need to compute the W(K) coefficients that are
        #     going to be used in the velocity interpolation.
        #     (Note, at this point, KS = 1, KS1 = 0.)
        for J in range(1, JX + 1):
            W[J + KS - 1] = self.getFC(J) * W[J + KS1 - 1] - WC[J - 1] * W[J + KS - 1]
        KS = KS - 1

        #     Perform velocity interpolation:
        for I in range(1, 3 + 1):
            KQQ = int(self.KQ[I])
            SUM = 0.0

            for J in range(KQQ, 0, -1):
                SUM = SUM + self.DT[J - 1+I*15] * W[J + KS - 1]

            STATE[I + 3] = self.REFVEL[I] + DELTA * SUM

        #     That's all folks.  We don't know why we did anything, but
        #     at least we can tell structurally what we did.
        return STATE
    def test(self,et0,N=100):
        """
        Compare the calculation of the state in pure python to that provided by spiceypy (which is an interface
        to the cspice library). This should be near bit-perfect.
        :param ET0: Time to begin at. Since the record only contains the end time, we need a beginning time. The
                    beginning time *might* be encoded in the DeltaT's, but this isn't known for sure yet
        :param N: Number of points to evaluate at. If you pass 100, you will get the point at the beginning,
                  99 points in between, and the point at the end, so a total of 101 points.
        :return: Sum of squares of differences of position components. Ideally should be zero, should be only a couple
                 of ulps.
        """
        from spiceypy import spkezr as cspice_spkezr
        import numpy as np
        et1=self.et
        result=0.0
        for i in range(N):
            et=et0+i*(et1-et0)/N
            #et=et1
            ref=cspice_spkezr("-189",et,"J2000","NONE","399")[0][0:3]
            test=np.array(self.spke01(et,self.e)[0:3])
            result+=np.sum((ref-test)**2)
        return result
    def spke01(self, ET, RECORD):
        import numpy as np
        """Compute position and velocity from a Modified Difference Array record

        Inputs:
            ET: Epoch time to evaluate position and velocity (seconds since J2000)
            RECORD: A record of Modified Difference Array
        Returns: STATE
            STATE: A numpy array which contains position and velocity
        """

        # This method has been translated from SPKE01 of SPICE Toolkit and
        # modified by Shushi Uetsuki.
        #
        # SPICE Toolkit for FORTRAN : http://naif.jpl.nasa.gov/naif/toolkit_FORTRAN.html
        # SPK Required Reading : http://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/req/spk.html
        #
        # Original FORTRAN code uses 'SAVE' directive, and it means all variable
        # should be saved for next call.  So i decided to make almost all
        # variables to be instance variable.  Some of them are initialized in
        # __init__ method.

        STATE = np.zeros(6)

        #     Variable  I/O  Description
        #     --------  ---  --------------------------------------------------
        #     ET         I   Target epoch.
        #     RECORD     I   Data record.
        #     STATE      O   State (position and velocity).
        #
        # $ Detailed_Input
        #
        #     ET          is a target epoch, at which a state vector is to
        #                 be computed.
        #
        #     RECORD      is a data record which, when evaluated at epoch ET,
        #                 will give the state (position and velocity) of some
        #                 body, relative to some center, in some inertial
        #                 reference frame.
        #
        # $ Detailed_Output
        #
        #     STATE       is the state. Units are km and km/sec.
        #
        # $ Parameters
        #
        #     None.
        #
        # $ Exceptions
        #
        #     None.
        #
        # $ Files
        #
        #     None.
        #
        # $ Particulars
        #
        #     The exact format and structure of type 1 (difference lines)
        #     segments are described in the SPK Required Reading file.
        #
        #     Difference lines (DL's) are generated by JPL navigation
        #     system programs P and PV. Each data record is equivalent
        #     to the (slightly rearranged) 'P' portion of a NAVIO PV file
        #     data record.
        #
        #     SPKE01 is a specialized version of Fred Krogh's subroutine DAINT.
        #     Only the calling sequence has been changed.
        #
        #     Because the original version was undocumented, only Fred
        #     knows how this really works.
        #
        # $ Examples
        #
        #     None.
        #
        # $ Restrictions
        #
        #     Unknown.
        #
        # $ Literature_References
        #
        #     NAIF Document 168.0, "S- and P- Kernel (SPK) Specification and
        #     User's Guide"
        #
        # $ Author_and_Institution
        #
        #     F.T. Krogh      (JPL)
        #     I.M. Underwood  (JPL)
        #
        # $ Version
        #
        # -    SPICELIB Version 1.1.0, 14-FEB-1997 (WLT)
        #
        #        The goto's were removed and loop and if structures
        #        revealed.  We still don't know exactly what's going
        #        on, but at least the bones of this routine have been
        #        cleaned off and are ready for assembly. (WLT)
        #
        # -    SPICELIB Version 1.0.4, 30-OCT-1996 (WLT)
        #
        #        Removed redundant SAVE statements from the declaration
        #        section.  Thanks to Steve Schlaifer for finding this
        #        error.
        #
        # -    SPICELIB Version 1.0.3, 10-MAR-1992 (WLT)
        #
        #        Comment section for permuted index source lines was added
        #        following the header.
        #
        # -    SPICELIB Version 1.0.2, 23-AUG-1991 (HAN)
        #
        #        SPK01 was removed from the Required_Reading section of the
        #        header. The information in the SPK01 Required Reading file
        #        is now part of the SPK Required Reading file.
        #
        # -    SPICELIB Version 1.0.1, 22-MAR-1990 (HAN)
        #
        #        Literature references added to the header.
        #
        # -    SPICELIB Version 1.0.0, 31-JAN-1990 (IMU) (FTK)
        #
        # -&
        #
        # $ Index_Entries
        #
        #     evaluate type_1 spk segment
        #
        # -&

        #
        #     Unpack the contents of the MDA array.
        #
        #        Name    Dimension  Description
        #        ------  ---------  -------------------------------
        #        TL              1  Final epoch of record
        #        G              15  Stepsize function vector
        #        REFPOS          3  Reference position vector
        #        REFVEL          3  Reference velocity vector
        #        DT         15,NTE  Modified divided difference arrays
        #        KQMAX1          1  Maximum integration order plus 1
        #        KQ            NTE  Integration order array
        #
        #     For our purposes, NTE is always 3.
        #
        self.TL = RECORD[0]
        self.G = RECORD[1:1 + 15]
        #
        #     Collect the reference position and velocity.
        #
        self.REFPOS[0] = RECORD[16]
        self.REFVEL[0] = RECORD[17]

        self.REFPOS[1] = RECORD[18]
        self.REFVEL[1] = RECORD[19]

        self.REFPOS[2] = RECORD[20]
        self.REFVEL[2] = RECORD[21]

        self.DT = np.reshape(RECORD[22:22 + 45], (15, 3), order='F')

        self.KQMAX1 = int(RECORD[67])
        self.KQ[0] = int(RECORD[68])
        self.KQ[1] = int(RECORD[69])
        self.KQ[2] = int(RECORD[70])
        #
        #     Next we set up for the computation of the various differences
        #
        self.DELTA = ET - self.TL
        self.TP = self.DELTA
        self.MQ2 = self.KQMAX1 - 2
        self.KS = self.KQMAX1 - 1
        #
        #     This is clearly collecting some kind of coefficients.
        #     The problem is that we have no idea what they are...
        #
        #     The G coefficients are supposed to be some kind of step size
        #     vector.
        #
        #     TP starts out as the delta t between the request time
        #     and the time for which we last had a state in the MDL file.
        #     We then change it from DELTA  by the components of the stepsize
        #     vector G.
        #
        for J in range(1, self.MQ2 + 1):
            self.FC[J] = self.TP / self.G[J - 1]
            self.WC[J - 1] = self.DELTA / self.G[J - 1]
            self.TP = self.DELTA + self.G[J - 1]
        #
        #     Collect KQMAX1 reciprocals.
        #
        for J in range(1, self.KQMAX1 + 1):
            self.W[J - 1] = 1.0 / float(J)
        #
        #     Compute the W(K) terms needed for the position interpolation
        #     (Note,  it is assumed throughout this routine that KS, which
        #     starts out as KQMAX1-1 (the ``maximum integration'')
        #     is at least 2.
        #
        self.JX = 0
        self.KS1 = self.KS - 1

        while self.KS >= 2:

            self.JX = self.JX + 1

            for J in range(1, self.JX + 1):
                self.W[J + self.KS - 1] = self.FC[J] * self.W[J + self.KS1 - 1] - self.WC[J - 1] * self.W[
                    J + self.KS - 1]

            self.KS = self.KS1
            self.KS1 = self.KS1 - 1
        #
        #     Perform position interpolation: (Note that KS = 1 right now.
        #     We don't know much more than that.)
        #
        for I in range(1, 3 + 1):

            self.KQQ = self.KQ[I - 1]
            self.SUM = 0.0

            for J in range(self.KQQ, 0, -1):
                self.SUM = self.SUM + self.DT[J - 1, I - 1] * self.W[J + self.KS - 1]

            #x=x0+delta*dx0+delta**2*sum
            STATE[I - 1] = self.REFPOS[I - 1] + self.DELTA * (self.REFVEL[I - 1] + self.DELTA * self.SUM)
        #
        #     Again we need to compute the W(K) coefficients that are
        #     going to be used in the velocity interpolation.
        #     (Note, at this point, KS = 1, KS1 = 0.)
        #
        for J in range(1, self.JX + 1):
            self.W[J + self.KS - 1] = self.FC[J] * self.W[J + self.KS1 - 1] - self.WC[J - 1] * self.W[J + self.KS - 1]

        self.KS = self.KS - 1

        #
        #     Perform velocity interpolation:
        #
        for I in range(1, 3 + 1):
            self.KQQ = self.KQ[I - 1]
            self.SUM = 0.0

            for J in range(self.KQQ, 0, -1):
                self.SUM = self.SUM + self.DT[J - 1, I - 1] * self.W[J + self.KS - 1]

            STATE[I + 3 - 1] = self.REFVEL[I - 1] + self.DELTA * self.SUM

        #
        #     That's all folks.  We don't know why we did anything, but
        #     at least we can tell structurally what we did.
        #

        return STATE

class daf_CK01line:
    def __init__(self,clknum,N,i,buf,rates):
        self.clknum=clknum
        self.rates=rates
        self.csvheader="SCLK,SCLK*,ET*,ETCAL*,q0,q1,q2,q3"
        if self.rates:
            self.n_elements = 7
            self.w = buf[i * self.n_elements+4:(i + 1) * self.n_elements]
            self.csvheader+=",wx,wy,wz"
        else:
            self.n_elements=4
        self.q = buf[i * self.n_elements:i * self.n_elements+4]

        self.sclk=buf[N*self.n_elements+i]
        self.et =cspice_sct2e(self.clknum,self.sclk)
    def __str__(self):
        result="%25.15e"%self.sclk
        try:
            et=cspice_sct2e(self.clknum, self.sclk)
            result+=",%s,%16.6f,%s" %(cspice_scdecd(self.clknum,self.sclk),et,cspice_etcal(et))
        except:
            result+=",,,"
        for i in range(4):
            result+=",%25.15e"%self.q[i]
        if self.rates:
            for i in range(3):
                result += ",%25.15e" % self.w[i]
        else:
            result+=",,,"
        return result

class daf_CK02line:
    def __init__(self,clknum,N,i,buf):
        self.clknum=clknum
        self.n_elements = 8
        self.w = buf[i * self.n_elements+4:i * self.n_elements+7]
        self.q = buf[i * self.n_elements:i * self.n_elements+4]
        self.rate=buf[i*self.n_elements+7]
        self.sclk=buf[N*self.n_elements+i]
        self.sclk1=buf[N*self.n_elements+N+i]
        self.et =cspice_sct2e(self.clknum,self.sclk)
        self.et1 =cspice_sct2e(self.clknum,self.sclk1)
        self.csvheader="SCLK,SCLK*,SCLK1,SCLK1*,ET*,ETCAL*,ET1*,ETCAL1*,q0,q1,q2,q3,wx,wy,wz,rate"
    def __str__(self):
        result="%25.15e"%self.sclk
        result+=",%s,%25.15e,%s,%16.6f,%s,%16.6f,%s" %(cspice_scdecd(self.clknum,self.sclk),self.sclk1,cspice_scdecd(self.clknum,self.sclk1),self.et,cspice_etcal(self.et),self.et1,cspice_etcal(self.et1))
        for i in range(4):
            result+=",%25.15e"%self.q[i]
        for i in range(3):
            result += ",%25.15e" % self.w[i]
        result += ",%25.15e" % self.rate
        return result

class daf_CK03line:
    def __init__(self,clknum,N,i,buf,rates):
        self.clknum=clknum
        self.rates=rates
        self.csvheader="SCLK,SCLK*,ET*,ETCAL*,q0,q1,q2,q3"
        if self.rates:
            self.n_elements = 7
            self.w = buf[i * self.n_elements+4:(i + 1) * self.n_elements]
            self.csvheader+=",wx,wy,wz"
        else:
            self.n_elements=4
        self.q = buf[i * self.n_elements:i * self.n_elements+4]

        self.sclk=buf[N*self.n_elements+i]
        self.et =cspice_sct2e(self.clknum,self.sclk)
    def __str__(self):
        result="%25.15e"%self.sclk
        try:
            et=cspice_sct2e(self.clknum, self.sclk)
            result+=",%s,%16.6f,%s" %(cspice_scdecd(self.clknum,self.sclk),et,cspice_etcal(et))
        except:
            result+=",,,"
        for i in range(4):
            result+=",%25.15e"%self.q[i]
        if self.rates:
            for i in range(3):
                result += ",%25.15e" % self.w[i]
        else:
            result+=",,,"
        return result

class daf_summary_record:
    def __init__(self,parent):
        """
        Read a record from a file-like object opened in binary mode
        :param parent: object of class daf
        """
        self.parent=parent
        self.dformat=self.parent.dformat
        self.buf=self.read(1024) #Read the whole block first
        self.namebuf=self.read(1024) #Read the name record next
        self.NEXT=int(struct.unpack(self.dformat,self.buf[ 0: 0+8])[0]) #Address of next summary record in file (zero if last)
        self.PREV=int(struct.unpack(self.dformat,self.buf[ 8: 8+8])[0]) #Address of prev summary record in file (zero if first)
        self.NSUM=int(struct.unpack(self.dformat,self.buf[16:16+8])[0]) #Number of summaries actually in this summary record
        self.SS=self.parent.SS
    def read(self,len):
        return self.parent.inf.read(len)
    def tell(self):
        return self.parent.inf.tell()
    def __iter__(self):
        for i in range(self.NSUM):
            result = daf_summary.make(self,self.parent,self.buf[8*(3+i*self.SS):8*(3+(i+1)*self.SS)],str(self.namebuf[i*self.parent.NC:(i+1)*self.parent.NC],encoding="cp437").strip(" \t\n\r\0"))
            yield result
    def __str__(self):
        result=(("NEXT: %d\n" +
                 "PREV: %s\n" +
                 "NSUM: %d") %
                (self.NEXT,self.PREV,self.NSUM))
        return result

class double_array_file:
    def __init__(self,inf):
        """
        Read a double-array file. This class reads the DAF header itself, and delegates reading the
        records, comments, and data. It therefore has fields representing the header, and lists
        of other class objects representing the data.
        :param inf: string filename or file-like object
        """
        try:
            self.inf=open(inf,"rb")
            self.needToClose=True
        except TypeError:
            self.inf=inf
            self.needToClose=False
        buf=self.inf.read(1024) #Read the whole header record first
        self.LOCIDW=str(buf[0:0+8],encoding='ASCII')
        if self.LOCIDW[0:4]=="NAIF":
            print("Old kernel type - guessing ")
            self.subtype="CK"
        elif self.LOCIDW[0:4]!="DAF/":
            raise ValueError("Not a DAF: Header is %s"%self.LOCIDW)
        else:
            self.subtype=self.LOCIDW[4:8].strip(" \t\n\r\0")
        self.LOCFMT=str(buf[88:88+8],encoding='ASCII') #IEEE format string
        if self.LOCFMT=="LTL-IEEE":
            self.endian="<"
        elif self.LOCFMT=="BIG-IEEE":
            self.endian=">"
        else:
            raise ValueError("Unknown double-precision format %s"%self.LOCFMT)
        self.iformat=self.endian+"i"
        self.dformat=self.endian+"d"
        self.ND=struct.unpack(self.iformat,buf[ 8: 8+4])[0] #Address  8, number of components in each array summary
        self.NI=struct.unpack(self.iformat,buf[12:12+4])[0] #Address 12, number of integer components in each array summary
        self.LOCIFN=str(buf[16:16+60],encoding='ASCII')     #Address 16, internal name or description of array file
        self.FWARD=struct.unpack(self.iformat,buf[76:76+4])[0] #Address 76, record number of initial summary in the file
        self.BWARD=struct.unpack(self.iformat,buf[80:80+4])[0] #Address 80, record number of final summary in the file
        self.FREE=struct.unpack(self.iformat,buf[84:84+4])[0]  #Address 84, first free address in file
        self.FTPSTR=str(buf[700:700+28],encoding='cp437') #Address 700, FTP validation string
        self.NC = 8 * (self.ND + (self.NI + 1) // 2)
        self.SS = self.ND + (self.NI + 1) // 2  # Size of a summary
        self.NS = 125 // self.SS  # Maximum number of summaries per summary record
    def __enter__(self):
        return self
    def __exit__(self,type,value,traceback):
        if self.needToClose:
            self.inf.close()
    def __str__(self):
        result=(("LOCIDW: %s\n"+
                 "ND:     %d\n"+
                 "NI:     %d\n" +
                 "NC*:    %d\n"+
                 "SS*:    %d\n" +
                 "NS*:    %d\n"+
                 "LOCIFN: %s\n" +
                 "FWARD:  %d\n" +
                 "BWARD:  %d\n" +
                 "FREE:   %d\n" +
                 "LOCFMT: %s") %
                (self.LOCIDW,self.ND,self.NI,self.NC,self.SS,self.NS,self.LOCIFN,self.FWARD,self.BWARD,self.FREE,self.LOCFMT))
        return result
    def __iter__(self):
        #Seek to the first summary record
        self.inf.seek((self.FWARD-1)*1024) #One for one-based, one for header record at beginning of file
        done=False
        while not done:
            result=daf_summary_record(self)
            yield result
            if result.NEXT==0:
                done=True
            else:
                self.inf.seek((result.NEXT-1)*1024)
    def dump(self,dump_lines=True,dump_comments=True):
        print(str(self))
        if dump_comments:
            for line in self.comments():
                print('"'+line+'"')
        for i,sr in enumerate(self):
            print("Summary record %d"%i)
            print(str(sr))
            for j,sum in enumerate(sr):
                print("Summary %d" % j)
                print(str(sum))
                print(str(sum.segment()))
                if dump_lines:
                    print("i*,dt," + sum.segment().csvheader)
                    et0=sum.et0
                    for k,line in enumerate(sum.segment()):
                        dt=line.et-et0
                        print(k,",%15.6f,"%dt,str(line))
#                        tval=line.test(et0)
#                        if True: #(tval>1e-10):
#                            print(line.et-et0,",",tval)
                        et0=line.et
    def comments(self):
        block=1
        done=False
        line=""
        while not done:
            self.inf.seek(block*1024)
            bb=self.inf.read(1000)
            for b in bb:
                if b==0:
                    yield line
                    line=""
                elif b==4:
                    yield line
                    return
                else:
                    line+=chr(b)
            block+=1

def spkdump():
    import argparse
    parser=argparse.ArgumentParser(description="Translate a Spice Double-Array File (DAF) into plain text")
    parser.add_argument('-x','--records',action='store_true',help="Dump all data records. May result in a very large amount of output!")
    parser.add_argument('-r','--comment',action='store_true',help="Dump file comments.")
    parser.add_argument('-a','--all',action='store_true',help="Dump everything that can be dumped")
    parser.add_argument('inf',nargs=1,help="Kernel to translate. Must be a bc or bsp kernel.")
    parser.add_argument('kernel',nargs='*',help="Kernels to furnish. Typically these are sclk and/or lsk files, needed to interpret bc timestamps")
    result=parser.parse_args()
    dump_lines=result.records or result.all
    dump_comments=result.comment or result.all
    inf=result.inf[0]

    from spiceypy import furnsh as cspice_furnsh
    for furnsh in result.kernel:
        cspice_furnsh(furnsh)
    with double_array_file(inf) as in_daf:
        cspice_furnsh(inf) #Just to run the test
        in_daf.dump(dump_lines=dump_lines,dump_comments=dump_comments)
#    with double_array_file("../../Data/spice/insight/insight_nom_2016e09o_edl_v1.bsp") as in_daf:
#        in_daf.dump(dump_lines)
#    with double_array_file("../../Data/spice/phoenix/phx_edl_rec_traj.bsp") as in_daf:
#        in_daf.dump(dump_lines)

if __name__=="__main__":
    spkdump()


import logging

import cffi

import numpy as np

from .matrix import CSR

_logger = logging.getLogger(__name__)

__mkl_syrk_defs = '''
typedef void* sparse_matrix_t;
struct matrix_descr {
    int  type;
    int  mode;
    int  diag;
};

int mkl_sparse_d_create_csr(sparse_matrix_t *A, int indexing, int rows, int cols,
                            int *rows_start, int *rows_end, int *col_indx, double *values);
int mkl_sparse_d_export_csr(const sparse_matrix_t source, int *indexing, int *rows, int *cols,
                            int **rows_start, int **rows_end, int **col_indx, double **values);
int mkl_sparse_order(sparse_matrix_t A);
int mkl_sparse_destroy(sparse_matrix_t A);

int mkl_sparse_syrk (int operation, const sparse_matrix_t A, sparse_matrix_t *C);
int mkl_sparse_d_mv (int operation, double alpha,
                     const sparse_matrix_t A, struct matrix_descr descr,
                     const double *x, double beta, double *y);
'''
_logger.debug('initializing CFFI interface')
_mkl_ffi = cffi.FFI()
_mkl_ffi.cdef(__mkl_syrk_defs)
try:
    _logger.debug('importing MKL')
    _mkl_lib = _mkl_ffi.dlopen('mkl_rt')
    _logger.info('Loaded MKL')
except OSError:
    _logger.info('Cannot load MKL')
    _mkl_lib = None

def _mkl_check_return(rv, call='<unknown>'):
    if rv:
        raise RuntimeError('MKL call {} failed with code {}'.format(call, rv))

def _mkl_basic_descr():
    desc = _mkl_ffi.new('struct matrix_descr*')
    desc.type = 20  # general matrix
    desc.mode = 0
    desc.diag = 0
    return desc

class SparseM:
    """
    Class encapsulating an MKL sparse matrix handle.
    """

    def __init__(self):
        self.h_ptr = _mkl_ffi.new('sparse_matrix_t*')

    @classmethod
    def from_csr(cls, csr):
        """
        Create an MKL sparse matrix from a LensKit CSR matrix.

        Args:
            csr(CSR): the input matrix.

        Returns:
            SparseM: a sparse matrix handle for the CSR matrix.
        """
        sp = np.require(csr.rowptrs, np.intc, 'C')
        ep = np.require(csr.rowptrs[1:], np.intc, 'C')
        cols = np.require(csr.colinds, np.intc, 'C')
        vals = np.require(csr.values, np.float_, 'C')

        m = SparseM()
        _sp = _mkl_ffi.cast('int*', sp.ctypes.data)
        _ep = _mkl_ffi.cast('int*', ep.ctypes.data)
        _cols = _mkl_ffi.cast('int*', cols.ctypes.data)
        _vals = _mkl_ffi.cast('double*', vals.ctypes.data)
        rv = _mkl_lib.mkl_sparse_d_create_csr(m.h_ptr, 0, csr.nrows, csr.ncols,
                                              _sp, _ep, _cols, _vals)
        _mkl_check_return(rv, 'mkl_sparse_d_create_csr')

        return m

    @property
    def handle(self):
        return self.h_ptr[0]

    def __del__(self):
        if self.h_ptr[0]:
            _logger.debug('destroying MKL sparse matrix')
            _mkl_lib.mkl_sparse_destroy(self.handle)

    def export(self):
        """
        Export an MKL sparse matrix as a LensKit CSR.

        Returns:
            CSR: the LensKit matrix.
        """
        indP = _mkl_ffi.new('int*')
        nrP = _mkl_ffi.new('int*')
        ncP = _mkl_ffi.new('int*')
        rsP = _mkl_ffi.new('int**')
        reP = _mkl_ffi.new('int**')
        ciP = _mkl_ffi.new('int**')
        vsP = _mkl_ffi.new('double**')
        rv = _mkl_lib.mkl_sparse_d_export_csr(self.handle, indP, nrP, ncP, rsP, reP, ciP, vsP)
        _mkl_check_return(rv, 'mkl_sparse_d_export_csr')
        if indP[0] != 0:
            raise ValueError('output index is not 0-indexed')
        nr = nrP[0]
        nc = ncP[0]
        reB = _mkl_ffi.buffer(reP[0], nr * _mkl_ffi.sizeof('int'))
        re = np.frombuffer(reB, np.intc)
        nnz = re[nr-1]
        ciB = _mkl_ffi.buffer(ciP[0], nnz * _mkl_ffi.sizeof('int'))
        vsB = _mkl_ffi.buffer(vsP[0], nnz * _mkl_ffi.sizeof('double'))

        cols = np.frombuffer(ciB, np.intc)[:nnz].copy()
        vals = np.frombuffer(vsB, np.float_)[:nnz].copy()
        rowptrs = np.zeros(nr + 1, dtype=np.int32)
        rowptrs[1:] = re

        return CSR(nr, nc, nnz, rowptrs, cols, vals)

    def mult_vec(self, alpha, x, beta, y):
        """
        Compute :math:`\\alpha A x + \\beta y`, where :math:`A` is this matrix.
        """
        desc = _mkl_basic_descr()
        x = np.require(x, np.float64, 'C')
        yout = np.require(y, np.float64, 'C')
        if yout is y:
            yout = yout.copy()

        _x = _mkl_ffi.cast('double*', x.ctypes.data)
        _y = _mkl_ffi.cast('double*', yout.ctypes.data)

        rv = _mkl_lib.mkl_sparse_d_mv(10, alpha, self.handle, desc[0], _x, beta, _y)
        _mkl_check_return(rv, 'mkl_sparse_d_mv')

        return yout

def csr_syrk(csr: CSR):
    """
    Interface to the ``mkl_sparse_syrk`` routine, with necessary setup and conversion.
    """

    _logger.debug('syrk: processing %dx%d matrix (%d nnz)', csr.nrows, csr.ncols, csr.nnz)

    src = SparseM.from_csr(csr)

    _logger.debug('syrk: ordering matrix')
    rv = _mkl_lib.mkl_sparse_order(src.handle)
    _mkl_check_return(rv, 'mkl_sparse_order')

    _logger.debug('syrk: multiplying matrix')
    mult = SparseM()
    rv = _mkl_lib.mkl_sparse_syrk(11, src.handle, mult.h_ptr)
    _mkl_check_return(rv, 'mkl_sparse_syrk')
    del src  # free a little memory

    _logger.debug('syrk: exporting matrix')
    result = mult.export()
    _logger.debug('syrk: received %dx%d matrix (%d nnz)',
                  result.nrows, result.ncols, result.nnz)
    return result


#!/usr/bin/env python

import threading
import cv2
import time
import rospy
import numpy as np
from bolt_msgs.msg import Control
from std_msgs.msg import Int32
from sensor_msgs.msg import Image

import sys
import os

sys.path.append('../neural_net/')
os.chdir('../neural_net/')

import const
from image_converter import ImageConverter
from drive_run import DriveRun
from config import Config
from image_process import ImageProcess

SHARP_TURN_MIN = 0.3
BRAKE_APPLY_SEC = 1.5
THROTTLE_DEFAULT = 0.2
THROTTLE_SHARP_TURN = 0.05

class NeuralControl:
    def __init__(self, weight_file_name):
        rospy.init_node('run_neural')
        self.ic = ImageConverter()
        self.image_process = ImageProcess()
        self.rate = rospy.Rate(30)
        self.drive= DriveRun(weight_file_name)
        rospy.Subscriber('/bolt/front_camera/image_raw', Image, self._controller_cb)
        self.image = None
        self.image_processed = False
        #self.config = Config()
        self.braking = False

    def _controller_cb(self, image):
        img = self.ic.imgmsg_to_opencv(image)
        cropped = img[Config.config['image_crop_y1']:Config.config['image_crop_y2'],
                      Config.config['image_crop_x1']:Config.config['image_crop_x2']]

        img = cv2.resize(cropped, (Config.config['input_image_width'],
                                   Config.config['input_image_height']))

        self.image = self.image_process.process(img)

        ## this is for CNN-LSTM net models
        if Config.config['lstm'] is True:
            self.image = np.array(self.image).reshape(1,
                                 Config.config['input_image_height'],
                                 Config.config['input_image_width'],
                                 Config.config['input_image_depth'])
        self.image_processed = True

    def timer_cb(self):
        self.braking = False

def main(weight_file_name):

    # ready for neural network
    neural_control = NeuralControl(weight_file_name)

    # ready for /bolt topic publisher
    joy_pub = rospy.Publisher('/bolt', Control, queue_size = 10)
    joy_data = Control()

    print('\nStart running. Vroom. Vroom. Vroooooom......')
    print('steer \tthrt: \tbrake')

    while not rospy.is_shutdown():

        if neural_control.image_processed is False:
            continue

        # predicted steering angle from an input image
        prediction = neural_control.drive.run(neural_control.image)
        joy_data.steer = prediction

        #############################
        ## TODO: you need to change the vehicle speed wisely
        ## e.g. not too fast in a curved road and not too slow in a straight road

        # if brake is not already applied and sharp turn
        if neural_control.braking is False:
            if abs(joy_data.steer) > SHARP_TURN_MIN:
                joy_data.throttle = THROTTLE_SHARP_TURN
                joy_data.brake = 0.5
                neural_control.braking = True
                timer = threading.Timer(BRAKE_APPLY_SEC, neural_control.timer_cb)
                timer.start()
            else:
                joy_data.throttle = THROTTLE_DEFAULT
                joy_data.brake = 0

        ## publish joy_data
        joy_pub.publish(joy_data)

        ## print out
        if Config.config['lstm'] is True:
            cur_output = '{0:.3f} \t{1:.2f} \t{2:.2f}\r'.format(prediction[0][0][0],
                          joy_data.throttle, joy_data.brake)
        else:
            cur_output = '{0:.3f} \t{1:.2f} \t{2:.2f}\r'.format(prediction[0][0],
                          joy_data.throttle, joy_data.brake)

        sys.stdout.write(cur_output)
        sys.stdout.flush()

        ## ready for processing a new input image
        neural_control.image_processed = False
        neural_control.rate.sleep()

if __name__ == "__main__":
    try:
        if len(sys.argv) != 2:
            exit('Usage:\n$ rosrun run_neural run_neural.py weight_file_name')

        main(sys.argv[1])

    except KeyboardInterrupt:
        print ('\nShutdown requested. Exiting...')


from pathlib import Path

import torch
from numpy.testing import assert_array_almost_equal

import pytest
from int_rew import rnd, vae
from int_rew.rollout import IntValueRolloutStorage
from rainy.envs import Atari, DummyParallelEnv, atari_parallel
from rainy.envs.testing import DummyEnv
from rainy.lib.rollout import RolloutSampler
from rainy.net.policy import CategoricalDist
from rainy.utils import Device

def config() -> rnd.RNDConfig:
    c = rnd.RNDConfig()
    c.nworkers = 4
    c.nsteps = 4
    c.set_env(lambda: Atari("Venture", cfg="rnd", frame_stack=False))
    c.set_parallel_env(atari_parallel())
    return c

@pytest.mark.parametrize(
    "irew_gen", [rnd.irew_gen_default(), vae.irew_gen_cnn_vae()],
)
def test_save_and_load(irew_gen) -> None:
    c = config()
    c._int_reward_gen = irew_gen
    agent = rnd.RNDAgent(c)
    agent.irew_gen.gen_rewards(torch.randn(4 * 4, 2, 84, 84))
    nonep = agent.irew_gen.rff_rms.mean.cpu().numpy()
    savedir = Path("Results/Test")
    if not savedir.exists():
        savedir.mkdir(parents=True)
    agent.save("agent.pth", savedir)
    agent.close()
    c.device = Device(use_cpu=True)
    agent = rnd.RNDAgent(c)
    agent.load("agent.pth", savedir)
    nonep_new = agent.irew_gen.rff_rms.mean.cpu().numpy()
    assert_array_almost_equal(nonep, nonep_new)
    agent.close()

def test_storage_and_irew() -> None:
    penv = DummyParallelEnv(lambda: DummyEnv(array_dim=(16, 16)), 6)
    NSTEPS = 4
    ACTION_DIM = 3
    NWORKERS = penv.nworkers
    states = penv.reset()
    storage = IntValueRolloutStorage(
        NSTEPS, NWORKERS, Device(), 0.99, Device(use_cpu=True)
    )
    storage.set_initial_state(states)
    policy_head = CategoricalDist(ACTION_DIM)
    for _ in range(NSTEPS):
        state, reward, done, _ = penv.step([None] * NWORKERS)
        value = torch.randn(NWORKERS)
        pvalue = torch.randn(NWORKERS)
        policy = policy_head(torch.randn(NWORKERS).view(-1, 1))
        storage.push(state, reward, done, value=value, policy=policy, pvalue=pvalue)
    rewards = torch.randn(NWORKERS * NSTEPS)
    storage.calc_int_returns(torch.randn(NWORKERS), rewards, gamma=0.99, lambda_=0.95)
    batch = storage.batch_states(penv)
    batch_shape = torch.Size((NSTEPS * NWORKERS,))
    assert batch.shape == torch.Size((*batch_shape, 16, 16))
    MINIBATCH = 12
    sampler = rnd.rollout.RNDRolloutSampler(
        RolloutSampler(storage, penv, MINIBATCH),
        storage,
        torch.randn(NSTEPS * NWORKERS),
        1.0,
        1.0,
    )
    assert sampler.int_returns.shape == batch_shape
    assert sampler.int_values.shape == batch_shape
    assert sampler.advantages.shape == batch_shape
    for batch in sampler:
        assert len(batch.states) == MINIBATCH
    penv.close()


# This file contains all submission functionality for
# the monophonic test data.
# Author: Stefan Kahl, 2018, Chemnitz University of Technology

import os
import json
import operator

import numpy as np

import config as cfg
import test
from model import lasagne_net as birdnet
from model import lasagne_io as io
from utils import audio
from utils import image
from utils import batch_generator as bg
from utils import log

################### DATASAT HANDLING ####################
def parseTestSet():

    # Status
    log.i('PARSING TEST SET...', new_line=False)
    t = []

    # Get all sound files
    wav_files = [os.path.join(cfg.TESTSET_PATH, f) for f in sorted(os.listdir(cfg.TESTSET_PATH)) if os.path.splitext(f)[1] in ['.wav']]

    # Parse files
    for f in wav_files:
        t.append((f, os.path.splitext(f)[0].split('_RN')[-1]))

    # Load class ids
    codes = []
    with open('metadata/labelset.txt', 'r') as lfile:
        for line in lfile.readlines():
            codes.append(line.replace('\r\n', '').replace('\n', ''))
    labels = []
    with open('metadata/labelset_latin.txt', 'r') as lfile:
        for line in lfile.readlines():
            labels.append(line.replace('\r\n', '').replace('\n', ''))

    # Status
    log.i(('Done!', len(t), 'TEST FILES'))

    return t, codes, labels

def getClassId(c):

    if c in LABELS:
        return CODES[LABELS.index(c)]
    else:
        print 'MISSING CLASS:', c
        return False

def runTest(SNAPSHOTS, TEST):

    # Do we have more than one snapshot?
    if not isinstance(SNAPSHOTS, (list, tuple)):
        SNAPSHOTS = [SNAPSHOTS]

    # Load snapshots
    test_functions = []
    for s in SNAPSHOTS:

        # Settings
        NET = s['net']
        cfg.CLASSES = s['classes']
        cfg.IM_DIM = s['im_dim']
        cfg.IM_SIZE = s['im_size']

        # Compile test function
        test_net = birdnet.test_function(NET, hasTargets=False, layer_index=-1)
        test_functions.append(test_net)

    # Status
    log.i('START TESTING...')

    # Make predictions
    submission = ''
    cnt = 1
    for spec_batch, mediaid, filename in bg.threadedGenerator(test.getSpecBatches(TEST)):

        try:

            # Prediction
            prediction_batch = []
            for test_func in test_functions:
                if len(prediction_batch) == 0:
                    prediction_batch = test_func(spec_batch)
                else:
                    prediction_batch += test_func(spec_batch)
            prediction_batch /= len(test_functions)

            # Eliminate the scores for 'Noise'
            if 'Noise' in cfg.CLASSES:
                prediction_batch[: , cfg.CLASSES.index('Noise')] = np.min(prediction_batch)

            # Prediction pooling
            p_pool = test.predictionPooling(prediction_batch)

            # Get class labels
            p_labels = {}
            for i in range(p_pool.shape[0]):
                p_labels[cfg.CLASSES[i]] = p_pool[i]

            # Sort by score
            p_sorted =  sorted(p_labels.items(), key=operator.itemgetter(1), reverse=True)

            # Add scores to submission
            for i in range(min(100, len(p_sorted))):
                if getClassId(p_sorted[i][0]):
                    submission += mediaid + ';' + getClassId(p_sorted[i][0]) + ';' + str(p_sorted[i][1]) + ';' + str(i + 1) + '\n'

            # Show sample stats
            log.i((cnt, filename), new_line=False)
            log.i(('TOP PREDICTION:', p_sorted[0][0], int(p_sorted[0][1] * 1000) / 10.0, '%'), new_line=True)
            cnt += 1

        except KeyboardInterrupt:
            cfg.DO_BREAK = True
            break

    # Status
    log.i('DONE TESTING!')

    return submission

if __name__ == '__main__':

    # Parse Testset
    TEST, CODES, LABELS = parseTestSet()

    # Load trained models
    if not isinstance(cfg.TEST_MODELS, (list, tuple)):
        cfg.TEST_MODELS = [cfg.TEST_MODELS]
    SNAPSHOTS = []
    for test_model in cfg.TEST_MODELS:
        SNAPSHOTS.append(io.loadModel(test_model))

    # Generate submission
    submission = runTest(SNAPSHOTS, TEST)

    # Write submission to file
    log.i('WRITING SUBMISSION...', new_line=False)
    with open(cfg.RUN_NAME + '_MONOPHONE_SUBMISSION.txt', 'w') as sfile:
        sfile.write(submission)
    log.i('DONE!')


from enum import Flag
import glob
import sys
sys.path += ['../utils']
import csv
from tqdm import tqdm
import collections
import gzip
import pickle
import numpy as np
import faiss
import os
import pytrec_eval
import json
from msmarco_eval import quality_checks_qids, compute_metrics, load_reference
import time
import logging
from metric import Metric
import argparse

def get_gpu_index(cpu_index):
    gpu_resources = []
    ngpu = faiss.get_num_gpus()
    tempmem = -1
    for i in range(ngpu):
        res = faiss.StandardGpuResources()
        if tempmem >= 0:
            res.setTempMemory(tempmem)
        gpu_resources.append(res)

    def make_vres_vdev(i0=0, i1=-1):
        " return vectors of device ids and resources useful for gpu_multiple"
        vres = faiss.GpuResourcesVector()
        vdev = faiss.IntVector()
        if i1 == -1:
            i1 = ngpu
        for i in range(i0, i1):
            vdev.push_back(i)
            vres.push_back(gpu_resources[i])
        return vres, vdev

    co = faiss.GpuMultipleClonerOptions()
    co.shard = True
    gpu_vector_resources, gpu_devices_vector = make_vres_vdev(0, ngpu)
    gpu_index = faiss.index_cpu_to_gpu_multiple(gpu_vector_resources, gpu_devices_vector, cpu_index, co)
    return gpu_index

def document_split_faiss_index(emb_dir,checkpoint_postfix,dev_query_emb,gpu_index=False,top_k=1000,emb_prefix="passage_"):
    logging.info(f"***** processing faiss indexing in split-mode *****")

    emb_file_pattern = os.path.join(emb_dir,f'{emb_prefix}{checkpoint_postfix}__emb_p__data_obj_*.pb')
    emb_file_lists = glob.glob(emb_file_pattern)
    emb_file_lists = sorted(emb_file_lists, key=lambda name: int(name.split('_')[-1].replace('.pb',''))) # sort with split num
    logging.info(f"pattern {emb_file_pattern}\n file lists: {emb_file_lists}")

    # [[scores-nparray,scores-nparray..],[ANCE_ids-nparray,ANCE_ids-nparray,...]]
    merged_candidate_pair_dev = {"D":None,"I":None}

    top_k_dev = top_k

    index_offset=0

    for emb_file in emb_file_lists:
        with open(emb_file,'rb') as handle:
            sub_passage_embedding = pickle.load(handle)
        # embid_file = emb_file.replace('emb_p','embid_p')
        # with open(embid_file,'rb') as handle:
        #     sub_passage_embedding2id = pickle.load(handle)
        logging.info(f"loaded {emb_file} embeddings")
        logging.info(f"sub_passage_embedding size {sub_passage_embedding.shape}")
        dim = sub_passage_embedding.shape[1]
        faiss.omp_set_num_threads(16)
        cpu_index = faiss.IndexFlatIP(dim)
        logging.info("***** Faiss: total {} gpus *****".format(faiss.get_num_gpus()))
        index = get_gpu_index(cpu_index) if gpu_index else cpu_index
        index.add(sub_passage_embedding)

        D, dev_I = index.search(dev_query_emb, top_k_dev) # [n_dev,top]; [n_dev,top]
        dev_I = dev_I + index_offset
        if merged_candidate_pair_dev["D"] is None:
            merged_candidate_pair_dev["D"] = D
            merged_candidate_pair_dev["I"] = dev_I
        else:
            merged_candidate_pair_dev["D"] = np.concatenate([merged_candidate_pair_dev["D"],D],axis=1) # [n_dev,topk_dev *2]
            merged_candidate_pair_dev["I"] = np.concatenate([merged_candidate_pair_dev["I"],dev_I],axis=1) # [n_dev,topk_dev *2]
            sorted_ind = np.flip(np.argsort(merged_candidate_pair_dev["D"],axis=1),axis=1) # descent sort along topk_dev *2 scores for each row in n_dev
            merged_candidate_pair_dev["D"]=np.take_along_axis(merged_candidate_pair_dev["D"], sorted_ind, axis=1)[:,:top_k_dev] # [n_dev,topk_dev *2]
            merged_candidate_pair_dev["I"]=np.take_along_axis(merged_candidate_pair_dev["I"], sorted_ind, axis=1)[:,:top_k_dev] # [n_dev,topk_dev *2]

        index_offset = index_offset + sub_passage_embedding.shape[0]
        index.reset()
        sub_passage_embedding=None

    return merged_candidate_pair_dev["D"],merged_candidate_pair_dev["I"]

def save_trec_file(query_embedding2id, passage_embedding2id,
                I_nearest_neighbor,D_nearest_neighbor,
                trec_save_path,topN=1000):

    print(f"saving trec file to {trec_save_path}")
    output_f = open(trec_save_path,"w")
    # qids_to_ranked_candidate_passages = {}

    # select_top_K = topN if "MaxP" not in task_name else  topN*4
    select_top_K = topN

    for query_idx in range(len(I_nearest_neighbor)):
        seen_pid = set() # handle MaxP
        query_id = query_embedding2id[query_idx]
        # prediction[query_id] = {}

        top_ann_pid = I_nearest_neighbor[query_idx].copy()
        selected_ann_idx = top_ann_pid[:select_top_K]

        rank = 0
        for matrix_rank, idx in enumerate(selected_ann_idx):
            pred_pid = passage_embedding2id[idx]

            if not pred_pid in seen_pid:
                # this check handles multiple vector per document
                output_f.write(f"{str(query_id)} Q0 {str(pred_pid)} {str(rank+1)} {str(D_nearest_neighbor[query_idx][matrix_rank])} ance\n")
                rank += 1
                seen_pid.add(pred_pid)

                if len(seen_pid) >= topN:
                    break

    output_f.close()

# memory saving embedding loading
def get_p_emb_idx_by_pid_list(pid_list,embbeding_dir,checkpoint_step,emb_prefix="passage_"):
    emb_file_pattern = os.path.join(embbeding_dir,f'{emb_prefix}{checkpoint_step}__emb_p__data_obj_*.pb')
    emb_file_lists = glob.glob(emb_file_pattern)
    emb_file_lists = sorted(emb_file_lists, key=lambda name: int(name.split('_')[-1].replace('.pb',''))) # sort with split num

    embedding2id = []
    for emb_file in emb_file_lists:
        embid_file= emb_file.replace('emb_p','embid_p')
        with open(embid_file,'rb') as handle:
            embedding2id.append(pickle.load(handle))

    embedding2id = np.concatenate(embedding2id, axis=0)

    pid_idxs=[]
    p_emb2id=[]
    for pid in pid_list:
        p_idx=np.asarray(np.where(embedding2id==pid)).flatten()
        pid_idxs.append(p_idx)
        p_emb2id.extend([pid for _ in range(p_idx.shape[0])])
    pid_idxs = np.concatenate(pid_idxs,axis=0)
    p_emb2id = np.asarray(p_emb2id).flatten()

    assert p_emb2id.shape[0] == pid_idxs.shape[0]

    return pid_idxs.reshape([-1,]), p_emb2id

def split_load_doc_emb_by_pid_embidx_list(pid_embidx_list,embbeding_dir,checkpoint_step,emb_prefix="passage_"):

    pid_embidx_list=np.asarray(pid_embidx_list).reshape([-1,])

    emb_file_pattern = os.path.join(embbeding_dir,f'{emb_prefix}{checkpoint_step}__emb_p__data_obj_*.pb')
    emb_file_lists = glob.glob(emb_file_pattern)
    emb_file_lists = sorted(emb_file_lists, key=lambda name: int(name.split('_')[-1].replace('.pb',''))) # sort with split num

    idx_lower_bound=0
    p_embedding=np.ones([pid_embidx_list.shape[0],768])
    for emb_file in emb_file_lists:
        with open(emb_file,'rb') as handle:
            embedding=pickle.load(handle)
            dt=embedding.dtype

        idx_upper_bound = idx_lower_bound + embedding.shape[0]
        sub_relative_idx = np.intersect1d(np.where(pid_embidx_list>=idx_lower_bound),
                                np.where(pid_embidx_list<idx_upper_bound))

        p_embedding[sub_relative_idx]=embedding[pid_embidx_list[sub_relative_idx] - idx_lower_bound]
        idx_lower_bound += embedding.shape[0]

    assert p_embedding.shape[0] == pid_embidx_list.shape[0]
    return p_embedding.astype(dt)

def load_passage_embeddings(load_emb=True,emb_prefix='passage_'):
    print(f'loading passage embeddings {load_emb} and ids')
    emb_file_pattern = os.path.join(checkpoint_path,f'{emb_prefix}{checkpoint}__emb_p__data_obj_*.pb')
    emb_file_lists = glob.glob(emb_file_pattern)
    emb_file_lists = sorted(emb_file_lists, key=lambda name: int(name.split('_')[-1].replace('.pb',''))) # sort with split num
    print(f"pattern {emb_file_pattern}\n file lists: {emb_file_lists}")

    passage_embedding = []
    passage_embedding2id = []
    for emb_file in emb_file_lists:
        try:
            if load_emb:
                with open(emb_file,'rb') as handle:
                    passage_embedding.append(pickle.load(handle))
            embid_file = emb_file.replace('emb_p','embid_p')
            with open(embid_file,'rb') as handle:
                passage_embedding2id.append(pickle.load(handle))
        except:
            break

    if load_emb:
        passage_embedding = np.concatenate(passage_embedding, axis=0)
    passage_embedding2id = np.concatenate(passage_embedding2id, axis=0)
    if load_emb:
        print(f'passge embedding shape {passage_embedding.shape}, passage offset-id shape {passage_embedding2id.shape}')
    print(f'passge embedding [:3] {passage_embedding[:3]}, passage offset-id [:3] {passage_embedding2id[:3]}')

    return passage_embedding, passage_embedding2id

def convert_trec_to_MARCO_id(processed_data_dir,tmp_path,output_path,queryset_prefix,data_type=0,d2q_trec=False):

    # trec_save_path_list = glob.glob(f"{task_name}data-type-{data_type}_test-set-{test_set}_ckpt-{specified_checkpoint}.trec")

    qidmap_path = os.path.join(processed_data_dir,"{}-query_qid2offset.pickle".format(queryset_prefix))
    with open(qidmap_path, 'rb') as handle:
        qid2offset = pickle.load(handle)
    pidmap_path = os.path.join(processed_data_dir,"pid2offset.pickle")
    with open(pidmap_path, 'rb') as handle:
        pid2offset = pickle.load(handle)

    offset2qid = {}
    for k in qid2offset:
        offset2qid[qid2offset[k]]=k
    offset2pid = {}
    for k in pid2offset:
        offset2pid[pid2offset[k]]=k

    # for path in trec_save_path_list:
    print("processing file:",tmp_path)
    with open(tmp_path) as f:
        lines=f.readlines()
    with open(output_path,"w") as f:
        for line in tqdm(lines):
            qid , Q0, pid, rank, score, tag = line.strip().split(' ')
            # print(offset2qid[int(qid)] , Q0, pid, rank, score.replace('-',''), tag)
            if not d2q_trec:
                qid_inrun=str(offset2qid[int(qid)])
                pid_inrun="D"+str(offset2pid[int(pid)]) if data_type==0 or data_type==-1 else str(offset2pid[int(pid)])
            else:
                qid_inrun="D"+str(offset2pid[int(qid)]) if data_type==0 or data_type==-1 else str(offset2pid[int(qid)])
                pid_inrun=str(offset2qid[int(pid)])

            f.write(f"{qid_inrun} Q0 {pid_inrun} {rank} {score} {tag}\n")

def retrieval_specify(processed_data_dir,queryset_prefix,trec_save_path,emb_dir,checkpoint_postfix,data_type,topN,gpu_index):

    # trec_save_path = f"{task_name}_data-type-{data_type}_test-set-{test_set}_ckpt-{checkpoint}.trec"
    dev_query_embedding = []
    dev_query_embedding2id = []

    for i in range(100):
        try:
            # load the embedding of queries
            with open(
                os.path.join(emb_dir,"{}_query_{}__emb_p__data_obj_{}.pb".format(queryset_prefix,str(checkpoint_postfix),str(i))),
                 'rb') as handle:
                dev_query_embedding.append(pickle.load(handle))

            with open(
                os.path.join(emb_dir,"{}_query_{}__embid_p__data_obj_{}.pb".format(queryset_prefix,str(checkpoint_postfix),
                str(i))), 'rb') as handle:
                dev_query_embedding2id.append(pickle.load(handle))
        except:
            break
    if (not dev_query_embedding) or (not dev_query_embedding2id):
        print("No query data found for checkpoint: ",emb_dir)
    dev_query_embedding = np.concatenate(dev_query_embedding, axis=0)
    dev_query_embedding2id = np.concatenate(dev_query_embedding2id, axis=0)

    print(f'query embedding shape {dev_query_embedding.shape}, query offset-id shape {dev_query_embedding2id.shape}')
    print(f'query embedding [:3] {dev_query_embedding[:3]}, query offset-id [:3] {dev_query_embedding2id[:3]}')

    # # full ranking metrics

    dev_D, dev_I = document_split_faiss_index(dev_query_embedding,top_k=topN,gpu_index=gpu_index)

    _,passage_embedding2id=load_passage_embeddings(load_emb=False,emb_prefix='passage_')
    save_trec_file(dev_query_embedding2id, passage_embedding2id,
                    dev_I,dev_D,"tmp_intermediate.trec",topN=topN)

    convert_trec_to_MARCO_id(
        processed_data_dir=processed_data_dir,
        tmp_path="tmp_intermediate.trec",
        output_path=trec_save_path,
        test_set=queryset_prefix,
        data_type=data_type)
    os.remove("tmp_intermediate.trec")
    return None

def get_arguments():
    parser = argparse.ArgumentParser()

    # Required parameters
    parser.add_argument(
        "--emb_dir",
        default="../task_folder/ann/",
        type=str,
        required=True,
    )
    parser.add_argument(
        "--checkpoint_postfix",
        default= "1000",
        type=str,
        required=True,
    )
    parser.add_argument(
        "--emb_dir",
        default="../task_folder/ann/",
        type=str,
        required=True,
    )
    parser.add_argument(
        "--processed_data_dir",
        default="../data/preprocessed_data",
        type=str,
        required=True,
    )

    parser.add_argument(
        "--trec_path",
        default="./result.trec",
        type=str,
        required=True,
    )

    parser.add_argument(
        "--queryset_prefix",
        default="dev",
        type=str,
        required=True,
    )
    parser.add_argument(
        "--data_type",
        type=int,
        default=0,
        help="document or passage",
    )
    parser.add_argument(
        "--topN",
        type=int,
        default=100,
        help="document or passage",
    )
    parser.add_argument(
        "--gpu_index",
        type=bool,
        default=False,
        help="document or passage",
    )
    args = parser.parse_args()

    return args

def main():

    # task_name = "custom name"
    # emb_dir = "../task_folder/ann/"# location for dumpped query and passage/document embeddings which is output_dir
    # checkpoint_postfix = "1000"
    # data_type = 0 # 0 for document, 1 for passage
    # processed_data_dir = "../data/preprocessed_data"# preprocessed data folder
    # queryset_prefix = "dev"
    # gpu_index=False
    # topN = 100
    # # raw_data_dir = "../data/raw_data/" # folder for the qrels

    args = get_arguments()
    emb_dir = args.emb_dir # location for dumpped query and passage/document embeddings which is output_dir
    checkpoint_postfix = args.checkpoint_postfix
    data_type = args.data_type # 0 for document, 1 for passage
    processed_data_dir = args.processed_data_dir # preprocessed data folder
    trec_path = args.trec_path
    queryset_prefix = args.queryset_prefix
    gpu_index = args.gpu_index
    topN = args.topN

    print("start time:",time.asctime( time.localtime(time.time()) ))
    # trec_save_path = f"{task_name}_data-type-{data_type}_query-set-{queryset_prefix}_ckpt-{checkpoint_postfix}.trec"
    retrieval_specify(
        processed_data_dir=processed_data_dir,
        queryset_prefix=queryset_prefix,trec_save_path=trec_path,
        emb_dir=emb_dir,checkpoint_postfix=checkpoint_postfix,data_type=data_type,topN=topN,gpu_index=gpu_index)

if __name__ == '__main__':
    main()


import face_recognition
import picamera
import numpy as np
known_face_encodings = []
names = []
def load_face_encoding(name, file_name):
   image = face_recognition.load_image_file(file_name)
   face_encoding = face_recognition.face_encodings(image)[0]
   known_face_encodings.append(face_encoding)
   names.append(name)
# Get a reference to the Raspberry Pi camera.
# If this fails, make sure you have a camera connected to the RPi and that you
# enabled your camera in raspi-config and rebooted first.
camera = picamera.PiCamera()
camera.resolution = (320, 240)
output = np.empty((240, 320, 3), dtype=np.uint8)
# Load a sample picture and learn how to recognize it.
print("Loading known face image(s)")
load_face_encoding("Wang Lei", "wanglei.jpg")
load_face_encoding("Obama", "obama.jpg")
load_face_encoding("Messi","messi.jpg")
load_face_encoding("Ronaldo","ronaldo.jpg")
load_face_encoding("Biden", "biden.jpg")
# Initialize some variables
face_locations = []
face_encodings = []
while True:
   print("Capturing image.")
   # Grab a single frame of video from the RPi camera as a numpy array
   camera.capture(output, format="rgb")
   # Find all the faces and face encodings in the current frame of video
   face_locations = face_recognition.face_locations(output)
   print("Found {} faces in image.".format(len(face_locations)))
   face_encodings = face_recognition.face_encodings(output, face_locations)
   # Loop over each face found in the frame to see if it's someone we know.
   for face_encoding in face_encodings:
       # See if the face is a match for the known face(s)
       matches = face_recognition.face_distance(known_face_encodings, face_encoding)
       name = "<Unknown Person>"
       min_distance = min(matches)
       if min_distance < 0.6:
           i = matches.argmin()
           name = names[i]
       print("I see someone named {}!".format(name))


# Module for flux calibrating spectra
import numpy as np
import os
import matplotlib.pyplot as plt
from astropy.io import fits

from pypeit import msgs
from pypeit.spectrographs.util import load_spectrograph
from pypeit import sensfunc
from pypeit import specobjs
from astropy import table
from IPython import embed

class FluxCalibrate(object):
    """
    Class for flux calibrating spectra.

    Args:
        spec1dfiles (list):
        sensfiles (list):
        par (pypeit.par.pypeitpar.FluxCalibrate, optional):
        outfiles (list, optional):
            Names of the outfiles.  If None, this is set to spec1dfiles and those are overwritten
    """
    # Superclass factory method generates the subclass instance
    @classmethod
    def get_instance(cls, spec1dfiles, sensfiles, par=None, debug=False):
        pypeline = fits.getheader(spec1dfiles[0])['PYPELINE'] + 'FC'
        return next(c for c in cls.__subclasses__() if c.__name__ == pypeline)(
            spec1dfiles, sensfiles, par=par, debug=debug)

    def __init__(self, spec1dfiles, sensfiles, par=None, debug=False, outfiles=None):

        self.spec1dfiles = spec1dfiles
        self.sensfiles = sensfiles

        # Output file names
        self.outfiles = spec1dfiles if outfiles is None else outfiles

        # Load the spectrograph
        header = fits.getheader(spec1dfiles[0])
        self.spectrograph = load_spectrograph(header['PYP_SPEC'])
        self.par = self.spectrograph.default_pypeit_par()['fluxcalib'] if par is None else par
        self.debug = debug

        sens_last = None
        for spec1, sens, outfile in zip(self.spec1dfiles, self.sensfiles, self.outfiles):
            # Read in the data
            sobjs = specobjs.SpecObjs.from_fitsfile(spec1)
            if sens != sens_last:
                wave, sensfunction, meta_table, out_table, header_sens = sensfunc.SensFunc.load(sens)
            self.flux_calib(sobjs, wave, sensfunction, meta_table)
            sobjs.write_to_fits(sobjs.header, outfile, overwrite=True)

    def flux_calib(self, sobjs, wave, sensfunction, meta_table):
        """
        Dummy method overloaded by subclass

        Args:
            sobjs:
            wave:
            sensfunction:
            meta_table:

        """
        pass

class MultiSlitFC(FluxCalibrate):
    """
    Child of FluxSpec for Multislit and Longslit reductions
    """

    def __init__(self, spec1dfiles, sensfiles, par=None, debug=False, outfiles=None):
        super().__init__(spec1dfiles, sensfiles, par=par, debug=debug, outfiles=outfiles)

    def flux_calib(self, sobjs, wave, sensfunction, meta_table):
        """
        Apply sensitivity function to all the spectra in an sobjs object.

        Args:
            sobjs (object):
               SpecObjs object
            wave (ndarray):
               wavelength array for sensitivity function (nspec,)
            sensfunction (ndarray):
               sensitivity function
            meta_table (table):
               astropy table containing meta data for sensitivity function

        """

        # Run
        for sci_obj in sobjs:
            sci_obj.apply_flux_calib(wave, sensfunction,
                                     sobjs.header['EXPTIME'],
                                     extinct_correct=self.par['extinct_correct'],
                                     longitude=self.spectrograph.telescope['longitude'],
                                     latitude=self.spectrograph.telescope['latitude'],
                                     extrap_sens=self.par['extrap_sens'],
                                     airmass=float(sobjs.header['AIRMASS']))

class EchelleFC(FluxCalibrate):
    """
    Child of FluxSpec for Echelle reductions
    """

    def __init__(self, spec1dfiles, sensfiles, par=None, debug=False):
        super().__init__(spec1dfiles, sensfiles, par=par, debug=debug)

    def flux_calib(self, sobjs, wave, sensfunction, meta_table):
        """
        Apply sensitivity function to all the spectra in an sobjs object.

        Args:
            sobjs (object):
               SpecObjs object
            wave (ndarray):
               wavelength array for sensitivity function (nspec,)
            sensfunction (ndarray):
               sensitivity function
            meta_table (table):
               astropy table containing meta data for sensitivity function

        """

        # Flux calibrate the orders that are mutually in the meta_table and in the sobjs. This allows flexibility
        # for applying to data for cases where not all orders are present in the data as in the sensfunc, etc.,
        # i.e. X-shooter with the K-band blocking filter.
        ech_orders = np.array(meta_table['ECH_ORDERS']).flatten()
        #norders = ech_orders.size
        for sci_obj in sobjs:
            # JFH Is there a more elegant pythonic way to do this without looping over both orders and sci_obj?
            indx = np.where(ech_orders == sci_obj.ECH_ORDER)[0]
            if indx.size==1:
                sci_obj.apply_flux_calib(wave[:, indx[0]],sensfunction[:,indx[0]],
                                         sobjs.header['EXPTIME'],
                                         extinct_correct=self.par['extinct_correct'],
                                         extrap_sens = self.par['extrap_sens'],
                                         longitude=self.spectrograph.telescope['longitude'],
                                         latitude=self.spectrograph.telescope['latitude'],
                                         airmass=float(sobjs.header['AIRMASS']))
            elif indx.size == 0:
                msgs.info('Unable to flux calibrate order = {:} as it is not in your sensitivity function. '
                          'Something is probably wrong with your sensitivity function.'.format(sci_obj.ECH_ORDER))
            else:
                msgs.error('This should not happen')