#!/usr/bin/env python3
"""Extract and combine all training data."""

import re
from pathlib import Path

datasets_dir = Path("C:\\MONIKA\\datasets")
output_file = Path("C:\\MONIKA\\training_sentences.txt")

all_sentences = []

# Process Alice in Wonderland
alice = datasets_dir / "alice_wonderland.txt"
if alice.exists():
    print(f"Processing {alice.name}...")
    text = alice.read_text(encoding='utf-8', errors='ignore')
    if "*** START" in text:
        text = text.split("*** START", 1)[1]
    if "*** END" in text:
        text = text.split("*** END", 1)[0]
    sentences = re.split(r'[.!?]+\s+', text)
    sentences = [s.strip() for s in sentences if 30 < len(s.strip()) < 200]
    all_sentences.extend(sentences[:500])
    print(f"  Added {min(500, len(sentences))} sentences")

# Process Wikipedia
wikipedia = datasets_dir / "wikipedia.txt"
if wikipedia.exists():
    print(f"Processing {wikipedia.name}...")
    text = wikipedia.read_text(encoding='utf-8', errors='ignore')
    # Split into sentences
    sentences = re.split(r'[.!?]+\s+', text)
    sentences = [s.strip() for s in sentences if 30 < len(s.strip()) < 200]
    all_sentences.extend(sentences[:2000])  # Take first 2000
    print(f"  Added {min(2000, len(sentences))} sentences")

# Process Pride and Prejudice
pride = datasets_dir / "pride_and_prejudice.txt"
if pride.exists():
    print(f"Processing {pride.name}...")
    text = pride.read_text(encoding='utf-8', errors='ignore')
    if "*** START" in text:
        text = text.split("*** START", 1)[1]
    if "*** END" in text:
        text = text.split("*** END", 1)[0]
    sentences = re.split(r'[.!?]+\s+', text)
    sentences = [s.strip() for s in sentences if 30 < len(s.strip()) < 200]
    all_sentences.extend(sentences[:1000])
    print(f"  Added {min(1000, len(sentences))} sentences")

# Process Sherlock Holmes
sherlock = datasets_dir / "sherlock_holmes.txt"
if sherlock.exists():
    print(f"Processing {sherlock.name}...")
    text = sherlock.read_text(encoding='utf-8', errors='ignore')
    if "*** START" in text:
        text = text.split("*** START", 1)[1]
    if "*** END" in text:
        text = text.split("*** END", 1)[0]
    sentences = re.split(r'[.!?]+\s+', text)
    sentences = [s.strip() for s in sentences if 30 < len(s.strip()) < 200]
    all_sentences.extend(sentences[:1000])
    print(f"  Added {min(1000, len(sentences))} sentences")

print(f"\nTotal sentences: {len(all_sentences)}")

# Write all to file
with open(output_file, 'w', encoding='utf-8') as f:
    for sent in all_sentences:
        f.write(sent + '\n')

print(f"Saved to: {output_file}")

# Create batches for training
batch_size = 50
batch_dir = Path("C:\\MONIKA\\training_batches")
batch_dir.mkdir(exist_ok=True)

for i in range(0, len(all_sentences), batch_size):
    batch = all_sentences[i:i+batch_size]
    batch_file = batch_dir / f"batch_{i//batch_size + 1:03d}.txt"
    with open(batch_file, 'w', encoding='utf-8') as f:
        for sent in batch:
            f.write(sent + '\n')

num_batches = (len(all_sentences) + batch_size - 1) // batch_size
print(f"\nCreated {num_batches} batch files in {batch_dir}")
print(f"Each batch has ~{batch_size} sentences")