import streamlit as st
from PIL import Image
import torch
import json
from transformers import (
    DonutProcessor, 
    VisionEncoderDecoderModel,
    LayoutLMv3Processor, 
    LayoutLMv3ForSequenceClassification,
    BrosProcessor,
    BrosForTokenClassification,
    LlavaProcessor,
    LlavaForConditionalGeneration
)

# Cache the model loading to improve performance
@st.cache_resource
def load_model(model_name):
    """Load the selected model and processor"""
    try:
        if model_name == "Donut":
            processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base")
            model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base")
            # Configure Donut specific parameters
            model.config.decoder_start_token_id = processor.tokenizer.bos_token_id
            model.config.pad_token_id = processor.tokenizer.pad_token_id
            model.config.vocab_size = len(processor.tokenizer)
            
        elif model_name == "LayoutLMv3":
            processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base")
            model = LayoutLMv3ForSequenceClassification.from_pretrained("microsoft/layoutlmv3-base")
            
        elif model_name == "BROS":
            processor = BrosProcessor.from_pretrained("microsoft/bros-base")
            model = BrosForTokenClassification.from_pretrained("microsoft/bros-base")
            
        elif model_name == "LLaVA-1.5":
            processor = LlavaProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
            model = LlavaForConditionalGeneration.from_pretrained("llava-hf/llava-1.5-7b-hf")
        
        return model, processor
    except Exception as e:
        st.error(f"Error loading model {model_name}: {str(e)}")
        return None, None

def analyze_document(image, model_name, model, processor):
    """Analyze document using selected model"""
    try:
        # Process image according to model requirements
        if model_name == "Donut":
            # Prepare input with task prompt
            pixel_values = processor(image, return_tensors="pt").pixel_values
            task_prompt = "<s_cord>analyze the document and extract information</s_cord>"
            decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
            
            # Generate output with improved parameters
            outputs = model.generate(
                pixel_values,
                decoder_input_ids=decoder_input_ids,
                max_length=512,
                early_stopping=True,
                pad_token_id=processor.tokenizer.pad_token_id,
                eos_token_id=processor.tokenizer.eos_token_id,
                use_cache=True,
                num_beams=4,
                bad_words_ids=[[processor.tokenizer.unk_token_id]],
                return_dict_in_generate=True
            )
            
            # Process and clean the output
            sequence = processor.batch_decode(outputs.sequences)[0]
            sequence = sequence.replace(task_prompt, "").replace("</s_cord>", "").strip()
            
            # Try to parse as JSON, fallback to raw text
            try:
                result = json.loads(sequence)
            except json.JSONDecodeError:
                result = {"raw_text": sequence}
                
        elif model_name == "LayoutLMv3":
            inputs = processor(image, return_tensors="pt")
            outputs = model(**inputs)
            result = {"logits": outputs.logits.tolist()}  # Convert tensor to list for JSON serialization
            
        elif model_name == "BROS":
            inputs = processor(image, return_tensors="pt")
            outputs = model(**inputs)
            result = {"predictions": outputs.logits.tolist()}
            
        elif model_name == "LLaVA-1.5":
            inputs = processor(image, return_tensors="pt")
            outputs = model.generate(**inputs, max_length=256)
            result = {"generated_text": processor.decode(outputs[0], skip_special_tokens=True)}
        
        return result
        
    except Exception as e:
        error_msg = str(e)
        st.error(f"Error analyzing document: {error_msg}")
        return {"error": error_msg, "type": "analysis_error"}

# Set page config with improved layout
st.set_page_config(
    page_title="Document Analysis Comparison",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Add custom CSS for better styling
st.markdown("""
    <style>
        .stAlert {
            margin-top: 1rem;
        }
        .upload-text {
            font-size: 1.2rem;
            margin-bottom: 1rem;
        }
        .model-info {
            padding: 1rem;
            border-radius: 0.5rem;
            background-color: #f8f9fa;
        }
    </style>
""", unsafe_allow_html=True)

# Title and description
st.title("Document Understanding Model Comparison")
st.markdown("""
Compare different models for document analysis and understanding.
Upload an image and select a model to analyze it.
""")

# Create two columns for layout
col1, col2 = st.columns([1, 1])

with col1:
    # File uploader with improved error handling
    uploaded_file = st.file_uploader(
        "Choose a document image",
        type=['png', 'jpg', 'jpeg', 'pdf'],
        help="Supported formats: PNG, JPEG, PDF"
    )
    
    if uploaded_file is not None:
        try:
            # Display uploaded image
            image = Image.open(uploaded_file)
            st.image(image, caption='Uploaded Document', use_column_width=True)
        except Exception as e:
            st.error(f"Error loading image: {str(e)}")

with col2:
    # Model selection with detailed information
    model_info = {
        "Donut": {
            "description": "Best for structured OCR and document format understanding",
            "memory": "6-8GB",
            "strengths": ["Structured OCR", "Memory efficient", "Good with fixed formats"],
            "best_for": ["Invoices", "Forms", "Structured documents"]
        },
        "LayoutLMv3": {
            "description": "Strong layout understanding with reasoning capabilities",
            "memory": "12-15GB",
            "strengths": ["Layout understanding", "Reasoning", "Pre-trained knowledge"],
            "best_for": ["Complex layouts", "Mixed content", "Tables"]
        },
        "BROS": {
            "description": "Memory efficient with fast inference",
            "memory": "4-6GB",
            "strengths": ["Fast inference", "Memory efficient", "Easy fine-tuning"],
            "best_for": ["Simple documents", "Quick analysis", "Basic OCR"]
        },
        "LLaVA-1.5": {
            "description": "Comprehensive OCR with strong reasoning",
            "memory": "25-40GB",
            "strengths": ["Strong reasoning", "Zero-shot capable", "Visual understanding"],
            "best_for": ["Complex documents", "Natural language understanding", "Visual QA"]
        }
    }
    
    selected_model = st.selectbox(
        "Select Model",
        list(model_info.keys())
    )
    
    # Display enhanced model information
    st.markdown("### Model Details")
    with st.expander("Model Information", expanded=True):
        st.markdown(f"**Description:** {model_info[selected_model]['description']}")
        st.markdown(f"**Memory Required:** {model_info[selected_model]['memory']}")
        st.markdown("**Strengths:**")
        for strength in model_info[selected_model]['strengths']:
            st.markdown(f"- {strength}")
        st.markdown("**Best For:**")
        for use_case in model_info[selected_model]['best_for']:
            st.markdown(f"- {use_case}")

# Analysis section with improved error handling and progress tracking
if uploaded_file is not None and selected_model:
    if st.button("Analyze Document", help="Click to start document analysis"):
        with st.spinner('Loading model and analyzing document...'):
            try:
                # Create a progress bar
                progress_bar = st.progress(0)
                
                # Load model with progress update
                progress_bar.progress(25)
                st.info("Loading model...")
                model, processor = load_model(selected_model)
                
                if model is None or processor is None:
                    st.error("Failed to load model. Please try again.")
                else:
                    # Update progress
                    progress_bar.progress(50)
                    st.info("Analyzing document...")
                    
                    # Analyze document
                    results = analyze_document(image, selected_model, model, processor)
                    
                    # Update progress
                    progress_bar.progress(75)
                    
                    # Display results with proper formatting
                    st.markdown("### Analysis Results")
                    if isinstance(results, dict) and "error" in results:
                        st.error(f"Analysis Error: {results['error']}")
                    else:
                        # Pretty print the results
                        st.json(results)
                    
                    # Complete progress
                    progress_bar.progress(100)
                    st.success("Analysis completed!")
                
            except Exception as e:
                st.error(f"Error during analysis: {str(e)}")
                st.error("Please try with a different image or model.")

# Add improved information about usage and limitations
st.markdown("""
---
### Usage Notes:
- Different models excel at different types of documents
- Processing time and memory requirements vary by model
- Image quality significantly affects results
- Some models may require specific document formats
""")

# Add performance metrics section
if st.checkbox("Show Performance Metrics"):
    st.markdown("""
    ### Model Performance Metrics
    | Model | Avg. Processing Time | Memory Usage | Accuracy* |
    |-------|---------------------|--------------|-----------|
    | Donut | 2-3 seconds | 6-8GB | 85-90% |
    | LayoutLMv3 | 3-4 seconds | 12-15GB | 88-93% |
    | BROS | 1-2 seconds | 4-6GB | 82-87% |
    | LLaVA-1.5 | 4-5 seconds | 25-40GB | 90-95% |
    
    *Accuracy varies based on document type and quality
    """)

# Add a footer with version and contact information
st.markdown("---")
st.markdown("""
v1.1 - Created with Streamlit
\nFor issues or feedback, please visit our [GitHub repository](https://github.com/yourusername/doc-analysis)
""")

# Add model selection guidance
if st.checkbox("Show Model Selection Guide"):
    st.markdown("""
    ### How to Choose the Right Model
    1. **Donut**: Choose for structured documents with clear layouts
    2. **LayoutLMv3**: Best for documents with complex layouts and relationships
    3. **BROS**: Ideal for quick analysis and simple documents
    4. **LLaVA-1.5**: Perfect for complex documents requiring deep understanding
    """)