{"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":"from __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, Counter\nimport itertools\n\nimport numpy as np\nimport pandas as pd\nimport polars as pl\nfrom tqdm.auto import tqdm\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom matplotlib.patches import Rectangle\nimport matplotlib.patches as mpatches\n\n# Set visualization style\nplt.style.use('seaborn-v0_8-whitegrid')\nsns.set_palette(\"husl\")\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=\"%(ascus)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)\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\n# ========================\n# Core Functions\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    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\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    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\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    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\ndef _strip_mouse_prefix(s: str | int) -> str:\n    s = str(s)\n    return s[5:] if s.startswith(\"mouse\") else s\n\n# ========================\n# ENHANCED FEATURE EXTRACTION\n# ========================\n\ndef _pair_features_enhanced(df: pl.DataFrame, agent_raw: str, target_raw: str, downsample: int = 1) -> Optional[pl.DataFrame]:\n    \"\"\"Enhanced feature extraction with behavioral indicators\"\"\"\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    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:\n        return None\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    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    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    # Basic features\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    # Enhanced features\n    angle_between = np.arctan2(dy, dx)\n    angle_change = np.abs(np.diff(angle_between, prepend=angle_between[0]))\n    \n    # Acceleration\n    acc_a = np.diff(speed_a, prepend=speed_a[0])\n    acc_b = np.diff(speed_b, prepend=speed_b[0])\n    rel_acc = acc_a - acc_b\n    \n    # Proximity indicators\n    dist_pct = np.percentile(dist, [10, 25, 50, 75, 90])\n    very_close = (dist <= dist_pct[0]).astype(float)\n    close = (dist <= dist_pct[1]).astype(float)\n    approaching = (ddist < 0).astype(float)\n    \n    # Temporal smoothing\n    window_size = min(5, len(dist))\n    if window_size > 1:\n        from scipy.ndimage import uniform_filter1d\n        dist_smooth = uniform_filter1d(dist, size=window_size)\n        speed_trend_a = uniform_filter1d(speed_a, size=window_size)\n        speed_trend_b = uniform_filter1d(speed_b, size=window_size)\n    else:\n        dist_smooth = dist\n        speed_trend_a = speed_a\n        speed_trend_b = speed_b\n    \n    # Interaction intensity\n    interaction_score = (very_close * 2 + close) * (speed_a + speed_b) / 2\n    \n    feat = pl.DataFrame({\n        \"frame\": frames,\n        \"dist\": dist,\n        \"rel_speed\": rel_speed,\n        \"ddist\": ddist,\n        \"angle_between\": angle_between,\n        \"angle_change\": angle_change,\n        \"acc_a\": acc_a,\n        \"acc_b\": acc_b,\n        \"rel_acc\": rel_acc,\n        \"very_close\": very_close,\n        \"close\": close,\n        \"approaching\": approaching,\n        \"dist_smooth\": dist_smooth,\n        \"speed_trend_a\": speed_trend_a,\n        \"speed_trend_b\": speed_trend_b,\n        \"interaction_score\": interaction_score,\n        \"speed_a\": speed_a,\n        \"speed_b\": speed_b\n    }).sort(\"frame\")\n\n    return feat\n\n# ========================\n# BEHAVIOR-SPECIFIC DETECTION\n# ========================\n\ndef _make_behavior_specific_windows(feat: pl.DataFrame, behavior: str, min_len: int) -> List[Tuple[int,int]]:\n    \"\"\"Create windows tailored to specific behaviors\"\"\"\n    if len(feat) == 0:\n        return []\n    \n    # Behavior-specific parameters\n    behavior_configs = {\n        'approach': {\n            'dist_q': 0.70,\n            'speed_q': 0.40,\n            'use_angle': True,\n            'use_approaching': True,\n            'primary_cond': lambda: (pl.col(\"approaching\") == 1) & (pl.col(\"dist\") <= pl.col(\"dist\").quantile(0.70))\n        },\n        'chase': {\n            'dist_q': 0.60,\n            'speed_q': 0.70,\n            'use_acceleration': True,\n            'use_rel_speed': True,\n            'primary_cond': lambda: (pl.col(\"rel_speed\") >= pl.col(\"rel_speed\").quantile(0.70)) & (pl.col(\"acc_a\") >= pl.col(\"acc_a\").quantile(0.60))\n        },\n        'attack': {\n            'dist_q': 0.15,\n            'speed_q': 0.60,\n            'use_proximity': True,\n            'use_interaction': True,\n            'primary_cond': lambda: (pl.col(\"very_close\") == 1) & (pl.col(\"interaction_score\") >= pl.col(\"interaction_score\").quantile(0.80))\n        },\n        'mount': {\n            'dist_q': 0.10,\n            'speed_q': 0.30,\n            'use_proximity': True,\n            'use_low_speed': True,\n            'primary_cond': lambda: (pl.col(\"very_close\") == 1) & (pl.col(\"speed_a\") <= pl.col(\"speed_a\").quantile(0.40))\n        },\n        'avoid': {\n            'dist_q': 0.80,\n            'speed_q': 0.60,\n            'inverse_approach': True,\n            'primary_cond': lambda: (pl.col(\"approaching\") == 0) & (pl.col(\"rel_speed\") <= pl.col(\"rel_speed\").quantile(0.30))\n        },\n        'chaseattack': {\n            'dist_q': 0.30,\n            'speed_q': 0.75,\n            'use_acceleration': True,\n            'use_proximity': True,\n            'primary_cond': lambda: (pl.col(\"close\") == 1) & (pl.col(\"rel_speed\") >= pl.col(\"rel_speed\").quantile(0.75)) & (pl.col(\"acc_a\") >= pl.col(\"acc_a\").quantile(0.70))\n        },\n        'submit': {\n            'dist_q': 0.20,\n            'speed_q': 0.20,\n            'use_proximity': True,\n            'use_low_speed': True,\n            'primary_cond': lambda: (pl.col(\"close\") == 1) & (pl.col(\"speed_a\") <= pl.col(\"speed_a\").quantile(0.30)) & (pl.col(\"speed_b\") <= pl.col(\"speed_b\").quantile(0.30))\n        }\n    }\n    \n    config = behavior_configs.get(behavior, {\n        'dist_q': 0.50, \n        'speed_q': 0.60,\n        'primary_cond': lambda: (pl.col(\"dist\") <= pl.col(\"dist\").quantile(0.50)) & (pl.col(\"rel_speed\") >= pl.col(\"rel_speed\").quantile(0.60))\n    })\n    \n    # Apply behavior-specific condition\n    try:\n        cond = config['primary_cond']()\n        mask = feat.select(cond.alias(\"m\")).to_series().to_list()\n    except Exception:\n        # Fallback to simple condition\n        qd = float(feat[\"dist\"].quantile(config['dist_q']))\n        qs = float(feat[\"rel_speed\"].quantile(config['speed_q']))\n        mask = feat.select(((pl.col(\"dist\") <= qd) & (pl.col(\"rel_speed\") >= qs)).alias(\"m\")).to_series().to_list()\n    \n    frames = feat[\"frame\"].to_list()\n    \n    # Convert mask to windows\n    windows = []\n    run = 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    \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    \n    return merge_intervals(windows)\n\n# ========================\n# SEQUENCE MODELING\n# ========================\n\ndef _model_behavior_sequences(segments: List[Tuple[str,int,int]], video_frames: int) -> List[Tuple[str,int,int]]:\n    \"\"\"Model sequential dependencies between behaviors\"\"\"\n    if not segments:\n        return segments\n    \n    # Define realistic behavior transitions\n    transitions = {\n        'approach': ['chase', 'attack', 'mount', 'avoid'],\n        'chase': ['attack', 'chaseattack', 'mount', 'approach'],\n        'attack': ['mount', 'chase', 'chaseattack'],\n        'chaseattack': ['attack', 'mount', 'chase'],\n        'mount': ['submit', 'attack'],\n        'avoid': ['approach'],  # Can lead to approach if situation changes\n        'submit': []  # Usually terminal\n    }\n    \n    # Sort segments by start time\n    segments = sorted(segments, key=lambda x: x[1])\n    improved_segments = []\n    \n    for i, (action, start, end) in enumerate(segments):\n        should_include = True\n        \n        if i > 0 and improved_segments:\n            prev_action, prev_start, prev_end = improved_segments[-1]\n            gap = start - prev_end\n            \n            # Check for realistic sequences\n            if 0 <= gap <= 45:  # Within 45 frames (~1.5 seconds)\n                valid_transitions = transitions.get(prev_action, [])\n                \n                # Handle illogical sequences\n                if action == prev_action and gap <= 10:\n                    # Merge same behaviors that are very close\n                    improved_segments[-1] = (action, prev_start, max(prev_end, end))\n                    continue\n                elif valid_transitions and action not in valid_transitions:\n                    # Check for impossible transitions\n                    impossible_pairs = [\n                        ('attack', 'approach'), ('mount', 'approach'), \n                        ('submit', 'chase'), ('submit', 'attack')\n                    ]\n                    if (prev_action, action) in impossible_pairs:\n                        # Skip this segment as it's unlikely\n                        continue\n        \n        if should_include:\n            improved_segments.append((action, start, end))\n    \n    return improved_segments\n\n# ========================\n# ENSEMBLE METHODS\n# ========================\n\ndef merge_ensemble_predictions(predictions_list: List[pl.DataFrame], consensus_threshold: float = 0.6) -> pl.DataFrame:\n    \"\"\"Merge multiple prediction sets using consensus voting\"\"\"\n    if not predictions_list:\n        return pl.DataFrame(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    \n    if len(predictions_list) == 1:\n        return predictions_list[0]\n    \n    # Collect all predictions\n    all_records = []\n    for pred_df in predictions_list:\n        for record in pred_df.to_dicts():\n            all_records.append(record)\n    \n    # Group by key and find consensus\n    grouped = defaultdict(list)\n    for record in all_records:\n        key = (record['video_id'], record['agent_id'], record['target_id'], record['action'])\n        grouped[key].append((record['start_frame'], record['stop_frame']))\n    \n    # Apply consensus threshold\n    min_votes = max(1, int(len(predictions_list) * consensus_threshold))\n    final_records = []\n    \n    for (vid, agent, target, action), intervals in grouped.items():\n        if len(intervals) >= min_votes:\n            # Merge overlapping intervals\n            merged = merge_intervals(intervals)\n            for start, end in merged:\n                final_records.append((vid, agent, target, action, start, end))\n    \n    if not final_records:\n        # Fallback to best single prediction\n        return predictions_list[0]\n    \n    return pl.DataFrame(\n        final_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\n# ========================\n# PARAMETER OPTIMIZATION\n# ========================\n\ndef optimize_parameters(solution_df: pl.DataFrame, train_sample: pl.DataFrame, cfg: Config) -> Dict:\n    \"\"\"Simple parameter optimization based on data characteristics\"\"\"\n    \n    # Analyze data characteristics\n    solution_df = solution_df.with_columns((pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\"))\n    \n    behavior_stats = {}\n    for action in solution_df['action'].unique():\n        action_data = solution_df.filter(pl.col('action') == action)\n        behavior_stats[action] = {\n            'count': len(action_data),\n            'mean_duration': action_data['duration'].mean(),\n            'median_duration': action_data['duration'].median(),\n            'std_duration': action_data['duration'].std() or 1.0\n        }\n    \n    # Adaptive parameter selection based on data\n    avg_duration = solution_df['duration'].mean()\n    std_duration = solution_df['duration'].std() or 1.0\n    \n    if avg_duration < 20:  # Short behaviors\n        return {\n            'min_len': max(5, int(avg_duration * 0.3)),\n            'gap_close': 2,\n            'p_min': 0.03,\n            'cap': 0.04\n        }\n    elif avg_duration > 50:  # Long behaviors\n        return {\n            'min_len': max(15, int(avg_duration * 0.2)),\n            'gap_close': 8,\n            'p_min': 0.07,\n            'cap': 0.08\n        }\n    else:  # Medium behaviors\n        return {\n            'min_len': max(10, int(avg_duration * 0.25)),\n            'gap_close': 4,\n            'p_min': 0.05,\n            'cap': 0.06\n        }\n\n# ========================\n# MAIN PREDICTION FUNCTIONS\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    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_improved(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,\n                               use_enhanced_features: bool = True,\n                               use_behavior_specific: bool = True,\n                               use_sequences: bool = True) -> pl.DataFrame:\n    \"\"\"Improved prediction with all enhancements\"\"\"\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            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                # Get 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                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                # Generate windows for each behavior\n                all_windows = []\n                if use_windows:\n                    if use_enhanced_features:\n                        feat = _pair_features_enhanced(trk, _norm_mouse_id(agent), _norm_mouse_id(target))\n                    else:\n                        feat = _pair_features_basic(trk, _norm_mouse_id(agent), _norm_mouse_id(target))\n                    \n                    if feat is None:\n                        all_windows = [(start_frame, stop_frame)]\n                    else:\n                        if use_behavior_specific:\n                            # Create behavior-specific windows\n                            for action in actions:\n                                behavior_windows = _make_behavior_specific_windows(feat, action, min_len)\n                                all_windows.extend(behavior_windows)\n                        else:\n                            all_windows = _make_windows_basic(feat, min_len)\n                        \n                        if not all_windows:\n                            all_windows = [(start_frame, stop_frame)]\n                else:\n                    all_windows = [(start_frame, stop_frame)]\n\n                all_windows = merge_intervals(all_windows)\n                allowed_total = sum(e - s for s,e in all_windows)\n                if allowed_total <= 0:\n                    continue\n\n                # Allocate segments\n                segs = _allocate_segments_in_windows(\n                    windows=all_windows,\n                    ordered_actions=ordered_actions,\n                    weights=weights,\n                    med_dur=md_map,\n                    total_frames=allowed_total\n                )\n\n                # Smooth segments\n                segs = _smooth_segments(segs, min_len=min_len, gap_close=gap_close)\n                \n                # Apply sequence modeling\n                if use_sequences:\n                    segs = _model_behavior_sequences(segs, video_frames)\n\n                # Add to records\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        return pl.DataFrame(schema=cfg.submission_schema)\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\ndef _pair_features_basic(df: pl.DataFrame, agent_raw: str, target_raw: str, downsample: int = 1) -> Optional[pl.DataFrame]:\n    \"\"\"Basic feature extraction (fallback)\"\"\"\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    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:\n        return None\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    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    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\ndef _make_windows_basic(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    \"\"\"Basic window detection\"\"\"\n    if len(feat) == 0:\n        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: 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\ndef predict_with_ensemble(dataset: pl.DataFrame, data_split: str, cfg: Config, **kwargs) -> pl.DataFrame:\n    \"\"\"Ensemble prediction with multiple configurations\"\"\"\n    \n    # Multiple configurations optimized for different scenarios\n    configs = [\n        # Sensitive detection\n        {\n            'min_len': 8, 'gap_close': 2, 'p_min': 0.02, 'cap': 0.03, \n            'use_enhanced_features': True, 'use_behavior_specific': True, 'use_sequences': True\n        },\n        # Balanced detection  \n        {\n            'min_len': 12, 'gap_close': 3, 'p_min': 0.04, 'cap': 0.05,\n            'use_enhanced_features': True, 'use_behavior_specific': True, 'use_sequences': True\n        },\n        # Conservative detection\n        {\n            'min_len': 18, 'gap_close': 5, 'p_min': 0.06, 'cap': 0.07,\n            'use_enhanced_features': True, 'use_behavior_specific': False, 'use_sequences': True\n        },\n        # Fallback configuration\n        {\n            'min_len': 15, 'gap_close': 4, 'p_min': 0.05, 'cap': 0.06,\n            'use_enhanced_features': False, 'use_behavior_specific': False, 'use_sequences': False\n        }\n    ]\n    \n    predictions = []\n    \n    for i, config in enumerate(configs):\n        print(f\"Running ensemble model {i+1}/{len(configs)}...\")\n        try:\n            # Merge config with existing kwargs\n            merged_kwargs = {**kwargs}\n            merged_kwargs.update(config)\n            \n            pred = predict_without_ml_improved(dataset, data_split, cfg, **merged_kwargs)\n            if len(pred) > 0:\n                predictions.append(pred)\n        except Exception as e:\n            logger.warning(f\"Ensemble model {i+1} failed: {e}\")\n            continue\n    \n    if not predictions:\n        # Fallback to basic prediction\n        return predict_without_ml_improved(dataset, data_split, cfg, **kwargs)\n    \n    # Merge using consensus\n    return merge_ensemble_predictions(predictions, consensus_threshold=0.5)\n\n# ========================\n# MAIN EXECUTION\n# ========================\n\ndef main():\n    print(\"MABe Mouse Behavior Detection - Enhanced Model\")\n    print(\"=\" * 60)\n    \n    # Setup\n    setup_logging(verbosity=1)\n    warnings.filterwarnings(\"ignore\")\n    \n    cfg = Config()\n    \n    # Load data\n    print(\"\\nLoading datasets...\")\n    train = pl.read_csv(cfg.train_csv)\n    train_subset = train.filter(~pl.col(\"lab_id\").str.starts_with(\"MABe22\"))\n    print(f\"Loaded {len(train)} training samples ({len(train_subset)} after filtering)\")\n    \n    print(\"\\nBuilding solution dataframe...\")\n    solution = create_solution_df(train_subset, cfg)\n    print(f\"Processed {len(solution):,} annotations\")\n    \n    print(\"\\nComputing video spans...\")\n    spans = build_video_spans(train_subset, \"train\", cfg)\n    print(f\"Computed spans for {len(spans)} videos\")\n    \n    print(\"\\nComputing priors...\")\n    per_lab, global_w, med_lab, med_glob = compute_action_priors(solution, eps=0.3)\n    timing_lab, timing_glob = compute_timing_priors(solution, spans)\n    print(f\"Computed priors for {len(global_w)} actions across {len(per_lab)} labs\")\n    \n    print(\"\\nOptimizing parameters...\")\n    optimal_params = optimize_parameters(solution, train_subset.sample(min(500, len(train_subset))), cfg)\n    print(f\"Optimal parameters: {optimal_params}\")\n    \n    print(\"\\nLoading test data...\")\n    test = pl.read_csv(cfg.test_csv)\n    print(f\"Loaded {len(test)} test samples\")\n    \n    print(\"\\nGenerating predictions with ensemble...\")\n    submission_test = predict_with_ensemble(\n        test, \"test\", cfg,\n        priors_per_lab=per_lab, \n        priors_global=global_w,\n        meddur_per_lab=med_lab, \n        meddur_global=med_glob,\n        timing_lab=timing_lab, \n        timing_global=timing_glob,\n        prior_scope=\"mixed\",\n        use_windows=True,\n        **optimal_params\n    )\n    print(f\"Generated {len(submission_test):,} predictions\")\n    \n    # Add row_id and save\n    ordered = list(cfg.submission_schema.keys())\n    submission_test = submission_test.select(ordered).with_row_index(cfg.row_id_col)\n    \n    print(f\"\\nSaving submission to {cfg.submission_file}...\")\n    submission_test.write_csv(cfg.submission_file)\n    print(\"Submission saved successfully!\")\n    \n    # Print summary\n    if len(submission_test) > 0:\n        actions = submission_test.group_by(\"action\").count().sort(\"count\", descending=True)\n        print(f\"\\nPrediction Summary:\")\n        for row in actions.head(5).to_dicts():\n            print(f\"  {row['action']}: {row['count']} predictions\")\n        \n        durations = (submission_test[\"stop_frame\"] - submission_test[\"start_frame\"]).to_list()\n        print(f\"\\nDuration Statistics:\")\n        print(f\"  Mean: {np.mean(durations):.1f} frames\")\n        print(f\"  Median: {np.median(durations):.1f} frames\")\n        print(f\"  Videos covered: {submission_test['video_id'].n_unique()}\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-22T16:01:25.066348Z","iopub.execute_input":"2025-09-22T16:01:25.066760Z","iopub.status.idle":"2025-09-22T16:03:11.877516Z","shell.execute_reply.started":"2025-09-22T16:01:25.066707Z","shell.execute_reply":"2025-09-22T16:03:11.876578Z"}},"outputs":[],"execution_count":null}]}