"""
Cohort Definition Builder

Simplified interface for creating OMOP CDM cohort definitions programmatically.
Supports common clinical criteria patterns.

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

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


class Gender(Enum):
    """Standard gender concepts."""
    MALE = 8507
    FEMALE = 8532
    UNKNOWN = 0


@dataclass
class CohortDefinition:
    """
    OMOP CDM cohort definition structure.

    Represents a cohort definition compatible with OHDSI WebAPI format.
    """

    name: str
    description: str = ""
    expression: Dict[str, Any] = field(default_factory=dict)
    id: Optional[int] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert to WebAPI-compatible dictionary."""
        result = {
            "name": self.name,
            "description": self.description,
            "expression": self.expression
        }
        if self.id is not None:
            result["id"] = self.id
        return result


class CohortBuilder:
    """
    Builder for OMOP cohort definitions.

    Provides fluent interface for constructing cohort inclusion criteria.
    """

    def __init__(self, name: str, description: str = ""):
        """
        Initialize cohort builder.

        Args:
            name: Cohort name
            description: Cohort description
        """
        self.name = name
        self.description = description
        self.concept_sets: List[Dict[str, Any]] = []
        self.criteria: List[Dict[str, Any]] = []
        self.age_min: Optional[int] = None
        self.age_max: Optional[int] = None
        self.gender: Optional[Gender] = None
        self.observation_window_prior: int = 0
        self.observation_window_post: int = 0

    def add_concept_set(
        self,
        name: str,
        concept_ids: List[int],
        include_descendants: bool = True
    ) -> CohortBuilder:
        """
        Add concept set to cohort definition.

        Args:
            name: Concept set name
            concept_ids: List of OMOP concept IDs
            include_descendants: Include descendant concepts in hierarchy

        Returns:
            Self for method chaining
        """
        concept_set_id = len(self.concept_sets)

        concept_set = {
            "id": concept_set_id,
            "name": name,
            "expression": {
                "items": [
                    {
                        "concept": {
                            "CONCEPT_ID": cid,
                            "CONCEPT_NAME": f"Concept {cid}",
                            "STANDARD_CONCEPT": "S",
                            "INVALID_REASON": "V",
                            "CONCEPT_CODE": str(cid),
                            "DOMAIN_ID": "Condition",
                            "VOCABULARY_ID": "SNOMED",
                            "CONCEPT_CLASS_ID": "Clinical Finding"
                        },
                        "isExcluded": False,
                        "includeDescendants": include_descendants,
                        "includeMapped": False
                    }
                    for cid in concept_ids
                ]
            }
        }

        self.concept_sets.append(concept_set)
        return self

    def with_condition(
        self,
        concept_set_name: str,
        concept_ids: List[int]
    ) -> CohortBuilder:
        """
        Add condition occurrence criterion.

        Args:
            concept_set_name: Name for condition concept set
            concept_ids: Condition concept IDs

        Returns:
            Self for method chaining
        """
        concept_set_id = len(self.concept_sets)
        self.add_concept_set(concept_set_name, concept_ids)

        criterion = {
            "ConditionOccurrence": {
                "CodesetId": concept_set_id,
                "ConditionTypeExclude": False
            }
        }
        self.criteria.append(criterion)
        return self

    def with_age_range(
        self,
        min_age: Optional[int] = None,
        max_age: Optional[int] = None
    ) -> CohortBuilder:
        """
        Set age range criterion.

        Args:
            min_age: Minimum age in years
            max_age: Maximum age in years

        Returns:
            Self for method chaining
        """
        self.age_min = min_age
        self.age_max = max_age
        return self

    def with_gender(self, gender: Gender) -> CohortBuilder:
        """
        Set gender criterion.

        Args:
            gender: Gender requirement

        Returns:
            Self for method chaining
        """
        self.gender = gender
        return self

    def with_observation_window(
        self,
        prior_days: int = 0,
        post_days: int = 0
    ) -> CohortBuilder:
        """
        Set observation window requirements.

        Args:
            prior_days: Days of observation before index
            post_days: Days of observation after index

        Returns:
            Self for method chaining
        """
        self.observation_window_prior = prior_days
        self.observation_window_post = post_days
        return self

    def build(self) -> CohortDefinition:
        """
        Build cohort definition.

        Returns:
            Complete cohort definition
        """
        # Build primary criteria
        primary_criteria = {
            "CriteriaList": self.criteria,
            "ObservationWindow": {
                "PriorDays": self.observation_window_prior,
                "PostDays": self.observation_window_post
            },
            "PrimaryCriteriaLimit": {
                "Type": "First"
            }
        }

        # Add demographic criteria if specified
        if self.age_min is not None or self.age_max is not None or self.gender is not None:
            inclusion_rules = []

            if self.age_min is not None or self.age_max is not None:
                age_rule = {
                    "name": "Age Range",
                    "expression": {
                        "Type": "ALL",
                        "CriteriaList": [],
                        "DemographicCriteriaList": [
                            {
                                "Age": {
                                    "Value": self.age_min or 0,
                                    "Op": "gte"
                                }
                            }
                        ]
                    }
                }

                if self.age_max is not None:
                    age_rule["expression"]["DemographicCriteriaList"].append({
                        "Age": {
                            "Value": self.age_max,
                            "Op": "lte"
                        }
                    })

                inclusion_rules.append(age_rule)

            if self.gender is not None:
                gender_rule = {
                    "name": "Gender",
                    "expression": {
                        "Type": "ALL",
                        "CriteriaList": [],
                        "DemographicCriteriaList": [
                            {
                                "Gender": [
                                    {
                                        "CONCEPT_ID": self.gender.value,
                                        "CONCEPT_NAME": self.gender.name
                                    }
                                ]
                            }
                        ]
                    }
                }
                inclusion_rules.append(gender_rule)

            expression = {
                "ConceptSets": self.concept_sets,
                "PrimaryCriteria": primary_criteria,
                "InclusionRules": inclusion_rules,
                "QualifiedLimit": {
                    "Type": "First"
                },
                "ExpressionLimit": {
                    "Type": "First"
                },
                "CensoringCriteria": [],
                "CollapseSettings": {
                    "CollapseType": "ERA",
                    "EraPad": 0
                }
            }
        else:
            expression = {
                "ConceptSets": self.concept_sets,
                "PrimaryCriteria": primary_criteria,
                "QualifiedLimit": {
                    "Type": "First"
                },
                "ExpressionLimit": {
                    "Type": "First"
                },
                "CensoringCriteria": [],
                "CollapseSettings": {
                    "CollapseType": "ERA",
                    "EraPad": 0
                }
            }

        return CohortDefinition(
            name=self.name,
            description=self.description,
            expression=expression
        )
