import copy
import logging
import os
import time
import pandas as pd
import pickle
import psutil
import math
from abc import abstractmethod

from numpy import ndarray
from pandas import DataFrame, Series

from autogluon.common.utils.pandas_utils import get_approximate_df_mem_usage

from ...utils.exceptions import TimeLimitExceeded, NotEnoughMemoryError, NotEnoughCudaMemoryError
from ...utils import get_cpu_count, get_gpu_count_all
from ...utils.try_import import try_import_ray

logger = logging.getLogger(__name__)

TEXT_MODEL = 'TextPredictorModel'
IMAGE_MODEL = 'ImagePredictorModel'
TABULAR_MXNET_MODEL = 'TabularNeuralNetModel'
TABULAR_FASTAI_MODEL = 'NNFastAiTabularModel'

class AbstractFoldFittingStrategy():

    @abstractmethod
    def schedule_fold_model_fit(self, fold_ctx):
        """
        Schedule fold model training.
        By design this part is supposed to be 'lazy' evaluator,
        no actual training is performed here.
        Distributed fitters will handle jobs scheduling here.
        """

    @abstractmethod
    def after_all_folds_scheduled(self):
        """
        Method is called when all the folds are scheduled.
        Local fitters will perform training here.
        Distributed fitters will handle job handles and results retrieval here.
        """

    @abstractmethod
    def _fit(self, model_base, time_start_fold, time_limit_fold, fold_ctx, kwargs):
        """
        Method is called when a fold is ready to be fit
        """

class LocalFoldFittingStrategy(AbstractFoldFittingStrategy):
    """
    Provides some default implementation for AbstractFoldFittingStrategy

    Parameters
    ----------
        model_base: AbstractModel
            The template for the folds of model to be trained on.
        model_base_kwargs: dict
            kwargs required to initialize the model base when training the model.
        bagged_ensemble_model : BaggedEnsembleModel
            The ensemble model that holds all the trained folds.
        X : DataFrame
            The training data the model will be trained on.
        y: Series
            The labels of the training data.
        sample_weight:
            The sample weight of the training data.
        time_limit: float
            Approximately how long(in seconds) the fold fitting should be run for.
            If None, no time-constraint will be enforced allowing the folds to fully train.
        time_start: float
            Time starts to fit the BaggedEnsembleModel.
        models: list
            List of models that will be trained.
        oof_pred_proba: ndarray
            Out of folds predict probabilities that are already calculated.
        oof_pred_model_repeats: ndarray,
            Number of repeats the out of folds predict probabilities has been done.
        save_folds: bool,
            Whether to save the folds to disk.
        time_limit_fold_ratio: float, default=0.8
            The ratio of max time limit for each fold.
            If the estimated time required exceeds this ratio, will raise TimeLimitExceed error
    Attributes
    ----------
        X : DataFrame
            The training data the model will be trained on.
        y: Series
            The labels of the training data.
        sample_weight:
            The sample weight of the training data.
        time_limit: float
            Approximately how long(in seconds) the fold fitting should be run for.
            If None, no time-constraint will be enforced allowing the folds to fully train.
        time_start: float
            Time starts to fit the BaggedEnsembleModel.
        models: list
            List of models that will be trained.
        oof_pred_proba: ndarray
            Out of folds predict probabilities that are already calculated.
        oof_pred_model_repeats: ndarray,
            Number of repeats the out of folds predict probabilities has been done.
        jobs: list
            List of jobs that will be scheduled.
        save_folds: bool,
            Whether to save the folds to disk.
        time_limit_fold_ratio: float
            The ratio of max time limit for each fold.
    """

    def __init__(self, model_base, model_base_kwargs, bagged_ensemble_model,
                 X: DataFrame, y: Series, X_pseudo: DataFrame, y_pseudo: Series,
                 sample_weight, time_limit: float, time_start: float,
                 models: list, oof_pred_proba: ndarray, oof_pred_model_repeats: ndarray,
                 save_folds: bool, time_limit_fold_ratio=0.8):
        self.model_base = model_base
        self.model_base_kwargs = model_base_kwargs
        self.X = X
        self.y = y
        self.X_pseudo = X_pseudo
        self.y_pseudo = y_pseudo
        self.sample_weight = sample_weight
        self.time_limit = time_limit
        self.time_start = time_start
        self.models = models
        self.oof_pred_proba = oof_pred_proba
        self.oof_pred_model_repeats = oof_pred_model_repeats
        self.bagged_ensemble_model = bagged_ensemble_model
        self.jobs = []
        self.save_folds = save_folds
        self.time_limit_fold_ratio = time_limit_fold_ratio

    def schedule_fold_model_fit(self, fold_ctx):
        raise NotImplementedError

    def after_all_folds_scheduled(self):
        raise NotImplementedError

    def _get_fold_time_limit(self, fold_ctx):
        _, folds_finished, folds_left, folds_to_fit, _, _ = self._get_fold_properties(fold_ctx)
        time_elapsed = time.time() - self.time_start
        if self.time_limit is not None:
            time_left = self.time_limit - time_elapsed
            required_time_per_fold = time_left / folds_left
            time_limit_fold = required_time_per_fold * self.time_limit_fold_ratio
            if folds_finished > 0:
                expected_time_required = time_elapsed * folds_to_fit / folds_finished
                expected_remaining_time_required = expected_time_required * folds_left / folds_to_fit
                if expected_remaining_time_required > time_left:
                    raise TimeLimitExceeded
            if time_left <= 0:
                raise TimeLimitExceeded
        else:
            time_limit_fold = None
        return time_limit_fold

    def _update_bagged_ensemble(self, fold_model, pred_proba, fold_ctx):
        _, val_index = fold_ctx['fold']
        model_to_append = fold_model
        if not self.save_folds:
            fold_model.model = None
        if self.bagged_ensemble_model.low_memory:
            self.bagged_ensemble_model.save_child(fold_model, verbose=False)
            model_to_append = fold_model.name
        self.models.append(model_to_append)
        self.oof_pred_proba[val_index] += pred_proba
        self.oof_pred_model_repeats[val_index] += 1
        self.bagged_ensemble_model._add_child_times_to_bag(model=fold_model)

    def _predict_oof(self, fold_model, fold_ctx):
        time_train_end_fold = time.time()
        fold, folds_finished, folds_left, \
            folds_to_fit, is_last_fold, model_name_suffix = self._get_fold_properties(fold_ctx)
        _, val_index = fold
        X_val_fold = self.X.iloc[val_index, :]
        y_val_fold = self.y.iloc[val_index]
        # Check to avoid unnecessarily predicting and saving a model
        # when an Exception is going to be raised later
        if self.time_limit is not None:
            if not is_last_fold:
                time_elapsed = time.time() - self.time_start
                time_left = self.time_limit - time_elapsed
                expected_time_required = time_elapsed * folds_to_fit / (folds_finished + 1)
                expected_remaining_time_required = expected_time_required * (folds_left - 1) / folds_to_fit
                if expected_remaining_time_required > time_left:
                    raise TimeLimitExceeded
        pred_proba = fold_model.predict_proba(X_val_fold)
        fold_model.predict_time = time.time() - time_train_end_fold
        fold_model.val_score = fold_model.score_with_y_pred_proba(y=y_val_fold,
                                                                  y_pred_proba=pred_proba)
        fold_model.reduce_memory_size(remove_fit=True, remove_info=False, requires_save=True)
        if not self.bagged_ensemble_model.params.get('save_bag_folds', True):
            fold_model.model = None
        return fold_model, pred_proba

    @staticmethod
    def _get_fold_properties(fold_ctx):
        fold, folds_finished, folds_left, \
            folds_to_fit, is_last_fold, model_name_suffix = [
                fold_ctx[f] for f in ['fold', 'folds_finished', 'folds_left', 'folds_to_fit', 'is_last_fold', 'model_name_suffix']]
        return fold, folds_finished, folds_left, folds_to_fit, is_last_fold, model_name_suffix

class SequentialLocalFoldFittingStrategy(LocalFoldFittingStrategy):
    """
    This strategy fits the folds locally in a sequence.
    """
    def schedule_fold_model_fit(self, fold_ctx):
        self.jobs.append(fold_ctx)

    def after_all_folds_scheduled(self):
        for job in self.jobs:
            self._fit_fold_model(job)

    def _fit_fold_model(self, fold_ctx):
        time_start_fold = time.time()
        time_limit_fold = self._get_fold_time_limit(fold_ctx)
        fold_model = self._fit(self.model_base, time_start_fold, time_limit_fold, fold_ctx, self.model_base_kwargs)
        fold_model, pred_proba = self._predict_oof(fold_model, fold_ctx)
        self._update_bagged_ensemble(fold_model, pred_proba, fold_ctx)

    def _fit(self, model_base, time_start_fold, time_limit_fold, fold_ctx, kwargs):
        fold, folds_finished, folds_left, folds_to_fit, is_last_fold, model_name_suffix = self._get_fold_properties(fold_ctx)
        train_index, val_index = fold
        X_fold, X_val_fold = self.X.iloc[train_index, :], self.X.iloc[val_index, :]
        y_fold, y_val_fold = self.y.iloc[train_index], self.y.iloc[val_index]
        fold_model = copy.deepcopy(model_base)
        fold_model.name = f'{fold_model.name}{model_name_suffix}'
        fold_model.set_contexts(self.bagged_ensemble_model.path + fold_model.name + os.path.sep)
        kwargs_fold = kwargs.copy()
        is_pseudo = self.X_pseudo is not None and self.y_pseudo is not None
        if self.sample_weight is not None:
            kwargs_fold['sample_weight'] = self.sample_weight[train_index]
            kwargs_fold['sample_weight_val'] = self.sample_weight[val_index]

            if is_pseudo:
                # TODO: Add support for sample_weight when pseudo is present
                raise Exception('Sample weights given, but not used due to pseudo labelled data being given.')
            else:
                kwargs_fold['sample_weight'] = self.sample_weight[train_index]
                kwargs_fold['sample_weight_val'] = self.sample_weight[val_index]

        if is_pseudo:
            logger.log(15, f'{len(self.X_pseudo)} extra rows of pseudolabeled data added to training set for {fold_model.name}')
            X_fold = pd.concat([X_fold, self.X_pseudo], axis=0, ignore_index=True)
            y_fold = pd.concat([y_fold, self.y_pseudo], axis=0, ignore_index=True)

        fold_model.fit(X=X_fold, y=y_fold, X_val=X_val_fold, y_val=y_val_fold, time_limit=time_limit_fold, **kwargs_fold)
        fold_model.fit_time = time.time() - time_start_fold
        return fold_model

def _ray_fit(model_base, bagged_ensemble_model_path,
             X, y, X_pseudo, y_pseudo,
             fold_ctx, time_limit_fold, num_cpus,
             save_bag_folds, kwargs_fold):
    logger.log(10, 'ray worker training')
    time_start_fold = time.time()
    fold, folds_finished, folds_left, \
        folds_to_fit, is_last_fold, \
        model_name_suffix = LocalFoldFittingStrategy._get_fold_properties(fold_ctx)
    train_index, val_index = fold
    fold_model = copy.deepcopy(model_base)
    fold_model.name = f'{fold_model.name}{model_name_suffix}'
    fold_model.set_contexts(bagged_ensemble_model_path + fold_model.name + os.path.sep)
    if type(X) == str and type(y) == str:
        with open(X, 'rb') as X_f, open(y, 'rb') as y_f:
            X = pickle.load(X_f)
            y = pickle.load(y_f)
    is_pseudo = False
    if X_pseudo is not None and y_pseudo is not None:
        if type(X_pseudo) == str and type(y_pseudo) == str:
            with open(X_pseudo, 'rb') as X_pseudo_f, open(y_pseudo, 'rb') as y_pseudo_f:
                X_pseudo = pickle.load(X_pseudo_f)
                y_pseudo = pickle.load(y_pseudo_f)
        is_pseudo = True

    X_fold, X_val_fold = X.iloc[train_index, :], X.iloc[val_index, :]
    y_fold, y_val_fold = y.iloc[train_index], y.iloc[val_index]
    if is_pseudo:
            logger.log(15, f'{len(X_pseudo)} extra rows of pseudolabeled data added to training set for {fold_model.name}')
            X_fold = pd.concat([X_fold, X_pseudo], axis=0, ignore_index=True)
            y_fold = pd.concat([y_fold, y_pseudo], axis=0, ignore_index=True)
    fold_model.fit(X=X_fold, y=y_fold, X_val=X_val_fold, y_val=y_val_fold,
                   time_limit=time_limit_fold, num_cpus=num_cpus, **kwargs_fold)
    time_train_end_fold = time.time()
    fold_model.fit_time = time_train_end_fold - time_start_fold
    fold_model, pred_proba = _ray_predict_oof(fold_model, X_val_fold, y_val_fold,
                                              time_train_end_fold, num_cpus, save_bag_folds)
    fold_model.save()
    return fold_model.name, pred_proba, time_start_fold, \
        time_train_end_fold, fold_model.predict_time, fold_model.predict_1_time

def _ray_predict_oof(fold_model, X_val_fold, y_val_fold, time_train_end_fold,
                     num_cpus=-1, save_bag_folds=True):
    pred_proba = fold_model.predict_proba(X_val_fold, num_cpus=num_cpus)
    time_pred_end_fold = time.time()
    fold_model.predict_time = time_pred_end_fold - time_train_end_fold
    fold_model.val_score = fold_model.score_with_y_pred_proba(y=y_val_fold,
                                                              y_pred_proba=pred_proba)
    fold_model.reduce_memory_size(remove_fit=True, remove_info=False, requires_save=True)
    if not save_bag_folds:
        fold_model.model = None
    return fold_model, pred_proba

class ParallelLocalFoldFittingStrategy(LocalFoldFittingStrategy):
    """
    An implementation of LocalFoldFittingStrategy to train multiple folds in parallel.
    Folds are spread to cpu/gpu cores by ray tasks.
    Large data are stored in ray object store, which minimizes memory usage and unessary serializations.
    Trained models are saved to disk within each ray task.

    Parameters
    ----------
        num_folds_parallel: int
            Number of folds to train in parallel at once.
            Consider lower this parameter if encounter out of memory issue.
        max_memory_usage_ratio: float, default=0.8
            The ratio of max memory usage for parallel folding.
            If the estimated usage exceeds this ratio, will fall back to sequential folding.
    Attributes
    ----------
        num_cpus: int
            Number of cpu cores available.
        num_gpus: int
            Number of gpus available.
        num_parallel_jobs: int
            Number of parallel jobs to be executed once.
        max_memory_usage_ratio: float
            The ratio of max memory usage for parallel folding.
        time_start_fit: float
            The time of the first model starts training.
        time_end_fit: float
            The time of the last model finishes training.
        fit_time: float
            The amount of time used to fit all folds.
        predict_time: float
            The amount of time used to do out of folds predictions for all folds.
    """
    def __init__(self, num_folds_parallel: int, max_memory_usage_ratio=0.8, **kwargs):
        super().__init__(**kwargs)
        self.ray = try_import_ray()
        self.num_cpus = get_cpu_count()
        self.num_gpus = None
        self.num_parallel_jobs = num_folds_parallel
        self.max_memory_usage_ratio = min(max_memory_usage_ratio, 1.0)
        self.time_start_fit = None
        self.time_end_fit = None
        self.fit_time = 0
        self.predict_time = 0
        self.predict_1_time = None
        # max_calls to guarantee release of gpu resource
        self._ray_fit = self.ray.remote(max_calls=1)(_ray_fit)
        # initialize the model base to get necessary info for estimating memory usage
        self._initialized_model_base = copy.deepcopy(self.model_base)
        self._initialized_model_base.initialize(X=self.X, y=self.y, **self.model_base_kwargs)

    def is_mem_sufficient(self, num_jobs):
        '''Check if the memory is sufficient to do parallel training'''
        model_mem_est = self._initialized_model_base.estimate_memory_usage(X=self.X)
        _, _, num_parallel_jobs = self._get_resource_suggestions(num_jobs)
        total_model_mem_est = num_parallel_jobs * model_mem_est
        data_mem_est = self._estimate_data_memory_usage()
        total_data_mem_est = num_parallel_jobs * data_mem_est
        mem_available = psutil.virtual_memory().available
        return (mem_available * self.max_memory_usage_ratio) > (total_model_mem_est + total_data_mem_est)

    def _estimate_data_memory_usage(self):
        X_mem = get_approximate_df_mem_usage(self.X).sum()
        y_mem = get_approximate_df_mem_usage(self.y.to_frame()).sum()
        return X_mem + y_mem

    def _get_gpu_count(self):
        if not self.num_gpus:
            _user_gpu_count = self.model_base.get_params().get('hyperparameters', {}). \
                get('ag_args_fit', {}).get('num_gpus', 0)
            model_base_class_name = self.model_base.__class__.__name__
            # If user didn't specify num gpus and we are training text or image models,
            # use all gpus.
            if model_base_class_name in [TEXT_MODEL, IMAGE_MODEL] and _user_gpu_count == 0:
                self.num_gpus = get_gpu_count_all()
            else:
                self.num_gpus = min(_user_gpu_count, get_gpu_count_all())

    def schedule_fold_model_fit(self, fold_ctx):
        self._get_gpu_count()
        if not self.ray.is_initialized():
            if self.num_gpus > 0:
                self.ray.init(log_to_driver=False, num_gpus=self.num_gpus)
            else:
                self.ray.init(log_to_driver=False)
        self.jobs.append(fold_ctx)

    def after_all_folds_scheduled(self):
        num_jobs = len(self.jobs)
        resources, _, _ = self._get_resource_suggestions(num_jobs)
        job_refs = []
        job_fold_map = {}
        # prepare shared data
        X, y, X_pseudo, y_pseudo = self._prepare_data()
        model_base_ref = self.ray.put(self.model_base)
        time_limit_fold = self._get_fold_time_limit(num_jobs)
        # spread the task
        for job in self.jobs:
            fold_ctx = job
            ref = self._fit(model_base_ref, X, y, X_pseudo, y_pseudo, time_limit_fold, fold_ctx, resources, self.model_base_kwargs)
            job_fold_map[ref] = fold_ctx
            job_refs.append(ref)

        # update ensemble whenver a model return
        unfinished = job_refs
        while unfinished:
            finished, unfinished = self.ray.wait(unfinished, num_returns=1)
            finished = finished[0]
            try:
                fold_model, pred_proba, time_start_fit, \
                    time_end_fit, predict_time, predict_1_time = self.ray.get(finished)
                fold_ctx = job_fold_map.get(finished, None)
                assert fold_ctx is not None
                self._update_bagged_ensemble(fold_model=fold_model,
                                             pred_proba=pred_proba,
                                             time_start_fit=time_start_fit,
                                             time_end_fit=time_end_fit,
                                             predict_time=predict_time,
                                             predict_1_time=predict_1_time,
                                             fold_ctx=fold_ctx)
            except TimeLimitExceeded:
                # Terminate all ray tasks because a fold failed
                self.ray.shutdown()
                raise TimeLimitExceeded
            # NotEnoughMemoryError is an autogluon custom error,
            # it predict memory usage before hand
            # MemoryError is the actual python memory error if the process failed
            except (NotEnoughMemoryError, MemoryError):
                error_msg = 'Consider decrease folds trained in parallel \
                             by passing num_fold_parallel to ag_args_ensemble \
                             when calling tabular.fit.\n\
                             If none working, use sequential folding by passing \
                             SequentialLocalFoldFittingStrategy to ag_args_ensemble \
                             when calling tabular.fit and try again.'
                logger.warning(error_msg)
                # Terminate all ray tasks because a fold failed
                self.ray.shutdown()
                raise NotEnoughMemoryError
            except Exception as e:
                processed_exception = self._parse_ray_error(e)
                # Terminate all ray tasks because a fold failed
                self.ray.shutdown()
                raise processed_exception
        self.fit_time = 0
        if self.time_start_fit and self.time_end_fit:
            self.fit_time = self.time_end_fit - self.time_start_fit
        self.bagged_ensemble_model._add_parallel_child_times(fit_time=self.fit_time, predict_time=self.predict_time, predict_1_time=self.predict_1_time)

    def _fit(self, model_base_ref, X_ref, y_ref, X_pseudo_ref, y_pseudo_ref, time_limit_fold, fold_ctx, resources, kwargs):
        fold, folds_finished, folds_left, \
            folds_to_fit, is_last_fold, \
            model_name_suffix = self._get_fold_properties(fold_ctx)
        train_index, val_index = fold
        fold_ctx_ref = self.ray.put(fold_ctx)
        save_bag_folds = self.bagged_ensemble_model.params.get('save_bag_folds', True)
        kwargs_fold = kwargs.copy()
        is_pseudo = X_pseudo_ref is not None and y_pseudo_ref is not None
        if self.sample_weight is not None:
            if is_pseudo:
                # TODO: Add support for sample_weight when pseudo is present
                raise Exception('Sample weights given, but not used due to pseudo labelled data being given.')
            else:
                kwargs_fold['sample_weight'] = self.sample_weight[train_index]
                kwargs_fold['sample_weight_val'] = self.sample_weight[val_index]
        num_cpus = resources.get('num_cpus', 0)
        return self._ray_fit.options(**resources) \
            .remote(model_base_ref, self.bagged_ensemble_model.path,
                    X_ref, y_ref, X_pseudo_ref, y_pseudo_ref, fold_ctx_ref, time_limit_fold, num_cpus,
                    save_bag_folds, kwargs_fold)

    def _update_bagged_ensemble(self, fold_model, pred_proba, time_start_fit,
                                time_end_fit, predict_time, predict_1_time, fold_ctx):
        _, val_index = fold_ctx['fold']
        self.models.append(fold_model)
        self.oof_pred_proba[val_index] += pred_proba
        self.oof_pred_model_repeats[val_index] += 1
        if self.time_start_fit:
            self.time_start_fit = min(time_start_fit, self.time_start_fit)
        else:
            self.time_start_fit = time_start_fit
        if self.time_end_fit:
            self.time_end_fit = max(time_end_fit, self.time_end_fit)
        else:
            self.time_end_fit = time_end_fit
        self.predict_time += predict_time
        if predict_1_time is not None:
            if self.predict_1_time is None:
                self.predict_1_time = 0
            self.predict_1_time += predict_1_time

    def _get_fold_time_limit(self, num_jobs):
        _, batches, _ = self._get_resource_suggestions(num_jobs)
        time_elapsed = time.time() - self.time_start
        if self.time_limit is not None:
            time_left = self.time_limit - time_elapsed
            required_time_per_fold = time_left / batches
            time_limit_fold = required_time_per_fold * self.time_limit_fold_ratio
            if time_left <= 0:
                raise TimeLimitExceeded
        else:
            time_limit_fold = None
        return time_limit_fold

    def _get_resource_suggestions(self, num_jobs):
        # TODO: We need to handle user provide custom num_cpus
        cpu_per_job = max(1, int(self.num_cpus // self.num_parallel_jobs))
        gpu_per_job = 0
        resources = dict(num_cpus=cpu_per_job)
        num_parallel_jobs = min(self.num_parallel_jobs, self.num_cpus // cpu_per_job)
        self.num_parallel_jobs = num_parallel_jobs
        batches = math.ceil(num_jobs / num_parallel_jobs)
        if self.num_gpus and self.num_gpus > 0:
            gpu_per_job = self.num_gpus / self.num_parallel_jobs
            # For Image and Text model,
            # we always guarantee a task has at least one full gpu to use
            # FIXME: Avoid hardcoding model names.
            if self.model_base.__class__.__name__ in [TEXT_MODEL, IMAGE_MODEL]:
                gpu_per_job = max(1, math.ceil(gpu_per_job))
            resources = dict(num_cpus=cpu_per_job, num_gpus=gpu_per_job)
        return resources, batches, num_parallel_jobs

    def _prepare_data(self, in_mem=True):
        X_pseudo = None
        y_pseudo = None
        if in_mem:
            X = self.ray.put(self.X)
            y = self.ray.put(self.y)
            if self.X_pseudo is not None and self.y_pseudo is not None:
                X_pseudo = self.ray.put(self.X_pseudo)
                y_pseudo = self.ray.put(self.y_pseudo)
        else:
            X = 'X.pkl'
            y = 'y.pkl'
            utils = 'utils'
            X = os.path.join(self.bagged_ensemble_model.path, utils, X)
            y = os.path.join(self.bagged_ensemble_model.path, utils, y)
            with open(X, 'wb') as X_f, open(y, 'wb') as y_f:
                pickle.dump(self.X, X_f)
                pickle.dump(self.y, y_f)
            if self.X_pseudo is not None and self.y_pseudo is not None:
                X_pseudo = 'X_pseudo.pkl'
                y_pseudo = 'y_pseudo.pkl'
                X_pseudo = os.path.join(self.bagged_ensemble_model.path, utils, X_pseudo)
                y_pseudo = os.path.join(self.bagged_ensemble_model.path, utils, y_pseudo)
        return X, y, X_pseudo, y_pseudo

    def _parse_ray_error(self, e):
        error = str(e).lower()
        if 'cuda' in error and ('out of memory' in error or 'alloc' in error):
            default_error_msg = 'If none working, use sequential folding by passing \
                         SequentialLocalFoldFittingStrategy to ag_args_ensemble \
                         when calling tabular.fit and try again.'
            # FIXME: Avoid hardcoding model names.
            if self.model_base.__class__.__name__ in [TEXT_MODEL, IMAGE_MODEL]:
                error_msg = f'Out of CUDA memory while training \
                            {self.model_base.__class__.__name__}. \
                            Consider decrease batch size in hyperparameter and try again.\n\
                            Or decrease folds trained in parallel by passing num_fold_parallel \
                            to ag_args_ensemble when calling tabular.fit if you have multiple \
                            gpus and try again'
                logger.warning(error_msg)
            # FIXME: Avoid hardcoding model names.
            elif self.model_base.__class__.__name__ in [TABULAR_MXNET_MODEL, TABULAR_FASTAI_MODEL]:
                error_msg = f'Out of CUDA memory while training \
                            {self.model_base.__class__.__name__}. \
                            Consider decrease batch size in hyperparameter and try again.\n\
                            Or decrease folds trained in parallel by passing num_fold_parallel \
                            to ag_args_ensemble when calling tabular.fit and try again'
                logger.warning(error_msg)
            else:
                error_msg = f'Out of CUDA memory while training \
                            {self.model_base.__class__.__name__}. \
                            Consider decrease folds trained in parallel by passing \
                            num_fold_parallel to ag_args_ensemble when calling tabular.fit \
                            and try again'
                logger.warning(error_msg)
            logger.warning(default_error_msg)
            e = NotEnoughCudaMemoryError
        return e


import argparse
from pathlib import Path
from typing import Dict, List, Optional

import numpy as np
import shogi
import shogi.CSA
from shogi.CSA import Parser
from tqdm import tqdm

from library.core import GameResult, Notation, read_csa, split_sfen

class MakeSfen:
    __MAX_VALUE = 100000

    def __init__(self) -> None:
        # 最小レーティング
        self.__min_rate: int = 0

        # 最大レーティング
        self.__max_rate: int = 10000

        # 先手の最小評価値
        self.__min_black_value: int = -MakeSfen.__MAX_VALUE

        # 先手の最大評価値
        self.__max_black_value: int = MakeSfen.__MAX_VALUE

        # 後手の最小評価値
        self.__min_white_value: int = -MakeSfen.__MAX_VALUE

        # 後手の最大評価値
        self.__max_white_value: int = MakeSfen.__MAX_VALUE

        # 千日手での最小評価値
        # Noneの時は、先後の最小評価値に従う。
        self.__min_draw_value: Optional[int] = None

        # 千日手での最大評価値
        self.__max_draw_value: int = MakeSfen.__MAX_VALUE

        # 最大差分評価値
        self.__max_diff_value: int = MakeSfen.__MAX_VALUE

        # 出力対象とする最大評価値
        self.__end_value: int = MakeSfen.__MAX_VALUE

    def set_rate_limit(self, min_rate: int, max_rate: int) -> None:
        if min_rate < 0:
            raise ValueError(f'最低レーティングには、0以上の数値を指定してください。: {min_rate}')
        if max_rate < min_rate:
            raise ValueError(f'最高レーティングには、最低レーティング以上の数値を指定してください。: {max_rate}')

        self.__min_rate = min_rate
        self.__max_rate = max_rate

    def set_result_filter(self, result: GameResult, min_value: Optional[int], max_value: Optional[int]) -> None:
        if min_value is max_value is not None and min_value > max_value:
            raise ValueError(f'最小評価値は最大評価値以下の数値を指定してください。: {min_value}')

        if result == GameResult.BLACK_WIN:
            self.__min_black_value = -MakeSfen.__MAX_VALUE if min_value is None else min_value
            self.__max_black_value = MakeSfen.__MAX_VALUE if max_value is None else max_value
        elif result == GameResult.WHITE_WIN:
            self.__min_white_value = -MakeSfen.__MAX_VALUE if max_value is None else -max_value
            self.__max_white_value = MakeSfen.__MAX_VALUE if min_value is None else -min_value
        elif result == GameResult.SENNICHITE:
            self.__min_draw_value = min_value
            self.__max_draw_value = MakeSfen.__MAX_VALUE if max_value is None else max_value

    def set_diff_value_limit(self, diff: int) -> None:
        self.__max_diff_value = diff

    def set_value_limit(self, end_value: int) -> None:
        self.__end_value = end_value

    def run(self, csa_path: Path) -> List[str]:
        if csa_path.is_dir():
            csa_files: List[Path] = [x for x in tqdm(csa_path.glob('**/*.csa'))]
        elif csa_path.is_file():
            csa_files = [csa_path]
        else:
            raise ValueError(f'csaファイルが存在しません。: {csa_path}')

        sfens: Dict[str, int] = {'lnsgkgsnl/1r5b1/ppppppppp/9/9/9/PPPPPPPPP/1B5R1/LNSGKGSNL b -': 1}

        for csa_file in tqdm(csa_files):
            try:
                notation = read_csa(csa_file)
            except ValueError:
                continue

            x = self.__worker(notation)

            if x is None:
                continue

            for sfen in x:
                position, game_ply = split_sfen(f'sfen {sfen}')
                sfens[position] = min(sfens.get(position, game_ply), game_ply)

        return [f'sfen {position} {game_ply}' for position, game_ply in sfens.items()]

    def __worker(self, notation: Notation) -> Optional[List[str]]:  # noqa: C901
        # 指定範囲以外のレーティングの場合は除外
        if not self.__min_rate <= max(notation.black.rate, notation.white.rate) <= self.__max_rate:
            return

        # 投了や入玉宣言勝ち、千日手、最大手数制限以外で終局した場合は除外
        if notation.result == GameResult.UNKNOWN:
            return

        # 評価値
        values = np.asarray([move.value for move in notation.moves], dtype=np.int32)
        black_values = values[::2]
        white_values = values[1::2]

        result = notation.result

        # 先手勝利
        if notation.result == GameResult.BLACK_WIN and not self.__is_output_black(black_values):
            return

        # 後手勝利
        elif notation.result == GameResult.WHITE_WIN and not self.__is_output_white(white_values):
            return

        # 千日手
        elif notation.result == GameResult.SENNICHITE and not self.__is_output_sennichite(values, black_values, white_values):
            return

        # 最大手数制限
        elif notation.result == GameResult.MAX_MOVES:
            # 先手勝勢
            if self.__is_output_black(black_values):
                result = GameResult.BLACK_WIN

            # 後手勝勢
            elif self.__is_output_white(white_values):
                result = GameResult.WHITE_WIN

            # その他
            else:
                return

        board = shogi.Board()
        sfens = []

        # 棋譜出力
        for i, move_data in enumerate(notation.moves):
            _, move = Parser.parse_move_str(move_data.move, board)
            board.push_usi(move)

            # 勝利（最大手数制限の場合は勝勢）側の棋譜のみを出力
            # 千日手の場合は両方
            x = 1 if i % 2 == 0 else -1
            if (result == GameResult.BLACK_WIN and x == 1) or (result == GameResult.WHITE_WIN and x == -1) or result == GameResult.SENNICHITE:
                # 評価値上限の場合は、それ以降の棋譜を出力しない。
                if move_data.value * x > self.__end_value:
                    break

                sfens.append(board.sfen())

        return sfens

    def __is_output_black(self, black_values) -> bool:
        # 全ての評価値が0の場合
        if np.all(black_values == 0):
            return False

        nonzero_black_values = black_values[black_values != 0]

        # 評価値0以外の指し手が手数の半分未満の場合
        if np.size(nonzero_black_values) < np.size(black_values) * 0.5:
            return False

        # 指定値より不利になった場合や投了時の評価値が指定値未満の場合
        if np.min(black_values) < self.__min_black_value or np.max(black_values) < self.__max_black_value:
            return False

        # 直前の評価値との差が指定値以上に不利になった場合
        diff = np.diff(nonzero_black_values)
        if np.min(diff) <= -self.__max_diff_value:
            return False

        return True

    def __is_output_white(self, white_values) -> bool:
        # 全ての評価値が0の場合
        if np.all(white_values == 0):
            return False

        nonzero_white_values = white_values[white_values != 0]

        # 評価値0以外の指し手が手数の半分未満の場合
        if np.size(nonzero_white_values) < np.size(white_values) * 0.5:
            return False

        # 指定値より不利になった場合や投了時の評価値が指定値未満の場合
        if np.max(white_values) > self.__max_white_value or np.min(white_values) > self.__min_white_value:
            return False

        # 直前の評価値との差が指定値以上に不利になった場合
        diff = np.diff(nonzero_white_values)
        if np.max(diff) >= self.__max_diff_value:
            return False

        return True

    def __is_output_sennichite(self, values, black_values, white_values) -> bool:
        # 全ての評価値が0の場合
        if np.all(values == 0):
            return False

        # 先手: 指定値より形勢に差がある場合（有利や不利になった場合）
        min_black_value: int = self.__min_black_value if self.__min_draw_value is None else self.__min_draw_value
        if np.min(black_values) < min_black_value or np.max(black_values) > self.__max_draw_value:
            return False

        # 後手: 指定値より形勢に差がある場合（有利や不利になった場合）
        max_white_value: int = self.__max_white_value if self.__min_draw_value is None else -self.__min_draw_value  # pylint: disable=invalid-unary-operand-type
        if np.max(white_values) > max_white_value or np.min(white_values) < -self.__max_draw_value:
            return False

        return True

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('input', type=Path, help='csaファイルのディレクトリ')
    parser.add_argument('-o', '--output', type=Path, default='book.sfen', help='sfenの出力先')

    rate_group = parser.add_argument_group('レーティング')
    rate_group.add_argument('--min_rate', type=int, default=0, help='レーティング最小値')
    rate_group.add_argument('--max_rate', type=int, default=5000, help='レーティング最大値')

    value_group = parser.add_argument_group('評価値')
    value_group.add_argument('--min_black_value', type=int, default=-10, help='先手勝利時における、先手側から見た最小評価値')
    value_group.add_argument('--max_black_value', type=int, default=100000, help='先手勝利時における、先手側から見た最大評価値')
    value_group.add_argument('--min_white_value', type=int, default=-150, help='後手勝利時における、後手側から見た最小評価値')
    value_group.add_argument('--max_white_value', type=int, default=100000, help='後手勝利時における、後手側から見た最大評価値')
    value_group.add_argument('--draw_value', type=int, default=150, help='千日手における、手番側から見た最大評価値')
    value_group.add_argument('--diff_value', type=int, default=100000, help='直前の評価値との差')
    value_group.add_argument('--end_value', type=int, default=0, help='勝利側から見た終局時の最小評価値')

    args = parser.parse_args()
    output_path: Path = args.output

    make = MakeSfen()
    make.set_rate_limit(args.min_rate, args.max_rate)
    make.set_result_filter(GameResult.BLACK_WIN, args.min_black_value, args.max_black_value)
    make.set_result_filter(GameResult.WHITE_WIN, args.min_white_value, args.max_white_value)
    make.set_result_filter(GameResult.SENNICHITE, None, args.draw_value)
    make.set_diff_value_limit(args.diff_value)
    make.set_value_limit(args.end_value)

    sfens: List[str] = make.run(args.input)

    # ファイル出力
    with output_path.open('w', encoding='utf_8') as f:
        for sfen in tqdm(sfens):
            f.write(f'{sfen}\n')


from scipy.misc import imread, imresize, imsave, fromimage, toimage
from sklearn.feature_extraction.image import reconstruct_from_patches_2d, extract_patches_2d
from PIL import Image
import numpy as np
import os

from keras import backend as K

aspect_ratio = 0
img_WIDTH = img_HEIGHT = 0, 0

# Util function to open, resize and format pictures into appropriate tensors
def preprocess_image(image_path, img_width=256, img_height=256, load_dims=False, resize=True, size_multiple=4):
    '''
    Preprocess the image so that it can be used by Keras.

    Args:
        image_path: path to the image
        img_width: image width after resizing. Optional: defaults to 256
        img_height: image height after resizing. Optional: defaults to 256
        load_dims: decides if original dimensions of image should be saved,
                   Optional: defaults to False
        vgg_normalize: decides if vgg normalization should be applied to image.
                       Optional: defaults to False
        resize: whether the image should be resided to new size. Optional: defaults to True
        size_multiple: Deconvolution network needs precise input size so as to
                       divide by 4 ("shallow" model) or 8 ("deep" model).

    Returns: an image of shape (3, img_width, img_height) for dim_ordering = "th",
             else an image of shape (img_width, img_height, 3) for dim ordering = "tf"

    '''
    img = imread(image_path, mode="RGB")  # Prevents crashes due to PNG images (ARGB)
    if load_dims:
        global img_WIDTH, img_HEIGHT, aspect_ratio
        img_WIDTH = img.shape[0]
        img_HEIGHT = img.shape[1]
        aspect_ratio = img_HEIGHT / img_WIDTH

    if resize:
        if img_width < 0 or img_height < 0: # We have already loaded image dims
            img_width = (img_WIDTH // size_multiple) * size_multiple # Make sure width is a multiple of 4
            img_height = (img_HEIGHT // size_multiple) * size_multiple # Make sure width is a multiple of 4
        img = imresize(img, (img_width, img_height))

    if K.image_dim_ordering() == "th":
        img = img.transpose((2, 0, 1)).astype(np.float32)
    else:
        img = img.astype(np.float32)

    img = np.expand_dims(img, axis=0)
    return img

# Util function to convert a tensor into a valid image
def deprocess_image(x):
    '''
    Removes the pre processing steps applied to image.

    Args:
        x: input image of shape (3, img_width, img_height) [th],
           or input image of shape (img_width, img_height, 3) [tf]
        denormalize_vgg: whether vgg normalization should be reversed

    Returns: image of same shape as input shape

    '''
    if K.image_dim_ordering() == "th":
        x = x.transpose((1, 2, 0))

    x = np.clip(x, 0, 255).astype('uint8')
    return x

# Util function to preserve image color
def original_color_transform(content, generated):
    '''
    Applies the color space of content image to the generated image

    Args:
        content: input image of shape (img_width, img_height, 3)
        generated: input image of shape (img_width, img_height, 3)

    Returns: image of same shape as input shape

    '''
    generated = fromimage(toimage(generated), mode='YCbCr')  # Convert to YCbCr color space
    generated[:, :, 1:] = content[:, :, 1:]  # Generated CbCr = Content CbCr
    generated = fromimage(toimage(generated, mode='YCbCr'), mode='RGB')  # Convert to RGB color space
    return generated

# Util function to save intermediate images
def save_result(img, fname, img_width=None, img_height=None, preserve_color=False, content_img_path=None, directory=None):
    '''
    Save the resultant image

    Args:
        img: input image of shape (img_width, img_height, 3)
        fname: filename of output image
        img_width: resize dimension
        img_height: resize dimension
        preserve_color: whether to preserve original color of the content image
        content_img_path: path to content image. Optional, but required if color preservation is required
        directory: base directory where image will be stored

    '''
    if directory is not None:
        if not os.path.exists(directory):
            os.makedirs(directory)

        fname = directory + fname

    # We require original image if we are to preserve color in YCbCr mode
    if preserve_color:
        assert content_img_path is not None, \
            "If color is to be preserved, then content image path must be given as well."

        content = imread(content_img_path, mode="YCbCr")
        if img_width is not None and img_height is not None:
            content = imresize(content, (img_width, img_height))
        img = original_color_transform(content, img)

    imsave(fname, img)

def _check_image(path, i, nb_images):
    '''
    Test if image can be loaded by PIL. If image cannot be loaded, delete it from dataset.

    Args:
        path: path to image
        i: iteration number
        nb_images: total number of images

    '''
    try:
        im = Image.open(path)
        im.verify()
        im = Image.open(path)
        im.load()
        if i % 1000 == 0: print('%0.2f percent images are checked.' % (i * 100 / nb_images))
    except:
        os.remove(path)
        print("Image number %d is corrupt (path = %s). Deleting from dataset." % (i, path))

def check_dataset(path):
    '''
    Use to check the dataset for any jpg corruptions.
    If there exists corruption in an image, then it is deleted.

    Note:
        If due to some reason the corrupted file cannot be deleted via os.remove(),
        it will print the path of the file. Please delete the file manually in such a case.

    Args:
        path: Path to the dataset of images
    '''
    from multiprocessing.pool import Pool
    pool = Pool()

    nb_images = len([name for name in os.listdir(path)])
    print("Checking %d images" % nb_images)

    for i, file in enumerate(os.listdir(path)):
        pool.apply_async(_check_image, args=(path + "\\" + file, i, nb_images))

        if i % 1000 == 0: print('%0.2f percent images are added to queue.' % (i * 100 / nb_images))

    pool.close()
    pool.join()

    new_nb_images = len([name for name in os.listdir(path)])
    print()
    print("New size of dataset : %d. Number of images deleted = %d" % (new_nb_images, nb_images - new_nb_images))

if __name__ == "__main__":
    '''
    Run this script to check for corrupt images in an image dataset whose path is provided
    '''
    ms_coco_path = r""

    '''
    Note:
        If due to some reason the corrupted file cannot be deleted via os.remove(),
        it will print the path of the file. Please delete the file manually in such a case.
    '''
    check_dataset(ms_coco_path)


#!/usr/bin/env python2
from __future__ import print_function
import sys
sys.path.append('../lib')
import os
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.path import Path
from matplotlib.patches import PathPatch
import string

import protocols
import model as m
from parameters import simvc, get_qc
from parameters import simvc_fix, simvc_fix_typical_values
from parameters import simvc_typical_values
from releakcorrect import I_releak, score_leak, protocol_leak_check

from scipy.optimize import fmin
# Set seed
np.random.seed(101)

savedir = './figs'
if not os.path.isdir(savedir):
    os.makedirs(savedir)

#refcell = 'D19'

def get_fix_param(var, val):
    """
        var: variable name.
        val: variable value to fix.
    """
    out = {}
    for i, j in zip(var, val):
        out[i] = j
    return out

def rmsd_compute(t1, t2):
    # Normalised RMSD value between trace 1 ``t1`` and trace 2 ``t2``
    #
    # Note, usually normalise to data, so
    # - ``t2`` data (or anything as reference)
    # - ``t1`` simulation (or anything for comparison)
    return np.sqrt(np.mean((t1 - t2) ** 2)) / np.sqrt(np.mean(t2 ** 2))

#
# Protocol info
#
protocol_funcs = {
    'staircaseramp': protocols.leak_staircase,
    'pharma': protocols.pharma,  # during drug application
    'apab': 'protocol-apab.csv',
    'apabv3': 'protocol-apabv3.csv',
    'ap05hz': 'protocol-ap05hz.csv',
    'ap1hz': 'protocol-ap1hz.csv',
    'ap2hz': 'protocol-ap2hz.csv',
    'sactiv': protocols.sactiv,
    'sinactiv': protocols.sinactiv,
}
protocol_dir = '../protocol-time-series'
protocol_list = [
    'staircaseramp',
    # 'sactiv',
    # 'sinactiv',
    'pharma',
    'apab',
    'apabv3',
    # 'ap05hz',
    'ap1hz',
    'ap2hz',
]
validation_idx = [
    None,
    # 1,
    # 2,
    3,
    4,
    5,
    # 6,
    7,
    8,
]

# IV protocol special treatment
protocol_iv = [
    'sactiv',
    'sinactiv',
]
protocol_iv_times = {
    'sactiv': protocols.sactiv_times,
    'sinactiv': protocols.sinactiv_times,
}
protocol_iv_convert = {
    'sactiv': protocols.sactiv_convert,
    'sinactiv': protocols.sinactiv_convert,
}
protocol_iv_args = {
    'sactiv': protocols.sactiv_iv_arg,
    'sinactiv': protocols.sinactiv_iv_arg,
}
protocol_iv_v = {
    'sactiv': protocols.sactiv_v,
    'sinactiv': protocols.sinactiv_v,
}

data_dir_staircase = '../data'
data_dir = '../data-autoLC'
file_dir = '../../hERGRapidCharacterisation/room-temperature-only/out'
file_dir2 = './out'
file_list = [
        'herg25oc1',
        ]
temperatures = np.array([25.0])
temperatures += 273.15  # in K
fit_seed = '542811797'
fit_seed2 = '717354021'

file_name = file_list[0]
temperature = temperatures[0]

# Load RMSD matrix
rmsd_matrix_file = '../../hERGRapidCharacterisation/room-temperature-only/figs/rmsd-hist-%s-autoLC-releak/rmsd-matrix.txt' \
                   % file_name
rmsd_cells_file = '../../hERGRapidCharacterisation/room-temperature-only/figs/rmsd-hist-%s-autoLC-releak/rmsd-matrix-cells.txt' \
                  % file_name

rmsd_matrix = np.loadtxt(rmsd_matrix_file)

with open(rmsd_matrix_file, 'r') as f:
    rmsd_prt = f.readline().strip('\n').strip('#').split()

rmsd_cells = []
with open(rmsd_cells_file, 'r') as f:
    for l in f:
        if not l.startswith('#'):
            rmsd_cells.append(l.strip('\n').split('-')[1])

rmsd_matrix_file2 = './out/rmsd-hist-%s-fixkinetics-simvclinleak-scheme3/rmsd-matrix.txt' \
                   % file_name
rmsd_cells_file2 = './out/rmsd-hist-%s-fixkinetics-simvclinleak-scheme3/rmsd-matrix-cells.txt' \
                  % file_name

rmsd_matrix2 = np.loadtxt(rmsd_matrix_file2)

with open(rmsd_matrix_file2, 'r') as f:
    rmsd_prt2 = f.readline().strip('\n').strip('#').split()

rmsd_cells2 = []
with open(rmsd_cells_file2, 'r') as f:
    for l in f:
        if not l.startswith('#'):
            rmsd_cells2.append(l.strip('\n').split('-')[1])

rankedlabels = [r'$*$',
                u'\u2021',
                r'#',
                u'\u2666']

#
# Do a very very tailored version........ :(
#
fig = plt.figure(figsize=(16, 15))
bigxgap = 12
n_xgrid = 84
bigygap = 5
n_ygrid = 31
grid = plt.GridSpec(2 * n_ygrid + 1 * bigygap, 3 * n_xgrid + 2 * bigxgap,
                    hspace=0.0, wspace=0.0)
axes = np.empty([10, int(len(protocol_list) / 2)], dtype=object)
# long list here:
for i in range(int(len(protocol_list) / 2)):
    i_grid = i * (n_xgrid + bigxgap)
    f_grid = (i + 1) * n_xgrid + i * bigxgap

    # First 'row'
    axes[0, i] = fig.add_subplot(grid[0:3, i_grid:f_grid])
    axes[0, i].set_xticklabels([])
    axes[1, i] = fig.add_subplot(grid[3:9, i_grid:f_grid])
    axes[1, i].set_xticklabels([])
    axes[2, i] = fig.add_subplot(grid[9:15, i_grid:f_grid])
    axes[2, i].set_xticklabels([])
    axes[3, i] = fig.add_subplot(grid[15:21, i_grid:f_grid])
    # Histogram
    axes[4, i] = fig.add_subplot(grid[24:31, i_grid:f_grid])

    # Second 'row'
    n_shift = n_ygrid + bigygap
    axes[5, i] = fig.add_subplot(grid[n_shift+0:n_shift+3, i_grid:f_grid])
    axes[5, i].set_xticklabels([])
    axes[6, i] = fig.add_subplot(grid[n_shift+3:n_shift+9, i_grid:f_grid])
    axes[6, i].set_xticklabels([])
    axes[7, i] = fig.add_subplot(grid[n_shift+9:n_shift+15, i_grid:f_grid])
    axes[7, i].set_xticklabels([])
    axes[8, i] = fig.add_subplot(grid[n_shift+15:n_shift+21, i_grid:f_grid])
    # Histogram
    axes[9, i] = fig.add_subplot(grid[n_shift+24:n_shift+31, i_grid:f_grid])

    # Set x-labels
    axes[3, i].set_xlabel('Time (s)', fontsize=14)
    axes[4, i].set_xlabel('RRMSE', fontsize=14)
    axes[8, i].set_xlabel('Time (s)', fontsize=14)
    axes[9, i].set_xlabel('RRMSE', fontsize=14)

# Set labels
axes[0, 0].set_ylabel('Voltage\n(mV)', fontsize=14)
axes[1, 0].set_ylabel(u'Best\n(*)', fontsize=14, color='#d95f02')
axes[2, 0].set_ylabel(u'Median\n(\u2021)', fontsize=14, color='#d95f02')
axes[3, 0].set_ylabel(u'90%ile\n(#)', fontsize=14, color='#d95f02')
axes[4, 0].set_ylabel('Frequency\n(N=%s)' % len(rmsd_cells), fontsize=14)
axes[5, 0].set_ylabel('Voltage\n(mV)', fontsize=14)
axes[6, 0].set_ylabel(u'Best\n(*)', fontsize=14, color='#d95f02')
axes[7, 0].set_ylabel(u'Median\n(\u2021)', fontsize=14, color='#d95f02')
axes[8, 0].set_ylabel(u'90%ile\n(#)', fontsize=14, color='#d95f02')
axes[9, 0].set_ylabel('Frequency\n(N=%s)' % len(rmsd_cells), fontsize=14)

axes[2, 0].text(-0.3, 0.5, 'Current (pA)', rotation=90, fontsize=18,
        transform=axes[2, 0].transAxes, ha='center', va='center')
axes[7, 0].text(-0.275, 0.5, 'Current (pA)', rotation=90, fontsize=18,
        transform=axes[7, 0].transAxes, ha='center', va='center')

#
# Model
#
prt2model = {}
prt2fixkineticsmodel = {}
for prt in protocol_list:

    protocol_def = protocol_funcs[prt]
    if type(protocol_def) is str:
        protocol_def = '%s/%s' % (protocol_dir, protocol_def)

    prt2model[prt] = m.Model('../mmt-model-files/ideal-ikr.mmt',
                        protocol_def=protocol_def,
                        temperature=temperature,  # K
                        transform=None,
                        useFilterCap=False)  # ignore capacitive spike

    prt2fixkineticsmodel[prt] = m.Model(
                '../mmt-model-files/simplified-voltage-clamp-ikr-linleak.mmt',
                protocol_def=protocol_def,
                temperature=temperature,  # K
                transform=None,
                useFilterCap=False)  # ignore capacitive spike

#
# Plot
#
mid = []
upper = []
lower = []
mid2 = []
upper2 = []
lower2 = []
for i_prt, prt in enumerate(protocol_list):

    # Calculate axis index
    ai, aj = 5 * int(i_prt / 3), i_prt % 3

    # Title
    if prt == 'staircaseramp':
        axes[ai, aj].set_title('Calibration', fontsize=16)
    else:
        axes[ai, aj].set_title('Validation %s' % validation_idx[i_prt],
                fontsize=16)

    # Add label!
    axes[ai, aj].text(-0.1, 1.4, string.ascii_uppercase[i_prt],
                      transform=axes[ai, aj].transAxes, size=20,
                      weight='bold')

    # Time point
    times = np.loadtxt('%s/%s-%s-times.csv' % (data_dir, file_name,
        prt), delimiter=',', skiprows=1) * 1e3  # s -> ms

    # Protocol
    model = prt2model[prt]
    modelfixkinetics = prt2fixkineticsmodel[prt]
    # Set which parameters to be inferred
    modelfixkinetics.set_parameters([
        'ikr.g',
        #'voltageclamp.rseries',
        'voltageclamp.voffset_eff',
        'voltageclamp.gLeak'])
    if prt not in protocol_iv:
        times_sim = np.copy(times)
        voltage = model.voltage(times_sim)
    else:
        times_sim = protocol_iv_times[prt](times[1] - times[0])
        voltage = model.voltage(times_sim)
        voltage, t = protocol_iv_convert[prt](voltage, times_sim)
        assert(np.mean(np.abs(t - times)) < 1e-8)
    axes[ai, aj].set_ylim((np.min(voltage) - 10, np.max(voltage) + 15))

    # Plot protocol
    if prt not in protocol_iv:
        axes[ai, aj].plot(times * 1e-3, voltage, c='#696969')
    else:
        # protocol
        for i in range(voltage.shape[1]):
            axes[ai, aj].plot(times * 1e-3, voltage[:, i], c='#696969')

    # Calculate ranking
    rmsd = rmsd_matrix[:, rmsd_prt.index(prt)]
    best_cell = np.argmin(rmsd)
    median_cell = np.argsort(rmsd)[len(rmsd)//2]
    p90_cell = np.argsort(rmsd)[int(len(rmsd)*0.9)]
    rankedcells = [rmsd_cells[best_cell],
                   rmsd_cells[median_cell],
                   rmsd_cells[p90_cell]]
    rankedvalues = [rmsd[best_cell],
                    rmsd[median_cell],
                    rmsd[p90_cell]]
                    #rmsd[rmsd_cells.index(refcell)]]

    rmsd2 = rmsd_matrix2[:, rmsd_prt2.index(prt)]
    #best_cell2 = np.argmin(rmsd2)
    #median_cell2 = np.argsort(rmsd2)[len(rmsd2)//2]
    #p90_cell2 = np.argsort(rmsd2)[int(len(rmsd2)*0.9)]
    #NOTE Compare with the 'red' cells; not its own ranking!!
    best_cell2 = rmsd_cells2.index(rmsd_cells[best_cell])
    median_cell2 = rmsd_cells2.index(rmsd_cells[median_cell])
    p90_cell2 = rmsd_cells2.index(rmsd_cells[p90_cell])
    rankedcells2 = [rmsd_cells2[best_cell2],
                   rmsd_cells2[median_cell2],
                   rmsd_cells2[p90_cell2]]
    rankedvalues2 = [rmsd2[best_cell2],
                    rmsd2[median_cell2],
                    rmsd2[p90_cell2]]
                    #rmsd2[rmsd_cells2.index(refcell)]]

    # Parameters
    fn = '%s/%s-scheme3-simvclinleak/%s-cells-%s.txt' % \
         (file_dir2, file_name, file_name, fit_seed2)
    scheme3_cell_list = []
    with open(fn, 'r') as f:
        for l in f:
            if not l.startswith('#'):
                scheme3_cell_list.append(l.split()[0])

    param_file = '%s/%s-scheme3-simvclinleak/%s-solution_i-%s.txt' % \
            (file_dir2, file_name, file_name, fit_seed2)
    obtained_parameters_all = np.loadtxt(param_file)

    ikr_param = [
            'ikr.p1', 'ikr.p2', 'ikr.p3', 'ikr.p4',
            'ikr.p5', 'ikr.p6', 'ikr.p7', 'ikr.p8',
            ]
    p_ikr = np.loadtxt('%s/%s-scheme3-simvclinleak/%s-solution-%s.txt' % \
            (file_dir2, file_name, file_name, fit_seed2))

    for i_cell, cell in enumerate(rankedcells):
        # Data
        if prt == 'staircaseramp':
            data = np.loadtxt('%s/%s-%s-%s.csv' % (data_dir_staircase,
                    file_name, prt, cell), delimiter=',', skiprows=1)
        elif prt not in protocol_iv:
            data = np.loadtxt('%s/%s-%s-%s.csv' % (data_dir, file_name,
                    prt, cell), delimiter=',', skiprows=1)
            # Re-leak correct the leak corrected data...
            g_releak = fmin(score_leak, [0.0], args=(data, voltage, times,
                                protocol_leak_check[prt]), disp=False)
            data = I_releak(g_releak[0], data, voltage)
        else:
            data = np.loadtxt('%s/%s-%s-%s.csv' % (data_dir, file_name,
                    prt, cell), delimiter=',', skiprows=1)
            # Re-leak correct the leak corrected data...
            for i in range(data.shape[1]):
                g_releak = fmin(score_leak, [0.0], args=(data[:, i],
                                    voltage[:, i], times,
                                    protocol_leak_check[prt]), disp=False)
                data[:, i] = I_releak(g_releak[0], data[:, i], voltage[:, i])
        assert(len(data) == len(times))

        # Fitted parameters
        param_file = '%s/%s/%s-staircaseramp-%s-solution-%s.txt' % \
                (file_dir, file_name, file_name, cell, fit_seed)
        obtained_parameters = np.loadtxt(param_file) * 1e-3  # V, s -> mV, ms

        # For fix kinetics model
        rseal, cm, rseries = get_qc('../qc', file_name, cell)
        #print('Est. Rseal, Cm, Rseries:', rseal, cm, rseries, '(GOhm, pF, GOhm)')
        alpha = 0.8  # rseries %compensation
        simvc_fix_values = [cm, rseries * alpha, rseries]
        extra_fix = ['voltageclamp.rseries']
        updateELeakCorrection = False
        if updateELeakCorrection:
            leakbeforeparam = np.loadtxt('../qc/' + file_name + '-staircaseramp-leak_before.txt')
            leakafterparam = np.loadtxt('../qc/' + file_name + '-staircaseramp-leak_after.txt')
            cell_id_file = '../qc/%s-staircaseramp-cell_id.txt' % file_name
            cell_ids = []
            with open(cell_id_file, 'r') as f:
                for l in f:
                    if not l.startswith('#'):
                        cell_ids.append(l.split()[0])
            cell_idx = cell_ids.index(cell)
            ga, Ea = leakbeforeparam[cell_idx]
            gb, Eb = leakafterparam[cell_idx]
            ELeakCorrection = - (ga * Ea - gb * Eb) / (gb - ga)
            #print('E_Leak correction: ', ELeakCorrection, ' (mV)')
            if np.abs(ELeakCorrection) > 200: print('==' * 30, ga, Ea, gb, Eb)
            extra_fix += ['voltageclamp.ELeak']
            simvc_fix_values += [ELeakCorrection]
        fix_p = get_fix_param(ikr_param + simvc_fix + extra_fix,
                np.append(p_ikr, simvc_fix_values))
        modelfixkinetics.set_fix_parameters(fix_p)

        scheme3_cell_idx = scheme3_cell_list.index(cell)
        obtained_parameters2 = obtained_parameters_all[scheme3_cell_idx]

        # Simulation
        simulation = model.simulate(obtained_parameters, times_sim)
        simulationfixkinetics = modelfixkinetics.simulate(obtained_parameters2, times_sim)
        if prt != 'staircaseramp' and prt not in protocol_iv:
            # Re-leak correct the leak corrected simulationfixkinetics... TODO?
            g_releak_simulationfixkinetics = fmin(score_leak, [0.1], args=(simulationfixkinetics, voltage, times,
                                protocol_leak_check[prt]), disp=False)
            simulationfixkinetics = I_releak(g_releak_simulationfixkinetics[0], simulationfixkinetics, voltage)
        if prt in protocol_iv:
            simulation, t = protocol_iv_convert[prt](simulation, times_sim)
            assert(np.mean(np.abs(t - times)) < 1e-6)
            simulationfixkinetics, t = protocol_iv_convert[prt](
                    simulationfixkinetics, times_sim)
            assert(np.mean(np.abs(t - times)) < 1e-6)
            # Re-leak correct the leak corrected simulationfixkinetics... TODO?
            for i in range(simulationfixkinetics.shape[1]):
                g_releak_simulationfixkinetics = fmin(score_leak, [0.1], args=(simulationfixkinetics[:, i],
                                    voltage[:, i], times,
                                    protocol_leak_check[prt]), disp=False)
                simulationfixkinetics[:, i] = I_releak(g_releak_simulationfixkinetics[0], simulationfixkinetics[:, i], voltage[:, i])

        # Work out ylim
        maximum = np.percentile(simulation, 100)
        minimum = np.percentile(simulation, 0.0)
        amplitude = maximum - minimum
        if prt in ['apabv3', 'ap05hz']:
            maximum += 0.6 * amplitude
            minimum -= 0.6 * amplitude
        elif prt in ['apab', 'ap1hz']:
            maximum += 0.3 * amplitude
            minimum -= 0.3 * amplitude
        else:
            maximum += 0.15 * amplitude
            minimum -= 0.15 * amplitude
        axes[ai + i_cell + 1, aj].set_ylim([minimum, maximum])

        # Plot
        if prt not in protocol_iv:
            # recording
            axes[ai + i_cell + 1, aj].plot(times * 1e-3, data, lw=1, alpha=0.5,
                    c='#9ecae1', label='data')
            # simulation
            if prt == 'staircaseramp':
                axes[ai + i_cell + 1, aj].plot(times * 1e-3, simulation, lw=2,
                        c='#d95f02', label='model fit to data')
                axes[ai + i_cell + 1, aj].plot(times * 1e-3, simulationfixkinetics, lw=2,
                        c='#1b9e77', label='model (fix kinetics) fit to data')
            else:
                axes[ai + i_cell + 1, aj].plot(times * 1e-3, simulation, lw=2,
                        c='#d95f02', label='model prediction')
                axes[ai + i_cell + 1, aj].plot(times * 1e-3, simulationfixkinetics, lw=2,
                        c='#1b9e77', label='model (fix kinetics) prediction')
        else:
            iv_v = protocol_iv_v[prt]()  # mV
            # recording
            iv_i = protocols.get_corrected_iv(data, times,
                                              *protocol_iv_args[prt]())
            axes[ai + i_cell + 1, aj].plot(iv_v, iv_i / np.max(iv_i), lw=2,
                    alpha=0.25, c='#9ecae1', label='data')
            # simulation
            iv_i = protocols.get_corrected_iv(simulation, times,
                                              *protocol_iv_args[prt]())
            axes[ai + i_cell + 1, aj].plot(iv_v, iv_i / np.max(iv_i), lw=2,
                    alpha=1, c='#d95f02', label='model prediction')
            # simulationfixkinetics
            iv_i_fixkinetics = protocols.get_corrected_iv(simulationfixkinetics, times,
                                              *protocol_iv_args[prt]())
            axes[ai + i_cell + 1, aj].plot(iv_v, iv_i_fixkinetics / np.max(iv_i_fixkinetics), lw=2,
                    alpha=1, c='#1b9e77', label='model (fix kinetics) prediction')
            if prt == 'sactiv':
                axes[ai + i_cell + 1, aj].set_ylim([-0.05, 1.05])
            elif prt == 'sinactiv':
                axes[ai + i_cell + 1, aj].set_ylim([-5, 1.05])

        if False:
            print(prt, i_cell, cell)
            print('red', rmsd_compute(simulation, data))
            print('green', rmsd_compute(simulationfixkinetics, data))

    # Plot rmsd histogram
    rmse_min = min(np.min(rmsd), np.min(rmsd2))
    rmse_max = max(np.max(rmsd), np.max(rmsd2))
    rmse_range = rmse_max - rmse_min
    bins = np.linspace(rmse_min - 0.1 * rmse_range,
            rmse_max + 0.1 * rmse_range, 20)
    n, b, _ = axes[ai + 4, aj].hist(rmsd, bins=bins, color='#d95f02', alpha=0.25)
    n2, b2, _ = axes[ai + 4, aj].hist(rmsd2, bins=bins, color='#2ca02c', alpha=0.25)

    mid.append(np.percentile(rmsd, 50))
    upper.append(np.percentile(rmsd, 90))
    lower.append(np.percentile(rmsd, 10))
    mid2.append(np.percentile(rmsd2, 50))
    upper2.append(np.percentile(rmsd2, 90))
    lower2.append(np.percentile(rmsd2, 10))

    # Add labels
    rankedidx = []
    for i, v in enumerate(rankedvalues):
        idx = np.where(b <= v)[0][-1]
        if idx in rankedidx:
            print('Ref. marker might clash with other markers...')
            shift = 4
        else:
            shift = 0
        axes[ai + 4, aj].text((b[idx] + b[idx + 1]) / 2., n[idx] + 3 + shift,
                rankedlabels[i], fontsize=16, color='#d95f02',
                ha='center', va='center')
        if n[idx] > 0.8 * max(np.max(n), np.max(n2)):
            axes[ai + 4, aj].set_ylim([0, max(np.max(n2), np.max(n)) + 8 + shift])
        rankedidx.append(idx)

    rankedidx2 = []
    for i, v in enumerate(rankedvalues2):
        idx = np.where(b2 <= v)[0][-1]
        if idx in rankedidx2:
            print('Ref. marker might clash with other markers...')
            shift = 4
        elif idx in rankedidx:
            diff = np.abs(n[idx] - n2[idx])
            if diff < max(np.max(n2), np.max(n)) * 0.1:
                shift = max(np.max(n2), np.max(n)) * 0.2
            else:
                shift = 0
        else:
            shift = 0
        axes[ai + 4, aj].text((b2[idx] + b2[idx + 1]) / 2., n2[idx] + 3 + shift,
                rankedlabels[i], fontsize=16, color='#2ca02c',
                ha='center', va='center')
        if n2[idx] > 0.8 * max(np.max(n2), np.max(n)):
            axes[ai + 4, aj].set_ylim([0, max(np.max(n2), np.max(n)) + 8 + shift])
        rankedidx2.append(idx)

#
# Final adjustment and save
#
#axes[1, 0].legend()
#axes[1, 1].legend()
import matplotlib.patches as mpatches
data_patch = mpatches.Patch(color='#9ecae1', label='Data')
h1_patch = mpatches.Patch(color='#d95f02', label='Hypothesis 1: independent kinetics models')
h2_patch = mpatches.Patch(color='#1b9e77', label='Hypothesis 2: identical kinetics models')
axes[0, 0].legend(handles=[data_patch, h1_patch, h2_patch], loc='upper left', bbox_to_anchor=(-.025, 2.75), fontsize=14, ncol=3, columnspacing=5.5)
#grid.tight_layout(fig, pad=0.6, rect=(0.02, 0.0, 1, 0.99))  # not working...
#grid.update(wspace=0.2, hspace=0.0)
plt.savefig('%s/rmsd-hist-fix-kinetics-simvclinleak-scheme3-part1.pdf' % (savedir), bbox_inch='tight',
        pad_inches=0, format='pdf')
plt.savefig('%s/rmsd-hist-fix-kinetics-simvclinleak-scheme3-part1.png' % (savedir), bbox_inch='tight',
        pad_inches=0, dpi=300)

print('Done')

#
# Table
#
tex = ''
tex += '\\begin{tabularx}{\\textwidth}{@{}l' \
    + 'XXc' * (len(protocol_list) - 1) + 'XX@{}}\n'
tex += '\\toprule\n'
tex += '   '
for i_prt, prt in enumerate(protocol_list):
    # span 2 columns
    if prt == 'staircaseramp':
        tex += ' & \multicolumn{2}{c}{Cal.}'
    else:
        tex += ' & \multicolumn{2}{c}{Val.~%s}' % validation_idx[i_prt]

    if i_prt < len(protocol_list) - 1:
        tex += ' & \phantom{}'
tex += ' \\\\\n'

ii = 1
for i_prt, prt in enumerate(protocol_list):
    ii += 1
    tex += '\\cmidrule{%s-%s}' % (ii, ii + 1)
    ii += 2
tex += '\n'

tex += '   '
for i_prt, prt in enumerate(protocol_list):
    tex += ' & H1 & H2'
    if i_prt < len(protocol_list) - 1:
        tex += ' &'
tex += ' \\\\\n'

tex += '\\midrule\n'

tex += 'Median'
for i_prt, prt in enumerate(protocol_list):
    tex += ' & %.2f & %.2f' % (mid[i_prt], mid2[i_prt])
    if i_prt < len(protocol_list) - 1:
        tex += ' &'
tex += ' \\\\\n'

tex += '10\\textsuperscript{th} \\%ile'
for i_prt, prt in enumerate(protocol_list):
    tex += ' & %.2f & %.2f' % (lower[i_prt], lower2[i_prt])
    if i_prt < len(protocol_list) - 1:
        tex += ' &'
tex += ' \\\\\n'

tex += '90\\textsuperscript{th} \\%ile'
for i_prt, prt in enumerate(protocol_list):
    tex += ' & %.2f & %.2f' % (upper[i_prt], upper2[i_prt])
    if i_prt < len(protocol_list) - 1:
        tex += ' &'
tex += ' \\\\\n'
tex += '\\bottomrule\n'
tex += '\\end{tabularx}'

print(tex)


'''Trains a simple convnet on the MNIST dataset.

Gets to 99.25% test accuracy after 12 epochs
(there is still a lot of margin for parameter tuning).
16 seconds per epoch on a GRID K520 GPU.
'''
from __future__ import print_function
import argparse
import numpy as np
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import tensorflow as tf
import keras
from keras import backend as K
from aiter import *
from keras.backend.tensorflow_backend import set_session
from importlib import import_module
import matplotlib.pyplot as plt
plt.switch_backend('agg')

batch_size = 100000

parser=argparse.ArgumentParser()
parser.add_argument("--end",type=float,default=1.,help='end ratio')
parser.add_argument("--epoch",type=int,default=10,help='epoch')
parser.add_argument("--save",type=str,default="ten100grucnn",help='rch')
parser.add_argument("--rc",type=str,default='rc',help='rnn or cnn')
parser.add_argument("--gpu",type=int,default=0,help='gpu')
parser.add_argument("--pt",type=int,default=100,help='pt range pt~pt*1.1')
parser.add_argument("--isz",type=int,default=0,help='0 or z or not')
parser.add_argument("--channel",type=int,default=30,help='sequence channel')
parser.add_argument("--order",type=int,default=1,help='pt ordering')
parser.add_argument("--eta",type=float,default=0.,help='end ratio')
parser.add_argument("--etabin",type=float,default=2.4,help='end ratio')

args=parser.parse_args()

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]=str(args.gpu)
config =tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction=0.3

set_session(tf.Session(config=config))

# input image dimensions
img_rows, img_cols = 33, 33

# the data, shuffled and split between train and test sets
#(x_train, y_train), (x_test, y_test) = mnist.load_data()

#if K.image_data_format() == 'channels_first':
#    x_train = x_train.reshape(x_train.shape[0], 1, img_rows, img_cols)
#    x_test = x_test.reshape(x_test.shape[0], 1, img_rows, img_cols)
#    input_shape = (1, img_rows, img_cols)
#else:
#x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)
#x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1)
input_shape1= (9,33,33)
input_shape2= (20,10)

#model.compile(loss=keras.losses.categorical_crossentropy,
#              optimizer=keras.optimizers.Adadelta(),
#              metrics=['accuracy'])
#model=keras.models.load_model('save/fullydijetsame_10')
savename="save/"+str(args.save)
history=open(savename+"/history").readlines()
try:
  hist=eval(history[0])
  a=hist['val1_auc']
except:
  hist=eval(history[1])
vzjdata="Data/zj_pt_{0}_{1}.root".format(args.pt,int(args.pt*1.1))
vjjdata="Data/jj_pt_{0}_{1}.root".format(args.pt,int(args.pt*1.1))
vzqdata="Data/zq_pt_{0}_{1}.root".format(args.pt,int(args.pt*1.1))
vzgdata="Data/zg_pt_{0}_{1}.root".format(args.pt,int(args.pt*1.1))
vqqdata="Data/qq_pt_{0}_{1}.root".format(args.pt,int(args.pt*1.1))
vggdata="Data/gg_pt_{0}_{1}.root".format(args.pt,int(args.pt*1.1))
from sklearn.metrics import roc_auc_score, auc, roc_curve
if(args.isz==0):iii=1
if(args.isz==1):iii=2
if(args.isz==-1):iii=3
epoch=hist['val{}_auc'.format(iii)].index(max(hist['val{}_auc'.format(iii)]))+1
model=keras.models.load_model(savename+"/check_"+str(epoch))
rc=""
for sha in model._feed_inputs:
  if(len(sha._keras_shape)==4):
    rc+="c"
  if(len(sha._keras_shape)==3):
    rc+="r"
onehot=0
test2=wkiter([vzqdata,vzgdata],batch_size=batch_size,begin=args.end*0.2,end=args.end*0.6,rc=rc,onehot=onehot,channel=args.channel,order=args.order,eta=args.eta,etabin=args.etabin,pt=args.pt)
test3=wkiter([vqqdata,vggdata],batch_size=batch_size,begin=args.end*0.2,end=args.end*0.6,rc=rc,onehot=onehot,channel=args.channel,order=args.order,eta=args.eta,etabin=args.etabin,pt=args.pt)
entries=test2.totalnum()
print ("test   ",entries)
#epoch=eval(open(savename+"/history").readline())+1
test2.reset()
test3.reset()
for ii in range(2,4):
  if(ii==2):
    if(args.isz==-1):
      continue
    gen=test2
  if(ii==3):
    if(args.isz==1):
      continue
    gen=test3
  x=[]
  y=[]
  g=[]
  q=[]
  outname=savename+"/v{}t{}".format(iii,ii)
  for j in range(1):
    p=model.predict(gen.gjetset,verbose=0)[:,0]
    pt=gen.gptset
    eta=gen.getaset
    f=open(outname+"ptg.dat",'w')
    f.write(str(p.tolist())+"\n")
    f.write(str(pt.tolist())+"\n")
    f.write(str(eta.tolist()))
    f.close()
    p=model.predict(gen.qjetset,verbose=0)[:,0]
    pt=gen.qptset
    eta=gen.qetaset
    f=open(outname+"ptq.dat",'w')
    f.write(str(p.tolist())+"\n")
    f.write(str(pt.tolist())+"\n")
    f.write(str(eta.tolist()))
    f.close()
    print("###number",args.pt,args.eta)


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

from torchsac.agent import Agent
from torchsac import utils
from torchsac.agent.critic import DoubleQCritic
from torchsac.agent.actor import DiagGaussianActor

class SACAgent(Agent):
    """SAC algorithm."""

    def __init__(self, cfg, device):
        super().__init__()

        obs_dim = cfg.obs_dim
        action_dim = cfg.action_dim

        self.action_range = cfg.action_range
        self.device = device
        self.discount = cfg.discount
        self.critic_tau = cfg.critic_tau
        self.actor_update_frequency = cfg.actor_update_frequency
        self.critic_target_update_frequency = cfg.critic_target_update_frequency
        self.batch_size = cfg.batch_size
        self.learnable_temperature = cfg.learnable_temperature

        self.critic = DoubleQCritic(
            obs_dim, action_dim, cfg.critic.hidden_dim, cfg.critic.hidden_depth
        ).to(self.device)
        self.critic_target = DoubleQCritic(
            obs_dim, action_dim, cfg.critic.hidden_dim, cfg.critic.hidden_depth
        ).to(self.device)
        self.critic_target.load_state_dict(self.critic.state_dict())

        self.actor = DiagGaussianActor(
            obs_dim,
            action_dim,
            cfg.actor.hidden_dim,
            cfg.actor.hidden_depth,
            cfg.actor.log_std_bounds,
        ).to(self.device)

        self.log_alpha = torch.tensor(np.log(cfg.init_temperature)).to(
            self.device
        )
        self.log_alpha.requires_grad = True
        # set target entropy to -|A|
        self.target_entropy = -action_dim

        # optimizers
        self.actor_optimizer = torch.optim.Adam(
            self.actor.parameters(), lr=cfg.actor_lr, betas=cfg.actor_betas
        )

        self.critic_optimizer = torch.optim.Adam(
            self.critic.parameters(), lr=cfg.critic_lr, betas=cfg.critic_betas
        )

        self.log_alpha_optimizer = torch.optim.Adam(
            [self.log_alpha], lr=cfg.alpha_lr, betas=cfg.alpha_betas
        )

        self.train()
        self.critic_target.train()

    def train(self, training=True):
        self.training = training
        self.actor.train(training)
        self.critic.train(training)

    @property
    def alpha(self):
        return self.log_alpha.exp()

    def act(self, obs, sample=False):
        obs = torch.FloatTensor(obs).to(self.device)
        obs = obs.unsqueeze(0)
        dist = self.actor(obs)
        action = dist.sample() if sample else dist.mean
        action = action.clamp(*self.action_range)
        assert action.ndim == 2 and action.shape[0] == 1
        return utils.to_np(action[0])

    def update_critic(
        self, obs, action, reward, next_obs, not_done, logger, step
    ):
        dist = self.actor(next_obs)
        next_action = dist.rsample()
        log_prob = dist.log_prob(next_action).sum(-1, keepdim=True)
        target_Q1, target_Q2 = self.critic_target(next_obs, next_action)
        target_V = (
            torch.min(target_Q1, target_Q2) - self.alpha.detach() * log_prob
        )
        target_Q = reward + (not_done * self.discount * target_V)
        target_Q = target_Q.detach()

        # get current Q estimates
        current_Q1, current_Q2 = self.critic(obs, action)
        critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(
            current_Q2, target_Q
        )
        logger.log("train_critic/loss", critic_loss, step)

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

        self.critic.log(logger, step)

    def update_actor_and_alpha(self, obs, logger, step):
        dist = self.actor(obs)
        action = dist.rsample()
        log_prob = dist.log_prob(action).sum(-1, keepdim=True)
        actor_Q1, actor_Q2 = self.critic(obs, action)

        actor_Q = torch.min(actor_Q1, actor_Q2)
        actor_loss = (self.alpha.detach() * log_prob - actor_Q).mean()

        logger.log("train_actor/loss", actor_loss, step)
        logger.log("train_actor/target_entropy", self.target_entropy, step)
        logger.log("train_actor/entropy", -log_prob.mean(), step)

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

        self.actor.log(logger, step)

        if self.learnable_temperature:
            self.log_alpha_optimizer.zero_grad()
            alpha_loss = (
                self.alpha * (-log_prob - self.target_entropy).detach()
            ).mean()
            logger.log("train_alpha/loss", alpha_loss, step)
            logger.log("train_alpha/value", self.alpha, step)
            alpha_loss.backward()
            self.log_alpha_optimizer.step()

    def update(self, replay_buffer, logger, step):
        (
            obs,
            action,
            reward,
            next_obs,
            not_done,
            not_done_no_max,
        ) = replay_buffer.sample(self.batch_size)

        logger.log("train/batch_reward", reward.mean(), step)

        self.update_critic(
            obs, action, reward, next_obs, not_done_no_max, logger, step
        )

        if step % self.actor_update_frequency == 0:
            self.update_actor_and_alpha(obs, logger, step)

        if step % self.critic_target_update_frequency == 0:
            utils.soft_update_params(
                self.critic, self.critic_target, self.critic_tau
            )

    def save(self, model_dir, step):
        torch.save(
            self.actor.state_dict(), "%s/actor_%s.pt" % (model_dir, step)
        )

    def load(self, model_dir, step):
        self.actor.load_state_dict(
            torch.load("%s/actor_%s.pt" % (model_dir, step))
        )
        print(f"Successfully loaded weights from {model_dir}.")


'''
Hardware abstraction layer for communicating with the Arduino.
'''

import time
import serial
import numpy as np

# Hardware constants
_speed_max = 2000  # True max: 2000
_speed_min = 700
_speed_range = _speed_max - _speed_min
_baud_rate = 9600

# Module globals
_serial_port = None
_setpoint_left = None
_setpoint_right = None

def init(port='/dev/ttyACM1'):
    '''
    Initialize the hardware.

    Parameters
    ----------
    port : str
        Port to setup serial for.
    '''
    global _serial_port
    _serial_port = serial.Serial(port, baudrate=_baud_rate)

def calibrate():
    '''
    Calibrate the hardware.
    '''
    set_speeds(_speed_max, _speed_max)
    time.sleep(5)
    set_speeds(_speed_min, _speed_min)
    time.sleep(3)
    set_speeds(1000, 1000)
    time.sleep(1)
    set_speeds(_speed_min, _speed_min)

def cleanup():
    '''
    Shutdown the hardware.
    '''
    global _serial_port
    assert _serial_port
    _serial_port.close()

def set_speeds(speed_left, speed_right):
    '''
    Set left and right motor speeds.

    Parameters
    ----------
    speed_left : int in [700, 2000]
        Forward speed of left motor, from stopped to maximum.
    speed_right : int in [700, 2000]
        Forward speed of right motor, from stopped to maximum.
    '''
    global _serial_port
    global _setpoint_left
    global _setpoint_right

    assert _serial_port
    assert _serial_port.is_open
    assert _speed_min <= speed_left <= _speed_max
    assert _speed_min <= speed_right <= _speed_max

    if (speed_left == _setpoint_left and speed_right == _setpoint_right):
        return

    _setpoint_left = speed_left
    _setpoint_right = speed_right

    serial_string = '{} {}\n'.format(speed_left, speed_right)
    _serial_port.write(str.encode(serial_string))

    print('Sending string:', serial_string)

def set_speeds_ramp(speed_left, speed_right, ramp_rate=200, pt_per_sec=10):
    '''
    Set left and right motor speeds, ramping from the current value.

    Parameters
    ----------
    speed_left : int in [700, 2000]
        Forward speed of left motor, from stopped to maximum.
    speed_right : int in [700, 2000]
        Forward speed of right motor, from stopped to maximum.
    ramp_rate : int
        Microseconds per second to ramp by.
    pts_per_sec : int
        Number of interpolation points traversed per second.
    '''
    if speed_left < _setpoint_left and speed_right < _setpoint_right:
        set_speeds(speed_left, speed_right)
        return

    vals_left = np.arange(_setpoint_left, speed_left, ramp_rate / pt_per_sec)
    vals_right = np.arange(_setpoint_right, speed_left, ramp_rate / pt_per_sec)
    delay_per_tick = 1. / pt_per_sec

    for new_speed_left, new_speed_right in zip(vals_left, vals_right):
        set_speeds(new_speed_left, new_speed_right)
        time.sleep(delay_per_tick)


import inspect
import logging
import sys
from copy import deepcopy

import numpy as np
from sklearn.model_selection import StratifiedKFold
from sklearn.utils import check_random_state

from snips_nlu_metrics.utils.constants import (
    DATA,
    ENTITIES,
    ENTITY,
    FALSE_NEGATIVE,
    FALSE_POSITIVE,
    INTENTS,
    NONE_INTENT_NAME,
    SLOT_NAME,
    TEXT,
    TRUE_POSITIVE,
    UTTERANCES,
    EXACT_PARSINGS,
)
from snips_nlu_metrics.utils.dataset_utils import (
    get_utterances_subset,
    input_string_from_chunks,
    update_entities_with_utterances,
)
from snips_nlu_metrics.utils.exception import NotEnoughDataError

logger = logging.getLogger(__name__)

INITIAL_METRICS = {TRUE_POSITIVE: 0, FALSE_POSITIVE: 0, FALSE_NEGATIVE: 0}

def create_shuffle_stratified_splits(
    dataset,
    n_splits,
    train_size_ratio=1.0,
    drop_entities=False,
    seed=None,
    out_of_domain_utterances=None,
    intents_filter=None,
):
    if train_size_ratio > 1.0 or train_size_ratio < 0:
        error_msg = "Invalid value for train size ratio: %s" % train_size_ratio
        logger.exception(error_msg)
        raise ValueError(error_msg)

    nb_utterances = {
        intent: len(data[UTTERANCES]) for intent, data in dataset[INTENTS].items()
    }
    if any((nb * train_size_ratio < n_splits for nb in nb_utterances.values())):
        raise NotEnoughDataError(dataset, n_splits, train_size_ratio)

    if drop_entities:
        dataset = deepcopy(dataset)
        for entity, data in dataset[ENTITIES].items():
            data[DATA] = []
    else:
        dataset = update_entities_with_utterances(dataset)

    utterances = np.array(
        [
            (intent_name, utterance)
            for intent_name, intent_data in dataset[INTENTS].items()
            for utterance in intent_data[UTTERANCES]
        ]
    )
    intents = np.array([u[0] for u in utterances])
    X = np.zeros(len(intents))
    random_state = check_random_state(seed)
    sss = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=random_state)
    splits = []
    for train_index, test_index in sss.split(X, intents):
        train_utterances = utterances[train_index].tolist()
        train_utterances = get_utterances_subset(train_utterances, train_size_ratio)
        test_utterances = utterances[test_index].tolist()
        train_dataset = deepcopy(dataset)
        train_dataset[INTENTS] = dict()
        for intent_name, utterance in train_utterances:
            if intent_name not in train_dataset[INTENTS]:
                train_dataset[INTENTS][intent_name] = {UTTERANCES: []}
            train_dataset[INTENTS][intent_name][UTTERANCES].append(deepcopy(utterance))
        splits.append((train_dataset, test_utterances))

    if intents_filter is not None:
        filtered_splits = []
        for train_dataset, test_utterances in splits:
            test_utterances = [
                (intent_name, utterance)
                for intent_name, utterance in test_utterances
                if intent_name in intents_filter
            ]
            filtered_splits.append((train_dataset, test_utterances))
        splits = filtered_splits

    if out_of_domain_utterances is not None:
        additional_test_utterances = [
            [NONE_INTENT_NAME, {DATA: [{TEXT: utterance}]}]
            for utterance in out_of_domain_utterances
        ]
        for split in splits:
            split[1].extend(additional_test_utterances)

    return splits

def compute_split_metrics(
    engine_class,
    split,
    intent_list,
    include_slot_metrics,
    slot_matching_lambda,
    intents_filter,
):
    """Fit and run engine on a split specified by train_dataset and
        test_utterances"""
    train_dataset, test_utterances = split
    engine = engine_class()
    engine.fit(train_dataset)
    return compute_engine_metrics(
        engine,
        test_utterances,
        intent_list,
        include_slot_metrics,
        slot_matching_lambda,
        intents_filter,
    )

def compute_engine_metrics(
    engine,
    test_utterances,
    intent_list,
    include_slot_metrics,
    slot_matching_lambda=None,
    intents_filter=None,
):
    if slot_matching_lambda is None:
        slot_matching_lambda = exact_match
    metrics = dict()
    intent_list = intent_list + [NONE_INTENT_NAME]
    confusion_matrix = dict(
        intents=intent_list,
        matrix=[[0 for _ in range(len(intent_list))] for _ in range(len(intent_list))],
    )
    intents_idx = {intent_name: idx for idx, intent_name in enumerate(intent_list)}

    errors = []
    for actual_intent, utterance in test_utterances:
        actual_slots = [chunk for chunk in utterance[DATA] if SLOT_NAME in chunk]
        input_string = input_string_from_chunks(utterance[DATA])
        if has_filter_param(engine):
            parsing = engine.parse(input_string, intents_filter=intents_filter)
        else:
            if intents_filter:
                logger.warning(
                    "The provided NLU engine (%r) does not support "
                    "intents filter through its `parse` API, "
                    "however one has been passed (%s)",
                    engine,
                    intents_filter,
                )
            parsing = engine.parse(input_string)

        if parsing["intent"] is not None:
            predicted_intent = parsing["intent"]["intentName"]
            if predicted_intent is None:
                predicted_intent = NONE_INTENT_NAME
        else:
            # Use a string here to avoid having a None key in the metrics dict
            predicted_intent = NONE_INTENT_NAME

        predicted_slots = [] if parsing["slots"] is None else parsing["slots"]

        i = intents_idx.get(actual_intent)
        j = intents_idx.get(predicted_intent)

        if i is None or j is None:
            continue

        confusion_matrix["matrix"][i][j] += 1

        utterance_metrics = compute_utterance_metrics(
            predicted_intent,
            predicted_slots,
            actual_intent,
            actual_slots,
            include_slot_metrics,
            slot_matching_lambda,
        )
        for intent in utterance_metrics:
            utterance_metrics[intent][EXACT_PARSINGS] = 0
        if contains_errors(utterance_metrics, include_slot_metrics):
            if not include_slot_metrics:
                parsing.pop("slots")
            errors.append(
                {
                    "nlu_output": parsing,
                    "expected_output": format_expected_output(
                        actual_intent, utterance, include_slot_metrics
                    ),
                }
            )
        else:
            utterance_metrics[actual_intent][EXACT_PARSINGS] = 1
        metrics = aggregate_metrics(metrics, utterance_metrics, include_slot_metrics)
    return metrics, errors, confusion_matrix

def has_filter_param(engine):
    if sys.version_info[0] == 2:
        parse_args = inspect.getargspec(engine.parse).args
    else:
        parse_args = inspect.signature(engine.parse).parameters
    return "intents_filter" in parse_args

def compute_utterance_metrics(
    predicted_intent,
    predicted_slots,
    actual_intent,
    actual_slots,
    include_slot_metrics,
    slot_matching_lambda,
):
    # initialize metrics
    intent_names = {predicted_intent, actual_intent}
    slot_names = set(
        [(predicted_intent, s["slotName"]) for s in predicted_slots]
        + [(actual_intent, u[SLOT_NAME]) for u in actual_slots]
    )

    metrics = dict()
    for intent in intent_names:
        metrics[intent] = {"intent": deepcopy(INITIAL_METRICS)}
        if include_slot_metrics:
            metrics[intent]["slots"] = dict()

    if include_slot_metrics:
        for (intent_name, slot_name) in slot_names:
            metrics[intent_name]["slots"][slot_name] = deepcopy(INITIAL_METRICS)

    if predicted_intent == actual_intent:
        metrics[predicted_intent]["intent"][TRUE_POSITIVE] += 1
    else:
        metrics[predicted_intent]["intent"][FALSE_POSITIVE] += 1
        metrics[actual_intent]["intent"][FALSE_NEGATIVE] += 1
        return metrics

    if not include_slot_metrics:
        return metrics

    # Check if expected slots have been parsed
    for slot in actual_slots:
        slot_name = slot[SLOT_NAME]
        slot_metrics = metrics[actual_intent]["slots"][slot_name]
        if any(
            s["slotName"] == slot_name and slot_matching_lambda(slot, s)
            for s in predicted_slots
        ):
            slot_metrics[TRUE_POSITIVE] += 1
        else:
            slot_metrics[FALSE_NEGATIVE] += 1

    # Check if there are unexpected parsed slots
    for slot in predicted_slots:
        slot_name = slot["slotName"]
        slot_metrics = metrics[predicted_intent]["slots"][slot_name]
        if all(
            s[SLOT_NAME] != slot_name or not slot_matching_lambda(s, slot)
            for s in actual_slots
        ):
            slot_metrics[FALSE_POSITIVE] += 1
    return metrics

def aggregate_metrics(lhs_metrics, rhs_metrics, include_slot_metrics):
    acc_metrics = deepcopy(lhs_metrics)
    for (intent, intent_metrics) in rhs_metrics.items():
        if intent not in acc_metrics:
            acc_metrics[intent] = deepcopy(intent_metrics)
        else:
            acc_metrics[intent]["intent"] = add_count_metrics(
                acc_metrics[intent]["intent"], intent_metrics["intent"]
            )
            acc_metrics[intent][EXACT_PARSINGS] += intent_metrics[EXACT_PARSINGS]
            if not include_slot_metrics:
                continue
            acc_slot_metrics = acc_metrics[intent]["slots"]
            for (slot, slot_metrics) in intent_metrics["slots"].items():
                if slot not in acc_slot_metrics:
                    acc_slot_metrics[slot] = deepcopy(slot_metrics)
                else:
                    acc_slot_metrics[slot] = add_count_metrics(
                        acc_slot_metrics[slot], slot_metrics
                    )
    return acc_metrics

def aggregate_matrices(lhs_matrix, rhs_matrix):
    if lhs_matrix is None:
        return rhs_matrix
    if rhs_matrix is None:
        return lhs_matrix
    acc_matrix = deepcopy(lhs_matrix)
    matrix_size = len(acc_matrix["matrix"])
    for i in range(matrix_size):
        for j in range(matrix_size):
            acc_matrix["matrix"][i][j] += rhs_matrix["matrix"][i][j]
    return acc_matrix

def add_count_metrics(lhs, rhs):
    return {
        TRUE_POSITIVE: lhs[TRUE_POSITIVE] + rhs[TRUE_POSITIVE],
        FALSE_POSITIVE: lhs[FALSE_POSITIVE] + rhs[FALSE_POSITIVE],
        FALSE_NEGATIVE: lhs[FALSE_NEGATIVE] + rhs[FALSE_NEGATIVE],
    }

def compute_average_metrics(metrics, ignore_none_intent=True):
    metrics = deepcopy(metrics)
    if ignore_none_intent:
        metrics = {
            intent: intent_metrics
            for intent, intent_metrics in metrics.items()
            if intent and intent != NONE_INTENT_NAME
        }

    nb_intents = len(metrics)
    if not nb_intents:
        return None

    average_intent_f1 = (
        sum(
            intent_metrics["intent"]["f1"] for intent, intent_metrics in metrics.items()
        )
        / nb_intents
    )
    average_intent_precision = (
        sum(
            intent_metrics["intent"]["precision"]
            for intent, intent_metrics in metrics.items()
        )
        / nb_intents
    )
    average_intent_recall = (
        sum(
            intent_metrics["intent"]["recall"]
            for intent, intent_metrics in metrics.items()
        )
        / nb_intents
    )

    average_metrics = {
        "intent": {
            "f1": average_intent_f1,
            "precision": average_intent_precision,
            "recall": average_intent_recall,
        },
    }

    nb_slots = sum(
        1
        for intent_metrics in metrics.values()
        for _ in intent_metrics.get("slots", dict()).values()
    )
    if nb_slots == 0:
        return average_metrics

    average_slot_f1 = (
        sum(
            slot_metrics["f1"]
            for intent_metrics in metrics.values()
            for slot_metrics in intent_metrics["slots"].values()
        )
        / nb_slots
    )
    average_slot_precision = (
        sum(
            slot_metrics["precision"]
            for intent_metrics in metrics.values()
            for slot_metrics in intent_metrics["slots"].values()
        )
        / nb_slots
    )
    average_slot_recall = (
        sum(
            slot_metrics["recall"]
            for intent_metrics in metrics.values()
            for slot_metrics in intent_metrics["slots"].values()
        )
        / nb_slots
    )

    average_metrics["slot"] = {
        "f1": average_slot_f1,
        "precision": average_slot_precision,
        "recall": average_slot_recall,
    }
    return average_metrics

def compute_precision_recall_f1(metrics):
    for intent_metrics in metrics.values():
        prec_rec_metrics = _compute_precision_recall_f1(intent_metrics["intent"])
        intent_metrics["intent"].update(prec_rec_metrics)
        if "slots" in intent_metrics:
            for slot_metrics in intent_metrics["slots"].values():
                prec_rec_metrics = _compute_precision_recall_f1(slot_metrics)
                slot_metrics.update(prec_rec_metrics)
    return metrics

def _compute_precision_recall_f1(count_metrics):
    tp = count_metrics[TRUE_POSITIVE]
    fp = count_metrics[FALSE_POSITIVE]
    fn = count_metrics[FALSE_NEGATIVE]
    precision = 0.0 if tp == 0 else float(tp) / float(tp + fp)
    recall = 0.0 if tp == 0 else float(tp) / float(tp + fn)
    if precision == 0.0 or recall == 0.0:
        f1 = 0.0
    else:
        f1 = 2 * (precision * recall) / (precision + recall)
    return {"precision": precision, "recall": recall, "f1": f1}

def contains_errors(utterance_metrics, check_slots):
    for metrics in utterance_metrics.values():
        intent_metrics = metrics["intent"]
        if intent_metrics.get(FALSE_POSITIVE, 0) > 0:
            return True
        if intent_metrics.get(FALSE_NEGATIVE, 0) > 0:
            return True
        if not check_slots:
            continue
        for slot_metrics in metrics["slots"].values():
            if slot_metrics.get(FALSE_POSITIVE, 0) > 0:
                return True
            if slot_metrics.get(FALSE_NEGATIVE, 0) > 0:
                return True
    return False

def format_expected_output(intent_name, utterance, include_slots):
    char_index = 0
    ranges = []
    for chunk in utterance[DATA]:
        range_end = char_index + len(chunk[TEXT])
        ranges.append({"start": char_index, "end": range_end})
        char_index = range_end

    expected_output = {
        "input": "".join(chunk[TEXT] for chunk in utterance[DATA]),
        "intent": {"intentName": intent_name, "probability": 1.0},
    }
    if include_slots:
        expected_output["slots"] = [
            {
                "rawValue": chunk[TEXT],
                "entity": chunk[ENTITY],
                "slotName": chunk[SLOT_NAME],
                "range": ranges[chunk_index],
            }
            for chunk_index, chunk in enumerate(utterance[DATA])
            if ENTITY in chunk
        ]
    return expected_output

def exact_match(lhs_slot, rhs_slot):
    return lhs_slot[TEXT] == rhs_slot["rawValue"]