#!/usr/bin/env python3
"""
Comprehensive parameter sweep validation: Compare Python vs R PRIME model.
Tests all parameter combinations and reports differences for ALL outputs.
Uses the Hugging Face Space (prime_py) code for validation.
"""

import json
import subprocess
import sys
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np
import pandas as pd

# Add src to path - USE PRIME_PY CODE (Hugging Face Space)
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

from primepy.demo import DemoParams, run_demo_and_export


def run_r_model(params_dict: dict, output_dir: Path) -> Dict:
    """Run R PRIME model and return all results."""
    r_script = Path(__file__).parent / "validate_r_python_parity.R"
    base_dir = Path(__file__).parent.parent

    output_dir.mkdir(parents=True, exist_ok=True)

    # Prepare R arguments
    args = [
        "Rscript",
        str(r_script),
        params_dict["country_iso3"],
        str(params_dict["vaccine_efficacy_before"]),
        str(params_dict["vaccine_efficacy_after"]),
        str(params_dict["coverage_vac"]),
        str(params_dict["age_vaccination"]),
        params_dict["dataset_year"],
        "TRUE" if params_dict["use_discounting"] else "FALSE",
        str(params_dict["discount_rate_cost"]),
        str(params_dict["discount_rate_benefit"]),
        params_dict["vaccine"],
        str(output_dir),
        "TRUE",  # analyseCosts
        str(params_dict.get("vaccine_price_total", 26.96)),
        str(params_dict.get("vaccine_delivery_cost_total", 15.0)),
        str(params_dict.get("cancer_treatment_cost", 1815.0)),
    ]

    result = subprocess.run(args, capture_output=True, text=True, cwd=base_dir)

    if result.returncode != 0:
        return {"error": f"R script failed: {result.stderr}"}

    # Check status
    status_file = output_dir / "r_status.txt"
    if not status_file.exists() or "ERROR" in status_file.read_text():
        return {
            "error": f"R model error: {status_file.read_text() if status_file.exists() else 'No status file'}"
        }

    # Read cost-effectiveness results
    ce_file = output_dir / "r_cost_effectiveness.csv"
    if not ce_file.exists():
        return {"error": f"R CE file not found: {ce_file}"}

    r_ce = pd.read_csv(ce_file)

    # Convert to dict with standardized keys
    results = {}
    for _, row in r_ce.iterrows():
        metric = row["metric"]
        value = row["value"]
        results[metric] = value

    return results


def run_python_model(params: DemoParams, output_dir: Path) -> Dict:
    """Run Python model and return all results."""
    base_dir = Path(__file__).parent.parent
    output_dir.mkdir(parents=True, exist_ok=True)

    try:
        output = run_demo_and_export(
            base_dir=base_dir,
            params=params,
            out_dir=output_dir,
            save_intermediates=False,
        )

        if "ce_results" not in output:
            return {"error": "Python model did not return CE results"}

        ce = output["ce_results"]

        # Convert to dict with R-compatible keys
        results = {
            "Vac cohort size": ce.vaccinated_cohort_size,
            "CeCx prevented": ce.cases_prevented,
            "Deaths prevented": ce.deaths_prevented,
            "Life years saved": ce.life_years_saved,
            "Nonfatal DALYs prevented": ce.nonfatal_dalys_prevented,  # Just YLD, not total
            "Vaccine cost": ce.vaccination_costs,
            "Costs saved": ce.treatment_costs_saved,
            "Net cost": ce.net_costs,
            "Cost/death prevented": ce.cost_per_death_prevented,
            "Cost/life year saved": ce.cost_per_life_year_saved,
            "Cost/DALY prevented": ce.cost_per_daly_prevented,
            "GDP/capita": ce.gdp_per_capita,
            "CE at 1xGDP/capita?": 1.0 if ce.ce_at_1xgdp else 0.0,
            "CE at 3xGDP/capita?": 1.0 if ce.ce_at_3xgdp else 0.0,
        }

        return results

    except Exception as e:
        return {"error": f"Python model failed: {str(e)}"}


def compare_results(
    r_results: Dict, py_results: Dict, tolerance: float = 0.01
) -> Tuple[bool, List[Dict]]:
    """
    Compare R and Python results.
    Returns (all_match, differences_list)
    """
    if "error" in r_results:
        return False, [{"metric": "R_ERROR", "error": r_results["error"]}]

    if "error" in py_results:
        return False, [{"metric": "PY_ERROR", "error": py_results["error"]}]

    # Metrics to compare (exclude GDP, Cohort size, Cut-off price which may differ)
    metrics_to_compare = [
        "Vac cohort size",
        "CeCx prevented",
        "Deaths prevented",
        "Life years saved",
        "Nonfatal DALYs prevented",
        "Vaccine cost",
        "Costs saved",
        "Net cost",
        "Cost/death prevented",
        "Cost/life year saved",
        "Cost/DALY prevented",
        "CE at 1xGDP/capita?",
        "CE at 3xGDP/capita?",
    ]

    differences = []
    all_match = True

    # Compare only the metrics we care about
    for key in metrics_to_compare:
        r_val = r_results.get(key)
        py_val = py_results.get(key)

        if r_val is None:
            differences.append(
                {
                    "metric": key,
                    "r_value": "MISSING",
                    "py_value": py_val,
                    "abs_diff": None,
                    "rel_diff": None,
                }
            )
            all_match = False
            continue

        if py_val is None:
            differences.append(
                {
                    "metric": key,
                    "r_value": r_val,
                    "py_value": "MISSING",
                    "abs_diff": None,
                    "rel_diff": None,
                }
            )
            all_match = False
            continue

        # Handle NaN/inf
        if np.isnan(r_val) and np.isnan(py_val):
            continue
        if np.isinf(r_val) and np.isinf(py_val) and np.sign(r_val) == np.sign(py_val):
            continue

        # Calculate differences
        abs_diff = abs(py_val - r_val)

        if abs(r_val) > 1e-10:
            rel_diff = abs_diff / abs(r_val)
        else:
            rel_diff = abs_diff

        if rel_diff > tolerance:
            differences.append(
                {
                    "metric": key,
                    "r_value": r_val,
                    "py_value": py_val,
                    "abs_diff": abs_diff,
                    "rel_diff": rel_diff,
                    "rel_diff_pct": rel_diff * 100,
                }
            )
            all_match = False

    return all_match, differences


def generate_parameter_combinations() -> List[Tuple[str, Dict]]:
    """Generate all parameter combinations to test."""
    base_params = {
        "country_iso3": "TUR",
        "dataset_year": "2020",
        "use_country_sexual_debut": True,
        "use_proxy_fallback": False,
        "use_new_turkish_population": False,  # Use standard data for R comparison
        "current_year": 2025,
        "oldest_year": 2100,
    }

    # Parameter sweep ranges
    param_sweep = {
        "vaccine": ["4vHPV", "9vHPV"],
        "age_vaccination": [9, 10, 12],
        "coverage_vac": [0.5, 0.7, 0.9],
        "vaccine_efficacy_before": [0.8, 0.9, 1.0],
        "vaccine_efficacy_after": [0.0, 0.1],
        # protection_duration_years: Only test None (lifelong) for R vs Python comparison
        # R's RunCountry doesn't support limited protection
        "use_discounting": [True, False],
        "discount_rate_cost": [0.0, 0.03, 0.05],
        "discount_rate_benefit": [0.0, 0.03, 0.05],
        # Cost parameters
        "vaccine_price_total": [13.48, 26.96, 40.0],
        "vaccine_delivery_cost_total": [7.5, 15.0, 30.0],
        "cancer_treatment_cost": [1000.0, 1815.0, 450.0],
    }

    # Generate combinations (with some filtering to reduce total)
    test_cases = []

    # Test each parameter individually (one-at-a-time)
    for param_name, param_values in param_sweep.items():
        for value in param_values:
            # Skip when value is the default (already tested in base case)
            defaults = {
                "vaccine": "4vHPV",
                "age_vaccination": 10,
                "coverage_vac": 0.9,
                "vaccine_efficacy_before": 1.0,
                "vaccine_efficacy_after": 0.0,
                "protection_duration_years": None,
                "use_discounting": True,
                "discount_rate_cost": 0.03,
                "discount_rate_benefit": 0.03,
                "vaccine_price_total": 26.96,
                "vaccine_delivery_cost_total": 15.0,
                "cancer_treatment_cost": 1815.0,
                "next_n_years": 1,
            }

            if param_name in defaults and value == defaults[param_name]:
                continue  # Skip default value, will be tested in base case

            params = base_params.copy()
            params.update(defaults)
            # Change the parameter being tested
            params[param_name] = value

            test_name = f"{param_name}={value}"
            test_cases.append((test_name, params))

    # Add some critical combinations
    critical_combos = [
        (
            "base_case",
            {
                **base_params,
                "vaccine": "4vHPV",
                "age_vaccination": 10,
                "coverage_vac": 0.9,
                "vaccine_efficacy_before": 1.0,
                "vaccine_efficacy_after": 0.0,
                "protection_duration_years": None,
                "use_discounting": True,
                "discount_rate_cost": 0.03,
                "discount_rate_benefit": 0.03,
                "vaccine_price_total": 26.96,
                "vaccine_delivery_cost_total": 15.0,
                "cancer_treatment_cost": 1815.0,
                "next_n_years": 1,
            },
        ),
        # next_n_years_scaling_test REMOVED - not comparable with R (R is single-cohort only)
        (
            "9vHPV_age12_cov70",
            {
                **base_params,
                "vaccine": "9vHPV",
                "age_vaccination": 12,
                "coverage_vac": 0.7,
                "vaccine_efficacy_before": 1.0,
                "vaccine_efficacy_after": 0.0,
                "protection_duration_years": None,
                "use_discounting": True,
                "discount_rate_cost": 0.03,
                "discount_rate_benefit": 0.03,
                "vaccine_price_total": 26.96,
                "vaccine_delivery_cost_total": 15.0,
                "cancer_treatment_cost": 1815.0,
                "next_n_years": 1,
            },
        ),
        # limited_protection_no_discount REMOVED - R doesn't support limited protection
        (
            "high_cost_scenario",
            {
                **base_params,
                "vaccine": "9vHPV",
                "age_vaccination": 10,
                "coverage_vac": 0.9,
                "vaccine_efficacy_before": 1.0,
                "vaccine_efficacy_after": 0.0,
                "protection_duration_years": None,  # Lifelong
                "use_discounting": True,
                "discount_rate_cost": 0.05,
                "discount_rate_benefit": 0.05,
                "vaccine_price_total": 40.0,
                "vaccine_delivery_cost_total": 30.0,
                "cancer_treatment_cost": 2000.0,
                "next_n_years": 1,
            },
        ),
        # low_cost_1dose_scenario REMOVED - R doesn't support limited protection (20yr duration)
    ]

    test_cases.extend(critical_combos)

    # Remove duplicates
    seen = set()
    unique_cases = []
    for name, params in test_cases:
        key = json.dumps(params, sort_keys=True)
        if key not in seen:
            seen.add(key)
            unique_cases.append((name, params))

    return unique_cases


def main():
    """Run comprehensive parameter sweep validation."""
    base_dir = Path(__file__).parent.parent
    validation_output = base_dir / "validation_output" / "parameter_sweep"
    validation_output.mkdir(parents=True, exist_ok=True)

    # Generate test cases
    test_cases = generate_parameter_combinations()

    print(f"{'=' * 80}")
    print("COMPREHENSIVE PARAMETER SWEEP VALIDATION")
    print(f"{'=' * 80}")
    print(f"Total test cases: {len(test_cases)}")
    print(f"Output directory: {validation_output}")
    print(f"{'=' * 80}\n")

    all_results = []
    summary_stats = {
        "total": 0,
        "passed": 0,
        "failed": 0,
        "errors": 0,
    }

    for i, (test_name, params_dict) in enumerate(test_cases, 1):
        print(f"\n[{i}/{len(test_cases)}] Testing: {test_name}")
        print(f"{'-' * 80}")

        test_output = validation_output / test_name.replace("=", "_").replace(" ", "_")
        test_output.mkdir(parents=True, exist_ok=True)

        # Save test parameters
        with open(test_output / "params.json", "w") as f:
            json.dump(params_dict, f, indent=2)

        # Create DemoParams
        params = DemoParams(**params_dict)

        # Run Python model
        print("  Running Python model...", end=" ", flush=True)
        py_results = run_python_model(params, test_output / "python")
        if "error" in py_results:
            print("❌ ERROR")
            print(f"    {py_results['error']}")
            summary_stats["errors"] += 1
            all_results.append(
                {
                    "test": test_name,
                    "status": "error",
                    "error_source": "python",
                    "error": py_results["error"],
                    "params": params_dict,
                }
            )
            continue
        print("✓")

        # Run R model
        print("  Running R model...", end=" ", flush=True)
        r_results = run_r_model(params_dict, test_output / "r")
        if "error" in r_results:
            print("❌ ERROR")
            print(f"    {r_results['error']}")
            summary_stats["errors"] += 1
            all_results.append(
                {
                    "test": test_name,
                    "status": "error",
                    "error_source": "r",
                    "error": r_results["error"],
                    "params": params_dict,
                }
            )
            continue
        print("✓")

        # Compare results
        print("  Comparing results...", end=" ", flush=True)
        is_match, differences = compare_results(r_results, py_results, tolerance=0.01)

        summary_stats["total"] += 1

        if is_match:
            print("✅ PASS")
            summary_stats["passed"] += 1
            all_results.append(
                {
                    "test": test_name,
                    "status": "pass",
                    "params": params_dict,
                    "r_results": r_results,
                    "py_results": py_results,
                }
            )
        else:
            print(f"❌ FAIL ({len(differences)} differences)")
            summary_stats["failed"] += 1

            # Show top differences
            sorted_diffs = sorted(
                differences,
                key=lambda x: x.get("rel_diff", 0)
                if x.get("rel_diff") is not None
                else 0,
                reverse=True,
            )
            for diff in sorted_diffs[:3]:
                if "rel_diff_pct" in diff:
                    print(
                        f"    • {diff['metric']}: {diff['rel_diff_pct']:.2f}% diff (R={diff['r_value']:.4f}, Py={diff['py_value']:.4f})"
                    )
                else:
                    print(f"    • {diff['metric']}: {diff}")

            all_results.append(
                {
                    "test": test_name,
                    "status": "fail",
                    "params": params_dict,
                    "r_results": r_results,
                    "py_results": py_results,
                    "differences": differences,
                }
            )

    # Save detailed results
    results_file = validation_output / "validation_results.json"
    with open(results_file, "w") as f:
        json.dump(all_results, f, indent=2, default=str)

    # Create summary DataFrame
    summary_rows = []
    for result in all_results:
        row = {
            "test": result["test"],
            "status": result["status"],
        }

        if result["status"] == "fail" and "differences" in result:
            # Add metrics with largest differences
            diffs = result["differences"]
            sorted_diffs = sorted(
                diffs,
                key=lambda x: x.get("rel_diff", 0)
                if x.get("rel_diff") is not None
                else 0,
                reverse=True,
            )
            for i, diff in enumerate(sorted_diffs[:3], 1):
                row[f"top_diff_{i}_metric"] = diff.get("metric", "")
                row[f"top_diff_{i}_pct"] = diff.get("rel_diff_pct", "")

        summary_rows.append(row)

    summary_df = pd.DataFrame(summary_rows)
    summary_df.to_csv(validation_output / "summary.csv", index=False)

    # Print summary
    print(f"\n{'=' * 80}")
    print("VALIDATION SUMMARY")
    print(f"{'=' * 80}")
    print(f"Total tests:    {summary_stats['total']}")
    print(
        f"Passed:         {summary_stats['passed']} ({summary_stats['passed'] / max(summary_stats['total'], 1) * 100:.1f}%)"
    )
    print(
        f"Failed:         {summary_stats['failed']} ({summary_stats['failed'] / max(summary_stats['total'], 1) * 100:.1f}%)"
    )
    print(f"Errors:         {summary_stats['errors']}")
    print(f"{'=' * 80}")

    if summary_stats["failed"] > 0:
        print(
            f"\n❌ VALIDATION FAILED: {summary_stats['failed']} test(s) with mismatches"
        )
        print("\nFailed tests:")
        for result in all_results:
            if result["status"] == "fail":
                print(f"  • {result['test']}")
        print(f"\nDetailed results saved to: {results_file}")
        return 1
    elif summary_stats["errors"] > 0:
        print(
            f"\n⚠️  VALIDATION INCOMPLETE: {summary_stats['errors']} test(s) with errors"
        )
        return 1
    else:
        print("\n✅ ALL TESTS PASSED")
        return 0


if __name__ == "__main__":
    sys.exit(main())