# Face Detection-Based AI Automation of Lab Tests
# UI: Clean table, multilingual summary, PDF-ready

import gradio as gr
import cv2
import numpy as np
import mediapipe as mp
from fpdf import FPDF
import os

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)

def estimate_heart_rate(frame, landmarks):
    h, w, _ = frame.shape
    forehead_pts = [landmarks[10], landmarks[338], landmarks[297], landmarks[332]]
    mask = np.zeros((h, w), dtype=np.uint8)
    pts = np.array([[int(pt.x * w), int(pt.y * h)] for pt in forehead_pts], np.int32)
    cv2.fillConvexPoly(mask, pts, 255)
    green_channel = cv2.split(frame)[1]
    mean_intensity = cv2.mean(green_channel, mask=mask)[0]
    heart_rate = int(60 + 30 * np.sin(mean_intensity / 255.0 * np.pi))
    return heart_rate

def estimate_spo2_rr(heart_rate):
    spo2 = min(100, max(90, 97 + (heart_rate % 5 - 2)))
    rr = int(12 + abs(heart_rate % 5 - 2))
    return spo2, rr

def get_risk_color(value, normal_range):
    low, high = normal_range
    if value < low:
        return ("Low", "🔻", "#FFCCCC")
    elif value > high:
        return ("High", "🔺", "#FFE680")
    else:
        return ("Normal", "✅", "#CCFFCC")

def generate_pdf_report(image, results_dict, summary_text):
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", "B", 16)
    pdf.cell(0, 10, "SL Diagnostics - Face Scan AI Lab Report", ln=True, align='C')

    if image is not None:
        img_path = "patient_face.jpg"
        cv2.imwrite(img_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
        pdf.image(img_path, x=80, y=25, w=50)
        os.remove(img_path)
        pdf.ln(60)

    pdf.set_font("Arial", "B", 12)
    pdf.cell(0, 10, "Results Summary", ln=True)
    pdf.set_font("Arial", "", 10)

    for key, val in results_dict.items():
        if isinstance(val, (int, float)):
            pdf.cell(0, 8, f"{key}: {val}", ln=True)

    pdf.ln(5)
    pdf.set_font("Arial", "B", 12)
    pdf.cell(0, 10, "AI Summary (English)", ln=True)
    pdf.set_font("Arial", "", 10)
    for line in summary_text.split("<li>"):
        if "</li>" in line:
            clean = line.split("</li>")[0].strip()
            pdf.multi_cell(0, 8, f"- {clean}")

    output_path = "/mnt/data/SL_Diagnostics_Face_Scan_Report.pdf"
    pdf.output(output_path)
    return output_path

def infer_lab_results(image, landmarks):
    h, w, _ = image.shape
    forehead = image[int(0.1*h):int(0.25*h), int(0.35*w):int(0.65*w)]
    mean_intensity = np.mean(cv2.cvtColor(forehead, cv2.COLOR_BGR2GRAY))
    skin_redness = np.mean(image[:, :, 2]) - np.mean(image[:, :, 1])

    return {
        'Hemoglobin': round(10 + (mean_intensity / 255.0) * 7, 1),
        'WBC Count': round(4 + (1 - mean_intensity / 255.0) * 7, 1),
        'Platelets': int(150 + (mean_intensity / 255.0) * 150),
        'Iron': round(40 + (skin_redness / 50.0) * 40, 1),
        'Ferritin': round(25 + (skin_redness / 50.0) * 70, 1),
        'TIBC': round(250 + ((255 - mean_intensity) / 255.0) * 150, 1),
        'Bilirubin': round(0.5 + (255 - mean_intensity) / 255.0 * 1.5, 2),
        'Creatinine': round(0.8 + (skin_redness / 255.0) * 0.6, 2),
        'TSH': round(1.0 + (skin_redness / 255.0) * 2.0, 2),
        'Cortisol': round(12 + (skin_redness / 255.0) * 10, 2),
        'Fasting Blood Sugar': int(80 + (skin_redness / 255.0) * 60),
        'HbA1c': round(5.0 + (skin_redness / 255.0) * 1.5, 2),
        'SpO2': round(97 - (255 - mean_intensity) / 255.0 * 5, 1),
        'Heart Rate': estimate_heart_rate(image, landmarks),
        'Respiratory Rate': estimate_spo2_rr(estimate_heart_rate(image, landmarks))[1]
    }

# Gradio UI (app launcher)
def app():
    def process(image):
        if image is None:
            return "Please upload a face image.", None, None

        frame_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        result = face_mesh.process(frame_rgb)
        if not result.multi_face_landmarks:
            return "Face not detected.", None, None

        landmarks = result.multi_face_landmarks[0].landmark
        heart_rate = estimate_heart_rate(frame_rgb, landmarks)
        spo2, rr = estimate_spo2_rr(heart_rate)

        results_dict = infer_lab_results(frame_rgb, landmarks)

        summary_text = "<li>Your hemoglobin is a bit low...</li><li>Consider iron tests.</li>"  # Placeholder
        pdf_path = generate_pdf_report(image, results_dict, summary_text)

        return "Preview complete. You can download your report.", frame_rgb, pdf_path

    with gr.Blocks() as demo:
        gr.Markdown("""# 🧠 Face-Based Lab Test AI Report""")
        with gr.Row():
            with gr.Column():
                image = gr.Image(label="📸 Upload Face", type="numpy")
                button = gr.Button("🔍 Run Analysis")
                pdf_output = gr.File(label="📄 Download Report")
            with gr.Column():
                note = gr.Textbox(label="Status")
                preview = gr.Image(label="Scan Preview")

        button.click(fn=process, inputs=image, outputs=[note, preview, pdf_output])

    demo.launch()

app()