from datetime import timedelta

import plotly.express as px

import src.constants as constants
from src.hub import load_model_card
from src.requests import load_request


async def get_env_impact(data):
    total_evaluation_time_seconds = data.get("total_evaluation_time_seconds")
    if total_evaluation_time_seconds:
        total_evaluation_time_seconds = float(total_evaluation_time_seconds)
    env_impact = {
        "co2_emissions": calculate_co2_emissions(total_evaluation_time_seconds),
        "total_evaluation_time": str(timedelta(seconds=total_evaluation_time_seconds)),
        "num_parameters_billions": data.get("config", {}).get("model_num_parameters") / 10**9,
        "precision": data.get("config", {}).get("model_dtype"),
    }
    request = await load_request(data["model_name"], env_impact["precision"])
    if request:
        model_type_label = request.get("model_type", "unknown")
        env_impact["model_type"] = constants.MODEL_TYPE_LABEL_TO_TYPE.get(model_type_label, model_type_label)
        env_impact["architecture"] = request.get("architectures", "Unknown")
    # MoE
    model_card = await load_model_card(data["model_name"])
    model_tags = get_moe_model_tags(model_card.data, data["model_name"])
    moe = "moe" in model_tags or "moe" in data["model_name"].lower()
    env_impact["moe"] = moe
    return env_impact


# Source: https://huggingface.co/docs/leaderboards/open_llm_leaderboard/emissions#function-for-c02-calculation
def calculate_co2_emissions(total_evaluation_time_seconds: float | None) -> float:
    if total_evaluation_time_seconds is None or total_evaluation_time_seconds <= 0:
        return -1
    # Power consumption for 8 H100 SXM GPUs in kilowatts (kW)
    power_consumption_kW = 5.6
    # Carbon intensity in grams CO₂ per kWh in Virginia
    carbon_intensity_g_per_kWh = 269.8
    # Convert evaluation time to hours
    total_evaluation_time_hours = total_evaluation_time_seconds / 3600
    # Calculate energy consumption in kWh
    energy_consumption_kWh = power_consumption_kW * total_evaluation_time_hours
    # Calculate CO₂ emissions in grams
    co2_emissions_g = energy_consumption_kWh * carbon_intensity_g_per_kWh
    # Convert grams to kilograms
    return co2_emissions_g / 1000


# Source: https://huggingface.co/spaces/open-llm-leaderboard/open_llm_leaderboard_parser/blob/main/src/submission/check_validity.py#L33
def get_moe_model_tags(model_card, model_id):
    # is_merge_from_metadata = False
    is_moe_from_metadata = False
    is_moe_from_model_card = False
    # is_merge_from_model_card = False
    tags = []
    if model_card is None:
        return tags
    if model_card.tags:
        # is_merge_from_metadata = any(tag in model_card.tags for tag in ["merge", "moerge", "mergekit", "lazymergekit"])
        is_moe_from_metadata = any(tag in model_card.tags for tag in ["moe", "moerge", "mixtral"])
    if model_card.get("text", False):
        # is_merge_from_model_card = any(
        #     keyword in model_card.text.lower() for keyword in ["merged model", "merge model", "moerge"]
        # )
        is_moe_from_model_card = any(keyword in model_card.text.lower() for keyword in ["moe", "mixtral"])
    # if is_merge_from_model_card or is_merge_from_metadata:
    #     tags.append("merge")
    is_moe_from_name = any(
        key in model_id.lower().replace("/", "-").replace("_", "-").split("-") for key in ["moe", "mixtral"]
    )
    # Hardcoded check for "rhymes-ai/Aria" model
    if model_id == "rhymes-ai/Aria":
        tags.append("moe")
    elif is_moe_from_model_card or is_moe_from_name or is_moe_from_metadata:
        tags.append("moe")
    return tags


def plot_env_impact(df):
    if df is None:
        return None, None
    fig_1 = px.scatter(
        df.rename_axis(index="Model").reset_index(),
        x="env_impact.num_parameters_billions",
        y="env_impact.co2_emissions",
        color="Model",
        title="Evaluation CO₂ Emissions (kg) vs. #Params (B)",
        labels={
            "env_impact.num_parameters_billions": "#Params (B)",
            "env_impact.co2_emissions": "Evaluation CO₂ Emissions (kg)",
        },
        color_discrete_sequence=px.colors.qualitative.Safe,  # TODO: https://plotly.com/python/discrete-color/
    )
    fig_2 = px.scatter(
        df.rename_axis(index="Model").reset_index(),
        x="results.leaderboard.acc_norm,none",
        y="env_impact.co2_emissions",
        color="Model",
        title="Evaluation CO₂ Emissions (kg) vs. Score",
        labels={
            "results.leaderboard.acc_norm,none": "Mean Score",
            "env_impact.co2_emissions": "Evaluation CO₂ Emissions (kg)",
        },
        color_discrete_sequence=px.colors.qualitative.Safe,  # TODO: https://plotly.com/python/discrete-color/
    )
    fig_2.update_xaxes(range=[0, 1])
    return fig_1, fig_2