#!/usr/bin/env python3
"""
Gradio demo for T5 Email Summarizer with better preprocessing and separate fields
Deployed on HuggingFace Spaces with T4 GPU
"""
import gradio as gr
import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer
import time
import re

# Load model and tokenizer
print("Loading T5 Email Summarizer model...")
model_name = "wordcab/t5-small-email-summarizer"

tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(
    model_name,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)

# Move to GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()

print(f"Model loaded successfully on {device}!")

def normalize_titles(text):
    """
    Normalize titles by removing periods to avoid tokenization issues.
    This is a general solution that handles Mr. Ms. Dr. Prof. etc.
    """
    # List of common titles that cause issues when followed by a period
    titles_with_period = [
        'Mr.', 'Ms.', 'Mrs.', 'Dr.', 'Prof.', 
        'Sr.', 'Jr.', 'Ph.D.', 'M.D.', 'B.A.', 'M.A.', 'B.S.', 'M.S.',
        'Rev.', 'Hon.', 'Pres.', 'Gov.', 'Ofc.', 'Msgr.',
        'Fr.', 'Br.', 'Sr.', 'Mx.'
    ]
    
    normalized = text
    for title in titles_with_period:
        # Replace the title with period with the title without period
        title_no_period = title.rstrip('.')
        # Use word boundary to avoid replacing parts of words
        normalized = re.sub(r'\b' + re.escape(title) + r'\b', title_no_period, normalized)
    
    return normalized

def clean_unicode(text):
    """Clean up special unicode characters that can cause issues"""
    # Normalize quotes
    text = text.replace('"', '"').replace('"', '"')
    text = text.replace(''', "'").replace(''', "'")
    text = text.replace('–', '-').replace('—', '-')
    text = text.replace('…', '...')
    
    # Remove zero-width spaces and other invisible characters
    text = re.sub(r'[\u200b\u200c\u200d\ufeff]', '', text)
    
    return text

def preprocess_email(subject, body, mode="brief"):
    """
    Preprocess email with general normalization
    """
    # Clean unicode in both subject and body
    if subject:
        subject = clean_unicode(subject)
        subject = normalize_titles(subject)
        
        # For brief mode, simplify long subjects with names
        # These confuse the model in brief mode
        if mode == "brief" and len(subject) > 100:
            # If it's a RE: or FW: with a long chain, try to simplify
            if subject.startswith(('RE:', 'Re:', 'FW:', 'Fw:')):
                # Extract key parts (hotel name, booking number, date)
                parts = []
                if 'Mia Saigon' in subject:
                    parts.append('Mia Saigon Hotel')
                if 'birthday' in subject.lower() or 'Birthday' in subject:
                    parts.append('Birthday Celebration')
                elif 'booking' in subject.lower():
                    parts.append('Booking')
                
                # Extract date if present
                import re
                date_match = re.search(r'\d{1,2}\s+\w+\s+\d{4}', subject)
                if date_match:
                    parts.append(date_match.group())
                
                if parts:
                    subject = ' - '.join(parts)
                else:
                    # Fallback: just take first 50 chars
                    subject = subject[:50] + '...'
    
    if body:
        body = clean_unicode(body)
        body = normalize_titles(body)
        
        # For brief mode, remove greeting lines that cause issues
        if mode == "brief":
            lines = body.strip().split('\n')
            result_lines = []
            skip_mode = False
            
            for i, line in enumerate(lines):
                line_stripped = line.strip()
                
                # Check if this is a greeting line at the beginning
                if i == 0 and line_stripped.lower().startswith(('dear', 'hi', 'hello', 'good morning', 'good afternoon', 'good evening')):
                    skip_mode = True
                    continue
                
                # Skip empty lines right after greeting
                if skip_mode and not line_stripped:
                    continue
                
                # Once we hit real content, stop skipping
                if line_stripped and skip_mode:
                    skip_mode = False
                
                if not skip_mode:
                    result_lines.append(line)
            
            if result_lines:
                body = '\n'.join(result_lines).strip()
    
    return subject, body

def summarize_email(subject, body, summary_type, temperature=0.7, max_length=150):
    """
    Generate email summary based on selected type
    """
    # Check if we have content
    if not body and not subject:
        return "Please enter email content (subject and/or body) to summarize.", 0, ""
    
    # If only subject is provided
    if subject and not body:
        body = subject
        subject = ""
    
    start_time = time.time()
    
    # Determine mode and parameters
    if summary_type == "Brief (1-2 sentences)":
        mode = "brief"
        prefix = "summarize_brief:"
        max_gen_length = 50
    elif summary_type == "Full (detailed)":
        mode = "full"
        prefix = "summarize_full:"
        max_gen_length = max_length
    else:  # Auto
        # Use brief for short emails, full for longer ones
        total_words = len((subject + " " + body).split())
        if total_words < 100:
            mode = "brief"
            prefix = "summarize_brief:"
            max_gen_length = 50
        else:
            mode = "full"
            prefix = "summarize_full:"
            max_gen_length = max_length
    
    # Preprocess the email
    original_subject = subject
    original_body = body
    processed_subject, processed_body = preprocess_email(subject, body, mode)
    
    # Track what preprocessing was done
    preprocessing_notes = []
    if original_subject != processed_subject:
        if len(original_subject) > 100 and len(processed_subject) < len(original_subject):
            preprocessing_notes.append("Simplified long subject")
        else:
            preprocessing_notes.append("Normalized titles in subject")
    
    if original_body != processed_body:
        if original_body.lower().startswith(('dear', 'hi', 'hello')) and not processed_body.lower().startswith(('dear', 'hi', 'hello')):
            preprocessing_notes.append("Removed greeting line")
        else:
            preprocessing_notes.append("Normalized titles in body")
    
    # Format input for the model
    if processed_subject:
        input_text = f"{prefix} Subject: {processed_subject}. Body: {processed_body}"
    else:
        input_text = f"{prefix} Subject: Email. Body: {processed_body}"
    
    # Tokenize
    inputs = tokenizer(
        input_text,
        max_length=512,
        truncation=True,
        return_tensors="pt"
    ).to(device)
    
    # Generate summary
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=max_gen_length,
            min_length=10,
            temperature=temperature,
            do_sample=temperature > 0,
            top_p=0.9,
            num_beams=2 if temperature == 0 else 1,
            early_stopping=True,
            no_repeat_ngram_size=3
        )
    
    # Decode
    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Calculate metrics
    processing_time = time.time() - start_time
    input_tokens = len(inputs['input_ids'][0])
    output_tokens = len(outputs[0])
    
    # Add metadata
    metadata = f"\n\n---\n📊 **Metrics:**\n"
    metadata += f"- Processing time: {processing_time:.2f}s\n"
    metadata += f"- Input tokens: {input_tokens}/512\n"
    metadata += f"- Output tokens: {output_tokens}\n"
    metadata += f"- Summary type: {mode.title()}\n"
    
    # Note about preprocessing
    if preprocessing_notes:
        metadata += f"- Preprocessing: {', '.join(preprocessing_notes)}\n"
    
    return summary, processing_time, metadata

# Example emails
examples = [
    [
        "Quarterly Budget Review Meeting",
        """Dear Team,

I hope this email finds you well. I wanted to remind everyone about our quarterly budget review meeting scheduled for next Tuesday, March 15th at 2:00 PM EST in Conference Room A.

Please come prepared with:
- Q1 expense reports
- Updated project timelines
- Resource allocation requests for Q2

We'll be discussing the 15% budget increase for digital marketing initiatives and the proposed headcount expansion for the engineering team.

If you cannot attend in person, please join via Zoom using the link in the calendar invite.

Best regards,
Sarah Johnson
Finance Director""",
        "Auto-detect",
        0.7,
        150
    ],
    [
        "",
        """hey team,

quick update - cant make the meeting tmrw bc im stuck at the airport (flight delayed AGAIN ugh). 

jim said we need to finalize teh proposal by friday or we'll miss the deadline... can someone take over? also dont forget to include the budget numbers from last months report.

btw has anyone seen my laptop charger? left it somewhere in the office yesterday lol

thx
mike""",
        "Brief (1-2 sentences)",
        0.7,
        150
    ],
    [
        "Research Collaboration Opportunity",
        """Dear Dr. Williams,

I hope this message finds you well. I'm writing to follow up on our recent discussion about the research collaboration opportunity.

As we discussed, our lab has extensive experience in computational biology and we believe there could be significant synergies with your work in genomics. We have secured funding for a 3-year project and are looking for partners.

Would you be available for a call next week to discuss the details? I can share the full proposal and budget breakdown then.

Looking forward to your response.

Best regards,
Prof. Sarah Chen
Department of Computer Science""",
        "Full (detailed)",
        0.7,
        150
    ]
]

# Create Gradio interface
with gr.Blocks(title="T5 Email Summarizer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 📧 T5 Email Summarizer - Brief & Full (v3)
    
    This model can generate both **brief** (1-2 sentences) and **full** (detailed) summaries of emails.
    It's robust to messy, informal text with typos and abbreviations.
    
    **🔧 v3 Updates:** 
    - Separate Subject/Body fields for better structure
    - General title normalization (Mr. → Mr, Dr. → Dr, etc.)
    - Improved unicode handling
    - Better preprocessing for all edge cases
    
    🤖 **Model:** [wordcab/t5-small-email-summarizer](https://huggingface.co/wordcab/t5-small-email-summarizer)
    | 📊 **Dataset:** [argilla/FinePersonas-Conversations-Email-Summaries](https://huggingface.co/datasets/argilla/FinePersonas-Conversations-Email-Summaries)
    | 🚀 **Running on:** CPU (Free tier)
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            subject_input = gr.Textbox(
                label="📌 Subject Line (Optional)",
                placeholder="e.g., Meeting Tomorrow, Project Update, etc.",
                lines=1
            )
            
            body_input = gr.Textbox(
                label="📝 Email Body",
                placeholder="Paste or type your email content here...\n\nThe model handles formal/informal, clean/messy text equally well.",
                lines=10
            )
            
            with gr.Row():
                summary_type = gr.Radio(
                    choices=["Auto-detect", "Brief (1-2 sentences)", "Full (detailed)"],
                    value="Auto-detect",
                    label="📊 Summary Type"
                )
            
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                temperature = gr.Slider(
                    minimum=0,
                    maximum=1,
                    value=0.7,
                    step=0.1,
                    label="Temperature (0 = deterministic, 1 = creative)"
                )
                max_length = gr.Slider(
                    minimum=50,
                    maximum=200,
                    value=150,
                    step=10,
                    label="Max Length (for full summaries)"
                )
            
            summarize_btn = gr.Button("✨ Generate Summary", variant="primary")
        
        with gr.Column(scale=1):
            output = gr.Textbox(
                label="📋 Summary",
                lines=8,
                interactive=False
            )
            
            processing_time = gr.Number(
                label="⏱️ Processing Time (seconds)",
                precision=2,
                interactive=False,
                visible=False
            )
            
            info_box = gr.Markdown(
                label="📊 Processing Info",
                value=""
            )
    
    gr.Markdown("### 💡 Try these examples:")
    
    gr.Examples(
        examples=examples,
        inputs=[subject_input, body_input, summary_type, temperature, max_length],
        outputs=[output, processing_time, info_box],
        fn=summarize_email,
        cache_examples=False
    )
    
    summarize_btn.click(
        fn=summarize_email,
        inputs=[subject_input, body_input, summary_type, temperature, max_length],
        outputs=[output, processing_time, info_box]
    )
    
    gr.Markdown("""
    ---
    ### 📖 How to use:
    1. **Enter Subject** (optional) and **Email Body** separately for best results
    2. **Select summary type** or use Auto-detect
    3. **Click Generate Summary** to get your summary
    
    ### 🎯 Features:
    - **Dual-mode**: Get brief or detailed summaries on demand
    - **Robust**: Handles typos, abbreviations, and informal language
    - **Smart normalization**: Automatically handles titles (Mr., Dr., Prof., etc.)
    - **Fast**: Processes emails quickly even on CPU
    
    ### 🔧 Preprocessing Features:
    - **Title Normalization**: Converts "Mr." → "Mr", "Dr." → "Dr" to avoid tokenization issues
    - **Unicode Cleaning**: Handles special quotes, dashes, and invisible characters
    - **Smart Structure**: Separate subject/body fields for optimal processing
    
    ### 🔧 API Usage:
    ```python
    from transformers import pipeline
    
    summarizer = pipeline("summarization", model="wordcab/t5-small-email-summarizer")
    
    # For production, normalize titles first:
    import re
    def normalize_titles(text):
        titles = ['Mr.', 'Ms.', 'Dr.', 'Prof.']
        for title in titles:
            text = text.replace(title, title.rstrip('.'))
        return text
    
    email = normalize_titles(your_email)
    
    # Brief summary
    result = summarizer(f"summarize_brief: Subject: {subject}. Body: {body}")
    
    # Full summary  
    result = summarizer(f"summarize_full: Subject: {subject}. Body: {body}")
    ```
    
    ### 📚 Citation:
    ```bibtex
    @misc{wordcab2025t5email,
      title={T5 Email Summarizer - Brief & Full},
      author={Wordcab Team},
      year={2025},
      publisher={HuggingFace}
    }
    ```
    """)

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