# adapted from detectron2.structures.masks.py (v0.5)
# Copyright (c) Facebook, Inc. and its affiliates.

import copy
import itertools
import numpy as np
from typing import Any, Iterator, List, Union
import torch

from detectron2.structures import Boxes
from detectron2.layers.roi_align import ROIAlign
from detectron2.structures import ROIMasks, PolygonMasks
from detectron2.utils.memory import retry_if_cuda_oom

class BitMasksMulticlass:
    """
    This class stores the segmentation masks for all objects in one image, in
    the form of bitmaps.

    Attributes:
        tensor: bool Tensor of N,C,H,W, representing N instances in the image.
    """

    def __init__(self, tensor: Union[torch.Tensor, np.ndarray]):
        """
        Args:
            tensor: bool Tensor of N,H,W, representing N instances in the image.
        """
        device = tensor.device if isinstance(tensor, torch.Tensor) else torch.device("cpu")

        if type(tensor) == torch.Tensor:
            self.image_size = tensor.shape[2:]
        else:
            tensor = torch.as_tensor(tensor, dtype=torch.uint8, device=device)
            self.image_size = tensor.shape[1:]

        self.tensor = tensor

    @torch.jit.unused
    def to(self, *args: Any, **kwargs: Any) -> "BitMasksMulticlass":
        return BitMasksMulticlass(self.tensor.to(*args, **kwargs))

    @property
    def device(self) -> torch.device:
        return self.tensor.device

    @torch.jit.unused
    def __getitem__(self, item: Union[int, slice, torch.ByteTensor]) -> "BitMasksMulticlass":
        if isinstance(item, int):
            return BitMasksMulticlass(self.tensor[item].view(1, -1))
        m = self.tensor[item]

        return BitMasksMulticlass(m)

    @torch.jit.unused
    def __iter__(self) -> torch.Tensor:
        yield from self.tensor

    @torch.jit.unused
    def __repr__(self) -> str:
        s = self.__class__.__name__ + "("
        s += "num_instances={})".format(len(self.tensor))
        return s

    def __len__(self) -> int:
        return self.tensor.shape[0]

    def nonempty(self) -> torch.Tensor:
        return self.tensor.flatten(1).any(dim=1)

    def crop_and_resize(self, boxes: torch.Tensor, mask_size: int) -> torch.Tensor:
        """
        Crop each bitmask by the given box, and resize results to (mask_size, mask_size).
        This can be used to prepare training targets for Mask R-CNN.
        It has less reconstruction error compared to rasterization with polygons.
        However we observe no difference in accuracy,
        but BitMasks requires more memory to store all the masks.

        Args:
            boxes (Tensor): Nx4 tensor storing the boxes for each mask
            mask_size (int): the size of the rasterized mask.

        Returns:
            Tensor:
                A bool tensor of shape (N, mask_size, mask_size), where
                N is the number of predicted boxes for this image.
        """
        assert len(boxes) == len(self), "{} != {}".format(len(boxes), len(self))
        device = self.tensor.device

        batch_inds = torch.arange(len(boxes), device=device).to(dtype=boxes.dtype)[:, None]
        rois = torch.cat([batch_inds, boxes], dim=1)  # Nx5

        bit_masks = self.tensor.to(dtype=torch.float32)
        rois = rois.to(device=device)
        output = (
            ROIAlign((mask_size, mask_size), 1.0, 0, aligned=True)
            .forward(bit_masks[:, None, :, :], rois)
            .squeeze(1)
        )

        return output

    def get_bounding_boxes(self) -> Boxes:
        """
        Returns:
            Boxes: tight bounding boxes around bitmasks.
            If a mask is empty, it's bounding box will be all zero.
        """
        boxes = torch.zeros(self.tensor.shape[0], 4, dtype=torch.float32)
        x_any = torch.any(self.tensor, dim=1)
        y_any = torch.any(self.tensor, dim=2)
        for idx in range(self.tensor.shape[0]):
            x = torch.where(x_any[idx, :])[0]
            y = torch.where(y_any[idx, :])[0]
            if len(x) > 0 and len(y) > 0:
                boxes[idx, :] = torch.as_tensor(
                    [x[0], y[0], x[-1] + 1, y[-1] + 1], dtype=torch.float32
                )
        return Boxes(boxes)

    @staticmethod
    def cat(bitmasks_list: List["BitMasksMulticlass"]) -> "BitMasksMulticlass":
        """
        Concatenates a list of BitMasks into a single BitMasks

        Arguments:
            bitmasks_list (list[BitMasks])

        Returns:
            BitMasks: the concatenated BitMasks
        """
        assert isinstance(bitmasks_list, (list, tuple))
        assert len(bitmasks_list) > 0
        assert all(isinstance(bitmask, BitMasksMulticlass) for bitmask in bitmasks_list)

        cat_bitmasks = type(bitmasks_list[0])(torch.cat([bm.tensor for bm in bitmasks_list], dim=0))
        return cat_bitmasks

class ROIMasksMulticlass(ROIMasks):
    """
    Represent masks by N smaller masks defined in some ROIs. Once ROI boxes are given,
    full-image bitmask can be obtained by "pasting" the mask on the region defined
    by the corresponding ROI box.
    """

    def __init__(self, tensor: torch.Tensor):
        """
        Args:
            tensor: (N, M, M) mask tensor that defines the mask within each ROI.
        """
        if tensor.dim() != 4:
            raise ValueError("ROIMasksMulticlass must take a masks of 4 dimension.")
        self.tensor = tensor

    def to(self, device: torch.device) -> "ROIMasksMulticlass":
        return ROIMasksMulticlass(self.tensor.to(device))

    @property
    def device(self) -> "device":
        return self.tensor.device

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

    def __getitem__(self, item) -> "ROIMasksMulticlass":
        """
        Returns:
            ROIMasksMulticlass: Create a new :class:`ROIMasksMulticlass` by indexing.

        The following usage are allowed:

        1. `new_masks = masks[2:10]`: return a slice of masks.
        2. `new_masks = masks[vector]`, where vector is a torch.BoolTensor
           with `length = len(masks)`. Nonzero elements in the vector will be selected.

        Note that the returned object might share storage with this object,
        subject to Pytorch's indexing semantics.
        """
        t = self.tensor[item]
        if t.dim() != 4:
            raise ValueError(
                f"Indexing on ROIMasks with {item} returns a tensor with shape {t.shape}!"
            )
        return ROIMasksMulticlass(t)

    @torch.jit.unused
    def __repr__(self) -> str:
        s = self.__class__.__name__ + "("
        s += "num_instances={})".format(len(self.tensor))
        return s

    @torch.jit.unused
    def to_bitmasks(self, boxes: torch.Tensor, height, width, threshold=0.5):
        """
        Args:

        """
        from .ops import paste_masks_in_image

        paste = retry_if_cuda_oom(paste_masks_in_image)
        bitmasks = paste(
            self.tensor,
            boxes,
            (height, width),
            threshold=threshold,
        )
        return BitMasksMulticlass(bitmasks)


# -*- coding:utf-8 -*-
"""
This is a class that is used for creating auto-ml object that
we could use for training step.

author: Guangqiang.lu
"""

import numpy as np
import pandas as pd
import time
from sklearn.base import BaseEstimator
from sklearn.exceptions import NotFittedError
from sklearn.model_selection import train_test_split

from auto_ml.utils.backend_obj import Backend
from auto_ml.base.classifier_algorithms import *
from auto_ml.metrics.scorer import *
from auto_ml.utils.CONSTANT import *
from auto_ml.pipelines.pipeline_training import ClassificationPipeline,  RegressionPipeline
from auto_ml.utils.logger import create_logger

logger = create_logger(__file__)

class AutoML(BaseEstimator):
    def __init__(self, models_path=None,
                 time_left_for_this_task=3600,
                 n_ensemble=10,
                 n_best_model=5,
                 include_estimators=None,
                 exclude_estimators=None,
                 include_preprocessors=None,
                 exclude_preprocessors=None,
                 keep_models=True,
                 model_dir=None,
                 precision=32,
                 delete_models=True
                 ):
        """
        this is to init automl class, whole thing should be ininstanted
        in this class, like what algorithms to use, how many models to
        be selected, etc.
        :param backend: backend object used to save and load models
        :param time_left_for_this_task: how long for this models to be trained.
        :param n_ensemble: how many models to be selected to be ensemble
        :param n_best_model: how many models to be keeped during training.
        :param include_estimators: what algorithms to be included
        :param exclude_estimators: what algorithms to be excluded
        :param include_preprocessors: what preprocessing step to be included
        :param exclude_preprocessors: what preprocessing step to be excluded
        :param keep_models: whether or not to keep trained models
        :param model_dir: keep model folder, if None use backend to create one folder
        :param precision: precision of data, to save memory
        """
        super(AutoML, self).__init__()
        # Change backend object creation as this class is entry class also Backend class support singleton, so output_folder will work.
        models_path = OUTPUT_FOLDER if models_path is None else models_path
        # Add with `models_path` parameters, as the only output for the framework is only models files, so to store models into
        # folder that we would like will be nice!
        self.models_path = models_path

        self.backend = Backend(output_folder=models_path)
        self.time_left_for_this_taks = 3600 if time_left_for_this_task is None else time_left_for_this_task
        self.n_ensemble = n_ensemble
        self.n_best_model = n_best_model
        self.include_estimators = include_estimators
        self.exclude_estimators = exclude_estimators
        self.include_preprocessors = include_preprocessors
        self.exclude_preprocessors = exclude_preprocessors
        self.keep_models = keep_models
        self.model_dir = model_dir
        self.precision = precision

        self.estimator = None

        # as we will use `ensemble` to combine models so the last will just be one model
        self.best_model = None

        # Add with what type of the problem
        self.type_of_problem = None

        self.delete_models = delete_models

    def fit(self, x=None, y=None, file_load=None, \
             xval=None, yval=None, val_split=None, n_jobs=None, use_neural_network=True, *args, **kwargs):
        """
        Type of the problem attribute should be added, so that for `score`,
        we could get metrics based on different problem.
        :param xtrain:
        :param ytrain:
        :param n_jobs:
        :return:
        """
        start_time = time.time()

        # Add logic here is that we should clean models' folder, so that every time we could get a clean folder for next time running!
        if self.backend and self.delete_models:
            self.backend.clean_folder()

        x, y = self._get_data_and_label(file_load, x, y)

        if val_split is not None:
            # if do need to do validation based on current train data, then just split current data into validation as well
            x, xval, y, yval = train_test_split(x, y, test_size=val_split)
        else:
            # Here I think if and only if the data length is over a threashold, then to do validation,even user haven't provided val data.
            if len(x) > VALIDATION_THRESHOLD:
                val_split = .2
                # if do need to do validation based on current train data, then just split current data into validation as well
                x, xval, y, yval = train_test_split(x, y, test_size=val_split)

        # after the checking process, then we need to create the Pipeline for whole process.
        # Here should use a Pipeline object to do real training, also with `ensemble`
        self.estimator.fit(x, y, n_jobs=n_jobs, use_neural_network=use_neural_network)

        # load trained models for prediction and scoring for testing data.
        # after we have fitted the trained models, then next step is to load whole of them from disk
        # and sort them based on score, and get scores for `test data` and `test label`
        # WE could get them from parent function, so that we could also use this for `regression`

        # TODO: This should be changed for regression! As regression may have negative values, so here should change
        # on different type.
        # After change regression metrics to r2, this is workable
        self.models_list = self._load_trained_models_ordered_by_score(higher_best=True)

        self._validation_models(xval, yval)

        logger.info("Whole training pipeline takes: {} seconds!".format(round(time.time() - start_time, 2)))

        return self

    def predict(self, x=None, file_load=None, **kwargs):
        """
        Most of the prediction logic should happen here, as for the score should based on
        prediction result.

        :param x:
        :param kwargs:
        :return:
        """
        x = self._get_training_data(x, file_load)

        logger.info("Start to get prediction based on best trained model.")
        prediction = self.estimator.predict(x)
        logger.info("Prediction step finishes.")

        return prediction

    def predict_proba(self, x=None, file_load=None, **kwargs):
        """
        Should support with probability if supported.
        :param x:
        :param kwargs:
        :return:
        """
        # should check the estimator should have function: `predict_proba`!
        if not hasattr(self.estimator, 'predict_proba'):
            raise NotImplementedError("Best fitted model:{} doesn't support `predict_proba`".format(self.best_model))

        x = self._get_training_data(x, file_load)

        logger.info("Start to get probability based on best trained model.")
        prob = self.estimator.predict_proba(x)
        logger.info("Prediction step finishes.")

        return prob

    def score(self, x=None, y=None, file_load=None, **kwargs):
        self._check_fitted()

        logger.info("Start to get prediction based on best trained model!")

        self._check_param(file_load, x, y)

        if file_load is not None:
            x, y = self.__get_file_load_data_label(file_load, use_for_pred=False)

        # ensure we have array type
        if isinstance(x, pd.DataFrame) or isinstance(x, pd.Series):
            x = x.values
            if len(x.shape) == 1:
                x = x.reshape(-1, 1)
        if isinstance(y, pd.DataFrame) or isinstance(y, pd.Series):
            y = y.values

            if len(y.shape) == 1:
                y = y.reshape(-1, 1)

        # Use child func, child should implement with score based on different type of problem.
        score = self.estimator.score(x, y)

        logger.info("Get score: {} based on best model!".format(score))

        return score

    def _load_trained_models_ordered_by_score(self, higher_best=True):
        """
        To load whole trained model from disk and sorted them based on `higher_best`.
        :param higher_best:
        :return:
            models_list: a list of trained models sorted with score. [("lr-0.9.pkl", lr-0.9.pkl), ...]
        """
        models_list = self.backend.load_models_combined_with_model_name()

        # let's sort these models based on score, this is a tuple:(model_name, model_instance)
        if higher_best:
            reverse = True
        else:
            reverse = False

        # sort models by `training` score with diff type of problem
        models_list = sorted(models_list,
                             key=lambda model: float(model[0].split("_")[1].replace(".pkl", '').replace(".h5", '')),
                             reverse=reverse)

        return models_list

    def get_sorted_models_scores(self, x, y, **kwargs):
        """
        Add this func to get whole score based for each trained models, so that
        we could get the result that we have taken that times and for each models,
        how about the testing result.

        But this should be implemented by pipeline!
        :param x:
        :param y:
        :param kwargs:
        :return:
            a list of tuple: [(model_name, model_score), ...]
        """
        raise NotImplementedError

    def _check_fitted(self):
        if not self.models_list:
            logger.error("When try to get prediction with `automl`, Please fit the model first")
            raise NotFittedError("When try to get prediction with `automl`, Please fit the model first")

        return True

    @staticmethod
    def _check_param(*args):
        all_None = all([arg is None for arg in args])

        if all_None:
            raise ValueError("Please provide at least one parameter!")

    @staticmethod
    def __get_file_load_data_label(file_load, use_for_pred=True):
        """
        Get data and label from original obj.
        """
        data, label = file_load.data, file_load.label

        return data, label

    @classmethod
    def reconstruct(cls, models_path=None, *args, **kwargs):
        """
        Used for Restful API to create
        """
        return cls(models_path, *args, **kwargs)

    @staticmethod
    def _get_data_and_label(file_load, x, y):
        """
        Ensure could get data and label.
        """
        if file_load is None and x is None and y is None:
            raise ValueError("When do real training, please provide at least a " +
                 "`file_load` or train data with `xtrain, ytrain`!")

        if file_load is not None:
            # with container, then just query the attribute then we could keep other as same.
            x, y = file_load.data, file_load.label
        else:
            if x is None or y is None:
                raise ValueError("When to do training, please provide both `x` and `y`!")

        # Let's try to make DF into array, so later will be easier!
        if isinstance(x, pd.DataFrame):
            x = x.values
        if isinstance(y, pd.DataFrame):
            y = y.values

        return x, y

    def _validation_models(self, xval, yval):
        if xval is not None and yval is not None:
            score_dict = self.estimator.get_sorted_models_scores(xval, yval)
            print(score_dict)

            score_log_str = self.__format_trained_model_scores(score_dict)
            print(score_log_str)
        else:
            logger.warning("No need to validation!")
            pass

    @staticmethod
    def __format_trained_model_scores(score_dict, n_space=35):
        out_str_format = '{{0:{0}}}|{{1:{0}}}|{{2:{0}}}'.format(n_space)

        score_log_str = out_str_format.format("Model name", "Train score", "Validation score")
        logger.info(score_log_str)

        for model_name, test_score in score_dict.items():
            try:
                model_name_split = model_name.split('_')
                # model_name = model_name_split[0]
                train_score = model_name_split[1]
                train_score = train_score[:train_score.rindex(".")]
                # must convert to str, otherwise with not wanted result.
                test_score = str(test_score)

                log_str = out_str_format.format(model_name, train_score, test_score)
                score_log_str += '\n' + log_str
            except Exception as e:
                logger.warning("Get invalidate model name: {}".format(model_name))
                continue

            logger.info(log_str)

        return score_log_str

    def get_sorted_models_scores(self, xtest=None, ytest=None, file_load=None, reverse=True, **kwargs):
        """
        To get some best trained model's score for `test` data with ordered.

        So that we could get the list of the best scores for later front end show case.
        :param x:
        :param y:
        :param kwargs:
        :return:
        """
        self._check_param(file_load, xtest, ytest)

        if file_load is not None:
            xtest, ytest = file_load.data, file_load.label

        score_dict = self.estimator.get_sorted_models_scores(xtest, ytest, reverse=reverse)

        return score_dict

    def _get_training_data(self, x, file_load):
        """Get training data based on `file_load` or `x`,
        if x is provided, then just return x, otherwise should get data
        from file_load.

        Either of thems should be provided.

        Args:
            x ([type]): [description]
            file_load ([type]): [description]

        Returns:
            [type]: [description]
        """
        self._check_param(file_load, x)

        if x is not None:
            if isinstance(x, pd.DataFrame):
                x = x.values
            return x
        elif file_load is not None:
            x, _ = self.__get_file_load_data_label(file_load)
            return x

class ClassificationAutoML(AutoML):
    def __init__(self, models_path=None,
            include_estimators=None,
            exclude_estimators=None,
            include_preprocessors=None,
            exclude_preprocessors=None,
            **kwargs):
        """Added with algorithm selection and processing selection, even with others in case we need.

        Args:
            models_path (Str, optional): Where to store our models. Defaults to None.
        """
        super(ClassificationAutoML, self).__init__(models_path=models_path,
                                            include_estimators=include_estimators,
                                            exclude_estimators=exclude_estimators,
                                            include_preprocessors=include_preprocessors,
                                            exclude_preprocessors=exclude_preprocessors, **kwargs)
        # after pipeline has finished, then we should use `ensemble` to combine these models
        # action should happen here.
        self.estimator = ClassificationPipeline(backend=self.backend,
                                            include_estimators=include_estimators,
                                            exclude_estimators=exclude_estimators,
                                            include_preprocessors=include_preprocessors,
                                            exclude_preprocessors=exclude_preprocessors,
                                            **kwargs)

class RegressionAutoML(AutoML):
    def __init__(self, models_path=None,
            include_estimators=None,
            exclude_estimators=None,
            include_preprocessors=None,
            exclude_preprocessors=None,
            **kwargs):
        """Added with algorithm selection and processing selection, even with others in case we need.

        Args:
            models_path (Str, optional): Where to store our models. Defaults to None.
        """
        super(RegressionAutoML, self).__init__(models_path=models_path,
                                            include_estimators=include_estimators,
                                            exclude_estimators=exclude_estimators,
                                            include_preprocessors=include_preprocessors,
                                            exclude_preprocessors=exclude_preprocessors, **kwargs)

        # after pipeline has finished, then we should use `ensemble` to combine these models
        # action should happen here.
        self.estimator = RegressionPipeline(backend=self.backend,
                                            include_estimators=include_estimators,
                                            exclude_estimators=exclude_estimators,
                                            include_preprocessors=include_preprocessors,
                                            exclude_preprocessors=exclude_preprocessors,
                                            **kwargs)

class FileLoad:
    """Load data from file, support with local file also with GCS.

    Make this class as a container for later use case.
    """
    def __init__(self, file_name, file_path=None, file_sep=',', label_name='label', use_for_pred=False,
            service_account_file_name=None, service_account_file_path=None, except_columns=None):
        """Main container for file-like dataset.

        Args:
            file_name (str): Name of file
            label_name (str, optional): What is `label` column's name?. Defaults to 'label'.
            file_path (str, optional): Where file located?. Defaults to None.
            file_sep (str, optional): File seprator. Defaults to ','.
            use_for_pred (Boolean, optional): Whether to use this for prediction?
                Noted: If file doesn't contain label column, do need set this parameter to `True`. Defaults to False.
            service_account_file_name (str, optional): SA file name. Defaults to None.
            service_account_file_path (str, optional): SA file path. Defaults to None.
            except_columns (List, optional): Columns are needed to be used. Defaults to None.

        Raises:
            ValueError: [description]
        """
        if not file_name.endswith('csv'):
            raise ValueError("Currently only support CSV file! Please provide with a CSV file.")

        self.file_name = file_name
        self.file_sep = file_sep
        self.file_path = file_path if file_path is not None else '.'
        self.label_name = label_name
        self.use_for_pred = use_for_pred
        self.service_account_file_name = service_account_file_name
        self.service_account_file_path = service_account_file_path
        if except_columns is not None and isinstance(except_columns, str):
            except_columns = [except_columns]
        self.except_columns = except_columns
        self.data, self.label = self._load_data_file()

        self.columns = None

    def _get_file_location(self):
        if self.file_path is not None:
            if self.file_path.startswith("gs://"):
                return 'gcs'
            else:
                return 'local'

    def _get_gcs_file(self):
        # first service account
        self._check_service_account()

        bucket_name = self.file_path.split('/')[2]
        if bucket_name.find(".") > 0:
            # file_path also contain file name with `.`
            logger.error("When to get bucket name from file path: {}, couldn't get the bucket name.".format(self.file_path))
            raise ValueError("When to get bucket name from file path: {}, couldn't get the bucket name.".format(self.file_path))

        from google.cloud import storage

        # init client with service account file
        client = storage.Client.from_service_account_json(os.path.join(self.service_account_file_path, self.service_account_file_name))

        bucket = client.bucket(bucket_name)
        blob = bucket.blob(self.file_name)

        try:
            # download as a file into temperate folder
            blob.download_to_filename(os.path.join(TMP_FOLDER, self.file_name))
            logger.info("File: {} has been downloaded from GCS to local.".format(self.file_name))

        except Exception as e:
            logger.error("Error downloading file with service account: {}".format(self.service_account_file_name))
            raise e

    def _check_service_account(self):
        # Check service account related
        if not self.service_account_file_name:
            raise ValueError("When try to use GCS, please must provide with service account file to interact with GCS!")

        if not self.service_account_file_name.endswith('json'):
            raise ValueError("When try to use GCS, Service acount file must be a JSON file, " +\
                 "but provided with {}".format(self.service_account_file_name))

        if not os.path.exists(os.path.join(self.service_account_file_path, self.service_account_file_name)):
            raise FileNotFoundError("When try to find service account file couldn't find file: {} in path: {}, " +
                    "please check it.".format(self.service_account_file_name, self.service_account_file_path))

    def _load_data_file(self):
        file_location = self._get_file_location()

        file_path = self.file_path

        if file_location == 'gcs':
            logger.info("Start to download file:{} from GCS.".format(self.file_name))
            self._get_gcs_file()
            file_path = TMP_FOLDER
            logger.info("Download file from GCS has finished.")

        try:
            df = pd.read_csv(os.path.join(file_path, self.file_name), sep=self.file_sep)

            df_cols = df.columns
            self.columns = df_cols

            # In case that we don't need some columns and these cols should be in DF columns
            if self.except_columns is not None:
                except_common_cols = set(self.except_columns).intersection(list(df.columns))

                logger.info("Need to delete columns:{}  from original DataFrame.".format('\t'.join(list(except_common_cols))))

                if except_common_cols:
                    df.drop(except_common_cols, axis=1, inplace=True)

            if self.use_for_pred:
                # if just for prediction, then we don't get label data
                data = df.values
                label = None

                return data, label

            # let's inference data and label from DataFrame

            if self.label_name not in df_cols:
                # if `label` col not in the original file, then make let first column as `label` column
                label = df.iloc[:, 0]
                data = df.iloc[:, 1:]
            else:
                label = df[self.label_name]
                data = df.drop([self.label_name], axis=1)

            # convert into array type will be easier
            data, label = data.values, label.values.reshape(-1, 1)

            return data, label
        except IOError as e:
            logger.error("Error load data file: {} from path: {}".format(self.file_name, file_path))
            raise e

if __name__ == '__main__':
    # Test with `FileLoad` class
    file_name = 'train.csv'
    file_path = r"C:\Users\guangqiiang.lu\Documents\lugq\real_code\auto_ml\auto_ml\test"
    # file_path = "gs://cloud_sch_test"
    # service_account_file_path = r"C:\Users\guangqiiang.lu\Downloads"
    # service_account_name = "buoyant-sum-302208-4542dcd74629.json"

    # file_load = FileLoad(file_name, file_path, file_sep=',',  label_name='Survived',
    #     service_account_file_name=service_account_name, service_account_file_path=service_account_file_path)

    file_load = FileLoad(file_name, file_path, file_sep=',',  label_name='Survived')
    models_path = r"C:\Users\guangqiiang.lu\Downloads\test_automl"

    auto_est = ClassificationAutoML(models_path=models_path)
    auto_est = RegressionAutoML(models_path=models_path)

    # # Start to train processing for `FileLoad`
    # auto_cl.fit(file_load=file_load, val_split=.2)
    # print(auto_cl.models_list)
    # # print(auto_cl.score(file_load_pred))

    # file_load_pred = FileLoad("test.csv", file_path, file_sep=',', use_for_pred=True)
    # print("Process based on new test file.")
    # print('*' * 20)
    # print(auto_cl.predict(file_load=file_load_pred)[:10])
    # print('*'*20)
    # print(auto_cl.predict_proba(file_load=file_load_pred)[:10])

    # print("Process based on original file.")
    # print("*" * 20)
    # print(auto_cl.predict_proba(file_load=file_load)[:10])
    # print("*" * 20)
    # print(auto_cl.score(file_load=file_load))
    # print(auto_cl.get_sorted_models_scores(file_load=file_load))

    # try to use sklearn iris dataset
    from sklearn.datasets import load_boston, load_iris
    from sklearn.model_selection import train_test_split

    x, y = load_iris(return_X_y=True)
    xtrain, xtest, ytrain, ytest = train_test_split(x, y, test_size=.2)

    # auto_cl.fit(xtrain, ytrain)

    regression_file_path = r"C:\Users\guangqiiang.lu\Documents\lugq\Kaggle\HousePricePredict"
    file_load = FileLoad('train.csv', file_path=regression_file_path, label_name='SalePrice')
    auto_est.fit(file_load=file_load)

    print(auto_est.models_list)
    print(auto_est.score(xtest, ytest))
    print('*' * 20)
    print(auto_est.predict(xtest)[:10])
    print('*' * 20)
    pred = auto_est.predict(xtest)
    print(auto_est.predict_proba(xtest)[:10])
    print("Truth data: ")
    print(ytest)

    # import matplotlib.pyplot as plt

    # plt.plot(pred, label='pred')
    # plt.plot(ytest, label='ytest')
    # plt.legend()

    # plt.show()

    # # get model score
    print(auto_cl.get_sorted_models_scores(xtest, ytest))

    # This is used for the submition for Kaggle
    # pred_data = pd.read_csv(os.path.join(file_path, test_file_name))

    # pass_id_df = pred_data[['PassengerId']]

    # prediction = auto_cl.predict(file_load_pred)
    # pred_df = pd.DataFrame(prediction, columns=['Survived'])

    # pred_df = pd.concat([pass_id_df, pred_df], axis=1)

    # pred_df.to_csv(os.path.join(file_path,"Submition_with_automl.csv"), index=False)

    # this is tested on the new computer


# All Rights Reserved.

"""Classes to load KITTI and Cityscapes data."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import glob
import json
import os
import sys
import shutil
import re

from absl import logging
import numpy as np
import imageio
from PIL import Image
import cv2
import torch
from torchvision.models.detection import maskrcnn_resnet50_fpn

from .preprocess import DynamicObjectDetector

# Color ids for masks
colors = np.arange(1, 256, dtype=np.uint8).reshape(-1, 1, 1)

class BaseLoader(object):
    r"""
    Base dataloader. We expect that all loaders inherit this class.
    """

    def __init__(self,
                 dataset_dir,
                 img_height=128,
                 img_width=416,
                 seq_length=3,
                 data_format='mono2',
                 mask='none',
                 batch_size=32,
                 threshold=0.5):
        self.dataset_dir = dataset_dir
        self.img_height = img_height
        self.img_width = img_width
        self.seq_length = seq_length
        self.data_format=data_format
        self.gen_mask = mask != 'none'
        self.mask = mask
        self.batch_size = batch_size
        self.threshold = threshold

        if self.gen_mask:
            self._initialize_mrcnn_model()

    def _initialize_mrcnn_model(self):
        model = maskrcnn_resnet50_fpn(pretrained=True)
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        model.eval()
        self.model = model.to(device)

    def run_mrcnn_model(self, images):
        r"""
        Run Mask-RCNN model on images.
        """
        images = list(images)
        with torch.no_grad():
            results = self.model(images)
        return results

    def generate_mask(self, mrcnn_result, dynamic_map=None):
        r"""
        Generate mask based on the output of Mask-RCNN model.
        """
        # Only consider objects with predicted scores higher than the threshold
        score = mrcnn_result['scores'].detach().cpu().numpy()
        valid = (score > self.threshold).sum()
        masks = (mrcnn_result['masks'] > self.threshold).squeeze(1).detach().cpu().numpy()
        labels = mrcnn_result['labels'].detach().cpu().numpy()
        if valid > 0:
            masks = masks[:valid] # (N, H, W)
            labels = labels[:valid]
        else:
            masks = np.zeros_like(masks[:1])
            labels = np.zeros_like(labels[:1])
        masks = masks.astype(np.uint8)

        # Throw away the masks that are not pedestrians or vehicles
        masks[labels == 0] *= 0 # __background__
        masks[labels == 5] *= 0 # airplane
        masks[labels == 7] *= 0 # train
        masks[labels > 8] *= 0

        mask_img = np.ones_like(masks, dtype=np.uint8)
        if self.mask == 'mono':
            mask_img = masks * mask_img
            mask_img = np.sum(mask_img, axis=0)
            mask_img = (mask_img > 0).astype(np.uint8) * 255
            return mask_img
        elif self.mask == 'color':
            for i in range(masks.shape[0]-1):
                masks[i+1:] *= 1 - masks[i]
            # ignore this step when masks is empty
            if masks.shape[0] != 0:
                # for non-background objects
                # sample colors evenly between 1 and 255
                mask_img = masks * mask_img * colors[
                        np.linspace(0, 254, num= masks.shape[0], dtype=np.uint8)
                        ]
            mask_img = np.sum(mask_img, axis=0)
            return mask_img

    def is_bad_sample(self, target_index):
        r"""
        Check whether this frame fails to satisfy certain criteria.
        This method can be overridden to customize the criteria.
        """
        return False

    def is_valid_sample(self, target_index):
        r"""
        Check whether we can find a valid sequence around this frame.
        Every loader should implement its own criterion.
        """
        raise NotImplementedError

    def get_example_with_index(self, target_index):

        add_to_file = False
        example = None

        #if self.is_bad_sample(target_index):
        #    return add_to_file, example

        # A frame without required adjacent frames
        if not self.is_valid_sample(target_index):
            # For mono2, the current frame is returned but it will not be
            # added to the train or val list
            if self.data_format == 'mono2':
                example = self.load_example(target_index)
                return add_to_file, example
            elif self.data_format == 'struct2depth':
                return add_to_file, example

        add_to_file = True
        example = self.load_example(target_index)

        return add_to_file, example

    def load_image_sequence(self, target_index):
        r"""
        Return a sequence with requested target frame.
        Every loader shoud implement its own method.
        """
        raise NotImplementedError

    def load_example(self, target_index):
        r"""
        Return a sequence with requested target frame.
        Every loader should implement its own method.
        """
        raise NotImplementedError

    def load_image_raw(self, infos):
        r"""
        Load an raw image given its id.
        Every loader should implement its own method.
        """
        raise NotImplementedError

    def load_intrinsics(self, infos):
        r"""
        Load the intrinsic matrix given its id.
        Every loader should implement its own method.
        """
        raise NotImplementedError

    def scale_intrinsics(self, mat, sx, sy, crop_top=0, crop_left=0):
        """Adjust intrinsics after resizing and then cropping
        """
        out = np.copy(mat)
        out[0, 0] *= sx
        out[0, 2] *= sx
        out[1, 1] *= sy
        out[1, 2] *= sy
        if crop_top:
            out[1, 2] -= crop_top
        if crop_left:
            out[0, 2] -= crop_left
        return out

def get_resource_path(relative_path):
    return relative_path

def get_seq_start_end(target_index, seq_length, sample_every=1):
    r"""
    Return absolute seq start and end indices for a given target frame.
    """
    if seq_length == 1:
        start_index = end_index = target_index
    else:
        half_offset = int((seq_length - 1) / 2) * sample_every
        end_index = target_index + half_offset
        start_index = end_index - (seq_length - 1) * sample_every
    return start_index, end_index

def atoi(text):
    return int(text) if text.isdigit() else text

def natural_keys(text):
    return [atoi(c) for c in re.split(r'(\d+)', text)]


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

import logging

import os.path as osp
from pathlib import Path
from typing import List, Dict, Union, Tuple, Callable

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.utils.data.distributed

import ctypes
import json
import numpy as np
import pandas as pd
from copy import deepcopy
from collections import OrderedDict
from natsort import natsorted

from nncf.debug import is_debug, DEBUG_LOG_DIR
from nncf.nncf_logger import logger
from nncf.hw_config import HWConfigType
from nncf.initialization import PartialDataLoader
from nncf.quantization.layers import BaseQuantizer
from nncf.quantization.algo import QuantizationController, NNCFNetwork, ExperimentalQuantizationController
from nncf.quantization.precision_constraints import HardwareQuantizationConstraints
from nncf.quantization.quantizer_id import QuantizerId, WeightQuantizerId, \
    NonWeightQuantizerId, InputQuantizerId, FunctionQuantizerId
from nncf.quantization.layers import QuantizerConfig

from sklearn.preprocessing import MinMaxScaler

from nncf.quantization.quantizer_setup import QuantizationPointId
from nncf.common.os import safe_open

def find_qid_by_str(qctrl: QuantizationController, qid_str: str) -> QuantizerId:
    for _qid, _q in qctrl.all_quantizations.items():
        if qid_str == str(_qid):
            return _qid
    return None

class ModelSizeCalculator:
    FLOAT_BITWIDTH = ctypes.sizeof(ctypes.c_float) * 8
    def __init__(self, qmodel: NNCFNetwork, per_quantizer_config_space: Dict[QuantizerId, List[QuantizerConfig]]):
        self._bw_space_map = OrderedDict()
        self._nparam_map = OrderedDict()
        for qid, qconfig_space in per_quantizer_config_space.items():
            if isinstance(qid, WeightQuantizerId):
                self._bw_space_map[qid] = [qconf.bits for qconf in qconfig_space]
                m = qmodel.get_module_by_scope(qid.scope)
                self._nparam_map[qid] = np.prod(m.weight.size())

        self.min_model_size, self.max_model_size = \
            self._calc_min_max_model_size()

        self.fp_model_size = self.get_uniform_bit_model_size(ModelSizeCalculator.FLOAT_BITWIDTH)

    def __call__(self, per_quantizer_bw):
        return self.get_model_size(per_quantizer_bw)

    def _calc_min_max_model_size(self) -> Tuple[np.int64, np.int64]:
        nparam_array = np.array(list(self._nparam_map.values()))
        min_size = np.sum(nparam_array * np.array(list(map(min, self._bw_space_map.values()))))
        max_size = np.sum(nparam_array * np.array(list(map(max, self._bw_space_map.values()))))
        return min_size, max_size

    def get_uniform_bit_model_size(self, uniform_bitwidth: int) -> np.int64:
        return np.sum(np.array(list(self._nparam_map.values())*uniform_bitwidth))

    def get_model_size(self, per_quantizer_bw: Dict[QuantizerId, int]) -> np.int64:
        model_size = 0
        for qid, nparam in self._nparam_map.items():
            if qid in per_quantizer_bw:
                model_size += nparam * per_quantizer_bw[qid]
            else:
                logger.warning("[ModelSizeCalculator] Missing Bitwidth of QID: {}, using {} bits"
                               .format(str(qid), ModelSizeCalculator.FLOAT_BITWIDTH))
                model_size += nparam * ModelSizeCalculator.FLOAT_BITWIDTH
        return model_size

    def get_model_size_ratio(self, per_quantizer_bw: Dict[QuantizerId, int]) -> np.float64:
        return self.get_model_size(per_quantizer_bw)/self.fp_model_size

class QuantizationEnvParams:
    def __init__(self, compression_ratio: float,
                 eval_subset_ratio: float,
                 skip_constraint: bool,
                 finetune: bool,
                 bits: List[int],
                 dump_init_precision_data: bool = False,
                 log_dir: str = None):
        self.compression_ratio = compression_ratio
        self.eval_subset_ratio = eval_subset_ratio
        self.skip_constraint = skip_constraint
        self.finetune = finetune
        self.bits = bits
        self.dump_init_precision_data = dump_init_precision_data
        self.log_dir = log_dir

class QuantizationEnv:
    # pylint:disable=too-many-branches,too-many-statements
    def __init__(self,
                 model: NNCFNetwork,
                 quantization_controller: ExperimentalQuantizationController,
                 hw_precision_constraints: HardwareQuantizationConstraints,
                 eval_loader: torch.utils.data.DataLoader,
                 eval_fn: Callable[[nn.Module, torch.utils.data.DataLoader], float],
                 hw_config_type: HWConfigType,
                 params: QuantizationEnvParams
                 ):

        logger.info("[Q.Env] Instantiating NNCF Quantization Environment")
        self.qctrl = quantization_controller
        self.qmodel = model
        self.eval_loader = eval_loader
        self.eval_fn = eval_fn
        self._hw_precision_constraints = hw_precision_constraints

        self.model_name = self.qmodel.nncf_module.__class__.__name__

        # Check and only proceed if target device is supported by Q.Env
        self.hw_cfg_type = hw_config_type
        assert self.hw_cfg_type in [None, HWConfigType.VPU]

        # Set target compression ratio
        self.compression_ratio = params.compression_ratio

        self.eval_loader = PartialDataLoader(self.eval_loader, iter_ratio=params.eval_subset_ratio)

        # Bool to disable hard resource constraint
        self.skip_constraint = params.skip_constraint

        # Bool to enable fine-tuning in each episode. Placeholder for now
        self.finetune = params.skip_constraint

        # Configure search space for precision according to target device
        if self.hw_cfg_type is None:
            self.model_bitwidth_space = params.bits
        elif self.hw_cfg_type is HWConfigType.VPU:
            self.model_bitwidth_space = self._hw_precision_constraints.get_all_unique_bits()
        self.model_bitwidth_space = sorted(list(self.model_bitwidth_space))

        # Create mapping of QuantizerId to the space of the corresponding quantizer's allowed qconfigs
        #pylint:disable=line-too-long
        self.qconfig_space_map = OrderedDict.fromkeys(self.qctrl.all_quantizations.keys())  # type: Dict[QuantizerId, List[QuantizerConfig]]
        if self.hw_cfg_type is None:
            for qid in self.qconfig_space_map.keys():
                conf = self.qctrl.all_quantizations[qid].get_current_config()
                conf_list_to_set = []
                for bit in self.model_bitwidth_space:
                    bit_adjusted_conf = deepcopy(conf)
                    bit_adjusted_conf.bits = bit
                    conf_list_to_set.append(bit_adjusted_conf)
                self.qconfig_space_map[qid] = conf_list_to_set
        else:
            for qid in self.qconfig_space_map:
                self.qconfig_space_map[qid] = self._hw_precision_constraints.get(qid)

        # Quantizer Master Table Creation
        self._groups_of_adjacent_quantizers = self.qctrl._groups_of_adjacent_quantizers
        self.quantizer_table = self._create_quantizer_table()

        # Create master dataframe to keep track of quantizable layers and their attributes
        self.master_df, self.state_list = self._get_state_space(self.qctrl, self.qmodel, self.quantizer_table)
        if self.master_df.isnull().values.any():
            raise ValueError("Q.Env Master Dataframe has null value(s)")

        assert len(self.quantizer_table) == len(self.qctrl.all_quantizations), \
            "Number of Quantizer is not tally between quantizer table and quantization controller"

        # MinMaxScaler for State Embedding
        self.state_scaler = MinMaxScaler()
        self.state_scaler.fit(self.master_df[self.state_list])

        # Model Size Calculation
        self.model_size_calculator = ModelSizeCalculator(self.qmodel, self.qconfig_space_map)
        self.orig_model_size = self.model_size_calculator.fp_model_size
        self.min_model_size = self.model_size_calculator.min_model_size
        self.max_model_size = self.model_size_calculator.max_model_size
        self.target_model_size = self.orig_model_size * self.compression_ratio

        if self.target_model_size < self.min_model_size and self.target_model_size > self.max_model_size:
            raise ValueError("Model Size Ratio {} is out of bound ({}, {})"
                             .format(self.compression_ratio,
                                     self.min_model_size / self.orig_model_size,
                                     self.max_model_size / self.orig_model_size))

        # Evaluate and store metric score of pretrained model
        self._evaluate_pretrained_model()
        self.qmodel_init_sd = deepcopy(self.qmodel.state_dict())

        self.reset()

        self._dump_autoq_data = params.dump_init_precision_data
        if self._dump_autoq_data or is_debug():
            dump_dir = params.log_dir
            if dump_dir is None:
                dump_dir = DEBUG_LOG_DIR
            self.dump_dir = Path(dump_dir) / Path("autoq_env_dump")
            self.dump_dir.mkdir(parents=True, exist_ok=True)
            # Serialize Q.Env information. Note that these functions should be at the end of Q.Env Initialization.
            self._dump_master_df()
            self._dump_quantized_graph()
            self._dump_groups_of_adjacent_quantizers()

        # End of QuantizationEnv.__init__()
        # --------------------------------------------------------------------------------------------------------------

    def reset(self):
        self.collected_strategy = []
        self.master_df['action'] = max(self.model_bitwidth_space)
        self.master_df['prev_action'] = 0
        self.master_df['unconstrained_action'] = 0

    def _create_quantizer_table(self) -> pd.DataFrame:
        # Create a mapping of qid to its adjacent quantizer group id
        adjq_gid_map = OrderedDict.fromkeys(self.qctrl.all_quantizations.keys())
        for qid in self.qctrl.all_quantizations:
            adjq_gid_map[qid] = self._groups_of_adjacent_quantizers.get_group_id_for_quantizer(qid)

        assert len(set(self.qconfig_space_map.keys()) - set(adjq_gid_map.keys())) == 0, \
            "both qconfig_space_map and adjq_gid_map must have exact keys."

        # Create a mapping of qid to its nodekey in NNCFGraph
        qid_nodekey_map = self._generate_qid_nodekey_map(self.qctrl, self.qmodel)

        assert len(set(qid_nodekey_map.keys()) - set(adjq_gid_map.keys())) == 0, \
            "both qid_nodekey_map and adjq_gid_map must have exact keys."

        d = OrderedDict()
        for qid in self.qctrl.all_quantizations:
            nodekey = qid_nodekey_map[qid]
            q_nx_nodeid = nodekey.split()[0]
            idx_str = '-'.join([q_nx_nodeid, str(qid)])
            gid = adjq_gid_map[qid]

            d[idx_str] = OrderedDict()
            d[idx_str]['qid'] = str(qid)
            d[idx_str]['q_nx_nodeid'] = q_nx_nodeid
            d[idx_str]['q_nx_nodekey'] = nodekey
            d[idx_str]['gid'] = gid
            d[idx_str]['qconf_space'] = self.qconfig_space_map[qid]
            d[idx_str]['qp_id_set'] = self.qctrl.module_id_to_qp_id_translation_dict[qid]

            if isinstance(qid, WeightQuantizerId):
                d[idx_str]['state_scope'] = qid.scope
            elif isinstance(qid, NonWeightQuantizerId):
                d[idx_str]['state_scope'] = qid.ia_op_exec_context.scope_in_model
            else:
                raise NotImplementedError("QuantizerId: {} of {} class is not supported."
                                          .format(str(qid), str(qid.__class__.__name__)))

        # quantizer_table index is QuantizerId in string prepended with its quantize node id in NNCFGraph
        df = pd.DataFrame.from_dict(d, orient='index')
        df['qid_obj'] = df['qid'].apply(lambda x: find_qid_by_str(self.qctrl, x))
        df['qmodule'] = df['qid_obj'].apply(lambda x: self.qctrl.all_quantizations[x])
        df['is_wt_quantizer'] = df['qmodule'].apply(lambda x: x.is_weights)
        df['state_module'] = df['state_scope'].apply(self.qmodel.get_module_by_scope)

        quantizer_table = df.loc[natsorted(df.index)]
        return quantizer_table

    def _get_state_space(self,
                         quantization_controller: QuantizationController,
                         quantized_model: NNCFNetwork,
                         quantizer_table: pd.DataFrame) -> Tuple[pd.DataFrame, List[str]]:
        def annotate_learnable_module_io_shape(model):
            def annotate_io_shape(module, input_, output):
                if hasattr(module, 'weight') or isinstance(module, BaseQuantizer):
                    module.input_shape_ = input_[0].shape
                    module.output_shape_ = output.shape

            hook_list = [m.register_forward_hook(annotate_io_shape) for n, m in model.named_modules()]
            model.do_dummy_forward(force_eval=True)
            for h in hook_list:
                h.remove()

        annotate_learnable_module_io_shape(quantized_model)

        # State Embedding Extraction
        #---------------------------
        df = quantizer_table
        layer_attr_df = df.apply(self._get_layer_attr, axis=1)
        layer_attr_df['layer_idx'] = np.array(range(len(layer_attr_df)))
        layer_attr_df['weight_quantizer'] = df['is_wt_quantizer'].astype('float')
        state_list = layer_attr_df.columns.to_list()

        # create master dataframe
        master_df = pd.concat([df, layer_attr_df], axis='columns')

        # Annotate a min and a max value in prev_action before minmaxscaler fitting
        master_df.loc[master_df.index[0], 'prev_action'] = max(self.model_bitwidth_space)
        master_df.loc[master_df.index[-1], 'prev_action'] = min(self.model_bitwidth_space)

        return master_df, state_list

    def _get_layer_attr(self, row: pd.Series) -> pd.Series:
        m = row.state_module
        qid = row.qid_obj
        feature = OrderedDict()

        if isinstance(qid, WeightQuantizerId):
            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
                feature['conv_dw'] = int(m.weight.shape[1] == m.groups) # 1.0 for depthwise, 0.0 for other conv2d
                feature['cin'] = m.weight.shape[1]
                feature['cout'] = m.weight.shape[0]
                feature['stride'] = m.stride[0]
                feature['kernel'] = m.kernel_size[0]
                feature['param'] = np.prod(m.weight.size())
                feature['ifm_size'] = np.prod(m.input_shape_[-2:]) # H*W
                feature['prev_action'] = 0.0 # placeholder

            elif isinstance(m, nn.Linear):
                feature['conv_dw'] = 0.0
                feature['cin'] = m.in_features
                feature['cout'] = m.out_features
                feature['stride'] = 0.0
                feature['kernel'] = 1.0
                feature['param'] = np.prod(m.weight.size())
                feature['ifm_size'] = np.prod(m.input_shape_[-1]) # feature nodes
                feature['prev_action'] = 0.0 # placeholder

            else:
                raise NotImplementedError("State embedding extraction of {}".format(m.__class__.__name__))

        elif isinstance(qid, NonWeightQuantizerId):
            qmod = self.qctrl.all_quantizations[qid]
            input_shape = qmod.input_shape_
            output_shape = qmod.output_shape_
            feature['cin'] = input_shape[1] if len(input_shape) == 4 else input_shape[-1]
            feature['cout'] = output_shape[1] if len(output_shape) == 4 else output_shape[-1]
            feature['ifm_size'] = np.prod(input_shape[-2:]) if len(input_shape) == 4 else input_shape[-1]
            feature['conv_dw'] = 0.0
            feature['stride'] = 0.0
            feature['kernel'] = 0.0
            feature['param'] = 0.0
            feature['prev_action'] = 0.0

            if len(input_shape) != 4 and len(input_shape) != 2:
                raise NotImplementedError("A design is required to cater this scenario. Pls. report to maintainer")

        elif isinstance(qid, InputQuantizerId):
            raise NotImplementedError("InputQuantizerId is supported, quantizer of nncf model input "
                                      "is of type NonWeightQuantizerId")

        elif isinstance(qid, FunctionQuantizerId):
            raise NotImplementedError("FunctionQuantizerId is supported, Pls. report to maintainer")

        else:
            raise ValueError("qid is an instance of unexpected class {}".format(qid.__class__.__name__))

        return pd.Series(feature)

    def _evaluate_pretrained_model(self):
        logger.info("[Q.Env] Evaluating Pretrained Model")
        self.qctrl.disable_weight_quantization()
        self.qctrl.disable_activation_quantization()

        with torch.no_grad():
            self.pretrained_score = self.eval_fn(self.qmodel, self.eval_loader)
            logger.info("Pretrained Score: {:.3f}".format(self.pretrained_score))

        self.qctrl.enable_weight_quantization()
        self.qctrl.enable_activation_quantization()
        self.qmodel.rebuild_graph()

    def _run_quantization_pipeline(self, finetune=False) -> float:
        self.qctrl.run_batchnorm_adaptation(self.qctrl.quantization_config)

        if finetune:
            raise NotImplementedError("Post-Quantization fine tuning is not implemented.")
        with torch.no_grad():
            quantized_score = self.eval_fn(self.qmodel, self.eval_loader)
            logger.info("[Q.Env] Quantized Score: {:.3f}".format(quantized_score))
        return quantized_score

    def _get_quantizer_bitwidth(self) -> Dict[BaseQuantizer, int]:
        assert len(set(self.model_bitwidth_space) - set(self.master_df.action.values)) >= 0, \
            "there is bitwidth choice not within model bitwidth space"
        return OrderedDict(zip(self.master_df.qid_obj, self.master_df.action))

    def _constrain_model_size(self, collected_strategy: List, skip=False) -> List:
        def lower_bitwidth(bw: int, qconf_space: List[QuantizerConfig]) -> int:
            bw_space = [qconf.bits for qconf in qconf_space]
            assert bw in bw_space
            sorted_bw_space = sorted(bw_space)
            return sorted_bw_space[sorted_bw_space.index(bw)-1] if sorted_bw_space.index(bw) > 0 else bw

        # This function acts on self.master_df['action']
        self.master_df['action'] = collected_strategy

        if skip is not True:
            self.master_df['unconstrained_action'] = self.master_df['action']

            current_model_size = self.model_size_calculator(self._get_quantizer_bitwidth())

            while self.min_model_size < current_model_size and self.target_model_size < current_model_size:
                for _, nodestr in enumerate(reversed(self.master_df.index.tolist())):
                    if self.master_df.loc[nodestr, "is_wt_quantizer"]:
                        bw_choice, qconf_space = self.master_df.loc[nodestr, ['action', 'qconf_space']]
                        new_bw = lower_bitwidth(bw_choice, qconf_space)
                        self.master_df.loc[nodestr, "action"] = new_bw if new_bw != bw_choice else bw_choice

                    current_model_size = self.model_size_calculator(self._get_quantizer_bitwidth())
                    if current_model_size <= self.target_model_size:
                        break
        else:
            logger.info("[Q.Env] Skipping Model Size Constraint")

        return self.master_df['action'].tolist()

    def reward(self, acc: float, model_ratio: float) -> float:
        return (acc - self.pretrained_score) * 0.1

    def step(self, action: Union[int, float]) -> Tuple:
        currently_processed_qconf_idx = len(self.collected_strategy)

        def is_final_step():
            return len(self.collected_strategy) == len(self.master_df)

        # Ensure action is in the quantizer's bitwidth space
        current_qconf_space = self.master_df.qconf_space[currently_processed_qconf_idx]
        current_bw_space = [qconf.bits for qconf in current_qconf_space]
        if action not in current_bw_space:
            closest_bw_idx = np.argmin(np.abs(action - np.array(current_bw_space)))
            action = current_bw_space[closest_bw_idx]

        self.collected_strategy.append(action)

        if not is_final_step():
            info_set = {}
            reward = 0
            self.set_next_step_prev_action(len(self.collected_strategy), action)
            obs = self.get_normalized_obs(len(self.collected_strategy))
            done = False
            return obs, reward, done, info_set

        return self.evaluate_strategy(self.collected_strategy, skip_constraint=self.skip_constraint)

    def select_config_for_actions(self, actions) -> Dict[QuantizationPointId, QuantizerConfig]:
        retval = OrderedDict()  # type: Dict[QuantizationPointId, QuantizerConfig]
        for action, qp_id_set, qconf_space in zip(actions,
                                            self.master_df['qp_id_set'],
                                            self.master_df['qconf_space']):
            matches = []
            for qconf in qconf_space:
                if qconf.bits == action:
                    matches.append(qconf)
            assert len(matches) == 1
            for qp_id in qp_id_set:
                retval[qp_id] = matches[0]
        return retval

    def evaluate_strategy(self, collected_strategy: List, skip_constraint=True) -> Tuple:
        assert len(collected_strategy) == len(self.master_df)
        if skip_constraint is not True:
            collected_strategy = self._constrain_model_size(collected_strategy)
        self.master_df['action'] = collected_strategy # This must be after constraint

        configs_to_set = self.select_config_for_actions(self.master_df['action'])
        self._apply_quantizer_configs_to_model(configs_to_set)

        for idx, qid in zip(self.master_df.index, self.master_df['qid']):
            logger.info("[Q.Env] {:50} | {}".format(
                str(self.qctrl.all_quantizations[find_qid_by_str(self.qctrl, qid)]),
                idx))

        quantized_score = self._run_quantization_pipeline(finetune=self.finetune)

        current_model_size = self.model_size_calculator(self._get_quantizer_bitwidth())
        current_model_ratio = self.model_size_calculator.get_model_size_ratio(self._get_quantizer_bitwidth())

        reward = self.reward(quantized_score, current_model_ratio)

        info_set = {'model_ratio': current_model_ratio, 'accuracy': quantized_score, 'model_size': current_model_size}

        obs = self.get_normalized_obs(len(collected_strategy) - 1)
        done = True

        return obs, reward, done, info_set

    def set_next_step_prev_action(self, idx, action):
        self.master_df.loc[self.master_df.index[idx], 'prev_action'] = action

    def get_normalized_obs(self, idx: int) -> pd.Series:
        _df = self.master_df.loc[self.master_df.index, self.state_list]
        _df.loc[_df.index, self.state_list] = self.state_scaler.transform(_df[self.state_list])
        return _df.iloc[idx]

    def _apply_quantizer_configs_to_model(self, qid_vs_qconfig_map: Dict[QuantizationPointId, QuantizerConfig]):
        new_quantizer_setup = self.qctrl.get_quantizer_setup_for_current_state()
        for qp_id, qconf in qid_vs_qconfig_map.items():
            new_quantizer_setup.quantization_points[qp_id].qconfig = qconf
        self.qmodel.load_state_dict(self.qmodel_init_sd)
        had_to_regenerate = self.qctrl.is_new_setup_requires_regeneration(new_quantizer_setup)
        self.qctrl, self.qmodel = self.qctrl.apply_new_quantizer_setup(new_quantizer_setup)
        if had_to_regenerate:
            self.qmodel_init_sd = deepcopy(self.qmodel.state_dict())

        # The QuantizerId's may have changed after the new quantizer setup application, but
        # QuantizationPointId's should not have. Will use this to update the qids in the master table.
        for qid, qp_id_set in self.qctrl.module_id_to_qp_id_translation_dict.items():
            self.master_df.loc[self.master_df.qp_id_set == qp_id_set].qid = str(qid)

    def _generate_qid_nodekey_map(self,
                                  quantization_controller: QuantizationController,
                                  quantized_network: NNCFNetwork) -> Dict[QuantizerId, str]:
        """
        Create a lookup mapping for each QuantizerId to its corresponding quantize node in network graph
        :param quantization_controller:
        :param quantized_network:
        :return: dict with key of QuantizerId and value of node key string
        """
        # Map Non Weight Qid to its nodes in nxgraph
        weight_quantize_nodekeys = []
        non_weight_quantize_nodekeys = []
        qid_nodekey_map = OrderedDict()

        quantized_network.rebuild_graph()
        g = quantized_network.get_graph()

        for nodekey in g.get_all_node_keys():
            if 'symmetric_quantize' in nodekey and 'UpdateWeight' in nodekey:
                weight_quantize_nodekeys.append(nodekey)
            if 'symmetric_quantize' in nodekey and 'UpdateWeight' not in nodekey:
                non_weight_quantize_nodekeys.append(nodekey)

        # Find nodekey of Weight Quantizer
        for qid, _ in quantization_controller.weight_quantizers.items():
            quantize_nodekeys = []
            for nodekey in weight_quantize_nodekeys:
                if str(qid.scope) in nodekey:
                    quantize_nodekeys.append(nodekey)

            if len(quantize_nodekeys) == 1:
                qid_nodekey_map[qid] = quantize_nodekeys[0]
            else:
                raise ValueError("Quantize Node not found or More Nodes are found for WQid: {}".format(qid))

        # Find nodekey of Non-Weight Quantizer
        for qid, _ in quantization_controller.non_weight_quantizers.items():
            quantize_nodekeys = []
            for nodekey in non_weight_quantize_nodekeys:
                if str(qid.ia_op_exec_context.scope_in_model) in nodekey:
                    quantize_nodekeys.append(nodekey)

            if len(quantize_nodekeys) > 0:
                qid_nodekey_map[qid] = quantize_nodekeys[qid.ia_op_exec_context.call_order]
            else:
                raise ValueError("Quantize Node not found for NWQid: {}".format(qid))

        if logger.level == logging.DEBUG:
            for qid, nodekey in qid_nodekey_map.items():
                logger.debug("QuantizerId: {}".format(qid))
                logger.debug("\tnodekey: {}".format(nodekey))

        return qid_nodekey_map

    def _dump_master_df(self):
        self.master_df.drop('state_module', axis=1).to_csv(
            osp.join(self.dump_dir, self.model_name + "_quantizable_state_table.csv"), index_label="nodestr")

    def _dump_quantized_graph(self):
        self.qmodel.get_graph().visualize_graph(osp.join(self.dump_dir, "qenv_graph.dot"))

    def _dump_groups_of_adjacent_quantizers(self):
        adj_quantizer_groups = []

        for i, _ in enumerate(self._groups_of_adjacent_quantizers):
            group_members = []
            for _, aq in enumerate(self._groups_of_adjacent_quantizers[i].activation_quantizers):
                group_members.append(self.master_df.index[self.master_df.qid == str(aq[0])][0])
            for _, wq in enumerate(self._groups_of_adjacent_quantizers[i].weight_quantizers):
                group_members.append(self.master_df.index[self.master_df.qid == str(wq[0])][0])
            adj_quantizer_groups.append(natsorted(group_members))

        with safe_open(self.dump_dir / self.model_name / "_groups_of_adjacent_quantizers.json", "w") as DUMP_FH:
            json.dump(natsorted(adj_quantizer_groups), DUMP_FH, indent=4)


"""
Run Prototypical Soft Actor Critic on HalfCheetahEnv.

"""
import os
import numpy as np
import click
import datetime
import pathlib
from gym.envs.mujoco import HalfCheetahEnv

from rlkit.envs.half_cheetah_vel import HalfCheetahVelEnv
from rlkit.envs.wrappers import NormalizedBoxEnv
from rlkit.launchers.launcher_util import setup_logger
from rlkit.torch.sac.policies import TanhGaussianPolicy
from rlkit.torch.networks import FlattenMlp, MlpEncoder, RecurrentEncoder
from rlkit.torch.sac.sac import ProtoSoftActorCritic
from rlkit.torch.sac.proto import ProtoAgent
import rlkit.torch.pytorch_util as ptu

def datetimestamp(divider=''):
    now = datetime.datetime.now()
    return now.strftime('%Y-%m-%d-%H-%M-%S-%f').replace('-', divider)

def experiment(variant):
    task_params = variant['task_params']
    env = NormalizedBoxEnv(HalfCheetahVelEnv(n_tasks=task_params['n_tasks']))
    ptu.set_gpu_mode(variant['use_gpu'], variant['gpu_id'])

    tasks = env.get_all_task_idx()

    obs_dim = int(np.prod(env.observation_space.shape))
    action_dim = int(np.prod(env.action_space.shape))
    latent_dim = 5
    task_enc_output_dim = latent_dim * 2 if variant['algo_params']['use_information_bottleneck'] else latent_dim
    reward_dim = 1

    net_size = variant['net_size']
    recurrent = variant['algo_params']['recurrent']
    encoder_model = RecurrentEncoder if recurrent else MlpEncoder
    task_enc = encoder_model(
            hidden_sizes=[200, 200, 200], # deeper net + higher dim space generalize better
            input_size=obs_dim + action_dim + reward_dim,
            output_size=task_enc_output_dim,
    )
    qf1 = FlattenMlp(
        hidden_sizes=[net_size, net_size, net_size],
        input_size=obs_dim + action_dim + latent_dim,
        output_size=1,
    )
    qf2 = FlattenMlp(
        hidden_sizes=[net_size, net_size, net_size],
        input_size=obs_dim + action_dim + latent_dim,
        output_size=1,
    )
    vf = FlattenMlp(
        hidden_sizes=[net_size, net_size, net_size],
        input_size=obs_dim + latent_dim,
        output_size=1,
    )
    policy = TanhGaussianPolicy(
        hidden_sizes=[net_size, net_size, net_size],
        obs_dim=obs_dim + latent_dim,
        latent_dim=latent_dim,
        action_dim=action_dim,
    )

    agent = ProtoAgent(
        latent_dim,
        [task_enc, policy, qf1, qf2, vf],
        **variant['algo_params']
    )

    algorithm = ProtoSoftActorCritic(
        env=env,
        train_tasks=tasks[:-30],
        eval_tasks=tasks[-30:],
        nets=[agent, task_enc, policy, qf1, qf2, vf],
        latent_dim=latent_dim,
        **variant['algo_params']
    )
    algorithm.to()
    algorithm.train()

@click.command()
@click.argument('gpu', default=0)
@click.option('--docker', default=0)
def main(gpu, docker):
    max_path_length = 200
    # noinspection PyTypeChecker
    variant = dict(
        task_params=dict(
            n_tasks=130,
            randomize_tasks=True,
        ),
        algo_params=dict(
            meta_batch=10,
            num_iterations=500, # meta-train epochs
            num_tasks_sample=5,
            num_steps_per_task=2 * max_path_length,
            num_train_steps_per_itr=2000,
            num_evals=2, # number of evals with separate task encodings
            num_steps_per_eval=2 * max_path_length,
            batch_size=256, # to compute training grads from
            embedding_batch_size=100,
            embedding_mini_batch_size=100,
            max_path_length=max_path_length,
            discount=0.99,
            soft_target_tau=0.005,
            policy_lr=3E-4,
            qf_lr=3E-4,
            vf_lr=3E-4,
            context_lr=3e-4,
            reward_scale=5.,
            sparse_rewards=False,
            reparameterize=True,
            kl_lambda=.1,
            use_information_bottleneck=True,
            train_embedding_source='online_exploration_trajectories',
            eval_embedding_source='online_exploration_trajectories',
            recurrent=False, # recurrent or averaging encoder
            dump_eval_paths=False,
        ),
        net_size=300,
        use_gpu=True,
        gpu_id=gpu,
    )
    exp_name = 'pearl'

    log_dir = '/mounts/output' if docker == 1 else 'output'
    exp_id = 'half-cheetah-vel'
    os.makedirs(os.path.join(log_dir, exp_id), exist_ok=True)
    experiment_log_dir = setup_logger(exp_name, variant=variant, exp_id=exp_id, base_log_dir=log_dir)

    # creates directories for pickle outputs of trajectories (point mass)
    pickle_dir = experiment_log_dir + '/eval_trajectories'
    pathlib.Path(pickle_dir).mkdir(parents=True, exist_ok=True)
    variant['algo_params']['output_dir'] = pickle_dir

    # debugging triggers a lot of printing
    DEBUG = 0
    os.environ['DEBUG'] = str(DEBUG)

    experiment(variant)

if __name__ == "__main__":
    main()


from pathlib import Path
import json
import random
import os

import numpy as np
import torch
from torch.nn import CrossEntropyLoss
from torch.optim import SGD, lr_scheduler
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.backends import cudnn
import torchvision

from opts import parse_opts
from model import (generate_model, load_pretrained_model, make_data_parallel,
                   get_fine_tuning_parameters)
from mean import get_mean_std
from spatial_transforms import (Compose, Normalize, Resize, CenterCrop,
                                CornerCrop, MultiScaleCornerCrop,
                                RandomResizedCrop, RandomHorizontalFlip,
                                ToTensor, ScaleValue, ColorJitter,
                                PickFirstChannels)
from temporal_transforms import (LoopPadding, TemporalRandomCrop,
                                 TemporalCenterCrop, TemporalEvenCrop,
                                 SlidingWindow, TemporalSubsampling)
from temporal_transforms import Compose as TemporalCompose
from dataset import get_training_data, get_validation_data, get_inference_data
from utils import Logger, worker_init_fn, get_lr
from training import train_epoch
# from validation import val_epoch
import inference

# from advertorch.attacks import LinfPGDAttack

## ----------------------- APE-GAN ----------------------- ##
from gan_models import GeneratorUCF3D as Generator, DiscriminatorUCF3D as Discriminator

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)
## ---------------------------- // ----------------------- ##
from scipy.stats import loguniform, beta
# from matplotlib import pyplot as plt

def json_serial(obj):
    if isinstance(obj, Path):
        return str(obj)

def get_opt(args=None):
    opt = parse_opts(args)

    if opt.root_path is not None:
        opt.video_path = opt.root_path / opt.video_path
        opt.annotation_path = opt.root_path / opt.annotation_path
        opt.result_path = opt.root_path / opt.result_path
        if opt.resume_path is not None:
            opt.resume_path = opt.root_path / opt.resume_path
        if opt.pretrain_path is not None:
            opt.pretrain_path = opt.root_path / opt.pretrain_path

    if opt.pretrain_path is not None:
        opt.n_finetune_classes = opt.n_classes
        opt.n_classes = opt.n_pretrain_classes

    if opt.output_topk <= 0:
        opt.output_topk = opt.n_classes

    if opt.inference_batch_size == 0:
        opt.inference_batch_size = opt.batch_size

    opt.arch = '{}-{}'.format(opt.model, opt.model_depth)
    opt.begin_epoch = 1
    opt.mean, opt.std = get_mean_std(opt.value_scale, dataset=opt.mean_dataset)
    opt.n_input_channels = 3
    if opt.input_type == 'flow':
        opt.n_input_channels = 2
        opt.mean = opt.mean[:2]
        opt.std = opt.std[:2]

#     if opt.distributed:
#         opt.dist_rank = int(os.environ["OMPI_COMM_WORLD_RANK"])

#         if opt.dist_rank == 0:
#             print(opt)
#             with (opt.result_path / 'opts.json').open('w') as opt_file:
#                 json.dump(vars(opt), opt_file, default=json_serial)
#     else:
#         print(opt)
#         with (opt.result_path / 'opts.json').open('w') as opt_file:
#             json.dump(vars(opt), opt_file, default=json_serial)

    return opt

def resume_model(resume_path, arch, model):
    print('loading checkpoint {} model'.format(resume_path))
    checkpoint = torch.load(resume_path, map_location='cpu')

    assert arch == checkpoint['arch']

    if hasattr(model, 'module'):
        model.module.load_state_dict(checkpoint['state_dict'])
    else:
        model.load_state_dict(checkpoint['state_dict'])

    return model

def resume_train_utils(resume_path, begin_epoch, optimizer, scheduler):
    print('loading checkpoint {} train utils'.format(resume_path))
    checkpoint = torch.load(resume_path, map_location='cpu')
    begin_epoch = checkpoint['epoch'] + 1
#     if optimizer is not None and 'optimizer' in checkpoint:
#         optimizer.load_state_dict(checkpoint['optimizer'])
#     if scheduler is not None and 'scheduler' in checkpoint:
#         scheduler.load_state_dict(checkpoint['scheduler'])

    return begin_epoch, optimizer, scheduler

def get_normalize_method(mean, std, no_mean_norm, no_std_norm):
    if no_mean_norm:
        if no_std_norm:
            return Normalize([0, 0, 0], [1, 1, 1])
        else:
            return Normalize([0, 0, 0], std)
    else:
        if no_std_norm:
            return Normalize(mean, [1, 1, 1])
        else:
            return Normalize(mean, std)

def get_train_utils(opt, model_parameters):
    assert opt.train_crop in ['random', 'corner', 'center']
    spatial_transform = []
    if opt.train_crop == 'random':
        spatial_transform.append(
            RandomResizedCrop(
                opt.sample_size, (opt.train_crop_min_scale, 1.0),
                (opt.train_crop_min_ratio, 1.0 / opt.train_crop_min_ratio)))
    elif opt.train_crop == 'corner':
        scales = [1.0]
        scale_step = 1 / (2**(1 / 4))
        for _ in range(1, 5):
            scales.append(scales[-1] * scale_step)
        spatial_transform.append(MultiScaleCornerCrop(opt.sample_size, scales))
    elif opt.train_crop == 'center':
        spatial_transform.append(Resize(opt.sample_size))
        spatial_transform.append(CenterCrop(opt.sample_size))
    normalize = get_normalize_method(opt.mean, opt.std, opt.no_mean_norm,
                                     opt.no_std_norm)
    if not opt.no_hflip:
        spatial_transform.append(RandomHorizontalFlip())
    if opt.colorjitter:
        spatial_transform.append(ColorJitter())
    spatial_transform.append(ToTensor())
    if opt.input_type == 'flow':
        spatial_transform.append(PickFirstChannels(n=2))
    spatial_transform.append(ScaleValue(opt.value_scale))
    spatial_transform.append(normalize)
    spatial_transform = Compose(spatial_transform)

    assert opt.train_t_crop in ['random', 'center']
    temporal_transform = []
    if opt.sample_t_stride > 1:
        temporal_transform.append(TemporalSubsampling(opt.sample_t_stride))
    if opt.train_t_crop == 'random':
        temporal_transform.append(TemporalRandomCrop(opt.sample_duration))
    elif opt.train_t_crop == 'center':
        temporal_transform.append(TemporalCenterCrop(opt.sample_duration))
    temporal_transform = TemporalCompose(temporal_transform)

    train_data = get_training_data(opt.video_path, opt.annotation_path,
                                   opt.dataset, opt.input_type, opt.file_type,
                                   spatial_transform, temporal_transform)
    if opt.distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(
            train_data)
    else:
        train_sampler = None
    train_loader = torch.utils.data.DataLoader(train_data,
                                               batch_size=opt.batch_size,
                                               shuffle=(train_sampler is None),
                                               num_workers=opt.n_threads,
                                               pin_memory=True,
                                               sampler=train_sampler,
                                               worker_init_fn=worker_init_fn)

    if opt.is_master_node:
        train_logger = Logger(opt.result_path / 'train.log',
                              ['epoch', 'loss', 'acc', 'lr'])
        train_batch_logger = Logger(
            opt.result_path / 'train_batch.log',
            ['epoch', 'batch', 'iter', 'loss', 'acc', 'lr'])
    else:
        train_logger = None
        train_batch_logger = None

    if opt.nesterov:
        dampening = 0
    else:
        dampening = opt.dampening
    optimizer = SGD(model_parameters,
                    lr=opt.learning_rate,
                    momentum=opt.momentum,
                    dampening=dampening,
                    weight_decay=opt.weight_decay,
                    nesterov=opt.nesterov)

    assert opt.lr_scheduler in ['plateau', 'multistep']
    assert not (opt.lr_scheduler == 'plateau' and opt.no_val)
    if opt.lr_scheduler == 'plateau':
        scheduler = lr_scheduler.ReduceLROnPlateau(
            optimizer, 'min', patience=opt.plateau_patience)
    else:
        scheduler = lr_scheduler.MultiStepLR(optimizer,
                                             opt.multistep_milestones, gamma=0.1) #, gamma=0.1

    return (train_loader, train_sampler, train_logger, train_batch_logger,
            optimizer, scheduler)

def get_val_utils(opt):
    normalize = get_normalize_method(opt.mean, opt.std, opt.no_mean_norm,
                                     opt.no_std_norm)
    spatial_transform = [
        Resize(opt.sample_size),
        CenterCrop(opt.sample_size),
        ToTensor()
    ]
    if opt.input_type == 'flow':
        spatial_transform.append(PickFirstChannels(n=2))
    spatial_transform.extend([ScaleValue(opt.value_scale), normalize])
    spatial_transform = Compose(spatial_transform)

    temporal_transform = []
    if opt.sample_t_stride > 1:
        temporal_transform.append(TemporalSubsampling(opt.sample_t_stride))
    temporal_transform.append(
        TemporalEvenCrop(opt.sample_duration, opt.n_val_samples))
    temporal_transform = TemporalCompose(temporal_transform)

    val_data, collate_fn = get_validation_data(opt.video_path,
                                               opt.annotation_path, opt.dataset,
                                               opt.input_type, opt.file_type,
                                               spatial_transform,
                                               temporal_transform)
    if opt.distributed:
        val_sampler = torch.utils.data.distributed.DistributedSampler(
            val_data, shuffle=False)
    else:
        val_sampler = None
    val_loader = torch.utils.data.DataLoader(val_data,
                                             batch_size=(opt.batch_size //
                                                         opt.n_val_samples),
                                             shuffle=False,
                                             num_workers=opt.n_threads,
                                             pin_memory=True,
                                             sampler=val_sampler,
                                             worker_init_fn=worker_init_fn,
                                             collate_fn=collate_fn)

    if opt.is_master_node:
        val_logger = Logger(opt.result_path / 'val.log',
                            ['epoch', 'loss', 'acc'])
    else:
        val_logger = None

    return val_loader, val_logger

def get_inference_utils(opt):
    assert opt.inference_crop in ['center', 'nocrop']

    normalize = get_normalize_method(opt.mean, opt.std, opt.no_mean_norm,
                                     opt.no_std_norm)

    spatial_transform = [Resize(opt.sample_size)]
    if opt.inference_crop == 'center':
        spatial_transform.append(CenterCrop(opt.sample_size))
    spatial_transform.append(ToTensor())
    if opt.input_type == 'flow':
        spatial_transform.append(PickFirstChannels(n=2))
    spatial_transform.extend([ScaleValue(opt.value_scale), normalize])
    spatial_transform = Compose(spatial_transform)

    temporal_transform = []
    if opt.sample_t_stride > 1:
        temporal_transform.append(TemporalSubsampling(opt.sample_t_stride))
    temporal_transform.append(
        SlidingWindow(opt.sample_duration, opt.inference_stride))
    temporal_transform = TemporalCompose(temporal_transform)

    inference_data, collate_fn = get_inference_data(
        opt.video_path, opt.annotation_path, opt.dataset, opt.input_type,
        opt.file_type, opt.inference_subset, spatial_transform,
        temporal_transform)

    inference_loader = torch.utils.data.DataLoader(
        inference_data,
        batch_size=opt.inference_batch_size,
        shuffle=False,
        num_workers=opt.n_threads,
        pin_memory=True,
        worker_init_fn=worker_init_fn,
        collate_fn=collate_fn)

    return inference_loader, inference_data.class_names

def save_checkpoint(save_file_path, epoch, arch, model, optimizer, scheduler, use_ape=False, D=None):

    if use_ape:
        APE = model[0]
        model = model[1]

        if hasattr(APE, 'module'):
            ape_state_dict = APE.module.state_dict()
        else:
            ape_state_dict = APE.state_dict()

        if hasattr(D, 'module'):
            D_state_dict = D.module.state_dict()
        else:
            D_state_dict = D.state_dict()

    if hasattr(model, 'module'):
        model_state_dict = model.module.state_dict()
    else:
        model_state_dict = model.state_dict()

    save_states = {
        'epoch': epoch,
        'arch': arch,
        'state_dict': model_state_dict,
        'generator': ape_state_dict if use_ape else None,
        'discriminator': D_state_dict if D is not None else None,
        'optimizer': optimizer.state_dict(),
        'scheduler': scheduler.state_dict()
    }
    torch.save(save_states, save_file_path)

def main_worker(index, opt):
    random.seed(opt.manual_seed)
    np.random.seed(opt.manual_seed)
    torch.manual_seed(opt.manual_seed)

    if index >= 0 and opt.device.type == 'cuda':
        opt.device = torch.device(f'cuda:{index}')

    if opt.distributed:
        opt.dist_rank = opt.dist_rank * opt.ngpus_per_node + index
        dist.init_process_group(backend='nccl',
                                init_method=opt.dist_url,
                                world_size=opt.world_size,
                                rank=opt.dist_rank)
        opt.batch_size = int(opt.batch_size / opt.ngpus_per_node)
        opt.n_threads = int(
            (opt.n_threads + opt.ngpus_per_node - 1) / opt.ngpus_per_node)
    opt.is_master_node = not opt.distributed or opt.dist_rank == 0

    model = generate_model(opt)
    D = None
    ## 3D-APE-GAN
    if opt.use_ape:
        in_ch = 3
        G = Generator(in_ch).to(opt.device)
        D = Discriminator(in_ch).to(opt.device)

    if opt.batchnorm_sync:
        assert opt.distributed, 'SyncBatchNorm only supports DistributedDataParallel.'
        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
    if opt.pretrain_path:
        model = load_pretrained_model(model, opt.pretrain_path, opt.model,
                                      opt.n_finetune_classes)
    if opt.resume_path is not None:
        model = resume_model(opt.resume_path, opt.arch, model)

    model = make_data_parallel(model, opt.distributed, opt.device)

    if opt.pretrain_path:
        parameters = get_fine_tuning_parameters(model, opt.ft_begin_module)
    else:
        parameters = model.parameters()

    ## 3D-APE-GAN
    if opt.use_ape:
        G = make_data_parallel(G, opt.distributed, opt.device)
        D = make_data_parallel(D, opt.distributed, opt.device)
        if opt.ape_path is not None:
            checkpoint = torch.load(os.path.join(opt.ape_path))
            G.load_state_dict(checkpoint['generator'])
            D.load_state_dict(checkpoint['discriminator'])
#             G.eval()
        model = torch.nn.Sequential(G, model)

    # end-to-end G+C
#     parameters = model.parameters()

    if opt.is_master_node:
        print(model)

    criterion = CrossEntropyLoss().to(opt.device)

    if not opt.no_train:
        (train_loader, train_sampler, train_logger, train_batch_logger,
         optimizer, scheduler) = get_train_utils(opt, parameters)
        if opt.resume_path is not None:
            opt.begin_epoch, optimizer, scheduler = resume_train_utils(
                opt.resume_path, opt.begin_epoch, optimizer, scheduler)
            if opt.overwrite_milestones:
                scheduler.milestones = opt.multistep_milestones
    if not opt.no_val:
        val_loader, val_logger = get_val_utils(opt)

    if opt.tensorboard and opt.is_master_node:
        from torch.utils.tensorboard import SummaryWriter
        if opt.begin_epoch == 1:
            tb_writer = SummaryWriter(log_dir=opt.result_path)
        else:
            tb_writer = SummaryWriter(log_dir=opt.result_path,
                                      purge_step=opt.begin_epoch)
    else:
        tb_writer = None

    prev_val_loss = None

    ## Check if eps_range is set
    if opt.eps_range:
          eps_rvs = opt.eps - (loguniform.rvs(0.5, opt.eps+1, size=opt.n_epochs))
#         eps_rvs = opt.eps - (loguniform.rvs(1, opt.eps+1, size=opt.n_epochs)-1)
#         eps_rvs_1 = loguniform.rvs(1, opt.eps+1, size=opt.n_epochs)-1
#         eps_rvs_2 = 12 - (loguniform.rvs(1, opt.eps+1, size=opt.n_epochs)-1)
#         eps_rvs = np.append(eps_rvs_1[opt.n_epochs//2:], eps_rvs_2[:opt.n_epochs//2])

        # Beta distribution
#         a, b = 1.5,1.5
#         eps_rvs = beta.rvs(a, b, size=1000)*12
#         eps_rvs[eps_rvs<0.7] = 0
#         plt.hist(eps_rvs, density=True, histtype='stepfilled', alpha=0.6)
#         plt.savefig('beta_'+str(a)+'_'+str(b)+'.png')
        # Uniform, triangular distribution
#         eps_rvs = np.random.uniform(0, opt.eps, size=opt.n_epochs)
#         eps_rvs = np.random.triangular(0, 8, opt.eps, size=opt.n_epochs)
#         eps_rvs = opt.eps - eps_rvs

    for i in range(opt.begin_epoch, opt.n_epochs + 1):
        if not opt.no_train:
            '''
            Adaptive eps values
            '''
#             if opt.eps_range:

#                 if i == opt.begin_epoch:
#                     val_loader, val_logger = get_val_utils(opt)
#                     linf_attack_range = np.linspace(0, opt.eps, 10)
#                     (inputs, targets) = next(iter(val_loader))
#                     inputs = inputs.to(opt.device)
#                     targets = targets.to(opt.device, non_blocking=True)
#                     adversary = LinfPGDAttack(predict=model, loss_fn=criterion,
#                               eps=float(0), nb_iter=opt.attack_iter, eps_iter=float(opt.step_size))

#                 if i == opt.begin_epoch or (i-1) % 5 == 0:
#                     model.eval()
#                     eps_loss = []
#                     eps_prob = []
#                     for eps in linf_attack_range:

#                         adversary.eps = float(eps/255)
#                         adv_inputs = adversary.perturb(inputs, targets)

#                         outputs = model(adv_inputs)
#                         loss = criterion(outputs, targets).item()

#                         del adv_inputs; del outputs;

#                         eps_loss.append(loss)
# #                     eps_loss[-1] = eps_loss[-1] * 2
# #                     eps_loss[-2] = eps_loss[-2] * 2
#                     eps_loss[-1] = 0.6*sum(eps_loss[:-1])
#                     eps_prob = [ x / sum(eps_loss) for x in eps_loss]
#                     print(sum(eps_loss))
#                     model.train()

#                 eps_rvs = np.random.choice(linf_attack_range, p=eps_prob, size=len(train_loader))
#                 plt.clf()
#                 plt.hist(eps_rvs, density=True, histtype='stepfilled', alpha=1.0)
#                 plt.title(str(";".join([str(round(k,3)) for k in eps_loss])))
#                 plt.savefig(os.path.join(os.path.dirname(opt.result_path), 'adaptive_' + str(i) + '.png'))

            if opt.distributed:
                train_sampler.set_epoch(i)
            current_lr = get_lr(optimizer)
            train_epoch(i, train_loader, model, criterion, optimizer,
                        opt.device, current_lr, train_logger,
                        train_batch_logger, tb_writer, opt.distributed,
                        opt.attack_type, eps_rvs if opt.eps_range else opt.eps, opt.step_size, opt.attack_iter, opt.use_ape, D)

            if i % opt.checkpoint == 0 and opt.is_master_node:
                save_file_path = opt.result_path / 'save_{}.pth'.format(i)
                save_checkpoint(save_file_path, i, opt.arch, model, optimizer,
                                scheduler, opt.use_ape, D)

        if not opt.no_val:
            eps_rvs = np.random.choice(linf_attack_range, p=eps_prob, size=len(val_loader))
            prev_val_loss = val_epoch(i, val_loader, model, criterion,
                                      opt.device, val_logger, tb_writer,
                                      opt.distributed, opt.attack_type,
                                      eps_rvs if opt.eps_range else opt.eps, opt.step_size, opt.attack_iter)

        if not opt.no_train and opt.lr_scheduler == 'multistep':
            scheduler.step()
        elif not opt.no_train and opt.lr_scheduler == 'plateau':
            scheduler.step(prev_val_loss)

    if opt.inference:
        inference_loader, inference_class_names = get_inference_utils(opt)
        inference_result_path = opt.result_path / '{}.json'.format(
            opt.inference_subset)

        inference.inference(inference_loader, model, inference_result_path,
                            inference_class_names, opt.inference_no_average,
                            opt.output_topk, opt.attack_type, opt.eps,
                            opt.step_size, opt.attack_iter)

if __name__ == '__main__':
    opt = get_opt()

    opt.device = torch.device('cpu' if opt.no_cuda else 'cuda')
    if not opt.no_cuda:
        cudnn.benchmark = True
    if opt.accimage:
        torchvision.set_image_backend('accimage')

    opt.ngpus_per_node = torch.cuda.device_count()
    if opt.distributed:
        opt.world_size = opt.ngpus_per_node * opt.world_size
        mp.spawn(main_worker, nprocs=opt.ngpus_per_node, args=(opt,))
    else:
        main_worker(-1, opt)


import argparse
import os
import tensorflow as tf
import numpy as np
import sys

from pathlib import Path
from skimage import io

from toolbox import GPUselect

def main(imagePath, outputPath):
    """
    """
    model_path = Path(__file__).parent.joinpath('models', 'artefact', 'artefact_segmentation.h5')

    model = tf.keras.models.load_model(model_path)

    iname = imagePath.split('/')[-1]

    img = io.imread(imagePath)

    if img.ndim == 2 and img.shape == (256, 256):
        ch1 = np.zeros([256, 256], dtype=np.uint8)
        ch2 = np.zeros([256, 256], dtype=np.uint8)
        img = np.dstack([img, ch1, ch2])

    if img.ndim == 3:
        if img.shape[0] == 2:
            np.reshape(img, (1, 2, 0))
        if img.shape[2] == 2:
            ch = np.zeros([256, 256], dtype=np.uint8)
            img = np.dstack([img, ch1, ch2])

    predictions = model.predict(np.expand_dims(img, 0))[0]

    pred = np.zeros((predictions.shape[0], predictions.shape[1]))
    pred[predictions.argmax(2)==0] = 1

    img_masks = io.imread('/Users/guq/projects/hack2022-02-segmentation/probability_maps/%s_NucleiPM_1.tif' % (iname.rsplit('.', 1)[0]))

    masks = np.multiply(pred, img_masks)

    io.imsave(outputPath, masks)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("imagePath", help="path to the .tif file")
    parser.add_argument("--model",  help="type of model. For example, nuclei vs cytoplasm",default = 'nucleiDAPILAMIN')
    parser.add_argument("--outputPath", help="output path of probability map")
    parser.add_argument("--channel", help="channel to perform inference on",  nargs = '+', default=[0])
    # parser.add_argument("--channel2", help="channel2 to perform inference on", type=int, default=-1)
    parser.add_argument("--classOrder", help="background, contours, foreground", type = int, nargs = '+', default=-1)
    parser.add_argument("--mean", help="mean intensity of input image. Use -1 to use model", type=float, default=-1)
    parser.add_argument("--std", help="mean standard deviation of input image. Use -1 to use model", type=float, default=-1)
    parser.add_argument("--scalingFactor", help="factor by which to increase/decrease image size by", type=float,
                        default=1)
    parser.add_argument("--stackOutput", help="save probability maps as separate files", action='store_true')
    parser.add_argument("--GPU", help="explicitly select GPU", type=int, default = -1)
    parser.add_argument("--outlier", help="map percentile intensity to max when rescaling intensity values. Max intensity as default", type=float, default=-1)
    args = parser.parse_args()

    logPath = ''
    scriptPath = os.path.dirname(os.path.realpath(__file__))
    modelPath = os.path.join(scriptPath, 'models', args.model)
    # modelPath = os.path.join(scriptPath, 'models/cytoplasmINcell')
    # modelPath = os.path.join(scriptPath, 'cytoplasmZeissNikon')
    pmPath = ''
    if os.system('nvidia-smi') == 0:
        if args.GPU == -1:
            print("automatically choosing GPU")
            GPU = GPUselect.pick_gpu_lowest_memory()
        else:
            GPU = args.GPU
        print('Using GPU ' + str(GPU))

    else:
        if sys.platform == 'win32':  # only 1 gpu on windows
            if args.GPU==-1:
                GPU = 0
                print('using default GPU')
            else:
                GPU = args.GPU
            print('Using GPU ' + str(GPU))
        else:
            GPU=0
            print('Using CPU')
    os.environ['CUDA_VISIBLE_DEVICES'] = '%d' % GPU

    main(args.imagePath, args.outputPath)