# -*- coding: utf-8 -*-
"""
COMPREHENSIVE AUDIT of litdata_pretrain_final (5B tokens)

Checks:
  1. Binary integrity — all 308 chunks readable, headers valid, correct dims
  2. Token range validation — all tokens in [0, vocab_size), no garbage
  3. EOS placement — correct document boundaries
  4. Duplicate detection — sample chunks for near-duplicate documents
  5. Token distribution — check for anomalous frequency spikes (noise)
  6. Decode quality — random sample of 50 documents, decode + inspect
  7. Training readiness — correct block_size, total tokens match config

This script is READ-ONLY. It does NOT modify data.
"""

import json
import os
import re
import time
import hashlib
from pathlib import Path
from collections import Counter, defaultdict

import numpy as np

ROOT = Path(__file__).resolve().parent.parent.parent
FINAL_DIR = ROOT / "Base" / "data" / "litdata_pretrain_final"
TOKENIZER_PATH = str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
BLOCK_SIZE = 1025
DTYPE = np.int32
EOS_TOKEN_ID = 0
VOCAB_SIZE = 50277

# How many chunks to deeply scan for duplicates (all 308 is fine for audit)
MAX_CHUNKS_DEEP = 308
# Number of random documents to decode and display
DECODE_SAMPLES = 50
# Min hash length for near-duplicate detection (chars)
HASH_WINDOW = 200


def read_chunk(filepath):
    """Read a single chunk binary file. Returns (blocks, num_blocks)."""
    with open(filepath, "rb") as f:
        raw = f.read()

    # Parse header
    num_blocks = np.frombuffer(raw[:4], dtype=np.uint32)[0]
    header_size = 4 + (num_blocks + 1) * 4  # 1 uint32 + (num_blocks+1) offsets
    data_bytes = raw[header_size:]

    expected_tokens = num_blocks * BLOCK_SIZE
    expected_bytes = expected_tokens * DTYPE().itemsize

    tokens = np.frombuffer(data_bytes[:expected_bytes], dtype=DTYPE)
    return tokens, int(num_blocks)


def extract_documents(tokens):
    """Split a token array at EOS boundaries into individual documents."""
    eos_positions = np.where(tokens == EOS_TOKEN_ID)[0]
    docs = []
    start = 0
    for eos_pos in eos_positions:
        if eos_pos > start:
            doc_tokens = tokens[start:eos_pos]
            if len(doc_tokens) > 0:
                docs.append(doc_tokens)
        start = eos_pos + 1
    # Trailing tokens (no final EOS — partial doc carried across chunks)
    if start < len(tokens):
        remaining = tokens[start:]
        if len(remaining) > 5:  # ignore tiny fragments
            docs.append(remaining)
    return docs


def token_hash(doc_tokens, window=200):
    """Hash first N tokens for near-duplicate detection."""
    key = doc_tokens[:window].tobytes()
    return hashlib.md5(key).hexdigest()


def main():
    t_start = time.time()

    # Load index
    with open(FINAL_DIR / "index.json") as f:
        index = json.load(f)
    chunks_meta = index["chunks"]
    config = index.get("config", {})
    num_chunks = len(chunks_meta)

    print(f"{'='*75}")
    print(f"  COMPREHENSIVE AUDIT — litdata_pretrain_final")
    print(f"{'='*75}")
    print(f"  Chunks in index: {num_chunks}")
    print(f"  Config: {config}")
    print(f"  Expected BLOCK_SIZE: {BLOCK_SIZE}")
    print(f"  Expected VOCAB_SIZE: {VOCAB_SIZE}")
    print(f"  Expected EOS: {EOS_TOKEN_ID}")
    print()

    # ═════════════════════════════════════════════════════════════════════
    # CHECK 1: Binary integrity + token range
    # ═════════════════════════════════════════════════════════════════════
    print(f"  CHECK 1: BINARY INTEGRITY + TOKEN RANGE")
    print(f"  {'-'*60}")

    total_tokens = 0
    total_blocks = 0
    missing_files = []
    corrupt_chunks = []
    out_of_range_chunks = []
    eos_counts = []
    chunk_token_counts = []

    # Global token frequency (sample every 10th chunk for speed)
    global_freq = Counter()
    FREQ_SAMPLE_INTERVAL = 3  # sample every 3rd chunk

    # For duplicate detection
    doc_hashes = defaultdict(list)  # hash -> [(chunk_idx, doc_idx)]
    total_docs = 0
    doc_lengths = []

    # For decode sampling: collect random docs
    sample_docs = []
    np.random.seed(42)

    for ci, meta in enumerate(chunks_meta):
        filename = meta["filename"]
        filepath = FINAL_DIR / filename
        expected_dim = meta["dim"]

        if not filepath.exists():
            missing_files.append(filename)
            print(f"    MISSING: {filename}")
            continue

        try:
            tokens, num_blocks = read_chunk(filepath)
        except Exception as e:
            corrupt_chunks.append((filename, str(e)))
            print(f"    CORRUPT: {filename} — {e}")
            continue

        actual_dim = len(tokens)
        if actual_dim != expected_dim:
            corrupt_chunks.append((filename, f"dim mismatch: expected {expected_dim}, got {actual_dim}"))
            print(f"    DIM MISMATCH: {filename} expected {expected_dim}, got {actual_dim}")

        total_tokens += actual_dim
        total_blocks += num_blocks
        chunk_token_counts.append(actual_dim)

        # Token range check
        min_tok = int(tokens.min())
        max_tok = int(tokens.max())
        if min_tok < 0 or max_tok >= VOCAB_SIZE:
            out_of_range_chunks.append((filename, min_tok, max_tok))
            print(f"    OUT OF RANGE: {filename} min={min_tok} max={max_tok}")

        # EOS count
        eos_count = int(np.sum(tokens == EOS_TOKEN_ID))
        eos_counts.append(eos_count)

        # Token frequency (sampled)
        if ci % FREQ_SAMPLE_INTERVAL == 0:
            unique, counts = np.unique(tokens, return_counts=True)
            for tok, cnt in zip(unique, counts):
                global_freq[int(tok)] += int(cnt)

        # Extract documents for duplicate + quality check
        if ci < MAX_CHUNKS_DEEP:
            docs = extract_documents(tokens)
            for di, doc in enumerate(docs):
                total_docs += 1
                doc_lengths.append(len(doc))
                if len(doc) >= 50:
                    h = token_hash(doc)
                    doc_hashes[h].append((ci, di))

                # Random sample for decode
                if len(sample_docs) < DECODE_SAMPLES and np.random.random() < 0.0005:
                    sample_docs.append(doc)

        if (ci + 1) % 50 == 0:
            print(f"    Scanned {ci+1}/{num_chunks} chunks... ({total_tokens:,} tokens)")

    print(f"    Scanned all {num_chunks} chunks: {total_tokens:,} total tokens")
    print()

    # Results for Check 1
    issues = []
    if missing_files:
        issues.append(f"MISSING FILES: {len(missing_files)}")
    if corrupt_chunks:
        issues.append(f"CORRUPT CHUNKS: {len(corrupt_chunks)}")
    if out_of_range_chunks:
        issues.append(f"OUT-OF-RANGE TOKENS: {len(out_of_range_chunks)} chunks")

    if not issues:
        print(f"    ✓ All {num_chunks} chunks intact, all tokens in [0, {VOCAB_SIZE})")
    else:
        for issue in issues:
            print(f"    ✗ {issue}")
    print()

    # ═════════════════════════════════════════════════════════════════════
    # CHECK 2: EOS + DOCUMENT BOUNDARIES
    # ═════════════════════════════════════════════════════════════════════
    print(f"  CHECK 2: EOS & DOCUMENT BOUNDARIES")
    print(f"  {'-'*60}")

    total_eos = sum(eos_counts)
    avg_eos = total_eos / max(num_chunks, 1)
    min_eos = min(eos_counts) if eos_counts else 0
    max_eos = max(eos_counts) if eos_counts else 0

    print(f"    Total EOS tokens: {total_eos:,}")
    print(f"    Avg EOS per chunk: {avg_eos:.1f}")
    print(f"    Min/Max EOS per chunk: {min_eos} / {max_eos}")
    print(f"    Total documents found: {total_docs:,}")

    if doc_lengths:
        dl = np.array(doc_lengths)
        print(f"    Doc length (tokens): min={int(dl.min())}, median={int(np.median(dl))}, "
              f"mean={int(dl.mean())}, max={int(dl.max())}")
        tiny_docs = int(np.sum(dl < 20))
        short_docs = int(np.sum(dl < 50))
        long_docs = int(np.sum(dl > 50000))
        print(f"    Tiny docs (<20 tok): {tiny_docs:,} ({100*tiny_docs/len(dl):.2f}%)")
        print(f"    Short docs (<50 tok): {short_docs:,} ({100*short_docs/len(dl):.2f}%)")
        print(f"    Very long docs (>50K tok): {long_docs:,}")

    # Flag if too many tiny docs (noise)
    if doc_lengths and tiny_docs / len(dl) > 0.05:
        print(f"    ⚠ WARNING: {100*tiny_docs/len(dl):.1f}% tiny docs — possible noise")
    else:
        print(f"    ✓ Document boundaries look healthy")
    print()

    # ═════════════════════════════════════════════════════════════════════
    # CHECK 3: DUPLICATE DETECTION
    # ═════════════════════════════════════════════════════════════════════
    print(f"  CHECK 3: NEAR-DUPLICATE DETECTION")
    print(f"  {'-'*60}")

    dup_groups = {h: locs for h, locs in doc_hashes.items() if len(locs) > 1}
    dup_doc_count = sum(len(locs) - 1 for locs in dup_groups.values())

    print(f"    Unique doc hashes: {len(doc_hashes):,}")
    print(f"    Duplicate groups: {len(dup_groups):,}")
    print(f"    Duplicate docs (extra copies): {dup_doc_count:,}")
    dup_pct = 100 * dup_doc_count / max(total_docs, 1)
    print(f"    Duplication rate: {dup_pct:.3f}%")

    if dup_pct > 5.0:
        print(f"    ⚠ WARNING: High duplication rate ({dup_pct:.1f}%). Consider deduplication.")
    elif dup_pct > 1.0:
        print(f"    ⚠ MODERATE: {dup_pct:.2f}% duplicates. Acceptable but not ideal.")
    else:
        print(f"    ✓ Very low duplication ({dup_pct:.3f}%). Excellent.")

    # Show a few duplicate examples
    if dup_groups:
        print(f"\n    Top 5 duplicate groups (by copy count):")
        sorted_dups = sorted(dup_groups.items(), key=lambda x: -len(x[1]))[:5]
        for h, locs in sorted_dups:
            print(f"      hash={h[:12]}... : {len(locs)} copies in chunks {[l[0] for l in locs[:6]]}")
    print()

    # ═════════════════════════════════════════════════════════════════════
    # CHECK 4: TOKEN DISTRIBUTION (anomaly detection)
    # ═════════════════════════════════════════════════════════════════════
    print(f"  CHECK 4: TOKEN DISTRIBUTION ANALYSIS")
    print(f"  {'-'*60}")

    total_sampled = sum(global_freq.values())
    print(f"    Sampled tokens: {total_sampled:,} (from every {FREQ_SAMPLE_INTERVAL}rd chunk)")

    # Top 20 most frequent tokens
    most_common = global_freq.most_common(30)
    print(f"    Top 30 tokens by frequency:")
    for tok_id, count in most_common:
        pct = 100 * count / total_sampled
        print(f"      token {tok_id:>6}: {count:>12,} ({pct:>5.2f}%)")

    # Check for suspicious spikes
    # If any single non-EOS token is >10% of all tokens, it's suspicious
    suspicious = []
    for tok_id, count in most_common:
        pct = count / total_sampled
        if tok_id != EOS_TOKEN_ID and pct > 0.10:
            suspicious.append((tok_id, pct))

    if suspicious:
        print(f"\n    ⚠ SUSPICIOUS: These non-EOS tokens appear in >10% of data:")
        for tok_id, pct in suspicious:
            print(f"      token {tok_id}: {100*pct:.2f}%")
    else:
        print(f"\n    ✓ No anomalous token frequency spikes detected")

    # Vocab coverage
    unique_tokens = len(global_freq)
    coverage_pct = 100 * unique_tokens / VOCAB_SIZE
    print(f"    Unique tokens seen: {unique_tokens:,} / {VOCAB_SIZE:,} ({coverage_pct:.1f}% vocab coverage)")

    # Check for dead zones (large ranges of unused tokens)
    used_set = set(global_freq.keys())
    unused_ranges = []
    start_unused = None
    for i in range(VOCAB_SIZE):
        if i not in used_set:
            if start_unused is None:
                start_unused = i
        else:
            if start_unused is not None:
                gap = i - start_unused
                if gap > 500:
                    unused_ranges.append((start_unused, i - 1, gap))
                start_unused = None

    if unused_ranges:
        print(f"    Large unused token ranges (>500):")
        for s, e, g in unused_ranges[:5]:
            print(f"      tokens {s}-{e} ({g} unused)")
    print()

    # ═════════════════════════════════════════════════════════════════════
    # CHECK 5: DECODE QUALITY — sample random documents
    # ═════════════════════════════════════════════════════════════════════
    print(f"  CHECK 5: DECODED DOCUMENT SAMPLES")
    print(f"  {'-'*60}")

    # If we didn't get enough samples randomly, grab from specific chunks
    if len(sample_docs) < DECODE_SAMPLES:
        # Grab from spread-out chunks
        sample_chunks = np.linspace(0, num_chunks - 1, min(DECODE_SAMPLES - len(sample_docs), 25), dtype=int)
        for sci in sample_chunks:
            if len(sample_docs) >= DECODE_SAMPLES:
                break
            meta = chunks_meta[sci]
            filepath = FINAL_DIR / meta["filename"]
            try:
                tokens, _ = read_chunk(filepath)
                docs = extract_documents(tokens)
                if docs:
                    # Pick a random doc from this chunk
                    idx = np.random.randint(0, len(docs))
                    sample_docs.append(docs[idx])
            except:
                pass

    # Now decode
    from tokenizers import Tokenizer
    tokenizer = Tokenizer.from_file(TOKENIZER_PATH)

    quality_issues = 0
    noise_docs = 0
    non_english_docs = 0

    RE_CJK = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7af]{3,}')
    RE_ARABIC = re.compile(r'[\u0600-\u06ff]{5,}')
    RE_CYRILLIC = re.compile(r'[\u0400-\u04ff]{5,}')

    for si, doc_tokens in enumerate(sample_docs[:DECODE_SAMPLES]):
        text = tokenizer.decode(doc_tokens.tolist(), skip_special_tokens=False)

        # Quality checks
        words = text.split()
        word_count = len(words)
        alpha_ratio = sum(c.isalpha() for c in text) / max(len(text), 1)
        unique_words = len(set(w.lower() for w in words)) / max(word_count, 1)

        is_noisy = False
        is_non_english = False
        flags = []

        if alpha_ratio < 0.50:
            flags.append(f"low-alpha({alpha_ratio:.2f})")
            is_noisy = True
        if word_count > 30 and unique_words < 0.15:
            flags.append(f"repetitive({unique_words:.2f})")
            is_noisy = True
        if RE_CJK.search(text) or RE_ARABIC.search(text) or RE_CYRILLIC.search(text):
            flags.append("non-English")
            is_non_english = True
        if word_count < 10:
            flags.append("very-short")
            is_noisy = True

        if is_noisy:
            noise_docs += 1
        if is_non_english:
            non_english_docs += 1
        if flags:
            quality_issues += 1

        # Print preview for flagged + a few clean ones
        if flags or si < 5 or si % 10 == 0:
            preview = text[:300].replace('\n', ' ↵ ')
            status = f"⚠ {','.join(flags)}" if flags else "✓ clean"
            print(f"\n    Sample {si+1}/{len(sample_docs[:DECODE_SAMPLES])} [{status}] ({word_count} words, alpha={alpha_ratio:.2f}, unique={unique_words:.2f})")
            print(f"      \"{preview}...\"")

    print(f"\n    DECODE SUMMARY:")
    print(f"      Samples checked: {min(len(sample_docs), DECODE_SAMPLES)}")
    print(f"      Clean: {min(len(sample_docs), DECODE_SAMPLES) - quality_issues}")
    print(f"      Noisy: {noise_docs}")
    print(f"      Non-English: {non_english_docs}")
    print(f"      Quality issues: {quality_issues}")

    if quality_issues / max(len(sample_docs), 1) > 0.1:
        print(f"      ⚠ WARNING: >10% of sampled docs have quality issues")
    else:
        print(f"      ✓ Sample quality looks good")
    print()

    # ═════════════════════════════════════════════════════════════════════
    # CHECK 6: TRAINING READINESS
    # ═════════════════════════════════════════════════════════════════════
    print(f"  CHECK 6: TRAINING READINESS")
    print(f"  {'-'*60}")

    index_total = sum(c["dim"] for c in chunks_meta)
    print(f"    Index total tokens: {index_total:,}")
    print(f"    Actual total tokens: {total_tokens:,}")
    if index_total == total_tokens:
        print(f"    ✓ Index matches actual data perfectly")
    else:
        print(f"    ⚠ MISMATCH: index says {index_total:,} but files have {total_tokens:,}")

    # Check all blocks are BLOCK_SIZE aligned
    misaligned = [c["filename"] for c in chunks_meta if c["dim"] % BLOCK_SIZE != 0]
    if misaligned:
        print(f"    ⚠ {len(misaligned)} chunks not BLOCK_SIZE-aligned: {misaligned[:5]}")
    else:
        print(f"    ✓ All chunks perfectly BLOCK_SIZE-aligned ({BLOCK_SIZE})")

    # Config check
    if config.get("block_size") == BLOCK_SIZE:
        print(f"    ✓ Config block_size matches: {BLOCK_SIZE}")
    else:
        print(f"    ⚠ Config block_size: {config.get('block_size')} (expected {BLOCK_SIZE})")

    if config.get("vocab_size") == VOCAB_SIZE:
        print(f"    ✓ Config vocab_size matches: {VOCAB_SIZE}")
    else:
        print(f"    ⚠ Config vocab_size: {config.get('vocab_size')} (expected {VOCAB_SIZE})")

    # Tokens per epoch for batch size calculation
    print(f"\n    TRAINING PARAMETERS:")
    seq_len = 1024
    steps = total_tokens // (120 * seq_len)  # global_batch=120
    print(f"      Total tokens: {total_tokens:,}")
    print(f"      Global batch size 120 × seq 1024 = {120*1024:,} tokens/step")
    print(f"      Total steps for 1 epoch: {steps:,}")
    print(f"      At ~10 steps/sec (A100): ~{steps/10/60:.0f} min ≈ {steps/10/3600:.1f} hours")
    print()

    # ═════════════════════════════════════════════════════════════════════
    # OVERALL VERDICT
    # ═════════════════════════════════════════════════════════════════════
    all_issues = []
    if missing_files: all_issues.append(f"{len(missing_files)} missing files")
    if corrupt_chunks: all_issues.append(f"{len(corrupt_chunks)} corrupt chunks")
    if out_of_range_chunks: all_issues.append(f"{len(out_of_range_chunks)} out-of-range chunks")
    if dup_pct > 5.0: all_issues.append(f"High duplication: {dup_pct:.1f}%")
    if suspicious: all_issues.append("Suspicious token spikes")
    if doc_lengths and tiny_docs / len(dl) > 0.05: all_issues.append("Too many tiny docs")
    if quality_issues / max(len(sample_docs), 1) > 0.1: all_issues.append("Quality issues in samples")

    elapsed = time.time() - t_start

    print(f"{'='*75}")
    print(f"  AUDIT VERDICT")
    print(f"{'='*75}")
    if not all_issues:
        print(f"  ✓ ALL CHECKS PASSED — Dataset is TRAINING-READY")
        print(f"    {total_tokens:,} clean tokens across {num_chunks} chunks")
        print(f"    No corruption, minimal duplicates, good quality")
    else:
        print(f"  ⚠ ISSUES FOUND:")
        for issue in all_issues:
            print(f"    - {issue}")
    print(f"\n  Audit completed in {elapsed:.1f}s")
    print(f"{'='*75}")


if __name__ == "__main__":
    main()