fix this python code:

import json
from typing import Dict, List, Tuple

import numpy as np
import torch


class BaseballFeaturizer:
    def __init__(self):
        # Define feature dimensions for consistent tensor shapes
        self.game_state_dim = 13  # Corrected count: 1+2+1+1+1+1+1+1+1+2+1+1 = 13
        self.context_dim = 51  # Corrected count: 9+6+20+16 = 51
        self.full_feature_dim = self.game_state_dim + self.context_dim  # 13 + 51 = 64

    def extract_game_state_tensor(self, state: Dict) -> torch.Tensor:
        """
        Extract the core game state that will be predicted by the model.
        This is the 'target' for next-state prediction.
        """
        features = [
            state["inning"],  # 1 feature
            *state["inning_topbot_onehot"],  # 2 features [1,0] or [0,1]
            state["outs"],  # 1 feature
            state["on_1b"],  # 1 feature (0/1)
            state["on_2b"],  # 1 feature (0/1)
            state["on_3b"],  # 1 feature (0/1)
            state["home_score"],  # 1 feature
            state["away_score"],  # 1 feature
            state["score_diff"],  # 1 feature
            *state["batting_team_onehot"],  # 2 features [1,0] or [0,1]
            state["batter_id"],  # 1 feature - WHO is batting (key for prediction!)
            state["pitcher_id"],  # 1 feature - WHO is pitching (key for prediction!)
            # Note: catcher_id could also be included if you want to model defensive effects
            # Model will learn embeddings for these player IDs
            # Batting order + lineup context helps model understand player relationships
        ]

        # Debug: Check actual length (only print once)
        if len(features) != self.game_state_dim:
            if not hasattr(self, "_game_state_debug_printed"):
                print(
                    f"WARNING: Game state features length {len(features)} != expected {self.game_state_dim}"
                )
                print(
                    f"inning_topbot_onehot length: {len(state['inning_topbot_onehot'])}"
                )
                print(
                    f"batting_team_onehot length: {len(state['batting_team_onehot'])}"
                )
                self._game_state_debug_printed = True

        return torch.tensor(features, dtype=torch.float32)

    def extract_context_tensor(
        self, context: Dict, batter_stats: Dict, pitcher_stats: Dict, game_context: Dict
    ) -> torch.Tensor:
        """
        Extract context features that help prediction but don't get predicted.
        """
        features = []

        # PA Context (9 features)
        pa_context_features = [
            context["batting_order_position"],  # 1 feature
            context["pitcher_pitch_count"],  # 1 feature
            context["times_through_order"],  # 1 feature
            context["is_starter"],  # 1 feature (0/1)
            *context["pitcher_hand_onehot"],  # 2 features
            *context["batter_hand_onehot"],  # 3 features
        ]
        features.extend(pa_context_features)

        # Game Context (6 features)
        game_context_features = [
            game_context["game_year"] / 2024.0,  # Normalized year
            game_context["game_month"] / 12.0,  # Normalized month
            game_context["game_day"] / 31.0,  # Normalized day
            *game_context["roof_status_onehot"][:3],  # 3 features
        ]
        features.extend(game_context_features)

        # Batter Stats (20 features)
        batter_career = batter_stats["rate_stats"]["overall"]
        batter_features = [
            batter_career["ba"],  # Batting average
            batter_career["obp"],  # On-base percentage
            batter_career["slg"],  # Slugging percentage
            batter_career["ops"],  # OPS
            batter_career["k_pct"],  # Strikeout rate
            batter_career["bb_pct"],  # Walk rate
            batter_career["hr_pct"],  # Home run rate
            batter_career["gb_pct"],  # Ground ball rate
            batter_career["fb_pct"],  # Fly ball rate
            batter_career["ld_pct"],  # Line drive rate
            batter_career["barrel_pct"],  # Barrel rate
            batter_career["hard_hit_pct"],  # Hard hit rate
            batter_career["swing_pct"],  # Swing rate
            batter_career["contact_pct"],  # Contact rate
            batter_career["whiff_pct"],  # Whiff rate
            batter_stats["time_stats"]["season"]["overall"]["pa"] / 700.0,  # Season PAs
            batter_stats["time_stats"]["last15"]["overall"]["pa"] / 75.0,  # Recent PAs
            batter_stats["time_stats"]["last30"]["overall"]["pa"] / 150.0,  # Recent PAs
            0.0,  # Padding
            0.0,  # Padding
        ]
        features.extend(batter_features)

        # Pitcher Stats (16 features)
        pitcher_career = pitcher_stats["rate_stats"]["overall"]
        pitcher_features = [
            pitcher_career["era"],  # ERA
            pitcher_career["fip"],  # FIP
            pitcher_career["whip"],  # WHIP
            pitcher_career["k_pct"],  # Strikeout rate
            pitcher_career["bb_pct"],  # Walk rate
            pitcher_career["hr_pct"],  # Home run rate
            pitcher_career["gb_pct"],  # Ground ball rate
            pitcher_career["fb_pct"],  # Fly ball rate
            pitcher_career["ld_pct"],  # Line drive rate
            pitcher_career["barrel_pct_allowed"],  # Barrel rate allowed
            pitcher_career["hard_hit_pct_allowed"],  # Hard hit rate allowed
            pitcher_stats["time_stats"]["season"]["overall"]["pa"]
            / 800.0,  # Season PAs
            pitcher_stats["time_stats"]["last15"]["overall"]["pa"] / 90.0,  # Recent PAs
            pitcher_stats["time_stats"]["last30"]["overall"]["pa"]
            / 180.0,  # Recent PAs
            0.0,  # Padding
            0.0,  # Padding
        ]
        features.extend(pitcher_features)

        # Debug feature counts on first call only
        if not hasattr(self, "_debug_printed"):
            print(f"PA context features: {len(pa_context_features)}")
            print(f"Game context features: {len(game_context_features)}")
            print(f"Batter features: {len(batter_features)}")
            print(f"Pitcher features: {len(pitcher_features)}")
            print(f"Total context features: {len(features)}")
            self._debug_printed = True

        # Total: 9 + 6 + 20 + 16 = 51 features
        # Ensure exact dimension
        if len(features) != self.context_dim:
            print(
                f"WARNING: Context features length {len(features)} != expected {self.context_dim}"
            )
            # Pad or truncate to exact context_dim
            if len(features) > self.context_dim:
                features = features[: self.context_dim]
            else:
                features.extend([0.0] * (self.context_dim - len(features)))

        return torch.tensor(features, dtype=torch.float32)

    def create_sequence_tensors(
        self, game_data: Dict
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Create input/target tensors for the entire game sequence.

        Returns:
            input_features: [seq_len, full_feature_dim] - game_state + context
            target_states: [seq_len, game_state_dim] - next game states
            game_states: [seq_len, game_state_dim] - current game states
        """
        pa_sequence = game_data["pa_sequence"]
        batter_stats_map = dict(game_data["batter_stats_map"])
        pitcher_stats_map = dict(game_data["pitcher_stats_map"])
        game_context = game_data["game_context"]

        seq_len = len(pa_sequence)

        # Initialize tensors
        input_features = torch.zeros(seq_len, self.full_feature_dim)
        target_states = torch.zeros(seq_len, self.game_state_dim)
        game_states = torch.zeros(seq_len, self.game_state_dim)

        for i, pa in enumerate(pa_sequence):
            state = pa["state"]
            context = pa["context"]

            # Get player stats
            batter_id = state["batter_id"]
            pitcher_id = state["pitcher_id"]
            batter_stats = batter_stats_map[batter_id]
            pitcher_stats = pitcher_stats_map[pitcher_id]

            # Extract features
            game_state_tensor = self.extract_game_state_tensor(state)
            context_tensor = self.extract_context_tensor(
                context, batter_stats, pitcher_stats, game_context
            )

            # Debug tensor sizes only on first iteration
            if i == 0:
                print(f"Game state tensor shape: {game_state_tensor.shape}")
                print(f"Context tensor shape: {context_tensor.shape}")
                print(f"Expected combined shape: {self.full_feature_dim}")
                combined_size = game_state_tensor.shape[0] + context_tensor.shape[0]
                print(f"Actual combined size: {combined_size}")

            # Combine for input
            input_features[i] = torch.cat([game_state_tensor, context_tensor])
            game_states[i] = game_state_tensor

            # Target is the next game state (shifted by 1)
            if i < seq_len - 1:
                next_state = pa_sequence[i + 1]["state"]
                target_states[i] = self.extract_game_state_tensor(next_state)
            else:
                # For the last PA, target could be end-of-game state or padding
                target_states[i] = game_state_tensor  # or special end token

        return input_features, target_states, game_states


# Usage example
def process_jsonl_file(jsonl_path: str, max_games: int = None):
    """Process JSONL file and create training tensors."""
    featurizer = BaseballFeaturizer()

    all_input_features = []
    all_target_states = []
    all_game_states = []

    with open(jsonl_path, "r") as f:
        for i, line in enumerate(f):
            if max_games and i >= max_games:
                break

            game_data = json.loads(line)

            # Create tensors for this game
            input_feat, target_states, game_states = featurizer.create_sequence_tensors(
                game_data
            )

            all_input_features.append(input_feat)
            all_target_states.append(target_states)
            all_game_states.append(game_states)

            if i % 100 == 0:
                print(f"Processed {i} games...")

    # Stack all games (you might want to pad sequences to same length)
    return all_input_features, all_target_states, all_game_states


# For Mamba training
def create_mamba_batch(
    input_features_list: List[torch.Tensor],
    target_states_list: List[torch.Tensor],
    max_seq_len: int = 200,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Create padded batches for Mamba training.
    """
    batch_size = len(input_features_list)
    feature_dim = input_features_list[0].shape[1]
    state_dim = target_states_list[0].shape[1]

    # Pad sequences to max_seq_len
    padded_inputs = torch.zeros(batch_size, max_seq_len, feature_dim)
    padded_targets = torch.zeros(batch_size, max_seq_len, state_dim)

    for i, (inp, tgt) in enumerate(zip(input_features_list, target_states_list)):
        seq_len = min(inp.shape[0], max_seq_len)
        padded_inputs[i, :seq_len] = inp[:seq_len]
        padded_targets[i, :seq_len] = tgt[:seq_len]

    return padded_inputs, padded_targets


# Example usage:
if __name__ == "__main__":
    # Process your JSONL file
    input_feat, target_states, game_states = process_jsonl_file(
        "mamba_training_data.jsonl", max_games=1000
    )

    # Create a batch for training
    batch_inputs, batch_targets = create_mamba_batch(
        input_feat[:32], target_states[:32]
    )

    print(
        f"Batch input shape: {batch_inputs.shape}"
    )  # [32, 200, 64] (13 game state + 51 context)
    print(
        f"Batch target shape: {batch_targets.shape}"
    )  # [32, 200, 13] (game state only)

    # Note: In your Mamba model, you'll want to handle player IDs specially:
    # 1. Extract batter_id and pitcher_id from positions [-2:] of the game state
    # 2. Pass them through embedding layers to get dense representations
    # 3. Combine embeddings with the other numerical features