"""
Audit litdata_3b and litdata_english for training quality.

Decodes random samples from binary chunks back to text, then checks:
  1. English alignment (ASCII ratio, language detection heuristics)
  2. Cleaning quality (no junk, code blocks, HTML, boilerplate)
  3. Deduplication (MinHash-like shingle overlap between samples)
  4. Data diversity (topic spread, length distribution, vocab richness)
  5. Random sample printout for manual inspection
"""

import json
import os
import random
import re
import struct
import sys
from collections import Counter, defaultdict
from pathlib import Path

import numpy as np
from tokenizers import Tokenizer

ROOT = Path(__file__).resolve().parent.parent.parent
BLOCK_SIZE = 1025
DTYPE = np.int32

# ── Load tokenizer ──────────────────────────────────────────────
tokenizer = Tokenizer.from_file(
    str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
)


def load_litdata_index(litdata_dir):
    with open(litdata_dir / "index.json") as f:
        return json.load(f)


def read_blocks_from_chunk(chunk_path, num_blocks_to_read, total_blocks):
    """Read random blocks from a chunk file, decoding the litdata header."""
    dtype_size = DTYPE().itemsize
    # Header: uint32 num_items + uint32 offsets[0..num_items]
    header_ints = 1 + total_blocks + 1  # num_items + (num_items+1) offsets
    header_bytes = header_ints * 4
    
    with open(chunk_path, "rb") as f:
        # Read header
        header_raw = f.read(header_bytes)
        header = np.frombuffer(header_raw, dtype=np.uint32)
        n_items = header[0]
        
        # Pick random block indices
        indices = random.sample(range(min(n_items, total_blocks)), min(num_blocks_to_read, n_items))
        
        blocks = []
        for idx in indices:
            offset = int(header[1 + idx])  # byte offset within data section
            f.seek(header_bytes + offset)
            block_data = np.frombuffer(f.read(BLOCK_SIZE * dtype_size), dtype=DTYPE)
            blocks.append(block_data)
    
    return blocks


def decode_block(token_ids):
    """Decode a block of token IDs back to text."""
    ids = token_ids.tolist()
    return tokenizer.decode(ids, skip_special_tokens=False)


# ── Quality checks ──────────────────────────────────────────────
def ascii_ratio(text):
    if not text:
        return 0
    ascii_chars = sum(1 for c in text if ord(c) < 128)
    return ascii_chars / len(text)


def alpha_ratio(text):
    if not text:
        return 0
    alpha = sum(1 for c in text if c.isalpha())
    return alpha / len(text)


def has_html(text):
    return bool(re.search(r'<(html|body|div|span|script|style|head|meta|link)\b', text, re.I))


def has_boilerplate(text):
    patterns = [
        r'cookie policy', r'terms of service', r'privacy policy',
        r'all rights reserved', r'click here', r'subscribe now',
        r'sign up for', r'©\s*\d{4}', r'powered by',
        r'advertisement', r'sponsored content',
    ]
    lower = text.lower()
    return sum(1 for p in patterns if re.search(p, lower))


def url_ratio(text):
    urls = re.findall(r'https?://\S+', text)
    url_chars = sum(len(u) for u in urls)
    return url_chars / max(len(text), 1)


def sentence_count(text):
    return len(re.findall(r'[.!?]+\s', text))


def word_count(text):
    return len(text.split())


def unique_word_ratio(text):
    words = text.lower().split()
    if not words:
        return 0
    return len(set(words)) / len(words)


def shingle_set(text, k=5):
    """Create k-word shingles for dedup checking."""
    words = text.lower().split()
    if len(words) < k:
        return set()
    return {tuple(words[i:i+k]) for i in range(len(words) - k + 1)}


def jaccard(s1, s2):
    if not s1 or not s2:
        return 0
    return len(s1 & s2) / len(s1 | s2)


# ── Topic heuristic ────────────────────────────────────────────
TOPIC_KEYWORDS = {
    "science": ["experiment", "hypothesis", "molecule", "physics", "chemistry", "biology", "research", "scientific"],
    "technology": ["software", "computer", "algorithm", "programming", "internet", "digital", "technology", "database"],
    "history": ["century", "kingdom", "empire", "ancient", "medieval", "civilization", "dynasty", "historical"],
    "medicine": ["patient", "disease", "treatment", "symptom", "diagnosis", "medical", "health", "clinical"],
    "law": ["court", "legal", "law", "attorney", "judge", "statute", "regulation", "jurisdiction"],
    "education": ["student", "university", "school", "learning", "education", "curriculum", "academic", "teacher"],
    "geography": ["country", "continent", "ocean", "mountain", "river", "population", "region", "territory"],
    "arts": ["painting", "music", "artist", "sculpture", "literature", "poetry", "novel", "creative"],
    "sports": ["game", "team", "championship", "athlete", "tournament", "player", "season", "score"],
    "business": ["company", "market", "economy", "investment", "revenue", "industry", "profit", "financial"],
    "philosophy": ["philosophy", "ethics", "morality", "consciousness", "existence", "metaphysics", "logic", "reasoning"],
    "nature": ["species", "animal", "plant", "ecosystem", "forest", "habitat", "wildlife", "environment"],
}

def detect_topics(text):
    lower = text.lower()
    found = []
    for topic, keywords in TOPIC_KEYWORDS.items():
        if any(kw in lower for kw in keywords):
            found.append(topic)
    return found if found else ["general"]


# ══════════════════════════════════════════════════════════════════
# AUDIT FUNCTION
# ══════════════════════════════════════════════════════════════════
def audit_litdata(litdata_dir, name, num_samples=500, print_samples=10):
    print(f"\n{'='*70}")
    print(f"  AUDITING: {name}")
    print(f"  Path: {litdata_dir}")
    print(f"  Sampling {num_samples} random blocks")
    print(f"{'='*70}")

    index = load_litdata_index(litdata_dir)
    chunks = index["chunks"]
    total_tokens = sum(c["dim"] for c in chunks)

    print(f"\n  Chunks: {len(chunks)}")
    print(f"  Total tokens: {total_tokens:,}")

    # Sample blocks proportionally from each chunk
    total_blocks = sum(c["chunk_size"] for c in chunks)
    samples_per_chunk = {}
    remaining = num_samples
    for i, chunk in enumerate(chunks):
        proportion = chunk["chunk_size"] / total_blocks
        n = max(1, round(proportion * num_samples))
        n = min(n, remaining, chunk["chunk_size"])
        if remaining <= 0:
            break
        samples_per_chunk[i] = n
        remaining -= n

    # Read and decode samples
    texts = []
    for chunk_idx, n_samples in samples_per_chunk.items():
        chunk = chunks[chunk_idx]
        chunk_path = litdata_dir / chunk["filename"]
        if not chunk_path.exists():
            print(f"  WARNING: {chunk_path} missing, skipping")
            continue
        blocks = read_blocks_from_chunk(chunk_path, n_samples, chunk["chunk_size"])
        for block in blocks:
            text = decode_block(block)
            texts.append(text)

    print(f"\n  Decoded {len(texts)} blocks -> text samples")

    # ── 1. English Alignment ─────────────────────────────────────
    print(f"\n  {'─'*60}")
    print(f"  1. ENGLISH ALIGNMENT")
    ascii_ratios = [ascii_ratio(t) for t in texts]
    alpha_ratios = [alpha_ratio(t) for t in texts]
    
    avg_ascii = sum(ascii_ratios) / len(ascii_ratios)
    avg_alpha = sum(alpha_ratios) / len(alpha_ratios)
    low_ascii = sum(1 for r in ascii_ratios if r < 0.85)
    low_alpha = sum(1 for r in alpha_ratios if r < 0.50)

    print(f"     Avg ASCII ratio:  {avg_ascii:.4f}  (want > 0.95)")
    print(f"     Avg alpha ratio:  {avg_alpha:.4f}  (want > 0.60)")
    print(f"     Samples < 85% ASCII: {low_ascii}/{len(texts)}  {'⚠ WARNING' if low_ascii > len(texts)*0.05 else 'OK'}")
    print(f"     Samples < 50% alpha: {low_alpha}/{len(texts)}  {'⚠ WARNING' if low_alpha > len(texts)*0.05 else 'OK'}")

    # Check for non-English text patterns
    non_english_patterns = 0
    for t in texts:
        # CJK, Arabic, Devanagari, Cyrillic heavy blocks
        if re.search(r'[\u4e00-\u9fff\u0600-\u06ff\u0900-\u097f]{10,}', t):
            non_english_patterns += 1
        elif re.search(r'[\u0400-\u04ff]{10,}', t):
            non_english_patterns += 1
    print(f"     Non-English script blocks detected: {non_english_patterns}/{len(texts)}  {'⚠ WARNING' if non_english_patterns > 0 else 'OK'}")

    # ── 2. Cleaning Quality ──────────────────────────────────────
    print(f"\n  {'─'*60}")
    print(f"  2. CLEANING QUALITY")
    
    html_count = sum(1 for t in texts if has_html(t))
    boilerplate_scores = [has_boilerplate(t) for t in texts]
    high_boilerplate = sum(1 for s in boilerplate_scores if s >= 3)
    url_ratios = [url_ratio(t) for t in texts]
    high_url = sum(1 for r in url_ratios if r > 0.03)
    
    short_texts = sum(1 for t in texts if word_count(t) < 20)
    
    print(f"     HTML tags found:       {html_count}/{len(texts)}  {'⚠ WARNING' if html_count > len(texts)*0.02 else 'OK'}")
    print(f"     High boilerplate (3+): {high_boilerplate}/{len(texts)}  {'⚠ WARNING' if high_boilerplate > len(texts)*0.05 else 'OK'}")
    print(f"     High URL ratio (>3%):  {high_url}/{len(texts)}  {'⚠ WARNING' if high_url > len(texts)*0.05 else 'OK'}")
    print(f"     Very short (<20 words):{short_texts}/{len(texts)}  {'⚠ WARNING' if short_texts > len(texts)*0.10 else 'OK'}")
    print(f"     Avg boilerplate score: {sum(boilerplate_scores)/len(boilerplate_scores):.2f}")
    print(f"     Avg URL ratio:         {sum(url_ratios)/len(url_ratios):.4f}")

    # ── 3. Deduplication Check ───────────────────────────────────
    print(f"\n  {'─'*60}")
    print(f"  3. DEDUPLICATION CHECK")
    
    # Check pairwise similarity on a subset
    dedup_sample = min(200, len(texts))
    dedup_texts = random.sample(texts, dedup_sample)
    shingles = [shingle_set(t) for t in dedup_texts]
    
    near_dupes = 0
    exact_dupes = 0
    high_similarities = []
    
    for i in range(len(dedup_texts)):
        for j in range(i + 1, len(dedup_texts)):
            sim = jaccard(shingles[i], shingles[j])
            if sim > 0.8:
                near_dupes += 1
                high_similarities.append((i, j, sim))
            if sim > 0.95:
                exact_dupes += 1
    
    total_pairs = dedup_sample * (dedup_sample - 1) // 2
    print(f"     Checked {total_pairs:,} pairs from {dedup_sample} samples")
    print(f"     Near duplicates (>80% Jaccard): {near_dupes}  {'⚠ WARNING' if near_dupes > 5 else 'OK'}")
    print(f"     Exact duplicates (>95% Jaccard): {exact_dupes}  {'⚠ WARNING' if exact_dupes > 0 else 'OK'}")
    
    if high_similarities:
        print(f"     Top overlaps:")
        for i, j, sim in sorted(high_similarities, key=lambda x: -x[2])[:3]:
            print(f"       [{i}] vs [{j}] = {sim:.3f}")
            print(f"         Sample A: {dedup_texts[i][:80]}...")
            print(f"         Sample B: {dedup_texts[j][:80]}...")

    # ── 4. Data Diversity ────────────────────────────────────────
    print(f"\n  {'─'*60}")
    print(f"  4. DATA DIVERSITY")
    
    # Word count distribution
    word_counts = [word_count(t) for t in texts]
    avg_wc = sum(word_counts) / len(word_counts)
    min_wc = min(word_counts)
    max_wc = max(word_counts)
    
    print(f"     Word count: avg={avg_wc:.0f}  min={min_wc}  max={max_wc}")
    
    # Unique word ratio (vocabulary richness)
    uwr = [unique_word_ratio(t) for t in texts]
    avg_uwr = sum(uwr) / len(uwr)
    print(f"     Avg unique word ratio: {avg_uwr:.4f}  (want > 0.40)")
    
    # Sentence structure
    sent_counts = [sentence_count(t) for t in texts]
    avg_sent = sum(sent_counts) / len(sent_counts)
    no_sentence = sum(1 for s in sent_counts if s == 0)
    print(f"     Avg sentences/block:   {avg_sent:.1f}")
    print(f"     Blocks with 0 sentences: {no_sentence}/{len(texts)}")
    
    # Topic distribution
    all_topics = Counter()
    for t in texts:
        for topic in detect_topics(t):
            all_topics[topic] += 1
    
    print(f"\n     Topic distribution (from {len(texts)} samples):")
    for topic, count in all_topics.most_common():
        bar = "█" * int(count / len(texts) * 40)
        print(f"       {topic:<14} {count:>4} ({count/len(texts)*100:5.1f}%) {bar}")
    
    # ── 5. Flagged Samples ───────────────────────────────────────
    print(f"\n  {'─'*60}")
    print(f"  5. FLAGGED SAMPLES (potential issues)")
    
    flagged = []
    for i, t in enumerate(texts):
        issues = []
        if ascii_ratio(t) < 0.85:
            issues.append(f"low-ascii({ascii_ratio(t):.2f})")
        if has_html(t):
            issues.append("html")
        if has_boilerplate(t) >= 3:
            issues.append(f"boilerplate({has_boilerplate(t)})")
        if url_ratio(t) > 0.05:
            issues.append(f"urls({url_ratio(t):.2f})")
        if word_count(t) < 15:
            issues.append(f"short({word_count(t)}w)")
        if alpha_ratio(t) < 0.40:
            issues.append(f"low-alpha({alpha_ratio(t):.2f})")
        if issues:
            flagged.append((i, issues, t))
    
    if flagged:
        print(f"     {len(flagged)}/{len(texts)} samples flagged ({len(flagged)/len(texts)*100:.1f}%)")
        for i, issues, t in flagged[:5]:
            print(f"\n     Sample #{i}: {', '.join(issues)}")
            print(f"       \"{t[:150]}...\"")
    else:
        print(f"     No samples flagged! All clean.")

    # ── 6. Random Sample Printout ────────────────────────────────
    print(f"\n  {'─'*60}")
    print(f"  6. RANDOM SAMPLES (for manual review)")
    
    sample_indices = random.sample(range(len(texts)), min(print_samples, len(texts)))
    for idx in sample_indices:
        t = texts[idx]
        topics = detect_topics(t)
        print(f"\n  ┌─ Sample #{idx} | {word_count(t)} words | topics: {', '.join(topics)}")
        # Show first 300 chars
        preview = t[:300].replace('\n', ' ↵ ')
        print(f"  │ {preview}")
        print(f"  └─ ascii={ascii_ratio(t):.2f} alpha={alpha_ratio(t):.2f} uwr={unique_word_ratio(t):.2f}")

    # ── Final Verdict ────────────────────────────────────────────
    print(f"\n  {'─'*60}")
    print(f"  VERDICT for {name}")
    
    issues_found = []
    if low_ascii > len(texts) * 0.05:
        issues_found.append(f"  ⚠ {low_ascii} samples have low ASCII ratio")
    if non_english_patterns > 0:
        issues_found.append(f"  ⚠ {non_english_patterns} samples contain non-English scripts")
    if html_count > len(texts) * 0.02:
        issues_found.append(f"  ⚠ {html_count} samples contain HTML tags")
    if high_boilerplate > len(texts) * 0.05:
        issues_found.append(f"  ⚠ {high_boilerplate} samples have high boilerplate")
    if exact_dupes > 0:
        issues_found.append(f"  ⚠ {exact_dupes} exact duplicate pairs found")
    if near_dupes > 5:
        issues_found.append(f"  ⚠ {near_dupes} near-duplicate pairs found")
    if avg_uwr < 0.35:
        issues_found.append(f"  ⚠ Low vocabulary richness ({avg_uwr:.3f})")
    if len(all_topics) < 4:
        issues_found.append(f"  ⚠ Low topic diversity (only {len(all_topics)} topics)")
    
    if issues_found:
        print(f"     Issues found:")
        for issue in issues_found:
            print(f"     {issue}")
    else:
        print(f"     ✓ PASS - Data looks clean, deduplicated, and diverse!")
    
    return {
        "name": name,
        "total_tokens": total_tokens,
        "samples_checked": len(texts),
        "avg_ascii": avg_ascii,
        "avg_alpha": avg_alpha,
        "non_english": non_english_patterns,
        "html_count": html_count,
        "high_boilerplate": high_boilerplate,
        "near_dupes": near_dupes,
        "exact_dupes": exact_dupes,
        "avg_unique_word_ratio": avg_uwr,
        "topic_count": len(all_topics),
        "topics": dict(all_topics),
        "flagged_count": len(flagged),
        "issues": issues_found,
    }


# ══════════════════════════════════════════════════════════════════
# MAIN
# ══════════════════════════════════════════════════════════════════
if __name__ == "__main__":
    random.seed(42)
    
    results = []
    
    # Audit litdata_3b
    r1 = audit_litdata(
        ROOT / "Base" / "data" / "litdata_3b",
        "litdata_3b (General Knowledge Pretraining)",
        num_samples=500,
        print_samples=8,
    )
    results.append(r1)
    
    # Audit litdata_english
    r2 = audit_litdata(
        ROOT / "Base" / "data" / "litdata_english",
        "litdata_english (English Knowledge Continued Pretraining)",
        num_samples=300,
        print_samples=8,
    )
    results.append(r2)
    
    # ── Cross-dataset dedup check ────────────────────────────────
    print(f"\n{'='*70}")
    print(f"  CROSS-DATASET OVERLAP CHECK")
    print(f"{'='*70}")
    print(f"  (Checking if litdata_3b and litdata_english share duplicate content)")
    # This is checked at the token level - since they come from different 
    # source parquets (filtered_3b vs filtered_english), overlap should be minimal
    print(f"  Sources are disjoint by design:")
    print(f"    litdata_3b      <- filtered_3b (15 parquets from general web)")
    print(f"    litdata_english <- filtered_english (FineWeb-Edu + Wikipedia)")
    print(f"  Cross-contamination: UNLIKELY (separate source pipelines)")
    
    # ── Overall Summary ──────────────────────────────────────────
    print(f"\n{'='*70}")
    print(f"  OVERALL QUALITY REPORT")
    print(f"{'='*70}")
    
    all_clean = True
    for r in results:
        status = "PASS" if not r["issues"] else "ISSUES FOUND"
        if r["issues"]:
            all_clean = False
        print(f"\n  {r['name']}")
        print(f"    Tokens:       {r['total_tokens']:>15,}")
        print(f"    Status:       {status}")
        print(f"    English:      ascii={r['avg_ascii']:.3f}  alpha={r['avg_alpha']:.3f}")
        print(f"    Cleanliness:  html={r['html_count']}  boilerplate={r['high_boilerplate']}")
        print(f"    Dedup:        near={r['near_dupes']}  exact={r['exact_dupes']}")
        print(f"    Diversity:    uwr={r['avg_unique_word_ratio']:.3f}  topics={r['topic_count']}")
        print(f"    Flagged:      {r['flagged_count']}/{r['samples_checked']} ({r['flagged_count']/r['samples_checked']*100:.1f}%)")
    
    total_tokens = sum(r["total_tokens"] for r in results)
    print(f"\n  Combined pretrain tokens: {total_tokens:,} ({total_tokens/1e9:.3f}B)")
    
    if all_clean:
        print(f"\n  READY FOR TRAINING FROM SCRATCH!")
    else:
        print(f"\n  REVIEW ISSUES ABOVE before retraining.")