#!/usr/bin/env python3
"""
Comprehensive Multi-ECG Testing and Evaluation
Tests multiple ECG samples and analyzes results systematically
"""

import requests
import numpy as np
import pandas as pd
import json
import time
import os
from typing import Dict, Any, List
from datetime import datetime

# Configuration
API_BASE_URL = "https://mystic-cbk-ecg-fm-api.hf.space"
ECG_DIR = "../ecg_uploads_greenwich/"
INDEX_FILE = "../Greenwichschooldata.csv"

def load_ecg_data(csv_file: str) -> List[List[float]]:
    """Load ECG data from CSV file"""
    try:
        df = pd.read_csv(csv_file)
        ecg_data = []
        for lead in df.columns:
            ecg_data.append(df[lead].astype(float).tolist())
        return ecg_data
    except Exception as e:
        print(f"❌ Error loading ECG data: {e}")
        return None

def test_individual_endpoints(ecg_data: List[List[float]], patient_info: Dict[str, Any]) -> Dict[str, Any]:
    """Test all individual endpoints with comprehensive analysis"""
    results = {}
    
    print(f"🧪 Testing individual endpoints for {patient_info.get('Patient Name', 'Unknown')}")
    
    # Test 1: Extract Features (Physiological measurements)
    print("1️⃣ Testing /extract_features endpoint...")
    try:
        payload = {"signal": ecg_data, "fs": 500}
        response = requests.post(f"{API_BASE_URL}/extract_features", json=payload, timeout=60)
        
        if response.status_code == 200:
            result = response.json()
            physio = result.get('physiological_parameters', {})
            
            print(f"   ✅ Features extracted successfully")
            print(f"   📊 Feature count: {result.get('features', {}).get('count', 'Unknown')}")
            print(f"   📊 Feature dimension: {result.get('features', {}).get('dimension', 'Unknown')}")
            print(f"   💓 Heart Rate: {physio.get('heart_rate')} BPM")
            print(f"   📏 QRS Duration: {physio.get('qrs_duration')} ms")
            print(f"   ⏱️  QT Interval: {physio.get('qt_interval')} ms")
            print(f"   🔗 PR Interval: {physio.get('pr_interval')} ms")
            print(f"   🧭 QRS Axis: {physio.get('qrs_axis')}°")
            
            results['extract_features'] = {"status": "success", "data": result}
        else:
            print(f"   ❌ Failed: {response.status_code}")
            results['extract_features'] = {"status": "error", "error": response.text}
    except Exception as e:
        print(f"   ❌ Error: {e}")
        results['extract_features'] = {"status": "error", "error": str(e)}
    
    # Test 2: Assess Quality
    print("2️⃣ Testing /assess_quality endpoint...")
    try:
        payload = {"signal": ecg_data, "fs": 500}
        response = requests.post(f"{API_BASE_URL}/assess_quality", json=payload, timeout=60)
        
        if response.status_code == 200:
            result = response.json()
            print(f"   ✅ Quality assessment completed")
            print(f"   🔍 Overall Quality: {result.get('quality', 'Unknown')}")
            metrics = result.get('metrics', {})
            print(f"   📊 SNR: {metrics.get('signal_to_noise_ratio', 'Unknown')}")
            print(f"   📊 Baseline Wander: {metrics.get('baseline_wander', 'Unknown')}")
            print(f"   📊 Standard Deviation: {metrics.get('standard_deviation', 'Unknown')}")
            
            results['assess_quality'] = {"status": "success", "data": result}
        else:
            print(f"   ❌ Failed: {response.status_code}")
            results['assess_quality'] = {"status": "error", "error": response.text}
    except Exception as e:
        print(f"   ❌ Error: {e}")
        results['assess_quality'] = {"status": "error", "error": str(e)}
    
    # Test 3: Predict endpoint
    print("3️⃣ Testing /predict endpoint...")
    try:
        payload = {"signal": ecg_data, "fs": 500}
        response = requests.post(f"{API_BASE_URL}/predict", json=payload, timeout=60)
        
        if response.status_code == 200:
            result = response.json()
            print(f"   ✅ Prediction completed")
            print(f"   🧬 Model Type: {result.get('model_type', 'Unknown')}")
            print(f"   📊 Confidence: {result.get('confidence', 'Unknown')}")
            
            results['predict'] = {"status": "success", "data": result}
        else:
            print(f"   ❌ Failed: {response.status_code}")
            results['assess_quality'] = {"status": "error", "error": response.text}
    except Exception as e:
        print(f"   ❌ Error: {e}")
        results['predict'] = {"status": "error", "error": str(e)}
    
    return results

def analyze_physiological_consistency(all_results: Dict[str, Any]) -> Dict[str, Any]:
    """Analyze consistency of physiological measurements across samples"""
    print(f"\n📊 PHYSIOLOGICAL MEASUREMENT CONSISTENCY ANALYSIS")
    print(f"=" * 70)
    
    # Extract all heart rate values
    hr_values = []
    qrs_values = []
    qt_values = []
    pr_values = []
    axis_values = []
    
    for patient_id, result in all_results.items():
        if 'endpoint_tests' in result and 'extract_features' in result['endpoint_tests']:
            physio = result['endpoint_tests']['extract_features'].get('data', {}).get('physiological_parameters', {})
            
            if physio.get('heart_rate') is not None:
                hr_values.append(physio['heart_rate'])
            if physio.get('qrs_duration') is not None:
                qrs_values.append(physio['qrs_duration'])
            if physio.get('qt_interval') is not None:
                qt_values.append(physio['qt_interval'])
            if physio.get('pr_interval') is not None:
                pr_values.append(physio['pr_interval'])
            if physio.get('qrs_axis') is not None:
                axis_values.append(physio['qrs_axis'])
    
    analysis = {}
    
    # Heart Rate Analysis
    if len(hr_values) > 0:
        hr_std = np.std(hr_values)
        hr_range = np.ptp(hr_values)
        print(f"💓 Heart Rate Analysis:")
        print(f"   Values: {hr_values}")
        print(f"   Mean: {np.mean(hr_values):.1f} BPM")
        print(f"   Standard Deviation: {hr_std:.2f} BPM")
        print(f"   Range: {hr_range:.1f} BPM")
        print(f"   Consistency: {'✅ Excellent' if hr_std < 2.0 else '⚠️ Good' if hr_std < 5.0 else '❌ Poor'}")
        
        analysis['heart_rate'] = {
            'values': hr_values,
            'mean': np.mean(hr_values),
            'std': hr_std,
            'range': hr_range,
            'consistency': 'Excellent' if hr_std < 2.0 else 'Good' if hr_std < 5.0 else 'Poor'
        }
    
    # QRS Duration Analysis
    if len(qrs_values) > 0:
        qrs_std = np.std(qrs_values)
        qrs_range = np.ptp(qrs_values)
        print(f"\n📏 QRS Duration Analysis:")
        print(f"   Values: {qrs_values}")
        print(f"   Mean: {np.mean(qrs_values):.1f} ms")
        print(f"   Standard Deviation: {qrs_std:.2f} ms")
        print(f"   Range: {qrs_range:.1f} ms")
        print(f"   Consistency: {'✅ Excellent' if qrs_std < 5.0 else '⚠️ Good' if qrs_std < 10.0 else '❌ Poor'}")
        
        analysis['qrs_duration'] = {
            'values': qrs_values,
            'mean': np.mean(qrs_values),
            'std': qrs_std,
            'range': qrs_range,
            'consistency': 'Excellent' if qrs_std < 5.0 else 'Good' if qrs_std < 10.0 else 'Poor'
        }
    
    # Overall Assessment
    print(f"\n🎯 OVERALL PHYSIOLOGICAL ASSESSMENT:")
    working_params = len([k for k in analysis.keys() if analysis[k]['consistency'] in ['Excellent', 'Good']])
    total_params = len(analysis)
    
    print(f"   Working Parameters: {working_params}/{total_params}")
    print(f"   Success Rate: {(working_params/total_params)*100:.1f}%")
    
    if working_params == total_params:
        print(f"   🎉 All physiological parameters working consistently!")
    elif working_params > total_params // 2:
        print(f"   ⚠️  Most parameters working - some inconsistencies")
    else:
        print(f"   ❌ Many parameters showing inconsistencies")
    
    return analysis

def main():
    """Main test function"""
    print("🚀 COMPREHENSIVE MULTI-ECG TESTING AND EVALUATION")
    print("=" * 80)
    print(f"🌐 API URL: {API_BASE_URL}")
    print(f"📁 ECG Directory: {ECG_DIR}")
    print(f"📋 Index File: {INDEX_FILE}")
    print()
    
    # Check if files exist
    if not os.path.exists(INDEX_FILE):
        print(f"❌ Index file not found: {INDEX_FILE}")
        return
    
    if not os.path.exists(ECG_DIR):
        print(f"❌ ECG directory not found: {ECG_DIR}")
        return
    
    # Load index file
    try:
        print("📁 Loading patient index file...")
        index_df = pd.read_csv(INDEX_FILE)
        print(f"✅ Loaded {len(index_df)} patient records")
    except Exception as e:
        print(f"❌ Error loading index file: {e}")
        return
    
    # Select multiple ECG files for testing
    test_files = [
        "ecg_98408931-6f8e-47cc-954a-ba0c058a0f3d.csv",  # Bharathi M K Teacher, 31, F
        "ecg_fc6d2ecb-7eb3-4eec-9281-17c24b7902b5.csv",  # Sayida thasmiya Bhanu Teacher, 29, F
        "ecg_022a3f3a-7060-4ff8-b716-b75d8e0637c5.csv",  # Afzal, 46, M
        # Add more files if available
    ]
    
    print(f"\n🚀 Testing with {len(test_files)} ECG samples...")
    print("=" * 80)
    
    all_results = {}
    
    for i, ecg_file in enumerate(test_files, 1):
        try:
            print(f"\n📊 Processing {i}/{len(test_files)}: {ecg_file}")
            
            # Find patient info in index
            patient_row = index_df[index_df['ECG File Path'].str.contains(ecg_file, na=False)]
            if len(patient_row) == 0:
                print(f"   ⚠️  Patient info not found for {ecg_file}")
                continue
            
            patient_info = patient_row.iloc[0]
            print(f"   👤 Patient: {patient_info['Patient Name']} ({patient_info['Age']} {patient_info['Gender']})")
            
            # Check if ECG file exists
            ecg_path = os.path.join(ECG_DIR, ecg_file)
            if not os.path.exists(ecg_path):
                print(f"   ❌ ECG file not found: {ecg_path}")
                continue
            
            # Load ECG data
            ecg_data = load_ecg_data(ecg_path)
            if ecg_data is None:
                print(f"   ❌ Failed to load ECG data")
                continue
            
            # Test individual endpoints
            endpoint_results = test_individual_endpoints(ecg_data, patient_info)
            
            # Store results
            all_results[ecg_file] = {
                "patient_info": patient_info.to_dict(),
                "endpoint_tests": endpoint_results
            }
            
            print(f"   ✅ Completed analysis for {ecg_file}")
            
        except Exception as e:
            print(f"   ❌ Error processing {ecg_file}: {e}")
            all_results[ecg_file] = {"error": str(e)}
    
    # Analyze physiological consistency
    if len(all_results) > 0:
        physiological_analysis = analyze_physiological_consistency(all_results)
        
        # Summary report
        print(f"\n📊 COMPREHENSIVE TEST SUMMARY")
        print(f"=" * 80)
        
        successful_tests = 0
        total_tests = len(test_files)
        
        for ecg_file, result in all_results.items():
            if "error" not in result:
                endpoint_status = result.get("endpoint_tests", {})
                working_endpoints = sum(1 for ep in endpoint_status.values() if ep.get("status") == "success")
                total_endpoints = len(endpoint_status)
                
                if working_endpoints == total_endpoints:
                    successful_tests += 1
                    print(f"✅ {ecg_file}: All endpoints working")
                else:
                    print(f"⚠️  {ecg_file}: {working_endpoints}/{total_endpoints} endpoints working")
            else:
                print(f"❌ {ecg_file}: {result['error']}")
        
        print(f"\n🎯 OVERALL RESULTS:")
        print(f"   Successful ECG Analysis: {successful_tests}/{total_tests}")
        print(f"   Success Rate: {(successful_tests/total_tests)*100:.1f}%")
        
        # Save detailed results
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        results_file = f"comprehensive_ecg_test_results_{timestamp}.json"
        
        try:
            with open(results_file, 'w') as f:
                json.dump(all_results, f, indent=2, default=str)
            print(f"\n💾 Detailed results saved to: {results_file}")
        except Exception as e:
            print(f"\n⚠️  Could not save results: {e}")
    
    print(f"\n🎉 Comprehensive ECG testing completed!")
    print(f"💡 Check the results above to evaluate system performance")

if __name__ == "__main__":
    main()