process_AD.py

import os as _os
import numpy as _np
import pandas as _pd
from collections import defaultdict as _defaultdict
from progressbar import progressbar as _progressbar
from .utils import align_stat_data as _align_stat_data
from .process_AD_helpers import (
    PROTEIN as _PROTEIN,
    DIMER as _DIMER,
    _apply_rxn_update,
    __read_A_D_line,
    read_survP_from_assoc as _read_survP_from_assoc
)


__all__ = [
    'read_assoc_dissoc_freq',
    'readResT_from_NERDSS',
    'readInnerStateTransitions', 
    'read_results_from_assoc_cache',
    'read_remaining_copies'
]


def readResT_from_NERDSS(
        file_assoc_dissoc_times:str, dt:float, proteins=['P'], substrates=['S','N'],
        debug:bool=False, startT:float=0
):
    """Read residence time and search time from assoc_dissoc_time.dat. Return numpy arrays of residence time and search time
    
    Input:
        file_assoc_dissoc_times: path to assoc_dissoc_time.dat
        dt: time step (s)
        proteins: list of protein names
        substrates: list of substrate names
        debug: print out the information of each protein
    Output:
        resTimesAll: numpy array of residence time
        searchTimesAll: numpy array of search time
    """

    def __maxComplexID(complexList):
        if not complexList:
            # if the complex list is empty
            return 0
        else:
            return max(complexList.keys())
    
    def __add_newComplex(complexList, protein):
        # create a new multimer
        newMultimerID = __maxComplexID(complexList) + 1
        newMultimer = _DIMER(newMultimerID)
        newMultimer.addComponent(protein, currt=-1, record=False)
        complexList[newMultimerID] = newMultimer
    
    def __add_protein_to_list(proteinList, complexList, pro, proID):
        if proID not in proteinList:
            # create a new protein
            proteinList[proID] = _PROTEIN(pro, proID)
            # add this protein to complex list, the id is the next integer
            __add_newComplex(complexList, proteinList[proID])

    with open(file_assoc_dissoc_times, 'r') as f:
        complexList:dict[int,_DIMER] = {} # complex id: DIMER
        proteinList:dict[int,_PROTEIN] = {} # protein id: PROTEIN
        for line in f:
            if debug: print(line.strip())
            rxntype, rxnInfo = __read_A_D_line(line, dt, proteins, substrates)
            if debug: print(rxntype, rxnInfo)
            currt = rxnInfo[-1]
            # always apply updates so pre-start structure/state is correct,
            # but only record residence/search times for events at/after startT
            currt = _apply_rxn_update(
                rxntype, rxnInfo, proteinList, complexList, __add_protein_to_list, 
                __add_newComplex, debug, record=(currt >= startT)
            )
            if debug and rxntype == 'rxnPRO':
                pro1, pro1ID, pro2, pro2ID, _, _ = rxnInfo
                print(pro1ID, 'startT: %.3f'%proteinList[pro1ID].boundStartTime, 'endT: %.3f'%proteinList[pro1ID].boundEndTime)
                print(pro2ID, 'startT: %.3f'%proteinList[pro2ID].boundStartTime, 'endT: %.3f'%proteinList[pro2ID].boundEndTime)
            elif debug and rxntype == 'rxnSUB':
                pro, proID, *_ = rxnInfo
                for memberPro in complexList[proteinList[proID].multimerID].components:
                    print('startT: %.3f'%memberPro.boundStartTime, 'endT: %.3f'%memberPro.boundEndTime)
        
        if debug: print()
    # finished reading the file
    # get the residence time and search time of each protein
    resTimesAll = []
    searchTimesAll = []
    for proID in proteinList:
        if debug: print(proID, proteinList[proID].resTimeList)
        resTimesAll.extend(proteinList[proID].resTimeList)
        searchTimesAll.extend(proteinList[proID].searchTimeList)
    # return numpy arrays of residence time and search time
    return _np.array(resTimesAll), _np.array(searchTimesAll)


def readInnerStateTransitions(
    file_assoc_dissoc_times: str,
    dt: float,
    proteins=['P'],
    substrates=['S','N'],
    startT: float = 0.0,
    exclude_pp_to_ppn: bool = True,
    debug: bool = False,
):
    def __maxComplexID(complexList):
        return max(complexList.keys()) if complexList else 0
    def __add_newComplex(complexList, protein):
        newMultimerID = __maxComplexID(complexList) + 1
        newMultimer = _DIMER(newMultimerID)
        newMultimer.addComponent(protein, currt=-1, record=False)
        complexList[newMultimerID] = newMultimer
    def __add_protein_to_list(proteinList, complexList, pro, proID):
        if proID not in proteinList:
            proteinList[proID] = _PROTEIN(pro, proID)
            __add_newComplex(complexList, proteinList[proID])

    complexList: dict[int, _DIMER] = {}
    proteinList: dict[int, _PROTEIN] = {}
    last_state = {}
    last_time = {}
    transitions = []

    with open(file_assoc_dissoc_times, 'r') as f:
        started = False
        for line in f:
            rxntype, rxnInfo = __read_A_D_line(line, dt, proteins, substrates)
            currt = rxnInfo[-1]
            # always apply updates so pre-start complexes/states are correct
            currt = _apply_rxn_update(rxntype, rxnInfo, proteinList, complexList, __add_protein_to_list, __add_newComplex, debug, record=(currt >= startT))
            # once we hit startT, initialize baseline states for existing dimers at startT
            if currt >= startT and not started:
                started = True
                for cid, cmplx in list(complexList.items()):
                    if len(cmplx.components) != 2:
                        continue
                    # compute state at startT
                    substrates_bound = cmplx.substrates
                    nbound_S = substrates_bound.count('S')
                    nbound_N = substrates_bound.count('N')
                    if nbound_S == 1:
                        if nbound_N == 0:
                            state = 'PPS'
                        elif nbound_N == 1:
                            state = 'PSPN'
                        else:
                            state = 'PNPSN'
                    else:
                        if nbound_N == 0:
                            state = 'PP'
                        elif nbound_N == 1:
                            state = 'PPN'
                        else:
                            state = 'PNPN'
                    last_state[cid] = state
                    last_time[cid] = startT
            # only evaluate & record transitions for times at/after startT
            if currt < startT:
                continue
            # evaluate inner state of dimers only
            for cid, cmplx in list(complexList.items()):
                # only consider dimers (2 proteins)
                if len(cmplx.components) != 2:
                    last_state.pop(cid, None)
                    last_time.pop(cid, None)
                    continue
                # get the substrates this complex binds to
                substrates_bound = cmplx.substrates
                # map to PP, PPN, PNPN, PPS, PSPN, only consider these states
                nbound_S = substrates_bound.count('S')
                nbound_N = substrates_bound.count('N')
                if nbound_S == 1:
                    if nbound_N == 0:
                        state = 'PPS'
                    elif nbound_N == 1:
                        state = 'PSPN'
                    else:
                        state = 'PNPSN' # both N bound
                else:
                    # no S bound
                    if nbound_N == 0:
                        state = 'PP'
                    elif nbound_N == 1:
                        state = 'PPN'
                    else:
                        state = 'PNPN'  # both N bound
                # record transition if state changed
                if cid not in last_state:
                    last_state[cid] = state
                    last_time[cid] = currt
                elif state != last_state[cid]:
                    fr, to = last_state[cid], state
                    t0, t1 = last_time[cid], currt
                    if not (exclude_pp_to_ppn and fr == 'PP'):
                        transitions.append((fr, to, t0, t1, t1 - t0))
                    last_state[cid] = state
                    last_time[cid] = currt
    return transitions


def read_results_from_assoc_cache(
    workdir, 
    dt=None, 
    t_eval=None, 
    indexList=[], 
    n_resample=500,
    n_reactions=1,
    cached=True, 
    cachFileName='/survival.csv'
):
    """
    Read or compute survival probability results with caching.
    
    Parameters
    ----------
    workdir : str
        Working directory
    dt : float, optional
        Time step in seconds (required if cached=False)
    t_eval : array-like, optional
        Time points for evaluation in seconds (required if cached=False)
    indexList : list, default=[]
        List of simulation indices
    n_resample : int, default=500
        Number of bootstrap resamples
    n_reactions : int, default=1
        Total number of expected association reactions.
        For multi-body problems with N irreversible associations, set n_reactions=N.
    cached : bool, default=True
        Whether to read from cache or recompute
    cachFileName : str, default='/survival.csv'
        Cache file name
        
    Returns
    -------
    t_eval : array
        Time points
    obs_mean : array
        Mean observed survival probability
    obs_ste : array
        Standard error of observed survival probability
    """
    cachefile = workdir + cachFileName
    
    if not cached:
        print("Computing survival probability and caching results...")
        # compute survival probability
        obs_mean, obs_ste = _read_survP_from_assoc(workdir, dt, t_eval, indexList, n_resample, n_reactions)
        # cache the results
        with open(cachefile, 'w') as f:
            f.write('t(s),survP_obs,stdErr_obs\n')
            for t, mean, err in zip(t_eval, obs_mean, obs_ste):
                f.write(f'{t},{mean},{err}\n')
        print(f"Results cached to {cachefile}. Final time point: {t_eval[-1]} s")
    else:
        try:
            df = _pd.read_csv(cachefile)
            t_eval = df['t(s)'].values
            obs_mean = df['survP_obs'].values
            obs_ste = df['stdErr_obs'].values
        except Exception as e:
            print(f"Cache read failed: {e}. Recomputing...")
            return read_results_from_assoc_cache(
                workdir, dt, t_eval, indexList, n_resample, n_reactions,
                cached=False, cachFileName=cachFileName
            )

    return t_eval, obs_mean, obs_ste


def read_remaining_copies_from_assoc(
    file_assoc_dissoc_times: str,
    species: str,
    total_copies: int,
    dt: float | None = None,
    debug: bool = False,
):
    """
    Compute remaining (unbound) copies of a given species over time.

    Parameters
    ----------
    file_assoc_dissoc_times : str
        Path to assoc_dissoc_time.dat
    species : str
        Species name (e.g., 'AF')
    total_copies : int
        Total number of copies for the species
    dt : float | None, optional
        Time step in seconds. If provided, time = ITR * dt, else time = ITR
    debug : bool, default=False
        Print per-event updates

    Returns
    -------
    times : np.ndarray
        Event times (ITR*dt if dt is provided, else ITR)
    remaining : np.ndarray
        Remaining (unbound) copies after each event time
    """
    remaining = total_copies
    times = [0.0]
    remaining_list = [total_copies]

    with open(file_assoc_dissoc_times, 'r') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            parts = [p.strip() for p in line.split(',')]
            if len(parts) < 3:
                continue

            # Parse ITR
            itr_token = parts[0]
            if not itr_token.startswith('ITR:'):
                continue
            try:
                itr = int(itr_token.split(':', 1)[1])
            except ValueError:
                continue

            time_val = itr * dt if dt is not None else itr

            rxn = parts[1].upper()
            if rxn not in {'BOND', 'BREAK'}:
                continue

            # Species appears in either reactant position in this file format
            if species not in parts:
                continue

            prev_remaining = remaining
            if rxn == 'BOND':
                remaining -= 1
            else:
                remaining += 1

            if remaining < 0:
                raise ValueError(
                    f"Remaining copies for {species} became negative at {itr_token}."
                )
            if remaining > total_copies:
                remaining = total_copies

            # enforce a step change for interpolation alignment
            if time_val == times[-1]:
                remaining_list[-1] = remaining
                continue
            if time_val > times[-1]:
                times.append(time_val)
                remaining_list.append(prev_remaining)

            times.append(time_val)
            remaining_list.append(remaining)

            if debug:
                print(f"{itr_token},{rxn} -> remaining {species}: {remaining}")

    return _np.array(times), _np.array(remaining_list)


def read_remaining_copies(
    parent_folder: str,
    iter_folders: list[str],
    species: str,
    total_copies: int,
    target_time_points: _np.ndarray,
    dt: float | None = None,
    assoc_rel_path: str = 'DATA/assoc_dissoc_time.dat',
    debug: bool = False,
    show_progress: bool = False,
    truncation: bool = True,
):
    """
    Read remaining copies from multiple trajectories and merge them.

    Each trajectory is aligned to the same target time points using
    `utils.align_stat_data`, then mean and standard deviation are computed.

    Parameters
    ----------
    parent_folder : str
        Parent folder that contains trajectory subfolders
    iter_folders : list[str]
        List of trajectory folder names (e.g., ['0','1','2'])
    species : str
        Species name (e.g., 'AF')
    total_copies : int
        Total number of copies for the species
    target_time_points : np.ndarray
        Common time points for alignment
    dt : float | None, optional
        Time step in seconds. If provided, time = ITR * dt, else time = ITR
    assoc_rel_path : str, default='DATA/assoc_dissoc_time.dat'
        Relative path to assoc_dissoc_time.dat inside each trajectory folder
    debug : bool, default=False
        Print per-event updates
    show_progress : bool, default=False
        Show a progress bar over trajectories.
    truncation : bool, default=True
        If True, truncate target_time_points to the earliest trajectory
        end time.

    Returns
    -------
    target_time_points : np.ndarray
        Time points used for alignment
    mean_remaining : np.ndarray
        Mean remaining copies across trajectories
    std_remaining : np.ndarray
        Standard error across trajectories
    """
    # First pass: read all trajectories and find the earliest end time
    raw_trajectories = []
    _iter = _progressbar(iter_folders) if show_progress else iter_folders

    for itr in _iter:
        assoc_path = _os.path.join(parent_folder, str(itr), assoc_rel_path)
        if not _os.path.isfile(assoc_path):
            raise FileNotFoundError(f"Missing assoc file: {assoc_path}")

        t_points, remaining = read_remaining_copies_from_assoc(
            assoc_path,
            species,
            total_copies,
            dt=dt,
            debug=debug,
        )
        raw_trajectories.append((t_points, remaining))

    if len(raw_trajectories) == 0:
        raise ValueError("No trajectories provided.")

    # Truncate target time points to the earliest trajectory end time
    if truncation:
        earliest_end = min(t[-1] for t, _ in raw_trajectories)
        target_time_points = target_time_points[target_time_points <= earliest_end]

    # Second pass: align each trajectory to the (possibly truncated) time points
    aligned_all = []
    for t_points, remaining in raw_trajectories:
        aligned_y, _ = _align_stat_data(
            t_points,
            remaining,
            _np.zeros_like(remaining),
            target_time_points,
        )
        aligned_all.append(aligned_y)

    aligned_all = _np.vstack(aligned_all)
    mean_remaining = _np.mean(aligned_all, axis=0)
    std_remaining = _np.std(aligned_all, axis=0) / _np.sqrt(aligned_all.shape[0])

    return target_time_points, mean_remaining, std_remaining


def calc_assoc_dissoc_freq(filepath, dt=1.0, start_t=0.0, end_t=_np.inf, track_mol=None,
                           return_durations=False):
    """Calculate association and dissociation frequencies from a NERDSS
    ``assoc_dissoc_time.dat`` file.

    Definitions
    -----------
    * **Free (unbound) duration** – The time between a molecule's BREAK event
      and its next BOND event.  This is the "search time".
    * **Bound duration** – The time between a molecule's BOND event and its
      next BREAK event.

    * **Association frequency** = 1 / mean(free durations)
    * **Dissociation frequency** = 1 / mean(bound durations)

    Parameters
    ----------
    filepath : str
        Path to the ``assoc_dissoc_time.dat`` file.
    dt : float, optional
        Physical time per iteration in **seconds** (``timeStep`` in
        ``parms.inp``).  Default is 1.0.
    start_t : float, optional
        Discard events whose physical time (in **seconds**) is earlier
        than this value.  Useful for ignoring the transient /
        pre-equilibrium period.  Default is 0.0.
    end_t : float, optional
        Discard events whose physical time (in **seconds**) is at or
        after this value.  Default is ``np.inf`` (use all data).
    track_mol : str or None, optional
        Molecule type to track (e.g. ``'AF'``).  If *None*, the first
        molecule type encountered in the file is used.
    return_durations : bool, optional
        If *True*, include ``'free_durations'`` and ``'bound_durations'``
        arrays in the returned dict.  Default is *False* to save RAM.

    Returns
    -------
    dict
        ``'assoc_freq'``  – float, mean association frequency (1/s).
        ``'dissoc_freq'`` – float, mean dissociation frequency (1/s).
        ``'free_durations'``  – np.ndarray, all individual free durations (s).
            Only present when ``return_durations=True``.
        ``'bound_durations'`` – np.ndarray, all individual bound durations (s).
            Only present when ``return_durations=True``.
    """
    # ------------------------------------------------------------------
    # 1. Parse the file
    # ------------------------------------------------------------------
    # Line format: ITR:<itr>,<BOND|BREAK>,<type1>,<id1>,<site1>,<type2>,<id2>,<site2>
    mol_events = _defaultdict(list)   # (mol_type, mol_id) -> [(time_s, event), ...]
    first_mol_type = None

    with open(filepath, 'r') as fh:
        for line in fh:
            line = line.strip()
            if not line:
                continue
            parts = line.split(',')
            itr        = int(parts[0].split(':')[1])
            time_s     = itr * dt              # physical time in seconds
            event_type = parts[1]              # 'BOND' or 'BREAK'
            mol_type1  = parts[2]
            mol_id1    = int(parts[3])
            # parts[4] is site1
            mol_type2  = parts[5]
            mol_id2    = int(parts[6])
            # parts[7] is site2

            if first_mol_type is None:
                first_mol_type = mol_type1

            # Record events for both molecules involved (time in seconds)
            mol_events[(mol_type1, mol_id1)].append((time_s, event_type))
            mol_events[(mol_type2, mol_id2)].append((time_s, event_type))

    # Decide which molecule type to analyse
    if track_mol is None:
        track_mol = first_mol_type

    # ------------------------------------------------------------------
    # 2. Compute durations per molecule
    # ------------------------------------------------------------------
    bound_durations = []
    free_durations  = []

    for (mol_type, mol_id), evts in mol_events.items():
        if mol_type != track_mol:
            continue

        evts.sort(key=lambda x: x[0])  # chronological order

        for i in range(len(evts) - 1):
            t_curr, type_curr = evts[i]
            t_next, type_next = evts[i + 1]

            # skip events outside the [start_t, end_t) window (seconds)
            if t_curr < start_t or t_curr >= end_t:
                continue

            duration = t_next - t_curr  # duration in seconds

            if type_curr == 'BOND' and type_next == 'BREAK':
                # molecule was bound for this duration
                bound_durations.append(duration)
            elif type_curr == 'BREAK' and type_next == 'BOND':
                # molecule was free (searching) for this duration
                free_durations.append(duration)
            # BOND->BOND or BREAK->BREAK pairs are ignored (partner swap
            # at same time step, etc.)

    bound_durations = _np.array(bound_durations)
    free_durations  = _np.array(free_durations)

    # ------------------------------------------------------------------
    # 3. Frequencies (in 1/s)
    # ------------------------------------------------------------------
    assoc_freq  = (1.0 / _np.mean(free_durations)
                   if len(free_durations) > 0 else _np.nan)
    dissoc_freq = (1.0 / _np.mean(bound_durations)
                   if len(bound_durations) > 0 else _np.nan)

    out = {
        'assoc_freq':  assoc_freq,
        'dissoc_freq': dissoc_freq,
    }
    if return_durations:
        out['free_durations']  = free_durations
        out['bound_durations'] = bound_durations

    return out


def read_assoc_dissoc_freq(
    parent_folder: str,
    iter_folders: list[str],
    dt: float = 1.0,
    start_t: float = 0.0,
    end_t: float = _np.inf,
    track_mol: str | None = None,
    assoc_rel_path: str = 'DATA/assoc_dissoc_time.dat',
    return_durations: bool = False,
    show_progress: bool = False,
):
    """Compute association / dissociation frequencies across multiple
    trajectories.

    Calls :func:`calc_assoc_dissoc_freq` on each trajectory, pools the
    individual durations, and returns aggregate statistics.

    Parameters
    ----------
    parent_folder : str
        Parent folder that contains trajectory subfolders.
    iter_folders : list[str]
        List of trajectory folder names (e.g. ``['0', '1', '2']``).
    dt : float, optional
        Physical time per iteration in **seconds**.  Default is 1.0.
    start_t : float, optional
        Discard events before this time in **seconds**.  Default is 0.0.
    end_t : float, optional
        Discard events at or after this time in **seconds**.
        Default is ``np.inf``.
    track_mol : str or None, optional
        Molecule type to track (e.g. ``'AF'``).  If *None*, the first
        molecule type encountered in each file is used.
    assoc_rel_path : str, default ``'DATA/assoc_dissoc_time.dat'``
        Relative path to ``assoc_dissoc_time.dat`` inside each trajectory
        folder.
    return_durations : bool, optional
        If *True*, include pooled ``'free_durations_all'`` and
        ``'bound_durations_all'`` arrays in the returned dict.
        Default is *False* to save RAM.
    show_progress : bool, optional
        Show a progress bar over trajectories.  Default is *False*.

    Returns
    -------
    dict
        ``'assoc_freq_mean'``   – float, mean association frequency (1/s)
        across trajectories.
        ``'assoc_freq_ste'``    – float, standard error of the per-trajectory
        association frequencies.
        ``'dissoc_freq_mean'``  – float, mean dissociation frequency (1/s)
        across trajectories.
        ``'dissoc_freq_ste'``   – float, standard error of the per-trajectory
        dissociation frequencies.
        ``'assoc_freq_all'``    – np.ndarray, per-trajectory association
        frequencies.
        ``'dissoc_freq_all'``   – np.ndarray, per-trajectory dissociation
        frequencies.
        ``'free_durations_all'``  – np.ndarray, pooled free durations (s)
        from all trajectories.  Only present when ``return_durations=True``.
        ``'bound_durations_all'`` – np.ndarray, pooled bound durations (s)
        from all trajectories.  Only present when ``return_durations=True``.
    """
    assoc_freqs  = []
    dissoc_freqs = []
    all_free     = [] if return_durations else None
    all_bound    = [] if return_durations else None
    _iter = _progressbar(iter_folders) if show_progress else iter_folders

    for itr in _iter:
        assoc_path = _os.path.join(parent_folder, str(itr), assoc_rel_path)
        if not _os.path.isfile(assoc_path):
            raise FileNotFoundError(f"Missing assoc file: {assoc_path}")

        result = calc_assoc_dissoc_freq(
            assoc_path, dt=dt, start_t=start_t, end_t=end_t,
            track_mol=track_mol, return_durations=return_durations,
        )

        assoc_freqs.append(result['assoc_freq'])
        dissoc_freqs.append(result['dissoc_freq'])
        if return_durations:
            all_free.append(result['free_durations'])
            all_bound.append(result['bound_durations'])

    assoc_freqs  = _np.array(assoc_freqs)
    dissoc_freqs = _np.array(dissoc_freqs)

    n = len(iter_folders)

    out = {
        'assoc_freq_mean':  _np.nanmean(assoc_freqs),
        'assoc_freq_ste':   _np.nanstd(assoc_freqs) / _np.sqrt(n),
        'dissoc_freq_mean': _np.nanmean(dissoc_freqs),
        'dissoc_freq_ste':  _np.nanstd(dissoc_freqs) / _np.sqrt(n),
        'assoc_freq_all':   assoc_freqs,
        'dissoc_freq_all':  dissoc_freqs,
    }
    if return_durations:
        out['free_durations_all']  = _np.concatenate(all_free)  if all_free  else _np.array([])
        out['bound_durations_all'] = _np.concatenate(all_bound) if all_bound else _np.array([])

    return out