gpu.py

import gradio as gr
import pandas as pd
import requests
import internetarchive
from datetime import datetime
import re
import os
import shutil
import time
import random
import json
import torch
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
import numpy as np
import nest_asyncio
import sys

# --- SYSTEM FIXES ---
try:
    nest_asyncio.apply()
except Exception as e:
    print(f"Warning: Could not apply nest_asyncio: {e}")

# --- CONFIGURATION ---
DATASET_DIR = "dataset_ml_final_v2"
BOOKS_DIR = os.path.join(DATASET_DIR, "books")
MODEL_DIR = "trained_models"
os.makedirs(MODEL_DIR, exist_ok=True)

# --- TOKENIZER & MODEL ---
TOKENIZER = None
MODEL = None
# Check for CUDA support for GPU, otherwise use CPU
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

try:
    from transformers import AutoTokenizer, AutoModelForSequenceClassification, get_linear_schedule_with_warmup, logging
    from torch.optim import AdamW

    logging.set_verbosity_error()

    print("Attempting to load Longformer Tokenizer...")
    TOKENIZER = AutoTokenizer.from_pretrained("allenai/longformer-base-4096")
    print("✅ Tokenizer loaded successfully.")
except Exception as e:
    print(f"⚠️ Tokenizer loading error: {e}")
    AdamW = None

# --- ERAS (10 Distinct Periods) ---
# DATASET FIX: Updated Search Hints for better boundary distinction
ERAS = [
    (500, 1200, "0_Medieval", "Medieval OR Latin manuscript OR Anglo-Saxon prose"),
    (1200, 1470, "1_Late_Medieval", "Chaucer OR Middle English OR morality play"),
    (1470, 1650, "2_Early_Modern_Renaissance", "Shakespeare OR Bacon OR Protestant theology OR Early English Bible"),
    (1650, 1800, "3_Enlightenment_Classical", "Pope couplets OR Swift satire OR Neoclassical OR reason science"),
    (1800, 1850, "4_Romantic", "Byron OR Keats OR Shelley OR nature sublime emotion"),
    (1850, 1920, "5_Industrial_Victorian", "Dickens OR Industrial Age OR Darwinism OR social novel"),
    (1920, 1945, "6_Modernist", "Modernism OR stream of consciousness OR avant-garde fiction"),
    (1945, 1960, "7_Postwar_Early_Modern", "Postwar OR Early Cold War OR existentialism"),
    (1960, 1990, "8_Late_20th_Century", "Late 20th Century OR Postmodern OR Vietnam War"),
    (1990, 2024, "9_Contemporary_Information_Age", "Contemporary OR Digital era OR internet culture")
]

ERA_LABELS = [era[2] for era in ERAS]
LABEL_TO_ID = {label: idx for idx, label in enumerate(ERA_LABELS)}
ID_TO_LABEL = {idx: label for idx, label in enumerate(ERA_LABELS)}

# --- RESCUE KEYWORDS (Unchanged) ---
RESCUE_KEYWORDS = {
    "0_Medieval": [
        "Beowulf", "Bede", "Anglo Saxon Chronicle", "Cynewulf", "Caedmon",
        "Old English Homilies", "Aelfric", "Boethius", "Alfred the Great",
        "Venerable Bede", "Old English", "Anglo-Saxon poetry"
    ],
    "1_Late_Medieval": [
        "Chaucer", "Canterbury Tales", "Piers Plowman", "Langland",
        "Gower", "Malory", "Morte d'Arthur", "Wycliffe",
        "Julian Norwich", "Margery Kempe", "Froissart", "Everyman",
        "Gawain", "Pearl Poet", "Lydgate", "Troilus Criseyde",
        "Book Duchess", "Parliament Fowls", "Legend Good Women",
        "Christine Pizan", "Romance Rose", "Confessio Amantis",
        "mystery plays", "miracle plays", "morality plays",
        "Middle English", "medieval romance", "medieval literature",
        "14th century literature", "15th century literature",
        "medieval poetry", "medieval drama", "Arthurian legend",
        "Chivalric romance", "Courtly love", "medieval manuscript",
        "Caxton", "medieval texts", "English medieval", "French medieval"
    ]
}

LATE_MEDIEVAL_COLLECTIONS = [
    "gutenberg", "opensource", "medievaltexts", "earlyenglishbooksonline",
    "englishliterature", "medievalmanuscripts", "britishlibrary"
]

# DATASET FIX: Added more contemporary-friendly topics
TOPICS = [
    "History", "Philosophy", "Science", "Mathematics", "Medicine", "Astronomy",
    "Physics", "Chemistry", "Biology", "Fiction", "Poetry", "Drama", "Mythology",
    "Folklore", "Religion", "Theology", "Biography", "Politics", "Economics", "Law",
    "Sociology", "Technology", "Engineering", "Travel", "War", "Military", "Art",
    "Psychology", "Anthropology", "Literature", "Essays", "Memoirs", "Education",
    "Computer Programming", "Digital Culture", "Current Affairs"
]

# ============================================================================
# TAB 1: DATASET GENERATION
# ============================================================================

def setup_dirs():
    if os.path.exists(DATASET_DIR):
        try: shutil.rmtree(DATASET_DIR)
        except: pass
    os.makedirs(BOOKS_DIR, exist_ok=True)

def text_quality_check(text):
    """
    A light-weight quality check to filter out poor scan or boilerplate text.
    """
    if len(text) < 3000: return False
    alpha_count = sum(c.isalpha() for c in text)
    total_count = len(text)
    if alpha_count / (total_count + 1e-6) < 0.5: return False

    start_snippet = text[:1000].lower()
    boilerplate_indicators = ["table of contents", "chapter i", "preface", "index", "list of figures"]
    if any(indicator in start_snippet for indicator in boilerplate_indicators):
        if len(text) < 10000: return False

    lines = text.split('\n')
    from collections import Counter
    line_counts = Counter(l.strip() for l in lines if l.strip())

    if len(line_counts) < 50: return False

    frequent_lines = sum(1 for count in line_counts.values() if count >= 3)
    if frequent_lines / len(line_counts) > 0.1: return False

    return True


def chunk_text_robust(text):
    MAX_TOKENS = 3500
    STRIDE = 500
    MAX_CHUNKS_PER_BOOK = 40
    chunks = []

    if TOKENIZER:
        try:
            tokens = TOKENIZER.encode(text, add_special_tokens=False)
            i = 0
            while i < len(tokens) and len(chunks) < MAX_CHUNKS_PER_BOOK:
                chunk_ids = tokens[i : i + MAX_TOKENS]
                chunk_str = TOKENIZER.decode(chunk_ids, skip_special_tokens=True)
                chunks.append(chunk_str)
                i += (MAX_TOKENS - STRIDE)
            return chunks
        except: pass

    WORDS_PER_CHUNK = 2700
    WORD_STRIDE = 400
    words = text.split()
    i = 0
    while i < len(words) and len(chunks) < MAX_CHUNKS_PER_BOOK:
        chunk_words = words[i : i + WORDS_PER_CHUNK]
        chunk_str = " ".join(chunk_words)
        if len(chunk_str) > 300:
            chunks.append(chunk_str)
        i += (WORDS_PER_CHUNK - WORD_STRIDE)
    return chunks

# ⭐️ FIX: Ensuring clean_text_content is defined before download_book
def clean_text_content(text):
    markers = [("*** START OF", "*** END OF")]
    for start_m, end_m in markers:
        s = text.find(start_m)
        e = text.find(end_m)
        if s != -1 and e != -1:
            text = text[s+len(start_m):e]
            break
    return text.strip()

# MODIFIED download_book to accept a bypass flag
def download_book(identifier, title, year, era_label, min_char_limit=5000, bypass_quality_check=False):
    urls = [
        f"https://archive.org/download/{identifier}/{identifier}_djvu.txt",
        f"https://archive.org/download/{identifier}/{identifier}.txt"
    ]
    content = ""
    for url in urls:
        try:
            r = requests.get(url, timeout=15)
            if r.status_code == 200:
                content = r.text
                break
        except: pass

    content = clean_text_content(content) # <-- The line that was failing

    if len(content) < min_char_limit:
        return None

    if not bypass_quality_check:
        if not text_quality_check(content):
            return None

    safe_title = re.sub(r'[^a-zA-Z0-9]', '_', title)[:40]
    filename = f"{year}_{era_label}_{safe_title}_{identifier}.txt"
    with open(os.path.join(BOOKS_DIR, filename), "w", encoding="utf-8") as f:
        f.write(content)

    return {
        "title": title, "year": int(year), "era_label": era_label,
        "filename": filename, "char_count": len(content), "source": "Internet Archive"
    }

def generate_dataset(total_books_needed, progress=gr.Progress()):
    setup_dirs()
    records = []

    books_per_era = max(1, int(total_books_needed / len(ERAS)))

    for start_year, end_year, era_label, search_hint in ERAS:
        collected = 0
        attempts = 0
        era_topics = TOPICS.copy()
        random.shuffle(era_topics)

        rescue_list = RESCUE_KEYWORDS.get(era_label, [])
        is_hard_era = len(rescue_list) > 0

        min_chars = 5000
        bypass_qc = False

        # FIX: Specialized rules for Contemporary Era
        if era_label == "9_Contemporary_Information_Age":
            min_chars = 2000 # Lower character requirement
            bypass_qc = True # Disable strict quality check
            max_attempts = 40 # Increase max attempts for this hard era
        elif is_hard_era:
             min_chars = 1000
             max_attempts = 50 if era_label == "1_Late_Medieval" else 20
        else:
             max_attempts = 20

        rescue_threshold = 0 if era_label == "1_Late_Medieval" else 3

        progress(0, desc=f"Scraping Era: {era_label}")
        print(f"\n{'='*60}")
        print(f"Starting Era: {era_label} (Target: {books_per_era} books | Min Chars: {min_chars})")
        print(f"{'='*60}")

        while collected < books_per_era and attempts < max_attempts:
            attempts += 1
            using_rescue = False

            if is_hard_era and attempts > rescue_threshold:
                using_rescue = True
                kw = random.choice(rescue_list)

                if era_label == "1_Late_Medieval":
                    query_type = attempts % 6
                    if query_type == 0: query = f"title:({kw}) AND mediatype:texts"
                    elif query_type == 1: query = f"({kw}) AND mediatype:texts AND language:eng"
                    elif query_type == 2: query = f"subject:({kw}) AND mediatype:texts"
                    elif query_type == 3:
                        col = random.choice(LATE_MEDIEVAL_COLLECTIONS)
                        query = f"({kw}) AND collection:({col}) AND mediatype:texts"
                    elif query_type == 4: query = f"({kw}) AND date:[1200 TO 1900] AND mediatype:texts AND language:eng"
                    else: query = f"{kw} mediatype:texts"
                else:
                    if attempts % 3 == 0: query = f"title:({kw}) AND mediatype:texts"
                    elif attempts % 3 == 1: query = f"({kw}) AND mediatype:texts AND language:eng"
                    else: query = f"subject:({kw}) AND mediatype:texts"
                print(f"  > 🛡️ Rescue Search #{attempts} ({era_label}): {kw}")
            else:
                if not era_topics:
                    era_topics = TOPICS.copy()
                    random.shuffle(era_topics)
                topic = era_topics.pop()
                query = f"(subject:{topic} OR {search_hint}) AND date:[{start_year} TO {end_year}] AND mediatype:texts AND language:eng"
                if end_year > 1928:
                    query += " AND (licenseurl:* OR rights:creative commons OR collection:opensourcemedia)"
                print(f"  > Standard Search #{attempts}: {topic} | Hint: {search_hint.split(' OR ')[0]}...")

            try:
                search_generator = internetarchive.search_items(
                    query,
                    sorts=['downloads desc'],
                    fields=['identifier', 'title', 'date', 'year']
                )

                search_results_batch = []
                max_check_per_query = (50 if is_hard_era or era_label == "9_Contemporary_Information_Age" else 15)
                for i, item in enumerate(search_generator):
                    search_results_batch.append(item)
                    if i >= max_check_per_query: break

                results_found = len(search_results_batch)

                for res in search_results_batch:
                    if collected >= books_per_era: break

                    id_ = res.get('identifier')
                    raw_date = res.get('date') or res.get('year')
                    year = str(raw_date)[:4] if raw_date else "0000"

                    if not year.isdigit(): year = "0000"

                    if not using_rescue:
                        if not (start_year <= int(year) <= end_year):
                            continue

                    if any(r['filename'].endswith(f"{id_}.txt") for r in records):
                        continue

                    rec = download_book(
                        id_, res.get('title', 'Unknown'), year, era_label,
                        min_char_limit=min_chars,
                        bypass_quality_check=bypass_qc
                    )

                    if rec:
                        rec['topic'] = "Classic" if using_rescue else topic
                        records.append(rec)
                        collected += 1
                        print(f"  ✅ Saved ({collected}/{books_per_era}): {rec['title']} ({year}) | Chars: {rec['char_count']}")

                if results_found == 0:
                    print(f"  ⚠️ No results found for this query")

            except Exception as e:
                print(f"  ❌ Search error: {e}")
                time.sleep(1)

        print(f"Completed {era_label}: {collected}/{books_per_era} books collected")

        # ... (Fallback logic for Late Medieval remains) ...
        if era_label == "1_Late_Medieval" and collected < books_per_era * 0.3:
            print(f"\n⚠️ EMERGENCY FALLBACK MODE for {era_label}")
            fallback_attempts = 0
            fallback_terms = [
                "medieval english", "middle english", "chaucer OR malory OR gower",
                "14th century OR 15th century", "medieval literature english",
                "arthurian romance", "medieval poetry english"
            ]

            while collected < books_per_era and fallback_attempts < len(fallback_terms):
                term = fallback_terms[fallback_attempts]
                fallback_attempts += 1
                query = f"({term}) AND mediatype:texts"
                print(f"  > 🚨 Fallback #{fallback_attempts}: {term}")

                try:
                    search_generator = internetarchive.search_items(query, sorts=['downloads desc'], fields=['identifier', 'title', 'date', 'year'])

                    fallback_batch = []
                    for i, item in enumerate(search_generator):
                        fallback_batch.append(item)
                        if i >= 100: break

                    checked = 0
                    for res in fallback_batch:
                        if collected >= books_per_era:
                            break
                        checked += 1

                        id_ = res.get('identifier')
                        if any(r['filename'].endswith(f"{id_}.txt") for r in records):
                            continue

                        raw_date = res.get('date') or res.get('year')
                        year = str(raw_date)[:4] if raw_date else "0000"
                        if not year.isdigit(): year = "0000"

                        rec = download_book(
                            id_, res.get('title', 'Unknown'), year, era_label,
                            min_char_limit=min_chars,
                            bypass_quality_check=bypass_qc
                        )
                        if rec:
                            rec['topic'] = "Medieval"
                            records.append(rec)
                            collected += 1
                            print(f"  ✅ FALLBACK Success ({collected}/{books_per_era}): {rec['title']} | Chars: {rec['char_count']}")
                except Exception as e:
                    print(f"  ❌ Fallback error: {e}")
                    time.sleep(1)

    if not records: return None, pd.DataFrame(), pd.DataFrame()

    print("\n" + "="*60)
    print("Starting Robust Chunking...")
    print("="*60)
    progress(0.9, desc="Chunking Text...")
    longformer_rows = []

    for r in records:
        file_path = os.path.join(BOOKS_DIR, r["filename"])
        try:
            with open(file_path, "r", encoding="utf-8") as f:
                raw_text = f.read()
            chunks = chunk_text_robust(raw_text)
            for idx, chunk in enumerate(chunks):
                longformer_rows.append({
                    "text": chunk,
                    "era_label": r["era_label"],
                    "year": r["year"],
                    "chunk_id": idx
                })
            print(f"  ✅ Chunked {r['title']}: {len(chunks)} chunks")
        except Exception as e:
            print(f"  ❌ Error processing {r['filename']}: {e}")

    df_rows = pd.DataFrame(longformer_rows)
    if not df_rows.empty:
        split_stats = df_rows['era_label'].value_counts().reset_index()
        split_stats.columns = ['Era Label', 'Total Chunks']
        split_stats['Est. Train (80%)'] = (split_stats['Total Chunks'] * 0.8).astype(int)
        split_stats['Est. Val (10%)'] = (split_stats['Total Chunks'] * 0.1).astype(int)
        split_stats['Est. Test (10%)'] = (split_stats['Total Chunks'] * 0.1).astype(int)
        split_stats['Status'] = split_stats['Est. Val (10%)'].apply(lambda x: "⚠️ LOW DATA" if x < 5 else "✅ OK")
    else:
        split_stats = pd.DataFrame()

    total_chunks = len(longformer_rows)
    avg_chunks = total_chunks / len(records) if records else 0
    general_stats_df = pd.DataFrame({
        "Metric": ["Total Books", "Total Training Examples", "Avg Examples/Book"],
        "Value": [len(records), total_chunks, f"{avg_chunks:.1f}"]
    })

    pd.DataFrame(records).to_csv(os.path.join(DATASET_DIR, "metadata.csv"), index=False)
    jsonl_path = os.path.join(DATASET_DIR, "longformer_dataset.jsonl")
    with open(jsonl_path, "w", encoding="utf-8") as f:
        for row in longformer_rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")

    timestamp = int(datetime.now().timestamp())
    zip_filename = f"Analyzed_ML_Dataset_{timestamp}"
    shutil.make_archive(zip_filename, 'zip', DATASET_DIR)

    print("\n" + "="*60)
    print("Dataset Generation Complete! READY FOR RETRAINING.")
    print("="*60)

    return f"{zip_filename}.zip", general_stats_df, split_stats

# ============================================================================
# TAB 2: TRAINING (No changes needed, already optimized for 4080 Super)
# ============================================================================

class LongformerDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length=4096):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = str(self.texts[idx])
        label = self.labels[idx]

        encoding = self.tokenizer(
            text,
            add_special_tokens=True,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )

        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'labels': torch.tensor(label, dtype=torch.long)
        }

def train_model(dataset_path, epochs, batch_size, learning_rate, gradient_accumulation_steps, progress=gr.Progress()):
    global MODEL, TOKENIZER

    if not TOKENIZER:
        return "❌ Tokenizer not loaded. Please install transformers library.", None, None

    if not os.path.exists(dataset_path):
        return "❌ Dataset file not found. Please generate a dataset first.", None, None

    if batch_size < 1:
        return "❌ Error: Batch Size must be at least 1.", None, None
    if gradient_accumulation_steps < 1:
        return "❌ Error: Gradient Accumulation Steps must be at least 1.", None, None

    scaler = torch.cuda.amp.GradScaler() if DEVICE == "cuda" else None

    try:
        progress(0.1, desc="Loading dataset...")
        data = []
        with open(dataset_path, 'r', encoding='utf-8') as f:
            for line in f:
                data.append(json.loads(line))

        df = pd.DataFrame(data)
        texts = df['text'].tolist()
        labels = [LABEL_TO_ID[label] for label in df['era_label'].tolist()]

        progress(0.2, desc="Splitting data...")
        X_train, X_temp, y_train, y_temp = train_test_split(texts, labels, test_size=0.2, random_state=42, stratify=labels)
        X_val, X_test, y_val, y_test = train_test_split(X_temp, y_temp, test_size=0.5, random_state=42, stratify=y_temp)

        train_dataset = LongformerDataset(X_train, y_train, TOKENIZER)
        val_dataset = LongformerDataset(X_val, y_val, TOKENIZER)

        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
        val_loader = DataLoader(val_dataset, batch_size=batch_size)

        progress(0.3, desc="Initializing model...")
        MODEL = AutoModelForSequenceClassification.from_pretrained(
            "allenai/longformer-base-4096",
            num_labels=len(LABEL_TO_ID)
        )
        MODEL.to(DEVICE)

        optimizer = AdamW(MODEL.parameters(), lr=learning_rate)

        total_batches = len(train_loader)
        total_training_steps = (total_batches // gradient_accumulation_steps) * epochs
        scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=total_training_steps)

        train_losses = []
        val_accuracies = []
        step_count = 0

        for epoch in range(epochs):
            MODEL.train()
            total_loss = 0

            for batch_idx, batch in enumerate(train_loader):
                progress_val = (0.3 + (epoch / epochs) * 0.6) + ((batch_idx / total_batches) / epochs * 0.6)
                progress(progress_val, desc=f"Training Epoch {epoch+1}/{epochs} (Batch {batch_idx+1}/{total_batches})")

                input_ids = batch['input_ids'].to(DEVICE)
                attention_mask = batch['attention_mask'].to(DEVICE)
                labels = batch['labels'].to(DEVICE)

                with torch.cuda.amp.autocast(enabled=(DEVICE == "cuda")):
                    outputs = MODEL(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
                    loss = outputs.loss
                    loss = loss / gradient_accumulation_steps

                if scaler:
                    scaler.scale(loss).backward()
                else:
                    loss.backward()

                total_loss += loss.item() * gradient_accumulation_steps
                step_count += 1

                if step_count % gradient_accumulation_steps == 0 or batch_idx == total_batches - 1:
                    if scaler:
                        scaler.unscale_(optimizer)
                        torch.nn.utils.clip_grad_norm_(MODEL.parameters(), 1.0)
                        scaler.step(optimizer)
                        scaler.update()
                    else:
                        torch.nn.utils.clip_grad_norm_(MODEL.parameters(), 1.0)
                        optimizer.step()

                    scheduler.step()
                    optimizer.zero_grad()

            MODEL.eval()
            correct = 0
            total = 0

            with torch.no_grad():
                with torch.cuda.amp.autocast(enabled=(DEVICE == "cuda")):
                    for batch in val_loader:
                        input_ids = batch['input_ids'].to(DEVICE)
                        attention_mask = batch['attention_mask'].to(DEVICE)
                        labels = batch['labels'].to(DEVICE)

                        outputs = MODEL(input_ids=input_ids, attention_mask=attention_mask)
                        predictions = torch.argmax(outputs.logits, dim=1)

                        correct += (predictions == labels).sum().item()
                        total += labels.size(0)

            avg_loss = total_loss / total_batches
            val_acc = correct / total

            train_losses.append(avg_loss)
            val_accuracies.append(val_acc)

            print(f"Epoch {epoch+1}/{epochs} - Loss: {avg_loss:.4f}, Val Acc: {val_acc:.4f}")

        progress(0.95, desc="Saving model...")
        timestamp = int(datetime.now().timestamp())
        model_path = os.path.join(MODEL_DIR, f"longformer_era_classifier_{timestamp}")
        MODEL.save_pretrained(model_path)
        TOKENIZER.save_pretrained(model_path)

        metrics_df = pd.DataFrame({
            "Epoch": list(range(1, epochs + 1)),
            "Training Loss": train_losses,
            "Validation Accuracy": [f"{acc:.4f}" for acc in val_accuracies]
        })

        summary = f"✅ Training Complete!\nFinal Val Acc: {val_accuracies[-1]:.4f}\nModel saved to: {model_path}"

        return summary, metrics_df, model_path

    except RuntimeError as e:
        if 'out of memory' in str(e):
            if DEVICE == "cuda": torch.cuda.empty_cache()
            return f"❌ Training error: CUDA Out Of Memory. Try reducing the 'Batch Size' slider to 1, or increase 'Gradient Accumulation Steps'. Error: {str(e)}", None, None
        return f"❌ Training error: {str(e)}", None, None
    except Exception as e:
        return f"❌ Training error: {str(e)}", None, None

# ============================================================================
# TAB 3: TESTING (No changes needed)
# ============================================================================

def load_trained_model(model_path):
    global MODEL, TOKENIZER

    try:
        TOKENIZER = AutoTokenizer.from_pretrained(model_path)
        MODEL = AutoModelForSequenceClassification.from_pretrained(model_path)
        MODEL.to(DEVICE)
        MODEL.eval()
        return f"✅ Model loaded successfully from {model_path}"
    except Exception as e:
        return f"❌ Error loading model: {str(e)}"

def predict_era(text, model_path):
    global MODEL, TOKENIZER

    if not MODEL or not TOKENIZER:
        if model_path and os.path.exists(model_path):
            load_result = load_trained_model(model_path)
            if "Error" in load_result:
                return load_result, None
        else:
            return "❌ No model loaded. Please train a model first or provide a valid model path.", None

    try:
        encoding = TOKENIZER(
            text,
            add_special_tokens=True,
            max_length=4096,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )

        input_ids = encoding['input_ids'].to(DEVICE)
        attention_mask = encoding['attention_mask'].to(DEVICE)

        with torch.no_grad():
            with torch.cuda.amp.autocast(enabled=(DEVICE == "cuda")):
                outputs = MODEL(input_ids=input_ids, attention_mask=attention_mask)
                logits = outputs.logits
                probabilities = torch.softmax(logits, dim=1)[0]
                predicted_class = torch.argmax(probabilities).item()

        top_3_probs, top_3_indices = torch.topk(probabilities, 3)

        results = []
        for idx, prob in zip(top_3_indices, top_3_probs):
            era_label = ID_TO_LABEL[idx.item()]
            confidence = prob.item() * 100
            results.append({
                "Era": era_label,
                "Confidence": f"{confidence:.2f}%"
            })

        predicted_era = ID_TO_LABEL[predicted_class]
        result_text = f"🎯 **Predicted Era:** {predicted_era}\n\n**Confidence:** {probabilities[predicted_class].item()*100:.2f}%"

        return result_text, pd.DataFrame(results)

    except Exception as e:
        return f"❌ Prediction error: {str(e)}", None

# ============================================================================
# GRADIO UI
# ============================================================================

with gr.Blocks(title="Complete ML Pipeline") as demo:
    gr.Markdown("# 📚 Complete ML Pipeline: Dataset Generation, Training & Testing (RTX 4080 Super Optimized)")

    with gr.Tabs():
        # TAB 1: Dataset Generation
        with gr.Tab("📊 Dataset Generation"):
            gr.Markdown("## Generate Historical Text Dataset")
            gr.Markdown("""
            **DATA QUALITY FIX:** Contemporary Era (`9_...`) now has lower length requirements and a less strict quality check to compensate for scarce open-source post-1990 data.
            """)

            with gr.Row():
                dataset_slider = gr.Slider(10, 500, step=10, value=100, label="Total Books to Collect (Max 500)")
                generate_btn = gr.Button("🚀 Generate Dataset (New Data Quality)", variant="primary", size="lg")

            dataset_download = gr.File(label="📥 Download Dataset ZIP")

            with gr.Row():
                with gr.Column():
                    gr.Markdown("### General Summary")
                    gen_stats = gr.Dataframe()
                with gr.Column():
                    gr.Markdown("### Class Balance Check")
                    split_stats = gr.Dataframe()

            generate_btn.click(
                generate_dataset,
                inputs=[dataset_slider],
                outputs=[dataset_download, gen_stats, split_stats]
            )

        # TAB 2: Training
        with gr.Tab("🎓 Model Training"):
            gr.Markdown("## Train Longformer Era Classifier")
            gr.Markdown(f"""
            **GPU OPTIMIZED:** Training now uses **Automatic Mixed Precision (FP16/AMP)** for the RTX 4080 Super.
            With 16GB VRAM, you can use a higher **Batch Size** (e.g., 4 or 8) and often set **Gradient Accumulation Steps** to 1.
            """)

            with gr.Row():
                with gr.Column():
                    train_dataset_path = gr.Textbox(
                        label="Dataset Path",
                        value=os.path.join(DATASET_DIR, "longformer_dataset.jsonl"),
                        placeholder="Path to dataset JSONL file"
                    )
                    train_epochs = gr.Slider(1, 10, step=1, value=3, label="Epochs")
                    train_batch = gr.Slider(1, 16, step=1, value=4, label="Batch Size (Memory Control)")
                    train_accum = gr.Slider(1, 16, step=1, value=1, label="Gradient Accumulation Steps (Effective Batch Size)")
                    train_lr = gr.Number(value=2e-5, label="Learning Rate")
                    train_btn = gr.Button("🏋️ Start Training", variant="primary", size="lg")

                with gr.Column():
                    train_output = gr.Textbox(label="Training Status", lines=8)
                    train_metrics = gr.Dataframe(label="Training Metrics")
                    model_path_output = gr.Textbox(label="Saved Model Path")

            train_btn.click(
                train_model,
                inputs=[train_dataset_path, train_epochs, train_batch, train_lr, train_accum],
                outputs=[train_output, train_metrics, model_path_output]
            )

        # TAB 3: Testing
        with gr.Tab("🧪 Model Testing"):
            gr.Markdown("## Test Era Classification (FP16/AMP Inference)")

            with gr.Row():
                with gr.Column():
                    test_model_path = gr.Textbox(
                        label="Model Path (optional - uses last trained model if empty)",
                        placeholder="trained_models/longformer_era_classifier_..."
                    )
                    test_input = gr.Textbox(
                        label="Input Text",
                        lines=10,
                        placeholder="Paste historical text here...\n\nExample: 'When that Aprille with his shoures soote, The droghte of Marche hath perced to the roote...'"
                    )
                    test_btn = gr.Button("🔍 Predict Era", variant="primary", size="lg")

                with gr.Column():
                    test_result = gr.Markdown(label="Prediction Result")
                    test_probabilities = gr.Dataframe(label="Top 3 Predictions")

            # Sample texts
            gr.Markdown("### 📝 Try Sample Texts")
            with gr.Row():
                sample1 = gr.Button("Medieval Sample")
                sample2 = gr.Button("Victorian Sample")
                sample3 = gr.Button("Contemporary Sample")

            def load_medieval():
                return "Hwæt! We Gardena in geardagum, þeodcyninga, þrym gefrunon, hu ða æþelingas ellen fremedon."

            def load_victorian():
                return "It is a truth universally acknowledged, that a single man in possession of a good fortune, must be in want of a wife."

            def load_contemporary():
                return "The internet has fundamentally transformed how we communicate, work, and access information in the digital age."

            sample1.click(load_medieval, outputs=[test_input])
            sample2.click(load_victorian, outputs=[test_input])
            sample3.click(load_contemporary, outputs=[test_input])

            test_btn.click(
                predict_era,
                inputs=[test_input, test_model_path],
                outputs=[test_result, test_probabilities]
            )

    gr.Markdown("---")
    gr.Markdown(f"**Device:** {DEVICE} | **Status:** {'✅ CUDA/FP16 Ready' if DEVICE == 'cuda' else '⚠️ CPU Mode'} | **Model:** Longformer-base-4096")

if __name__ == "__main__":
    demo.launch(ssr_mode=False)