# -*- coding: utf-8 -*-
"""
Reclean & normalize all pretraining data for optimal 100M model learning.

Reads litdata_3b and litdata_english, decodes all documents back to text,
applies comprehensive English cleaning/normalization, re-tokenizes, and
writes new litdata chunks.
"""

import json
import os
import re
import sys
import time
import unicodedata
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
DTYPE_INDEX = 16
CHUNK_BYTES_TARGET = 64 * 1024 * 1024
EOS_TOKEN_ID = 0

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

# ==============================================================================
# TEXT CLEANING PIPELINE
# ==============================================================================

# Null / control chars to strip
CONTROL_CHARS = [
    "\x00", "\x01", "\x02", "\x03", "\x04", "\x05", "\x06", "\x07",
    "\x08", "\x0b", "\x0c", "\x0e", "\x0f", "\x10", "\x11", "\x12",
    "\x13", "\x14", "\x15", "\x16", "\x17", "\x18", "\x19", "\x1a",
    "\x1b", "\x1c", "\x1d", "\x1e", "\x1f", "\x7f",
    "\ufeff", "\ufffd",
]

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

# URL/email/path patterns
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+')

# HTML tag leftovers
RE_HTML_TAG = re.compile(r'</?[a-zA-Z][a-zA-Z0-9]*(?:\s[^>]*)?\s*/?>')
RE_HTML_COMMENT = re.compile(r'<!--.*?-->', re.DOTALL)

# Code/programming artifacts
RE_CODE_BLOCK = re.compile(r'```[\s\S]*?```')
RE_IMPORT = re.compile(r'^(?:import |from \S+ import |#include |using namespace |require\()', re.M)

# Repeated content
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)

# Whitespace
RE_MULTI_NEWLINE = re.compile(r'\n{4,}')
RE_MULTI_SPACE = re.compile(r'[ \t]{2,}')
RE_TRAILING_SPACE = re.compile(r'[ \t]+$', re.M)

# Sentence fixing
RE_NO_SPACE_AFTER_PERIOD = re.compile(r'([.!?])([A-Z])')
RE_DOUBLE_PERIOD = re.compile(r'\.{2}(?!\.)')  # .. but not ...
RE_SPACE_BEFORE_PUNCT = re.compile(r'\s+([.,;:!?])')


def clean_text(text):
    """Apply full cleaning pipeline to a single document text."""
    if not text or len(text.strip()) < 30:
        return None

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

    # 2. Strip control characters
    for ch in CONTROL_CHARS:
        text = text.replace(ch, "")

    # 3. Fix HTML entities
    for old, new in HTML_ENTITIES:
        text = text.replace(old, new)

    # 4. Remove HTML tags and comments
    text = RE_HTML_COMMENT.sub("", text)
    text = RE_HTML_TAG.sub("", text)

    # 5. Remove URLs, emails, file paths
    text = RE_URL.sub("", text)
    text = RE_EMAIL.sub("", text)
    text = RE_FILE_PATH.sub("", text)

    # 6. Remove code blocks
    text = RE_CODE_BLOCK.sub("", text)

    # 7. Fix repeated content
    text = RE_REPEATED_LINE.sub(r'\1', text)
    text = RE_REPEATED_PUNCT.sub(r'\1\1\1', text)
    text = RE_REPEATED_CHAR.sub(r'\1\1\1', text)
    text = RE_REPEATED_WORD.sub(r'\1', text)

    # 8. Normalize whitespace
    text = text.replace('\t', ' ')
    text = RE_TRAILING_SPACE.sub('', text)
    text = RE_MULTI_SPACE.sub(' ', text)
    text = RE_MULTI_NEWLINE.sub('\n\n\n', text)

    # 9. Fix punctuation
    text = RE_DOUBLE_PERIOD.sub('.', text)
    text = RE_NO_SPACE_AFTER_PERIOD.sub(r'\1 \2', text)
    text = RE_SPACE_BEFORE_PUNCT.sub(r'\1', text)

    # 10. Normalize smart quotes and special punctuation to ASCII
    text = text.replace('\u2018', "'").replace('\u2019', "'")
    text = text.replace('\u201c', '"').replace('\u201d', '"')
    text = text.replace('\u2013', '-').replace('\u2014', ' - ')
    text = text.replace('\u2026', '...')
    text = text.replace('\u2022', '- ')
    text = text.replace('\u00b7', ' ')
    text = text.replace('\u00a0', ' ')  # non-breaking space

    # 11. Process line by line: capitalize sentence starts, remove junk lines
    lines = text.split('\n')
    clean_lines = []
    for line in lines:
        line = line.strip()
        if not line:
            clean_lines.append('')
            continue

        # Skip lines that are mostly non-alphabetic (tables, code, etc.)
        if len(line) > 10:
            alpha_count = sum(1 for c in line if c.isalpha())
            if alpha_count / len(line) < 0.40:
                continue

        # Skip lines with too many special chars (tables, markup)
        if line.count('|') > 3 or line.count('{') > 2 or line.count('}') > 2:
            continue

        # Skip lines that look like code imports
        if RE_IMPORT.match(line):
            continue

        # Capitalize first letter of sentences
        if line and line[0].isalpha() and line[0].islower():
            if not clean_lines or clean_lines[-1] == '' or clean_lines[-1].rstrip().endswith(('.', '!', '?', ':')):
                line = line[0].upper() + line[1:]

        clean_lines.append(line)

    text = '\n'.join(clean_lines)

    # 12. Remove leading/trailing whitespace
    text = text.strip()

    # 13. Remove duplicate paragraphs
    paragraphs = text.split('\n\n')
    seen = set()
    unique_paragraphs = []
    for p in paragraphs:
        p_stripped = p.strip()
        if not p_stripped:
            continue
        p_key = ' '.join(p_stripped.lower().split())
        if p_key not in seen:
            seen.add(p_key)
            unique_paragraphs.append(p_stripped)
    text = '\n\n'.join(unique_paragraphs)

    # 14. Final quality gate
    text = text.strip()
    if len(text) < 50:
        return None
    word_count = len(text.split())
    if word_count < 10:
        return None
    # Must be mostly ASCII/English
    ascii_count = sum(1 for c in text if ord(c) < 128)
    if ascii_count / max(len(text), 1) < 0.85:
        return None

    return text


# ==============================================================================
# LITDATA I/O
# ==============================================================================

def read_all_tokens(litdata_dir):
    """Read all chunks and return the full flat token stream as numpy array."""
    with open(litdata_dir / "index.json") as f:
        index = json.load(f)

    chunks = index["chunks"]
    total_tokens = sum(c["dim"] for c in chunks)
    print(f"  Reading {len(chunks)} chunks ({total_tokens:,} tokens)...")

    all_tokens = np.empty(total_tokens, dtype=DTYPE)
    pos = 0

    for i, chunk in enumerate(chunks):
        chunk_path = litdata_dir / chunk["filename"]
        n_blocks = chunk["chunk_size"]
        header_ints = 1 + n_blocks + 1
        header_bytes = header_ints * 4

        with open(chunk_path, "rb") as f:
            f.seek(header_bytes)
            data = np.fromfile(f, dtype=DTYPE, count=chunk["dim"])

        all_tokens[pos:pos + len(data)] = data
        pos += len(data)

        if (i + 1) % 50 == 0 or i == len(chunks) - 1:
            print(f"    Read {i+1}/{len(chunks)} chunks ({pos:,} tokens)")

    return all_tokens[:pos]


def split_documents(token_stream):
    """Split token stream by EOS token (0) into individual documents."""
    eos_positions = np.where(token_stream == EOS_TOKEN_ID)[0]
    docs = []
    start = 0
    for eos_pos in eos_positions:
        if eos_pos > start:
            docs.append(token_stream[start:eos_pos])
        start = eos_pos + 1
    if start < len(token_stream):
        docs.append(token_stream[start:])
    return docs


def write_litdata_chunks(output_dir, token_stream, config):
    """Write token stream as litdata chunks, returns index metadata."""
    os.makedirs(output_dir, exist_ok=True)

    dtype_size = DTYPE().itemsize
    tokens_per_chunk = CHUNK_BYTES_TARGET // dtype_size
    tokens_per_chunk = (tokens_per_chunk // BLOCK_SIZE) * BLOCK_SIZE

    chunks_metadata = []
    pos = 0
    chunk_idx = 0

    while pos < len(token_stream):
        remaining = len(token_stream) - pos
        chunk_tokens = min(tokens_per_chunk, remaining)
        num_blocks = chunk_tokens // BLOCK_SIZE
        if num_blocks == 0:
            break
        actual_tokens = num_blocks * BLOCK_SIZE

        chunk_data = token_stream[pos:pos + actual_tokens]
        filename = f"chunk-0-{chunk_idx}.bin"
        filepath = os.path.join(output_dir, filename)

        # Header: [num_items(uint32)] + [offsets 0..num_blocks(uint32)]
        header_num_items = np.array([num_blocks], dtype=np.uint32)
        offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * dtype_size)
        header = np.concatenate([header_num_items, offsets])

        with open(filepath, "wb") as f:
            header.tofile(f)
            chunk_data.tofile(f)

        meta = {
            "chunk_bytes": int(header.nbytes + chunk_data.nbytes),
            "chunk_size": num_blocks,
            "dim": int(actual_tokens),
            "filename": filename,
        }
        chunks_metadata.append(meta)
        pos += actual_tokens
        chunk_idx += 1

        if chunk_idx % 25 == 0 or pos >= len(token_stream):
            print(f"    Written chunk {chunk_idx} ({pos:,}/{len(token_stream):,} tokens)")

    # Write index.json
    index = {
        "chunks": chunks_metadata,
        "config": config,
        "updated_at": str(time.time()),
    }
    with open(os.path.join(output_dir, "index.json"), "w") as f:
        json.dump(index, f, indent=2)

    return chunks_metadata


# ==============================================================================
# MAIN PROCESSING
# ==============================================================================

def process_litdata(input_dir, output_dir, name):
    print(f"\n{'='*65}")
    print(f"  PROCESSING: {name}")
    print(f"  Input:  {input_dir}")
    print(f"  Output: {output_dir}")
    print(f"{'='*65}")

    # 1. Read all tokens
    t0 = time.time()
    token_stream = read_all_tokens(input_dir)
    print(f"  Read {len(token_stream):,} tokens in {time.time()-t0:.1f}s")

    # 2. Split into documents
    t1 = time.time()
    doc_tokens = split_documents(token_stream)
    print(f"  Found {len(doc_tokens):,} documents in {time.time()-t1:.1f}s")
    del token_stream

    # 3. Decode all documents to text (batch for speed)
    t2 = time.time()
    print(f"  Decoding documents back to text...")
    raw_texts = []
    BATCH = 5000
    for i in range(0, len(doc_tokens), BATCH):
        batch = doc_tokens[i:i+BATCH]
        for doc in batch:
            text = tokenizer.decode(doc.tolist(), skip_special_tokens=False)
            raw_texts.append(text)
        done = min(i + BATCH, len(doc_tokens))
        if done % 100000 == 0 or done == len(doc_tokens):
            print(f"    Decoded {done:,}/{len(doc_tokens):,}")
    del doc_tokens
    print(f"  Decoded in {time.time()-t2:.1f}s")

    # 4. Clean each document
    t3 = time.time()
    print(f"  Cleaning {len(raw_texts):,} documents...")
    cleaned_texts = []
    dropped = 0
    for i, text in enumerate(raw_texts):
        result = clean_text(text)
        if result is not None:
            cleaned_texts.append(result)
        else:
            dropped += 1
        if (i + 1) % 200000 == 0 or i == len(raw_texts) - 1:
            print(f"    Processed {i+1:,}/{len(raw_texts):,} | kept={len(cleaned_texts):,} | dropped={dropped:,}")
    del raw_texts
    print(f"  Cleaning done in {time.time()-t3:.1f}s")
    print(f"  Kept {len(cleaned_texts):,} docs | Dropped {dropped:,} ({dropped/(max(dropped+len(cleaned_texts),1))*100:.1f}%)")

    # 5. Re-tokenize cleaned texts
    t4 = time.time()
    print(f"  Re-tokenizing {len(cleaned_texts):,} documents...")
    all_token_ids = []
    total_new_tokens = 0
    ENCODE_BATCH = 10000
    for i in range(0, len(cleaned_texts), ENCODE_BATCH):
        batch = cleaned_texts[i:i+ENCODE_BATCH]
        encoded = tokenizer.encode_batch(batch, add_special_tokens=False)
        for enc in encoded:
            ids = enc.ids
            all_token_ids.extend(ids)
            all_token_ids.append(EOS_TOKEN_ID)
            total_new_tokens += len(ids) + 1
        done = min(i + ENCODE_BATCH, len(cleaned_texts))
        if done % 200000 == 0 or done == len(cleaned_texts):
            print(f"    Tokenized {done:,}/{len(cleaned_texts):,} ({total_new_tokens:,} tokens)")
    del cleaned_texts
    print(f"  Tokenized in {time.time()-t4:.1f}s")
    print(f"  New total: {total_new_tokens:,} tokens")

    # 6. Convert to numpy and write chunks
    t5 = time.time()
    print(f"  Building token stream array...")
    new_stream = np.array(all_token_ids, dtype=DTYPE)
    del all_token_ids

    with open(input_dir / "index.json") as f:
        config = json.load(f)["config"]

    print(f"  Writing litdata chunks...")
    chunks = write_litdata_chunks(str(output_dir), new_stream, config)
    print(f"  Written {len(chunks)} chunks in {time.time()-t5:.1f}s")

    total_in = sum(c["dim"] for c in json.load(open(input_dir / "index.json"))["chunks"])
    total_out = sum(c["dim"] for c in chunks)
    print(f"\n  SUMMARY for {name}:")
    print(f"    Input tokens:  {total_in:,}")
    print(f"    Output tokens: {total_out:,}")
    print(f"    Difference:    {total_in - total_out:,} ({(total_in-total_out)/total_in*100:.2f}% removed)")

    return total_in, total_out


if __name__ == "__main__":
    t_start = time.time()
    data_dir = ROOT / "Base" / "data"

    # Process litdata_3b
    orig_3b, clean_3b = process_litdata(
        data_dir / "litdata_3b",
        data_dir / "litdata_3b_clean",
        "litdata_3b (General Knowledge)",
    )

    # Process litdata_english
    orig_en, clean_en = process_litdata(
        data_dir / "litdata_english",
        data_dir / "litdata_english_clean",
        "litdata_english (English Knowledge)",
    )

    # Final report
    print(f"\n{'='*65}")
    print(f"  FINAL REPORT")
    print(f"{'='*65}")
    print(f"  litdata_3b:      {orig_3b:>15,} -> {clean_3b:>15,} tokens")
    print(f"  litdata_english: {orig_en:>15,} -> {clean_en:>15,} tokens")
    print(f"  ---------------------------------------------------------")
    total_orig = orig_3b + orig_en
    total_clean = clean_3b + clean_en
    print(f"  TOTAL:           {total_orig:>15,} -> {total_clean:>15,} tokens")
    print(f"  Removed:         {total_orig - total_clean:,} ({(total_orig-total_clean)/total_orig*100:.2f}%)")
    print(f"\n  Total time: {time.time()-t_start:.0f}s")
    print(f"\n  Clean data ready at:")
    print(f"    {data_dir / 'litdata_3b_clean'}")
    print(f"    {data_dir / 'litdata_english_clean'}")
    print(f"\n  Update your training configs to point to the _clean directories!")