Domain Entities¶

def add_ko(self, ko: KO) -> None:
    """
    Adds a KO to the sample with duplicate validation.

    Parameters
    ----------
    ko : KO
        KO to be added.

    Notes
    -----
    Duplicate KOs are automatically ignored.
    """
    if ko not in self.ko_list:
        self.ko_list.append(ko)
        logger.debug(
            "KO added to sample",
            extra={
                "sample_id": str(self.id),
                "ko": str(ko),
                "total_kos": len(self.ko_list),
            },
        )
    else:
        logger.debug(
            "Duplicate KO ignored", extra={"sample_id": str(self.id), "ko": str(ko)}
        )

def remove_ko(self, ko: KO) -> None:
    """
    Removes a KO from the sample if it exists.

    Parameters
    ----------
    ko : KO
        KO to be removed.

    Notes
    -----
    If KO does not exist in the list, operation is silently ignored.
    """
    if ko in self.ko_list:
        self.ko_list.remove(ko)

def has_ko(self, ko: KO) -> bool:
    """
    Checks if sample has a specific KO.

    Parameters
    ----------
    ko : KO
        KO to be checked.

    Returns
    -------
    bool
        True if KO is present in the sample.
    """
    return ko in self.ko_list

def get_unique_kos(self) -> List[KO]:
    """
    Returns list of unique KOs (without duplicates).

    Returns
    -------
    List[KO]
        List of unique KOs.

    Notes
    -----
    In practice, ko_list should not contain duplicates due to
    add_ko() method, but this method ensures uniqueness.
    """
    return list(set(self.ko_list))

def validate(self) -> None:
    """
    Validates entity business rules.

    Raises
    ------
    ValueError
        If sample does not have at least one KO.

    Notes
    -----
    This validation ensures the business invariant: every processed
    sample must contain at least one valid KO.
    """
    if self.ko_count == 0:
        logger.error(
            "Sample validation failed: No KOs", extra={"sample_id": str(self.id)}
        )
        raise ValueError(f"Sample {self.id} must have at least one KO")

    logger.debug(
        "Sample validation successful",
        extra={"sample_id": str(self.id), "ko_count": self.ko_count},
    )

def __str__(self) -> str:
    """
    Returns string representation of sample.

    Returns
    -------
    str
        String in format "Sample(id) with X KOs".
    """
    return f"Sample({self.id}) with {self.ko_count} KOs"

def __repr__(self) -> str:
    """
    Returns debug representation of sample.

    Returns
    -------
    str
        Detailed representation.
    """
    return (
        f"Sample(id={self.id}, "
        f"ko_count={self.ko_count}, "
        f"created_at={self.created_at})"
    )

def __eq__(self, other: object) -> bool:
    """
    Compares samples by identity (SampleId).

    Parameters
    ----------
    other : object
        Object to be compared.

    Returns
    -------
    bool
        True if both samples have the same ID.
    """
    if not isinstance(other, Sample):
        return False
    return self.id == other.id

def __hash__(self) -> int:
    """
    Hash based on sample identifier.

    Returns
    -------
    int
        Hash of SampleId.
    """
    return hash(self.id)

def add_sample(self, sample: Sample) -> None:
    """
    Adds validated sample to dataset.

    Parameters
    ----------
    sample : Sample
        Sample to be added.

    Raises
    ------
    ValueError
        If sample does not pass validation.

    Notes
    -----
    Sample is validated before being added ensuring
    only valid samples enter the dataset.
    """
    try:
        sample.validate()
        self.samples.append(sample)
        logger.debug(
            "Sample added to dataset",
            extra={
                "sample_id": str(sample.id),
                "ko_count": sample.ko_count,
                "total_samples": len(self.samples),
            },
        )
    except ValueError as e:
        logger.error(
            "Failed to add sample to dataset",
            extra={"sample_id": str(sample.id), "error": str(e)},
        )
        raise

def remove_sample(self, sample_id: SampleId) -> bool:
    """
    Removes sample from dataset by ID.

    Parameters
    ----------
    sample_id : SampleId
        ID of sample to be removed.

    Returns
    -------
    bool
        True if sample was removed, False if not found.
    """
    sample = self.get_sample_by_id(sample_id)
    if sample:
        self.samples.remove(sample)
        return True
    return False

def get_sample_by_id(self, sample_id: SampleId) -> Optional[Sample]:
    """
    Searches for sample by ID.

    Parameters
    ----------
    sample_id : SampleId
        ID of sample to be searched.

    Returns
    -------
    Optional[Sample]
        Found sample or None.
    """
    for sample in self.samples:
        if sample.id == sample_id:
            return sample
    return None

def get_all_kos(self) -> List[KO]:
    """
    Returns list of all unique KOs from dataset.

    Returns
    -------
    List[KO]
        List of unique KOs.
    """
    unique_kos = set()
    for sample in self.samples:
        unique_kos.update(sample.ko_list)
    return list(unique_kos)

def get_ko_distribution(self) -> Dict[KO, int]:
    """
    Returns KO distribution across samples.

    Returns
    -------
    Dict[KO, int]
        Dictionary mapping each KO to the number of samples
        in which it appears.

    Examples
    --------
    >>> distribution = dataset.get_ko_distribution()
    >>> distribution[KO('K00001')]
    5  # KO appears in 5 samples
    """
    distribution: Dict[KO, int] = {}
    for sample in self.samples:
        for ko in sample.ko_list:
            distribution[ko] = distribution.get(ko, 0) + 1
    return distribution

def get_samples_with_ko(self, ko: KO) -> List[Sample]:
    """
    Returns samples containing a specific KO.

    Parameters
    ----------
    ko : KO
        KO to be searched.

    Returns
    -------
    List[Sample]
        List of samples containing the KO.
    """
    return [sample for sample in self.samples if sample.has_ko(ko)]

def to_dict(self) -> Dict[str, List[str]]:
    """
    Converts dataset to dictionary format.

    Returns
    -------
    Dict[str, List[str]]
        Dictionary with format {'sample': [...], 'ko': [...]}.

    Notes
    -----
    This format is useful for later conversion to DataFrame.
    """
    records = {"sample": [], "ko": []}

    for sample in self.samples:
        for ko in sample.ko_list:
            records["sample"].append(str(sample.id))
            records["ko"].append(str(ko))

    return records

def is_empty(self) -> bool:
    """
    Checks if dataset is empty.

    Returns
    -------
    bool
        True if there are no samples.
    """
    return self.total_samples == 0

def validate(self) -> None:
    """
    Validates entire dataset.

    Raises
    ------
    ValueError
        If dataset is empty or if any sample is invalid.
    """
    if self.is_empty():
        raise ValueError("Dataset cannot be empty")

    for sample in self.samples:
        sample.validate()

def __str__(self) -> str:
    """
    Returns string representation of dataset.

    Returns
    -------
    str
        Descriptive string.
    """
    return (
        f"Dataset with {self.total_samples} samples and {self.total_kos} unique KOs"
    )

def __repr__(self) -> str:
    """
    Returns debug representation of dataset.

    Returns
    -------
    str
        Detailed representation.
    """
    return f"Dataset(samples={self.total_samples}, unique_kos={self.total_kos})"

def get_merge_status(self) -> Dict[str, bool]:
    """
    Returns the status of all merges.

    Returns
    -------
    Dict[str, bool]
        Dictionary with the status of each database.

    Examples
    --------
    >>> status = merged.get_merge_status()
    >>> status
    {
        'biorempp': True,
        'kegg': True,
        'hadeg': True,
        'toxcsm': False
    }
    """
    return {
        "biorempp": self.is_biorempp_merged,
        "kegg": self.is_kegg_merged,
        "hadeg": self.is_hadeg_merged,
        "toxcsm": self.is_toxcsm_merged,
    }

def validate(self) -> None:
    """
    Validates the merge state.

    Raises
    ------
    ValueError
        If the BioRemPP merge was not executed (mandatory).

    Notes
    -----
    Only BioRemPP is validated as it is the fundamental database.
    KEGG and HADEG may be optional depending on the context of use.
    """
    if not self.is_biorempp_merged:
        raise ValueError("BioRemPP merge is required")

def __str__(self) -> str:
    """
    Returns the string representation of the merged data.

    Returns
    -------
    str
        Descriptive string of the merge status.
    """
    status = self.get_merge_status()
    merged_count = sum(status.values())
    return f"MergedData ({merged_count}/4 databases merged)"

def __repr__(self) -> str:
    """
    Returns the debug representation of the merged data.

    Returns
    -------
    str
        Detailed representation.
    """
    return (
        f"MergedData("
        f"biorempp={self.is_biorempp_merged}, "
        f"kegg={self.is_kegg_merged}, "
        f"hadeg={self.is_hadeg_merged}, "
        f"toxcsm={self.is_toxcsm_merged})"
    )

def start(self) -> None:
    """
    Marks the analysis as started.

    Notes
    -----
    Updates status to RUNNING and records the start timestamp.
    """
    self.status = AnalysisStatus.RUNNING
    self.started_at = datetime.now()

def complete(self, metadata: Optional[Dict[str, Any]] = None) -> None:
    """
    Marks the analysis as successfully completed.

    Parameters
    ----------
    metadata : Optional[Dict[str, Any]], optional
        Result metadata, by default None.

    Notes
    -----
    Updates status to COMPLETED and records the completion timestamp.
    """
    self.status = AnalysisStatus.COMPLETED
    self.completed_at = datetime.now()
    if metadata:
        self.result_metadata.update(metadata)

def fail(self, error_message: str) -> None:
    """
    Marks the analysis as failed.

    Parameters
    ----------
    error_message : str
        Error message describing the failure.

    Notes
    -----
    Updates status to FAILED and records the error message.
    """
    self.status = AnalysisStatus.FAILED
    self.completed_at = datetime.now()
    self.error_message = error_message

def mark_from_cache(self) -> None:
    """
    Marks the analysis as loaded from cache.

    Notes
    -----
    The CACHED status indicates that the result was retrieved from
    the cache without the need for reprocessing.
    """
    self.status = AnalysisStatus.CACHED
    self.completed_at = datetime.now()

def validate(self) -> None:
    """
    Validates the analysis metadata.

    Raises
    ------
    ValueError
        If ID, name, or category are empty.
    """
    if not self.id or not self.id.strip():
        raise ValueError("Analysis ID cannot be empty")

    if not self.name or not self.name.strip():
        raise ValueError("Analysis name cannot be empty")

    if not self.category or not self.category.strip():
        raise ValueError("Analysis category cannot be empty")

def __str__(self) -> str:
    """
    Returns the string representation of the analysis.

    Returns
    -------
    str
        Descriptive string.
    """
    return f"Analysis({self.id}: {self.name}) - {self.status.value}"

def __repr__(self) -> str:
    """
    Returns the debug representation of the analysis.

    Returns
    -------
    str
        Detailed representation.
    """
    return (
        f"Analysis(id='{self.id}', "
        f"name='{self.name}', "
        f"status={self.status})"
    )