from __future__ import absolute_import
import xarray as xr
import numpy as np
import pandas as pd
import datetime
import os
#turn off warnings so i can use the progressbar
import warnings
warnings.filterwarnings('ignore')

class GPMDPR():

    """
    Author: Randy J. Chase. This class is intended to help with the efficient processing of GPM-DPR radar files.
    Xarray now supports direct reading the NASA hdf5 files. Here is the major refactor from the previous version
    of DRpy.

    Feel free to reach out to me on twitter (@dopplerchase) or email randychase@ou.edu

    For your reference, please check out GPM-DPR's ATBD: https://pps.gsfc.nasa.gov/GPMprelimdocs.html
    """

    def __init__(self,filename=[],bounding_box=None,outer_swath=False,auto_run=True,heavy=True):
        """
        Initializes things

        params::
        filename: str, path to GPM-DPR file
        boundingbox: list of floats, if you would like to cut the gpm to a lat lon box
        send in a list of [lon_min,lon_mat,lat_min,lat_max]
        """
        self.filename = filename
        self.corners = bounding_box
        self.heavy=heavy

        if auto_run:
            #this reads the hdf5 file
            self.read()
            #this gets a datetime obj for each scan
            self.parse_dtime()

    def read(self):
        """
        This method unfolds all the groups into one combined xarray dataset for
        each radar (e.g., KuPR and KaPR). It will be lazily loaded to save on RAM.
        Note that this code was primarily developed for V7 DPR products. It will not
        work for V6 data

        """
        #######################################################################
        ################################ KuPR #################################
        #######################################################################

        prefix = '/FS/'
        geo = xr.open_dataset(self.filename,group=prefix,engine='netcdf4',decode_cf=False)
        pre = xr.open_dataset(self.filename,group=prefix+'PRE',engine='netcdf4',decode_cf=False)
        slv = xr.open_dataset(self.filename,group=prefix+'SLV',engine='netcdf4',decode_cf=False)
        tim = xr.open_dataset(self.filename,group=prefix+'ScanTime',engine='netcdf4',decode_cf=False)

        #rename dims to proper names
        bad_dims = list(geo.dims)
        geo = geo.rename_dims({bad_dims[0]:'nscan',
                            bad_dims[1]:'nrayNS'})
        bad_dims = list(pre.dims)
        pre = pre.rename_dims({bad_dims[0]:'nscan',
                            bad_dims[1]:'nrayNS',
                            bad_dims[2]:'nfreq',
                            bad_dims[3]:'nbin'})
        bad_dims = list(slv.dims)
        slv = slv.rename_dims({bad_dims[0]:'nscan',
                            bad_dims[1]:'nrayNS',
                            bad_dims[2]:'nbin',
                            bad_dims[3]:'nfreq',
                            bad_dims[4]:'nNUBF'})
        bad_dims = list(tim.dims)
        tim = tim.rename_dims({bad_dims[0]:'nscan'})

        #if you want to load the whole thing in, turn the heavy flag on
        if self.heavy:
            ver = xr.open_dataset(self.filename,group=prefix+'VER',engine='netcdf4',decode_cf=False)
            srt = xr.open_dataset(self.filename,group=prefix+'SRT',engine='netcdf4',decode_cf=False)
            csf = xr.open_dataset(self.filename,group=prefix+'CSF',engine='netcdf4',decode_cf=False)
            exp = xr.open_dataset(self.filename,group=prefix+'Experimental',engine='netcdf4',decode_cf=False)
            flg = xr.open_dataset(self.filename,group=prefix+'FLG',engine='netcdf4',decode_cf=False)
            #rename dims to proper names
            bad_dims = list(ver.dims)
            ver = ver.rename_dims({bad_dims[0]:'nscan',
                                bad_dims[1]:'nrayNS',
                                bad_dims[2]:'nbin',
                                bad_dims[3]:'nfreq',
                                bad_dims[4]:'nNP'})
            bad_dims = list(srt.dims)
            srt = srt.rename_dims({bad_dims[0]:'nscan',
                                bad_dims[1]:'nrayNS',
                                bad_dims[2]:'method',
                                bad_dims[3]:'foreBack',
                                bad_dims[4]:'nearFar',
                                bad_dims[5]:'nsdew'})
            bad_dims = list(csf.dims)
            csf = csf.rename_dims({bad_dims[0]:'nscan',
                                bad_dims[1]:'nrayNS',
                                bad_dims[2]:'nfreqHI'})
            bad_dims = list(exp.dims)
            exp = exp.rename_dims({bad_dims[0]:'nscan',
                                bad_dims[1]:'nrayNS',
                                bad_dims[2]:'nbinSZP',
                                bad_dims[3]:'nfreq'})
            bad_dims = list(flg.dims)
            flg = flg.rename_dims({bad_dims[0]:'nscan',
                                bad_dims[1]:'nrayNS',
                                bad_dims[2]:'nbin',
                                bad_dims[3]:'nfreq'})

            #MERGE into one ds
            self.ds = xr.merge([geo,pre,slv,ver,srt,csf,exp,flg,tim])
        else:
            #MERGE into one ds
            self.ds = xr.merge([geo,pre,slv,tim])

        #close uneeded xr datasets
        geo.close()
        pre.close()
        slv.close()
        tim.close()
        if self.heavy:
          ver.close()
          srt.close()
          csf.close()
          exp.close()
          flg.close()

        #set lat,lon,height as the coords to allow for easy xr slicing
        self.ds = self.ds.set_coords(['Latitude','Longitude','height'])

    def setboxcoords(self):
        """
        This method sets all points outside the box to nan.
        """
        if len(self.corners) > 0:
            self.ll_lon = self.corners[0]
            self.ur_lon = self.corners[1]
            self.ll_lat = self.corners[2]
            self.ur_lat = self.corners[3]
            self.ds = self.ds.where((self.ds.Longitude >= self.ll_lon) & (self.ds.Longitude <= self.ur_lon) & (self.ds.Latitude >= self.ll_lat)  & (self.ds.Latitude <= self.ur_lat),drop=False)
        else:
            print('ERROR, no boxcoods set...did you mean to do this?')
    def parse_dtime(self):
        year = self.ds.Year.values
        ind = np.where(year == -9999)[0]
        year = np.asarray(year,dtype=str)
        year = list(year)

        month = self.ds.Month.values
        month = np.asarray(month,dtype=str)
        month = np.char.rjust(month, 2, fillchar='0')
        month = list(month)

        day = self.ds.DayOfMonth.values
        day = np.asarray(day,dtype=str)
        day = np.char.rjust(day, 2, fillchar='0')
        day = list(day)

        hour = self.ds.Hour.values
        hour = np.asarray(hour,dtype=str)
        hour = np.char.rjust(hour, 2, fillchar='0')
        hour = list(hour)

        minute = self.ds.Minute.values
        minute = np.asarray(minute,dtype=str)
        minute = np.char.rjust(minute, 2, fillchar='0')
        minute = list(minute)

        second = self.ds.Second.values
        second = np.asarray(second,dtype=str)
        second = np.char.rjust(second, 2, fillchar='0')
        second = list(second)

        millisecond = self.ds.MilliSecond.values
        millisecond = np.asarray(millisecond,dtype=str)
        millisecond = np.char.rjust(millisecond, 3, fillchar='0')
        millisecond = list(millisecond)

        datestr  = [year[i] +"-"+ month[i]+ "-" + day[i] + \
                    ' ' + hour[i] + ':' + minute[i] + ':' + second[i] + '.' + millisecond[i] for i in range(len(year))]
        datestr = np.asarray(datestr,dtype=str)
        datestr[ind] = '1970-01-01 00:00:00'
        datestr = np.reshape(datestr,[len(datestr),1])
        datestr = np.tile(datestr,(1,49))

        self.ds['time'] = xr.DataArray(np.asarray(datestr,dtype=np.datetime64),dims=['nscan','nrayNS'])
        self.ds = self.ds.set_coords('time')
        #drop the variables we dont need
        self.ds = self.ds.drop(['Year','Month','DayOfMonth','Hour','Minute','Second','MilliSecond','DayOfYear','SecondOfDay'])

    def get_physcial_distance(self,reference_point = None):
        """
        This method uses pyproj to calcualte distances between lats and lons.
        reference_point is a list or array conisting of two entries, [Longitude,Latitude]

        Please note that this intentionally uses an older version of pyproj (< version 2.0, i used 1.9.5.1)
        This is because it preserves how the function is called.
        """

        if reference_point is None and self.reference_point is None:
            print('Error, no reference point found...please enter one')
        else:
            #this is envoke the pyproj package. Please note this must be an old version** < 2.0
            from pyproj import Proj
            p = Proj(proj='aeqd', ellps='WGS84', datum='WGS84', lat_0=reference_point[1], lon_0=reference_point[0])
            #double check to make sure this returns 0 meters
            x,y = p(reference_point[0],reference_point[1])
            if np.sqrt(x**2 + y**2) != 0:
                'something isnt right with the projection. investigate'
            else:
                ind = np.isnan(self.ds.zFactorFinalNearSurface.values[:,:,0])
                x = np.zeros(self.ds.Longitude.values.shape)
                y = np.zeros(self.ds.Latitude.values.shape)
                x[~ind],y[~ind] = p(self.ds.Longitude.values[~ind],self.ds.Latitude.values[~ind])
                x[ind] = np.nan
                y[ind] = np.nan
                da = xr.DataArray(np.sqrt(x**2 + y**2)/1000, dims=['nscan', 'nrayNS'])
                da.attrs['units'] = 'km'
                da.attrs['standard_name'] = 'distance, way of the crow (i.e. direct), to the reference point'
                self.ds['distance'] = da
    def run_Chase2021(self,models_path='../models/'):

        """ Method to run Chase et al. (2021) JAMC Neural Network retrieval. THIS NEEDS TENSORFLOW!

        Outputs:
        ========

        This retrieval will add the following variables to the dataset

        Dml_nn:2d np.array, dim = (along_track,range), unit = mm, name= liquid eq. mass weighted mean diameter
        Dms_nn: 2d np.array, dim = (along_track,range), unit = mm, name= solid phase mass weighted mean diameter
        Nw_nn: 2d np.array, dim = (along_track,range), unit = log(m^-4), name= liquid eq. normalized intercept parameter
        IWC_nn: 2d np.array, dim = (along_track,range), unit = g m^{-3}, name= ice water content

        ========
        """
        import pickle
        import copy

        try:
            import tensorflow as tf
        except:
            print("No Tensorflow. Please go install it with conda")

        with open(models_path+ 'scaler_X.pkl', 'rb') as inp:
            scaler_X = pickle.load(inp)

        with open(models_path+ 'scaler_y.pkl', 'rb') as inp:
            scaler_y = pickle.load(inp)

        Ku = copy.deepcopy(self.ds.zFactorFinal[:,:,:,0].values) #use the corrected Ku
        Ka = copy.deepcopy(self.ds.zFactorMeasured[:,:,:,1].values) #use the raw Ka
        T = copy.deepcopy(self.ds.airTemperature.values) #grab JMA temperature
        T = T - 273.15 #convert to degC

        #supress warnings. skrews up my progress bar when running in parallel
        def warn(*args, **kwargs):
            pass
        import warnings
        warnings.warn = warn
        #

        #load the trained model
        tf.config.run_functions_eagerly(True)
        model = tf.keras.models.load_model(models_path + 'NN_6by8.h5',custom_objects=None,compile=True)

        #now we have to reshape things to make sure they are in the right shape for the NN model [n_samples,n_features]
        shape_step1 = Ku.shape
        Ku = Ku.reshape([Ku.shape[0]*Ku.shape[1]*Ku.shape[2]])
        Ka = Ka.reshape([Ka.shape[0]*Ka.shape[1]*Ka.shape[2]])
        T = T.reshape([T.shape[0]*T.shape[1]*T.shape[2]])

        #make sure that Ku and Ka match
        Ku[np.isnan(Ka)] = np.nan
        Ka[np.isnan(Ku)] = np.nan
        #
        #mask gates with temperatures > 0
        Ku[T>0] = np.nan
        #mask gates with DFR < -0.5, shouldnt be ever be lower than this (plus/minus Cal uncert.)
        Ku[(Ku-Ka) < -0.5] = np.nan

        #Make sure we only run in on non-nan values to save time
        ind_masked = np.isnan(Ku)
        ind_masked2 = np.isnan(Ka)
        Ku_nomask = np.ones(Ku.shape)*-9999.
        Ka_nomask = np.ones(Ka.shape)*-9999.
        T_nomask = np.ones(T.shape)*-9999.
        Ku_nomask[~ind_masked] = Ku[~ind_masked]
        Ka_nomask[~ind_masked] = Ka[~ind_masked]
        T_nomask[~ind_masked] = T[~ind_masked]

        ind = np.where(Ku_nomask!=-9999.)[0]

        #scale the input vectors by the mean that it was trained with
        X = np.zeros([Ku_nomask.shape[0],3])
        X[:,0] = (Ku_nomask - scaler_X.mean_[0])/scaler_X.scale_[0] #Ku
        X[:,1] = ((Ku_nomask - Ka_nomask)- scaler_X.mean_[1])/scaler_X.scale_[1] #DFR Ku - Ka
        X[:,2] = (T_nomask - scaler_X.mean_[2])/scaler_X.scale_[2] #T
        #

        #conduct the retrieval
        yhat = model.predict(X[ind,0:3],batch_size=len(X[ind,0]))
        #scale it properly
        yhat[:,0] = (yhat[:,0]*scaler_y.scale_[0]) + scaler_y.mean_[0]
        yhat[:,1] = (yhat[:,1]*scaler_y.scale_[1]) + scaler_y.mean_[1]
        yhat[:,2] = (yhat[:,2]*scaler_y.scale_[2]) + scaler_y.mean_[2]
        yhat[:,1] = 10**yhat[:,1] #unlog Dm liquid
        yhat[:,2] = 10**yhat[:,2] #unlog Dm solid

        #fill the Nw array
        ind = np.where(Ku_nomask!=-9999.)[0]
        Nw = np.zeros(Ku_nomask.shape)
        Nw[ind] = np.squeeze(yhat[:,0])
        #reshape it back to input shape
        Nw = Nw.reshape([shape_step1[0],shape_step1[1],shape_step1[2]])
        Nw[Nw==0.0] = np.nan
        #shove it into the xarray dataset
        self.ds['Nw_nn'] = xr.DataArray(Nw,dims=['nscan','nrayNS','nbin'])

        Dm = np.zeros(Ku_nomask.shape)
        Dm[ind] = np.squeeze(yhat[:,1])
        Dm = Dm.reshape([shape_step1[0],shape_step1[1],shape_step1[2]])
        Dm[Dm==0.0] = np.nan
        self.ds['Dml_nn'] = xr.DataArray(Dm,dims=['nscan','nrayNS','nbin'])

        Dm_frozen = np.zeros(Ku_nomask.shape)
        Dm_frozen[ind] = np.squeeze(yhat[:,2])
        Dm_frozen = Dm_frozen.reshape([shape_step1[0],shape_step1[1],shape_step1[2]])
        Dm_frozen[Dm_frozen==0.0] = np.nan
        self.ds['Dms_nn'] = xr.DataArray(Dm_frozen,dims=['nscan','nrayNS','nbin'])

        #calculate IWC
        Nw = 10**Nw #undo log, should be in m^-4
        Dm = Dm/1000. # convert to m ^4
        IWC = (Nw*(Dm)**4*1000*np.pi)/4**(4) # the 1000 is density of water (kg/m^3)
        IWC = IWC*1000 #convert to g/m^3
        self.ds['IWC_nn'] = xr.DataArray(IWC,dims=['nscan','nrayNS','nbin'])

        #calculate R (following Chase et al. 2022 paramaterization between log(R) - log(IWC)
        R = 3.64*(IWC**1.06)
        self.ds['R_nn'] = xr.DataArray(R,dims=['nscan','nrayNS','nbin'])

        return


import cv2
import glob
import numpy as np
import scipy as sp
from matplotlib import *
from pylab import *
import time
from keras.models import Model, model_from_json
import numpy as np
from keras import backend as K
import tensorflow as tf
import os
from sklearn.utils import shuffle
from keras.models import model_from_json
from sklearn.metrics import mean_squared_error
import csv

def load_model(model_path):
    with open(model_path + 'auto_model.json', 'r') as jfile:
            model_svdd = model_from_json(jfile.read())
    model_svdd.load_weights(model_path + 'auto_model.h5')
    return model_svdd

def mse(imageA, imageB):
	# the 'Mean Squared Error' between the two images is the
	# sum of the squared difference between the two images;
	# NOTE: the two images must have the same dimension
	err = np.sum((imageA.astype("float") - imageB.astype("float")) ** 2)
	err /= float(imageA.shape[0] * imageA.shape[1])

	# return the MSE, the lower the error, the more "similar"
	# the two images are
	return err

#Load complete input images without shuffling
def load_images(paths,calib_path):
    numImages = 0
    inputs = []
    comp_inp = []
    for path in paths:
        print(path)
        numFiles = len(glob.glob1(path,'*.png'))
        numImages += numFiles
        for img in glob.glob(path+'*.png'):
            img = cv2.imread(img)
            img = cv2.resize(img, (224, 224))
            #img = img / 255.
            inputs.append(img)
    #inpu = shuffle(inputs)
    #print("Total number of images:%d" %(numImages))
    j=0
    for i in range(0,len(inputs),5):
        cv2.imwrite(calib_path + "/frame%d.png"%j,inputs[i])
        j+=1
            #comp_inp.append(calib_path+ "/",inputs[i])
    print("Total number of images:%d" %j)
    # return inputs, comp_inp
    return inputs

def createFolderPaths(path,folders):
    paths = []
    for folder in folders:
        data_path = path + folder + '/'
        paths.append(data_path)
    return paths

def load_training_images(path,trainingFolders,calib_path):
    paths = createFolderPaths(path,trainingFolders)
    return load_images(paths,calib_path)

def load_calib_images(calib_path):
    numImages = 0
    inputs = []
    for img in glob.glob(calib_path+'*.png'):
        img = cv2.imread(img)
        img = cv2.resize(img, (224, 224))
        img = img / 255.
        inputs.append(img)
    #inpu = shuffle(inputs)
    print("Total number of images:%d" %(numImages))
    return inputs

def calib_data_generation(model_vae,calib_path,model_path):
    calib_images = load_calib_images(calib_path)
    calib_images = np.array(calib_images)
    calib_images = np.reshape(calib_images, [-1, calib_images.shape[1],calib_images.shape[2],calib_images.shape[3]])
    for i in range(0,len(calib_images)):
        dist_val=[]
        img = np.array(calib_images[i])[np.newaxis]
        predicted_reps = model_vae.predict(img)
        dist = mse(predicted_reps, img)
        print(dist)
        dist_val.append(dist)
        with open(model_path + 'test-noon.csv', 'a') as file:
            writer = csv.writer(file)
            writer.writerow(dist_val)

if __name__ == '__main__':
    path = "/home/scope/Carla/autopilot_Carla_ad/sample_data/new-trial/"
    # list of folders used in training
    trainingFolders = ["clear_noon1","clear_noon2","cloudy_noon1","cloudy_noon2"]
    calib_path = "/home/scope/Carla/autopilot_Carla_ad/sample_data/new-trial/clear_noon1/"
    model_path = "/home/scope/Carla/autopilot_Carla_ad/leaderboard/team_code/detector_code/new-trial-100/"
    model_vae=load_model(model_path)
    calib_data_generation(model_vae,calib_path,model_path)


"""
Programmer: Dahsom Jang
Description: Face Detection and Segmentation
"""

import tensorflow as tf
import numpy as np
import cv2

class ImageSegmentation:
    # Class that recognizes faces and provides segmentation
    def __init__(self, width, height, no_drop=True):

        self.width, self.height = width, height
        self.input_size = (256, 256)

        # Image Segmentation Model Load
        if no_drop:
            self.model = tf.keras.models.load_model('./models/unet_no_drop.h5')
        else:
            self.model = tf.keras.models.load_model('./models/unet.h5')

    def predict(self, image):
        # Function to pre-process an image,
        # perform segmentation prediction using a model,
        # and generate and return a mask as a result
        input_image = self.image_ready(image)
        seg_mask = self.model.predict(input_image)
        mask = self.generate_mask(seg_mask)
        return mask

    def image_ready(self, image):
        # Image pre-processing function
        if type(image) != np.ndarray:
            pil_image = image.convert('RGB')
            open_cv_image = np.array(pil_image)
            open_cv_image = open_cv_image[:, :, ::-1].copy()
            image = open_cv_image
        else:
            image = image

        img = image
        width, height = self.input_size

        im = np.zeros((width, height, 3), dtype=np.uint8)
        if img.shape[0] >= img.shape[1]:
            scale = img.shape[0] / height
            new_width = int(img.shape[1] / scale)
            diff = (width - new_width) // 2
            img = cv2.resize(img, (new_width, height))
            im[:, diff:diff + new_width, :] = img

        else:
            scale = img.shape[1] / width
            new_height = int(img.shape[0] / scale)
            diff = (height - new_height) // 2
            img = cv2.resize(img, (width, new_height))
            im[diff:diff + new_height, :, :] = img

        input_image = im.reshape((width, height, 3))
        input_image = image_sharpening(input_image)
        input_image = input_image.reshape((1, width, height, 3)).astype(np.float32) / 255.

        return input_image

    def generate_mask(self, seg_mask, threshold=0.1):
        # Function to generate mask with predicted segmentation information
        height, width = self.height, self.width
        prediction = seg_mask
        mask_ori = (prediction.squeeze()[:, :, 1] > threshold).astype(np.uint8)
        max_size = max(width, height)
        result_mask = cv2.resize(mask_ori, dsize=(max_size, max_size))
        if height >= width:
            diff = (max_size - width) // 2
            if diff > 0:
                result_mask = result_mask[:, diff:-diff]
        else:
            diff = (max_size - height) // 2
            if diff > 0:
                result_mask = result_mask[diff:-diff, :]
        result_mask = cv2.resize(result_mask, dsize=(width, height))
        return result_mask

def image_sharpening(image):
    # Sharpening the image
    sharpening = np.array([[-1, -1, -1, -1, -1],
                           [-1, 2, 2, 2, -1],
                           [-1, 2, 9, 2, -1],
                           [-1, 2, 2, 2, -1],
                           [-1, -1, -1, -1, -1]]) / 9.0

    dst = cv2.filter2D(image, -1, sharpening)
    return dst

# if __name__ == '__main__':
#
#     img = cv2.imread('images/1.png', cv2.IMREAD_COLOR)
#     cv2.resize(img, (640, 480))
#     image_segmentation = ImageSegmentation(640, 480)
#     copy_img = img.copy()
#     seg_mask = image_segmentation.predict(copy_img)
#     mask = cv2.cvtColor(seg_mask, cv2.COLOR_GRAY2RGB)
#     cv2.imshow('1', mask)
#
#     cv2.waitKey(0)
#     cv2.destroyAllWindows()


# -*- coding: utf-8 -*-
import numpy as np

from qulab.Driver import visaDriver, QList, QOption, QReal

class Driver(visaDriver):
    support_models = ['AWG70001A', 'AWG70002A']

    CHs=[1,2]
    quants = [
        # Sample Rate set_cmd is block cmd
        #
        QReal('Sample Rate',
              unit='S/s',
              set_cmd='CLOC:SRAT %(value).10e; *WAI;',
              get_cmd='CLOC:SRAT?'),
        QOption('Run Mode',
                value='CONT',
                ch=1,
                set_cmd='SOUR%(ch)d:RMOD %(option)s',
                get_cmd='SOUR%(ch)d:RMOD?',
                options=[('Continuous', 'CONT'), ('Triggered', 'TRIG'),
                         ('TContinuous', 'TCON')]),
        QOption('Clock Source',
                value='INT',
                set_cmd='CLOC:SOUR %(option)s',
                get_cmd='CLOC:SOUR?',
                options=[('Internal', 'INT'), ('External', 'EXT'),
                         ('Efixed', 'EFIX'), ('Evariable', 'EVAR')]),

        # QOption('Reference Source', set_cmd='SOUR:ROSC:SOUR %(option)s', get_cmd='SOUR:ROSC:SOUR?',
        #   options = [('Internal', 'INT'), ('External', 'EXT')]),
        QReal('Multiplier Rate ',
              value=1,
              set_cmd='CLOC:EREF:MULT %(value)d',
              get_cmd='CLOC:EREF:MULT?'),
        QReal('Divider Rate ',
              value=1,
              set_cmd='CLOC:EREF:DIV %(value)d',
              get_cmd='CLOC:EREF:DIV?'),

        # 以下四种只在 output path 为 Direct 时有效
        # Vpp range: 250-500mVpp
        QReal('Vpp',
              unit='Vpp',
              ch=1,
              set_cmd='SOUR%(ch)d:VOLT %(value)f%(unit)s',
              get_cmd='SOUR%(ch)d:VOLT?'),
        QReal('Offset',
              unit='V',
              ch=1,
              set_cmd='SOUR%(ch)d:VOLT:OFFS %(value)f%(unit)s',
              get_cmd='SOUR%(ch)d:VOLT:OFFS?'),
        QReal('Volt Low',
              unit='V',
              ch=1,
              set_cmd='SOUR%(ch)d:VOLT:LOW %(value)f%(unit)s',
              get_cmd='SOUR%(ch)d:VOLT:LOW?'),
        QReal('Volt High',
              unit='V',
              ch=1,
              set_cmd='SOUR%(ch)d:VOLT:HIGH %(value)f%(unit)s',
              get_cmd='SOUR%(ch)d:VOLT:HIGH?'),

        # output delay in time
        QReal('timeDelay',
              unit='s',
              ch=1,
              set_cmd='SOUR%(ch)d:DEL:ADJ %(value)f%(unit)s',
              get_cmd='SOUR%(ch)d:DEL:ADJ?'),
        # output delay in point
        QReal('timeDelay',
              unit='point',
              ch=1,
              set_cmd='SOUR%(ch)d:DEL:POIN %(value)d',
              get_cmd='SOUR%(ch)d:DELy:POIN?'),
        QOption('Run',
                value='Play',
                set_cmd='AWGC:%(option)s; *WAI;',
                get_cmd='AWGC:RST?',
                options=[('Play', 'RUN'), ('Stop', 'STOP')]),
        QOption('Output',
                ch=1,
                set_cmd='OUTP%(ch)d %(option)s',
                get_cmd='OUTP%(ch)d?',
                options=[('ON', 1), ('OFF', 0), (1, 'ON'), (0, 'OFF')]),
        QList('WList'),
        QList('SList'),

        # INSTrument MODE : AWG or FG
        QOption('instMode',
                value='AWG',
                set_cmd='INST:MODE %(option)s',
                get_cmd='INST:MODE?',
                options=[
                    ('AWG', 'AWG'),
                    ('FG', 'FGEN'),
                ]),
        QOption('FG Type',
                ch=1,
                value='Sin',
                set_cmd='FGEN:CHAN%(ch)d:TYPE %(option)s',
                get_cmd='FGEN:CHAN%(ch)d:TYPE?',
                options=[
                    ('Sin', 'SINE'),
                    ('Square', 'SQU'),
                    ('Triangle', 'TRI'),
                    ('Noise', 'NOIS'),
                    ('DC', 'DC'),
                    ('Gaussian', 'GAUS'),
                    ('ExpRise', 'EXPR'),
                    ('ExpDecay', 'EXPD'),
                    ('None', 'NONE'),
                ]),
        QReal('FG Amplitude',
              ch=1,
              value=0.5,
              unit='V',
              set_cmd='FGEN:CHAN%(ch)d:AMPL:VOLT %(value)f%(unit)s',
              get_cmd='FGEN:CHAN%(ch)d:AMPL:VOLT?'),

        #FG Offset -150mV~150mV, min Unit 1mV
        QReal('FG Offset',
              ch=1,
              value=0,
              unit='V',
              set_cmd='FGEN:CHAN%(ch)d:OFFS %(value)f%(unit)s',
              get_cmd='FGEN:CHAN%(ch)d:OFFS?'),
        QReal('FG Phase',
              ch=1,
              value=0,
              unit='deg',
              set_cmd='FGEN:CHAN%(ch)d:PHAS %(value)f',
              get_cmd='FGEN:CHAN%(ch)d:PHAS?'),
        QReal('FG Frequency',
              ch=1,
              value=1e6,
              unit='Hz',
              set_cmd='FGEN:CHAN%(ch)d:FREQ %(value)f%(unit)s',
              get_cmd='FGEN:CHAN%(ch)d:FREQ?'),
        QReal('FG Period',
            ch=1,
            unit='s',
            # 周期与频率绑定，无法设置周期，但可读
            set_cmd='',
            get_cmd='FGEN:CHAN%(ch)d:PER?'),

        # DC Level Range: –250 mV to 250 mV
        QReal('FG DC',
              ch=1,
              value=0,
              unit='V',
              set_cmd='FGEN:CHAN%(ch)d:DCL %(value)f%(unit)s',
              get_cmd='FGEN:CHAN%(ch)d:DCL?'),
        QReal('FG High',
              ch=1,
              value=0.25,
              unit='V',
              set_cmd='FGEN:CHAN%(ch)d:HIGH %(value)f%(unit)s',
              get_cmd='FGEN:CHAN%(ch)d:HIGH?'),
        QReal('FG Low',
              ch=1,
              value=-0.25,
              unit='V',
              set_cmd='FGEN:CHAN%(ch)d:LOW %(value)f%(unit)s',
              get_cmd='FGEN:CHAN%(ch)d:LOW?'),

        #coupling mode: DIR, DCAM, AC
        QOption('FG Path',
                ch=1,
                value='Direct',
                set_cmd='FGEN:CHAN%(ch)d:PATH %(option)s',
                get_cmd='FGEN:CHAN%(ch)d:PATH?',
                options=[
                    ('Direct', 'DIR'),
                    ('DCAmplified', 'DCAM'),
                    ('AC', 'AC'),
                ]),
    ]

    def performOpen(self):
        super(Driver,self).performOpen()
        self.waveform_list = self.get_waveform_list()
        try:  #没有sequence模块的仪器会产生一个错误
            self.sequence_list = self.get_sequence_list()
        except:
            self.sequence_list = None

    def performSetValue(self, quant, value, **kw):
        super(Driver,self).performSetValue(self, quant, value, **kw)

    def performGetValue(self, quant, **kw):
        if quant.name == 'WList':
            self.waveform_list = self.get_waveform_list()
            return self.waveform_list
        elif quant.name == 'SList':
            self.sequence_list = self.get_sequence_list()
            return self.sequence_list
        else:
            return super(Driver,self).performGetValue(self, quant, **kw)

    def get_waveform_list(self):
        return self.query('WLIS:LIST?').strip("\"\n' ").split(',')

    def get_sequence_list(self):
        ret = []
        slist_size = int(self.query("SLIS:SIZE?"))
        for i in range(slist_size):
            ret.append(self.query("SLIS:NAME? %d" % i).strip("\"\n '"))
        return ret

    def create_waveform(self, name, length, format='REAL'):
        '''format: REAL or IQ'''
        if name in self.waveform_list:
            return
        self.write('WLIS:WAV:NEW "%s",%d,%s;*WAI;' % (name, length, format))
        self.waveform_list = self.get_waveform_list()

    def remove_waveform(self, name=None, all=False):
        if all:
            self.write('WLIS:WAV:DEL ALL; *WAI;')
            self.waveform_list.clear()
        elif name not in self.waveform_list:
            return
        else:
            self.write('WLIS:WAV:DEL "%s"; *WAI;' % name)
            self.waveform_list = self.get_waveform_list()

    def get_waveform_length(self, name):
        size = int(self.query('WLIS:WAV:LENGTH? "%s"' % name))
        return size

    def use_waveform(self, name, ch=1, type=None):
        '''type: I or Q'''
        if type is not None:
            self.write('SOUR%d:CASS:WAV "%s",%s' % (ch, name, type))
        else:
            self.write('SOUR%d:CASS:WAV "%s"' % (ch, name))
        self.write('*WAI;')

    # 关于RUN的设置和状态询问，建议使用Quantity：Run的方法
    def run_state(self):
        return int(self.query('AWGC:RST?'))

    def run(self):
        self.write('AWGC:RUN')
        self.write('*WAI')

    def stop(self):
        self.write('AWGC:STOP')

    # 关于Output的设置和状态询问，建议使用Quantity：Output的方法
    def output_on(self, ch=1):
        self.write('OUTP%d:STAT 1' % ch)

    def output_off(self, ch=1):
        self.write('OUTP%d:STAT 0' % ch)

    def get_current_asset(self, ch=1):
        current_type = self.query('SOUR%d:CASS:TYPE?' % ch)
        current_asset = self.query('SOUR%d:CASS?' % ch)
        return current_type, current_asset

    def update_waveform(self, points, name='ABS', IQ='I', start=0, size=None):
        w_type = self.query('WLIS:WAV:TYPE? "%s"' % name).strip()
        if w_type == 'REAL':
            self._update_waveform_float(points, name, IQ, start, size)
        elif w_type == 'IQ':
            self._update_waveform_float(points[0], name, 'I', start, size)
            self._update_waveform_float(points[1], name, 'Q', start, size)
        self.write('*WAI;')
        # else:
        #     self._update_waveform_int(points, name, start, size)

    # def _update_waveform_int(self, points, name='ABS', start=0, size=None):
    #     """
    #     points : a 1D numpy.array which values between -1 and 1.
    #     """
    #     message = 'WLIST:WAVEFORM:DATA "%s",%d,' % (name, start)
    #     if size is not None:
    #         message = message + ('%d,' % size)
    #     points = points.clip(-1,1)
    #     values = (points * 0x1fff).astype(int) + 0x1fff
    #     self.write_binary_values(message, values, datatype=u'H',
    #                              is_big_endian=False,
    #                              termination=None, encoding=None)

    def _update_waveform_float(self,
                               points,
                               name='ABS',
                               IQ='I',
                               start=0,
                               size=None):
        message = 'WLIST:WAVEFORM:DATA:%s "%s",%d,' % (IQ, name, start)
        if size is not None:
            message = message + ('%d,' % size)
        values = points.clip(-1, 1)
        self.write_binary_values(message,
                                 values,
                                 datatype=u'f',
                                 is_big_endian=False,
                                 termination=None,
                                 encoding=None)

    # def update_marker(self, name, mk1, mk2=None, mk3=None, mk4=None, start=0, size=None):
    #     def format_marker_data(markers, bits):
    #         values = 0
    #         for i, v in markers:
    #             v = 0 if v is None else np.asarray(v)
    #             values += v << bits[i]
    #         return values
    #
    #     if self.model in ['AWG5014C']:
    #         values = format_marker_data([mk1, mk2], [5,6])
    #     elif self.model in ['AWG5208']:
    #         values = format_marker_data([mk1, mk2, mk3, mk4], [7,6,5,4])
    #     if size is None:
    #         message = 'WLIST:WAVEFORM:MARKER:DATA "%s",%d,' % (name, start)
    #     else:
    #         message = 'WLIST:WAVEFORM:MARKER:DATA "%s",%d,%d,' % (name, start, size)
    #     self.write_binary_values(message, values, datatype=u'B',
    #                              is_big_endian=False,
    #                              termination=None, encoding=None)

    def create_sequence(self, name, steps, tracks=1):
        if name in self.sequence_list:
            return
        self.write('SLIS:SEQ:NEW "%s", %d, %d; *WAI;' % (name, steps, tracks))
        self.sequence_list = self.get_sequence_list()

    def remove_sequence(self, name=None, all=False):
        if all:
            self.write('SLIS:SEQ:DEL ALL; *WAI;')
            self.sequence_list.clear()
        elif name not in self.sequence_list:
            return
        else:
            self.write('SLIS:SEQ:DEL "%s"; *WAI;' % name)
            self.sequence_list = self.get_sequence_list()

    #
    # def set_sequence_step(self, name, sub_name, step, wait='OFF', goto='NEXT', repeat=1, jump=None):
    #     """set a step of sequence
    #
    #     name: sequence name
    #     sub_name: subsequence name or list of waveforms for every tracks
    #     wait: ATRigger | BTRigger | ITRigger | OFF
    #     goto: <NR1> | LAST | FIRSt | NEXT | END
    #     repeat: ONCE | INFinite | <NR1>
    #     jump: a tuple (jump_input, jump_to)
    #         jump_input: ATRigger | BTRigger | OFF | ITRigger
    #         jump_to: <NR1> | NEXT | FIRSt | LAST | END
    #     """
    #     if isinstance(sub_name, str):
    #         self.write('SLIS:SEQ:STEP%d:TASS:SEQ "%s","%s"' % (step, name, sub_name))
    #     else:
    #         for i, wav in enumerate(sub_name):
    #             self.write('SLIS:SEQ:STEP%d:TASS%d:WAV "%s","%s"' % (step, i+1, name, wav))
    #     self.write('SLIS:SEQ:STEP%d:WINP "%s", %s' % (step, name, wait))
    #     self.write('SLIS:SEQ:STEP%d:GOTO "%s", %s' % (step, name, goto))
    #     self.write('SLIS:SEQ:STEP%d:RCO "%s", %s' % (step, name, repeat))
    #     if jump is not None:
    #         self.write('SLIS:SEQ:STEP%d:EJIN "%s", %s' % (step, name, jump[0]))
    #         self.write('SLIS:SEQ:STEP%d:EJUM "%s", %s' % (step, name, jump[1]))

    def use_sequence(self, name, ch=1, track=1, type=None):
        '''type: I or Q'''
        if type is not None:
            self.write('SOUR%d:CASS:SEQ "%s", %d, %s' %
                       (ch, name, track, type))
        else:
            self.write('SOUR%d:CASS:SEQ "%s", %d' % (ch, name, track))
        self.write('*WAI;')


"""

=== DESCRIPTION
This is a basic seq2seq translation system without any
bells or whistles. The model is in good shape, but this main
file is pretty horrid I admit. That's fine, though. If you run
it as-is, the model will overfit on a subset of the provided data.

=== USAGE
python main.py data/processed/ en vi tmp/checkpoint.ckpt-103 load

"""

from dataset import Dataset
from model import Seq2SeqV3
import sys

class config:
    src_vocab_size = 5000 + 1 # +1 for unk
    max_source_len = 50
    embedding_size = 64
    hidden_size = 128
    dropout_rate = 0.5
    num_layers = 3
    target_vocab_size = 5000 + 1 # +1 for unk
    max_target_len = 50
    learning_rate = 1.0    # sgd
#    learning_rate = 0.001  # adam

data_loc = sys.argv[1]
lang1 = sys.argv[2]
lang2 = sys.argv[3]
model_path = sys.argv[4]
restore = sys.argv[5] if len(sys.argv) == 6 else None

batch_size = 5
print 'building dataset...'
d = Dataset(data_loc, lang1, lang2)
'dataset done'
d.subset(6)    # take only 2k sentances

print d.get_start_token_indices()

c = config()

# if restore is not None:
#     print 'restoring model...'
#     model = Seq2Seq(c, batch_size, testing=True, model_path=model_path)
#     print 'model restored.'
# else:
#     print 'no model found. building model...'
#     model = Seq2Seq(c, batch_size)
#     print 'model built.'
print 'building model...'
model = Seq2SeqV3(c, batch_size, d, testing=False)
print 'model built.'

#batch = d.next_batch(batch_size)    # extract x's
#pred, logits = model.predict_on_batch(*batch[:2])
#print batch[2]
#print pred
#print logits
#d.reset()
#quit()
import numpy as np

print 'training...'
lr = c.learning_rate
for epoch in range(20000):
    epoch_loss = 0.0
    i = 0
    while d.has_next_batch(batch_size):
        batch = d.next_batch(batch_size)
        pred, loss = model.train_on_batch(*batch, learning_rate=lr)
        epoch_loss += loss
        i += 1
        if i % 750 == 0:
            lr = lr * 0.5

#        print batch[2]
#        print pred
#        print
#        print
    print 'epoch', epoch, 'loss', (epoch_loss / (i * batch_size))
    print batch[2]
    print pred
    d.reset()

d.reset()

batch = d.next_batch(batch_size)    # extract x's
pred, _ = model.predict_on_batch(*batch)

print batch[2]
print pred

quit()

####################################################
# FANCIER TRAINING STUFF THAT I SCRAPPED DURING DEBUGGING IS AFTER HERE
####################################################

try:
    print 'training...'
    for epoch in range(20):
        epoch_loss = 0.0
        i = 0
        while d.has_next_batch(batch_size):
            loss = model.train_on_batch(*d.next_batch(batch_size))
            epoch_loss += loss
            i += 1
        print 'epoch={}\t mean batch loss={:.4f}'.format(epoch, epoch_loss / (i * batch_size))
        d.reset()

except KeyboardInterrupt:
    print 'saving model to', model_path
    save_path = model.save('./%s' % model_path)
    print 'saved at ', save_path

finally:
        print 'testing...'
        d.reset()

        test_batch = d.next_batch(batch_size)
        pred_indices = model.predict_on_batch(*test_batch)

        for (y, y_hat) in zip(test_batch[0], pred_indices):
            print d.reconstruct(y, lang2)
            print d.reconstruct(y_hat, lang2)


import numpy as np
import argparse, time, os

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.optim import lr_scheduler
import torch.optim as optim

from datasets import BalancedBatchSampler, DeepFashionDataset
from networks import ResNetbasedNet
from losses import OnlineTripletLoss
from utils import read_data, pdist, RandomNegativeTripletSelector, HardestNegativeTripletSelector
from metrics import AverageNonzeroTripletsMetric
from trainer_clf import fit

def main(args):
    model_save_dir = args.save_dir
    vec_dim = 128

    data_type = ['train', 'validation']
    img_list, base_path, item_dict = read_data("DeepFashion2", bbox_gt=True, type_list=data_type)
    clf_inst_num = len(set(img_list['train'][:, 1]))
    clf_cate_num = len(set(img_list['train'][:, 2]))

    model = ResNetbasedNet(vec_dim=vec_dim, max_pool=True, clf1_num=clf_inst_num, clf2_num=clf_cate_num)
    # model.load_state_dict(torch.load('/home/jayeon/Documents/siamese-triplet/model_clf/00006.pth'))
    # model.load_state_dict(torch.load(args.model_path))
    # model.fc = nn.Linear(2048, vec_dim)
    # model.load_state_dict(torch.load('./model_adap_hardes_continue_norm69/00005.pth'))
    # model.load_state_dict(torch.load(args.model_path))

    # domain_adap = args.domain_adap
    is_cud = torch.cuda.is_available()
    device = torch.device("cuda" if is_cud else "cpu")
    if is_cud:
        if torch.cuda.device_count() > 1:
            model = nn.DataParallel(model)
        model.to(device)
    kwargs = {'num_workers': 8, 'pin_memory': True} if is_cud else {}

    if args.phase == 'train':
        train_dataset = DeepFashionDataset(img_list['train'], root=base_path, augment=True, clf=True)
        online_train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=256, shuffle=True, **kwargs)

        test_dataset = DeepFashionDataset(img_list['validation'], root=base_path, clf=True)
        online_test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=256, shuffle=False, **kwargs)

        criterion1 = nn.CrossEntropyLoss()
        criterion2 = nn.CrossEntropyLoss()
        grad_norm = 1.
        optimizer = optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)
        nn.utils.clip_grad_norm_(model.parameters(), grad_norm)
        scheduler = lr_scheduler.StepLR(optimizer, 4, gamma=0.9, last_epoch=-1)
        n_epochs = 500
        log_interval = 200

        fit(online_train_loader, online_test_loader, model, criterion1, criterion2, optimizer, scheduler, n_epochs,
            is_cud, log_interval, model_save_dir, start_epoch=0)

    else:
        with torch.no_grad():
            model.eval()
            test_dataset = DeepFashionDataset(img_list['validation'], root=base_path)
            test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=256, shuffle=False, num_workers=4)
            cate_preds = np.zeros(len(test_dataset))
            labels = np.zeros(len(test_dataset))
            top_k = 500
            idx_ = 0
            start_time = time.time()
            for idx, (data, target, _, _) in enumerate(test_loader):
                label, cate = model(data.cuda())
                cate_preds[idx_:idx_ + len(data)] = torch.argmax(cate, dim=1).cpu().numpy()
                labels[idx_:idx_ + len(data)] = np.asarray(target)
                idx_ += len(data)
                if idx % 20 == 0:
                    print(
                        'processing {}/{}... elapsed time {}s'.format(idx + 1, len(test_loader), time.time() - start_time))
        np.save(os.path.join(args.save_dir, 'cate_pred.npy'), cate_preds)
        with open(os.path.join(args.save_dir, 'file_info.txt'), 'w') as f:
            for i in range(len(test_dataset)):
                f.write('{},{},{},{}\n'.format(img_list['validation'][i][0], test_dataset[i][1], test_dataset[i][2], test_dataset[i][3]))
        print('save files!')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='image retrieval system')
    parser.add_argument('-d', '--dataset', required=True, help='DeepFashion / DeepFashion2')
    # parser.add_argument('-b', '--bbox', required=True, help='gt / pred')
    parser.add_argument('--phase', required=False, default='train', help='Train or Inference')
    parser.add_argument('--domain', dest='domain_adap', required=False, default=False, help='Train or Inference')
    parser.add_argument('-mp', '--model_path', dest='model_path',  required=False)
    parser.add_argument('-s', dest='save_dir', required=False, help='Directory to save the results')
    args = parser.parse_args()

    main(args)


import os
import zipfile
import gzip
import shutil
import numpy as np
import nltk
import fasttext
import gensim
import urllib.request
import networkx as nx
nltk.download('punkt') # one time execution
from nltk.tokenize import sent_tokenize
from lexrank import LexRank
from gensim.models.callbacks import CallbackAny2Vec
from summa.summarizer import summarize as summa_summarizer
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfVectorizer

example_text = """Mae Erthygl 25 o Ddatganiad Cyffredinol Hawliau Dynol 1948 y Cenhedloedd Unedig yn nodi: "Mae gan bawb yr hawl i safon byw sy'n ddigonol ar gyfer iechyd a lles ei hun a'i deulu, gan gynnwys bwyd, dillad, tai a gofal meddygol a gwasanaethau cymdeithasol angenrheidiol. "Mae'r Datganiad Cyffredinol yn cynnwys lletyaeth er mwyn diogelu person ac mae hefyd yn sôn yn arbennig am y gofal a roddir i'r rheini sydd mewn mamolaeth neu blentyndod. Ystyrir mai Datganiad Cyffredinol o Hawliau Dynol fel y datganiad rhyngwladol cyntaf o hawliau dynol sylfaenol. Dywedodd Uchel Gomisiynydd y Cenhedloedd Unedig dros Hawliau Dynol Navanethem Pillay fod y Datganiad Cyffredinol o Hawliau Dynol "yn ymgorffori gweledigaeth sy'n gofyn am gymryd yr holl hawliau dynol - sifil, gwleidyddol, economaidd, cymdeithasol neu ddiwylliannol - fel cyfanwaith anwahanadwy ac organig, anwahanadwy a rhyngddibynnol."""
example_summary = """Mae Datganiad Cyffredinol Hawliau Dynol 1948 yn dweud bod gan bawb yr hawl i safon byw digonol. Mae hynny yn cynnwys mynediad at fwyd a dillad a gofal meddygol i bob unigolyn. Dyma’r datganiad cyntaf o hawliau dynol"""

## Define summarizer models

# bottomline: return 1st sentence
def first_sent_summarize(article):
  return sent_tokenize(article)[0]

# lex_rank
def lex_rank_summarize(article):
  sentences = sent_tokenize(article)
  summary = LexRank(sentences).get_summary(sentences,
                              summary_size=int(len(sentences)/2), threshold=.1)
  return "\n".join(summary)

# text_rank
def text_rank_summarize(article):
  return summa_summarizer(article, ratio=0.5) #text

## Define Topline summarizers

def tfidf_summarize(article, ratio=0.5):
  sentences = sent_tokenize(article)
  # get similarity matrix
  sim_mat = cosine_similarity(TfidfVectorizer().fit_transform(sentences))
  scores = nx.pagerank_numpy(nx.from_numpy_array(sim_mat))
  top_ranked = sorted(scores.items(), key=lambda x: x[1],
                      reverse=True)[:int(len(scores)*ratio)]
  summary = [sentences[i] for i,_ in top_ranked]
  return "\n".join(summary)

# generate sentence vectors
def getSentenceVectors(embeddings, sents) -> list:
  sent_vectors = []
  for sent in sents:
    if len(sent) != 0:
      vecs = []
      for w in sent.split():
        try:
          vecs.append(embeddings[w])
        except AttributeError:
          vecs.append(np.zeros((300,)))
      sent_vectors.append(np.mean(vecs, axis=0).reshape(1,300))
  return sent_vectors

def embed_summarize(article, embedding, ratio=0.5):
  sentences = sent_tokenize(article)
  sentence_vectors = getSentenceVectors(embedding, sentences)
  nx_graph = nx.from_numpy_array(gen_similarity_matrix(sentence_vectors))
  scores = nx.pagerank_numpy(nx_graph)
  ranked_sentences = sorted(((scores[i],s) for i,s in enumerate(sentences)), reverse=True)
  return "\n".join([s for _, s in ranked_sentences[:int(len(ranked_sentences)*ratio)]])

# build similarity matrix
def gen_similarity_matrix(sents):
  sim_mat = np.zeros([len(sents), len(sents)])
  for i in range(len(sents)):
    for j in range(len(sents)):
      if i != j:
        sim_mat[i][j] = float(cosine_similarity(sents[i], sents[j]))
        # sim_mat[i][j] = float(cosine_similarity(sents[i].reshape(1,300),
        #                                         sents[j].reshape(1,300)))
  return sim_mat

class EpochLogger(CallbackAny2Vec):
    """
    Callback to log information about training

    ** THIS MUST BE DEFINED BEFORE BEING ABLE TO LOAD IN MODELS **
    """
    def __init__(self):
        self.epoch = 0

    def on_epoch_begin(self, model):
        print("Epoch #{} start".format(self.epoch))

    def on_epoch_end(self, model):
        print("Epoch #{} end".format(self.epoch))
        self.epoch += 1

# Unzipping files
def unzip(zipfilename):
  try:
    if zipfilename.endswith(".zip"):
      with zipfile.ZipFile(zipfilename, 'r') as zip_ref:
        zip_ref.extractall(zipfilename[:-4])
        print(f"'{zipfilename}' unzipped!")
    elif zipfilename.endswith(".gz"):
      with gzip.open(zipfilename, 'rb') as f_in:
        with open(zipfilename[:-3], 'wb') as f_out:
            shutil.copyfileobj(f_in, f_out)
    else:
      print(f"Invalid file name: {zipfilename}! Must be a '.zip' or a '.gz'")
    os.remove(zipfilename)
  except FileNotFoundError:
    print(f"Cannot find '{zipfilename}' file")


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

import copy
import os
import re

from absl import app
from absl import flags
from absl import logging
import numpy as np

import tensorflow as tf
from tensorflow._api.v2 import data

import cflags
import datasets
import effnetv2_configs
import effnetv2_model
import hparams
import utils
import wandb
from wandb.keras import WandbCallback

import tlc
from tlc.keras import TLCCallback

tlc.init("tlc_dir")

# 1. Start a new run
wandb.init(project='kostya-playground', entity='bluware')

# 2. Save model inputs and hyperparameters
config = wandb.config
# config.learning_rate = 0.01

FLAGS = flags.FLAGS

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  try:
    for gpu in gpus:
      # tf.config.experimental.set_memory_growth(gpu, True)
      tf.config.experimental.set_virtual_device_configuration(gpu, [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=8024)])
  except RuntimeError as e:
    print(e)

def build_tf2_optimizer(learning_rate,
                        optimizer_name='rmsprop',
                        decay=0.9,
                        epsilon=0.001,
                        momentum=0.9):
  """Build optimizer."""
  if optimizer_name == 'sgd':
    logging.info('Using SGD optimizer')
    optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate)
  elif optimizer_name == 'momentum':
    logging.info('Using Momentum optimizer')
    optimizer = tf.keras.optimizers.SGD(
        learning_rate=learning_rate, momentum=momentum)
  elif optimizer_name == 'rmsprop':
    logging.info('Using RMSProp optimizer')
    optimizer = tf.keras.optimizers.RMSprop(learning_rate, decay, momentum,
                                            epsilon)
  elif optimizer_name == 'adam':
    logging.info('Using Adam optimizer')
    optimizer = tf.keras.optimizers.Adam(learning_rate)
  else:
    raise Exception('Unknown optimizer: %s' % optimizer_name)

  return optimizer

class TrainableModel(effnetv2_model.EffNetV2Model):
  """Wraps efficientnet to make a keras trainable model.

  Handles efficientnet's multiple outputs and adds weight decay.
  """

  def __init__(self,
               model_name='efficientnetv2-s',
               model_config=None,
               name=None,
               weight_decay=0.0):
    super().__init__(
        model_name=model_name,
        model_config=model_config,
        name=name or model_name)

    self.weight_decay = weight_decay

  def _reg_l2_loss(self, weight_decay, regex=r'.*(kernel|weight):0$'):
    """Return regularization l2 loss loss."""
    var_match = re.compile(regex)
    return weight_decay * tf.add_n([
        tf.nn.l2_loss(v)
        for v in self.trainable_variables
        if var_match.match(v.name)
    ])

  def train_step(self, data):
    features, labels = data
    images, labels = features['image'], labels['label']

    with tf.GradientTape() as tape:
      pred = self(images, training=True)
      pred = tf.cast(pred, tf.float32)
      loss = self.compiled_loss(
          labels,
          pred,
          regularization_losses=[self._reg_l2_loss(self.weight_decay)])

    self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
    self.compiled_metrics.update_state(labels, pred)
    return {m.name: m.result() for m in self.metrics}

  def test_step(self, data):
    features, labels = data
    images, labels = features['image'], labels['label']
    pred = self(images, training=False)
    pred = tf.cast(pred, tf.float32)

    self.compiled_loss(
        labels,
        pred,
        regularization_losses=[self._reg_l2_loss(self.weight_decay)])

    self.compiled_metrics.update_state(labels, pred)
    return {m.name: m.result() for m in self.metrics}

def main(_) -> None:
  config = copy.deepcopy(hparams.base_config)
  config.override(effnetv2_configs.get_model_config(FLAGS.model_name))
  config.override(datasets.get_dataset_config(FLAGS.dataset_cfg))
  config.override(FLAGS.hparam_str)
  config.model.num_classes = config.data.num_classes
  strategy = config.runtime.strategy
  if strategy == 'tpu' and not config.model.bn_type:
    config.model.bn_type = 'tpu_bn'

  # log and save config.
  logging.info('config=%s', str(config))
  if 'train' in FLAGS.mode:
    if not tf.io.gfile.exists(FLAGS.model_dir):
      tf.io.gfile.makedirs(FLAGS.model_dir)
    config.save_to_yaml(os.path.join(FLAGS.model_dir, 'config.yaml'))

  if strategy == 'tpu':
    tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
        FLAGS.tpu, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
    tf.config.experimental_connect_to_cluster(tpu_cluster_resolver)
    tf.tpu.experimental.initialize_tpu_system(tpu_cluster_resolver)
    ds_strategy = tf.distribute.TPUStrategy(tpu_cluster_resolver)
    logging.info('All devices: %s', tf.config.list_logical_devices('TPU'))
  elif strategy == 'gpus':
    ds_strategy = tf.distribute.MirroredStrategy()
    logging.info('All devices: %s', tf.config.list_physical_devices('GPU'))
  else:
    if tf.config.list_physical_devices('GPU'):
      ds_strategy = tf.distribute.MirroredStrategy(['GPU:0'])
    else:
      ds_strategy = tf.distribute.MirroredStrategy(['CPU:0'])

  with ds_strategy.scope():
    train_split = config.train.split or 'train'
    eval_split = config.eval.split or 'eval'
    num_train_images = config.data.splits[train_split].num_images
    num_eval_images = config.data.splits[eval_split].num_images

    train_size = config.train.isize
    eval_size = config.eval.isize
    if train_size <= 16.:
      train_size = int(eval_size * train_size) // 16 * 16

    image_dtype = None
    if config.runtime.mixed_precision:
      image_dtype = 'bfloat16' if strategy == 'tpu' else 'float16'
      precision = 'mixed_bfloat16' if strategy == 'tpu' else 'mixed_float16'
      policy = tf.keras.mixed_precision.Policy(precision)
      tf.keras.mixed_precision.set_global_policy(policy)

    model = TrainableModel(
        config.model.model_name,
        config.model,
        weight_decay=config.train.weight_decay)

    if config.train.ft_init_ckpt:  # load pretrained ckpt for finetuning.
      model(tf.keras.Input([None, None, 3]))
      ckpt = config.train.ft_init_ckpt
      utils.restore_tf2_ckpt(model, ckpt, exclude_layers=('_fc', 'optimizer'))

    steps_per_epoch = num_train_images // config.train.batch_size
    # steps_per_epoch = 300
    total_steps = steps_per_epoch * config.train.epochs

    scaled_lr = config.train.lr_base * (config.train.batch_size / 256.0)
    scaled_lr_min = config.train.lr_min * (config.train.batch_size / 256.0)
    learning_rate = utils.WarmupLearningRateSchedule(
        scaled_lr,
        steps_per_epoch=steps_per_epoch,
        decay_epochs=config.train.lr_decay_epoch,
        warmup_epochs=config.train.lr_warmup_epoch,
        decay_factor=config.train.lr_decay_factor,
        lr_decay_type=config.train.lr_sched,
        total_steps=total_steps,
        minimal_lr=scaled_lr_min)

    optimizer = build_tf2_optimizer(
        learning_rate, optimizer_name=config.train.optimizer)

    model.compile(
        optimizer=optimizer,
        loss=tf.keras.losses.CategoricalCrossentropy(
            label_smoothing=config.train.label_smoothing, from_logits=True),
        metrics=[
            tf.keras.metrics.TopKCategoricalAccuracy(k=1, name='acc_top1'),
            tf.keras.metrics.TopKCategoricalAccuracy(k=5, name='acc_top5')
        ],
    )

    ckpt_callback = tf.keras.callbacks.ModelCheckpoint(
        os.path.join(FLAGS.model_dir, 'ckpt-{epoch:d}'),
        verbose=1,
        save_weights_only=True)
    tb_callback = tf.keras.callbacks.TensorBoard(
        log_dir=FLAGS.model_dir, update_freq=100)
    rstr_callback = utils.ReuableBackupAndRestore(backup_dir=FLAGS.model_dir)

    def filter_callbacks(callbacks):
      if strategy == 'tpu' and not FLAGS.model_dir.startswith('gs://'):
        return list(filter(lambda callback: isinstance(callback, tf.keras.callbacks.ModelCheckpoint), callbacks))
      return callbacks

    def get_dataset(training, image_size, config):
      """A shared utility to get input dataset."""
      if training:
        return ds_strategy.distribute_datasets_from_function(
            datasets.build_dataset_input(
                True, image_size, image_dtype, FLAGS.data_dir, train_split,
                config.data).distribute_dataset_fn(config.train.batch_size))
      else:
        return ds_strategy.distribute_datasets_from_function(
            datasets.build_dataset_input(
                False, image_size, image_dtype, FLAGS.data_dir, eval_split,
                config.data).distribute_dataset_fn(config.eval.batch_size))

    if FLAGS.mode == 'traineval':

      # eval_results = model.evaluate(
      #       get_dataset(training=False, image_size=eval_size, config=config),
      #       batch_size=config.eval.batch_size,
      #       steps=num_eval_images // config.eval.batch_size,
      #       callbacks=filter_callbacks([tb_callback, rstr_callback, TLCCallback()]),
      #       verbose=2 if strategy == 'tpu' else 1,
      #   )
      prediction_batch_size = 128
      dataset = datasets.build_dataset_input(
                is_training=False, is_predict=True, image_size=train_size, image_dtype=image_dtype, data_dir=FLAGS.data_dir,
                split=train_split, data_config=config.data)
      prediction_dataset = dataset._input_fn(prediction_batch_size, 0, 1, name="Predict Images")

      # iter_tlc = iter(prediction_dataset)
      # while True:
      #   a_tfr = next(iter_tlc)
      #   rgb_tfr = a_tfr[0]["image"].numpy()
      #   label_tfr = a_tfr[1]["label"].numpy()
      #   for i in range(128):
      #     rgb = rgb_tfr[i]
      #     image_tfr = tf.keras.preprocessing.image.array_to_img(rgb)
      #     image_tfr.show()
      #     label = label_tfr[i]
      #     print(np.where(label == 1))

      model.fit(
          get_dataset(training=True, image_size=train_size, config=config),
          epochs=config.train.epochs,
          steps_per_epoch=steps_per_epoch,
          validation_data=get_dataset(
              training=False, image_size=eval_size, config=config),
          validation_steps=num_eval_images // config.eval.batch_size,
          callbacks=filter_callbacks([ckpt_callback, tb_callback, rstr_callback,
          WandbCallback(),
          TLCCallback(prediction_dataset)]),
          # don't log spam if running on tpus
          verbose=2 if strategy == 'tpu' else 1,
      )

      # loss_function = tf.keras.losses.CategoricalCrossentropy(reduction=tf.keras.losses.Reduction.NONE)

      # it = iter(dataset)
      # while it:
      #   features,labels  = next(it)
      #   images = features["image"]
      #   labels = labels["label"]
      #   predicted = model.predict_on_batch(images)
      #   a = loss_function(labels, predicted).numpy()
      #   # print(a)

      # it = iter(predicted)
      # itl = iter(labels)

      # def _gen():
      #   prediscted_label = next(it)
      #   label = next(itl)[0]
      #   yield prediscted_label, label

      # d = tf.data.Dataset.from_generator(_gen, output_types=(np.float32, np.float32), output_shapes=((10), (10)))
      # d.batch(20)

      # def calc(prediscted_label, label):
      #   l = loss(label, prediscted_label)
      #   return l

      # d = d.map(calc, num_parallel_calls=tf.data.experimental.AUTOTUNE)

      # itd = iter(d)
      # while True:
      #   print(next(itd))

      # while True:
      #   prediscted_label = next(it)
      #   label = next(itl).numpy()[0]

      #   l = loss(label, prediscted_label).numpy()
      #   print(l)

      # images = images.numpy()
      # n_images = images.shape[0]
      # loss = n_images*[None]

      # it = iter(images)
      # itl = iter(labels)
      # while True:
      #   image = next(it).numpy()
      #   label = next(itl).numpy()
      #   l = model.evaluate(x=image, y=label,
      #                             batch_size=1,
      #                             verbose=1,
      #                             steps=1
      #                             )

      # it = iter(dataset)
      # while True:
      #   b = next(it)
      #   i = b[0]["image"]
      #   l = b[1]["label"]

      #   def _gen():
      #       yield i, l

      #   d = tf.data.Dataset.from_generator(_gen, output_types=(np.float16, np.float32), output_shapes=((None, None, 3), (10)))

      #   def p(image, onehot_label):
      #     features={"image": image}
      #     labels = {"label": onehot_label}

      #     return features, labels

      #   d = d.map(p,
      #     num_parallel_calls=tf.data.experimental.AUTOTUNE)

      #   # d = tf.data.Dataset.from_tensor_slices([i, l ])

      #   # def gen():
      #   #   yield b

      #   # d = tf.data.Dataset.from_generator(gen)
      #   l = model.evaluate(d,
      #                             batch_size=1,
      #                             verbose=1,
      #                             steps=1
      #                             )
      #   print(l)

    elif FLAGS.mode == 'train':
      if not config.train.stages:
        model.fit(
            get_dataset(training=True, image_size=train_size, config=config),
            epochs=config.train.epochs,
            steps_per_epoch=steps_per_epoch,
            callbacks=filter_callbacks([ckpt_callback, tb_callback, rstr_callback,
            WandbCallback(),
            TLCCallback()]),
            verbose=2 if strategy == 'tpu' else 1,
        )
      else:
        total_stages = config.train.stages
        ibase = config.data.ibase or (train_size / 2)

        if config.train.sched:
          ram_list = np.linspace(5, config.data.ram, total_stages)
          mixup_list = np.linspace(0, config.data.mixup_alpha, total_stages)
          cutmix_list = np.linspace(0, config.data.cutmix_alpha, total_stages)

        for stage in range(total_stages):
          ratio = float(stage + 1) / float(total_stages)
          start_epoch = int(
              float(stage) / float(total_stages) * config.train.epochs)
          end_epoch = int(ratio * config.train.epochs)
          image_size = int(ibase + (train_size - ibase) * ratio)

          if config.train.sched:
            config.data.ram = ram_list[stage]
            config.data.mixup_alpha = mixup_list[stage]
            config.data.cutmix_alpha = cutmix_list[stage]

          model.fit(
              get_dataset(training=True, image_size=image_size, config=config),
              initial_epoch=start_epoch,
              epochs=end_epoch,
              steps_per_epoch=steps_per_epoch,
              callbacks=filter_callbacks([ckpt_callback, tb_callback, rstr_callback,
              WandbCallback(),
              TLCCallback()]),
              verbose=2 if strategy == 'tpu' else 1,
          )
    elif FLAGS.mode == 'eval':
      for ckpt in tf.train.checkpoints_iterator(
          FLAGS.model_dir, timeout=60 * 60 * 24):
        model.load_weights(ckpt)
        eval_results = model.evaluate(
            get_dataset(training=False, image_size=eval_size, config=config),
            batch_size=config.eval.batch_size,
            steps=num_eval_images // config.eval.batch_size,
            callbacks=filter_callbacks([tb_callback, rstr_callback]),
            verbose=2 if strategy == 'tpu' else 1,
        )

        try:
          current_epoch = int(os.path.basename(ckpt).split('-')[1])
        except IndexError:
          logging.info('%s has no epoch info: stop!', ckpt)
          break

        logging.info('Epoch: %d, total %d', current_epoch, config.train.epochs)
        if current_epoch >= config.train.epochs:
          break
    else:
      raise ValueError(f'Invalid mode {FLAGS.mode}')

if __name__ == '__main__':
  cflags.define_flags()
  app.run(main)


import datetime as dt
from typing import List
import numpy as np
import pandas as pd
import requests
import urllib, json
from finrl_meta.data_processors.basic_processor import BasicProcessor
from .func import download_n_unzip_file, convert_to_date_object, get_path

class BinanceProcessor(BasicProcessor):
    def __init__(self, data_source: str, **kwargs):
        BasicProcessor.__init__(self, data_source, **kwargs)
        self.url = "https://api.binance.com/api/v3/klines"

    #main functions
    def download_data(self, ticker_list: List[str], start_date: str, end_date: str, time_interval: str):
        startTime = dt.datetime.strptime(start_date, '%Y-%m-%d')
        endTime = dt.datetime.strptime(end_date, '%Y-%m-%d')

        self.start_time = self.stringify_dates(startTime)
        self.end_time = self.stringify_dates(endTime)
        self.interval = time_interval
        self.limit = 1440

        #1s for now, will add support for variable time and variable tick soon
        if time_interval == "1s":
            # as per https://binance-docs.github.io/apidocs/spot/en/#compressed-aggregate-trades-list
            self.limit = 1000
            final_df = self.fetch_n_combine(start_date, end_date, ticker_list)
        else:
            final_df = pd.DataFrame()
            for i in ticker_list:
                hist_data = self.dataframe_with_limit(symbol=i)
                df = hist_data.iloc[:-1].dropna()
                df['tic'] = i
                final_df = final_df.append(df)
        self.dataframe = final_df

    # def clean_data(self, df):
    #     df = df.dropna()
    #     return df

    # def add_technical_indicator(self, df, tech_indicator_list):
    #     print('Adding self-defined technical indicators is NOT supported yet.')
    #     print('Use default: MACD, RSI, CCI, DX.')
    #     self.tech_indicator_list = ['open', 'high', 'low', 'close', 'volume',
    #                                 'macd', 'macd_signal', 'macd_hist',
    #                                 'rsi', 'cci', 'dx']
    #     final_df = pd.DataFrame()
    #     for i in df.tic.unique():
    #         tic_df = df[df.tic==i]
    #         tic_df['macd'], tic_df['macd_signal'], tic_df['macd_hist'] = MACD(tic_df['close'], fastperiod=12,
    #                                                                             slowperiod=26, signalperiod=9)
    #         tic_df['rsi'] = RSI(tic_df['close'], timeperiod=14)
    #         tic_df['cci'] = CCI(tic_df['high'], tic_df['low'], tic_df['close'], timeperiod=14)
    #         tic_df['dx'] = DX(tic_df['high'], tic_df['low'], tic_df['close'], timeperiod=14)
    #         final_df = final_df.append(tic_df)
    #
    #     return final_df

    # def add_turbulence(self, df):
    #     print('Turbulence not supported yet. Return original DataFrame.')
    #
    #     return df

    # def add_vix(self, df):
    #     print('VIX is not applicable for cryptocurrencies. Return original DataFrame')
    #
    #     return df

    # def df_to_array(self, df, tech_indicator_list, if_vix):
    #     unique_ticker = df.tic.unique()
    #     price_array = np.column_stack([df[df.tic==tic].close for tic in unique_ticker])
    #     tech_array = np.hstack([df.loc[(df.tic==tic), tech_indicator_list] for tic in unique_ticker])
    #     assert price_array.shape[0] == tech_array.shape[0]
    #     return price_array, tech_array, np.array([])

    # helper functions
    def stringify_dates(self, date:dt.datetime):
        return str(int(date.timestamp()*1000))

    def get_binance_bars(self, last_datetime, symbol):
        '''
        klines api returns data in the following order:
        open_time, open_price, high_price, low_price, close_price,
        volume, close_time, quote_asset_volume, n_trades,
        taker_buy_base_asset_volume, taker_buy_quote_asset_volume,
        ignore
        '''
        req_params = {"symbol": symbol, 'interval': self.interval,
                      'startTime': last_datetime, 'endTime': self.end_time,
                      'limit': self.limit}
        # For debugging purposes, uncomment these lines and if they throw an error
        # then you may have an error in req_params
        # r = requests.get(self.url, params=req_params)
        # print(r.text)
        df = pd.DataFrame(requests.get(self.url, params=req_params).json())

        if df.empty:
            return None

        df = df.iloc[:, 0:6]
        df.columns = ['datetime','open','high','low','close','volume']

        df[['open','high','low','close','volume']] = df[['open','high','low','close','volume']].astype(float)

        # No stock split and dividend announcement, hence adjusted close is the same as close
        df['adj_close'] = df['close']
        df['datetime'] = df.datetime.apply(lambda x: dt.datetime.fromtimestamp(x/1000.0))
        df.reset_index(drop=True, inplace=True)

        return df

    def dataframe_with_limit(self, symbol):
        final_df = pd.DataFrame()
        last_datetime = self.start_time
        while True:
            new_df = self.get_binance_bars(last_datetime, symbol)
            if new_df is None:
                break
            final_df = final_df.append(new_df)
            last_datetime = max(new_df.datetime) + dt.timedelta(days=1)
            last_datetime = self.stringify_dates(last_datetime)

        date_value = final_df['datetime'].apply(lambda x: x.strftime('%Y-%m-%d %H:%M:%S'))
        final_df.insert(0, 'time', date_value)
        final_df.drop('datetime', inplace=True, axis=1)
        return final_df

    #helpers for manipulating tick level data (1s intervals)
    def download_daily_aggTrades(self, symbols, num_symbols, dates, start_date, end_date):
        current = 0
        trading_type="spot"
        date_range = start_date + " " + end_date
        start_date = convert_to_date_object(start_date)
        end_date = convert_to_date_object(end_date)

        print("Found {} symbols".format(num_symbols))

        map = {}
        for symbol in symbols:
            map[symbol] = []
            print("[{}/{}] - start download daily {} aggTrades ".format(current+1, num_symbols, symbol))
            for date in dates:
                current_date = convert_to_date_object(date)
                if current_date >= start_date and current_date <= end_date:
                    path = get_path(trading_type, "aggTrades", "daily", symbol)
                    file_name = "{}-aggTrades-{}.zip".format(symbol.upper(), date)
                    fhandle = download_n_unzip_file(path, file_name, date_range)
                    map[symbol] += fhandle
            current += 1
        return map

    def fetch_aggTrades(self, startDate: str, endDate:str, tickers: List[str]):
        #all valid symbols traded on v3 api
        response = urllib.request.urlopen("https://api.binance.com/api/v3/exchangeInfo").read()
        valid_symbols = list(map(lambda symbol: symbol['symbol'], json.loads(response)['symbols']))

        for tic in tickers:
            if tic not in valid_symbols:
                print(tic+" not a valid ticker, removing from download")
        tickers = list(set(tickers) & set(valid_symbols))
        num_symbols = len(tickers)
        #not adding tz yet
        #for ffill missing data on starting on first day 00:00:00 (if any)
        tminus1 = (convert_to_date_object(startDate) - dt.timedelta(1)).strftime('%Y-%m-%d')
        dates = pd.date_range(start=tminus1, end=endDate)
        dates = [date.strftime("%Y-%m-%d") for date in dates]
        return self.download_daily_aggTrades(tickers, num_symbols, dates, tminus1, endDate)

    #Dict[str]:List[str] -> pd.DataFrame
    def combine_raw(self, map):
        #same format as jingyang's current data format
        final_df = pd.DataFrame()
        # using AggTrades with headers from https://github.com/binance/binance-public-data/
        colNames = ["AggregatetradeId","Price","volume","FirsttradeId","LasttradeId","time","buyerWasMaker","tradeWasBestPriceMatch"]
        for tic in map.keys():
            security = pd.DataFrame()
            for i, csv in enumerate(map[tic]):
                dailyticks = pd.read_csv(csv,
                                names=colNames,
                                index_col=["time"],
                                parse_dates=['time'],
                                date_parser=lambda epoch: pd.to_datetime(epoch, unit='ms'))
                dailyfinal = dailyticks.resample('1s').agg({'Price': 'ohlc', 'volume': 'sum'})
                dailyfinal.columns = dailyfinal.columns.droplevel(0)
                # favor continuous series
                #dailyfinal.dropna(inplace=True)

                # implemented T-1 day ffill day start missing values
                # guaranteed first csv is tminus1 day
                if i == 0:
                    tmr = dailyfinal.index[0].date() + dt.timedelta(1)
                    tmr_dt = dt.datetime.combine(tmr, dt.time.min)
                    last_time_stamp_dt = dailyfinal.index[-1].to_pydatetime()
                    s_delta = (tmr_dt-last_time_stamp_dt).seconds
                    lastsample = dailyfinal.iloc[-1:]
                    lastsample.index = lastsample.index.shift(s_delta, 's')
                else:
                    day_dt = dailyfinal.index[0].date()
                    day_str = day_dt.strftime("%Y-%m-%d")
                    nextday_str = (day_dt + dt.timedelta(1)).strftime("%Y-%m-%d")
                    if dailyfinal.index[0].second != 0:
                        #append last sample
                        dailyfinal = lastsample.append(dailyfinal)
                    #otherwise, just reindex and ffill
                    dailyfinal = dailyfinal.reindex(pd.date_range(day_str, nextday_str, freq="1s")[:-1], method='ffill')
                    #save reference info (guaranteed to be :59)
                    lastsample = dailyfinal.iloc[-1:]
                    lastsample.index = lastsample.index.shift(1, 's')

                    if dailyfinal.shape[0] != 86400:
                        raise ValueError("everyday should have 86400 datapoints")

                    #only save real startDate - endDate
                    security = security.append(dailyfinal)

            security.ffill(inplace=True)
            security['tic'] = tic
            final_df = final_df.append(security)
        return final_df

    def fetch_n_combine(self, startDate, endDate, tickers):
        #return combine_raw(fetchAggTrades(startDate, endDate, tickers))
        mapping = self.fetch_aggTrades(startDate, endDate, tickers)
        combined_df = self.combine_raw(mapping)
        return combined_df


import argparse
import time
import requests
import numpy as np
import logging
import string
import random
from math import exp

from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from pyvirtualdisplay import Display

logging.basicConfig(level=logging.INFO)

parser = argparse.ArgumentParser(description='HTTPS-Client Simulation.')

parser.add_argument('-ip', dest='server_ip', action='store', type=str, required=True,
                    help='The IP address of the target server')
parser.add_argument('-u', dest='username', action='store',
                    type=str, required=True, help='The username of the client')
parser.add_argument('-p', dest='password', action='store',
                    type=str, required=True, help='The password of the client')

args = parser.parse_args()

# Disable requests warnings (caused by self signed server certificate)
requests.packages.urllib3.disable_warnings()

# Virtual display to run chrome-browser
display = Display(visible=0, size=(800, 800))
display.start()

# Headless chrome-browser settings
chrome_options = Options()
chrome_options.add_argument("--no-sandbox")
chrome_options.add_argument('--ignore-certificate-errors')
browser = webdriver.Chrome(chrome_options=chrome_options)

# probability to enter valid input into the sql_injection form
sqli_valid_input_probability = 0.75
# if no valid value is send, whats the probability of it been a random string
sqli_random_string_probability = 0.5

# states:
logged_off = 0
logged_in = 1

# initially the client is logged off
client_state = logged_off

# links to exclude when randomly choosing one
exclude_links = ['/logout.php', '/captcha/',
                 '/security.php', '/phpinfo.php', '.pdf']

def generate_random_string(size=6, chars=string.printable):
    """
    generates a random printable string of the given size using the given characters
    :param size: length of the string to be generated
    :param chars: set of characters to be used
    :return: a random string
    """
    return ''.join(random.choice(chars) for _ in range(size))

def wait_for_dvwa():
    done = False
    logging.info('wait till dvwa is up und running...')
    while not done:
        # wait till dvwa is up und running
        url = 'http://' + args.server_ip + '/login.php'
        logging.info("get url: " + url)
        browser.get(url)
        if 'Login :: Damn Vulnerable Web Application' in browser.title:
            done = True
            logging.info('dvwa is ready...')
        else:
            logging.info(
                'dvwa not ready yet, waiting 2 seconds, current title: \n' + browser.title)
            time.sleep(2)

def random_browsing():
    """
    main loop of random browsing normal behaviour
    """
    wait_for_dvwa()
    # now start the "normal routine"
    scheduler_sync(normal_step)

def normal_step(arg):
    global client_state
    # is the client logged in?
    # yes, its logged in
    if client_state == logged_in:
        # do some random stuff, but never log out
        follow_link()
        do_things()
    # no, its not logged in:
    else:
        log_in()

def do_things():
    """
    do things depending on the current active site
    - "vulnerabilities/sqli/" in path:
        <input type="text" size="15" name="id">
        <input type="submit" name="Submit" value="Submit">
        Values 1 to 5 are ok
    """
    if 'vulnerabilities/sqli/' in browser.current_url:
        logging.info(' SQL injection form found')
        sending_value = ''
        if np.random.random() <= sqli_valid_input_probability:
            # enter a valid integer value (valid is from 1 to 5)
            sending_value = str(np.random.randint(low=1, high=6))
        else:
            if np.random.random() <= sqli_random_string_probability:
                # enter a random string
                sending_value = generate_random_string(
                    np.random.randint(low=1, high=40))
            else:
                # enter a "not" valid number
                sending_value = str(np.random.randint(low=-10000, high=10000))

        logging.info(' sending: %s', sending_value)
        browser.find_element_by_name('id').send_keys(sending_value)
        browser.find_element_by_name('Submit').click()

def follow_link():
    """
    iterates all "internal" links and chooses a random one to follow
    """
    # list all available local links
    logging.info('follow link...')
    link_list = list()
    for link in browser.find_elements_by_xpath('.//a'):
        link_url = link.get_attribute('href')
        if args.server_ip in link_url:
            if not any(e in link_url for e in exclude_links):
                link_list.append(link)
                logging.info('%d. %s', len(link_list)-1, link_url)
    if len(link_list) == 0:
        logging.info('    no link found, returning to index.php')
        browser.get('http://' + args.server_ip + '/index.php')
    else:
        # randomly select one link to follow
        i = np.random.randint(0, high=len(link_list))
        logging.info('    selected: [%d] of %d links', i, len(link_list))
        selected_link = link_list[i]
        selected_link.click()

def log_in():
    """
    logs into dvwa with the given username and password (args.username and args.password)
    changes client_state to logged_in
    """
    global client_state
    url = 'http://' + args.server_ip + '/login.php'
    logging.info('login... ' + url)
    browser.get(url)
    logging.info('    got response')
    browser.find_element_by_name('username').send_keys(args.username)
    browser.find_element_by_name('password').send_keys(args.password)
    logging.info('    filled form and click')
    browser.find_element_by_name('Login').click()
    logging.info('    logged in')
    client_state = logged_in

def uniform_K():
    """
    Returns a uniformly chosen k-parameter value from the Values recommended from
    >>Empirical Model of WWW Document Arrivals at Access Link<<
    """
    return np.random.uniform(
        low=.77,
        high=.91
    )

def uniform_THETA():
    """
    Returns a uniformly chosen θ-parameter value from the Values recommended from
    >>Empirical Model of WWW Document Arrivals at Access Link<<
    """
    return np.random.uniform(
        low=exp(4.4),
        high=exp(4.6)
    )

def scheduler_sync(fn):
    """
    synchronous scheduler -> waits until fn returns before calling the next one
    """

    while True:
        """
        block duration: weibul
        next block: pareto
        """
        on_time_scale, on_time_coefficient = uniform_THETA(), uniform_K()
        block_duration = on_time_scale * (np.random.weibull(on_time_coefficient))
        block_start_time = time.time()
        block_end_time = block_start_time + block_duration
        off_time = np.random.pareto(0.9)
        next_inner_time = block_start_time
        while True:
            # call fn with the current timestamp
            fn(time.time())
            # get the next inner block event timestamp
            next_inner_time = 1.5 * np.random.weibull(0.5) + next_inner_time
            # now wait until the next inner event time
            current_time = time.time()
            inner_sleep_time = next_inner_time - current_time
            # check for block end
            if current_time > block_end_time:
                break
            # check for next event time
            # do not sleep if it will end after the block end
            if inner_sleep_time > 0:
                if current_time + inner_sleep_time > block_end_time:
                    break
                time.sleep(inner_sleep_time)

        # sleep off time
        sleep_until = block_end_time + off_time
        time_to_sleep = sleep_until - time.time()
        if time_to_sleep < 0:
            time_to_sleep = 1
        time.sleep(time_to_sleep)

random_browsing()


import torch
import torch.nn
from maddpg_agent_changed import Agent
import numpy as np
import make_env as gym_env
from utils import visdom_line_plotter
import time
import os
import pickle

def create_agents(env, obs_shape_n):
    agents = []
    for i in range(0, env.n):
        agents.append(Agent("agent_%d" % i, state_size=obs_shape_n, action_size=5, agent_index=i,
                            random_seed=0))  # change after starts working
    return agents

def train(display):
    env = gym_env.make_env(scenario_name="simple_tag", benchmark=True)
    obs_shape_n = [env.observation_space[i].shape[0] for i in range(env.n)]
    agents = create_agents(env, obs_shape_n)
    # plotter = visdom_line_plotter.VisdomLinePlotter()

    plots_dir = './learning_curves/'
    benchmark_iters = 600000
    benchmark_dir = './benchmark_files_'
    exp_name = 'maddpg_pytorch'

    benchmark = True
    TAU = 1.0 - 1e-2
    max_episode_len = 25

    if display:
        for i in range(len(agents)):
            actor_local = agents[i].actor_local
            actor_local_ckpt = torch.load('./checkpoints/checkpoint_actor_{}.pth'.format(i), map_location='cpu')
            actor_local.load_state_dict(actor_local_ckpt)
            actor_target = agents[i].actor_target
            actor_target_ckpt = torch.load('./checkpoints/checkpoint_actor_target_{}.pth'.format(i), map_location='cpu')
            actor_target.load_state_dict(actor_target_ckpt)
            critic_local = agents[i].critic_local
            critic_local_ckpt = torch.load('./checkpoints/checkpoint_critic_{}.pth'.format(i), map_location='cpu')
            critic_local.load_state_dict(critic_local_ckpt)
            critic_target = agents[i].critic_target
            critic_target_ckpt = torch.load('./checkpoints/checkpoint_critic_target_{}.pth'.format(i),
                                            map_location='cpu')
            critic_target.load_state_dict(critic_target_ckpt)

    final_ep_rewards = []
    final_ep_ag_rewards = []
    episode_rewards = [0.0]
    agent_rewards = [[0.0] for _ in range(env.n)]
    agent_info = [[[]]]
    obs_n = env.reset()
    episode_step = 0
    train_step = 0
    t_start = time.time()

    print('Starting iterations...')
    while True:
        # get action
        for agent in agents:
            agent.reset()
        action_n = [agent.act(obs, add_noise=True) for agent, obs in zip(agents, obs_n)]

        # environment step
        new_obs_n, rew_n, done_n, info_n = env.step(action_n)
        episode_step += 1
        done = all(done_n)
        terminal = (episode_step >= max_episode_len)
        # collect experience
        for i, agent in enumerate(agents):
            agent.experience(obs_n[i], action_n[i], rew_n[i], new_obs_n[i], done_n[i], terminal)
        obs_n = new_obs_n
        # print(rew_n)
        for i, rew in enumerate(rew_n):
            episode_rewards[-1] += rew
            agent_rewards[i][-1] += rew

        if done or terminal:
            obs_n = env.reset()
            episode_step = 0
            episode_rewards.append(0)
            # print(episode_rewards)
            for a in agent_rewards:
                a.append(0)
            agent_info.append([[]])

        # increment global step counter
        train_step += 1
        # for benchmarking learned policies
        if benchmark:
            for i, info in enumerate(info_n):
                agent_info[-1][i].append(info_n['n'])
            if train_step > benchmark_iters and (done or terminal) and (len(episode_rewards) % 1000 == 0):
                file_name = benchmark_dir + exp_name + '.pkl'
                print('Finished benchmarking, now saving...')
                with open(file_name, 'wb') as fp:
                    pickle.dump(agent_info[:-1], fp)
                # break
            # continue

        # for displaying learned policies
        if display:
            time.sleep(0.1)
            env.render()
            continue

        # update all trainers, if not in display or benchmark mode
        loss = None
        for agent in agents:
            agent.preupdate()
        for agent in agents:
            loss = agent.step(agents, train_step, terminal)

        # save model, display training output
        if terminal and (len(episode_rewards) % 1000 == 0):  # 25 and 1000

            print("steps: {}, episodes: {}, mean episode reward: {}, time: {}".format(
                train_step, len(episode_rewards), np.mean(episode_rewards[-1000:]), round(time.time() - t_start, 3)))
            # plotter.plot('Episode Rewards', 'Rewards', 'Training', len(episode_rewards), np.mean(episode_rewards[-1000:]))
            i = 0
            for agt in agents:
                torch.save(agt.actor_local.state_dict(), './checkpoints/checkpoint_actor_{}.pth'.format(i))
                torch.save(agt.actor_target.state_dict(), './checkpoints/checkpoint_actor_target_{}.pth'.format(i))
                torch.save(agt.critic_local.state_dict(), './checkpoints/checkpoint_critic_{}.pth'.format(i))
                torch.save(agt.critic_target.state_dict(), './checkpoints/checkpoint_critic_target_{}.pth'.format(i))

                i += 1

            t_start = time.time()
            # Keep track of final episode reward
            final_ep_rewards.append(np.mean(episode_rewards[-1000:]))
            for rew in agent_rewards:
                final_ep_ag_rewards.append(np.mean(rew[-1000:]))

        # if len(episode_rewards) > 60000:
        #     break

        # saves final episode reward for plotting training curve later
        if len(episode_rewards) > 20000:
            rew_file_name = plots_dir + exp_name + '_rewards.pkl'
            os.makedirs(os.path.dirname(rew_file_name), exist_ok=True)
            with open(rew_file_name, 'wb') as fp:
                pickle.dump(final_ep_rewards, fp)
            agrew_file_name = plots_dir + exp_name + '_agrewards.pkl'
            os.makedirs(os.path.dirname(agrew_file_name), exist_ok=True)
            with open(agrew_file_name, 'wb') as fp:
                pickle.dump(final_ep_ag_rewards, fp)
            print('...Finished total of {} episodes.'.format(len(episode_rewards)))
            break

if __name__ == '__main__':
    train(display=False)