scripts/download_data.py

#!/usr/bin/env python3
# pyright: reportAttributeAccessIssue=false
# pyright: reportArgumentType=false
# pyright: reportCallIssue=false
"""
Dataset download script for LexiMind.

Focus: Books, Academic Papers, Technical Writing
- NO news articles (overdone, dated)
- YES literary text, research, technical writing

Datasets:
- Goodreads + Gutenberg for book descriptions (back-cover style blurbs)
- arXiv for academic paper summarization (abstracts)
- Project Gutenberg for literary language modeling
- GoEmotions for emotion classification (28 labels)
- Custom topic classification: Fiction, Science, Technology, etc.

Key Design Decision:
Book "summarization" uses Goodreads descriptions which describe what the book
is ABOUT (like a back cover), not plot summaries. This trains the model to
answer "What is this book about?" rather than "What happens in this book?"

Usage:
    python scripts/download_data.py              # Download all
    python scripts/download_data.py --task summarization  # Just summarization
    python scripts/download_data.py --max-arxiv 50000

Author: Oliver Perrin
Date: January 2026
"""

from __future__ import annotations

import argparse
import json
import random
import re
from pathlib import Path
from typing import Any

from datasets import load_dataset  # type: ignore[import-untyped]
from tqdm import tqdm

# Output directory
OUTPUT_DIR = Path(__file__).parent.parent / "data" / "processed"

# ============== LABEL DEFINITIONS ==============

# 28 emotions from GoEmotions - works for all text types
EMOTION_LABELS = [
    "admiration", "amusement", "anger", "annoyance", "approval", "caring",
    "confusion", "curiosity", "desire", "disappointment", "disapproval",
    "disgust", "embarrassment", "excitement", "fear", "gratitude", "grief",
    "joy", "love", "nervousness", "optimism", "pride", "realization",
    "relief", "remorse", "sadness", "surprise", "neutral",
]

# New topic labels for books + papers + blogs
TOPIC_LABELS = [
    "Fiction",           # Novels, short stories, literary fiction
    "Science",           # Physics, chemistry, biology, nature
    "Technology",        # CS, engineering, programming, AI/ML
    "Philosophy",        # Ethics, logic, metaphysics, epistemology
    "History",           # Historical texts, biographies, memoirs
    "Psychology",        # Mind, behavior, self-help, mental health
    "Business",          # Economics, finance, entrepreneurship
    "Arts",              # Music, visual arts, film, architecture
]

# arXiv category → our topic mapping
ARXIV_CATEGORY_MAP = {
    # Computer Science
    "cs.AI": "Technology", "cs.CL": "Technology", "cs.CV": "Technology",
    "cs.LG": "Technology", "cs.NE": "Technology", "cs.RO": "Technology",
    "cs.SE": "Technology", "cs.PL": "Technology", "cs.DB": "Technology",
    "cs.DS": "Technology", "cs.CR": "Technology", "cs.DC": "Technology",
    "cs.HC": "Technology", "cs.IR": "Technology", "cs.IT": "Technology",
    "cs.MA": "Technology", "cs.MM": "Technology", "cs.NI": "Technology",
    "cs.OS": "Technology", "cs.PF": "Technology", "cs.SY": "Technology",
    # Physics
    "physics": "Science", "astro-ph": "Science", "cond-mat": "Science",
    "gr-qc": "Science", "hep-ex": "Science", "hep-lat": "Science",
    "hep-ph": "Science", "hep-th": "Science", "math-ph": "Science",
    "nlin": "Science", "nucl-ex": "Science", "nucl-th": "Science",
    "quant-ph": "Science",
    # Math
    "math": "Science",
    # Biology/Medicine
    "q-bio": "Science", "stat": "Science",
    # Economics/Finance
    "econ": "Business", "q-fin": "Business",
    # Electrical Engineering
    "eess": "Technology",
}

# Gutenberg subject → our topic mapping
GUTENBERG_SUBJECT_MAP = {
    "fiction": "Fiction", "novel": "Fiction", "stories": "Fiction",
    "poetry": "Arts", "drama": "Arts", "plays": "Arts",
    "science": "Science", "physics": "Science", "chemistry": "Science",
    "biology": "Science", "nature": "Science", "astronomy": "Science",
    "philosophy": "Philosophy", "ethics": "Philosophy", "logic": "Philosophy",
    "history": "History", "biography": "History", "memoir": "History",
    "psychology": "Psychology", "mind": "Psychology",
    "economics": "Business", "business": "Business", "finance": "Business",
    "art": "Arts", "music": "Arts", "architecture": "Arts",
    "technology": "Technology", "engineering": "Technology",
}


def write_jsonl(records: list[dict[str, Any]], path: Path, desc: str = "Writing") -> None:
    """Write records to JSONL file."""
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        for record in tqdm(records, desc=desc, leave=False):
            f.write(json.dumps(record, ensure_ascii=False) + "\n")
    print(f"  ✓ {len(records):,} samples → {path}")


# ============== ENGLISH LANGUAGE FILTER ==============

# Common English words for detection
ENGLISH_WORDS = {
    "the", "and", "of", "to", "a", "in", "that", "is", "was", "he", "she", "it",
    "for", "with", "as", "his", "her", "they", "be", "at", "on", "have", "had",
    "this", "but", "not", "from", "by", "or", "an", "said", "were", "been",
    "would", "could", "which", "their", "there", "what", "when", "who", "will",
    "more", "if", "no", "out", "so", "up", "into", "than", "them", "can", "only",
    "other", "new", "some", "very", "just", "over", "such", "also", "its", "then",
}

# Non-English language patterns
NON_ENGLISH_PATTERNS = [
    # French
    r"\b(le|la|les|un|une|des|du|et|est|sont|dans|pour|avec|sur|qui|que|ce|cette|nous|vous|ils|elles|je|tu|il|elle|être|avoir)\b",
    # German
    r"\b(der|die|das|ein|eine|und|ist|nicht|mit|von|zu|den|dem|auf|für|als|auch|oder|nach|bei|nur|noch|wie|mehr|aber|wenn|hat|kann|ich|sie|er|wir|ihr|es|sich|sein)\b",
    # Spanish
    r"\b(el|la|los|las|un|una|que|por|para|con|del|al|es|en|se|no|más|como|pero|su|sus|le|lo|te|me|nos)\b",
    # Italian
    r"\b(il|lo|la|gli|le|che|per|con|del|della|di|da|non|sono|anche|più|ma|se|mi|ti|ci)\b",
    # Latin
    r"\b(et|in|ad|cum|de|ex|per|pro|sub|ab|ante|post|inter|contra|super|trans|apud)\b",
]

# ============== TEXT QUALITY FILTERS ==============

# Patterns that indicate garbage/metadata text
GARBAGE_PATTERNS = [
    r"^Page \d+:",           # Page corrections
    r"changed to",           # Errata
    r"Punctuation has been", # Editorial notes
    r"^\[.*\]$",             # Bracketed notes
    r"^Note\.?[-—]",         # Notes
    r"^follows:",            # "as follows:"
    r"CHAPTER [IVXLC]+\.",   # Chapter headers only
    r"^\*\*\*",              # Project Gutenberg markers
    r"^End of.*Project",     # End markers
    r"^Produced by",         # Production credits
    r"transcriber",          # Transcriber notes
    r"eBook",                # eBook references
    r"©|copyright",          # Copyright notices
    r"^INDEX",               # Index pages
    r"^\d+\.\s+\w+,\s+\d+",  # Index entries like "1. Name, 234"
    r"(syn\.|var\.|sp\.)",   # Botanical abbreviations
    r"[A-Z][a-z]+aceae",     # Botanical family names
    r"\(\s*syn\s+",          # Synonym references
]

# Patterns that indicate technical manuals/instructions (not narrative)
TECHNICAL_PATTERNS = [
    r"\d+\.\s+It\s+(is|has|can)",  # Numbered features "1. It is a..."
    r"^\d+(st|nd|rd|th)\.",        # "1st. 2nd. 3rd."
    r"Mesh\.?\s*\d+",              # Mesh sizes (pottery)
    r"\d+\s*(oz|lb|kg|g|ml|mm|cm|inch)",  # Measurements
    r"Parts?\s*:?\s*\d+",          # "Parts: 50"
    r"Method of Using",            # Instructions
    r"How to\s+\w+",               # How-to guides
    r"Step\s+\d+",                 # Step-by-step
    r"wire.*address",              # Business instructions
    r"orders?\s+should\s+be",      # Order instructions  
    r"specifications?",            # Technical specs
    r"(Front|Back)\s+Focus",       # Camera terms
    r"Rack and Pinion",            # Mechanical terms
]

# Shakespeare and plays to exclude (model hallucinates on Early Modern English)
EXCLUDED_TITLES = {
    # Shakespeare
    "King Lear", "Hamlet", "Macbeth", "Othello", "Romeo and Juliet",
    "A Midsummer Night's Dream", "The Tempest", "Julius Caesar",
    "The Merchant of Venice", "Twelfth Night", "Much Ado About Nothing",
    "As You Like It", "The Taming of the Shrew", "Antony and Cleopatra",
    "Coriolanus", "Cymbeline", "Timon of Athens", "Troilus and Cressida",
    "Measure for Measure", "All's Well That Ends Well", "Pericles",
    "The Winter's Tale", "The Comedy of Errors", "Two Gentlemen of Verona",
    "Love's Labour's Lost", "The Merry Wives of Windsor", "Henry IV",
    "Henry V", "Henry VI", "Henry VIII", "Richard II", "Richard III",
    "King John", "Titus Andronicus",
    # French plays
    "Tartuffe", "Phaedra", "Cyrano de Bergerac", "Cyrano De Bergerac",
    "Le Misanthrope", "The School for Wives", "The Miser", "The Imaginary Invalid",
    "Andromaque", "Britannicus", "Bérénice", "Le Cid",
    # Greek/Roman plays
    "Oedipus Rex", "Oedipus the King", "Antigone", "Electra", "Medea",
    "The Bacchae", "The Oresteia", "Agamemnon", "Prometheus Bound",
    # Other classic plays
    "The Importance of Being Earnest", "Pygmalion", "Doctor Faustus",
    "Waiting for Godot", "Death of a Salesman", "A Streetcar Named Desire",
    "The Glass Menagerie", "Our Town", "Long Day's Journey Into Night",
    "Who's Afraid of Virginia Woolf", "The Crucible", "Cat on a Hot Tin Roof",
    # Verse/poetic epics
    "Idylls of the King", "Paradise Lost", "Paradise Regained",
    "The Divine Comedy", "Inferno", "Purgatorio", "Paradiso",
    "The Faerie Queene", "Beowulf",
}


def is_technical_manual(text: str) -> bool:
    """Check if text appears to be a technical manual/instructions."""
    for pattern in TECHNICAL_PATTERNS:
        if re.search(pattern, text, re.IGNORECASE | re.MULTILINE):
            return True
    return False


def is_quality_text(text: str) -> bool:
    """Check if text is quality content (not metadata/garbage)."""
    # Check for garbage patterns
    for pattern in GARBAGE_PATTERNS:
        if re.search(pattern, text, re.IGNORECASE | re.MULTILINE):
            return False
    
    # Reject technical manuals/instructions
    if is_technical_manual(text):
        return False
    
    # Must have reasonable length
    if len(text) < 300:
        return False
    
    # Must have sentences (not just fragments)
    sentences = re.split(r'[.!?]+', text)
    if len(sentences) < 4:
        return False
    
    # Check for too many special characters
    special_ratio = len(re.findall(r'[^\w\s.,!?\'"()-]', text)) / max(len(text), 1)
    if special_ratio > 0.08:
        return False
    
    return True


def is_excluded_title(title: str) -> bool:
    """Check if title should be excluded (plays, epics, etc.)."""
    title_lower = title.lower()
    return any(excluded.lower() in title_lower for excluded in EXCLUDED_TITLES)


def is_play_text(text: str) -> bool:
    """Check if text appears to be a play/dramatic format."""
    play_patterns = [
        r"^(SCENE|ACT|Enter|Exit|Exeunt)\s",
        r"^\[.*\]$",  # Stage directions
        r"^[A-Z]{2,}\.\s",  # Character names like "HAMLET."
        r"Alarum|Flourish|Sennet",  # Stage directions
    ]
    lines = text.split('\n')[:10]
    play_indicators = 0
    for line in lines:
        for pattern in play_patterns:
            if re.search(pattern, line.strip()):
                play_indicators += 1
    return play_indicators >= 2


def is_english_text(text: str, min_ratio: float = 0.08, max_foreign: int = 5) -> bool:
    """
    Check if text is primarily English.
    
    Args:
        text: Text to check
        min_ratio: Minimum ratio of common English words
        max_foreign: Maximum number of foreign word matches before rejecting
    
    Returns:
        True if text appears to be English
    """
    if not text or len(text) < 100:
        return False
    
    text_lower = text.lower()
    words = text_lower.split()
    
    if len(words) < 20:
        return False
    
    # Check for excessive non-English words
    for pattern in NON_ENGLISH_PATTERNS:
        matches = len(re.findall(pattern, text_lower))
        if matches > max_foreign:
            return False
    
    # Check for sufficient English words
    english_count = sum(1 for w in words if w.strip(".,!?;:'\"") in ENGLISH_WORDS)
    ratio = english_count / len(words)
    
    return ratio >= min_ratio


def normalize_title(title: str) -> str:
    """Normalize a book title for matching."""
    # Remove common prefixes/suffixes
    title = re.sub(r'^(The|A|An)\s+', '', title, flags=re.IGNORECASE)
    title = re.sub(r'\s*\([^)]*\)\s*', '', title)  # Remove parentheticals
    title = re.sub(r'\s*:.+$', '', title)  # Remove subtitles
    title = re.sub(r'[^\w\s]', '', title)  # Remove punctuation
    return title.lower().strip()


# ============== SUMMARIZATION: BOOKS + ARXIV ==============

def download_goodreads_descriptions() -> dict[str, dict]:
    """
    Download Goodreads book descriptions - back-cover style blurbs.
    
    These are "what the book is about" descriptions, not plot summaries.
    Returns dict mapping normalized title -> {title, description}
    """
    print("\n📚 Loading Goodreads book descriptions...")
    
    descriptions = {}
    
    # Try multiple sources
    datasets_to_try = [
        "booksouls/goodreads-book-descriptions",
        "Skelebor/book_titles_and_descriptions_en_clean",
    ]
    
    for ds_name in datasets_to_try:
        try:
            print(f"    Loading {ds_name}...")
            ds = load_dataset(ds_name, split="train")
            
            for item in tqdm(ds, desc="Goodreads", leave=False):
                title = item.get("title", "")
                description = item.get("description", "")
                
                if not title or not description:
                    continue
                
                # Skip very short descriptions (not useful for training)
                if len(description) < 100:
                    continue
                
                # Skip very long descriptions (truncate later)
                if len(description) > 2000:
                    description = description[:2000]
                
                # Skip plays and excluded titles
                if is_excluded_title(title):
                    continue
                
                # Skip non-English descriptions
                if not is_english_text(description):
                    continue
                
                norm_title = normalize_title(title)
                if norm_title and norm_title not in descriptions:
                    descriptions[norm_title] = {
                        "title": title,
                        "description": description,
                    }
            
            print(f"    Loaded {len(descriptions):,} descriptions from {ds_name}")
        except Exception as e:
            print(f"    {ds_name} failed: {e}")
    
    print(f"    Total: {len(descriptions):,} unique book descriptions")
    return descriptions


def download_book_descriptions(
    goodreads_descriptions: dict[str, dict],
    max_samples: int = 20000
) -> list[dict[str, Any]]:
    """
    Download book description data by matching Gutenberg texts with Goodreads descriptions.
    
    This gives us (book_excerpt, book_description) training pairs where descriptions
    are back-cover style "what is this book about" blurbs, not plot summaries.
    """
    print("\n📖 Matching Gutenberg books with Goodreads descriptions...")
    
    try:
        gutenberg = load_dataset("sedthh/gutenberg_english", split="train")
    except Exception:
        gutenberg = load_dataset("pg19", split="train")
    
    records: list[dict[str, Any]] = []
    matched_titles = set()
    skipped_quality = 0
    skipped_play = 0
    
    indices = list(range(len(gutenberg)))
    random.shuffle(indices)
    
    for i in tqdm(indices, desc="Matching books", leave=False):
        if len(records) >= max_samples:
            break
        
        item = gutenberg[i]
        text = item.get("TEXT", "") or item.get("text", "")
        metadata_raw = item.get("METADATA", "") or "{}"
        
        # Parse metadata
        try:
            metadata = json.loads(metadata_raw) if isinstance(metadata_raw, str) else metadata_raw
        except (json.JSONDecodeError, TypeError):
            metadata = {}
        
        # Get title
        title = metadata.get("title", "") if isinstance(metadata, dict) else ""
        if not title:
            continue
        
        # Check if we have a Goodreads description for this book
        norm_title = normalize_title(title)
        if norm_title not in goodreads_descriptions:
            continue
        
        # Skip if already matched this book
        if norm_title in matched_titles:
            continue
        
        goodreads_data = goodreads_descriptions[norm_title]
        
        # Skip plays and excluded titles
        if is_excluded_title(title):
            skipped_play += 1
            continue
        
        if not text or len(text) < 2000:
            continue
        
        # Get a clean excerpt from the book (skip front matter)
        paragraphs = re.split(r'\n\s*\n', text)
        excerpt_parts = []
        total_len = 0
        
        for para in paragraphs[10:]:  # Skip front matter
            para = para.strip()
            if len(para) < 100:
                continue
            
            # Quality check on paragraph
            if not is_english_text(para):
                continue
            if is_play_text(para):
                skipped_play += 1
                break
            if not is_quality_text(para) and len(para) > 300:
                skipped_quality += 1
                continue
            
            excerpt_parts.append(para)
            total_len += len(para)
            
            if total_len >= 3000:
                break
        
        if total_len < 1000:
            continue
        
        book_excerpt = "\n\n".join(excerpt_parts)[:4000]
        matched_titles.add(norm_title)
        
        records.append({
            "source": book_excerpt,
            "summary": goodreads_data["description"][:800],  # Back-cover blurbs are shorter
            "type": "literary",
            "title": goodreads_data["title"],
        })
    
    print(f"    Matched {len(records):,} books with descriptions")
    print(f"    Skipped: {skipped_quality} quality, {skipped_play} plays")
    
    return records


# Keep BookSum for additional literary training (chapter summaries are still useful)
def download_booksum(max_samples: int = 20000) -> list[dict[str, Any]]:
    """Download BookSum - literary chapter summarization (English only, quality filtered).
    
    Note: These are chapter-level plot summaries, useful as supplementary training data.
    The primary book training comes from Goodreads descriptions (back-cover style).
    """
    print("\n📖 Loading BookSum (supplementary literary data)...")
    
    all_records: list[dict[str, Any]] = []
    booksum = load_dataset("kmfoda/booksum")
    
    for split_name in booksum.keys():
        split = str(split_name)
        data = booksum[split_name]
        limit = max_samples if "train" in split else max_samples // 10
        indices = random.sample(range(len(data)), min(len(data), limit))
        
        records = []
        skipped_language = 0
        skipped_excluded = 0
        skipped_play = 0
        
        for i in tqdm(indices, desc=f"BookSum {split}", leave=False):
            item = data[i]
            chapter = item.get("chapter", "")
            summary = item.get("summary_text") or item.get("summary", "")
            
            # Extract book title from book_id (e.g., "The Last of the Mohicans.chapters 1-2")
            book_id = item.get("book_id", "")
            book_title = book_id.split(".")[0] if "." in book_id else book_id
            chapter_name = item.get("summary_id", "") or item.get("summary_name", "")
            
            if not (chapter and summary and len(chapter) > 300):
                continue
            
            # Filter: excluded titles (Shakespeare, plays, etc.)
            if is_excluded_title(book_title):
                skipped_excluded += 1
                continue
            
            # Filter: play text format
            if is_play_text(chapter):
                skipped_play += 1
                continue
            
            # Filter: English only
            if not is_english_text(chapter):
                skipped_language += 1
                continue
            
            # Filter: quality text
            if not is_quality_text(chapter):
                continue
            
            records.append({
                "source": chapter[:4000],
                "summary": summary,
                "type": "literary",
                "split": split,
                "title": book_title,
                "chapter": chapter_name,
            })
        all_records.extend(records)
        print(f"    {split}: {len(records):,} (skipped {skipped_language} non-English, {skipped_excluded} excluded, {skipped_play} plays)")
    
    return all_records


def clean_arxiv_text(text: str) -> str:
    """Clean arXiv LaTeX-style text to make it more readable."""
    import re
    # Remove LaTeX math placeholders
    text = re.sub(r'@xmath\d+', '', text)
    text = re.sub(r'@xcite', '', text)
    # Remove excessive whitespace
    text = re.sub(r'\s+', ' ', text)
    # Remove LaTeX commands
    text = re.sub(r'\\[a-zA-Z]+\{[^}]*\}', '', text)
    text = re.sub(r'\\[a-zA-Z]+', '', text)
    return text.strip()


def extract_paper_title(abstract: str) -> str:
    """Extract a meaningful title from the first sentence of an abstract."""
    # Clean the abstract first
    abstract = clean_arxiv_text(abstract)
    
    # Get the first sentence (up to first period, question mark, or newline)
    first_sentence = re.split(r'[.!?\n]', abstract)[0].strip()
    
    # Truncate if too long
    if len(first_sentence) > 100:
        # Try to cut at a natural word boundary
        first_sentence = first_sentence[:100].rsplit(' ', 1)[0] + '...'
    
    # Capitalize first letter
    if first_sentence:
        first_sentence = first_sentence[0].upper() + first_sentence[1:]
    
    return first_sentence or "Untitled Paper"


def download_arxiv_summarization(max_samples: int = 50000) -> list[dict[str, Any]]:
    """
    Download arXiv papers for academic summarization only (English only).
    Note: This dataset doesn't have categories, so can't be used for topic classification.
    
    Returns: summarization_records
    """
    print("\n🎓 Loading arXiv (academic papers for summarization)...")
    
    print("  Loading dataset (this may take a minute)...")
    arxiv = load_dataset("ccdv/arxiv-summarization", split="train")
    
    summ_records: list[dict[str, Any]] = []
    skipped_language = 0
    
    indices = list(range(len(arxiv)))
    random.shuffle(indices)
    
    print("  Processing papers...")
    for i in tqdm(indices[:max_samples * 2], desc="arXiv", leave=False):
        if len(summ_records) >= max_samples:
            break
            
        item = arxiv[i]
        
        # Get abstract and article
        abstract = item.get("abstract", "")
        article = item.get("article", "")
        
        if not abstract or len(abstract) < 100:
            continue
        
        # Clean LaTeX artifacts
        abstract = clean_arxiv_text(abstract)
        article = clean_arxiv_text(article)
        
        # Skip if still has too many weird characters after cleaning
        if '@' in abstract or '@' in article[:500]:
            continue
        
        # Filter: English only
        if not is_english_text(article[:1000]):
            skipped_language += 1
            continue
        
        # Summarization: article → abstract
        if article and len(article) > 500:
            # Extract title from abstract
            paper_title = extract_paper_title(abstract)
            
            summ_records.append({
                "source": article[:4000],
                "summary": abstract,
                "type": "academic",
                "title": paper_title,
            })
    
    print(f"    Summarization: {len(summ_records):,} (skipped {skipped_language} non-English)")
    
    return summ_records


def download_topics_from_datasets(max_samples: int = 50000) -> list[dict[str, Any]]:
    """
    Download topic classification data from multiple sources with real categories.
    
    Sources:
    - 20 Newsgroups (classic topic classification)
    - Wikipedia (article categories)
    """
    print("\n📂 Loading topic classification datasets...")
    
    records: list[dict[str, Any]] = []
    
    # 20 Newsgroups - classic topic dataset
    print("  Loading 20 Newsgroups...")
    try:
        newsgroups = load_dataset("SetFit/20_newsgroups", split="train")
        
        # Map 20 newsgroups categories to our 8 topics
        newsgroup_map = {
            # Science
            "sci.crypt": "Science", "sci.electronics": "Science",
            "sci.med": "Science", "sci.space": "Science",
            # Technology  
            "comp.graphics": "Technology", "comp.os.ms-windows.misc": "Technology",
            "comp.sys.ibm.pc.hardware": "Technology", "comp.sys.mac.hardware": "Technology",
            "comp.windows.x": "Technology",
            # Philosophy/Religion
            "alt.atheism": "Philosophy", "soc.religion.christian": "Philosophy",
            "talk.religion.misc": "Philosophy",
            # History/Politics
            "talk.politics.guns": "History", "talk.politics.mideast": "History",
            "talk.politics.misc": "History",
            # Business
            "misc.forsale": "Business",
            # Sports/Recreation
            "rec.autos": "Arts", "rec.motorcycles": "Arts",
            "rec.sport.baseball": "Arts", "rec.sport.hockey": "Arts",
        }
        
        for item in tqdm(newsgroups, desc="20 Newsgroups", leave=False):
            if len(records) >= max_samples:
                break
            label_name = item.get("label_text", "")
            text = item.get("text", "")
            
            if label_name in newsgroup_map and text and len(text) > 100:
                records.append({
                    "text": text[:1500],
                    "topic": newsgroup_map[label_name],
                    "source": "newsgroups",
                })
        
        print(f"    20 Newsgroups: {len(records):,}")
    except Exception as e:
        print(f"    20 Newsgroups failed: {e}")
    
    # Add from Gutenberg for Fiction
    gutenberg_topics = download_gutenberg_topics(max_samples // 4)
    records.extend(gutenberg_topics)
    
    # Add from scientific papers abstract dataset for more Science/Tech
    print("  Loading scientific papers...")
    try:
        sci_papers = load_dataset("scientific_papers", "arxiv", split="train", streaming=True)
        sci_count = 0
        for item in tqdm(sci_papers, desc="Scientific papers", leave=False, total=max_samples//4):
            if sci_count >= max_samples // 4:
                break
            abstract = item.get("abstract", "")
            if abstract and len(abstract) > 100:
                # Alternate between Science and Technology
                topic = "Science" if sci_count % 2 == 0 else "Technology"
                records.append({
                    "text": abstract[:1500],
                    "topic": topic,
                    "source": "scientific_papers",
                })
                sci_count += 1
        print(f"    Scientific papers: {sci_count:,}")
    except Exception as e:
        print(f"    Scientific papers failed: {e}")
    
    return records


def download_summarization(max_books: int = 20000, max_arxiv: int = 50000) -> None:
    """Download all summarization data (books + arxiv, NO news).
    
    Book data now uses Goodreads descriptions (back-cover blurbs) instead of
    plot summaries. This trains the model to describe "what the book is about"
    rather than summarizing the plot.
    """
    print("\n📝 Downloading Summarization Data...")
    out_dir = OUTPUT_DIR / "summarization"
    
    all_records: list[dict[str, Any]] = []
    
    # Goodreads descriptions - primary book training data (back-cover style)
    goodreads_descriptions = download_goodreads_descriptions()
    book_records = download_book_descriptions(goodreads_descriptions, max_books)
    all_records.extend(book_records)
    
    # Optional: Add some BookSum for additional literary variety
    # These are chapter summaries, not back-cover style, so keep limited
    # booksum_records = download_booksum(max_books // 4)
    # all_records.extend(booksum_records)
    
    # arXiv - academic (abstracts are already "what is this paper about")
    arxiv_summ = download_arxiv_summarization(max_arxiv)
    all_records.extend(arxiv_summ)
    
    # Shuffle and split
    random.shuffle(all_records)
    
    # Split by original split if available, else 90/5/5
    train_records = [r for r in all_records if r.get("split", "train") == "train" or "split" not in r]
    val_records = [r for r in all_records if r.get("split") == "validation"]
    test_records = [r for r in all_records if r.get("split") == "test"]
    
    # If no split info, do 90/5/5
    if len(val_records) < 100:
        n = len(train_records)
        random.shuffle(train_records)
        val_records = train_records[int(n*0.9):int(n*0.95)]
        test_records = train_records[int(n*0.95):]
        train_records = train_records[:int(n*0.9)]
    
    # Remove split key before saving
    for r in train_records + val_records + test_records:
        r.pop("split", None)
    
    write_jsonl(train_records, out_dir / "train.jsonl", "train")
    write_jsonl(val_records, out_dir / "validation.jsonl", "val")
    write_jsonl(test_records, out_dir / "test.jsonl", "test")
    
    # Print breakdown
    literary_count = sum(1 for r in train_records + val_records + test_records if r.get("type") == "literary")
    academic_count = sum(1 for r in train_records + val_records + test_records if r.get("type") == "academic")
    print(f"\n  ✓ Total summarization: {len(train_records) + len(val_records) + len(test_records):,}")
    print(f"    Literary (book descriptions): {literary_count:,}")
    print(f"    Academic (paper abstracts): {academic_count:,}")


# ============== TOPIC CLASSIFICATION ==============

def download_topics(max_samples: int = 50000) -> None:
    """
    Download topic classification data from multiple sources.
    
    Sources:
    - 20 Newsgroups (classic topic dataset)
    - Gutenberg books (Fiction)
    - Scientific papers (Science, Technology)
    """
    print("\n📂 Downloading Topic Classification...")
    out_dir = OUTPUT_DIR / "topic"
    
    # Get topic records from various sources
    all_records = download_topics_from_datasets(max_samples)
    
    # Balance topics
    topic_counts: dict[str, list] = {t: [] for t in TOPIC_LABELS}
    for r in all_records:
        topic = r.get("topic")
        if topic in topic_counts:
            topic_counts[topic].append(r)
    
    # Print distribution before balancing
    print("\n  Topic distribution (before balancing):")
    for topic, records in topic_counts.items():
        print(f"    {topic}: {len(records):,}")
    
    # Balance to min count (with some tolerance) - only from topics that have data
    counts_with_data = [len(v) for v in topic_counts.values() if v]
    if not counts_with_data:
        print("  ⚠️ No topic data found!")
        return
    
    min_count = min(counts_with_data)
    target_count = min(min_count, max_samples // len(TOPIC_LABELS))
    
    balanced: list[dict[str, Any]] = []
    for topic, records in topic_counts.items():
        if records:
            random.shuffle(records)
            balanced.extend(records[:target_count])
    
    random.shuffle(balanced)
    
    # Split 90/5/5
    n = len(balanced)
    train_records = balanced[:int(n*0.9)]
    val_records = balanced[int(n*0.9):int(n*0.95)]
    test_records = balanced[int(n*0.95):]
    
    write_jsonl(train_records, out_dir / "train.jsonl", "train")
    write_jsonl(val_records, out_dir / "validation.jsonl", "val")
    write_jsonl(test_records, out_dir / "test.jsonl", "test")
    
    # Save labels - only labels that have data
    used_labels = [t for t in TOPIC_LABELS if topic_counts.get(t)]
    (out_dir / "labels.json").write_text(json.dumps(used_labels, indent=2))
    print(f"\n  ✓ {len(used_labels)} topic labels with data: {used_labels}")


def download_gutenberg_topics(max_samples: int = 30000) -> list[dict[str, Any]]:
    """Extract topic-labeled samples from Gutenberg books (English only)."""
    print("\n📚 Loading Gutenberg for topic classification...")
    
    try:
        gutenberg = load_dataset("sedthh/gutenberg_english", split="train")
    except Exception:
        print("  Trying pg19...")
        gutenberg = load_dataset("pg19", split="train")
    
    records: list[dict[str, Any]] = []
    skipped_language = 0
    
    indices = list(range(len(gutenberg)))
    random.shuffle(indices)
    
    for i in tqdm(indices, desc="Gutenberg topics", leave=False):
        if len(records) >= max_samples:
            break
        
        item = gutenberg[i]
        text = item.get("TEXT", "") or item.get("text", "")
        metadata = item.get("METADATA", {}) or {}
        
        if not text or len(text) < 1000:
            continue
        
        # Try to determine topic from metadata
        subjects = ""
        if isinstance(metadata, dict):
            subjects = str(metadata.get("subjects", "")).lower()
            subjects += " " + str(metadata.get("subject", "")).lower()
            subjects += " " + str(metadata.get("category", "")).lower()
        
        topic = None
        for keyword, mapped_topic in GUTENBERG_SUBJECT_MAP.items():
            if keyword in subjects:
                topic = mapped_topic
                break
        
        # Default fiction for novels without clear subject
        if not topic and ("novel" in subjects or not subjects.strip()):
            topic = "Fiction"
        
        if topic:
            # Get a clean paragraph as sample
            paragraphs = re.split(r'\n\s*\n', text)
            for para in paragraphs[5:]:  # Skip front matter
                para = para.strip()
                if 200 < len(para) < 1500 and para.count('.') >= 2:
                    # Filter: English only
                    if not is_english_text(para):
                        skipped_language += 1
                        break
                    
                    records.append({
                        "text": para,
                        "topic": topic,
                        "source": "gutenberg",
                    })
                    break
    
    print(f"    Gutenberg topics: {len(records):,} (skipped {skipped_language} non-English)")
    return records


# ============== EMOTIONS (unchanged) ==============

def download_emotions() -> None:
    """Download GoEmotions for emotion classification."""
    print("\n😊 Downloading Emotions (GoEmotions)...")
    out_dir = OUTPUT_DIR / "emotion"
    
    ds = load_dataset("google-research-datasets/go_emotions", "simplified")
    
    for split_name in ds.keys():
        split = str(split_name)
        data = ds[split_name]
        
        records: list[dict[str, Any]] = []
        for item in tqdm(data, desc=split, leave=False):
            text = item.get("text", "")
            label_ids = item.get("labels", [])
            if text and label_ids:
                emotions = [EMOTION_LABELS[i] for i in label_ids if 0 <= i < len(EMOTION_LABELS)]
                if emotions:
                    records.append({"text": text, "emotions": emotions})
        write_jsonl(records, out_dir / f"{split}.jsonl", split)
    
    (out_dir / "labels.json").write_text(json.dumps(EMOTION_LABELS, indent=2))
    print(f"  ✓ {len(EMOTION_LABELS)} emotion labels saved")


# ============== GUTENBERG BOOKS (for language modeling) ==============

GUTENBERG_JUNK_PATTERNS = [
    r"Project Gutenberg", r"www\.gutenberg\.org", r"This ebook is for",
    r"Gutenberg License", r"^\*\*\* START OF", r"^\*\*\* END OF",
    r"Produced by", r"Transcriber's Note", r"TABLE OF CONTENTS",
    r"^\s*CHAPTER\s+[IVXLC\d]+", r"^\s*Chapter\s+[IVXLC\d]+",
    r"^\s*BOOK\s+[IVXLC\d]+", r"^\s*PREFACE\s*$", r"^\s*INTRODUCTION\s*$",
    r"E-text prepared by", r"Internet Archive", r"Distributed Proofreaders",
]
GUTENBERG_JUNK_REGEX = re.compile("|".join(GUTENBERG_JUNK_PATTERNS), re.IGNORECASE)


def is_clean_prose(text: str) -> bool:
    """Check if text is clean literary prose (English only)."""
    if len(text) < 300 or len(text) > 3000:
        return False
    if GUTENBERG_JUNK_REGEX.search(text):
        return False
    if text.count('.') < 2:
        return False
    uppercase_ratio = sum(1 for c in text if c.isupper()) / max(len(text), 1)
    if uppercase_ratio > 0.3:
        return False
    digit_ratio = sum(1 for c in text if c.isdigit()) / max(len(text), 1)
    if digit_ratio > 0.1:
        return False
    # English filter
    if not is_english_text(text):
        return False
    return True


def download_gutenberg(max_samples: int = 30000) -> None:
    """Download Gutenberg books for language modeling (English only)."""
    print("\n📚 Downloading Gutenberg Books (English only)...")
    out_dir = OUTPUT_DIR / "books"
    out_dir.mkdir(parents=True, exist_ok=True)
    
    try:
        gutenberg = load_dataset("sedthh/gutenberg_english", split="train")
    except Exception:
        gutenberg = load_dataset("pg19", split="train")
    
    records: list[dict[str, Any]] = []
    indices = list(range(len(gutenberg)))
    random.shuffle(indices)
    
    for i in tqdm(indices, desc="Books", leave=False):
        if len(records) >= max_samples:
            break
        
        item = gutenberg[i]
        text = item.get("TEXT", "") or item.get("text", "")
        metadata_raw = item.get("METADATA", "") or "{}"
        
        # Parse metadata - it's stored as JSON string
        try:
            metadata = json.loads(metadata_raw) if isinstance(metadata_raw, str) else metadata_raw
        except (json.JSONDecodeError, TypeError):
            metadata = {}
        
        # Extract title and author
        title = metadata.get("title", "") if isinstance(metadata, dict) else ""
        author = metadata.get("author", "") if isinstance(metadata, dict) else ""
        if not title:
            title = item.get("title", f"Unknown Book #{i}")
        
        if not text or len(text) < 1000:
            continue
        
        paragraphs = re.split(r'\n\s*\n', text)
        for para in paragraphs:
            para = para.strip()
            if is_clean_prose(para):
                records.append({
                    "text": para, 
                    "title": title, 
                    "author": author,
                    "type": "gutenberg"
                })
                if len(records) >= max_samples:
                    break
    
    random.shuffle(records)
    n = len(records)
    write_jsonl(records[:int(n*0.9)], out_dir / "train.jsonl", "train")
    write_jsonl(records[int(n*0.9):int(n*0.95)], out_dir / "validation.jsonl", "val")
    write_jsonl(records[int(n*0.95):], out_dir / "test.jsonl", "test")


# ============== MAIN ==============

def main() -> None:
    parser = argparse.ArgumentParser(description="Download LexiMind datasets")
    parser.add_argument(
        "--task",
        choices=["all", "summarization", "emotion", "topic", "gutenberg"],
        default="all",
        help="Dataset to download"
    )
    parser.add_argument("--max-books", type=int, default=40000, help="Max BookSum samples")
    parser.add_argument("--max-arxiv", type=int, default=50000, help="Max arXiv samples")
    parser.add_argument("--max-gutenberg", type=int, default=30000, help="Max Gutenberg chunks")
    parser.add_argument("--max-topics", type=int, default=50000, help="Max topic samples")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    args = parser.parse_args()
    
    random.seed(args.seed)
    
    print("=" * 60)
    print("LexiMind Dataset Download")
    print("Books + Academic Papers + Topic Classification")
    print("=" * 60)
    
    if args.task in ["all", "summarization"]:
        download_summarization(args.max_books, args.max_arxiv)
    if args.task in ["all", "emotion"]:
        download_emotions()
    if args.task in ["all", "topic"]:
        download_topics(args.max_topics)
    if args.task in ["all", "gutenberg"]:
        download_gutenberg(args.max_gutenberg)
    
    print("\n" + "=" * 60)
    print("✅ Download complete!")
    print("=" * 60)


if __name__ == "__main__":
    main()