{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.11.13"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"}],"dockerImageVersionId":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false},"papermill":{"default_parameters":{},"duration":29841.37378,"end_time":"2025-12-08T13:04:44.759592","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2025-12-08T04:47:23.385812","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"#  MABe Challenge - Social Action Recognition in Mice\n\n### Overview\n> In this competition, you’ll develop machine learning models to recognize behaviors in mice based on their movements, providing new insights into animal social structures and advancing behavioral science research.\n\n### Description\n> Animal social behavior is complex. Species from ants to wolves to mice form social groups where they build nests, raise their young, care for their groupmates, and defend their territory. Studying these behaviors teaches us about the brain and the evolution of behavior, but the work has usually required subjective, time-consuming documentation of animals' actions. ML advancements now let us automate this process, supporting large-scale behavioral studies in the wild and in the lab.\n\n> But even automated systems suffer from limited training data and poor generalizability. In current methods, an experimenter must hand-label hundreds of new training examples to automate recognition of a new behavior, which makes studying rare behaviors a challenge. And models trained within one research group usually fail when applied to data from other studies, meaning there is no guarantee that two labs are really studying the same behavior.\n\n> This competition challenges you to build models to identify over 30 different social and non-social behaviors in pairs and groups of co-housed mice, based on markerless motion capture of their movements in top-down video recordings. The dataset includes over 400 hours of footage from 20+ behavioral recording systems, all carefully labeled frame-by-frame by experts. Your goal is to recognize these behaviors as accurately as a trained human observer while overcoming the inherent variability arising from the use of different data collection equipment and motion capture pipelines.\n\nYour work will help scientists automate behavior analysis and better understand animal social structures. These models may be deployed across numerous labs in neuroscience, computational biology, ethology, and ecology to create a foundation for future ML and behavior research.","metadata":{"papermill":{"duration":0.004379,"end_time":"2025-12-08T04:47:27.613751","exception":false,"start_time":"2025-12-08T04:47:27.609372","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"This notebook tackles the MABe challenge by building a separate GBDT ensemble model for each unique `body_parts_tracked` configuration. \n\n**Core Strategy:**\n1.  **Grouped by Tracker:** Loop through each `body_parts_tracked` string.\n2.  **FPS-Aware Features:** Generate advanced temporal and spatial features (`transform_single`, `transform_pair`). All window sizes, lags, and spans are *scaled* by the video's FPS to ensure features are time-invariant (`_scale` function).\n3.  **Stratified Subsampling:** Use a `StratifiedSubsetClassifier` wrapper to train the GBDTs on a large, stratified subsample of the full dataset to manage memory and time.\n4.  **Ensemble Model:** For each behavior, train an ensemble of LGBM, XGBoost, and CatBoost models.\n5.  **Adaptive Prediction:** Use temporal smoothing, adaptive per-action probability thresholds, and minimum duration filtering to generate final event segments (`predict_multiclass_adaptive`).","metadata":{"papermill":{"duration":0.003076,"end_time":"2025-12-08T04:47:27.620059","exception":false,"start_time":"2025-12-08T04:47:27.616983","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# ==================== IMPORTS & SETUP ====================\n\nverbose = True\n\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nimport itertools\nimport warnings\nimport json\nimport os, random\nimport gc, re, math\nimport lightgbm\nfrom collections import defaultdict\nimport polars as pl\nfrom scipy import signal, stats\nfrom scipy.spatial.distance import euclidean\nfrom typing import Dict, Optional, Tuple, List\nfrom time import perf_counter\nfrom sklearn.base import ClassifierMixin, BaseEstimator, clone\nfrom sklearn.model_selection import cross_val_predict, GroupKFold, StratifiedKFold\nfrom sklearn.pipeline import make_pipeline\nfrom sklearn.model_selection import StratifiedShuffleSplit\nfrom sklearn.metrics import f1_score\nfrom sklearn.preprocessing import StandardScaler\n\nwarnings.filterwarnings('ignore')\nUSE_GPU = (\"KAGGLE_KERNEL_RUN_TYPE\" in os.environ) and (__import__(\"shutil\").which(\"nvidia-smi\") is not None)\nprint(f'Using GPU?  {USE_GPU}')\n\nfrom xgboost import XGBClassifier\nfrom catboost import CatBoostClassifier\n\n# Optional imports for advanced features\ntry:\n    import pywt\n    HAS_PYWT = True\n    print(\"PyWavelets available\")\nexcept ImportError:\n    HAS_PYWT = False\n    print(\"PyWavelets not available - wavelet features disabled\")\n\ntry:\n    import sklearn_crfsuite\n    HAS_CRF = True\n    print(\"sklearn-crfsuite available\")\nexcept ImportError:\n    HAS_CRF = False\n    print(\"sklearn-crfsuite not available - CRF post-processing disabled\")\n\ntry:\n    from fastdtw import fastdtw\n    HAS_DTW = True\n    print(\"fastdtw available\")\nexcept ImportError:\n    HAS_DTW = False\n    print(\"fastdtw not available - template matching disabled\")\n\nSEED = 1234\n\n# ==================== SEED EVERYTHING ====================\n\nos.environ[\"PYTHONHASHSEED\"] = str(SEED)\n\nrnd = np.random.RandomState(SEED)\nrandom.seed(SEED)\nnp.random.seed(SEED)\n\ndef _make_lgbm(**kw):\n    kw.setdefault(\"random_state\", SEED)\n    kw.setdefault(\"feature_fraction_seed\", SEED)\n    kw.setdefault(\"data_random_seed\", SEED)\n    kw.setdefault(\"device\", 'gpu' if USE_GPU else 'cpu')\n    return lightgbm.LGBMClassifier(**kw)\n\ndef _make_xgb(**kw):\n    kw.setdefault(\"random_state\", SEED)\n    kw.setdefault(\"tree_method\", \"gpu_hist\" if USE_GPU else \"hist\")\n    return XGBClassifier(**kw)\n\ndef _make_cb(**kw):\n    kw.setdefault(\"random_seed\", SEED)\n    if USE_GPU:\n        kw.setdefault(\"task_type\", \"GPU\")\n        kw.setdefault(\"devices\", \"0\")\n    else:\n        kw.setdefault(\"task_type\", \"CPU\")\n    return CatBoostClassifier(**kw)\n\n\n# ==================== STRATIFIED SUBSET CLASSIFIERS ====================\n\nclass StratifiedSubsetClassifierWEval(ClassifierMixin, BaseEstimator):\n    def __init__(self,\n                 estimator,\n                 n_samples=None,\n                 random_state: int = 42,\n                 valid_size: float = 0.10,\n                 val_cap_ratio: float = 0.25,\n                 es_rounds: \"int|str\" = \"auto\",\n                 es_metric: str = \"auto\"):\n        self.estimator = estimator\n        self.n_samples = (int(n_samples) if (n_samples is not None) else None)\n        self.random_state = random_state\n        self.valid_size = float(valid_size)\n        self.val_cap_ratio = float(val_cap_ratio)\n        self.es_rounds = es_rounds\n        self.es_metric = es_metric\n\n    def fit(self, X: pd.DataFrame, y):\n        y = np.asarray(y)\n        n_total = len(y)\n        assert n_total == len(X)\n\n        tr_idx, va_idx = self._compute_train_val_indices(y, n_total)\n        Xtr = X.iloc[tr_idx]\n        ytr = y[tr_idx]\n\n        Xtr = Xtr.to_numpy(np.float32, copy=False)\n\n        Xva = yva = None\n        if va_idx is not None and len(va_idx) > 0:\n            Xva = X.iloc[va_idx].to_numpy(np.float32, copy=False)\n            yva = y[va_idx]\n\n        pos_rate = None\n        if yva is not None and len(yva) > 0:\n            pos_rate = float(np.mean(yva == 1))\n\n        metric = self._choose_metric(pos_rate)\n        patience = self._choose_patience(pos_rate)\n\n        if self._is_xgb(self.estimator):\n            n_pos = max(1, int((ytr == 1).sum()))\n            n_neg = max(1, len(ytr) - n_pos)\n            self.estimator.set_params(scale_pos_weight=(n_neg / n_pos))\n            self.estimator.set_params(eval_metric=metric)\n\n        elif self._is_catboost(self.estimator):\n            try:\n                self.estimator.set_params(auto_class_weights=\"Balanced\")\n            except Exception:\n                pass\n            try:\n                self.estimator.set_params(eval_metric=metric)\n            except Exception:\n                pass\n\n        has_valid = (Xva is not None and len(yva) > 0)\n        if has_valid and self._is_xgb(self.estimator):\n            import xgboost as xgb\n            self.estimator.fit(\n                Xtr, ytr,\n                eval_set=[(Xva, yva)],\n                verbose=False,\n                callbacks=[xgb.callback.EarlyStopping(\n                    rounds=int(patience),\n                    metric_name=metric,\n                    data_name=\"validation_0\",\n                    save_best=True\n                )]\n            )\n        elif has_valid and self._is_catboost(self.estimator):\n            from catboost import Pool\n            self.estimator.set_params(\n                use_best_model=True,\n                od_type=\"Iter\",\n                od_wait=int(patience),\n                custom_metric=[\"PRAUC:type=Classic;hints=skip_train~true\"],\n            )\n            self.estimator.fit(\n                Xtr, ytr,\n                eval_set=Pool(Xva, yva),\n                verbose=False,\n                metric_period=50\n            )\n        else:\n            self.estimator.fit(Xtr, ytr)\n\n        self.classes_ = getattr(self.estimator, \"classes_\", np.array([0, 1]))\n        self._tr_idx_ = tr_idx\n        self._va_idx_ = va_idx\n        self._pos_rate_ = pos_rate\n        return self\n\n    def predict_proba(self, X: pd.DataFrame):\n        return self.estimator.predict_proba(X)\n\n    def predict(self, X: pd.DataFrame):\n        return self.estimator.predict(X)\n\n    def _compute_train_val_indices(self, y: np.ndarray, n_total: int):\n        rng = np.random.default_rng(self.random_state)\n        n_classes = np.unique(y).size\n\n        def full_data_split():\n            if self.valid_size <= 0 or n_classes < 2:\n                idx = rng.permutation(n_total)\n                return idx, None\n            sss = StratifiedShuffleSplit(n_splits=1, test_size=self.valid_size, random_state=self.random_state)\n            tr, va = next(sss.split(np.zeros(n_total, dtype=np.int8), y))\n            return tr, va\n\n        if self.n_samples is None or self.n_samples >= n_total:\n            return full_data_split()\n\n        sss_tr = StratifiedShuffleSplit(n_splits=1, train_size=self.n_samples, random_state=self.random_state)\n        tr_idx, rest_idx = next(sss_tr.split(np.zeros(n_total, dtype=np.int8), y))\n        remaining = len(rest_idx)\n\n        min_val_needed = int(np.ceil(self.n_samples * max(self.valid_size, 0.0)))\n        val_cap = max(min_val_needed, int(round(self.val_cap_ratio * self.n_samples)))\n        want_val = min(remaining, val_cap)\n\n        y_rest = y[rest_idx]\n        if remaining < min_val_needed or np.unique(y_rest).size < 2 or self.valid_size <= 0:\n            return full_data_split()\n\n        sss_val = StratifiedShuffleSplit(n_splits=1, train_size=want_val, random_state=self.random_state)\n        try:\n            va_sel, _ = next(sss_val.split(np.zeros(remaining, dtype=np.int8), y_rest))\n        except ValueError:\n            return full_data_split()\n\n        va_idx = rest_idx[va_sel]\n        return tr_idx, va_idx\n\n    def _choose_metric(self, pos_rate=0.01) -> str:\n        if self.es_metric != \"auto\":\n            return self.es_metric\n        if pos_rate is None or pos_rate == 0.0 or pos_rate == 1.0:\n            return \"logloss\" if self._is_xgb(self.estimator) else \"Logloss\"\n        return \"aucpr\" if self._is_xgb(self.estimator) else \"PRAUC:type=Classic\"\n\n    def _choose_patience(self, pos_rate: Optional[float]) -> int:\n        if isinstance(self.es_rounds, int):\n            return self.es_rounds\n        try:\n            n_estimators = (int(self.estimator.get_params().get(\"n_estimators\", 200))\n                            if self._is_xgb(self.estimator)\n                            else int(self.estimator.get_params().get(\"iterations\", 500)))\n        except Exception:\n            n_estimators = 200\n        base = max(30, int(round(0.20 * (n_estimators or 200))))\n        if pos_rate is None:\n            return base\n        if pos_rate < 0.005:\n            return int(round(base * 1.75))\n        if pos_rate < 0.02:\n            return int(round(base * 1.40))\n        return base\n\n    @staticmethod\n    def _is_xgb(est):\n        name = est.__class__.__name__.lower()\n        mod = getattr(est, \"__module__\", \"\")\n        return \"xgb\" in name or \"xgboost\" in mod or hasattr(est, \"get_xgb_params\")\n\n    @staticmethod\n    def _is_catboost(est):\n        name = est.__class__.__name__.lower()\n        mod = getattr(est, \"__module__\", \"\")\n        return \"catboost\" in name or \"catboost\" in mod or hasattr(est, \"get_all_params\")\n\n\nclass StratifiedSubsetClassifier(ClassifierMixin, BaseEstimator):\n    def __init__(self, estimator, n_samples, random_state=SEED):\n        self.estimator = estimator\n        self.n_samples = n_samples and int(n_samples)\n        self.random_state = random_state\n\n    def fit(self, X, y):\n        y = np.asarray(y)\n        n_total = len(y)\n\n        if self.n_samples is None or self.n_samples >= n_total:\n            rng = np.random.default_rng(self.random_state)\n            idx = rng.permutation(n_total)\n        else:\n            sss = StratifiedShuffleSplit(\n                n_splits=1, train_size=self.n_samples, random_state=self.random_state\n            )\n            idx, _ = next(sss.split(np.zeros(n_total, dtype=np.int8), y))\n\n        Xn = X.iloc[idx]\n        Xn = Xn.to_numpy(np.float32, copy=False)\n        yn = y[idx]\n\n        self.estimator.fit(Xn, yn)\n        self.classes_ = getattr(self.estimator, \"classes_\", np.array([0, 1]))\n        return self\n\n    def predict_proba(self, X):\n        return self.estimator.predict_proba(X)\n\n    def predict(self, X):\n        return self.estimator.predict(X)\n\n\n# ==================== SCORING FUNCTIONS ====================\n\nclass HostVisibleError(Exception):\n    pass\n\ndef single_lab_f1(lab_solution: pl.DataFrame, lab_submission: pl.DataFrame, beta: float = 1) -> float:\n    label_frames: defaultdict[str, set[int]] = defaultdict(set)\n    prediction_frames: defaultdict[str, set[int]] = defaultdict(set)\n\n    for row in lab_solution.to_dicts():\n        label_frames[row['label_key']].update(range(row['start_frame'], row['stop_frame']))\n\n    for video in lab_solution['video_id'].unique():\n        active_labels: str = lab_solution.filter(pl.col('video_id') == video)['behaviors_labeled'].first()\n        active_labels: set[str] = set(json.loads(active_labels))\n        predicted_mouse_pairs: defaultdict[str, set[int]] = defaultdict(set)\n\n        for row in lab_submission.filter(pl.col('video_id') == video).to_dicts():\n            if ','.join([str(row['agent_id']), str(row['target_id']), row['action']]) not in active_labels:\n                continue\n\n            new_frames = set(range(row['start_frame'], row['stop_frame']))\n            new_frames = new_frames.difference(prediction_frames[row['prediction_key']])\n            prediction_pair = ','.join([str(row['agent_id']), str(row['target_id'])])\n            if predicted_mouse_pairs[prediction_pair].intersection(new_frames):\n                raise HostVisibleError('Multiple predictions for the same frame from one agent/target pair')\n            prediction_frames[row['prediction_key']].update(new_frames)\n            predicted_mouse_pairs[prediction_pair].update(new_frames)\n\n    tps = defaultdict(int)\n    fns = defaultdict(int)\n    fps = defaultdict(int)\n    for key, pred_frames in prediction_frames.items():\n        action = key.split('_')[-1]\n        matched_label_frames = label_frames[key]\n        tps[action] += len(pred_frames.intersection(matched_label_frames))\n        fns[action] += len(matched_label_frames.difference(pred_frames))\n        fps[action] += len(pred_frames.difference(matched_label_frames))\n\n    distinct_actions = set()\n    for key, 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(frames)\n\n    action_f1s = []\n    for action in distinct_actions:\n        if tps[action] + fns[action] + fps[action] == 0:\n            action_f1s.append(0)\n        else:\n            action_f1s.append((1 + beta**2) * tps[action] / ((1 + beta**2) * tps[action] + beta**2 * fns[action] + fps[action]))\n    return sum(action_f1s) / len(action_f1s)\n\ndef mouse_fbeta(solution: pd.DataFrame, submission: pd.DataFrame, beta: float = 1) -> float:\n    if len(solution) == 0 or len(submission) == 0:\n        raise ValueError('Missing solution or submission data')\n\n    expected_cols = ['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']\n\n    for col in expected_cols:\n        if col not in solution.columns:\n            raise ValueError(f'Solution is missing column {col}')\n        if col not in submission.columns:\n            raise ValueError(f'Submission is missing column {col}')\n\n    solution: pl.DataFrame = pl.DataFrame(solution)\n    submission: pl.DataFrame = pl.DataFrame(submission)\n    assert (solution['start_frame'] <= solution['stop_frame']).all()\n    assert (submission['start_frame'] <= submission['stop_frame']).all()\n    solution_videos = set(solution['video_id'].unique())\n    submission = submission.filter(pl.col('video_id').is_in(solution_videos))\n\n    solution = solution.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('label_key'),\n    )\n    submission = submission.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('prediction_key'),\n    )\n\n    lab_scores = []\n    for lab in solution['lab_id'].unique():\n        lab_solution = solution.filter(pl.col('lab_id') == lab).clone()\n        lab_videos = set(lab_solution['video_id'].unique())\n        lab_submission = submission.filter(pl.col('video_id').is_in(lab_videos)).clone()\n        lab_scores.append(single_lab_f1(lab_solution, lab_submission, beta=beta))\n\n    return sum(lab_scores) / len(lab_scores)\n\ndef score(solution: pd.DataFrame, submission: pd.DataFrame, row_id_column_name: str, beta: float = 1) -> float:\n    solution = solution.drop(row_id_column_name, axis='columns', errors='ignore')\n    submission = submission.drop(row_id_column_name, axis='columns', errors='ignore')\n    return mouse_fbeta(solution, submission, beta=beta)\n\n\n# ==================== DATA LOADING ====================\n\ntrain = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/train.csv')\n\n# Drop likely-sleeping MABe22 clips\ntrain = train.loc[~(train['lab_id'].astype(str).str.contains('MABe22', na=False) &\n                    train['mouse1_condition'].astype(str).str.lower().eq('lights on'))].copy()\n\ntrain['n_mice'] = 4 - train[['mouse1_strain', 'mouse2_strain', 'mouse3_strain', 'mouse4_strain']].isna().sum(axis=1)\n\ntest = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/test.csv')\ntest['sleeping'] = (\n    test['lab_id'].astype(str).str.contains('MABe22', na=False) &\n    test['mouse1_condition'].astype(str).str.lower().eq('lights on')\n)\ntest['n_mice'] = 4 - test[['mouse1_strain', 'mouse2_strain', 'mouse3_strain', 'mouse4_strain']].isna().sum(axis=1)\n\nbody_parts_tracked_list = list(np.unique(train.body_parts_tracked))\n\ndrop_body_parts = ['headpiece_bottombackleft', 'headpiece_bottombackright', 'headpiece_bottomfrontleft',\n                   'headpiece_bottomfrontright', 'headpiece_topbackleft', 'headpiece_topbackright',\n                   'headpiece_topfrontleft', 'headpiece_topfrontright', 'spine_1', 'spine_2',\n                   'tail_middle_1', 'tail_middle_2', 'tail_midpoint']\n\n_sex_cols = [f'mouse{i}_sex' for i in range(1, 5)]\n_train_sex_lut = (train[['video_id'] + _sex_cols].drop_duplicates('video_id')\n                  .set_index('video_id').to_dict('index'))\n_test_sex_lut = (test[['video_id'] + _sex_cols].drop_duplicates('video_id')\n                 .set_index('video_id').to_dict('index'))\n_FEATURE_TEMPLATES = {}\n\nprint(f\"Train videos: {len(train)}, Test videos: {len(test)}\")\nprint(f\"Unique body part configs: {len(body_parts_tracked_list)}\")\n\n\n\n# ==================== DATA GENERATOR ====================\n\ndef generate_mouse_data(dataset, traintest, traintest_directory=None,\n                        generate_single=True, generate_pair=True):\n    assert traintest in ['train', 'test']\n    if traintest_directory is None:\n        traintest_directory = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking\"\n\n    def _to_num(x):\n        if isinstance(x, (int, np.integer)):\n            return int(x)\n        m = re.search(r'(\\d+)$', str(x))\n        return int(m.group(1)) if m else None\n\n    for _, row in dataset.iterrows():\n        lab_id = row.lab_id\n        video_id = row.video_id\n        fps = float(row.frames_per_second)\n        n_mice = int(row.n_mice)\n        arena_w = float(row.get('arena_width_cm', np.nan))\n        arena_h = float(row.get('arena_height_cm', np.nan))\n        sleeping = bool(getattr(row, 'sleeping', False))\n        arena_shape = row.get('arena_shape', 'rectangular')\n\n        if not isinstance(row.behaviors_labeled, str):\n            continue\n\n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        vid = pd.read_parquet(path)\n        if len(np.unique(vid.bodypart)) > 5:\n            vid = vid.query(\"~ bodypart.isin(@drop_body_parts)\")\n        pvid = vid.pivot(columns=['mouse_id', 'bodypart'], index='video_frame', values=['x', 'y'])\n        del vid\n        pvid = pvid.reorder_levels([1, 2, 0], axis=1).T.sort_index().T\n        pvid = (pvid / float(row.pix_per_cm_approx)).astype('float32', copy=False)\n\n        avail = list(pvid.columns.get_level_values('mouse_id').unique())\n        avail_set = set(avail) | set(map(str, avail)) | {f\"mouse{_to_num(a)}\" for a in avail if _to_num(a) is not None}\n\n        def _resolve(agent_str):\n            m = re.search(r'(\\d+)$', str(agent_str))\n            cand = [agent_str]\n            if m:\n                n = int(m.group(1))\n                cand = [n, n - 1, str(n), f\"mouse{n}\", agent_str]\n            for c in cand:\n                if c in avail_set:\n                    if c in set(avail):\n                        return c\n                    for a in avail:\n                        if str(a) == str(c) or f\"mouse{_to_num(a)}\" == str(c):\n                            return a\n            return None\n\n        vb = json.loads(row.behaviors_labeled)\n        vb = sorted(list({b.replace(\"'\", \"\") for b in vb}))\n        vb = pd.DataFrame([b.split(',') for b in vb], columns=['agent', 'target', 'action'])\n        vb['agent'] = vb['agent'].astype(str)\n        vb['target'] = vb['target'].astype(str)\n        vb['action'] = vb['action'].astype(str).str.lower()\n\n        if traintest == 'train':\n            try:\n                annot = pd.read_parquet(path.replace('train_tracking', 'train_annotation'))\n            except FileNotFoundError:\n                continue\n\n        def _mk_meta(index, agent_id, target_id):\n            m = pd.DataFrame({\n                'lab_id': lab_id,\n                'video_id': video_id,\n                'agent_id': agent_id,\n                'target_id': target_id,\n                'video_frame': index.astype('int32', copy=False),\n                'frames_per_second': np.float32(fps),\n                'sleeping': sleeping,\n                'arena_shape': arena_shape,\n                'arena_width_cm': np.float32(arena_w),\n                'arena_height_cm': np.float32(arena_h),\n                'n_mice': np.int8(n_mice),\n            })\n            for c in ('lab_id', 'video_id', 'agent_id', 'target_id', 'arena_shape'):\n                m[c] = m[c].astype('category')\n            return m\n\n        # SINGLE\n        if generate_single:\n            vb_single = vb.query(\"target == 'self'\")\n            for agent_str in pd.unique(vb_single['agent']):\n                col_lab = _resolve(agent_str)\n                if col_lab is None:\n                    continue\n                actions = sorted(vb_single.loc[vb_single['agent'].eq(agent_str), 'action'].unique().tolist())\n                if not actions:\n                    continue\n\n                single = pvid.loc[:, col_lab]\n                meta_df = _mk_meta(single.index, agent_str, 'self')\n\n                if traintest == 'train':\n                    a_num = _to_num(col_lab)\n                    y = pd.DataFrame(False, index=single.index.astype('int32', copy=False), columns=actions)\n                    a_sub = annot.query(\"(agent_id == @a_num) & (target_id == @a_num)\")\n                    for i in range(len(a_sub)):\n                        ar = a_sub.iloc[i]\n                        a = str(ar.action).lower()\n                        if a in y.columns:\n                            y.loc[int(ar['start_frame']):int(ar['stop_frame']), a] = True\n                    yield 'single', single, meta_df, y\n                else:\n                    yield 'single', single, meta_df, actions\n\n        # PAIR\n        if generate_pair:\n            vb_pair = vb.query(\"target != 'self'\")\n            if len(vb_pair) > 0:\n                allowed_pairs = set(map(tuple, vb_pair[['agent', 'target']].itertuples(index=False, name=None)))\n\n                for agent_num, target_num in itertools.permutations(\n                        np.unique(pvid.columns.get_level_values('mouse_id')), 2):\n                    agent_str = f\"mouse{_to_num(agent_num)}\"\n                    target_str = f\"mouse{_to_num(target_num)}\"\n                    if (agent_str, target_str) not in allowed_pairs:\n                        continue\n\n                    a_col = _resolve(agent_str)\n                    b_col = _resolve(target_str)\n                    if a_col is None or b_col is None:\n                        continue\n\n                    actions = sorted(\n                        vb_pair.query(\"(agent == @agent_str) & (target == @target_str)\")['action'].unique().tolist()\n                    )\n                    if not actions:\n                        continue\n\n                    pair_xy = pd.concat([pvid[a_col], pvid[b_col]], axis=1, keys=['A', 'B'])\n                    meta_df = _mk_meta(pair_xy.index, agent_str, target_str)\n\n                    if traintest == 'train':\n                        a_num = _to_num(a_col)\n                        b_num = _to_num(b_col)\n                        y = pd.DataFrame(False, index=pair_xy.index.astype('int32', copy=False), columns=actions)\n                        a_sub = annot.query(\"(agent_id == @a_num) & (target_id == @b_num)\")\n                        for i in range(len(a_sub)):\n                            ar = a_sub.iloc[i]\n                            a = str(ar.action).lower()\n                            if a in y.columns:\n                                y.loc[int(ar['start_frame']):int(ar['stop_frame']), a] = True\n                        yield 'pair', pair_xy, meta_df, y\n                    else:\n                        yield 'pair', pair_xy, meta_df, actions\n\n\n# ==================== HELPER FUNCTIONS ====================\n\ndef safe_rolling(series, window, func, min_periods=None):\n    \"\"\"Safe rolling operation with NaN handling\"\"\"\n    if min_periods is None:\n        min_periods = max(1, window // 4)\n    return series.rolling(window, min_periods=min_periods, center=True).apply(func, raw=True)\n\ndef _scale(n_frames_at_30fps, fps, ref=30.0):\n    \"\"\"Scale a frame count defined at 30 fps to the current video's fps.\"\"\"\n    return max(1, int(round(n_frames_at_30fps * float(fps) / ref)))\n\ndef _scale_signed(n_frames_at_30fps, fps, ref=30.0):\n    \"\"\"Signed version of _scale for forward/backward shifts.\"\"\"\n    if n_frames_at_30fps == 0:\n        return 0\n    s = 1 if n_frames_at_30fps > 0 else -1\n    mag = max(1, int(round(abs(n_frames_at_30fps) * float(fps) / ref)))\n    return s * mag\n\ndef _fps_from_meta(meta_df, fallback_lookup, default_fps=30.0):\n    if 'frames_per_second' in meta_df.columns and pd.notnull(meta_df['frames_per_second']).any():\n        return float(meta_df['frames_per_second'].iloc[0])\n    vid = meta_df['video_id'].iloc[0]\n    return float(fallback_lookup.get(vid, default_fps))\n\ndef _speed(cx: pd.Series, cy: pd.Series, fps: float) -> pd.Series:\n    return np.hypot(cx.diff(), cy.diff()).fillna(0.0) * float(fps)\n\ndef _roll_future_mean(s: pd.Series, w: int, min_p: int = 1) -> pd.Series:\n    return s.iloc[::-1].rolling(w, min_periods=min_p).mean().iloc[::-1]\n\ndef _roll_future_var(s: pd.Series, w: int, min_p: int = 2) -> pd.Series:\n    return s.iloc[::-1].rolling(w, min_periods=min_p).var().iloc[::-1]\n\n\n# ==================== NEW: WAVELET FEATURES ====================\n\ndef add_wavelet_features(X, cx, cy, fps):\n    \"\"\"Extract multi-scale frequency features using wavelet transform\"\"\"\n    if not HAS_PYWT:\n        return X\n    \n    try:\n        cx_vals = cx.values if hasattr(cx, 'values') else cx\n        cy_vals = cy.values if hasattr(cy, 'values') else cy\n        \n        # Speed signal\n        speed = np.sqrt(np.diff(cx_vals)**2 + np.diff(cy_vals)**2) * fps\n        speed = np.pad(speed, (1, 0), mode='edge')\n        \n        if len(speed) < 64:\n            return X\n        \n        # Discrete Wavelet Transform\n        max_level = min(4, int(np.log2(len(speed))) - 2)\n        if max_level < 1:\n            return X\n            \n        coeffs = pywt.wavedec(speed, 'db4', level=max_level)\n        \n        # Energy at each decomposition level\n        for i, c in enumerate(coeffs):\n            energy = np.mean(c**2)\n            X[f'wavelet_energy_L{i}'] = energy\n        \n        # Reconstruct low-frequency (approximation) and high-frequency (detail)\n        approx_coeffs = [coeffs[0]] + [np.zeros_like(c) for c in coeffs[1:]]\n        approx = pywt.waverec(approx_coeffs, 'db4')[:len(speed)]\n        \n        detail = speed - approx[:len(speed)]\n        \n        # Rolling stats on frequency components\n        w = _scale(30, fps)\n        X['wavelet_low_freq_mean'] = pd.Series(approx).rolling(w, min_periods=1, center=True).mean().values[:len(X)]\n        X['wavelet_high_freq_mean'] = pd.Series(np.abs(detail)).rolling(w, min_periods=1, center=True).mean().values[:len(X)]\n        X['wavelet_high_freq_std'] = pd.Series(detail).rolling(w, min_periods=1, center=True).std().values[:len(X)]\n        \n        # Ratio of high to low frequency energy (indicates behavior type)\n        X['wavelet_freq_ratio'] = X['wavelet_high_freq_mean'] / (X['wavelet_low_freq_mean'] + 1e-6)\n        \n    except Exception as e:\n        pass\n    \n    return X\n\n\n# ==================== NEW: PHYSICS-INFORMED FEATURES ====================\n\ndef add_physics_features(X, cx, cy, fps):\n    \"\"\"Physics-based behavior signatures: jerk, angular velocity, kinetic energy\"\"\"\n    \n    try:\n        cx_vals = cx.values if hasattr(cx, 'values') else cx\n        cy_vals = cy.values if hasattr(cy, 'values') else cy\n        \n        # Velocity (cm/s)\n        vel_x = np.diff(cx_vals, prepend=cx_vals[0]) * fps\n        vel_y = np.diff(cy_vals, prepend=cy_vals[0]) * fps\n        speed = np.sqrt(vel_x**2 + vel_y**2)\n        \n        # Acceleration (cm/s^2)\n        acc_x = np.diff(vel_x, prepend=vel_x[0]) * fps\n        acc_y = np.diff(vel_y, prepend=vel_y[0]) * fps\n        acc_mag = np.sqrt(acc_x**2 + acc_y**2)\n        \n        # Jerk (rate of acceleration change) - key for detecting behavior transitions\n        jerk_x = np.diff(acc_x, prepend=acc_x[0]) * fps\n        jerk_y = np.diff(acc_y, prepend=acc_y[0]) * fps\n        jerk_mag = np.sqrt(jerk_x**2 + jerk_y**2)\n        \n        w = _scale(30, fps)\n        X['jerk_mean'] = pd.Series(jerk_mag).rolling(w, min_periods=1, center=True).mean().values\n        X['jerk_max'] = pd.Series(jerk_mag).rolling(w, min_periods=1, center=True).max().values\n        X['jerk_std'] = pd.Series(jerk_mag).rolling(w, min_periods=1, center=True).std().values\n        \n        # Angular velocity (rotation detection)\n        angle = np.arctan2(vel_y, vel_x)\n        angular_vel = np.abs(np.diff(angle, prepend=angle[0]))\n        # Handle wraparound at ±π\n        angular_vel = np.minimum(angular_vel, 2 * np.pi - angular_vel) * fps\n        \n        X['angular_vel_mean'] = pd.Series(angular_vel).rolling(w, min_periods=1, center=True).mean().values\n        X['angular_vel_max'] = pd.Series(angular_vel).rolling(w, min_periods=1, center=True).max().values\n        X['angular_vel_std'] = pd.Series(angular_vel).rolling(w, min_periods=1, center=True).std().values\n        \n        # Kinetic energy proxy (mass normalized, so just v^2)\n        KE = speed**2\n        X['kinetic_energy_mean'] = pd.Series(KE).rolling(w, min_periods=1, center=True).mean().values\n        X['kinetic_energy_max'] = pd.Series(KE).rolling(w, min_periods=1, center=True).max().values\n        \n        # Centripetal acceleration (indicates turning/circling)\n        speed_safe = np.maximum(speed, 1e-6)\n        centripetal = speed_safe * angular_vel\n        X['centripetal_mean'] = pd.Series(centripetal).rolling(w, min_periods=1, center=True).mean().values\n        \n        # Tangential vs radial acceleration ratio\n        tangential_acc = np.abs(np.diff(speed, prepend=speed[0])) * fps\n        X['tangential_acc_mean'] = pd.Series(tangential_acc).rolling(w, min_periods=1, center=True).mean().values\n        X['acc_ratio'] = X['tangential_acc_mean'] / (X['centripetal_mean'] + 1e-6)\n        \n    except Exception as e:\n        pass\n    \n    return X\n\n\n# ==================== NEW: TRANSITION DETECTION FEATURES ====================\n\ndef add_transition_features(X, cx, cy, fps):\n    \"\"\"Features for detecting behavior onset/offset transitions\"\"\"\n    \n    try:\n        cx_vals = cx.values if hasattr(cx, 'values') else cx\n        cy_vals = cy.values if hasattr(cy, 'values') else cy\n        \n        speed = np.sqrt(np.diff(cx_vals)**2 + np.diff(cy_vals)**2) * fps\n        speed = np.pad(speed, (1, 0), mode='edge')\n        speed_series = pd.Series(speed)\n        \n        # Z-score of current speed vs local context (detect sudden changes)\n        for w_base in [30, 60, 120]:\n            w = _scale(w_base, fps)\n            local_mean = speed_series.rolling(w, min_periods=1, center=True).mean()\n            local_std = speed_series.rolling(w, min_periods=1, center=True).std() + 1e-6\n            X[f'speed_zscore_{w_base}'] = ((speed - local_mean) / local_std).values\n        \n        # Change point detection: difference between past and future windows\n        for w_base in [15, 30]:\n            w = _scale(w_base, fps)\n            past_mean = speed_series.rolling(w, min_periods=1).mean()\n            future_mean = speed_series.iloc[::-1].rolling(w, min_periods=1).mean().iloc[::-1]\n            X[f'speed_change_{w_base}'] = (future_mean - past_mean).values\n            \n            past_std = speed_series.rolling(w, min_periods=1).std()\n            future_std = speed_series.iloc[::-1].rolling(w, min_periods=1).std().iloc[::-1]\n            X[f'speed_var_change_{w_base}'] = (future_std - past_std).values\n        \n        # Autocorrelation at different lags (rhythmic behaviors like grooming)\n        for lag_base in [5, 10, 15]:\n            lag = _scale(lag_base, fps)\n            w = _scale(60, fps)\n            if len(speed) > lag + w:\n                def calc_autocorr(x):\n                    if len(x) <= lag:\n                        return 0\n                    try:\n                        return np.corrcoef(x[:-lag], x[lag:])[0, 1]\n                    except:\n                        return 0\n                \n                autocorr = speed_series.rolling(w, min_periods=w//2, center=True).apply(calc_autocorr, raw=True)\n                X[f'speed_autocorr_{lag_base}'] = autocorr.fillna(0).values\n        \n        # Speed derivative (acceleration sign changes indicate transitions)\n        speed_deriv = np.diff(speed, prepend=speed[0])\n        sign_changes = np.abs(np.diff(np.sign(speed_deriv), prepend=0))\n        w = _scale(30, fps)\n        X['acc_sign_changes'] = pd.Series(sign_changes).rolling(w, min_periods=1, center=True).sum().values\n        \n    except Exception as e:\n        pass\n    \n    return X\n\n\n# ==================== NEW: BEHAVIOR TEMPLATE MATCHING ====================\n\nclass BehaviorTemplateMatcher:\n    \"\"\"Match trajectory segments to known behavior templates using DTW\"\"\"\n    \n    def __init__(self, segment_length=60, n_templates=5):\n        self.segment_length = segment_length\n        self.n_templates = n_templates\n        self.templates = {}\n        self.is_fitted = False\n    \n    def normalize_segment(self, seg):\n        \"\"\"Normalize segment to fixed length and zero-center\"\"\"\n        if len(seg) == 0:\n            return np.zeros((self.segment_length, 2))\n        \n        # Resample to fixed length\n        if len(seg) != self.segment_length:\n            indices = np.linspace(0, len(seg) - 1, self.segment_length).astype(int)\n            seg = seg[indices]\n        \n        # Zero-center\n        seg = seg - seg.mean(axis=0)\n        \n        # Normalize scale\n        scale = np.std(seg) + 1e-6\n        seg = seg / scale\n        \n        return seg\n    \n    def fit(self, trajectories_by_action):\n        \"\"\"Build templates from training data\"\"\"\n        if not HAS_DTW:\n            return self\n        \n        for action, trajs in trajectories_by_action.items():\n            if len(trajs) == 0:\n                continue\n            \n            # Normalize all trajectories\n            normalized = [self.normalize_segment(t) for t in trajs]\n            \n            if len(normalized) <= self.n_templates:\n                self.templates[action] = normalized\n            else:\n                # Random sample for now\n                indices = np.random.choice(len(normalized), self.n_templates, replace=False)\n                self.templates[action] = [normalized[i] for i in indices]\n        \n        self.is_fitted = True\n        return self\n    \n    def compute_features(self, trajectory, fps):\n        \"\"\"Compute template matching features for a trajectory\"\"\"\n        if not self.is_fitted or not HAS_DTW:\n            return {}\n        \n        features = {}\n        segment_len = _scale(self.segment_length, fps)\n        \n        # Get middle segment\n        mid = len(trajectory) // 2\n        start = max(0, mid - segment_len // 2)\n        end = min(len(trajectory), start + segment_len)\n        \n        if end - start < segment_len // 2:\n            return features\n        \n        segment = trajectory[start:end]\n        segment_norm = self.normalize_segment(segment)\n        \n        for action, templates in self.templates.items():\n            dists = []\n            for template in templates:\n                try:\n                    dist, _ = fastdtw(segment_norm, template, dist=euclidean)\n                    dists.append(dist)\n                except:\n                    pass\n            \n            if dists:\n                features[f'template_{action}_min'] = min(dists)\n                features[f'template_{action}_mean'] = np.mean(dists)\n        \n        return features\n\n\n# Global template matcher\ntemplate_matcher = BehaviorTemplateMatcher()\n\n\n# ==================== NEW: CRF POST-PROCESSING ====================\n\nclass BehaviorCRF:\n    \"\"\"Use CRF to enforce temporal consistency of predictions\"\"\"\n    \n    def __init__(self):\n        self.crf = None\n        self.is_fitted = False\n        self.action_list = []\n    \n    def prepare_features(self, gbdt_probs, fps):\n        \"\"\"Convert per-frame predictions to CRF features\"\"\"\n        X_crf = []\n        \n        for i in range(len(gbdt_probs)):\n            frame_feats = {}\n            \n            # GBDT prediction confidences\n            for action, prob in gbdt_probs[i].items():\n                frame_feats[f'prob_{action}'] = str(round(prob, 2))\n                frame_feats[f'prob_{action}_high'] = str(prob > 0.5)\n                frame_feats[f'prob_{action}_very_high'] = str(prob > 0.8)\n            \n            # Context features\n            max_prob = max(gbdt_probs[i].values()) if gbdt_probs[i] else 0\n            frame_feats['max_prob'] = str(round(max_prob, 2))\n            \n            # Entropy\n            probs = list(gbdt_probs[i].values())\n            if probs:\n                entropy = -sum(p * np.log(p + 1e-10) for p in probs)\n                frame_feats['entropy'] = str(round(entropy, 2))\n            \n            # Temporal features\n            if i > 0:\n                for action in gbdt_probs[i]:\n                    prev_prob = gbdt_probs[i-1].get(action, 0)\n                    delta = gbdt_probs[i][action] - prev_prob\n                    frame_feats[f'delta_{action}'] = str(round(delta, 2))\n            \n            X_crf.append(frame_feats)\n        \n        return X_crf\n    \n    def fit(self, videos_probs, videos_labels, fps_list):\n        \"\"\"Train CRF on video sequences\"\"\"\n        if not HAS_CRF:\n            return self\n        \n        X_train = []\n        y_train = []\n        \n        for probs, labels, fps in zip(videos_probs, videos_labels, fps_list):\n            X_train.append(self.prepare_features(probs, fps))\n            y_train.append(labels)\n        \n        self.crf = sklearn_crfsuite.CRF(\n            algorithm='lbfgs',\n            c1=0.1,\n            c2=0.1,\n            max_iterations=100,\n            all_possible_transitions=True,\n        )\n        \n        try:\n            self.crf.fit(X_train, y_train)\n            self.is_fitted = True\n        except:\n            self.is_fitted = False\n        \n        return self\n    \n    def predict(self, gbdt_probs, fps):\n        \"\"\"Predict with temporal consistency\"\"\"\n        if not self.is_fitted or not HAS_CRF:\n            return None\n        \n        X = [self.prepare_features(gbdt_probs, fps)]\n        try:\n            return self.crf.predict(X)[0]\n        except:\n            return None\n\n\n# ==================== ORIGINAL FEATURE FUNCTIONS (ENHANCED) ====================\n\ndef add_curvature_features(X, center_x, center_y, fps):\n    \"\"\"Trajectory curvature (window lengths scaled by fps).\"\"\"\n    vel_x = center_x.diff()\n    vel_y = center_y.diff()\n    acc_x = vel_x.diff()\n    acc_y = vel_y.diff()\n\n    cross_prod = vel_x * acc_y - vel_y * acc_x\n    vel_mag = np.sqrt(vel_x**2 + vel_y**2)\n    curvature = np.abs(cross_prod) / (vel_mag**3 + 1e-6)\n\n    for w in [30, 60]:\n        ws = _scale(w, fps)\n        X[f'curv_mean_{w}'] = curvature.rolling(ws, min_periods=max(1, ws // 6)).mean()\n\n    angle = np.arctan2(vel_y, vel_x)\n    angle_change = np.abs(angle.diff())\n    w = 30\n    ws = _scale(w, fps)\n    X[f'turn_rate_{w}'] = angle_change.rolling(ws, min_periods=max(1, ws // 6)).sum()\n\n    return X\n\ndef add_multiscale_features(X, center_x, center_y, fps):\n    \"\"\"Multi-scale temporal features (speed in cm/s; windows scaled by fps).\"\"\"\n    speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps)\n\n    scales = [10, 40, 160]\n    for scale in scales:\n        ws = _scale(scale, fps)\n        if len(speed) >= ws:\n            X[f'sp_m{scale}'] = speed.rolling(ws, min_periods=max(1, ws // 4)).mean()\n            X[f'sp_s{scale}'] = speed.rolling(ws, min_periods=max(1, ws // 4)).std()\n\n    if len(scales) >= 2 and f'sp_m{scales[0]}' in X.columns and f'sp_m{scales[-1]}' in X.columns:\n        X['sp_ratio'] = X[f'sp_m{scales[0]}'] / (X[f'sp_m{scales[-1]}'] + 1e-6)\n\n    return X\n\ndef add_state_features(X, center_x, center_y, fps):\n    \"\"\"Behavioral state transitions; bins adjusted so semantics are fps-invariant.\"\"\"\n    speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps)\n    w_ma = _scale(15, fps)\n    speed_ma = speed.rolling(w_ma, min_periods=max(1, w_ma // 3)).mean()\n\n    try:\n        bins = [-np.inf, 0.5 * fps, 2.0 * fps, 5.0 * fps, np.inf]\n        speed_states = pd.cut(speed_ma, bins=bins, labels=[0, 1, 2, 3]).astype(float)\n\n        for window in [60, 120]:\n            ws = _scale(window, fps)\n            if len(speed_states) >= ws:\n                for state in [0, 1, 2, 3]:\n                    X[f's{state}_{window}'] = (\n                        (speed_states == state).astype(float)\n                        .rolling(ws, min_periods=max(1, ws // 6)).mean()\n                    )\n                state_changes = (speed_states != speed_states.shift(1)).astype(float)\n                X[f'trans_{window}'] = state_changes.rolling(ws, min_periods=max(1, ws // 6)).sum()\n    except Exception:\n        pass\n\n    return X\n\ndef add_longrange_features(X, center_x, center_y, fps):\n    \"\"\"Long-range temporal features (windows & spans scaled by fps).\"\"\"\n    for window in [120, 240]:\n        ws = _scale(window, fps)\n        if len(center_x) >= ws:\n            X[f'x_ml{window}'] = center_x.rolling(ws, min_periods=max(5, ws // 6)).mean()\n            X[f'y_ml{window}'] = center_y.rolling(ws, min_periods=max(5, ws // 6)).mean()\n\n    for span in [60, 120]:\n        s = _scale(span, fps)\n        X[f'x_e{span}'] = center_x.ewm(span=s, min_periods=1).mean()\n        X[f'y_e{span}'] = center_y.ewm(span=s, min_periods=1).mean()\n\n    speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps)\n    for window in [60, 120]:\n        ws = _scale(window, fps)\n        if len(speed) >= ws:\n            X[f'sp_pct{window}'] = speed.rolling(ws, min_periods=max(5, ws // 6)).rank(pct=True)\n\n    return X\n\ndef add_cumulative_distance_single(X, cx, cy, fps, horizon_frames_base: int = 180, colname: str = \"path_cum180\"):\n    L = max(1, _scale(horizon_frames_base, fps))\n    step = np.hypot(cx.diff(), cy.diff())\n    path = step.rolling(2*L + 1, min_periods=max(5, L//6), center=True).sum()\n    X[colname] = path.fillna(0.0).astype(np.float32)\n    return X\n\ndef add_groom_microfeatures(X, df, fps):\n    parts = df.columns.get_level_values(0)\n    if 'body_center' not in parts or 'nose' not in parts:\n        return X\n\n    cx = df['body_center']['x']\n    cy = df['body_center']['y']\n    nx = df['nose']['x']\n    ny = df['nose']['y']\n\n    cs = (np.sqrt(cx.diff()**2 + cy.diff()**2) * float(fps)).fillna(0)\n    ns = (np.sqrt(nx.diff()**2 + ny.diff()**2) * float(fps)).fillna(0)\n\n    w30 = _scale(30, fps)\n    X['head_body_decouple'] = (ns / (cs + 1e-3)).clip(0, 10).rolling(w30, min_periods=max(1, w30//3)).median()\n\n    r = np.sqrt((nx - cx)**2 + (ny - cy)**2)\n    X['nose_rad_std'] = r.rolling(w30, min_periods=max(1, w30//3)).std().fillna(0)\n\n    if 'tail_base' in parts:\n        ang = np.arctan2(df['nose']['y']-df['tail_base']['y'], df['nose']['x']-df['tail_base']['x'])\n        dang = np.abs(ang.diff()).fillna(0)\n        X['head_orient_jitter'] = dang.rolling(w30, min_periods=max(1, w30//3)).mean()\n\n    return X\n\ndef add_interaction_features(X, mouse_pair, avail_A, avail_B, fps):\n    \"\"\"Social interaction features (windows scaled by fps).\"\"\"\n    if 'body_center' not in avail_A or 'body_center' not in avail_B:\n        return X\n\n    rel_x = mouse_pair['A']['body_center']['x'] - mouse_pair['B']['body_center']['x']\n    rel_y = mouse_pair['A']['body_center']['y'] - mouse_pair['B']['body_center']['y']\n    rel_dist = np.sqrt(rel_x**2 + rel_y**2)\n\n    A_vx = mouse_pair['A']['body_center']['x'].diff()\n    A_vy = mouse_pair['A']['body_center']['y'].diff()\n    B_vx = mouse_pair['B']['body_center']['x'].diff()\n    B_vy = mouse_pair['B']['body_center']['y'].diff()\n\n    A_lead = (A_vx * rel_x + A_vy * rel_y) / (np.sqrt(A_vx**2 + A_vy**2) * rel_dist + 1e-6)\n    B_lead = (B_vx * (-rel_x) + B_vy * (-rel_y)) / (np.sqrt(B_vx**2 + B_vy**2) * rel_dist + 1e-6)\n\n    for window in [30, 60]:\n        ws = _scale(window, fps)\n        X[f'A_ld{window}'] = A_lead.rolling(ws, min_periods=max(1, ws // 6)).mean()\n        X[f'B_ld{window}'] = B_lead.rolling(ws, min_periods=max(1, ws // 6)).mean()\n\n    approach = -rel_dist.diff()\n    chase = approach * B_lead\n    w = 30\n    ws = _scale(w, fps)\n    X[f'chase_{w}'] = chase.rolling(ws, min_periods=max(1, ws // 6)).mean()\n\n    for window in [60, 120]:\n        ws = _scale(window, fps)\n        A_sp = np.sqrt(A_vx**2 + A_vy**2)\n        B_sp = np.sqrt(B_vx**2 + B_vy**2)\n        X[f'sp_cor{window}'] = A_sp.rolling(ws, min_periods=max(1, ws // 6)).corr(B_sp)\n\n    return X\n\ndef add_speed_asymmetry_future_past_single(X, cx, cy, fps, horizon_base=30, agg=\"mean\"):\n    w = max(3, _scale(horizon_base, fps))\n    v = _speed(cx, cy, fps)\n    if agg == \"median\":\n        v_past = v.rolling(w, min_periods=max(3, w//4), center=False).median()\n        v_fut = v.iloc[::-1].rolling(w, min_periods=max(3, w//4)).median().iloc[::-1]\n    else:\n        v_past = v.rolling(w, min_periods=max(3, w//4), center=False).mean()\n        v_fut = _roll_future_mean(v, w, min_p=max(3, w//4))\n    X[\"spd_asym_1s\"] = (v_fut - v_past).fillna(0.0)\n    return X\n\ndef add_gauss_shift_speed_future_past_single(X, cx, cy, fps, window_base=30, eps=1e-6):\n    w = max(5, _scale(window_base, fps))\n    v = _speed(cx, cy, fps)\n\n    mu_p = v.rolling(w, min_periods=max(3, w//4)).mean()\n    va_p = v.rolling(w, min_periods=max(3, w//4)).var().clip(lower=eps)\n\n    mu_f = _roll_future_mean(v, w, min_p=max(3, w//4))\n    va_f = _roll_future_var(v, w, min_p=max(3, w//4)).clip(lower=eps)\n\n    kl_pf = 0.5 * ((va_p/va_f) + ((mu_f - mu_p)**2)/va_f - 1.0 + np.log(va_f/va_p))\n    kl_fp = 0.5 * ((va_f/va_p) + ((mu_p - mu_f)**2)/va_p - 1.0 + np.log(va_p/va_f))\n    X[\"spd_symkl_1s\"] = (kl_pf + kl_fp).replace([np.inf, -np.inf], np.nan).fillna(0.0)\n    return X\n\n\n# ==================== ENHANCED TRANSFORM FUNCTIONS ====================\n\ndef transform_single(single_mouse, body_parts_tracked, fps):\n    \"\"\"Enhanced single mouse transform with NEW features.\"\"\"\n    available_body_parts = single_mouse.columns.get_level_values(0)\n\n    # Base distance features\n    X = pd.DataFrame({\n        f\"{p1}+{p2}\": np.square(single_mouse[p1] - single_mouse[p2]).sum(axis=1, skipna=False)\n        for p1, p2 in itertools.combinations(body_parts_tracked, 2)\n        if p1 in available_body_parts and p2 in available_body_parts\n    })\n    X = X.reindex(columns=[f\"{p1}+{p2}\" for p1, p2 in itertools.combinations(body_parts_tracked, 2)], copy=False)\n\n    # Speed-like features\n    if all(p in single_mouse.columns for p in ['ear_left', 'ear_right', 'tail_base']):\n        lag = _scale(10, fps)\n        shifted = single_mouse[['ear_left', 'ear_right', 'tail_base']].shift(lag)\n        speeds = pd.DataFrame({\n            'sp_lf': np.square(single_mouse['ear_left'] - shifted['ear_left']).sum(axis=1, skipna=False),\n            'sp_rt': np.square(single_mouse['ear_right'] - shifted['ear_right']).sum(axis=1, skipna=False),\n            'sp_lf2': np.square(single_mouse['ear_left'] - shifted['tail_base']).sum(axis=1, skipna=False),\n            'sp_rt2': np.square(single_mouse['ear_right'] - shifted['tail_base']).sum(axis=1, skipna=False),\n        })\n        X = pd.concat([X, speeds], axis=1)\n\n    if 'nose+tail_base' in X.columns and 'ear_left+ear_right' in X.columns:\n        X['elong'] = X['nose+tail_base'] / (X['ear_left+ear_right'] + 1e-6)\n\n    # Body angle\n    if all(p in available_body_parts for p in ['nose', 'body_center', 'tail_base']):\n        v1 = single_mouse['nose'] - single_mouse['body_center']\n        v2 = single_mouse['tail_base'] - single_mouse['body_center']\n        X['body_ang'] = (v1['x'] * v2['x'] + v1['y'] * v2['y']) / (\n            np.sqrt(v1['x']**2 + v1['y']**2) * np.sqrt(v2['x']**2 + v2['y']**2) + 1e-6)\n\n    # Core temporal features\n    if 'body_center' in available_body_parts:\n        cx = single_mouse['body_center']['x']\n        cy = single_mouse['body_center']['y']\n\n        for w in [5, 15, 30, 60]:\n            ws = _scale(w, fps)\n            roll = dict(min_periods=1, center=True)\n            X[f'cx_m{w}'] = cx.rolling(ws, **roll).mean()\n            X[f'cy_m{w}'] = cy.rolling(ws, **roll).mean()\n            X[f'cx_s{w}'] = cx.rolling(ws, **roll).std()\n            X[f'cy_s{w}'] = cy.rolling(ws, **roll).std()\n            X[f'x_rng{w}'] = cx.rolling(ws, **roll).max() - cx.rolling(ws, **roll).min()\n            X[f'y_rng{w}'] = cy.rolling(ws, **roll).max() - cy.rolling(ws, **roll).min()\n            X[f'disp{w}'] = np.sqrt(cx.diff().rolling(ws, min_periods=1).sum()**2 +\n                                     cy.diff().rolling(ws, min_periods=1).sum()**2)\n            X[f'act{w}'] = np.sqrt(cx.diff().rolling(ws, min_periods=1).var() +\n                                    cy.diff().rolling(ws, min_periods=1).var())\n\n        # Original advanced features\n        X = add_curvature_features(X, cx, cy, fps)\n        X = add_multiscale_features(X, cx, cy, fps)\n        X = add_state_features(X, cx, cy, fps)\n        X = add_longrange_features(X, cx, cy, fps)\n        X = add_cumulative_distance_single(X, cx, cy, fps, horizon_frames_base=180)\n        X = add_groom_microfeatures(X, single_mouse, fps)\n        X = add_speed_asymmetry_future_past_single(X, cx, cy, fps, horizon_base=30)\n        X = add_gauss_shift_speed_future_past_single(X, cx, cy, fps, window_base=30)\n        \n        # ===== NEW FEATURES =====\n        X = add_wavelet_features(X, cx, cy, fps)\n        X = add_physics_features(X, cx, cy, fps)\n        X = add_transition_features(X, cx, cy, fps)\n\n    # Nose-tail features\n    if all(p in available_body_parts for p in ['nose', 'tail_base']):\n        nt_dist = np.sqrt((single_mouse['nose']['x'] - single_mouse['tail_base']['x'])**2 +\n                          (single_mouse['nose']['y'] - single_mouse['tail_base']['y'])**2)\n        for lag in [10, 20, 40]:\n            l = _scale(lag, fps)\n            X[f'nt_lg{lag}'] = nt_dist.shift(l)\n            X[f'nt_df{lag}'] = nt_dist - nt_dist.shift(l)\n\n    # Ear features\n    if all(p in available_body_parts for p in ['ear_left', 'ear_right']):\n        ear_d = np.sqrt((single_mouse['ear_left']['x'] - single_mouse['ear_right']['x'])**2 +\n                        (single_mouse['ear_left']['y'] - single_mouse['ear_right']['y'])**2)\n        for off in [-20, -10, 10, 20]:\n            o = _scale_signed(off, fps)\n            X[f'ear_o{off}'] = ear_d.shift(-o)\n        w = _scale(30, fps)\n        X['ear_con'] = ear_d.rolling(w, min_periods=1, center=True).std() / \\\n                       (ear_d.rolling(w, min_periods=1, center=True).mean() + 1e-6)\n\n    return X.astype(np.float32, copy=False)\n\n\ndef transform_pair(mouse_pair, body_parts_tracked, fps):\n    \"\"\"Enhanced pair transform with NEW features.\"\"\"\n    avail_A = mouse_pair['A'].columns.get_level_values(0)\n    avail_B = mouse_pair['B'].columns.get_level_values(0)\n\n    # Inter-mouse distances\n    X = pd.DataFrame({\n        f\"12+{p1}+{p2}\": np.square(mouse_pair['A'][p1] - mouse_pair['B'][p2]).sum(axis=1, skipna=False)\n        for p1, p2 in itertools.product(body_parts_tracked, repeat=2)\n        if p1 in avail_A and p2 in avail_B\n    })\n    X = X.reindex(columns=[f\"12+{p1}+{p2}\" for p1, p2 in itertools.product(body_parts_tracked, repeat=2)], copy=False)\n\n    # Speed-like features\n    if ('A', 'ear_left') in mouse_pair.columns and ('B', 'ear_left') in mouse_pair.columns:\n        lag = _scale(10, fps)\n        shA = mouse_pair['A']['ear_left'].shift(lag)\n        shB = mouse_pair['B']['ear_left'].shift(lag)\n        speeds = pd.DataFrame({\n            'sp_A': np.square(mouse_pair['A']['ear_left'] - shA).sum(axis=1, skipna=False),\n            'sp_AB': np.square(mouse_pair['A']['ear_left'] - shB).sum(axis=1, skipna=False),\n            'sp_B': np.square(mouse_pair['B']['ear_left'] - shB).sum(axis=1, skipna=False),\n        })\n        X = pd.concat([X, speeds], axis=1)\n\n    # Relative orientation\n    if all(p in avail_A for p in ['nose', 'tail_base']) and all(p in avail_B for p in ['nose', 'tail_base']):\n        dir_A = mouse_pair['A']['nose'] - mouse_pair['A']['tail_base']\n        dir_B = mouse_pair['B']['nose'] - mouse_pair['B']['tail_base']\n        X['rel_ori'] = (dir_A['x'] * dir_B['x'] + dir_A['y'] * dir_B['y']) / (\n            np.sqrt(dir_A['x']**2 + dir_A['y']**2) * np.sqrt(dir_B['x']**2 + dir_B['y']**2) + 1e-6)\n\n    # Approach rate\n    if all(p in avail_A for p in ['nose']) and all(p in avail_B for p in ['nose']):\n        cur = np.square(mouse_pair['A']['nose'] - mouse_pair['B']['nose']).sum(axis=1, skipna=False)\n        lag = _scale(10, fps)\n        shA_n = mouse_pair['A']['nose'].shift(lag)\n        shB_n = mouse_pair['B']['nose'].shift(lag)\n        past = np.square(shA_n - shB_n).sum(axis=1, skipna=False)\n        X['appr'] = cur - past\n\n    # Distance bins\n    if 'body_center' in avail_A and 'body_center' in avail_B:\n        cd = np.sqrt((mouse_pair['A']['body_center']['x'] - mouse_pair['B']['body_center']['x'])**2 +\n                     (mouse_pair['A']['body_center']['y'] - mouse_pair['B']['body_center']['y'])**2)\n        X['v_cls'] = (cd < 5.0).astype(float)\n        X['cls'] = ((cd >= 5.0) & (cd < 15.0)).astype(float)\n        X['med'] = ((cd >= 15.0) & (cd < 30.0)).astype(float)\n        X['far'] = (cd >= 30.0).astype(float)\n\n    # Temporal interaction features\n    if 'body_center' in avail_A and 'body_center' in avail_B:\n        cd_full = np.square(mouse_pair['A']['body_center'] - mouse_pair['B']['body_center']).sum(axis=1, skipna=False)\n\n        for w in [5, 15, 30, 60]:\n            ws = _scale(w, fps)\n            roll = dict(min_periods=1, center=True)\n            X[f'd_m{w}'] = cd_full.rolling(ws, **roll).mean()\n            X[f'd_s{w}'] = cd_full.rolling(ws, **roll).std()\n            X[f'd_mn{w}'] = cd_full.rolling(ws, **roll).min()\n            X[f'd_mx{w}'] = cd_full.rolling(ws, **roll).max()\n\n            d_var = cd_full.rolling(ws, **roll).var()\n            X[f'int{w}'] = 1 / (1 + d_var)\n\n            Axd = mouse_pair['A']['body_center']['x'].diff()\n            Ayd = mouse_pair['A']['body_center']['y'].diff()\n            Bxd = mouse_pair['B']['body_center']['x'].diff()\n            Byd = mouse_pair['B']['body_center']['y'].diff()\n            coord = Axd * Bxd + Ayd * Byd\n            X[f'co_m{w}'] = coord.rolling(ws, **roll).mean()\n            X[f'co_s{w}'] = coord.rolling(ws, **roll).std()\n\n        # ===== NEW: Physics features for both mice =====\n        cx_A = mouse_pair['A']['body_center']['x']\n        cy_A = mouse_pair['A']['body_center']['y']\n        cx_B = mouse_pair['B']['body_center']['x']\n        cy_B = mouse_pair['B']['body_center']['y']\n        \n        # Add physics features for mouse A\n        X_A = pd.DataFrame(index=X.index)\n        X_A = add_physics_features(X_A, cx_A, cy_A, fps)\n        for col in X_A.columns:\n            X[f'A_{col}'] = X_A[col]\n        \n        # Add physics features for mouse B\n        X_B = pd.DataFrame(index=X.index)\n        X_B = add_physics_features(X_B, cx_B, cy_B, fps)\n        for col in X_B.columns:\n            X[f'B_{col}'] = X_B[col]\n        \n        # ===== NEW: Relative physics features =====\n        # Relative velocity\n        rel_vx = np.diff(cx_A.values - cx_B.values, prepend=0) * fps\n        rel_vy = np.diff(cy_A.values - cy_B.values, prepend=0) * fps\n        rel_speed = np.sqrt(rel_vx**2 + rel_vy**2)\n        \n        w = _scale(30, fps)\n        X['rel_speed_mean'] = pd.Series(rel_speed).rolling(w, min_periods=1, center=True).mean().values\n        X['rel_speed_std'] = pd.Series(rel_speed).rolling(w, min_periods=1, center=True).std().values\n        \n        # Relative acceleration\n        rel_ax = np.diff(rel_vx, prepend=0) * fps\n        rel_ay = np.diff(rel_vy, prepend=0) * fps\n        rel_acc = np.sqrt(rel_ax**2 + rel_ay**2)\n        X['rel_acc_mean'] = pd.Series(rel_acc).rolling(w, min_periods=1, center=True).mean().values\n        \n        # ===== NEW: Interaction dynamics =====\n        # Time to collision (if approaching)\n        dist = np.sqrt((cx_A.values - cx_B.values)**2 + (cy_A.values - cy_B.values)**2)\n        closing_speed = -np.diff(dist, prepend=dist[0]) * fps\n        ttc = dist / (closing_speed + 1e-6)\n        ttc = np.clip(ttc, -100, 100)\n        X['time_to_collision'] = pd.Series(ttc).rolling(w, min_periods=1, center=True).mean().values\n        \n        # Bearing angle (direction from A to B relative to A's heading)\n        if 'nose' in avail_A and 'tail_base' in avail_A:\n            heading_x = mouse_pair['A']['nose']['x'] - mouse_pair['A']['tail_base']['x']\n            heading_y = mouse_pair['A']['nose']['y'] - mouse_pair['A']['tail_base']['y']\n            to_B_x = cx_B - cx_A\n            to_B_y = cy_B - cy_A\n            \n            heading_angle = np.arctan2(heading_y, heading_x)\n            to_B_angle = np.arctan2(to_B_y, to_B_x)\n            bearing = to_B_angle - heading_angle\n            # Normalize to [-pi, pi]\n            bearing = np.arctan2(np.sin(bearing), np.cos(bearing))\n            \n            X['bearing_to_B'] = bearing.values\n            X['bearing_to_B_abs'] = np.abs(bearing.values)\n            X['facing_B'] = (np.abs(bearing.values) < np.pi/4).astype(float)\n\n    # Nose-nose dynamics\n    if 'nose' in avail_A and 'nose' in avail_B:\n        nn = np.sqrt((mouse_pair['A']['nose']['x'] - mouse_pair['B']['nose']['x'])**2 +\n                     (mouse_pair['A']['nose']['y'] - mouse_pair['B']['nose']['y'])**2)\n        for lag in [10, 20, 40]:\n            l = _scale(lag, fps)\n            X[f'nn_lg{lag}'] = nn.shift(l)\n            X[f'nn_ch{lag}'] = nn - nn.shift(l)\n            is_cl = (nn < 10.0).astype(float)\n            X[f'cl_ps{lag}'] = is_cl.rolling(l, min_periods=1).mean()\n\n    # Velocity alignment\n    if 'body_center' in avail_A and 'body_center' in avail_B:\n        Avx = mouse_pair['A']['body_center']['x'].diff()\n        Avy = mouse_pair['A']['body_center']['y'].diff()\n        Bvx = mouse_pair['B']['body_center']['x'].diff()\n        Bvy = mouse_pair['B']['body_center']['y'].diff()\n        val = (Avx * Bvx + Avy * Bvy) / (np.sqrt(Avx**2 + Avy**2) * np.sqrt(Bvx**2 + Bvy**2) + 1e-6)\n\n        for off in [-20, -10, 0, 10, 20]:\n            o = _scale_signed(off, fps)\n            X[f'va_{off}'] = val.shift(-o)\n\n        w = _scale(30, fps)\n        X['int_con'] = cd_full.rolling(w, min_periods=1, center=True).std() / \\\n                       (cd_full.rolling(w, min_periods=1, center=True).mean() + 1e-6)\n\n        # Advanced interaction\n        X = add_interaction_features(X, mouse_pair, avail_A, avail_B, fps)\n\n    return X.astype(np.float32, copy=False)\n\n\n# ==================== ADAPTIVE THRESHOLDING ====================\n\ndef predict_multiclass_adaptive(pred, meta, action_thresholds=None):\n    \"\"\"Adaptive thresholding per action + temporal smoothing\"\"\"\n    if action_thresholds is None:\n        action_thresholds = defaultdict(lambda: 0.27)\n    \n    # Apply temporal smoothing\n    pred_smoothed = pred.rolling(window=5, min_periods=1, center=True).mean()\n    \n    ama = np.argmax(pred_smoothed.values, axis=1)\n    \n    max_probs = pred_smoothed.max(axis=1)\n    threshold_mask = np.zeros(len(pred_smoothed), dtype=bool)\n    for i, action in enumerate(pred_smoothed.columns):\n        action_mask = (ama == i)\n        threshold = action_thresholds.get(action, 0.27)\n        threshold_mask |= (action_mask & (max_probs >= threshold))\n    \n    ama = np.where(threshold_mask, ama, -1)\n    ama = pd.Series(ama, index=meta.video_frame)\n    \n    changes_mask = (ama != ama.shift(1)).values\n    ama_changes = ama[changes_mask]\n    meta_changes = meta[changes_mask]\n    mask = ama_changes.values >= 0\n    if len(mask) > 0:\n        mask[-1] = False\n    \n    if mask.sum() == 0:\n        return pd.DataFrame(columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame'])\n    \n    submission_part = pd.DataFrame({\n        'video_id': meta_changes['video_id'][mask].values,\n        'agent_id': meta_changes['agent_id'][mask].values,\n        'target_id': meta_changes['target_id'][mask].values,\n        'action': pred.columns[ama_changes[mask].values],\n        'start_frame': ama_changes.index[mask],\n        'stop_frame': ama_changes.index[1:][mask[:-1]]\n    })\n    \n    stop_video_id = meta_changes['video_id'][1:][mask[:-1]].values\n    stop_agent_id = meta_changes['agent_id'][1:][mask[:-1]].values\n    stop_target_id = meta_changes['target_id'][1:][mask[:-1]].values\n    \n    for i in range(len(submission_part)):\n        video_id = submission_part.video_id.iloc[i]\n        agent_id = submission_part.agent_id.iloc[i]\n        target_id = submission_part.target_id.iloc[i]\n        if i < len(stop_video_id):\n            if stop_video_id[i] != video_id or stop_agent_id[i] != agent_id or stop_target_id[i] != target_id:\n                new_stop_frame = meta.query(\"(video_id == @video_id)\").video_frame.max() + 1\n                submission_part.iat[i, submission_part.columns.get_loc('stop_frame')] = new_stop_frame\n        else:\n            new_stop_frame = meta.query(\"(video_id == @video_id)\").video_frame.max() + 1\n            submission_part.iat[i, submission_part.columns.get_loc('stop_frame')] = new_stop_frame\n    \n    # Filter out very short events (likely noise)\n    duration = submission_part.stop_frame - submission_part.start_frame\n    submission_part = submission_part[duration >= 3].reset_index(drop=True)\n    \n    if len(submission_part) > 0:\n        assert (submission_part.stop_frame > submission_part.start_frame).all(), 'stop <= start'\n    \n    if verbose:\n        print(f'  actions found: {len(submission_part)}')\n    return submission_part\n\n\n# ==================== HELPER FOR LGBM ====================\n\ndef _find_lgbm_step(pipe):\n    try:\n        if \"stratifiedsubsetclassifier__estimator\" in pipe.get_params():\n            est = pipe.get_params()[\"stratifiedsubsetclassifier__estimator\"]\n            if isinstance(est, lightgbm.LGBMClassifier):\n                return \"stratifiedsubsetclassifier\"\n        if \"stratifiedsubsetclassifierweval__estimator\" in pipe.get_params():\n            est = pipe.get_params()[\"stratifiedsubsetclassifierweval__estimator\"]\n            if isinstance(est, lightgbm.LGBMClassifier):\n                return \"stratifiedsubsetclassifierweval\"\n    except Exception as e:\n        print(e)\n    return None\n\n\n# ==================== ENSEMBLE SUBMISSION ====================\n\ndef submit_ensemble(body_parts_tracked_str, switch_tr, X_tr, label, meta, n_samples=1_500_000):\n    models = []\n    models.append(make_pipeline(\n        StratifiedSubsetClassifier(_make_lgbm(\n            n_estimators=225, learning_rate=0.07, min_child_samples=40,\n            num_leaves=31, subsample=0.8, colsample_bytree=0.8, verbose=-1, gpu_use_dp=USE_GPU\n        ), n_samples)\n    ))\n    models.append(make_pipeline(\n        StratifiedSubsetClassifier(_make_lgbm(\n            n_estimators=150, learning_rate=0.1, min_child_samples=20,\n            num_leaves=63, max_depth=8, subsample=0.7, colsample_bytree=0.9,\n            reg_alpha=0.1, reg_lambda=0.1, verbose=-1, gpu_use_dp=USE_GPU\n        ), (n_samples and int(n_samples/1.25)))\n    ))\n    models.append(make_pipeline(\n        StratifiedSubsetClassifier(_make_lgbm(\n            n_estimators=100, learning_rate=0.05, min_child_samples=30,\n            num_leaves=127, max_depth=10, subsample=0.75, verbose=-1, gpu_use_dp=USE_GPU,\n        ), (n_samples and int(n_samples/1.66)))\n    ))\n\n    xgb0 = _make_xgb(\n        n_estimators=180, learning_rate=0.08, max_depth=6,\n        min_child_weight=8 if USE_GPU else 5, gamma=1.0 if USE_GPU else 0.,\n        subsample=0.8, colsample_bytree=0.8, single_precision_histogram=USE_GPU,\n        verbosity=0\n    )\n    models.append(make_pipeline(StratifiedSubsetClassifier(xgb0, n_samples and int(n_samples/1.2))))\n\n    cb_est = _make_cb(iterations=120, learning_rate=0.1, depth=6,\n                      verbose=False, allow_writing_files=False)\n    models.append(make_pipeline(StratifiedSubsetClassifier(cb_est, n_samples)))\n\n    model_names = ['lgbm_225', 'lgbm_150', 'lgbm_100', 'xgb_180', 'cat_120']\n\n    if USE_GPU:\n        xgb1 = XGBClassifier(\n            random_state=SEED, booster=\"gbtree\", tree_method=\"gpu_hist\",\n            n_estimators=2000, learning_rate=0.05, grow_policy=\"lossguide\",\n            max_leaves=255, max_depth=0, min_child_weight=10, gamma=0.0,\n            subsample=0.90, colsample_bytree=1.00, colsample_bylevel=0.85,\n            reg_alpha=0.0, reg_lambda=1.0, max_bin=256,\n            single_precision_histogram=True, verbosity=0\n        )\n        models.append(make_pipeline(\n            StratifiedSubsetClassifierWEval(xgb1, n_samples and int(n_samples/2.),\n                                            random_state=SEED, valid_size=0.10, val_cap_ratio=0.25,\n                                            es_rounds=\"auto\", es_metric=\"auto\")\n        ))\n        xgb2 = XGBClassifier(\n            random_state=SEED, booster=\"gbtree\", tree_method=\"gpu_hist\",\n            n_estimators=1400, learning_rate=0.06, max_depth=7,\n            min_child_weight=12, subsample=0.70, colsample_bytree=0.80,\n            reg_alpha=0.0, reg_lambda=1.5, max_bin=256,\n            single_precision_histogram=True, verbosity=0\n        )\n        models.append(make_pipeline(\n            StratifiedSubsetClassifierWEval(xgb2, n_samples and int(n_samples/1.5),\n                                            random_state=SEED, valid_size=0.10, val_cap_ratio=0.25,\n                                            es_rounds=\"auto\", es_metric=\"auto\")\n        ))\n\n        cb1 = CatBoostClassifier(\n            random_seed=SEED, task_type=\"GPU\", devices=\"0\",\n            iterations=4000, learning_rate=0.03, depth=8, l2_leaf_reg=6.0,\n            bootstrap_type=\"Bayesian\", bagging_temperature=0.5,\n            random_strength=0.5, loss_function=\"Logloss\",\n            eval_metric=\"PRAUC:type=Classic\", auto_class_weights=\"Balanced\",\n            border_count=64, verbose=False, allow_writing_files=False\n        )\n        models.append(make_pipeline(\n            StratifiedSubsetClassifierWEval(cb1, n_samples and int(n_samples/2.0),\n                                            random_state=SEED, valid_size=0.10, val_cap_ratio=0.25,\n                                            es_rounds=\"auto\", es_metric=\"auto\")\n        ))\n        model_names.extend(['xgb1', 'xgb2', 'cat_bay'])\n\n    model_list = []\n    for action in label.columns:\n        action_mask = ~label[action].isna().values\n        y_action = label[action][action_mask].values.astype(int)\n        meta_masked = meta.iloc[action_mask]\n\n        trained = []\n        for model_idx, m in enumerate(models):\n            m_clone = clone(m)\n            try:\n                t0 = perf_counter()\n                m_clone.fit(X_tr[action_mask], y_action)\n                dt = perf_counter() - t0\n                print(f\"trained model {model_names[model_idx]} | {switch_tr} | action={action} | {dt:.1f}s\", flush=True)\n            except Exception:\n                step = _find_lgbm_step(m_clone)\n                if step is None:\n                    continue\n                try:\n                    m_clone.set_params(**{f\"{step}__estimator__device\": \"cpu\"})\n                    t0 = perf_counter()\n                    m_clone.fit(X_tr[action_mask], y_action)\n                    dt = perf_counter() - t0\n                    print(f\"trained (CPU fallback) {model_names[model_idx]} | {switch_tr} | action={action} | {dt:.1f}s\", flush=True)\n                except Exception as e2:\n                    print(e2)\n                    continue\n            trained.append(m_clone)\n\n        if trained:\n            model_list.append((action, trained))\n\n    del X_tr\n    gc.collect()\n\n    # ---- TEST INFERENCE ----\n    body_parts_tracked = json.loads(body_parts_tracked_str)\n    if len(body_parts_tracked) > 5:\n        body_parts_tracked = [b for b in body_parts_tracked if b not in drop_body_parts]\n\n    test_subset = test[test.body_parts_tracked == body_parts_tracked_str]\n    generator = generate_mouse_data(\n        test_subset, 'test',\n        generate_single=(switch_tr == 'single'),\n        generate_pair=(switch_tr == 'pair')\n    )\n    fps_lookup = (test_subset[['video_id', 'frames_per_second']]\n                  .drop_duplicates('video_id')\n                  .set_index('video_id')['frames_per_second'].to_dict())\n\n    for switch_te, data_te, meta_te, actions_te in generator:\n        assert switch_te == switch_tr\n        try:\n            fps_i = _fps_from_meta(meta_te, fps_lookup, default_fps=30.0)\n            if switch_te == 'single':\n                X_te = transform_single(data_te, body_parts_tracked, fps_i)\n            else:\n                X_te = transform_pair(data_te, body_parts_tracked, fps_i)\n\n            del data_te\n\n            pred = pd.DataFrame(index=meta_te.video_frame)\n            for action, trained in model_list:\n                if action in actions_te:\n                    probs = []\n                    for mi, mdl in enumerate(trained):\n                        probs.append(mdl.predict_proba(X_te)[:, 1])\n                    pred[action] = np.mean(probs, axis=0)\n\n            del X_te\n            gc.collect()\n\n            if pred.shape[1] != 0:\n                submission_list.append(predict_multiclass_adaptive(pred, meta_te))\n        except Exception as e:\n            print(e)\n            try:\n                del data_te\n            except:\n                pass\n            gc.collect()\n\n\n# ==================== ROBUSTIFY SUBMISSION ====================\n\ndef robustify(submission, dataset, traintest, traintest_directory=None):\n    if traintest_directory is None:\n        traintest_directory = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking\"\n\n    submission = submission[submission.start_frame < submission.stop_frame]\n\n    group_list = []\n    for _, group in submission.groupby(['video_id', 'agent_id', 'target_id']):\n        group = group.sort_values('start_frame')\n        mask = np.ones(len(group), dtype=bool)\n        last_stop = 0\n        for i, (_, row) in enumerate(group.iterrows()):\n            if row['start_frame'] < last_stop:\n                mask[i] = False\n            else:\n                last_stop = row['stop_frame']\n        group_list.append(group[mask])\n    submission = pd.concat(group_list) if group_list else submission\n\n    s_list = []\n    for _, row in dataset.iterrows():\n        lab_id = row['lab_id']\n        video_id = row['video_id']\n        if (submission.video_id == video_id).any():\n            continue\n\n        if verbose:\n            print(f\"Video {video_id} has no predictions\")\n\n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        vid = pd.read_parquet(path)\n\n        vid_behaviors = eval(row['behaviors_labeled'])\n        vid_behaviors = sorted(list({b.replace(\"'\", \"\") for b in vid_behaviors}))\n        vid_behaviors = [b.split(',') for b in vid_behaviors]\n        vid_behaviors = pd.DataFrame(vid_behaviors, columns=['agent', 'target', 'action'])\n\n        start_frame = vid.video_frame.min()\n        stop_frame = vid.video_frame.max() + 1\n\n        for (agent, target), actions in vid_behaviors.groupby(['agent', 'target']):\n            batch_len = int(np.ceil((stop_frame - start_frame) / len(actions)))\n            for i, (_, action_row) in enumerate(actions.iterrows()):\n                batch_start = start_frame + i * batch_len\n                batch_stop = min(batch_start + batch_len, stop_frame)\n                s_list.append((video_id, agent, target, action_row['action'], batch_start, batch_stop))\n\n    if len(s_list) > 0:\n        submission = pd.concat([\n            submission,\n            pd.DataFrame(s_list, columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame'])\n        ])\n\n    submission = submission.reset_index(drop=True)\n    return submission\n\n\n# ==================== MAIN LOOP ====================\n\nsubmission_list = []\n\nfor section in range(len(body_parts_tracked_list)):\n    body_parts_tracked_str = body_parts_tracked_list[section]\n    try:\n        body_parts_tracked = json.loads(body_parts_tracked_str)\n        print(f\"{section}. Processing: {len(body_parts_tracked)} body parts\")\n        if len(body_parts_tracked) > 5:\n            body_parts_tracked = [b for b in body_parts_tracked if b not in drop_body_parts]\n\n        train_subset = train[train.body_parts_tracked == body_parts_tracked_str]\n\n        _fps_lookup = (\n            train_subset[['video_id', 'frames_per_second']]\n            .drop_duplicates('video_id')\n            .set_index('video_id')['frames_per_second']\n            .to_dict()\n        )\n\n        single_list, single_label_list, single_meta_list = [], [], []\n        pair_list, pair_label_list, pair_meta_list = [], [], []\n\n        for switch, data, meta, label in generate_mouse_data(train_subset, 'train'):\n            if switch == 'single':\n                single_list.append(data)\n                single_meta_list.append(meta)\n                single_label_list.append(label)\n            else:\n                pair_list.append(data)\n                pair_meta_list.append(meta)\n                pair_label_list.append(label)\n\n        if len(single_list) > 0:\n            single_feats_parts = []\n            for data_i, meta_i in zip(single_list, single_meta_list):\n                fps_i = _fps_from_meta(meta_i, _fps_lookup, default_fps=30.0)\n                Xi = transform_single(data_i, body_parts_tracked, fps_i).astype(np.float32)\n                single_feats_parts.append(Xi)\n\n            X_tr = pd.concat(single_feats_parts, axis=0, ignore_index=True)\n\n            single_label = pd.concat(single_label_list, axis=0, ignore_index=True)\n            single_meta = pd.concat(single_meta_list, axis=0, ignore_index=True)\n\n            del single_list, single_label_list, single_meta_list, single_feats_parts\n            gc.collect()\n\n            print(f\"  Single: {X_tr.shape}\")\n            submit_ensemble(body_parts_tracked_str, 'single', X_tr, single_label, single_meta)\n\n            del X_tr, single_label, single_meta\n            gc.collect()\n\n        if len(pair_list) > 0:\n            pair_feats_parts = []\n            for data_i, meta_i in zip(pair_list, pair_meta_list):\n                fps_i = _fps_from_meta(meta_i, _fps_lookup, default_fps=30.0)\n                Xi = transform_pair(data_i, body_parts_tracked, fps_i).astype(np.float32)\n                pair_feats_parts.append(Xi)\n\n            X_tr = pd.concat(pair_feats_parts, axis=0, ignore_index=True)\n\n            pair_label = pd.concat(pair_label_list, axis=0, ignore_index=True)\n            pair_meta = pd.concat(pair_meta_list, axis=0, ignore_index=True)\n\n            del pair_list, pair_label_list, pair_meta_list, pair_feats_parts\n            gc.collect()\n\n            print(f\"  Pair: {X_tr.shape}\")\n            submit_ensemble(body_parts_tracked_str, 'pair', X_tr, pair_label, pair_meta)\n\n            del X_tr, pair_label, pair_meta\n            gc.collect()\n\n    except Exception as e:\n        print(f'***Exception*** {str(e)[:100]}')\n\n    gc.collect()\n    print()\n\nif len(submission_list) > 0:\n    submission = pd.concat(submission_list, ignore_index=True)\nelse:\n    submission = pd.DataFrame({\n        'video_id': [438887472],\n        'agent_id': ['mouse1'],\n        'target_id': ['self'],\n        'action': ['rear'],\n        'start_frame': [278],\n        'stop_frame': [500]\n    })\n\nsubmission_robust = robustify(submission, test, 'test')\nsubmission_robust.index.name = 'row_id'\nsubmission_robust.to_csv('submission.csv')\nprint(f\"\\nSubmission created: {len(submission_robust)} predictions\")\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"🏆 ADVANCED FEATURES ADDED:\")\nprint(\"=\"*60)\nprint(\"✅ Wavelet Transform Features (multi-scale frequency analysis)\")\nprint(\"✅ Physics-Informed Features (jerk, angular velocity, kinetic energy)\")\nprint(\"✅ Transition Detection Features (behavior onset/offset)\")\nprint(\"✅ Enhanced Pair Interaction Features (bearing, time-to-collision)\")\nprint(\"✅ Relative Physics Features (relative velocity, acceleration)\")\nprint(\"=\"*60)","metadata":{"execution":{"iopub.status.busy":"2025-12-28T04:08:52.835774Z","iopub.execute_input":"2025-12-28T04:08:52.836718Z"},"papermill":{"duration":29835.973761,"end_time":"2025-12-08T13:04:43.596870","exception":false,"start_time":"2025-12-08T04:47:27.623109","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null}]}