use anyhow::Result;
use pmcore::bestdose::{BestDosePosterior, DoseRange, Target};
use pmcore::prelude::*;
use pmcore::structs::theta::Theta;
use pmcore::structs::weights::Weights;

/// Test that infusions are properly included in the dose optimization mask
/// This test verifies that infusions with amount=0 are treated as optimizable doses
#[test]
fn test_infusion_mask_inclusion() -> Result<()> {
    // Create a simple one-compartment model
    let eq = equation::ODE::new(
        |x, p, _t, dx, b, _rateiv, _cov| {
            fetch_params!(p, ke, _v);
            dx[0] = -ke * x[0] + b[0];
        },
        |_p, _, _| lag! {},
        |_p, _, _| fa! {},
        |_p, _t, _cov, _x| {},
        |x, p, _t, _cov, y| {
            fetch_params!(p, _ke, v);
            y[0] = x[0] / v;
        },
    );

    let params = Parameters::new().add("ke", 0.1, 0.5).add("v", 40.0, 60.0);

    let ems = AssayErrorModels::new().add(
        0,
        AssayErrorModel::additive(ErrorPoly::new(0.0, 0.20, 0.0, 0.0), 0.0),
    )?;

    let mut settings = Settings::builder()
        .set_algorithm(Algorithm::NPAG)
        .set_parameters(params)
        .set_error_models(ems.clone())
        .build();

    settings.disable_output();
    settings.set_cycles(0);

    // Create a target subject with an optimizable infusion
    // Use reasonable target concentrations that match typical PK behavior
    let target = Subject::builder("test_patient")
        .infusion(0.0, 0.0, 0, 1.0) // Optimizable 1-hour infusion
        .observation(2.0, 2.0, 0) // Target concentration at 2h
        .observation(4.0, 1.5, 0) // Target concentration at 4h
        .build();

    // Create a prior with reasonable PK parameters
    let prior_theta = {
        let mat = faer::Mat::from_fn(1, 2, |_r, c| match c {
            0 => 0.3,  // ke
            1 => 50.0, // v
            _ => 0.0,
        });
        Theta::from_parts(mat, settings.parameters().clone())?
    };
    let prior_weights = Weights::uniform(1);

    // Create BestDose posterior
    let posterior = BestDosePosterior::compute(
        &prior_theta,
        &prior_weights,
        None,
        eq.clone(),
        settings.clone(),
    )?;

    // Count optimizable doses in the target
    let mut optimizable_infusions = 0;
    for occasion in target.occasions() {
        for event in occasion.events() {
            if let Event::Infusion(inf) = event {
                if inf.amount() == 0.0 {
                    optimizable_infusions += 1;
                }
            }
        }
    }

    assert_eq!(
        optimizable_infusions, 1,
        "Should have 1 optimizable infusion"
    );

    // Run optimization - it should not panic and should handle infusion
    let result = posterior.optimize(
        target.clone(),
        None,
        DoseRange::new(10.0, 300.0),
        0.5,
        Target::Concentration,
    );

    // The optimization should succeed
    assert!(
        result.is_ok(),
        "Optimization should succeed with infusions: {:?}",
        result.err()
    );

    let result = result?;

    // We should get back 1 optimized dose (the infusion placeholder)
    assert_eq!(
        result.doses().len(),
        1,
        "Should have 1 optimized dose (the infusion)"
    );

    let optinf = result.doses();

    // The optimized dose should be reasonable (not NaN, not infinite)
    assert!(
        optinf[0].is_finite(),
        "Optimized dose should be finite, got {}",
        optinf[0]
    );

    Ok(())
}

/// Test that fixed infusions are preserved during optimization
#[test]
fn test_fixed_infusion_preservation() -> Result<()> {
    let eq = equation::ODE::new(
        |x, p, _t, dx, b, _rateiv, _cov| {
            fetch_params!(p, ke, _v);
            dx[0] = -ke * x[0] + b[0];
        },
        |_p, _, _| lag! {},
        |_p, _, _| fa! {},
        |_p, _t, _cov, _x| {},
        |x, p, _t, _cov, y| {
            fetch_params!(p, _ke, v);
            y[0] = x[0] / v;
        },
    );

    let params = Parameters::new()
        .add("ke", 0.001, 3.0)
        .add("v", 25.0, 250.0);

    let ems = AssayErrorModels::new().add(
        0,
        AssayErrorModel::additive(ErrorPoly::new(0.0, 0.20, 0.0, 0.0), 0.0),
    )?;

    let mut settings = Settings::builder()
        .set_algorithm(Algorithm::NPAG)
        .set_parameters(params)
        .set_error_models(ems.clone())
        .build();

    settings.disable_output();
    settings.set_cycles(0);

    // Create past data with a fixed infusion
    let past = Subject::builder("test_patient")
        .infusion(0.0, 200.0, 0, 1.0) // Fixed past infusion
        .observation(2.0, 3.5, 0)
        .build();

    // Create target with a future optimizable dose
    let target = Subject::builder("test_patient")
        .bolus(0.0, 0.0, 0) // Future dose to optimize
        .observation(2.0, 5.0, 0)
        .build();

    let prior_theta = {
        let mat = faer::Mat::from_fn(1, 2, |_r, c| match c {
            0 => 0.3,
            1 => 50.0,
            _ => 0.0,
        });
        Theta::from_parts(mat, settings.parameters().clone())?
    };
    let prior_weights = Weights::uniform(1);

    // Use current_time to separate past and future
    let posterior = BestDosePosterior::compute(
        &prior_theta,
        &prior_weights,
        Some(past),
        eq.clone(),
        settings.clone(),
    )?;

    let result = posterior.optimize(
        target,
        Some(0.0), // No gap after past (past ends at t=2.0)
        DoseRange::new(0.0, 500.0),
        0.5,
        Target::Concentration,
    )?;

    // With time_offset, past doses are concatenated with future target.
    // Result should have 2 doses: fixed past infusion + optimized future bolus.
    let doses = result.doses();
    eprintln!("Optimized doses: {:?}", doses);
    assert_eq!(
        doses.len(),
        2,
        "Should have 2 doses (past infusion + future bolus)"
    );
    assert!(
        (doses[0] - 200.0).abs() < 1e-6,
        "Past infusion should remain fixed at 200.0, got {}",
        doses[0]
    );
    assert!(doses[1] > 0.0, "Future bolus dose should be optimized");

    Ok(())
}

/// Test that dose count validation works
#[test]
fn test_dose_count_validation() -> Result<()> {
    let eq = equation::ODE::new(
        |x, p, _t, dx, b, _rateiv, _cov| {
            fetch_params!(p, ke, _v);
            dx[0] = -ke * x[0] + b[0];
        },
        |_p, _, _| lag! {},
        |_p, _, _| fa! {},
        |_p, _t, _cov, _x| {},
        |x, p, _t, _cov, y| {
            fetch_params!(p, _ke, v);
            y[0] = x[0] / v;
        },
    );

    let params = Parameters::new().add("ke", 0.1, 0.5).add("v", 40.0, 60.0);
    let ems = AssayErrorModels::new().add(
        0,
        AssayErrorModel::additive(ErrorPoly::new(0.0, 0.20, 0.0, 0.0), 0.0),
    )?;

    let mut settings = Settings::builder()
        .set_algorithm(Algorithm::NPAG)
        .set_parameters(params)
        .set_error_models(ems.clone())
        .build();
    settings.disable_output();
    settings.set_cycles(0);

    // Create target with 2 optimizable doses
    let target = Subject::builder("test_patient")
        .bolus(0.0, 0.0, 0)
        .bolus(6.0, 0.0, 0)
        .observation(2.0, 5.0, 0)
        .observation(8.0, 3.0, 0)
        .build();

    let prior_theta = {
        let mat = faer::Mat::from_fn(1, 2, |_r, c| match c {
            0 => 0.3,
            1 => 50.0,
            _ => 0.0,
        });
        Theta::from_parts(mat, settings.parameters().clone())?
    };
    let prior_weights = Weights::uniform(1);

    let posterior = BestDosePosterior::compute(&prior_theta, &prior_weights, None, eq, settings)?;

    // Optimize with the correct target (2 optimizable doses, 2 observations) - should succeed
    let result = posterior.optimize(
        target,
        None,
        DoseRange::new(10.0, 300.0),
        0.5,
        Target::Concentration,
    );
    assert!(
        result.is_ok(),
        "Should succeed with correct target: {:?}",
        result.err()
    );
    let result = result?;
    assert_eq!(result.doses().len(), 2, "Should have 2 optimized doses");

    Ok(())
}

/// Test that empty observations are caught
#[test]
fn test_empty_observations_validation() -> Result<()> {
    let eq = equation::ODE::new(
        |x, p, _t, dx, b, _rateiv, _cov| {
            fetch_params!(p, ke, _v);
            dx[0] = -ke * x[0] + b[0];
        },
        |_p, _, _| lag! {},
        |_p, _, _| fa! {},
        |_p, _t, _cov, _x| {},
        |x, p, _t, _cov, y| {
            fetch_params!(p, _ke, v);
            y[0] = x[0] / v;
        },
    );

    let params = Parameters::new().add("ke", 0.1, 0.5).add("v", 40.0, 60.0);
    let ems = AssayErrorModels::new().add(
        0,
        AssayErrorModel::additive(ErrorPoly::new(0.0, 0.20, 0.0, 0.0), 0.0),
    )?;

    let mut settings = Settings::builder()
        .set_algorithm(Algorithm::NPAG)
        .set_parameters(params)
        .set_error_models(ems.clone())
        .build();
    settings.disable_output();
    settings.set_cycles(0);

    // Create target with doses but NO observations
    let target = Subject::builder("test_patient").bolus(0.0, 0.0, 0).build(); // No observations!

    let prior_theta = {
        let mat = faer::Mat::from_fn(1, 2, |_r, c| match c {
            0 => 0.3,
            1 => 50.0,
            _ => 0.0,
        });
        Theta::from_parts(mat, settings.parameters().clone())?
    };
    let prior_weights = Weights::uniform(1);

    let posterior = BestDosePosterior::compute(&prior_theta, &prior_weights, None, eq, settings)?;

    // Try to optimize - should fail with no observations
    let result = posterior.optimize(
        target,
        None,
        DoseRange::new(10.0, 300.0),
        0.5,
        Target::Concentration,
    );
    assert!(result.is_err(), "Should fail with no observations");
    assert!(result.unwrap_err().to_string().contains("no observations"));

    Ok(())
}

/// Test basic AUC mode with bolus (simpler test)
#[test]
fn test_basic_auc_mode() -> Result<()> {
    let eq = equation::ODE::new(
        |x, p, _t, dx, b, _rateiv, _cov| {
            fetch_params!(p, ke, _v);
            dx[0] = -ke * x[0] + b[0];
        },
        |_p, _, _| lag! {},
        |_p, _, _| fa! {},
        |_p, _t, _cov, _x| {},
        |x, p, _t, _cov, y| {
            fetch_params!(p, _ke, v);
            y[0] = x[0] / v;
        },
    );

    let params = Parameters::new().add("ke", 0.1, 0.5).add("v", 40.0, 60.0);

    let ems = AssayErrorModels::new().add(
        0,
        AssayErrorModel::additive(ErrorPoly::new(0.0, 0.20, 0.0, 0.0), 0.0),
    )?;

    let mut settings = Settings::builder()
        .set_algorithm(Algorithm::NPAG)
        .set_parameters(params)
        .set_error_models(ems.clone())
        .build();

    settings.disable_output();
    settings.set_idelta(30.0);
    settings.set_cycles(0);

    let target = Subject::builder("test_patient")
        .bolus(0.0, 0.0, 0) // Optimizable bolus
        .observation(6.0, 50.0, 0) // Target AUC at 6h
        .build();

    let prior_theta = {
        let mat = faer::Mat::from_fn(1, 2, |_r, c| match c {
            0 => 0.3,
            1 => 50.0,
            _ => 0.0,
        });
        Theta::from_parts(mat, settings.parameters().clone())?
    };
    let prior_weights = Weights::uniform(1);

    let posterior = BestDosePosterior::compute(&prior_theta, &prior_weights, None, eq, settings)?;

    let result = posterior.optimize(
        target,
        None,
        DoseRange::new(100.0, 2000.0),
        0.8,
        Target::AUCFromZero,
    );

    assert!(
        result.is_ok(),
        "AUC optimization should succeed: {:?}",
        result.err()
    );

    let result = result?;
    let doses = result.doses();
    assert_eq!(doses.len(), 1);

    assert!(result.auc_predictions().is_some());

    let auc_preds = result.auc_predictions().unwrap();
    eprintln!("Basic AUC test - AUC predictions: {:?}", auc_preds);
    assert_eq!(auc_preds.len(), 1);

    let (_time, auc) = auc_preds[0];
    assert!(
        auc.is_finite() && auc > 0.0,
        "AUC should be positive and finite, got {}",
        auc
    );

    Ok(())
}

/// Test that infusions work correctly in AUC mode
#[test]
fn test_infusion_auc_mode() -> Result<()> {
    let eq = equation::ODE::new(
        |x, p, _t, dx, b, rateiv, _cov| {
            fetch_params!(p, ke, _v);
            dx[0] = -ke * x[0] + b[0] + rateiv[0]; // Include infusion rate!
        },
        |_p, _, _| lag! {},
        |_p, _, _| fa! {},
        |_p, _t, _cov, _x| {},
        |x, p, _t, _cov, y| {
            fetch_params!(p, _ke, v);
            y[0] = x[0] / v;
        },
    );

    let params = Parameters::new().add("ke", 0.1, 0.5).add("v", 40.0, 60.0);

    let ems = AssayErrorModels::new().add(
        0,
        AssayErrorModel::additive(ErrorPoly::new(0.0, 0.20, 0.0, 0.0), 0.0),
    )?;

    let mut settings = Settings::builder()
        .set_algorithm(Algorithm::NPAG)
        .set_parameters(params)
        .set_error_models(ems.clone())
        .build();

    settings.disable_output();
    settings.set_idelta(30.0); // 30-minute intervals for AUC calculation
    settings.set_cycles(0);

    // Create a target with an optimizable infusion and AUC targets
    let target = Subject::builder("test_patient")
        .infusion(0.0, 0.0, 0, 2.0) // Optimizable 2-hour infusion
        .observation(6.0, 50.0, 0) // Target AUC at 6h
        .observation(12.0, 80.0, 0) // Target AUC at 12h
        .build();

    let prior_theta = {
        let mat = faer::Mat::from_fn(1, 2, |_r, c| match c {
            0 => 0.3,
            1 => 50.0,
            _ => 0.0,
        });
        Theta::from_parts(mat, settings.parameters().clone())?
    };
    let prior_weights = Weights::uniform(1);

    // Create BestDose posterior and optimize in AUC mode
    let posterior = BestDosePosterior::compute(&prior_theta, &prior_weights, None, eq, settings)?;

    // Run optimization
    let result = posterior.optimize(
        target,
        None,
        DoseRange::new(100.0, 2000.0),
        0.8,                 // Higher bias weight typically works better for AUC targets
        Target::AUCFromZero, // AUC mode!
    );

    assert!(
        result.is_ok(),
        "AUC optimization with infusion should succeed: {:?}",
        result.err()
    );

    let result = result?;
    let doses = result.doses();

    eprintln!("Optimized dose: {:?}", doses);

    // Should have 1 optimized dose (the infusion)
    assert_eq!(doses.len(), 1, "Should have 1 optimized dose");

    // Should have AUC predictions
    assert!(
        result.auc_predictions().is_some(),
        "Should have AUC predictions"
    );

    let auc_preds = result.auc_predictions().unwrap();
    eprintln!("AUC predictions: {:?}", auc_preds);
    assert_eq!(auc_preds.len(), 2, "Should have 2 AUC predictions");

    // AUC values should be reasonable (finite and non-negative)
    // Note: AUC could be very small but shouldn't be exactly 0 if dose is non-zero
    for (time, auc) in &auc_preds {
        assert!(auc.is_finite(), "AUC at time {} should be finite", time);
        // Be more lenient - just check it's not NaN
    }

    Ok(())
}

#[test]
fn test_multi_outeq_auc_mode() -> Result<()> {
    // Test that AUC optimization works correctly with multiple output equations
    // This validates that predictions are properly separated by outeq before AUC calculation

    // SIMPLIFIED TEST: Just verify cost calculation doesn't crash with multi-outeq
    // Don't run full optimization (too slow for unit test)

    // Create a simple one-compartment model with two output equations
    let eq = equation::ODE::new(
        |x, p, _t, dx, b, _rateiv, _cov| {
            fetch_params!(p, ke, _v);
            dx[0] = -ke * x[0] + b[0];
        },
        |_p, _, _| lag! {},
        |_p, _, _| fa! {},
        |_p, _t, _cov, _x| {},
        |x, p, _t, _cov, y| {
            fetch_params!(p, _ke, v);
            y[0] = x[0] / v; // outeq 0: concentration
            y[1] = x[0]; // outeq 1: amount
        },
    );

    let params = Parameters::new().add("ke", 0.1, 0.5).add("v", 40.0, 60.0);

    let error_model = AssayErrorModel::additive(ErrorPoly::new(0.0, 5.0, 0.0, 0.0), 0.0);
    let ems = AssayErrorModels::new()
        .add(0, error_model.clone())?
        .add(1, error_model)?;

    let mut settings = Settings::builder()
        .set_algorithm(Algorithm::NPAG)
        .set_parameters(params.clone())
        .set_error_models(ems.clone())
        .build();

    settings.disable_output();
    settings.set_cycles(0);
    settings.set_idelta(30.0); // 30-minute intervals for AUC calculation

    // Subject with fixed dose and target observations at multiple outeqs
    let target = Subject::builder("test")
        .bolus(0.0, 500.0, 0) // FIXED dose (not optimizable)
        .observation(2.0, 40.0, 0) // Target AUC at outeq 0 (concentration)
        .observation(4.0, 200.0, 1) // Target AUC at outeq 1 (amount)
        .build();

    // Create prior with reasonable PK parameters
    let prior_theta = {
        let mat = faer::Mat::from_fn(1, 2, |_r, c| match c {
            0 => 0.2,  // ke
            1 => 50.0, // v
            _ => 0.0,
        });
        Theta::from_parts(mat, settings.parameters().clone())?
    };
    let prior_weights = Weights::uniform(1);

    let posterior = BestDosePosterior::compute(&prior_theta, &prior_weights, None, eq, settings)?;

    let _result = posterior.optimize(
        target,
        None,
        DoseRange::new(0.0, 2000.0),
        0.5,
        Target::AUCFromZero,
    )?;

    // Just verify that posterior compute and optimize succeed
    // This tests that cost calculation works with multi-outeq

    Ok(())
}

#[test]
#[ignore] // Mark as ignored - full optimization test is too slow
fn test_multi_outeq_auc_optimization() -> Result<()> {
    // Full optimization test - only run when explicitly requested
    let eq = equation::ODE::new(
        |x, p, _t, dx, b, _rateiv, _cov| {
            fetch_params!(p, ke, _v);
            dx[0] = -ke * x[0] + b[0];
        },
        |_p, _, _| lag! {},
        |_p, _, _| fa! {},
        |_p, _t, _cov, _x| {},
        |x, p, _t, _cov, y| {
            fetch_params!(p, _ke, v);
            y[0] = x[0] / v;
            y[1] = x[0];
        },
    );

    let params = Parameters::new().add("ke", 0.1, 0.5).add("v", 40.0, 60.0);
    let error_model = AssayErrorModel::additive(ErrorPoly::new(0.0, 5.0, 0.0, 0.0), 0.0);
    let ems = AssayErrorModels::new()
        .add(0, error_model.clone())?
        .add(1, error_model)?;

    let mut settings = Settings::builder()
        .set_algorithm(Algorithm::NPAG)
        .set_parameters(params.clone())
        .set_error_models(ems.clone())
        .build();
    settings.disable_output();
    settings.set_cycles(3);

    let target = Subject::builder("test")
        .bolus(0.0, 0.0, 0)
        .observation(2.0, 40.0, 0)
        .observation(4.0, 200.0, 1)
        .build();

    let prior_theta = {
        let mat = faer::Mat::from_fn(1, 2, |_r, c| match c {
            0 => 0.2,
            1 => 50.0,
            _ => 0.0,
        });
        Theta::from_parts(mat, settings.parameters().clone())?
    };
    let prior_weights = Weights::uniform(1);

    let posterior = BestDosePosterior::compute(&prior_theta, &prior_weights, None, eq, settings)?;

    let result = posterior.optimize(
        target,
        None,
        DoseRange::new(0.0, 2000.0),
        0.5,
        Target::AUCFromZero,
    );
    assert!(
        result.is_ok(),
        "Multi-outeq AUC optimization failed: {:?}",
        result.err()
    );

    let best_dose_result = result?;

    let doses = best_dose_result.doses();

    assert_eq!(doses.len(), 1);
    assert!(doses[0] > 0.0);
    assert!(best_dose_result.objf().is_finite());

    assert!(best_dose_result.auc_predictions().is_some());
    let auc_preds = best_dose_result.auc_predictions().unwrap();
    assert_eq!(
        auc_preds.len(),
        2,
        "Should have 2 AUC predictions (one per outeq)"
    );

    Ok(())
}

// ============================================================================
// AUC MODE TESTS - Comprehensive testing for both AUC calculation modes
// ============================================================================

/// Test AUCFromZero: Verify cumulative AUC calculation from time 0
#[test]
fn test_auc_from_zero_single_dose() -> Result<()> {
    let eq = equation::ODE::new(
        |x, p, _t, dx, b, _rateiv, _cov| {
            fetch_params!(p, ke, _v);
            dx[0] = -ke * x[0] + b[0];
        },
        |_p, _, _| lag! {},
        |_p, _, _| fa! {},
        |_p, _t, _cov, _x| {},
        |x, p, _t, _cov, y| {
            fetch_params!(p, _ke, v);
            y[0] = x[0] / v;
        },
    );

    let params = Parameters::new().add("ke", 0.2, 0.4).add("v", 40.0, 60.0);

    let ems = AssayErrorModels::new().add(
        0,
        AssayErrorModel::additive(ErrorPoly::new(0.0, 0.20, 0.0, 0.0), 0.0),
    )?;

    let mut settings = Settings::builder()
        .set_algorithm(Algorithm::NPAG)
        .set_parameters(params)
        .set_error_models(ems.clone())
        .build();

    settings.disable_output();
    settings.set_cycles(0);
    settings.set_idelta(10.0); // 10-minute intervals for AUC calculation

    // Target: Single dose, cumulative AUC from 0 to 12h
    let target = Subject::builder("patient_auc_zero")
        .bolus(0.0, 0.0, 0) // Dose to optimize
        .observation(12.0, 150.0, 0) // Target: AUC₀₋₁₂ = 150 mg·h/L
        .build();

    let prior_theta = {
        let mat = faer::Mat::from_fn(1, 2, |_r, c| match c {
            0 => 0.3,  // ke
            1 => 50.0, // v
            _ => 0.0,
        });
        Theta::from_parts(mat, settings.parameters().clone())?
    };
    let prior_weights = Weights::uniform(1);

    let posterior = BestDosePosterior::compute(&prior_theta, &prior_weights, None, eq, settings)?;

    let result = posterior.optimize(
        target,
        None,
        DoseRange::new(100.0, 1000.0),
        0.8,
        Target::AUCFromZero, // Cumulative AUC from time 0
    )?;

    let doses: Vec<f64> = result.doses();

    // Verify we got a result
    assert_eq!(doses.len(), 1);
    assert!(doses[0] > 0.0);
    assert!(result.objf().is_finite());

    // Verify we have AUC predictions
    assert!(result.auc_predictions().is_some());
    let auc_preds = result.auc_predictions().unwrap();
    assert_eq!(auc_preds.len(), 1);

    let (time, auc) = auc_preds[0];
    assert!((time - 12.0).abs() < 0.01);
    assert!(auc > 0.0 && auc.is_finite());

    eprintln!("AUCFromZero test:");
    eprintln!("  Optimal dose: {:.1} mg", doses[0]);
    eprintln!("  Predicted AUC₀₋₁₂: {:.2} mg·h/L", auc);
    eprintln!("  Target AUC₀₋₁₂: 150.0 mg·h/L");

    Ok(())
}

/// Test AUCFromLastDose: Verify interval AUC calculation from last dose
#[test]
fn test_auc_from_last_dose_maintenance() -> Result<()> {
    let eq = equation::ODE::new(
        |x, p, _t, dx, b, _rateiv, _cov| {
            fetch_params!(p, ke, _v);
            dx[0] = -ke * x[0] + b[0];
        },
        |_p, _, _| lag! {},
        |_p, _, _| fa! {},
        |_p, _t, _cov, _x| {},
        |x, p, _t, _cov, y| {
            fetch_params!(p, _ke, v);
            y[0] = x[0] / v;
        },
    );

    let params = Parameters::new().add("ke", 0.2, 0.4).add("v", 40.0, 60.0);

    let ems = AssayErrorModels::new().add(
        0,
        AssayErrorModel::additive(ErrorPoly::new(0.0, 0.20, 0.0, 0.0), 0.0),
    )?;

    let mut settings = Settings::builder()
        .set_algorithm(Algorithm::NPAG)
        .set_parameters(params)
        .set_error_models(ems.clone())
        .build();

    settings.disable_output();
    settings.set_cycles(0);
    settings.set_idelta(10.0);

    // Target: Loading dose (fixed) + maintenance dose (optimize)
    // Target interval AUC from t=12 to t=24
    let target = Subject::builder("patient_auc_interval")
        .bolus(0.0, 300.0, 0) // Loading dose (fixed at 300 mg)
        .bolus(12.0, 0.0, 0) // Maintenance dose to optimize
        .observation(24.0, 80.0, 0) // Target: AUC₁₂₋₂₄ = 80 mg·h/L
        .build();

    let prior_theta = {
        let mat = faer::Mat::from_fn(1, 2, |_r, c| match c {
            0 => 0.3,  // ke
            1 => 50.0, // v
            _ => 0.0,
        });
        Theta::from_parts(mat, settings.parameters().clone())?
    };
    let prior_weights = Weights::uniform(1);

    let posterior = BestDosePosterior::compute(&prior_theta, &prior_weights, None, eq, settings)?;

    let result = posterior.optimize(
        target,
        None,
        DoseRange::new(50.0, 500.0),
        0.8,
        Target::AUCFromLastDose, // Interval AUC from last dose
    )?;
    let doses = result.doses();

    // Verify we got a result
    assert_eq!(doses.len(), 2, "Should be 2 doses (loading + maintenance)");
    // Very first one is fixed loading dose, second is optimized maintenance dose
    assert_eq!(doses[0], 300.0);
    assert!(doses[0] > 0.0);
    assert!(result.objf().is_finite());

    // Verify we have AUC predictions
    assert!(result.auc_predictions().is_some());
    let auc_preds = result.auc_predictions().unwrap();
    assert_eq!(auc_preds.len(), 1);

    let (time, auc) = auc_preds[0];
    assert!((time - 24.0).abs() < 0.01);
    assert!(auc > 0.0 && auc.is_finite());

    eprintln!("AUCFromLastDose test:");
    eprintln!("  Loading dose (fixed): 300.0 mg at t=0");
    eprintln!("  Optimal maintenance dose: {:.1} mg at t=12", doses[0]);
    eprintln!("  Predicted AUC₁₂₋₂₄: {:.2} mg·h/L", auc);
    eprintln!("  Target AUC₁₂₋₂₄: 80.0 mg·h/L");

    Ok(())
}

/// Test comparison: AUCFromZero vs AUCFromLastDose should give different results
#[test]
fn test_auc_modes_comparison() -> Result<()> {
    let eq = equation::ODE::new(
        |x, p, _t, dx, b, _rateiv, _cov| {
            fetch_params!(p, ke, _v);
            dx[0] = -ke * x[0] + b[0];
        },
        |_p, _, _| lag! {},
        |_p, _, _| fa! {},
        |_p, _t, _cov, _x| {},
        |x, p, _t, _cov, y| {
            fetch_params!(p, _ke, v);
            y[0] = x[0] / v;
        },
    );

    let params = Parameters::new().add("ke", 0.3, 0.3).add("v", 50.0, 50.0);

    let ems = AssayErrorModels::new().add(
        0,
        AssayErrorModel::additive(ErrorPoly::new(0.0, 0.20, 0.0, 0.0), 0.0),
    )?;

    let mut settings = Settings::builder()
        .set_algorithm(Algorithm::NPAG)
        .set_parameters(params)
        .set_error_models(ems.clone())
        .build();

    settings.disable_output();
    settings.set_cycles(0);
    settings.set_idelta(10.0);

    let prior_theta = {
        let mat = faer::Mat::from_fn(1, 2, |_r, c| match c {
            0 => 0.3,  // ke
            1 => 50.0, // v
            _ => 0.0,
        });
        Theta::from_parts(mat, settings.parameters().clone())?
    };
    let prior_weights = Weights::uniform(1);

    // Scenario: Two doses, observation after second dose
    // Target same AUC value (100 mg·h/L) but different interpretation

    // Mode 1: AUCFromZero - target is cumulative AUC from t=0 to t=24
    let target_zero = Subject::builder("patient_zero")
        .bolus(0.0, 200.0, 0) // First dose fixed
        .bolus(12.0, 0.0, 0) // Second dose to optimize
        .observation(24.0, 100.0, 0) // Target: AUC₀₋₂₄ = 100
        .build();

    let posterior_zero = BestDosePosterior::compute(
        &prior_theta,
        &prior_weights,
        None,
        eq.clone(),
        settings.clone(),
    )?;

    let result_zero = posterior_zero.optimize(
        target_zero,
        None,
        DoseRange::new(10.0, 2000.0),
        0.8,
        Target::AUCFromZero,
    )?;
    // Extract only the second dose (the optimized one at t=12)
    let dose_zero = result_zero.doses()[1];

    // Mode 2: AUCFromLastDose - target is interval AUC from t=12 to t=24
    let target_last = Subject::builder("patient_last")
        .bolus(0.0, 200.0, 0) // First dose fixed
        .bolus(12.0, 0.0, 0) // Second dose to optimize
        .observation(24.0, 100.0, 0) // Target: AUC₁₂₋₂₄ = 100
        .build();

    let posterior_last =
        BestDosePosterior::compute(&prior_theta, &prior_weights, None, eq, settings)?;

    let result_last = posterior_last.optimize(
        target_last,
        None,
        DoseRange::new(10.0, 2000.0),
        0.8,
        Target::AUCFromLastDose,
    )?;
    // Extract only the second dose (the optimized one at t=12)
    let dose_last = result_last.doses()[1];

    // The two modes should recommend DIFFERENT doses for the same target value
    // because they're measuring different things
    eprintln!("\nAUC Mode Comparison:");
    eprintln!("  Scenario: 200mg at t=0 (fixed), optimize dose at t=12");
    eprintln!("  Target value: 100 mg·h/L (same number, different meaning)");
    eprintln!("  ");
    eprintln!("  AUCFromZero (cumulative 0→24h):");
    eprintln!("    Optimal 2nd dose: {:.1} mg", dose_zero);
    eprintln!(
        "    AUC prediction: {:.2}",
        result_zero.auc_predictions().as_ref().unwrap()[0].1
    );
    eprintln!("  ");
    eprintln!("  AUCFromLastDose (interval 12→24h):");
    eprintln!("    Optimal 2nd dose: {:.1} mg", dose_last);
    eprintln!(
        "    AUC prediction: {:.2}",
        result_last.auc_predictions().as_ref().unwrap()[0].1
    );

    // Verify both modes work
    assert!(dose_zero > 0.0);
    assert!(dose_last > 0.0);

    // The doses should be different (cumulative includes first dose effect,
    // interval only measures second dose)
    // We expect AUCFromZero to recommend a smaller second dose since it includes
    // the AUC contribution from the first dose
    assert_ne!(
        (dose_zero * 10.0).round() / 10.0,
        (dose_last * 10.0).round() / 10.0,
        "AUCFromZero and AUCFromLastDose should recommend different doses"
    );

    Ok(())
}

/// Test AUCFromLastDose with multiple observations
#[test]
fn test_auc_from_last_dose_multiple_observations() -> Result<()> {
    let eq = equation::ODE::new(
        |x, p, _t, dx, b, _rateiv, _cov| {
            fetch_params!(p, ke, _v);
            dx[0] = -ke * x[0] + b[0];
        },
        |_p, _, _| lag! {},
        |_p, _, _| fa! {},
        |_p, _t, _cov, _x| {},
        |x, p, _t, _cov, y| {
            fetch_params!(p, _ke, v);
            y[0] = x[0] / v;
        },
    );

    let params = Parameters::new().add("ke", 0.2, 0.4).add("v", 40.0, 60.0);

    let ems = AssayErrorModels::new().add(
        0,
        AssayErrorModel::additive(ErrorPoly::new(0.0, 0.20, 0.0, 0.0), 0.0),
    )?;

    let mut settings = Settings::builder()
        .set_algorithm(Algorithm::NPAG)
        .set_parameters(params)
        .set_error_models(ems.clone())
        .build();

    settings.disable_output();
    settings.set_cycles(0);
    settings.set_idelta(10.0);

    // Multiple doses and observations - each observation measures AUC from its preceding dose
    let target = Subject::builder("patient_multi")
        .bolus(0.0, 0.0, 0) // Dose 1 to optimize
        .observation(12.0, 50.0, 0) // AUC₀₋₁₂ = 50
        .bolus(12.0, 0.0, 0) // Dose 2 to optimize
        .observation(24.0, 50.0, 0) // AUC₁₂₋₂₄ = 50
        .build();

    let prior_theta = {
        let mat = faer::Mat::from_fn(1, 2, |_r, c| match c {
            0 => 0.3,  // ke
            1 => 50.0, // v
            _ => 0.0,
        });
        Theta::from_parts(mat, settings.parameters().clone())?
    };
    let prior_weights = Weights::uniform(1);

    let posterior = BestDosePosterior::compute(&prior_theta, &prior_weights, None, eq, settings)?;

    let result = posterior.optimize(
        target,
        None,
        DoseRange::new(50.0, 500.0),
        0.8,
        Target::AUCFromLastDose,
    )?;
    let doses: Vec<f64> = result.doses();

    // Should optimize 2 doses
    assert_eq!(doses.len(), 2);
    assert!(doses[0] > 0.0);
    assert!(doses[1] > 0.0);

    // Should have 2 AUC predictions
    assert!(result.auc_predictions().is_some());
    let auc_preds = result.auc_predictions().unwrap();
    assert_eq!(auc_preds.len(), 2);

    // First observation measures AUC from t=0 (first dose) to t=12
    let (time1, auc1) = auc_preds[0];
    assert!((time1 - 12.0).abs() < 0.01);

    // Second observation measures AUC from t=12 (second dose) to t=24
    let (time2, auc2) = auc_preds[1];
    assert!((time2 - 24.0).abs() < 0.01);

    eprintln!("AUCFromLastDose multiple observations test:");
    eprintln!(
        "  Dose 1 (t=0): {:.1} mg → AUC₀₋₁₂ = {:.2} (target: 50.0)",
        doses[0], auc1
    );
    eprintln!(
        "  Dose 2 (t=12): {:.1} mg → AUC₁₂₋₂₄ = {:.2} (target: 50.0)",
        doses[1], auc2
    );

    Ok(())
}

/// Test edge case: observation before any dose (should integrate from time 0)
#[test]
fn test_auc_from_last_dose_no_prior_dose() -> Result<()> {
    let eq = equation::ODE::new(
        |x, p, _t, dx, b, _rateiv, _cov| {
            fetch_params!(p, ke, _v);
            dx[0] = -ke * x[0] + b[0];
        },
        |_p, _, _| lag! {},
        |_p, _, _| fa! {},
        |_p, _t, _cov, _x| {},
        |x, p, _t, _cov, y| {
            fetch_params!(p, _ke, v);
            y[0] = x[0] / v;
        },
    );

    let params = Parameters::new().add("ke", 0.2, 0.4).add("v", 40.0, 60.0);

    let ems = AssayErrorModels::new().add(
        0,
        AssayErrorModel::additive(ErrorPoly::new(0.0, 0.20, 0.0, 0.0), 0.0),
    )?;

    let mut settings = Settings::builder()
        .set_algorithm(Algorithm::NPAG)
        .set_parameters(params)
        .set_error_models(ems.clone())
        .build();

    settings.disable_output();
    settings.set_cycles(0);
    settings.set_idelta(10.0);

    // Edge case: observation at t=6, but dose is at t=12 (after the observation)
    let target = Subject::builder("patient_edge")
        .observation(6.0, 30.0, 0) // Observation before any dose
        .bolus(12.0, 0.0, 0) // Dose after observation
        .build();

    let prior_theta = {
        let mat = faer::Mat::from_fn(1, 2, |_r, c| match c {
            0 => 0.3,  // ke
            1 => 50.0, // v
            _ => 0.0,
        });
        Theta::from_parts(mat, settings.parameters().clone())?
    };
    let prior_weights = Weights::uniform(1);

    let posterior = BestDosePosterior::compute(&prior_theta, &prior_weights, None, eq, settings)?;

    let result = posterior.optimize(
        target,
        None,
        DoseRange::new(50.0, 500.0),
        0.8,
        Target::AUCFromLastDose,
    )?;
    let doses: Vec<f64> = result.doses();

    assert_eq!(doses.len(), 1);
    assert!(doses[0] > 0.0);

    assert!(result.auc_predictions().is_some());
    let auc_preds = result.auc_predictions().unwrap();
    assert_eq!(auc_preds.len(), 1);

    let (_time, auc) = auc_preds[0];

    eprintln!("AUCFromLastDose edge case (no prior dose):");
    eprintln!("  Observation at t=6 (before any dose)");
    eprintln!("  Dose at t=12: {:.1} mg", doses[0]);
    eprintln!("  AUC₀₋₆: {:.2} (should be ~0, no drug yet)", auc);

    assert!(
        auc.abs() < 1.0,
        "AUC before any dose should be nearly zero, got {}",
        auc
    );

    Ok(())
}

// ============================================================================
// DOSE RANGE BOUNDS TESTS - Verify optimizer respects DoseRange constraints
// ============================================================================

/// Test that optimizer respects DoseRange bounds
#[test]
fn test_dose_range_bounds_respected() -> Result<()> {
    // Create a simple one-compartment model
    let eq = equation::ODE::new(
        |x, p, _t, dx, b, _rateiv, _cov| {
            fetch_params!(p, ke, _v);
            dx[0] = -ke * x[0] + b[0];
        },
        |_p, _, _| lag! {},
        |_p, _, _| fa! {},
        |_p, _t, _cov, _x| {},
        |x, p, _t, _cov, y| {
            fetch_params!(p, _ke, v);
            y[0] = x[0] / v;
        },
    );

    let params = Parameters::new().add("ke", 0.1, 0.5).add("v", 40.0, 60.0);

    let ems = AssayErrorModels::new().add(
        0,
        AssayErrorModel::additive(ErrorPoly::new(0.0, 0.20, 0.0, 0.0), 0.0),
    )?;

    let mut settings = Settings::builder()
        .set_algorithm(Algorithm::NPAG)
        .set_parameters(params)
        .set_error_models(ems.clone())
        .build();

    settings.disable_output();
    settings.set_cycles(0);

    // Target with high concentration requiring large dose
    let target = Subject::builder("test_patient")
        .bolus(0.0, 0.0, 0) // Dose to optimize
        .observation(2.0, 20.0, 0) // High target concentration
        .build();

    let prior_theta = {
        let mat = faer::Mat::from_fn(1, 2, |_r, c| match c {
            0 => 0.3,  // ke
            1 => 50.0, // v
            _ => 0.0,
        });
        Theta::from_parts(mat, settings.parameters().clone())?
    };
    let prior_weights = Weights::uniform(1);

    // Set a narrow dose range: 50-200 mg
    let dose_range = DoseRange::new(50.0, 200.0);

    let posterior = BestDosePosterior::compute(
        &prior_theta,
        &prior_weights,
        None,
        eq.clone(),
        settings.clone(),
    )?;

    let result =
        posterior.optimize(target.clone(), None, dose_range, 0.0, Target::Concentration)?;
    let doses: Vec<f64> = result.doses();

    println!("Optimal dose: {:.1} mg", doses[0]);
    println!("Dose range: 50-200 mg");

    // Verify dose is within bounds
    assert!(
        doses[0] >= 50.0,
        "Dose {} is below minimum bound 50.0",
        doses[0]
    );
    assert!(
        doses[0] <= 200.0,
        "Dose {} is above maximum bound 200.0",
        doses[0]
    );

    // The optimal dose should hit the upper bound (200 mg) since the target is high
    // Allow small tolerance for numerical precision
    assert!(
        (doses[0] - 200.0).abs() < 1.0,
        "Expected dose near upper bound (200 mg), got {:.1} mg",
        doses[0]
    );

    Ok(())
}

// ═════════════════════════════════════════════════════════════════════════════
// Tests for time_offset behavior
// ═════════════════════════════════════════════════════════════════════════════

/// Helper to build a simple one-compartment model used by multiple tests
fn one_compartment_model() -> pharmsol::ODE {
    equation::ODE::new(
        |x, p, _t, dx, b, _rateiv, _cov| {
            fetch_params!(p, ke, _v);
            dx[0] = -ke * x[0] + b[0];
        },
        |_p, _, _| lag! {},
        |_p, _, _| fa! {},
        |_p, _t, _cov, _x| {},
        |x, p, _t, _cov, y| {
            fetch_params!(p, _ke, v);
            y[0] = x[0] / v;
        },
    )
}

/// Helper to build minimal settings for tests (no posterior refinement)
fn minimal_settings() -> Settings {
    let params = Parameters::new()
        .add("ke", 0.001, 3.0)
        .add("v", 25.0, 250.0);
    let ems = AssayErrorModels::new()
        .add(
            0,
            AssayErrorModel::additive(ErrorPoly::new(0.0, 0.20, 0.0, 0.0), 0.0),
        )
        .unwrap();
    let mut settings = Settings::builder()
        .set_algorithm(Algorithm::NPAG)
        .set_parameters(params)
        .set_error_models(ems)
        .build();
    settings.disable_output();
    settings.set_cycles(0);
    settings
}

/// Helper to build a simple prior (single support point)
fn simple_prior(settings: &Settings) -> (Theta, Weights) {
    let mat = faer::Mat::from_fn(1, 2, |_r, c| match c {
        0 => 0.3,  // ke
        1 => 50.0, // v
        _ => 0.0,
    });
    let theta = Theta::from_parts(mat, settings.parameters().clone()).unwrap();
    let weights = Weights::uniform(1);
    (theta, weights)
}

/// Test that gap=0 and gap=12 produce different results
///
/// When time_offset is applied as a gap after the last past event,
/// different gaps change when the future dose is given relative to
/// the past, affecting the PK simulation outcome.
#[test]
fn test_time_offset_zero_vs_nonzero_differ() -> Result<()> {
    let eq = one_compartment_model();
    let settings = minimal_settings();
    let (theta, weights) = simple_prior(&settings);

    // Past data: dose at t=0, observation at t=6
    let past = Subject::builder("patient")
        .bolus(0.0, 500.0, 0)
        .observation(6.0, 5.0, 0)
        .build();

    let posterior =
        BestDosePosterior::compute(&theta, &weights, Some(past), eq.clone(), settings.clone())?;

    // Target: optimizable dose at t=0 (relative), target conc at t=1 (relative)
    // Short observation window so residual from past dose matters
    let target = Subject::builder("patient")
        .bolus(0.0, 0.0, 0)
        .observation(1.0, 5.0, 0) // target: 5 mg/L at 1h after the future dose
        .build();

    // gap=0: target dose at t=6 absolute (right after past), obs at t=7
    // Past dose (500mg at t=0): C(7) = 500/50 * e^(-0.3*7) ≈ 1.22 mg/L residual
    let result_gap0 = posterior.optimize(
        target.clone(),
        Some(0.0),
        DoseRange::new(10.0, 1000.0),
        0.5,
        Target::Concentration,
    )?;

    // gap=12: target dose at t=18 absolute, obs at t=19
    // Past dose (500mg at t=0): C(19) = 500/50 * e^(-0.3*19) ≈ 0.003 mg/L (negligible)
    let result_gap12 = posterior.optimize(
        target,
        Some(12.0),
        DoseRange::new(10.0, 1000.0),
        0.5,
        Target::Concentration,
    )?;

    let doses_gap0 = result_gap0.doses();
    let doses_gap12 = result_gap12.doses();

    eprintln!("Gap=0  doses: {:?}", doses_gap0);
    eprintln!("Gap=12 doses: {:?}", doses_gap12);

    // With gap=0, there's still significant residual from the past dose (~1.2 mg/L),
    // so the optimizer needs less future dose. With gap=12, the past dose is negligible,
    // so it needs more future dose. The optimizable doses should differ.
    assert!(
        (doses_gap0.last().unwrap() - doses_gap12.last().unwrap()).abs() > 1e-3,
        "gap=0 and gap=12 must produce different optimizable doses, \
         but got {:.4} vs {:.4}",
        doses_gap0.last().unwrap(),
        doses_gap12.last().unwrap()
    );

    Ok(())
}

/// Test that the first target event lands at last_past_time + gap
/// and subsequent target times are shifted correctly.
#[test]
fn test_time_offset_event_placement() -> Result<()> {
    let eq = one_compartment_model();
    let settings = minimal_settings();
    let (theta, weights) = simple_prior(&settings);

    // Past: dose at t=0, observation at t=6 (last event at t=6)
    let past = Subject::builder("patient")
        .bolus(0.0, 500.0, 0)
        .observation(6.0, 5.0, 0)
        .build();

    let posterior =
        BestDosePosterior::compute(&theta, &weights, Some(past), eq.clone(), settings.clone())?;

    // Target: dose at t=0, dose at t=12, obs at t=24 (all relative)
    let target = Subject::builder("patient")
        .bolus(0.0, 0.0, 0)
        .bolus(12.0, 0.0, 0)
        .observation(24.0, 5.0, 0)
        .build();

    // gap=0: future starts immediately after last past event (t=6)
    // effective_offset = 6 + 0 = 6
    let gap = 0.0;
    let result = posterior.optimize(
        target,
        Some(gap),
        DoseRange::new(10.0, 500.0),
        0.5,
        Target::Concentration,
    )?;

    // After concatenation we should have:
    //   past dose at t=0 (fixed 500mg)
    //   target dose at t=0+6=6 (optimizable)
    //   target dose at t=12+6=18 (optimizable)
    //   target obs at t=24+6=30

    let optimal_subject = result.optimal_subject();
    let mut dose_times = Vec::new();
    let mut obs_times = Vec::new();

    for occ in optimal_subject.occasions() {
        for event in occ.events() {
            match event {
                Event::Bolus(b) => dose_times.push(b.time()),
                Event::Infusion(i) => dose_times.push(i.time()),
                Event::Observation(o) => obs_times.push(o.time()),
            }
        }
    }

    eprintln!("Dose times: {:?}", dose_times);
    eprintln!("Obs times: {:?}", obs_times);

    // Past dose at t=0
    assert!(
        (dose_times[0] - 0.0).abs() < 1e-10,
        "First dose (past) should be at t=0, got {}",
        dose_times[0]
    );
    // First target dose at t = 0 + 6 = 6
    assert!(
        (dose_times[1] - 6.0).abs() < 1e-10,
        "Second dose should be at t=0+effective_offset=6, got {}",
        dose_times[1]
    );
    // Second target dose at t = 12 + 6 = 18
    assert!(
        (dose_times[2] - 18.0).abs() < 1e-10,
        "Third dose should be at t=12+effective_offset=18, got {}",
        dose_times[2]
    );
    // Observation at t = 24 + 6 = 30
    assert!(
        (obs_times[0] - 30.0).abs() < 1e-10,
        "Observation should be at t=24+effective_offset=30, got {}",
        obs_times[0]
    );

    // Past dose should remain fixed at 500
    let doses = result.doses();
    assert!(
        (doses[0] - 500.0).abs() < 1e-6,
        "Past dose should be fixed at 500, got {}",
        doses[0]
    );

    Ok(())
}

/// Test that time_offset=None leaves target events unchanged
#[test]
fn test_time_offset_none_no_shift() -> Result<()> {
    let eq = one_compartment_model();
    let settings = minimal_settings();
    let (theta, weights) = simple_prior(&settings);

    let posterior =
        BestDosePosterior::compute(&theta, &weights, None, eq.clone(), settings.clone())?;

    let target = Subject::builder("patient")
        .bolus(0.0, 0.0, 0)
        .bolus(12.0, 0.0, 0)
        .observation(24.0, 5.0, 0)
        .build();

    let result = posterior.optimize(
        target,
        None, // No offset
        DoseRange::new(10.0, 500.0),
        0.5,
        Target::Concentration,
    )?;

    let optimal_subject = result.optimal_subject();
    let mut dose_times = Vec::new();
    let mut obs_times = Vec::new();

    for occ in optimal_subject.occasions() {
        for event in occ.events() {
            match event {
                Event::Bolus(b) => dose_times.push(b.time()),
                Event::Infusion(i) => dose_times.push(i.time()),
                Event::Observation(o) => obs_times.push(o.time()),
            }
        }
    }

    // Without offset, times should be exactly as specified in target
    assert!((dose_times[0] - 0.0).abs() < 1e-10);
    assert!((dose_times[1] - 12.0).abs() < 1e-10);
    assert!((obs_times[0] - 24.0).abs() < 1e-10);

    Ok(())
}

// ═════════════════════════════════════════════════════════════════════════════
// Tests for multi-target / multi-dose optimization
// ═════════════════════════════════════════════════════════════════════════════

/// Test that multiple optimizable doses all get meaningful values
#[test]
fn test_multi_dose_all_optimized() -> Result<()> {
    let eq = one_compartment_model();
    let settings = minimal_settings();
    let (theta, weights) = simple_prior(&settings);

    let posterior =
        BestDosePosterior::compute(&theta, &weights, None, eq.clone(), settings.clone())?;

    // Two optimizable doses, two target concentrations
    let target = Subject::builder("patient")
        .bolus(0.0, 0.0, 0)
        .bolus(12.0, 0.0, 0)
        .observation(6.0, 5.0, 0) // Target 5 mg/L at t=6
        .observation(18.0, 5.0, 0) // Target 5 mg/L at t=18
        .build();

    let result = posterior.optimize(
        target,
        None,
        DoseRange::new(10.0, 500.0),
        0.5,
        Target::Concentration,
    )?;

    let doses = result.doses();
    eprintln!("Multi-dose optimization: {:?}", doses);

    assert_eq!(doses.len(), 2, "Should optimize 2 doses");

    // Both doses should be meaningful (not collapsed to minimum)
    assert!(
        doses[0] > 10.0 + 1.0,
        "Dose 1 should be above minimum bound, got {}",
        doses[0]
    );
    assert!(
        doses[1] > 10.0 + 1.0,
        "Dose 2 should be above minimum bound, got {}",
        doses[1]
    );

    Ok(())
}

/// Test that changing target for dose 2 changes dose 2's result
#[test]
fn test_multi_target_second_dose_responds_to_target_change() -> Result<()> {
    let eq = one_compartment_model();
    let settings = minimal_settings();
    let (theta, weights) = simple_prior(&settings);

    let posterior =
        BestDosePosterior::compute(&theta, &weights, None, eq.clone(), settings.clone())?;

    // Scenario A: second target is LOW (2 mg/L)
    let target_low = Subject::builder("patient")
        .bolus(0.0, 0.0, 0)
        .bolus(12.0, 0.0, 0)
        .observation(6.0, 5.0, 0)
        .observation(18.0, 2.0, 0) // Low second target
        .build();

    // Scenario B: second target is HIGH (15 mg/L)
    let target_high = Subject::builder("patient")
        .bolus(0.0, 0.0, 0)
        .bolus(12.0, 0.0, 0)
        .observation(6.0, 5.0, 0)
        .observation(18.0, 15.0, 0) // High second target
        .build();

    let result_low = posterior.optimize(
        target_low,
        None,
        DoseRange::new(10.0, 1000.0),
        0.5,
        Target::Concentration,
    )?;

    let result_high = posterior.optimize(
        target_high,
        None,
        DoseRange::new(10.0, 1000.0),
        0.5,
        Target::Concentration,
    )?;

    let doses_low = result_low.doses();
    let doses_high = result_high.doses();

    eprintln!("Low second target:  doses = {:?}", doses_low);
    eprintln!("High second target: doses = {:?}", doses_high);

    // The second dose should be higher when the second target is higher
    assert!(
        doses_high[1] > doses_low[1],
        "Higher second target ({}) should produce higher second dose, \
         but got low={:.2} vs high={:.2}",
        15.0,
        doses_low[1],
        doses_high[1]
    );

    Ok(())
}

// ═════════════════════════════════════════════════════════════════════════════
// Tests for BestDosePosterior and BestDoseResult API surface
// ═════════════════════════════════════════════════════════════════════════════

/// Test BestDosePosterior accessor methods
#[test]
fn test_posterior_accessors() -> Result<()> {
    let eq = one_compartment_model();
    let settings = minimal_settings();
    let (theta, weights) = simple_prior(&settings);

    let posterior =
        BestDosePosterior::compute(&theta, &weights, None, eq.clone(), settings.clone())?;

    // n_support_points should match the prior (no filtering with 0 cycles and no data)
    assert!(
        posterior.n_support_points() > 0,
        "Posterior should have at least 1 support point"
    );

    // theta() should return a valid Theta with the correct number of rows
    assert_eq!(
        posterior.theta().matrix().nrows(),
        posterior.n_support_points()
    );

    // posterior_weights() should sum to ~1
    let weight_sum: f64 = posterior.posterior_weights().iter().sum();
    assert!(
        (weight_sum - 1.0).abs() < 1e-6,
        "Posterior weights should sum to 1.0, got {}",
        weight_sum
    );

    // population_weights() should also sum to ~1
    let pop_weight_sum: f64 = posterior.population_weights().iter().sum();
    assert!(
        (pop_weight_sum - 1.0).abs() < 1e-6,
        "Population weights should sum to 1.0, got {}",
        pop_weight_sum
    );

    Ok(())
}

/// Test BestDoseResult accessor methods
#[test]
fn test_result_accessors() -> Result<()> {
    let eq = one_compartment_model();
    let settings = minimal_settings();
    let (theta, weights) = simple_prior(&settings);

    let posterior =
        BestDosePosterior::compute(&theta, &weights, None, eq.clone(), settings.clone())?;

    let target = Subject::builder("patient")
        .bolus(0.0, 0.0, 0)
        .observation(6.0, 5.0, 0)
        .build();

    let result = posterior.optimize(
        target,
        None,
        DoseRange::new(10.0, 500.0),
        0.5,
        Target::Concentration,
    )?;

    // doses() should return 1 dose
    assert_eq!(result.doses().len(), 1);
    assert!(result.doses()[0].is_finite());

    // objf() should be finite and non-negative
    assert!(result.objf().is_finite());
    assert!(result.objf() >= 0.0, "Cost should be non-negative");

    // status() should be Converged (1000 iterations is usually enough for 1D)
    assert_eq!(
        *result.status(),
        pmcore::bestdose::BestDoseStatus::Converged
    );

    // predictions() should have predictions
    assert!(
        !result.predictions().predictions().is_empty(),
        "Predictions should not be empty"
    );

    // optimization_method() should be Posterior or Uniform
    let method = result.optimization_method();
    assert!(
        method == pmcore::bestdose::OptimalMethod::Posterior
            || method == pmcore::bestdose::OptimalMethod::Uniform
    );

    // auc_predictions() should be None for concentration targets
    assert!(
        result.auc_predictions().is_none(),
        "AUC predictions should be None for concentration targets"
    );

    // optimal_subject() should have the optimized dose
    let optimal_subj = result.optimal_subject();
    let mut found_dose = false;
    for occ in optimal_subj.occasions() {
        for event in occ.events() {
            if let Event::Bolus(b) = event {
                assert!(b.amount() > 0.0, "Optimized dose should be > 0");
                found_dose = true;
            }
        }
    }
    assert!(
        found_dose,
        "Should find at least one dose in optimal subject"
    );

    Ok(())
}

/// Test that negative time_offset is rejected
#[test]
fn test_negative_time_offset_rejected() -> Result<()> {
    let eq = one_compartment_model();
    let settings = minimal_settings();
    let (theta, weights) = simple_prior(&settings);

    let posterior =
        BestDosePosterior::compute(&theta, &weights, None, eq.clone(), settings.clone())?;

    let target = Subject::builder("patient")
        .bolus(0.0, 0.0, 0)
        .observation(6.0, 5.0, 0)
        .build();

    let result = posterior.optimize(
        target,
        Some(-1.0), // Negative offset should be rejected
        DoseRange::new(10.0, 500.0),
        0.5,
        Target::Concentration,
    );

    assert!(result.is_err(), "Negative time_offset should be rejected");
    assert!(
        result.unwrap_err().to_string().contains("negative"),
        "Error message should mention negative time_offset"
    );

    Ok(())
}

/// Test that posterior can be reused for multiple optimizations
/// This is the key new feature of the two-stage API
#[test]
fn test_posterior_reuse() -> Result<()> {
    let eq = one_compartment_model();
    let settings = minimal_settings();
    let (theta, weights) = simple_prior(&settings);

    // Compute posterior once
    let posterior =
        BestDosePosterior::compute(&theta, &weights, None, eq.clone(), settings.clone())?;

    // Optimize with different dose ranges
    let target = Subject::builder("patient")
        .bolus(0.0, 0.0, 0)
        .observation(6.0, 5.0, 0)
        .build();

    let result_narrow = posterior.optimize(
        target.clone(),
        None,
        DoseRange::new(10.0, 100.0),
        0.5,
        Target::Concentration,
    )?;

    let result_wide = posterior.optimize(
        target.clone(),
        None,
        DoseRange::new(10.0, 1000.0),
        0.5,
        Target::Concentration,
    )?;

    // Both should succeed
    assert!(result_narrow.doses()[0].is_finite());
    assert!(result_wide.doses()[0].is_finite());

    // Wide range should allow a potentially better (lower cost) result
    assert!(
        result_wide.objf() <= result_narrow.objf() + 1e-6,
        "Wider dose range should give equal or better cost: wide={:.6} vs narrow={:.6}",
        result_wide.objf(),
        result_narrow.objf()
    );

    // Optimize with different bias weights
    let result_personal = posterior.optimize(
        target.clone(),
        None,
        DoseRange::new(10.0, 500.0),
        0.0, // Full personalization
        Target::Concentration,
    )?;

    let result_population = posterior.optimize(
        target,
        None,
        DoseRange::new(10.0, 500.0),
        1.0, // Full population weighting
        Target::Concentration,
    )?;

    // Both should succeed
    assert!(result_personal.doses()[0].is_finite());
    assert!(result_population.doses()[0].is_finite());

    Ok(())
}
