# 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.


import numpy as np
from typing import List, Optional, Union
from langchain_core.language_models.chat_models import BaseChatModel

from uqlm.judges.judge import LLMJudge
from uqlm.scorers.baseclass.uncertainty import UncertaintyQuantifier, UQResult


class LLMPanel(UncertaintyQuantifier):
    def __init__(
        self,
        judges: List[Union[LLMJudge, BaseChatModel]],
        llm: Optional[BaseChatModel] = None,
        system_prompt: str = "You are a helpful assistant.",
        max_calls_per_min: Optional[int] = None,
    ) -> None:
        """
        Class for aggregating multiple instances of LLMJudge using min, max, or majority voting

        Parameters
        ----------
        judges: list of LLMJudge or BaseChatModel
            Judges to use. If BaseChatModel, LLMJudge is instantiated using default 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. Will be used for original llm and any judges constructed
            from instances of BaseChatModel in judges

        system_prompt : str or None, default="You are a helpful assistant."
            Optional argument for user to provide custom system prompt
        """
        super().__init__(
            llm=llm,
            max_calls_per_min=max_calls_per_min,
            system_prompt=system_prompt,
        )
        self.judges = []
        for judge in judges:
            if isinstance(judge, BaseChatModel):
                judge = LLMJudge(llm=judge, max_calls_per_min=max_calls_per_min)
            elif not isinstance(judge, LLMJudge):
                raise ValueError(
                    "judges must be a list containing instances of either LLMJudge or BaseChatModel"
                )
            self.judges.append(judge)

    async def generate_and_score(
        self,
        prompts: List[str],
    ) -> UQResult:
        """
        Generate LLM responses to provided prompts and use panel of judges to score responses for correctness.

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

        Returns
        -------
        UQResult
            UQResult containing prompts, responses, Q/A concatenations, judge responses, and judge scores
        """
        responses = await self.generate_original_responses(prompts)
        return await self.score(prompts=prompts, responses=responses)
    
    async def score(
        self,
        prompts: List[str],
        responses: Optional[List[str]] = None,
    ) -> UQResult:
        """
        Use panel to of judges to score provided responses for correctness. Use if responses are already generated. Otherwise,
        use generate_and_score.

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

        responses: list of str, default = None
            A list of LLM responses for the corresponding to the provided prompts. 

        Returns
        -------
        UQResult
            UQResult containing prompts, responses, Q/A concatenations, judge responses, and judge scores
        """
        self.prompts = prompts
        self.responses = responses
        data = {
            "prompts": prompts,
            "responses": responses,
        }
        
        judge_count = 1
        scores_lists = []
        for judge in self.judges:
            tmp = await judge.judge_responses(prompts=prompts, responses=responses)
            scores_lists.append(tmp["scores"])
            data[f"judge_{judge_count}"] = tmp["scores"]
            judge_count += 1

        scores_dict = {
            "avg": [np.mean(scores) for scores in zip(*scores_lists)],
            "max": [np.max(scores) for scores in zip(*scores_lists)],
            "min": [np.min(scores) for scores in zip(*scores_lists)],
            "median": [np.median(scores) for scores in zip(*scores_lists)],
        }
        data.update(scores_dict)
        result = {
            "data": data,
            "metadata": {
                "num_judges": len(self.judges),
                "temperature": None if not self.llm else self.llm.temperature,
            },
        }
        return UQResult(result)