{"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.9"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"}],"dockerImageVersionId":31192,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# CV Training and Inference Notebook\n\nThis notebook combines the best extracted code snippets from the MABe mouse behavior detection project to create a complete end-to-end workflow with:\n\n- **FPS-aware feature engineering** for handling variable frame rates\n- **Cross-validation training** with proper fold handling and metrics\n- **Mode-aware thresholding** for single vs pair behaviors\n- **Robust inference** with post-processing\n\n## Features\n\n- GroupKFold cross-validation (5 folds) with proper edge case handling\n- PRAUC and F1 metrics for comprehensive evaluation\n- OOF (Out-of-Fold) prediction collection\n- Two-stage prediction aggregation (models → folds)\n- FPS-aware temporal features\n- Mode-aware action thresholds\n","metadata":{}},{"cell_type":"markdown","source":"## 1. Imports and Configuration\n\nImport all required libraries and set up environment configuration.\n","metadata":{}},{"cell_type":"code","source":"# MABe Mouse Behavior Detection - CV Training and Inference Notebook\n# Built from extracted snippets with best practices\n\nvalidate_or_submit = 'submit'  # 'validate' or 'submit'\nverbose = True\n\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nimport itertools\nimport warnings\nimport json\nimport os\nimport random\nimport gc\nimport lightgbm\nfrom collections import defaultdict\nimport polars as pl\nfrom scipy import signal, stats\n\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, average_precision_score\n\nwarnings.filterwarnings('ignore')\n\n# Try importing additional models\ntry:\n    from xgboost import XGBClassifier\n    XGBOOST_AVAILABLE = True\nexcept:\n    XGBOOST_AVAILABLE = False\n    \ntry:\n    from catboost import CatBoostClassifier\n    CATBOOST_AVAILABLE = True\nexcept:\n    CATBOOST_AVAILABLE = False\n\n# ==================== GPU DETECTION ====================\nUSE_GPU = True  # Enable GPU by default - will be verified by detection\nGPU_DEVICE = 'cuda'  # Default to cuda for RTX 5090\n\ntry:\n    import torch\n    if torch.cuda.is_available():\n        USE_GPU = True\n        GPU_DEVICE = 'cuda'\n        if verbose:\n            print(f\"[INFO] GPU detected: {torch.cuda.get_device_name(0)}\")\n            print(f\"[INFO] CUDA version: {torch.version.cuda}\")\nexcept:\n    pass\n\nif not USE_GPU:\n    try:\n        import subprocess\n        result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)\n        if result.returncode == 0:\n            USE_GPU = True\n            GPU_DEVICE = 'gpu'\n            if verbose:\n                print(\"[INFO] GPU detected via nvidia-smi\")\n    except:\n        pass\n\nif not USE_GPU and verbose:\n    print(\"[INFO] No GPU detected - training will use CPU\")\n\n# ==================== LIBRARY-SPECIFIC GPU DETECTION ====================\n# Conservative approach: Default to False, only enable if GPU actually works\ndef check_lightgbm_gpu():\n    \"\"\"Check if LightGBM actually supports GPU\"\"\"\n    if not USE_GPU:\n        return False\n    try:\n        import lightgbm as lgb\n        import numpy as np\n        # Try to actually train with GPU\n        X = np.array([[1, 2], [3, 4], [5, 6]])\n        y = np.array([0, 1, 0])\n        train_data = lgb.Dataset(X, label=y)\n        params = {'device': 'gpu', 'verbose': -1, 'num_threads': 1}\n        # This will fail if GPU not available\n        model = lgb.train(params, train_data, num_boost_round=1)\n        return True\n    except Exception as e:\n        if verbose:\n            print(f\"  LightGBM GPU check failed: {str(e)[:100]}\")\n        return False\n\ndef check_xgboost_gpu():\n    \"\"\"Check if XGBoost actually supports GPU\"\"\"\n    if not USE_GPU:\n        return False\n    try:\n        from xgboost import XGBClassifier\n        import numpy as np\n        # Try to actually fit with gpu_hist\n        X = np.array([[1, 2], [3, 4]])\n        y = np.array([0, 1])\n        model = XGBClassifier(tree_method='gpu_hist', n_estimators=1, verbosity=0)\n        model.fit(X, y)\n        return True\n    except Exception as e:\n        if verbose:\n            print(f\"  XGBoost GPU check failed: {str(e)[:100]}\")\n        return False\n\ndef check_catboost_gpu():\n    \"\"\"Check if CatBoost supports GPU (it auto-falls back, so we can be optimistic)\"\"\"\n    if not USE_GPU:\n        return False\n    # CatBoost handles GPU fallback automatically, so we can enable it\n    # It will use CPU if GPU not available\n    return True\n\n# Set GPU availability flags (default to False for safety)\nLIGHTGBM_GPU_AVAILABLE = check_lightgbm_gpu()\nXGBOOST_GPU_AVAILABLE = check_xgboost_gpu()\nCATBOOST_GPU_AVAILABLE = check_catboost_gpu()\n\nif verbose:\n    print(f\"[INFO] GPU Support: LightGBM={LIGHTGBM_GPU_AVAILABLE}, XGBoost={XGBOOST_GPU_AVAILABLE}, CatBoost={CATBOOST_GPU_AVAILABLE}\")\n\n# ==================== ENVIRONMENT DETECTION ====================\nIS_KAGGLE = os.path.exists('/kaggle/input')\nif IS_KAGGLE:\n    BASE_PATH = '/kaggle/input/MABe-mouse-behavior-detection'\n    TRACKING_BASE_PATH = BASE_PATH\n    if verbose:\n        print(\"[INFO] Running on Kaggle - using Kaggle data paths\")\nelse:\n    current_dir = os.getcwd()\n    if os.path.basename(current_dir) == 'notebook':\n        BASE_PATH = os.path.join('..', 'data', '01_raw')\n        project_root = os.path.abspath('..')\n    else:\n        BASE_PATH = os.path.join('data', '01_raw')\n        project_root = os.path.abspath('.')\n    \n    BASE_PATH = os.path.abspath(BASE_PATH)\n    TRACKING_BASE_PATH = BASE_PATH  # Use same path as CSV files (all data in data/01_raw)\n    if verbose:\n        print(f\"[INFO] Running locally - CSV data path: {BASE_PATH}\")\n        print(f\"[INFO] Running locally - Tracking data path: {TRACKING_BASE_PATH}\")\n\n# ==================== CV CONFIGURATION ====================\nUSE_CV = True  # Toggle CV on/off\nCV_N_SPLITS = 5  # Number of folds\nCV_GROUP_COLUMN = 'video_id'  # Grouping column for GroupKFold\n\n# ==================== LOCAL TESTING CONFIGURATION ====================\n# Set to False to use full dataset parameters (same as Kaggle)\nLOCAL_TESTING = False  # Disabled to use full dataset parameters\nMAX_VIDEOS_LOCAL = 5\nLOCAL_N_SAMPLES = 5000\n\nif LOCAL_TESTING:\n    validate_or_submit = 'submit'\n    if verbose:\n        print(f\"[INFO] Local testing mode enabled: max_videos={MAX_VIDEOS_LOCAL}, n_samples={LOCAL_N_SAMPLES}\")\n\nelse:\n    if verbose and not IS_KAGGLE:\n        print(\"[INFO] Running locally with full dataset parameters (same as Kaggle)\")\n\n# ==================== SEED EVERYTHING ====================\nSEED = 1234\nos.environ[\"PYTHONHASHSEED\"] = str(SEED)\nrnd = np.random.RandomState(SEED)\nrandom.seed(SEED)\nnp.random.seed(SEED)\n\nprint(f\"\\n[INFO] Configuration:\")\nprint(f\"  - USE_CV: {USE_CV}\")\nprint(f\"  - CV_N_SPLITS: {CV_N_SPLITS}\")\nprint(f\"  - CV_GROUP_COLUMN: {CV_GROUP_COLUMN}\")\nprint(f\"  - XGBoost: {XGBOOST_AVAILABLE}, CatBoost: {CATBOOST_AVAILABLE}\")\nprint(f\"  - GPU: {USE_GPU} ({GPU_DEVICE})\")\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. Data Loading Functions\n\nGenerator function to load and process video tracking data for both training and testing.\n","metadata":{}},{"cell_type":"code","source":"# Body parts to drop (optional, high-dimensional parts)\ndrop_body_parts = ['headpiece_bottombackleft', 'headpiece_bottombackright', 'headpiece_bottomfrontleft', 'headpiece_bottomfrontright', \n                   'headpiece_topbackleft', 'headpiece_topbackright', 'headpiece_topfrontleft', 'headpiece_topfrontright', \n                   'spine_1', 'spine_2', 'tail_middle_1', 'tail_middle_2', 'tail_midpoint']\n\ndef generate_mouse_data(dataset, traintest, traintest_directory=None, generate_single=True, generate_pair=True):\n    \"\"\"\n    Generator function to load and process video data one by one.\n    Yields data for single mice (agent=target) and mouse pairs (agent!=target).\n    \"\"\"\n    assert traintest in ['train', 'test']\n    if traintest_directory is None:\n        if IS_KAGGLE:\n            traintest_directory = os.path.join(BASE_PATH, f\"{traintest}_tracking\")\n        else:\n            traintest_directory = os.path.join(TRACKING_BASE_PATH, f\"{traintest}_tracking\")\n    \n    for _, row in dataset.iterrows():\n        lab_id = row.lab_id\n        video_id = row.video_id\n\n        if type(row.behaviors_labeled) != str:\n            if verbose: \n                print('No labeled behaviors:', lab_id, video_id)\n            continue\n\n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        try:\n            vid = pd.read_parquet(path)\n        except FileNotFoundError:\n            if verbose:\n                print(f\"File not found: {path}\")\n            continue\n        \n        if len(np.unique(vid.bodypart)) > 5:\n            vid = vid.query(\"~ bodypart.isin(@drop_body_parts)\")\n        \n        pvid = vid.pivot(columns=['mouse_id', 'bodypart'], index='video_frame', values=['x', 'y'])\n        if pvid.isna().any().any():\n            if verbose and traintest == 'test': \n                print('video with missing values', video_id, traintest, len(vid), 'frames')\n        else:\n            if verbose and traintest == 'test': \n                print('video with all values', video_id, traintest, len(vid), 'frames')\n        del vid\n        pvid = pvid.reorder_levels([1, 2, 0], axis=1).T.sort_index().T\n        pvid /= row.pix_per_cm_approx\n\n        vid_behaviors = json.loads(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        annot = None\n        if traintest == 'train':\n            try:\n                if IS_KAGGLE:\n                    annot_path = path.replace('train_tracking', 'train_annotation')\n                else:\n                    # Simple replacement - works for both Windows and Unix paths\n                    annot_path = path.replace('train_tracking', 'train_annotation')\n                annot = pd.read_parquet(annot_path)\n            except FileNotFoundError:\n                if verbose:\n                    print(f\"  WARNING: Annotation file not found: {annot_path}\")\n                    print(f\"    Tracking path was: {path}\")\n                continue\n\n        if generate_single:\n            vid_behaviors_subset = vid_behaviors.query(\"target == 'self'\")\n            for mouse_id_str in np.unique(vid_behaviors_subset.agent):\n                try:\n                    mouse_id = int(mouse_id_str[-1])\n                    vid_agent_actions = np.unique(vid_behaviors_subset.query(\"agent == @mouse_id_str\").action)\n                    single_mouse = pvid.loc[:, mouse_id]\n                    assert len(single_mouse) == len(pvid)\n                    single_mouse_meta = pd.DataFrame({\n                        'video_id': video_id,\n                        'agent_id': mouse_id_str,\n                        'target_id': 'self',\n                        'video_frame': single_mouse.index,\n                        'frames_per_second': row.frames_per_second\n                    })\n                    if traintest == 'train' and annot is not None:\n                        single_mouse_label = pd.DataFrame(0.0, columns=vid_agent_actions, index=single_mouse.index)\n                        if len(annot) > 0 and 'agent_id' in annot.columns and 'target_id' in annot.columns:\n                            annot_subset = annot[(annot['agent_id'] == mouse_id) & (annot['target_id'] == mouse_id)]\n                        else:\n                            annot_subset = pd.DataFrame()\n                        for i in range(len(annot_subset)):\n                            annot_row = annot_subset.iloc[i]\n                            if annot_row['action'] in vid_agent_actions:\n                                single_mouse_label.loc[annot_row['start_frame']:annot_row['stop_frame'], annot_row['action']] = 1.0\n                        yield 'single', single_mouse, single_mouse_meta, single_mouse_label\n                    else:\n                        if verbose: \n                            print('- test single', video_id, mouse_id)\n                        yield 'single', single_mouse, single_mouse_meta, vid_agent_actions\n                except KeyError:\n                    pass\n\n        if generate_pair:\n            vid_behaviors_subset = vid_behaviors.query(\"target != 'self'\")\n            if len(vid_behaviors_subset) > 0:\n                for agent, target in itertools.permutations(np.unique(pvid.columns.get_level_values('mouse_id')), 2):\n                    agent_str = f\"mouse{agent}\"\n                    target_str = f\"mouse{target}\"\n                    vid_agent_actions = np.unique(vid_behaviors_subset.query(\"(agent == @agent_str) & (target == @target_str)\").action)\n                    mouse_pair = pd.concat([pvid[agent], pvid[target]], axis=1, keys=['A', 'B'])\n                    assert len(mouse_pair) == len(pvid)\n                    mouse_pair_meta = pd.DataFrame({\n                        'video_id': video_id,\n                        'agent_id': agent_str,\n                        'target_id': target_str,\n                        'video_frame': mouse_pair.index,\n                        'frames_per_second': row.frames_per_second\n                    })\n                    if traintest == 'train' and annot is not None:\n                        mouse_pair_label = pd.DataFrame(0.0, columns=vid_agent_actions, index=mouse_pair.index)\n                        if len(annot) > 0 and 'agent_id' in annot.columns and 'target_id' in annot.columns:\n                            annot_subset = annot[\n                                (annot['agent_id'] == agent) &\n                                (annot['target_id'] == target)\n                            ]\n                        else:\n                            annot_subset = pd.DataFrame()\n                        for i in range(len(annot_subset)):\n                            annot_row = annot_subset.iloc[i]\n                            if annot_row['action'] in vid_agent_actions:\n                                mouse_pair_label.loc[annot_row['start_frame']:annot_row['stop_frame'], annot_row.action] = 1.0\n                        yield 'pair', mouse_pair, mouse_pair_meta, mouse_pair_label\n                    else:\n                        if verbose: \n                            print('- test pair', video_id, agent, target)\n                        yield 'pair', mouse_pair, mouse_pair_meta, vid_agent_actions\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3. Feature Engineering Utilities\n\nFPS-aware utility functions for scaling temporal windows and extracting features.\n","metadata":{}},{"cell_type":"code","source":"def 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 (keeps at least 1 frame when |n|>=1).\"\"\"\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    \"\"\"Extract FPS from metadata DataFrame with fallback lookup.\"\"\"\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 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_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","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4. Transform Functions\n\nFPS-aware feature transformation functions for single mouse and mouse pair behaviors.\n","metadata":{}},{"cell_type":"code","source":"def transform_single(single_mouse, body_parts_tracked, fps):\n    \"\"\"Enhanced single mouse transform (FPS-aware windows/lags; distances in cm).\"\"\"\n    available_body_parts = single_mouse.columns.get_level_values(0)\n\n    # Base distance features (squared distances across body parts)\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 via lagged displacements (duration-aware lag)\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 (orientation)\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 (windows scaled by fps)\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        # Advanced features (fps-scaled)\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\n    # Nose-tail features with duration-aware lags\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 with duration-aware offsets\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\ndef transform_pair(mouse_pair, body_parts_tracked, fps):\n    \"\"\"Enhanced pair transform (FPS-aware windows/lags; distances in cm).\"\"\"\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 (squared distances across all part pairs)\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 via lagged displacements (duration-aware lag)\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    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    # 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 (duration-aware lag)\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 (cm; unchanged by fps)\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 (fps-adjusted windows)\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    # Nose-nose dynamics (duration-aware lags)\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 (duration-aware offsets)\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 (fps-adjusted internals)\n        X = add_interaction_features(X, mouse_pair, avail_A, avail_B, fps)\n\n    return X.astype(np.float32, copy=False)\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"class StratifiedSubsetClassifier(ClassifierMixin, BaseEstimator):\n    \"\"\"A wrapper to train an estimator on a stratified subsample of data.\"\"\"\n    def __init__(self, estimator, n_samples=None):\n        self.estimator = estimator\n        self.n_samples = n_samples  # if None → no subsampling\n\n    def _to_numpy(self, X):\n        try:\n            return X.to_numpy(np.float32, copy=False)\n        except AttributeError:\n            return np.asarray(X, dtype=np.float32)\n\n    def fit(self, X, y):\n        Xn = self._to_numpy(X)\n        y = np.asarray(y).ravel()\n\n        # Handle rare cases where labels might be [0, 2]\n        uniq = np.unique(y[~pd.isna(y)])\n        if set(uniq.tolist()) == {0, 2}:\n            y = (y > 0).astype(np.int8)\n\n        # Calculate class weights for imbalance handling\n        n_pos = int((y == 1).sum())\n        n_neg = len(y) - n_pos\n        class_weight = n_neg / n_pos if n_pos > 0 and n_neg > 0 else 1.0\n        \n        # Set class imbalance handling for different model types\n        if hasattr(self.estimator, 'set_params'):\n            # LightGBM\n            if hasattr(self.estimator, 'is_unbalance'):\n                self.estimator.set_params(is_unbalance=True)\n            # XGBoost\n            if hasattr(self.estimator, 'scale_pos_weight'):\n                self.estimator.set_params(scale_pos_weight=class_weight)\n            # CatBoost\n            if hasattr(self.estimator, 'auto_class_weights'):\n                try:\n                    self.estimator.set_params(auto_class_weights='Balanced')\n                except Exception:\n                    pass\n\n        # Calculate minimum positive samples required\n        min_pos = max(100, int(self.n_samples * 0.01)) if self.n_samples else 100\n        n_pos_actual = int((y == 1).sum())\n\n        # If n_samples is None or data is small, fit on full data\n        if self.n_samples is None or len(Xn) <= int(self.n_samples):\n            self.estimator.fit(Xn, y)\n        else:\n            # If too few positives, use all data instead of subsampling\n            if n_pos_actual < min_pos:\n                self.estimator.fit(Xn, y)\n            else:\n                # Fit on a stratified subset that preserves class balance\n                sss = StratifiedShuffleSplit(n_splits=1, train_size=int(self.n_samples), random_state=42)\n                try:\n                    idx, _ = next(sss.split(np.zeros_like(y), y))\n                    # Verify we have positive samples in the subset\n                    if (y[idx] == 1).sum() > 0:\n                        self.estimator.fit(Xn[idx], y[idx])\n                    else:\n                        # Fallback: use all data if stratified split loses all positives\n                        self.estimator.fit(Xn, y)\n                except Exception:\n                    # Fallback for cases where stratification fails\n                    step = max(len(Xn) // int(self.n_samples), 1)\n                    self.estimator.fit(Xn[::step], y[::step])\n\n        try:\n            self.classes_ = np.asarray(self.estimator.classes_)\n        except Exception:\n            self.classes_ = np.unique(y)\n        return self\n\n    def predict_proba(self, X):\n        Xn = self._to_numpy(X)\n        try:\n            P = self.estimator.predict_proba(Xn)\n        except Exception:\n            if len(self.classes_) == 1:\n                n = len(Xn)\n                c = int(self.classes_[0])\n                if c == 1:\n                    return np.column_stack([np.zeros(n, dtype=np.float32), np.ones(n, dtype=np.float32)])\n                else:\n                    return np.column_stack([np.ones(n, dtype=np.float32), np.zeros(n, dtype=np.float32)])\n            return np.full((len(Xn), 2), 0.5, dtype=np.float32)\n\n        P = np.asarray(P)\n        if P.ndim == 1:\n            P1 = P.astype(np.float32)\n            return np.column_stack([1.0 - P1, P1])\n        if P.shape[1] == 1 and len(self.classes_) == 2:\n            P1 = P[:, 0].astype(np.float32)\n            return np.column_stack([1.0 - P1, P1])\n        return P\n\n    def predict(self, X):\n        Xn = self._to_numpy(X)\n        try:\n            return self.estimator.predict(Xn)\n        except Exception:\n            return np.argmax(self.predict_proba(Xn), axis=1)\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"action_thresholds = {\n    \"default\": 0.27,\n    \"single_default\": 0.27,\n    \"pair_default\": 0.27,\n    \"single\": {\n        \"rear\": 0.30,\n    },\n}\n\ndef _select_threshold_map(thresholds, mode: str):\n    \"\"\"Select threshold map based on mode (single or pair).\"\"\"\n    if isinstance(thresholds, dict):\n        # mode-aware?\n        if (\"single\" in thresholds) or (\"pair\" in thresholds) or \\\n           (\"single_default\" in thresholds) or (\"pair_default\" in thresholds):\n            base_default = float(thresholds.get(\"default\", 0.27))\n            mode_default = float(thresholds.get(f\"{mode}_default\", base_default))\n            mode_overrides = thresholds.get(mode, {}) or {}\n            out = defaultdict(lambda: mode_default)\n            out.update({str(k): float(v) for k, v in mode_overrides.items()})\n            return out\n        # plain per-action dict\n        out = defaultdict(lambda: float(thresholds.get(\"default\", 0.27)))\n        out.update({str(k): float(v) for k, v in thresholds.items() if k != \"default\"})\n        return out\n    return defaultdict(lambda: 0.27)\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 7. Cross-Validation Training Function\n\nGroupKFold CV training with proper fold handling, OOF predictions, and metrics.\n","metadata":{}},{"cell_type":"code","source":"def train_ensemble_single_split(X_tr, label, models, n_samples):\n    \"\"\"Original single-split training (fallback when CV is disabled).\"\"\"\n    X_tr_np = X_tr.to_numpy(np.float32, copy=False)\n    model_list = []\n    \n    for action in label.columns:\n        y_raw = label[action].to_numpy()\n        mask = ~pd.isna(y_raw)\n        y_action = y_raw[mask].astype(int)\n        if not (y_action == 0).all() and np.sum(y_action) >= 5:\n            trained = []\n            idx = np.flatnonzero(mask)\n            for m in models:\n                m_clone = clone(m)\n                m_clone.fit(X_tr_np[idx], y_action)\n                trained.append(m_clone)\n            model_list.append((action, [trained]))  # Wrap in list for compatibility\n    \n    return model_list\n\ndef optimize_threshold(y_true, y_pred_proba):\n    \"\"\"Optimize threshold for F1 score on validation set.\"\"\"\n    thresholds = np.arange(0.1, 0.9, 0.01)\n    best_threshold = 0.5\n    best_f1 = 0\n    for thresh in thresholds:\n        y_pred = (y_pred_proba >= thresh).astype(int)\n        f1 = f1_score(y_true, y_pred, zero_division=0)\n        if f1 > best_f1:\n            best_f1 = f1\n            best_threshold = thresh\n    return best_threshold, best_f1\n\ndef train_ensemble_with_cv(X_tr, label, meta, models, n_samples, switch_tr, body_parts_tracked_str):\n    \"\"\"\n    Train ensemble models using GroupKFold cross-validation.\n    \n    Returns:\n        model_list: List of (action, list_of_fold_models) tuples\n        oof_predictions: Dict mapping action -> DataFrame with OOF predictions\n        cv_metrics: Dict mapping action -> CV metrics (mean ± std)\n    \"\"\"\n    if not USE_CV:\n        # Fall back to single-split training\n        return train_ensemble_single_split(X_tr, label, models, n_samples), {}, {}\n    \n    # Get grouping column\n    if CV_GROUP_COLUMN not in meta.columns:\n        if verbose:\n            print(f\"Warning: {CV_GROUP_COLUMN} not in meta, falling back to single-split training\")\n        return train_ensemble_single_split(X_tr, label, models, n_samples), {}, {}\n    \n    groups = meta[CV_GROUP_COLUMN].values\n    X_tr_np = X_tr.to_numpy(np.float32, copy=False)\n    \n    all_models = {}\n    all_oof_preds = {}\n    all_cv_metrics = {}\n    \n    if verbose:\n        print(f\"\\n{'='*60}\")\n        print(f\"Training with {CV_N_SPLITS}-fold GroupKFold CV\")\n        print(f\"Grouping by: {CV_GROUP_COLUMN}\")\n        print(f\"{'='*60}\\n\")\n    \n    for action in label.columns:\n        y_raw = label[action].to_numpy()\n        mask = ~pd.isna(y_raw)\n        y_action = y_raw[mask].astype(int)\n        X_action = X_tr_np[mask]\n        groups_action = groups[mask]\n        \n        # Skip if no positive examples or insufficient data\n        if (y_action == 0).all() or np.sum(y_action) < 5:\n            continue\n        \n        if len(np.unique(y_action)) < 2:\n            continue\n        \n        # Calculate positive rate for adaptive CV folds\n        pos_rate = float(np.mean(y_action == 1))\n        n_pos_total = int(np.sum(y_action == 1))\n        \n        # Adaptive CV folds based on class distribution\n        if pos_rate < 0.01:  # Very rare (< 1%)\n            n_splits = 3\n        elif pos_rate < 0.05:  # Rare (< 5%)\n            n_splits = 5\n        else:\n            n_splits = CV_N_SPLITS\n        \n        # Ensure we have enough groups for the number of splits\n        n_unique_groups = len(np.unique(groups_action))\n        n_splits = min(n_splits, n_unique_groups)\n        \n        if verbose:\n            print(f\"Training CV for action: {action} (pos_rate={pos_rate:.4f}, n_splits={n_splits})\")\n        \n        # Create GroupKFold with adaptive splits\n        gkf = GroupKFold(n_splits=n_splits)\n        \n        # Generate CV splits\n        try:\n            splits = list(gkf.split(X_action, y_action, groups_action))\n        except ValueError as e:\n            if verbose:\n                print(f\"  Warning: GroupKFold failed for {action}: {e}. Using single-split training.\")\n            # Fall back to single-split for this action\n            idx = np.flatnonzero(mask)\n            trained = []\n            for m in models:\n                m_clone = clone(m)\n                m_clone.fit(X_tr_np[idx], y_action)\n                trained.append(m_clone)\n            all_models[action] = [trained]  # Single \"fold\"\n            continue\n        \n        action_models = []  # List of model lists, one per fold\n        action_oof_preds = []  # List of OOF prediction arrays\n        action_oof_indices = []  # List of validation indices\n        action_fold_metrics = []  # List of metrics per fold\n        \n        # Train on each fold\n        for fold_idx, (train_idx, val_idx) in enumerate(splits):\n            # Edge case check 1: Empty validation fold\n            if len(val_idx) == 0:\n                if verbose:\n                    print(f\"  Fold {fold_idx + 1}: Skipping (no validation samples)\")\n                continue\n            \n            X_train_fold = X_action[train_idx]\n            y_train_fold = y_action[train_idx]\n            X_val_fold = X_action[val_idx]\n            y_val_fold = y_action[val_idx]\n            \n            # Edge case check 2: No positive examples in validation\n            val_pos = np.sum(y_val_fold == 1)\n            train_pos = np.sum(y_train_fold == 1)\n            \n            # Minimum positive sample checks for rare classes\n            min_train_pos = 50  # Minimum positive samples in training\n            min_val_pos = 10   # Minimum positive samples in validation\n            \n            if val_pos < min_val_pos:\n                if verbose:\n                    print(f\"  Fold {fold_idx + 1}: Skipping (insufficient positive examples in validation: {val_pos} < {min_val_pos})\")\n                continue\n            \n            if train_pos < min_train_pos:\n                if verbose:\n                    print(f\"  Fold {fold_idx + 1}: Skipping (insufficient positive examples in training: {train_pos} < {min_train_pos})\")\n                continue\n            \n            # Edge case check 3: Single class in training\n            if len(np.unique(y_train_fold)) < 2:\n                if verbose:\n                    print(f\"  Fold {fold_idx + 1}: Skipping (only one class in training)\")\n                continue\n            \n            if verbose:\n                print(f\"  Fold {fold_idx + 1}/{len(splits)}: train={len(train_idx)}, val={len(val_idx)} (pos={val_pos})\")\n            \n            # Train models on this fold\n            fold_models = []\n            fold_oof_preds = []\n            \n            for m in models:\n                m_clone = clone(m)\n                try:\n                    # Fit model (early stopping is handled by model parameters and StratifiedSubsetClassifier)\n                    # The improved model parameters (is_unbalance, min_child_samples, etc.) already help with training stability\n                    m_clone.fit(X_train_fold, y_train_fold)\n                    \n                    # Get OOF predictions\n                    val_proba = m_clone.predict_proba(X_val_fold)[:, 1]\n                    fold_oof_preds.append(val_proba)\n                    fold_models.append(m_clone)\n                except Exception as e:\n                    if verbose:\n                        print(f\"    Warning: Model training failed: {e}\")\n                    continue\n            \n            if not fold_models:\n                if verbose:\n                    print(f\"  Fold {fold_idx + 1}: No models trained successfully\")\n                continue\n            \n            # Average OOF predictions across models for this fold\n            fold_oof_mean = np.mean(fold_oof_preds, axis=0)\n            action_oof_preds.append(fold_oof_mean)\n            action_oof_indices.append(val_idx)\n            action_models.append(fold_models)\n            \n            # Compute fold metrics\n            fold_metrics = {\n                'prauc': average_precision_score(y_val_fold, fold_oof_mean),\n                'f1': f1_score(y_val_fold, (fold_oof_mean >= 0.5).astype(int), zero_division=0),\n            }\n            action_fold_metrics.append(fold_metrics)\n            \n            if verbose:\n                print(f\"    PRAUC={fold_metrics['prauc']:.4f}, F1={fold_metrics['f1']:.4f}\")\n        \n        if not action_models:\n            if verbose:\n                print(f\"  Warning: No valid folds for {action}, skipping\")\n            continue\n        \n        # Store models (list of lists, one inner list per fold)\n        all_models[action] = action_models\n        \n        # Aggregate OOF predictions\n        oof_df = pd.DataFrame(index=np.where(mask)[0])\n        for fold_idx, (val_idx, oof_pred) in enumerate(zip(action_oof_indices, action_oof_preds)):\n            original_val_idx = np.where(mask)[0][val_idx]\n            oof_df.loc[original_val_idx, f'fold_{fold_idx}'] = oof_pred\n        \n        fold_cols = [c for c in oof_df.columns if c.startswith('fold_')]\n        if fold_cols:\n            oof_df['oof_mean'] = oof_df[fold_cols].mean(axis=1)\n        all_oof_preds[action] = oof_df\n        \n        # Aggregate CV metrics and optimize threshold\n        if action_fold_metrics:\n            # Optimize threshold using aggregated OOF predictions\n            if 'oof_mean' in oof_df.columns and len(oof_df) > 0:\n                # Get true labels for OOF predictions - oof_df.index contains original indices\n                # Map back to y_action using the mask\n                oof_indices = oof_df.index.values\n                # oof_indices are already in the space of mask, so we can index y_action directly\n                oof_true = y_action[oof_indices]\n                oof_pred_mean = oof_df['oof_mean'].values\n                # Only optimize if we have both classes\n                if len(np.unique(oof_true)) >= 2:\n                    opt_threshold, opt_f1 = optimize_threshold(oof_true, oof_pred_mean)\n                else:\n                    opt_threshold, opt_f1 = 0.5, 0.0\n            else:\n                opt_threshold, opt_f1 = 0.5, 0.0\n            \n            cv_metrics = {\n                'prauc_mean': np.mean([m['prauc'] for m in action_fold_metrics]),\n                'prauc_std': np.std([m['prauc'] for m in action_fold_metrics]),\n                'f1_mean': np.mean([m['f1'] for m in action_fold_metrics]),\n                'f1_std': np.std([m['f1'] for m in action_fold_metrics]),\n                'n_folds': len(action_fold_metrics),\n                'opt_threshold': opt_threshold,\n                'opt_f1': opt_f1,\n            }\n            all_cv_metrics[action] = cv_metrics\n            \n            if verbose:\n                print(f\"  CV Summary: PRAUC={cv_metrics['prauc_mean']:.4f}±{cv_metrics['prauc_std']:.4f}, \"\n                      f\"F1={cv_metrics['f1_mean']:.4f}±{cv_metrics['f1_std']:.4f} ({cv_metrics['n_folds']} folds)\\n\")\n    \n    # Convert to list format compatible with submit_ensemble\n    model_list = [(action, fold_models_list) for action, fold_models_list in all_models.items()]\n    \n    return model_list, all_oof_preds, all_cv_metrics\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"def predict_multiclass_adaptive(pred, meta, action_thresholds):\n    \"\"\"Adaptive thresholding per action + temporal smoothing\"\"\"\n    # Apply temporal smoothing\n    pred_smoothed = pred.rolling(window=5, min_periods=1, center=True).mean()\n\n    mode = 'pair'\n    try:\n        if 'target_id' in meta.columns and meta['target_id'].eq('self').all():\n            mode = 'single'\n    except Exception:\n        pass\n\n    ama = np.argmax(pred_smoothed, axis=1)\n    th_map = _select_threshold_map(action_thresholds, mode)\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 = th_map[action]\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    mask[-1] = False\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","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 9. Robustify Function\n\nPost-processing function to handle missing predictions and clean overlapping segments.\n","metadata":{}},{"cell_type":"code","source":"def robustify(submission, dataset, traintest, traintest_directory=None):\n    \"\"\"Ensures all videos have at least one prediction and cleans overlaps.\"\"\"\n    if traintest_directory is None:\n        if IS_KAGGLE:\n            traintest_directory = os.path.join(BASE_PATH, f\"{traintest}_tracking\")\n        else:\n            traintest_directory = os.path.join(TRACKING_BASE_PATH, f\"{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        try:\n            vid = pd.read_parquet(path)\n        except FileNotFoundError:\n            continue\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","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 10. Training Execution\n\nLoad training data, apply feature engineering, and train models with CV.\n","metadata":{}},{"cell_type":"code","source":"# Load training data\ntrain = pd.read_csv(os.path.join(BASE_PATH, 'train.csv'))\ntrain['n_mice'] = 4 - train[['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\n# Store trained models, CV metrics, and OOF predictions\nall_trained_models = {}\nall_cv_metrics_global = {}\nall_oof_predictions_global = {}\n\nprint(f\"Found {len(body_parts_tracked_list)} body part configurations\\n\")\n\n# Process each body part configuration\nfor section in range(1, 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        if LOCAL_TESTING and len(train_subset) > MAX_VIDEOS_LOCAL:\n            train_subset = train_subset.head(MAX_VIDEOS_LOCAL)\n            if verbose:\n                print(f\"  [Local testing] Using first {MAX_VIDEOS_LOCAL} videos\")\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        # Process single mouse behaviors\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        # Train single mouse models\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            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            \n            # Create models - 3 diverse LightGBM models for better ensemble\n            models = []\n            # Model 1: Conservative, balanced\n            models.append(make_pipeline(\n                StratifiedSubsetClassifier(\n                    lightgbm.LGBMClassifier(\n                        n_estimators=300, learning_rate=0.05, max_depth=6,\n                        num_leaves=31, min_child_samples=20, min_data_in_leaf=10,\n                        subsample=0.8, colsample_bytree=0.8,\n                        reg_alpha=0.1, reg_lambda=0.1,\n                        is_unbalance=True,\n                        random_state=SEED, verbose=-1, \n                        device='gpu' if LIGHTGBM_GPU_AVAILABLE else 'cpu'\n                    ), n_samples=LOCAL_N_SAMPLES if LOCAL_TESTING else 2_000_000\n                )\n            ))\n            # Model 2: Deeper, more complex\n            models.append(make_pipeline(\n                StratifiedSubsetClassifier(\n                    lightgbm.LGBMClassifier(\n                        n_estimators=200, learning_rate=0.08, max_depth=8,\n                        num_leaves=63, min_child_samples=15, min_data_in_leaf=8,\n                        subsample=0.75, colsample_bytree=0.85,\n                        reg_alpha=0.2, reg_lambda=0.2,\n                        is_unbalance=True,\n                        random_state=SEED, verbose=-1, \n                        device='gpu' if LIGHTGBM_GPU_AVAILABLE else 'cpu'\n                    ), n_samples=LOCAL_N_SAMPLES if LOCAL_TESTING else int(2_000_000 * 0.9)\n                )\n            ))\n            # Model 3: Larger leaves, more regularization\n            models.append(make_pipeline(\n                StratifiedSubsetClassifier(\n                    lightgbm.LGBMClassifier(\n                        n_estimators=150, learning_rate=0.1, max_depth=10,\n                        num_leaves=127, min_child_samples=30, min_data_in_leaf=15,\n                        subsample=0.7, colsample_bytree=0.9,\n                        reg_alpha=0.3, reg_lambda=0.3,\n                        is_unbalance=True,\n                        random_state=SEED, verbose=-1, \n                        device='gpu' if LIGHTGBM_GPU_AVAILABLE else 'cpu'\n                    ), n_samples=LOCAL_N_SAMPLES if LOCAL_TESTING else int(2_000_000 * 0.8)\n                )\n            ))\n            \n            if XGBOOST_AVAILABLE:\n                models.append(make_pipeline(\n                    StratifiedSubsetClassifier(\n                        XGBClassifier(\n                            n_estimators=200, learning_rate=0.1, max_depth=7,\n                            random_state=SEED, \n                            tree_method='gpu_hist' if XGBOOST_GPU_AVAILABLE else 'hist', \n                            verbosity=0\n                        ), n_samples=LOCAL_N_SAMPLES if LOCAL_TESTING else int(2_000_000/1.5)\n                    )\n                ))\n            \n            if CATBOOST_AVAILABLE:\n                models.append(make_pipeline(\n                    StratifiedSubsetClassifier(\n                        CatBoostClassifier(\n                            iterations=120, learning_rate=0.1, depth=6,\n                            verbose=False, allow_writing_files=False, \n                            task_type='GPU' if CATBOOST_GPU_AVAILABLE else 'CPU',\n                            random_seed=SEED\n                        ), n_samples=LOCAL_N_SAMPLES if LOCAL_TESTING else 2_000_000\n                    )\n                ))\n            \n            # Train with CV\n            model_list, oof_preds, cv_metrics = train_ensemble_with_cv(\n                X_tr, single_label, single_meta, models, \n                LOCAL_N_SAMPLES if LOCAL_TESTING else 2_000_000,\n                'single', body_parts_tracked_str\n            )\n            \n            all_trained_models[(body_parts_tracked_str, 'single')] = model_list\n            all_cv_metrics_global[(body_parts_tracked_str, 'single')] = cv_metrics\n            all_oof_predictions_global[(body_parts_tracked_str, 'single')] = oof_preds\n\n            del X_tr, single_label, single_meta\n            gc.collect()\n\n        # Train pair mouse models\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            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            \n            # Create models - 3 diverse LightGBM models for better ensemble\n            models = []\n            # Model 1: Conservative, balanced\n            models.append(make_pipeline(\n                StratifiedSubsetClassifier(\n                    lightgbm.LGBMClassifier(\n                        n_estimators=300, learning_rate=0.05, max_depth=6,\n                        num_leaves=31, min_child_samples=20, min_data_in_leaf=10,\n                        subsample=0.8, colsample_bytree=0.8,\n                        reg_alpha=0.1, reg_lambda=0.1,\n                        is_unbalance=True,\n                        random_state=SEED, verbose=-1, \n                        device='gpu' if LIGHTGBM_GPU_AVAILABLE else 'cpu'\n                    ), n_samples=LOCAL_N_SAMPLES if LOCAL_TESTING else 900_000\n                )\n            ))\n            # Model 2: Deeper, more complex\n            models.append(make_pipeline(\n                StratifiedSubsetClassifier(\n                    lightgbm.LGBMClassifier(\n                        n_estimators=200, learning_rate=0.08, max_depth=8,\n                        num_leaves=63, min_child_samples=15, min_data_in_leaf=8,\n                        subsample=0.75, colsample_bytree=0.85,\n                        reg_alpha=0.2, reg_lambda=0.2,\n                        is_unbalance=True,\n                        random_state=SEED, verbose=-1, \n                        device='gpu' if LIGHTGBM_GPU_AVAILABLE else 'cpu'\n                    ), n_samples=LOCAL_N_SAMPLES if LOCAL_TESTING else int(900_000 * 0.9)\n                )\n            ))\n            # Model 3: Larger leaves, more regularization\n            models.append(make_pipeline(\n                StratifiedSubsetClassifier(\n                    lightgbm.LGBMClassifier(\n                        n_estimators=150, learning_rate=0.1, max_depth=10,\n                        num_leaves=127, min_child_samples=30, min_data_in_leaf=15,\n                        subsample=0.7, colsample_bytree=0.9,\n                        reg_alpha=0.3, reg_lambda=0.3,\n                        is_unbalance=True,\n                        random_state=SEED, verbose=-1, \n                        device='gpu' if LIGHTGBM_GPU_AVAILABLE else 'cpu'\n                    ), n_samples=LOCAL_N_SAMPLES if LOCAL_TESTING else int(900_000 * 0.8)\n                )\n            ))\n            \n            if XGBOOST_AVAILABLE:\n                models.append(make_pipeline(\n                    StratifiedSubsetClassifier(\n                        XGBClassifier(\n                            n_estimators=200, learning_rate=0.1, max_depth=7,\n                            random_state=SEED, \n                            tree_method='gpu_hist' if XGBOOST_GPU_AVAILABLE else 'hist', \n                            verbosity=0\n                        ), n_samples=LOCAL_N_SAMPLES if LOCAL_TESTING else int(900_000/1.5)\n                    )\n                ))\n            \n            if CATBOOST_AVAILABLE:\n                models.append(make_pipeline(\n                    StratifiedSubsetClassifier(\n                        CatBoostClassifier(\n                            iterations=120, learning_rate=0.1, depth=6,\n                            verbose=False, allow_writing_files=False, \n                            task_type='GPU' if CATBOOST_GPU_AVAILABLE else 'CPU',\n                            random_seed=SEED\n                        ), n_samples=LOCAL_N_SAMPLES if LOCAL_TESTING else 900_000\n                    )\n                ))\n            \n            # Train with CV\n            model_list, oof_preds, cv_metrics = train_ensemble_with_cv(\n                X_tr, pair_label, pair_meta, models,\n                LOCAL_N_SAMPLES if LOCAL_TESTING else 900_000,\n                'pair', body_parts_tracked_str\n            )\n            \n            all_trained_models[(body_parts_tracked_str, 'pair')] = model_list\n            all_cv_metrics_global[(body_parts_tracked_str, 'pair')] = cv_metrics\n            all_oof_predictions_global[(body_parts_tracked_str, 'pair')] = oof_preds\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        import traceback\n        if verbose:\n            traceback.print_exc()\n\n    gc.collect()\n    print()\n\n# Display CV metrics summary\nif verbose and all_cv_metrics_global:\n    print(\"\\n\" + \"=\"*60)\n    print(\"CV Metrics Summary\")\n    print(\"=\"*60)\n    for key, metrics in all_cv_metrics_global.items():\n        if metrics:\n            print(f\"\\n{key[0][:30]}... ({key[1]}):\")\n            for action, cv_metrics in metrics.items():\n                print(f\"  {action}: PRAUC={cv_metrics['prauc_mean']:.4f}±{cv_metrics['prauc_std']:.4f}, \"\n                      f\"F1={cv_metrics['f1_mean']:.4f}±{cv_metrics['f1_std']:.4f} ({cv_metrics['n_folds']} folds)\")\n\n# Save CV metrics to CSV file\nif all_cv_metrics_global:\n    cv_results_list = []\n    for key, metrics in all_cv_metrics_global.items():\n        body_parts_str, mode = key\n        if metrics:\n            for action, cv_metrics in metrics.items():\n                cv_results_list.append({\n                    'body_parts_tracked': body_parts_str,\n                    'mode': mode,\n                    'action': action,\n                    'prauc_mean': cv_metrics['prauc_mean'],\n                    'prauc_std': cv_metrics['prauc_std'],\n                    'f1_mean': cv_metrics['f1_mean'],\n                    'f1_std': cv_metrics['f1_std'],\n                    'n_folds': cv_metrics['n_folds']\n                })\n    \n    if cv_results_list:\n        cv_results_df = pd.DataFrame(cv_results_list)\n        cv_results_df = cv_results_df.sort_values(['body_parts_tracked', 'mode', 'action'])\n        \n        # Save to CSV\n        cv_output_path = 'cv_results.csv' if IS_KAGGLE else os.path.join('..', 'cv_results.csv')\n        cv_results_df.to_csv(cv_output_path, index=False)\n        \n        if verbose:\n            print(f\"\\n{'='*60}\")\n            print(f\"CV Results saved to: {cv_output_path}\")\n            print(f\"Total actions evaluated: {len(cv_results_df)}\")\n            print(f\"Average PRAUC: {cv_results_df['prauc_mean'].mean():.4f}\")\n            print(f\"Average F1: {cv_results_df['f1_mean'].mean():.4f}\")\n            print(f\"{'='*60}\\n\")\n\n# Save OOF (Out-of-Fold) predictions\nif all_oof_predictions_global:\n    oof_output_path = 'oof_predictions.parquet' if IS_KAGGLE else os.path.join('..', 'oof_predictions.parquet')\n    oof_csv_path = 'oof_predictions.csv' if IS_KAGGLE else os.path.join('..', 'oof_predictions.csv')\n    \n    # Collect all OOF predictions with metadata\n    oof_data_list = []\n    for key, oof_preds_dict in all_oof_predictions_global.items():\n        body_parts_str, mode = key\n        if oof_preds_dict:\n            for action, oof_df in oof_preds_dict.items():\n                # Add metadata columns\n                oof_df_copy = oof_df.copy()\n                oof_df_copy['body_parts_tracked'] = body_parts_str\n                oof_df_copy['mode'] = mode\n                oof_df_copy['action'] = action\n                oof_df_copy = oof_df_copy.reset_index()\n                oof_df_copy.rename(columns={'index': 'sample_index'}, inplace=True)\n                oof_data_list.append(oof_df_copy)\n    \n    if oof_data_list:\n        # Combine all OOF predictions\n        oof_combined = pd.concat(oof_data_list, ignore_index=True)\n        \n        # Save as Parquet (more efficient for large data)\n        try:\n            oof_combined.to_parquet(oof_output_path, index=False, engine='pyarrow')\n            parquet_saved = True\n        except Exception as e:\n            if verbose:\n                print(f\"Warning: Could not save Parquet file ({str(e)[:50]}), saving as CSV only\")\n            parquet_saved = False\n        \n        # Also save as CSV for easy viewing (may be large)\n        # Only save CSV if not too large (limit to avoid memory issues)\n        if len(oof_combined) < 10_000_000:  # 10M rows limit\n            oof_combined.to_csv(oof_csv_path, index=False)\n            if verbose:\n                if parquet_saved:\n                    print(f\"OOF Predictions saved to:\")\n                    print(f\"  - Parquet: {oof_output_path}\")\n                    print(f\"  - CSV: {oof_csv_path}\")\n                else:\n                    print(f\"OOF Predictions saved to: {oof_csv_path}\")\n        else:\n            if verbose:\n                if parquet_saved:\n                    print(f\"OOF Predictions saved to: {oof_output_path} (CSV skipped - too large: {len(oof_combined):,} rows)\")\n                else:\n                    print(f\"OOF Predictions saved to: {oof_csv_path} (Parquet failed, CSV only)\")\n        \n        if verbose:\n            print(f\"Total OOF samples: {len(oof_combined):,}\")\n            print(f\"Actions with OOF predictions: {oof_combined['action'].nunique()}\")\n            print(f\"Body part configurations: {oof_combined['body_parts_tracked'].nunique()}\")\n            print(f\"{'='*60}\\n\")\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 11. Inference Execution\n\nLoad test data, generate predictions using trained CV models, and apply post-processing.\n","metadata":{}},{"cell_type":"code","source":"if validate_or_submit == 'submit':\n    # Load test data\n    test = pd.read_csv(os.path.join(BASE_PATH, 'test.csv'))\n    submission_list = []\n    \n    print(f\"Generating predictions for {len(test)} test videos\\n\")\n    \n    # Process each body part configuration\n    for section in range(1, 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            test_subset = test[test.body_parts_tracked == body_parts_tracked_str]\n            \n            if LOCAL_TESTING and len(test_subset) > MAX_VIDEOS_LOCAL:\n                test_subset = test_subset.head(MAX_VIDEOS_LOCAL)\n                if verbose:\n                    print(f\"  [Local testing] Using first {MAX_VIDEOS_LOCAL} videos\")\n\n            fps_lookup = (\n                test_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            # Process single mouse behaviors\n            for switch_te, data_te, meta_te, actions_te in generate_mouse_data(\n                test_subset, 'test',\n                generate_single=True,\n                generate_pair=False\n            ):\n                if switch_te != 'single':\n                    continue\n                \n                key = (body_parts_tracked_str, 'single')\n                if key not in all_trained_models:\n                    continue\n                \n                model_list = all_trained_models[key]\n                \n                try:\n                    fps_i = _fps_from_meta(meta_te, fps_lookup, default_fps=30.0)\n                    X_te = transform_single(data_te, body_parts_tracked, fps_i).astype(np.float32)\n                    X_te_np = X_te.to_numpy(np.float32, copy=False)\n                    del X_te, data_te\n                    gc.collect()\n\n                    pred = pd.DataFrame(index=meta_te.video_frame)\n                    for action, trained in model_list:\n                        if action in actions_te:\n                            # Handle CV models: trained is a list of fold model lists\n                            if len(trained) > 1:\n                                # CV mode: average across folds, then across models within each fold\n                                fold_probs = []\n                                for fold_models in trained:\n                                    fold_model_probs = [m.predict_proba(X_te_np)[:, 1] for m in fold_models]\n                                    fold_probs.append(np.mean(fold_model_probs, axis=0))\n                                # Average across folds\n                                pred[action] = np.mean(fold_probs, axis=0)\n                            else:\n                                # Single-split mode: unwrap the single fold and average across models\n                                fold_models = trained[0]\n                                probs = [m.predict_proba(X_te_np)[:, 1] for m in fold_models]\n                                pred[action] = np.mean(probs, axis=0)\n\n                    del X_te_np\n                    gc.collect()\n\n                    if pred.shape[1] != 0:\n                        sub_part = predict_multiclass_adaptive(pred, meta_te, action_thresholds)\n                        submission_list.append(sub_part)\n                    else:\n                        if verbose:\n                            print(\"  ERROR: no training data\")\n\n                except Exception as e:\n                    if verbose:\n                        print(f\"  ERROR: {str(e)[:50]}\")\n                    try:\n                        del data_te\n                    except:\n                        pass\n                    gc.collect()\n\n            # Process pair mouse behaviors\n            for switch_te, data_te, meta_te, actions_te in generate_mouse_data(\n                test_subset, 'test',\n                generate_single=False,\n                generate_pair=True\n            ):\n                if switch_te != 'pair':\n                    continue\n                \n                key = (body_parts_tracked_str, 'pair')\n                if key not in all_trained_models:\n                    continue\n                \n                model_list = all_trained_models[key]\n                \n                try:\n                    fps_i = _fps_from_meta(meta_te, fps_lookup, default_fps=30.0)\n                    X_te = transform_pair(data_te, body_parts_tracked, fps_i).astype(np.float32)\n                    X_te_np = X_te.to_numpy(np.float32, copy=False)\n                    del X_te, data_te\n                    gc.collect()\n\n                    pred = pd.DataFrame(index=meta_te.video_frame)\n                    for action, trained in model_list:\n                        if action in actions_te:\n                            # Handle CV models: trained is a list of fold model lists\n                            if len(trained) > 1:\n                                # CV mode: average across folds, then across models within each fold\n                                fold_probs = []\n                                for fold_models in trained:\n                                    fold_model_probs = [m.predict_proba(X_te_np)[:, 1] for m in fold_models]\n                                    fold_probs.append(np.mean(fold_model_probs, axis=0))\n                                # Average across folds\n                                pred[action] = np.mean(fold_probs, axis=0)\n                            else:\n                                # Single-split mode: unwrap the single fold and average across models\n                                fold_models = trained[0]\n                                probs = [m.predict_proba(X_te_np)[:, 1] for m in fold_models]\n                                pred[action] = np.mean(probs, axis=0)\n\n                    del X_te_np\n                    gc.collect()\n\n                    if pred.shape[1] != 0:\n                        sub_part = predict_multiclass_adaptive(pred, meta_te, action_thresholds)\n                        submission_list.append(sub_part)\n                    else:\n                        if verbose:\n                            print(\"  ERROR: no training data\")\n\n                except Exception as e:\n                    if verbose:\n                        print(f\"  ERROR: {str(e)[:50]}\")\n                    try:\n                        del data_te\n                    except:\n                        pass\n                    gc.collect()\n\n        except Exception as e:\n            print(f'***Exception*** {str(e)[:100]}')\n            import traceback\n            if verbose:\n                traceback.print_exc()\n\n        gc.collect()\n        print()\n\n    # Combine all predictions\n    if len(submission_list) > 0:\n        submission = pd.concat(submission_list, ignore_index=True)\n    else:\n        # Fallback: create minimal submission\n        submission = pd.DataFrame({\n            'video_id': [test.video_id.iloc[0] if len(test) > 0 else 438887472],\n            'agent_id': ['mouse1'],\n            'target_id': ['self'],\n            'action': ['rear'],\n            'start_frame': [278],\n            'stop_frame': [500]\n        })\n\n    # Apply robustify post-processing\n    submission = robustify(submission, test, 'test')\n    \n    print(f\"\\nGenerated {len(submission)} predictions\")\n    print(f\"Unique videos: {submission.video_id.nunique()}\")\n    print(f\"Unique actions: {submission.action.nunique()}\")\n    \nelse:\n    print(\"Skipping inference (validate_or_submit='validate')\")\n    submission = pd.DataFrame()\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 12. Submission Generation\n\nFormat and save the final submission file.\n","metadata":{}},{"cell_type":"code","source":"if validate_or_submit == 'submit' and len(submission) > 0:\n    # Format submission for Kaggle (add row_id, ensure proper types and order)\n    submission['row_id'] = range(len(submission))\n    \n    # Ensure proper column order (Kaggle format)\n    required_columns = ['row_id', 'video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']\n    submission = submission[required_columns]\n    \n    # Cast numeric columns to int (required for Kaggle)\n    submission['row_id'] = submission['row_id'].astype(int)\n    submission['video_id'] = submission['video_id'].astype(int)\n    submission['start_frame'] = submission['start_frame'].astype(int)\n    submission['stop_frame'] = submission['stop_frame'].astype(int)\n    \n    # Sort by video_id and start_frame for consistency\n    submission = submission.sort_values(['video_id', 'start_frame']).reset_index(drop=True)\n    \n    # Reset row_id to be sequential after sorting\n    submission['row_id'] = range(len(submission))\n    \n    # Save submission\n    output_path = 'submission.csv'  # Always save to current directory for Kaggle compatibility\n    submission.to_csv(output_path, index=False)\n    \n    print(f\"\\nSubmission saved to: {output_path}\")\n    print(f\"\\nSubmission preview:\")\n    print(submission.head(10))\n    print(f\"\\nSubmission statistics:\")\n    print(f\"  Total predictions: {len(submission)}\")\n    print(f\"  Unique videos: {submission.video_id.nunique()}\")\n    print(f\"  Unique actions: {submission.action.nunique()}\")\n    print(f\"  Actions distribution:\")\n    print(submission.action.value_counts().head(10))\n    print(f\"\\nSubmission format validation:\")\n    print(f\"  Columns: {list(submission.columns)}\")\n    print(f\"  Data types:\")\n    print(submission.dtypes)\nelse:\n    print(\"No submission generated (validate mode or empty submission)\")\n","metadata":{},"outputs":[],"execution_count":null}]}