import argparse
import gradio as gr
import bittensor as bt
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass
import wandb
import math
import os
import datetime
import time
import json
from dotenv import load_dotenv
from huggingface_hub import HfApi
from apscheduler.schedulers.background import BackgroundScheduler
import pandas as pd 
from tqdm import tqdm
import numpy as np
from substrateinterface import Keypair


load_dotenv()

FONT = (
    """<link href="https://fonts.cdnfonts.com/css/jmh-typewriter" rel="stylesheet">"""
)
TITLE = """<h1 align="center" id="space-title" class="typewriter">Subnet 32 Leaderboard</h1>"""
HEADER = """<h2 align="center" class="typewriter"><a href="https://github.com/It-s-AI/llm-detection" target="_blank">Subnet 32</a> is a <a href="https://bittensor.com/" target="_blank">Bittensor</a> subnet that incentivizes the development of distributed solutions aimed at identifying LLM-generated content. Reward calculation integrates F1 score, False Positive score, and Average Precision score to accurately evaluate model performance.</h3>"""
EVALUATION_DETAILS = """<ul><li><b>UID:</b> the Bittensor UID of the miner</li><li><b>Rewards:</b> result number that is average of 3 metrics below.</li><li><b>F1 Score:</b> f-score metric</li><li><b>FP:</b> False Positive metric</li><li><b>AP:</b> average precision metric</li></ul><br/>More stats on <a href="https://x.taostats.io/subnet/32" target="_blank">taostats</a>."""
EVALUATION_HEADER = """<h3 align="center">Shows the latest internal evaluation statistics as calculated by the "OpenTensor Foundation" validator</h3>"""
VALIDATOR_WANDB_PROJECT = "itsai-dev/subnet32"
H4_TOKEN = os.environ.get("H4_TOKEN", None)
API = HfApi(token=H4_TOKEN)
WANDB_TOKEN = os.environ.get("WANDB_API_KEY", None)
SUBTENSOR_ENDPOINT=os.environ.get("SUBTENSOR_ENDPOINT", None)
REPO_ID = "Infin/ai-detection-leaderboard"
MAX_AVG_LOSS_POINTS = 1
RETRIES = 5
DELAY_SECS = 3
NETUID = 32
UID_MAIN_VALIDATOR = 33
BENCHMARK_TOP_AMOUNT = 3
EVALUATION_STATS_AMOUNT = 5


@dataclass
class ModelData:
    uid: int
    coldkey: str
    hotkey: str
    incentive: float
    emission: float

    @classmethod
    def from_compressed_str(
        cls,
        uid: int,
        coldkey: str,
        hotkey: str,
        incentive: float,
        emission: float,
    ):
        """Returns an instance of this class from a compressed string representation"""
        return ModelData(
            uid=uid,
            coldkey=coldkey,
            hotkey=hotkey,
            incentive=incentive,
            emission=emission,
        )


def run_with_retries(func, *args, **kwargs):
    for i in range(0, RETRIES):
        try:
            return func(*args, **kwargs)
        except (Exception, RuntimeError):
            if i == RETRIES - 1:
                raise
            time.sleep(DELAY_SECS)
    raise RuntimeError("Should never happen")


def get_subtensor_and_metagraph() -> Tuple[bt.subtensor, bt.metagraph]:
    def _internal() -> Tuple[bt.subtensor, bt.metagraph]:
        if SUBTENSOR_ENDPOINT:
            parser = argparse.ArgumentParser()
            bt.subtensor.add_args(parser)            
            subtensor = bt.subtensor(config=bt.config(parser=parser, args=["--subtensor.chain_endpoint", SUBTENSOR_ENDPOINT]))
        else:
            subtensor = bt.subtensor("finney")
        metagraph = subtensor.metagraph(NETUID, lite=False)
        return subtensor, metagraph

    return run_with_retries(_internal)


def get_validator_weights(
    metagraph: bt.metagraph,
) -> Dict[int, Tuple[float, int, Dict[int, float]]]:
    """Returns a dictionary of validator UIDs to (vtrust, stake, {uid: weight})."""
    ret = {}
    for uid in metagraph.uids.tolist():
        vtrust = metagraph.validator_trust[uid].item()
        if vtrust > 0:
            ret[uid] = (vtrust, metagraph.S[uid].item(), {})
            for ouid in metagraph.uids.tolist():
                if ouid == uid:
                    continue
                weight = round(metagraph.weights[uid][ouid].item(), 6)
                if weight > 0:
                    ret[uid][-1][ouid] = weight
    return ret


def get_subnet_data(
    metagraph: bt.metagraph
) -> List[ModelData]:
    result = []
    for uid in tqdm(metagraph.uids.tolist()):
        if metagraph.validator_trust[uid] != 0:
            continue

        coldkey = metagraph.coldkeys[uid]
        hotkey = metagraph.hotkeys[uid]
        incentive = metagraph.incentive[uid]
        emission = (
            metagraph.emission[uid] * 20
        )  # convert to daily TAO

        model_data = None
        try:
            model_data = ModelData.from_compressed_str(
                uid, coldkey, hotkey, incentive, emission
            )
        except:
            continue

        result.append(model_data)
    return result


def is_floatable(x) -> bool:
    return (
        isinstance(x, float) and not math.isnan(x) and not math.isinf(x)
    ) or isinstance(x, int)


def get_wandb_runs(
        project: str, filters: Dict[str, Any]
    ) -> List:
    """Get the latest runs from Wandb, retrying infinitely until we get them."""
    while True:
        api = wandb.Api(api_key=WANDB_TOKEN)

        runs = list(
            api.runs(
                project,
                order="-created_at",
                filters=filters,
            )
        )

        print('Runs amount: ', len(runs))
        if len(runs) > 0:
            return runs
        # WandDB API is quite unreliable. Wait another minute and try again.
        print("Failed to get runs from Wandb. Trying again in 60 seconds.")
        time.sleep(10)


def is_in_last_7_days(current_time, unix_timestamp):
    seven_days_ago = current_time - 7 * 24 * 60 * 60
    return seven_days_ago <= unix_timestamp <= current_time


def is_hash_repeated(current_hash: str, current_index: int, wandb_runs: List):
    for run in wandb_runs[current_index+1:]:
        if 'signed_msg' not in run.summary:
            continue

        if run.summary['signed_msg'] == current_hash:
            return True
    return False


def get_scores(
    wandb_runs: List, hotkeys: List[str]
) -> Dict[int, Dict[str, Optional[float]]]:
    result = {}
    # result = []
    previous_timestamp = None
    # Iterate through the runs until we've processed all the uids.
    
    current_time = time.time()
    for i, run in enumerate(wandb_runs):
        config_data = run.config
        vali_uid = config_data['uid']

        # get only last run of each vali except OTF
        if int(vali_uid) == UID_MAIN_VALIDATOR:
            if len(result.get(vali_uid, [])) >= 5:
                continue
        else:
            if vali_uid in result:
                continue

        if run.history().empty:
            continue
        
        data = json.loads(run.summary["original_format_json"])
        if vali_uid > len(hotkeys):
            print('VALI UID EXCEEDS NUMBER OF HOTKEYS')
            continue

        if not is_in_last_7_days(current_time, int(data['timestamp'])):
            print("TIMESTAMP NOT IN LAST 7 DAYS", data['timestamp'])
            continue

        keypair = Keypair(ss58_address=hotkeys[vali_uid])

        if 'signed_msg' not in run.summary:
            continue

        s = time.time()
        if is_hash_repeated(run.summary['signed_msg'], i, wandb_runs):
            print("THIS HASH ALREADY BEEN SEEN IN RUNS", run.summary['signed_msg'])
            continue
        print("TIME CONSUMED FOR HASH CHECKING", int(time.time() - s))

        verify_result = keypair.verify(run.summary["original_format_json"], run.summary['signed_msg'])
        print(vali_uid, hotkeys[vali_uid])
        print('Signature is correct: ', verify_result)
        if not verify_result:
            print("SIGNATURE IS BROKEN")
            continue

        timestamp = data["timestamp"]

        # Make sure runs are indeed in descending time order.
        # assert (
        #     previous_timestamp is None or timestamp < previous_timestamp
        # ), f"Timestamps are not in descending order: {timestamp} >= {previous_timestamp}"
        previous_timestamp = timestamp

        if vali_uid not in result:
            result[vali_uid] = []
        local_vali_uid_iter = {}
        for miner_data in list(data['uid_metrics'].values()):
            local_vali_uid_iter[miner_data['uid']] = {}
            local_vali_uid_iter[miner_data['uid']].update(miner_data)
        result[vali_uid].append(local_vali_uid_iter)
    return result


def average_scores(data: List):
    stats = {}

    for item in data:
        for key, values in item.items():
            if key not in stats:
                stats[key] = {'sums': {'reward': 0, 'fp_score': 0, 'f1_score': 0, 'ap_score': 0, 'penalty': 0}, 'count': 0}
                stats[key]['uid'] = values['uid']
                stats[key]['weight'] = values['weight']

            stats[key]['sums']['reward'] += values['reward']
            stats[key]['sums']['fp_score'] += values['fp_score']
            stats[key]['sums']['f1_score'] += values['f1_score']
            stats[key]['sums']['ap_score'] += values['ap_score']
            stats[key]['sums']['penalty'] += values['penalty']
            stats[key]['count'] += 1


    averages = {}
    for key, data in stats.items():
        averages[key] = {field: data['sums'][field] / data['count'] for field in data['sums']}
        averages[key] = {**averages[key], 'uid': data['uid'], 'weight': data['weight']}
    return averages


def format_score(uid: int, scores, key) -> Optional[float]:
    if uid in scores:
        if key in scores[uid]:
            point = scores[uid][key]
            if is_floatable(point):
                return round(scores[uid][key], 6)
    return None


def next_epoch(subtensor: bt.subtensor, block: int) -> int:
    return (
        block
        + subtensor.get_subnet_hyperparameters(NETUID).tempo
        - subtensor.blocks_since_epoch(NETUID, block)
    )


def get_last_updated_div() -> str:
    return f"""<div>Last Updated: {datetime.datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")} (UTC)</div>"""


def leaderboard_data(
    leaderboard: List[ModelData],
    scores: Dict[int, Dict[str, Optional[float]]],
    show_stale: bool,
) -> List[List[Any]]:
    """Returns the leaderboard data, based on models data and UID scores."""
    # headers=["Top", "UID", "Reward", "F1 Score", "FP Score", "AP Score"],

    rows = [
        [
                f"{c.coldkey[:12]}",
                c.uid,
                format_score(c.uid, scores, "reward"),
                format_score(c.uid, scores, "f1_score"),
                format_score(c.uid, scores, "fp_score"),
                format_score(c.uid, scores, "ap_score"),
        ] for c in leaderboard if c.uid in scores
    ]
   
    sorted_rows = sorted(rows, key=lambda x: (x[2], x[0]), reverse=True)
    if not show_stale:
        sorted_rows = sorted_rows[:EVALUATION_STATS_AMOUNT]
    return sorted_rows


def restart_space():
    API.restart_space(repo_id=REPO_ID, token=H4_TOKEN)


def main():
    # To avoid leaderboard failures, infinitely try until we get all data
    # needed to populate the dashboard
    while True:
        # try:
            subtensor, metagraph = get_subtensor_and_metagraph()

            model_data: List[ModelData] = get_subnet_data(metagraph)
            model_data.sort(key=lambda x: x.incentive, reverse=False)
            time_now = datetime.datetime.now(datetime.timezone.utc)
            n_days_ago = time_now - datetime.timedelta(hours=24*7)
 
            vali_runs = get_wandb_runs(project=VALIDATOR_WANDB_PROJECT, filters={
                "$and": [
                    {
                        'created_at': {
                            '$gte': n_days_ago.isoformat()
                        }
                    },
                    {
                        'state': {
                            "$in": ["finished"]
                        }
                    },
                    # {
                    #     'config.version': {
                    #         '$in': ["2.5.0", "2.6.0", "3.0.0", "3.0.1"]
                    # #     }
                    # },
                    {
                        'config.uid': {
                            '$nin': [86, 50, 106]
                        }
                    }
                ]
            })   
            s = time.time() 
            scores = get_scores(vali_runs, metagraph.hotkeys)
            print('TIME FOR get_scores()', time.time() - s)
            averaged_scores = {}
            for vali_uid, vali_score in scores.items():
                if len(vali_score) > 1:
                    averaged_scores.update({vali_uid: average_scores(vali_score)})
                else:
                    averaged_scores.update({vali_uid: vali_score[0]})

            scores = averaged_scores

            rows = []
            for validator_uid, miners in scores.items():                 
                for miner_uid, stats in miners.items():
                    row = {'validator_uid': validator_uid}
                    row.update({k: v for k, v in stats.items()})
                    rows.append(row)

            miners_stats = pd.DataFrame(rows)

            miners_stats['stake_value'] = miners_stats['validator_uid'].apply(lambda x: metagraph.S[x].item())
            miners_stats = miners_stats.sort_values(by='stake_value', ascending=False)
            miners_stats = miners_stats.drop(columns='stake_value')
            miners_stats.to_csv('miners_stats.csv', index=False)

            main_validator_scores = scores[UID_MAIN_VALIDATOR]

            sorted_main_validator_scores = sorted(main_validator_scores.items(), key=lambda x: x[1]['reward'], reverse=True)

            axons_info = metagraph.axons
            top_miners = []
            top_keys = []
            for score_i in sorted_main_validator_scores:
                if len(top_miners) >= BENCHMARK_TOP_AMOUNT:
                    break

                coldkey_i = axons_info[score_i[0]].coldkey
                if coldkey_i in top_keys:
                    continue

                top_keys.append(coldkey_i)
                top_miners.append(score_i)
                
            top_miner_score = dict(top_miners)
            data_list = [{'Model': f'miner_{key}', 
                        'Average': value['reward'],
                        'F1 score': value['f1_score'], 
                        'FP score': value['fp_score'], 
                        'AP score': value['ap_score']} for key, value in top_miner_score.items()]

            df = pd.DataFrame(data_list)
            baseline_data = [
                {'Model': 'baseline: deberta', 'F1 score': 0.863, 'FP score': 0.896, 'AP score': 0.789, 'Average': 0.849}
            ]

            benchmarks = pd.concat([df, pd.DataFrame(baseline_data)], ignore_index=True)
            benchmarks = benchmarks.sort_values(by='Average', ascending=False)

            validator_df = get_validator_weights(metagraph)
            break
        # except Exception as e:
        #     print(f"Failed to get data: {e}")
        #     time.sleep(30)            

    demo = gr.Blocks(css=".typewriter {font-family: 'JMH Typewriter', sans-serif;}")
    with demo:
        gr.HTML(FONT)
        gr.HTML(TITLE)
        gr.HTML(HEADER)

        if benchmarks is not None:
            with gr.Accordion("Top Model Benchmarks"):
                gr.components.Dataframe(benchmarks)
                gr.HTML("""<div>Fore more information about baseline miner architecture see <a href='https://github.com/It-s-AI/llm-detection/blob/main/neurons/miner.py'>here</a> for the full code.</div>""")

        with gr.Accordion("Evaluation Stats"):
            gr.HTML(EVALUATION_HEADER)
            show_stale = gr.Checkbox(label="Show All Miners", interactive=True)
            leaderboard_table = gr.components.Dataframe(
                value=leaderboard_data(model_data, main_validator_scores, show_stale.value),
                headers=["Coldkey", "UID", "Reward", "F1 Score", "FP Score", "AP Score"],
                datatype=["markdown", "number", "number", "number", "number", "number"],
                elem_id="leaderboard-table",
                interactive=False,
                visible=True,
            )

            gr.HTML(EVALUATION_DETAILS)
            show_stale.change(
                lambda stale: leaderboard_data(model_data, main_validator_scores, stale),
                inputs=[show_stale],
                outputs=leaderboard_table,
            )

        with gr.Accordion("Validator Stats"):
            model_data.sort(key=lambda x: x.uid, reverse=False)

            values = [
                    [uid, int(validator_df[uid][1]), round(validator_df[uid][0], 6)]
                    + [
                        validator_df[uid][-1].get(c.uid)
                        for c in model_data
                        if c.incentive
                    ]
                    for uid, _ in sorted(
                        zip(
                            validator_df.keys(),
                            [validator_df[x][1] for x in validator_df.keys()],
                        ),
                        key=lambda x: x[1],
                        reverse=True,
                    )
                ]
            
            print("VALUES:", values)
            averages = np.nanmean(np.array(values, dtype=float)[:, 3:], axis=0)
            averages = np.around(averages, decimals=6)

            values.append(["Average", 0, 0] + averages.tolist())

            gr.components.Dataframe(
                value=values,
                headers=["UID", "Stake (τ)", "V-Trust"]
                + [
                    f"{c.uid}/reward"
                    for c in model_data
                    if c.incentive
                ]
                ,
                datatype=["markdown", "number", "number"]
                + ["number" for c in model_data if c.incentive]
                ,
                interactive=False,
                visible=True,
            )

        def get_miner_stats(uid):
            local_stats = miners_stats[miners_stats['uid'] == int(uid)]
            local_stats.drop('penalty', axis=1, inplace=True)
            local_stats = local_stats[["validator_uid","uid","weight", "reward","fp_score","f1_score","ap_score"]]
            local_stats.columns = ["Validator UID", "UID", "Weight", "Reward", "FP Score", "F1 Score", "AP Score"]
            # local_stats = local_stats[["Validator UID", "UID", "Weight", "Reward", "FP Score", "F1 Score", "AP Score"]]
            return local_stats

        with gr.Accordion("Get your miner stats"):
            with gr.Row():
                input_text = gr.Textbox(label="Enter your Miner ID")
                submit_button = gr.Button("Get Stats")

            output_df = gr.components.DataFrame(
                headers=["Validator UID", "Miner UID"] + ["Weight", "Reward"] + ["FP Score", "F1 Score", "AP Score"],
                datatype=["number", "number"] + ["number", "number", "number", "number", "number"],
                interactive=False,
                visible=True
            )

            submit_button.click(
                fn=get_miner_stats,
                inputs=input_text,
                outputs=output_df
            )

        gr.HTML(value=get_last_updated_div())

    scheduler = BackgroundScheduler()
    scheduler.add_job(
        restart_space, "interval", seconds=60 * 60 * 24
    )  # restart every 45 minutes
    scheduler.start()

    demo.launch()

main()