# 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 math
import torch

from typing import Any

import torch

from torchcp.utils.registry import Registry

METRICS_REGISTRY_LLM = Registry("METRICS")


@METRICS_REGISTRY_LLM.register()
def average_size(prediction_sets):
    """the average size of the prediction sets"""
    size_avg = torch.mean(prediction_sets)
    return size_avg.cpu().item()


@METRICS_REGISTRY_LLM.register()
def average_sample_size(prediction_sets):
    """the average number of sample size"""
    max_indices = prediction_sets.shape[1] - 1 - torch.argmax(prediction_sets.flip(1), dim=1)
    return torch.mean(max_indices.float()).cpu().item()


@METRICS_REGISTRY_LLM.register()
def average_set_loss(prediction_sets, prediction_set_loss):
    """
    Calculate the average set loss.

    Args:
        prediction_sets (torch.Tensor): The prediction sets generated by CP algorithms.
        prediction_set_loss (torch.Tensor): The loss associated with each prediction set.

    Returns:
        torch.Tensor: The average set loss.
    """

    max_indices = prediction_sets.shape[1] - 1 - torch.argmax(prediction_sets.flip(1).to(torch.int), dim=1)

    losses = prediction_set_loss[torch.arange(prediction_sets.shape[0]), max_indices]
    return torch.mean(losses.float()).cpu().item()


@METRICS_REGISTRY_LLM.register()
def SSCL(prediction_sets, prediction_set_loss, num_bins=20):
    """
    Size-stratified conditional loss.
    
    Paper: Conformal Language Modeling (Victor Quach et al., ICLR'24)
    
    Args:
        prediction_sets (torch.Tensor): The prediction sets generated by CP algorithms.
        prediction_set_loss (torch.Tensor): The loss associated with each prediction set.
        num_bins (int, optional): The number of bins for stratification. Default is 20.

    Returns:
        torch.Tensor: The size-stratified conditional loss.
    """

    prediction_sets = prediction_sets.to(torch.float32).clone().detach()
    prediction_sizes = torch.sum(prediction_sets, dim=0)
    bins = torch.quantile(prediction_sizes, torch.linspace(0, 1, num_bins, dtype=prediction_sets.dtype))

    binids = torch.bucketize(prediction_sizes, torch.cat([torch.tensor([0]), torch.unique(bins)]))

    L_worst_avg = -1
    for binid in torch.unique(binids):
        kept = binids == binid
        num_kept_examples = torch.maximum(torch.sum(kept), torch.tensor(1, device=kept.device))
        Ls_mask_avg = torch.sum(prediction_set_loss * kept) / num_kept_examples
        L_worst_avg = max(L_worst_avg, Ls_mask_avg)
    return L_worst_avg.cpu().item()


class Metrics:
    def __call__(self, metric) -> Any:
        if metric not in METRICS_REGISTRY_LLM.registered_names():
            raise NameError(f"The metric: {metric} is not defined in TorchCP.")
        return METRICS_REGISTRY_LLM.get(metric)
