#!/usr/bin/env python3
"""
Deep + Smart cleaning for the 1B English corpus.

Two-phase pipeline:
  PHASE 1 — Deep Clean (text-level):
    - Unicode normalization (NFKC)
    - HTML entity & tag stripping
    - URL / email / file-path removal
    - Control character removal
    - Smart quote / dash / ellipsis normalization
    - Fix spacing, punctuation, sentence boundaries
    - Strip code blocks, import lines
    - Fix repeated words, characters, punctuation

  PHASE 2 — Smart Filter (document-level):
    - Remove docs with non-English scripts (CJK, Arabic, Cyrillic, Devanagari)
    - Remove docs under 50 words
    - Remove very repetitive docs (unique word ratio < 0.22)
    - Remove docs with heavy residual code (5+ code patterns)
    - Strip boilerplate (cookies, subscribe, social, login, nav, copyright)
    - Strip number-only table lines
    - Cross-document deduplication (blake2b fingerprints)

Input:  Base/data/english_1b_raw/ (parquet files from build_english_1b.py)
Output: Base/data/english_1b_clean/ (cleaned parquet files + CLEANING_REPORT.txt)

Usage:
    python Base/scripts/clean_english_1b.py
    python Base/scripts/clean_english_1b.py --input_dir Base/data/english_1b_raw --output_dir Base/data/english_1b_clean
"""

import argparse
import hashlib
import json
import multiprocessing as mp
import re
import time
import unicodedata
from collections import Counter
from functools import partial
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq


# ─── Constants ────────────────────────────────────────────────────────────────

TOKENS_PER_WORD = 1.3
FLUSH_DOCS = 5_000

# ─── Deep Clean Patterns ─────────────────────────────────────────────────────

CTRL_CHARS = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f\ufeff\ufffd]")

HTML_ENTITIES = [
    ("&amp;", "&"), ("&lt;", "<"), ("&gt;", ">"),
    ("&quot;", '"'), ("&#39;", "'"), ("&apos;", "'"),
    ("&nbsp;", " "), ("&mdash;", " - "), ("&ndash;", "-"),
    ("&hellip;", "..."), ("&laquo;", '"'), ("&raquo;", '"'),
    ("&bull;", "- "), ("&middot;", " "), ("&copy;", "(c)"),
    ("&reg;", "(R)"), ("&trade;", "(TM)"), ("&deg;", " degrees"),
]

RE_URL = re.compile(r"https?://\S+|www\.\S+", re.I)
RE_EMAIL = re.compile(r"\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b")
RE_FILE_PATH = re.compile(r"(?:[A-Z]:\\|/(?:home|usr|var|etc|opt)/)\S+")
RE_HTML_TAG = re.compile(r"</?[a-zA-Z][a-zA-Z0-9]*(?:\s[^>]*)?\s*/?>")
RE_HTML_COMMENT = re.compile(r"<!--.*?-->", re.DOTALL)
RE_CODE_BLOCK = re.compile(r"```[\s\S]*?```")
RE_IMPORT_LINE = re.compile(
    r"^(?:import |from \S+ import |#include |using namespace |require\(|"
    r"const \w+ = require|module\.exports|export default).*$",
    re.M,
)
RE_REPEATED_LINE = re.compile(r"^(.{20,})\n(?:\1\n?)+", re.M)
RE_REPEATED_PUNCT = re.compile(r"([!?.])\1{3,}")
RE_REPEATED_CHAR = re.compile(r"(.)\1{5,}")
RE_REPEATED_WORD = re.compile(r"\b(\w+)(?:\s+\1){2,}\b", re.I)
RE_MULTI_NL = re.compile(r"\n{4,}")
RE_MULTI_SP = re.compile(r"[ \t]{2,}")
RE_TRAIL_SP = re.compile(r"[ \t]+$", re.M)
RE_NO_SP_PERIOD = re.compile(r"([.!?])([A-Z])")
RE_SP_PUNCT = re.compile(r"\s+([.,;:!?])")
RE_DBL_PERIOD = re.compile(r"\.{2}(?!\.)")
RE_ORPHAN_PAREN = re.compile(r"\(\s*\)|^\s*\)\s*$|^\s*\(\s*$", re.M)

# ─── Smart Filter Patterns ───────────────────────────────────────────────────

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,}")
RE_DEVANAGARI = re.compile(r"[\u0900-\u097f]{5,}")

RE_RESIDUAL_CODE = re.compile(
    r"(function\s*\(|var\s+\w+\s*=|console\.log|document\.get|"
    r"if\s*\(\s*\w+\s*[!=]==|\.addEventListener|"
    r"def \w+\(|class \w+:|self\.\w+|print\(|"
    r"\.forEach|\.map\(|\.filter\(|async\s+function)",
    re.I,
)

# Boilerplate line patterns to STRIP (not remove doc, just strip these lines)
BOILERPLATE_PATTERNS = [
    re.compile(r"^.*(?:cookie|cookies)\s+(?:policy|consent|notice|preferences|settings).*$", re.I | re.M),
    re.compile(r"^.*(?:subscribe|sign\s*up\s+(?:for|to)\s+(?:our|the)\s+newsletter|unsubscribe|opt[\s-]*out).*$", re.I | re.M),
    re.compile(r"^.*(?:you\s+won'?t\s+believe|click\s+here|read\s+more\s*\.{0,3}$|share\s+this|trending\s+now|sponsored\s+content|advertisement).*$", re.I | re.M),
    re.compile(r"^.*(?:follow\s+us\s+on|share\s+on\s+(?:facebook|twitter|linkedin|instagram)|like\s+us\s+on|tweet\s+this).*$", re.I | re.M),
    re.compile(r"^.*(?:skip\s+to\s+(?:main\s+)?content|back\s+to\s+top|previous\s+article|next\s+article|related\s+(?:articles|posts)).*$", re.I | re.M),
    re.compile(r"^.*(?:log\s*in\s+to\s+(?:your|an)\s+account|create\s+(?:a\s+)?(?:free\s+)?account|forgot\s+(?:your\s+)?password|already\s+(?:a\s+)?member).*$", re.I | re.M),
    re.compile(r"^.*(?:leave\s+a\s+(?:comment|reply)|post\s+a\s+comment|\d+\s+comments?$|logged\s+in\s+as).*$", re.I | re.M),
    re.compile(r"^.*(?:all\s+rights\s+reserved|\(c\)\s*\d{4}|copyright\s+\d{4}).*$", re.I | re.M),
]
RE_NUMBER_TABLE = re.compile(r"^[\d\s,.\-+/%$]+$", re.M)


# ─── Stats tracker ────────────────────────────────────────────────────────────

class Stats:
    def __init__(self):
        self.total_input = 0
        self.total_output = 0
        self.tokens_in = 0
        self.tokens_out = 0
        # Deep clean
        self.ctrl_removed = 0
        self.html_entities = 0
        self.html_tags = 0
        self.urls_removed = 0
        self.emails_removed = 0
        self.code_blocks = 0
        self.import_lines = 0
        self.repeated_lines = 0
        self.repeated_punct = 0
        self.repeated_words = 0
        self.punct_fixed = 0
        self.smart_quotes = 0
        self.orphan_parens = 0
        # Smart filter
        self.dropped_short = 0
        self.dropped_nonenglish = 0
        self.dropped_repetitive = 0
        self.dropped_code = 0
        self.dropped_lowquality = 0
        self.dropped_dup = 0
        self.boilerplate_stripped = 0
        self.number_tables_stripped = 0


# ─── Phase 1: Deep Clean ─────────────────────────────────────────────────────

def deep_clean(text: str, st: Stats) -> str:
    """Thorough text-level cleaning."""
    if not text or len(text.strip()) < 30:
        return ""

    # Unicode normalization
    text = unicodedata.normalize("NFKC", text)

    # Control characters
    new = CTRL_CHARS.sub("", text)
    if new != text:
        st.ctrl_removed += len(text) - len(new)
        text = new

    # HTML entities
    for old, repl in HTML_ENTITIES:
        if old in text:
            c = text.count(old)
            st.html_entities += c
            text = text.replace(old, repl)

    # HTML tags
    new = RE_HTML_TAG.sub("", text)
    if new != text:
        st.html_tags += 1
        text = new

    # HTML comments
    text = RE_HTML_COMMENT.sub("", text)

    # Code blocks
    new = RE_CODE_BLOCK.sub("", text)
    if new != text:
        st.code_blocks += 1
        text = new

    # Import / require lines
    new = RE_IMPORT_LINE.sub("", text)
    if new != text:
        st.import_lines += 1
        text = new

    # URLs
    new = RE_URL.sub("", text)
    if new != text:
        st.urls_removed += 1
        text = new

    # Emails
    new = RE_EMAIL.sub("", text)
    if new != text:
        st.emails_removed += 1
        text = new

    # File paths
    text = RE_FILE_PATH.sub("", text)

    # Repeated lines
    new = RE_REPEATED_LINE.sub(r"\1", text)
    if new != text:
        st.repeated_lines += 1
        text = new

    # Repeated punctuation (!!!! -> !)
    new = RE_REPEATED_PUNCT.sub(r"\1", text)
    if new != text:
        st.repeated_punct += 1
        text = new

    # Repeated characters (aaaaaa -> aa)
    new = RE_REPEATED_CHAR.sub(r"\1\1", text)
    if new != text:
        text = new

    # Repeated words (the the the -> the)
    new = RE_REPEATED_WORD.sub(r"\1", text)
    if new != text:
        st.repeated_words += 1
        text = new

    # Smart quotes & dashes -> ASCII
    changes = [
        ("\u2018", "'"), ("\u2019", "'"),
        ("\u201c", '"'), ("\u201d", '"'),
        ("\u2013", "-"), ("\u2014", " - "),
        ("\u2026", "..."), ("\u00a0", " "),
        ("\u200b", ""), ("\u200d", ""), ("\u200c", ""),
        ("\ufeff", ""),
    ]
    for old, repl in changes:
        if old in text:
            st.smart_quotes += 1
            text = text.replace(old, repl)

    # Fix double periods
    new = RE_DBL_PERIOD.sub(".", text)
    if new != text:
        st.punct_fixed += 1
        text = new

    # Fix missing space after sentence end
    new = RE_NO_SP_PERIOD.sub(r"\1 \2", text)
    if new != text:
        st.punct_fixed += 1
        text = new

    # Fix space before punctuation
    new = RE_SP_PUNCT.sub(r"\1", text)
    if new != text:
        st.punct_fixed += 1
        text = new

    # Orphan parentheses
    new = RE_ORPHAN_PAREN.sub("", text)
    if new != text:
        st.orphan_parens += 1
        text = new

    # Collapse whitespace
    text = RE_MULTI_SP.sub(" ", text)
    text = RE_MULTI_NL.sub("\n\n\n", text)
    text = RE_TRAIL_SP.sub("", text)

    # Strip each line
    text = "\n".join(line.strip() for line in text.splitlines())
    return text.strip()


# ─── Phase 2: Smart Filter ───────────────────────────────────────────────────

def smart_filter(text: str, st: Stats) -> tuple[str | None, str | None]:
    """Document-level quality filter. Returns (cleaned_text, reason_if_dropped)."""
    words = text.split()
    wc = len(words)

    if wc < 50:
        st.dropped_short += 1
        return None, "short"

    # Non-English script check
    scripts = []
    if RE_CJK.search(text):
        scripts.append("CJK")
    if RE_ARABIC.search(text):
        scripts.append("Arabic")
    if RE_CYRILLIC.search(text):
        scripts.append("Cyrillic")
    if RE_DEVANAGARI.search(text):
        scripts.append("Devanagari")
    if scripts:
        st.dropped_nonenglish += 1
        return None, f"non-english:{','.join(scripts)}"

    # Alphabetic ratio
    alpha = sum(c.isalpha() for c in text)
    if alpha / max(len(text), 1) < 0.60:
        st.dropped_lowquality += 1
        return None, "low_alpha"

    # Repetitive content
    if wc > 50:
        unique_ratio = len(set(w.lower() for w in words)) / wc
        if unique_ratio < 0.22:
            st.dropped_repetitive += 1
            return None, f"repetitive:{unique_ratio:.3f}"

    # Residual code detection
    code_hits = len(RE_RESIDUAL_CODE.findall(text))
    if code_hits >= 5:
        st.dropped_code += 1
        return None, f"code:{code_hits}"

    # Strip boilerplate lines
    cleaned = text
    for pat in BOILERPLATE_PATTERNS:
        new = pat.sub("", cleaned)
        if new != cleaned:
            st.boilerplate_stripped += 1
            cleaned = new

    # Strip pure number table lines
    new = RE_NUMBER_TABLE.sub("", cleaned)
    if new != cleaned:
        st.number_tables_stripped += 1
        cleaned = new

    # Re-collapse whitespace after stripping
    cleaned = re.sub(r"\n{3,}", "\n\n", cleaned).strip()

    # Re-check length after stripping
    if len(cleaned.split()) < 40:
        st.dropped_short += 1
        return None, "short_after_strip"

    return cleaned, None


# ─── Deduplication ────────────────────────────────────────────────────────────

def fingerprint(text: str) -> int:
    norm = " ".join(text.lower().split())[:8192]
    return int.from_bytes(
        hashlib.blake2b(norm.encode(), digest_size=8).digest(),
        byteorder="little", signed=False,
    )


# ─── Parallel worker (top-level for pickle) ──────────────────────────────────

def _clean_one(raw_text: str) -> dict:
    """Clean a single document. Returns dict with result + local stats.
    Must be a top-level function for multiprocessing.
    """
    local = {
        "text": None, "fp": None, "drop_reason": None,
        "tokens_in": 0, "tokens_out": 0,
        # deep clean counts
        "ctrl_removed": 0, "html_entities": 0, "html_tags": 0,
        "urls_removed": 0, "emails_removed": 0, "code_blocks": 0,
        "import_lines": 0, "repeated_lines": 0, "repeated_punct": 0,
        "repeated_words": 0, "punct_fixed": 0, "smart_quotes": 0,
        "orphan_parens": 0,
        # smart filter counts
        "dropped_short": 0, "dropped_nonenglish": 0, "dropped_repetitive": 0,
        "dropped_code": 0, "dropped_lowquality": 0,
        "boilerplate_stripped": 0, "number_tables_stripped": 0,
    }
    est_in = max(1, int(len((raw_text or "").split()) * TOKENS_PER_WORD))
    local["tokens_in"] = est_in

    # Use a throwaway Stats for counting
    st = Stats()
    cleaned = deep_clean(raw_text or "", st)

    # Copy stats
    for k in ["ctrl_removed", "html_entities", "html_tags", "urls_removed",
              "emails_removed", "code_blocks", "import_lines", "repeated_lines",
              "repeated_punct", "repeated_words", "punct_fixed", "smart_quotes",
              "orphan_parens"]:
        local[k] = getattr(st, k)

    if not cleaned or len(cleaned) < 50:
        local["drop_reason"] = "short"
        local["dropped_short"] = 1
        return local

    final, reason = smart_filter(cleaned, st)
    for k in ["dropped_short", "dropped_nonenglish", "dropped_repetitive",
              "dropped_code", "dropped_lowquality", "boilerplate_stripped",
              "number_tables_stripped"]:
        local[k] = getattr(st, k)

    if final is None:
        local["drop_reason"] = reason
        return local

    local["text"] = final
    local["fp"] = fingerprint(final)
    local["tokens_out"] = max(1, int(len(final.split()) * TOKENS_PER_WORD))
    return local


# ─── Main pipeline ────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Deep + Smart clean English 1B corpus")
    parser.add_argument("--input_dir", default="Base/data/english_1b_raw")
    parser.add_argument("--output_dir", default="Base/data/english_1b_clean")
    parser.add_argument("--workers", type=int, default=0,
                        help="Number of parallel workers (0 = all CPUs)")
    args = parser.parse_args()

    in_dir = Path(args.input_dir)
    out_dir = Path(args.output_dir)
    n_workers = args.workers or mp.cpu_count()

    if not in_dir.exists():
        print(f"ERROR: Input directory not found: {in_dir}")
        print("Run build_english_1b.py first!")
        return

    out_dir.mkdir(parents=True, exist_ok=True)

    parquet_files = sorted(in_dir.glob("*.parquet"))
    if not parquet_files:
        print(f"ERROR: No parquet files in {in_dir}")
        return

    print("=" * 72)
    print("  DEEP + SMART CLEAN — English 1B Corpus  ** PARALLEL MODE **")
    print("=" * 72)
    print(f"  Input:   {in_dir} ({len(parquet_files)} files)")
    print(f"  Output:  {out_dir}")
    print(f"  Workers: {n_workers} (CPU cores: {mp.cpu_count()})")
    print("=" * 72)

    st = Stats()
    seen_fps: set[int] = set()
    out_buf: dict[str, list] = {}
    out_idx: Counter = Counter()
    t0 = time.time()

    def flush_buf(source: str):
        bucket = out_buf.get(source)
        if not bucket:
            return
        i = out_idx[source]
        path = out_dir / f"{source}_{i:04d}.parquet"
        pq.write_table(pa.table({"text": bucket}), str(path))
        out_idx[source] += 1
        out_buf[source] = []

    # Use a persistent pool for all files
    pool = mp.Pool(n_workers)

    try:
        for pf in parquet_files:
            source = pf.stem.rsplit("_", 1)[0]
            print(f"\n  Processing: {pf.name}")
            table = pq.read_table(str(pf), columns=["text"])
            texts = table.column("text").to_pylist()
            del table

            # Parallel clean: map all docs across workers
            results = pool.map(_clean_one, texts, chunksize=256)

            file_kept = 0
            for r in results:
                st.total_input += 1
                st.tokens_in += r["tokens_in"]
                # Merge deep clean stats
                st.ctrl_removed += r["ctrl_removed"]
                st.html_entities += r["html_entities"]
                st.html_tags += r["html_tags"]
                st.urls_removed += r["urls_removed"]
                st.emails_removed += r["emails_removed"]
                st.code_blocks += r["code_blocks"]
                st.import_lines += r["import_lines"]
                st.repeated_lines += r["repeated_lines"]
                st.repeated_punct += r["repeated_punct"]
                st.repeated_words += r["repeated_words"]
                st.punct_fixed += r["punct_fixed"]
                st.smart_quotes += r["smart_quotes"]
                st.orphan_parens += r["orphan_parens"]
                # Merge smart filter stats
                st.dropped_short += r["dropped_short"]
                st.dropped_nonenglish += r["dropped_nonenglish"]
                st.dropped_repetitive += r["dropped_repetitive"]
                st.dropped_code += r["dropped_code"]
                st.dropped_lowquality += r["dropped_lowquality"]
                st.boilerplate_stripped += r["boilerplate_stripped"]
                st.number_tables_stripped += r["number_tables_stripped"]

                if r["text"] is None:
                    continue

                # Dedup (must be serial)
                fpi = r["fp"]
                if fpi in seen_fps:
                    st.dropped_dup += 1
                    continue
                seen_fps.add(fpi)

                # Accept
                st.total_output += 1
                st.tokens_out += r["tokens_out"]
                bucket = out_buf.setdefault(source, [])
                bucket.append(r["text"])
                file_kept += 1
                if len(bucket) >= FLUSH_DOCS:
                    flush_buf(source)

            elapsed = time.time() - t0
            rate = st.total_input / max(elapsed, 1)
            print(f"    kept={file_kept:,} / {len(texts):,}  "
                  f"total_out={st.total_output:,}  "
                  f"~{st.tokens_out:,} tokens  "
                  f"elapsed={elapsed:.0f}s  "
                  f"rate={rate:,.0f} docs/s")

    finally:
        pool.close()
        pool.join()

        elapsed = time.time() - t0
        print(f"    kept={file_kept:,} / {len(texts):,}  "
              f"total_out={st.total_output:,}  "
              f"~{st.tokens_out:,} tokens  "
              f"elapsed={elapsed:.0f}s")

    # Flush remaining
    for source in list(out_buf):
        flush_buf(source)

    elapsed = time.time() - t0

    # ── Report ────────────────────────────────────────────────────────────
    kept_pct = st.total_output / max(st.total_input, 1) * 100
    tok_pct = st.tokens_out / max(st.tokens_in, 1) * 100

    lines = [
        "=" * 72,
        "  CLEANING REPORT — English 1B Corpus",
        "=" * 72, "",
        f"Input:       {in_dir}",
        f"Output:      {out_dir}",
        f"Files in:    {len(parquet_files)}",
        f"Files out:   {sum(out_idx.values())}",
        f"Elapsed:     {elapsed:.0f}s ({elapsed/60:.1f} min)", "",
        "─── Document Stats ───",
        f"  Total input:    {st.total_input:,}",
        f"  Total output:   {st.total_output:,}",
        f"  Kept:           {kept_pct:.1f}%",
        f"  Unique FPs:     {len(seen_fps):,}", "",
        "─── Token Estimates ───",
        f"  Tokens in:      {st.tokens_in:,}",
        f"  Tokens out:     {st.tokens_out:,}",
        f"  Retained:       {tok_pct:.1f}%", "",
        "─── Deep Clean Stats ───",
        f"  Control chars removed:  {st.ctrl_removed:,}",
        f"  HTML entities fixed:    {st.html_entities:,}",
        f"  HTML tags stripped:     {st.html_tags:,}",
        f"  URLs removed:           {st.urls_removed:,}",
        f"  Emails removed:         {st.emails_removed:,}",
        f"  Code blocks removed:    {st.code_blocks:,}",
        f"  Import lines removed:   {st.import_lines:,}",
        f"  Repeated lines fixed:   {st.repeated_lines:,}",
        f"  Repeated punct fixed:   {st.repeated_punct:,}",
        f"  Repeated words fixed:   {st.repeated_words:,}",
        f"  Punctuation fixed:      {st.punct_fixed:,}",
        f"  Smart quotes fixed:     {st.smart_quotes:,}",
        f"  Orphan parens removed:  {st.orphan_parens:,}", "",
        "─── Smart Filter Stats ───",
        f"  Dropped (too short):    {st.dropped_short:,}",
        f"  Dropped (non-English):  {st.dropped_nonenglish:,}",
        f"  Dropped (repetitive):   {st.dropped_repetitive:,}",
        f"  Dropped (code):         {st.dropped_code:,}",
        f"  Dropped (low quality):  {st.dropped_lowquality:,}",
        f"  Dropped (duplicate):    {st.dropped_dup:,}",
        f"  Boilerplate stripped:   {st.boilerplate_stripped:,}",
        f"  Number tables stripped: {st.number_tables_stripped:,}", "",
        "Next: python Base/scripts/prepare_litdata.py \\",
        "        --filtered_dir Base/data/english_1b_clean \\",
        "        --output_dir Base/data/litdata_english_1b \\",
        "        --label ENGLISH_1B", "",
        "=" * 72,
    ]
    report = "\n".join(lines)
    print("\n" + report)
    (out_dir / "CLEANING_REPORT.txt").write_text(report, encoding="utf-8")
    (out_dir / "clean_stats.json").write_text(json.dumps({
        "docs_in": st.total_input,
        "docs_out": st.total_output,
        "tokens_in": st.tokens_in,
        "tokens_out": st.tokens_out,
        "dropped": {
            "short": st.dropped_short,
            "non_english": st.dropped_nonenglish,
            "repetitive": st.dropped_repetitive,
            "code": st.dropped_code,
            "low_quality": st.dropped_lowquality,
            "duplicate": st.dropped_dup,
        },
    }, indent=2), encoding="utf-8")

    print(f"\nCleaning complete. {st.total_output:,} docs, ~{st.tokens_out:,} tokens retained.")


if __name__ == "__main__":
    main()