# Copyright 2025 CVS Health and/or one of its affiliates
#
# 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.

from typing import Any, Dict, List, Optional
import math
import numpy as np
from langchain_core.language_models.chat_models import BaseChatModel
from rich import print as rprint

from uqlm.scorers.baseclass.uncertainty import UncertaintyQuantifier, UQResult


class WhiteBoxUQ(UncertaintyQuantifier):
    def __init__(self, llm: Optional[BaseChatModel] = None, system_prompt: str = "You are a helpful assistant.", max_calls_per_min: Optional[int] = None, scorers: Optional[List[str]] = None) -> None:
        """
        Class for computing white-box UQ confidence scores. This class offers two confidence scores, normalized
        probability :footcite:`malinin2021uncertaintyestimationautoregressivestructured` and minimum probability :footcite:`manakul2023selfcheckgptzeroresourceblackboxhallucination`.

        Parameters
        ----------
        llm : BaseChatModel
            A langchain llm object to get passed to chain constructor. User is responsible for specifying
            temperature and other relevant parameters to the constructor of their `llm` object.

        max_calls_per_min : int, default=None
            Used to control rate limiting.

        system_prompt : str or None, default="You are a helpful assistant."
            Optional argument for user to provide custom system prompt

        scorers : subset of {
            "imperplexity", "geometric_mean_probability", "min_probability", "max_probability",
        }, default=None
            Specifies which black box (consistency) scorers to include. If None, defaults to all.
        """
        super().__init__(llm=llm, max_calls_per_min=max_calls_per_min, system_prompt=system_prompt)
        self.scorers = scorers if scorers else self.white_box_names

    async def generate_and_score(self, prompts: List[str], progress_bar: Optional[bool] = True) -> UQResult:
        """
        Generate responses and compute white-box confidence scores based on extracted token probabilities.

        Parameters
        ----------
        prompts : list of str
            A list of input prompts for the model.

        progress_bar : bool, default=True
            If True, displays a progress bar while generating and scoring responses

        Returns
        -------
        UQResult
            UQResult containing prompts, responses, logprobs, and white-box UQ scores
        """
        assert hasattr(self.llm, "logprobs"), """
        BaseChatModel must have logprobs attribute and have logprobs=True
        """
        if progress_bar:
            rprint("🤖📈 Generation & Scoring")
        self.llm.logprobs = True
        responses = await self.generate_original_responses(prompts, progress_bar=progress_bar)
        return self.score(prompts=prompts, responses=responses, logprobs_results=self.logprobs)

    def score(self, logprobs_results: List[List[Dict[str, Any]]], prompts: Optional[List[str]] = None, responses: Optional[List[str]] = None) -> UQResult:
        """
        Compute white-box confidence scores from provided logprobs.

        Parameters
        ----------
        logprobs_results : list of logprobs_result
            List of dictionaries, each returned by BaseChatModel.agenerate

        prompts : list of str, default=None
            A list of input prompts for the model.

        responses : list of str, default=None
            A list of model responses for the prompts.

        Returns
        -------
        UQResult
            UQResult containing prompts, responses, logprobs, and white-box UQ scores
        """

        self.logprobs = logprobs_results
        self.prompts = prompts
        self.responses = responses

        data = {}
        if self.prompts:
            data["prompts"] = self.prompts
        if self.responses:
            data["responses"] = self.responses

        data["logprobs"] = self.logprobs
        scores = self._compute_scores(self.logprobs)
        for key in self.scorers:
            data[key] = scores[key]

        result = {"data": data, "metadata": {"temperature": None if not self.llm else self.llm.temperature}}
        return UQResult(result)

    def _compute_scores(self, logprobs_results: List[List[Dict[str, Any]]]) -> List[float]:
        """
        This method computes token-probability-based confidence scores.
        """
        return {"normalized_probability": [np.nan if not r else self._norm_prob(r) for r in logprobs_results], "min_probability": [np.nan if not r else self._min_prob(r) for r in logprobs_results]}

    def _norm_prob(self, logprobs: List[Dict[str, Any]]) -> float:
        """Compute normalized token probability"""
        return math.exp(self.avg_logprob(logprobs))

    def _min_prob(self, logprobs: List[Dict[str, Any]]) -> float:
        """Compute minimum token probability"""
        return min(self._get_probs(logprobs))

    def avg_logprob(self, logprobs: List[Dict[str, Any]]) -> float:
        "Compute average logprob"
        return np.mean(self.get_logprobs(logprobs))

    @staticmethod
    def _get_probs(logprobs):
        """Extract token probabilities"""
        return [math.exp(d["logprob"]) for d in logprobs]

    @staticmethod
    def get_logprobs(logprobs):
        """Extract log token probabilities"""
        return [d["logprob"] for d in logprobs]
