Cohort
======

.. currentmodule:: corr_vars.core.cohort

.. autoclass:: Cohort
   :members:
   :undoc-members:
   :show-inheritance:

Cohort Workflow
----------------

Initialization
^^^^^^^^^^^^^^^

.. code-block:: python

   cohort = Cohort(
       obs_level="icu_stay",  # One of: "hospital_stay", "icu_stay", "procedure"
       database="db_hypercapnia_prepared",
       merge_consecutive=True,  # Optional, defaults to True. Only for obs_level="icu_stay"
       load_default_vars=False  # Optional, defaults to True
   )

Adding Variables
^^^^^^^^^^^^^^^^

.. code-block:: python

   # Use pre-defined variables
   cohort.add_variable("pf_ratio")

   # Create a custom variable on the fly
   cohort.add_variable(
       NativeStatic(
           var_name="median_sodium_before_hn",
           select="!median value",
           base_var="blood_sodium",
           tmin="hospital_admission",
           tmax=cohort.t_eligible
       )
   )

   # Perform manual operations on the dataframe
   cohort.obs['idx_hypernatremia_was_on_admission'] = (
       cohort.obs['first_sodium_recordtime'] == cohort.obs['first_severe_hypernatremia_recordtime']
   )

   cohort.obs['hn_origin'] = np.where(
       cohort.obs['idx_hypernatremia_was_on_admission'],
       'community_acquired',
       'hospital_acquired'
   )



Inclusion/Exclusion
^^^^^^^^^^^^^^^^^^^

.. code-block:: python


   # Add multiple inclusion criteria
   cohort.add_inclusion([
       {
           "variable": "age",
           "operation": ">= 18",
           "label": "Adult patients"
       },
       {
           "variable": "icu_length_of_stay",
           "operation": "> 2",
           "label": "ICU stay > 2 days"
       }
   ])

   # Add exclusion criteria
   cohort.add_exclusion([
       {
           "variable": "any_dx_covid_19",
           "operation": "== True",
           "label": "Exclude COVID-19 patients"
       }
   ])

   # Visualize the exclusion criteria
   cohort.exclude_ct.plot_flowchart()


Exploration
^^^^^^^^^^^

.. code-block:: python

   # Display the cohort dataframe
   print(cohort.obs)


   # Create a TableOne object
   tableone = cohort.tableone()
   print(tableone)

Data Export
^^^^^^^^^^^

.. code-block:: python

   # Save to file
   cohort.save("my_cohort.corr")
   
   # Load from file
   cohort = Cohort.load("my_cohort.corr")
   
   # Export to CSV
   cohort.to_csv("path/to/output_folder")
   
   # Convert to AnnData
   adata = cohort.to_adata()