import tempfile
from pathlib import Path
import os
from argparse import Namespace
import time
import dlib
import os
import sys
import numpy as np
from PIL import Image
import torch
import torchvision.transforms as transforms
import tensorflow as tf
import numpy as np
import torch
import clip
from PIL import Image
import pickle
import copy
import matplotlib.pyplot as plt
from MapTS import GetFs, GetBoundary, GetDt
from manipulate import Manipulator
from dnnlib import tflib

sys.path.insert(0, "/content")
sys.path.insert(0, "/content/encoder4editing")

from encoder4editing.utils.common import tensor2im
from encoder4editing.utils.alignment import align_face
from encoder4editing.models.psp import pSp

import cog

class Model(cog.Model):
    def setup(self):
        print("starting setup")

        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model, self.preprocess = clip.load(
            "ViT-B/32", device=self.device, jit=False
        )

        self.graph = tf.get_default_graph()
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
        self.sess = tf.Session(
            graph=self.graph, config=tf.ConfigProto(gpu_options=gpu_options)
        )

        experiment_type = "ffhq_encode"

        self.experiment_args = {"model_path": "e4e_ffhq_encode.pt"}
        self.experiment_args["transform"] = transforms.Compose(
            [
                transforms.Resize((256, 256)),
                transforms.ToTensor(),
                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
            ]
        )
        self.resize_dims = (256, 256)

        model_path = self.experiment_args["model_path"]

        ckpt = torch.load(model_path, map_location="cpu")
        opts = ckpt["opts"]
        # pprint.pprint(opts)  # Display full options used
        # update the training options
        opts["checkpoint_path"] = model_path
        opts = Namespace(**opts)

        self.net = pSp(opts)
        self.net.eval()
        self.net.cuda()

        self.shape_predictor = dlib.shape_predictor(
            "/content/shape_predictor_68_face_landmarks.dat"
        )

        with self.graph.as_default(), self.sess.as_default():
            #tflib.init_tf()

            self.M = Manipulator(dataset_name="ffhq", sess=self.sess)
            self.fs3 = np.load("./npy/ffhq/fs3.npy")
            np.set_printoptions(suppress=True)

        print("setup complete")

    @cog.input("input", type=Path, help="Input image")
    @cog.input("neutral", type=str, help="Neutral image description")
    @cog.input("target", type=str, help="Target image description")
    @cog.input(
        "manipulation_strength",
        type=float,
        min=-10,
        max=10,
        default=4.1,
        help="The higher the manipulation strength, the closer the generated image becomes to the target description. Negative values moves the generated image further from the target description",
    )
    @cog.input(
        "disentanglement_threshold",
        type=float,
        min=0.08,
        max=0.3,
        default=0.15,
        help="The higher the disentanglement threshold, the more specific the changes are to the target attribute. Lower values mean that broader changes are made to the input image",
    )
    def predict(
        self,
        input,
        neutral,
        target,
        manipulation_strength,
        disentanglement_threshold,
    ):
        # @title Align image
        original_image = Image.open(str(input))
        original_image = original_image.convert("RGB")
        input_image = self.run_alignment(str(input))
        input_image = original_image
        input_image.resize(self.resize_dims)

        img_transforms = self.experiment_args["transform"]
        transformed_image = img_transforms(input_image)

        with torch.no_grad():
            images, latents = self.run_on_batch(transformed_image.unsqueeze(0))
            result_image, latent = images[0], latents[0]

            print("latents", latents)

        print(transformed_image.shape, result_image.shape)

        w_plus = latents.cpu().detach().numpy()
        with self.graph.as_default(), self.sess.as_default():
            dlatents_loaded = self.M.W2S(w_plus)

        #print("w_plus, dlatents_loaded", w_plus, dlatents_loaded)

        img_index = 0
        w_plus=latents.cpu().detach().numpy()
        with self.graph.as_default(), self.sess.as_default():
            dlatents_loaded=self.M.W2S(w_plus)

        img_indexs=[img_index]
        dlatent_tmp=[tmp[img_indexs] for tmp in dlatents_loaded]
        with self.graph.as_default(), self.sess.as_default():
            self.M.num_images = len(img_indexs)
            self.M.alpha = [0]
            self.M.manipulate_layers = [0]

        with self.graph.as_default(), self.sess.as_default():
            codes, out = self.M.EditOneC(0, dlatent_tmp)

        original = Image.fromarray(out[0, 0]).resize((512, 512))
        with self.graph.as_default(), self.sess.as_default():
            self.M.manipulate_layers = None

        classnames = [target, neutral]
        dt = GetDt(classnames, self.model)

        with self.graph.as_default(), self.sess.as_default():
            self.M.alpha = [manipulation_strength]
            boundary_tmp2, c = GetBoundary(
                self.fs3, dt, self.M, threshold=disentanglement_threshold
            )
            codes = self.M.MSCode(dlatent_tmp, boundary_tmp2)
            out = self.M.GenerateImg(codes)
        generated = Image.fromarray(out[0, 0])  # .resize((512,512))

        out_path = Path(tempfile.mkdtemp()) / "out.jpg"
        generated.save(str(out_path))

        return out_path

    def run_alignment(self, image_path):
        aligned_image = align_face(filepath=image_path, predictor=self.shape_predictor)
        print("Aligned image has shape: {}".format(aligned_image.size))
        return aligned_image

    def run_on_batch(self, inputs):
        images, latents = self.net(
            inputs.to("cuda").float(), randomize_noise=False, return_latents=True
        )
        return images, latents

def concat_images(*images):
    width = 0
    for im in images:
        width += im.width
    height = max([im.height for im in images])
    concat = Image.new("RGB", (width, height))
    offset = 0
    for im in images:
        concat.paste(im, (offset, 0))
        offset += im.width
    return concat


"""
format:

    pointfit [parameters] <datafile> <sourcelist> [<outputfile>]

where:
    <datafile>: Name of a file, including glob-style wild cards, of photon data in FITS
        format or a merit ROOT tuple, or a FITS file with pixel data. If preceded
        by an at-sign, the remainder is the name of an ascii file containing a list of
        files to open
    <sourcelist>: a file, ascii or FITS, containing a list of seed points for fit.
        Expected fields name, ra, dec. There should be a field (name TBD) with TS
        values allowing a ranking of significance or flux. pointfit will use this
        to sort nearby pairs. If not found, default behavior is to not sort.
        [Not currently working, will be fixed]
    <outputfile>: Optional parameter specifying file name to write results, in
        format: name, input_ra, input_dec, fit_ra, fit_dec TS, sigma, pull, [name of background source(s)]
        where sigma is the fit error circle radius in degrees: >90 indicates a convergence error
              pull is the difference between fit and seed position in units of sigma

Optional parameters:
    --exposure=: Either the name of a gtexpcube-generated file, or a value, in cm^2 s,
        to apply to all energies and directions. (3e10 is appropriate for a year.)
        If not specified, use the GTI information from the datafile.
        [This is TODO: Is there a function member to sum the GTI intervals?]
        Note that this is only actually needed for overlapped sources in the Galactic
        plane, if spectral information is not required.
    --diffuse=:  Define a diffuse file (see the flag --galdiffuse to define it in the context of the science tools)
    --galdiffuse: Flag to use the galprop-generated galactic diffuse file which is distributed with the science tools)
    --eventtype= [-1] Event selection if datafile is event data. -1 means front and back,
        0/1 for front/back specifically.
    --write=<output>: if set, and the datafile is event data, write a pixelfile for
        subsequent input
    -v or --verbose [0] set verbosity
    --quiet [False] Disable some output
    --binsperdecade [0] default is 2.35 ratio. otherwise energy binning is set.
    --emin= [200]  minimum energy, used to select bands
    --TSmin= [10]
    --[no]lsq  use (or not) "lsq" mode to find maximum
    --region=: define a region file
    --TSmin [10]  minimum TS

 $Header: /nfs/slac/g/glast/ground/cvs/pointlike/python/pointfit.py,v 1.20 2009/06/04 20:24:39 burnett Exp $
"""
#    --ltcube=: [None] If specified, use for exposure [not implemented]

import os, sys, types
from numpy import arange

from skymaps import DiffuseFunction, Background
from pointlike import SourceList, Source, PointSourceLikelihood, Data
#TODO from pointlike import pixeldata

class PointFit(object):

    def __init__(self):
        pass

#----------------------------------------------------------------------------------------
class Fitter(object):
    """ this is an attempt to be consistent with an older interface that is used by ASP
    """
    def __init__(self, source, data,  background=None, skip=2, verbose=0, localize=True):
        psl = PointSourceLikelihood(data.map(), source.name(), source.dir())
        if background is not None: PointSourceLikelihood.set_diffuse(background)
        psl.set_verbose(verbose)
        self.TS = psl.TS()
        self.name = source.name()
        if localize:
            self.sigma = psl.localize(skip)
            self.delta = math.degrees(psl.dir().difference(source.dir()))
        else:
            self.sigma = self.delta=-1
        self.ra = psl.dir().ra()
        self.dec= psl.dir().dec()
        self.sdir = psl.dir()
        self.like = psl # access to the likelihood functions

#----------------------------------------------------------------------------------------

def photonmap(filename, eventtype=-1, pixeloutput=False, tstart=0, tstop=0, ignorepix=False):
    """ return a Data object, determined one of 3 ways:
        * the name of a file containing a list of photon files, preceded by an @
        * the name of a photon file, that will be expanded by glob
        * the name of a photon map file.
        if pixeloutput is set, write the photonmap to the file
        eventtype is -1 for all events, 0/1 for front back
    """
    data = None
    if filename[0]=='@':
        # it is a list of data files
        filelist = [line.strip() for line in file(filename[1:]) if len(line)>0 and line[0]!='#']
        #print (filelist)
        data =  Data(filelist, eventtype, tstart, tstop)
    elif type(filename) is types.ListType:
        # it is a (Python) list of data files
        data = Data(filename, eventtype, tstart, tstop)
    elif filename[-5:]=='.fits' or filename[-4:]=='.fit' :
        # a fits file: either data to read, or a photonmap
        import astropy.io.fits as pyfits
        import glob
        files = glob.glob(filename)
        if len(files)==0:
            raise Exception('no such file(s): %s' %filename)
        #Check to see if there is a pixel file
        # compact way to get the first file name
        pixelfilename = os.path.splitext(os.path.split(files[0])[1])[0]+'_bands.fits'

        if os.path.exists(pixelfilename) and not ignorepix:
            files[0] = pixelfilename
        hd = pyfits.open(files[0])
        if len(hd)==1:
            raise Exception('Invalid data file, apparent image file with primary only')
        if hd[1].name=='PHOTONMAP' or hd[1].name=='BANDS':
            hd.close()
            if len(files)>1: print ('Warning: more than one photonmap file not supported')
            data = Data(files[0], hd[1].name)
        else:
            from types import ListType
            if type(files) is not ListType: files = [files]
            hd.close()
            data = Data(files, eventtype, tstart , tstop)
    else:
        raise Exception('filename %s not a valid list of files or fits file' % filename)
    if pixeloutput:
        data.map().write(pixelfilename)
        print ('created a BinnedPhotonData file: %s' % pixelfilename)
    return data
#--------------------------------------------------------

def set_diffuse(diffusefilename='galdiffuse', exposure=1e10):
    if 'GLAST_EXT' in os.environ and diffusefilename=='galdiffuse':
        diffusefilename = os.path.join(os.environ['GLAST_EXT'],'extFiles','v0r7','galdiffuse', 'GP_gamma.fits')
    diffuse = DiffuseFunction(diffusefilename)
    background = Background(diffuse, exposure)
    PointSourceLikelihood.set_diffuse(background)
    return (diffuse, background) # return to keep references

#--------------------------------------------------------

def main():

    def help(msg=None):
        if msg: print ('\nError: %s' % msg)
        print (__doc__)
        sys.exit(0)

    from getopt import getopt, GetoptError
    import sys

    options = 'b:w:v'
    long_options= [ 'diffuse=','write=', 'verbose', 'galdiffuse',
                    'eventtype=', 'exposure=', 'binsperdecade=', 'emin=', 'region=',
                    'TSmin=', 'lsq', 'nolsq', 'quiet', 'history=']

    try:
        (opts, args) = getopt(sys.argv[1:], options, long_options )
    except GetoptError as msg:
        help(msg)

    outputpixelfile= background=regfile=region=None
    diffusefilename='galdiffuse' # wire in for now
    verbose=0
    exposure=3e10 # this is appropriate for 1 year.
    eventtype=-1  # all events
    binsperdecade=0 # default binning
    emin = 500
    TSmin= 10  # minimum TS for final list
    lsq = False
    quiet = False
    history = None

    for (opt,val) in opts:
        if   opt=='-b' or opt=='--diffuse'  : diffusefilename = val
        elif opt=='-w' or opt=='--write'    : outputpixelfile = val
        elif opt=='-v' or opt=='--verbose'  : verbose =1
        elif opt=='--galdiffuse'            : diffusefilename='galdiffuse' #flag
        elif opt=='--eventtype'             : eventtype= int(val)
        elif opt=='--emin'                  : emin = float(val)
        elif opt=='--binsperdecade'         : binsperdecade=float(val)
        elif opt=='--region'                : region = val
        elif opt=='--TSmin'                 : TSmin = float(val)
        elif opt=='--lsq'                   : lsq = True
        elif opt=='--nolsq'                 : lsq = False
        elif opt=='--quiet'                 : quiet = True
        elif opt=='--exposure'              :
            try: exposure= float(val)
            except: exposure = val
        elif opt=='--history'               : history = val

    if len(args)>0: eventfilename=args[0]
    else: help('No event file name specified')
    if binsperdecade>0:
        bins = 10**arange(2,5,1./binsperdecade)
        Data.setEnergyBins(bins)

    if history is None:
        data = photonmap(eventfilename, pixeloutput=outputpixelfile, eventtype=eventtype)
        if not quiet: data.info()
    else:
        pass #TODO pdata = pixeldata.PixelData()

    if len(args)>1: sourcefilename= args[1]
    else: help('No source file list')

    if diffusefilename is not None:
        print ('setting up background from file %s' % diffusefilename)
        (t1,t2)=set_diffuse(diffusefilename, exposure)
    PointSourceLikelihood.set_energy_range(emin)
    PointSourceLikelihood.set_verbose(verbose)
    #PointSourceLikelihood.set_fitlsq(lsq)
    SourceList.set_uselsq(lsq)
    SourceList.set_data(data.map())
    sourcelist = SourceList(sourcefilename)
    sourcelist.sort_TS()
    sourcelist.refit()
    if TSmin>0:  sourcelist.filter_TS(TSmin)
    sourcelist.sort_ra()
    sourcelist.dump()
    if len(args)>2:
        sourcelist.dump(args[2])
    if region is not None:
       sourcelist.createRegFile(region, 'red', TSmin) # todo: color?

#--------------------------------------------------------

if __name__=='__main__':
    main()


import os,sys
import re
import json
import urllib.request
import zipfile
import hashlib
import subprocess
import uuid
import time
import zipfile

from PIL import Image
import matplotlib.pyplot as plt

from sklearn.metrics import confusion_matrix
import seaborn as sns

os.environ['CUDA_VISIBLE_DEVICES'] = sys.argv[1]
os.environ["TF_KERAS"] = "1"

from vis.utils import utils
from vis.visualization import visualize_cam, overlay

import tensorflow as tf
import numpy as np
import skimage

from tensorflow.python import keras

from tensorflow.contrib.data import assert_element_shape
from tensorflow.python.keras import backend as K, Sequential
from tensorflow.python.keras.layers.core import Dense, Activation
from tensorflow.python.keras.layers import ZeroPadding2D, GlobalAveragePooling2D
from tensorflow.python.keras.optimizers import Adam
from tensorflow.python.keras.metrics import categorical_crossentropy
from tensorflow.python.keras.preprocessing.image import ImageDataGenerator
from tensorflow.python.keras.preprocessing import image
from tensorflow.python.keras.models import Model
from tensorflow.python.keras.applications import imagenet_utils
from tensorflow.python.keras.layers import Dense, Flatten, Dropout
from tensorflow.python.keras.applications import MobileNet
from tensorflow.python.keras.applications.mobilenet import preprocess_input
from tensorflow.python.keras.optimizers import Adam
from tensorflow.python.keras.utils import multi_gpu_model
from tensorflow.python.keras.models import load_model
from tensorflow.python.keras.preprocessing.image import ImageDataGenerator
from tensorflow.python.keras import backend as K
from tensorflow.python.keras import utils
from tensorflow.python.keras.datasets.cifar import load_batch
from tensorflow.python.keras.utils.data_utils import get_file
from tensorflow.python.keras.models import load_model

from tensorflow.python.keras import optimizers
from tensorflow.python.keras.utils import multi_gpu_model
from tensorflow.python.keras import callbacks
from tensorflow.python.keras.callbacks import History, TensorBoard, LearningRateScheduler

from keras_radam import RAdam

import oss2

import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.image import MIMEImage

from tools.utils import Helper, create_loss_fn, INFO, ERROR, NOTE
from tools.custom import Yolo_Precision, Yolo_Recall
from models.yolonet import *

auth = oss2.Auth('MEOWMEOWMEOW', 'MEOWMEOWMEOW')
bucket = oss2.Bucket(auth, 'MEOWMEOWMEOW', 'MEOWMEOWMEOW')

upstreamServerAddress = "MEOWMEOWMEOW"
upstreamMagic = "MEOWMEOWMEOW"

aliyunOSSAddress = 'MEOWMEOWMEOW'

localSSDLoc = ''
nncaseLoc = f'MEOWMEOWMEOW'

senderAddr = "MEOWMEOWMEOW"
senderPass = "MEOWMEOWMEOW"

bootFileName = "MEOWMEOWMEOW"

detectionBootFileName = "MEOWMEOWMEOW"
detectionBinFile = "MEOWMEOWMEOW"
flashListFileName = 'MEOWMEOWMEOW'

completedTaskHeader = ""

def loss_softmax_cross_entropy_with_logits_v2(correct, predicted):
        return tf.nn.softmax_cross_entropy_with_logits_v2(labels=correct,logits=predicted)

def getDataset(address, uuid):
    try:
        urllib.request.urlretrieve(address, f"{localSSDLoc}dataset/{uuid}_dataset.zip")
        md5 = hashlib.md5(open(f"{localSSDLoc}dataset/{uuid}_dataset.zip", 'rb').read()).hexdigest()
        with zipfile.ZipFile(f"{localSSDLoc}dataset/{uuid}_dataset.zip", 'r') as zip_ref:
            zip_ref.extractall(f"{localSSDLoc}dataset_tmp/{uuid}_dataset")

        return(0, f"{localSSDLoc}dataset_tmp/{uuid}_dataset", md5)
    except Exception as e:
        return(-1, f"Failed to fetch the dataset, dut to {e}", 'NaN')

def chkWaterMark(filename, magic):
    m = hashlib.sha256()
    m.update(str.encode(f"{magic} is super cool!"))

    try:
        fsize = os.path.getsize(filename)
        f = open(filename, 'rb')
        f.seek(fsize - 32)

        wtmData = f.read()
        f.close()
    except Exception as e:
        return(-1, f"Failed to fetch the dataset, dut to {e}")

    if wtmData == m.digest():
        return(0, wtmData)
    else:
        return(-7, "Sorry, We currently only serivce M5StickV Customers.")

def chkDataset(dirname):

    listDatasetDir = os.listdir(dirname)

    if "v-training.config" in listDatasetDir:
        print('Found config file, using detection mode.')
        try:
            vdetConfig = json.load(open(os.path.join(dirname, "v-training.config")))
            numOfClass = vdetConfig['classes']
            #key = vdetConfig['key']
        except Exception as e:
            return(-34, 'Failed to read config file, please make sure you have the correct format. Err:', e)

        numOfDataset = 0
        fileHash = []
        for fileName in listDatasetDir:
            if '.jpg' in fileName and (fileName.replace('.jpg', '.txt')) not in listDatasetDir:
                return(-35, f"Unmatched pair found in dataset. Found {fileName}, but {fileName.replace('.jpg', '.txt')} is not exist.")

            if '.jpg' in fileName:
                fileMD5 = hashlib.md5(open(os.path.join(dirname, fileName), 'rb').read()).hexdigest()
                if fileMD5 in fileHash:
                    return(-36, f"Repeated image {fileName} found in dataset. You should remove the repeated image for ensuring the quaility of model.")

                try:
                    im = Image.open(os.path.join(dirname, fileName))
                    #if(im.size != (320, 240)):
                    #    return(-2, f"Unexpected Image Size, only accpet 320x240 from M5Stick series products, but you have {im.size}")
                    im.close()
                except:
                    return(-37, f"Failed to open image {fileName}, plase make sure the images are not corrupted.")

                try:
                    labelFile = open(os.path.join(dirname, fileName.replace('.jpg', '.txt')), 'r')
                    for labelLine in labelFile:
                        labels = labelLine.split(' ')

                        if int(labels[0]) >= numOfClass:
                            return(-38, f"Label index overflowed, in {labelFile}, you claimed {numOfClass}, but it index {labels[0]}.")

                        if float(labels[1]) > 1.0 or float(labels[1]) < 0.0 or float(labels[2]) > 1.0 or float(labels[2]) < 0.0:
                            return(-39, f"Label bbox overflowed, the x,y,w,h should be in reletive mode and in the range of [0, 1].")

                    labelFile.close()
                except Exception as e:
                    return(-40, f"Failed to open label file {fileName.replace('.jpg', '.txt')}, plase make sure the label files are not corrupted. Err:", e)

                numOfDataset = numOfDataset + 1

        if numOfDataset < 30:
            return(-37, f"We require you upload more than 30 labled image for each training requres, but you only have {numOfDataset}")

        return (0, numOfClass, 'detection', vdetConfig)

    else:
        print('Config file not found, using classification mode.')

        if "train" not in listDatasetDir or ("vaild" not in listDatasetDir and "valid" not in listDatasetDir):
            return(-8, "Train or Valid folder not found. If you are using the M5StickV software, make sure you reach enough image counts of 35 per class.")

        if "valid" in listDatasetDir:
            os.rename(f"{dirname}/valid", f"{dirname}/vaild")

        TrainImageNum = 0
        VaildImageNum = 0

        NumOfClass = len(os.listdir(os.path.join(dirname, "train")))
        NumOfClassVaild = len(os.listdir(os.path.join(dirname, "vaild")))

        if NumOfClass != NumOfClassVaild:
            return(-11, "Number of Classes presented in Train and Vaild dataset is not equal.")

        print("[INFO]", f"Total {NumOfClass} Classes Found")

        if NumOfClass < 2:
            return(-16, "Number of Classes should larger or equal to two.")

        for file in listDatasetDir:
            if(os.path.splitext(file)[1] == '.dpf'):
                isConfigFilePresent = 1
                with open(os.path.join(dirname, file), 'r') as f:
                    datasetCfg = json.load(f)
        try:
            for dirpath, _ , filenames in os.walk(f"{dirname}/train", topdown=False):
                for file in filenames:
                    if(os.path.splitext(file)[1] == '.jpg'):
                        im=Image.open(os.path.join(dirpath, file))
                        #if(im.size != (320, 240)):
                        #    return(-2, f"Unexpected Image Size, only accpet 320x240, you have {im.size}")
                        im.close()
                        TrainImageNum = TrainImageNum + 1
                    #else:
                    #   return(-3, f"Unexpected File Present: {file}")

            for dirpath, _ , filenames in os.walk(f"{dirname}/vaild", topdown=False):
                for file in filenames:
                    if(os.path.splitext(file)[1] == '.jpg'):
                        im=Image.open(os.path.join(dirpath, file))
                        #if(im.size != (320, 240)):
                        #    return(-2, f"Unexpected Image Size, only accpet 320x240, you have {im.size}")
                        im.close()
                        VaildImageNum = VaildImageNum + 1
                    #else:
                    #   return(-3, f"Unexpected File Present: {file}")
        except Exception as e:
            return(-6, f"Unexpected error happened during checking dataset, {e}")

        if TrainImageNum < NumOfClass * 30:
            return(-4, f"Lake of Enough Train Dataset, Only {TrainImageNum} pictures found, but you need {NumOfClass * 30} in total.")

        if VaildImageNum < NumOfClass * 5:
            return(-4, f"Lake of Enough Valid Dataset, Only {VaildImageNum} pictures found, but you need {NumOfClass * 5} in total.")

        #if isConfigFilePresent == 0: #Not enable isConfigFilePresent Yet
        #    return(-5, "Unable to find dpf dataset description file.")

        return (0, NumOfClass, 'classification')

def runTrainingDetection(uuid, datasetDir, numOfClass, obj_thresh = 0.7, iou_thresh = 0.5, obj_weight = 1.0, noobj_weight = 1.0, wh_weight = 1.0, max_nrof_epochs = 50, batch_size = 96, vaildation_split = 0.2):
    config = tf.ConfigProto()
    sess = tf.Session(config=config)
    keras.backend.set_session(sess)

    datasetList = [os.path.join(datasetDir, f) for f in os.listdir(datasetDir)]
    image_list = []

    #img_ann_f = open(os.path.join(datasetDir, 'dataset_img_ann.txt'), 'w+')

    for fileName in datasetList:
        if '.jpg' in fileName:
    #        print('/home/m5stack/VTrainingService/' + fileName, file=img_ann_f)
            image_list.append('/home/m5stack/VTrainingService/' + fileName)

    #img_ann_f.close()

    image_path_list = np.array(image_list)#np.loadtxt(os.path.join(datasetDir, 'dataset_img_ann.txt'), dtype=str)

    ann_list = list(image_path_list)
    ann_list = [re.sub(r'JPEGImages', 'labels', s) for s in ann_list]
    ann_list = [re.sub(r'.jpg', '.txt', s) for s in ann_list]

    lines = np.array([
        np.array([
            image_path_list[i],
            np.loadtxt(ann_list[i], dtype=float, ndmin=2),
            np.array(skimage.io.imread(image_path_list[i]).shape[0:2])]
        ) for i in range(len(ann_list))])

    np.save(os.path.join(datasetDir, 'dataset_img_ann.npy'), lines)

    #print('dataset npu>>>', os.path.join(datasetDir, 'dataset_img_ann.npy'))

    h = Helper(os.path.join(datasetDir, 'dataset_img_ann.npy'), numOfClass, 'voc_anchor.npy',
        np.reshape(np.array((224, 320)), (-1, 2)), np.reshape(np.array((7, 10, 14, 20)), (-1, 2)), vaildation_split)

    h.set_dataset(batch_size, 6)

    network = eval('yolo_mobilev1')  # type :yolo_mobilev2
    yolo_model, train_model = network([224, 320, 3], len(h.anchors[0]), numOfClass, alpha=0.5)

    train_model.compile(
        RAdam(),
        loss=[create_loss_fn(h, obj_thresh, iou_thresh, obj_weight, noobj_weight, wh_weight, layer)
            for layer in range(len(train_model.output) if isinstance(train_model.output, list) else 1)],
        metrics=[Yolo_Precision(obj_thresh, name='p'), Yolo_Recall(obj_thresh, name='r')])

    shapes = (train_model.input.shape, tuple(h.output_shapes))
    h.train_dataset = h.train_dataset.apply(assert_element_shape(shapes))
    h.test_dataset = h.test_dataset.apply(assert_element_shape(shapes))

    #print('train', h.train_dataset, '\n\r\n\rtest', h.test_dataset)

    try:
        train_model.fit(h.train_dataset, epochs=max_nrof_epochs,
                    steps_per_epoch=10,
                    validation_data=h.test_dataset, validation_steps=1)
    except Exception as e:
        return(-45, 'Unexpected error found during training, err:', e)

    keras.models.save_model(yolo_model, f'{localSSDLoc}trained_h5_file/{uuid}_mbnet5_yolov3.h5')

    converter = tf.lite.TFLiteConverter.from_keras_model_file(f'{localSSDLoc}trained_h5_file/{uuid}_mbnet5_yolov3.h5', custom_objects={'RAdam': RAdam, 'loss_softmax_cross_entropy_with_logits_v2':loss_softmax_cross_entropy_with_logits_v2})
    tflite_model = converter.convert()
    open(f'{localSSDLoc}trained_tflite_file/{uuid}_mbnet5_yolov3_quant.tflite', "wb").write(tflite_model)

    subprocess.run([f'{nncaseLoc}/ncc', f'{localSSDLoc}trained_tflite_file/{uuid}_mbnet5_yolov3_quant.tflite', f'{localSSDLoc}trained_kmodel_file/{uuid}_mbnet5_yolov3.kmodel', '-i', 'tflite', '-o', 'k210model', '--dataset', datasetDir])

    if os.path.isfile(f'{localSSDLoc}trained_kmodel_file/{uuid}_mbnet5_yolov3.kmodel'):
        return (0, f'{localSSDLoc}trained_kmodel_file/{uuid}_mbnet5_yolov3.kmodel')
    else:
        return (-16, 'Unexpected Error Found During generating Kendryte k210model.')

def runTrainingClassification(uuid, datasetDir, validDir, classNum, dropoutValue = 0.2, batch_size = 128, nb_epoch = 20, step_size_train = 10, alphaVal = 0.75, depthMul = 1):

    imageGen = ImageDataGenerator(
        rotation_range=30,
        width_shift_range=0.35,
        height_shift_range=0.35,
        zoom_range=0.35,
        shear_range=0.35,
        vertical_flip=False,
        horizontal_flip=False,
        brightness_range = [0.65, 1.35],
        rescale=1. / 255)

    trainSet=imageGen.flow_from_directory(datasetDir,
                                                    target_size=(224,224),
                                                    color_mode='rgb',
                                                    batch_size=batch_size,
                                                    class_mode='categorical', shuffle=True)
    validSet=imageGen.flow_from_directory(validDir,
                                                    target_size=(224,224),
                                                    color_mode='rgb',
                                                    batch_size=32,
                                                    class_mode='categorical', shuffle=True)

    class EarlyStoppingAtMinLoss(tf.keras.callbacks.Callback):
        def __init__(self, patience=3):
            super(EarlyStoppingAtMinLoss, self).__init__()
            self.patience = patience
            self.best_weights = None

        def on_train_begin(self, logs=None):
            self.wait = 0
            self.stopped_epoch = 0
            self.best = np.Inf
            self.last_acc = 0
            self.atleastepoc = 0

        def on_epoch_end(self, epoch, logs=None):
            current = logs.get('val_loss')
            val_acc = logs.get('val_acc')
            self.atleastepoc = self.atleastepoc + 1
            if np.less(current, self.best) or self.last_acc < 0.95 or self.atleastepoc < 25:
                self.best = current
                self.wait = 0
                self.last_acc = val_acc
                self.best_weights = self.model.get_weights()
            else:
                self.wait += 1
                if self.wait >= self.patience:
                    self.stopped_epoch = epoch
                    self.model.stop_training = True
                    print('\nRestoring model weights from the end of the best epoch.')
                    self.model.set_weights(self.best_weights)

        def on_train_end(self, logs=None):
            if self.stopped_epoch > 0:
                print('Epoch %05d: early stopping' % (self.stopped_epoch + 1))

    base_model = tf.keras.applications.MobileNet(input_shape=(224, 224, 3), alpha = alphaVal,depth_multiplier = depthMul, dropout = dropoutValue, pooling='avg',include_top = False, weights = "imagenet", classes = classNum)

    mbnetModel = Sequential([
        base_model,
        Dropout(dropoutValue, name='dropout'),
        Dense(classNum, activation='softmax')
    ])

    if classNum == 2:
        mbnetModel.compile(loss='binary_crossentropy',
                optimizer=RAdam(),
                metrics=['accuracy'])
    else:
        mbnetModel.compile(loss='categorical_crossentropy',#loss_softmax_cross_entropy_with_logits_v2,
                optimizer=RAdam(),
                metrics=['accuracy'])

    history = History()

    try:
        mbnetModel.fit_generator(generator=trainSet,steps_per_epoch=step_size_train,callbacks=[EarlyStoppingAtMinLoss(), history],epochs=50, validation_data=validSet)
    except Exception as e:
        return(-14, f'Unexpected Error Found During Triaining, {e}')

    mbnetModel.save(f'{localSSDLoc}trained_h5_file/{uuid}_mbnet10.h5')

    converter = tf.lite.TFLiteConverter.from_keras_model_file(f'{localSSDLoc}trained_h5_file/{uuid}_mbnet10.h5', custom_objects={'RAdam': RAdam, 'loss_softmax_cross_entropy_with_logits_v2':loss_softmax_cross_entropy_with_logits_v2})
    tflite_model = converter.convert()
    open(f'{localSSDLoc}trained_tflite_file/{uuid}_mbnet10_quant.tflite', "wb").write(tflite_model)

    subprocess.run([f'{nncaseLoc}/ncc', f'{localSSDLoc}trained_tflite_file/{uuid}_mbnet10_quant.tflite', f'{localSSDLoc}trained_kmodel_file/{uuid}_mbnet10_quant.kmodel', '-i', 'tflite', '-o', 'k210model', '--dataset', validDir])

    if os.path.isfile(f'{localSSDLoc}trained_kmodel_file/{uuid}_mbnet10_quant.kmodel'):
        return (0, f'{localSSDLoc}trained_kmodel_file/{uuid}_mbnet10_quant.kmodel', history, validSet, mbnetModel)
    else:
        return (-16, 'Unexpected Error Found During generating Kendryte k210model.')

def packImages(uuid, history, validSet, mbnetModel):
    acc = history.history['acc']
    val_acc = history.history['val_acc']
    loss = history.history['loss']
    val_loss = history.history['val_loss']

    epochs = range(len(acc))

    plt.figure()
    plt.ylim(0, 1)
    plt.plot(epochs, acc, 'bo', label='Training acc')
    plt.plot(epochs, val_acc, 'b', label='Validation acc')
    plt.autoscale()
    plt.title('Training and validation accuracy')
    plt.legend()

    plt.savefig(f'{localSSDLoc}imageAssest/{uuid}_acc_graph.png')
    plt.figure()
    plt.plot(epochs, loss, 'bo', label='Training loss')
    plt.plot(epochs, val_loss, 'b', label='Validation loss')
    plt.autoscale()
    plt.title('Training and validation loss')
    plt.legend()
    plt.savefig(f'{localSSDLoc}imageAssest/{uuid}_loss_graph.png')

def uploadFile(filename):
    try:
        with open(filename, 'rb') as fileobj:
            bucket.put_object(os.path.basename(filename), fileobj)
    except:
        return(-13, "Fail to upload file to server.")
    return (0, f'{aliyunOSSAddress}/{os.path.basename(filename)}')

def sendErrorEmail(mailaddr, userid, content):
    msg_content = 'Hi!\n\rSorry for that, there is some error happened during the training process. '\
        '\nWe attached the reason below, if you have any questions, welcome to report to us through Twitter, Facebook, or Forum.\n'\
        'Or, check out our docs here: https://docs.m5stack.com/#/en/related_documents/v-training'\
        f'\n\rUSERID: {userid}\nCONTENT: {content}'\

    message = MIMEText(msg_content, 'plain')

    message['From'] = f'V-Trainer <{senderAddr}>'
    message['To'] = f'{mailaddr} <{mailaddr}>'
    message['Subject'] = f'[V-Trainer] {userid} Online Training Request Failed'

    msg_full = message.as_string()

    server = smtplib.SMTP_SSL()
    #server.set_debuglevel(1)
    server.connect(host='smtp.qiye.aliyun.com',port=465)
    server.login(senderAddr, senderPass)
    server.sendmail(senderAddr,
                    [mailaddr],
                    msg_full)
    server.quit()

def sendSuccessEmail(mailaddr, userid, ossAddr, typeOfRunning):
    msgRoot = MIMEMultipart('related')
    msgRoot['From'] = f'V-Trainer <{senderAddr}>'
    msgRoot['To'] = f'{mailaddr} <{mailaddr}>'
    msgRoot['Subject'] = f'[V-Trainer] {userid} Online Training Request Finished'

    msgAlternative = MIMEMultipart('alternative')
    msgRoot.attach(msgAlternative)

    if typeOfRunning == 'classification':
        msg_content = f'Hi!\n\nYour training request have been successfully processed, you can download the kmodel & sample program files here: \n{ossAddr}\n\r\n\r'\
                    "Model: Classification MobileNetV1 Alpha: 0.7 Depth: 1"

        msgText = MIMEText(msg_content, 'plain')
        msgAlternative.attach(msgText)
    # '''
        msgText = MIMEText(msg_content.replace("\n", "<br>") + '<br><img src="cid:acc_graph"><br><img src="cid:loss_graph"><br><img src="cid:conf_matrix"><br>', 'html')
        msgAlternative.attach(msgText)

        with open(f'{localSSDLoc}imageAssest/{userid}_loss_graph.png', 'rb') as imageFile:
            msgImage = MIMEImage(imageFile.read())
            msgImage.add_header('Content-ID', '<loss_graph>')
            msgRoot.attach(msgImage)

        with open(f'{localSSDLoc}imageAssest/{userid}_acc_graph.png', 'rb') as imageFile:
            msgImage = MIMEImage(imageFile.read())
            msgImage.add_header('Content-ID', '<acc_graph>')
            msgRoot.attach(msgImage)
    else:
        msg_content = f'Hi!\n\nYour training request have been successfully processed, you can download the kmodel & sample program files here: \n{ossAddr}\n\r\n\r'\
                        "Model: Detection YoloV3 Backbone: MobileNetV1 Alpha: 0.5 Depth: 1\n\r\n\r"\
                        "In order to run the model without crash, you need to download the attached flash to your M5 devices.\n\r"\
                        "(To be awared, although every time the attached maixpy.bin should be the same, we may update it without notificaiton.)\n\r\n\r"\
                        "Have fun with the new v-training detectio :p\n\r"\
                        "if you have any questions, welcome to report to us through Twitter, Facebook, or Forum.\n\r"

        msgText = MIMEText(msg_content, 'plain')
        msgAlternative.attach(msgText)

    server = smtplib.SMTP_SSL()
    server.set_debuglevel(1)
    server.connect(host='smtp.qiye.aliyun.com',port=465)
    server.login(senderAddr, senderPass)
    server.sendmail(senderAddr,
                    [mailaddr],
                    msgRoot.as_string())
    server.quit()

bootFileContent = open(bootFileName, "r").read()
blackListEmail = []
fileMD5Hash = {}

while 1:
    try:

    #print("===================REV REQUEST=====================")

        try:
            if completedTaskHeader != "":
                print(">>> Req With TaskHeader: " + completedTaskHeader)
            req = urllib.request.Request(upstreamServerAddress)
            req.add_header('Basic', upstreamMagic)
            req.add_header('X-Complete-Task', completedTaskHeader)
            r = urllib.request.urlopen(req).read()

            completedTaskHeader = ""

            jsonRequest = json.loads(r.decode('utf-8'))

            userid = jsonRequest['id']

        except:
            print("[FATAL]", "Failed to fetch request from server")
            continue

        if jsonRequest['id'] == "":
            #print("[INFO]", "No pending works...")
            time.sleep(1)
        else:
            print("===================REV REQUEST=====================")
            print("[INFO]", "Get work! Json String:", r)

            if jsonRequest['email'] in blackListEmail:
                sendErrorEmail(jsonRequest['email'], jsonRequest['id'], 'Service not available.')
                completedTaskHeader = str(userid)
                continue

            print("[INFO]", "Start Downloading Dataset, Address: ", jsonRequest['url'])
            ret = getDataset(jsonRequest['url'], jsonRequest['id'])
            datasetDir = ret[1]

            md5 = ret[2]

            if md5 in fileMD5Hash and ret[0] == 0:
                fileMD5Hash[md5] = fileMD5Hash[md5] + 1
                if fileMD5Hash[md5] == 3:
                    sendErrorEmail(jsonRequest['email'], jsonRequest['id'], 'Dataset has been uploaded too many times.')
                    completedTaskHeader = str(userid)
                    continue
            else:
                fileMD5Hash[md5] = 0

            if ret[0] != 0:
                print("[FATAL]", ret[1])
                sendErrorEmail(jsonRequest['email'], jsonRequest['id'], ret[1])
                completedTaskHeader = str(userid)
                continue

            print("[INFO]", "Start Checking Dataset, dirname: ", ret[1])
            ret = chkDataset(ret[1])
            if ret[0] != 0:
                print("[FATAL]", ret[1])
                sendErrorEmail(jsonRequest['email'], jsonRequest['id'], ret[1])
                completedTaskHeader = str(userid)
                continue

            numOfClass = ret[1]
            typeOfRunning = ret[2]

            print("[INFO]", f"Start {typeOfRunning} Training")
            if typeOfRunning == 'classification':
                ret = runTrainingClassification(jsonRequest['id'], f"{datasetDir}/train", f"{datasetDir}/vaild", numOfClass, 1/numOfClass)
                if ret[0] != 0:
                    print("[FATAL]", ret[1])
                    sendErrorEmail(jsonRequest['email'], jsonRequest['id'], ret[1])
                    completedTaskHeader = str(userid)
                    continue

                try:
                    packImages(jsonRequest['id'], ret[2], ret[3], ret[4])
                except Exception as e:
                    print("[FATAL]", "Unable to generate Images," + str(e))
                    sendErrorEmail(jsonRequest['email'], jsonRequest['id'], "Unable to generate Images," + str(e))
                    completedTaskHeader = str(userid)
                    continue

                customBootFile = open(bootFileName, 'r').read()

            else:
                vdetConfig = ret[3]

                vdetDict = ['obj_thresh', 'iou_thresh', 'obj_weight', 'noobj_weight', 'wh_weight', 'max_nrof_epochs', 'batch_size', 'vaildation_split']
                vdetDebugDict = {}
                for vdetKey in vdetDict:
                    if vdetKey in vdetConfig:
                        print(f'[INFO] Found override parameter {vdetKey} = {vdetConfig[vdetKey]}')
                        vdetDebugDict[vdetKey] = vdetConfig[vdetKey]

                ret = runTrainingDetection(jsonRequest['id'], datasetDir, numOfClass, **vdetDebugDict)
                if ret[0] != 0:
                    print("[FATAL]", ret[1])
                    sendErrorEmail(jsonRequest['email'], jsonRequest['id'], ret[1])
                    completedTaskHeader = str(userid)
                    continue

                customBootFile = open(detectionBootFileName, 'r').read()

            print("[INFO]", "Start Creating Custom Bootfile.")
            customBootFile = customBootFile.replace("{MODEL_NAME}", f"{os.path.basename(ret[1])}")
            labelsStr = ""

            for i in range(1, numOfClass):
                labelsStr = labelsStr + f"\"{i}\","
            labelsStr = labelsStr + f"\"{numOfClass}\""

            customBootFile = customBootFile.replace("{MODEL_LABELS}", labelsStr)

            os.mkdir(f"{localSSDLoc}custom_bootfile/{jsonRequest['id']}")
            open(f"{localSSDLoc}custom_bootfile/{jsonRequest['id']}/boot.py", "w+").write(customBootFile)

            if typeOfRunning == 'detection':
                customFlashListFile = open(flashListFileName, 'r').read()
                customFlashListFile = customFlashListFile.replace('{MODEL_NAME}', os.path.basename(ret[1]))
                open(f"{localSSDLoc}custom_bootfile/{jsonRequest['id']}/flash-list.json", "w+").write(customFlashListFile)

            userid = jsonRequest['id']

            vtrainer_output = f'{localSSDLoc}output_file/{str(uuid.uuid4())}_{userid}_vtrainer.zip'

            print("[INFO]", f"Start Compressing All Files, \nzipdir: {vtrainer_output}\n"\
                f"bootfile Dir: {localSSDLoc}custom_bootfile/{jsonRequest['id']}/boot.py"\
                f"kmodel Dir: {ret[1]}")

            with zipfile.ZipFile(vtrainer_output, 'w') as myzip:
                myzip.write(f"{localSSDLoc}custom_bootfile/{jsonRequest['id']}/boot.py", "boot.py")

                if typeOfRunning == 'detection':
                    vtrainer_kfpkg = f'{localSSDLoc}output_file/{str(uuid.uuid4())}_{userid}_vtrainer.kfpkg'
                    with zipfile.ZipFile(vtrainer_kfpkg, 'w') as mykfpkg:
                        mykfpkg.write(f"{localSSDLoc}custom_bootfile/{jsonRequest['id']}/flash-list.json", "flash-list.json")
                        mykfpkg.write(ret[1], os.path.basename(ret[1]))
                        mykfpkg.write(detectionBinFile, "maixpy_with_yolov3_patch.bin")

                    myzip.write(vtrainer_kfpkg, f"{userid}_vtrainer.kfpkg")
                else:
                    myzip.write(f"{localSSDLoc}startup.jpg", "startup.jpg")
                    myzip.write(ret[1], os.path.basename(ret[1]))

            print("[INFO]", "Start Uploading File, filename:", ret[1])
            ret = uploadFile(vtrainer_output)
            if ret[0] != 0:
                print("[FATAL]", ret[1])
                sendErrorEmail(jsonRequest['email'], jsonRequest['id'], ret[1])
                completedTaskHeader = str(userid)
                continue

            print("[INFO]", "Start Sending Success Email, address:", ret[1])
            sendSuccessEmail(jsonRequest['email'], jsonRequest['id'], ret[1], typeOfRunning)
            completedTaskHeader = str(userid)

    except Exception as e:
        completedTaskHeader = str(userid)
        print("[FATAL] Unexpcted error happened during training. Err:", e)
        try:
            sendErrorEmail(jsonRequest['email'], jsonRequest['id'], 'Unexpcted error happened during training. Err: ' + str(e))
        except:
            pass


#!/usr/bin/env python3

import math
from itertools import chain
from collections.abc import MutableMapping
from copy import deepcopy

import numpy as np
import torch

from espnet.nets.pytorch_backend.transformer.mask import target_mask
import lattice
from utils import sym2idx

class Cache(dict):
    """A customised dictionary as the cache for lattice rescoring.
    Input list of string is first concatenated and then loop up.
    The key value pair is ngram: (state, pred, post).
    """
    __slots__ = 'model', 'dict', 'type', 'sp', 'h', 'cache_locality', 'acronyms'

    @staticmethod
    def _process_args(mapping=(), **kwargs):
        if hasattr(mapping, 'items'):
            mapping = getattr(mapping, 'items')()
        return ((' '.join(k), v)
                for k, v in chain(mapping, getattr(kwargs, 'items')()))

    def __init__(self, model, dictionary, feat=None, model_type='rnnlm',
                 sp=None, cache_locality=9, acronyms={}, mapping=(), **kwargs):
        super(Cache, self).__init__(self._process_args(mapping, **kwargs))
        self.model = model
        self.dict = dictionary
        self.type = model_type
        self.sp = sp
        self.cache_locality = cache_locality
        self.acronyms = acronyms
        self.h = None
        if self.type == 'rnnlm':
            self.init_rnnlm()
        elif self.type == 'las':
            assert self.sp is not None, 'sentencepiece model required'
            self.init_las(feat)
        elif self.type == 'tfm':
            assert self.sp is not None, 'sentencepiece model required'
            self.init_tfm(feat)

    def init_rnnlm(self):
        assert hasattr(self.model, 'predict'), 'model must have predict method'
        # RNNLM state, output distribution, posterior, previou ngram
        self.__setitem__(([], 0),
                         (None, torch.zeros(len(self.dict)), [], ([], 0)))

    def init_las(self, feat):
        assert feat is not None, 'for LAS model, acoustic feature must be given'
        with torch.no_grad():
            self.h = self.model.encode(feat).unsqueeze(0)
        c_list = [self.model.dec.zero_state(self.h)
                  for _ in range(self.model.dec.dlayers)]
        z_list = [self.model.dec.zero_state(self.h)
                  for _ in range(self.model.dec.dlayers)]
        att_w = None
        self.model.dec.att[0].reset()
        self.__setitem__(([], 0),
                         ((att_w, c_list, z_list), torch.zeros(len(self.dict)),
                          [], ([], 0), {}))

    def init_tfm(self, feat):
        assert feat is not None, 'for TFM model, acoustic feature must be given'
        with torch.no_grad():
            self.h = self.model.encode(feat).unsqueeze(0)
        self.__setitem__(([], 0), (('', 0.0), None, [], ([], 0), {}))

    def get_value_by_locality(self, timed_cache, key):
        cached_keys = list(timed_cache.keys())
        distance = [abs(key - i) for i in cached_keys]
        min_idx = np.argmin(distance)
        if distance[min_idx] <= self.cache_locality:
            return timed_cache[cached_keys[min_idx]]
        else:
            raise KeyError

    def __getitem__(self, k):
        ngram, timestamp = k
        ngram_dict = super(Cache, self).__getitem__(' '.join(ngram))
        return self.get_value_by_locality(ngram_dict, timestamp)

    def __setitem__(self, k, v):
        ngram, timestamp = k
        try:
            ngram_dict = super(Cache, self).__getitem__(' '.join(ngram))
            return ngram_dict.__setitem__(timestamp, v)
        except KeyError:
            return super(Cache, self).__setitem__(
                ' '.join(ngram), {timestamp: v})

    def __delitem__(self, k):
        ngram, timestamp = k
        return super(Cache, self).__delitem__(' '.join(ngram))

    def get(self, k, default=None):
        try:
            return self.__getitem__(k)
        except KeyError:
            return default

    def setdefault(self, k, default=None):
        try:
            return self.__getitem__(k)
        except KeyError:
            return self.__setitem__(k, default)

    def pop(self, k, v=object()):
        ngram, timestamp = k
        if v is object():
            return super(Cache, self).pop(' '.join(ngram))
        return super(Cache, self).pop(' '.join(ngram), v)

    def update(self, mapping=(), **kwargs):
        super(Cache, self).update(self._process_args(mapping, **kwargs))

    def __contains__(self, k):
        ngram, timestamp = k
        contain_ngram = super(Cache, self).__contains__(' '.join(ngram))
        if not contain_ngram:
            return False
        try:
            _ = self.__getitem__(k)
            return True
        except KeyError:
            return False

    def copy(self):
        return type(self)(self, self.model, self.dict)

    @classmethod
    def fromkeys(cls, keys, v=None):
        return super(Cache, cls).fromkeys((' '.join(k) for k in keys), v)

    def __repr__(self):
        return '{0}({1})'.format(
            type(self).__name__, super(Cache, self).__repr__())

    def get_state(self, k):
        return deepcopy(self.__getitem__(k)[0])

    def get_pred(self, k, word):
        if self.type == 'rnnlm':
            return self.__getitem__(k)[1][sym2idx(self.dict, word)].item()
        elif self.type == 'las':
            with torch.no_grad():
                att_w, c_list, z_list = self.get_state(k)
                if word == lattice.SOS:
                    y = torch.LongTensor([self.model.sos])
                elif word == lattice.EOS:
                    y = torch.LongTensor([self.model.eos])
                else:
                    # This also works for the first word
                    mapped_word = self.acronyms.get(word, word)
                    hyp = self.sp.encode_as_pieces(' ' + mapped_word)
                    y = torch.LongTensor(
                        [sym2idx(self.dict, char) for char in hyp])
                if word in self.get_word_dict(k):
                    return self.get_word_dict(k)[word]
                score = self.__getitem__(k)[1][y[0]].item()
                for char_idx in range(len(y) - 1):
                    att_c, att_w = self.model.dec.att[0](
                        self.h,
                        [self.h.size(1)],
                        self.model.dec.dropout_dec[0](z_list[0]),
                        att_w,
                    )
                    ey = self.model.dec.dropout_emb(
                        self.model.dec.embed(y[char_idx].unsqueeze(0)))
                    ey = torch.cat((ey, att_c), dim=1)
                    z_list, c_list = self.model.dec.rnn_forward(
                        ey, z_list, c_list, z_list, c_list)
                    input_to_layerout = self.model.dec.dropout_dec[-1](
                        z_list[-1])
                    if self.model.dec.context_residual:
                        input_to_layerout = torch.cat(
                            (input_to_layerout, att_c), dim=-1)
                    out = self.model.dec.output(input_to_layerout)
                    pred = torch.nn.functional.log_softmax(out, dim=1)
                    score += pred[0][y[char_idx + 1]].item()
            self.get_word_dict(k)[word] = score
            return score
        elif self.type == 'tfm':
            with torch.no_grad():
                hist, hist_score = self.get_state(k)
                if word == lattice.SOS:
                    y = torch.LongTensor([self.model.sos])
                    score = 0.0
                else:
                    if word in self.get_word_dict(k):
                        return self.get_word_dict(k)[word]
                    if word == lattice.EOS:
                        hyp = self.sp.encode_as_pieces(hist)
                        y = torch.LongTensor(
                            [self.model.sos]
                            + [sym2idx(self.dict, char) for char in hyp]
                            + [self.model.eos])
                    else:
                        mapped_word = self.acronyms.get(word, word)
                        hyp = self.sp.encode_as_pieces(hist + ' ' + mapped_word)
                        y = torch.LongTensor(
                            [self.model.sos]
                            + [sym2idx(self.dict, char) for char in hyp])
                    y_in = y[:-1].unsqueeze(0)
                    y_mask = target_mask(y_in, self.model.ignore_id)
                    pred = self.model.decoder(y_in, y_mask, self.h, None)[0]
                    pred = torch.nn.functional.log_softmax(pred[0], dim=1)
                    score = torch.sum(
                        pred[torch.arange(pred.size(0)), y[1:]]).item()
            self.get_word_dict(k)[word] = score - hist_score
            return score - hist_score

    def get_post(self, k):
        return self.__getitem__(k)[2]

    def get_prev_ngram(self, k):
        return self.__getitem__(k)[3]

    def get_word_dict(self, k):
        # this is only for ISCA subword rescoring
        return self.__getitem__(k)[4]

    def get_timestamp(self, k):
        ngram, timestamp = k
        ngram_dict = super(Cache, self).__getitem__(' '.join(ngram))
        cached_keys = list(ngram_dict.keys())
        distance = [abs(timestamp - i) for i in cached_keys]
        min_idx = np.argmin(distance)
        return cached_keys[min_idx]

    def renew(self, prev_ngram, new_ngram, post):
        word = new_ngram[0][-1]
        if word in [lattice.OOV, lattice.UNK]:
            # skip oov and unk
            value = self.__getitem__(prev_ngram)
        else:
            with torch.no_grad():
                if self.type == 'rnnlm':
                    state, pred = self.model.predict(
                        self.get_state(prev_ngram),
                        torch.LongTensor([sym2idx(self.dict, word)])
                    )
                    value = (state, pred[0], post, prev_ngram)
                elif self.type == 'las':
                    att_w, c_list, z_list = self.get_state(prev_ngram)
                    if word == lattice.SOS:
                        y = torch.LongTensor([self.model.sos])
                    elif word == lattice.EOS:
                        y = torch.LongTensor([self.model.eos])
                    else:
                        mapped_word = self.acronyms.get(word, word)
                        hyp = self.sp.encode_as_pieces(mapped_word)
                        y = torch.LongTensor(
                            [sym2idx(self.dict, char) for char in hyp])
                    score = self.__getitem__(prev_ngram)[1][y[0]].item()
                    for char_idx in range(len(y)):
                        att_c, att_w = self.model.dec.att[0](
                            self.h,
                            [self.h.size(1)],
                            self.model.dec.dropout_dec[0](z_list[0]),
                            att_w,
                        )
                        ey = self.model.dec.dropout_emb(
                            self.model.dec.embed(y[char_idx].unsqueeze(0)))
                        ey = torch.cat((ey, att_c), dim=1)
                        z_list, c_list = self.model.dec.rnn_forward(
                            ey, z_list, c_list, z_list, c_list)
                        input_to_layerout = self.model.dec.dropout_dec[-1](
                            z_list[-1])
                        if self.model.dec.context_residual:
                            input_to_layerout = torch.cat(
                                (input_to_layerout, att_c), dim=-1)
                        out = self.model.dec.output(input_to_layerout)
                        pred = torch.nn.functional.log_softmax(out, dim=1)
                        if char_idx < len(y) - 1:
                            score += pred[0][y[char_idx + 1]].item()
                    self.get_word_dict(prev_ngram)[word] = score
                    state = (att_w, c_list, z_list)
                    value = (state, pred[0], post, prev_ngram, {})
                elif self.type == 'tfm':
                    hist, hist_score = self.get_state(prev_ngram)
                    if word == lattice.SOS:
                        y = torch.LongTensor([self.model.sos])
                        score = 0.0
                        state = ('', 0.0)
                    else:
                        if word == lattice.EOS:
                            hyp = self.sp.encode_as_pieces(hist)
                            y = torch.LongTensor(
                                [self.model.sos]
                                + [sym2idx(self.dict, char) for char in hyp]
                                + [self.model.eos])
                        else:
                            mapped_word = self.acronyms.get(word, word)
                            hyp = self.sp.encode_as_pieces(
                                hist + ' ' + mapped_word)
                            y = torch.LongTensor(
                                [self.model.sos]
                                + [sym2idx(self.dict, char) for char in hyp])
                        y_in = y[:-1].unsqueeze(0)
                        y_mask = target_mask(y_in, self.model.ignore_id)
                        pred = self.model.decoder(y_in, y_mask, self.h, None)[0]
                        pred = torch.nn.functional.log_softmax(pred[0], dim=1)
                        score = torch.sum(
                            pred[torch.arange(pred.size(0)), y[1:]]).item()
                        state = (hist + ' ' + mapped_word, score)
                    self.get_word_dict(prev_ngram)[word] = score - hist_score
                    value = (state, None, post, prev_ngram, {})
        self.__setitem__(new_ngram, value)

def rnnlm_rescore(lat, lm, dictionary, ngram, replace=False, cache_locality=9):
    """Lattice rescoring with RNNLM with on-the-fly lattice expansion
    using n-gram based history clustering.
    Optionally, run forward-backward before calling this function,
    so the cache can be updated based on lattice node posterior.

    :param lat: Word lattice object.
    :type lat: lattice.Lattice
    :param lm: RNNLM.
    :type lm: torch.nn.Module
    :param dictionary: Mapping from word to index.
    :type dictionary: dict
    :param ngram: Number of n-gram for history clustering.
    :type ngram: int
    :param replace: Replace existing scores if True, otherwise append.
    :type replace: bool
    :param cache_locality: Only use cache if around given number of frames.
    :type cache_locality: int
    :return: An expanded lattice with RNNLM score on each arc.
    :rtype: lattice.Lattice
    """
    # setup ngram cache
    cache = Cache(
        lm, dictionary, model_type='rnnlm', cache_locality=cache_locality)
    cache_hit = 0
    cache_miss_new = 0
    cache_miss_update = 0
    # initialise expanded node & outbound arc list
    for node in lat.nodes:
        node.expnodes = []
        node.exparcs = []
    lat.start.expnodes.append(lat.start.subnode())
    # lattice traversal
    for n_i in lat.nodes:
        for n_j in n_i.expnodes:
            for a_k in n_i.exits:
                # find the destination node n_k of arc a_k
                n_k = a_k.dest
                # find the LM state phi(h_{n_0}^{n_j}) of expanded node n_j
                phi_nj = n_j.lmstate
                # find a new LM state phi(h_{n_0}^{n_k}) for node n_k
                phi_nk = ((phi_nj[0] + [n_k.sym])[-ngram:], n_k.entry)
                try:
                    # check if the destination node needs to be expanded
                    idx = [' '.join(i.lmstate[0]) for i in n_k.expnodes].index(
                        ' '.join(phi_nk[0]))
                    n_l = n_k.expnodes[idx]
                except ValueError:
                    # create a new node for expansion
                    n_l = n_k.subnode()
                    n_l.lmstate = deepcopy(phi_nk)
                    n_k.expnodes.append(n_l)
                new_arc = a_k.subarc(n_j, n_l)

                # update cache except for the final node
                if n_k.sym != lattice.EOS:
                    phi_nk_post = (cache.get_post(phi_nj) + [a_k.post])[-ngram:]
                    # compute LM probability P(n_k|phi(h_{n_0}^{n_j}))
                    if phi_nk not in cache:
                        # create new entry in cache for unseen ngram
                        cache.renew(phi_nj, phi_nk, phi_nk_post)
                        cache_miss_new += 1
                    else:
                        timestamp_condition = (
                            abs(cache.get_timestamp(phi_nk) - n_k.entry)
                            <= cache.cache_locality)
                        # if the new ngram phi_nk in already cache
                        if (cache.get_prev_ngram(phi_nk)[0] == phi_nj
                            and timestamp_condition):
                            # if the previous ngram phi_nj is the same
                            # then cache hit, do not forward again
                            # note that same ngram can have different posterior
                            # because of different timestamps
                            cache_hit += 1
                        else:
                            # if the previous ngram phi_nj is different
                            posterior_condition = (
                                sum(phi_nk_post) > sum(cache.get_post(phi_nk)))
                            if posterior_condition or not timestamp_condition:
                                # renew the cache with higher posterior
                                cache.renew(phi_nj, phi_nk, phi_nk_post)
                                cache_miss_update += 1
                            else:
                                cache_hit += 1

                if n_k.sym in [lattice.OOV, lattice.UNK]:
                    if replace:
                        new_arc.nscr = [0.0]
                    else:
                        new_arc.nscr.append(0.0)
                else:
                    if replace:
                        new_arc.nscr = [cache.get_pred(phi_nj, n_k.sym)]
                    else:
                        new_arc.nscr.append(cache.get_pred(phi_nj, n_k.sym))
                n_j.exits.append(new_arc)
                n_l.entries.append(new_arc)
    # Rebuild lattice from expanded nodes & arcs
    new_lat = lattice.Lattice(header=lat.header, nframes=lat.nframes)
    for node in lat.nodes:
        if node != lat.end:
            for expnode in node.expnodes:
                new_lat.nodes.append(expnode)
                new_lat.arcs.extend(expnode.entries)
        else:
            # merge the final node
            end_node = node.expnodes[0]
            for expnode in node.expnodes[1:]:
                for exparc in expnode.entries:
                    exparc.dest = end_node
                    end_node.entries.append(exparc)
            new_lat.nodes.append(end_node)
            new_lat.arcs.extend(end_node.entries)
    new_lat.start = new_lat.nodes[0]
    new_lat.end = new_lat.nodes[-1]
    # print('Cache hit: %d' % cache_hit)
    # print('Cache miss new: %d' % cache_miss_new)
    # print('Cache miss update: %d' % cache_miss_update)
    # print('Cache hit rate: %.4f' % (
    #     cache_hit / (cache_hit + cache_miss_new + cache_miss_update)))
    return new_lat

def isca_rescore(lat, feat, model, char_dict, ngram, sp, model_type='las',
                 replace=False, cache_locality=9, acronyms={}):
    """Lattice rescoring with LAS model with on-the-fly lattice expansion
    using n-gram based history clustering.
    Optionally, run forward-backward before calling this function,
    so the cache can be updated based on lattice node posterior.

    :param lat: Word lattice object.
    :type lat: lattice.Lattice
    :param feat: Acoustic feature for the utterance.
    :type feat: np.ndarray
    :param model: ESPnet RNN-based LAS model.
    :type model: torch.nn.Module
    :param char_dict: Mapping from character to index.
    :type char_dict: dict
    :param ngram: Number of n-gram for history clustering.
    :type ngram: int
    :param sp: Sentencepiece model.
    :type sp: sentencepiece model object
    :param model_type: Type of the end-to-end model, one of ['las', 'tfm'].
    :type model_type: str
    :param replace: Replace existing scores if True, otherwise append.
    :type replace: bool
    :param cache_locality: Only use cache if around given number of frames.
    :type cache_locality: int
    :param acronyms: Mapping to match vocabulary.
    :type acronyms: dict
    :return: An expanded lattice with ISCA score on each arcs.
    :rtype: lattice.Lattice
    """
    # setup ngram cache
    cache = Cache(
        model, char_dict, feat=feat, model_type=model_type, sp=sp,
        cache_locality=cache_locality, acronyms=acronyms
    )
    cache_hit = 0
    cache_miss_new = 0
    cache_miss_update = 0
    # initialise expanded node & outbound arc list
    for node in lat.nodes:
        node.expnodes = []
        node.exparcs = []
    lat.start.expnodes.append(lat.start.subnode())
    # lattice traversal
    for n_i in lat.nodes:
        for n_j in n_i.expnodes:
            for a_k in n_i.exits:
                # find the destination node n_k of arc a_k
                n_k = a_k.dest
                # find the LM state phi(h_{n_0}^{n_j}) of expanded node n_j
                phi_nj = n_j.lmstate
                # find a new LM state phi(h_{n_0}^{n_k}) for node n_k
                phi_nk = ((phi_nj[0] + [n_k.sym])[-ngram:], n_k.entry)
                # check if the destination node needs to be expanded
                try:
                    idx = [' '.join(i.lmstate[0]) for i in n_k.expnodes].index(
                        ' '.join(phi_nk[0]))
                    n_l = n_k.expnodes[idx]
                except ValueError:
                    # create a new node for expansion
                    n_l = n_k.subnode()
                    n_l.lmstate = deepcopy(phi_nk)
                    n_k.expnodes.append(n_l)
                new_arc = a_k.subarc(n_j, n_l)

                # update cache except for the final node
                if n_k.sym != lattice.EOS:
                    phi_nk_post = (cache.get_post(phi_nj) + [a_k.post])[-ngram:]
                    # compute LM probability P(n_k|phi(h_{n_0}^{n_j}))
                    if phi_nk not in cache:
                        # create new entry in cache for unseen ngram
                        cache.renew(phi_nj, phi_nk, phi_nk_post)
                        cache_miss_new += 1
                    else:
                        timestamp_condition = (
                            abs(cache.get_timestamp(phi_nk) - n_k.entry)
                            <= cache.cache_locality)
                        if (cache.get_prev_ngram(phi_nk)[0] == phi_nj
                            and timestamp_condition):
                            # if the previous ngram phi_nj is the same
                            # then cache hit, do not forward again
                            # note that same ngram can have different posterior
                            # because of different timestamps
                            cache_hit += 1
                        else:
                            posterior_condition = (
                                sum(phi_nk_post) > sum(cache.get_post(phi_nk)))
                            # if the previous ngram phi_nj is different
                            if posterior_condition or not timestamp_condition:
                                # renew the cache with higher posterior
                                cache.renew(phi_nj, phi_nk, phi_nk_post)
                                cache_miss_update += 1
                            else:
                                cache_hit += 1

                if n_k.sym in [lattice.OOV, lattice.UNK]:
                    if replace:
                        new_arc.iscr = [0.0]
                    else:
                        new_arc.iscr.append(0.0)
                else:
                    if replace:
                        new_arc.iscr = cache.get_pred(phi_nj, n_k.sym)
                    else:
                        new_arc.iscr.append(cache.get_pred(phi_nj, n_k.sym))
                n_j.exits.append(new_arc)
                n_l.entries.append(new_arc)
    # Rebuild lattice from expanded nodes & arcs
    new_lat = lattice.Lattice(header=lat.header, nframes=lat.nframes)
    for node in lat.nodes:
        if node != lat.end:
            for expnode in node.expnodes:
                new_lat.nodes.append(expnode)
                new_lat.arcs.extend(expnode.entries)
        else:
            # merge the final node
            end_node = node.expnodes[0]
            for expnode in node.expnodes[1:]:
                for exparc in expnode.entries:
                    exparc.dest = end_node
                    end_node.entries.append(exparc)
            new_lat.nodes.append(end_node)
            new_lat.arcs.extend(end_node.entries)
    new_lat.start = new_lat.nodes[0]
    new_lat.end = new_lat.nodes[-1]
    # print('Cache hit: %d' % cache_hit)
    # print('Cache miss new: %d' % cache_miss_new)
    # print('Cache miss update: %d' % cache_miss_update)
    # print('Cache hit rate: %.4f' % (
    #     cache_hit / (cache_hit + cache_miss_new + cache_miss_update)))
    return new_lat