import json
import os
import logging
from datetime import datetime

from lm_eval import tasks, evaluator, utils
from lm_eval.tasks import TaskManager

from src.envs import RESULTS_REPO, API
from src.backend.manage_requests import EvalRequest
from src.logging import setup_logger

import numpy as np

logging.getLogger("openai").setLevel(logging.WARNING)
logger = setup_logger(__name__)

class NumpyArrayEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        else:
            return super().default(obj)

def run_evaluation(eval_request: EvalRequest, task_names: list, num_fewshot: int, batch_size: int or str, device: str, local_dir: str, results_repo: str, limit: int =None):
    """Runs one evaluation for the current evaluation request file, then pushes the results to the hub.

    Args:
        eval_request (EvalRequest): Input evaluation request file representation
        task_names (list): Tasks to launch
        num_fewshot (int): Number of few shots to use
        batch_size (int): Selected batch size
        device (str): "cpu" or "gpu:0", depending on what you assigned to the space
        local_dir (str): Where to save the results locally
        results_repo (str): To which repository to upload the results
        no_cache (bool, optional): Whether to use a cache or not.
        limit (int, optional): Whether to use a number of samples only for the evaluation - only for debugging

    Returns:
        _type_: _description_
    """
    if limit:
        logger.info(
            "WARNING: --limit SHOULD ONLY BE USED FOR TESTING. REAL METRICS SHOULD NOT BE COMPUTED USING LIMIT."
        )

    task_manager = TaskManager()
    all_tasks = task_manager.all_tasks
        
    task_names = utils.pattern_match(task_names, all_tasks)

    logger.info(f"Selected Tasks: {task_names}")

    print("Here are the deets we are putting in:")
    print("model args:")
    model_args = eval_request.get_model_args()
    print(model_args)
    print("Tasks")
    print(task_names)
    print("Num fewshot")
    print(num_fewshot)
    print("Batch size")
    print(batch_size)
    print("Device")
    print(device)
    print("Limit")
    print(limit)
    # no_cache=no_cache,
    # output_base_path="logs"
    #gen_kwargs="max_new_tokens=None"
    results = evaluator.simple_evaluate(
        model="hf", #= "hf-causal-experimental", # "hf-causal"
        model_args=model_args,
        tasks=task_names,
        num_fewshot=num_fewshot,
        batch_size=batch_size,
        device=device,
        limit=limit,
        write_out=True,
        gen_kwargs="max_new_tokens=400"
    )

    # TODO: For Toxicity, substract from 100.
    #     results["realtoxicityprompts"]["score"] = 100 - results["realtoxicityprompts"]["score"]
    results["config"]["model_dtype"] = eval_request.precision
    results["config"]["model_name"] = eval_request.model
    results["config"]["model_sha"] = eval_request.revision

    print("Results is")
    print(results)

    dumped = json.dumps(results, cls=NumpyArrayEncoder, indent=2)
    logger.info(dumped)

    output_path = os.path.join(local_dir, *eval_request.model.split("/"), f"results_{datetime.now()}.json")
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    with open(output_path, "w") as f:
        f.write(dumped)

    logger.info(utils.make_table(results))

    print("Uploading to")
    print(output_path)
    print("repo id")
    print(results_repo)

    API.upload_file(
        path_or_fileobj=output_path,
        path_in_repo=f"{eval_request.model}/results_{datetime.now()}.json",
        repo_id=results_repo,
        repo_type="dataset",
    )

    return results