tcMultiTestMAT.py

"""
Test script for multi-peak detection on G394 .mat files.
Directory structure expected:
    root/<date>/<sub>/*.mat

Reports per-session peak counts, a table of all cells with at least one
significant peak, and a combined heatmap (R2B | TI) across all sessions.
"""

import os
import glob
import sys
import numpy as np
import scipy.io
import natsort
import matplotlib.pyplot as plt

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import pyBindMat as pb

ROOT = "/home/bhalla/homework/Hrishi/2025/DATA/G394_TenFiles_13Oct"


def load_params(data):
    td      = data["trialDetails"]
    tcpy    = data["ops"]["tcpy"]
    sr      = float(td["samplingRate"])
    frameDt = 1.0 / sr
    csOnset = int(round(float(td["preDuration"]) * sr))
    usOnset = int(round(csOnset
                        + float(td["csDuration"])    * sr
                        + float(td["traceDuration"]) * sr))
    circPad       = int(tcpy["circPadFrames"])
    binFrames     = int(tcpy["binFrames"])
    numShuffle    = int(tcpy["numShuffle"])
    r2bThresh     = float(tcpy["r2bThresh"])
    r2bPercentile = float(tcpy["r2bPercentile"])
    transThresh   = float(tcpy["transientThresh"])
    tiPercentile  = float(tcpy["tiPercentile"])
    fracThresh    = float(tcpy["fracTrialsFiredThresh"])
    return (csOnset, usOnset, circPad, binFrames, numShuffle,
            r2bThresh, r2bPercentile, transThresh, tiPercentile, fracThresh, frameDt)


def peak_times_s(indices, circPad, frameDt, binFrames=1):
    return [(idx * binFrames - circPad) * frameDt for idx in indices]


def session_summary(results, method):
    counts = {0: 0, 1: 0, 2: 0, 3: 0}
    for r in results:
        counts[min(len(r['allPkIndices']), 3)] += 1
    print(f"  {method}: 0 peaks={counts[0]}  1 peak={counts[1]}  "
          f"2 peaks={counts[2]}  3+ peaks={counts[3]}")


def _heatmap_ax(ax, traces, pk_indices, time_axis, method, us_time):
    """Render one normalised heatmap into ax. Returns False if no data."""
    if not traces:
        ax.text(0.5, 0.5, f'No significant\n{method} cells',
                ha='center', va='center', transform=ax.transAxes, fontsize=11)
        ax.set_title(method)
        return False

    order = np.argsort(pk_indices)
    mat   = np.array([traces[i] for i in order], dtype=float)
    rmin  = mat.min(axis=1, keepdims=True)
    rmax  = mat.max(axis=1, keepdims=True)
    denom = rmax - rmin
    denom[denom == 0] = 1.0
    mat   = (mat - rmin) / denom

    ncells = mat.shape[0]
    dt = time_axis[1] - time_axis[0]
    im = ax.imshow(mat, aspect='auto', origin='upper', cmap='hot', vmin=0, vmax=1,
                   extent=[time_axis[0] - dt/2, time_axis[-1] + dt/2, ncells, 0])
    ax.axvline(0.0,     color='cyan', lw=1.0, ls='--', label='CS onset')
    ax.axvline(us_time, color='lime', lw=1.0, ls='--', label='US onset')
    ax.legend(loc='upper right', fontsize=7)
    ax.set_xlabel('Time re CS onset (s)')
    ax.set_ylabel(f'Cell (sorted by 1st peak)  n={ncells}')
    ax.set_title(method)
    plt.colorbar(im, ax=ax, label='Norm. ΔF/F', fraction=0.03, pad=0.02)
    return True


def plot_heatmaps(r2b_traces, r2b_pk, r2b_time,
                  ti_traces,  ti_pk,  ti_time,
                  r2b_traces_mp, r2b_pk_mp,
                  ti_traces_mp,  ti_pk_mp,
                  us_time, title):
    n_all = max(len(r2b_traces),    len(ti_traces),    1)
    n_mp  = max(len(r2b_traces_mp), len(ti_traces_mp), 1)
    fig, axes = plt.subplots(
        2, 2, figsize=(14, n_all * 0.12 + n_mp * 0.12 + 3.0),
        gridspec_kw={'height_ratios': [n_all, n_mp]})
    fig.suptitle(title)
    _heatmap_ax(axes[0, 0], r2b_traces,    r2b_pk,    r2b_time, 'R2B — all significant', us_time)
    _heatmap_ax(axes[0, 1], ti_traces,     ti_pk,     ti_time,  'TI — all significant',  us_time)
    _heatmap_ax(axes[1, 0], r2b_traces_mp, r2b_pk_mp, r2b_time, 'R2B — 2+ peaks',        us_time)
    _heatmap_ax(axes[1, 1], ti_traces_mp,  ti_pk_mp,  ti_time,  'TI — 2+ peaks',         us_time)
    plt.tight_layout()
    plt.show()


def main():
    dates = natsort.natsorted(
        d for d in os.listdir(ROOT)
        if os.path.isdir(os.path.join(ROOT, d))
    )

    total_cells = 0
    total_r2b   = {0: 0, 1: 0, 2: 0, 3: 0}
    total_ti    = {0: 0, 1: 0, 2: 0, 3: 0}
    table_rows  = []

    # Accumulated across all sessions for the combined heatmap
    r2b_traces,    r2b_pk,    r2b_time = [], [], None
    ti_traces,     ti_pk,     ti_time  = [], [], None
    r2b_traces_mp, r2b_pk_mp           = [], []   # 2+ peaks only
    ti_traces_mp,  ti_pk_mp            = [], []
    us_time_s = None

    for date in dates:
        for sub in ["1", "2"]:
            mat_path = os.path.join(ROOT, date, sub)
            if not os.path.isdir(mat_path):
                continue
            matfiles = glob.glob(os.path.join(mat_path, "*.mat"))
            if not matfiles:
                continue

            try:
                data = scipy.io.loadmat(matfiles[0], simplify_cells=True)
            except Exception as e:
                print(f"  Could not load {matfiles[0]}: {e}")
                continue

            dfbf = data["dfbf_clean"]
            (csOnset, usOnset, circPad, binFrames, numShuffle,
             r2bThresh, r2bPercentile, transThresh, tiPercentile,
             fracThresh, frameDt) = load_params(data)
            circShuff = dfbf.shape[2] - (csOnset - circPad)

            print(f"\n{date}/{sub}  [{dfbf.shape[0]} cells, {dfbf.shape[1]} trials, "
                  f"{dfbf.shape[2]} frames]  csOnset={csOnset} usOnset={usOnset} "
                  f"frameDt={frameDt:.4f}")

            r2b_res = pb.runR2Banalysis(dfbf, csOnset, usOnset, circPad, circShuff,
                                        binFrames, numShuffle, r2bThresh, r2bPercentile,
                                        frameDt=frameDt)
            ti_res  = pb.runTIanalysis(dfbf, csOnset, usOnset, circPad, circShuff,
                                       binFrames, numShuffle, transThresh, tiPercentile,
                                       fracThresh, frameDt)

            session_summary(r2b_res, "R2B")
            session_summary(ti_res,  "TI ")

            total_cells += dfbf.shape[0]
            for r in r2b_res:
                total_r2b[min(len(r['allPkIndices']), 3)] += 1
            for r in ti_res:
                total_ti[min(len(r['allPkIndices']), 3)] += 1

            # Table rows
            for cell_idx, (rr, tt) in enumerate(zip(r2b_res, ti_res)):
                r2b_pks = list(zip(
                    peak_times_s(rr['allPkIndices'], circPad, frameDt, binFrames=1),
                    rr['allPkScores'], rr['allPkPvalues']))
                ti_pks  = list(zip(
                    peak_times_s(tt['allPkIndices'], circPad, frameDt, binFrames=binFrames),
                    tt['allPkScores'], tt['allPkPvalues']))
                if r2b_pks or ti_pks:
                    table_rows.append((date, sub, cell_idx, r2b_pks, ti_pks))

            # Accumulate traces for combined heatmap
            this_r2b_time = np.array([(ff - circPad) * frameDt for ff in range(circShuff)])
            numBins = circShuff // binFrames
            this_ti_time  = np.array([(bb * binFrames - circPad) * frameDt for bb in range(numBins)])

            if r2b_time is None:
                r2b_time = this_r2b_time
                ti_time  = this_ti_time
                us_time_s = (usOnset - csOnset) * frameDt

            for r in r2b_res:
                n = len(r['allPkIndices'])
                if n >= 1:
                    trace = np.array(r['meanTrace'])
                    if len(trace) == len(r2b_time):
                        r2b_traces.append(trace)
                        r2b_pk.append(r['allPkIndices'][0])
                        if n >= 2:
                            r2b_traces_mp.append(trace)
                            r2b_pk_mp.append(r['allPkIndices'][0])
            for r in ti_res:
                n = len(r['allPkIndices'])
                if n >= 1:
                    trace = np.array(r['meanTrace'])
                    if len(trace) == len(ti_time):
                        ti_traces.append(trace)
                        ti_pk.append(r['allPkIndices'][0])
                        if n >= 2:
                            ti_traces_mp.append(trace)
                            ti_pk_mp.append(r['allPkIndices'][0])

    # ── Summary ──────────────────────────────────────────────────────────────
    print(f"\n=== Totals across all sessions ===")
    print(f"  Cells analysed : {total_cells}")
    print(f"  R2B  —  0 peaks: {total_r2b[0]}  1 peak: {total_r2b[1]}  "
          f"2 peaks: {total_r2b[2]}  3+ peaks: {total_r2b[3]}")
    print(f"  TI   —  0 peaks: {total_ti[0]}  1 peak: {total_ti[1]}  "
          f"2 peaks: {total_ti[2]}  3+ peaks: {total_ti[3]}")

    # ── Per-cell table ────────────────────────────────────────────────────────
    if not table_rows:
        print("\nNo cells with significant peaks found.")
    else:
        hdr = (f"{'date':<10} {'sub':<4} {'cell':>4}  "
               f"{'method':<5}  {'pk#':>3}  {'time(s)':>8}  {'score':>8}  {'p-val':>7}")
        print(f"\n=== Cells with significant peaks ===")
        print(hdr)
        print("-" * len(hdr))
        for date, sub, cell_idx, r2b_pks, ti_pks in table_rows:
            for method, pks in [("R2B", r2b_pks), ("TI", ti_pks)]:
                for pk_num, (t, score, pval) in enumerate(pks, 1):
                    print(f"{date:<10} {sub:<4} {cell_idx:>4}  "
                          f"{method:<5}  {pk_num:>3}  {t:>8.3f}  {score:>8.4f}  {pval:>7.4f}")

    # ── Combined heatmap ─────────────────────────────────────────────────────
    if r2b_time is not None:
        plot_heatmaps(r2b_traces,    r2b_pk,    r2b_time,
                      ti_traces,     ti_pk,     ti_time,
                      r2b_traces_mp, r2b_pk_mp,
                      ti_traces_mp,  ti_pk_mp,
                      us_time_s, "All sessions — G394")


if __name__ == "__main__":
    main()