# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Module to handle observations based on observation modes."""

# STDLIB
import re
import warnings

# THIRD-PARTY
import numpy as np

# ASTROPY
from astropy import log
from astropy import units as u
from astropy.utils.decorators import lazyproperty
from astropy.utils.exceptions import AstropyUserWarning

# SYNPHOT
from synphot import exceptions as synexceptions
from synphot import units
from synphot.models import Empirical1D
from synphot.spectrum import SourceSpectrum, SpectralElement
from synphot.thermal import ThermalSpectralElement
from synphot.utils import merge_wavelengths

# LOCAL
from . import exceptions, stio
from .config import conf
from .spectrum import Vega, interpolate_spectral_element
from .tables import GraphTable, CompTable
from .wavetable import WAVECAT

__all__ = ['reset_cache', 'Component', 'ThermalComponent',
           'BaseObservationMode', 'ObservationMode', 'ThermalObservationMode']

_band_patt = re.compile(r'band\((.*?)\)', re.IGNORECASE)

# Cache previously loaded graph, component, thermal, and detector tables
_GRAPHDICT = {}
_COMPDICT = {}
_THERMDICT = {}
_DETECTORDICT = {}

def reset_cache():
    """Empty the table dictionaries cache."""
    global _GRAPHDICT, _COMPDICT, _THERMDICT, _DETECTORDICT
    _GRAPHDICT.clear()
    _COMPDICT.clear()
    _THERMDICT.clear()
    _DETECTORDICT.clear()

class Component:
    """Class to handle individual components in `BaseObservationMode`.

    Parameters
    ----------
    throughput_name : str
        Component filename.

    interpval : float or `None`
        If not `None`, interpolate to the given value.
        See :func:`stsynphot.spectrum.interpolate_spectral_element`.

    Attributes
    ----------
    throughput_name : str
        Component filename.

    throughput : `synphot.spectrum.SpectralElement` or `None`
        Component spectrum object.

    """
    def __init__(self, throughput_name, interpval=None):
        self.throughput_name = throughput_name

        # Extract bandpass unless component is a CLEAR filter.
        if throughput_name != conf.clear_filter:
            if interpval is None:
                self.throughput = SpectralElement.from_file(throughput_name)
            else:
                self.throughput = interpolate_spectral_element(
                    throughput_name, interpval)
        else:
            self.throughput = None

    @property
    def empty(self):
        """`True` if ``self.throughput`` is empty."""
        return self.throughput is None

    def __str__(self):
        return str(self.throughput)

class ThermalComponent(Component):
    """Class to handle thermal components in `BaseObservationMode`.

    Parameters
    ----------
    throughput_name, thermal_name : str
        Optical and thermal component filenames.

    interpval : float or `None`
        If not `None`, interpolate to the given value (optical only).
        See :func:`stsynphot.spectrum.interpolate_spectral_element`.

    Attributes
    ----------
    throughput_name, thermal_name : str
        Optical and thermal component filenames.

    throughput : `synphot.spectrum.SpectralElement` or `None`
        Optical component spectrum object.

    emissivity : `synphot.thermal.ThermalSpectralElement` or `None`
        Thermal component spectrum object.

    """
    def __init__(self, throughput_name, thermal_name, interpval=None):
        super(ThermalComponent, self).__init__(
            throughput_name, interpval=interpval)
        self.thermal_name = thermal_name

        # Extract thermal spectrum unless component is a CLEAR filter.
        if thermal_name != conf.clear_filter:
            self.emissivity = ThermalSpectralElement.from_file(thermal_name)
        else:
            self.emissivity = None

    @property
    def empty(self):
        """`True` if component is empty."""
        return self.throughput is None and self.emissivity is None

    def __str__(self):
        return str(self.emissivity)

def _process_graphtable(graphtable):
    """Load and cache graphtable. Get primary area."""
    global _GRAPHDICT

    # Use default graph table if not given
    if graphtable is None:
        gtname = conf.graphtable
    else:
        gtname = graphtable

    # Load and cache graphtable, if not in cache
    if gtname in _GRAPHDICT:
        gt = _GRAPHDICT[gtname]
    else:
        gt = GraphTable(gtname)
        _GRAPHDICT[gtname] = gt

    # Set telescope collecting area
    if gt.primary_area is None:
        area = conf.area
    else:
        area = gt.primary_area
    primary_area = units.validate_quantity(area, units.AREA)

    return gt, gtname, primary_area

class BaseObservationMode:
    """Base class to handle an observation that uses the graph
    and optical component tables, common to both optical and
    thermal observation modes.

    .. note::

        ``modes`` for parameterized mode is set such that the parameterized
        value is stripped away; e.g., ``'acs,wfc1,f555w,mjd#53000'``
        becomes ``['acs', 'wfc1', 'f555w', 'mjd#']``. Instead, the
        parameterized value is kept in ``pardict``; e.g., ``'mjd#53000'``
        becomes ``{'mjd': 53000.0}``.

        ``pixscale`` is set from ``stsynphot.config.conf.detectorfile``,
        which is parsed with :func:`stsynphot.stio.read_detector_pars`.

        ``binset`` is set by ``stsynphot.wavetable.WAVECAT``.

    Parameters
    ----------
    obsmode : str
        Observation mode.

    graphtable : str or `None`
        Graph table filename.
        If `None`, uses ``stsynphot.config.conf.graphtable``.

    comptable : str or `None`
        Optical component table filename.
        If `None`, uses ``stsynphot.config.conf.comptable``.

    Attributes
    ----------
    modes : list
        List of individual modes within observation mode.

    pardict : dict
        Maps parameterized mode to its value.

    gtname, ctname : str
        Graph and component table filenames.

    compnames, thcompnames : list of str
        Optical and thermal components.

    components : list of obj
        List of component objects.

    primary_area : `astropy.units.quantity.Quantity`
        Telescope collecting area.

    pixscale : `astropy.units.quantity.Quantity`
        Detector pixel scale.

    binset : str
        Wavelength table filename/param string from matching obsmode.

    bandwave : `astropy.units.quantity.Quantity`
        Wavelength set defined by ``binset``.

    """
    def __init__(self, obsmode, graphtable=None, comptable=None):
        global _GRAPHDICT, _COMPDICT

        # Strip "band()" syntax if present, and force lowercase
        tmatch = _band_patt.search(obsmode)
        if tmatch:
            self._obsmode = tmatch.group(1).lower()
        else:
            self._obsmode = obsmode.lower()

        # Split obsmode and separate parameterized modes
        modes = self._obsmode.replace(' ', '').split(',')
        self.pardict = {}
        if '#' in self._obsmode:
            self.modes = []
            for m in modes:
                if '#' in m:
                    key, val = m.split('#')
                    self.pardict[key] = float(val)
                    self.modes.append(f'{key:s}#')
                else:
                    self.modes.append(m)
        else:
            self.modes = modes

        # Get graph table and primary area
        gt, self.gtname, self.primary_area = _process_graphtable(graphtable)

        # Get optical and thermal components
        self.compnames, self.thcompnames = gt.get_comp_from_gt(self.modes, 1)

        # Use default optical component table if not given
        if comptable is None:
            self.ctname = conf.comptable
        else:
            self.ctname = comptable

        # Load and cache comptable, if not in cache
        if self.ctname in _COMPDICT:
            ct = _COMPDICT[self.ctname]
        else:
            ct = CompTable(self.ctname)
            _COMPDICT[self.ctname] = ct

        # Set by sub-classes
        self.components = None

        # Get optical component filenames
        self._throughput_filenames = ct.get_filenames(self.compnames)

        # Set detector pixel scale
        self._set_pixscale()

        # For sensitivity calculations
        self._constant = self.primary_area / units.HC

        # Get wavelength set
        try:
            self.binset, self.bandwave = WAVECAT.load_waveset(self._obsmode)
        except (KeyError, exceptions.AmbiguousObsmode):
            self.binset = ''
            self.bandwave = None

    def _set_pixscale(self):
        """Set pixel scale.
        If multiple matches found, only first match is used.

        """
        global _DETECTORDICT

        det_file = stio.irafconvert(conf.detectorfile)

        if det_file in _DETECTORDICT:
            data = _DETECTORDICT[det_file]
        else:
            data = stio.read_detector_pars(det_file)
            _DETECTORDICT[det_file] = data

        obsmode = ','.join(self._obsmode.split(',')[:2])
        pixscales = data[data['OBSMODE'] == obsmode]['SCALE'].data

        if pixscales.size < 1:
            self.pixscale = None
        else:
            self.pixscale = pixscales[0] * u.arcsec

    def _get_components(self):
        raise NotImplementedError('To be implemented by subclasses.')

    def __str__(self):
        return self._obsmode

    def __len__(self):
        return len(self.components)

    @staticmethod
    def _parkey_from_filename(filename):
        """Extract parkey from component filename in the format of
        ``name[parkey#]``.

        """
        if filename.endswith('#]'):
            x = filename.split('[')
            parkey = x[1][:-2]
        else:
            parkey = None
        return parkey

    def showfiles(self):  # pragma: no cover
        """Display optical component filenames.

        .. note:: Similar to IRAF SYNPHOT SHOWFILES.

        """
        info_str = '#Throughput table names:\n'
        for name in self._throughput_filenames:
            if name != conf.clear_filter:
                info_str += f'{name}\n'
        log.info(info_str.rstrip())

class ObservationMode(BaseObservationMode):
    """Class to handle an observation that uses the graph
    and optical component tables.

    See `BaseObservationMode` for additional attributes.

    Parameters
    ----------
    obsmode, graphtable, comptable
        See `BaseObservationMode`.

    component_dict : dict
        Maps component filename to corresponding `Component`.

    """
    def __init__(self, obsmode, graphtable=None, comptable=None,
                 component_dict={}):
        super(ObservationMode, self).__init__(
            obsmode, graphtable=graphtable, comptable=comptable)
        self._component_dict = component_dict
        self.components = self._get_components()

    def _get_components(self):
        """Get optical components."""
        components = []

        for throughput_name in self._throughput_filenames:
            parkey = self._parkey_from_filename(throughput_name)

            # Hack for WFPC2 ramp filters
            if parkey == 'wave' and self.modes[0] == 'wfpc2':
                parkey = 'lrf'

            cdict_key = (throughput_name, self.pardict.get(parkey))

            if cdict_key not in self._component_dict:
                self._component_dict[cdict_key] = Component(
                    cdict_key[0], interpval=cdict_key[1])

            component = self._component_dict[cdict_key]

            if not component.empty:
                components.append(component)

        return components

    def _mul_thru(self, index):
        """Multiply all component spectra starting at given index."""
        product = self.components[index].throughput

        if len(self.components) > index:
            for component in self.components[index + 1:]:
                if not component.empty:
                    product = product * component.throughput

        product.meta['header'] = ''  # Clean up messy header
        return product

    @lazyproperty
    def throughput(self):
        """Combined throughput from multiplying all the components together."""
        try:
            thru = self._mul_thru(0)
        except IndexError as e:  # pragma: no cover
            thru = None
            warnings.warn(
                f'Graph table is broken: {repr(e)}', AstropyUserWarning)
        return thru

    @lazyproperty
    def sensitivity(self):
        """Sensitivity spectrum to convert flux in
        :math:`erg \\; cm^{-2} \\; s^{-1} \\; \\AA^{-1}` to
        :math:`count s^{-1} \\AA^{-1}`. Calculation is done by
        combining the throughput curves with
        :math:`\\frac{h \\; c}{\\lambda}` .

        """
        x = self.throughput.waveset
        y = self.throughput(x)
        thru = y.value * x.value * self._constant.value
        meta = {'expr': f'Sensitivity for {self._obsmode}'}
        return SpectralElement(
            Empirical1D, points=x, lookup_table=thru, meta=meta)

    def thermal_spectrum(self, thermtable=None):
        """Calculate thermal spectrum using
        :func:`ThermalObservationMode.to_spectrum`.

        Parameters
        ----------
        thermtable : str or `None`
            Thermal component table filename.
            If `None`, uses ``stsynphot.config.conf.thermtable``.

        Returns
        -------
        sp : `synphot.spectrum.SourceSpectrum`
            Thermal spectrum in PHOTLAM.

        Raises
        ------
        synphot.exceptions.SynphotError
            Calculation failed.

        """
        thom = ThermalObservationMode(
            self._obsmode, graphtable=self.gtname, comptable=self.ctname,
            thermtable=thermtable)

        try:
            sp = thom.to_spectrum()
        except IndexError as e:  # pragma: no cover
            raise synexceptions.SynphotError(f'Broken graph table: {repr(e)}')

        return sp

class ThermalObservationMode(BaseObservationMode):
    """Class to handle an observation that uses the graph,
    optical component, and thermal component tables.

    See `BaseObservationMode` for additional attributes.

    Parameters
    ----------
    obsmode, graphtable, comptable
        See `BaseObservationMode`.

    thermtable : str or `None`
        Thermal component table filename.
        If `None`, uses ``stsynphot.config.conf.thermtable``.

    Attributes
    ----------
    thname : str
        Thermal component table filename.

    Raises
    ------
    NotImplementedError
        Clear filters only not supported.

    """
    def __init__(self, obsmode, graphtable=None, comptable=None,
                 thermtable=None):
        global _GRAPHDICT, _COMPDICT, _THERMDICT

        super(ThermalObservationMode, self).__init__(
            obsmode, graphtable=graphtable, comptable=comptable)

        # Check here to see if there are any valid filters.
        # "0.0" was added in tae17277m_tmt.fits (Apr 2017).
        if set(self.thcompnames).issubset(set([conf.clear_filter, '', '0.0'])):
            raise NotImplementedError(
                f'No thermal support provided for {self._obsmode}')

        # Use default thermal component table if not given
        if thermtable is None:
            self.thname = conf.thermtable
        else:
            self.thname = thermtable

        # Load and cache thermtable, if not in cache
        if self.thname in _THERMDICT:
            thct = _THERMDICT[self.thname]
        else:
            thct = CompTable(self.thname)
            _THERMDICT[self.thname] = thct

        # Get thermal component filenames
        self._thermal_filenames = thct.get_filenames(self.thcompnames)

        self.components = self._get_components()

    def _get_components(self):
        """Get thermal components."""
        components = []

        for throughput_name, thermal_name in zip(
                self._throughput_filenames, self._thermal_filenames):
            parkey = self._parkey_from_filename(throughput_name)
            component = ThermalComponent(throughput_name, thermal_name,
                                         interpval=self.pardict.get(parkey))
            if not component.empty:
                components.append(component)

        return components

    def __str__(self):
        return f'{self._obsmode} (thermal)'

    def _merge_em_wave(self):
        """Merge emissivity wavelength sets in Angstrom."""
        index = 1

        # Find first wavelength set
        for component in self.components:
            emissivity = component.emissivity
            if emissivity is None:
                index += 1
            else:
                result = emissivity.waveset.value
                break

        # Merge subsequent wavelength sets
        for component in self.components[index:]:
            emissivity = component.emissivity
            if emissivity is not None:
                result = merge_wavelengths(result, emissivity.waveset.value)

        return result

    def _get_wave_intersection(self):
        """Find wavelengths in Angstrom where
        ``stsynphot.config.conf.waveset_array`` intersects with
        ``stsynphot.spectrum.Vega``.

        .. note:: No one knows why Vega is the chosen one.

        """
        def_wave = conf.waveset_array  # Angstrom
        minw = min(def_wave)
        maxw = max(def_wave)

        # Refine min and max wavelengths using thermal components
        for component in self.components[1:]:
            emissivity = component.emissivity
            if emissivity is not None:
                w = emissivity.waveset.value  # Angstrom
                minw = max(minw, w.min())
                maxw = min(maxw, w.max())

        w = self._merge_em_wave()
        result = w[(w > minw) & (w < maxw)]

        # Intersect with Vega
        if Vega is None:
            raise synexceptions.SynphotError('Missing Vega spectrum.')
        w = Vega.waveset.value  # Angstrom

        return result[(result > w.min()) & (result < w.max())]

    def to_spectrum(self):
        """Get thermal spectrum.

        The calculations start with zero-flux spectrum.
        Then for each component:

        #. Multiply with optical throughput.
        #. Add in thermal source spectrum from
           :meth:`synphot.thermal.ThermalSpectralElement.thermal_source`.

        Returns
        -------
        sp : `synphot.spectrum.SourceSpectrum`
            Thermal spectrum in PHOTLAM.

        """
        # Create zero-flux spectrum
        x = self._get_wave_intersection()  # Angstrom
        y = np.zeros_like(x, dtype=np.float64)  # PHOTLAM
        minw, maxw = x[([0, -1], )]
        sp = SourceSpectrum(Empirical1D, points=x, lookup_table=y)

        for component in self.components:
            # Transmissive section (optical passband)
            if component.throughput is not None:
                sp = sp * component.throughput

            # Thermal section
            if component.emissivity is not None:
                sp = sp + component.emissivity.thermal_source()

                # Trim spectrum
                w = sp.waveset.value
                mask = (w >= minw) & (w <= maxw)
                x = w[mask]
                y = sp(x)
                sp = SourceSpectrum(Empirical1D, points=x, lookup_table=y)

        meta = {'expr': f'{self._obsmode} ThermalSpectrum'}
        sp.meta.update(meta)
        return sp


"""Training a face recognizer with TensorFlow based on the FaceNet paper
FaceNet: A Unified Embedding for Face Recognition and Clustering: http://arxiv.org/abs/1503.03832
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from datetime import datetime
import os.path
import time
import sys
import tensorflow as tf
import numpy as np
import importlib
import itertools
import argparse
import facenet
import cv2
import random
import pdb
import h5py
import math

def main(args):
    network = importlib.import_module(args.model_def)

    subdir = datetime.strftime(datetime.now(), '%Y%m%d-%H%M%S') + args.experiment_name
    log_dir = os.path.join(os.path.expanduser(args.logs_base_dir), subdir)
    if not os.path.isdir(log_dir):  # Create the log directory if it doesn't exist
        os.makedirs(log_dir)
    model_dir = os.path.join(os.path.expanduser(args.models_base_dir), subdir)
    if not os.path.isdir(model_dir):  # Create the model directory if it doesn't exist
        os.makedirs(model_dir)

    # Write arguments to a text file
    facenet.write_arguments_to_file(args, os.path.join(log_dir, 'arguments.txt'))

    # Store some git revision info in a text file in the log directory
    src_path,_ = os.path.split(os.path.realpath(__file__))
    facenet.store_revision_info(src_path, log_dir, ' '.join(sys.argv))

    np.random.seed(seed=args.seed)
    class_name = ['smile','oval_face','5_ocloc_shadow','bald','archied_eyebrows','big_lips', 'big_Nose']
    class_num = len(class_name)
    class_index = [31,25,0,4,1,6,7]
    all_image_list = []
    all_label_list = []
    for i in range(class_num):

        image_list = []
        label_list = []

        train_set = facenet.get_sub_category_dataset(args.data_dir, class_index[i])

        image_list_p, label_list_p = facenet.get_image_paths_and_labels_triplet(train_set[0], args)
        image_list_n, label_list_n = facenet.get_image_paths_and_labels_triplet(train_set[1], args)

        image_list.append(image_list_p)
        image_list.append(image_list_n)
        label_list.append(label_list_p)
        label_list.append(label_list_n)

        all_image_list.append(image_list)
        all_label_list.append(label_list)

    print('Model directory: %s' % model_dir)
    print('Log directory: %s' % log_dir)
    if args.pretrained_model:
        print('Pre-trained model: %s' % os.path.expanduser(args.pretrained_model))

    image_size = args.image_size
    batch_size = args.batch_size
    with tf.Graph().as_default():
        tf.set_random_seed(args.seed)
        global_step = tf.Variable(0, trainable=False)

        # Placeholder for the learning rate
        learning_rate_placeholder = tf.placeholder(tf.float32, name='learning_rate')

        batch_size_placeholder = tf.placeholder(tf.int32, name='batch_size')

        phase_train_placeholder = tf.placeholder(tf.bool, name='phase_train')

        # image_paths_placeholder = tf.placeholder(tf.string, shape=(None,3), name='image_paths')
        image_placeholder = tf.placeholder(tf.float32, shape=(batch_size, image_size, image_size,3), name='images')
        labels_placeholder = tf.placeholder(tf.int64, shape=(batch_size,3), name='labels')

        code_placeholder = tf.placeholder(tf.float32, shape=(batch_size,class_num,1,1), name='code')

        image_batch = normalized_image(image_placeholder)
        code_batch = code_placeholder

        # for i in range(class_num):
        #     mask_array[:,i,(args.embedding_size/class_num)*i:(args.embedding_size/class_num)*(i+1)] = 1

        prelogits, _ = network.inference(image_batch, args.keep_probability,
            phase_train=phase_train_placeholder, bottleneck_layer_size=args.embedding_size,
            weight_decay=args.weight_decay)

        embeddings = tf.nn.l2_normalize(prelogits, 1, 1e-10, name='embeddings')
        anchor_index = list(range(0,batch_size,3))
        positive_index = list(range(1,batch_size,3))
        negative_index = list(range(2,batch_size,3))

        a_indice = tf.constant(np.array(anchor_index))
        p_indice = tf.constant(np.array(positive_index))

        n_indice = tf.constant(np.array(negative_index))

        anchor = tf.gather(embeddings,a_indice)
        positive = tf.gather(embeddings,p_indice)
        negative = tf.gather(embeddings,n_indice)

        triplet_loss = facenet.triplet_loss(anchor, positive, negative, args.alpha)

        learning_rate = tf.train.exponential_decay(learning_rate_placeholder, global_step,
            args.learning_rate_decay_epochs*args.epoch_size, args.learning_rate_decay_factor, staircase=True)
        tf.summary.scalar('learning_rate', learning_rate)

        # Calculate the total losses
        regularization_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
        total_loss = tf.add_n([triplet_loss] + regularization_losses, name='total_loss')

        # Build a Graph that trains the model with one batch of examples and updates the model parameters
        train_op = facenet.train(total_loss, global_step, args.optimizer,
            learning_rate, args.moving_average_decay, tf.global_variables())

        # Create a saver
        trainable_variables = tf.trainable_variables()
        saver = tf.train.Saver(trainable_variables, max_to_keep=3)

        # Build the summary operation based on the TF collection of Summaries.
        summary_op = tf.summary.merge_all()

        # Start running operations on the Graph.
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=args.gpu_memory_fraction)
        sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

        # Initialize variables
        sess.run(tf.global_variables_initializer(), feed_dict={phase_train_placeholder:True})
        sess.run(tf.local_variables_initializer(), feed_dict={phase_train_placeholder:True})

        summary_writer = tf.summary.FileWriter(log_dir, sess.graph)
        coord = tf.train.Coordinator()
        tf.train.start_queue_runners(coord=coord, sess=sess)

        with sess.as_default():

            if args.pretrained_model:
                print('Restoring pretrained model: %s' % args.pretrained_model)
                saver.restore(sess, os.path.expanduser(args.pretrained_model))

            # Training and validation loop
            epoch = 0
            Accuracy = [0]
            while epoch < args.max_nrof_epochs:
                step = sess.run(global_step, feed_dict=None)
                # epoch = step // args.epoch_size
                # Train for one epoch
                code_list = []
                triplets_list = []
                if (epoch+1)%args.lr_epoch == 0:
                    args.learning_rate = 0.1*args.learning_rate

                for i in range(class_num):

                    code = np.zeros((batch_size, class_num, 1, 1), np.float32)
                    _class = i
                    code[:, _class, :, :] = 1

                    Triplets = triplet_random(args, sess, all_image_list[i], all_image_list, epoch, image_placeholder,
                                              batch_size_placeholder, learning_rate_placeholder,
                                              phase_train_placeholder, global_step,
                                              embeddings, total_loss, train_op, summary_op, summary_writer,
                                              args.embedding_size, anchor, positive, negative, triplet_loss)
                    triplets_list.append(Triplets)
                    code_list.append(code)

                train(args, sess, image_list, epoch, image_placeholder, code_placeholder,
                    batch_size_placeholder, learning_rate_placeholder, phase_train_placeholder, global_step,
                    embeddings, total_loss, train_op, summary_op, summary_writer,
                    args.embedding_size, anchor, positive, negative, triplet_loss, triplets_list, code_list, model_dir, Accuracy)

                if (epoch+1)%2 == 0:
                    Accuracy = test(args, sess, image_list, epoch, image_placeholder,code_placeholder,
                          batch_size_placeholder, learning_rate_placeholder, phase_train_placeholder, global_step,
                          embeddings, total_loss, train_op, summary_op, summary_writer,
                          args.embedding_size, anchor, positive, negative, triplet_loss, triplets_list, Accuracy)

                # Save variables and the metagraph if it doesn't exist already
                model_name = 'epoch' + str(epoch+1)
                print(model_dir)
                if (epoch+1) > 0 :
                    if (epoch +1)%2 == 0:
                        save_variables_and_metagraph(sess, saver, summary_writer, model_dir, model_name, step)
                        print('models are saved in ', os.path.join(model_dir, model_name))
                epoch = epoch + 1
    sess.close()
    return model_dir

def normalized_image(img_sym, resnet_mean = [102.9801, 115.9465, 122.7717]):
    return tf.scalar_mul(0.0078125, tf.subtract(tf.cast(img_sym, tf.float32), tf.constant(resnet_mean)))

def triplet_random(args, sess, dataset, image_list, epoch, image_placeholder,
          batch_size_placeholder, learning_rate_placeholder, phase_train_placeholder, global_step,
          embeddings, loss, train_op, summary_op, summary_writer,
          embedding_size, anchor, positive, negative, triplet_loss):

    batch_number = 0

    images_data = h5py.File('./data/train_classification.h5')
    images = images_data['data']

    image_list_p = dataset[0]
    image_list_n = dataset[1]

    random.shuffle(image_list_n)

    start_time = time.time()
    nrof_examples = len(image_list_p)
    triplets = []
    nrof_neg = len(image_list_n)

    for i in range(nrof_examples):

        a_idx = i
        p_idx = np.random.randint(i,nrof_examples)
        n_idx = np.random.randint(nrof_neg)
        triplets.append((image_list_p[a_idx], image_list_p[p_idx], image_list_n[n_idx]))

    return triplets

def train(args, sess, dataset, epoch, image_placeholder, code_placeholder,
          batch_size_placeholder, learning_rate_placeholder, phase_train_placeholder, global_step,
          embeddings, loss, train_op, summary_op, summary_writer,
          embedding_size, anchor, positive, negative, triplet_loss, triplets_list, code_list,model_dir,Accuracy):

    batch_number = 0

    images_data = h5py.File('./data/train_classification.h5')
    images = images_data['data']

    lr = args.learning_rate

    start_time = time.time()

    # sess.run(enqueue_op, {image_paths_placeholder: image_paths_array, labels_placeholder: labels_array})
    nrof_examples = 0
    class_num = len(triplets_list)
    # for i in range(class_num):
    nrof_examples = max(len(triplets_list[0]), nrof_examples)

    nrof_batches = int(np.ceil(nrof_examples / (args.batch_size/3/class_num)))

    print('load mean image done!')

    train_time = 0
    batch_number = 0
    batch_size = args.batch_size
        # emb_array = np.zeros((nrof_examples, embedding_size))
        # loss_array = np.zeros((nrof_triplets,))
    while batch_number < nrof_batches:

            start_time = time.time()

            class_num = len(triplets_list)

            image_array = np.zeros((batch_size,150,150,3),np.float32)
            code_array = np.zeros((0,class_num,1,1), np.float32)
            length = []
            for t in range(class_num):

              triplets = triplets_list[t]
              length.append(len(triplets))

              start_index = t*int(batch_size/class_num)

              image_array = facenet.get_triplet_image_batch1(images, triplets, batch_size/class_num, batch_number, image_array, start_index, args)
              code = code_list[t]

              code_array_tem = code[0:int(batch_size/class_num)]

              code_array = np.vstack((code_array, code_array_tem))

            time1 = time.time() - start_time
            feed_dict = {batch_size_placeholder: batch_size, learning_rate_placeholder: lr, phase_train_placeholder: True, image_placeholder: image_array, code_placeholder: code_array}
            err, _, step, emb = sess.run([loss, train_op, global_step, embeddings], feed_dict=feed_dict)
            # emb_array[lab,:] = emb
            # loss_array[i] = err
            emb_feature = emb
            nrof_pair = emb_feature.shape[0]/3
            correct_num = 0

            for pair_index in range(0, int(nrof_pair/class_num)):
                anchor_index = pair_index*3
                pos_index = anchor_index + 1
                neg_index = anchor_index + 2
                an_feature = emb_feature[anchor_index,:]
                pos_feature = emb_feature[pos_index, :]
                neg_feature = emb_feature[neg_index, :]
                if np.sum(np.square(an_feature - pos_feature),0) < np.sum(np.square(an_feature - neg_feature),0):
                    correct_num = correct_num + 1

            p_dist = np.sum(np.square(an_feature - pos_feature), 0)
            n_dist = np.sum(np.square(an_feature - neg_feature),0)
            print('triplet number:', length)
            print('pos_distance and neg_distance', p_dist, n_dist, 'model_dir:',model_dir,'last accuracy:',Accuracy[-1])
            accuracy = correct_num/(nrof_pair/class_num)
            duration = time.time() - start_time
            print('Epoch: [%d][%d/%d]\tTime %.3f\tread_time %.3f\tAccuracy %2.3f\tlr %2.3f\tLoss %2.3f' %
                  (epoch, batch_number+1, nrof_batches, duration, time1, accuracy, lr, err))
            batch_number += 1
            train_time += duration

        # Add validation loss and accuracy to summary
    summary = tf.Summary()
    #pylint: disable=maybe-no-member
    #summary.value.add(tag='time/selection', simple_value=train_time)
    summary_writer.add_summary(summary, step)
    return step
def get_code_batch(code, triplets, batch_size, batch_index):

    nrof_examples = len(triplets)

    j = batch_index * batch_size % nrof_examples

    if j + batch_size < nrof_examples:
        code_array = code[j: j+ batch_size]

    else:
        code_array1 = code[j:nrof_examples]
        code_array2 = code[0:batch_size - (nrof_examples-j)]
        code_array = np.vstack((code_array1,code_array2))
    return code_array
def test(args, sess, dataset, epoch, image_placeholder, code_placeholder,
          batch_size_placeholder, learning_rate_placeholder, phase_train_placeholder, global_step,
          embeddings, loss, train_op, summary_op, summary_writer, learning_rate_schedule_file,
          embedding_size, anchor, positive, negative, triplet_loss, triplet_list, Accuracy):

    images_data = h5py.File('./data/test_triplet_smile.h5')
    images_a = images_data['anchors']
    images_p = images_data['positive']
    images_n = images_data['negative']

    nrof_images = images_a.shape[0]
    nrof_batches = int(math.ceil(1.0 * nrof_images / (args.batch_size/3)))

    batch_number = 0
    correct_num = 0
    class_num = len(triplet_list)
    batch_size = args.batch_size
    print('load mean image done!')
    # fail_dir = '../fail_smile_pair'
    code = np.zeros((batch_size, class_num, 1, 1), np.float32)
    _class = 0
    if _class > 3:
        _class = 3

    code[:, _class, :, :] = 1

    check = 0
    while batch_number < nrof_batches:

            # batch_size = min(nrof_images - batch_number * args.batch_size, args.batch_size)
            batch_inter = int(batch_size/3)
            if (nrof_images - batch_number*batch_inter) < batch_inter:
                check = check + 1
                anchor_array1 = images_a[batch_number * batch_inter:nrof_images]
                anchor_array2 = images_a[0:batch_inter - (nrof_images - batch_number * batch_inter)]
                anchor_array = np.vstack((anchor_array1,anchor_array2))

                positive_array1 = images_p[batch_number * batch_inter:nrof_images]
                positive_array2 = images_p[0:batch_inter - (nrof_images - batch_number * batch_inter)]
                positive_array = np.vstack((positive_array1, positive_array2))

                negative_array1 = images_n[batch_number * batch_inter:nrof_images]
                negative_array2 = images_n[0:batch_inter - (nrof_images - batch_number * batch_inter)]
                negative_array = np.vstack((negative_array1, negative_array2))
            else:
                anchor_array = images_a[batch_number*batch_inter:batch_number*batch_inter + batch_inter]
                positive_array = images_p[batch_number * batch_inter:batch_number * batch_inter + batch_inter]
                negative_array = images_n[batch_number * batch_inter:batch_number * batch_inter + batch_inter]

            image_array = np.vstack((anchor_array,positive_array,negative_array))

            feed_dict = {batch_size_placeholder: batch_size, phase_train_placeholder: False, image_placeholder: image_array, code_placeholder: code}

            emb = sess.run(embeddings, feed_dict=feed_dict)

            batch_number += 1

            # emb_array[lab,:] = emb
            # loss_array[i] = er

            for i in range(int(batch_size/3)):

                an_feature = emb[i,:]
                pos_feature = emb[i+batch_size/3, :]
                neg_feature = emb[i+2*batch_size/3, :]

                if np.sum(np.square(an_feature - pos_feature),0) < np.sum(np.square(an_feature - neg_feature),0):
                    correct_num = correct_num + 1
    nrof_images = nrof_batches*batch_inter
    accuracy = correct_num/nrof_images
    Accuracy.append(accuracy)
    print('accuracy:', Accuracy)
    return Accuracy

def save_variables_and_metagraph(sess, saver, summary_writer, model_dir, model_name, step):
    # Save the model checkpoint
    print('Saving variables')
    start_time = time.time()
    checkpoint_path = os.path.join(model_dir, 'model-%s.ckpt' % model_name)
    saver.save(sess, checkpoint_path, global_step=step, write_meta_graph=False)
    save_time_variables = time.time() - start_time
    print('Variables saved in %.2f seconds' % save_time_variables)
    metagraph_filename = os.path.join(model_dir, 'model-%s.meta' % model_name)
    save_time_metagraph = 0
    if not os.path.exists(metagraph_filename):
        print('Saving metagraph')
        start_time = time.time()
        saver.export_meta_graph(metagraph_filename)
        save_time_metagraph = time.time() - start_time
        print('Metagraph saved in %.2f seconds' % save_time_metagraph)
    summary = tf.Summary()
    #pylint: disable=maybe-no-member
    summary.value.add(tag='time/save_variables', simple_value=save_time_variables)
    summary.value.add(tag='time/save_metagraph', simple_value=save_time_metagraph)
    summary_writer.add_summary(summary, step)

def get_learning_rate_from_file(filename, epoch):
    with open(filename, 'r') as f:
        for line in f.readlines():
            line = line.split('#', 1)[0]
            if line:
                par = line.strip().split(':')
                e = int(par[0])
                lr = float(par[1])
                if e <= epoch:
                    learning_rate = lr
                else:
                    return learning_rate

def parse_arguments(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('--is_training', type=bool,default= True)
    parser.add_argument('--lr_epoch', type=int,default= 30)
    parser.add_argument('--random_trip', type=bool, default= False)
    parser.add_argument('--experiment_name', type=str,default='triplet')
    parser.add_argument('--category', type=int, default=31)
    parser.add_argument('--learned_mask', type=bool, default=True)
    parser.add_argument('--logs_base_dir', type=str,
        help='Directory where to write event logs.', default='~/logs/facenet')
    parser.add_argument('--models_base_dir', type=str,
        help='Directory where to write trained models and checkpoints.', default='~/models/facenet')
    parser.add_argument('--gpu_memory_fraction', type=float,
        help='Upper bound on the amount of GPU memory that will be used by the process.', default=1.0)
    parser.add_argument('--pretrained_model', type=str,
        help='Load a pretrained model before training starts.', default = '')
    parser.add_argument('--data_dir', type=str,
        help='Path to the data directory containing aligned face patches. Multiple directories are separated with colon.',
        default='')
    parser.add_argument('--model_def', type=str,
        help='Model definition. Points to a module containing the definition of the inference graph.', default='models.inception_resnet_v1')
    parser.add_argument('--max_nrof_epochs', type=int,
        help='Number of epochs to run.', default=40)
    parser.add_argument('--batch_size', type=int,
        help='Number of images to process in a batch.', default=90)
    parser.add_argument('--image_size', type=int,
        help='Image size (height, width) in pixels.', default=160)
    parser.add_argument('--epoch_size', type=int,
        help='Number of batches per epoch.', default=1000)
    parser.add_argument('--alpha', type=float,
        help='Positive to negative triplet distance margin.', default=0.2)
    parser.add_argument('--num_attribute', type=int,
        help='', default=1)
    parser.add_argument('--embedding_size', type=int,
        help='Dimensionality of the embedding.', default=128)
    parser.add_argument('--keep_probability', type=float,
        help='Keep probability of dropout for the fully connected layer(s).', default=1.0)
    parser.add_argument('--weight_decay', type=float,
        help='L2 weight regularization.', default=0.0)
    parser.add_argument('--optimizer', type=str, choices=['ADAGRAD', 'ADADELTA', 'ADAM', 'RMSPROP', 'MOM'],
        help='The optimization algorithm to use', default='ADAGRAD')
    parser.add_argument('--learning_rate', type=float,
        help='Initial learning rate. If set to a negative value a learning rate ' +
        'schedule can be specified in the file "learning_rate_schedule.txt"', default=0.1)
    parser.add_argument('--learning_rate_decay_epochs', type=int,
        help='Number of epochs between learning rate decay.', default=100)
    parser.add_argument('--learning_rate_decay_factor', type=float,
        help='Learning rate decay factor.', default=1.0)
    parser.add_argument('--moving_average_decay', type=float,
        help='Exponential decay for tracking of training parameters.', default=0.9999)
    parser.add_argument('--seed', type=int,
        help='Random seed.', default=666)
    return parser.parse_args(argv)

if __name__ == '__main__':
    main(parse_arguments(sys.argv[1:]))


# Copyright (c) 2019-2020, NVIDIA CORPORATION.

import datetime as dt

import numpy as np
import pandas as pd
import pyarrow as pa

import cudf._lib as libcudf
from cudf._lib.nvtx import annotate
from cudf.core.buffer import Buffer
from cudf.core.column import column
from cudf.utils import utils
from cudf.utils.dtypes import is_scalar, np_to_pa_dtype

# nanoseconds per time_unit
_numpy_to_pandas_conversion = {
    "ns": 1,
    "us": 1000,
    "ms": 1000000,
    "s": 1000000000,
    "m": 60000000000,
    "h": 3600000000000,
    "D": 1000000000 * 86400,
}

class DatetimeColumn(column.ColumnBase):
    def __init__(
        self, data, dtype, mask=None, size=None, offset=0, null_count=None
    ):
        """
        Parameters
        ----------
        data : Buffer
            The datetime values
        dtype : np.dtype
            The data type
        mask : Buffer; optional
            The validity mask
        """
        dtype = np.dtype(dtype)
        if data.size % dtype.itemsize:
            raise ValueError("Buffer size must be divisible by element size")
        if size is None:
            size = data.size // dtype.itemsize
            size = size - offset
        super().__init__(
            data,
            size=size,
            dtype=dtype,
            mask=mask,
            offset=offset,
            null_count=null_count,
        )
        assert self.dtype.type is np.datetime64
        self._time_unit, _ = np.datetime_data(self.dtype)

    def __contains__(self, item):
        # Handles improper item types
        try:
            item = np.datetime64(item, self._time_unit)
        except Exception:
            return False
        return item.astype("int_") in self.as_numerical

    @property
    def time_unit(self):
        return self._time_unit

    @property
    def year(self):
        return self.get_dt_field("year")

    @property
    def month(self):
        return self.get_dt_field("month")

    @property
    def day(self):
        return self.get_dt_field("day")

    @property
    def hour(self):
        return self.get_dt_field("hour")

    @property
    def minute(self):
        return self.get_dt_field("minute")

    @property
    def second(self):
        return self.get_dt_field("second")

    @property
    def weekday(self):
        return self.get_dt_field("weekday")

    def get_dt_field(self, field):
        return libcudf.datetime.extract_datetime_component(self, field)

    def normalize_binop_value(self, other):
        if isinstance(other, dt.datetime):
            other = np.datetime64(other)

        if isinstance(other, pd.Timestamp):
            m = _numpy_to_pandas_conversion[self.time_unit]
            ary = utils.scalar_broadcast_to(
                other.value * m, size=len(self), dtype=self.dtype
            )
        elif isinstance(other, np.datetime64):
            other = other.astype(self.dtype)
            ary = utils.scalar_broadcast_to(
                other, size=len(self), dtype=self.dtype
            )
        else:
            raise TypeError("cannot broadcast {}".format(type(other)))

        return column.build_column(data=Buffer(ary), dtype=self.dtype)

    @property
    def as_numerical(self):
        from cudf.core.column import build_column

        return build_column(
            data=self.base_data,
            dtype=np.int64,
            mask=self.base_mask,
            offset=self.offset,
            size=self.size,
        )

    def as_datetime_column(self, dtype, **kwargs):
        dtype = np.dtype(dtype)
        if dtype == self.dtype:
            return self
        return libcudf.unary.cast(self, dtype=dtype)

    def as_numerical_column(self, dtype, **kwargs):
        return self.as_numerical.astype(dtype)

    def as_string_column(self, dtype, **kwargs):
        from cudf.core.column import string

        if len(self) > 0:
            return string._numeric_to_str_typecast_functions[
                np.dtype(self.dtype)
            ](self, **kwargs)
        else:
            return column.column_empty(0, dtype="object", masked=False)

    def to_pandas(self, index=None):
        return pd.Series(
            self.to_array(fillna="pandas").astype(self.dtype), index=index
        )

    def to_arrow(self):
        mask = None
        if self.nullable:
            mask = pa.py_buffer(self.mask_array_view.copy_to_host())
        data = pa.py_buffer(self.as_numerical.data_array_view.copy_to_host())
        pa_dtype = np_to_pa_dtype(self.dtype)
        return pa.Array.from_buffers(
            type=pa_dtype,
            length=len(self),
            buffers=[mask, data],
            null_count=self.null_count,
        )

    def default_na_value(self):
        """Returns the default NA value for this column
        """
        dkind = self.dtype.kind
        if dkind == "M":
            return np.datetime64("nat", self.time_unit)
        else:
            raise TypeError(
                "datetime column of {} has no NaN value".format(self.dtype)
            )

    def binary_operator(self, op, rhs, reflect=False):
        lhs, rhs = self, rhs

        if op in ("eq", "ne", "lt", "gt", "le", "ge"):
            out_dtype = np.bool
        else:
            raise TypeError(
                f"Series of dtype {self.dtype} cannot perform "
                f" the operation {op}"
            )
        return binop(lhs, rhs, op=op, out_dtype=out_dtype)

    def fillna(self, fill_value):
        if is_scalar(fill_value):
            fill_value = np.datetime64(fill_value, self.time_unit)
        else:
            fill_value = column.as_column(fill_value, nan_as_null=False)

        result = libcudf.replace.replace_nulls(self, fill_value)
        result = column.build_column(
            result.base_data,
            result.dtype,
            mask=None,
            offset=result.offset,
            size=result.size,
        )

        return result

    def min(self, dtype=None):
        return libcudf.reduce.reduce("min", self, dtype=dtype)

    def max(self, dtype=None):
        return libcudf.reduce.reduce("max", self, dtype=dtype)

    def find_first_value(self, value, closest=False):
        """
        Returns offset of first value that matches
        """
        value = pd.to_datetime(value)
        value = column.as_column(value).as_numerical[0]
        return self.as_numerical.find_first_value(value, closest=closest)

    def find_last_value(self, value, closest=False):
        """
        Returns offset of last value that matches
        """
        value = pd.to_datetime(value)
        value = column.as_column(value).as_numerical[0]
        return self.as_numerical.find_last_value(value, closest=closest)

    @property
    def is_unique(self):
        return self.as_numerical.is_unique

    def can_cast_safely(self, to_dtype):
        if np.issubdtype(to_dtype, np.datetime64):

            to_res, _ = np.datetime_data(to_dtype)
            self_res, _ = np.datetime_data(self.dtype)

            max_int = np.iinfo(np.dtype("int64")).max

            max_dist = self.max().astype(np.timedelta64, copy=False)
            min_dist = self.min().astype(np.timedelta64, copy=False)

            self_delta_dtype = np.timedelta64(0, self_res).dtype

            if max_dist <= np.timedelta64(max_int, to_res).astype(
                self_delta_dtype
            ) and min_dist <= np.timedelta64(max_int, to_res).astype(
                self_delta_dtype
            ):
                return True
            else:
                return False
        elif to_dtype == np.dtype("int64") or to_dtype == np.dtype("O"):
            # can safely cast to representation, or string
            return True
        else:
            return False

@annotate("BINARY_OP", color="orange", domain="cudf_python")
def binop(lhs, rhs, op, out_dtype):
    out = libcudf.binaryop.binaryop(lhs, rhs, op, out_dtype)
    return out

def infer_format(element, **kwargs):
    """
    Infers datetime format from a string, also takes cares for `ms` and `ns`
    """
    import re

    fmt = pd.core.tools.datetimes._guess_datetime_format(element, **kwargs)

    if fmt is not None:
        return fmt

    element_parts = element.split(".")
    if len(element_parts) != 2:
        raise ValueError("Unable to infer the timestamp format from the data")

    # There is possibility that the element is of following format
    # '00:00:03.333333 2016-01-01'
    second_part = re.split(r"(\D+)", element_parts[1], maxsplit=1)
    subsecond_fmt = ".%" + str(len(second_part[0])) + "f"

    first_part = pd.core.tools.datetimes._guess_datetime_format(
        element_parts[0], **kwargs
    )
    # For the case where first_part is '00:00:03'
    if first_part is None:
        tmp = "1970-01-01 " + element_parts[0]
        first_part = pd.core.tools.datetimes._guess_datetime_format(
            tmp, **kwargs
        ).split(" ", 1)[1]
    if first_part is None:
        raise ValueError("Unable to infer the timestamp format from the data")

    if len(second_part) > 1:
        second_part = pd.core.tools.datetimes._guess_datetime_format(
            "".join(second_part[1:]), **kwargs
        )
    else:
        second_part = ""

    try:
        fmt = first_part + subsecond_fmt + second_part
    except Exception:
        raise ValueError("Unable to infer the timestamp format from the data")

    return fmt


import torch
import time
import numpy as np
from a2c_ppo_acktr import algo as algorithm
from a2c_ppo_acktr.arguments import get_args
from a2c_ppo_acktr.envs import *
from a2c_ppo_acktr.model import Policy
from a2c_ppo_acktr.storage import RolloutStorage
from a2c_ppo_acktr.utils import get_vec_normalize, update_linear_schedule

from baselines import logger

class a2c:

    def __init__(self,envs,algo='a2c',lr=7e-4,eps=1e-5,alpha=0.99,gamma=0.9,use_gae=False,
        tau=0.95,entropy_coef=0.01,value_loss_coef=0.5,max_grad_norm=0.5,seed=1,
        num_processes=1,num_steps=4,log_interval=100,save_interval=100,log_dir='/tmp/gym',
        save_dir='/trained_models',add_timestep=False,recurrent_policy=False,use_linear_lr_decay=False,
        use_linear_clip_decay = False,device=torch.device("cpu")):

        self.log_interval=log_interval
        self.gamma=gamma
        self.use_gae=use_gae
        self.tau=tau
        self.num_steps=num_steps
        self.num_processes=num_processes
        self.envs=envs
        self.num_updates=0

        #Initialize directory, cuda et compute the total number of update
        #self.num_updates = int(total_timesteps) // num_steps // num_processes  # compute number of updates from number of steps, each process make numsteps moves for one update

        # create policy model neural network and distribution
        self.actor_critic = Policy(envs.observation_space.shape, envs.action_space,
                              base_kwargs={'recurrent': recurrent_policy})
        self.actor_critic.to(device)

        #We want to use A2C
        self.agent = algorithm.A2C_ACKTR(self.actor_critic, value_loss_coef,
                               entropy_coef, lr=lr,
                               eps=eps, alpha=alpha,
                               max_grad_norm=max_grad_norm)

        # Object which generate batch of learning
        self.rollouts = RolloutStorage(num_steps, num_processes,
                                  envs.observation_space.shape, envs.action_space,
                                  self.actor_critic.recurrent_hidden_state_size)
        self.rollouts.to(device)
        self.init_env()

        self.episode_rewards = [0.]
        self.start = time.time()
        self.freq_timer=self.start

    def init_env(self,render=False):
        # Reste environment and
        obs = self.envs.reset()
        self.rollouts.obs[0].copy_(obs) #init the rollout
        if render :
            self.envs.render()

    def step(self,render=False):
        step = self.rollouts.step
        with torch.no_grad():
            value, action, action_log_prob, recurrent_hidden_states = self.actor_critic.act(
                self.rollouts.obs[step],
                self.rollouts.recurrent_hidden_states[step],
                self.rollouts.masks[step])

        # Obser reward and next obs
        obs, reward, done, infos = self.envs.step(action)
        if render:
            self.envs.render()

        #For logging mean rewards for the the last interactions
        self.episode_rewards[-1] += reward[0][0].item()
        if done[0]:
            self.episode_rewards.append(0.)

        # If done then clean the history of observations.
        masks = torch.FloatTensor([[0.0] if done_ else [1.0]
                                   for done_ in done])

        # add interaction in rollout to learn after
        self.rollouts.insert(obs, recurrent_hidden_states, action, action_log_prob, value, reward, masks)
        return obs,action,reward,done

    def nstep(self,render=False):
        #num steps is the number of steps for one update
        for _ in range(self.num_steps):
            obs,action,reward,done= self.step(render=render)
        return obs,action,reward,done

    def init_steps(self):
        # init the first observation
        self.rollouts.after_update()

    def learn(self):
        assert len(self.rollouts.obs)>1,"Must interact with environment before learning"
        self.num_updates+=1

        # compute value function for the last interaction only
        with torch.no_grad():
            next_value = self.actor_critic.get_value(self.rollouts.obs[-1],
                                                self.rollouts.recurrent_hidden_states[-1],
                                                self.rollouts.masks[-1]).detach()

        # compute the target for each interaction, the last is value
        self.rollouts.compute_returns(next_value, self.use_gae, self.gamma, self.tau)

        # Do learning step
        value_loss, action_loss, dist_entropy = self.agent.update(self.rollouts)
        self.init_steps()

        if self.num_updates % self.log_interval == 0:
            # Calculate the fps (frame per second)
            fps = int((self.num_steps * self.num_processes * self.log_interval) / (time.time() - self.freq_timer))

            # Calculates if value function is a good predicator of the returns (ev > 1)
            # or if it's just worse than predicting nothing (ev =< 0)
            # ev = explained_variance(values, rewards)
            logger.record_tabular("mean_rewards", np.mean(self.episode_rewards[-self.log_interval:-2]))
            logger.record_tabular("nupdates", self.num_updates)
            logger.record_tabular("total_timesteps", self.num_steps * self.num_processes * self.num_updates)
            logger.record_tabular("fps", fps)
            logger.record_tabular("policy_entropy", float(dist_entropy))
            logger.record_tabular("value_loss", float(value_loss))
            logger.record_tabular("total_time", float(time.time() - self.start))
            logger.dump_tabular()
            self.freq_timer = time.time()

        return self.actor_critic


import os

import cv2
import numpy as np
import pandas as pd
import tensorflow as tf

def read_image(path, size):
    """
    Load image from local storage
    :param path: image path
    :param size: image size
    :return: loaded image
    """
    image = cv2.imread(path, cv2.IMREAD_COLOR)
    image = cv2.resize(image, (size, size))
    image = image / 255.0
    image = image.astype(np.float32)
    return image

def recognize_image(img_path):
    """
    Recognize image
    Taking image and predicting breed basing on trained model
    :param img_path: Image Path
    :return: top 4 matching breeds, most similar breed
    """
    path = "input/"
    train_path = os.path.join(path, "train/*")
    test_path = os.path.join(path, "test/*")
    labels_path = os.path.join(path, "labels.csv")

    labels_df = pd.read_csv(labels_path)
    breed = labels_df["breed"].unique()
    id2breed = {i: name for i, name in enumerate(breed)}

    ## Model
    model = tf.keras.models.load_model("model.h5")

    image = read_image(img_path, 224)
    image = np.expand_dims(image, axis=0)
    pred = model.predict(image)[0]
    label_idx = np.argmax(pred)
    top3 = np.argsort(pred)[-4:][::-1]
    possible_breed = list()
    print(str(id2breed[top3[0]]).replace("_", " "))
    possible_breed.append(str(id2breed[top3[0]]).replace("_", " "))
    possible_breed.append(str(id2breed[top3[1]]).replace("_", " "))
    possible_breed.append(str(id2breed[top3[2]]).replace("_", " "))
    possible_breed.append(str(id2breed[top3[3]]).replace("_", " "))
    return str(id2breed[label_idx]).replace("_", " "), possible_breed

if __name__ == "__main__":
    print(recognize_image())