"""Implementation of FbaResults class.

Support functions to analyze FBA and TFA results.
TFA models constructred by f2xba package.

Peter Schubert, HHU Duesseldorf, CCB, November 2024
"""

import re
import pandas as pd
from collections import defaultdict

import f2xba.prefixes as pf
from .results import Results


class FbaResults(Results):
    """Process optimization solution for FBA and TFA models.

    Using the FbaOptimization class, one or several model optimization solutions can be determined,
    e.g. across different media conditions. Solutions can be collected in a dictionary.

    Example: Prepare the processing of optimization solutions for two different conditions.

    .. code-block:: python

        fr = FbaResults(fo, {'glucose': solution_glc, 'acetate': solution_ac})
    """


    def __init__(self, optim, results):
        """Instantiate the FbaResults instance.

        :param optim: FbaOptimization instance
        :type optim: :class:`FbaOptimization`
        :param dict results: conditions with respective optimization solutions (:class:`Solution`)
        """
        super().__init__(optim, results)

    def get_fluxes(self, solution):
        """Extract fluxes of an optimization solution.

        :meta private:
        """
        fluxes = {}
        for rid, mmol_per_gdwh in solution.fluxes.items():
            if re.match(pf.V_, rid) is None:
                reaction_str = self.optim.rdata[rid]['reaction_str']
                gpr = self.optim.rdata[rid]['gpr']
                groups = self.optim.rdata[rid]['groups']
                fluxes[rid] = [reaction_str, rid, gpr, groups, mmol_per_gdwh, abs(mmol_per_gdwh)]
        cols = ['reaction_str', 'net_rid', 'gpr', 'groups', 'mmol_per_gDWh', 'abs mmol_per_gDWh']
        df_fluxes = pd.DataFrame(fluxes.values(), index=list(fluxes), columns=cols)
        df_fluxes.index.name = 'reaction'
        return df_fluxes

    def get_net_fluxes(self, solution):
        """Extract net fluxes of an optimization solution.

        :meta private:
        """
        net_fluxes = defaultdict(float)
        for rid, val in solution.fluxes.items():
            if re.match(pf.V_, rid) is None:
                net_rid = re.sub('_REV$', '', rid)
                if re.search('_REV$', rid):
                    net_fluxes[net_rid] -= val
                else:
                    net_fluxes[net_rid] += val

        # add reaction string and gene product relations
        net_flux_data = {}
        for rid, mmol_per_gdwh in net_fluxes.items():
            rdata = self.optim.net_rdata.get(rid)
            if rdata:
                net_flux_data[rid] = [rdata['reaction_str'], rdata['gpr'],
                                      rdata['groups'], mmol_per_gdwh, abs(mmol_per_gdwh)]
            else:
                net_flux_data[rid] = ['', '', '', mmol_per_gdwh, abs(mmol_per_gdwh)]

        cols = ['reaction_str', 'gpr', 'groups', 'mmol_per_gDWh', 'abs mmol_per_gDWh']
        df_net_fluxes = pd.DataFrame(net_flux_data.values(), index=list(net_flux_data), columns=cols)
        df_net_fluxes.index.name = 'rid'
        return df_net_fluxes

    def collect_protein_results(self):
        """Return empty DataFrame.

        :meta private:
        """
        return pd.DataFrame()

    def get_predicted_protein_data(self, solution):
        """Return empty DataFram.

        :meta private:
        """
        return pd.DataFrame()
