#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
timeline_en.py
Create behavior-state timelines (and stacked composition) from a DLC CSV.

- If the CSV has a 3-level DLC header (scorer/bodyparts/coords), the script
  infers states (Idle/Grooming/Moving) using nose velocity and nose–ear distance.
- If the CSV is a flat table that already contains a 'state' column, that column
  is used directly (values are case-insensitive; e.g., 'idle', 'grooming', 'moving').

Outputs:
  <input>_timeline.png
  <input>_stacked_<bin>s.png
  <input>_timeline_segments.csv

Usage:
  python timeline_en.py /path/to/dlc_or_state_csv.csv --fps 30 --bin-sec 60
"""

import argparse
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# ===== Thresholds (adjust if necessary) =====
PCUTOFF = 0.90        # minimum likelihood for the "Nose" bodypart
V_GROOM = 3.0         # velocity threshold (px/frame) for Grooming
V_MOVE  = 8.0         # velocity threshold (px/frame) for Moving
DIST_EAR_NOSE = 15.0  # Grooming candidate if nose–ear distance is short
MIN_FRAMES = 30       # minimum consecutive frames for a valid state
# ============================================

BODYPOINTS_16 = [
    "Nose", "Left_eye", "Right_eye", "Left_forelimb", "Right_forelimb",
    "Left_hindlimb", "Right_hindlimb", "Left_ear", "Right_ear", "Neck",
    "Body_anterior", "Body_middle", "Body_posterior",
    "Tail_base", "Tail_middle", "Tail_tip"
]

def read_csv_any(csv_path: Path) -> pd.DataFrame:
    """
    Try to read as DLC multi-index CSV first; fall back to flat CSV.
    """
    try:
        df = pd.read_csv(csv_path, header=[0, 1, 2], low_memory=False)
        # if it didn't actually become a MultiIndex, treat as flat
        if not isinstance(df.columns, pd.MultiIndex):
            df = pd.read_csv(csv_path, low_memory=False)
        return df
    except Exception:
        return pd.read_csv(csv_path, low_memory=False)

def get_xy_lik(df: pd.DataFrame, bp: str):
    """
    Return x, y, likelihood for a given bodypart in DLC-format DataFrame.
    """
    x = df.loc[:, (slice(None), bp, "x")].iloc[:, 0]
    y = df.loc[:, (slice(None), bp, "y")].iloc[:, 0]
    l = df.loc[:, (slice(None), bp, "likelihood")].iloc[:, 0]
    return x, y, l

def smooth_min_run(labels: pd.Series, min_len: int) -> pd.Series:
    """
    Replace runs shorter than min_len with the nearest neighboring state.
    """
    a = labels.to_numpy().copy()
    n = len(a); i = 0
    while i < n:
        j = i + 1
        while j < n and a[j] == a[i]:
            j += 1
        if (j - i) < min_len:
            left = a[i - 1] if i > 0 else None
            right = a[j] if j < n else None
            a[i:j] = right if right is not None else left
        i = j
    return pd.Series(a, index=labels.index)

def infer_states_from_dlc(df: pd.DataFrame, fps: float) -> pd.Series:
    """
    Infer Idle/Grooming/Moving from DLC coordinates using:
      - nose velocity (px/frame)
      - minimum nose–ear distance (Left_ear or Right_ear)
    """
    # Choose bodyparts (require "Nose"; ears optional)
    if "Nose" not in df.columns.get_level_values(1):
        # fallback: use the first bodypart as "Nose"
        nose_bp = df.columns.get_level_values(1)[0]
    else:
        nose_bp = "Nose"

    x_nose, y_nose, l_nose = get_xy_lik(df, nose_bp)
    vx = x_nose.diff().fillna(0.0)
    vy = y_nose.diff().fillna(0.0)
    v = np.hypot(vx, vy).to_numpy()

    # Compute nose–ear distance
    dist_ne = np.full(len(x_nose), np.inf)
    for ear in ["Left_ear", "Right_ear"]:
        if ear in df.columns.get_level_values(1):
            x_e, y_e, _ = get_xy_lik(df, ear)
            d = np.hypot((x_nose - x_e), (y_nose - y_e)).to_numpy()
            dist_ne = np.minimum(dist_ne, d)

    valid = (l_nose >= PCUTOFF).to_numpy()

    labels = np.full(len(x_nose), "Idle", dtype=object)
    labels[(v >= V_MOVE) & valid] = "Moving"
    groom_mask = (v < V_GROOM) & valid & (
        dist_ne < (DIST_EAR_NOSE if np.isfinite(dist_ne).any() else V_GROOM)
    )
    labels[groom_mask] = "Grooming"

    return smooth_min_run(pd.Series(labels), MIN_FRAMES)

def contiguous_segments(states: np.ndarray, time_sec: np.ndarray):
    """
    Convert per-frame states into continuous segments (state, t_start, t_end).
    """
    change_idx = np.where(states[1:] != states[:-1])[0] + 1
    starts = np.r_[0, change_idx]
    ends = np.r_[change_idx, len(states)]
    segs = [(states[s], time_sec[s], time_sec[e - 1]) for s, e in zip(starts, ends) if e > s]
    return segs

def plot_timeline(segments, order, colors, out_path: Path):
    """
    Plot a hypnogram-like timeline of states.
    """
    row = {st: i for i, st in enumerate(order[::-1])}
    fig, ax = plt.subplots(figsize=(12, 2.8))
    for st, t0, t1 in segments:
        if st in row:
            ax.broken_barh([(t0, t1 - t0)], (row[st] - 0.4, 0.8),
                           facecolors=colors[st], edgecolors="none")
    ax.set_yticks([row[s] for s in order])
    ax.set_yticklabels(order)
    ax.set_xlabel("Time (s)")
    if segments:
        ax.set_xlim(segments[0][1], max(s[2] for s in segments))
    ax.set_title("Behavior state timeline")
    plt.tight_layout()
    plt.savefig(out_path, dpi=200)
    plt.close()

def plot_stacked(states: pd.Series, fps: float, bin_sec: int, order, colors, out_path: Path):
    """
    Plot stacked proportions of states in non-overlapping windows of bin_sec.
    """
    n = len(states)
    bin_frames = max(1, int(bin_sec * fps))
    bins = np.arange(n) // bin_frames
    counts = pd.crosstab(bins, states).reindex(columns=order, fill_value=0)
    props = counts.div(counts.sum(axis=1), axis=0).fillna(0.0)

    fig, ax = plt.subplots(figsize=(12, 3.2))
    x = np.arange(len(props)) * (bin_sec / 60.0)
    bottom = np.zeros(len(props))
    width = (bin_sec / 60.0) * 0.95
    for st in order:
        ax.bar(x, props[st].values, bottom=bottom, width=width,
               label=st, color=colors[st])
        bottom += props[st].values
    ax.set_xlabel("Time (min)")
    ax.set_ylabel("Proportion")
    ax.set_ylim(0, 1)
    ax.legend(loc="upper right", ncol=len(order), frameon=False)
    ax.set_title(f"Behavior composition per {bin_sec}s")
    plt.tight_layout()
    plt.savefig(out_path, dpi=200)
    plt.close()

def main():
    ap = argparse.ArgumentParser(
        description="Create behavior timeline/stacked plots from DLC CSV or a CSV with a 'state' column."
    )
    ap.add_argument("csv_path")
    ap.add_argument("--fps", type=float, default=30.0, help="Video frame rate.")
    ap.add_argument("--bin-sec", type=int, default=60, help="Bin size in seconds for stacked plot.")
    ap.add_argument("--order", default="Idle,Grooming,Moving", help="Comma-separated state order.")
    args = ap.parse_args()

    csv_path = Path(args.csv_path)
    fps = float(args.fps)
    bin_sec = int(args.bin_sec)
    order = [s.strip() for s in args.order.split(",") if s.strip()]
    colors = {"Idle": "#4daf4a", "Grooming": "#377eb8", "Moving": "#ff7f00"}

    df = read_csv_any(csv_path)

    # Determine states
    if isinstance(df.columns, pd.MultiIndex):
        states = infer_states_from_dlc(df, fps)
    else:
        if "state" not in df.columns:
            raise ValueError("Flat CSV detected but no 'state' column found.")
        states = df["state"].astype(str).str.capitalize()

    # Time axis and segments
    time_sec = np.arange(len(states)) / max(1e-9, fps)
    segments = contiguous_segments(states.to_numpy(), time_sec)

    # Outputs
    prefix = csv_path.with_suffix("")
    out_timeline = prefix.parent / f"{prefix.name}_timeline.png"
    out_stacked = prefix.parent / f"{prefix.name}_stacked_{bin_sec}s.png"
    out_segments = prefix.parent / f"{prefix.name}_timeline_segments.csv"

    # Render
    plot_timeline(segments, order, colors, out_timeline)
    plot_stacked(states, fps, bin_sec, order, colors, out_stacked)
    pd.DataFrame(segments, columns=["state", "t_start_sec", "t_end_sec"]).to_csv(out_segments, index=False)

    print(f"saved: {out_timeline}")
    print(f"saved: {out_stacked}")
    print(f"saved: {out_segments}")

if __name__ == "__main__":
    main()