import gradio as gr
import pandas as pd

from chain_data import WEIGHTS_BY_MINER, get_neurons, sync_chain, Weight
from wandb_data import Key, get_current_runs


def get_color_by_weight(weight: float) -> str:
    if weight < 0.001:
        return "gray"
    elif weight < 0.3:
        r = int(255)
        g = int((weight / 0.3) * 165)
        return f"rgb({r}, {g}, 0)"
    elif weight < 0.8:
        progress = (weight - 0.3) / 0.5
        r = int(255 - (progress * 255))
        g = int(165 + (progress * 90))
        return f"rgb({r}, {g}, 0)"
    else:
        progress = (weight - 0.8) / 0.2
        g = int(255 - ((1 - progress) * 50))
        return f"rgb(0, {g}, 0)"


def get_active_weights() -> dict[Key, list[tuple[Key, Weight]]]:
    runs = get_current_runs()
    weights: dict[Key, list[tuple[Key, Weight]]] = {}

    for hotkey, validator_weights in WEIGHTS_BY_MINER.items():
        new_weights: list[tuple[Key, Weight]] = []
        for validator_hotkey, weight in validator_weights:
            if validator_hotkey in [run.hotkey for run in runs]:
                new_weights.append((validator_hotkey, weight))
        weights[hotkey] = new_weights

    return weights


def create_weights(include_inactive: bool) -> gr.Dataframe:
    data: list[list] = []
    sync_chain()

    headers = ["Miner UID", "Incentive"]
    datatype = ["number", "markdown"]

    weights = WEIGHTS_BY_MINER if include_inactive else get_active_weights()

    neurons = get_neurons()
    validator_uids = set()
    for _, validator_weights in weights.items():
        for hotkey, _ in validator_weights:
            validator_uids.add(neurons[hotkey].uid)

    for validator_uid in sorted(validator_uids):
        headers.append(str(validator_uid))
        datatype.append("markdown")

    for hotkey, validator_weights in weights.items():
        if not hotkey in neurons:
            continue
        incentive = neurons[hotkey].incentive
        row = [neurons[hotkey].uid, f"<span style='color: {get_color_by_weight(incentive)}'>{incentive:.{3}f}</span>"]
        for _, weight in validator_weights:
            row.append(f"<span style='color: {get_color_by_weight(weight)}'>{weight:.{3}f}</span>")
        data.append(row)

    data.sort(key=lambda val: float(val[1].split(">")[1].split("<")[0]), reverse=True)

    return gr.Dataframe(
        pd.DataFrame(data, columns=headers),
        datatype=datatype,
        interactive=False,
        max_height=800,
    )