import os
import glob
import cv2
import random
import numpy as np
import rasterio
import slidingwindow as sw

#Plotting and polygon overlap
from shapely.geometry import box, shape, Point
from rtree import index
from PIL import Image
import geopandas as gp

#NEON recall rate
import pandas as pd
from matplotlib import pyplot

from keras_retinanet.utils.visualization import draw_detections, draw_annotations
from keras_retinanet.utils.eval import _get_detections

#DeepForest
from DeepForest import Lidar
from DeepForest import postprocessing
from DeepForest import onthefly_generator
import copy

def neonRecall(
    sites,
    generator,
    model,
    score_threshold=0.05,
    max_detections=100,
    suppression_threshold=0.15,
    save_path=None,
    experiment=None):

    point_contains = [ ]

    site_data_dict = {}
    for site in sites:
        #Container for recall pts.

        #load field data
        field_data = pd.read_csv("data/field_data.csv")
        field_data = field_data[field_data['UTM_E'].notnull()]

        #select site
        site_data = field_data[field_data["siteID"]==site]

        #select tree species
        specieslist = pd.read_csv("data/AcceptedSpecies.csv",encoding="utf-8")
        specieslist =  specieslist[specieslist["siteID"] == site]

        site_data = site_data[site_data["scientificName"].isin(specieslist["scientificName"].values)]

        #Single bole individuals as representitve, no individualID ending in non-digits
        site_data = site_data[site_data["individualID"].str.contains("\d$")]
        site_data_dict[site] = site_data

        #Only data within the last two years, sites can be hand managed
        #site_data=site_data[site_data["eventID"].str.contains("2015|2016|2017|2018")]

    for i in range(generator.size()):

        #Load image
        raw_image    = generator.load_image(i)
        plot_image = copy.deepcopy(raw_image)

        #Skip if missing a component data source
        if raw_image is False:
            print("Empty image, skipping")
            continue

        #Store plotting images.
        plot_rgb = plot_image[:,:,:3].copy()
        plot_chm = plot_image[:,:,3]

        image        = generator.preprocess_image(raw_image)
        image, scale = generator.resize_image(image)

        # run network
        boxes, scores, labels = model.predict_on_batch(np.expand_dims(image, axis=0))[:3]

        # correct boxes for image scale
        boxes /= scale

        # select indices which have a score above the threshold
        indices = np.where(scores[0, :] > score_threshold)[0]

        # select those scores
        scores = scores[0][indices]

        # find the order with which to sort the scores
        scores_sort = np.argsort(-scores)[:max_detections]

        # select detections
        image_boxes      = boxes[0, indices[scores_sort], :]
        image_scores     = scores[scores_sort]
        image_labels     = labels[0, indices[scores_sort]]
        image_detections = np.concatenate([image_boxes, np.expand_dims(image_scores, axis=1), np.expand_dims(image_labels, axis=1)], axis=1)

        #Find geographic bounds
        base_dir = generator.DeepForest_config[generator.row["site"]][generator.name]["RGB"]
        tile_path = os.path.join(base_dir, generator.image_data[i]["tile"])

        with rasterio.open(tile_path) as dataset:
            tile_bounds = dataset.bounds

        #drape boxes
        #get lidar cloud if a new tile, or if not the same tile as previous image.
        if i == 0:
            generator.load_lidar_tile()
        elif not generator.image_data[i]["tile"] == generator.image_data[i-1]["tile"]:
            generator.load_lidar_tile()

        #The tile could be the full tile, so let's check just the 400 pixel crop we are interested
        #Not the best structure, but the on-the-fly generator always has 0 bounds
        if hasattr(generator, 'hf'):
            bounds = generator.hf["utm_coords"][generator.row["window"]]
        else:
            bounds=[]

        density = Lidar.check_density(generator.lidar_tile, bounds=bounds)

        #print("Point density is {:.2f}".format(density))

        if density > 100:
            #find window utm coordinates
            #print("Bounds for image {}, window {}, are {}".format(generator.row["tile"], generator.row["window"], bounds))
            pc = postprocessing.drape_boxes(boxes=image_boxes, pc = generator.lidar_tile, bounds=bounds)

            #Get new bounding boxes
            new_boxes = postprocessing.cloud_to_box(pc, bounds)
            new_scores = image_scores[:new_boxes.shape[0]]
            new_labels = image_labels[:new_boxes.shape[0]]
            image_detections = np.concatenate([new_boxes, np.expand_dims(new_scores, axis=1), np.expand_dims(new_labels, axis=1)], axis=1)

        else:
            #print("Point density of {:.2f} is too low, skipping image {}".format(density, generator.row["tile"]))
            pass

        #add spatial NEON points
        site_data =site_data_dict[generator.row["site"]]
        plotID = os.path.splitext(generator.image_data[i]["tile"])[0]
        plot_data = site_data[site_data.plotID == plotID]

        #Save image and send it to logger
        if save_path is not None:

            x = (plot_data.UTM_E - tile_bounds.left).values / 0.1
            y = (tile_bounds.top - plot_data.UTM_N).values / 0.1

            for j in np.arange(len(x)):
                cv2.circle(plot_image,(int(x[j]),int(y[j])), 2, (0,0,255), -1)

            #Write RGB
            draw_detections(plot_rgb, image_boxes, image_scores, image_labels, label_to_name=generator.label_to_name,score_threshold=score_threshold)

            #name image
            image_name=generator.image_names[i]
            row=generator.image_data[image_name]
            fname=os.path.splitext(row["tile"])[0] + "_" + str(row["window"])

            #Write RGB
            cv2.imwrite(os.path.join(save_path, '{}_NeonPlot.png'.format(fname)), plot_rgb)

            plot_chm = plot_chm/plot_chm.max() * 255
            chm = np.uint8(plot_chm)
            draw_detections(chm, image_boxes, image_scores, image_labels, label_to_name=generator.label_to_name, score_threshold=score_threshold, color = (80,127,255))
            cv2.imwrite(os.path.join(save_path, '{}_Lidar_NeonPlot.png'.format(plotID)), chm)

            #Format name and save
            if experiment:
                experiment.log_image(os.path.join(save_path, '{}_NeonPlot.png'.format(plotID)),file_name=str(plotID))
                experiment.log_image(os.path.join(save_path, '{}_Lidar_NeonPlot.png'.format(plotID)),file_name=str("Lidar_" + plotID))

        #calculate recall
            s = gp.GeoSeries(map(Point, zip(plot_data.UTM_E, plot_data.UTM_N)))

        #Calculate recall
        projected_boxes = []

        for row in  image_boxes:
            #Add utm bounds and create a shapely polygon
            pbox=create_polygon(row, tile_bounds, cell_size=0.1)
            projected_boxes.append(pbox)

        #for each point, is it within a prediction?
        for index, tree in plot_data.iterrows():
            p=Point(tree.UTM_E, tree.UTM_N)

            within_polygon=[]
            for prediction in projected_boxes:
                within_polygon.append(p.within(prediction))

            #Check for overlapping polygon, add it to list
            is_within = sum(within_polygon) > 0
            point_contains.append(is_within)

    #sum recall across plots
    if len(point_contains)==0:
        recall = None
    else:
        ## Recall rate for plot
        recall = sum(point_contains)/len(point_contains)

    return(recall)

#IoU for non-rectangular polygons
def compute_windows(numpy_image, pixels=400, overlap=0.05):
    windows = sw.generate(numpy_image, sw.DimOrder.HeightWidthChannel, pixels,overlap )
    return(windows)

def retrieve_window(numpy_image,window):
    crop=numpy_image[window.indices()]
    return(crop)

def non_max_suppression(boxes, overlapThresh):
    # if there are no boxes, return an empty list
    if len(boxes) == 0:
        return []

    # if the bounding boxes integers, convert them to floats --
    # this is important since we'll be doing a bunch of divisions
    if boxes.dtype.kind == "i":
        boxes = boxes.astype("float")

    # initialize the list of picked indexes
    pick = []

    # grab the coordinates of the bounding boxes
    x1 = boxes[:,0]
    y1 = boxes[:,1]
    x2 = boxes[:,2]
    y2 = boxes[:,3]

    # compute the area of the bounding boxes and sort the bounding
    # boxes by the bottom-right y-coordinate of the bounding box
    area = (x2 - x1 + 1) * (y2 - y1 + 1)
    idxs = np.argsort(y2)

    # keep looping while some indexes still remain in the indexes
    # list
    while len(idxs) > 0:
        # grab the last index in the indexes list and add the
        # index value to the list of picked indexes
        last = len(idxs) - 1
        i = idxs[last]
        pick.append(i)

        # find the largest (x, y) coordinates for the start of
        # the bounding box and the smallest (x, y) coordinates
        # for the end of the bounding box
        xx1 = np.maximum(x1[i], x1[idxs[:last]])
        yy1 = np.maximum(y1[i], y1[idxs[:last]])
        xx2 = np.minimum(x2[i], x2[idxs[:last]])
        yy2 = np.minimum(y2[i], y2[idxs[:last]])

        # compute the width and height of the bounding box
        w = np.maximum(0, xx2 - xx1 + 1)
        h = np.maximum(0, yy2 - yy1 + 1)

        # compute the ratio of overlap
        overlap = (w * h) / area[idxs[:last]]

        # delete all indexes from the index list that have
        idxs = np.delete(idxs, np.concatenate(([last],
                                               np.where(overlap > overlapThresh)[0])))

    # return the indices for only the bounding boxes that were picked using the
    # integer data type
    return pick

def create_polygon(row, bounds, cell_size):

    #boxes are in form x1, y1, x2, y2, add the origin utm extent
    x1= (row[0]*cell_size) + bounds.left
    y1 = bounds.top - (row[1]*cell_size)
    x2 =(row[2]*cell_size) + bounds.left
    y2 = bounds.top - (row[3]*cell_size)

    b = box(x1, y1, x2, y2)

    return(b)

def IoU_polygon(a, b):

    #Area of predicted box
    predicted_area=b.area

    #Area of ground truth polygon
    polygon_area=a.area

    #Intersection
    intersection_area=a.intersection(b).area
    iou = intersection_area / float(predicted_area + polygon_area - intersection_area)

    return iou

def calculateIoU(itcs, predictions):

    '''
    1) Find overlap among polygons efficiently
    2) Calulate a cost matrix of overlap, with rows as itcs and columns as predictions
    3) Hungarian matching for pairing
    4) Calculate intersection over union (IoU)
    5) Mean IoU returned.
    '''
    # Populate R-tree index with bounds of prediction boxes
    idx = index.Index()

    for pos, cell in enumerate(predictions):
        # assuming cell is a shapely object
        idx.insert(pos, cell.bounds)

    #Create polygons
    itc_polygons=[shape(x["geometry"]) for x in itcs["data"]]

    overlap_dict={}

    #select predictions that overlap with the polygons
    matched=[predictions[x] for x in idx.intersection(itcs["bounds"])]

    #Create a container
    cost_matrix=np.zeros((len(itc_polygons),len(matched)))

    for x,poly in enumerate(itc_polygons):
        for y,match in enumerate(matched):
            cost_matrix[x,y]= poly.intersection(match).area

    #Assign polygon pairs
    assignments=linear_sum_assignment(-1 *cost_matrix)

    iou_list=[]

    for i in np.arange(len(assignments[0])):
        a=itc_polygons[assignments[0][i]]
        b=matched[assignments[1][i]]
        iou=IoU_polygon(a,b)
        iou_list.append(iou)

    return(iou_list)


from __future__ import division
import collections
import mxnet as mx
import numpy as np
from numpy.linalg import norm
import mxnet.ndarray as nd
from ..model_zoo import model_zoo
from ..utils import face_align

__all__ = ['FaceAnalysis', 'Face']

Face = collections.namedtuple('Face', [
    'bbox', 'landmark', 'det_score', 'embedding', 'gender', 'age',
    'embedding_norm', 'normed_embedding'
])

Face.__new__.__defaults__ = (None, ) * len(Face._fields)

class FaceAnalysis:
    def __init__(self,
                 det_name='retinaface_mnet025_v2',
                 rec_name='arcface_r100_v1',
                 ga_name='genderage_v1'):
        assert det_name is not None
        self.det_model = model_zoo.get_model(det_name)
        if rec_name is not None:
            self.rec_model = model_zoo.get_model(rec_name)
        else:
            self.rec_model = None
        if ga_name is not None:
            self.ga_model = model_zoo.get_model(ga_name)
        else:
            self.ga_model = None

    def prepare(self, ctx_id, nms=0.4):
        self.det_model.prepare(ctx_id, nms)
        if self.rec_model is not None:
            self.rec_model.prepare(ctx_id)
        if self.ga_model is not None:
            self.ga_model.prepare(ctx_id)

    def get(self, img, det_thresh=0.8, det_scale=1.0, max_num=0):
        bboxes, landmarks = self.det_model.detect(img,
                                                  threshold=det_thresh,
                                                  scale=det_scale)
        if bboxes.shape[0] == 0:
            return []
        if max_num > 0 and bboxes.shape[0] > max_num:
            area = (bboxes[:, 2] - bboxes[:, 0]) * (bboxes[:, 3] -
                                                    bboxes[:, 1])
            img_center = img.shape[0] // 2, img.shape[1] // 2
            offsets = np.vstack([
                (bboxes[:, 0] + bboxes[:, 2]) / 2 - img_center[1],
                (bboxes[:, 1] + bboxes[:, 3]) / 2 - img_center[0]
            ])
            offset_dist_squared = np.sum(np.power(offsets, 2.0), 0)
            values = area - offset_dist_squared * 2.0  # some extra weight on the centering
            bindex = np.argsort(
                values)[::-1]  # some extra weight on the centering
            bindex = bindex[0:max_num]
            bboxes = bboxes[bindex, :]
            landmarks = landmarks[bindex, :]
        ret = []
        for i in range(bboxes.shape[0]):
            bbox = bboxes[i, 0:4]
            det_score = bboxes[i, 4]
            landmark = landmarks[i]
            _img = face_align.norm_crop(img, landmark=landmark)
            embedding = None
            embedding_norm = None
            normed_embedding = None
            gender = None
            age = None
            if self.rec_model is not None:
                embedding = self.rec_model.get_embedding(_img).flatten()
                embedding_norm = norm(embedding)
                normed_embedding = embedding / embedding_norm
            if self.ga_model is not None:
                gender, age = self.ga_model.get(_img)
            face = Face(bbox=bbox,
                        landmark=landmark,
                        det_score=det_score,
                        embedding=embedding,
                        gender=gender,
                        age=age,
                        normed_embedding=normed_embedding,
                        embedding_norm=embedding_norm)
            ret.append(face)
        return ret


import argparse
import torch
import time
import utils.pytorch_util as ptu
from replay_buffer import ReplayBuffer, ReplayBufferCount
from replay_buffer_no_resampling import ReplayBufferNoResampling
from utils.env_utils import NormalizedBoxEnv, domain_to_epoch, env_producer
from utils.rng import set_global_pkg_rng_state
from launcher_util import run_experiment_here
from path_collector import MdpPathCollector, RemoteMdpPathCollector
from trainer.policies import TanhGaussianPolicy, MakeDeterministic, EnsemblePolicy, TanhGaussianMixturePolicy
from trainer.policies import GaussianPolicy
from trainer.trainer import SACTrainer
from trainer.ddpg_trainer import DDPGTrainer
from trainer.particle_trainer import ParticleTrainer
from trainer.gaussian_trainer import GaussianTrainer
from trainer.gaussian_trainer_ts import GaussianTrainerTS
from trainer.particle_trainer_ts import ParticleTrainerTS
from trainer.particle_trainer_oac import ParticleTrainer as ParticleTrainerOAC
from networks import FlattenMlp
from rl_algorithm import BatchRLAlgorithm
import numpy as np

# ray.init(
#     # If true, then output from all of the worker processes on all nodes will be directed to the driver.
#     log_to_driver=True,
#     logging_level=logging.WARNING,
#
#     # # The amount of memory (in bytes)
#     # object_store_memory=1073741824, # 1g
#     # redis_max_memory=1073741824 # 1g
# )

def get_current_branch(dir):
    from git import Repo

    repo = Repo(dir)
    return repo.active_branch.name

def get_policy_producer(obs_dim, action_dim, hidden_sizes, clip=True, std=None):
    def policy_producer(deterministic=False, bias=None, ensemble=False, n_policies=1, n_components=1,
                        approximator=None, share_layers=False):
        if ensemble:
            policy = EnsemblePolicy(approximator=approximator,
                                    hidden_sizes=hidden_sizes,
                                    obs_dim=obs_dim,
                                    action_dim=action_dim,
                                    n_policies=n_policies,
                                    bias=bias,
                                    share_layers=share_layers)
        else:
            if not clip:
                policy = GaussianPolicy(
                    obs_dim=obs_dim,
                    action_dim=action_dim,
                    hidden_sizes=hidden_sizes,
                    bias=bias
                )
            elif n_components == 1:
                policy = TanhGaussianPolicy(
                    obs_dim=obs_dim,
                    action_dim=action_dim,
                    hidden_sizes=hidden_sizes,
                    bias=bias,
                    std=std
                )

            else:
                policy = TanhGaussianMixturePolicy(
                    obs_dim=obs_dim,
                    action_dim=action_dim,
                    hidden_sizes=hidden_sizes,
                    bias=bias,
                    n_components=n_components
                )
            '''
            policy = TanhGaussianMixturePolicy(
                obs_dim=obs_dim,
                action_dim=action_dim,
                hidden_sizes=hidden_sizes,
                bias=bias,
                n_components=n_components
            )
            '''
            if deterministic:
                policy = MakeDeterministic(policy)

        return policy

    return policy_producer

def get_q_producer(obs_dim, action_dim, hidden_sizes, output_size=1):
    def q_producer(bias=None, positive=False, train_bias=True):
        return FlattenMlp(input_size=obs_dim + action_dim,
                          output_size=output_size,
                          hidden_sizes=hidden_sizes,
                          bias=bias,
                          positive=positive,
                          train_bias=train_bias)

    return q_producer

def experiment(variant, prev_exp_state=None):
    domain = variant['domain']
    seed = variant['seed']
    if seed == 0:
        np.random.seed()
        seed = np.random.randint(0, 1000000)
    torch.manual_seed(seed)
    np.random.seed(seed)
    env_args = {}
    if domain in ['riverswim']:
        env_args['dim'] = variant['dim']
    if domain in ['point']:
        env_args['difficulty'] = variant['difficulty']
        env_args['clip_state'] = variant['clip_state']
        env_args['terminal'] = variant['terminal']
        env_args['max_state'] = variant['max_state']
    if 'cliff' in domain:
        env_args['sigma_noise'] = variant['sigma_noise']

    expl_env = env_producer(domain, seed, **env_args)
    eval_env = env_producer(domain, seed * 10 + 1, **env_args)
    obs_dim = expl_env.observation_space.low.size
    action_dim = expl_env.action_space.low.size

    # Get producer function for policy and value functions
    M = variant['layer_size']
    N = variant['num_layers']
    n_estimators = variant['n_estimators']

    if variant['share_layers']:
        output_size = n_estimators
        # n_estimators = 1
    else:
        output_size = 1
    q_producer = get_q_producer(obs_dim, action_dim, hidden_sizes=[M] * N, output_size=output_size)
    std = None
    if variant['algorithm_kwargs']['ddpg_noisy']:
        std = variant['std']
        std = np.ones(action_dim) * std
    policy_producer = get_policy_producer(obs_dim, action_dim, hidden_sizes=[M] * N, clip=variant['clip_action'],
                                          std=std)
    # Finished getting producer

    remote_eval_path_collector = MdpPathCollector(
        eval_env
    )

    # remote_eval_path_collector = RemoteMdpPathCollector.remote(
    #     domain, seed * 10 + 1,
    #     policy_producer
    # )

    expl_path_collector = MdpPathCollector(
        expl_env,
    )

    if variant['alg'] in ['p-oac', 'g-oac', 'g-tsac', 'p-tsac'] and variant['trainer_kwargs']["counts"]:
        replay_buffer = ReplayBufferCount(
            variant['replay_buffer_size'],
            ob_space=expl_env.observation_space,
            action_space=expl_env.action_space,
            priority_sample=variant['priority_sample']
        )
    elif variant["no_resampling"] and variant['alg'] in ['p-oac', 'g-oac', 'g-tsac', 'p-tsac']:
        replay_buffer = ReplayBufferNoResampling(
            variant['replay_buffer_size'],
            ob_space=expl_env.observation_space,
            action_space=expl_env.action_space
        )
    else:
        replay_buffer = ReplayBuffer(
            variant['replay_buffer_size'],
            ob_space=expl_env.observation_space,
            action_space=expl_env.action_space
        )
    if variant['alg'] in ['ddpg']:
        trainer = DDPGTrainer(
            policy_producer,
            q_producer,
            action_space=expl_env.action_space,
            **variant['trainer_kwargs']
        )
    elif variant['alg'] in ['oac', 'sac']:
        trainer = SACTrainer(
            policy_producer,
            q_producer,
            action_space=expl_env.action_space,
            **variant['trainer_kwargs']
        )
    elif variant['alg'] == 'p-oac':
        if variant['optimistic_exp']['should_use']:
            trainer_class = ParticleTrainerOAC
            variant['trainer_kwargs']['deterministic'] = False
        else:
            trainer_class = ParticleTrainer
            variant['trainer_kwargs']['deterministic'] = not variant['stochastic']
        q_min = variant['r_min'] / (1 - variant['trainer_kwargs']['discount'])
        q_max = variant['r_max'] / (1 - variant['trainer_kwargs']['discount'])
        trainer = trainer_class(
            policy_producer,
            q_producer,
            n_estimators=n_estimators,
            delta=variant['delta'],
            q_min=q_min,
            q_max=q_max,
            action_space=expl_env.action_space,
            ensemble=variant['ensemble'],
            n_policies=variant['n_policies'],
            **variant['trainer_kwargs']
        )
    elif variant['alg'] == 'g-oac':
        q_min = variant['r_min'] / (1 - variant['trainer_kwargs']['discount'])
        q_max = variant['r_max'] / (1 - variant['trainer_kwargs']['discount'])
        trainer = GaussianTrainer(
            policy_producer,
            q_producer,
            n_estimators=n_estimators,
            delta=variant['delta'],
            q_min=q_min,
            q_max=q_max,
            action_space=expl_env.action_space,
            pac=variant['pac'],
            ensemble=variant['ensemble'],
            n_policies=variant['n_policies'],
            **variant['trainer_kwargs']
        )
    elif variant['alg'] == 'g-tsac':
        q_min = variant['r_min'] / (1 - variant['trainer_kwargs']['discount'])
        q_max = variant['r_max'] / (1 - variant['trainer_kwargs']['discount'])
        q_posterior_producer = None
        if variant['share_layers']:
            q_posterior_producer = get_q_producer(obs_dim, action_dim, hidden_sizes=[M] * N, output_size=1)
        trainer = GaussianTrainerTS(
            policy_producer,
            q_producer,
            n_estimators=n_estimators,
            delta=variant['delta'],
            q_min=q_min,
            q_max=q_max,
            action_space=expl_env.action_space,
            n_components=variant['n_components'],
            q_posterior_producer=q_posterior_producer,
            **variant['trainer_kwargs']
        )
    elif variant['alg'] == 'p-tsac':
        q_min = variant['r_min'] / (1 - variant['trainer_kwargs']['discount'])
        q_max = variant['r_max'] / (1 - variant['trainer_kwargs']['discount'])
        q_posterior_producer = None
        if variant['share_layers']:
            q_posterior_producer = get_q_producer(obs_dim, action_dim, hidden_sizes=[M] * N, output_size=1)
        trainer = ParticleTrainerTS(
            policy_producer,
            q_producer,
            n_estimators=n_estimators,
            delta=variant['delta'],
            q_min=q_min,
            q_max=q_max,
            action_space=expl_env.action_space,
            n_components=variant['n_components'],
            q_posterior_producer=q_posterior_producer,
            **variant['trainer_kwargs']
        )
    else:
        raise ValueError("Algorithm no implemented:" + variant['alg'])

    algorithm = BatchRLAlgorithm(
        trainer=trainer,
        exploration_data_collector=expl_path_collector,
        remote_eval_data_collector=remote_eval_path_collector,
        replay_buffer=replay_buffer,
        optimistic_exp_hp=variant['optimistic_exp'],
        deterministic=variant['alg'] == 'p-oac',
        **variant['algorithm_kwargs']
    )

    algorithm.to(ptu.device)

    if prev_exp_state is not None:
        expl_path_collector.restore_from_snapshot(
            prev_exp_state['exploration'])

        remote_eval_path_collector.restore_from_snapshot(
            prev_exp_state['evaluation_remote'])

        replay_buffer.restore_from_snapshot(prev_exp_state['replay_buffer'])

        trainer.restore_from_snapshot(prev_exp_state['trainer'])

        set_global_pkg_rng_state(prev_exp_state['global_pkg_rng_state'])

    start_epoch = prev_exp_state['epoch'] + \
                  1 if prev_exp_state is not None else 0

    algorithm.train(start_epoch)

def get_cmd_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--seed', type=int, default=0, help='Random seed')
    parser.add_argument('--domain', type=str, default='mountain')
    parser.add_argument('--dim', type=int, default=25)
    parser.add_argument('--pac', action="store_true")
    parser.add_argument('--ensemble', action="store_true")
    parser.add_argument('--n_policies', type=int, default=1)
    parser.add_argument('--gamma', type=float, default=0.99)
    parser.add_argument('--alg', type=str, default='oac', choices=['oac', 'p-oac', 'sac', 'g-oac', 'g-tsac', 'p-tsac',
                                                                   'ddpg'])
    parser.add_argument('--no_gpu', default=False, action='store_true')
    parser.add_argument('--base_log_dir', type=str, default='./data')
    parser.add_argument('--load_dir', type=str, default='')
    parser.add_argument('--num_layers', type=int, default=1)
    parser.add_argument('--layer_size', type=int, default=16)
    parser.add_argument('--n_estimators', type=int, default=2)
    parser.add_argument('--share_layers', action="store_true")
    parser.add_argument('--mean_update', action="store_true")
    parser.add_argument('--counts', action="store_true", help="count the samples in replay buffer")
    parser.add_argument('--log_dir', type=str, default='./data')
    parser.add_argument('--suffix', type=str, default='')
    parser.add_argument('--max_path_length', type=int, default=100)
    parser.add_argument('--replay_buffer_size', type=float, default=1e4)
    parser.add_argument('--epochs', type=int, default=200)
    parser.add_argument('--batch_size', type=int, default=32)
    parser.add_argument('--r_min', type=float, default=0.)
    parser.add_argument('--r_max', type=float, default=1.)
    parser.add_argument('--r_mellow_max', type=float, default=1.)
    parser.add_argument('--mellow_max', action="store_true")
    parser.add_argument('--priority_sample', action="store_true")
    parser.add_argument('--global_opt', action="store_true")
    parser.add_argument('--save_sampled_data', default=False, action='store_true')
    parser.add_argument('--n_components', type=int, default=1)
    parser.add_argument('--snapshot_gap', type=int, default=10)
    parser.add_argument('--save_fig', action='store_true')
    parser.add_argument('--snapshot_mode', type=str, default='last_every_gap', choices=['last_every_gap',
                                                                                'all',
                                                                                'last',
                                                                                "gap",
                                                                                'gap_and_last'])
    parser.add_argument('--difficulty', type=str, default='hard', choices=['easy',
                                                                           'medium',
                                                                           'hard',
                                                                           "harder",
                                                                           'maze',
                                                                           'maze_easy']
                        , help='only for point environment')
    parser.add_argument('--policy_lr', type=float, default=3E-4)
    parser.add_argument('--qf_lr', type=float, default=3E-4)
    parser.add_argument('--std_lr', type=float, default=3E-5)
    parser.add_argument('--sigma_noise', type=float, default=0.0)

    parser.add_argument('--std_soft_update', action="store_true")
    parser.add_argument('--clip_state', action="store_true", help='only for point environment')
    parser.add_argument('--terminal', action="store_true", help='only for point environment')
    parser.add_argument('--max_state', type=float, default=500., help='only for point environment')

    # optimistic_exp_hyper_param
    parser.add_argument('--beta_UB', type=float, default=0.0)
    parser.add_argument('--trainer_UB', action='store_true')
    parser.add_argument('--delta', type=float, default=0.95)
    parser.add_argument('--delta_oac', type=float, default=20.53)
    parser.add_argument('--deterministic_optimistic_exp', action='store_true')
    parser.add_argument('--no_resampling', action="store_true",
                        help="Samples are removed from replay buffer after being used once")

    # Training param
    parser.add_argument('--num_expl_steps_per_train_loop',
                        type=int, default=2000)
    parser.add_argument('--num_trains_per_train_loop', type=int, default=1000)
    parser.add_argument('--num_train_loops_per_epoch', type=int, default=1)
    parser.add_argument('--num_eval_steps_per_epoch', type=int, default=1000)
    parser.add_argument('--min_num_steps_before_training', type=int, default=1000)
    parser.add_argument('--clip_action', dest='clip_action', action='store_true')
    parser.add_argument('--no_clip_action', dest='clip_action', action='store_false')
    parser.set_defaults(clip_action=True)
    parser.add_argument('--entropy_tuning', dest='entropy_tuning', action='store_true')
    parser.add_argument('--no_entropy_tuning', dest='entropy_tuning', action='store_false')
    parser.set_defaults(entropy_tuning=True)
    parser.add_argument('--load_from', type=str, default='')
    parser.add_argument('--train_bias', dest='train_bias', action='store_true')
    parser.add_argument('--no_train_bias', dest='train_bias', action='store_false')
    parser.add_argument('--should_use',  action='store_true')
    parser.add_argument('--stochastic',  action='store_true')
    parser.set_defaults(train_bias=True)
    parser.add_argument('--soft_target_tau', type=float, default=5E-3)
    parser.add_argument('--ddpg', action='store_true', help='use a ddpg version of the algorithms')
    parser.add_argument('--ddpg_noisy', action='store_true', help='use noisy exploration policy')
    parser.add_argument('--std', type=float, default=0.1, help='use noisy exploration policy for ddpg')
    parser.add_argument('--use_target_policy', action='store_true', help='use a target policy in ddpg')
    parser.add_argument('--rescale_targets_around_mean', action='store_true', help='use a target policy in ddpg')

    args = parser.parse_args()

    return args

def get_log_dir(args, should_include_base_log_dir=True, should_include_seed=True, should_include_domain=True):
    start_time = time.time()
    if args.load_dir != '':
        log_dir = args.load_dir
    else:
        if args.n_policies > 1:
            el = str(args.n_policies)
        elif args.n_components > 1:
            el = str(args.n_components)
        else:
            el = ''
        log_dir = args.log_dir + '/' + args.domain + '/' + \
                  (args.difficulty + '/' if args.domain == 'point' else '') + \
                  ('terminal' + '/' if args.terminal and  args.domain == 'point' else '') + \
                  (str(args.dim) + '/' if args.domain == 'riverswim' else '') + \
                  ('global/' if args.global_opt else '') + \
                  ('ddpg/' if args.ddpg else '') + \
                  ('mean_update_' if args.mean_update else '') + \
                  ('_priority_' if args.priority_sample else '') + \
                  ('counts/' if args.counts else '') + \
                  ('/' if args.mean_update and not args.counts else '') + \
                   args.alg + ('_std' if args.std_soft_update else '') + '_' + el + '/' +\
                   args.suffix + '/' + str(start_time) + '/'

    return log_dir

if __name__ == "__main__":

    # Parameters for the experiment are either listed in variant below
    # or can be set through cmdline args and will be added or overrided
    # the corresponding attributein variant

    variant = dict(
        algorithm="SAC",
        version="normal",
        layer_size=256,
        replay_buffer_size=int(1E4),
        algorithm_kwargs=dict(
            num_eval_steps_per_epoch=5000,
            num_trains_per_train_loop=None,
            num_expl_steps_per_train_loop=None,
            min_num_steps_before_training=1000,
            max_path_length=100,
            batch_size=32,
        ),
        trainer_kwargs=dict(
            discount=0.99,
            soft_target_tau=5e-3,
            target_update_period=1,
            policy_lr=3E-4,
            qf_lr=3E-4,
            reward_scale=1,
        ),
        optimistic_exp={}
    )

    args = get_cmd_args()

    variant['log_dir'] = get_log_dir(args)
    if args.load_from != '':
        variant['log_dir'] = args.load_from

    variant['seed'] = args.seed
    variant['domain'] = args.domain
    variant['num_layers'] = args.num_layers
    variant['layer_size'] = args.layer_size
    variant['share_layers'] = args.share_layers
    variant['n_estimators'] = args.n_estimators if args.alg in ['p-oac', 'p-tsac'] else 2
    variant['replay_buffer_size'] = int(args.replay_buffer_size)
    variant['algorithm_kwargs']['num_epochs'] = domain_to_epoch(args.domain) if args.epochs <= 0 else args.epochs
    variant['algorithm_kwargs']['num_trains_per_train_loop'] = args.num_trains_per_train_loop
    variant['algorithm_kwargs']['num_expl_steps_per_train_loop'] = args.num_expl_steps_per_train_loop
    variant['algorithm_kwargs']['max_path_length'] = args.max_path_length
    variant['algorithm_kwargs']['num_eval_steps_per_epoch'] = args.num_eval_steps_per_epoch
    variant['algorithm_kwargs']['min_num_steps_before_training'] = args.min_num_steps_before_training
    variant['algorithm_kwargs']['batch_size'] = args.batch_size
    variant['algorithm_kwargs']['save_sampled_data'] = args.save_sampled_data
    variant['algorithm_kwargs']['num_train_loops_per_epoch'] = args.num_train_loops_per_epoch
    variant['algorithm_kwargs']['trainer_UB'] = args.trainer_UB

    variant['delta'] = args.delta
    variant['std'] = args.std
    variant['optimistic_exp']['should_use'] = args.beta_UB > 0 or args.delta > 0 and not args.alg in ['p-oac', 'sac',
                                                                                'g-oac', 'g-tsac',
                                                                                'p-tsac', 'ddpg']
    if not variant['optimistic_exp']['should_use']:
        variant['optimistic_exp']['should_use'] = args.should_use
    variant['optimistic_exp']['beta_UB'] = args.beta_UB if args.alg == 'oac' else 0
    variant['optimistic_exp']['delta'] = args.delta if args.alg in ['p-oac', 'oac', 'g-oac'] else 0
    variant['optimistic_exp']['share_layers'] = False
    if args.alg in ['p-oac']:
        variant['optimistic_exp']['share_layers'] = args.share_layers
    if args.should_use and args.alg in ['p-oac']:
        variant['optimistic_exp']['delta'] = args.delta_oac
    variant['optimistic_exp']['deterministic'] = args.deterministic_optimistic_exp
    if args.alg not in ['ddpg']:
        variant['trainer_kwargs']['use_automatic_entropy_tuning'] = args.entropy_tuning
    variant['trainer_kwargs']['discount'] = args.gamma
    variant['trainer_kwargs']['policy_lr'] = args.policy_lr
    variant['trainer_kwargs']['qf_lr'] = args.qf_lr
    variant['ensemble'] = args.ensemble
    variant['n_policies'] = args.n_policies if args.ensemble else 1
    variant['n_components'] = args.n_components
    variant['priority_sample'] = False
    variant['clip_action'] = args.clip_action
    variant['stochastic'] = args.stochastic
    if args.domain == 'lqg':
        variant['clip_action'] = False
    if args.alg in ['p-oac', 'g-oac', 'g-tsac', 'p-tsac']:
        variant['trainer_kwargs']['std_soft_update'] = args.std_soft_update
        variant['trainer_kwargs']['share_layers'] = args.share_layers
        variant['trainer_kwargs']['mean_update'] = args.mean_update
        variant['trainer_kwargs']['counts'] = args.counts
        variant['priority_sample'] = args.priority_sample
        variant['trainer_kwargs']['global_opt'] = args.global_opt
        if args.alg in ['p-oac', 'g-oac']:
            variant['trainer_kwargs']['r_mellow_max'] = args.r_mellow_max
            variant['trainer_kwargs']['mellow_max'] = args.mellow_max
            variant['algorithm_kwargs']['global_opt'] = args.global_opt
            variant['algorithm_kwargs']['save_fig'] = args.save_fig
            variant['trainer_kwargs']['train_bias'] = args.train_bias
            variant['trainer_kwargs']['rescale_targets_around_mean'] = args.rescale_targets_around_mean
            variant['trainer_kwargs']['use_target_policy'] = args.use_target_policy
        if args.alg in ['g-oac']:
            variant['trainer_kwargs']['std_lr'] = args.std_lr

    variant['alg'] = args.alg
    variant['dim'] = args.dim
    variant['difficulty'] = args.difficulty
    variant['max_state'] = args.max_state
    variant['clip_state'] = args.clip_state
    variant['terminal'] = args.terminal
    variant['pac'] = args.pac
    variant['no_resampling'] = args.no_resampling
    variant['r_min'] = args.r_min
    variant['r_max'] = args.r_max
    variant['sigma_noise'] = args.sigma_noise

    variant['trainer_kwargs']['soft_target_tau'] = args.soft_target_tau
    variant['algorithm_kwargs']['ddpg_noisy'] = args.ddpg_noisy
    if args.alg in ['p-oac', 'g-oac', 'g-tsac', 'p-tsac']:
        N_expl = variant['algorithm_kwargs']['num_expl_steps_per_train_loop']
        N_train = variant['algorithm_kwargs']['num_trains_per_train_loop']
        B = variant['algorithm_kwargs']['batch_size']
        N_updates = (N_train * B) / N_expl
        std_soft_update_prob = 2 / (N_updates * (N_updates + 1))
        variant['trainer_kwargs']['std_soft_update_prob'] = std_soft_update_prob
    if args.ddpg or args.alg == 'ddpg':
        variant['algorithm_kwargs']['num_trains_per_train_loop'] = 1
        variant['algorithm_kwargs']['num_expl_steps_per_train_loop'] = 4
        variant['algorithm_kwargs']['num_train_loops_per_epoch'] = args.num_expl_steps_per_train_loop // 4
        variant['trainer_kwargs']['use_target_policy'] = args.use_target_policy
        variant['algorithm_kwargs']['ddpg'] = args.ddpg
        if args.alg == 'ddpg':
            variant['algorithm_kwargs']['ddpg_noisy'] = True
        else:
            variant['algorithm_kwargs']['ddpg_noisy'] = args.ddpg_noisy

    #print("Prob %s" % variant['trainer_kwargs']['std_soft_update_prob'])

    if args.no_resampling:
        variant['algorithm_kwargs']['num_trains_per_train_loop'] = 500
        variant['algorithm_kwargs']['num_expl_steps_per_train_loop'] = 500 * args.batch_size
        variant['algorithm_kwargs']['min_num_steps_before_training'] = 4 * 500 * args.batch_size
        variant['algorithm_kwargs']['batch_size'] = args.batch_size
        variant['replay_buffer_size'] = 5 * 500 * args.batch_size
    if torch.cuda.is_available():
        gpu_id = int(args.seed % torch.cuda.device_count())
    else:
        gpu_id = None
    if not args.no_gpu:
        try:
            torch.set_default_tensor_type('torch.cuda.FloatTensor')
        except:
            pass
    run_experiment_here(experiment, variant,
                        seed=args.seed,
                        use_gpu=not args.no_gpu and torch.cuda.is_available(),
                        gpu_id=gpu_id,

                        # Save the params every snapshot_gap and override previously saved result
                        snapshot_gap=args.snapshot_gap,
                        snapshot_mode=args.snapshot_mode,

                        log_dir=variant['log_dir']
                        )


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

import warnings
import numpy as np
import tensorflow as tf

from stable_baselines.common.distributions import make_proba_dist_type, spaces, \
    DiagGaussianProbabilityDistributionType
from stable_baselines.common.policies import FeedForwardPolicy, nature_cnn, mlp_extractor, linear

def make_proba_dist_type(ac_space):
    """
    return an instance of ProbabilityDistributionType for the correct type of action space

    :param ac_space: (Gym Space) the input action space
    :return: (ProbabilityDistributionType) the appropriate instance of a ProbabilityDistributionType
    """
    if isinstance(ac_space, spaces.Box):
        assert len(ac_space.shape) == 1, "Error: the action space must be a vector"
        return DiagGaussianFixedVarProbabilityDistributionType(ac_space.shape[0])
    else:
        return make_proba_dist_type(ac_space)

class DiagGaussianFixedVarProbabilityDistributionType(DiagGaussianProbabilityDistributionType):
    def __init__(self, size):
        super(DiagGaussianFixedVarProbabilityDistributionType, self).__init__(size)
        return

    def proba_distribution_from_latent(self, pi_latent_vector, vf_latent_vector,
                                        pi_init_scale=1.0, pi_init_bias=0.0, pi_init_std=1.0,
                                        vf_init_scale=1.0, vf_init_bias=0.0):
        mean = linear(pi_latent_vector, 'pi', self.size, init_scale=pi_init_scale, init_bias=pi_init_bias)
        logstd = tf.get_variable(name='pi/logstd', shape=[1, self.size], initializer=tf.constant_initializer(np.log(pi_init_std)), trainable=False)
        pdparam = tf.concat([mean, mean * 0.0 + logstd], axis=1)
        q_values = linear(vf_latent_vector, 'q', self.size, init_scale=vf_init_scale, init_bias=vf_init_bias)
        return self.proba_distribution_from_flat(pdparam), mean, q_values

class ImitationPolicy(FeedForwardPolicy):
    """
    Policy object that implements actor critic, using a feed forward neural network.

    :param sess: (TensorFlow session) The current TensorFlow session
    :param ob_space: (Gym Space) The observation space of the environment
    :param ac_space: (Gym Space) The action space of the environment
    :param n_env: (int) The number of environments to run
    :param n_steps: (int) The number of steps to run for each environment
    :param n_batch: (int) The number of batch to run (n_envs * n_steps)
    :param reuse: (bool) If the policy is reusable or not
    :param layers: ([int]) (deprecated, use net_arch instead) The size of the Neural network for the policy
        (if None, default to [64, 64])
    :param net_arch: (list) Specification of the actor-critic policy network architecture (see mlp_extractor
        documentation for details).
    :param act_fun: (tf.func) the activation function to use in the neural network.
    :param cnn_extractor: (function (TensorFlow Tensor, ``**kwargs``): (TensorFlow Tensor)) the CNN feature extraction
    :param feature_extraction: (str) The feature extraction type ("cnn" or "mlp")
    :param kwargs: (dict) Extra keyword arguments for the nature CNN feature extraction
    """

    def __init__(self, sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=False, layers=None, net_arch=None,
                 act_fun=tf.tanh, cnn_extractor=nature_cnn, feature_extraction="mlp", **kwargs):
        super(FeedForwardPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=reuse,
                                                scale=(feature_extraction == "cnn"))

        self._pdtype = make_proba_dist_type(ac_space)

        self._kwargs_check(feature_extraction, kwargs)

        if layers is not None:
            warnings.warn("Usage of the `layers` parameter is deprecated! Use net_arch instead "
                          "(it has a different semantics though).", DeprecationWarning)
            if net_arch is not None:
                warnings.warn("The new `net_arch` parameter overrides the deprecated `layers` parameter!",
                              DeprecationWarning)

        if net_arch is None:
            if layers is None:
                layers = [64, 64]
            net_arch = [dict(vf=layers, pi=layers)]

        with tf.variable_scope("model", reuse=reuse):
            if feature_extraction == "cnn":
                pi_latent = vf_latent = cnn_extractor(self.processed_obs, **kwargs)
            else:
                pi_latent, vf_latent = mlp_extractor(tf.layers.flatten(self.processed_obs), net_arch, act_fun)

            self._value_fn = linear(vf_latent, 'vf', 1)

            self._proba_distribution, self._policy, self.q_value = \
                self.pdtype.proba_distribution_from_latent(pi_latent, vf_latent,
                                                           pi_init_scale=1.0, pi_init_bias=0.0, pi_init_std=0.125,
                                                           vf_init_scale=1.0, vf_init_bias=0.0)

        self._setup_init()
        return


""" This file contains the EfficientNet facial emotion classifier """

from typing import Dict

import numpy as np
import tensorflow as tf

from src.classification.image.image_emotion_classifier import (
    ImageEmotionClassifier,
)
from src.data.data_reader import Set

class MultiTaskEfficientNetB2Classifier(ImageEmotionClassifier):
    """
    Class that implements an efficient net emotion classifier.
    """

    def __init__(self, parameters: Dict = None):
        """
        Initialize the EfficientNet emotion classifier

        :param name: The name for the classifier
        :param parameters: Some configuration parameters for the classifier
        """
        super().__init__("efficientnet", parameters)
        tf.get_logger().setLevel("ERROR")
        self.model = None

    def initialize_model(self, parameters: Dict) -> None:
        """
        Initializes a new and pretrained version of the EfficientNetB2 model
        """
        extra_layer = parameters.get("extra_layer", None)
        input = tf.keras.layers.Input(
            shape=(48, 48, 3), dtype=tf.float32, name="image"
        )
        input = tf.keras.applications.efficientnet.preprocess_input(input)

        model = tf.keras.applications.EfficientNetB2(
            include_top=False,
            weights="imagenet",
            input_tensor=input,
            input_shape=(48, 48, 3),
        )
        for layer in model.layers[: parameters.get("frozen_layers", -10)]:
            layer.trainable = False
        output = model(input)
        flat = tf.keras.layers.Flatten()(output)
        if extra_layer:
            flat = tf.keras.layers.Dense(
                extra_layer,
                activation="relu",
                kernel_regularizer=tf.keras.regularizers.L1L2(
                    l1=0.001, l2=0.001
                ),
            )(flat)
        top = tf.keras.layers.Dense(
            7, activation="softmax", name="classifier"
        )(flat)
        self.model = tf.keras.Model(input, top)

    def train(self, parameters: Dict = None, **kwargs) -> None:
        """
        Training method for EfficientNet model

        :param parameters: Parameter dictionary used for training
        :param kwargs: Additional kwargs parameters
        """
        parameters = self.init_parameters(parameters, **kwargs)
        epochs = parameters.get("epochs", 20)

        if not self.model:
            self.initialize_model(parameters)
        self.prepare_training(parameters)
        self.model.compile(
            optimizer=self.optimizer, loss=self.loss, metrics=self.metrics
        )
        self.prepare_data(parameters)

        _ = self.model.fit(
            x=self.train_data,
            validation_data=self.val_data,
            epochs=epochs,
            callbacks=[self.callback],
            class_weight=self.class_weights,
        )
        self.is_trained = True

    def load(self, parameters: Dict = None, **kwargs) -> None:
        """
        Loading method that loads a previously trained model from disk.

        :param parameters: Parameters required for loading the model
        :param kwargs: Additional kwargs parameters
        """
        parameters = self.init_parameters(parameters, **kwargs)
        save_path = parameters.get("save_path", "models/image/efficientnet")
        self.model = tf.keras.models.load_model(save_path)

    def save(self, parameters: Dict = None, **kwargs) -> None:
        """
        Saving method that saves a previously trained model on disk.

        :param parameters: Parameters required for storing the model
        :param kwargs: Additional kwargs parameters
        """
        if not self.is_trained:
            raise RuntimeError(
                "Model needs to be trained in order to save it!"
            )
        parameters = self.init_parameters(parameters, **kwargs)
        save_path = parameters.get("save_path", "models/image/efficientnet")
        self.model.save(save_path, include_optimizer=False)

    def classify(self, parameters: Dict = None, **kwargs) -> np.array:
        """
        Classification method used to classify emotions from images

        :param parameters: Parameter dictionary used for classification
        :param kwargs: Additional kwargs parameters
        :return: An array with predicted emotion indices
        """
        parameters = self.init_parameters(parameters, **kwargs)
        which_set = parameters.get("which_set", Set.TEST)
        batch_size = parameters.get("batch_size", 64)
        dataset = self.data_reader.get_emotion_data(
            self.emotions, which_set, batch_size
        ).map(lambda x, y: (tf.image.grayscale_to_rgb(x), y))

        if not self.model:
            raise RuntimeError(
                "Please load or train the model before inference!"
            )
        results = self.model.predict(dataset)
        return np.argmax(results, axis=1)

if __name__ == "__main__":  # pragma: no cover
    classifier = MultiTaskEfficientNetB2Classifier()
    classifier.train(
        {
            "epochs": 50,
            "batch_size": 256,
            "patience": 15,
            "learning_rate": 0.003,
            "frozen_layers": 0,
            "extra_layer": 2048,
            "augment": True,
            "weighted": True,
        }
    )
    classifier.save()
    classifier.load()
    emotions = classifier.classify()
    labels = classifier.data_reader.get_labels(Set.TEST)
    print(f"Labels Shape: {labels.shape}")
    print(f"Emotions Shape: {emotions.shape}")
    print(f"Accuracy: {np.sum(emotions == labels) / labels.shape[0]}")


# Load LSTM network and generate text
import sys
import numpy

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Dropout
from tensorflow.keras.layers import LSTM
from tensorflow.keras.callbacks import ModelCheckpoint
from tensorflow.keras import utils

# load ascii text and covert to lowercase
filename = "(ASMR) Dragonology-FQ_TfYAbQMM.en.txt"
raw_text = open(filename, 'r', encoding='utf-8').read()
raw_text = raw_text.lower()

# create mapping of unique chars to integers, and a reverse mapping
chars = sorted(list(set(raw_text)))
char_to_int = dict((c, i) for i, c in enumerate(chars))
int_to_char = dict((i, c) for i, c in enumerate(chars))

# summarize the loaded data
n_chars = len(raw_text)
n_vocab = len(chars)
print ("Total Characters: ", n_chars)
print ("Total Vocab: ", n_vocab)

# prepare the dataset of input to output pairs encoded as integers
seq_length = 100
dataX = []
dataY = []
for i in range(0, n_chars - seq_length, 1):
	seq_in = raw_text[i:i + seq_length]
	seq_out = raw_text[i + seq_length]
	dataX.append([char_to_int[char] for char in seq_in])
	dataY.append(char_to_int[seq_out])
n_patterns = len(dataX)
print ("Total Patterns: ", n_patterns)

# reshape X to be [samples, time steps, features]
X = numpy.reshape(dataX, (n_patterns, seq_length, 1))

# normalize
X = X / float(n_vocab)

# one hot encode the output variable
y = utils.to_categorical(dataY)

# define the LSTM model
model = Sequential()
model.add(LSTM(256, input_shape=(X.shape[1], X.shape[2])))
model.add(Dropout(0.2))
model.add(Dense(y.shape[1], activation='softmax'))

# load the network weights
# absolute path or filepath from directory
filename2 = "weights-improvement-49-0.3137-bigger.hdf5"
model.load_weights(filename2)
model.compile(loss='categorical_crossentropy', optimizer='adam')

# pick a random seed
start = numpy.random.randint(0, len(dataX)-1)
pattern = dataX[start]
print ("Seed:")
print ("\"", ''.join([int_to_char[value] for value in pattern]), "\"")

# generate characters
for i in range(250):
	x = numpy.reshape(pattern, (1, len(pattern), 1))
	x = x / float(n_vocab)
	prediction = model.predict(x, verbose=0)
	index = numpy.argmax(prediction)
	result = int_to_char[index]
	seq_in = [int_to_char[value] for value in pattern]
	sys.stdout.write(result)
	pattern.append(index)
	pattern = pattern[1:len(pattern)]
print ("\nDone.")


import os
import argparse
import numpy as np
import pandas as pd
from keras.models import load_model
from keras import backend as K
from sklearn.neighbors import NearestNeighbors

import utils
preprocess import preprocess

parser = argparse.ArgumentParser(description='Malconv-keras classifier training')
parser.add_argument('--save_path', type=str, default='../saved/adversarial_samples',    help="Directory for saving adv samples")
parser.add_argument('--model_path', type=str, default='../saved/malconv.h5',            help='Path to target model')
parser.add_argument('--log_path', type=str, default='../saved/adversarial_log.csv',     help="[csv file] Adv sample generation log")
parser.add_argument('--pad_percent', type=float, default=0.1,                        help="padding percentage to origin file")
parser.add_argument('--targetClass', type=int, default=1,                            help="target class. default:1")
parser.add_argument('--step_size', type=float, default=0.01,                         help="optimiztion step size for fgsm, senitive")
parser.add_argument('--limit', type=float, default = 0.,                             help="limit gpu memory percentage")
parser.add_argument('csv', type=str,                                                 help="[csv file] Filenames")

def fgsm(model, inp, pad_idx, pad_len, e, step_size=0.001, target_class=1):
    adv = inp.copy()
    loss = K.mean(model.output[:, target_class])
    grads = K.gradients(loss, model.layers[1].output)[0]
    grads /= (K.sqrt(K.mean(K.square(grads))) + 1e-8)

    mask = np.zeros(model.layers[1].output.shape[1:]) # embedding layer output shape
    mask[pad_idx:pad_idx+pad_len] = 1
    grads *= K.constant(mask)

    iterate = K.function([model.layers[1].output], [loss, grads])
    g = 0.
    step = int(1/step_size)*10
    for _ in range(step):
        loss_value, grads_value = iterate([adv])
        grads_value *= step_size
        g += grads_value
        adv += grads_value
        #print (e, loss_value, grads_value.mean(), end='\r')
        if loss_value >= 0.9:
            break

    return adv, g, loss_value

def gen_adv_samples(model, fn_list, pad_percent=0.1, step_size=0.001, target_class=1):

    ###   search for nearest neighbor in embedding space ###
    def emb_search(org, adv, pad_idx, pad_len, neigh):
        out = org.copy()
        for idx in range(pad_idx, pad_idx+pad_len):
            target = adv[idx].reshape(1, -1)
            best_idx = neigh.kneighbors(target, 1, False)[0][0]
            out[0][idx] = best_idx
        return out

    max_len = int(model.input.shape[1])
    emb_layer = model.layers[1]
    emb_weight = emb_layer.get_weights()[0]
    inp2emb = K.function([model.input]+ [K.learning_phase()], [emb_layer.output]) # [function] Map sequence to embedding

    # Build neighbor searches
    neigh = NearestNeighbors(1)
    neigh.fit(emb_weight)

    log = utils.logger()
    adv_samples = []

    for e, fn in enumerate(fn_list):

        ###   run one file at a time due to different padding length, [slow]
        inp, len_list = preprocess([fn], max_len)
        inp_emb = np.squeeze(np.array(inp2emb([inp, False])), 0)

        pad_idx = len_list[0]
        pad_len = max(min(int(len_list[0]*pad_percent), max_len-pad_idx), 0)
        org_pred = np.argmax(model.predict(inp)[0])    ### origianl score, 0 -> malicious, 1 -> benign
        loss, pred = float('nan'), float('nan')

        if pad_len > 0:

            if np.argmax(org_pred) != target_class:
                adv_emb, gradient, loss = fgsm(model, inp_emb, pad_idx, pad_len, e, step_size, target_class)
                adv = emb_search(inp, adv_emb[0], pad_idx, pad_len, neigh)
                pred = np.argmax(model.predict(adv)[0]).astype(float)
                final_adv = adv[0][:pad_idx+pad_len]

            else: # use origin file
                final_adv = inp[0][:pad_idx]

        log.write(fn, org_pred, pad_idx, pad_len, loss, pred)

        # sequence to bytes
        bin_adv = bytes(list(final_adv))
        adv_samples.append(bin_adv)

    return adv_samples, log

if __name__ == '__main__':
    args = parser.parse_args()

    # limit gpu memory
    if args.limit > 0:
        utils.limit_gpu_memory(args.limit)

    df = pd.read_csv(args.csv, header=None)
    fn_list = df[0].values
    model = load_model(args.model_path)

    adv_samples, log = gen_adv_samples(model, fn_list, args.pad_percent, args.step_size, args.targetClass)

    # write to file
    log.save(args.log_path)
    for fn, adv in zip(fn_list, adv_samples):
        _fn = fn.split('/')[-1]
        dst = os.path.join(args.save_path, _fn)
        with open(dst, 'wb') as f:
            f.write(adv)


"""

================================================
Pandas querying and metadata with Epochs objects
================================================

Demonstrating pandas-style string querying with Epochs metadata.
For related uses of :class:`mne.Epochs`, see the starting tutorial
:ref:`sphx_glr_auto_tutorials_plot_object_epochs.py`.

Sometimes you may have a complex trial structure that cannot be easily
summarized as a set of unique integers. In this case, it may be useful to use
the ``metadata`` attribute of :class:`mne.Epochs` objects. This must be a
:class:`pandas.DataFrame` where each row corresponds to an epoch, and each
column corresponds to a metadata attribute of each epoch. Columns must
contain either strings, ints, or floats.

In this dataset, subjects were presented with individual words
on a screen, and the EEG activity in response to each word was recorded.
We know which word was displayed in each epoch, as well as
extra information about the word (e.g., word frequency).

Loading the data
----------------
First we'll load the data. If metadata exists for an :class:`mne.Epochs`
fif file, it will automatically be loaded in the ``metadata`` attribute.
"""

# Authors: Chris Holdgraf <choldgraf@gmail.com>
#          Jona Sassenhagen <jona.sassenhagen@gmail.com>
#          Eric Larson <larson.eric.d@gmail.com>

# License: BSD (3-clause)

import mne
import numpy as np
import matplotlib.pyplot as plt

# Load the data from the internet
path = mne.datasets.kiloword.data_path() + '/kword_metadata-epo.fif'
epochs = mne.read_epochs(path)

# The metadata exists as a Pandas DataFrame
print(epochs.metadata.head(10))

###############################################################################
# We can use this metadata attribute to select subsets of Epochs. This
# uses the Pandas :meth:`pandas.DataFrame.query` method under the hood.
# Any valid query string will work. Below we'll make two plots to compare
# between them:

av1 = epochs['Concreteness < 5 and WordFrequency < 2'].average()
av2 = epochs['Concreteness > 5 and WordFrequency > 2'].average()

joint_kwargs = dict(ts_args=dict(time_unit='s'),
                    topomap_args=dict(time_unit='s'))
av1.plot_joint(show=False, **joint_kwargs)
av2.plot_joint(show=False, **joint_kwargs)

###############################################################################
# Next we'll choose a subset of words to keep.
words = ['film', 'cent', 'shot', 'cold', 'main']
epochs['WORD in {}'.format(words)].plot_image(show=False)

###############################################################################
# Note that traditional epochs sub-selection still works. The traditional
# MNE methods for selecting epochs will supersede the rich metadata querying.
epochs['cent'].average().plot(show=False, time_unit='s')

###############################################################################
# Below we'll show a more involved example that leverages the metadata
# of each epoch. We'll create a new column in our metadata object and use
# it to generate averages for many subsets of trials.

# Create two new metadata columns
metadata = epochs.metadata
is_concrete = metadata["Concreteness"] > metadata["Concreteness"].median()
metadata["is_concrete"] = np.where(is_concrete, 'Concrete', 'Abstract')
is_long = metadata["NumberOfLetters"] > 5
metadata["is_long"] = np.where(is_long, 'Long', 'Short')
epochs.metadata = metadata

###############################################################################
# Now we can quickly extract (and plot) subsets of the data. For example, to
# look at words split by word length and concreteness:

query = "is_long == '{0}' & is_concrete == '{1}'"
evokeds = dict()
for concreteness in ("Concrete", "Abstract"):
    for length in ("Long", "Short"):
        subset = epochs[query.format(length, concreteness)]
        evokeds["/".join((concreteness, length))] = list(subset.iter_evoked())

# For the actual visualisation, we store a number of shared parameters.
style_plot = dict(
    colors={"Long": "Crimson", "Short": "Cornflowerblue"},
    linestyles={"Concrete": "-", "Abstract": ":"},
    split_legend=True,
    ci=.68,
    show_sensors='lower right',
    show_legend='lower left',
    truncate_yaxis="max_ticks",
    picks=epochs.ch_names.index("Pz"),
)

fig, ax = plt.subplots(figsize=(6, 4))
mne.viz.plot_compare_evokeds(evokeds, axes=ax, **style_plot)
plt.show()

###############################################################################
# To compare words which are 4, 5, 6, 7 or 8 letters long:

letters = epochs.metadata["NumberOfLetters"].unique().astype(int).astype(str)

evokeds = dict()
for n_letters in letters:
    evokeds[n_letters] = epochs["NumberOfLetters == " + n_letters].average()

style_plot["colors"] = {n_letters: int(n_letters)
                        for n_letters in letters}
style_plot["cmap"] = ("# of Letters", "viridis_r")
del style_plot['linestyles']

fig, ax = plt.subplots(figsize=(6, 4))
mne.viz.plot_compare_evokeds(evokeds, axes=ax, **style_plot)
plt.show()

###############################################################################
# And finally, for the interaction between concreteness and continuous length
# in letters:
evokeds = dict()
query = "is_concrete == '{0}' & NumberOfLetters == {1}"
for concreteness in ("Concrete", "Abstract"):
    for n_letters in letters:
        subset = epochs[query.format(concreteness, n_letters)]
        evokeds["/".join((concreteness, n_letters))] = subset.average()

style_plot["linestyles"] = {"Concrete": "-", "Abstract": ":"}

fig, ax = plt.subplots(figsize=(6, 4))
mne.viz.plot_compare_evokeds(evokeds, axes=ax, **style_plot)
plt.show()

###############################################################################
# .. note::
#
#    Creating an :class:`mne.Epochs` object with metadata is done by passing
#    a :class:`pandas.DataFrame` to the ``metadata`` kwarg as follows:

data = epochs.get_data()
metadata = epochs.metadata.copy()
epochs_new = mne.EpochsArray(data, epochs.info, metadata=metadata)