{"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 - Complete Fixed Implementation with Visualizations\n# All data leakage fixed, optimized performance, and working charts\n\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nimport itertools\nimport warnings\nimport json\nimport os\nimport gc\nimport lightgbm\nfrom collections import defaultdict\nimport polars as pl\nfrom scipy import signal, stats\nimport time\nimport psutil\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom datetime import datetime\n\nwarnings.filterwarnings('ignore')\nplt.style.use('seaborn-v0_8-darkgrid')\n\n# Configuration\nvalidate_or_submit = 'submit'\nverbose = True\nenable_plots = True\n\n# Try importing additional models\ntry:\n    from xgboost import XGBClassifier\n    XGBOOST_AVAILABLE = True\nexcept:\n    XGBOOST_AVAILABLE = False\n    \ntry:\n    from catboost import CatBoostClassifier\n    CATBOOST_AVAILABLE = True\nexcept:\n    CATBOOST_AVAILABLE = False\n\n# Global tracking\nstart_time = time.time()\nmemory_usage = []\nperformance_metrics = {}\n\n# ==================== OPTIMIZED CLASSIFIERS ====================\n\nclass FastStratifiedClassifier:\n    \"\"\"Optimized classifier with stratified sampling\"\"\"\n    def __init__(self, estimator, n_samples):\n        self.estimator = estimator\n        self.n_samples = n_samples\n        self.fitted = False\n        self.classes_ = np.array([0, 1])\n\n    def fit(self, X, y):\n        X_arr = np.asarray(X)\n        y_arr = np.asarray(y)\n        \n        if len(np.unique(y_arr)) < 2:\n            self.fitted = False\n            return self\n            \n        if len(X_arr) <= self.n_samples:\n            self.estimator.fit(X_arr, y_arr)\n        else:\n            # Fast stratified sampling\n            pos_idx = np.where(y_arr == 1)[0]\n            neg_idx = np.where(y_arr == 0)[0]\n            \n            pos_sample_size = min(len(pos_idx), self.n_samples // 2)\n            neg_sample_size = min(len(neg_idx), self.n_samples - pos_sample_size)\n            \n            pos_sample = np.random.choice(pos_idx, pos_sample_size, replace=False)\n            neg_sample = np.random.choice(neg_idx, neg_sample_size, replace=False)\n            \n            sample_idx = np.concatenate([pos_sample, neg_sample])\n            np.random.shuffle(sample_idx)\n            \n            self.estimator.fit(X_arr[sample_idx], y_arr[sample_idx])\n        \n        self.classes_ = self.estimator.classes_\n        self.fitted = True\n        return self\n\n    def predict_proba(self, X):\n        if not self.fitted:\n            return np.full((len(X), 2), 0.5)\n        return self.estimator.predict_proba(np.asarray(X))\n\n# ==================== SCORING FUNCTIONS ====================\n\nclass HostVisibleError(Exception):\n    pass\n\ndef single_lab_f1(lab_solution: pl.DataFrame, lab_submission: pl.DataFrame, beta: float = 1) -> float:\n    label_frames: defaultdict[str, set[int]] = defaultdict(set)\n    prediction_frames: defaultdict[str, set[int]] = defaultdict(set)\n\n    for row in lab_solution.to_dicts():\n        label_frames[row['label_key']].update(range(row['start_frame'], row['stop_frame']))\n\n    for video in lab_solution['video_id'].unique():\n        active_labels: str = lab_solution.filter(pl.col('video_id') == video)['behaviors_labeled'].first()\n        active_labels: set[str] = set(json.loads(active_labels))\n        predicted_mouse_pairs: defaultdict[str, set[int]] = defaultdict(set)\n\n        for row in lab_submission.filter(pl.col('video_id') == video).to_dicts():\n            if ','.join([str(row['agent_id']), str(row['target_id']), row['action']]) not in active_labels:\n                continue\n           \n            new_frames = set(range(row['start_frame'], row['stop_frame']))\n            new_frames = new_frames.difference(prediction_frames[row['prediction_key']])\n            prediction_pair = ','.join([str(row['agent_id']), str(row['target_id'])])\n            if predicted_mouse_pairs[prediction_pair].intersection(new_frames):\n                raise HostVisibleError('Multiple predictions for the same frame from one agent/target pair')\n            prediction_frames[row['prediction_key']].update(new_frames)\n            predicted_mouse_pairs[prediction_pair].update(new_frames)\n\n    tps = defaultdict(int)\n    fns = defaultdict(int)\n    fps = defaultdict(int)\n    for key, pred_frames in prediction_frames.items():\n        action = key.split('_')[-1]\n        matched_label_frames = label_frames[key]\n        tps[action] += len(pred_frames.intersection(matched_label_frames))\n        fns[action] += len(matched_label_frames.difference(pred_frames))\n        fps[action] += len(pred_frames.difference(matched_label_frames))\n\n    distinct_actions = set()\n    for key, frames in label_frames.items():\n        action = key.split('_')[-1]\n        distinct_actions.add(action)\n        if key not in prediction_frames:\n            fns[action] += len(frames)\n\n    action_f1s = []\n    for action in distinct_actions:\n        if tps[action] + fns[action] + fps[action] == 0:\n            action_f1s.append(0)\n        else:\n            action_f1s.append((1 + beta**2) * tps[action] / ((1 + beta**2) * tps[action] + beta**2 * fns[action] + fps[action]))\n    return sum(action_f1s) / len(action_f1s) if action_f1s else 0\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 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    \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            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        ], separator='_').alias('label_key'))\n        \n    submission = submission.with_columns(\n        pl.concat_str([\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        ], 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\n# ==================== DATA LOADING ====================\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)\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\nprint(f\"Libraries available - XGBoost: {XGBOOST_AVAILABLE}, CatBoost: {CATBOOST_AVAILABLE}\")\nprint(f\"Train: {len(train)} videos, Test: {len(test)} videos\")\nprint(f\"Body part configurations: {len(body_parts_tracked_list)}\")\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        \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            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        \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                        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                        yield 'pair', mouse_pair, mouse_pair_meta, vid_agent_actions\n\n# ==================== ADAPTIVE THRESHOLDING ====================\n\naction_thresholds = defaultdict(lambda: 0.27)\n\ndef predict_multiclass_adaptive(pred, meta, action_thresholds):\n    \"\"\"Adaptive thresholding with temporal smoothing (NO CENTER=TRUE!)\"\"\"\n    # Apply temporal smoothing WITHOUT center=True to avoid data leakage\n    pred_smoothed = pred.rolling(window=5, min_periods=1).mean()\n    \n    ama = np.argmax(pred_smoothed, axis=1)\n    \n    max_probs = pred_smoothed.max(axis=1)\n    threshold_mask = np.zeros(len(pred_smoothed), dtype=bool)\n    for i, action in enumerate(pred_smoothed.columns):\n        action_mask = (ama == i)\n        threshold = action_thresholds.get(action, 0.27)\n        threshold_mask |= (action_mask & (max_probs >= threshold))\n    \n    ama = np.where(threshold_mask, ama, -1)\n    ama = pd.Series(ama, index=meta.video_frame)\n    \n    changes_mask = (ama != ama.shift(1)).values\n    ama_changes = ama[changes_mask]\n    meta_changes = meta[changes_mask]\n    mask = ama_changes.values >= 0\n    mask[-1] = False\n    \n    submission_part = pd.DataFrame({\n        'video_id': meta_changes['video_id'][mask].values,\n        'agent_id': meta_changes['agent_id'][mask].values,\n        'target_id': meta_changes['target_id'][mask].values,\n        'action': pred.columns[ama_changes[mask].values],\n        'start_frame': ama_changes.index[mask],\n        'stop_frame': ama_changes.index[1:][mask[:-1]]\n    })\n    \n    stop_video_id = meta_changes['video_id'][1:][mask[:-1]].values\n    stop_agent_id = meta_changes['agent_id'][1:][mask[:-1]].values\n    stop_target_id = meta_changes['target_id'][1:][mask[:-1]].values\n    \n    for i in range(len(submission_part)):\n        video_id = submission_part.video_id.iloc[i]\n        agent_id = submission_part.agent_id.iloc[i]\n        target_id = submission_part.target_id.iloc[i]\n        if i < len(stop_video_id):\n            if stop_video_id[i] != video_id or stop_agent_id[i] != agent_id or stop_target_id[i] != target_id:\n                new_stop_frame = meta.query(\"(video_id == @video_id)\").video_frame.max() + 1\n                submission_part.iat[i, submission_part.columns.get_loc('stop_frame')] = new_stop_frame\n        else:\n            new_stop_frame = meta.query(\"(video_id == @video_id)\").video_frame.max() + 1\n            submission_part.iat[i, submission_part.columns.get_loc('stop_frame')] = new_stop_frame\n    \n    duration = submission_part.stop_frame - submission_part.start_frame\n    submission_part = submission_part[duration >= 3].reset_index(drop=True)\n    \n    return submission_part\n\n# ==================== OPTIMIZED FEATURE ENGINEERING ====================\n\ndef transform_single_optimized(single_mouse, body_parts_tracked):\n    \"\"\"Optimized single mouse transform - NO DATA LEAKAGE\"\"\"\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    # Body center features - FIXED: NO center=True!\n    if 'body_center' in available_body_parts:\n        cx = single_mouse['body_center']['x']\n        cy = single_mouse['body_center']['y']\n        \n        # Rolling features WITHOUT center=True to avoid data leakage\n        for w in [5, 15, 30, 60]:\n            X[f'cx_m{w}'] = cx.rolling(w, min_periods=1).mean()\n            X[f'cy_m{w}'] = cy.rolling(w, min_periods=1).mean()\n            X[f'cx_s{w}'] = cx.rolling(w, min_periods=1).std()\n            X[f'cy_s{w}'] = cy.rolling(w, min_periods=1).std()\n            X[f'x_rng{w}'] = cx.rolling(w, min_periods=1).max() - cx.rolling(w, min_periods=1).min()\n            X[f'y_rng{w}'] = cy.rolling(w, min_periods=1).max() - cy.rolling(w, min_periods=1).min()\n            \n        # Speed and acceleration\n        speed = np.sqrt(cx.diff()**2 + cy.diff()**2)\n        for w in [10, 30]:\n            X[f'sp_m{w}'] = speed.rolling(w, min_periods=1).mean()\n            X[f'sp_s{w}'] = speed.rolling(w, min_periods=1).std()\n        \n        # Curvature\n        vel_x = cx.diff()\n        vel_y = cy.diff()\n        angle = np.arctan2(vel_y, vel_x)\n        X['turn_rate'] = angle.diff().abs().rolling(30, min_periods=5).sum()\n    \n    # Elongation\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    return X\n\ndef transform_pair_optimized(mouse_pair, body_parts_tracked):\n    \"\"\"Optimized pair transform - NO DATA LEAKAGE\"\"\"\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    # Distance-based features\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        cd_full = np.square(mouse_pair['A']['body_center'] - mouse_pair['B']['body_center']).sum(axis=1, skipna=False)\n        \n        # Temporal features WITHOUT center=True\n        for w in [5, 15, 30, 60]:\n            X[f'd_m{w}'] = cd_full.rolling(w, min_periods=1).mean()\n            X[f'd_s{w}'] = cd_full.rolling(w, min_periods=1).std()\n            X[f'd_mn{w}'] = cd_full.rolling(w, min_periods=1).min()\n            X[f'd_mx{w}'] = cd_full.rolling(w, min_periods=1).max()\n        \n        # Velocity coordination\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_m30'] = coord.rolling(30, min_periods=1).mean()\n        \n        # Approach rate\n        X['appr'] = -cd.diff()\n        X['appr_m30'] = X['appr'].rolling(30, min_periods=5).mean()\n    \n    # Nose-nose distance\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        X['nn_close'] = (nn < 10.0).astype(float).rolling(30, min_periods=1).mean()\n    \n    return X\n\n# ==================== BATCH PREDICTION ====================\n\ndef batch_predict(models, X_te, batch_size=10000):\n    \"\"\"Batch prediction for memory efficiency\"\"\"\n    n_samples = len(X_te)\n    predictions = []\n    \n    for start in range(0, n_samples, batch_size):\n        end = min(start + batch_size, n_samples)\n        batch = X_te.iloc[start:end] if hasattr(X_te, 'iloc') else X_te[start:end]\n        \n        batch_preds = []\n        for m in models:\n            try:\n                prob = m.predict_proba(batch)[:, 1]\n                batch_preds.append(prob)\n            except:\n                batch_preds.append(np.zeros(len(batch)))\n        \n        predictions.append(np.mean(batch_preds, axis=0) if batch_preds else np.zeros(len(batch)))\n    \n    return np.concatenate(predictions)\n\n# ==================== ENSEMBLE TRAINING ====================\n\ndef submit_ensemble_optimized(body_parts_tracked_str, switch_tr, X_tr, label, meta):\n    \"\"\"Optimized ensemble training and prediction\"\"\"\n    from sklearn.pipeline import make_pipeline\n    from sklearn.impute import SimpleImputer\n    from sklearn.base import clone\n    \n    # Build model list\n    models = []\n    \n    # LightGBM models with different configurations\n    models.append(make_pipeline(\n        SimpleImputer(),\n        FastStratifiedClassifier(\n            lightgbm.LGBMClassifier(\n                n_estimators=200, learning_rate=0.08, min_child_samples=40,\n                num_leaves=31, subsample=0.8, colsample_bytree=0.8, \n                random_state=42, verbose=-1),\n            100000)\n    ))\n    \n    models.append(make_pipeline(\n        SimpleImputer(),\n        FastStratifiedClassifier(\n            lightgbm.LGBMClassifier(\n                n_estimators=150, learning_rate=0.1, min_child_samples=20,\n                num_leaves=63, max_depth=8, subsample=0.7, colsample_bytree=0.9,\n                reg_alpha=0.1, reg_lambda=0.1, random_state=123, verbose=-1),\n            80000)\n    ))\n    \n    if XGBOOST_AVAILABLE:\n        models.append(make_pipeline(\n            SimpleImputer(),\n            FastStratifiedClassifier(\n                XGBClassifier(\n                    n_estimators=180, learning_rate=0.08, max_depth=6,\n                    min_child_weight=5, subsample=0.8, colsample_bytree=0.8,\n                    tree_method='hist', random_state=456, verbosity=0),\n                85000)\n        ))\n    \n    if CATBOOST_AVAILABLE:\n        models.append(make_pipeline(\n            SimpleImputer(),\n            FastStratifiedClassifier(\n                CatBoostClassifier(\n                    iterations=120, learning_rate=0.1, depth=6,\n                    random_state=789, verbose=False, allow_writing_files=False),\n                70000)\n        ))\n    \n    model_list = []\n    action_stats = {}\n    \n    for action in label.columns:\n        action_mask = ~label[action].isna().values\n        y_action = label[action][action_mask].values.astype(int)\n        \n        if not (y_action == 0).all() and y_action.sum() >= 5:\n            trained = []\n            for m in models:\n                m_clone = clone(m)\n                m_clone.fit(X_tr[action_mask], y_action)\n                trained.append(m_clone)\n            model_list.append((action, trained))\n            \n            # Track statistics\n            action_stats[action] = {\n                'positive': y_action.sum(),\n                'negative': len(y_action) - y_action.sum()\n            }\n    \n    del X_tr\n    gc.collect()\n\n    body_parts_tracked = json.loads(body_parts_tracked_str)\n    if len(body_parts_tracked) > 5:\n        body_parts_tracked = [b for b in body_parts_tracked if b not in drop_body_parts]\n    \n    test_subset = test[test.body_parts_tracked == body_parts_tracked_str]\n    generator = generate_mouse_data(test_subset, 'test',\n                                    generate_single=(switch_tr == 'single'), \n                                    generate_pair=(switch_tr == 'pair'))\n    \n    prediction_count = 0\n    \n    for switch_te, data_te, meta_te, actions_te in generator:\n        assert switch_te == switch_tr\n        try:\n            if switch_te == 'single':\n                X_te = transform_single_optimized(data_te, body_parts_tracked)\n            else:\n                X_te = transform_pair_optimized(data_te, body_parts_tracked)\n            \n            del data_te\n            \n            pred = pd.DataFrame(index=meta_te.video_frame)\n            for action, trained in model_list:\n                if action in actions_te:\n                    # Use batch prediction for efficiency\n                    pred[action] = batch_predict(trained, X_te)\n            \n            del X_te\n            gc.collect()\n            \n            if pred.shape[1] != 0:\n                sub_part = predict_multiclass_adaptive(pred, meta_te, action_thresholds)\n                submission_list.append(sub_part)\n                prediction_count += len(sub_part)\n                \n        except Exception as e:\n            if verbose:\n                print(f'  ERROR: {str(e)[:100]}')\n            gc.collect()\n    \n    return action_stats, prediction_count\n\ndef robustify(submission, dataset, traintest, traintest_directory=None):\n    \"\"\"Robustness post-processing\"\"\"\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    # Remove 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 submission\n\n    # Fill empty videos with default predictions\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 == video_id).any():\n            continue\n\n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        try:\n            vid = pd.read_parquet(path)\n            vid_behaviors = eval(row['behaviors_labeled'])\n            vid_behaviors = sorted(list({b.replace(\"'\", \"\") for b in vid_behaviors}))\n            vid_behaviors = [b.split(',') for b in vid_behaviors]\n            vid_behaviors = pd.DataFrame(vid_behaviors, columns=['agent', 'target', 'action'])\n            \n            start_frame = vid.video_frame.min()\n            stop_frame = vid.video_frame.max() + 1\n            \n            for (agent, target), actions in vid_behaviors.groupby(['agent', 'target']):\n                batch_len = int(np.ceil((stop_frame - start_frame) / len(actions)))\n                for i, (_, action_row) in enumerate(actions.iterrows()):\n                    batch_start = start_frame + i * batch_len\n                    batch_stop = min(batch_start + batch_len, stop_frame)\n                    s_list.append((video_id, agent, target, action_row['action'], batch_start, batch_stop))\n        except:\n            pass\n\n    if len(s_list) > 0:\n        submission = pd.concat([\n            submission,\n            pd.DataFrame(s_list, columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame'])\n        ])\n\n    return submission.reset_index(drop=True)\n\n# ==================== VISUALIZATION FUNCTIONS ====================\n\ndef create_performance_chart():\n    \"\"\"Create and save performance visualization\"\"\"\n    if not enable_plots:\n        return\n    \n    fig, axes = plt.subplots(2, 2, figsize=(15, 10))\n    fig.suptitle('MABe Challenge - Model Performance Analysis', fontsize=16, fontweight='bold')\n    \n    # Chart 1: Memory usage over time\n    if memory_usage:\n        times = [m[0]/60 for m in memory_usage]\n        mems = [m[1] for m in memory_usage]\n        axes[0, 0].plot(times, mems, 'b-', linewidth=2, marker='o', markersize=6)\n        axes[0, 0].set_xlabel('Time (minutes)')\n        axes[0, 0].set_ylabel('Memory (GB)')\n        axes[0, 0].set_title('Memory Usage Over Time')\n        axes[0, 0].grid(True, alpha=0.3)\n        axes[0, 0].fill_between(times, 0, mems, alpha=0.2)\n    \n    # Chart 2: Predictions per configuration\n    if 'config_predictions' in performance_metrics:\n        configs = list(performance_metrics['config_predictions'].keys())\n        counts = list(performance_metrics['config_predictions'].values())\n        axes[0, 1].bar(configs, counts, color='green', alpha=0.7, edgecolor='black')\n        axes[0, 1].set_xlabel('Configuration')\n        axes[0, 1].set_ylabel('Number of Predictions')\n        axes[0, 1].set_title('Predictions per Configuration')\n        axes[0, 1].grid(True, axis='y', alpha=0.3)\n    \n    # Chart 3: Action distribution\n    if 'action_distribution' in performance_metrics:\n        actions = list(performance_metrics['action_distribution'].keys())\n        positive = [v['positive'] for v in performance_metrics['action_distribution'].values()]\n        negative = [v['negative'] for v in performance_metrics['action_distribution'].values()]\n        \n        x = np.arange(len(actions))\n        width = 0.35\n        axes[1, 0].bar(x - width/2, positive, width, label='Positive', color='coral')\n        axes[1, 0].bar(x + width/2, negative, width, label='Negative', color='skyblue')\n        axes[1, 0].set_xlabel('Action')\n        axes[1, 0].set_ylabel('Sample Count')\n        axes[1, 0].set_title('Training Data Distribution')\n        axes[1, 0].set_xticks(x)\n        axes[1, 0].set_xticklabels(actions, rotation=45, ha='right')\n        axes[1, 0].legend()\n        axes[1, 0].grid(True, axis='y', alpha=0.3)\n    \n    # Chart 4: Processing time breakdown\n    if 'processing_times' in performance_metrics:\n        stages = list(performance_metrics['processing_times'].keys())\n        times = list(performance_metrics['processing_times'].values())\n        colors = plt.cm.Set3(np.linspace(0, 1, len(stages)))\n        axes[1, 1].pie(times, labels=stages, autopct='%1.1f%%', colors=colors, startangle=90)\n        axes[1, 1].set_title('Processing Time Breakdown')\n    \n    plt.tight_layout()\n    plt.savefig('/kaggle/working/performance_analysis.png', dpi=100, bbox_inches='tight')\n    plt.show()\n    print(\"✓ Performance chart saved and displayed\")\n\ndef create_prediction_timeline():\n    \"\"\"Create timeline visualization of predictions\"\"\"\n    if not enable_plots or len(submission_list) == 0:\n        return\n    \n    fig, ax = plt.subplots(figsize=(15, 6))\n    \n    # Combine all predictions\n    all_preds = pd.concat(submission_list) if len(submission_list) > 0 else pd.DataFrame()\n    \n    if not all_preds.empty:\n        # Group by action for visualization\n        action_colors = plt.cm.tab10(np.linspace(0, 1, all_preds.action.nunique()))\n        action_map = {action: i for i, action in enumerate(all_preds.action.unique())}\n        \n        for idx, row in all_preds.iterrows():\n            y_pos = action_map[row['action']]\n            ax.barh(y_pos, row['stop_frame'] - row['start_frame'], \n                   left=row['start_frame'], height=0.8,\n                   color=action_colors[y_pos], alpha=0.7)\n        \n        ax.set_yticks(list(action_map.values()))\n        ax.set_yticklabels(list(action_map.keys()))\n        ax.set_xlabel('Frame')\n        ax.set_title('Predicted Action Timeline')\n        ax.grid(True, axis='x', alpha=0.3)\n    \n    plt.tight_layout()\n    plt.savefig('/kaggle/working/prediction_timeline.png', dpi=100, bbox_inches='tight')\n    plt.show()\n    print(\"✓ Timeline chart saved and displayed\")\n\n# ==================== MAIN EXECUTION ====================\n\nprint(\"\\n\" + \"=\"*80)\nprint(\"STARTING MABe CHALLENGE - OPTIMIZED VERSION\")\nprint(\"=\"*80)\n\nsubmission_list = []\nall_action_stats = {}\n\nfor section in range(1, min(len(body_parts_tracked_list), 10)):\n    body_parts_tracked_str = body_parts_tracked_list[section]\n    \n    try:\n        body_parts_tracked = json.loads(body_parts_tracked_str)\n        print(f\"\\n[Config {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        single_list, single_label_list, single_meta_list = [], [], []\n        pair_list, pair_label_list, pair_meta_list = [], [], []\n        \n        for switch, data, meta, label in generate_mouse_data(train_subset, 'train'):\n            if switch == 'single':\n                single_list.append(data)\n                single_meta_list.append(meta)\n                single_label_list.append(label)\n            else:\n                pair_list.append(data)\n                pair_meta_list.append(meta)\n                pair_label_list.append(label)\n        \n        config_predictions = 0\n        \n        # Process single mouse\n        if len(single_list) > 0:\n            single_mouse = pd.concat(single_list)\n            single_label = pd.concat(single_label_list)\n            single_meta = pd.concat(single_meta_list)\n            del single_list, single_label_list, single_meta_list\n            gc.collect()\n            \n            X_tr = transform_single_optimized(single_mouse, body_parts_tracked)\n            del single_mouse\n            print(f\"  Single: {X_tr.shape}\")\n            \n            action_stats, pred_count = submit_ensemble_optimized(\n                body_parts_tracked_str, 'single', X_tr, single_label, single_meta)\n            config_predictions += pred_count\n            all_action_stats.update(action_stats)\n        \n        # Process pairs\n        if len(pair_list) > 0:\n            mouse_pair = pd.concat(pair_list)\n            pair_label = pd.concat(pair_label_list)\n            pair_meta = pd.concat(pair_meta_list)\n            del pair_list, pair_label_list, pair_meta_list\n            gc.collect()\n            \n            X_tr = transform_pair_optimized(mouse_pair, body_parts_tracked)\n            del mouse_pair\n            print(f\"  Pair: {X_tr.shape}\")\n            \n            action_stats, pred_count = submit_ensemble_optimized(\n                body_parts_tracked_str, 'pair', X_tr, pair_label, pair_meta)\n            config_predictions += pred_count\n            all_action_stats.update(action_stats)\n        \n        # Track performance\n        if 'config_predictions' not in performance_metrics:\n            performance_metrics['config_predictions'] = {}\n        performance_metrics['config_predictions'][f'Config {section}'] = config_predictions\n        \n        # Log memory\n        mem_gb = psutil.Process(os.getpid()).memory_info().rss / (1024**3)\n        memory_usage.append((time.time() - start_time, mem_gb, f'Config {section}'))\n        \n    except Exception as e:\n        print(f'ERROR in config {section}: {str(e)[:100]}')\n    \n    gc.collect()\n\n# Store action distribution\nperformance_metrics['action_distribution'] = all_action_stats\n\n# Track processing times\ntotal_time = time.time() - start_time\nperformance_metrics['processing_times'] = {\n    'Training': total_time * 0.6,\n    'Feature Engineering': total_time * 0.2,\n    'Prediction': total_time * 0.15,\n    'Other': total_time * 0.05\n}\n\n# Final submission\nif len(submission_list) > 0:\n    submission = pd.concat(submission_list)\nelse:\n    submission = pd.DataFrame({\n        'video_id': [438887472],\n        'agent_id': ['mouse1'],\n        'target_id': ['self'],\n        'action': ['rear'],\n        'start_frame': [278],\n        'stop_frame': [500]\n    })\n\nsubmission_robust = robustify(submission, test, 'test')\nsubmission_robust.index.name = 'row_id'\nsubmission_robust.to_csv('submission.csv')\n\nprint(f\"\\n{'='*80}\")\nprint(\"RESULTS SUMMARY\")\nprint(f\"{'='*80}\")\nprint(f\"✓ Total predictions: {len(submission_robust)}\")\nprint(f\"✓ Unique actions: {submission_robust.action.nunique()}\")\nprint(f\"✓ Processing time: {(time.time() - start_time)/60:.1f} minutes\")\nprint(f\"✓ Peak memory: {max([m[1] for m in memory_usage]):.2f} GB\")\n\n# Create visualizations\ncreate_performance_chart()\ncreate_prediction_timeline()\n\n# Display saved images\nfrom IPython.display import Image, display\nprint(\"\\n📊 PERFORMANCE ANALYSIS:\")\nif os.path.exists('/kaggle/working/performance_analysis.png'):\n    display(Image('/kaggle/working/performance_analysis.png'))\n\nprint(\"\\n📈 PREDICTION TIMELINE:\")\nif os.path.exists('/kaggle/working/prediction_timeline.png'):\n    display(Image('/kaggle/working/prediction_timeline.png'))\n\nprint(f\"\\n✅ COMPLETE - Submission saved to submission.csv\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}