import cv2
import math
import numpy as np
import os
import os.path as osp
import random
import time
import torch
from basicsr.data.degradations import circular_lowpass_kernel, random_mixed_kernels
from basicsr.data.transforms import augment
from basicsr.utils import FileClient, get_root_logger, imfrombytes, img2tensor
from basicsr.utils.registry import DATASET_REGISTRY
from torch.utils import data as data

@DATASET_REGISTRY.register()
class RealESRGANDataset(data.Dataset):
    """Dataset used for Real-ESRGAN model:
    Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.

    It loads gt (Ground-Truth) images, and augments them.
    It also generates blur kernels and sinc kernels for generating low-quality images.
    Note that the low-quality images are processed in tensors on GPUS for faster processing.

    Args:
        opt (dict): Config for train datasets. It contains the following keys:
            dataroot_gt (str): Data root path for gt.
            meta_info (str): Path for meta information file.
            io_backend (dict): IO backend type and other kwarg.
            use_hflip (bool): Use horizontal flips.
            use_rot (bool): Use rotation (use vertical flip and transposing h and w for implementation).
            Please see more options in the codes.
    """

    def __init__(self, opt):
        super(RealESRGANDataset, self).__init__()
        self.opt = opt
        self.file_client = None
        self.io_backend_opt = opt['io_backend']
        self.gt_folder = opt['dataroot_gt']

        # file client (lmdb io backend)
        if self.io_backend_opt['type'] == 'lmdb':
            self.io_backend_opt['db_paths'] = [self.gt_folder]
            self.io_backend_opt['client_keys'] = ['gt']
            if not self.gt_folder.endswith('.lmdb'):
                raise ValueError(f"'dataroot_gt' should end with '.lmdb', but received {self.gt_folder}")
            with open(osp.join(self.gt_folder, 'meta_info.txt')) as fin:
                self.paths = [line.split('.')[0] for line in fin]
        else:
            # disk backend with meta_info
            # Each line in the meta_info describes the relative path to an image
            with open(self.opt['meta_info']) as fin:
                paths = [line.strip().split(' ')[0] for line in fin]
                self.paths = [os.path.join(self.gt_folder, v) for v in paths]

        # blur settings for the first degradation
        self.blur_kernel_size = opt['blur_kernel_size']
        self.kernel_list = opt['kernel_list']
        self.kernel_prob = opt['kernel_prob']  # a list for each kernel probability
        self.blur_sigma = opt['blur_sigma']
        self.betag_range = opt['betag_range']  # betag used in generalized Gaussian blur kernels
        self.betap_range = opt['betap_range']  # betap used in plateau blur kernels
        self.sinc_prob = opt['sinc_prob']  # the probability for sinc filters

        # blur settings for the second degradation
        self.blur_kernel_size2 = opt['blur_kernel_size2']
        self.kernel_list2 = opt['kernel_list2']
        self.kernel_prob2 = opt['kernel_prob2']
        self.blur_sigma2 = opt['blur_sigma2']
        self.betag_range2 = opt['betag_range2']
        self.betap_range2 = opt['betap_range2']
        self.sinc_prob2 = opt['sinc_prob2']

        # a final sinc filter
        self.final_sinc_prob = opt['final_sinc_prob']

        self.kernel_range = [2 * v + 1 for v in range(3, 11)]  # kernel size ranges from 7 to 21
        # TODO: kernel range is now hard-coded, should be in the configure file
        self.pulse_tensor = torch.zeros(21, 21).float()  # convolving with pulse tensor brings no blurry effect
        self.pulse_tensor[10, 10] = 1

    def __getitem__(self, index):
        if self.file_client is None:
            self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)

        # -------------------------------- Load gt images -------------------------------- #
        # Shape: (h, w, c); channel order: BGR; image range: [0, 1], float32.
        gt_path = self.paths[index]
        # avoid errors caused by high latency in reading files
        retry = 3
        while retry > 0:
            try:
                img_bytes = self.file_client.get(gt_path, 'gt')
                img_gt = imfrombytes(img_bytes, float32=True)
            except (IOError, OSError) as e:
                logger = get_root_logger()
                logger.warn(f'File client error: {e}, remaining retry times: {retry - 1}')
                # change another file to read
                index = random.randint(0, self.__len__())
                gt_path = self.paths[index]
                time.sleep(1)  # sleep 1s for occasional server congestion
            except (AttributeError) as e:
                logger = get_root_logger()
                logger.warn(f'Attribute error: {e} for file {gt_path}, remaining retry times: {retry - 1}')
                # change another file to read
                index = random.randint(0, self.__len__())
                gt_path = self.paths[index]
                time.sleep(1)  # sleep 1s for occasional server congestion
            else:
                break
            finally:
                retry -= 1

        # -------------------- Do augmentation for training: flip, rotation -------------------- #
        img_gt = augment(img_gt, self.opt['use_hflip'], self.opt['use_rot'])

        # crop or pad to 400
        # TODO: 400 is hard-coded. You may change it accordingly
        h, w = img_gt.shape[0:2]
        crop_pad_size = 400
        # pad
        if h < crop_pad_size or w < crop_pad_size:
            pad_h = max(0, crop_pad_size - h)
            pad_w = max(0, crop_pad_size - w)
            img_gt = cv2.copyMakeBorder(img_gt, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT_101)
        # crop
        if img_gt.shape[0] > crop_pad_size or img_gt.shape[1] > crop_pad_size:
            h, w = img_gt.shape[0:2]
            # randomly choose top and left coordinates
            top = random.randint(0, h - crop_pad_size)
            left = random.randint(0, w - crop_pad_size)
            img_gt = img_gt[top:top + crop_pad_size, left:left + crop_pad_size, ...]

        # ------------------------ Generate kernels (used in the first degradation) ------------------------ #
        kernel_size = random.choice(self.kernel_range)
        if np.random.uniform() < self.opt['sinc_prob']:
            # this sinc filter setting is for kernels ranging from [7, 21]
            if kernel_size < 13:
                omega_c = np.random.uniform(np.pi / 3, np.pi)
            else:
                omega_c = np.random.uniform(np.pi / 5, np.pi)
            kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False)
        else:
            kernel = random_mixed_kernels(
                self.kernel_list,
                self.kernel_prob,
                kernel_size,
                self.blur_sigma,
                self.blur_sigma, [-math.pi, math.pi],
                self.betag_range,
                self.betap_range,
                noise_range=None)
        # pad kernel
        pad_size = (21 - kernel_size) // 2
        kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size)))

        # ------------------------ Generate kernels (used in the second degradation) ------------------------ #
        kernel_size = random.choice(self.kernel_range)
        if np.random.uniform() < self.opt['sinc_prob2']:
            if kernel_size < 13:
                omega_c = np.random.uniform(np.pi / 3, np.pi)
            else:
                omega_c = np.random.uniform(np.pi / 5, np.pi)
            kernel2 = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False)
        else:
            kernel2 = random_mixed_kernels(
                self.kernel_list2,
                self.kernel_prob2,
                kernel_size,
                self.blur_sigma2,
                self.blur_sigma2, [-math.pi, math.pi],
                self.betag_range2,
                self.betap_range2,
                noise_range=None)

        # pad kernel
        pad_size = (21 - kernel_size) // 2
        kernel2 = np.pad(kernel2, ((pad_size, pad_size), (pad_size, pad_size)))

        # ------------------------------------- the final sinc kernel ------------------------------------- #
        if np.random.uniform() < self.opt['final_sinc_prob']:
            kernel_size = random.choice(self.kernel_range)
            omega_c = np.random.uniform(np.pi / 3, np.pi)
            sinc_kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=21)
            sinc_kernel = torch.FloatTensor(sinc_kernel)
        else:
            sinc_kernel = self.pulse_tensor

        # BGR to RGB, HWC to CHW, numpy to tensor
        img_gt = img2tensor([img_gt], bgr2rgb=True, float32=True)[0]
        kernel = torch.FloatTensor(kernel)
        kernel2 = torch.FloatTensor(kernel2)

        return_d = {'gt': img_gt, 'kernel1': kernel, 'kernel2': kernel2, 'sinc_kernel': sinc_kernel, 'gt_path': gt_path}
        return return_d

    def __len__(self):
        return len(self.paths)


import os
import cv2
import numpy as np
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import img_to_array
import datetime
from src.utils.GenderAndAgePredictor import GenderAndAgePredictor
from src.utils.Loader import Loader

if __name__ == '__main__':
    loader = Loader()
    images = loader.load_images_from_folder("input")

    predictor = GenderAndAgePredictor(
        gender_model_path=os.path.abspath("models/gender_model.model"),
        age_model_path=os.path.abspath("models/age_model.model")
    )

    prototxt = os.path.sep.join([os.path.abspath("models"), "deploy.prototxt"])
    weights = os.path.sep.join([os.path.abspath("models"), "res10_300x300_ssd_iter_140000.caffemodel"])
    net = cv2.dnn.readNet(prototxt, weights)

    MaskModel = load_model(os.path.abspath("models/mask_detector.model"))

    for img in images:
        orig = img.copy()
        (h, w) = img.shape[:2]

        blob = cv2.dnn.blobFromImage(img, 1.0, (300, 300),
                                     (104.0, 177.0, 123.0))

        net.setInput(blob)
        detections = net.forward()

        confidence = detections[0, 0, 0, 2]
        if confidence > 0.5:
            dt_string = datetime.datetime.now().strftime("%d-%m-%Y_%H-%M-%S")
            box = detections[0, 0, 0, 3:7] * np.array([w, h, w, h])
            (startX, startY, endX, endY) = box.astype("int")
            (startX, startY) = (max(0, startX), max(0, startY))
            (endX, endY) = (min(w - 1, endX), min(h - 1, endY))

            face = img[startY:endY, startX:endX]
            try:
                face_for_mask = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
            except BaseException as e:
                print(e)
                continue

            face_for_mask = cv2.resize(face_for_mask, (224, 224))
            face_for_mask = img_to_array(face_for_mask)
            face_for_mask = preprocess_input(face_for_mask)
            face_for_mask = np.expand_dims(face_for_mask, axis=0)

            (mask, withoutMask) = MaskModel.predict(face_for_mask)[0]
            label = "Mask" if mask > withoutMask else "Intruder"
            color = (0, 255, 0) if label == "Mask" else (0, 0, 255)

            if mask < withoutMask:
                face_for_age_gender = cv2.resize(face, (48, 48))
                face_for_age_gender = cv2.cvtColor(face_for_age_gender, cv2.COLOR_BGR2GRAY)
                data = ''
                for i in range(48):
                    for j in range(48):
                        data += " " + str(face_for_age_gender[i][j])
                x = data[1:]
                img_arr = np.array(x.split(), dtype="float32") / 255
                X = img_arr.reshape(1, 48, 48, 1)
                gender, age = predictor.predict_gender_and_age(X)
                label += f"Gender: {gender} Age: {age}"
                label = "{}: {:.2f}%".format(label, max(mask, withoutMask) * 100)
                cv2.putText(img, label, (startX, startY - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.45, color, 3)
                cv2.rectangle(img, (startX, startY), (endX, endY), color, 2)
                cv2.imwrite(os.path.join(os.path.abspath("intruder"), f"intruder_{gender}_{int(age)}.jpg"), img)
            else:
                cv2.putText(img, label, (startX, startY - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.45, color, 3)
                cv2.rectangle(img, (startX, startY), (endX, endY), color, 2)
                cv2.imwrite(os.path.join(os.path.abspath("in_mask"), f"in_mask_{dt_string}.jpg"), img)


import os
import numpy as np

import logging

from app_globals import *
from alad_support import *
from r_support import matrix, cbind

from forest_aad_detector import *
from forest_aad_support import prepare_forest_aad_debug_args
from results_support import write_sequential_results_to_csv
from data_stream import *

"""
To debug:
    pythonw pyalad/forest_aad_stream.py
"""

logger = logging.getLogger(__name__)

class StreamingAnomalyDetector(object):
    """
    Attributes:
        model: AadForest
        stream: DataStream
        max_buffer: int
            Determines the window size
        buffer_instances_x: list
    """
    def __init__(self, stream, model, labeled_x=None, labeled_y=None,
                 unlabeled_x=None, unlabeled_y=None, opts=None, max_buffer=512):
        self.model = model
        self.stream = stream
        self.max_buffer = max_buffer
        self.opts = opts

        self.buffer_x = None
        self.buffer_y = None

        self.unlabeled_x = unlabeled_x
        self.unlabeled_y = unlabeled_y

        self.labeled_x = labeled_x
        self.labeled_y = labeled_y

        self.qstate = None

    def reset_buffer(self):
        self.buffer_x = None
        self.buffer_y = None

    def add_buffer_xy(self, x, y):
        if self.buffer_x is None:
            self.buffer_x = x
        else:
            self.buffer_x = rbind(self.buffer_x, x)

        if self.buffer_y is None:
            self.buffer_y = y
        else:
            if y is not None:
                self.buffer_y = append(self.buffer_y, y)

    def move_buffer_to_unlabeled(self):
        self.unlabeled_x = self.buffer_x
        self.unlabeled_y = self.buffer_y
        self.reset_buffer()

    def get_num_instances(self):
        """Returns the total number of labeled and unlabeled instances that will be used for weight inference"""
        n = 0
        if self.unlabeled_x is not None:
            n += nrow(self.unlabeled_x)
        if self.labeled_x is not None:
            # logger.debug("labeled_x: %s" % str(self.labeled_x.shape))
            n += nrow(self.labeled_x)
        return n

    def init_query_state(self, opts):
        n = self.get_num_instances()
        bt = get_budget_topK(n, opts)
        self.qstate = Query.get_initial_query_state(opts.qtype, opts=opts, qrank=bt.topK,
                                                    a=1., b=1., budget=bt.budget)

    def get_next_from_stream(self, n=0):
        if n == 0:
            n = self.max_buffer
        x, y = self.stream.read_next_from_stream(n)

        if x is None:
            return x, y

        if False:
            if self.buffer_x is not None:
                logger.debug("buffer shape: %s" % str(self.buffer_x.shape))
            logger.debug("x.shape: %s" % str(x.shape))

        self.add_buffer_xy(x, y)

        self.model.add_samples(x, current=False)

        return x, y

    def update_model_from_buffer(self):
        self.model.update_model_from_stream_buffer()

    def get_next_transformed(self, n=1):
        x, y = self.get_next_from_stream(n)
        if x is None:
            return x, y
        x_new = self.model.transform_to_region_features(x, dense=False)
        return x_new, y

    def stream_buffer_empty(self):
        return self.stream.empty()

    def get_anomaly_scores(self, x):
        x_new = self.model.transform_to_region_features(x, dense=False)
        scores = self.model.get_score(x_new)
        return scores

    def setup_data_for_feedback(self):
        """
        Prepares the input matrices/data structures for weight update. The format
        is such that the top rows of data matrix are labeled and below are unlabeled.

        :return: (np.ndarray, np.array, np.array, np.array)
            (x, y, ha, hn)
            x - data matrix, y - labels (np.nan for unlabeled),
            ha - indexes of labeled anomalies, hn - indexes of labeled nominals
        """
        x = None
        y = None
        if self.labeled_x is not None:
            x = self.labeled_x.copy()
            y = self.labeled_y.copy()
            ha = np.where(self.labeled_y == 1)[0]
            hn = np.where(self.labeled_y == 0)[0]
        else:
            ha = np.zeros(0, dtype=int)
            hn = np.zeros(0, dtype=int)
        if self.unlabeled_x is not None:
            if x is None:
                x = self.unlabeled_x.copy()
            else:
                x = np.append(x, self.unlabeled_x, axis=0)
            if self.unlabeled_y is not None:
                if y is not None:
                    y = np.append(y, self.unlabeled_y)
                else:
                    y = self.unlabeled_y.copy()
            else:
                if y is not None:
                    y = np.append(y, np.ones(nrow(self.unlabeled_x), dtype=int) * -1)
                else:
                    y = np.ones(nrow(self.unlabeled_x), dtype=int) * -1
        if False:
            logger.debug("x: %d, y: %d, ha: %d, hn:%d" % (nrow(x), len(y), len(ha), len(hn)))
        return x, y, ha, hn

    def get_instance_stats(self):
        nha = nhn = nul = 0
        if self.labeled_y is not None:
            nha = len(np.where(self.labeled_y == 1)[0])
            nhn = len(np.where(self.labeled_y == 0)[0])
        if self.unlabeled_x is not None:
            nul = nrow(self.unlabeled_x)
        return nha, nhn, nul

    def get_num_labeled(self):
        """Returns the number of instances for which we already have label feedback"""
        if self.labeled_y is not None:
            return len(self.labeled_y)
        return 0

    def get_query_data(self, x=None, y=None, ha=None, hn=None, unl=None, w=None, n_query=1):
        """Returns the best instance that should be queried, along with other data structures

        Args:
            x: np.ndarray
                input instances (labeled + unlabeled)
            y: np.array
                labels for instances which are already labeled, else some dummy values
            ha: np.array
                indexes of labeled anomalies
            hn: np.array
                indexes of labeled nominals
            unl: np.array
                unlabeled instances that should be ignored for query
            w: np.array
                current weight vector
            n_query: int
                number of instances to query
        """
        n = self.get_num_instances()
        n_feedback = self.get_num_labeled()
        if False:
            logger.debug("get_query_data() n: %d, n_feedback: %d" % (n, n_feedback))
        if n == 0:
            raise ValueError("No instances available")
        if x is None:
            x, y, ha, hn = self.setup_data_for_feedback()
        if w is None:
            w = self.model.w
        if unl is None:
            unl = np.zeros(0, dtype=int)
        # the top n_feedback instances in the instance list are the labeled items
        queried_items = append(np.arange(n_feedback), unl)
        x_transformed = self.model.transform_to_region_features(x, dense=False)
        order_anom_idxs, anom_score = self.model.order_by_score(x_transformed)
        xi = self.qstate.get_next_query(maxpos=n, ordered_indexes=order_anom_idxs,
                                        queried_items=queried_items,
                                        x=x_transformed, lbls=y, anom_score=anom_score,
                                        w=w, hf=append(ha, hn),
                                        remaining_budget=self.opts.budget - n_feedback,
                                        n=n_query)
        if False:
            logger.debug("ordered instances[%d]: %s\nha: %s\nhn: %s\nxi: %s" %
                         (self.opts.budget, str(list(order_anom_idxs[0:self.opts.budget])),
                          str(list(ha)), str(list(hn)), str(list(xi))))
        return xi, x, y, x_transformed, ha, hn, order_anom_idxs, anom_score

    def move_unlabeled_to_labeled(self, xi, yi):
        unlabeled_idx = xi - self.get_num_labeled()

        self.labeled_x = rbind(self.labeled_x, matrix(self.unlabeled_x[unlabeled_idx], nrow=1))
        if self.labeled_y is None:
            self.labeled_y = np.array([yi], dtype=int)
        else:
            self.labeled_y = np.append(self.labeled_y, [yi])
        mask = np.ones(self.unlabeled_x.shape[0], dtype=bool)
        mask[unlabeled_idx] = False
        self.unlabeled_x = self.unlabeled_x[mask]
        self.unlabeled_y = self.unlabeled_y[mask]

    def update_weights_with_feedback(self, xi, yi, x, y, x_transformed, ha, hn, opts):
        """Relearns the optimal weights from feedback and updates internal labeled and unlabeled matrices

        IMPORTANT:
            This API assumes that the input x, y, x_transformed are consistent with
            the internal labeled/unlabeled matrices, i.e., the top rows/values in
            these matrices are from labeled data and bottom ones are from internally
            stored unlabeled data.
        """

        # Add the newly labeled instance to the corresponding list of labeled
        # instances and remove it from the unlabeled set.
        self.move_unlabeled_to_labeled(xi, yi)

        if yi == 1:
            ha = append(ha, [xi])
        else:
            hn = append(hn, [xi])

        self.model.update_weights(x_transformed, y, ha, hn, opts)

    def get_score_variance(self, x, n_instances, opts, transform=False):
        """Computes variance in scores of top ranked instances
        """
        w = self.model.w
        if w is None:
            raise ValueError("Model not trained")
        if transform:
            x = self.model.transform_to_region_features(x, dense=False)
        ordered_indexes, scores = self.model.order_by_score(x, w=w)
        bt = get_budget_topK(n_instances, opts)
        tn = min(10, nrow(x))
        vars = np.zeros(tn, dtype=float)
        for i in np.arange(tn):
            vars[i] = get_linear_score_variance(x[ordered_indexes[i], :], w)
        # logger.debug("top %d vars:\n%s" % (tn, str(list(vars))))
        return vars

def get_rearranging_indexes(add_pos, move_pos, n):
    """Creates an array 0...n-1 and moves value at 'move_pos' to 'add_pos', and shifts others back

    Useful to reorder data when we want to move instances from unlabeled set to labeled.
    TODO:
        Use this to optimize the API StreamingAnomalyDetector.get_query_data()
        since it needs to repeatedly convert the data to transformed [node] features.

    Example:
        get_rearranging_indexes(2, 2, 10):
            array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

        get_rearranging_indexes(0, 1, 10):
            array([1, 0, 2, 3, 4, 5, 6, 7, 8, 9])

        get_rearranging_indexes(2, 9, 10):
            array([0, 1, 9, 2, 3, 4, 5, 6, 7, 8])

    :param add_pos:
    :param move_pos:
    :param n:
    :return:
    """
    if add_pos > move_pos:
        raise ValueError("add_pos must be less or equal to move_pos")
    rearr_idxs = np.arange(n)
    if add_pos == move_pos:
        return rearr_idxs
    rearr_idxs[(add_pos + 1):(move_pos + 1)] = rearr_idxs[add_pos:move_pos]
    rearr_idxs[add_pos] = move_pos
    return rearr_idxs

def read_data(opts):
    data = DataFrame.from_csv(opts.datafile, header=0, sep=',', index_col=None)
    X_train = np.zeros(shape=(data.shape[0], data.shape[1] - 1))
    for i in range(X_train.shape[1]):
        X_train[:, i] = data.iloc[:, i + 1]
    labels = np.array([1 if data.iloc[i, 0] == "anomaly" else 0 for i in range(data.shape[0])], dtype=int)
    return X_train, labels

def train_aad_model(opts, X_train):
    rng = np.random.RandomState(opts.randseed + opts.fid * opts.reruns + opts.runidx)
    # fit the model
    model = AadForest(n_estimators=opts.forest_n_trees,
                      max_samples=min(opts.forest_n_samples, X_train.shape[0]),
                      score_type=opts.forest_score_type, random_state=rng,
                      add_leaf_nodes_only=opts.forest_add_leaf_nodes_only,
                      max_depth=opts.forest_max_depth,
                      ensemble_score=opts.ensemble_score,
                      detector_type=opts.detector_type, n_jobs=opts.n_jobs)
    model.fit(X_train)
    return model

def prepare_aad_model(X, y, opts):
    if opts.load_model and opts.modelfile != "" and os.path.isfile(opts.modelfile):
        logger.debug("Loading model from file %s" % opts.modelfile)
        model = load_aad_model(opts.modelfile)
    else:
        model = train_aad_model(opts, X)

    logger.debug("total #nodes: %d" % (len(model.all_regions)))
    if False:
        if model.w is not None:
            logger.debug("w:\n%s" % str(list(model.w)))
        else:
            logger.debug("model weights are not set")
    return model

def run_feedback(sad, min_feedback, max_feedback, opts):
    """

    :param sad: StreamingAnomalyDetector
    :param max_feedback: int
    :param opts: Opts
    :return:
    """

    if False:
        # get baseline metrics
        x_transformed = sad.model.transform_to_region_features(sad.unlabeled_x, dense=False)
        ordered_idxs, _ = sad.model.order_by_score(x_transformed)
        seen_baseline = sad.unlabeled_y[ordered_idxs[0:max_feedback]]
        num_seen_baseline = np.cumsum(seen_baseline)
        logger.debug("num_seen_baseline:\n%s" % str(list(num_seen_baseline)))

    # baseline scores
    w_unif = sad.model.get_uniform_weights()
    x_transformed_baseline = sad.model.transform_to_region_features(sad.unlabeled_x, dense=False)
    order_baseline, scores_baseline = sad.model.order_by_score(x_transformed_baseline, w_unif)
    n_seen_baseline = min(max_feedback, len(sad.unlabeled_y))
    queried_baseline = order_baseline[0:n_seen_baseline]
    seen_baseline = sad.unlabeled_y[queried_baseline]
    # seen_baseline = min(max_feedback, len(sad.unlabeled_y))
    # found_baseline = np.sum(sad.unlabeled_y[order_baseline[0:seen_baseline]])

    seen = np.zeros(0, dtype=int)
    queried = np.zeros(0, dtype=int)
    unl = np.zeros(0, dtype=int)
    i = 0
    while i < max_feedback:
        i += 1
        # scores based on current weights
        xi_, x, y, x_transformed, ha, hn, order_anom_idxs, anom_score = \
            sad.get_query_data(unl=unl, n_query=max_feedback)

        order_anom_idxs_minus_ha_hn = get_first_vals_not_marked(
            order_anom_idxs, append(ha, hn), n=len(order_anom_idxs))

        bt = get_budget_topK(x_transformed.shape[0], opts)

        # Note: We will ensure that the tau-th instance is atleast 10-th (or lower) ranked
        tau_rank = min(max(bt.topK, 10), x.shape[0])

        xi = xi_[0]
        means = vars = qpos = m_tau = v_tau = None
        if opts.query_confident:
            # get the mean score and its variance for the top ranked instances
            # excluding the instances which have already been queried
            means, vars, test, v_eval, _ = get_score_variances(x_transformed, sad.model.w,
                                                               n_test=tau_rank,
                                                               ordered_indexes=order_anom_idxs,
                                                               queried_indexes=append(ha, hn))
            # get the mean score and its variance for the tau-th ranked instance
            m_tau, v_tau, _, _, _ = get_score_variances(x_transformed[order_anom_idxs_minus_ha_hn[tau_rank]],
                                                        sad.model.w, n_test=1,
                                                        test_indexes=np.array([0], dtype=int))
            qpos = np.where(test == xi)[0]  # top-most ranked instance

        if False and opts.query_confident:
            logger.debug("tau score:\n%s (%s)" % (str(list(m_tau)), str(list(v_tau))))
            strmv = ",".join(["%f (%f)" % (means[j], vars[j]) for j in np.arange(len(means))])
            logger.debug("scores:\n%s" % strmv)

        # check if we are confident that this is larger than the tau-th ranked instance
        if (not opts.query_confident) or (i <= min_feedback or
                                          means[qpos] - 3. * np.sqrt(vars[qpos]) >= m_tau):
            seen = append(seen, [y[xi]])
            queried = append(queried, xi)
            # seen += 1
            # found += y[xi]
            tm_update = Timer()
            sad.update_weights_with_feedback(xi, y[xi], x, y, x_transformed, ha, hn, opts)
            tm_update.end()
            # reset the list of queried test instances because their scores would have changed
            unl = np.zeros(0, dtype=int)
            if True:
                nha, nhn, nul = sad.get_instance_stats()
                # logger.debug("xi:%d, test indxs: %s, qpos: %d" % (xi, str(list(test)), qpos))
                # logger.debug("orig scores:\n%s" % str(list(anom_score[order_anom_idxs[0:tau_rank]])))
                logger.debug("[%d] #feedback: %d; ha: %d; hn: %d, mnw: %d, mxw: %d; update: %f sec(s)" %
                             (i, nha + nhn, nha, nhn, min_feedback, max_feedback, tm_update.elapsed()))
        else:
            # ignore this instance from query
            unl = append(unl, [xi])
            # logger.debug("skipping feedback for xi=%d at iter %d; unl: %s" % (xi, i, str(list(unl))))
            continue
        # logger.debug("y:\n%s" % str(list(y)))
    # logger.debug("w:\n%s" % str(list(sad.model.w)))
    # logger.debug("\nseen   : %s\nqueried: %s" % (str(list(seen)), str(list(queried))))
    return seen, seen_baseline, None, None

def main():

    if False:
        # DEBUG
        args = prepare_forest_aad_debug_args()
    else:
        # PRODUCTION
        args = get_command_args(debug=False)
    # print "log file: %s" % args.log_file
    configure_logger(args)

    opts = Opts(args)
    # print opts.str_opts()
    logger.debug(opts.str_opts())

    if not opts.streaming:
        raise ValueError("Only streaming supported")

    X_full, y_full = read_data(opts)
    # X_train = X_train[0:10, :]
    # labels = labels[0:10]

    logger.debug("loaded file: (%s) %s" % (str(X_full.shape), opts.datafile))
    logger.debug("results dir: %s" % opts.resultsdir)

    all_num_seen = None
    all_num_seen_baseline = None
    all_window = None
    all_window_baseline = None

    aucs = np.zeros(0, dtype=float)

    opts.fid = 1
    for runidx in opts.get_runidxs():
        tm_run = Timer()
        opts.set_multi_run_options(opts.fid, runidx)

        stream = DataStream(X_full, y_full)
        X_train, y_train = stream.read_next_from_stream(opts.stream_window)

        # logger.debug("X_train:\n%s\nlabels:\n%s" % (str(X_train), str(list(labels))))

        model = prepare_aad_model(X_train, y_train, opts)  # initial model training
        sad = StreamingAnomalyDetector(stream, model, unlabeled_x=X_train, unlabeled_y=y_train,
                                       max_buffer=opts.stream_window, opts=opts)
        sad.init_query_state(opts)

        if False:
            # use for DEBUG only
            run_feedback(sad, 0, opts.budget, opts)
            print "This is experimental/demo code for streaming integration and will be application specific." + \
                  " Exiting after reading max %d instances from stream and iterating for %d feedback..." % \
                    (opts.stream_window, opts.budget)
            exit(0)

        all_scores = np.zeros(0)
        all_y = np.zeros(0, dtype=int)

        scores = sad.get_anomaly_scores(X_train)
        # auc = fn_auc(cbind(y_train, -scores))
        all_scores = np.append(all_scores, scores)
        all_y = np.append(all_y, y_train)
        iter = 0
        seen = np.zeros(0, dtype=int)
        seen_baseline = np.zeros(0, dtype=int)
        stream_window_tmp = np.zeros(0, dtype=int)
        stream_window_baseline = np.zeros(0, dtype=int)
        stop_iter = False
        while not stop_iter:
            iter += 1

            tm = Timer()
            seen_, seen_baseline_, queried_, queried_baseline_ = run_feedback(sad,
                                                                              opts.min_feedback_per_window,
                                                                              opts.max_feedback_per_window,
                                                                              opts)
            seen = append(seen, seen_)
            seen_baseline = append(seen_baseline, seen_baseline_)
            stream_window_tmp = append(stream_window_tmp, np.ones(len(seen_)) * iter)
            stream_window_baseline = append(stream_window_baseline, np.ones(len(seen_baseline_)) * iter)
            # queried = append(queried, queried_)
            # queried_baseline = append(queried_baseline, queried_baseline_)
            # logger.debug("seen:\n%s;\nbaseline:\n%s" % (str(list(seen)), str(list(seen_baseline))))

            x_eval, y_eval = sad.get_next_from_stream(sad.max_buffer)
            if x_eval is None or iter >= opts.max_windows:
                if iter >= opts.max_windows:
                    logger.debug("Exceeded %d iters; exiting stream read..." % opts.max_windows)
                stop_iter = True
            else:
                scores = sad.get_anomaly_scores(x_eval)  # compute scores before updating the model

                all_scores = np.append(all_scores, scores)
                all_y = np.append(all_y, y_eval)

                if opts.allow_stream_update:
                    sad.update_model_from_buffer()

                sad.move_buffer_to_unlabeled()

            logger.debug(tm.message("Stream window [%d]: algo [%d/%d]; baseline [%d/%d]: " %
                                    (iter, np.sum(seen), len(seen), np.sum(seen_baseline), len(seen_baseline))))

        auc = fn_auc(cbind(all_y, -all_scores))
        # logger.debug("AUC: %f" % auc)
        aucs = append(aucs, [auc])

        # queried_baseline = order(all_scores, decreasing=True)[0:opts.budget]
        num_seen_tmp = np.cumsum(seen)  # np.cumsum(all_y[queried])
        # logger.debug("\nnum_seen    : %s" % (str(list(num_seen_tmp)),))

        num_seen_baseline = np.cumsum(seen_baseline)  # np.cumsum(all_y[queried_baseline])
        # logger.debug("Numseen in %d budget (overall):\n%s" % (opts.budget, str(list(num_seen_baseline))))

        stream_window_baseline = append(np.array([opts.fid, opts.runidx],
                                                 dtype=stream_window_baseline.dtype),
                                        stream_window_baseline)
        stream_window = np.ones(len(stream_window_baseline) + 2, dtype=stream_window_tmp.dtype) * -1
        stream_window[0:2] = [opts.fid, opts.runidx]
        stream_window[2:(2+len(stream_window_tmp))] = stream_window_tmp

        # queried = append(np.array([opts.fid, opts.runidx], dtype=queried.dtype), queried)
        # queried_baseline = append(np.array([opts.fid, opts.runidx], dtype=queried_baseline.dtype), queried_baseline)

        # num_seen_baseline has the uniformly maximum number of queries.
        # the number of queries in num_seen will vary under the query confidence mode
        num_seen = np.ones(len(num_seen_baseline) + 2, dtype=num_seen_tmp.dtype) * -1
        num_seen[0:2] = [opts.fid, opts.runidx]
        num_seen[2:(2+len(num_seen_tmp))] = num_seen_tmp

        num_seen_baseline = append(np.array([opts.fid, opts.runidx], dtype=num_seen_baseline.dtype), num_seen_baseline)

        # all_queried = rbind(all_queried, matrix(queried, nrow=1))
        # all_queried_baseline = rbind(all_queried_baseline, matrix(queried_baseline, nrow=1))

        all_num_seen = rbind(all_num_seen, matrix(num_seen, nrow=1))
        all_num_seen_baseline = rbind(all_num_seen_baseline, matrix(num_seen_baseline, nrow=1))
        all_window = rbind(all_window, matrix(stream_window, nrow=1))
        all_window_baseline = rbind(all_window_baseline, matrix(stream_window_baseline, nrow=1))

        logger.debug(tm_run.message("Completed runidx: %d" % runidx))

    results = SequentialResults(num_seen=all_num_seen,
                                # true_queried_indexes=all_queried,
                                num_seen_baseline=all_num_seen_baseline,
                                # true_queried_indexes_baseline=all_queried_baseline,
                                stream_window=all_window,
                                stream_window_baseline=all_window_baseline,
                                aucs=aucs)
    write_sequential_results_to_csv(results, opts)

if __name__ == "__main__":
    main()


"""elbb.engine"""

import os
import time
import numpy
import scipy
import random
import pyautogui
import threading
import pytesseract

from scipy import interpolate

from elbb.queue import log
from elbb.resources import UI

pyautogui.FAILSAFE = False
COOLDOWNS = dict()

def launch_client():
    game_client = os.environ['GAMECLIENT']
    if not os.popen(f'pgrep -f {game_client}').read():
        os.system(f'{game_client} &')

def locate(image):
    location = pyautogui.locateOnScreen(
        image.path,
        grayscale=True
    )
    return location

def randomize(location):
    x = random.randint(location.left, location.left + location.width)
    y = random.randint(location.top, location.top + location.height)
    return (x, y)

def move_mouse(x_dest, y_dest):
    """
    Semi-natural mouse movement using Bézier curves
    https://stackoverflow.com/a/44666798
    """

    x_curr, y_curr = pyautogui.position()

    control = random.randint(3, 5)
    x = numpy.linspace(x_curr, x_dest, num=control, dtype='int')
    y = numpy.linspace(y_curr, y_dest, num=control, dtype='int')

    R = 10
    x_rand = scipy.random.randint(-R, R, size=control)
    y_rand = scipy.random.randint(-R, R, size=control)

    x_rand[0] = y_rand[0] = x_rand[-1] = y_rand[-1] = 0
    x += x_rand
    y += y_rand

    degree = 3 if control > 3 else control - 1
    tck, u = interpolate.splprep([x, y], k=degree)
    u = numpy.linspace(0, 1, num=1920)

    points = interpolate.splev(u, tck)
    duration = 0.2
    timeout = duration / len(points[0])

    for point in zip(*(i.astype(int) for i in points)):
        pyautogui.platformModule._moveTo(*point)
        time.sleep(timeout)

def get_delta(x, y):
    x_curr, y_curr = pyautogui.position()
    return x_curr + x, y_curr + y

def click(n=1):
    for _ in range(n):
        pyautogui.click()
        time.sleep(random.uniform(0.05, 0.3))

def move_to(image, click_mode=False):
    location = locate(image)
    if not location:
        return False
    x, y = randomize(location)
    move_mouse(x, y)
    if click_mode:
        click()
    return True

def click_in_region(x1, y1, x2, y2, clicks=1):
    dx, dy = x2 - x1, y2 - y1
    rx, ry = random.randint(0, dx), random.randint(0, dy)
    pyautogui.moveTo(x1+rx+128, y1+ry+128)
    click(n=clicks)

def select_walk():
    click_in_region(1, 737, 29, 766, clicks=2)

def select_sit():
    click_in_region(33, 737, 61, 766, clicks=2)

def select_look():
    click_in_region(65, 737, 93, 766, clicks=2)

def select_use():
    click_in_region(97, 737, 125, 766, clicks=2)

def select_fist():
    click_in_region(129, 737, 157, 766, clicks=2)

def select_trade():
    click_in_region(161, 737, 189, 766, clicks=2)

def select_combat():
    click_in_region(193, 737, 221, 766, clicks=2)

def select_inventory():
    click_in_region(225, 737, 253, 766)

def select_spells():
    click_in_region(257, 737, 285, 766)

def select_manufacture():
    click_in_region(289, 737, 317, 766)

def select_emotes():
    click_in_region(321, 737, 349, 766)

def select_quest_log():
    click_in_region(353, 737, 381, 766)

#def select_map():
#    click_in_region(385, 737, 413, 766)

def select_notepad():
    click_in_region(417, 737, 445, 766)

def select_buddy():
    click_in_region(449, 737, 477, 766)

def select_map():
    click_in_region(481, 737, 509, 766)

def select_stats():
    click_in_region(513, 737, 541, 766)

def select_console():
    click_in_region(545, 737, 573, 766)

def select_help():
    click_in_region(577, 737, 605, 766)

def select_ranging():
    click_in_region(609, 737, 637, 766)

def select_minimap():
    click_in_region(641, 737, 669, 766)

def select_options():
    click_in_region(673, 737, 701, 766)

def select_book_status():
    click_in_region(963, 605, 1021, 616)

def select_item_1():
    click_in_region(995, 67, 1020, 92)

def select_item_2():
    click_in_region(995, 97, 1020, 122)

def select_item_3():
    click_in_region(995, 127, 1020, 152)

def select_item_4():
    click_in_region(995, 157, 1020, 182)

def select_item_5():
    click_in_region(995, 187, 1020, 212)

def select_item_6():
    click_in_region(995, 217, 1020, 242)

def select_logo():
    click_in_region(968, 4, 1020, 60)

def write(text):
    pyautogui.write(text)

def move_inventory_to_storage(item, n=1):
    log(f'Moving {item.__name__} from inventory to storage...')

    select_walk()

    # click the item in the inventory
    success = move_to(item.Inventory, click_mode=True)
    if not success:
        log(f'Could not find {item.__name__} in the inventory!', category='bad')
        return False

    # move the item to the storage
    storage_location = locate(UI.Storage.Banner)
    x_sto, y_sto = randomize(storage_location)
    rx, ry = random.randint(50, 100), random.randint(50, 100)
    move_mouse(x_sto+rx, y_sto+ry)
    click(n)
    pyautogui.click(button='right')

    log(f'{item.__name__} moved to storage.', category='good')
    return True

def move_storage_to_inventory(item, n=1):
    log(f'Moving {item.__name__} from storage to inventory...')

    select_walk()

    # click the correct storage menu for the item
    success = move_to(item.storage_menu)
    if success:
        click()

    # click the item in the storage
    success = move_to(item.Storage)
    if not success:
        log(f'Could not find {item.__name__} in the storage!', category='bad')
        return False
    # small sleep here to reduce misclicking
    time.sleep(random.uniform(0.1, 0.2))
    click()

    # move the item to the inventory
    inventory_location = locate(UI.Inventory.Banner)
    x_inv, y_inv = randomize(inventory_location)
    rx, ry = random.randint(0, 10), random.randint(50, 100)
    move_mouse(x_inv+rx, y_inv+ry)
    click(n)
    pyautogui.click(button='right')

    log(f'{item.__name__} moved to inventory.', category='good')
    return True

def run_cooldown(name, seconds):
    COOLDOWNS[name] = seconds
    for _ in range(seconds):
        time.sleep(1)
        COOLDOWNS[name] -= 1
    COOLDOWNS.pop(name)

def start_cooldown(name, seconds):
    t = threading.Thread(
        target=run_cooldown,
        kwargs={'name': name, 'seconds': seconds}
    )
    t.daemon = True
    t.start()

def eat(food, storage_mode=False):
    can_eat = False

    if locate(food.Inventory):
        can_eat = True
    elif storage_mode:
        success = move_storage_to_inventory(food, n=1)
        if success:
            can_eat = True
        else:
            return False

    if can_eat:
        select_use()
        move_to(food.Inventory)
        click()
        start_cooldown(food.__name__, food.cooldown)
        log(f'Ate {food.__name__}.', category='good')
    else:
        return False

def toggle_console():
    pyautogui.press('f1')

def toggle_map():
    # TODO: check if map is already loaded here?
    pyautogui.press('tab')

def close_map():
    if locate(UI.Map.Legend):
        toggle_map()

def zoom_in(n=1):
    for _ in range(n):
        pyautogui.press('pageup')

def zoom_out(n=1):
    for _ in range(n):
        pyautogui.press('pagedown')

def pan_down(n=1):
    for _ in range(n):
        pyautogui.press('up')

def pan_up(n=1):
    for _ in range(n):
        pyautogui.press('down')

def click_flag(shitty=False):
    select_use()
    zoom_in(n=30)
    pan_down(n=10)

    if shitty:
        zoom_out()

    def scan_for_flag():
        for _ in range(100):
            pyautogui.press('left')
            img = pyautogui.screenshot(
                region=(590, 462, 100, 100)
            )
            pixels = img.load()
            for y in range(0, 100, 10):
                for x in range(0, 100, 10):
                    _, _, b = pixels[x, y]
                    if b > 75:
                        return (x, y)
        return None

    offset_coords = scan_for_flag()
    if not offset_coords:
        return False

    x, y = offset_coords

    # attempt 1
    move_mouse(590+x, 462+y)
    click()
    # attempt 2
    move_mouse(590+x+5, 462+y+5)
    click()
    # attempt 3
    move_mouse(640, 512)
    click()

    # wait for new map to load
    time.sleep(3)

    return True

def click_storage():
    select_use()
    zoom_in(n=30)
    pan_down(n=10)

    def scan_for_storage():
        for _ in range(100):
            pyautogui.press('left')
            img = pyautogui.screenshot(
                region=(590, 462, 100, 100)
            )
            pixels = img.load()
            for y in range(0, 100, 10):
                for x in range(0, 100, 10):
                    _, _, b = pixels[x, y]
                    if b < 10:
                        return (x, y)
        return None

    offset_coords = scan_for_storage()
    if not offset_coords:
        return False

    x, y = offset_coords
    move_mouse(590+x, 462+y)
    click(n=2)
    return True

def click_on_map(x, y):
    toggle_map()
    # wait a bit for the map to load
    time.sleep(5)
    move_mouse(x, y)
    click()
    toggle_map()

def go_to_isla_prima():
    pyautogui.write('#beam me')
    pyautogui.press('enter')

def go_to_white_stone():
    close_map()
    go_to_isla_prima()
    click_on_map(362, 776)
    # wait to walk to ship
    time.sleep(50)
    click_flag()

def go_to_desert_pines():
    go_to_white_stone()
    click_on_map(889, 724)
    # wait to walk to ship
    time.sleep(20)
    click_flag(shitty=True)

def go_to_valley_of_the_dwarves():
    go_to_desert_pines()
    click_on_map(900, 230)
    # wait to walk to ship
    time.sleep(20)
    click_flag()

def go_to_valley_of_the_dwarves_storage():
    go_to_valley_of_the_dwarves()
    click_on_map(434, 380)
    # wait to walk to storage
    time.sleep(20)
    select_sit()

def go_to_portland():
    go_to_desert_pines()
    click_on_map(902, 262)
    # wait to walk to ship
    time.sleep(30)
    click_flag()

def go_to_crystal_cave():
    go_to_desert_pines()
    click_on_map(364, 290)
    # wait to walk across map
    time.sleep(110)
    # enter cave
    move_mouse(620, 219)
    click()

def go_to_crystal_cave_sulfur():
    go_to_crystal_cave()
    # wait a sec for cave to load
    time.sleep(1)
    click_on_map(524, 172)
    # wait to walk to sulfur
    time.sleep(25)
    select_sit()

def spin_compass(comp_x, comp_y):
    for _ in range(100):
        pyautogui.press('left')
        img = pyautogui.screenshot(
            region=(comp_x, comp_y, 16, 16)
        )
        pixels = img.load()
        for y in range(16):
            for x in range(16):
                r, g, b = pixels[x, y]
                if (r == 255) and (g == 255):
                    return

def point_north():
    spin_compass(1116, 840)

def point_north_east():
    spin_compass(1135, 852)

def point_east():
    spin_compass(1135, 858)

def point_south_east():
    spin_compass(1132, 876)

def point_south():
    spin_compass(1117, 882)

def point_south_west():
    spin_compass(1096, 872)

def point_west():
    spin_compass(1092, 853)

def point_north_west():
    spin_compass(1101, 839)

def open_inventory():
    # inventory already open
    if locate(UI.Inventory.Banner):
        return
    select_inventory()

def close_everything():
    while move_to(UI.Generic.X, click_mode=True):
        pass

def sto_items(items):
    close_everything()

    # fix for storage
    zoom_in(n=30)
    pan_up(n=10)
    point_south()
    point_north()

    # click storage
    move_mouse(530, 680)
    click()

    # open storage
    move_mouse(720+128, 366+128)
    move_to(UI.Storage.Open)
    click(n=2) # not sure why, but gotta double click here?

    open_inventory()

    # select 100
    move_to(UI.Inventory.OneHundred, click_mode=True)

    for item in items:
        move_inventory_to_storage(item, n=2)

def ocr(x, y, dx, dy):
    ss = pyautogui.screenshot(region=(x, y, dx, dy))
    return pytesseract.image_to_string(ss)

def get_load_status():
    ss = pyautogui.screenshot(region=(546, 849, 24, 12))
    ss_pixels = ss.load()

    for y in range(12):
        for x in range(24):
            r, g, b = ss_pixels[x, y]

            if (r, g, b) == (0, 0, 0):
                ss.putpixel((x, y), (255, 255, 255))
            elif r < 60:
                ss.putpixel((x, y), (255, 255, 255))
            elif (-5 < (r - g) < 5):
                ss.putpixel((x, y), (0, 0, 0))
            else:
                ss.putpixel((x, y), (255, 255, 255))

    text = pytesseract.image_to_string(ss, config='--psm 7 -c tessedit_char_whitelist=0123456789')
    return int(text)

def get_manufacture_status():
    m_location = locate(UI.Manufacture.Banner)
    m_text = ocr(m_location.left-159, m_location.top+115, 255, 16)
    log(f'Got manufacture status: {m_text}')
    return m_text

def get_inventory_text():
    log('Scanning inventory...')

    select_look()

    inventory = list()
    inventory_location = locate(UI.Inventory.Banner)
    init_x = x = inventory_location.left - 150
    y = inventory_location.top + 45

    for i in range(36):
        move_mouse(x, y)
        click()

        text = ocr(inventory_location.left-175, inventory_location.top+325, 450, 100)
        inventory.append({
            'text': text,
            'location': (x, y)
        })

        # TODO: Fix this later? (breaks on duplicate item)
        if i != 0:
            if inventory[i]['text'] == inventory[i-1]['text']:
                break

        if (i+1) % 6 == 0:
            x = init_x
            y += 50
        else:
            x += 50

    log('Finished scanning inventory.', category='good')
    return inventory

def get_book_status():
    move_mouse(992+128, 610+128)
    click()
    text = ocr(4+128, 24+128, 353, 17)
    log(f'Got book status: {text}')
    return text


import torch
from model.base_model import BaseModel
from model.networks import base_function, external_function
import model.networks as network
from util import task, util, openpose_utils
import itertools
import data as Dataset
import numpy as np
from itertools import islice
import random
import os, ntpath
import numpy as np
import matplotlib.pyplot as plt
from collections import OrderedDict
import glob
import cv2

class Dance(BaseModel):
    def name(self):
        return "Pose-Guided Person Image Animation"

    @staticmethod
    def modify_options(parser, is_train=True):
        parser.add_argument('--attn_layer', action=util.StoreList, metavar="VAL1,VAL2...")
        parser.add_argument('--kernel_size', action=util.StoreDictKeyPair, metavar="KEY1=VAL1,KEY2=VAL2...")

        parser.add_argument('--layers', type=int, default=3, help='number of layers in G')
        parser.add_argument('--netG', type=str, default='dance', help='The name of net Generator')
        parser.add_argument('--netD', type=str, default='res', help='The name of net Discriminator')
        parser.add_argument('--netD_V', type=str, default='temporal', help='The name of net Discriminator')
        parser.add_argument('--init_type', type=str, default='orthogonal', help='Initial type')

        parser.add_argument('--ratio_g2d', type=float, default=0.1, help='learning rate ratio G to D')
        parser.add_argument('--lambda_rec', type=float, default=5.0, help='weight for image reconstruction loss')
        parser.add_argument('--lambda_g', type=float, default=2.0, help='weight for generation loss')
        parser.add_argument('--lambda_correct', type=float, default=5.0, help='weight for generation loss')
        parser.add_argument('--lambda_style', type=float, default=500.0, help='weight for generation loss')
        parser.add_argument('--lambda_content', type=float, default=0.5, help='weight for generation loss')
        parser.add_argument('--lambda_regularization', type=float, default=0.0025, help='weight for generation loss')
        parser.add_argument('--frames_D_V', type=int, default=6, help='number of frames of D_V')

        parser.add_argument('--use_spect_g', action='store_false')
        parser.add_argument('--use_spect_d', action='store_false')
        parser.add_argument('--write_ext', type=str, help='png | jpg')

        parser.set_defaults(use_spect_g=False)
        parser.set_defaults(use_spect_d=True)

        # display
        parser.set_defaults(display_freq=2)
        parser.set_defaults(eval_iters_freq=500)
        parser.set_defaults(print_freq=2)
        parser.set_defaults(save_latest_freq=2)
        parser.set_defaults(save_iters_freq=1000)

        parser.set_defaults(write_ext='png')
        return parser

    def __init__(self, opt):
        BaseModel.__init__(self, opt)
        self.loss_names = ['app_gen','correctness_p', 'correctness_r',
                           'content_gen','style_gen',
                           'regularization_p','regularization_r',
                           'ad_gen', 'dis_img_gen',
                           'ad_gen_v', 'dis_img_gen_v']

        self.visual_names = ['ref_image', 'ref_skeleton', 'BP_step', 'P_step','img_gen']
        self.model_names = ['G', 'D', 'D_V']

        self.FloatTensor = torch.cuda.FloatTensor if len(self.gpu_ids)>0 \
            else torch.FloatTensor
        self.ByteTensor = torch.cuda.ByteTensor if len(self.gpu_ids)>0 \
            else torch.ByteTensor

        self.net_G = network.define_g(opt, filename='generator', image_nc=opt.image_nc, structure_nc=opt.structure_nc, ngf=64, img_f=512,
                                      layers=opt.layers, num_blocks=2, use_spect=opt.use_spect_g, attn_layer=opt.attn_layer,
                                      norm='instance', activation='LeakyReLU', extractor_kz=opt.kernel_size)
        if len(opt.gpu_ids) > 1:
            self.net_G = torch.nn.DataParallel(self.net_G, device_ids=self.gpu_ids)

        self.flow2color = util.flow2color()
        self.convert2skeleton = openpose_utils.tensor2skeleton(spatial_draw=True)

        self.net_D = network.define_d(opt, ndf=32, img_f=128, layers=4, use_spect=opt.use_spect_d)
        if len(opt.gpu_ids) > 1:
            self.net_D = torch.nn.DataParallel(self.net_D, device_ids=self.gpu_ids)

        self.net_D_V = network.define_d(opt, name=opt.netD_V, input_length=opt.frames_D_V, ndf=32, img_f=128, layers=4, use_spect=opt.use_spect_d)
        if len(opt.gpu_ids) > 1:
            self.net_D_V = torch.nn.DataParallel(self.net_D_V, device_ids=self.gpu_ids)

        if self.isTrain:
            self.GANloss = external_function.AdversarialLoss(opt.gan_mode).to(opt.device)
            self.L1loss = torch.nn.L1Loss()
            self.Correctness = external_function.PerceptualCorrectness().to(opt.device)

            self.Regularization = external_function.MultiAffineRegularizationLoss(kz_dic=opt.kernel_size).to(opt.device)
            self.Vggloss = external_function.VGGLoss().to(opt.device)

            # define the optimizer
            self.optimizer_G = torch.optim.Adam(itertools.chain(
                                               filter(lambda p: p.requires_grad, self.net_G.parameters())),
                                               lr=opt.lr, betas=(0.0, 0.999))
            self.optimizers.append(self.optimizer_G)

            self.optimizer_D = torch.optim.Adam(itertools.chain(
                                filter(lambda p: p.requires_grad, self.net_D.parameters()),
                                filter(lambda p: p.requires_grad, self.net_D_V.parameters())),
                                lr=opt.lr*opt.ratio_g2d, betas=(0.0, 0.999))
            self.optimizers.append(self.optimizer_D)

            if self.opt.use_mask:
                # use mask to calculate the correctness loss for foreground content
                self.opt.lambda_correct = 2.0
        else:
            self.results_dir_base = self.opt.results_dir
        self.esp=1e-6
        self.setup(opt)

    def set_input(self, data):
        # move to GPU and change data types
        opt = self.opt
        if self.isTrain:
            _, n_frames_total, height, width = data['gen_images'].size()  #
            self.n_frames_total = n_frames_total // opt.image_nc
            self.n_frames_load  = opt.max_frames_per_gpu

            gen_images = data['gen_images']
            gen_masks = data['gen_masks'] if self.opt.use_mask else None
            gen_skeletons = data['gen_skeleton']
            self.gen_images, self.gen_skeletons, self.gen_masks=[],[],[]
            for i in range(0, self.n_frames_total, self.n_frames_load):
                P_step = gen_images[:, i*opt.image_nc:(i+self.n_frames_load)*opt.image_nc]
                P_step = P_step.view(-1, self.n_frames_load, opt.image_nc, height, width)
                self.gen_images.append(P_step)

                skeleton_step  = gen_skeletons[:, i*opt.structure_nc:(i+self.n_frames_load)*opt.structure_nc]
                skeleton_step  = skeleton_step.view(-1, self.n_frames_load, opt.structure_nc, height, width)
                self.gen_skeletons.append(skeleton_step)

                if self.opt.use_mask:
                    mask_step = gen_masks[:, i*opt.mask_nc:(i+self.n_frames_load)*opt.mask_nc]
                    mask_step = mask_step.view(-1, self.n_frames_load, opt.mask_nc, height, width)
                    self.gen_masks.append(mask_step)

            self.ref_image = data['ref_image'].cuda()
            self.ref_skeleton = data['ref_skeleton'].cuda()
            self.pre_gt_image = self.ref_image
            self.pre_image = None
            self.pre_skeleton = None
            self.image_paths = data['gen_paths']
        else:
            _, n_frames_total, height, width = data['gen_images'].size()  #
            self.n_frames_total = n_frames_total // opt.image_nc
            self.n_frames_load = self.n_frames_total

            gen_images = data['gen_images']
            gen_skeletons = data['gen_skeleton']
            self.openpose_kp = data['gen_kps_noise']
            self.video2d_kp = data['gen_kps_clean']
            self.gen_images = gen_images.view(-1, self.n_frames_load, opt.image_nc, height, width).cuda()
            self.gen_skeletons = gen_skeletons.view(-1, self.n_frames_load, opt.structure_nc, height, width).cuda()
            self.frame_idx = data['frame_idx']

            if self.frame_idx == self.opt.start_frame+self.opt.n_frames_pre_load_test:
                self.ref_image = data['ref_image'].cuda()
                self.ref_skeleton = data['ref_skeleton'].cuda()
                self.pre_image = None
                self.pre_skeleton = None

            self.change_seq = data['change_seq']
            self.image_paths = data['gen_paths']
            self.ref_paths = data["ref_path"]
            if not self.if_cross_eval(self.image_paths, self.ref_paths):
                self.opt.results_dir = os.path.join(self.results_dir_base,
                                                    self.image_paths[0][0].split('/')[-2])
            else:
                name = self.image_paths[0][0].split('/')[-2] + '_with_' + self.ref_paths[0].split('/')[-2]
                self.opt.results_dir = os.path.join(self.results_dir_base, name)

    def if_cross_eval(self,image_paths, ref_paths):
        video_driven = image_paths[0][0].split('/')[-2]
        video_ref = ref_paths[0].split('/')[-2]
        return False if video_ref == video_driven else True

    def write2video(self, name_list):
        images=[]
        for name in name_list:
            images.append(sorted(glob.glob(self.opt.results_dir+'/*_'+name+'.'+self.opt.write_ext)))

        image_array=[]
        for i in range(len(images[0])):
            cat_im=None
            for image_list in images:
                im = cv2.imread(image_list[i])
                if cat_im is not None:
                    cat_im = np.concatenate((cat_im, im), axis=1)
                else:
                    cat_im = im
            image_array.append(cat_im)

        res=''
        for name in name_list:
            res += (name +'_')
        out_name = self.opt.results_dir+'_'+res+'.mp4'
        print('write video %s'%out_name)
        height, width, layers = cat_im.shape
        size = (width,height)
        out = cv2.VideoWriter(out_name, cv2.VideoWriter_fourcc(*'mp4v'), 15, size)

        for i in range(len(image_array)):
            out.write(image_array[i])
        out.release()

    def get_current_visuals(self):
        """Return visualization images"""
        visual_ret = OrderedDict()
        for name in self.visual_names:
            if isinstance(name, str):
                value = getattr(self, name)
                if '_step' in name:
                    value = value[0]
                    list_value=[]
                    for i in range(value.size(0)):
                        _value = value[i, -3:,...]
                        list_value.append(_value.unsqueeze(0))
                    value=list_value
                if '_skeleton' in name:
                    value = value[0, -3:,...]
                if 'flow_field' in name or 'occlusion' in name:
                    list_value = [item for sub_list in value for item in sub_list]
                    value = list_value

                if isinstance(value, list):
                    # visual multi-scale ouputs
                    for i in range(len(value)):
                        visual_ret[name + str(i)] = self.convert2im(value[i], name)
                    # visual_ret[name] = util.tensor2im(value[-1].data)
                else:
                    visual_ret[name] = self.convert2im(value, name)
        return visual_ret

    def test(self, save_features=False, save_all=False, generate_edge=True):
        """Forward function used in test time"""
        self.img_gen, flow_fields, \
        occlusion, P_previous_recoder \
                        = self.net_G(self.gen_skeletons,
                                    self.ref_image,
                                    self.ref_skeleton,
                                    self.pre_image,
                                    self.pre_skeleton)
        self.pre_image = self.img_gen[-1]
        self.pre_skeleton = self.gen_skeletons[:,-1,...]

        self.image_paths = self.check_image_paths()
        if self.frame_idx == self.opt.start_frame+self.opt.n_frames_pre_load_test:
            self.save_results(self.ref_image, data_name='ref', data_ext=self.opt.write_ext)

        # save the generated image
        gen_image = torch.cat(self.img_gen, 0)
        self.save_results(gen_image, data_name='vis', data_ext=self.opt.write_ext)

        # save the gt image
        gt_image = self.gen_images.squeeze(0)
        self.save_results(gt_image, data_name='gt', data_ext=self.opt.write_ext)

        # save the skeleton image
        current_skeleton = self.convert2skeleton(self.openpose_kp[0,...], 'COCO_17')
        obtained_skeleton = self.convert2skeleton(self.video2d_kp[0], 'human36m_17')
        for i in range(len(current_skeleton)):
            short_path = ntpath.basename(self.image_paths[i])
            name = os.path.splitext(short_path)[0]
            util.mkdir(self.opt.results_dir)

            in_img_name = '%s_%s.%s' % (name, 'skeleton_in', self.opt.write_ext)
            img_path = os.path.join(self.opt.results_dir, in_img_name)
            skeleton_in = current_skeleton[i].astype(np.uint8)
            util.save_image(skeleton_in, img_path)

            out_img_name = '%s_%s.%s' % (name, 'skeleton_out', self.opt.write_ext)
            img_path = os.path.join(self.opt.results_dir, out_img_name)
            skeleton_out = obtained_skeleton[i].astype(np.uint8)
            util.save_image(skeleton_out, img_path)

        if self.change_seq:
            name_list=['gt', 'vis', 'skeleton_in', 'skeleton_out']
            self.write2video(name_list)

    def check_image_paths(self):
        names = []
        for name in self.image_paths:
            if isinstance(name, tuple):
                name = name[0]
            names.append(name)
        return names

    def update(self):
        """Run forward processing to get the inputs"""
        for i in range(len(self.gen_images)):
            self.P_step = self.gen_images[i].cuda()
            self.BP_step = self.gen_skeletons[i].cuda()
            self.mask_step = self.gen_masks[i].cuda() if self.opt.use_mask else None
            self.P_gt_previous_recoder = torch.cat((self.pre_gt_image.unsqueeze(1), self.P_step[:,:-1,...]), 1)

            self.img_gen, self.flow_fields, \
            self.occlusion, self.P_previous_recoder \
                            = self.net_G(self.BP_step,
                                        self.ref_image,
                                        self.ref_skeleton,
                                        self.pre_image,
                                        self.pre_skeleton)

            self.optimizer_D.zero_grad()
            self.backward_D()
            self.optimizer_D.step()

            self.optimizer_G.zero_grad()
            self.backward_G()
            self.optimizer_G.step()

            self.pre_image  = self.img_gen[-1].detach()
            self.pre_skeleton = self.BP_step[:,-1,...].detach()
            self.pre_gt_image = self.P_step[:,-1,...]

    def backward_D_basic(self, netD, real, fake):
        """Calculate GAN loss for the discriminator"""
        # Real
        D_real = netD(real)
        D_real_loss = self.GANloss(D_real, True, True)
        # fake
        D_fake = netD(fake.detach())
        D_fake_loss = self.GANloss(D_fake, False, True)
        # loss for discriminator
        D_loss = (D_real_loss + D_fake_loss) * 0.5

        if self.opt.gan_mode == 'wgangp':
            gradient_penalty, gradients = external_function.cal_gradient_penalty(netD, real, fake.detach())
            D_loss += gradient_penalty
        D_loss.backward()

        return D_loss

    def backward_D(self):
        """Calculate the GAN loss for the discriminators"""

        # Spatial GAN Loss
        base_function._unfreeze(self.net_D)
        i = np.random.randint(len(self.img_gen))
        fake = self.img_gen[i]
        real = self.P_step[:,i,...]
        self.loss_dis_img_gen = self.backward_D_basic(self.net_D, real, fake)

        # Temporal GAN Loss
        base_function._unfreeze(self.net_D_V)
        i = np.random.randint(len(self.img_gen)-self.opt.frames_D_V+1)
        fake = []
        real = []
        for frame in range(self.opt.frames_D_V):
            fake.append(self.img_gen[i+frame].unsqueeze(2))
            real.append(self.P_step[:,i+frame,...].unsqueeze(2))
        fake = torch.cat(fake, dim=2)
        real = torch.cat(real, dim=2)
        self.loss_dis_img_gen_v = self.backward_D_basic(self.net_D_V, real, fake)

    def backward_G(self):
        """Calculate training loss for the generator"""
        loss_style_gen, loss_content_gen, loss_app_gen, loss_pose=0,0,0,0

        # Calculate the Reconstruction Loss
        for i in range(len(self.img_gen)):
            gen = self.img_gen[i]
            gt = self.P_step[:,i,...]
            loss_app_gen += self.L1loss(gen, gt)

            content_gen, style_gen = self.Vggloss(gen, gt)
            loss_style_gen += style_gen
            loss_content_gen += content_gen

        self.loss_style_gen = loss_style_gen * self.opt.lambda_style
        self.loss_content_gen = loss_content_gen * self.opt.lambda_content
        self.loss_app_gen = loss_app_gen * self.opt.lambda_rec

        loss_correctness_p, loss_regularization_p=0, 0
        loss_correctness_r, loss_regularization_r=0, 0

        # Calculate the Sampling Correctness Loss and Regularization Loss
        for i in range(len(self.flow_fields)):
            flow_field_i = self.flow_fields[i]
            flow_p, flow_r=[],[]
            for j in range(0, len(flow_field_i), 2):
                flow_p.append(flow_field_i[j])
                flow_r.append(flow_field_i[j+1])

            mask = self.mask_step[:,i,...] if self.opt.use_mask else None
            correctness_r = self.Correctness(self.P_step[:,i,...], self.ref_image,
                                            flow_r, self.opt.attn_layer, mask)
            correctness_p = self.Correctness(self.P_step[:,i,...], self.P_gt_previous_recoder[:,i,...],
                                            flow_p, self.opt.attn_layer, mask)
            loss_correctness_p += correctness_p
            loss_correctness_r += correctness_r

            loss_regularization_p += self.Regularization(flow_p)
            loss_regularization_r += self.Regularization(flow_r)

        self.loss_correctness_p = loss_correctness_p * self.opt.lambda_correct
        self.loss_correctness_r = loss_correctness_r * self.opt.lambda_correct
        self.loss_regularization_p = loss_regularization_p * self.opt.lambda_regularization
        self.loss_regularization_r = loss_regularization_r * self.opt.lambda_regularization

        # Spatial GAN Loss
        base_function._freeze(self.net_D)
        i = np.random.randint(len(self.img_gen))
        fake = self.img_gen[i]
        D_fake = self.net_D(fake)
        self.loss_ad_gen = self.GANloss(D_fake, True, False) * self.opt.lambda_g

        # Temporal GAN Loss
        base_function._freeze(self.net_D_V)
        i = np.random.randint(len(self.img_gen)-self.opt.frames_D_V+1)
        fake = []
        for frame in range(self.opt.frames_D_V):
            fake.append(self.img_gen[i+frame].unsqueeze(2))
        fake = torch.cat(fake, dim=2)
        D_fake = self.net_D_V(fake)
        self.loss_ad_gen_v = self.GANloss(D_fake, True, False) * self.opt.lambda_g

        total_loss = 0
        for name in self.loss_names:
            if name != 'dis_img_gen_v' and name != 'dis_img_gen':
                total_loss += getattr(self, "loss_" + name)
        total_loss.backward()

    def optimize_parameters(self):
        self.update()


from typing import Tuple, List, Dict
import argparse
from collections import OrderedDict
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import pytorch_lightning as pl
# from .common import wrappers
from lib.sepsis_simulator.policy import SepsisOptimalSolver
from lib.sepsis_simulator.sepsisSimDiabetes.State import State
from lib.sepsis_simulator.sepsisSimDiabetes.Action import Action
from lib.mimic3.dataset import HypotensionDataset
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence, pad_packed_sequence

from lib.sepsis_simulator.utils import run_policy_to_get_exp
import pickle
import numpy as np
from os.path import exists as pexists
from lib.sepsis_simulator.utils import run_policy
from lib.utils import Timer
from lib.lightning.utils import Squeeze, MinMaxNormalize
from lib.random_search import RandomSearch
from lib import nodegam
from sklearn.metrics import roc_auc_score
from ..utils import evaluating
from lib.lightning.utils import create_fc_layer

class LinearDisc(nn.Module):
    def __init__(self, hparams) -> None:
        '''
        experiment: choose which experiment to run. If select sepsis,
            the the architecture includes a Min-Max normalization layer.
            If select hypotension, then no normalization is performed
        '''
        super().__init__()
        self.hparams = hparams
        self.build()

        # Backward compatability
        # lr = self.hparams.get('disc_lr', self.hparams.get('lr', None))
        # wd = self.hparams.get('disc_wd', self.hparams.get('wd', 0))

        self.optimizer = optim.Adam(self.parameters(), lr=self.hparams.disc_lr,
                                    betas=(0.5, 0.99), weight_decay=self.hparams.disc_wd)

    def build(self) -> None:
        # g is the reward net, and h is the shaping net. Combined would be advantage
        self.g = torch.nn.Sequential(
            MinMaxNormalize(min=State.PHI_MIN, max=State.PHI_MAX),
            torch.nn.Linear(State.PHI_DIM, 1),
            Squeeze(dim=-1),
        )
        if self.hparams.shaping:
            self.h = torch.nn.Sequential(
                MinMaxNormalize(min=State.PHI_MIN, max=State.PHI_MAX),
                torch.nn.Linear(State.PHI_DIM, 1),
                Squeeze(dim=-1),
            )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.g(x)

    def cal_loss(self, expert_batch, generator, epoch, step):
        gen_batch = generator.gen_exp(expert_batch)

        # Concatenate
        cur_device = next(self.parameters()).device
        o_next, o = (
            torch.cat([torch.tensor(gen_batch[k], device=cur_device), expert_batch[k]], dim=0)
            for k in ['o_next', 'o'])

        # Add noise to inputs. Linearly decreasing noise
        def add_noise(x):
            if self.hparams.noise == 0 or self.hparams.noise_epochs == 0:
                return x
            n_ratio = max(1 - (epoch / (self.hparams.noise_epochs * self.hparams.epochs)), 0)
            return x + n_ratio * self.hparams.noise * torch.randn_like(x)

        o_next, o = (add_noise(tmp) for tmp in (o_next, o))

        # Do a special AIRL optimization for D
        # Note D = exp(f) / exp(f) + pi
        # Note D's logits is g(s') + (\gamma * h(s') - h(s)) - logpi(s, a, s')
        reward = self.get_reward(o, o_next)

        np_s, np_a = (np.concatenate([gen_batch[k], expert_batch[k].cpu().numpy()])
                      for k in ['s', 'a'])
        prob = generator.get_action_probs(np_s, np_a)

        logpi = torch.from_numpy(np.log(prob + 1e-8)).to(reward.device)
        # D_logit = reward - prob
        ## Should I use AIRL reward design?
        # D_logit = reward
        return self.cal_loss_by_logits(reward, logpi, epoch)

    def cal_loss_by_logits(self, reward, logpi, epoch):
        D_logit = reward
        if self.hparams.airl_obj > 0.:
            coeff = min(epoch / (self.hparams.airl_anneal_epochs * self.hparams.epochs + 1e-8), 1.)
            D_logit -= coeff * logpi

        num_exp = D_logit.shape[0] // 2
        ## (1) expert loss = -log(D(expert))
        eps = self.hparams.label_smoothing
        exp_logp = (1. - eps) * F.logsigmoid(D_logit[num_exp:]) + eps * F.logsigmoid(-D_logit[num_exp:])
        exp_loss = -torch.sum(exp_logp)
        ## (2) generator loss = log(1 - D(gen)). Don't do smoothing
        gen_logp = F.logsigmoid(-D_logit[:num_exp])
        gen_loss = -torch.sum(gen_logp)

        ## Should I do trick with label smoothing?
        # gen_loss = torch.sum(F.logsigmoid(D_logit[:num_exp]))
        # gen_loss = torch.sum(-F.logsigmoid(-D_logit[:num_exp]))
        loss = (exp_loss + gen_loss) / D_logit.shape[0]

        ## Calculate the AUC
        y_true = np.concatenate([np.zeros(num_exp), np.ones(num_exp)]).astype(int)
        y_score = D_logit.detach().cpu().numpy()
        auc = roc_auc_score(y_true, y_score)

        # exp_acc = (D_logit[num_exp:] >= 0).sum() / num_exp
        # gen_acc = (D_logit[:num_exp] < 0).sum() / num_exp]

        return {'loss': loss, 'auc': auc}

    def get_reward(self, s, s_next, dones=0):
        ''' Return the reward of shape B instead of (B, 1) '''
        the_s = s_next if self.hparams.disc_state_time == 'next' else s

        reward = self.g(the_s)
        if self.hparams.shaping:
            shaping = ((1 - dones) * self.hparams.model_gamma * self.h(s_next) - self.h(s))
            reward += shaping

        return reward

    @classmethod
    def get_rs_loader(cls, args, rs=None):
        if rs is None:
            rs = RandomSearch(hparams=args, seed=args.seed)
        rs.add_rs_hparams('disc_lr', short_name='dlr', chose_from=[2e-4, 4e-4])
        # rs.add_rs_hparams('airl_obj', short_name='aob', chose_from=[0, 1])
        rs.add_rs_hparams('shaping', short_name='sh', chose_from=[0, 1])
        # rs.add_rs_hparams('label_smoothing', short_name='ls', chose_from=[0, 0.05])
        return rs

    @classmethod
    def add_model_specific_args(cls, parser) -> argparse.ArgumentParser:
        """
        Adds arguments for DQN model
        Note: these params are fine tuned for Pong env
        Args:
            parent
        """
        # Optimization
        parser.add_argument('--disc_lr', type=float, default=3e-4)
        parser.add_argument('--disc_wd', default=0., type=float)

        # Objective
        parser.add_argument('--shaping', type=int, default=1,
                            help='Model shaping reward or not.')
        parser.add_argument("--label_smoothing", type=float, default=0.)
        parser.add_argument('--airl_obj', type=float, default=0,
                            help='Use AIRL objective or not')
        parser.add_argument('--airl_anneal_epochs', type=float, default=0.5,
                            help='Annealing the airl objective. 0 means no annealing'
                                 'and 10 means increase to 1 in first 10 epochs')
        parser.add_argument("--disc_state_time", type=str, default='next',
                            choices=['current', 'next'])
        return parser

    def temp_callback(self, global_step):
        pass

class NODEGAM_Disc(LinearDisc):
    def build(self) -> None:
        self.choice_fn = nodegam.EM15Temp(
            max_temp=1., min_temp=0.01, steps=self.hparams.anneal_steps)

        the_dict = dict(
            input_dim=State.PHI_DIM,
            layer_dim=self.hparams.num_trees,
            num_layers=self.hparams.num_layers,
            num_classes=1,
            addi_tree_dim=self.hparams.addi_tree_dim,
            depth=self.hparams.depth,
            flatten_output=False,
            choice_function=self.choice_fn,
            bin_function=nodegam.entmoid15,
            output_dropout=self.hparams.output_dropout,
            last_dropout=self.hparams.last_dropout,
            colsample_bytree=self.hparams.colsample_bytree,
            add_last_linear=self.hparams.add_last_linear,
        )

        self.g = torch.nn.Sequential(
            # Normalize the input features
            MinMaxNormalize(min=State.PHI_MIN, max=State.PHI_MAX),
            nodegam.GAMBlock(**the_dict),
        )
        if self.hparams.shaping:
            self.h = torch.nn.Sequential(
                MinMaxNormalize(min=State.PHI_MIN, max=State.PHI_MAX),
                nodegam.GAMBlock(**the_dict),
            )

    @classmethod
    def get_rs_loader(cls, args, rs=None):
        rs = super().get_rs_loader(args, rs=rs)
        rs.add_rs_hparams('num_layers', short_name='nl', chose_from=[1, 2])
        rs.add_rs_hparams('num_trees', short_name='nt', chose_from=[100, 200, 400])
        rs.add_rs_hparams('addi_tree_dim', short_name='ad', chose_from=[0, 1])
        rs.add_rs_hparams('depth', short_name='td', chose_from=[1, 2])
        rs.add_rs_hparams('output_dropout', short_name='od', chose_from=[0., 0.1])
        rs.add_rs_hparams('last_dropout', short_name='ld', chose_from=[0., 0.3])
        rs.add_rs_hparams('colsample_bytree', short_name='cs', chose_from=[1., 0.5])
        rs.add_rs_hparams('anneal_steps', short_name='an', chose_from=[3000])
        # rs.add_rs_hparams('l2_lambda', short_name='la', chose_from=[10, 20])
        return rs

    @classmethod
    def add_model_specific_args(cls, p):
        p = super().add_model_specific_args(p)
        # p.add_argument("--arch", type=str, default='GAM',
        #                choices=['ODST', 'GAM', 'GAMAtt', 'GAMAtt3'])
        p.add_argument("--anneal_steps", type=int, default=500)
        p.add_argument("--num_trees", type=int, default=1024)
        p.add_argument("--num_layers", type=int, default=2)
        p.add_argument("--addi_tree_dim", type=int, default=0)
        p.add_argument("--depth", type=int, default=1)
        p.add_argument("--output_dropout", type=float, default=0.)
        p.add_argument("--l2_lambda", type=float, default=0.)
        p.add_argument("--last_dropout", type=float, default=0.)
        p.add_argument("--colsample_bytree", type=float, default=1.)
        p.add_argument("--add_last_linear", type=int, default=1)
        return p

class FCNN_Disc(LinearDisc):
    def build(self) -> None:
        self.g = torch.nn.Sequential(
            # Normalize the input features
            MinMaxNormalize(min=State.PHI_MIN, max=State.PHI_MAX),
            create_fc_layer(
                n_in=State.PHI_DIM,
                n_layer=self.hparams.disc_n_layer,
                n_hidden=self.hparams.disc_n_hidden,
                n_out=1,
                dropout=self.hparams.disc_dropout,
            ),
            Squeeze(dim=-1),
        )
        if self.hparams.shaping:
            self.h = torch.nn.Sequential(
                # Normalize the input features
                MinMaxNormalize(min=State.PHI_MIN, max=State.PHI_MAX),
                create_fc_layer(
                    n_in=State.PHI_DIM,
                    n_layer=self.hparams.disc_n_layer,
                    n_hidden=self.hparams.disc_n_hidden,
                    n_out=1,
                    dropout=self.hparams.disc_dropout,
                ),
                Squeeze(dim=-1),
            )

    @classmethod
    def get_rs_loader(cls, args, rs=None):
        rs = super().get_rs_loader(args, rs=rs)
        rs.add_rs_hparams('disc_n_layer', short_name='dnl', chose_from=[2, 3, 4, 5])
        rs.add_rs_hparams('disc_n_hidden', short_name='dnh', chose_from=[32, 64, 128, 256])
        rs.add_rs_hparams('disc_dropout', short_name='ddr', chose_from=[0.1, 0.3, 0.5])
        return rs

    @classmethod
    def add_model_specific_args(cls, p):
        p = super().add_model_specific_args(p)
        p.add_argument("--disc_n_layer", type=int, default=2)
        p.add_argument("--disc_n_hidden", type=int, default=64)
        p.add_argument("--disc_dropout", type=float, default=0.5)
        return p

'''
Below are the discriminators for MIMIC3 datasets
'''
class MIMIC3_LinearDisc(LinearDisc):
    def __init__(self, hparams):
        # Add a discriminator type input
        # if 'disc_state_type' not in hparams:
        #     hparams['disc_state_type'] = 'states'

        self.input_dim = HypotensionDataset.get_state_dim_by_type(hparams['disc_state_type'])
        super().__init__(hparams)

    def build(self) -> None:
        # g is the reward net, and h is the shaping net. Combined would be advantage
        self.g = torch.nn.Sequential(
            torch.nn.Linear(self.input_dim, 1),
            Squeeze(dim=-1),
        )
        if self.hparams.shaping:
            self.h = torch.nn.Sequential(
                torch.nn.Linear(self.input_dim, 1),
                Squeeze(dim=-1),
            )

    def cal_loss(self, expert_batch, generator, epoch, step):
        with torch.no_grad(), evaluating(generator):
            gen_s, gen_a, gen_next_s, gen_dones = generator.gen_exp(
                expert_batch, actions='agent', state_type='all',
                same_act_use_exp=self.hparams.same_act_use_exp)

            exp_s, exp_a, exp_next_s, exp_dones = generator.gen_exp(
                expert_batch, actions='expert', state_type='all',
                same_act_use_exp=self.hparams.same_act_use_exp)

        s = torch.cat([gen_s, exp_s], dim=0)
        a = torch.cat([gen_a, exp_a], dim=0)
        next_s = torch.cat([gen_next_s, exp_next_s], dim=0)
        dones = torch.cat([gen_dones, exp_dones], dim=0)

        # Generator takes all features to generate log probs...
        with torch.no_grad(), evaluating(generator):
            logpi = generator.get_action_log_probs(s, a)

        # Add noise to inputs. Linearly decreasing noise
        def add_noise(x):
            if self.hparams.noise == 0 or self.hparams.noise_epochs == 0:
                return x
            n_ratio = max(1 - (epoch / (self.hparams.noise_epochs * self.hparams.epochs + 1e-8)), 0)
            return x + n_ratio * self.hparams.noise * torch.randn_like(x)

        s, next_s = (add_noise(tmp) for tmp in (s, next_s))
        reward = self.get_reward(s, next_s, dones)

        return self.cal_loss_by_logits(reward, logpi, epoch)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.shape[1] != self.input_dim:
            x = HypotensionDataset.extract_s_by_state_type(x, self.hparams.disc_state_type)
        reward = self.g(x)
        return reward

    def get_reward(self, s, s_next, dones=0, shaping=None):
        ''' Return the reward of shape B instead of (B, 1) '''
        s = HypotensionDataset.extract_s_by_state_type(s, self.hparams.disc_state_type)
        s_next = HypotensionDataset.extract_s_by_state_type(s_next, self.hparams.disc_state_type)

        the_s = s_next if self.hparams.disc_state_time == 'next' else s
        reward = self.g(the_s)

        if shaping is None:
            shaping = self.hparams.shaping

        if shaping:
            next_V = self.hparams.model_gamma * self.h(s_next)
            next_V[dones] = 0.
            shaping = (next_V - self.h(s))
            reward += shaping

        return reward

    def eval_step(self, batch, batch_idx):
        states, actions, next_states, dones = HypotensionDataset.extract_s_and_a_pairs(batch['x_list'])

        D_logit = self.get_reward(states, next_states, dones)

        exp_logp = F.logsigmoid(D_logit)
        exp_loss = -torch.sum(exp_logp)
        exp_acc = torch.sum((D_logit > 0).float())
        return {f'exp_loss': exp_loss, 'exp_acc': exp_acc, f'total': states.shape[0]}

    def eval_epoch_end(self, outputs, prefix='val'):
        """Log the avg of the test results"""
        exp_loss = sum([x['exp_loss'] for x in outputs])
        exp_acc = sum([x['exp_acc'] for x in outputs])
        total = sum([x['total'] for x in outputs])

        avg_exp_loss = exp_loss / total
        avg_exp_acc = exp_acc / total

        logs = {f'{prefix}_exp_loss': avg_exp_loss, f'{prefix}_exp_acc': avg_exp_acc}
        return logs

    @classmethod
    def get_rs_loader(cls, args, rs=None):
        if rs is None:
            rs = RandomSearch(hparams=args, seed=args.seed)
        rs.add_rs_hparams('disc_lr', short_name='dlr', chose_from=[5e-4, 8e-4, 1e-3])
        # rs.add_rs_hparams('disc_wd', short_name='dwd', chose_from=[0.])
        # rs.add_rs_hparams('shaping', short_name='shp', chose_from=[0, 1])
        rs.add_rs_hparams('airl_obj', short_name='aob', chose_from=[0, 1])

        rs.add_rs_hparams('label_smoothing', short_name='ls', chose_from=[0.005, 0.01])
        # rs.add_rs_hparams('same_act_use_exp', short_name='sa', chose_from=[0, 0.5, 1])
        # rs.add_rs_hparams('disc_state_type', short_name='', chose_from=['features'])
        return rs

    @classmethod
    def add_model_specific_args(cls, parser) -> argparse.ArgumentParser:
        parser = super().add_model_specific_args(parser)
        parser.add_argument("--disc_state_type", type=str, default='features',
                            choices=['features', 'states', 'all'])
        parser.add_argument("--same_act_use_exp", type=float, default=1.,
                            help='If 1, use orig exp to distinguish. If 0, use gru pred. '
                                 'It may increase difficulty.')
        return parser

    def temp_callback(self, global_step):
        pass

class MIMIC3_NODEGAM_Disc(MIMIC3_LinearDisc):
    def build(self) -> None:
        self.choice_fn = nodegam.EM15Temp(
            max_temp=1., min_temp=0.01, steps=self.hparams.anneal_steps)

        the_dict = dict(
            input_dim=self.input_dim,
            layer_dim=self.hparams.num_trees,
            num_layers=self.hparams.num_layers,
            num_classes=1,
            addi_tree_dim=self.hparams.addi_tree_dim,
            depth=self.hparams.depth,
            flatten_output=False,
            choice_function=self.choice_fn,
            bin_function=nodegam.entmoid15,
            output_dropout=self.hparams.output_dropout,
            last_dropout=self.hparams.last_dropout,
            colsample_bytree=self.hparams.colsample_bytree,
            add_last_linear=self.hparams.add_last_linear,
            ga2m=self.hparams.ga2m,
        )
        if self.hparams.disc_arch.startswith('GAMAtt'):
            the_dict['dim_att'] = self.hparams.dim_att

        arch_cls = getattr(nodegam, self.hparams.disc_arch + 'Block')
        self.g = arch_cls(**the_dict)
        if self.hparams.shaping:
            self.h = arch_cls(**the_dict)

    @classmethod
    def get_rs_loader(cls, args, rs=None):
        rs = super().get_rs_loader(args, rs=rs)
        # rs.add_rs_hparams('airl_anneal_epochs', short_name='ane',
        #                   gen=lambda hparams: -1 if hparams.airl_obj == 0
        #                   else rs.np_gen.choice([0.5]))
        # rs.add_rs_hparams('ga2m', short_name='ga2m', chose_from=[0, 1])
        rs.add_rs_hparams('disc_arch', short_name='a', chose_from=['GAM'])
        rs.add_rs_hparams('num_layers', short_name='nl', chose_from=[1, 2, 3])
        rs.add_rs_hparams('num_trees', short_name='nt',
                          chose_from=[200, 300])
                          # gen=(lambda args: rs.np_gen.choice(
                          #      [500, 1000] if args.num_layers >= 4 else [500, 1000])))
        rs.add_rs_hparams('addi_tree_dim', short_name='ad', #chose_from=[0, 1],
                          gen=lambda args: rs.np_gen.choice([0, 1]) if args.num_layers > 1 else 0
                          )
        rs.add_rs_hparams('depth', short_name='td', #chose_from=[2, 3],
                          gen=lambda args: rs.np_gen.choice([2, 3, 4])
                          )
        rs.add_rs_hparams('last_dropout', short_name='ld', chose_from=[0.3, 0.5])
        rs.add_rs_hparams('colsample_bytree', short_name='cs', chose_from=[0.5])
        rs.add_rs_hparams('anneal_steps', short_name='an', chose_from=[3000])
        rs.add_rs_hparams('l2_lambda', short_name='la', chose_from=[0])
        rs.add_rs_hparams('add_last_linear', short_name='ll',
                          gen=lambda args: 0 if args.disc_arch == 'GAMAtt' else rs.np_gen.choice([1]))
        rs.add_rs_hparams('output_dropout', short_name='od',
                          gen=lambda args: 0 if args.add_last_linear == 0 else rs.np_gen.choice([0.1, 0.2]))
        # rs.add_rs_hparams('dim_att', short_name='da',
        #                   gen=lambda args: rs.np_gen.choice([8, 16, 32])
        #                   if args.disc_arch.startswith('GAMAtt') else -1)
        return rs

    @classmethod
    def add_model_specific_args(cls, p):
        p = super().add_model_specific_args(p)
        p.add_argument("--disc_arch", type=str, default='GAM',
                       choices=['ODST', 'GAM', 'GAMAtt', 'GAMAtt3'])
        p.add_argument("--anneal_steps", type=int, default=500)
        p.add_argument("--num_trees", type=int, default=128)
        p.add_argument("--num_layers", type=int, default=2)
        p.add_argument("--addi_tree_dim", type=int, default=0)
        p.add_argument("--depth", type=int, default=2)
        p.add_argument("--output_dropout", type=float, default=0.)
        p.add_argument("--l2_lambda", type=float, default=0.)
        p.add_argument("--last_dropout", type=float, default=0.)
        p.add_argument("--colsample_bytree", type=float, default=1.)
        p.add_argument("--add_last_linear", type=int, default=1)
        p.add_argument("--ga2m", type=int, default=1)
        p.add_argument("--dim_att", type=int, default=8)
        return p

class MIMIC3_FCNN_Disc(MIMIC3_LinearDisc):
    def build(self) -> None:
        self.g = torch.nn.Sequential(
            # Normalize the input features
            create_fc_layer(
                n_in=self.input_dim,
                n_layer=self.hparams.disc_n_layer,
                n_hidden=self.hparams.disc_n_hidden,
                n_out=1,
                dropout=self.hparams.disc_dropout,
            ),
            Squeeze(dim=-1),
        )
        if self.hparams.shaping:
            self.h = torch.nn.Sequential(
                # Normalize the input features
                create_fc_layer(
                    n_in=self.input_dim,
                    n_layer=self.hparams.disc_n_layer,
                    n_hidden=self.hparams.disc_n_hidden,
                    n_out=1,
                    dropout=self.hparams.disc_dropout,
                ),
                Squeeze(dim=-1),
            )

    @classmethod
    def get_rs_loader(cls, args, rs=None):
        rs = super().get_rs_loader(args, rs=rs)
        rs.add_rs_hparams('disc_n_layer', short_name='dnl', chose_from=[2, 3, 4, 5])
        rs.add_rs_hparams('disc_n_hidden', short_name='dnh', chose_from=[32, 64, 128, 256])
        rs.add_rs_hparams('disc_dropout', short_name='ddr', chose_from=[0.1, 0.3, 0.5])
        return rs

    @classmethod
    def add_model_specific_args(cls, p):
        p = super().add_model_specific_args(p)
        p.add_argument("--disc_n_layer", type=int, default=2)
        p.add_argument("--disc_n_hidden", type=int, default=64)
        p.add_argument("--disc_dropout", type=float, default=0.5)
        return p


""" script for generating samples from a trained model """

import argparse
import os
from math import sqrt

import cv2
import numpy as np
from scipy.ndimage import gaussian_filter
from tqdm import tqdm

import torch as th

# turn fast mode on
from utils import adjust_dynamic_range
from torch.backends import cudnn

cudnn.benchmark = True

# define the device for the training script
device = th.device("cuda" if th.cuda.is_available() else "cpu")

def parse_arguments():
    """
    command line arguments parser
    :return: args => parsed command line arguments
    """
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--generator_file",
        action="store",
        type=str,
        help="pretrained weights file for generator",
        required=True,
    )

    parser.add_argument(
        "--latent_size",
        action="store",
        type=int,
        default=512,
        help="latent size for the generator",
    )

    parser.add_argument(
        "--depth",
        action="store",
        type=int,
        default=9,
        help="latent size for the generator",
    )

    parser.add_argument(
        "--out_depth",
        action="store",
        type=int,
        default=6,
        help="output depth of images. **Starts from 0",
    )

    parser.add_argument(
        "--time",
        action="store",
        type=float,
        default=300,
        help="Number of seconds for the video to make",
    )

    parser.add_argument(
        "--fps",
        action="store",
        type=int,
        default=60,
        help="Frames per second in the video",
    )

    parser.add_argument(
        "--smoothing",
        action="store",
        type=float,
        default=0.75,
        help="Smoothing amount in transition frames",
    )

    parser.add_argument(
        "--out_dir",
        action="store",
        type=str,
        default="interp_animation_frames/",
        help="path to the output directory for the frames",
    )

    parser.add_argument(
        "--video_only",
        action="store_true",
        help="Pass this to skip saving of individual frames.",
    )

    parser.add_argument(
        "--video_name", action="store", type=str, default="", help="Filename of video"
    )

    args = parser.parse_args()

    return args

def get_image(gen, point, depth, alpha):
    image = gen(point, depth, alpha).detach()
    image = adjust_dynamic_range(image).squeeze(dim=0)
    return image.cpu().numpy().transpose(1, 2, 0)

def main(args):
    """
    Main function of the script
    :param args: parsed commandline arguments
    :return: None
    """
    from networks import Generator

    # create generator object:
    print("Creating a generator object ...")
    generator = th.nn.DataParallel(
        Generator(depth=args.depth, latent_size=args.latent_size).to(device)
    )

    # load the trained generator weights
    print("loading the trained generator weights ...")
    generator.load_state_dict(th.load(args.generator_file, str(device)))

    # total_frames in the video:
    total_frames = int(args.time * args.fps)

    # Let's create the animation video from the latent space interpolation
    # all latent vectors:
    all_latents = th.randn(total_frames, args.latent_size).to(device)
    all_latents = gaussian_filter(all_latents.cpu(), [args.smoothing * args.fps, 0])
    all_latents = th.from_numpy(all_latents)
    all_latents = (
        all_latents / all_latents.norm(dim=-1, keepdim=True) * (sqrt(args.latent_size))
    )

    # create output directory
    os.makedirs(args.out_dir, exist_ok=True)

    global_frame_counter = 1

    # If we're saving a video, make the video object
    out_file, video_out = None, None
    if args.video_name:
        width = 2 ** (args.depth + 1)
        out_file = os.path.join(args.out_dir, args.video_name)
        video_out = cv2.VideoWriter(
            args.video_name, cv2.VideoWriter_fourcc(*"mp4v"), args.fps, (width, width)
        )

    # Run the main loop for the interpolation:
    print("Generating the video frames ...")
    for latent in tqdm(all_latents):
        latent = th.unsqueeze(latent, dim=0)

        # generate the image for this point:
        img = get_image(generator, latent, args.out_depth, 1) * 255

        if not args.video_only:
            cv2.imwrite(
                os.path.join(
                    args.out_dir, "frames", "{:05d}.png".format(global_frame_counter)
                ),
                img,
            )

        # Make an image of unsigned 8-bit integers for OpenCV
        if args.video_name:
            img_int = img.astype(np.uint8)
            video_out.write(img_int)

        # Increment the counter
        global_frame_counter += 1

    # video frames have been generated
    if not args.video_only:
        print("Video frames have been generated at:", args.out_dir)

    if args.video_name:
        print("Video saved to {}".format(out_file))
        video_out.release()

if __name__ == "__main__":
    main(parse_arguments())