import json
import os
from dataclasses import dataclass

import numpy as np

from src.display.utils import Tasks


@dataclass
class EvalResult:
    """Represents one full evaluation. Built from a combination of the result and request file for a given run.
    """
    model_name: str
    org: str
    results: dict
    model_type: str # Pretrained, fine tuned, ...
    num_params: int = 0
    date: str = "" # submission date of request file
    still_on_hub: bool = False

    @classmethod
    def init_from_json_file(self, json_filepath):
        """Inits the result from the specific model result file"""
        with open(json_filepath) as fp:
            data = json.load(fp)

        config = data.get("config")

        num_params = config.get("params")
        org = config.get("developer", "Unknown")
        model_type = config.get("model_type", "Unknown")
        model_name = config.get("model_name")

        # Extract results available in this file (some results are split in several files)
        results = {}
        for task in Tasks:
            task = task.value

            # We average all scores of a given metric (not all metrics are present in all files)
            accs = np.array([v.get(task.metric, None) for k, v in data["results"].items() if task.benchmark == k])
            if accs.size == 0 or any([acc is None for acc in accs]):
                continue

            mean_acc = np.mean(accs)
            results[task.benchmark] = mean_acc

        return self(
            model_name=model_name,
            org=org,
            results=results,
            num_params=num_params,
            model_type=model_type
        )
    
    def to_dict(self):
        """Converts the Eval Result to a dict compatible with our dataframe display"""
        data_dict = {
            "Model": self.model_name,
            "Type": self.model_type,
            "#Params (B)": self.num_params
        }

        for task in Tasks:
            data_dict[task.value.col_name] = self.results[task.value.benchmark]

        return data_dict

def get_raw_eval_results(results_path: str) -> list[EvalResult]:
    """From the path of the results folder root, extract all needed info for results"""
    model_result_filepaths = []

    for fn in os.listdir(results_path):
        # We should only have json files in model results
        if fn.endswith(".json"):
            model_result_filepaths.append(os.path.join(results_path, fn))

    eval_results = {}
    for model_result_filepath in model_result_filepaths:
        # Creation of result
        eval_result = EvalResult.init_from_json_file(model_result_filepath)
        # Store results of same eval together
        model_name = eval_result.model_name
        eval_results[model_name] = eval_result

    results = []
    for v in eval_results.values():
        try:
            v.to_dict() # we test if the dict version is complete
            results.append(v)
        except KeyError:  # not all eval values present
            continue

    return results