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

import calibration as cal
import dill
import numpy as np

from fastestimator.summary.system import System
from fastestimator.trace.trace import Trace
from fastestimator.util.data import Data
from fastestimator.util.traceability_util import traceable
from fastestimator.util.util import to_list, to_number


@traceable()
class PBMCalibrator(Trace):
    """A trace to generate a PlattBinnerMarginalCalibrator given a set of predictions.

    Unlike many common calibration error correction algorithms, this one has actual theoretical bounds on the quality
    of its output: https://arxiv.org/pdf/1909.10155v1.pdf. This trace is commonly used together with the Calibrate
    NumpyOp for postprocessing. This trace will collect data from whichever `mode` it is set to run on in order to
    perform empirical probability calibration. The calibrated predictions will be output on epoch end. The trained
    calibration function will also be saved if `save_path` is provided.

    Args:
        true_key: Name of the key that corresponds to ground truth in the batch dictionary.
        pred_key: Name of the key that corresponds to predicted score in the batch dictionary.
        mode: What mode(s) to execute this Trace in. For example, "train", "eval", "test", or "infer". To execute
            regardless of mode, pass None. To execute in all modes except for a particular one, you can pass an argument
            like "!infer" or "!train".
        ds_id: What dataset id(s) to execute this Trace in. To execute regardless of ds_id, pass None. To execute in all
            ds_ids except for a particular one, you can pass an argument like "!ds1".
        output_name: What to call the output from this trace. If None, the default will be '<pred_key>_calibrated'.
        save_path: Where to save the calibrator generated by this Trace. If None, then no saving will be performed.
        save_if_key: Name of a key to control whether to save the calibrator. For example "since_best_acc". If provided,
            then the calibrator will only be saved when the save_if_key value is zero.

    Raises:
        ValueError: If 'save_if_key' is provided but no 'save_path' is given.
    """
    system: System

    def __init__(self,
                 true_key: str,
                 pred_key: str,
                 output_name: Optional[str] = None,
                 save_path: Optional[str] = None,
                 save_if_key: Optional[str] = None,
                 mode: Union[str, Set[str]] = "eval",
                 ds_id: Union[None, str, Iterable[str]] = None) -> None:
        if output_name is None:
            output_name = f"{pred_key}_calibrated"
        if save_if_key is not None and save_path is None:
            raise ValueError("If 'save_if_key' is provided, then a 'save_path' must also be provided.")
        super().__init__(inputs=[true_key, pred_key] + to_list(save_if_key),
                         outputs=output_name,
                         mode=mode,
                         ds_id=ds_id)
        self.y_true = []
        self.y_pred = []
        if save_path is not None:
            save_path = os.path.abspath(os.path.normpath(save_path))
        self.save_path = save_path

    @property
    def true_key(self) -> str:
        return self.inputs[0]

    @property
    def pred_key(self) -> str:
        return self.inputs[1]

    @property
    def save_key(self) -> Optional[str]:
        if len(self.inputs) == 3:
            return self.inputs[2]
        return None

    def on_epoch_begin(self, data: Data) -> None:
        self.y_true = []
        self.y_pred = []

    def on_batch_end(self, data: Data) -> None:
        y_true, y_pred = to_number(data[self.true_key]), to_number(data[self.pred_key])
        if y_true.shape[-1] > 1 and y_true.ndim > 1:
            y_true = np.argmax(y_true, axis=-1)
        assert y_pred.shape[0] == y_true.shape[0]
        self.y_true.extend(y_true)
        self.y_pred.extend(y_pred)

    def on_epoch_end(self, data: Data) -> None:
        self.y_true = np.squeeze(np.stack(self.y_true))
        self.y_pred = np.stack(self.y_pred)
        calibrator = cal.PlattBinnerMarginalCalibrator(num_calibration=len(self.y_true), num_bins=10)
        calibrator.train_calibration(probs=self.y_pred, labels=self.y_true)
        if self.save_path:
            if not self.save_key or (self.save_key and to_number(data[self.save_key]) == 0):
                with open(self.save_path, 'wb') as f:
                    dill.dump(calibrator.calibrate, file=f)
                print(f"FastEstimator-PBMCalibrator: Calibrator written to {self.save_path}")
        data.write_without_log(self.outputs[0], calibrator.calibrate(self.y_pred))
