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
import pandas as pd
from collections import defaultdict
import re

# 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;
    }
    .timeline-container {
        margin: 2rem 0;
        padding: 1rem;
        background-color: #fff;
        border-radius: 0.5rem;
        border: 1px solid #e9ecef;
    }
    .timepoint-card {
        background-color: #f8f9fa;
        padding: 1rem;
        margin: 0.5rem 0;
        border-radius: 0.5rem;
        border-left: 3px solid #2E86C1;
    }
    </style>
""", unsafe_allow_html=True)

# System prompts
CORNEA_ANALYSIS_PROMPT = """You are an expert ophthalmologist specializing in corneal diseases. Analyze these Pentacam scans and patient data with focus on:

1. Corneal Parameters Analysis:
   • Thickness mapping and progression
   • Topographic changes
   • Elevation data (anterior and posterior)
   • Keratoconus indices and classification

2. Disease Assessment:
   • ABCD Keratoconus staging
   • Fuchs Endothelial Corneal Dystrophy evaluation
   • Subclinical corneal edema (Sun criteria)
   • Risk assessment

3. Clinical Interpretation:
   • Pattern recognition
   • Disease progression markers
   • Treatment implications

Please provide a detailed clinical assessment."""

PROGRESSION_ANALYSIS_PROMPT = """Analyze the progression of corneal parameters across multiple timepoints, focusing on:

1. Temporal Changes:
   • Progressive changes in corneal thickness
   • Evolution of topographic patterns
   • Changes in elevation maps
   • Progression of keratoconus indices

2. Rate of Progression:
   • Quantify changes between timepoints
   • Identify acceleration or stabilization periods
   • Compare with expected disease progression

3. Risk Assessment:
   • Current status evaluation
   • Future progression risk
   • Treatment recommendations

4. Timeline Analysis:
   • Key changes between each timepoint
   • Overall progression pattern
   • Critical periods of change

Please provide a comprehensive progression analysis with clinical recommendations."""

def extract_patient_data(uploaded_file):
    """Extract and process patient data from uploaded file"""
    patient_data = {}
    
    if uploaded_file.type == "application/pdf":
        pdf_reader = PyPDF2.PdfReader(uploaded_file)
        text = ""
        for page in pdf_reader.pages:
            text += page.extract_text()
        patient_data['raw_text'] = text
    else:
        # Handle other file types if needed
        patient_data['raw_text'] = "File type not supported for detailed extraction"
    
    return patient_data

def analyze_timepoint(images, date, patient_data=None):
    """Analyze a single timepoint"""
    prompt = f"{CORNEA_ANALYSIS_PROMPT}\n\nTimepoint: {date}\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 analyze_progression(timepoints_data):
    """Analyze progression across multiple timepoints"""
    prompt = f"{PROGRESSION_ANALYSIS_PROMPT}\n\n"
    prompt += "Timepoints for analysis:\n"
    
    # Add all timepoints to the prompt
    all_images = []
    for date, images in timepoints_data.items():
        prompt += f"\n- {date}:"
        all_images.extend(images)
    
    prompt += "\n\nPlease analyze the progression across these timepoints:"
    content = [prompt] + all_images
    response = model.generate_content(content)
    return response.text

def extract_date_from_filename(filename):
    """Extract date from filename using common patterns"""
    # Common date patterns (add more patterns if needed)
    patterns = [
        r'(\d{4}[-_/]\d{2}[-_/]\d{2})',  # YYYY-MM-DD, YYYY_MM_DD
        r'(\d{2}[-_/]\d{2}[-_/]\d{4})',  # DD-MM-YYYY, DD_MM_YYYY
        r'(\d{8})',  # YYYYMMDD
    ]
    
    for pattern in patterns:
        match = re.search(pattern, filename)
        if match:
            date_str = match.group(1)
            try:
                # Try different date formats
                for fmt in ['%Y-%m-%d', '%Y_%m_%d', '%d-%m-%Y', '%d_%m_%Y', '%Y%m%d']:
                    try:
                        return datetime.strptime(date_str.replace('/', '-'), fmt).strftime('%Y-%m-%d')
                    except ValueError:
                        continue
            except ValueError:
                continue
    return None

def organize_scans_by_date(files):
    """Organize uploaded files by their dates"""
    organized_files = defaultdict(list)
    unorganized_files = []
    
    for file in files:
        date = extract_date_from_filename(file.name)
        if date:
            organized_files[date].append(file)
        else:
            unorganized_files.append(file)
    
    return organized_files, unorganized_files

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)
    
    # Patient Information Section
    st.markdown("### Patient Information")
    st.markdown("""
        <div class="info-box">
            Upload patient information including:
            • Clinical history
            • Previous diagnoses
            • Current symptoms
            • Family history
            • Previous treatments
            • Current medications
            • Other relevant medical conditions
        </div>
    """, unsafe_allow_html=True)
    
    patient_file = st.file_uploader("Upload Patient Information (PDF/Text)", type=['pdf', 'txt'])
    patient_data = None
    if patient_file:
        patient_data = extract_patient_data(patient_file)
        with st.expander("View Extracted Patient Information"):
            st.text(patient_data.get('raw_text', 'No text extracted'))

    # Scan Analysis Section
    st.markdown("### Pentacam Scan Analysis")
    analysis_type = st.radio("Select Analysis Type", ["Single Timepoint", "Progression Analysis"])

    if analysis_type == "Single Timepoint":
        st.markdown('<div class="upload-section">', unsafe_allow_html=True)
        uploaded_files = st.file_uploader("Upload Pentacam Scans", type=['png', 'jpg', 'jpeg'], accept_multiple_files=True)
        
        if uploaded_files:
            images = []
            for file in uploaded_files:
                image = Image.open(file)
                images.append(image)
            
            if images:
                st.markdown("#### Preview 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_timepoint(
                            images,
                            datetime.now().strftime("%Y-%m-%d"),  # Current date for reference
                            patient_data.get('raw_text') if patient_data else None
                        )
                        st.markdown("### Analysis Results")
                        st.markdown('<div class="analysis-container">', unsafe_allow_html=True)
                        st.markdown(analysis)
                        st.markdown('</div>', unsafe_allow_html=True)
        
        st.markdown('</div>', unsafe_allow_html=True)
    
    else:  # Progression Analysis
        st.markdown('<div class="upload-section">', unsafe_allow_html=True)
        st.info("""Upload all your Pentacam scans at once. The system will automatically organize them by date and analyze progression.
                   For best results, ensure your scan filenames include dates (e.g., 'scan_2023-01-15.jpg' or 'pentacam_20230115.png')""")
        
        uploaded_files = st.file_uploader(
            "Upload All Pentacam Scans",
            type=['png', 'jpg', 'jpeg'],
            accept_multiple_files=True
        )
        
        if uploaded_files:
            organized_files, unorganized_files = organize_scans_by_date(uploaded_files)
            
            if organized_files:
                st.markdown("### Organized Scans by Date")
                st.markdown('<div class="timeline-container">', unsafe_allow_html=True)
                
                timepoints_data = defaultdict(list)
                dates = sorted(organized_files.keys())
                
                for date in dates:
                    st.markdown(f'<div class="timepoint-card">', unsafe_allow_html=True)
                    st.markdown(f"#### Timepoint: {date}")
                    
                    files = organized_files[date]
                    images = []
                    cols = st.columns(len(files))
                    
                    for idx, (file, col) in enumerate(zip(files, cols)):
                        with col:
                            image = Image.open(file)
                            images.append(image)
                            st.image(image, caption=f"Scan {idx + 1}", use_column_width=True)
                    
                    timepoints_data[date].extend(images)
                    st.markdown('</div>', unsafe_allow_html=True)
                
                if unorganized_files:
                    st.warning(f"{len(unorganized_files)} files couldn't be automatically dated. Please ensure filenames include dates.")
                    with st.expander("Manually Assign Dates"):
                        for file in unorganized_files:
                            col1, col2 = st.columns([2, 1])
                            with col1:
                                st.text(file.name)
                            with col2:
                                date = st.date_input(f"Date for {file.name}", key=f"manual_{file.name}")
                                image = Image.open(file)
                                timepoints_data[date.strftime("%Y-%m-%d")].append(image)
                
                if len(timepoints_data) >= 2:
                    if st.button("Analyze Progression"):
                        with st.spinner("Analyzing progression across timepoints..."):
                            progression_analysis = analyze_progression(timepoints_data)
                            st.markdown("### Progression Analysis Results")
                            st.markdown('<div class="analysis-container">', unsafe_allow_html=True)
                            st.markdown(progression_analysis)
                            st.markdown('</div>', unsafe_allow_html=True)
                else:
                    st.warning("Please upload scans from at least 2 different timepoints for progression analysis.")
                
                st.markdown('</div>', unsafe_allow_html=True)
            else:
                st.error("No dated scans found. Please ensure your filenames include dates (e.g., 'scan_2023-01-15.jpg').")
        
        st.markdown('</div>', unsafe_allow_html=True)

if __name__ == "__main__":
    main()