from __future__ import print_function
from __future__ import division

import argparse
import random
import torch
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
from torch.autograd import Variable
import numpy as np
from warpctc_pytorch import CTCLoss
import os
import models.utils as utils
from loader import DatasetLoader
from multiprocessing import cpu_count
from tqdm import tqdm
from torchsummary import summary

import models.crnn as crnn

parser = argparse.ArgumentParser()
parser.add_argument('--root', required=True, help='path to root folder')
parser.add_argument('--train', required=True, help='path to dataset')
parser.add_argument('--val', required=True, help='path to dataset')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=4)
parser.add_argument('--batch_size', type=int, default=64, help='input batch size')
parser.add_argument('--imgH', type=int, default=32, help='the height of the input image to network')
parser.add_argument('--imgW', type=int, default=512, help='the width of the input image to network')
parser.add_argument('--nh', type=int, default=256, help='size of the lstm hidden state')
parser.add_argument('--nepoch', type=int, default=100, help='number of epochs to train for')
parser.add_argument('--cuda', action='store_true', help='enables cuda')
parser.add_argument('--gpu', type=int, default=1, help='number of GPUs to use')
parser.add_argument('--pretrained', default='', help="path to pretrained model (to continue training)")
parser.add_argument('--alphabet', type=str, required=True, help='path to char in labels')
parser.add_argument('--expr_dir', required=True, type=str, help='Where to store samples and models')
parser.add_argument('--displayInterval', type=int, default=1, help='Interval to be displayed')
parser.add_argument('--n_test_disp', type=int, default=10, help='Number of samples to display when test')
parser.add_argument('--valInterval', type=int, default=1, help='Interval to be displayed')
parser.add_argument('--saveInterval', type=int, default=1, help='Interval to be displayed')
parser.add_argument('--lr', type=float, default=0.001, help='learning rate for Critic, not used by adadealta')
parser.add_argument('--manualSeed', type=int, default=1234, help='reproduce experiemnt')
opt = parser.parse_args()
print(opt)
os.environ['CUDA_VISIBLE_DEVICES'] = str(opt.gpu)

if not os.path.exists(opt.expr_dir):
    os.makedirs(opt.expr_dir)

random.seed(opt.manualSeed)
np.random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)

cudnn.benchmark = True

if torch.cuda.is_available() and not opt.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")

loader = DatasetLoader(opt.root, opt.train, opt.val, opt.imgW, opt.imgH)
train_loader = loader.train_loader(opt.batch_size, num_workers=opt.workers)
test_loader = loader.test_loader(opt.batch_size, num_workers=opt.workers)

alphabet = open(os.path.join(opt.root, opt.alphabet)).read().rstrip()
nclass = len(alphabet) + 1
nc = 3

print(len(alphabet), alphabet)
converter = utils.strLabelConverter(alphabet, ignore_case=False)
criterion = CTCLoss()

crnn = crnn.CRNN(opt.imgH, nc, nclass, opt.nh)
if opt.pretrained != '':
    print('loading pretrained model from %s' % opt.pretrained)
    pretrain = torch.load(opt.pretrained)
    crnn.load_state_dict(pretrain, strict=False)

image = torch.FloatTensor(opt.batch_size, 3, opt.imgH, opt.imgH)
text = torch.IntTensor(opt.batch_size * 5)
length = torch.IntTensor(opt.batch_size)

if opt.cuda:
    crnn.cuda()
    image = image.cuda()
    criterion = criterion.cuda()

summary(crnn.cnn, (3, opt.imgH, opt.imgW))

image = Variable(image)
text = Variable(text)
length = Variable(length)
train_loss_avg = utils.averager()
train_cer_avg = utils.averager()

# setup optimizer
optimizer = optim.Adam(crnn.parameters(), lr=opt.lr)

def val(net, data_loader, criterion, max_iter=1000):
    print('Start val')

    for p in crnn.parameters():
        p.requires_grad = False

    net.eval()

    val_iter = iter(data_loader)

    val_loss_avg = utils.averager()
    val_cer_avg = utils.averager()
    max_iter = min(max_iter, len(data_loader))
    with torch.no_grad():
        for i in range(max_iter):
            data = val_iter.next()
            cpu_images, cpu_texts = data
            batch_size = cpu_images.size(0)
            utils.loadData(image, cpu_images)
            t, l = converter.encode(cpu_texts)
            utils.loadData(text, t)
            utils.loadData(length, l)

            preds = crnn(image)
            preds_size = Variable(torch.IntTensor([preds.size(0)] * batch_size))
            cost = criterion(preds, text, preds_size, length)/batch_size
            cost = cost.detach().item()
            val_loss_avg.add(cost)

            _, preds = preds.max(2)
            preds = preds.transpose(1, 0).contiguous().view(-1)
            sim_preds = converter.decode(preds.data, preds_size.data, raw=False)
            cer_loss = utils.cer_loss(sim_preds, cpu_texts)
            val_cer_avg.add(cer_loss)

    raw_preds = converter.decode(preds.data, preds_size.data, raw=True)[:opt.n_test_disp]
    for raw_pred, pred, gt in zip(raw_preds, sim_preds, cpu_texts):
        print('%-30s => %-30s, gt: %-30s' % (raw_pred, pred, gt))

    print('Test loss: %f - cer loss %f' % (val_loss_avg.val(), val_cer_avg.val()))

def trainBatch(net, data, criterion, optimizer):
    cpu_images, cpu_texts = data
    batch_size = cpu_images.size(0)
    utils.loadData(image, cpu_images)
    t, l = converter.encode(cpu_texts)
    utils.loadData(text, t)
    utils.loadData(length, l)

    preds = crnn(image)
    preds_size = Variable(torch.IntTensor([preds.size(0)] * batch_size))
    cost = criterion(preds, text, preds_size, length)/batch_size
    crnn.zero_grad()
    cost.backward()
    optimizer.step()
    cost = cost.detach().item()

    _, preds = preds.max(2)
    preds = preds.transpose(1, 0).contiguous().view(-1)
    sim_preds = converter.decode(preds.data, preds_size.data, raw=False)
    cer_loss = utils.cer_loss(sim_preds, cpu_texts)
    return cost, cer_loss, len(cpu_images)

for epoch in range(1, opt.nepoch+1):
    t = tqdm(iter(train_loader), total=len(train_loader), desc='Epoch {}'.format(epoch))
    for i, data in enumerate(t):
        for p in crnn.parameters():
            p.requires_grad = True
        crnn.train()

        cost, cer_loss, n = trainBatch(crnn, data, criterion, optimizer)

        train_loss_avg.add(cost)
        train_cer_avg.add(cer_loss)

    print('[%d/%d] Loss: %f - cer loss: %f' %
            (epoch, opt.nepoch, train_loss_avg.val(), train_cer_avg.val()))
    train_loss_avg.reset()
    train_cer_avg.reset()

    if epoch % opt.valInterval == 0:
        val(crnn, test_loader, criterion)

    # do checkpointing
    if epoch % opt.saveInterval == 0:
        torch.save(
            crnn.state_dict(), '{}/netCRNN_{}.pth'.format(opt.expr_dir, epoch))


import pandas as pd
import hashlib
import statsmodels.stats.multitest as multitest
import scipy as sp
import scipy.stats
import numpy as np
from collections import namedtuple
import warnings

from tqdm import tqdm

#  from .annotations import ProteinId, PeptideId, PeptideVariantId
#  from tqdm import tqdm

#  from modlamp.sequences import Helices
#  from modlamp.descriptors import PeptideDescriptor, GlobalDescriptor

################################################################################
# namedtuples
################################################################################
TTestResult = namedtuple("TTestResult", ("pvalue", "statistic", "ratio"))
PeptideVariantId = namedtuple('PeptideVariantId',
                              ('protein_id', 'start', 'stop', 'mod_seq', 'origin'))
PeptideId = namedtuple('PeptideId', ('protein_id', 'start', 'stop'))
ProteinId = namedtuple('ProteinId', ('protein_id',))

################################################################################
# Pandas peptidomics Excensions
################################################################################
@pd.api.extensions.register_index_accessor("peptidomics")
class PandasPeptidomicsIndex:
    def __init__(self, index):
        self.index = index

        id_classes = (PeptideVariantId, PeptideId, ProteinId)
        id_levels = ("peptide_variants", "peptides", "proteins")
        for id_class, level in zip(id_classes, id_levels):
            if id_class._fields == self.index.names[1:]:
                self.id_class = id_class
                self.level = level
                break
        else:  # no break
            raise ValueError("index does not seem to be 'proteins', 'peptides' or "
                             "peptide_variants")

    def to_proteins(self):
        if self.level == 'peptides':
            drop = ("start", "stop")
        elif self.level == "peptide_variants":
            drop = ("start", "stop", "mod_seq", "origin")
        else:
            raise ValueError("index does not seem to be 'peptides' or 'peptide_variants'")
        return self._drop_index(drop)

    def to_peptides(self):
        if self.level != 'peptide_variants':
            raise ValueError("can only craete peptide index from a peptide_variant index")
        return self._drop_index(("mod_seq", "origin"))

    def _drop_index(self, drop):
        index = self.index.copy()
        for level in drop:
            index = index.droplevel(level)
        return index.unique()

    def iter_index(self):
        for row in self.index:
            yield row[0], self.id_class(*row[1:])

    def create_annotation_df(self):
        annotation_columns = pd.MultiIndex.from_tuples((), names=['annotation', 'group'])
        return pd.DataFrame(columns=annotation_columns, index=self.index)

@pd.api.extensions.register_series_accessor("peptidomics")
class PandasPeptidomicsSerie:
    def __init__(self, series):
        self.series = series

    def fdr(self, alpha=0.05, method='fdr_bh'):
        if method != "fdr_bh":
            raise NotImplementedError("TODO!!!")

        mask = self.series == self.series
        corrected = np.full(self.series.shape, np.nan)
        corrected[mask] = multitest.multipletests(self.series[mask], method=method, alpha=alpha)[1]
        return pd.Series(corrected, index=self.series.index, name=("FDR", self.series.name[1]))

@pd.api.extensions.register_dataframe_accessor("peptidomics")
class PandasPeptidomics:
    rename_header = {"end": "stop", "begin": "start", "prot_acc": "protein_id",
                     "pep_start": "start", "pep_end": "stop", "pep_mod_seq": "mod_seq",
                     "pep_seq": "seq"}

    def __init__(self, df):
        self.df = df

    def create_annotation(self):
        annotation_columns = pd.MultiIndex.from_tuples((), names=['annotation', 'group'])
        return pd.DataFrame(columns=annotation_columns, index=self.df.index)

    def add_annotation(self, annotation_series):
        self.df[annotation_series.name] = annotation_series

    @classmethod
    def load_ppv_file(cls, ppv_file, campaign_id, index=None):
        print('Load ppv file')
        ppv = pd.read_csv(ppv_file, sep="\t")
        ppv.rename(columns=cls.rename_header, inplace=True)
        ppv["campaign_id"] = campaign_id
        ppv["origin"] = "collapsed"
        if index is None:
            index = ['campaign_id'] + list(PeptideVariantId._fields)
        return pd.pivot_table(ppv, index=index, values="score")

    def _ppv_pval_combine(self, subset, compare, pvalue_name="P-Value", ratio_name="Ratio"):
        gt = subset[subset[ratio_name, compare] >= 0][pvalue_name, compare] / 2
        lt = subset[subset[ratio_name, compare] <= 0][pvalue_name, compare] / 2
        gt_flip = 1 - gt
        lt_flip = 1 - lt
        pval_gt = np.append(gt.values, lt_flip.values)
        pval_lt = np.append(lt.values, gt_flip.values)
        p1 = sp.stats.combine_pvalues(pval_gt)[1]
        p2 = sp.stats.combine_pvalues(pval_lt)[1]
        pval = min(p1, p2) * 2
        return pval

    def ppv_pval(self, df_observed, compare, method="fisher",
                 annotation_name="P-Value", progresbar=True):
        if method != "fisher":
            raise NotImplementedError("Only fishers method is implemented")
        _iter = (self._ppv_pval_combine(subset, compare)
                 for subset in self._iter_ppvs(df_observed))
        if progresbar:
            _iter = tqdm(_iter, "Collapsing p-values from UPF", total=self.df.shape[0])
        return pd.Series(list(_iter), index=self.df.index, name=(annotation_name, compare))

    def _ppv_intensity_combine(self, subset, group):
        return pd.Series(subset[group].values.flatten()).mean()

    def ppv_intensity(self, df_observed, group, annotation_name="Intensity", progresbar=True):
        _iter = (self._ppv_intensity_combine(subset, group)
                 for subset in self._iter_ppvs(df_observed))
        if progresbar:
            _iter = tqdm(_iter, "Collapsing intensities from UPF", total=self.df.shape[0])

        return pd.Series(list(_iter), index=self.df.index, name=(annotation_name, group))

    def _iter_ppvs(self, df):
        for index, ppv_data in self.df.iterrows():
            campaign_id, protein_id, start, stop, mod_seq, origin = index
            subset = df.query('protein_id == "{}" and start >= {} and stop <= {}'.format(
                              protein_id, start, stop))
            yield subset

    def ppv_ratio(self, df_observed, compare, annotation_name="Ratio", progresbar=True):
        _iter = (subset[annotation_name, compare].median()
                 for subset in self._iter_ppvs(df_observed))
        if progresbar:
            _iter = tqdm(_iter, "Annotating ratio from UPF", total=self.df.shape[0])
        return pd.Series(list(_iter), index=self.df.index, name=(annotation_name, compare))

    @classmethod
    def load_upf_meta(cls, upf_file, meta_file, campaign_id, upf_file_sep="\t", meta_file_sep="\t",
                      left_on='rs_acc', right_on='accno', pivot=True, add_meta_defaults=True,
                      **kwargs):
        # TODO: campaing ID should be infered from dataset.meta
        #  print('Load, merge and create peptide id')
        upf = pd.read_csv(upf_file, sep=upf_file_sep)
        upf.rename(columns=cls.rename_header, inplace=True)
        meta = cls._load_meta_from_file(meta_file, meta_file_sep, add_meta_defaults)
        upf_sample = pd.merge(meta, upf, left_on=left_on, right_on=right_on)
        upf_sample["campaign_id"] = campaign_id
        upf_sample["origin"] = "observed"
        if pivot:
            return cls.ms_pivot_table(upf_sample, **kwargs)
        return upf_sample

    @classmethod
    def _load_meta_from_file(cls, meta_file, sep, add_defaults=True):
        meta = pd.read_csv(meta_file, sep=sep)
        if not add_defaults:
            return meta
        if 'subject' not in meta.columns or meta['subject'].dropna().shape[0] == 0:
            # if there is no subject info assume that each sample is a unique subject
            meta["subject"] = meta["ps_acc"]
        if "qc" not in meta.columns or meta['qc'].dropna().shape[0] == 0:
            # if there are no qc info, all samples were probbably fine :D
            meta["qc"] = "OK"
        return meta

    @classmethod
    def ms_pivot_table(cls, df, values='intensity',
                       index=['campaign_id', 'protein_id', 'start', 'stop', 'mod_seq', 'origin'],
                       columns=['subject', 'group', 'qc']):
        return pd.pivot_table(df, values=values, index=index, columns=columns)

    def normalize(self, full_rank='auto', full_rank_cutoff=100):
        raw_data = self.df.replace(0, np.nan)
        log10_data = np.log10(raw_data)
        log10_data_full_rank = log10_data.dropna()
        rank = log10_data_full_rank.shape[0]
        if full_rank == 'auto':
            full_rank = rank >= full_rank_cutoff
        if full_rank:
            if rank < full_rank_cutoff:
                msg = ("The full rank of the normalization matrix is {}, consider "
                       "consider running df.peptidomics.normalize(full_rank=False) instead")
                warnings.warn(msg.format(rank))
            median_intensities = log10_data_full_rank.median()
        else:
            median_intensities = log10_data.median()
        norm_target = median_intensities.mean()
        normalization_factors = norm_target - median_intensities
        #  ax = full_rank_median_intensities.plot(kind="hist",
        #                                         title="Median Intensity across samples")
        return 10 ** (log10_data + normalization_factors)

    @classmethod
    def _ttest_return(cls, t_test, index, ratios):
        pvalue = pd.Series(t_test.pvalue.data, index=index, name="P-Value")
        statistic = pd.Series(t_test.statistic.data, index=index, name="T-Statistic")
        ratio = pd.Series(ratios, index=index, name="Ratio")
        return TTestResult(pvalue, statistic, ratio)

    def _paired_ttest(self, c1, c2, n_min):
        counts = (((c1 == c1).astype(int) + (c2 == c2).astype(int)) == 2).sum(axis=1)
        mask = n_min <= counts
        t_test = sp.stats.ttest_rel(c1[mask], c2[mask], axis=1, nan_policy='omit')
        ratios = (c1[mask] - c2[mask]).mean(axis=1)
        return self._ttest_return(t_test, self.df.index[mask], ratios)

    def _unpaired_ttest(self, c1, c2, n_min, equal_var):
        _data = {'c1': (c1 == c1).sum(axis=1), 'c2': (c2 == c2).sum(axis=1)}
        counts = pd.DataFrame(_data).min(axis=1)
        mask = n_min <= counts
        t_test = sp.stats.ttest_ind(c1[mask], c2[mask], axis=1, nan_policy='omit',
                                    equal_var=equal_var)
        ratios = c1[mask].mean(axis=1) - c2[mask].mean(axis=1)
        return self._ttest_return(t_test, self.df.index[mask], ratios)

    def calculate_t_test(self, cond1, cond2, ttype='unpaired', level='group', n_min=2,
                         equal_var=False):
        c1 = self.df.xs(cond1, level=level, axis=1)
        c2 = self.df.xs(cond2, level=level, axis=1)

        if ttype == 'paired':
            return self._paired_ttest(c1, c2, n_min)
        elif ttype == 'unpaired':
            return self._unpaired_ttest(c1, c2, n_min, equal_var)
        raise ValueError("ttype must be either 'paired' or 'unpaired' not {}".format(ttype))

    @classmethod
    def peptide_id_label(cls, acc, start, end, mod_seq):
        md5sum = hashlib.md5((mod_seq).encode('utf-8')).hexdigest().upper()
        return "%s:%s:%s:%s" % (acc, start, end, md5sum)

    def is_annotation(self):
        return self.df.columns.names == ['annotation', 'group']

    def iterrows(self):
        id_class = self.df.index.peptidomics.id_class
        for id_tuple, data in self.df.iterrows():
            yield id_class(*id_tuple[1:]), data

        #  for id_tuple, data in self.

import numpy,import scipy,import statsmodels


import numpy as np
import matplotlib.pyplot as plt
import sys
from scipy import ndimage
#from databroker import get_table, db
#from skimage.filters.rank import median
from skimage.morphology import disk
from skimage import io

'''
from hxntools.handlers import register
import filestore
register()

import yaml
with open("./hxn.yml","r") as read_file:
    data=yaml.load(read_file)

from metadatastore.mds import MDS
from databroker import Broker
from databroker.core import register_builtin_handlers
from filestore.fs import FileStore

_mds_config = {'host': data['metadatastore']['config']['host'],
               'port': 27017,
               'database': data['metadatastore']['config']['database'],
               'timezone': 'US/Eastern'}
mds = MDS(_mds_config)
_fs_config = {'host': data['assets']['config']['host'],
              'port': 27017,
              'database': data['assets']['config']['database']}
db = Broker(mds, FileStore(_fs_config))

from hxntools.handlers.timepix import TimepixHDF5Handler
from hxntools.handlers.xspress3 import Xspress3HDF5Handler
db.fs.register_handler(TimepixHDF5Handler._handler_name, TimepixHDF5Handler, overwrite=True)
'''

def rm_pixel(data,ix,iy):
    data[ix,iy] = np.median(data[ix-1:ix+1,iy-1:iy+1])
    return data

def display_frame(index):
    #t = np.flipud(np.squeeze(np.asarray(filestore.api.get_data(df['merlin1'][np.ceil(index)]))))
    #print(index)
    t = diff_array[:,:,index]
    fig3 = plt.figure(3)
    plt.clf()
    cl = np.percentile(diff_array,15)
    ch = np.percentile(diff_array,99.95)
    #im2 = plt.imshow(np.flipud(np.log10(t+.001)), cmap = 'spectral', interpolation = 'none')
    im2 = plt.imshow(np.flipud(t.T), cmap = 'hot', interpolation = 'none',clim=[cl,ch])
    #im2 = plt.imshow(np.flipud(np.log10(t+0.001).T), cmap = 'hot', interpolation = 'none')
    iy = np.floor(index/num_x)
    ix = np.mod(index, num_x)
    index_new = (num_y-iy) * num_x + (num_x - ix)
    #plt.title('x: '+np.str(x_data[index])+'um, y: '+np.str(y_data[index])+'um')
    plt.title('x: %.3f' % x_data[index]+' um, y: %.3f' %y_data[index]+' um')
    plt.colorbar()
    plt.draw()

    '''
    t2 = diff_array_2[:,:,index]
    fig4 = plt.figure(4)
    plt.clf()
    #im2 = plt.imshow(np.flipud(np.log10(t+.001)), cmap = 'spectral', interpolation = 'none')
    im2 = plt.imshow(np.flipud(np.log10(t2+0.001).T), cmap = 'hot', interpolation = 'none')#,clim=[cl2,ch2])
    iy = np.floor(index/num_x)
    ix = np.mod(index, num_x)
    index_new = (num_y-iy) * num_x + (num_x - ix)
    #plt.title('x: '+np.str(x_data[index])+'um, y: '+np.str(y_data[index])+'um')
    plt.title('x: %.3f' % x_data[index]+' um, y: %.3f' %y_data[index]+' um')
    plt.colorbar()
    plt.draw()
    '''

def onclick(event):
    #print(w,l,event.xdata,event.ydata)
    index = w*(int(np.round(event.ydata)))  + int(np.round(event.xdata))
    fig = plt.figure(0)
    plt.clf()
    plt.imshow(xrf,interpolation = 'none',aspect='auto')
    plt.scatter(int(np.round(event.xdata)),int(np.round(event.ydata)),zorder=1)
    plt.title('#'+scan_num +' '+elem)
    plt.draw()

    fig1 = plt.figure(1)
    plt.clf()
    plt.imshow(roi,interpolation = 'none',aspect='auto')
    plt.scatter(int(np.round(event.xdata)),int(np.round(event.ydata)),zorder=1)
    plt.title('#'+scan_num +' Diff ROI')
    plt.draw()

    '''
    fig2 = plt.figure(2)
    plt.clf()
    im1 = plt.imshow(roi2,interpolation = 'none')
    plt.scatter(int(np.round(event.xdata)), int(np.round(event.ydata)),zorder=1)
    plt.title('#'+scan_num +' Trans ROI')
    plt.draw()
    '''
    #print(index)
    display_frame(index)

def onclick_roi(event):
    #print(w,l,event.xdata,event.ydata)
    index = w*(int(np.round(event.ydata)))  + int(np.round(event.xdata))
    fig1 = plt.figure(1)
    plt.clf()
    plt.imshow(roi,interpolation = 'none',aspect='auto')
    plt.scatter(int(np.round(event.xdata)),int(np.round(event.ydata)),zorder=1)
    plt.title('#'+scan_num +' Diff ROI')
    plt.draw()

    fig1 = plt.figure(0)
    plt.clf()
    plt.imshow(xrf,interpolation = 'none',aspect='auto')
    plt.scatter(int(np.round(event.xdata)),int(np.round(event.ydata)),zorder=1)
    plt.title('#'+scan_num +' '+elem)
    plt.draw()

    '''
    fig1 = plt.figure(2)
    plt.clf()
    im1 = plt.imshow(roi2,interpolation = 'none')
    plt.scatter(int(np.round(event.xdata)),int(np.round(event.ydata)),zorder=1)
    plt.title('#'+scan_num +' Trans ROI')
    plt.draw()
    '''
    #print(index)
    display_frame(index)

def onclick_roi2(event):
    #print(w,l,event.xdata,event.ydata)
    index = w*(int(np.round(event.ydata)))  + int(np.round(event.xdata))
    fig1 = plt.figure(2)
    plt.clf()
    im1 = plt.imshow(roi2,interpolation = 'none')
    plt.scatter(int(np.round(event.xdata)),int(np.round(event.ydata)),zorder=1)
    plt.title('#'+scan_num +' tran ROI')
    plt.draw()

    fig1 = plt.figure(0)
    plt.clf()
    im1 = plt.imshow(xrf,interpolation = 'none')
    plt.scatter(int(np.round(event.xdata)),int(np.round(event.ydata)),zorder=1)
    plt.title('#'+scan_num +' '+elem)
    plt.draw()

    fig1 = plt.figure(1)
    plt.clf()
    im1 = plt.imshow(roi,interpolation = 'none')
    plt.scatter(int(np.round(event.xdata)),int(np.round(event.ydata)),zorder=1)
    plt.title('#'+scan_num +' Diff ROI')
    plt.draw()

    #print(index)
    display_frame(index)

def onclick_fermat(event):
    index = (np.abs(x-event.xdata)**2+np.abs(y-event.ydata)**2).argmin()+1
    fig1 = plt.figure(1)
    plt.clf()
    props = dict(alpha=0.8, edgecolors='none' )
    plt.scatter(x,y,c=xrf,s=50,marker='s',**props)
    plt.axes().set_aspect('equal')
    plt.gca().invert_yaxis()
    plt.scatter(event.xdata,event.ydata,zorder=1)
    plt.draw()
    display_frame(index)

def show_diff_data(sid,element,det_name='merlin1',fermat_flag=False, save_flag=False,zp_flag=False):

    #scan_num = sys.argv[1]
    #sid = int(scan_num)
    global elem
    elem = element
    global scan_num

    '''
    #scan_num = np.str(sid)
    #scan_id, df = _load_scan(sid, fill_events=False)
    scan_num = np.str(sid)
    df = db.get_table(db[sid],fill=False)
    num_frame, count = np.shape(df)
    hdr = db[sid]
    #fermat_flag = int(sys.argv[3])
    #elem = sys.argv[2]
    #det_name = sys.argv[4]
    '''
    scan_num = np.str(sid)
    hdr = db[sid]
    bl = db.get_table(db[sid],stream_name='baseline')
    df = db.get_table(db[sid],fill=False)
    #images = db.get_images(db[sid],name=det_name)
    h = db[sid]
    images = list(h.data(det_name))
    #print('image size:',np.shape(images))
    images = np.array(np.squeeze(images))
    plan_args = db[sid].start['plan_args']

    ic = np.asfarray(df['sclr1_ch4'])
    #ic_0 = 153000

    #images = db.get_images(db[sid],name=det_name)
    images = np.array(np.squeeze(list(hdr.data(det_name))))
    print(np.shape(images))
    num_frame,nnx,nny = np.shape(images)
    #mask = 1-io.imread('/data/users/2020Q3/Huang_2020Q3/NPO_mask.tif')
    mm = np.load('/data/users/2022Q1/May_2022Q1/small_particle/mask.npy')
    #mm2 = np.load('/data/users/2021Q2/Huang_2021Q2/TMA_LCO_pristine/mask2.npy')
    #index = np.where(mask == 1)
    #mx = index[0]
    #my = index[1]
    #print(mx)
    #print(my)
    #m_num = np.shape(mx)
    #print('load mask 2')
    #mm = np.load('/data/users/2022Q1/Singer_2022Q1/S1/mask.npy')
    #mm2 = np.load('/data/users/2021Q2/Liu_2021Q2/NZ150_2/mask2.npy')
    #mm3 = np.load('/data/users/2021Q2/Liu_2021Q2/Z150_1/mask.npy')
    for i in range(num_frame):
        if np.mod(i,500) ==0:
            print('load frame ',i, '/', num_frame)
        #t = np.flipud(images.get_frame(i)[0]).T
        t = np.flipud(images[i,:,:]).T
        #t = t * mask
        t = t*ic[0] / ic[i]
        t = t * mm
        #t *= (1-mm)
        #t *= (1-mm2)
        ##t = np.flipud(t)
        #t[126,124] = 0
        #t[126,124] = 0
        #t *= (1-mm3)
        #t *= (1-mm3)
        #t *= (1-mm4)
        #t[mm2 == 1.] == 0
        #t[mm3 == 1.] == 0
        #t[mm4 == 1.] == 0
        #for jj in range(m_num[0]):
        #    t = rm_pixel(t,mx[jj],my[jj])
        '''
        plt.figure()
        plt.imshow(t)
        plt.show()
        ddd
        '''
        #if i == 0:
        #    index = np.where(t >= 5)
        #'''
        #t[96,222] = 0.
        #t[140,94] = 0.
        #t[57,138] = 0.
        #'''
        if i == 0:
            nx,ny = np.shape(t)
            global diff_array
            diff_array = np.zeros((nx,ny,num_frame))
        #t[index] = 0
        #t[mask == 1] = 0
        #t[164,107] = 0
        diff_array[:,:,i] = t #* mask

    #diff_array[22,255,:] = 0
    #diff_array[296,281,:] = 0
    #diff_array[91,221,:] = 0
    #diff_array[440,381,:] = 0

    #diff_array[diff_array > 200000] = 0
    '''
    #diff_array[diff_array > 1e4] = 0
    diff_array[419,412,:] = 0
    diff_array[431,408,:] = 0
    diff_array[118,370,:] = 0
    diff_array[145,357,:] = 0
    diff_array[142,345,:] = 0
    diff_array[57,462,:] = 0
    diff_array[206,83,:] = 0
    diff_array[115,348,:] = 0
    diff_array[118,359,:] = 0
    diff_array[446,385,:] = 0
    '''
    '''
    diff_array[191,87,:] = 0
    diff_array[124,57,:] = 0
    diff_array[51,152,:] = 0
    diff_array[114,113,:] = 0
    diff_array[111,133,:] = 0
    diff_array[134,122,:] = 0
    '''
    #if elem == 'roi':
    for i in range(num_frame):
        if i == 0:
            global roi
            roi = np.zeros(num_frame)
        roi[i] = np.sum(diff_array[:,:,i])

    global xrf
    if elem in df:
        xrf = np.asarray(df[elem])
    else:
        xrf = np.asfarray(eval('df.Det1_'+elem)) + np.asarray(eval('df.Det2_'+elem)) + np.asarray(eval('df.Det3_'+elem))
    xrf = xrf * ic[0] / (ic + 1.e-9)

    if fermat_flag:
        x = np.asarray(df.dssx)
        y = np.asarray(df.dssy)
        fig = plt.figure(1)
        ax = fig.add_subplot(111)
        props = dict(alpha=0.8, edgecolors='none' )
        im = ax.scatter(x,y,c=xrf,s=50,marker='s',**props)
        plt.axes().set_aspect('equal')
        plt.gca().invert_yaxis()

        cid = fig.canvas.mpl_connect('button_press_event', onclick_fermat)
        fig2 = plt.figure(2)
        ax2 = fig2.add_subplot(111)
    else:
        global y_data
        global x_data
        global num_x
        global num_y
        try:
            hdr.start.plan_args['num']
            #print(hdr.start.plan_args['num'])
            xrf = np.reshape(xrf,(1,hdr.start.plan_args['num']))
            roi = np.reshape(roi,(1,hdr.start.plan_args['num']))
            #roi2 = np.reshape(roi2,(1,hdr.start.plan_args['num']))
        except:
            if hdr.start.plan_name == 'grid_scan':
                xrf = np.reshape(xrf,(hdr.start.shape[0],hdr.start.shape[1]))
                roi = np.reshape(roi,(hdr.start.shape[0],hdr.start.shape[1]))
                x_data = df[hdr.start.motors[1]]
                y_data = df[hdr.start.motors[0]]
                num_x = hdr.start.shape[0]
                num_y = hdr.start.shape[1]
                #roi2 = np.reshape(roi2,(hdr.start.shape[0],hdr.start.shape[1]))
                extent = (hdr.start.plan_args['args'][2], hdr.start.plan_args['args'][1],hdr.start.plan_args['args'][6],hdr.start.plan_args['args'][5])
            elif hdr.start.plan_name == 'FlyPlan2D':
                xrf = np.reshape(xrf,(hdr.start.shape[1],hdr.start.shape[0]))
                roi = np.reshape(roi,(hdr.start.shape[1],hdr.start.shape[0]))
                #roi2 = np.reshape(roi2,(hdr.start.shape[1],hdr.start.shape[0]))
                extent = (hdr.start.plan_args['scan_end1'], hdr.start.plan_args['scan_start1'],hdr.start.plan_args['scan_end2'],hdr.start.plan_args['scan_start2'])
                #x_motor = hdr['motor1']
                if zp_flag:
                    x_motor = hdr.start['motors'][0]
                    x_data = np.asarray(df[x_motor])
                    #y_motor = hdr['motor2']
                    y_motor = hdr.start['motors'][1]
                    y_data = np.asarray(df[y_motor])
                else:
                    x_motor = hdr.start['motors'][0]
                    x_data = np.asarray(df[x_motor])
                    y_motor = hdr.start['motors'][1]
                    y_data = np.asarray(df[y_motor])
                #global num_x
                #global num_y
                num_x = hdr.start.shape[0]
                num_y = hdr.start.shape[1]
            else:
                x_motor = hdr['motor1']
                x_data = np.asarray(df[x_motor])
                y_motor = hdr['motor2']
                y_data = np.asarray(df[y_motor])
                xrf = np.reshape(xrf,(hdr.start.shape[1],hdr.start.shape[0]))
                roi = np.reshape(roi,(hdr.start.shape[1],hdr.start.shape[0]))
                #roi2 = np.reshape(roi2,(hdr.start.shape[1],hdr.start.shape[0]))
                extent = (np.nanmin(x_data), np.nanmax(x_data),np.nanmax(y_data), np.nanmin(y_data))
                #print('no num')

    fig=plt.figure(0)
    ax = fig.add_subplot(111)
    im = ax.imshow(xrf,interpolation = 'none',aspect='auto')
    plt.title('#'+scan_num +' '+ elem)
    global w
    w = xrf.shape[1]
    cid = fig.canvas.mpl_connect('button_press_event', onclick)

    fig1 = plt.figure(1)
    ax1 = fig1.add_subplot(111)
    im = ax1.imshow(roi,interpolation = 'none',aspect='auto')
    plt.title('#'+scan_num+' Diff ROI')
    #w = roi.shape[1]
    cid = fig1.canvas.mpl_connect('button_press_event', onclick_roi)

    fig3 = plt.figure(3)
    ax3 = fig3.add_subplot(111)

    plt.show()

    if save_flag:
        io.imsave('/data/users/2022Q1/May_2022Q1/small_particle/rock_'+scan_num+'_roi.tif',roi.astype(np.float32))
        io.imsave('/data/users/2022Q1/May_2022Q1/small_particle/rock_'+scan_num+'_xrf.tif',xrf.astype(np.float32))
        io.imsave('/data/users/2022Q1/May_2022Q1/small_particle/rock_'+scan_num+'_diff_data.tif',diff_array.astype(np.float32))


import os
import glob
import sys
import numpy as np
import tensorflow as tf
import DSR
import utils
import cv2
import scipy.io as sio
import time

tf.app.flags.DEFINE_string('model_path', '/where/your/model/folder', '')
tf.app.flags.DEFINE_string('run_gpu', '0', '')
tf.app.flags.DEFINE_integer('sr_scale',2,'Super resolutio scale')
FLAGS = tf.app.flags.FLAGS

def load_model(model_path):
    input_low_images = tf.placeholder(tf.float32, shape=[1,None,None,3],name = 'input_low_images')
    model_builder = DSR.Model_Build()
    main_hf,main_lf= model_builder.generator(input_low_images)
    main_out_clip = tf.cast(tf.round(tf.clip_by_value(main_hf+main_lf,0,255)),tf.uint8)
    all_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
    gen_vars = [var for var in all_vars if var.name.startswith('generator')]
    saver = tf.train.Saver(gen_vars)
    sess = tf.Session(config = tf.ConfigProto(allow_soft_placement=True))
    ckpt_path = utils.get_last_ckpt_path(model_path)
    print("model_ path : ",ckpt_path)
    saver.restore(sess,ckpt_path)

    return input_low_images,main_out_clip,main_hf,main_lf,sess

if __name__ == '__main__':
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.run_gpu
    input_image,main_output,hf_output,lf_output,sess = load_model(FLAGS.model_path)

    #######################################
    #Where is your test low-resolution images
    ########################################
    test_image_path = './Test/image/path'

    output_path = ''
    test_image_list = utils.get_image_paths(test_image_path)
    temp_time = 0

    for test_idx,test_image in enumerate(test_image_list):
        loaded_image = cv2.imread(test_image)
        feed_dict ={input_image:[loaded_image[:,:,::-1]]}

        main_out,hf_out,lf_out = sess.run([main_output,hf_output,lf_output]feed_dict=feed_dict)

        main_image = main_out[0,:,:,:]
        hf_image = hf_out[0,:,:,:]
        lf_image = lf_out[0,:,:,:]
        image_name = os.path.basename(test_image).split('.')[0]
        ################################
        #Where to Save the output images
        ################################
        #For examples

        #cv2.imwrite('Your main output path'+image_name+'.png',main_image[:,:,::-1])
        #cv2.imwrite('Your hf output path'+image_name+'.png',hf_image[:,:,::-1])
        #cv2.imwrite('Your lf output path'+image_name+'.png',lf_image[:,:,::-1])

        #If you need full negative value of high frequnecy(hf_image)
        #you can use sio.savemat() and check the document page.


import numpy as np
from PIL import Image
import cv2
import imutils
import tensorflow as tf
import tflearn
from tflearn.layers.conv import conv_2d,max_pool_2d
from tflearn.layers.core import input_data,dropout,fully_connected
from tflearn.layers.estimator import regression

# global variables
bcg = None

def resizeImage(imageName):
    basewidth = 100
    img = Image.open(imageName)
    weight_percent = (basewidth/float(img.size[0]))
    height_size = int((float(img.size[1])*float(weight_percent)))
    img = img.resize((basewidth,height_size), Image.ANTIALIAS)
    img.save(imageName)

def run_avg(image, weight):
    global bcg
    # initializing the background
    if bcg is None:
        bcg = image.copy().astype("float")
        return

    # update the background based on  weighted average
    cv2.accumulateWeighted(image, bcg, weight)

def segment(image, threshold=20):
    global bcg
    # find  diference between background and  frame
    dif = cv2.absdiff(bcg.astype("uint8"), image)

    # threshold the dif image so that we get the foreground
    img_threshold = cv2.threshold(dif,
                                threshold,
                                255,
                                cv2.THRESH_BINARY)[1]

    # get the contours in the img_threshold image
    (contour, _) = cv2.findContours(img_threshold.copy(),
                                    cv2.RETR_EXTERNAL,
                                    cv2.CHAIN_APPROX_SIMPLE)

    # return None, if no contours detected
    if len(contour) == 0:
        return
    else:
        #  getting the maximum contour which is the hand area
        img_segment = max(contour, key=cv2.contourArea)
        return (img_threshold, img_segment)

def main():
    # initialize weight
    weight = 0.6

    # (ROI) coordinates
    top, right, bottom, left = 15, 300, 200, 550

    # get the reference to the webcam
    camera = cv2.VideoCapture(0)

    # initialize num of frames
    frame_num = 0
    start_recording = False

    # looping
    while(True):
        # get the current frame
        (observed, frame) = camera.read()

        # resize the frame
        frame = imutils.resize(frame, width = 700)

        # flip the frame so that it is not the mirror view
        frame = cv2.flip(frame, 1)

        # clone the frame
        clone = frame.copy()

        # get the height and width of the frame
        (height, width) = frame.shape[:2]

        # get the ROI
        roi = frame[top:bottom, right:left]

        # convert ROI to grayscale and blurring

        gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
        gray = cv2.GaussianBlur(gray, (7, 7), 0)

        # look till threshold is reached
        if frame_num < 30:
            run_avg(gray, weight)
        else:
            # segment the hand region
            hand = segment(gray)

            # check whether hand region is img_segment
            if hand is not None:
                # if yes, unpack
                (img_threshold, img_segment) = hand

                # draw the img_segment region and display the frame
                cv2.drawContours(clone, [img_segment + (right, top)], -1, (0, 0, 255))
                if start_recording:
                    cv2.imwrite('Temp.png', img_threshold)
                    resizeImage('Temp.png')
                    predictedClass, confidence = getPredictedClass()
                    showStatistics(predictedClass, confidence)
                cv2.imshow("Thesholded", img_threshold)

        # draw the img_segment hand
        cv2.rectangle(clone, (left, top), (right, bottom), (0,255,0), 2)

        # increment the number of frames
        frame_num += 1

        # display the frame with img_segment hand
        cv2.imshow("Sign video", clone)

        # observe the keypress by the user
        keypress = cv2.waitKey(1) & 0xFF

        # if the user pressed "s", then stop
        if keypress == ord("s"):
            break

        if keypress == ord("b"):
            start_recording = True

def getPredictedClass():
    # Predict
    image = cv2.imread('Temp.png')
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    prediction = model.predict([gray_image.reshape(74, 100, 1)])
    return np.argmax(prediction), (np.amax(prediction) / (prediction[0][0] + prediction[0][1] + prediction[0][2]))

def showStatistics(predictedClass, confidence):

    textImage = np.zeros((300,512,3), np.uint8)
    className = ""

    if predictedClass == 0:
        className = "C"
    elif predictedClass == 1:
        className = "N"
    elif predictedClass == 2:
        className = "A"

    cv2.putText(textImage,"Pedicted Letter : " + className,
    (30, 30),
    cv2.FONT_HERSHEY_SIMPLEX,
    1,
    (255, 255, 255),
    2)

    cv2.putText(textImage,"Confidence percentage : " + str(confidence * 100) + '%',
    (30, 100),
    cv2.FONT_HERSHEY_SIMPLEX,
    1,
    (255, 255, 255),
    2)
    cv2.imshow("Statistics", textImage)

# Model defined
tf.compat.v1.reset_default_graph()
convnet=input_data(shape=[None,74,100,1],name='input')
convnet=conv_2d(convnet,32,2,activation='relu')
convnet=max_pool_2d(convnet,2)
convnet=conv_2d(convnet,64,2,activation='relu')
convnet=max_pool_2d(convnet,2)

convnet=conv_2d(convnet,128,2,activation='relu')
convnet=max_pool_2d(convnet,2)

convnet=conv_2d(convnet,256,2,activation='relu')
convnet=max_pool_2d(convnet,2)

convnet=conv_2d(convnet,256,2,activation='relu')
convnet=max_pool_2d(convnet,2)

convnet=conv_2d(convnet,128,2,activation='relu')
convnet=max_pool_2d(convnet,2)

convnet=conv_2d(convnet,64,2,activation='relu')
convnet=max_pool_2d(convnet,2)

convnet=fully_connected(convnet,1000,activation='relu')
convnet=dropout(convnet,0.75)

convnet=fully_connected(convnet,3,activation='softmax')

convnet=regression(convnet,optimizer='adam',learning_rate=0.001,loss='categorical_crossentropy',name='regression')

model=tflearn.DNN(convnet,tensorboard_verbose=0)

# Load Saved Model
model.load("Model/SignRecog.tfl")

main()


import os
import torch
from archmanager.finder import myArchManager
from archmanager.utils import getconfig,random_sample
from torchvision import transforms, datasets
import numpy as np
import time
import random
import math
import tensorrt as trt
from torch.autograd import Variable
from ofa.model_zoo import ofa_net
from ofa.utils import download_url,accuracy,AverageMeter
import common
from ofa.tutorial import AccuracyPredictor, FLOPsTable, LatencyTable, EvolutionFinder
import torch.backends.cudnn as cudnn
import csv

# set random seed
TRT_LOGGER = trt.Logger()
random_seed = 1
random.seed(random_seed)
np.random.seed(random_seed)
torch.manual_seed(random_seed)
print('Successfully imported all packages and configured random seed to %d!'%random_seed)
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
device = 'cuda:3' if torch.cuda.is_available() else 'cpu'

cuda_available = torch.cuda.is_available()
if cuda_available:
    torch.backends.cudnn.enabled = True
    torch.backends.cudnn.benchmark = True
    torch.cuda.manual_seed(random_seed)
    print('Using GPU.')
else:
    print('Using CPU.')

ofa_network = ofa_net('ofa_mbv3_d234_e346_k357_w1.2', pretrained=True)
print('The OFA Network is ready.')

def get_engine(onnx_file_path):
    """Attempts to load a serialized engine if available, otherwise builds a new TensorRT engine and saves it."""
    # def build_engine():
    """Takes an ONNX file and creates a TensorRT engine to run inference with"""

    with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
            common.EXPLICIT_BATCH) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
        builder.max_workspace_size = 1 << 32  # 4G
        builder.max_batch_size = 1
        with open(onnx_file_path, 'rb') as model:
            print('Beginning ONNX file parsing')
            if not parser.parse(model.read()):
                print ('ERROR: Failed to parse the ONNX file.')
                for error in range(parser.num_errors):
                    print (parser.get_error(error))
                return None

        network.get_input(0).shape = [1, 3, net_config['r'][0], net_config['r'][0]]
        # print('Completed parsing of ONNX file')
        print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))
        engine = builder.build_cuda_engine(network)
        print("Completed creating Engine")
        return engine

if cuda_available:
    # path to the ImageNet dataset
    print("Please input the path to the ImageNet dataset.\n")
    imagenet_data_path = '../ofa/imagenet_codebase/dataset/mini_imagenet'
    # if 'imagenet_data_path' is empty, download a subset of ImageNet containing 2000 images (~250M) for test
    if not os.path.isdir(imagenet_data_path):
        os.makedirs(imagenet_data_path, exist_ok=True)
        print('%s is empty. Download a subset of ImageNet for test.' % imagenet_data_path)
    print('The ImageNet dataset files are ready.')

    data_loader = torch.utils.data.DataLoader(
        datasets.ImageFolder(
            root=os.path.join(imagenet_data_path, 'val'),
            # transform=build_val_transform(224)
        ),
        batch_size=1,  # test batch size
        shuffle=True,
        num_workers=16,  # number of workers for the data loader
        pin_memory=True,
        drop_last=False,
    )
    print('The ImageNet dataloader is ready.')
else:
    print('Since GPU is not found in the environment, we skip all scripts related to ImageNet evaluation.')

gpu_ava_delay = AverageMeter()
cpu_ava_delay = AverageMeter()
resolution = [160, 176, 192, 208, 224]
onnxpath = './onnxs/tmp2.onnx'
csv_f = open('./dataset/tmep.csv', 'w', encoding='utf-8', newline='')
csv_writer = csv.writer(csv_f)
csv_writer.writerow(['arch_config', 'gpu latency', 'cpu latency',])

archmanager = myArchManager()
# configs = getconfig(archmanager.depths,archmanager.resolutions)
configs = []
for item in csv.reader(open('latency_predictor/dataset/configs/high_latency_config1.csv', 'r')):
    configs.append(list(map(int,item)))

for config in configs:
    net_config = random_sample(depths=archmanager.depths, kernel_sizes=archmanager.kernel_sizes,
                               num_blocks=archmanager.num_blocks,
                               expand_ratios=archmanager.expand_ratios, num_stages=archmanager.num_stages)
    net_config = {'wid': None,
                  'ks': [7, 7, 7, 7, 5, 3, 5, 3, 5, 3, 7, 5, 7, 5, 3, 5, 5, 7, 5, 3, 5, 5, 5, 3, 5, 5, 3, 5, 3, 3, 3, 5, 3, 7, 3, 3, 3, 7, 5, 7],
                  'e': [6, 6, 4, 3, 6, 4, 4, 6, 6, 4, 6, 3, 6, 4, 4, 6, 4, 3, 3, 3, 4, 4, 6, 4, 4, 6, 6, 6, 3, 6, 6, 6, 4, 4, 6, 6, 6, 3, 6, 6],
                  'd': [4, 3, 3, 1, 4, 3, 2, 1, 3, 4],
                  'r': [1024]}
    # "{'wid': None,
    # 'ks': [7, 7, 7, 7, 5, 3, 5, 3, 5, 3, 7, 5, 7, 5, 3, 5, 5, 7, 5, 3, 5, 5, 5, 3, 5, 5, 3, 5, 3, 3, 3, 5, 3, 7, 3, 3, 3, 7, 5, 7],
    # 'e': [6, 6, 4, 3, 6, 4, 4, 6, 6, 4, 6, 3, 6, 4, 4, 6, 4, 3, 3, 3, 4, 4, 6, 4, 4, 6, 6, 6, 3, 6, 6, 6, 4, 4, 6, 6, 6, 3, 6, 6],
    # 'd': [4, 3, 3, 1, 4, 3, 2, 1, 3, 4],
    # 'r': [1024]}",44.3389,0

    # config = [4,2,4,4,4,3,4,2,3,4,1024]
    # config = [4]*10
    # config.append(1024)
    # net_config['d'] = config[:-1]
    # net_config['r'] = [config[-1]]
    print(net_config)

    data_loader.dataset.transform = transforms.Compose([
        transforms.Resize(int(math.ceil(int(net_config['r'][0])/0.875))),
        transforms.CenterCrop(net_config['r'][0]),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )])
    cudnn.benchmark = True
    assert 'ks' in net_config and 'd' in net_config and 'e' in net_config
    assert len(net_config['ks']) == 40 and len(net_config['e']) == 40 and len(net_config['d']) == 10

    ofa_network.set_active_subnet(ks=net_config['ks'], d=net_config['d'], e=net_config['e'])
    subnet = ofa_network.get_active_subnet().to(device)

    input_name = ['input']
    output_name = ['output']
    input = Variable(torch.randn(1, 3, net_config['r'][0], net_config['r'][0])).cuda(device)
    torch.onnx.ExportTypes()
    torch.onnx.export(subnet, input, onnxpath, export_params=True,input_names=input_name, output_names=output_name, verbose=False)

    gpu_ava_delay.reset()
    cpu_ava_delay.reset()

    with get_engine(onnxpath) as engine, engine.create_execution_context() as context:
        inputs, outputs, bindings, stream = common.allocate_buffers(engine)
        # Set host input to the image. The common.do_inference function will copy the input to the GPU before executing.
        for i, (images, labels) in enumerate(data_loader):
            images, labels = images.numpy(), labels.numpy()
            inputs[0].host = images.astype(np.float32)
            t1 = time.time()
            trt_outputs = common.do_inference_v2(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
            t2 =time.time()
            print((t2-t1)*1000)
            if i >5:
                gpu_delay = (t2-t1) * 1000
                gpu_ava_delay.update(gpu_delay)
        csv_array = [net_config, round(gpu_ava_delay.avg, 4), round(cpu_ava_delay.avg, 4)]
        csv_writer.writerow(csv_array)
        csv_f.flush()


from rich import print
from rich.console import Console
import glob
import os
from spacy.lang.de import German
from spacy.lang.en import English
from spacy.lang.fr import French
import time
import pyaudio
import keyboard
from rich.progress import Progress
from rich.progress import track
from rich.table import Table
import wave
from datetime import datetime
import librosa
import struct
import numpy as np
import webrtcvad
from scipy.ndimage.morphology import binary_dilation
import soundfile
import sys
import csv

# Audio processing
chunk = 1024
sample_format = pyaudio.paInt16
channels = 1
frame_rate = 22050
timeout = 60

# For silence trimming
vad_moving_average_width = 8
int16_max = (2 ** 15) - 1
vad_window_length = 30
sample_rate = 16000
vad_max_silence_length = 6

# Source: CorentinJ - Real Time Voice Cloning - Thanks for this one!
def trim_long_silences(wav):

	# Compute the voice detection window size
	samples_per_window = (vad_window_length * sample_rate) // 1000

	# Trim the end of the audio to have a multiple of the window size
	wav = wav[:len(wav) - (len(wav) % samples_per_window)]

	# Convert the float waveform to 16-bit mono PCM
	pcm_wave = struct.pack("%dh" % len(wav), *(np.round(wav * int16_max)).astype(np.int16))

	# Perform voice activation detection
	voice_flags = []
	vad = webrtcvad.Vad(mode=3)
	for window_start in range(0, len(wav), samples_per_window):
		window_end = window_start + samples_per_window
		voice_flags.append(vad.is_speech(pcm_wave[window_start * 2:window_end * 2], sample_rate=sample_rate))
	voice_flags = np.array(voice_flags)

    # Smooth the voice detection with a moving average
	def moving_average(array, width):
		array_padded = np.concatenate((np.zeros((width - 1) // 2), array, np.zeros(width // 2)))
		ret = np.cumsum(array_padded, dtype=float)
		ret[width:] = ret[width:] - ret[:-width]
		return ret[width - 1:] / width

	audio_mask = moving_average(voice_flags, vad_moving_average_width)
	audio_mask = np.round(audio_mask).astype(np.bool)

	# Dilate the voiced regions
	audio_mask = binary_dilation(audio_mask, np.ones(vad_max_silence_length + 1))
	audio_mask = np.repeat(audio_mask, samples_per_window)

	return wav[audio_mask == True]

if __name__ == '__main__':

	console = Console()

	table = Table()
	table.add_column("TSS Dataset Creator", style="cyan")
	table.add_row("2021 - padmalcom")
	table.add_row("www.stonedrum.de")

	console.print(table)

	console.print("\nPlease select your [red]microphone[/red] (enter the device number).")
	# Initialisiere pyaudio
	pyaudio = pyaudio.PyAudio()

	# select input device
	info = pyaudio.get_host_api_info_by_index(0)
	numdevices = info.get('deviceCount')
	for i in range(0, numdevices):
		if (pyaudio.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0:
			print ("Input Device id ", i, " - ", pyaudio.get_device_info_by_host_api_device_index(0, i).get('name'))
	in_mic_id = int(input())
	console.print("You have selected [red]%s[/red] as input device." % pyaudio.get_device_info_by_host_api_device_index(0, in_mic_id).get('name'))

	# Select a project folder
	app_folder = os.path.dirname(os.path.realpath(__file__))
	project_directories = glob.glob(os.path.join(app_folder, 'project*/'))

	if len(project_directories) == 0:
		console.print("There are no project directories. Create a folder starting with 'project' that contains a 'metadata.csv'. This works best using main_generate_csv.py.")
		sys.exit(0)

	console.print("Please select a [red]project folder[/red] by entering its id (default 0 (%s))" % project_directories[0])
	for index, pd in enumerate(project_directories):
		console.print("%d: %s" % (index, pd))

	in_folder_id = input()
	if not in_folder_id:
		in_folder_id = 0
	project_folder = project_directories[int(in_folder_id)]

	metadata_csv_file = os.path.join(project_folder, 'metadata.csv')
	if not os.path.exists(metadata_csv_file):
		console.print("Project folder does not contain a metadata.csv file. Exiting.")
		sys.exit(0)

	console.print("wavs will be saved in [red]%s" % project_folder)

	# 0.1 select languages
	console.print("Please select a [red]language[/red] (default [i]%s[/i])." % 'de')
	in_lang = input()
	if not in_lang:
		in_lang = 'de'
	console.print("Language is set to [red]%s[/red]." % in_lang)

	# load csv
	with open(metadata_csv_file, encoding = "utf-8") as csv_file:
		csv_reader = csv.reader(csv_file, delimiter='|')
		csv_data = list(csv_reader)

	console.print("Found %d sentences." % len(csv_data))

	if len(csv_data) == 0:
		console.print("You need to add some sentences to your text files first. Exiting.")
		sys.exit(0)

	# 1.1 Hit Space to skip to save+next, hit backspace to discard
	console.print("[green]n[/green] = next sentence, [yellow]d[/yellow] = discard and repeat last recording, [red]e[/red] = exit recording.")

	console.print("Ready?", style="green")
	input("Press Enter to start")

	i = 0
	cancelled = False
	while i < len(csv_data):
		console.clear()

		current_sentence = csv_data[i][1]
		current_sentence = current_sentence.replace("\n", " ")
		current_sentence = current_sentence.replace("\t", " ")

		# If this wav file has been recorded before, continue
		wav_file_name = csv_data[i][0]
		if os.path.exists(os.path.join(project_folder, wav_file_name)):
			i += 1
			continue

		console.print("\n\n" + current_sentence + "\n\n", style = "black on white", justify="center")
		console.print("(%d/%d) [green]n[/green] = next sentence, [yellow]d[/yellow] = discard and repeat last recording, [blue]s[/blue] = skip, [red]e[/red] = exit recording." % ((i+1), len(csv_data)))

		start_time = time.time()
		current_time = time.time()
		frames = []
		stream = pyaudio.open(input_device_index = in_mic_id, format=sample_format, channels=channels, rate=frame_rate, frames_per_buffer=chunk, input=True)

		with Progress() as recording_progress:
			recording_task = recording_progress.add_task("[red]Recording...", total=timeout)

			#while (current_time - start_time) < timeout:
			while not recording_progress.finished:
				recording_progress.update(recording_task, completed = current_time - start_time)

				data = stream.read(chunk)
				frames.append(data)
				current_time = time.time()

				if keyboard.is_pressed('n'):
					while keyboard.is_pressed('n'):
						time.sleep(0.1)

					# Write the wav file
					data = stream.read(chunk)
					frames.append(data)
					stream.close()
					wf = wave.open(os.path.join(project_folder, wav_file_name), 'wb')
					wf.setnchannels(channels)
					wf.setsampwidth(pyaudio.get_sample_size(sample_format))
					wf.setframerate(frame_rate)
					wf.writeframes(b''.join(frames))
					wf.close()

					# trim silence?
					wav, source_sr = librosa.load(str(os.path.join(project_folder, wav_file_name)), sr=None)
					wav = trim_long_silences(wav)
					soundfile.write(str(os.path.join(project_folder, wav_file_name)), wav, source_sr)

					i += 1
					break
				elif keyboard.is_pressed("s"):
					while keyboard.is_pressed('s'):
						time.sleep(0.1)
					i += 1
					stream.close()
					break
				elif keyboard.is_pressed("d"):
					while keyboard.is_pressed('d'):
						time.sleep(0.1)
					stream.close()
					break
				elif keyboard.is_pressed("e"):
					while keyboard.is_pressed('e'):
						time.sleep(0.1)
					stream.close()
					cancelled = True
					break

			if cancelled:
				break

	pyaudio.terminate()


# -*- coding: utf-8 -*
import os
from typing import Dict

import cv2
import numpy as np

from ..monitor_base import VOS_MONITORS, MonitorBase
from .utils import (AverageMeter, fast_hist, labelcolormap, mask_colorize,
                    np_sigmoid)

@VOS_MONITORS.register
class SegMetric(MonitorBase):
    r"""
    metrics for segmentation

    Hyper-parameters
    ----------------
    gt_name: str
        gt name in training data
    img_name: str
        img name in training data
    num_classes: int
        number of classes
    ignore_label: int
        ignore label
    show_items: list
        metric names for show, ["mean_iou", "acc", "acc_cls", "fwavacc"]
    show_predict: bool
        whether show the predict mask
    show_gt: bool
        whether show the gt mask
    max_show_num: int
        the max number of images to show at all
    interval: int
        the interval for claculation
    result_path:
        the path of dir for show image
    """

    default_hyper_params = {
        "gt_name": "seg_mask",
        "img_name": "seg_img",
        "num_classes": 2,
        "ignore_label": 255,
        "show_items": ["mean_iou"],
        "avg_range": 5,
        "show_predict": True,
        "show_gt": True,
        "max_show_num": 10,
        "interval": 1,
        "result_path": ""
    }

    def __init__(self, ):
        r"""
        Arguments
        ---------
        """
        super(SegMetric, self).__init__()
        self.show_id = 0
        self.outputs = {}
        self.metric_dict = {}
        self.color_map = None

    def init(self, engine_state: Dict):
        super(SegMetric, self).init(engine_state)

    def _fast_hist(self, label_pred, label_true, num_classes, ignore_label=255):
        mask = (label_true >= 0) & (label_true < num_classes) & (label_true !=
                                                                 ignore_label)
        hist = np.bincount(
            num_classes * label_true[mask].astype(int) + label_pred[mask],
            minlength=num_classes**2).reshape(num_classes, num_classes)
        return hist

    def _draw_predict_mask(self, image, predict, num_classes):
        if self.color_map is None:
            self.color_map = labelcolormap(num_classes)
        mask = mask_colorize(predict, num_classes, self.color_map)
        mask = cv2.resize(mask, (image.shape[1], image.shape[0]),
                          interpolation=cv2.INTER_NEAREST)
        result_image = cv2.addWeighted(image, 0.6, mask, 0.4, 0)
        return result_image

    def update(self, engine_data: Dict):
        r"""
        """
        iter = engine_data["iter"]
        if iter % self._hyper_params["interval"] != 0:
            engine_data["extras"]["seg_metric"] = self.metric_dict
        else:
            # data
            seg_mask = engine_data["training_data"][
                self._hyper_params["gt_name"]]
            images = engine_data["training_data"][
                self._hyper_params["img_name"]]
            predict_data = engine_data["predict_data"][-1]
            extras = engine_data["extras"]
            if "seg_metric" not in extras:
                extras["seg_metric"] = {}
            num_classes = self._hyper_params["num_classes"]
            hist = np.zeros((num_classes, num_classes))
            ignore_label = self._hyper_params["ignore_label"]
            result_dict = {}
            for i, (image, ro,
                    lt) in enumerate(zip(images, predict_data, seg_mask)):
                image = image.cpu().numpy()
                ro = ro.cpu().detach().numpy()
                lt = lt.cpu().numpy()
                ro = ro.transpose(1, 2, 0)
                if ro.shape[0] != lt.shape[0] or ro.shape[1] != lt.shape[1]:
                    ro = cv2.resize(ro, (lt.shape[1], lt.shape[0]))
                lp = np.zeros_like(ro, dtype=np.uint8)
                lp[ro > 0.5] = 1
                lp = lp.squeeze()
                hist += fast_hist(lp.flatten(), lt.flatten(), num_classes,
                                  ignore_label)
                result_dir = self._hyper_params["result_path"]
                if not result_dir:
                    result_dir = "tmp/"
                if not os.path.exists(result_dir):
                    os.makedirs(result_dir)
                if self._hyper_params["show_predict"]:
                    image_show = image.transpose(1, 2, 0).astype(np.uint8)
                    show_image = self._draw_predict_mask(
                        image_show, lp, num_classes)
                    cv2.imwrite(
                        os.path.join(result_dir,
                                     "predict{}.png".format(self.show_id)),
                        show_image)
                if self._hyper_params["show_gt"]:
                    image_show = image.transpose(1, 2, 0).astype(np.uint8)
                    show_image = self._draw_predict_mask(
                        image_show, lt, num_classes)
                    cv2.imwrite(
                        os.path.join(result_dir,
                                     "gt{}.png".format(self.show_id)),
                        show_image)
                self.show_id += 1
                if self.show_id >= self._hyper_params["max_show_num"]:
                    self.show_id = 0

            result_dict["acc"] = np.diag(hist).sum() / hist.sum()
            acc_cls = np.diag(hist) / hist.sum(axis=1)
            result_dict["acc_cls"] = np.nanmean(acc_cls)
            iu = np.diag(hist) / (hist.sum(axis=1) + hist.sum(axis=0) -
                                  np.diag(hist))
            result_dict["mean_iou"] = np.nanmean(iu)
            freq = hist.sum(axis=1) / hist.sum()
            result_dict["fwavacc"] = (freq[freq > 0] * iu[freq > 0]).sum()
            for show_key in self._hyper_params["show_items"]:
                if show_key not in self.outputs:
                    self.outputs[show_key] = AverageMeter(
                        self._hyper_params["avg_range"])
                self.outputs[show_key].update(result_dict[show_key])
            for key in self.outputs.keys():
                extras["seg_metric"][key] = self.outputs[key].get_mean()
            self.metric_dict = extras["seg_metric"]


#!/home/pi/.virtualenvs/owl/bin/python3
from algorithms import exg, exg_standardised, exg_standardised_hue, hsv, exgr, gndvi
from button_inputs import Selector, Recorder
from image_sampler import image_sample
from imutils.video import VideoStream, FileVideoStream, FPS
from relay_control import Controller
from queue import Queue
from time import strftime
import subprocess
import imutils # reference PyImageSearch
import shutil
import numpy as np # reference
import time
import sys
import cv2 # reference
import os

def nothing(x):
    pass

def green_on_brown(image, exgMin=30, exgMax=250, hueMin=30, hueMax=90, brightnessMin=5, brightnessMax=200, saturationMin=30,
                   saturationMax=255, minArea=1, headless=True, algorithm='exg'):
    '''
    Uses a provided algorithm and contour detection to determine green objects in the image. Min and Max
    thresholds are provided.
    :param image: input image to be analysed
    :param exgMin:
    :param exgMax:
    :param hueMin:
    :param hueMax:
    :param brightnessMin:
    :param brightnessMax:
    :param saturationMin:
    :param saturationMax:
    :param minArea: minimum area for the detection - used to filter out small detections
    :param headless: True: no windows display; False: watch what the algorithm does
    :param algorithm: the algorithm to use. Defaults to ExG if not correct
    :return:
    '''

    # different algorithm options, add in your algorithm here if you make a new one!
    threshedAlready = False
    if algorithm == 'exg':
        output = exg(image)

    elif algorithm == 'exgr':
        output = exgr(image)

    elif algorithm == 'nexg':
        output = exg_standardised(image)

    elif algorithm == 'exhsv':
        output = exg_standardised_hue(image, hueMin=hueMin, hueMax=hueMax,
                                      brightnessMin=brightnessMin, brightnessMax=brightnessMax,
                                      saturationMin=saturationMin, saturationMax=saturationMax)

    elif algorithm == 'hsv':
        output, threshedAlready = hsv(image, hueMin=hueMin, hueMax=hueMax,
                                      brightnessMin=brightnessMin, brightnessMax=brightnessMax,
                                      saturationMin=saturationMin, saturationMax=saturationMax)

    elif algorithm == 'gndvi':
        output = gndvi(image)

    else:
        output = exg(image)
        print('[WARNING] DEFAULTED TO EXG')

    if not headless:
        cv2.imshow("Threshold", output)

    # run the thresholds provided
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    # if not a binary image, run an adaptive threshold on the area that fits within the thresholded bounds.
    if not threshedAlready:
        output = np.where(output > exgMin, output, 0)
        output = np.where(output > exgMax, 0, output)
        output = np.uint8(np.abs(output))
        if not headless:
            cv2.imshow("post", output)

        thresholdOut = cv2.adaptiveThreshold(output, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 31, 2)
        thresholdOut = cv2.morphologyEx(thresholdOut, cv2.MORPH_CLOSE, kernel, iterations=1)

    # if already binary, run morphological operations to remove any noise
    if threshedAlready:
        thresholdOut = cv2.morphologyEx(output, cv2.MORPH_CLOSE, kernel, iterations=5)

    if not headless:
        cv2.imshow("Threshold", thresholdOut)

    # find all the contours on the binary images
    cnts = cv2.findContours(thresholdOut.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)
    weedCenters = []
    boxes = []
    # loop over all the detected contours and calculate the centres and bounding boxes
    for c in cnts:
        # filter based on total area of contour
        if cv2.contourArea(c) > minArea:
            # calculate the min bounding box
            startX, startY, boxW, boxH = cv2.boundingRect(c)
            endX = startX + boxW
            endY = startY + boxH
            cv2.rectangle(image, (int(startX), int(startY)), (endX, endY), (0, 0, 255), 2)
            # save the bounding box
            boxes.append([startX, startY, boxW, boxH])
            # compute box center
            centerX = int(startX + (boxW / 2))
            centerY = int(startY + (boxH / 2))
            weedCenters.append([centerX, centerY])

    # returns the contours, bounding boxes, centroids and the image on which the boxes have been drawn
    return cnts, boxes, weedCenters, image

# the
class Owl:
    def __init__(self, video=False, videoFile=None, recording=False, nozzleNum=4, headless=True,
                 exgMin=30, exgMax=180, hueMin=30,hueMax=92, brightnessMin=5, brightnessMax=200,
                 saturationMin=30, saturationMax=255, resolution=(832, 624), framerate=32):

        # different detection parameters
        self.headless = headless
        self.recording = recording
        self.resolution = resolution
        self.framerate = framerate

        # threshold parameters for different algorithms
        self.exgMin = exgMin
        self.exgMax = exgMax
        self.hueMin = hueMin
        self.hueMax = hueMax
        self.saturationMin = saturationMin
        self.saturationMax = saturationMax
        self.brightnessMin = brightnessMin
        self.brightnessMax = brightnessMax

        # setup the track bars if headless is False
        if not self.headless:
            # create trackbars for the threshold calculation
            cv2.namedWindow("Params")
            cv2.createTrackbar("thresholdMin", "Params", self.exgMin, 255, nothing)
            cv2.createTrackbar("thresholdMax", "Params", self.exgMax, 255, nothing)

        # instantiate the recorder if recording is True
        if self.recording:
            self.fourcc = cv2.VideoWriter_fourcc(*'MJPG')
            self.writer = None

        else:
            self.record = False
            self.saveRecording = False

        # check if test video or videostream from camera
        if video:
            self.cam = FileVideoStream(videoFile).start()
        # if no video, start the camera with the provided parameters
        else:
            try:
                self.cam = VideoStream(usePiCamera=True, resolution=self.resolution, framerate=self.framerate).start()
            except ModuleNotFoundError:
                self.cam = VideoStream(src=0).start()
            time.sleep(1.0)
        # set the sprayqueue size
        self.sprayQueue = Queue(maxsize=10)

        # nozzleDict maps the reference nozzle number to a boardpin on the embedded device
        self.nozzleDict = {
            0: 13,
            1: 15,
            2: 16,
            3: 18
            }

        ### Data collection only ###
        # algorithmDict maps pins to algorithms for data collection
        self.algorithmDict = {
            "exg": 29,
            "nexg": 31,
            "hsv": 33,
            "exhsv": 35,
        }
        # this is where the recording button can be added. Currently set to pin 37
        if self.recording:
            self.recorderButton = Recorder(recordGPIO=37)
        ############################

        # instantiate the nozzle controller - successful start should beep the buzzer
        self.controller = Controller(nozzleDict=self.nozzleDict)

        # instantiate the logger
        self.logger = self.controller.logger

        # sensitivity and weed size to be added
        self.sensitivity = None
        self.laneCoords = {}

        # add the total number of nozzles. This can be changed easily, but the nozzleDict and physical relays would need
        # to be updated too. Fairly straightforward, so an opportunity for more precise application
        self.nozzleNum = nozzleNum

    def hoot(self, sprayDur, sample=False, sampleDim=400, saveDir='output', camera_name='cam1', algorithm='exg', selectorEnabled=False, minArea=10):
        # track FPS and framecount
        fps = FPS().start()
        if selectorEnabled:
            self.selector = Selector(switchDict=self.algorithmDict)

        try:
            while True:
                frame = self.cam.read()
                if selectorEnabled:
                    algorithm, newAlgorithm = self.selector.algorithm_selector(algorithm)
                    if newAlgorithm:
                        self.logger.log_line('[NEW ALGO] {}'.format(algorithm))

                if self.recording:
                    self.record = self.recorderButton.record
                    self.saveRecording = self.recorderButton.saveRecording

                if frame is None:
                    fps.stop()
                    print("[INFO] Stopped. Approximate FPS: {:.2f}".format(fps.fps()))
                    self.stop()
                    break

                if self.record and self.writer is None:
                    saveDir = os.path.join(saveDir, strftime("%Y%m%d-{}-{}".format(camera_name, algorithm)))
                    if not os.path.exists(saveDir):
                        os.makedirs(saveDir)

                    self.baseName = os.path.join(saveDir, strftime("%Y%m%d-%H%M%S-{}-{}".format(camera_name, algorithm)))
                    videoName = self.baseName + '.avi'
                    self.logger.new_video_logfile(name=self.baseName + '.txt')
                    self.writer = cv2.VideoWriter(videoName, self.fourcc, 30, (frame.shape[1], frame.shape[0]), True)

                # retrieve the trackbar positions for thresholds
                if not self.headless:
                    self.exgMin = cv2.getTrackbarPos("thresholdMin", "Params")
                    self.exgMax = cv2.getTrackbarPos("thresholdMax", "Params")

                else:
                    # this leaves it open to adding dials for sensitivity. Static at the moment, but could be dynamic
                    self.update(exgMin=self.exgMin, exgMax=self.exgMax) # add in update values here

                # pass image, thresholds to green_on_brown function
                cnts, boxes, weedCentres, imageOut = green_on_brown(frame.copy(), exgMin=self.exgMin,
                                                                    exgMax=self.exgMax,
                                                                    hueMin=self.hueMin,
                                                                    hueMax=self.hueMax,
                                                                    saturationMin=self.saturationMin,
                                                                    saturationMax=self.saturationMax,
                                                                    brightnessMin=self.brightnessMin,
                                                                    brightnessMax=self.brightnessMax,
                                                                    headless=self.headless,
                                                                    algorithm=algorithm, minArea=minArea)

                ##### IMAGE SAMPLER #####
                # record sample images if required of weeds detected
                # uncomment if needed
                # if frameCount % 60 == 0 and sample is True:
                #     saveFrame = frame.copy()
                #     sampleThread = Thread(target=image_sample, args=[saveFrame, weedCentres, saveDir, sampleDim])
                #     sampleThread.start()
                #########################

                # activation region limit - once weed crosses this line, nozzle is activated
                self.yAct = int((0.2) * frame.shape[0])
                laneWidth = imageOut.shape[1] / self.nozzleNum

                # calculate lane coords and draw on frame
                for i in range(self.nozzleNum):
                    laneX = int(i * laneWidth)
                    # cv2.line(displayFrame, (laneX, 0), (laneX, imageOut.shape[0]), (0, 255, 255), 2)
                    self.laneCoords[i] = laneX

                # loop over the ID/weed centres from contours
                for ID, centre in enumerate(weedCentres):
                    # if they are in activation region the spray them
                    if centre[1] > self.yAct:
                        sprayTime = time.time()
                        for i in range(self.nozzleNum):
                            # determine which lane needs to be activated
                            if int(self.laneCoords[i]) <= centre[0] < int(self.laneCoords[i] + laneWidth):
                                # log a spray job with the controller using the nozzle, timestamp and spray duration
                                self.controller.receive(nozzle=i, timeStamp=sprayTime, duration=sprayDur)

                # update the framerate counter
                fps.update()

                if not self.headless:
                    cv2.imshow("Output", imutils.resize(imageOut, width=600))

                if self.record and not self.saveRecording:
                    self.writer.write(frame)

                if self.saveRecording and not self.record:
                    self.writer.release()
                    self.controller.solenoid.beep(duration=0.1)
                    self.recorderButton.saveRecording = False
                    fps.stop()
                    self.writer = None
                    self.logger.log_line_video("[INFO] {}. Approximate FPS: {:.2f}".format(self.baseName, fps.fps()), verbose=True)
                    fps = FPS().start()

                k = cv2.waitKey(1) & 0xFF
                if k == 27:
                    fps.stop()
                    self.logger.log_line_video("[INFO] Stopped. Approximate FPS: {:.2f}".format(fps.fps()), verbose=True)
                    self.stop()
                    break

        except KeyboardInterrupt:
            fps.stop()
            self.logger.log_line_video("[INFO] Stopped. Approximate FPS: {:.2f}".format(fps.fps()), verbose=True)
            self.stop()

        except Exception as e:
            self.controller.solenoid.beep(duration=0.5, repeats=5)
            self.logger.log_line("[CRITICAL ERROR] STOPPED: {}".format(e))

    # still in development
    def update_software(self):
        USBDir, USBConnected = check_for_usb()
        if USBConnected:
            files = os.listdir(USBDir)
            workingDir = '/home/pi'

            # move old version to version control directory first
            oldVersionDir = strftime(workingDir + "/%Y%m%d-%H%M%S_update")
            os.mkdir(oldVersionDir)

            currentDir = '/home/pi/owl'
            shutil.move(currentDir, oldVersionDir)

            # move new directory to working directory
            for item in files:
                if 'owl' in item:
                    shutil.move()

    def stop(self):
        self.controller.running = False
        self.controller.solenoid.all_off()
        self.controller.solenoid.beep(duration=0.1)
        self.controller.solenoid.beep(duration=0.1)
        self.cam.stop()
        if self.record:
            self.writer.release()
            self.recorderButton.running = False

        if not self.headless:
            cv2.destroyAllWindows()

        sys.exit()

    def update(self, exgMin=30, exgMax=180):
        self.exgMin = exgMin
        self.exgMax = exgMax

def check_for_usb():
    try:
        nanoMediaFolder = 'ls /media/pi'
        proc = subprocess.Popen(nanoMediaFolder, shell=True, preexec_fn=os.setsid, stdout=subprocess.PIPE)
        usbName = proc.stdout.readline().rstrip().decode('utf-8')

        if len(usbName) > 0:
            print('[INFO] Saving to {} usb'.format(usbName))
            saveDir = '/media/pi/{}/'.format(usbName)
            return saveDir, True

        else:
            print('[INFO] No USB connected. Saving to videos')
            saveDir = '/home/pi/owl/videos'
            return saveDir, False

    except AttributeError:
        print('[INFO] Windows computer detected...')
        saveDir = '/videos/'
        return saveDir, False

# business end of things
if __name__ == "__main__":
    owl = Owl(video=False,
              videoFile=r'',
              headless=True,
              recording=False,
              exgMin=25,
              exgMax=200,
              hueMin=39,
              hueMax=83,
              saturationMin=50,
              saturationMax=220,
              brightnessMin=60,
              brightnessMax=190,
              framerate=32,
              resolution=(416, 320))

    # start the targeting!
    owl.hoot(sprayDur=0.15,
             sample=False,
             sampleDim=1000,
             saveDir='/home/pi',
             algorithm='exhsv',
             selectorEnabled=False,
             camera_name='hsv',
             minArea=10)


import hydra
import numpy as np
import pandas as pd
import torch
from botorch.utils.multi_objective import pareto

from lambo.candidate import FoldedCandidate
from lambo.tasks.base_task import BaseTask

class ProxyRFPTask(BaseTask):
    def __init__(self, tokenizer, candidate_pool, obj_dim, transform=lambda x: x,
                 num_start_examples=1024, **kwargs):
        super().__init__(tokenizer, candidate_pool, obj_dim, transform, **kwargs)
        self.op_types = ["sub"]
        self.num_start_examples = num_start_examples

    def task_setup(self, config, project_root=None, *args, **kwargs):
        project_root = hydra.utils.get_original_cwd() if project_root is None else project_root
        work_dir = f'{project_root}/{config.log_dir}/{config.job_name}/{config.timestamp}/foldx'
        rfp_known_structures = pd.read_csv(
            f'{project_root}/lambo/assets/fpbase/rfp_known_structures.csv'
        )

        all_seqs = rfp_known_structures.foldx_seq.values
        all_targets = np.stack([
            -rfp_known_structures.SASA.values,
            rfp_known_structures.foldx_total_energy.values,
        ], axis=-1)

        seed_data = pd.read_csv(
            f'{project_root}/lambo/assets/fpbase/proxy_rfp_seed_data.csv'
        )
        seed_data = seed_data.sample(self.num_start_examples)
        sample_batch_targets = np.stack([
            -seed_data.SASA.values,
            -seed_data.stability.values,
        ], axis=-1)

        all_seqs = np.concatenate((all_seqs, seed_data.foldx_seq.values))
        all_targets = np.concatenate((all_targets, sample_batch_targets))

        seq_len_mask = np.array([len(x) <= self.max_len for x in all_seqs])
        all_seqs = all_seqs[seq_len_mask]
        all_targets = all_targets[seq_len_mask]

        # filter candidate sequences by length
        foldx_seq_len = rfp_known_structures.foldx_seq.apply(lambda x: len(x))
        rfp_known_structures = rfp_known_structures[foldx_seq_len <= self.max_len]
        rfp_known_structures.reset_index(inplace=True)

        # find valid, non-dominated starting candidates
        valid_targets = np.stack([
            -rfp_known_structures.SASA.values,
            rfp_known_structures.foldx_total_energy.values,
        ], axis=-1)
        pareto_mask = pareto.is_non_dominated(-torch.tensor(valid_targets))
        base_targets = valid_targets[pareto_mask]

        base_candidates = []
        for row_idx, datum in rfp_known_structures.iterrows():
            if not pareto_mask[row_idx]:
                continue
            print(f'{datum.Name} is non-dominated, adding to start pool')
            pdb_id = datum.pdb_id.lower()
            chain_id = datum.longest_chain
            parent_pdb_path = f'{project_root}/lambo/assets/foldx/{pdb_id}_{chain_id}/wt_input_Repair.pdb'
            base_candidates.append(
                FoldedCandidate(work_dir, parent_pdb_path, [], self.tokenizer,
                                skip_minimization=True, chain=chain_id, wild_name=datum.Name)
            )
        base_candidates = np.array(base_candidates).reshape(-1)

        return base_candidates, base_targets, all_seqs, all_targets

    def make_new_candidates(self, base_candidates, new_seqs):
        assert base_candidates.shape[0] == new_seqs.shape[0]
        new_candidates = []
        for b_cand, n_seq in zip(base_candidates, new_seqs):
            b_seq = b_cand.mutant_residue_seq
            assert len(b_seq) == len(n_seq), 'FoldX only accepts substitutions'
            mutation_ops = []
            for i, (b_char, n_char) in enumerate(zip(b_seq, n_seq)):
                if not b_char == n_char:
                    mutation_ops.append(b_cand.new_mutation(i, n_char, 'sub'))
            # mutation_ops = mutation_list(b_cand.mutant_residue_seq, n_seq)
            new_candidates.append(b_cand.new_candidate(mutation_ops))
        return np.stack(new_candidates)

    def _evaluate(self, x, out, *args, **kwargs):
        assert x.ndim == 2
        x_cands, x_seqs, f_vals = [], [], []
        for query_pt in x:
            cand_idx, mut_pos, mut_res_idx, _ = query_pt
            base_candidate = self.candidate_pool[cand_idx]
            mut_res = self.tokenizer.sampling_vocab[mut_res_idx]
            mut_list = [base_candidate.new_mutation(mut_pos, mut_res, mutation_type='sub')]
            candidate = base_candidate.new_candidate(mut_list)
            x_cands.append(candidate)
            x_seqs.append(candidate.mutant_residue_seq)
        x_seqs = np.array(x_seqs).reshape(-1)
        x_cands = np.array(x_cands).reshape(-1)

        out["X_cand"] = x_cands
        out["X_seq"] = x_seqs
        norm_scores = self.transform(self.score(x_cands))
        out["F"] = norm_scores

    def score(self, candidates):
        f_vals = []
        for cand in candidates:
            f_vals.append(np.array([
                -cand.mutant_surface_area,
                cand.mutant_total_energy,
            ]))
        return np.stack(f_vals)


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

import numpy as np
import matplotlib.pyplot as plt
import glob
import time
from collections import deque

import os
# os.environ['CUDA_VISIBLE_DEVICES']='-1'

from tracker import tracker,detector,helpers
import cv2
import tqdm
import argparse
import config
from data_management import get_mapped_dir_lists,sort_frames
from data_utils import box_to_str

class single_tracker(object):
    def __init__(self,max_age=65,frame_count=0,min_hits=1):
        self.tracker_list=[]
        self.track_id_list= deque([str(i) for i in range(100)])
        self.max_age=max_age
        self.frame_count=frame_count
        self.min_hits=min_hits
        self.detect_track_iou_thres=0.6
        self.box_ios_thres=0.8
        self.otsu_iou_thres=0.5
        self.current_tracker=None

# Global variables to be used by funcitons of VideoFileClop
frame_count = 0 # frame counter

max_age = 20  # no.of consecutive unmatched detection before
             # a track is deleted

min_hits =1  # no. of consecutive matches needed to establish a track

tracker_list =[] # list for trackers
# list for track ID
track_id_list= deque(['1', '2', '3', '4', '5', '6', '7', '7', '8', '9', '10'])

debug = False
'''
Function to track single person

'''
def pipeline_single_tracker(det,img,otsu_box,track,draw=False):
    track.frame_count+=1
    org_im=img.copy()
    img_dim = (img.shape[1], img.shape[0])
    if debug:
        print('\nFrame:', track.frame_count,' \n')
    #Detect person in the image
    detect_box = det.get_localization(img) # measurement
    final_box=[]
    improved=False
    #check for small box
    if len(detect_box)!=0:
        detect_box=helpers.remove_small_box(detect_box[0],height_limit=150,width_limit=150)

    #If Detected
    if len(detect_box)!=0:
        if debug:
            print("Detection found")
        # detect_box=detect_box[0]
        if draw:
           img1= helpers.draw_box_label('Org Det.',img, detect_box, box_color=(255, 0, 0))
        #Tracker alive or not
        if track.current_tracker!=None:
        #If alive
            if debug:
                print("Tracker Alive")
            track_box=track.current_tracker.box
            #Track result matches,detection or not
            #------------------------------------
            #If matches
            if helpers.box_iou2(track_box,detect_box)>track.detect_track_iou_thres:

                #Check abnormal detect box
                #Abnormal, use previous box NOTE can be improved
                # detect_area=helpers.find_area(detect_box)
                # track_area=helpers.find_area(track_box)

                height_d,width_d=helpers.find_dim(detect_box)
                height_t,width_t=helpers.find_dim(track_box)
                delta=0.2
                delta2=0.3
                if height_d<(1-delta)*height_t or width_d<(1-delta)*width_t or height_d>(1+delta2)*height_t or width_d>(1+delta2)*width_t:
                    if debug:
                        print("Detection improved by tracker")
                    improved=True
                    detect_box=track.current_tracker.box
                # detect_box=helpers.union_box(track_box,detect_box)
            #Track box does not matched
            elif otsu_box==True and helpers.box_ios(detect_box,track_box)>track.box_ios_thres and helpers.box_iou2(track_box,helpers.largest_contour(img))>track.otsu_iou_thres:
                if debug:
                    print("Detect box is subset of track. Track box and Otsu are similar.")
                    print("Detection improved by tracker")
                detect_box=track.current_tracker.box
                improved=True
            else:
                if debug:
                    print("Tracker lost deleting the current tracker")
                track.tracker_list.append(track.current_tracker)
                track.current_tracker=None

        #Improve detect_box by Otsu or any other way
        if otsu_box==True:
            ret,detect_box=helpers.detection_otsu(img,detect_box,draw=True,threshold=track.otsu_iou_thres)
        #Update or create tracker
        #Update if exist or matched
        if track.current_tracker!=None:
            final_box = detect_box
            z = np.expand_dims(detect_box, axis=0).T
            track.current_tracker.kalman_filter(z)
            xx = track.current_tracker.x_state.T[0].tolist()
            xx =[xx[0], xx[2], xx[4], xx[6]]
            if improved:
                final_box = xx
            track.current_tracker.box =xx
            track.current_tracker.hits += 1
            track.current_tracker.no_losses =0
        else:
            final_box = detect_box
            z = np.expand_dims(detect_box, axis=0).T
            track.current_tracker = tracker.Tracker() # Create a new tracker
            x = np.array([[z[0], 0, z[1], 0, z[2], 0, z[3], 0]]).T
            track.current_tracker.x_state = x
            track.current_tracker.predict_only()
            xx = track.current_tracker.x_state
            xx = xx.T[0].tolist()
            xx =[xx[0], xx[2], xx[4], xx[6]]
            track.current_tracker.box =xx

            track.current_tracker.id = track.track_id_list.popleft() # assign an ID for the tracker
            if debug:
                print("New Tracker\n ID: ",track.current_tracker.id)

    #Not Detection
    else:
        #Tracker alive or not
        #alive
        if track.current_tracker!=None:
            if debug:
                print("Tracker Alive")
            track.current_tracker.predict_only()
            xx = track.current_tracker.x_state
            xx = xx.T[0].tolist()
            xx =[xx[0], xx[2], xx[4], xx[6]]

            if otsu_box==True:
            # if False:
                current_state=xx
                flag,current_otsu=helpers.tracker_otsu(img,current_state,draw=True,threshold=track.otsu_iou_thres)
                if not flag:
                    if debug:
                        print("Tracker does not matched with Otsu box, Tracker id",track.current_tracker.id)
                    xx=helpers.remove_small_box(xx,height_limit=150,width_limit=150)

                    if len(xx)==0:
                        if debug:
                            print("Small track box. Deleting...............")
                        track.tracker_list.append(track.current_tracker)
                        track.current_tracker=None
                        final_box = []
                    else:
                        track.current_tracker.no_losses+=1
                        track.current_tracker.box =xx
                        final_box = xx
                else:
                    if debug:
                        print("Tracker box matched with Otsu box, Tracker id",track.current_tracker.id)

                    track.current_tracker.no_losses += 0.5
                    final_box = current_otsu
                    current_otsu = np.expand_dims(current_otsu, axis=0).T

                    track.current_tracker.kalman_filter(current_otsu)
                    xx = track.current_tracker.x_state.T[0].tolist()
                    xx =[xx[0], xx[2], xx[4], xx[6]]

                    track.current_tracker.box =xx
            else:
                if debug:
                    print("No Otsu")
                xx=helpers.remove_small_box(xx,height_limit=150,width_limit=150)
                if len(xx)==0:
                    if debug:
                        print("Small track box. Deleting...............")
                    track.tracker_list.append(track.current_tracker)
                    track.current_tracker=None
                    final_box = []
                else:
                    track.current_tracker.no_losses += 1
                    track.current_tracker.box =xx
                    final_box = xx

            #---------------------
            #Person left the frames or not
                #If left
                #Not left, no detection
        #Not active tracker
        else:
            if debug:
                print("No tracked Box ")

    #Final box
    if track.current_tracker!=None:
        if ((track.current_tracker.hits >= min_hits) and (track.current_tracker.no_losses <=max_age)):
             # final_box = track.current_tracker.box
             if debug:
                 print('updated box: ', final_box)
                 print()
             if draw:
                 img= helpers.draw_box_label("Final",img, final_box,show_label=True) # Draw the bounding boxes on the
                 img= helpers.draw_box_label(track.current_tracker.id,img, track.current_tracker.box,box_color=(255, 255,0),show_label=False) # Draw the bounding boxes on the
        elif track.current_tracker.no_losses >max_age:
            if debug:
                print('Tracker age criteria is not satisfied. Deleting..........')
            track.tracker_list.append(track.current_tracker)
            track.current_tracker=None
        else:
            if debug:
                print('Tracker zero hit')
            if draw:
                img= helpers.draw_box_label("Final",img, final_box,show_label=True) # Draw the bounding boxes on the
    if debug:
        print("Final Box")
        print(final_box)
    return final_box,img

def tracking_frames(detection_threshold,frames,numbers,otsu_box=True):
    """
        Function to track and return boxes
    """
    detector_ = detector.PersonDetector(threshold=float(detection_threshold),model_path=config.detector_model_path)
    track=single_tracker()
    total=len(frames)
    frame_num=[]
    boxes=[]
    for i in tqdm.tqdm(range(total)):
        frame=frames[i]
        number=numbers[i]
        img=cv2.imread(frame)

        box,new_img = pipeline_single_tracker(detector_,img,otsu_box,track,draw=False)
        if len(box)!=0:
            frame_num.append(number)
            boxes.append(box)
    return boxes,frame_num

def tracking_frames_to_csv(detector,frames,numbers,output_path,otsu_box=True,visualize=False):
    """
    Function to track and save the corner coordinates in csv
    input:
        frames:- list of frames path in sorted order
        numbers: the sorted numbes as mentioned in the image name
        Output_path: output csv path
        otsu_box: True or False
        visualize:  True or False
    """
    track=single_tracker()
    filename=os.path.basename(output_path)
    output_csvname=os.path.join(output_path,filename+'.csv')
    outFile = open(output_csvname, 'w')
    print("----------------------")
    print("csv output path-",output_csvname)
    print("----------------------")
    outFile.write(','.join(['Frame number', 'Track box (y_up x_left y_down x_right)']) + '\n');
    total=len(frames)
    for i in tqdm.tqdm(range(total)):
        frame=frames[i]
        number=numbers[i]
        img=cv2.imread(frame)
        np.asarray(img)

        box,new_img = pipeline_single_tracker(detector,img,otsu_box,track,draw=visualize)
        if len(box)!=0:
            outFile.write(','.join([str(number),box_to_str(box)]) + '\n');

        if visualize:
            cv2.imshow("frame",new_img)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    outFile.close()

    cv2.destroyAllWindows()

"""
Function to track and save the video visualizations
input:
    frames:- list of frames path in sorted order
    Output_path:
    otsu_box:
    visualize:
"""
def tracking_frames_to_video(detector,frames,output_path,frame_rate=10.0,otsu_box=True,visualize=False):
    track=single_tracker()
    output_filename=os.path.join(output_path,'Track.avi')

    print("----------------------")
    print("Video output path-",output_filename)
    print("----------------------")
    # start=time.time()

    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter(output_filename,fourcc, frame_rate, (640,480))
    for frame in tqdm.tqdm(frames):
        img=cv2.imread(frame)
        np.asarray(img)

        # new_img = pipeline(det,img,otsu_box)

        _,new_img = pipeline_single_tracker(detector,img,otsu_box,track,draw=True)
        # cv2.imshow("frame",new_img)
        out.write(new_img)
        if visualize:
            cv2.imshow("frame",new_img)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    cv2.destroyAllWindows()
    out.release()

if __name__ == "__main__":

    # print("Available GPU devices")
    parser = argparse.ArgumentParser(description='Person tracking on TSF dataset with Tensorflow API model and Kalman Filtering')
    parser.add_argument('--detection_threshold', default=0.3,
                        help='Threshold for object detection')
    parser.add_argument('--output_type', default='csv',
                             help='Type of output-csv or video')
    parser.add_argument('--otsu_box',default=True,
                help='Set True or False for contour box localization')
    parser.add_argument('--visualize',default=False,
                help='It will show the tracking boxes at run time')
    parser.add_argument('--debug',default=False,
                help='Print the intermediate steps for each image')
    args = parser.parse_args()

    otsu_box=False
    if args.otsu_box=='True':
        print("-------------------------")
        print("Using Otsu Contour box")
        print("-------------------------")
        otsu_box=True

    if args.debug=='True':
        debug=True

    visualize=False
    if args.visualize=='True':
        print("-------------------------")
        print("Viusal mode: All the frames will be shown")
        print("-------------------------")
        visualize=True
        #Initialization of tracker and detector
    detector_ = detector.PersonDetector(threshold=float(args.detection_threshold),model_path=config.detector_model_path)

    #Output video frame rate
    frame_rate=10.0
    root_drive = config.root_drive
    track_root_folder=config.track_root_folder
    output_type=args.output_type
    dset=config.root_folder#The name of dataset should match the folder name

    output_dir=root_drive+'/'+track_root_folder+'/'+output_type
    #Getting input and output folder mapping
    ADL_list,Fall_list=get_mapped_dir_lists(dset,output_dir=output_dir,d_type='frame')
    if len(ADL_list)==0:
        print("Dataset directory not found")

    #Tracking fall videos
    for input_path,output_path in tqdm.tqdm(Fall_list):
        print("------------------")
        print("input_path:",input_path)
        print("------------------")
        os.makedirs(output_path, exist_ok=True)
        frames = glob.glob(input_path+'/*.jpg') + glob.glob(input_path+'/*.png')
        frames,numbers = sort_frames(frames, dset)
        if output_type=='video':
            tracking_frames_to_video(detector_,frames,output_path,frame_rate,otsu_box,visualize)
        elif output_type=='csv':
            tracking_frames_to_csv(detector_,frames,numbers,output_path,otsu_box,visualize)
        else:
            print("Invalid output_type argument")
            sys.exit()

    #Tracking ADL videos
    for input_path,output_path in tqdm.tqdm(ADL_list):
        os.makedirs(output_path, exist_ok=True)
        print("------------------")
        print("input_path:",input_path)
        print("------------------")
        frames = glob.glob(input_path+'/*.jpg') + glob.glob(input_path+'/*.png')
        frames,numbers = sort_frames(frames, dset)
        if output_type=='video':
            #tracking_frames_to_video(tracker,detector,frames,output_path,frame_rate=10.0,otsu_box=True,visualize=False)
            tracking_frames_to_video(detector,frames,output_path,frame_rate,otsu_box,visualize)
        elif output_type=='csv':
            tracking_frames_to_csv(detector,frames,numbers,output_path,otsu_box,visualize)
        else:
            print("Invalid output_type argument")
            sys.exit()