from __future__ import division, absolute_import, print_function

from builtins import range

import numpy as np
import scipy.interpolate as interpolate
import scipy.integrate as integrate
import os
import sys
from pkg_resources import resource_exists
from pkg_resources import resource_filename
from pkg_resources import resource_listdir

try:
    import fitsio
    fits_package = 'fitsio'
except ImportError:
    try:
        import astropy.io.fits as pyfits
        fits_package = 'pyfits'
    except ImportError:
        try:
            import pyfits
            fits_package = 'pyfits'
        except ImportError:
            pass

from .modtranGenerator import ModtranGenerator

from .sharedNumpyMemManager import SharedNumpyMemManager as snmm
from .fgcmLogger import FgcmLogger

class FgcmAtmosphereTable(object):
    """
    Class to generate, store, and interpolate pre-computed atmospheres.

    parameters
    ----------
    atmConfig: dict
       Dictionary with FgcmAtmosphereTable config variables
    atmosphereTableFile: str, optional
       Name of the table file
    fgcmLog: FgcmLog, optional
       Logger
    """

    def __init__(self, atmConfig, atmosphereTableFile=None, fgcmLog=None):

        if fgcmLog is None:
            self.fgcmLog = FgcmLogger('dummy.log', 'INFO', printLogger=True)
        else:
            self.fgcmLog = fgcmLog

        if atmosphereTableFile is not None:
            # We have an atmosphereTableFile, so read that
            self.atmosphereTableFile = atmosphereTableFile
            self.atmConfig = atmConfig
        else:
            # get modTran stuff ready since we're going to be generating
            # also check the config
            requiredKeys = ['elevation',
                            'pmbRange','pmbSteps',
                            'pwvRange','pwvSteps',
                            'o3Range','o3Steps',
                            'tauRange','tauSteps',
                            'alphaRange','alphaSteps',
                            'zenithRange','zenithSteps',
                            'pmbStd','pwvStd','o3Std',
                            'tauStd','alphaStd','airmassStd']

            for key in requiredKeys:
                if key not in atmConfig:
                    raise ValueError("Required %s not in atmConfig" % (key))
                if 'Range' in key:
                    if len(atmConfig[key]) != 2:
                        raise ValueError("%s must have 2 elements" % (key))

            self.atmConfig = atmConfig

            self.elevation = self.atmConfig['elevation']

            self.modGen = ModtranGenerator(self.elevation)
            self.pmbElevation = self.modGen.pmbElevation

            self.pmbStd = self.atmConfig['pmbStd']
            self.pwvStd = self.atmConfig['pwvStd']
            self.lnPwvStd = np.log(self.pwvStd)
            self.o3Std = self.atmConfig['o3Std']
            self.tauStd = self.atmConfig['tauStd']
            self.lnTauStd = np.log(self.tauStd)
            self.alphaStd = self.atmConfig['alphaStd']
            self.secZenithStd = self.atmConfig['airmassStd']
            self.zenithStd = np.arccos(1./self.secZenithStd)*180./np.pi

            if ('lambdaRange' in self.atmConfig):
                self.lambdaRange = np.array(self.atmConfig['lambdaRange'])
            else:
                self.lambdaRange = np.array([3000.0,11000.0])

            if ('lambdaStep' in self.atmConfig):
                self.lambdaStep = self.atmConfig['lambdaStep']
            else:
                self.lambdaStep = 0.5

            if ('lambdaNorm' in self.atmConfig):
                self.lambdaNorm = self.atmConfig['lambdaNorm']
            else:
                self.lambdaNorm = 7750.0

        self.o2Interpolator = None
        self.o3Interpolator = None
        self.rayleighInterpolator = None
        self.lnPwvInterpolator = None

    @staticmethod
    def getAvailableTables():
        """
        Get list of installed tables

        parameters
        ----------
        None

        returns
        -------
        List of installed table names
        """

        try:
            files = resource_listdir(__name__,'data/tables/')
        except:
            raise IOError("Could not find associated data/tables path!")

        # build a dictionary and put in the key details of each...

        availableTables = {}

        for f in files:
            availableTables[os.path.basename(f)] = FgcmAtmosphereTable.getInfoDict(f)

        return availableTables

    @staticmethod
    def getInfoDict(atmosphereTableFile):
        """
        Get information dictionary on a table file

        parameters
        ----------
        atmosphereTableFile: str
           Name of the table file

        returns
        -------
        infoDict: dict
           Dictionary with key parameters describing the table
        """

        if fits_package == 'fitsio':
            parStruct = fitsio.read(atmosphereTableFile, ext='PARS')
        elif fits_package == 'pyfits':
            parStruct = pyfits.getdata(atmosphereTableFile, ext=('PARS', 1))
        else:
            raise IOError("Reading atmosphere tables not supported without fitsio/astropy.io.fits/pyfits")

        infoDict = {'elevation':parStruct['ELEVATION'][0],
                    'pmbRange':[parStruct['PMB'][0][0], parStruct['PMB'][0][-1]],
                    'pmbSteps':parStruct['PMB'][0].size,
                    'pwvRange':[np.exp(parStruct['LNPWV'][0][0]),
                                np.exp(parStruct['LNPWV'][0][-1])],
                    'pwvSteps':parStruct['LNPWV'][0].size,
                    'o3Range':[parStruct['O3'][0][0], parStruct['O3'][0][-1]],
                    'o3Steps':parStruct['O3'][0].size,
                    'tauRange':[np.exp(parStruct['LNTAU'][0][0]), np.exp(parStruct['LNTAU'][0][-1])],
                    'tauSteps':parStruct['LNTAU'][0].size,
                    'alphaRange':[parStruct['ALPHA'][0][0], parStruct['ALPHA'][0][-1]],
                    'alphaSteps':parStruct['ALPHA'][0].size,
                    'zenithRange':[np.rad2deg(np.arccos(1./parStruct['SECZENITH'][0][0])),
                                   np.rad2deg(np.arccos(1./parStruct['SECZENITH'][0][-1]))],
                    'zenithSteps':parStruct['SECZENITH'][0].size,
                    'pmbStd':parStruct['PMBSTD'][0],
                    'pwvStd':parStruct['PWVSTD'][0],
                    'o3Std':parStruct['O3STD'][0],
                    'tauStd':parStruct['TAUSTD'][0],
                    'alphaStd':parStruct['ALPHASTD'][0],
                    'airmassStd':parStruct['AIRMASSSTD'][0]}

        return infoDict

    @classmethod
    def initWithTableName(cls, atmosphereTableName):
        """
        Initialize FgcmAtmosphereTable with a table name

        parameters
        ----------
        atmosphereTableName: str
           Name of atmosphere table

        returns
        -------
        FgcmAtmosphereTable: FgcmAtmosphereTable object

        notes
        -----
        Raises IOError if atmosphereTableName couldn't be found
        """

        # first, check if we have something in the path...
        if os.path.isfile(atmosphereTableName):
            atmosphereTableFile = os.path.abspath(atmosphereTableName)
            print("Found atmosphereTableName: %s" % (atmosphereTableName))
        else:
            # allow for a name with or without .fits extension
            if resource_exists(__name__,'data/tables/%s' % (atmosphereTableName)):
                testFile = 'data/tables/%s' % (atmosphereTableName)
            elif resource_exists(__name__,'data/tables/%s.fits' % (atmosphereTableName)):
                testFile = 'data/tables/%s.fits' % (atmosphereTableName)
            else:
                raise IOError("Could not find atmosphereTableName (%s) in the path or in data/tables/" % (atmosphereTableName))
            try:
                atmosphereTableFile = resource_filename(__name__,testFile)
            except:
                raise IOError("Error finding atmosphereTableName (%s)" % (atmosphereTableName))

        # will set self.atmConfig
        print("Found atmosphere file: %s" % (atmosphereTableFile))
        atmConfig = FgcmAtmosphereTable.getInfoDict(atmosphereTableFile)

        return cls(atmConfig, atmosphereTableFile=atmosphereTableFile)

    def loadTable(self):
        """
        Load atmosphere table

        parameters
        ----------
        None

        returns
        -------
        None
        """

        # note that at this point these are the only options
        if fits_package == 'fitsio':
            parStruct = fitsio.read(self.atmosphereTableFile, ext='PARS')
        else:
            parStruct = pyfits.getdata(self.atmosphereTableFile, ext=('PARS', 1))

        self.elevation = parStruct['ELEVATION'][0]
        self.pmbElevation = parStruct['PMBELEVATION'][0]
        self.pmbStd = parStruct['PMBSTD'][0]
        self.pwvStd = parStruct['PWVSTD'][0]
        self.o3Std = parStruct['O3STD'][0]
        self.tauStd = parStruct['TAUSTD'][0]
        self.alphaStd = parStruct['ALPHASTD'][0]
        self.secZenithStd = parStruct['AIRMASSSTD'][0]
        self.lambdaRange = parStruct['LAMBDARANGE'][0]
        self.lambdaStep = parStruct['LAMBDASTEP'][0]
        self.lambdaNorm = parStruct['LAMBDANORM'][0]

        self.atmLambda = parStruct['ATMLAMBDA'][0]
        self.atmStdTrans = parStruct['ATMSTDTRANS'][0]

        self.pmb = parStruct['PMB'][0]
        self.pmbDelta = parStruct['PMBDELTA'][0]
        self.lnPwv = parStruct['LNPWV'][0]
        self.lnPwvDelta = parStruct['LNPWVDELTA'][0]
        self.o3 = parStruct['O3'][0]
        self.o3Delta = parStruct['O3DELTA'][0]
        self.lnTau = parStruct['LNTAU'][0]
        self.lnTauDelta = parStruct['LNTAUDELTA'][0]
        self.alpha = parStruct['ALPHA'][0]
        self.alphaDelta = parStruct['ALPHADELTA'][0]
        self.secZenith = parStruct['SECZENITH'][0]
        self.secZenithDelta = parStruct['SECZENITHDELTA'][0]

        if fits_package == 'fitsio':
            self.pwvAtmTable = fitsio.read(self.atmosphereTableFile, ext='PWVATM')
            self.o3AtmTable = fitsio.read(self.atmosphereTableFile, ext='O3ATM')
            self.o2AtmTable = fitsio.read(self.atmosphereTableFile, ext='O2ATM')
            self.rayleighAtmTable = fitsio.read(self.atmosphereTableFile, ext='RAYATM')
        else:
            self.pwvAtmTable = pyfits.getdata(self.atmosphereTableFile, ext=('PWVATM', 1))
            self.o3AtmTable = pyfits.getdata(self.atmosphereTableFile, ext=('O3ATM', 1))
            self.o2AtmTable = pyfits.getdata(self.atmosphereTableFile, ext=('O2ATM', 1))
            self.rayleighAtmTable = pyfits.getdata(self.atmosphereTableFile, ext=('RAYATM', 1))

    def generateTable(self):
        """
        Generate atmosphere table using MODTRAN

        parameters
        ----------
        None

        returns
        -------
        None
        """
        self.atmStd = self.modGen(pmb=self.pmbStd,pwv=self.pwvStd,
                                  o3=self.o3Std,tau=self.tauStd,
                                  alpha=self.alphaStd,zenith=self.zenithStd,
                                  lambdaRange=self.lambdaRange/10.0,
                                  lambdaStep=self.lambdaStep)
        self.atmLambda = self.atmStd['LAMBDA']
        self.atmStdTrans = self.atmStd['COMBINED']

        # get all the steps
        self.pmb = np.linspace(self.atmConfig['pmbRange'][0],
                               self.atmConfig['pmbRange'][1],
                               num=self.atmConfig['pmbSteps'])
        self.pmbDelta = self.pmb[1] - self.pmb[0]
        pmbPlus = np.append(self.pmb, self.pmb[-1] + self.pmbDelta)

        self.lnPwv = np.linspace(np.log(self.atmConfig['pwvRange'][0]),
                                 np.log(self.atmConfig['pwvRange'][1]),
                                 num=self.atmConfig['pwvSteps'])
        self.lnPwvDelta = self.lnPwv[1] - self.lnPwv[0]
        lnPwvPlus = np.append(self.lnPwv, self.lnPwv[-1] + self.lnPwvDelta)
        pwvPlus = np.exp(lnPwvPlus)

        self.o3 = np.linspace(self.atmConfig['o3Range'][0],
                               self.atmConfig['o3Range'][1],
                               num=self.atmConfig['o3Steps'])
        self.o3Delta = self.o3[1] - self.o3[0]
        o3Plus = np.append(self.o3, self.o3[-1] + self.o3Delta)

        self.lnTau = np.linspace(np.log(self.atmConfig['tauRange'][0]),
                                 np.log(self.atmConfig['tauRange'][1]),
                                 num=self.atmConfig['tauSteps'])
        self.lnTauDelta = self.lnTau[1] - self.lnTau[0]
        self.tau = np.exp(self.lnTau)
        lnTauPlus = np.append(self.lnTau, self.lnTau[-1] + self.lnTauDelta)
        tauPlus = np.exp(lnTauPlus)

        self.alpha = np.linspace(self.atmConfig['alphaRange'][0],
                               self.atmConfig['alphaRange'][1],
                               num=self.atmConfig['alphaSteps'])
        self.alphaDelta = self.alpha[1] - self.alpha[0]
        alphaPlus = np.append(self.alpha, self.alpha[-1] + self.alphaDelta)

        self.secZenith = np.linspace(1./np.cos(self.atmConfig['zenithRange'][0]*np.pi/180.),
                                     1./np.cos(self.atmConfig['zenithRange'][1]*np.pi/180.),
                                     num=self.atmConfig['zenithSteps'])
        self.secZenithDelta = self.secZenith[1]-self.secZenith[0]
        self.zenith = np.arccos(1./self.secZenith)*180./np.pi
        secZenithPlus = np.append(self.secZenith, self.secZenith[-1] + self.secZenithDelta)
        zenithPlus = np.arccos(1./secZenithPlus)*180./np.pi

        # run MODTRAN a bunch of times

        self.fgcmLog.info("Generating %d*%d=%d PWV atmospheres..." % (pwvPlus.size,zenithPlus.size,pwvPlus.size*zenithPlus.size))
        self.pwvAtmTable = np.zeros((pwvPlus.size,zenithPlus.size,self.atmLambda.size))

        for i in range(pwvPlus.size):
            sys.stdout.write('%d' % (i))
            sys.stdout.flush()
            for j in range(zenithPlus.size):
                sys.stdout.write('.')
                sys.stdout.flush()
                atm=self.modGen(pwv=pwvPlus[i],zenith=zenithPlus[j],
                                lambdaRange=self.lambdaRange/10.0,
                                lambdaStep=self.lambdaStep)
                self.pwvAtmTable[i,j,:] = atm['H2O']

        self.fgcmLog.info("\nGenerating %d*%d=%d O3 atmospheres..." % (o3Plus.size,zenithPlus.size,o3Plus.size*zenithPlus.size))
        self.o3AtmTable = np.zeros((o3Plus.size, zenithPlus.size, self.atmLambda.size))

        for i in range(o3Plus.size):
            sys.stdout.write('%d' % (i))
            sys.stdout.flush()
            for j in range(zenithPlus.size):
                sys.stdout.write('.')
                sys.stdout.flush()
                atm=self.modGen(o3=o3Plus[i],zenith=zenithPlus[j],
                                lambdaRange=self.lambdaRange/10.0,
                                lambdaStep=self.lambdaStep)
                self.o3AtmTable[i,j,:] = atm['O3']

        self.fgcmLog.info("\nGenerating %d O2/Rayleigh atmospheres..." % (zenithPlus.size))

        self.o2AtmTable = np.zeros((zenithPlus.size, self.atmLambda.size))
        self.rayleighAtmTable = np.zeros((zenithPlus.size, self.atmLambda.size))

        for j in range(zenithPlus.size):
            sys.stdout.write('.')
            sys.stdout.flush()
            atm=self.modGen(zenith=zenithPlus[j],
                            lambdaRange=self.lambdaRange/10.0,
                            lambdaStep=self.lambdaStep)
            self.o2AtmTable[j,:] = atm['O2']
            self.rayleighAtmTable[j,:] = atm['RAYLEIGH']

        self.fgcmLog.info("\nDone.")

    def saveTable(self, fileName, clobber=False):
        """
        Save atmosphere table to a file

        parameters
        ----------
        fileName: str
           File name to save file
        clobber: bool, default=False
           Clobber fileName if it exists
        """
        # at the moment, only work with fitsio...

        import fitsio

        if os.path.isfile(fileName) and not clobber:
            raise IOError("File %s already exists and clobber is set to False" % (fileName))

        fits = fitsio.FITS(fileName, 'rw')

        parStruct = np.zeros(1, dtype=[('ELEVATION','f4'),
                                       ('PMBELEVATION','f4'),
                                       ('PMBSTD','f4'),
                                       ('PWVSTD','f4'),
                                       ('O3STD','f4'),
                                       ('TAUSTD','f4'),
                                       ('ALPHASTD','f4'),
                                       ('AIRMASSSTD','f4'),
                                       ('LAMBDARANGE','f4',2),
                                       ('LAMBDASTEP','f4'),
                                       ('LAMBDANORM','f4'),
                                       ('ATMLAMBDA','f4',self.atmLambda.size),
                                       ('ATMSTDTRANS','f4',self.atmStdTrans.size),
                                       ('PMB','f4',self.pmb.size),
                                       ('PMBDELTA','f4'),
                                       ('LNPWV','f4',self.lnPwv.size),
                                       ('LNPWVDELTA','f4'),
                                       ('O3','f4',self.o3.size),
                                       ('O3DELTA','f4'),
                                       ('LNTAU','f4',self.lnTau.size),
                                       ('LNTAUDELTA','f4'),
                                       ('ALPHA','f4',self.alpha.size),
                                       ('ALPHADELTA','f4'),
                                       ('SECZENITH','f4',self.secZenith.size),
                                       ('SECZENITHDELTA','f4')])

        parStruct['ELEVATION'] = self.elevation
        parStruct['PMBELEVATION'] = self.pmbElevation
        parStruct['PMBSTD'] = self.pmbStd
        parStruct['PWVSTD'] = self.pwvStd
        parStruct['O3STD'] = self.o3Std
        parStruct['TAUSTD'] = self.tauStd
        parStruct['ALPHASTD'] = self.alphaStd
        parStruct['AIRMASSSTD'] = self.secZenithStd
        parStruct['LAMBDARANGE'][:] = self.lambdaRange
        parStruct['LAMBDASTEP'][:] = self.lambdaStep
        parStruct['LAMBDANORM'][:] = self.lambdaNorm

        parStruct['ATMLAMBDA'][:] = self.atmLambda
        parStruct['ATMSTDTRANS'][:] = self.atmStdTrans

        parStruct['PMB'][:] = self.pmb
        parStruct['PMBDELTA'] = self.pmbDelta
        parStruct['LNPWV'][:] = self.lnPwv
        parStruct['LNPWVDELTA'] = self.lnPwvDelta
        parStruct['O3'][:] = self.o3
        parStruct['O3DELTA'] = self.o3Delta
        parStruct['LNTAU'][:] = self.lnTau
        parStruct['LNTAUDELTA'] = self.lnTauDelta
        parStruct['ALPHA'][:] = self.alpha
        parStruct['ALPHADELTA'] = self.alphaDelta
        parStruct['SECZENITH'][:] = self.secZenith
        parStruct['SECZENITHDELTA'] = self.secZenithDelta

        fits.write_table(parStruct, extname='PARS')

        fits.write_image(self.pwvAtmTable, extname='PWVATM')
        fits.write_image(self.o3AtmTable, extname='O3ATM')
        fits.write_image(self.o2AtmTable, extname='O2ATM')
        fits.write_image(self.rayleighAtmTable, extname='RAYATM')

        fits.close()

    def interpolateAtmosphere(self, lam=None, pmb=None, pwv=None, o3=None, tau=None,
                              alpha=None, zenith=None, ctranslamstd=None):
        """
        Interpolate the atmosphere table to generate an atmosphere without
        requiring MODTRAN be installed.

        parameters
        ----------
        lam: float array, optional
           wavelengths.  Default is self.atmLambda
        pmb: float, optional
           pressure in millibar.  Default to pmbStd
        pwv: float, optional
           pwv in mm.  Default to pwvStd
        o3: float, optional
           ozone in Dobson.  Default to o3Std
        tau: float, optional
           Aerosol optical depth.  Default to tauStd
        alpha: float, optional
           Aerosol slope.  Default to alphaStd
        zenith: float, optional
           Zenith angle (degrees).  Default to zenithStd
        ctranslamstd: [ctrans, lamstd], optional
           Transmission adjustment constant and lambdastd.
           Default to [0.0, lamnorm]

        returns
        -------
        atmInterpolated: float array
           Interpolated atmosphere at self.atmLambda wavelengths
        """

        if lam is None:
            lam = self.atmLambda
        if pmb is None:
            pmb = self.pmbStd
        if pwv is None:
            pwv = self.pwvStd
        if o3 is None:
            o3 = self.o3Std
        if tau is None:
            tau = self.tauStd
        if alpha is None:
            alpha = self.alphaStd
        if zenith is None:
            zenith = self.zenithStd
        if ctranslamstd is None:
            ctranslamstd = [0.0, 7750.0]

        if self.o2Interpolator is None:
            secZenithPlus = np.append(self.secZenith, self.secZenith[-1] + self.secZenithDelta)
            self.o2Interpolator = interpolate.RegularGridInterpolator((secZenithPlus, self.atmLambda), self.o2AtmTable)

        if self.rayleighInterpolator is None:
            secZenithPlus = np.append(self.secZenith, self.secZenith[-1] + self.secZenithDelta)
            self.rayleighInterpolator = interpolate.RegularGridInterpolator((secZenithPlus, self.atmLambda), self.rayleighAtmTable)

        if self.lnPwvInterpolator is None:
            lnPwvPlus = np.append(self.lnPwv, self.lnPwv[-1] + self.lnPwvDelta)
            secZenithPlus = np.append(self.secZenith, self.secZenith[-1] + self.secZenithDelta)
            self.lnPwvInterpolator = interpolate.RegularGridInterpolator((lnPwvPlus, secZenithPlus, self.atmLambda), self.pwvAtmTable)

        if self.o3Interpolator is None:
            o3Plus = np.append(self.o3, self.o3[-1] + self.o3Delta)
            secZenithPlus = np.append(self.secZenith, self.secZenith[-1] + self.secZenithDelta)
            self.o3Interpolator = interpolate.RegularGridInterpolator((o3Plus, secZenithPlus, self.atmLambda), self.o3AtmTable)

        # And finally the pressure correction
        pmbMolecularScattering = np.exp(-(pmb - self.pmbElevation) / self.pmbElevation)
        pmbMolecularAbsorption = pmbMolecularScattering ** 0.6
        pmbFactor = pmbMolecularScattering * pmbMolecularAbsorption

        _secZenith = np.clip(1./np.cos(np.radians(zenith)), self.secZenith[0], self.secZenith[-1])
        _lnPwv = np.clip(np.log(pwv), self.lnPwv[0], self.lnPwv[-1])
        _o3 = np.clip(o3, self.o3[0], self.o3[-1])
        atmInterpolated = (pmbFactor *
                           self.o2Interpolator((_secZenith, lam)) *
                           self.rayleighInterpolator((_secZenith, lam)) *
                           self.lnPwvInterpolator((_lnPwv, _secZenith, lam)) *
                           self.o3Interpolator((_o3, _secZenith, lam)) *
                           (1.0 + ctranslamstd[0] * (lam - ctranslamstd[1]) / ctranslamstd[1]) *
                           np.exp(-1.0 * tau * _secZenith * (lam / self.lambdaNorm)**(-alpha)))

        return atmInterpolated

import numpy,import scipy,import astropy


# --------------------------------------------------------
# Tensorflow TIN
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------

"""
Generating training instance
"""

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

import numpy as np
import json
import pickle
import random
from random import randint
import tensorflow as tf
import cv2

#from config import cfg # cause error because of "from __future__ import absolute_import"
from ult import config # use absolute import and config.cfg

list_no_inter = [10,24,31,46,54,65,76,86,92,96,107,111,
129,146,160,170,174,186,194,198,208,214,224,232,235,239,
243,247,252,257,264,273,283,290,295,305,313,325,330,336,
342,348,352,356,363,368,376,383,389,393,397,407,414,418,
429,434,438,445,449,453,463,474,483,488,502,506,516,528,
533,538,546,550,558,562,567,576,584,588,595,600]

OBJECT_MASK = [
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 1],
    [0, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 1, 0, 0, 1, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1, 1],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 1, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 1, 1, 0],
    [1, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 1, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 1, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0]
]

print("*************data path:*************")
print(config.cfg.DATA_DIR)
print("************************************")

def bbox_trans(human_box_ori, object_box_ori, ratio, size = 64):

    human_box  = human_box_ori.copy()
    object_box = object_box_ori.copy()

    InteractionPattern = [min(human_box[0], object_box[0]), min(human_box[1], object_box[1]), max(human_box[2], object_box[2]), max(human_box[3], object_box[3])]

    height = InteractionPattern[3] - InteractionPattern[1] + 1
    width = InteractionPattern[2] - InteractionPattern[0] + 1

    if height > width:
        ratio = 'height'
    else:
        ratio = 'width'

    # shift the top-left corner to (0,0)

    human_box[0] -= InteractionPattern[0]
    human_box[2] -= InteractionPattern[0]
    human_box[1] -= InteractionPattern[1]
    human_box[3] -= InteractionPattern[1]
    object_box[0] -= InteractionPattern[0]
    object_box[2] -= InteractionPattern[0]
    object_box[1] -= InteractionPattern[1]
    object_box[3] -= InteractionPattern[1]

    if ratio == 'height': # height is larger than width

        human_box[0] = 0 + size * human_box[0] / height
        human_box[1] = 0 + size * human_box[1] / height
        human_box[2] = (size * width / height - 1) - size * (width  - 1 - human_box[2]) / height
        human_box[3] = (size - 1)                  - size * (height - 1 - human_box[3]) / height

        object_box[0] = 0 + size * object_box[0] / height
        object_box[1] = 0 + size * object_box[1] / height
        object_box[2] = (size * width / height - 1) - size * (width  - 1 - object_box[2]) / height
        object_box[3] = (size - 1)                  - size * (height - 1 - object_box[3]) / height

        # Need to shift horizontally
        InteractionPattern = [min(human_box[0], object_box[0]), min(human_box[1], object_box[1]), max(human_box[2], object_box[2]), max(human_box[3], object_box[3])]
        #assert (InteractionPattern[0] == 0) & (InteractionPattern[1] == 0) & (InteractionPattern[3] == 63) & (InteractionPattern[2] <= 63)
        if human_box[3] > object_box[3]:
            human_box[3] = size - 1
        else:
            object_box[3] = size - 1

        shift = size / 2 - (InteractionPattern[2] + 1) / 2
        human_box += [shift, 0 , shift, 0]
        object_box += [shift, 0 , shift, 0]

    else: # width is larger than height

        human_box[0] = 0 + size * human_box[0] / width
        human_box[1] = 0 + size * human_box[1] / width
        human_box[2] = (size - 1)                  - size * (width  - 1 - human_box[2]) / width
        human_box[3] = (size * height / width - 1) - size * (height - 1 - human_box[3]) / width

        object_box[0] = 0 + size * object_box[0] / width
        object_box[1] = 0 + size * object_box[1] / width
        object_box[2] = (size - 1)                  - size * (width  - 1 - object_box[2]) / width
        object_box[3] = (size * height / width - 1) - size * (height - 1 - object_box[3]) / width

        # Need to shift vertically
        InteractionPattern = [min(human_box[0], object_box[0]), min(human_box[1], object_box[1]), max(human_box[2], object_box[2]), max(human_box[3], object_box[3])]

        #assert (InteractionPattern[0] == 0) & (InteractionPattern[1] == 0) & (InteractionPattern[2] == 63) & (InteractionPattern[3] <= 63)

        if human_box[2] > object_box[2]:
            human_box[2] = size - 1
        else:
            object_box[2] = size - 1

        shift = size / 2 - (InteractionPattern[3] + 1) / 2

        human_box = human_box + [0, shift, 0 , shift]
        object_box = object_box + [0, shift, 0 , shift]

    return np.round(human_box), np.round(object_box)

def Get_next_sp(human_box, object_box):

    InteractionPattern = [min(human_box[0], object_box[0]), min(human_box[1], object_box[1]), max(human_box[2], object_box[2]), max(human_box[3], object_box[3])]
    height = InteractionPattern[3] - InteractionPattern[1] + 1
    width = InteractionPattern[2] - InteractionPattern[0] + 1
    if height > width:
        H, O = bbox_trans(human_box, object_box, 'height')
    else:
        H, O  = bbox_trans(human_box, object_box, 'width')

    Pattern = np.zeros((64,64,2))
    Pattern[int(H[1]):int(H[3]) + 1,int(H[0]):int(H[2]) + 1,0] = 1
    Pattern[int(O[1]):int(O[3]) + 1,int(O[0]):int(O[2]) + 1,1] = 1

    return Pattern

def bb_IOU(boxA, boxB):

    ixmin = np.maximum(boxA[0], boxB[0])
    iymin = np.maximum(boxA[1], boxB[1])
    ixmax = np.minimum(boxA[2], boxB[2])
    iymax = np.minimum(boxA[3], boxB[3])
    iw = np.maximum(ixmax - ixmin + 1., 0.)
    ih = np.maximum(iymax - iymin + 1., 0.)
    inters = iw * ih

    # union
    uni = ((boxB[2] - boxB[0] + 1.) * (boxB[3] - boxB[1] + 1.) +
           (boxA[2] - boxA[0] + 1.) *
           (boxA[3] - boxA[1] + 1.) - inters)

    overlaps = inters / uni
    return overlaps

def Augmented_box(bbox, shape, image_id, augment = 15, break_flag = True):

    thres_ = 0.7

    box = np.array([0, bbox[0],  bbox[1],  bbox[2],  bbox[3]]).reshape(1,5)
    box = box.astype(np.float64)

    count = 0
    time_count = 0
    while count < augment:

        time_count += 1
        height = bbox[3] - bbox[1]
        width  = bbox[2] - bbox[0]

        height_cen = (bbox[3] + bbox[1]) / 2
        width_cen  = (bbox[2] + bbox[0]) / 2

        ratio = 1 + randint(-10,10) * 0.01

        height_shift = randint(-np.floor(height),np.floor(height)) * 0.1
        width_shift  = randint(-np.floor(width),np.floor(width)) * 0.1

        H_0 = max(0, width_cen + width_shift - ratio * width / 2)
        H_2 = min(shape[1] - 1, width_cen + width_shift + ratio * width / 2)
        H_1 = max(0, height_cen + height_shift - ratio * height / 2)
        H_3 = min(shape[0] - 1, height_cen + height_shift + ratio * height / 2)

        if bb_IOU(bbox, np.array([H_0, H_1, H_2, H_3])) > thres_:
            box_ = np.array([0, H_0, H_1, H_2, H_3]).reshape(1,5)
            box  = np.concatenate((box,     box_),     axis=0)
            count += 1
        if break_flag == True and time_count > 150:
            return box

    return box

def Generate_action(action_list):
    action_ = np.zeros(29)
    for GT_idx in action_list:
        action_[GT_idx] = 1
    action_ = action_.reshape(1,29)
    return action_

def Generate_action_HICO(action_list):
    action_ = np.zeros(600)
    for GT_idx in action_list:
        action_[GT_idx] = 1
    action_ = action_.reshape(1,600)
    return action_

##############################################################################
##############################################################################
# tools added
##############################################################################
##############################################################################

def Generate_action_30(action_list):
    action_ = np.zeros(30)
    for GT_idx in action_list:
        action_[GT_idx] = 1
    action_ = action_.reshape(1,30)
    return action_

def draw_relation(human_pattern, joints, size = 64):

    joint_relation = [[1,3],[2,4],[0,1],[0,2],[0,17],[5,17],[6,17],[5,7],[6,8],[7,9],[8,10],[11,17],[12,17],[11,13],[12,14],[13,15],[14,16]]
    color = [0.15,0.2,0.25,0.3,0.35,0.4,0.45,0.5,0.55,0.6,0.65,0.7,0.75,0.8,0.85,0.9,0.95]
    skeleton = np.zeros((size, size, 1), dtype="float32")

    for i in range(len(joint_relation)):
        cv2.line(skeleton, tuple(joints[joint_relation[i][0]]), tuple(joints[joint_relation[i][1]]), (color[i]))

    # cv2.rectangle(skeleton, (int(human_pattern[0]), int(human_pattern[1])), (int(human_pattern[2]), int(human_pattern[3])), (255))
    # cv2.imshow("Joints", skeleton)
    # cv2.waitKey(0)
    # print(skeleton[:,:,0])

    return skeleton

def get_skeleton(human_box, human_pose, human_pattern, num_joints = 17, size = 64):
    width = human_box[2] - human_box[0] + 1
    height = human_box[3] - human_box[1] + 1
    pattern_width = human_pattern[2] - human_pattern[0] + 1
    pattern_height = human_pattern[3] - human_pattern[1] + 1
    joints = np.zeros((num_joints + 1, 2), dtype='int32')

    for i in range(num_joints):
        joint_x, joint_y, joint_score = human_pose[3 * i : 3 * (i + 1)]
        x_ratio = (joint_x - human_box[0]) / float(width)
        y_ratio = (joint_y - human_box[1]) / float(height)
        joints[i][0] = min(size - 1, int(round(x_ratio * pattern_width + human_pattern[0])))
        joints[i][1] = min(size - 1, int(round(y_ratio * pattern_height + human_pattern[1])))
    joints[num_joints] = (joints[5] + joints[6]) / 2

    return draw_relation(human_pattern, joints)

def Get_next_sp_with_pose(human_box, object_box, human_pose, num_joints=17):

    InteractionPattern = [min(human_box[0], object_box[0]), min(human_box[1], object_box[1]), max(human_box[2], object_box[2]), max(human_box[3], object_box[3])]
    height = InteractionPattern[3] - InteractionPattern[1] + 1
    width = InteractionPattern[2] - InteractionPattern[0] + 1
    if height > width:
        H, O = bbox_trans(human_box, object_box, 'height')
    else:
        H, O  = bbox_trans(human_box, object_box, 'width')

    Pattern = np.zeros((64,64,2), dtype='float32')
    Pattern[int(H[1]):int(H[3]) + 1,int(H[0]):int(H[2]) + 1,0] = 1
    Pattern[int(O[1]):int(O[3]) + 1,int(O[0]):int(O[2]) + 1,1] = 1

    if human_pose != None and len(human_pose) == 51:
        skeleton = get_skeleton(human_box, human_pose, H, num_joints)
    else:
        skeleton = np.zeros((64,64,1), dtype='float32')
        skeleton[int(H[1]):int(H[3]) + 1,int(H[0]):int(H[2]) + 1,0] = 0.05

    Pattern = np.concatenate((Pattern, skeleton), axis=2)

    return Pattern

##############################################################################
# for vcoco with pose pattern
##############################################################################

def Get_Next_Instance_HO_Neg_pose_pattern_version2(trainval_GT, Trainval_Neg, iter, Pos_augment, Neg_select, Data_length):

    GT       = trainval_GT[iter%Data_length]
    image_id = GT[0][0]
    im_file  = config.cfg.DATA_DIR + '/' + 'v-coco/coco/images/train2014/COCO_train2014_' + (str(image_id)).zfill(12) + '.jpg'
    im       = cv2.imread(im_file)
    im_orig  = im.astype(np.float32, copy=True)
    im_orig -= config.cfg.PIXEL_MEANS
    im_shape = im_orig.shape
    im_orig  = im_orig.reshape(1, im_shape[0], im_shape[1], 3)

    Pattern, Human_augmented, Human_augmented_solo, Object_augmented, action_HO, action_H, mask_HO, mask_H, binary_label = Augmented_HO_Neg_pose_pattern_version2(GT, Trainval_Neg, im_shape, Pos_augment, Neg_select)

    blobs = {}
    blobs['image']       = im_orig
    blobs['H_boxes_solo']= Human_augmented_solo
    blobs['H_boxes']     = Human_augmented
    blobs['O_boxes']     = Object_augmented
    blobs['gt_class_HO'] = action_HO
    blobs['gt_class_H']  = action_H
    blobs['Mask_HO']     = mask_HO
    blobs['Mask_H']      = mask_H
    blobs['sp']          = Pattern
    blobs['H_num']       = len(action_H)
    blobs['binary_label'] = binary_label

    return blobs

def Augmented_HO_Neg_pose_pattern_version2(GT, Trainval_Neg, shape, Pos_augment, Neg_select):
    image_id = GT[0][0]
    GT_count = len(GT)
    aug_all = int(Pos_augment / GT_count)
    aug_last = Pos_augment - aug_all * (GT_count - 1)

    Human_augmented, Object_augmented, action_HO, Human_augmented_solo, action_H = [], [], [], [], []
    Pattern   = np.empty((0, 64, 64, 3), dtype=np.float32)

    for i in range(GT_count - 1):
        Human    = GT[i][2]
        Object   = GT[i][3]

        Human_augmented_temp  = Augmented_box(Human,  shape, image_id, aug_all)
        Object_augmented_temp = Augmented_box(Object, shape, image_id, aug_all)
        Human_augmented_solo.extend(Human_augmented_temp)

        length_min = min(len(Human_augmented_temp),len(Object_augmented_temp))

        Human_augmented_temp  =  Human_augmented_temp[:length_min]
        Object_augmented_temp = Object_augmented_temp[:length_min]

        action_H__temp = Generate_action(GT[i][4])
        action_H_temp  = action_H__temp
        for j in range(length_min - 1):
            action_H_temp = np.concatenate((action_H_temp, action_H__temp), axis=0)

        action_HO__temp = Generate_action(GT[i][1])
        action_HO_temp  = action_HO__temp
        for j in range(length_min - 1):
            action_HO_temp = np.concatenate((action_HO_temp, action_HO__temp), axis=0)

        for j in range(length_min):
            Pattern_ = Get_next_sp_with_pose(Human_augmented_temp[j][1:], Object_augmented_temp[j][1:], GT[i][5]).reshape(1, 64, 64, 3)
            Pattern  = np.concatenate((Pattern, Pattern_), axis=0)

        Human_augmented.extend(Human_augmented_temp)
        Object_augmented.extend(Object_augmented_temp)
        action_HO.extend(action_HO_temp)
        action_H.extend(action_H_temp)

    Human    = GT[GT_count - 1][2]
    Object   = GT[GT_count - 1][3]

    Human_augmented_temp  = Augmented_box(Human,  shape, image_id, aug_last)
    Object_augmented_temp = Augmented_box(Object, shape, image_id, aug_last)
    Human_augmented_solo.extend(Human_augmented_temp)

    length_min = min(len(Human_augmented_temp),len(Object_augmented_temp))

    Human_augmented_temp  =  Human_augmented_temp[:length_min]
    Object_augmented_temp = Object_augmented_temp[:length_min]

    action_H__temp = Generate_action(GT[GT_count - 1][4])
    action_H_temp  = action_H__temp
    for j in range(length_min - 1):
        action_H_temp = np.concatenate((action_H_temp, action_H__temp), axis=0)

    action_HO__temp = Generate_action(GT[GT_count - 1][1])
    action_HO_temp  = action_HO__temp
    for j in range(length_min - 1):
        action_HO_temp = np.concatenate((action_HO_temp, action_HO__temp), axis=0)

    for j in range(length_min):
        Pattern_ = Get_next_sp_with_pose(Human_augmented_temp[j][1:], Object_augmented_temp[j][1:], GT[GT_count - 1][5]).reshape(1, 64, 64, 3)
        Pattern  = np.concatenate((Pattern, Pattern_), axis=0)

    Human_augmented.extend(Human_augmented_temp)
    Object_augmented.extend(Object_augmented_temp)
    action_HO.extend(action_HO_temp)
    action_H.extend(action_H_temp)

    num_pos = len(Human_augmented)

    if image_id in Trainval_Neg:

        if len(Trainval_Neg[image_id]) < Neg_select:
            for Neg in Trainval_Neg[image_id]:
                Human_augmented  = np.concatenate((Human_augmented,  np.array([0, Neg[2][0], Neg[2][1], Neg[2][2], Neg[2][3]]).reshape(1,5)), axis=0)
                Object_augmented = np.concatenate((Object_augmented, np.array([0, Neg[3][0], Neg[3][1], Neg[3][2], Neg[3][3]]).reshape(1,5)), axis=0)
                Pattern_ = Get_next_sp_with_pose(np.array(Neg[2], dtype='float64'), np.array(Neg[3], dtype='float64'), Neg[7]).reshape(1, 64, 64, 3)
                Pattern  = np.concatenate((Pattern, Pattern_), axis=0)
        else:
            List = random.sample(range(len(Trainval_Neg[image_id])), len(Trainval_Neg[image_id]))
            for i in range(Neg_select):
                Neg = Trainval_Neg[image_id][List[i]]
                Human_augmented  = np.concatenate((Human_augmented,  np.array([0, Neg[2][0], Neg[2][1], Neg[2][2], Neg[2][3]]).reshape(1,5)), axis=0)
                Object_augmented = np.concatenate((Object_augmented, np.array([0, Neg[3][0], Neg[3][1], Neg[3][2], Neg[3][3]]).reshape(1,5)), axis=0)
                Pattern_ = Get_next_sp_with_pose(np.array(Neg[2], dtype='float64'), np.array(Neg[3], dtype='float64'), Neg[7]).reshape(1, 64, 64, 3)
                Pattern  = np.concatenate((Pattern, Pattern_), axis=0)

    num_pos_neg = len(Human_augmented)

    mask_HO_   = np.asarray([1,1,1,0,1,1,1,1,1,1,1,1,1,1,1,1,1,0,1,1,1,1,0,1,1,1,1,0,1]).reshape(1,29) # some actions do not have objects
    mask_H_    = np.asarray([1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]).reshape(1,29)
    mask_HO   = mask_HO_
    mask_H    = mask_H_

    for i in range(num_pos - 1):
        mask_H    = np.concatenate((mask_H,    mask_H_),    axis=0)

    for i in range(num_pos_neg - 1):
        mask_HO   = np.concatenate((mask_HO,   mask_HO_),   axis=0)

    for i in range(num_pos_neg - num_pos):
        action_HO = np.concatenate((action_HO, np.zeros(29).reshape(1,29)), axis=0)

    Human_augmented   = np.array(Human_augmented, dtype = 'float32')
    Human_augmented_solo = np.array(Human_augmented_solo, dtype = 'float32')
    Object_augmented  = np.array(Object_augmented, dtype = 'float32')
    action_HO         = np.array(action_HO, dtype = 'int32')
    action_H          = np.array(action_H, dtype = 'int32')

    Pattern           = Pattern.reshape( num_pos_neg, 64, 64, 3)
    Human_augmented   = Human_augmented.reshape( num_pos_neg, 5)
    Human_augmented_solo = Human_augmented_solo.reshape( -1, 5)
    Object_augmented  = Object_augmented.reshape(num_pos_neg, 5)
    action_HO         = action_HO.reshape(num_pos_neg, 29)
    action_H          = action_H.reshape( num_pos, 29)
    mask_HO           = mask_HO.reshape(  num_pos_neg, 29)
    mask_H            = mask_H.reshape(   num_pos, 29)

    # binary label for vcoco, different from hico, vcoco does not have the no-interaction label, so give it binary label accord to the locations of pos and neg directly
    binary_label = np.zeros((num_pos_neg, 2), dtype = 'int32')
    for i in range(num_pos):
        binary_label[i][0] = 1 # pos is at 0
    for i in range(num_pos, num_pos_neg):
        binary_label[i][1] = 1 # neg is at 1

    return Pattern, Human_augmented, Human_augmented_solo, Object_augmented, action_HO, action_H, mask_HO, mask_H, binary_label

def Get_Next_Instance_HO_spNeg_pose_pattern_version2(trainval_GT, Trainval_Neg, iter, Pos_augment, Neg_select, Data_length):

    GT       = trainval_GT[iter%Data_length]
    image_id = GT[0][0]
    im_file  = config.cfg.DATA_DIR + '/' + 'v-coco/coco/images/train2014/COCO_train2014_' + (str(image_id)).zfill(12) + '.jpg'
    im       = cv2.imread(im_file)
    im_orig  = im.astype(np.float32, copy=True)
    im_orig -= config.cfg.PIXEL_MEANS
    im_shape = im_orig.shape
    im_orig  = im_orig.reshape(1, im_shape[0], im_shape[1], 3)

    Pattern, Human_augmented_sp, Human_augmented, Object_augmented, action_sp, action_HO, action_H, mask_sp, mask_HO, mask_H, binary_label = Augmented_HO_spNeg_pose_pattern_version2(GT, Trainval_Neg, im_shape, Pos_augment, Neg_select)

    blobs = {}
    blobs['image']       = im_orig
    blobs['H_boxes']     = Human_augmented
    blobs['Hsp_boxes']   = Human_augmented_sp
    blobs['O_boxes']     = Object_augmented
    blobs['gt_class_sp'] = action_sp
    blobs['gt_class_HO'] = action_HO
    blobs['gt_class_H']  = action_H
    blobs['Mask_sp']     = mask_sp
    blobs['Mask_HO']     = mask_HO
    blobs['Mask_H']      = mask_H
    blobs['sp']          = Pattern
    blobs['H_num']       = len(action_H)
    blobs['binary_label'] = binary_label

    return blobs

def Augmented_HO_spNeg_pose_pattern_version2(GT, Trainval_Neg, shape, Pos_augment, Neg_select):
    image_id = GT[0][0]
    GT_count = len(GT)
    aug_all = int(Pos_augment / GT_count)
    aug_last = Pos_augment - aug_all * (GT_count - 1)

    Human_augmented, Object_augmented, action_HO, action_H = [], [], [], []
    Pattern   = np.empty((0, 64, 64, 3), dtype=np.float32)

    for i in range(GT_count - 1):
        Human    = GT[i][2]
        Object   = GT[i][3]

        Human_augmented_temp  = Augmented_box(Human,  shape, image_id, aug_all)
        Object_augmented_temp = Augmented_box(Object, shape, image_id, aug_all)

        length_min = min(len(Human_augmented_temp),len(Object_augmented_temp))

        Human_augmented_temp  =  Human_augmented_temp[:length_min]
        Object_augmented_temp = Object_augmented_temp[:length_min]

        action_H__temp = Generate_action(GT[i][4])
        action_H_temp  = action_H__temp
        for j in range(length_min - 1):
            action_H_temp = np.concatenate((action_H_temp, action_H__temp), axis=0)

        action_HO__temp = Generate_action(GT[i][1])
        action_HO_temp  = action_HO__temp
        for j in range(length_min - 1):
            action_HO_temp = np.concatenate((action_HO_temp, action_HO__temp), axis=0)

        for j in range(length_min):
            Pattern_ = Get_next_sp_with_pose(Human_augmented_temp[j][1:], Object_augmented_temp[j][1:], GT[i][5]).reshape(1, 64, 64, 3)
            Pattern  = np.concatenate((Pattern, Pattern_), axis=0)

        Human_augmented.extend(Human_augmented_temp)
        Object_augmented.extend(Object_augmented_temp)
        action_HO.extend(action_HO_temp)
        action_H.extend(action_H_temp)

    Human    = GT[GT_count - 1][2]
    Object   = GT[GT_count - 1][3]

    Human_augmented_temp  = Augmented_box(Human,  shape, image_id, aug_last)
    Object_augmented_temp = Augmented_box(Object, shape, image_id, aug_last)

    length_min = min(len(Human_augmented_temp),len(Object_augmented_temp))

    Human_augmented_temp  =  Human_augmented_temp[:length_min]
    Object_augmented_temp = Object_augmented_temp[:length_min]

    action_H__temp = Generate_action(GT[GT_count - 1][4])
    action_H_temp  = action_H__temp
    for j in range(length_min - 1):
        action_H_temp = np.concatenate((action_H_temp, action_H__temp), axis=0)

    action_HO__temp = Generate_action(GT[GT_count - 1][1])
    action_HO_temp  = action_HO__temp
    for j in range(length_min - 1):
        action_HO_temp = np.concatenate((action_HO_temp, action_HO__temp), axis=0)

    for j in range(length_min):
        Pattern_ = Get_next_sp_with_pose(Human_augmented_temp[j][1:], Object_augmented_temp[j][1:], GT[GT_count - 1][5]).reshape(1, 64, 64, 3)
        Pattern  = np.concatenate((Pattern, Pattern_), axis=0)

    Human_augmented.extend(Human_augmented_temp)
    Object_augmented.extend(Object_augmented_temp)
    action_HO.extend(action_HO_temp)
    action_H.extend(action_H_temp)

    action_sp = np.array(action_HO).copy()

    num_pos = len(Human_augmented)

    if image_id in Trainval_Neg:

        if len(Trainval_Neg[image_id]) < Neg_select:
            for Neg in Trainval_Neg[image_id]:
                Human_augmented  = np.concatenate((Human_augmented,  np.array([0, Neg[2][0], Neg[2][1], Neg[2][2], Neg[2][3]]).reshape(1,5)), axis=0)
                Object_augmented = np.concatenate((Object_augmented, np.array([0, Neg[3][0], Neg[3][1], Neg[3][2], Neg[3][3]]).reshape(1,5)), axis=0)
                Pattern_ = Get_next_sp_with_pose(np.array(Neg[2], dtype='float64'), np.array(Neg[3], dtype='float64'), Neg[7]).reshape(1, 64, 64, 3)
                Pattern  = np.concatenate((Pattern, Pattern_), axis=0)
        else:
            List = random.sample(range(len(Trainval_Neg[image_id])), len(Trainval_Neg[image_id]))
            for i in range(Neg_select):
                Neg = Trainval_Neg[image_id][List[i]]
                Human_augmented  = np.concatenate((Human_augmented,  np.array([0, Neg[2][0], Neg[2][1], Neg[2][2], Neg[2][3]]).reshape(1,5)), axis=0)
                Object_augmented = np.concatenate((Object_augmented, np.array([0, Neg[3][0], Neg[3][1], Neg[3][2], Neg[3][3]]).reshape(1,5)), axis=0)
                Pattern_ = Get_next_sp_with_pose(np.array(Neg[2], dtype='float64'), np.array(Neg[3], dtype='float64'), Neg[7]).reshape(1, 64, 64, 3)
                Pattern  = np.concatenate((Pattern, Pattern_), axis=0)

    num_pos_neg = len(Human_augmented)

    mask_sp_   = np.asarray([1,1,1,0,1,1,1,1,1,1,1,1,1,1,1,1,1,0,1,1,1,1,0,1,1,1,1,0,1]).reshape(1,29)
    mask_HO_   = np.asarray([1,1,1,0,1,1,1,1,1,1,1,1,1,1,1,1,1,0,1,1,1,1,0,1,1,1,1,0,1]).reshape(1,29)
    mask_H_    = np.asarray([1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]).reshape(1,29)
    mask_sp   = mask_sp_
    mask_HO   = mask_HO_
    mask_H    = mask_H_

    for i in range(num_pos - 1):
        mask_HO   = np.concatenate((mask_HO,   mask_HO_),   axis=0)
        mask_H    = np.concatenate((mask_H,    mask_H_),    axis=0)

    for i in range(num_pos_neg - 1):
        mask_sp   = np.concatenate((mask_sp,   mask_sp_),   axis=0)

    for i in range(num_pos_neg - num_pos):
        action_sp = np.concatenate((action_sp, np.zeros(29).reshape(1,29)), axis=0) # neg has all 0 in 29 label

    Human_augmented   = np.array(Human_augmented, dtype = 'float32')
    Object_augmented  = np.array(Object_augmented, dtype = 'float32')
    action_HO         = np.array(action_HO, dtype = 'int32')
    action_sp         = np.array(action_sp, dtype = 'int32')
    action_H          = np.array(action_H, dtype = 'int32')

    Pattern           = Pattern.reshape( num_pos_neg, 64, 64, 3)
    Human_augmented_sp= Human_augmented.reshape( num_pos_neg, 5)
    Human_augmented   = Human_augmented[:num_pos].reshape( num_pos, 5)
    Object_augmented  = Object_augmented[:num_pos].reshape(num_pos, 5)
    action_sp         = action_sp.reshape(num_pos_neg, 29)
    action_HO         = action_HO.reshape(num_pos, 29)
    action_H          = action_H.reshape( num_pos, 29)
    mask_sp           = mask_sp.reshape(num_pos_neg, 29)
    mask_HO           = mask_HO.reshape(  num_pos, 29)
    mask_H            = mask_H.reshape(   num_pos, 29)

    binary_label = np.zeros((num_pos_neg, 2), dtype = 'int32')
    for i in range(num_pos):
        binary_label[i][0] = 1 # pos is at 0
    for i in range(num_pos, num_pos_neg):
        binary_label[i][1] = 1 # neg is at 1

    return Pattern, Human_augmented_sp, Human_augmented, Object_augmented, action_sp, action_HO, action_H, mask_sp, mask_HO, mask_H, binary_label

##############################################################################
# for hico with pose map
##############################################################################

def Get_Next_Instance_HO_Neg_HICO_pose_pattern_version2(trainval_GT, Trainval_Neg, iter, Pos_augment, Neg_select, Data_length):

    GT       = trainval_GT[iter%Data_length]
    image_id = GT[0][0]
    im_file = config.cfg.DATA_DIR + '/' + 'hico_20160224_det/images/train2015/HICO_train2015_' + (str(image_id)).zfill(8) + '.jpg'
    im       = cv2.imread(im_file)
    im_orig  = im.astype(np.float32, copy=True)
    im_orig -= config.cfg.PIXEL_MEANS
    im_shape = im_orig.shape
    im_orig  = im_orig.reshape(1, im_shape[0], im_shape[1], 3)

    Pattern, Human_augmented, Object_augmented, action_HO, num_pos, binary_label = Augmented_HO_Neg_HICO_pose_pattern_version2(GT, Trainval_Neg, im_shape, Pos_augment, Neg_select)

    blobs = {}
    blobs['image']       = im_orig
    blobs['H_boxes']     = Human_augmented
    blobs['O_boxes']     = Object_augmented
    blobs['gt_class_HO'] = action_HO
    blobs['sp']          = Pattern
    blobs['H_num']       = num_pos
    blobs['binary_label'] = binary_label

    return blobs

def Augmented_HO_Neg_HICO_pose_pattern_version2(GT, Trainval_Neg, shape, Pos_augment, Neg_select):
    GT_count = len(GT)
    aug_all = int(Pos_augment / GT_count)
    aug_last = Pos_augment - aug_all * (GT_count - 1)
    image_id = GT[0][0]

    Human_augmented, Object_augmented, action_HO = [], [], []
    Pattern   = np.empty((0, 64, 64, 3), dtype=np.float32)

    for i in range(GT_count - 1):
        Human    = GT[i][2]
        Object   = GT[i][3]

        Human_augmented_temp  = Augmented_box(Human,  shape, image_id, aug_all)
        Object_augmented_temp = Augmented_box(Object, shape, image_id, aug_all)

        length_min = min(len(Human_augmented_temp),len(Object_augmented_temp))

        Human_augmented_temp  =  Human_augmented_temp[:length_min]
        Object_augmented_temp = Object_augmented_temp[:length_min]

        action_HO__temp = Generate_action_HICO(GT[i][1])
        action_HO_temp  = action_HO__temp
        for j in range(length_min - 1):
            action_HO_temp = np.concatenate((action_HO_temp, action_HO__temp), axis=0)

        for j in range(length_min):
            Pattern_ = Get_next_sp_with_pose(Human_augmented_temp[j][1:], Object_augmented_temp[j][1:], GT[i][5]).reshape(1, 64, 64, 3)
            Pattern  = np.concatenate((Pattern, Pattern_), axis=0)

        Human_augmented.extend(Human_augmented_temp)
        Object_augmented.extend(Object_augmented_temp)
        action_HO.extend(action_HO_temp)

    Human    = GT[GT_count - 1][2]
    Object   = GT[GT_count - 1][3]

    Human_augmented_temp  = Augmented_box(Human,  shape, image_id, aug_last)
    Object_augmented_temp = Augmented_box(Object, shape, image_id, aug_last)

    length_min = min(len(Human_augmented_temp),len(Object_augmented_temp))

    Human_augmented_temp  =  Human_augmented_temp[:length_min]
    Object_augmented_temp = Object_augmented_temp[:length_min]

    action_HO__temp = Generate_action_HICO(GT[GT_count - 1][1])
    action_HO_temp  = action_HO__temp
    for j in range(length_min - 1):
        action_HO_temp = np.concatenate((action_HO_temp, action_HO__temp), axis=0)

    for j in range(length_min):
        Pattern_ = Get_next_sp_with_pose(Human_augmented_temp[j][1:], Object_augmented_temp[j][1:], GT[GT_count - 1][5]).reshape(1, 64, 64, 3)
        Pattern  = np.concatenate((Pattern, Pattern_), axis=0)

    Human_augmented.extend(Human_augmented_temp)
    Object_augmented.extend(Object_augmented_temp)
    action_HO.extend(action_HO_temp)

    num_pos = len(Human_augmented)

    if image_id in Trainval_Neg:

        if len(Trainval_Neg[image_id]) < Neg_select:
            for Neg in Trainval_Neg[image_id]:

                Human_augmented  = np.concatenate((Human_augmented,  np.array([0, Neg[2][0], Neg[2][1], Neg[2][2], Neg[2][3]]).reshape(1,5)), axis=0)
                Object_augmented = np.concatenate((Object_augmented, np.array([0, Neg[3][0], Neg[3][1], Neg[3][2], Neg[3][3]]).reshape(1,5)), axis=0)
                action_HO        = np.concatenate((action_HO, Generate_action_HICO([Neg[1]])), axis=0)
                Pattern_ = Get_next_sp_with_pose(np.array(Neg[2], dtype='float64'), np.array(Neg[3], dtype='float64'), Neg[7]).reshape(1, 64, 64, 3)
                Pattern  = np.concatenate((Pattern, Pattern_), axis=0)
        else:
            List = random.sample(range(len(Trainval_Neg[image_id])), len(Trainval_Neg[image_id]))
            for i in range(Neg_select):
                Neg = Trainval_Neg[image_id][List[i]]

                Human_augmented  = np.concatenate((Human_augmented,  np.array([0, Neg[2][0], Neg[2][1], Neg[2][2], Neg[2][3]]).reshape(1,5)), axis=0)
                Object_augmented = np.concatenate((Object_augmented, np.array([0, Neg[3][0], Neg[3][1], Neg[3][2], Neg[3][3]]).reshape(1,5)), axis=0)
                action_HO        = np.concatenate((action_HO, Generate_action_HICO([Neg[1]])), axis=0)
                Pattern_ = Get_next_sp_with_pose(np.array(Neg[2], dtype='float64'), np.array(Neg[3], dtype='float64'), Neg[7]).reshape(1, 64, 64, 3)
                Pattern  = np.concatenate((Pattern, Pattern_), axis=0)

    num_pos_neg = len(Human_augmented)

    Pattern           = Pattern.reshape(num_pos_neg, 64, 64, 3)
    Human_augmented   = np.array(Human_augmented, dtype='float32').reshape(num_pos_neg, 5)
    Object_augmented  = np.array(Object_augmented, dtype='float32').reshape(num_pos_neg, 5)
    action_HO         = np.array(action_HO, dtype='int32').reshape(num_pos_neg, 600)

    binary_label = np.zeros((num_pos_neg, 2), dtype = 'int32')
    for i in range(num_pos):
        # if hoi_id belong to 80 then is 0 1, else 1 0
        binary_label[i][0] = 1 # pos is at 0 #############problem: here in pos, it contains the no interaction 80, so we may give this 80 '0 1' and 520 '1 0', and test
    for i in range(num_pos, num_pos_neg):
        binary_label[i][1] = 1 # neg is at 1

    return Pattern, Human_augmented, Object_augmented, action_HO, num_pos, binary_label


# -*- coding: utf-8 -*-
from __future__ import division
import os
import plotly
from plotly.graph_objs import Scatter
from plotly.graph_objs.scatter import Line
import torch
from tqdm import tqdm, trange
import numpy as np

import atari_py
from env import Env

import argparse
from agent import Agent

def get_args():
  parser = argparse.ArgumentParser(description='Rainbow')
  parser.add_argument('--model-dir', type=str, default='results', help='Experiment ID')
  parser.add_argument('--seed', type=int, default=1, help='Random seed')
  parser.add_argument('--game', default='battle_zone', choices=atari_py.list_games())
  parser.add_argument('--start-ckpt', default=100000, type=int)
  parser.add_argument('--env-type', default='atari', choices=['atari', 'sepsis', 'hiv', 'gym'])
  parser.add_argument('--deploy-policy', default=None, choices=['fixed', 'exp', 'dqn-feature', "reset_feature", 'q-value', 'dqn-feature-min',
                                                                'reset', 'policy', 'policy_adapt', 'policy_diverge', 'reset_policy', 'visited', "info-matrix",
                                                                'reset_feature_force', 'feature_lowf'])
  parser.add_argument("--info-matrix-interval", default=10, type=int)
  parser.add_argument("--info-matrix-ratio", default=2, type=float)
  parser.add_argument("--use-gradient-weight", default=False, action="store_true")
  parser.add_argument("--adaptive-softmax", default=False, action="store_true")
  parser.add_argument('--record-action-diff', default=False, action="store_true")
  parser.add_argument('--record-feature-sim', default=False, action="store_true")
  parser.add_argument('--switch-memory-priority', default=True, type=eval)
  parser.add_argument('--switch-bsz', default=32, type=int)
  parser.add_argument('--switch-sample-strategy', default=None, choices=['uniform', 'recent'], type=str, help="only useful when switch-memory-priority is False")
  parser.add_argument('--switch-memory-capcacity', default=1000000, type=int)
  parser.add_argument('--diverge-threshold', default=0.2, type=float)
  parser.add_argument('--delploy-interval', default=1, type=int)
  parser.add_argument('--force-interval', default=1000, type=int, help='For reset_feature_force strategy')
  parser.add_argument('--max-episode-length', type=int, default=int(108e3), metavar='LENGTH', help='Max episode length in game frames (0 to disable)')
  parser.add_argument('--history-length', type=int, default=4, metavar='T', help='Number of consecutive states processed')
  parser.add_argument('--architecture', type=str, default='canonical', choices=['canonical', 'data-efficient'], metavar='ARCH', help='Network architecture')
  parser.add_argument('--hidden-size', type=int, default=512, metavar='SIZE', help='Network hidden size')
  parser.add_argument('--noisy-std', type=float, default=0.1, metavar='σ', help='Initial standard deviation of noisy linear layers')
  parser.add_argument('--atoms', type=int, default=51, metavar='C', help='Discretised size of value distribution')
  parser.add_argument('--V-min', type=float, default=-10, metavar='V', help='Minimum of value distribution support')
  parser.add_argument('--V-max', type=float, default=10, metavar='V', help='Maximum of value distribution support')
  parser.add_argument('--model', type=str, metavar='PARAMS', help='Pretrained model (state dict)')
  parser.add_argument('--memory-capacity', type=int, default=int(1e6), metavar='CAPACITY', help='Experience replay memory capacity')
  parser.add_argument("--gradient-steps", default=1, type=int)
  parser.add_argument('--replay-frequency', type=int, default=4, metavar='k', help='Frequency of sampling from memory')
  parser.add_argument('--priority-exponent', type=float, default=0.5, metavar='ω', help='Prioritised experience replay exponent (originally denoted α)')
  parser.add_argument('--priority-weight', type=float, default=0.4, metavar='β', help='Initial prioritised experience replay importance sampling weight')
  parser.add_argument('--multi-step', type=int, default=3, metavar='n', help='Number of steps for multi-step return')
  parser.add_argument('--discount', type=float, default=0.99, metavar='γ', help='Discount factor')
  # parser.add_argument('--target-update', type=int, default=int(8e3), metavar='τ', help='Number of steps after which to update target network')
  parser.add_argument('--reward-clip', type=int, default=1, metavar='VALUE', help='Reward clipping (0 to disable)')
  parser.add_argument('--learning-rate', type=float, default=0.0000625, metavar='η', help='Learning rate')
  parser.add_argument('--adam-eps', type=float, default=1.5e-4, metavar='ε', help='Adam epsilon')
  parser.add_argument('--batch-size', type=int, default=32, metavar='SIZE', help='Batch size')
  parser.add_argument('--norm-clip', type=float, default=10, metavar='NORM', help='Max L2 norm for gradient clipping')
  # parser.add_argument('--learn-start', type=int, default=int(20e3), metavar='STEPS', help='Number of steps before starting training')
  # parser.add_argument('--evaluate', action='store_true', help='Evaluate only')
  parser.add_argument('--evaluate-dir', default=None, type=str)
  parser.add_argument('--evaluation-interval', type=int, default=100000, metavar='STEPS', help='Number of training steps between evaluations')
  parser.add_argument('--evaluation-episodes', type=int, default=1, metavar='N', help='Number of evaluation episodes to average over')
  # parser.add_argument('--state-visitation-episodes', type=int, default=10, help='Number of evaluation episodes to measure state visitation')
  # parser.add_argument('--result-dir', default='results/', type=str)
  # TODO: Note that DeepMind's evaluation method is running the latest agent for 500K frames ever every 1M steps
  # parser.add_argument('--evaluation-size', type=int, default=500, metavar='N', help='Number of transitions to use for validating Q')
  parser.add_argument('--render', action='store_true', help='Display screen (testing only)')
  # parser.add_argument('--enable-cudnn', action='store_true', help='Enable cuDNN (faster but nondeterministic)')
  # parser.add_argument('--checkpoint-interval', default=0, type=int, help='How often to checkpoint the model, defaults to 0 (never checkpoint)')
  # parser.add_argument('--memory', help='Path to save/load the memory from')
  parser.add_argument('--disable-bzip-memory', action='store_true', help='Don\'t zip the memory file. Not recommended (zipping is a bit slower and much, much smaller)')

  return parser.parse_args()

# Test DQN

def test(args, T, dqn, metrics, env_class):
  env = env_class(args, training=False)
  metrics['steps'].append(T)

  T_rewards = []

  # Test performance over several episodes
  done = True
  for _ in range(args.evaluation_episodes):
    state, reward_sum, done = env.reset(), 0, False
    for step in range(args.max_episode_length):
      # if args.explore_eps is None:
      action = dqn.act(state)  # Choose an action greedily (with noisy weights)
      # else:
      #   action = dqn.act_e_greedy(state)
      state, reward, done, _ = env.step(action)  # Step
      reward_sum += reward
      if args.render:
        env.render()
      if getattr(env, 'life_termination', None) and args.game == "breakout":
        env.ale.act(1)
        env.life_termination = False
      if done:
        T_rewards.append(reward_sum)
        break
  env.close()

  avg_reward = sum(T_rewards) / len(T_rewards)

  metrics['rewards'].append(T_rewards)

  assert len(metrics['steps']) == len(metrics['rewards']) == len(metrics['nums_deploy'])

  return avg_reward

if __name__ == "__main__":
  args = get_args()
  args.device = torch.device('cpu')
  args.env_type = 'atari'
  print(args.evaluation_interval)
  for game in sorted(atari_py.list_games()):
    print("*" * 50, game, "*" * 50)
    if not os.path.exists(os.path.join(args.model_dir, game)) or game == "montezuma_revenge":
      continue
    for method in ['none', 'fix_1000', 'feature0.98', 'feature0.99', 'info', 'visited', 'reset_feature_force', 'reset_feature_force10k', 'policy', "policy0.25"]:
      for seed in [1, 2, 3]:
        model_dir = f'{args.model_dir}/{game}/{method}.{seed}'

        if os.path.exists(os.path.join(model_dir, 'new_metrics_{}.pth'.format(args.evaluation_interval) if args.evaluation_interval != 100000 else 'new_metrics.pth' )):
          check_metrics = torch.load(os.path.join(model_dir, 'new_metrics_{}.pth'.format(args.evaluation_interval) if args.evaluation_interval != 100000 else 'new_metrics.pth'))
          if check_metrics['steps'][-1] >= 3e6:
            print('already exist', game, method, seed)
            continue
        args.game = game
        env = Env(args)
        dqn = Agent(args, env)
        try:
          old_metrics = torch.load(os.path.join(model_dir, 'metrics.pth'))
          if old_metrics['steps'][-1] <= 3e6:
            continue
        except:
          continue
        metrics = {'steps': [], 'rewards': [], 'nums_deploy': []}

        for ckpt in tqdm([20000] + list(range(args.start_ckpt, int(3.1e6), args.evaluation_interval)), ncols=50):
          idx = old_metrics['steps'].index(ckpt)
          metrics['nums_deploy'].append(old_metrics['nums_deploy'][idx])

          model_path = os.path.join(model_dir, f'checkpoint_{ckpt}.pth')
          try:
            dqn.load(model_path)
          except FileNotFoundError:
            continue
          test(args, ckpt, dqn, metrics, Env)

          tqdm.write('game : {}, method: {}, seed, {}, Step: {:.2f}, Rewards: {:.2f}, Num_deploy: {}'.format(game, method, seed, metrics['steps'][-1], np.mean(metrics['rewards'][-1]), metrics['nums_deploy'][-1]))
          torch.save(metrics, os.path.join(model_dir, 'new_metrics_{}.pth'.format(args.evaluation_interval) if args.evaluation_interval != 100000 else 'new_metrics.pth'))


import numpy as np
import colorsys
import matplotlib
#matplotlib.use('Agg') ## this calls matplotlib without a X-windows GUI

# vectorize these operations for convenience below
rgb_to_hls_v = np.vectorize(colorsys.rgb_to_hls)
hls_to_rgb_v = np.vectorize(colorsys.hls_to_rgb)

def rgb_to_hls(r,g,b):
    return rgb_to_hls_v(r,g,b);

def hls_to_rgb(h,l,s):
    return hls_to_rgb_v(h,l,s);

def load_my_custom_color_tables():
    fna='heat_red'
    cdict_tmp={\
    'red':   ((0., 0.0416, 0.0416),(0.365079, 1.000000, 1.000000),(1.0, 1.0, 1.0)),\
    'green': ((0., 0., 0.),(0.71, 0.000000, 0.000000),\
            (1., 1.000000, 1.000000),(1.0, 1.0, 1.0)),\
    'blue':  ((0., 0., 0.),(0.71, 0.000000, 0.000000),(1.0, 1.0, 1.0))}
    my_cmap = matplotlib.colors.LinearSegmentedColormap(fna,cdict_tmp,256);
    matplotlib.cm.register_cmap(name=fna,cmap=my_cmap);

    fna='heat_blue'
    cdict_tmp={\
    'blue':   ((0., 0.0416, 0.0416),(0.365079, 1.000000, 1.000000),(1.0, 1.0, 1.0)),\
    'green': ((0., 0., 0.),(0.365079, 0.000000, 0.000000),\
            (0.746032, 1.000000, 1.000000),(1.0, 1.0, 1.0)),\
    'red':  ((0., 0., 0.),(0.746032, 0.000000, 0.000000),(1.0, 1.0, 1.0))}
    my_cmap = matplotlib.colors.LinearSegmentedColormap(fna,cdict_tmp,256);
    matplotlib.cm.register_cmap(name=fna,cmap=my_cmap);

    fna='heat_green'
    cdict_tmp={\
    'green':   ((0., 0.0416, 0.0416),(0.365079, 1.000000, 1.000000),(1.0, 1.0, 1.0)),\
    'red': ((0., 0., 0.),(0.365079, 0.000000, 0.000000),\
            (1.000000, 1.000000, 1.000000),(1.0, 1.0, 1.0)),\
    'blue':  ((0., 0., 0.),(0.1, 0.000000, 0.000000),(1.0, 1.0, 1.0))}
    my_cmap = matplotlib.colors.LinearSegmentedColormap(fna,cdict_tmp,256);
    matplotlib.cm.register_cmap(name=fna,cmap=my_cmap);

    fna='heat_redyellow'
    cdict_tmp={\
    'red':   ((0., 0.0416, 0.0416),(0.365079, 1.000000, 1.000000),(1.0, 1.0, 1.0)),\
    'green': ((0., 0., 0.),(0.365079, 0.000000, 0.000000),\
            (0.746032, 1.000000, 1.000000),(1.0, 1.0, 1.0)),\
    'blue':  ((0., 0., 0.),(0.746032, 0.000000, 0.000000),(1.0, 1.0, 1.0))}
    my_cmap = matplotlib.colors.LinearSegmentedColormap(fna,cdict_tmp,256);
    matplotlib.cm.register_cmap(name=fna,cmap=my_cmap);

    fna='heat_yellow'
    cdict_tmp={\
    'red':   ((0., 0.0416, 0.0416),(0.365079, 1.000000, 1.000000),(1.0, 1.0, 1.0)),\
    'green':   ((0., 0.0416, 0.0416),(0.365079, 1.000000, 1.000000),(1.0, 1.0, 1.0)),\
    'blue':  ((0., 0., 0.),(0.746032, 0.000000, 0.000000),(1.0, 1.0, 1.0))}
    my_cmap = matplotlib.colors.LinearSegmentedColormap(fna,cdict_tmp,256);
    matplotlib.cm.register_cmap(name=fna,cmap=my_cmap);

    fna='heat_purple'
    cdict_tmp={\
    'red':   ((0., 0.0416, 0.0416),(0.565079, 1.000000, 1.000000),(1.0, 1.0, 1.0)),\
    'green': ((0., 0., 0.),(0.565079, 0.000000, 0.000000),\
            (0.946032, 1.000000, 1.000000),(1.0, 1.0, 1.0)),\
    'blue':   ((0., 0.0416, 0.0416),(0.565079, 1.000000, 1.000000),(1.0, 1.0, 1.0))}
    my_cmap = matplotlib.colors.LinearSegmentedColormap(fna,cdict_tmp,256);
    matplotlib.cm.register_cmap(name=fna,cmap=my_cmap);

    fna='heat_orange'
    cdict_tmp={\
    'red':   ((0., 0.0416, 0.0416),(0.365079, 1.000000, 1.000000),(1.0, 1.0, 1.0)),\
    'green':   ((0., 0.0416, 0.0416),(0.365079, 0.400000, 0.400000),(1.0, 1.0, 1.0)),\
    'blue':  ((0., 0., 0.),(0.746032, 0.000000, 0.000000),(1.0, 1.0, 1.0))}
    my_cmap = matplotlib.colors.LinearSegmentedColormap(fna,cdict_tmp,256);
    matplotlib.cm.register_cmap(name=fna,cmap=my_cmap);

    fna='rainbow'
    cdict_tmp = {'red':   ((0.0, 80/256., 80/256.),
                       (0.2, 0.0, 0.0),
                       (0.4, 0.0, 0.0),
                       (0.6, 256/256., 256/256.),
                       (0.95, 256/256., 256/256.),
                       (1.0, 150/256., 150/256.)),
             'green': ((0.0, 0/256., 0/256.),
                       (0.2, 0/256., 0/256.),
                       (0.4, 130/256., 130/256.),
                       (0.6, 256/256., 256/256.),
                       (1.0, 0.0, 0.0)),
             'blue':  ((0.0, 80/256., 80/256.),
                       (0.2, 220/256., 220/256.),
                       (0.4, 0.0, 0.0),
                       (0.6, 20/256., 20/256.),
                       (1.0, 0.0, 0.0))}
    my_cmap = matplotlib.colors.LinearSegmentedColormap(fna,cdict_tmp,256);
    matplotlib.cm.register_cmap(name=fna,cmap=my_cmap);

def hue_minus_convert(hue,lightness,saturation):
    h=hue-90.; h[h<0.]+=360.; h[h>360.]-=360.;
    s=saturation*1.0; s[s<0.]=0.; s[s>1.]=1.;
    l=lightness*1.0; l[l<0.]=0.; l[l>1.]=1.;
    l=l**0.1; l=(l-np.min(l))/(np.max(l)-np.min(l));
    return h,l,s;

def invertreverse_rgb(R,G,B):
    ## invert+huereverse
    r=1.-R; g=1.-G; b=1.-B ## rgb inversion -- *different* from photoshop (cmyk inversion)
    ## convert to hls
    h,l,s=rgb_to_hls(r,g,b);
    ## huereverse
    h+=0.5; h[h>1.]-=1.; h[h<0.]+=1.;
    s*=1.1; s[s<0.]=0.; s[s>1.]=1.;
    l*=1.15; l[l<0.]=0.; l[l>1.]=1.;
    l=l**0.43; l=(l-np.min(l))/(np.max(l)-np.min(l));
    ## convert back to rgb
    r,g,b=hls_to_rgb(h,l,s);
    return r,g,b;

def temperature_map_color_index(mass_pic, temp, set_temp_max=0, set_temp_min=0,
        huem100=0, invertreverse=0):

    ## custom indexing of colors for temperature maps
    ## Assuming Temperature is the weighting :: interesting temp scale is from 1.0d4 to 1.0d7
    temp_max=np.max(temp); temp_min=np.median(temp)/(np.max(temp)/np.median(temp));
    if (set_temp_max != 0): temp_max=set_temp_max;
    if (set_temp_min != 0): temp_min=set_temp_min;
    tmp_scale=np.log10(temp/temp_min)/np.log10(temp_max/temp_min);
    print 'Temp. Scale Min/Max = ',np.min(tmp_scale),temp_min,np.max(tmp_scale),temp_max

    N0=temp.shape[0]; image24=np.zeros([N0,N0,3]);
	## different kernels we can use for the temp weighting
	## use a simple HLS (hue, lightness, saturation) model and use that to get the RGB colors
	## hue runs 0=red - 60=yellow - 120=green - 240=blue - 320=magenta

    ## get something like Volker's scheme with
    tmp_scale[tmp_scale<0.]=0.; tmp_scale[tmp_scale>1.]=1.;
    hue = 240.+tmp_scale*180.; hue[hue >= 360.]-=360.; hue[hue<0.]+=360.;
    saturation = (0.*tmp_scale + 0.99) * 0.5  ## richer colors -> 1, more gray -> 0
    L = mass_pic / 256.;
    lightness = 0.01+0.985*L;  ## 0=black, 1=white
    lightness=0.5 + (lightness-0.5)*0.75;
    lightness[lightness<0.]=0.; lightness[lightness>1.]=1.;
    if (huem100==1): hue,lightness,saturation = hue_minus_convert(hue,lightness,saturation);
    hue /= 360.; ## python uses units of 0-1 for colors

    R,G,B = hls_to_rgb(hue,lightness,saturation);
    if (invertreverse==1): R,G,B = invertreverse_rgb(R,G,B);

    image24[:,:,0]=R; image24[:,:,1]=G; image24[:,:,2]=B
    return image24;

class CustomLightSource(object):
    """
    Create a light source coming from the specified azimuth and elevation.
    Angles are in degrees, with the azimuth measured
    clockwise from north and elevation up from the zero plane of the surface.
    The :meth:`shade` is used to produce rgb values for a shaded relief image
    given a data array.
    (Modified version of matplotlib 'lightsource' that allows for fixed limits
     on shading, important for consistency between images in a series)
    """
    def __init__(self,azdeg=315,altdeg=45,\
                 hsv_min_val=0,hsv_max_val=1,hsv_min_sat=1,hsv_max_sat=0):
       """
       Specify the azimuth (measured clockwise from south) and altitude
       (measured up from the plane of the surface) of the light source
       in degrees.

       The color of the resulting image will be darkened
       by moving the (s,v) values (in hsv colorspace) toward
       (hsv_min_sat, hsv_min_val) in the shaded regions, or
       lightened by sliding (s,v) toward
       (hsv_max_sat hsv_max_val) in regions that are illuminated.
       The default extremes are chose so that completely shaded points
       are nearly black (s = 1, v = 0) and completely illuminated points
       are nearly white (s = 0, v = 1).
       """
       self.azdeg = azdeg
       self.altdeg = altdeg
       self.hsv_min_val = hsv_min_val
       self.hsv_max_val = hsv_max_val
       self.hsv_min_sat = hsv_min_sat
       self.hsv_max_sat = hsv_max_sat

    def shade(self,data,cmap,vmin=None,vmax=None,limit_elevation=False):
        """
        Take the input data array, convert to HSV values in the
        given colormap, then adjust those color values
        to given the impression of a shaded relief map with a
        specified light source.
        RGBA values are returned, which can then be used to
        plot the shaded image with imshow.

        Parameters
        ----------
        data            Input data array
        vmin            Minimum data value for colormap.  Default min(data).
        vmax            Maximum data value for colormap.  Default max(data).
        limit_elevation Limit the elevation in the shading routine? Default False.
                        If true, the elevation will be limited by vmin and vmax.

        Returns
        -------
        rgb             Shaded RGBA values, suitable for use with imshow.

        """

        if (vmin is not None) or (vmax is not None):
            limitschanged = True
        else:
            limitschanged = False
        if vmin is None:
            vmin = np.percentile(data, 5.0)   #min(data)
        if vmax is None:
            vmax = np.percentile(data, 95.0)  #max(data)
        if (np.float(vmax-vmin) ) > 0:
            rgb0 = cmap( (data-vmin) / (np.float(vmax-vmin) ) )
        else:
            rgb0 = cmap( (data-vmin) )

        if limitschanged and limit_elevation:
            d = data.copy()
            d[d<vmin] = vmin
            d[d>vmax] = vmax
            rgb1 = self.shade_rgb(rgb0, elevation=d)
        else:
            rgb1 = self.shade_rgb(rgb0, elevation=data)
        rgb0[:,:,0] = rgb1[:,:, 0]
        rgb0[:,:,1] = rgb1[:,:, 1]
        rgb0[:,:,2] = rgb1[:,:, 2]
        return rgb0

    def shade_rgb(self,rgb, elevation, fraction=1., vmin=None,vmax=None):
        """
        Take the input RGB array (ny*nx*3) adjust their color values
        to given the impression of a shaded relief map with a
        specified light source using the elevation (ny*nx).
        A new RGB array ((ny*nx*3)) is returned.
        """
        # imagine an artificial sun placed at infinity in
        # some azimuth and elevation position illuminating our surface. The parts of
        # the surface that slope toward the sun should brighten while those sides
        # facing away should become darker.
        # convert alt, az to radians
        az = self.azdeg*np.pi/180.0
        alt = self.altdeg*np.pi/180.0
        # gradient in x and y directions
        dx, dy = np.gradient(elevation)

        # clip according to the imposed limits:
        if vmin is None: vmin=np.min(dx);
        if vmax is None: vmax=np.max(dx);
        dx[dx<vmin]=vmin; dx[dx>vmax]=vmax;
        if vmin is None: vmin=np.min(dy);
        if vmax is None: vmax=np.max(dy);
        dy[dy<vmin]=vmin; dy[dy>vmax]=vmax;

        slope = 0.5*np.pi - np.arctan(np.hypot(dx, dy))
        aspect = np.arctan2(dx, dy)
        intensity = np.sin(alt)*np.sin(slope) + np.cos(alt)*np.cos(slope)*np.cos(-az -\
                aspect - 0.5*np.pi)
        # rescale to interval -1,1
        # +1 means maximum sun exposure and -1 means complete shade.
        imin = np.min(intensity)
        imax = np.max(intensity)
        imin1 = (np.sin(alt) + np.cos(alt)*vmin) / np.sqrt(1.+vmin*vmin)
        imin2 = (np.sin(alt) - np.cos(alt)*vmin) / np.sqrt(1.+vmin*vmin)
        imin3 = (np.sin(alt) + np.cos(alt)*vmax) / np.sqrt(1.+vmax*vmax)
        imin4 = (np.sin(alt) - np.cos(alt)*vmax) / np.sqrt(1.+vmax*vmax)
        imin = np.min(np.array([imin1,imin2,imin3,imin4]))
        imax = np.max(np.array([imin1,imin2,imin3,imin4]))
        print ' intensity limits being used in this map == ',\
            np.min(elevation),np.max(elevation),np.min(intensity),np.max(intensity),\
            np.min(dx),np.max(dx),np.min(dy),np.max(dy),np.min(slope),np.max(slope),\
            np.min(aspect),np.max(aspect),imin,imax,az,alt

        intensity[intensity<imin] = imin
        intensity[intensity>imax] = imax
        if (imax - imin) != 0:
            intensity = (intensity - imin)/(imax - imin)

        intensity = (2.*intensity - 1.)*fraction
        # convert to rgb, then rgb to hsv
        hsv = matplotlib.colors.rgb_to_hsv(rgb[:,:,0:3])
        # modify hsv values to simulate illumination.
        hsv[:,:,1] = np.where(np.logical_and(np.abs(hsv[:,:,1])>1.e-10,intensity>0),\
                (1.-intensity)*hsv[:,:,1]+intensity*self.hsv_max_sat, hsv[:,:,1])
        hsv[:,:,2] = np.where(intensity > 0, (1.-intensity)*hsv[:,:,2] +\
                intensity*self.hsv_max_val, hsv[:,:,2])
        hsv[:,:,1] = np.where(np.logical_and(np.abs(hsv[:,:,1])>1.e-10,intensity<0),\
                (1.+intensity)*hsv[:,:,1]-intensity*self.hsv_min_sat, hsv[:,:,1])
        hsv[:,:,2] = np.where(intensity < 0, (1.+intensity)*hsv[:,:,2] -\
                intensity*self.hsv_min_val, hsv[:,:,2])
        hsv[:,:,1:] = np.where(hsv[:,:,1:]<0.,0,hsv[:,:,1:])
        hsv[:,:,1:] = np.where(hsv[:,:,1:]>1.,1,hsv[:,:,1:])
        # convert modified hsv back to rgb.
        return matplotlib.colors.hsv_to_rgb(hsv)


#!/usr/bin/env python
import argparse
import datetime
import numpy as np
import pandas as pd
import pysam

def calculate_variant_quality(variant):
    try:
        return int(min(100, round(-10 * np.log10(1 - (variant.contig_support / float(variant.contig_depth))) * np.log(variant.contig_depth), 0)))
    except ZeroDivisionError:
        return 0

def save_calls_as_vcf(vcf_filename, calls):
    # Save genotypes as tab-delimited file.
    with open(vcf_filename, "w") as vcf:
        vcf.write("##fileformat=VCFv4.2\n")
        vcf.write("##fileDate=%s\n" % datetime.date.strftime(datetime.date.today(), "%Y%m%d"))
        vcf.write("##source=SMRT_SV\n")
        vcf.write('##INFO=<ID=SAMPLES,Number=1,Type=String,Description="Samples with the given variant">' + "\n")
        vcf.write('##INFO=<ID=DP,Number=1,Type=Integer,Description="Mean depth of raw reads">' + "\n")
        vcf.write('##INFO=<ID=CONTIG_DEPTH,Number=1,Type=Integer,Description="Total depth of local assemblies">' + "\n")
        vcf.write('##INFO=<ID=CONTIG_SUPPORT,Number=1,Type=Integer,Description="Depth of local assemblies supporting variant">' + "\n")
        vcf.write('##INFO=<ID=SVTYPE,Number=1,Type=String,Description="Type of structural variant">' + "\n")
        vcf.write('##INFO=<ID=SVLEN,Number=1,Type=Integer,Description="Difference in length between REF and ALT alleles">' + "\n")
        vcf.write('##INFO=<ID=END,Number=1,Type=Integer,Description="End coordinate of this variant">' + "\n")
        vcf.write('##INFO=<ID=CONTIG,Number=1,Type=String,Description="Name of alternate assembly contig">' + "\n")
        vcf.write('##INFO=<ID=CONTIG_START,Number=1,Type=Integer,Description="Start coordinate of this variant in the alternate assembly contig">' + "\n")
        vcf.write('##INFO=<ID=CONTIG_END,Number=1,Type=Integer,Description="End coordinate of this variant in the alternate assembly contig">' + "\n")
        vcf.write('##INFO=<ID=REPEAT_TYPE,Number=1,Type=String,Description="Repeat classification of variant content">' + "\n")
        vcf.write('##INFO=<ID=SEQ,Number=1,Type=String,Description="Sequence associated with variant">' + "\n")

        if len(calls) > 0:
            calls.to_csv(vcf, sep="\t", index=False)

def convert_bed_to_vcf(bed_filename, reference_filename, vcf_filename, sample, variant_type):
    # Get variants.
    if variant_type == "sv":
        columns = (0, 1, 2, 3, 4, 5, 8, 11, 12, 15, 16, 20)
        names = ("chr", "start", "end", "sv_call", "event_size", "sv_sequence", "contig", "contig_start", "contig_end", "contig_support", "contig_depth", "repeat_type")
    elif variant_type == "indel":
        # chr1    94824   94827   3       Cttttcttttttttt 1       1       29.04   deletion
        columns = (0, 1, 2, 3, 4, 5, 6, 7, 8)
        names = ("chr", "start", "end", "event_size", "sv_sequence", "contig_support", "contig_depth", "depth", "sv_call")
    elif variant_type == "inversion":
        columns = (0, 1, 2, 3, 4, 5)
        names = ("chr", "start", "end", "sv_call", "contig_support", "contig_depth")
    else:
        raise Exception("Unsupported variant type: %s" % variant_type)

    calls = pd.read_table(bed_filename, header=None, usecols=columns, names=names)

    # Handle the special case when there are no calls.
    if len(calls) == 0:
        save_calls_as_vcf(vcf_filename, calls)
        return

    calls["sample_name"] = sample
    calls["call_id"] = "."
    calls["quality"] = calls.apply(calculate_variant_quality, axis=1)
    calls["filter"] = "PASS"

    # Get the reference base at the position of the variant start.
    reference = pysam.FastaFile(reference_filename)
    calls["reference"] = calls.apply(lambda row: reference.fetch(row.chr, row.start, row.start + 1).upper(), axis=1)

    # Update start position to be 1-based.
    calls["start"] = calls["start"] + 1

    # Build an INFO field for each call.
    if variant_type == "sv":
        calls["alt"] = calls.apply(lambda row: "<%s>" % row.sv_call[:3].upper(), axis=1)
        calls["info"] = calls.apply(
            lambda row: ";".join(
                ["=".join(map(str, item))
                 for item in (
                        ("END", row.end),
                        ("SVTYPE", row.sv_call),
                        ("SVLEN", row.event_size),
                        ("CONTIG", row.contig),
                        ("CONTIG_START", row.contig_start),
                        ("CONTIG_END", row.contig_end),
                        ("REPEAT_TYPE", row.repeat_type),
                        ("CONTIG_SUPPORT", row.contig_support),
                        ("CONTIG_DEPTH", row.contig_depth),
                        ("SAMPLES", row.sample_name),
                        ("SEQ", row.sv_sequence)
                 )]
            ),
            axis=1
        )
    elif variant_type == "indel":
        calls["alt"] = calls.apply(lambda row: "<%s>" % row.sv_call[:3].upper(), axis=1)
        calls["info"] = calls.apply(
            lambda row: ";".join(
                ["=".join(map(str, item))
                 for item in (
                        ("END", row.end),
                        ("SVTYPE", row.sv_call),
                        ("SVLEN", row.event_size),
                        ("CONTIG_SUPPORT", row.contig_support),
                        ("CONTIG_DEPTH", row.contig_depth),
                        ("DP", row.depth),
                        ("SAMPLES", row.sample_name),
                        ("SEQ", row.sv_sequence)
                 )]
            ),
            axis=1
        )
    elif variant_type == "inversion":
        calls["alt"] = "<INV>"
        calls["info"] = calls.apply(
            lambda row: ";".join(
                ["=".join(map(str, item))
                 for item in (
                        ("END", row.end),
                        ("SVTYPE", row.sv_call),
                        ("SVLEN", row.end - row.start),
                        ("CONTIG_SUPPORT", row.contig_support),
                        ("CONTIG_DEPTH", row.contig_depth),
                        ("SAMPLES", row.sample_name),
                 )]
            ),
            axis=1
        )

    simple_calls = calls[["chr", "start", "call_id", "reference", "alt", "quality", "filter", "info"]].rename_axis({"chr": "#CHROM", "start": "POS", "reference": "REF", "call_id": "ID", "quality": "QUAL", "info": "INFO", "alt": "ALT", "filter": "FILTER"}, axis=1)

    # Save calls as VCF.
    save_calls_as_vcf(vcf_filename, simple_calls)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("bed", help="input BED file of variant calls")
    parser.add_argument("reference", help="FASTA file for reference used with SMRT SV")
    parser.add_argument("vcf", help="output VCF file of variant calls")
    parser.add_argument("sample", help="name of sample with variants")
    parser.add_argument("type", help="variant call type", choices=("sv", "indel", "inversion"))
    args = parser.parse_args()

    convert_bed_to_vcf(args.bed, args.reference, args.vcf, args.sample, args.type)