import gradio as gr
from transformers import T5Tokenizer, T5ForConditionalGeneration
import torch
import re

print("Loading Green Prompt Optimizer model...")
MODEL_PATH = "./models/prompt_optimizer"

tokenizer = T5Tokenizer.from_pretrained(MODEL_PATH)
model = T5ForConditionalGeneration.from_pretrained(MODEL_PATH)
model.eval()

print(f"✓ Model loaded successfully")
print(f"✓ Tokenizer vocab size: {len(tokenizer)}")
print(f"✓ Model parameters: {model.num_parameters()}")

def smart_fallback_optimize(prompt):
    """Fallback optimization when model fails"""
    fillers = {
        'please', 'could', 'would', 'can', 'you', 'the', 'a', 'an',
        'very', 'really', 'quite', 'just', 'actually', 'basically',
        'literally', 'honestly', 'i think', 'in my opinion',
        'it seems', 'kind of', 'sort of', 'i want', 'i would like'
    }
    
    replacements = {
        r'\bcan you please\b': '',
        r'\bcould you please\b': '',
        r'\bwould you please\b': '',
        r'\bi would like to\b': '',
        r'\bi want to\b': '',
        r'\bhelp me\b': '',
        r'\bfor me\b': '',
    }
    
    optimized = prompt.lower()
    
    for pattern, replacement in replacements.items():
        optimized = re.sub(pattern, replacement, optimized)
    
    words = optimized.split()
    words = [w for w in words if w not in fillers]
    optimized = ' '.join(words).strip()
    
    if optimized:
        optimized = optimized[0].upper() + optimized[1:]
    
    if not optimized or len(optimized) < 5:
        words = prompt.split()
        optimized = ' '.join(words[:8])
    
    return optimized

def optimize_prompt(prompt, preserve_meaning=True):
    """Optimize prompt to reduce tokens"""
    if not prompt or not prompt.strip():
        return "", "0 → 0", 0, "0%", "0.000000 Wh", "0.000000 g"
    
    print(f"\n=== OPTIMIZING ===")
    print(f"Input: {prompt[:100]}")
    
    try:
        input_text = "optimize: " + prompt
        input_ids = tokenizer.encode(input_text, return_tensors="pt", truncation=True, max_length=512)
        
        with torch.no_grad():
            outputs = model.generate(
                input_ids,
                max_new_tokens=128,
                num_beams=4,
                no_repeat_ngram_size=2,
                early_stopping=True,
                temperature=0.7,
                do_sample=False,
                repetition_penalty=1.2
            )
        
        optimized = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        original_tokens = len(tokenizer.encode(prompt))
        optimized_tokens = len(tokenizer.encode(optimized))
        reduction_rate = (original_tokens - optimized_tokens) / original_tokens
        
        if reduction_rate < 0.1 or not optimized or len(optimized.strip()) < 3:
            print("⚠ Model weak, using fallback")
            optimized = smart_fallback_optimize(prompt)
        else:
            print(f"✓ Model success")
        
    except Exception as e:
        print(f"⚠ Error: {e}, using fallback")
        optimized = smart_fallback_optimize(prompt)
    
    print(f"Output: {optimized}")
    
    original_tokens = len(tokenizer.encode(prompt))
    optimized_tokens = len(tokenizer.encode(optimized))
    tokens_saved = max(0, original_tokens - optimized_tokens)
    reduction = (tokens_saved / max(original_tokens, 1)) * 100
    energy = tokens_saved * 0.000002
    co2 = energy * 475 / 1000
    
    print(f"Tokens: {original_tokens} → {optimized_tokens} (saved {tokens_saved})")
    print(f"=== DONE ===\n")
    
    return (
        optimized,
        f"{original_tokens} → {optimized_tokens}",
        tokens_saved,
        f"{reduction:.2f}%",
        f"{energy:.6f} Wh",
        f"{co2:.6f} g"
    )

# Load custom CSS
with open("custom.css", "r") as f:
    custom_css = f.read()

# Create interface
with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("# 🌱 Green Prompts Optimizer")
    gr.Markdown("### Reduce AI carbon footprint by optimizing your prompts")
    
    with gr.Row():
        with gr.Column():
            prompt_box = gr.Textbox(
                label="📝 Enter Your Prompt",
                placeholder="Type or paste your AI prompt here...",
                lines=6
            )
            preserve_box = gr.Checkbox(
                label="Preserve Meaning",
                value=True,
                info="Maintain semantic meaning while reducing tokens"
            )
            btn = gr.Button("🚀 Optimize Prompt", variant="primary")
    
    optimized_box = gr.Textbox(
        label="✨ Optimized Prompt",
        lines=6,
        interactive=False
    )
    
    gr.Markdown("### 📊 Optimization Results")
    
    with gr.Row():
        token_info = gr.Textbox(label="📈 Token Count", interactive=False)
        tokens_saved = gr.Textbox(label="💾 Tokens Saved", interactive=False)
    
    with gr.Row():
        reduction = gr.Textbox(label="📉 Reduction %", interactive=False)
        energy = gr.Textbox(label="⚡ Energy Saved", interactive=False)
    
    co2 = gr.Textbox(label="🌍 CO₂ Reduced", interactive=False)
    
    btn.click(
        optimize_prompt,
        inputs=[prompt_box, preserve_box],
        outputs=[optimized_box, token_info, tokens_saved, reduction, energy, co2]
    )
    
    gr.Markdown("---")
    gr.Markdown("✅ **Built with 💚 for a greener AI future | Powered by T5**")

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