export.py

"""Per-bacterium output: one comprehensive TSV per organism.

Each file has one row per time point and columns:
    time, biomass, conc_<metabolite> ..., flux_<reaction> ...
where concentrations are the shared extracellular pool and fluxes are that
organism's reactions. (Flux columns have no value at the final time point - a
flux describes the interval starting at that time - so the last row is blank
for fluxes.)
"""
from __future__ import annotations

import os
import math
import pandas as pd


def _pad_flux(values, n_times):
    """A flux exists for each interval (len = n_times-1); pad the last row."""
    col = list(values)
    while len(col) < n_times:
        col.append(math.nan)
    return col[:n_times]


def _frames_joint(result):
    cm = result.cm
    times = result.times
    mets = result.get_metabolites()
    bm = result.get_biomasses()
    frames = {}
    for org in cm.organisms:
        data = {"time": times, "biomass": bm[org]}
        for mid in sorted(mets):
            data[f"conc_{mid}"] = mets[mid]
        for rid in cm.organism_reactions(org):
            vals = [result.fluxes[t][rid] for t in range(len(result.fluxes))]
            data[f"flux_{rid}"] = _pad_flux(vals, len(times))
        frames[org] = pd.DataFrame(data)
    return frames


def _frames_parallel(result):
    times = result.times
    mets = result.get_metabolites()
    bm = result.get_biomasses()
    frames = {}
    for org in result.organisms:
        data = {"time": times, "biomass": bm[org]}
        for mid in sorted(mets):
            data[f"conc_{mid}"] = mets[mid]
        for rid in result.initial_bounds[org].keys():
            vals = [result.fluxes[org][t][rid] for t in range(len(result.fluxes[org]))]
            data[f"flux_{rid}"] = _pad_flux(vals, len(times))
        frames[org] = pd.DataFrame(data)
    return frames


def _frames_endpoint(result):
    cm = result.cm
    times = result.times
    mets = result.get_metabolites()
    bm = result.get_biomasses()
    frames = {}
    for org in cm.organisms:
        data = {"time": times, "biomass": bm[org]}
        for mid in sorted(mets):
            data[f"conc_{mid}"] = mets[mid]
        for rid in cm.organism_reactions(org):
            data[f"flux_{rid}"] = _pad_flux(result.flux_series(rid), len(times))
        frames[org] = pd.DataFrame(data)
    return frames


def per_organism_frames(result) -> dict[str, "pd.DataFrame"]:
    """Return {organism_id: DataFrame} of comprehensive per-organism output."""
    from .joint import DynamicJointFBA
    from .parallel import DynamicParallelFBA
    from .endpoint import EndPointFBA
    if isinstance(result, DynamicJointFBA):
        return _frames_joint(result)
    if isinstance(result, DynamicParallelFBA):
        return _frames_parallel(result)
    if isinstance(result, EndPointFBA):
        return _frames_endpoint(result)
    raise TypeError(f"don't know how to export {type(result).__name__}")


def export_tsv(result, out_dir=".", prefix="") -> list[str]:
    """Write one comprehensive TSV per organism; return the file paths."""
    os.makedirs(out_dir, exist_ok=True)
    paths = []
    for org, df in per_organism_frames(result).items():
        path = os.path.join(out_dir, f"{prefix}{org}.tsv")
        df.to_csv(path, sep="\t", index=False)
        paths.append(path)
    return paths