{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"}],"dockerImageVersionId":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# MABe Challenge - Memory-Optimized Neural Network\n# Reduced memory footprint while maintaining core functionality\n\nvalidate_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\nfrom collections import defaultdict, deque\nimport polars as pl\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.impute import SimpleImputer\nfrom sklearn.preprocessing import StandardScaler, PolynomialFeatures\nfrom sklearn.metrics import f1_score\nfrom scipy.stats import entropy, skew, kurtosis\nfrom scipy.fft import fft, fftfreq\nfrom scipy.signal import welch, find_peaks\n\nwarnings.filterwarnings('ignore')\nnp.seterr(all='ignore')\npd.options.mode.chained_assignment = None\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\n\n# ==================== CONFIGURATION (MEMORY OPTIMIZED) ====================\nclass Config:\n    \"\"\"Memory-optimized configuration\"\"\"\n    # Chunk processing - REDUCED\n    chunk_size = 5\n    replay_buffer_size = 3000  # Down from 15000\n    max_chunks_per_action = 10  # Down from 20\n    diversity_k = 30  # Down from 100\n    \n    # Validation\n    val_ratio = 0.2\n    convergence_patience = 4\n    min_improvement = 0.005\n    \n    # Neural network architecture - REDUCED\n    mlp_hidden_dims = [256, 128]  # Down from [512, 256, 128]\n    \n    # Temporal convolution - REDUCED\n    use_temporal_conv = True\n    temporal_conv_channels = [32, 64, 32]  # Down from [64, 128, 64]\n    temporal_kernel_sizes = [5, 5, 3]\n    \n    # LSTM - REDUCED\n    use_lstm = True\n    lstm_hidden_size = 64  # Down from 128\n    lstm_num_layers = 2\n    \n    # Attention - REDUCED\n    use_attention = True\n    attention_heads = 2  # Down from 4\n    attention_dim = 128\n    \n    # Graph network - DISABLED\n    use_graph_net = False\n    graph_hidden_dim = 64\n    graph_layers = 2\n    \n    dropout = 0.35\n    learning_rate = 0.0005\n    batch_size = 128  # Down from 256\n    epochs_per_chunk = 15\n    early_stopping_patience = 5\n    \n    # EWC - REDUCED\n    ewc_lambda = 2000\n    fisher_sample_size = 500  # Down from 2000\n    \n    # Advanced feature engineering - DISABLED\n    use_polynomial_features = False  # Was True\n    poly_degree = 2\n    poly_interaction_only = True\n    \n    use_wavelets = False  # Was True\n    use_phase_features = True\n    use_higher_order_stats = False  # Was True\n    \n    # Feature windows - REDUCED\n    temporal_windows = [5, 15, 30, 60, 120]  # Down from 11 windows\n    long_range_windows = [60, 120, 240]  # Down from 8 windows\n    skip_distances = [15, 30, 60, 120]  # Down from 8 distances\n    entropy_windows = [10, 30, 60]  # Down from 5\n    spectral_windows = [30, 60]  # Down from 3\n    \n    drop_body_parts = [\n        'headpiece_bottombackleft', 'headpiece_bottombackright',\n        'headpiece_bottomfrontleft', 'headpiece_bottomfrontright',\n        'headpiece_topbackleft', 'headpiece_topbackright',\n        'headpiece_topfrontleft', 'headpiece_topfrontright',\n        'spine_1', 'spine_2', 'tail_middle_1', 'tail_middle_2', 'tail_midpoint'\n    ]\n\nconfig = Config()\n\n# ==================== SCORING FUNCTIONS ====================\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(set)\n    prediction_frames = 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 = lab_solution.filter(pl.col('video_id') == video)['behaviors_labeled'].first()\n        active_labels = set(json.loads(active_labels))\n        predicted_mouse_pairs = 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 same frame')\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    for col in expected_cols:\n        if col not in solution.columns:\n            raise ValueError(f'Solution missing {col}')\n        if col not in submission.columns:\n            raise ValueError(f'Submission missing {col}')\n\n    solution = pl.DataFrame(solution)\n    submission = pl.DataFrame(submission)\n    \n    solution = solution.with_columns(\n        pl.concat_str([pl.col('video_id').cast(pl.Utf8), pl.col('agent_id').cast(pl.Utf8),\n                      pl.col('target_id').cast(pl.Utf8), pl.col('action')], separator='_').alias('label_key'))\n    submission = submission.with_columns(\n        pl.concat_str([pl.col('video_id').cast(pl.Utf8), pl.col('agent_id').cast(pl.Utf8),\n                      pl.col('target_id').cast(pl.Utf8), pl.col('action')], separator='_').alias('prediction_key'))\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# ==================== ADVANCED NEURAL NETWORKS ====================\n\nclass TemporalConvBlock(nn.Module):\n    \"\"\"1D Convolutional block for temporal features\"\"\"\n    def __init__(self, in_channels, out_channels, kernel_size=5, dropout=0.3):\n        super().__init__()\n        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, padding=kernel_size//2)\n        self.bn = nn.BatchNorm1d(out_channels)\n        self.dropout = nn.Dropout(dropout)\n        \n    def forward(self, x):\n        x = self.conv(x)\n        x = self.bn(x)\n        x = F.relu(x)\n        x = self.dropout(x)\n        return x\n\nclass MultiHeadAttention(nn.Module):\n    \"\"\"Multi-head self-attention for feature importance\"\"\"\n    def __init__(self, embed_dim, num_heads=4):\n        super().__init__()\n        self.attention = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)\n        self.norm = nn.LayerNorm(embed_dim)\n        \n    def forward(self, x):\n        if x.dim() == 2:\n            x = x.unsqueeze(1)\n        \n        attn_out, _ = self.attention(x, x, x)\n        x = self.norm(x + attn_out)\n        x = x.mean(dim=1)\n        return x\n\nclass AdvancedBehaviorNet(nn.Module):\n    \"\"\"Advanced neural network with multiple architectural components\"\"\"\n    \n    def __init__(self, input_dim, \n                 mlp_hidden=[256, 128],\n                 use_temporal_conv=True, temporal_channels=[32, 64, 32],\n                 use_lstm=True, lstm_hidden=64, lstm_layers=2,\n                 use_attention=True, attention_heads=2,\n                 dropout=0.35):\n        super().__init__()\n        \n        self.use_temporal_conv = use_temporal_conv\n        self.use_lstm = use_lstm\n        self.use_attention = use_attention\n        \n        current_dim = input_dim\n        \n        if use_temporal_conv:\n            conv_layers = []\n            in_ch = 1\n            for out_ch in temporal_channels:\n                conv_layers.append(TemporalConvBlock(in_ch, out_ch, dropout=dropout))\n                in_ch = out_ch\n            self.temporal_conv = nn.Sequential(*conv_layers)\n            self.temporal_pool = nn.AdaptiveAvgPool1d(1)\n            current_dim += temporal_channels[-1]\n        \n        if use_lstm:\n            self.lstm = nn.LSTM(input_dim, lstm_hidden, lstm_layers, \n                               batch_first=True, dropout=dropout if lstm_layers > 1 else 0)\n            current_dim += lstm_hidden\n        \n        if use_attention:\n            self.attention_proj = nn.Linear(input_dim, attention_heads * 32)\n            self.attention = MultiHeadAttention(attention_heads * 32, attention_heads)\n            current_dim += attention_heads * 32\n        \n        mlp_layers = []\n        prev_dim = current_dim\n        for hidden_dim in mlp_hidden:\n            mlp_layers.extend([\n                nn.Linear(prev_dim, hidden_dim),\n                nn.BatchNorm1d(hidden_dim),\n                nn.ReLU(),\n                nn.Dropout(dropout)\n            ])\n            prev_dim = hidden_dim\n        \n        mlp_layers.append(nn.Linear(prev_dim, 1))\n        self.mlp = nn.Sequential(*mlp_layers)\n        \n    def forward(self, x):\n        features = [x]\n        \n        if self.use_temporal_conv:\n            x_temp = x.unsqueeze(1)\n            x_temp = self.temporal_conv(x_temp)\n            x_temp = self.temporal_pool(x_temp).squeeze(-1)\n            features.append(x_temp)\n        \n        if self.use_lstm:\n            x_lstm = x.unsqueeze(1)\n            _, (h_n, _) = self.lstm(x_lstm)\n            x_lstm = h_n[-1]\n            features.append(x_lstm)\n        \n        if self.use_attention:\n            x_attn = self.attention_proj(x)\n            x_attn = self.attention(x_attn)\n            features.append(x_attn)\n        \n        x_combined = torch.cat(features, dim=1)\n        output = self.mlp(x_combined)\n        \n        return output\n\nclass EWCModel:\n    \"\"\"Neural network with EWC for continual learning\"\"\"\n    \n    def __init__(self, input_dim, \n                 mlp_hidden=[256, 128],\n                 use_temporal_conv=True, temporal_channels=[32, 64, 32],\n                 use_lstm=True, lstm_hidden=64, lstm_layers=2,\n                 use_attention=True, attention_heads=2,\n                 dropout=0.35, learning_rate=0.0005, \n                 ewc_lambda=2000, device='cpu'):\n        \n        self.device = device\n        self.model = AdvancedBehaviorNet(\n            input_dim, mlp_hidden, use_temporal_conv, temporal_channels,\n            use_lstm, lstm_hidden, lstm_layers, use_attention, attention_heads,\n            dropout\n        ).to(device)\n        \n        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=learning_rate)\n        self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n            self.optimizer, mode='min', factor=0.5, patience=3, verbose=False\n        )\n        self.ewc_lambda = ewc_lambda\n        \n        self.fisher_matrices = {}\n        self.optimal_params = {}\n        self.task_count = 0\n        \n        self.scaler = StandardScaler()\n        self.scaler_fitted = False\n        \n    def compute_fisher_matrix(self, dataloader, sample_size=500):\n        \"\"\"Compute Fisher Information Matrix\"\"\"\n        self.model.eval()\n        fisher = {n: torch.zeros_like(p) for n, p in self.model.named_parameters()}\n        \n        samples_seen = 0\n        for inputs, targets in dataloader:\n            if samples_seen >= sample_size:\n                break\n                \n            inputs = inputs.to(self.device)\n            targets = targets.to(self.device)\n            \n            self.model.zero_grad()\n            outputs = self.model(inputs).squeeze()\n            probs = torch.sigmoid(outputs)\n            \n            log_probs = torch.log(probs + 1e-8) * (targets == 1).float() + \\\n                       torch.log(1 - probs + 1e-8) * (targets == 0).float()\n            loss = -log_probs.mean()\n            loss.backward()\n            \n            for n, p in self.model.named_parameters():\n                if p.grad is not None:\n                    fisher[n] += p.grad.data ** 2 / len(dataloader.dataset)\n            \n            samples_seen += len(inputs)\n        \n        return fisher\n    \n    def ewc_loss(self):\n        \"\"\"Compute EWC regularization\"\"\"\n        loss = 0\n        for task_id in range(self.task_count):\n            for n, p in self.model.named_parameters():\n                fisher = self.fisher_matrices[task_id][n]\n                optimal = self.optimal_params[task_id][n]\n                loss += (fisher * (p - optimal) ** 2).sum()\n        return self.ewc_lambda * loss\n    \n    def fit(self, X, y, val_X=None, val_y=None, epochs=15, batch_size=128, \n            early_stopping_patience=5):\n        \"\"\"Train model\"\"\"\n        \n        if not self.scaler_fitted:\n            X_scaled = self.scaler.fit_transform(X)\n            self.scaler_fitted = True\n        else:\n            X_scaled = self.scaler.transform(X)\n        \n        dataset = torch.utils.data.TensorDataset(\n            torch.FloatTensor(X_scaled),\n            torch.FloatTensor(y)\n        )\n        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n        \n        if val_X is not None and val_y is not None:\n            val_X_scaled = self.scaler.transform(val_X)\n            val_dataset = torch.utils.data.TensorDataset(\n                torch.FloatTensor(val_X_scaled),\n                torch.FloatTensor(val_y)\n            )\n            val_dataloader = DataLoader(val_dataset, batch_size=batch_size)\n        else:\n            val_dataloader = None\n        \n        best_val_loss = float('inf')\n        patience_counter = 0\n        \n        for epoch in range(epochs):\n            self.model.train()\n            train_loss = 0\n            \n            for inputs, targets in dataloader:\n                inputs = inputs.to(self.device)\n                targets = targets.to(self.device)\n                \n                self.optimizer.zero_grad()\n                outputs = self.model(inputs).squeeze()\n                \n                bce_loss = F.binary_cross_entropy_with_logits(outputs, targets, reduction='none')\n                pt = torch.exp(-bce_loss)\n                focal_loss = ((1 - pt) ** 2 * bce_loss).mean()\n                \n                loss = focal_loss + self.ewc_loss()\n                \n                loss.backward()\n                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)\n                self.optimizer.step()\n                \n                train_loss += loss.item()\n            \n            if val_dataloader is not None:\n                val_loss = self.evaluate(val_dataloader)\n                self.scheduler.step(val_loss)\n                \n                if val_loss < best_val_loss:\n                    best_val_loss = val_loss\n                    patience_counter = 0\n                else:\n                    patience_counter += 1\n                \n                if patience_counter >= early_stopping_patience:\n                    if verbose:\n                        print(f\"      Early stop at epoch {epoch+1}\")\n                    break\n        \n        self.update_ewc(dataloader)\n    \n    def update_ewc(self, dataloader):\n        \"\"\"Update Fisher and optimal params\"\"\"\n        fisher = self.compute_fisher_matrix(dataloader, config.fisher_sample_size)\n        optimal = {n: p.data.clone() for n, p in self.model.named_parameters()}\n        \n        self.fisher_matrices[self.task_count] = fisher\n        self.optimal_params[self.task_count] = optimal\n        self.task_count += 1\n    \n    def evaluate(self, dataloader):\n        \"\"\"Evaluate\"\"\"\n        self.model.eval()\n        total_loss = 0\n        \n        with torch.no_grad():\n            for inputs, targets in dataloader:\n                inputs = inputs.to(self.device)\n                targets = targets.to(self.device)\n                \n                outputs = self.model(inputs).squeeze()\n                loss = F.binary_cross_entropy_with_logits(outputs, targets)\n                total_loss += loss.item()\n        \n        return total_loss / len(dataloader)\n    \n    def predict_proba(self, X):\n        \"\"\"Predict probabilities\"\"\"\n        self.model.eval()\n        \n        X_scaled = self.scaler.transform(X)\n        X_tensor = torch.FloatTensor(X_scaled).to(self.device)\n        \n        with torch.no_grad():\n            outputs = self.model(X_tensor).squeeze()\n            probs = torch.sigmoid(outputs).cpu().numpy()\n        \n        return np.column_stack([1 - probs, probs])\n\n# ==================== REPLAY BUFFER ====================\nclass DiverseReplayBuffer:\n    \"\"\"Diverse replay buffer\"\"\"\n    def __init__(self, max_size: int, diversity_k: int):\n        self.buffer = deque(maxlen=max_size)\n        self.max_size = max_size\n        self.diversity_k = diversity_k\n        self.action_counts = defaultdict(int)\n    \n    def add(self, X: pd.DataFrame, y: np.ndarray, video_id: int):\n        \"\"\"Add samples\"\"\"\n        unique_actions = np.unique(y)\n        for action_val in unique_actions:\n            mask = (y == action_val)\n            X_action = X[mask]\n            \n            if len(X_action) == 0:\n                continue\n            \n            if len(X_action) > self.diversity_k:\n                indices = np.linspace(0, len(X_action)-1, self.diversity_k, dtype=int)\n                X_selected = X_action.iloc[indices]\n            else:\n                X_selected = X_action\n            \n            for idx, row in X_selected.iterrows():\n                self.buffer.append({\n                    'X': row.to_dict(),\n                    'y': action_val,\n                    'video_id': video_id\n                })\n                self.action_counts[action_val] += 1\n    \n    def sample(self, n_samples: int, columns: list):\n        \"\"\"Sample from buffer\"\"\"\n        if len(self.buffer) == 0:\n            return None, None\n        \n        sample_size = min(n_samples, len(self.buffer))\n        step = max(1, len(self.buffer) // sample_size)\n        indices = list(range(0, len(self.buffer), step))[:sample_size]\n        \n        samples = [self.buffer[i] for i in indices]\n        X_list, y_list = [], []\n        \n        for sample in samples:\n            try:\n                row = [sample['X'].get(col, np.nan) for col in columns]\n                X_list.append(row)\n                y_list.append(sample['y'])\n            except:\n                continue\n        \n        if len(X_list) > 0:\n            return pd.DataFrame(X_list, columns=columns), np.array(y_list)\n        return None, None\n\n# ==================== FEATURE ENGINEERING (MEMORY OPTIMIZED) ====================\n\ndef safe_entropy(x):\n    x_clean = x[~np.isnan(x)]\n    if len(x_clean) < 2:\n        return 0.0\n    try:\n        hist, _ = np.histogram(x_clean, bins=10)\n        return entropy(hist + 1e-10)\n    except:\n        return 0.0\n\ndef safe_fft_features(x):\n    \"\"\"Extract FFT features\"\"\"\n    x_clean = x[~np.isnan(x)]\n    if len(x_clean) < 4:\n        return 0.0, 0.0, 0.0\n    try:\n        fft_vals = np.abs(fft(x_clean)[:len(x_clean)//2])\n        freqs = fftfreq(len(x_clean), 1.0)[:len(x_clean)//2]\n        \n        peak_freq = freqs[np.argmax(fft_vals)] if len(fft_vals) > 0 else 0\n        peak_power = fft_vals.max()\n        spectral_centroid = np.sum(freqs * fft_vals) / (np.sum(fft_vals) + 1e-10)\n        \n        return peak_freq, peak_power, spectral_centroid\n    except:\n        return 0.0, 0.0, 0.0\n\ndef safe_welch_features(x, window=30):\n    \"\"\"Power spectral density via Welch method\"\"\"\n    x_clean = x[~np.isnan(x)]\n    if len(x_clean) < window:\n        return 0.0, 0.0\n    try:\n        f, psd = welch(x_clean, nperseg=min(window, len(x_clean)))\n        \n        dominant_freq = f[np.argmax(psd)]\n        spectral_entropy = entropy(psd + 1e-10)\n        \n        return dominant_freq, spectral_entropy\n    except:\n        return 0.0, 0.0\n\ndef add_spectral_features(X, data, windows=[30, 60]):\n    \"\"\"Add spectral features - REDUCED\"\"\"\n    for part in ['nose', 'body_center', 'tail_base']:\n        if part in data.columns.get_level_values(0):\n            for axis in ['x', 'y']:\n                if axis in data[part].columns:\n                    for window in windows:\n                        peak_freq = data[part][axis].rolling(window, min_periods=1).apply(\n                            lambda x: safe_fft_features(x)[0], raw=True\n                        )\n                        X[f'{part}_{axis}_peak_freq_{window}'] = peak_freq.astype(np.float32)\n                        \n                        peak_power = data[part][axis].rolling(window, min_periods=1).apply(\n                            lambda x: safe_fft_features(x)[1], raw=True\n                        )\n                        X[f'{part}_{axis}_peak_power_{window}'] = peak_power.astype(np.float32)\n    \n    return X\n\ndef add_phase_features(X, data):\n    \"\"\"Add phase and circular statistics\"\"\"\n    if 'nose' in data.columns.get_level_values(0) and 'tail_base' in data.columns.get_level_values(0):\n        dx = data['nose']['x'] - data['tail_base']['x']\n        dy = data['nose']['y'] - data['tail_base']['y']\n        angle = np.arctan2(dy, dx)\n        \n        X['body_angle_rad'] = angle.astype(np.float32)\n        X['body_angle_sin'] = np.sin(angle).astype(np.float32)\n        X['body_angle_cos'] = np.cos(angle).astype(np.float32)\n        \n        X['angular_vel'] = angle.diff().astype(np.float32)\n        X['angular_accel'] = X['angular_vel'].diff().astype(np.float32)\n        \n        for window in [10, 30, 60]:\n            angle_var = angle.rolling(window, min_periods=1).var()\n            X[f'phase_coherence_{window}'] = (1 / (1 + angle_var)).astype(np.float32)\n    \n    return X\n\ndef add_interaction_features(X, data):\n    \"\"\"Add interaction terms between key body parts\"\"\"\n    available_parts = data.columns.get_level_values(0).unique()\n    \n    if 'nose' in available_parts and 'tail_base' in available_parts:\n        nose_speed = np.sqrt(data['nose']['x'].diff()**2 + data['nose']['y'].diff()**2)\n        tail_speed = np.sqrt(data['tail_base']['x'].diff()**2 + data['tail_base']['y'].diff()**2)\n        \n        X['nose_tail_speed_product'] = (nose_speed * tail_speed).astype(np.float32)\n        X['nose_tail_speed_ratio'] = (nose_speed / (tail_speed + 1e-6)).astype(np.float32)\n    \n    if 'ear_left' in available_parts and 'ear_right' in available_parts:\n        ear_dist = np.sqrt(\n            (data['ear_left']['x'] - data['ear_right']['x'])**2 +\n            (data['ear_left']['y'] - data['ear_right']['y'])**2\n        )\n        X['ear_dist'] = ear_dist.astype(np.float32)\n        X['ear_dist_change'] = ear_dist.diff().astype(np.float32)\n        X['ear_dist_accel'] = ear_dist.diff().diff().astype(np.float32)\n    \n    return X\n\ndef add_autoregressive_features(X, data, lags=[1, 3, 5, 10]):  # REDUCED from [1,2,3,5,10]\n    \"\"\"Add autoregressive features - REDUCED\"\"\"\n    if 'body_center' in data.columns.get_level_values(0):\n        center_x = data['body_center']['x']\n        center_y = data['body_center']['y']\n        \n        for lag in lags:\n            X[f'ar_x_lag{lag}'] = center_x.shift(lag).astype(np.float32)\n            X[f'ar_y_lag{lag}'] = center_y.shift(lag).astype(np.float32)\n            \n            X[f'autocorr_x_{lag}'] = (center_x * center_x.shift(lag)).astype(np.float32)\n            X[f'autocorr_y_{lag}'] = (center_y * center_y.shift(lag)).astype(np.float32)\n    \n    return X\n\ndef transform_single(single_mouse, body_parts_tracked):\n    \"\"\"Comprehensive single mouse features - MEMORY OPTIMIZED\"\"\"\n    available_body_parts = single_mouse.columns.get_level_values(0)\n    \n    X = pd.DataFrame()\n    \n    # Base pairwise distances\n    for part1, part2 in itertools.combinations(body_parts_tracked, 2):\n        if part1 in available_body_parts and part2 in available_body_parts:\n            X[f\"{part1}_{part2}\"] = np.square(single_mouse[part1] - single_mouse[part2]).sum(axis=1, skipna=False).astype(np.float32)\n    \n    # Speed features - REDUCED\n    if 'ear_left' in single_mouse.columns and 'ear_right' in single_mouse.columns:\n        for dt in [1, 5, 10]:  # REDUCED from [1, 3, 5, 10, 20]\n            shifted = single_mouse[['ear_left', 'ear_right']].shift(dt)\n            X[f'speed_left_{dt}'] = np.square(single_mouse['ear_left'] - shifted['ear_left']).sum(axis=1, skipna=False).astype(np.float32)\n            X[f'speed_right_{dt}'] = np.square(single_mouse['ear_right'] - shifted['ear_right']).sum(axis=1, skipna=False).astype(np.float32)\n    \n    # Body geometry\n    if 'nose' in available_body_parts and 'body_center' in available_body_parts and 'tail_base' in available_body_parts:\n        v1 = single_mouse['nose'] - single_mouse['body_center']\n        v2 = single_mouse['tail_base'] - single_mouse['body_center']\n        \n        dot_product = (v1['x'] * v2['x'] + v1['y'] * v2['y'])\n        norm_v1 = np.sqrt(v1['x']**2 + v1['y']**2)\n        norm_v2 = np.sqrt(v2['x']**2 + v2['y']**2)\n        \n        X['body_angle'] = (dot_product / (norm_v1 * norm_v2 + 1e-6)).astype(np.float32)\n        \n        body_angle = np.arctan2(v1['y'], v1['x']) - np.arctan2(v2['y'], v2['x'])\n        X['angle_change'] = body_angle.diff().astype(np.float32)\n        X['angle_std'] = body_angle.rolling(15, min_periods=1).std().astype(np.float32)\n    \n    # Body length dynamics\n    if 'nose' in available_body_parts and 'tail_base' in available_body_parts:\n        body_length = np.sqrt(\n            (single_mouse['nose']['x'] - single_mouse['tail_base']['x'])**2 + \n            (single_mouse['nose']['y'] - single_mouse['tail_base']['y'])**2\n        ).astype(np.float32)\n        \n        X['body_length'] = body_length\n        X['length_change'] = body_length.diff().astype(np.float32)\n        X['length_std'] = body_length.rolling(20, min_periods=1).std().astype(np.float32)\n        X['length_accel'] = body_length.diff().diff().astype(np.float32)\n    \n    # Multi-scale temporal features\n    if 'body_center' in available_body_parts:\n        center_x = single_mouse['body_center']['x']\n        center_y = single_mouse['body_center']['y']\n        \n        for dt in [1, 3, 5, 10, 15]:  # REDUCED from [1,2,3,5,10,15]\n            velocity = np.sqrt(\n                (center_x - center_x.shift(dt))**2 + \n                (center_y - center_y.shift(dt))**2\n            ) / dt\n            X[f'velocity_{dt}'] = velocity.astype(np.float32)\n            X[f'accel_{dt}'] = velocity.diff().astype(np.float32)\n        \n        for window in config.temporal_windows:\n            X[f'cx_mean_{window}'] = center_x.rolling(window, min_periods=1, center=True).mean().astype(np.float32)\n            X[f'cy_mean_{window}'] = center_y.rolling(window, min_periods=1, center=True).mean().astype(np.float32)\n            X[f'cx_std_{window}'] = center_x.rolling(window, min_periods=1, center=True).std().astype(np.float32)\n            X[f'cy_std_{window}'] = center_y.rolling(window, min_periods=1, center=True).std().astype(np.float32)\n            \n            X[f'x_range_{window}'] = (center_x.rolling(window, min_periods=1, center=True).max() - \n                                      center_x.rolling(window, min_periods=1, center=True).min()).astype(np.float32)\n            X[f'y_range_{window}'] = (center_y.rolling(window, min_periods=1, center=True).max() - \n                                      center_y.rolling(window, min_periods=1, center=True).min()).astype(np.float32)\n        \n        for window in [15, 30, 60]:  # REDUCED from [15, 30, 60, 90]\n            dist_traveled = np.sqrt(center_x.diff()**2 + center_y.diff()**2).rolling(window, min_periods=1).sum()\n            displacement = np.sqrt(\n                (center_x - center_x.shift(window))**2 + \n                (center_y - center_y.shift(window))**2\n            )\n            X[f'tortuosity_{window}'] = (dist_traveled / (displacement + 1e-6)).astype(np.float32)\n    \n    # Head dynamics\n    if 'nose' in available_body_parts:\n        nose_x = single_mouse['nose']['x']\n        nose_y = single_mouse['nose']['y']\n        \n        X['head_vel'] = np.sqrt(nose_x.diff()**2 + nose_y.diff()**2).astype(np.float32)\n        X['head_accel'] = X['head_vel'].diff().astype(np.float32)\n        X['head_jerk'] = X['head_accel'].diff().astype(np.float32)\n        \n        head_dir = np.arctan2(nose_y.diff(), nose_x.diff())\n        X['head_dir_change'] = head_dir.diff().astype(np.float32)\n        X['head_dir_std'] = head_dir.rolling(10, min_periods=1).std().astype(np.float32)\n    \n    # Long-range features\n    for window in config.long_range_windows:\n        if 'body_center' in available_body_parts:\n            center_x = single_mouse['body_center']['x']\n            center_y = single_mouse['body_center']['y']\n            \n            X[f'long_disp_{window}'] = np.sqrt(\n                (center_x - center_x.shift(window))**2 + \n                (center_y - center_y.shift(window))**2\n            ).astype(np.float32)\n            \n            dx = center_x.diff()\n            dy = center_y.diff()\n            speed = np.sqrt(dx**2 + dy**2)\n            \n            X[f'burst_{window}'] = (\n                speed.rolling(window, min_periods=1).max() / \n                (speed.rolling(window, min_periods=1).mean() + 1e-6)\n            ).astype(np.float32)\n    \n    # Skip-gram features\n    for skip in config.skip_distances:\n        if 'nose' in available_body_parts:\n            X[f'nose_skip_{skip}'] = np.sqrt(\n                (single_mouse['nose']['x'] - single_mouse['nose']['x'].shift(skip))**2 +\n                (single_mouse['nose']['y'] - single_mouse['nose']['y'].shift(skip))**2\n            ).astype(np.float32)\n    \n    # Behavior transitions\n    if 'body_center' in available_body_parts:\n        center_x = single_mouse['body_center']['x']\n        center_y = single_mouse['body_center']['y']\n        vx = center_x.diff()\n        vy = center_y.diff()\n        speed = np.sqrt(vx**2 + vy**2)\n        acceleration = speed.diff()\n        \n        for window in [30, 60, 120]:\n            X[f'move_var_{window}'] = speed.rolling(window, min_periods=1).std().astype(np.float32)\n            \n            active_threshold = speed.quantile(0.7) if len(speed) > 0 and not speed.isna().all() else 0\n            is_active = (speed > active_threshold).astype(float)\n            transitions = is_active.diff().abs()\n            X[f'transitions_{window}'] = transitions.rolling(window, min_periods=1).sum().astype(np.float32)\n            X[f'active_persist_{window}'] = is_active.rolling(window, min_periods=1).mean().astype(np.float32)\n            X[f'accel_changes_{window}'] = acceleration.rolling(window, min_periods=1).std().astype(np.float32)\n    \n    # Advanced features (conditional)\n    if config.use_phase_features:\n        X = add_phase_features(X, single_mouse)\n    \n    X = add_spectral_features(X, single_mouse, config.spectral_windows)\n    X = add_interaction_features(X, single_mouse)\n    X = add_autoregressive_features(X, single_mouse)\n    \n    # Entropy features\n    for window in config.entropy_windows:\n        if 'body_center' in single_mouse.columns.get_level_values(0):\n            for axis in ['x', 'y']:\n                if axis in single_mouse['body_center'].columns:\n                    ent = single_mouse['body_center'][axis].rolling(window, min_periods=1).apply(\n                        safe_entropy, raw=True\n                    )\n                    X[f'entropy_{axis}_{window}'] = ent.astype(np.float32)\n    \n    return X\n\ndef transform_pair(mouse_pair, body_parts_tracked):\n    \"\"\"Comprehensive pair features - MEMORY OPTIMIZED\"\"\"\n    available_body_parts_A = mouse_pair['A'].columns.get_level_values(0)\n    available_body_parts_B = mouse_pair['B'].columns.get_level_values(0)\n    \n    X = pd.DataFrame()\n    \n    # Pairwise distances\n    for part1, part2 in itertools.product(body_parts_tracked, repeat=2):\n        if part1 in available_body_parts_A and part2 in available_body_parts_B:\n            X[f\"AB_{part1}_{part2}\"] = np.square(mouse_pair['A'][part1] - mouse_pair['B'][part2]).sum(axis=1, skipna=False).astype(np.float32)\n    \n    # Relative speeds - REDUCED\n    if ('A', 'ear_left') in mouse_pair.columns and ('B', 'ear_left') in mouse_pair.columns:\n        for dt in [1, 5, 10]:  # REDUCED from [1, 3, 5, 10, 20]\n            shifted_A = mouse_pair['A']['ear_left'].shift(dt)\n            shifted_B = mouse_pair['B']['ear_left'].shift(dt)\n            X[f'speed_A_{dt}'] = np.square(mouse_pair['A']['ear_left'] - shifted_A).sum(axis=1, skipna=False).astype(np.float32)\n            X[f'speed_B_{dt}'] = np.square(mouse_pair['B']['ear_left'] - shifted_B).sum(axis=1, skipna=False).astype(np.float32)\n            X[f'speed_diff_{dt}'] = (X[f'speed_A_{dt}'] - X[f'speed_B_{dt}']).astype(np.float32)\n    \n    # Relative orientation\n    if 'nose' in available_body_parts_A and 'tail_base' in available_body_parts_A and \\\n       'nose' in available_body_parts_B and 'tail_base' in available_body_parts_B:\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        \n        dot_product = (dir_A['x'] * dir_B['x'] + dir_A['y'] * dir_B['y'])\n        norm_A = np.sqrt(dir_A['x']**2 + dir_A['y']**2)\n        norm_B = np.sqrt(dir_B['x']**2 + dir_B['y']**2)\n        \n        X['rel_orient'] = (dot_product / (norm_A * norm_B + 1e-6)).astype(np.float32)\n        \n        A_to_B = mouse_pair['B']['body_center'] - mouse_pair['A']['body_center'] if 'body_center' in available_body_parts_A else mouse_pair['B']['nose'] - mouse_pair['A']['nose']\n        dot_A = (dir_A['x'] * A_to_B['x'] + dir_A['y'] * A_to_B['y'])\n        dot_B = (-dir_B['x'] * A_to_B['x'] + -dir_B['y'] * A_to_B['y'])\n        \n        X['A_facing'] = (dot_A / (norm_A * np.sqrt(A_to_B['x']**2 + A_to_B['y']**2) + 1e-6)).astype(np.float32)\n        X['B_facing'] = (dot_B / (norm_B * np.sqrt(A_to_B['x']**2 + A_to_B['y']**2) + 1e-6)).astype(np.float32)\n        X['mutual_face'] = (X['A_facing'] * X['B_facing']).astype(np.float32)\n    \n    # Approach dynamics - REDUCED\n    if 'nose' in available_body_parts_A and 'nose' in available_body_parts_B:\n        nose_dist = np.sqrt(\n            (mouse_pair['A']['nose']['x'] - mouse_pair['B']['nose']['x'])**2 +\n            (mouse_pair['A']['nose']['y'] - mouse_pair['B']['nose']['y'])**2\n        ).astype(np.float32)\n        \n        for dt in [1, 5, 10]:  # REDUCED from [1, 3, 5, 10, 20]\n            X[f'nose_change_{dt}'] = (nose_dist - nose_dist.shift(dt)).astype(np.float32)\n            X[f'approach_{dt}'] = (-X[f'nose_change_{dt}'] / dt).astype(np.float32)\n        \n        # Proximity zones\n        for thresh in [3, 5, 10, 15, 20]:\n            is_close = (nose_dist < thresh).astype(np.float32)\n            for window in [5, 10, 30, 60]:\n                X[f'close_{thresh}_{window}'] = is_close.rolling(window, min_periods=1).mean().astype(np.float32)\n    \n    # Distance features\n    if 'body_center' in available_body_parts_A and 'body_center' in available_body_parts_B:\n        center_dist = np.sqrt(\n            (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        ).astype(np.float32)\n        \n        X['very_close'] = (center_dist < 5.0).astype(np.float32)\n        X['close'] = ((center_dist >= 5.0) & (center_dist < 10.0)).astype(np.float32)\n        X['medium'] = ((center_dist >= 10.0) & (center_dist < 20.0)).astype(np.float32)\n        X['far'] = (center_dist >= 20.0).astype(np.float32)\n        \n        for window in config.temporal_windows:\n            X[f'dist_mean_{window}'] = center_dist.rolling(window, min_periods=1, center=True).mean().astype(np.float32)\n            X[f'dist_std_{window}'] = center_dist.rolling(window, min_periods=1, center=True).std().astype(np.float32)\n            X[f'dist_min_{window}'] = center_dist.rolling(window, min_periods=1, center=True).min().astype(np.float32)\n            X[f'dist_max_{window}'] = center_dist.rolling(window, min_periods=1, center=True).max().astype(np.float32)\n            \n            dist_var = center_dist.rolling(window, min_periods=1, center=True).var()\n            X[f'interact_{window}'] = (1 / (1 + dist_var)).astype(np.float32)\n        \n        # Coordinated movement\n        A_vel_x = mouse_pair['A']['body_center']['x'].diff()\n        A_vel_y = mouse_pair['A']['body_center']['y'].diff()\n        B_vel_x = mouse_pair['B']['body_center']['x'].diff()\n        B_vel_y = mouse_pair['B']['body_center']['y'].diff()\n        \n        vel_alignment = (A_vel_x * B_vel_x + A_vel_y * B_vel_y) / (\n            np.sqrt(A_vel_x**2 + A_vel_y**2) * np.sqrt(B_vel_x**2 + B_vel_y**2) + 1e-6\n        )\n        \n        X['vel_align'] = vel_alignment.astype(np.float32)\n        \n        for window in [5, 10, 30, 60]:\n            X[f'vel_align_mean_{window}'] = vel_alignment.rolling(window, min_periods=1, center=True).mean().astype(np.float32)\n            X[f'vel_align_std_{window}'] = vel_alignment.rolling(window, min_periods=1, center=True).std().astype(np.float32)\n        \n        # Following behavior - REDUCED\n        for lag in [5, 10, 15]:  # REDUCED from [3, 5, 10, 15, 20]\n            A_vel_x_lag = A_vel_x.shift(lag)\n            A_vel_y_lag = A_vel_y.shift(lag)\n            \n            follow_score = (A_vel_x_lag * B_vel_x + A_vel_y_lag * B_vel_y) / (\n                np.sqrt(A_vel_x_lag**2 + A_vel_y_lag**2) * np.sqrt(B_vel_x**2 + B_vel_y**2) + 1e-6\n            )\n            X[f'follow_{lag}'] = follow_score.astype(np.float32)\n        \n        # Long-range interaction\n        for long_window in config.long_range_windows:\n            X[f'dist_long_{long_window}'] = center_dist.rolling(long_window, min_periods=1, center=True).mean().astype(np.float32)\n    \n    # Add features for mouse A\n    X = add_spectral_features(X, mouse_pair['A'], config.spectral_windows)\n    X = add_interaction_features(X, mouse_pair['A'])\n    X = add_autoregressive_features(X, mouse_pair['A'])\n    \n    if config.use_phase_features:\n        X = add_phase_features(X, mouse_pair['A'])\n    \n    return X\n\n# ==================== DATA GENERATOR ====================\ndef generate_mouse_data(dataset, traintest, traintest_directory=None, generate_single=True, generate_pair=True):\n    \"\"\"Generate mouse data\"\"\"\n    assert traintest in ['train', 'test']\n    if traintest_directory is None:\n        traintest_directory = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking\"\n    \n    for _, 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\n        if type(row.behaviors_labeled) != str:\n            continue\n\n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        try:\n            vid = pd.read_parquet(path)\n        except:\n            continue\n            \n        if len(np.unique(vid.bodypart)) > 5:\n            vid = vid.query(\"~ bodypart.isin(@config.drop_body_parts)\")\n        \n        pvid = vid.pivot(columns=['mouse_id', 'bodypart'], index='video_frame', values=['x', 'y'])\n        del vid\n        gc.collect()\n        \n        pvid = pvid.reorder_levels([1, 2, 0], axis=1).T.sort_index().T\n        pvid = (pvid / row.pix_per_cm_approx).astype(np.float32)\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                    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, video_id\n                    else:\n                        yield 'single', single_mouse, single_mouse_meta, vid_agent_actions, video_id\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                    if len(vid_agent_actions) == 0:\n                        continue\n                    mouse_pair = pd.concat([pvid[agent], pvid[target]], axis=1, keys=['A', 'B'])\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, video_id\n                    else:\n                        yield 'pair', mouse_pair, mouse_pair_meta, vid_agent_actions, video_id\n\n# ==================== CHUNK LEARNING MANAGER ====================\nclass ChunkLearningManager:\n    \"\"\"Manages chunk-based learning with advanced neural networks\"\"\"\n    \n    def __init__(self):\n        self.imputer = SimpleImputer(strategy='median')\n        self.action_models = {}\n        self.replay_buffers = {}\n        self.feature_columns = {}\n        self.action_thresholds = {}\n        self.convergence_history = defaultdict(list)\n        self.train_videos = set()\n        self.val_videos = set()\n        self.imputer_fitted = False\n    \n    def split_videos(self, video_ids):\n        \"\"\"Split for validation\"\"\"\n        video_ids = list(video_ids)\n        np.random.shuffle(video_ids)\n        split_idx = int(len(video_ids) * (1 - config.val_ratio))\n        self.train_videos = set(video_ids[:split_idx])\n        self.val_videos = set(video_ids[split_idx:])\n        if verbose:\n            print(f\"      Train/Val: {len(self.train_videos)}/{len(self.val_videos)}\")\n    \n    def check_convergence(self, action):\n        \"\"\"Check convergence\"\"\"\n        history = self.convergence_history[action]\n        if len(history) < config.convergence_patience + 1:\n            return False\n        \n        recent_best = max(history[-config.convergence_patience:])\n        overall_best = max(history[:-config.convergence_patience]) if len(history) > config.convergence_patience else 0\n        improvement = recent_best - overall_best\n        \n        return improvement < config.min_improvement\n    \n    def optimize_threshold(self, y_true, y_pred_proba):\n        \"\"\"Optimize threshold\"\"\"\n        best_f1, best_thresh = 0, 0.27\n        for thresh in np.arange(0.15, 0.5, 0.05):\n            f1 = f1_score(y_true, y_pred_proba >= thresh, zero_division=0)\n            if f1 > best_f1:\n                best_f1, best_thresh = f1, thresh\n        return best_thresh, best_f1\n    \n    def train_in_chunks(self, mode_type: str, body_parts: list, train_dataset: pd.DataFrame):\n        \"\"\"Train with chunk-based processing\"\"\"\n        \n        print(f\"  Training {mode_type} models with advanced neural networks...\")\n        \n        # Collect actions and videos\n        all_actions = set()\n        video_ids = []\n        for switch, data, meta, label, vid_id in generate_mouse_data(train_dataset, 'train',\n                                                                      generate_single=(mode_type == 'single'),\n                                                                      generate_pair=(mode_type == 'pair')):\n            if switch == mode_type:\n                all_actions.update(label.columns)\n                video_ids.append(vid_id)\n                del data, meta, label\n                gc.collect()\n        \n        print(f\"    Found {len(all_actions)} actions, {len(set(video_ids))} videos\")\n        \n        if len(set(video_ids)) == 0:\n            return\n        \n        self.split_videos(set(video_ids))\n        \n        # Train each action\n        for action in all_actions:\n            if verbose:\n                print(f\"    Training: {action}\")\n            \n            self.replay_buffers[action] = DiverseReplayBuffer(config.replay_buffer_size, config.diversity_k)\n            \n            chunk_count = 0\n            val_X_list, val_y_list = [], []\n            ewc_model = None\n            \n            for switch, data, meta, label, vid_id in generate_mouse_data(train_dataset, 'train',\n                                                                          generate_single=(mode_type == 'single'),\n                                                                          generate_pair=(mode_type == 'pair')):\n                if switch != mode_type or action not in label.columns:\n                    continue\n                \n                try:\n                    # Extract features\n                    if mode_type == 'single':\n                        X = transform_single(data, body_parts)\n                    else:\n                        X = transform_pair(data, body_parts)\n                    \n                    del data\n                    gc.collect()\n                    \n                    # Get labels\n                    mask = ~label[action].isna().values\n                    X_action = X[mask]\n                    y_action = label[action][mask].values.astype(int)\n                    \n                    if len(X_action) == 0:\n                        continue\n                    \n                    # Store feature columns\n                    if action not in self.feature_columns:\n                        self.feature_columns[action] = X_action.columns.tolist()\n                        if not self.imputer_fitted:\n                            self.imputer.fit(X_action)\n                            self.imputer_fitted = True\n                    else:\n                        X_action = X_action.reindex(columns=self.feature_columns[action], fill_value=0)\n                    \n                    # Impute\n                    X_imputed = pd.DataFrame(\n                        self.imputer.transform(X_action),\n                        columns=self.feature_columns[action],\n                        index=X_action.index\n                    )\n                    \n                    # Split train/val\n                    if vid_id in self.val_videos:\n                        val_X_list.append(X_imputed)\n                        val_y_list.append(y_action)\n                        del X, X_action, X_imputed\n                        gc.collect()\n                        continue\n                    \n                    if vid_id not in self.train_videos:\n                        continue\n                    \n                    # Initialize model\n                    if ewc_model is None:\n                        input_dim = len(self.feature_columns[action])\n                        ewc_model = EWCModel(\n                            input_dim=input_dim,\n                            mlp_hidden=config.mlp_hidden_dims,\n                            use_temporal_conv=config.use_temporal_conv,\n                            temporal_channels=config.temporal_conv_channels,\n                            use_lstm=config.use_lstm,\n                            lstm_hidden=config.lstm_hidden_size,\n                            lstm_layers=config.lstm_num_layers,\n                            use_attention=config.use_attention,\n                            attention_heads=config.attention_heads,\n                            dropout=config.dropout,\n                            learning_rate=config.learning_rate,\n                            ewc_lambda=config.ewc_lambda,\n                            device=device\n                        )\n                    \n                    # Get replay\n                    X_replay, y_replay = self.replay_buffers[action].sample(1500, self.feature_columns[action])  # REDUCED from 3000\n                    \n                    # Combine\n                    if X_replay is not None and len(X_replay) > 0:\n                        X_combined = pd.concat([X_imputed, X_replay], ignore_index=True)\n                        y_combined = np.concatenate([y_action, y_replay])\n                    else:\n                        X_combined = X_imputed\n                        y_combined = y_action\n                    \n                    # Train\n                    if len(val_X_list) > 0:\n                        val_X_concat = pd.concat(val_X_list, ignore_index=True)\n                        val_y_concat = np.concatenate(val_y_list)\n                        ewc_model.fit(\n                            X_combined.values, y_combined,\n                            val_X=val_X_concat.values, val_y=val_y_concat,\n                            epochs=config.epochs_per_chunk,\n                            batch_size=config.batch_size,\n                            early_stopping_patience=config.early_stopping_patience\n                        )\n                    else:\n                        ewc_model.fit(\n                            X_combined.values, y_combined,\n                            epochs=config.epochs_per_chunk,\n                            batch_size=config.batch_size\n                        )\n                    \n                    # Validate\n                    if len(val_X_list) > 0:\n                        try:\n                            val_pred = ewc_model.predict_proba(val_X_concat.values)[:, 1]\n                            _, val_f1 = self.optimize_threshold(val_y_concat, val_pred)\n                            self.convergence_history[action].append(val_f1)\n                            \n                            if verbose:\n                                print(f\"      Chunk {chunk_count+1}: Val F1={val_f1:.3f}\")\n                        except:\n                            pass\n                    \n                    # Add to replay\n                    self.replay_buffers[action].add(X_imputed, y_action, vid_id)\n                    \n                    del X, X_action, X_imputed, X_combined, y_combined\n                    gc.collect()\n                    \n                    chunk_count += 1\n                    \n                    if self.check_convergence(action):\n                        if verbose:\n                            print(f\"      Converged at chunk {chunk_count}\")\n                        break\n                    \n                    if chunk_count >= config.max_chunks_per_action:\n                        break\n                        \n                except Exception as e:\n                    if verbose:\n                        print(f\"      Error: {e}\")\n                    continue\n            \n            # Optimize threshold\n            if len(val_X_list) > 0 and ewc_model is not None:\n                val_X_concat = pd.concat(val_X_list, ignore_index=True)\n                val_y_concat = np.concatenate(val_y_list)\n                \n                try:\n                    val_pred = ewc_model.predict_proba(val_X_concat.values)[:, 1]\n                    best_thresh, best_f1 = self.optimize_threshold(val_y_concat, val_pred)\n                    self.action_thresholds[action] = best_thresh\n                    if verbose:\n                        print(f\"      Threshold: {best_thresh:.2f}, F1: {best_f1:.3f}, Tasks: {ewc_model.task_count}\")\n                except:\n                    self.action_thresholds[action] = 0.27\n            else:\n                self.action_thresholds[action] = 0.27\n            \n            if ewc_model is not None:\n                self.action_models[action] = ewc_model\n            \n            # Aggressive cleanup\n            del val_X_list, val_y_list\n            gc.collect()\n            if torch.cuda.is_available():\n                torch.cuda.empty_cache()\n    \n    def predict(self, mode_type: str, body_parts: list, test_dataset: pd.DataFrame) -> list:\n        \"\"\"Predict on test data\"\"\"\n        submission_parts = []\n        \n        for switch, data, meta, actions, vid_id in generate_mouse_data(test_dataset, 'test',\n                                                                        generate_single=(mode_type == 'single'),\n                                                                        generate_pair=(mode_type == 'pair')):\n            if switch != mode_type:\n                continue\n            \n            try:\n                if mode_type == 'single':\n                    X = transform_single(data, body_parts)\n                else:\n                    X = transform_pair(data, body_parts)\n                \n                del data\n                gc.collect()\n                \n                X_imputed = pd.DataFrame(\n                    self.imputer.transform(X),\n                    columns=X.columns,\n                    index=X.index\n                )\n                del X\n                gc.collect()\n                \n                probs = pd.DataFrame(index=meta.video_frame)\n                \n                for action in actions:\n                    if action in self.action_models:\n                        model = self.action_models[action]\n                        \n                        if action in self.feature_columns:\n                            cols = self.feature_columns[action]\n                            X_aligned = X_imputed.reindex(columns=cols, fill_value=0)\n                        else:\n                            X_aligned = X_imputed\n                        \n                        try:\n                            pred = model.predict_proba(X_aligned.values)\n                            probs[action] = pred[:, 1]\n                        except:\n                            pass\n                \n                del X_imputed\n                gc.collect()\n                \n                if probs.shape[1] > 0:\n                    sub_part = self._predictions_to_submission(probs, meta)\n                    submission_parts.append(sub_part)\n            \n            except Exception as e:\n                if verbose:\n                    print(f'  ERROR: {e}')\n        \n        return submission_parts\n    \n    def _predictions_to_submission(self, probs: pd.DataFrame, meta: pd.DataFrame) -> pd.DataFrame:\n        \"\"\"Convert to submission\"\"\"\n        ama = np.argmax(probs.values, axis=1)\n        max_probs = probs.values.max(axis=1)\n        thresh_array = np.array([self.action_thresholds.get(probs.columns[i], 0.27) for i in ama])\n        ama = np.where(max_probs >= thresh_array, ama, -1)\n        \n        ama = pd.Series(ama, index=meta.video_frame)\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': probs.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        return submission_part\n\n# ==================== ROBUSTIFY ====================\ndef robustify(submission, dataset):\n    \"\"\"Clean submission\"\"\"\n    if len(submission) == 0:\n        return submission\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_frame = 0\n        for i, (_, row) in enumerate(group.iterrows()):\n            if row['start_frame'] < last_stop_frame:\n                mask[i] = False\n            else:\n                last_stop_frame = row['stop_frame']\n        group_list.append(group[mask])\n    \n    submission = pd.concat(group_list) if group_list else pd.DataFrame()\n    submission = submission.reset_index(drop=True)\n    return submission\n\n# ==================== MAIN ====================\ndef main():\n    \"\"\"Main execution\"\"\"\n    print(\"Loading data...\")\n    train = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/train.csv')\n    test = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/test.csv')\n    \n    body_parts_list = list(np.unique(train.body_parts_tracked))\n    submission_list = []\n    \n    for section in range(1, len(body_parts_list)):\n        body_parts_tracked_str = body_parts_list[section]\n        \n        try:\n            body_parts_tracked = json.loads(body_parts_tracked_str)\n            print(f\"\\n{section}. Processing {len(body_parts_tracked)} body parts\")\n            \n            if len(body_parts_tracked) > 5:\n                body_parts_tracked = [b for b in body_parts_tracked if b not in config.drop_body_parts]\n            \n            train_subset = train[train.body_parts_tracked == body_parts_tracked_str]\n            test_subset = test[test.body_parts_tracked == body_parts_tracked_str]\n            \n            if len(test_subset) == 0:\n                print(\"  No test videos\")\n                continue\n            \n            # Single mouse\n            if len(train_subset) > 0:\n                manager_single = ChunkLearningManager()\n                manager_single.train_in_chunks('single', body_parts_tracked, train_subset)\n                \n                if len(manager_single.action_models) > 0:\n                    print(f\"  Predicting on {len(test_subset)} test videos\")\n                    parts = manager_single.predict('single', body_parts_tracked, test_subset)\n                    submission_list.extend(parts)\n                \n                del manager_single\n                gc.collect()\n                if torch.cuda.is_available():\n                    torch.cuda.empty_cache()\n            \n            # Pairs\n            if len(train_subset) > 0:\n                manager_pair = ChunkLearningManager()\n                manager_pair.train_in_chunks('pair', body_parts_tracked, train_subset)\n                \n                if len(manager_pair.action_models) > 0:\n                    print(f\"  Predicting on {len(test_subset)} test videos\")\n                    parts = manager_pair.predict('pair', body_parts_tracked, test_subset)\n                    submission_list.extend(parts)\n                \n                del manager_pair\n                gc.collect()\n                if torch.cuda.is_available():\n                    torch.cuda.empty_cache()\n        \n        except Exception as e:\n            print(f'Exception: {e}')\n            import traceback\n            traceback.print_exc()\n    \n    # Final submission\n    if len(submission_list) > 0:\n        submission = pd.concat(submission_list, ignore_index=True)\n        submission_robust = robustify(submission, test)\n    else:\n        submission_robust = pd.DataFrame({\n            'video_id': [438887472],\n            'agent_id': ['mouse1'],\n            'target_id': ['self'],\n            'action': ['rear'],\n            'start_frame': [278],\n            'stop_frame': [500]\n        })\n    \n    submission_robust.index.name = 'row_id'\n    submission_robust.to_csv('/kaggle/working/submission.csv')\n    print(f\"\\nFinal submission: {len(submission_robust)} predictions\")\n\nif __name__ == '__main__':\n    main()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}