<reponame>jfaccioni/clovars
from __future__ import annotations

from typing import Generator, TYPE_CHECKING

from matplotlib import pyplot as plt

if TYPE_CHECKING:
    from pathlib import Path
    from clovars.bio import Treatment


class TreatmentDrawer:
    """Class containing functions to draw and display Treatments."""
    def __init__(
            self,
            treatment_data: dict[tuple[str, int], Treatment],
    ) -> None:
        """Initializes a TreeDrawer instance."""
        self.treatment_data = treatment_data

    def display(
            self,
            show_division: bool,
            show_death: bool,
    ) -> None:
        """Displays the Division and Death Gaussians for each Treatment in the Simulation."""
        for figure, _ in self.yield_curves(show_death=show_death, show_division=show_division):
            plt.show()

    def render(
            self,
            show_division: bool,
            show_death: bool,
            folder_path: Path,
            file_name: str,
            file_extension: str,
    ) -> None:
        """Renders the Division and Death Gaussians for each Treatment in the Simulation."""
        for figure, label in self.yield_curves(show_death=show_death, show_division=show_division):
            fname = folder_path / f'{file_name}_{label}.{file_extension}'
            figure.savefig(fname)
            plt.close(figure)

    def yield_curves(
            self,
            show_death: bool,
            show_division: bool,
    ) -> Generator[tuple[plt.Figure, str], None, None]:
        """Sequentially yields gaussian Figures from the Simulation view."""
        for (colony_name, treatment_frame), treatment in self.treatment_data.items():
            fig, ax = plt.subplots()
            suffix = ''
            if show_death is True:
                suffix += 'death'
                treatment.death_curve.plot_pdf(ax=ax, x_steps=100_000, color='#E96F00', label='Death')
            if show_division is True:
                suffix += 'div'
                treatment.division_curve.plot_pdf(ax=ax, x_steps=100_000, color='#0098B1', label='Division')
            label = f'{treatment.name}_{suffix}'
            fig.suptitle(
                f'Treatment {treatment.name} added on frame {treatment_frame}'
                f'\nfor colonies starting with {colony_name}'
            )
            plt.legend()
            yield fig, label

from __future__ import annotations

from typing import Generator, TYPE_CHECKING

from matplotlib import pyplot as plt

if TYPE_CHECKING:
    from pathlib import Path
    from clovars.bio import Treatment


class TreatmentDrawer:
    """Class containing functions to draw and display Treatments."""
    def __init__(
            self,
            treatment_data: dict[tuple[str, int], Treatment],
    ) -> None:
        """Initializes a TreeDrawer instance."""
        self.treatment_data = treatment_data

    def display(
            self,
            show_division: bool,
            show_death: bool,
    ) -> None:
        """Displays the Division and Death Gaussians for each Treatment in the Simulation."""
        for figure, _ in self.yield_curves(show_death=show_death, show_division=show_division):
            plt.show()

    def render(
            self,
            show_division: bool,
            show_death: bool,
            folder_path: Path,
            file_name: str,
            file_extension: str,
    ) -> None:
        """Renders the Division and Death Gaussians for each Treatment in the Simulation."""
        for figure, label in self.yield_curves(show_death=show_death, show_division=show_division):
            fname = folder_path / f'{file_name}_{label}.{file_extension}'
            figure.savefig(fname)
            plt.close(figure)

    def yield_curves(
            self,
            show_death: bool,
            show_division: bool,
    ) -> Generator[tuple[plt.Figure, str], None, None]:
        """Sequentially yields gaussian Figures from the Simulation view."""
        for (colony_name, treatment_frame), treatment in self.treatment_data.items():
            fig, ax = plt.subplots()
            suffix = ''
            if show_death is True:
                suffix += 'death'
                treatment.death_curve.plot_pdf(ax=ax, x_steps=100_000, color='#E96F00', label='Death')
            if show_division is True:
                suffix += 'div'
                treatment.division_curve.plot_pdf(ax=ax, x_steps=100_000, color='#0098B1', label='Division')
            label = f'{treatment.name}_{suffix}'
            fig.suptitle(
                f'Treatment {treatment.name} added on frame {treatment_frame}'
                f'\nfor colonies starting with {colony_name}'
            )
            plt.legend()
            yield fig, label
Class containing functions to draw and display Treatments. Initializes a TreeDrawer instance. Displays the Division and Death Gaussians for each Treatment in the Simulation. Renders the Division and Death Gaussians for each Treatment in the Simulation. Sequentially yields gaussian Figures from the Simulation view.