{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":59156,"databundleVersionId":13719397,"sourceType":"competition"}],"dockerImageVersionId":31090,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## MABe Challenge — Dataset-Driven Heuristic Baseline (Polars + Visual QA)\n\nThis notebook:\n- Loads metadata and annotations (train), computes priors and timing statistics.\n- Summarizes the dataset with compact, informative visualizations.\n- Generates heuristic predictions on test via priors + proximity windows.\n- Saves a valid submission.csv.\n\nFocus: tight runtime, clear outputs, minimal moving parts to establish a strong baseline and robust analysis.","metadata":{}},{"cell_type":"code","source":"from __future__ import annotations\n\nimport os\nimport json\nimport logging\nimport warnings\nfrom dataclasses import dataclass\nfrom pathlib import Path\nfrom typing import Dict, Iterable, List, Optional, Set, Tuple\nfrom collections import defaultdict\n\nimport numpy as np\nimport pandas as pd\nimport polars as pl\nfrom tqdm.auto import tqdm\n\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom matplotlib.patches import Rectangle\n\n# Style and warnings\nwarnings.filterwarnings(\"ignore\")\nplt.style.use('seaborn-v0_8-whitegrid')\nsns.set_palette(\"husl\")\nplt.rcParams[\"figure.dpi\"] = 120\n\n# Logging\nlogger = logging.getLogger(__name__)\ndef setup_logging(verbosity: int = 1) -> None:\n    level = logging.WARNING if verbosity <= 0 else logging.INFO if verbosity == 1 else logging.DEBUG\n    logging.basicConfig(level=level, format=\"%(asctime)s | %(levelname)s | %(name)s | %(message)s\", force=True)\n\n@dataclass(frozen=True)\nclass Config:\n    data_root: Path = Path(os.getenv(\"MABE_DATA_ROOT\", \"/kaggle/input/MABe-mouse-behavior-detection\"))\n    submission_file: str = os.getenv(\"MABE_SUBMISSION\", \"submission.csv\")\n    row_id_col: str = os.getenv(\"MABE_ROW_ID_COL\", \"row_id\")\n\n    @property\n    def train_csv(self) -> Path: return self.data_root / \"train.csv\"\n    @property\n    def test_csv(self) -> Path: return self.data_root / \"test.csv\"\n    @property\n    def train_annot_dir(self) -> Path: return self.data_root / \"train_annotation\"\n    @property\n    def train_track_dir(self) -> Path: return self.data_root / \"train_tracking\"\n    @property\n    def test_track_dir(self) -> Path: return self.data_root / \"test_tracking\"\n\n    @property\n    def submission_schema(self) -> Dict[str, pl.DataType]:\n        return {\n            \"video_id\": pl.Int64, \"agent_id\": pl.Utf8, \"target_id\": pl.Utf8,\n            \"action\": pl.Utf8, \"start_frame\": pl.Int64, \"stop_frame\": pl.Int64,\n        }\n\n    @property\n    def solution_schema(self) -> Dict[str, pl.DataType]:\n        return {\n            \"video_id\": pl.Int64, \"agent_id\": pl.Utf8, \"target_id\": pl.Utf8,\n            \"action\": pl.Utf8, \"start_frame\": pl.Int64, \"stop_frame\": pl.Int64,\n            \"lab_id\": pl.Utf8, \"behaviors_labeled\": pl.Utf8,\n        }\n\nclass HostVisibleError(Exception): pass\n\ndef safe_json_loads(s: Optional[str]) -> List[str]:\n    if s is None: return []\n    if isinstance(s, list): return [str(x) for x in s]\n    if not isinstance(s, str): return []\n    s = s.strip()\n    if not s: return []\n    try:\n        return json.loads(s)\n    except Exception:\n        try: return json.loads(s.replace(\"'\", '\"'))\n        except Exception: return []\n\ndef validate_schema(df: pl.DataFrame, schema: Dict[str, pl.DataType], name: str) -> pl.DataFrame:\n    missing = set(schema.keys()) - set(df.columns)\n    if missing: raise ValueError(f\"{name} is missing columns: {missing}\")\n    casts = [pl.col(col).cast(dtype) for col, dtype in schema.items() if df[col].dtype != dtype]\n    return df.with_columns(casts) if casts else df\n\ndef validate_frame_ranges(df: pl.DataFrame, name: str) -> None:\n    if not (df[\"start_frame\"] <= df[\"stop_frame\"]).all():\n        raise ValueError(f\"{name}: start_frame > stop_frame detected\")\n\ndef _norm_mouse_id(x: str | int) -> str:\n    s = str(x)\n    return s if s.startswith(\"mouse\") else f\"mouse{s}\"\n\ndef merge_intervals(intervals: List[Tuple[int,int]]) -> List[Tuple[int,int]]:\n    if not intervals: return []\n    intervals = sorted(intervals)\n    merged = [intervals[0]]\n    for s,e in intervals[1:]:\n        ps,pe = merged[-1]\n        if s <= pe: merged[-1] = (ps, max(pe, e))\n        else: merged.append((s,e))\n    return merged","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-22T19:08:17.695261Z","iopub.execute_input":"2025-09-22T19:08:17.695613Z","iopub.status.idle":"2025-09-22T19:08:17.742466Z","shell.execute_reply.started":"2025-09-22T19:08:17.695587Z","shell.execute_reply":"2025-09-22T19:08:17.741455Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"setup_logging(verbosity=1)\ncfg = Config()\n\nprint(\"Loading train/test metadata...\")\ntrain = pl.read_csv(cfg.train_csv)\ntest = pl.read_csv(cfg.test_csv)\ntrain_subset = train.filter(~pl.col(\"lab_id\").str.starts_with(\"MABe22\"))\nprint(f\"Train rows: {len(train)} (after filter: {len(train_subset)}) | Test rows: {len(test)}\")\n\ndef create_solution_df(dataset: pl.DataFrame, cfg: Optional[Config] = None) -> pl.DataFrame:\n    cfg = cfg or Config()\n    records: List[pl.DataFrame] = []\n    for row in tqdm(dataset.to_dicts(), total=len(dataset), desc=\"Building solution\"):\n        lab_id: str = row[\"lab_id\"]\n        video_id: int = row[\"video_id\"]\n        annot_path = cfg.train_annot_dir / lab_id / f\"{video_id}.parquet\"\n        if not annot_path.exists():\n            continue\n        try:\n            annot = pl.read_parquet(annot_path).with_columns(\n                [\n                    pl.lit(lab_id).alias(\"lab_id\"),\n                    pl.lit(video_id).alias(\"video_id\"),\n                    pl.lit(row[\"behaviors_labeled\"]).alias(\"behaviors_labeled\"),\n                    pl.concat_str([pl.lit(\"mouse\"), pl.col(\"agent_id\").cast(pl.Utf8)]).alias(\"agent_id\"),\n                    pl.concat_str([pl.lit(\"mouse\"), pl.col(\"target_id\").cast(pl.Utf8)]).alias(\"target_id\"),\n                ]\n            )\n            # keep only required columns, cast if needed\n            keep_cols = list(Config().solution_schema.keys())\n            annot = annot.select([c for c in keep_cols if c in annot.columns])\n            annot = validate_schema(annot, Config().solution_schema, \"Solution\")\n            records.append(annot)\n        except Exception as e:\n            logger.warning(\"Failed to load %s: %s\", annot_path, e)\n            continue\n    if not records: raise ValueError(\"No annotation files loaded.\")\n    solution = pl.concat(records, how=\"vertical\")\n    return solution\n\ndef build_video_spans(dataset: pl.DataFrame, split: str, cfg: Optional[Config] = None) -> Dict[int, Tuple[int,int]]:\n    cfg = cfg or Config()\n    track_dir = cfg.train_track_dir if split == \"train\" else cfg.test_track_dir\n    spans: Dict[int, Tuple[int,int]] = {}\n    for row in tqdm(dataset.to_dicts(), total=len(dataset), desc=\"Scanning spans\"):\n        lab_id = row[\"lab_id\"]\n        if str(lab_id).startswith(\"MABe22\"): continue\n        vid = row[\"video_id\"]\n        path = track_dir / lab_id / f\"{vid}.parquet\"\n        if not path.exists(): continue\n        try:\n            df = pl.read_parquet(path).select([\"video_frame\"])\n            s = int(df[\"video_frame\"].min())\n            e = int(df[\"video_frame\"].max()) + 1\n            spans[int(vid)] = (s,e)\n        except Exception as e:\n            continue\n    return spans\n\ndef compute_action_priors(solution: pl.DataFrame, eps: float = 1.0):\n    sol = solution.with_columns((pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"dur\"))\n    by_lab = sol.group_by([\"lab_id\", \"action\"]).agg(pl.col(\"dur\").sum().alias(\"dur_sum\"))\n    global_ = sol.group_by([\"action\"]).agg(pl.col(\"dur\").sum().alias(\"dur_sum\"))\n    actions = set(global_[\"action\"].to_list())\n\n    per_lab_weight: Dict[str, Dict[str, float]] = defaultdict(dict)\n    for lab in by_lab[\"lab_id\"].unique():\n        sub = by_lab.filter(pl.col(\"lab_id\") == lab)\n        dmap = {r[\"action\"]: float(r[\"dur_sum\"]) for r in sub.to_dicts()}\n        for a in actions: dmap[a] = dmap.get(a, 0.0) + eps\n        total = sum(dmap.values()) or 1.0\n        per_lab_weight[str(lab)] = {a: dmap[a]/total for a in actions}\n\n    gmap = {r[\"action\"]: float(r[\"dur_sum\"]) for r in global_.to_dicts()}\n    for a in actions: gmap[a] = gmap.get(a, 0.0) + eps\n    gtotal = sum(gmap.values()) or 1.0\n    global_weight = {a: gmap[a]/gtotal for a in actions}\n\n    med_by_lab = sol.group_by([\"lab_id\", \"action\"]).median().select([\"lab_id\",\"action\",\"dur\"])\n    per_lab_med_dur: Dict[str, Dict[str, int]] = defaultdict(dict)\n    for r in med_by_lab.to_dicts():\n        per_lab_med_dur[str(r[\"lab_id\"])][str(r[\"action\"])] = int(r[\"dur\"])\n    med_global = sol.group_by([\"action\"]).median().select([\"action\",\"dur\"])\n    global_med_dur: Dict[str, int] = {r[\"action\"]: int(r[\"dur\"]) for r in med_global.to_dicts()}\n    return per_lab_weight, global_weight, per_lab_med_dur, global_med_dur\n\ndef compute_timing_priors(solution: pl.DataFrame, video_spans: Dict[int, Tuple[int,int]]):\n    def start_pct_func(row) -> float:\n        vid = int(row[\"video_id\"])\n        if vid not in video_spans:\n            return 0.5\n        s, e = video_spans[vid]\n        denom = max(1, e - s)\n        return float(max(0, min(1, (int(row[\"start_frame\"]) - s) / denom)))\n    rows = [\n        {\"lab_id\": r[\"lab_id\"], \"action\": r[\"action\"], \"start_pct\": start_pct_func(r)}\n        for r in solution.select([\"lab_id\", \"action\", \"video_id\", \"start_frame\"]).to_dicts()\n    ]\n    df = pl.DataFrame(rows)\n    by_lab = df.group_by([\"lab_id\", \"action\"]).median().select([\"lab_id\", \"action\", \"start_pct\"])\n    per_lab = defaultdict(dict)\n    for r in by_lab.to_dicts():\n        per_lab[str(r[\"lab_id\"])][str(r[\"action\"])] = float(r[\"start_pct\"])\n    g = df.group_by([\"action\"]).median().select([\"action\", \"start_pct\"])\n    global_ = {r[\"action\"]: float(r[\"start_pct\"]) for r in g.to_dicts()}\n    return per_lab, global_\n\nprint(\"Building solution from annotations...\")\nsolution = create_solution_df(train_subset, cfg)\nprint(f\"Annotations loaded: {len(solution):,}\")\n\nprint(\"Computing frame spans...\")\nspans = build_video_spans(train_subset, \"train\", cfg)\nprint(f\"Video spans computed: {len(spans)}\")\n\nprint(\"Computing priors...\")\nper_lab, global_w, med_lab, med_glob = compute_action_priors(solution, eps=0.5)\ntiming_lab, timing_glob = compute_timing_priors(solution, spans)\nprint(f\"Actions in priors: {len(global_w)} | Labs with priors: {len(per_lab)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-22T19:10:29.481545Z","iopub.execute_input":"2025-09-22T19:10:29.481840Z","iopub.status.idle":"2025-09-22T19:11:09.331117Z","shell.execute_reply.started":"2025-09-22T19:10:29.481819Z","shell.execute_reply":"2025-09-22T19:11:09.330276Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def visualize_data_overview(train_df: pl.DataFrame, solution_df: pl.DataFrame):\n    fig, axes = plt.subplots(2, 3, figsize=(16, 9))\n\n    # 1) Top Labs\n    ax = axes[0,0]\n    lab_counts = train_df.group_by(\"lab_id\").count().sort(\"count\", descending=True)\n    top = lab_counts.head(12)\n    ax.barh(top[\"lab_id\"].to_list()[::-1], top[\"count\"].to_list()[::-1], color=\"#4C72B0\")\n    ax.set_title(\"Videos per Lab (Top 12)\"); ax.set_xlabel(\"Count\"); ax.grid(axis=\"x\", alpha=0.3)\n\n    # 2) Action counts (solution)\n    ax = axes[0,1]\n    act_counts = solution_df.group_by(\"action\").count().sort(\"count\", descending=True)\n    acts = act_counts[\"action\"].to_list()[:20][::-1]\n    vals = act_counts[\"count\"].to_list()[:20][::-1]\n    ax.barh(acts, vals, color=\"#64B5CD\")\n    ax.set_title(\"Top Behaviors (Train)\"); ax.set_xlabel(\"Count\"); ax.grid(axis=\"x\", alpha=0.3)\n\n    # 3) Duration histogram\n    ax = axes[0,2]\n    sd = solution_df.with_columns((pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\"))[\"duration\"].to_list()\n    if sd:\n        ax.hist(sd, bins=40, color=\"#C44E52\", edgecolor=\"black\", linewidth=0.3)\n        ax.axvline(np.median(sd), color=\"black\", linestyle=\"--\", alpha=0.7, label=f\"Median={np.median(sd):.1f}\")\n        ax.legend()\n    ax.set_title(\"Behavior Duration (frames)\"); ax.set_xlabel(\"Frames\"); ax.set_ylabel(\"Count\"); ax.grid(axis=\"y\", alpha=0.3)\n\n    # 4) Pairs\n    ax = axes[1,0]\n    pair_counts = solution_df.group_by([\"agent_id\", \"target_id\"]).count().sort(\"count\", descending=True).head(10)\n    labels = [f\"{r['agent_id']}->{r['target_id']}\" for r in pair_counts.to_dicts()][::-1]\n    values = pair_counts[\"count\"].to_list()[::-1]\n    ax.barh(labels, values, color=\"#E17C05\")\n    ax.set_title(\"Top Pairs (Train)\"); ax.set_xlabel(\"Count\"); ax.grid(axis=\"x\", alpha=0.3)\n\n    # 5) Behaviors per video\n    ax = axes[1,1]\n    per_video = solution_df.group_by(\"video_id\").count()[\"count\"].to_list()\n    if per_video:\n        ax.hist(per_video, bins=30, color=\"#8172B2\", edgecolor=\"black\", linewidth=0.3)\n        ax.axvline(np.mean(per_video), color=\"black\", linestyle=\"--\", alpha=0.7, label=f\"Mean={np.mean(per_video):.1f}\")\n        ax.legend()\n    ax.set_title(\"Behaviors per Video\"); ax.set_xlabel(\"Count\"); ax.set_ylabel(\"Videos\"); ax.grid(axis=\"y\", alpha=0.3)\n\n    # 6) Summary\n    ax = axes[1,2]; ax.axis(\"off\")\n    n_vid = int(train_df[\"video_id\"].n_unique())\n    n_lab = int(train_df[\"lab_id\"].n_unique())\n    n_anns = len(solution_df)\n    n_act = int(solution_df[\"action\"].n_unique())\n    dur = solution_df.with_columns((pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\"))[\"duration\"]\n    text = f\"\"\"DATASET SUMMARY\nVideos: {n_vid}\nLabs: {n_lab}\nAnnotations: {n_anns:,}\nUnique behaviors: {n_act}\nMean duration: {float(dur.mean() if len(dur)>0 else 0):.1f} frames\"\"\"\n    ax.text(0.05, 0.95, text, transform=ax.transAxes, fontsize=10,\n            va=\"top\", fontfamily=\"monospace\", bbox=dict(boxstyle=\"round\", facecolor=\"lightgray\", alpha=0.3))\n\n    plt.tight_layout(); plt.show()\n\nprint(\"Overview visuals...\")\nvisualize_data_overview(train_subset, solution)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-22T19:11:24.641566Z","iopub.execute_input":"2025-09-22T19:11:24.641921Z","iopub.status.idle":"2025-09-22T19:11:26.228742Z","shell.execute_reply.started":"2025-09-22T19:11:24.641886Z","shell.execute_reply":"2025-09-22T19:11:26.227927Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def _strip_mouse_prefix(s: str | int) -> str:\n    s = str(s)\n    return s[5:] if s.startswith(\"mouse\") else s\n\ndef _pair_features(df: pl.DataFrame, agent_raw: str, target_raw: str, downsample: int = 1) -> Optional[pl.DataFrame]:\n    frame_candidates = [\"video_frame\",\"frame\",\"frame_idx\"]\n    id_candidates = [\"mouse_id\",\"id\",\"track_id\",\"agent_id\"]\n    x_candidates = [\"x\",\"x_pos\",\"x_position\",\"x_mm\",\"centroid_x\",\"cx\"]\n    y_candidates = [\"y\",\"y_pos\",\"y_position\",\"y_mm\",\"centroid_y\",\"cy\"]\n\n    cols = set(df.columns)\n    frame_col = next((c for c in frame_candidates if c in cols), None)\n    id_col = next((c for c in id_candidates if c in cols), None)\n    x_col = next((c for c in x_candidates if c in cols), None)\n    y_col = next((c for c in y_candidates if c in cols), None)\n    if not all([frame_col, id_col, x_col, y_col]): return None\n\n    a_id = _strip_mouse_prefix(agent_raw)\n    t_id = _strip_mouse_prefix(target_raw)\n\n    pdf = df.select([frame_col, id_col, x_col, y_col]).to_pandas()\n    pdf[frame_col] = pdf[frame_col].astype(np.int64, copy=False)\n    pdf[id_col] = pdf[id_col].astype(str, copy=False)\n\n    a = pdf[pdf[id_col] == a_id].copy()\n    b = pdf[pdf[id_col] == t_id].copy()\n    if a.empty or b.empty: return None\n\n    a.drop_duplicates(subset=[frame_col], inplace=True)\n    b.drop_duplicates(subset=[frame_col], inplace=True)\n\n    merged = a.merge(b, on=frame_col, how=\"inner\", suffixes=(\"_a\", \"_b\"))\n    if merged.empty: return None\n    merged.sort_values(frame_col, inplace=True)\n\n    ax = merged[f\"{x_col}_a\"].to_numpy(dtype=np.float64)\n    ay = merged[f\"{y_col}_a\"].to_numpy(dtype=np.float64)\n    bx = merged[f\"{x_col}_b\"].to_numpy(dtype=np.float64)\n    by = merged[f\"{y_col}_b\"].to_numpy(dtype=np.float64)\n    frames = merged[frame_col].to_numpy(dtype=np.int64)\n\n    if downsample and downsample > 1:\n        sl = slice(0, None, int(downsample))\n        ax, ay, bx, by, frames = ax[sl], ay[sl], bx[sl], by[sl], frames[sl]\n        if ax.size == 0: return None\n\n    dx, dy = ax - bx, ay - by\n    dist = np.sqrt(dx*dx + dy*dy)\n    dax, day = np.diff(ax, prepend=ax[0]), np.diff(ay, prepend=ay[0])\n    dbx, dby = np.diff(bx, prepend=bx[0]), np.diff(by, prepend=by[0])\n    rel_speed = np.sqrt(dax*dax + day*day) - np.sqrt(dbx*dbx + dby*dby)\n    ddist = np.diff(dist, prepend=dist[0])\n\n    return pl.DataFrame({\"frame\": frames, \"dist\": dist, \"rel_speed\": rel_speed, \"ddist\": ddist}).sort(\"frame\")\n\ndef _make_windows(feat: pl.DataFrame, min_len: int, q_dist: float = 0.40, q_rel: float = 0.60, q_ddist: float = 0.40) -> List[Tuple[int,int]]:\n    if len(feat) == 0: return []\n    qd = float(feat[\"dist\"].quantile(q_dist))\n    qr = float(feat[\"rel_speed\"].quantile(q_rel))\n    qdd = float(feat[\"ddist\"].quantile(q_ddist))\n    cond = (pl.col(\"dist\") <= qd) | ((pl.col(\"rel_speed\") >= qr) & (pl.col(\"ddist\") <= qdd))\n    mask = feat.select(cond.alias(\"m\")).to_series().to_list()\n    frames = feat[\"frame\"].to_list()\n\n    windows: List[Tuple[int,int]] = []\n    run = None\n    for i, flag in enumerate(mask):\n        if flag and run is None: run = [frames[i], frames[i]]\n        elif flag and run is not None: run[1] = frames[i]\n        elif (not flag) and run is not None:\n            s,e = run[0], run[1]+1\n            if e - s >= min_len: windows.append((s,e))\n            run = None\n    if run is not None:\n        s,e = run[0], run[1]+1\n        if e - s >= min_len: windows.append((s,e))\n    return merge_intervals(windows)\n\ndef _order_actions_by_timing(actions: List[str], lab_id: str,\n                             timing_lab: Dict[str, Dict[str, float]],\n                             timing_global: Dict[str, float],\n                             canonical: Dict[str,int]) -> List[str]:\n    def score(a: str) -> float:\n        if lab_id in timing_lab and a in timing_lab[lab_id]:\n            return timing_lab[lab_id][a]\n        return timing_global.get(a, 0.5)\n    return sorted(actions, key=lambda a: (score(a), canonical.get(a, 99)))\n\ndef _clip_rare_actions(weights_map: Dict[str,float], actions: List[str], p_min: float, cap: float) -> Dict[str,float]:\n    w = {a: max(0.0, float(weights_map.get(a, 0.0))) for a in actions}\n    for a in actions:\n        if w[a] < p_min:\n            w[a] = min(w[a], cap)\n    s = sum(w.values()) or 1.0\n    return {a: w[a]/s for a in actions}\n\ndef _allocate_segments_in_windows(windows: List[Tuple[int,int]],\n                                  ordered_actions: List[str],\n                                  weights: Dict[str,float],\n                                  med_dur: Dict[str,int],\n                                  total_frames: int) -> List[Tuple[str,int,int]]:\n    win_idx = 0\n    cur_s, cur_e = (windows[0] if windows else (0,0))\n    remain = sum(e-s for s,e in windows)\n    out: List[Tuple[str,int,int]] = []\n\n    for a in ordered_actions:\n        if remain <= 0: break\n        want = int(weights.get(a, 0.0) * total_frames)\n        want = max(want, int(med_dur.get(a, 0) or 0))\n        want = min(want, remain)\n        got = 0\n        while got < want and win_idx < len(windows):\n            s,e = cur_s, cur_e\n            if s >= e:\n                win_idx += 1\n                if win_idx >= len(windows): break\n                cur_s, cur_e = windows[win_idx]\n                continue\n            take = min(want - got, e - s)\n            out.append((a, s, s+take))\n            got += take\n            remain -= take\n            cur_s = s + take\n            if cur_s >= e and win_idx < len(windows):\n                win_idx += 1\n                if win_idx < len(windows):\n                    cur_s, cur_e = windows[win_idx]\n    return out\n\ndef _smooth_segments(segments: List[Tuple[str,int,int]], min_len: int, gap_close: int) -> List[Tuple[str,int,int]]:\n    if not segments: return []\n    segments = sorted(segments, key=lambda x: (x[1], x[2], x[0]))\n    segments = [seg for seg in segments if seg[2] - seg[1] >= min_len]\n    if not segments: return []\n    out = [segments[0]]\n    for a,s,e in segments[1:]:\n        pa,ps,pe = out[-1]\n        if a == pa and s - pe <= gap_close:\n            out[-1] = (a, ps, e)\n        else:\n            out.append((a,s,e))\n    return out\n\ndef predict_without_ml(dataset: pl.DataFrame, data_split: str, cfg: Optional[Config] = None,\n                       priors_per_lab: Optional[Dict[str, Dict[str, float]]] = None,\n                       priors_global: Optional[Dict[str, float]] = None,\n                       meddur_per_lab: Optional[Dict[str, Dict[str, int]]] = None,\n                       meddur_global: Optional[Dict[str, int]] = None,\n                       timing_lab: Optional[Dict[str, Dict[str, float]]] = None,\n                       timing_global: Optional[Dict[str, float]] = None,\n                       prior_scope: str = \"mixed\",\n                       use_windows: bool = True,\n                       min_len: int = 15,\n                       gap_close: int = 3,\n                       p_min: float = 0.05,\n                       cap: float = 0.05) -> pl.DataFrame:\n    cfg = cfg or Config()\n    track_dir = cfg.test_track_dir if data_split == \"test\" else cfg.train_track_dir\n    records: List[Tuple[int, str, str, str, int, int]] = []\n    canonical = {\"approach\": 0, \"avoid\": 1, \"chase\": 2, \"chaseattack\": 3, \"attack\": 4, \"mount\": 5, \"submit\": 6}\n\n    for row in tqdm(dataset.to_dicts(), total=len(dataset), desc=f\"Predicting ({data_split})\"):\n        lab_id: str = row[\"lab_id\"]\n        video_id: int = row[\"video_id\"]\n        path = track_dir / lab_id / f\"{video_id}.parquet\"\n        if not path.exists():\n            continue\n\n        try:\n            trk = pl.read_parquet(path)\n            start_frame = int(trk[\"video_frame\"].min())\n            stop_frame = int(trk[\"video_frame\"].max()) + 1\n            video_frames = stop_frame - start_frame\n            if video_frames <= 0: continue\n\n            raw_list = safe_json_loads(row[\"behaviors_labeled\"])\n            triples: List[List[str]] = []\n            for b in raw_list:\n                parts = [p.strip() for p in str(b).replace(\"'\", \"\").split(\",\")]\n                if len(parts) == 3:\n                    triples.append(parts)\n            if not triples:\n                continue\n\n            beh_df = pl.DataFrame(triples, schema=[\"agent\",\"target\",\"action\"], orient=\"row\").with_columns(\n                [pl.col(\"agent\").cast(pl.Utf8), pl.col(\"target\").cast(pl.Utf8), pl.col(\"action\").cast(pl.Utf8)]\n            )\n\n            for (agent, target), group in beh_df.group_by([\"agent\",\"target\"]):\n                actions = sorted(list(set(group[\"action\"].to_list())), key=lambda a: canonical.get(a, 99))\n                if not actions: continue\n\n                if prior_scope == \"lab\" and priors_per_lab is not None:\n                    w_map = priors_per_lab.get(str(lab_id), {})\n                    md_map = meddur_per_lab.get(str(lab_id), {}) if meddur_per_lab else {}\n                elif prior_scope == \"global\" and priors_global is not None:\n                    w_map = priors_global\n                    md_map = meddur_global or {}\n                else:\n                    w_map = (priors_per_lab or {}).get(str(lab_id), {}) or (priors_global or {})\n                    md_map = (meddur_per_lab or {}).get(str(lab_id), {}) or (meddur_global or {})\n\n                weights = _clip_rare_actions(w_map, actions, p_min=p_min, cap=cap)\n                ordered_actions = _order_actions_by_timing(\n                    actions, str(lab_id), timing_lab or {}, timing_global or {}, canonical\n                )\n\n                if use_windows:\n                    feat = _pair_features(trk, _norm_mouse_id(agent), _norm_mouse_id(target))\n                    if feat is None:\n                        windows = [(start_frame, stop_frame)]\n                    else:\n                        windows = _make_windows(feat, min_len=min_len)\n                        if not windows:\n                            windows = [(start_frame, stop_frame)]\n                else:\n                    windows = [(start_frame, stop_frame)]\n\n                windows = merge_intervals(windows)\n                allowed_total = sum(e - s for s,e in windows)\n                if allowed_total <= 0: continue\n\n                segs = _allocate_segments_in_windows(\n                    windows=windows,\n                    ordered_actions=ordered_actions,\n                    weights=weights,\n                    med_dur=md_map,\n                    total_frames=allowed_total\n                )\n                segs = _smooth_segments(segs, min_len=min_len, gap_close=gap_close)\n\n                for a, s, e in segs:\n                    if e > s:\n                        records.append((video_id, _norm_mouse_id(agent), _norm_mouse_id(target), a, int(s), int(e)))\n\n        except Exception as e:\n            logger.warning(\"Error processing %s: %s\", path, e)\n            continue\n\n    if not records:\n        raise ValueError(\"No predictions generated.\")\n\n    df = pl.DataFrame(\n        records,\n        schema={\n            \"video_id\": pl.Int64, \"agent_id\": pl.Utf8, \"target_id\": pl.Utf8,\n            \"action\": pl.Utf8, \"start_frame\": pl.Int64, \"stop_frame\": pl.Int64,\n        },\n        orient=\"row\",\n    )\n    df = validate_schema(df, cfg.submission_schema, \"Submission\")\n    validate_frame_ranges(df, \"Submission\")\n    return df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-22T19:11:39.372750Z","iopub.execute_input":"2025-09-22T19:11:39.373542Z","iopub.status.idle":"2025-09-22T19:11:39.410272Z","shell.execute_reply.started":"2025-09-22T19:11:39.373516Z","shell.execute_reply":"2025-09-22T19:11:39.409437Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Computing overview and priors completed above.\")\nprint(\"Loading test and generating predictions...\")\n\nprediction_params = dict(\n    prior_scope=\"mixed\",\n    use_windows=True,\n    min_len=15,\n    gap_close=3,\n    p_min=0.05,\n    cap=0.05\n)\n\nsubmission_test = predict_without_ml(\n    test, \"test\", cfg,\n    priors_per_lab=per_lab, priors_global=global_w,\n    meddur_per_lab=med_lab, meddur_global=med_glob,\n    timing_lab=timing_lab, timing_global=timing_glob,\n    **prediction_params\n)\nprint(f\"Predictions generated: {len(submission_test):,}\")\n\n# Quick prediction QA visuals (safe for small test)\ndef visualize_predictions(pred_df: pl.DataFrame):\n    if len(pred_df) == 0:\n        print(\"No predictions to visualize.\"); return\n    fig, axes = plt.subplots(1, 2, figsize=(14, 4.5))\n    # Action histogram\n    ax = axes[0]\n    pc = pred_df.group_by(\"action\").count().sort(\"count\", descending=True)\n    a = pc[\"action\"].to_list()[:20][::-1]; c = pc[\"count\"].to_list()[:20][::-1]\n    ax.barh(a, c, color=\"#64B5CD\"); ax.set_title(\"Predicted Actions\"); ax.set_xlabel(\"Count\"); ax.grid(axis=\"x\", alpha=0.3)\n    # Duration\n    ax = axes[1]\n    pdur = pred_df.with_columns((pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\"))[\"duration\"].to_list()\n    ax.hist(pdur, bins=30, color=\"#C44E52\", edgecolor=\"black\", linewidth=0.3)\n    if pdur:\n        ax.axvline(np.median(pdur), color=\"black\", linestyle=\"--\", alpha=0.7, label=f\"Median={np.median(pdur):.1f}\")\n        ax.legend()\n    ax.set_title(\"Predicted Duration (frames)\"); ax.set_xlabel(\"Frames\"); ax.set_ylabel(\"Count\"); ax.grid(axis=\"y\", alpha=0.3)\n    plt.tight_layout(); plt.show()\n\nvisualize_predictions(submission_test)\n\n# Save submission\nordered = list(cfg.submission_schema.keys())\nsubmission_out = submission_test.select(ordered).with_row_index(cfg.row_id_col)\nsubmission_out.write_csv(cfg.submission_file)\nprint(f\"Saved submission: {cfg.submission_file} | rows: {len(submission_out)}\")\n\n# Final concise summary\nprint(\"Summary\")\nprint(f\"- Train videos (filtered): {int(train_subset['video_id'].n_unique())}\")\nprint(f\"- Labs: {int(train_subset['lab_id'].n_unique())}\")\nprint(f\"- Annotations: {len(solution)}\")\nprint(f\"- Unique behaviors: {int(solution['action'].n_unique())}\")\nprint(f\"- Predictions: {len(submission_out)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-22T19:11:55.420439Z","iopub.execute_input":"2025-09-22T19:11:55.421033Z","iopub.status.idle":"2025-09-22T19:12:04.503099Z","shell.execute_reply.started":"2025-09-22T19:11:55.421010Z","shell.execute_reply":"2025-09-22T19:12:04.502465Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Requires: `submission_out` (from Cell 6), `test`\nqa_actions = (\n    submission_out\n    .with_columns((pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\"))\n    .group_by(\"action\")\n    .agg([\n        pl.count().alias(\"num_events\"),\n        pl.col(\"duration\").mean().alias(\"mean_dur\"),\n        pl.col(\"duration\").median().alias(\"median_dur\"),\n    ])\n    .sort(\"num_events\", descending=True)\n)\nqa_videos = (\n    submission_out\n    .group_by(\"video_id\")\n    .agg([\n        pl.count().alias(\"num_events\"),\n        pl.col(\"action\").n_unique().alias(\"num_actions\"),\n    ])\n    .sort(\"num_events\", descending=True)\n)\n\nprint(\"Per-action summary (top 20):\")\ndisplay(qa_actions.head(20).to_pandas())\n\nprint(\"\\nPer-video summary (top 20):\")\ndisplay(qa_videos.head(20).to_pandas())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-22T19:13:59.779616Z","iopub.execute_input":"2025-09-22T19:13:59.780404Z","iopub.status.idle":"2025-09-22T19:13:59.810992Z","shell.execute_reply.started":"2025-09-22T19:13:59.780378Z","shell.execute_reply":"2025-09-22T19:13:59.810402Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Requires: predict_without_ml, cfg, test, priors (per_lab, global_w, med_lab, med_glob, timing_lab, timing_glob)\n\ndef try_params(min_len: int, gap_close: int):\n    preds = predict_without_ml(\n        test, \"test\", cfg,\n        priors_per_lab=per_lab, priors_global=global_w,\n        meddur_per_lab=med_lab, meddur_global=med_glob,\n        timing_lab=timing_lab, timing_global=timing_glob,\n        prior_scope=\"mixed\", use_windows=True,\n        min_len=min_len, gap_close=gap_close,\n        p_min=0.05, cap=0.05\n    )\n    return preds.select(pl.count()).item()\n\ngrid = [(12,3), (10,3), (10,5), (8,5), (6,5)]\nresults = []\nfor ml, gc_ in grid:\n    try:\n        n = try_params(ml, gc_)\n        results.append((ml, gc_, int(n)))\n        print(f\"min_len={ml}, gap_close={gc_} -> events={n}\")\n    except Exception as e:\n        print(f\"min_len={ml}, gap_close={gc_} -> failed: {e}\")\n\nprint(\"\\nGrid results (sorted by events desc):\")\nfor ml, gc_, n in sorted(results, key=lambda x: x[2], reverse=True):\n    print(f\"- min_len={ml}, gap_close={gc_}, events={n}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-22T19:14:13.163380Z","iopub.execute_input":"2025-09-22T19:14:13.163692Z","iopub.status.idle":"2025-09-22T19:14:55.376719Z","shell.execute_reply.started":"2025-09-22T19:14:13.163670Z","shell.execute_reply":"2025-09-22T19:14:55.375883Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Parameter Tuning: Segment Consolidation\n\nWe lightly tune two post-processing parameters that govern segment quality:\n- min_len: minimum event length in frames (filters spurious short segments).\n- gap_close: maximum gap to merge same-action neighboring segments (smooths fragmented detections).\n\nProcedure:\n1) Sweep a small grid of (min_len, gap_close) to gauge event counts.\n2) Pick the configuration that best matches a target density: not too sparse (risking recall) and not too dense (risking precision).\n3) Regenerate submission with the chosen configuration.","metadata":{}},{"cell_type":"code","source":"# If you already ran the light sweep cell, reuse `results`. Otherwise, define a default grid fallback.\ntry:\n    assert 'results' in globals() and len(results) > 0\nexcept Exception:\n    grid = [(12,3), (10,3), (10,5), (8,5), (6,5)]\n    results = []\n    for ml, gc_ in grid:\n        try:\n            n = predict_without_ml(\n                test, \"test\", cfg,\n                priors_per_lab=per_lab, priors_global=global_w,\n                meddur_per_lab=med_lab, meddur_global=med_glob,\n                timing_lab=timing_lab, timing_global=timing_glob,\n                prior_scope=\"mixed\", use_windows=True,\n                min_len=ml, gap_close=gc_, p_min=0.05, cap=0.05\n            ).select(pl.count()).item()\n            results.append((ml, gc_, int(n)))\n        except Exception as e:\n            results.append((ml, gc_, -1))\n\n# Heuristic selector: prefer the top third by event count, then choose the middle of that subset\nvalid = [(ml, gc_, n) for (ml, gc_, n) in results if n >= 0]\nif not valid:\n    print(\"No valid results from sweep; keeping previous submission.\")\nelse:\n    valid_sorted = sorted(valid, key=lambda x: x[2], reverse=True)\n    top_k = max(1, len(valid_sorted) // 3)\n    candidates = valid_sorted[:top_k]\n    chosen = candidates[len(candidates)//2]  # middle of the top slice\n    chosen_min_len, chosen_gap = chosen[0], chosen[1]\n\n    print(\"Sweep results (top 10 by events):\")\n    for r in valid_sorted[:10]:\n        print(f\"min_len={r[0]}, gap_close={r[1]} -> events={r[2]}\")\n\n    print(f\"\\nChosen params: min_len={chosen_min_len}, gap_close={chosen_gap}\")\n\n    # Re-run predictions with chosen params\n    submission_tuned = predict_without_ml(\n        test, \"test\", cfg,\n        priors_per_lab=per_lab, priors_global=global_w,\n        meddur_per_lab=med_lab, meddur_global=med_glob,\n        timing_lab=timing_lab, timing_global=timing_glob,\n        prior_scope=\"mixed\", use_windows=True,\n        min_len=chosen_min_len, gap_close=chosen_gap, p_min=0.05, cap=0.05\n    )\n\n    # Quick QA\n    qa = (\n        submission_tuned\n        .with_columns((pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\"))\n        .group_by(\"action\")\n        .agg([pl.count().alias(\"num_events\"), pl.col(\"duration\").median().alias(\"median_dur\")])\n        .sort(\"num_events\", descending=True)\n    )\n    print(\"\\nPer-action (top 15):\")\n    display(qa.head(15).to_pandas())\n\n    # Save tuned submission\n    ordered_cols = list(cfg.submission_schema.keys())\n    submission_tuned_out = submission_tuned.select(ordered_cols).with_row_index(cfg.row_id_col)\n    tuned_path = \"submission.csv\"  # overwrite as final\n    submission_tuned_out.write_csv(tuned_path)\n    print(f\"\\nSaved tuned submission: {tuned_path} | rows: {len(submission_tuned_out)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-22T19:16:11.567440Z","iopub.execute_input":"2025-09-22T19:16:11.568005Z","iopub.status.idle":"2025-09-22T19:16:20.182090Z","shell.execute_reply.started":"2025-09-22T19:16:11.567982Z","shell.execute_reply":"2025-09-22T19:16:20.181419Z"}},"outputs":[],"execution_count":null}]}