batch_utils.py

"""
Batch processing and PDF generation utilities for Smartwatch Normative Z-Score Calculator.

Author: Lars Masanneck 2026
"""
import pandas as pd
import numpy as np
from io import BytesIO
from reportlab.lib import colors
from reportlab.lib.pagesizes import A4
from reportlab.platypus import SimpleDocTemplate, Table, TableStyle, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from reportlab.graphics.shapes import Drawing, Rect, Line, String

# Import from the main normalizer model
import normalizer_model

# Friendly biomarker labels (matching the main app)
BIOMARKER_LABELS = {
    "nb_steps": "Number of Steps",
    "max_steps": "Maximum Steps",
    "mean_active_time": "Mean Active Time",
    "sbp": "Systolic Blood Pressure",
    "dbp": "Diastolic Blood Pressure",
    "sleep_duration": "Sleep Duration",
    "avg_night_hr": "Average Night Heart Rate",
    "nb_moderate_active_minutes": "Moderate Active Minutes",
    "nb_vigorous_active_minutes": "Vigorous Active Minutes",
    "weight": "Weight",
    "pwv": "Pulse Wave Velocity",
}

# Biomarkers available for batch processing (excluding disabled ones)
AVAILABLE_BIOMARKERS = [
    "nb_steps",
    "max_steps", 
    "mean_active_time",
    "sleep_duration",
    "avg_night_hr",
    "nb_moderate_active_minutes",
]


def get_batch_template_df():
    """Return a template DataFrame for batch upload."""
    return pd.DataFrame({
        "patient_id": ["P001", "P002", "P003"],
        "age": [45, 62, 38],
        "gender": ["Man", "Woman", "Man"],
        "region": ["Western Europe", "Western Europe", "North America"],
        "bmi": [24.5, 28.1, 22.3],
        "nb_steps": [7500, 4200, 9800],
        "sleep_duration": [7.2, 6.5, 8.1],
        "avg_night_hr": [62, 68, 58],
    })


def process_batch_data(df: pd.DataFrame, normative_df: pd.DataFrame, 
                       biomarkers_to_process: list = None) -> pd.DataFrame:
    """
    Process batch data and add z-score and percentile columns for selected biomarkers.
    
    Parameters
    ----------
    df : pd.DataFrame
        Input data with patient demographics and biomarker values
    normative_df : pd.DataFrame
        Normative reference table
    biomarkers_to_process : list, optional
        List of biomarker columns to process. If None, auto-detect from data.
    
    Returns
    -------
    pd.DataFrame
        Results with z-scores and percentiles added
    """
    results = []
    
    # Auto-detect biomarkers if not specified
    if biomarkers_to_process is None:
        biomarkers_to_process = [col for col in df.columns if col in AVAILABLE_BIOMARKERS]
    
    for _, row in df.iterrows():
        result = row.to_dict()
        
        # Process each biomarker
        for biomarker in biomarkers_to_process:
            if pd.notna(row.get(biomarker)):
                try:
                    res = normalizer_model.compute_normative_position(
                        value=float(row[biomarker]),
                        biomarker=biomarker,
                        age_group=int(row['age']) if pd.notna(row.get('age')) else 45,
                        region=row.get('region', 'Western Europe'),
                        gender=row.get('gender', 'Man'),
                        bmi=float(row.get('bmi', 24.0)) if pd.notna(row.get('bmi')) else 24.0,
                        normative_df=normative_df,
                    )
                    result[f'{biomarker}_z'] = round(res['z_score'], 2)
                    result[f'{biomarker}_percentile'] = round(res['percentile'], 1)
                    
                    # Add interpretation
                    z = res['z_score']
                    if z < -2:
                        result[f'{biomarker}_interpretation'] = 'Very Low'
                    elif z < -1:
                        result[f'{biomarker}_interpretation'] = 'Below Average'
                    elif z < 1:
                        result[f'{biomarker}_interpretation'] = 'Average'
                    elif z < 2:
                        result[f'{biomarker}_interpretation'] = 'Above Average'
                    else:
                        result[f'{biomarker}_interpretation'] = 'Very High'
                        
                except Exception as e:
                    result[f'{biomarker}_z'] = 'N/A'
                    result[f'{biomarker}_percentile'] = 'N/A'
                    result[f'{biomarker}_interpretation'] = f'Error: {str(e)[:30]}'
            else:
                result[f'{biomarker}_z'] = 'N/A'
                result[f'{biomarker}_percentile'] = 'N/A'
                result[f'{biomarker}_interpretation'] = 'No data'
        
        results.append(result)
    
    return pd.DataFrame(results)


def create_z_score_gauge(z_score: float, label: str, width: float = 350, height: float = 100) -> Drawing:
    """Create a horizontal gauge showing z-score position with orange theme."""
    d = Drawing(width, height)
    
    gauge_y = 35
    gauge_height = 25
    gauge_left = 50
    gauge_width = width - 100
    
    # Color zones - orange themed
    zone_colors = [
        (colors.HexColor('#2ecc71'), -2),   # Green - very low (good for some metrics)
        (colors.HexColor('#27ae60'), -1),   # Darker green
        (colors.HexColor('#f39c12'), 0),    # Orange - average
        (colors.HexColor('#e67e22'), 1),    # Darker orange
        (colors.HexColor('#d35400'), 2),    # Deep orange
        (colors.HexColor('#c0392b'), 3),    # Red - extreme
    ]
    
    zone_width = gauge_width / 6
    for i, (color, _) in enumerate(zone_colors):
        d.add(Rect(gauge_left + i * zone_width, gauge_y, zone_width, gauge_height, 
                   fillColor=color, strokeColor=None))
    
    # Border
    d.add(Rect(gauge_left, gauge_y, gauge_width, gauge_height, 
               fillColor=None, strokeColor=colors.black, strokeWidth=1))
    
    # Marker position (clamp z to -3, 3)
    clamped_z = max(-3, min(3, z_score))
    marker_x = gauge_left + ((clamped_z + 3) / 6) * gauge_width
    
    # Marker line
    d.add(Line(marker_x, gauge_y - 8, marker_x, gauge_y + gauge_height + 8, 
               strokeColor=colors.black, strokeWidth=3))
    
    # Scale labels
    for i, val in enumerate([-3, -2, -1, 0, 1, 2, 3]):
        x = gauge_left + (i / 6) * gauge_width
        d.add(String(x, gauge_y - 15, str(val), fontSize=9, textAnchor='middle'))
    
    # Title
    d.add(String(width / 2, height - 8, label, fontSize=11, textAnchor='middle', fontName='Helvetica-Bold'))
    
    # Z-score value
    d.add(String(width / 2, gauge_y + gauge_height + 18, f"Z = {z_score:.2f}", 
                 fontSize=10, textAnchor='middle', fontName='Helvetica-Bold'))
    
    return d


def generate_pdf_report(patient_info: dict, measurements: dict, z_scores: dict = None) -> BytesIO:
    """
    Generate a PDF report for a patient with Z-scores and graphs.
    
    Parameters
    ----------
    patient_info : dict
        Patient demographics (age, gender, region, bmi)
    measurements : dict
        Biomarker measurements (biomarker_code: value)
    z_scores : dict
        Z-score results for each biomarker
    
    Returns
    -------
    BytesIO
        PDF buffer ready for download
    """
    buffer = BytesIO()
    doc = SimpleDocTemplate(buffer, pagesize=A4, topMargin=0.5*inch, bottomMargin=0.5*inch)
    
    styles = getSampleStyleSheet()
    
    # Orange-themed styles
    title_style = ParagraphStyle(
        'Title', 
        parent=styles['Heading1'], 
        fontSize=18, 
        spaceAfter=12, 
        alignment=1,
        textColor=colors.HexColor('#d35400')
    )
    heading_style = ParagraphStyle(
        'Heading', 
        parent=styles['Heading2'], 
        fontSize=14, 
        spaceAfter=8, 
        spaceBefore=12,
        textColor=colors.HexColor('#e67e22')
    )
    normal_style = styles['Normal']
    
    elements = []
    
    # Title
    elements.append(Paragraph("Smartwatch Normative Z-Score Report", title_style))
    elements.append(Spacer(1, 0.2*inch))
    
    # Patient Information
    elements.append(Paragraph("Demographics", heading_style))
    patient_data = [
        ["Age:", f"{patient_info.get('age', 'N/A')} years"],
        ["Gender:", patient_info.get('gender', 'N/A')],
        ["Region:", patient_info.get('region', 'N/A')],
        ["BMI:", f"{patient_info.get('bmi', 'N/A')}"],
    ]
    patient_table = Table(patient_data, colWidths=[2*inch, 4*inch])
    patient_table.setStyle(TableStyle([
        ('FONTNAME', (0, 0), (0, -1), 'Helvetica-Bold'),
        ('ALIGN', (0, 0), (-1, -1), 'LEFT'),
        ('VALIGN', (0, 0), (-1, -1), 'MIDDLE'),
        ('BOTTOMPADDING', (0, 0), (-1, -1), 6),
    ]))
    elements.append(patient_table)
    elements.append(Spacer(1, 0.2*inch))
    
    # Measurements
    if measurements:
        elements.append(Paragraph("Measurements", heading_style))
        measurements_data = []
        for biomarker, value in measurements.items():
            label = BIOMARKER_LABELS.get(biomarker, biomarker.replace('_', ' ').title())
            measurements_data.append([f"{label}:", f"{value}"])
        
        if measurements_data:
            meas_table = Table(measurements_data, colWidths=[2.5*inch, 3.5*inch])
            meas_table.setStyle(TableStyle([
                ('FONTNAME', (0, 0), (0, -1), 'Helvetica-Bold'),
                ('ALIGN', (0, 0), (-1, -1), 'LEFT'),
                ('VALIGN', (0, 0), (-1, -1), 'MIDDLE'),
                ('BOTTOMPADDING', (0, 0), (-1, -1), 6),
            ]))
            elements.append(meas_table)
        elements.append(Spacer(1, 0.2*inch))
    
    # Z-Score Analysis
    if z_scores:
        elements.append(Paragraph("Z-Score Analysis", heading_style))
        elements.append(Paragraph(
            "Z-scores indicate how many standard deviations a measurement is from the population mean. "
            "Values between -2 and +2 are typically considered within normal range.",
            ParagraphStyle('ZInfo', parent=normal_style, fontSize=9, textColor=colors.grey, spaceAfter=8)
        ))
        
        # Z-score table
        z_data = [["Biomarker", "Value", "Z-Score", "Percentile", "Interpretation"]]
        
        for biomarker, data in z_scores.items():
            if isinstance(data, dict) and 'z_score' in data:
                z = data['z_score']
                pct = data['percentile']
                value = measurements.get(biomarker, 'N/A')
                label = BIOMARKER_LABELS.get(biomarker, biomarker.replace('_', ' ').title())
                
                # Interpretation
                if z < -2:
                    interp = "Very Low"
                elif z < -1:
                    interp = "Below Average"
                elif z < 1:
                    interp = "Average"
                elif z < 2:
                    interp = "Above Average"
                else:
                    interp = "Very High"
                
                z_data.append([label, str(value), f"{z:.2f}", f"{pct:.1f}%", interp])
        
        if len(z_data) > 1:
            z_table = Table(z_data, colWidths=[1.5*inch, 1*inch, 0.8*inch, 1*inch, 1.2*inch])
            z_table.setStyle(TableStyle([
                ('BACKGROUND', (0, 0), (-1, 0), colors.HexColor('#e67e22')),
                ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
                ('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),
                ('FONTSIZE', (0, 0), (-1, -1), 9),
                ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
                ('VALIGN', (0, 0), (-1, -1), 'MIDDLE'),
                ('GRID', (0, 0), (-1, -1), 0.5, colors.grey),
                ('BOTTOMPADDING', (0, 0), (-1, -1), 6),
                ('TOPPADDING', (0, 0), (-1, -1), 6),
            ]))
            elements.append(z_table)
            elements.append(Spacer(1, 0.15*inch))
            
            # Add Z-score gauges
            for biomarker, data in z_scores.items():
                if isinstance(data, dict) and 'z_score' in data:
                    label = BIOMARKER_LABELS.get(biomarker, biomarker.replace('_', ' ').title())
                    gauge = create_z_score_gauge(data['z_score'], label)
                    elements.append(gauge)
                    elements.append(Spacer(1, 0.1*inch))
        
        elements.append(Spacer(1, 0.2*inch))
    
    # Cohort Information
    elements.append(Paragraph("Reference Population", heading_style))
    cohort_text = (
        f"Z-scores calculated using normative data from Withings users in "
        f"{patient_info.get('region', 'Western Europe')}, filtered by gender "
        f"({patient_info.get('gender', 'N/A')}), age group, and BMI category."
    )
    elements.append(Paragraph(cohort_text, normal_style))
    elements.append(Spacer(1, 0.3*inch))
    
    # Disclaimer
    disclaimer = Paragraph(
        "<i>This report is for educational and research purposes only. Z-scores are based on "
        "Withings population data and may not reflect clinical reference ranges. For detailed "
        "questions regarding personal health data, contact your healthcare professionals.</i>",
        ParagraphStyle('Disclaimer', parent=normal_style, fontSize=8, textColor=colors.grey)
    )
    elements.append(disclaimer)
    
    # Footer
    elements.append(Spacer(1, 0.2*inch))
    footer = Paragraph(
        "Built with ❤️ in Düsseldorf. © Lars Masanneck 2026.",
        ParagraphStyle('Footer', parent=normal_style, fontSize=8, textColor=colors.grey, alignment=1)
    )
    elements.append(footer)
    
    doc.build(elements)
    buffer.seek(0)
    return buffer