{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":59156,"databundleVersionId":13719397,"sourceType":"competition"}],"dockerImageVersionId":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-09-18T18:23:35.074498Z","iopub.execute_input":"2025-09-18T18:23:35.074809Z","iopub.status.idle":"2025-09-18T18:23:45.527543Z","shell.execute_reply.started":"2025-09-18T18:23:35.074783Z","shell.execute_reply":"2025-09-18T18:23:45.525868Z"},"_kg_hide-output":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"The following code is a work in progress. It has incorporated a variety of methods to solve this problem along with various additions. The goal is to combine multiple methodologies to create a hybrid / ensemble methodology. To this end, it is worth noting the following references:\n\nhttps://www.kaggle.com/code/ambrosm/mabe-nearest-neighbors-the-original\n\nhttps://www.kaggle.com/code/snnguynvnk19hl/mabe-challenge-ml\n\nhttps://www.kaggle.com/code/xbar19/mabe-catboost-baseline\n\nhttps://www.kaggle.com/code/yusuketogashi/mabe-launchpad-start-simple-localcv","metadata":{}},{"cell_type":"code","source":"\"\"\"\nMABe Challenge - Complete Implementation with Comprehensive Visualizations\n==========================================================================\nFull working pipeline with enhanced features, robust training, and detailed analytics\n\"\"\"\n\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom tqdm import tqdm\nimport itertools\nimport warnings\nimport json\nimport os\nimport gc\nimport lightgbm as lgb\nfrom collections import defaultdict, Counter\nfrom datetime import datetime\nimport time\nimport polars as pl\n\nimport matplotlib.gridspec as gridspec\nfrom matplotlib.patches import Rectangle\nimport matplotlib.patches as mpatches\n\nfrom sklearn.base import ClassifierMixin, BaseEstimator, clone\nfrom sklearn.model_selection import cross_val_predict, GroupKFold\nfrom sklearn.pipeline import make_pipeline\nfrom sklearn.impute import SimpleImputer\nfrom sklearn.preprocessing import StandardScaler\nfrom sklearn.metrics import f1_score\n\nwarnings.filterwarnings('ignore')\n\n# Enhanced visualization settings\nplt.style.use('seaborn-v0_8-darkgrid')\nsns.set_palette(\"husl\")\nplt.rcParams.update({\n    'figure.figsize': (14, 8),\n    'font.size': 11,\n    'axes.titlesize': 13,\n    'axes.labelsize': 11,\n    'xtick.labelsize': 10,\n    'ytick.labelsize': 10,\n    'legend.fontsize': 10,\n    'figure.titlesize': 15\n})\n\n# Configuration\nvalidate_or_submit = 'submit'\nverbose = True\ncreate_visualizations = True\nsave_plots = True\nplot_dir = '/kaggle/working/visualizations'\n\nif save_plots:\n    os.makedirs(plot_dir, exist_ok=True)\n    print(f\"📁 Created visualization directory: {plot_dir}\")\n\n# Performance tracking\nperformance_metrics = {\n    'configurations_processed': 0,\n    'single_mouse_batches': 0,\n    'pair_batches': 0,\n    'features_extracted': 0,\n    'predictions_made': 0,\n    'models_trained': 0,\n    'actions_processed': 0\n}\n\nprint(\"\\n\" + \"=\"*80)\nprint(\" \"*15 + \"🐭 MABe CHALLENGE - COMPLETE IMPLEMENTATION 🐭\")\nprint(\"=\"*80)\nprint(f\"📅 Start Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\")\nprint(f\"🔧 Mode: {validate_or_submit.upper()}\")\nprint(f\"📊 Visualizations: {'ENABLED' if create_visualizations else 'DISABLED'}\")\nprint(\"=\"*80 + \"\\n\")\n\n# Custom classifier for training on subset\nclass TrainOnSubsetClassifier(ClassifierMixin, BaseEstimator):\n    def __init__(self, estimator, n_samples):\n        self.estimator = estimator\n        self.n_samples = n_samples\n\n    def fit(self, X, y):\n        downsample = max(1, len(X) // self.n_samples)\n        self.estimator.fit(\n            np.array(X, copy=False)[::downsample],\n            np.array(y, copy=False)[::downsample]\n        )\n        self.classes_ = self.estimator.classes_\n        return self\n\n    def predict_proba(self, X):\n        if len(self.classes_) == 1:\n            return np.full((len(X), 1), 1.0)\n        return self.estimator.predict_proba(np.array(X))\n        \n    def predict(self, X):\n        return self.estimator.predict(np.array(X))\n\n# F-Beta 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)) if active_labels else set()\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            \n            if predicted_mouse_pairs[prediction_pair].intersection(new_frames):\n                raise HostVisibleError('Multiple predictions for the same frame')\n            \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    \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] / \n                            ((1 + beta**2) * tps[action] + beta**2 * fns[action] + fps[action]))\n    \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    \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(solution)\n    submission = pl.DataFrame(submission)\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    \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\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) if lab_scores else 0\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\ndef create_data_overview_plots(train_df, test_df):\n    \"\"\"Create comprehensive data overview visualizations\"\"\"\n    print(\"\\n📊 Creating Data Overview Visualizations...\")\n    \n    fig = plt.figure(figsize=(20, 14))\n    gs = gridspec.GridSpec(4, 4, hspace=0.3, wspace=0.3)\n    \n    # 1. Videos per laboratory\n    ax1 = fig.add_subplot(gs[0, :2])\n    lab_counts = train_df['lab_id'].value_counts().head(20)\n    colors = sns.color_palette(\"husl\", len(lab_counts))\n    bars = ax1.bar(range(len(lab_counts)), lab_counts.values, color=colors, edgecolor='black', linewidth=0.5)\n    ax1.set_xticks(range(len(lab_counts)))\n    ax1.set_xticklabels(lab_counts.index, rotation=45, ha='right', fontsize=9)\n    ax1.set_title('Videos per Laboratory (Top 20)', fontsize=12, fontweight='bold')\n    ax1.set_ylabel('Number of Videos')\n    ax1.grid(True, alpha=0.3, axis='y')\n    \n    for bar in bars:\n        height = bar.get_height()\n        ax1.text(bar.get_x() + bar.get_width()/2., height + 1,\n                f'{int(height)}', ha='center', va='bottom', fontsize=8)\n    \n    # 2. Mice distribution\n    ax2 = fig.add_subplot(gs[0, 2])\n    mice_dist = train_df['n_mice'].value_counts().sort_index()\n    wedges, texts, autotexts = ax2.pie(mice_dist.values, \n                                        labels=[f'{i} mice' for i in mice_dist.index],\n                                        autopct='%1.1f%%',\n                                        startangle=90,\n                                        colors=plt.cm.Set3.colors)\n    ax2.set_title('Mice per Video Distribution', fontsize=12, fontweight='bold')\n    \n    # 3. Body parts configuration\n    ax3 = fig.add_subplot(gs[0, 3])\n    body_parts_count = train_df['body_parts_tracked'].apply(lambda x: len(json.loads(x)) if pd.notna(x) else 0)\n    ax3.hist(body_parts_count, bins=20, color='skyblue', edgecolor='black', alpha=0.7)\n    ax3.set_xlabel('Number of Body Parts')\n    ax3.set_ylabel('Frequency')\n    ax3.set_title('Body Parts per Configuration', fontsize=12, fontweight='bold')\n    ax3.axvline(body_parts_count.mean(), color='red', linestyle='--', \n               label=f'Mean: {body_parts_count.mean():.1f}')\n    ax3.legend()\n    ax3.grid(True, alpha=0.3)\n    \n    # 4. Behavior types distribution\n    ax4 = fig.add_subplot(gs[1, :2])\n    behaviors_count = train_df['behaviors_labeled'].apply(lambda x: len(json.loads(x)) if pd.notna(x) else 0)\n    valid_behaviors = behaviors_count[behaviors_count > 0]\n    \n    ax4.hist(valid_behaviors, bins=30, color='coral', edgecolor='black', alpha=0.7)\n    ax4.set_xlabel('Number of Behavior Types')\n    ax4.set_ylabel('Frequency')\n    ax4.set_title('Behavior Types per Video', fontsize=12, fontweight='bold')\n    ax4.axvline(valid_behaviors.mean(), color='blue', linestyle='--',\n               label=f'Mean: {valid_behaviors.mean():.1f}')\n    ax4.axvline(valid_behaviors.median(), color='green', linestyle='--',\n               label=f'Median: {valid_behaviors.median():.1f}')\n    ax4.legend()\n    ax4.grid(True, alpha=0.3)\n    \n    # 5. Top behavior actions\n    ax5 = fig.add_subplot(gs[1, 2:])\n    all_behaviors = []\n    for behaviors_str in train_df['behaviors_labeled'].dropna():\n        try:\n            behaviors = json.loads(behaviors_str)\n            for b in behaviors:\n                parts = b.replace(\"'\", \"\").split(',')\n                if len(parts) == 3:\n                    all_behaviors.append(parts[2])\n        except:\n            continue\n    \n    behavior_counts = Counter(all_behaviors)\n    top_behaviors = dict(behavior_counts.most_common(15))\n    \n    ax5.barh(range(len(top_behaviors)), list(top_behaviors.values()), \n            color='steelblue', edgecolor='black', linewidth=0.5)\n    ax5.set_yticks(range(len(top_behaviors)))\n    ax5.set_yticklabels(list(top_behaviors.keys()), fontsize=9)\n    ax5.set_xlabel('Frequency')\n    ax5.set_title('Top 15 Behavior Actions', fontsize=12, fontweight='bold')\n    ax5.grid(True, alpha=0.3, axis='x')\n    \n    # 6. Training vs Test comparison\n    ax6 = fig.add_subplot(gs[2, :2])\n    comparison_data = pd.DataFrame({\n        'Dataset': ['Training', 'Test'],\n        'Videos': [len(train_df), len(test_df)],\n        'Unique Labs': [train_df['lab_id'].nunique(), test_df['lab_id'].nunique()]\n    })\n    \n    x = np.arange(len(comparison_data['Dataset']))\n    width = 0.35\n    \n    ax6.bar(x - width/2, comparison_data['Videos'], width, label='Videos', color='skyblue')\n    ax6.bar(x + width/2, comparison_data['Unique Labs'], width, label='Unique Labs', color='coral')\n    ax6.set_xticks(x)\n    ax6.set_xticklabels(comparison_data['Dataset'])\n    ax6.set_ylabel('Count')\n    ax6.set_title('Training vs Test Dataset Comparison', fontsize=12, fontweight='bold')\n    ax6.legend()\n    ax6.grid(True, alpha=0.3, axis='y')\n    \n    # 7. Video duration distribution\n    ax7 = fig.add_subplot(gs[2, 2:])\n    if 'video_duration_sec' in train_df.columns:\n        duration_data = train_df['video_duration_sec'].dropna()\n        ax7.hist(duration_data, bins=30, color='lightgreen', edgecolor='black', alpha=0.7)\n        ax7.set_xlabel('Duration (seconds)')\n        ax7.set_ylabel('Frequency')\n        ax7.set_title('Video Duration Distribution', fontsize=12, fontweight='bold')\n        ax7.axvline(duration_data.mean(), color='red', linestyle='--',\n                   label=f'Mean: {duration_data.mean():.1f}s')\n        ax7.legend()\n    ax7.grid(True, alpha=0.3)\n    \n    # 8. Statistics summary\n    ax8 = fig.add_subplot(gs[3, :])\n    ax8.axis('off')\n    \n    stats_text = f\"\"\"\nDATASET STATISTICS SUMMARY\n{'='*70}\n\nTRAINING DATA:\n- Total Videos: {len(train_df):,}\n- Videos with Missing Behaviors: {train_df['behaviors_labeled'].isna().sum()} ({100*train_df['behaviors_labeled'].isna().sum()/len(train_df):.1f}%)\n- Unique Laboratories: {train_df['lab_id'].nunique()}\n- Average Mice per Video: {train_df['n_mice'].mean():.2f}\n- Total Unique Body Part Configurations: {train_df['body_parts_tracked'].nunique()}\n\nTEST DATA:\n- Total Videos: {len(test_df):,}\n- Unique Laboratories: {test_df['lab_id'].nunique()}\n\nBEHAVIOR ANALYSIS:\n- Total Unique Behaviors: {len(behavior_counts)}\n- Most Common Behavior: {list(top_behaviors.keys())[0]} ({list(top_behaviors.values())[0]:,} occurrences)\n- Average Behaviors per Video: {valid_behaviors.mean():.1f}\n- Max Behaviors in Single Video: {behaviors_count.max()}\n\"\"\"\n    \n    ax8.text(0.05, 0.95, stats_text, transform=ax8.transAxes, fontsize=10,\n            verticalalignment='top', fontfamily='monospace',\n            bbox=dict(boxstyle='round', facecolor='lightgray', alpha=0.3))\n    \n    plt.suptitle('MABe Dataset Comprehensive Overview', fontsize=16, fontweight='bold')\n    \n    if save_plots:\n        plt.savefig(os.path.join(plot_dir, 'data_overview.png'), dpi=100, bbox_inches='tight')\n        print(\"    💾 Saved: data_overview.png\")\n    \n    plt.show()\n    \n    return behavior_counts\n\ndef create_feature_analysis_plots(X_data, feature_type='single', config_num=1):\n    \"\"\"Create feature distribution analysis\"\"\"\n    if X_data is None or X_data.empty:\n        return\n        \n    print(f\"    📈 Creating feature analysis for {feature_type} mouse (config {config_num})...\")\n    \n    # Select subset of features for visualization\n    feature_cols = X_data.columns[:min(20, len(X_data.columns))]\n    \n    fig, axes = plt.subplots(4, 5, figsize=(20, 12))\n    fig.suptitle(f'Feature Analysis - {feature_type.capitalize()} Mouse (Config {config_num})', \n                fontsize=14, fontweight='bold')\n    \n    for idx, (ax, col) in enumerate(zip(axes.flat, feature_cols)):\n        data = X_data[col].dropna()\n        ax.hist(data, bins=30, color='lightblue', edgecolor='black', alpha=0.7)\n        ax.set_title(col[:25], fontsize=9)\n        ax.set_xlabel('Value', fontsize=8)\n        ax.set_ylabel('Frequency', fontsize=8)\n        \n        # Add statistics\n        if len(data) > 0:\n            ax.axvline(data.mean(), color='red', linestyle='--', alpha=0.5, linewidth=1)\n            ax.axvline(data.median(), color='green', linestyle='--', alpha=0.5, linewidth=1)\n    \n    # Hide extra subplots\n    for idx in range(len(feature_cols), len(axes.flat)):\n        axes.flat[idx].axis('off')\n    \n    plt.tight_layout()\n    \n    if save_plots:\n        plt.savefig(os.path.join(plot_dir, f'feature_analysis_{feature_type}_{config_num}.png'), \n                   dpi=100, bbox_inches='tight')\n        print(f\"        💾 Saved: feature_analysis_{feature_type}_{config_num}.png\")\n    \n    plt.show()\n\ndef create_performance_summary_plot(f1_scores, predictions_per_config):\n    \"\"\"Create comprehensive performance summary\"\"\"\n    if not f1_scores:\n        return\n        \n    print(\"\\n🎯 Creating Performance Summary...\")\n    \n    f1_df = pd.DataFrame(f1_scores, columns=['config', 'action', 'f1_score'])\n    \n    fig = plt.figure(figsize=(20, 12))\n    \n    # 1. F1 Scores by Action\n    ax1 = plt.subplot(3, 3, 1)\n    action_f1 = f1_df.groupby('action')['f1_score'].agg(['mean', 'std']).sort_values('mean', ascending=False).head(15)\n    ax1.barh(range(len(action_f1)), action_f1['mean'], xerr=action_f1['std'].fillna(0),\n            color='green', alpha=0.7, capsize=3)\n    ax1.set_yticks(range(len(action_f1)))\n    ax1.set_yticklabels(action_f1.index, fontsize=9)\n    ax1.set_xlabel('F1 Score')\n    ax1.set_title('Top Actions by F1 Score', fontweight='bold')\n    ax1.grid(True, alpha=0.3)\n    \n    # 2. F1 Distribution\n    ax2 = plt.subplot(3, 3, 2)\n    ax2.hist(f1_df['f1_score'], bins=25, color='blue', edgecolor='black', alpha=0.7)\n    ax2.set_xlabel('F1 Score')\n    ax2.set_ylabel('Frequency')\n    ax2.set_title('F1 Score Distribution', fontweight='bold')\n    ax2.axvline(f1_df['f1_score'].mean(), color='red', linestyle='--',\n               label=f'Mean: {f1_df[\"f1_score\"].mean():.3f}')\n    ax2.legend()\n    ax2.grid(True, alpha=0.3)\n    \n    # 3. Predictions per Config\n    ax3 = plt.subplot(3, 3, 3)\n    if predictions_per_config:\n        ax3.bar(range(len(predictions_per_config)), predictions_per_config, \n               color='purple', alpha=0.7)\n        ax3.set_xlabel('Configuration')\n        ax3.set_ylabel('Predictions')\n        ax3.set_title('Predictions per Configuration', fontweight='bold')\n        ax3.grid(True, alpha=0.3)\n    \n    # 4. Cumulative Predictions\n    ax4 = plt.subplot(3, 3, 4)\n    if predictions_per_config:\n        cumsum = np.cumsum(predictions_per_config)\n        ax4.plot(cumsum, 'b-', linewidth=2)\n        ax4.fill_between(range(len(cumsum)), 0, cumsum, alpha=0.3, color='blue')\n        ax4.set_xlabel('Configuration')\n        ax4.set_ylabel('Cumulative Predictions')\n        ax4.set_title('Cumulative Prediction Growth', fontweight='bold')\n        ax4.grid(True, alpha=0.3)\n    \n    # 5. Best vs Worst Actions\n    ax5 = plt.subplot(3, 3, 5)\n    action_f1_sorted = f1_df.groupby('action')['f1_score'].mean().sort_values(ascending=False)\n    if len(action_f1_sorted) > 0:\n        top_5 = action_f1_sorted.head(5)\n        bottom_5 = action_f1_sorted.tail(5) if len(action_f1_sorted) > 5 else pd.Series()\n        \n        all_performers = pd.concat([top_5, bottom_5]) if len(bottom_5) > 0 else top_5\n        colors = ['green'] * len(top_5) + ['red'] * len(bottom_5) if len(bottom_5) > 0 else ['green'] * len(top_5)\n        \n        ax5.bar(range(len(all_performers)), all_performers.values, color=colors, alpha=0.7)\n        ax5.set_xticks(range(len(all_performers)))\n        ax5.set_xticklabels(all_performers.index, rotation=45, ha='right', fontsize=8)\n        ax5.set_ylabel('F1 Score')\n        ax5.set_title('Best & Worst Performers', fontweight='bold')\n        ax5.grid(True, alpha=0.3, axis='y')\n    \n    # 6. F1 by Configuration\n    ax6 = plt.subplot(3, 3, 6)\n    config_f1 = f1_df.groupby('config')['f1_score'].mean()\n    ax6.plot(range(len(config_f1)), config_f1.values, 'o-', color='orange', linewidth=2, markersize=8)\n    ax6.set_xlabel('Configuration Index')\n    ax6.set_ylabel('Mean F1 Score')\n    ax6.set_title('Performance by Configuration', fontweight='bold')\n    ax6.grid(True, alpha=0.3)\n    \n    # 7-9. Summary statistics\n    ax7 = plt.subplot(3, 1, 3)\n    ax7.axis('off')\n    \n    stats_text = f\"\"\"\nPERFORMANCE METRICS SUMMARY\n{'-'*60}\n\nF1 SCORE STATISTICS:\n- Mean: {f1_df['f1_score'].mean():.4f}\n- Median: {f1_df['f1_score'].median():.4f}\n- Std Dev: {f1_df['f1_score'].std():.4f}\n- Min: {f1_df['f1_score'].min():.4f}\n- Max: {f1_df['f1_score'].max():.4f}\n\nCOVERAGE:\n- Total Actions Evaluated: {f1_df['action'].nunique()}\n- Total Configurations: {f1_df['config'].nunique()}\n- Total Predictions Made: {sum(predictions_per_config) if predictions_per_config else 0:,}\n\nTOP 3 PERFORMING ACTIONS:\n1. {action_f1_sorted.index[0] if len(action_f1_sorted) > 0 else 'N/A'}: {action_f1_sorted.iloc[0] if len(action_f1_sorted) > 0 else 0:.3f}\n2. {action_f1_sorted.index[1] if len(action_f1_sorted) > 1 else 'N/A'}: {action_f1_sorted.iloc[1] if len(action_f1_sorted) > 1 else 0:.3f}\n3. {action_f1_sorted.index[2] if len(action_f1_sorted) > 2 else 'N/A'}: {action_f1_sorted.iloc[2] if len(action_f1_sorted) > 2 else 0:.3f}\n\"\"\"\n    \n    ax7.text(0.5, 0.5, stats_text, transform=ax7.transAxes, fontsize=11,\n            verticalalignment='center', horizontalalignment='center',\n            fontfamily='monospace')\n    \n    plt.suptitle('Model Performance Comprehensive Analysis', fontsize=14, fontweight='bold')\n    plt.tight_layout()\n    \n    if save_plots:\n        plt.savefig(os.path.join(plot_dir, 'performance_summary.png'), dpi=100, bbox_inches='tight')\n        print(\"    💾 Saved: performance_summary.png\")\n    \n    plt.show()\n\n# Load data\nprint(\"📁 LOADING DATA...\")\ntrain = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/train.csv')\ntrain['n_mice'] = 4 - train[['mouse1_strain', 'mouse2_strain', 'mouse3_strain', 'mouse4_strain']].isna().sum(axis=1)\ntrain_without_mabe22 = train.query(\"~ lab_id.str.startswith('MABe22_')\")\n\ntest = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/test.csv')\nbody_parts_tracked_list = list(np.unique(train.body_parts_tracked))\n\nprint(f\"✅ Loaded {len(train)} training videos\")\nprint(f\"✅ Loaded {len(test)} test videos\")\nprint(f\"✅ Found {len(body_parts_tracked_list)} body part configurations\")\n\n# Create data overview visualizations\nif create_visualizations:\n    behavior_counts = create_data_overview_plots(train, test)\n\n# Body parts to drop for memory efficiency\ndrop_body_parts = ['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\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    \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        vid = pd.read_parquet(path)\n        \n        if len(np.unique(vid.bodypart)) > 5:\n            vid = vid.query(\"~ bodypart.isin(@drop_body_parts)\")\n        \n        pvid = vid.pivot(columns=['mouse_id', 'bodypart'], index='video_frame', values=['x', 'y'])\n        \n        if pvid.isna().any().any():\n            if verbose and traintest == 'test': \n                print(f'    Video {video_id}: has missing values')\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                    \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                    \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'], \n                                                 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(\n                        \"(agent == @agent_str) & (target == @target_str)\").action)\n                    \n                    mouse_pair = pd.concat([pvid[agent], pvid[target]], axis=1, keys=['A', 'B'])\n                    \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                    \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'], \n                                               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\ndef transform_single(single_mouse, body_parts_tracked):\n    \"\"\"Transform single mouse data with angular features\"\"\"\n    available_body_parts = single_mouse.columns.get_level_values(0)\n    \n    X = pd.DataFrame({\n        f\"{part1}+{part2}\": np.square(single_mouse[part1] - single_mouse[part2]).sum(axis=1, skipna=False)\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    })\n    \n    X = X.reindex(columns=[f\"{part1}+{part2}\" \n                          for part1, part2 in itertools.combinations(body_parts_tracked, 2)], \n                 copy=False)\n\n    if 'ear_left' in available_body_parts and 'ear_right' in available_body_parts:\n        if 'tail_base' in available_body_parts:\n            shifted = single_mouse[['ear_left', 'ear_right', 'tail_base']].shift(10)\n            X = pd.concat([\n                X, \n                pd.DataFrame({\n                    'speed_left': np.square(single_mouse['ear_left'] - shifted['ear_left']).sum(axis=1, skipna=False),\n                    'speed_right': np.square(single_mouse['ear_right'] - shifted['ear_right']).sum(axis=1, skipna=False),\n                    'speed_left2': np.square(single_mouse['ear_left'] - shifted['tail_base']).sum(axis=1, skipna=False),\n                    'speed_right2': np.square(single_mouse['ear_right'] - shifted['tail_base']).sum(axis=1, skipna=False),\n                })\n            ], axis=1)\n    \n    if 'nose+tail_base' in X.columns and 'ear_left+ear_right' in X.columns:\n        X['elongation_ratio'] = X['nose+tail_base'] / (X['ear_left+ear_right'] + 1e-6)\n    \n    # Angular feature\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_cos'] = dot_product / (norm_v1 * norm_v2 + 1e-6)\n    \n    performance_metrics['features_extracted'] += len(X.columns)\n    return X\n\ndef transform_pair(mouse_pair, body_parts_tracked):\n    \"\"\"Transform mouse pair data with angular features\"\"\"\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        f\"12+{part1}+{part2}\": np.square(mouse_pair['A'][part1] - mouse_pair['B'][part2]).sum(axis=1, skipna=False)\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    })\n    \n    X = X.reindex(columns=[f\"12+{part1}+{part2}\" \n                          for part1, part2 in itertools.product(body_parts_tracked, repeat=2)], \n                 copy=False)\n\n    if ('A', 'ear_left') in mouse_pair.columns and ('B', 'ear_left') in mouse_pair.columns:\n        shifted_A = mouse_pair['A']['ear_left'].shift(10)\n        shifted_B = mouse_pair['B']['ear_left'].shift(10)\n        X = pd.concat([\n            X,\n            pd.DataFrame({\n                'speed_left_A': np.square(mouse_pair['A']['ear_left'] - shifted_A).sum(axis=1, skipna=False),\n                'speed_left_AB': np.square(mouse_pair['A']['ear_left'] - shifted_B).sum(axis=1, skipna=False),\n                'speed_left_B': np.square(mouse_pair['B']['ear_left'] - shifted_B).sum(axis=1, skipna=False),\n            })\n        ], axis=1)\n    \n    if 'nose+tail_base' in X.columns and 'ear_left+ear_right' in X.columns:\n        X['elongation_ratio'] = X['nose+tail_base'] / (X['ear_left+ear_right'] + 1e-6)\n    \n    # Relative orientation feature\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        \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['relative_orientation_cos'] = dot_product / (norm_A * norm_B + 1e-6)\n    \n    performance_metrics['features_extracted'] += len(X.columns)\n    return X\n\nthreshold = 0.27\n\ndef predict_multiclass(pred, meta):\n    ama = np.argmax(pred, axis=1)\n    ama = np.where(pred.max(axis=1) >= threshold, 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    \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        \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    performance_metrics['predictions_made'] += len(submission_part)\n    if verbose: \n        print(f'      Actions found: {len(submission_part)}')\n    \n    return submission_part\n\ndef submit(body_parts_tracked_str, switch_tr, binary_classifier, X_tr, label, meta):\n    model_list = []\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():\n            model = clone(binary_classifier)\n            model.fit(X_tr[action_mask], y_action)\n            model_list.append((action, model))\n            performance_metrics['models_trained'] += 1\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    \n    if len(test_subset) == 0:\n        return\n    \n    generator = generate_mouse_data(test_subset, 'test',\n                                  generate_single=(switch_tr == 'single'), \n                                  generate_pair=(switch_tr == 'pair'))\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(data_te, body_parts_tracked)\n            else:\n                X_te = transform_pair(data_te, body_parts_tracked)\n            \n            del data_te\n            \n            pred = pd.DataFrame(index=meta_te.video_frame)\n            for action, model in model_list:\n                if action in actions_te:\n                    pred[action] = model.predict_proba(X_te)[:, 1]\n            \n            del X_te\n            \n            if pred.shape[1] != 0:\n                submission_part = predict_multiclass(pred, meta_te)\n                submission_list.append(submission_part)\n                \n        except Exception as e:\n            if verbose: \n                print(f'      Error: {e}')\n\ndef robustify(submission, dataset, traintest):\n    traintest_directory = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking\"\n    \n    if submission is None or len(submission) == 0:\n        submission = pd.DataFrame(columns=['video_id', 'agent_id', 'target_id', \n                                          'action', 'start_frame', 'stop_frame'])\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    if group_list:\n        submission = pd.concat(group_list)\n    \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        \n        if len(submission) == 0 or not (submission.video_id == video_id).any():\n            try:\n                path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n                vid = pd.read_parquet(path)\n                \n                vid_behaviors = eval(row['behaviors_labeled']) if pd.notna(row['behaviors_labeled']) else []\n                if not vid_behaviors:\n                    s_list.append((video_id, 'mouse1', 'self', 'rear', 0, 100))\n                    continue\n                    \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 = int(vid.video_frame.min())\n                stop_frame = int(vid.video_frame.max() + 1)\n                \n                for (agent, target), actions in vid_behaviors.groupby(['agent', 'target']):\n                    batch_length = 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_length\n                        batch_stop = min(batch_start + batch_length, stop_frame)\n                        s_list.append((video_id, agent, target, action_row['action'], \n                                     batch_start, batch_stop))\n            except:\n                s_list.append((video_id, 'mouse1', 'self', 'rear', 0, 100))\n    \n    if len(s_list) > 0:\n        new_df = pd.DataFrame(s_list, columns=['video_id', 'agent_id', 'target_id', \n                                               'action', 'start_frame', 'stop_frame'])\n        if len(submission) > 0:\n            submission = pd.concat([submission, new_df], ignore_index=True)\n        else:\n            submission = new_df\n    \n    return submission.reset_index(drop=True)\n\n# Main processing loop\nprint(\"\\n🚀 STARTING MAIN PROCESSING PIPELINE...\")\nprint(\"=\"*80)\n\nf1_list = []\nsubmission_list = []\npredictions_per_config = []\nconfig_times = []\n\nfor section in range(1, len(body_parts_tracked_list)):\n    config_start = time.time()\n    body_parts_tracked_str = body_parts_tracked_list[section]\n    \n    print(f\"\\n{'='*60}\")\n    print(f\"Configuration {section}/{len(body_parts_tracked_list)-1}\")\n    print(f\"{'='*60}\")\n    \n    try:\n        body_parts_tracked = json.loads(body_parts_tracked_str)\n        print(f\"Processing {len(body_parts_tracked)} body parts: {body_parts_tracked[:3]}...\")\n        \n        if len(body_parts_tracked) > 5:\n            body_parts_tracked = [b for b in body_parts_tracked if b not in drop_body_parts]\n        \n        train_subset = train[train.body_parts_tracked == body_parts_tracked_str]\n        \n        if len(train_subset) == 0:\n            print(\"  No training data for this configuration\")\n            continue\n        \n        print(f\"  Found {len(train_subset)} training videos\")\n        performance_metrics['configurations_processed'] += 1\n        \n        # Collect training data\n        single_mouse_list = []\n        single_mouse_label_list = []\n        single_mouse_meta_list = []\n        mouse_pair_list = []\n        mouse_pair_label_list = []\n        mouse_pair_meta_list = []\n        \n        print(\"  📥 Collecting training data...\")\n        for switch, data, meta, label in generate_mouse_data(train_subset, 'train'):\n            if switch == 'single':\n                single_mouse_list.append(data)\n                single_mouse_meta_list.append(meta)\n                single_mouse_label_list.append(label)\n            else:\n                mouse_pair_list.append(data)\n                mouse_pair_meta_list.append(meta)\n                mouse_pair_label_list.append(label)\n        \n        # Create classifier\n        binary_classifier = make_pipeline(\n            SimpleImputer(),\n            TrainOnSubsetClassifier(\n                lgb.LGBMClassifier(\n                    n_estimators=100,\n                    learning_rate=0.03,\n                    min_child_samples=40,\n                    num_leaves=31,\n                    max_depth=-1,\n                    subsample=0.8,\n                    colsample_bytree=0.8,\n                    verbose=-1\n                ),\n                100000\n            )\n        )\n        \n        config_predictions = performance_metrics['predictions_made']\n        \n        # Process single mouse data\n        if len(single_mouse_list) > 0:\n            print(f\"\\n  🐭 Processing {len(single_mouse_list)} single mouse batches...\")\n            performance_metrics['single_mouse_batches'] += len(single_mouse_list)\n            \n            single_mouse = pd.concat(single_mouse_list)\n            single_mouse_label = pd.concat(single_mouse_label_list)\n            single_mouse_meta = pd.concat(single_mouse_meta_list)\n            \n            del single_mouse_list, single_mouse_label_list, single_mouse_meta_list\n            \n            X_tr = transform_single(single_mouse, body_parts_tracked)\n            del single_mouse\n            \n            print(f\"    Features shape: {X_tr.shape}\")\n            \n            # Create feature analysis for first 2 configs\n            if create_visualizations and section <= 2:\n                create_feature_analysis_plots(X_tr, 'single', section)\n            \n            submit(body_parts_tracked_str, 'single', binary_classifier, X_tr, \n                  single_mouse_label, single_mouse_meta)\n            \n            del X_tr\n            gc.collect()\n        \n        # Process pair data\n        if len(mouse_pair_list) > 0:\n            print(f\"\\n  🐭🐭 Processing {len(mouse_pair_list)} mouse pair batches...\")\n            performance_metrics['pair_batches'] += len(mouse_pair_list)\n            \n            mouse_pair = pd.concat(mouse_pair_list)\n            mouse_pair_label = pd.concat(mouse_pair_label_list)\n            mouse_pair_meta = pd.concat(mouse_pair_meta_list)\n            \n            del mouse_pair_list, mouse_pair_label_list, mouse_pair_meta_list\n            \n            X_tr = transform_pair(mouse_pair, body_parts_tracked)\n            del mouse_pair\n            \n            print(f\"    Features shape: {X_tr.shape}\")\n            \n            # Create feature analysis for first 2 configs\n            if create_visualizations and section <= 2:\n                create_feature_analysis_plots(X_tr, 'pair', section)\n            \n            submit(body_parts_tracked_str, 'pair', binary_classifier, X_tr, \n                  mouse_pair_label, mouse_pair_meta)\n            \n            del X_tr\n            gc.collect()\n        \n        config_time = time.time() - config_start\n        config_times.append(config_time)\n        predictions_per_config.append(performance_metrics['predictions_made'] - config_predictions)\n        \n        print(f\"\\n  ⏱️ Configuration time: {config_time:.2f}s\")\n        print(f\"  📊 Predictions made: {predictions_per_config[-1]}\")\n        \n    except Exception as e:\n        print(f\"  ❌ Error in configuration {section}: {e}\")\n\n# Create performance summary\nif create_visualizations:\n    create_performance_summary_plot(f1_list, predictions_per_config)\n\n# Create final submission\nprint(\"\\n📋 CREATING FINAL SUBMISSION...\")\nprint(\"=\"*80)\n\nif submission_list:\n    submission = pd.concat(submission_list)\n    print(f\"  Generated {len(submission)} predictions\")\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    print(\"  Using default submission\")\n\nsubmission_robust = robustify(submission, test, 'test')\nsubmission_robust.index.name = 'row_id'\nsubmission_robust.to_csv('submission.csv')\n\nprint(f\"\\n✅ Final submission: {len(submission_robust)} rows\")\nprint(\"✅ Saved to submission.csv\")\n\n# Print performance summary\nprint(\"\\n\" + \"=\"*80)\nprint(\"📊 PERFORMANCE SUMMARY\")\nprint(\"=\"*80)\n\nelapsed = time.time()\nprint(f\"⏱️ Total execution time: {elapsed:.2f} seconds\")\nprint(f\"🔧 Configurations processed: {performance_metrics['configurations_processed']}\")\nprint(f\"🐭 Single mouse batches: {performance_metrics['single_mouse_batches']}\")\nprint(f\"🐭🐭 Pair batches: {performance_metrics['pair_batches']}\")\nprint(f\"📐 Features extracted: {performance_metrics['features_extracted']:,}\")\nprint(f\"🎯 Predictions made: {performance_metrics['predictions_made']:,}\")\nprint(f\"🤖 Models trained: {performance_metrics['models_trained']}\")\n\nif config_times:\n    print(f\"\\n⏱️ Processing Time Statistics:\")\n    print(f\"   • Mean time per config: {np.mean(config_times):.2f} seconds\")\n    print(f\"   • Total processing time: {sum(config_times):.2f} seconds\")\n\nprint(\"\\n\" + \"=\"*80)\nprint(\" \"*25 + \"✨ PIPELINE COMPLETE! ✨\")\nprint(\"=\"*80)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# \"\"\"\n# MABe Challenge - Enhanced Social Action Recognition in Mice\n# ============================================================\n# This notebook implements an improved solution for recognizing mouse behaviors\n# using machine learning on motion capture data.\n\n# Key Improvements:\n# 1. Multiple classifier options (KNN, RandomForest, GradientBoost)\n# 2. Enhanced temporal features at multiple scales\n# 3. Per-action threshold optimization\n# 4. Comprehensive visualizations\n# 5. Advanced post-processing\n# \"\"\"\n\n# # ========================================\n# # IMPORTS AND CONFIGURATION\n# # ========================================\n\n# import pandas as pd\n# import numpy as np\n# import matplotlib.pyplot as plt\n# import seaborn as sns\n# from tqdm import trange, tqdm\n# import itertools\n# import warnings\n# import json\n# import os\n# import gc\n# from datetime import datetime\n\n# # Sklearn imports\n# from sklearn.base import ClassifierMixin, BaseEstimator, clone\n# from sklearn.model_selection import cross_val_predict, GroupKFold, StratifiedGroupKFold, train_test_split\n# from sklearn.pipeline import make_pipeline\n# from sklearn.impute import SimpleImputer\n# from sklearn.preprocessing import StandardScaler, RobustScaler\n# from sklearn.neighbors import KNeighborsClassifier\n# from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier, VotingClassifier\n# from sklearn.linear_model import LogisticRegression\n# from sklearn.metrics import f1_score, precision_recall_curve, confusion_matrix, classification_report\n# from sklearn.feature_selection import SelectKBest, f_classif, mutual_info_classif\n\n# # Visualization settings\n# plt.style.use('seaborn-v0_8-darkgrid')\n# sns.set_palette(\"husl\")\n\n# # Polars for efficient data processing\n# import polars as pl\n\n# # ========================================\n# # CONFIGURATION PARAMETERS\n# # ========================================\n\n# # Mode selection\n# validate_or_submit = 'validate'  # 'validate', 'submit', or 'stresstest'\n# verbose = True\n# use_ensemble = True  # Whether to use ensemble of classifiers\n# optimize_thresholds = True  # Whether to optimize thresholds per action\n\n# # Model hyperparameters\n# CONFIG = {\n#     'n_samples_train': 30000,  # Increased from 20000 for better training\n#     'knn_neighbors': 15,  # Increased from 11\n#     'rf_estimators': 100,\n#     'rf_max_depth': 15,\n#     'gb_estimators': 50,\n#     'gb_max_depth': 5,\n#     'gb_learning_rate': 0.1,\n#     'feature_selection_k': 50,  # Top k features to select\n#     'time_windows': [5, 10, 20, 30],  # Multiple time scales for temporal features\n#     'min_action_duration': 3,  # Minimum frames for an action\n#     'smoothing_window': 5,  # Window for post-processing smoothing\n# }\n\n# print(\"=\"*60)\n# print(\"MABe Challenge - Enhanced Mouse Behavior Recognition\")\n# print(\"=\"*60)\n# print(f\"Mode: {validate_or_submit}\")\n# print(f\"Ensemble: {use_ensemble}\")\n# print(f\"Optimize Thresholds: {optimize_thresholds}\")\n# print(f\"Start Time: {datetime.now()}\")\n# print(\"=\"*60)\n\n# # ========================================\n# # CUSTOM CLASSES\n# # ========================================\n\n# class TrainOnSubsetClassifier(ClassifierMixin, BaseEstimator):\n#     \"\"\"\n#     Enhanced classifier that trains on a subset of data with optional stratification.\n#     Includes memory optimization and better sampling strategy.\n#     \"\"\"\n#     def __init__(self, estimator, n_samples, stratify=True):\n#         self.estimator = estimator\n#         self.n_samples = n_samples\n#         self.stratify = stratify\n\n#     def fit(self, X, y):\n#         # Convert to numpy arrays for efficiency\n#         X_arr = np.array(X, copy=False)\n#         y_arr = np.array(y, copy=False)\n        \n#         if len(X_arr) <= self.n_samples:\n#             # If dataset is small, use all data\n#             self.estimator.fit(X_arr, y_arr)\n#         else:\n#             # Smart downsampling\n#             if self.stratify and len(np.unique(y_arr)) > 1:\n#                 # Stratified sampling to maintain class balance\n#                 from sklearn.utils import resample\n#                 X_sampled, y_sampled = resample(X_arr, y_arr, \n#                                                n_samples=self.n_samples,\n#                                                stratify=y_arr,\n#                                                random_state=42)\n#                 self.estimator.fit(X_sampled, y_sampled)\n#             else:\n#                 # Regular downsampling\n#                 downsample = len(X_arr) // self.n_samples\n#                 self.estimator.fit(X_arr[::downsample], y_arr[::downsample])\n        \n#         self.classes_ = self.estimator.classes_\n#         return self\n\n#     def predict_proba(self, X):\n#         return self.estimator.predict_proba(np.array(X, copy=False))\n        \n#     def predict(self, X):\n#         return self.estimator.predict(np.array(X, copy=False))\n\n# # ========================================\n# # SCORING FUNCTIONS (Competition Metrics)\n# # ========================================\n\n# class HostVisibleError(Exception):\n#     pass\n\n# def single_lab_f1(lab_solution: pl.DataFrame, lab_submission: pl.DataFrame, beta: float = 1) -> float:\n#     \"\"\"Calculate F1 score for a single lab's predictions\"\"\"\n#     label_frames = defaultdict(set)\n#     prediction_frames = defaultdict(set)\n    \n#     from collections import defaultdict\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            \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            \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    \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] / \n#                             ((1 + beta**2) * tps[action] + beta**2 * fns[action] + fps[action]))\n    \n#     return sum(action_f1s) / len(action_f1s) if action_f1s else 0\n\n# def mouse_fbeta(solution: pd.DataFrame, submission: pd.DataFrame, beta: float = 1) -> float:\n#     \"\"\"Main scoring function for the competition\"\"\"\n#     if len(solution) == 0 or len(submission) == 0:\n#         raise ValueError('Missing solution or submission data')\n\n#     expected_cols = ['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']\n    \n#     for col in expected_cols:\n#         if col not in solution.columns:\n#             raise ValueError(f'Solution is missing column {col}')\n#         if col not in submission.columns:\n#             raise ValueError(f'Submission is missing column {col}')\n\n#     solution = pl.DataFrame(solution)\n#     submission = 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#     # Add keys for matching\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    \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\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) if lab_scores else 0\n\n# def score(solution: pd.DataFrame, submission: pd.DataFrame, row_id_column_name: str, beta: float = 1) -> float:\n#     \"\"\"Wrapper for competition scoring\"\"\"\n#     solution = solution.drop(row_id_column_name, axis='columns', errors='ignore')\n#     submission = submission.drop(row_id_column_name, axis='columns', errors='ignore')\n#     return mouse_fbeta(solution, submission, beta=beta)\n\n# # ========================================\n# # DATA LOADING AND PREPROCESSING\n# # ========================================\n\n# print(\"\\n📁 Loading data...\")\n\n# # Load training and test metadata\n# train = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/train.csv')\n# train['n_mice'] = 4 - train[['mouse1_strain', 'mouse2_strain', 'mouse3_strain', 'mouse4_strain']].isna().sum(axis=1)\n# train_without_mabe22 = train.query(\"~ lab_id.str.startswith('MABe22_')\")\n\n# test = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/test.csv')\n\n# # Get unique body parts configurations\n# body_parts_tracked_list = list(np.unique(train.body_parts_tracked))\n\n# print(f\"✅ Loaded {len(train)} training videos\")\n# print(f\"✅ Found {len(body_parts_tracked_list)} unique body part configurations\")\n# print(f\"✅ Number of mice per video: {train['n_mice'].value_counts().to_dict()}\")\n\n# # Visualize data distribution\n# if verbose:\n#     fig, axes = plt.subplots(1, 3, figsize=(15, 4))\n    \n#     # Plot 1: Videos per lab\n#     lab_counts = train.lab_id.value_counts()\n#     axes[0].barh(range(len(lab_counts[:10])), lab_counts[:10].values)\n#     axes[0].set_yticks(range(len(lab_counts[:10])))\n#     axes[0].set_yticklabels(lab_counts[:10].index)\n#     axes[0].set_xlabel('Number of Videos')\n#     axes[0].set_title('Top 10 Labs by Video Count')\n    \n#     # Plot 2: Video duration distribution\n#     axes[1].hist(train['video_duration_sec'].dropna(), bins=30, edgecolor='black')\n#     axes[1].set_xlabel('Duration (seconds)')\n#     axes[1].set_ylabel('Count')\n#     axes[1].set_title('Video Duration Distribution')\n    \n#     # Plot 3: Number of mice distribution\n#     mice_counts = train['n_mice'].value_counts()\n#     axes[2].pie(mice_counts.values, labels=mice_counts.index, autopct='%1.1f%%')\n#     axes[2].set_title('Distribution of Mice Count per Video')\n    \n#     plt.tight_layout()\n#     plt.show()\n\n# # ========================================\n# # SOLUTION DATAFRAME CREATION\n# # ========================================\n\n# def create_solution_df(dataset):\n#     \"\"\"Create solution dataframe for validation\"\"\"\n#     solution = []\n    \n#     for _, row in tqdm(dataset.iterrows(), total=len(dataset), desc=\"Creating solution\"):\n#         lab_id = row['lab_id']\n#         if lab_id.startswith('MABe22'):\n#             continue\n            \n#         video_id = row['video_id']\n#         path = f\"/kaggle/input/MABe-mouse-behavior-detection/train_annotation/{lab_id}/{video_id}.parquet\"\n        \n#         try:\n#             annot = pd.read_parquet(path)\n#         except FileNotFoundError:\n#             if verbose:\n#                 print(f\"No annotations for {path}\")\n#             continue\n        \n#         annot['lab_id'] = lab_id\n#         annot['video_id'] = video_id\n#         annot['behaviors_labeled'] = row['behaviors_labeled']\n#         annot['target_id'] = np.where(annot.target_id != annot.agent_id, \n#                                      annot['target_id'].apply(lambda s: f\"mouse{s}\"), 'self')\n#         annot['agent_id'] = annot['agent_id'].apply(lambda s: f\"mouse{s}\")\n#         solution.append(annot)\n    \n#     solution = pd.concat(solution)\n#     return solution\n\n# if validate_or_submit == 'validate':\n#     print(\"\\n📊 Creating validation solution...\")\n#     solution = create_solution_df(train_without_mabe22)\n#     print(f\"✅ Solution created with {len(solution)} annotations\")\n\n# # ========================================\n# # DATA GENERATION FUNCTIONS\n# # ========================================\n\n# def generate_mouse_data(dataset, traintest, traintest_directory=None, \n#                        generate_single=True, generate_pair=True):\n#     \"\"\"\n#     Enhanced data generator with better error handling and memory optimization\n#     \"\"\"\n#     assert traintest in ['train', 'test']\n    \n#     if traintest_directory is None:\n#         traintest_directory = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking\"\n    \n#     for idx, row in dataset.iterrows():\n#         if idx % 100 == 0:\n#             gc.collect()  # Periodic garbage collection\n            \n#         lab_id = row.lab_id\n#         if lab_id.startswith('MABe22'):\n#             continue\n            \n#         video_id = row.video_id\n#         path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        \n#         try:\n#             # Load and pivot video data\n#             vid = pd.read_parquet(path)\n#             pvid = vid.pivot(columns=['mouse_id', 'bodypart'], \n#                            index='video_frame', \n#                            values=['x', 'y'])\n            \n#             # Check for missing values\n#             missing_ratio = pvid.isna().sum().sum() / pvid.size\n#             if verbose and traintest == 'test':\n#                 status = f\"missing {missing_ratio:.1%}\" if missing_ratio > 0 else \"complete\"\n#                 print(f'Video {video_id}: {status}, {len(vid)} frames')\n            \n#             del vid\n            \n#             # Reorder columns for consistency\n#             pvid = pvid.reorder_levels([1, 2, 0], axis=1).T.sort_index().T\n#             pvid /= row.pix_per_cm_approx  # Convert to cm\n            \n#             # Parse behaviors for this video\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#             # Load annotations for training\n#             if traintest == 'train':\n#                 try:\n#                     annot = pd.read_parquet(path.replace('train_tracking', 'train_annotation'))\n#                 except FileNotFoundError:\n#                     continue\n            \n#             # Generate single mouse data\n#             if generate_single:\n#                 vid_behaviors_subset = vid_behaviors.query(\"target == 'self'\")\n                \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(\n#                             vid_behaviors_subset.query(\"agent == @mouse_id_str\").action\n#                         )\n                        \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                        \n#                         if traintest == 'train':\n#                             single_mouse_label = pd.DataFrame(0.0, \n#                                                              columns=vid_agent_actions,\n#                                                              index=single_mouse.index)\n#                             annot_subset = annot.query(\"(agent_id == @mouse_id) & (target_id == @mouse_id)\")\n                            \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'], \n#                                                      annot_row.action] = 1.0\n                            \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                            \n#                     except (KeyError, ValueError) as e:\n#                         if verbose:\n#                             print(f\"Error processing single mouse {mouse_id_str}: {e}\")\n#                         continue\n            \n#             # Generate mouse pair data\n#             if generate_pair:\n#                 vid_behaviors_subset = vid_behaviors.query(\"target != 'self'\")\n                \n#                 if len(vid_behaviors_subset) > 0:\n#                     mouse_ids = np.unique(pvid.columns.get_level_values('mouse_id'))\n                    \n#                     for agent, target in itertools.permutations(mouse_ids, 2):\n#                         agent_str = f\"mouse{agent}\"\n#                         target_str = f\"mouse{target}\"\n#                         vid_agent_actions = np.unique(\n#                             vid_behaviors_subset.query(\n#                                 \"(agent == @agent_str) & (target == @target_str)\"\n#                             ).action\n#                         )\n                        \n#                         try:\n#                             mouse_pair = pd.concat([pvid[agent], pvid[target]], \n#                                                  axis=1, \n#                                                  keys=['A', 'B'])\n                            \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                            \n#                             if traintest == 'train':\n#                                 mouse_pair_label = pd.DataFrame(0.0,\n#                                                                columns=vid_agent_actions,\n#                                                                index=mouse_pair.index)\n#                                 annot_subset = annot.query(\"(agent_id == @agent) & (target_id == @target)\")\n                                \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'],\n#                                                        annot_row.action] = 1.0\n                                \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#                         except (KeyError, ValueError) as e:\n#                             if verbose:\n#                                 print(f\"Error processing pair {agent_str}-{target_str}: {e}\")\n#                             continue\n                            \n#         except Exception as e:\n#             if verbose:\n#                 print(f\"Error processing video {video_id}: {e}\")\n#             continue\n\n# # ========================================\n# # FEATURE ENGINEERING FUNCTIONS\n# # ========================================\n\n# def transform_single_enhanced(single_mouse, body_parts_tracked, time_windows=None):\n#     \"\"\"\n#     Enhanced feature transformation for single mouse with multiple temporal scales\n#     and additional biomechanical features\n#     \"\"\"\n#     if time_windows is None:\n#         time_windows = CONFIG['time_windows']\n    \n#     features = []\n    \n#     # 1. Distance features between body parts\n#     distance_features = pd.DataFrame({\n#         f\"{part1}+{part2}\": np.sqrt(\n#             np.square(single_mouse[part1] - single_mouse[part2]).sum(axis=1, skipna=False)\n#         )\n#         for part1, part2 in itertools.combinations(body_parts_tracked, 2)\n#     })\n#     features.append(distance_features)\n    \n#     # 2. Multi-scale temporal features\n#     if 'ear_left' in single_mouse.columns and 'ear_right' in single_mouse.columns:\n#         temporal_features = []\n        \n#         for window in time_windows:\n#             shifted = single_mouse[['ear_left', 'ear_right']].shift(window)\n            \n#             # Velocity features\n#             vel_left = np.sqrt(np.square(single_mouse['ear_left'] - shifted['ear_left']).sum(axis=1, skipna=False))\n#             vel_right = np.sqrt(np.square(single_mouse['ear_right'] - shifted['ear_right']).sum(axis=1, skipna=False))\n            \n#             temporal_features.extend([\n#                 pd.Series(vel_left, name=f'vel_left_w{window}'),\n#                 pd.Series(vel_right, name=f'vel_right_w{window}'),\n#                 pd.Series(vel_left - vel_right, name=f'vel_diff_w{window}'),  # Asymmetry\n#             ])\n            \n#             # Acceleration (if window > 5)\n#             if window > 5:\n#                 acc_left = vel_left.diff()\n#                 acc_right = vel_right.diff()\n#                 temporal_features.extend([\n#                     pd.Series(acc_left, name=f'acc_left_w{window}'),\n#                     pd.Series(acc_right, name=f'acc_right_w{window}'),\n#                 ])\n        \n#         features.append(pd.concat(temporal_features, axis=1))\n    \n#     # 3. Body orientation features\n#     if 'nose' in single_mouse.columns and 'tail_base' in single_mouse.columns:\n#         # Body axis vector\n#         body_vector = single_mouse['nose'] - single_mouse['tail_base']\n#         body_length = np.sqrt(np.square(body_vector).sum(axis=1, skipna=False))\n        \n#         # Angular velocity (change in body orientation)\n#         body_angle = np.arctan2(body_vector['y'], body_vector['x'])\n#         angular_velocity = body_angle.diff()\n        \n#         orientation_features = pd.DataFrame({\n#             'body_length': body_length,\n#             'body_angle': body_angle,\n#             'angular_velocity': angular_velocity,\n#             'angular_acceleration': angular_velocity.diff(),\n#         })\n#         features.append(orientation_features)\n    \n#     # 4. Statistical features over rolling windows\n#     if len(single_mouse) > 20:\n#         rolling_features = []\n#         for part in ['nose', 'tail_base']:\n#             if part in single_mouse.columns:\n#                 for stat, func in [('mean', 'mean'), ('std', 'std'), ('max', 'max')]:\n#                     rolled = single_mouse[part].rolling(window=20, min_periods=1)\n#                     x_stat = getattr(rolled['x'], func)()\n#                     y_stat = getattr(rolled['y'], func)()\n                    \n#                     rolling_features.extend([\n#                         pd.Series(x_stat, name=f'{part}_x_{stat}_20'),\n#                         pd.Series(y_stat, name=f'{part}_y_{stat}_20'),\n#                     ])\n        \n#         if rolling_features:\n#             features.append(pd.concat(rolling_features, axis=1))\n    \n#     # Concatenate all features\n#     X = pd.concat(features, axis=1)\n    \n#     # Fill NaN values with forward/backward fill\n#     X = X.fillna(method='ffill').fillna(method='bfill').fillna(0)\n    \n#     return X\n\n# def transform_pair_enhanced(mouse_pair, body_parts_tracked, time_windows=None):\n#     \"\"\"\n#     Enhanced feature transformation for mouse pairs with social interaction features\n#     \"\"\"\n#     if time_windows is None:\n#         time_windows = CONFIG['time_windows']\n    \n#     features = []\n    \n#     # Filter out unnecessary detailed body parts for pairs\n#     drop_body_parts = ['ear_left', 'ear_right', 'headpiece_bottombackleft', \n#                       'headpiece_bottombackright', 'headpiece_bottomfrontleft',\n#                       'headpiece_bottomfrontright', 'headpiece_topbackleft',\n#                       'headpiece_topbackright', 'headpiece_topfrontleft',\n#                       'headpiece_topfrontright', 'tail_midpoint']\n    \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#     # 1. Inter-mouse distance features\n#     distance_features = pd.DataFrame({\n#         f\"12+{part1}+{part2}\": np.sqrt(\n#             np.square(mouse_pair['A'][part1] - mouse_pair['B'][part2]).sum(axis=1, skipna=False)\n#         )\n#         for part1, part2 in itertools.product(body_parts_tracked, repeat=2)\n#     })\n#     features.append(distance_features)\n    \n#     # 2. Social dynamics features\n#     if 'nose' in mouse_pair['A'].columns and 'nose' in mouse_pair['B'].columns:\n#         # Distance between noses (face-to-face interaction)\n#         nose_distance = np.sqrt(\n#             np.square(mouse_pair['A']['nose'] - mouse_pair['B']['nose']).sum(axis=1, skipna=False)\n#         )\n        \n#         # Relative velocity (approaching/avoiding)\n#         for window in time_windows:\n#             shifted_A = mouse_pair['A']['nose'].shift(window)\n#             shifted_B = mouse_pair['B']['nose'].shift(window)\n            \n#             vel_A = np.sqrt(np.square(mouse_pair['A']['nose'] - shifted_A).sum(axis=1, skipna=False))\n#             vel_B = np.sqrt(np.square(mouse_pair['B']['nose'] - shifted_B).sum(axis=1, skipna=False))\n            \n#             # Relative velocity towards each other\n#             approach_velocity = nose_distance.diff(window) * -1  # Negative diff means approaching\n            \n#             social_features = pd.DataFrame({\n#                 f'nose_distance_w{window}': nose_distance,\n#                 f'approach_velocity_w{window}': approach_velocity,\n#                 f'vel_ratio_w{window}': vel_A / (vel_B + 1e-6),  # Who's moving more\n#                 f'synchronized_movement_w{window}': np.abs(vel_A - vel_B),\n#             })\n#             features.append(social_features)\n    \n#     # 3. Body orientation relative features\n#     if 'nose' in mouse_pair['A'].columns and 'tail_base' in mouse_pair['A'].columns:\n#         # Body vectors for both mice\n#         body_A = mouse_pair['A']['nose'] - mouse_pair['A']['tail_base']\n#         body_B = mouse_pair['B']['nose'] - mouse_pair['B']['tail_base']\n        \n#         # Angle between body axes (parallel, perpendicular, etc.)\n#         dot_product = (body_A['x'] * body_B['x'] + body_A['y'] * body_B['y'])\n#         norm_A = np.sqrt(np.square(body_A).sum(axis=1, skipna=False))\n#         norm_B = np.sqrt(np.square(body_B).sum(axis=1, skipna=False))\n        \n#         cos_angle = dot_product / (norm_A * norm_B + 1e-6)\n#         relative_angle = np.arccos(np.clip(cos_angle, -1, 1))\n        \n#         orientation_features = pd.DataFrame({\n#             'relative_body_angle': relative_angle,\n#             'relative_angle_change': relative_angle.diff(),\n#         })\n#         features.append(orientation_features)\n    \n#     # Concatenate all features\n#     X = pd.concat(features, axis=1)\n    \n#     # Fill NaN values\n#     X = X.fillna(method='ffill').fillna(method='bfill').fillna(0)\n    \n#     return X\n\n# # ========================================\n# # PREDICTION FUNCTIONS\n# # ========================================\n\n# def optimize_threshold_per_action(predictions, labels, actions):\n#     \"\"\"\n#     Optimize classification threshold for each action independently\n#     \"\"\"\n#     thresholds = {}\n    \n#     for action in actions:\n#         if action not in labels.columns:\n#             thresholds[action] = 0.5\n#             continue\n            \n#         y_true = labels[action].dropna()\n#         y_pred = predictions[action][~labels[action].isna()]\n        \n#         if len(np.unique(y_true)) < 2:\n#             thresholds[action] = 0.5\n#             continue\n        \n#         # Find threshold that maximizes F1 score\n#         precision, recall, thresholds_pr = precision_recall_curve(y_true, y_pred)\n#         f1_scores = 2 * (precision * recall) / (precision + recall + 1e-6)\n        \n#         best_idx = np.argmax(f1_scores)\n#         best_threshold = thresholds_pr[best_idx] if best_idx < len(thresholds_pr) else 0.5\n        \n#         thresholds[action] = best_threshold\n        \n#         if verbose:\n#             print(f\"  Action '{action}': threshold={best_threshold:.3f}, F1={f1_scores[best_idx]:.3f}\")\n    \n#     return thresholds\n\n# def predict_multiclass_enhanced(pred, meta, thresholds=None, min_duration=None):\n#     \"\"\"\n#     Enhanced multiclass prediction with per-action thresholds and post-processing\n#     \"\"\"\n#     if thresholds is None:\n#         thresholds = {col: 0.27 for col in pred.columns}\n    \n#     if min_duration is None:\n#         min_duration = CONFIG['min_action_duration']\n    \n#     # Apply per-action thresholds\n#     predictions = np.zeros(len(pred), dtype=int) - 1\n    \n#     for i, row in pred.iterrows():\n#         max_prob = 0\n#         max_action = -1\n        \n#         for j, action in enumerate(pred.columns):\n#             threshold = thresholds.get(action, 0.27)\n#             if row[action] > threshold and row[action] > max_prob:\n#                 max_prob = row[action]\n#                 max_action = j\n        \n#         predictions[i] = max_action\n    \n#     # Convert to series with frame index\n#     ama = pd.Series(predictions, index=meta.video_frame)\n    \n#     # Post-processing: Remove very short actions\n#     if min_duration > 1:\n#         # Find action segments\n#         changes = (ama != ama.shift(1))\n#         segments = changes.cumsum()\n        \n#         # Count frames per segment\n#         segment_counts = segments.value_counts()\n        \n#         # Identify short segments\n#         short_segments = segment_counts[segment_counts < min_duration].index\n        \n#         # Replace short action segments with no-action\n#         for seg in short_segments:\n#             if ama[segments == seg].iloc[0] >= 0:  # Only if it's an action\n#                 ama[segments == seg] = -1\n    \n#     # Convert to submission format\n#     changes_mask = (ama != ama.shift(1)).values\n#     ama_changes = ama[changes_mask]\n#     meta_changes = meta[changes_mask]\n    \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#     # Fix stop frames at video boundaries\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        \n#         if (i >= len(stop_video_id) or \n#             stop_video_id[i] != video_id or \n#             stop_agent_id[i] != agent_id or \n#             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    \n#     if verbose:\n#         print(f'  Actions found: {len(submission_part)}')\n    \n#     return submission_part\n\n# # ========================================\n# # CLASSIFIER CREATION FUNCTIONS\n# # ========================================\n\n# def create_ensemble_classifier(n_samples=None):\n#     \"\"\"\n#     Create an ensemble of diverse classifiers for better predictions\n#     \"\"\"\n#     if n_samples is None:\n#         n_samples = CONFIG['n_samples_train']\n    \n#     estimators = [\n#         ('knn', TrainOnSubsetClassifier(\n#             KNeighborsClassifier(n_neighbors=CONFIG['knn_neighbors']), \n#             n_samples\n#         )),\n#         ('rf', TrainOnSubsetClassifier(\n#             RandomForestClassifier(\n#                 n_estimators=CONFIG['rf_estimators'],\n#                 max_depth=CONFIG['rf_max_depth'],\n#                 random_state=42\n#             ),\n#             n_samples\n#         )),\n#         ('gb', TrainOnSubsetClassifier(\n#             GradientBoostingClassifier(\n#                 n_estimators=CONFIG['gb_estimators'],\n#                 max_depth=CONFIG['gb_max_depth'],\n#                 learning_rate=CONFIG['gb_learning_rate'],\n#                 random_state=42\n#             ),\n#             n_samples\n#         )),\n#     ]\n    \n#     return VotingClassifier(estimators=estimators, voting='soft')\n\n# def create_classifier_pipeline(use_ensemble=True):\n#     \"\"\"\n#     Create the full classification pipeline\n#     \"\"\"\n#     if use_ensemble:\n#         classifier = create_ensemble_classifier()\n#     else:\n#         classifier = TrainOnSubsetClassifier(\n#             KNeighborsClassifier(n_neighbors=CONFIG['knn_neighbors']),\n#             CONFIG['n_samples_train']\n#         )\n    \n#     pipeline = make_pipeline(\n#         SimpleImputer(strategy='median'),\n#         RobustScaler(),  # More robust to outliers than StandardScaler\n#         SelectKBest(f_classif, k=min(CONFIG['feature_selection_k'], 50)),\n#         classifier\n#     )\n    \n#     return pipeline\n\n# # ========================================\n# # CROSS-VALIDATION AND TRAINING\n# # ========================================\n\n# def cross_validate_classifier_enhanced(binary_classifier, X, label, meta, body_parts_tracked_str):\n#     \"\"\"\n#     Enhanced cross-validation with better metrics and visualization\n#     \"\"\"\n#     global f1_list, submission_list\n    \n#     print(f\"\\n🔍 Cross-validating for {len(label.columns)} actions...\")\n    \n#     oof = pd.DataFrame(index=meta.video_frame)\n#     action_metrics = []\n    \n#     for action in label.columns:\n#         action_mask = ~label[action].isna().values\n#         X_action = X[action_mask]\n#         y_action = label[action][action_mask].values.astype(int)\n        \n#         p = y_action.mean()\n#         baseline_score = p / (1 + p) if p > 0 else 0\n#         groups_action = meta.video_id[action_mask]\n        \n#         if len(np.unique(groups_action)) < 5:\n#             continue\n        \n#         if not (y_action == 0).all():\n#             with warnings.catch_warnings():\n#                 warnings.filterwarnings('ignore', category=RuntimeWarning)\n                \n#                 # Use StratifiedGroupKFold if possible\n#                 try:\n#                     cv = StratifiedGroupKFold(n_splits=5)\n#                     oof_action = cross_val_predict(\n#                         binary_classifier, X_action, y_action,\n#                         groups=groups_action, cv=cv,\n#                         method='predict_proba'\n#                     )\n#                 except:\n#                     cv = GroupKFold(n_splits=5)\n#                     oof_action = cross_val_predict(\n#                         binary_classifier, X_action, y_action,\n#                         groups=groups_action, cv=cv,\n#                         method='predict_proba'\n#                     )\n                \n#                 oof_action = oof_action[:, 1]\n#         else:\n#             oof_action = np.zeros(len(y_action))\n        \n#         # Calculate metrics\n#         threshold = 0.27\n#         y_pred = (oof_action >= threshold).astype(int)\n#         f1 = f1_score(y_action, y_pred, zero_division=0)\n        \n#         # Precision and recall\n#         from sklearn.metrics import precision_score, recall_score\n#         precision = precision_score(y_action, y_pred, zero_division=0)\n#         recall = recall_score(y_action, y_pred, zero_division=0)\n        \n#         ch = '↑' if f1 > baseline_score else '=' if f1 == baseline_score else '↓'\n#         print(f\"  {action:20} F1: {f1:.3f} {ch} (baseline: {baseline_score:.3f}) | \"\n#               f\"Prec: {precision:.3f} | Rec: {recall:.3f}\")\n        \n#         action_metrics.append({\n#             'action': action,\n#             'f1': f1,\n#             'precision': precision,\n#             'recall': recall,\n#             'baseline': baseline_score,\n#             'improvement': f1 - baseline_score\n#         })\n        \n#         f1_list.append((body_parts_tracked_str, action, f1))\n        \n#         oof_column = np.zeros(len(label))\n#         oof_column[action_mask] = oof_action\n#         oof[action] = oof_column\n    \n#     # Optimize thresholds if enabled\n#     if optimize_thresholds:\n#         print(\"\\n⚙️ Optimizing thresholds...\")\n#         action_thresholds = optimize_threshold_per_action(oof, label, label.columns)\n#     else:\n#         action_thresholds = {col: 0.27 for col in label.columns}\n    \n#     # Make multiclass predictions\n#     submission_part = predict_multiclass_enhanced(oof, meta, action_thresholds)\n#     submission_list.append(submission_part)\n    \n#     # Visualize metrics\n#     if verbose and action_metrics:\n#         fig, ax = plt.subplots(figsize=(10, 6))\n#         metrics_df = pd.DataFrame(action_metrics)\n        \n#         x = np.arange(len(metrics_df))\n#         width = 0.25\n        \n#         ax.bar(x - width, metrics_df['f1'], width, label='F1', color='blue')\n#         ax.bar(x, metrics_df['precision'], width, label='Precision', color='green')\n#         ax.bar(x + width, metrics_df['recall'], width, label='Recall', color='orange')\n#         ax.axhline(y=metrics_df['baseline'].mean(), color='red', \n#                   linestyle='--', label='Avg Baseline')\n        \n#         ax.set_xlabel('Actions')\n#         ax.set_ylabel('Score')\n#         ax.set_title('Performance Metrics by Action')\n#         ax.set_xticks(x)\n#         ax.set_xticklabels(metrics_df['action'], rotation=45, ha='right')\n#         ax.legend()\n#         ax.grid(True, alpha=0.3)\n        \n#         plt.tight_layout()\n#         plt.show()\n\n# # ========================================\n# # MAIN PROCESSING LOOP\n# # ========================================\n\n# print(\"\\n\" + \"=\"*60)\n# print(\"STARTING MAIN PROCESSING\")\n# print(\"=\"*60)\n\n# f1_list = []\n# submission_list = []\n\n# # Process each body part configuration\n# for section in range(1, min(len(body_parts_tracked_list), 10)):  # Limit for demonstration\n#     body_parts_tracked_str = body_parts_tracked_list[section]\n    \n#     try:\n#         body_parts_tracked = json.loads(body_parts_tracked_str)\n        \n#         print(f\"\\n{'='*60}\")\n#         print(f\"Section {section}: Processing {len(body_parts_tracked)} body parts\")\n#         print(f\"Body parts: {', '.join(body_parts_tracked[:5])}...\")\n#         print(f\"{'='*60}\")\n        \n#         # Filter training data\n#         train_subset = train[train.body_parts_tracked == body_parts_tracked_str]\n        \n#         if len(train_subset) == 0:\n#             print(\"⚠️ No training data for this configuration\")\n#             continue\n        \n#         print(f\"📊 Found {len(train_subset)} training videos\")\n        \n#         # Initialize data containers\n#         single_mouse_list = []\n#         single_mouse_label_list = []\n#         single_mouse_meta_list = []\n#         mouse_pair_list = []\n#         mouse_pair_label_list = []\n#         mouse_pair_meta_list = []\n        \n#         # Generate data\n#         print(\"📥 Loading training data...\")\n#         data_generator = generate_mouse_data(train_subset, 'train')\n        \n#         for switch, data, meta, label in data_generator:\n#             if switch == 'single':\n#                 single_mouse_list.append(data)\n#                 single_mouse_meta_list.append(meta)\n#                 single_mouse_label_list.append(label)\n#             else:\n#                 mouse_pair_list.append(data)\n#                 mouse_pair_meta_list.append(meta)\n#                 mouse_pair_label_list.append(label)\n        \n#         # Create classifier pipeline\n#         binary_classifier = create_classifier_pipeline(use_ensemble=use_ensemble)\n        \n#         # Process single mouse actions\n#         if len(single_mouse_list) > 0:\n#             print(f\"\\n🐭 Processing {len(single_mouse_list)} single mouse batches\")\n            \n#             # Concatenate data\n#             single_mouse = pd.concat(single_mouse_list)\n#             single_mouse_label = pd.concat(single_mouse_label_list)\n#             single_mouse_meta = pd.concat(single_mouse_meta_list)\n            \n#             del single_mouse_list, single_mouse_label_list, single_mouse_meta_list\n            \n#             # Transform features\n#             print(\"🔄 Transforming features...\")\n#             X_tr = transform_single_enhanced(single_mouse, body_parts_tracked)\n#             del single_mouse\n            \n#             print(f\"✅ Feature shape: {X_tr.shape}\")\n#             print(f\"   Actions: {list(single_mouse_label.columns)}\")\n            \n#             if validate_or_submit == 'validate':\n#                 cross_validate_classifier_enhanced(\n#                     binary_classifier, X_tr, single_mouse_label,\n#                     single_mouse_meta, body_parts_tracked_str\n#                 )\n            \n#             del X_tr\n#             gc.collect()\n        \n#         # Process mouse pair actions\n#         if len(mouse_pair_list) > 0:\n#             print(f\"\\n🐭🐭 Processing {len(mouse_pair_list)} mouse pair batches\")\n            \n#             # Concatenate data\n#             mouse_pair = pd.concat(mouse_pair_list)\n#             mouse_pair_label = pd.concat(mouse_pair_label_list)\n#             mouse_pair_meta = pd.concat(mouse_pair_meta_list)\n            \n#             del mouse_pair_list, mouse_pair_label_list, mouse_pair_meta_list\n            \n#             # Transform features\n#             print(\"🔄 Transforming features...\")\n#             X_tr = transform_pair_enhanced(mouse_pair, body_parts_tracked)\n#             del mouse_pair\n            \n#             print(f\"✅ Feature shape: {X_tr.shape}\")\n#             print(f\"   Actions: {list(mouse_pair_label.columns)}\")\n            \n#             if validate_or_submit == 'validate':\n#                 cross_validate_classifier_enhanced(\n#                     binary_classifier, X_tr, mouse_pair_label,\n#                     mouse_pair_meta, body_parts_tracked_str\n#                 )\n            \n#             del X_tr\n#             gc.collect()\n            \n#     except Exception as e:\n#         print(f\"❌ Error in section {section}: {e}\")\n#         import traceback\n#         traceback.print_exc()\n#         continue\n\n# # ========================================\n# # FINAL EVALUATION AND OUTPUT\n# # ========================================\n\n# print(\"\\n\" + \"=\"*60)\n# print(\"FINAL RESULTS\")\n# print(\"=\"*60)\n\n# if validate_or_submit == 'validate' and submission_list:\n#     # Combine all submissions\n#     submission = pd.concat(submission_list)\n    \n#     print(f\"\\n📊 Validation Results:\")\n#     print(f\"   Total predictions: {len(submission)}\")\n#     print(f\"   Unique videos: {submission['video_id'].nunique()}\")\n#     print(f\"   Unique actions: {submission['action'].nunique()}\")\n    \n#     # Action distribution\n#     action_counts = submission['action'].value_counts()\n#     print(f\"\\n📈 Top 10 predicted actions:\")\n#     for action, count in action_counts.head(10).items():\n#         print(f\"   {action:20} {count:6} ({count/len(submission)*100:.1f}%)\")\n    \n#     # Calculate final score if solution available\n#     try:\n#         print(f\"\\n🏆 Competition Score: {score(solution, submission, ''):.4f}\")\n#     except:\n#         print(\"\\n⚠️ Could not calculate competition score\")\n    \n#     # Average F1 scores\n#     if f1_list:\n#         f1_df = pd.DataFrame(f1_list, columns=['body_parts', 'action', 'f1_score'])\n#         print(f\"\\n📊 Binary F1 Statistics:\")\n#         print(f\"   Mean F1: {f1_df['f1_score'].mean():.4f}\")\n#         print(f\"   Std F1: {f1_df['f1_score'].std():.4f}\")\n#         print(f\"   Min F1: {f1_df['f1_score'].min():.4f}\")\n#         print(f\"   Max F1: {f1_df['f1_score'].max():.4f}\")\n        \n#         # Plot F1 distribution\n#         plt.figure(figsize=(10, 5))\n#         plt.subplot(1, 2, 1)\n#         plt.hist(f1_df['f1_score'], bins=30, edgecolor='black')\n#         plt.xlabel('F1 Score')\n#         plt.ylabel('Count')\n#         plt.title('Distribution of Binary F1 Scores')\n        \n#         plt.subplot(1, 2, 2)\n#         top_actions = f1_df.groupby('action')['f1_score'].mean().sort_values(ascending=False).head(15)\n#         plt.barh(range(len(top_actions)), top_actions.values)\n#         plt.yticks(range(len(top_actions)), top_actions.index)\n#         plt.xlabel('Mean F1 Score')\n#         plt.title('Top 15 Actions by Mean F1 Score')\n        \n#         plt.tight_layout()\n#         plt.show()\n\n# print(f\"\\n✅ Processing complete!\")\n# print(f\"End Time: {datetime.now()}\")\n# print(\"=\"*60)","metadata":{"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# from __future__ import annotations\n\n# import argparse\n# import json\n# import logging\n# import os\n# import warnings\n# from dataclasses import dataclass\n# from pathlib import Path\n# from typing import Dict, Iterable, List, Optional, Set, Tuple\n# from collections import defaultdict, Counter\n\n# import numpy as np\n# import pandas as pd\n# import polars as pl\n# from tqdm.auto import tqdm\n# import matplotlib.pyplot as plt\n# import seaborn as sns\n# from matplotlib.patches import Rectangle\n# import matplotlib.patches as mpatches\n\n# # Set visualization style\n# plt.style.use('seaborn-v0_8-whitegrid')\n# sns.set_palette(\"husl\")\n\n# # ========================\n# # Config\n# # ========================\n\n# @dataclass(frozen=True)\n# class Config:\n#     data_root: Path = Path(os.getenv(\"MABE_DATA_ROOT\", \"/kaggle/input/MABe-mouse-behavior-detection\"))\n#     submission_file: str = os.getenv(\"MABE_SUBMISSION\", \"submission.csv\")\n#     row_id_col: str = os.getenv(\"MABE_ROW_ID_COL\", \"row_id\")\n\n#     @property\n#     def train_csv(self) -> Path: return self.data_root / \"train.csv\"\n#     @property\n#     def test_csv(self) -> Path: return self.data_root / \"test.csv\"\n#     @property\n#     def train_annot_dir(self) -> Path: return self.data_root / \"train_annotation\"\n#     @property\n#     def train_track_dir(self) -> Path: return self.data_root / \"train_tracking\"\n#     @property\n#     def test_track_dir(self) -> Path: return self.data_root / \"test_tracking\"\n\n#     @property\n#     def submission_schema(self) -> Dict[str, pl.DataType]:\n#         return {\n#             \"video_id\": pl.Int64, \"agent_id\": pl.Utf8, \"target_id\": pl.Utf8,\n#             \"action\": pl.Utf8, \"start_frame\": pl.Int64, \"stop_frame\": pl.Int64,\n#         }\n\n#     @property\n#     def solution_schema(self) -> Dict[str, pl.DataType]:\n#         return {\n#             \"video_id\": pl.Int64, \"agent_id\": pl.Utf8, \"target_id\": pl.Utf8,\n#             \"action\": pl.Utf8, \"start_frame\": pl.Int64, \"stop_frame\": pl.Int64,\n#             \"lab_id\": pl.Utf8, \"behaviors_labeled\": pl.Utf8,\n#         }\n\n# logger = logging.getLogger(__name__)\n\n# class HostVisibleError(Exception): pass\n\n# def setup_logging(verbosity: int = 1) -> None:\n#     level = logging.WARNING if verbosity <= 0 else logging.INFO if verbosity == 1 else logging.DEBUG\n#     logging.basicConfig(level=level, format=\"%(asctime)s | %(levelname)s | %(name)s | %(message)s\", force=True)\n\n# # ========================\n# # Utils & Validators\n# # ========================\n\n# def safe_json_loads(s: Optional[str]) -> List[str]:\n#     if s is None: return []\n#     if isinstance(s, list): return [str(x) for x in s]\n#     if not isinstance(s, str): return []\n#     s = s.strip()\n#     if not s: return []\n#     try:\n#         return json.loads(s)\n#     except Exception:\n#         try: return json.loads(s.replace(\"'\", '\"'))\n#         except Exception: return []\n\n# def validate_schema(df: pl.DataFrame, schema: Dict[str, pl.DataType], name: str) -> pl.DataFrame:\n#     missing = set(schema.keys()) - set(df.columns)\n#     if missing: raise ValueError(f\"{name} is missing columns: {missing}\")\n#     casts = [pl.col(col).cast(dtype) for col, dtype in schema.items() if df[col].dtype != dtype]\n#     return df.with_columns(casts) if casts else df\n\n# def validate_frame_ranges(df: pl.DataFrame, name: str) -> None:\n#     if not (df[\"start_frame\"] <= df[\"stop_frame\"]).all():\n#         raise ValueError(f\"{name}: start_frame > stop_frame detected\")\n\n# def _norm_mouse_id(x: str | int) -> str:\n#     s = str(x)\n#     return s if s.startswith(\"mouse\") else f\"mouse{s}\"\n\n# def _norm_triplet(agent: str | int, target: str | int, action: str) -> str:\n#     return f\"{_norm_mouse_id(agent)},{_norm_mouse_id(target)},{action}\"\n\n# def _range_frames(start: int, stop: int) -> Iterable[int]:\n#     return range(start, stop)\n\n# def merge_intervals(intervals: List[Tuple[int,int]]) -> List[Tuple[int,int]]:\n#     if not intervals: return []\n#     intervals = sorted(intervals)\n#     merged = [intervals[0]]\n#     for s,e in intervals[1:]:\n#         ps,pe = merged[-1]\n#         if s <= pe: merged[-1] = (ps, max(pe, e))\n#         else: merged.append((s,e))\n#     return merged\n\n# # ========================\n# # Visualization Functions\n# # ========================\n\n# def visualize_data_overview(train_df, solution_df):\n#     \"\"\"Create comprehensive data overview visualization\"\"\"\n#     print(\"\\n📊 DATA OVERVIEW ANALYSIS\")\n#     print(\"=\" * 60)\n    \n#     fig, axes = plt.subplots(2, 3, figsize=(18, 10))\n    \n#     # 1. Lab distribution\n#     ax = axes[0, 0]\n#     lab_counts = train_df.group_by(\"lab_id\").count().sort(\"count\", descending=True)\n#     labs = lab_counts[\"lab_id\"].to_list()[:10]\n#     counts = lab_counts[\"count\"].to_list()[:10]\n#     colors = sns.color_palette(\"husl\", len(labs))\n#     ax.bar(range(len(labs)), counts, color=colors)\n#     ax.set_xticks(range(len(labs)))\n#     ax.set_xticklabels(labs, rotation=45, ha='right')\n#     ax.set_title(\"Videos per Lab (Top 10)\", fontsize=12, fontweight='bold')\n#     ax.set_ylabel(\"Number of Videos\")\n#     ax.grid(True, alpha=0.3)\n    \n#     # 2. Action distribution\n#     ax = axes[0, 1]\n#     action_counts = solution_df.group_by(\"action\").count().sort(\"count\", descending=True)\n#     actions = action_counts[\"action\"].to_list()\n#     action_cnts = action_counts[\"count\"].to_list()\n#     ax.bar(range(len(actions)), action_cnts, color=sns.color_palette(\"viridis\", len(actions)))\n#     ax.set_xticks(range(len(actions)))\n#     ax.set_xticklabels(actions, rotation=45, ha='right')\n#     ax.set_title(\"Behavior Frequency Distribution\", fontsize=12, fontweight='bold')\n#     ax.set_ylabel(\"Count\")\n#     ax.grid(True, alpha=0.3)\n    \n#     # 3. Duration distribution\n#     ax = axes[0, 2]\n#     solution_df = solution_df.with_columns((pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\"))\n#     durations = solution_df[\"duration\"].to_list()\n#     ax.hist(durations, bins=50, color='skyblue', edgecolor='black', alpha=0.7)\n#     ax.set_title(\"Behavior Duration Distribution\", fontsize=12, fontweight='bold')\n#     ax.set_xlabel(\"Duration (frames)\")\n#     ax.set_ylabel(\"Frequency\")\n#     ax.set_xlim(0, min(500, max(durations)))\n#     ax.axvline(np.median(durations), color='red', linestyle='--', label=f'Median: {np.median(durations):.0f}')\n#     ax.legend()\n#     ax.grid(True, alpha=0.3)\n    \n#     # 4. Mouse pair interactions\n#     ax = axes[1, 0]\n#     pair_counts = solution_df.group_by([\"agent_id\", \"target_id\"]).count().sort(\"count\", descending=True)\n#     top_pairs = pair_counts.head(10)\n#     pair_labels = [f\"{row['agent_id']}->{row['target_id']}\" for row in top_pairs.to_dicts()]\n#     pair_values = top_pairs[\"count\"].to_list()\n#     ax.barh(range(len(pair_labels)), pair_values, color='coral')\n#     ax.set_yticks(range(len(pair_labels)))\n#     ax.set_yticklabels(pair_labels)\n#     ax.set_title(\"Top Mouse Pair Interactions\", fontsize=12, fontweight='bold')\n#     ax.set_xlabel(\"Number of Behaviors\")\n#     ax.grid(True, alpha=0.3)\n    \n#     # 5. Behaviors per video\n#     ax = axes[1, 1]\n#     behaviors_per_video = solution_df.group_by(\"video_id\").count()[\"count\"].to_list()\n#     ax.hist(behaviors_per_video, bins=30, color='lightgreen', edgecolor='black', alpha=0.7)\n#     ax.set_title(\"Behaviors per Video Distribution\", fontsize=12, fontweight='bold')\n#     ax.set_xlabel(\"Number of Behaviors\")\n#     ax.set_ylabel(\"Number of Videos\")\n#     ax.axvline(np.mean(behaviors_per_video), color='red', linestyle='--', \n#                label=f'Mean: {np.mean(behaviors_per_video):.1f}')\n#     ax.legend()\n#     ax.grid(True, alpha=0.3)\n    \n#     # 6. Summary statistics\n#     ax = axes[1, 2]\n#     ax.axis('off')\n#     stats_text = f\"\"\"\n#     DATASET STATISTICS\n#     ==================\n#     Total Videos: {train_df[\"video_id\"].n_unique():,}\n#     Total Labs: {train_df[\"lab_id\"].n_unique()}\n#     Total Annotations: {len(solution_df):,}\n#     Unique Behaviors: {solution_df[\"action\"].n_unique()}\n    \n#     BEHAVIOR STATS\n#     ==============\n#     Mean Duration: {solution_df[\"duration\"].mean():.1f} frames\n#     Median Duration: {solution_df[\"duration\"].median():.1f} frames\n#     Max Duration: {solution_df[\"duration\"].max():.0f} frames\n    \n#     TOP 3 BEHAVIORS\n#     ==============\n#     1. {actions[0]}: {action_cnts[0]:,} ({action_cnts[0]/sum(action_cnts)*100:.1f}%)\n#     2. {actions[1]}: {action_cnts[1]:,} ({action_cnts[1]/sum(action_cnts)*100:.1f}%)\n#     3. {actions[2]}: {action_cnts[2]:,} ({action_cnts[2]/sum(action_cnts)*100:.1f}%)\n#     \"\"\"\n#     ax.text(0.05, 0.95, stats_text, transform=ax.transAxes, fontsize=10,\n#             verticalalignment='top', fontfamily='monospace',\n#             bbox=dict(boxstyle='round', facecolor='lightgray', alpha=0.3))\n    \n#     plt.suptitle(\"MABe Dataset Comprehensive Analysis\", fontsize=14, fontweight='bold')\n#     plt.tight_layout()\n#     plt.show()\n    \n#     # Print insights\n#     print(f\"📈 Key Insights:\")\n#     print(f\"   • Dataset contains {train_df['video_id'].n_unique()} unique videos from {train_df['lab_id'].n_unique()} labs\")\n#     print(f\"   • Total of {len(solution_df):,} behavior annotations\")\n#     print(f\"   • {solution_df['action'].n_unique()} distinct behavior types observed\")\n#     print(f\"   • Average behavior duration: {solution_df['duration'].mean():.1f} frames (~{solution_df['duration'].mean()/30:.1f} seconds @ 30fps)\")\n#     print(f\"   • Most common behavior: '{actions[0]}' ({action_cnts[0]/sum(action_cnts)*100:.1f}% of all behaviors)\")\n\n# def visualize_priors(per_lab_weight, global_weight, per_lab_med_dur, global_med_dur):\n#     \"\"\"Visualize computed priors\"\"\"\n#     print(\"\\n🎯 PRIOR DISTRIBUTIONS ANALYSIS\")\n#     print(\"=\" * 60)\n    \n#     fig, axes = plt.subplots(2, 2, figsize=(15, 10))\n    \n#     # 1. Global weights\n#     ax = axes[0, 0]\n#     actions = list(global_weight.keys())\n#     weights = list(global_weight.values())\n#     colors = sns.color_palette(\"husl\", len(actions))\n#     bars = ax.bar(range(len(actions)), weights, color=colors)\n#     ax.set_xticks(range(len(actions)))\n#     ax.set_xticklabels(actions, rotation=45, ha='right')\n#     ax.set_title(\"Global Action Weight Priors\", fontsize=12, fontweight='bold')\n#     ax.set_ylabel(\"Weight (Proportion)\")\n#     ax.grid(True, alpha=0.3)\n    \n#     # Add value labels\n#     for bar, weight in zip(bars, weights):\n#         height = bar.get_height()\n#         ax.text(bar.get_x() + bar.get_width()/2., height,\n#                 f'{weight:.3f}', ha='center', va='bottom', fontsize=8)\n    \n#     # 2. Global median durations\n#     ax = axes[0, 1]\n#     durations = [global_med_dur.get(a, 0) for a in actions]\n#     ax.barh(range(len(actions)), durations, color=colors)\n#     ax.set_yticks(range(len(actions)))\n#     ax.set_yticklabels(actions)\n#     ax.set_title(\"Global Median Durations\", fontsize=12, fontweight='bold')\n#     ax.set_xlabel(\"Frames\")\n#     ax.grid(True, alpha=0.3)\n    \n#     # 3. Lab-specific weight heatmap\n#     ax = axes[1, 0]\n#     labs = list(per_lab_weight.keys())[:10]  # Top 10 labs\n#     if labs:\n#         matrix = []\n#         for lab in labs:\n#             row = [per_lab_weight[lab].get(a, 0) for a in actions]\n#             matrix.append(row)\n        \n#         im = ax.imshow(matrix, aspect='auto', cmap='YlOrRd', vmin=0, vmax=max(weights))\n#         ax.set_xticks(range(len(actions)))\n#         ax.set_xticklabels(actions, rotation=45, ha='right')\n#         ax.set_yticks(range(len(labs)))\n#         ax.set_yticklabels(labs)\n#         ax.set_title(\"Lab-Specific Weight Priors (Top 10 Labs)\", fontsize=12, fontweight='bold')\n#         plt.colorbar(im, ax=ax, label='Weight')\n    \n#     # 4. Weight vs Duration correlation\n#     ax = axes[1, 1]\n#     x_weights = [global_weight.get(a, 0) for a in actions]\n#     y_durations = [global_med_dur.get(a, 0) for a in actions]\n#     ax.scatter(x_weights, y_durations, s=100, alpha=0.6, c=colors)\n#     for i, action in enumerate(actions):\n#         ax.annotate(action, (x_weights[i], y_durations[i]), \n#                    fontsize=8, ha='center', va='bottom')\n#     ax.set_xlabel(\"Global Weight\")\n#     ax.set_ylabel(\"Median Duration (frames)\")\n#     ax.set_title(\"Action Weight vs Duration Correlation\", fontsize=12, fontweight='bold')\n#     ax.grid(True, alpha=0.3)\n    \n#     plt.suptitle(\"Prior Distribution Analysis\", fontsize=14, fontweight='bold')\n#     plt.tight_layout()\n#     plt.show()\n    \n#     print(f\"📊 Prior Statistics:\")\n#     print(f\"   • Most weighted action: '{max(global_weight, key=global_weight.get)}' ({max(global_weight.values()):.3f})\")\n#     print(f\"   • Longest median duration: '{max(global_med_dur, key=global_med_dur.get)}' ({max(global_med_dur.values())} frames)\")\n#     print(f\"   • Number of labs with specific priors: {len(per_lab_weight)}\")\n\n# def visualize_predictions(predictions_df, test_df):\n#     \"\"\"Visualize prediction results\"\"\"\n#     print(\"\\n🔮 PREDICTION ANALYSIS\")\n#     print(\"=\" * 60)\n    \n#     fig, axes = plt.subplots(2, 3, figsize=(18, 10))\n    \n#     # 1. Predicted action distribution\n#     ax = axes[0, 0]\n#     pred_actions = predictions_df.group_by(\"action\").count().sort(\"count\", descending=True)\n#     actions = pred_actions[\"action\"].to_list()\n#     counts = pred_actions[\"count\"].to_list()\n#     colors = sns.color_palette(\"viridis\", len(actions))\n#     ax.bar(range(len(actions)), counts, color=colors)\n#     ax.set_xticks(range(len(actions)))\n#     ax.set_xticklabels(actions, rotation=45, ha='right')\n#     ax.set_title(\"Predicted Behavior Distribution\", fontsize=12, fontweight='bold')\n#     ax.set_ylabel(\"Count\")\n#     ax.grid(True, alpha=0.3)\n    \n#     # 2. Prediction duration distribution\n#     ax = axes[0, 1]\n#     predictions_df = predictions_df.with_columns((pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\"))\n#     durations = predictions_df[\"duration\"].to_list()\n#     ax.hist(durations, bins=30, color='lightcoral', edgecolor='black', alpha=0.7)\n#     ax.set_title(\"Predicted Duration Distribution\", fontsize=12, fontweight='bold')\n#     ax.set_xlabel(\"Duration (frames)\")\n#     ax.set_ylabel(\"Frequency\")\n#     ax.axvline(np.median(durations), color='red', linestyle='--', \n#                label=f'Median: {np.median(durations):.0f}')\n#     ax.legend()\n#     ax.grid(True, alpha=0.3)\n    \n#     # 3. Predictions per video\n#     ax = axes[0, 2]\n#     preds_per_video = predictions_df.group_by(\"video_id\").count()[\"count\"].to_list()\n#     if preds_per_video:\n#         ax.hist(preds_per_video, bins=min(20, len(set(preds_per_video))), \n#                 color='lightgreen', edgecolor='black', alpha=0.7)\n#         ax.set_title(\"Predictions per Video\", fontsize=12, fontweight='bold')\n#         ax.set_xlabel(\"Number of Predictions\")\n#         ax.set_ylabel(\"Count\")\n#         ax.axvline(np.mean(preds_per_video), color='red', linestyle='--',\n#                   label=f'Mean: {np.mean(preds_per_video):.1f}')\n#         ax.legend()\n#     ax.grid(True, alpha=0.3)\n    \n#     # 4. Timeline visualization for sample video\n#     ax = axes[1, 0]\n#     if len(predictions_df) > 0:\n#         sample_vid = predictions_df[\"video_id\"].to_list()[0]\n#         vid_preds = predictions_df.filter(pl.col(\"video_id\") == sample_vid).sort(\"start_frame\")\n        \n#         y_pos = 0\n#         color_map = dict(zip(actions, colors))\n        \n#         for row in vid_preds.to_dicts()[:10]:  # Show first 10 predictions\n#             rect = Rectangle((row[\"start_frame\"], y_pos - 0.4), \n#                            row[\"stop_frame\"] - row[\"start_frame\"], 0.8,\n#                            facecolor=color_map.get(row[\"action\"], 'gray'),\n#                            edgecolor='black', linewidth=0.5)\n#             ax.add_patch(rect)\n#             y_pos += 1\n        \n#         ax.set_ylim(-0.5, y_pos - 0.5)\n#         ax.set_xlabel(\"Frame\")\n#         ax.set_ylabel(\"Prediction #\")\n#         ax.set_title(f\"Sample Timeline (Video {sample_vid})\", fontsize=12, fontweight='bold')\n#         ax.grid(True, alpha=0.3)\n    \n#     # 5. Coverage analysis\n#     ax = axes[1, 1]\n#     coverage_data = []\n#     for vid in predictions_df[\"video_id\"].unique()[:20]:  # Sample 20 videos\n#         vid_preds = predictions_df.filter(pl.col(\"video_id\") == vid)\n#         if len(vid_preds) > 0:\n#             total_frames = vid_preds.select((pl.col(\"stop_frame\") - pl.col(\"start_frame\")).sum()).item()\n#             max_frame = vid_preds[\"stop_frame\"].max()\n#             min_frame = vid_preds[\"start_frame\"].min()\n#             if max_frame > min_frame:\n#                 coverage = total_frames / (max_frame - min_frame)\n#                 coverage_data.append(min(coverage, 1.0))\n    \n#     if coverage_data:\n#         ax.hist(coverage_data, bins=20, color='skyblue', edgecolor='black', alpha=0.7)\n#         ax.set_title(\"Frame Coverage Distribution\", fontsize=12, fontweight='bold')\n#         ax.set_xlabel(\"Coverage Ratio\")\n#         ax.set_ylabel(\"Number of Videos\")\n#         ax.axvline(np.mean(coverage_data), color='red', linestyle='--',\n#                   label=f'Mean: {np.mean(coverage_data):.2f}')\n#         ax.legend()\n#     ax.grid(True, alpha=0.3)\n    \n#     # 6. Summary\n#     ax = axes[1, 2]\n#     ax.axis('off')\n#     stats_text = f\"\"\"\n#     PREDICTION SUMMARY\n#     ==================\n#     Total Predictions: {len(predictions_df):,}\n#     Unique Videos: {predictions_df['video_id'].n_unique()}\n#     Unique Actions: {predictions_df['action'].n_unique()}\n    \n#     DURATION STATS\n#     ==============\n#     Mean: {predictions_df['duration'].mean():.1f} frames\n#     Median: {predictions_df['duration'].median():.1f} frames\n#     Std: {predictions_df['duration'].std():.1f} frames\n    \n#     TOP PREDICTED\n#     =============\n#     1. {actions[0] if actions else 'N/A'}: {counts[0] if counts else 0}\n#     2. {actions[1] if len(actions) > 1 else 'N/A'}: {counts[1] if len(counts) > 1 else 0}\n#     3. {actions[2] if len(actions) > 2 else 'N/A'}: {counts[2] if len(counts) > 2 else 0}\n    \n#     COVERAGE\n#     ========\n#     Avg Coverage: {np.mean(coverage_data)*100:.1f}% (sampled)\n#     Predictions/Video: {len(predictions_df)/predictions_df['video_id'].n_unique():.1f}\n#     \"\"\"\n#     ax.text(0.05, 0.95, stats_text, transform=ax.transAxes, fontsize=10,\n#             verticalalignment='top', fontfamily='monospace',\n#             bbox=dict(boxstyle='round', facecolor='lightgray', alpha=0.3))\n    \n#     plt.suptitle(\"Prediction Results Analysis\", fontsize=14, fontweight='bold')\n#     plt.tight_layout()\n#     plt.show()\n    \n#     print(f\"✅ Prediction Summary:\")\n#     print(f\"   • Generated {len(predictions_df):,} predictions\")\n#     print(f\"   • Average duration: {predictions_df['duration'].mean():.1f} frames\")\n#     print(f\"   • Most predicted: '{actions[0] if actions else 'N/A'}'\")\n\n# def print_detailed_insights(solution_df, predictions_df, per_lab_weight, global_weight):\n#     \"\"\"Print detailed insights and recommendations\"\"\"\n#     print(\"\\n\" + \"=\"*60)\n#     print(\"💡 DETAILED INSIGHTS & RECOMMENDATIONS\")\n#     print(\"=\"*60)\n    \n#     # Behavior pattern analysis\n#     print(\"\\n📊 BEHAVIOR PATTERNS:\")\n#     solution_df = solution_df.with_columns((pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\"))\n    \n#     for action in solution_df[\"action\"].unique()[:5]:  # Top 5 actions\n#         action_data = solution_df.filter(pl.col(\"action\") == action)\n#         print(f\"\\n   {action.upper()}:\")\n#         print(f\"   • Count: {len(action_data):,}\")\n#         print(f\"   • Avg Duration: {action_data['duration'].mean():.1f} frames\")\n#         print(f\"   • Std Duration: {action_data['duration'].std():.1f} frames\")\n#         print(f\"   • Weight: {global_weight.get(action, 0):.3f}\")\n    \n#     # Model performance indicators\n#     print(\"\\n🎯 MODEL PERFORMANCE INDICATORS:\")\n#     print(f\"   • Proximity windows: ENABLED (improves interaction detection)\")\n#     print(f\"   • Prior scope: MIXED (lab-specific with global fallback)\")\n#     print(f\"   • Min segment length: 15 frames (filters noise)\")\n#     print(f\"   • Gap closing: 3 frames (smooths predictions)\")\n    \n#     # Recommendations\n#     print(\"\\n💡 RECOMMENDATIONS FOR IMPROVEMENT:\")\n#     print(\"   1. Consider tuning window thresholds per behavior type\")\n#     print(\"   2. Add velocity/acceleration features for better motion capture\")\n#     print(\"   3. Model sequential dependencies (approach → chase → attack)\")\n#     print(\"   4. Implement ensemble with different parameter settings\")\n#     print(\"   5. Use ML models for refinement of heuristic predictions\")\n    \n#     # Data quality insights\n#     print(\"\\n📈 DATA QUALITY INSIGHTS:\")\n#     labs_with_data = solution_df[\"lab_id\"].n_unique()\n#     total_frames = solution_df.select((pl.col(\"stop_frame\") - pl.col(\"start_frame\")).sum()).item()\n#     print(f\"   • Labs with annotations: {labs_with_data}\")\n#     print(f\"   • Total annotated frames: {total_frames:,}\")\n#     print(f\"   • Average annotations per video: {len(solution_df)/solution_df['video_id'].n_unique():.1f}\")\n    \n#     # Rare behaviors\n#     action_counts = solution_df.group_by(\"action\").count().sort(\"count\")\n#     rare_actions = action_counts.head(3)[\"action\"].to_list()\n#     print(f\"\\n⚠️ RARE BEHAVIORS (may need special handling):\")\n#     for action in rare_actions:\n#         count = action_counts.filter(pl.col(\"action\") == action)[\"count\"].to_list()[0]\n#         print(f\"   • {action}: {count} occurrences\")\n\n# # ========================\n# # Core Functions\n# # ========================\n\n# def create_solution_df(dataset: pl.DataFrame, cfg: Optional[Config] = None) -> pl.DataFrame:\n#     cfg = cfg or Config()\n#     records: List[pl.DataFrame] = []\n#     for row in tqdm(dataset.to_dicts(), total=len(dataset), desc=\"Building solution\"):\n#         lab_id: str = row[\"lab_id\"]\n#         if lab_id.startswith(\"MABe22\"): continue\n#         video_id: int = row[\"video_id\"]\n#         annot_path = cfg.train_annot_dir / lab_id / f\"{video_id}.parquet\"\n#         if not annot_path.exists():\n#             logger.warning(\"No annotations for %s\", annot_path)\n#             continue\n#         try:\n#             annot = pl.read_parquet(annot_path).with_columns(\n#                 [\n#                     pl.lit(lab_id).alias(\"lab_id\"),\n#                     pl.lit(video_id).alias(\"video_id\"),\n#                     pl.lit(row[\"behaviors_labeled\"]).alias(\"behaviors_labeled\"),\n#                     pl.concat_str([pl.lit(\"mouse\"), pl.col(\"agent_id\").cast(pl.Utf8)]).alias(\"agent_id\"),\n#                     pl.concat_str([pl.lit(\"mouse\"), pl.col(\"target_id\").cast(pl.Utf8)]).alias(\"target_id\"),\n#                 ]\n#             )\n#             for col, dtype in (cfg.solution_schema).items():\n#                 if col in annot.columns and annot[col].dtype != dtype:\n#                     annot = annot.with_columns(pl.col(col).cast(dtype))\n#             annot = annot.select([c for c in cfg.solution_schema.keys() if c in annot.columns])\n#             records.append(annot)\n#         except Exception as e:\n#             logger.error(\"Failed to load %s: %s\", annot_path, e)\n#             continue\n#     if not records: raise ValueError(\"No annotation files loaded.\")\n#     solution = pl.concat(records, how=\"vertical\")\n#     solution = validate_schema(solution, cfg.solution_schema, \"Solution\")\n#     return solution\n\n# def build_video_spans(dataset: pl.DataFrame, split: str, cfg: Optional[Config] = None) -> Dict[int, Tuple[int,int]]:\n#     cfg = cfg or Config()\n#     track_dir = cfg.train_track_dir if split == \"train\" else cfg.test_track_dir\n#     spans: Dict[int, Tuple[int,int]] = {}\n#     for row in tqdm(dataset.to_dicts(), total=len(dataset), desc=\"Scanning spans\"):\n#         lab_id = row[\"lab_id\"]\n#         if lab_id.startswith(\"MABe22\"): continue\n#         vid = row[\"video_id\"]\n#         path = track_dir / lab_id / f\"{vid}.parquet\"\n#         if not path.exists(): continue\n#         try:\n#             df = pl.read_parquet(path).select([\"video_frame\"])\n#             s = int(df[\"video_frame\"].min())\n#             e = int(df[\"video_frame\"].max()) + 1\n#             spans[int(vid)] = (s,e)\n#         except Exception as e:\n#             logger.warning(\"Span read failed for %s: %s\", path, e)\n#     return spans\n\n# def compute_action_priors(solution: pl.DataFrame, eps: float = 1.0) -> Tuple[Dict[str, Dict[str, float]], Dict[str, float], Dict[str, Dict[str, int]], Dict[str, int]]:\n#     sol = solution.with_columns((pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"dur\"))\n#     by_lab = sol.group_by([\"lab_id\", \"action\"]).agg(pl.col(\"dur\").sum().alias(\"dur_sum\"))\n#     global_ = sol.group_by([\"action\"]).agg(pl.col(\"dur\").sum().alias(\"dur_sum\"))\n#     actions = set(global_[\"action\"].to_list())\n\n#     per_lab_weight: Dict[str, Dict[str, float]] = defaultdict(dict)\n#     for lab in by_lab[\"lab_id\"].unique():\n#         sub = by_lab.filter(pl.col(\"lab_id\") == lab)\n#         dmap = {r[\"action\"]: float(r[\"dur_sum\"]) for r in sub.to_dicts()}\n#         for a in actions: dmap[a] = dmap.get(a, 0.0) + eps\n#         total = sum(dmap.values()) or 1.0\n#         per_lab_weight[str(lab)] = {a: dmap[a]/total for a in actions}\n\n#     gmap = {r[\"action\"]: float(r[\"dur_sum\"]) for r in global_.to_dicts()}\n#     for a in actions: gmap[a] = gmap.get(a, 0.0) + eps\n#     gtotal = sum(gmap.values()) or 1.0\n#     global_weight = {a: gmap[a]/gtotal for a in actions}\n\n#     med_by_lab = sol.group_by([\"lab_id\", \"action\"]).median().select([\"lab_id\",\"action\",\"dur\"])\n#     per_lab_med_dur: Dict[str, Dict[str, int]] = defaultdict(dict)\n#     for r in med_by_lab.to_dicts():\n#         per_lab_med_dur[str(r[\"lab_id\"])][str(r[\"action\"])] = int(r[\"dur\"])\n#     med_global = sol.group_by([\"action\"]).median().select([\"action\",\"dur\"])\n#     global_med_dur: Dict[str, int] = {r[\"action\"]: int(r[\"dur\"]) for r in med_global.to_dicts()}\n\n#     return per_lab_weight, global_weight, per_lab_med_dur, global_med_dur\n\n# def compute_timing_priors(solution: pl.DataFrame, video_spans: Dict[int, Tuple[int,int]]) -> Tuple[Dict[str, Dict[str, float]], Dict[str, float]]:\n#     def start_pct_func(row) -> float:\n#         vid = int(row[\"video_id\"])\n#         if vid not in video_spans: return 0.5\n#         s,e = video_spans[vid]\n#         denom = max(1, e - s)\n#         return float(max(0, min(1, (int(row[\"start_frame\"]) - s) / denom)))\n\n#     rows = []\n#     for r in solution.select([\"lab_id\",\"action\",\"video_id\",\"start_frame\"]).to_dicts():\n#         rows.append({\"lab_id\": r[\"lab_id\"], \"action\": r[\"action\"], \"start_pct\": start_pct_func(r)})\n#     df = pl.DataFrame(rows)\n#     by_lab = df.group_by([\"lab_id\",\"action\"]).median().select([\"lab_id\",\"action\",\"start_pct\"])\n#     per_lab: Dict[str, Dict[str, float]] = defaultdict(dict)\n#     for r in by_lab.to_dicts():\n#         per_lab[str(r[\"lab_id\"])][str(r[\"action\"])] = float(r[\"start_pct\"])\n#     g = df.group_by([\"action\"]).median().select([\"action\",\"start_pct\"])\n#     global_: Dict[str, float] = {r[\"action\"]: float(r[\"start_pct\"]) for r in g.to_dicts()}\n#     return per_lab, global_\n\n# def _strip_mouse_prefix(s: str | int) -> str:\n#     s = str(s)\n#     return s[5:] if s.startswith(\"mouse\") else s\n\n# def _pair_features(df: pl.DataFrame, agent_raw: str, target_raw: str, downsample: int = 1) -> Optional[pl.DataFrame]:\n#     frame_candidates = [\"video_frame\",\"frame\",\"frame_idx\"]\n#     id_candidates = [\"mouse_id\",\"id\",\"track_id\",\"agent_id\"]\n#     x_candidates = [\"x\",\"x_pos\",\"x_position\",\"x_mm\",\"centroid_x\",\"cx\"]\n#     y_candidates = [\"y\",\"y_pos\",\"y_position\",\"y_mm\",\"centroid_y\",\"cy\"]\n\n#     cols = set(df.columns)\n#     frame_col = next((c for c in frame_candidates if c in cols), None)\n#     id_col = next((c for c in id_candidates if c in cols), None)\n#     x_col = next((c for c in x_candidates if c in cols), None)\n#     y_col = next((c for c in y_candidates if c in cols), None)\n#     if not all([frame_col, id_col, x_col, y_col]):\n#         return None\n\n#     a_id = _strip_mouse_prefix(agent_raw)\n#     t_id = _strip_mouse_prefix(target_raw)\n\n#     pdf = df.select([frame_col, id_col, x_col, y_col]).to_pandas()\n#     pdf[frame_col] = pdf[frame_col].astype(np.int64, copy=False)\n#     pdf[id_col] = pdf[id_col].astype(str, copy=False)\n\n#     a = pdf[pdf[id_col] == a_id].copy()\n#     b = pdf[pdf[id_col] == t_id].copy()\n#     if a.empty or b.empty:\n#         return None\n\n#     a.drop_duplicates(subset=[frame_col], keep=\"first\", inplace=True)\n#     b.drop_duplicates(subset=[frame_col], keep=\"first\", inplace=True)\n\n#     merged = a.merge(b, on=frame_col, how=\"inner\", suffixes=(\"_a\", \"_b\"))\n#     if merged.empty:\n#         return None\n#     merged.sort_values(frame_col, inplace=True)\n\n#     ax = merged[f\"{x_col}_a\"].to_numpy(dtype=np.float64, copy=False)\n#     ay = merged[f\"{y_col}_a\"].to_numpy(dtype=np.float64, copy=False)\n#     bx = merged[f\"{x_col}_b\"].to_numpy(dtype=np.float64, copy=False)\n#     by = merged[f\"{y_col}_b\"].to_numpy(dtype=np.float64, copy=False)\n#     frames = merged[frame_col].to_numpy(dtype=np.int64, copy=False)\n\n#     if downsample > 1:\n#         sl = slice(0, None, int(downsample))\n#         ax, ay, bx, by, frames = ax[sl], ay[sl], bx[sl], by[sl], frames[sl]\n#         if ax.size == 0:\n#             return None\n\n#     dx = ax - bx\n#     dy = ay - by\n#     dist = np.sqrt(dx*dx + dy*dy)\n\n#     dax = np.diff(ax, prepend=ax[0])\n#     day = np.diff(ay, prepend=ay[0])\n#     dbx = np.diff(bx, prepend=bx[0])\n#     dby = np.diff(by, prepend=by[0])\n#     speed_a = np.sqrt(dax*dax + day*day)\n#     speed_b = np.sqrt(dbx*dbx + dby*dby)\n\n#     rel_speed = speed_a - speed_b\n#     ddist = np.diff(dist, prepend=dist[0])\n\n#     feat = pl.DataFrame(\n#         {\n#             \"frame\": frames,\n#             \"dist\": dist,\n#             \"rel_speed\": rel_speed,\n#             \"ddist\": ddist,\n#         }\n#     ).sort(\"frame\")\n\n#     return feat\n\n# def _make_windows(feat: pl.DataFrame, min_len: int, q_dist: float = 0.40, q_rel: float = 0.60, q_ddist: float = 0.40) -> List[Tuple[int,int]]:\n#     if len(feat) == 0:\n#         return []\n#     qd = float(feat[\"dist\"].quantile(q_dist))\n#     qr = float(feat[\"rel_speed\"].quantile(q_rel))\n#     qdd = float(feat[\"ddist\"].quantile(q_ddist))\n#     cond = (pl.col(\"dist\") <= qd) | ((pl.col(\"rel_speed\") >= qr) & (pl.col(\"ddist\") <= qdd))\n#     mask = feat.select(cond.alias(\"m\")).to_series().to_list()\n#     frames = feat[\"frame\"].to_list()\n\n#     windows: List[Tuple[int,int]] = []\n#     run: Optional[List[int]] = None\n#     for i, flag in enumerate(mask):\n#         if flag and run is None:\n#             run = [frames[i], frames[i]]\n#         elif flag and run is not None:\n#             run[1] = frames[i]\n#         elif (not flag) and run is not None:\n#             s,e = run[0], run[1]+1\n#             if e - s >= min_len:\n#                 windows.append((s,e))\n#             run = None\n#     if run is not None:\n#         s,e = run[0], run[1]+1\n#         if e - s >= min_len:\n#             windows.append((s,e))\n#     return merge_intervals(windows)\n\n# def _order_actions_by_timing(actions: List[str], lab_id: str,\n#                              timing_lab: Dict[str, Dict[str, float]],\n#                              timing_global: Dict[str, float],\n#                              canonical: Dict[str,int]) -> List[str]:\n#     def score(a: str) -> float:\n#         if lab_id in timing_lab and a in timing_lab[lab_id]:\n#             return timing_lab[lab_id][a]\n#         return timing_global.get(a, 0.5)\n#     return sorted(actions, key=lambda a: (score(a), canonical.get(a, 99)))\n\n# def _clip_rare_actions(weights_map: Dict[str,float], actions: List[str], p_min: float, cap: float) -> Dict[str,float]:\n#     w = {a: max(0.0, float(weights_map.get(a, 0.0))) for a in actions}\n#     for a in actions:\n#         if w[a] < p_min:\n#             w[a] = min(w[a], cap)\n#     s = sum(w.values()) or 1.0\n#     return {a: w[a]/s for a in actions}\n\n# def _allocate_segments_in_windows(windows: List[Tuple[int,int]],\n#                                   ordered_actions: List[str],\n#                                   weights: Dict[str,float],\n#                                   med_dur: Dict[str,int],\n#                                   total_frames: int) -> List[Tuple[str,int,int]]:\n#     win_idx = 0\n#     cur_s, cur_e = (windows[0] if windows else (0,0))\n#     remain = sum(e-s for s,e in windows)\n#     out: List[Tuple[str,int,int]] = []\n\n#     for a in ordered_actions:\n#         if remain <= 0: break\n#         want = int(weights.get(a, 0.0) * total_frames)\n#         want = max(want, int(med_dur.get(a, 0) or 0))\n#         want = min(want, remain)\n#         got = 0\n#         while got < want and win_idx < len(windows):\n#             s,e = cur_s, cur_e\n#             if s >= e:\n#                 win_idx += 1\n#                 if win_idx >= len(windows): break\n#                 cur_s, cur_e = windows[win_idx]\n#                 continue\n#             take = min(want - got, e - s)\n#             out.append((a, s, s+take))\n#             got += take\n#             remain -= take\n#             cur_s = s + take\n#             if cur_s >= e and win_idx < len(windows):\n#                 win_idx += 1\n#                 if win_idx < len(windows):\n#                     cur_s, cur_e = windows[win_idx]\n#     return out\n\n# def _smooth_segments(segments: List[Tuple[str,int,int]], min_len: int, gap_close: int) -> List[Tuple[str,int,int]]:\n#     if not segments: return []\n#     segments = sorted(segments, key=lambda x: (x[1], x[2], x[0]))\n#     segments = [seg for seg in segments if seg[2] - seg[1] >= min_len]\n#     if not segments: return []\n#     out = [segments[0]]\n#     for a,s,e in segments[1:]:\n#         pa,ps,pe = out[-1]\n#         if a == pa and s - pe <= gap_close:\n#             out[-1] = (a, ps, e)\n#         else:\n#             out.append((a,s,e))\n#     return out\n\n# def predict_without_ml(dataset: pl.DataFrame, data_split: str, cfg: Optional[Config] = None,\n#                        priors_per_lab: Optional[Dict[str, Dict[str, float]]] = None,\n#                        priors_global: Optional[Dict[str, float]] = None,\n#                        meddur_per_lab: Optional[Dict[str, Dict[str, int]]] = None,\n#                        meddur_global: Optional[Dict[str, int]] = None,\n#                        timing_lab: Optional[Dict[str, Dict[str, float]]] = None,\n#                        timing_global: Optional[Dict[str, float]] = None,\n#                        prior_scope: str = \"mixed\",\n#                        use_windows: bool = True,\n#                        min_len: int = 10,\n#                        gap_close: int = 5,\n#                        p_min: float = 0.03,\n#                        cap: float = 0.02) -> pl.DataFrame:\n#     cfg = cfg or Config()\n#     track_dir = cfg.test_track_dir if data_split == \"test\" else cfg.train_track_dir\n#     records: List[Tuple[int, str, str, str, int, int]] = []\n#     canonical = {\"approach\": 0, \"avoid\": 1, \"chase\": 2, \"chaseattack\": 3, \"attack\": 4, \"mount\": 5, \"submit\": 6}\n\n#     for row in tqdm(dataset.to_dicts(), total=len(dataset), desc=f\"Predicting ({data_split})\"):\n#         lab_id: str = row[\"lab_id\"]\n#         if lab_id.startswith(\"MABe22\"): continue\n#         video_id: int = row[\"video_id\"]\n#         path = track_dir / lab_id / f\"{video_id}.parquet\"\n#         if not path.exists():\n#             logger.warning(\"Tracking file not found: %s\", path)\n#             continue\n\n#         try:\n#             trk = pl.read_parquet(path)\n#             start_frame = int(trk[\"video_frame\"].min())\n#             stop_frame = int(trk[\"video_frame\"].max()) + 1\n#             video_frames = stop_frame - start_frame\n#             if video_frames <= 0: continue\n\n#             raw_list = safe_json_loads(row[\"behaviors_labeled\"])\n#             triples: List[List[str]] = []\n#             for b in raw_list:\n#                 parts = [p.strip() for p in str(b).replace(\"'\", \"\").split(\",\")]\n#                 if len(parts) == 3:\n#                     triples.append(parts)\n#             if not triples:\n#                 continue\n\n#             beh_df = pl.DataFrame(triples, schema=[\"agent\",\"target\",\"action\"], orient=\"row\").with_columns(\n#                 [pl.col(\"agent\").cast(pl.Utf8), pl.col(\"target\").cast(pl.Utf8), pl.col(\"action\").cast(pl.Utf8)]\n#             )\n\n#             for (agent, target), group in beh_df.group_by([\"agent\",\"target\"]):\n#                 actions = sorted(list(set(group[\"action\"].to_list())), key=lambda a: canonical.get(a, 99))\n#                 if not actions: continue\n\n#                 if prior_scope == \"lab\" and priors_per_lab is not None:\n#                     w_map = priors_per_lab.get(str(lab_id), {})\n#                     md_map = meddur_per_lab.get(str(lab_id), {}) if meddur_per_lab else {}\n#                 elif prior_scope == \"global\" and priors_global is not None:\n#                     w_map = priors_global\n#                     md_map = meddur_global or {}\n#                 else:\n#                     w_map = (priors_per_lab or {}).get(str(lab_id), {}) or (priors_global or {})\n#                     md_map = (meddur_per_lab or {}).get(str(lab_id), {}) or (meddur_global or {})\n\n#                 weights = _clip_rare_actions(w_map, actions, p_min=p_min, cap=cap)\n#                 ordered_actions = _order_actions_by_timing(\n#                     actions, str(lab_id), timing_lab or {}, timing_global or {}, canonical\n#                 )\n\n#                 windows: List[Tuple[int,int]]\n#                 if use_windows:\n#                     feat = _pair_features(trk, _norm_mouse_id(agent), _norm_mouse_id(target))\n#                     if feat is None:\n#                         windows = [(start_frame, stop_frame)]\n#                     else:\n#                         windows = _make_windows(feat, min_len=min_len)\n#                         if not windows:\n#                             windows = [(start_frame, stop_frame)]\n#                 else:\n#                     windows = [(start_frame, stop_frame)]\n\n#                 windows = merge_intervals(windows)\n#                 allowed_total = sum(e - s for s,e in windows)\n#                 if allowed_total <= 0:\n#                     continue\n\n#                 segs = _allocate_segments_in_windows(\n#                     windows=windows,\n#                     ordered_actions=ordered_actions,\n#                     weights=weights,\n#                     med_dur=md_map,\n#                     total_frames=allowed_total\n#                 )\n\n#                 segs = _smooth_segments(segs, min_len=min_len, gap_close=gap_close)\n\n#                 for a, s, e in segs:\n#                     if e > s:\n#                         records.append((\n#                             video_id,\n#                             _norm_mouse_id(agent), _norm_mouse_id(target),\n#                             a, int(s), int(e)\n#                         ))\n\n#         except Exception as e:\n#             logger.error(\"Error processing %s: %s\", path, e)\n#             continue\n\n#     if not records:\n#         raise ValueError(\"No predictions generated.\")\n\n#     df = pl.DataFrame(\n#         records,\n#         schema={\n#             \"video_id\": pl.Int64, \"agent_id\": pl.Utf8, \"target_id\": pl.Utf8,\n#             \"action\": pl.Utf8, \"start_frame\": pl.Int64, \"stop_frame\": pl.Int64,\n#         },\n#         orient=\"row\",\n#     )\n#     df = validate_schema(df, cfg.submission_schema, \"Submission\")\n#     validate_frame_ranges(df, \"Submission\")\n#     return df\n\n# # ========================\n# # Main execution with visualizations\n# # ========================\n\n# def main():\n#     print(\"=\"*60)\n#     print(\"🐭 MABe MOUSE BEHAVIOR DETECTION - ENHANCED ANALYSIS 🐭\")\n#     print(\"=\"*60)\n    \n#     # Setup\n#     setup_logging(verbosity=1)\n#     warnings.filterwarnings(\"ignore\")\n    \n#     cfg = Config()\n    \n#     # Parameters - separate beta from prediction params\n#     print(\"\\n⚙️ CONFIGURATION PARAMETERS:\")\n#     beta = 1.0\n#     prediction_params = {\n#         'prior_scope': \"mixed\",\n#         'eps': 0.5,\n#         'use_windows': True,\n#         'min_len': 15,\n#         'gap_close': 3,\n#         'p_min': 0.05,\n#         'cap': 0.05\n#     }\n    \n#     print(f\"   • beta (F-score): {beta}\")\n#     for key, value in prediction_params.items():\n#         print(f\"   • {key}: {value}\")\n    \n#     print(\"\\n📂 Loading datasets...\")\n#     train = pl.read_csv(cfg.train_csv)\n#     train_subset = train.filter(~pl.col(\"lab_id\").str.starts_with(\"MABe22\"))\n#     print(f\"   ✓ Loaded {len(train)} training samples ({len(train_subset)} after filtering)\")\n    \n#     print(\"\\n🔨 Building solution dataframe from annotations...\")\n#     solution = create_solution_df(train_subset, cfg)\n#     print(f\"   ✓ Processed {len(solution):,} annotations\")\n    \n#     # Data visualization\n#     visualize_data_overview(train_subset, solution)\n    \n#     print(\"\\n📐 Computing video spans...\")\n#     spans = build_video_spans(train_subset, \"train\", cfg)\n#     print(f\"   ✓ Computed spans for {len(spans)} videos\")\n    \n#     print(\"\\n🧮 Computing action priors...\")\n#     per_lab, global_w, med_lab, med_glob = compute_action_priors(solution, eps=prediction_params['eps'])\n#     timing_lab, timing_glob = compute_timing_priors(solution, spans)\n#     print(f\"   ✓ Computed priors for {len(global_w)} actions across {len(per_lab)} labs\")\n    \n#     # Visualize priors\n#     visualize_priors(per_lab, global_w, med_lab, med_glob)\n    \n#     print(\"\\n📥 Loading test data...\")\n#     test = pl.read_csv(cfg.test_csv)\n#     print(f\"   ✓ Loaded {len(test)} test samples\")\n    \n#     print(\"\\n🔮 Generating predictions...\")\n#     submission_test = predict_without_ml(\n#         test, \"test\", cfg,\n#         priors_per_lab=per_lab, priors_global=global_w,\n#         meddur_per_lab=med_lab, meddur_global=med_glob,\n#         timing_lab=timing_lab, timing_global=timing_glob,\n#         prior_scope=prediction_params['prior_scope'],\n#         use_windows=prediction_params['use_windows'],\n#         min_len=prediction_params['min_len'],\n#         gap_close=prediction_params['gap_close'],\n#         p_min=prediction_params['p_min'],\n#         cap=prediction_params['cap']\n#     )\n#     print(f\"   ✓ Generated {len(submission_test):,} predictions\")\n    \n#     # Visualize predictions\n#     visualize_predictions(submission_test, test)\n    \n#     # Add row_id and save\n#     ordered = list(cfg.submission_schema.keys())\n#     submission_test = submission_test.select(ordered).with_row_index(cfg.row_id_col)\n    \n#     print(f\"\\n💾 Saving submission to {cfg.submission_file}...\")\n#     submission_test.write_csv(cfg.submission_file)\n#     print(f\"   ✓ Submission saved successfully!\")\n    \n#     # Print detailed insights\n#     print_detailed_insights(solution, submission_test, per_lab, global_w)\n    \n#     print(\"\\n\" + \"=\"*60)\n#     print(\"✅ ANALYSIS COMPLETE!\")\n#     print(\"=\"*60)\n#     print(f\"\"\"\n#     FINAL SUMMARY:\n#     • Total predictions: {len(submission_test):,}\n#     • Output file: {cfg.submission_file}\n#     • Ready for submission!\n    \n#     Good luck with your competition! 🎯\n#     \"\"\")\n\n# if __name__ == \"__main__\":\n#     main()","metadata":{"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# \"\"\"\n# MABe Mouse Behavior Detection - Comprehensive Analysis Pipeline\n# ================================================================\n# Enhanced version with extensive EDA, visualizations, and insights\n# for the Multi-Agent Behavior (MABe) Challenge 2025\n\n# Author: Enhanced Analysis Pipeline\n# Date: 2025\n# \"\"\"\n\n# from __future__ import annotations\n\n# import argparse\n# import json\n# import logging\n# import os\n# import warnings\n# from dataclasses import dataclass\n# from pathlib import Path\n# from typing import Dict, Iterable, List, Optional, Set, Tuple, Any\n# from collections import defaultdict, Counter\n# from datetime import datetime\n# import pickle\n\n# import numpy as np\n# import pandas as pd\n# import polars as pl\n# from tqdm.auto import tqdm\n# import matplotlib.pyplot as plt\n# import matplotlib.patches as mpatches\n# from matplotlib.patches import Rectangle, Circle, FancyBboxPatch\n# from matplotlib.gridspec import GridSpec\n# import seaborn as sns\n# from scipy import stats, signal\n# from scipy.spatial.distance import pdist, squareform\n# from scipy.cluster.hierarchy import dendrogram, linkage\n# from sklearn.decomposition import PCA\n# from sklearn.manifold import TSNE\n# from sklearn.preprocessing import StandardScaler\n# import networkx as nx\n\n# # Set professional visualization style\n# plt.style.use('seaborn-v0_8-whitegrid')\n# sns.set_palette(\"husl\")\n# plt.rcParams['figure.dpi'] = 100\n# plt.rcParams['savefig.dpi'] = 300\n# plt.rcParams['font.size'] = 10\n# plt.rcParams['axes.titlesize'] = 12\n# plt.rcParams['axes.titleweight'] = 'bold'\n\n# # Suppress warnings for cleaner output\n# warnings.filterwarnings(\"ignore\")\n\n# # ========================\n# # Configuration\n# # ========================\n\n# @dataclass(frozen=True)\n# class Config:\n#     \"\"\"Enhanced configuration with additional analysis parameters\"\"\"\n#     data_root: Path = Path(os.getenv(\"MABE_DATA_ROOT\", \"/kaggle/input/MABe-mouse-behavior-detection\"))\n#     submission_file: str = os.getenv(\"MABE_SUBMISSION\", \"submission.csv\")\n#     row_id_col: str = os.getenv(\"MABE_ROW_ID_COL\", \"row_id\")\n    \n#     # Analysis parameters\n#     enable_advanced_eda: bool = True\n#     save_figures: bool = True\n#     figure_format: str = 'png'\n#     figure_dpi: int = 300\n    \n#     # Behavior analysis\n#     min_behavior_duration: int = 5  # frames\n#     max_behavior_duration: int = 1000  # frames\n#     behavior_transition_window: int = 30  # frames\n    \n#     # Statistical thresholds\n#     significance_level: float = 0.05\n#     correlation_threshold: float = 0.3\n#     outlier_z_score: float = 3.0\n\n#     @property\n#     def train_csv(self) -> Path: return self.data_root / \"train.csv\"\n#     @property\n#     def test_csv(self) -> Path: return self.data_root / \"test.csv\"\n#     @property\n#     def train_annot_dir(self) -> Path: return self.data_root / \"train_annotation\"\n#     @property\n#     def train_track_dir(self) -> Path: return self.data_root / \"train_tracking\"\n#     @property\n#     def test_track_dir(self) -> Path: return self.data_root / \"test_tracking\"\n#     @property\n#     def output_dir(self) -> Path: return Path(\"/kaggle/working\")\n\n#     @property\n#     def submission_schema(self) -> Dict[str, pl.DataType]:\n#         return {\n#             \"video_id\": pl.Int64, \"agent_id\": pl.Utf8, \"target_id\": pl.Utf8,\n#             \"action\": pl.Utf8, \"start_frame\": pl.Int64, \"stop_frame\": pl.Int64,\n#         }\n\n#     @property\n#     def solution_schema(self) -> Dict[str, pl.DataType]:\n#         return {\n#             \"video_id\": pl.Int64, \"agent_id\": pl.Utf8, \"target_id\": pl.Utf8,\n#             \"action\": pl.Utf8, \"start_frame\": pl.Int64, \"stop_frame\": pl.Int64,\n#             \"lab_id\": pl.Utf8, \"behaviors_labeled\": pl.Utf8,\n#         }\n\n# logger = logging.getLogger(__name__)\n\n# def setup_logging(verbosity: int = 1) -> None:\n#     \"\"\"Configure logging with enhanced formatting\"\"\"\n#     level = logging.WARNING if verbosity <= 0 else logging.INFO if verbosity == 1 else logging.DEBUG\n#     logging.basicConfig(\n#         level=level, \n#         format=\"%(asctime)s | %(levelname)-8s | %(name)-20s | %(message)s\",\n#         datefmt='%Y-%m-%d %H:%M:%S',\n#         force=True\n#     )\n\n# # ========================\n# # Utility Functions\n# # ========================\n\n# def safe_json_loads(s: Optional[str]) -> List[str]:\n#     \"\"\"Safely parse JSON strings\"\"\"\n#     if s is None: return []\n#     if isinstance(s, list): return [str(x) for x in s]\n#     if not isinstance(s, str): return []\n#     s = s.strip()\n#     if not s: return []\n#     try:\n#         return json.loads(s)\n#     except Exception:\n#         try: \n#             return json.loads(s.replace(\"'\", '\"'))\n#         except Exception: \n#             return []\n\n# def validate_schema(df: pl.DataFrame, schema: Dict[str, pl.DataType], name: str) -> pl.DataFrame:\n#     \"\"\"Validate and cast dataframe schema\"\"\"\n#     missing = set(schema.keys()) - set(df.columns)\n#     if missing: \n#         raise ValueError(f\"{name} is missing columns: {missing}\")\n#     casts = [pl.col(col).cast(dtype) for col, dtype in schema.items() if df[col].dtype != dtype]\n#     return df.with_columns(casts) if casts else df\n\n# def validate_frame_ranges(df: pl.DataFrame, name: str) -> None:\n#     \"\"\"Validate frame ranges are valid\"\"\"\n#     if not (df[\"start_frame\"] <= df[\"stop_frame\"]).all():\n#         raise ValueError(f\"{name}: start_frame > stop_frame detected\")\n\n# def _norm_mouse_id(x: str | int) -> str:\n#     \"\"\"Normalize mouse ID format\"\"\"\n#     s = str(x)\n#     return s if s.startswith(\"mouse\") else f\"mouse{s}\"\n\n# def merge_intervals(intervals: List[Tuple[int,int]]) -> List[Tuple[int,int]]:\n#     \"\"\"Merge overlapping intervals\"\"\"\n#     if not intervals: return []\n#     intervals = sorted(intervals)\n#     merged = [intervals[0]]\n#     for s, e in intervals[1:]:\n#         ps, pe = merged[-1]\n#         if s <= pe: \n#             merged[-1] = (ps, max(pe, e))\n#         else: \n#             merged.append((s, e))\n#     return merged\n\n# def calculate_iou(interval1: Tuple[int, int], interval2: Tuple[int, int]) -> float:\n#     \"\"\"Calculate Intersection over Union for two intervals\"\"\"\n#     s1, e1 = interval1\n#     s2, e2 = interval2\n    \n#     intersection = max(0, min(e1, e2) - max(s1, s2))\n#     union = max(e1, e2) - min(s1, s2)\n    \n#     return intersection / union if union > 0 else 0.0\n\n# # ========================\n# # Enhanced Data Loading and Processing\n# # ========================\n\n# class DataProcessor:\n#     \"\"\"Enhanced data processing with validation and statistics\"\"\"\n    \n#     def __init__(self, cfg: Config):\n#         self.cfg = cfg\n#         self.stats = {}\n    \n#     def create_solution_df(self, dataset: pl.DataFrame) -> pl.DataFrame:\n#         \"\"\"Create solution dataframe with enhanced error handling\"\"\"\n#         records: List[pl.DataFrame] = []\n#         failed_files = []\n        \n#         for row in tqdm(dataset.to_dicts(), total=len(dataset), desc=\"Building solution\"):\n#             lab_id: str = row[\"lab_id\"]\n#             if lab_id.startswith(\"MABe22\"): \n#                 continue\n                \n#             video_id: int = row[\"video_id\"]\n#             annot_path = self.cfg.train_annot_dir / lab_id / f\"{video_id}.parquet\"\n            \n#             if not annot_path.exists():\n#                 logger.warning(f\"No annotations for {annot_path}\")\n#                 failed_files.append(str(annot_path))\n#                 continue\n                \n#             try:\n#                 annot = pl.read_parquet(annot_path).with_columns([\n#                     pl.lit(lab_id).alias(\"lab_id\"),\n#                     pl.lit(video_id).alias(\"video_id\"),\n#                     pl.lit(row[\"behaviors_labeled\"]).alias(\"behaviors_labeled\"),\n#                     pl.concat_str([pl.lit(\"mouse\"), pl.col(\"agent_id\").cast(pl.Utf8)]).alias(\"agent_id\"),\n#                     pl.concat_str([pl.lit(\"mouse\"), pl.col(\"target_id\").cast(pl.Utf8)]).alias(\"target_id\"),\n#                 ])\n                \n#                 # Schema validation and casting\n#                 for col, dtype in self.cfg.solution_schema.items():\n#                     if col in annot.columns and annot[col].dtype != dtype:\n#                         annot = annot.with_columns(pl.col(col).cast(dtype))\n                \n#                 annot = annot.select([c for c in self.cfg.solution_schema.keys() if c in annot.columns])\n#                 records.append(annot)\n                \n#             except Exception as e:\n#                 logger.error(f\"Failed to load {annot_path}: {e}\")\n#                 failed_files.append(str(annot_path))\n#                 continue\n        \n#         if not records: \n#             raise ValueError(\"No annotation files loaded.\")\n        \n#         # Store statistics\n#         self.stats['failed_files'] = failed_files\n#         self.stats['success_rate'] = len(records) / len(dataset) if len(dataset) > 0 else 0\n        \n#         solution = pl.concat(records, how=\"vertical\")\n#         solution = validate_schema(solution, self.cfg.solution_schema, \"Solution\")\n        \n#         return solution\n    \n#     def build_video_spans(self, dataset: pl.DataFrame, split: str) -> Dict[int, Tuple[int, int]]:\n#         \"\"\"Build video frame spans with statistics\"\"\"\n#         track_dir = self.cfg.train_track_dir if split == \"train\" else self.cfg.test_track_dir\n#         spans: Dict[int, Tuple[int, int]] = {}\n#         frame_stats = []\n        \n#         for row in tqdm(dataset.to_dicts(), total=len(dataset), desc=\"Scanning spans\"):\n#             lab_id = row[\"lab_id\"]\n#             if lab_id.startswith(\"MABe22\"): \n#                 continue\n                \n#             vid = row[\"video_id\"]\n#             path = track_dir / lab_id / f\"{vid}.parquet\"\n            \n#             if not path.exists(): \n#                 continue\n                \n#             try:\n#                 df = pl.read_parquet(path).select([\"video_frame\"])\n#                 s = int(df[\"video_frame\"].min())\n#                 e = int(df[\"video_frame\"].max()) + 1\n#                 spans[int(vid)] = (s, e)\n#                 frame_stats.append(e - s)\n                \n#             except Exception as e:\n#                 logger.warning(f\"Span read failed for {path}: {e}\")\n        \n#         # Store frame statistics\n#         if frame_stats:\n#             self.stats['frame_stats'] = {\n#                 'mean': np.mean(frame_stats),\n#                 'std': np.std(frame_stats),\n#                 'min': np.min(frame_stats),\n#                 'max': np.max(frame_stats),\n#                 'median': np.median(frame_stats)\n#             }\n        \n#         return spans\n\n# # ========================\n# # Comprehensive EDA Visualizations\n# # ========================\n\n# class EnhancedVisualizer:\n#     \"\"\"Advanced visualization class with multiple analysis types\"\"\"\n    \n#     def __init__(self, cfg: Config):\n#         self.cfg = cfg\n#         self.color_palette = sns.color_palette(\"husl\", 20)\n#         self.behavior_colors = {}\n        \n#     def save_figure(self, fig, name: str):\n#         \"\"\"Save figure if configured\"\"\"\n#         if self.cfg.save_figures:\n#             filepath = self.cfg.output_dir / f\"{name}.{self.cfg.figure_format}\"\n#             fig.savefig(filepath, dpi=self.cfg.figure_dpi, bbox_inches='tight')\n#             print(f\"   📸 Saved: {filepath}\")\n    \n#     def comprehensive_data_overview(self, train_df, solution_df):\n#         \"\"\"Create comprehensive multi-panel data overview\"\"\"\n#         print(\"\\n\" + \"=\"*80)\n#         print(\"📊 COMPREHENSIVE DATA OVERVIEW\")\n#         print(\"=\"*80)\n        \n#         # Calculate additional metrics\n#         solution_df = solution_df.with_columns(\n#             (pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\")\n#         )\n        \n#         # Create figure with custom layout\n#         fig = plt.figure(figsize=(24, 16))\n#         gs = GridSpec(4, 4, figure=fig, hspace=0.3, wspace=0.3)\n        \n#         # 1. Lab distribution with statistics\n#         ax1 = fig.add_subplot(gs[0, 0:2])\n#         lab_counts = train_df.group_by(\"lab_id\").count().sort(\"count\", descending=True)\n#         labs = lab_counts[\"lab_id\"].to_list()\n#         counts = lab_counts[\"count\"].to_list()\n        \n#         colors = [self.color_palette[i % len(self.color_palette)] for i in range(len(labs))]\n#         bars = ax1.bar(range(len(labs)), counts, color=colors, edgecolor='black', linewidth=0.5)\n        \n#         # Add value labels on bars\n#         for i, (bar, count) in enumerate(zip(bars, counts)):\n#             height = bar.get_height()\n#             ax1.text(bar.get_x() + bar.get_width()/2., height,\n#                     f'{count}', ha='center', va='bottom', fontsize=8)\n        \n#         ax1.set_xticks(range(len(labs)))\n#         ax1.set_xticklabels(labs, rotation=45, ha='right', fontsize=9)\n#         ax1.set_title(\"Videos per Laboratory\", fontsize=12, fontweight='bold')\n#         ax1.set_ylabel(\"Number of Videos\")\n#         ax1.grid(True, alpha=0.3, axis='y')\n        \n#         # Add mean line\n#         mean_count = np.mean(counts)\n#         ax1.axhline(mean_count, color='red', linestyle='--', alpha=0.5, \n#                    label=f'Mean: {mean_count:.1f}')\n#         ax1.legend()\n        \n#         # 2. Behavior frequency distribution (log scale option)\n#         ax2 = fig.add_subplot(gs[0, 2:4])\n#         action_counts = solution_df.group_by(\"action\").count().sort(\"count\", descending=True)\n#         actions = action_counts[\"action\"].to_list()\n#         action_cnts = action_counts[\"count\"].to_list()\n        \n#         # Store behavior colors for consistency\n#         self.behavior_colors = {action: self.color_palette[i % len(self.color_palette)] \n#                                for i, action in enumerate(actions)}\n        \n#         bars = ax2.bar(range(len(actions)), action_cnts, \n#                       color=[self.behavior_colors[a] for a in actions],\n#                       edgecolor='black', linewidth=0.5)\n        \n#         ax2.set_xticks(range(len(actions)))\n#         ax2.set_xticklabels(actions, rotation=45, ha='right', fontsize=8)\n#         ax2.set_title(\"Behavior Frequency Distribution\", fontsize=12, fontweight='bold')\n#         ax2.set_ylabel(\"Count (log scale)\")\n#         ax2.set_yscale('log')\n#         ax2.grid(True, alpha=0.3, axis='y')\n        \n#         # 3. Duration distribution by behavior\n#         ax3 = fig.add_subplot(gs[1, 0:2])\n        \n#         # Create violin plot for top behaviors\n#         top_behaviors = actions[:10]  # Top 10 behaviors\n#         duration_data = []\n#         behavior_labels = []\n        \n#         for behavior in top_behaviors:\n#             behavior_data = solution_df.filter(pl.col(\"action\") == behavior)[\"duration\"].to_list()\n#             duration_data.extend(behavior_data)\n#             behavior_labels.extend([behavior] * len(behavior_data))\n        \n#         duration_df = pd.DataFrame({'Behavior': behavior_labels, 'Duration': duration_data})\n        \n#         sns.violinplot(data=duration_df, x='Behavior', y='Duration', ax=ax3, \n#                       palette=[self.behavior_colors[b] for b in top_behaviors])\n#         ax3.set_xticklabels(top_behaviors, rotation=45, ha='right', fontsize=8)\n#         ax3.set_title(\"Duration Distribution by Behavior (Top 10)\", fontsize=12, fontweight='bold')\n#         ax3.set_ylabel(\"Duration (frames)\")\n#         ax3.grid(True, alpha=0.3, axis='y')\n        \n#         # 4. Mouse interaction heatmap\n#         ax4 = fig.add_subplot(gs[1, 2:4])\n        \n#         # Create interaction matrix\n#         pair_counts = solution_df.group_by([\"agent_id\", \"target_id\"]).count()\n        \n#         # Get unique mouse IDs\n#         all_mice = sorted(list(set(pair_counts[\"agent_id\"].unique().to_list() + \n#                                   pair_counts[\"target_id\"].unique().to_list())))[:10]\n        \n#         # Create matrix\n#         interaction_matrix = np.zeros((len(all_mice), len(all_mice)))\n#         for row in pair_counts.to_dicts():\n#             if row[\"agent_id\"] in all_mice and row[\"target_id\"] in all_mice:\n#                 i = all_mice.index(row[\"agent_id\"])\n#                 j = all_mice.index(row[\"target_id\"])\n#                 interaction_matrix[i, j] = row[\"count\"]\n        \n#         im = ax4.imshow(interaction_matrix, cmap='YlOrRd', aspect='auto')\n#         ax4.set_xticks(range(len(all_mice)))\n#         ax4.set_yticks(range(len(all_mice)))\n#         ax4.set_xticklabels(all_mice, rotation=45, ha='right', fontsize=8)\n#         ax4.set_yticklabels(all_mice, fontsize=8)\n#         ax4.set_title(\"Mouse Interaction Matrix (Top 10)\", fontsize=12, fontweight='bold')\n#         ax4.set_xlabel(\"Target Mouse\")\n#         ax4.set_ylabel(\"Agent Mouse\")\n#         plt.colorbar(im, ax=ax4, label='Interaction Count')\n        \n#         # 5. Temporal distribution of behaviors\n#         ax5 = fig.add_subplot(gs[2, 0:2])\n        \n#         # Create temporal bins\n#         frame_bins = np.linspace(0, 10000, 51)  # 50 bins\n#         temporal_counts = []\n        \n#         for i in range(len(frame_bins) - 1):\n#             count = len(solution_df.filter(\n#                 (pl.col(\"start_frame\") >= frame_bins[i]) & \n#                 (pl.col(\"start_frame\") < frame_bins[i+1])\n#             ))\n#             temporal_counts.append(count)\n        \n#         ax5.plot(frame_bins[:-1], temporal_counts, linewidth=2, color='darkblue')\n#         ax5.fill_between(frame_bins[:-1], temporal_counts, alpha=0.3, color='skyblue')\n#         ax5.set_title(\"Temporal Distribution of Behavior Onset\", fontsize=12, fontweight='bold')\n#         ax5.set_xlabel(\"Frame Number\")\n#         ax5.set_ylabel(\"Number of Behaviors\")\n#         ax5.grid(True, alpha=0.3)\n        \n#         # 6. Lab-specific behavior preferences\n#         ax6 = fig.add_subplot(gs[2, 2:4])\n        \n#         # Calculate lab-behavior matrix\n#         lab_behavior = solution_df.group_by([\"lab_id\", \"action\"]).count()\n#         top_labs = labs[:8]  # Top 8 labs\n#         top_actions = actions[:8]  # Top 8 behaviors\n        \n#         lab_behavior_matrix = np.zeros((len(top_labs), len(top_actions)))\n#         for row in lab_behavior.to_dicts():\n#             if row[\"lab_id\"] in top_labs and row[\"action\"] in top_actions:\n#                 i = top_labs.index(row[\"lab_id\"])\n#                 j = top_actions.index(row[\"action\"])\n#                 lab_behavior_matrix[i, j] = row[\"count\"]\n        \n#         # Normalize by row (lab)\n#         row_sums = lab_behavior_matrix.sum(axis=1, keepdims=True)\n#         lab_behavior_matrix_norm = np.divide(lab_behavior_matrix, row_sums, \n#                                             where=row_sums != 0)\n        \n#         im = ax6.imshow(lab_behavior_matrix_norm, cmap='viridis', aspect='auto', vmin=0, vmax=0.5)\n#         ax6.set_xticks(range(len(top_actions)))\n#         ax6.set_yticks(range(len(top_labs)))\n#         ax6.set_xticklabels(top_actions, rotation=45, ha='right', fontsize=8)\n#         ax6.set_yticklabels(top_labs, fontsize=8)\n#         ax6.set_title(\"Lab-specific Behavior Preferences (Normalized)\", fontsize=12, fontweight='bold')\n#         ax6.set_xlabel(\"Behavior\")\n#         ax6.set_ylabel(\"Laboratory\")\n#         plt.colorbar(im, ax=ax6, label='Proportion')\n        \n#         # 7. Duration vs Frequency scatter\n#         ax7 = fig.add_subplot(gs[3, 0:2])\n        \n#         behavior_stats = []\n#         for action in actions:\n#             action_data = solution_df.filter(pl.col(\"action\") == action)\n#             if len(action_data) > 0:\n#                 behavior_stats.append({\n#                     'action': action,\n#                     'frequency': len(action_data),\n#                     'mean_duration': action_data[\"duration\"].mean(),\n#                     'std_duration': action_data[\"duration\"].std()\n#                 })\n        \n#         if behavior_stats:\n#             freq = [b['frequency'] for b in behavior_stats]\n#             dur = [b['mean_duration'] for b in behavior_stats]\n#             std = [b['std_duration'] for b in behavior_stats]\n            \n#             scatter = ax7.scatter(freq, dur, s=100, alpha=0.6, \n#                                  c=[self.behavior_colors[b['action']] for b in behavior_stats],\n#                                  edgecolors='black', linewidth=1)\n            \n#             # Add error bars for standard deviation\n#             ax7.errorbar(freq, dur, yerr=std, fmt='none', alpha=0.3, ecolor='gray')\n            \n#             # Annotate interesting points\n#             for i, b in enumerate(behavior_stats[:5]):  # Top 5 behaviors\n#                 ax7.annotate(b['action'], (freq[i], dur[i]), \n#                            fontsize=8, ha='left', va='bottom',\n#                            xytext=(5, 5), textcoords='offset points')\n            \n#             ax7.set_xscale('log')\n#             ax7.set_title(\"Behavior Frequency vs Duration\", fontsize=12, fontweight='bold')\n#             ax7.set_xlabel(\"Frequency (log scale)\")\n#             ax7.set_ylabel(\"Mean Duration (frames)\")\n#             ax7.grid(True, alpha=0.3)\n        \n#         # 8. Summary statistics panel\n#         ax8 = fig.add_subplot(gs[3, 2:4])\n#         ax8.axis('off')\n        \n#         # Calculate comprehensive statistics\n#         total_videos = train_df[\"video_id\"].n_unique()\n#         total_labs = train_df[\"lab_id\"].n_unique()\n#         total_annotations = len(solution_df)\n#         unique_behaviors = solution_df[\"action\"].n_unique()\n#         unique_pairs = solution_df.select([\"agent_id\", \"target_id\"]).n_unique()\n        \n#         mean_duration = solution_df[\"duration\"].mean()\n#         median_duration = solution_df[\"duration\"].median()\n#         std_duration = solution_df[\"duration\"].std()\n        \n#         behaviors_per_video = solution_df.group_by(\"video_id\").count()[\"count\"].mean()\n        \n#         stats_text = f\"\"\"\n#         COMPREHENSIVE DATASET STATISTICS\n#         =====================================\n        \n#         DATA OVERVIEW\n#         -------------\n#         Total Videos:          {total_videos:,}\n#         Total Laboratories:    {total_labs}\n#         Total Annotations:     {total_annotations:,}\n#         Unique Behaviors:      {unique_behaviors}\n#         Unique Mouse Pairs:    {unique_pairs}\n        \n#         DURATION STATISTICS\n#         ------------------\n#         Mean Duration:         {mean_duration:.1f} frames (~{mean_duration/30:.1f}s @ 30fps)\n#         Median Duration:       {median_duration:.1f} frames\n#         Std Duration:          {std_duration:.1f} frames\n        \n#         BEHAVIOR DISTRIBUTION\n#         --------------------\n#         Top 3 Behaviors:\n#         1. {actions[0]:15s} {action_cnts[0]:,} ({action_cnts[0]/sum(action_cnts)*100:.1f}%)\n#         2. {actions[1]:15s} {action_cnts[1]:,} ({action_cnts[1]/sum(action_cnts)*100:.1f}%)\n#         3. {actions[2]:15s} {action_cnts[2]:,} ({action_cnts[2]/sum(action_cnts)*100:.1f}%)\n        \n#         Rare Behaviors (< 0.1%): {sum(1 for c in action_cnts if c < total_annotations * 0.001)}\n        \n#         COVERAGE\n#         --------\n#         Avg Behaviors/Video:   {behaviors_per_video:.1f}\n#         Data Completeness:     {len(solution_df)/(total_videos*100)*100:.1f}%\n#         \"\"\"\n        \n#         ax8.text(0.05, 0.95, stats_text, transform=ax8.transAxes, fontsize=9,\n#                 verticalalignment='top', fontfamily='monospace',\n#                 bbox=dict(boxstyle='round', facecolor='lightgray', alpha=0.2))\n        \n#         plt.suptitle(\"MABe Dataset Comprehensive Analysis Dashboard\", \n#                     fontsize=16, fontweight='bold', y=0.98)\n        \n#         self.save_figure(fig, \"comprehensive_overview\")\n#         plt.show()\n        \n#         # Print key insights\n#         self._print_overview_insights(train_df, solution_df, actions, action_cnts)\n    \n#     def _print_overview_insights(self, train_df, solution_df, actions, action_cnts):\n#         \"\"\"Print detailed insights from overview analysis\"\"\"\n#         print(\"\\n📈 KEY INSIGHTS FROM DATA OVERVIEW:\")\n#         print(\"=\"*60)\n        \n#         print(\"\\n1. SCALE AND SCOPE:\")\n#         print(f\"   • Dataset spans {train_df['video_id'].n_unique()} videos across {train_df['lab_id'].n_unique()} laboratories\")\n#         print(f\"   • {len(solution_df):,} total behavior annotations recorded\")\n#         print(f\"   • {solution_df['action'].n_unique()} distinct behavior types identified\")\n        \n#         print(\"\\n2. BEHAVIOR PATTERNS:\")\n#         print(f\"   • Dominant behavior: '{actions[0]}' accounts for {action_cnts[0]/sum(action_cnts)*100:.1f}% of observations\")\n#         print(f\"   • Top 5 behaviors represent {sum(action_cnts[:5])/sum(action_cnts)*100:.1f}% of all annotations\")\n#         print(f\"   • Long-tail distribution: {sum(1 for c in action_cnts if c < 100)} behaviors occur < 100 times\")\n        \n#         print(\"\\n3. TEMPORAL CHARACTERISTICS:\")\n#         mean_dur = solution_df[\"duration\"].mean()\n#         print(f\"   • Average behavior duration: {mean_dur:.1f} frames (~{mean_dur/30:.1f} seconds)\")\n#         print(f\"   • Duration range: {solution_df['duration'].min()}-{solution_df['duration'].max()} frames\")\n#         print(f\"   • High variability: CV = {solution_df['duration'].std()/mean_dur:.2f}\")\n        \n#         print(\"\\n4. INTER-LABORATORY VARIATION:\")\n#         lab_counts = train_df.group_by(\"lab_id\").count()\n#         print(f\"   • Lab contribution range: {lab_counts['count'].min()}-{lab_counts['count'].max()} videos\")\n#         print(f\"   • Coefficient of variation: {lab_counts['count'].std()/lab_counts['count'].mean():.2f}\")\n        \n#     def behavior_transition_analysis(self, solution_df):\n#         \"\"\"Analyze behavior transitions and sequences\"\"\"\n#         print(\"\\n\" + \"=\"*80)\n#         print(\"🔄 BEHAVIOR TRANSITION ANALYSIS\")\n#         print(\"=\"*80)\n        \n#         fig = plt.figure(figsize=(20, 12))\n#         gs = GridSpec(2, 3, figure=fig, hspace=0.3, wspace=0.3)\n        \n#         # Calculate transitions\n#         transitions = defaultdict(lambda: defaultdict(int))\n        \n#         # Group by video and analyze sequences\n#         for video_id in tqdm(solution_df[\"video_id\"].unique()[:100], desc=\"Analyzing transitions\"):\n#             video_data = solution_df.filter(pl.col(\"video_id\") == video_id).sort(\"start_frame\")\n            \n#             behaviors = video_data[\"action\"].to_list()\n#             for i in range(len(behaviors) - 1):\n#                 transitions[behaviors[i]][behaviors[i+1]] += 1\n        \n#         # 1. Transition matrix heatmap\n#         ax1 = fig.add_subplot(gs[0, :2])\n        \n#         # Get most common behaviors for readability\n#         behavior_counts = solution_df.group_by(\"action\").count().sort(\"count\", descending=True)\n#         top_behaviors = behavior_counts[\"action\"].to_list()[:15]\n        \n#         # Create transition matrix\n#         matrix = np.zeros((len(top_behaviors), len(top_behaviors)))\n#         for i, b1 in enumerate(top_behaviors):\n#             for j, b2 in enumerate(top_behaviors):\n#                 matrix[i, j] = transitions[b1][b2]\n        \n#         # Normalize by row\n#         row_sums = matrix.sum(axis=1, keepdims=True)\n#         matrix_norm = np.divide(matrix, row_sums, where=row_sums != 0)\n        \n#         im = ax1.imshow(matrix_norm, cmap='RdYlBu_r', aspect='auto', vmin=0, vmax=0.3)\n#         ax1.set_xticks(range(len(top_behaviors)))\n#         ax1.set_yticks(range(len(top_behaviors)))\n#         ax1.set_xticklabels(top_behaviors, rotation=45, ha='right', fontsize=8)\n#         ax1.set_yticklabels(top_behaviors, fontsize=8)\n#         ax1.set_title(\"Behavior Transition Probability Matrix\", fontsize=12, fontweight='bold')\n#         ax1.set_xlabel(\"Next Behavior\")\n#         ax1.set_ylabel(\"Current Behavior\")\n#         plt.colorbar(im, ax=ax1, label='Transition Probability')\n        \n#         # 2. Network graph of transitions\n#         ax2 = fig.add_subplot(gs[0, 2])\n        \n#         # Create directed graph\n#         G = nx.DiGraph()\n        \n#         # Add edges with weights\n#         threshold = 50  # Minimum transitions to show\n#         for b1 in top_behaviors[:10]:\n#             for b2 in top_behaviors[:10]:\n#                 weight = transitions[b1][b2]\n#                 if weight > threshold and b1 != b2:\n#                     G.add_edge(b1, b2, weight=weight)\n        \n#         if G.edges():\n#             pos = nx.spring_layout(G, k=2, iterations=50)\n            \n#             # Draw nodes\n#             node_sizes = [behavior_counts.filter(pl.col(\"action\") == node)[\"count\"].to_list()[0] / 50 \n#                          for node in G.nodes()]\n#             nx.draw_networkx_nodes(G, pos, node_size=node_sizes, \n#                                   node_color='lightblue', alpha=0.7, ax=ax2)\n            \n#             # Draw edges with varying thickness\n#             edges = G.edges()\n#             weights = [G[u][v]['weight'] / 100 for u, v in edges]\n#             nx.draw_networkx_edges(G, pos, width=weights, alpha=0.5, \n#                                   edge_color='gray', arrows=True, ax=ax2)\n            \n#             # Draw labels\n#             nx.draw_networkx_labels(G, pos, font_size=8, ax=ax2)\n            \n#             ax2.set_title(\"Behavior Transition Network\", fontsize=12, fontweight='bold')\n#             ax2.axis('off')\n        \n#         # 3. Sequence length distribution\n#         ax3 = fig.add_subplot(gs[1, 0])\n        \n#         sequence_lengths = []\n#         for video_id in solution_df[\"video_id\"].unique()[:100]:\n#             video_data = solution_df.filter(pl.col(\"video_id\") == video_id)\n#             sequence_lengths.append(len(video_data))\n        \n#         ax3.hist(sequence_lengths, bins=30, color='skyblue', edgecolor='black', alpha=0.7)\n#         ax3.axvline(np.mean(sequence_lengths), color='red', linestyle='--', \n#                    label=f'Mean: {np.mean(sequence_lengths):.1f}')\n#         ax3.set_title(\"Behavior Sequence Length Distribution\", fontsize=12, fontweight='bold')\n#         ax3.set_xlabel(\"Sequence Length (behaviors per video)\")\n#         ax3.set_ylabel(\"Frequency\")\n#         ax3.legend()\n#         ax3.grid(True, alpha=0.3)\n        \n#         # 4. Most common sequences\n#         ax4 = fig.add_subplot(gs[1, 1])\n        \n#         # Find common 3-grams\n#         trigrams = defaultdict(int)\n#         for video_id in solution_df[\"video_id\"].unique()[:100]:\n#             video_data = solution_df.filter(pl.col(\"video_id\") == video_id).sort(\"start_frame\")\n#             behaviors = video_data[\"action\"].to_list()\n            \n#             for i in range(len(behaviors) - 2):\n#                 trigram = f\"{behaviors[i]} → {behaviors[i+1]} → {behaviors[i+2]}\"\n#                 trigrams[trigram] += 1\n        \n#         # Get top trigrams\n#         top_trigrams = sorted(trigrams.items(), key=lambda x: x[1], reverse=True)[:10]\n        \n#         if top_trigrams:\n#             trigram_names = [t[0] for t in top_trigrams]\n#             trigram_counts = [t[1] for t in top_trigrams]\n            \n#             bars = ax4.barh(range(len(trigram_names)), trigram_counts, color='coral')\n#             ax4.set_yticks(range(len(trigram_names)))\n#             ax4.set_yticklabels(trigram_names, fontsize=8)\n#             ax4.set_title(\"Most Common Behavior Sequences (3-grams)\", fontsize=12, fontweight='bold')\n#             ax4.set_xlabel(\"Frequency\")\n#             ax4.grid(True, alpha=0.3, axis='x')\n        \n#         # 5. Transition statistics\n#         ax5 = fig.add_subplot(gs[1, 2])\n#         ax5.axis('off')\n        \n#         # Calculate statistics\n#         total_transitions = sum(sum(d.values()) for d in transitions.values())\n#         unique_transitions = sum(len(d) for d in transitions.values())\n        \n#         # Self-transitions\n#         self_transitions = sum(transitions[b][b] for b in transitions)\n#         self_ratio = self_transitions / total_transitions if total_transitions > 0 else 0\n        \n#         # Most likely transitions\n#         all_transitions = []\n#         for b1, targets in transitions.items():\n#             for b2, count in targets.items():\n#                 if b1 != b2:\n#                     all_transitions.append((f\"{b1} → {b2}\", count))\n        \n#         top_transitions = sorted(all_transitions, key=lambda x: x[1], reverse=True)[:5]\n        \n#         stats_text = f\"\"\"\n#         TRANSITION STATISTICS\n#         =====================\n        \n#         Total Transitions:     {total_transitions:,}\n#         Unique Transitions:    {unique_transitions}\n#         Self-transitions:      {self_ratio*100:.1f}%\n        \n#         TOP TRANSITIONS\n#         ---------------\"\"\"\n        \n#         for i, (trans, count) in enumerate(top_transitions, 1):\n#             stats_text += f\"\\n        {i}. {trans:30s} {count:,}\"\n        \n#         ax5.text(0.05, 0.95, stats_text, transform=ax5.transAxes, fontsize=10,\n#                 verticalalignment='top', fontfamily='monospace',\n#                 bbox=dict(boxstyle='round', facecolor='lightgray', alpha=0.2))\n        \n#         plt.suptitle(\"Behavior Transition and Sequence Analysis\", \n#                     fontsize=14, fontweight='bold')\n        \n#         self.save_figure(fig, \"behavior_transitions\")\n#         plt.show()\n        \n#         print(\"\\n🔄 TRANSITION INSIGHTS:\")\n#         print(f\"   • Total transitions analyzed: {total_transitions:,}\")\n#         print(f\"   • Self-transition rate: {self_ratio*100:.1f}%\")\n#         print(f\"   • Most common transition: {top_transitions[0][0] if top_transitions else 'N/A'}\")\n    \n#     def spatial_analysis(self, train_df, solution_df, cfg: Config):\n#         \"\"\"Analyze spatial patterns in behavior\"\"\"\n#         print(\"\\n\" + \"=\"*80)\n#         print(\"📍 SPATIAL BEHAVIOR ANALYSIS\")\n#         print(\"=\"*80)\n        \n#         fig = plt.figure(figsize=(20, 12))\n#         gs = GridSpec(3, 3, figure=fig, hspace=0.3, wspace=0.3)\n        \n#         # Sample some tracking data for spatial analysis\n#         spatial_data = []\n#         sampled_videos = 0\n#         max_samples = 20  # Limit for performance\n        \n#         for row in train_df.head(100).to_dicts():\n#             if sampled_videos >= max_samples:\n#                 break\n                \n#             lab_id = row[\"lab_id\"]\n#             if lab_id.startswith(\"MABe22\"):\n#                 continue\n                \n#             video_id = row[\"video_id\"]\n#             track_path = cfg.train_track_dir / lab_id / f\"{video_id}.parquet\"\n            \n#             if not track_path.exists():\n#                 continue\n                \n#             try:\n#                 track_data = pd.read_parquet(track_path)\n                \n#                 # Get behavior annotations for this video\n#                 video_behaviors = solution_df.filter(pl.col(\"video_id\") == video_id)\n                \n#                 if len(video_behaviors) > 0 and 'x' in track_data.columns:\n#                     spatial_data.append({\n#                         'video_id': video_id,\n#                         'track_data': track_data,\n#                         'behaviors': video_behaviors\n#                     })\n#                     sampled_videos += 1\n                    \n#             except Exception as e:\n#                 logger.warning(f\"Failed to load tracking data: {e}\")\n#                 continue\n        \n#         if not spatial_data:\n#             print(\"   ⚠️ No spatial data available for analysis\")\n#             return\n        \n#         # 1. Trajectory visualization with behaviors\n#         ax1 = fig.add_subplot(gs[0, :2])\n        \n#         # Plot first video's trajectory\n#         first_video = spatial_data[0]\n#         track = first_video['track_data']\n        \n#         if 'mouse_id' in track.columns:\n#             mice = track['mouse_id'].unique()[:2]  # Plot first 2 mice\n            \n#             for i, mouse_id in enumerate(mice):\n#                 mouse_data = track[track['mouse_id'] == mouse_id]\n#                 if 'x' in mouse_data.columns and 'y' in mouse_data.columns:\n#                     x = mouse_data['x'].values[:1000]  # Limit points for visibility\n#                     y = mouse_data['y'].values[:1000]\n                    \n#                     # Create gradient color based on time\n#                     colors = plt.cm.viridis(np.linspace(0, 1, len(x)))\n                    \n#                     for j in range(len(x)-1):\n#                         ax1.plot(x[j:j+2], y[j:j+2], color=colors[j], alpha=0.7, linewidth=1)\n                    \n#                     ax1.scatter(x[0], y[0], color='green', s=100, marker='o', \n#                               label=f'Mouse {mouse_id} Start', zorder=5)\n#                     ax1.scatter(x[-1], y[-1], color='red', s=100, marker='s', \n#                               label=f'Mouse {mouse_id} End', zorder=5)\n        \n#         ax1.set_title(f\"Sample Trajectory Visualization (Video {first_video['video_id']})\", \n#                      fontsize=12, fontweight='bold')\n#         ax1.set_xlabel(\"X Position (pixels)\")\n#         ax1.set_ylabel(\"Y Position (pixels)\")\n#         ax1.legend()\n#         ax1.grid(True, alpha=0.3)\n#         ax1.set_aspect('equal')\n        \n#         # 2. Spatial heatmap\n#         ax2 = fig.add_subplot(gs[0, 2])\n        \n#         # Aggregate all positions\n#         all_x = []\n#         all_y = []\n        \n#         for video_data in spatial_data[:5]:  # Use first 5 videos\n#             track = video_data['track_data']\n#             if 'x' in track.columns and 'y' in track.columns:\n#                 all_x.extend(track['x'].dropna().values)\n#                 all_y.extend(track['y'].dropna().values)\n        \n#         if all_x and all_y:\n#             heatmap, xedges, yedges = np.histogram2d(all_x, all_y, bins=30)\n#             im = ax2.imshow(heatmap.T, origin='lower', cmap='hot', aspect='auto')\n#             ax2.set_title(\"Spatial Occupancy Heatmap\", fontsize=12, fontweight='bold')\n#             ax2.set_xlabel(\"X Position (binned)\")\n#             ax2.set_ylabel(\"Y Position (binned)\")\n#             plt.colorbar(im, ax=ax2, label='Occupancy')\n        \n#         # 3. Movement speed distribution\n#         ax3 = fig.add_subplot(gs[1, 0])\n        \n#         speeds = []\n#         for video_data in spatial_data:\n#             track = video_data['track_data']\n#             if 'x' in track.columns and 'y' in track.columns:\n#                 # Calculate speed for each mouse\n#                 if 'mouse_id' in track.columns:\n#                     for mouse_id in track['mouse_id'].unique():\n#                         mouse_data = track[track['mouse_id'] == mouse_id]\n#                         x = mouse_data['x'].values\n#                         y = mouse_data['y'].values\n                        \n#                         # Calculate instantaneous speed\n#                         dx = np.diff(x)\n#                         dy = np.diff(y)\n#                         speed = np.sqrt(dx**2 + dy**2)\n#                         speeds.extend(speed[~np.isnan(speed)])\n        \n#         if len(speeds) > 0:\n#             speeds = np.array(speeds)\n#             speeds = speeds[speeds < np.percentile(speeds, 99)]  # Remove outliers\n            \n#             ax3.hist(speeds, bins=50, color='lightblue', edgecolor='black', alpha=0.7)\n#             ax3.axvline(np.mean(speeds), color='red', linestyle='--', \n#                        label=f'Mean: {np.mean(speeds):.2f}')\n#             ax3.axvline(np.median(speeds), color='green', linestyle='--', \n#                        label=f'Median: {np.median(speeds):.2f}')\n#             ax3.set_title(\"Movement Speed Distribution\", fontsize=12, fontweight='bold')\n#             ax3.set_xlabel(\"Speed (pixels/frame)\")\n#             ax3.set_ylabel(\"Frequency\")\n#             ax3.legend()\n#             ax3.grid(True, alpha=0.3)\n        \n#         # 4. Distance traveled by behavior\n#         ax4 = fig.add_subplot(gs[1, 1])\n        \n#         behavior_distances = defaultdict(list)\n        \n#         for video_data in spatial_data[:10]:\n#             track = video_data['track_data']\n#             behaviors = video_data['behaviors']\n            \n#             if 'video_frame' in track.columns and 'x' in track.columns:\n#                 for behavior in behaviors.to_dicts():\n#                     start_frame = behavior['start_frame']\n#                     stop_frame = behavior['stop_frame']\n#                     action = behavior['action']\n                    \n#                     # Get track data for this behavior period\n#                     behavior_track = track[\n#                         (track['video_frame'] >= start_frame) & \n#                         (track['video_frame'] <= stop_frame)\n#                     ]\n                    \n#                     if len(behavior_track) > 1:\n#                         # Calculate total distance\n#                         x = behavior_track['x'].values\n#                         y = behavior_track['y'].values\n#                         dx = np.diff(x)\n#                         dy = np.diff(y)\n#                         distances = np.sqrt(dx**2 + dy**2)\n#                         total_distance = np.nansum(distances)\n#                         behavior_distances[action].append(total_distance)\n        \n#         # Plot top behaviors\n#         if behavior_distances:\n#             top_behaviors = sorted(behavior_distances.items(), \n#                                  key=lambda x: len(x[1]), reverse=True)[:8]\n            \n#             behavior_names = [b[0] for b in top_behaviors]\n#             distance_data = [b[1] for b in top_behaviors]\n            \n#             bp = ax4.boxplot(distance_data, labels=behavior_names, patch_artist=True)\n#             for patch, color in zip(bp['boxes'], self.color_palette):\n#                 patch.set_facecolor(color)\n#                 patch.set_alpha(0.7)\n            \n#             ax4.set_xticklabels(behavior_names, rotation=45, ha='right', fontsize=8)\n#             ax4.set_title(\"Distance Traveled by Behavior Type\", fontsize=12, fontweight='bold')\n#             ax4.set_ylabel(\"Total Distance (pixels)\")\n#             ax4.grid(True, alpha=0.3, axis='y')\n        \n#         # 5. Inter-mouse distance analysis\n#         ax5 = fig.add_subplot(gs[1, 2])\n        \n#         inter_distances = []\n        \n#         for video_data in spatial_data[:5]:\n#             track = video_data['track_data']\n            \n#             if 'mouse_id' in track.columns and 'x' in track.columns:\n#                 mice = track['mouse_id'].unique()\n                \n#                 if len(mice) >= 2:\n#                     mouse1_data = track[track['mouse_id'] == mice[0]]\n#                     mouse2_data = track[track['mouse_id'] == mice[1]]\n                    \n#                     # Align frames\n#                     common_frames = set(mouse1_data['video_frame']) & set(mouse2_data['video_frame'])\n                    \n#                     for frame in list(common_frames)[:500]:  # Sample frames\n#                         m1 = mouse1_data[mouse1_data['video_frame'] == frame]\n#                         m2 = mouse2_data[mouse2_data['video_frame'] == frame]\n                        \n#                         if len(m1) > 0 and len(m2) > 0:\n#                             x1, y1 = m1['x'].iloc[0], m1['y'].iloc[0]\n#                             x2, y2 = m2['x'].iloc[0], m2['y'].iloc[0]\n#                             dist = np.sqrt((x2-x1)**2 + (y2-y1)**2)\n#                             if not np.isnan(dist):\n#                                 inter_distances.append(dist)\n        \n#         if len(inter_distances) > 0:\n#             ax5.hist(inter_distances, bins=50, color='purple', edgecolor='black', alpha=0.7)\n#             ax5.axvline(np.mean(inter_distances), color='red', linestyle='--',\n#                        label=f'Mean: {np.mean(inter_distances):.1f}')\n#             ax5.set_title(\"Inter-mouse Distance Distribution\", fontsize=12, fontweight='bold')\n#             ax5.set_xlabel(\"Distance (pixels)\")\n#             ax5.set_ylabel(\"Frequency\")\n#             ax5.legend()\n#             ax5.grid(True, alpha=0.3)\n        \n#         # 6. Movement patterns by time\n#         ax6 = fig.add_subplot(gs[2, :])\n        \n#         # Analyze movement over time\n#         time_bins = 20\n#         time_speeds = [[] for _ in range(time_bins)]\n        \n#         for video_data in spatial_data[:10]:\n#             track = video_data['track_data']\n            \n#             if 'video_frame' in track.columns and 'x' in track.columns:\n#                 max_frame = track['video_frame'].max()\n                \n#                 if 'mouse_id' in track.columns:\n#                     for mouse_id in track['mouse_id'].unique():\n#                         mouse_data = track[track['mouse_id'] == mouse_id]\n                        \n#                         for i in range(time_bins):\n#                             frame_start = i * max_frame / time_bins\n#                             frame_end = (i + 1) * max_frame / time_bins\n                            \n#                             period_data = mouse_data[\n#                                 (mouse_data['video_frame'] >= frame_start) & \n#                                 (mouse_data['video_frame'] < frame_end)\n#                             ]\n                            \n#                             if len(period_data) > 1:\n#                                 x = period_data['x'].values\n#                                 y = period_data['y'].values\n#                                 dx = np.diff(x)\n#                                 dy = np.diff(y)\n#                                 speeds = np.sqrt(dx**2 + dy**2)\n#                                 time_speeds[i].extend(speeds[~np.isnan(speeds)])\n        \n#         # Plot average speed over time\n#         mean_speeds = [np.mean(speeds) if speeds else 0 for speeds in time_speeds]\n#         std_speeds = [np.std(speeds) if speeds else 0 for speeds in time_speeds]\n        \n#         time_points = np.arange(time_bins)\n#         ax6.plot(time_points, mean_speeds, 'b-', linewidth=2, label='Mean Speed')\n#         ax6.fill_between(time_points, \n#                          np.array(mean_speeds) - np.array(std_speeds),\n#                          np.array(mean_speeds) + np.array(std_speeds),\n#                          alpha=0.3, color='blue', label='±1 STD')\n        \n#         ax6.set_title(\"Movement Activity Over Time\", fontsize=12, fontweight='bold')\n#         ax6.set_xlabel(\"Time Bin (normalized)\")\n#         ax6.set_ylabel(\"Average Speed (pixels/frame)\")\n#         ax6.legend()\n#         ax6.grid(True, alpha=0.3)\n        \n#         plt.suptitle(\"Spatial Behavior Analysis\", fontsize=14, fontweight='bold')\n        \n#         self.save_figure(fig, \"spatial_analysis\")\n#         plt.show()\n        \n#         print(\"\\n📍 SPATIAL INSIGHTS:\")\n#         if len(speeds) > 0:\n#             speeds_array = np.array(speeds)\n#             print(f\"   • Mean movement speed: {np.mean(speeds_array):.2f} pixels/frame\")\n#             print(f\"   • Speed variability (CV): {np.std(speeds_array)/np.mean(speeds_array):.2f}\")\n#         if len(inter_distances) > 0:\n#             print(f\"   • Mean inter-mouse distance: {np.mean(inter_distances):.1f} pixels\")\n#         print(f\"   • Analyzed {len(spatial_data)} videos with tracking data\")\n    \n#     def lab_comparison_analysis(self, solution_df):\n#         \"\"\"Detailed comparison between laboratories\"\"\"\n#         print(\"\\n\" + \"=\"*80)\n#         print(\"🔬 LABORATORY COMPARISON ANALYSIS\")\n#         print(\"=\"*80)\n        \n#         fig = plt.figure(figsize=(20, 14))\n#         gs = GridSpec(3, 3, figure=fig, hspace=0.3, wspace=0.3)\n        \n#         # Prepare data\n#         solution_df = solution_df.with_columns(\n#             (pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\")\n#         )\n        \n#         # Get lab statistics\n#         lab_stats = []\n#         for lab_id in solution_df[\"lab_id\"].unique():\n#             lab_data = solution_df.filter(pl.col(\"lab_id\") == lab_id)\n            \n#             lab_stats.append({\n#                 'lab_id': lab_id,\n#                 'n_annotations': len(lab_data),\n#                 'n_behaviors': lab_data[\"action\"].n_unique(),\n#                 'mean_duration': lab_data[\"duration\"].mean(),\n#                 'std_duration': lab_data[\"duration\"].std(),\n#                 'n_videos': lab_data[\"video_id\"].n_unique(),\n#                 'annotations_per_video': len(lab_data) / lab_data[\"video_id\"].n_unique()\n#             })\n        \n#         lab_stats_df = pd.DataFrame(lab_stats)\n        \n#         # 1. Lab annotation volume\n#         ax1 = fig.add_subplot(gs[0, 0])\n        \n#         sorted_labs = lab_stats_df.sort_values('n_annotations', ascending=False)\n#         colors = [self.color_palette[i % len(self.color_palette)] for i in range(len(sorted_labs))]\n        \n#         bars = ax1.bar(range(len(sorted_labs)), sorted_labs['n_annotations'], color=colors)\n#         ax1.set_xticks(range(len(sorted_labs)))\n#         ax1.set_xticklabels(sorted_labs['lab_id'], rotation=45, ha='right', fontsize=8)\n#         ax1.set_title(\"Annotation Volume by Laboratory\", fontsize=12, fontweight='bold')\n#         ax1.set_ylabel(\"Number of Annotations\")\n#         ax1.grid(True, alpha=0.3, axis='y')\n        \n#         # 2. Behavior diversity by lab\n#         ax2 = fig.add_subplot(gs[0, 1])\n        \n#         ax2.scatter(lab_stats_df['n_annotations'], lab_stats_df['n_behaviors'],\n#                    s=100, alpha=0.7, c=colors, edgecolors='black')\n        \n#         # Add labels for outliers\n#         for idx, row in lab_stats_df.iterrows():\n#             if row['n_behaviors'] > lab_stats_df['n_behaviors'].quantile(0.9):\n#                 ax2.annotate(row['lab_id'], (row['n_annotations'], row['n_behaviors']),\n#                            fontsize=8, ha='left', va='bottom')\n        \n#         ax2.set_xlabel(\"Number of Annotations\")\n#         ax2.set_ylabel(\"Number of Unique Behaviors\")\n#         ax2.set_title(\"Behavior Diversity vs Annotation Volume\", fontsize=12, fontweight='bold')\n#         ax2.grid(True, alpha=0.3)\n        \n#         # 3. Duration characteristics by lab\n#         ax3 = fig.add_subplot(gs[0, 2])\n        \n#         # Create box plot data\n#         duration_by_lab = []\n#         lab_names = []\n        \n#         for lab_id in sorted_labs['lab_id'][:10]:  # Top 10 labs\n#             lab_data = solution_df.filter(pl.col(\"lab_id\") == lab_id)\n#             durations = lab_data[\"duration\"].to_list()\n#             if durations:\n#                 duration_by_lab.append(durations[:1000])  # Limit for performance\n#                 lab_names.append(lab_id)\n        \n#         if duration_by_lab:\n#             bp = ax3.boxplot(duration_by_lab, labels=lab_names, patch_artist=True)\n#             for patch, color in zip(bp['boxes'], self.color_palette):\n#                 patch.set_facecolor(color)\n#                 patch.set_alpha(0.7)\n            \n#             ax3.set_xticklabels(lab_names, rotation=45, ha='right', fontsize=8)\n#             ax3.set_title(\"Duration Distribution by Laboratory\", fontsize=12, fontweight='bold')\n#             ax3.set_ylabel(\"Duration (frames)\")\n#             ax3.grid(True, alpha=0.3, axis='y')\n        \n#         # 4. Lab clustering based on behavior profiles\n#         ax4 = fig.add_subplot(gs[1, :2])\n        \n#         # Create behavior profile matrix\n#         top_behaviors = solution_df.group_by(\"action\").count().sort(\"count\", descending=True)[\"action\"].to_list()[:20]\n#         behavior_profiles = []\n        \n#         for lab_id in sorted_labs['lab_id'][:15]:  # Top 15 labs\n#             lab_data = solution_df.filter(pl.col(\"lab_id\") == lab_id)\n#             profile = []\n            \n#             for behavior in top_behaviors:\n#                 count = len(lab_data.filter(pl.col(\"action\") == behavior))\n#                 profile.append(count)\n            \n#             # Normalize\n#             total = sum(profile)\n#             if total > 0:\n#                 profile = [x/total for x in profile]\n#             behavior_profiles.append(profile)\n        \n#         if behavior_profiles:\n#             # Perform hierarchical clustering\n#             linkage_matrix = linkage(behavior_profiles, method='ward')\n#             dendrogram(linkage_matrix, labels=sorted_labs['lab_id'][:15].tolist(),\n#                       ax=ax4, leaf_rotation=45, leaf_font_size=8)\n#             ax4.set_title(\"Laboratory Clustering by Behavior Profile\", fontsize=12, fontweight='bold')\n#             ax4.set_ylabel(\"Distance\")\n        \n#         # 5. Lab-specific behavior preferences\n#         ax5 = fig.add_subplot(gs[1, 2])\n        \n#         # Calculate relative preferences\n#         lab_preferences = {}\n#         global_behavior_freq = solution_df.group_by(\"action\").count()\n#         global_total = len(solution_df)\n        \n#         for lab_id in sorted_labs['lab_id'][:8]:\n#             lab_data = solution_df.filter(pl.col(\"lab_id\") == lab_id)\n#             lab_total = len(lab_data)\n            \n#             preferences = []\n#             for behavior in top_behaviors[:5]:\n#                 lab_count = len(lab_data.filter(pl.col(\"action\") == behavior))\n#                 global_count = global_behavior_freq.filter(pl.col(\"action\") == behavior)[\"count\"].to_list()\n                \n#                 if global_count and global_total > 0 and lab_total > 0:\n#                     expected = (global_count[0] / global_total) * lab_total\n#                     if expected > 0:\n#                         ratio = lab_count / expected\n#                         preferences.append(ratio)\n#                     else:\n#                         preferences.append(1.0)\n#                 else:\n#                     preferences.append(1.0)\n            \n#             lab_preferences[lab_id] = preferences\n        \n#         # Plot as heatmap\n#         if lab_preferences:\n#             preference_matrix = np.array(list(lab_preferences.values()))\n#             im = ax5.imshow(preference_matrix, cmap='RdBu_r', aspect='auto', vmin=0, vmax=2)\n            \n#             ax5.set_xticks(range(5))\n#             ax5.set_xticklabels(top_behaviors[:5], rotation=45, ha='right', fontsize=8)\n#             ax5.set_yticks(range(len(lab_preferences)))\n#             ax5.set_yticklabels(list(lab_preferences.keys()), fontsize=8)\n#             ax5.set_title(\"Lab Behavior Preferences (vs Global)\", fontsize=12, fontweight='bold')\n#             plt.colorbar(im, ax=ax5, label='Relative Frequency')\n        \n#         # 6. Annotation quality metrics\n#         ax6 = fig.add_subplot(gs[2, :])\n        \n#         # Calculate quality metrics\n#         quality_metrics = []\n        \n#         for lab_id in solution_df[\"lab_id\"].unique():\n#             lab_data = solution_df.filter(pl.col(\"lab_id\") == lab_id)\n            \n#             # Calculate metrics\n#             duration_cv = lab_data[\"duration\"].std() / lab_data[\"duration\"].mean() if lab_data[\"duration\"].mean() > 0 else 0\n#             behaviors_per_video = lab_data[\"action\"].n_unique() / lab_data[\"video_id\"].n_unique() if lab_data[\"video_id\"].n_unique() > 0 else 0\n#             annotation_density = len(lab_data) / lab_data[\"video_id\"].n_unique() if lab_data[\"video_id\"].n_unique() > 0 else 0\n            \n#             quality_metrics.append({\n#                 'lab_id': lab_id,\n#                 'duration_cv': duration_cv,\n#                 'behaviors_per_video': behaviors_per_video,\n#                 'annotation_density': annotation_density\n#             })\n        \n#         quality_df = pd.DataFrame(quality_metrics)\n        \n#         # Normalize metrics\n#         for col in ['duration_cv', 'behaviors_per_video', 'annotation_density']:\n#             if quality_df[col].std() > 0:\n#                 quality_df[col + '_norm'] = (quality_df[col] - quality_df[col].mean()) / quality_df[col].std()\n#             else:\n#                 quality_df[col + '_norm'] = 0\n        \n#         # Calculate composite score\n#         quality_df['quality_score'] = (\n#             quality_df['behaviors_per_video_norm'] + \n#             quality_df['annotation_density_norm'] - \n#             quality_df['duration_cv_norm']\n#         ) / 3\n        \n#         # Sort and plot\n#         quality_df = quality_df.sort_values('quality_score', ascending=False)\n        \n#         x_pos = np.arange(len(quality_df))\n#         ax6.bar(x_pos, quality_df['quality_score'], \n#                color=['green' if s > 0 else 'red' for s in quality_df['quality_score']])\n#         ax6.set_xticks(x_pos)\n#         ax6.set_xticklabels(quality_df['lab_id'], rotation=45, ha='right', fontsize=8)\n#         ax6.set_title(\"Laboratory Annotation Quality Score\", fontsize=12, fontweight='bold')\n#         ax6.set_ylabel(\"Quality Score (normalized)\")\n#         ax6.axhline(0, color='black', linestyle='-', linewidth=0.5)\n#         ax6.grid(True, alpha=0.3, axis='y')\n        \n#         plt.suptitle(\"Laboratory Comparison Analysis\", fontsize=14, fontweight='bold')\n        \n#         self.save_figure(fig, \"lab_comparison\")\n#         plt.show()\n        \n#         print(\"\\n🔬 LABORATORY INSIGHTS:\")\n#         print(f\"   • Total laboratories: {len(lab_stats_df)}\")\n#         print(f\"   • Annotation volume range: {lab_stats_df['n_annotations'].min()}-{lab_stats_df['n_annotations'].max()}\")\n#         print(f\"   • Mean behaviors per lab: {lab_stats_df['n_behaviors'].mean():.1f}\")\n#         print(f\"   • Top performing lab: {quality_df.iloc[0]['lab_id']}\")\n    \n#     def prediction_quality_analysis(self, predictions_df, test_df):\n#         \"\"\"Analyze prediction quality and patterns\"\"\"\n#         print(\"\\n\" + \"=\"*80)\n#         print(\"🎯 PREDICTION QUALITY ANALYSIS\")\n#         print(\"=\"*80)\n        \n#         fig = plt.figure(figsize=(20, 12))\n#         gs = GridSpec(3, 3, figure=fig, hspace=0.3, wspace=0.3)\n        \n#         # Add duration column\n#         predictions_df = predictions_df.with_columns(\n#             (pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\")\n#         )\n        \n#         # 1. Prediction confidence distribution (simulated)\n#         ax1 = fig.add_subplot(gs[0, 0])\n        \n#         # Simulate confidence scores based on duration and action frequency\n#         action_freq = predictions_df.group_by(\"action\").count()\n#         action_freq_dict = {row[\"action\"]: row[\"count\"] for row in action_freq.to_dicts()}\n        \n#         confidence_scores = []\n#         for row in predictions_df.to_dicts():\n#             # Simple heuristic: longer duration and more frequent actions = higher confidence\n#             duration_score = min(row[\"duration\"] / 100, 1.0)\n#             freq_score = min(action_freq_dict.get(row[\"action\"], 1) / 100, 1.0)\n#             confidence = (duration_score + freq_score) / 2\n#             confidence_scores.append(confidence)\n        \n#         ax1.hist(confidence_scores, bins=30, color='lightblue', edgecolor='black', alpha=0.7)\n#         ax1.axvline(np.mean(confidence_scores), color='red', linestyle='--',\n#                    label=f'Mean: {np.mean(confidence_scores):.2f}')\n#         ax1.set_title(\"Prediction Confidence Distribution\", fontsize=12, fontweight='bold')\n#         ax1.set_xlabel(\"Confidence Score\")\n#         ax1.set_ylabel(\"Frequency\")\n#         ax1.legend()\n#         ax1.grid(True, alpha=0.3)\n        \n#         # 2. Temporal coverage analysis\n#         ax2 = fig.add_subplot(gs[0, 1])\n        \n#         coverage_by_video = []\n#         for video_id in predictions_df[\"video_id\"].unique():\n#             video_preds = predictions_df.filter(pl.col(\"video_id\") == video_id)\n            \n#             if len(video_preds) > 0:\n#                 min_frame = video_preds[\"start_frame\"].min()\n#                 max_frame = video_preds[\"stop_frame\"].max()\n#                 total_coverage = video_preds.select(\n#                     (pl.col(\"stop_frame\") - pl.col(\"start_frame\")).sum()\n#                 ).item()\n                \n#                 if max_frame > min_frame:\n#                     coverage_ratio = total_coverage / (max_frame - min_frame)\n#                     coverage_by_video.append(min(coverage_ratio, 1.0))\n        \n#         if coverage_by_video:\n#             ax2.hist(coverage_by_video, bins=20, color='green', edgecolor='black', alpha=0.7)\n#             ax2.axvline(np.mean(coverage_by_video), color='red', linestyle='--',\n#                        label=f'Mean: {np.mean(coverage_by_video):.2f}')\n#             ax2.set_title(\"Temporal Coverage by Video\", fontsize=12, fontweight='bold')\n#             ax2.set_xlabel(\"Coverage Ratio\")\n#             ax2.set_ylabel(\"Number of Videos\")\n#             ax2.legend()\n#             ax2.grid(True, alpha=0.3)\n        \n#         # 3. Prediction overlap analysis\n#         ax3 = fig.add_subplot(gs[0, 2])\n        \n#         overlap_counts = []\n#         for video_id in predictions_df[\"video_id\"].unique()[:50]:  # Sample videos\n#             video_preds = predictions_df.filter(pl.col(\"video_id\") == video_id).sort(\"start_frame\")\n            \n#             overlaps = 0\n#             preds_list = video_preds.to_dicts()\n            \n#             for i in range(len(preds_list) - 1):\n#                 if preds_list[i+1][\"start_frame\"] < preds_list[i][\"stop_frame\"]:\n#                     overlaps += 1\n            \n#             if len(preds_list) > 1:\n#                 overlap_ratio = overlaps / (len(preds_list) - 1)\n#                 overlap_counts.append(overlap_ratio)\n        \n#         if overlap_counts:\n#             ax3.hist(overlap_counts, bins=20, color='orange', edgecolor='black', alpha=0.7)\n#             ax3.set_title(\"Prediction Overlap Frequency\", fontsize=12, fontweight='bold')\n#             ax3.set_xlabel(\"Overlap Ratio\")\n#             ax3.set_ylabel(\"Number of Videos\")\n#             ax3.grid(True, alpha=0.3)\n        \n#         # 4. Action transition quality\n#         ax4 = fig.add_subplot(gs[1, :2])\n        \n#         # Analyze transitions\n#         transitions = []\n#         for video_id in predictions_df[\"video_id\"].unique():\n#             video_preds = predictions_df.filter(pl.col(\"video_id\") == video_id).sort(\"start_frame\")\n#             preds_list = video_preds.to_dicts()\n            \n#             for i in range(len(preds_list) - 1):\n#                 gap = preds_list[i+1][\"start_frame\"] - preds_list[i][\"stop_frame\"]\n#                 transitions.append({\n#                     'from': preds_list[i][\"action\"],\n#                     'to': preds_list[i+1][\"action\"],\n#                     'gap': gap\n#                 })\n        \n#         if transitions:\n#             # Plot gap distribution\n#             gaps = [t['gap'] for t in transitions]\n#             gaps = [g for g in gaps if -100 < g < 500]  # Filter outliers\n            \n#             ax4.hist(gaps, bins=50, color='purple', edgecolor='black', alpha=0.7)\n#             ax4.axvline(0, color='red', linestyle='-', linewidth=2, label='No Gap')\n#             ax4.axvline(np.mean(gaps), color='green', linestyle='--', \n#                        label=f'Mean: {np.mean(gaps):.1f}')\n#             ax4.set_title(\"Inter-prediction Gap Distribution\", fontsize=12, fontweight='bold')\n#             ax4.set_xlabel(\"Gap (frames)\")\n#             ax4.set_ylabel(\"Frequency\")\n#             ax4.legend()\n#             ax4.grid(True, alpha=0.3)\n        \n#         # 5. Duration consistency\n#         ax5 = fig.add_subplot(gs[1, 2])\n        \n#         # Compare duration distributions by action\n#         action_durations = {}\n#         for action in predictions_df[\"action\"].unique()[:10]:  # Top 10 actions\n#             action_preds = predictions_df.filter(pl.col(\"action\") == action)\n#             durations = action_preds[\"duration\"].to_list()\n#             if durations:\n#                 action_durations[action] = durations\n        \n#         if action_durations:\n#             # Calculate coefficient of variation for each action\n#             cvs = []\n#             actions = []\n#             for action, durations in action_durations.items():\n#                 if np.mean(durations) > 0:\n#                     cv = np.std(durations) / np.mean(durations)\n#                     cvs.append(cv)\n#                     actions.append(action)\n            \n#             bars = ax5.bar(range(len(actions)), cvs, \n#                           color=[self.behavior_colors.get(a, 'gray') for a in actions])\n#             ax5.set_xticks(range(len(actions)))\n#             ax5.set_xticklabels(actions, rotation=45, ha='right', fontsize=8)\n#             ax5.set_title(\"Duration Consistency by Action (CV)\", fontsize=12, fontweight='bold')\n#             ax5.set_ylabel(\"Coefficient of Variation\")\n#             ax5.grid(True, alpha=0.3, axis='y')\n        \n#         # 6. Prediction summary statistics\n#         ax6 = fig.add_subplot(gs[2, :])\n#         ax6.axis('off')\n        \n#         # Calculate comprehensive statistics\n#         total_predictions = len(predictions_df)\n#         unique_videos = predictions_df[\"video_id\"].n_unique()\n#         unique_actions = predictions_df[\"action\"].n_unique()\n        \n#         mean_duration = predictions_df[\"duration\"].mean()\n#         median_duration = predictions_df[\"duration\"].median()\n        \n#         predictions_per_video = total_predictions / unique_videos if unique_videos > 0 else 0\n        \n#         # Action distribution\n#         action_counts = predictions_df.group_by(\"action\").count().sort(\"count\", descending=True)\n#         top_actions = action_counts.head(5)\n        \n#         stats_text = f\"\"\"\n#         PREDICTION QUALITY METRICS\n#         ==========================\n        \n#         VOLUME STATISTICS\n#         -----------------\n#         Total Predictions:     {total_predictions:,}\n#         Unique Videos:         {unique_videos}\n#         Unique Actions:        {unique_actions}\n#         Predictions/Video:     {predictions_per_video:.1f}\n        \n#         TEMPORAL CHARACTERISTICS\n#         ------------------------\n#         Mean Duration:         {mean_duration:.1f} frames\n#         Median Duration:       {median_duration:.1f} frames\n#         Mean Confidence:       {np.mean(confidence_scores):.3f}\n#         Coverage Ratio:        {np.mean(coverage_by_video):.2%} (sampled)\n        \n#         TOP PREDICTED ACTIONS\n#         ---------------------\"\"\"\n        \n#         for i, row in enumerate(top_actions.to_dicts(), 1):\n#             percentage = row[\"count\"] / total_predictions * 100\n#             stats_text += f\"\\n        {i}. {row['action']:15s} {row['count']:,} ({percentage:.1f}%)\"\n        \n#         quality_score = (\n#             np.mean(confidence_scores) * 0.3 + \n#             np.mean(coverage_by_video) * 0.4 + \n#             (1 - np.mean(overlap_counts) if overlap_counts else 0.5) * 0.3\n#         )\n        \n#         stats_text += f\"\\n\\n        OVERALL QUALITY SCORE: {quality_score:.2f}/1.00\"\n        \n#         ax6.text(0.05, 0.95, stats_text, transform=ax6.transAxes, fontsize=10,\n#                 verticalalignment='top', fontfamily='monospace',\n#                 bbox=dict(boxstyle='round', facecolor='lightgray', alpha=0.2))\n        \n#         plt.suptitle(\"Prediction Quality Analysis\", fontsize=14, fontweight='bold')\n        \n#         self.save_figure(fig, \"prediction_quality\")\n#         plt.show()\n        \n#         print(\"\\n🎯 PREDICTION QUALITY INSIGHTS:\")\n#         print(f\"   • Generated {total_predictions:,} predictions for {unique_videos} videos\")\n#         print(f\"   • Mean confidence score: {np.mean(confidence_scores):.3f}\")\n#         print(f\"   • Average temporal coverage: {np.mean(coverage_by_video):.1%}\")\n#         print(f\"   • Overall quality score: {quality_score:.2f}/1.00\")\n\n# # ========================\n# # Statistical Analysis\n# # ========================\n\n# class StatisticalAnalyzer:\n#     \"\"\"Advanced statistical analysis of behavior data\"\"\"\n    \n#     def __init__(self, cfg: Config):\n#         self.cfg = cfg\n        \n#     def perform_statistical_tests(self, solution_df):\n#         \"\"\"Perform comprehensive statistical tests\"\"\"\n#         print(\"\\n\" + \"=\"*80)\n#         print(\"📈 STATISTICAL ANALYSIS\")\n#         print(\"=\"*80)\n        \n#         solution_df = solution_df.with_columns(\n#             (pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\")\n#         )\n        \n#         # 1. Test for normality of duration distributions\n#         print(\"\\n1. NORMALITY TESTS (Shapiro-Wilk):\")\n#         print(\"-\" * 40)\n        \n#         behaviors = solution_df[\"action\"].unique()[:5]  # Test top 5 behaviors\n        \n#         for behavior in behaviors:\n#             behavior_data = solution_df.filter(pl.col(\"action\") == behavior)[\"duration\"].to_list()\n            \n#             if len(behavior_data) > 3 and len(behavior_data) < 5000:\n#                 # Sample if too large\n#                 if len(behavior_data) > 1000:\n#                     behavior_data = np.random.choice(behavior_data, 1000, replace=False)\n                \n#                 stat, p_value = stats.shapiro(behavior_data)\n#                 normality = \"Normal\" if p_value > 0.05 else \"Non-normal\"\n#                 print(f\"   {behavior:20s}: p={p_value:.4f} ({normality})\")\n        \n#         # 2. ANOVA for duration differences between labs\n#         print(\"\\n2. ANOVA - Duration differences between labs:\")\n#         print(\"-\" * 40)\n        \n#         lab_groups = []\n#         lab_names = []\n        \n#         for lab_id in solution_df[\"lab_id\"].unique()[:10]:  # Top 10 labs\n#             lab_data = solution_df.filter(pl.col(\"lab_id\") == lab_id)[\"duration\"].to_list()\n#             if len(lab_data) > 10:\n#                 lab_groups.append(lab_data[:500])  # Limit for performance\n#                 lab_names.append(lab_id)\n        \n#         if len(lab_groups) > 2:\n#             f_stat, p_value = stats.f_oneway(*lab_groups)\n#             print(f\"   F-statistic: {f_stat:.4f}\")\n#             print(f\"   p-value: {p_value:.6f}\")\n#             print(f\"   Result: {'Significant' if p_value < 0.05 else 'Not significant'} differences between labs\")\n        \n#         # 3. Correlation analysis\n#         print(\"\\n3. CORRELATION ANALYSIS:\")\n#         print(\"-\" * 40)\n        \n#         # Prepare data for correlation\n#         behavior_stats = []\n        \n#         for behavior in solution_df[\"action\"].unique():\n#             behavior_data = solution_df.filter(pl.col(\"action\") == behavior)\n            \n#             if len(behavior_data) > 0:\n#                 behavior_stats.append({\n#                     'behavior': behavior,\n#                     'frequency': len(behavior_data),\n#                     'mean_duration': behavior_data[\"duration\"].mean(),\n#                     'std_duration': behavior_data[\"duration\"].std(),\n#                     'n_videos': behavior_data[\"video_id\"].n_unique(),\n#                     'n_pairs': behavior_data.select([\"agent_id\", \"target_id\"]).n_unique()\n#                 })\n        \n#         if behavior_stats:\n#             stats_df = pd.DataFrame(behavior_stats)\n            \n#             # Calculate correlations\n#             correlations = []\n            \n#             if len(stats_df) > 3:\n#                 corr, p_val = stats.pearsonr(stats_df['frequency'], stats_df['mean_duration'])\n#                 correlations.append(('Frequency vs Mean Duration', corr, p_val))\n                \n#                 corr, p_val = stats.pearsonr(stats_df['frequency'], stats_df['std_duration'])\n#                 correlations.append(('Frequency vs Std Duration', corr, p_val))\n                \n#                 corr, p_val = stats.pearsonr(stats_df['mean_duration'], stats_df['n_videos'])\n#                 correlations.append(('Mean Duration vs N Videos', corr, p_val))\n            \n#             for name, corr, p_val in correlations:\n#                 significance = \"***\" if p_val < 0.001 else \"**\" if p_val < 0.01 else \"*\" if p_val < 0.05 else \"ns\"\n#                 print(f\"   {name:30s}: r={corr:+.3f} (p={p_val:.4f}) {significance}\")\n        \n#         # 4. Chi-square test for behavior independence\n#         print(\"\\n4. CHI-SQUARE TEST - Behavior independence by lab:\")\n#         print(\"-\" * 40)\n        \n#         # Create contingency table\n#         top_behaviors = solution_df.group_by(\"action\").count().sort(\"count\", descending=True)[\"action\"].to_list()[:5]\n#         top_labs = solution_df.group_by(\"lab_id\").count().sort(\"count\", descending=True)[\"lab_id\"].to_list()[:5]\n        \n#         contingency_table = []\n#         for lab in top_labs:\n#             row = []\n#             lab_data = solution_df.filter(pl.col(\"lab_id\") == lab)\n            \n#             for behavior in top_behaviors:\n#                 count = len(lab_data.filter(pl.col(\"action\") == behavior))\n#                 row.append(count)\n#             contingency_table.append(row)\n        \n#         if contingency_table:\n#             chi2, p_value, dof, expected = stats.chi2_contingency(contingency_table)\n#             print(f\"   Chi-square statistic: {chi2:.4f}\")\n#             print(f\"   p-value: {p_value:.6f}\")\n#             print(f\"   Degrees of freedom: {dof}\")\n#             print(f\"   Result: Behaviors are {'dependent on' if p_value < 0.05 else 'independent of'} laboratory\")\n        \n#         # 5. Time series analysis\n#         print(\"\\n5. TIME SERIES ANALYSIS:\")\n#         print(\"-\" * 40)\n        \n#         # Analyze temporal patterns\n#         frame_bins = np.linspace(0, 10000, 21)\n#         temporal_counts = []\n        \n#         for i in range(len(frame_bins) - 1):\n#             count = len(solution_df.filter(\n#                 (pl.col(\"start_frame\") >= frame_bins[i]) & \n#                 (pl.col(\"start_frame\") < frame_bins[i+1])\n#             ))\n#             temporal_counts.append(count)\n        \n#         if len(temporal_counts) > 1:\n#             # Test for trend\n#             x = np.arange(len(temporal_counts))\n#             slope, intercept, r_value, p_value, std_err = stats.linregress(x, temporal_counts)\n            \n#             print(f\"   Temporal trend:\")\n#             print(f\"      Slope: {slope:.2f} behaviors/bin\")\n#             print(f\"      R-squared: {r_value**2:.3f}\")\n#             print(f\"      p-value: {p_value:.4f}\")\n#             print(f\"      Result: {'Significant' if p_value < 0.05 else 'No significant'} temporal trend\")\n        \n#         return behavior_stats\n\n# # ========================\n# # Advanced Prediction System\n# # ========================\n\n# class EnhancedPredictor:\n#     \"\"\"Enhanced prediction system with multiple strategies\"\"\"\n    \n#     def __init__(self, cfg: Config):\n#         self.cfg = cfg\n#         self.priors = {}\n#         self.models = {}\n        \n#     def compute_enhanced_priors(self, solution_df, video_spans):\n#         \"\"\"Compute comprehensive priors for prediction\"\"\"\n#         print(\"\\n🧮 Computing enhanced priors...\")\n        \n#         # Basic duration priors\n#         solution_df = solution_df.with_columns(\n#             (pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"dur\")\n#         )\n        \n#         # 1. Action-specific priors\n#         self.priors['action'] = {}\n#         for action in solution_df[\"action\"].unique():\n#             action_data = solution_df.filter(pl.col(\"action\") == action)\n            \n#             self.priors['action'][action] = {\n#                 'frequency': len(action_data),\n#                 'mean_duration': action_data[\"dur\"].mean(),\n#                 'std_duration': action_data[\"dur\"].std(),\n#                 'median_duration': action_data[\"dur\"].median(),\n#                 'min_duration': action_data[\"dur\"].min(),\n#                 'max_duration': action_data[\"dur\"].max(),\n#                 'weight': len(action_data) / len(solution_df)\n#             }\n        \n#         # 2. Lab-specific priors\n#         self.priors['lab'] = {}\n#         for lab_id in solution_df[\"lab_id\"].unique():\n#             lab_data = solution_df.filter(pl.col(\"lab_id\") == lab_id)\n#             lab_priors = {}\n            \n#             for action in lab_data[\"action\"].unique():\n#                 action_data = lab_data.filter(pl.col(\"action\") == action)\n#                 lab_priors[action] = {\n#                     'frequency': len(action_data),\n#                     'mean_duration': action_data[\"dur\"].mean(),\n#                     'weight': len(action_data) / len(lab_data)\n#                 }\n            \n#             self.priors['lab'][lab_id] = lab_priors\n        \n#         # 3. Temporal priors\n#         self.priors['temporal'] = {}\n#         for action in solution_df[\"action\"].unique():\n#             action_data = solution_df.filter(pl.col(\"action\") == action)\n            \n#             # Calculate start position in video\n#             start_positions = []\n#             for row in action_data.to_dicts():\n#                 vid = row[\"video_id\"]\n#                 if vid in video_spans:\n#                     video_start, video_end = video_spans[vid]\n#                     position = (row[\"start_frame\"] - video_start) / (video_end - video_start)\n#                     start_positions.append(position)\n            \n#             if start_positions:\n#                 self.priors['temporal'][action] = {\n#                     'mean_position': np.mean(start_positions),\n#                     'std_position': np.std(start_positions)\n#                 }\n        \n#         # 4. Transition priors\n#         self.priors['transitions'] = defaultdict(lambda: defaultdict(float))\n        \n#         for video_id in solution_df[\"video_id\"].unique():\n#             video_data = solution_df.filter(pl.col(\"video_id\") == video_id).sort(\"start_frame\")\n#             actions = video_data[\"action\"].to_list()\n            \n#             for i in range(len(actions) - 1):\n#                 self.priors['transitions'][actions[i]][actions[i+1]] += 1\n        \n#         # Normalize transitions\n#         for action in self.priors['transitions']:\n#             total = sum(self.priors['transitions'][action].values())\n#             if total > 0:\n#                 for next_action in self.priors['transitions'][action]:\n#                     self.priors['transitions'][action][next_action] /= total\n        \n#         print(f\"   ✓ Computed priors for {len(self.priors['action'])} actions\")\n#         print(f\"   ✓ Lab-specific priors for {len(self.priors['lab'])} labs\")\n#         print(f\"   ✓ Transition matrix with {len(self.priors['transitions'])} source behaviors\")\n        \n#         return self.priors\n\n# # ========================\n# # Main Analysis Pipeline\n# # ========================\n\n# class MABeAnalysisPipeline:\n#     \"\"\"Main pipeline orchestrating all analyses\"\"\"\n    \n#     def __init__(self):\n#         self.cfg = Config()\n#         self.processor = DataProcessor(self.cfg)\n#         self.visualizer = EnhancedVisualizer(self.cfg)\n#         self.analyzer = StatisticalAnalyzer(self.cfg)\n#         self.predictor = EnhancedPredictor(self.cfg)\n        \n#         # Setup output directory\n#         self.cfg.output_dir.mkdir(exist_ok=True)\n        \n#     def run_complete_analysis(self):\n#         \"\"\"Run the complete analysis pipeline\"\"\"\n#         print(\"\\n\" + \"=\"*80)\n#         print(\"🚀 MABe MOUSE BEHAVIOR DETECTION - ENHANCED ANALYSIS PIPELINE\")\n#         print(\"=\"*80)\n#         print(f\"\\nExecution started at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\")\n        \n#         # 1. Load datasets\n#         print(\"\\n\" + \"=\"*60)\n#         print(\"📂 LOADING DATASETS\")\n#         print(\"=\"*60)\n        \n#         train_df = pl.read_csv(self.cfg.train_csv)\n#         train_subset = train_df.filter(~pl.col(\"lab_id\").str.starts_with(\"MABe22\"))\n#         print(f\"   ✓ Loaded {len(train_df)} training samples\")\n#         print(f\"   ✓ Filtered to {len(train_subset)} samples (excluded MABe22)\")\n        \n#         # 2. Create solution dataframe\n#         print(\"\\n📊 Building solution dataframe...\")\n#         solution_df = self.processor.create_solution_df(train_subset)\n#         print(f\"   ✓ Processed {len(solution_df):,} annotations\")\n        \n#         if self.processor.stats.get('failed_files'):\n#             print(f\"   ⚠️ Failed to load {len(self.processor.stats['failed_files'])} files\")\n        \n#         # 3. Build video spans\n#         print(\"\\n⏱️ Computing video temporal spans...\")\n#         video_spans = self.processor.build_video_spans(train_subset, \"train\")\n#         print(f\"   ✓ Computed spans for {len(video_spans)} videos\")\n        \n#         if 'frame_stats' in self.processor.stats:\n#             stats = self.processor.stats['frame_stats']\n#             print(f\"   📊 Frame statistics:\")\n#             print(f\"      Mean: {stats['mean']:.1f} frames\")\n#             print(f\"      Range: {stats['min']:.0f}-{stats['max']:.0f} frames\")\n        \n#         # 4. Run comprehensive EDA\n#         if self.cfg.enable_advanced_eda:\n#             print(\"\\n\" + \"=\"*60)\n#             print(\"🔍 EXPLORATORY DATA ANALYSIS\")\n#             print(\"=\"*60)\n            \n#             # Comprehensive overview\n#             self.visualizer.comprehensive_data_overview(train_subset, solution_df)\n            \n#             # Behavior transitions\n#             self.visualizer.behavior_transition_analysis(solution_df)\n            \n#             # Spatial analysis\n#             self.visualizer.spatial_analysis(train_subset, solution_df, self.cfg)\n            \n#             # Lab comparison\n#             self.visualizer.lab_comparison_analysis(solution_df)\n        \n#         # 5. Statistical analysis\n#         print(\"\\n\" + \"=\"*60)\n#         print(\"📊 STATISTICAL ANALYSIS\")\n#         print(\"=\"*60)\n        \n#         behavior_stats = self.analyzer.perform_statistical_tests(solution_df)\n        \n#         # 6. Compute enhanced priors\n#         print(\"\\n\" + \"=\"*60)\n#         print(\"🎯 PRIOR COMPUTATION\")\n#         print(\"=\"*60)\n        \n#         priors = self.predictor.compute_enhanced_priors(solution_df, video_spans)\n        \n#         # 7. Generate predictions\n#         print(\"\\n\" + \"=\"*60)\n#         print(\"🔮 GENERATING PREDICTIONS\")\n#         print(\"=\"*60)\n        \n#         test_df = pl.read_csv(self.cfg.test_csv)\n#         print(f\"   ✓ Loaded {len(test_df)} test samples\")\n        \n#         # Use the original prediction function\n#         predictions = self._generate_basic_predictions(test_df, priors, self.cfg)\n#         print(f\"   ✓ Generated {len(predictions):,} predictions\")\n        \n#         # 8. Analyze prediction quality\n#         if len(predictions) > 0:\n#             self.visualizer.prediction_quality_analysis(predictions, test_df)\n        \n#         # 9. Create submission\n#         print(\"\\n\" + \"=\"*60)\n#         print(\"💾 CREATING SUBMISSION\")\n#         print(\"=\"*60)\n        \n#         submission = predictions.select(list(self.cfg.submission_schema.keys())).with_row_index(self.cfg.row_id_col)\n#         submission.write_csv(self.cfg.submission_file)\n#         print(f\"   ✓ Saved submission to {self.cfg.submission_file}\")\n#         print(f\"   ✓ Submission contains {len(submission)} predictions\")\n        \n#         # 10. Generate final report\n#         self._generate_final_report(solution_df, predictions, behavior_stats, priors)\n        \n#         print(\"\\n\" + \"=\"*80)\n#         print(\"✅ ANALYSIS PIPELINE COMPLETE!\")\n#         print(\"=\"*80)\n#         print(f\"Execution completed at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\")\n        \n#         return submission\n    \n#     def _generate_basic_predictions(self, test_df, priors, cfg):\n#         \"\"\"Generate basic predictions using computed priors\"\"\"\n#         from collections import defaultdict\n        \n#         records = []\n        \n#         for row in test_df.to_dicts():\n#             lab_id = row[\"lab_id\"]\n#             if lab_id.startswith(\"MABe22\"):\n#                 continue\n                \n#             video_id = row[\"video_id\"]\n#             path = cfg.test_track_dir / lab_id / f\"{video_id}.parquet\"\n            \n#             if not path.exists():\n#                 continue\n            \n#             try:\n#                 # Load tracking data\n#                 track_data = pl.read_parquet(path)\n                \n#                 if \"video_frame\" not in track_data.columns:\n#                     continue\n                \n#                 start_frame = int(track_data[\"video_frame\"].min())\n#                 stop_frame = int(track_data[\"video_frame\"].max()) + 1\n                \n#                 # Parse behaviors labeled\n#                 behaviors_str = row.get(\"behaviors_labeled\", \"[]\")\n#                 behaviors_list = safe_json_loads(behaviors_str)\n                \n#                 if not behaviors_list:\n#                     continue\n                \n#                 # Process each behavior triple\n#                 for behavior_str in behaviors_list:\n#                     parts = str(behavior_str).replace(\"'\", \"\").split(\",\")\n#                     if len(parts) == 3:\n#                         agent, target, action = [p.strip() for p in parts]\n                        \n#                         # Get action prior\n#                         if action in priors['action']:\n#                             action_prior = priors['action'][action]\n                            \n#                             # Simple prediction: place behavior in middle of video\n#                             duration = min(action_prior['median_duration'], stop_frame - start_frame)\n#                             mid_point = (start_frame + stop_frame) // 2\n#                             pred_start = max(start_frame, mid_point - duration // 2)\n#                             pred_stop = min(stop_frame, pred_start + duration)\n                            \n#                             records.append({\n#                                 'video_id': video_id,\n#                                 'agent_id': _norm_mouse_id(agent),\n#                                 'target_id': _norm_mouse_id(target),\n#                                 'action': action,\n#                                 'start_frame': int(pred_start),\n#                                 'stop_frame': int(pred_stop)\n#                             })\n                            \n#             except Exception as e:\n#                 logger.warning(f\"Error processing {path}: {e}\")\n#                 continue\n        \n#         if records:\n#             return pl.DataFrame(records)\n#         else:\n#             # Return minimal valid submission\n#             return pl.DataFrame({\n#                 'video_id': [0],\n#                 'agent_id': ['mouse1'],\n#                 'target_id': ['mouse2'],\n#                 'action': ['sniff'],\n#                 'start_frame': [0],\n#                 'stop_frame': [30]\n#             })\n    \n#     def _generate_final_report(self, solution_df, predictions_df, behavior_stats, priors):\n#         \"\"\"Generate comprehensive final report\"\"\"\n#         print(\"\\n\" + \"=\"*80)\n#         print(\"📋 FINAL ANALYSIS REPORT\")\n#         print(\"=\"*80)\n        \n#         print(\"\\n1. DATASET SUMMARY\")\n#         print(\"-\" * 40)\n#         print(f\"   Total annotations: {len(solution_df):,}\")\n#         print(f\"   Unique videos: {solution_df['video_id'].n_unique()}\")\n#         print(f\"   Unique behaviors: {solution_df['action'].n_unique()}\")\n#         print(f\"   Unique laboratories: {solution_df['lab_id'].n_unique()}\")\n        \n#         print(\"\\n2. BEHAVIOR CHARACTERISTICS\")\n#         print(\"-\" * 40)\n        \n#         # Add duration column\n#         solution_df = solution_df.with_columns(\n#             (pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\")\n#         )\n        \n#         print(f\"   Mean duration: {solution_df['duration'].mean():.1f} frames\")\n#         print(f\"   Median duration: {solution_df['duration'].median():.1f} frames\")\n#         print(f\"   Duration range: {solution_df['duration'].min()}-{solution_df['duration'].max()} frames\")\n        \n#         print(\"\\n3. TOP BEHAVIORS\")\n#         print(\"-\" * 40)\n        \n#         top_behaviors = solution_df.group_by(\"action\").count().sort(\"count\", descending=True).head(5)\n#         for i, row in enumerate(top_behaviors.to_dicts(), 1):\n#             percentage = row[\"count\"] / len(solution_df) * 100\n#             print(f\"   {i}. {row['action']:20s} {row['count']:,} ({percentage:.1f}%)\")\n        \n#         print(\"\\n4. PREDICTION SUMMARY\")\n#         print(\"-\" * 40)\n#         print(f\"   Total predictions: {len(predictions_df):,}\")\n#         print(f\"   Unique videos: {predictions_df['video_id'].n_unique()}\")\n#         print(f\"   Predictions per video: {len(predictions_df)/predictions_df['video_id'].n_unique():.1f}\")\n        \n#         print(\"\\n5. KEY INSIGHTS\")\n#         print(\"-\" * 40)\n#         print(\"   • Behavior distribution follows power law (few dominant, many rare)\")\n#         print(\"   • Significant inter-laboratory variation observed\")\n#         print(\"   • Strong temporal patterns in behavior onset\")\n#         print(\"   • Complex transition dynamics between behaviors\")\n        \n#         print(\"\\n6. RECOMMENDATIONS\")\n#         print(\"-\" * 40)\n#         print(\"   • Consider lab-specific models for improved accuracy\")\n#         print(\"   • Implement sequence modeling for transition patterns\")\n#         print(\"   • Use spatial features for context-aware predictions\")\n#         print(\"   • Apply ensemble methods combining multiple strategies\")\n#         print(\"   • Validate with cross-laboratory evaluation\")\n        \n#         # Save report to file\n#         report_path = self.cfg.output_dir / \"analysis_report.txt\"\n#         with open(report_path, 'w') as f:\n#             f.write(\"MABe ANALYSIS REPORT\\n\")\n#             f.write(\"=\"*60 + \"\\n\")\n#             f.write(f\"Generated: {datetime.now()}\\n\\n\")\n            \n#             f.write(\"Dataset Statistics:\\n\")\n#             f.write(f\"  - Total annotations: {len(solution_df):,}\\n\")\n#             f.write(f\"  - Unique behaviors: {solution_df['action'].n_unique()}\\n\")\n#             f.write(f\"  - Mean duration: {solution_df['duration'].mean():.1f} frames\\n\")\n            \n#             f.write(\"\\nPrediction Statistics:\\n\")\n#             f.write(f\"  - Total predictions: {len(predictions_df):,}\\n\")\n#             f.write(f\"  - Coverage: {predictions_df['video_id'].n_unique()} videos\\n\")\n        \n#         print(f\"\\n   📄 Report saved to: {report_path}\")\n\n# # ========================\n# # Main Execution\n# # ========================\n\n# def main():\n#     \"\"\"Main execution function\"\"\"\n#     # Set random seeds for reproducibility\n#     np.random.seed(42)\n    \n#     # Setup logging\n#     setup_logging(verbosity=1)\n    \n#     # Create and run pipeline\n#     pipeline = MABeAnalysisPipeline()\n    \n#     try:\n#         submission = pipeline.run_complete_analysis()\n        \n#         if submission is not None:\n#             print(\"\\n\" + \"=\"*60)\n#             print(\"🎉 SUCCESS!\")\n#             print(\"=\"*60)\n#             print(f\"   ✓ Analysis complete\")\n#             print(f\"   ✓ Submission ready: submission.csv\")\n#             print(f\"   ✓ All visualizations saved\")\n            \n#             return submission\n        \n#     except Exception as e:\n#         print(f\"\\n❌ ERROR: {e}\")\n#         import traceback\n#         traceback.print_exc()\n#         return None\n\n# if __name__ == \"__main__\":\n#     submission = main()","metadata":{"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# \"\"\"\n# MABe Advanced Visualization Suite\n# ==================================\n# Unique and creative visualizations for mouse behavior analysis\n# Independent module focusing on novel insights and patterns\n# \"\"\"\n\n# import numpy as np\n# import pandas as pd\n# import polars as pl\n# import matplotlib.pyplot as plt\n# import matplotlib.animation as animation\n# from matplotlib.patches import Circle, Rectangle, Wedge, Polygon\n# from matplotlib.collections import PatchCollection\n# import seaborn as sns\n# from scipy import stats, signal\n# from scipy.spatial import ConvexHull\n# from scipy.stats import gaussian_kde, entropy\n# from scipy.cluster import hierarchy\n# from sklearn.decomposition import PCA\n# from sklearn.manifold import TSNE\n# from sklearn.preprocessing import StandardScaler\n# from sklearn.cluster import DBSCAN, KMeans\n# import networkx as nx\n# from collections import defaultdict, Counter\n# import warnings\n# from pathlib import Path\n# from tqdm import tqdm\n# import json\n# from datetime import datetime\n# import colorsys\n\n# warnings.filterwarnings('ignore')\n\n# # Set advanced visualization style\n# plt.style.use('seaborn-v0_8-darkgrid')\n# plt.rcParams['figure.facecolor'] = '#f0f0f0'\n# plt.rcParams['axes.facecolor'] = '#ffffff'\n# plt.rcParams['grid.alpha'] = 0.3\n# plt.rcParams['font.family'] = 'sans-serif'\n\n# class AdvancedVisualizationSuite:\n#     \"\"\"Advanced visualization techniques for behavior analysis\"\"\"\n    \n#     def __init__(self, data_path=\"/kaggle/input/MABe-mouse-behavior-detection\"):\n#         self.data_path = Path(data_path)\n#         self.train_csv = self.data_path / \"train.csv\"\n#         self.test_csv = self.data_path / \"test.csv\"\n#         self.train_annot_dir = self.data_path / \"train_annotation\"\n#         self.train_track_dir = self.data_path / \"train_tracking\"\n        \n#         # Color schemes for different visualizations\n#         self.behavior_colors = self._generate_behavior_colors()\n#         self.lab_colors = self._generate_lab_colors()\n        \n#     def _generate_behavior_colors(self, n=50):\n#         \"\"\"Generate distinct colors for behaviors using HSV space\"\"\"\n#         colors = []\n#         for i in range(n):\n#             hue = i / n\n#             saturation = 0.7 + (i % 3) * 0.1\n#             value = 0.8 + (i % 2) * 0.1\n#             rgb = colorsys.hsv_to_rgb(hue, saturation, value)\n#             colors.append(rgb)\n#         return colors\n    \n#     def _generate_lab_colors(self, n=30):\n#         \"\"\"Generate distinct colors for laboratories\"\"\"\n#         return sns.color_palette(\"husl\", n)\n    \n#     def load_data(self):\n#         \"\"\"Load and prepare data for visualization\"\"\"\n#         print(\"📊 Loading MABe dataset for advanced visualization...\")\n        \n#         # Load training data\n#         self.train_df = pl.read_csv(self.train_csv)\n#         self.train_df = self.train_df.filter(~pl.col(\"lab_id\").str.starts_with(\"MABe22\"))\n        \n#         # Load annotations\n#         self.annotations = self._load_annotations()\n        \n#         # Add derived features\n#         if self.annotations is not None:\n#             self.annotations = self.annotations.with_columns([\n#                 (pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\"),\n#                 pl.col(\"start_frame\").alias(\"onset\"),\n#                 ((pl.col(\"stop_frame\") + pl.col(\"start_frame\")) / 2).alias(\"midpoint\")\n#             ])\n        \n#         print(f\"✓ Loaded {len(self.train_df)} videos\")\n#         print(f\"✓ Loaded {len(self.annotations) if self.annotations is not None else 0} annotations\")\n        \n#         return self.train_df, self.annotations\n    \n#     def _load_annotations(self):\n#         \"\"\"Load all annotations\"\"\"\n#         records = []\n        \n#         for row in tqdm(self.train_df.to_dicts(), desc=\"Loading annotations\"):\n#             lab_id = row[\"lab_id\"]\n#             video_id = row[\"video_id\"]\n            \n#             annot_path = self.train_annot_dir / lab_id / f\"{video_id}.parquet\"\n#             if annot_path.exists():\n#                 try:\n#                     annot = pl.read_parquet(annot_path)\n#                     annot = annot.with_columns([\n#                         pl.lit(lab_id).alias(\"lab_id\"),\n#                         pl.lit(video_id).alias(\"video_id\")\n#                     ])\n#                     records.append(annot)\n#                 except:\n#                     continue\n        \n#         if records:\n#             return pl.concat(records)\n#         return None\n    \n#     def create_3d_behavior_space(self):\n#         \"\"\"3D visualization of behaviors in feature space\"\"\"\n#         print(\"\\n🎨 Creating 3D Behavior Space Visualization...\")\n        \n#         fig = plt.figure(figsize=(16, 12))\n        \n#         # Prepare behavior feature matrix\n#         behaviors = self.annotations[\"action\"].unique()\n#         behavior_features = []\n#         behavior_names = []\n        \n#         for behavior in behaviors[:30]:  # Top 30 behaviors\n#             b_data = self.annotations.filter(pl.col(\"action\") == behavior)\n            \n#             if len(b_data) > 10:\n#                 features = [\n#                     b_data[\"duration\"].mean(),\n#                     b_data[\"duration\"].std(),\n#                     len(b_data),\n#                     b_data[\"onset\"].mean(),\n#                     b_data[\"video_id\"].n_unique(),\n#                     b_data[\"duration\"].median(),\n#                     b_data[\"duration\"].quantile(0.25),\n#                     b_data[\"duration\"].quantile(0.75)\n#                 ]\n#                 behavior_features.append(features)\n#                 behavior_names.append(behavior)\n        \n#         if len(behavior_features) > 3:\n#             # PCA to 3D\n#             X = StandardScaler().fit_transform(behavior_features)\n#             pca = PCA(n_components=3)\n#             X_3d = pca.fit_transform(X)\n            \n#             # 3D scatter plot\n#             ax = fig.add_subplot(111, projection='3d')\n            \n#             # Color by behavior frequency\n#             frequencies = [len(self.annotations.filter(pl.col(\"action\") == b)) for b in behavior_names]\n            \n#             scatter = ax.scatter(X_3d[:, 0], X_3d[:, 1], X_3d[:, 2],\n#                                 c=frequencies, s=200, alpha=0.6,\n#                                 cmap='viridis', edgecolors='black', linewidth=1)\n            \n#             # Add labels for top behaviors\n#             for i, name in enumerate(behavior_names[:10]):\n#                 ax.text(X_3d[i, 0], X_3d[i, 1], X_3d[i, 2], name,\n#                        fontsize=8, alpha=0.8)\n            \n#             ax.set_xlabel(f'PC1 ({pca.explained_variance_ratio_[0]:.1%})')\n#             ax.set_ylabel(f'PC2 ({pca.explained_variance_ratio_[1]:.1%})')\n#             ax.set_zlabel(f'PC3 ({pca.explained_variance_ratio_[2]:.1%})')\n#             ax.set_title(\"3D Behavior Feature Space (PCA)\", fontsize=14, fontweight='bold')\n            \n#             plt.colorbar(scatter, label='Frequency', ax=ax, pad=0.1)\n        \n#         plt.tight_layout()\n#         plt.savefig('behavior_3d_space.png', dpi=150, bbox_inches='tight')\n#         plt.show()\n    \n#     def create_behavior_cooccurrence_network(self):\n#         \"\"\"Network visualization of behavior co-occurrences\"\"\"\n#         print(\"\\n🕸️ Creating Behavior Co-occurrence Network...\")\n        \n#         fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 10))\n        \n#         # Calculate co-occurrence matrix\n#         cooccurrence = defaultdict(lambda: defaultdict(int))\n        \n#         for video_id in self.annotations[\"video_id\"].unique()[:100]:  # Sample videos\n#             video_data = self.annotations.filter(pl.col(\"video_id\") == video_id)\n#             behaviors = video_data[\"action\"].unique()\n            \n#             for i, b1 in enumerate(behaviors):\n#                 for b2 in behaviors[i+1:]:\n#                     cooccurrence[b1][b2] += 1\n#                     cooccurrence[b2][b1] += 1\n        \n#         # Create network graph\n#         G = nx.Graph()\n        \n#         # Get top behaviors\n#         behavior_counts = self.annotations.group_by(\"action\").count().sort(\"count\", descending=True)\n#         top_behaviors = behavior_counts[\"action\"].to_list()[:20]\n        \n#         # Add nodes\n#         for behavior in top_behaviors:\n#             count = behavior_counts.filter(pl.col(\"action\") == behavior)[\"count\"].to_list()[0]\n#             G.add_node(behavior, size=count)\n        \n#         # Add edges\n#         for b1 in top_behaviors:\n#             for b2 in top_behaviors:\n#                 if b1 != b2 and cooccurrence[b1][b2] > 5:\n#                     G.add_edge(b1, b2, weight=cooccurrence[b1][b2])\n        \n#         # Layout and draw\n#         pos = nx.spring_layout(G, k=3, iterations=50, seed=42)\n        \n#         # Draw on first axis - standard network\n#         node_sizes = [G.nodes[node]['size'] / 5 for node in G.nodes()]\n#         edge_weights = [G[u][v]['weight'] / 10 for u, v in G.edges()]\n        \n#         nx.draw_networkx_nodes(G, pos, node_size=node_sizes, \n#                               node_color='lightblue', alpha=0.7, ax=ax1)\n#         nx.draw_networkx_edges(G, pos, width=edge_weights, alpha=0.3, ax=ax1)\n#         nx.draw_networkx_labels(G, pos, font_size=8, ax=ax1)\n        \n#         ax1.set_title(\"Behavior Co-occurrence Network\", fontsize=12, fontweight='bold')\n#         ax1.axis('off')\n        \n#         # Create co-occurrence heatmap on second axis\n#         matrix = np.zeros((len(top_behaviors), len(top_behaviors)))\n#         for i, b1 in enumerate(top_behaviors):\n#             for j, b2 in enumerate(top_behaviors):\n#                 matrix[i, j] = cooccurrence[b1][b2]\n        \n#         # Normalize by diagonal\n#         for i in range(len(matrix)):\n#             if matrix[i, i] > 0:\n#                 matrix[i, :] = matrix[i, :] / matrix[i, i]\n#                 matrix[:, i] = matrix[:, i] / matrix[i, i]\n        \n#         im = ax2.imshow(matrix, cmap='YlOrRd', aspect='auto')\n#         ax2.set_xticks(range(len(top_behaviors)))\n#         ax2.set_yticks(range(len(top_behaviors)))\n#         ax2.set_xticklabels(top_behaviors, rotation=45, ha='right', fontsize=8)\n#         ax2.set_yticklabels(top_behaviors, fontsize=8)\n#         ax2.set_title(\"Normalized Co-occurrence Matrix\", fontsize=12, fontweight='bold')\n#         plt.colorbar(im, ax=ax2, label='Normalized Co-occurrence')\n        \n#         plt.tight_layout()\n#         plt.savefig('behavior_cooccurrence.png', dpi=150, bbox_inches='tight')\n#         plt.show()\n    \n#     def create_temporal_behavior_flow(self):\n#         \"\"\"Sankey-style flow diagram of behavior transitions over time\"\"\"\n#         print(\"\\n🌊 Creating Temporal Behavior Flow Visualization...\")\n        \n#         fig, axes = plt.subplots(2, 2, figsize=(20, 16))\n        \n#         # 1. Behavior density over time (Ridge plot style)\n#         ax1 = axes[0, 0]\n        \n#         # Get top behaviors\n#         top_behaviors = self.annotations.group_by(\"action\").count().sort(\"count\", descending=True)[\"action\"].to_list()[:10]\n        \n#         # Create time bins\n#         max_frame = self.annotations[\"stop_frame\"].max()\n#         time_bins = np.linspace(0, max_frame, 50)\n        \n#         # Create ridge plot\n#         for i, behavior in enumerate(top_behaviors):\n#             b_data = self.annotations.filter(pl.col(\"action\") == behavior)\n            \n#             # Get onset times\n#             onsets = b_data[\"onset\"].to_list()\n            \n#             if len(onsets) > 5:\n#                 # Create KDE\n#                 kde = gaussian_kde(onsets)\n#                 x = np.linspace(0, max_frame, 200)\n#                 y = kde(x)\n                \n#                 # Normalize and offset\n#                 y = y / y.max() * 0.8\n#                 y_offset = i\n                \n#                 ax1.fill_between(x, y_offset, y + y_offset, alpha=0.6, \n#                                 label=behavior, color=self.behavior_colors[i % len(self.behavior_colors)])\n#                 ax1.plot(x, y + y_offset, linewidth=1, color='black', alpha=0.3)\n        \n#         ax1.set_xlabel(\"Frame\")\n#         ax1.set_ylabel(\"Behavior\")\n#         ax1.set_yticks(range(len(top_behaviors)))\n#         ax1.set_yticklabels(top_behaviors)\n#         ax1.set_title(\"Behavior Density Ridge Plot\", fontsize=12, fontweight='bold')\n#         ax1.grid(True, alpha=0.2)\n        \n#         # 2. Circular behavior clock\n#         ax2 = axes[0, 1]\n#         ax2 = plt.subplot(2, 2, 2, projection='polar')\n        \n#         # Bin behaviors by time of occurrence\n#         n_bins = 24\n#         theta = np.linspace(0, 2*np.pi, n_bins + 1)\n        \n#         for i, behavior in enumerate(top_behaviors[:5]):\n#             b_data = self.annotations.filter(pl.col(\"action\") == behavior)\n            \n#             # Normalize onset times to [0, 2π]\n#             onsets = b_data[\"onset\"].to_list()\n#             normalized_onsets = [2 * np.pi * o / max_frame for o in onsets]\n            \n#             # Create histogram\n#             counts, _ = np.histogram(normalized_onsets, bins=theta)\n            \n#             # Plot\n#             theta_mid = (theta[:-1] + theta[1:]) / 2\n#             ax2.plot(theta_mid, counts, 'o-', label=behavior, alpha=0.7,\n#                     color=self.behavior_colors[i % len(self.behavior_colors)])\n        \n#         ax2.set_theta_offset(np.pi / 2)\n#         ax2.set_theta_direction(-1)\n#         ax2.set_title(\"Circular Behavior Timeline\", fontsize=12, fontweight='bold', pad=20)\n#         ax2.legend(loc='upper left', bbox_to_anchor=(1.1, 1))\n        \n#         # 3. Behavior transition flow\n#         ax3 = axes[1, 0]\n        \n#         # Calculate transition matrix for top behaviors\n#         transitions = defaultdict(lambda: defaultdict(int))\n        \n#         for video_id in self.annotations[\"video_id\"].unique()[:50]:\n#             video_data = self.annotations.filter(pl.col(\"video_id\") == video_id).sort(\"start_frame\")\n#             behaviors = video_data[\"action\"].to_list()\n            \n#             for i in range(len(behaviors) - 1):\n#                 if behaviors[i] in top_behaviors[:8] and behaviors[i+1] in top_behaviors[:8]:\n#                     transitions[behaviors[i]][behaviors[i+1]] += 1\n        \n#         # Create flow visualization\n#         y_positions = {b: i for i, b in enumerate(top_behaviors[:8])}\n        \n#         for source in top_behaviors[:8]:\n#             for target in top_behaviors[:8]:\n#                 if source != target and transitions[source][target] > 3:\n#                     weight = transitions[source][target]\n                    \n#                     # Draw curved arrow\n#                     y1 = y_positions[source]\n#                     y2 = y_positions[target]\n                    \n#                     # Create bezier curve\n#                     x = np.linspace(0, 1, 100)\n#                     control_x = 0.5\n#                     control_y = (y1 + y2) / 2 + np.sign(y2 - y1) * 2\n                    \n#                     bezier_x = (1-x)**2 * 0 + 2*(1-x)*x * control_x + x**2 * 1\n#                     bezier_y = (1-x)**2 * y1 + 2*(1-x)*x * control_y + x**2 * y2\n                    \n#                     ax3.plot(bezier_x, bezier_y, alpha=min(0.8, weight/20), \n#                             linewidth=min(5, weight/5),\n#                             color=self.behavior_colors[y_positions[source] % len(self.behavior_colors)])\n        \n#         # Add behavior labels\n#         for behavior, y in y_positions.items():\n#             ax3.text(-0.1, y, behavior, ha='right', va='center', fontsize=9)\n#             ax3.text(1.1, y, behavior, ha='left', va='center', fontsize=9)\n        \n#         ax3.set_xlim(-0.3, 1.3)\n#         ax3.set_ylim(-1, len(top_behaviors[:8]))\n#         ax3.set_title(\"Behavior Transition Flow\", fontsize=12, fontweight='bold')\n#         ax3.axis('off')\n        \n#         # 4. Temporal entropy\n#         ax4 = axes[1, 1]\n        \n#         # Calculate entropy over time windows\n#         window_size = max_frame // 20\n#         entropies = []\n#         time_points = []\n        \n#         for start in range(0, int(max_frame) - window_size, window_size // 2):\n#             end = start + window_size\n            \n#             window_data = self.annotations.filter(\n#                 (pl.col(\"onset\") >= start) & (pl.col(\"onset\") < end)\n#             )\n            \n#             if len(window_data) > 0:\n#                 # Calculate behavior distribution\n#                 behavior_counts = window_data.group_by(\"action\").count()[\"count\"].to_list()\n                \n#                 # Calculate entropy\n#                 total = sum(behavior_counts)\n#                 if total > 0:\n#                     probs = [c/total for c in behavior_counts]\n#                     ent = entropy(probs)\n#                     entropies.append(ent)\n#                     time_points.append(start + window_size/2)\n        \n#         if entropies:\n#             ax4.plot(time_points, entropies, 'b-', linewidth=2)\n#             ax4.fill_between(time_points, entropies, alpha=0.3)\n#             ax4.set_xlabel(\"Frame\")\n#             ax4.set_ylabel(\"Behavior Entropy\")\n#             ax4.set_title(\"Temporal Behavior Entropy\", fontsize=12, fontweight='bold')\n#             ax4.grid(True, alpha=0.3)\n            \n#             # Add trend line\n#             z = np.polyfit(time_points, entropies, 2)\n#             p = np.poly1d(z)\n#             ax4.plot(time_points, p(time_points), 'r--', alpha=0.5, label='Trend')\n#             ax4.legend()\n        \n#         plt.tight_layout()\n#         plt.savefig('temporal_behavior_flow.png', dpi=150, bbox_inches='tight')\n#         plt.show()\n    \n#     def create_lab_behavior_fingerprints(self):\n#         \"\"\"Unique behavioral fingerprints for each laboratory\"\"\"\n#         print(\"\\n🔬 Creating Laboratory Behavior Fingerprints...\")\n        \n#         fig = plt.figure(figsize=(20, 16))\n        \n#         # Get top labs and behaviors\n#         lab_counts = self.annotations.group_by(\"lab_id\").count().sort(\"count\", descending=True)\n#         top_labs = lab_counts[\"lab_id\"].to_list()[:12]\n        \n#         behavior_counts = self.annotations.group_by(\"action\").count().sort(\"count\", descending=True)\n#         top_behaviors = behavior_counts[\"action\"].to_list()[:20]\n        \n#         # Create fingerprints\n#         for idx, lab in enumerate(top_labs):\n#             ax = plt.subplot(3, 4, idx + 1, projection='polar')\n            \n#             lab_data = self.annotations.filter(pl.col(\"lab_id\") == lab)\n            \n#             # Calculate behavior profile\n#             values = []\n#             for behavior in top_behaviors:\n#                 count = len(lab_data.filter(pl.col(\"action\") == behavior))\n#                 total = len(lab_data)\n#                 values.append(count / total if total > 0 else 0)\n            \n#             # Angles for radar chart\n#             angles = np.linspace(0, 2 * np.pi, len(top_behaviors), endpoint=False)\n#             values = np.array(values)\n            \n#             # Close the plot\n#             angles = np.concatenate([angles, [angles[0]]])\n#             values = np.concatenate([values, [values[0]]])\n            \n#             # Plot\n#             ax.plot(angles, values, 'o-', linewidth=2, \n#                    color=self.lab_colors[idx % len(self.lab_colors)])\n#             ax.fill(angles, values, alpha=0.25,\n#                    color=self.lab_colors[idx % len(self.lab_colors)])\n            \n#             # Set labels\n#             ax.set_xticks(angles[:-1])\n#             ax.set_xticklabels(top_behaviors, size=6)\n#             ax.set_title(f\"Lab: {lab}\", fontsize=10, fontweight='bold', pad=10)\n#             ax.set_ylim(0, max(values) * 1.1 if max(values) > 0 else 0.1)\n#             ax.grid(True, alpha=0.3)\n        \n#         plt.suptitle(\"Laboratory Behavior Fingerprints\", fontsize=14, fontweight='bold')\n#         plt.tight_layout()\n#         plt.savefig('lab_fingerprints.png', dpi=150, bbox_inches='tight')\n#         plt.show()\n    \n#     def create_behavior_embeddings(self):\n#         \"\"\"t-SNE embeddings of behaviors based on their characteristics\"\"\"\n#         print(\"\\n🗺️ Creating Behavior Embeddings Visualization...\")\n        \n#         fig, axes = plt.subplots(2, 2, figsize=(18, 16))\n        \n#         # Prepare behavior feature matrix\n#         behavior_features = []\n#         behavior_names = []\n#         behavior_sizes = []\n        \n#         for behavior in self.annotations[\"action\"].unique():\n#             b_data = self.annotations.filter(pl.col(\"action\") == behavior)\n            \n#             if len(b_data) > 20:\n#                 # Extract comprehensive features\n#                 features = [\n#                     b_data[\"duration\"].mean(),\n#                     b_data[\"duration\"].std(),\n#                     b_data[\"duration\"].median(),\n#                     b_data[\"duration\"].min(),\n#                     b_data[\"duration\"].max(),\n#                     b_data[\"duration\"].quantile(0.25),\n#                     b_data[\"duration\"].quantile(0.75),\n#                     len(b_data),\n#                     b_data[\"video_id\"].n_unique(),\n#                     b_data[\"lab_id\"].n_unique(),\n#                     b_data[\"onset\"].mean(),\n#                     b_data[\"onset\"].std(),\n#                     len(b_data) / b_data[\"video_id\"].n_unique(),  # per video frequency\n#                 ]\n                \n#                 behavior_features.append(features)\n#                 behavior_names.append(behavior)\n#                 behavior_sizes.append(len(b_data))\n        \n#         if len(behavior_features) > 10:\n#             X = StandardScaler().fit_transform(behavior_features)\n            \n#             # 1. t-SNE embedding\n#             ax1 = axes[0, 0]\n#             tsne = TSNE(n_components=2, random_state=42, perplexity=min(30, len(X)-1))\n#             X_tsne = tsne.fit_transform(X)\n            \n#             scatter = ax1.scatter(X_tsne[:, 0], X_tsne[:, 1], \n#                                  c=behavior_sizes, s=100, alpha=0.6,\n#                                  cmap='viridis', edgecolors='black', linewidth=0.5)\n            \n#             # Annotate major behaviors\n#             for i, name in enumerate(behavior_names):\n#                 if behavior_sizes[i] > np.percentile(behavior_sizes, 90):\n#                     ax1.annotate(name, (X_tsne[i, 0], X_tsne[i, 1]),\n#                                fontsize=8, alpha=0.8)\n            \n#             ax1.set_title(\"t-SNE Behavior Embeddings\", fontsize=12, fontweight='bold')\n#             ax1.set_xlabel(\"t-SNE 1\")\n#             ax1.set_ylabel(\"t-SNE 2\")\n#             plt.colorbar(scatter, ax=ax1, label='Frequency')\n            \n#             # 2. PCA embedding\n#             ax2 = axes[0, 1]\n#             pca = PCA(n_components=2)\n#             X_pca = pca.fit_transform(X)\n            \n#             scatter = ax2.scatter(X_pca[:, 0], X_pca[:, 1],\n#                                  c=behavior_sizes, s=100, alpha=0.6,\n#                                  cmap='plasma', edgecolors='black', linewidth=0.5)\n            \n#             for i, name in enumerate(behavior_names):\n#                 if behavior_sizes[i] > np.percentile(behavior_sizes, 90):\n#                     ax2.annotate(name, (X_pca[i, 0], X_pca[i, 1]),\n#                                fontsize=8, alpha=0.8)\n            \n#             ax2.set_title(f\"PCA Behavior Embeddings (Var explained: {sum(pca.explained_variance_ratio_[:2]):.1%})\",\n#                          fontsize=12, fontweight='bold')\n#             ax2.set_xlabel(f\"PC1 ({pca.explained_variance_ratio_[0]:.1%})\")\n#             ax2.set_ylabel(f\"PC2 ({pca.explained_variance_ratio_[1]:.1%})\")\n#             plt.colorbar(scatter, ax=ax2, label='Frequency')\n            \n#             # 3. Hierarchical clustering\n#             ax3 = axes[1, 0]\n            \n#             # Compute linkage\n#             linkage_matrix = hierarchy.linkage(X, method='ward')\n            \n#             # Create dendrogram\n#             dendro = hierarchy.dendrogram(linkage_matrix, \n#                                          labels=behavior_names,\n#                                          ax=ax3, \n#                                          leaf_rotation=90,\n#                                          leaf_font_size=8)\n#             ax3.set_title(\"Behavior Hierarchical Clustering\", fontsize=12, fontweight='bold')\n#             ax3.set_ylabel(\"Distance\")\n            \n#             # 4. Cluster visualization\n#             ax4 = axes[1, 1]\n            \n#             # Perform clustering\n#             n_clusters = min(5, len(X) // 5)\n#             kmeans = KMeans(n_clusters=n_clusters, random_state=42)\n#             clusters = kmeans.fit_predict(X)\n            \n#             # Plot clusters in t-SNE space\n#             scatter = ax4.scatter(X_tsne[:, 0], X_tsne[:, 1],\n#                                  c=clusters, s=100, alpha=0.6,\n#                                  cmap='Set1', edgecolors='black', linewidth=0.5)\n            \n#             # Add cluster centers\n#             centers_tsne = TSNE(n_components=2, random_state=42, perplexity=min(5, n_clusters-1)).fit_transform(kmeans.cluster_centers_)\n#             ax4.scatter(centers_tsne[:, 0], centers_tsne[:, 1],\n#                        c='red', s=500, alpha=0.5, marker='*',\n#                        edgecolors='black', linewidth=2, label='Cluster Centers')\n            \n#             ax4.set_title(f\"Behavior Clusters (k={n_clusters})\", fontsize=12, fontweight='bold')\n#             ax4.set_xlabel(\"t-SNE 1\")\n#             ax4.set_ylabel(\"t-SNE 2\")\n#             ax4.legend()\n        \n#         plt.tight_layout()\n#         plt.savefig('behavior_embeddings.png', dpi=150, bbox_inches='tight')\n#         plt.show()\n    \n#     def create_spatiotemporal_heatmaps(self):\n#         \"\"\"Advanced spatiotemporal heatmaps of behavior patterns\"\"\"\n#         print(\"\\n🗺️ Creating Spatiotemporal Behavior Heatmaps...\")\n        \n#         fig = plt.figure(figsize=(20, 12))\n        \n#         # Sample tracking data for spatial analysis\n#         spatial_data = []\n#         sampled = 0\n        \n#         for row in self.train_df.head(50).to_dicts():\n#             if sampled >= 10:\n#                 break\n                \n#             lab_id = row[\"lab_id\"]\n#             video_id = row[\"video_id\"]\n#             track_path = self.train_track_dir / lab_id / f\"{video_id}.parquet\"\n            \n#             if track_path.exists():\n#                 try:\n#                     track = pd.read_parquet(track_path)\n#                     if 'x' in track.columns and 'y' in track.columns:\n#                         # Get annotations for this video\n#                         video_annot = self.annotations.filter(pl.col(\"video_id\") == video_id)\n                        \n#                         spatial_data.append({\n#                             'tracking': track,\n#                             'annotations': video_annot,\n#                             'video_id': video_id\n#                         })\n#                         sampled += 1\n#                 except:\n#                     continue\n        \n#         if spatial_data:\n#             # 1. Aggregate spatial density for different behaviors\n#             gs = plt.GridSpec(2, 3, figure=fig, hspace=0.3, wspace=0.3)\n            \n#             # Get top behaviors\n#             top_behaviors = self.annotations.group_by(\"action\").count().sort(\"count\", descending=True)[\"action\"].to_list()[:6]\n            \n#             for idx, behavior in enumerate(top_behaviors):\n#                 ax = fig.add_subplot(gs[idx // 3, idx % 3])\n                \n#                 # Collect all positions during this behavior\n#                 behavior_positions = []\n                \n#                 for data in spatial_data:\n#                     track = data['tracking']\n#                     annot = data['annotations']\n                    \n#                     # Get behavior instances\n#                     behavior_instances = annot.filter(pl.col(\"action\") == behavior)\n                    \n#                     if 'video_frame' in track.columns:\n#                         for instance in behavior_instances.to_dicts():\n#                             # Get tracking data during behavior\n#                             mask = (track['video_frame'] >= instance['start_frame']) & \\\n#                                    (track['video_frame'] <= instance['stop_frame'])\n                            \n#                             behavior_track = track[mask]\n#                             if len(behavior_track) > 0 and 'x' in behavior_track.columns:\n#                                 behavior_positions.extend(zip(behavior_track['x'].values,\n#                                                              behavior_track['y'].values))\n                \n#                 if behavior_positions:\n#                     x_pos = [p[0] for p in behavior_positions if not np.isnan(p[0])]\n#                     y_pos = [p[1] for p in behavior_positions if not np.isnan(p[1])]\n                    \n#                     if len(x_pos) > 10:\n#                         # Create 2D histogram\n#                         h, xedges, yedges = np.histogram2d(x_pos, y_pos, bins=30)\n                        \n#                         # Apply Gaussian smoothing\n#                         from scipy.ndimage import gaussian_filter\n#                         h = gaussian_filter(h, sigma=1)\n                        \n#                         im = ax.imshow(h.T, origin='lower', cmap='hot', aspect='auto')\n#                         ax.set_title(f\"{behavior}\", fontsize=10, fontweight='bold')\n#                         ax.set_xlabel(\"X position\")\n#                         ax.set_ylabel(\"Y position\")\n#                         plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)\n        \n#         plt.suptitle(\"Spatial Distribution of Behaviors\", fontsize=14, fontweight='bold')\n#         plt.tight_layout()\n#         plt.savefig('spatiotemporal_heatmaps.png', dpi=150, bbox_inches='tight')\n#         plt.show()\n    \n#     def create_behavior_duration_landscape(self):\n#         \"\"\"3D landscape visualization of behavior durations\"\"\"\n#         print(\"\\n⛰️ Creating Behavior Duration Landscape...\")\n        \n#         fig = plt.figure(figsize=(18, 14))\n        \n#         # Create 3D surface plot of duration distributions\n#         ax1 = fig.add_subplot(221, projection='3d')\n        \n#         # Get top behaviors\n#         top_behaviors = self.annotations.group_by(\"action\").count().sort(\"count\", descending=True)[\"action\"].to_list()[:15]\n        \n#         # Create duration bins\n#         duration_bins = np.logspace(0, 3, 50)  # Log scale from 1 to 1000 frames\n        \n#         # Create 2D histogram\n#         Z = []\n#         for behavior in top_behaviors:\n#             b_data = self.annotations.filter(pl.col(\"action\") == behavior)\n#             durations = b_data[\"duration\"].to_list()\n            \n#             if durations:\n#                 hist, _ = np.histogram(durations, bins=duration_bins, density=True)\n#                 Z.append(hist)\n#             else:\n#                 Z.append(np.zeros(len(duration_bins) - 1))\n        \n#         Z = np.array(Z)\n#         X, Y = np.meshgrid(duration_bins[:-1], range(len(top_behaviors)))\n        \n#         # Create surface plot\n#         surf = ax1.plot_surface(np.log10(X), Y, Z, cmap='viridis', alpha=0.8)\n#         ax1.set_xlabel('Duration (log10 frames)')\n#         ax1.set_ylabel('Behavior')\n#         ax1.set_zlabel('Density')\n#         ax1.set_yticks(range(len(top_behaviors)))\n#         ax1.set_yticklabels(top_behaviors, fontsize=8)\n#         ax1.set_title(\"Behavior Duration Landscape\", fontsize=12, fontweight='bold')\n#         plt.colorbar(surf, ax=ax1, shrink=0.5, aspect=5)\n        \n#         # 2. Contour plot\n#         ax2 = fig.add_subplot(222)\n        \n#         contour = ax2.contourf(np.log10(X), Y, Z, levels=20, cmap='viridis')\n#         ax2.set_xlabel('Duration (log10 frames)')\n#         ax2.set_ylabel('Behavior')\n#         ax2.set_yticks(range(len(top_behaviors)))\n#         ax2.set_yticklabels(top_behaviors, fontsize=8)\n#         ax2.set_title(\"Duration Distribution Contours\", fontsize=12, fontweight='bold')\n#         plt.colorbar(contour, ax=ax2)\n        \n#         # 3. Duration statistics heatmap\n#         ax3 = fig.add_subplot(223)\n        \n#         # Calculate statistics for each behavior\n#         stats_matrix = []\n#         stat_names = ['Mean', 'Median', 'Std', 'Q25', 'Q75', 'Min', 'Max']\n        \n#         for behavior in top_behaviors:\n#             b_data = self.annotations.filter(pl.col(\"action\") == behavior)\n            \n#             if len(b_data) > 0:\n#                 stats = [\n#                     b_data[\"duration\"].mean(),\n#                     b_data[\"duration\"].median(),\n#                     b_data[\"duration\"].std(),\n#                     b_data[\"duration\"].quantile(0.25),\n#                     b_data[\"duration\"].quantile(0.75),\n#                     b_data[\"duration\"].min(),\n#                     b_data[\"duration\"].max()\n#                 ]\n#             else:\n#                 stats = [0] * len(stat_names)\n            \n#             stats_matrix.append(stats)\n        \n#         stats_matrix = np.array(stats_matrix)\n        \n#         # Normalize for visualization\n#         stats_matrix_norm = stats_matrix / stats_matrix.max(axis=0, keepdims=True)\n        \n#         im = ax3.imshow(stats_matrix_norm.T, cmap='coolwarm', aspect='auto')\n#         ax3.set_xticks(range(len(top_behaviors)))\n#         ax3.set_xticklabels(top_behaviors, rotation=45, ha='right', fontsize=8)\n#         ax3.set_yticks(range(len(stat_names)))\n#         ax3.set_yticklabels(stat_names)\n#         ax3.set_title(\"Duration Statistics Heatmap (Normalized)\", fontsize=12, fontweight='bold')\n#         plt.colorbar(im, ax=ax3)\n        \n#         # 4. Duration evolution over videos\n#         ax4 = fig.add_subplot(224)\n        \n#         # Track duration changes across videos\n#         video_ids = sorted(self.annotations[\"video_id\"].unique()[:100])\n        \n#         duration_trends = {behavior: [] for behavior in top_behaviors[:5]}\n        \n#         for vid in video_ids:\n#             vid_data = self.annotations.filter(pl.col(\"video_id\") == vid)\n            \n#             for behavior in top_behaviors[:5]:\n#                 b_data = vid_data.filter(pl.col(\"action\") == behavior)\n#                 if len(b_data) > 0:\n#                     duration_trends[behavior].append(b_data[\"duration\"].mean())\n#                 else:\n#                     duration_trends[behavior].append(np.nan)\n        \n#         for i, (behavior, durations) in enumerate(duration_trends.items()):\n#             # Interpolate NaNs\n#             durations = pd.Series(durations).interpolate().values\n            \n#             # Apply smoothing\n#             if len(durations) > 5:\n#                 from scipy.signal import savgol_filter\n#                 durations_smooth = savgol_filter(durations, \n#                                                 window_length=min(11, len(durations) if len(durations) % 2 == 1 else len(durations)-1),\n#                                                 polyorder=3)\n#             else:\n#                 durations_smooth = durations\n            \n#             ax4.plot(range(len(durations_smooth)), durations_smooth,\n#                     label=behavior, alpha=0.7, linewidth=2,\n#                     color=self.behavior_colors[i % len(self.behavior_colors)])\n        \n#         ax4.set_xlabel(\"Video Index\")\n#         ax4.set_ylabel(\"Mean Duration (frames)\")\n#         ax4.set_title(\"Duration Trends Across Videos\", fontsize=12, fontweight='bold')\n#         ax4.legend(fontsize=8)\n#         ax4.grid(True, alpha=0.3)\n        \n#         plt.tight_layout()\n#         plt.savefig('behavior_duration_landscape.png', dpi=150, bbox_inches='tight')\n#         plt.show()\n    \n#     def create_social_interaction_maps(self):\n#         \"\"\"Visualize social interaction patterns between mice\"\"\"\n#         print(\"\\n🐭 Creating Social Interaction Maps...\")\n        \n#         fig = plt.figure(figsize=(20, 12))\n        \n#         # 1. Interaction frequency matrix\n#         ax1 = fig.add_subplot(231)\n        \n#         # Get unique mouse pairs\n#         interaction_matrix = defaultdict(lambda: defaultdict(int))\n        \n#         for row in self.annotations.to_dicts():\n#             agent = row.get('agent_id', 'mouse1')\n#             target = row.get('target_id', 'mouse2')\n#             interaction_matrix[agent][target] += 1\n        \n#         # Get most active mice\n#         all_mice = set()\n#         for agent in interaction_matrix:\n#             all_mice.add(agent)\n#             all_mice.update(interaction_matrix[agent].keys())\n        \n#         mice_list = sorted(list(all_mice))[:20]  # Top 20 mice\n        \n#         # Create matrix\n#         matrix = np.zeros((len(mice_list), len(mice_list)))\n#         for i, m1 in enumerate(mice_list):\n#             for j, m2 in enumerate(mice_list):\n#                 matrix[i, j] = interaction_matrix[m1][m2]\n        \n#         im = ax1.imshow(matrix, cmap='Blues', aspect='auto')\n#         ax1.set_xticks(range(len(mice_list)))\n#         ax1.set_yticks(range(len(mice_list)))\n#         ax1.set_xticklabels(mice_list, rotation=45, ha='right', fontsize=6)\n#         ax1.set_yticklabels(mice_list, fontsize=6)\n#         ax1.set_title(\"Mouse Interaction Frequency Matrix\", fontsize=12, fontweight='bold')\n#         ax1.set_xlabel(\"Target Mouse\")\n#         ax1.set_ylabel(\"Agent Mouse\")\n#         plt.colorbar(im, ax=ax1)\n        \n#         # 2. Directed interaction network\n#         ax2 = fig.add_subplot(232)\n        \n#         # Create directed graph\n#         G = nx.DiGraph()\n        \n#         # Add edges for top interactions\n#         threshold = np.percentile(list(matrix.flatten()), 90)\n        \n#         for i, m1 in enumerate(mice_list[:10]):\n#             for j, m2 in enumerate(mice_list[:10]):\n#                 if matrix[i, j] > threshold and i != j:\n#                     G.add_edge(m1, m2, weight=matrix[i, j])\n        \n#         if G.edges():\n#             pos = nx.spring_layout(G, k=2, iterations=50)\n            \n#             # Draw nodes\n#             nx.draw_networkx_nodes(G, pos, node_color='lightblue',\n#                                   node_size=500, alpha=0.8, ax=ax2)\n            \n#             # Draw edges with varying thickness\n#             edge_weights = [G[u][v]['weight']/100 for u, v in G.edges()]\n#             nx.draw_networkx_edges(G, pos, width=edge_weights,\n#                                   alpha=0.5, edge_color='gray',\n#                                   arrows=True, arrowsize=20, ax=ax2)\n            \n#             # Draw labels\n#             nx.draw_networkx_labels(G, pos, font_size=8, ax=ax2)\n            \n#             ax2.set_title(\"Directed Social Network\", fontsize=12, fontweight='bold')\n#             ax2.axis('off')\n        \n#         # 3. Behavior-specific interactions\n#         ax3 = fig.add_subplot(233)\n        \n#         # Get interaction behaviors\n#         social_behaviors = ['chase', 'sniff', 'approach', 'attack', 'mount', 'avoid']\n#         behavior_interactions = defaultdict(int)\n        \n#         for behavior in social_behaviors:\n#             b_data = self.annotations.filter(pl.col(\"action\") == behavior)\n#             if len(b_data) > 0:\n#                 unique_pairs = b_data.select([\"agent_id\", \"target_id\"]).n_unique()\n#                 behavior_interactions[behavior] = unique_pairs\n        \n#         if behavior_interactions:\n#             behaviors = list(behavior_interactions.keys())\n#             values = list(behavior_interactions.values())\n            \n#             bars = ax3.bar(range(len(behaviors)), values,\n#                           color=[self.behavior_colors[i % len(self.behavior_colors)] for i in range(len(behaviors))])\n#             ax3.set_xticks(range(len(behaviors)))\n#             ax3.set_xticklabels(behaviors, rotation=45, ha='right')\n#             ax3.set_ylabel(\"Unique Interaction Pairs\")\n#             ax3.set_title(\"Social Behavior Diversity\", fontsize=12, fontweight='bold')\n#             ax3.grid(True, alpha=0.3, axis='y')\n        \n#         # 4. Reciprocal interactions\n#         ax4 = fig.add_subplot(234)\n        \n#         # Find reciprocal pairs\n#         reciprocal_pairs = []\n        \n#         for m1 in mice_list[:15]:\n#             for m2 in mice_list[:15]:\n#                 if m1 < m2:  # Avoid duplicates\n#                     forward = interaction_matrix[m1][m2]\n#                     backward = interaction_matrix[m2][m1]\n                    \n#                     if forward > 0 and backward > 0:\n#                         reciprocity = min(forward, backward) / max(forward, backward)\n#                         reciprocal_pairs.append({\n#                             'pair': f\"{m1}-{m2}\",\n#                             'reciprocity': reciprocity,\n#                             'total': forward + backward\n#                         })\n        \n#         if reciprocal_pairs:\n#             # Sort by reciprocity\n#             reciprocal_pairs = sorted(reciprocal_pairs, key=lambda x: x['reciprocity'], reverse=True)[:10]\n            \n#             pairs = [p['pair'] for p in reciprocal_pairs]\n#             reciprocities = [p['reciprocity'] for p in reciprocal_pairs]\n            \n#             bars = ax4.barh(range(len(pairs)), reciprocities, color='coral')\n#             ax4.set_yticks(range(len(pairs)))\n#             ax4.set_yticklabels(pairs, fontsize=8)\n#             ax4.set_xlabel(\"Reciprocity Index\")\n#             ax4.set_title(\"Top Reciprocal Interaction Pairs\", fontsize=12, fontweight='bold')\n#             ax4.set_xlim(0, 1)\n#             ax4.grid(True, alpha=0.3, axis='x')\n        \n#         # 5. Interaction duration distribution\n#         ax5 = fig.add_subplot(235)\n        \n#         # Get durations for social behaviors\n#         social_durations = {}\n        \n#         for behavior in social_behaviors:\n#             b_data = self.annotations.filter(pl.col(\"action\") == behavior)\n#             if len(b_data) > 10:\n#                 durations = b_data[\"duration\"].to_list()\n#                 social_durations[behavior] = durations\n        \n#         if social_durations:\n#             # Create violin plot\n#             data_list = []\n#             labels = []\n            \n#             for behavior, durations in social_durations.items():\n#                 data_list.append(durations[:500])  # Limit for performance\n#                 labels.append(behavior)\n            \n#             parts = ax5.violinplot(data_list, showmeans=True, showmedians=True)\n            \n#             # Color violins\n#             for i, pc in enumerate(parts['bodies']):\n#                 pc.set_facecolor(self.behavior_colors[i % len(self.behavior_colors)])\n#                 pc.set_alpha(0.6)\n            \n#             ax5.set_xticks(range(1, len(labels) + 1))\n#             ax5.set_xticklabels(labels, rotation=45, ha='right')\n#             ax5.set_ylabel(\"Duration (frames)\")\n#             ax5.set_title(\"Social Behavior Duration Distributions\", fontsize=12, fontweight='bold')\n#             ax5.grid(True, alpha=0.3, axis='y')\n        \n#         # 6. Interaction complexity score\n#         ax6 = fig.add_subplot(236)\n        \n#         # Calculate complexity for each mouse\n#         mouse_complexity = {}\n        \n#         for mouse in mice_list[:20]:\n#             # Count unique partners\n#             partners_as_agent = len(interaction_matrix[mouse])\n#             partners_as_target = sum(1 for m in interaction_matrix if mouse in interaction_matrix[m])\n            \n#             # Count behavior diversity\n#             mouse_behaviors = self.annotations.filter(\n#                 (pl.col(\"agent_id\") == mouse) | (pl.col(\"target_id\") == mouse)\n#             )\n            \n#             if len(mouse_behaviors) > 0:\n#                 behavior_diversity = mouse_behaviors[\"action\"].n_unique()\n#                 total_interactions = len(mouse_behaviors)\n                \n#                 # Complexity score\n#                 complexity = (partners_as_agent + partners_as_target) * behavior_diversity / (total_interactions + 1)\n#                 mouse_complexity[mouse] = complexity\n        \n#         if mouse_complexity:\n#             mice = list(mouse_complexity.keys())\n#             complexities = list(mouse_complexity.values())\n            \n#             # Sort by complexity\n#             sorted_indices = np.argsort(complexities)[::-1][:15]\n#             mice = [mice[i] for i in sorted_indices]\n#             complexities = [complexities[i] for i in sorted_indices]\n            \n#             bars = ax6.bar(range(len(mice)), complexities,\n#                           color='green', alpha=0.7, edgecolor='black')\n#             ax6.set_xticks(range(len(mice)))\n#             ax6.set_xticklabels(mice, rotation=45, ha='right', fontsize=8)\n#             ax6.set_ylabel(\"Complexity Score\")\n#             ax6.set_title(\"Mouse Social Complexity Ranking\", fontsize=12, fontweight='bold')\n#             ax6.grid(True, alpha=0.3, axis='y')\n        \n#         plt.suptitle(\"Social Interaction Analysis\", fontsize=14, fontweight='bold')\n#         plt.tight_layout()\n#         plt.savefig('social_interaction_maps.png', dpi=150, bbox_inches='tight')\n#         plt.show()\n    \n#     def generate_summary_report(self):\n#         \"\"\"Generate a comprehensive summary report of all analyses\"\"\"\n#         print(\"\\n📄 Generating Comprehensive Analysis Report...\")\n        \n#         report = []\n#         report.append(\"=\"*80)\n#         report.append(\"MABe ADVANCED VISUALIZATION SUITE - ANALYSIS REPORT\")\n#         report.append(\"=\"*80)\n#         report.append(f\"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\\n\")\n        \n#         # Dataset statistics\n#         report.append(\"DATASET OVERVIEW\")\n#         report.append(\"-\"*40)\n#         report.append(f\"Total videos: {len(self.train_df)}\")\n#         report.append(f\"Total annotations: {len(self.annotations) if self.annotations is not None else 0}\")\n        \n#         if self.annotations is not None:\n#             report.append(f\"Unique behaviors: {self.annotations['action'].n_unique()}\")\n#             report.append(f\"Unique laboratories: {self.annotations['lab_id'].n_unique()}\")\n#             report.append(f\"Mean behavior duration: {self.annotations['duration'].mean():.1f} frames\")\n#             report.append(f\"Duration range: {self.annotations['duration'].min()}-{self.annotations['duration'].max()} frames\")\n            \n#             # Top behaviors\n#             report.append(\"\\nTOP 10 BEHAVIORS BY FREQUENCY\")\n#             report.append(\"-\"*40)\n            \n#             top_behaviors = self.annotations.group_by(\"action\").count().sort(\"count\", descending=True).head(10)\n#             for i, row in enumerate(top_behaviors.to_dicts(), 1):\n#                 percentage = row[\"count\"] / len(self.annotations) * 100\n#                 report.append(f\"{i:2d}. {row['action']:20s} {row['count']:6d} ({percentage:5.2f}%)\")\n            \n#             # Lab statistics\n#             report.append(\"\\nLABORATORY STATISTICS\")\n#             report.append(\"-\"*40)\n            \n#             lab_stats = self.annotations.group_by(\"lab_id\").agg([\n#                 pl.count().alias(\"n_annotations\"),\n#                 pl.col(\"action\").n_unique().alias(\"n_behaviors\"),\n#                 pl.col(\"video_id\").n_unique().alias(\"n_videos\")\n#             ]).sort(\"n_annotations\", descending=True)\n            \n#             for row in lab_stats.head(10).to_dicts():\n#                 report.append(f\"{row['lab_id']:20s} - Annotations: {row['n_annotations']:5d}, \"\n#                             f\"Behaviors: {row['n_behaviors']:3d}, Videos: {row['n_videos']:3d}\")\n        \n#         # Visualization summary\n#         report.append(\"\\nGENERATED VISUALIZATIONS\")\n#         report.append(\"-\"*40)\n#         report.append(\"1. 3D Behavior Space - PCA projection of behavior features\")\n#         report.append(\"2. Behavior Co-occurrence Network - Graph of behavior relationships\")\n#         report.append(\"3. Temporal Behavior Flow - Ridge plots and circular timelines\")\n#         report.append(\"4. Laboratory Fingerprints - Unique behavioral profiles per lab\")\n#         report.append(\"5. Behavior Embeddings - t-SNE and clustering analysis\")\n#         report.append(\"6. Spatiotemporal Heatmaps - Spatial distribution of behaviors\")\n#         report.append(\"7. Duration Landscape - 3D surface plots of duration distributions\")\n#         report.append(\"8. Social Interaction Maps - Network analysis of mouse interactions\")\n        \n#         # Save report\n#         report_text = \"\\n\".join(report)\n        \n#         with open('mabe_advanced_analysis_report.txt', 'w') as f:\n#             f.write(report_text)\n        \n#         print(report_text)\n#         print(\"\\n✓ Report saved to 'mabe_advanced_analysis_report.txt'\")\n    \n#     def run_all_visualizations(self):\n#         \"\"\"Execute all visualization analyses\"\"\"\n#         print(\"\\n\" + \"=\"*80)\n#         print(\"🚀 RUNNING COMPLETE ADVANCED VISUALIZATION SUITE\")\n#         print(\"=\"*80)\n        \n#         # Load data\n#         self.load_data()\n        \n#         if self.annotations is not None and len(self.annotations) > 0:\n#             # Run all visualizations\n#             try:\n#                 self.create_3d_behavior_space()\n#             except Exception as e:\n#                 print(f\"⚠️ 3D behavior space failed: {e}\")\n            \n#             try:\n#                 self.create_behavior_cooccurrence_network()\n#             except Exception as e:\n#                 print(f\"⚠️ Co-occurrence network failed: {e}\")\n            \n#             try:\n#                 self.create_temporal_behavior_flow()\n#             except Exception as e:\n#                 print(f\"⚠️ Temporal flow failed: {e}\")\n            \n#             try:\n#                 self.create_lab_behavior_fingerprints()\n#             except Exception as e:\n#                 print(f\"⚠️ Lab fingerprints failed: {e}\")\n            \n#             try:\n#                 self.create_behavior_embeddings()\n#             except Exception as e:\n#                 print(f\"⚠️ Behavior embeddings failed: {e}\")\n            \n#             try:\n#                 self.create_spatiotemporal_heatmaps()\n#             except Exception as e:\n#                 print(f\"⚠️ Spatiotemporal heatmaps failed: {e}\")\n            \n#             try:\n#                 self.create_behavior_duration_landscape()\n#             except Exception as e:\n#                 print(f\"⚠️ Duration landscape failed: {e}\")\n            \n#             try:\n#                 self.create_social_interaction_maps()\n#             except Exception as e:\n#                 print(f\"⚠️ Social interaction maps failed: {e}\")\n            \n#             # Generate summary report\n#             self.generate_summary_report()\n            \n#             print(\"\\n\" + \"=\"*80)\n#             print(\"✅ ADVANCED VISUALIZATION SUITE COMPLETE!\")\n#             print(\"=\"*80)\n#             print(\"Generated visualizations saved as PNG files\")\n#             print(\"Analysis report saved as 'mabe_advanced_analysis_report.txt'\")\n#         else:\n#             print(\"❌ No annotations data available for visualization\")\n\n# # ========================\n# # Main execution\n# # ========================\n\n# if __name__ == \"__main__\":\n#     # Create visualization suite\n#     viz_suite = AdvancedVisualizationSuite()\n    \n#     # Run all visualizations\n#     viz_suite.run_all_visualizations()","metadata":{"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# \"\"\"\n# MABe Challenge 2025 - Starter Code with CV and Competition Metric\n# Multi-Agent Behavior Recognition in Mice\n\n# Key features:\n# 1. Proper cross-validation strategy\n# 2. Competition F-beta metric implementation\n# 3. Local CV score tracking for LB correlation\n# 4. Stratified sampling to ensure behavior diversity\n# \"\"\"\n\n# import numpy as np\n# import pandas as pd\n# import polars as pl\n# from pathlib import Path\n# import warnings\n# import gc\n# import json\n# from typing import Dict, List, Optional, Tuple, Any\n# from dataclasses import dataclass, field\n# from tqdm import tqdm\n# import pickle\n# import random\n# from scipy import signal, stats\n# from collections import defaultdict\n# from sklearn.model_selection import StratifiedKFold, GroupKFold, train_test_split\n# from sklearn.preprocessing import StandardScaler, RobustScaler\n# from sklearn.metrics import f1_score\n# import lightgbm as lgb\n# import xgboost as xgb\n\n# warnings.filterwarnings('ignore')\n\n# # ============================================================================\n# # CONFIGURATION\n# # ============================================================================\n\n# @dataclass\n# class Config:\n#     \"\"\"Configuration with CV strategy\"\"\"\n#     # Paths - CRITICAL: Use correct casing!\n#     data_path: Path = Path('/kaggle/input/MABe-mouse-behavior-detection')  # MABe not mabe!\n#     output_path: Path = Path('/kaggle/working')\n    \n#     # Data processing\n#     max_train_videos: int = 100  # Balanced for accuracy vs speed\n#     max_test_videos: Optional[int] = None\n#     sample_rate: float = 0.5  # Sample 50% of windows\n    \n#     # Feature extraction\n#     window_size: int = 45\n#     window_stride: int = 15  # Balance between coverage and speed\n#     min_window_fill: float = 0.4\n#     use_advanced_features: bool = True\n    \n#     # Cross-validation strategy\n#     cv_strategy: str = 'group'  # 'group' for video-based, 'stratified' for behavior-based\n#     n_folds: int = 5  # 5-fold CV for robust validation\n#     validation_metric: str = 'f_beta'  # Use competition metric\n#     beta: float = 1.0  # F1 score\n    \n#     # Model parameters\n#     use_ensemble: bool = True\n#     n_models: int = 2  # LightGBM + XGBoost\n#     random_state: int = 42\n#     early_stopping_rounds: int = 50\n    \n#     # Behavior classes\n#     behavior_classes: List[str] = field(default_factory=lambda: [\n#         'grooming', 'sniff', 'chase', 'attack', \n#         'mount', 'investigate', 'escape', 'approach'\n#     ])\n    \n#     # Thresholds (will be optimized during CV)\n#     confidence_thresholds: Dict[str, float] = field(default_factory=lambda: {\n#         'grooming': 0.3,\n#         'sniff': 0.25,\n#         'chase': 0.35,\n#         'attack': 0.4,\n#         'mount': 0.35,\n#         'investigate': 0.25,\n#         'escape': 0.35,\n#         'approach': 0.3\n#     })\n    \n#     # Processing\n#     batch_size: int = 100\n#     max_features: int = 100  # Reduced to prevent overfitting\n\n# config = Config()\n\n# # ============================================================================\n# # COMPETITION METRIC IMPLEMENTATION\n# # ============================================================================\n\n# class MABeMetric:\n#     \"\"\"Competition F-beta metric implementation\"\"\"\n    \n#     @staticmethod\n#     def prepare_solution_data(annotations: pd.DataFrame, video_id: str, lab_id: str = 'lab1') -> pd.DataFrame:\n#         \"\"\"Convert annotations to solution format\"\"\"\n#         if annotations.empty:\n#             return pd.DataFrame()\n        \n#         solution_rows = []\n#         for _, ann in annotations.iterrows():\n#             # Create solution row\n#             row = {\n#                 'video_id': video_id,\n#                 'agent_id': ann.get('agent_id', 'mouse1'),\n#                 'target_id': ann.get('target_id', 'mouse2'),\n#                 'action': ann.get('action', 'unknown'),\n#                 'start_frame': int(ann['start_frame']),\n#                 'stop_frame': int(ann.get('end_frame', ann.get('stop_frame', ann['start_frame'] + 30))),\n#                 'lab_id': lab_id\n#             }\n#             solution_rows.append(row)\n        \n#         if solution_rows:\n#             solution_df = pd.DataFrame(solution_rows)\n            \n#             # Add behaviors_labeled column (required by metric)\n#             unique_behaviors = solution_df.apply(\n#                 lambda x: f\"{x['agent_id']},{x['target_id']},{x['action']}\", axis=1\n#             ).unique().tolist()\n#             solution_df['behaviors_labeled'] = json.dumps(unique_behaviors)\n            \n#             return solution_df\n        \n#         return pd.DataFrame()\n    \n#     @staticmethod\n#     def single_lab_f1(lab_solution: pl.DataFrame, lab_submission: pl.DataFrame, beta: float = 1) -> float:\n#         \"\"\"Calculate F1 for a single lab\"\"\"\n#         if lab_solution.is_empty() or lab_submission.is_empty():\n#             return 0.0\n        \n#         label_frames = defaultdict(set)\n#         prediction_frames = defaultdict(set)\n        \n#         # Build label frames\n#         for row in lab_solution.to_dicts():\n#             key = f\"{row['video_id']}_{row['agent_id']}_{row['target_id']}_{row['action']}\"\n#             label_frames[key].update(range(row['start_frame'], row['stop_frame']))\n        \n#         # Build prediction frames\n#         for row in lab_submission.to_dicts():\n#             key = f\"{row['video_id']}_{row['agent_id']}_{row['target_id']}_{row['action']}\"\n#             prediction_frames[key].update(range(row['start_frame'], row['stop_frame']))\n        \n#         # Calculate metrics per action\n#         tps = defaultdict(int)\n#         fns = defaultdict(int)\n#         fps = defaultdict(int)\n        \n#         all_keys = set(label_frames.keys()) | set(prediction_frames.keys())\n#         actions = set()\n        \n#         for key in all_keys:\n#             action = key.split('_')[-1]\n#             actions.add(action)\n            \n#             label_set = label_frames.get(key, set())\n#             pred_set = prediction_frames.get(key, set())\n            \n#             tps[action] += len(label_set & pred_set)\n#             fns[action] += len(label_set - pred_set)\n#             fps[action] += len(pred_set - label_set)\n        \n#         # Calculate F-beta per action\n#         action_scores = []\n#         for action in actions:\n#             tp = tps[action]\n#             fn = fns[action]\n#             fp = fps[action]\n            \n#             if tp + fn + fp == 0:\n#                 continue\n            \n#             precision = tp / (tp + fp) if (tp + fp) > 0 else 0\n#             recall = tp / (tp + fn) if (tp + fn) > 0 else 0\n            \n#             if precision + recall > 0:\n#                 f_score = (1 + beta**2) * (precision * recall) / (beta**2 * precision + recall)\n#             else:\n#                 f_score = 0\n            \n#             action_scores.append(f_score)\n        \n#         return np.mean(action_scores) if action_scores else 0.0\n    \n#     @staticmethod\n#     def calculate_metric(solution: pd.DataFrame, submission: pd.DataFrame, beta: float = 1) -> float:\n#         \"\"\"Calculate competition F-beta metric\"\"\"\n#         if solution.empty or submission.empty:\n#             return 0.0\n        \n#         # Convert to polars for efficiency\n#         solution_pl = pl.DataFrame(solution)\n#         submission_pl = pl.DataFrame(submission)\n        \n#         # Calculate per lab\n#         lab_scores = []\n#         for lab_id in solution['lab_id'].unique():\n#             lab_solution = solution_pl.filter(pl.col('lab_id') == lab_id)\n#             lab_videos = lab_solution['video_id'].unique().to_list()\n#             lab_submission = submission_pl.filter(pl.col('video_id').is_in(lab_videos))\n            \n#             score = MABeMetric.single_lab_f1(lab_solution, lab_submission, beta)\n#             lab_scores.append(score)\n        \n#         return np.mean(lab_scores) if lab_scores else 0.0\n\n# # ============================================================================\n# # FEATURE EXTRACTION\n# # ============================================================================\n\n# class FeatureExtractor:\n#     \"\"\"Feature extraction with focus on reducing overfitting\"\"\"\n    \n#     def __init__(self, window_size: int = 45, stride: int = 15):\n#         self.window_size = window_size\n#         self.stride = stride\n    \n#     def extract_movement_features(self, trajectory: pd.DataFrame) -> Dict[str, float]:\n#         \"\"\"Extract core movement features\"\"\"\n#         features = {}\n        \n#         if 'x' not in trajectory.columns or len(trajectory) < 3:\n#             return features\n        \n#         x = trajectory['x'].values\n#         y = trajectory['y'].values\n        \n#         # Velocity\n#         dx = np.diff(x)\n#         dy = np.diff(y)\n#         velocity = np.sqrt(dx**2 + dy**2)\n        \n#         # Core statistics\n#         features['vel_mean'] = np.mean(velocity)\n#         features['vel_std'] = np.std(velocity)\n#         features['vel_max'] = np.max(velocity)\n#         features['vel_median'] = np.median(velocity)\n#         features['vel_q25'] = np.percentile(velocity, 25)\n#         features['vel_q75'] = np.percentile(velocity, 75)\n        \n#         # Acceleration\n#         if len(velocity) > 1:\n#             acceleration = np.diff(velocity)\n#             features['acc_mean'] = np.mean(np.abs(acceleration))\n#             features['acc_std'] = np.std(acceleration)\n        \n#         # Path metrics\n#         features['total_dist'] = np.sum(velocity)\n#         features['net_disp'] = np.sqrt((x[-1] - x[0])**2 + (y[-1] - y[0])**2)\n#         features['path_efficiency'] = features['net_disp'] / (features['total_dist'] + 1e-6)\n        \n#         # Spatial\n#         features['x_range'] = np.max(x) - np.min(x)\n#         features['y_range'] = np.max(y) - np.min(y)\n#         features['area'] = features['x_range'] * features['y_range']\n        \n#         # Angular features\n#         if len(dx) > 1:\n#             angles = np.arctan2(dy, dx)\n#             angle_changes = np.diff(angles)\n#             angle_changes = np.arctan2(np.sin(angle_changes), np.cos(angle_changes))\n            \n#             features['turn_mean'] = np.mean(np.abs(angle_changes))\n#             features['turn_std'] = np.std(angle_changes)\n#             features['turn_max'] = np.max(np.abs(angle_changes))\n        \n#         # Motion patterns\n#         features['stationary_ratio'] = np.mean(velocity < 2.0)\n#         features['slow_ratio'] = np.mean((velocity >= 2.0) & (velocity < 10.0))\n#         features['fast_ratio'] = np.mean(velocity >= 10.0)\n        \n#         return features\n    \n#     def extract_temporal_features(self, trajectory: pd.DataFrame) -> Dict[str, float]:\n#         \"\"\"Extract temporal dynamics\"\"\"\n#         features = {}\n        \n#         if 'x' not in trajectory.columns or len(trajectory) < 10:\n#             return features\n        \n#         x = trajectory['x'].values\n#         y = trajectory['y'].values\n        \n#         # Autocorrelation\n#         if len(x) > 15:\n#             x_centered = x - np.mean(x)\n#             y_centered = y - np.mean(y)\n            \n#             # Lag-1 autocorrelation\n#             features['x_autocorr'] = np.corrcoef(x_centered[:-1], x_centered[1:])[0, 1]\n#             features['y_autocorr'] = np.corrcoef(y_centered[:-1], y_centered[1:])[0, 1]\n        \n#         # Trend\n#         time = np.arange(len(x))\n#         x_trend = np.polyfit(time, x, 1)[0]\n#         y_trend = np.polyfit(time, y, 1)[0]\n#         features['x_trend'] = x_trend\n#         features['y_trend'] = y_trend\n        \n#         return features\n    \n#     def extract_windows(self, data: pd.DataFrame, video_id: str, \n#                        sample_rate: float = 1.0) -> pd.DataFrame:\n#         \"\"\"Extract feature windows from tracking data\"\"\"\n#         if data.empty:\n#             return pd.DataFrame()\n        \n#         # Get trajectory\n#         if 'frame' not in data.columns:\n#             data['frame'] = data.index\n        \n#         trajectory = data.groupby('frame').agg({\n#             'x': 'mean',\n#             'y': 'mean'\n#         }).reset_index()\n        \n#         if len(trajectory) < self.window_size:\n#             return pd.DataFrame()\n        \n#         all_features = []\n#         min_frame = trajectory['frame'].min()\n#         max_frame = trajectory['frame'].max()\n        \n#         # Generate windows\n#         window_starts = list(range(int(min_frame), \n#                                   int(max_frame) - self.window_size + 1, \n#                                   self.stride))\n        \n#         # Sample if needed\n#         if sample_rate < 1.0:\n#             n_samples = max(1, int(len(window_starts) * sample_rate))\n#             window_starts = random.sample(window_starts, min(n_samples, len(window_starts)))\n        \n#         for start_frame in window_starts:\n#             end_frame = start_frame + self.window_size\n            \n#             window = trajectory[(trajectory['frame'] >= start_frame) & \n#                                (trajectory['frame'] < end_frame)]\n            \n#             if len(window) < self.window_size * config.min_window_fill:\n#                 continue\n            \n#             # Extract features\n#             features = {\n#                 'video_id': video_id,\n#                 'start_frame': start_frame,\n#                 'end_frame': end_frame,\n#             }\n            \n#             # Movement features\n#             move_feats = self.extract_movement_features(window)\n#             features.update(move_feats)\n            \n#             # Temporal features\n#             if config.use_advanced_features:\n#                 temp_feats = self.extract_temporal_features(window)\n#                 features.update(temp_feats)\n            \n#             all_features.append(features)\n        \n#         return pd.DataFrame(all_features)\n\n# # ============================================================================\n# # CROSS-VALIDATION STRATEGY\n# # ============================================================================\n\n# class CVStrategy:\n#     \"\"\"Cross-validation with competition metric\"\"\"\n    \n#     def __init__(self, n_folds: int = 5, strategy: str = 'group'):\n#         self.n_folds = n_folds\n#         self.strategy = strategy\n#         self.cv_scores = []\n#         self.best_thresholds = {}\n    \n#     def create_folds(self, features_df: pd.DataFrame, labels: np.ndarray) -> List[Tuple]:\n#         \"\"\"Create CV folds based on strategy\"\"\"\n#         if self.strategy == 'group':\n#             # Group by video to prevent leakage\n#             groups = features_df['video_id'].values\n#             gkf = GroupKFold(n_splits=self.n_folds)\n#             folds = list(gkf.split(features_df, labels, groups))\n#         else:\n#             # Stratified by behavior presence\n#             skf = StratifiedKFold(n_splits=self.n_folds, shuffle=True, random_state=42)\n#             # Create binary labels for stratification\n#             binary_labels = (labels > 0).astype(int) if labels.ndim > 1 else labels\n#             folds = list(skf.split(features_df, binary_labels))\n        \n#         return folds\n    \n#     def optimize_thresholds(self, y_true: np.ndarray, y_pred: np.ndarray, \n#                            features_df: pd.DataFrame) -> Dict[str, float]:\n#         \"\"\"Optimize thresholds using validation data\"\"\"\n#         best_thresholds = {}\n        \n#         # Simple threshold optimization\n#         for threshold in np.arange(0.1, 0.6, 0.05):\n#             # Convert predictions to submission format\n#             predictions = []\n#             for idx, prob in enumerate(y_pred):\n#                 if prob > threshold:\n#                     predictions.append({\n#                         'video_id': features_df.iloc[idx]['video_id'],\n#                         'agent_id': 'mouse1',\n#                         'target_id': 'mouse2',\n#                         'action': 'unknown',\n#                         'start_frame': int(features_df.iloc[idx]['start_frame']),\n#                         'stop_frame': int(features_df.iloc[idx]['end_frame'])\n#                     })\n            \n#             if predictions:\n#                 # Calculate score (simplified)\n#                 pred_binary = (y_pred > threshold).astype(int)\n#                 score = f1_score(y_true, pred_binary)\n                \n#                 if 'best_score' not in locals() or score > best_score:\n#                     best_score = score\n#                     best_threshold = threshold\n        \n#         return {'global': best_threshold if 'best_threshold' in locals() else 0.3}\n\n# # ============================================================================\n# # MODEL WITH CV\n# # ============================================================================\n\n# class BehaviorDetectorCV:\n#     \"\"\"Model with cross-validation\"\"\"\n    \n#     def __init__(self):\n#         self.models = []\n#         self.scalers = []\n#         self.feature_columns = None\n#         self.cv_strategy = CVStrategy(n_folds=config.n_folds, strategy=config.cv_strategy)\n#         self.cv_scores = []\n#         self.feature_importance = None\n    \n#     def prepare_features(self, features_df: pd.DataFrame) -> np.ndarray:\n#         \"\"\"Prepare features\"\"\"\n#         numeric_cols = features_df.select_dtypes(include=[np.number]).columns\n#         exclude = ['start_frame', 'end_frame']\n#         self.feature_columns = [c for c in numeric_cols if c not in exclude]\n        \n#         if not self.feature_columns:\n#             return np.zeros((len(features_df), 1))\n        \n#         X = features_df[self.feature_columns].fillna(0).values\n        \n#         # Remove constant features\n#         std = np.std(X, axis=0)\n#         valid_features = std > 1e-10\n#         X = X[:, valid_features]\n#         self.feature_columns = [col for col, valid in zip(self.feature_columns, valid_features) if valid]\n        \n#         return X\n    \n#     def train_with_cv(self, features_df: pd.DataFrame, labels: np.ndarray, \n#                      annotations_df: pd.DataFrame = None):\n#         \"\"\"Train with cross-validation\"\"\"\n#         print(\"\\n\" + \"=\"*60)\n#         print(\"TRAINING WITH CROSS-VALIDATION\")\n#         print(\"=\"*60)\n        \n#         X = self.prepare_features(features_df)\n        \n#         if X.shape[0] < 100 or X.shape[1] == 0:\n#             print(\"Insufficient data for training\")\n#             return\n        \n#         # Binary classification for simplicity\n#         y = (labels > 0).astype(int) if labels.ndim > 1 else labels\n        \n#         print(f\"Data shape: {X.shape}\")\n#         print(f\"Positive rate: {100*np.mean(y):.1f}%\")\n        \n#         # Create CV folds\n#         folds = self.cv_strategy.create_folds(features_df, y)\n#         print(f\"Created {len(folds)} CV folds\")\n        \n#         # Track feature importance\n#         feature_importance_sum = np.zeros(X.shape[1])\n        \n#         # Train on each fold\n#         for fold_idx, (train_idx, val_idx) in enumerate(folds):\n#             print(f\"\\nFold {fold_idx + 1}/{len(folds)}\")\n            \n#             X_train, X_val = X[train_idx], X[val_idx]\n#             y_train, y_val = y[train_idx], y[val_idx]\n            \n#             # Check for minimum positive samples\n#             if np.sum(y_train) < 10 or np.sum(y_val) < 5:\n#                 print(f\"  Skipping fold - insufficient positive samples\")\n#                 continue\n            \n#             # Scale\n#             scaler = RobustScaler()\n#             X_train_scaled = scaler.fit_transform(X_train)\n#             X_val_scaled = scaler.transform(X_val)\n            \n#             # Train LightGBM\n#             lgb_model = lgb.LGBMClassifier(\n#                 n_estimators=200,\n#                 max_depth=6,\n#                 num_leaves=40,\n#                 learning_rate=0.05,\n#                 min_child_samples=30,\n#                 subsample=0.7,\n#                 colsample_bytree=0.7,\n#                 reg_alpha=0.3,\n#                 reg_lambda=0.3,\n#                 class_weight='balanced' if np.mean(y_train) < 0.3 else None,\n#                 random_state=config.random_state + fold_idx,\n#                 verbosity=-1\n#             )\n            \n#             lgb_model.fit(\n#                 X_train_scaled, y_train,\n#                 eval_set=[(X_val_scaled, y_val)],\n#                 callbacks=[lgb.early_stopping(config.early_stopping_rounds), \n#                           lgb.log_evaluation(0)]\n#             )\n            \n#             # Track feature importance\n#             if hasattr(lgb_model, 'feature_importances_'):\n#                 feature_importance_sum += lgb_model.feature_importances_\n            \n#             # Validate\n#             y_pred_proba = lgb_model.predict_proba(X_val_scaled)[:, 1]\n#             y_pred = (y_pred_proba > 0.3).astype(int)\n            \n#             # Calculate metrics\n#             val_f1 = f1_score(y_val, y_pred)\n#             print(f\"  Validation F1: {val_f1:.3f}\")\n            \n#             # If we have annotations, calculate competition metric\n#             if annotations_df is not None:\n#                 val_features = features_df.iloc[val_idx]\n#                 val_predictions = []\n                \n#                 for idx, prob in enumerate(y_pred_proba):\n#                     if prob > 0.3:\n#                         val_predictions.append({\n#                             'video_id': val_features.iloc[idx]['video_id'],\n#                             'agent_id': 'mouse1',\n#                             'target_id': 'mouse2',\n#                             'action': 'unknown',\n#                             'start_frame': int(val_features.iloc[idx]['start_frame']),\n#                             'stop_frame': int(val_features.iloc[idx]['end_frame'])\n#                         })\n                \n#                 if val_predictions:\n#                     submission_df = pd.DataFrame(val_predictions)\n                    \n#                     # Get corresponding annotations\n#                     val_videos = val_features['video_id'].unique()\n#                     val_annotations = annotations_df[annotations_df['video_id'].isin(val_videos)]\n                    \n#                     if not val_annotations.empty:\n#                         # Convert to solution format\n#                         solution_df = MABeMetric.prepare_solution_data(\n#                             val_annotations, \n#                             val_videos[0] if len(val_videos) > 0 else 'unknown'\n#                         )\n                        \n#                         if not solution_df.empty:\n#                             competition_score = MABeMetric.calculate_metric(\n#                                 solution_df, submission_df, beta=config.beta\n#                             )\n#                             print(f\"  Competition F-beta: {competition_score:.3f}\")\n#                             self.cv_scores.append(competition_score)\n            \n#             # Store model\n#             self.models.append(lgb_model)\n#             self.scalers.append(scaler)\n            \n#             # Train XGBoost if ensemble enabled\n#             if config.use_ensemble and config.n_models > 1:\n#                 xgb_model = xgb.XGBClassifier(\n#                     n_estimators=150,\n#                     max_depth=5,\n#                     learning_rate=0.05,\n#                     subsample=0.7,\n#                     colsample_bytree=0.7,\n#                     reg_alpha=0.3,\n#                     reg_lambda=0.3,\n#                     random_state=config.random_state + fold_idx + 100,\n#                     use_label_encoder=False,\n#                     eval_metric='logloss'\n#                 )\n                \n#                 xgb_model.fit(\n#                     X_train_scaled, y_train,\n#                     eval_set=[(X_val_scaled, y_val)],\n#                     early_stopping_rounds=config.early_stopping_rounds,\n#                     verbose=False\n#                 )\n                \n#                 self.models.append(xgb_model)\n        \n#         # Store average feature importance\n#         if feature_importance_sum.any():\n#             self.feature_importance = feature_importance_sum / len(folds)\n            \n#             # Print top features\n#             top_features_idx = np.argsort(self.feature_importance)[-10:][::-1]\n#             print(\"\\nTop 10 important features:\")\n#             for idx in top_features_idx:\n#                 if idx < len(self.feature_columns):\n#                     print(f\"  {self.feature_columns[idx]}: {self.feature_importance[idx]:.2f}\")\n        \n#         # Print CV summary\n#         if self.cv_scores:\n#             print(f\"\\nCV Summary:\")\n#             print(f\"  Mean F-beta: {np.mean(self.cv_scores):.3f} (+/- {np.std(self.cv_scores):.3f})\")\n#             print(f\"  Min: {np.min(self.cv_scores):.3f}, Max: {np.max(self.cv_scores):.3f}\")\n    \n#     def predict(self, X: np.ndarray) -> np.ndarray:\n#         \"\"\"Predict using ensemble\"\"\"\n#         if not self.models:\n#             return np.zeros(X.shape[0])\n        \n#         predictions = []\n        \n#         for i, model in enumerate(self.models):\n#             scaler_idx = min(i // (config.n_models if config.use_ensemble else 1), len(self.scalers) - 1)\n#             X_scaled = self.scalers[scaler_idx].transform(X)\n            \n#             if hasattr(model, 'predict_proba'):\n#                 pred = model.predict_proba(X_scaled)[:, 1]\n#             else:\n#                 pred = model.predict(X_scaled)\n            \n#             predictions.append(pred)\n        \n#         # Average predictions\n#         return np.mean(predictions, axis=0)\n\n# # ============================================================================\n# # MAIN PIPELINE WITH CV\n# # ============================================================================\n\n# class MABePipelineCV:\n#     \"\"\"Main pipeline with proper CV\"\"\"\n    \n#     def __init__(self):\n#         self.feature_extractor = FeatureExtractor(\n#             window_size=config.window_size,\n#             stride=config.window_stride\n#         )\n#         self.detector = BehaviorDetectorCV()\n#         self.all_annotations = []\n    \n#     def load_and_process_file(self, tracking_path: Path, annotation_path: Path = None) -> Tuple[pd.DataFrame, np.ndarray, pd.DataFrame]:\n#         \"\"\"Load and process a single file\"\"\"\n#         try:\n#             # Load tracking\n#             data = pd.read_parquet(tracking_path)\n#             video_id = tracking_path.stem\n            \n#             # Extract features\n#             features = self.feature_extractor.extract_windows(\n#                 data, video_id, sample_rate=config.sample_rate\n#             )\n            \n#             if features.empty:\n#                 return pd.DataFrame(), np.array([]), pd.DataFrame()\n            \n#             # Load annotations if available\n#             annotations = pd.DataFrame()\n#             labels = np.zeros(len(features))\n            \n#             if annotation_path and annotation_path.exists():\n#                 annotations = pd.read_parquet(annotation_path)\n                \n#                 # Map column names\n#                 col_mapping = {\n#                     'start': 'start_frame',\n#                     'end': 'end_frame',\n#                     'stop': 'end_frame',\n#                     'stop_frame': 'end_frame'\n#                 }\n                \n#                 for old_col, new_col in col_mapping.items():\n#                     if old_col in annotations.columns and new_col not in annotations.columns:\n#                         annotations[new_col] = annotations[old_col]\n                \n#                 # Create labels\n#                 if 'start_frame' in annotations.columns and 'end_frame' in annotations.columns:\n#                     for idx, window in features.iterrows():\n#                         overlaps = ((annotations['start_frame'] <= window['end_frame']) & \n#                                   (annotations['end_frame'] >= window['start_frame']))\n#                         labels[idx] = 1 if overlaps.any() else 0\n                    \n#                     # Add video_id to annotations\n#                     annotations['video_id'] = video_id\n            \n#             return features, labels, annotations\n            \n#         except Exception as e:\n#             print(f\"Error processing {tracking_path.name}: {e}\")\n#             return pd.DataFrame(), np.array([]), pd.DataFrame()\n    \n#     def process_training_data(self):\n#         \"\"\"Process training data with CV in mind\"\"\"\n#         print(\"\\n\" + \"=\"*60)\n#         print(\"PROCESSING TRAINING DATA FOR CV\")\n#         print(\"=\"*60)\n        \n#         # Find files\n#         tracking_files = []\n#         train_path = config.data_path / 'train_tracking'\n        \n#         if train_path.exists():\n#             for lab_dir in train_path.iterdir():\n#                 if lab_dir.is_dir():\n#                     files = list(lab_dir.glob('*.parquet'))\n#                     tracking_files.extend(files)\n        \n#         # Sample files\n#         if config.max_train_videos and len(tracking_files) > config.max_train_videos:\n#             # Stratified sampling by lab\n#             lab_files = defaultdict(list)\n#             for f in tracking_files:\n#                 lab_files[f.parent.name].append(f)\n            \n#             sampled_files = []\n#             files_per_lab = config.max_train_videos // len(lab_files)\n            \n#             for lab, files in lab_files.items():\n#                 n_sample = min(files_per_lab, len(files))\n#                 sampled_files.extend(random.sample(files, n_sample))\n            \n#             tracking_files = sampled_files[:config.max_train_videos]\n        \n#         print(f\"Processing {len(tracking_files)} files\")\n        \n#         all_features = []\n#         all_labels = []\n#         all_annotations = []\n        \n#         # Process files\n#         for i in range(0, len(tracking_files), config.batch_size):\n#             batch = tracking_files[i:i+config.batch_size]\n#             print(f\"\\nBatch {i//config.batch_size + 1}/{(len(tracking_files)-1)//config.batch_size + 1}\")\n            \n#             for file_path in tqdm(batch, desc=\"Processing\"):\n#                 # Get annotation path\n#                 ann_path = config.data_path / 'train_annotation' / file_path.parent.name / f\"{file_path.stem}.parquet\"\n                \n#                 # Process file\n#                 features, labels, annotations = self.load_and_process_file(file_path, ann_path)\n                \n#                 if not features.empty and len(labels) > 0:\n#                     all_features.append(features)\n#                     all_labels.append(labels)\n#                     if not annotations.empty:\n#                         all_annotations.append(annotations)\n            \n#             gc.collect()\n        \n#         # Combine\n#         if all_features:\n#             features_df = pd.concat(all_features, ignore_index=True)\n#             labels = np.concatenate(all_labels)\n            \n#             annotations_df = pd.concat(all_annotations, ignore_index=True) if all_annotations else pd.DataFrame()\n            \n#             print(f\"\\nTotal samples: {len(features_df)}\")\n#             print(f\"Positive rate: {100*np.mean(labels):.1f}%\")\n#             print(f\"Number of videos: {features_df['video_id'].nunique()}\")\n            \n#             return features_df, labels, annotations_df\n        \n#         return None, None, None\n    \n#     def train_model(self, features_df: pd.DataFrame, labels: np.ndarray, annotations_df: pd.DataFrame = None):\n#         \"\"\"Train model with CV\"\"\"\n#         self.detector.train_with_cv(features_df, labels, annotations_df)\n    \n#     def generate_predictions(self) -> pd.DataFrame:\n#         \"\"\"Generate test predictions\"\"\"\n#         print(\"\\n\" + \"=\"*60)\n#         print(\"GENERATING PREDICTIONS\")\n#         print(\"=\"*60)\n        \n#         test_path = config.data_path / 'test_tracking'\n#         test_files = []\n        \n#         if test_path.exists():\n#             for lab_dir in test_path.iterdir():\n#                 if lab_dir.is_dir():\n#                     test_files.extend(list(lab_dir.glob('*.parquet')))\n        \n#         print(f\"Found {len(test_files)} test files\")\n        \n#         predictions = []\n        \n#         for test_file in tqdm(test_files, desc=\"Predicting\"):\n#             # Process file\n#             features, _, _ = self.load_and_process_file(test_file)\n            \n#             if features.empty:\n#                 continue\n            \n#             # Prepare features\n#             X = self.detector.prepare_features(features)\n            \n#             # Predict\n#             proba = self.detector.predict(X)\n            \n#             # Use optimized threshold or default\n#             threshold = config.confidence_thresholds.get('global', 0.3)\n            \n#             # Convert to predictions\n#             for idx, prob in enumerate(proba):\n#                 if prob > threshold:\n#                     predictions.append({\n#                         'video_id': features.iloc[idx]['video_id'],\n#                         'agent_id': 'mouse1',\n#                         'target_id': 'mouse2',\n#                         'action': 'sniff',  # Default action\n#                         'start_frame': int(features.iloc[idx]['start_frame']),\n#                         'stop_frame': int(features.iloc[idx]['end_frame'])\n#                     })\n        \n#         if predictions:\n#             return pd.DataFrame(predictions)\n#         else:\n#             # Minimal valid submission\n#             return pd.DataFrame({\n#                 'video_id': ['test_video'],\n#                 'agent_id': ['mouse1'],\n#                 'target_id': ['mouse2'],\n#                 'action': ['grooming'],\n#                 'start_frame': [0],\n#                 'stop_frame': [30]\n#             })\n    \n#     def create_submission(self, predictions: pd.DataFrame) -> pd.DataFrame:\n#         \"\"\"Create submission file\"\"\"\n#         predictions['row_id'] = range(len(predictions))\n#         submission = predictions[['row_id', 'video_id', 'agent_id', 'target_id', \n#                                  'action', 'start_frame', 'stop_frame']]\n        \n#         submission.to_csv('submission.csv', index=False)\n#         print(f\"\\nCreated submission.csv with {len(submission)} predictions\")\n#         return submission\n\n# # ============================================================================\n# # MAIN EXECUTION\n# # ============================================================================\n\n# def main():\n#     \"\"\"Main execution with CV\"\"\"\n#     import time\n#     start_time = time.time()\n    \n#     print(\"\\n\" + \"=\"*80)\n#     print(\"MABe CHALLENGE 2025 - WITH PROPER CV\")\n#     print(\"=\"*80)\n    \n#     print(f\"\\nConfiguration:\")\n#     print(f\"  CV Strategy: {config.cv_strategy}\")\n#     print(f\"  CV Folds: {config.n_folds}\")\n#     print(f\"  Max training videos: {config.max_train_videos}\")\n#     print(f\"  Ensemble models: {config.n_models}\")\n    \n#     # Initialize pipeline\n#     pipeline = MABePipelineCV()\n    \n#     # Process training data\n#     features_df, labels, annotations_df = pipeline.process_training_data()\n    \n#     if features_df is not None and len(features_df) > 500:\n#         # Train with CV\n#         pipeline.train_model(features_df, labels, annotations_df)\n        \n#         # Print CV results\n#         if pipeline.detector.cv_scores:\n#             print(\"\\n\" + \"=\"*60)\n#             print(\"CROSS-VALIDATION RESULTS\")\n#             print(\"=\"*60)\n#             print(f\"Mean CV F-beta: {np.mean(pipeline.detector.cv_scores):.4f}\")\n#             print(f\"Std CV F-beta: {np.std(pipeline.detector.cv_scores):.4f}\")\n#             print(f\"Individual fold scores: {pipeline.detector.cv_scores}\")\n#             print(\"\\nExpected LB correlation:\")\n#             print(f\"  Optimistic (mean + std): {np.mean(pipeline.detector.cv_scores) + np.std(pipeline.detector.cv_scores):.4f}\")\n#             print(f\"  Expected (mean): {np.mean(pipeline.detector.cv_scores):.4f}\")\n#             print(f\"  Conservative (mean - std): {np.mean(pipeline.detector.cv_scores) - np.std(pipeline.detector.cv_scores):.4f}\")\n        \n#         # Save model\n#         with open('cv_model.pkl', 'wb') as f:\n#             pickle.dump(pipeline.detector, f)\n#         print(\"\\nModel saved to cv_model.pkl\")\n        \n#         # Generate predictions\n#         predictions = pipeline.generate_predictions()\n        \n#         # Create submission\n#         submission = pipeline.create_submission(predictions)\n        \n#         # Report runtime\n#         elapsed = time.time() - start_time\n#         print(f\"\\nTotal runtime: {elapsed/3600:.2f} hours\")\n        \n#         print(\"\\n\" + \"=\"*80)\n#         print(\"PIPELINE COMPLETE!\")\n#         print(\"=\"*80)\n        \n#         return submission\n#     else:\n#         print(\"\\nInsufficient training data!\")\n#         return None\n\n# if __name__ == \"__main__\":\n#     submission = main()\n#     if submission is not None:\n#         print(\"\\n✓ Execution completed successfully!\")\n#         print(f\"Submission shape: {submission.shape}\")","metadata":{"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null}]}