import streamlit as st
import google.generativeai as genai
from PIL import Image
import os
from dotenv import load_dotenv
import PyPDF2
import io
from datetime import datetime

# Page configuration
st.set_page_config(
    page_title="Cornea AI Pentacam Analyzer",
    page_icon="👁️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Load environment variables
load_dotenv()

# Configure Gemini API
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
model = genai.GenerativeModel("gemini-2.0-flash-exp")

# Custom CSS
st.markdown("""
    <style>
    .main {
        padding: 2rem;
    }
    .stButton>button {
        width: 100%;
        background-color: #2E86C1;
        color: white;
        padding: 0.5rem;
        margin-top: 1rem;
    }
    .credit-box {
        background-color: #f0f2f6;
        padding: 1.5rem;
        border-radius: 0.5rem;
        margin: 1rem 0;
        border-left: 5px solid #2E86C1;
    }
    .header-box {
        background: linear-gradient(135deg, #2E86C1, #3498DB);
        padding: 2rem;
        border-radius: 0.5rem;
        color: white;
        margin-bottom: 2rem;
        text-align: center;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    }
    .image-container {
        margin: 1rem 0;
        padding: 1rem;
        border-radius: 0.5rem;
        background-color: #f8f9fa;
        border: 1px solid #e9ecef;
    }
    .analysis-container {
        margin-top: 1rem;
        padding: 1.5rem;
        border-radius: 0.5rem;
        background-color: #f8f9fa;
        border: 1px solid #e9ecef;
    }
    .patient-info {
        background-color: #fff;
        padding: 1.5rem;
        border-radius: 0.5rem;
        border: 1px solid #e9ecef;
        margin-bottom: 1rem;
    }
    .upload-section {
        background-color: #f8f9fa;
        padding: 1.5rem;
        border-radius: 0.5rem;
        border: 1px dashed #2E86C1;
        margin: 1rem 0;
    }
    .info-box {
        background-color: #e1f5fe;
        padding: 1rem;
        border-radius: 0.5rem;
        margin: 0.5rem 0;
        border-left: 3px solid #03a9f4;
    }
    </style>
""", unsafe_allow_html=True)

# System prompts
CORNEA_ANALYSIS_PROMPT = """You are an expert ophthalmologist specializing in corneal diseases. A user will ask you questions about interpreting corneal imaging, specifically Pentacam, and about subclinical corneal edema.

**Part 1: Pentacam Interpretation**
• Analyze Pentacam imaging for diagnosing and monitoring corneal diseases
• Assess key parameters: corneal shape, thickness, and elevation
• Evaluate anterior and posterior corneal surfaces
• Consider ABCD staging system and Cruces University Hospital progression score

**Part 2: Subclinical Corneal Edema Analysis**
• Evaluate for Fuchs Endothelial Corneal Dystrophy (FECD)
• Apply Sun criteria for subclinical corneal edema
• Assess displacement of thinnest point
• Check for loss of parallel isopachs
• Evaluate focal depression of posterior surface

Please analyze the provided images and provide a detailed assessment."""

COMPARISON_PROMPT = """Compare the provided Pentacam scans from different timepoints, focusing on:

1. Changes in corneal parameters:
   • Thickness variations
   • Topographic changes
   • Elevation differences
   • Progression indicators

2. Clinical significance:
   • Disease progression/stability
   • Treatment response
   • Risk assessment

3. Recommendations:
   • Follow-up interval
   • Treatment modifications
   • Additional testing needs"""

def extract_pdf_text(pdf_file):
    """Extract text from uploaded PDF file"""
    pdf_reader = PyPDF2.PdfReader(pdf_file)
    text = ""
    for page in pdf_reader.pages:
        text += page.extract_text()
    return text

def analyze_cornea_images(images, timepoint=None, patient_data=None):
    """Analyze corneal images with optional timepoint and patient data"""
    if timepoint:
        prompt = f"{CORNEA_ANALYSIS_PROMPT}\n\nTimepoint: {timepoint}\n"
    else:
        prompt = f"{CORNEA_ANALYSIS_PROMPT}\n"
    
    if patient_data:
        prompt += f"\nPatient Information:\n{patient_data}\n"
    
    prompt += "\nPlease analyze these corneal scans:"
    content = [prompt] + images
    response = model.generate_content(content)
    return response.text

def compare_timepoints(images1, date1, images2, date2, patient_data=None):
    """Compare corneal images from two timepoints"""
    prompt = f"{COMPARISON_PROMPT}\n\nTimepoint 1: {date1}\nTimepoint 2: {date2}\n"
    
    if patient_data:
        prompt += f"\nPatient Information:\n{patient_data}\n"
    
    prompt += "\nPlease compare these corneal scans:"
    content = [prompt] + images1 + images2
    response = model.generate_content(content)
    return response.text

def main():
    # Header
    st.markdown("""
        <div class="header-box">
            <h1>Cornea AI Pentacam Analyzer</h1>
            <p style="font-size: 1.2em; margin-top: 1rem;">Advanced Corneal Analysis & Diagnostics</p>
        </div>
    """, unsafe_allow_html=True)

    # Credits
    st.markdown("""
        <div class="credit-box">
            <h3>About</h3>
            <p>Developed by Dr. Verónica Gómez Calleja</p>
            <p>Cornea Specialist</p>
            <p>This advanced tool assists in the analysis of Pentacam scans and corneal conditions using state-of-the-art AI technology. 
            It provides comprehensive analysis of corneal parameters and supports clinical decision-making in keratoconus, FECD, and other corneal conditions.</p>
            <p><strong>Note:</strong> This tool is for assistance only and should not replace professional medical judgment.</p>
        </div>
    """, unsafe_allow_html=True)
    
    # Main content
    col1, col2 = st.columns([1, 1])
    
    with col1:
        # Patient Information Section
        st.markdown("### Patient Information")
        with st.container():
            st.markdown('<div class="patient-info">', unsafe_allow_html=True)
            patient_name = st.text_input("Patient Name")
            patient_id = st.text_input("Patient ID")
            dob = st.date_input("Date of Birth")
            diagnosis = st.text_input("Previous Diagnosis (if any)")
            
            # Clinical History
            st.markdown("#### Clinical History")
            history_text = st.text_area("Enter relevant clinical history")
            
            # Multiple PDF uploads for patient records
            st.markdown("#### Additional Patient Records (PDF)")
            patient_pdfs = st.file_uploader("Upload Patient Records", type=['pdf'], accept_multiple_files=True)
            if patient_pdfs:
                for pdf in patient_pdfs:
                    with st.expander(f"View {pdf.name}"):
                        st.text(extract_pdf_text(pdf))
            st.markdown('</div>', unsafe_allow_html=True)

        # Pentacam Scan Upload Section
        st.markdown("### Upload Pentacam Scans")
        st.markdown('<div class="upload-section">', unsafe_allow_html=True)
        scan_type = st.radio("Select Analysis Type", ["Single Timepoint", "Compare Timepoints"])
        
        if scan_type == "Single Timepoint":
            scan_date = st.date_input("Scan Date")
            uploaded_files = st.file_uploader("Upload Pentacam Scans", type=['png', 'jpg', 'jpeg', 'pdf'], accept_multiple_files=True)
            
            if uploaded_files:
                images = []
                for file in uploaded_files:
                    if file.type.startswith('image'):
                        image = Image.open(file)
                        images.append(image)
                    elif file.type == 'application/pdf':
                        # Handle PDF files here if needed
                        st.warning(f"PDF processing for {file.name} will be implemented soon")
                
                if images:
                    st.markdown("#### Preview Uploaded Scans")
                    cols = st.columns(len(images))
                    for idx, (col, img) in enumerate(zip(cols, images)):
                        with col:
                            st.image(img, caption=f"Scan {idx + 1}", use_column_width=True)
                    
                    if st.button("Analyze Scans"):
                        with st.spinner("Analyzing..."):
                            analysis = analyze_cornea_images(
                                images,
                                timepoint=scan_date.strftime("%Y-%m-%d"),
                                patient_data=f"Name: {patient_name}\nID: {patient_id}\nDOB: {str(dob)}\nDiagnosis: {diagnosis}\nHistory: {history_text}"
                            )
                            st.markdown("### Analysis Results")
                            st.markdown('<div class="analysis-container">', unsafe_allow_html=True)
                            st.markdown(analysis)
                            st.markdown('</div>', unsafe_allow_html=True)
        
        else:  # Compare Timepoints
            col1, col2 = st.columns(2)
            
            with col1:
                st.markdown("#### First Timepoint")
                date1 = st.date_input("Date of First Scan")
                files1 = st.file_uploader("Upload First Scans", type=['png', 'jpg', 'jpeg', 'pdf'], accept_multiple_files=True)
            
            with col2:
                st.markdown("#### Second Timepoint")
                date2 = st.date_input("Date of Second Scan")
                files2 = st.file_uploader("Upload Second Scans", type=['png', 'jpg', 'jpeg', 'pdf'], accept_multiple_files=True)
            
            if files1 and files2:
                images1 = [Image.open(f) for f in files1 if f.type.startswith('image')]
                images2 = [Image.open(f) for f in files2 if f.type.startswith('image')]
                
                if images1 and images2:
                    st.markdown("#### Preview First Timepoint")
                    cols1 = st.columns(len(images1))
                    for idx, (col, img) in enumerate(zip(cols1, images1)):
                        with col:
                            st.image(img, caption=f"Scan {idx + 1}", use_column_width=True)
                    
                    st.markdown("#### Preview Second Timepoint")
                    cols2 = st.columns(len(images2))
                    for idx, (col, img) in enumerate(zip(cols2, images2)):
                        with col:
                            st.image(img, caption=f"Scan {idx + 1}", use_column_width=True)
                    
                    if st.button("Compare Scans"):
                        with st.spinner("Analyzing..."):
                            comparison = compare_timepoints(
                                images1, date1.strftime("%Y-%m-%d"),
                                images2, date2.strftime("%Y-%m-%d"),
                                patient_data=f"Name: {patient_name}\nID: {patient_id}\nDOB: {str(dob)}\nDiagnosis: {diagnosis}\nHistory: {history_text}"
                            )
                            st.markdown("### Comparison Results")
                            st.markdown('<div class="analysis-container">', unsafe_allow_html=True)
                            st.markdown(comparison)
                            st.markdown('</div>', unsafe_allow_html=True)
        
        st.markdown('</div>', unsafe_allow_html=True)

if __name__ == "__main__":
    main()