state.py

"""
state.py — Streamlit session-state management and UI callbacks.

Covers:
  - init_session_state()
  - Speaker rename helpers (get_display_name, apply_speaker_renames_to_df)
  - Category callbacks (addCategory, removeCategory, updateCategoryOptions)
  - Global rename callbacks (addGlobalRename, removeGlobalRename, on_grename_change,
    apply_inline_rename)
  - File-switch callback (updateMultiSelect)
  - analyze() — builds and caches all DataFrames for a single file
  - convert_df(), printV()
"""

import copy
import traceback

import pandas as pd
import streamlit as st

import sonogram_utility as su
import utils


# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------

verbosity = 4  # 0=None 1=Low 2=Medium 3=High 4=Debug

def printV(message, level):
    if verbosity >= level:
        print(message)


# ---------------------------------------------------------------------------
# Session state initialisation
# ---------------------------------------------------------------------------

def init_session_state():
    """Idempotently initialise every session-state key the app needs."""
    defaults = {
        "results":          {},   # {filename: (annotations, totalSeconds)}
        "speakerRenames":   {},   # {filename: {speaker: name}}
        "summaries":        {},   # {filename: {df2, df3, ...}}
        "categories":       ["Instructor", "Students"],
        "categorySelect":   [[], []],   # [[token, ...], ...]  one list per category, tokens = "fname: SPEAKER_##"; starts with 2 lists for Instructor/Students
        "removeCategory":   None,
        "resetResult":      False,
        "unusedSpeakers":   {},   # {filename: [speaker, ...]}
        "file_names":       [],
        "valid_files":      [],
        "file_paths":       {},   # {filename: path}
        "showSummary":      "No",
        "speakerClips":     {},   # {filename: {speaker: wav_bytes}}
        "speakerSegments":  {},   # {filename: {speaker: [(start,end), ...]}}
        "speakerWaveforms": {},   # {filename: (waveform_tensor, sample_rate)}
        "globalRenames":    [],   # [{"name": str, "speakers": ["file: SPEAKER_##", ...]}]
        "analyzeAllToggle": False,
    }
    for key, value in defaults.items():
        if key not in st.session_state:
            st.session_state[key] = value


# ---------------------------------------------------------------------------
# Display-name helpers
# ---------------------------------------------------------------------------

def get_display_name(speaker, fileName):
    """Return the user-assigned display name for a speaker, or the original label.
    Role assignments (categorySelect) are intentionally excluded — roles are for
    grouping in charts, not for renaming speakers.
    """
    return st.session_state.speakerRenames.get(fileName, {}).get(speaker, speaker)


def apply_speaker_renames_to_df(df, fileName, column="task"):
    """Replace SPEAKER_## labels in a DataFrame column with display names."""
    if column not in df.columns:
        return df
    df = df.copy()
    df[column] = df[column].apply(lambda s: get_display_name(s, fileName))
    return df


@st.cache_data
def convert_df(df):
    return df.to_csv(index=False).encode("utf-8")


def build_all_csv_zip():
    """Build an in-memory ZIP containing one CSV per analyzed file.

    Applies the same transformations as the single-file download:
    drop Task, rename Resource -> Speaker, sort by Start, add Role.
    Returns raw ZIP bytes ready for st.download_button.
    """
    import io
    import zipfile

    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
        for fname, result in st.session_state.results.items():
            if len(result) != 2:
                continue
            try:
                annotation, _ = result
                currDF, _ = su.annotationToSimpleDataFrame(annotation)

                # Add Role column against raw SPEAKER_## labels BEFORE renames
                raw_to_role = {
                    token.split(": ", 1)[1]: st.session_state.categories[i]
                    for i, tokens in enumerate(st.session_state.categorySelect)
                    for token in tokens
                    if token.startswith(f"{fname}: ")
                }
                currDF = currDF.copy()
                currDF["Role"] = currDF["Resource"].map(raw_to_role).fillna("")

                # Apply speaker renames after Role is set
                renames = st.session_state.speakerRenames.get(fname, {})
                if "Resource" in currDF.columns:
                    currDF["Resource"] = currDF["Resource"].apply(
                        lambda s: renames.get(s, s)
                    )
                currDF = currDF.drop(columns=["Task"], errors="ignore")
                currDF = currDF.rename(columns={"Resource": "Speaker"})
                if "Start" in currDF.columns:
                    currDF = currDF.sort_values("Start").reset_index(drop=True)

                plain_name = fname.rsplit(".", 1)[0]
                zf.writestr(
                    f"sonogram-analysis-{plain_name}.csv",
                    currDF.to_csv(index=False)
                )
            except Exception as e:
                print(f"build_all_csv_zip: skipping {fname} — {e}")

    buf.seek(0)
    return buf.read()


# ---------------------------------------------------------------------------
# Category callbacks
# ---------------------------------------------------------------------------

def addCategory():
    new = st.session_state.categoryInput.strip()
    if not new:
        return
    st.toast(f"Adding {new}")
    st.session_state.categories.append(new)
    st.session_state.categorySelect.append([])
    st.session_state.pop(f"multiselect_{new}", None)
    st.session_state.categoryInput = ""


def removeCategory(index):
    name = st.session_state.categories[index]
    st.toast(f"Removing {name}")
    st.session_state.pop(f"multiselect_{name}", None)
    del st.session_state.categories[index]
    del st.session_state.categorySelect[index]


def updateCategoryOptions(token_display_map=None):
    """Store tokens ('fname: SPEAKER_##') in the global categorySelect lists.

    token_display_map: dict {display_label -> raw_token} passed from ui.py.
    Widget keys hold display labels; categorySelect must hold raw tokens.
    """
    if st.session_state.resetResult:
        return
    display_to_raw = token_display_map or {}

    # Guard: ensure categorySelect has exactly one slot per category.
    # Rapid interactions (e.g. addCategory firing mid-callback) can leave
    # the two lists temporarily out of sync.
    while len(st.session_state.categorySelect) < len(st.session_state.categories):
        st.session_state.categorySelect.append([])
    while len(st.session_state.categorySelect) > len(st.session_state.categories):
        st.session_state.categorySelect.pop()

    for i, category in enumerate(st.session_state.categories):
        ms_key = f"multiselect_{category}"
        display_vals = list(st.session_state.get(ms_key, []))
        raw_vals = [display_to_raw.get(t, t) for t in display_vals]
        st.session_state.categorySelect[i] = raw_vals

    # Recompute unusedSpeakers for all files
    all_assigned_tokens = {
        token
        for tokens in st.session_state.categorySelect
        for token in tokens
    }
    for fname, result in st.session_state.results.items():
        if len(result) != 2:
            continue
        try:
            annotation, _ = result
            unused = [
                sp for sp in annotation.labels()
                if f"{fname}: {sp}" not in all_assigned_tokens
            ]
            st.session_state.unusedSpeakers[fname] = unused
        except Exception:
            pass


# ---------------------------------------------------------------------------
# Global rename callbacks
# ---------------------------------------------------------------------------

# ---------------------------------------------------------------------------
# Global rename helpers
# ---------------------------------------------------------------------------

def _global_rename_key(index):
    return f"grename_speakers_{index}"


def _write_rename(token, name):
    """Write name into speakerRenames for a single token (fname: SPEAKER_##).
    If name is empty, clears the entry (revert to raw label).
    Silently ignores tokens that don't match a known file.
    """
    if ": " not in token:
        return
    fname, raw_sp = token.split(": ", 1)
    if fname not in st.session_state.speakerRenames:
        return  # token references an unknown file — ignore
    if name:
        st.session_state.speakerRenames[fname][raw_sp] = name
    else:
        st.session_state.speakerRenames[fname].pop(raw_sp, None)


def addGlobalRename():
    new_name = st.session_state.globalRenameInput.strip()
    if not new_name:
        return
    for entry in st.session_state.globalRenames:
        if entry["name"] == new_name:
            st.toast(f"'{new_name}' already exists in the rename list")
            st.session_state.globalRenameInput = ""
            return
    st.toast(f"Adding rename '{new_name}'")
    st.session_state.globalRenames.append({"name": new_name, "speakers": []})
    st.session_state.globalRenameInput = ""


def removeGlobalRename(index):
    entry = st.session_state.globalRenames[index]
    st.toast(f"Removing rename '{entry['name']}'")
    # Revert every speaker that belonged to this entry
    for token in entry["speakers"]:
        _write_rename(token, "")
    st.session_state.pop(_global_rename_key(index), None)
    del st.session_state.globalRenames[index]
    # Shift remaining widget keys down — ui.py will re-sync them to display
    # labels on the next render, so just clear them to force a clean re-seed.
    for i in range(index, len(st.session_state.globalRenames)):
        st.session_state.pop(_global_rename_key(i), None)


def apply_inline_rename(currFile, raw_sp, new_name):
    """Write a rename from the Rename Speaker tab into speakerRenames and globalRenames."""
    new_name = new_name.strip()
    token = f"{currFile}: {raw_sp}"

    for idx, entry in enumerate(st.session_state.globalRenames):
        if token in entry["speakers"]:
            entry["speakers"].remove(token)
            st.session_state.pop(_global_rename_key(idx), None)

    if new_name:
        _write_rename(token, new_name)
        for idx, entry in enumerate(st.session_state.globalRenames):
            if entry["name"] == new_name:
                entry["speakers"].append(token)
                st.session_state.pop(_global_rename_key(idx), None)
                return
        st.session_state.globalRenames.append({"name": new_name, "speakers": [token]})
    else:
        _write_rename(token, "")
        st.toast(f"Reverted {raw_sp} to original label")


def on_grename_change(idx, token_display_map=None):
    """Callback for the sidebar rename multiselect at position idx.

    token_display_map: dict {display_label -> raw_token} passed from ui.py.
    Widget keys hold display labels; entry["speakers"] must hold raw tokens.
    """
    # Guard: the entry may have been deleted (e.g. trash button fired just
    # before Streamlit re-fired this multiselect callback for the same index).
    if idx >= len(st.session_state.globalRenames):
        return
    grkey = _global_rename_key(idx)
    entry = st.session_state.globalRenames[idx]
    name  = entry["name"]

    display_to_raw = token_display_map or {}
    raw_to_display = {v: k for k, v in display_to_raw.items()}

    prev     = list(entry["speakers"])   # raw tokens in data model
    # Widget reports display labels — translate back to raw
    reported_display = list(st.session_state.get(grkey, []))
    reported = [display_to_raw.get(t, t) for t in reported_display]

    # Build the set of tokens that are legitimately available for this entry
    # right now (not claimed by any OTHER entry).
    other_claimed = {
        t
        for other_idx, other_entry in enumerate(st.session_state.globalRenames)
        if other_idx != idx
        for t in other_entry["speakers"]
    }

    # Spurious-empty guard: Streamlit sometimes re-fires this callback with []
    # when available_tokens shrinks (e.g. another entry just claimed a token).
    # Only treat it as spurious when prev had MORE than 1 token — if prev had
    # exactly 1, the user may genuinely be deselecting it, so always let it through.
    if not reported and len(prev) > 1:
        all_still_valid = all(t not in other_claimed for t in prev)
        if all_still_valid:
            # Restore widget key as display labels
            st.session_state[grkey] = [raw_to_display.get(t, t) for t in prev]
            return

    prev_set = set(prev)
    new_set  = set(reported)
    added    = new_set - prev_set
    removed  = prev_set - new_set

    # Write the new speakers list to the data model
    kept = [t for t in prev if t in new_set]
    entry["speakers"] = kept + [t for t in added]
    # Sync the widget key to match the data model.
    # Do NOT pop the key — popping causes Streamlit to re-initialise the widget
    # to [] on the next render (because no `default=` is passed), erasing the
    # selection visually even though the data model is correct.
    # Sync widget key as display labels
    st.session_state[grkey] = [raw_to_display.get(t, t) for t in entry["speakers"]]

    # Enforce exclusivity: a speaker can only belong to one rename entry at a time.
    # Remove the token from any other entry before writing the new name.
    for token in added:
        for other_idx, other_entry in enumerate(st.session_state.globalRenames):
            if other_idx == idx:
                continue
            if token in other_entry["speakers"]:
                other_entry["speakers"].remove(token)
                st.session_state[_global_rename_key(other_idx)] = [
                    raw_to_display.get(t, t) for t in other_entry["speakers"]
                ]
                _write_rename(token, "")
        _write_rename(token, name)

    # Revert speakerRenames for tokens genuinely removed from this entry
    for token in removed:
        _write_rename(token, "")


# ---------------------------------------------------------------------------
# File-switch callback
# ---------------------------------------------------------------------------

def updateMultiSelect():
    fileName = st.session_state["select_currFile"]
    st.session_state.resetResult = True
    result = st.session_state.results.get(fileName)
    if not result:
        return
    # Pop category widget keys so they re-seed from categorySelect data
    for category in st.session_state.categories:
        st.session_state.pop(f"multiselect_{category}", None)
    # Pop globalRenames widget keys so they re-seed from entry["speakers"] data
    for i in range(len(st.session_state.globalRenames)):
        st.session_state.pop(_global_rename_key(i), None)


# ---------------------------------------------------------------------------
# Speaker-clip session-state helpers
# ---------------------------------------------------------------------------

def store_speaker_clips(fname, annotations, waveform, sample_rate):
    """Generate samples & segments and write them into session state."""
    clips, segments = utils.build_speaker_clips(annotations, waveform, sample_rate)
    st.session_state.speakerClips[fname]    = clips
    st.session_state.speakerSegments[fname] = segments
    st.session_state.speakerWaveforms[fname] = (waveform, sample_rate)
    print(f"Generated {len(clips)} speaker samples for {fname}")


def randomize_speaker_clip(file_index, speaker):
    """Replace a speaker's audio sample with a freshly randomized one."""
    segs         = st.session_state.speakerSegments.get(file_index, {}).get(speaker)
    waveform_data = st.session_state.speakerWaveforms.get(file_index)
    if not segs or waveform_data is None:
        return
    waveform, sample_rate = waveform_data
    new_clip = utils.get_randomized_clip(waveform, sample_rate, segs)
    st.session_state.speakerClips[file_index][speaker] = new_clip
    print(f"Randomized sample for {speaker} in {file_index}")


# ---------------------------------------------------------------------------
# Per-file registration helper (keeps Demo / upload code DRY)
# ---------------------------------------------------------------------------

def register_file(fname):
    """Ensure all session-state dicts have an entry for fname."""
    st.session_state.results.setdefault(fname, [])
    st.session_state.summaries.setdefault(fname, {})
    st.session_state.unusedSpeakers.setdefault(fname, [])
    # Ensure categorySelect has one list per category (global, not per-file)
    while len(st.session_state.categorySelect) < len(st.session_state.categories):
        st.session_state.categorySelect.append([])
    st.session_state.speakerRenames.setdefault(fname, {})
    st.session_state.speakerClips.setdefault(fname, {})
    if fname not in st.session_state.file_names:
        st.session_state.file_names.append(fname)


# ---------------------------------------------------------------------------
# File loading helpers
# ---------------------------------------------------------------------------

def load_annotation_file(fname, fpath):
    """Load an annotation-only file (.txt / .rttm / .csv) into session state."""
    ext = fpath.lower()
    if ext.endswith(".txt"):
        _, annotations = su.loadAudioTXT(fpath)
    elif ext.endswith(".rttm"):
        _, annotations = su.loadAudioRTTM(fpath)
    elif ext.endswith(".csv"):
        _, annotations = su.loadAudioCSV(fpath)
    else:
        raise ValueError(f"Unsupported annotation format: {fpath}")
    totalSeconds = max((s.end for s in annotations.itersegments()), default=0)
    st.session_state.results[fname]        = (annotations, totalSeconds)
    st.session_state.summaries[fname]      = {}
    st.session_state.unusedSpeakers[fname] = list(annotations.labels())
    return annotations, totalSeconds


def load_demo_single(demo_path):
    """Register and load a single RTTM demo file, then run analyze()."""
    import time
    dname = demo_path.split("/")[-1]
    register_file(dname)
    st.session_state.file_paths[dname] = demo_path
    start_time = time.time()
    with st.spinner("Loading Demo Sample"):
        load_annotation_file(dname, demo_path)
    with st.spinner("Analyzing Demo Data"):
        analyze(dname)
    st.success(f"Took {time.time() - start_time:.1f}s to analyze the demo file!")
    st.session_state.select_currFile = dname
    return dname


def load_demo_single_sample(sample_path):
    """Register and load the pre-made short RTTM demo file, then run analyze()."""
    import time
    dname = sample_path.split("/")[-1]
    register_file(dname)
    st.session_state.file_paths[dname] = sample_path
    start_time = time.time()
    with st.spinner("Loading Sample Demo"):
        load_annotation_file(dname, sample_path)
    with st.spinner("Analyzing Sample Demo Data"):
        analyze(dname)
    st.success(f"Took {time.time() - start_time:.1f}s to analyze the sample demo!")
    st.session_state.select_currFile = dname
    return dname


def load_demo_multi(demo_paths):
    """Register and load multiple RTTM demo files."""
    for demo_path in demo_paths:
        dname = demo_path.split("/")[-1]
        register_file(dname)
        st.session_state.file_paths[dname] = demo_path
        with st.spinner(f"Loading: {dname}"):
            load_annotation_file(dname, demo_path)
    st.session_state.analyzeAllToggle = True


def run_analysis_loop(file_names, file_paths_dict, pipeline,
                      enable_denoise, early_cleanup,
                      gain_window, minimum_gain, maximum_gain,
                      df_model, df_state, atten_lim_db):
    """Process only new (not yet analyzed) files and populate session state."""
    import time
    import utils as _utils
    start_time = time.time()

    # Only process files that haven't been fully analyzed yet.
    # A file is considered done only when both results AND summaries are
    # populated — load_annotation_file sets results but not summaries, so
    # demo/annotation-only files correctly appear in pending until analyze()
    # has actually run.
    pending = [
        fname for fname in file_names
        if not (
            fname in st.session_state.results
            and len(st.session_state.results[fname]) == 2
            and st.session_state.summaries.get(fname, {}).get("speakers_dataFrame") is not None
        )
    ]

    if not pending:
        st.info("All files have already been analyzed.")
        st.session_state.analyzeAllToggle = False
        return

    totalFiles = len(pending)

    for i, fname in enumerate(pending):
        fpath = file_paths_dict.get(fname, "")
        ext   = fpath.lower()

        if ext.endswith((".txt", ".rttm", ".csv")):
            label = ext.rsplit(".", 1)[-1].upper()
            with st.spinner(f"Loading {label} {i+1}/{totalFiles}"):
                load_annotation_file(fname, fpath)
        else:
            with st.spinner(f"Processing Audio {i+1}/{totalFiles}"):
                annotations, totalSeconds, waveform, sample_rate = _utils.processFile(
                    fpath, pipeline, enable_denoise, early_cleanup,
                    gain_window, minimum_gain, maximum_gain,
                    df_model, df_state, atten_lim_db,
                )
                st.session_state.results[fname]        = (annotations, totalSeconds)
                st.session_state.summaries[fname]      = {}
                st.session_state.unusedSpeakers[fname] = list(annotations.labels())
            with st.spinner(f"Generating audio samples {i+1}/{totalFiles}"):
                store_speaker_clips(fname, annotations, waveform, sample_rate)
                del waveform

        with st.spinner(f"Analyzing {i+1}/{totalFiles}"):
            analyze(fname)

    st.success(f"Analyzed {totalFiles} new file(s) in {time.time() - start_time:.1f}s")
    st.session_state.analyzeAllToggle = False
    # Rotate uploader key to clear the file uploader widget
    st.session_state.uploader_key = st.session_state.get("uploader_key", 0) + 1


def build_table_df(displayDF):
    """Return a display-only copy of displayDF with cosmetic transforms applied:
       - Rename 'Resource' -> 'Speaker'
       - Drop 'Task' column if present
       - Format Start / Finish as HH:MM:SS.cs strings
    """
    def _fmt(val):
        try:
            secs = float(val)
        except (TypeError, ValueError):
            return str(val)
        h  = int(secs // 3600)
        m  = int(secs % 3600 // 60)
        s  = int(secs % 60)
        cs = round((secs % 1) * 100)
        return f"{h:02d}:{m:02d}:{s:02d}.{cs:02d}"

    df = displayDF.copy()
    if "Task" in df.columns:
        df = df.drop(columns=["Task"])
    if "Start" in df.columns:
        df["Start"] = df["Start"].apply(_fmt)
    if "Finish" in df.columns:
        df["Finish"] = df["Finish"].apply(_fmt)
    return df.rename(columns={"Resource": "Speaker"})


# ---------------------------------------------------------------------------
# analyze() — build and cache all DataFrames for one file
# ---------------------------------------------------------------------------

def analyze(inFileName):
    """Compute and store all summary DataFrames for inFileName."""
    try:
        printV(f"Start analyzing {inFileName}", 4)
        st.session_state.resetResult = False

        if not (
            inFileName in st.session_state.results
            and inFileName in st.session_state.summaries
            and len(st.session_state.results[inFileName]) > 0
        ):
            return

        currAnnotation, currTotalTime = st.session_state.results[inFileName]
        speakerNames      = currAnnotation.labels()
        # categorySelect is global tokens ("fname: SPEAKER_##"); extract raw IDs for this file
        prefix = inFileName + ": "
        categorySelections = [
            [token[len(prefix):] for token in tokens if token.startswith(prefix)]
            for tokens in st.session_state.categorySelect
        ]
        printV("Loaded results", 4)

        noVoice, oneVoice, multiVoice = su.calcSpeakingTypes(currAnnotation, currTotalTime)
        sumNoVoice    = su.sumTimes(noVoice)
        sumOneVoice   = su.sumTimes(oneVoice)
        sumMultiVoice = su.sumTimes(multiVoice)

        # df3
        df3 = utils.build_df3(noVoice, oneVoice, multiVoice)
        st.session_state.summaries[inFileName]["df3"] = df3
        printV("Set df3", 4)

        # df4
        df4, nameList, valueList, extraNames, extraValues = utils.build_df4(
            speakerNames, categorySelections, st.session_state.categories, currAnnotation
        )
        st.session_state.summaries[inFileName]["df4"] = df4
        printV("Set df4", 4)

        # df5
        df5 = utils.build_df5(
            oneVoice, multiVoice,
            sumNoVoice, sumOneVoice, sumMultiVoice,
            currTotalTime,
        )
        st.session_state.summaries[inFileName]["df5"] = df5
        printV("Set df5", 4)

        # speakers_dataFrame, df2
        speakers_dataFrame, speakers_times = su.annotationToDataFrame(currAnnotation)
        st.session_state.summaries[inFileName]["speakers_dataFrame"] = speakers_dataFrame
        st.session_state.summaries[inFileName]["speakers_times"]     = speakers_times

        df2 = utils.build_df2(
            nameList + extraNames,
            valueList + extraValues,
            currTotalTime,
        )
        st.session_state.summaries[inFileName]["df2"] = df2
        printV("Set df2", 4)

    except Exception as e:
        print(f"Error in analyze: {e}")
        traceback.print_exc()
        st.error(f"Debug - analyze() failed: {e}")