#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
This is a sub-module of smili2 handling image data.
'''
__author__ = "Smili Developer Team"
from ...util.xarrds import XarrayDataset
from ... import util
import copy
import numpy as np
import itertools
import scipy.ndimage as sn
import h5py
import tqdm

# Logger
from logging import getLogger
logger = getLogger(__name__)

class Image(XarrayDataset):
    '''
    The Class to handle five dimensional images

    Attributes:
        ds (xarray.Dataset): Image data
    '''
    # Data type name
    datatype = "ehtdata_image"

    # Default Group Name for the I/O
    group = "image"

    # Data Set
    #   This supposed to include the data set
    ds = None

    # data set dimension
    ds_dims = ("time", "freq", "pol", "y", "x")

    def init_attrs(self):
        '''
        Initialize the metadata.
        '''
        self.ds.attrs["source"] = "No Name"
        self.ds.attrs["instrument"] = "No Name"
        self.ds.attrs["ra"] = 0.
        self.ds.attrs["dec"] = 0.
        self.ds.attrs["coordsys"] = "ICRS"
        self.ds.attrs["equinox"] = -1.
        self.ds.attrs["dx"] = 1.
        self.ds.attrs["dy"] = 1.
        self.ds.attrs["ixref"] = 1.
        self.ds.attrs["iyref"] = 1.
        self.ds.attrs["beam_maj"] = 0.
        self.ds.attrs["beam_min"] = 0.
        self.ds.attrs["beam_pa"] = 0.
        self.ds.attrs["pulse_type"] = "rect"

    def init_pol(self):
        '''
        Add the polarization information to the image data array.
        '''
        ns = self.ds.dims["pol"]

        if ns == 1:
            self.ds.coords["pol"] = ("pol", ["I"])
        elif ns == 4:
            self.ds.coords["pol"] = ("pol", ["I", "Q", "U", "V"])
        else:
            raise ValueError("Current version of this library accepts only"
                             "single or full polarization images.")

    def init_xygrid(self):
        '''
        Add the xy coordinates (in radians) to the image data array.
        '''
        from numpy import arange

        axis_list = ["x", "y"]
        sign_list = [-1, 1]

        for iaxis in range(2):
            axis = axis_list[iaxis]
            sign = sign_list[iaxis]

            # get some meta info
            nx = self.ds.dims[axis]
            dx = self.ds.attrs["d"+axis]
            ixref = self.ds.attrs["i{}ref".format(axis)]

            # compute coordinates
            self.ds.coords[axis] = (axis, sign*dx*(1+arange(nx)-ixref))

    def chunk(self, chunks=None, **args):
        if chunks is None:
            chunks = dict(
                time=1,
                freq=1,
                pol=1,
                y=self.ds.dims["y"],
                x=self.ds.dims["x"]
            )
        self.ds = self.ds.chunk(chunks=chunks, **args)

    def set_source(self, source="M87", srccoord=None):
        '''
        Set the source name and the source coordinate to the metadata.
        If source coordinate is not given, it will be taken from the CDS.

        Args:
            source (str; default="M87"):
                Source Name
            srccoord (astropy.coordinates.Skycoord object; default=None):
                Source position. If not specified, it is automatically pulled
                from the CDS.
        '''
        from astropy.coordinates import SkyCoord

        if source is None and srccoord is None:
            raise ValueError("source or srccoord must be specified.")

        if source is not None:
            self.ds.attrs["source"] = source
            if srccoord is None:
                srccoord = SkyCoord.from_name(source)

        if srccoord is not None:
            # parse coordinates
            self.ds.attrs["ra"] = srccoord.ra.rad
            self.ds.attrs["dec"] = srccoord.dec.rad
            # coordinate system
            self.ds.attrs["coordsys"] = srccoord.name
            # equinox
            if "icrs" in srccoord.name:
                self.ds.attrs["equinox"] = -1
            elif "fk" in srccoord.name:
                self.ds.attrs["equinox"] = srccoord.equinox.jyear
            elif "gcrs" in srccoord.name:
                self.ds.attrs["equinox"] = srccoord.obstime.jyear
            else:
                raise ValueError(
                    "Coordinate System %s is not supported currently." % (srccoord.name))

    def set_instrument(self, instrument="EHT"):
        """
        Set the metadata for the instrument with a
        specified name of the instrument.

        Args:
            instrument (str): The instrument name. Defaults to "EHT".
        """
        if instrument is not None:
            self.ds.attrs["instrument"] = instrument

    def set_mjd(self, mjd):
        """
        Set the MJD infromation for the image data array.

        Args:
            mjd (float or array): MJD
        """
        from numpy import isscalar, asarray

        if isscalar(mjd):
            mjd_arr = asarray([mjd], dtype="float64")
        else:
            mjd_arr = mjd.copy()

        self.ds.coords["mjd"] = ("time", mjd_arr)

    def set_freq(self, freq):
        """
        Set the frequency infromation for the image data array.

        Args:
            freq (float or array): Frequency in Hz
        """
        from numpy import isscalar, asarray

        if isscalar(freq):
            freq_arr = asarray([freq], dtype="float64")
        else:
            freq_arr = freq.copy()

        self.ds.coords["freq"] = ("freq", freq_arr)

    def set_angunit(self, angunit="auto"):
        from numpy import amax
        from ...util.units import Unit, DEG

        if angunit != "auto":
            self.ds.attrs["angunit"] = angunit
            return

        angunits = ["deg", "arcmin", "arcsec", "mas", "uas"]
        xmax = amax(self.get_imextent(angunit="deg"))*DEG

        for angunit in angunits:
            self.ds.attrs["angunit"] = angunit
            if xmax < 0.1*Unit(angunit):
                continue
            else:
                break

    def set_beam(self, majsize=0., minsize=0., pa=0., scale=1., angunit=None):
        '''
        Set beam parameters into headers.

        Args:
            majsize, minsize(float, default=0):
                major/minor-axis FWHM size
            scale(float, default=1):
                scaling factor that will be multiplied to maj/min size.
            pa(float, default=0):
                position angle in deg
        '''
        from ...util.units import conv
        from numpy import deg2rad

        # get the conversion factor for the angular unit
        if angunit is None:
            angunit = self.angunit
        factor = conv(angunit, "rad")

        # beam paramters
        self.ds.attrs["beam_maj"] = majsize * factor * scale
        self.ds.attrs["beam_min"] = minsize * factor * scale
        self.ds.attrs["beam_pa"] = deg2rad(pa)

    def get_bconv(self, fluxunit="Jy", saunit="pixel"):
        from ...util.units import conv, Unit
        from numpy import log, pi

        # pixel size (in radian)
        dx = self.ds.attrs["dx"]
        dy = self.ds.attrs["dy"]
        dxdy = dx*dy  # solid angle of the pixel

        if "k" in fluxunit.lower():
            from astropy.constants import c, k_B
            nu = self.ds.coords["freq"]
            Jy2K = c.si.value ** 2 / \
                (2 * k_B.si.value * nu ** 2) / dxdy * 1e-26
            return Jy2K * conv("K", fluxunit)
        else:
            fluxconv = conv("Jy", fluxunit)

            if "pix" in saunit.lower() or "px" in saunit.lower():
                saconv = 1
            elif "beam" in saunit.lower():
                # get beamsize (in radian)
                bmaj = self.ds.attrs["beam_maj"]
                bmin = self.ds.attrs["beam_min"]

                # beamsolid angle
                beamsa = bmaj*bmin*pi/(4*log(2))
                saconv = conv(dxdy*Unit("rad**2"), beamsa*Unit("rad**2"))
            else:
                saconv = conv(dxdy*Unit("rad**2"), saunit)

            return fluxconv/saconv

    def get_xygrid(self, twodim=False, angunit=None):
        '''
        Get the xy coordinates of the image

        Args:
            angunit(string): Angular unit(uas, mas, asec or arcsec,
                             amin or arcmin, degree)
            twodim(boolean): It True, the 2D grids will be returned.
                             Otherwise, the 1D arrays will be returned
        Returns:
            x, y: x, y coordinates in the specified unit.
        '''
        from ...util.units import conv
        from numpy import meshgrid

        # get the conversion factor for the angular unit
        if angunit is None:
            angunit = self.ds.attrs["angunit"]
        factor = conv("rad", angunit)

        x = self.ds.coords["x"].data * factor
        y = self.ds.coords["y"].data * factor
        if twodim:
            x, y = meshgrid(x, y)
        return x, y

    def get_imextent(self, angunit=None):
        '''
        Get the field of view of the image for the pyplot.imshow function.

        Args:
          angunit(string): Angular unit

        Returns:
          [xmax, xmin, ymin, ymax]: extent of the image
        '''
        from ...util.units import conv
        from numpy import asarray

        # get the conversion factor for the angular unit
        if angunit is None:
            angunit = self.ds.attrs["angunit"]
        factor = conv("rad", angunit)

        dx = self.ds.attrs["dx"]
        dy = self.ds.attrs["dy"]
        nx = self.ds.dims["x"]
        ny = self.ds.dims["y"]
        ixref = self.ds.attrs["ixref"]
        iyref = self.ds.attrs["iyref"]

        xmax = -dx * (1 - ixref - 0.5)
        xmin = -dx * (nx - ixref + 0.5)
        ymax = dy * (ny - iyref + 0.5)
        ymin = dy * (1 - iyref - 0.5)

        return asarray([xmax, xmin, ymin, ymax]) * factor

    def get_imarray(self, fluxunit="Jy", saunit="pixel"):
        bconv = self.get_bconv(fluxunit=fluxunit, saunit=saunit)
        converted = self.ds.image * bconv
        return converted.data.copy()

    def get_uvgrid(self, twodim=False):
        """
        Get the uv coordinates of the image on the Fourier domain

        Args:
            twodim(boolean): It True, the 2D grids will be returned.
                             Otherwise, the 1D arrays will be returned
        Returns:
            u, v: u, v coordinates.
        """
        from numpy import meshgrid
        from numpy.fft import fftfreq, ifftshift

        # get the shape of array
        nx = self.ds.dims["x"]
        ny = self.ds.dims["y"]

        # pixel size
        dxrad = self.ds.attrs["dx"]
        dyrad = self.ds.attrs["dy"]

        # create uv grids
        ug = ifftshift(fftfreq(nx, d=-dxrad))
        vg = ifftshift(fftfreq(ny, d=dyrad))

        if twodim:
            return meshgrid(ug, vg)
        else:
            return ug, vg

    def get_uvextent(self):
        """
        Get the field of view of the image on the Fourier transform
        for the pyplot.imshow function. Here we assume that the Fourier image
        is created by the get_visarr method.

        Returns:
          [umax, umin, vmin, vmax]: extent of the Fourier transformed image.
        """
        from numpy import asarray
        ug, vg = self.get_uvgrid()
        du_half = (ug[1] - ug[0])/2
        dv_half = (vg[1] - vg[0])/2
        return asarray([ug[0]-du_half, ug[-1]+du_half, vg[0]-dv_half,
                        vg[-1]+dv_half])

    def get_vis(self, idx=(0, 0, 0), apply_pulsefunc=True):
        """
        Get an array of visibilities computed from the image.

        Args:
            idx (tuple, optional):
                An index, or a list of indice of the image.
                The index should be in the form of (time, freq, pol).
                Defaults to (0, 0, 0). If you specify None, then
                visibilities will be computed for every index of images.
            apply_pulsefunc (bool, optional):
                If True, the pulse function specified in the meta data
                will be applied. Defaults to True.

        Returns:
            numpy.ndarray:
                full complex visibilities. The array shape will
                depend on idx.
        """
        from numpy import pi, exp, asarray, unravel_index
        from numpy.fft import fftshift, ifftshift, fft2

        # get array
        imarr = self.get_imarray()
        nt, nf, ns, ny, nx = imarr.shape
        ixref = self.ds.attrs["ixref"]
        iyref = self.ds.attrs["iyref"]
        dxrad = self.ds.attrs["dx"]
        dyrad = self.ds.attrs["dy"]

        # get uv grid
        ug, vg = self.get_uvgrid(twodim=True)

        # adjust phase
        ix_cen = nx//2 + nx % 2  # image center index used in np.fft.fft2
        iy_cen = ny//2 + ny % 2
        dix = ix_cen - ixref  # shift in the pixel unit
        diy = iy_cen - iyref
        viskernel = exp(1j*2*pi*(-dxrad*dix*ug + dyrad*diy*vg))

        # mutiply the pulse function
        # if apply_pulsefunc:
        #    viskernel *= self.get_pulsefunc()(ug, vg)

        # define FFT function
        def dofft(imarr2d):
            return ifftshift(fft2(fftshift(imarr2d)))

        # compute full complex visibilities
        if idx is None:
            shape3d = (nt, nf, ns)
            vis = asarray([dofft(imarr[unravel_index(i, shape=shape3d)])
                           for i in range(nt*nf*ns)]).reshape([nt, nf, ns,
                                                               ny, nx])
            vis[:, :, :] *= viskernel
        else:
            ndim = asarray(idx).ndim
            if ndim == 1:
                vis = dofft(imarr[idx])*viskernel
            elif ndim == 2:
                nidx = len(idx)
                vis = asarray([dofft(imarr[tuple(idx[i])])
                               for i in range(nidx)]).reshape([nidx, ny, nx])
                vis[:] *= viskernel
            else:
                raise ValueError("Invalid dimension of the input index.")
        return vis

    def get_skycoord(self):
        from astropy.coordinates import SkyCoord
        from astropy.time import Time
        from ...util.units import RAD

        args = dict(
            ra=self.ds.attrs["x"] * RAD,
            dec=self.ds.attrs["y"] * RAD,
            frame=self.ds.attrs["coordsys"]
        )
        if "fk" in args["frame"]:
            args["equinox"] = Time(self.ds.attrs["equinox"], format="jyear")
        elif args["frame"] != "gcrs":
            args["obstime"] = Time(self.ds.attrs["equinox"], format="jyear")

        return SkyCoord(**args)

    def get_beam(self, angunit=None):
        '''
        Get beam parameters.

        Args:
            angunit(string): Angular Unit
        Return:
            dic: the beam parameter information
        '''
        from ...util.units import conv
        from numpy import rad2deg

        if angunit is None:
            angunit = self.ds.attrs["angunit"]

        factor = conv("rad", angunit)

        outdic = {}
        outdic["scale"] = 1.0
        outdic["majsize"] = self.ds.attrs["beam_maj"] * factor
        outdic["minsize"] = self.ds.attrs["beam_min"] * factor
        outdic["pa"] = rad2deg(self.ds.attrs["beam_pa"])
        outdic["angunit"] = angunit
        return outdic

    def get_pulsefunc(self):
        ptype = self.ds.attrs["ptype"].lower()

        if "delta" in ptype:
            def pulsefunc(u, v): return 1
        elif "rect" in ptype:
            from ...geomodel import Rectangular
            dxrad = self.ds.attrs["dx"]
            dyrad = self.ds.attrs["dy"]
            pulsefunc = Rectangular(
                Lx=dxrad, Ly=dyrad, dx=dxrad, dy=dyrad, angunit="rad").V
        else:
            raise ValueError("unknown pulse type: %s" % (ptype))
        return pulsefunc

    def add_geomodel(self, geomodel, idx=(0, 0, 0), inplace=False):
        """
        Add a specified geometric model to the image

        Args:
            geomodel (geomodel.GeoModel):
                The input geometric model.
            idx (tuble):
                An index of the image where the Gaussians to be added.
                Should be in the format of (time, freq, pol).
                Default to (0,0,0).
            inplace (bool, optional):
                If False, return the convolved image. Defaults to False.

        Returns:
            imdata.Image: if inplace==False.
        """
        if inplace:
            outimage = self
        else:
            outimage = self.copy()

        # get x,y coordinates
        xg, yg = self.get_xygrid(angunit="rad", twodim=True)

        # compute the intensity
        imarr = geomodel.Img(xg, yg)

        # add data
        outimage.data[idx] += imarr

        if not inplace:
            return outimage

    def add_gauss(self, totalflux=1., x0=0., y0=0., majsize=1., minsize=None,
                  pa=0., scale=1., angunit="uas", inplace=False):
        """
        Add a Gaussian to the image.

        Args:
            totalflux (float, optional):
                total flux density in Jy.
            x0, y0 (float, optional):
                The centoroid position of the kernel.
            majsize, minsize (float, optional):
                The major- and minor- axis FWHM size of the kernel.
            pa (int, optional):
                The position angle of the kernel.
            scale (int, optional):
                The scaling factor to be applied to the kernel size.
            angunit (str, optional):
                The angular unit for the centroid location and kernel size.
                Defaults to "mas".
            inplace (bool, optional):
                If False, return the convolved image. Defaults to False.

        Returns:
            imdata.Image: the convolved image if inplace==False.
        """
        from ...geomodel import Gaussian
        from ..util.units import conv

        # scale the axis size
        majsize_scaled = majsize * scale
        if minsize is None:
            minsize_scaled = majsize_scaled
        else:
            minsize_scaled = minsize * scale

        factor = conv("rad", angunit)
        dx = self.ds.attrs["dx"] * factor
        dy = self.ds.attrs["dy"] * factor

        # initialize Gaussian model
        geomodel = Gaussian(x0=x0, y0=y0, dx=dx, dy=dy, majsize=majsize_scaled,
                            minsize=minsize_scaled, pa=pa, angunit=angunit)

        # run convolution
        if inplace:
            self.add_geomodel(geomodel=geomodel, inplace=inplace)
        else:
            return self.add_geomodel(geomodel=geomodel, inplace=inplace)

    #  Convolution with Geometric Models

    def convolve_geomodel(self, geomodel, inplace=False):
        """
        Convolve the image with an input geometrical model

        Args:
            geomodel (geomodel.GeoModel):
                The input geometric model.
            inplace (bool, optional):
                If False, return the convolved image. Defaults to False.

        Returns:
            imdata.Image: the convolved image if inplace==False.
        """
        from numpy import real, unravel_index, conj
        from numpy.fft import fftshift, ifftshift, fft2, ifft2
        #import jax
        from numpy import asarray
        # import dask

        if inplace:
            outimage = self
        else:
            outimage = self.copy()

        # get the array shape
        nt, nf, ns, ny, nx = self.ds.image.shape
        self.chunk()
        imarr = self.ds.image.data
        shape3d = (nt, nf, ns)

        # get uv coordinates and compute kernel
        # conj is applied because radio astronomy uses
        # "+" for the fourier exponent
        ug, vg = self.get_uvgrid(twodim=True)
        convkernel = conj(geomodel.Vis(ug, vg))

        # define convolve functions
        def dofft(imarr2d):
            return ifftshift(fft2(fftshift(imarr2d)))

        def doifft(vis2d):
            return real(ifftshift(ifft2(fftshift(vis2d))))

        def convolve2d(imarr2d):
            #imarr2d = jax.lax.stop_gradient(imarr2d)
            return doifft(dofft(imarr2d)*convkernel)

        #convolve2d = jax.jit(convolve2d)

        # run fft convolve
        # output = [dask.delayed(convolve2d)(imarr[unravel_index(i, shape=shape3d)])
        #          for i in range(nt*nf*ns)]
        # output = dask.array.array(dask.compute(output))
        output = asarray(
            [convolve2d(imarr[unravel_index(i, shape=shape3d)])
             for i in range(nt*nf*ns)]
        )
        outimage.ds.image.data = output.reshape((nt, nf, ns, ny, nx))

        # return the output image
        if inplace is False:
            return outimage

    def convolve_gauss(self, x0=0., y0=0., majsize=1., minsize=None, pa=0.,
                       scale=1., angunit="uas", inplace=False):
        """
        Gaussian Convolution.

        Args:
            x0, y0 (float, optional):
                The centoroid position of the kernel.
            majsize, minsize (float, optional):
                The major- and minor- axis FWHM size of the kernel.
            pa (int, optional):
                The position angle of the kernel.
            scale (int, optional):
                The scaling factor to be applied to the kernel size.
            angunit (str, optional):
                The angular unit for the centroid location and kernel size.
                Defaults to "mas".
            inplace (bool, optional):
                If False, return the convolved image. Defaults to False.

        Returns:
            imdata.Image: the convolved image if inplace==False.
        """
        from ...geomodel import Gaussian

        # scale the axis size
        majsize_scaled = majsize * scale
        if minsize is None:
            minsize_scaled = majsize_scaled
        else:
            minsize_scaled = minsize * scale

        # initialize Gaussian model
        geomodel = Gaussian(x0=x0, y0=y0, majsize=majsize_scaled,
                            minsize=minsize_scaled, pa=pa, angunit=angunit)
        # run convolution
        if inplace:
            self.convolve_geomodel(geomodel=geomodel, inplace=inplace)
        else:
            return self.convolve_geomodel(geomodel=geomodel, inplace=inplace)

    def convolve_rectangular(self, x0=0., y0=0., Lx=1., Ly=None, angunit="mas",
                             inplace=False):
        """
        Convolution with a rectangular kernel

        Args:
            x0, y0 (float, optional):
                The centoroid position of the kernel.
            Lx, Ly (float, optional):
                The width of the kernel.
            angunit (str, optional):
                The angular unit for the centroid location and kernel size.
                Defaults to "mas".
            inplace (bool, optional):
                If False, return the convolved image. Defaults to False.

        Returns:
            imdata.Image: the convolved image if inplace==False.
        """
        from ...geomodel import Rectangular
        from ..util.units import conv

        # get the pixel size of the image
        rad2aunit = conv("rad", angunit)
        dx = self.ds.attrs["dx"] * rad2aunit
        dy = self.ds.attrs["dy"] * rad2aunit

        # initialize Gaussian model
        geomodel = Rectangular(x0=x0, y0=y0, Lx=Lx, Ly=Ly, dx=dx, dy=dy,
                               angunit=angunit)

        # run convolution
        if inplace:
            self.convolve_geomodel(geomodel=geomodel, inplace=inplace)
        else:
            return self.convolve_geomodel(geomodel=geomodel, inplace=inplace)

    def imshow(self,
               scale="linear",
               dyrange=100,
               gamma=0.5,
               vmax=None,
               vmin=None,
               relative=False,
               angunit=None,
               fluxunit="Jy",
               saunit="pixel",
               axisoff=False,
               axislabel=True,
               cmap="afmhot",
               idx=(0, 0, 0),
               interpolation="bilinear",
               **imshow_args):
        '''
        Plot the image.
        To change the angular unit, please change IMFITS.angunit.

        Args:
            scale(str; default="linear"):
                Transfar function. Availables are "linear", "log", "gamma"
            dyrange(float; default=100):
                Dynamic range of the log color contour.
            gamma(float; default=1/2.):
                Gamma parameter for scale = "gamma".
            vmax(float):
                The maximum value of the color contour.
            vmin(float):
                The minimum value of the color contour.
                If logscale = True, dyrange will be used to set vmin.
            relative(boolean, default=True):
                If True, vmin will be the relative value to the peak or vmax.
            fluxunit(string):
                Unit for the flux desity(Jy, mJy, uJy, K, si, cgs)
            saunit(string):
                Angular Unit for the solid angle(pixel, uas, mas, asec or arcsec,
                amin or arcmin, degree, beam).
            axisoff(boolean, default=False):
                If True, plotting without any axis label, ticks, and lines.
                This option is superior to the axislabel option.
            axislabel(boolean, default=True):
                If True, plotting the axislabel.
            colorbar(boolean, default=False):
                If True, the colorbar will be shown.
            colorbarprm(dic, default={}):
                parameters for pyplot.colorbar
            index(integer):
            **imshow_args: Args will be input in matplotlib.pyplot.imshow

        Returns:
            output from plt.imshow
        '''
        from .plot import plot_imshow
        return plot_imshow(
            image=self,
            scale=scale, dyrange=dyrange, gamma=gamma,
            vmax=vmax, vmin=vmin,
            relative=relative,
            fluxunit=fluxunit, saunit=saunit, angunit=angunit,
            axisoff=axisoff, axislabel=axislabel,
            cmap=cmap,
            idx=idx,
            interpolation=interpolation,
            **imshow_args
        )

    def plot_xylabel(self, **labelargs):
        """
        Add xy label in the existing image plot.

        Args:
            **labelargs:
                Arguments for pyplot.xlabel / ylable functions.
        """
        from .plot import plot_xylabel
        plot_xylabel(self, **labelargs)

    def plot_beam(self, boxfc=None, boxec=None, beamfc=None, beamec="white",
                  lw=1., alpha=0.5, x0=0.05, y0=0.05, boxsize=1.5, zorder=None):
        '''
        Plot beam in the header.
        To change the angular unit, please change IMFITS.angunit.

        Args:
            x0, y0(float, default=0.05):
                leftmost, lowermost location of the box
                if relative = True, the value is on transAxes coordinates
            relative(boolean, default=True):
                If True, the relative coordinate to the current axis will be
                used to plot data
            boxsize(float, default=1.5):
                Relative size of the box to the major axis size.
            boxfc, boxec(color formatter):
                Face and edge colors of the box
            beamfc, beamec(color formatter):
                Face and edge colors of the beam
            lw(float, default=1): linewidth
            alpha(float, default=0.5): transparency parameter (0 < 1)
                                       (0 < 1) for the face color

        Returns:
            dictionary of outputs from plt.plot and plt.fill
        '''
        from .plot import plot_beam
        return plot_beam(
            image=self,
            boxfc=boxfc, boxec=boxec, beamfc=beamfc, beamec=beamec, lw=lw,
            alpha=alpha, x0=x0, y0=y0, boxsize=boxsize, zorder=zorder
        )

    @ staticmethod
    def plot_scalebar(x, y, length, ha="center", color="white", lw=1,
                      **plotargs):
        '''
        Plot a scale bar

        Args:
            x, y ( in the unit of the current plot):
                x, y coordinates of the scalebar
            length (in the unit of the current plot):
                length of the scale bar
            ha(str, default="center"):
                The horizontal alignment of the bar.
                Available options is ["center", "left", "right"]
            plotars:
                Arbital arguments for pyplot.plot.
        '''
        from .plot import plot_scalebar
        return plot_scalebar(
            x=x, y=y, length=length, ha=ha, color=color, lw=lw, **plotargs
        )

    def to_fits(self, outfits=None, idx=None, fitsfmt="casa", overwrite=True):
        '''
        save the image(s) to the image FITS file or HDUList

        Args:
            outfits (string; default is None):
                FITS file name. If not specified, then HDUList object will be
                returned.
            idx (list):
                Index for images. For ehtim & CASA formats,
                it should be (MJD, FREQ).
            fitsfmr (str):
                The format for FITS file. Default is "casa".
                Available formats are ["casa", "ehtim"]
            overwrite (boolean):
                It True, an existing file will be overwritten.
        Returns:
            HDUList object if outfits is None
        '''
        from .io.fits_ehtim import save_image_to_fits_ehtim
        from .io.fits_casa import save_image_to_fits_casa

        if fitsfmt.lower() == "casa":
            output = save_image_to_fits_casa(
                self, outfits=outfits, overwrite=overwrite, idx=idx
            )
        elif fitsfmt.lower() == "ehtim":
            output = save_image_to_fits_ehtim(
                self, outfits=outfits, overwrite=overwrite, idx=idx
            )
        else:
            raise ValueError("fitsfmt=%s is not supported." % (fitsfmt))

        return output

    def to_hdf5_ehtim(self, file_name):
        '''
        Save image object in a hdf5 format

        Args:
            file_name (string): output hdf5 file name
        '''
        Ntime, Nfreq, Npol, Nx, Ny = self.ds.image.shape

        if Nx != Ny:
            raise ValueError(
                "Grid numbers of x, y coordinates are different Nx(=%d) != Ny(=%d)" % (Nx, Ny))
        if len(list(set(self.ds.freq.data))) > 1:
            raise ValueError(
                "This function cannot be used for multi epoch hdf5")
        if len(list(set(self.ds.pol.data))) > 1 or self.ds.pol.data[0] != "I":
            raise ValueError("This function is only used for Stokes I")

        mjd = int(self.ds.mjd.data[0])
        times = (self.ds.mjd.data - mjd)*24

        with h5py.File(file_name, "w") as file:
            head = file.create_dataset('header', (0,), dtype="S10")
            head.attrs['mjd'] = np.string_(str(mjd))
            head.attrs['psize'] = np.string_(str(Nx))
            head.attrs['source'] = np.string_(str(self.ds.source))
            head.attrs['ra'] = np.string_(str(self.ds.ra))
            head.attrs['dec'] = np.string_(str(self.ds.dec))
            head.attrs['rf'] = np.string_(str(self.ds.freq.data[0]))
            # If including polarization the following term should be updated
            head.attrs['polrep'] = np.string_('stokes')
            head.attrs['pol_prim'] = np.string_(self.ds.pol.data[0])
            dset = file.create_dataset("times", data=times, dtype='f8')

            # image table (Nt * Ny * Nx)
            name = "I"
            frames = np.stack(
                [np.flipud(self.ds.image.data[it, 0, 0, :, :]) for it in range(Ntime)])
            dset = file.create_dataset(name, data=frames, dtype='f8')

    def cpimage(self, im, save_totalflux=False, order=1):
        outim = copy.deepcopy(self)

        Ntime, Nfreq, Npol, Nx0, Ny0 = im.ds["image"].data.shape
        dx0 = im.ds.dx
        dy0 = im.ds.dy
        Nxr0 = im.ds.ixref
        Nyr0 = im.ds.iyref

        dummy, dummy, dummy, Nx1, Ny1 = outim.ds["image"].data.shape
        dx1 = outim.ds.dx
        dy1 = outim.ds.dy
        Nxr1 = outim.ds.ixref
        Nyr1 = outim.ds.iyref

        # output grid information
        coord = np.zeros([2, Nx1 * Ny1])
        xgrid = (np.arange(Nx1) + 1 - Nxr1) * dx1 / dx0 + Nxr0 - 1
        ygrid = (np.arange(Ny1) + 1 - Nyr1) * dy1 / dy0 + Nyr0 - 1
        x, y = np.meshgrid(xgrid, ygrid)
        coord[0, :] = y.reshape(Nx1 * Ny1)
        coord[1, :] = x.reshape(Nx1 * Ny1)

        for itime, ifreq, ipol in itertools.product(range(Ntime), range(Nfreq), range(Npol)):
            data = im.ds["image"].data[itime, ifreq, ipol]

            outim.ds["image"].data[itime, ifreq, ipol] = sn.map_coordinates(
                data, coord, order=order,
                mode='constant', cval=0.0, prefilter=True).reshape([Ny1, Nx1]
                                                                   ) * dx1 * dy1 / dx0 / dy0

            # Flux Scaling
            if save_totalflux:
                totalflux = im.totalflux(
                    itime=itime, ifreq=ifreq, ipol=ipol)
                outim.ds["image"].data[0, 0, 0] *= totalflux / \
                    outim.totalflux(itime=itime, ifreq=ifreq, ipol=ipol)

        # outim.update_fits()
        return outim

    def totalflux(self, fluxunit="Jy", itime=0, ifreq=0, ipol=0):
        '''
        Calculate the total flux density of the image

        Args:
          itime (integer): index for the image frame
          ipol (integer): index for polarization parameter at which the total flux will be calculated
          ifreq (integer): index for Frequency at which the total flux will be calculated
        '''
        return self.ds["image"].data[itime, ifreq, ipol].sum() * util.fluxconv("Jy", fluxunit)

    def lightcurve(self, fluxunit="Jy", ifreq=0, ipol=0):
        Ntime, dummy, dummy, dummy, dummy = self.ds["image"].data.shape
        return np.array([self.totalflux(fluxunit, itime, ifreq, ipol) for itime in range(Ntime)])

    ###########

def gen_image(
        nx=128, ny=None,
        dx=2., dy=None,
        ixref=None, iyref=None,
        angunit="uas",
        mjd=[0.],
        freq=[230e9],
        ns=1,
        source="M87",
        srccoord=None,
        instrument="EHT",
        **args):
    """
    Generate a blank image.

    Args:
        nx (int, optional): [description]. Defaults to 128.
        ny ([type], optional): [description]. Defaults to None.
        dx ([type], optional): [description]. Defaults to 2..
        dy ([type], optional): [description]. Defaults to None.
        ixref ([type], optional): [description]. Defaults to None.
        iyref ([type], optional): [description]. Defaults to None.
        angunit (str, optional): [description]. Defaults to "uas".
        mjd (list, optional): [description]. Defaults to [0.].
        freq (list, optional): [description]. Defaults to [230e9].
        ns (int, optional): [description]. Defaults to 1.
        source (str, optional): [description]. Defaults to "M87".
        srccoord ([type], optional): [description]. Defaults to None.
        instrument (str, optional): [description]. Defaults to "VLBI".
    """
    from numpy import float64, int32, abs, isscalar, zeros
    from xarray import Dataset
    from ...util.units import conv
    #  dx & dy
    factor = conv(angunit, "rad")
    dxrad = abs(float64(dx*factor))
    if dy is None:
        dyrad = dxrad
    else:
        dyrad = abs(float64(dy*factor))

    #  nx & ny
    nx = int32(nx)
    if ny is None:
        ny = nx
    else:
        ny = int32(ny)

    # ixref & iyref
    if ixref is None:
        ixref = nx/2 + 0.5
    else:
        ixref = float64(ixref)
    if iyref is None:
        iyref = ny/2 + 0.5
    else:
        iyref = float64(iyref)

    # ns
    ns = int32(ns)

    # nf
    if isscalar(freq):
        nf = int32(1)
    else:
        nf = len(freq)

    # nt
    if isscalar(mjd):
        nt = int32(1)
    else:
        nt = len(mjd)

    # create an empty image
    outim = Image(
        ds=Dataset(
            data_vars=dict(
                image=(Image.ds_dims, zeros(
                    [nt, nf, ns, ny, nx], dtype=float64))
            )
        )
    )

    # initialize meta data
    outim.init_attrs()
    outim.ds.attrs["dx"] = dxrad
    outim.ds.attrs["dy"] = dyrad
    outim.ds.attrs["ixref"] = ixref
    outim.ds.attrs["iyref"] = iyref
    outim.ds.attrs["angunit"] = angunit
    outim.set_source(source, srccoord)
    outim.set_instrument(instrument)

    # add coordinates
    outim.init_pol()
    outim.init_xygrid()
    outim.set_mjd(mjd)
    outim.set_freq(freq)

    return outim

def load_hdf5(hdf5file, angunit="uas"):
    '''
    Read data from the movie hdf5 file
    Args:
        hdf5file (string): hdf5 file name. currently the format is based on 2021 ehtim library.
    Returns:
        imdata.Image object
    '''

    # Read hdf5
    with h5py.File(hdf5file, "r") as file:
        head = file['header']
        mjd = int(head.attrs['mjd'].astype(str))
        times = np.float64(file['times'][:]/24)+mjd
        dx = np.float64(head.attrs['psize'].astype(
            str)) * util.angconv("rad", angunit)
        ra = np.float64(head.attrs['ra'].astype(str))
        dec = np.float64(head.attrs['dec'].astype(str))
        freq = np.float64(head.attrs['rf'].astype(str))
        polrep = head.attrs['polrep'].astype(str)
        pol_prim = head.attrs['pol_prim'].astype(str)
        source = head.attrs["source"].astype(str)
        data = file[pol_prim][:]  # time, y, x
        Ntime, Ny, Nx = data.shape

    # intialization of image object
    mov = gen_image(nx=Nx, ny=Ny, dx=dx, dy=dx, ixref=None, iyref=None,
                    angunit=angunit, mjd=times, freq=[freq], ns=1,
                    source=source, srccoord=None, instrument="EHT")

    logger.info("Construct movie frame")
    logger.info("Note: so far hdf5 includes single pol (=I) and single band")
    # logger.info("Note: need to invert the y direction")
    for itime in tqdm.tqdm(range(Ntime)):
        mov.ds["image"].data[itime, 0, 0] = data[itime, ::-1]

    return mov

def load_fits(infits, fitsfmt="casa"):
    """
    [summary]

    Args:
        infits (string or HDUList):
            Input FITS file name or HDUList object.
        fitsfmr (str):
            The format for FITS file. Default is "casa".
            Available formats are ["casa", "ehtim"]

    Returns:
        imdata.Image object
    """
    from .io.fits_ehtim import load_fits_ehtim
    from .io.fits_casa import load_fits_casa

    if fitsfmt.lower() == "casa":
        return load_fits_casa(infits)
    elif fitsfmt.lower() == "ehtim":
        return load_fits_ehtim(infits)
    else:
        raise ValueError("fitsfmt=%s is not supported." % (fitsfmt))

import numpy,from numpy,import scipy,from astropy,import jax


from styx_msgs.msg import TrafficLight
import tensorflow as tf
import numpy as np
import cv2
import rospy
import time

traffic_light_state = ['Green', 'Yellow', 'Red', 'Unknown']
"""
Traffic Light Classifier that uses a trained model given in the path and its information to detect
the light and hence the state of the traffic light.

"""
class TLClassifier(object):
    def __init__(self, is_site):
        if is_site:
            PATH_TO_FROZEN_GRAPH = 'light_classification/model/site_model/'
        else:
            PATH_TO_FROZEN_GRAPH = 'light_classification/model/sim_model/'

        #Used frozen_inference_graph.pb based on mentor instructions
        FROZEN_GRAPH = PATH_TO_FROZEN_GRAPH + 'frozen_inference_graph.pb'
        rospy.logerr(FROZEN_GRAPH)

        self.graph = tf.Graph()
        with self.graph.as_default():
            od_graph_def = tf.GraphDef()
            with tf.gfile.GFile(FROZEN_GRAPH, 'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                rospy.logerr("After parsing before Importing graph")
                tf.import_graph_def(od_graph_def, name='')

            # tensor information below
            # Image, detection box, detection score and detection class
            # create a tensor flow session for detection

            self.image_tensor = self.graph.get_tensor_by_name('image_tensor:0')
            self.detection_boxes = self.graph.get_tensor_by_name('detection_boxes:0')
            self.detection_scores = self.graph.get_tensor_by_name('detection_scores:0')
            self.detection_classes = self.graph.get_tensor_by_name('detection_classes:0')
        self.sess = tf.Session(graph=self.graph)

    def to_image_coords(self, boxes, height, width):
        """
        Function    : Convert the normalized coordinates to original coordinate based on the image size.

        Input Args  : traffic image Width, height and box information
        Output      : None
        """
        box_coords = np.zeros_like(boxes)
        box_coords[:, 0] = boxes[:, 0] * height
        box_coords[:, 1] = boxes[:, 1] * width
        box_coords[:, 2] = boxes[:, 2] * height
        box_coords[:, 3] = boxes[:, 3] * width

        return box_coords

    def draw_boxes(self, image, boxes, classes, scores):
        """
        Function    : Draw bounding boxes based on the input information

        Input Args  : traffic image, box information, scores and class detaails
        Output      : None
        """
        for i in range(len(boxes)):
            top, left, bot, right = boxes[i, ...]
            cv2.rectangle(image, (left, top), (right, bot), (255,0,0), 3)
            text = traffic_light_state[int(classes[i])-1] + ': ' + str(int(scores[i]*100)) + '%'
            cv2.putText(image , text, (left, int(top - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (200,0,0), 1, cv2.LINE_AA)

    def filter_boxes(self, min_score, boxes, scores, classes):
        """
        Function    : Perform bounding box filtering based on a minimum score
                      Return bounding box information where confidence >= a minimum score

        Input Args  : minimum confidence scores, bounding box info, scores and classes
        Output      : filtered box,filtered scores and filtered class detaails
        """

        n = len(classes)
        idxs = []
        for i in range(n):
            if scores[i] >= min_score:
                idxs.append(i)

        filtered_boxes = boxes[idxs, ...]
        filtered_scores = scores[idxs, ...]
        filtered_classes = classes[idxs, ...]
        return filtered_boxes, filtered_scores, filtered_classes

    def get_classification(self, image, is_site):
        """
        Function    : Determines the color of the traffic light in the image

        Input Args  : image containing the traffic light
        Output      : ID of traffic light color (specified in styx_msgs/TrafficLight)
        """
        timestr = time.strftime("%Y%m%d-%H%M%S")
        image = np.dstack((image[:, :, 2], image[:, :, 1], image[:, :, 0]))
        if is_site:
            width = image.shape[1]
            height = image.shape[0]
            #rospy.loginfo("Width: %r, height: %r" % (width, height))
            image = image[:int(height/2), :, :]
        image_np = np.expand_dims(np.asarray(image, dtype=np.uint8), 0)

        with tf.Session(graph=self.graph) as sess:
            # Actual detection.
            (boxes, scores, classes) = sess.run([self.detection_boxes, self.detection_scores, self.detection_classes],
                                                feed_dict={self.image_tensor: image_np})
            # Dimension reduction for performance efficiency and simplicity
            boxes = np.squeeze(boxes)
            scores = np.squeeze(scores)
            classes = np.squeeze(classes)

            confidence_cutoff = 0.7
            # Filter boxes with a confidence score less than `confidence_cutoff`
            boxes, scores, classes = self.filter_boxes(confidence_cutoff, boxes, scores, classes)

        # Save boxed images ; filename will be created uniquely based on timestamp
        write = True
        if write:
            image = np.dstack((image[:, :, 2], image[:, :, 1], image[:, :, 0]))
            width, height = image.shape[1], image.shape[0]
            box_coords = self.to_image_coords(boxes, height, width)
            self.draw_boxes(image, box_coords, classes, scores)
            #cv2.imwrite('/home/workspace/CarND-Capstone/imgs/img.jpg', image)
            fname = '/home/workspace/CarND-Capstone/imgs/img_'+timestr+'.jpg'
            cv2.imwrite(fname, image)

        if len(scores)>0:
            this_class = int(classes[np.argmax(scores)])
        else:
            this_class = 4

        if this_class == 1:
            return TrafficLight.GREEN
        elif this_class == 2:
             return TrafficLight.YELLOW
        elif this_class == 3:
             return TrafficLight.RED

        return TrafficLight.UNKNOWN


import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]= '0, 1, 2, 3'
import sys
import argparse
import logging
import time
import datetime
from collections import defaultdict

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.autograd import Variable

import torch.distributed as dist

from datasets.build import make_dataloader
import numpy as np
from pytorch_i3d import InceptionI3d
from torchvision.models.video import r3d_18, mc3_18, r2plus1d_18
from pytorch_c3d import C3D

# from charades_dataset import Charades as Dataset
from tqdm import tqdm
from model_serialization import load_state_dict
from detection.utils.logger import Logger
from detection.utils.metric_logger import MetricLogger
from detection.utils.comm import synchronize, get_rank

from mpi4py import MPI
import apex
from apex.parallel import DistributedDataParallel as DDP
from detection.utils.comm import get_world_size
from detection.utils.comm import is_main_process, all_gather, synchronize

from datasets.evaluation.evaluation import ActionClassificationEvaluator
from sklearn.metrics import ConfusionMatrixDisplay
import pickle as pkl
import pdb

np.set_printoptions(precision=3)

def reduce_loss_dict(loss_dict):
    """
    Reduce the loss dictionary from all processes so that process with rank
    0 has the averaged results. Returns a dict with the same fields as
    loss_dict, after reduction.
    """
    world_size = get_world_size()
    if world_size < 2:
        return loss_dict
    with torch.no_grad():
        loss_names = []
        all_losses = []
        for k in sorted(loss_dict.keys()):
            loss_names.append(k)
            all_losses.append(loss_dict[k])
        # print("all_losses:", all_losses)
        # print("\n")
        all_losses = torch.stack(all_losses, dim=0)
        if torch.isnan(torch.sum(all_losses)):
            pdb.set_trace()
        dist.reduce(all_losses, dst=0)
        if torch.isnan(torch.sum(all_losses)):
            pdb.set_trace()
        if dist.get_rank() == 0:
            # only main process gets accumulated, so only divide by
            # world_size in this case
            all_losses /= world_size

        reduced_losses = {k: v for k, v in zip(loss_names, all_losses)}
    return reduced_losses

def loss_func(pred, target):
    # sigmoid
    # bce_loss = F.binary_cross_entropy_with_logits(pred, target)
    # softmax
    ce_loss = F.cross_entropy(pred, target)
    return ce_loss

def do_train(model_name,
             model,
             train_dataloader,
             val_dataloader,
             device,
             checkpoint_peroid=1000,
             save_model='',
             logger=None,
             distributed=False,
             evaluator=None):

    lr = 0.01 #0.003 #init_lr
    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=0.0000001)
    # NOTE: maybe the weight decay is too soon?
    lr_sched = optim.lr_scheduler.MultiStepLR(optimizer, [4000, 8000])

    num_steps_per_update = 1 #4 # accum gradient
    steps = 0

    tot_loc_loss = 0.0
    tot_cls_loss = 0.0

    if logger is None:
        logger = logging.getLogger("I3D.trainer")
    logger.info("Start training")
    meters = MetricLogger(delimiter="  ")
    end = time.time()

    max_iters = len(train_dataloader)
    model.train()
    optimizer.zero_grad()
    # label_counts = defaultdict(int)
    for iters, data in enumerate(tqdm(train_dataloader)):
        data_time = time.time() - end
        iters += 1

        # get the inputs
        # inputs, labels, video_names, _, _ = data
        inputs, labels, video_names = data
        # wrap them in Variable
        inputs = Variable(inputs.to(device))
        t = inputs.size(2)
        labels = Variable(labels.to(device))
        per_frame_logits = model(inputs) # inputs: B X C X T X H X W
        # pdb.set_trace()
        if len(per_frame_logits.shape) == 3:
            per_frame_logits = per_frame_logits.mean(dim=-1) # B X C

        loss = loss_func(per_frame_logits, labels)
        # track time
        batch_time = time.time() - end

        # reduce losses over all GPUs for logging purposes
        loss_dict = {"loss_cls": loss} #{"loss_loc": loc_loss, "loss_cls": cls_loss}
        loss_dict_reduced = reduce_loss_dict(loss_dict)
        losses_reduced = loss_dict_reduced['loss_cls'] #0.5 * loss_dict_reduced['loss_loc'] + 0.5 * loss_dict_reduced['loss_cls']

        loss.backward()
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 5)

        meters.update(loss=losses_reduced, **loss_dict_reduced)
        meters.update(time=batch_time, data=data_time)

        # estimate the rest of the running time
        eta_seconds = meters.time.global_avg * (max_iters - iters)
        eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))

        if iters % num_steps_per_update == 0 or iters == len(train_dataloader):
            steps += 1
            optimizer.step()
            optimizer.zero_grad()

            lr_sched.step()
            if steps % 1 == 0:
                # NOTE: Add log file
                info = meters.delimiter.join(
                        [
                            "eta: {eta}",
                            "iter: {iter}",
                            "{meters}",
                            "lr: {lr:.6f}",
                            "max mem: {memory:.0f}",
                        ]
                    ).format(
                        eta=eta_string,
                        iter=iters,
                        meters=str(meters),
                        lr=optimizer.param_groups[0]["lr"],
                        memory=torch.cuda.max_memory_allocated() / 1024.0 / 1024.0,
                    )

                if hasattr(logger, 'log_values'):
                    logger.info(info)
                    for name, meter in meters.meters.items():
                        logger.log_values({name: meter.median}, step=iters)
                    logger.log_values({"grad_norm": grad_norm}, step=iters)
                else:
                    print(info)

            if steps % checkpoint_peroid == 0:
                del inputs, loss
                model.eval()
                do_val(model_name,
                       model,
                       val_dataloader,
                       device,
                       distributed,
                       logger,
                       output_dir=os.path.join(save_model, 'inference'),
                       train_iters=iters,
                       evaluator=evaluator)
                model.train()

                if get_rank() == 0: # only save model for rank0
                    if not os.path.isdir(save_model):
                        os.makedirs(save_model)

                    save_dir = os.path.join(save_model, str(steps).zfill(6)+'.pt')
                    if hasattr(model, 'module'):
                        torch.save(model.module.state_dict(), save_dir)
                    else:
                        torch.save(model.state_dict(), save_dir)

        end = time.time()
def do_val(model_name,
           model,
           val_dataloader,
           device,
           distributed=False,
           logger=None,
           output_dir='',
           train_iters=0,
           evaluator=None):
    if logger is None:
        logger = logging.getLogger("I3D.trainer")

    torch.cuda.empty_cache()  # TODO check if it helps

    tot_loc_loss = 0.0
    tot_cls_loss = 0.0

    # run on dataset
    results = defaultdict(list)
    target_labels = {}
    for iters, data in enumerate(tqdm(val_dataloader)):
        # get the inputs
        # inputs, labels, video_names, start, end = data
        inputs, labels, video_names= data
        # wrap them in Variable
        inputs = inputs.to(device)
        t = inputs.size(2)
        labels = labels.to(device)

        per_frame_logits = model(inputs)
        if len(per_frame_logits.shape) == 3:
            per_frame_logits = per_frame_logits.mean(dim=-1)

        loss = loss_func(per_frame_logits, labels)
        loss = loss.item() #(0.5*loc_loss + 0.5*cls_loss)
        tot_cls_loss += loss
        # collect results
        per_frame_logits = F.softmax(per_frame_logits, dim=1).detach().cpu()
        for batch_id, vid in enumerate(video_names):
            # frame_id = int((start[batch_id] + end[batch_id])/2)
            # pred = per_frame_logits[batch_id]
            # if vid not in results:
            #     results[vid] = {}
            # assert frame_id not in results[vid]
            # results[vid][frame_id] = pred
            pred = per_frame_logits[batch_id]
            results[vid].append(pred)
            target_labels[vid] = int(labels[batch_id].detach().cpu())

    if hasattr(logger, 'log_values'):
        logger.log_values({"loss_cls_val": tot_cls_loss/(iters+1)}, step=train_iters)

    results = _accumulate_from_multiple_gpus(results)
    if not is_main_process():
        return

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    # torch.save(results, os.path.join(output_dir, 'predictions.pth'))
    # Run video-level evaluation
    per_vid_confusion_matrix = np.zeros((16, 16))
    per_clip_confusion_matrix = np.zeros((16, 16))
    per_vid_top3_acc = np.zeros(16)
    per_clip_top3_acc = np.zeros(16)
    per_class_clip_num = np.zeros(16)
    # Save per video prediction
    all_per_vid_results = {}
    for vid in results.keys():
        # ----------
        # NOTE: add clip-level accuracy
        clip_results = torch.stack(results[vid], dim=0)
        for per_clip_result in clip_results:
            _, sorted_pred = torch.sort(per_clip_result, descending=True)
            top1_pred = sorted_pred[0]
            top3_pred = sorted_pred[:3]
            y_true = target_labels[vid]
            per_class_clip_num[y_true] += 1
            per_clip_confusion_matrix[y_true, top1_pred] += 1
            if y_true in top3_pred:
                per_clip_top3_acc[y_true] += 1
        # ---------

        # ---------
        # per video results
        per_vid_results = torch.stack(results[vid], dim=0).mean(dim=0)
        all_per_vid_results[vid] = per_vid_results
        _, sorted_pred = torch.sort(per_vid_results, descending=True)
        top1_pred = sorted_pred[0]
        top3_pred = sorted_pred[:3]
        y_true = target_labels[vid]
        per_vid_confusion_matrix[y_true, top1_pred] += 1
        if y_true in top3_pred:
            per_vid_top3_acc[y_true] += 1
        # ---------

    # Save per video prediction
    pkl.dump(all_per_vid_results, open(os.path.join('/u/bryao/work/DATA/i3d_outputs', model_name, 'per_video_results.pkl'), 'wb'))

    # per clip results
    per_clip_top_1_acc = per_clip_confusion_matrix.diagonal() / (per_clip_confusion_matrix.sum(axis=1) + 1e-6)
    per_clip_top_3_acc = np.array([num/per_class_clip_num[i] for i, num in enumerate(per_clip_top3_acc)])
    logger.info("Clip-level evalutaion:")
    per_clip_result = "Top 1 acc: {}, Top 3 acc: {}, per_cls_acc: {}".format(np.around(per_clip_top_1_acc.mean(), 3),
                                                                       np.around(per_clip_top_3_acc.mean(), 3),
                                                                       np.around(per_clip_top_1_acc, 3))
    logger.info(per_clip_result)
    print(per_clip_result)
    print("confusion matrix:", np.around(per_clip_confusion_matrix, 3))
    # np.save('/u/bryao/work/DATA/i3d_outputs/r2plus1d_18/confusion_matrix_7000.npy', per_clip_confusion_matrix)
    if hasattr(logger, 'log_values'):
        logger.log_values({'Per clip Top 1 Acc': per_clip_top_1_acc.mean()}, step=train_iters)
        logger.log_values({'Per clip Top 3 Acc': per_clip_top_3_acc.mean()}, step=train_iters)

    # per video results
    data_category_stats = val_dataloader.dataset.data_category_stats
    top_1_acc = per_vid_confusion_matrix.diagonal() / (per_vid_confusion_matrix.sum(axis=1) + 1e-6)
    top_3_acc = np.array([num/data_category_stats[i] for i, num in enumerate(per_vid_top3_acc)])
    logger.info("Video-level evalutaion:")
    per_vid_result = "Top 1 acc: {}, Top 3 acc: {}, per_cls_acc: {}".format(np.around(top_1_acc.mean(), 3),
                                                                       np.around(top_3_acc.mean(), 3),
                                                                       np.around(top_1_acc, 3))
    logger.info(per_vid_result)
    print(per_vid_result)
    if hasattr(logger, 'log_values'):
        logger.log_values({'Per vid Top 1 Acc': top_1_acc.mean()}, step=train_iters)
        logger.log_values({'Per vid Top 3 Acc': top_3_acc.mean()}, step=train_iters)

    # # Run video-level evaluation
    # eval_results = evaluator.evaluate(results)
    # for k, v in eval_results.items():
    #     logger.info('{}:{}'.format(k, v))
    #     if isinstance(v, (float, int)) and hasattr(logger, 'log_values'):
    #         logger.log_values(eval_results, step=train_iters)

    # if 'confusion_matrix' in eval_results:
    #     cm_display = ConfusionMatrixDisplay(eval_results['confusion_matrix'],
    #                         display_labels=val_dataloader.dataset.name_shorts)
    #     ret = cm_display.plot(fontsize=6)
    #     logger.log_plot(ret.figure_, label='Confusion Matrix',step=train_iters)

def _accumulate_from_multiple_gpus(item_per_gpu):
    # all_keys
    all_items = all_gather(item_per_gpu)
    if not is_main_process():
        return
    # merge the list of dicts
    predictions = {}
    for p in all_items:
        predictions.update(p)
    return predictions

def run(model_name='i3d',
        init_lr=0.1,
        gpu_id=0,
        max_steps=64e3,
        batch_per_gpu=4,
        mode='rgb',
        root='/home/data/vision7/A3D_2.0/frames/', #'/ssd/Charades_v1_rgb',
        train_split='A3D_2.0_train.json', #'charades/charades.json',
        val_split='A3D_2.0_val.json',
        checkpoint_peroid=1000,
        save_model='',
        with_normal=True):

    num_gpus = MPI.COMM_WORLD.Get_size()
    distributed = False
    if num_gpus > 1:
        distributed = True

    local_rank = MPI.COMM_WORLD.Get_rank() % torch.cuda.device_count()

    if distributed:
        torch.cuda.set_device(local_rank)
        host = os.environ["MASTER_ADDR"] if "MASTER_ADDR" in os.environ else "127.0.0.1"
        torch.distributed.init_process_group(
            backend="nccl",
            init_method='tcp://{}:12345'.format(host),
            rank=MPI.COMM_WORLD.Get_rank(),
            world_size=MPI.COMM_WORLD.Get_size()
        )

        synchronize()
    # logger must be initialized after distributed!
    cfg = {'PROJECT': 'i3d_a3d'}
    if args.use_wandb:
        logger = Logger("I3D",
                        cfg,#convert_to_dict(cfg, []),
                        project = 'i3d_a3d',
                        viz_backend="wandb"
                        )
    else:
        logger = logging.Logger('I3D')

    logger.info("Using {} GPUs".format(num_gpus))
    # setup dataset

    train_dataloader = make_dataloader(root,
                                       train_split,
                                       mode,
                                       model_name,
                                       seq_len=16, #64,
                                       overlap=15, #32,
                                       phase='train',
                                       max_iters=10000,
                                       batch_per_gpu=batch_per_gpu, #8,
                                       num_workers=16,
                                       shuffle=True,
                                       distributed=distributed,
                                       with_normal=with_normal)

    val_dataloader = make_dataloader(root,
                                     val_split,
                                     mode,
                                     model_name,
                                     seq_len=16, #64,
                                     overlap=15, #32,
                                     phase='val',
                                     max_iters=None,
                                     batch_per_gpu=batch_per_gpu, #8,
                                     num_workers=16,
                                     shuffle=False,
                                     distributed=distributed,
                                     with_normal=with_normal)
    # setup the model
    # set  dropout_keep_prob=0.0 for overfit
    logger.info("Running {} model".format(model_name))
    if model_name == 'i3d':
        if mode == 'flow':
            model = InceptionI3d(train_dataloader.dataset.num_classes, in_channels=2, dropout_keep_prob=0.5)
            load_state_dict(model, torch.load('models/flow_imagenet.pt'), ignored_prefix='logits')
            logger.info("Loaded pretrained I3D")
        else:
            model = InceptionI3d(train_dataloader.dataset.num_classes, in_channels=3, dropout_keep_prob=0.5)
            load_state_dict(model, torch.load('models/rgb_imagenet.pt'), ignored_prefix='logits')
            logger.info("Loaded pretrained I3D")
        model.replace_logits(train_dataloader.dataset.num_classes)
    elif model_name == 'r3d_18':
        model = r3d_18(pretrained=True, num_classes=train_dataloader.dataset.num_classes)
    elif model_name == 'mc3_18':
        model = mc3_18(pretrained=True, num_classes=train_dataloader.dataset.num_classes)
    elif model_name == 'r2plus1d_18':
        model = r2plus1d_18(pretrained=True, num_classes=train_dataloader.dataset.num_classes)
    elif model_name == 'c3d':
        model = C3D(pretrained=True, num_classes=train_dataloader.dataset.num_classes)
    else:
        raise NameError('unknown model name:{}'.format(model_name))
    if hasattr(logger, 'run_id'):
        run_id = logger.run_id
    else:
        run_id = 'no_wandb'

    save_model = os.path.join(save_model, model_name, run_id)

    # Create evaluator
    evaluator = ActionClassificationEvaluator(cfg=None,
                                              dataset=val_dataloader.dataset,
                                              split='val',
                                              mode='accuracy',#'mAP',
                                              output_dir=save_model,
                                              with_normal=with_normal)

    # device = torch.device('cuda')
    device = torch.device('cuda:{}'.format(gpu_id))
    model.to(device)
    if distributed:
        model = apex.parallel.convert_syncbn_model(model)
        model = DDP(model.to(device), delay_allreduce=True)
    # ckpt = '/home/data/vision7/brianyao/DATA/i3d_outputs/i3d/tg9s4ff5/004500.pt'
    # load_state_dict(model, torch.load(ckpt))

    # #######NOTE: only for testing checkpoint is correct or not ##
    # model.eval()
    # do_val(model,
    #         val_dataloader,
    #         device,
    #         distributed,
    #         logger,
    #         output_dir=os.path.join(save_model, 'inference'),
    #         train_iters=0,
    #         evaluator=evaluator)
    # pdb.set_trace()
    # #########
    do_train(model_name,
             model,
             train_dataloader,
             val_dataloader,
             device=device,
             checkpoint_peroid=checkpoint_peroid,
             save_model=save_model,
             logger=logger,
             evaluator=evaluator)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-model_name', type=str, help='name of the model to run')
    parser.add_argument('-mode', type=str, help='rgb or flow')
    parser.add_argument('-save_model', type=str)
    parser.add_argument('-root', type=str)
    parser.add_argument('-train_split', type=str)
    parser.add_argument('-val_split', type=str)
    parser.add_argument('-batch_per_gpu', type=int)
    parser.add_argument('-gpu', type=int)
    parser.add_argument('-checkpoint_peroid', type=int)
    parser.add_argument('-use_wandb', const=True, nargs='?')
    args = parser.parse_args()

    # need to add argparse
    run(model_name=args.model_name,
        mode=args.mode,
        gpu_id=args.gpu,
        batch_per_gpu=args.batch_per_gpu,
        train_split=args.train_split,
        val_split=args.val_split,
        root=args.root,
        save_model=args.save_model,
        checkpoint_peroid=args.checkpoint_peroid,
        with_normal=False
        )


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

# Lint as: python3
"""Model components for an end-to-end-trainable graph/automaton hybrid.

The components defined in this module share a common interface

  (input graph, node embeddings, edge embeddings)
    -> (node embeddings, edge embeddings)

which allows them to be composed with each other. Note that most components
either modify node embeddings or edge embeddings but not both. Also note that
the output edge embeddings are allowed to be of a different size; in particular,
the components that add new edge types use SharedGraphContext.edges_are_embedded
to determine how to modify the edge embeddings.
"""

from typing import Dict, List, Optional, Tuple

import dataclasses
import flax
import gin
import jax
import jax.numpy as jnp

from gfsa import automaton_builder
from gfsa import jax_util
from gfsa.datasets import graph_bundle
from gfsa.model import automaton_layer
from gfsa.model import edge_supervision_models
from gfsa.model import graph_layers
from gfsa.model import model_util
from gfsa.model import side_outputs

# TODO(ddjohnson) Move common layers out of `edge_supervision_models`.

# Flax adds name keyword arguments.
# pylint: disable=unexpected-keyword-arg

NodeAndEdgeEmbeddings = Tuple[jax_util.NDArray, jax_util.NDArray]

@dataclasses.dataclass
class SharedGraphContext:
  """Shared information about the input graph.

  Attributes:
    bundle: The input graph.
    static_metadata: Padded size of the graph.
    edge_types_to_indices: Mapping from string edge names to edge type indices.
    builder: Automaton builder associated with this graph.
    edges_are_embedded: Whether the "edge_embeddings" represent edge types that
      are embedded into vectors (True), or just edge type adjacency matrices
      that are concatenated together.
  """
  bundle: graph_bundle.GraphBundle
  static_metadata: automaton_builder.EncodedGraphMetadata
  edge_types_to_indices: Dict[str, int]
  builder: automaton_builder.AutomatonBuilder
  edges_are_embedded: bool

def _add_edges(old_edge_array,
               new_edge_types,
               edges_are_embedded,
               add_reverse = True):
  """Helper function to add edges of a new edge type.

  If edges_are_embedded=True, we assume `old_edge_dim` is an embedding matrix;
  the new edges are embedded and then added into this matrix. Otherwise, we
  assume `old_edge_dim` is a stacked set of adjacency matrices, and concatenate
  the new types.

  Args:
    old_edge_array: <float32[num_nodes, num_nodes, old_edge_dim]>
    new_edge_types: <float32[num_nodes, num_nodes, new_edge_types]>, which
      should be between 0 and 1, one for each added edge type.
    edges_are_embedded: Whether edge types are embedded.
    add_reverse: Whether to add reverse edges as well with a different type.

  Returns:
    <float32[num_nodes, num_nodes, output_edge_dim]>, where
    output_edge_dim = old_edge_dim if edges_are_embedded=True, and otherwise
    output_edge_dim = old_edge_dim + new_edge_types
  """
  if add_reverse:
    new_edge_types = jnp.concatenate(
        [new_edge_types, new_edge_types.transpose((1, 0, 2))], -1)
  if edges_are_embedded:
    # Project the outputs into new edge embeddings.
    # (No bias is used so that an absorbing probability of 0 produces no change
    # in the edge embeddings.)
    new_edge_type_embeddings = flax.deprecated.nn.Dense(
        new_edge_types,
        features=old_edge_array.shape[-1],
        bias=False,
        name="new_edge_type_embeddings")
    output_edge_array = old_edge_array + new_edge_type_embeddings
  else:
    # Concatenate new embedding.
    output_edge_array = jnp.concatenate([old_edge_array, new_edge_types],
                                        axis=-1)

  return output_edge_array

def _shared_automaton_logic(
    graph_context, node_embeddings,
    edge_embeddings,
    variant_weights):
  """Helper function for shared automaton logic."""

  # Run the automaton.
  edge_weights = automaton_layer.FiniteStateGraphAutomaton(
      encoded_graph=graph_context.bundle.automaton_graph,
      variant_weights=variant_weights,
      dynamic_metadata=graph_context.bundle.graph_metadata,
      static_metadata=graph_context.static_metadata,
      builder=graph_context.builder)

  return (node_embeddings,
          _add_edges(edge_embeddings, edge_weights.transpose([1, 2, 0]),
                     graph_context.edges_are_embedded))

@flax.deprecated.nn.module
@gin.configurable
def variantless_automaton(
    graph_context, node_embeddings,
    edge_embeddings):
  """Runs an automaton without variants.

  Args:
    graph_context: Input graph for this example.
    node_embeddings: Current node embeddings, as <float32[num_nodes,
      node_embedding_dim]>
    edge_embeddings: Current edge embeddings, as <float32[num_nodes, num_nodes,
      edge_embedding_dim]>

  Returns:
    New node and edge embeddings. Node embeddings will not be modified. Edge
    embeddings will be modified by adding a new edge type (either embedded or
    concatenated based on graph_context.edges_are_embedded).
  """
  return _shared_automaton_logic(
      graph_context, node_embeddings, edge_embeddings, variant_weights=None)

@flax.deprecated.nn.module
@gin.configurable
def edge_variant_automaton(
    graph_context,
    node_embeddings,
    edge_embeddings,
    variant_edge_types = gin.REQUIRED):
  """Runs an automaton with variants based on edges in the input graph.

  Args:
    graph_context: Input graph for this example.
    node_embeddings: Current node embeddings, as <float32[num_nodes,
      node_embedding_dim]>
    edge_embeddings: Current edge embeddings, as <float32[num_nodes, num_nodes,
      edge_embedding_dim]>
    variant_edge_types: List of edge types used as variants.

  Returns:
    New node and edge embeddings. Node embeddings will not be modified. Edge
    embeddings will be modified by adding a new edge type (either embedded or
    concatenated based on graph_context.edges_are_embedded).
  """
  # Set up variants from edge types.
  variant_edge_type_indices = [
      graph_context.edge_types_to_indices[type_str]
      for type_str in variant_edge_types
  ]
  num_edge_types = len(graph_context.edge_types_to_indices)
  variant_weights = edge_supervision_models.variants_from_edges(
      graph_context.bundle, graph_context.static_metadata,
      variant_edge_type_indices, num_edge_types)

  return _shared_automaton_logic(graph_context, node_embeddings,
                                 edge_embeddings, variant_weights)

@flax.deprecated.nn.module
@gin.configurable
def embedding_variant_automaton(
    graph_context,
    node_embeddings,
    edge_embeddings,
    num_variants = gin.REQUIRED):
  """Runs an automaton with variants based on node embeddings.

  Args:
    graph_context: Input graph for this example.
    node_embeddings: Current node embeddings, as <float32[num_nodes,
      node_embedding_dim]>
    edge_embeddings: Current edge embeddings, as <float32[num_nodes, num_nodes,
      edge_embedding_dim]>
    num_variants: How many variants to use.

  Returns:
    New node and edge embeddings. Node embeddings will not be modified. Edge
    embeddings will be modified by adding a new edge type (either embedded or
    concatenated based on graph_context.edges_are_embedded).
  """
  if num_variants <= 1:
    raise ValueError(
        "Must have at least one variant to use embedding_variant_automaton.")
  # Generate variants using a pairwise readout of the node embeddings.
  variant_logits = graph_layers.BilinearPairwiseReadout(
      node_embeddings, num_variants, name="variant_logits")

  variant_logits = side_outputs.encourage_discrete_logits(
      variant_logits, distribution_type="categorical", name="variant_logits")
  variant_weights = jax.nn.softmax(variant_logits)

  return _shared_automaton_logic(graph_context, node_embeddings,
                                 edge_embeddings, variant_weights)

@flax.deprecated.nn.module
@gin.configurable
def nri_encoder_readout(
    graph_context,
    node_embeddings,
    edge_embeddings,
    num_edge_types = gin.REQUIRED):
  """Modifies edge embeddings using an NRI encoder.

  Note that we use a sigmoid rather than a softmax, because we don't necessarily
  want to enforce having exactly one edge type per pair of nodes.

  Args:
    graph_context: Input graph for this example.
    node_embeddings: Current node embeddings, as <float32[num_nodes,
      node_embedding_dim]>
    edge_embeddings: Current edge embeddings, as <float32[num_nodes, num_nodes,
      edge_embedding_dim]>
    num_edge_types: How many edge types to produce.

  Returns:
    New node and edge embeddings. Node embeddings will not be modified. Edge
    embeddings will be modified by adding a new edge type (either embedded or
    concatenated based on graph_context.edges_are_embedded).
  """
  # Run the NRI readout layer.
  logits = graph_layers.NRIReadout(
      node_embeddings=node_embeddings, readout_dim=num_edge_types)
  new_edge_weights = jax.nn.sigmoid(logits)

  mask = (
      jnp.arange(new_edge_weights.shape[0]) <
      graph_context.bundle.graph_metadata.num_nodes)
  new_edge_weights = jnp.where(mask[:, None, None], new_edge_weights,
                               jnp.zeros_like(new_edge_weights))

  return (node_embeddings,
          _add_edges(edge_embeddings, new_edge_weights,
                     graph_context.edges_are_embedded))

class UniformRandomWalk(flax.deprecated.nn.Module):
  """Adds edges according to a uniform random walk along the graph."""

  @gin.configurable("UniformRandomWalk")
  def apply(
      self,
      graph_context,
      node_embeddings,
      edge_embeddings,
      forward_edge_types = gin.REQUIRED,
      reverse_edge_types = gin.REQUIRED,
      walk_length_log2 = gin.REQUIRED,
  ):
    """Modifies edge embeddings using a uniform random walk.

    Uses an efficient repeated-squaring technique to compute the absorbing
    distribution.

    Args:
      graph_context: Input graph for this example.
      node_embeddings: Current node embeddings, as <float32[num_nodes,
        node_embedding_dim]>
      edge_embeddings: Current edge embeddings, as <float32[num_nodes,
        num_nodes, edge_embedding_dim]>
      forward_edge_types: Edge types to use in the forward direction. As a list
        of lists to allow configuring groups of edges in config files; this will
        be flattened before use.
      reverse_edge_types: Edge types to use in the reverse direction. Note that
        reversed edge types are given a separate embedding from forward edge
        types; undirected edges should be represented by adding two edges in
        opposite directions and then only using `forward_edge_types`. Also a
        list of lists, as above.
      walk_length_log2: Base-2 logarithm of maximum walk length; this determines
        how many times we will square the transition matrix (doubling the walk
        length).

    Returns:
      New node and edge embeddings. Node embeddings will not be modified. Edge
      embeddings will be modified by adding a new edge type (either embedded or
      concatenated based on graph_context.edges_are_embedded).
    """
    num_nodes = node_embeddings.shape[0]

    # pylint: disable=g-complex-comprehension
    forward_edge_type_indices = [
        graph_context.edge_types_to_indices[type_str]
        for group in forward_edge_types
        for type_str in group
    ]
    reverse_edge_type_indices = [
        graph_context.edge_types_to_indices[type_str]
        for group in reverse_edge_types
        for type_str in group
    ]
    # pylint: enable=g-complex-comprehension

    adjacency = graph_layers.edge_mask(
        edges=graph_context.bundle.edges,
        num_nodes=num_nodes,
        num_edge_types=len(graph_context.edge_types_to_indices),
        forward_edge_type_indices=forward_edge_type_indices,
        reverse_edge_type_indices=reverse_edge_type_indices)
    adjacency = jnp.maximum(adjacency, jnp.eye(num_nodes))

    absorbing_logit = self.param(
        "absorbing_logit",
        shape=(),
        initializer=lambda *_: jax.scipy.special.logit(0.1))

    absorbing_prob = jax.nn.sigmoid(absorbing_logit)
    nonabsorbing_prob = jax.nn.sigmoid(-absorbing_logit)
    walk_matrix = nonabsorbing_prob * adjacency / jnp.sum(
        adjacency, axis=1, keepdims=True)

    # A, I
    # A^2, A + I
    # (A^2)^2 = A^4, (A + I)A^2 + (A + I) = A^3 + A^2 + A + I
    # ...

    def step(state, _):
      nth_power, nth_partial_sum = state
      return (nth_power @ nth_power,
              nth_power @ nth_partial_sum + nth_partial_sum), None

    (_, partial_sum), _ = jax.lax.scan(
        step, (walk_matrix, jnp.eye(num_nodes)), None, length=walk_length_log2)

    approx_visits = absorbing_prob * partial_sum
    logits = model_util.safe_logit(approx_visits)
    logits = model_util.ScaleAndShift(logits)
    edge_weights = jax.nn.sigmoid(logits)

    return (node_embeddings,
            _add_edges(edge_embeddings, edge_weights[:, :, None],
                       graph_context.edges_are_embedded))

@flax.deprecated.nn.module
def ggnn_adapter(graph_context,
                 node_embeddings,
                 edge_embeddings):
  """Adapter function to run GGNN steps.

  Args:
    graph_context: Input graph for this example.
    node_embeddings: Current node embeddings, as <float32[num_nodes,
      node_embedding_dim]>
    edge_embeddings: Current edge embeddings, as <float32[num_nodes, num_nodes,
      edge_embedding_dim]>

  Returns:
    New node and edge embeddings. Node embeddings are processed by a GGNN,
    and edge embeddings are returned unchanged.
  """
  del graph_context
  return (
      edge_supervision_models.ggnn_steps(node_embeddings, edge_embeddings),
      edge_embeddings,
  )

@flax.deprecated.nn.module
def transformer_adapter(
    graph_context, node_embeddings,
    edge_embeddings):
  """Adapter function to run transformer blocks.

  Args:
    graph_context: Input graph for this example.
    node_embeddings: Current node embeddings, as <float32[num_nodes,
      node_embedding_dim]>
    edge_embeddings: Current edge embeddings, as <float32[num_nodes, num_nodes,
      edge_embedding_dim]>

  Returns:
    New node and edge embeddings. Node embeddings are processed by a
    transformer, and edge embeddings are returned unchanged.
  """
  return (
      edge_supervision_models.transformer_steps(
          node_embeddings,
          edge_embeddings,
          neighbor_mask=None,
          num_real_nodes_per_graph=(
              graph_context.bundle.graph_metadata.num_nodes),
          mask_to_neighbors=False),
      edge_embeddings,
  )

@flax.deprecated.nn.module
def nri_adapter(graph_context,
                node_embeddings,
                edge_embeddings):
  """Adapter function to run NRI blocks.

  Args:
    graph_context: Input graph for this example.
    node_embeddings: Current node embeddings, as <float32[num_nodes,
      node_embedding_dim]>
    edge_embeddings: Current edge embeddings, as <float32[num_nodes, num_nodes,
      edge_embedding_dim]>

  Returns:
    New node and edge embeddings. Node embeddings are processed by a NRI-style
    model, and edge embeddings are returned unchanged.
  """
  return (
      edge_supervision_models.nri_steps(
          node_embeddings,
          edge_embeddings,
          num_real_nodes_per_graph=(
              graph_context.bundle.graph_metadata.num_nodes)),
      edge_embeddings,
  )

ALL_COMPONENTS = {
    "variantless_automaton": variantless_automaton,
    "edge_variant_automaton": edge_variant_automaton,
    "embedding_variant_automaton": embedding_variant_automaton,
    "nri_encoder_readout": nri_encoder_readout,
    "ggnn_adapter": ggnn_adapter,
    "transformer_adapter": transformer_adapter,
    "nri_adapter": nri_adapter,
    "UniformRandomWalk": UniformRandomWalk,
}


import os
import sys
import logging

import cv2
from keras.models import load_model
import numpy as np

from utils.datasets import get_labels
from utils.inference import detect_faces
from utils.inference import draw_text
from utils.inference import draw_bounding_box
from utils.inference import apply_offsets
from utils.inference import load_detection_model
from utils.inference import load_image
from utils.preprocessor import preprocess_input

def process_image(image):

    try:
        # parameters for loading data and images
        detection_model_path = './trained_models/detection_models/haarcascade_frontalface_default.xml'
        emotion_model_path = './trained_models/emotion_models/fer2013_35a93d1a30cf52b8722837f3e096b7b1006949325de8d00f26595d5a418241d3.hdf5'
        emotion_labels = get_labels('fer2013')
        font = cv2.FONT_HERSHEY_SIMPLEX

        # hyper-parameters for bounding boxes shape
        emotion_offsets = (20, 40)
        emotion_offsets = (0, 0)

        # loading models
        face_detection = load_detection_model(detection_model_path)
        emotion_classifier = load_model(emotion_model_path, compile=False)

        # getting input model shapes for inference
        emotion_target_size = emotion_classifier.input_shape[1:3]

        # loading images
        image_array = np.fromstring(image, np.uint8)
        unchanged_image = cv2.imdecode(image_array, cv2.IMREAD_UNCHANGED)

        rgb_image = cv2.cvtColor(unchanged_image, cv2.COLOR_BGR2RGB)
        gray_image = cv2.cvtColor(unchanged_image, cv2.COLOR_BGR2GRAY)

        faces = detect_faces(face_detection, gray_image)
        for face_coordinates in faces:
            x1, x2, y1, y2 = apply_offsets(face_coordinates, emotion_offsets)
            gray_face = gray_image[y1:y2, x1:x2]

            try:
                gray_face = cv2.resize(gray_face, (emotion_target_size))
            except:
                continue

            rgb_face = preprocess_input(rgb_face, False)
            rgb_face = np.expand_dims(rgb_face, 0)

            gray_face = preprocess_input(gray_face, True)
            gray_face = np.expand_dims(gray_face, 0)
            gray_face = np.expand_dims(gray_face, -1)
            emotion_label_arg = np.argmax(emotion_classifier.predict(gray_face))
            emotion_text = emotion_labels[emotion_label_arg]

            color = (255, 0, 0)

            draw_bounding_box(face_coordinates, rgb_image, color)
            draw_text(face_coordinates, rgb_image, emotion_text, color, 0, -50, 1, 2)
    except Exception as err:
        logging.error('Error in emotion gender processor: "{0}"'.format(err))

    bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)

    dirname = 'result'
    if not os.path.exists(dirname):
        os.mkdir(dirname)

    cv2.imwrite(os.path.join(dirname, 'predicted_image.png'), bgr_image)


"""DQN Agent."""

import tensorflow as tf
import numpy as np
from tf_agents.trajectories import trajectory, time_step as ts
from tf_agents.networks import q_network
from tf_agents.networks import q_rnn_network
from tf_agents.agents.dqn import dqn_agent
from tf_agents.replay_buffers import tf_uniform_replay_buffer
import random
from tf_agents.policies import random_tf_policy
from tf_agents.policies import fixed_policy

class DQNAgent:

  def __init__(self, config, *args, **kwargs):
    """Initialize the agent."""
    self.observation_spec = config['observation_spec']
    self.time_step_spec = ts.time_step_spec(self.observation_spec)
    self.action_spec = config['action_spec']
    self.environment_batch_size = config['environment_batch_size']

    self.q_net = q_rnn_network.QRnnNetwork(
      self.observation_spec,
      self.action_spec,
      input_fc_layer_params=(256, ),
      lstm_size=(256, 256))

    self.optimizer = tf.keras.optimizers.Adam()

    self.agent = dqn_agent.DqnAgent(
      self.time_step_spec,
      self.action_spec,
      q_network=self.q_net,
      optimizer=self.optimizer,
      epsilon_greedy=0.3)

    self.agent.initialize()

    self.replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
      data_spec=self.agent.collect_data_spec,
      batch_size=1,
      dataset_drop_remainder=True)

    self.dataset = self.replay_buffer.as_dataset(
      sample_batch_size=1,
      num_steps=2,
      single_deterministic_pass=True)

    self.intention_classifier = tf.keras.models.Sequential(
      [tf.keras.layers.Dense(1, input_shape=self.observation_spec.shape, activation='sigmoid')])

    self.intention_classifier.compile(optimizer=self.optimizer, loss=tf.keras.losses.BinaryCrossentropy())

    self.intention_dataset_input = []
    self.intention_dataset_true = []

    self.num_hinted = 0

  def initial_policy_state(self):
    return self.agent.policy.get_initial_state(self.environment_batch_size)

  def time_step_converter(self, time_step):
    if time_step.is_first():
      tf_time_step = ts.restart(
        tf.convert_to_tensor(time_step.observation, dtype=tf.int32),
        self.environment_batch_size)
    elif time_step.is_last():
      tf_time_step = ts.termination(
        tf.convert_to_tensor(time_step.observation, dtype=tf.int32),
        tf.convert_to_tensor(time_step.reward, dtype=tf.float32))
    else:
      tf_time_step = ts.transition(
        tf.convert_to_tensor(time_step.observation, dtype=tf.int32),
        tf.convert_to_tensor(time_step.reward, dtype=tf.float32))

    return tf_time_step

  def act(self, time_step, policy_state, prev_knowledge, current_knowledge, card_playable):

    time_step = self.time_step_converter(time_step)

    hinted = []
    candidate_action = -1
    for i, cards in enumerate(zip(prev_knowledge, current_knowledge)):
      if cards[0] != cards[1]:
        hinted.append(i)

    if len(hinted) == 1:
      candidate_action = 5 + hinted[0]

    self.intention_dataset_input.append(time_step.observation[0])
    self.intention_dataset_true.append(int(card_playable))

    confidence_of_hint = self.intention_classifier(time_step.observation).numpy()[0,0]

    if (confidence_of_hint > 0.85) and (candidate_action != -1):
      print("Hinted card is played!")
      print(current_knowledge)
      self.num_hinted += 1
      candidate_action = tf.convert_to_tensor(candidate_action, dtype=tf.int32)
      policy = fixed_policy.FixedPolicy(candidate_action, self.time_step_spec, self.action_spec)
      time_step = ts.restart(tf.ones(self.time_step_spec.observation.shape, dtype=tf.int32))
      action_step = policy.action(time_step)
    else:
      action_step = self.agent.collect_policy.action(time_step, policy_state)

    return action_step

  def collect(self, last_time_step, action_step, next_time_step):
    last_time_step = self.time_step_converter(last_time_step)
    next_time_step = self.time_step_converter(next_time_step)
    traj = trajectory.from_transition(last_time_step, action_step, next_time_step)
    self.replay_buffer.add_batch(traj)

  def learn(self):
    if self.replay_buffer.num_frames() > 1:
      for batched_experience in self.dataset:
        experience, _ = batched_experience
        self.agent.train(experience)

    self.replay_buffer.clear()

    self.intention_classifier.fit(np.asarray(self.intention_dataset_input), np.asarray(self.intention_dataset_true), batch_size=1)
    self.intention_dataset_input = []
    self.intention_dataset_true = []

    self.num_hinted = 0


"""
Generate speaker embeddings and metadata for training
"""
import os
import pickle
from model_bl import D_VECTOR
from collections import OrderedDict
import numpy as np
import torch
import pandas as pd

class Metadata(object):

    def __init__(self, config):
        """Initialize configurations."""

        self.speaker_embed = config.speaker_embed
        self.main_dir = config.main_dir
        self.model_type = config.model_type
        self.root_dir = self.main_dir+'/'+self.model_type # containing spmel or stft spects
        self.num_uttrs = 10
        #we use spmel for metadata speaker encoding regardless
        self.len_crop = 128

        self.subject_conversions = [ #((original speaker, sentence), target speaker)
                        #(('p002','020'),'p002'), #(identic)
                        #(('p002','020'),'p227'), #seen sentence, between different genders
                        #(('p002', '020'), 'p003'), # unseen subject, seen sentence
                        #(('p003', '4019'), 'p002'), # unseen cubject and sentence
                        (('p225', '001'), 'p225'),
                        #(('p225', '001'), 'p227'), # like autovc paper
                        #(('p227', '001'), 'p225'),
                        #(('p227', '003'), 'p002')
                        ]
        self.main_dir = config.main_dir

        self.speaker_info = pd.read_csv('speaker_info.txt', delim_whitespace=True)

    def metadata(self):
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        C = D_VECTOR(dim_input=80, dim_cell=768, dim_emb=256).eval().to(device)

        c_checkpoint = torch.load('3000000-BL.ckpt',map_location=device)
        new_state_dict = OrderedDict()
        for key, val in c_checkpoint['model_b'].items():
            new_key = key[7:]
            new_state_dict[new_key] = val
        C.load_state_dict(new_state_dict)

        num_uttrs = self.num_uttrs
        len_crop = self.len_crop

        # speaker embedding is created from mel! (not stft, since checkpoint 300000-BL.pth is created for mel)
        self.mel_dir = self.main_dir+'/spmel'
        dirName, subdirList, _ = next(os.walk(self.mel_dir))
        print('Found directory: %s' % dirName)

        speakers = []
        for speaker in sorted(subdirList):
            print('Processing speaker: %s' % speaker)
            utterances = []
            utterances.append(speaker)
            _, _, fileList = next(os.walk(os.path.join(dirName,speaker)))

            # make speaker embedding
            assert len(fileList) >= num_uttrs
            idx_uttrs = np.random.choice(len(fileList), size=num_uttrs, replace=False)
            embs = []
            for i in range(num_uttrs):
                tmp = np.load(os.path.join(dirName, speaker, fileList[idx_uttrs[i]]))
                candidates = np.delete(np.arange(len(fileList)), idx_uttrs)
                # choose another utterance if the current one is too short
                while tmp.shape[0] < len_crop:
                    idx_alt = np.random.choice(candidates)
                    tmp = np.load(os.path.join(dirName, speaker, fileList[idx_alt]))
                    candidates = np.delete(candidates, np.argwhere(candidates==idx_alt))
                left = np.random.randint(0, tmp.shape[0]-len_crop)
                melsp = torch.from_numpy(tmp[np.newaxis, left:left+len_crop, :]).to(device)
                emb = C(melsp)
                embs.append(emb.detach().squeeze().cpu().numpy())
            utterances.append(np.mean(embs, axis=0)) # average speaker embedding

            # create file list
            for fileName in sorted(fileList):
                utterances.append(os.path.join(speaker,fileName))
            speakers.append(utterances)

        with open(os.path.join(self.root_dir, 'train.pkl'), 'wb') as handle:
            pickle.dump(speakers, handle)

        ######### Our modification: Reading numpy files in speaker embedding ##########
        with open(os.path.join(self.root_dir, 'train.pkl'), 'rb') as file:
            train = pickle.load(file)
        #construct hash map for easy retrieval
        subject_speaker_embedding = {}
        for embed in train:
            subject_speaker_embedding[embed[0]] = embed[1]

        with open(os.path.join(self.root_dir, 'metadata.log'), 'w') as log:
        #with open('metadata.log', 'w') as log:
            log_ref_int = 0
            metadata = [] #format is list of conversion metadata [log ref int(eventual filename), [from subject.sentence, from embedding, from sound input], [to subject, to embedding]
            for conversion in self.subject_conversions:
                log.write('CONVERSION FILENAME: '+str(log_ref_int) +' ' + '#'*40 + '\n\n')
                with open(os.path.join(self.main_dir,'txt',conversion[0][0],conversion[0][0] + '_' + conversion[0][1]+'.txt',), 'r') as sentence_file:
                    sentence = "\"" + sentence_file.readline().rstrip('\n').rstrip() + "\""
                    log.write(f'Converting from sentence no. {conversion[0][1]} : {sentence} \n')

                log.write('Uttered by the speaker:\n')
                log.write(self.speaker_info[self.speaker_info['ID']==conversion[0][0]].to_string(index=False))
                log.write('\n')

                log.write('To the speaker:\n')
                log.write(self.speaker_info[self.speaker_info['ID']==conversion[1]].to_string(index=False))
                log.write('\n\n')

                #bad use for try/except, checking for mic2. I DON'T CARE. YOLO

                try:
                    sound_input = np.load(self.root_dir+'/'+conversion[0][0] + '/' + conversion[0][0] + '_' + conversion[0][1] + '_mic2.npy')
                except:
                    sound_input = np.load(self.root_dir+'/'+conversion[0][0] + '/' + conversion[0][0] + '_' + conversion[0][1] + '.npy')

                metadata.append([log_ref_int,
                                [conversion[0][0]+'_'+conversion[0][1], subject_speaker_embedding[conversion[0][0]], sound_input],
                                [conversion[1], subject_speaker_embedding[conversion[1]]]
                                ])

                log_ref_int = log_ref_int + 1

            with open(os.path.join(self.root_dir, 'metadata.pkl'), 'wb') as handle:
                pickle.dump(metadata, handle)
        print('Finished generating metadata')


# encoding: utf-8
"""
@author:  liaoxingyu
@contact: sherlockliao01@gmail.com
"""

import numpy as np
import pytrec_eval

def eval_func(distmat, q_pids, g_pids, q_camids, g_camids, max_rank=100):
    """Evaluation with market1501 metric
        Key: for each query identity, its gallery images from the same camera view are discarded.
        """
    num_q, num_g = distmat.shape
    if num_g < max_rank:
        max_rank = num_g
        print("Note: number of gallery samples is quite small, got {}".format(num_g))
    indices = np.argsort(distmat, axis=1)
    matches = (g_pids[indices] == q_pids[:, np.newaxis]).astype(np.int32)

    # compute cmc curve for each query
    all_cmc = []
    all_AP = []
    num_valid_q = 0.  # number of valid query
    for q_idx in range(num_q):
        # get query pid and camid
        q_pid = q_pids[q_idx]
        q_camid = q_camids[q_idx]

        # remove gallery samples that have the same pid and camid with query
        order = indices[q_idx]
        remove = (g_pids[order] == q_pid) & (g_camids[order] == q_camid)
        keep = np.invert(remove)

        # compute cmc curve
        # binary vector, positions with value 1 are correct matches
        orig_cmc = matches[q_idx][keep]
        if not np.any(orig_cmc):
            # this condition is true when query identity does not appear in gallery
            continue

        cmc = orig_cmc.cumsum()
        cmc[cmc > 1] = 1

        all_cmc.append(cmc[:max_rank])
        num_valid_q += 1.

        # compute average precision
        # reference: https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Average_precision
        num_rel = orig_cmc.sum()
        tmp_cmc = orig_cmc.cumsum()
        tmp_cmc = [x / (i + 1.) for i, x in enumerate(tmp_cmc)]
        tmp_cmc = np.asarray(tmp_cmc) * orig_cmc
        AP = tmp_cmc.sum() / num_rel
        all_AP.append(AP)

    assert num_valid_q > 0, "Error: all query identities do not appear in gallery"

    all_cmc = np.asarray(all_cmc).astype(np.float32)
    all_cmc = all_cmc.sum(0) / num_valid_q
    # mAP = np.mean(all_AP)
    mAP = pytrec_mAP(distmat, q_pids, g_pids, q_camids, g_camids, topk=max_rank)

    return all_cmc, mAP

def pytrec_mAP(distmat, query_ids=None, gallery_ids=None, query_cams=None, gallery_cams=None, topk=100):
    m, n = distmat.shape
    if query_ids is None:
        query_ids = np.arange(m)
    if gallery_ids is None:
        gallery_ids = np.arange(n)
    if query_cams is None:
        query_cams = np.zeros(m).astype(np.int32)
    if gallery_cams is None:
        gallery_cams = np.ones(n).astype(np.int32)
    query_ids = np.asarray(query_ids)
    gallery_ids = np.asarray(gallery_ids)
    query_cams = np.asarray(query_cams)
    gallery_cams = np.asarray(gallery_cams)
    gt = dict()
    pt = dict()
    for i in range(m):
        gt[str(i)] = dict()
        pt[str(i)] = dict()
        valid = ((gallery_ids != query_ids[i]) |
                 (gallery_cams != query_cams[i]))
        valid_mat = distmat[i, valid]
        match = (gallery_ids[valid] == query_ids[i])
        indices = np.argsort(valid_mat)
        for j in range(valid_mat.shape[0]):
            if match[j]: gt[str(i)][str(j)] = 1
            else: gt[str(i)][str(j)] = 0
        for j, idx in enumerate(indices):
            if j < 100:
                pt[str(i)][str(idx)] = 100 - j
    evaluator = pytrec_eval.RelevanceEvaluator(gt, {'map_cut'})
    res = evaluator.evaluate(pt)
    map_list = [v['map_cut_{}'.format(topk)] for k, v in res.items()]
    return np.mean(map_list)