import evaluate
import lm_eval
from typing import Union, List, Optional
from dmx.compressor.dmx import config_rules, DmxModel
import datasets
import torch

_DESCRIPTION = """
Evaluation function using lm-eval with d-Matrix integration.
This function allows for the evaluation of language models across various tasks, 
with the option to use d-Matrix compressed models.
"""

_KWARGS_DESCRIPTION = """
Args:
    model (str): The name or path of the model to evaluate.
    tasks (Union[str, List[str]]): The task or list of tasks to evaluate on.
    dmx_config (Optional[str]): Configuration string for d-Matrix transformations, defaults to None.
    num_fewshot (Optional[int]): Number of examples in few-shot context, defaults to None.
    batch_size (Optional[Union[int, str]]): Batch size for model, defaults to None.
    max_batch_size (Optional[int]): Maximum batch size to try with automatic batch size detection, defaults to None.
    limit (Optional[Union[int, float]]): Limit the number of examples per task, defaults to None.
    device (Optional[str]): Device to run on. If None, defaults to 'cuda' if available, otherwise 'cpu'.
    revision (str): Model revision to use, defaults to 'main'.
    trust_remote_code (bool): Whether to trust remote code, defaults to False.
    log_samples (bool): If True, logs all model outputs and documents, defaults to True.
    verbosity (str): Logging verbosity level, defaults to 'INFO'.
    **kwargs: Additional keyword arguments to pass to `lm_eval.evaluate`.

Returns:
    dict: A dictionary containing the evaluation results.
"""

@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class  DmxMetric(evaluate.Metric):
    def _info(self):
        return evaluate.MetricInfo(
            module_type="metric",
            description=_DESCRIPTION,
            citation="",
            inputs_description=_KWARGS_DESCRIPTION,
            features=datasets.Features(
                {
                    "references": datasets.Value("string"),
                }
            ),
            reference_urls=["https://github.com/EleutherAI/lm-evaluation-harness"],
        )

    def _compute(
        self,
        model: str,
        tasks: Union[str, List[str]],
        dmx_config: Optional[str] = None,
        num_fewshot: Optional[int] = None,
        batch_size: Optional[Union[int, str]] = None,
        max_batch_size: Optional[int] = None,
        limit: Optional[Union[int, float]] = None,
        device: Optional[str] = None,
        revision: str = "main",
        trust_remote_code: bool = False, 
        log_samples: bool = True,
        verbosity: str = "INFO",
        **kwargs
    ):
        """
        Evaluate a model on multiple tasks and metrics using lm-eval with optional d-Matrix integration.
        """
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        
        model_args = f"pretrained={model},revision={revision},trust_remote_code={str(trust_remote_code)},device={device}"

        lm = lm_eval.api.registry.get_model("hf").create_from_arg_string(
            model_args,
            {
                "batch_size": batch_size,
                "max_batch_size": max_batch_size,
            }
        )
        
        if dmx_config:
            lm._model = DmxModel.from_torch(lm._model)
            lm._model.transform(lm._model.dmx_config, *eval(f"config_rules.{dmx_config}"))
        
        task_dict = lm_eval.tasks.get_task_dict(tasks if isinstance(tasks, list) else [tasks])
    
        for task in task_dict.values():
            if num_fewshot is not None:
                task.set_config(key="num_fewshot", value=num_fewshot)
        
        eval_params = {
            'lm': lm,
            'task_dict': task_dict,
            'limit': limit,
            'log_samples': log_samples,
            'verbosity': verbosity,
            **kwargs
        }
        
        results = lm_eval.evaluate(**eval_params)
        return results.get('results', {})