# Verify sample document exists
import os

# Path to your pdf (using local file)
pdf_path = "patient-packet-completed.pdf"

if os.path.exists(pdf_path):
    print(f"✅ Found: {pdf_path}")
    print(f"   Size: {os.path.getsize(pdf_path) / 1024:.1f} KB")
else:
    print(f"❌ File not found: {pdf_path}")
    print("   Please ensure patient-packet-completed.pdf is in the working directory")

# List available sample files in the workspace
!ls -la *.pdf 2>/dev/null || echo "No PDF files found in current directory"

# Initialize Mistral client with API key
import os
from mistralai import Mistral
from google.colab import userdata
import requests

api_key = userdata.get('MISTRAL_API_KEY') # Replace with your way to retrieve API key

if not api_key:
    print("⚠️  WARNING: No API key found!")
else:
    client = Mistral(api_key=api_key)
    print("✅ Mistral client initialized")

print(json.dumps(response_dict, indent=4)[0:1000]) # check the first 1000 characters

# Display Page 1 - Form Elements
print("📄 PAGE 1: Patient Admission Form")
print("Notice: Checkboxes rendered as ☐ (unchecked) and ☑ (checked)\n")
display_page_with_tables(0, response_dict, pdf_response)

# Display Page 2 - Vital Signs Flowsheet with HTML table
print("📄 PAGE 2: Vital Signs Flowsheet")
print("Notice: Tables output as HTML with rowspan/colspan preserved\n")
display_page_with_tables(1, response_dict, pdf_response)

# Display Page 3 - X-ray Image
print("📄 PAGE 3: Foot X-ray")
print("Notice: Images are detected and embedded with base64 encoding\n")

page3 = pdf_response.pages[2]

# Show image metadata
print(f"Images detected on this page: {len(page3.images)}")
for img in page3.images:
    print(f"  - Image ID: {img.id}")
    print(f"    Dimensions: ({img.top_left_x}, {img.top_left_y}) to ({img.bottom_right_x}, {img.bottom_right_y})")

print("\n" + "="*60)
print("Rendered Output (with embedded X-ray image):")
print("="*60)

# Display page with images and tables
display_page_with_tables(2, response_dict, pdf_response)

import base64
import json
import time
import uuid
import os
from datetime import datetime
from enum import Enum
from typing import List, Dict, Any, Optional
from dataclasses import dataclass
from collections import Counter

from pydantic import BaseModel, Field, JsonValue
from mistralai import Mistral
from mistralai.extra import response_format_from_pydantic_model
from google.colab import userdata

api_key = userdata.get('MISTRAL_API_KEY') # Replace with your way to retrieve API key

if not api_key:
    print("⚠️  WARNING: No API key found!")
    print("   Set MISTRAL_API_KEY environment variable, or")
    print("   Uncomment and set api_key directly above")
else:
    client = Mistral(api_key=api_key)
    print("✅ Mistral client initialized")

def encode_pdf(pdf_path: str) -> Optional[str]:
    """Encode a PDF file to base64 string."""
    try:
        with open(pdf_path, "rb") as pdf_file:
            return base64.b64encode(pdf_file.read()).decode('utf-8')
    except FileNotFoundError:
        print(f"Error: The file {pdf_path} was not found.")
        return None
    except Exception as e:
        print(f"Error: {e}")
        return None

print("✅ Production patterns setup complete")

# Classify the patient packet document

base64_packet = encode_pdf(pdf_path)

# First-pass classification using only page 1
classification_response = client.ocr.process(
    model="mistral-ocr-latest",
    document={
        "type": "document_url",
        "document_url": f"data:application/pdf;base64,{base64_packet}"
    },
    pages=list(range(8)),  # Document Annotations has a limit of 8 pages, we recommend spliting your documents when using it; bbox annotations does not have the same limit
    document_annotation_format=response_format_from_pydantic_model(DocumentClassification),
    include_image_base64=False
)

# Parse and display classification
classification = json.loads(classification_response.document_annotation)
print("📋 Document Classification Results")
print("=" * 50)
print(f"Type:        {classification['document_type']}")
print(f"Confidence:  {classification['confidence']:.0%}")
print(f"Route to:    {classification['routing_department']}")
print(f"Urgency:     {classification['urgency']}")
print(f"Identifiers: {', '.join(classification['key_identifiers_found'])}")
print(f"Signature:   {'Yes' if classification['requires_signature'] else 'No'}")
print(f"\nSummary: {classification['summary']}")

# Define type-specific extraction schemas for each document type

class PatientDemographics(BaseModel):
    """Schema for patient demographics and intake forms."""
    patient_name: str = Field(..., description="Full patient name")
    date_of_birth: str = Field(..., description="DOB in MM/DD/YYYY format")
    gender: Optional[str] = Field(None, description="Patient gender")
    address: Optional[str] = Field(None, description="Patient address")
    phone: Optional[str] = Field(None, description="Contact phone number")
    insurance_id: Optional[str] = Field(None, description="Insurance member ID")
    emergency_contact: Optional[str] = Field(None, description="Emergency contact info")

class VitalsFlowsheet(BaseModel):
    """Schema for vital signs flowsheets."""
    date_recorded: str = Field(..., description="Date vitals were recorded")
    blood_pressure: Optional[str] = Field(None, description="Blood pressure reading (systolic/diastolic)")
    heart_rate: Optional[str] = Field(None, description="Heart rate in BPM")
    temperature: Optional[str] = Field(None, description="Body temperature")
    respiratory_rate: Optional[str] = Field(None, description="Respiratory rate")
    oxygen_saturation: Optional[str] = Field(None, description="SpO2 percentage")
    weight: Optional[str] = Field(None, description="Patient weight")
    height: Optional[str] = Field(None, description="Patient height")

# Map document types to their extraction schemas
EXTRACTION_SCHEMAS: Dict[str, type] = {
    "patient_demographics": PatientDemographics,
    "vitals_flowsheet": VitalsFlowsheet,
    # Add more mappings as needed
}

print("✅ Defined extraction schemas for:", list(EXTRACTION_SCHEMAS.keys()))

@dataclass
class PageResult:
    """Result container for a single processed page"""
    page_index: int
    document_type: str
    classification: Dict[str, Any]
    extracted_data: Optional[Dict[str, Any]]
    markdown_content: str
    status: str  # "success", "error", "skipped"
    error_message: Optional[str] = None
    processing_time_ms: float = 0

def process_patient_packet(pdf_path: str, rate_limit_delay: float = 0.5) -> List[PageResult]:
    """
    Process a multi-page patient packet with per-page classification and extraction.

    Args:
        pdf_path: Path to the PDF file
        rate_limit_delay: Delay between API calls to respect rate limits

    Returns:
        List of PageResult objects with classification and extracted data
    """
    base64_pdf = encode_pdf(pdf_path)
    results = []

    # First, get the full document to know page count
    full_response = client.ocr.process(
        model="mistral-ocr-latest",
        document={
            "type": "document_url",
            "document_url": f"data:application/pdf;base64,{base64_pdf}"
        },
        include_image_base64=False
    )

    num_pages = len(full_response.pages)
    print(f"📄 Processing {num_pages} pages from {pdf_path}")
    print("-" * 50)

    for page_idx in range(num_pages):
        start_time = time.time()

        try:
            # Step 1: Classify this page
            time.sleep(rate_limit_delay)  # Rate limiting

            classify_response = client.ocr.process(
                model="mistral-ocr-latest",
                document={
                    "type": "document_url",
                    "document_url": f"data:application/pdf;base64,{base64_pdf}"
                },
                pages=[page_idx],
                document_annotation_format=response_format_from_pydantic_model(DocumentClassification),
                include_image_base64=False
            )

            page_classification = json.loads(classify_response.document_annotation)
            doc_type = page_classification["document_type"]

            # Step 2: Extract with type-specific schema if available
            extracted_data = None
            if doc_type in EXTRACTION_SCHEMAS:
                time.sleep(rate_limit_delay)

                extract_response = client.ocr.process(
                    model="mistral-ocr-latest",
                    document={
                        "type": "document_url",
                        "document_url": f"data:application/pdf;base64,{base64_pdf}"
                    },
                    pages=[page_idx],
                    document_annotation_format=response_format_from_pydantic_model(EXTRACTION_SCHEMAS[doc_type]),
                    include_image_base64=False
                )
                extracted_data = json.loads(extract_response.document_annotation)

            processing_time = (time.time() - start_time) * 1000

            result = PageResult(
                page_index=page_idx,
                document_type=doc_type,
                classification=page_classification,
                extracted_data=extracted_data,
                markdown_content=full_response.pages[page_idx].markdown,
                status="success",
                processing_time_ms=processing_time
            )

            print(f"  ✅ Page {page_idx + 1}: {doc_type} ({processing_time:.0f}ms)")

        except Exception as e:
            result = PageResult(
                page_index=page_idx,
                document_type="unknown",
                classification={},
                extracted_data=None,
                markdown_content="",
                status="error",
                error_message=str(e),
                processing_time_ms=(time.time() - start_time) * 1000
            )
            print(f"  ❌ Page {page_idx + 1}: Error - {str(e)[:50]}")

        results.append(result)

    return results

# Process the patient packet
batch_results = process_patient_packet("patient-packet-completed.pdf")

# Summary statistics
print("\n" + "=" * 50)
print("📊 BATCH PROCESSING SUMMARY")
print("=" * 50)
successful = [r for r in batch_results if r.status == "success"]
failed = [r for r in batch_results if r.status == "error"]

print(f"Total pages:     {len(batch_results)}")
print(f"Successful:      {len(successful)}")
print(f"Failed:          {len(failed)}")
print(f"Total time:      {sum(r.processing_time_ms for r in batch_results):.0f}ms")

# Group by document type
doc_types = Counter(r.document_type for r in successful)
print(f"\nDocument types found:")
for doc_type, count in doc_types.items():
    print(f"  • {doc_type}: {count} page(s)")

# Display extracted data for each page
print("📋 EXTRACTED DATA BY PAGE")
print("=" * 50)

for result in batch_results:
    if result.status == "success":
        print(f"\n🔹 Page {result.page_index + 1}: {result.document_type}")
        print(f"   Confidence: {result.classification.get('confidence', 'N/A'):.0%}")
        print(f"   Route to: {result.classification.get('routing_department', 'N/A')}")

        if result.extracted_data:
            print("   Extracted fields:")
            for key, value in result.extracted_data.items():
                if value:  # Only show non-null values
                    print(f"     • {key}: {value}")

# FHIR R4 Resource Generation from extracted OCR data

def generate_fhir_patient(demographics: Dict[str, Any]) -> Dict[str, Any]:
    """Convert extracted demographics to FHIR R4 Patient resource"""
    # Parse name (assumes "Last, First" or "First Last" format)
    name_parts = demographics.get("patient_name", "Unknown").replace(",", " ").split()

    patient = {
        "resourceType": "Patient",
        "id": str(uuid.uuid4()),
        "meta": {
            "profile": ["http://hl7.org/fhir/us/core/StructureDefinition/us-core-patient"]
        },
        "identifier": [{
            "system": "urn:oid:2.16.840.1.113883.4.1",  # Example OID
            "value": demographics.get("insurance_id", "UNKNOWN")
        }],
        "name": [{
            "use": "official",
            "family": name_parts[0] if name_parts else "Unknown",
            "given": name_parts[1:] if len(name_parts) > 1 else []
        }],
        "birthDate": convert_date_to_fhir(demographics.get("date_of_birth")),
        "gender": map_gender(demographics.get("gender")),
    }

    # Add address if present
    if demographics.get("address"):
        patient["address"] = [{
            "use": "home",
            "text": demographics["address"]
        }]

    # Add phone if present
    if demographics.get("phone"):
        patient["telecom"] = [{
            "system": "phone",
            "value": demographics["phone"],
            "use": "home"
        }]

    return patient

def generate_fhir_vitals(vitals: Dict[str, Any], patient_id: str) -> List[Dict[str, Any]]:
    """Convert extracted vitals to FHIR R4 Observation resources"""
    observations = []

    # LOINC codes for common vitals
    vital_mappings = {
        "blood_pressure": {"code": "85354-9", "display": "Blood pressure panel"},
        "heart_rate": {"code": "8867-4", "display": "Heart rate", "unit": "/min"},
        "temperature": {"code": "8310-5", "display": "Body temperature", "unit": "Cel"},
        "respiratory_rate": {"code": "9279-1", "display": "Respiratory rate", "unit": "/min"},
        "oxygen_saturation": {"code": "2708-6", "display": "Oxygen saturation", "unit": "%"},
        "weight": {"code": "29463-7", "display": "Body weight", "unit": "kg"},
        "height": {"code": "8302-2", "display": "Body height", "unit": "cm"}
    }

    effective_date = convert_date_to_fhir(vitals.get("date_recorded")) or datetime.now().strftime("%Y-%m-%d")

    for vital_key, loinc in vital_mappings.items():
        value = vitals.get(vital_key)
        if value:
            observation = {
                "resourceType": "Observation",
                "id": str(uuid.uuid4()),
                "status": "final",
                "category": [{
                    "coding": [{
                        "system": "http://terminology.hl7.org/CodeSystem/observation-category",
                        "code": "vital-signs",
                        "display": "Vital Signs"
                    }]
                }],
                "code": {
                    "coding": [{
                        "system": "http://loinc.org",
                        "code": loinc["code"],
                        "display": loinc["display"]
                    }]
                },
                "subject": {"reference": f"Patient/{patient_id}"},
                "effectiveDateTime": effective_date,
                "valueString": str(value)  # Using string for flexibility; production would parse numeric
            }
            observations.append(observation)

    return observations

# Helper functions
def convert_date_to_fhir(date_str: Optional[str]) -> Optional[str]:
    """Convert various date formats to FHIR format (YYYY-MM-DD)"""
    if not date_str:
        return None
    # Handle MM/DD/YYYY format
    try:
        parts = date_str.replace("-", "/").split("/")
        if len(parts) == 3:
            if len(parts[0]) == 4:  # Already YYYY-MM-DD
                return date_str
            return f"{parts[2]}-{parts[0].zfill(2)}-{parts[1].zfill(2)}"
    except:
        pass
    return date_str

def map_gender(gender_str: Optional[str]) -> str:
    """Map various gender representations to FHIR values"""
    if not gender_str:
        return "unknown"
    g = gender_str.lower().strip()
    if g in ["m", "male"]:
        return "male"
    elif g in ["f", "female"]:
        return "female"
    return "unknown"

print("✅ FHIR resource generators ready")