from datetime import datetime
import numpy as np
import json
import re
import heapq
from collections import defaultdict
import tempfile
from typing import Dict, Tuple, List, Literal
import gradio as gr
from datatrove.utils.stats import MetricStatsDict

from src.logic.graph_settings import Grouping

PARTITION_OPTIONS = Literal["Top", "Bottom", "Most frequent (n_docs)"]

def prepare_for_non_grouped_plotting(metric: Dict[str, MetricStatsDict], normalization: bool, rounding: int) -> Dict[float, float]:
    keys = np.array([float(key) for key in metric.keys()])
    values = np.array([value.total for value in metric.values()])
    
    rounded_keys = np.round(keys, rounding)
    unique_keys, indices = np.unique(rounded_keys, return_inverse=True)
    metrics_rounded = np.zeros_like(unique_keys, dtype=float)
    np.add.at(metrics_rounded, indices, values)
    
    if normalization:
        normalizer = np.sum(metrics_rounded)
        metrics_rounded /= normalizer
    
    return dict(zip(unique_keys, metrics_rounded))

def prepare_for_group_plotting(metric: Dict[str, MetricStatsDict], top_k: int, direction: PARTITION_OPTIONS, regex: str | None, rounding: int) -> Tuple[List[str], List[float], List[float]]:
    regex_compiled = re.compile(regex) if regex else None
    filtered_metric = {key: value for key, value in metric.items() if not regex or regex_compiled.match(key)}
    
    keys = np.array(list(filtered_metric.keys()))
    means = np.array([float(value.mean) for value in filtered_metric.values()])
    stds = np.array([value.standard_deviation for value in filtered_metric.values()])
    
    rounded_means = np.round(means, rounding)
    
    if direction == "Top":
        top_indices = np.argsort(rounded_means)[-top_k:][::-1]
    elif direction == "Most frequent (n_docs)":
        totals = np.array([int(value.n) for value in filtered_metric.values()])
        top_indices = np.argsort(totals)[-top_k:][::-1]
    else:
        top_indices = np.argsort(rounded_means)[:top_k]
    
    top_keys = keys[top_indices]
    top_means = rounded_means[top_indices]
    top_stds = stds[top_indices]
    
    return top_keys.tolist(), top_means.tolist(), top_stds.tolist()

def export_data(exported_data: Dict[str, MetricStatsDict], metric_name: str, grouping: Grouping):
    if not exported_data:
        return None

    file_name = f"{metric_name}_{grouping}_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.json"
    with open(file_name, 'w') as f:
        json.dump({
            name: sorted([{"value": key, **value} for key, value in dt.to_dict().items()], key=lambda x: x["value"])
            for name, dt in exported_data.items()
        }, f, indent=2)
    return gr.File(value=file_name, visible=True)