#!/usr/bin/env python3
import os
import glob
import base64
import time
import shutil
import streamlit as st
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
from diffusers import StableDiffusionPipeline
from torch.utils.data import Dataset, DataLoader
import csv
import fitz  # PyMuPDF
import requests
from PIL import Image
import cv2
import numpy as np
import logging
import asyncio
import aiofiles
from io import BytesIO
from dataclasses import dataclass
from typing import Optional, Tuple
import zipfile
import math
import random
import re

# Logging setup with custom buffer
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
log_records = []

class LogCaptureHandler(logging.Handler):
    def emit(self, record):
        log_records.append(record)

logger.addHandler(LogCaptureHandler())

# Page Configuration
st.set_page_config(
    page_title="AI Vision & SFT Titans 🚀",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded",
    menu_items={
        'Get Help': 'https://huggingface.co/awacke1',
        'Report a Bug': 'https://huggingface.co/spaces/awacke1',
        'About': "AI Vision & SFT Titans: PDFs, OCR, Image Gen, Line Drawings, Custom Diffusion, and SFT on CPU! 🌌"
    }
)

# Initialize st.session_state
if 'history' not in st.session_state:
    st.session_state['history'] = []  # Flat list for history
if 'builder' not in st.session_state:
    st.session_state['builder'] = None
if 'model_loaded' not in st.session_state:
    st.session_state['model_loaded'] = False
if 'processing' not in st.session_state:
    st.session_state['processing'] = {}
if 'pdf_checkboxes' not in st.session_state:
    st.session_state['pdf_checkboxes'] = {}  # Shared cache for PDF checkboxes
if 'downloaded_pdfs' not in st.session_state:
    st.session_state['downloaded_pdfs'] = {}  # Cache for downloaded PDF paths

# Model Configuration Classes
@dataclass
class ModelConfig:
    name: str
    base_model: str
    size: str
    domain: Optional[str] = None
    model_type: str = "causal_lm"
    @property
    def model_path(self):
        return f"models/{self.name}"

@dataclass
class DiffusionConfig:
    name: str
    base_model: str
    size: str
    @property
    def model_path(self):
        return f"diffusion_models/{self.name}"

# Datasets
class SFTDataset(Dataset):
    def __init__(self, data, tokenizer, max_length=128):
        self.data = data
        self.tokenizer = tokenizer
        self.max_length = max_length
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        prompt = self.data[idx]["prompt"]
        response = self.data[idx]["response"]
        full_text = f"{prompt} {response}"
        full_encoding = self.tokenizer(full_text, max_length=self.max_length, padding="max_length", truncation=True, return_tensors="pt")
        prompt_encoding = self.tokenizer(prompt, max_length=self.max_length, padding=False, truncation=True, return_tensors="pt")
        input_ids = full_encoding["input_ids"].squeeze()
        attention_mask = full_encoding["attention_mask"].squeeze()
        labels = input_ids.clone()
        prompt_len = prompt_encoding["input_ids"].shape[1]
        if prompt_len < self.max_length:
            labels[:prompt_len] = -100
        return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}

class DiffusionDataset(Dataset):
    def __init__(self, images, texts):
        self.images = images
        self.texts = texts
    def __len__(self):
        return len(self.images)
    def __getitem__(self, idx):
        return {"image": self.images[idx], "text": self.texts[idx]}

class TinyDiffusionDataset(Dataset):
    def __init__(self, images):
        self.images = [torch.tensor(np.array(img.convert("RGB")).transpose(2, 0, 1), dtype=torch.float32) / 255.0 for img in images]
    def __len__(self):
        return len(self.images)
    def __getitem__(self, idx):
        return self.images[idx]

# Custom Tiny Diffusion Model
class TinyUNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=3):
        super(TinyUNet, self).__init__()
        self.down1 = nn.Conv2d(in_channels, 32, 3, padding=1)
        self.down2 = nn.Conv2d(32, 64, 3, padding=1, stride=2)
        self.mid = nn.Conv2d(64, 128, 3, padding=1)
        self.up1 = nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1)
        self.up2 = nn.Conv2d(64 + 32, 32, 3, padding=1)
        self.out = nn.Conv2d(32, out_channels, 3, padding=1)
        self.time_embed = nn.Linear(1, 64)

    def forward(self, x, t):
        t_embed = F.relu(self.time_embed(t.unsqueeze(-1)))
        t_embed = t_embed.view(t_embed.size(0), t_embed.size(1), 1, 1)
        
        x1 = F.relu(self.down1(x))
        x2 = F.relu(self.down2(x1))
        x_mid = F.relu(self.mid(x2)) + t_embed
        x_up1 = F.relu(self.up1(x_mid))
        x_up2 = F.relu(self.up2(torch.cat([x_up1, x1], dim=1)))
        return self.out(x_up2)

class TinyDiffusion:
    def __init__(self, model, timesteps=100):
        self.model = model
        self.timesteps = timesteps
        self.beta = torch.linspace(0.0001, 0.02, timesteps)
        self.alpha = 1 - self.beta
        self.alpha_cumprod = torch.cumprod(self.alpha, dim=0)

    def train(self, images, epochs=50):
        dataset = TinyDiffusionDataset(images)
        dataloader = DataLoader(dataset, batch_size=1, shuffle=True)
        optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-4)
        device = torch.device("cpu")
        self.model.to(device)
        for epoch in range(epochs):
            total_loss = 0
            for x in dataloader:
                x = x.to(device)
                t = torch.randint(0, self.timesteps, (x.size(0),), device=device).float()
                noise = torch.randn_like(x)
                alpha_t = self.alpha_cumprod[t.long()].view(-1, 1, 1, 1)
                x_noisy = torch.sqrt(alpha_t) * x + torch.sqrt(1 - alpha_t) * noise
                pred_noise = self.model(x_noisy, t)
                loss = F.mse_loss(pred_noise, noise)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                total_loss += loss.item()
            logger.info(f"Epoch {epoch + 1}/{epochs}, Loss: {total_loss / len(dataloader):.4f}")
        return self

    def generate(self, size=(64, 64), steps=100):
        device = torch.device("cpu")
        x = torch.randn(1, 3, size[0], size[1], device=device)
        for t in reversed(range(steps)):
            t_tensor = torch.full((1,), t, device=device, dtype=torch.float32)
            alpha_t = self.alpha_cumprod[t].view(-1, 1, 1, 1)
            pred_noise = self.model(x, t_tensor)
            x = (x - (1 - self.alpha[t]) / torch.sqrt(1 - alpha_t) * pred_noise) / torch.sqrt(self.alpha[t])
            if t > 0:
                x += torch.sqrt(self.beta[t]) * torch.randn_like(x)
        x = torch.clamp(x * 255, 0, 255).byte()
        return Image.fromarray(x.squeeze(0).permute(1, 2, 0).cpu().numpy())

    def upscale(self, image, scale_factor=2):
        img_tensor = torch.tensor(np.array(image.convert("RGB")).transpose(2, 0, 1), dtype=torch.float32).unsqueeze(0) / 255.0
        upscaled = F.interpolate(img_tensor, scale_factor=scale_factor, mode='bilinear', align_corners=False)
        upscaled = torch.clamp(upscaled * 255, 0, 255).byte()
        return Image.fromarray(upscaled.squeeze(0).permute(1, 2, 0).cpu().numpy())

# Model Builders
class ModelBuilder:
    def __init__(self):
        self.config = None
        self.model = None
        self.tokenizer = None
        self.sft_data = None
        self.jokes = ["Why did the AI go to therapy? Too many layers to unpack! 😂", "Training complete! Time for a binary coffee break. ☕"]
    def load_model(self, model_path: str, config: Optional[ModelConfig] = None):
        with st.spinner(f"Loading {model_path}... ⏳"):
            self.model = AutoModelForCausalLM.from_pretrained(model_path)
            self.tokenizer = AutoTokenizer.from_pretrained(model_path)
            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token
            if config:
                self.config = config
            self.model.to("cuda" if torch.cuda.is_available() else "cpu")
        st.success(f"Model loaded! 🎉 {random.choice(self.jokes)}")
        return self
    def fine_tune_sft(self, csv_path: str, epochs: int = 3, batch_size: int = 4):
        self.sft_data = []
        with open(csv_path, "r") as f:
            reader = csv.DictReader(f)
            for row in reader:
                self.sft_data.append({"prompt": row["prompt"], "response": row["response"]})
        dataset = SFTDataset(self.sft_data, self.tokenizer)
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
        optimizer = torch.optim.AdamW(self.model.parameters(), lr=2e-5)
        self.model.train()
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(device)
        for epoch in range(epochs):
            with st.spinner(f"Training epoch {epoch + 1}/{epochs}... ⚙️"):
                total_loss = 0
                for batch in dataloader:
                    optimizer.zero_grad()
                    input_ids = batch["input_ids"].to(device)
                    attention_mask = batch["attention_mask"].to(device)
                    labels = batch["labels"].to(device)
                    outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
                    loss = outputs.loss
                    loss.backward()
                    optimizer.step()
                    total_loss += loss.item()
                st.write(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}")
        st.success(f"SFT Fine-tuning completed! 🎉 {random.choice(self.jokes)}")
        return self
    def save_model(self, path: str):
        with st.spinner("Saving model... 💾"):
            os.makedirs(os.path.dirname(path), exist_ok=True)
            self.model.save_pretrained(path)
            self.tokenizer.save_pretrained(path)
        st.success(f"Model saved at {path}! ✅")
    def evaluate(self, prompt: str, status_container=None):
        self.model.eval()
        if status_container:
            status_container.write("Preparing to evaluate... 🧠")
        try:
            with torch.no_grad():
                inputs = self.tokenizer(prompt, return_tensors="pt", max_length=128, truncation=True).to(self.model.device)
                outputs = self.model.generate(**inputs, max_new_tokens=50, do_sample=True, top_p=0.95, temperature=0.7)
                return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        except Exception as e:
            if status_container:
                status_container.error(f"Oops! Something broke: {str(e)} 💥")
            return f"Error: {str(e)}"

class DiffusionBuilder:
    def __init__(self):
        self.config = None
        self.pipeline = None
    def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None):
        with st.spinner(f"Loading diffusion model {model_path}... ⏳"):
            self.pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32).to("cpu")
            if config:
                self.config = config
        st.success(f"Diffusion model loaded! 🎨")
        return self
    def fine_tune_sft(self, images, texts, epochs=3):
        dataset = DiffusionDataset(images, texts)
        dataloader = DataLoader(dataset, batch_size=1, shuffle=True)
        optimizer = torch.optim.AdamW(self.pipeline.unet.parameters(), lr=1e-5)
        self.pipeline.unet.train()
        for epoch in range(epochs):
            with st.spinner(f"Training diffusion epoch {epoch + 1}/{epochs}... ⚙️"):
                total_loss = 0
                for batch in dataloader:
                    optimizer.zero_grad()
                    image = batch["image"][0].to(self.pipeline.device)
                    text = batch["text"][0]
                    latents = self.pipeline.vae.encode(torch.tensor(np.array(image)).permute(2, 0, 1).unsqueeze(0).float().to(self.pipeline.device)).latent_dist.sample()
                    noise = torch.randn_like(latents)
                    timesteps = torch.randint(0, self.pipeline.scheduler.num_train_timesteps, (latents.shape[0],), device=latents.device)
                    noisy_latents = self.pipeline.scheduler.add_noise(latents, noise, timesteps)
                    text_embeddings = self.pipeline.text_encoder(self.pipeline.tokenizer(text, return_tensors="pt").input_ids.to(self.pipeline.device))[0]
                    pred_noise = self.pipeline.unet(noisy_latents, timesteps, encoder_hidden_states=text_embeddings).sample
                    loss = torch.nn.functional.mse_loss(pred_noise, noise)
                    loss.backward()
                    optimizer.step()
                    total_loss += loss.item()
                st.write(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}")
        st.success("Diffusion SFT Fine-tuning completed! 🎨")
        return self
    def save_model(self, path: str):
        with st.spinner("Saving diffusion model... 💾"):
            os.makedirs(os.path.dirname(path), exist_ok=True)
            self.pipeline.save_pretrained(path)
        st.success(f"Diffusion model saved at {path}! ✅")
    def generate(self, prompt: str):
        return self.pipeline(prompt, num_inference_steps=20).images[0]

# Utility Functions
def generate_filename(sequence, ext="png"):
    timestamp = time.strftime("%d%m%Y%H%M%S")
    return f"{sequence}_{timestamp}.{ext}"

def pdf_url_to_filename(url):
    # Convert full URL to filename, replacing illegal characters
    safe_name = re.sub(r'[<>:"/\\|?*]', '_', url)
    return f"{safe_name}.pdf"

def get_download_link(file_path, mime_type="application/pdf", label="Download"):
    with open(file_path, 'rb') as f:
        data = f.read()
    b64 = base64.b64encode(data).decode()
    return f'<a href="data:{mime_type};base64,{b64}" download="{os.path.basename(file_path)}">{label}</a>'

def zip_directory(directory_path, zip_path):
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, _, files in os.walk(directory_path):
            for file in files:
                zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), os.path.dirname(directory_path)))

def get_model_files(model_type="causal_lm"):
    path = "models/*" if model_type == "causal_lm" else "diffusion_models/*"
    return [d for d in glob.glob(path) if os.path.isdir(d)]

def get_gallery_files(file_types=["png"]):
    return sorted([f for ext in file_types for f in glob.glob(f"*.{ext}")])

def get_pdf_files():
    return sorted(glob.glob("*.pdf"))

def download_pdf(url, output_path):
    try:
        response = requests.get(url, stream=True, timeout=10)
        if response.status_code == 200:
            with open(output_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
            return True
    except requests.RequestException as e:
        logger.error(f"Failed to download {url}: {e}")
    return False

# Async Processing Functions
async def process_pdf_snapshot(pdf_path, mode="single"):
    start_time = time.time()
    status = st.empty()
    status.text(f"Processing PDF Snapshot ({mode})... (0s)")
    try:
        doc = fitz.open(pdf_path)
        output_files = []
        if mode == "single":
            page = doc[0]
            pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))  # High-res: 200% scale
            output_file = generate_filename("single", "png")
            pix.save(output_file)
            output_files.append(output_file)
        elif mode == "twopage":
            for i in range(min(2, len(doc))):
                page = doc[i]
                pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))  # High-res: 200% scale
                output_file = generate_filename(f"twopage_{i}", "png")
                pix.save(output_file)
                output_files.append(output_file)
        elif mode == "allthumbs":
            for i in range(len(doc)):
                page = doc[i]
                pix = page.get_pixmap(matrix=fitz.Matrix(0.5, 0.5))  # Thumbnail: 50% scale
                output_file = generate_filename(f"thumb_{i}", "png")
                pix.save(output_file)
                output_files.append(output_file)
        doc.close()
        elapsed = int(time.time() - start_time)
        status.text(f"PDF Snapshot ({mode}) completed in {elapsed}s!")
        update_gallery()
        return output_files
    except Exception as e:
        status.error(f"Failed to process PDF: {str(e)}")
        return []

async def process_ocr(image, output_file):
    start_time = time.time()
    status = st.empty()
    status.text("Processing GOT-OCR2_0... (0s)")
    tokenizer = AutoTokenizer.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True)
    model = AutoModel.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True, torch_dtype=torch.float32).to("cpu").eval()
    result = model.chat(tokenizer, image, ocr_type='ocr')
    elapsed = int(time.time() - start_time)
    status.text(f"GOT-OCR2_0 completed in {elapsed}s!")
    async with aiofiles.open(output_file, "w") as f:
        await f.write(result)
    update_gallery()
    return result

async def process_image_gen(prompt, output_file):
    start_time = time.time()
    status = st.empty()
    status.text("Processing Image Gen... (0s)")
    pipeline = StableDiffusionPipeline.from_pretrained("OFA-Sys/small-stable-diffusion-v0", torch_dtype=torch.float32).to("cpu")
    gen_image = pipeline(prompt, num_inference_steps=20).images[0]
    elapsed = int(time.time() - start_time)
    status.text(f"Image Gen completed in {elapsed}s!")
    gen_image.save(output_file)
    update_gallery()
    return gen_image

async def process_custom_diffusion(images, output_file, model_name):
    start_time = time.time()
    status = st.empty()
    status.text(f"Training {model_name}... (0s)")
    unet = TinyUNet()
    diffusion = TinyDiffusion(unet)
    diffusion.train(images)
    gen_image = diffusion.generate()
    upscaled_image = diffusion.upscale(gen_image, scale_factor=2)
    elapsed = int(time.time() - start_time)
    status.text(f"{model_name} completed in {elapsed}s!")
    upscaled_image.save(output_file)
    update_gallery()
    return upscaled_image

# Mock Search Tool for RAG
def mock_search(query: str) -> str:
    if "superhero" in query.lower():
        return "Latest trends: Gold-plated Batman statues, VR superhero battles."
    return "No relevant results found."

def mock_duckduckgo_search(query: str) -> str:
    if "superhero party trends" in query.lower():
        return """
        Latest trends for 2025: 
        - Luxury decorations: Gold-plated Batman statues, holographic Avengers displays.
        - Entertainment: Live stunt shows with Iron Man suits, VR superhero battles.
        - Catering: Gourmet kryptonite-green cocktails, Thor’s hammer-shaped appetizers.
        """
    return "No relevant results found."

# Agent Classes
class PartyPlannerAgent:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
    def generate(self, prompt: str) -> str:
        self.model.eval()
        with torch.no_grad():
            inputs = self.tokenizer(prompt, return_tensors="pt", max_length=128, truncation=True).to(self.device)
            outputs = self.model.generate(**inputs, max_new_tokens=100, do_sample=True, top_p=0.95, temperature=0.7)
            return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
    def plan_party(self, task: str) -> pd.DataFrame:
        search_result = mock_duckduckgo_search("latest superhero party trends")
        prompt = f"Given this context: '{search_result}'\n{task}"
        plan_text = self.generate(prompt)
        locations = {
            "Wayne Manor": (42.3601, -71.0589),
            "New York": (40.7128, -74.0060),
            "Los Angeles": (34.0522, -118.2437),
            "London": (51.5074, -0.1278)
        }
        wayne_coords = locations["Wayne Manor"]
        travel_times = {loc: calculate_cargo_travel_time(coords, wayne_coords) for loc, coords in locations.items() if loc != "Wayne Manor"}
        catchphrases = ["To the Batmobile!", "Avengers, assemble!", "I am Iron Man!", "By the power of Grayskull!"]
        data = [
            {"Location": "New York", "Travel Time (hrs)": travel_times["New York"], "Luxury Idea": "Gold-plated Batman statues", "Catchphrase": random.choice(catchphrases)},
            {"Location": "Los Angeles", "Travel Time (hrs)": travel_times["Los Angeles"], "Luxury Idea": "Holographic Avengers displays", "Catchphrase": random.choice(catchphrases)},
            {"Location": "London", "Travel Time (hrs)": travel_times["London"], "Luxury Idea": "Live stunt shows with Iron Man suits", "Catchphrase": random.choice(catchphrases)},
            {"Location": "Wayne Manor", "Travel Time (hrs)": 0.0, "Luxury Idea": "VR superhero battles", "Catchphrase": random.choice(catchphrases)},
            {"Location": "New York", "Travel Time (hrs)": travel_times["New York"], "Luxury Idea": "Gourmet kryptonite-green cocktails", "Catchphrase": random.choice(catchphrases)},
            {"Location": "Los Angeles", "Travel Time (hrs)": travel_times["Los Angeles"], "Luxury Idea": "Thor’s hammer-shaped appetizers", "Catchphrase": random.choice(catchphrases)},
        ]
        return pd.DataFrame(data)

class CVPartyPlannerAgent:
    def __init__(self, pipeline):
        self.pipeline = pipeline
    def generate(self, prompt: str) -> Image.Image:
        return self.pipeline(prompt, num_inference_steps=20).images[0]
    def plan_party(self, task: str) -> pd.DataFrame:
        search_result = mock_search("superhero party trends")
        prompt = f"Given this context: '{search_result}'\n{task}"
        data = [
            {"Theme": "Batman", "Image Idea": "Gold-plated Batman statue"},
            {"Theme": "Avengers", "Image Idea": "VR superhero battle scene"}
        ]
        return pd.DataFrame(data)

def calculate_cargo_travel_time(origin_coords: Tuple[float, float], destination_coords: Tuple[float, float], cruising_speed_kmh: float = 750.0) -> float:
    def to_radians(degrees: float) -> float:
        return degrees * (math.pi / 180)
    lat1, lon1 = map(to_radians, origin_coords)
    lat2, lon2 = map(to_radians, destination_coords)
    EARTH_RADIUS_KM = 6371.0
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    a = (math.sin(dlat / 2) ** 2 + math.cos(lat1) * math.cos(lat2) * math.sin(dlon / 2) ** 2)
    c = 2 * math.asin(math.sqrt(a))
    distance = EARTH_RADIUS_KM * c
    actual_distance = distance * 1.1
    flight_time = (actual_distance / cruising_speed_kmh) + 1.0
    return round(flight_time, 2)

# Main App
st.title("AI Vision & SFT Titans 🚀")

# Sidebar
st.sidebar.header("Captured Files 📜")
gallery_size = st.sidebar.slider("Gallery Size", 1, 10, 2)  # Default to 2
def update_gallery():
    media_files = get_gallery_files(["png"])
    pdf_files = get_pdf_files()
    if media_files or pdf_files:
        st.sidebar.subheader("Images 📸")
        cols = st.sidebar.columns(2)
        for idx, file in enumerate(media_files[:gallery_size * 2]):  # Limit by gallery size
            with cols[idx % 2]:
                st.image(Image.open(file), caption=os.path.basename(file), use_container_width=True)
        st.sidebar.subheader("PDF Downloads 📖")
        for pdf_file in pdf_files[:gallery_size * 2]:  # Limit by gallery size
            st.markdown(get_download_link(pdf_file, "application/pdf", f"📥 Grab {os.path.basename(pdf_file)}"), unsafe_allow_html=True)
update_gallery()

st.sidebar.subheader("Model Management 🗂️")
model_type = st.sidebar.selectbox("Model Type", ["Causal LM", "Diffusion"], key="sidebar_model_type")
model_dirs = get_model_files(model_type)
selected_model = st.sidebar.selectbox("Select Saved Model", ["None"] + model_dirs, key="sidebar_model_select")
if selected_model != "None" and st.sidebar.button("Load Model 📂"):
    builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder()
    config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(name=os.path.basename(selected_model), base_model="unknown", size="small")
    builder.load_model(selected_model, config)
    st.session_state['builder'] = builder
    st.session_state['model_loaded'] = True
    st.rerun()

st.sidebar.subheader("Action Logs 📜")
log_container = st.sidebar.empty()
with log_container:
    for record in log_records:
        st.write(f"{record.asctime} - {record.levelname} - {record.message}")

st.sidebar.subheader("History 📜")
history_container = st.sidebar.empty()
with history_container:
    for entry in st.session_state['history'][-gallery_size * 2:]:  # Limit by gallery size
        st.write(entry)

# Tabs
tab1, tab2, tab3, tab4, tab5, tab6, tab7, tab8, tab9 = st.tabs([
    "Camera Snap 📷", "Download PDFs 📥", "Build Titan 🌱", "Fine-Tune Titan 🔧", 
    "Test Titan 🧪", "Agentic RAG Party 🌐", "Test OCR 🔍", "Test Image Gen 🎨", "Custom Diffusion 🎨🤓"
])

with tab1:
    st.header("Camera Snap 📷")
    st.subheader("Single Capture")
    cols = st.columns(2)
    with cols[0]:
        cam0_img = st.camera_input("Take a picture - Cam 0", key="cam0")
        if cam0_img:
            filename = generate_filename("cam0")
            with open(filename, "wb") as f:
                f.write(cam0_img.getvalue())
            entry = f"Snapshot from Cam 0: {filename}"
            if entry not in st.session_state['history']:
                st.session_state['history'] = [e for e in st.session_state['history'] if not e.startswith("Snapshot from Cam 0:")] + [entry]
            st.image(Image.open(filename), caption="Camera 0", use_container_width=True)
            logger.info(f"Saved snapshot from Camera 0: {filename}")
            update_gallery()
    with cols[1]:
        cam1_img = st.camera_input("Take a picture - Cam 1", key="cam1")
        if cam1_img:
            filename = generate_filename("cam1")
            with open(filename, "wb") as f:
                f.write(cam1_img.getvalue())
            entry = f"Snapshot from Cam 1: {filename}"
            if entry not in st.session_state['history']:
                st.session_state['history'] = [e for e in st.session_state['history'] if not e.startswith("Snapshot from Cam 1:")] + [entry]
            st.image(Image.open(filename), caption="Camera 1", use_container_width=True)
            logger.info(f"Saved snapshot from Camera 1: {filename}")
            update_gallery()

with tab2:
    st.header("Download PDFs 📥")
    # Examples button with arXiv PDF links from README.md
    if st.button("Examples 📚"):
        example_urls = [
            "https://arxiv.org/pdf/2308.03892",  # Streamlit
            "https://arxiv.org/pdf/1912.01703",  # PyTorch
            "https://arxiv.org/pdf/2408.11039",  # Qwen2-VL
            "https://arxiv.org/pdf/2109.10282",  # TrOCR
            "https://arxiv.org/pdf/2112.10752",  # LDM
            "https://arxiv.org/pdf/2308.11236",  # OpenCV
            "https://arxiv.org/pdf/1706.03762",  # Attention is All You Need
            "https://arxiv.org/pdf/2006.11239",  # DDPM
            "https://arxiv.org/pdf/2305.11207",  # Pandas
            "https://arxiv.org/pdf/2106.09685",  # LoRA
            "https://arxiv.org/pdf/2005.11401",  # RAG
            "https://arxiv.org/pdf/2106.10504"   # Fine-Tuning Vision Transformers
        ]
        st.session_state['pdf_urls'] = "\n".join(example_urls)
    
    # Robo-Downloader
    url_input = st.text_area("Enter PDF URLs (one per line)", value=st.session_state.get('pdf_urls', ""), height=200)
    if st.button("Robo-Download 🤖"):
        urls = url_input.strip().split("\n")
        progress_bar = st.progress(0)
        status_text = st.empty()
        total_urls = len(urls)
        existing_pdfs = get_pdf_files()
        for idx, url in enumerate(urls):
            if url:
                output_path = pdf_url_to_filename(url)
                status_text.text(f"Fetching {idx + 1}/{total_urls}: {os.path.basename(output_path)}...")
                if output_path not in existing_pdfs:
                    if download_pdf(url, output_path):
                        st.session_state['downloaded_pdfs'][url] = output_path
                        logger.info(f"Downloaded PDF from {url} to {output_path}")
                        entry = f"Downloaded PDF: {output_path}"
                        if entry not in st.session_state['history']:
                            st.session_state['history'].append(entry)
                    else:
                        st.error(f"Failed to nab {url} 😿")
                else:
                    st.info(f"Already got {os.path.basename(output_path)}! Skipping... 🐾")
                    st.session_state['downloaded_pdfs'][url] = output_path
                progress_bar.progress((idx + 1) / total_urls)
        status_text.text("Robo-Download complete! 🚀")
        update_gallery()

    # PDF Gallery with Thumbnails and Checkboxes
    st.subheader("PDF Gallery 📖")
    downloaded_pdfs = list(st.session_state['downloaded_pdfs'].values())
    if downloaded_pdfs:
        cols_per_row = 3
        for i in range(0, len(downloaded_pdfs), cols_per_row):
            cols = st.columns(cols_per_row)
            for j, pdf_path in enumerate(downloaded_pdfs[i:i + cols_per_row]):
                with cols[j]:
                    doc = fitz.open(pdf_path)
                    page = doc[0]
                    pix = page.get_pixmap(matrix=fitz.Matrix(0.5, 0.5))  # Thumbnail at 50% scale
                    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                    st.image(img, caption=os.path.basename(pdf_path), use_container_width=True)
                    # Checkbox for SFT/Input use
                    checkbox_key = f"pdf_{pdf_path}"
                    st.session_state['pdf_checkboxes'][checkbox_key] = st.checkbox(
                        "Use for SFT/Input", 
                        value=st.session_state['pdf_checkboxes'].get(checkbox_key, False), 
                        key=checkbox_key
                    )
                    # Download and Delete Buttons
                    st.markdown(get_download_link(pdf_path, "application/pdf", "Snag It! 📥"), unsafe_allow_html=True)
                    if st.button("Zap It! 🗑️", key=f"delete_{pdf_path}"):
                        os.remove(pdf_path)
                        url_key = next((k for k, v in st.session_state['downloaded_pdfs'].items() if v == pdf_path), None)
                        if url_key:
                            del st.session_state['downloaded_pdfs'][url_key]
                        del st.session_state['pdf_checkboxes'][checkbox_key]
                        st.success(f"PDF {os.path.basename(pdf_path)} vaporized! 💨")
                        st.rerun()
                    doc.close()
    else:
        st.info("No PDFs captured yet. Feed the robo-downloader some URLs! 🤖")

    mode = st.selectbox("Snapshot Mode", ["Single Page (High-Res)", "Two Pages (High-Res)", "All Pages (Thumbnails)"], key="download_mode")
    if st.button("Snapshot Selected 📸"):
        selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)]
        if selected_pdfs:
            for pdf_path in selected_pdfs:
                mode_key = {"Single Page (High-Res)": "single", "Two Pages (High-Res)": "twopage", "All Pages (Thumbnails)": "allthumbs"}[mode]
                snapshots = asyncio.run(process_pdf_snapshot(pdf_path, mode_key))
                for snapshot in snapshots:
                    st.image(Image.open(snapshot), caption=snapshot, use_container_width=True)
        else:
            st.warning("No PDFs selected for snapshotting! Check some boxes first. 📝")

with tab3:
    st.header("Build Titan 🌱")
    model_type = st.selectbox("Model Type", ["Causal LM", "Diffusion"], key="build_type")
    base_model = st.selectbox("Select Tiny Model", 
        ["HuggingFaceTB/SmolLM-135M", "Qwen/Qwen1.5-0.5B-Chat"] if model_type == "Causal LM" else 
        ["OFA-Sys/small-stable-diffusion-v0", "stabilityai/stable-diffusion-2-base"])
    model_name = st.text_input("Model Name", f"tiny-titan-{int(time.time())}")
    domain = st.text_input("Target Domain", "general")
    if st.button("Download Model ⬇️"):
        config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(name=model_name, base_model=base_model, size="small", domain=domain)
        builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder()
        builder.load_model(base_model, config)
        builder.save_model(config.model_path)
        st.session_state['builder'] = builder
        st.session_state['model_loaded'] = True
        entry = f"Built {model_type} model: {model_name}"
        if entry not in st.session_state['history']:
            st.session_state['history'].append(entry)
        st.success(f"Model downloaded and saved to {config.model_path}! 🎉")
        st.rerun()

with tab4:
    st.header("Fine-Tune Titan 🔧")
    if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
        st.warning("Please build or load a Titan first! ⚠️")
    else:
        if isinstance(st.session_state['builder'], ModelBuilder):
            if st.button("Generate Sample CSV 📝"):
                sample_data = [
                    {"prompt": "What is AI?", "response": "AI is artificial intelligence, simulating human smarts in machines."},
                    {"prompt": "Explain machine learning", "response": "Machine learning is AI’s gym where models bulk up on data."},
                ]
                csv_path = f"sft_data_{int(time.time())}.csv"
                with open(csv_path, "w", newline="") as f:
                    writer = csv.DictWriter(f, fieldnames=["prompt", "response"])
                    writer.writeheader()
                    writer.writerows(sample_data)
                st.markdown(get_download_link(csv_path, "text/csv", "Download Sample CSV"), unsafe_allow_html=True)
                st.success(f"Sample CSV generated as {csv_path}! ✅")

            uploaded_csv = st.file_uploader("Upload CSV for SFT", type="csv")
            if uploaded_csv and st.button("Fine-Tune with Uploaded CSV 🔄"):
                csv_path = f"uploaded_sft_data_{int(time.time())}.csv"
                with open(csv_path, "wb") as f:
                    f.write(uploaded_csv.read())
                new_model_name = f"{st.session_state['builder'].config.name}-sft-{int(time.time())}"
                new_config = ModelConfig(name=new_model_name, base_model=st.session_state['builder'].config.base_model, size="small", domain=st.session_state['builder'].config.domain)
                st.session_state['builder'].config = new_config
                st.session_state['builder'].fine_tune_sft(csv_path)
                st.session_state['builder'].save_model(new_config.model_path)
                zip_path = f"{new_config.model_path}.zip"
                zip_directory(new_config.model_path, zip_path)
                entry = f"Fine-tuned Causal LM: {new_model_name}"
                if entry not in st.session_state['history']:
                    st.session_state['history'].append(entry)
                st.markdown(get_download_link(zip_path, "application/zip", "Download Fine-Tuned Titan"), unsafe_allow_html=True)
                st.rerun()
        elif isinstance(st.session_state['builder'], DiffusionBuilder):
            captured_files = get_gallery_files(["png"])
            selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)]
            if len(captured_files) + len(selected_pdfs) >= 2:
                demo_data = [{"image": img, "text": f"Superhero {os.path.basename(img).split('.')[0]}"} for img in captured_files]
                for pdf_path in selected_pdfs:
                    demo_data.append({"image": pdf_path, "text": f"PDF {os.path.basename(pdf_path)}"})
                edited_data = st.data_editor(pd.DataFrame(demo_data), num_rows="dynamic")
                if st.button("Fine-Tune with Dataset 🔄"):
                    images = [Image.open(row["image"]) if row["image"].endswith('.png') else Image.frombytes("RGB", fitz.open(row["image"])[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)).size, fitz.open(row["image"])[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)).samples) for _, row in edited_data.iterrows()]
                    texts = [row["text"] for _, row in edited_data.iterrows()]
                    new_model_name = f"{st.session_state['builder'].config.name}-sft-{int(time.time())}"
                    new_config = DiffusionConfig(name=new_model_name, base_model=st.session_state['builder'].config.base_model, size="small")
                    st.session_state['builder'].config = new_config
                    st.session_state['builder'].fine_tune_sft(images, texts)
                    st.session_state['builder'].save_model(new_config.model_path)
                    zip_path = f"{new_config.model_path}.zip"
                    zip_directory(new_config.model_path, zip_path)
                    entry = f"Fine-tuned Diffusion: {new_model_name}"
                    if entry not in st.session_state['history']:
                        st.session_state['history'].append(entry)
                    st.markdown(get_download_link(zip_path, "application/zip", "Download Fine-Tuned Diffusion Model"), unsafe_allow_html=True)
                csv_path = f"sft_dataset_{int(time.time())}.csv"
                with open(csv_path, "w", newline="") as f:
                    writer = csv.writer(f)
                    writer.writerow(["image", "text"])
                    for _, row in edited_data.iterrows():
                        writer.writerow([row["image"], row["text"]])
                st.markdown(get_download_link(csv_path, "text/csv", "Download SFT Dataset CSV"), unsafe_allow_html=True)

with tab5:
    st.header("Test Titan 🧪")
    if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
        st.warning("Please build or load a Titan first! ⚠️")
    else:
        if isinstance(st.session_state['builder'], ModelBuilder):
            if st.session_state['builder'].sft_data:
                st.write("Testing with SFT Data:")
                for item in st.session_state['builder'].sft_data[:3]:
                    prompt = item["prompt"]
                    expected = item["response"]
                    status_container = st.empty()
                    generated = st.session_state['builder'].evaluate(prompt, status_container)
                    st.write(f"**Prompt**: {prompt}")
                    st.write(f"**Expected**: {expected}")
                    st.write(f"**Generated**: {generated}")
                    st.write("---")
                    status_container.empty()
            test_prompt = st.text_area("Enter Test Prompt", "What is AI?")
            if st.button("Run Test ▶️"):
                status_container = st.empty()
                result = st.session_state['builder'].evaluate(test_prompt, status_container)
                entry = f"Causal LM Test: {test_prompt} -> {result}"
                if entry not in st.session_state['history']:
                    st.session_state['history'].append(entry)
                st.write(f"**Generated Response**: {result}")
                status_container.empty()
        elif isinstance(st.session_state['builder'], DiffusionBuilder):
            test_prompt = st.text_area("Enter Test Prompt", "Neon Batman")
            selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)]
            if st.button("Run Test ▶️"):
                image = st.session_state['builder'].generate(test_prompt)
                output_file = generate_filename("diffusion_test", "png")
                image.save(output_file)
                entry = f"Diffusion Test: {test_prompt} -> {output_file}"
                if entry not in st.session_state['history']:
                    st.session_state['history'].append(entry)
                st.image(image, caption="Generated Image")
                update_gallery()

with tab6:
    st.header("Agentic RAG Party 🌐")
    if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
        st.warning("Please build or load a Titan first! ⚠️")
    else:
        if isinstance(st.session_state['builder'], ModelBuilder):
            if st.button("Run NLP RAG Demo 🎉"):
                agent = PartyPlannerAgent(st.session_state['builder'].model, st.session_state['builder'].tokenizer)
                task = "Plan a luxury superhero-themed party at Wayne Manor."
                plan_df = agent.plan_party(task)
                entry = f"NLP RAG Demo: Planned party at Wayne Manor"
                if entry not in st.session_state['history']:
                    st.session_state['history'].append(entry)
                st.dataframe(plan_df)
        elif isinstance(st.session_state['builder'], DiffusionBuilder):
            if st.button("Run CV RAG Demo 🎉"):
                agent = CVPartyPlannerAgent(st.session_state['builder'].pipeline)
                task = "Generate images for a luxury superhero-themed party."
                plan_df = agent.plan_party(task)
                entry = f"CV RAG Demo: Generated party images"
                if entry not in st.session_state['history']:
                    st.session_state['history'].append(entry)
                st.dataframe(plan_df)
                for _, row in plan_df.iterrows():
                    image = agent.generate(row["Image Idea"])
                    output_file = generate_filename(f"cv_rag_{row['Theme'].lower()}", "png")
                    image.save(output_file)
                    st.image(image, caption=f"{row['Theme']} - {row['Image Idea']}")
                update_gallery()

with tab7:
    st.header("Test OCR 🔍")
    captured_files = get_gallery_files(["png"])
    selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)]
    all_files = captured_files + selected_pdfs
    if all_files:
        selected_file = st.selectbox("Select Image or PDF", all_files, key="ocr_select")
        if selected_file:
            if selected_file.endswith('.png'):
                image = Image.open(selected_file)
            else:
                doc = fitz.open(selected_file)
                pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
                image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                doc.close()
            st.image(image, caption="Input Image", use_container_width=True)
            if st.button("Run OCR 🚀", key="ocr_run"):
                output_file = generate_filename("ocr_output", "txt")
                st.session_state['processing']['ocr'] = True
                result = asyncio.run(process_ocr(image, output_file))
                entry = f"OCR Test: {selected_file} -> {output_file}"
                if entry not in st.session_state['history']:
                    st.session_state['history'].append(entry)
                st.text_area("OCR Result", result, height=200, key="ocr_result")
                st.success(f"OCR output saved to {output_file}")
                st.session_state['processing']['ocr'] = False
    else:
        st.warning("No images or PDFs captured yet. Use Camera Snap or Download PDFs first!")

with tab8:
    st.header("Test Image Gen 🎨")
    captured_files = get_gallery_files(["png"])
    selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)]
    all_files = captured_files + selected_pdfs
    if all_files:
        selected_file = st.selectbox("Select Image or PDF", all_files, key="gen_select")
        if selected_file:
            if selected_file.endswith('.png'):
                image = Image.open(selected_file)
            else:
                doc = fitz.open(selected_file)
                pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
                image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                doc.close()
            st.image(image, caption="Reference Image", use_container_width=True)
            prompt = st.text_area("Prompt", "Generate a similar superhero image", key="gen_prompt")
            if st.button("Run Image Gen 🚀", key="gen_run"):
                output_file = generate_filename("gen_output", "png")
                st.session_state['processing']['gen'] = True
                result = asyncio.run(process_image_gen(prompt, output_file))
                entry = f"Image Gen Test: {prompt} -> {output_file}"
                if entry not in st.session_state['history']:
                    st.session_state['history'].append(entry)
                st.image(result, caption="Generated Image", use_container_width=True)
                st.success(f"Image saved to {output_file}")
                st.session_state['processing']['gen'] = False
    else:
        st.warning("No images or PDFs captured yet. Use Camera Snap or Download PDFs first WAV!")

with tab9:
    st.header("Custom Diffusion 🎨🤓")
    st.write("Unleash your inner artist with our tiny diffusion models!")
    captured_files = get_gallery_files(["png"])
    selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)]
    all_files = captured_files + selected_pdfs
    if all_files:
        st.subheader("Select Images or PDFs to Train")
        selected_files = st.multiselect("Pick Images or PDFs", all_files, key="diffusion_select")
        images = []
        for file in selected_files:
            if file.endswith('.png'):
                images.append(Image.open(file))
            else:
                doc = fitz.open(file)
                pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
                images.append(Image.frombytes("RGB", [pix.width, pix.height], pix.samples))
                doc.close()
        
        model_options = [
            ("PixelTickler 🎨✨", "OFA-Sys/small-stable-diffusion-v0"),
            ("DreamWeaver 🌙🖌️", "stabilityai/stable-diffusion-2-base"),
            ("TinyArtBot 🤖🖼️", "custom")
        ]
        model_choice = st.selectbox("Choose Your Diffusion Dynamo", [opt[0] for opt in model_options], key="diffusion_model")
        model_name = next(opt[1] for opt in model_options if opt[0] == model_choice)

        if st.button("Train & Generate 🚀", key="diffusion_run"):
            output_file = generate_filename("custom_diffusion", "png")
            st.session_state['processing']['diffusion'] = True
            if model_name == "custom":
                result = asyncio.run(process_custom_diffusion(images, output_file, model_choice))
            else:
                builder = DiffusionBuilder()
                builder.load_model(model_name)
                result = builder.generate("A superhero scene inspired by captured images")
                result.save(output_file)
            entry = f"Custom Diffusion: {model_choice} -> {output_file}"
            if entry not in st.session_state['history']:
                st.session_state['history'].append(entry)
            st.image(result, caption=f"{model_choice} Masterpiece", use_container_width=True)
            st.success(f"Image saved to {output_file}")
            st.session_state['processing']['diffusion'] = False
    else:
        st.warning("No images or PDFs captured yet. Use Camera Snap or Download PDFs first!")

# Initial Gallery Update
update_gallery()