"""
Tests for cohort builders

Tests simple and full cohort building.
"""

import pytest
from smart_omop import (
    CohortBuilder,
    Gender,
    CohortBuilderFull,
    AgeOperator,
    create_simple_cohort
)


def test_simple_cohort_builder():
    """Test simple cohort builder."""
    builder = CohortBuilder("Test Cohort", "Test description")
    builder.with_condition("COPD", [255573])
    builder.with_age_range(min_age=60)
    builder.with_gender(Gender.FEMALE)

    cohort_def = builder.build()

    assert cohort_def.name == "Test Cohort"
    assert cohort_def.description == "Test description"
    assert isinstance(cohort_def.expression, dict)
    assert 'ConceptSets' in cohort_def.expression


def test_full_cohort_builder():
    """Test full cohort builder with CIRCE syntax."""
    builder = CohortBuilderFull("Complex Cohort", "Complex cohort definition")

    # Add concept set
    copd = builder.add_concept_set("COPD")
    copd.add_concept(255573, "COPD", include_descendants=True)

    # Add primary criterion
    builder.add_primary_condition(concept_set_id=0)

    # Add observation window
    builder.set_observation_window(prior_days=365, post_days=0)

    # Add inclusion rule
    rule = builder.add_inclusion_rule("Demographics")
    rule.add_age_criterion(AgeOperator.GTE, 60)
    rule.add_gender_criterion(Gender.FEMALE)

    cohort_def = builder.build()

    assert cohort_def['name'] == "Complex Cohort"
    assert len(cohort_def['expression']['ConceptSets']) == 1
    assert len(cohort_def['expression']['InclusionRules']) == 1
    assert cohort_def['expression']['PrimaryCriteria']['ObservationWindow']['PriorDays'] == 365


def test_create_simple_cohort():
    """Test create_simple_cohort convenience function."""
    cohort_def = create_simple_cohort(
        name="Simple COPD",
        description="Simple cohort",
        concept_ids=[255573],
        age_min=60,
        genders=[Gender.FEMALE]
    )

    assert cohort_def['name'] == "Simple COPD"
    assert len(cohort_def['expression']['ConceptSets']) == 1
    assert len(cohort_def['expression']['InclusionRules']) == 1


def test_multiple_concept_sets():
    """Test builder with multiple concept sets."""
    builder = CohortBuilderFull("Multi-Concept", "Multiple concept sets")

    copd = builder.add_concept_set("COPD")
    copd.add_concept(255573)

    htn = builder.add_concept_set("Hypertension")
    htn.add_concept(316866)

    builder.add_primary_condition(0)

    cohort_def = builder.build()

    assert len(cohort_def['expression']['ConceptSets']) == 2
    assert cohort_def['expression']['ConceptSets'][0]['name'] == "COPD"
    assert cohort_def['expression']['ConceptSets'][1]['name'] == "Hypertension"


def test_age_operators():
    """Test different age operators."""
    builder = CohortBuilderFull("Age Test", "Testing age operators")

    copd = builder.add_concept_set("COPD")
    copd.add_concept(255573)
    builder.add_primary_condition(0)

    rule = builder.add_inclusion_rule("Age Criteria")
    rule.add_age_criterion(AgeOperator.GTE, 60)
    rule.add_age_criterion(AgeOperator.LTE, 85)

    cohort_def = builder.build()

    demo_criteria = cohort_def['expression']['InclusionRules'][0]['expression']['DemographicCriteriaList']
    assert len(demo_criteria) == 2
    assert demo_criteria[0]['Age']['Op'] == 'gte'
    assert demo_criteria[0]['Age']['Value'] == 60
    assert demo_criteria[1]['Age']['Op'] == 'lte'
    assert demo_criteria[1]['Age']['Value'] == 85
