"""
WorldSense: A Synthetic Benchmark for Grounded Reasoning
in Large Language Models

Youssef Benchekroun, Megi Dervishi, Mark Ibrahim, Jean-Baptiste Gaya, Xavier Martinet, Grégoire Mialon, Thomas Scialom, Emmanuel Dupoux, Dieuwke Hupkes, Pascal Vincent
https://arxiv.org/pdf/2311.15930

# eval all problemnames w/ 500 randomly selected samples
inspect eval worldsense --limit 500

# add chain of thought
inspect eval worldsense --limit 500

# eval selected problem
inspect eval worldsense -T problemnames=Compl.normal
inspect eval worldsense -T problemnames=Consist.normal
inspect eval worldsense -T problemnames=Infer.normal,Infer.trivial
"""  # noqa: D205

import bz2
import json
from typing import Any, Callable, Dict

import requests
from inspect_ai import Task, task
from inspect_ai.dataset import Dataset, MemoryDataset, Sample
from inspect_ai.model import GenerateConfig
from inspect_ai.scorer import (
    Metric,
    SampleScore,
    Score,
    Scorer,
    Target,
    accuracy,
    metric,
    pattern,
    scorer,
    stderr,
)
from inspect_ai.solver import TaskState, generate

from ._utils import compute_accuracy, compute_bias, preprocess_scores


@task
def worldsense(problemnames: str | list[str] = []) -> Task:
    """
    Task for evaluating reasoning related to world descriptions. There are three problem types ("Infer", "Compl", "Consist") and two grades ("trivial", "normal"). A problemname is formed by concatenation "<type>.<grade>". See README for details.

    Args:
        problemnames (str | list[str], optional): A string or list of strings specifying the names of problems or tasks to filter the dataset. If provided, it filters the dataset to samples that contain matching metadata for the specified problem names.

    Returns:
        Task: A task object configured with a dataset filtered by problem names (if specified), a solver, a scoring pattern for evaluating task responses, and custom metrics.
    """
    dataset = load_worldsense_dataset(sample_fields=record_to_sample, shuffle=True)

    # filter dataset if requested
    problemnames = problemnames if isinstance(problemnames, list) else [problemnames]
    if len(problemnames) > 0:
        task_dataset = dataset.filter(
            name=f"{dataset.name}-{'-'.join(problemnames)}",
            predicate=lambda sample: sample.metadata is not None
            and sample.metadata.get("problemname") in problemnames,
        )
    else:
        task_dataset = dataset

    return Task(
        dataset=task_dataset,
        solver=generate(),
        scorer=pattern_with_metadata(
            r"^\(?\s*(1|2|3|TRUE|FALSE|IMPOSSIBLE|POSSIBLE)\s*\)?"
        ),
        metrics=[accuracy(), stderr(), ws_accuracy(), ws_bias()],
        config=GenerateConfig(temperature=0.0),
    )


def record_to_sample(record: Dict[str, Any]) -> Sample:
    goldresp_mapping: Dict[str, str] = {
        "Emmanuel": "TRUE",
        "Megi": "FALSE",
        "Dieuwke": "POSSIBLE",
        "Pascal": "IMPOSSIBLE",
        "Mark": "1",
        "Youssef": "2",
        "Yoda": "3",
    }

    return Sample(
        input=record["text"],
        choices=record["expectedresp"],
        target=goldresp_mapping.get(record["goldresp_obfusc"], "None"),
        metadata={
            "tuple_ID": record["tuple_ID"],
            "problemname": record["problemname"],
            "problemsize": record["problemsize"],
        },
    )


def load_worldsense_dataset(
    sample_fields: Callable[[dict[str, Any]], Sample],
    shuffle: bool = True,
) -> Dataset:
    """
    Load the worldsense dataset from a bz2 file directly into memory and return a Dataset.

    Args:
        sample_fields (Callable): Function to map records to samples.
        shuffle (bool): Whether to shuffle the dataset. Default is True.

    Returns:
        Dataset: The loaded and decompressed dataset in memory.
    """
    url = "https://github.com/facebookresearch/worldsense/raw/bd81d945077f169cf95ff39207f788f86e4645e9/data/worldsense/test_set/trials.jsonl.bz2"

    # Download and decompress the bz2 data directly into memory
    response = requests.get(url)
    response.raise_for_status()
    decompressed_data = bz2.decompress(response.content).decode("utf-8")

    # Parse the decompressed data into records
    lines = decompressed_data.strip().splitlines()
    samples = []
    for line in lines:
        record = json.loads(line)
        sample = sample_fields(record)
        if isinstance(sample, list):
            samples.extend(sample)
        else:
            samples.append(sample)

    # Create a MemoryDataset
    dataset = MemoryDataset(samples=samples, name="worldsense", location=url)

    # Shuffle if needed
    if shuffle:
        dataset.shuffle()

    return dataset


@scorer(metrics=[accuracy(), stderr()])
def pattern_with_metadata(
    pattern_str: str, ignore_case: bool = True, match_all: bool = False
) -> Scorer:
    base_scorer = pattern(pattern_str, ignore_case, match_all)

    async def score(state: TaskState, target: Target) -> Score:
        base_score = await base_scorer(state, target)
        base_score.metadata = state.metadata
        return base_score

    return score


@metric
def ws_accuracy() -> Metric:
    """Compute weighted accuracy metric."""

    def metric(scores: list[SampleScore]) -> float:
        grouped_scores = preprocess_scores(
            [sample_score.score for sample_score in scores]
        )
        return compute_accuracy(grouped_scores)

    return metric


@metric
def ws_bias() -> Metric:
    """Compute weighted bias metric."""

    def metric(scores: list[SampleScore]) -> float:
        grouped_scores = preprocess_scores(
            [sample_score.score for sample_score in scores]
        )
        return compute_bias(grouped_scores)

    return metric
