#!/usr/bin/env python3
#encoding: utf-8

#%%
import os ,sys
import argparse
import datetime
from os.path import join as osp
from os.path import exists as ose
from shutil import copytree, rmtree
from math import ceil
import json
import numpy as np
from models import encode_sentences
from models import build_pretrained_models
import pandas as pd
from keras.preprocessing.text import Tokenizer
from keras.preprocessing.sequence import pad_sequences
from keras_image_caption_data_generator import MultimodalInputDataGenerator as datagen
from keras.preprocessing.image import ImageDataGenerator as IDG
from models import concept_detector
from keras import backend  as K
from keras.models import load_model, model_from_json
from layers import L2Normalize
import json
import tensorflow as tf
from utils import caption_image

try:
    import cPickle as pkl
except ImportError:
    import pickle as pkl

if 'tensorflow' == K.backend():

    from keras.backend.tensorflow_backend import set_session

def mil_squared_error(y_true, y_pred):
    return K.tile(K.square(K.max(y_pred) - K.max(y_true)), 16)
#%%
if __name__ == '__main__':

    parser = argparse.ArgumentParser('Visual semantic embeddings')
    parser.add_argument('--model_file', type=str,default = None)
    parser.add_argument('--test_csv_file', type=str)
    parser.add_argument('--source_dataset', type=str)
    parser.add_argument('--tokenizer_pkl_file_id', dest="train_file_id", type=str)
    parser.add_argument('--exp_id', dest="exp_id", type=str)
    parser.add_argument('--model_train_timestamp', dest="train_timestamp", type=str)

    parser.add_argument('--length', type=int, default=None)
    parser.add_argument('--dataaug', default=False,  action="store_true")
    parser.add_argument('--maxtokencount', type=int, default=32198)
    parser.add_argument('--fix_gpu', type=int, default=-1)
    parser.add_argument('--verbose', default=False,  action="store_true")
    parser.add_argument('--image_only_model', default=False,  action="store_true")
    parser.add_argument('--restore_checkpoint', default=False,  action="store_true")
    parser.add_argument('--class_ct_threshold', default = 120, type = int)

    args = parser.parse_args()

    output_id =  args.exp_id +"_"+args.train_file_id
    debug = False
    verbose = args.verbose
    K.set_floatx('float32')
    batch_size = 32

    model_fname = "{}_keras_vse_model-{}".format(args.train_file_id,args.train_timestamp)

    #Depending on the source data copy the images to local storage  a subdir of /export/u10
    dataset_localized = False
    timestamp = datetime.datetime.now().strftime("%Y_%m_%d_%H_%M")
    if args.source_dataset=="GI":
        KERAS_DATAGEN_DIR = "/nfs/mercury-11/u113/projects/AIDA/GoogleImageDownload_Rus_Scenario/image_data_links"
        regex_exp = r'/nfs/mercury-11/u113/projects/AIDA/GoogleImageDownload_Rus_Scenario/image_data_links(.*)'
        LOCAL_STORAGE_DIR = "/export/u10/sadali/AIDA/images/GoogleImageDownload_Rus_Scenario/squared"
        replace_regex_exp = r'/export/u10/sadali/AIDA/images/GoogleImageDownload_Rus_Scenario/squared\1'
    # try:
        if not os.path.exists (osp(LOCAL_STORAGE_DIR,"successful_local_clone")):
            print ("trying to copy to local storage")
            try:
              os.system("resize_and_copy_local.sh valid_images_unique.txt" )
              # copytree(KERAS_DATAGEN_DIR,LOCAL_STORAGE_DIR)
              dataset_localized = True
              with open(osp(LOCAL_STORAGE_DIR,"successful_local_clone"),"w") as fh:
                fh.write(timestamp+"\n")

            except Exception as e:
                dataset_localized = False

    elif args.source_dataset=="VG":
        KERAS_DATAGEN_DIR = "/nfs/mercury-11/u113/projects/AIDA/VisualGenomeData/image_data"
        regex_exp = r'/nfs/mercury-11/u113/projects/AIDA/VisualGenomeData/image_data(.*)'
        LOCAL_STORAGE_DIR = "/export/u10/sadali/AIDA/images/VisualGenomeData/image_data"
        replace_regex_exp = r'/export/u10/sadali/AIDA/images/VisualGenomeData/image_data\1'

        if not os.path.exists (osp(LOCAL_STORAGE_DIR,"VG_100K")):
            print ("copyying VG data from ",KERAS_DATAGEN_DIR,LOCAL_STORAGE_DIR)
            try:
                copytree(KERAS_DATAGEN_DIR,LOCAL_STORAGE_DIR)
                dataset_localized = True
            except Exception as e:
                print (e)
                print ("Unable to copy image files for {} ".format(args.source_dataset) )
                dataset_localized = False
        else:
            dataset_localized = True
    elif args.source_dataset=="AIDASeedling":
        KERAS_DATAGEN_DIR = "/nfs/raid66/u12/users/rbock/aida/image_captions/data_collection/images"
        regex_exp = r'/nfs/raid66/u12/users/rbock/aida/image_captions/data_collection/images(.*)'
        LOCAL_STORAGE_DIR = "/export/u10/sadali/AIDA/images/AIDASeedling/image_data"
        replace_regex_exp = r'/export/u10/sadali/AIDA/images/AIDASeedling/image_data\1'
        if not os.path.exists (osp(LOCAL_STORAGE_DIR,"LDC2018E01")):
            print ("copyying VG data from ",KERAS_DATAGEN_DIR,LOCAL_STORAGE_DIR)
            try:
                copytree(KERAS_DATAGEN_DIR,LOCAL_STORAGE_DIR)
                dataset_localized = True
            except Exception as e:
                print (e)
                print ("Unable to copy image files for {} ".format(args.source_dataset) )
                dataset_localized = False
                print ("Removing any localized dirs")
                try:
                    rmtree(osp(LOCAL_STORAGE_DIR,"LDC2018E01"))
                except Exception as e:
                    print (e)
                try:
                    rmtree(osp(LOCAL_STORAGE_DIR,"LDC2018E52"))
                except Exception as e:
                    print(e)
        else:
            dataset_localized = True

    has_labels = True

    gpu_id = 1
    gpu_id_str = str(int(gpu_id))
    if args.fix_gpu >= 0:
        print ("Overwriting gpu id from cfg file with given arg {}".format(args.fix_gpu))
        gpu_id_str = str(args.fix_gpu)
    #%%

    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    config.gpu_options.visible_device_list = gpu_id_str
    def check_gpu_availability():
        print ("checking if gpus  are available and seen by keras/tf")
        from tensorflow.python.client import device_lib
        assert 'GPU' in str(device_lib.list_local_devices())

        # confirm Keras sees the GPU

        assert len(K.tensorflow_backend._get_available_gpus()) > 0

    #check_gpu_availability()
    #session = tf.Session(config=config)
    set_session(tf.Session(config=config))
    test_df = pd.read_csv(args.test_csv_file, encoding='utf8')
    if verbose:
        print("CSV file columnar data types")
        print( test_df.apply(lambda x: pd.lib.infer_dtype(x.values)))
    texts = test_df["image_captions"].values.tolist()
    if has_labels:
        class_names_pd =  pd.unique(test_df["class"].values)
        init_classnames =class_names_pd.tolist()

        class_counts = test_df["class"].value_counts()
        class_counts.to_csv("class_counts_test_orig.csv")
        class_ct_threshold = args.class_ct_threshold

        untrainable_classes    = class_counts < class_ct_threshold
        untrainable_classnames = untrainable_classes[untrainable_classes].index.tolist()
        if verbose:
            print ("Removed classes:\n",  untrainable_classnames)
            print ("length of test_df",len(test_df))
        test_df = test_df.loc[~test_df['class'].isin(untrainable_classnames),:]

    #Update the filepaths if images were copied to local storage
        if dataset_localized :
            test_df =test_df.replace(KERAS_DATAGEN_DIR,LOCAL_STORAGE_DIR,regex= True)
        print ("new examplar count {}".format(len(test_df)))

        classnames= [k for k in init_classnames if k not in untrainable_classnames]
        if verbose:
            print("Num of classes ")
            print (len(classnames))
        new_class_counts = test_df["class"].value_counts()
        new_class_counts.to_csv("class_counts_test.csv")
    try:
        if not args.restore_checkpoint:

            class_indices_json = "./models_dir/{}/{}_{}_class_indices.json".format(output_id,model_fname,args.class_ct_threshold)
        else:
            class_indices_json = "./models_dir/GI_image_only_softmax_GI_keras_train_qa_image_only_epoch_3/GI_keras_train_qa_image_only_epoch_3_keras_vse_model-2019_01_09_23_31_class_indices.json"
        print("class_indices file is located ",class_indices_json)
        assert os.path.exists(class_indices_json)
        with open (class_indices_json,"r") as json_fh:
            class_indices_for_model = json.load(json_fh, encoding="utf8")

    except Exception as e:
        print (e)
        #class_dirs=os.listdir(KERAS_DATAGEN_DIR)

        with open("GI_class_counts.csv","r") as fh:
            classnames_orig =[ line.strip().split()[0] for line in  fh.readlines() ]
        classnames_ordered = np.sort(np.array(classnames_orig)).tolist()
        print ("This is a temp hack.Should not be necessary if class_indices.json is available")
        #classnames_ordered = ["class_{}".format(i) for i in range(854)]
        class_indices_for_model = dict(zip(classnames_ordered,range(len(classnames_ordered))))
        #print(class_indices_for_model)

    if dataset_localized:
        test_df =test_df.replace(KERAS_DATAGEN_DIR, LOCAL_STORAGE_DIR, regex= True)

    if verbose:
        print("Dimensions of training set dataframe")
        print(test_df.shape)

    texts_ascii = [k.encode('ascii','ignore').decode() for k in texts]
    test_tokenizer = Tokenizer(num_words=args.maxtokencount)

    if ose('./{}/keras_captiontokenizer_{}.pkl'.format(output_id,args.train_file_id)):
        tokenizer_pkl_file = './{}/keras_captiontokenizer_{}.pkl'.format(output_id,args.train_file_id)
    else:
        tokenizer_pkl_file = 'keras_captiontokenizer_{}.pkl'.format(args.train_file_id)

    with open(tokenizer_pkl_file,"rb")  as kfh:
        tokenizer = pkl.load(kfh)
    tokenizer.fit_on_texts(texts_ascii)

    print (type(texts[0]))
    word_index = tokenizer.word_index
    print('Found %s unique tokens.' % len(word_index))

    if  not args.restore_checkpoint:
        end2endmodel = load_model(args.model_file,custom_objects={'L2Normalize':L2Normalize})

        if debug:
            #print("dense_1 wts: \n", end2endmodel.get_layer('dense_1').get_weights())
            if not args.image_only_model:
                print("gru_1 wts: \n", end2endmodel.get_layer('gru_1').get_weights())
                print("dense_2 wts: \n", end2endmodel.get_layer('dense_2').get_weights())
        with open("layer_weights_at_test_time.txt","w") as fh:
            fh.write("dense_1 wts: \n")
            fh.write(str( end2endmodel.get_layer('dense_1').get_weights()))
            if not args.image_only_model:
                fh.write("\ngru_1 wts: \n")
                fh.write(str( end2endmodel.get_layer('gru_1').get_weights()))
                fh.write("\ndense_2 wts: \n")
                fh.write(str( end2endmodel.get_layer('dense_2').get_weights()))
            fh.write("\nblock5_conv4  wts: \n")
            fh.write(str(end2endmodel.get_layer('block5_conv4').get_weights()))
        end2endmodel.compile(optimizer='nadam', loss="categorical_crossentropy")
    else:
        try:

            timestamp = "2018_12_20_22_03"
            timestamp = args.train_timestamp

            #model_fname = "GI_keras_train_qa"+"_keras_vse_model-{}".format(timestamp)
            model_fname = "{}_keras_vse_model-{}".format(args.train_file_id,timestamp)
            print (model_fname)
            model_fpath = "./models_dir/{}.json".format(model_fname)
            print (model_fpath)
            model_json = None
            if ose(model_fpath):
                pass
            else:
                model_fpath = "./models_dir/{}/{}.json".format(output_id,model_fname)
            with open(model_fpath,"r") as fh:
                model_json_str = fh.read()
            end2endmodel = model_from_json(model_json_str,
                                            custom_objects={'L2Normalize':L2Normalize} )
            print ("loaded from json")
            end2endmodel.load_weights(args.model_file)
            print ("loaded weights")
            print("\ndense_1 wts: \n", end2endmodel.get_layer('dense_1').get_weights())
        except Exception as e:
            print (e)

    #print ("Model summary",end2endmodel.summary())
    #end2endmodel.compile(optimizer='nadam', loss="binary_crossentropy")
    model_output_dim = end2endmodel.outputs[0].shape[1]
    print ("model output layer shape",end2endmodel.outputs[0].shape)
    print ("model output layer dim",model_output_dim)
    #sys.exit(0)
    test_datagen = None
    if args.image_only_model:
        if args.dataaug:
            test_datagen = IDG(width_shift_range = 0.2,zoom_range=0.2,rotation_range=25, height_shift_range=0.3 )
        else:
            test_datagen = IDG()
    else:
        if args.dataaug:
            test_datagen = datagen(width_shift_range = 0.2,zoom_range=0.2,rotation_range=25, height_shift_range=0.3 )
        else:
            test_datagen = datagen()
    print ("setting up image data generator")
    print("type(test_datagen)" ,type(test_datagen))
    if  args.image_only_model:
        test_data_it = test_datagen.flow_from_dataframe(
                                                        dataframe= test_df,
                                                        directory= None,
                                                        x_col="filenames",
                                                        y_col="class" if has_labels else None,
                                                        has_ext=True,
                                                        target_size=(256, 256), color_mode='rgb',
                                                        class_mode=None,
                                                        batch_size=batch_size, shuffle=False, seed=None,
                                                        save_to_dir=None,
                                                        save_prefix='',
                                                        save_format='png',
                                                        subset=None,
                                                        interpolation='nearest',
                                                        sort=False,
                                                        follow_links= True)
    else:
        test_data_it = test_datagen.flow_from_dataframe(
                                                        dataframe= test_df,
                                                        directory= None,
                                                        x_col=["filenames","image_captions"],
                                                        y_col="class" if has_labels else None,
                                                        has_ext=True,
                                                        target_size=(256, 256), color_mode='rgb',
                                                        class_mode=None,
                                                        batch_size=batch_size , shuffle=False, seed=None,
                                                        save_to_dir=None,
                                                        save_prefix='',
                                                        save_format='png',
                                                        subset=None,
                                                        interpolation='nearest',
                                                        sort=False,
                                                        cap_token_vocab = word_index,
                                                        num_tokens = len(word_index),
                                                        drop_duplicates = True,
                                                        follow_links    = True)
    # predictions = end2endmodel.predict_generator(test_data_it)
    # preds_out = open("preds_out.txt","w")
    # for pr in predictions:
    #   print(pr)
    #   preds_out.write("{}\n".format(pr))
    batch_ctr  = 0
    output_dir = "/export/u10/sadali/AIDA/images/captioned_images/{}-{}".format(output_id, model_fname)

    if not os.path.exists (output_dir):
        os.makedirs(output_dir)

    end_of_samples = False
    #while True:
    b_it = 0
    sample_count = len(test_data_it.filenames)
    batch_ctr_max = ceil(sample_count/batch_size )
    while not end_of_samples:
        batch_indices = []
        batch_idx =None
        for batch_idx in next(test_data_it.index_generator, -1):
            if batch_idx < 0:
                print(batch_idx)
                end_of_samples = True
                break
            #print (batch_ctr)

            batch_indices.append(batch_idx)
            #files_in_batch = test_df["filenames"][example_it:batch_end].values.tolist()

            b_it += 1
            if b_it == batch_size:
                break
        if batch_idx <0 :
            end_of_samples = True
        #print(batch_indices)
        test_batch     = test_data_it._get_batches_of_transformed_samples(batch_indices)
        files_in_batch = [test_data_it.filenames[k] for k in batch_indices]

        #print("files_in_batch:",str(files_in_batch))
        preds_out = end2endmodel.predict_on_batch(test_batch)
        print("predictions tensor shape", preds_out.shape)
        if not os.path.exists("./{}".format(output_id)):
            os.makedirs(output_id)
        preds_out_file = open("./{}/{}_{}_{}_preds_out.txt".format(output_id,
                               args.train_file_id,
                               args.train_timestamp,timestamp),"w")
        print("predictions")
        for pr in preds_out:
            preds_out_file.write("{}\n".format(pr))
        preds_out_file.close()
        print("starting captioning")
        highest_idx = np.argmax(preds_out, axis=1)
        found_highest_idx = False
        for b_i,f in enumerate(files_in_batch):
            concept_score_triples = []

            highest_class = ""
            for k,v in class_indices_for_model.items():
                new_tri = (k, preds_out[b_i, v], preds_out[b_i, v])
                #new_tri= (k, preds_out[v,b_i], preds_out[v,b_i])
                if v == highest_idx[b_i]:
                  highest_class = k
                  highest_score = preds_out[b_i, v]
                  found_highest_idx =True
                concept_score_triples.append(new_tri)
            if not found_highest_idx:
                print( "{} the class name not found".format(highest_idx[b_i] ) )
            caption_image(f, concept_score_triples, output_dir,
                    caption_threshold = 0.08 , trans_dict = None, true_classname= None,
                    highest_pair     = [highest_class, highest_score] if found_highest_idx else None  )
        batch_ctr += 1
        if batch_ctr == batch_ctr_max:
            break
        if batch_ctr % 200 == 0 :
            print ("{}th batch of images used on model" .format(batch_ctr))
        # if end_of_samples :
        #     break


import utils
from pycocotools.coco import COCO
import os
import torch
import numpy as np
from config import Config
import random
from coco_mod import *
from model import *
import pandas as pd
import cv2

class CocoConfig(Config):
    NAME = "coco"
    IMAGES_PER_GPU = 1
    # GPU_COUNT = 8
    NUM_CLASSES = 2   # COCO has 80 classes

def load_image(image_id):
    image = cv2.imread(image_id)
    if image.shape[-1] !=3:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    return image

def load_mask(image_id):
    mask = np.empty([0, 0, 0])
    class_ids = np.empty([0], np.int32)
    return mask, class_ids

def load_image_gt_mod(config, image_id, augment=False, use_mini_mask=False):
    image = load_image(image_id)
    mask, class_ids = load_mask(image_id)
    shape = image.shape
    image, window, scale, padding = utils.resize_image(
        image,
        min_dim=config.IMAGE_MIN_DIM,
        max_dim=config.IMAGE_MAX_DIM,
        padding=config.IMAGE_PADDING)
    mask = utils.resize_mask(mask, scale, padding)

    # Random horizontal flips.
    if augment:
        if random.randint(0, 1):
            image = np.fliplr(image)
            mask = np.fliplr(mask)

    bbox = utils.extract_bboxes(mask)
    active_class_ids = np.zeros([config.NUM_CLASSES], dtype=np.int32)
    source_class_ids = dataset.source_class_ids[dataset.image_info[image_id]["source"]]
    active_class_ids[source_class_ids] = 1
    # Resize masks to smaller size to reduce memory usage
    if use_mini_mask:
        mask = utils.minimize_mask(bbox, mask, config.MINI_MASK_SHAPE)
    # Image meta data
    image_meta = compose_image_meta(image_id, shape, window, active_class_ids)

    return image, image_meta, class_ids, bbox, mask

class Dataset(torch.utils.data.Dataset):
    def __init__(self, df_path, config, augment=True):
        df = pd.read_csv(df_path)
        self.b = 0  # batch item index
        self.image_index = -1
        self.image_ids = np.copy(df["file_name"].to_list())
        self.error_count = 0

        self.config = config
        self.augment = augment

        # Anchors
        # [anchor_count, (y1, x1, y2, x2)]
        self.anchors = utils.generate_pyramid_anchors(config.RPN_ANCHOR_SCALES,
                                                 config.RPN_ANCHOR_RATIOS,
                                                 config.BACKBONE_SHAPES,
                                                 config.BACKBONE_STRIDES,
                                                 config.RPN_ANCHOR_STRIDE)

    def __getitem__(self, image_index):
        image_id = self.image_ids[image_index]
        image, image_metas, gt_class_ids, gt_boxes, gt_masks = \
            load_image_gt_mod(self.config, image_id, augment=self.augment,
                          use_mini_mask=self.config.USE_MINI_MASK)
        if not np.any(gt_class_ids > 0):
            return None

        # RPN Targets
        rpn_match, rpn_bbox = build_rpn_targets(image.shape, self.anchors,
                                                gt_class_ids, gt_boxes, self.config)

        # If more instances than fits in the array, sub-sample from them.
        if gt_boxes.shape[0] > self.config.MAX_GT_INSTANCES:
            ids = np.random.choice(
                np.arange(gt_boxes.shape[0]), self.config.MAX_GT_INSTANCES, replace=False)
            gt_class_ids = gt_class_ids[ids]
            gt_boxes = gt_boxes[ids]
            gt_masks = gt_masks[:, :, ids]

        # Add to batch
        rpn_match = rpn_match[:, np.newaxis]
        images = mold_image(image.astype(np.float32), self.config)

        # Convert
        images = torch.from_numpy(images.transpose(2, 0, 1)).float()
        image_metas = torch.from_numpy(image_metas)
        rpn_match = torch.from_numpy(rpn_match)
        rpn_bbox = torch.from_numpy(rpn_bbox).float()
        gt_class_ids = torch.from_numpy(gt_class_ids)
        gt_boxes = torch.from_numpy(gt_boxes).float()
        gt_masks = torch.from_numpy(gt_masks.astype(int).transpose(2, 0, 1)).float()

        return images, image_metas, rpn_match, rpn_bbox, gt_class_ids, gt_boxes, gt_masks

    def __len__(self):
        return self.image_ids.shape[0]

dataset_dir = "/home/nitheesh/Documents/projects_3/maskrcnnpytorch/pytorch-mask-rcnn/data/anno_cell_new_set.csv"
subset = "train"
config = CocoConfig()
train_set = Dataset(dataset_dir, config, augment=True)


import torch
import torch.nn.functional as F
import numpy as np
from torch.optim import Adam, SGD
from gym import spaces
from itertools import chain
from utils.misc import soft_update, hard_update, enable_gradients, disable_gradients, pol_kl, RunningMeanStd, apply_to_all_elements
from utils.agents import Agent
from utils.policies import HeadSelector
from utils.critics import CentralCritic

MSELoss = torch.nn.MSELoss()
SmoothL1Loss = torch.nn.SmoothL1Loss()
CrossEntropyLoss = torch.nn.CrossEntropyLoss()

class SAC(object):
    """
    Training decentralized policies w/ centralized critic using
    (discrete action) Soft Actor-Critic
    """
    def __init__(self, nagents, obs_shape,
                 state_shape, action_size, gamma_e=0.95, gamma_i=0.95, tau=0.01,
                 hard_update_interval=None,
                 pi_lr=0.01, q_lr=0.01, phi_lr=0.1,
                 adam_eps=1e-8,
                 q_decay=1e-3, phi_decay=1e-4, reward_scale=10.,
                 head_reward_scale=25.,
                 pol_hidden_dim=64, critic_hidden_dim=64, nonlin=F.relu,
                 n_intr_rew_types=0,
                 sep_extr_head=False, beta=0.5, **kwargs):
        """
        Inputs:
            obs_shape (int): Dimensions of vector observations
            state_shape (int): Dimensions of vector global state
            gamma (float): Discount factor
            tau (float): Target update rate
            pi_lr (float): Learning rate for policy
            q_lr (float): Learning rate for critic
            reward_scale (float): Scaling for reward (has effect of optimal
                                  policy entropy)
            hidden_dim (int): Number of hidden dimensions for networks
        """
        self.nagents = nagents
        n_pol_heads = n_intr_rew_types + int(sep_extr_head)

        self.agents = [Agent(obs_shape, action_size,
                             lr=pi_lr, adam_eps=adam_eps,
                             hidden_dim=pol_hidden_dim,
                             nonlin=nonlin,
                             n_pol_heads=n_pol_heads)
                       for _ in range(nagents)]

        self.critic = CentralCritic(state_shape[0], action_size,
                                    nagents, hidden_dim=critic_hidden_dim,
                                    n_intr_rew_heads=n_intr_rew_types,
                                    sep_extr_head=sep_extr_head)
        self.target_critic = CentralCritic(state_shape[0], action_size,
                                           nagents,
                                           hidden_dim=critic_hidden_dim,
                                           n_intr_rew_heads=n_intr_rew_types,
                                           sep_extr_head=sep_extr_head)
        hard_update(self.target_critic, self.critic)
        self.critic_optimizer = Adam(self.critic.parameters(), lr=q_lr,
                                     eps=adam_eps, weight_decay=q_decay)

        self.gamma_e = gamma_e
        self.gamma_i = gamma_i
        self.tau = tau
        self.hard_update_interval = hard_update_interval
        self.pi_lr = pi_lr
        self.q_lr = q_lr
        self.reward_scale = reward_scale
        self.head_reward_scale = head_reward_scale
        self.n_intr_rew_types = n_intr_rew_types
        self.sep_extr_head = sep_extr_head  # separate policy head only trained on extr rews
        self.n_pol_heads = n_pol_heads
        self.beta = beta
        self.head_selector = HeadSelector(nagents, n_pol_heads)
        self.head_selector_optimizer = SGD(self.head_selector.parameters(),
                                           lr=phi_lr, weight_decay=phi_decay)
        self.curr_pol_heads = self.sample_pol_heads()
        self.grad_norm_clip = 10
        self.niter = 0

    @property
    def policies(self):
        return [a.policy for a in self.agents]

    @property
    def target_policies(self):
        return [a.target_policy for a in self.agents]

    def update_heads_onpol(self, mc_rets, ret_rms, soft=True, logger=None):
        """
        Update policy head selector(s) using monte carlo rollouts
        """
        curr_device = next(self.head_selector.parameters()).device
        heads = torch.tensor(self.curr_pol_heads[:1], device=curr_device).view(-1, 1)

        # get log_pi for all heads
        _, all_probs, all_log_probs, entropy = self.head_selector(
            return_all_probs=True, return_all_log_probs=True, return_entropy=True)
        all_probs = all_probs.flatten()
        log_pi = all_log_probs.gather(1, heads)

        q = torch.tensor(mc_rets.mean(), dtype=torch.float32,
                         device=curr_device)
        v = torch.tensor([rrms.mean for rrms in ret_rms], dtype=torch.float32,
                         device=curr_device)
        v = (v * all_probs).sum()

        if soft:
            loss = (-log_pi * ((-log_pi / self.head_reward_scale) + q - v).detach()).mean()
        else:
            loss = (-log_pi * (q - v).detach()).mean()
        self.head_selector_optimizer.zero_grad()
        loss.backward()
        self.head_selector_optimizer.step()
        if logger is not None:
            logger.add_scalar('head_entropy', entropy[0], self.niter)

            for i in range(self.n_pol_heads):
                logger.add_scalar('bandits/head%i' % i,
                                  all_probs[i], self.niter)

    def sample_pol_heads(self, uniform=False):
        """
        Sample new policy heads
        """
        if uniform:
            heads = [np.random.randint(self.n_pol_heads)]
        else:
            heads = self.head_selector().flatten().tolist()
        heads *= self.nagents
        return heads

    def set_pol_heads(self, pol_heads):
        """
        Set policy heads for rollouts
        """
        self.curr_pol_heads = pol_heads

    def step(self, observations, explore=False):
        """
        Take a step forward in environment with all agents
        Inputs:
            observations: List of observations for each agent
        Outputs:
            actions: List of actions for each agent
        """
        return [a.step(obs, explore=explore, head=h)
                for a, obs, h in zip(self.agents, observations, self.curr_pol_heads)]

    def update_critic(self, sample, soft=True, logger=None, intr_rews=None,
                      **kwargs):
        """
        Update central critic for all agents
        """
        state, obs, acs, rews, next_state, next_obs, dones = sample
        # Q loss
        next_acs = []
        next_log_pis = []
        for pi, ob in zip(self.target_policies, next_obs):
            pi_outs = pi(ob, return_log_pi=True)
            curr_next_ac, curr_next_log_pi = list(zip(*pi_outs))  # these contain entries per policy head
            next_acs.append(curr_next_ac)
            next_log_pis.append(curr_next_log_pi)
        trgt_critic_in = (next_state, next_obs, next_acs)
        critic_in = (state, obs, acs)
        next_qs = self.target_critic(trgt_critic_in)
        pred_qs = self.critic(critic_in)
        q_loss = 0

        for a_i, (neqs, niqs), log_pi, (peqs, piqs) in zip(
                range(self.nagents), next_qs, next_log_pis, pred_qs):
            if self.sep_extr_head:
                target_q = (rews.view(-1, 1) +
                            self.gamma_e * neqs[0] *
                            (1 - dones.view(-1, 1)))
                if soft:
                    target_q -= log_pi[0] / self.reward_scale
                q_loss += MSELoss(peqs[0], target_q.detach())
                neqs = neqs[1:]
                peqs = peqs[1:]
                log_pi = log_pi[1:]

            for i, neq, peq in zip(range(self.n_intr_rew_types), neqs, peqs):
                target_q = (rews.view(-1, 1) +
                            self.gamma_e * neq *
                            (1 - dones.view(-1, 1)))
                if soft:
                    target_q -= log_pi[i] / self.reward_scale
                q_loss += MSELoss(peq, target_q.detach())

            for i, niq, piq, in zip(range(self.n_intr_rew_types), niqs, piqs):
                target_q = (intr_rews[i][a_i].view(-1, 1) +
                            self.gamma_i * niq *
                            (1 - dones.view(-1, 1)))
                if soft:
                    target_q -= log_pi[i] / self.reward_scale
                q_loss += SmoothL1Loss(piq, target_q.detach())  # more robust to outliers

        q_loss.backward()
        self.critic.scale_shared_grads()
        grad_norm = torch.nn.utils.clip_grad_norm_(self.critic.parameters(), self.grad_norm_clip)
        self.critic_optimizer.step()
        self.critic_optimizer.zero_grad()

        if logger is not None:
            logger.add_scalar('grad_norms/critic', grad_norm, self.niter)
            logger.add_scalar('losses/q_loss', q_loss, self.niter)
        self.niter += 1

    def update_policies(self, sample, soft=True, logger=None,
                        **kwargs):
        state, obs, acs, rews, next_state, next_obs, dones = sample
        samp_acs = []
        all_probs = []
        all_log_probs = []
        all_log_pis = []
        all_pol_regs = []
        n_pol_heads = self.n_intr_rew_types + int(self.sep_extr_head)
        for a_i, pi, ob in zip(range(self.nagents), self.policies, obs):
            pi_outs = pi(
                ob, return_all_probs=True, return_all_log_probs=True, return_log_pi=True,
                regularize=True, return_entropy=True, head=None)
            curr_ac, probs, log_probs, log_pi, pol_regs, ent = list(zip(*pi_outs))
            for j in range(n_pol_heads):
                logger.add_scalar('agent%i/pol%i_entropy' % (a_i, j), ent[j],
                                  self.niter)
            samp_acs.append(curr_ac)
            all_probs.append(probs)
            all_log_probs.append(log_probs)
            all_log_pis.append(log_pi)
            all_pol_regs.append(pol_regs)

        critic_in = (state, obs, samp_acs)
        critic_rets = self.critic(critic_in, return_all_q=True)
        for a_i, probs, log_pi, pol_regs, ((eqs, iqs), (all_eqs, all_iqs)) in zip(
                range(self.nagents), all_probs, all_log_pis, all_pol_regs,
                      critic_rets):
            curr_agent = self.agents[a_i]
            pol_loss = 0.0
            if self.sep_extr_head:
                all_q = all_eqs[0]
                q = eqs[0]
                v = (all_q * probs[0]).sum(dim=1, keepdim=True)
                pol_target = q - v
                if soft:
                    pol_loss += (log_pi[0] * (log_pi[0] / self.reward_scale - pol_target).detach()).mean()
                else:
                    pol_loss += (log_pi[0] * (-pol_target).detach()).mean()
                for reg in pol_regs[0]:
                    pol_loss += 1e-3 * reg  # policy regularization
                eqs = eqs[1:]
                all_eqs = all_eqs[1:]
                probs = probs[1:]
                log_pi = log_pi[1:]
                pol_regs = pol_regs[1:]
            for j in range(self.n_intr_rew_types):
                q = eqs[j] + self.beta * iqs[j]
                all_q = all_eqs[j] + self.beta * all_iqs[j]
                v = (all_q * probs[j]).sum(dim=1, keepdim=True)
                pol_target = q - v
                if soft:
                    pol_loss += (log_pi[j] * (log_pi[j] / self.reward_scale - pol_target).detach()).mean()
                else:
                    pol_loss += (log_pi[j] * (-pol_target).detach()).mean()
                for reg in pol_regs[j]:
                    pol_loss += 1e-3 * reg  # policy regularization
            # don't want critic to accumulate gradients from policy loss
            disable_gradients(self.critic)
            pol_loss.backward()
            curr_agent.policy.scale_shared_grads()
            enable_gradients(self.critic)

            grad_norm = torch.nn.utils.clip_grad_norm_(curr_agent.policy.parameters(), self.grad_norm_clip/100.)
            curr_agent.policy_optimizer.step()
            curr_agent.policy_optimizer.zero_grad()

            if logger is not None:
                logger.add_scalar('grad_norms/agent%i_policy' % a_i,
                                  grad_norm, self.niter)
                logger.add_scalar('agent%i/losses/pol_loss' % a_i,
                                  pol_loss, self.niter)

    def update_all_targets(self):
        """
        Update all target networks (called after normal updates have been
        performed for each agent)
        """
        if self.hard_update_interval is None:
            soft_update(self.target_critic, self.critic, self.tau)
            for a in self.agents:
                soft_update(a.target_policy, a.policy, self.tau)
        elif self.niter % self.hard_update_interval == 0:
            hard_update(self.target_critic, self.critic)
            for a in self.agents:
                hard_update(a.target_policy, a.policy)

    def prep_training(self, device='cuda'):
        self.critic.train()
        self.target_critic.train()
        self.critic = self.critic.to(device)
        self.target_critic = self.target_critic.to(device)
        for a in self.agents:
            a.policy.train()
            a.target_policy.train()
            a.policy = a.policy.to(device)
            a.target_policy = a.target_policy.to(device)
        self.head_selector.train()
        self.head_selector.to(device)

    def prep_rollouts(self, device='cpu'):
        for a in self.agents:
            a.policy.eval()
            a.policy = a.policy.to(device)
        self.head_selector.eval()
        self.head_selector.to(device)

    def save(self, filename):
        """
        Save trained parameters of all agents into one file
        """
        self.prep_training(device='cpu')  # move parameters to CPU before saving
        save_dict = {'init_dict': self.init_dict,
                     'agent_params': [a.get_params() for a in self.agents],
                     'critic_params': {'critic': self.critic.state_dict(),
                                       'target_critic': self.target_critic.state_dict(),
                                       'critic_optimizer': self.critic_optimizer.state_dict()},
                     'head_selector_params': {'head_selector': self.head_selector.state_dict(),
                                              'head_selector_optimizer': self.head_selector_optimizer.state_dict()}}
        torch.save(save_dict, filename)

    @classmethod
    def init_from_env(cls, env, nagents=2, gamma_e=0.95, gamma_i=0.95,
                      tau=0.01, hard_update_interval=None,
                      pi_lr=0.01, q_lr=0.01, phi_lr=0.1,
                      adam_eps=1e-8,
                      q_decay=1e-3, phi_decay=1e-4, reward_scale=10.,
                      head_reward_scale=25.,
                      pol_hidden_dim=128, critic_hidden_dim=128,
                      nonlin=F.relu, n_intr_rew_types=0, beta=0.5,
                      sep_extr_head=False, **kwargs):
        """
        Instantiate instance of this class from multi-agent environment

        env: Multi-agent Gym environment
        gamma: discount factor
        tau: rate of update for target networks
        lr: learning rate for networks
        hidden_dim: number of hidden dimensions for networks
        """
        if type(env.observation_space) is spaces.Tuple:
            obs_shape = tuple(space.shape for space in
                              env.observation_space.spaces)
        else:
            obs_shape = env.observation_space.shape
        state_shape = env.state_space.shape
        action_size = env.action_space.n

        init_dict = {'nagents': nagents,
                     'gamma_e': gamma_e, 'gamma_i': gamma_i, 'tau': tau,
                     'hard_update_interval': hard_update_interval,
                     'pi_lr': pi_lr, 'q_lr': q_lr, 'phi_lr': phi_lr,
                     'adam_eps': adam_eps,
                     'q_decay': q_decay, 'phi_decay': phi_decay,
                     'reward_scale': reward_scale,
                     'head_reward_scale': head_reward_scale,
                     'pol_hidden_dim': pol_hidden_dim,
                     'critic_hidden_dim': critic_hidden_dim,
                     'nonlin': nonlin,
                     'n_intr_rew_types': n_intr_rew_types,
                     'sep_extr_head': sep_extr_head,
                     'beta': beta,
                     'obs_shape': obs_shape,
                     'state_shape': state_shape,
                     'action_size': action_size}
        instance = cls(**init_dict)
        instance.init_dict = init_dict
        return instance

    @classmethod
    def init_from_save(cls, filename, load_critic=False, load_ir=False):
        """
        Instantiate instance of this class from file created by 'save' method
        """
        save_dict = torch.load(filename)
        instance = cls(**save_dict['init_dict'])
        instance.init_dict = save_dict['init_dict']
        for a, params in zip(instance.agents, save_dict['agent_params']):
            a.load_params(params, load_ir=load_ir)

        head_selector_params = save_dict['head_selector_params']
        instance.head_selector.load_state_dict(head_selector_params['head_selector'])
        instance.head_selector_optimizer.load_state_dict(head_selector_params['head_selector_optimizer'])
        if load_critic:
            critic_params = save_dict['critic_params']
            instance.critic.load_state_dict(critic_params['critic'])
            instance.target_critic.load_state_dict(critic_params['target_critic'])
            instance.critic_optimizer.load_state_dict(critic_params['critic_optimizer'])
        return instance


#!/usr/env/python

## Import General Tools
import sys
import os
import argparse
import logging
from datetime import datetime as dt
from datetime import timedelta as tdelta

import callhorizons

import numpy as np
from astropy import units as u
from astropy.coordinates import SkyCoord
from astropy.table import Table

##-------------------------------------------------------------------------
## Parse Command Line Arguments
##-------------------------------------------------------------------------
## create a parser object for understanding command-line arguments
p = argparse.ArgumentParser(description='''
''')
## add flags
p.add_argument("-v", "--verbose", dest="verbose",
    default=False, action="store_true",
    help="Be verbose! (default = False)")
p.add_argument("-f", "--from", dest="fromdate",
    help="From date for the starlist table (in ISO format)")
p.add_argument("-t", "--to", dest="todate",
    help="To date for the starlist table (in ISO format)")
p.add_argument("-s", "--spacing", dest="spacing",
    default="1h",
    help="The spacing for each starlist entry (e.g. 1h or 15m)")
## add arguments
p.add_argument('name', type=str,
               help="The name of the target compatible with JPL horizons")
args = p.parse_args()

if args.fromdate is None:
    fromdate = dt.utcnow()
else:
    try:
        fromdate = dt.strptime(args.fromdate, '%Y-%m-%dT%H:%M:%S')
    except:
        print('Could not parse from date')
        raise

if args.todate is None:
    todate = dt.utcnow() + tdelta(1, 0)
else:
    try:
        todate = dt.strptime(args.todate, '%Y-%m-%dT%H:%M:%S')
    except:
        print('Could not parse to date')
        raise

##-------------------------------------------------------------------------
## Create logger object
##-------------------------------------------------------------------------
log = logging.getLogger('MyLogger')
log.setLevel(logging.DEBUG)
## Set up console output
LogConsoleHandler = logging.StreamHandler()
if args.verbose is True:
    LogConsoleHandler.setLevel(logging.DEBUG)
else:
    LogConsoleHandler.setLevel(logging.INFO)
LogFormat = logging.Formatter('%(asctime)s %(levelname)8s: %(message)s',
                              datefmt='%Y-%m-%d %H:%M:%S')
LogConsoleHandler.setFormatter(LogFormat)
log.addHandler(LogConsoleHandler)
## Set up file output
# LogFileName = None
# LogFileHandler = logging.FileHandler(LogFileName)
# LogFileHandler.setLevel(logging.DEBUG)
# LogFileHandler.setFormatter(LogFormat)
# log.addHandler(LogFileHandler)

##-------------------------------------------------------------------------
## Main Program
##-------------------------------------------------------------------------
def main(name, fromdate, todate, obscode=568, spacing='1h'):
    log.debug(f'Querying horizons for: "{name}"')
    target = callhorizons.query(name)
    fromstr = fromdate.strftime('%Y-%m-%d %H:%M')
    tostr = todate.strftime('%Y-%m-%d %H:%M')
    log.debug(f'  From: "{fromstr}"')
    log.debug(f'  To:   "{tostr}"')
    log.debug(f'  Increment: "{spacing}"')
    target.set_epochrange(fromstr, tostr, spacing)
    try:
        target.get_ephemerides(obscode)
    except ValueError as e:
        log.error("Failed to get ephemerides")
        print(e.args[0])
    tab = Table(target.data)
    if len(tab) == 0:
        log.error("ephemerides have zero length")

    name = name.replace("/", "")
    name = name.replace(" ", "_")

    wasup = True
    for entry in tab:
        time = (entry['datetime'].split(' ')[1]).replace(':', '')
        starlistname = f'{name[0:9]:s}_{time:s}'
        if np.isnan(float(entry['airmass'])) == True:
            if wasup == True:
                print(f"# Target {name} is down at {time}")
                wasup = False
        else:
            line = [f'{starlistname:15s}']
            coord = SkyCoord(entry['RA'], entry['DEC'], frame='fk5', unit=(u.deg, u.deg))
            line.append(f'{coord.to_string("hmsdms", sep=" ", precision=2)}')
            line.append(f'{coord.equinox.jyear:.2f}')
            line.append(f'dra={float(entry["RA_rate"])/15*3600:.3f}')
            line.append(f'ddec={float(entry["DEC_rate"])*3600:.3f}')
            line.append(f'vmag={float(entry["V"]):.2f}')
            line.append(f'# airmass={float(entry["airmass"]):.2f}')
            print(' '.join(line))

if __name__ == '__main__':
    main(args.name, fromdate, todate, spacing=args.spacing)


import svg.path as sp
import numpy as np
import json
import time
import subprocess

# Copy the "d" attribute of "path" SVG elements here
# A quick-and-dirty way to extract this from an image:
#  1. open the svg in chrome
#  2. F12 to open console
#  3. run this command: $$('path[d]').map(function(e){return e.getAttribute("d")})
#  4. right-click the output -> copy object
#  5. paste it here

# Data below is: https://www.coffeestainstudios.com/assets/images/coffeestain-logo.svg
path_text = [
    "M92.4,345.7c-6.6,0-10.9,5.7-12.2,11.8l0.6,0.2c1.2-1.2,2.8-1.8,4.4-1.7c2.6,0.2,3.6,2.8,3.5,5.4  c-0.3,6.4-8,16.7-17.4,16.5c-12.9-0.2-4.4-19-1.6-24.9c2.6-5.5,8.1-12.5,8.1-18.9c0-4.7-3.2-8.6-8.1-8.5c-8.4,0.2-16.4,7-21.7,12.9  c0.5-4.7-3.4-8.5-7.8-8.5c-4.4-0.1-8.7,2.8-10.2,6.9c1.9,1.1,2.1,4.5,1.4,8.8c-2.6,7.4-10.1,23.8-18.3,24.2c-8.4,0.4,6.7-28,9.3-31  c-1-2.7-4.5-4.3-7.8-4.4c-2,0-4,0.8-5.4,1.8c-5.4,3.8-11.4,18.9-13.2,26c-2.9,6.7-6.3,14.1-10.3,14.2c-1.5,0.1-2.4-1.5-1.5-5.5  c0.8-4,3.7-9.3,3.6-12.6c-0.1-3.1-2.9-4.7-6-4.7c-5.5,0.1-8.4,7.4-10.8,11.4c-3.2,5.5-8.3,12.4-12.4,12c-9.5-0.9,9.5-33.8,21.2-34.1  c3.7-0.1,2.9,1.5,2.4,5.4c-0.6,3.9,10,0.9,9.9-6.8c0-6.3-6.5-9.5-12.1-9.3c-17.8,0.6-34.7,23.4-40.4,42.1c-2.9,3.8-7.4,7.8-10.8,7.9  c-5.1,0.1-2.6-8.2-1.6-10.7c4.2-10.4,11.4-24.7,19.1-37.1c28.6-7,61.1-9.3,89.9-12.1c16-1.6,38.2-3.1,48.4-17.6  c4.3-6.2,4.9-11.9,4.2-19.1c3.9-3.6,1.8-9.1-2.4-11.5c-9.7-5.5-24.4-2.2-32.7,4c-2.6,1.9-2.1,6-0.8,6.2c4.7-1.3,21.4-5.3,23.8,1.3  c0.9,2.5-0.9,5.8-3.7,8.4c-8,7.4-26.6,12.1-39.8,14.6c-14.4,2.7-50,7.4-78.1,12.9c1.5-1.9,2.9-3.7,4.3-5.3c-0.3-3.6-2.4-6.5-6.1-6.4  c-4,0-6.1,3-8.1,6c-1.7,2.5-3.8,5.7-6.3,9.3c-9.3,2.3-16.5,4.7-20,7.1c-2.7,1.9-4.6,4.9-5.4,8.3c-0.5,2.2,0.3,6.1,2.9,6.2  c3.3-2.8,7.8-5.1,12.8-7.1C-82.7,358-95.7,380.2-96.1,392c-0.2,6.3,2.8,11.7,9.7,11.5c8.5-0.2,18.3-9.5,24.5-17.3  c1.2,5,4.6,8.1,10.4,7.9c10.2-0.3,18-9.6,22.9-17.6c0.5,6.6,2.9,12.3,10,12.3c5.7,0,10-4,13.1-9.2c0.9,1.8,3.2,5.3,8.6,5.4  c6.4,0.2,14.6-6.4,18.5-11.4c0.1,3.2,3.6,6.4,6.9,6.5c4.8,0.1,7.1-3.8,8.7-7.6c2.5-5.9,14.6-32.7,21-33.1c4-0.3,1.3,5.7,0.6,7.7  c-3.3,8.8-10.6,17.4-11.4,29c-0.8,11.1,6.8,17,17.3,17.1c17.3,0.1,37.3-18.9,37.3-36.5C101.9,350.4,98.8,345.7,92.4,345.7z",
    "M120.1,235c-4,5.7-13.2,13.5-20.7,13.9c-10.6,0.5-9.8-11.4-4-17.3c1.1,3,4.2,3.5,7,3.4  c11.9-0.6,27.6-12.9,28.1-25.5c0.2-6.7-4-12-12.2-11.8c-16.5,0.3-32,18.1-39,32.6c-6.2,9.7-15.8,19.2-26.7,19.3  c-7.8,0-9.3-7.8-4.8-13.2c0.7,2.9,3.4,3.8,6.1,3.8c12,0,27.6-13.5,28-25.6c0.2-6-4-10.1-10.6-10.4c-18.3-0.7-40.4,26.7-41.4,44.7  c-0.5,9.2,4.3,15.7,13.8,16.3c9.7,0.5,21.8-7.9,28.9-15.7c0.6,9.8,7.5,14.4,16.5,14.5c12.9,0.1,28.3-9.7,34-24.2  C123.8,237.2,122.8,234.7,120.1,235z M117.7,209c8.8,0.4-4.9,15.7-13.7,15.8C95.1,224.9,111.6,208.8,117.7,209z M69.4,214.9  c1.7-0.1,2.5,0.9,2.5,2.4c-0.1,7.7-17.4,15.8-17.6,11.1C54.1,224.9,64.9,215,69.4,214.9z",
    "M55.3,398.5c-2.1-1.1-8.1-2.7-12.9-2.5c-6.8,0.2-8.2,6-3.2,9.8c7-0.6,13.3,0.3,20,1.3  C60,402.7,58.1,400.1,55.3,398.5z",
    "M43.5,277.3c-0.2-2.6-3.9-4.2-6.3-4.2c-9.2,0.2-16.1,10.2-15.9,19c0.2,5.4,3.5,10,10,10.3  c1.5-0.5,1.4-1.9,1.4-1.9c0.7-6.6,5.5-14.6,9.3-19.8C43.7,278.2,43.5,277.3,43.5,277.3z",
    "M-65,300.4c0.2-5.8-2.5-11-10.1-11.2c-18-0.3-26.4,31.6-30.2,45.6c-12.4,7.8-26.4,12.7-38.8,12.4  c-5.3-0.1-10.3-1.4-12.7-4.7c-2-2.8-2.8-6.8,0.1-9.3c-0.5-1-2.2-1.3-3.6-0.9c-13.1,3.6-14.1,35.1,14.2,34.5  c9.8-0.2,23.1-5.9,36.1-14.3c-10,30.4-31.9,71-61.3,71.2c-12.5,0.1-20.8-8.5-20.9-19.6c-0.2-17.1,19.5-33.9,34.5-21.4l0.9-1  c-0.2-10.9-12.9-14.8-21.9-14.8c-23.8,0.2-41.3,22-41.3,44.6c0,24.1,18.5,39,42.1,38.5c51.1-1.1,72.5-65.9,85.2-110.4  C-81.3,329.4-65.5,312.6-65,300.4z M-74.8,298c3.2,0.3-0.1,6-0.9,7.5c-2.9,5.1-6.8,10.2-11.5,14.9C-84.5,311.9-79.5,297.6-74.8,298z  ",
    "M23.5,396.3c-40.4-1.9-111.1,12-143.8,34.3c-5,3.9-7.8,8.5-7.4,14.9c0.3,4.6,3.5,9.3,7.9,10.5  c35-29.3,92.3-48.1,149.6-50C30.1,400.7,29.3,397,23.5,396.3z",
    "M-150.5,324.7c22.5-0.4,49.4-27.4,60.8-49.5c5.7,8.9,18.9,9.1,29-0.4c7.2-6.7,13.4-16.2,17.3-24.2  c7.5-3.1,12.9-8.3,12.9-8.2c-9.2,11.6-26.7,38.1-26.7,52.8c0,6,2.5,9.6,8.8,9.4c21.8-0.4,37.6-44.6,34.7-62.6c0,0,10.5-2.9,20.8-8.1  c-10.9,16.2-29.8,48.6-29.6,62.6c0.3,21.2,38.4,10.5,41.2-42c0.2-3.4,0.2-6.5,0-9.2c2.6-4.7,10.1-6.8,10.1-6.8  c2.6-7.2,6.1-11.3,6.1-11.3c-3.4-0.3-8.3,1.4-11.4,2.7c0.6-1.5,1.2-2.9,1.9-4.3c18.4-10.8,45.8-38.8,46.2-57  c0.1-5.4-2.2-10.5-8.9-10.6c-15.8-0.2-40.8,47.3-47.8,60.8c-6.5,3.8-16.1,8.3-25,8.9c19-16.8,43.6-44.2,43.6-60.7  c0-5-2.2-9.3-8.6-9.2c-18.6,0.2-40.1,48-48.9,69.4c-2,2.3-7.7,8-15.1,10.6c1.7-9.4,0.9-21.5-11.1-21.2c-6.6,0.2-11.3,5.4-14.1,11  c-14.9,1.9-27.3,20.9-28.3,34.6c-11.2,15.7-30.6,38.2-49.2,37.9c-40.8-0.8,19.9-121,58.4-122c20.7-0.5-4.9,51.8-14.8,52.6  c-1.7-0.1-2.4-1.8-2.5-4.3c-0.2-8.5,5.9-18.6,11.5-24.4c0.4-0.4-0.8-1-0.8-1c-16.2-0.9-33.6,24-34.1,39.9  c-0.2,6.1,2.1,13.9,9.6,13.8c21.9-0.2,53.6-49.7,53.5-70.4c-0.1-12.1-7-19.5-19.8-19.4c-46.3,0.6-97.3,82.2-97.9,124.8  C-178.9,308-171.4,325.1-150.5,324.7z M61.1,167.6c4.5-0.1,0.4,7.4-4.3,14.2c-6.9,10-14,18.7-23.5,25.8  C36.9,199.9,56.1,167.7,61.1,167.6z M-3.2,289c-5.2-0.9,7.2-26.3,15.8-41.9C12.8,255.8,3.2,290.2-3.2,289z M23.4,169  C29,170,7.7,198.6-2,208.7C2.7,198.3,17.5,167.9,23.4,169z M-19.7,247.4c0,0.1,0,0.1,0,0.2c-0.8,8.3-13.1,37.6-18.9,37.9  C-42.8,285.7-27,254.1-19.7,247.4z M-48.1,227.4c3.3,0,2.8,5.7,0.3,11.5c-2.4-0.2-4.5-1.2-4.7-4.3C-52.6,232.1-51,227.5-48.1,227.4z   M-67,241.1c0.7,6.7,4.6,11.1,11.5,11.8c-3.8,5.5-8.8,11-14.3,11.1C-79.1,264-76.3,248.7-67,241.1z"
]
paths = [sp.parse_path(t) for t in path_text]

template = """
                {{
                  "properties": [
                    {{
                      "name": "Location",
                      "type": "StructProperty",
                      "index": 0,
                      "value": {{
                        "type": "Vector",
                        "x": {lx},
                        "y": {ly},
                        "z": {lz}
                      }}
                    }},
                    {{
                      "name": "ArriveTangent",
                      "type": "StructProperty",
                      "index": 0,
                      "value": {{
                        "type": "Vector",
                        "x": {atx},
                        "y": {aty},
                        "z": {atz}
                      }}
                    }},
                    {{
                      "name": "LeaveTangent",
                      "type": "StructProperty",
                      "index": 0,
                      "value": {{
                        "type": "Vector",
                        "x": {ltx},
                        "y": {lty},
                        "z": {ltz}
                      }}
                    }}
                  ]
                }}

"""

class Placer:
    template_path = "pipe_template.json"
    id_names = ["circle_top"]
    def __init__(self, id_offset, x, y, z, color=1):
        with open(self.template_path, "r") as t:
            template_s = t.read()
            ids = {}
            buffer_0_id = None
            for i, id_name in enumerate(self.id_names):
                new_id = id_offset + i
                template_s = template_s.replace("{{"+id_name+"}}", str(new_id))
                if id_name == self.id_names[0]:
                    buffer_0_id = new_id

        self.template_json = json.loads(template_s)
        self.actors = []
        self.components = []
        for actor in self.template_json:
            if actor["type"] == 1:
                actor["transform"]["translation"][0] = x
                actor["transform"]["translation"][1] = y
                actor["transform"]["translation"][2] = z
                # Bind the pipe to a color-gun slot
                actor["entity"]["properties"][3]["value"]["value"] = color
                self.actors.append(actor)
            else:
                self.components.append(actor)

    def write_shape_points(self, points, arrive_tangents, leave_tangents):
        for actor in self.actors:
            actor["entity"]["properties"][0]["value"]["values"].clear()
            for c,at,lt in zip(points, arrive_tangents, leave_tangents):
                s = json.loads(template.format(lx=c[0], ly=c[1], lz=c[2], atx=at[0], aty=at[1], atz=at[2], ltx=lt[0], lty=lt[1], ltz=lt[2]))
                actor["entity"]["properties"][0]["value"]["values"].append(s)

    def write(self, save_json, debug=False):
        if debug:
            print(json.dumps(self.actors, indent=" "))
        save_json["actors"].extend(self.actors)
        save_json["components"].extend(self.components)

def tangents(points):
    diffs = (points[1:] - points[:-1]) / 2
    leave_tangents = np.append(diffs, np.array([0.,0.,0.]))
    arrive_tangents = np.insert(diffs, 0, np.array([0.,0.,0.]))
    return arrive_tangents.reshape((len(points), 3)), leave_tangents.reshape((len(points), 3))

def find_item_id_offset(save_json):
    max_id = None
    for actor in save_json["actors"]:
        if actor["className"] == "/Game/FactoryGame/Buildable/Factory/Pipeline/Build_Pipeline.Build_Pipeline_C":
            actor_id = int(actor["pathName"].split("_")[-1])
            if not max_id or actor_id > max_id:
                max_id = actor_id
    return max_id

# your source save file
# convert to json with https://ficsit-felix.netlify.app/#/ or
# https://github.com/ficsit-felix/satisfactory-json
json_save_path = "debug.json"
with open(json_save_path, 'rb') as f:
    save_json = json.load(f)

id_offset = find_item_id_offset(save_json) + 1

# where in the game to spawn
# all pipes have this location, SVG path gives the offset relative to this point
# NOTE: I haven't figured out rotation, the resulting pipes will be oriented
#  due East (you'll be looking at it facing due West)
translation = np.array([-186378.09375, -82377.40625, 24282.08203125])
# scale the image up or down, play with this till it looks good
script_scale = 50.
# how many points per line segment
# higher numbers give smoother lines, but too high and you'll get wrinkles
steps_per_path = 100.
# x,y offset in the SVG coordinate space
# if the SVG paths have a transformation in their own coordinate system you can plug that in here
point_translate = (0.,0.)

for path_hm in paths:
    # paths are composed of segments, spawn one pipe per segment
    # NOTE: you can use the same code to instead spawn one pipe per path, ymmv
    for path in path_hm:
        domain_points = np.arange(0.,1.,1./steps_per_path)
        path_coords = np.array([(0., -path.point(dp).real - point_translate[0], -path.point(dp).imag - point_translate[1]) for dp in domain_points])
        path_coords = path_coords * script_scale
        arrive_tangents, leave_tangents = tangents(path_coords)
        placer = Placer(id_offset, translation[0], translation[1], translation[2])
        placer.write_shape_points(path_coords, arrive_tangents, leave_tangents)
        placer.write(save_json)
        id_offset += 1

# output json save file
json_save_out_path = "debug-curve.json"
with open(json_save_out_path, 'w') as f:
    json.dump(obj=save_json, fp=f)

# convert back to a save with
# https://ficsit-felix.netlify.app/#/ or
# https://github.com/ficsit-felix/satisfactory-json
"""
print(subprocess.run(args=[
    "node", "lib\\cli\\json2sav.js", "debug-curve.json", "debug-curve.sav"],
               check=False, capture_output=True))
 """


#!/usr/bin/env python
# coding=utf-8
"""
Function:
This class is used to generate GUP dump data of the TensorFlow model.
Copyright Information:
Huawei Technologies Co., Ltd. All Rights Reserved © 2021
"""
import argparse
import sys
import os
import numpy as np
import tensorflow as tf
from tensorflow.python import debug as tf_debug
from common import utils, tf_common
from common.utils import AccuracyCompareException

class TfDebugRunner(object):
    """
    This class is used to generate GUP dump data of the TensorFlow model.
    """

    def __init__(self, arguments):
        self.args = arguments
        self.global_graph = None
        self.input_shapes = utils.parse_input_shape(self.args.input_shape)
        self.input_path = self.args.input_path
        self.dump_root = os.path.realpath(self.args.out_path)

    def _load_graph(self):
        try:
            with tf.io.gfile.GFile(self.args.model_path, 'rb') as f:
                global_graph_def = tf.compat.v1.GraphDef.FromString(f.read())
            self.global_graph = tf.Graph()
            with self.global_graph.as_default():
                tf.import_graph_def(global_graph_def, name='')
        except Exception as err:
            utils.print_error_log("Failed to load the model %s. %s" % (self.args.model_path, err))
            raise AccuracyCompareException(utils.ACCURACY_COMPARISON_OPEN_FILE_ERROR)
        utils.print_info_log("Load the model %s successfully." % self.args.model_path)

    def _get_outputs_tensor(self):
        outputs_tensor = []
        for tensor_name in self.args.output_nodes.split(';'):
            tensor = self.global_graph.get_tensor_by_name(tensor_name)
            outputs_tensor.append(tensor)
        return outputs_tensor

    def _get_inputs_data(self, inputs_tensor):
        inputs_map = {}
        input_path = self.args.input_path.split(",")
        for index, tensor in enumerate(inputs_tensor):
            try:
                input_data = np.fromfile(input_path[index], tf_common.convert_to_numpy_type(tensor.dtype))
                if tensor.shape:
                    input_data = input_data.reshape(tensor.shape)
                inputs_map[tensor] = input_data
                utils.print_info_log("load file name: {}, shape: {}, dtype: {}".format(
                    os.path.basename(input_path[index]), input_data.shape, input_data.dtype))
            except Exception as err:
                utils.print_error_log("Failed to load data %s. %s" % (input_path[index], err))
                raise AccuracyCompareException(utils.ACCURACY_COMPARISON_BIN_FILE_ERROR)
        return inputs_map

    def _run_model(self, inputs_map, outputs_tensor):
        config = tf.compat.v1.ConfigProto(log_device_placement=False, allow_soft_placement=True)
        with tf.compat.v1.Session(graph=self.global_graph, config=config) as sess:
            sess = tf_debug.LocalCLIDebugWrapperSession(sess, ui_type="readline", dump_root=self.dump_root)
            return sess.run(outputs_tensor, feed_dict=inputs_map)

    def run(self):
        """
        Function description:
            run TensorFlow model
        """
        self._load_graph()
        inputs_tensor = tf_common.get_inputs_tensor(self.global_graph, self.args.input_shape)
        inputs_map = self._get_inputs_data(inputs_tensor)
        outputs_tensor = self._get_outputs_tensor()
        self._run_model(inputs_map, outputs_tensor)

def _make_dump_data_parser(parser):
    parser.add_argument("-m", "--model-path", dest="model_path", default="",
                        help="<Required> The original model (.pb) file path", required=True)
    parser.add_argument("-i", "--input-path", dest="input_path", default="",
                        help="<Required> The input data path of the model. Separate multiple inputs with commas(,)."
                             " E.g: input_0.bin,input_1.bin", required=True)
    parser.add_argument("-o", "--out-path", dest="out_path", default="", help="<Required> The output path")
    parser.add_argument("-s", "--input-shape", dest="input_shape", default="",
                        help="<Required> Shape of input shape. Separate multiple nodes with semicolons(;)."
                             " E.g: input_name1:1,224,224,3;input_name2:3,300")
    parser.add_argument("--output-nodes", dest="output_nodes", default="", required=True,
                        help="<Required> Output nodes designated by user. Separate multiple nodes with semicolons(;)."
                             " E.g: node_name1:0;node_name2:1;node_name3:0")

def main():
    """
   Function Description:
       main process function
   Exception Description:
       exit the program when an AccuracyCompare Exception  occurs
   """
    parser = argparse.ArgumentParser()
    _make_dump_data_parser(parser)
    args = parser.parse_args(sys.argv[1:])
    TfDebugRunner(args).run()

if __name__ == '__main__':
    main()


import os
import numpy as np
import tensorflow as tf
import source.utils as utils
import whiteboxlayer.layers as wbl
import whiteboxlayer.extensions.utility as wblu

class Agent(object):

    def __init__(self, **kwargs):

        print("\nInitializing Neural Network...")

        self.dim_h = kwargs['dim_h']
        self.dim_w = kwargs['dim_w']
        self.dim_c = kwargs['dim_c']
        self.num_class = kwargs['num_class']
        self.ksize = kwargs['ksize']
        self.learning_rate = kwargs['learning_rate']
        self.path_ckpt = kwargs['path_ckpt']

        self.variables = {}

        self.__model = Neuralnet(\
            who_am_i="RepVGG", **kwargs, \
            filters=[1, 32, 64, 128])

        dummy = tf.zeros((1, self.dim_h, self.dim_w, self.dim_c), dtype=tf.float32)
        self.__model.forward(x=dummy, training=True, verbose=True)
        self.__model.forward(x=dummy, training=False, verbose=True)

        self.__init_propagation(path=self.path_ckpt)

    def __init_propagation(self, path):

        self.summary_writer = tf.summary.create_file_writer(self.path_ckpt)

        self.variables['trainable'] = []
        ftxt = open("list_parameters.txt", "w")
        for key in list(self.__model.layer.parameters.keys()):
            trainable = self.__model.layer.parameters[key].trainable
            text = "T: " + str(key) + str(self.__model.layer.parameters[key].shape)
            if(trainable):
                self.variables['trainable'].append(self.__model.layer.parameters[key])
            ftxt.write("%s\n" %(text))
        ftxt.close()

        self.optimizer = tf.optimizers.Adam(learning_rate=self.learning_rate)
        self.save_params()

    def __loss(self, y, y_hat):

        entropy_b = tf.nn.softmax_cross_entropy_with_logits(labels=y, logits=y_hat)
        entropy = tf.math.reduce_mean(entropy_b)

        return {'entropy_b': entropy_b, 'entropy': entropy}

    @tf.autograph.experimental.do_not_convert
    def step(self, minibatch, iteration=0, training=False):

        x, y = minibatch['x'], minibatch['y']

        with tf.GradientTape() as tape:
            logit, y_hat = self.__model.forward(x=x, training=training, verbose=False)
            losses = self.__loss(y=y, y_hat=logit)

        if(training):
            gradients = tape.gradient(losses['entropy'], self.variables['trainable'])
            self.optimizer.apply_gradients(zip(gradients, self.variables['trainable']))

            with self.summary_writer.as_default():
                tf.summary.scalar('%s/entropy' %(self.__model.who_am_i), losses['entropy'], step=iteration)

        return {'y_hat':y_hat, 'losses':losses}

    def save_params(self, model='base', tflite=False):

        if(tflite):
            # https://github.com/tensorflow/tensorflow/issues/42818
            conc_func = self.__model.__call__.get_concrete_function(\
                tf.TensorSpec(shape=(1, self.dim_h, self.dim_w, self.dim_c), dtype=tf.float32))
            converter = tf.lite.TFLiteConverter.from_concrete_functions([conc_func])

            converter.optimizations = [tf.lite.Optimize.DEFAULT]
            converter.experimental_new_converter = True
            converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]

            tflite_model = converter.convert()

            with open('model.tflite', 'wb') as f:
                f.write(tflite_model)
        else:
            vars_to_save = self.__model.layer.parameters.copy()
            vars_to_save["optimizer"] = self.optimizer

            ckpt = tf.train.Checkpoint(**vars_to_save)
            ckptman = tf.train.CheckpointManager(ckpt, directory=os.path.join(self.path_ckpt, model), max_to_keep=1)
            ckptman.save()

    def load_params(self, model):

        vars_to_load = self.__model.layer.parameters.copy()
        vars_to_load["optimizer"] = self.optimizer

        ckpt = tf.train.Checkpoint(**vars_to_load)
        latest_ckpt = tf.train.latest_checkpoint(os.path.join(self.path_ckpt, model))
        status = ckpt.restore(latest_ckpt)
        status.expect_partial()

class Neuralnet(tf.Module):

    def __init__(self, **kwargs):
        super(Neuralnet, self).__init__()

        self.who_am_i = kwargs['who_am_i']
        self.dim_h = kwargs['dim_h']
        self.dim_w = kwargs['dim_w']
        self.dim_c = kwargs['dim_c']
        self.ksize = kwargs['ksize']
        self.num_class = kwargs['num_class']
        self.filters = kwargs['filters']

        self.layer = wbl.Layers()

        self.forward = tf.function(self.__call__)

    @tf.function
    def __call__(self, x, training=False, verbose=False):

        logit = self.__nn(x=x, training=training, name=self.who_am_i, verbose=verbose)
        y_hat = tf.nn.softmax(logit, name="y_hat")

        return logit, y_hat

    def __nn(self, x, training, name='neuralnet', verbose=True):

        att = None
        for idx, _ in enumerate(self.filters[:-1]):
            if(idx == 0): continue
            x = self.__repvgg(x=x, ksize=self.ksize, filter_in=self.filters[idx-1], filter_out=self.filters[idx], stride=2, group=1, \
                activation='relu', name='%s-%d_repvgg1' %(name, idx), training=training, verbose=verbose)
            x = self.__repvgg(x=x, ksize=self.ksize, filter_in=self.filters[idx], filter_out=self.filters[idx], stride=1, group=2, \
                activation='relu', name='%s-%d_repvgg2' %(name, idx), training=training, verbose=verbose)

        x = tf.math.reduce_mean(x, axis=(1, 2))
        x = self.layer.fully_connected(x=x, c_out=self.num_class, \
                activation=None, name="%s-clf" %(name), verbose=verbose)

        return x

    def __repvgg(self, x, ksize, filter_in, filter_out, stride=1, group=1, \
        activation='relu', name="", training=False, verbose=False):

        if(training):
            x_main = self.layer.conv2d(x=x, stride=stride, \
                filter_size=[ksize, ksize, filter_in, filter_out], \
                batch_norm=True, activation='relu', name='%s_main' %(name), verbose=verbose)
            x_sub = self.layer.conv2d(x=x, stride=stride, \
                filter_size=[1, 1, filter_in, filter_out], \
                batch_norm=True, activation='relu', name='%s_sub' %(name), verbose=verbose)

            if(stride == 1):
                x_id = self.layer.batch_normalization(x=x, trainable=True, name='%s_id_bn' %(name), verbose=verbose)
                x = x_main + x_sub + x_id
            else:
                x = x_main + x_sub
        else:
            x = self.__fuse_to_deploy(x=x, stride=stride, group=group, name=name)

        return self.layer.activation(x, activation=activation, name='%s_act' %(name))

    def __fuse_to_deploy(self, x, stride=1, group=1, \
        dilations=[1, 1, 1, 1], padding='SAME', name=""):

        w_main = self.layer.parameters['%s_main_w' %(name)]
        b_main = self.layer.parameters['%s_main_b' %(name)]
        w_sub = self.layer.parameters['%s_sub_w' %(name)]
        b_sub = self.layer.parameters['%s_sub_b' %(name)]
        pad = int((self.ksize-1)/2)
        w_sub = tf.pad(w_sub, paddings=[[pad, pad], [pad, pad], [0, 0], [0, 0]], mode='CONSTANT', constant_values=0)

        w = w_main + w_sub
        b = b_main + b_sub

        if(stride == 1):
            list_shape = w_main.get_shape().as_list()
            w_id = np.zeros(list_shape)
            idx_k, num_c = int(list_shape[0]/2), list_shape[-2]
            for idx_c in range(num_c):
                w_id[idx_k, idx_k, idx_c, ::group] = 1
            w_id = tf.convert_to_tensor(w_id, dtype=tf.float32)

            w = w + w_id

        wx = tf.nn.conv2d(
            input=x,
            filters=w,
            strides=[1, stride, stride, 1],
            padding=padding,
            data_format='NHWC',
            dilations=dilations,
            name='%s_deploy_conv' %(name)
        )

        y = tf.math.add(wx, b, name='%s_deploy_add' %(name))

        mean_main, variance_main, offset_main, scale_main = self.__bn_extraction(x=y, name='%s_main_bn' %(name))
        mean_sub, variance_sub, offset_sub, scale_sub = self.__bn_extraction(x=y, name='%s_sub_bn' %(name))
        try:
            mean_id, variance_id, offset_id, scale_id = self.__bn_extraction(x=y, name='%s_id_bn' %(name))
        except:
            mean = mean_main + mean_sub
            variance = variance_main + variance_sub
            offset = offset_main + offset_sub
            scale = scale_main + scale_sub
        else:
            mean = mean_main + mean_sub + mean_id
            variance = variance_main + variance_sub + variance_id
            offset = offset_main + offset_sub + offset_id
            scale = scale_main + scale_sub + scale_id

        return (scale * (y - mean)) / variance + offset

    def __bn_extraction(self, x, name=""):

        mean, variance = tf.nn.moments(x=x, axes=[0], keepdims=True, name="%s_mmt" %(name))
        ofs, sce = self.layer.parameters['%s_ofs' %(name)], self.layer.parameters['%s_sce' %(name)]

        return mean, variance, ofs, sce


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

import numpy as np
import os
import argparse
from fastspeech2 import FastSpeech2
from loss import FastSpeech2Loss
from dataset import Dataset
import hparams as hp
import utils

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

def get_FastSpeech2(num):
    checkpoint_path = os.path.join(hp.checkpoint_path, "checkpoint_{}.pth.tar".format(num))
    model = nn.DataParallel(FastSpeech2())
    model.load_state_dict(torch.load(checkpoint_path)['model'])
    model.requires_grad = False
    model.eval()
    return model

def evaluate(model, step):
    model.eval()
    torch.manual_seed(0)

    eval_path = hp.eval_path
    if not os.path.exists(eval_path):
        os.makedirs(eval_path)

    # Get dataset
    dataset = Dataset("val.txt", sort=False)
    loader = DataLoader(dataset, batch_size=hp.batch_size, shuffle=False, collate_fn=dataset.collate_fn, drop_last=False, num_workers=0, )

    # Get loss function
    Loss = FastSpeech2Loss().to(device)

    # Evaluation
    d_l = []
    f_l = []
    e_l = []
    mel_l = []
    mel_p_l = []
    current_step = 0
    idx = 0
    for i, data_of_batch in enumerate(loader):
        # Get Data
        id_ = data_of_batch["id"]
        text = torch.from_numpy(data_of_batch["text"]).long().to(device)
        mel_target = torch.from_numpy(data_of_batch["mel_target"]).float().to(device)
        D = torch.from_numpy(data_of_batch["D"]).int().to(device)
        f0 = torch.from_numpy(data_of_batch["f0"]).float().to(device)
        energy = torch.from_numpy(data_of_batch["energy"]).float().to(device)
        src_len = torch.from_numpy(data_of_batch["src_len"]).long().to(device)
        mel_len = torch.from_numpy(data_of_batch["mel_len"]).long().to(device)
        max_src_len = np.int(np.max(data_of_batch["src_len"]))
        max_mel_len = np.int(np.max(data_of_batch["mel_len"]))
        emotion = torch.from_numpy(np.array(data_of_batch["emotion"])).float().cuda()

        with torch.no_grad():
            # Forward
            mel_output, mel_postnet_output, duration_output, f0_output, energy_output, src_mask, mel_mask, out_mel_len = model(
                    text, src_len, emotion, mel_len, D, f0, energy, max_src_len, max_mel_len)

            # Cal Loss
            mel_loss, mel_postnet_loss, d_loss, f_loss, e_loss = Loss(
                    duration_output, D, f0_output, f0, energy_output, energy, mel_output, mel_postnet_output, mel_target, ~src_mask, ~mel_mask)

            d_l.append(d_loss.item())
            f_l.append(f_loss.item())
            e_l.append(e_loss.item())
            mel_l.append(mel_loss.item())
            mel_p_l.append(mel_postnet_loss.item())

            if idx == 0:
                # Run vocoding and plotting spectrogram only when the vocoder is defined
                for k in range(1):
                    basename = id_[k]
                    gt_length = mel_len[k]
                    out_length = out_mel_len[k]
                    mel_target_ = mel_target[k, :gt_length]
                    mel_output_ = mel_output[k, :gt_length]
                    mel_postnet = mel_postnet_output[k, :out_length]

                    mel_target_ = mel_target_.transpose(0, 1).detach()
                    mel_output_ = mel_output_.transpose(0, 1).detach()
                    mel_postnet = mel_postnet.transpose(0, 1).detach()

                    np.save(os.path.join(hp.eval_path, 'eval_step_{}_{}_mel_postnet.npy'.format(step, basename)), mel_postnet.cpu().numpy())
                    np.save(os.path.join(hp.eval_path, 'eval_step_{}_{}_mel_gt.npy'.format(step, basename)), mel_target_.cpu().numpy())
                    np.save(os.path.join(hp.eval_path, 'eval_step_{}_{}_mel.npy'.format(step, basename)), mel_output_.cpu().numpy())

                    f0_ = f0[k, :gt_length]
                    f0_output_ = f0_output[k, :out_length]

                    energy_ = energy[k, :gt_length]
                    energy_output_ = energy_output[k, :out_length]

                    gt_duration = D[k, :gt_length].detach().cpu().numpy()
                    utils.plot_data2([(mel_postnet.cpu().numpy(), f0_output_, energy_output_, gt_duration), (mel_target_.cpu().numpy(), f0_, energy_, gt_duration)],
                        ['Synthesized Spectrogram', 'Ground-Truth Spectrogram'], filename=os.path.join(hp.eval_path, 'eval_step_{}_{}.png'.format(step, basename)))
                    idx += 1
                print("done")
        current_step += 1

    d_l = sum(d_l) / len(d_l)
    f_l = sum(f_l) / len(f_l)
    e_l = sum(e_l) / len(e_l)

    mel_l = sum(mel_l) / len(mel_l)
    mel_p_l = sum(mel_p_l) / len(mel_p_l)

    model.train()

    return d_l, f_l, e_l, mel_l, mel_p_l

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('--step', type=int, default=30000)
    args = parser.parse_args()

    # Get model
    model = get_FastSpeech2(args.step).to(device)
    print("Model Has Been Defined")
    num_param = utils.get_param_num(model)
    print('Number of FastSpeech2 Parameters:', num_param)

    # Init directories
    if not os.path.exists(hp.log_path):
        os.makedirs(hp.log_path)
    if not os.path.exists(hp.eval_path):
        os.makedirs(hp.eval_path)
    evaluate(model, args.step)


# Copyright Howto100M authors.
# Copyright (c) Facebook, Inc. All Rights Reserved

import torch as th
import torch.nn.functional as F
import math
import numpy as np
import argparse

from torch.utils.data import DataLoader
from model import get_model
from preprocessing import Preprocessing
from random_sequence_shuffler import RandomSequenceSampler

from tqdm import tqdm
from pathbuilder import PathBuilder
from videoreader import VideoLoader

parser = argparse.ArgumentParser(description='Easy video feature extractor')

parser.add_argument('--vdir', type=str)
parser.add_argument('--fdir', type=str)
parser.add_argument('--hflip', type=int, default=0)

parser.add_argument('--batch_size', type=int, default=64,
                            help='batch size')
parser.add_argument('--type', type=str, default='2d',
                            help='CNN type')
parser.add_argument('--half_precision', type=int, default=0,
                            help='output half precision float')
parser.add_argument('--num_decoding_thread', type=int, default=4,
                            help='Num parallel thread for video decoding')
parser.add_argument('--l2_normalize', type=int, default=1,
                            help='l2 normalize feature')
parser.add_argument('--resnext101_model_path', type=str, default='model/resnext101.pth',
                            help='Resnext model path')
parser.add_argument('--vmz_model_path', type=str, default='model/r2plus1d_34_clip8_ig65m_from_scratch-9bae36ae.pth',
                            help='vmz model path')

args = parser.parse_args()

# TODO: refactor all args into config. (current code is from different people.)
CONFIGS = {
    "2d": {
        "fps": 1,
        "size": 224,
        "centercrop": False,
        "shards": 0,
    },
    "3d": {
        "fps": 24,
        "size": 112,
        "centercrop": True,
        "shards": 0,
    },
    "s3d": {
        "fps": 30,
        "size": 224,
        "centercrop": True,
        "shards": 0,
    },
    "vmz": {
        "fps": 24,
        "size": 112,
        "centercrop": True,
        "shards": 0,
    },
    "vae": {
        "fps": 2,
        "size": 256,
        "centercrop": True,
        "shards": 100,
    }
}

config = CONFIGS[args.type]

video_dirs = args.vdir
feature_dir = args.fdir

video_dict = PathBuilder.build(video_dirs, feature_dir, ".npy", config["shards"])

dataset = VideoLoader(
    video_dict=video_dict,
    framerate=config["fps"],
    size=config["size"],
    centercrop=config["centercrop"],
    hflip=args.hflip
)
n_dataset = len(dataset)
sampler = RandomSequenceSampler(n_dataset, 10)
loader = DataLoader(
    dataset,
    batch_size=1,
    shuffle=False,
    num_workers=args.num_decoding_thread,
    sampler=sampler if n_dataset > 10 else None,
)
preprocess = Preprocessing(args.type)
model = get_model(args)

with th.no_grad():
    for k, data in tqdm(enumerate(loader), total=loader.__len__(), ascii=True):
        input_file = data['input'][0]
        output_file = data['output'][0]
        if len(data['video'].shape) > 3:
            video = data['video'].squeeze()
            if len(video.shape) == 4:
                video = preprocess(video)
                n_chunk = len(video)
                if args.type == 'vmz':
                    n_chunk = math.ceil(n_chunk/float(3))
                    features = th.cuda.FloatTensor(n_chunk, 512).fill_(0)
                elif args.type == 's3d':
                    features = th.cuda.FloatTensor(n_chunk, 512).fill_(0)
                elif args.type == "vae":
                    features = th.cuda.LongTensor(n_chunk, 1024).fill_(0)
                else:
                    features = th.cuda.FloatTensor(n_chunk, 2048).fill_(0)
                n_iter = int(math.ceil(n_chunk / float(args.batch_size)))
                for i in range(n_iter):
                    factor = 1
                    if args.type == 'vmz':
                        factor = 3
                    min_ind = factor * i * args.batch_size
                    max_ind = factor * (i + 1) * args.batch_size
                    video_batch = video[min_ind:max_ind:factor].cuda()
                    if args.type == '2d':
                        batch_features = model(video_batch) # (51, 487), (51, 512)
                    elif args.type == 's3d':
                        batch_features = model(video_batch)
                        batch_features = batch_features['video_embedding']
                    elif args.type == "vae":
                        # image_code.
                        batch_features = model(video_batch)
                    else:
                        batch_pred, batch_features = model(video_batch) # (51, 487), (51, 512)
                    if args.l2_normalize:
                        batch_features = F.normalize(batch_features, dim=1)
                    features[i*args.batch_size:(i+1)*args.batch_size] = batch_features
                features = features.cpu().numpy()
                if args.half_precision:
                    if args.type == "vae":
                        features = features.astype(np.int16)
                    else:
                        features = features.astype('float16')
                else:
                    if args.type == "vae":
                        features = features.astype(np.int32)
                    else:
                        features = features.astype('float32')
                np.save(output_file, features)
        else:
            print('Video {} error.'.format(input_file))


import numpy as np
from settings import *
import random
from keras.models import Sequential
from keras.layers import Activation, Dense, LSTM, BatchNormalization
from keras.optimizers import Adam

class DQNAgent:

<<<<<<< HEAD
class Agent:      # 매수, 매도, Hold 를 수행하는 주체 (policy_learner.init 에서 생성)
    # 에이전트 상태가 구성하는 값 개수
    STATE_DIM = 2  # 주식 보유 비율, 포트폴리오 가치 비율, policy_learner 에서 입력 data feature 로 추가

    # 매매 수수료 및 세금
    TRADING_CHARGE = 0.00015  # 거래 수수료(일반적으로 0.015%)
    TRADING_TAX = 0.003  # 거래세 0.3%

    # 행동
    ACTION_BUY = 0  # 매수
    ACTION_SELL = 1  # 매도
    ACTION_HOLD = 2  # 홀딩 (관망)
    ACTIONS = [ACTION_BUY, ACTION_SELL, ACTION_HOLD]  # 인공 신경망에서 softmax 로 구할 행동들
    NUM_ACTIONS = len(ACTIONS)  # 인공 신경망에서 고려할 출력값의 개수 (3)

    def __init__(
        self, environment, min_trading_unit=0, max_trading_unit=0,
        delayed_reward_threshold=0):
        # Environment 객체
        self.environment = environment  # 현재 주식 가격을 가져오기 위해 환경 참조

        # 최소 매매 단위, 최대 매매 단위, 지연보상 임계치
        self.min_trading_unit = min_trading_unit  # 최소 단일 거래 단위
        self.max_trading_unit = max_trading_unit  # 최대 단일 거래 단위
        self.delayed_reward_threshold = delayed_reward_threshold  # 지연보상 임계치 (손익률이 threshold 를 넘으면 지연보상 발생)

        # Agent 클래스의 속성
        self.initial_balance = 0  # 초기 자본금 (투자 시작 시점의 보유 현금)
        self.balance = 0  # 현재 현금 잔고
        self.num_stocks = 0  # 현재 보유 주식 수
        self.portfolio_value = 0  # balance + num_stocks * {현재 주식 가격} --> 포트폴리오 가치
        self.base_portfolio_value = 0  # 직전 학습 시점의 PV (기준 포트폴리오)
        self.base_stock_price = 0      # 직전 학습 시점의 주가
        self.num_buy = 0  # 매수 횟수
        self.num_sell = 0  # 매도 횟수
        self.num_hold = 0  # 홀딩 횟수
        self.immediate_reward = 0  # 즉시 보상 (행동을 수행한 시점에 수익발행 +1, 아니면 -1)
        self.buy_charge = 0   # 매수 수수료
        self.sell_charge = 0  # 매도 수수료
        self.sell_tax  = 0    # 거래세

        # Agent 클래스의 상태
        self.ratio_hold = 0  # 주식 보유 비율 (현재보유하고있는 주식수/최대로 보유할수 있는 주식수)
        self.ratio_portfolio_value = 0  # 포트폴리오 가치 비율 (현재 PV/직전지연보상시점의 PV)

    def reset(self):   # class 속성 초기화 (매 epoch 마다)
        self.balance = self.initial_balance
        self.num_stocks = 0
        self.portfolio_value = self.initial_balance
        self.base_portfolio_value = self.initial_balance
        self.num_buy = 0
        self.num_sell = 0
        self.num_hold = 0
        self.immediate_reward = 0
        self.buy_charge = 0
        self.sell_charge = 0
        self.sell_tax  = 0
        self.ratio_hold = 0
        self.ratio_portfolio_value = 0

    def set_balance(self, balance):
        self.initial_balance = balance     # 초기 자본금 설정

    def get_states(self):                 # agent 의 상태 반환
        self.ratio_hold = self.num_stocks / int(
            self.portfolio_value / self.environment.get_price())     # 보유주식수/(PV/현재주가)
        self.ratio_portfolio_value = self.portfolio_value / self.base_portfolio_value     # PV/기준PV
        return (
            self.ratio_hold,              # 주식보유비율 - 0: 주식 없음, 1: 최대보유
            self.ratio_portfolio_value    # 포트폴리오 가치 비율 - < 1 : 손실, 1 > : 수익발생
        )

    def decide_action(self, policy_network, sample, epsilon):
        # confidence = 0.
        # 탐험 결정
        if np.random.rand() < epsilon:
            exploration = True
            action = np.random.randint(self.NUM_ACTIONS)  # 무작위로 행동 결정
        else:
            exploration = False
            probs = policy_network.predict(sample)  # 각 행동에 대한 확률, 지금까지 훈련시킨 청책망으로 다음 action 결정
            action = np.argmax(probs)
            # confidence = probs[action]

        # return action, confidence, exploration
        return action, exploration     # 0:매수, 1:매도, 2:HOLD / False/True(탐험여부)

    def validate_action(self, action):
        validity = True
        if action == Agent.ACTION_BUY:  # 매수의 경우
            # 적어도 1주를 살 수 있는지 확인
            if self.balance < self.environment.get_price() * (
                1 + self.TRADING_CHARGE) * self.min_trading_unit:
                validity = False
        elif action == Agent.ACTION_SELL:   # 매도의 경우
            # 주식 잔고가 있는지 확인하여 없으면 False return
            if self.num_stocks <= 0:
                validity = False
        return validity

    def decide_trading_unit(self, action, curr_price):   # 매수 단위 결정
        if action == Agent.ACTION_BUY:    # 매수의 경우
            trading_unit = int(self.balance / (curr_price * (1 + self.TRADING_CHARGE)))   # 가능한 최대 주식 매수
        elif action == Agent.ACTION_SELL:  # 매도의 경우
            trading_unit = self.num_stocks  # 보유수량 전부 매도
        else:
            trading_unit = 0               # Hold
        return trading_unit
=======
    def __init__(self, env, state_size):
        # Environment 객체
        self.env = env # 현재 주식 가격을 가져오기 위해 환경 참조
        self.gamma = GAMMA
        self.epsilon = EPSILON
        self.epsilon_decay = EPSILON_DECAY
        self.learning_rate = LEARNING_RATE
        self.state_size = state_size

        self.model        = self._build_model()
        self.target_model = self._build_model()

        # # 최소 매매 단위, 최대 매매 단위, 지연보상 임계치
        # self.min_trading_unit = min_trading_unit  # 최소 단일 거래 단위
        # self.max_trading_unit = max_trading_unit  # 최대 단일 거래 단위
        # self.delayed_reward_threshold = delayed_reward_threshold  # 지연보상 임계치 (손익률이 threshold 를 넘으면 지연보상 발생)
        #
        # # Agent 클래스의 속성
        # self.initial_balance = 0  # 초기 자본금 (투자 시작 시점의 보유 현금)
        # self.balance = 0  # 현재 현금 잔고
        # self.num_stocks = 0  # 현재 보유 주식 수
        # self.portfolio_value = 0  # balance + num_stocks * {현재 주식 가격} --> 포트폴리오 가치
        # self.base_portfolio_value = 0  # 직전 학습 시점의 PV (기준 포트폴리오)
        # self.base_stock_price = 0      # 직전 학습 시점의 주가
        self.num_buy = 0  # 매수 횟수
        self.num_sell = 0  # 매도 횟수
        self.num_hold = 0  # 홀딩 횟수
        self.profitloss = 0  # 누적 손익
        # self.immediate_reward = 0  # 즉시 보상 (행동을 수행한 시점에 수익발행 +1, 아니면 -1)
        # self.buy_charge = 0   # 매수 수수료
        # self.sell_charge = 0  # 매도 수수료
        # self.sell_tax  = 0    # 거래세
        #
        # # Agent 클래스의 상태
        # self.ratio_hold = 0  # 주식 보유 비율 (현재보유하고있는 주식수/최대로 보유할수 있는 주식수)
        # self.ratio_portfolio_value = 0  # 포트폴리오 가치 비율 (현재 PV/직전지연보상시점의 PV)
    #

    def _build_model(self):
        model = Sequential()
        model.add(LSTM(128, input_shape=(None, self.state_size),
                return_sequences=False, dropout=0.2))
        # model.add(BatchNormalization())
        # model.add(LSTM(64, return_sequences=False, dropout=0.2))
        model.add(BatchNormalization())
        model.add(Dense(32))
        model.add(Dense(ACTION_SIZE, activation='linear'))

        model.compile(optimizer=Adam(lr=LEARNING_RATE), loss='mse')
        return model

    def replay(self, replay_buffer, batch_size):

        if len(replay_buffer) < batch_size:
            return

        minibatch = random.sample(replay_buffer, batch_size)

        for episode_rewarded in minibatch:
            y = np.zeros((len(episode_rewarded), ACTION_SIZE))

            for (state, action, reward, next_state, done) in episode_rewarded:

                target_f = self.target_model.predict(state)

                if done:
                    target_f[:, action] = reward
                    if target_f.shape[0] != 1:
                        print("Inavlid target_f shape : ", target_f.shape)
                else:
                    target_f[:, action] = (reward + self.gamma * \
                               np.amax(self.target_model.predict(next_state)[:, action]))
                    if target_f.shape[0] != 1:
                       print("Inavlid target_f shape : ", target_f.shape)

                self.model.fit(state, target_f, epochs=1, verbose=0)

    # def reset(self):   # class 속성 초기화 (매 epoch 마다)
    #     self.balance = self.initial_balance
    #     self.num_stocks = 0
    #     self.portfolio_value = self.initial_balance
    #     self.base_portfolio_value = self.initial_balance
    #     self.num_buy = 0
    #     self.num_sell = 0
    #     self.num_hold = 0
    #     self.immediate_reward = 0
    #     self.buy_charge = 0
    #     self.sell_charge = 0
    #     self.sell_tax  = 0
    #     self.ratio_hold = 0
    #     self.ratio_portfolio_value = 0
    #
    # def set_balance(self, balance):
    #     self.initial_balance = balance     # 초기 자본금 설정
    #
    # def get_states(self):                 # agent 의 상태 반환
    #     self.ratio_hold = self.num_stocks / int(
    #         self.portfolio_value / self.environment.get_price())     # 보유주식수/(PV/현재주가)
    #     self.ratio_portfolio_value = self.portfolio_value / self.base_portfolio_value     # PV/기준PV
    #     return (
    #         self.ratio_hold,              # 주식보유비율 - 0: 주식 없음, 1: 최대보유
    #         self.ratio_portfolio_value    # 포트폴리오 가치 비율 - < 1 : 손실, 1 > : 수익발생
    #     )

    # def decide_action(self, policy_network, sample, epsilon):
    #     # confidence = 0.
    #     # 탐험 결정
    #     if np.random.rand() < epsilon:
    #         exploration = True
    #         action = np.random.randint(self.NUM_ACTIONS)  # 무작위로 행동 결정
    #     else:
    #         exploration = False
    #         probs = policy_network.predict(sample)  # 각 행동에 대한 확률, 지금까지 훈련시킨 청책망으로 다음 action 결정
    #         action = np.argmax(probs)
    #         # confidence = probs[action]
    #
    #     # return action, confidence, exploration
    #     return action, exploration     # 0:매수, 1:매도, 2:HOLD / False/True(탐험여부)
    #
    # def validate_action(self, action):
    #     validity = True
    #     if action == Agent.ACTION_BUY:  # 매수의 경우
    #         # 적어도 1주를 살 수 있는지 확인
    #         if self.balance < self.environment.get_price() * (
    #             1 + self.TRADING_CHARGE) * self.min_trading_unit:
    #             validity = False
    #     elif action == Agent.ACTION_SELL:   # 매도의 경우
    #         # 주식 잔고가 있는지 확인하여 없으면 False return
    #         if self.num_stocks <= 0:
    #             validity = False
    #     return validity
    #
    # def decide_trading_unit(self, action, curr_price):   # 매수 단위 결정
    #     if action == Agent.ACTION_BUY:    # 매수의 경우
    #         trading_unit = int(self.balance / (curr_price * (1 + self.TRADING_CHARGE)))   # 가능한 최대 주식 매수
    #     elif action == Agent.ACTION_SELL:  # 매도의 경우
    #         trading_unit = self.num_stocks  # 보유수량 전부 매도
    #     else:
    #         trading_unit = 0               # Hold
    #     return trading_unit
>>>>>>> 6ad1fe99b5f6f32b80563112803f402ef76f9fa6

    # def decide_trading_unit(self, confidence):   # 확률이 높을수록 더 많은 주식 거래  --> 차라리 주식 변동 예상 폭이 클때로 변경이 좋을 것임(다음 version 고려)
    #     if np.isnan(confidence):
    #         return self.min_trading_unit
    #     added_traiding = max(min(
    #         int(confidence * (self.max_trading_unit - self.min_trading_unit)),
    #         self.max_trading_unit-self.min_trading_unit
    #     ), 0)
    #     return self.min_trading_unit + added_traiding

    #def act(self, action, confidence):
<<<<<<< HEAD
    def act(self, action):      # Agent 가 결정한 행동을 수행
        if not self.validate_action(action):     # action 을 할 수 없는 경우 관망
            action = Agent.ACTION_HOLD

        # 환경에서 현재 가격 얻기 (매일의 종가로 매매를 가정함)
        curr_price = self.environment.get_price()

        # 즉시 보상 초기화
        self.immediate_reward = 0

        # 매수/매도시 거래할 주식 수 결정
        trading_unit = self.decide_trading_unit(action, curr_price)

        # 매수
        if action == Agent.ACTION_BUY:
            balance = self.balance - curr_price * (1 + self.TRADING_CHARGE) * trading_unit
            # 보유 현금이 모자라거나 최대 trading_unit 을 초과한 경우 가능한 만큼 최대한 매수
            # if balance < 0:
            #     trading_unit = max(min(
            #         int(self.balance / (
            #             curr_price * (1 + self.TRADING_CHARGE))), self.max_trading_unit),
            #         self.min_trading_unit
            #     )
            # 수수료를 적용하여 총 매수 금액 산정
            buy_amount = curr_price * (1 + self.TRADING_CHARGE) * trading_unit
            self.balance -= buy_amount       # 보유 현금을 갱신
            self.num_stocks += trading_unit  # 보유 주식 수를 갱신
            self.num_buy += 1                # 매수 횟수 증가
            self.buy_charge += curr_price * self.TRADING_CHARGE * trading_unit   # 매수 수수료 누적
        # 매도
        elif action == Agent.ACTION_SELL:
            # 보유 주식이 모자랄 경우 가능한 만큼 최대한 매도
            #trading_unit = min(trading_unit, self.num_stocks)
            sell_amount = curr_price * (
                1 - (self.TRADING_TAX + self.TRADING_CHARGE)) * trading_unit
            self.num_stocks -= trading_unit  # 보유 주식 수를 갱신
            self.balance += sell_amount      # 보유 현금을 갱신
            self.num_sell += 1               # 매도 횟수 증가
            self.sell_charge += curr_price * self.TRADING_CHARGE * trading_unit   # 매도 수수료 누적
            self.sell_tax += curr_price * self.TRADING_TAX * trading_unit         # 거래세 누적
        # 홀딩
        elif action == Agent.ACTION_HOLD:
            self.num_hold += 1  # 홀딩 횟수 증가
        else:
            print("invalid action", action)
            quit()

        # 포트폴리오 가치 갱신
        self.portfolio_value = self.balance + curr_price * self.num_stocks
        # 주식 보유 비율 계산
        num_stocks_ratio = self.num_stocks / (self.balance / curr_price + self.num_stocks)
        #-------- 초기 포트폴리오 가치와 현재 포트폴리오 가치 비교하여 손익 비율 계산
        profitloss = (self.portfolio_value - self.base_portfolio_value) / self.base_portfolio_value

        # 즉시 보상 판단
        self.immediate_reward = 1 if profitloss >= 0 else -1

        # 지연 보상 판단
        if profitloss > self.delayed_reward_threshold:
            delayed_reward = 1
            # 목표 수익률 달성하여 기준 포트폴리오 가치 갱신
            self.base_portfolio_value = self.portfolio_value
        elif profitloss < -self.delayed_reward_threshold:
            delayed_reward = -1
            # 손실 기준치를 초과하여 기준 포트폴리오 가치 갱신
            self.base_portfolio_value = self.portfolio_value
        else:
            delayed_reward = 0

        self.base_stock_price = curr_price   # 기준 포트폴리오 가치 갱신 시점의 주가 save

        return self.immediate_reward, delayed_reward
=======
    # def act(self, action):      # Agent 가 결정한 행동을 수행
    #     if not self.validate_action(action):     # action 을 할 수 없는 경우 관망
    #         action = Agent.ACTION_HOLD
    #
    #     # 환경에서 현재 가격 얻기 (매일의 종가로 매매를 가정함)
    #     curr_price = self.environment.get_price()
    #
    #     # 즉시 보상 초기화
    #     self.immediate_reward = 0
    #
    #     # 매수/매도시 거래할 주식 수 결정
    #     trading_unit = self.decide_trading_unit(action, curr_price)
    #
    #     # 매수
    #     if action == Agent.ACTION_BUY:
    #         balance = self.balance - curr_price * (1 + self.TRADING_CHARGE) * trading_unit
    #         # 보유 현금이 모자라거나 최대 trading_unit 을 초과한 경우 가능한 만큼 최대한 매수
    #         # if balance < 0:
    #         #     trading_unit = max(min(
    #         #         int(self.balance / (
    #         #             curr_price * (1 + self.TRADING_CHARGE))), self.max_trading_unit),
    #         #         self.min_trading_unit
    #         #     )
    #         # 수수료를 적용하여 총 매수 금액 산정
    #         buy_amount = curr_price * (1 + self.TRADING_CHARGE) * trading_unit
    #         self.balance -= buy_amount       # 보유 현금을 갱신
    #         self.num_stocks += trading_unit  # 보유 주식 수를 갱신
    #         self.num_buy += 1                # 매수 횟수 증가
    #         self.buy_charge += curr_price * self.TRADING_CHARGE * trading_unit   # 매수 수수료 누적
    #     # 매도
    #     elif action == Agent.ACTION_SELL:
    #         # 보유 주식이 모자랄 경우 가능한 만큼 최대한 매도
    #         #trading_unit = min(trading_unit, self.num_stocks)
    #         sell_amount = curr_price * (
    #             1 - (self.TRADING_TAX + self.TRADING_CHARGE)) * trading_unit
    #         self.num_stocks -= trading_unit  # 보유 주식 수를 갱신
    #         self.balance += sell_amount      # 보유 현금을 갱신
    #         self.num_sell += 1               # 매도 횟수 증가
    #         self.sell_charge += curr_price * self.TRADING_CHARGE * trading_unit   # 매도 수수료 누적
    #         self.sell_tax += curr_price * self.TRADING_TAX * trading_unit         # 거래세 누적
    #     # 홀딩
    #     elif action == Agent.ACTION_HOLD:
    #         self.num_hold += 1  # 홀딩 횟수 증가
    #     else:
    #         print("invalid action", action)
    #         quit()
    #
    #     # 포트폴리오 가치 갱신
    #     self.portfolio_value = self.balance + curr_price * self.num_stocks
    #     # 주식 보유 비율 계산
    #     num_stocks_ratio = self.num_stocks / (self.balance / curr_price + self.num_stocks)
    #     #-------- 초기 포트폴리오 가치와 현재 포트폴리오 가치 비교하여 손익 비율 계산
    #     profitloss = (self.portfolio_value - self.base_portfolio_value) / self.base_portfolio_value
    #
    #     # 즉시 보상 판단
    #     self.immediate_reward = 1 if profitloss >= 0 else -1
    #
    #     # 지연 보상 판단
    #     if profitloss > self.delayed_reward_threshold:
    #         delayed_reward = 1
    #         # 목표 수익률 달성하여 기준 포트폴리오 가치 갱신
    #         self.base_portfolio_value = self.portfolio_value
    #     elif profitloss < -self.delayed_reward_threshold:
    #         delayed_reward = -1
    #         # 손실 기준치를 초과하여 기준 포트폴리오 가치 갱신
    #         self.base_portfolio_value = self.portfolio_value
    #     else:
    #         delayed_reward = 0
    #
    #     self.base_stock_price = curr_price   # 기준 포트폴리오 가치 갱신 시점의 주가 save
    #
    #     return self.immediate_reward, delayed_reward
>>>>>>> 6ad1fe99b5f6f32b80563112803f402ef76f9fa6