# uncompyle6 version 3.3.1
# Python bytecode 2.7 (62211)
# Decompiled from: Python 3.7.1 (default, Dec 14 2018, 19:28:38)
# [GCC 7.3.0]
# Embedded file name: /inna/quantize/quantize.py
# Compiled at: 2019-07-20 08:25:54
import numpy as np, os, sys, argparse, math, collections,  google.protobuf.text_format as tfmt, json #,scipy
from quantize_base import *
#from quantize_caffe import *
from quantize_tf import *
import cv2

class Frontend:

    def __init__(self, framework='tensorflow', deploy_model=None, output_json=None, weights=None, calibration_directory=None, calibration_size=None, calibration_seed=None, calibration_indices=None, bitwidths=None, dims=None, transpose=None, channel_swap=None, raw_scale=None, mean_value=None, input_scale=None):
        self.framework = framework
        self.deploy_model = deploy_model
        self.output_json = output_json
        self.weights = weights
        self.calibration_directory = calibration_directory
        self.calibration_size = calibration_size
        self.calibration_seed = calibration_seed
        self.calibration_indices = calibration_indices
        if bitwidths is None:
            self.bitwidths = [
             16, 16, 16]
        else:
            if not isinstance(bitwidths, list):
                self.bitwidths = eval(bitwidths)
            else:
                self.bitwidths = bitwidths
        if dims is not None and not isinstance(dims, list):
            self.dims = eval(dims)
        else:
            self.dims = dims
        if not isinstance(transpose, list):
            self.transpose = eval(transpose)
        else:
            self.transpose = transpose
        if not isinstance(channel_swap, list):
            self.channel_swap = eval(channel_swap)
        else:
            self.channel_swap = channel_swap
        self.raw_scale = raw_scale
        if not isinstance(mean_value, list):
            self.mean_value = eval(mean_value)
        else:
            self.mean_value = mean_value
        self.input_scale = input_scale
        return

    def quantize(self):
        if self.deploy_model is None:
            print 'Error: You must provide deploy model...'
            return
        if self.calibration_directory is None:
            print 'Error: You must provide calibration directory...'
            return
        calibration_filenames = list((os.path.isfile(os.path.join(self.calibration_directory, calibration_filename)) and calibration_filename for calibration_filename in os.listdir(self.calibration_directory)))
        for i in range(0, len(calibration_filenames)):
            calibration_filenames[i] = os.path.join(self.calibration_directory, calibration_filenames[i])

        calibration_size = None
        if self.calibration_size is not None:
            calibration_size = int(self.calibration_size)
            if calibration_size > len(calibration_filenames):
                print 'Error: Requested calibration size is greater than the number of available files in the specified calibration directory\n'
                return
        else:
            calibration_size = len(calibration_filenames)
        calibration_indices = None
        extra_calibration_indices = None
        np.random.seed(self.calibration_seed)
        if self.calibration_indices is not None:
            calibration_indices = np.array([ int(s) for s in eval(self.calibration_indices) ])
            remaining_array = np.setdiff1d(np.arange(0, len(calibration_filenames)), calibration_indices)
            extra_calibration_indices = np.sort(np.random.choice(remaining_array, calibration_size - len(calibration_indices), replace=False))
            calibration_indices = np.sort(np.append(calibration_indices, extra_calibration_indices))
        else:
            calibration_indices = np.sort(np.random.choice(list(range(len(calibration_filenames))), calibration_size, replace=False))
        bitwidths = self.bitwidths
        dims = self.dims
        transpose = self.transpose
        channel_swap = self.channel_swap
        raw_scale = float(self.raw_scale)
        mean_value = None
        if self.mean_value is not None:
            mean_value_list = [ float(s) for s in self.mean_value ]
            mean_value = np.array(mean_value_list)
        input_scale = float(self.input_scale) if self.input_scale else None
        print ('Mean :', mean_value)
        quant_params = QuantParam()
        #framework = 'caffe'
        if self.framework == 'caffe':
            caffeObj = Caffe(self.deploy_model, self.weights)
            caffeObj.declareTransformer(transpose, channel_swap, raw_scale, mean_value, input_scale, dims=None)
            caffeObj.initializeBatch(calibration_size, calibration_filenames, calibration_indices)
            quant_layers = caffeObj.executeCalibration(bitwidths, self.deploy_model, quant_params)
        else:
            if self.framework == 'tensorflow':
                tfobj = tf_Quantizer(self.deploy_model, bitwidths, quant_params, calibration_filenames, calibration_size, mean_value, self.calibration_seed, calibration_indices)
                quant_layers = tfobj.quantize('data', 'prob')
                #pass
            else:
                print 'Unknown framework %s.  Exiting.' % self.framework
                return
        ofile = self.deploy_model.replace('prototxt', 'json')
        if self.output_json is not None:
            ofile = self.output_json
        pathname = os.path.dirname(ofile)
        if pathname is '':
            pass
        else:
            if not os.path.exists(pathname):
                try:
                    os.makedirs(pathname)
                except:
                    print 'Error: The string given for the output_json contains a path that does not exist, and cannot be created.'
                    sys.exit(1)

        quant_params.saveToJson_per_layer(quant_layers, ofile) #lly
        #quant_params.saveToJson(quant_layers, ofile)
        print 'Arguments:'
        print [ (key, val) for key, val in list(self.__dict__.items()) if not key.startswith('__') ]
        return quant_params,quant_layers#True

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parameters = [(None, '--framework', str,  'tensorflow', 'store'),
     (
      None, '--deploy_model',
      str,
      '../models/tensorflow/resnet/resnet50_without_bn_test.pb',
      'store'),
     (
      None,
      '--output_json',
      str, './out/filter.bin', 'store'),
    (None, '--weights', str, '../models/resnet50_without_bn_test.pb', 'store'), (None, '--calibration_directory', str, '../runtime/images/', 'store'), (None, '--calibration_size', int, 8, 'store'), (None, '--calibration_seed', int, None, 'store'), (None, '--calibration_indices', str, None, 'store'), (None, '--bitwidths', str, '[8,8,8]', 'store'), (None, '--dims', str, '[3,224,224]', 'store'), (None, '--transpose', str, '[2,0,1]', 'store'), (None, '--channel_swap', str, '[2,1,0]', 'store'), (None, '--raw_scale', float, 255.0, 'store'), (None, '--mean_value', str, '[104.0,117.0,123.0]', 'store'), (None, '--input_scale', float, 1.0, 'store')]
    for x in parameters:
        if x[0] is not None:
            parser.add_argument(x[0], x[1], type=x[2], default=x[3], action=x[4])
        else:
            parser.add_argument(x[1], type=x[2], default=x[3], action=x[4])

    args = parser.parse_args()
    if not (os.path.isfile(args.deploy_model) and os.access(args.deploy_model, os.R_OK)):
        sys.exit('ERROR: Specified prototxt file does not exist or is not readable.')
    if not (os.path.isfile(args.weights) and os.access(args.weights, os.R_OK)):
        sys.exit('ERROR: Specified caffemodel file does not exist or is not readable.')
    if not (os.path.isdir(args.calibration_directory) and os.access(args.calibration_directory, os.R_OK)):
        sys.exit('ERROR: Specified calibration directory does not exist or is not readable.')
    quantizer = Frontend(args.framework, args.deploy_model, args.output_json, args.weights, args.calibration_directory, args.calibration_size, args.calibration_seed, args.calibration_indices, args.bitwidths, args.dims, args.transpose, args.channel_swap, args.raw_scale, args.mean_value, args.input_scale)
    params, layers = quantizer.quantize()
    params_json = params.jsontype(layers)

    weights_int8data, biases_int32data = quantization_tensorflow(args.deploy_model, params_json)

    filters = generate_filter(params_json, weights_int8data, biases_int32data)

    with open(os.path.dirname(os.path.abspath(__file__)) + '/filters.bin', 'wb') as fp:
        for filter in filters:
            fp.write(filter.tobytes())
# okay decompiling quantize.pyc


import os
import argparse
import cv2
import sys
import numpy as np
import insightface
import fuzzyextractor
import time
import math
from insightface.app import FaceAnalysis
from multiprocessing import Pool

if __name__=='__main__':

    # This is a super simple (but slow) example of running face recognition on live video from your webcam.
    # There's a second example that's a little more complicated but runs faster.

    # PLEASE NOTE: This example requires OpenCV (the `cv2` library) to be installed only to read from your webcam.
    # OpenCV is *not* required to use the face_recognition library. It's only required if you want to run this
    # specific demo. If you have trouble installing it, try any of the other demos that don't require it instead.

    # Get a reference to webcam #0 (the default one)
    video_capture = cv2.VideoCapture(0)

    n=512
    FE=fuzzyextractor.FuzzyExtractor(n,0.075)

    keys=[]
    helpers=[]
    names=[]
    pool=Pool(1)

    parser = argparse.ArgumentParser(description='insightface app test')
    parser.add_argument('--ctx', default=0, type=int, help='ctx id, <0 means using cpu')
    parser.add_argument('--det-size', default=640, type=int, help='detection size')
    args = parser.parse_args()
    app = FaceAnalysis(name='antelopev2')
    app.prepare(ctx_id=args.ctx, det_size=(args.det_size,args.det_size))

    for f in os.listdir('./RPs/'):
        if f[-3:]=='key':
            name=f[0:-4]

            with open('./RPs/'+name+'.key','rb') as fi:
                R=fi.read(32)
            with open('./RPs/'+name+'.helper','r') as fi:
                P=np.array([int(x) for x in fi.readline().split()])

            keys.append(R)
            helpers.append(P)
            names.append(name)

    proc=pool.apply_async(FE.Rep,(P,np.zeros(n)))
    name='Unknown'
    idx=0
    cnt=0

    while True:
        # Grab a single frame of video
        ret, img = video_capture.read()
        try:
            faces = app.get(img)
        except:
            continue

        rimg=img.copy()

        for face in faces[0:1]:
            img_embedding = face.normed_embedding.tolist()
            box = face.bbox.astype(np.int)
            color = (0, 0, 255)
            cv2.rectangle(rimg, (box[0], box[1]), (box[2], box[3]), color, 2)

            if proc.ready():
                R_=proc.get()
                if R_==keys[idx]:
                    print('Yes!')
                    name=names[idx]
                else:
                    cnt+=1
                    if cnt>len(names)+1:
                        cnt=0
                        name = "Unknown"
                idx=(idx+1)%len(names)
                print(idx)
                proc=pool.apply_async(FE.Rep,(helpers[idx],img_embedding))

            cv2.rectangle(rimg, (box[0], box[3] + 35), (box[2], box[3]), (0, 0, 255), cv2.FILLED)
            font = cv2.FONT_HERSHEY_DUPLEX
            cv2.putText(rimg, name, (box[0] , box[3] +30), font, 1.0, (255, 255, 255), 1)

        # Display the resulting image
        cv2.imshow('Video, Press Q to exit', rimg)

        # Hit 'q' on the keyboard to quit!
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # Release handle to the webcam
    video_capture.release()
    cv2.destroyAllWindows()


import functools
import glob
import itertools
import random

import cv2
import kornia
import numpy as np
import pytorch_ssim
import torch
import os

import torchvision
from torch.utils.data import DataLoader
from torchvision.transforms import Normalize, CenterCrop
from tqdm import tqdm

from data import util
# Builds a dataset created from a simple folder containing a list of training/test/validation images.
from data.image_corruptor import ImageCorruptor, kornia_color_jitter_numpy
from data.image_label_parser import VsNetImageLabeler
from utils.util import opt_get

def ndarray_center_crop(crop, img):
    y, x, c = img.shape
    startx = x // 2 - crop // 2
    starty = y // 2 - crop // 2
    return img[starty:starty + crop, startx:startx + crop, :]

class ImageFolderDataset:
    def __init__(self, opt):
        self.opt = opt
        self.corruptor = ImageCorruptor(opt)
        if 'center_crop_hq_sz' in opt.keys():
            self.center_crop = functools.partial(ndarray_center_crop, opt['center_crop_hq_sz'])
        self.target_hq_size = opt['target_size'] if 'target_size' in opt.keys() else None
        self.multiple = opt['force_multiple'] if 'force_multiple' in opt.keys() else 1
        self.scale = opt['scale']
        self.paths = opt['paths']
        self.corrupt_before_downsize = opt['corrupt_before_downsize'] if 'corrupt_before_downsize' in opt.keys() else False
        self.fetch_alt_image = opt['fetch_alt_image']  # If specified, this dataset will attempt to find a second image
                                                       # from the same video source. Search for 'fetch_alt_image' for more info.
        self.fetch_alt_tiled_image = opt['fetch_alt_tiled_image']  # If specified, this dataset will attempt to find anoter tile from the same source image
                                                                   #  Search for 'fetch_alt_tiled_image' for more info.
        assert not (self.fetch_alt_image and self.fetch_alt_tiled_image)  # These are mutually exclusive.
        self.skip_lq = opt_get(opt, ['skip_lq'], False)
        self.disable_flip = opt_get(opt, ['disable_flip'], False)
        self.rgb_n1_to_1 = opt_get(opt, ['rgb_n1_to_1'], False)
        self.force_square = opt_get(opt, ['force_square'], True)
        self.fixed_parameters = {k: torch.tensor(v) for k, v in opt_get(opt, ['fixed_parameters'], {}).items()}
        self.all_image_color_jitter = opt_get(opt, ['all_image_color_jitter'], 0)
        if 'normalize' in opt.keys():
            if opt['normalize'] == 'stylegan2_norm':
                self.normalize = Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
            elif opt['normalize'] == 'imagenet':
                self.normalize = Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225), inplace=True)
            else:
                raise Exception('Unsupported normalize')
        else:
            self.normalize = None
        if self.target_hq_size is not None:
            assert (self.target_hq_size // self.scale) % self.multiple == 0  # If we dont throw here, we get some really obscure errors.
        if not isinstance(self.paths, list):
            self.paths = [self.paths]
            self.weights = [1]
        else:
            self.weights = opt['weights']

        if 'labeler' in opt.keys():
            if opt['labeler']['type'] == 'patch_labels':
                self.labeler = VsNetImageLabeler(opt['labeler']['label_file'])
            assert len(self.paths) == 1   # Only a single base-path is supported for labeled images.
            self.image_paths = self.labeler.get_labeled_paths(self.paths[0])
        else:
            self.labeler = None

            # Just scan the given directory for images of standard types.
            supported_types = ['jpg', 'jpeg', 'png', 'gif']
            self.image_paths = []
            for path, weight in zip(self.paths, self.weights):
                cache_path = os.path.join(path, 'cache.pth')
                if os.path.exists(cache_path):
                    imgs = torch.load(cache_path)
                else:
                    print("Building image folder cache, this can take some time for large datasets..")
                    imgs = util.find_files_of_type('img', path)[0]
                    torch.save(imgs, cache_path)
                for w in range(weight):
                    self.image_paths.extend(imgs)
        self.len = len(self.image_paths)

    def get_paths(self):
        return self.image_paths

    # Given an HQ square of arbitrary size, resizes it to specifications from opt.
    def resize_hq(self, imgs_hq):
        # Enforce size constraints
        h, w, _ = imgs_hq[0].shape

        if self.target_hq_size is not None and self.target_hq_size != h:
            hqs_adjusted = []
            for hq in imgs_hq:
                # It is assumed that the target size is a square.
                target_size = (self.target_hq_size, self.target_hq_size)
                hqs_adjusted.append(cv2.resize(hq, target_size, interpolation=cv2.INTER_AREA))
            h, w = self.target_hq_size, self.target_hq_size
        else:
            hqs_adjusted = imgs_hq
        hq_multiple = self.multiple * self.scale  # Multiple must apply to LQ image.
        if h % hq_multiple != 0 or w % hq_multiple != 0:
            hqs_conformed = []
            for hq in hqs_adjusted:
                h, w = (h - h % hq_multiple), (w - w % hq_multiple)
                hqs_conformed.append(hq[:h, :w, :])
            return hqs_conformed
        return hqs_adjusted

    def synthesize_lq(self, hs):
        h, w, _ = hs[0].shape
        ls = []
        local_scale = self.scale
        if self.corrupt_before_downsize:
            # You can downsize to a specified scale, then corrupt, then continue the downsize further using this option.
            if 'corrupt_before_downsize_factor' in self.opt.keys():
                special_factor = self.opt['corrupt_before_downsize_factor']
                hs = [cv2.resize(h_, (h // special_factor, w // special_factor), interpolation=cv2.INTER_AREA) for h_ in hs]
                local_scale = local_scale // special_factor
            else:
                hs = [h.copy() for h in hs]
            hs, ent = self.corruptor.corrupt_images(hs, return_entropy=True)
        for hq in hs:
            h, w, _ = hq.shape
            ls.append(cv2.resize(hq, (h // local_scale, w // local_scale), interpolation=cv2.INTER_AREA))
        # Corrupt the LQ image (only in eval mode)
        if not self.corrupt_before_downsize:
            ls, ent = self.corruptor.corrupt_images(ls, return_entropy=True)
        return ls, ent

    def reset_random(self):
        self.corruptor.reset_random()

    def __len__(self):
        return self.len

    def __getitem__(self, item):
        hq = util.read_img(None, self.image_paths[item], rgb=True)
        if hasattr(self, 'center_crop'):
            hq = self.center_crop(hq)
        if not self.disable_flip and random.random() < .5:
            hq = hq[:, ::-1, :]

        if self.force_square:
            h, w, _ = hq.shape
            dim = min(h, w)
            hq = hq[(h - dim) // 2:dim + (h - dim) // 2, (w - dim) // 2:dim + (w - dim) // 2, :]

        # Perform color jittering on the HQ image if specified. The given value should be between [0,1].
        if self.all_image_color_jitter > 0:
            hq = kornia_color_jitter_numpy(hq, self.all_image_color_jitter)

        if self.labeler:
            assert hq.shape[0] == hq.shape[1]  # This just has not been accomodated yet.
            dim = hq.shape[0]

        hs = self.resize_hq([hq])
        if not self.skip_lq:
            for_lq = [hs[0]]

        # Convert to torch tensor
        hq = torch.from_numpy(np.ascontiguousarray(np.transpose(hs[0], (2, 0, 1)))).float()

        out_dict = {'hq': hq, 'LQ_path': self.image_paths[item], 'HQ_path': self.image_paths[item], 'has_alt': False}

        if self.fetch_alt_image:
            # This works by assuming a specific filename structure as would produced by ffmpeg. ex:
            # 'Candied Walnutsxjktqhr_SYc.webm_00000478.jpg` and
            # 'Candied Walnutsxjktqhr_SYc.webm_00000479.jpg` and
            # 'Candied Walnutsxjktqhr_SYc.webm_00000480.jpg`
            # The basic format is `<anything>%08d.<extension>`. This logic parses off that 8 digit number. If it is
            # not found, the 'alt_image' returned is just the current image. If it is found, the algorithm searches for
            # an image one number higher. If it is found - it is returned in the 'alt_hq' and 'alt_lq' keys, else the
            # current image is put in those keys.

            imname_parts = self.image_paths[item]
            while '.jpg.jpg' in imname_parts:
                imname_parts = imname_parts.replace(".jpg.jpg", ".jpg")  # Hack workaround to my own bug.
            imname_parts = imname_parts.split('.')
            if len(imname_parts) >= 2 and len(imname_parts[-2]) > 8:
                try:
                    imnumber = int(imname_parts[-2][-8:])
                    # When we're dealing with images in the 1M range, it's straight up faster to attempt to just open
                    # the file rather than searching the path list. Let the exception handler below do its work.
                    next_img = self.image_paths[item].replace(str(imnumber), str(imnumber+1))
                    alt_hq = util.read_img(None, next_img, rgb=True)
                    alt_hs = self.resize_hq([alt_hq])
                    alt_hq = torch.from_numpy(np.ascontiguousarray(np.transpose(alt_hs[0], (2, 0, 1)))).float()
                    out_dict['has_alt'] = True
                    if not self.skip_lq:
                        for_lq.append(alt_hs[0])
                except:
                    alt_hq = hq
                    if not self.skip_lq:
                        for_lq.append(hs[0])
            else:
                alt_hq = hq
                if not self.skip_lq:
                    for_lq.append(hs[0])
            out_dict['alt_hq'] = alt_hq

        if self.fetch_alt_tiled_image:
            # This assumes the output format generated by the tiled image generation scripts included with DLAS. Specifically,
            # all image read by this dataset are assumed to be in subfolders with other tiles from the same source image. When
            # this option is set, another random image from the same folder is selected and returned as the alt image.
            sel_path = os.path.dirname(self.image_paths[item])
            other_images = os.listdir(sel_path)
            # Assume that the directory contains at least <image>, <ref.jpg>, <centers.pt>
            try:
                if len(other_images) <= 3:
                    alt_hq = hq  # This is a fallback in case an alt image can't be found.
                else:
                    random.shuffle(other_images)
                    for oi in other_images:
                        if oi == os.path.basename(self.image_paths[item]) or 'ref.' in oi or 'centers.pt' in oi:
                            continue
                        alt_hq = util.read_img(None, os.path.join(sel_path, oi), rgb=True)
                        alt_hs = self.resize_hq([alt_hq])
                        alt_hq = torch.from_numpy(np.ascontiguousarray(np.transpose(alt_hs[0], (2, 0, 1)))).float()
            except:
                alt_hq = hq
                print(f"Error with {self.image_paths[item]}")
            out_dict['has_alt'] = True
            out_dict['alt_hq'] = alt_hq

        if not self.skip_lq:
            lqs, ent = self.synthesize_lq(for_lq)
            ls = lqs[0]
            out_dict['lq'] = torch.from_numpy(np.ascontiguousarray(np.transpose(ls, (2, 0, 1)))).float()
            out_dict['corruption_entropy'] = torch.tensor(ent)
            if len(lqs) > 1:
                alt_lq = lqs[1]
                out_dict['alt_lq'] = torch.from_numpy(np.ascontiguousarray(np.transpose(alt_lq, (2, 0, 1)))).float()

        if self.labeler:
            base_file = self.image_paths[item].replace(self.paths[0], "")
            while base_file.startswith("\\"):
                base_file = base_file[1:]
            assert dim % hq.shape[1] == 0
            lbls, lbl_masks, lblstrings = self.labeler.get_labels_as_tensor(hq, base_file, dim // hq.shape[1])
            out_dict['labels'] = lbls
            out_dict['labels_mask'] = lbl_masks
            out_dict['label_strings'] = lblstrings

        for k, v in out_dict.items():
            if isinstance(v, torch.Tensor) and len(v.shape) == 3:
                if self.normalize:
                    v = self.normalize(v)
                if self.rgb_n1_to_1:
                    v = v * 2 - 1
                out_dict[k] = v

        out_dict.update(self.fixed_parameters)
        return out_dict

if __name__ == '__main__':
    opt = {
        'name': 'amalgam',
        'paths': ['E:\\4k6k\\datasets\\ns_images\\imagesets\\256_only_humans_masked'],
        'weights': [1],
        'target_size': 256,
        'scale': 1,
        'corrupt_before_downsize': True,
        'fetch_alt_image': False,
        'fetch_alt_tiled_image': True,
        'disable_flip': True,
        'fixed_corruptions': ['lq_resampling', 'jpeg-medium', 'gaussian_blur', 'noise', 'color_jitter'],
        'num_corrupts_per_image': 0,
        'corruption_blur_scale': 1,
        'all_image_color_jitter': .1,
    }

    ds = DataLoader(ImageFolderDataset(opt), shuffle=True, num_workers=0, batch_size=64)
    import os
    output_path = 'F:\\tmp'
    os.makedirs(output_path, exist_ok=True)
    res = []
    for i, d in tqdm(enumerate(ds)):
        '''
        x = d['hq']
        b,c,h,w = x.shape
        x_c = x.view(c*b, h, w)
        x_c = torch.view_as_real(torch.fft.rfft(x_c))
        # Log-normalize spectrogram
        x_c = (x_c.abs() ** 2).clip(min=1e-8, max=1e16)
        x_c = torch.log(x_c)
        res.append(x_c)
        if i % 100 == 99:
            stacked = torch.cat(res, dim=0)
            print(stacked.mean(dim=[0,1,2]), stacked.std(dim=[0,1,2]))
        '''

        for k, v in d.items():
            if isinstance(v, torch.Tensor) and len(v.shape) >= 3:
                os.makedirs(f'{output_path}\\{k}', exist_ok=True)
                torchvision.utils.save_image(v, f'{output_path}\\{k}\\{i}.png')
        if i >= 200000:
            break


import joblib
import numpy as np
import os

from matplotlib import pyplot as plt
from matplotlib.colors import LogNorm

from arctor import Arctor, create_raw_lc_stddev
from arctor import debug_message, warning_message, info_message

from arctor.plotting import (plot_aperture_edges_with_angle,
                             plot_lightcurve,
                             plot_lightcurve_orbit,
                             plot_xcenter_vs_time,
                             plot_ycenter_vs_time,
                             plot_xcenter_vs_flux,
                             plot_ycenter_vs_flux,
                             plot_center_position_vs_scan_and_orbit,
                             plot_2D_stddev)

if __name__ == '__main__':
    from argparse import ArgumentParser

    parser = ArgumentParser()

    parser.add_argument('--planet_name', type=str, default=None)
    parser.add_argument('--file_type', type=str, default='flt.fits')
    parser.add_argument('--save_now', action='store_true')
    parser.add_argument('--fit_model_flag', action='store_true')
    parser.add_argument('--plot_verbose', action='store_true')

    clargs = parser.parse_args()

    planet_name = clargs.planet_name
    if clargs.planet_name is None:
        planet_name = 'HD80606'  # because I am unimaginative

    assert(planet_name is not None), \
        "Must provide planet name as --planet_name or -pn"

    file_type = clargs.file_type

    # HOME = os.environ['HOME']
    base_dir = os.path.join('/', 'path', 'to', 'base')
    data_dir = os.path.join(base_dir, 'path', 'to', 'data')
    working_dir = os.path.join(base_dir, 'analysis')
    save_dir = os.path.join(base_dir, 'savefiles')

    assert(os.path.exists(base_dir)), (
        '[WARNING] Please either edit the `base_dir` '
        'or configure the directory structure as above')
    assert(os.path.exists(data_dir)), (
        '[WARNING] Please either edit the `data_dir` '
        'or configure the directory structure as above')

    if not os.path.exists(working_dir):
        info_message(f'Creating {working_dir}')
        os.mkdir(working_dir)
    if not os.path.exists(save_dir):
        info_message(f'Creating {save_dir}')
        os.mkdir(save_dir)

    joblib_filename = (f'{planet_name}'
                       '_savedict_1683_1-200x1-200_x5.joblib.save')

    # csv_filename = f'{working_dir}/{csv_filename}'
    joblib_filename = f'{save_dir}/{joblib_filename}'

    assert(os.path.exists(joblib_filename)),\
        f'File not found:{joblib_filename}'

    planet = Arctor(
        planet_name=planet_name,
        data_dir=data_dir,
        working_dir=working_dir,
        file_type=file_type)

    # This forces the fits files to be opened in time order
    #   Very useful the *first* time we run the example
    # planet.rename_fits_files_by_time()

    if os.path.exists(joblib_filename):
        info_message('Loading Data from Save File')
        planet.load_dict(joblib_filename)
    else:
        info_message('Loading New Data Object')
        planet.load_data()  # sort_by_time=sort_by_time)

    if not hasattr(planet, 'gaussian_centers'):
        planet.clean_cosmic_rays()
        planet.calibration_trace_location(oversample=100)
        planet.identify_trace_direction()
        planet.simple_phots()
        planet.center_all_traces(
            stddev=2, notit_verbose=False, idx_buffer=10)
        # planet.fit_trace_slopes(stddev=2, notit_verbose=False)
        planet.compute_trace_slopes(
            stddev=2, notit_verbose=False, x_offset=100)
        planet.compute_sky_background(
            subpixels=32, positions=None, inner_width=75, outer_width=150,
            inner_height=225, outer_height=350, thetas=None,
            notit_verbose=False, done_it=False)
        planet.compute_columnwise_sky_background(inner_height=150, edge=10)

    # Set up the list of aperture widths and heights to search
    min_aper_width = 1
    max_aper_width = 200
    min_aper_height = 1
    max_aper_height = 200

    aper_widths = np.arange(min_aper_width, max_aper_width + 2, 5)
    aper_heights = np.arange(min_aper_height, max_aper_height + 2, 5)

    if 'yes' in input('Are you sure? (yes/no) ').lower()[:3]:
        planet.do_multi_phot(aper_widths, aper_heights)
    else:
        print('[INFO] Skipping `do_multi_phot` until you are ready')

    aper_width = 176  # lowest SDNR
    aper_height = 116  # lowest SDNR
    ax = None
    ax = plot_lightcurve(planet, aper_width, aper_height,
                         t0_base=planet.times.min(),
                         ax=ax,
                         include_orbits=False,
                         highlight_outliers=True,
                         include_uncertainties=True,
                         use_stat_ylim=True,
                         size=100,
                         n_sig=2)

    axs = None
    axs = plot_2D_stddev(planet, signal_max=350, axs=axs)

    ax = plot_center_position_vs_scan_and_orbit(planet, ax=ax)

    snr_lightcurves = create_raw_lc_stddev(planet)
    min_snr_colname = planet.photometry_df.columns[snr_lightcurves.argmin()]
    min_snr_col = planet.normed_photometry_df[min_snr_colname]
    temp = min_snr_colname.split('_')[-1].split('x')
    min_snr_aper_width, min_snr_aper_height = np.int32(temp)

    fine_buffer = 10
    fine_aper_widths = np.arange(min_snr_aper_width - fine_buffer,
                                 min_snr_aper_width + fine_buffer)

    fine_aper_heights = np.arange(min_snr_aper_height - fine_buffer,
                                  min_snr_aper_height + fine_buffer)

    if 'y' in input('Are you sure? ').lower()[:1]:
        planet.do_multi_phot(fine_aper_widths, fine_aper_heights)
    else:
        print('[INFO] Skipping `do_multi_phot` until you are ready')

    if clargs.plot_verbose:
        planet.plot_errorbars()

    if clargs.fit_model_flag:
        res = planet.do_fit(times, fluxes, errors,
                            init_params=[], static_params={})

    # csv_filename = f'{clargs.planet_name}_photometry.csv'
    # joblib_filename =
    # f'{planet_name}_savedict_not_finished_with_images.joblib.save'
    # f'{planet_name}_savedict_no_phots_with_images.joblib.save'

    if clargs.save_now:
        # joblib_filename = f'{planet_name}_savedict_NNN_NNNxNNN.joblib.save'
        joblib_filename = f'{planet_name}_savedict_1683_1-200x1-200_x5.joblib.save'

        # csv_filename = f'{working_dir}/{csv_filename}'
        joblib_filename = f'{save_dir}/{joblib_filename}'

        # planet.save_text_file(csv_filename)
        planet.save_dict(joblib_filename)

    aper_width = 176  # FINDME: Check these
    aper_height = 116  # FINDME: Check these
    ax = plot_lightcurve(planet, aper_width, aper_height,
                         t0_base=planet.times.min(), ax=ax,
                         include_orbits=True,
                         highlight_outliers=True)
    plot_ycenter_vs_flux(planet, aper_width, aper_height)

    plot_xcenter_vs_flux(planet, aper_width, aper_height, ax=ax,
                         include_orbits=True,
                         highlight_outliers=True)

    plot_ycenter_vs_flux(planet, aper_width, aper_height, ax=ax,
                         include_orbits=True,
                         highlight_outliers=True)
    plot_ycenter_vs_time(planet, aper_width, aper_height, ax=ax,
                         include_orbits=True,
                         highlight_outliers=True)
    plot_xcenter_vs_time(planet, aper_width, aper_height, ax=ax,
                         include_orbits=True,
                         highlight_outliers=True)

    plot_aperture_edges_with_angle(planet, img_id=42)
    plot_center_position_vs_scan_and_orbit(planet)
    plot_2D_stddev(planet.photometry_df,
                   aper_widths,
                   aper_heights,
                   signal_max=230)
    ax1, ax2 = None, None
    ax1 = plot_trace_over_time(
        planet, ax=ax1, adjust_xloc=False, metric=np.std, delta_y=10)
    ax2 = plot_trace_over_time(
        planet, ax=ax2, adjust_xloc=True, metric=np.std, delta_y=10)


#!/usr/bin/python
# -*- coding: utf-8 -*-

import Reader, pickle
import networkx as nx
import Participation_Coefficient_origal as MPC
# import ForestFireCross as ForestFire
import random_walk_sampler as MHRW
import numpy as np
from functools import partial
from multiprocessing import Pool as ThreadPool

import warnings
warnings.filterwarnings(action='ignore', category=UserWarning, module='gensim')
import gensim
from gensim.models import Word2Vec
from sklearn import metrics, model_selection, pipeline
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
import os
import copy
import warnings
assert gensim.models.doc2vec.FAST_VERSION > -1

edge_functions = {
    "hadamard": lambda a, b: a * b,
    "average": lambda a, b: 0.5 * (a + b),
    "l1": lambda a, b: np.abs(a - b),
    "l2": lambda a, b: np.abs(a - b) ** 2,
}

default_params = { # 创建一个算法需要用到的参数字典
    'num_walks': 10,                # Number of walks from each node
    'walk_length': 30,              # Walk length， karate datasets is 30
    'window_size': 10,              # Context size for word2vec
    'edge_function': "hadamard",    # Default edge function to use
    "prop_pos": 0.5,                # Proportion of edges to remove nad use as positive samples
    "prop_neg": 0.5,                # Number of non-edges to use as negative samples
    "Reflash_test_data": False        #  (as a proportion of existing edges, same as prop_pos)
}

class Mergeing_vec_N2V:
    def __init__(self, path, ps, j, sampling_p, p, q, num_walks, walk_length, r, num_partitions, dimensions, workers, ff):
        self.path = path
        self.ps = ps
        self.j = j
        self.s_p = sampling_p # 采样比例
        self.p = p
        self.q = q
        self.num_walks = num_walks
        self.walk_length = walk_length
        self.r = r
        self.exnum = num_partitions
        self.dimensions = dimensions
        self.workers = workers
        self.LG = None
        self.MG = None
        self._pos_edge_list = None
        self._neg_edge_list = None
        self.wvecs = None
        self._rnd = np.random.RandomState(seed=None)  # 随机一个种子
        self.ff = ff

    def generate_pos_neg_links(self):
        # Select n edges at random (positive samples)
        test_g = self.LG[-1] # 选择最后一层作为测试层
        share_L_nodes = set([node for g in self.LG[:-1] for node in g.nodes()]) # 选择训练层所存在的节点
        n_edges = test_g.number_of_edges()
        test_edges = test_g.edges()

        # 如果测试的边都在训练层中，作为最终要测试的边
        share_test_edges = [edge for edge in test_edges if (edge[0] in share_L_nodes) and (edge[1] in share_L_nodes)]
        share_edge = len(share_test_edges)
        npos = int(self.s_p * share_edge) #正采样的边数，默认为1
        nneg = npos                       #负采样的边数，默认等于正采样的边数

        # 负样例的边集
        non_edges = [e for e in nx.non_edges(test_g) if (e[0] in share_L_nodes) and (e[1] in share_L_nodes)]
        print("Finding %d of %d non-edges" % (nneg, len(non_edges)))
        # 随机采样负样例的边表
        rnd_inx = self._rnd.choice(len(non_edges), nneg, replace=False) # 基于之前的随机种子在范围len(non_edges)中生成size为nneg的边数，不用替换；最后形成下标
        neg_edge_list = [non_edges[ii] for ii in rnd_inx] #构建负采样的边

        # 判断网络中负采样的边个数是否满足需要
        if len(neg_edge_list) < nneg:
            raise RuntimeWarning("Only %d negative edges found" % (len(neg_edge_list)))
        print("Finding %d positive edges of %d total edges, which share edges with each layer %d" % (npos, n_edges, share_edge))
        print("Finding %d negtive edges of %d total edges, which share edges with each layer %d" % (npos, n_edges, share_edge))

        # 正例边的构造
        edges = share_test_edges
        pos_edge_list = []
        n_count = 0
        # 随机性的对一个排序进行重排序
        rnd_inx = self._rnd.permutation(share_edge)
        for eii in rnd_inx[:npos]:
            edge = edges[eii]
            pos_edge_list.append(edge)
            n_count += 1

        # 判断正例边的个数是否满足需要
        if len(pos_edge_list) < npos:
            print("Only %d positive edges found." % (n_count))

        self._pos_edge_list = pos_edge_list
        self._neg_edge_list = neg_edge_list

    def get_selected_edges(self):
        edges = self._pos_edge_list + self._neg_edge_list
        labels = np.zeros(len(edges))
        labels[:len(self._pos_edge_list)] = 1
        return edges, labels

    def edges_to_features(self, edge_list, edge_function, dimensions):
        n_tot = len(edge_list)
        feature_vec = np.empty((n_tot, dimensions), dtype='f')

        # Iterate over edges
        for ii in range(n_tot):
            v1, v2 = edge_list[ii]

            # Edge-node features
            emb1 = np.asarray(self.wvecs[str(v1)])
            emb2 = np.asarray(self.wvecs[str(v2)])

            # Calculate edge feature
            feature_vec[ii] = edge_function(emb1, emb2)

        return feature_vec

    def get_alias_edges(self, g, src, dest, p=1, q=1):
        probs = []
        for nei in sorted(g.neighbors(dest)):
            if nei == src:
                probs.append(1 / p)
            elif g.has_edge(nei, src):
                probs.append(1)
            else:
                probs.append(1 / q)
        norm_probs = [float(prob) / sum(probs) for prob in probs]
        return self.get_alias_nodes(norm_probs)

    def get_alias_nodes(self, probs):
        l = len(probs)
        a, b = np.zeros(l), np.zeros(l, dtype=np.int)
        small, large = [], []

        for i, prob in enumerate(probs):
            a[i] = l * prob
            if a[i] < 1.0:
                small.append(i)
            else:
                large.append(i)

        while small and large:
            sma, lar = small.pop(), large.pop()
            b[sma] = lar
            a[lar] += a[sma] - 1.0
            if a[lar] < 1.0:
                small.append(lar)
            else:
                large.append(lar)
        return b, a

    def preprocess_transition_probs(self, g, directed=False, p=1, q=1):
        alias_nodes, alias_edges = {}, {};
        for node in g.nodes():
            probs = [g[node][nei]['weight'] for nei in sorted(g.neighbors(node))]
            norm_const = sum(probs)
            norm_probs = [float(prob) / norm_const for prob in probs]
            alias_nodes[node] = self.get_alias_nodes(norm_probs)

        if directed:
            for edge in g.edges():
                alias_edges[edge] = self.get_alias_edges(g, edge[0], edge[1], p, q)
                # print(alias_edges[edge])
        else:
            for edge in g.edges():
                alias_edges[edge] = self.get_alias_edges(g, edge[0], edge[1], p, q)
                alias_edges[(edge[1], edge[0])] = self.get_alias_edges(g, edge[1], edge[0], p, q)

        return alias_nodes, alias_edges

    def node2vec_walk(self, g, alias_nodes, alias_edges, walk_length = 30, start = 1):
        path = [start]
        walk_length = self.walk_length
        while len(path) < walk_length:
            node = path[-1]
            neis = sorted(g.neighbors(node))
            if len(neis) > 0:
                if len(path) == 1:
                    l = len(alias_nodes[node][0])
                    idx = int(np.floor(np.random.rand() * l))
                    if np.random.rand() < alias_nodes[node][1][idx]:
                        path.append(neis[idx])
                    else:
                        path.append(neis[alias_nodes[node][0][idx]])
                else:
                    prev = path[-2]
                    l = len(alias_edges[(prev, node)][0])
                    idx = int(np.floor(np.random.rand() * l))
                    if np.random.rand() < alias_edges[(prev, node)][1][idx]:
                        path.append(neis[idx])
                    else:
                        path.append(neis[alias_edges[(prev, node)][0][idx]])
            else:
                break
        return path

    def learn_embeddings(self, walks, dimensions, window_size=10, niter=5):
        '''
        Learn embeddings by optimizing the Skipgram objective using SGD.
        '''
        # TODO: Python27 only
        # walks = [map(str, walk) for walk in walks]
        model = Word2Vec(walks,
                         size=dimensions,
                         window=window_size,
                         min_count=0,
                         sg=1,
                         workers=self.workers,
                         iter=niter)
        self.wvecs = model.wv

    def run(self):
        cparams = default_params.copy() # 基本参数设置
        path = self.path
        no_reflesh_testdata = bool(1 - cparams["Reflash_test_data"]) #是否刷新：测试-训练数据集
        num_partitions = self.exnum #实验次数
        # Step 1: reading and sampling graphs
        cached_fn = "Sampling_graph/"+ os.path.basename(path) + "%s.graph"
        if os.path.exists(cached_fn) and bool(no_reflesh_testdata):
            # 加载已分好的缓存数据
            print("Loading link prediction graphs from %s" % cached_fn)
            with open(cached_fn, 'rb') as f:
                cache_data = pickle.load(f)
                self.LG = cache_data['g_train']
                self.MG = cache_data['g_merg'] # 合并的网络，可以用于合并层的表示学习
                self._pos_edge_list = cache_data['remove_list']
                self._neg_edge_list = cache_data['ne_list']
        else:
            # 重新构建缓存数据:多层图，合并图，正例边，负例边
            self.MG, self.LG, _ = Reader.multi_readG_with_Merg(path)
            self.generate_pos_neg_links()
            cache_data = {'g_train': self.LG, 'g_merg': self.MG, 'remove_list': self._pos_edge_list, 'ne_list': self._neg_edge_list}
            with open(cached_fn, 'wb') as f:
                pickle.dump(cache_data, f)

        # 构造训练数据集和测试数据集
        nx_graphs_sampled = self.LG[:-1]

        # 构建有向的测试集图例，因为在下面的很多对比算法中有向图作为主要图
        m_graph_sampled = nx.DiGraph()
        for g in nx_graphs_sampled:
            for e in g.edges.data():
                if m_graph_sampled.has_edge(e[0], e[1]):
                    m_graph_sampled[e[0]][e[1]]['weight'] = float(m_graph_sampled[e[0]][e[1]]['weight']) + float(e[2]['weight'])
                else:
                    m_graph_sampled.add_edge(e[0], e[1], weight=e[2]['weight'])

        # 实验：多层网络的表示学习方法(MHRW, FFS)
        MHWL = [] # 跳转MHRW采样-避免随机游走对度大的节点偏执
        # FFSN = [] # 林火采样

        # 将网络中的边权转化为1，然后进行网络的重构
        allnodes = m_graph_sampled.nodes()
        nodeinfluence, node_exit = MPC.Multiplex_PartC(m_graph_sampled, nx_graphs_sampled) #重要性的定义
        ExpansionSample_MHRW = partial(MHRW.random_walk_sampler, nx_graphs_sampled, node_exit, self.walk_length, nodeinfluence, True)
        # ExpansionSample_FFS = partial(ForestFire.forest_fire_sampling, nx_graphs_sampled, node_exit, self.walk_length, nodeinfluence, self.ff)
        with ThreadPool(processes=4) as pool:
            for walks in range(self.num_walks):
                MHWL.extend(pool.map(ExpansionSample_MHRW, allnodes))
                # FFSN.extend(pool.map(ExpansionSample_FFS, allnodes)) # VS2: Forest Fire 采样

        visited = []
        visited.append(MHWL)
        # visited.append(FFSN)
        # Algorithm = ["FFnsME"]
        Algorithm = ["MHME"]
        alg_num = len(visited) # 算法的个数，用于下面算法的评估

        # 对生成的序列进行词向量的学习
        partitioner = model_selection.StratifiedKFold(num_partitions, shuffle=True)

        # 提取要测试的所有边和对应的标签
        edges_all, edge_labels_all = self.get_selected_edges()

        # 基于节点向量计算边权的函数--可以作为一种实验函数
        edge_fn = edge_functions[cparams['edge_function']]

        # 实验性能测试部分
        for steps in range(alg_num):
            auc_train = []
            f1_train = []
            recall_train = []
            acc_train = []
            auc_test = []
            f1_test= []
            recall_test = []
            acc_test = []
            NMI = []
            for train_inx, test_inx in partitioner.split(edges_all, edge_labels_all):
                edges_train = [edges_all[jj] for jj in train_inx]
                labels_train = [edge_labels_all[jj] for jj in train_inx]
                edges_test = [edges_all[jj] for jj in test_inx]
                labels_test = [edge_labels_all[jj] for jj in test_inx]
                # Learn embeddings with current parameter values
                self.learn_embeddings(visited[steps], self.dimensions, window_size=10, niter=5)
                # Calculate edge embeddings using binary function
                edge_features_train = self.edges_to_features(edges_train, edge_fn, self.dimensions)
                edge_features_test = self.edges_to_features(edges_test, edge_fn, self.dimensions)

                # Linear classifier
                scaler = StandardScaler()
                lin_clf = LogisticRegression(C=1)
                clf = pipeline.make_pipeline(scaler, lin_clf)

                # Train & validate classifier
                clf.fit(edge_features_train, labels_train)
                metrics.scorer.mutual_info_scorer(clf, edge_features_train, labels_train)
                auc_train.append(metrics.scorer.roc_auc_scorer(clf, edge_features_train, labels_train))
                f1_train.append(metrics.scorer.f1_scorer(clf, edge_features_train, labels_train))
                recall_train.append(metrics.scorer.recall_scorer(clf, edge_features_train, labels_train))
                acc_train.append(metrics.scorer.precision_scorer(clf, edge_features_train, labels_train))

                # Test classifier
                NMI.append(metrics.scorer.normalized_mutual_info_scorer(clf, edge_features_train, labels_train))
                auc_test.append(metrics.scorer.roc_auc_scorer(clf, edge_features_test, labels_test))
                f1_test.append(metrics.scorer.f1_scorer(clf, edge_features_test, labels_test))
                recall_test.append(metrics.scorer.recall_scorer(clf, edge_features_test, labels_test))
                acc_test.append(metrics.scorer.accuracy_scorer(clf, edge_features_test, labels_test))
            print("Algorithm %s; AUC train: %.2f|AUC test: %.4f---F1 train: %.2f|F1 test: %.4f---Recall train:%.2f|Recall test:%.4f---precision train:%.2f|precision test:%.4f : NMI=%.4f"% (Algorithm[steps], sum(auc_train)/num_partitions, sum(auc_test)/num_partitions,sum(f1_train)/num_partitions,sum(f1_test)/num_partitions,sum(recall_train)/num_partitions,sum(recall_test)/num_partitions,sum(acc_train)/num_partitions,sum(acc_test)/num_partitions, sum(NMI)/num_partitions))
        print("---------------END Link Prediction----------------")


from pathlib import Path
import json

import numpy as np

from library import logger_generate
from config import base

reverse_config = base.reverse_config()
rc = reverse_config
logger = logger_generate.generate(base.logger_config())
# ===== 環境OK =====

logger.info("Loading analysis_from_video.json ...")
aims_folder_path = Path(rc['aims_folder_path'])
output_sheet_path = (aims_folder_path /
                     Path(rc['output_sheet_path'])).resolve()
_temp = output_sheet_path / './analysis_from_video.json'

with open(_temp, mode='r', encoding='utf-8') as f:
    _analysis_from_video = f.read()

analysis_from_video = json.loads(_analysis_from_video)
frame_keyboards = json.loads(analysis_from_video['notes'])
fps = analysis_from_video['fps']
cool_down_time = rc['cool_down_time']  # 冷卻時間(單位 ms)
cool_down_frame = round((cool_down_time / 1000) / (1 / fps))  # 冷卻時間(單位 偵數)
logger.info("Loaded analysis_from_video.json.")

logger.info('generating original sheet...')
# 先來定義一下 kb_list 格式
kb_list = []
for i in range(0, len(frame_keyboards[0])):
    # logger.debug(frame_keyboards[0][i])
    kb_list.append([])

# 再來要找區域的像素數量總和
max_pixel_len_area = []
for i in range(0, len(frame_keyboards[0])):
    # i = 0
    # logger.debug(frame_keyboards[0][i])
    max_pixel_len = 0
    try:
        b = frame_keyboards[0][i][0]
        try:
            k = frame_keyboards[0][i][1]
        except Exception as e:
            e
            k = 0
        max_pixel_len = b + k
        max_pixel_len_area.append(max_pixel_len)
    except Exception as e:
        e
        pass

for i in frame_keyboards:
    for j in range(0, len(i)):
        kb_list[j].append(max_pixel_len_area[j] - i[j][0])

# 狀態器初始化
temp_state_list = []  # 狀態器陣列
for i in range(0, len(frame_keyboards[0])):
    temp_state = {
        "st_frame": 0,
        "refractory": False
    }
    temp_state_list.append(temp_state)

# 生成閥值陣列
trigger_valve = []
for i in kb_list:
    mean = int(np.mean(i))
    # logger.debug(mean)
    trigger_valve.append(mean / rc['trigger_valve_parameter'])

# 譜面生成
sheet = [].copy()
for track in range(0, len(kb_list)):
    # track = 1
    for m in range(0, len(frame_keyboards)):
        # m = 0
        # logger.debug('+')
        after_time = (m - temp_state_list[track]['st_frame'])
        refractory_timeout = after_time > cool_down_frame
        trigger = kb_list[track][m] < trigger_valve[track]
        if trigger and refractory_timeout:
            # logger.debug('.')
            temp_state_list[track]['st_frame'] = m
            temp_state_list[track]['refractory'] = True
            sheet.append({"type": "note", "frame": m, "keyboard": track})

sort_sheet = sorted(sheet, key=lambda s: s['frame'])
sort_sheet

# 加個正確時間(frame -> time)
# 順便給個編號
for j in range(0, len(sort_sheet)):
    sort_sheet[j]['time'] = sort_sheet[j]['frame'] / fps
    sort_sheet[j]['note'] = j
analysis_from_video['original_sheet'] = sort_sheet
logger.info('generated original sheet done.')

del analysis_from_video['notes']

analysis_from_video['trigger_valve'] = trigger_valve

# ============================================================
# 經評估後決定把原本負責 index "4_generate_sheet" 內容搬移過來

logger.info('generating native sheet...')

original_sheet = analysis_from_video['original_sheet']
# 基本設定讀取
sheet_formats = rc['sheet_formats'][rc['output_sheet_format']]
sync_area_time = rc['sync_area_time']
sync_symbol = rc['sync_symbol']
blank_symbol = rc['blank_symbol']

# 為了防止 list 在最後倒數14個搜尋 out of range 用的
def get_in_area(n, a, max):
    # n = 6
    # a = 15
    # max = 20
    d = 0
    # max += 1
    if (n + a) > max:
        d = max - (n + a)
    return n + a + d

# 開始將同時按的標註起來
original_sheet_len = len(original_sheet)
osl = original_sheet_len
index = 0
for i in range(0, osl):
    # i = 0
    # i = 1
    # i = 2
    for j in range(i, get_in_area(i, 15, osl)):
        # j = 0
        # j = 1
        # j = 2
        if j == (osl - 1):  # 防止 out of range
            break
        interval = original_sheet[j + 1]['time'] - original_sheet[i]['time']
        if interval * 1000 > sync_area_time:  # 高於差值的就直接下一個
            break
        else:
            if 'index' in original_sheet[i]:  # 初始 index 綁定
                original_sheet[j + 1]['index'] = original_sheet[i]['index']
            else:  # 被初始 index 綁定
                original_sheet[i]['index'] = index
                original_sheet[j + 1]['index'] = index
                index += 1

# 按照 index(同時) 分組
sheet = ""
index_st = ''
for i in range(0, osl):
    # i = 0
    # i = 9
    # i = 10

    if 'index' in original_sheet[i]:
        if original_sheet[i]['index'] == index_st:
            continue
        index_st = original_sheet[i]['index']
        _text_1 = str(sheet_formats[int(original_sheet[i]['keyboard'])])
        _text_2 = ""

        # 接下來在15個音符中搜尋哪個是同時按的
        # (這已被index標註，所以換句話說找接下來15個有沒有跟開頭的index一樣的)
        # ps 設15個是因為鍵盤最多15個，如果之後有增加數量要再改
        # TODO: 看看要不要把這個用base設定的鍵盤數動態生成，畢竟有8個的鍵盤
        for k in range(i, get_in_area(i, 15, osl)):
            # 如果有的話看看index一不一樣
            if 'index' in original_sheet[k]:
                # 一樣就標起來，組合字串
                if original_sheet[k]['index'] == index_st:
                    # 按他base中的譜面格式設定生成要被組合的字串
                    _a = original_sheet[k]['keyboard']
                    _text_2 += str(sheet_formats[_a]) + blank_symbol
                else:
                    break

        # 處理自動排序音符
        # _text_2 = 'B1 A4 C1 B2 '
        if rc['auto_sort_sync_note']:
            _k1 = _text_2[:-1].split(blank_symbol)
            _k2 = sorted(_k1)
            _k3 = ''
            for _k in _k2:
                _k3 += _k + blank_symbol
            _text_3 = _k3[:-1]

        else:
            _text_3 = _text_2[:-1]
        # 組合完畢就用組合符號括起來(預設是 【 】)
        note = f"{sync_symbol[0]}{_text_3}{sync_symbol[1]}"
    else:
        # 沒有index的就直接按base要求組起來就好
        note = str(sheet_formats[int(original_sheet[i]['keyboard'])])
    sheet += note + blank_symbol

logger.info('generated native sheet done.')
# ============================================================

# 先存 output_sheet
logger.info('generating output_sheet.')
output_sheet_path = (aims_folder_path /
                     Path(rc['output_sheet_path'])).resolve()
_temp = output_sheet_path / "penultimate_sheet.txt"
with open(_temp, mode='w', encoding='utf-8') as f:
    f.write(str(sheet))
logger.info('generated output_sheet done.')

# 再存 original_sheet
logger.info('generating native_sheet.')
output_sheet_path = (aims_folder_path /
                     Path(rc['output_sheet_path'])).resolve()
_temp = output_sheet_path / 'native_sheet.json'
_afv = analysis_from_video
with open(_temp, mode='w', encoding='utf-8') as f:
    new_data = {
        "original_sheet": original_sheet,
        "frame_end": _afv['frame_end'],
        "fps": _afv['fps'],
        "duration": _afv['duration'],
        "minute": _afv['minute'],
        "seconds": _afv['seconds'],
        "st_specify_count": _afv['st_specify_count'],
        "ed_specify_count": _afv['ed_specify_count'],
        "trigger_valve": _afv['trigger_valve'],
        "cool_down_frame": cool_down_frame,
    }
    f.write(json.dumps(new_data))
logger.info('generated native_sheet done.')

logger.info('save data done.')
logger.info('Please proceed to the next action.')
input('input any key to exit. 輸入任意值離開.')

#


#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import division

import struct
import onmt
import onmt.markdown
import torch
import argparse
import math
import numpy
import sys
import h5py as h5
import numpy as np
import apex
from scipy.io.wavfile import write
from onmt.model_factory import build_model, build_language_model, optimize_model
from onmt.inference.fast_translator import FastTranslator
from onmt.inference.stream_translator import StreamTranslator
import os

parser = argparse.ArgumentParser(description='translate.py')
onmt.markdown.add_md_help_argument(parser)

parser.add_argument('-model', required=True,
                    help='Path to model .pt file')

parser.add_argument('-src', required=True,
                    help='Source sequence to decode (one line per sequence)')

parser.add_argument('-src_lang', default='src',
                    help='Source language')

parser.add_argument('-concat', type=int, default=1,
                    help="Concate sequential audio features to decrease sequence length")

parser.add_argument('-stride', type=int, default=1,
                    help="Stride on input features")

parser.add_argument('-encoder_type', default='text',
                    help="Type of encoder to use. Options are [text|img|audio].")

parser.add_argument('-output',
                    help="Path to output")

parser.add_argument('-fp16', action='store_true',
                    help='To use floating point 16 in decoding')

parser.add_argument('-asr_format', default="scp", required=False,
                    help="Format of asr data h5 or scp")

parser.add_argument('-gpu', type=int, default=-1,
                    help="Device to run on")

parser.add_argument('-batch_size', type=int, default=1,
                    help='Batch size, currently support batch size 1')

def write_array(fd, array):
    size = 0
    assert isinstance(array, np.ndarray), type(array)
    fd.write(b'\0B')
    size += 2
    dt = array.dtype
    if dt == np.float32 or dt == np.float16:
        atype = b'FM ' if dt == np.float32 else b'HM '
        if len(array.shape) == 2:
            fd.write(atype)
            size += 3
            fd.write(b'\4')
            size += 1
            fd.write(struct.pack('<i', len(array)))  # Rows
            size += 4

            fd.write(b'\4')
            size += 1
            fd.write(struct.pack('<i', array.shape[1]))  # Cols
            size += 4
        fd.write(array.tobytes())
        size += array.nbytes
    else:
        raise ValueError('Unsupported array type: {}'.format(dt))
    return size

def write_ark(ark, dic, scp=None, append=False):
    # Write ark
    mode = 'ab' if append else 'wb'
    pos_list = []
    with open(ark, mode) as fd:
        pos = fd.tell() if append else 0
        for key in dic:
            encode_key = (key + ' ').encode()
            fd.write(encode_key)
            pos += len(encode_key)
            pos_list.append(pos)
            data = dic[key]
            pos += write_array(fd, data)

    # Write scp
    if scp is not None:
        mode = 'a' if append else 'w'
        with open(scp, mode) as fd:
            for key, position in zip(dic, pos_list):
                fd.write(key + u' ' + ark + ':' + str(position) + os.linesep)
def main():
    opt = parser.parse_args()
    opt.cuda = opt.gpu > -1
    if opt.cuda:
        torch.cuda.set_device(opt.gpu)

    # waveglow = torch.hub.load('nvidia/DeepLearningExamples:torchhub', 'nvidia_waveglow')
    # waveglow = waveglow.remove_weightnorm(waveglow)
    # waveglow.eval()

    model = opt.model
    checkpoint = torch.load(model,
                            map_location=lambda storage, loc: storage)

    model_opt = checkpoint['opt']

    model = build_model(model_opt, checkpoint['dicts'])
    optimize_model(model)
    model.load_state_dict(checkpoint['model'])

    if opt.fp16:
        model = model.half()

    if opt.cuda:
        model = model.cuda()
    else:
        model = model.cpu()

    model.eval()
    if opt.encoder_type == "audio" and opt.asr_format == "scp":
        import kaldiio
        from kaldiio import ReadHelper
        audio_data = iter(ReadHelper('scp:' + opt.src))

    in_file = None

    output_ark = opt.output + "/output.ark"
    output_scp =  opt.output + "/output.scp"

    if opt.encoder_type == "audio":

        i = 0
        while True:
            if opt.asr_format == "h5":
                if i == len(in_file):
                    break
                line = np.array(in_file[str(i)])
                i += 1
            elif opt.asr_format == "scp":
                try:
                    seg_name, line = next(audio_data)
                except StopIteration:
                    break

            if opt.stride != 1:
                line = line[0::opt.stride]
            line = torch.from_numpy(line)
            if opt.concat != 1:
                add = (opt.concat - line.size()[0] % opt.concat) % opt.concat
                z = torch.FloatTensor(add, line.size()[1]).zero_()
                line = torch.cat((line, z), 0)
                line = line.reshape((line.size()[0] // opt.concat, line.size()[1] * opt.concat))

            length, feature_size = line.size(0), line.size(1)

            tensor = torch.cat([torch.ones(length).unsqueeze(1),line],dim=-1)

            outputs = model.inference(tensor.unsqueeze(0).half().cuda())
            mel = outputs[0].squeeze(0).transpose(0,1).float().cpu().detach().numpy()
            print(mel.shape)
            dic = {seg_name: mel}
            write_ark(output_ark, dic, output_scp, append=True)

if __name__ == "__main__":
    main()


import odin
import numpy
from pg import DB

class db(DB):
    def __init__(self):
        DB.__init__(self, dbname='odin')

def djl(year, mon, day, hour, min, secs):
    dn = 367 * year - 7 * (year + (mon + 9) / 12) / 4 \
        - 3 * ((year + (mon - 9) / 7) / 100 + 1) / 4 + 275 * mon / 9 \
        + day + 1721029
    jd = float(dn) - 0.5 + (hour + (min + secs / 60.0) / 60.0) / 24.0
    return jd

con = db()
query = con.query(
    '''select year,mon,day,hour,min,secs,stw,orbit,qt,qa,qe,gps,acs
	           from attitude_level0''')
result = query.dictresult()

query = con.query('''select stw,inttime from ac_level0 where sig_type='SIG'
order by stw''')
sigresult = query.dictresult()

for sig in sigresult:
    query = con.query('''select year,mon,day,hour,min,secs,stw,
                           orbit,qt,qa,qe,gps,acs
	                   from attitude_level0 where
                           stw>{0}-200000 and stw<{0}+200000 order by stw'''.format(sig['stw'] + 2**32))
    result = query.dictresult()
    if len(result) > 0:
        stw = float(sig['stw']) + 2**32 - sig['inttime'] * 16.0 / 2.0
        cols = 2 + 2 * 4 + 3 + 6 + 1
        rows = len(result)
        lookup = numpy.zeros(shape=(rows, cols))
        attstw = numpy.zeros(shape=(rows,))
        for index, row in enumerate(result):
            JD = djl(row['year'], row['mon'], row['day'],
                     row['hour'], row['min'], row['secs'])
            attstw[index] = float(row['stw'])
            lookup[index, 0: 2] = numpy.array((JD, row['orbit']))
            qt = eval(row['qt'].replace('{', '(').replace('}', ')'))
            lookup[index, 2: 6] = numpy.array(qt)
            qa = eval(row['qa'].replace('{', '(').replace('}', ')'))
            lookup[index, 6:10] = numpy.array(qa)
            qe = eval(row['qe'].replace('{', '(').replace('}', ')'))
            lookup[index, 10:13] = numpy.array(qe)
            gps = eval(row['gps'].replace('{', '(').replace('}', ')'))
            lookup[index, 13:19] = numpy.array(gps)
            lookup[index, 19] = row['acs']
        ind = (attstw > stw).nonzero()[0]
        if len(ind) == 0:
            pass
            # continue
        if ind[0] == 0:
            pass
            # continue
        # ind=ind[0]
        i = 3
        dt = attstw[i] - attstw[i - 1]
        dt0 = attstw[i] - stw
        dt1 = stw - attstw[i - 1]
        att = (dt0 * lookup[i - 1, :] + dt1 * lookup[i, :]) / dt
        t = (att[0], long(stw), att[1],
             tuple(att[2:6]), tuple(att[6:10]),
             tuple(att[10:13]), tuple(att[13:19]), att[19])
        s = odin.Spectrum()
        s.stw = long(stw)
        s.Attitude(t)
        print s.longitude, s.latitude, s.altitude
        print s.gpspos, s.gpsvel, s.sunpos, s.sunzd
        print s.ra2000, s.dec2000, s.vgeo, s.skybeamhit
        print dir(s)
# (year,mon,day,hour,min,sec,stw,orbit,qt,qa,qe,gps,acs)


"""
Credit: https://github.com/sfujim/TD3
"""
import argparse
import copy
import os
import os.path as osp
import sys
from collections import deque

import gym
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

current_dir = osp.join(osp.abspath(osp.dirname(__file__)))
sys.path.append(current_dir)
sys.path.append(osp.dirname(current_dir))
print(current_dir)

from envs import make_envs

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class ReplayBuffer(object):
    def __init__(self, state_dim, action_dim, max_size=int(1e6)):
        self.max_size = max_size
        self.ptr = 0
        self.size = 0

        self.state = np.zeros((max_size, state_dim))
        self.action = np.zeros((max_size, action_dim))
        self.next_state = np.zeros((max_size, state_dim))
        self.reward = np.zeros((max_size, 1))
        self.not_done = np.zeros((max_size, 1))

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def add(self, state, action, next_state, reward, done):
        self.state[self.ptr] = state
        self.action[self.ptr] = action
        self.next_state[self.ptr] = next_state
        self.reward[self.ptr] = reward
        self.not_done[self.ptr] = 1. - done

        self.ptr = (self.ptr + 1) % self.max_size
        self.size = min(self.size + 1, self.max_size)

    def sample(self, batch_size):
        ind = np.random.randint(0, self.size, size=batch_size)

        return (
            torch.FloatTensor(self.state[ind]).to(self.device),
            torch.FloatTensor(self.action[ind]).to(self.device),
            torch.FloatTensor(self.next_state[ind]).to(self.device),
            torch.FloatTensor(self.reward[ind]).to(self.device),
            torch.FloatTensor(self.not_done[ind]).to(self.device)
        )

class TD3Actor(nn.Module):
    def __init__(self, state_dim, action_dim, max_action):
        super(TD3Actor, self).__init__()

        self.l1 = nn.Linear(state_dim, 256)
        self.l2 = nn.Linear(256, 256)
        self.l3 = nn.Linear(256, action_dim)

        self.max_action = max_action

    def forward(self, state):
        a = F.relu(self.l1(state))
        a = F.relu(self.l2(a))
        return self.max_action * torch.tanh(self.l3(a))

class TD3Critic(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(TD3Critic, self).__init__()

        # Q1 architecture
        self.l1 = nn.Linear(state_dim + action_dim, 256)
        self.l2 = nn.Linear(256, 256)
        self.l3 = nn.Linear(256, 1)

        # Q2 architecture
        self.l4 = nn.Linear(state_dim + action_dim, 256)
        self.l5 = nn.Linear(256, 256)
        self.l6 = nn.Linear(256, 1)

    def forward(self, state, action):
        sa = torch.cat([state, action], 1)

        q1 = F.relu(self.l1(sa))
        q1 = F.relu(self.l2(q1))
        q1 = self.l3(q1)

        q2 = F.relu(self.l4(sa))
        q2 = F.relu(self.l5(q2))
        q2 = self.l6(q2)
        return q1, q2

    def Q1(self, state, action):
        sa = torch.cat([state, action], 1)

        q1 = F.relu(self.l1(sa))
        q1 = F.relu(self.l2(q1))
        q1 = self.l3(q1)
        return q1

class TD3Trainer:
    def __init__(
            self,
            state_dim,
            action_dim,
            max_action,
            discount=0.99,
            tau=0.005,
            policy_noise=0.2,
            noise_clip=0.5,
            policy_freq=2,
            lr=5e-5  # Small LR in TD3 is important to train in MetaDrive!
    ):
        self.actor = TD3Actor(state_dim, action_dim, max_action).to(device)
        self.actor_target = copy.deepcopy(self.actor)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=lr)

        self.critic = TD3Critic(state_dim, action_dim).to(device)
        self.critic_target = copy.deepcopy(self.critic)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=lr)

        self.max_action = max_action
        self.discount = discount
        self.tau = tau
        self.policy_noise = policy_noise
        self.noise_clip = noise_clip
        self.policy_freq = policy_freq

        self.total_it = 0

    def select_action(self, state):
        state = torch.FloatTensor(state.reshape(1, -1)).to(device)
        return self.actor(state).cpu().data.numpy().flatten()

    def train(self, replay_buffer, batch_size=256):
        self.total_it += 1

        # Sample replay buffer
        state, action, next_state, reward, not_done = replay_buffer.sample(batch_size)

        # [TODO] Following the TODOs below to implement critic loss
        with torch.no_grad():
            # Compute the target Q value
            noise = (torch.randn_like(action) * self.policy_noise).clamp(-self.noise_clip, self.noise_clip)
            next_action = (self.actor_target(next_state) + noise).clamp(-self.max_action, self.max_action)
            Q1, Q2 = self.critic_target(next_state, next_action)
            Q = torch.min(Q1, Q2)
            target_Q = reward + not_done * self.discount * Q

        # Get current Q estimates
        current_Q1, current_Q2 = self.critic(state, action)

        # Compute critic loss
        critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(current_Q2, target_Q)

        # Optimize the critic
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        # Delayed policy updates
        if self.total_it % self.policy_freq == 0:

            # [TODO] Compute actor loss
            actor_loss = -self.critic.Q1(state, self.actor(state)).mean()

            # Optimize the actor
            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            self.actor_optimizer.step()

            # Update the frozen target models
            for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):
                target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)

            for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()):
                target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)

    def save(self, filename):
        torch.save(self.critic.state_dict(), filename + "_critic")
        torch.save(self.critic_optimizer.state_dict(), filename + "_critic_optimizer")

        torch.save(self.actor.state_dict(), filename + "_actor")
        torch.save(self.actor_optimizer.state_dict(), filename + "_actor_optimizer")

    def load(self, filename):
        self.critic.load_state_dict(torch.load(filename + "_critic"))
        self.critic_optimizer.load_state_dict(torch.load(filename + "_critic_optimizer"))
        self.critic_target = copy.deepcopy(self.critic)

        self.actor.load_state_dict(torch.load(filename + "_actor"))
        self.actor_optimizer.load_state_dict(torch.load(filename + "_actor_optimizer"))
        self.actor_target = copy.deepcopy(self.actor)

# Runs policy for X episodes and returns average reward
# A fixed seed is used for the eval environment
def eval_policy(policy, eval_env, seed, eval_episodes=10):

    eval_env.seed(seed + 100)

    avg_reward = 0.
    for _ in range(eval_episodes):
        state, done = eval_env.reset(), False
        while not done:
            action = policy.select_action(np.array(state))
            state, reward, done, _ = eval_env.step(action)
            avg_reward += reward

    avg_reward /= eval_episodes

    print("---------------------------------------")
    print(f"Evaluation over {eval_episodes} episodes: {avg_reward:.3f}")
    print("---------------------------------------")
    return avg_reward

def summary(array, name, extra_dict=None):
    array = np.asarray(array, dtype=np.float32)
    ret = {
        "{}_mean".format(name): float(np.mean(array)) if len(array) else np.nan,
        "{}_min".format(name): float(np.min(array)) if len(array) else np.nan,
        "{}_max".format(name): float(np.max(array)) if len(array) else np.nan,
    }
    if extra_dict:
        ret.update(extra_dict)
    return ret

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--env-id", default="MetaDrive-Tut-Easy-v0")
    parser.add_argument("--seed", default=0)
    parser.add_argument("-lr", default=5e-5, type=float)
    parser.add_argument(
        "--log-dir",
        default="data-MetaDrive-Tut-Easy-v0/",
        type=str,
        help="The path of directory that you want to store the data to. "
             "Default: ./data/"
    )
    parser.add_argument("--start_timesteps", default=1e4, type=int)  # Time steps initial random policy is used
    parser.add_argument("--save_freq", default=5e3, type=int)  # How often (time steps) we save model
    parser.add_argument("--max_timesteps", default=1e6, type=int)  # Max time steps to run environment
    parser.add_argument("--expl_noise", default=0.1)  # Std of Gaussian exploration noise
    parser.add_argument("--batch_size", default=256, type=int)  # Batch size for both actor and critic
    parser.add_argument("--discount", default=0.99)  # Discount factor
    parser.add_argument("--tau", default=0.005)  # Target network update rate
    parser.add_argument("--policy_noise", default=0.2)  # Noise added to target policy during critic update
    parser.add_argument("--noise_clip", default=0.5)  # Range to clip target policy noise
    parser.add_argument("--policy_freq", default=2, type=int)  # Frequency of delayed policy updates
    parser.add_argument("--load_model", action="store_true")
    args = parser.parse_args()

    log_dir = args.log_dir
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)

    if not os.path.exists(osp.join(log_dir, "results")):
        os.makedirs(osp.join(log_dir, "results"))

    if not os.path.exists(osp.join(log_dir, "models")):
        os.makedirs(osp.join(log_dir, "models"))

    environments = make_envs(
        env_id=args.env_id,
        log_dir=log_dir,
        num_envs=1,
        asynchronous=False,
    )
    env = environments.envs[0]

    # Set seeds
    env.seed(args.seed)
    env.action_space.seed(args.seed)
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.shape[0]
    max_action = float(env.action_space.high[0])

    kwargs = {
        "state_dim": state_dim,
        "action_dim": action_dim,
        "max_action": max_action,
        "discount": args.discount,
        "tau": args.tau,
    }

    # Target policy smoothing is scaled wrt the action scale
    kwargs["policy_noise"] = args.policy_noise * max_action
    kwargs["noise_clip"] = args.noise_clip * max_action
    kwargs["policy_freq"] = args.policy_freq
    kwargs["lr"] = args.lr
    policy = TD3Trainer(**kwargs)

    discrete = False
    max_size = 1e-6
    replay_buffer = ReplayBuffer(state_dim, action_dim)

    state, done = env.reset(), False
    episode_reward = 0
    episode_timesteps = 0
    episode_num = 0

    reward_recorder = deque(maxlen=100)

    for t in range(int(args.max_timesteps)):

        episode_timesteps += 1

        # Select action randomly or according to policy
        if t < args.start_timesteps:
            action = env.action_space.sample()
        else:
            action = (
                    policy.select_action(np.array(state))
                    + np.random.normal(0, max_action * args.expl_noise, size=action_dim)
            ).clip(-max_action, max_action)

        # Perform action
        next_state, reward, done, _ = env.step(action)
        done_bool = float(done)  # if episode_timesteps < env._max_episode_steps else 0

        if args.load_model:
            # Modify this to load proper models!
            policy.load(f"{log_dir}/models/default")

        # Store data in replay buffer
        replay_buffer.add(state, action, next_state, reward, done_bool)

        state = next_state
        episode_reward += reward

        # Train agent after collecting sufficient data
        if t >= args.start_timesteps:
            policy.train(replay_buffer, args.batch_size)

        if done:
            # record reward
            reward_recorder.append(episode_reward)

            # +1 to account for 0 indexing. +0 on ep_timesteps since it will increment +1 even if done=True
            print(
                f"Total T: {t + 1} Episode Num: {episode_num + 1} Episode T: {episode_timesteps} Reward: {episode_reward:.3f}")
            # Reset environment
            state, done = env.reset(), False
            episode_reward = 0
            episode_timesteps = 0
            episode_num += 1

        if (t + 1) % args.save_freq == 0:
            # evaluations.append(eval_policy(policy, args.env, args.seed))
            stats = {
                'training_episode_reward' : np.mean(reward_recorder),
                'total_step' : t
            }
            filename="step-{}".format(t)
            with open(f"{log_dir}/results/{filename}", 'wb') as f:
                np.save(f, stats)
            policy.save(f"{log_dir}/models/{filename}")