import os
import numbers
import numpy as np
from collections import OrderedDict
from tensorflow.keras.models import load_model
from pkg_resources import resource_filename
from distutils.spawn import find_executable

from astrorapid.process_light_curves import read_multiple_light_curves
from astrorapid.prepare_arrays import PrepareInputArrays

try:
    import matplotlib.pyplot as plt
    from matplotlib.ticker import MaxNLocator
    import matplotlib
    import matplotlib.animation as animation

    # Check if latex is installed
    if find_executable('latex'):
        plt.rcParams['text.usetex'] = True
    plt.rcParams['font.serif'] = ['Computer Modern Roman'] + plt.rcParams['font.serif']

except ImportError:
    print("Warning: You will need to install 'matplotlib' if you wish to plot the classifications.")

CLASS_COLOR = {'Pre-explosion': 'grey', 'SNIa-norm': 'tab:green', 'SNIbc': 'tab:orange', 'SNII': 'tab:blue',
               'SNIa-91bg': 'tab:red', 'SNIa-x': 'tab:purple', 'point-Ia': 'tab:brown', 'Kilonova': '#aaffc3',
               'SLSN-I': 'tab:olive', 'PISN': 'tab:cyan', 'ILOT': '#FF1493', 'CART': 'navy', 'TDE': 'tab:pink',
               'AGN': 'bisque'}
PB_COLOR = {'u': 'tab:blue', 'g': 'tab:blue', 'r': 'tab:orange', 'i': 'm', 'z': 'k', 'Y': 'y'}
PB_MARKER = {'g': 'o', 'r': 's'}
PB_ALPHA = {'g': 0.3, 'r': 1.}
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))

class Classify(object):
    def __init__(self, known_redshift=True, model_filepath='', passbands=('g', 'r'),
                 class_names=('Pre-explosion', 'SNIa-norm', 'SNIbc', 'SNII', 'SNIa-91bg', 'SNIa-x', 'Kilonova', 'SLSN-I', 'TDE'),
                 nobs=50, mintime=-70, maxtime=80, timestep=3.0, bcut=False, zcut=None, graph=None, model=None):
        """ Takes a list of photometric information and classifies light curves as a function of time

        Parameters
        ----------
        known_redshift : bool
            Different model to be used if redshift is not known.
        model_filepath : str
            Optional argument. The model is taken from the pre-trained model ZTF model if not specified.
        passbands : tuple
            Optional argument. A tuple listing each passband. E.g. ('g', 'r').
        class_names : tuple
            List of class names that the model has been trained on. Note that this must be in the same order
            as used in training for the model specified in the argument model_filepath.
            If you are using the default model, leave this argument out.
        nobs : int
            Number of points to use in interpolation of light curve between mintime and maxtime.
            Do not change this argument unless you are using your own model_filepath
            and have changed this during training.
        mintime : int
            Days from trigger (minimum) to extract from light curve.
        maxtime : int
            Days from trigger (maximum) to extract from light curve.
        timestep : float
            Time-step between interpolated points in light curve.
            Do not change this argument unless you are using your own model_filepath
            and have changed this during training.
        bcut : bool
            Cut on galactic latitude.
            Do not set unless you know what you are doing.
        zcut : float or None
            Remove redshifts above this value.
            Do not set unless you know what you are doing.
        graph : tensorflow graph
            Do not set unless you know what you are doing.
            If you are running astrorapid in multiple threads you may need to predefine this
            This would have been created with Tensorflow i.e. graph = tf.get_default_graph()
        model : tensorflow model
            Do not set unless you know what you are doing.
            If you are running astrorapid in multiple threads you may need to predefine this
            This would have been created with keras' load_model function e.g. model = load_model('keras_model.hdf5')

        """
        self.known_redshift = known_redshift
        self.passbands = passbands
        self.bcut = bcut
        self.zcut = zcut
        self.class_names = class_names
        self.nobs = nobs
        self.mintime = mintime
        self.maxtime = maxtime
        self.timestep = timestep

        if self.known_redshift:
            self.contextual_info = ('redshift',)
            filename = 'keras_model_with_redshift.hdf5'
        else:
            self.contextual_info = ()
            filename = 'keras_model_no_redshift.hdf5'

        if model_filepath != '' and os.path.exists(model_filepath):
            self.model_filepath = model_filepath
        else:
            if not (model_filepath == '' and os.path.exists(model_filepath)):
                print("Invalid keras model. Using default model...")
            self.model_filepath = os.path.join(SCRIPT_DIR, filename)
            if not os.path.exists(self.model_filepath):
                self.model_filepath = resource_filename(__name__, 'keras_model_with_redshift.hdf5')

        print(self.model_filepath)
        self.graph = graph
        if graph is not None and model is not None:
            self.model = model
        else:
            self.model = load_model(self.model_filepath)

    def process_light_curves(self, light_curves):
        processed_lightcurves = read_multiple_light_curves(light_curves, known_redshift=self.known_redshift,
                                                           training_set_parameters=None)
        prepareinputarrays = PrepareInputArrays(self.passbands, self.contextual_info, self.bcut, self.zcut,
                                                self.nobs, self.mintime, self.maxtime, self.timestep)
        X, orig_lc, timesX, objids_list, trigger_mjds = prepareinputarrays.prepare_input_arrays(processed_lightcurves)

        # # REMOVE CORRECTION FACTOR IF NOT USED
        # correction_factor = np.load('astrorapid/correction_factor.npy')
        # for i, pb in enumerate(self.passbands):
        #     X[:, :, i] = X[:, :, i] / correction_factor[i]

        return X, orig_lc, timesX, objids_list, trigger_mjds

    def _do_error_checks(self, light_curves):
        assert isinstance(light_curves, (list, np.ndarray))
        for light_curve in light_curves:
            assert len(light_curve) == 10
            mjd, flux, fluxerr, passband, photflag, ra, dec, objid, redshift, mwebv = light_curve
            _lcshape = len(mjd)
            assert _lcshape == len(flux)
            assert _lcshape == len(fluxerr)
            assert _lcshape == len(passband)
            assert _lcshape == len(photflag)
            assert all(isinstance(n, numbers.Number) and np.isfinite(n) for n in mjd)
            assert all(isinstance(n, numbers.Number) and np.isfinite(n) for n in flux)
            assert all(isinstance(n, numbers.Number) and np.isfinite(n) for n in fluxerr)
            assert all(isinstance(n, numbers.Number) and np.isfinite(n) for n in photflag)
            assert all(isinstance(n, str) for n in passband)
            assert isinstance(ra, numbers.Number)
            assert isinstance(dec, numbers.Number)
            assert isinstance(mwebv, numbers.Number)
            assert np.isfinite(ra)
            assert np.isfinite(dec)
            assert np.isfinite(mwebv)
            if self.known_redshift:
                assert isinstance(redshift, numbers.Number)
                assert np.isfinite(redshift)

    def get_predictions(self, light_curves, return_predictions_at_obstime=False, return_objids=False):
        """ Return the classification accuracies as a function of time for each class

        Parameters
        ----------
        light_curves : list
            Is a list of tuples. Each tuple contains the light curve information of a transient object in the form
            (mjd, flux, fluxerr, passband, photflag, ra, dec, objid, redshift, mwebv).
            Here, mjd, flux, fluxerr, passband, and photflag are arrays.
            And ra, dec, objid, redshift, and mwebv are floats.
                mjd: list of Modified Julian Dates of light curve

                flux: list of fluxes at each mjd

                fluxerr: list of flux errors

                passband: list of strings indicating the passband. 'r' or 'g' for r-band or g-band observations.

                photflag: list of flags identifying whether the observation is a detection (4096), non-detection (0),
                or the first detection (6144).

                ra: Right Ascension (float value).

                dec: Declination (float value).

                objid: Object Identifier (String).

                redshift: Cosmological redshift of object (float). Set to NoneType if redshift is unknown.

                mwebv: Milky way extinction.

        return_predictions_at_obstime: bool
            Return the predictions at the observation times instead of at the 50 interpolated timesteps.
        return_objids : bool, optional
            If True, also return the object IDs (objids) in the same order as the returned predictions.

        Returns
        -------
        y_predict: array
            Classification probability vector at each time step for each object.
            Array of shape (s, n, m) is returned.
            Where s is the number of obejcts that are classified,
            n is the number of times steps, and m is the number of classes.
        time_steps: array
            MJD time steps corresponding to the timesteps of the y_predict array.
        objids : array, optional
            The object ids (objids) that were input into light_curves are returned in the same order as y_predict.
            Only provided if return_objids is True.
        """

        self._do_error_checks(light_curves)

        self.X, self.orig_lc, self.timesX, self.objids, self.trigger_mjds = self.process_light_curves(light_curves)
        nobjects = len(self.objids)

        if nobjects == 0:
            print("No objects to classify. These may have been removed from the chosen selection cuts")
            if return_objids:
                return None, None, self.objids
            else:
                return None, None

        if self.graph is not None:
            with self.graph.as_default():
                self.y_predict = self.model.predict(self.X)
        else:
            self.y_predict = self.model.predict(self.X)

        argmax = self.timesX.argmax(axis=1) + 1

        if return_predictions_at_obstime:
            (s, n, m) = self.y_predict.shape  # (s, n, m) = (num light curves, num timesteps, num classes)
            y_predict = []
            time_steps = []
            for idx in range(s):
                obs_time = []
                lc_data = self.orig_lc[idx]
                used_passbands = [pb for pb in self.passbands if pb in lc_data['passband']]
                for pb in used_passbands:
                    pbmask = lc_data['passband'] == pb
                    if pb in self.orig_lc[idx]:
                        obs_time.append(lc_data[pbmask]['time'].data)
                obs_time = np.array(obs_time)
                obs_time = np.sort(obs_time[~np.isnan(obs_time)])
                y_predict_at_obstime = []
                for classnum, classname in enumerate(self.class_names):
                    y_predict_at_obstime.append(np.interp(obs_time, self.timesX[idx][:argmax[idx]], self.y_predict[idx][:, classnum][:argmax[idx]]))
                y_predict.append(np.array(y_predict_at_obstime).T)
                time_steps.append(obs_time + self.trigger_mjds[idx])
        else:
            y_predict = [self.y_predict[i][:argmax[i]] for i in range(nobjects)]
            time_steps = [self.timesX[i][:argmax[i]] + self.trigger_mjds[i] for i in range(nobjects)]

        if return_objids:
            return y_predict, time_steps, self.objids

        return y_predict, time_steps

    def plot_light_curves_and_classifications(self, indexes_to_plot=None, step=True, use_interp_flux=False, figdir='.',
                                              plot_matrix_input=False, light_curves=None):
        """
        Plot light curve (top panel) and classifications (bottom panel) vs time.

        Parameters
        ----------
        indexes_to_plot : tuple
            The indexes listed in the tuple will be plotted according to the order of the input light curves.
            E.g. (0, 1, 3, 5) will plot the zeroth, first, third and fifth light curves and classifications.
            If None or True, then all light curves will be plotted
        step : bool
            Plot step function along data points instead of interpolating classifications between data.
        use_interp_flux : bool
            Use all 50 timesteps when plotting classification probabilities rather than just at the timesteps with data.
        figdir : str
            Directory to save figure.
        plot_matrix_input : bool
            Plots the interpolated light curve passed into the neural network on top of the observations.
        light_curves : list
            This argument is only required if the get_predictions() method has not been run.
            Is a list of tuples. Each tuple contains the light curve information of a transient object in the form
            (mjd, flux, fluxerr, passband, photflag, ra, dec, objid, redshift, mwebv).
            Here, mjd, flux, fluxerr, passband, and photflag are arrays.
            ra, dec, objid, redshift, and mwebv are floats

        """

        font = {'family': 'normal',
                'size': 33}
        matplotlib.rc('font', **font)

        if indexes_to_plot is None or indexes_to_plot is True:
            indexes_to_plot = np.arange(len(self.y_predict))

        if not hasattr(self, 'y_predict'):
            self.get_predictions(light_curves)

        if plot_matrix_input:
            alpha_observations = 0.2
        else:
            alpha_observations = 1.
        for idx in indexes_to_plot:
            argmax = self.timesX[idx].argmax() + 1
            fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, figsize=(13, 15),
                                           num="classification_vs_time_{}".format(idx), sharex=True)
            # ax1.axvline(x=0, color='k', linestyle='-', linewidth=1)
            # ax2.axvline(x=0, color='k', linestyle='-', linewidth=1)

            lc_data = self.orig_lc[idx]
            used_passbands = [pb for pb in self.passbands if pb in lc_data['passband']]

            for pbidx, pb in enumerate(used_passbands):
                pbmask = lc_data['passband'] == pb
                ax1.errorbar(lc_data[pbmask]['time'], lc_data[pbmask]['flux'],
                             yerr=lc_data[pbmask]['fluxErr'], fmt=PB_MARKER[pb], label=pb,
                             c=PB_COLOR[pb], lw=3, markersize=10, alpha=alpha_observations)
                if plot_matrix_input:
                    ax1.plot(self.timesX[idx][:argmax], self.X[idx][:, pbidx][:argmax], c=PB_COLOR[pb], lw=3)

            new_t = np.concatenate([lc_data[lc_data['passband'] == pb]['time'].data for pb in used_passbands])
            new_t = np.sort(new_t[~np.isnan(new_t)])
            if not use_interp_flux and not plot_matrix_input:
                new_y_predict = []
                for classnum, classname in enumerate(self.class_names):
                    new_y_predict.append(np.interp(new_t, self.timesX[idx][:argmax], self.y_predict[idx][:, classnum][:argmax]))

            for classnum, classname in enumerate(self.class_names):
                if use_interp_flux or plot_matrix_input:
                    ax2.plot(self.timesX[idx][:argmax], self.y_predict[idx][:, classnum][:argmax], '-', label=classname,
                             color=CLASS_COLOR[classname], linewidth=3)
                else:
                    if step:
                        ax2.step(new_t, new_y_predict[classnum], '-', label=classname, color=CLASS_COLOR[classname], linewidth=3, where='post')
                    else:
                        ax2.plot(new_t, new_y_predict[classnum], '-', label=classname, color=CLASS_COLOR[classname], linewidth=3,)

            ax1.legend(frameon=True, fontsize=33)#, loc='lower right')
            ax2.legend(frameon=True, fontsize=21.5)#, loc='center right')  # , ncol=2)
            ax2.set_xlabel("Days since trigger (rest frame)")  # , fontsize=18)
            ax1.set_ylabel("Relative Flux")  # , fontsize=15)
            ax2.set_ylabel("Class Probability")  # , fontsize=18)
            # ax1.set_ylim(-0.1, 1.1)
            ax2.set_ylim(0, 1)
            ax1.set_xlim(left=min(new_t))  # ax1.set_xlim(-70, 80)
            # ax1.grid(True)
            # ax2.grid(True)
            plt.setp(ax1.get_xticklabels(), visible=False)
            ax2.yaxis.set_major_locator(MaxNLocator(nbins=6, prune='upper'))  # added
            plt.tight_layout()
            fig.subplots_adjust(hspace=0)
            savename = 'classification_vs_time_{}{}{}.pdf'.format(self.objids[idx], '_step' if step else '', '_no_interp' if not use_interp_flux else '')
            plt.savefig(os.path.join(figdir, savename))
            # plt.savefig("{}.png".format(savename))
            plt.close()

        return self.orig_lc, self.timesX, self.y_predict

    def plot_classification_animation(self, indexes_to_plot=None, figdir='.', light_curves=None):
        """ Plot light curve (top panel) and classifications (bottom panel) vs time as an mp4 animation.

        Parameters
        ----------
        indexes_to_plot : tuple
            The indexes listed in the tuple will be plotted according to the order of the input light curves.
            E.g. (0, 1, 3, 5) will plot the zeroth, first, third and fifth light curves and classifications.
            If None or True, then all light curves will be plotted
        figdir : str
            Directory to save figure.
        light_curves : list
            This argument is only required if the get_predictions() method has not been run.
            Is a list of tuples. Each tuple contains the light curve information of a transient object in the form
            (mjd, flux, fluxerr, passband, photflag, ra, dec, objid, redshift, mwebv).
            Here, mjd, flux, fluxerr, passband, and photflag are arrays.
            ra, dec, objid, redshift, and mwebv are floats

        """

        font = {'family': 'normal',
                'size': 33}
        matplotlib.rc('font', **font)

        if indexes_to_plot is None or indexes_to_plot is True:
            indexes_to_plot = np.arange(len(self.y_predict))

        if not hasattr(self, 'y_predict'):
            self.get_predictions(light_curves)

        for idx in indexes_to_plot:
            lc_data = self.orig_lc[idx]
            used_passbands = [pb for pb in self.passbands if pb in lc_data['passband']]

            new_t = np.concatenate([lc_data[lc_data['passband'] == pb]['time'].data for pb in used_passbands])
            all_flux = list(lc_data[lc_data['passband'] == 'g']['flux']) + list(lc_data[lc_data['passband'] == 'r']['flux'])

            argmax = self.timesX[idx].argmax() + 1
            fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, figsize=(13, 15),
                                           num="animation_classification_vs_time_{}".format(idx), sharex=True)

            ax1.legend(frameon=False, fontsize=33)
            ax2.legend(frameon=False, fontsize=23.5)  # , ncol=2)
            ax2.set_xlabel("Days since trigger (rest frame)")  # , fontsize=18)
            ax1.set_ylabel("Relative Flux")  # , fontsize=15)
            ax2.set_ylabel("Class Probability")  # , fontsize=18)
            ax1.set_ylim(min(all_flux), max(all_flux))
            ax2.set_ylim(0, 1)
            ax1.set_xlim(min(new_t), max(new_t))  # ax1.set_xlim(-70, 80)
            plt.setp(ax1.get_xticklabels(), visible=False)
            ax2.yaxis.set_major_locator(MaxNLocator(nbins=6, prune='upper'))  # added
            plt.tight_layout()
            fig.subplots_adjust(hspace=0)
            ax1.axvline(x=0, color='k', linestyle='-', linewidth=1)
            ax2.axvline(x=0, color='k', linestyle='-', linewidth=1)

            Writer = animation.writers['ffmpeg']
            writer = Writer(fps=5, bitrate=1800)

            def animate(i):
                for pbidx, pb in enumerate(used_passbands):
                    ax1.plot(self.timesX[idx][:argmax][:int(i + 1)], self.X[idx][:, pbidx][:argmax][:int(i + 1)],
                             label=pb, c=PB_COLOR[pb], lw=3)  # , markersize=10, marker=MARKPB[pb])

                for classnum, classname in enumerate(self.class_names):
                    ax2.plot(self.timesX[idx][:argmax][:int(i + 1)],
                             self.y_predict[idx][:, classnum][:argmax][:int(i + 1)],
                             '-', label=classname, color=CLASS_COLOR[classname], linewidth=3)

                # Don't repeat legend items
                ax1.legend(frameon=False, fontsize=33)
                ax2.legend(frameon=False, fontsize=23.5)
                handles1, labels1 = ax1.get_legend_handles_labels()
                handles2, labels2 = ax2.get_legend_handles_labels()
                by_label1 = OrderedDict(zip(labels1, handles1))
                by_label2 = OrderedDict(zip(labels2, handles2))
                ax1.legend(by_label1.values(), by_label1.keys(), frameon=False, fontsize=33, loc='lower right')
                ax2.legend(by_label2.values(), by_label2.keys(), frameon=False, fontsize=21.5, loc='center right')

            ani = animation.FuncAnimation(fig, animate, frames=50, repeat=True)
            savename = os.path.join(figdir,'classification_vs_time_{}.mp4'.format(self.objids[idx]))
            ani.save(savename, writer=writer)

    def plot_classification_animation_step(self, indexes_to_plot=None, figdir='.', light_curves=None):
        """
        Plot light curve (top panel) and classifications (bottom panel) vs time as an mp4 animation
        as step function.

        Parameters
        ----------
        indexes_to_plot : tuple
            The indexes listed in the tuple will be plotted according to the order of the input light curves.
            E.g. (0, 1, 3, 5) will plot the zeroth, first, third and fifth light curves and classifications.
            If None or True, then all light curves will be plotted
        figdir : str
            Directory to save figure.
        light_curves : list
            This argument is only required if the get_predictions() method has not been run.
            Is a list of tuples. Each tuple contains the light curve information of a transient object in the form
            (mjd, flux, fluxerr, passband, photflag, ra, dec, objid, redshift, mwebv).
            Here, mjd, flux, fluxerr, passband, and photflag are arrays.
            ra, dec, objid, redshift, and mwebv are floats

        """

        font = {'family': 'normal',
                'size': 33}
        matplotlib.rc('font', **font)

        if indexes_to_plot is None or indexes_to_plot is True:
            indexes_to_plot = np.arange(len(self.y_predict))

        if not hasattr(self, 'y_predict'):
            self.get_predictions()

        for idx in indexes_to_plot:
            lc_data = self.orig_lc[idx]
            used_passbands = [pb for pb in self.passbands if pb in lc_data['passband']]

            new_t = np.concatenate([lc_data[lc_data['passband'] == pb]['time'].data for pb in used_passbands])
            new_t = np.sort(new_t[~np.isnan(new_t)])
            new_y_predict = []
            all_flux = list(lc_data[lc_data['passband'] == 'g']['flux']) + list(lc_data[lc_data['passband'] == 'r']['flux'])

            argmax = self.timesX[idx].argmax() + 1
            fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, figsize=(13, 15),
                                           num="animation_classification_vs_time_{}".format(idx), sharex=True)

            ax1.legend(frameon=False, fontsize=33)
            ax2.legend(frameon=False, fontsize=23.5)  # , ncol=2)
            ax2.set_xlabel("Days since trigger (rest frame)")  # , fontsize=18)
            ax1.set_ylabel("Relative Flux")  # , fontsize=15)
            ax2.set_ylabel("Class Probability")  # , fontsize=18)
            ax1.set_ylim(min(all_flux), max(all_flux))
            ax2.set_ylim(0, 1)
            ax1.set_xlim(min(new_t), max(new_t))
            plt.setp(ax1.get_xticklabels(), visible=False)
            ax2.yaxis.set_major_locator(MaxNLocator(nbins=6, prune='upper'))  # added
            plt.tight_layout()
            fig.subplots_adjust(hspace=0)
            ax1.axvline(x=0, color='k', linestyle='-', linewidth=1)
            ax2.axvline(x=0, color='k', linestyle='-', linewidth=1)

            Writer = animation.writers['ffmpeg']
            writer = Writer(fps=2, bitrate=1800)

            for classnum, classname in enumerate(self.class_names):
                new_y_predict.append(np.interp(new_t, self.timesX[idx][:argmax], self.y_predict[idx][:, classnum][:argmax]))

            def animate(i):
                for pbidx, pb in enumerate(used_passbands):
                    # ax1.plot(self.timesX[idx][:argmax][:int(i + 1)], self.X[idx][:, pbidx][:argmax][:int(i + 1)],
                    #          label=pb, c=PB_COLOR[pb], lw=3)  # , markersize=10, marker=MARKPB[pb])

                    # dea = [self.orig_lc[idx][pb]['time'] < self.timesX[idx][:argmax][int(i)]]
                    if i + 1 >= len(new_t):
                        break

                    pbmask = lc_data['passband'] == pb
                    dea = [lc_data[pbmask]['time'] < new_t[int(i+1)]]

                    ax1.errorbar(np.array(lc_data[pbmask]['time'])[dea], np.array(lc_data[pbmask]['flux'])[dea],
                                 yerr=np.array(lc_data[pbmask]['fluxErr'])[dea], fmt=PB_MARKER[pb], label=pb,
                                 c=PB_COLOR[pb], lw=3, markersize=10)

                for classnum, classname in enumerate(self.class_names):
                    # ax2.plot(self.timesX[idx][:argmax][:int(i + 1)],
                    #          self.y_predict[idx][:, classnum][:argmax][:int(i + 1)],
                    #          '-', label=classname, color=CLASS_COLOR[classname], linewidth=3)
                    ax2.step(new_t[:int(i + 1)],
                             new_y_predict[classnum][:int(i + 1)],
                             '-', label=classname, color=CLASS_COLOR[classname], linewidth=3, where='post')

                # Don't repeat legend items
                ax1.legend(frameon=False, fontsize=33)
                ax2.legend(frameon=False, fontsize=23.5)
                handles1, labels1 = ax1.get_legend_handles_labels()
                handles2, labels2 = ax2.get_legend_handles_labels()
                by_label1 = OrderedDict(zip(labels1, handles1))
                by_label2 = OrderedDict(zip(labels2, handles2))
                ax1.legend(by_label1.values(), by_label1.keys(), frameon=False, fontsize=33, loc='lower right')
                ax2.legend(by_label2.values(), by_label2.keys(), frameon=False, fontsize=21.5, loc='center right')

            ani = animation.FuncAnimation(fig, animate, frames=len(new_t), repeat=True)
            savename = os.path.join(figdir, 'classification_vs_time_{}_step.mp4'.format(self.objids[idx]))
            ani.save(savename, writer=writer)


import scipy.special

# Old sph_* removed in SciPy 1.0:
# https://docs.scipy.org/doc/scipy/reference/release.1.0.0.html#backwards-incompatible-changes
old_special_funcs = [
    'sph_jn',
    'sph_yn',
    'sph_jnyn',
    'sph_in',
    'sph_kn',
    'sph_inkn',
]

original_version = scipy.__version__
try:

    # Fool how PyDSTool checks SciPy's version number:
    scipy.__version__ = '0.9'

    # PyDSTool tries to access `scipy.special.sph_*`; let's not fail
    # by that by setting them to None.  Those functions won't be
    # usable, but at least (hopefully) other PyDSTool functionalities
    # are usable:
    for name in old_special_funcs:
        if not hasattr(scipy.special, name):
            setattr(scipy.special, name, None)

    import PyDSTool
finally:
    scipy.__version__ = original_version
    for name in old_special_funcs:
        if getattr(scipy.special, name) is None:
            delattr(scipy.special, name)


import cv2
import numpy as np

import utils.tracking_utils as tu
from utils.config import cfg
from model.tracker_model import KalmanBoxTracker

class Tracker2D:

    def __init__(self,
                 dataset,
                 max_depth=150,
                 max_age=1,
                 min_hits=3,
                 affinity_threshold=0.3,
                 deep_sort=False,
                 kf2d=False,
                 verbose=False,
                 visualize=False):

        self.dataset = dataset
        self.max_depth = max_depth
        self.max_age = max_age
        self.min_hits = min_hits
        self.affinity_threshold = affinity_threshold
        self.trackers = []
        self.frame_count = 0

        self.current_frame = {}

        self.n_gt = 0
        self.n_FP = 0
        self.n_FN = 0
        self.n_mismatch = 0
        self.det_dim = 5  # [x1, y1, x2, y2, conf]
        self.feat_dim = 128  # feature dim = 128
        self.det_thresh = 0.9

        DATASET = cfg.GTA if dataset == 'gta' else cfg.KITTI
        self.H = DATASET.H
        self.W = DATASET.W
        self.FOCAL = DATASET.FOCAL_LENGTH

        self.kf2d = kf2d
        self.deep_sort = deep_sort
        self.verbose = verbose
        self.visualize = visualize

        self.iou_affinity_weight = 0.7 if deep_sort else 1.0
        self.feat_affinity_weight = 1.0 - self.iou_affinity_weight

        if verbose:
            print("IOU weight: {}, Feat weight: {}".format(
                self.iou_affinity_weight,
                self.feat_affinity_weight))

    def update(self, data):
        # frame information here
        ret = []
        self.frame_count += 1
        self.cam_rot = data['cam_rot'].squeeze() # In rad
        self.cam_coord = data['cam_loc'].squeeze()
        if self.frame_count == 1:
            self.init_coord = self.cam_coord.copy()
        self.cam_coord -= self.init_coord
        self.cam_calib = data['cam_calib'].squeeze()
        self.cam_pose = tu.Pose(self.cam_coord, self.cam_rot)

        # process information
        dets, feats, dims, alphas, single_depths, cens, roty, world_coords = \
            self.process_info(
                data,
                det_thresh=self.det_thresh,
                max_depth=self.max_depth,
                _nms=self.dataset=='gta',
                _valid=self.dataset=='gta',
                )

        # save to current frame
        self.current_frame = {
            'bbox': dets,
            'feat': feats,
            'dim': dims,
            'alpha': alphas, # in deg
            'roty': roty, # in rad
            'depth': single_depths,
            'center': cens,
            'location': world_coords,
            'n_obj': len(alphas),
        }

        # get predicted locations from existing trackers.
        trk_dim = self.feat_dim + self.det_dim
        trks = np.zeros((len(self.trackers), trk_dim))

        # Prediction
        for t, trk in enumerate(trks):
            if self.kf2d:
                pos = self.trackers[t].predict()[0]
            else:
                pos = self.trackers[t].predict_no_effect()[0]
            trk[:self.det_dim] = [pos[0], pos[1], pos[2], pos[3], 1.0]
            trk[self.det_dim:] = self.trackers[t].feat

        # Association
        self.affinity = np.zeros((len(dets), len(trks)), dtype=np.float32)

        for d, det in enumerate(dets):
            for t, trk in enumerate(trks):
                self.affinity[
                    d, t] += self.iou_affinity_weight * tu.compute_iou(det[:4],
                                                                       trk[:4])
        if self.deep_sort and len(dets) * len(trks) > 0:
            self.affinity += self.feat_affinity_weight * tu.compute_cos_dis(
                feats, trks[:, 5:])

        matched, unmatched_dets, unmatched_trks = \
            tu.associate_detections_to_trackers(
                dets, trks, self.affinity,
                self.affinity_threshold)

        # update matched trackers with assigned detections
        for t, trk in enumerate(self.trackers):
            if t in unmatched_trks:
                trk.lost = True
                continue

            d = matched[np.where(matched[:, 1] == t)[0], 0]
            trk.update(dets[d, :][0])
            trk.det = dets[d, :][0]
            trk.feat = self.current_frame['feat'][d[0]]
            trk.dim = self.current_frame['dim'][d[0]]
            trk.alpha = self.current_frame['alpha'][d[0]]
            trk.depth = self.current_frame['depth'][d[0]]
            trk.cen = self.current_frame['center'][d[0]]
            trk.rot = self.current_frame['roty'][d[0]]

        # create and initialise new trackers for unmatched detections
        for i in unmatched_dets:
            trk = KalmanBoxTracker(dets[i, :])
            trk.det = dets[i, :]
            trk.feat = self.current_frame['feat'][i]
            trk.dim = self.current_frame['dim'][i]
            trk.alpha = self.current_frame['alpha'][i]
            trk.depth = self.current_frame['depth'][i]
            trk.cen = self.current_frame['center'][i]
            trk.rot = self.current_frame['roty'][i]
            self.trackers.append(trk)

        # Check if boxes are correct
        if self.visualize:
            img = cv2.imread(data['im_path'][0])
            h, w, _ = img.shape
            img = cv2.putText(img, str(self.frame_count), (20, 30),
                    cv2.FONT_HERSHEY_COMPLEX, 1,
                    (200, 200, 200), 2)

            for idx, trk in enumerate(self.trackers):
                box_color = (0, 150, 150) if trk.lost else (0, 255, 0)
                box_bold = 2 if trk.lost else 4
                box = trk.get_state()[0].astype('int')
                print(trk.id+1,
                    'Lost' if trk.lost else 'Tracked',
                    #trk.aff_value,
                    trk.depth,
                    )
                if trk.depth < 0 or trk.depth > self.max_depth:
                    continue
                img = cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), box_color, box_bold-1)
                img = cv2.putText(img, str(int(trk.id+1)), (int(box[0]), int(box[1])+20),
                        cv2.FONT_HERSHEY_COMPLEX, 1,
                        (0, 0, 255), 2)
                img = cv2.putText(img, str(int(trk.depth)), (int(box[2])-14, int(box[3])),
                        cv2.FONT_HERSHEY_COMPLEX, 0.8,
                        (0, 0, 255), 2)
                xc = int(trk.cen[0])
                yc = int(trk.cen[1])
                img = cv2.rectangle(img, (xc-1, yc-1),(xc+1, yc+1), (0, 0, 255), box_bold)

            for idx, (det, cen) in enumerate(zip(dets, cens)):
                det = det.astype('int')
                cen = cen.astype('int')
                img = cv2.rectangle(img, (det[0], det[1]), (det[2], det[3]), (255, 0, 0), 2)
                img = cv2.rectangle(img, (cen[0]-1, cen[1]-1), (cen[0]+1, cen[1]+1), (255, 0, 0), 4)

            key = 0
            while(key not in [ord('q'), ord(' '), 27]):
                cv2.imshow('preview', cv2.resize(img, (0, 0), fx=0.5, fy=0.5))
                key = cv2.waitKey(1)

            if key == 27:
                cv2.destroyAllWindows()
                exit()

        # Get output returns and remove dead tracklet
        i = len(self.trackers)
        for trk in reversed(self.trackers):
            trk_box = trk.get_state()[0]
            if (trk.time_since_update < 1) and (
                    trk.hit_streak >= self.min_hits or self.frame_count <=
                    self.min_hits):
                # +1 as MOT benchmark requires positive
                height = float(trk.det[3] - trk.det[1])
                width = float(trk.det[2] - trk.det[0])
                hypo = {'height': height,
                        'width': width,
                        'trk_box': trk_box.tolist(),
                        'det_box': trk.det.tolist(),
                        'id': trk.id + 1,
                        'x': int(trk.cen[0]),
                        'y': int(trk.cen[1]),
                        'dim': trk.dim.tolist(),
                        'alpha': trk.alpha.item(),
                        'roty': trk.rot.item(),
                        'depth': trk.depth.item()
                        }
                ret.append(hypo)
            i -= 1
            # remove dead tracklet
            if trk.time_since_update > self.max_age:
                self.trackers.pop(i)
        return ret

    # UTILS
    def process_info(self, data, det_thresh=0.3, match_thres=0.5, max_depth=150,
                     _nms=True, _valid=True):
        """return the valid, matched index of bbox, and corresponding
        tracking ids, and corresponding info"""
        bbox = data['rois_pd']
        if len(bbox) > 0:
            if _nms:
                keep = tu.nms_cpu(bbox, 0.3)
            else:
                keep = np.arange(0, bbox.shape[0])
            bbox = bbox[keep]
            res_feat = data['feature'][keep]
            res_dim = data['dim_pd'][keep]
            res_alpha = data['alpha_pd'][keep]
            res_depth = data['depth_pd'][keep]
            res_cen = data['center_pd'][keep]

        gt_boxes = data['rois_gt']
        dim_gt = data['dim_gt']
        alpha_gt = data['alpha_gt']
        depth_gt = data['depth_gt']
        cen_gt = data['center_gt']
        ignores = data['ignore']
        tracking_ids = data['tid_gt']

        # Ignore if gt is max_depth meters away
        ignores[depth_gt > max_depth] = 1

        # Prune bbox if prediction is max_depth meters away
        if len(bbox) > 0:
            bbox = bbox[res_depth < max_depth]
            res_feat = res_feat[res_depth < max_depth]
            res_dim = res_dim[res_depth < max_depth]
            res_alpha = res_alpha[res_depth < max_depth]
            res_cen = res_cen[res_depth < max_depth]
            res_depth = res_depth[res_depth < max_depth]

        else:
            # print('*** Pure detecton, do not get pose information here!')
            res_feat = np.zeros([bbox.shape[0], self.feat_dim])
            res_dim = np.ones([bbox.shape[0], 3])
            res_alpha = np.ones([bbox.shape[0]])
            res_depth = np.ones([bbox.shape[0]])
            res_cen = np.zeros([bbox.shape[0], 2])

        # (x1, y1, x2, y2, conf)
        if gt_boxes.shape[0] > 0:
            gt_boxes = gt_boxes[:, 0:4]

        self.frame_annotation = tu.build_frame_annotation(
                                    gt_boxes,
                                    ignores,
                                    tracking_ids,
                                    dim_gt,
                                    alpha_gt,
                                    depth_gt,
                                    cen_gt,
                                    self.cam_calib,
                                    self.cam_rot,
                                    self.cam_coord)

        gt_boxes_ignored = [gb for i, gb in enumerate(gt_boxes) if ignores[i]]

        if _valid:
            valid_bbox_ind = []
            for i in range(bbox.shape[0]):
                box = bbox[i, :4]
                score = bbox[i, 4]
                if score > det_thresh:  # if valid
                    save = True
                    for bg in gt_boxes_ignored:
                        if tu.compute_iou(box, bg) > match_thres:
                            save = False
                            break
                    if save:
                        valid_bbox_ind.append(i)
        else:
            valid_bbox_ind = np.arange(0, bbox.shape[0])

        valid_bbox = bbox[valid_bbox_ind]
        valid_feat = res_feat[valid_bbox_ind]
        valid_dim = res_dim[valid_bbox_ind]
        valid_alpha = res_alpha[valid_bbox_ind].reshape(-1, 1)
        valid_depth = res_depth[valid_bbox_ind].reshape(-1, 1)
        valid_cen = res_cen[valid_bbox_ind]

        valid_rots = np.zeros_like(valid_alpha)
        for idx, (alpha, det, center) in enumerate(
                        zip(valid_alpha, valid_bbox, valid_cen)):
            valid_rots[idx] = tu.deg2rad(tu.alpha2rot_y(
                                    alpha,
                                    center[0] - self.W//2, #cam_calib[0][2],
                                    FOCAL_LENGTH=self.cam_calib[0][0]))

        # now build the world 3d coordinate
        valid_worldcoords = tu.point3dcoord(valid_cen, valid_depth, self.cam_calib, self.cam_pose)

        return valid_bbox, valid_feat, valid_dim, valid_alpha, valid_depth, \
               valid_cen, valid_rots, valid_worldcoords


import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.optim import lr_scheduler
from torch import optim
import torch.nn.functional as F
from utils.masked_cross_entropy import *
from utils.config import *
import random
import numpy as np
import datetime
from utils.measures import wer, moses_multi_bleu
import matplotlib
# matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn  as sns
import nltk
import os
import logging
from sklearn.metrics import f1_score

class Mem2Seq(nn.Module):
    def __init__(self, hidden_size, max_len, max_r, lang, path, lr, n_layers, dropout):
        super(Mem2Seq, self).__init__()
        self.name = "Mem2Seq"
        self.input_size = lang.n_words
        self.output_size = lang.n_words
        self.hidden_size = hidden_size
        self.max_len = max_len ## max input
        self.max_r = max_r ## max responce len
        self.lang = lang
        self.lr = lr
        self.n_layers = n_layers
        self.dropout = dropout

        if path:
            if USE_CUDA:
                logging.info("MODEL {} LOADED".format(str(path)))
                self.encoder = torch.load(str(path)+'/enc.th')
                self.decoder = torch.load(str(path)+'/dec.th')
            else:
                logging.info("MODEL {} LOADED".format(str(path)))
                self.encoder = torch.load(str(path)+'/enc.th',lambda storage, loc: storage)
                self.decoder = torch.load(str(path)+'/dec.th',lambda storage, loc: storage)
        else:
            self.encoder = EncoderMemNN(lang.n_words, hidden_size, n_layers, self.dropout)
            self.decoder = DecoderrMemNN(lang.n_words, hidden_size, n_layers, self.dropout)
        # Initialize optimizers and criterion
        self.encoder_optimizer = optim.Adam(self.encoder.parameters(), lr=lr)
        self.decoder_optimizer = optim.Adam(self.decoder.parameters(), lr=lr)
        self.scheduler = lr_scheduler.ReduceLROnPlateau(self.decoder_optimizer,mode='max',factor=0.5,patience=1,min_lr=0.0001, verbose=True)
        self.criterion = nn.MSELoss()
        self.loss = 0
        self.loss_gate = 0
        self.loss_ptr = 0
        self.loss_vac = 0
        self.print_every = 1
        self.batch_size = 0
        # Move models to GPU
        if USE_CUDA:
            self.encoder.cuda()
            self.decoder.cuda()

    def print_loss(self):
        print_loss_avg = self.loss / self.print_every
        print_loss_ptr =  self.loss_ptr / self.print_every
        print_loss_vac =  self.loss_vac / self.print_every
        self.print_every += 1
        return 'L:{:.2f}, VL:{:.2f}, PL:{:.2f}'.format(print_loss_avg,print_loss_vac,print_loss_ptr)

    def save_model(self, dec_type):
        directory = 'save/mem2seq_'+'HDD'+str(self.hidden_size)+'BSZ'+str(self.batch_size)+'DR'+str(self.dropout)+'L'+str(self.n_layers)+'lr'+str(self.lr)+str(dec_type)
        if not os.path.exists(directory):
            os.makedirs(directory)
        torch.save(self.encoder, directory+'/enc.th')
        torch.save(self.decoder, directory+'/dec.th')

    def train_batch(self, input_batches, input_lengths, target_batches,
                    target_lengths, target_index, batch_size, clip,
                    teacher_forcing_ratio,reset):
        if reset:
            self.loss = 0
            self.loss_gate = 0
            self.loss_ptr = 0
            self.loss_vac = 0
            self.print_every = 1

        self.batch_size = batch_size
        # Zero gradients of both optimizers
        self.encoder_optimizer.zero_grad()
        self.decoder_optimizer.zero_grad()
        loss_Vocab,loss_Ptr= 0,0

        # Run words through encoder
        decoder_hidden = self.encoder(input_batches.transpose(0,1)).unsqueeze(0)

        # load memories with input
        self.decoder.load_memory(input_batches.transpose(0,1))

        # Prepare input and output variables
        decoder_input = Variable(torch.LongTensor([SOS_token] * batch_size))

        max_target_length = max(target_lengths)
        all_decoder_outputs_vocab = Variable(torch.zeros(max_target_length, batch_size, self.output_size))
        all_decoder_outputs_ptr = Variable(torch.zeros(max_target_length, batch_size, input_batches.size(0)))

        # Move new Variables to CUDA
        if USE_CUDA:
            all_decoder_outputs_vocab = all_decoder_outputs_vocab.cuda()
            all_decoder_outputs_ptr = all_decoder_outputs_ptr.cuda()
            decoder_input = decoder_input.cuda()

        # Choose whether to use teacher forcing
        use_teacher_forcing = random.random() < teacher_forcing_ratio

        if use_teacher_forcing:
            # Run through decoder one time step at a time
            for t in range(max_target_length):
                decoder_ptr, decoder_vacab, decoder_hidden  = self.decoder.ptrMemDecoder(decoder_input, decoder_hidden)
                all_decoder_outputs_vocab[t] = decoder_vacab
                all_decoder_outputs_ptr[t] = decoder_ptr
                decoder_input = target_batches[t]# Chosen word is next input
                if USE_CUDA: decoder_input = decoder_input.cuda()
        else:
            for t in range(max_target_length):
                decoder_ptr, decoder_vacab, decoder_hidden = self.decoder.ptrMemDecoder(decoder_input, decoder_hidden)
                _, toppi = decoder_ptr.data.topk(1)
                _, topvi = decoder_vacab.data.topk(1)
                all_decoder_outputs_vocab[t] = decoder_vacab
                all_decoder_outputs_ptr[t] = decoder_ptr
                ## get the correspective word in input
                top_ptr_i = torch.gather(input_batches,0,Variable(toppi.view(1, -1)))
                next_in = [top_ptr_i.squeeze()[i].item() if(toppi.squeeze()[i] < input_lengths[i]-1) else topvi.squeeze()[i] for i in range(batch_size)]
                decoder_input = Variable(torch.LongTensor(next_in)) # Chosen word is next input
                if USE_CUDA: decoder_input = decoder_input.cuda()

        #Loss calculation and backpropagation
        loss_Vocab = masked_cross_entropy(
            all_decoder_outputs_vocab.transpose(0, 1).contiguous(), # -> batch x seq
            target_batches.transpose(0, 1).contiguous(), # -> batch x seq
            target_lengths
        )
        loss_Ptr = masked_cross_entropy(
            all_decoder_outputs_ptr.transpose(0, 1).contiguous(), # -> batch x seq
            target_index.transpose(0, 1).contiguous(), # -> batch x seq
            target_lengths
        )

        loss = loss_Vocab + loss_Ptr
        loss.backward()

        # Clip gradient norms
        ec = torch.nn.utils.clip_grad_norm_(self.encoder.parameters(), clip)
        dc = torch.nn.utils.clip_grad_norm_(self.decoder.parameters(), clip)
        # Update parameters with optimizers
        self.encoder_optimizer.step()
        self.decoder_optimizer.step()
        self.loss += loss.item()
        #self.loss_gate += loss_gate.item()
        self.loss_ptr += loss_Ptr.item()
        self.loss_vac += loss_Vocab.item()

    def evaluate_batch(self,batch_size,input_batches, input_lengths, target_batches, target_lengths, target_index,src_plain):
        # Set to not-training mode to disable dropout
        self.encoder.train(False)
        self.decoder.train(False)
        # Run words through encoder
        decoder_hidden = self.encoder(input_batches.transpose(0,1)).unsqueeze(0)
        self.decoder.load_memory(input_batches.transpose(0,1))

        # Prepare input and output variables
        decoder_input = Variable(torch.LongTensor([SOS_token] * batch_size))

        decoded_words = []
        all_decoder_outputs_vocab = Variable(torch.zeros(self.max_r, batch_size, self.output_size))
        all_decoder_outputs_ptr = Variable(torch.zeros(self.max_r, batch_size, input_batches.size(0)))
        # Move new Variables to CUDA
        if USE_CUDA:
            all_decoder_outputs_vocab = all_decoder_outputs_vocab.cuda()
            all_decoder_outputs_ptr = all_decoder_outputs_ptr.cuda()
            decoder_input = decoder_input.cuda()

        # Run through decoder one time step at a time
        for t in range(self.max_r):
            decoder_ptr,decoder_vacab, decoder_hidden = self.decoder.ptrMemDecoder(decoder_input, decoder_hidden)
            all_decoder_outputs_vocab[t] = decoder_vacab
            _, topvi = decoder_vacab.data.topk(1)
            all_decoder_outputs_ptr[t] = decoder_ptr
            _, toppi = decoder_ptr.data.topk(1)
            top_ptr_i = torch.gather(input_batches,0,Variable(toppi.view(1, -1)))
            next_in = [top_ptr_i.squeeze()[i].item() if(toppi.squeeze()[i] < input_lengths[i]-1) else topvi.squeeze()[i] for i in range(batch_size)]
            decoder_input = Variable(torch.LongTensor(next_in)) # Chosen word is next input
            if USE_CUDA: decoder_input = decoder_input.cuda()

            temp = []
            for i in range(batch_size):
                if(toppi.squeeze()[i] < len(src_plain[i])-1 ):
                    temp.append(src_plain[i][toppi.squeeze()[i]]) ## copy from the input
                else:
                    ind = topvi.squeeze()[i]
                    if ind.item() == EOS_token:
                        temp.append('<EOS>')
                    else:
                        temp.append(self.lang.index2word[ind.item()]) ## get from vocabulary
            decoded_words.append(temp)

        # Set back to training mode
        self.encoder.train(True)
        self.decoder.train(True)
        return decoded_words

    def evaluate(self,dev,avg_best,BLEU=False):
        logging.info("STARTING EVALUATION")
        acc_avg = 0.0
        wer_avg = 0.0
        ref = []
        hyp = []
        pbar = tqdm(enumerate(dev),total=len(dev))
        for j, data_dev in pbar:
            words = self.evaluate_batch(
                                        batch_size=len(data_dev[1]),
                                        input_batches=data_dev[0],
                                        input_lengths=data_dev[1],
                                        target_batches=data_dev[2],
                                        target_lengths=data_dev[3],
                                        target_index=data_dev[4],
                                        src_plain=data_dev[5])
            acc=0
            w = 0
            temp_gen = []
            for i, row in enumerate(np.transpose(words)):
                st = ''
                for e in row:
                    if e== '<EOS>': break
                    else: st+= e + ' '
                temp_gen.append(st)
                correct = " ".join(data_dev[6][i])
                ### IMPORTANT
                ### WE NEED TO COMPARE THE PLAIN STRING, BECAUSE WE COPY THE WORDS FROM THE INPUT
                ### ====>> the index in the output gold can be UNK
                if (correct.lstrip().rstrip() == st.lstrip().rstrip()):
                    acc+=1
                w += wer(correct.lstrip().rstrip(),st.lstrip().rstrip())
                ref.append(str(correct.lstrip().rstrip()))
                hyp.append(str(st.lstrip().rstrip()))

            acc_avg += acc/float(len(data_dev[1]))
            wer_avg += w/float(len(data_dev[1]))
            pbar.set_description("R:{:.4f},W:{:.4f}".format(acc_avg/float(len(dev)),wer_avg/float(len(dev))))

        bleu_score = moses_multi_bleu(np.array(hyp), np.array(ref), lowercase=True)
        logging.info("BLEU SCORE:"+str(bleu_score))

        if (bleu_score >= avg_best):
            self.save_model(str(self.name)+str(bleu_score))
            logging.info("MODEL SAVED")
        return bleu_score

class EncoderMemNN(nn.Module):
    def __init__(self, vocab, embedding_dim, hop, dropout):
        super(EncoderMemNN, self).__init__()
        self.num_vocab = vocab
        self.max_hops = hop
        self.embedding_dim = embedding_dim
        self.dropout = dropout
        for hop in range(self.max_hops+1):
            C = nn.Embedding(self.num_vocab, embedding_dim, padding_idx=PAD_token)
            C.weight.data.normal_(0, 0.1)
            self.add_module("C_{}".format(hop), C)
        self.C = AttrProxy(self, "C_")
        self.softmax = nn.Softmax(dim=1)

    def get_state(self,bsz):
        """Get cell states and hidden states."""
        if USE_CUDA:
            return Variable(torch.zeros(bsz, self.embedding_dim)).cuda()
        else:
            return Variable(torch.zeros(bsz, self.embedding_dim))

    def forward(self, story):
        u = [self.get_state(story.size(0))]
        for hop in range(self.max_hops):
            embed_A = self.C[hop](story.contiguous().view(story.size(0), -1).long()) # b * (m * s) * e
            u_temp = u[-1].unsqueeze(1).expand_as(embed_A)
            prob   = self.softmax(torch.sum(embed_A*u_temp, 2))
            embed_C = self.C[hop+1](story.contiguous().view(story.size(0), -1).long())
            prob = prob.unsqueeze(2).expand_as(embed_C)
            o_k  = torch.sum(embed_C*prob, 1)
            u_k = u[-1] + o_k
            u.append(u_k)
        return u_k

class DecoderrMemNN(nn.Module):
    def __init__(self, vocab, embedding_dim, hop, dropout):
        super(DecoderrMemNN, self).__init__()
        self.num_vocab = vocab
        self.max_hops = hop
        self.embedding_dim = embedding_dim
        self.dropout = dropout
        for hop in range(self.max_hops+1):
            C = nn.Embedding(self.num_vocab, embedding_dim, padding_idx=PAD_token)
            C.weight.data.normal_(0, 0.1)
            self.add_module("C_{}".format(hop), C)
        self.C = AttrProxy(self, "C_")
        self.softmax = nn.Softmax(dim=1)
        self.W = nn.Linear(embedding_dim,1)
        self.W1 = nn.Linear(2*embedding_dim,self.num_vocab)
        self.gru = nn.GRU(embedding_dim, embedding_dim, dropout=dropout)

    def load_memory(self, story):
        self.m_story = []
        for hop in range(self.max_hops):
            embed_A = self.C[hop](story.contiguous().view(story.size(0), -1))#.long()) # b * (m * s) * e
            m_A = embed_A
            embed_C = self.C[hop+1](story.contiguous().view(story.size(0), -1).long())
            m_C = embed_C
            self.m_story.append(m_A)
        self.m_story.append(m_C)

    def ptrMemDecoder(self, enc_query, last_hidden):
        embed_q = self.C[0](enc_query) # b * e
        _, hidden = self.gru(embed_q.unsqueeze(0), last_hidden)
        temp = []
        u = [hidden[0].squeeze()]
        for hop in range(self.max_hops):
            m_A = self.m_story[hop]
            u_temp = u[-1].unsqueeze(1).expand_as(m_A)
            prob_lg = torch.sum(m_A*u_temp, 2)
            prob_   = self.softmax(prob_lg)
            m_C = self.m_story[hop+1]
            temp.append(prob_)
            prob = prob_.unsqueeze(2).expand_as(m_C)
            o_k  = torch.sum(m_C*prob, 1)
            if (hop==0):
                p_vocab = self.W1(torch.cat((u[0], o_k),1))
            u_k = u[-1] + o_k
            u.append(u_k)
        p_ptr = prob_lg
        return p_ptr, p_vocab, hidden

class AttrProxy(object):
    """
    Translates index lookups into attribute lookups.
    To implement some trick which able to use list of nn.Module in a nn.Module
    see https://discuss.pytorch.org/t/list-of-nn-module-in-a-nn-module/219/2
    """
    def __init__(self, module, prefix):
        self.module = module
        self.prefix = prefix

    def __getitem__(self, i):
        return getattr(self.module, self.prefix + str(i))


#external modules
import cv2
import numpy as np
import os
from pathlib import Path

#local modules
from DuckieBotModes.driver import Driver

class Squid(Driver):
    '''
    driver mode but with ink on the screen

    Attributes:
        ink(cv2.Image): the ink image with transparent background
    '''
    def __init__(self):
        Driver.__init__(self)
        image_path = Path(os.path.realpath(__file__))/".."/"ink.png"
        self.ink = cv2.imread(str(image_path.resolve()), cv2.IMREAD_UNCHANGED)
        self.ink.reshape(self.ink.shape[0]*self.ink.shape[1], 4)

    def frame(self, frame):
        '''
        implementation of method in parent (overriding)
        '''
        #resizing and getting everything into an easy list
        ink_pixels = cv2.resize(self.ink, (frame.shape[1], frame.shape[0]))
        ink_pixels = ink_pixels.reshape(frame.shape[0]*frame.shape[1], 4)
        frame_pixels = frame.reshape(frame.shape[0]*frame.shape[1], 3)

        # handles the blending of the two images
        is_zero = (ink_pixels[:,3] == 0) # 0 is transparent
        is_zero = np.array([is_zero, is_zero, is_zero]).transpose() #makes three dimensional array
        frame_pixels = frame_pixels*(is_zero) + ink_pixels[:, 0:3]*(is_zero != True) #three rows, takes pixels

        frame = frame_pixels.reshape(frame.shape[0], frame.shape[1], 3)

        return frame