# -*- coding: utf-8 -*-
import numpy as np
import cPickle as pickle
import sys, os
import pandas as pd
from collections import defaultdict

if not os.path.exists('./data'):
    os.mkdir('./data')

def confine_pf(num, max_pos):
    return max(0, min(num + max_pos, max_pos + max_pos + 1))

def load_data(ftrain, frela, fvec, args, data_path):
    """
    Transform all information from file ./origin_data/train.txt into a list-dict structure
    """
    max_sent, max_pos = args.max_sent, args.max_pos
    rela2id, word2id = {}, {}

    with open(frela, "r") as f:
        #rela_size = f.readline().strip()
        while True:
            line = f.readline()
            if line == '':
                break
            rela2id[line.strip().split()[0]] = int(line.strip().split()[1])
    pickle.dump(rela2id, open(data_path + 'rela2id.pkl', 'wb'))

    with open(fvec, "r") as f:
        W_size, w_dim = f.readline().strip().split()[:2]
        W_size, w_dim = int(W_size), int(w_dim)

        # word embedding matrix
        W = np.zeros(shape=(W_size+2, w_dim), dtype='float32')
        word2id['BLANK'] = 0  # the padding vector
        for i in range(W_size):
            temp = f.readline().strip().split()
            word2id[temp[0]] = i+1
            for j in range(w_dim):
                W[i+1][j] = (float)(temp[j+1])
        word2id['UNK'] = W_size+1
        W[W_size+1] = np.random.uniform(-0.25, 0.25, w_dim)  # UNK
    pickle.dump(word2id, open(data_path + 'word2id.pkl', 'wb'))
    np.savez(data_path + 'WordMatrix', W)

    instances = list()
    with open(ftrain, "r") as f:
        while True:
            line = f.readline()
            if line == '':
                break

            line = line.strip().split()
            e1_id = line[0]
            e2_id = line[1]
            e1 = line[2]
            e2 = line[3]
            rela = line[4]
            rela = 'NA' if rela not in rela2id else rela
            sent = line[5:-1]

            # Entity position
            e1_pos, e2_pos = 0, 0
            for num in range(len(sent)):
                if sent[num] == e1:
                    e1_pos = num
                if sent[num] == e2:
                    e2_pos = num

            instance  = {"y": rela,
                      "y_id": rela2id[rela],
                      "text": " ".join(sent),
                      "sent": sent,
                      "e1": e1,
                      "e1_id": e1_id,
                      "e2": e2,
                      "e2_id": e2_id,
                      "pf1": e1_pos,
                      "pf2": e2_pos,
                      "num_words": len(sent)}
            instances.append(instance)

            if len(instances)%100000 == 0:
                print "%d sentences have beeen processed ..." % (len(instances))
    pickle.dump(instances, open(data_path + 'instances.pkl', 'wb'))

    #tripple2sents = defaultdict(list)
    rela2tripple = defaultdict(set)
    for num in range(len(instances)):
        instance = instances[num]
        tripple = instance['e1_id'] + "\t" +instance['e2_id'] + "\t" + str(instance['y_id'])
        #tripple2sents[tripple].append(num)
        rela2tripple[str(instance['y_id'])].add(tripple)
    #pickle.dump(tripple2sents, open(data_path + 'tripple2sents.pkl', 'wb'))
    pickle.dump(rela2tripple, open(data_path + 'rela2tripple.pkl', 'wb'))
    #print "There are %d tripples in training set" % (len(tripple2sents))

    data_size = len(instances)
    sent2id = np.zeros((data_size, max_sent), dtype = np.int32) #  Maxtrix of word id in sentences
    pos2id_1 = np.zeros((data_size, max_sent), dtype = np.int32)
    pos2id_2 = np.zeros((data_size, max_sent), dtype = np.int32)
    label2id = np.zeros((data_size), dtype = np.int32)

    for s in range(data_size):
        sent = instances[s]['sent']
        pf1 = instances[s]['pf1']
        pf2 = instances[s]['pf2']
        label2id[s] = int(instances[s]['y_id'])
        for i in range(max_sent):
            pos2id_1[s][i] = confine_pf(i-pf1, max_pos)
            pos2id_2[s][i] = confine_pf(i-pf2, max_pos)
            if i < len(sent):
                if sent[i] not in word2id:
                    sent2id[s][i] = word2id['UNK']
                else:
                    sent2id[s][i] = word2id[sent[i]]

    inputs = np.concatenate((sent2id, pos2id_1, pos2id_2), 1)
    np.savez(data_path + "trans_id", sentM=sent2id, labelM=label2id, pf1M=pos2id_1, pf2M=pos2id_2)
    np.savez(data_path + "inputs", inputs=inputs)

    max_l = np.max(pd.DataFrame(instances)["num_words"])
    print "data loaded!"
    print "number of sentences: " + str(sent2id.shape[0])
    print "Word Embedding matrix size: " + str(W.shape)
    print "max sentence length: " + str(max_l)
    print "sentM:%s, labelM:%s, pf1M:%s, pf2M:%s" % (str(sent2id.shape), str(label2id.shape), str(pos2id_1.shape), str(pos2id_2.shape))
    print "dataset created!"

    GAN_relations, rela2sents, sent2sents = process_data(instances, data_path, label2id)
    repetitive_tripples = find_repetitive_tripples(rela2id)

    return instances, inputs, label2id, W, GAN_relations, rela2sents, sent2sents, repetitive_tripples

def process_data(instances, data_path, label2id):
    """
    Generate GAN training dataset for each relation that
    has more than 1500 sentences
    ep2sents 293162
    sent2sents 510386
    label2id shape (570088,)
    """
    ep2sents, sent2sents, rela2sents = dict(), dict(), dict()
    for num in range(len(instances)):
        sent_id = num
        ep = str(instances[sent_id]["e1_id"]) + '-' + str(instances[sent_id]["e2_id"]) + '-' + str(instances[sent_id]["y_id"])
        if ep2sents.has_key(ep):
            ep2sents[ep].append(sent_id)
        else:
            ep2sents[ep] = [sent_id]
    pickle.dump(ep2sents, open(data_path + "ep2sents.pkl", "wb"))
    print "There are %d entity pairs in training set" % (len(ep2sents))

    # For repetitive sentences of one entity pair, we just treat them as one sentence during reinforcement training
    for ep, sids in ep2sents.items():
        if len(sids) == 0:
            sys.exit('Error')
        elif len(sids) == 1:
            sent2sents[sids[0]] = sids
        else:
            stat_rpt = dict()
            for sid in sids:
                sent_text = instances[sid]["text"]
                if stat_rpt.has_key(sent_text):
                    stat_rpt[sent_text].append(sid)
                else:
                    stat_rpt[sent_text] = [sid]
            for rpt_sids in stat_rpt.values():
                sent2sents[rpt_sids[0]] = rpt_sids
    pickle.dump(sent2sents, open(data_path + "sent2sents.pkl", "wb"))
    print "There are %d unique sentences in the traning set" % len(sent2sents)

    tripple2sents = defaultdict(list)
    for num in range(len(instances)):
        instance = instances[num]
        tripple = instance['e1_id'] + "\t" +instance['e2_id'] + "\t" + str(instance['y_id'])
        tripple2sents[tripple].append(num)
    pickle.dump(tripple2sents, open(data_path + 'tripple2sents.pkl', 'wb'))
    print "There are %d tripples in training set" % (len(tripple2sents))

    #label2id = np.load(data_path + "trans_id.npz")['labelM']
    for sid in sent2sents.keys():
        rela_id = label2id[sid]
        if rela2sents.has_key(rela_id):
            rela2sents[rela_id].append(sid)
        else:
            rela2sents[rela_id] = [sid]

    GAN_relations = list()
    for k, v in rela2sents.items():
        if len(v) > 1500 and k != 0:
            GAN_relations.append(k)
            print k, len(v)
    pickle.dump(GAN_relations, open(data_path + "GAN_relations.pkl", "wb"))
    pickle.dump(rela2sents, open(data_path + "rela2sents.pkl", "wb"))
    print 'There are %d relations in the training data' % (len(rela2sents))
    print '%d relations will be processed by this system, they are %s' % (len(GAN_relations), str(GAN_relations))

    return GAN_relations, rela2sents, sent2sents

def find_repetitive_tripples(rela2id):
    test = open('./origin_data/test.txt').readlines()
    test_dict = dict()
    for line in test:
        line = line.strip().split()
        e1_id, e2_id = line[0], line[1]
        rela = 'NA' if line[4] not in rela2id else line[4]
        key = e1_id + "\t" + e2_id + "\t" + str(rela2id[rela])
        test_dict[key] = 1

    train = open('./origin_data/train.txt').readlines()
    train_dict = dict()
    count = 0
    repetitive_tripples = dict()
    for line in train:
        line = line.strip().split()
        e1_id, e2_id = line[0], line[1]
        rela = 'NA' if line[4] not in rela2id else line[4]
        key = e1_id + "\t" + e2_id + "\t" + str(rela2id[rela])
        if test_dict.has_key(key):
            repetitive_tripples[key] = 1
            count += 1
    print "The size of repetitive_tripples: ", len(repetitive_tripples)
    pickle.dump(repetitive_tripples, open("./data/repetitive_tripples.pkl", "w"))

    return repetitive_tripples


import concurrent.futures
import math
import os
import pickle
import tempfile
from collections import OrderedDict
from itertools import repeat
from pathlib import Path

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import torch
from kmtools import sequence_tools

import pagnn
from pagnn import settings
from pagnn.prediction.gan import Args
from pagnn.training.gan import Args as ArgsTraining
from pagnn.training.gan.utils import generate_noise
from pagnn.types import DataSetGAN

settings.device = torch.device("cpu")

def main():
    args = Args.from_cli()

    # Training arguments
    args_training = ArgsTraining(root_path=args.work_path.parent)
    args_training.unique_name = args.work_path.name

    # Load network
    net_d = pagnn.models.AESeqAdjApplyExtra(
        "discriminator", hidden_size=args_training.hidden_size, bottleneck_size=1
    )
    net_g = pagnn.models.AESeqAdjApplyExtra(
        "generator",
        hidden_size=args_training.hidden_size,
        bottleneck_size=16,
        encoder_network=net_d,
    )

    net_g.load_state_dict(
        torch.load(
            args_training.work_path.joinpath("models")
            .joinpath(f"net_g-step_{args.step}.model")
            .as_posix(),
            map_location="cpu",
        )
    )
    net_d.load_state_dict(
        torch.load(
            args_training.work_path.joinpath("models")
            .joinpath(f"net_d-step_{args.step}.model")
            .as_posix(),
            map_location="cpu",
        )
    )

    # Load input data
    if args.input_file.suffix == ".pdb":
        raise NotImplementedError("PDB support coming soon...")
    elif args.input_file.suffix == ".parquet":
        dataset = read_parquet_input(args)
    else:
        raise NotImplementedError("Only PDB and Parquet inputs are supported!")

    # Generate sequences
    random_state = np.random.RandomState(int(os.getenv("SLURM_ARRAY_TASK_ID", "42")))
    with args.validation_dataset_file.open("rb") as fin:
        datagen_ref = pickle.load(fin)
    df = generate_examples(args, dataset, datagen_ref, net_d, net_g, random_state)

    # Add secondary structure
    if args.include_ss:
        add_secondary_structure(args, df, copy=False)

    # Write output
    pq.write_table(
        pa.Table.from_pandas(df, preserve_index=False),
        args.output_file,
        version="2.0",
        flavor="spark",
    )

def read_parquet_input(args):
    rowgen = pagnn.io.iter_datarows(args.input_file)
    rowgen = list(rowgen)
    assert len(rowgen) == 1
    row = rowgen[0]
    dataset = pagnn.dataset.row_to_dataset(row, target=1)
    dataset = pagnn.dataset.dataset_to_gan(dataset)
    return dataset

def get_bottleneck_indices(adjs):
    idxs = []
    start = 0
    for i, adj in enumerate(adjs):
        stop = start + adj[4].shape[1]
        idxs.append((math.floor(start / 4), math.ceil(stop / 4)))
        start = stop
    assert idxs[-1][1] == math.ceil(sum(adj[4].shape[1] for adj in adjs) / 4)
    return idxs

def generate_examples(
    args: Args, dataset: DataSetGAN, datagen_ref: DataSetGAN, net_d, net_g, random_state
) -> pd.DataFrame:
    datavar = net_d.dataset_to_datavar(dataset)

    # TODO: Eventually, eval mode should work better
    net_g.train()
    net_d.train()

    # DataFrame columns
    df_columns = OrderedDict(
        [
            ("discriminator_score", [np.nan]),
            ("blosum62_score", [np.nan]),
            ("edit_score", [np.nan]),
            ("sequence", [dataset.seqs[0].decode()]),
            ("sequence_type", ["ref"]),
        ]
    )

    # Randomly select a reference dataset for each batch
    ref_idxs = random_state.choice(np.arange(len(datagen_ref)), args.batch_size - 1, replace=False)
    dataset_ref = [datagen_ref[i] for i in ref_idxs]
    datavar_ref = [net_d.dataset_to_datavar(ds) for ds in dataset_ref]
    seqs_ref = [dv.seqs[0:1, :, :] for dv in datavar_ref]
    adjs_ref = [dv.adjs for dv in datavar_ref]

    # Create a batch
    seqs_batch = torch.cat(
        seqs_ref[: args.batch_size // 2] + [datavar.seqs] + seqs_ref[args.batch_size // 2 :], 2
    )
    adjs_batch = (
        adjs_ref[: args.batch_size // 2] + [datavar.adjs] + adjs_ref[args.batch_size // 2 :]
    )

    bottleneck_idxs = get_bottleneck_indices(adjs_batch)
    assert len(bottleneck_idxs) == args.batch_size == len(adjs_batch)

    noise = generate_noise(net_g, adjs_batch)

    for _ in range(args.nseqs):
        noisev = noise.normal_(0, 1)
        with torch.no_grad():
            fake_seq = net_g(noisev, adjs_batch).data
            assert fake_seq.shape[2] == seqs_batch.shape[2]
            fake_pred = net_d(fake_seq, adjs_batch)
            assert bottleneck_idxs[-1][1] <= fake_pred.shape[2] <= (bottleneck_idxs[-1][1] + 1)

        start = 0
        for i, adj in enumerate(adjs_batch):
            stop = start + adj[0].shape[1]
            if i == args.batch_size // 2:
                fake_seq_slice = fake_seq[:, :, start:stop]
            start = stop
        assert stop == fake_seq.shape[2]

        for i, (start, stop) in enumerate(bottleneck_idxs):
            if i == args.batch_size // 2:
                fake_pred_slice = fake_pred[:, :, start:stop]
        assert stop <= fake_pred.shape[2] <= (stop + 1)

        fake_seq_slice_onehot = pagnn.utils.argmax_onehot(fake_seq_slice)
        assert (fake_seq_slice_onehot.sum(dim=1) == 1).all()

        df_columns["discriminator_score"].append(float(fake_pred_slice.sigmoid().mean()))
        df_columns["blosum62_score"].append(
            float(pagnn.utils.score_blosum62(datavar.seqs[0, :, :].data, fake_seq_slice_onehot))
        )
        df_columns["edit_score"].append(
            float(pagnn.utils.score_edit(datavar.seqs[0, :, :].data, fake_seq_slice_onehot))
        )
        df_columns["sequence"].append(
            "".join(pagnn.AMINO_ACIDS[int(i)] for i in np.argmax(fake_seq_slice.data.numpy(), 1)[0])
        )
        df_columns["sequence_type"].append("gen")

    df = pd.DataFrame(df_columns, index=range(len(df_columns["sequence"])))
    assert (df["sequence"].str.len() == len(df["sequence"][0])).all()
    return df

def run_psipred(idx, sequence, hhblits_database, tmp_dir):
    fasta_file = Path(tmp_dir).joinpath(f"seq_{idx}.fasta")
    with fasta_file.open("wt") as fout:
        fout.write(f"> seq_{idx}\n")
        fout.write(sequence + "\n")
    psipred_file = sequence_tools.run_psipred(fasta_file, hhblits_database)
    psipred = sequence_tools.read_psipred(psipred_file)
    return psipred

def add_secondary_structure(args, df: pd.DataFrame, copy=True) -> pd.DataFrame:
    df = df.copy() if copy else df

    with concurrent.futures.ProcessPoolExecutor(
        args.nprocs
    ) as p, tempfile.TemporaryDirectory() as tmp_dir:
        df["sequence_ss"] = list(
            p.map(
                run_psipred,
                df.index,
                df["sequence"].values,
                repeat(args.hhblits_database),
                repeat(tmp_dir),
            )
        )

    return df

if __name__ == "__main__":
    main()


#Main file for namsless should splinter later in development

#UNITS
#time:s
#weight:g
#length:cm
#density: g/cm^3

import numpy as np
import scipy.signal
import scipy.optimize
import matplotlib.pyplot as plt
from collections import OrderedDict
import os
import datetime
import sys

SMALL_SIZE = 16
MEDIUM_SIZE = 18
BIGGER_SIZE = 24

plt.rc('font', size=SMALL_SIZE)          # controls default text sizes
plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title
plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize
plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title

def pull_spectrum(spec_file):
    """
    Function to extract the spectrum from a *.Spe file
    Parameters:
        spec_file - file name of *.Spe file
    Returns:
        spectrum - numpy array of the number of counts per bin
    """
    with open(spec_file) as data:
        spectrum = np.array([float(i.strip()) for i in data.readlines()[12:2060]])
    return spectrum

def smooth(y, box_pts):
    """
    Simple 1-D convolution function
    Parameters:
        y - 1-D numpy array to be smoothed
        box_pts - width of kernel
    Returns:
        y_smooth - smoothed version of y
    """
    box = np.ones(box_pts)/box_pts
    y_smooth = np.convolve(y, box, mode='same')
    return y_smooth

class ATen:
    def __init__(self,input_filepath):
        self.banner = """                     _______     _________ _______  _
                    (  ___  )    \__   __/(  ____ \( (    /|
                    | (   ) |       ) (   | (    \/|  \  ( |
                    | (___) | _____ | |   | (__    |   \ | |
                    |  ___  |(_____)| |   |  __)   | (\ \) |
                    | (   ) |       | |   | (      | | \   |
                    | )   ( |       | |   | (____/\| )  \  |
                    |/     \|       )_(   (_______/|/    )_)

               "The People's 1D Continuous Energy Attenuation Code"

                             ...................
                             .. . ..............
             ......................,.,~$OOZ7O+.....................
             .....................Z$ZOO8OOZDMZ~....................
             .....................I8DDD8O88D888I...................
             ......................MNMNDDD888888:..................
             .......................$MMMMMNN888O8..................
             ........................=MNNNDDDDOOOI.................
             .........................=DDDDD88OZ88.................
             .........................ID88DD88OO88$................
             ......................,.$8DD8DDDD8OO88I...............
             ....................~OD88DN8DDD8DDDOO88...............
             ...............,IOOZ8888M=+DD88O888DOO8Z..   .........
             ............?O$OO8DNNMZ,..$DD8OO8888D8OZZZ   .........
             .........:ZO8DDMM8........DD888O8888OOZM.... .........
             ....,=8MNNNM...... .......ND888888D8ZD8D..............
             ... ...?I?7,,.............DD8D8888D8D88N..............
             ...   .... .    ..........+DDDDDDDDDD8Z8:.............
             ..........................DDDDDDDDDDDD8O8,............
             .........................88DDNDNDDDNNN8O88+...........
             ........................O8D8DDNNNNNNNNDO88D...........
             .......................+8DDDNNNMNMNNNNDOO887..........
             ......................,NDDNNNM.MMNNODNNOO88D..........
             .....................=8DNMN=:8NNDOO8NMMDOO8N:.........
             ....................:8DN~.......?:.....NO88D,.........
             ....................Z88ND.... ..+......,NO8D..........
             ................... .D888M. ....7.......Z888..........
             ......................OD8DM,....7......7OO8:..........
             ........................DDOM...$8$:....OOO8:..........
             ................... ......IODMMMDIM7..+OO88...........
             .........................~N8NNN8M7DMNMOOODN...........
             ........................D8M$N=DD:Z$NNM8OODO...........
             ...................  .,NDONDDD$8:ZM..M888N:...........
             ....................  D8ZM..8.$O,8I.ZDOO8D............
             ................... .?DDO7..?D$8+D.INDO888............
             ....................,DNO.MI~.IID==.=.8O8DN+. .........
             ................  ..O8$7..,D,.DM?~M.+OODMN8...........
             ................  ..DND.D$...N$$ON.Z:O8=8DM  .........
             ................... 8NN .:$DDIZDON,~8ZD$D8M ..........
             ....................8DI.  , =DND$$$,88I.O8M...........
             ....................DDO $8?..N8DNZ.8ZDNID8D ..........
             ....................8N8 ...M+?=DM$.OO?.8DD$ ..........
             ....................MNN,7ZD $N$+N8ZZO:.DDM.. .........
             ....................ODO8D..=D.M.8+DON.Z8O+ ...........
             ................... .MN$=.ZI..O,N?8D~~NOM.............
             ...................  :DDDZ~I.8$:D,D$8NDM..............
             ..................... 7ND8D.,7..ODDDNNMO..............
             ................... . ..MDN8NM8MMMDDD7................
             ........................ :MNDDDDMNMI .................
             ...................... ....  ~$,:.... ................
                                   ................

"""

        self.input_filepath = input_filepath
        self.casename = self.input_filepath.split("/")[-1][:-3]
        self.working_directory = "./" + "/".join(self.input_filepath.split("/")[:-1]) + "/"
        self.script_dir = "/".join(os.path.realpath(__file__).split("/")[:-1]) + "/"

        os.chdir(self.working_directory)

        #Reading contecnts of input file
        with open(input_filepath) as input_file:
            self.raw_input = input_file.readlines()
            input_data = [line.strip() for line in self.raw_input if not (line[0] == "#" or line == "\n")]

        #Read raw input text into cards
        self.material_card = input_data[input_data.index('READ MATERIAL')+1:input_data.index('END MATERIAL')]
        self.geometry_card = input_data[input_data.index('READ GEOMETRY')+1:input_data.index('END GEOMETRY')]
        self.parameters_card = input_data[input_data.index('READ PARAMETERS')+1:input_data.index('END PARAMETERS')]
        self.paths_card = input_data[input_data.index('READ PATHS')+1:input_data.index('END PATHS')]

        #initializing material card
        self.material_process()

        #parsing geometry related inputs
        self.layer_materials = [a for a in self.geometry_card if a.split()[0] == "layer_material"][0].split()[2:]
        self.layer_thicknesses = np.array([a for a in self.geometry_card if a.split()[0] == "layer_thicknesses"][0].split()[2:]).astype(np.float32)

        #parsing parameter related inputs
        self.parameters = {x[0] : x[1] for x in [parameter.replace(" ", "").split("=") for parameter in self.parameters_card]}

        #parsing path related inputs
        self.paths = {x[0] : self.working_directory + x[1] for x in [parameter.replace(" ", "").split("=") for parameter in self.paths_card]}

        self.bin_calibration(self.paths["cs137spec_filepath"])

        self.source_spec = pull_spectrum(self.paths["source_filepath"])/float(self.parameters["source_time"])

        if "background_filepath" in self.paths:
            self.source_spec -= pull_spectrum(self.paths["background_filepath"])/float(self.parameters["background_time"])

        #identigy groups and strengths and assign "cross sections"
        self.id_groups(plot_spec_peaks = True, plot_norm_peaks = True)
        self.ac_process()

    def bin_calibration(self,cs137spec_filepath):
        """
        Function to set energy values to each bin. Must be based on cs137 spectrum.
        Parameters:
            self - does not pull any attributes from self
            cs137spec_filepath - file name  of a *.Spe file containing unattenuated spectrum from a Cs-137 source
        Returns:
            self.bin_energies - numpy array containing the estimated energies associated with each bin
        """
        def how_much_gaussian(n, x, y):
            def gaus(x,a0,mu,sigma):
                return a0*(sigma*np.sqrt(2*np.pi))**-1*np.exp(-(x-mu)**2/(2*sigma**2))

            a0 = np.trapz(y)
            mean0 =  x[y.argmax()]
            sigma0 =  sum(y*(x-mean0)**2)/n

            popt,pcov = scipy.optimize.curve_fit(gaus,x,y,p0=[a0,mean0,sigma0],maxfev=1000000)
            perr = np.sqrt(np.diag(pcov)).sum()

            if perr > 500:
                return False
            else:
                return True

        cs137_peak = 661.6 #keV

        #Smooth the spectrum and extract energy local maxima
        cs137_spec = pull_spectrum(cs137spec_filepath)
        plt.show()

        cs137_spec = smooth(cs137_spec,30)
        peaks, empty_dict = scipy.signal.find_peaks(cs137_spec)

        #calculate the prominence associated with each maxima and find the bin with the largest prominence
        prominences, left_bases, right_bases = scipy.signal.peak_prominences(cs137_spec, peaks)
        peaks = np.compress(prominences > prominences.mean(), peaks)
        left_bases = np.compress(prominences > prominences.mean(), left_bases)
        right_bases = np.compress(prominences > prominences.mean(), right_bases)

        gauss_errors = np.array([])
        for i in range(peaks.size):
            n = right_bases[i]-left_bases[i]
            x = np.linspace(left_bases[i],right_bases[i],n)
            y = cs137_spec[left_bases[i]:right_bases[i]]
            gauss_errors = np.append(gauss_errors, how_much_gaussian(n,x,y))

        peaks = np.compress(gauss_errors, peaks)
        prominences, left_bases, right_bases = scipy.signal.peak_prominences(cs137_spec, peaks)

        #Find bin of maximum energy
        maxEbin = peaks[prominences.argmax()]
        self.maxEbin = maxEbin
        self.cs137_spec = cs137_spec
        delE = cs137_peak/maxEbin

        #Create bin energy array
        self.bin_energies = np.arange(0, cs137_spec.size * delE, delE)

        return 1

    def id_groups(self, plot_spec_peaks = False, plot_norm_peaks = False, width_threshold = [30, 170], width_rel_height = 0.7):
        """
        Function to identify energy group and count rate of each energy group.
        Parameters:
            self - both bin_energies and source_spec must be defined
            plot_spec_peaks - boolean to plot full source spectrum with e group identified
            plot_norm_peaks = boolean to plot norm fits as well. Should only be true if plot_spec_peaks is true
            width_threshold - tuning parameter for acceptable peak width
        Returns:
            self.group_counts - count rate associated with each group
            self.group_energies - energy value associated with each group
        """
        #Define gaussian checker for later in function
        def is_gaussian(n, x, y, width_rel_height,be,plot_norm_peaks):
            def gaus(x,a0,mu,sigma):
                return a0*(sigma*np.sqrt(2*np.pi))**-1*np.exp(-(x-mu)**2/(2*sigma**2))

            a0 = np.trapz(y) #4253
            mean0 =  x[y.argmax()] #sum(x*y)/n #377.977
            sigma0 =  sum(y*(x-mean0)**2)/n # 114.967

            popt,pcov = scipy.optimize.curve_fit(gaus,x,y,p0=[a0,mean0,sigma0],maxfev=100000)
            perr = np.sqrt(np.diag(pcov)).sum()

            if perr > 30 or perr is np.nan:
                return (False, 0)
            else:
                if plot_norm_peaks:
                    plt.plot(be,gaus(x,*popt),'o', c = "mediumseagreen", MarkerSize=5, label = "Gaussian Approximation")
                return (True, popt[0])

        #initializing figure if necessary
        if plot_spec_peaks or plot_norm_peaks:
            plt.figure(figsize=[15,8])

        #Smooth signal to and easily identify spectrum peaks
        smoothed_spec = self.source_spec
        peaks, empty_dict = scipy.signal.find_peaks(smoothed_spec)

        #Apply criteria that prominences must be greater than the average promenance
        prominences, left_bases, right_bases = scipy.signal.peak_prominences(smoothed_spec, peaks)
        peaks = np.compress(prominences > prominences.mean(), peaks)

        #Apply criteria that widths must be greater than width_threshold
        widths, width_heights, leftips, rightips = scipy.signal.peak_widths(x = smoothed_spec, peaks = peaks, rel_height = width_rel_height)
        peaks = np.compress((widths > width_threshold[0]) * (widths < width_threshold[1]), peaks)
        leftips = np.floor(np.compress((widths > width_threshold[0]) * (widths < width_threshold[1]), leftips)).astype(np.int32)
        rightips = np.ceil(np.compress((widths > width_threshold[0]) * (widths < width_threshold[1]), rightips)).astype(np.int32)
        widths = np.compress((widths > width_threshold[0]) * (widths < width_threshold[1]), widths)

        #Apply criteria that shape must be appropriately gaussian
        #Hiding in here is getting the count rate under each peak as a return from the is_gaussian
        remove_i = np.array([]).astype(int)
        group_counts = np.array([])
        for i in range(peaks.size):
            n = rightips[i]-leftips[i]
            x = np.linspace(leftips[i],rightips[i],n)
            y = self.source_spec[leftips[i]:rightips[i]]
            gaussianality, group_count = is_gaussian(n,x,y, width_rel_height,self.bin_energies[leftips[i]:rightips[i]],plot_norm_peaks = plot_norm_peaks)
            if not gaussianality:
                remove_i = np.append(remove_i, i)
            else:
                group_counts = np.append(group_counts, group_count)

        peaks = np.delete(peaks,remove_i)
        leftips = np.delete(leftips,remove_i)
        rightips = np.delete(rightips,remove_i)

        #Add optional peak plotted to make sure peaks were done correctly
        if plot_spec_peaks:
            plt.plot(self.bin_energies,self.source_spec, c = "royalblue",label = "Source Spectrum")

            for i in range(leftips.size):
                plt.axvline(self.bin_energies[peaks[i]], c = "crimson", label = "Group Energy Value")

            plt.ylabel("Counts [#/s]")
            plt.xlabel("Energy [keV]")
            handles, labels = plt.gca().get_legend_handles_labels()
            by_label = OrderedDict(zip(labels, handles))
            plt.legend(by_label.values(), by_label.keys())
            plt.tight_layout()

        self.group_counts = group_counts
        self.group_energies = self.bin_energies[peaks]
        return 1

    def material_process(self):
        """
        Function to decompose entries in MATERIAL CARD from text into data structs
        Parameters:
            self - must have material_card specified
        Returns:
            self.materials - dictionary-containing two items, a density and a
            material composition list
        """
        self.materials = {}
        for entry in self.material_card:
            entry = entry.split()
            material = entry.pop(0)
            density = float(entry.pop(0))
            composition = [[int(entry[2*i]), float(entry[2*i-1])] for i in range(int(len(entry)/2))]
            self.materials[material] = [density, composition]

        return 1

    def ac_process(self):
        """
        Function to assign cross sections to each energy group for each material
        Parameters:
            self - must have self.materials defined
        Returns:
            self.ac - dictionary with keys of each material ID and values of list representing energies
                    - energy list as follows
                        - each row is a separate nuclide
                        - each column is an energy group
                        - sum each row for total ac for an energy group
        """
        def pull_ac(filename):
            #function to pull cross sections from libraries in workable form
            with open(filename, "r") as f:
                lib = np.array([s[3:].split() for s in f.readlines()])[:,:2].astype(np.float32)
            return lib

        self.ac = {}
        for key, value in self.materials.items():
            density = value[0]
            group_ac = np.empty((len(value[1])))
            for energy in self.group_energies:
                ac = np.array([])
                for znum, mass_frac in value[1]:
                    library = self.script_dir + "alib/" + [i for i in os.listdir(self.script_dir + "alib/") if "alib_" + str(znum) in i][0]
                    temp_ac = pull_ac(library)
                    ac = np.append(ac, np.interp(energy/1000,temp_ac[:,0],temp_ac[:,1])*mass_frac)

                group_ac = np.vstack((group_ac, ac*density))
            self.ac[key] = group_ac[1:,:] #some adjustment that needs to be made because of array initialization
        return 1

    def compute(self):
        def solve_master(master_ara, div_rows):
            for i in range(master_ara.shape[0]):
                if np.isclose(master_ara[i,1],0):
                    thickness = master_ara[i,2] - master_ara[i-1,2]
                    master_ara[i,1] = master_ara[i-1,1]*np.exp(-master_ara[i-1,3]*thickness)

        #setting up depth vector and n depending on "layer_mesh_divs"
        if "layer_mesh_divs" in self.parameters.keys():
            depth = np.array([0])
            material_list = []

            n = int(self.parameters["layer_mesh_divs"])

            for thickness, material in zip(self.layer_thicknesses,self.layer_materials):
                material_list += n*[material]
                start = depth[-1]
                stop = depth[-1]+thickness
                depth = np.append(depth, np.linspace(start + (stop-start)/(n+1) ,depth[-1]+thickness,n))
        else:
            n = 1
            depth = np.append(0,np.cumsum(self.layer_thicknesses))
            material_list = self.layer_materials

        #set up master array with columns energy, strength, depth, and atten coeff
        master_ara = np.zeros((depth.size*self.group_energies.size, 4))

        #create array with enery group repeated the same number of times as the length of depth vec
        working_grp = np.array([])
        for i in reversed(range(self.group_energies.size)):
            working_grp = np.append(working_grp, np.tile(self.group_energies[i],depth.size))

        #set up attenuation coeff vector
        ac_vec = np.array([])

        self.material_divisions = []

        for i in range(self.group_energies.size):
            ac_vec = np.append(0,ac_vec)
            for mat in material_list:
                ac_vec = np.append( self.ac[mat][i,:].sum(), ac_vec)
                self.material_divisions.append(mat)

        #place these vectors in master_ara
        master_ara[:,0] = working_grp
        master_ara[:,2] = np.tile(depth, self.group_energies.size)
        master_ara[:,3] = ac_vec

        #finding rows that start each energy group
        div_rows = np.array([0])
        for i in range(master_ara.shape[0]-1):
            if master_ara[i,3]==0:
                div_rows = np.append(div_rows, i+1)

        #initializing count rate for each energy group
        master_ara[div_rows,1] = self.group_counts

        #Fill in counts to each division in master_ara
        solve_master(master_ara, div_rows)

        self.master_ara = [master_ara,div_rows]

        return 1

    def print_output(self):
        def heading(string):
            return("\n" + 80*"-" + "\n" + string + "\n" + 80*"-" + "\n")

        #set up output and banner
        output = open(self.working_directory + self.casename + ".out", "w")
        output.write(self.banner + "\n" + str(datetime.datetime.now()) + "\n\n")

        #printing raw input
        output.write(heading("Raw Input"))
        output.write("".join(self.raw_input))

        #printing cards
        output.write(heading("Material Card"))
        output.write("\n".join(self.material_card) + "\n")
        output.write(heading("Geometry Card"))
        output.write("\n".join(self.geometry_card) + "\n")
        output.write(heading("Parameters Card"))
        output.write("\n".join(self.parameters_card) + "\n")
        output.write(heading("Paths Card"))
        output.write("\n".join(self.paths_card) + "\n")

        #printing materials
        output.write(heading("Materials"))
        for key, item in self.materials.items():
            #print(key,item[)
            output.write("Material: %s       Density: %.4f [g/cm^3]\n" % (key.ljust(15), item[0]))
            output.write("Atomic Number        Mass Fraction\n")
            for z,p in item[1]:
                output.write("%i                   %.5f\n"%(z,p))
            output.write("\n")

        #printing geomety
        output.write(heading("Geometry"))
        output.write("Material             Layer Thickness [cm]    Cumulative Thickness [cm]\n")
        accum = 0
        for i in range(len(self.layer_materials)):
            accum += self.layer_thicknesses[i]
            output.write("{:<17}".format(self.layer_materials[i]) + "    %.5f"%self.layer_thicknesses[i] + " "*17 + "%.5f"%(accum) + "\n")

        #printing group energy and strength
        output.write(heading("Energy Group Identification"))
        output.write("Group Energy [keV]                Group Count Rate [#/s]\n")
        for i in range(len(self.group_counts)):
            grp = "%.5f"%(self.group_energies[i])
            output.write(grp + (34 - len(grp))*" " + "%.5f\n"%(self.group_counts[i]))

        #printing master array
        output.write(heading("Master Array"))
        output.write("Grp Energy [keV]   Count Rate [#/s]   Depth [cm]   Atten Coeff [cm^2]   Material\n")

        i = 0
        for j, row in enumerate(self.master_ara[0]):
            e = "%.3f"%row[0]
            s = "%.3f"%row[1]
            d = "%.3f"%row[2]
            c = "%.5f"%row[3]

            if np.isclose(row[3],0):
                m = ""
            else:
                m = self.material_divisions[i]
                i += 1

            output.write(e + " "*(19-len(e)) + s + " "*(19-len(s)) + d + " "*(13-len(d))+ c + (21-len(c))*" " + m + "\n")

            if m == "":
                output.write(26*" - " + "\n")

        #attenuation coefficient libraries
        output.write(heading("Attenuation Coefficient Values"))
        for key, item in self.ac.items():
            output.write("Material = {mat}\n".format(mat = key))
            output.write("Atomic Number   Mass Fraction    ")
            for i, e in enumerate(self.group_energies):
                a = "E%i:%.2f"%(i,e)
                output.write(a + " "*(15-len(a)))
            output.write("\n")
            for i in range(item.T.shape[0]):
                anum = str(self.materials[key][1][i][0])
                frac = "%.4f"%self.materials[key][1][i][1]
                output.write(anum + " "*(16-len(anum)) + frac + " "*11 )
                for j in range(item.T.shape[1]):
                    const = "%.4f"%item.T[i,j]
                    output.write(const + " "*(15-len(const)))
                output.write("\n")
            output.write("\n")

        #printing calibration spectra with peak identified
        output.write(heading("Cs-137 Calibration Spectrum"))
        output.write("Bin Number           Count Number\n")
        for i in range(self.cs137_spec.size):
            output.write("{:<21}".format(str(i)) + str(self.cs137_spec[i]))

            if i == self.maxEbin:
                output.write("    *661.5 keV peak, used for calibration")

            output.write("\n")

        #printing source spectra
        output.write(heading("Unattenuated Spectra"))
        output.write("Bin Energy [keV]              Count Rate [#/s]\n")
        for i in range(self.source_spec.size):
            e = "%.5f"%self.bin_energies[i]
            s = "%.5f\n"%np.abs(self.source_spec[i])
            output.write(e +(30-len(e))*" " + s)

        output.close()

if __name__ == '__main__':
    ms_input = sys.argv[2]

    test = ATen(ms_input)
    test.compute()
    test.print_output()

    plt.show()


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

import time
import sys
import math
import rospy
from geometry_msgs.msg import PoseStamped, Point
from std_msgs.msg import Bool
import tf.transformations
from cv_bridge import CvBridge, CvBridgeError
from sensor_msgs.msg import Image
import cv2
from numpy import *

"""
Wripper for ros contolll of drone via python
"""

class drone_client():

    def __init__(self, ):
        # super(drone_client,self).__init__()
        """
        """
        self.current_position = None
        self.roll = 0.0
        self.pitch = 0.0
        self.yaw = 0.0
        self.cv_image = False

        self.bridge = CvBridge()
        rospy.Subscriber("/mavros/local_position/pose", PoseStamped, self.current_nav_clb)
        self.image_sub = rospy.Subscriber("/d400/color/image_raw", Image, self.image_clb)
        self.goal_pub = rospy.Publisher("/goal", PoseStamped,queue_size=10)
        self.gripper_pub = rospy.Publisher("gripper", Bool, queue_size=10)

        self.flag = 0
        self.basket = None
        self.cubes = None
        self.c = 0

    def current_nav_clb(self, data):
        """
        callback of current position
        :param data:
        :return:
        """
        self.current_position = data.pose.position

        self.roll, self.pitch, self.yaw = tf.transformations.euler_from_quaternion((data.pose.orientation.x,
                                                                                data.pose.orientation.y,
                                                                                data.pose.orientation.z,
                                                                                data.pose.orientation.w))

    def image_clb(self, data):
        """
        get image from camera
        :param data:
        :return:
        """
        try:
            self.cv_image = self.bridge.imgmsg_to_cv2(data, "bgr8")
        except CvBridgeError as e:
            print(e)

    def SetGoal(self, pose, yaw):
        """
        :param x: pose
        :param yaw:
        :return:
        """
        goal_msgs = PoseStamped()
        goal_msgs.header.frame_id = "map"
        goal_msgs.header.stamp = rospy.Time.now()
        goal_msgs.pose.position = pose

        quat = tf.transformations.quaternion_from_euler(0.,0.,yaw)
        goal_msgs.pose.orientation.x = quat[0]
        goal_msgs.pose.orientation.x = quat[1]
        goal_msgs.pose.orientation.x = quat[2]
        goal_msgs.pose.orientation.x = quat[3]
        self.goal_pub.publish(goal_msgs)

    def SetGripper(self, state):
        """
        open/close gripper

        :param state: true - is close, false - is open
        :return:
        """
        self.gripper_pub.publish(state)

    #def takeoff(self, height):
    def setup(self, x, y, z, w):
        dist = sqrt((x - self.current_position.x)**2 + (y - self.current_position.y)**2 + (z - self.current_position.z)**2)
        self.tp = dist/0.16
        self.t0 = time.time()
        self.x0 = self.current_position.x
	#print "1self.x0: ", self.x0
        self.y0 = self.current_position.y
        self.z0 = self.current_position.z

    def set_position(self, x, y, z, w):
        if (x!=self.current_position.x or y!=self.current_position.y or z!=self.current_position.z):
            if (time.time() - self.t0 < self.tp):
                t = time.time() - self.t0
                pose1 = Point()
                pose1.x = ((x - self.x0)/2)*(1 - cos(pi*t/self.tp)) + self.x0
                pose1.y = ((y - self.y0)/2)*(1 - cos(pi*t/self.tp)) + self.y0
                pose1.z = ((z - self.z0)/2)*(1 - cos(pi*t/self.tp)) + self.z0
                #print "pose1:", pose1
                #print "x:",pose1.x
                #print "self.x0:",self.x0
                print "flag: ",self.flag
                self.SetGoal(pose1, w)
            else:
                self.flag += 1
                time.sleep(1)
        else:
            self.flag+=1

    def get_basket(self):
        img = self.cv_image
        print "get basket"
        img = self.cv_image
        if img is None:
            print "no image"
            return
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        hsv_min = array((0, 0, 242), uint8)
        hsv_max = array((92, 255, 255), uint8)
        thresh = cv2.inRange(hsv, hsv_min, hsv_max)
        edged = cv2.Canny(thresh, 50, 100)
        edged = cv2.dilate(edged, None, iterations=1)
        edged = cv2.erode(edged, None, iterations=1)
        _,cnts, h = cv2.findContours(edged.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        r_meters = 0.2
        min_radius = 640*r_meters/(2*self.current_position.z*tan(radians(69/2)))
        min_radius = 75*min_radius/100
        x = -1
        y = -1
        for c in cnts:
            print "next cnt"
            if cv2.contourArea(c) < (2*pi*min_size):
                print "small"
                continue
            (x,y),radius = cv2.minEnclosingCircle(c)
            center = (int(x),int(y))
            radius = int(radius)
            cv2.circle(img,center,radius,(0,255,0),2)
            cv2.circle(img, center, 2, (0, 255, 0),2)
            if radius > min_radius:
                x = int(x) #- 640/2
                y = int(y) #- 480/2
                basket = [x, y]
                max_radius = radius
                cv2.circle(img, (basket[0], basket[1]), 3, (255, 0, 0), 2)
                basket[0] += 0.09
                basket[0] = (basket[0]/640)*2*h*tan(radians(69/2))
                basket[1] = (basket[1]/480)*2*h*tan(radians(42/2))
                h=self.current_position.z
                basket[0] = (basket[0])*h*0.00175
                basket[1] = (basket[1])*h*0.00175
        cv2.imwrite("basket.jpg", img)
        cv2.imwrite("thresh.jpg", thresh)
        cv2.imwrite("edged.jpg", edged)
        if x == -1:
            return None
        else:
            return (x - 640/2, y - 480/2)

    def get_cubes(self):
        self.c+=1
        print "get cubes"
        img = self.cv_image
        if img is None:
            print "no image"
            return
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        hsv_min = array((0, 0, 242), uint8)
        hsv_max = array((92, 255, 255), uint8)
        thresh = cv2.inRange(hsv, hsv_min, hsv_max)
        edged = cv2.Canny(thresh, 50, 100)
        edged = cv2.dilate(edged, None, iterations=1)
        edged = cv2.erode(edged, None, iterations=1)
        _,cnts, h = cv2.findContours(edged.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        centers = []
        num = -1
        r_meters = 0.06
        min_size = 640*r_meters/(2*self.current_position.z*tan(69/2))
        min_size = 85*min_size/100
        min_radius = 640*r_meters/(2*self.current_position.z*tan(radians(69/2)))
        min_radius = 85*min_radius/100
        for c in cnts:
            print "next cnt"
            if cv2.contourArea(c) < (2*pi*min_size):
                print "small"
                continue
            rect = cv2.minAreaRect(c)
            box = cv2.boxPoints(rect)
            box = int0(box)
            cv2.drawContours(img,[box],0,(0,0,255),2)
            area = sqrt((box[0][0] - box[1][0])**2 + (box[0][1] - box[1][1])**2) * sqrt((box[3][0] - box[0][1])**2 + (box[3][1] - box[1][1])**2)
            if (area > min_size) and (area < pi*min_radius**2):
                print "normal size"
                num += 1
                centers.append([(box[0][0] +  box[1][0] + box[2][0] + box[3][0])/4 - 640/2, (box[0][1] +  box[1][1] + box[2][1] + box[3][1])/4 - 480/2])
                cv2.circle(img,(centers[num][0]+640/2, centers[num][1]+480/2), 3,(0,0,0),2)
                centers[num][0] += 0.09
                h=self.current_position.z
                centers[num][0] *= -1
                centers[num][1] *= -1
                print "centers pixel" + str(centers[num][0]) + " " + str(centers[num][1])
                centers[num][0] = (centers[num][0])*h*0.00175
                centers[num][1] = (centers[num][1])*h*0.00175
                print "centers meters" + str(centers[num][0]) + " " + str(centers[num][1])
        cv2.imwrite("cubes"+str(self.c)+".jpg", img)
        cv2.imwrite("thresh"+str(self.c)+".jpg", thresh)
        cv2.imwrite("edged"+str(self.c)+".jpg", edged)
        if len(centers) == 0:
            return None
        else:
            return centers

    def rotate(self, x, y, z):
        print "rotete"
        roll = self.roll
        pitch = self.pitch
        yaw = self.yaw
        Rx = mat([[1, 0, 0],
             [0, cos(roll), -sin(roll)],
             [0, sin(roll), cos(roll)]])
        Ry = mat([[cos(pitch), 0, sin(pitch)],
             [0, 1, 0],
             [-sin(pitch), 0, cos(pitch)]])
        Rz = mat([[cos(yaw), -sin(yaw), 0],
             [sin(yaw), cos(yaw), 0],
             [0, 0, 1]])
        R = Rz*Ry*Rx
        X = R*mat([[x], [y], [z]])
        X[0] += self.current_position.x
        X[1] += self.current_position.y
        #X[2] += self.current_position.z
        return X


import cv2
import pyttsx3
import keyboard
import ImageToSpeech as ITS
import numpy as np
from keras.models import load_model

face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
cap = cv2.VideoCapture(0)
textEngine = pyttsx3.init()

stroke = 2
color = (255, 255, 255)
model = load_model('model_CNN_V2.h5')
font = cv2.FONT_HERSHEY_SIMPLEX
Gender = {0: "Male", 1: "Female"}

def normalizeImage(img):
    IMG_SIZ = 120
    new_img = cv2.resize(img, (IMG_SIZ, IMG_SIZ))
    image = new_img.reshape((120, 120, 1))
    image = image.astype('float32') / 255
    return image

i=0
while(True):
    i+=1
    ret, frame = cap.read()
    if keyboard.is_pressed('i'):
        cv2.imwrite('C:/Users/Jitender kumar/Desktop/Intel HackFury2/Image To String Data/{index}.jpg'.format(index=i), frame)
        ITS.convertImageToString('C:/Users/Jitender kumar/Desktop/Intel HackFury2/Image To String Data/{index}.jpg'.format(index=i))
    frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(frame_gray, 1.2, 5)
    for (x, y, w, h) in faces:
        cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
        roi_gray = frame_gray[y:y + h, x:x + w]     # region of interest
        roi_color = frame[y:y + h, x:x + w]
        img = np.array(roi_gray)
        img = normalizeImage(img)
        prediction = model.predict([[img]]).argmax()
        gender = Gender.get(prediction)
        text = "Some {} is coming towards you.".format(gender)
        if keyboard.is_pressed('p'):
            textEngine.say(text)
            textEngine.runAndWait()
            textEngine.stop()
        cv2.putText(frame, gender, (x, y), font, 1, color, stroke, cv2.LINE_AA)

    cv2.imshow('Eye For Blindness', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()


import serial, sys, time, os, gzip, math
import Sigma_koki
import numpy as np

class Stage(object):
    def __init__(self):
        #self.Range_Speed = 1
        self.Min_Speed = 50
        self.Max_Speed = 20000

        self.w_x = 1000 / 2. #/1mm SGSP 33-200
        self.w_z = 1000 / 6. #/1mm SGSP 26-200
        self.w_angle = 1 / 0.0025 #/deg OSMS-60YAW

    def Open(self):
        #LED-XZ stage
        self.gscXZ = Sigma_koki.SHOT702()
        self.gscXZ.open('/dev/tty.usbserial-FTRTNFQI')

        if self.gscXZ.getStatus() == '':
            raise ValueError('XZstage is invalid')

        #TARGET-rotation stage
        self.gscT = Sigma_koki.GSC02()
        self.gscT.open('/dev/tty.usbserial-FTRTNEFZ')

        if self.gscT.getStatus() == '':
            raise ValueError('Rotation stage is invalid')

        #self.gscT.returnToMechanicalOrigin('+', '-')

    def Set_Speed_parameter(self, rang, minS, maxS):
        #self.Range_Speed = rang
        self.Min_Speed = minS
        self.Max_Speed = maxS

    def stageX_Relative(self, n, AccelerationSpeed):
        self.gscXZ.setSpeed(self.Min_Speed, self.Max_Speed, AccelerationSpeed)

        self.gscXZ.move(0, n * self.w_x)
        time.sleep(6)

        a = self.gscXZ.getStatus()
        count = int(a[12:21])

        operation_vale = count / self.w_x
        print (" operation value of LED stageX = ", operation_vale, "(mm)")

    def stageX(n, AccelerationSpeed):
        self.gscXZ.returnToMechanicalOrigin(0, 1) # (0, Xstage) #

        while 1:
            aaa = self.gscXZ.getStatus()
            if aaa[-1] == 'R':
                break
        time.sleep(1)

        self.gscXZ.setSpeed(self.Min_Speed, self.Max_Speed, AccelerationSpeed)

        self.gscXZ.move(0, n * self.w_x)
        time.sleep(6)

        a = self.gscXZ.getStatus()
        count = int(a[12:21])

        operation_vale = count / self.w_x
        print (" operation value of LED stageX = ", operation_vale, "(mm)")

    def stageZ_Relative(self, n, AccelerationSpeed):
        self.gscXZ.setSpeed(self.Min_Speed, self.Max_Speed, AccelerationSpeed)

        self.gscXZ.move(n * self.w_z, 0)
        time.sleep(4)

        a = self.gscXZ.getStatus()
        count = int(a[1:10])

        operation_vale = count / self.w_z
        print (" operation value of LED stageZ = ", operation_vale, "(mm)")

    def stageZ(self, n, AccelerationSpeed):
        self.gscXZ.returnToMechanicalOrigin(1, 0) # (Zstage, 0) #
        while 1:
            aaa = self.gscXZ.getStatus()
            if aaa[-1] == 'R':
                break
        time.sleep(1)

        self.gscXZ.setSpeed(self.Min_Speed, self.Max_Speed, AccelerationSpeed)

        self.gscXZ.move(n * self.w_z, 0)
        time.sleep(4)

        a = self.gscXZ.getStatus()
        count = int(a[1:10])

        operation_vale = count / self.w_z
        print (" operation value of LED stageZ = ", operation_vale, "(mm)")

    def rotation_initialize(self):
        self.gscT.returnToMechanicalOrigin('-', 0) # (theta_stage, 0) #

        #print (" operation value of theta_stage = ", gscA.getStatus(), "(mm)")

    def rotation_initialize0deg(self):
        self.gscT.returnToMechanicalOrigin('-', 0) # (theta_stage, 0) #

        while 1:
            aaa = self.gscT.getStatus()
            if aaa[-1] == 'R':
                break
        time.sleep(2)

        self.rotation_Relative(90., 50) # must be set to TARGET

        while 1:
            aaa = self.gscT.getStatus()
            if aaa[-1] == 'R':
                break
        time.sleep(2)

        self.Set0pointAngle()
        self.rotation_Relative(0, 50)

    def rotation_Relative(self, angle, AccelerationSpeed):

        self.gscT.setSpeed(1, 50, 5000, AccelerationSpeed, 50, 5000, AccelerationSpeed) #MaxSpeed 20000 is bad for OSMS-60YAW

        self.gscT.move(angle * self.w_angle, 0)
        time.sleep(4)

        while 1:
            aaa = self.gscT.getStatus()
            if aaa[-1] == 'R':
                break
        time.sleep(2)

        a = self.gscT.getStatus()
        theta_value = int(a[1:10])

        if a[0] == '-':
            theta_value = - theta_value
        print (" Rotation Angle = ", theta_value / self.w_angle, "(deg)")

    def rotation(self, angle, AccelerationSpeed):

        self.gscT.returnToMechanicalOrigin('-', 0) # (theta_stage, 0) #

        while 1:
            aaa = self.gscT.getStatus()
            if aaa[-1] == 'R':
                break
        time.sleep(1)

        self.rotation_Relative(90., 50) # must be set to TARGET

        while 1:
            aaa = self.gscT.getStatus()
            if aaa[-1] == 'R':
                break
        time.sleep(1)

        self.Set0pointAngle()

        self.gscT.setSpeed(1, 50, 5000, AccelerationSpeed, 50, 5000, AccelerationSpeed) #MaxSpeed 20000 is bad for OSMS-60YAW

        self.gscT.move(angle * self.w_angle, 0)
        time.sleep(4)

        a = self.gscT.getStatus()
        theta_value = int(a[1:10])

        if a[0] == '-':
            theta_value = - theta_value
        print (" Rotation Angle = ", theta_value / self.w_angle, "(deg)")

    def Set0pointX(self):
        self.gscXZ.initializeOrigin(0, 1)

    def Set0pointZ(self):
        self.gscXZ.initializeOrigin(1, 0)

    def Set0pointAngle(self):
        self.gscT.initializeOrigin(1, 0)

    def stageX_initialize(self):
        '''
        changed by YK
        '''
        self.gscXZ.returnToMechanicalOrigin(0, 1) # (0, Xstage) #
        print (" operation value of theta_stage = ", self.gscXZ.getStatus(), "(mm)")

    def stageZ_initialize(self):
        '''
        changed by YK
        '''
        self.gscXZ.returnToMechanicalOrigin(1, 0) # (Zstage, 0) #
        print (" operation value of theta_stage = ", self.gscXZ.getStatus(), "(mm)")

    def GetX(self):
        a = self.gscXZ.getStatus()
        count = int(a[12:21])
        return round(count / self.w_x, 1)

    def GetZ(self):
        a = self.gscXZ.getStatus()
        count = int(a[1:10])
        return round(count / self.w_z, 1)

    def GetAngle(self):
        a = self.gscT.getStatus()
        theta_value = int(a[1:10])
        if a[0] == '-':
            theta_value = - theta_value
        return round(theta_value / self.w_angle, 2)

    def Wait_Rotationmove(self):
        while 1:
            aaa = self.gscT.getStatus()
            if aaa[-1] == 'R':
                break
        time.sleep(1)

if __name__ == '__main__':

    ##########################################################
    S = Stage()
    #gscA.initializeOrigin(1,0)
    S.Open()
    S.rotation_Relative(0,50) # rotation(angle[deg],speed[50])
    S.stageX_Relative(0,50)   # stageX(distance[mm],speed[50])

    print('A Rotation Status', S.GetX(), S.GetZ())
    print('B LEDstage Status', S.GetAngle())

    exit()
    ###########################################################


import numpy as np
import datetime
import timeit
import subprocess
import RPi.GPIO as gpio

print("Program initialized!")

# default alarm time variables
alarm_hour = 7 # X am default
alarm_minute = 0

Alarm_Lights = False
Alarm_Lights_Flash = False

Audio_Enable = True
Play_Music = False
Power_Speakers = False
Music_Start_OneShot = False

Change_Flag = False

Alarm_Set_Btn = False

## music files and info
music_files = ['/home/pi/Desktop/theintro.mp3', '/home/pi/Desktop/LifeCycles.mp3']
music_file_len = [4600, 103000] # approximate number of frames in each audio file

######## pin definitions ########
Lights_Pin = 26 # com to arduino
Flash_Pin = 19 # com to arduino
Speaker_Pwr_Pin = 6 # com to arduino
# _Pin = 13 # unused

####### gpio setup ###############
gpio.setmode(gpio.BCM) # sets pin reference scheme to gpio instead of pins
##gpio.setup(MPS_inp, gpio.IN) # N.O.
##gpio.setup(Start_btn, gpio.IN, gpio.PUD_UP) # N.C. with pull up resistor
gpio.setup(Lights_Pin, gpio.OUT) # N.O.
gpio.setup(Flash_Pin, gpio.OUT) # N.O.
gpio.setup(Speaker_Pwr_Pin, gpio.OUT) # N.O.

## temp counter vars
ii = 0

######## main loop ##############
try:
    while True:
        Now = datetime.datetime.now() # update time variable

        ## update alarm each day
        if Now.hour >= 11:
            try:
##                alarm = datetime.datetime(Now.year,Now.month,Now.day+1,alarm_hour,alarm_minute)
                alarm = datetime.datetime(Now.year,Now.month,Now.day,alarm_hour,alarm_minute)
            except ValueError:
                try:
                    alarm = datetime.datetime(Now.year,Now.month+1,1,alarm_hour,alarm_minute)
                except ValueError:
                    alarm = datetime.datetime(Now.year+1,1,1,alarm_hour,alarm_minute)
        else: alarm = datetime.datetime(Now.year,Now.month,Now.day,alarm_hour,alarm_minute)

        ########## alarm activation
        Alarm_Lights = False
        Play_Music = False
        Alarm_Lights_Flash = False

        alm_delt = alarm-Now # time till alarm
        if (alm_delt.days >= 0 and alm_delt.seconds/60 <= 3) or (alm_delt.days < 0 and abs(int(alm_delt.seconds/60-24*60)) < 5): # T - 3 to T + 5
            Alarm_Lights = True
            Play_Music = True
        elif alm_delt.days >= 0 and alm_delt.seconds/60 <= 15: # T - 15
            Alarm_Lights = True
        elif alm_delt.days < 0 and abs(int(alm_delt.seconds/60-24*60)) < 15: # T + 5-15
            Alarm_Lights = True
            Play_Music = True
            Alarm_Lights_Flash = True
        ############################

        ######### alarm button
        # if alarm set button high for
##        alarm_set_time = timeit.default_timer() # begin timer
##        while Alarm_Set_Btn:
##            # enter into alarm set mode
##            if timeit.default_timer()-alarm_set_time > 6000:
##                # enter setting mode
##                # every press moves the alarm time up by 0.25 hours
##                unpressed_time = timeit.default_timer()
##                while timeit.default_timer() - unpressed_time < 20000:
##                    if Alarm_Set_Btn:
##                        unpressed_time = timeit.default_timer()
##                        alarm_time_minute += 15
##                        if alarm_time_minute >= 60:
##                            alarm_time_minute = 0
##                            alarm_time_hour += 1
##                        if alarm_time_hour >= 8: # no later than 8 pm
##                            alarm_time_hour = 5
##                            alarm_time_minute = 0
        # audio button
        # if audio button latched then enable audio
        # if audio button unlatched then no audio
        # if latched and unlatched 4 times in 8 seconds then launch an audio test

        # if time t - 5 then set a digital pin high for 5 seconds to tell arduino
        # to begin lights
        # if time t = 0 then ''''
        # if t+ 5 then ''''

           ####### debug testing and overrides ###########
##	if Now.second - 30 <= 0:
##            Alarm_Lights = False
##            Play_Music = False
##            Alarm_Lights_Flash = False
##        else:
####            Alarm_Lights = True
####            Play_Music = True
##            Alarm_Lights_Flash = True
#	Play_Music = True
        ######### starting and stopping music subprocess ###################
        if Play_Music and not Music_Start_OneShot:
            Music_Start_OneShot = True
            Power_Speakers = True
            file_num = np.random.randint(0,2)
            if file_num == 0: # the intro
            	skip = '-k 1'
            else:
		if np.random.randint(0,2) == 0:
                	skip = '-k 1'
                else:
			# below means start position will be randomly between t+5 and t-5 on the music file
            		frame = int(np.random.rand(1)[0]*(music_file_len[file_num]-30000)+15000)
			skip = '-k ' + str(frame)
            print(music_files[file_num], 'start frame:',skip)
            music = subprocess.Popen(['mpg123',skip,music_files[file_num]])
        elif not Play_Music and Music_Start_OneShot:
            music.kill()
            Music_Start_OneShot = False
            Power_Speakers = False

        ######## gpio pin controls ################################
        if Alarm_Lights: gpio.output(Lights_Pin, gpio.LOW)
        else: gpio.output(Lights_Pin, gpio.HIGH)
        if Alarm_Lights_Flash: gpio.output(Flash_Pin, gpio.LOW)
        else: gpio.output(Flash_Pin, gpio.HIGH)
        if Power_Speakers: gpio.output(Speaker_Pwr_Pin, gpio.LOW)
        else: gpio.output(Speaker_Pwr_Pin, gpio.HIGH)

        ### debug printing
##        print("Lights:",Alarm_Lights,"Flash:",Alarm_Lights_Flash,"Music:",Play_Music)
##        if not Change_Flag and Alarm_Lights_Flash:
##            print(Alarm_Lights_Flash)
##            Change_Flag = True
##        elif Change_Flag and not Alarm_Lights_Flash:
##            print(Alarm_Lights_Flash)
##            Change_Flag = False

########## shutdown, interrupt, cleanup ###############
except KeyboardInterrupt: pass
gpio.cleanup()


import json
import joblib
import numpy as np
from azureml.core.model import Model
import pytorch_lightning as pl

# Called when the service is loaded
def init():
    global model
    # Get the path to the deployed model file and load it
    model_path = Model.get_model_path('../models/model.py')
    model = joblib.load(model_path)

# Called when a request is received
def run(raw_data):
    # Get the input data as a numpy array
    data = np.array(json.loads(raw_data)['data'])

    # Create a trainer
    trainer = pl.Trainer(max_epochs=3)

    # Get predictions from the trainer, TODO: data needs to be a dataloader
    predictions = trainer.predict(model, data) # returns a list of dictionaries, one for each provided dataloader containing their respective predictions
    predicted_classes = []
    classnames = ['yes', 'no']

    # TODO: needs reviewing
    # for prediction,  in predictions:
    #     predicted_classes.append(classnames[prediction])
    # Return the predictions as JSON
    return json.dumps(predicted_classes)


import re
import math
import torch
import torchaudio
import numpy as np
import argparse
import librosa
import warnings
import itertools as it
from jiwer import wer
from tqdm import tqdm
import pandas as pd
import os

try:
    from flashlight.lib.sequence.criterion import CpuViterbiPath, get_data_ptr_as_bytes
    from flashlight.lib.text.dictionary import create_word_dict, load_words
    from flashlight.lib.text.decoder import (
        CriterionType,
        LexiconDecoderOptions,
        KenLM,
        LM,
        LMState,
        SmearingMode,
        Trie,
        LexiconDecoder,
    )
    from flashlight.lib.text.decoder import LexiconFreeDecoder, LexiconFreeDecoderOptions

except:
    warnings.warn(
        "flashlight python bindings are required to use this KenLM. Please install from https://github.com/facebookresearch/flashlight/tree/master/bindings/python"
    )
    LM = object
    LMState = object

from datasets import load_dataset, load_metric, concatenate_datasets

from transformers import Wav2Vec2CTCTokenizer
from transformers import Wav2Vec2FeatureExtractor
from transformers import Wav2Vec2Processor
from transformers import Wav2Vec2ForCTC

from utils.generic_utils import load_config, load_vocab, calculate_wer

from utils.dataset_preprocessed import remove_extra_columns, parse_dataset_dict, vocab_to_string, DataColletor
from torch.utils.data import DataLoader

from transformers import GPT2LMHeadModel, GPT2TokenizerFast
# BERT
from transformers import BertTokenizer, BertForMaskedLM

def score_fun_linear(s1, s2, w1=1, w2=1):
  return s1*w1 + s2*w2

class Scorer:
    def __init__(self, model_name, kenLM_weight=1, externalLM_weight=1, score_fn=score_fun_linear):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.tokenizer = GPT2TokenizerFast.from_pretrained(model_name)
        self.model = GPT2LMHeadModel.from_pretrained(model_name).to(self.device)
        self.model.eval()
        self.kenLM_weight = kenLM_weight
        self.externalLM_weight = externalLM_weight

        self.score_fn = score_fn

        print('---->>> Testing Model.')
        self.test_model(['a capital da frança é paris', 'a capital da franca é paris', 'a capital da frança é parir'])
        print('---->>> Done testing model')

    def calculate_perplexity(self, text):
        encodings = self.tokenizer(text, return_tensors="pt")
        max_length = self.model.config.n_positions
        # stride = 1
        stride = 512

        lls = []
        for i in (range(0, encodings.input_ids.size(1), stride)):
            begin_loc = max(i + stride - max_length, 0)
            end_loc = min(i + stride, encodings.input_ids.size(1))
            trg_len = end_loc - i
            input_ids = encodings.input_ids[:,
                                            begin_loc:end_loc].to(self.device)
            target_ids = input_ids.clone()
            target_ids[:, :-trg_len] = -100

            with torch.no_grad():
                outputs = self.model(input_ids, labels=target_ids)
                log_likelihood = outputs[0] * trg_len
            if not math.isnan(log_likelihood):
                lls.append(log_likelihood)
        ppl = torch.exp(torch.stack(lls).sum() / end_loc).item()
        return ppl

    def test_model(self, candidates):
        for candidate in candidates:
            ppl = self.calculate_perplexity(candidate)
            print("{0} ---- {1}".format(candidate, ppl))

    def chose_best_candidate(self, candidates, candidate_scores):
        best_candidate = None
        best_candidate_id = None
        best_score = float('inf')

        for i in range(len(candidates)):
            # *-1 because in kenLM high is better
            kenlm_score = candidate_scores[i] * -1
            candidate = candidates[i]
            external_lm_score = self.calculate_perplexity(candidate)
            new_score = self.score_fn(kenlm_score, external_lm_score, self.kenLM_weight, self.externalLM_weight)
            # print(candidate, "-->", new_score)
            if new_score < best_score:
                best_candidate = candidate
                best_candidate_id = i
                best_score = new_score

        return (best_candidate_id, best_candidate, best_score)

def remove_invalid_characters(batch):
    text = batch[text_column].lower()
    text = re.sub("[^{}]".format(vocab_string), " ", text)
    text = re.sub("[ ]+", " ", text)
    batch[text_column] = text + " "
    return batch

def load_audio(batch):
    if dataset_base_path:
        batch[audio_path_column] = os.path.join(dataset_base_path, batch[audio_path_column])
    speech_array, sampling_rate = torchaudio.load(batch[audio_path_column])
    batch["speech"] = speech_array.squeeze().numpy()
    batch["sampling_rate"] = sampling_rate
    if text_column in batch:
        batch["target_text"] = batch[text_column]
    return batch

def resample_audio(batch):
    if batch["sampling_rate"] != config['sampling_rate']:
        batch["speech"] = librosa.resample(np.asarray(batch["speech"]),  batch["sampling_rate"], config['sampling_rate'])
        batch["sampling_rate"] = config['sampling_rate']
    return batch

def prepare_dataset(batch):
    batch['audio_path'] = batch[audio_path_column]
    batch["input_values"] = processor(batch["speech"], sampling_rate=config['sampling_rate']).input_values

    if "target_text" in batch:
        with processor.as_target_processor():
            batch["labels"] = processor(batch["target_text"]).input_ids
    return batch

class KenLMDecoder(object):
    def __init__(self, kenlm_args, vocab_dict, rescore_args=None, blank="<pad>", silence="|", unk="<unk>"):

        self.vocab_size = len(vocab_dict)
        self.blank_token = (vocab_dict[blank])
        self.silence_token = vocab_dict[silence]
        self.unk_token = vocab_dict[unk]

        self.nbest = kenlm_args['nbest']

        if kenlm_args['lexicon_path']:
            vocab_keys = vocab_dict.keys()
            self.lexicon = load_words(kenlm_args['lexicon_path'])
            self.word_dict = create_word_dict(self.lexicon)
            self.unk_word = self.word_dict.get_index(unk)

            self.lm = KenLM(kenlm_args['kenlm_model_path'], self.word_dict)
            self.trie = Trie(self.vocab_size, self.silence_token)

            start_state = self.lm.start(False)
            for i, (word, spellings) in enumerate(self.lexicon.items()):
                word_idx = self.word_dict.get_index(word)
                _, score = self.lm.score(start_state, word_idx)

                for spelling in spellings:
                    spelling_idxs = []
                    for token in spelling:
                        if token.upper() in vocab_keys:
                            spelling_idxs.append(vocab_dict[token.upper()])
                        elif token.lower() in vocab_keys:
                            spelling_idxs.append(vocab_dict[token.lower()])
                        else:
                            print("WARNING: The token", token, "not exist in your vocabulary, using <unk> token instead")
                            spelling_idxs.append(self.unk_token)

                    self.trie.insert(spelling_idxs, word_idx, score)
            self.trie.smear(SmearingMode.MAX)

            self.decoder_opts = LexiconDecoderOptions(
                beam_size=kenlm_args['beam'],
                beam_size_token=kenlm_args['beam_size_token'] if "beam_size_token" in kenlm_args else len(vocab_dict),
                beam_threshold=kenlm_args['beam_threshold'],
                lm_weight=kenlm_args['lm_weight'],
                word_score=kenlm_args['word_score'],
                unk_score=-math.inf,
                sil_score=kenlm_args['sil_weight'],
                log_add=False,
                criterion_type=CriterionType.CTC,
            )

            self.decoder = LexiconDecoder(
                self.decoder_opts,
                self.trie,
                self.lm,
                self.silence_token,
                self.blank_token,
                self.unk_word,
                [],
                False,
            )
        else:
            d = {w: [[w]] for w in vocab_dict.keys()}
            self.word_dict = create_word_dict(d)
            self.lm = KenLM(kenlm_args['kenlm_model_path'], self.word_dict)
            self.decoder_opts = LexiconFreeDecoderOptions(
                beam_size=kenlm_args['beam'],
                beam_size_token=kenlm_args['beam_size_token'] if "beam_size_token" in kenlm_args else len(vocab_dict),
                beam_threshold=kenlm_args['beam_threshold'],
                lm_weight=kenlm_args['lm_weight'],
                sil_score=kenlm_args['sil_weight'],
                log_add=False,
                criterion_type=CriterionType.CTC,
            )
            self.decoder = LexiconFreeDecoder(
                self.decoder_opts, self.lm, self.silence_token, self.blank_token, []
            )

    def get_tokens(self, idxs):
        """Normalize tokens by handling CTC blank"""
        idxs = (g[0] for g in it.groupby(idxs))
        idxs = filter(lambda x: x != self.blank_token, idxs)
        return torch.LongTensor(list(idxs))

    def decode(self, emissions):
        B, T, N = emissions.size()
        # print(emissions.shape)
        tokens = []
        scores = []
        for b in range(B):
            emissions_ptr = emissions.data_ptr() + 4 * b * emissions.stride(0)
            results = self.decoder.decode(emissions_ptr, T, N)
            nbest_results = results[: self.nbest]
            tokens_nbest = []
            scores_nbest = []
            for result in nbest_results:
                tokens_nbest.append(result.tokens)
                scores_nbest.append(result.score)
            tokens.append(tokens_nbest)
            scores.append(scores_nbest)

        token_array = np.array(tokens, dtype=object)
        scores_arrray = np.array(scores, dtype=object)
        return token_array, scores_arrray

def test(model, test_dataset, processor, kenlm, calcule_wer=True, return_predictions=False):
    model.eval()
    predictions = []
    tot_samples = 0
    tot_wer = 0
    tot_cer = 0
    with torch.no_grad():
        for batch in tqdm(test_dataset):
            input_values, attention_mask = batch['input_values'], batch['attention_mask']
            if calcule_wer:
                labels = batch['labels']

            if USE_CUDA:
                input_values = input_values.cuda(non_blocking=True)
                attention_mask = attention_mask.cuda(non_blocking=True)
                if calcule_wer:
                    labels = labels.cuda(non_blocking=True)

            logits = model(input_values, attention_mask=attention_mask).logits

            if kenlm:
                logits = torch.nn.functional.log_softmax(logits.float(), dim=-1)
                # get all candidates
                lm_tokens, lm_scores = kenlm.decode(logits.cpu().detach())
                # choise the best candidate

                if rescore_lm:
                    pred_ids = []
                    for b in range(logits.size(0)):
                        candidates_ids = []
                        scores = []
                        for c in range(len(lm_tokens[b])):
                            candidate = lm_tokens[b][c]
                            score = lm_scores[b][c]
                            candidates_ids.append(candidate)
                            scores.append(score)

                        candidates_text = processor.batch_decode(candidates_ids)
                        # if less than 3 tokens, ignore rescore
                        if len(candidates_text[0].split(' ')) < 3:
                            # use the best kenLM candidate
                            pred_id = candidates_ids[0]
                        else:
                            best_candidate_id, _, _ = rescore_lm.chose_best_candidate(candidates_text, scores)
                            pred_id = candidates_ids[best_candidate_id]

                        pred_ids.append(pred_id)
                else:

                    pred_ids = []
                    for b in range(logits.size(0)):
                        pred_ids.append(lm_tokens[b][0])

                pred_ids = np.array(pred_ids)
            else:
                pred_ids = np.argmax(logits.cpu().detach().numpy(), axis=-1)

            if calcule_wer:
                # compute metrics
                wer, cer = calculate_wer(pred_ids, labels.cpu().detach().numpy(), processor)
                tot_wer += wer
                tot_cer += cer

            if return_predictions:
                audios_path = batch['audio_path']
                # get text
                pred_string = processor.batch_decode(pred_ids)

                for i in range(len(audios_path)):
                    output_wav_path = audios_path[i]
                    if dataset_base_path:
                        output_wav_path = output_wav_path.replace(dataset_base_path, '').replace(dataset_base_path+'/', '')

                    predictions.append([output_wav_path, pred_string[i].lower()])

            tot_samples += input_values.size(0)
    if calcule_wer:
        # calculate avg of metrics
        avg_wer = tot_wer/tot_samples
        avg_cer = tot_cer/tot_samples
        print("\n\n --> TEST PERFORMANCE\n")
        print("     | > :   WER    ({:.5f})\n".format(avg_wer))
        print("     | > :   CER    ({:.5f})\n".format(avg_cer))

    return predictions

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config_path', type=str, required=True,
                        help="json file with configurations")
    parser.add_argument('--checkpoint_path_or_name', type=str, required=True,
                        help="path or name of checkpoints")
    parser.add_argument('--no_use_kenlm', default=False, action='store_true',
                        help="Not use KenLm during inference ?")
    parser.add_argument('--rescore', default=False, action='store_true',
                        help="Use a external LM to rescore?")
    parser.add_argument('--audio_path', type=str, default=None,
                        help="If it's passed the inference will be done in all audio files in this path and the dataset present in the config json will be ignored")
    parser.add_argument('--output_csv', type=str, default=None,
                        help="CSV for save all predictions")

    args = parser.parse_args()

    config = load_config(args.config_path)

    # Use CUDA
    USE_CUDA = torch.cuda.is_available()

    model = Wav2Vec2ForCTC.from_pretrained(args.checkpoint_path_or_name)

    feature_extractor = Wav2Vec2FeatureExtractor(feature_size=1, sampling_rate=config['sampling_rate'], padding_value=0.0, do_normalize=True, return_attention_mask=True)
    processor = Wav2Vec2Processor.from_pretrained(args.checkpoint_path_or_name)
    vocab_dict = processor.tokenizer.get_vocab()
    pad_token = processor.tokenizer.pad_token
    silence_token = processor.tokenizer.word_delimiter_token
    unk_token = processor.tokenizer.unk_token

    # if the model uses upper words in vocab force tokenizer lower case for compatibility with our data loader
    if list(vocab_dict.keys())[-1].isupper():
        processor.tokenizer.do_lower_case = True

    data_collator = DataColletor(processor=processor, padding=True, test=True)

    if USE_CUDA:
        model = model.cuda()

    if not args.no_use_kenlm:
        print("> Inference using KenLM")
        kenlm = KenLMDecoder(config.KenLM, vocab_dict, blank=pad_token, silence=silence_token, unk=unk_token)
    else:
        print("> Inference without KenLM")
        kenlm = None

    if args.rescore:
        rescore_args = config.rescore if "rescore" in config.keys() else None
        if rescore_args:
            print("> Inference with External LM rescoring")
            rescore_lm = Scorer(rescore_args['lm_path_or_name'], kenLM_weight=rescore_args['KenLM_weight'], externalLM_weight=rescore_args['ExternalLM_weight'])
        else:
            print("> Inference without External LM rescoring")
            rescore_lm = None
    else:
        print("> Inference without External LM rescoring")
        rescore_lm = None

    if not args.audio_path:
        # load dataset
        test_dataset_config = config.datasets['test']
        text_column, audio_path_column = parse_dataset_dict(test_dataset_config)

        dataset = load_dataset(**test_dataset_config)
        # made compatibility with csv load
        if isinstance(dataset, dict) and 'train' in dataset.keys():
            concat_list = []
            for k in dataset.keys():
                concat_list.append(dataset[k])
            dataset = concatenate_datasets(concat_list)

        if 'files_path' in config['datasets'].keys() and config.datasets['files_path']:
            if test_dataset_config['name'].lower() == 'csv':
                dataset_base_path = config.datasets['files_path']
            else:
                print("> Warning: datasets['files_path'] igonored because dataset is not CSV !")
                dataset_base_path = None
        else:
            dataset_base_path = None

        # preprocess dataset
        dataset = remove_extra_columns(dataset, text_column, audio_path_column)

        vocab_string = vocab_to_string(vocab_dict, pad_token, silence_token, unk_token).lower()

        print("\n\n> Remove invalid chars \n\n")
        # remove invalid chars
        dataset = dataset.map(remove_invalid_characters, num_proc=config['num_loader_workers'])

        # Load audio files
        dataset = dataset.map(load_audio)
        print("\n\n> Resample Audio Files \n\n")
        # resample audio files if necessary
        dataset = dataset.map(resample_audio, num_proc=config['num_loader_workers'])

        print("\n\n> Prepare dataset \n\n")
        # batched dataset
        dataset = dataset.map(prepare_dataset, remove_columns=dataset.column_names, batch_size=config['batch_size'], num_proc=1, batched=True)

        test_dataset = DataLoader(dataset=dataset,
                    batch_size=config['batch_size'],
                    collate_fn=data_collator,
                    shuffle=True,
                    num_workers=config['num_loader_workers'])

        print("\n\n> Starting Evaluation \n\n")
        preds = test(model, test_dataset, processor, kenlm,  calcule_wer=True, return_predictions=True)

    if args.output_csv:
        root_path = os.path.dirname(args.output_csv)
        os.makedirs(root_path, exist_ok=True)

        df = pd.DataFrame(preds, columns=["file_path", "transcription"])
        df.sort_values(by=['file_path'], inplace=True)
        df.to_csv(args.output_csv, index=False)
        print("\n\n> Evaluation outputs saved in: ", args.output_csv)


from json import loads
from pathlib import Path
import requests
import numpy as np
from imageio import imread
import param
import holoviews as hv
from holoviews.operation.datashader import rasterize
import torch
from torchvision import transforms
import matplotlib.pyplot as plt
import seaborn as sn
from neural_rock.data_models import ImageDataset, ModelZoo, LabelSets
import bokeh
from neural_rock.app.plot import create_holoviews_thinsection, create_holoviews_cam
hv.extension('bokeh')

import panel as pn
import logging
import pandas as pd
LOG_FORMAT = "%(levelname)s %(asctime)s - %(message)s"
file_handler = logging.FileHandler(filename='test.log', mode='w')
file_handler.setFormatter(logging.Formatter(LOG_FORMAT))
logger = logging.getLogger(__name__)
logger.addHandler(file_handler)
logger.setLevel(logging.INFO)

class ThinSectionViewer(param.Parameterized):
    """
    Base class for the thin-section viewer and the selection widgets.
    Builds on Holoviews Panel and Holoviews Parameters to create parameters that are then turned into Widgets
    throught the front end.
    """
    Labelset_Name = param.ObjectSelector(default="Lucia", objects=["Dunham", 'Lucia', 'DominantPore'], check_on_set=True)

    Model_Selector = param.ObjectSelector(default="resnet", objects=['vgg', 'resnet'], check_on_set=True)
    Frozen_Selector = param.ObjectSelector(default="True", objects=['True', 'False'], check_on_set=True)

    Network_Layer_Number = param.Integer(default=0, bounds=(0, 1))

    Class_Name = param.ObjectSelector(default="default", objects=["default"], check_on_set=True)

    Image_Name = param.ObjectSelector(default="default", objects=["default"], check_on_set=True)
    Show_CAM = param.Boolean(True, doc="Show CAM Map")
    Alpha = param.Number(0.3, bounds=(0.0, 1.0), doc="CAM Alpha")

    def __init__(self, server_address: Path,
                 image_dataset: ImageDataset,
                 **params):
        super(ThinSectionViewer, self).__init__(**params)
        self.cache = []
        self.server_address = server_address
        self.image_dataset = image_dataset
        self.model_zoo = self._get_model_zoo()
        self.label_sets = self._get_labelsets()
        self.current_sample_id = None
        self._get_layer_ranges()
        self._update_class_names()
        self._get_available_image_ids()

        self.map = self.load_symbol
        self.probs = self.make_map()[1]
        self.confusion_matrices = self._get_confusion_matrices(self.Labelset_Name, self.Model_Selector, self.Frozen_Selector)

    def _get_model_zoo(self) -> ModelZoo:
        logger.info("Get Model Zoo")
        """
        Calls API to retrieve available Models in Model Zoo
        """
        with requests.Session() as s:
            result = s.get(self.server_address + 'models')
            r = loads(result.text)
            model_zoo = ModelZoo(**r)
        return model_zoo

    def _get_labelsets(self) -> LabelSets:
        logger.info("Get Labelsets")
        """
        Calls API to retrieve available Labelsets
        """
        with requests.Session() as s:
            result = s.get(self.server_address + 'labelsets')
            r = loads(result.text)
            labelsets = LabelSets(**r)

        self.param['Labelset_Name'].default = list(labelsets.sets.keys())[0]
        self.param['Labelset_Name'].objects = list(labelsets.sets.keys())
        self.Labelset_Name = list(labelsets.sets.keys())[0]

        return labelsets

    @param.depends('Labelset_Name', watch=True)
    def _update_class_names(self):
        logger.info("Update Class Names")
        class_names = self.label_sets.sets[self.Labelset_Name].class_names
        self.param['Class_Name'].default = class_names[0]
        self.param['Class_Name'].objects = class_names
        self.Class_Name = class_names[0]

    @param.depends('Labelset_Name', 'Model_Selector', 'Frozen_Selector', watch=True)
    def _get_available_image_ids(self):
        logger.info("Get available image ids")
        """
        Calls API to retrieve available Images for a given Labelset and Model
        """
        with requests.Session() as s:
            result = s.get(self.server_address + 'dataset/sample_ids')
            r = loads(result.text)

        train_test_split = self._get_train_test_config()

        labels = self.label_sets.sets[self.Labelset_Name].sample_labels
        samples_text_map = {}
        for sample_id in r:
            if sample_id in train_test_split['train']:
                samples_text_map["{0:}-Train-True Class-{1:}".format(sample_id, labels[sample_id])] = sample_id
            elif sample_id in train_test_split['test']:
                samples_text_map["{0:}-Test-True Class-{1:}".format(sample_id, labels[sample_id])] = sample_id

        self.samples_text_map = samples_text_map
        if self.current_sample_id is None:
            logger.info("Setting default image")
            self.param['Image_Name'].default = list(self.samples_text_map.keys())[0]
            self.current_sample_id = int(self.param['Image_Name'].default.split("-")[0])
        else:
            logger.info("Setting image based on previous selection")
            for key in self.samples_text_map.keys():
                if int(key.split("-")[0]) == self.current_sample_id:
                    logger.info("Found image based on previous selection")
                    self.param['Image_Name'].default = key
                    self.current_sample_id = int(self.param['Image_Name'].default.split("-")[0])
        self.param['Image_Name'].objects = list(self.samples_text_map.keys())
        self.Image_Name = self.param['Image_Name'].objects[0]

    @param.depends('Labelset_Name', 'Model_Selector', 'Frozen_Selector')
    def _get_train_test_config(self):
        logger.info("Get _get_train_test_config")
        """
        Calls API to retrieve the train test split for a specific trained model.
        """
        with requests.Session() as s:
            result = s.get(self.server_address + 'get_train_test_images/{0:}/{1:}/{2:}'.format(self.Labelset_Name,
                                                                                self.Model_Selector,
                                                                                self.Frozen_Selector))
            r = loads(result.text)

        return r

    @param.depends('Model_Selector', watch=True)
    def _get_layer_ranges(self):
        logger.info("Get _get_layer_ranges")
        """
        Calls API to retrieve available layer numbers for the CAM computation
        """
        with requests.Session() as s:
            result = s.get(self.server_address + 'models/{0:}/valid_layer_range'.format(self.Model_Selector))
            r = loads(result.text)

        min_layer = r[0]
        max_layer = r[-1]
        self.param['Network_Layer_Number'].default = max_layer
        self.param['Network_Layer_Number'].bounds = (min_layer, max_layer)
        self.Network_Layer_Number = self.param['Network_Layer_Number'].default

    @param.depends('Image_Name')
    def load_image(self):
        logger.info("Get load_image")
        """
        Loads an image. Shoudld replace with other loading function in future.
        """
        image_id = self.samples_text_map[self.Image_Name]
        try:
            X_np = imread(self.image_dataset.image_paths[image_id])
            mask = imread(self.image_dataset.roi_paths[image_id])
            mask = (1 - (mask > 0).astype(np.uint8))
            X_np = X_np * np.repeat(np.expand_dims(mask, axis=-1), 3, axis=-1)
            X_np = np.transpose(X_np, (1, 0, 2))
        except KeyError:
            X_np = np.random.uniform(0, 1, size=(1000, 1000, 3))
        return X_np

    def check_cache(self, labelset_Name, Model_Selector, Network_Layer_Number, Frozen_Selector, sample_id, Class_Name):
        for row in self.cache:
            if row['labelset_Name'] == labelset_Name and row['Model_Selector'] == Model_Selector and row['Network_Layer_Number'] == Network_Layer_Number and row['Frozen_Selector'] == Frozen_Selector and row['sample_id'] == sample_id and row['Class_Name'] == Class_Name:
                return row
        return None

    def _get_confusion_matrices(self, labelset_Name, Model_Selector, Frozen_Selector):
        with requests.Session() as s:
            result = s.get(self.server_address + 'confusion_matrices/{0:}/{1:}/{2:}'.format(labelset_Name,
                                                                                Model_Selector,
                                                                                Frozen_Selector))
            r = loads(result.text)
            return r['confusion_matrices']

    @param.depends('Image_Name', 'Model_Selector', 'Labelset_Name', 'Frozen_Selector', 'Class_Name')
    def make_map(self):
        logger.info("Get make_map")
        """
        Calls API to generate a cam map for a given image, model, layer number and target class name.
        """
        sample_id = self.samples_text_map[self.Image_Name]

        cache_result = self.check_cache(self.Labelset_Name, self.Model_Selector, self.Network_Layer_Number, self.Frozen_Selector, sample_id, self.Class_Name)

        self.confusion_matrices = self._get_confusion_matrices(self.Labelset_Name,
                                                               self.Model_Selector,
                                                               self.Frozen_Selector)

        if cache_result is None:
            with requests.Session() as s:
                result = s.get(self.server_address + 'cam/{0:}/{1:}/{2:}/{3:}/{4:}/{5:}'.format(self.Labelset_Name,
                                                                                self.Model_Selector,
                                                                                self.Network_Layer_Number,
                                                                                self.Frozen_Selector,
                                                                                sample_id,
                                                                                self.Class_Name))
                r = loads(result.text)
                map = np.array(r['map'])
                probs = r['y_prob']

                self.cache.append({'labelset_Name': self.Labelset_Name,
                                   'Model_Selector': self.Model_Selector,
                                   'Network_Layer_Number': self.Network_Layer_Number,
                                   'Frozen_Selector': self.Frozen_Selector,
                                   'sample_id': sample_id,
                                   'Class_Name': self.Class_Name,
                                   'probs': probs,
                                   'map': map})
        else:
            probs = cache_result['probs']
            map = cache_result['map']

        self.probs = probs
        return map, probs

    @param.depends('Image_Name', 'Show_CAM', 'Alpha')
    def load_symbol(self):
        logger.info("Get load_symbol")
        """
        Creates the CAM and Thin section holoviews objects to plot in the viewer.
        """
        X_np = self.load_image()
        resize = transforms.Resize((X_np.shape[0], X_np.shape[1]))
        map, _ = self.make_map()
        map = resize(torch.from_numpy(map).unsqueeze(0)).numpy()[0]
        hv_thinsection = create_holoviews_thinsection(X_np[::-1])

        hv_cam = create_holoviews_cam(map.T).opts(alpha=self.Alpha, cmap='inferno')

        image = hv_thinsection

        if self.Show_CAM:
            image *= hv_cam

        return image

    def bar_plot(self):
        logger.info("Get bar_plot")
        """
        Creates the bar plot for image probabilities
        """
        data = [(clas, prob) for clas, prob in zip(self.param['Class_Name'].objects, self.probs)]
        bars = hv.Bars(data, hv.Dimension('Class Name'), 'Probability')

        return bars

    #@pn.depends('Class_Name')
    def plot_confusion_matrix_train(self):
        class_names = self.param['Class_Name'].objects

        fig, axarr = plt.subplots(2, 1, figsize=(6, 14))
        for ax, phase in zip(axarr, ["train", "test"]):
            df_cm = pd.DataFrame(self.confusion_matrices[phase], class_names, class_names)
            sn.set(font_scale=1.4)  # for label size
            sn.heatmap(df_cm, annot=True, annot_kws={"size": 16}, ax=ax)  # font size
            ax.set_ylabel("Ground Truth Label", fontsize=16)
            ax.set_xlabel("Predicted Label", fontsize=16)
            ax.set_title(phase, fontsize=18)
        fig.savefig("tmp.png", bbox_inches="tight")
        return pn.panel('tmp.png', width=300)

    def view(self):
        logger.info("Get view")
        """
        Initializes the viewer and provides the dynamically updating holoviews image viewer.
        """
        self.make_map()
        thin_section = hv.DynamicMap(self.map)

        thin_section = rasterize(thin_section).opts(data_aspect=1.0, frame_height=600, frame_width=1000,
                                                    active_tools=['xwheel_zoom', 'pan'])

        col = pn.Row(thin_section, pn.Column(self.bar_plot(), self.plot_confusion_matrix_train()))
        return col