"""
MCA SDK Example: Hospital Readmission Prediction Model

Demonstrates core SDK usage with Registry API integration.
Uses a real Random Forest model trained on Diabetes 130-US Hospitals dataset for 30-day readmission prediction.
"""

import os
import sys
import time
import pickle
import pandas as pd

# Add both sdk-examples (for shared.py) and mca-prototype root (for mca_sdk) to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))

from shared import (
    print_gcp_auth_status,
    print_collector_status,
    COLLECTOR_HTTP_ENDPOINT,
    REGISTRY_URL,
    INTERNAL_MODEL_ID,
)

from mca_sdk import MCAClient
from mca_sdk.registry import RegistryClient


def load_model():
    """Load trained hospital readmission prediction model."""
    model_path = os.path.join(os.path.dirname(__file__), "readmission_model.pkl")

    if not os.path.exists(model_path):
        print(f"Model not found at {model_path}")
        print("Please run: python train_model.py")
        sys.exit(1)

    with open(model_path, "rb") as f:
        artifacts = pickle.load(f)

    return artifacts["model"], artifacts["scaler"], artifacts["feature_names"], artifacts["metrics"]


def main():
    print("MCA SDK - Clinical Prediction Model Example\n")

    # Load trained model
    print("=== Loading Model ===")
    model, scaler, feature_names, metrics = load_model()
    print(f"✓ Model loaded (Accuracy: {metrics['accuracy']:.3f})")
    print(f"✓ Features: {len(feature_names)}")
    print()

    # Step 1-2: Verify prerequisites
    if not print_gcp_auth_status() or not print_collector_status():
        sys.exit(1)

    # Step 3-5: Registry integration
    print("\n=== Registry Integration ===")
    model_id = os.environ.get("MODEL_ID", INTERNAL_MODEL_ID)

    registry_client = RegistryClient(url=REGISTRY_URL, use_gcp_auth=True, timeout=10.0)

    # Fetch specific model config
    model_config = registry_client.fetch_model_config(model_id)
    print(f"✓ Using: {model_config.service_name} ({model_config.team_name})")
    if model_config.thresholds:
        print(f"✓ Thresholds: {model_config.thresholds}")
    print()

    # Step 6: Initialize SDK
    print("=== SDK Initialization ===")
    client = MCAClient(
        service_name=model_config.service_name,
        model_id=model_id,
        model_version=None,
        team_name=model_config.team_name,
        collector_endpoint=COLLECTOR_HTTP_ENDPOINT,
        allow_insecure_collector=True,
        registry_url=REGISTRY_URL,
        use_gcp_auth=True,
    )
    print("✓ SDK initialized\n")

    try:
        # Step 7: Input/Output Capture
        print("=== Input/Output Capture ===")

        @client.predict(capture_input=True, capture_output=True)
        def predict_readmission(encounter_data: dict) -> dict:
            """Predict 30-day hospital readmission risk using trained Random Forest model."""
            # Create DataFrame with proper feature names to match training
            feature_data = pd.DataFrame(
                [
                    {
                        "time_in_hospital": encounter_data.get("time_in_hospital", 5),
                        "num_lab_procedures": encounter_data.get("num_lab_procedures", 40),
                        "num_procedures": encounter_data.get("num_procedures", 2),
                        "num_medications": encounter_data.get("num_medications", 15),
                        "number_diagnoses": encounter_data.get("number_diagnoses", 7),
                        "age_numeric": encounter_data.get("age_numeric", 65),
                    }
                ],
                columns=feature_names,
            )

            # Scale and predict
            features_scaled = scaler.transform(feature_data)
            prediction = model.predict(features_scaled)[0]
            probability = model.predict_proba(features_scaled)[0]

            return {
                "prediction": "readmit_30_days" if prediction == 1 else "no_readmission",
                "readmission_risk": float(probability[1]),
                "confidence": float(max(probability)),
            }

        # Test with sample encounter
        sample_encounter = {
            "encounter_id": "sanitized-enc-001",
            "patient_nbr": "sanitized-pat-001",
            "time_in_hospital": 7,
            "num_lab_procedures": 55,
            "num_procedures": 3,
            "num_medications": 18,
            "number_diagnoses": 9,
            "age_numeric": 75,
        }
        result = predict_readmission(sample_encounter)
        print(
            f"✓ Prediction: {result['prediction']} (risk: {result['readmission_risk']:.2f}, confidence: {result['confidence']:.2f})\n"
        )

        # Step 8: Manual predictions with threshold checking
        print("=== Manual Predictions with Threshold Checking ===")
        latency_threshold_ms = client.thresholds.get("latency_ms", 1000)

        # Generate test encounters with varying characteristics
        test_encounters = [
            {
                "time_in_hospital": 3,
                "num_lab_procedures": 35,
                "num_procedures": 1,
                "num_medications": 10,
                "number_diagnoses": 5,
                "age_numeric": 55,
            },
            {
                "time_in_hospital": 10,
                "num_lab_procedures": 75,
                "num_procedures": 5,
                "num_medications": 25,
                "number_diagnoses": 12,
                "age_numeric": 82,
            },
            {
                "time_in_hospital": 5,
                "num_lab_procedures": 45,
                "num_procedures": 2,
                "num_medications": 15,
                "number_diagnoses": 7,
                "age_numeric": 48,
            },
        ]

        for i, encounter in enumerate(test_encounters):
            start_time = time.time()
            feature_data = pd.DataFrame([encounter], columns=feature_names)
            features_scaled = scaler.transform(feature_data)
            probability = model.predict_proba(features_scaled)[0]
            latency = time.time() - start_time

            client.record_prediction(
                latency=latency,
                prediction_id=f"pred-{i + 1:03d}",  # Auto-stripped by SDK
                prediction_category="high_risk"
                if probability[1] > 0.5
                else "low_risk",  # Low-cardinality metric label
                span_attributes={
                    "ml.prediction.readmission_risk": float(
                        probability[1]
                    ),  # High-cardinality trace attribute
                },
                department="internal_medicine",
            )

            # Check threshold
            if latency * 1000 > latency_threshold_ms:
                print(
                    f"  ⚠️  pred-{i + 1:03d}: High latency {latency * 1000:.0f}ms (threshold: {latency_threshold_ms}ms)"
                )
            else:
                print(
                    f"  ✓ pred-{i + 1:03d}: {latency * 1000:.0f}ms (readmission risk: {probability[1]:.2f})"
                )
        print()

        # Step 9: Custom operation with trace()
        print("=== Custom Operation with trace() ===")
        with client.trace("data_preprocessing") as span:
            time.sleep(0.03)
            span.set_attribute("records_processed", 1000)
            span.set_attribute("features_extracted", 42)
            print("✓ Data preprocessing traced\n")

        # Step 10: Custom metrics
        print("=== Custom Metrics ===")
        client.record_counter("model.cache_hits", 15, cache_type="redis")
        client.record_gauge("model.queue_size", 8, queue_name="predictions")
        client.record_metric("model.confidence", 0.87, metric_type="histogram")
        print("✓ Recorded custom metrics\n")

        # Step 10b: Evaluation metrics
        print("=== Evaluation Metrics ===")
        # Use actual metrics from trained model
        client.record_metric(
            "model.accuracy_ratio", metrics["accuracy"], dataset="validation", metric_type="gauge"
        )
        client.record_metric(
            "model.precision_ratio", metrics["precision"], dataset="validation", metric_type="gauge"
        )
        client.record_metric(
            "model.recall_ratio", metrics["recall"], dataset="validation", metric_type="gauge"
        )
        print(
            f"✓ Recorded classification metrics (accuracy={metrics['accuracy']:.3f}, precision={metrics['precision']:.3f}, recall={metrics['recall']:.3f})\n"
        )

        # Step 11: Error tracking
        print("=== Error Tracking ===")
        try:
            raise ValueError("Missing required feature: num_medications")
        except Exception as e:
            client.record_error(
                error=e,
                error_type="validation_error",
                severity="warning",
                context={
                    "encounter_id": "sanitized-enc-002",
                    "missing_features": ["num_medications"],
                },
            )
            print(f"✓ Error recorded: {type(e).__name__}\n")

        # Step 12: Flush
        print("=== Flush Telemetry ===")
        client.shutdown(timeout_millis=5000)
        registry_client.close()
        print("✓ Telemetry sent to collector\n")

        # Step 13: Verification
        print("=== Verify in GCP ===")
        print("Cloud Logging: https://console.cloud.google.com/logs/query")
        print('  Filter: logName="projects/bhsf-ai-team-projects/logs/emms-model-telemetry"')
        print(f'          AND labels."service.name"="{model_config.service_name}"')
        print("\nCloud Trace: https://console.cloud.google.com/traces/list")
        print(f'  Search: service.name="{model_config.service_name}"')
        print("  Look for: predict, data_preprocessing spans")

    except Exception as e:
        print(f"\n✗ Error: {e}")
        client.shutdown()
        registry_client.close()
        raise


if __name__ == "__main__":
    main()
