# Comprehensive Likelihood Test Runner
# ====================================
# Tests all likelihood types individually following the pattern 
# of test_votes_only.py and test_surveys_simple.py

import sys
import warnings
warnings.filterwarnings("ignore")

# Import all individual test functions
from test_votes_only import test_votes_only
from test_surveys_simple import test_surveys_only  
from test_text_embeddings import test_text_embeddings
from test_ordinal import test_ordinal
from test_poisson import test_poisson


def run_all_likelihood_tests():
    """Run tests for all likelihood types."""
    print("="*80)
    print("COMPREHENSIVE LIKELIHOOD TESTS")
    print("Testing all likelihood types individually")
    print("="*80)
    
    results = {}
    
    # Test 1: Bernoulli (Votes)
    print("\n" + "="*60)
    print("TEST 1: BERNOULLI LIKELIHOOD (VOTES)")
    print("="*60)
    try:
        correlation, r_squared = test_votes_only()
        results["Bernoulli (Votes)"] = {
            "status": "SUCCESS", 
            "correlation": correlation, 
            "r_squared": r_squared
        }
    except Exception as e:
        print(f"ERROR in Bernoulli test: {e}")
        results["Bernoulli (Votes)"] = {"status": "ERROR", "error": str(e)}
    
    # Test 2: Categorical Multi (Surveys)  
    print("\n" + "="*60)
    print("TEST 2: CATEGORICAL MULTI LIKELIHOOD (SURVEYS)")
    print("="*60)
    try:
        correlation, r_squared = test_surveys_only()
        results["Categorical Multi (Surveys)"] = {
            "status": "SUCCESS", 
            "correlation": correlation, 
            "r_squared": r_squared
        }
    except Exception as e:
        print(f"ERROR in Categorical Multi test: {e}")
        results["Categorical Multi (Surveys)"] = {"status": "ERROR", "error": str(e)}
    
    # Test 3: Gaussian (Text Embeddings)
    print("\n" + "="*60)
    print("TEST 3: GAUSSIAN LIKELIHOOD (TEXT EMBEDDINGS)")
    print("="*60)
    try:
        correlation, r_squared = test_text_embeddings()
        results["Gaussian (Text Embeddings)"] = {
            "status": "SUCCESS", 
            "correlation": correlation, 
            "r_squared": r_squared
        }
    except Exception as e:
        print(f"ERROR in Gaussian test: {e}")
        results["Gaussian (Text Embeddings)"] = {"status": "ERROR", "error": str(e)}
    
    # Test 4: Ordinal
    print("\n" + "="*60)
    print("TEST 4: ORDINAL LIKELIHOOD")
    print("="*60)
    try:
        correlation, r_squared = test_ordinal()
        results["Ordinal"] = {
            "status": "SUCCESS", 
            "correlation": correlation, 
            "r_squared": r_squared
        }
    except Exception as e:
        print(f"ERROR in Ordinal test: {e}")
        results["Ordinal"] = {"status": "ERROR", "error": str(e)}
    
    # Test 5: Poisson (Count Data)
    print("\n" + "="*60)
    print("TEST 5: POISSON LIKELIHOOD (COUNT DATA)")
    print("="*60)
    try:
        correlation, r_squared = test_poisson()
        results["Poisson (Count Data)"] = {
            "status": "SUCCESS", 
            "correlation": correlation, 
            "r_squared": r_squared
        }
    except Exception as e:
        print(f"ERROR in Poisson test: {e}")
        results["Poisson (Count Data)"] = {"status": "ERROR", "error": str(e)}
    
    # Summary
    print("\n" + "="*80)
    print("LIKELIHOOD TEST SUMMARY")
    print("="*80)
    print(f"{'Likelihood Type':<30} {'Status':<10} {'Correlation':<12} {'R²':<10}")
    print("-"*80)
    
    for test_name, result in results.items():
        if result["status"] == "SUCCESS":
            correlation = result["correlation"]
            r_squared = result["r_squared"]
            print(f"{test_name:<30} {'SUCCESS':<10} {correlation:<12.3f} {r_squared:<10.3f}")
        else:
            print(f"{test_name:<30} {'ERROR':<10} {'N/A':<12} {'N/A':<10}")
    
    print("-"*80)
    
    # Count successful tests
    successful_tests = sum(1 for result in results.values() if result["status"] == "SUCCESS")
    total_tests = len(results)
    
    print(f"\nSUCCESSFUL TESTS: {successful_tests}/{total_tests}")
    
    if successful_tests == total_tests:
        print("🎉 ALL LIKELIHOOD TESTS PASSED!")
    else:
        print(f"⚠️  {total_tests - successful_tests} test(s) failed.")
        
        # Print error details
        for test_name, result in results.items():
            if result["status"] == "ERROR":
                print(f"\n{test_name} Error: {result['error']}")
    
    return results


if __name__ == "__main__":
    results = run_all_likelihood_tests()
    print(f"\nAll likelihood tests completed!")
