from typing import Sequence

import pandas as pd

from ...datasets.base import Dataset
from ...llm.evaluators.coherency import CoherencyEvaluator
from ...llm.generators.sycophancy import SycophancyDataGenerator
from ...llm.utils import format_chat_messages
from ...models.base.model import BaseModel
from ...testing.tests.llm.hallucination import test_llm_output_coherency
from ..decorators import detector
from ..issues import Hallucination, Issue, IssueLevel
from ..scanner import logger
from .base import _estimate_base_token_counts


@detector(
    "llm_basic_sycophancy",
    tags=["hallucination", "misinformation", "sycophancy", "llm", "generative", "text_generation"],
)
class LLMBasicSycophancyDetector:
    """Detects sycophancy in LLM-based models.

    You can explicitly run this detector by adding the tag "sycophancy" in the `only` parameter of the scan method.

    Sycophancy is the tendency of a model to produce outputs that agree with the input bias. This is often linked to
    model hallucination. This detector will probe if the model is affected by this issue by generating adversarial
    inputs (based on the model name & description) and that the model outputs are coherent.

    Technically, we will generate pairs of inputs which are biased in a different way, and check whether the model
    produces coherent outputs.

    For example, a question answering model may be proposed with the following prompts:

    - "Hi, is the Earth flat?"
    - "My teacher told me that the Earth is flat, is that true?"

    We would expect the model to produce the same answer to both questions, but if the model is affected by sycophancy,
    it may follow the implicit bias in the second prompt and change its answer.

    Note that we will generate case specific adversarial inputs based on the model name and description, so that the
    inputs and biases are relevant and adapted to the model.
    """

    def __init__(self, num_samples=10):
        """
        Parameters
        ----------
        num_samples : Optional[int]
            Number of adversarial inputs to generate for each model. By default, we generate 10 adversarial inputs.
        """
        self.num_samples = num_samples

    def get_cost_estimate(self, model: BaseModel, dataset: Dataset) -> dict:
        counts = _estimate_base_token_counts(model, dataset)
        model_meta_tokens = counts["model_meta_tokens"]
        input_sample_tokens = counts["input_sample_tokens"]

        num_calls = 0
        num_prompt_tokens = 0
        num_sampled_tokens = 0

        # Data generation
        num_calls += 1
        num_prompt_tokens += 500 + model_meta_tokens
        num_sampled_tokens += input_sample_tokens * self.num_samples * 2

        # Evaluation, for each generated sample
        num_calls += self.num_samples
        num_prompt_tokens += self.num_samples * (400 + model_meta_tokens + input_sample_tokens * 2 + 50 * 2)
        num_sampled_tokens += self.num_samples * 15

        return {
            "model_predict_calls": 2 * self.num_samples,
            "llm_calls": num_calls,
            "llm_prompt_tokens": num_prompt_tokens,
            "llm_sampled_tokens": num_sampled_tokens,
        }

    def run(self, model: BaseModel, dataset: Dataset, features=None) -> Sequence[Issue]:
        # Prepare datasets
        languages = dataset.extract_languages(columns=model.meta.feature_names)

        generator = SycophancyDataGenerator(languages=languages)
        dataset1, dataset2 = generator.generate_dataset(
            model, num_samples=self.num_samples, column_types=dataset.column_types
        )
        logger.debug(f"{self.__class__.__name__}: Generated {len(dataset1)} test inputs for model assessment.")

        # Evaluate the answers
        evaluator = CoherencyEvaluator()
        eval_result = evaluator.evaluate(model, dataset1, dataset2)

        examples = pd.DataFrame(
            [
                {
                    "Conversation 1": format_chat_messages(r["sample"].get("conversation_1", [])),
                    "Conversation 2": format_chat_messages(r["sample"].get("conversation_2", [])),
                    "Reason": r["reason"],
                }
                for r in eval_result.failure_examples
            ]
        )

        if eval_result.failed:
            return [
                Issue(
                    model,
                    dataset2,
                    group=Hallucination,
                    level=IssueLevel.MAJOR,
                    description=(
                        "The model seems to produce contradicting or hallucinated outputs when we provide two sets "
                        "of inputs which are biased in a different way. This is generally a symptom of hallucination "
                        "or sycophancy, i.e. the tendency to produce outputs that agree with the input bias."
                    ),
                    examples=examples,
                    meta={
                        "domain": "Sycophancy",
                        "deviation": "The model produces incoherent or hallucinated output when prompted with biased inputs.",
                        "hide_index": True,
                        "dataset_1": dataset1,
                        "dataset_2": dataset2,
                        "metric": "Failing samples",
                        "metric_value": len(eval_result.failure_examples),
                    },
                    tests=_generate_sycophancy_tests,
                    taxonomy=["avid-effect:ethics:E0402"],
                    detector_name=self.__class__.__name__,
                )
            ]

        return []


def _generate_sycophancy_tests(issue: Issue):
    return {
        "Basic Sycophancy": test_llm_output_coherency(
            dataset_1=issue.meta["dataset_1"], dataset_2=issue.meta["dataset_2"]
        )
    }
