{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":13777770,"sourceType":"datasetVersion","datasetId":8667881}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"validate_or_submit = '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 gc\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.impute import SimpleImputer\nfrom sklearn.preprocessing import StandardScaler\nfrom collections import defaultdict\nimport polars as pl\nfrom scipy import signal\nimport math\nimport joblib\nimport glob\n\nwarnings.filterwarnings('ignore')\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\nTREE_MODELS_DIR = \"/kaggle/input/trained-action-2/saved_models\" \n\n# ======================== LSTM MODEL ========================\n\nclass LSTMDataset(Dataset):\n    def __init__(self, X, y, seq_len=30):\n        self.X = torch.FloatTensor(X)\n        self.y = torch.FloatTensor(y)\n        self.seq_len = seq_len\n\n    def __len__(self):\n        return max(1, len(self.X) - self.seq_len + 1)\n\n    def __getitem__(self, idx):\n        x_seq = self.X[idx:idx + self.seq_len]\n        y_val = self.y[idx + self.seq_len - 1]\n        return x_seq, y_val  # y_val is (1,)\n\nclass PositionalEncoding(nn.Module):\n    def __init__(self, d_model, dropout=0.1, max_len=5000):\n        super(PositionalEncoding, self).__init__()\n        self.dropout = nn.Dropout(p=dropout)\n\n        pe = torch.zeros(max_len, d_model)\n        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n        pe[:, 0::2] = torch.sin(position * div_term)\n        pe[:, 1::2] = torch.cos(position * div_term)\n        pe = pe.unsqueeze(0).transpose(0, 1)\n        self.register_buffer('pe', pe)\n\n    def forward(self, x):\n        x = x + self.pe[:x.size(0), :]\n        return self.dropout(x)\n\nclass LSTMClassifier(nn.Module):\n    def __init__(self, input_size, hidden_size=64, num_layers=2, dropout=0.3):\n        super().__init__()\n        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, \n                              batch_first=True, \n                              dropout=dropout if num_layers > 1 else 0,\n                              bidirectional=True)\n    \n        self.fc = nn.Linear(hidden_size * 2, 1)\n\n    def forward(self, x):\n        lstm_out, _ = self.lstm(x)\n        last_time_step_out = lstm_out[:, -1, :]\n        \n        out = self.fc(last_time_step_out)\n        return out\n\ndef train_lstm(X_tr, y_tr, input_size, epochs=1, batch_size=512, lr=0.001, save_path= None):\n    if len(X_tr) < 10:  # Too small\n        def predict_dummy(X_te):\n            return np.full(len(X_te), 0.5)\n        return predict_dummy, 1\n\n    seq_len_train = min(30, max(5, len(X_tr) // 4))  # Min 5 for small\n    \n    # FIX: Removed .values from y_tr because it's already a numpy array\n    dataset = LSTMDataset(X_tr, y_tr.reshape(-1, 1), seq_len_train)\n    \n    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0)\n\n    model = LSTMClassifier(input_size).to(device)\n    optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n    \n    # criterion = nn.BCELoss() not_perfect_define_values\n    \n    positives = y_tr.sum()\n    negatives = len(y_tr) - positives\n    pos_weight = torch.tensor([negatives / positives], device=device) if positives > 0 else torch.tensor([1.0], device=device)\n\n    # BCEWithLogitsLoss zyada stable hota hai aur pos_weight handle karta hai\n    criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)\n\n    model.train()\n    for epoch in range(epochs):\n        total_loss = 0\n        for seq_x, seq_y in dataloader:\n            seq_x = seq_x.to(device)\n            seq_y = seq_y.to(device)  # (batch, 1)\n            optimizer.zero_grad()\n            out = model(seq_x)  # (batch, 1) - NO SQUEEZE!\n            loss = criterion(out, seq_y)\n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # Stable\n            optimizer.step()\n            total_loss += loss.item()\n        if verbose and epoch % 5 == 0:\n            print(f\"Epoch {epoch}, loss: {total_loss / len(dataloader):.4f}\")\n\n    if save_path:\n        os.makedirs(os.path.dirname(save_path), exist_ok=True)\n        torch.save(model.state_dict(), save_path)\n        if verbose:\n            print(f\"✅ Model and Scaler saved successfully at: {save_path}\")\n\n    # The rest of the function remains the same...\n    def predict(X_te):\n        if len(X_te) == 0:\n            return np.array([])\n        model.eval()\n        with torch.no_grad():\n            seq_len_test = min(seq_len_train, len(X_te))\n            if seq_len_test < 1:\n                seq_len_test = 1\n            step = max(1, seq_len_test // 4)  # Overlap for better coverage\n            indices = np.arange(0, len(X_te) - seq_len_test + 1, step)\n            if len(indices) == 0:\n                # Very small, use full as seq\n                seq = X_te[:seq_len_test] if len(X_te) > 0 else X_te\n                if len(seq) < seq_len_train:\n                    pad_len = seq_len_train - len(seq)\n                    pad = np.tile(seq[-1:] if len(seq)>0 else np.zeros((1, input_size)), (pad_len, 1))\n                    seq = np.vstack([seq, pad])\n                seq_t = torch.FloatTensor(seq).unsqueeze(0).to(device)\n                prob = model(seq_t).cpu().numpy()[0, 0]\n                return np.full(len(X_te), prob)\n\n            batch_size_pred = 128\n            all_probs = []\n            for b_start in range(0, len(indices), batch_size_pred):\n                b_end = min(b_start + batch_size_pred, len(indices))\n                b_inds = indices[b_start:b_end]\n                batch_seqs = []\n                for j in b_inds:\n                    seq_slice = X_te[j:j + seq_len_test]\n                    if len(seq_slice) < seq_len_test:\n                        pad_len = seq_len_test - len(seq_slice)\n                        pad = np.tile(seq_slice[-1:], (pad_len, 1)) if len(seq_slice) > 0 else np.zeros((pad_len, input_size))\n                        seq_slice = np.vstack([seq_slice, pad])\n                    batch_seqs.append(seq_slice)\n                if batch_seqs:\n                    batch_t = torch.FloatTensor(np.stack(batch_seqs)).to(device)\n                    b_out = model(batch_t).cpu().numpy()[:, 0]\n                    all_probs.extend(b_out)\n\n            # Assign probs to frame positions (last of each window)\n            frame_probs = np.full(len(X_te), 0.5)\n            for k, j in enumerate(indices):\n                assign_frame = min(j + seq_len_test - 1, len(X_te) - 1)\n                frame_probs[assign_frame] = all_probs[k]\n            # Simple forward fill for gaps\n            frame_probs = pd.Series(frame_probs).fillna(method='ffill').fillna(0.5).values\n            return frame_probs\n\n    return predict, seq_len_train\n\n# ======================== SCORING FUNCTIONS (UNCHANGED) ========================\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# ======================== DATA LOADING (MINOR FIXES) ========================\n\ntrain = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/train.csv')\ntrain['n_mice'] = 4 - train[['mouse1_strain', 'mouse2_strain', 'mouse3_strain', 'mouse4_strain']].isna().sum(axis=1)\ntrain_without_mabe22 = train.query(\"~ lab_id.str.startswith('MABe22_')\")\n\ntest = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/test.csv')\nbody_parts_tracked_list = list(np.unique(train.body_parts_tracked))\n\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    assert traintest in ['train', 'test']\n    if traintest_directory is None:\n        traintest_directory = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking\"\n    for _, row in dataset.iterrows():\n        lab_id = row.lab_id\n        if lab_id.startswith('MABe22'): continue\n        video_id = row.video_id\n\n        if type(row.behaviors_labeled) != str:\n            if verbose: print('No labeled behaviors:', lab_id, video_id)\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        if pvid.isna().any().any():\n            if verbose and traintest == 'test': print('video with missing values', video_id, traintest, len(vid), 'frames')\n            pvid = pvid.fillna(method='ffill').fillna(pvid.mean())  # Impute\n        else:\n            if verbose and traintest == 'test': 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        if traintest == 'train':\n            try:\n                annot = pd.read_parquet(path.replace('train_tracking', 'train_annotation'))\n            except FileNotFoundError:\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                    })\n                    if traintest == 'train':\n                        single_mouse_label = pd.DataFrame(0.0, columns=vid_agent_actions, index=single_mouse.index)\n                        annot_subset = annot.query(\"(agent_id == @mouse_id) & (target_id == @mouse_id)\")\n                        for i in range(len(annot_subset)):\n                            annot_row = annot_subset.iloc[i]\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: 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                    })\n                    if traintest == 'train':\n                        mouse_pair_label = pd.DataFrame(0.0, columns=vid_agent_actions, index=mouse_pair.index)\n                        annot_subset = annot.query(\"(agent_id == @agent) & (target_id == @target)\")\n                        for i in range(len(annot_subset)):\n                            annot_row = annot_subset.iloc[i]\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: print('- test pair', video_id, agent, target)\n                        yield 'pair', mouse_pair, mouse_pair_meta, vid_agent_actions\n\n\n# ======================== ADVANCED FEATURE ENGINEERING FOR (TREES)========================\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 add_curvature_features(X, center_x, center_y):\n    \"\"\"Trajectory curvature\"\"\"\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 window in [30, 60]:\n        X[f'curv_mean_{window}'] = curvature.rolling(window, min_periods=5).mean()\n\n    angle = np.arctan2(vel_y, vel_x)\n    angle_change = np.abs(angle.diff())\n    X['turn_rate_30'] = angle_change.rolling(30, min_periods=5).sum()\n\n    return X\n\ndef add_multiscale_features(X, center_x, center_y):\n    \"\"\"Multi-scale temporal features\"\"\"\n    speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2)\n\n    scales = [10, 40, 160]\n    for scale in scales:\n        if len(speed) >= scale:\n            X[f'sp_m{scale}'] = speed.rolling(scale, min_periods=max(1, scale//4)).mean()\n            X[f'sp_s{scale}'] = speed.rolling(scale, min_periods=max(1, scale//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):\n    \"\"\"Behavioral state transitions\"\"\"\n    speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2)\n    speed_ma = speed.rolling(15, min_periods=5).mean()\n\n    try:\n        speed_states = pd.cut(speed_ma, bins=[-np.inf, 0.5, 2.0, 5.0, np.inf], labels=[0, 1, 2, 3]).astype(float)\n        \n        for window in [60, 120]:\n            if len(speed_states) >= window:\n                for state in [0, 1, 2, 3]:\n                    X[f's{state}_{window}'] = (speed_states == state).astype(float).rolling(window, min_periods=10).mean()\n                \n                state_changes = (speed_states != speed_states.shift(1)).astype(float)\n                X[f'trans_{window}'] = state_changes.rolling(window, min_periods=10).sum()\n    except:\n        pass\n\n    return X\n\ndef add_longrange_features(X, center_x, center_y):\n    \"\"\"Long-range temporal features\"\"\"\n    for window in [120, 240]:\n        if len(center_x) >= window:\n            X[f'x_ml{window}'] = center_x.rolling(window, min_periods=20).mean()\n            X[f'y_ml{window}'] = center_y.rolling(window, min_periods=20).mean()\n\n    for span in [60, 120]:\n        X[f'x_e{span}'] = center_x.ewm(span=span, min_periods=1).mean()\n        X[f'y_e{span}'] = center_y.ewm(span=span, min_periods=1).mean()\n\n    speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2)\n    for window in [60, 120]:\n        if len(speed) >= window:\n            X[f'sp_pct{window}'] = speed.rolling(window, min_periods=20).rank(pct=True)\n\n    return X\n\ndef add_interaction_features(X, mouse_pair, avail_A, avail_B):\n    \"\"\"Social interaction features\"\"\"\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        X[f'A_ld{window}'] = A_lead.rolling(window, min_periods=5).mean()\n        X[f'B_ld{window}'] = B_lead.rolling(window, min_periods=5).mean()\n\n    approach = -rel_dist.diff()\n    chase = approach * B_lead\n    X['chase_30'] = chase.rolling(30, min_periods=5).mean()\n\n    for window in [60, 120]:\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(window, min_periods=10).corr(B_sp)\n\n    return X\n\ndef transform_single(single_mouse, body_parts_tracked):\n    \"\"\"Enhanced single mouse transform\"\"\"\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 features\n    if all(p in single_mouse.columns for p in ['ear_left', 'ear_right', 'tail_base']):\n        shifted = single_mouse[['ear_left', 'ear_right', 'tail_base']].shift(10)\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            X[f'cx_m{w}'] = cx.rolling(w, min_periods=1, center=True).mean()\n            X[f'cy_m{w}'] = cy.rolling(w, min_periods=1, center=True).mean()\n            X[f'cx_s{w}'] = cx.rolling(w, min_periods=1, center=True).std()\n            X[f'cy_s{w}'] = cy.rolling(w, min_periods=1, center=True).std()\n            X[f'x_rng{w}'] = cx.rolling(w, min_periods=1, center=True).max() - cx.rolling(w, min_periods=1, center=True).min()\n            X[f'y_rng{w}'] = cy.rolling(w, min_periods=1, center=True).max() - cy.rolling(w, min_periods=1, center=True).min()\n            X[f'disp{w}'] = np.sqrt(cx.diff().rolling(w, min_periods=1).sum()**2 + cy.diff().rolling(w, min_periods=1).sum()**2)\n            X[f'act{w}'] = np.sqrt(cx.diff().rolling(w, min_periods=1).var() + cy.diff().rolling(w, min_periods=1).var())\n        \n        # Advanced features\n        X = add_curvature_features(X, cx, cy)\n        X = add_multiscale_features(X, cx, cy)\n        X = add_state_features(X, cx, cy)\n        X = add_longrange_features(X, cx, cy)\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            X[f'nt_lg{lag}'] = nt_dist.shift(lag)\n            X[f'nt_df{lag}'] = nt_dist - nt_dist.shift(lag)\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            X[f'ear_o{off}'] = ear_d.shift(-off)\n        X['ear_con'] = ear_d.rolling(30, min_periods=1, center=True).std() / (ear_d.rolling(30, min_periods=1, center=True).mean() + 1e-6)\n\n    return X\n\ndef transform_pair(mouse_pair, body_parts_tracked):\n    \"\"\"Enhanced pair transform\"\"\"\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 features\n    if ('A', 'ear_left') in mouse_pair.columns and ('B', 'ear_left') in mouse_pair.columns:\n        shA = mouse_pair['A']['ear_left'].shift(10)\n        shB = mouse_pair['B']['ear_left'].shift(10)\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\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        shA_n = mouse_pair['A']['nose'].shift(10)\n        shB_n = mouse_pair['B']['nose'].shift(10)\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\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            X[f'd_m{w}'] = cd_full.rolling(w, min_periods=1, center=True).mean()\n            X[f'd_s{w}'] = cd_full.rolling(w, min_periods=1, center=True).std()\n            X[f'd_mn{w}'] = cd_full.rolling(w, min_periods=1, center=True).min()\n            X[f'd_mx{w}'] = cd_full.rolling(w, min_periods=1, center=True).max()\n            \n            d_var = cd_full.rolling(w, min_periods=1, center=True).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(w, min_periods=1, center=True).mean()\n            X[f'co_s{w}'] = coord.rolling(w, min_periods=1, center=True).std()\n\n    # Nose-nose\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            X[f'nn_lg{lag}'] = nn.shift(lag)\n            X[f'nn_ch{lag}'] = nn - nn.shift(lag)\n            is_cl = (nn < 10.0).astype(float)\n            X[f'cl_ps{lag}'] = is_cl.rolling(lag, 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            X[f'va_{off}'] = val.shift(-off)\n        \n        X['int_con'] = cd_full.rolling(30, min_periods=1, center=True).std() / (cd_full.rolling(30, min_periods=1, center=True).mean() + 1e-6)\n        \n        # Advanced interaction\n        X = add_interaction_features(X, mouse_pair, avail_A, avail_B)\n\n    return X\n\n\n# ======================== FEATURE ENGINEERING (IMPROVED) FOR (LSTM)========================\ndef get_or_create_body_center(mouse_data):\n    \"\"\"\n    Yeh function check karta hai ki 'body_center' hai ya nahi.\n    Agar hai, to use return karta hai.\n    Agar nahi hai, to sabhi available body parts ka average lekar ek naya 'body_center' banata hai.\n    \"\"\"\n    available_body_parts = mouse_data.columns.get_level_values(0).unique()\n    \n    if 'body_center' in available_body_parts:\n        # Agar body_center pehle se hai, to use hi use karo\n        return mouse_data['body_center']\n    else:\n        # Agar nahi hai, to ek naya banao\n        # Saare x aur y coordinates ko alag karo\n        x_coords = mouse_data.loc[:, (slice(None), 'x')]\n        y_coords = mouse_data.loc[:, (slice(None), 'y')]\n        \n        # Unka average le lo\n        center_x = x_coords.mean(axis=1)\n        center_y = y_coords.mean(axis=1)\n        \n        # Ek naya DataFrame banakar return karo\n        return pd.DataFrame({'x': center_x, 'y': center_y})\n\ndef get_orientation(mouse_data):\n    if 'nose' in mouse_data and 'tail_base' in mouse_data:\n        dx = mouse_data['nose']['x'] - mouse_data['tail_base']['x']\n        dy = mouse_data['nose']['y'] - mouse_data['tail_base']['y']\n        return np.arctan2(dy, dx)\n    return pd.Series(0, index=mouse_data.index) # Agar data nahi hai\n                        \n# ======================== FEATURE ENGINEERING (SIMPLIFIED FOR LSTM) ========================\ndef transform_single_v2(single_mouse, row):\n    X = pd.DataFrame(index=single_mouse.index)\n    body_center = get_or_create_body_center(single_mouse)\n    available_parts = single_mouse.columns.get_level_values(0).unique()\n\n    # === Basic Kinematics (Aapke paas pehle se hain) ===\n    X['speed'] = np.sqrt(body_center['x'].diff()**2 + body_center['y'].diff()**2).fillna(0)\n    X['accel'] = X['speed'].diff().fillna(0)\n    X['jerk'] = X['accel'].diff().fillna(0) # Jerk for sudden movements\n\n    # === Posture & Position Features ===\n    # Feature #1: Rearing (खड़ा होना)\n    X['y_pos'] = body_center['y'] # Y-coordinate\n    X['y_vel'] = body_center['y'].diff().fillna(0) # Vertical velocity\n\n    # Feature #6: Body Orientation\n    if 'nose' in available_parts and 'tail_base' in available_parts:\n        dx = single_mouse['nose']['x'] - single_mouse['tail_base']['x']\n        dy = single_mouse['nose']['y'] - single_mouse['tail_base']['y']\n        X['orientation'] = np.arctan2(dy, dx)\n        X['orientation_change'] = X['orientation'].diff().fillna(0)\n        X['body_length'] = np.sqrt(dx**2 + dy**2) # For elongation\n\n    # Feature #4: Grooming / Scratching\n    if 'nose' in available_parts and 'forepaw_left' in available_parts and 'forepaw_right' in available_parts:\n        # Snout-to-forepaws distance\n        dist_left = np.sqrt((single_mouse['nose']['x'] - single_mouse['forepaw_left']['x'])**2 + (single_mouse['nose']['y'] - single_mouse['forepaw_left']['y'])**2)\n        dist_right = np.sqrt((single_mouse['nose']['x'] - single_mouse['forepaw_right']['x'])**2 + (single_mouse['nose']['y'] - single_mouse['forepaw_right']['y'])**2)\n        X['snout_to_paw_dist'] = np.minimum(dist_left, dist_right)\n        \n        # Speed of forepaws for rhythmic movement\n        X['paw_speed_left'] = np.sqrt(single_mouse['forepaw_left']['x'].diff()**2 + single_mouse['forepaw_left']['y'].diff()**2).fillna(0)\n\n    arena_width = 100  \n    arena_height = 100\n    \n    if 'arena_width_cm' in row:\n       arena_width = row['arena_width_cm']\n    if 'arena_height_cm' in row:\n       arena_height = row['arena_height_cm']\n    \n    # Ab in safe variables ka istemaal karo\n    dist_x = np.minimum(body_center['x'], arena_width - body_center['x'])\n    dist_y = np.minimum(body_center['y'], arena_height - body_center['y'])\n    X['dist_from_wall'] = np.minimum(dist_x, dist_y)\n    X['dist_from_wall'] = np.minimum(dist_x, dist_y)\n\n    return X\n\nfrom scipy.spatial.distance import cdist\ndef transform_pair_v2(mouse_pair):\n    X = pd.DataFrame(index=mouse_pair.index)\n    body_center_A = get_or_create_body_center(mouse_pair['A'])\n    body_center_B = get_or_create_body_center(mouse_pair['B'])\n    available_parts_A = mouse_pair['A'].columns.get_level_values(0).unique()\n    available_parts_B = mouse_pair['B'].columns.get_level_values(0).unique()\n\n    # === Basic Relative Kinematics (Aapke paas pehle se hain) ===\n    rel_x = body_center_A['x'] - body_center_B['x']\n    rel_y = body_center_A['y'] - body_center_B['y']\n    X['center_dist'] = np.sqrt(rel_x**2 + rel_y**2)\n\n    # NEW: Velocity (Speed)\n    X['rel_vel_x'] = rel_x.diff().fillna(0)\n    X['rel_vel_y'] = rel_y.diff().fillna(0)\n    X['dist_vel'] = X['center_dist'].diff().fillna(0) # Are they getting closer or farther?\n\n    # NEW: Angle (Magic feature for chasing/sniffing)\n    X['rel_angle'] = np.arctan2(rel_y, rel_x)\n    X['rel_angle_vel'] = X['rel_angle'].diff().fillna(0)\n\n    # NEW: Individual Speeds\n    X['A_speed'] = np.sqrt((body_center_A.diff()**2).sum(axis=1)).fillna(0)\n    X['B_speed'] = np.sqrt((body_center_B.diff()**2).sum(axis=1)).fillna(0)\n\n    # transform_pair_v2 mein:\n    X['A_orientation'] = get_orientation(mouse_pair['A'])\n    X['B_orientation'] = get_orientation(mouse_pair['B'])\n    \n    # Feature for Approach / Disengaged\n    X['dist_change_rate'] = X['center_dist'].diff().fillna(0)\n\n    speed_A = np.sqrt(body_center_A['x'].diff()**2 + body_center_A['y'].diff()**2).fillna(0)\n    speed_B = np.sqrt(body_center_B['x'].diff()**2 + body_center_B['y'].diff()**2).fillna(0)\n    X['speed_diff'] = speed_A - speed_B\n    X['speed_ratio'] = speed_A / (speed_B + 1e-6)\n\n    # === Behavior-Specific Features ===\n    # Feature for Attack / Push\n    X['rel_accel'] = X['speed_diff'].diff().fillna(0) # Relative acceleration\n\n    # Feature for Sniffing (Face & Anogenital)\n    if 'nose' in available_parts_A and 'nose' in available_parts_B:\n        X['nose_dist'] = np.sqrt((mouse_pair['A']['nose']['x'] - mouse_pair['B']['nose']['x'])**2 + (mouse_pair['A']['nose']['y'] - mouse_pair['B']['nose']['y'])**2)\n    \n    if 'nose' in available_parts_A and 'tail_base' in available_parts_B:\n        X['chase_dist'] = np.sqrt((mouse_pair['A']['nose']['x'] - mouse_pair['B']['tail_base']['x'])**2 + (mouse_pair['A']['nose']['y'] - mouse_pair['B']['tail_base']['y'])**2)\n\n    # Feature for Mounting\n    X['y_pos_diff'] = body_center_A['y'] - body_center_B['y'] # Vertical difference\n\n    if 'nose' in mouse_pair['A'].columns.get_level_values(0) and 'nose' in mouse_pair['B'].columns.get_level_values(0):\n       nose_dist_x = mouse_pair['A']['nose']['x'] - mouse_pair['B']['nose']['x']\n       nose_dist_y = mouse_pair['A']['nose']['y'] - mouse_pair['B']['nose']['y']\n       X['nose_to_nose_dist'] = np.sqrt(nose_dist_x**2 + nose_dist_y**2)\n\n    # Feature for Allogroom\n    if 'nose' in available_parts_A:\n      X['allogroom_dist'] = np.sqrt(\n        (mouse_pair['A']['nose']['x'] - body_center_B['x'])**2 + \n        (mouse_pair['A']['nose']['y'] - body_center_B['y'])**2\n    )\n     \n    return X\n    \naction_thresholds = defaultdict(lambda: 0.30)\n\n# ======================== NEW, IMPROVED PREDICTION-TO-EVENT FUNCTION ========================\ndef events_from_probabilities(probabilities, meta, action_thresholds, min_duration=5, merge_gap=3):\n    \"\"\"\n    Ek behtar tareeka probabilities se action events nikalne ka.\n    Yeh noise ko handle karta hai aur har action ko alag se process karta hai.\n    \"\"\"\n    submission_parts = []\n    \n    # Har action column ke liye alag se process karo\n    for action in probabilities.columns:\n        # 1. Threshold lagao\n        threshold = action_thresholds.get(action, 0.5)\n        preds = (probabilities[action].values > threshold).astype(int)\n        \n        if preds.sum() == 0: # Agar koi bhi frame threshold se upar nahi hai, to skip\n            continue\n            \n        # 2. Chote gaps (jaise 0,0,1,1,0,1,1,0,0) ko fill karo\n        # Isse toote hue events jud jaate hain\n        change_points = np.diff(preds, prepend=0, append=0)\n        start_indices = np.where(change_points == 1)[0]\n        stop_indices = np.where(change_points == -1)[0]\n        \n        for i in range(len(start_indices) - 1):\n            gap = start_indices[i+1] - stop_indices[i]\n            if gap <= merge_gap:\n                preds[stop_indices[i] : start_indices[i+1]] = 1 # Gap ko 1 se bhar do\n        \n        # 3. Final events (start aur stop frames) nikalo\n        change_points = np.diff(preds, prepend=0, append=0)\n        start_indices = np.where(change_points == 1)[0]\n        stop_indices = np.where(change_points == -1)[0]\n        \n        # 4. Chote events (jo min_duration se kam hain) ko filter karo\n        for start_idx, stop_idx in zip(start_indices, stop_indices):\n            duration = stop_idx - start_idx\n            if duration >= min_duration:\n                # Meta-data se video_id, agent_id, etc. lo\n                # Humein actual frame numbers chahiye, index nahi\n                start_frame = meta['video_frame'].iloc[start_idx]\n                # stop_idx exclusive hota hai, isliye -1\n                stop_frame = meta['video_frame'].iloc[stop_idx - 1] + 1 \n                \n                submission_parts.append({\n                    'video_id': meta['video_id'].iloc[0],\n                    'agent_id': meta['agent_id'].iloc[0],\n                    'target_id': meta['target_id'].iloc[0],\n                    'action': action,\n                    'start_frame': start_frame,\n                    'stop_frame': stop_frame\n                })\n\n    if not submission_parts:\n        return pd.DataFrame()\n    else:\n        return pd.DataFrame(submission_parts)\n\n\n# --- STEP 2: ADD THIS FUNCTION BEFORE 'submit_resnet_pipeline' ---\ndef predict_with_saved_trees(X_test, config_id, switch, action):\n    \"\"\"\n    Saved Tree models (0-4) ko load karta hai aur Average Prediction deta hai.\n    Features ko auto-align karta hai taaki error na aaye.\n    \"\"\"\n    all_preds = []\n    models_found = 0\n    \n    # Hum 5 folds (model0 se model4) check karenge\n    for i in range(5):\n        # Filename pattern tumhare bataye anusar:\n        filename = f\"model_config{config_id}_{switch}_{action}_model{i}.joblib\"\n        path = os.path.join(TREE_MODELS_DIR, filename)\n        \n        if os.path.exists(path):\n            try:\n                model = joblib.load(path)\n                \n                # --- Feature Alignment Logic (Crash rokne ke liye) ---\n                # Tree models ko wahi columns chahiye hote hain jinpar wo train huye.\n                # Hum check karte hain agar model me feature names saved hain.\n                try:\n                    if hasattr(model, 'feature_name_'):\n                        required_feats = model.feature_name_\n                    elif hasattr(model, 'get_booster'): # XGBoost\n                        required_feats = model.get_booster().feature_names\n                    else:\n                        required_feats = None\n\n                    X_aligned = X_test.copy()\n                    if required_feats:\n                        # Missing columns ko 0 se bhar do\n                        for col in required_feats:\n                            if col not in X_aligned.columns:\n                                X_aligned[col] = 0.0\n                        # Extra columns hata do aur order same karo\n                        X_aligned = X_aligned[required_feats]\n                except:\n                    X_aligned = X_test # Agar pata na chale to direct try karo\n                \n                # Prediction lo\n                if hasattr(model, \"predict_proba\"):\n                    probs = model.predict_proba(X_aligned)[:, 1]\n                else:\n                    probs = model.predict(X_aligned)\n                    \n                all_preds.append(probs)\n                models_found += 1\n            except Exception as e:\n                pass # Agar koi model corrupt hai to skip karo\n\n    if models_found == 0:\n        return None\n    \n    # Saare tree models ka average\n    return np.mean(all_preds, axis=0)\n\n# ======================== LSTM TRAINING & PREDICTION ========================\n\nMODELS_DIR = \"saved_models\"\n\ndef submit_lstm(config_id,body_parts_tracked_str, switch_tr, X_tr, label, meta):\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    if X_tr.shape[1] == 0:\n        print(\"No features, skip\")\n        return\n\n    assert len(X_tr) == len(label), f\"FATAL ERROR: Mismatch! X_tr length {len(X_tr)} != label length {len(label)}\"\n        \n    input_size = X_tr.shape[1]\n    scaler = StandardScaler()\n    X_tr_scaled = scaler.fit_transform(X_tr)\n\n    model_list = []\n    for action in label.columns:\n        action_mask = ~label[action].isna()\n        if action_mask.sum() >= 5 and label[action][action_mask].sum() > 0:  # Has positives\n            y_action = label[action][action_mask].values.astype(np.float32)\n            X_action = X_tr_scaled[action_mask.values]\n            model_filename = f\"model_config_{config_id}_{switch_tr}_{action}.pth\"\n            model_path = os.path.join(MODELS_DIR, model_filename)\n            predict_fn, _ = train_lstm(X_action, y_action, input_size, save_path=model_path)#,\n            model_list.append((action, predict_fn, scaler))\n            \n            scaler_path = f'scaler_config_{config_id}_{switch_tr}_{action}.joblib'\n            joblib.dump(scaler, scaler_path)\n\n    del X_tr, label, X_tr_scaled\n    gc.collect()\n\n    test_subset = test[test.body_parts_tracked == body_parts_tracked_str]\n    generator = generate_mouse_data(test_subset, 'test', generate_single=(switch_tr == 'single'), generate_pair=(switch_tr == 'pair'))\n\n    if verbose: print(f\"n_videos: {len(test_subset)}, models: {len(model_list)}\")\n\n    for switch_te, data_te, meta_te, actions_te in generator:\n        if switch_te != switch_tr: continue\n        try:\n            \n            # 1. Generate Features for LSTM (using v2)\n            if switch_te == 'single':\n                X_te_lstm = transform_single_v2(data_te, body_parts_tracked)\n                X_te_tree = transform_single(data_te, body_parts_tracked)\n                \n            else:\n                X_te_lstm = transform_pair_v2(data_te)\n                X_te_tree = transform_pair(data_te, body_parts_tracked)\n            \n            del data_te\n            gc.collect()\n\n            pred_df = pd.DataFrame(index=meta_te['video_frame'].values)  # Use values to avoid index issues\n            for action, predict_fn, scaler_te in model_list:\n                if action in actions_te:\n                    X_te_scaled = scaler_te.transform(X_te_lstm)\n                    probs_lstm = predict_fn(X_te_scaled)\n                    probs_tree = predict_with_saved_trees(X_te_tree, config_id, switch_te, action)\n                    if probs_tree is not None:\n                        min_len = min(len(probs_lstm), len(probs_tree))\n                        final_probs = (0.3 * probs_lstm[:min_len]) + (0.7 * probs_tree[:min_len])\n                        if len(final_probs) < len(probs_lstm):\n                            pad = np.zeros(len(probs_lstm) - len(final_probs))\n                            final_probs = np.concatenate([final_probs, pad])\n                    else:\n                        final_probs = probs_lstm \n                        \n                    pred_df[action] = final_probs\n            \n            if not pred_df.empty and pred_df.shape[1] > 0:\n                 pred_df_smooth = pred_df.rolling(window=5, min_periods=1, center=True).mean()\n                 sub_part = events_from_probabilities(pred_df_smooth, meta_te, action_thresholds, min_duration=5, merge_gap=3)\n            \n                 submission_list.append(sub_part)\n            else:\n               if verbose: \n                   print(f\"  No models/actions for {switch_te}\")\n                    \n        except Exception as e:\n            if verbose: print(f'  ERROR in {switch_te}: {str(e)[:100]}')\n            gc.collect()\n            \n# ======================== ROBUSTIFY (UNCHANGED) ========================\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    # No overlaps\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 pd.DataFrame()\n\n    # Fill empty\n    s_list = []\n    for idx, row in dataset.iterrows():\n        lab_id = row['lab_id']\n        if lab_id.startswith('MABe22'):\n            continue\n        video_id = row['video_id']\n        if submission['video_id'].eq(video_id).any():\n            continue\n\n        if verbose: 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 = 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        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            if len(actions) == 0: continue\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 s_list:\n        dummy_df = pd.DataFrame(s_list, columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame'])\n        submission = pd.concat([submission, dummy_df], ignore_index=True)\n\n    submission = submission.reset_index(drop=True)\n    return submission\n\n# ======================== MAIN LOOP ========================\n\nsubmission_list = []\n\nfor section in range(1, len(body_parts_tracked_list)):\n    body_parts_tracked_str = body_parts_tracked_list[section]\n    try:\n        print(f\"{section}. Processing: {body_parts_tracked_str}\")\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        train_subset = train[train.body_parts_tracked == body_parts_tracked_str]\n        all_data_for_config = list(generate_mouse_data(train_subset, 'train'))\n\n        # <<<<<<<<< STEP 2: Data ko type ke hisaab se alag-alag karo >>>>>>>>>\n        single_data_list = [(d, m, l) for sw, d, m, l in all_data_for_config if sw == 'single' and not d.empty and not l.empty]\n        pair_data_list = [(d, m, l) for sw, d, m, l in all_data_for_config if sw == 'pair' and not d.empty and not l.empty]\n\n        if single_data_list:\n            single_list = [item[0] for item in single_data_list]\n            single_meta_list = [item[1] for item in single_data_list]\n            single_label_list = [item[2] for item in single_data_list]\n\n            single_mouse = pd.concat(single_list)\n            single_meta = pd.concat(single_meta_list)\n            single_label = pd.concat(single_label_list)\n\n            # Ab reset_index karo\n            single_mouse.reset_index(drop=True, inplace=True)\n            single_meta.reset_index(drop=True, inplace=True)\n            single_label.reset_index(drop=True, inplace=True)\n            \n            del single_list, single_label_list, single_meta_list\n            gc.collect()\n            \n            X_tr_single = transform_single_v2(single_mouse, body_parts_tracked)\n            \n            del single_mouse\n            print(f\"  Single: {X_tr_single.shape}\")\n            submit_lstm(section,body_parts_tracked_str, 'single', X_tr_single, single_label, single_meta)\n                \n        if pair_data_list:\n           pair_list = [item[0] for item in pair_data_list]\n           pair_meta_list = [item[1] for item in pair_data_list]\n           pair_label_list = [item[2] for item in pair_data_list]\n\n           # Concat aab bilkul aaram se kaam karega\n           mouse_pair = pd.concat(pair_list)\n           pair_meta = pd.concat(pair_meta_list)\n           pair_label = pd.concat(pair_label_list)\n\n           # Reset index for alignment\n           mouse_pair.reset_index(drop=True, inplace=True)\n           pair_meta.reset_index(drop=True, inplace=True)\n           pair_label.reset_index(drop=True, inplace=True)\n\n           print(f\"  Final check after concat & reset: Mouse Pair length = {len(mouse_pair)}, Pair Label length = {len(pair_label)}\")\n           X_tr_pair = transform_pair_v2(mouse_pair)\n           del mouse_pair\n           if X_tr_pair.shape[1] == 0:\n               print(f\"  Skipping submission for pair data as no features were generated.\")\n           else:\n              print(f\"  Pair: {X_tr_pair.shape}\")\n              submit_lstm(section, body_parts_tracked_str, 'pair', X_tr_pair, pair_label, pair_meta)\n               \n        else:\n              print(\"  Skipping pair processing: All generated pair dataframes were empty.\")\n                \n    except Exception as e:\n        \n        print(f'***Exception*** {str(e)[:100]}')\n\n    gc.collect()\n    print()\n\n# ======================== FINAL SUBMISSION ========================\nif submission_list:\n    \n    submission_list_non_empty = [df for df in submission_list if not df.empty]\n    \n    if submission_list_non_empty:\n        submission = pd.concat(submission_list_non_empty, ignore_index=True)\n    else:\n        submission = pd.DataFrame()\nelse:\n    submission = pd.DataFrame()\n\nsubmission_robust = robustify(submission, test, 'test')\n\nif submission_robust.empty:\n    print(\"WARNING: Submission is still empty after robustify. Creating a minimal default submission.\")\n    submission_robust = pd.DataFrame({\n        'video_id': [438887472], 'agent_id': ['mouse1'], 'target_id': ['self'],\n        'action': ['rear'], 'start_frame': [278], 'stop_frame': [500]\n    })\n\nsubmission_robust.index.name = 'row_id'\nsubmission_robust.to_csv('submission.csv')\nprint(f\"\\nSubmission created: {len(submission_robust)} predictions\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}