{"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":"none","dataSources":[{"sourceId":59156,"databundleVersionId":13719397,"sourceType":"competition"}],"dockerImageVersionId":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"%%writefile baseline.py\n\nfrom __future__ import annotations\n\nimport argparse\nimport json\nimport logging\nimport os\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 polars as pl\nfrom tqdm.auto import tqdm\n\n# ========================\n# Config\n# ========================\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\nlogger = logging.getLogger(__name__)\n\nclass HostVisibleError(Exception): pass\n\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# ========================\n# Utils & Validators\n# ========================\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 _norm_triplet(agent: str | int, target: str | int, action: str) -> str:\n    return f\"{_norm_mouse_id(agent)},{_norm_mouse_id(target)},{action}\"\n\ndef _range_frames(start: int, stop: int) -> Iterable[int]:\n    return range(start, stop)  # [start, stop)\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\n\ndef split_interval(s: int, e: int, parts: int) -> List[Tuple[int,int]]:\n    if parts <= 1: return [(s,e)]\n    L = e - s\n    step = L // parts\n    rem = L % parts\n    out = []\n    cur = s\n    for i in range(parts):\n        extra = 1 if i < rem else 0\n        nxt = cur + step + extra\n        out.append((cur, min(nxt, e)))\n        cur = nxt\n    return out\n\ndef largest_remainder_allocation(total: int, weights: List[float]) -> List[int]:\n    if total <= 0 or not weights: return [0]*len(weights)\n    s = sum(weights) or 1.0\n    w = [x/s for x in weights]\n    raw = [total*x for x in w]\n    base = [int(v) for v in raw]\n    remainder = total - sum(base)\n    if remainder > 0:\n        fr = sorted([(i, raw[i]-base[i]) for i in range(len(w))], key=lambda x: x[1], reverse=True)\n        for i in range(remainder):\n            base[fr[i % len(w)][0]] += 1\n    return base\n\n# ========================\n# Metrics (F-beta)\n# ========================\n\ndef single_lab_f1(lab_solution: pl.DataFrame, lab_submission: pl.DataFrame, beta: float = 1.0) -> float:\n    label_frames: Dict[str, Set[int]] = defaultdict(set)\n    for row in lab_solution.to_dicts():\n        label_frames[row[\"label_key\"]].update(_range_frames(row[\"start_frame\"], row[\"stop_frame\"]))\n\n    active_by_video: Dict[int, Set[str]] = {}\n    for row in lab_solution.select([\"video_id\", \"behaviors_labeled\"]).unique().to_dicts():\n        s: Set[str] = set()\n        for item in safe_json_loads(row[\"behaviors_labeled\"]):\n            parts = [p.strip() for p in str(item).replace(\"'\", \"\").split(\",\")]\n            if len(parts) == 3:\n                a, t, act = parts\n                s.add(_norm_triplet(a, t, act))\n        active_by_video[int(row[\"video_id\"])] = s\n\n    prediction_frames: Dict[str, Set[int]] = defaultdict(set)\n    for video_id in lab_solution[\"video_id\"].unique():\n        active = active_by_video.get(int(video_id), set())\n        predicted_mouse_pairs: Dict[str, Set[int]] = defaultdict(set)\n        for row in lab_submission.filter(pl.col(\"video_id\") == video_id).to_dicts():\n            triple_norm = _norm_triplet(row[\"agent_id\"], row[\"target_id\"], row[\"action\"])\n            if triple_norm not in active:\n                continue\n            pred_key = row[\"prediction_key\"]\n            agent_target = f\"{row['agent_id']},{row['target_id']}\"\n            new_frames = set(_range_frames(row[\"start_frame\"], row[\"stop_frame\"]))\n            new_frames -= prediction_frames[pred_key]\n            if predicted_mouse_pairs[agent_target] & new_frames:\n                raise HostVisibleError(\"Multiple predictions for the same frame from one agent/target pair\")\n            prediction_frames[pred_key].update(new_frames)\n            predicted_mouse_pairs[agent_target].update(new_frames)\n\n    tps: Dict[str, int] = defaultdict(int)\n    fns: Dict[str, int] = defaultdict(int)\n    fps: Dict[str, int] = defaultdict(int)\n    distinct_actions: Set[str] = set()\n\n    for key, pred_frames in prediction_frames.items():\n        action = key.split(\"_\")[-1]\n        distinct_actions.add(action)\n        gt_frames = label_frames.get(key, set())\n        tps[action] += len(pred_frames & gt_frames)\n        fns[action] += len(gt_frames - pred_frames)\n        fps[action] += len(pred_frames - gt_frames)\n\n    for key, gt_frames in label_frames.items():\n        action = key.split(\"_\")[-1]\n        distinct_actions.add(action)\n        if key not in prediction_frames:\n            fns[action] += len(gt_frames)\n\n    if not distinct_actions:\n        return 0.0\n\n    beta2 = beta * beta\n    f_scores: List[float] = []\n    for action in distinct_actions:\n        tp, fn, fp = tps[action], fns[action], fps[action]\n        denom = (1 + beta2) * tp + beta2 * fn + fp\n        f_scores.append(0.0 if denom == 0 else (1 + beta2) * tp / denom)\n    return sum(f_scores) / len(f_scores)\n\ndef mouse_fbeta(solution: pl.DataFrame, submission: pl.DataFrame, beta: float = 1.0, cfg: Optional[Config] = None) -> float:\n    cfg = cfg or Config()\n    solution = validate_schema(solution, cfg.solution_schema, \"Solution\")\n    submission = validate_schema(submission, cfg.submission_schema, \"Submission\")\n    validate_frame_ranges(solution, \"Solution\")\n    validate_frame_ranges(submission, \"Submission\")\n\n    solution_videos = solution[\"video_id\"].unique()\n    submission = submission.filter(pl.col(\"video_id\").is_in(solution_videos))\n\n    def add_key(df: pl.DataFrame, col_name: str) -> pl.DataFrame:\n        return df.with_columns(\n            pl.concat_str(\n                [\n                    pl.col(\"video_id\").cast(pl.Utf8),\n                    pl.col(\"agent_id\").cast(pl.Utf8),\n                    pl.col(\"target_id\").cast(pl.Utf8),\n                    pl.col(\"action\"),\n                ],\n                separator=\"_\",\n            ).alias(col_name)\n        )\n\n    solution = add_key(solution, \"label_key\")\n    submission = add_key(submission, \"prediction_key\")\n\n    lab_scores: List[float] = []\n    for lab_id in solution[\"lab_id\"].unique():\n        lab_solution = solution.filter(pl.col(\"lab_id\") == lab_id)\n        lab_videos = lab_solution[\"video_id\"].unique()\n        lab_submission = submission.filter(pl.col(\"video_id\").is_in(lab_videos))\n        lab_scores.append(single_lab_f1(lab_solution, lab_submission, beta=beta))\n\n    return sum(lab_scores) / len(lab_scores) if lab_scores else 0.0\n\ndef score(solution: pl.DataFrame, submission: pl.DataFrame, row_id_column_name: str = \"\", beta: float = 1.0, cfg: Optional[Config] = None) -> float:\n    if row_id_column_name:\n        solution = solution.drop(row_id_column_name, strict=False)\n        submission = submission.drop(row_id_column_name, strict=False)\n    return mouse_fbeta(solution, submission, beta=beta, cfg=cfg)\n\n# ========================\n# Build solution + video spans\n# ========================\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        if lab_id.startswith(\"MABe22\"): continue\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            logger.warning(\"No annotations for %s\", annot_path)\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            for col, dtype in (cfg.solution_schema).items():\n                if col in annot.columns and annot[col].dtype != dtype:\n                    annot = annot.with_columns(pl.col(col).cast(dtype))\n            annot = annot.select([c for c in cfg.solution_schema.keys() if c in annot.columns])\n            records.append(annot)\n        except Exception as e:\n            logger.error(\"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    solution = validate_schema(solution, cfg.solution_schema, \"Solution\")\n    return solution\n\ndef build_video_spans(dataset: pl.DataFrame, split: str, cfg: Optional[Config] = None) -> Dict[int, Tuple[int,int]]:\n    \"\"\"\n    Map video_id -> (min_frame, max_frame+1).\n    \"\"\"\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 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            logger.warning(\"Span read failed for %s: %s\", path, e)\n    return spans\n\n# ========================\n# Priors (duration & timing)\n# ========================\n\ndef compute_action_priors(solution: pl.DataFrame, eps: float = 1.0) -> Tuple[Dict[str, Dict[str, float]], Dict[str, float], Dict[str, Dict[str, int]], Dict[str, int]]:\n    \"\"\"\n    Returns:\n      per_lab_weight: {lab: {action: weight_share}}\n      global_weight: {action: weight_share}\n      per_lab_med_dur: {lab: {action: median_duration_frames}}\n      global_med_dur: {action: median_duration_frames}\n    \"\"\"\n    sol = solution.with_columns((pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"dur\"))\n    # shares\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    # median durations\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\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]]) -> Tuple[Dict[str, Dict[str, float]], Dict[str, float]]:\n    \"\"\"\n    Median start percentile per (lab, action) and global.\n    start_pct = (start_frame - video_start)/(video_stop - video_start)\n    \"\"\"\n    # attach start_pct\n    def start_pct_func(row) -> float:\n        vid = int(row[\"video_id\"])\n        if vid not in video_spans: 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\n    rows = []\n    for r in solution.select([\"lab_id\",\"action\",\"video_id\",\"start_frame\"]).to_dicts():\n        rows.append({\"lab_id\": r[\"lab_id\"], \"action\": r[\"action\"], \"start_pct\": start_pct_func(r)})\n    df = pl.DataFrame(rows)\n    by_lab = df.group_by([\"lab_id\",\"action\"]).median().select([\"lab_id\",\"action\",\"start_pct\"])\n    per_lab: Dict[str, Dict[str, float]] = 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_: Dict[str, float] = {r[\"action\"]: float(r[\"start_pct\"]) for r in g.to_dicts()}\n    return per_lab, global_\n\n# ========================\n# Tracking features → windows\n# ========================\n\ndef _detect_tracking_schema(df: pl.DataFrame) -> Optional[Tuple[str,str,str,str]]:\n    \"\"\"\n    Try to detect columns: id_col, frame_col, x_col, y_col.\n    Returns None if cannot detect.\n    \"\"\"\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\"]\n    y_candidates = [\"y\",\"y_pos\",\"y_position\",\"y_mm\"]\n\n    frame_col = next((c for c in frame_candidates if c in df.columns), None)\n    id_col = next((c for c in id_candidates if c in df.columns), None)\n    x_col = next((c for c in x_candidates if c in df.columns), None)\n    y_col = next((c for c in y_candidates if c in df.columns), None)\n    if all([frame_col, id_col, x_col, y_col]):\n        return id_col, frame_col, x_col, y_col\n    return None\n\nimport numpy as np  # đặt ở phần import đầu file\n\ndef _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    \"\"\"\n    Tính features cho một cặp (agent,target) bằng NumPy:\n      - dist  = sqrt((ax-bx)^2 + (ay-by)^2)\n      - rel_speed = speed_a - speed_b, với speed = sqrt(dx^2 + dy^2) khung-kề-khung\n      - ddist = diff(dist)\n    Trả về Polars DataFrame: [\"frame\",\"dist\",\"rel_speed\",\"ddist\"] đã sort theo frame.\n    downsample: lấy mỗi N khung (N>=1). Giá trị 2–3 giúp tăng tốc đáng kể.\n    \"\"\"\n    # --- auto-detect schema ---\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]):\n        return None\n\n    # --- chuẩn hoá ID ---\n    a_id = _strip_mouse_prefix(agent_raw)\n    t_id = _strip_mouse_prefix(target_raw)\n\n    # --- lấy tối thiểu các cột cần thiết & chuyển sang pandas để dùng NumPy ---\n    # (to_pandas trên 4 cột nhỏ rất nhanh; tránh join nhiều lần trong Polars)\n    pdf = df.select([frame_col, id_col, x_col, y_col]).to_pandas()\n    # ép kiểu an toàn\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:\n        return None\n\n    # ưu tiên một bản ghi / frame / mouse (nếu trùng) để merge gọn\n    a.drop_duplicates(subset=[frame_col], keep=\"first\", inplace=True)\n    b.drop_duplicates(subset=[frame_col], keep=\"first\", inplace=True)\n\n    merged = a.merge(b, on=frame_col, how=\"inner\", suffixes=(\"_a\", \"_b\"))\n    if merged.empty:\n        return None\n    merged.sort_values(frame_col, inplace=True)\n\n    # --- NumPy vectors ---\n    ax = merged[f\"{x_col}_a\"].to_numpy(dtype=np.float64, copy=False)\n    ay = merged[f\"{y_col}_a\"].to_numpy(dtype=np.float64, copy=False)\n    bx = merged[f\"{x_col}_b\"].to_numpy(dtype=np.float64, copy=False)\n    by = merged[f\"{y_col}_b\"].to_numpy(dtype=np.float64, copy=False)\n    frames = merged[frame_col].to_numpy(dtype=np.int64, copy=False)\n\n    # downsample nếu cần\n    if 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:\n            return None\n\n    dx = ax - bx\n    dy = ay - by\n    dist = np.sqrt(dx*dx + dy*dy)\n\n    dax = np.diff(ax, prepend=ax[0])\n    day = np.diff(ay, prepend=ay[0])\n    dbx = np.diff(bx, prepend=bx[0])\n    dby = np.diff(by, prepend=by[0])\n    speed_a = np.sqrt(dax*dax + day*day)\n    speed_b = np.sqrt(dbx*dbx + dby*dby)\n\n    rel_speed = speed_a - speed_b\n    ddist = np.diff(dist, prepend=dist[0])\n\n    feat = pl.DataFrame(\n        {\n            \"frame\": frames,\n            \"dist\": dist,\n            \"rel_speed\": rel_speed,\n            \"ddist\": ddist,\n        }\n    ).sort(\"frame\")\n\n    return feat\n\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    \"\"\"\n    Local percentiles per pair for robust thresholds.\n    Condition:\n      (dist <= Pq_dist) or (rel_speed >= Pq_rel and ddist <= Pq_ddist)\n    \"\"\"\n    if len(feat) == 0:\n        return []\n    # quantiles\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))  # typically negative\n    # boolean mask\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: Optional[List[int]] = None\n    for i, flag in enumerate(mask):\n        if flag and run is None:\n            run = [frames[i], frames[i]]\n        elif flag and run is not None:\n            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:\n                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:\n            windows.append((s,e))\n    return merge_intervals(windows)\n\n# ========================\n# Advanced baseline prediction\n# ========================\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    # earlier → smaller\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    # if an action extremely rare, cap its share\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    \"\"\"\n    Allocate contiguous segments across union(windows) sequentially following ordered_actions.\n    Length per action ~ max(weight*total, median_duration) but clipped by remaining frames.\n    \"\"\"\n    # flatten windows into a sequence of positions\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    # sort\n    segments = sorted(segments, key=lambda x: (x[1], x[2], x[0]))\n    # remove too short\n    segments = [seg for seg in segments if seg[2] - seg[1] >= min_len]\n    if not segments: return []\n    # merge same-action with small gap\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 = 10,\n                       gap_close: int = 5,\n                       p_min: float = 0.03,\n                       cap: float = 0.02) -> pl.DataFrame:\n    \"\"\"\n    Advanced heuristic:\n      - Optionally create proximity windows from tracking per (agent,target) using pairwise features.\n      - Allocate within union(windows) following action order by timing prior; lengths by weight & median duration.\n      - Smooth & gap-close small segments; rare-action clipping.\n    \"\"\"\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        if lab_id.startswith(\"MABe22\"): continue\n        video_id: int = row[\"video_id\"]\n        path = track_dir / lab_id / f\"{video_id}.parquet\"\n        if not path.exists():\n            logger.warning(\"Tracking file not found: %s\", path)\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            # parse behaviors\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:  # no actions labeled for this video\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            # per (agent,target)\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                # choose priors\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                # rare-action clipping & renorm\n                weights = _clip_rare_actions(w_map, actions, p_min=p_min, cap=cap)\n\n                # order actions by timing prior\n                ordered_actions = _order_actions_by_timing(\n                    actions, str(lab_id), timing_lab or {}, timing_global or {}, canonical\n                )\n\n                # windows\n                windows: List[Tuple[int,int]]\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:\n                    continue\n\n                # allocate segments along windows\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\n                segs = _smooth_segments(segs, min_len=min_len, gap_close=gap_close)\n\n                # emit rows\n                for a, s, e in segs:\n                    if e > s:\n                        records.append((\n                            video_id,\n                            _norm_mouse_id(agent), _norm_mouse_id(target),\n                            a, int(s), int(e)\n                        ))\n\n        except Exception as e:\n            logger.error(\"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\n\n# ========================\n# CLI / Main\n# ========================\n\ndef run_validate(cfg: Config, beta: float, prior_scope: str, eps: float,\n                 use_windows: bool, min_len: int, gap_close: int, p_min: float, cap: float) -> float:\n    logger.info(\"Loading train data for validation: %s\", cfg.train_csv)\n    train = pl.read_csv(cfg.train_csv)\n    train_subset = train.filter(~pl.col(\"lab_id\").str.starts_with(\"MABe22\"))\n\n    logger.info(\"Building solution dataframe & spans...\")\n    solution = create_solution_df(train_subset, cfg)\n    spans = build_video_spans(train_subset, \"train\", cfg)\n\n    logger.info(\"Computing priors (eps=%.2f) & timing...\", eps)\n    per_lab, global_w, med_lab, med_glob = compute_action_priors(solution, eps=eps)\n    timing_lab, timing_glob = compute_timing_priors(solution, spans)\n\n    logger.info(\"Generating predictions (advanced)...\")\n    submission_train = predict_without_ml(\n        train_subset, \"train\", 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=prior_scope,\n        use_windows=use_windows, min_len=min_len, gap_close=gap_close,\n        p_min=p_min, cap=cap\n    )\n\n    logger.info(\"Scoring (beta=%.3f)...\", beta)\n    val_score = score(solution, submission_train, cfg.row_id_col, beta=beta, cfg=cfg)\n    print(f\"[RESULT] Validation F1: {val_score:.6f}\")\n    return val_score\n\ndef run_submit(cfg: Config, prior_scope: str, eps: float,\n               use_windows: bool, min_len: int, gap_close: int, p_min: float, cap: float) -> None:\n    logger.info(\"Loading train (for priors) and test data...\")\n    train = pl.read_csv(cfg.train_csv)\n    train_subset = train.filter(~pl.col(\"lab_id\").str.starts_with(\"MABe22\"))\n    solution = create_solution_df(train_subset, cfg)\n    spans = build_video_spans(train_subset, \"train\", cfg)\n    per_lab, global_w, med_lab, med_glob = compute_action_priors(solution, eps=eps)\n    timing_lab, timing_glob = compute_timing_priors(solution, spans)\n\n    test = pl.read_csv(cfg.test_csv)\n\n    logger.info(\"Generating predictions (advanced, test)...\")\n    submission_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        prior_scope=prior_scope,\n        use_windows=use_windows, min_len=min_len, gap_close=gap_close,\n        p_min=p_min, cap=cap\n    )\n\n    ordered = list(cfg.submission_schema.keys())\n    submission_test = submission_test.select(ordered).with_row_index(cfg.row_id_col)\n\n    logger.info(\"Saving submission to %s\", cfg.submission_file)\n    submission_test.write_csv(cfg.submission_file)\n\n    try:\n        with open(cfg.submission_file, \"r\") as f:\n            for _ in range(10):\n                line = f.readline()\n                if not line: break\n                logger.info(\"SUBMISSION PREVIEW | %s\", line.strip())\n    except Exception as e:\n        logger.warning(\"Preview failed: %s\", e)\n\ndef main() -> None:\n    parser = argparse.ArgumentParser(description=\"MABe Mouse Behavior: baseline++ (priors + windows + timing)\")\n    parser.add_argument(\"--data-root\", type=str, default=None, help=\"Dataset root directory\")\n    parser.add_argument(\"--beta\", type=float, default=1.0, help=\"F-beta value\")\n    parser.add_argument(\"--mode\", choices=[\"validate\", \"submit\", \"all\"], default=\"all\", help=\"Run mode\")\n    parser.add_argument(\"--submission\", type=str, default=None, help=\"Submission CSV output path\")\n    parser.add_argument(\"--prior-scope\", choices=[\"lab\", \"global\", \"mixed\"], default=\"mixed\",\n                        help=\"Use lab-level priors, global priors, or per-lab with global fallback (default)\")\n    parser.add_argument(\"--eps\", type=float, default=1.0, help=\"Laplace smoothing (frame units) for priors\")\n    parser.add_argument(\"--no-windows\", action=\"store_true\", help=\"Disable proximity windows (use full video)\")\n    parser.add_argument(\"--min-len\", type=int, default=15, help=\"Minimum segment/window length (frames)\")\n    parser.add_argument(\"--gap-close\", type=int, default=3, help=\"Merge same-action gaps up to this many frames\")\n    parser.add_argument(\"--p-min\", type=float, default=0.05, help=\"Rare-action min prior threshold\")\n    parser.add_argument(\"--cap\", type=float, default=0.05, help=\"Rare-action maximum share if below p-min\")\n    parser.add_argument(\"-v\", \"--verbose\", action=\"count\", default=1, help=\"Increase verbosity (-v, -vv)\")\n    args = parser.parse_args()\n\n    setup_logging(args.verbose)\n    warnings.filterwarnings(\"ignore\")\n\n    cfg = Config(\n        data_root=Path(args.data_root) if args.data_root else Config().data_root,\n        submission_file=args.submission if args.submission else Config().submission_file,\n        row_id_col=Config().row_id_col,\n    )\n\n    val = None\n    if args.mode in (\"validate\", \"all\"):\n        val = run_validate(cfg, beta=args.beta, prior_scope=args.prior_scope, eps=args.eps,\n                           use_windows=not args.no_windows, min_len=args.min_len, gap_close=args.gap_close,\n                           p_min=args.p_min, cap=args.cap)\n        logger.info(\"Validation F1: %.6f\", val)\n    if args.mode in (\"submit\", \"all\"):\n        run_submit(cfg, prior_scope=args.prior_scope, eps=args.eps,\n                   use_windows=not args.no_windows, min_len=args.min_len, gap_close=args.gap_close,\n                   p_min=args.p_min, cap=args.cap)\n\nif __name__ == \"__main__\":\n    main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-20T20:34:17.483537Z","iopub.execute_input":"2025-09-20T20:34:17.483993Z","iopub.status.idle":"2025-09-20T20:34:17.510606Z","shell.execute_reply.started":"2025-09-20T20:34:17.483961Z","shell.execute_reply":"2025-09-20T20:34:17.509565Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python baseline.py --mode submit --prior-scope mixed --eps 0.5 -vv\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-20T20:45:25.765266Z","iopub.execute_input":"2025-09-20T20:45:25.765712Z"}},"outputs":[],"execution_count":null}]}