import json
import numpy as np
from pathlib import Path
from collections import Counter

def compute_token_frequencies(tokens, top_n=10000):
    if len(tokens) == 0:
        raise ValueError("Cannot compute distribution on empty token list")
    
    counter = Counter(tokens)
    most_common = counter.most_common(top_n)
    
    freq_dict = {token: count for token, count in most_common}
    return freq_dict

def save_token_frequencies(freq_dict, version_dir):
    version_dir = Path(version_dir)
    version_dir.mkdir(parents=True, exist_ok=True)
    
    freq_path = version_dir / "token_freq.json"
    freq_path.write_text(json.dumps(freq_dict, indent=2), encoding="utf-8")
    
    total_unique = len(freq_dict)
    total_tokens = sum(freq_dict.values())
    
    top_10 = []
    for i, (token, count) in enumerate(freq_dict.items()):
        if i >= 10:
            break
        pct = (count / total_tokens) * 100 if total_tokens > 0 else 0
        top_10.append({
            "token": token,
            "count": count,
            "percentage": round(pct, 4)
        })
        
    summary = {
        "total_unique_tokens": total_unique,
        "total_token_count": total_tokens,
        "top_10_tokens": top_10
    }
    
    summary_path = version_dir / "token_freq_summary.json"
    summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")

def load_token_frequencies(version_dir):
    version_dir = Path(version_dir)
    freq_path = version_dir / "token_freq.json"
    
    if not freq_path.exists():
        version_id = version_dir.name
        raise FileNotFoundError(f"token_freq.json not found for version {version_id} — was it created with an older datavault?")
        
    return json.loads(freq_path.read_text(encoding="utf-8"))

def kl_divergence(p_counts, q_counts):
    vocab = set(p_counts.keys()) | set(q_counts.keys())
    vocab_list = list(vocab)
    
    p_arr = np.array([p_counts.get(w, 0) for w in vocab_list], dtype=np.float64)
    q_arr = np.array([q_counts.get(w, 0) for w in vocab_list], dtype=np.float64)
    
    epsilon = 1e-10
    p_arr += epsilon
    q_arr += epsilon
    
    p_arr /= p_arr.sum()
    q_arr /= q_arr.sum()
    
    kl_pq = np.sum(p_arr * np.log(p_arr / q_arr))
    kl_qp = np.sum(q_arr * np.log(q_arr / p_arr))
    sym_kl = (kl_pq + kl_qp) / 2
    
    return {
        "kl_pq": round(float(kl_pq), 4),
        "kl_qp": round(float(kl_qp), 4),
        "symmetric_kl": round(float(sym_kl), 4)
    }

def get_distribution_shift(p_counts, q_counts, top_n=15):
    vocab = set(p_counts.keys()) | set(q_counts.keys())
    
    p_total = sum(p_counts.values()) or 1
    q_total = sum(q_counts.values()) or 1
    
    shifts = []
    for token in vocab:
        prob_p = p_counts.get(token, 0) / p_total
        prob_q = q_counts.get(token, 0) / q_total
        shift = prob_q - prob_p
        shifts.append({
            "token": token,
            "prob_v1": round(prob_p, 4),
            "prob_v2": round(prob_q, 4),
            "shift": round(shift, 4),
            "abs_shift": abs(shift)
        })
        
    shifts.sort(key=lambda x: x["abs_shift"], reverse=True)
    
    gained = []
    lost = []
    
    for item in shifts:
        out = {
            "token": item["token"],
            "prob_v1": item["prob_v1"],
            "prob_v2": item["prob_v2"],
            "shift": item["shift"]
        }
        if item["shift"] > 0 and len(gained) < top_n:
            gained.append(out)
        elif item["shift"] < 0 and len(lost) < top_n:
            lost.append(out)
            
        if len(gained) >= top_n and len(lost) >= top_n:
            break

    return {
        "gained": gained,
        "lost": lost
    }

def generate_distribution_report(v1_dir, v2_dir):
    p_counts = load_token_frequencies(v1_dir)
    q_counts = load_token_frequencies(v2_dir)
    
    kl_res = kl_divergence(p_counts, q_counts)
    sym_kl = kl_res["symmetric_kl"]
    
    if sym_kl < 0.1:
        interp = "LOW — distributions nearly identical"
    elif sym_kl <= 0.5:
        interp = "MEDIUM — moderate distribution shift"
    elif sym_kl <= 1.0:
        interp = "HIGH — distributions meaningfully different"
    else:
        interp = "CRITICAL — distributions are fundamentally different"
        
    kl_res["interpretation"] = interp
    
    shift_res = get_distribution_shift(p_counts, q_counts)
    
    def get_top_10(counts):
        return [{"token": k, "count": v} for k, v in list(counts.items())[:10]]
        
    report = {
        "kl_divergence": kl_res,
        "top_gained": shift_res["gained"],
        "top_lost": shift_res["lost"],
        "v1_top10": get_top_10(p_counts),
        "v2_top10": get_top_10(q_counts)
    }
    
    p_total = sum(p_counts.values())
    q_total = sum(q_counts.values())
    if p_total < 100 or q_total < 100:
        report["warning"] = "Corpus too small for reliable distribution analysis (< 100 tokens)"
        
    return report
