import os
import pytest
from cvasl.harmonizers import (
    NeuroHarmonize, 
    Covbat, 
    NeuroCombat, 
    NestedComBat, 
    ComscanNeuroCombat,
    AutoCombat,
    RELIEF,
    CombatPlusPlus
)
from cvasl.dataset import MRIdataset, encode_cat_features


def load_datasets(shared_datadir):
    """Load test datasets and preprocess them."""
    input_paths = [
        os.path.realpath(shared_datadir / "TestingData_Site1_fake.csv"),
        os.path.realpath(shared_datadir / "TestingData_Site2_fake.csv"),
        os.path.realpath(shared_datadir / "TrainingData_Site1_fake.csv")
    ]
    # Using duplicate site_id (1, 2, 1) to test that harmonizers handle multiple datasets from the same site correctly
    input_sites = [1, 2, 1]

    mri_datasets = [
        MRIdataset(input_path, input_site, "participant_id", features_to_drop=[])
        for input_site, input_path in zip(input_sites, input_paths)
    ]
    
    for mri_dataset in mri_datasets:
        mri_dataset.preprocess()
    
    features_to_map = ['sex']
    encode_cat_features(mri_datasets, features_to_map)
    
    return mri_datasets


def test_neurocombat(shared_datadir):
    """Test whether the NeuroCombat harmonizer runs."""
    datasets = load_datasets(shared_datadir)
    
    features_to_harmonize = [
        'ACA_B_CoV', 'MCA_B_CoV', 'PCA_B_CoV', 'TotalGM_B_CoV',
        'ACA_B_CBF', 'MCA_B_CBF', 'PCA_B_CBF', 'TotalGM_B_CBF'
    ]
    discrete_covariates = ['sex']
    continuous_covariates = ['age']
    patient_identifier = 'participant_id'
    site_indicator = 'site'

    harmonizer = NeuroCombat(
        features_to_harmonize, 
        discrete_covariates, 
        continuous_covariates, 
        patient_identifier, 
        site_indicator
    )
    harmonized_data = harmonizer.harmonize(datasets)
    
    # Verify harmonization succeeded
    assert harmonized_data is not None
    assert len(harmonized_data) == len(datasets)
    
    # Verify data structure is preserved
    for original, harmonized in zip(datasets, harmonized_data):
        assert harmonized.data is not None
        assert len(harmonized.data.columns) > 0


def test_neuroharmonize(shared_datadir):
    """Test whether the NeuroHarmonize harmonizer runs."""
    datasets = load_datasets(shared_datadir)
    
    features_to_harmonize = [
        'aca_b_cov', 'mca_b_cov', 'pca_b_cov', 'totalgm_b_cov',
        'aca_b_cbf', 'mca_b_cbf', 'pca_b_cbf', 'totalgm_b_cbf'
    ]
    covariates = ['age', 'sex', 'site']
    site_indicator = 'site'
    
    harmonizer = NeuroHarmonize(
        features_to_harmonize=features_to_harmonize,
        covariates=covariates,
        smooth_terms=[],
        site_indicator=site_indicator,
        empirical_bayes=True
    )
    harmonized_data = harmonizer.harmonize(datasets)
    
    assert harmonized_data is not None
    assert len(harmonized_data) == len(datasets)
    
    # Check that harmonized features exist in output
    for dataset in harmonized_data:
        for feature in features_to_harmonize:
            assert feature in dataset.data.columns


def test_covbat(shared_datadir):
    """Test whether the CovBat harmonizer runs."""
    datasets = load_datasets(shared_datadir)
    
    patient_identifier = 'participant_id'
    site_indicator = 'site'
    features_to_harmonize = [
        'aca_b_cov', 'mca_b_cov', 'pca_b_cov', 'totalgm_b_cov',
        'aca_b_cbf', 'mca_b_cbf', 'pca_b_cbf', 'totalgm_b_cbf'
    ]
    numerical_covariates = ['age']
    covariates = ['age', 'sex']
    
    harmonizer = Covbat(
        features_to_harmonize=features_to_harmonize,
        covariates=covariates,
        site_indicator=site_indicator,
        patient_identifier=patient_identifier,
        numerical_covariates=numerical_covariates
    )
    harmonized_data = harmonizer.harmonize(datasets)
    
    assert harmonized_data is not None
    assert len(harmonized_data) == len(datasets)


@pytest.mark.skip(reason="Requires multiple samples per covariate value (test data too small)")
def test_comscan_neurocombat(shared_datadir):
    """Test whether the ComscanNeuroCombat harmonizer runs."""
    datasets = load_datasets(shared_datadir)
    
    features_to_harmonize = [
        'aca_b_cov', 'mca_b_cov', 'pca_b_cov', 'totalgm_b_cov',
        'aca_b_cbf', 'mca_b_cbf', 'pca_b_cbf', 'totalgm_b_cbf'
    ]
    discrete_covariates = ['sex']
    continuous_covariates = ['age']
    site_indicator = 'site'
    
    harmonizer = ComscanNeuroCombat(
        features_to_harmonize=features_to_harmonize,
        discrete_covariates=discrete_covariates,
        continuous_covariates=continuous_covariates,
        site_indicator=site_indicator
    )
    harmonized_data = harmonizer.harmonize(datasets)
    
    # ComscanNeuroCombat may return None if no features to harmonize
    if harmonized_data is not None:
        assert len(harmonized_data) == len(datasets)


@pytest.mark.skip(reason="Not enough samples for clustering algorithm (test data too small)")
def test_autocombat(shared_datadir):
    """Test whether the AutoCombat harmonizer runs."""
    datasets = load_datasets(shared_datadir)
    
    features_to_harmonize = [
        'aca_b_cov', 'mca_b_cov', 'pca_b_cov', 'totalgm_b_cov',
        'aca_b_cbf', 'mca_b_cbf', 'pca_b_cbf', 'totalgm_b_cbf'
    ]
    data_subset = [
        'aca_b_cov', 'mca_b_cov', 'pca_b_cov', 'totalgm_b_cov',
        'aca_b_cbf', 'mca_b_cbf', 'pca_b_cbf', 'totalgm_b_cbf',
        'site', 'sex', 'age'
    ]
    discrete_covariates = ['sex']
    continuous_covariates = ['age']
    site_indicator = ['site']
    
    harmonizer = AutoCombat(
        data_subset=data_subset,
        features_to_harmonize=features_to_harmonize,
        site_indicator=site_indicator,
        discrete_covariates=discrete_covariates,
        continuous_covariates=continuous_covariates
    )
    harmonized_data = harmonizer.harmonize(datasets)
    
    if harmonized_data is not None:
        assert len(harmonized_data) == len(datasets)


@pytest.mark.skipif(
    not hasattr(pytest, 'R_AVAILABLE') or not getattr(pytest, 'R_AVAILABLE', False),
    reason="R dependencies (rpy2) not available"
)
@pytest.mark.skip(reason="Requires R packages (denoiseR, RcppCNPy) which may not be available in all environments")
def test_relief(shared_datadir):
    """Test whether the RELIEF harmonizer runs (requires R)."""
    datasets = load_datasets(shared_datadir)
    
    features_to_harmonize = [
        'aca_b_cov', 'mca_b_cov', 'pca_b_cov', 'totalgm_b_cov',
        'aca_b_cbf', 'mca_b_cbf', 'pca_b_cbf', 'totalgm_b_cbf'
    ]
    covariates = ['sex', 'age']
    patient_identifier = 'participant_id'
    
    harmonizer = RELIEF(
        features_to_harmonize=features_to_harmonize,
        covariates=covariates,
        patient_identifier=patient_identifier
    )
    harmonized_data = harmonizer.harmonize(datasets)
    
    assert harmonized_data is not None
    assert len(harmonized_data) == len(datasets)


@pytest.mark.skip(reason="File path issue trips up test on Windows")
@pytest.mark.skipif(
    not hasattr(pytest, 'R_AVAILABLE') or not getattr(pytest, 'R_AVAILABLE', False),
    reason="R dependencies (rpy2) not available"
)
def test_combat_plusplus(shared_datadir):
    """Test whether the CombatPlusPlus harmonizer runs (requires R)."""
    datasets = load_datasets(shared_datadir)
    
    features_to_harmonize = [
        'aca_b_cov', 'mca_b_cov', 'pca_b_cov', 'totalgm_b_cov',
        'aca_b_cbf', 'mca_b_cbf', 'pca_b_cbf', 'totalgm_b_cbf'
    ]
    discrete_covariates = ['sex']
    continuous_covariates = ['age']
    discrete_covariates_to_remove = []
    continuous_covariates_to_remove = []
    site_indicator = 'site'
    patient_identifier = 'participant_id'
    
    harmonizer = CombatPlusPlus(
        features_to_harmonize=features_to_harmonize,
        discrete_covariates=discrete_covariates,
        continuous_covariates=continuous_covariates,
        discrete_covariates_to_remove=discrete_covariates_to_remove,
        continuous_covariates_to_remove=continuous_covariates_to_remove,
        site_indicator=site_indicator,
        patient_identifier=patient_identifier
    )
    harmonized_data = harmonizer.harmonize(datasets)
    
    assert harmonized_data is not None
    assert len(harmonized_data) == len(datasets)


def test_nested_combat(shared_datadir):
    """Test whether the NestedComBat harmonizer runs."""
    datasets = load_datasets(shared_datadir)
    
    features_to_harmonize = [
        'ACA_B_CoV', 'MCA_B_CoV', 'PCA_B_CoV', 'TotalGM_B_CoV',
        'ACA_B_CBF', 'MCA_B_CBF', 'PCA_B_CBF', 'TotalGM_B_CBF'
    ]
    batch_list_harmonisations = ['site']  # Simplified for test data
    site_indicator = ['site']
    discrete_covariates = ['sex']
    continuous_covariates = ['age']
    patient_identifier = 'participant_id'
    
    harmonizer = NestedComBat(
        features_to_harmonize=features_to_harmonize,
        batch_list_harmonisations=batch_list_harmonisations,
        site_indicator=site_indicator,
        discrete_covariates=discrete_covariates,
        continuous_covariates=continuous_covariates,
        patient_identifier=patient_identifier,
        intermediate_results_path='.',
        return_extended=False,
        use_gmm=False
    )
    harmonized_data = harmonizer.harmonize(datasets)
    
    assert harmonized_data is not None
    assert len(harmonized_data) == len(datasets)


# Test for error handling
def test_neurocombat_with_invalid_input(shared_datadir):
    """Test that NeuroCombat raises appropriate errors with invalid input."""
    datasets = load_datasets(shared_datadir)
    
    # Test with empty features list
    with pytest.raises((ValueError, TypeError)):
        harmonizer = NeuroCombat(
            features_to_harmonize=[],
            discrete_covariates=['sex'],
            continuous_covariates=['age'],
            patient_identifier='participant_id',
            site_indicator='site'
        )


def test_harmonizer_preserves_patient_ids(shared_datadir):
    """Test that harmonization preserves patient identifiers."""
    datasets = load_datasets(shared_datadir)
    
    # Store original patient IDs
    original_ids = [set(dataset.data['participant_id'].values) for dataset in datasets]
    
    features_to_harmonize = [
        'ACA_B_CoV', 'MCA_B_CoV', 'PCA_B_CoV', 'TotalGM_B_CoV',
        'ACA_B_CBF', 'MCA_B_CBF', 'PCA_B_CBF', 'TotalGM_B_CBF'
    ]
    
    harmonizer = NeuroCombat(
        features_to_harmonize=features_to_harmonize,
        discrete_covariates=['sex'],
        continuous_covariates=['age'],
        patient_identifier='participant_id',
        site_indicator='site'
    )
    harmonized_data = harmonizer.harmonize(datasets)
    
    # Check that patient IDs are preserved
    for original_id_set, harmonized_dataset in zip(original_ids, harmonized_data):
        harmonized_id_set = set(harmonized_dataset.data['participant_id'].values)
        assert original_id_set == harmonized_id_set, "Patient IDs should be preserved"


def test_harmonizer_output_shape(shared_datadir):
    """Test that harmonized data has the expected shape."""
    datasets = load_datasets(shared_datadir)
    
    # Store original shapes
    original_shapes = [(len(dataset.data), len(dataset.data.columns)) 
                       for dataset in datasets]
    
    features_to_harmonize = [
        'ACA_B_CoV', 'MCA_B_CoV', 'PCA_B_CoV', 'TotalGM_B_CoV',
        'ACA_B_CBF', 'MCA_B_CBF', 'PCA_B_CBF', 'TotalGM_B_CBF'
    ]
    
    harmonizer = NeuroCombat(
        features_to_harmonize=features_to_harmonize,
        discrete_covariates=['sex'],
        continuous_covariates=['age'],
        patient_identifier='participant_id',
        site_indicator='site'
    )
    harmonized_data = harmonizer.harmonize(datasets)
    
    # Check that shapes match (rows should be same, columns should be preserved)
    for (orig_rows, orig_cols), harmonized_dataset in zip(original_shapes, harmonized_data):
        harm_rows = len(harmonized_dataset.data)
        harm_cols = len(harmonized_dataset.data.columns)
        
        # Number of rows should be preserved
        assert harm_rows == orig_rows, f"Row count mismatch: {harm_rows} vs {orig_rows}"
        # Column count should be similar (allowing for minor differences in processing)
        assert abs(harm_cols - orig_cols) <= 2, f"Column count significantly different: {harm_cols} vs {orig_cols}"
