1from pathlib import Path
  2
  3import pandas as pd
  4from loguru import logger
  5
  6from enpi_api.l1 import openapi_client
  7from enpi_api.l2.client.api.file_api import FileApi
  8from enpi_api.l2.events.workflow_execution_task_waitable import WorkflowExecutionTaskWaitable
  9from enpi_api.l2.types.api_error import ApiError, ApiErrorContext
 10from enpi_api.l2.types.cluster import ClusterId, ClusterRunId
 11from enpi_api.l2.types.execution import Execution
 12from enpi_api.l2.types.log import LogLevel
 13from enpi_api.l2.types.phylogeny import Phylogeny
 14from enpi_api.l2.types.task import TaskState
 15from enpi_api.l2.types.workflow import WorkflowExecutionId, WorkflowExecutionTaskId, WorkflowTaskTemplateName
 16
 17
 18class PhylogenyApi:
 19    _inner_api_client: openapi_client.ApiClient
 20    _log_level: LogLevel
 21
 22    def __init__(self, inner_api_client: openapi_client.ApiClient, log_level: LogLevel):
 23        """@private"""
 24        self._inner_api_client = inner_api_client
 25        self._log_level = log_level
 26
 27    def get_phylogeny_by_task_id(self, task_id: WorkflowExecutionTaskId) -> list[Phylogeny]:
 28        """Get the calculated trees for a phylogenetic calculation.
 29
 30        Args:
 31            task_id (enpi_api.l2.types.task.TaskId): The phylogeny task id
 32
 33        Returns:
 34            list[enpi_api.l2.types.phylogeny.Phylogeny]: The calculated trees.
 35
 36        Raises:
 37            enpi_api.l2.types.api_error.ApiError: If API request fails.
 38
 39        Example:
 40            ```python
 41                task_id = TaskId(456)
 42                phylogeny = self.get_phylogeny_by_task_id(task_id=task_id)
 43            ```
 44        """
 45
 46        phylogeny_api_instance = openapi_client.PhylogenyApi(self._inner_api_client)
 47
 48        try:
 49            return [Phylogeny.from_raw(i) for i in phylogeny_api_instance.get_phylogeny_by_task_id(task_id=task_id).phylogeny]
 50        except openapi_client.ApiException as e:
 51            raise ApiError(e)
 52
 53    def start(self, cluster_run_id: ClusterRunId, cluster_id: ClusterId) -> Execution[list[Phylogeny]]:
 54        """Start the calculation of phylogenetic trees for a specified cluster.
 55
 56        Args:
 57            cluster_run_id (enpi_api.l2.types.cluster.ClusterRunId): The unique identifier of the cluster run.
 58            cluster_id (enpi_api.l2.types.cluster.ClusterId): The unique identifier of the cluster.
 59
 60        Returns:
 61            enpi_api.l2.types.execution.Execution[list[enpi_api.l2.types.phylogeny.Phylogeny]]: An awaitable that returns the calculated trees.
 62
 63        Raises:
 64            enpi_api.l2.types.api_error.ApiError: If API request fails.
 65
 66        Example:
 67            ```python
 68                cluster_run_id = ClusterRunId("b35a7864-0887-44e5-896a-feffdcc9d022")
 69                cluster_id = ClusterId(123)
 70                phylogeny = client.phylogeny_api.start(cluster_run_id=cluster_run_id, cluster_id=cluster_id).wait()
 71            ```
 72        """
 73
 74        phylogeny_api_instance = openapi_client.PhylogenyApi(self._inner_api_client)
 75
 76        try:
 77            phylogeny_work = openapi_client.PhylogenyWork(cluster_run_id=cluster_run_id, cluster_id=cluster_id)
 78            data = phylogeny_api_instance.start_phylogeny(phylogeny_work)
 79
 80            assert data.workflow_execution_id is not None
 81
 82            workflow_execution_id = WorkflowExecutionId(int(data.workflow_execution_id))
 83
 84            def on_complete(task_id: WorkflowExecutionTaskId, task_state: TaskState) -> list[Phylogeny]:
 85                phylogeny: list[Phylogeny] = self.get_phylogeny_by_task_id(task_id=task_id)
 86
 87                logger.success(f"Phylogeny compute with task ID: {task_id} in workflow execution with ID: {workflow_execution_id} has successfully finished.")
 88
 89                return phylogeny
 90
 91            waitable = WorkflowExecutionTaskWaitable[list[Phylogeny]](
 92                workflow_execution_id=workflow_execution_id, task_template_name=WorkflowTaskTemplateName.ENPI_APP_PHYLOGENY, on_complete=on_complete
 93            )
 94            return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state)
 95        except openapi_client.ApiException as e:
 96            raise ApiError(e)
 97
 98    def export_phylogeny_as_tsv(self, phylogeny: Phylogeny, output_directory: str | Path | None = None) -> Execution[Path]:
 99        """Export the sequences of the specified tree into a TSV file and download it to the specified directory.
100
101        Args:
102            phylogeny (enpi_api.l2.types.phylogeny.Phylogeny): The phylogenetic tree to export.
103            output_directory (str | Path | None): The directory where to download the TSV file. Defaults to a
104              unique temporary directory.
105
106        Returns:
107            enpi_api.l2.types.execution.Execution[Path]: An awaitable that returns the path to the downloaded TSV file containing the sequences.
108
109        Raises:
110            enpi_api.l2.types.api_error.ApiError: If API request fails.
111
112        Example:
113            ```python
114            # Assuming phylogenetic trees were computed and fetched already
115            trees = ...
116
117            # Export the sequences of the Amino Acid tree
118            amino_acid_tree: Phylogeny = next(tree for tree in trees if tree.type == PhylogenyType.AMINO_ACID)
119
120            # Export the phylogenentic tree as a TSV file
121            logger.info("Exporting the phylogenentic tree as a TSV file")
122            phylogeny_tsv_path: Path = enpi_client.phylogeny_api.export_phylogeny_as_tsv(
123                phylogeny=amino_acid_tree,
124            ).wait()
125            ```
126        """
127
128        phylogeny_api_instance = openapi_client.PhylogenyApi(self._inner_api_client)
129
130        with ApiErrorContext():
131            export_request = openapi_client.StartPhylogenyExportRequest(tree_id=phylogeny.tree_id)
132            export_job = phylogeny_api_instance.start_phylogeny_export(export_request)
133            assert export_job.workflow_execution_id is not None
134
135            workflow_execution_id = WorkflowExecutionId(export_job.workflow_execution_id)
136
137            def on_complete(task_id: WorkflowExecutionTaskId, task_state: TaskState) -> Path:
138                file_api = FileApi(self._inner_api_client, self._log_level)
139                file_path = file_api.download_export_by_workflow_execution_task_id(task_id=task_id, output_directory=output_directory)
140
141                return file_path
142
143            waitable = WorkflowExecutionTaskWaitable[Path](
144                workflow_execution_id=workflow_execution_id, on_complete=on_complete, task_template_name=WorkflowTaskTemplateName.ENPI_APP_PHYLOGENY_EXPORT
145            )
146            return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state)
147
148    def export_phylogeny_as_df(self, phylogeny: Phylogeny) -> Execution[pd.DataFrame]:
149        """Export the sequences of the specified phylogeny into a pandas DataFrame.
150
151        Args:
152            phylogeny (enpi_api.l2.types.phylogeny.Phylogeny): The phylogenetic tree to export.
153
154        Returns:
155            enpi_api.l2.types.execution.Execution[pd.DataFrame]: An awaitable that returns a pandas DataFrame containing the sequences.
156
157        Raises:
158            enpi_api.l2.types.api_error.ApiError: If API request fails.
159
160        Example:
161            ```python
162            # Assuming trees were computed and fetched already
163            trees = ...
164
165            # Export the sequences of the Amino Acid tree
166            amino_acid_tree: Phylogeny = next(tree for tree in trees if tree.type == PhylogenyType.AMINO_ACID)
167
168            # Export the phylogenetic tree as a Pandas DataFrame
169            tree_df: pd.DataFrame = enpi_client.phylogeny_api.export_phylogeny_as_df(
170                tree=amino_acid_tree,
171            ).wait()
172        ```
173        """
174
175        export_tsv = self.export_phylogeny_as_tsv(phylogeny)
176
177        def wait() -> pd.DataFrame:
178            tsv_path = export_tsv.wait()
179            return pd.read_csv(tsv_path, sep="\t")
180
181        return Execution(wait=wait, check_execution_state=export_tsv.check_execution_state)