import hydra
import hydra.experimental
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import RobertaTokenizer
from omegaconf import DictConfig

from torchfly.text.decode import TransformerDecoder
from torchfly.common import set_random_seed
from torchfly.text.rl import TextRLRewardFunc

from configure_dataloader import DataLoaderHandler, TextRLCollator

from model import TextGAILModel
from textgail_discriminator import TextGAILDiscriminator
from textgail_trainerloop import TextGAILTrainerLoop

import logging

@hydra.main(config_path="config/config.yaml", strict=False)
def main(config=None):
    set_random_seed(config.training.random_seed)
    dataloader_handler = DataLoaderHandler(config)

    model = TextGAILModel(config)
    # model_weights = torch.load(config.task.weights_path)
    # model.generator.load_state_dict(model_weights, strict=False)
    model = model.cuda()

    # Register your transformer for decoding
    tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
    decoder_helper = TransformerDecoder(config.decode)
    decoder_helper.register_generator(model.generator.decoder)
    decoder_helper.register_tokenizer(tokenizer)

    reward_func = TextGAILDiscriminator(config, tokenizer, model.discriminator)

    trainer = TextGAILTrainerLoop(config=config,
                                reward_func=reward_func,
                                decoder_helper=decoder_helper,
                                model=model,
                                train_dataloader_fn=dataloader_handler.train_dataloader,
                                valid_dataloader_fn=dataloader_handler.test_dataloader)

    trainer.train()

if __name__ == "__main__":
    main()


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

"""
    File name: predict_2.py
    Function Des: ...
    ~~~~~~~~~~

    author: Skyduy <cuteuy@gmail.com> <http://skyduy.me>

"""
import urllib
import numpy as np
from utils import asc_chr
from PIL import Image
from keras.models import Graph
from keras.layers.core import Dense, Flatten, Dropout
from keras.layers.convolutional import Convolution2D, MaxPooling2D

print '... loading data'

def get_data():
    pic_name = 'tmp.jpg'
    urllib.urlretrieve('http://passport2.chaoxing.com/img/code', 'tmp/%s' % pic_name)
    results = {}
    raw_img = Image.open('tmp/tmp.jpg')
    img = raw_img.convert('1')
    x_size, y_size = img.size
    for x in xrange(x_size):
        results.setdefault(x, 0)
        for y in xrange(y_size):
            pixel = img.getpixel((x, y))
            if pixel == 0:
                results[x] += 1
    begin = end = -1
    for i in range(150):
        if results[i] > 0:
            begin = i
            break
    threshold_list = []
    for i in range(131, 134):
        threshold_list.append(results[i])
    threshold = max(threshold_list)
    i = 150
    stop = False
    while not stop:
        i -= 1
        if i == -1:
            begin, end = 10, 135
            break
        v = results[i]
        if v >= max((3.5, 2 * threshold)):
            stop = True
            tmp_value = v
            while True:
                i += 1
                if results[i] <= tmp_value:
                    if results[i] <= threshold:
                        end = i
                        break
                    else:
                        tmp_value = results[i]
                else:
                    end = i - 1
    centers = [((2 * i + 1) * end + (9 - 2 * i) * begin) / 10 for i in range(5)]
    data = np.empty((5, 3, 40, 40), dtype="float32")
    for index, center in enumerate(centers):
        a = center - 19
        b = center + 21
        img = raw_img.crop((a, 0, b, y_size))
        arr = np.asarray(img, dtype="float32") / 255.0
        data[index, :, :, :] = np.rollaxis(arr, 2)
    return data

X = get_data()

print '... create model'

graph = Graph()
graph.add_input(name='input', input_shape=(3, 40, 40))
graph.add_node(Convolution2D(22, 5, 5, activation='relu'), name='conv1', input='input')
graph.add_node(MaxPooling2D(pool_size=(2, 2)), name='pool1', input='conv1')
graph.add_node(Convolution2D(44, 3, 3, activation='relu'), name='conv2', input='pool1')
graph.add_node(MaxPooling2D(pool_size=(2, 2)), name='pool2', input='conv2')
graph.add_node(Dropout(0.25), name='drop', input='pool2')
graph.add_node(Flatten(), name='flatten', input='drop')
graph.add_node(Dense(256, activation='relu'), name='ip', input='flatten')
graph.add_node(Dropout(0.5), name='drop_out', input='ip')
graph.add_node(Dense(19, activation='softmax'), name='result', input='drop_out')

graph.add_output(name='out', input='result')

print '... compiling'
graph.compile(
    optimizer='adadelta',
    loss={
        'out': 'categorical_crossentropy',
    }
)

graph.load_weights('model/model_2.hdf5')
print '... predicting'
result = graph.predict({'input': X})
out = [np.argmax(result['out'][0]), np.argmax(result['out'][1]),
       np.argmax(result['out'][2]), np.argmax(result['out'][3]),
       np.argmax(result['out'][4])]

print out

r = [asc_chr(i) for i in out]

print r


import sys
import threading
import logging

import numpy as np

import soundfile as sf
from common.continuousBuffer import ContinuousBuffer

class TargetSignal(threading.Thread):
    """
    Represents an output signal such as the speaker
    Runs in its own thread to push data onto the buffer

    Attributes
    ----------
    buffer : ContinuousBuffer
        A thread-safe ContinuousBuffer that reads the underlying data stream (stereo sound)
    stepSize: int
        The number of signal frames or matrix rows used in the transfer of data to the server
    size: int
        The number of frames to read from the target signal file
    targetFile: String
        The name of the target signal data file
    targetSignal: numpy.ndarray
        The sound data read in from targetFile
    stopped: Boolean
        A flag indicating if the algorithm is running
    """
    def __init__(self, targetFile, buffer, stepSize, size, threadName):
        """
        Parameters
        ----------
        targetFile: String
            The name of the target signal data file
        buffer : ContinousBuffer
            A thread-safe ContinuousBUffer that reads the underlying data stream (stereo sound)
        stepSize: int
            The number of signal frames or matrix rows used in the transfer of data to the server
        size: int
            The number of frames to read from the target signal file
        threadName : str
            The name of the thread. Utilized for debugging.
        """
        logging.debug('Initialize Target File thread')
        threading.Thread.__init__(self, name=threadName, daemon=True)
        self.buffer = buffer
        self.stepSize = stepSize
        self.size = size
        self.targetFile = targetFile
        self.targetSignal = None
        self.stopped = False

    def stop(self):
        """
        Called when the algorithm is stopped and sets the stopped flag to False

        Parameters
        ----------
        None

        Returns
        -------
        None

        Raises
        ------
        None
        """
        self.stopped = True

    def run(self):
        """
        The primary thread function that initializes the
        TargetStream. This is called once to load the buffer
        and then the algorithm can use the ContinuousBuffer
        to load "new" information into the algorithm

        Parameters
        ----------
        None

        Returns
        -------
        None

        Raises
        ------
        None
        """
        try:
            # this thread does not keep running. it will end as soon as the buffer is populated
            logging.debug('Running Target File thread')
            self.targetSignal, _ = sf.read(
                self.targetFile, dtype='float32')
            self.buffer.push(self.targetSignal[0:self.size])

        except Exception as e:
            logging.exception(f'Exception thrown: {e}')


"""
The coronagraph model relies on numerous parameters describing the telescope,
planet, and star used for each calculation. Below :class:`Telescope`,
:class:`Planet`, and :class:`Star` classes are listed, which can be instantiated
and passed along to noise calculations.
"""
from __future__ import (division as _, print_function as _,
                absolute_import as _, unicode_literals as _)

import numpy as np
from .noise_routines import lambertPhaseFunction

__all__ = ['Telescope', 'Planet', 'Star']

################################################################################
# TELESCOPE
################################################################################

class Telescope(object):
    """
    A class to represent a telescope object and all design specifications
    therein

    Parameters
    ----------
    mode : str
        Telescope observing modes: 'IFS', 'Imaging'
    lammin : float
        Minimum wavelength (um)
    lammax : float
        Maximum wavelength (um)
    R : float
        Spectral resolution (lambda / delta-lambda)
    Tsys : float
        Telescope temperature (K)
    D : float
        Telescope diameter (m)
    emis : float
        Telescope emissivity
    IWA : float
        Inner Working Angle (lambda/D)
    OWA : float
        Outer Working Angle (lambda/D)
    Tput : float
        Telescope throughput
    C : float
        Raw Contrast
    De : float
        Dark current (s**-1)
    DNHpix : float
        Horizontal pixel spread of IFS spectrum
    Re : float
        Read noise per pixel
    Rc : float, optional
        Clock induced charge [counts/pixel/photon]
    Dtmax : float
        Maximum exposure time (hr)
    X : float
        Size of photometric aperture (lambda/D)
    q : float
        Quantum efficiency
    filter_wheel : Wheel (optional)
        Wheel object containing imaging filters
    aperture : str
        Aperture type ("circular" or "square")
    A_collect : float
        Mirror collecting area (m**2) if different than :math:`\\pi(D/2)^2`
    diam_circumscribed : float, optional
        Circumscribed telescope diameter [m] used for IWA and OWA (uses `diam`
        if `None` provided)
    diam_inscribed : float, optional
        Inscribed telescope diameter [m] used for lenslet calculations
        (uses `diam` if `None` provided)
    Tput_lam : tuple of arrays
        Wavelength-dependent throughput e.g. ``(wls, tputs)``. Note that if
        ``Tput_lam`` is used the end-to-end throughput will equal the
        convolution of ``Tput_lam[1]`` with ``Tput``.
    qe_lam : tuple of arrays
        Wavelength-dependent throughput e.g. ``(wls, qe)``. Note that if
        ``qe_lam`` is used the total quantum efficiency will equal the
        convolution of ``qe_lam[1]`` with ``q``.
    lammin_lenslet : float, optional
        Minimum wavelength to use for lenslet calculation (default is ``lammin``)
    lam : array-like, optional
        Wavelength grid for spectrograph [microns] (uses ``lammin``, ``lammax``,
        and ``resolution`` to determine if ``None`` provided)
    dlam : array-like, optional
        Wavelength grid `widths` for spectrograph [microns] (uses ``lammin``, ``lammax``,
        and ``resolution`` to determine if ``None`` provided)

    Methods
    -------
    default_luvoir()
        Initialize telescope object using current LUVOIR parameters (Not decided!)
    default_habex()
        Initialize telescope object using current HabEx parameters (Not decided!)
    default_wfirst()
        Initialize telescope object using current WFIRST parameters (Not decided!)
    """

    # Define a constructor
    def __init__(self, mode='IFS', lammin=0.3, lammax=2.0, R=70., Tput=0.2,\
                 D=8.0, Tsys=260., Tdet=50., IWA=0.5, OWA=30000., emis=0.9,\
                 C=1e-10, De=1e-4, DNHpix=3., Re=0.1, Rc=0.0, Dtmax=1.0, X=0.7, q=0.9,\
                 filter_wheel=None, aperture = "circular", A_collect = None,
                 Tput_lam = None, qe_lam = None, lammin_lenslet = None,
                 diam_circumscribed = None, diam_inscribed = None, lam = None,
                 dlam = None):
        self._mode=mode
        self.lammin=lammin
        self.lammax=lammax
        self.resolution=R
        self.throughput=Tput
        self.diameter=D
        self.Tsys=Tsys
        self.Tdet=Tdet
        self.IWA=IWA
        self.OWA=OWA
        self.emissivity=emis
        self.contrast=C
        self.aperture = aperture
        self.diam_circumscribed = diam_circumscribed
        self.diam_inscribed = diam_inscribed
        self.lam = lam
        self.dlam = dlam

        self.darkcurrent=De
        self.DNHpix=DNHpix
        self.readnoise=Re
        self.Rc=Rc
        self.Dtmax=Dtmax
        self.X=X
        self.qe=q
        self.A_collect = A_collect
        self.Tput_lam = Tput_lam
        self.qe_lam = qe_lam
        self.lammin_lenslet = lammin_lenslet

        self._filter_wheel=filter_wheel

        if self._mode == 'Imaging':
            from filters.imager import johnson_cousins
            self._filter_wheel = johnson_cousins()

    @classmethod
    def default_luvoir(cls):
        # Return new class instance
        return cls(mode="IFS", lammin=0.5, lammax=1.0, R=70.,
                   Tput=0.05, D=12., Tsys=150., Tdet=50., IWA=3.0,
                   OWA=20.0, emis=0.9, C=1e-10, De=1e-4,
                   DNHpix=3.0, Re=0.1, Dtmax=1.0, X=1.5,
                   q=0.9, filter_wheel=None)

    @classmethod
    def default_habex(cls):
        print("These HabEx parameters are not confirmed yet!")
        # Return new class instance
        return cls(mode="IFS", lammin=0.4, lammax=2.5, R=70.,
                   Tput=0.05, D=6., Tsys=150., Tdet=50., IWA=3.0,
                   OWA=20.0, emis=0.9, C=1e-10, De=1e-4,
                   DNHpix=3.0, Re=0.1, Dtmax=1.0, X=1.5,
                   q=0.9, filter_wheel=None)

    @classmethod
    def default_wfirst(cls):
        print("These WFIRST parameters are not confirmed yet!")
        # Return new class instance
        return cls(mode="IFS", lammin=0.6, lammax=1.0, R=70.,
                   Tput=0.05, D=2.4, Tsys=150., Tdet=50., IWA=3.0,
                   OWA=20.0, emis=0.9, C=1e-9, De=1e-4,
                   DNHpix=3.0, Re=0.1, Dtmax=1.0, X=1.5,
                   q=0.9, filter_wheel=None)

    @property
    def mode(self):
        return self._mode

    @mode.setter
    def mode(self, value):
        self._mode = value
        if value == 'Imaging':
            from filters.imager import johnson_cousins
            self._filter_wheel = johnson_cousins()
        else:
            self._filter_wheel = None

    @property
    def filter_wheel(self):
        return self._filter_wheel

    @filter_wheel.setter
    def filter_wheel(self, value):
        if (value.__class__.__name__ == 'Wheel') or (value.__class__.__base__.__name__ == 'Wheel'):
            self._filter_wheel = value
        else:
            print("Error in Telescope: Specified filter wheel is not of type 'Wheel'")
            self._filter_wheel = None

    def __str__(self):
        string = 'Coronagraph: \n------------\n'+\
            '- Telescope observing mode : '+"%s" % (self.mode)+'\n'+\
            '- Minimum wavelength (um) : '+"%s" % (self.lammin)+'\n'+\
            '- Maximum wavelength (um)  : '+"%s" % (self.lammax)+'\n'+\
            '- Spectral resolution (lambda / delta-lambda)  : '+"%s" % (self.resolution)+' \n'+\
            '- Telescope/System temperature (K)  : '+"%s" % (self.Tsys)+' \n'+\
            '- Detector temperature (K)  : '+"%s" % (self.Tdet)+' \n'+\
            '- Telescope diameter (m)  : '+"%s" % (self.diameter)+' \n'+\
            '- Telescope emissivity  : '+"%s" % (self.emissivity)+' \n'+\
            '- Inner Working Angle (lambda/D)  : '+"%s" % (self.IWA)+' \n'+\
            '- Outer Working Angle (lambda/D)  : '+"%s" % (self.OWA)+' \n'+\
            '- Telescope throughput  : '+"%s" % (self.throughput)+' \n'+\
            '- Raw Contrast  : '+"%s" % (self.contrast)+' \n'+\
            '- Dark current (s**-1)  : '+"%s" % (self.darkcurrent)+' \n'+\
            '- Horizontal pixel spread of IFS spectrum  : '+"%s" % (self.DNHpix)+' \n'+\
            '- Read noise per pixel  : '+"%s" % (self.readnoise)+' \n'+\
            '- Maximum exposure time (hr)  : '+"%s" % (self.Dtmax)+' \n'+\
            '- Size of photometric aperture (lambda/D)  : '+"%s" % (self.X)+' \n'+\
            '- Quantum efficiency  : '+"%s" % (self.qe)+''
        return string

################################################################################
# PLANET
################################################################################

class Planet(object):
    """
    A class to represent a planet and all associated parameters of the planet
    to be observed.

    Parameters
    ----------
    name : string
        Planet name from database
    star : string
        Stellar type of planet host star
    d : float
        Distance to system (pc)
    Nez : float
        Number of exzodis (zodis)
    Rp : float
        Radius of planet (Earth Radii)
    a : float
        Semi-major axis (AU)
    alpha : float
        Phase angle (deg)
    Phi : float
        Lambertian phase function
    MzV : float
        Zodiacal light surface brightness (mag/arcsec**2)
    MezV : float
        exozodiacal light surface brightness (mag/arcsec**2)

    Methods
    -------
    from_file()
        Initialize object using planet parameters in the Input file
    """

    # Define a constructor
    def __init__(self, name='earth', star='sun', d=10.0,Nez=1.0,\
                 Rp=1.0, a=1.0, alpha=90.,\
                 MzV=23.0, MezV=22.0):
        self.name=name
        self.star=star
        self.distance=d
        self.Nez=Nez
        self.Rp=Rp
        self.a=a
        self._alpha=alpha
        self._Phi = None
        self.MzV  = MzV     # zodiacal light surface brightness (mag/arcsec**2)
        self.MezV = MezV     # exozodiacal light surface brightness (mag/arcsec**2)

        if self._Phi is None:
            self._Phi = lambertPhaseFunction(self._alpha)
        else:
            raise Exception("Error in Planet Phase Function (Phi)")

    @property
    def alpha(self):
        return self._alpha

    @alpha.setter
    def alpha(self, value):
        self._alpha = value
        self._Phi = lambertPhaseFunction(value)

    @property
    def Phi(self):
        return self._Phi

    @Phi.setter
    def Phi(self, value):
        self._Phi = value

    def __str__(self):
        string = 'Planet: \n-------\n'+\
            '- Planet name  : '+"%s" % (self.name)+'\n'+\
            '- Stellar type of planet host star : '+"%s" % (self.star)+'\n'+\
            '- Distance to system (pc) : '+"%s" % (self.distance)+'\n'+\
            '- Number of exzodis (zodis) : '+"%s" % (self.Nez)+'\n'+\
            '- Radius of planet (Earth Radii) : '+"%s" % (self.Rp)+'\n'+\
            '- Semi-major axis (AU) : '+"%s" % (self.a)+'\n'+\
            '- Phase angle (deg) : '+"%s" % (self.alpha)+'\n'+\
            '- Lambertian phase function : '+"%s" % (self.Phi)+'\n'+\
            '- Zodiacal light surface brightness (mag/arcsec**2) : '+"%s" % (self.MzV)+'\n'+\
            '- Exozodiacal light surface brightness (mag/arcsec**2) : '+"%s" % (self.MezV)
        return string

################################################################################
# STAR
################################################################################

class Star(object):
    """
    A class to represent the stellar host for an exoplanet observation

    Parameters
    ----------
    Teff : float
        Stellar effective temperature [K]
    Rs : float
        Stellar radius [Solar Radii]
    """

    def __init__(self, Teff=5780.0, Rs=1.0):
        self.Teff=Teff
        self.Rs=Rs

    def __str__(self):
        string = 'Star: \n-----\n'+\
            '- Effective Temperature (K) : '+"%s" % (self.Teff)+'\n'+\
            '- Radius (Solar Radii) : '+"%s" % (self.Rs)
        return string


"""
Usage:
    makelevel (-h|--help)
    makelevel [-v|--verbose] <levelname> <texturemap> <depthmap>

Options:
    -h, --help      Show a brief usage summary.
    -v, --verbose   Increase verbosity of logging
"""
import docopt
import logging
import os

# HACK: pymclevel does not appear to install data correctly. Use our own copy.
os.environ['PYMCLEVEL_YAML_ROOT'] = os.path.join(os.path.dirname(__file__), 'data')
from pymclevel import mclevel
import numpy as np
from PIL import Image

log = logging.getLogger(__name__)

def main():
    # parse options
    opts = docopt.docopt(__doc__)
    logging.basicConfig(level=logging.INFO if opts['--verbose'] else logging.WARN)
    levelname = opts['<levelname>']
    texturemap_fn = opts['<texturemap>']
    depthmap_fn = opts['<depthmap>']

    log.info('Loading texture map...')
    tm_im = Image.open(texturemap_fn)

    log.info('Loading depth map...')
    dm_im = Image.open(depthmap_fn)

    tm_im = tm_im.resize((1920//2,1080//2), Image.ANTIALIAS)
    dm_im = dm_im.resize((1920//2,1080//2), Image.ANTIALIAS)

    # Flip images to correspond to MC's co-ord syste,
    tm_im = tm_im.transpose(Image.FLIP_LEFT_RIGHT)
    dm_im = dm_im.transpose(Image.FLIP_LEFT_RIGHT)

    alpha = np.asarray(tm_im.convert('RGBA'))[:,:,3].astype(np.float32)
    alpha /= alpha.max()

    texture = np.asarray(tm_im.convert('L')).astype(np.float32)
    texture /= texture.max()

    depth = np.asarray(dm_im.convert('L')).astype(np.float32)
    depth /= depth.max()

    if texture.shape != depth.shape:
        log.error('Texture map is {0[1]}x{0[0]} whereas depth map is {1[1]}x{1[0]}'.format(
            texture.shape, depth.shape))
        return 1

    log.info('Loading world...')
    world = mclevel.fromFile(levelname)

    xchunks, zchunks = tuple(s//16 for s in texture.shape)

    log.info('Creating empty chunks...')
    chunks = []
    dx, dz = (xchunks>>1), (zchunks>>1)
    for xchunk in np.arange(xchunks) - dx:
        for zchunk in np.arange(zchunks) - dz:
            chunks.append((xchunk, zchunk))
    world.createChunks(chunks)

    log.info('Processing chunks...')
    n_processed = 0
    for xchunk, zchunk in chunks:
        n_processed += 1
        if n_processed % 100 == 0:
            log.info('Processing chunk {0}/{1}'.format(n_processed, len(chunks)))
        chunk = world.getChunk(xchunk, zchunk)
        px, pz = (xchunk+dx)*16, (zchunk+dz)*16

        # Reset chunk
        chunk.Blocks[:] = world.materials.Air.ID
        chunk.Data[:] = 0

        colour = texture[px:(px+16), pz:(pz+16)]
        height = depth[px:(px+16), pz:(pz+16)]
        mask = alpha[px:(px+16), pz:(pz+16)]

        # Scale height
        height = np.where(mask > 0, height * 64 + 70, 10)

        # Get floor value of height / 16
        block_height = np.floor(height).astype(np.int32)
        drift_height = np.floor((height - block_height) * 9).astype(np.uint8)

        for x in range(colour.shape[0]):
            for z in range(colour.shape[0]):
                h = block_height[x,z]
                chunk.Blocks[x,z,:h] = world.materials.Snow.ID
                if drift_height[x,z] != 0:
                    chunk.Blocks[x,z,h] = world.materials.SnowLayer.ID
                    chunk.Data[x,z,h] = drift_height[x,z] - 1

        chunk.chunkChanged()

    world.generateLights()
    world.saveInPlace()

    return 0 # success

if __name__ == '__main__':
    import sys
    sys.exit(main())


"""
 Copyright (c) 2020 Intel Corporation

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
"""

import logging

import numpy as np

from segmentoly.rcnn.openvino_net import OpenVINONet

class TextMaskRCNNOpenVINO(OpenVINONet):
    def __init__(self, detector_model_path, encoder_model_path, decoder_model_path, *args,
                 **kwargs):
        super(TextMaskRCNNOpenVINO, self).__init__(detector_model_path[:-3] + 'xml',
                                                   detector_model_path[:-3] + 'bin',
                                                   *args, **kwargs)
        required_input_keys = {'im_data', 'im_info'}
        assert required_input_keys == set(self.net.inputs.keys())
        required_output_keys = {'boxes', 'scores', 'classes', 'raw_masks', 'text_features'}
        assert required_output_keys.issubset(self.net.outputs.keys())

        self.n, self.c, self.h, self.w = self.net.inputs['im_data'].shape
        assert self.n == 1, 'Only batch 1 is supported.'

        self.text_encoder = OpenVINONet(encoder_model_path[:-3] + 'xml',
                                        encoder_model_path[:-3] + 'bin')

        self.text_decoder = OpenVINONet(decoder_model_path[:-3] + 'xml',
                                        decoder_model_path[:-3] + 'bin')

    def __call__(self, im_data, im_info, **kwargs):
        logging.warning('Please do not use this code for performance evaluation purposes. '
                        'This is for accuracy validation only. If you would like to see how fast it '
                        'is please refer to https://github.com/opencv/open_model_zoo/blob/develop/'
                        'demos/python_demos/text_spotting_demo/README.md.')

        im_data = im_data[0].cpu().numpy()
        im_info = im_info[0].cpu().numpy()
        if (self.h - im_data.shape[1] < 0) or (self.w - im_data.shape[2] < 0):
            raise ValueError('Input image should resolution of {}x{} or less, '
                             'got {}x{}.'.format(self.w, self.h, im_data.shape[2],
                                                 im_data.shape[1]))
        im_data = np.pad(im_data, ((0, 0),
                                   (0, self.h - im_data.shape[1]),
                                   (0, self.w - im_data.shape[2])),
                         mode='constant', constant_values=0).reshape(1, self.c, self.h, self.w)
        im_info = im_info.reshape(1, *im_info.shape)
        output = super().__call__(dict(im_data=im_data, im_info=im_info))

        classes = output['classes']
        valid_detections_mask = classes > 0
        classes = classes[valid_detections_mask]
        boxes = output['boxes'][valid_detections_mask]
        scores = output['scores'][valid_detections_mask]
        masks = output['raw_masks'][valid_detections_mask]
        text_features = output['text_features'][valid_detections_mask, :, :]

        texts = []
        for feature in text_features:
            feature = self.text_encoder({'input': feature})['output']
            feature = np.reshape(feature, (feature.shape[0], feature.shape[1], -1))
            feature = np.transpose(feature, (0, 2, 1))

            hidden = np.zeros([1, 1, 256])
            prev_symbol = np.zeros((1,))

            eos = 1
            max_seq_len = 28
            vocab_size = 38
            per_feature_outputs = np.zeros([max_seq_len, vocab_size])

            for i in range(max_seq_len):
                o = self.text_decoder(
                    {'prev_symbol': prev_symbol, 'prev_hidden': hidden, 'encoder_outputs': feature})
                output = o['output']
                per_feature_outputs[i] = output
                prev_symbol = np.argmax(output, axis=1)
                if prev_symbol == eos:
                    break
                hidden = o['hidden']

            texts.append(per_feature_outputs)

        texts = np.array(texts)

        return boxes, classes, scores, np.full(len(classes), 0, dtype=np.int32), masks, texts


# Copyright 2018 Jetperch LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# https://www.usb.org/document-library/usbet20
# "C:\Program Files (x86)\USB-IF Test Suite\USBET20\USBET20.exe"

import numpy as np
import tempfile
import os
import subprocess
import logging

log = logging.getLogger(__name__)
USBET20_PATH = r"C:\Program Files (x86)\USB-IF Test Suite\USBET20\USBET20.exe"

PLUGIN = {
    'name': 'USB Inrush',
    'description': 'Run the USBET tool validate USB inrush current',
    # config: additional configuration options for this plugin
}

def is_available():
    return os.path.isfile(USBET20_PATH)

class UsbInrush:

    def run(self, data):  # RangeToolInvocation
        duration = data.sample_count / data.sample_frequency
        if not is_available():
            return f'USBET tool not found.'
        if not 0.1 < duration < 0.5:  # todo: confirm range
            return f'Invalid duration {duration:.2f}, must be between 0.1 and 0.5 seconds.'
        usbet20_path = USBET20_PATH
        d = data.samples_get()
        current = d['signals']['current']['value']
        voltage = d['signals']['voltage']['value']
        x = np.arange(len(current), dtype=np.float) * (1.0 / data.sample_frequency)
        valid = np.isfinite(current)
        voltage = np.mean(voltage[valid])
        x = x[valid].reshape((-1, 1))
        i_mean = current[valid].reshape((-1, 1))
        values = np.hstack((x, i_mean))

        with tempfile.TemporaryDirectory(prefix='js_') as tempdir:
            filename = os.path.join(tempdir, 'inrush.csv')
            with open(filename, 'wt') as f:
                np.savetxt(f, values, ['%.8f', '%.3f'], delimiter=',')
            args = ','.join(['usbinrushcheck', filename, '%.3f' % voltage])
            log.info('Running USBET20')
            rv = subprocess.run([usbet20_path, args], capture_output=True)
            log.info('USBET returned %s\nSTDERR: %s\nSTDOUT: %s', rv.returncode, rv.stderr, rv.stdout)

def plugin_register(api):
    """Register the USB Inrush plugin.

    :param api: The :class:`PluginServiceAPI` instance.
    :return: True on success any other value on failure.
    """
    if not is_available():
        log.info('USBET20 tool not found - skip usb_inrush plugin')
        return True  # not an error, normal operation

    api.range_tool_register('Analysis/USB Inrush', UsbInrush)
    return True


import os, cx_Oracle
from datetime import *
import requests
import MySQLdb
import numpy as np
import pandas as pd
from fn import *
from oDT import *

livedb = os.getcwd () + "\\robi_live.csv"
db = os.getcwd () + "\\OMDB.csv"
semcol = os.getcwd () + "\\semcols.txt"
CAT = os.getcwd () + "\\CATdef.txt"
try:
    mysqlconn = MySQLdb.connect ("localhost", "root", "admin", "om2")
except:
    mysqlconn = ""

n = datetime.now ()
tm = n.strftime("%H:%M") + " on " + n.strftime ("%m-%d-%Y")

def hr_minus(diff):
    x = datetime.now ()
    d = x - timedelta (hours=diff)
    str_d = d.strftime ("%m-%d-%Y %H:%M:%S")
    return str_d

def timedelt(diff):
    x = datetime.now ()
    d = x + timedelta (hours=diff)
    str_d = d.strftime ("%d-%m-%Y %H:%M:%S")
    return str_d

def text2list(pth):
    f = open (pth, 'r+')
    ls = []
    for i in f.readlines ():
        ls.append (i.replace ('\n', ''))
    return ls

def text2dic(pth):
    f = open (pth, 'r+')
    dc = {}
    for i in f.readlines():
        a1 = i.replace ('\n', '')
        a2 = a1.split (':')
        dc[a2[0]] = a2[1]
    return dc

def getkey(my_dict, ky):
    if ky is not None:
        for key, value in my_dict.items ():
            if key in str (ky):
                return value
        else:
            return "other"

DRCAT = lambda x: '2H' if (x < 120) \
    else ('4H' if (x < 240)\
    else ('6H' if (x < 360)\
    else ('12H' if (x < 720)\
    else ('24H' if (x < 1440)\
    else ('48H' if (x < 2880)\
    else ('72H'))))))

TS = lambda x: '2G' if ('2G' in x) \
    else ('3G' if ('3G' in x) \
    else ('4G' if ('4G' in x) \
    else ('OML' if ('2G' in x) \
    else "other")))

def extrafeat(xdf, tmdelta = 0):
    xdf = xdf.rename (columns=str.upper)
    df = xdf.assign (DURCAT='0')
    df = df.assign (LO='0')
    df = df.assign (CDLO='0')
    df = df.assign (CDLOTECH='0')
    df['DURCAT'] = df.apply (lambda x: DRCAT (x.DUR), axis=1)
    df['LO'] = df.apply (lambda x: pd.to_datetime (x['LASTOCCURRENCE'], errors='coerce', cache=True).strftime("%d%m%y%H%M"), axis=1)
    df['CDLO'] = df['CUSTOMATTR15'].str.cat (df['LO'])
    df['CDLOTECH'] = df['CDLO'].str.cat (df['CATX'])
    print('done duration')
    return df

def catmap_mod(df):
    print("strart operation..............")
    dfdb1 = pd.read_csv (db)
    dfdb = dfdb1[['Code', 'Zone']]
    df0 = df.rename (columns=str.upper)
    ls = text2list (semcol)
    df1 = df0[ls]
    dc = text2dic (CAT)
    df1 = df1.assign (CAT='0')
    df1 = df1.assign (CATX='0')
    df1 = df1.assign (Code='0')
    df1['CAT'] = df1.apply (lambda x: getkey (dc, x.SUMMARY), axis=1)
    df1['CATX'] = df1.apply (lambda x: TS (x.SUMMARY), axis=1)
    df1['Code'] = df1.apply (lambda x: x.CUSTOMATTR15[0:5], axis=1)
    df2 = df1.merge (dfdb, on='Code')
    try:
        df3 = DateDiff(df2, "DUR", "LASTOCCURRENCE")
    except:
        df3 = datediff_ondf(df2, "DUR", 'LASTOCCURRENCE')
    df4 = extrafeat(df3)
    xdf = df4.replace (np.nan, 0)
    ndf = countifs(xdf, xdf['CUSTOMATTR15'], xdf['CUSTOMATTR15'], xdf['DURCAT'], xdf['DURCAT'])
    odf = countifs(ndf, xdf['EQUIPMENTKEY'], xdf['EQUIPMENTKEY'], xdf['DURCAT'], xdf['DURCAT'])
    odf.to_csv (os.getcwd () + "\\FINAL12.csv", index=False)
    return odf

def sort_rvmdup(df):
    df1 = df[~df['CATX'].isin(['other']) & ~df['CAT'].isin(['other'])]
    df1 = df1.sort_values(by=['CAT','CDLO'], ascending=True)
    df1 = df1.drop_duplicates(subset=['CDLOTECH'], inplace=False, ignore_index=True)
    #df2 = df1.groupby(['DURCAT','EQUIPMENTKEY','CAT'])['CUSTOMATTR15'].count()
    pvt = df1.pivot_table(index=['CUSTOMATTR15','CAT'], columns='DURCAT', values='cnt_x', aggfunc='sum').reset_index()
    ndf = pvt[(pvt['72H'] > 10) & (pvt['48H'] > 2)]
    return ndf

def fmtmsg_techwise(df, name_thread_col, ls_datacol, name_catcol, cat_text):
    lss = []
    heap = ''
    hp = ""
    hpx = ""
    for n in range(len(df)):
        code = df.loc[n, name_thread_col]
        cat = df.loc[n, name_catcol]
        if str(cat) == str(cat_text):
            for i in range(len(ls_datacol)):
                if hp == "":
                    hp = df.loc[n, ls_datacol[i]]
                else:
                    hp = hp + " | " + df.loc[n, ls_datacol[i]]
            hpx = code + ": " + hp
            if heap == "":
                heap = hpx
                hp = ""
            else:
                heap = heap + chr(10) + hpx
                hp = ""
    return heap

def tmsg(chatid,msg):
    TOK = "1176189570:AAEfPi9TIZIbnhWi4Ko6KQev2Iv7UbMw5js"
    url = "https://api.telegram.org/bot" + TOK + "/sendMessage?chat_id=" + str(chatid) + "&text=" + msg
    requests.get(url)
    return ""

def semqry():
    conn = cx_Oracle.connect ('SOC_READ','soc_read', 'ossam-cluster-scan.robi.com.bd:1721/RBPB.robi.com.bd')
    print (conn.version)
    agent = ['U2000 TX','Ericsson OSS','EricssonOSS','Huawei U2000 vEPC','Huawei U2020','LTE_BR1_5','MV36-PFM3-MIB','BusinessRule14','BusinessRule14_ERI_ABIP']
    cols = "SERIAL,NODE,AGENT,ALERTGROUP,SEVERITY,LOCALSECOBJ,X733EVENTTYPE,X733SPECIFICPROB,MANAGEDOBJCLASS,GEOINFO,CUSTOMATTR3,CUSTOMATTR5,CUSTOMATTR25,TTSEQUENCE,TTSTATUS,SRCDOMAIN,CUSTOMATTR26,OUTAGEDURATION,TALLY,ALARMDETAILS,EQUIPMENTKEY,CUSTOMATTR15,SUMMARY,LASTOCCURRENCE,CLEARTIMESTAMP"
    q1 = "SELECT " +  cols + " FROM SEMHEDB.ALERTS_STATUS WHERE "
    STDT = timedelt(-22)
    ENDT = timedelt(1)
    q2 = "LASTOCCURRENCE BETWEEN TO_DATE('" + STDT + "','DD-MM-YYYY HH24:MI:SS') AND TO_DATE('" + ENDT + "','DD-MM-YYYY HH24:MI:SS')"
    q3 = q1 + q2
    print(q3)
    print('starts: ', datetime.now())
    df = pd.read_sql(q3, con=conn)
    print ('ends: ', datetime.now())
    print(df.shape[0])
    print(df.columns)
    df1 = df[df['AGENT'].isin([agent])]
    print (df.shape[0])
    df1.to_csv(os.getcwd () + "\\SEMQRY.csv")
    return df1

def main(df):
    ls = ['2H', '12H']
    df1 = catmap_mod(df)
    df2 = sort_rvmdup(df1)
    print('2')
    df2.to_csv(os.getcwd () + "\\pvt.csv")
    df2 = df2.astype (str)
    G2 = "2G:" + chr (10) + fmtmsg_techwise (df2, 'CUSTOMATTR15', ls, 'CAT', '2') + chr (10) + chr (10)
    G3 = "3G:" + chr (10) + fmtmsg_techwise (df2, 'CUSTOMATTR15', ls, 'CAT', '3') + chr (10) + chr (10)
    G4 = "3G:" + chr (10) + fmtmsg_techwise (df2, 'CUSTOMATTR15', ls, 'CAT', '4') + chr (10) + chr (10)
    HD1 = "SITE FLUCTUATION COUNT" + chr (10) + "at " + tm + chr (10) + chr (10)
    HD2 = "Code : 2Hr | 12Hr" + chr (10) + chr (10)
    TR1 = "Note: sites fluctuates 2+ in last 2hr and 10+ in last 12hr"
    GG2 = "2G " + HD1 + HD2 + G2 + TR1
    msk = '-407548960'
    q = tmsg(msk, GG2)
    GG3 = "3G " + HD1 + HD2 + G3 + TR1
    q = tmsg (msk, GG3)
    GG4 = "4G " + HD1 + HD2 + G4 + TR1
    q = tmsg (msk, GG4)
    print('done')

#svpt = os.getcwd () + "\\OMTX.csv"
df = semqry()
xxx = main(df)
#svpt = os.getcwd () + "\\pvt.csv"
#df = pd.read_csv (svpt, low_memory=False)


import ast
import math
import pickle
import random
from copy import deepcopy
from typing import List, Tuple, Union

import numpy as np
import pandas as pd
from rl_coach.core_types import CsvDataset, Episode, Transition
from rl_coach.filters.filter import InputFilter
from rl_coach.logger import screen
from rl_coach.memories.memory import Memory, MemoryGranularity, MemoryParameters
from rl_coach.utils import ProgressBar, ReaderWriterLock

class DeepRacerMemoryParameters(MemoryParameters):
    def __init__(self):
        super().__init__()
        self.max_size = (MemoryGranularity.Transitions, 1000000)
        self.n_step = -1
        self.train_to_eval_ratio = 1  # for OPE we'll want a value < 1

    @property
    def path(self):
        return "markov.memories.deepracer_memory:DeepRacerMemory"

class DeepRacerMemory(Memory):
    """
    A replay buffer that stores episodes of transitions. The additional structure allows performing various
    calculations of total return and other values that depend on the sequential behavior of the transitions
    in the episode.
    """

    def __init__(
        self,
        max_size: Tuple[MemoryGranularity, int] = (MemoryGranularity.Transitions, 1000000),
        n_step=-1,
        train_to_eval_ratio: int = 1,
    ):
        """
        :param max_size: the maximum number of transitions or episodes to hold in the memory
        """
        super().__init__(max_size)
        self.n_step = n_step
        self._buffer = [Episode(n_step=self.n_step)]  # list of episodes
        self.transitions = []
        self._length = 1  # the episodic replay buffer starts with a single empty episode
        self._num_transitions = 0
        self._num_transitions_in_complete_episodes = 0
        self.last_training_set_episode_id = None  # used in batch-rl
        self.last_training_set_transition_id = None  # used in batch-rl
        self.train_to_eval_ratio = train_to_eval_ratio  # used in batch-rl
        self.evaluation_dataset_as_episodes = None
        self.evaluation_dataset_as_transitions = None

        self.frozen = False

    def length(self, lock: bool = False) -> int:
        """
        Get the number of episodes in the ER (even if they are not complete)
        """
        length = self._length
        if self._length is not 0 and self._buffer[-1].is_empty():
            length = self._length - 1

        return length

    def num_complete_episodes(self):
        """Get the number of complete episodes in ER"""
        length = self._length - 1

        return length

    def num_transitions(self):
        return self._num_transitions

    def num_transitions_in_complete_episodes(self):
        return self._num_transitions_in_complete_episodes

    def get_last_training_set_episode_id(self):
        return self.last_training_set_episode_id

    def sample(self, size: int, is_consecutive_transitions=False) -> List[Transition]:
        """
        Sample a batch of transitions from the replay buffer. If the requested size is larger than the number
        of samples available in the replay buffer then the batch will return empty.
        :param size: the size of the batch to sample
        :param is_consecutive_transitions: if set True, samples a batch of consecutive transitions.
        :return: a batch (list) of selected transitions from the replay buffer
        """
        if self.num_complete_episodes() >= 1:
            if is_consecutive_transitions:
                episode_idx = np.random.randint(0, self.num_complete_episodes())
                if self._buffer[episode_idx].length() <= size:
                    batch = self._buffer[episode_idx].transitions
                else:
                    transition_idx = np.random.randint(size, self._buffer[episode_idx].length())
                    batch = self._buffer[episode_idx].transitions[
                        transition_idx - size : transition_idx
                    ]
            else:
                transitions_idx = np.random.randint(
                    self.num_transitions_in_complete_episodes(), size=size
                )
                batch = [self.transitions[i] for i in transitions_idx]

        else:
            raise ValueError(
                "The episodic replay buffer cannot be sampled since there are no complete episodes yet. "
                "There is currently 1 episodes with {} transitions".format(self._buffer[0].length())
            )

        return batch

    def get_episode_for_transition(self, transition: Transition) -> (int, Episode):
        """
        Get the episode from which that transition came from.
        :param transition: The transition to lookup the episode for
        :return: (Episode number, the episode) or (-1, None) if could not find a matching episode.
        """

        for i, episode in enumerate(self._buffer):
            if transition in episode.transitions:
                return i, episode
        return -1, None

    def shuffle_episodes(self):
        """
        Shuffle all the complete episodes in the replay buffer, while deleting the last non-complete episode
        :return:
        """
        self.assert_not_frozen()

        # unlike the standard usage of the EpisodicExperienceReplay, where we always leave an empty episode after
        # the last full one, so that new transitions will have where to be added, in this case we delibrately remove
        # that empty last episode, as we are about to shuffle the memory, and we don't want it to be shuffled in
        self.remove_last_episode(lock=False)

        random.shuffle(self._buffer)
        self.transitions = [t for e in self._buffer for t in e.transitions]

        # create a new Episode for the next transitions to be placed into
        self._buffer.append(Episode(n_step=self.n_step))
        self._length += 1

    def get_shuffled_training_data_generator(self, size: int) -> List[Transition]:
        """
        Get an generator for iterating through the shuffled replay buffer, for processing the data in epochs.
        If the requested size is larger than the number of samples available in the replay buffer then the batch will
        return empty. The last returned batch may be smaller than the size requested, to accommodate for all the
        transitions in the replay buffer.

        :param size: the size of the batch to return
        :return: a batch (list) of selected transitions from the replay buffer
        """
        shuffled_transition_indices = list(range(self.last_training_set_transition_id))
        random.shuffle(shuffled_transition_indices)

        # The last batch drawn will usually be < batch_size (=the size variable)
        for i in range(math.ceil(len(shuffled_transition_indices) / size)):
            sample_data = [
                self.transitions[j] for j in shuffled_transition_indices[i * size : (i + 1) * size]
            ]

            yield sample_data

    def get_all_complete_episodes_transitions(self) -> List[Transition]:
        """
        Get all the transitions from all the complete episodes in the buffer
        :return: a list of transitions
        """
        return self.transitions[: self.num_transitions_in_complete_episodes()]

    def get_all_complete_episodes(self) -> List[Episode]:
        """
        Get all the transitions from all the complete episodes in the buffer
        :return: a list of transitions
        """
        return self.get_all_complete_episodes_from_to(0, self.num_complete_episodes())

    def get_all_complete_episodes_from_to(self, start_episode_id, end_episode_id) -> List[Episode]:
        """
        Get all the transitions from all the complete episodes in the buffer matching the given episode range
        :return: a list of transitions
        """
        return self._buffer[start_episode_id:end_episode_id]

    def _enforce_max_length(self) -> None:
        """
        Make sure that the size of the replay buffer does not pass the maximum size allowed.
        If it passes the max size, the oldest episode in the replay buffer will be removed.
        :return: None
        """
        granularity, size = self.max_size
        if granularity == MemoryGranularity.Transitions:
            while size != 0 and self.num_transitions() > size:
                self.remove_first_episode(lock=False)
        elif granularity == MemoryGranularity.Episodes:
            while self.length() > size:
                self.remove_first_episode(lock=False)

    def _update_episode(self, episode: Episode) -> None:
        episode.update_transitions_rewards_and_bootstrap_data()

    def verify_last_episode_is_closed(self) -> None:
        """
        Verify that there is no open episodes in the replay buffer
        :return: None
        """
        last_episode = self.get(-1, False)
        if last_episode and last_episode.length() > 0:
            self.close_last_episode(lock=False)

    def close_last_episode(self, lock=True) -> None:
        """
        Close the last episode in the replay buffer and open a new one
        :return: None
        """

        last_episode = self._buffer[-1]

        self._num_transitions_in_complete_episodes += last_episode.length()
        self._length += 1

        # create a new Episode for the next transitions to be placed into
        self._buffer.append(Episode(n_step=self.n_step))

        # if update episode adds to the buffer, a new Episode needs to be ready first
        # it would be better if this were less state full
        self._update_episode(last_episode)

        self._enforce_max_length()

    def store(self, transition: Transition) -> None:
        """
        Store a new transition in the memory. If the transition game_over flag is on, this closes the episode and
        creates a new empty episode.
        Warning! using the episodic memory by storing individual transitions instead of episodes will use the default
        Episode class parameters in order to create new episodes.
        :param transition: a transition to store
        :return: None
        """
        self.assert_not_frozen()

        # Calling super.store() so that in case a memory backend is used, the memory backend can store this transition.
        super().store(transition)

        if len(self._buffer) == 0:
            self._buffer.append(Episode(n_step=self.n_step))
        last_episode = self._buffer[-1]
        last_episode.insert(transition)
        self.transitions.append(transition)
        self._num_transitions += 1
        if transition.game_over:
            self.close_last_episode(False)

        self._enforce_max_length()

    def store_episode(self, episode: Episode, lock: bool = False) -> None:
        """
        Store a new episode in the memory.
        :param episode: the new episode to store
        :return: None
        """
        self.assert_not_frozen()

        # Calling super.store() so that in case a memory backend is used, the memory backend can store this episode.
        super().store_episode(episode)

        if self._buffer[-1].length() == 0:
            self._buffer[-1] = episode
        else:
            self._buffer.append(episode)
        self.transitions.extend(episode.transitions)
        self._num_transitions += episode.length()
        self.close_last_episode(False)

    def _remove_episode(self, episode_index: int) -> None:
        """
        Remove either the first or the last index
        :param episode_index: the index of the episode to remove (either 0 or -1)
        :return: None
        """
        self.assert_not_frozen()
        assert episode_index == 0 or episode_index == -1, (
            "_remove_episode only supports removing the first or the last " "episode"
        )

        if len(self._buffer) > 0:
            episode_length = self._buffer[episode_index].length()
            self._length -= 1
            self._num_transitions -= episode_length
            self._num_transitions_in_complete_episodes -= episode_length
            if episode_index == 0:
                del self.transitions[:episode_length]
            else:  # episode_index = -1
                del self.transitions[-episode_length:]
            del self._buffer[episode_index]

    def remove_first_episode(self, lock: bool = True) -> None:
        """
        Remove the first episode (even if it is not complete yet)
        :param lock: if true, will lock the readers writers lock. this can cause a deadlock if an inheriting class
                     locks and then calls store with lock = True
        :return: None
        """
        self._remove_episode(0)

    def remove_last_episode(self, lock: bool = True) -> None:
        """
        Remove the last episode (even if it is not complete yet)
        :param lock: if true, will lock the readers writers lock. this can cause a deadlock if an inheriting class
                     locks and then calls store with lock = True
        :return: None
        """
        self._remove_episode(-1)

    # for API compatibility
    def get(self, episode_index: int, lock: bool = True) -> Union[None, Episode]:
        """
        Returns the episode in the given index. If the episode does not exist, returns None instead.
        :param episode_index: the index of the episode to return
        :return: the corresponding episode
        """
        if self.length() == 0 or episode_index >= self.length():
            episode = None
        else:
            episode = self._buffer[episode_index]

        return episode

    def get_last_complete_episode(self) -> Union[None, Episode]:
        """
        Returns the last complete episode in the memory or None if there are no complete episodes
        :return: None or the last complete episode
        """
        last_complete_episode_index = self.num_complete_episodes() - 1
        episode = None
        if last_complete_episode_index >= 0:
            episode = self.get(last_complete_episode_index)

        return episode

    def clean(self) -> None:
        """
        Clean the memory by removing all the episodes
        :return: None
        """
        self.assert_not_frozen()

        self.transitions = []
        self._buffer = [Episode(n_step=self.n_step)]
        self._length = 1
        self._num_transitions = 0
        self._num_transitions_in_complete_episodes = 0

    def mean_reward(self) -> np.ndarray:
        """
        Get the mean reward in the replay buffer
        :return: the mean reward
        """
        mean = np.mean([transition.reward for transition in self.transitions])

        return mean

    def load_csv(self, csv_dataset: CsvDataset, input_filter: InputFilter) -> None:
        """
        Restore the replay buffer contents from a csv file.
        The csv file is assumed to include a list of transitions.
        :param csv_dataset: A construct which holds the dataset parameters
        :param input_filter: A filter used to filter the CSV data before feeding it to the memory.
        """
        self.assert_not_frozen()

        df = pd.read_csv(csv_dataset.filepath)
        if len(df) > self.max_size[1]:
            screen.warning(
                "Warning! The number of transitions to load into the replay buffer ({}) is "
                "bigger than the max size of the replay buffer ({}). The excessive transitions will "
                "not be stored.".format(len(df), self.max_size[1])
            )

        episode_ids = df["episode_id"].unique()
        progress_bar = ProgressBar(len(episode_ids))
        state_columns = [col for col in df.columns if col.startswith("state_feature")]

        for e_id in episode_ids:
            progress_bar.update(e_id)
            df_episode_transitions = df[df["episode_id"] == e_id]
            input_filter.reset()

            if len(df_episode_transitions) < 2:
                # we have to have at least 2 rows in each episode for creating a transition
                continue

            episode = Episode()
            transitions = []
            for (_, current_transition), (_, next_transition) in zip(
                df_episode_transitions[:-1].iterrows(), df_episode_transitions[1:].iterrows()
            ):
                state = np.array([current_transition[col] for col in state_columns])
                next_state = np.array([next_transition[col] for col in state_columns])

                transitions.append(
                    Transition(
                        state={"observation": state},
                        action=int(current_transition["action"]),
                        reward=current_transition["reward"],
                        next_state={"observation": next_state},
                        game_over=False,
                        info={
                            "all_action_probabilities": ast.literal_eval(
                                current_transition["all_action_probabilities"]
                            )
                        },
                    ),
                )

            transitions = input_filter.filter(transitions, deep_copy=False)
            for t in transitions:
                episode.insert(t)

            # Set the last transition to end the episode
            if csv_dataset.is_episodic:
                episode.get_last_transition().game_over = True

            self.store_episode(episode)

        # close the progress bar
        progress_bar.update(len(episode_ids))
        progress_bar.close()

    def freeze(self):
        """
        Freezing the replay buffer does not allow any new transitions to be added to the memory.
        Useful when working with a dataset (e.g. batch-rl or imitation learning).
        :return: None
        """
        self.frozen = True

    def assert_not_frozen(self):
        """
        Check that the memory is not frozen, and can be changed.
        :return:
        """
        assert self.frozen is False, "Memory is frozen, and cannot be changed."

    def prepare_evaluation_dataset(self):
        """
        Gather the memory content that will be used for off-policy evaluation in episodes and transitions format
        :return:
        """
        self._split_training_and_evaluation_datasets()
        self.evaluation_dataset_as_episodes = deepcopy(
            self.get_all_complete_episodes_from_to(
                self.get_last_training_set_episode_id() + 1, self.num_complete_episodes()
            )
        )

        if len(self.evaluation_dataset_as_episodes) == 0:
            raise ValueError(
                "train_to_eval_ratio is too high causing the evaluation set to be empty. "
                "Consider decreasing its value."
            )

        self.evaluation_dataset_as_transitions = [
            t for e in self.evaluation_dataset_as_episodes for t in e.transitions
        ]

    def _split_training_and_evaluation_datasets(self):
        """
        If the data in the buffer was not split to training and evaluation yet, split it accordingly.
        :return: None
        """

        if self.last_training_set_transition_id is None:
            if self.train_to_eval_ratio < 0 or self.train_to_eval_ratio >= 1:
                raise ValueError("train_to_eval_ratio should be in the (0, 1] range.")

            transition = self.transitions[
                round(self.train_to_eval_ratio * self.num_transitions_in_complete_episodes())
            ]
            episode_num, episode = self.get_episode_for_transition(transition)
            self.last_training_set_episode_id = episode_num
            self.last_training_set_transition_id = len(
                [
                    t
                    for e in self.get_all_complete_episodes_from_to(
                        0, self.last_training_set_episode_id + 1
                    )
                    for t in e
                ]
            )

    def save(self, file_path: str) -> None:
        """
        Save the replay buffer contents to a pickle file
        :param file_path: the path to the file that will be used to store the pickled transitions
        """
        with open(file_path, "wb") as file:
            pickle.dump(self.get_all_complete_episodes(), file)

    def load_pickled(self, file_path: str) -> None:
        """
        Restore the replay buffer contents from a pickle file.
        The pickle file is assumed to include a list of transitions.
        :param file_path: The path to a pickle file to restore
        """
        self.assert_not_frozen()

        with open(file_path, "rb") as file:
            episodes = pickle.load(file)
            num_transitions = sum([len(e.transitions) for e in episodes])
            if num_transitions > self.max_size[1]:
                screen.warning(
                    "Warning! The number of transition to load into the replay buffer ({}) is "
                    "bigger than the max size of the replay buffer ({}). The excessive transitions will "
                    "not be stored.".format(num_transitions, self.max_size[1])
                )

            progress_bar = ProgressBar(len(episodes))
            for episode_idx, episode in enumerate(episodes):
                self.store_episode(episode)

                # print progress
                progress_bar.update(episode_idx)

            progress_bar.close()


# coding: utf-8
import numpy as np
from PIL import Image, ImageDraw
from math import sin,cos
from numpngw import write_apng

W,H = 1024,1024

COLOR_BLACK   = (0x00, 0x00, 0x00, 0x00)
COLOR_WHITE   = (0xF0, 0xF0, 0xE0)
COLOR_BLUE    = (0x0D, 0x36, 0xFF)
COLOR_BLYNK   = (0x2E, 0xFF, 0xB9)
COLOR_RED     = (0xFF, 0x10, 0x08)
COLOR_MAGENTA = (0xA7, 0x00, 0xFF)

def fr(color):
    im = Image.new('RGBA', (W,H))
    surface = ImageDraw.Draw(im)
    surface.ellipse((10,10,W-20,H-20), fill=color)
    del surface
    return np.array(im.resize((16,16), Image.ANTIALIAS))

def fade(color):
    frames = []
    for i in range(25):
        frames.append(fr(color + tuple([int(i*(255.0/25))])))
    for i in range(25):
        frames.append(fr(color + tuple([255-int(i*(255.0/25))])))
    return frames

path = "./docs/images/states/"
write_apng(path + "0.png", map(fr,[COLOR_BLUE, COLOR_BLACK]), delay=[50, 500])
write_apng(path + "1.png", map(fr,[COLOR_BLUE, COLOR_BLACK]), delay=[200, 200])
write_apng(path + "2.png", map(fr,[COLOR_BLYNK, COLOR_BLACK]), delay=[50, 500])
write_apng(path + "3.png", map(fr,[COLOR_BLYNK, COLOR_BLACK]), delay=[100, 100])
write_apng(path + "4.png", fade(COLOR_BLYNK), delay=100)
write_apng(path + "5.png", map(fr,[COLOR_MAGENTA, COLOR_BLACK]), delay=[50, 50])
write_apng(path + "6.png", map(fr,[COLOR_RED, COLOR_BLACK, COLOR_RED, COLOR_BLACK]), delay=[80, 100, 80, 1000])
write_apng(path + "7.png", fade(COLOR_WHITE), delay=50)
write_apng(path + "8.png", map(fr,[COLOR_WHITE, COLOR_BLACK]), delay=[100, 100])


import argparse
import torch
import time
import json
import numpy as np
import math
import random
import xml.etree.ElementTree as ET
from subprocess import check_output
from subprocess import CalledProcessError

np.random.seed(1337)
random.seed(1337)
torch.manual_seed(1337)
# torch.cuda.manual_seed(1337)

class Model(torch.nn.Module):
    def __init__(self, gen_emb, domain_emb, num_classes=3, dropout=0.5, crf=False):
        super(Model, self).__init__()
        self.gen_embedding = torch.nn.Embedding(gen_emb.shape[0], gen_emb.shape[1])
        self.gen_embedding.weight=torch.nn.Parameter(torch.from_numpy(gen_emb), requires_grad=False)
        self.domain_embedding = torch.nn.Embedding(domain_emb.shape[0], domain_emb.shape[1])
        self.domain_embedding.weight=torch.nn.Parameter(torch.from_numpy(domain_emb), requires_grad=False)

        self.conv1=torch.nn.Conv1d(gen_emb.shape[1]+domain_emb.shape[1], 128, 5, padding=2 )
        self.conv2=torch.nn.Conv1d(gen_emb.shape[1]+domain_emb.shape[1], 128, 3, padding=1 )
        self.dropout=torch.nn.Dropout(dropout)

        self.conv3=torch.nn.Conv1d(256, 256, 5, padding=2)
        self.conv4=torch.nn.Conv1d(256, 256, 5, padding=2)
        self.conv5=torch.nn.Conv1d(256, 256, 5, padding=2)
        self.linear_ae=torch.nn.Linear(256, num_classes)
        self.crf_flag=crf
        if self.crf_flag:
            from allennlp.modules import ConditionalRandomField
            self.crf=ConditionalRandomField(num_classes)

    def forward(self, x, x_len, x_mask, x_tag=None, testing=False):
        x_emb=torch.cat((self.gen_embedding(x), self.domain_embedding(x) ), dim=2)
        x_emb=self.dropout(x_emb).transpose(1, 2)
        x_conv=torch.nn.functional.relu(torch.cat((self.conv1(x_emb), self.conv2(x_emb)), dim=1) )
        x_conv=self.dropout(x_conv)
        x_conv=torch.nn.functional.relu(self.conv3(x_conv) )
        x_conv=self.dropout(x_conv)
        x_conv=torch.nn.functional.relu(self.conv4(x_conv) )
        x_conv=self.dropout(x_conv)
        x_conv=torch.nn.functional.relu(self.conv5(x_conv) )
        x_conv=x_conv.transpose(1, 2)
        x_logit=self.linear_ae(x_conv)
        if testing:
            if self.crf_flag:
                score=self.crf.viterbi_tags(x_logit, x_mask)
            else:
                x_logit=x_logit.transpose(2, 0)
                score=torch.nn.functional.log_softmax(x_logit).transpose(2, 0)
        else:
            if self.crf_flag:
                score=-self.crf(x_logit, x_tag, x_mask)
            else:
                x_logit=torch.nn.utils.rnn.pack_padded_sequence(x_logit, x_len, batch_first=True)
                score=torch.nn.functional.nll_loss(torch.nn.functional.log_softmax(x_logit.data), x_tag.data)
        return score

def label_rest_xml(fn, output_fn, corpus, label):
    dom=ET.parse(fn)
    root=dom.getroot()
    pred_y=[]
    for zx, sent in enumerate(root.iter("sentence") ) :
        tokens=corpus[zx]
        lb=label[zx]
        opins=ET.Element("Opinions")
        token_idx, pt, tag_on=0, 0, False
        start, end=-1, -1
        for ix, c in enumerate(sent.find('text').text):
            if token_idx<len(tokens) and pt>=len(tokens[token_idx] ):
                pt=0
                token_idx+=1

            if token_idx<len(tokens) and lb[token_idx]==1 and pt==0 and c!=' ':
                if tag_on:
                    end=ix
                    tag_on=False
                    opin=ET.Element("Opinion")
                    opin.attrib['target']=sent.find('text').text[start:end]
                    opin.attrib['from']=str(start)
                    opin.attrib['to']=str(end)
                    opins.append(opin)
                start=ix
                tag_on=True
            elif token_idx<len(tokens) and lb[token_idx]==2 and pt==0 and c!=' ' and not tag_on:
                start=ix
                tag_on=True
            elif token_idx<len(tokens) and (lb[token_idx]==0 or lb[token_idx]==1) and tag_on and pt==0:
                end=ix
                tag_on=False
                opin=ET.Element("Opinion")
                opin.attrib['target']=sent.find('text').text[start:end]
                opin.attrib['from']=str(start)
                opin.attrib['to']=str(end)
                opins.append(opin)
            elif token_idx>=len(tokens) and tag_on:
                end=ix
                tag_on=False
                opin=ET.Element("Opinion")
                opin.attrib['target']=sent.find('text').text[start:end]
                opin.attrib['from']=str(start)
                opin.attrib['to']=str(end)
                opins.append(opin)
            if c==' ':
                pass
            elif tokens[token_idx][pt:pt+2]=='``' or tokens[token_idx][pt:pt+2]=="''":
                pt+=2
            else:
                pt+=1
        if tag_on:
            tag_on=False
            end=len(sent.find('text').text)
            opin=ET.Element("Opinion")
            opin.attrib['target']=sent.find('text').text[start:end]
            opin.attrib['from']=str(start)
            opin.attrib['to']=str(end)
            opins.append(opin)
        sent.append(opins )
    dom.write(output_fn)

def label_laptop_xml(fn, output_fn, corpus, label):
    dom=ET.parse(fn)
    root=dom.getroot()
    pred_y=[]
    for zx, sent in enumerate(root.iter("sentence") ) :
        tokens=corpus[zx]
        lb=label[zx]
        opins=ET.Element("aspectTerms")
        token_idx, pt, tag_on=0, 0, False
        start, end=-1, -1
        for ix, c in enumerate(sent.find('text').text):
            if token_idx<len(tokens) and pt>=len(tokens[token_idx] ):
                pt=0
                token_idx+=1

            if token_idx<len(tokens) and lb[token_idx]==1 and pt==0 and c!=' ':
                if tag_on:
                    end=ix
                    tag_on=False
                    opin=ET.Element("aspectTerm")
                    opin.attrib['term']=sent.find('text').text[start:end]
                    opin.attrib['from']=str(start)
                    opin.attrib['to']=str(end)
                    opins.append(opin)
                start=ix
                tag_on=True
            elif token_idx<len(tokens) and lb[token_idx]==2 and pt==0 and c!=' ' and not tag_on:
                start=ix
                tag_on=True
            elif token_idx<len(tokens) and (lb[token_idx]==0 or lb[token_idx]==1) and tag_on and pt==0:
                end=ix
                tag_on=False
                opin=ET.Element("aspectTerm")
                opin.attrib['term']=sent.find('text').text[start:end]
                opin.attrib['from']=str(start)
                opin.attrib['to']=str(end)
                opins.append(opin)
            elif token_idx>=len(tokens) and tag_on:
                end=ix
                tag_on=False
                opin=ET.Element("aspectTerm")
                opin.attrib['term']=sent.find('text').text[start:end]
                opin.attrib['from']=str(start)
                opin.attrib['to']=str(end)
                opins.append(opin)
            if c==' ' or ord(c)==160:
                pass
            elif tokens[token_idx][pt:pt+2]=='``' or tokens[token_idx][pt:pt+2]=="''":
                pt+=2
            else:
                pt+=1
        if tag_on:
            tag_on=False
            end=len(sent.find('text').text)
            opin=ET.Element("aspectTerm")
            opin.attrib['term']=sent.find('text').text[start:end]
            opin.attrib['from']=str(start)
            opin.attrib['to']=str(end)
            opins.append(opin)
        sent.append(opins )
    dom.write(output_fn)

def test(model, test_X, raw_X, domain, command, template, batch_size=128, crf=False):
    pred_y=np.zeros((test_X.shape[0], 83), np.int16)
    model.eval()
    for offset in range(0, test_X.shape[0], batch_size):
        batch_test_X_len=np.sum(test_X[offset:offset+batch_size]!=0, axis=1)
        batch_idx=batch_test_X_len.argsort()[::-1]
        batch_test_X_len=batch_test_X_len[batch_idx]
        batch_test_X_mask=(test_X[offset:offset+batch_size]!=0)[batch_idx].astype(np.uint8)
        batch_test_X=test_X[offset:offset+batch_size][batch_idx]
        # batch_test_X_mask=torch.autograd.Variable(torch.from_numpy(batch_test_X_mask).long().cuda() )
        batch_test_X_mask=torch.autograd.Variable(torch.from_numpy(batch_test_X_mask).long())
        # batch_test_X = torch.autograd.Variable(torch.from_numpy(batch_test_X).long().cuda() )
        batch_test_X = torch.autograd.Variable(torch.from_numpy(batch_test_X).long())
        batch_pred_y=model(batch_test_X, batch_test_X_len, batch_test_X_mask, testing=True)
        r_idx=batch_idx.argsort()
        if crf:
            batch_pred_y=[batch_pred_y[idx] for idx in r_idx]
            for ix in range(len(batch_pred_y) ):
                for jx in range(len(batch_pred_y[ix]) ):
                    pred_y[offset+ix,jx]=batch_pred_y[ix][jx]
        else:
            batch_pred_y=batch_pred_y.data.cpu().numpy().argmax(axis=2)[r_idx]
            pred_y[offset:offset+batch_size,:batch_pred_y.shape[1]]=batch_pred_y
    model.train()
    assert len(pred_y)==len(test_X)

    command=command.split()
    print("COMMAND: ")
    print(command)
    if domain=='restaurant':
        label_rest_xml(template, command[8], raw_X, pred_y)
        acc=check_output(command ).split()
        print(acc)
        return float(acc[9][10:])
    elif domain=='laptop':
        label_laptop_xml(template, command[6], raw_X, pred_y)
        try:
            acc=check_output(command).split()
        except CalledProcessError as e:
            print(e.output)
        print(acc)
        return float(acc[15])

def evaluate(runs, data_dir, model_dir, domain, command, template):
    ae_data=np.load(data_dir+domain+".npz")
    with open(data_dir+domain+"_raw_test.json") as f:
        raw_X=json.load(f)
    results=[]
    for r in range(runs):
        model=torch.load(model_dir+domain+str(r) )
        result=test(model, ae_data['test_X'], raw_X, domain, command, template, crf=False)
        results.append(result)
    print(sum(results)/len(results) )

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--runs', type=int, default=2)
    parser.add_argument('--data_dir', type=str, default="/sem8/COL772-NLP/DE-CNN/data/prep_data/")
    parser.add_argument('--model_dir', type=str, default="/sem8/COL772-NLP/DE-CNN/script/model/")
    parser.add_argument('--domain', type=str, default="laptop")

    args = parser.parse_args()

    if args.domain=='restaurant':
        command="java --add-modules java.xml.bind -cp /sem8/COL772-NLP/DE-CNN/script/A.jar absa16.Do Eval -prd /sem8/COL772-NLP/DE-CNN/data/official_data/pred.xml -gld /sem8/COL772-NLP/DE-CNN/data/official_data/EN_REST_SB1_TEST.xml.gold -evs 2 -phs A -sbt SB1"
        template="/sem8/COL772-NLP/DE-CNN/data/official_data/EN_REST_SB1_TEST.xml.A"
    elif args.domain=='laptop':
        command="java --add-modules java.xml.bind -cp /sem8/COL772-NLP/DE-CNN/script/eval.jar Main.Aspects /sem8/COL772-NLP/DE-CNN/data/official_data/pred.xml /sem8/COL772-NLP/DE-CNN/data/official_data/Laptops_Test_Gold.xml"
        template="/sem8/COL772-NLP/DE-CNN/data/official_data/Laptops_Test_Data_PhaseA.xml"

    evaluate(args.runs, args.data_dir, args.model_dir, args.domain, command, template)


from __future__ import print_function, division
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.models import load_model as tf_load_model
from tensorflow.keras.preprocessing.image import img_to_array
from utils.box_utils import decode, decode_landm
from layers.functions.prior_box import PriorBox
from utils.nms.py_cpu_nms import py_cpu_nms
from models.retinaface import RetinaFace
import torch.backends.cudnn as cudnn
from matplotlib import pyplot as plt
from data import cfg_mnet, cfg_re50
from torch.autograd import Variable
from torchvision import transforms
from scipy.signal import lfilter
from itertools import chain
from PIL import Image
import torch.nn as nn
from tqdm import tqdm
import pandas as pd
import numpy as np
import torchvision
import argparse
import torch
import cv2
import gc

parser = argparse.ArgumentParser()
parser.add_argument('--img', dest='img', action='store', )
args = parser.parse_args()

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Torch CUDA available: {}".format(torch.cuda.is_available()))
cudnn.benchmark = True

def free_cache(model=None):
    print('Cleaning CUDA cache...')
    if model: del model
    gc.collect()
    torch.cuda.empty_cache()

def check_keys(model, pretrained_state_dict):
    ckpt_keys = set(pretrained_state_dict.keys())
    model_keys = set(model.state_dict().keys())
    used_pretrained_keys = model_keys & ckpt_keys
    unused_pretrained_keys = ckpt_keys - model_keys
    assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'
    return True

def remove_prefix(state_dict, prefix):
    f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x
    return {f(key): value for key, value in state_dict.items()}

def load_retina_model():
    print("Loading face detector model...")
    torch.set_grad_enabled(False)
    cfg = cfg_re50
    model = RetinaFace(cfg=cfg, phase='test')
    if not torch.cuda.is_available():
        pretrained_dict = torch.load('models/Resnet50_Final.pth', map_location=lambda storage, loc: storage)
    else:
        dev = torch.cuda.current_device()
        pretrained_dict = torch.load('models/Resnet50_Final.pth', map_location=lambda storage, loc: storage.cuda(dev))
    if "state_dict" in pretrained_dict.keys():
        pretrained_dict = remove_prefix(pretrained_dict['state_dict'], 'module.')
    else:
        pretrained_dict = remove_prefix(pretrained_dict, 'module.')
    check_keys(model, pretrained_dict)
    model.load_state_dict(pretrained_dict, strict=False)
    model = model.to(device)
    model.eval()
    return model

def load_mask_model():
    print("Loading mask detector model...")
    model = tf_load_model("models/mask_detector.model")
    return model

def face_detect(model, image):
    print("Detecting faces...")
    cfg = cfg_re50
    resize = 1

    img_raw = image
    img = np.float32(img_raw)

    im_height, im_width, _ = img.shape
    scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
    img -= (104, 117, 123)
    img = img.transpose(2, 0, 1)
    img = torch.from_numpy(img).unsqueeze(0)
    img = img.to(device)
    scale = scale.to(device)

    loc, conf, _ = model(img)  # forward pass

    priorbox = PriorBox(cfg, image_size=(im_height, im_width))
    priors = priorbox.forward()
    priors = priors.to(device)
    prior_data = priors.data
    boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance'])
    boxes = boxes * scale / resize
    boxes = boxes.cpu().numpy()
    scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
    scale1 = torch.Tensor([img.shape[3], img.shape[2], img.shape[3], img.shape[2],
                            img.shape[3], img.shape[2], img.shape[3], img.shape[2],
                            img.shape[3], img.shape[2]])
    scale1 = scale1.to(device)

    # ignore low scores
    inds = np.where(scores > 0.02)[0]
    boxes = boxes[inds]
    scores = scores[inds]

    # keep top-K before NMS
    order = scores.argsort()[::-1][:5000]
    boxes = boxes[order]
    scores = scores[order]

    # do NMS
    dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
    keep = py_cpu_nms(dets, 0.4)
    dets = dets[keep, :]

    # keep top-K faster NMS
    dets = dets[:750, :]

    # show image
    bb = []
    for b in dets:
        if b[4] < 0.6:
            continue
        b = list(map(int, b))
        bb.append(b)
    return bb

def mask_detect(model, image, bboxes):
    print('Detecting masks...')
    mask_list = []
    for bb in tqdm(bboxes):
        for e, b in enumerate(bb):
            if b < 0:
                bb[e] = 0
        raw_face = image[bb[1]:bb[3], bb[0]:bb[2]]
        img_UMat = cv2.UMat(raw_face)
        face = cv2.cvtColor(img_UMat, cv2.COLOR_BGR2RGB)
        face = cv2.resize(face, (224, 224))
        face = cv2.UMat.get(face)
        face = img_to_array(face)
        face = preprocess_input(face)
        face = np.expand_dims(face, axis=0)
        (mask, withoutMask) = model.predict(face)[0]
        label = "Mask" if mask > withoutMask else "No Mask"
        mask_list.append(label)
    return mask_list

def load_fairface_model():
    print("Loading age, gender, and race models...")
    model = torchvision.models.resnet34(pretrained=True)
    model.fc = nn.Linear(model.fc.in_features, 18)
    model.load_state_dict(torch.load('models/res34_fair_align_multi_7_20190809.pt'))
    model = model.to(device)
    model.eval()
    return model

def predict_age_gender_race(model, img, bboxes):
    print('Classifying genders, ages, and races...')
    trans = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    # list within a list. Each sublist contains scores for all races. Take max for predicted race
    face_names = []
    race_scores_fair = []
    gender_scores_fair = []
    age_scores_fair = []
    race_preds_fair = []
    gender_preds_fair = []
    age_preds_fair = []

    for i, bb in enumerate(tqdm(bboxes)):
        for e, b in enumerate(bb):
            if b < 0:
                bb[e] = 0
        raw_face = img[bb[1]:bb[3], bb[0]:bb[2]]
        #img_UMat = cv2.UMat(raw_face)
        #raw_face = cv2.cvtColor(img_UMat, cv2.COLOR_BGR2RGB)
        #face = cv2.UMat.get(raw_face)

        image = trans(raw_face)
        image = image.view(1, 3, 224, 224)  # reshape image to match model dimensions (1 batch size)
        image = image.to(device)

        # fair
        outputs = model(image)
        outputs = outputs.cpu().detach().numpy()
        outputs = np.squeeze(outputs)

        race_outputs = outputs[:7]
        gender_outputs = outputs[7:9]
        age_outputs = outputs[9:18]

        race_score = np.exp(race_outputs) / np.sum(np.exp(race_outputs))
        gender_score = np.exp(gender_outputs) / np.sum(np.exp(gender_outputs))
        age_score = np.exp(age_outputs) / np.sum(np.exp(age_outputs))

        face_names.append(i)
        race_pred = np.argmax(race_score)
        gender_pred = np.argmax(gender_score)
        age_pred = np.argmax(age_score)

        race_scores_fair.append(race_score)
        gender_scores_fair.append(gender_score)
        age_scores_fair.append(age_score)

        race_preds_fair.append(race_pred)
        gender_preds_fair.append(gender_pred)
        age_preds_fair.append(age_pred)

    result = pd.DataFrame([face_names,
                           race_preds_fair,
                           gender_preds_fair,
                           age_preds_fair,
                           race_scores_fair,
                           gender_scores_fair,
                           age_scores_fair,
                           bboxes]).T
    result.columns = ['index',
                      'race_preds_fair',
                      'gender_preds_fair',
                      'age_preds_fair',
                      'race_scores_fair',
                      'gender_scores_fair',
                      'age_scores_fair',
                      "bbox"]

    result.loc[result['race_preds_fair'] == 0, 'race'] = 'White'
    result.loc[result['race_preds_fair'] == 1, 'race'] = 'Black'
    result.loc[result['race_preds_fair'] == 2, 'race'] = 'Latino_Hispanic'
    result.loc[result['race_preds_fair'] == 3, 'race'] = 'East Asian'
    result.loc[result['race_preds_fair'] == 4, 'race'] = 'Southeast Asian'
    result.loc[result['race_preds_fair'] == 5, 'race'] = 'Indian'
    result.loc[result['race_preds_fair'] == 6, 'race'] = 'Middle Eastern'

    # gender
    result.loc[result['gender_preds_fair'] == 0, 'gender'] = 'Male'
    result.loc[result['gender_preds_fair'] == 1, 'gender'] = 'Female'

    # age
    result.loc[result['age_preds_fair'] == 0, 'age'] = '2'
    result.loc[result['age_preds_fair'] == 1, 'age'] = '7'
    result.loc[result['age_preds_fair'] == 2, 'age'] = '15'
    result.loc[result['age_preds_fair'] == 3, 'age'] = '25'
    result.loc[result['age_preds_fair'] == 4, 'age'] = '35'
    result.loc[result['age_preds_fair'] == 5, 'age'] = '45'
    result.loc[result['age_preds_fair'] == 6, 'age'] = '55'
    result.loc[result['age_preds_fair'] == 7, 'age'] = '65'
    result.loc[result['age_preds_fair'] == 8, 'age'] = '75'
    return result

def create_log_and_image(id_, img, bboxes, agr_table, mask_list):
    print("Creating a bounding box image...")
    shape = img.shape
    new_img = cv2.UMat(img)
    for i, bb in enumerate(tqdm(bboxes)):
        color = (0, 255, 0) if mask_list[i] == "Mask" else (0, 0, 255)
        cv2.rectangle(new_img, (bb[0], bb[1]), (bb[2], bb[3]), color, thickness=2)
        contour_top, contour_bot = get_contours(bb, shape)
        age = get_age(agr_table['age_scores_fair'][i])
        cv2.putText(img=new_img, text='{} {}'.format(agr_table['gender'][i], age), org=(contour_top[0], contour_top[1]),
                    fontScale=1, color=color, thickness=2, lineType=cv2.LINE_4, fontFace=cv2.FONT_HERSHEY_PLAIN)
        cv2.putText(img=new_img, text=agr_table['race'][i], org=(contour_bot[0], contour_bot[1]), fontScale=1,
                    color=color, thickness=2, lineType=cv2.LINE_4, fontFace=cv2.FONT_HERSHEY_PLAIN)

    final_img = cv2.UMat.get(new_img)
    cv2.imwrite("output/output_{}.jpg".format(id_), final_img)
    print("Saved image at output/output_{}".format(id_))

    print("Saving log...")
    mask_table = pd.DataFrame(mask_list, columns=['mask'])
    result = pd.concat([agr_table, mask_table], axis=1)
    result[['index', 'race', 'gender', 'age', 'bbox', 'mask']].to_csv("output/log_{}.csv".format(id_), index=False)
    print("Saved results at log_{}.csv".format(id_))

def get_age(results):
    res = [np.linspace(results[0], results[1], 5), np.linspace(results[1], results[2], 10),
           np.linspace(results[2], results[3], 10), np.linspace(results[3], results[4], 10),
           np.linspace(results[4], results[5], 10), np.linspace(results[5], results[6], 10),
           np.linspace(results[6], results[7], 10), np.linspace(results[7], results[8], 10)]
    arr = np.array(list(chain.from_iterable(res)))

    n = 3
    b = [1.0 / n] * n
    a = 1
    yy = lfilter(b, a, arr)
    age = np.argmax(yy)
    #x = np.arange(0, 75)
    #plt.plot(x, yy)
    #plt.show()
    if age < 30 and age > 20: age = age - 3
    elif age > 13 and age < 17: age = age + 4
    return age

def get_contours(bbox, shape):
    x = bbox[0] + 5
    y = bbox[3] - 10
    contourTop = [x, y]
    y = bbox[1] + 20
    contourBot = [x, y]
    return contourTop, contourBot

def open_image():
    print('Opening image...')
    dir_ = os.path.abspath(os.path.dirname(__file__))
    img = cv2.imread("{}/{}".format(dir_, args.img))
    outputs = os.listdir("{}/output".format(dir_))
    id_ = int(len(outputs)/2+1)
    return id_, img

if __name__ == "__main__":
    free_cache()
    id_, img = open_image()

    model = load_retina_model()
    bboxes = face_detect(model, img)
    print("Faces number: {}".format(len(bboxes)))
    free_cache(model)

    model = load_mask_model()
    mask_list = mask_detect(model, img, bboxes)
    free_cache(model)

    model = load_fairface_model()
    agr_table = predict_age_gender_race(model, img, bboxes)
    free_cache(model)

    create_log_and_image(id_, img, bboxes, agr_table, mask_list)
    print("All done!")


#        data = pygame.surfarray.array3d(pygame.display.get_surface())
from itertools import cycle
import random
import sys
import matplotlib.pyplot as plt
import pygame
import math
import numpy as np
import cv2
from skimage import color as colour
from skimage.transform import resize, downscale_local_mean
from pygame.locals import *

class FlappyGame:
    def __init__(self):
        self.FPS = 30
        self.SCREENWIDTH  = 288
        self.SCREENHEIGHT = 512
        self.PIPEGAPSIZE  = 100 # gap between upper and lower part of pipe
        self.BASEY        = self.SCREENHEIGHT * 0.79
        # image, sound and hitmask  dicts
        self.IMAGES, self.SOUNDS, self.HITMASKS = {}, {}, {}

        # list of all possible players (tuple of 3 positions of flap)
        self.PLAYERS_LIST = [
            # red bird
            (
                'assets/sprites/redbird-upflap.png',
                'assets/sprites/redbird-midflap.png',
                'assets/sprites/redbird-downflap.png',
            )
        ]

        # list of backgrounds
        self.BACKGROUNDS_LIST = ['assets/sprites/background-day.png']

        # list of pipes
        self.PIPES_LIST = ['assets/sprites/pipe-green.png']
        pygame.init()
        self.FPSCLOCK = pygame.time.Clock()
        self.SCREEN = pygame.display.set_mode((self.SCREENWIDTH, self.SCREENHEIGHT))
        pygame.display.set_caption('Flappy Bird')

        # numbers sprites for score display
        self.IMAGES['numbers'] = (
            pygame.image.load('assets/sprites/0.png').convert_alpha(),
            pygame.image.load('assets/sprites/1.png').convert_alpha(),
            pygame.image.load('assets/sprites/2.png').convert_alpha(),
            pygame.image.load('assets/sprites/3.png').convert_alpha(),
            pygame.image.load('assets/sprites/4.png').convert_alpha(),
            pygame.image.load('assets/sprites/5.png').convert_alpha(),
            pygame.image.load('assets/sprites/6.png').convert_alpha(),
            pygame.image.load('assets/sprites/7.png').convert_alpha(),
            pygame.image.load('assets/sprites/8.png').convert_alpha(),
            pygame.image.load('assets/sprites/9.png').convert_alpha()
        )

        # game over sprite
        self.IMAGES['gameover'] = pygame.image.load('assets/sprites/gameover.png').convert_alpha()
        # message sprite for welcome screen
        self.IMAGES['message'] = pygame.image.load('assets/sprites/message.png').convert_alpha()
        # base (ground) sprite
        self.IMAGES['base'] = pygame.image.load('assets/sprites/base.png').convert_alpha()

        # sounds
        if 'win' in sys.platform:
            soundExt = '.wav'
        else:
            soundExt = '.ogg'

        self.SOUNDS['die']    = pygame.mixer.Sound('assets/audio/die' + soundExt)
        self.SOUNDS['hit']    = pygame.mixer.Sound('assets/audio/hit' + soundExt)
        self.SOUNDS['point']  = pygame.mixer.Sound('assets/audio/point' + soundExt)
        self.SOUNDS['swoosh'] = pygame.mixer.Sound('assets/audio/swoosh' + soundExt)
        self.SOUNDS['wing']   = pygame.mixer.Sound('assets/audio/wing' + soundExt)
        randBg = 0
        self.IMAGES['background'] = pygame.image.load(self.BACKGROUNDS_LIST[randBg]).convert()

        # select random player sprites
        # randPlayer = random.randint(0, len(PLAYERS_LIST) - 1)
        randPlayer = 0
        self.IMAGES['player'] = (
            pygame.image.load(self.PLAYERS_LIST[randPlayer][0]).convert_alpha(),
            pygame.image.load(self.PLAYERS_LIST[randPlayer][1]).convert_alpha(),
            pygame.image.load(self.PLAYERS_LIST[randPlayer][2]).convert_alpha(),
        )

        # select random pipe sprites
        # pipeindex = random.randint(0, len(PIPES_LIST) - 1)
        pipeindex = 0
        self.IMAGES['pipe'] = (
            pygame.transform.flip(
                pygame.image.load(self.PIPES_LIST[pipeindex]).convert_alpha(), False, True),
            pygame.image.load(self.PIPES_LIST[pipeindex]).convert_alpha(),
        )

        # hismask for pipes
        self.HITMASKS['pipe'] = (
            self.getHitmask(self.IMAGES['pipe'][0]),
            self.getHitmask(self.IMAGES['pipe'][1]),
        )

        # hitmask for player
        self.HITMASKS['player'] = (
            self.getHitmask(self.IMAGES['player'][0]),
            self.getHitmask(self.IMAGES['player'][1]),
            self.getHitmask(self.IMAGES['player'][2]),
        )
        self.done = False
        self.time = 0
        self.reset()

    def reset(self):
        # player velocity, max velocity, downward accleration, accleration on flap
        self.pipeVelX = -4
        self.playerVelY    =  -9   # player's velocity along Y, default same as self.playerFlapped
        self.playerMaxVelY =  10   # max vel along Y, max descend speed
        self.playerMinVelY =  -8   # min vel along Y, max ascend speed
        self.playerAccY    =   1   # players downward accleration
        self.playerRot     =  45   # player's rotation
        self.playerVelRot  =   3   # angular speed
        self.playerRotThr  =  20   # rotation threshold
        self.playerFlapAcc =  -9   # players speed on flapping
        self.playerFlapped = False # True when player flaps
        self.playerx = int(self.SCREENWIDTH * 0.2)
        self.playery = int((self.SCREENHEIGHT - self.IMAGES['player'][0].get_height()) / 2)
        self.basex = 0
        self.playerIndexGen = cycle([0, 1, 2, 1])
        playerShmVals = {'val': 0, 'dir': 1}
        self.done = False
        self.time = 0
        self.score = self.playerIndex = self.loopIter = 0
        self.baseShift = self.IMAGES['base'].get_width() - self.IMAGES['background'].get_width()

        # get 2 new pipes to add to upperPipes lowerPipes list
        newPipe1 = self.getRandomPipe()
        newPipe2 = self.getRandomPipe()

        # list of upper pipes
        self.upperPipes = [
            {'x': self.SCREENWIDTH + 200, 'y': newPipe1[0]['y']},
            {'x': self.SCREENWIDTH + 200 + (self.SCREENWIDTH / 2), 'y': newPipe2[0]['y']},
        ]

        # list of lowerpipe
        self.lowerPipes = [
            {'x': self.SCREENWIDTH + 200, 'y': newPipe1[1]['y']},
            {'x': self.SCREENWIDTH + 200 + (self.SCREENWIDTH / 2), 'y': newPipe2[1]['y']},
        ]
        # return self._getScreen()
        return self._getValues()

    def step(self, action):
        reward = 1
        reward += self.act(action)
        # for x in range(3):
        #     reward += self.act(0)
        # return self._getScreen(), reward, self.done, self.time, self.score
        return self._getValues(), reward, self.done, self.time, self.score

    def act(self, action):
        #DEBUGGING
        #----------
        # for event in pygame.event.get():
        #     if event.type == KEYDOWN and (event.key == K_SPACE or event.key == K_UP):
        #         if self.playery > -2 * self.IMAGES['player'][0].get_height():
        #             self.playerVelY = self.playerFlapAcc
        #             self.playerFlapped = True
        #             self.SOUNDS['wing'].play()
        #--------

        #ACTUAL
        #--------
        pygame.event.get()
        self.time += 1
        reward = 0
        if action == 1 and not self.done:
            if self.playery > -2 * self.IMAGES['player'][0].get_height():
                        self.playerVelY = self.playerFlapAcc
                        self.playerFlapped = True
                        # self.SOUNDS['wing'].play()
        #--------

        if self.done:
            return 0
        crashTest, upper = self.checkCrash({'x': self.playerx, 'y': self.playery, 'index': self.playerIndex},
                            self.upperPipes, self.lowerPipes)
        if crashTest[0]:
            self.done = True
            if upper:
                return -1015
            return -1000
            # return a
        # check for score
        playerMidPos = self.playerx + self.IMAGES['player'][0].get_width() / 2
        for pipe in self.upperPipes:
            pipeMidPos = pipe['x'] + self.IMAGES['pipe'][0].get_width() / 2
            if pipeMidPos <= playerMidPos < pipeMidPos + 4:
                self.score += 1
                # self.SOUNDS['point'].play()
                reward += 10

        if (self.loopIter + 1) % 3 == 0:
            self.playerIndex = next(self.playerIndexGen)
        self.loopIter = (self.loopIter + 1) % 30
        self.basex = -((-self.basex + 100) % self.baseShift)

        # rotate the player
        if self.playerRot > -90:
            self.playerRot -= self.playerVelRot

        # player's movement
        if self.playerVelY < self.playerMaxVelY and not self.playerFlapped:
            self.playerVelY += self.playerAccY
        if self.playerFlapped:
            self.playerFlapped = False

            # more rotation to cover the threshold (calculated in visible rotation)
            self.playerRot = 45

        playerHeight = self.IMAGES['player'][self.playerIndex].get_height()
        self.playery += min(self.playerVelY, self.BASEY - self.playery - playerHeight)

        # move pipes to left
        for uPipe, lPipe in zip(self.upperPipes, self.lowerPipes):
            uPipe['x'] += self.pipeVelX
            lPipe['x'] += self.pipeVelX

        # add new pipe when first pipe is about to touch left of screen
        if len(self.upperPipes) > 0 and 0 < self.upperPipes[0]['x'] < 5:
            newPipe = self.getRandomPipe()
            self.upperPipes.append(newPipe[0])
            self.lowerPipes.append(newPipe[1])

        # remove first pipe if its out of the screen
        if len(self.upperPipes) > 0 and self.upperPipes[0]['x'] < -self.IMAGES['pipe'][0].get_width():
            self.upperPipes.pop(0)
            self.lowerPipes.pop(0)

        # draw sprites
        self.SCREEN.blit(self.IMAGES['background'], (0,0))

        for uPipe, lPipe in zip(self.upperPipes, self.lowerPipes):
            self.SCREEN.blit(self.IMAGES['pipe'][0], (uPipe['x'], uPipe['y']))
            self.SCREEN.blit(self.IMAGES['pipe'][1], (lPipe['x'], lPipe['y']))

        self.SCREEN.blit(self.IMAGES['base'], (self.basex, self.BASEY))
        # print score so player overlaps the score
        self.showScore(self.score)

        # Player rotation has a threshold
        self.visibleRot = self.playerRotThr
        if self.playerRot <= self.playerRotThr:
            self.visibleRot = self.playerRot

        playerSurface = pygame.transform.rotate(self.IMAGES['player'][self.playerIndex], self.visibleRot)
        self.SCREEN.blit(playerSurface, (self.playerx, self.playery))
        pygame.display.update()
        self.FPSCLOCK.tick(self.FPS)
        return reward

    def _getValues(self):
        nextPipe = 0
        playerPosY = self.playery + self.IMAGES['player'][0].get_height() / 2
        pipeWidth = self.IMAGES['pipe'][0].get_width()
        while nextPipe < len(self.upperPipes) and self.upperPipes[nextPipe]['x'] + 2 * pipeWidth / 3 < self.playerx:
            nextPipe += 1
        targetTop = self.upperPipes[nextPipe]['y'] + self.IMAGES['pipe'][nextPipe].get_height()
        targetBottom = self.lowerPipes[nextPipe]['y']

        targetX = self.upperPipes[nextPipe]['x']
        targetY = (targetTop + targetBottom) / 2
        if nextPipe + 1 < len(self.lowerPipes):
            nextY = self.lowerPipes[nextPipe + 1]['y'] - self.PIPEGAPSIZE / 2
        else:
            nextY = 0
        return np.array([playerPosY, self.playerVelY, targetX, targetY])

    def _getScreen(self):
        img = pygame.surfarray.array3d(self.SCREEN)
        img = colour.rgb2gray(img.swapaxes(0,1))[40:-124]
        img = downscale_local_mean(img, (8,8))
        img = resize(img, (128,76))
        return img

    def playerShm(self, playerShm):
        """oscillates the value of playerShm['val'] between 8 and -8"""
        if abs(playerShm['val']) == 8:
            playerShm['dir'] *= -1

        if playerShm['dir'] == 1:
            playerShm['val'] += 1
        else:
            playerShm['val'] -= 1

    def getRandomPipe(self):
        """returns a randomly generated pipe"""
        # y of gap between upper and lower pipe
        gapY = random.randrange(0, int(self.BASEY * 0.6 - self.PIPEGAPSIZE))
        gapY += int(self.BASEY * 0.2)
        pipeHeight = self.IMAGES['pipe'][0].get_height()
        pipeX = self.SCREENWIDTH + 10

        return [
            {'x': pipeX, 'y': gapY - pipeHeight},  # upper pipe
            {'x': pipeX, 'y': gapY + self.PIPEGAPSIZE}, # lower pipe
        ]

    def showScore(self, score):
        """displays score in center of screen"""
        scoreDigits = [int(x) for x in list(str(score))]
        totalWidth = 0 # total width of all numbers to be printed

        for digit in scoreDigits:
            totalWidth += self.IMAGES['numbers'][digit].get_width()

        Xoffset = (self.SCREENWIDTH - totalWidth) / 2

        for digit in scoreDigits:
            self.SCREEN.blit(self.IMAGES['numbers'][digit], (Xoffset, self.SCREENHEIGHT * 0.1))
            Xoffset += self.IMAGES['numbers'][digit].get_width()

    def checkCrash(self, player, upperPipes, lowerPipes):
        """returns True if player collders with base or pipes."""
        pi = player['index']
        player['w'] = self.IMAGES['player'][0].get_width()
        player['h'] = self.IMAGES['player'][0].get_height()

        # if player crashes into ground
        if player['y'] + player['h'] >= self.BASEY - 1:
            return [True, True], False
        else:
            playerRect = pygame.Rect(player['x'], player['y'],
                        player['w'], player['h'])
            pipeW = self.IMAGES['pipe'][0].get_width()
            pipeH = self.IMAGES['pipe'][0].get_height()

            for uPipe, lPipe in zip(upperPipes, lowerPipes):
                # upper and lower pipe rects
                uPipeRect = pygame.Rect(uPipe['x'], uPipe['y'], pipeW, pipeH)
                lPipeRect = pygame.Rect(lPipe['x'], lPipe['y'], pipeW, pipeH)

                # player and upper/lower pipe hitmasks
                pHitMask = self.HITMASKS['player'][pi]
                uHitmask = self.HITMASKS['pipe'][0]
                lHitmask = self.HITMASKS['pipe'][1]

                # if bird collided with upipe or lpipe
                uCollide = self.pixelCollision(playerRect, uPipeRect, pHitMask, uHitmask)
                lCollide = self.pixelCollision(playerRect, lPipeRect, pHitMask, lHitmask)

                if uCollide or lCollide:
                    return [True, False], uCollide

        return [False, False], False

    def pixelCollision(self, rect1, rect2, hitmask1, hitmask2):
        """Checks if two objects collide and not just their rects"""
        rect = rect1.clip(rect2)

        if rect.width == 0 or rect.height == 0:
            return False

        x1, y1 = rect.x - rect1.x, rect.y - rect1.y
        x2, y2 = rect.x - rect2.x, rect.y - rect2.y

        for x in range(rect.width):
            for y in range(rect.height):
                if hitmask1[x1+x][y1+y] and hitmask2[x2+x][y2+y]:
                    return True
        return False

    def getHitmask(self, image):
        """returns a hitmask using an image's alpha."""
        mask = []
        for x in range(image.get_width()):
            mask.append([])
            for y in range(image.get_height()):
                mask[x].append(bool(image.get_at((x,y))[3]))
        return mask