data.py

#!/usr/bin/env python3

import pandas as pd
import numpy as np
import os
import glob
import gc
import re
import argparse
from tqdm import tqdm
from preprocess import safe_pickle_load, rotate_coordinates
from sklearn.preprocessing import StandardScaler
import joblib


def load_and_split_data(data_root, processed_path, train_week_threshold=15, use_memmap=True,
                        sumer_coverage_path=None, extend_y_target=False):
    def ensure_scalar(value):
        if isinstance(value, pd.Series):
            return value.iloc[0] if len(value) > 0 else None
        elif hasattr(value, 'item'):
            try:
                return value.item()
            except (ValueError, AttributeError):
                return value
        return value
    
    load_mode = 'r' if use_memmap else None
    X_Players = np.load(os.path.join(processed_path, "X_Players_exp15.npy"), mmap_mode=load_mode)
    X_Static = np.load(os.path.join(processed_path, "X_Static_exp15.npy"), mmap_mode=load_mode)
    Y_Target = np.load(os.path.join(processed_path, "Y_Target_exp15.npy"), mmap_mode=load_mode)
    
    if X_Players.shape[1] != 22 or X_Players.shape[2] != 12:
        raise ValueError(f"X_Players 维度错误: {X_Players.shape}")
    if X_Static.shape[1] != 22 or X_Static.shape[2] < 18:
        raise ValueError(f"X_Static 维度错误: {X_Static.shape}")
    if Y_Target.shape[1] != 60:
        raise ValueError(f"Y_Target 维度错误: {Y_Target.shape}")
    
    meta_files = [
        os.path.join(processed_path, "frames_metadata_temp.pkl"),
        os.path.join(processed_path, "_cache_df_frames.pkl"),
        os.path.join(processed_path, "_checkpoint_frame_data_exp15.pkl")
    ]
    
    df_meta = None
    for meta_file in meta_files:
        if os.path.exists(meta_file):
            try:
                loaded_data, success = safe_pickle_load(meta_file, default=None)
                if success and isinstance(loaded_data, pd.DataFrame):
                    df_meta = loaded_data
                    break
            except Exception:
                continue
    
    if df_meta is None:
        raise FileNotFoundError(f"元数据文件未找到: {meta_files}")
    
    required_cols = ['game_id', 'play_id', 'frame_id']
    missing_cols = [col for col in required_cols if col not in df_meta.columns]
    if missing_cols:
        raise ValueError(f"Metadata missing required columns: {missing_cols}")
    
    cols_to_keep = required_cols.copy()
    if 'nfl_ids' in df_meta.columns:
        cols_to_keep.append('nfl_ids')
    df_meta = df_meta[cols_to_keep].copy()
    
    if len(df_meta) != len(X_Players):
        raise ValueError(
            f"元数据长度 ({len(df_meta)}) 与数组长度 ({len(X_Players)}) 不匹配！"
        )
    
    week_map = {}
    input_files = sorted(glob.glob(os.path.join(data_root, 'input_*.csv')))
    if not input_files:
        raise FileNotFoundError(f"No input_*.csv in {data_root}")
    
    encodings = ['utf-8', 'latin-1', 'iso-8859-1', 'cp1252', 'gbk']
    
    for f in input_files:
        match = re.search(r'w(\d+)', os.path.basename(f))
        if match:
            week = int(match.group(1))
            df_temp = None
            for encoding in encodings:
                try:
                    df_temp = pd.read_csv(f, usecols=['game_id', 'play_id'], low_memory=False, encoding=encoding)
                    break
                except (UnicodeDecodeError, UnicodeError):
                    continue
                except Exception:
                    break
            
            if df_temp is None:
                continue
            
            df_temp = df_temp.drop_duplicates()
            for _, row in df_temp.iterrows():
                game_id = ensure_scalar(row['game_id'])
                play_id = ensure_scalar(row['play_id'])
                key = (game_id, play_id)
                if key not in week_map:
                    week_map[key] = week
    
    def get_week(row, wm=week_map):
        game_id = ensure_scalar(row['game_id'])
        play_id = ensure_scalar(row['play_id'])
        return wm.get((game_id, play_id), None)
    df_meta['week'] = df_meta.apply(get_week, axis=1)
    
    missing_weeks = df_meta['week'].isna().sum()
    if missing_weeks > 0:
        if missing_weeks > len(df_meta) * 0.1:
            raise ValueError(f"缺失week的帧占比 {missing_weeks/len(df_meta)*100:.1f}%，数据不完整")
        df_meta['week'] = df_meta['week'].fillna(0).astype(int)
    else:
        df_meta['week'] = df_meta['week'].astype(int)
    
    train_mask = (df_meta['week'] <= train_week_threshold).values
    val_mask = (df_meta['week'] > train_week_threshold).values
    
    train_count = train_mask.sum()
    val_count = val_mask.sum()
    
    if train_count == 0:
        raise ValueError("训练集为空")
    if val_count == 0:
        raise ValueError("验证集为空")
    
    X_P_train = np.array(X_Players[train_mask], dtype=np.float32)
    X_S_train = np.array(X_Static[train_mask], dtype=np.float32)
    Y_train = np.array(Y_Target[train_mask], dtype=np.float32)
    
    X_P_val = np.array(X_Players[val_mask], dtype=np.float32)
    X_S_val = np.array(X_Static[val_mask], dtype=np.float32)
    Y_val = np.array(Y_Target[val_mask], dtype=np.float32)
    
    has_coverage = bool(sumer_coverage_path and os.path.exists(sumer_coverage_path))
    cache_suffix = f"_w{train_week_threshold}_coverage{int(has_coverage)}"
    cache_dir = os.path.join(processed_path, "_cache_expanded_data")
    os.makedirs(cache_dir, exist_ok=True)
    
    cache_files = {
        'X_P_train': os.path.join(cache_dir, f"X_P_train_expanded{cache_suffix}.npy"),
        'X_S_train': os.path.join(cache_dir, f"X_S_train_expanded{cache_suffix}.npy"),
        'Y_train': os.path.join(cache_dir, f"Y_train_expanded{cache_suffix}.npy"),
        'X_P_val': os.path.join(cache_dir, f"X_P_val_expanded{cache_suffix}.npy"),
        'X_S_val': os.path.join(cache_dir, f"X_S_val_expanded{cache_suffix}.npy"),
        'Y_val': os.path.join(cache_dir, f"Y_val_expanded{cache_suffix}.npy"),
        'player_indices_train': os.path.join(cache_dir, f"player_indices_train_expanded{cache_suffix}.npy"),
        'player_indices_val': os.path.join(cache_dir, f"player_indices_val_expanded{cache_suffix}.npy"),
    }
    
    cache_exists = all(os.path.exists(f) for f in cache_files.values())
    
    if cache_exists:
        try:
            X_P_train = np.load(cache_files['X_P_train'], mmap_mode=load_mode)
            X_S_train = np.load(cache_files['X_S_train'], mmap_mode=load_mode)
            Y_train = np.load(cache_files['Y_train'], mmap_mode=load_mode)
            X_P_val = np.load(cache_files['X_P_val'], mmap_mode=load_mode)
            X_S_val = np.load(cache_files['X_S_val'], mmap_mode=load_mode)
            Y_val = np.load(cache_files['Y_val'], mmap_mode=load_mode)
            player_indices_train = np.load(cache_files['player_indices_train'], mmap_mode=load_mode)
            player_indices_val = np.load(cache_files['player_indices_val'], mmap_mode=load_mode)
            
            if use_memmap:
                X_P_train = np.array(X_P_train, dtype=np.float32)
                X_S_train = np.array(X_S_train, dtype=np.float32)
                Y_train = np.array(Y_train, dtype=np.float32)
                X_P_val = np.array(X_P_val, dtype=np.float32)
                X_S_val = np.array(X_S_val, dtype=np.float32)
                Y_val = np.array(Y_val, dtype=np.float32)
                player_indices_train = np.array(player_indices_train, dtype=np.int32)
                player_indices_val = np.array(player_indices_val, dtype=np.int32)
            
            skip_expansion = True
        except Exception:
            skip_expansion = False
    else:
        skip_expansion = False
    
    if not skip_expansion:
        X_NFL_ids = None
        
        if 'nfl_ids' in df_meta.columns:
            nfl_ids_list = df_meta['nfl_ids'].tolist()
            X_NFL_ids = np.array([np.array(ids, dtype=np.int32) if isinstance(ids, (list, np.ndarray)) else np.zeros(22, dtype=np.int32) for ids in nfl_ids_list], dtype=np.int32)
        else:
            nfl_ids_file_candidates = [
                os.path.join(processed_path, "frames_data_temp_nfl_ids.npy"),
                os.path.join(processed_path, "_checkpoint_frame_data_exp15_nfl_ids.npy"),
                os.path.join(processed_path, "frames_metadata_temp_nfl_ids.npy"),
                os.path.join(processed_path, "_nfl_ids.npy"),
            ]
            for nfl_ids_file in nfl_ids_file_candidates:
                if os.path.exists(nfl_ids_file):
                    try:
                        X_NFL_ids = np.load(nfl_ids_file, mmap_mode=load_mode)
                        break
                    except Exception:
                        continue
            
            if X_NFL_ids is None:
                raise FileNotFoundError(f"未找到 nfl_ids: {nfl_ids_file_candidates}")
        
        if X_NFL_ids.shape != (len(X_Players), 22):
            raise ValueError(f"nfl_ids 形状不匹配！期望 ({len(X_Players)}, 22)，实际 {X_NFL_ids.shape}")
        
        output_files = sorted(glob.glob(os.path.join(data_root, 'output_*.csv')))
        if not output_files:
            raise FileNotFoundError(f"未找到 output_*.csv: {data_root}")
        
        encodings = ['utf-8', 'latin-1', 'iso-8859-1', 'cp1252', 'gbk']
        all_outputs = []
        for f in tqdm(output_files, desc="加载Output文件"):
            df_temp = None
            for encoding in encodings:
                try:
                    df_temp = pd.read_csv(f, usecols=['game_id', 'play_id', 'nfl_id', 'frame_id', 'x', 'y'], 
                                         low_memory=False, encoding=encoding)
                    break
                except (UnicodeDecodeError, UnicodeError):
                    continue
                except Exception:
                    break
            if df_temp is not None:
                all_outputs.append(df_temp)
        
        if not all_outputs:
            raise FileNotFoundError("无法加载 output CSV 文件")
        
        df_output = pd.concat(all_outputs, ignore_index=True)
        
        input_files = sorted(glob.glob(os.path.join(data_root, 'input_*.csv')))
        if not input_files:
            raise FileNotFoundError(f"未找到 input_*.csv: {data_root}")
        
        all_inputs = []
        for f in tqdm(input_files[:5], desc="加载Input文件"):
            df_temp = None
            for encoding in encodings:
                try:
                    df_temp = pd.read_csv(f, usecols=['game_id', 'play_id', 'play_direction'], 
                                         low_memory=False, encoding=encoding)
                    break
                except (UnicodeDecodeError, UnicodeError):
                    continue
                except Exception:
                    break
            if df_temp is not None:
                all_inputs.append(df_temp)
        
        if not all_inputs:
            raise FileNotFoundError("无法加载 input CSV 文件")
        
        df_input_sample = pd.concat(all_inputs, ignore_index=True)
        play_direction_map = df_input_sample[['game_id', 'play_id', 'play_direction']].drop_duplicates(subset=['game_id', 'play_id'])
        del df_input_sample, all_inputs
        gc.collect()
        
        df_output = pd.merge(df_output, play_direction_map, on=['game_id', 'play_id'], how='left')
        df_output = rotate_coordinates(df_output)
        del play_direction_map
        gc.collect()
        
        output_dict = {}
        for (game_id, play_id, nfl_id), group in tqdm(df_output.groupby(['game_id', 'play_id', 'nfl_id']), 
                                                       desc="构建输出索引"):
            sorted_group = group.sort_values('frame_id')
            frame_coords = {}
            for _, row in sorted_group.iterrows():
                frame_coords[int(row['frame_id'])] = (float(row['x']), float(row['y']))
            output_dict[(game_id, play_id, nfl_id)] = frame_coords
        
        del df_output, all_outputs
        gc.collect()
        
        def build_all_players_targets(X_P, X_S, Y, mask, nfl_ids_subset, df_meta_subset, output_dict_ref, dataset_name):
            n_frames = len(X_P)
            expanded_X_P = []
            expanded_X_S = []
            expanded_Y = []
            expanded_player_indices = []
            
            matched_count = 0
            total_players = 0
            
            for frame_idx in tqdm(range(n_frames), desc=f"扩展{dataset_name}集"):
                frame_nfl_ids = nfl_ids_subset[frame_idx]
                meta_row = df_meta_subset.iloc[frame_idx]
                game_id = ensure_scalar(meta_row['game_id'])
                play_id = ensure_scalar(meta_row['play_id'])
                
                frame_X_P = X_P[frame_idx]
                frame_X_S = X_S[frame_idx]
                
                for player_idx in range(22):
                    nfl_id = int(frame_nfl_ids[player_idx])
                    total_players += 1
                    
                    key = (game_id, play_id, nfl_id)
                    current_frame_id = ensure_scalar(meta_row['frame_id'])
                    
                    if key in output_dict_ref:
                        frame_coords = output_dict_ref[key]
                        future_30 = np.zeros(60, dtype=np.float32)
                        found_frames = 0
                        for offset in range(30):
                            target_frame_id = current_frame_id + offset + 1
                            if target_frame_id in frame_coords:
                                x, y = frame_coords[target_frame_id]
                                future_30[offset*2] = x
                                future_30[offset*2 + 1] = y
                                found_frames += 1
                        if found_frames > 0:
                            matched_count += 1
                    else:
                        future_30 = np.zeros(60, dtype=np.float32)
                    
                    expanded_X_P.append(frame_X_P)
                    expanded_X_S.append(frame_X_S)
                    expanded_Y.append(future_30)
                    expanded_player_indices.append(player_idx)
            
            return (np.array(expanded_X_P, dtype=np.float32),
                    np.array(expanded_X_S, dtype=np.float32),
                    np.array(expanded_Y, dtype=np.float32),
                    np.array(expanded_player_indices, dtype=np.int32))
    
        X_P_train_expanded, X_S_train_expanded, Y_train_expanded, player_indices_train = build_all_players_targets(
            X_P_train, X_S_train, Y_train, train_mask, 
            X_NFL_ids[train_mask], df_meta[train_mask], output_dict, "训练"
        )
        
        X_P_val_expanded, X_S_val_expanded, Y_val_expanded, player_indices_val = build_all_players_targets(
            X_P_val, X_S_val, Y_val, val_mask,
            X_NFL_ids[val_mask], df_meta[val_mask], output_dict, "验证"
        )
        
        X_P_train = X_P_train_expanded
        X_S_train = X_S_train_expanded
        Y_train = Y_train_expanded
        X_P_val = X_P_val_expanded
        X_S_val = X_S_val_expanded
        Y_val = Y_val_expanded
        
        try:
            np.save(cache_files['X_P_train'], X_P_train)
            np.save(cache_files['X_S_train'], X_S_train)
            np.save(cache_files['Y_train'], Y_train)
            np.save(cache_files['X_P_val'], X_P_val)
            np.save(cache_files['X_S_val'], X_S_val)
            np.save(cache_files['Y_val'], Y_val)
            np.save(cache_files['player_indices_train'], player_indices_train)
            np.save(cache_files['player_indices_val'], player_indices_val)
        except Exception:
            pass
        
        del X_NFL_ids, output_dict, X_P_train_expanded, X_S_train_expanded, Y_train_expanded
        del X_P_val_expanded, X_S_val_expanded, Y_val_expanded
        gc.collect()
    
    if sumer_coverage_path and os.path.exists(sumer_coverage_path):
        try:
            df_sumer = pd.read_parquet(sumer_coverage_path)
            
            coverage_col = None
            potential_cols = [c for c in df_sumer.columns if 'coverage' in c.lower() or 'cover' in c.lower()]
            if 'coverage_scheme' in df_sumer.columns:
                coverage_col = 'coverage_scheme'
            elif 'defense_coverage_type' in df_sumer.columns:
                coverage_col = 'defense_coverage_type'
            elif potential_cols:
                coverage_col = potential_cols[0]
            else:
                coverage_col = None
            
            if coverage_col is not None:
                coverage_map = df_sumer[['game_id', 'play_id', coverage_col]].drop_duplicates(subset=['game_id', 'play_id']).set_index(['game_id', 'play_id'])[coverage_col]
                
                all_coverages = sorted(df_sumer[coverage_col].dropna().unique())
                coverage_to_idx = {cov: i for i, cov in enumerate(all_coverages)}
                num_coverage_classes = len(all_coverages)
                
                def generate_coverage_features(mask):
                    subset_meta = df_meta[mask].copy()
                    n_samples = len(subset_meta)
                    cov_features = np.zeros((n_samples, 22, num_coverage_classes), dtype=np.float32)
                    
                    mask_indices = np.where(mask)[0]
                    for local_idx, global_idx in enumerate(mask_indices):
                        row = df_meta.iloc[global_idx]
                        game_id = ensure_scalar(row['game_id'])
                        play_id = ensure_scalar(row['play_id'])
                        key = (game_id, play_id)
                        if key in coverage_map.index:
                            coverage_type = coverage_map.loc[key]
                            if pd.notna(coverage_type) and coverage_type in coverage_to_idx:
                                label_idx = coverage_to_idx[coverage_type]
                                cov_features[local_idx, :, label_idx] = 1.0
                    
                    return cov_features
                
                cov_train = generate_coverage_features(train_mask)
                X_S_train = np.concatenate([X_S_train, cov_train], axis=2)
                
                cov_val = generate_coverage_features(val_mask)
                X_S_val = np.concatenate([X_S_val, cov_val], axis=2)
                
                del cov_train, cov_val, df_sumer, coverage_map
                gc.collect()
        except Exception:
            pass
    
    if 'X_NFL_ids' not in locals() or X_NFL_ids is None:
        if 'nfl_ids' in df_meta.columns:
            nfl_ids_list = df_meta['nfl_ids'].tolist()
            X_NFL_ids = np.array([np.array(ids, dtype=np.int32) if isinstance(ids, (list, np.ndarray)) else np.zeros(22, dtype=np.int32) for ids in nfl_ids_list], dtype=np.int32)
        else:
            nfl_ids_file_candidates = [
                os.path.join(processed_path, "frames_data_temp_nfl_ids.npy"),
                os.path.join(processed_path, "_checkpoint_frame_data_exp15_nfl_ids.npy"),
                os.path.join(processed_path, "frames_metadata_temp_nfl_ids.npy"),
                os.path.join(processed_path, "_nfl_ids.npy"),
            ]
            for nfl_ids_file in nfl_ids_file_candidates:
                if os.path.exists(nfl_ids_file):
                    try:
                        X_NFL_ids = np.load(nfl_ids_file, mmap_mode=load_mode)
                        if use_memmap:
                            X_NFL_ids = np.array(X_NFL_ids, dtype=np.int32)
                        break
                    except Exception:
                        continue
            if 'X_NFL_ids' not in locals() or X_NFL_ids is None:
                raise FileNotFoundError("未找到 nfl_ids")
        
        if 'output_dict' not in locals() or output_dict is None:
            output_files = sorted(glob.glob(os.path.join(data_root, 'output_*.csv')))
            if not output_files:
                raise FileNotFoundError(f"未找到 output_*.csv: {data_root}")
            
            encodings = ['utf-8', 'latin-1', 'iso-8859-1', 'cp1252', 'gbk']
            all_outputs = []
            for f in tqdm(output_files, desc="加载Output文件"):
                df_temp = None
                for encoding in encodings:
                    try:
                        df_temp = pd.read_csv(f, usecols=['game_id', 'play_id', 'nfl_id', 'frame_id', 'x', 'y'], 
                                             low_memory=False, encoding=encoding)
                        break
                    except (UnicodeDecodeError, UnicodeError):
                        continue
                    except Exception:
                        break
                if df_temp is not None:
                    all_outputs.append(df_temp)
            
            if not all_outputs:
                raise FileNotFoundError("无法加载 output CSV 文件")
            
            df_output = pd.concat(all_outputs, ignore_index=True)
            
            input_files = sorted(glob.glob(os.path.join(data_root, 'input_*.csv')))
            if not input_files:
                raise FileNotFoundError(f"未找到 input_*.csv: {data_root}")
            
            all_inputs = []
            for f in tqdm(input_files[:5], desc="加载Input文件"):
                df_temp = None
                for encoding in encodings:
                    try:
                        df_temp = pd.read_csv(f, usecols=['game_id', 'play_id', 'play_direction'], 
                                             low_memory=False, encoding=encoding)
                        break
                    except (UnicodeDecodeError, UnicodeError):
                        continue
                    except Exception:
                        break
                if df_temp is not None:
                    all_inputs.append(df_temp)
            
            if not all_inputs:
                raise FileNotFoundError("无法加载 input CSV 文件")
            
            df_input_sample = pd.concat(all_inputs, ignore_index=True)
            play_direction_map = df_input_sample[['game_id', 'play_id', 'play_direction']].drop_duplicates(subset=['game_id', 'play_id'])
            del df_input_sample, all_inputs
            gc.collect()
            
            df_output = pd.merge(df_output, play_direction_map, on=['game_id', 'play_id'], how='left')
            df_output = rotate_coordinates(df_output)
            del play_direction_map
            gc.collect()
            
            output_dict = {}
            for (game_id, play_id, nfl_id), group in tqdm(df_output.groupby(['game_id', 'play_id', 'nfl_id']), 
                                                           desc="构建输出索引"):
                sorted_group = group.sort_values('frame_id')
                frame_coords = {}
                for _, row in sorted_group.iterrows():
                    frame_coords[int(row['frame_id'])] = (float(row['x']), float(row['y']))
                output_dict[(game_id, play_id, nfl_id)] = frame_coords
            
            del df_output, all_outputs
            gc.collect()
        
        receivers_dict = {}
        for i in range(len(df_meta)):
            game_id = ensure_scalar(df_meta.iloc[i]['game_id'])
            play_id = ensure_scalar(df_meta.iloc[i]['play_id'])
            key = (game_id, play_id)
            if key not in receivers_dict:
                target_receiver_nfl_id = int(X_NFL_ids[i, 0])
                receivers_dict[key] = target_receiver_nfl_id
        
        is_expanded = len(Y_train) > len(df_meta[train_mask])
        if not is_expanded:
            # 修复训练集（原始数据，未扩展）
            Y_train_fixed = np.zeros_like(Y_train, dtype=np.float32)
            train_indices = np.where(train_mask)[0]
            matched_train = 0
            for local_idx, global_idx in enumerate(tqdm(train_indices, desc="修复训练集 Y_Target")):
                row = df_meta.iloc[global_idx]
                game_id = ensure_scalar(row['game_id'])
                play_id = ensure_scalar(row['play_id'])
                frame_id = int(row['frame_id'])
                
                nfl_id = receivers_dict.get((game_id, play_id))
                if nfl_id is not None:
                    key = (game_id, play_id, nfl_id)
                    if key in output_dict:
                        frame_coords = output_dict[key]
                        future_30 = np.zeros(60, dtype=np.float32)
                        found_frames = 0
                        for offset in range(30):
                            target_frame_id = frame_id + offset + 1  # 未来帧
                            if target_frame_id in frame_coords:
                                x, y = frame_coords[target_frame_id]
                                future_30[offset*2] = x
                                future_30[offset*2+1] = y
                                found_frames += 1
                        if found_frames > 0:
                            Y_train_fixed[local_idx] = future_30
                            matched_train += 1
            
            # 修复验证集（原始数据，未扩展）
            Y_val_fixed = np.zeros_like(Y_val, dtype=np.float32)
            val_indices = np.where(val_mask)[0]
            matched_val = 0
            for local_idx, global_idx in enumerate(tqdm(val_indices, desc="修复验证集 Y_Target")):
                row = df_meta.iloc[global_idx]
                game_id = ensure_scalar(row['game_id'])
                play_id = ensure_scalar(row['play_id'])
                frame_id = int(row['frame_id'])
                
                nfl_id = receivers_dict.get((game_id, play_id))
                if nfl_id is not None:
                    key = (game_id, play_id, nfl_id)
                    if key in output_dict:
                        frame_coords = output_dict[key]
                        future_30 = np.zeros(60, dtype=np.float32)
                        found_frames = 0
                        for offset in range(30):
                            target_frame_id = frame_id + offset + 1  # 未来帧
                            if target_frame_id in frame_coords:
                                x, y = frame_coords[target_frame_id]
                                future_30[offset*2] = x
                                future_30[offset*2+1] = y
                                found_frames += 1
                        if found_frames > 0:
                            Y_val_fixed[local_idx] = future_30
                            matched_val += 1
            
            Y_train = Y_train_fixed
            Y_val = Y_val_fixed
            
            # 清理内存
            del Y_train_fixed, Y_val_fixed
        
        # 清理内存
        del receivers_dict
        if 'output_dict' in locals() and output_dict is not None:
            # 如果扩展数据时还需要使用 output_dict，暂时不删除
            pass
        gc.collect()
    
    
    players_scaler_path = os.path.join(processed_path, "players_scaler_exp15.joblib")
    if not os.path.exists(players_scaler_path):
        raise FileNotFoundError(f"未找到 players_scaler: {players_scaler_path}")
    players_scaler = joblib.load(players_scaler_path)
    
    
    try:
        is_expanded_for_delta = player_indices_train is not None and len(player_indices_train) > 0
    except NameError:
        is_expanded_for_delta = False
    if is_expanded_for_delta:
        n_train = len(X_P_train)
        n_val = len(X_P_val)
        P0_train_scaled = np.zeros((n_train, 2), dtype=np.float32)
        P0_val_scaled = np.zeros((n_val, 2), dtype=np.float32)
        for i in range(n_train):
            P0_train_scaled[i] = X_P_train[i, player_indices_train[i], 0:2]
        for i in range(n_val):
            P0_val_scaled[i] = X_P_val[i, player_indices_val[i], 0:2]
        
    else:
        P0_train_scaled = X_P_train[:, 0, 0:2].copy()
        P0_val_scaled = X_P_val[:, 0, 0:2].copy()
    
    P0_train = P0_train_scaled * players_scaler.scale_[:2] + players_scaler.mean_[:2]
    P0_val = P0_val_scaled * players_scaler.scale_[:2] + players_scaler.mean_[:2]
    P0_train = P0_train.reshape(-1, 1, 2)
    P0_val = P0_val.reshape(-1, 1, 2)
    
    Y_train_reshaped = Y_train.reshape(-1, 30, 2)
    Y_val_reshaped = Y_val.reshape(-1, 30, 2)
    
    zero_mask_train = (Y_train_reshaped[:, 0, 0] == 0.0) & (Y_train_reshaped[:, 0, 1] == 0.0)
    zero_mask_val = (Y_val_reshaped[:, 0, 0] == 0.0) & (Y_val_reshaped[:, 0, 1] == 0.0)
    
    
    def to_pure_delta(P0_batch, Y_seq_batch):
        full_seq = np.concatenate([P0_batch, Y_seq_batch], axis=1)
        deltas = full_seq[:, 1:, :] - full_seq[:, :-1, :]
        return deltas
    
    
    Y_train = to_pure_delta(P0_train, Y_train_reshaped).reshape(-1, 60)
    Y_val = to_pure_delta(P0_val, Y_val_reshaped).reshape(-1, 60)
    
    
    
    del X_Players, X_Static, week_map, P0_train, P0_val
    if 'df_meta' in locals():
        del df_meta
    gc.collect()
    
    y_scaler = StandardScaler()
    nonzero_mask_train_scaler = ~zero_mask_train
    
    if nonzero_mask_train_scaler.sum() > 0:
        y_scaler.fit(Y_train[nonzero_mask_train_scaler])
    else:
        y_scaler.fit(Y_train)
    
    y_scaler_path = os.path.join(processed_path, "y_target_scaler_exp15.joblib")
    joblib.dump(y_scaler, y_scaler_path)
    
    Y_train_scaled = y_scaler.transform(Y_train)
    Y_val_scaled = y_scaler.transform(Y_val)
    
    return X_P_train, X_S_train, Y_train_scaled, X_P_val, X_S_val, Y_val_scaled, player_indices_train, player_indices_val


def generate_decoder_input(Y_Target):
    """生成解码器输入（Teacher Forcing）"""
    Y_Target_reshaped = Y_Target.reshape(-1, 30, 2)
    Y_Decoder_Input = np.zeros_like(Y_Target_reshaped, dtype=np.float32)
    Y_Decoder_Input[:, 0, :] = [0.0, 0.0]
    Y_Decoder_Input[:, 1:, :] = Y_Target_reshaped[:, :-1, :]
    
    return Y_Decoder_Input


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='测试数据加载工具')
    parser.add_argument('--data_root', type=str, required=True,
                        help='原始 CSV 数据路径')
    parser.add_argument('--processed_path', type=str, required=True,
                        help='预处理后的数据路径')
    parser.add_argument('--train_week_threshold', type=int, default=15,
                        help='训练集 week 阈值')
    
    args = parser.parse_args()
    
    try:
        X_P_train, X_S_train, Y_train, X_P_val, X_S_val, Y_val = load_and_split_data(
            args.data_root, args.processed_path, args.train_week_threshold
        )
        print("\n 数据加载测试成功！")
    except Exception as e:
        print(f"\n 数据加载测试失败: {e}")
        import traceback
        traceback.print_exc()