import streamlit as st
import cv2
import numpy as np
import mediapipe as mp
from sklearn.linear_model import LinearRegression
import random
import base64
import joblib
from datetime import datetime
import io
from reportlab.lib.pagesizes import A4
from reportlab.platypus import SimpleDocTemplate, Table, TableStyle, Paragraph, Spacer, PageBreak, PageTemplate, Frame, Image
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from reportlab.lib import colors
from reportlab.lib.enums import TA_CENTER, TA_LEFT, TA_RIGHT
from reportlab.graphics.shapes import Drawing, Line
from PIL import Image as PILImage
import tempfile
import os
import logging
import re
import pandas as pd
import torch
from torchvision import models, transforms
import PyPDF2

# Set up logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

# Configure Streamlit page
st.set_page_config(
    page_title="AI Health Report Generator",
    page_icon="🧠",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for enhanced UI with DataFrame styling
st.markdown("""
<style>
    body {
        font-family: 'Helvetica', Arial, sans-serif;
    }
    .main-header {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        padding: 1.5rem;
        border-radius: 10px;
        color: white;
        text-align: center;
        margin-bottom: 1.5rem;
        box-shadow: 0 8px 20px rgba(0,0,0,0.15);
    }
    .main-header h1 {
        font-size: 2.5rem;
        font-weight: 500;
        margin-bottom: 0.5rem;
    }
    .main-header p {
        font-size: 1rem;
        font-weight: 400;
    }
    .patient-form {
        background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
        padding: 1.5rem;
        border-radius: 10px;
        color: white;
        margin-bottom: 1.5rem;
    }
    .patient-form input, .patient-form select {
        border-radius: 8px;
        margin-bottom: 0.5rem;
        font-size: 0.9rem;
    }
    .upload-area {
        background: linear-gradient(135deg, #4facfe 0%, #00f2fe 100%);
        padding: 1.5rem;
        border-radius: 10px;
        border: 1.5px dashed white;
        text-align: center;
        color: white;
        margin-bottom: 1.5rem;
    }
    .upload-area h3 {
        font-size: 1.5rem;
        font-weight: 500;
        margin-bottom: 0.5rem;
    }
    .upload-area p {
        font-size: 0.9rem;
    }
    .health-card {
        background: linear-gradient(135deg, #E6E6FA 0%, #F0F8FF 100%);
        border: 1px solid #CCCCCC;
        border-radius: 12px;
        padding: 1.5rem;
        margin: 0.5rem 0;
        box-shadow: 0 6px 20px rgba(0, 0, 0, 0.1);
    }
    .stButton > button {
        background: linear-gradient(135deg, #2E8B57, #228B22);
        color: white;
        border: 1px solid #228B22;
        border-radius: 8px;
        padding: 0.5rem 1.5rem;
        font-weight: 500;
        font-size: 14px;
        box-shadow: 0 3px 10px rgba(46, 139, 87, 0.2);
        transition: all 0.3s;
        width: 100%;
    }
    .stButton > button:hover {
        background: linear-gradient(135deg, #228B22, #1B6B1B);
        transform: translateY(-1px);
        box-shadow: 0 4px 12px rgba(46, 139, 87, 0.3);
    }
    .metric-card {
        background: white;
        padding: 0.8rem;
        border-radius: 8px;
        border-left: 3px solid #2E8B57;
        margin: 0.5rem 0;
        box-shadow: 0 2px 6px rgba(0,0,0,0.1);
        transition: transform 0.2s;
    }
    .metric-card:hover {
        transform: scale(1.02);
    }
    .metric-card p {
        font-size: 0.9rem;
        margin: 0.3rem 0;
    }
    .category-container {
        border: 1px solid transparent;
        border-image: linear-gradient(135deg, #2E8B57, #228B22) 1;
        border-radius: 10px;
        padding: 0.5rem;
        margin-bottom: 1rem;
        background-color: #E6E6FA;
        box-shadow: 0 2px 8px rgba(0,0,0,0.1);
        border-top: 2px solid #2E8B57;
    }
    .table-container {
        width: 100%;
        overflow-x: auto;
        display: block;
    }
    .stDataFrame table {
        width: 100 !important;
        border-collapse: collapse !important;
        font-size: 12px !important;
        background-color: #FFFFFF !important;
    }
    .stDataFrame thead tr th {
        background-color: #2E8B57 !important;
        color: white !important;
        font-weight: bold !important;
        font-size: 14px !important;
        padding: 0.5rem !important;
        border: 1px solid #228B22 !important;
        text-align: center !important;
    }
    .stDataFrame thead tr th:first-child {
        text-align: left !important;
        padding-left: 0.7rem !important;
    }
    .stDataFrame tbody tr td {
        padding: 0.5rem !important;
        border: 1px solid #CCCCCC !important;
        text-align: center !important;
    }
    .stDataFrame tbody tr td:first-child {
        text-align: left !important;
        padding-left: 0.7rem !important;
    }
    .stDataFrame tbody tr:nth-child(even) {
        background-color: #F8F8FF !important;
    }
    .stDataFrame tbody tr:nth-child(odd) {
        background-color: #F0F8FF !important;
    }
    .stDataFrame tbody tr:hover {
        background-color: #E0E0FF !important;
    }
    .status-normal {
        color: #2E8B57 !important;
    }
    .status-low {
        color: #ffca28 !important;
    }
    .status-high {
        color: #d32f2f !important;
    }
    .summary-card {
        background-color: #E6E6FA;
        border-left: 3px solid #2E8B57;
        border-radius: 8px;
        padding: 1rem;
        margin: 1rem 0;
        box-shadow: 0 2px 6px rgba(0,0,0,0.1);
    }
    .summary-card p {
        font-size: 0.95rem;
        margin: 0.3rem 0;
        color: #333;
    }
    .summary-card b {
        color: #2E8B57;
    }
    .help-report {
        background: linear-gradient(135deg, #E6E6FA 0%, #F0F8FF 100%);
        border: 1px solid #CCCCCC;
        border-radius: 12px;
        padding: 1.5rem;
        margin: 1rem 0;
        box-shadow: 0 6px 20px rgba(0, 0, 0, 0.1);
    }
    .help-report-content {
        display: flex;
        align-items: center;
    }
    .help-report img {
        max-width: 150px;
        max-height: 150px;
        border-radius: 8px;
        margin-right: 1rem;
    }
    .help-report-details {
        display: inline-block;
    }
    .help-report-details p {
        font-size: 0.9rem;
        margin: 0.3rem 0;
        color: #333;
    }
    .help-report-details b {
        color: #2E8B57;
    }
    .xray-analysis {
        margin-top: 1rem;
        padding: 1rem;
        background-color: #ffffff;
        border-radius: 10px;
        box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
    }
</style>
""", unsafe_allow_html=True)

# Initialize MediaPipe
@st.cache_resource
def load_face_mesh():
    try:
        mp_face_mesh = mp.solutions.face_mesh
        face_mesh = mp_face_mesh.FaceMesh(
            static_image_mode=True,
            max_num_faces=1,
            refine_landmarks=True,
            min_detection_confidence=0.5
        )
        logger.info("MediaPipe FaceMesh initialized successfully")
        return face_mesh
    except Exception as e:
        st.error(f"Failed to initialize MediaPipe: {str(e)}")
        logger.error(f"MediaPipe initialization failed: {str(e)}")
        return None

# Initialize MediaPipe drawing utilities
@st.cache_resource
def load_drawing_utils():
    try:
        mp_drawing = mp.solutions.drawing_utils
        mp_drawing_styles = mp.solutions.drawing_styles
        logger.info("MediaPipe drawing utilities initialized successfully")
        return mp_drawing, mp_drawing_styles
    except Exception as e:
        logger.error(f"MediaPipe drawing utilities initialization failed: {str(e)}")
        return None, None

# Load models
@st.cache_resource
def load_models():
    try:
        hemoglobin_model = joblib.load("hemoglobin_model_from_anemia_dataset.pkl")
    except FileNotFoundError:
        st.warning("Hemoglobin model not found. Training a temporary model.")
        hemoglobin_model = train_model((13.5, 17.5))
    try:
        spo2_model = joblib.load("spo2_model_simulated.pkl")
    except FileNotFoundError:
        st.warning("SpO2 model not found. Training a temporary model.")
        spo2_model = train_model((95, 100))
    try:
        hr_model = joblib.load("heart_rate_model.pkl")
    except FileNotFoundError:
        st.warning("Heart rate model not found. Training a temporary model.")
        hr_model = train_model((60, 100))
    logger.info("Models loaded or trained successfully")
    return hemoglobin_model, spo2_model, hr_model

@st.cache_resource
def load_xray_model():
    try:
        model = models.densenet121(pretrained=True)
        model.eval()
        logger.info("X-ray model (DenseNet121) loaded successfully")
        return model
    except Exception as e:
        st.error(f"Failed to load X-ray model: {str(e)}")
        logger.error(f"X-ray model loading failed: {str(e)}")
        return None

def train_model(output_range):
    try:
        X = [[random.uniform(0.2, 0.5), random.uniform(0.05, 0.2), random.uniform(0.05, 0.2),
              random.uniform(0.2, 0.5), random.uniform(0.2, 0.5), random.uniform(0.2, 0.5),
              random.uniform(0.2, 0.5)] for _ in range(100)]
        y = [random.uniform(*output_range) for _ in X]
        model = LinearRegression().fit(X, y)
        logger.info(f"Model trained for range {output_range}")
        return model
    except Exception as e:
        st.error(f"Failed to train model: {str(e)}")
        logger.error(f"Model training failed: {str(e)}")
        return None

def extract_features(image, landmarks):
    try:
        if len(image.shape) < 3 or image.shape[2] != 3:
            st.error("Invalid image format: Expected RGB image.")
            logger.error("Invalid image format: Not RGB")
            return None
        red_channel = image[:, :, 2]
        green_channel = image[:, :, 1]
        blue_channel = image[:, :, 0]
        red_percent = 100 * np.mean(red_channel) / 255
        green_percent = 100 * np.mean(green_channel) / 255
        blue_percent = 100 * np.mean(blue_channel) / 255
        logger.info("Features extracted successfully")
        return [red_percent, green_percent, blue_percent]
    except Exception as e:
        st.error(f"Failed to extract features: {str(e)}")
        logger.error(f"Feature extraction failed: {str(e)}")
        return None

def get_risk_level(value, normal_range):
    try:
        low, high = normal_range
        if value < low:
            return "Low", "#ffca28"
        elif value > high:
            return "High", "#d32f2f"
        else:
            return "Normal", "#2E8B57"
    except Exception as e:
        st.error(f"Failed to determine risk level: {str(e)}")
        logger.error(f"Risk level determination failed: {str(e)}")
        return "Unknown", "#ffffff"

def draw_analyzed_image(image, landmarks):
    try:
        mp_drawing, mp_drawing_styles = load_drawing_utils()
        if mp_drawing is None or mp_drawing_styles is None:
            logger.error("Drawing utilities not initialized")
            return image
        annotated_image = image.copy()
        h, w = annotated_image.shape[:2]
        # Draw facial landmarks
        mp_drawing.draw_landmarks(
            image=annotated_image,
            landmark_list=landmarks,
            connections=mp.solutions.face_mesh.FACEMESH_TESSELATION,
            landmark_drawing_spec=None,
            connection_drawing_spec=mp_drawing_styles.get_default_face_mesh_tesselation_style()
        )
        # Highlight analysis regions (cheeks, forehead, nose)
        cheek_left_points = [landmarks.landmark[i] for i in range(50, 151)]
        cheek_right_points = [landmarks.landmark[i] for i in range(280, 381)]
        forehead_points = [landmarks.landmark[i] for i in range(10, 51)]
        nose_points = [landmarks.landmark[i] for i in range(1, 5)]
        def normalize_to_pixel(landmark):
            return (int(landmark.x * w), int(landmark.y * h))
        def get_bounding_box(points):
            x_coords = [p.x * w for p in points]
            y_coords = [p.y * h for p in points]
            x_min, x_max = int(min(x_coords)), int(max(x_coords))
            y_min, y_max = int(min(y_coords)), int(max(y_coords))
            return x_min, y_min, x_max, y_max
        # Draw semi-transparent colored rectangles
        overlay = annotated_image.copy()
        # Left cheek (red)
        x_min, y_min, x_max, y_max = get_bounding_box(cheek_left_points)
        cv2.rectangle(overlay, (x_min, y_min), (x_max, y_max), (255, 0, 0), -1)
        logger.debug(f"Left cheek box: ({x_min}, {y_min}, {x_max}, {y_max})")
        # Right cheek (red)
        x_min, y_min, x_max, y_max = get_bounding_box(cheek_right_points)
        cv2.rectangle(overlay, (x_min, y_min), (x_max, y_max), (255, 0, 0), -1)
        logger.debug(f"Right cheek box: ({x_min}, {y_min}, {x_max}, {y_max})")
        # Forehead (green)
        x_min, y_min, x_max, y_max = get_bounding_box(forehead_points)
        cv2.rectangle(overlay, (x_min, y_min), (x_max, y_max), (0, 255, 0), -1)
        logger.debug(f"Forehead box: ({x_min}, {y_min}, {x_max}, {y_max})")
        # Nose (blue)
        x_min, y_min, x_max, y_max = get_bounding_box(nose_points)
        cv2.rectangle(overlay, (x_min, y_min), (x_max, y_max), (0, 0, 255), -1)
        logger.debug(f"Nose box: ({x_min}, {y_min}, {x_max}, {y_max})")
        # Apply transparency
        alpha = 0.4  # 40% opacity
        cv2.addWeighted(overlay, alpha, annotated_image, 1 - alpha, 0, annotated_image)
        logger.info("Analyzed image generated with landmarks and region highlights")
        return annotated_image
    except Exception as e:
        logger.error(f"Failed to draw analyzed image: {str(e)}")
        st.error(f"Analyzed image generation failed: {str(e)}")
        return image

def create_pdf_report(patient_data, test_results, profile_image):
    try:
        logger.info("Starting PDF generation")
        buffer = io.BytesIO()
        doc = SimpleDocTemplate(buffer, pagesize=A4, rightMargin=30, leftMargin=30, topMargin=50, bottomMargin=80)
        elements = []
        styles = getSampleStyleSheet()
        # Define custom styles
        header_logo_style = ParagraphStyle(
            'HeaderLogoStyle',
            parent=styles['Heading1'],
            fontName='Helvetica-Bold',
            fontSize=16,
            spaceAfter=8,
            alignment=TA_LEFT,
            textColor=colors.HexColor('#2E8B57'),
            leftIndent=0,
            spaceBefore=0
        )
        test_report_badge_style = ParagraphStyle(
            'TestReportBadgeStyle',
            parent=styles['Normal'],
            fontName='Helvetica-Bold',
            fontSize=12,
            alignment=TA_RIGHT,
            textColor=colors.HexColor('#2E8B57'),
            borderWidth=1,
            borderColor=colors.HexColor('#2E8B57'),
            borderPadding=8,
            spaceBefore=0,
            spaceAfter=0
        )
        patient_info_style = ParagraphStyle(
            'PatientInfoStyle',
            parent=styles['Normal'],
            fontName='Helvetica',
            fontSize=9,
            spaceAfter=6,
            textColor=colors.black,
            alignment=TA_LEFT
        )
        section_header_style = ParagraphStyle(
            'SectionHeaderStyle',
            parent=styles['Heading2'],
            fontName='Helvetica-Bold',
            fontSize=12,
            spaceAfter=8,
            spaceBefore=12,
            textColor=colors.black,
            alignment=TA_CENTER,
            backColor=colors.HexColor('#f0f0f0'),
            borderWidth=1,
            borderColor=colors.black,
            borderPadding=6
        )
        footer_style = ParagraphStyle(
            'FooterStyle',
            parent=styles['Normal'],
            fontName='Helvetica',
            fontSize=8,
            textColor=colors.black,
            alignment=TA_CENTER,
            spaceAfter=4
        )
        signatory_style = ParagraphStyle(
            'SignatoryStyle',
            parent=styles['Normal'],
            fontName='Helvetica-Bold',
            fontSize=9,
            spaceAfter=6,
            textColor=colors.black,
            alignment=TA_RIGHT
        )
        # Footer content
        footer_text = """
        Sathkrutha Tech Solutions Pvt. Ltd Registered Office: H.No: 2-3-685/5/1, Flat N Venkateshwara Nagar, Amberpet, Hyderabad, Telangana 500013, INDIA<br/>
        T: +91 4027264141 F: +91 4027263667 E: helpdesk@sathkrutha.com
        """
        def add_page_footer(canvas, doc):
            canvas.saveState()
            canvas.setLineWidth(2)
            canvas.setStrokeColor(colors.black)
            canvas.rect(20, 20, A4[0]-40, A4[1]-40)
            canvas.setFont('Helvetica', 8)
            page_num = canvas.getPageNumber()
            if hasattr(doc, '_total_pages'):
                page_text = f"Page {page_num} of {doc._total_pages}"
            else:
                page_text = f"Page {page_num}"
            canvas.drawRightString(A4[0]-40, 25, page_text)
            footer_para = Paragraph(footer_text, footer_style)
            w, h = footer_para.wrap(A4[0]-60, 40)
            footer_para.drawOn(canvas, 30, 30)
            canvas.restoreState()
        doc.addPageTemplates([
            PageTemplate(id='AllPages', frames=[Frame(30, 80, A4[0]-60, A4[1]-130)], onPage=add_page_footer)
        ])
        # Header
        header_table_data = [
            [Paragraph("<b>Sathkrutha</b><br/><font color='#FF8C00'>Clinical Diagnostics</font>", header_logo_style), 
             Paragraph("Test Report", test_report_badge_style)]
        ]
        header_table = Table(header_table_data, colWidths=[4*inch, 2*inch])
        header_table.setStyle(TableStyle([
            ('VALIGN', (0, 0), (-1, -1), 'TOP'),
            ('ALIGN', (0, 0), (0, 0), 'LEFT'),
            ('ALIGN', (1, 0), (1, 0), 'RIGHT'),
            ('LEFTPADDING', (0, 0), (-1, -1), 0),
            ('RIGHTPADDING', (0, 0), (-1, -1), 0),
            ('TOPPADING', (0, 0), (-1, -1), 0),
            ('BOTTOMPADDING', (0, 0), (-1, -1), 12),
        ]))
        elements.append(header_table)
        elements.append(Spacer(1, 10))
        # Issued to section
        issued_to_data = [["Issued to :", ""]]
        issued_table = Table(issued_to_data, colWidths=[1*inch, 5*inch])
        issued_table.setStyle(TableStyle([
            ('FONT', (0, 0), (-1, -1), 'Helvetica-Bold', 9),
            ('GRID', (0, 0), (-1, -1), 1, colors.black),
            ('BACKGROUND', (0, 0), (-1, -1), colors.HexColor('#e6e6e6')),
            ('LEFTPADDING', (0, 0), (-1, -1), 4),
            ('TOPPADING', (0, 0), (-1, -1), 4),
            ('BOTTOMPADDING', (0, 0), (-1, -1), 4),
        ]))
        elements.append(issued_table)
        # Patient details with photo
        img_buffer = io.BytesIO()
        profile_image.save(img_buffer, format="PNG")
        img_buffer.seek(0)
        img = Image(img_buffer, width=1.5*inch, height=1.5*inch)
        patient_details = [
            [img, Paragraph(f"Name: {patient_data.get('name', 'Unknown Patient')}<br/>"
                            f"Age: {patient_data.get('age', 'N/A')} Years<br/>"
                            f"Gender: {patient_data.get('gender', 'Male')}<br/>"
                            f"ID: {patient_data.get('id', 'N/A')}<br/>"
                            f"Date: {datetime.now().strftime('%d-%b-%Y %H:%M')}", patient_info_style)]
        ]
        patient_table = Table(patient_details, colWidths=[1.5*inch, 5.5*inch])
        patient_table.setStyle(TableStyle([
            ('VALIGN', (0, 0), (-1, -1), 'MIDDLE'),
            ('LEFTPADDING', (0, 0), (-1, -1), 10),
            ('RIGHTPADDING', (0, 0), (-1, -1), 10),
            ('TOPPADING', (0, 0), (-1, -1), 10),
            ('BOTTOMPADING', (0, 0), (-1, -1), 10),
        ]))
        elements.append(patient_table)
        elements.append(Spacer(1, 15))
        # Test categories
        for category_index, (category, tests) in enumerate(test_results.items()):
            if category_index > 0:
                elements.append(PageBreak())
            clean_category = category.replace("■", "").strip().upper()
            elements.append(Paragraph(clean_category, section_header_style))
            elements.append(Spacer(1, 10))
            table_data = [["Test Description", "Value Observed", "Unit", "Biological Reference Interval"]]
            for test_name, result, range_val, level_info in tests:
                level, _ = level_info
                status_indicator = " L" if level == "Low" else " H" if level == "High" else ""
                if "Count" in test_name or test_name == "Respiratory Rate":
                    value_str = f"{result:.0f}{status_indicator}"
                elif test_name in ["Temperature", "SpO2"]:
                    value_str = f"{result:.1f}{status_indicator}"
                else:
                    value_str = f"{result:.1f}{status_indicator}"
                unit = "" if "BP" in test_name else ("g/dL" if "Hemoglobin" in test_name else
                                                   "cu/mm" if "WBC Count" in test_name else
                                                   "Thousand/µL" if "Platelet Count" in test_name else
                                                   "µg/dL" if "Iron" in test_name or "TIBC" in test_name else
                                                   "ng/mL" if "Ferritin" in test_name else
                                                   "mg/dL" if "Bilirubin" in test_name or "Creatinine" in test_name or "Urea" in test_name else
                                                   "mEq/L" if "Sodium" in test_name or "Potassium" in test_name else
                                                   "%" if "SpO2" in test_name else
                                                   "bpm" if "Heart Rate" in test_name else
                                                   "/min" if "Respiratory Rate" in test_name else
                                                   "°F" if "Temperature" in test_name else "mmHg")
                range_str = f"{range_val[0]:.0f} - {range_val[1]:.0f}" if "Count" in test_name or test_name == "Respiratory Rate" else f"{range_val[0]:.1f} - {range_val[1]:.1f}"
                table_data.append([test_name, value_str, unit, range_str])
            test_table = Table(table_data, colWidths=[2.5*inch, 1.2*inch, 0.8*inch, 1.5*inch])
            test_table.setStyle(TableStyle([
                ('FONT', (0, 0), (-1, 0), 'Helvetica-Bold', 9),
                ('FONT', (0, 1), (-1, -1), 'Helvetica', 9),
                ('BACKGROUND', (0, 0), (-1, 0), colors.white),
                ('GRID', (0, 0), (-1, -1), 0.5, colors.black),
                ('BOX', (0, 0), (-1, -1), 1, colors.black),
                ('LEFTPADDING', (0, 0), (-1, -1), 6),
                ('RIGHTPADDING', (0, 0), (-1, -1), 6),
                ('TOPPADING', (0, 0), (-1, -1), 6),
                ('BOTTOMPADING', (0, 0), (-1, -1), 6),
                ('VALIGN', (0, 0), (-1, -1), 'MIDDLE'),
                ('ALIGN', (0, 0), (0, -1), 'LEFT'),
                ('ALIGN', (1, 1), (-1, -1), 'CENTER'),
            ]))
            elements.append(test_table)
            elements.append(Spacer(1, 30))
            signatory_table_data = [["", ""], ["", "DR.SATHAIAH BEGARI"], ["", "MBBS,DCP, Clinical Pathologist"], ["", "AUTHORISED SIGNATORY"]]
            signatory_table = Table(signatory_table_data, colWidths=[4*inch, 2*inch])
            signatory_table.setStyle(TableStyle([
                ('FONT', (1, 1), (1, -1), 'Helvetica-Bold', 9),
                ('ALIGN', (1, 1), (1, -1), 'RIGHT'),
                ('VALIGN', (1, 1), (1, -1), 'TOP'),
                ('LEFTPADDING', (0, 0), (-1, -1), 0),
                ('RIGHTPADDING', (0, 0), (-1, -1), 0),
                ('TOPPADING', (0, 0), (-1, -1), 2),
                ('BOTTOMPADING', (0, 0), (-1, -1), 2),
            ]))
            elements.append(signatory_table)
        try:
            doc.build(elements)
        except MemoryError as me:
            st.error(f"PDF generation failed due to memory issue: {str(me)}")
            logger.error(f"Memory error during PDF build: {str(me)}")
            return None
        except Exception as e:
            st.error(f"PDF generation failed: {str(e)}")
            logger.error(f"PDF building failed: {str(e)}")
            return None
        buffer.seek(0)
        logger.info("PDF buffer ready")
        return buffer
    except Exception as e:
        st.error(f"Unexpected error in PDF generation: {str(e)}")
        logger.error(f"Unexpected PDF generation error: {str(e)}")
        return None

def process_input(input_data):
    if input_data is None:
        return None, None
    if input_data.name.endswith(('.jpg', '.jpeg', '.png')):
        try:
            image = PILImage.open(input_data)
            logger.info("Image processed successfully")
            return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR), image
        except Exception as e:
            st.error(f"Failed to process image: {str(e)}")
            logger.error(f"Image processing failed: {str(e)}")
            return None, None
    elif input_data.name.endswith(('.mp4', '.avi', '.mov')):
        tmp_path = None
        try:
            with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp:
                tmp.write(input_data.read())
                tmp_path = tmp.name
            cap = cv2.VideoCapture(tmp_path)
            ret, frame = cap.read()
            cap.release()
            if ret:
                image = PILImage.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                logger.info("Video frame extracted successfully")
                return frame, image
            else:
                st.error("Failed to extract frame from video.")
                logger.error("Failed to extract video frame")
                return None, None
        except Exception as e:
            st.error(f"Failed to process video: {str(e)}")
            logger.error(f"Video processing failed: {str(e)}")
            return None, None
        finally:
            if tmp_path and os.path.exists(tmp_path):
                try:
                    os.unlink(tmp_path)
                    logger.info(f"Temporary file {tmp_path} cleaned up")
                except Exception as e:
                    logger.warning(f"Failed to clean up temporary file {tmp_path}: {str(e)}")
    return None, None

def analyze_face(image, patient_data):
    face_mesh = load_face_mesh()
    if face_mesh is None:
        return None, None, None, "Failed to initialize face mesh."
    hemoglobin_model, spo2_model, hr_model = load_models()
    try:
        frame = cv2.resize(image, (640, 480))
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        result = face_mesh.process(frame_rgb)
        logger.info("Image processed for face detection")
    except Exception as e:
        st.error(f"Image processing failed: {str(e)}")
        logger.error(f"Image processing failed: {str(e)}")
        return None, None, None, "Image processing error."
    if not result.multi_face_landmarks:
        return None, None, None, "Face not detected. Please try another image or video."
    landmarks = result.multi_face_landmarks[0]
    features = extract_features(frame_rgb, landmarks.landmark)
    if features is None:
        return None, None, None, "Failed to extract image features."
    analyzed_image = draw_analyzed_image(frame_rgb, landmarks)
    models = {
        "Hemoglobin": hemoglobin_model,
        "WBC Count": train_model((4.0, 11.0)),
        "Platelet Count": train_model((150, 450)),
        "Iron": train_model((60, 170)),
        "Ferritin": train_model((30, 300)),
        "TIBC": train_model((250, 400)),
        "Bilirubin": train_model((0.3, 1.2)),
        "Creatinine": train_model((0.6, 1.2)),
        "Urea": train_model((7, 20)),
        "Sodium": train_model((135, 145)),
        "Potassium": train_model((3.5, 5.1)),
        "Temperature": train_model((97, 99)),
        "BP Systolic": train_model((90, 120)),
        "BP Diastolic": train_model((60, 80))
    }
    test_values = {}
    for label in models:
        if models[label] is None:
            st.error(f"Model for {label} is not initialized.")
            logger.error(f"Model not initialized for {label}")
            return None, None, None, f"Model error for {label}."
        try:
            if label == "Hemoglobin":
                prediction = models[label].predict([features])[0]
                test_values[label] = prediction
            else:
                value = models[label].predict([[random.uniform(0.2, 0.5) for _ in range(7)]])[0]
                test_values[label] = value
        except Exception as e:
            st.error(f"Prediction failed for {label}: {str(e)}")
            logger.error(f"Prediction failed for {label}: {str(e)}")
            return None, None, None, f"Prediction error for {label}."
    try:
        gray = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2GRAY)
        green_std = np.std(frame_rgb[:, :, 1]) / 255
        brightness_std = np.std(gray) / 255
        tone_index = np.mean(frame_rgb[100:150, 100:150]) / 255 if frame_rgb[100:150, 100:150].size else 0.5
        hr_features = [brightness_std, green_std, tone_index]
        heart_rate = float(np.clip(hr_model.predict([hr_features])[0], 60, 100))
        skin_patch = frame_rgb[100:150, 100:150]
        skin_tone_index = np.mean(skin_patch) / 255 if skin_patch.size else 0.5
        brightness_variation = np.std(cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2GRAY)) / 255
        spo2_features = [heart_rate, brightness_variation, skin_tone_index]
        spo2 = spo2_model.predict([spo2_features])[0]
        rr = int(12 + abs(heart_rate % 5 - 2))
        logger.info(f"Vitals calculated: heart_rate={heart_rate}, spo2={spo2}, rr={rr}")
    except Exception as e:
        st.error(f"Vitals calculation failed: {str(e)}")
        logger.error(f"Vitals calculation failed: {str(e)}")
        return None, None, None, "Vitals calculation error."
    test_results = {
        "Hematology": [
            ("Hemoglobin", test_values["Hemoglobin"], (13.5, 17.5), get_risk_level(test_values["Hemoglobin"], (13.5, 17.5))),
            ("WBC Count", test_values["WBC Count"], (4.0, 11.0), get_risk_level(test_values["WBC Count"], (4.0, 11.0))),
            ("Platelet Count", test_values["Platelet Count"], (150, 450), get_risk_level(test_values["Platelet Count"], (150, 450)))
        ],
        "Iron Panel": [
            ("Iron", test_values["Iron"], (60, 170), get_risk_level(test_values["Iron"], (60, 170))),
            ("Ferritin", test_values["Ferritin"], (30, 300), get_risk_level(test_values["Ferritin"], (30, 300))),
            ("TIBC", test_values["TIBC"], (250, 400), get_risk_level(test_values["TIBC"], (250, 400)))
        ],
        "Liver & Kidney": [
            ("Bilirubin", test_values["Bilirubin"], (0.3, 1.2), get_risk_level(test_values["Bilirubin"], (0.3, 1.2))),
            ("Creatinine", test_values["Creatinine"], (0.6, 1.2), get_risk_level(test_values["Creatinine"], (0.6, 1.2))),
            ("Urea", test_values["Urea"], (7, 20), get_risk_level(test_values["Urea"], (7, 20)))
        ],
        "Electrolytes": [
            ("Sodium", test_values["Sodium"], (135, 145), get_risk_level(test_values["Sodium"], (135, 145))),
            ("Potassium", test_values["Potassium"], (3.5, 5.1), get_risk_level(test_values["Potassium"], (3.5, 5.1)))
        ],
        "Vitals": [
            ("SpO2", spo2, (95, 100), get_risk_level(spo2, (95, 100))),
            ("Heart Rate", heart_rate, (60, 100), get_risk_level(heart_rate, (60, 100))),
            ("Respiratory Rate", rr, (12, 20), get_risk_level(rr, (12, 20))),
            ("Temperature", test_values["Temperature"], (97, 99), get_risk_level(test_values["Temperature"], (97, 99))),
            ("BP Systolic", test_values["BP Systolic"], (90, 120), get_risk_level(test_values["BP Systolic"], (90, 120))),
            ("BP Diastolic", test_values["BP Diastolic"], (60, 80), get_risk_level(test_values["BP Diastolic"], (60, 80)))
        ]
    }
    logger.info(f"Test results generated: {test_results}")
    return test_results, frame_rgb, analyzed_image

def preprocess_image(image):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
    ])
    return transform(image).unsqueeze(0)

def predict_xray(image):
    xray_model = load_xray_model()
    if xray_model is None:
        return "Error: X-ray model not loaded.", "", ""
    image_tensor = preprocess_image(image)
    with torch.no_grad():
        outputs = xray_model(image_tensor)
    probs = torch.nn.functional.softmax(outputs[0], dim=0)

    conditions = ["Normal", "Pneumonia", "Cancer", "TB", "Other"]
    results = {conditions[i]: float(probs[i]) for i in range(len(conditions))}

    most_likely_condition = max(results, key=results.get)
    confidence = results[most_likely_condition] * 100

    summary = f"**Summary**: Based on the X-ray analysis, the most likely diagnosis is: <b>{most_likely_condition}</b> with a confidence of <b>{confidence:.2f}%</b>."

    condition_details = {
        "Normal": {
            "description": "The X-ray shows no abnormal signs, and the lungs appear healthy.",
            "recommendation": "No further tests are required. Continue with regular health check-ups."
        },
        "Pneumonia": {
            "description": "Pneumonia is an infection that causes inflammation in the lungs.",
            "recommendation": "Consult a healthcare provider for treatment."
        },
        "Cancer": {
            "description": "Lung cancer may appear as abnormal growths in the lungs.",
            "recommendation": "Consult an oncologist for further diagnostic procedures."
        },
        "TB": {
            "description": "Tuberculosis is a bacterial infection that affects the lungs.",
            "recommendation": "Seek immediate medical attention for treatment."
        },
        "Other": {
            "description": "There may be other conditions requiring investigation.",
            "recommendation": "Consult a radiologist for further analysis."
        }
    }

    detailed_results = "<ul>"
    for condition, prob in results.items():
        detailed_results += f"<li><b>{condition}:</b> {prob*100:.2f}%</li>"
    detailed_results += "</ul>"

    additional_feedback = condition_details.get(most_likely_condition, "Consult a doctor for more details.")

    return summary, detailed_results, additional_feedback

def analyze_report(file):
    text = ""
    if file.name.endswith(".pdf"):
        pdf_reader = PyPDF2.PdfReader(file)
        for page in pdf_reader.pages:
            text += page.extract_text()
    report_summary = f"Patient Report (Preview): {text[:300]}..."
    return report_summary

def main():
    st.markdown("""
    <div class="main-header">
        <h1>🧠 AI Health Report Generator</h1>
        <p>Advanced face-based health analysis using AI technology</p>
    </div>
    """, unsafe_allow_html=True)
    with st.sidebar:
        st.markdown("""
        <div class="patient-form">
            <h2>👤 Patient Information</h2>
        </div>
        """, unsafe_allow_html=True)
        patient_name = st.text_input("👤 Patient Name", placeholder="Enter full name")
        patient_age = st.number_input("🎂 Age", min_value=0, max_value=150, step=1, placeholder="Enter age")
        patient_gender = st.selectbox("⚧ Gender", ["Male", "Female", "Other"])
        patient_id = st.text_input("🆔 Patient ID", placeholder="Optional")
        st.markdown("---")
    col1, col2 = st.columns([1, 1], gap="medium")
    with col1:
        st.markdown("""
        <div class="upload-area">
            <h3>📸 Capture or Upload Media</h3>
            <p>Use camera (with permission) or upload a file</p>
        </div>
        """, unsafe_allow_html=True)
        use_camera = st.checkbox("Enable Camera Capture", value=False)
        if use_camera:
            camera_input = st.camera_input(
                "Take a photo",
                help="Capture a clear front-facing photo for analysis",
                key="camera_input"
            )
        else:
            camera_input = None
        uploaded_file = st.file_uploader(
            "Or upload an image or video",
            type=['jpg', 'jpeg', 'png', 'mp4', 'avi', 'mov'],
            help="Upload a clear front-facing photo or video for analysis"
        )
        st.info("👆 Enable camera or upload a photo/video and click 'Generate Health Report' to see results.", icon="ℹ️")
        input_data = camera_input if camera_input else uploaded_file
        if input_data is not None:
            opencv_image, pil_image = process_input(input_data)
            if opencv_image is not None:
                st.image(pil_image, caption="Captured/Uploaded Media", use_container_width=True)
                if st.button("🔬 Generate Health Report", type="primary"):
                    with st.spinner("Analyzing media and generating report..."):
                        if not patient_name.strip():
                            st.error("Please enter a valid patient name.")
                            return
                        patient_data = {
                            'name': patient_name or "Unknown Patient",
                            'age': patient_age if patient_age > 0 else 'N/A',
                            'gender': patient_gender or "Male",
                            'id': patient_id or "N/A",
                            'date_time': datetime.now().strftime("%d %b %Y %H:%M"),
                            'analysis_method': 'AI Face-Based Health Scan'
                        }
                        test_results, processed_frame, analyzed_image = analyze_face(opencv_image, patient_data)
                        if test_results:
                            analyzed_pil = PILImage.fromarray(cv2.cvtColor(analyzed_image, cv2.COLOR_RGB2BGR))
                            st.image(analyzed_pil, caption="Analyzed Face with Highlighted Regions", use_container_width=True)
                            # Health Summary
                            critical_count = 0
                            key_metrics = {}
                            for category, tests in test_results.items():
                                for test_name, result, range_val, level_info in tests:
                                    level, _ = level_info
                                    if level in ["Low", "High"]:
                                        critical_count += 1
                                    if test_name in ["Hemoglobin", "SpO2", "Heart Rate"]:
                                        key_metrics[test_name] = (result, level)
                            overall_status = "Requires Attention" if critical_count > 2 else "Normal"
                            current_date = datetime.now().strftime("%B %d, %Y %I:%M %p IST")  # 10:42 AM IST, Thursday, July 03, 2025
                            summary_text = f"""
                            <div class='summary-card'>
                                <p><b>📊 Health Summary (as of {current_date}):</b></p>
                                <p><b>Hemoglobin:</b> {key_metrics.get('Hemoglobin', ('N/A', 'N/A'))[0]} g/dL ({key_metrics.get('Hemoglobin', ('N/A', 'N/A'))[1]})</p>
                                <p><b>SpO2:</b> {key_metrics.get('SpO2', ('N/A', 'N/A'))[0]}% ({key_metrics.get('SpO2', ('N/A', 'N/A'))[1]})</p>
                                <p><b>Heart Rate:</b> {key_metrics.get('Heart Rate', ('N/A', 'N/A'))[0]} bpm ({key_metrics.get('Heart Rate', ('N/A', 'N/A'))[1]})</p>
                                <p><b>Critical Values:</b> {critical_count}</p>
                                <p><b>Overall Status:</b> <span style='color: {"#d32f2f" if critical_count > 2 else "#2E8B57"}'>{overall_status}</span></p>
                            </div>
                            """
                            st.markdown(summary_text, unsafe_allow_html=True)
                            st.session_state.test_results = test_results
                            st.session_state.patient_data = patient_data
                            st.session_state.processed_frame = processed_frame
                            st.session_state.pil_image = pil_image
                            st.session_state.analyzed_image = analyzed_pil
                            st.success("✅ Analysis complete! View results in the next column.")
                            logger.info(f"Stored test_results in session_state: {test_results}")
                        else:
                            st.error(f"❌ {analyzed_image}")
                            logger.error(f"Analysis failed: {analyzed_image}")
            else:
                st.error("❌ Failed to process media. Please try again.")
    with col2:
        if 'test_results' in st.session_state:
            st.markdown("""
            <div class="health-card">
                <h2>📋 Health Report</h2>
            </div>
            """, unsafe_allow_html=True)
            patient_data = st.session_state.patient_data
            current_date = datetime.now().strftime("%B %d, %Y %I:%M %p IST")  # 10:42 AM IST, Thursday, July 03, 2025
            st.markdown("""
            <div class='metric-card'>
                <p><b>👤 Patient:</b> {}</p>
                <p><b>🎂 Age:</b> {}</p>
                <p><b>⚧ Gender:</b> {}</p>
                <p><b>🆔 ID:</b> {}</p>
                <p><b>📅 Date:</b> {}</p>
            </div>
            """.format(patient_data['name'], patient_data['age'], patient_data['gender'], patient_data['id'], current_date), unsafe_allow_html=True)
            st.markdown("---")
            # Category icons
            category_icons = {
                "Hematology": "🩺",
                "Iron Panel": "🧪",
                "Liver & Kidney": "🫀",
                "Electrolytes": "⚡",
                "Vitals": "📈"
            }
            for category, tests in st.session_state.test_results.items():
                st.subheader(f"{category_icons.get(category, '📋')} {category}")
                with st.container():
                    st.markdown('<div class="category-container"><div class="table-container">', unsafe_allow_html=True)
                    try:
                        table_data = []
                        for test_name, result, range_val, level_info in tests:
                            level, color = level_info
                            status_indicator = " L" if level == "Low" else " H" if level == "High" else ""
                            if "Count" in test_name or test_name == "Respiratory Rate":
                                value_str = f"{result:.0f}{status_indicator}"
                            elif test_name in ["Temperature", "SpO2"]:
                                value_str = f"{result:.1f}{status_indicator}"
                            else:
                                value_str = f"{result:.1f}{status_indicator}"
                            unit = "" if "BP" in test_name else ("g/dL" if "Hemoglobin" in test_name else
                                                               "cu/mm" if "WBC Count" in test_name else
                                                               "Thousand/µL" if "Platelet Count" in test_name else
                                                               "µg/dL" if "Iron" in test_name or "TIBC" in test_name else
                                                               "ng/mL" if "Ferritin" in test_name else
                                                               "mg/dL" if "Bilirubin" in test_name or "Creatinine" in test_name or "Urea" in test_name else
                                                               "mEq/L" if "Sodium" in test_name or "Potassium" in test_name else
                                                               "%" if "SpO2" in test_name else
                                                               "bpm" if "Heart Rate" in test_name else
                                                               "/min" if "Respiratory Rate" in test_name else
                                                               "°F" if "Temperature" in test_name else "mmHg")
                            range_str = f"{range_val[0]:.0f} - {range_val[1]:.0f}" if "Count" in test_name or test_name == "Respiratory Rate" else f"{range_val[0]:.1f} - {range_val[1]:.1f}"
                            table_data.append({
                                "Test Description": f"{category_icons.get(category, '📋')} {test_name}",
                                "Value Observed": value_str,
                                "Unit": unit,
                                "Biological Reference Interval": range_str,
                                "_color": color
                            })
                        df = pd.DataFrame(table_data)
                        df_display = df.drop(columns=['_color'])
                        def style_row(row):
                            idx = row.name
                            color = df.loc[idx, '_color']
                            return [
                                'text-align: left; padding-left: 0.7rem;',
                                f'color: {color}; text-align: center;',
                                'text-align: center;',
                                'text-align: center;'
                            ]
                        styled_df = df_display.style.set_properties(**{
                            'font-family': 'Helvetica',
                            'font-size': '12px',
                            'border': '1px solid #CCCCCC'
                        }).apply(style_row, axis=1).set_table_styles([
                            {'selector': 'th', 'props': [
                                ('background-color', '#2E8B57'),
                                ('color', 'white'),
                                ('font-weight', 'bold'),
                                ('font-size', '14px'),
                                ('padding', '0.5rem'),
                                ('border', '1px solid #228B22'),
                                ('text-align', 'center')
                            ]},
                            {'selector': 'th:first-child', 'props': [
                                ('text-align', 'left'),
                                ('padding-left', '0.7rem')
                            ]},
                            {'selector': 'tr:nth-child(even)', 'props': [
                                ('background-color', '#F8F8FF')
                            ]},
                            {'selector': 'tr:nth-child(odd)', 'props': [
                                ('background-color', '#F0F8FF')
                            ]},
                            {'selector': 'tr:hover', 'props': [
                                ('background-color', '#E0E0FF')
                            ]}
                        ])
                        st.dataframe(styled_df, use_container_width=True)
                    except Exception as e:
                        st.error(f"Failed to render table for {category}: {str(e)}")
                        logger.error(f"Table rendering failed for {category}: {str(e)}")
                        st.table(table_data)
                    st.markdown('</div></div>', unsafe_allow_html=True)
                    st.markdown("---")
            if st.button("Download PDF Report"):
                with st.spinner("Generating PDF..."):
                    try:
                        pdf_buffer = create_pdf_report(
                            st.session_state.patient_data,
                            st.session_state.test_results,
                            st.session_state.pil_image
                        )
                        if pdf_buffer:
                            st.download_button(
                                label="Download PDF Report",
                                data=pdf_buffer,
                                file_name=f"health_report_{re.sub(r'[^a-zA-Z0-9]', '_', patient_data['name'])}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf",
                                mime="application/pdf"
                            )
                            st.success("✅ PDF generated successfully!")
                            logger.info("PDF download button created")
                        else:
                            st.error("❌ Failed to generate PDF.")
                            logger.error("PDF buffer is None")
                    except Exception as e:
                        st.error(f"❌ Error generating PDF: {str(e)}")
                        logger.error(f"PDF generation error: {str(e)}")
        # Help Report Section with Patient Photo and X-ray Analysis
        if 'pil_image' in st.session_state:
            st.markdown("""
            <div class="help-report">
                <h2>ℹ️ Help Report</h2>
                <div class="help-report-content">
                    <img src="data:image/png;base64,{}" alt="Patient Photo">
                    <div class="help-report-details">
                        <p><b>👤 Name:</b> {}</p>
                        <p><b>🎂 Age:</b> {}</p>
                        <p><b>⚧ Gender:</b> {}</p>
                        <p><b>🆔 ID:</b> {}</p>
                        <p><b>📅 Date:</b> {}</p>
                    </div>
                </div>
            </div>
            """.format(
                base64.b64encode(st.session_state.pil_image.tobytes()).decode(),
                st.session_state.patient_data['name'],
                st.session_state.patient_data['age'],
                st.session_state.patient_data['gender'],
                st.session_state.patient_data['id'],
                current_date
            ), unsafe_allow_html=True)
            # X-ray Analysis
            if st.button("Analyze X-ray"):
                with st.spinner("Analyzing X-ray..."):
                    summary, detailed_results, additional_feedback = predict_xray(st.session_state.pil_image)
                    st.markdown(f'<div class="xray-analysis"><h3>X-ray Diagnosis</h3>{summary}</div>', unsafe_allow_html=True)
                    st.markdown(f'<div class="xray-analysis">{detailed_results}</div>', unsafe_allow_html=True)
                    st.markdown(f'<div class="xray-analysis"><p><b>Additional Feedback:</b> {additional_feedback}</p></div>', unsafe_allow_html=True)

if __name__ == "__main__":
    main()