import os
import math

import pandas as pd
import numpy as np

class Data:
    """
    class acts as a convenient data feeder for few shots learning

    """

    def __init__(self, datapath=None, seed=3, remove_unk=False):

        self.remove_unk = remove_unk

        if datapath is None:
            datapath = os.path.join(os.path.dirname(os.path.abspath(__file__)), "./preprocessed")

        np.random.seed(seed)

        # loading vocab
        def return_vocab(filename):

            keys = [l.strip() for l in open(filename).readlines()]
            values = range(0, len(keys))
            return dict(zip(keys, values))

        def return_inv_vocab(filename):

            keys = [l.strip() for l in open(filename).readlines()]
            values = range(0, len(keys))
            return dict(zip(values, keys))

        self.entityvocab = return_vocab(os.path.join(datapath, "entity.vocab"))
        self.propertyvocab = return_vocab(os.path.join(datapath, "property.vocab"))
        self.wordvocab = return_vocab(os.path.join(datapath, "word.vocab"))

        self.inv_entityvocab = return_inv_vocab(os.path.join(datapath, "entity.vocab"))
        self.inv_propertyvocab = return_inv_vocab(os.path.join(datapath, "property.vocab"))
        self.inv_wordvocab = return_inv_vocab(os.path.join(datapath, "word.vocab"))

        # loading data files names
        self.datafile = {"train":os.path.join(datapath, "train.ids"),
                         "valid": os.path.join(datapath, "valid.ids"),
                         "test": os.path.join(datapath, "valid.ids")
                         }

        self.data = {}

    def read_data(self, mode):
        # modes = ["train", "test", "valid"]
        data = []

        f = self.datafile[mode]
        x = pd.read_csv(f, names=["sub", "pred", "obj", "question", "subtype", "objtype", "dep", "direction", "placeholder_dict"])

        if self.remove_unk:
            unkdep = self.wordvocab["_UNK_DEP_"] if "_UNK_DEP_" in self.wordvocab else None
            x = x[x.dep != unkdep]
            x = x[x.apply(lambda i: str(self.wordvocab["_PLACEHOLDER_SUB_"]) in i['question'].split(), axis=1)]

        x.reset_index(inplace=True)

        tmp = [[], [], [], []]
        for l in x.iterrows():

            tmp[0].append([int(i) for i in l[1]['question'].split()])
            tmp[1].append([int(i) for i in l[1]['subtype'].split()])
            tmp[2].append([int(i) for i in l[1]['objtype'].split()])
            tmp[3].append([int(i) for i in l[1]['dep'].split()])

        x['question'] = tmp[0]
        x['subtype'] = tmp[1]
        x['objtype'] = tmp[2]
        x['dep'] = tmp[3]

        x['question_length'] = x.apply(lambda l: len(l['question']), axis=1)
        x['subtype_length'] = x.apply(lambda l: len(l['subtype']), axis=1)
        x['objtype_length'] = x.apply(lambda l: len(l['objtype']), axis=1)
        x['dep_length'] = x.apply(lambda l: len(l['dep']), axis=1)
        x['triple'] = x.apply(lambda l: [l['sub'], l['pred'], l['obj']], axis=1)

        return x

    def datafeed(self, mode, config, shuffle=True):
        """

        :param mode: train, valid, test
        :param config: config object
        :param shot_percentage: float between 0 and 1 indicating the percentage of the training data taken into consideration
        :param min_count: int indicating the minimum count of the predicates of the examples being taken in to consideration
        :param shuffle: whether to shuffle the training data or not
        :param kfold: a number between 1 and 10
        :return:
        """

        x = self.read_data(mode)
        self.data[mode] = x

        dataids = x.index

        if shuffle:
            np.random.shuffle(dataids)

        return self.yield_datafeed(mode, dataids, x, config)

    def yield_datafeed(self, mode, dataids, x, config):
        """
        given a dataframe and selected ids and a mode yield data for experiments
        :param mode:
        :param dataids:
        :param x:
        :param config:
        :return:
        """

        if mode == "train":

            for epoch in range(config.MAX_EPOCHS):

                def chunks(l, n):
                    """Yield successive n-sized chunks from l."""
                    for i in range(0, len(l), n):
                        yield l[i:i + n]

                for bn, batchids in enumerate(chunks(dataids, config.BATCH_SIZE)):

                    batch = x.iloc[batchids]

                    max_length = max([batch['subtype_length'].values.max(),
                         batch['objtype_length'].values.max(),
                         batch['dep_length'].values.max(),
                         ])

                    yield (
                        np.array([i for i in batch['triple'].values]),
                        self.pad(batch['subtype'].values, max_length=max_length),
                        batch['subtype_length'].values,
                        self.pad(batch['objtype'].values, max_length=max_length),
                        batch['objtype_length'].values,
                        self.pad(batch['dep'].values, max_length=max_length),
                        batch['dep_length'].values,
                        self.pad(batch['question'].values),
                        batch['question_length'].values,
                        batch['direction'].values,
                        {"epoch": epoch, "batch_id": bn, "ids": batchids, "placeholder_dict":[eval(i) for i in batch["placeholder_dict"].values]}  # meta info
                    )

        if mode == "test" or mode == "valid":
            # in case of test of validation batch of size 1 and no shuffle
            # takes longer computation time but allows variable lengths

            for id in dataids:

                batch = x.iloc[[id]]

                max_length = max([batch['subtype_length'].values.max(),
                                  batch['objtype_length'].values.max(),
                                  batch['dep_length'].values.max(),
                                  ])

                yield (
                    np.array([i for i in batch['triple']]),
                    self.pad(batch['subtype'], max_length=max_length),
                    batch['subtype_length'].values,
                    self.pad(batch['objtype'], max_length=max_length),
                    batch['objtype_length'].values,
                    self.pad(batch['dep'].values, max_length=max_length),
                    batch['dep_length'].values,
                    self.pad(batch['question'].values),
                    batch['question_length'].values,
                    batch['direction'].values,
                    {"ids": id, "placeholder_dict": [eval(i) for i in batch["placeholder_dict"].values]}  # meta info
                )

    def pad(self, x, pad_char=0, max_length=None):
        """
        helper function to add padding to a batch
        :param x: array of arrays of variable length
        :return:  x with padding of max length
        """

        if max_length is None:
            max_length = max([len(i) for i in x])

        y = np.ones([len(x), max_length]) * pad_char

        for c, i in enumerate(x):
            y[c, :len(i)] = i

        return y

class FewShotsDataFeeder:
    """
    class acts as a convenient data feeder for few shots learning

    """

    def __init__(self, datapath=None, seed=3, train_percent=0.7, test_percent=0.2, remove_unk=False):

        assert 0 < train_percent + test_percent <= 1

        self.train_percent = train_percent
        self.test_percent = test_percent
        self.valid_percent = 1 - (train_percent+test_percent)
        self.remove_unk = remove_unk

        if datapath is None:
            datapath = os.path.join(os.path.dirname(os.path.abspath(__file__)), "./preprocessed")

        np.random.seed(seed)

        # loading vocab
        def return_vocab(filename):

            keys = [l.strip() for l in open(filename).readlines()]
            values = range(0, len(keys))
            return dict(zip(keys, values))

        def return_inv_vocab(filename):

            keys = [l.strip() for l in open(filename).readlines()]
            values = range(0, len(keys))
            return dict(zip(values, keys))

        self.entityvocab = return_vocab(os.path.join(datapath, "entity.vocab"))
        self.propertyvocab = return_vocab(os.path.join(datapath, "property.vocab"))
        self.wordvocab = return_vocab(os.path.join(datapath, "word.vocab"))

        self.inv_entityvocab = return_inv_vocab(os.path.join(datapath, "entity.vocab"))
        self.inv_propertyvocab = return_inv_vocab(os.path.join(datapath, "property.vocab"))
        self.inv_wordvocab = return_inv_vocab(os.path.join(datapath, "word.vocab"))

        # loading data files names
        self.datafile = {"train":os.path.join(datapath, "train.ids"),
                         "valid": os.path.join(datapath, "valid.ids"),
                         "test": os.path.join(datapath, "valid.ids")
                         }

        self.data = {}

    def read_data(self):
        modes = ["train", "test", "valid"]
        data = []
        for m in modes:
            f = self.datafile[m]
            data.append(pd.read_csv(f, names=["sub", "pred", "obj", "question", "subtype", "objtype", "dep", "direction", "placeholder_dict"]))

        x = pd.concat(data)

        if self.remove_unk:
            unkdep = self.wordvocab["_UNK_DEP_"] if "_UNK_DEP_" in self.wordvocab else None
            x = x[x.dep != unkdep]
            x = x[x.apply(lambda i: str(self.wordvocab["_PLACEHOLDER_SUB_"]) in i['question'].split(), axis=1)]

        x.reset_index(inplace=True)

        tmp = [[], [], [], []]
        for l in x.iterrows():

            tmp[0].append([int(i) for i in l[1]['question'].split()])
            tmp[1].append([int(i) for i in l[1]['subtype'].split()])
            tmp[2].append([int(i) for i in l[1]['objtype'].split()])
            tmp[3].append([int(i) for i in l[1]['dep'].split()])

        x['question'] = tmp[0]
        x['subtype'] = tmp[1]
        x['objtype'] = tmp[2]
        x['dep'] = tmp[3]

        x['question_length'] = x.apply(lambda l: len(l['question']), axis=1)
        x['subtype_length'] = x.apply(lambda l: len(l['subtype']), axis=1)
        x['objtype_length'] = x.apply(lambda l: len(l['objtype']), axis=1)
        x['dep_length'] = x.apply(lambda l: len(l['dep']), axis=1)
        x['triple'] = x.apply(lambda l: [l['sub'], l['pred'], l['obj']], axis=1)

        return x

    def filter_data(self, x, mode, shot_percentage, min_count):

        # removing predicates with less than min_count examples in the labeled set
        x = x.groupby("pred").filter(lambda x: len(x) >= min_count)
        ids = x.groupby("pred").indices

        keep_ids = np.array([], dtype=np.int)
        if mode == "train":
            for v in ids.values():
                start = 0
                end = int(math.ceil(len(v) * shot_percentage * self.train_percent))
                keep_ids = np.append(keep_ids, v[start:end])

        elif mode == "test":
            for v in ids.values():
                start = int(math.ceil(len(v) * shot_percentage * self.train_percent))
                end = int(start + math.ceil(len(v) * self.test_percent))
                keep_ids = np.append(keep_ids, v[start:end])

        elif mode == "valid":

            for v in ids.values():
                start = int(math.ceil(len(v) * shot_percentage * self.train_percent) + math.ceil(len(v) * self.test_percent))
                end = int(start + math.ceil(len(v) * self.valid_percent))
                keep_ids = np.append(keep_ids, v[start:end])

        return keep_ids, x

    def datafeed(self, mode, config, shot_percentage=1, min_count=10, shuffle=True):
        """

        :param mode: train, valid, test
        :param config: config object
        :param shot_percentage: float between 0 and 1 indicating the percentage of the training data taken into consideration
        :param min_count: int indicating the minimum count of the predicates of the examples being taken in to consideration
        :param shuffle: whether to shuffle the training data or not
        :param kfold: a number between 1 and 10
        :return:
        """

        x = self.read_data()
        self.data[mode] = x

        dataids, x = self.filter_data(x, mode, shot_percentage, min_count)
        dataids = [i for i in dataids if i in x.index]

        if shuffle:
            np.random.shuffle(dataids)

        return self.yield_datafeed(mode, dataids, x, config)

    def yield_datafeed(self, mode, dataids, x, config):
        """
        given a dataframe and selected ids and a mode yield data for experiments
        :param mode:
        :param dataids:
        :param x:
        :param config:
        :return:
        """

        if mode == "train":

            for epoch in range(config.MAX_EPOCHS):

                def chunks(l, n):
                    """Yield successive n-sized chunks from l."""
                    for i in range(0, len(l), n):
                        yield l[i:i + n]

                for bn, batchids in enumerate(chunks(dataids, config.BATCH_SIZE)):

                    batch = x.iloc[batchids]

                    max_length = max([batch['subtype_length'].values.max(),
                         batch['objtype_length'].values.max(),
                         batch['dep_length'].values.max(),
                         ])

                    yield (
                        np.array([i for i in batch['triple'].values]),
                        self.pad(batch['subtype'].values, max_length=max_length),
                        batch['subtype_length'].values,
                        self.pad(batch['objtype'].values, max_length=max_length),
                        batch['objtype_length'].values,
                        self.pad(batch['dep'].values, max_length=max_length),
                        batch['dep_length'].values,
                        self.pad(batch['question'].values),
                        batch['question_length'].values,
                        batch['direction'].values,
                        {"epoch": epoch, "batch_id": bn, "ids": batchids, "placeholder_dict":[eval(i) for i in batch["placeholder_dict"].values]}  # meta info
                    )

        if mode == "test" or mode == "valid":
            # in case of test of validation batch of size 1 and no shuffle
            # takes longer computation time but allows variable lengths

            for id in dataids:

                batch = x.iloc[[id]]

                max_length = max([batch['subtype_length'].values.max(),
                                  batch['objtype_length'].values.max(),
                                  batch['dep_length'].values.max(),
                                  ])

                yield (
                    np.array([i for i in batch['triple']]),
                    self.pad(batch['subtype'], max_length=max_length),
                    batch['subtype_length'].values,
                    self.pad(batch['objtype'], max_length=max_length),
                    batch['objtype_length'].values,
                    self.pad(batch['dep'].values, max_length=max_length),
                    batch['dep_length'].values,
                    self.pad(batch['question'].values),
                    batch['question_length'].values,
                    batch['direction'].values,
                    {"ids": id, "placeholder_dict": [eval(i) for i in batch["placeholder_dict"].values]}  # meta info
                )

    def pad(self, x, pad_char=0, max_length=None):
        """
        helper function to add padding to a batch
        :param x: array of arrays of variable length
        :return:  x with padding of max length
        """

        if max_length is None:
            max_length = max([len(i) for i in x])

        y = np.ones([len(x), max_length]) * pad_char

        for c, i in enumerate(x):
            y[c, :len(i)] = i

        return y

class ZeroShotsDataFeeder(FewShotsDataFeeder):

    def filter_data(self, x, mode, criteria="pred", min_count=10, kfold=10, cv=0):
        """

        :param x:
        :param mode:
        :param criteria:
        :param min_count:
        :param kfold:
        :param cv:
        :return:
        """

        # removing predicates with less than min_count examples in the labeled set
        criteria_hash = criteria + "_hash"
        x[criteria_hash] = x.apply(lambda a: str(a[criteria]), axis=1)  # make criteria hashable
        x = x.groupby(criteria_hash).filter(lambda i: len(i) >= min_count)
        ids = x.groupby(criteria_hash).indices
        ids = sorted(ids.items(), key=lambda a: len(a[1]), reverse=True)

        keep_ids = np.array([], dtype=np.int)

        if mode == "train":

            start = cv
            pos = [(i + start) % kfold for i in range(int(math.ceil(kfold * self.train_percent)))]

            for c, i in enumerate(ids):
                if c % kfold in pos:
                    keep_ids = np.append(keep_ids, i[1])

        elif mode == "test":

            start = cv
            start = [(i + start + 1) % kfold for i in range(int(math.ceil(kfold * self.train_percent)))][-1]
            pos = [(i + start) % kfold for i in range(int(math.ceil(kfold * self.test_percent)))]

            for c, i in enumerate(ids):
                if c % kfold in pos:
                    keep_ids = np.append(keep_ids, i[1])

        elif mode == "valid":

            start = cv
            start = [(i + start + 1) % kfold for i in range(int(math.ceil(kfold * (self.train_percent + self.test_percent))))][-1]
            pos = [(i + start) % kfold for i in range(int(math.ceil(kfold * self.valid_percent)))]

            for c, i in enumerate(ids):
                if c % kfold in pos:
                    keep_ids = np.append(keep_ids, i[1])

        return keep_ids, x

    def datafeed(self, mode, config, criteria="pred", min_count=10, shuffle=True, kfold=10, cv=1):
        """

        :param mode: train, valid, test
        :param config: config object
        :param criteria: the column label to do zero shot on "pred" "subtype" "objtype"
        :param shot_percentage: float between 0 and 1 indicating the percentage of the training data taken into consideration
        :param min_count: int indicating the minimum count of the predicates of the examples being taken in to consideration
        :param shuffle: whether to shuffle the training data or not
        :param kfold:
        :param cv:
        :return:
        """

        x = self.read_data()

        dataids, x = self.filter_data(x, mode, criteria, min_count, kfold, cv)
        dataids = [i for i in dataids if i in x.index]

        if shuffle:
            np.random.shuffle(dataids)

        return self.yield_datafeed(mode, dataids, x, config)


import cv2
import numpy as np
import sqlite3
import datetime
from recognize_func import FaceRecognition
from AUTOMATED_EMAIL_FINAL import sendAutomatedMail

def Logout():

    conn = sqlite3.connect('Students.db')
    db = conn.cursor()

    now_time = datetime.datetime.now()
    end = datetime.datetime(now_time.year, now_time.month, now_time.day, 17, 0, 0)

    if now_time >= end:
        sql = """SELECT * from Students"""
        db.execute(sql)
        rows = db.fetchall()
        for Srow in rows:

            mobile_no = Srow[0]
            EmpTable = str(Srow[1]) + "_" + str(Srow[0])
            sql = """SELECT Date, Login, Logout, TimeDelta, TimedOut from {} WHERE Date = ?""".format(EmpTable)
            db.execute(sql, [(str(datetime.date.today()))])
            row = db.fetchone()

            if row:
                if row[2] == None and row[3] == None and row[4] == None:

                    date = row[0]
                    login = datetime.datetime.strptime(row[1], '%Y-%m-%d %H:%M:%S.%f')
                    logout = end
                    timeDelta = logout - login
                    min = timeDelta.total_seconds()/60.0
                    sql = """UPDATE {} set Logout = ?, TimeDelta = ?, TimedOut = ?""".format(EmpTable)
                    db.execute(sql, [(str(end)), str(min), ('True')])
                    conn.commit()
                    # NOTE: Code for AutomatedEmail is called through fuction for modularity
                    name = Srow[1]
                    email_address = Srow[2]
                    print(name, email_address)
                    try:
                        sendAutomatedMail(name=name, email_address=email_address)
                        print('Automatic Mail has been sent...')
                    except Exception as e:
                        print('No internet service for sending email...')

    elif now_time < end:
        # NOTE: input mobile which is already Registered in previous page
        mobile_no = input("Enter mobile NO: ")

        if mobile_no:
            try:
                sql = """SELECT * from Students WHERE Mobile = ?"""
                db.execute(sql, [(mobile_no)])
                rows = db.fetchone()
                if rows:
                    if mobile_no == str(rows[0]):
                        name = FaceRecognition()
                        print("Name : ", name)
                        Name = rows[1].replace("_"," ")
                        EmpTable = str(Name) + " " +str(rows[0])
                        print(EmpTable)
                        if EmpTable == name:
                            name = name.replace(" ","_")
                            try:
                                sql = """SELECT Logout from {} WHERE Date = ?""".format(name)
                                db.execute(sql, [(str(datetime.date.today()))])
                                row = db.fetchone()
                                if row[0] == None:
                                    print(f"{Name} hasn't Logged out yet, Proceeding to Logout")

                                    sql = """UPDATE {} SET Logout = ?""".format(name)
                                    db.execute(sql, [(str(datetime.datetime.now()))])
                                    conn.commit()

                                    sql = """SELECT Date, Login, Logout, TimeDelta, TimedOut from {} WHERE Date = ?""".format(name)
                                    db.execute(sql,[(str(datetime.date.today()))])
                                    row = db.fetchone()
                                    if row:
                                        date = row[0]
                                        login = datetime.datetime.strptime(row[1], '%Y-%m-%d %H:%M:%S.%f')
                                        logout = datetime.datetime.strptime(row[2], '%Y-%m-%d %H:%M:%S.%f')
                                        timeDelta = logout - login
                                        min = timeDelta.total_seconds()/60.0
                                        if row[3] == None and row[4] == None:
                                            print(f"Saving total time spent by {Name} Employee on Date {date}...")
                                            sql = """UPDATE {} SET timeDelta = ?, TimedOut = ?""".format(name)
                                            db.execute(sql, [(str(min)), ('False')])
                                            conn.commit()
                                            print(f"{Name} Employee total time spent in office on Date {date} is : {timeDelta} or {min} Minutes")
                                            print(f"{Name} your Logout Successful...Data has been saved successfully...Thank you for your service...Have a nice day")

                                else:
                                    print(f'{Name} Employee is Already successfully logged out and data has been saved...')
                            except Exception as e:
                                print("User doesn't exist...or Database retrieval error")
                        else:
                            print("User doesn't exist...Please register first or Wrong person identified i.e: person and mobile number mismatch")

            except Exception as e:
                print("User doesn't exist...or Database retrieval error")
        else:
            print("Please enter mobile No. It can't be blank")

    conn.commit()
    conn.close()

Logout()


import sys
sys.path.append ('./model/AV_model')
sys.path.append ('./data/utils')
from keras.models import load_model
from option import ModelMGPU
import os
import scipy.io.wavfile as wavfile
import numpy as np
import utils
import tensorflow as tf
import librosa

#parameters
people = 2
num_gpu=1

#path
model_path = './saved_AV_models/AVmodel-2p-065-0.40568.h5'
result_path = './predict/'
os.makedirs(result_path,exist_ok=True)

database = '/Work19/2020/lijunjie/LRS3/AV_model_database/mix/test/'
face_emb_path = '/Work19/2020/lijunjie/LRS3/face1022_emb/test/'

single_stft_path ='/Work19/2020/lijunjie/LRS3/AV_model_database/single/test/'
single_wav_path = '/CCALab/lijunjie/LRS3_process/audio_norm/test/'
single_wav_filename = ''

print('Initialing Parameters......')

#loading data
print('Loading data ......')
test_file = []
with open('/Work19/2020/lijunjie/LRS3/AV_model_database/AVdataset_test.txt','r') as f:
    test_file = f.readlines()

def get_data_name(line,people=people,database=database,face_emb=face_emb_path,single_stft_path=single_stft_path):
    parts = line.split() # get each name of file for one testset
    mix_str = parts[0]
    name_list = mix_str.replace('.npy','')
    name_list = name_list.replace('mix-','',1)
    names = name_list.split('-')
    single_idxs = []
    for i in range(people):
        single_idxs.append(names[i])
    file_path = database + mix_str
    mix = np.load(file_path)
    face_embs = np.zeros((1,75,1,1792,people))
    for i in range(people):
        face_embs[0,:,:,:,i] = np.load(face_emb+"%s_faceemb.npy"%single_idxs[i])

    single_stft = np.zeros((1,people,16000*3))
    for i in range(people):
        for file in os.listdir(os.path.join(single_stft_path,single_idxs[i])):
           single_wav_filename = 'norm_'+os.path.splitext(file.split('_')[1])[0]+'.wav'
           x,_ = librosa.load(os.path.join(single_wav_path,single_idxs[i],single_wav_filename),sr=16000)
           single_stft[0,i,:] = x[0:16000*3]

    return mix,single_idxs,face_embs,single_stft

#result predict
av_model = load_model(model_path,custom_objects={'tf':tf},compile=False)
if num_gpu>1:
    parallel = ModelMGPU(av_model,num_gpu)
    for line in test_file:
        mix,single_idxs,face_emb,single_stft = get_data_name(line,people,database,face_emb_path,single_stft_path)
        mix_ex = np.expand_dims(mix,axis=0)
        cRMs = parallel.predict([mix_ex,face_emb])
        cRMs = cRMs[0]
        prefix =''
        for idx in single_idxs:
            prefix +=idx+'-'
        for i in range(len(cRMs)):
            cRM =cRMs[:,:,:,i]
            assert cRM.shape ==(298,257,2)
            F = utils.fast_icRM(mix,cRM)
            si_snr = utils.cal_si_snr(F,single_stft[:,:,:,i])
            print('si_snr:',si_snr)
            T = utils.fase_istft(F,power=False)
            filename = result_path+str(single_idxs[i])+'.wav'
            wavfile.write(filename,16000,T)

if num_gpu<=1:
    for line in test_file:
        mix,single_idxs,face_emb,single_stft = get_data_name(line,people,database,face_emb_path,single_stft_path)
        mix_ex = np.expand_dims(mix,axis=0)
        cRMs = av_model.predict([mix_ex,face_emb])
        cRMs = cRMs[0]
        prefix =''
        for idx in single_idxs:
            prefix +=idx+'-'
        T_=np.zeros((1,people,3*16000))
        for i in range(cRMs.shape[-1]):
            cRM =cRMs[:,:,:,i]
            assert cRM.shape ==(298,257,2)
            F = utils.fast_icRM(mix,cRM)
            # # si_snr = utils.cal_si_snr(F,F)
            T = utils.fast_istft(F,power=False)
            T_[0,i,:] = T
            filename = result_path+str(single_idxs[i])+'.wav'
            wavfile.write(filename,16000,T)
        si_snr,_,si_snr_idx = utils.cal_si_snr_with_pit(T_,single_stft)
        print('si_snr:',si_snr)


from __future__ import print_function
from os.path import expandvars
from icecube import icetray
import math

"""
Convenience functions for configuring CLSim components.
"""

def configureOpenCLDevices(UseGPUs=True, UseCPUs=False, OverrideApproximateNumberOfWorkItems=None, DoNotParallelize=True, UseOnlyDeviceNumber=None):
    # get OpenCL devices
    from icecube.clsim import I3CLSimOpenCLDevice
    from icecube.icetray import logging
    import string

    openCLDevicesRaw = [device for device in I3CLSimOpenCLDevice.GetAllDevices() if (device.gpu and UseGPUs) or (device.cpu and UseCPUs)]
    openCLDevices = []

    # (auto-)configure OpenCL devices
    for i, device in enumerate(openCLDevicesRaw):
        if UseOnlyDeviceNumber is not None and i != UseOnlyDeviceNumber:
            # skip all devices except for the selected one (if there is a selection)
            continue

        if  'Tesla' in device.device or 'GTX' in device.device or 'TITAN' in device.device:
            # assume these are "fast", all others are "slow"
            device.useNativeMath=True
            if '1080' in device.device or 'TITAN' in device.device:
                # These things have 8GB of memory, so use more of it (because why not)
                device.approximateNumberOfWorkItems=1024000*2
            elif 'Tesla' in device.device or '580' in device.device or '680' in device.device or '980' in device.device:
                # these cards should have enough ram to support this
                device.approximateNumberOfWorkItems=1024000
            else:
                device.approximateNumberOfWorkItems=102400
        elif 'Graphics Device' in device.device:
            # with nVidia's driver version 367.18, a GTX 1080 shows up as just "Graphics Devices"
            # These things have 8GB of memory, so use more of it (because why not)
            device.useNativeMath=True
            device.approximateNumberOfWorkItems=1024000*2
        elif 'Tahiti' in device.device:
            device.useNativeMath=True
            device.approximateNumberOfWorkItems=102400*2
        else:
            device.useNativeMath=False
            device.approximateNumberOfWorkItems=10240

        if OverrideApproximateNumberOfWorkItems is not None:
            device.approximateNumberOfWorkItems=OverrideApproximateNumberOfWorkItems

        if DoNotParallelize and device.cpu:
            # check if we can split this device into individual cores
            try:
                if device.platform == "Intel(R) OpenCL":
                    # device fission seems to cause serious segfaults in the Intel OpenCL driver.
                    # do not use it.
                    subDevices = []
                else:
                    subDevices = device.SplitDevice()

                if len(subDevices) > 0:
                    if OverrideApproximateNumberOfWorkItems is not None:
                        subDevices[0].approximateNumberOfWorkItems=OverrideApproximateNumberOfWorkItems
                    openCLDevices.append(subDevices[0])
                else:
                    logging.log_warn("failed to split CPU device into individual cores %s %s [using full device with minimal number of work-items to (hopefully) disable parallelization]" % (device.platform, device.device), unit="clsim")
                    device.approximateNumberOfWorkItems=1
                    openCLDevices.append(device)
            except:
                logging.log_error("failed to split CPU device into individual cores %s %s [using full device with minimal number of work-items to (hopefully) disable parallelization]" % (device.platform, device.device), unit="clsim")
                device.approximateNumberOfWorkItems=1
                openCLDevices.append(device)

        else:
            openCLDevices.append(device)

    return openCLDevices

def parseIceModel(IceModelLocation, disableTilt=False):
    from os.path import exists, isdir, isfile, expandvars
    from icecube.clsim.MakeIceCubeMediumProperties import MakeIceCubeMediumProperties
    from icecube.clsim.MakeAntaresMediumProperties import MakeAntaresMediumProperties
    from icecube.clsim.MakeIceCubeMediumPropertiesPhotonics import MakeIceCubeMediumPropertiesPhotonics

    if IceModelLocation=="ANTARES":
        return MakeAntaresMediumProperties()

    if not exists(IceModelLocation):
        raise RuntimeError("The specified ice model path \"%s\" does not exist" % IceModelLocation)

    if isdir(IceModelLocation):
        # it's a PPC ice description directory
        mediumProperties = MakeIceCubeMediumProperties(iceDataDirectory=IceModelLocation, useTiltIfAvailable=not disableTilt)
    elif isfile(IceModelLocation):
        # it's a photonics ice description file
        mediumProperties = MakeIceCubeMediumPropertiesPhotonics(tableFile=IceModelLocation)
    else:
        raise RuntimeError("The specified ice model path \"%s\" is neither a directory nor a file." % IceModelLocation)

    return mediumProperties

def setupDetector(GCDFile,
                  SimulateFlashers=False,
                  IceModelLocation=expandvars("$I3_SRC/clsim/resources/ice/spice_mie"),
                  DisableTilt=False,
                  UnWeightedPhotons=False,
                  UnWeightedPhotonsScalingFactor=None,
                  MMCTrackListName="MMCTrackList",
                  UseGeant4=False,
                  CrossoverEnergyEM=None,
                  CrossoverEnergyHadron=None,
                  UseCascadeExtension=True,
                  StopDetectedPhotons=True,
                  DOMOversizeFactor=5.,
                  UnshadowedFraction=0.9,
                  HoleIceParameterization=expandvars("$I3_SRC/ice-models/resources/models/angsens/as.h2-50cm"),
                  WavelengthAcceptance=None,
                  DOMRadius=0.16510*icetray.I3Units.m, # 13" diameter
                  IgnoreSubdetectors=['IceTop']):
    """
    Set up data structures used in N different places in clsim
    """

    from icecube import clsim
    from icecube.clsim import GetDefaultParameterizationList
    from icecube.clsim import GetFlasherParameterizationList
    from icecube.clsim import GetHybridParameterizationList
    import numpy

    def harvest_detector_parameters(GCDFile):

        from icecube import dataio, dataclasses
        from I3Tray import I3Tray

        tray = I3Tray()
        tray.Add(dataio.I3Reader, Filenamelist=[GCDFile])

        # make sure the geometry is updated to the new granular format
        tray.AddModule("I3GeometryDecomposer",
                       If=lambda frame: ("I3OMGeoMap" not in frame) and ("I3ModuleGeoMap" not in frame))

        def pluck_geo(frame):
            pluck_geo.frame = frame
        pluck_geo.frame = None
        tray.Add(pluck_geo, Streams=[icetray.I3Frame.Geometry])

        def pluck_calib(frame):
            pluck_calib.frame = frame
        pluck_calib.frame = None
        tray.Add(pluck_calib, Streams=[icetray.I3Frame.Calibration])

        tray.Execute()

        geometry = clsim.I3CLSimSimpleGeometryFromI3Geometry(
            DOMRadius, DOMOversizeFactor, pluck_geo.frame,
            ignoreSubdetectors=dataclasses.ListString(IgnoreSubdetectors),
            # NB: we trust advanced users to properly label subdetectors, and disable any
            # min/max string/om numbers and strings/oms to ignore
            ignoreStrings=dataclasses.ListInt(), ignoreDomIDs=dataclasses.ListUInt(),
            ignoreStringIDsSmallerThan=1, ignoreStringIDsLargerThan=numpy.iinfo(numpy.int32).max,
            ignoreDomIDsSmallerThan=1, ignoreDomIDsLargerThan=numpy.iinfo(numpy.uint32).max,
            splitIntoPartsAccordingToPosition=False, useHardcodedDeepCoreSubdetector=False
            )

        rde = dict()
        for k, domcal in pluck_calib.frame['I3Calibration'].dom_cal.iteritems():
            rde[k] = domcal.relative_dom_eff

        return geometry, rde

    geometry, rde = harvest_detector_parameters(GCDFile)

    # ice properties
    if isinstance(IceModelLocation, str):
        mediumProperties = parseIceModel(IceModelLocation, disableTilt=DisableTilt)
    else:
        # get ice model directly if not a string
        mediumProperties = IceModelLocation

    # detector properties
    if WavelengthAcceptance is None:
        # the hole ice acceptance curve peaks at a value different than 1
        peak = numpy.loadtxt(HoleIceParameterization)[0] # The value at which the hole ice acceptance curve peaks
        domEfficiencyCorrection = UnshadowedFraction*peak

        acceptance = {
            'IceCube': clsim.GetIceCubeDOMAcceptance(domRadius = DOMRadius*DOMOversizeFactor, efficiency=domEfficiencyCorrection),
            'DeepCore': clsim.GetIceCubeDOMAcceptance(domRadius = DOMRadius*DOMOversizeFactor, efficiency=domEfficiencyCorrection, highQE=True),
        }

        domAcceptanceEnvelope = clsim.I3CLSimFunctionFromTable(acceptance['IceCube'].GetMinWlen(), acceptance['IceCube'].GetWavelengthStepping(), [max((f.GetEntryValue(i) for f in acceptance.values())) for i in range(acceptance['IceCube'].GetNumEntries())])

        domAcceptance = clsim.I3CLSimFunctionMap()
        #loop over every event in the geometry and determine what type of DOM it is
        #based on the relative efficiency in the detector status
        for string_id,dom_id in zip(geometry.stringIDs,geometry.domIDs):
            k = icetray.OMKey(string_id,dom_id,0)
            releff = rde.get(k,float('nan'))
            if releff == 1:
                domAcceptance[k] = acceptance['IceCube']
            elif round(releff, 6) == 1.35:
                domAcceptance[k] = acceptance['DeepCore']
            elif math.isnan(releff):
                #DOMs where rde is nan are bad DOMs which are unplugged
                #call them IceCube for now, hits generated by them will be
                #removed later by detector simulation
                domAcceptance[k] = acceptance['IceCube']
            else:
                raise ValueError("Relative DOM efficiency is neither 1 nor 1.35. You probably need to add support for individual DOM efficiencies")

    else:
        domAcceptance = WavelengthAcceptance
        domAcceptanceEnvelope = WavelengthAcceptance

    angularAcceptance = clsim.GetIceCubeDOMAngularSensitivity(holeIce=HoleIceParameterization)

    # photon generation wavelength bias
    if not UnWeightedPhotons:
        wavelengthGenerationBias = domAcceptanceEnvelope
        if UnWeightedPhotonsScalingFactor is not None:
            raise RuntimeError("UnWeightedPhotonsScalingFactor should not be set when UnWeightedPhotons is not set")
    else:
        if UnWeightedPhotonsScalingFactor is not None:
            print("***** running unweighted simulation with a photon pre-scaling of", UnWeightedPhotonsScalingFactor)
            wavelengthGenerationBias = clsim.I3CLSimFunctionConstant(UnWeightedPhotonsScalingFactor)
        else:
            wavelengthGenerationBias = None

    # create wavelength generators
    wavelengthGenerators = [clsim.makeCherenkovWavelengthGenerator(wavelengthGenerationBias, UnWeightedPhotons, mediumProperties)]

    # flasher parameterizations
    if SimulateFlashers:
        # this needs a spectrum table in order to pass spectra to OpenCL
        spectrumTable = clsim.I3CLSimSpectrumTable()
        for spectrum in spectrumTable:
            wavelengthGenerators.append(makeWavelengthGenerator(spectrum, wavelengthGenerationBias, mediumProperties))
        print("number of spectra (1x Cherenkov + Nx flasher):", len(spectrumTable))
    else:
        # no spectrum table is necessary when only using the Cherenkov spectrum
        spectrumTable = None

    # muon&cascade parameterizations
    ppcConverter = clsim.I3CLSimLightSourceToStepConverterPPC(photonsPerStep=200)
    ppcConverter.SetUseCascadeExtension(UseCascadeExtension)
    if not UseGeant4:
        particleParameterizations = GetDefaultParameterizationList(ppcConverter, muonOnly=False)
    else:
        if CrossoverEnergyEM>0 or CrossoverEnergyHadron>0:
            particleParameterizations = GetHybridParameterizationList(ppcConverter, CrossoverEnergyEM=CrossoverEnergyEM, CrossoverEnergyHadron=CrossoverEnergyHadron)
        elif MMCTrackListName is None or MMCTrackListName=="":
            particleParameterizations = [] # make sure absolutely **no** parameterizations are used
        else:
            # use no parameterizations except for muons with lengths assigned to them
            # (those are assumed to have been generated by MMC)
            particleParameterizations = GetDefaultParameterizationList(ppcConverter, muonOnly=True)

    return dict(Geometry=geometry,
                MediumProperties=mediumProperties,
                IceModelLocation=IceModelLocation,
                WavelengthGenerationBias=wavelengthGenerationBias,
                SpectrumTable=spectrumTable,
                GenerateCherenkovPhotonsWithoutDispersion=UnWeightedPhotons,
                WavelengthGenerators=wavelengthGenerators,
                DOMRadius=DOMRadius,
                DOMOversizeFactor=DOMOversizeFactor,
                DOMPancakeFactor=DOMOversizeFactor,
                UnshadowedFraction=UnshadowedFraction,
                AngularAcceptance=angularAcceptance,
                WavelengthAcceptance=domAcceptance,
                ParameterizationList=particleParameterizations,
                IgnoreSubdetectors=IgnoreSubdetectors)

def setupPropagators(RandomService,
                     DetectorParams,
                     UseCPUs=False,
                     UseGPUs=True,
                     UseOnlyDeviceNumber=None,
                     DoNotParallelize=False,
                     OverrideApproximateNumberOfWorkItems=None,
                     ):
    """
    Create a collection of photon propagators suitable for use with I3CLSimServer
    """
    from icecube import clsim

    geometry = DetectorParams['Geometry']
    mediumProperties = DetectorParams['MediumProperties']
    wavelengthGenerationBias = DetectorParams['WavelengthGenerationBias']
    wavelengthGenerators = DetectorParams['WavelengthGenerators']

    openCLDevices = configureOpenCLDevices(
        UseGPUs=UseGPUs,
        UseCPUs=UseCPUs,
        OverrideApproximateNumberOfWorkItems=OverrideApproximateNumberOfWorkItems,
        DoNotParallelize=DoNotParallelize,
        UseOnlyDeviceNumber=UseOnlyDeviceNumber
    )

    def create_converter(device):
        return clsim.initializeOpenCL(device, RandomService, geometry,
            mediumProperties, wavelengthGenerationBias, clsim.I3CLSimRandomValuePtrSeries(wavelengthGenerators),
            pancakeFactor=DetectorParams['DOMPancakeFactor'],
            )

    return map(create_converter, openCLDevices)


import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__),'..'))

from os.path import isfile
import numpy as np
import torch
from .config import config
from Sensation7 import Sensation
from Sensation7.sensation_models import Encoder
from .oshaberi_model import OshaberiText
from .output import Output
from Sensation7.config import config

import sentencepiece as spm
import os
from typing import Tuple

from TrainBase import TrainBase

class Train(TrainBase):
    LogTitle:str = 'TrainOshaberi'
    dtype:torch.dtype = torch.float16
    torchdtype:torch.dtype = torch.float32
    def TrainProcess(self) -> None:
        # ------ Additional Trainings ------
        current,memory,answer= self.GetOshaberiTextData()
        self.log('shape of current,memory,answer',current.shape,memory.shape,answer.shape)

        model = self.LoadPytorchModel(OshaberiText(),Output.OshaberiText_params,dtype=torch.float32,eval_mode=False)
        batch_size = Output.OshaberiTextBatchSize
        epochs = Output.OshaberiTextEpochs
        criterion = torch.nn.CrossEntropyLoss()
        optimizer = torch.optim.Adam(model.parameters(),lr=Output.OshaberiTextLearningRate)
        self.Train(
            model=model,
            epochs=epochs,
            batch_size=batch_size,
            optimizer=optimizer,
            criterion=criterion,
            device=self.device,
            train_x=[current,memory],
            train_y=answer,
            metrics=self.fit.CE_Accuracy,
        )
        torch.save(model.state_dict(),Output.OshaberiText_params)
        del model,current,memory,answer
        self.log('trained Text Generator.')
        self.release_system_memory()
        # --- end of Additional Training ---

        self.log('Train process was finished')

    @torch.no_grad()
    def GetOshaberiTextData(self) -> Tuple[torch.Tensor,...]:
        self.FTmodel = self.load_python_obj(Sensation.FastText_params)
        self.separator = spm.SentencePieceProcessor()
        self.separator.Load(Sensation.Separator_params)
        self.bos_id = self.separator.bos_id()
        self.eos_id = self.separator.eos_id()
        self.unk_id = self.separator.unk_id()
        self.bos = self.separator.IdToPiece(self.bos_id)
        self.eos = self.separator.IdToPiece(self.eos_id)

        self.encoder = self.LoadPytorchModel(Encoder(),Sensation.Encoder_params,dtype=self.dtype,device=self.device)
        self.padvec = np.zeros((config.word_dim,),dtype='float16')
        self.log('loaded FastText Sentencepiece,TextEncoder')

        with open(Sensation.Corpus_file,'r',encoding='utf-8') as f:
            corpus = f.read().split('\n')
        idx = np.random.permutation(len(corpus))[:Output.CorpusUseLength]
        corpus = [corpus[i] for i in idx]
        self.log('Using corpus length is',len(corpus))
        cur,mem,ans = [],[],[]
        for c in corpus:
            i,j,k = self.get_a_OshaberiText_data(c)
            cur.append(i)
            mem.append(j)
            ans.append(k)
        cur = torch.cat(cur)[:Output.OshaberiTextDataSize]
        mem = torch.cat(mem)[:Output.OshaberiTextDataSize]
        ans = torch.cat(ans)[:Output.OshaberiTextDataSize]
        del self.FTmodel,self.separator,corpus,self.encoder,self.padvec
        self.release_system_memory()
        return cur,mem,ans

    def get_a_OshaberiText_data(self,indata:str) -> Tuple[torch.Tensor,...]:
        pieces = [self.bos,*self.separator.EncodeAsPieces(indata),self.eos]
        piecesid,vectors = [],[]
        for i in pieces:
            pid =self.separator.PieceToId(i)
            if i in self.FTmodel.wv and pid != self.separator.unk_id():
                piecesid.append(pid)
                vectors.append(self.FTmodel.wv[i])
        answer = np.array(piecesid[1:],dtype='int64')
        answer = torch.from_numpy(answer)

        current = []
        veclen = len(vectors)
        for i in range(1,veclen):
            vec = vectors[:i]
            if len(vec) < config.generate_max_words:
                vec += [self.padvec]*(config.generate_max_words - len(vec))
            else:
                vec = vec[-config.generate_max_words:]
            current.append(np.stack(vec))
        current = np.stack(current)
        current = torch.from_numpy(current).type(self.dtype)

        memvec = []
        for i in range(veclen):
            vec = vectors[i:i+config.text_seq_len]
            if len(vec) < config.text_seq_len:
                vec += [self.padvec] * (config.text_seq_len- len(vec))
            memvec.append(np.stack(vec))
        memvec = torch.from_numpy(np.stack(memvec)).type(self.dtype)
        memvec = self.fit.Predict(self.encoder,memvec,Output.MaxSamples,self.device)
        mvl = memvec.size(0)
        if mvl < config.use_mem_len:
            memvec = memvec.repeat(((config.use_mem_len//mvl)+1,1))
            mvl = memvec.size(0)
        memory = torch.stack([memvec[np.random.permutation(mvl)[:config.use_mem_len]] for _ in range(current.size(0))])
        sample_idx = np.random.permutation(answer.size(0))[:Output.MaxSamples]
        memory = memory[sample_idx]
        current = current[sample_idx]
        answer = answer[sample_idx]
        return current,memory,answer


import ROOT
import numpy as np
from collections import defaultdict
from gna import datapath
import gna.parameters.ibd
import gna.parameters.oscillation
import itertools
from gna.exp import baseexp
from gna.env import env
import gna.constructors as C

year=365*24*60*60.0

spectrumfiles = {
    'Pu239': 'reactor_anu_spectra/Huber/Huber_smooth_extrap_Pu239_13MeV0.01MeVbin.dat',
    'Pu241': 'reactor_anu_spectra/Huber/Huber_smooth_extrap_Pu241_13MeV0.01MeVbin.dat',
    'U235': 'reactor_anu_spectra/Huber/Huber_smooth_extrap_U235_13MeV0.01MeVbin.dat',
    'U238': 'reactor_anu_spectra/Mueller/Mueller_smooth_extrap_U238_13MeV0.01MeVbin.dat',
}

eperfission = {
    'Pu239': (209.99, 0.60),
    'Pu241': (213.60, 0.65),
    'U235': (201.92, 0.46),
    'U238': (205.52, 0.96),
}

geo_flux_files = {
    'U238': 'AntineutrinoSpectrum_238U.knt',
    'Th232': 'AntineutrinoSpectrum_232Th.knt',
}

geo_flux_normalizations = {
    'U238': (2.7e3/5, 0.3),
    'Th232': (0.8e3/5, 0.3),
}

class Reactor(object):
    def __init__(self, name, location, fission_fractions,
                 power=None, power_rate=None):
        self.name = name
        self.location = location

        self.power = power

        self.power_rate = power_rate

        self.fission_fractions = {isoname: frac
                                  for isoname, frac in fission_fractions.items()}

    def assign(self, ns):
        self.ns = ns("reactors")(self.name)
        self.ns.reqparameter("ThermalPower", central=self.power, sigma=0)

    def calcdistance(self, detector):
        return np.sqrt(np.sum((np.array(self.location) - np.array(detector.location))**2))

    def initdistance(self, ns, detector):
        dist = self.calcdistance(detector)
        pair_ns = detector.ns(self.name)
        pair_ns.reqparameter("L_{0}".format(self.name), central=dist, sigma=0)
        pair_ns.defparameter("L", target=pair_ns.ref("L_{0}".format(self.name)))

class ReactorGroup(object):
    def __init__(self, name, reactors):
        self.name = name
        self.reactors = reactors

        self.power = 1
        self.power_rate = reactors[0].power_rate

        self.fission_fractions = reactors[0].fission_fractions

    def assign(self, ns):
        self.ns = ns("reactorgroups")(self.name)

    def initdistance(self, ns, detector):
        pair_ns = detector.ns(self.name)
        bindings = {}
        for i, reactor in enumerate(self.reactors):
            bindings["L_{0}".format(i)] = detector.ns(reactor.name)["L"]
            bindings["P_{0}".format(i)] = reactor.ns["ThermalPower"]
        ROOT.ReactorGroup(len(self.reactors), ns=pair_ns, bindings=bindings)

        pair_ns.defparameter("L", target=pair_ns.ref("Lavg"))
        pair_ns.defparameter("ThermalPower", target=pair_ns.ref("Pavg"))

class Detector(object):
    def __init__(self, name, edges, location,
                 orders=None, protons=None, livetime=None, edges_final=None):
        self.name = name
        self.edges = edges
        self.edges_final = edges_final
        self.orders = orders

        self.location = location

        self.protons = protons

        self.livetime = livetime

        self.unoscillated = None
        self.components = defaultdict(lambda: defaultdict(ROOT.Sum))
        self.backgrounds = defaultdict(ROOT.Sum)
        self.intermediates = {}
        self.intermediates_bkg = {}
        self.hists = defaultdict(dict)
        self.back_hists = defaultdict(dict)

    def assign(self, ns):
        self.ns = ns("detectors")(self.name)

        if self.protons is not None:
            self.ns.defparameter("TargetProtons", central=self.protons, relsigma=0.01)

        self.ns.reqparameter("Eres_a", central=0.0, sigma=0)
        self.ns.reqparameter("Eres_b", central=0.03, sigma=0)
        self.ns.reqparameter("Eres_c", central=0.0, sigma=0)
        self.ns.reqparameter("rho_C14", central=1e-16, sigma=1e-19)
        self.ns.reqparameter("CoincidenceWindow", central=300, sigma=1)

class GeoNeutrinoIsotope(object):
    def __init__(self, name):
        self.name = name

        try:
            Es_keV, self.ys = np.loadtxt(datapath(geo_flux_files[name]), unpack=True, skiprows=5)
        except OSError:
            raise Exception("Failed to load spectrum of {0} geo isotope from {1}".format(name, geo_flux_files[name]))
        self.Es = Es_keV*1e-3
        self.spectrum = ROOT.LinearInterpolator(len(self.Es), self.Es.copy(), self.ys.copy(), "use_zero")

class Isotope(object):
    def __init__(self, ns, name):
        self.ns = ns("isotopes")(name)
        self.name = name

        try:
            self.Es, self.ys = np.loadtxt(datapath(spectrumfiles[name]), unpack=True)
        except OSError:
            raise Exception("Failed to load spectrum of {0} reactor isotope from {1}".format(name, datapath(spectrumfiles[name])))

        self.spectrum = ROOT.LinearInterpolator(len(self.Es), self.Es.copy(), self.ys.copy(), "use_zero")

class ReactorExperimentModel(baseexp):
    oscprob_classes = {
            'standard': ROOT.OscProbPMNS,
            'decoh': ROOT.OscProbPMNSDecoh,
            }

    @classmethod
    def initparser(self, parser, env):
        """Initializes arguments parser with args common to all reactor experiments"""
        #  def OscProb(name):
            #  return {
                #  'standard': ROOT.OscProbPMNS,
                #  'decoh': ROOT.OscProbPMNSDecoh,
            #  }[name]
        parser.add_argument('--name', required=True)
        parser.add_argument('--ibd', choices=['zero', 'first'], default='zero')
        parser.add_argument('--backgrounds', choices=['geo'], action='append',
                            default=[], help='Choose backgrounds you want to add')
        parser.add_argument('--oscprob', choices=self.oscprob_classes.keys(),
                            default='standard')
        parser.add_argument('--binning', nargs=4, metavar=('DETECTOR', 'EMIN', 'EMAX', 'NBINS'), action='append', default=[])
        parser.add_argument('--binning-final', nargs=6, metavar=('DETECTOR', 'EMIN', 'E0', 'E1', 'NBINS01', 'EMAX'), action='append', default=[])
        parser.add_argument('--integration-order', type=int, default=4)
        parser.add_argument('--no-reactor-groups', action='store_true')
        parser.add_argument('--with-C14', action='store_true')

    def __init__(self, opts, ns=None, reactors=None, detectors=None):
        """Initialize a reactor experiment

        opts -- object with parsed common arguments returned by argparse
        ns -- namespace where experiment will create missing parameters, if is not provided ``opts.name`` will be used
        reactors -- iterable over Reactor objects, self.makereactors() will be called if None
        detectors -- iterable over Detector objects, self.makedetoctrs() will be called if None
        """
        super(ReactorExperimentModel, self).__init__(None, opts)
        self._oscprobs = {}
        self._oscprobcls = self.oscprob_classes[self.opts.oscprob]
        self._isotopes = defaultdict(list)
        self._Enu_inputs = defaultdict(set)
        self.oscprobs_comps = defaultdict(dict)
        self._geo_isotopes = defaultdict(list)

        self.ns = ns or env.ns(opts.name)
        self.reqparameters(self.ns)

        self.detectors = list(detectors if detectors is not None else self.makedetectors())
        for binopt in self.opts.binning:
            for det in self.detectors:
                if det.name == binopt[0]:
                    break
            else:
                raise Exception("can't find detector {}".format(binopt[0]))
            det.edges = np.linspace(float(binopt[1]), float(binopt[2]), int(binopt[3]))
        for (detname, emin, e0, e1, ne, emax) in self.opts.binning_final:
            for det in self.detectors:
                if det.name == detname:
                    break
            else:
                raise Exception("can't find detector {}".format(detname))
            det.edges_final = np.concatenate(([float(emin)],
                                              np.linspace(float(e0), float(e1), int(ne)),
                                              [float(emax)]))
        for det in self.detectors:
            det.assign(self.ns)
            if det.orders is None:
                det.orders = np.full(len(det.edges)-1, self.opts.integration_order, int)

        self.reactors = list(reactors if reactors is not None else self.makereactors())
        for reactor in self.reactors:
            reactor.assign(self.ns)

        self.linkpairs(self.reactors, self.detectors)

        for detector in self.detectors:
            self.make_backgrounds(detector, opts.backgrounds)
            self.setibd(detector, opts.ibd)
            self.setupobservations(detector)

    def make_backgrounds(self, detector, background_list):
        """Produce backgrounds and store them to detector components,
        currently only geoneutrino flux. Normalize it latter in makeibd"""

        #  pass
        for bkg in background_list:
            if bkg == 'geo':
                for iso_name in geo_flux_files.keys():
                    geo_isotope = GeoNeutrinoIsotope(iso_name)
                    self._isotopes["geo_"+iso_name].append(geo_isotope)
                    self._Enu_inputs[detector].add(geo_isotope.spectrum.f.inputs.x)
                    detector.intermediates_bkg['geo_flux_{}'.format(iso_name)] = geo_isotope.spectrum

    def _getoscprobcls(self, reactor, detector):
        """Returns oscillation probability class for given (reactor, detector) pair"""
        if isinstance(reactor, ReactorGroup):
            return (ROOT.OscProbPMNSMult, ROOT.OscProbPMNS)
        keys = [frozenset([reactor, detector]), reactor, detector, None]
        for key in keys:
            try:
                return (self._oscprobs[key], self._oscprobs[key])
            except KeyError:
                pass
        return (self._oscprobcls, self._oscprobcls)

    def _getnormtype(self, reactor, detector):
        """Returns normalization type applicable for given reactor and detector

        if power, power rate, protons and livetime are provided, 'calc' is returned to calculate normalization
        otherwise 'manual' is returned meaning that normalization should be done by the caller with external Norm parameters
        """
        if all(x is not None for x in [reactor.power, reactor.power_rate, detector.protons, detector.livetime]):
            return 'calc'
        return 'manual'

    def groupreactors(self, reactors, detector, precision=1e-2):
        """Merges some Reactors of reactors to ReactorGroup if their distances to detector are similar enough

        Faster approximated formula is used. Only reactors with standard OscProbPMNS formula are considered for grouping.
        Approximation of the probability is of the order of order precision**3.
        """
        groups = [[reactor] for reactor in reactors]
        groupable = []
        if not self.opts.no_reactor_groups:
            for reactor in reactors:
                if self._getoscprobcls(reactor, detector)[0] is ROOT.OscProbPMNS:
                    if self._getnormtype(reactor, detector) == 'calc':
                        groupable.append(reactor)
                        groups.remove([reactor])
        if not groupable:
            return reactors
        groups.append([])
        for reactor in sorted(groupable, key=lambda r: r.calcdistance(detector)):
            groups[-1].append(reactor)
            L = np.array([r.calcdistance(detector) for r in groups[-1]])
            P = np.array([r.power for r in groups[-1]])
            Lavg = sum(P/L)/sum(P/L**2)
            if not all(abs(L/Lavg-1) < precision):
                groups.append([groups[-1].pop()])
        grouped = []
        cnt = 0
        for group in groups:
            if len(group) > 1:
                rgroup = ReactorGroup("group{}".format(cnt), group)
                rgroup.assign(self.ns)
                rgroup.initdistance(self.ns, detector)
                grouped.append(rgroup)
                cnt += 1
            else:
                grouped.append(group[0])
        return grouped

    def normalization(self, reactor, detector, normtype):
        """Returns normalization object for given reactor/detector pair and norm type"""
        def vec(lst):
            v = ROOT.vector("std::string")()
            for s in lst:
                v.push_back(s)
            return v

        if normtype == 'calc':
            bindings = {}
            for isoname in reactor.fission_fractions:
                bindings["EnergyPerFission_{0}".format(isoname)] = self.ns("isotopes")(isoname)["EnergyPerFission"]
            norm = ROOT.ReactorNorm(vec(reactor.fission_fractions.keys()), bindings=bindings)
            lt=C.Points(detector.livetime)
            lt.points.setLabel('livetime: | '+detector.name)
            norm.isotopes.livetime(lt)
            pr=C.Points(reactor.power_rate)
            pr.points.setLabel('power: | '+reactor.name)
            norm.isotopes.power_rate(pr)
            norm.isotopes.setLabel('norm: | {} to {}'.format(reactor.name, detector.name))
            for isoname, frac in reactor.fission_fractions.items():
                ff=C.Points(frac)
                ff.points.setLabel('fission frac: | {} at {}'.format(isoname, reactor.name))
                norm.isotopes['fission_fraction_{0}'.format(isoname)](ff)
        elif normtype == 'manual':
            norm = ROOT.ReactorNormAbsolute(vec(reactor.fission_fractions.keys()))
            for isoname, frac in reactor.fission_fractions.items():
                norm.isotopes['fission_fraction_{0}'.format(isoname)](frac)
        return norm

    def linkpairs(self, reactors, detectors):
        """Setup oscillations and normalizations of all reactors and detectors.

        Resulting components are stored in detector.components.
        """
        for reactor, detector in itertools.product(reactors, detectors):
            reactor.initdistance(self.ns, detector)

        for detector in detectors:
            detector.unoscillated = ROOT.Sum()
            detector.unoscillated.sum.setLabel('unosc flux: | '+detector.name)
            grouped = self.groupreactors(reactors, detector)
            for rgroup in grouped:
                pair_ns = detector.ns(rgroup.name)
                normtype = self._getnormtype(rgroup, detector)
                with detector.ns, rgroup.ns, pair_ns:
                    norm = self.norm = self.normalization(rgroup, detector, normtype)
                    oscprobcls, weightscls = self._getoscprobcls(rgroup, detector)
                    with self.ns("oscillation"):
                        oscprob = oscprobcls(ROOT.Neutrino.ae(), ROOT.Neutrino.ae())

                normedflux = ROOT.Sum()
                normedflux.sum.setLabel('normed flux: | {}'.format(rgroup.name))

                for isoname in rgroup.fission_fractions.keys():
                    isotope = Isotope(self.ns, isoname)
                    isotope.spectrum.f.setLabel('spectrum: | {} at {}'.format(isoname, rgroup.name))
                    self._isotopes[(detector, rgroup)].append(isotope)
                    self._Enu_inputs[detector].add(isotope.spectrum.f.inputs.x)
                    subflux = ROOT.Product()
                    subflux.product.setLabel('flux | {}'.format(isoname))
                    subflux.multiply(isotope.spectrum)
                    subflux.multiply(norm.isotopes['norm_{0}'.format(isoname)])
                    detector.intermediates['flux_{}'.format(isoname)] = subflux
                    normedflux.add(subflux)
                detector.intermediates['flux'] = normedflux

                compnames = set(oscprob.probsum.inputs)

                detector.unoscillated.add(normedflux)
                if 'comp0' in compnames:
                    csum = detector.components['rate'][(weightscls, 'comp0')]
                    csum.add(normedflux)
                    csum.sum.setLabel('normed flux | comp0 at '+detector.name)
                    ones = ROOT.FillLike(1.0)
                    ones.fill.setLabel('1: comp0')
                    ones.fill.inputs(normedflux)
                    opsum=detector.components['oscprob'][(weightscls, 'comp0')]
                    opsum.add(ones)
                    opsum.sum.setLabel('osc flux | {} at {}'.format('comp0', detector.name))
                    self.oscprobs_comps[(detector, rgroup)][(weightscls, 'comp0')] = ones
                    compnames.remove('comp0')
                for osccomps in oscprob.transformations.values():
                    osccomps.setLabel('oscprob {} | {}-\\>{}'.format(osccomps.label(), rgroup.name, detector.name))
                    for compname, osccomp in osccomps.outputs.items():
                        if compname not in compnames:
                            continue
                        product = ROOT.Product()
                        product.product.setLabel('osc flux: | {} from {}'.format(compname, rgroup.name))
                        product.multiply(normedflux)
                        product.multiply(osccomp)
                        csum=detector.components['rate'][(weightscls, compname)]
                        csum.add(product)
                        csum.sum.setLabel('normed flux | {} at {}'.format(compname, detector.name))
                        opsum=detector.components['oscprob'][(weightscls, compname)]
                        opsum.add(osccomp)
                        opsum.sum.setLabel('osc flux | {} at {}'.format(compname, detector.name))
                        self.oscprobs_comps[(detector, rgroup)][(weightscls, compname)] = osccomp
                        if compname not in compnames:
                            raise Exception("overriden component {}".format(compname))
                        compnames.remove(compname)

                        if 'Enu' in osccomps.inputs:
                            self._Enu_inputs[detector].add(osccomps.inputs.Enu)
                if compnames:
                    raise Exception("components not found: {}".format(compnames))

    def setibd(self, detector, ibdtype):
        """Setup input energy values, integration and IBD calculation object for the given detector according to the ibdtype .
        ibdtype may be 'zero' or 'first' for the corresponding order.
        Resulting components are stored in detector.components.
        The backgrounds histos are also constructed here for now.
        """
        Evis_edges = detector.edges
        orders = detector.orders
        with self.ns("ibd"):
            econv = self.econv = ROOT.EvisToEe()
        if ibdtype == 'zero':
            with self.ns("ibd"):
                ibd = self.ibd = ROOT.IbdZeroOrder()
            integrator = self.integrator = ROOT.GaussLegendre(Evis_edges, orders, len(orders))
            integrator.points.setLabel('integrator 1d')
            histcls = ROOT.GaussLegendreHist
            econv.Ee.Evis(integrator.points.x)
            ibd.xsec.Ee(econv.Ee.Ee)
            ibd.Enu.Ee(econv.Ee.Ee)
            eventsparts = [ibd.xsec]
        elif ibdtype == 'first':
            with self.ns("ibd"):
                ibd = self.ibd = ROOT.IbdFirstOrder()
            integrator = self.integrator = ROOT.GaussLegendre2d(Evis_edges, orders, len(orders), -1.0, 1.0, 5)
            integrator.points.setLabel('integrator 2d')
            histcls = ROOT.GaussLegendre2dHist
            econv.Ee.Evis(integrator.points.x)
            ibd.Enu.Ee(econv.Ee.Ee)
            ibd.Enu.ctheta(integrator.points.y)
            detector.intermediates['ctheta'] = integrator.points.y
            ibd.xsec.Enu(ibd.Enu)
            ibd.xsec.ctheta(integrator.points.y)
            ibd.jacobian.Enu(ibd.Enu)
            # ibd.jacobian.Ee(integrator.points.x) # legacy: incorrect. X is Evis, not Ee
            ibd.jacobian.Ee(econv.Ee.Ee)
            ibd.jacobian.ctheta(integrator.points.y)
            eventsparts = [ibd.xsec, ibd.jacobian]
        else:
            raise Exception("unknown ibd type {0!r}".format(ibdtype))
        ibd.xsec.setLabel('IBD xsec')

        detector.intermediates['Enu'] = ibd.Enu
        detector.intermediates['xsec'] = ibd.xsec
        for inp in self._Enu_inputs.get(detector, []):
            inp.connect(ibd.Enu.Enu)

        for detector in self.detectors:
            for resname, comps in detector.components.items():
                for compid, comp in comps.items():
                    res = None
                    if resname == 'rate':
                        res = ROOT.Product()
                        res.product.setLabel('count rate: | {} at {}'.format(compid[1], detector.name))
                        res.multiply(comp)
                        for part in eventsparts:
                            res.multiply(part)
                    elif resname == 'oscprob':
                        res = comp
                    hist = detector.hists[resname][compid] = histcls(integrator)
                    hist.hist.setLabel('{} hist: | {} at {}'.format(resname, compid[1], detector.name))
                    detector.hists[resname][compid].hist.inputs(res)

            # Below we construct backgrounds histos
            if self.opts.backgrounds:
                bkg_summary = ROOT.Sum()
                bkg_summary.sum.setLabel('bkg: | '+detector.name)
                for bkg_name, bkg in detector.intermediates_bkg.items():
                    prod = ROOT.Product()
                    prod.product.setLabel('bkg: | '+detector.name)
                    prod.multiply(bkg)
                    unosc_bkg_name = 'unosc_' + bkg_name
                    for part in eventsparts:
                        prod.multiply(part)
                    hist = detector.back_hists[unosc_bkg_name] = histcls(integrator)
                    hist.hist.setLabel('bkg hist: | '+detector.name)
                    detector.back_hists[unosc_bkg_name].hist.inputs(prod)

                    # normalize isotope flux now
                    iso_name = bkg_name.split('_')[-1]
                    with self.ns("geo_isotopes")(iso_name):
                        norm_geo = ROOT.GeoNeutrinoFluxNormed(detector.livetime[0]/year)
                    norm_geo.flux_norm.flux(detector.back_hists[unosc_bkg_name].hist)

                    with self.ns("oscillation"):
                        aver_oscprob = ROOT.OscProbAveraged(ROOT.Neutrino.ae(), ROOT.Neutrino.ae())
                    aver_oscprob.average_oscillations.inputs(norm_geo.flux_norm)
                    detector.back_hists['geo_'+iso_name] = norm_geo.flux_norm.normed_flux
                    bkg_summary.add(aver_oscprob.average_oscillations.flux_averaged_osc)

            detector.unoscillated_hist = histcls(integrator)
            res = ROOT.Product()
            res.product.setLabel('unosc IBD: | '+detector.name)
            detector.unoscillated_hist.hist.setLabel('unoscillated hist: | '+detector.name)
            res.multiply(detector.unoscillated)
            for part in eventsparts:
                res.multiply(part)
            detector.unoscillated_hist.hist.inputs(res)

            if self.opts.backgrounds:
                detector.unoscillated_with_bkg = ROOT.Sum()
                detector.unoscillated_with_bkg.sum.setLabel('noosc+bkg | '+detector.name)
                detector.unoscillated_with_bkg.add(bkg_summary)
                detector.unoscillated_with_bkg.add(detector.unoscillated_hist)
                detector.back_hists['sum_bkg'] = bkg_summary

    def _sumcomponents(self, components, name):
        oscprobs = {}
        for oscprobcls, compname in components.keys():
            if oscprobcls in oscprobs:
                continue
            with self.ns("oscillation"):
                oscprob = oscprobcls(ROOT.Neutrino.ae(), ROOT.Neutrino.ae(),
                                     freevars=['L'])
            oscprobs[oscprobcls] = oscprob

        for compid, comp in components.items():
            ps = oscprobs[compid[0]].probsum
            ps.setLabel('osc prob: | '+name)
            ps[compid[1]](comp)

        probsums = [oscprob.probsum for oscprob in oscprobs.values()]
        if len(probsums) > 1:
            finalsum = ROOT.Sum()
            finalsum.sum.setLabel('finalsum | '+name)
            for probsum in probsums:
                finalsum.add(probsum)
            return finalsum
        else:
            return probsums[0]

    def setupobservations(self, detector):
        """Sum over the components and setup the observables to namespace"""

        self.ns.addobservable("{0}_unoscillated".format(detector.name), detector.unoscillated_hist.hist, export=False)
        if self.opts.backgrounds:
            self.ns.addobservable("{0}_unoscillated_with_bkg".format(detector.name),
                           detector.unoscillated_with_bkg, export=False)

        sums = {resname: self._sumcomponents(detector.hists[resname], detector.name)
                for resname in detector.hists}

        if 'oscprob' in sums:
            detector.intermediates["oscprob"] = sums['oscprob']
        if 'rate' in sums:
            inter_sum = ROOT.Sum()
            inter_sum.sum.setLabel('Observed IBD: | '+detector.name)
            sum_without_bkg = sums['rate']
            inter_sum.add(sum_without_bkg)
            if self.opts.backgrounds:
                inter_sum.add(detector.back_hists['sum_bkg'])

            self.ns.addobservable("{0}_noeffects".format(detector.name), inter_sum, export=False)

            if self.opts.with_C14:
                with self.ns('ibd'), detector.ns:
                    detector.c14 = ROOT.C14Spectrum(8, 5)
                    detector.c14.smear.inputs(inter_sum)
                    finalsum = detector.c14.smear

                    self.ns.addobservable("{0}_c14".format(detector.name),
                                          finalsum, export=True)
            else:
                finalsum = inter_sum

            with detector.ns:
                detector.eres = ROOT.EnergyResolution()
            detector.eres.matrix.Edges(self.integrator.points.xhist)
            detector.eres.smear.setLabel('Eres')
            detector.eres.smear.Ntrue(finalsum)
            finalsum = detector.eres.smear

            if detector.edges_final is not None:
                self.ns.addobservable("{0}_fine".format(detector.name), finalsum)

                detector.rebin = C.Rebin(detector.edges_final, 6)
                finalsum >> detector.rebin
                finalsum = detector.rebin
                self.ns.addobservable("{0}".format(detector.name), finalsum.single())
            else:
                self.ns.addobservable("{0}".format(detector.name), finalsum)

        det_ns = self.ns("detectors")(detector.name)

        for name in detector.intermediates:
            det_ns.addobservable(name, detector.intermediates[name], export=False)

        for name, bkg in detector.back_hists.items():
            det_ns.addobservable("bkg_{}".format(name), bkg, export=False)

        for pair, comps in self.oscprobs_comps.items():
            pair_ns = det_ns(pair[1].name)
            pair_ns.addobservable('oscprob', self._sumcomponents(comps, detector.name), export=False)

    @classmethod
    def reqparameters(cls, ns):
        """Setup the parameters"""
        gna.parameters.ibd.reqparameters(ns("ibd"))
        gna.parameters.oscillation.reqparameters(ns("oscillation"))

        for isoname, (central, sigma) in eperfission.items():
            isons = ns("isotopes")(isoname)
            isons.reqparameter("EnergyPerFission", central=central, sigma=sigma)
        for geo_isoname, (central, relsigma) in geo_flux_normalizations.items():
            geo_isons = ns("geo_isotopes")(geo_isoname)
            geo_isons.reqparameter("FluxNorm", central=central, relsigma=sigma)


import numpy as np
import pybullet as pb
import cv2
import time
import os

import nuro_arm.transformation_utils as tfm
from nuro_arm.camera import camera_utils
from nuro_arm.camera.gui import GUI
from nuro_arm.camera.capturer import Capturer, SimCapturer
from nuro_arm import constants

class Camera:
    def __init__(self,
                 camera_type='sim',
                 camera_id=None,
                 pose_mtx=None,
                 free_floating=True,
                 run_async=True,
                 pb_client=0):
        '''Camera class that controls real or simulated camera feed, and is
        calibrated to allow for conversion from pixel to cartesian space

        Parameters
        ----------
        camera_type : str, default to 'real'
            Whether to connect to a image feed from a real or simulated camera.
            Must be either 'real' or 'sim'
        camera_id : int, optional
            Camera number to read frames from.  The default value is None,
            meaning that the camera number will be taken from the config file.
            Use this parameter to override the config file.  This is **not**
            used for simulated camera
        pose_mtx : array_like, optional
            4x4 Transformation matrix describing pose of camera in base frame.
            Only used for simulated camera; for real camera, pose is calculated
            during calibration.  If not provided, then the default camera
            pose from `nuro_arm.constants` is used.
        free_floating : bool, default to False
            If True, then the collision objects of the camera and camera stand
            are not added to simulator.  This can be used for a real or simulated
            camera, and will impact collision detection by the robot.  Only set to
            False if the camera is well out of range of arm.
        pb_client : int, optional
            Physics client id number in which the camera will be placed.  This is
            used even for real camera, where it controls which simulator the
            collision objects are added to.

        Attributes
        ----------
        cap : obj
            Capturer instance used to read frames from camera
        gui : obj
            GUI instance used for plotting frames
        '''
        self._camera_id = 0
        self._pb_client = pb_client

        self.free_floating = free_floating
        self.camera_collision_obj = None
        self.rod_collision_obj = None

        self.camera_type = camera_type
        assert camera_type in ('real', 'sim'), 'Invalid argument for camera_type'
        if camera_type == 'real':
            self.cap = Capturer()
            self.load_configs()
            # override camera_id config
            if camera_id is not None:
                self._camera_id = camera_id
            self.add_collision_objects()
        else:
            self.cap = SimCapturer()
            if pose_mtx is None:
                pose_mtx = constants.DEFAULT_CAM_POSE_MTX
            self.set_location(pose_mtx)

        self.gui = GUI(self.cap)

        if not self.connect(run_async):
            print(f'[ERROR] Failed to connect to camera{camera_id}.')

    def connect(self, run_async):
        '''Sets up connection to camera

        Returns
        -------
        bool
            True if connection was successful, False otherwise
        '''
        is_connected = self.cap.set_camera_id(self._camera_id, run_async)
        return is_connected

    def change_camera_id(self, cam_id):
        '''Change connection to another camera id.  If already connected to
        specified camera id, then nothing happens

        Parameters
        ----------
        cam_id : int
            camera id

        Returns
        -------
        bool
            True if connection was made, False if connection failed
        '''
        if cam_id == self._camera_id:
            return True
        self._camera_id = cam_id
        return self.cap.set_camera_id(self._camera_id)

    @property
    def frame_rate(self):
        '''Frame rate of the Capturer attribute

        Returns
        -------
        int
            Frames per second
        '''
        return self.cap._frame_rate

    def calc_location(self):
        '''Determine camera pose in world frame by localizing checkerboard pattern

        Returns
        -------
        bool
            True if location was determined, False otherwise.  Failure will occur
            if camera did not return frame or if checkerboard pattern could
            not be identified in the image.
        tuple
            rvec : ndarray
                see cv2.Rodrigues for info on rotation vector
            tvec : ndarray
                translation vector
            world2cam : ndarray
                4x4 transformation matrix that transforms homogeneous vector from
                world coordinate frame to camera coordinate frame
            cam2world : ndarray
                4x4 transformation matrix that transforms homogeneous vector from
                camera coordinate frame to world coordinate frame
        '''
        if self.camera_type == 'sim':
            print('[WARNING:] Calibration is not supported for simulated camera. '
                  'Simulated cameras location must be set by user.')
            return False, None

        gh, gw = constants.CALIBRATION_GRIDSHAPE
        gsize = constants.CALIBRATION_GRIDSIZE

        img = self.get_image()
        gray = camera_utils.convert_gray(img)

        ret, corners = cv2.findChessboardCorners(gray, (gh, gw), None)

        if ret:
            # coordinates of grid corners in world coordinates
            objp = np.zeros((gh*gw, 3), np.float32)
            objp[:, :2] = gsize * np.dstack(np.mgrid[1:gw+1, -1:gh-1]).reshape(-1, 2)
            objp += constants.TVEC_WORLD2RIGHTFOOT

            criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
            corners2 = cv2.cornerSubPix(gray, corners, (11,11),(-1,-1), criteria)

            ret, rvec, tvec = cv2.solvePnP(objp, corners2,
                                           constants.NEW_CAM_MTX,
                                           np.zeros(5))

            world2cam = tfm.transformation_matrix(rvec, tvec)
            cam2world = tfm.invert_transformation_matrix(world2cam)

            return True, (rvec, tvec, world2cam, cam2world)

        return False, None

    def set_location(self, pose_mtx):
        '''Set camera pose in world frame, only valid for simulated camera

        Parameters
        ----------
        pose_mtx : array_like
            4x4 pose matrix of camera in world frame
        '''
        if self.camera_type == 'real':
            print("[WARNING:] Real camera's pose cannot be set by user.")
            return

        self.cap.set_pose_mtx(pose_mtx)
        self._cam2world = pose_mtx
        self._world2cam = tfm.invert_transformation_matrix(pose_mtx)
        self._rvec, self._tvec = tfm.unpack_rvec_tvec(self._world2cam)

        # update collision object location
        self.add_collision_objects()

    def add_collision_objects(self):
        if self.free_floating:
            return

        self.remove_collision_objects()

        cam_pos, cam_quat, rod_pos, rod_quat = self._unpack_camera_pose(self._cam2world)

        camera_urdf_path = os.path.join(constants.URDF_DIR, 'camera.urdf')
        rod_urdf_path = os.path.join(constants.URDF_DIR, 'camera_rod.urdf')
        self.camera_collision_obj = pb.loadURDF(camera_urdf_path, cam_pos, cam_quat,
                                                physicsClientId=self._pb_client)
        self.rod_collision_obj = pb.loadURDF(rod_urdf_path, rod_pos, rod_quat,
                                             physicsClientId=self._pb_client)

    def remove_collision_objects(self):
        if self.camera_collision_obj is not None:
            pb.removeBody(self.camera_collision_obj,
                          physicsClientId=self._pb_client)
            self.camera_collision_obj = None
        if self.rod_collision_obj is not None:
            pb.removeBody(self.rod_collision_obj,
                          physicsClientId=self._pb_client)
            self.rod_collision_obj = None

    def start_recording(self, duration):
        '''Starts recording on camera

        Parameters
        ----------
        duration : float
            Seconds to record camera feed
        '''
        self.cap.start_recording(duration)

    def end_recording(self):
        '''Ends recording on camera
        '''
        self.cap.end_recording()

    def wait_for_recording(self):
        '''Waits until camera is done recording, then returns recorded frames

        Returns
        -------
        ndarray
            sequence of images
        '''
        while self.cap.is_recording():
            time.sleep(0.1)
        return self.cap.get_recording()

    def project_world_points(self, pts_wframe):
        '''Projects 3D world points to pixel locations in camera feed

        Parameters
        ----------
        pts_wframe : ndarray
            sequence of 3D arrays representing x,y,z location in world
            coordinate frame

        Returns
        -------
        ndarray
            sequence of 2D pixel indices; shape=(*,2); dtype=float
        '''
        return camera_utils.project_to_pixels(pts_wframe, self._cam2world)

    def load_configs(self):
        '''Reads config file writing to private attributes

        Returns
        -------
        bool
            True if configs were loaded succesfully, False otherwise
        '''
        if os.path.exists(constants.CAMERA_CONFIG_FILE):
            data = np.load(constants.CAMERA_CONFIG_FILE, allow_pickle=True).item()
            try:
                self._camera_id = data.get('camera_id')
                self._rvec = data.get('rvec')
                self._tvec = data.get('tvec')
                self._world2cam = data.get('world2cam')
                self._cam2world = data.get('cam2world')
                return True
            except IndexError:
                pass

        print('[WARNING] Camera config file not found. '
              ' Calibration should be performed.')
        return False

    def get_image(self):
        '''Get frame from capturer

        Returns
        -------
        img : ndarray
        '''
        img = self.cap.read()
        return img

    @property
    def configs(self):
        return {'camera_id': self._camera_id,
                'rvec': self._rvec,
                'tvec': self._tvec,
                'world2cam': self._world2cam,
                'cam2world': self._cam2world,
                }

    def _unpack_camera_pose(self, pose_mtx):
        '''Get params for positioning camera and rod based on pose of camera

        Parameters
        ----------
        pose_mtx: ndarray
            Transformation matrix from world frame to camera frame; shape=(4,4);
            dtype=float

        Returns
        -------
        cam_pos : array_like
            3d postion vector of camera body
        cam_quat : array_like
            quaternion of camera body
        rod_pos : array_like
            3d postion vector of rod body
        rod_quat : array_like
            quaternion of rod body
        '''
        cam_pos = pose_mtx[:3,3]
        cam_rotmat = pose_mtx[:3,:3]
        cam_quat = tfm.rotmat_to_quaternion(cam_rotmat)

        rod_offset_vec = np.array((0.026, -0.012, -0.013))
        rod_pos = cam_pos + np.dot(cam_rotmat, rod_offset_vec)
        rod_pos[2] = 0
        rod_quat = (0,0,0,1)

        return cam_pos, cam_quat, rod_pos, rod_quat

    def find_cubes(self, cube_size=None, tag_size=None):
        '''Identify cubes in image, given that an Aruco tag is located in the
        center of one of its visible faces.

        Parameters
        ----------
        tag_size: float, optional
            Size of aruco tag in meters.  If not provided, the default value
            from `nuro_arm.constants` will be used.
        cube_size: float, optional
            Side length of cube in meters.  If not provided, the default value
            from `nuro_arm.constants` will be used.

        Returns
        -------
        list
            List of ArucoCube namedtuples, see `nuro_arm/camera/camera_utils.py`
        '''
        img = self.get_image()
        return camera_utils.find_cubes(img, self._cam2world, cube_size, tag_size)

    def get_pb_client(self):
        return self._pb_client

    def __call__(self):
        return self.get_image()


from __future__ import print_function
from builtins import zip
from builtins import str
from .. import ParallelSampler
import numpy as np
import ctypes as ct
import os
import sys
import collections

MINUIT_INI_SECTION = "minuit"

loglike_type = ct.CFUNCTYPE(
    ct.c_double, #likelihood
    ct.POINTER(ct.c_double),  #parameter cube
)

class MinuitOptionsStruct(ct.Structure):
    _fields_ = [
        ("max_evals", ct.c_int),
        ("strategy", ct.c_int),
        ("algorithm", ct.c_int),
        ("save_cov", ct.c_char_p),
        ("tolerance", ct.c_double),
        ("width_estimate", ct.c_double),
        ("do_master_output", ct.c_int),
    ]

libname=os.path.join(os.path.split(__file__)[0],"minuit_wrapper.so")

class MinuitSampler(ParallelSampler):
    needs_output = False
    needs_parallel_output = False
    libminuit = None

    def config(self):
        self.converged = False
        if MinuitSampler.libminuit is None:
            if not os.path.exists(libname):
                raise ValueError("The CosmoSIS minuit2 wrapper was not compiled.  If you installed CosmoSIS manually you need the library minuit2 to use this sampler. See the wiki for more details.")
            MinuitSampler.libminuit = ct.cdll.LoadLibrary(libname)
        self.maxiter = self.read_ini("maxiter", int, 1000)
        self.save_dir = self.read_ini("save_dir", str, "")
        self.save_cov = self.read_ini("save_cov", str, "")
        self.output_ini = self.read_ini("output_ini", str, "")
        self.verbose = self.read_ini("verbose", bool, False)
        self.iterations = 0

        #Minuit options
        self.strategy = self.read_ini("strategy", str, "medium").lower()
        self.algorithm = self.read_ini("algorithm", str, "migrad").lower()

        self.width_estimate = self.read_ini("width_estimate", float, 0.05)
        self.tolerance = self.read_ini("tolerance", float, 50.0)
        self.neval = 0
        self.param_vector = self.pipeline.start_vector()  #initial value

        strategy = {
            "fast":0,
            "medium":1,
            "safe":2,
            "0":0,
            "1":1,
            "2":2,
        }.get(self.strategy)
        if strategy is None:
            raise ValueError("Minuit 'strategy' parameter must be one of fast, medium, safe (default=medium) not '%s'"%self.strategy)
        self.strategy = strategy

        algorithm = {
            "migrad":0,
            "fallback":1,
            "combined":1,
            "simplex":2,
            # "scan":3,
        }.get(self.algorithm)
        if algorithm is None:
            raise ValueError("Minuit 'algorithm' parameter must be one of migrad, fallback, simplex (default=migrad) not '%s'"%strategy)
        self.algorithm = algorithm

        self._run = MinuitSampler.libminuit.cosmosis_minuit2_wrapper
        self._run.restype = ct.c_int
        vec = ct.POINTER(ct.c_double)
        self.ndim = len(self.pipeline.varied_params)

        cube_type = ct.c_double*self.ndim

        @loglike_type
        def wrapped_likelihood(cube_p):
            vector = np.frombuffer(cube_type.from_address(ct.addressof(cube_p.contents)))
            try:
                like, extra = self.pipeline.posterior(vector)
            except KeyboardInterrupt:
                sys.exit(1)
            self.iterations += 1
            if self.verbose:
                print(self.iterations, like, "   ",  "    ".join(str(v) for v in vector))
            return -like

        self.wrapped_likelihood = wrapped_likelihood

    def execute(self):
        #Run an iteration of minuit
        param_vector, param_names, like, data, status, made_cov, cov_vector = self.sample()

        #update the current parameters
        self.param_vector = param_vector.copy()
        self.neval += status

        if status == 0:
            print()
            print("SUCCESS: Minuit has converged!")
            print()
            self.save_results(param_vector, param_names, like, data)
            self.converged = True
            self.distribution_hints.set_peak(param_vector)

            if made_cov:
                cov_matrix = cov_vector.reshape((self.ndim, self.ndim))
                self.distribution_hints.set_cov(cov_matrix)
        elif self.neval > self.maxiter:
            print()
            print("MINUIT has failed to converge properly in the max number of iterations.  Sorry.")
            print("Saving the best fitting parameters of the ones we trid, though beware: these are probably not the best-fit")
            print()
            self.save_results(param_vector, param_names, like, data)
            #we actually just use self.converged to indicate that the
            #sampler should stop now
            self.converged = True

        else:
            print()
            print("Minuit did not converge this run; trying again")
            print("until we run out of iterations.")
            print()

    def save_results(self, param_vector, param_names, like, data):
        section = None

        if self.pool is not None:
            print()
            print("Note that the # of function calls printed above is not the total count for all cores, just for one core.")
            print()

        if self.output_ini:
          self.pipeline.create_ini(param_vector, self.output_ini)

        for name, value in zip(param_names, param_vector):
            sec,name=name.split('--')
            if section!=sec:
                print()
                print("[%s]" % sec)
                section=sec
            print("%s = %g" % (name,value))
        print()
        print("Likelihood = ", like)

        if self.save_dir:
            print("Saving best-fit model cosmology to ", self.save_dir)
            data.save_to_directory(self.save_dir, clobber=True)

        self.converged = True

    def sample(self):

        param_names = [str(p) for p in self.pipeline.varied_params]
        param_names_array = (ct.c_char_p * len(param_names))()
        param_names_array[:] = [p.encode('ascii') for p in param_names]

        master = 1 if self.pool is None or self.pool.is_master() else 0

        param_max = self.pipeline.max_vector()
        param_min = self.pipeline.min_vector()

        param_vector = self.param_vector.copy()

        cov_vector = np.zeros(self.ndim*self.ndim)
        made_cov = ct.c_int()

        options = MinuitOptionsStruct(
            #allow for more loops
            max_evals=self.maxiter-self.neval, strategy=self.strategy,
            algorithm=self.algorithm, save_cov=self.save_cov.encode('ascii'),
            tolerance=self.tolerance, width_estimate=self.width_estimate,
            do_master_output = master
            )

        status = self._run(self.ndim,
            param_vector.ctypes.data_as(ct.POINTER(ct.c_double)),
            param_min.ctypes.data_as(ct.POINTER(ct.c_double)),
            param_max.ctypes.data_as(ct.POINTER(ct.c_double)),
            self.wrapped_likelihood,
            param_names_array,
            cov_vector.ctypes.data_as(ct.POINTER(ct.c_double)),
            ct.byref(made_cov),
            options
            )

        #Run the pipeline one last time ourselves, so we can save the
        #likelihood and cosmology
        like, _, data = self.pipeline.likelihood(param_vector, return_data=True)
        return param_vector, param_names, like, data, status, made_cov, cov_vector

    def worker(self):
        self.sample()

    def is_converged(self):
        return self.converged