# Copyright (c) 2023-present, SUSTech-ML.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#

import torch

from torchcp.regression.predictor.base import BasePredictor


class SplitPredictor(BasePredictor):
    """
    Split Conformal Prediction for Regression.
    
    This predictor allows for the construction of a prediction band for the response 
    variable using any estimator of the regression function.
        
    Args:
        score_function (torchcp.regression.scores): A class that implements the score function.
        model (torch.nn.Module): A pytorch regression model that can output predicted point.
        alpha (float, optional): The significance level. Default is 0.1.
        device (torch.device, optional): The device on which the model is located. Default is None.
        
    Reference:
        Paper: Distribution-Free Predictive Inference For Regression (Lei et al., 2017)
        Link: https://arxiv.org/abs/1604.04173
        Github: https://github.com/ryantibs/conformal
    """

    def __init__(self, score_function, model=None, alpha=0.1, device=None):
        super().__init__(score_function, model, alpha, device)

    def train(self, train_dataloader, **kwargs):
        """
        Trains the model using the provided train_dataloader and score_function.

        Args:
            train_dataloader (DataLoader): DataLoader for the training data.
            **kwargs: Additional keyword arguments for training configuration.
                - model (nn.Module, optional): The model to be trained.
                - epochs (int, optional): Number of training epochs.
                - criterion (nn.Module, optional): Loss function.
                - lr (float, optional): Learning rate for the optimizer.
                - optimizer (torch.optim.Optimizer, optional): Optimizer for training.
                - verbose (bool, optional): If True, prints training progress.

        .. note::
            This function is optional but recommended, because the training process for each score_function is different. 
            We provide a default training method, and users can change the hyperparameters :attr:`kwargs` to modify the training process.
            If the train function is not used, users should pass the trained model to the predictor at the beginning.
        """
        model = kwargs.pop('model', None)
        device = kwargs.pop('device', self._device)
        self._device = device

        if model is not None:
            self._model = self.score_function.train(
                train_dataloader, model=model, device=device, **kwargs
            )
        elif self._model is not None:
            self._model = self.score_function.train(
                train_dataloader, model=self._model, device=device, **kwargs
            )
        else:
            self._model = self.score_function.train(
                train_dataloader, device=device, **kwargs
            )
            
        if self.score_function.__class__.__name__ == 'NorABS':
            self.score_function.calibrate(self._model)

    def calculate_score(self, predicts, y_truth, x_batch=None):
        """
        Calculate the nonconformity scores based on predictions and true values.

        The exact score calculation is determined by the configured score_function.

        Args:
            predicts (torch.Tensor): Model predictions.
            y_truth (torch.Tensor): Ground truth values.
            x_batch (torch.Tensor, optional): The corresponding batch of input features.
                This is only required by normalized score functions (e.g., 'NorABS')
                to compute adaptive nonconformity scores. Defaults to None.

        Returns:
            torch.Tensor: The computed nonconformity score for each sample.
        """
        if self.score_function.__class__.__name__ == 'NorABS':
            return self.score_function(predicts, y_truth, x_batch)
        else:
            return self.score_function(predicts, y_truth)

    def generate_intervals(self, predicts_batch, q_hat, x_batch=None):
        """
        Generate prediction intervals based on the model's predictions and the conformal value.

        Args:
            predicts_batch (torch.Tensor): Batch of predictions from the model.
            q_hat (float): Conformal value computed during calibration.

        Returns:
            torch.Tensor: Prediction intervals.
        """
        
        if self.score_function.__class__.__name__ == 'NorABS':
            return self.score_function.generate_intervals(predicts_batch, q_hat, x_batch)
        else:
            return self.score_function.generate_intervals(predicts_batch, q_hat)

    def calibrate(self, cal_dataloader, alpha=None):
        if alpha is None:
            alpha = self.alpha

        self._model.eval()
        x_list, predicts_list, y_truth_list = [], [], []
        with torch.no_grad():
            for tmp_x, tmp_labels in cal_dataloader:
                tmp_x, tmp_labels = tmp_x.to(self._device), tmp_labels.to(self._device)
                tmp_predicts = self._model(tmp_x).detach()
                x_list.append(tmp_x)
                predicts_list.append(tmp_predicts)
                y_truth_list.append(tmp_labels)

        predicts = torch.cat(predicts_list).float().to(self._device)
        y_truth = torch.cat(y_truth_list).to(self._device)
        x_batch = torch.cat(x_list).float().to(self._device)
        if self.score_function.__class__.__name__ == 'NorABS':
            self.scores = self.calculate_score(predicts, y_truth, x_batch)
        else:
            self.scores = self.calculate_score(predicts, y_truth)
        self.q_hat = self._calculate_conformal_value(self.scores, alpha)

    def predict(self, x_batch):
        self._model.eval()
        x_batch = x_batch.to(self._device)
        with torch.no_grad():
            predicts_batch = self._model(x_batch)
            return self.generate_intervals(predicts_batch, self.q_hat, x_batch)

    def evaluate(self, data_loader):
        y_list, predict_list = [], []
        with torch.no_grad():
            for tmp_x, tmp_y in data_loader:
                tmp_x, tmp_y = tmp_x.to(self._device), tmp_y.to(self._device)
                tmp_prediction_intervals = self.predict(tmp_x)
                y_list.append(tmp_y)
                predict_list.append(tmp_prediction_intervals)

        predicts = torch.cat(predict_list, dim=0).to(self._device)
        test_y = torch.cat(y_list).to(self._device)

        res_dict = {
            "coverage_rate": self._metric('coverage_rate')(predicts, test_y),
            "average_size": self._metric('average_size')(predicts)
        }

        return res_dict

    def predict_p(self, x_batch, y_batch, smooth=False):
        """
        Compute p-values for conformal prediction.
        
        Args:
            x_batch (torch.Tensor): A batch of instances.
            y_batch (torch.Tensor): A batch of labels for instances.
            smooth (bool): Whether to apply randomized smoothing when calibration scores equal test scores.

        Returns:
            Tensor: p-values for each test sample and class, shape (n_test, k)
        """

        if self._model is None:
            raise ValueError("Model is not defined. Please provide a valid model.")

        self._model.eval()
        with torch.no_grad():
            x_batch = self._model(x_batch.to(self._device))
            test_scores = self.calculate_score(x_batch, y_batch)
        
        cal_scores = self.scores

        n_cal = cal_scores.size(0)
        n_test = test_scores.size(0)

        cal_scores_expanded = cal_scores.view(1, n_cal)
        test_scores_expanded = test_scores.view(n_test, 1)

        greater = (cal_scores_expanded > test_scores_expanded).sum(dim=1)
        equal = (cal_scores_expanded == test_scores_expanded).sum(dim=1)

        if smooth:
            tau = torch.rand_like(equal, dtype=torch.float)
            p_values = (greater + tau * (equal + 1)) / (n_cal + 1)
        else:
            p_values = (greater + (equal + 1)) / (n_cal + 1)

        return p_values
