"""
OHDSI Cohort Definition Builder

CIRCE expression syntax for cohort definitions:
- Multiple concept sets with hierarchy and mapping
- Primary criteria (Condition, Procedure, Drug, Measurement, etc.)
- Inclusion/exclusion rules with demographics
- Observation windows and temporal constraints
- Expression format variants (PascalCase/camelCase)

Author: Ankur Lohachab
Department of Advanced Computing Sciences, Maastricht University
"""

from __future__ import annotations
from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, Any, List, Optional


class CriteriaType(Enum):
    """Primary criteria domain types."""
    CONDITION = "ConditionOccurrence"
    PROCEDURE = "ProcedureOccurrence"
    DRUG = "DrugExposure"
    MEASUREMENT = "Measurement"
    OBSERVATION = "Observation"
    VISIT = "VisitOccurrence"
    DEVICE = "DeviceExposure"
    DEATH = "Death"


class AgeOperator(Enum):
    """Age comparison operators."""
    GT = "gt"   # Greater than
    GTE = "gte"  # Greater than or equal
    LT = "lt"   # Less than
    LTE = "lte"  # Less than or equal
    EQ = "eq"   # Equal


class Gender(Enum):
    """OMOP gender concept IDs."""
    MALE = 8507
    FEMALE = 8532
    UNKNOWN = 8551


@dataclass
class ConceptItem:
    """Individual concept in a concept set."""
    concept_id: int
    concept_name: Optional[str] = None
    include_descendants: bool = True
    include_mapped: bool = True
    is_excluded: bool = False

    def to_dict(self) -> Dict[str, Any]:
        return {
            "concept": {
                "CONCEPT_ID": self.concept_id,
                "CONCEPT_NAME": self.concept_name
            },
            "isExcluded": self.is_excluded,
            "includeDescendants": self.include_descendants,
            "includeMapped": self.include_mapped
        }


@dataclass
class ConceptSet:
    """OMOP concept set with expression."""
    id: int
    name: str
    concepts: List[ConceptItem] = field(default_factory=list)

    def add_concept(
        self,
        concept_id: int,
        concept_name: Optional[str] = None,
        include_descendants: bool = True,
        include_mapped: bool = True
    ) -> ConceptSet:
        """
        Add concept to this concept set.

        Args:
            concept_id: OMOP concept ID
            concept_name: Optional concept name
            include_descendants: Include descendant concepts
            include_mapped: Include mapped concepts

        Returns:
            Self for method chaining

        Raises:
            ValueError: If concept_id is invalid
        """
        if not isinstance(concept_id, int) or concept_id <= 0:
            raise ValueError(f"Concept ID must be positive integer, got: {concept_id}")

        self.concepts.append(ConceptItem(
            concept_id=concept_id,
            concept_name=concept_name,
            include_descendants=include_descendants,
            include_mapped=include_mapped
        ))
        return self

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "name": self.name,
            "expression": {
                "items": [c.to_dict() for c in self.concepts]
            }
        }


@dataclass
class AgeCriterion:
    """Age demographic criterion."""
    value: int
    operator: AgeOperator

    def to_dict(self) -> Dict[str, Any]:
        return {
            "Age": {
                "Value": self.value,
                "Op": self.operator.value
            }
        }


@dataclass
class GenderCriterion:
    """Gender demographic criterion."""
    genders: List[Gender]

    def to_dict(self) -> Dict[str, Any]:
        return {
            "Gender": [{"CONCEPT_ID": g.value} for g in self.genders]
        }


@dataclass
class ObservationWindow:
    """Observation window for primary criteria."""
    prior_days: int = 0
    post_days: int = 0

    def to_dict(self) -> Dict[str, Any]:
        return {
            "PriorDays": self.prior_days,
            "PostDays": self.post_days
        }


@dataclass
class InclusionRule:
    """Inclusion rule with demographics and additional criteria."""
    name: str
    description: Optional[str] = None
    demographics: List[Any] = field(default_factory=list)
    additional_criteria: List[Any] = field(default_factory=list)

    def add_age_criterion(self, operator: AgeOperator, value: int) -> InclusionRule:
        """Add age criterion to this rule."""
        self.demographics.append(AgeCriterion(value, operator))
        return self

    def add_gender_criterion(self, *genders: Gender) -> InclusionRule:
        """Add gender criterion to this rule."""
        self.demographics.append(GenderCriterion(list(genders)))
        return self

    def to_dict(self) -> Dict[str, Any]:
        demo_list = []
        for d in self.demographics:
            demo_list.append(d.to_dict())

        return {
            "name": self.name,
            "description": self.description or "",
            "expression": {
                "Type": "ALL",
                "CriteriaList": [],
                "DemographicCriteriaList": demo_list,
                "Groups": []
            }
        }


class CohortBuilderFull:
    """
    Builder for OHDSI cohort definitions with CIRCE syntax support.

    Supports cohorts with multiple concept sets, inclusion rules,
    demographics, and various criteria types.
    """

    def __init__(self, name: str, description: str = ""):
        self.name = name
        self.description = description
        self.concept_sets: List[ConceptSet] = []
        self.primary_criteria: List[Dict[str, Any]] = []
        self.observation_window = ObservationWindow()
        self.primary_criteria_limit = "All"
        self.inclusion_rules: List[InclusionRule] = []
        self.cdm_version_range = ">=5.0.0"
        self._next_concept_set_id = 0

    def add_concept_set(self, name: str) -> ConceptSet:
        """Create and add a new concept set."""
        cs = ConceptSet(id=self._next_concept_set_id, name=name)
        self._next_concept_set_id += 1
        self.concept_sets.append(cs)
        return cs

    def add_primary_condition(self, concept_set_id: int) -> CohortBuilderFull:
        """Add condition occurrence as primary criterion."""
        self.primary_criteria.append({
            "ConditionOccurrence": {"CodesetId": concept_set_id}
        })
        return self

    def add_primary_procedure(self, concept_set_id: int) -> CohortBuilderFull:
        """Add procedure occurrence as primary criterion."""
        self.primary_criteria.append({
            "ProcedureOccurrence": {"CodesetId": concept_set_id}
        })
        return self

    def add_primary_drug(self, concept_set_id: int) -> CohortBuilderFull:
        """Add drug exposure as primary criterion."""
        self.primary_criteria.append({
            "DrugExposure": {"CodesetId": concept_set_id}
        })
        return self

    def add_primary_measurement(self, concept_set_id: int) -> CohortBuilderFull:
        """Add measurement as primary criterion."""
        self.primary_criteria.append({
            "Measurement": {"CodesetId": concept_set_id}
        })
        return self

    def set_observation_window(
        self,
        prior_days: int = 0,
        post_days: int = 0
    ) -> CohortBuilderFull:
        """Set observation window for primary criteria."""
        self.observation_window = ObservationWindow(prior_days, post_days)
        return self

    def set_primary_criteria_limit(
        self,
        limit_type: str = "All"
    ) -> CohortBuilderFull:
        """Set primary criteria limit (All, First)."""
        self.primary_criteria_limit = limit_type
        return self

    def add_inclusion_rule(self, name: str, description: str = "") -> InclusionRule:
        """Create and add new inclusion rule."""
        rule = InclusionRule(name, description)
        self.inclusion_rules.append(rule)
        return rule

    def build(self) -> Dict[str, Any]:
        """
        Build complete cohort definition in OHDSI format.

        Returns:
            Complete cohort definition dictionary

        Raises:
            RuntimeError: If cohort definition is incomplete or invalid
        """
        if not self.name:
            raise RuntimeError("Cohort name is required")

        if not self.concept_sets:
            raise RuntimeError(
                "At least one concept set is required. "
                "Use add_concept_set() to add concepts."
            )

        if not self.primary_criteria:
            raise RuntimeError(
                "At least one primary criterion is required. "
                "Use add_primary_condition() or similar methods."
            )

        for cs in self.concept_sets:
            if not cs.concepts:
                raise RuntimeError(
                    f"Concept set '{cs.name}' is empty. "
                    f"Use concept_set.add_concept() to add concepts."
                )

        return {
            "name": self.name,
            "description": self.description,
            "expressionType": "SIMPLE_EXPRESSION",
            "expression": {
                "ConceptSets": [cs.to_dict() for cs in self.concept_sets],
                "PrimaryCriteria": {
                    "CriteriaList": self.primary_criteria,
                    "ObservationWindow": self.observation_window.to_dict(),
                    "PrimaryCriteriaLimit": {"Type": self.primary_criteria_limit}
                },
                "AdditionalCriteria": {
                    "Type": "ALL",
                    "CriteriaList": [],
                    "DemographicCriteriaList": [],
                    "Groups": []
                },
                "QualifiedLimit": {"Type": "First"},
                "ExpressionLimit": {"Type": "First"},
                "InclusionRules": [rule.to_dict() for rule in self.inclusion_rules],
                "CensoringCriteria": [],
                "CollapseSettings": {"CollapseType": "ERA", "EraPad": 0},
                "CensorWindow": {},
                "cdmVersionRange": self.cdm_version_range
            }
        }


def create_simple_cohort(
    name: str,
    description: str,
    concept_ids: List[int],
    include_descendants: bool = True,
    age_min: Optional[int] = None,
    age_max: Optional[int] = None,
    genders: Optional[List[Gender]] = None
) -> Dict[str, Any]:
    """
    Convenience function to create simple cohort with common criteria.

    Args:
        name: Cohort name
        description: Cohort description
        concept_ids: List of OMOP concept IDs for primary condition
        include_descendants: Include descendant concepts
        age_min: Minimum age (optional)
        age_max: Maximum age (optional)
        genders: List of genders to include (optional)

    Returns:
        Complete cohort definition dict
    """
    builder = CohortBuilderFull(name, description)

    # Add concept set for primary condition
    cs = builder.add_concept_set("Primary Condition")
    for cid in concept_ids:
        cs.add_concept(cid, include_descendants=include_descendants)

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

    # Add inclusion rule if demographics specified
    if age_min or age_max or genders:
        rule = builder.add_inclusion_rule("Demographics", "Age and gender filters")

        if age_min:
            rule.add_age_criterion(AgeOperator.GTE, age_min)
        if age_max:
            rule.add_age_criterion(AgeOperator.LTE, age_max)
        if genders:
            rule.add_gender_criterion(*genders)

    return builder.build()
