"""This module provides the functions to calculate and filter the derived traits.

This module implements only one function: `calcu_derived_trait`.
"""

import pandas as pd

from glytrait.formula import TraitFormula
from glytrait.data_type import AbundanceTable, MetaPropertyTable, DerivedTraitTable

__all__ = ["calcu_derived_trait"]


def calcu_derived_trait(
    abund_df: AbundanceTable,
    meta_prop_df: MetaPropertyTable,
    formulas: list[TraitFormula],
) -> DerivedTraitTable:
    """Calculate the derived trait values.

    Args:
        abund_df (AbundanceTable): The abundance table, with samples as index and Compositions
            as columns.
        meta_prop_df (MetaPropertyTable): The table of meta properties generated by
            `build_meta_property_table`.
        formulas (list[TraitFormula]): The trait formulas.

    Returns:
        DerivedTraitTable: The trait values, with samples as index and trait names as columns.
    """
    # The meta-property table the formulas use for initialization must have
    # the same order of index as the abundance table's columns (glycans).
    # `TraitFormula.calcu_trait` will raise an assertion error if the orders
    # of the glycans are different.
    # # (See `TraitFormula`)
    meta_prop_df_ordered = MetaPropertyTable(meta_prop_df.reindex(abund_df.columns))
    trait_series: list[pd.Series] = []
    for formula in formulas:
        formula.initialize(meta_prop_df_ordered)
        trait_series.append(formula.calcu_trait(abund_df))
    derived_trait_df = pd.concat(trait_series, axis=1)
    derived_trait_df = derived_trait_df.round(6)
    return DerivedTraitTable(derived_trait_df)
