deepcfr.py

"""
deepcfr.py - Deep CFR with Enhanced Features for Production Use
================================================================

DEEP REINFORCEMENT LEARNING POKER AI using Counterfactual Regret Minimization
with neural network function approximation.

KEY FEATURES (Enhanced from cfr7.py):
-------------------------------------
✅ Neural network function approximation (advantage + strategy networks)
✅ Experience replay buffers for stable learning
✅ Enhanced 30-dimensional feature extraction
✅ External sampling CFR (most efficient variant)
✅ Fixed bet size abstractions (0.5x, 1x, 2x pot + all-in)
✅ GPU acceleration with CUDA support
✅ Batch normalization and dropout for regularization
✅ Gradient clipping for stability
✅ Comprehensive training metrics and logging
✅ Model checkpointing and resumption
✅ Compatible with existing src.py infrastructure

IMPROVEMENTS OVER cfr7.py:
--------------------------
1. SCALABILITY: Constant memory (neural nets) vs growing dictionary
2. GENERALIZATION: Networks generalize to unseen states
3. FULL GAME: Can handle full Texas Hold'em without abstractions
4. GPU USAGE: True GPU utilization with neural network training
5. LEARNING: Adapts and improves with more training data

USAGE:
------
# Training (default 10,000 iterations):
python deepcfr.py

# Training with custom iterations:
python deepcfr.py --iterations 50000 --save-every 5000

# Testing trained model:
python deepcfr.py --test-only --model-path poker_ai_dev/model/iteration_10000

# In code:
from deepcfr import DeepCFRBot
bot = DeepCFRBot("poker_ai_dev/model/iteration_10000", player_idx=0)
action = bot.get_action(table)

ARCHITECTURE:
-------------
- Input: 30-dimensional feature vector (hand, board, betting, position)
- Hidden: 256 → 128 neurons with ReLU, BatchNorm, Dropout
- Output: 7-dimensional action vector (fold, check/call, 5 bet sizes)
- Two networks per player: Advantage (regrets) + Strategy (policy)

TRAINING ALGORITHM:
-------------------
1. External Sampling CFR: traverse all actions for target player, sample opponent
2. Store experiences (state, regrets) in replay buffer
3. Periodically train networks on batched experiences
4. Use regret matching to convert advantages to strategy
5. Repeat for thousands of iterations until convergence

Compatible with: src.py (Table, Player, Cards, evaluate_best_hand)
"""

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import random
import copy
import os
import json
import time
from collections import deque
from typing import Dict, List, Tuple, Optional
from tqdm import tqdm

# Import your existing classes
try:
    from src import Table, Player, Card, Cards, evaluate_best_hand
except ImportError:
    print("ERROR: Cannot import from src.py - ensure src.py exists in the same directory!")
    exit()

# Set seeds for reproducibility
torch.manual_seed(42)
np.random.seed(42)
random.seed(42)

# GPU setup
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

#--------------------------------------------------------------------
# CONFIGURATION
#--------------------------------------------------------------------
# Deep CFR Training Parameters
LEARNING_RATE = 0.001
BATCH_SIZE = 64
BUFFER_SIZE = 50000
HIDDEN_SIZE_1 = 256
HIDDEN_SIZE_2 = 128
TRAIN_FREQUENCY = 50  # Train networks every N iterations
MAX_DEPTH = 15  # Maximum recursion depth

# Feature extraction settings
FEATURE_SIZE = 30  # Enhanced feature vector size
NUM_ACTIONS = 7  # fold, check/call, bet/raise (0.5x, 1x, 2x pot, allin)

#--------------------------------------------------------------------
# NEURAL NETWORKS - Enhanced Architecture
#--------------------------------------------------------------------
class AdvantageNetwork(nn.Module):
    """Network to approximate cumulative regrets (counterfactual values)"""
    def __init__(self, input_size: int = FEATURE_SIZE, num_actions: int = NUM_ACTIONS):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_size, HIDDEN_SIZE_1),
            nn.ReLU(),
            nn.BatchNorm1d(HIDDEN_SIZE_1),
            nn.Dropout(0.1),
            nn.Linear(HIDDEN_SIZE_1, HIDDEN_SIZE_2),
            nn.ReLU(),
            nn.BatchNorm1d(HIDDEN_SIZE_2),
            nn.Dropout(0.1),
            nn.Linear(HIDDEN_SIZE_2, num_actions)
        )
    
    def forward(self, x):
        if x.dim() == 1:
            x = x.unsqueeze(0)
        return self.net(x)

class StrategyNetwork(nn.Module):
    """Network to approximate average strategy (policy)"""
    def __init__(self, input_size: int = FEATURE_SIZE, num_actions: int = NUM_ACTIONS):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_size, HIDDEN_SIZE_1),
            nn.ReLU(),
            nn.BatchNorm1d(HIDDEN_SIZE_1),
            nn.Dropout(0.1),
            nn.Linear(HIDDEN_SIZE_1, HIDDEN_SIZE_2),
            nn.ReLU(),
            nn.BatchNorm1d(HIDDEN_SIZE_2),
            nn.Dropout(0.1),
            nn.Linear(HIDDEN_SIZE_2, num_actions)
        )
    
    def forward(self, x):
        if x.dim() == 1:
            x = x.unsqueeze(0)
        logits = self.net(x)
        return torch.softmax(logits, dim=-1)

#--------------------------------------------------------------------
# REPLAY BUFFER
#--------------------------------------------------------------------
class ReplayBuffer:
    """Experience replay buffer for storing and sampling training data"""
    def __init__(self, max_size: int = BUFFER_SIZE):
        self.buffer = deque(maxlen=max_size)
        self.priorities = deque(maxlen=max_size)  # For prioritized replay
        
    def add(self, experience, priority: float = 1.0):
        self.buffer.append(experience)
        self.priorities.append(priority)
    
    def sample(self, batch_size: int = BATCH_SIZE):
        if len(self.buffer) < batch_size:
            return None
        return random.sample(self.buffer, batch_size)
    
    def __len__(self):
        return len(self.buffer)

#--------------------------------------------------------------------
# FEATURE EXTRACTION - Enhanced for better state representation
#--------------------------------------------------------------------
class EnhancedFeatureExtractor:
    """Extract rich features from game state for neural network input"""
    def __init__(self):
        self.rank_map = {'2':2,'3':3,'4':4,'5':5,'6':6,'7':7,'8':8,'9':9,'10':10,'J':11,'Q':12,'K':13,'A':14}
        
    def get_hand_strength_bucket(self, hole_cards: List[Card], community: List[Card]) -> int:
        """Bucket hands into strength categories (0-20)"""
        if len(hole_cards) < 2:
            return 0
            
        if not community:  # Pre-flop
            r = sorted([self.rank_map[c.rank] for c in hole_cards], reverse=True)
            suited = hole_cards[0].suit == hole_cards[1].suit
            gap = r[0] - r[1]
            
            # Granular preflop buckets
            if r[0] == r[1]:  # Pairs
                if r[0] >= 14: return 15  # AA
                elif r[0] >= 13: return 14  # KK
                elif r[0] >= 11: return 13  # QQ,JJ
                elif r[0] >= 9: return 12   # TT,99
                else: return 10
            elif r[0] >= 14:  # Ace hands
                if suited and r[1] >= 12: return 15  # AKs, AQs
                elif r[1] >= 12: return 14  # AK, AQ
                elif suited: return 12
                else: return 10
            elif suited and gap <= 1:  # Suited connectors
                return 11 if r[0] >= 10 else 9
            else:
                return max(1, min(8, r[0] - 6))
        else:  # Post-flop
            try:
                all_cards = list(hole_cards) + list(community)
                strength = evaluate_best_hand(all_cards)
                hand_type = int(strength[0])
                
                if hand_type >= 8: return 20  # Straight flush+
                elif hand_type == 7: return 19  # Full house
                elif hand_type == 6: return 18  # Flush
                elif hand_type == 5: return 17  # Straight
                elif hand_type == 4: return 16  # Three of a kind
                elif hand_type == 3: return 14  # Two pair
                elif hand_type == 2: return 12  # One pair
                else: return 8  # High card
            except:
                return 8
        
        return 5
        
    def extract_features(self, table: Table, player_idx: int) -> torch.FloatTensor:
        """Extract comprehensive features (30 dimensions)"""
        features = []
        player = table.players[player_idx]
        opponent = table.players[1 - player_idx]
        
        # 1. Hand features (8 dimensions)
        if len(player.hole_cards) == 2:
            c1, c2 = player.hole_cards
            r1, r2 = self.rank_map[c1.rank], self.rank_map[c2.rank]
            features.extend([r1/14.0, r2/14.0])  # Normalized ranks
            features.append(float(c1.suit == c2.suit))  # Suited
            features.append(float(abs(r1-r2) <= 1))  # Connected
            features.append(max(r1, r2) / 14.0)  # High card
            features.append(min(r1, r2) / 14.0)  # Low card
            features.append(float(r1 == r2))  # Pocket pair
            
            # Hand strength bucket
            strength_bucket = self.get_hand_strength_bucket(
                list(player.hole_cards), 
                list(table.community)
            )
            features.append(strength_bucket / 20.0)  # Normalized strength
        else:
            features.extend([0.0] * 8)
            
        # 2. Stage features (4 dimensions)
        stage_encoding = {
            "pre-flop": [1,0,0,0], 
            "flop": [0,1,0,0], 
            "turn": [0,0,1,0], 
            "river": [0,0,0,1]
        }
        features.extend(stage_encoding.get(table.stage, [0,0,0,0]))
        
        # 3. Betting features (6 dimensions)
        features.append(min(table.pot / 100.0, 10.0))  # Pot size
        features.append(min(table.current_bet / 50.0, 5.0))  # Current bet
        features.append(min(table.amount_to_call() / 50.0, 5.0))  # To call
        features.append(float(table.bet_occurred))  # Betting occurred
        
        # Stack-to-pot ratio (SPR)
        stack_size = player.chips + player.total_committed
        spr = stack_size / max(table.pot, 1)
        features.append(min(spr / 20.0, 1.0))  # Normalized SPR
        
        # Position
        features.append(float(player_idx == table.dealer))  # 1 if dealer (BTN)
        
        # 4. Player features (6 dimensions)
        features.append(min(player.chips / 1000.0, 2.0))  # Stack size
        features.append(min(player.total_committed / 100.0, 5.0))  # Committed
        features.append(float(player.in_hand))  # Still in hand
        features.append(min(player.current_bet / 50.0, 5.0))  # Current bet
        
        # Pot odds
        to_call = table.amount_to_call()
        pot_odds = to_call / max(table.pot + to_call, 1)
        features.append(min(pot_odds, 1.0))
        
        # Bet sizing (last action)
        if hasattr(table, 'action_history') and table.action_history:
            last_action = table.action_history[-1]
            if 'raise' in last_action or 'bet' in last_action:
                features.append(1.0)
            else:
                features.append(0.0)
        else:
            features.append(0.0)
        
        # 5. Opponent features (6 dimensions)
        features.append(min(opponent.chips / 1000.0, 2.0))  # Opponent stack
        features.append(min(opponent.total_committed / 100.0, 5.0))  # Opponent committed
        features.append(float(opponent.in_hand))  # Opponent still in
        features.append(min(opponent.current_bet / 50.0, 5.0))  # Opponent bet
        
        # Opponent SPR
        opp_stack = opponent.chips + opponent.total_committed
        opp_spr = opp_stack / max(table.pot, 1)
        features.append(min(opp_spr / 20.0, 1.0))
        
        # Opponent aggression indicator
        features.append(float(opponent.current_bet > table.bb * 2))
        
        return torch.FloatTensor(features).to(device)

#--------------------------------------------------------------------
# DEEP CFR TRAINER - Main training class
#--------------------------------------------------------------------
class DeepCFRTrainer:
    """Deep CFR trainer with neural function approximation"""
    def __init__(self, sb: int = 5, bb: int = 10):
        self.feature_extractor = EnhancedFeatureExtractor()
        self.sb = sb
        self.bb = bb
        
        # Networks (moved to GPU if available)
        self.advantage_net_p0 = AdvantageNetwork().to(device)
        self.advantage_net_p1 = AdvantageNetwork().to(device)
        self.strategy_net_p0 = StrategyNetwork().to(device)
        self.strategy_net_p1 = StrategyNetwork().to(device)
        
        # Optimizers
        self.adv_opt_p0 = optim.Adam(self.advantage_net_p0.parameters(), lr=LEARNING_RATE)
        self.adv_opt_p1 = optim.Adam(self.advantage_net_p1.parameters(), lr=LEARNING_RATE)
        self.strat_opt_p0 = optim.Adam(self.strategy_net_p0.parameters(), lr=LEARNING_RATE)
        self.strat_opt_p1 = optim.Adam(self.strategy_net_p1.parameters(), lr=LEARNING_RATE)
        
        # Replay buffers
        self.adv_buffer_p0 = ReplayBuffer()
        self.adv_buffer_p1 = ReplayBuffer()
        self.strat_buffer_p0 = ReplayBuffer()
        self.strat_buffer_p1 = ReplayBuffer()
        
        # Action mapping for fixed bet sizes (like cfr7.py)
        self.action_map = {
            "fold": 0, 
            "check": 1, 
            "call": 1,
            "bet_0.5x": 2,
            "bet_1.0x": 3,
            "bet_2.0x": 4,
            "bet_allin": 5,
            "raise_0.5x": 2,
            "raise_1.0x": 3,
            "raise_2.0x": 4,
            "raise_allin": 5,
            "raise": 6
        }
        
        # Training statistics
        self.iteration = 0
        self.losses_adv = []
        self.losses_strat = []
        
    def get_legal_actions_with_betsizes(self, table: Table) -> List[str]:
        """Get legal actions with fixed bet size abstractions (like cfr7.py)"""
        if table.to_act >= len(table.players):
            return []
            
        player = table.players[table.to_act]
        to_call = table.amount_to_call()
        
        actions = []
        
        if to_call == 0:  # Can check or bet
            actions.append("check")
            if player.chips > 0:
                # Add fixed bet sizes
                for size in [0.5, 1.0, 2.0]:
                    bet_amt = int(table.pot * size)
                    if bet_amt <= player.chips:
                        actions.append(f"bet_{size}x")
                if player.chips > 0:
                    actions.append("bet_allin")
        else:  # Must call or fold
            actions.append("fold")
            if player.chips >= to_call:
                actions.append("call")
            
            # Add fixed raise sizes
            if player.chips > to_call:
                for size in [0.5, 1.0, 2.0]:
                    raise_amt = int(table.pot * size)
                    if raise_amt + to_call <= player.chips:
                        actions.append(f"raise_{size}x")
                if player.chips > to_call:
                    actions.append("raise_allin")
        
        return actions
    
    def get_strategy(self, table: Table, player_idx: int) -> Dict[str, float]:
        """Get strategy using regret matching on neural network advantages"""
        features = self.feature_extractor.extract_features(table, player_idx)
        valid_actions = self.get_legal_actions_with_betsizes(table)
        
        if not valid_actions:
            return {}
        
        # Get advantage network
        adv_net = self.advantage_net_p0 if player_idx == 0 else self.advantage_net_p1
        adv_net.eval()
        
        with torch.no_grad():
            advantages = adv_net(features).squeeze()
            
        # Regret matching: max(0, regret) for each action
        action_probs = {}
        total_positive = 0.0
        
        for action in valid_actions:
            idx = self.action_map.get(action, 1)
            if idx < len(advantages):
                prob = max(0.0, advantages[idx].item())
            else:
                prob = 0.0
            action_probs[action] = prob
            total_positive += prob
            
        # Normalize to probability distribution
        if total_positive > 0:
            for action in action_probs:
                action_probs[action] /= total_positive
        else:
            # Uniform random if no positive advantages
            uniform_prob = 1.0 / len(action_probs)
            for action in action_probs:
                action_probs[action] = uniform_prob
                
        return action_probs
    
    def apply_action(self, table: Table, action: str) -> bool:
        """Apply action to table, return success status"""
        try:
            player = table.players[table.to_act]
            to_call = table.amount_to_call()
            
            if action == "fold":
                player.fold()
            elif action == "check":
                if to_call != 0:
                    return False  # Invalid check
            elif action == "call":
                table.pot += player.call(to_call)
                table.bet_occurred = True
            elif action.startswith("bet_") or action.startswith("raise_"):
                # Parse bet size
                if "allin" in action:
                    bet_amt = player.chips
                elif "_" in action:
                    try:
                        size_str = action.split("_")[1].replace("x", "")
                        size = float(size_str)
                        bet_amt = int(table.pot * size)
                    except:
                        bet_amt = table.bb * 2
                else:
                    bet_amt = table.bb * 2
                
                if action.startswith("raise_"):
                    # Raise = call + additional bet
                    bet_amt = min(bet_amt, player.chips - to_call)
                    table.pot += player.raise_bet(to_call, bet_amt)
                else:
                    # Bet (no one has bet yet)
                    bet_amt = min(bet_amt, player.chips)
                    table.pot += player.bet(bet_amt)
                
                table.current_bet = player.current_bet
                table.bet_occurred = True
            
            return True
        except:
            return False
    
    def is_terminal(self, table: Table) -> bool:
        """Check if game state is terminal"""
        alive = [p for p in table.players if p.in_hand]
        return len(alive) <= 1 or table.stage == "showdown"
    
    def get_payoff(self, table: Table, player_idx: int) -> float:
        """Calculate payoff for player at terminal state"""
        alive = [i for i, p in enumerate(table.players) if p.in_hand]
        
        # Fold win
        if len(alive) == 1:
            winner = alive[0]
            total_pot = sum(p.total_committed for p in table.players)
            if player_idx == winner:
                return total_pot - table.players[player_idx].total_committed
            else:
                return -table.players[player_idx].total_committed
        
        # Showdown
        if len(alive) >= 2 and len(table.community) == 5:
            try:
                # Evaluate hands
                hands = []
                for i in alive:
                    p = table.players[i]
                    if len(p.hole_cards) == 2:
                        all_cards = list(p.hole_cards) + list(table.community)
                        strength = evaluate_best_hand(all_cards)
                        hands.append((strength, i))
                
                hands.sort(key=lambda x: x[0], reverse=True)
                best_strength = hands[0][0]
                winners = [i for strength, i in hands if strength == best_strength]
                
                total_pot = sum(p.total_committed for p in table.players)
                if player_idx in winners:
                    return (total_pot / len(winners)) - table.players[player_idx].total_committed
                else:
                    return -table.players[player_idx].total_committed
            except:
                pass
        
        return 0.0
    
    def advance_to_next_player(self, table: Table):
        """Move to next active player"""
        table.to_act = (table.to_act + 1) % len(table.players)
        attempts = 0
        while attempts < len(table.players):
            if table.players[table.to_act].in_hand and table.players[table.to_act].chips > 0:
                return
            table.to_act = (table.to_act + 1) % len(table.players)
            attempts += 1
    
    def check_advance_stage(self, table: Table):
        """Check if betting round is over and advance stage"""
        alive = [p for p in table.players if p.in_hand]
        if len(alive) <= 1:
            return
        
        active_with_chips = [p for p in alive if p.chips > 0]
        
        # Check if betting complete
        all_checked = not table.bet_occurred and all(
            p.current_bet == 0 for p in alive
        )
        all_matched = table.bet_occurred and all(
            p.current_bet == table.current_bet or p.chips == 0 
            for p in alive
        )
        
        if all_checked or all_matched:
            # Advance stage
            if table.stage == "pre-flop":
                table.stage = "flop"
                table.community = table.comm_buffer[:3]
            elif table.stage == "flop":
                table.stage = "turn"
                table.community = table.comm_buffer[:4]
            elif table.stage == "turn":
                table.stage = "river"
                table.community = table.comm_buffer[:5]
            elif table.stage == "river":
                table.stage = "showdown"
            
            # Reset betting
            if table.stage != "showdown":
                for p in table.players:
                    p.current_bet = 0
                table.current_bet = 0
                table.bet_occurred = False
                
                # Start from dealer+1
                table.to_act = (table.dealer + 1) % len(table.players)
                while not table.players[table.to_act].in_hand or table.players[table.to_act].chips == 0:
                    table.to_act = (table.to_act + 1) % len(table.players)
    
    def cfr_iteration(self, target_player: int = 0) -> float:
        """Run one CFR iteration for target player (external sampling)"""
        # Create fresh game
        players = [Player("P0", 1000), Player("P1", 1000)]
        table = Table(players, self.sb, self.bb)
        table.auto_advance = False
        table.start_hand()
        
        # Run traversal
        return self._traverse(table, target_player, 0)
    
    def _traverse(self, table: Table, target_player: int, depth: int) -> float:
        """External sampling CFR traversal"""
        if depth > MAX_DEPTH:
            return 0.0
        
        if self.is_terminal(table):
            return self.get_payoff(table, target_player)
        
        current_player = table.to_act
        strategy = self.get_strategy(table, current_player)
        valid_actions = self.get_legal_actions_with_betsizes(table)
        
        if not valid_actions:
            return 0.0
        
        if current_player == target_player:
            # Target player: traverse ALL actions
            utilities = []
            for action in valid_actions:
                # Deep copy table
                new_table = copy.deepcopy(table)
                
                if self.apply_action(new_table, action):
                    self.advance_to_next_player(new_table)
                    self.check_advance_stage(new_table)
                    
                    util = -self._traverse(new_table, target_player, depth + 1)
                    utilities.append(util)
                else:
                    utilities.append(-1.0)  # Invalid action penalty
            
            # Calculate node value
            node_value = sum(strategy.get(a, 0.0) * u for a, u in zip(valid_actions, utilities))
            
            # Store regrets for training
            features = self.feature_extractor.extract_features(table, current_player)
            regrets = torch.zeros(NUM_ACTIONS).to(device)
            
            for i, action in enumerate(valid_actions):
                idx = self.action_map.get(action, 1)
                if idx < NUM_ACTIONS:
                    regret = utilities[i] - node_value
                    regrets[idx] = regret
            
            buffer = self.adv_buffer_p0 if current_player == 0 else self.adv_buffer_p1
            buffer.add((features, regrets))
            
            return node_value
        else:
            # Opponent: sample ONE action
            acts_list = list(valid_actions)
            probs = [strategy.get(a, 0.0) for a in acts_list]
            
            # Normalize
            prob_sum = sum(probs)
            if prob_sum > 0:
                probs = [p / prob_sum for p in probs]
            else:
                probs = [1.0 / len(acts_list)] * len(acts_list)
            
            chosen_action = np.random.choice(acts_list, p=probs)
            
            # Apply and continue
            if self.apply_action(table, chosen_action):
                self.advance_to_next_player(table)
                self.check_advance_stage(table)
                return -self._traverse(table, target_player, depth + 1)
            else:
                return 0.0
    
    def train_networks(self):
        """Train advantage and strategy networks"""
        total_loss = 0.0
        
        for player_idx in [0, 1]:
            # Train advantage network
            buffer = self.adv_buffer_p0 if player_idx == 0 else self.adv_buffer_p1
            optimizer = self.adv_opt_p0 if player_idx == 0 else self.adv_opt_p1
            network = self.advantage_net_p0 if player_idx == 0 else self.advantage_net_p1
            
            batch = buffer.sample(BATCH_SIZE)
            if batch is None:
                continue
            
            network.train()
            
            features = torch.stack([exp[0] for exp in batch])
            targets = torch.stack([exp[1] for exp in batch])
            
            optimizer.zero_grad()
            predictions = network(features)
            loss = nn.MSELoss()(predictions, targets)
            loss.backward()
            
            # Gradient clipping
            torch.nn.utils.clip_grad_norm_(network.parameters(), max_norm=1.0)
            
            optimizer.step()
            
            total_loss += loss.item()
            self.losses_adv.append(loss.item())
        
        return total_loss / 2.0

#--------------------------------------------------------------------
# TRAINING FUNCTION
#--------------------------------------------------------------------
def train_deep_cfr(iterations: int = 10000, save_every: int = 1000, sb: int = 5, bb: int = 10):
    """Main training function with comprehensive logging and saving"""
    print("="*70)
    print("DEEP CFR TRAINING - Neural Network Poker AI")
    print("="*70)
    print(f"Device: {device}")
    print(f"Iterations: {iterations:,}")
    print(f"Batch size: {BATCH_SIZE}")
    print(f"Buffer size: {BUFFER_SIZE}")
    print(f"Learning rate: {LEARNING_RATE}")
    print(f"Feature size: {FEATURE_SIZE}")
    print(f"Hidden layers: {HIDDEN_SIZE_1} → {HIDDEN_SIZE_2}")
    print("="*70)
    
    # Create models directory
    model_dir = os.path.join("poker_ai_dev", "model")
    os.makedirs(model_dir, exist_ok=True)
    
    trainer = DeepCFRTrainer(sb=sb, bb=bb)
    
    start_time = time.time()
    utilities_p0 = []
    utilities_p1 = []
    
    # Training loop with progress bar
    with tqdm(total=iterations, desc="Deep CFR Training", unit="iter") as pbar:
        for iteration in range(iterations):
            trainer.iteration = iteration + 1
            
            try:
                # Run CFR iteration for both players (external sampling)
                util_p0 = trainer.cfr_iteration(target_player=0)
                util_p1 = trainer.cfr_iteration(target_player=1)
                
                utilities_p0.append(util_p0)
                utilities_p1.append(util_p1)
                
                # Train networks periodically
                if iteration > 100 and iteration % TRAIN_FREQUENCY == 0:
                    loss = trainer.train_networks()
                    pbar.set_postfix({
                        'loss': f'{loss:.4f}',
                        'util_p0': f'{util_p0:.2f}',
                        'buf_p0': len(trainer.adv_buffer_p0),
                        'buf_p1': len(trainer.adv_buffer_p1)
                    })
                
                # Save checkpoints
                if iteration % save_every == 0 and iteration > 0:
                    save_models(trainer, iteration, utilities_p0, utilities_p1)
                
                pbar.update(1)
                    
            except Exception as e:
                print(f"\nError in iteration {iteration}: {e}")
                import traceback
                traceback.print_exc()
                continue
    
    # Final save
    elapsed = time.time() - start_time
    print(f"\n{'='*70}")
    print(f"Training complete in {elapsed:.1f}s ({iterations/elapsed:.1f} iter/s)")
    print(f"Buffer sizes: P0={len(trainer.adv_buffer_p0)}, P1={len(trainer.adv_buffer_p1)}")
    
    save_models(trainer, iterations, utilities_p0, utilities_p1)
    
    return trainer

def save_models(trainer: DeepCFRTrainer, iteration: int, 
                utilities_p0: List[float] = None, utilities_p1: List[float] = None):
    """Save models and training statistics to disk"""
    save_dir = os.path.join("poker_ai_dev", "model", f"iteration_{iteration}")
    os.makedirs(save_dir, exist_ok=True)
    
    # Save advantage networks
    torch.save(trainer.advantage_net_p0.state_dict(), 
               os.path.join(save_dir, "advantage_p0.pth"))
    torch.save(trainer.advantage_net_p1.state_dict(), 
               os.path.join(save_dir, "advantage_p1.pth"))
    
    # Save strategy networks  
    torch.save(trainer.strategy_net_p0.state_dict(), 
               os.path.join(save_dir, "strategy_p0.pth"))
    torch.save(trainer.strategy_net_p1.state_dict(), 
               os.path.join(save_dir, "strategy_p1.pth"))
    
    # Save training metadata
    metadata = {
        'iteration': iteration,
        'device': str(device),
        'feature_size': FEATURE_SIZE,
        'num_actions': NUM_ACTIONS,
        'hidden_sizes': [HIDDEN_SIZE_1, HIDDEN_SIZE_2],
        'learning_rate': LEARNING_RATE,
        'batch_size': BATCH_SIZE,
        'buffer_size': BUFFER_SIZE,
        'buffer_lengths': {
            'adv_p0': len(trainer.adv_buffer_p0),
            'adv_p1': len(trainer.adv_buffer_p1),
        },
        'timestamp': time.strftime('%Y-%m-%d %H:%M:%S')
    }
    
    if utilities_p0:
        metadata['avg_utility_p0'] = np.mean(utilities_p0[-1000:])
        metadata['avg_utility_p1'] = np.mean(utilities_p1[-1000:])
    
    with open(os.path.join(save_dir, "metadata.json"), 'w') as f:
        json.dump(metadata, f, indent=2)
    
    print(f"✅ Models saved to: {os.path.abspath(save_dir)}")

#--------------------------------------------------------------------
# BOT CLASS FOR INFERENCE
#--------------------------------------------------------------------
class DeepCFRBot:
    """Trained Deep CFR bot for live play"""
    def __init__(self, model_path: str, player_idx: int = 0):
        self.feature_extractor = EnhancedFeatureExtractor()
        self.player_idx = player_idx
        
        # Load advantage network (used for strategy via regret matching)
        self.advantage_net = AdvantageNetwork().to(device)
        adv_file = os.path.join(model_path, f"advantage_p{player_idx}.pth")
        
        if os.path.exists(adv_file):
            self.advantage_net.load_state_dict(torch.load(adv_file, map_location=device))
            self.advantage_net.eval()
            print(f"✅ Loaded advantage network from {adv_file}")
        else:
            print(f"❌ Model file not found: {adv_file}")
            
        # Load strategy network
        self.strategy_net = StrategyNetwork().to(device)
        strat_file = os.path.join(model_path, f"strategy_p{player_idx}.pth")
        
        if os.path.exists(strat_file):
            self.strategy_net.load_state_dict(torch.load(strat_file, map_location=device))
            self.strategy_net.eval()
            print(f"✅ Loaded strategy network from {strat_file}")
        
        self.action_map = {
            "fold": 0, 
            "check": 1, 
            "call": 1,
            "bet_0.5x": 2,
            "bet_1.0x": 3,
            "bet_2.0x": 4,
            "bet_allin": 5,
            "raise_0.5x": 2,
            "raise_1.0x": 3,
            "raise_2.0x": 4,
            "raise_allin": 5,
            "raise": 6
        }
    
    def get_legal_actions_with_betsizes(self, table: Table) -> List[str]:
        """Get legal actions with bet size abstractions"""
        player = table.players[self.player_idx]
        to_call = table.amount_to_call()
        
        actions = []
        
        if to_call == 0:
            actions.append("check")
            if player.chips > 0:
                for size in [0.5, 1.0, 2.0]:
                    bet_amt = int(table.pot * size)
                    if bet_amt <= player.chips:
                        actions.append(f"bet_{size}x")
                actions.append("bet_allin")
        else:
            actions.append("fold")
            if player.chips >= to_call:
                actions.append("call")
            if player.chips > to_call:
                for size in [0.5, 1.0, 2.0]:
                    raise_amt = int(table.pot * size)
                    if raise_amt + to_call <= player.chips:
                        actions.append(f"raise_{size}x")
                actions.append("raise_allin")
        
        return actions
    
    def get_action(self, table: Table, use_strategy_net: bool = False) -> str:
        """Get action for live play
        
        Args:
            table: Current game state
            use_strategy_net: If True, use strategy network; if False, use advantage network with regret matching
        """
        features = self.feature_extractor.extract_features(table, self.player_idx)
        valid_actions = self.get_legal_actions_with_betsizes(table)
        
        if not valid_actions:
            return "fold"
        
        with torch.no_grad():
            if use_strategy_net:
                # Use strategy network directly
                probs = self.strategy_net(features).squeeze()
            else:
                # Use advantage network with regret matching
                advantages = self.advantage_net(features).squeeze()
                probs = torch.relu(advantages)  # Regret matching
                
                # Normalize
                if probs.sum() > 0:
                    probs = probs / probs.sum()
                else:
                    probs = torch.ones_like(probs) / len(probs)
        
        # Map network outputs to valid actions
        action_probs = []
        for action in valid_actions:
            idx = self.action_map.get(action, 1)
            if idx < len(probs):
                action_probs.append(probs[idx].item())
            else:
                action_probs.append(0.0)
        
        # Normalize and sample
        action_probs = np.array(action_probs)
        if action_probs.sum() > 0:
            action_probs = action_probs / action_probs.sum()
        else:
            action_probs = np.ones(len(action_probs)) / len(action_probs)
        
        # Sample action (or take argmax for deterministic play)
        chosen_action = np.random.choice(valid_actions, p=action_probs)
        
        # Convert abstract actions to concrete ones
        if chosen_action.startswith("bet_") or chosen_action.startswith("raise_"):
            return "raise"  # Simplified for basic table interface
        
        return chosen_action
    
    def get_strategy_probabilities(self, table: Table) -> Dict[str, float]:
        """Get full strategy distribution (for analysis)"""
        features = self.feature_extractor.extract_features(table, self.player_idx)
        valid_actions = self.get_legal_actions_with_betsizes(table)
        
        with torch.no_grad():
            advantages = self.advantage_net(features).squeeze()
            probs = torch.relu(advantages)
            
            if probs.sum() > 0:
                probs = probs / probs.sum()
            else:
                probs = torch.ones_like(probs) / len(probs)
        
        strategy = {}
        for action in valid_actions:
            idx = self.action_map.get(action, 1)
            if idx < len(probs):
                strategy[action] = probs[idx].item()
            else:
                strategy[action] = 0.0
        
        # Normalize
        total = sum(strategy.values())
        if total > 0:
            strategy = {k: v/total for k, v in strategy.items()}
        
        return strategy

#--------------------------------------------------------------------
# MAIN EXECUTION
#--------------------------------------------------------------------
if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='Deep CFR Poker AI Training')
    parser.add_argument('--iterations', type=int, default=10000, 
                       help='Number of training iterations (default: 10000)')
    parser.add_argument('--save-every', type=int, default=1000, 
                       help='Save model every N iterations (default: 1000)')
    parser.add_argument('--test-only', action='store_true', 
                       help='Only test a trained model without training')
    parser.add_argument('--model-path', type=str, 
                       default='poker_ai_dev/model/iteration_10000',
                       help='Path to trained model for testing')
    
    args = parser.parse_args()
    
    if args.test_only:
        # Test mode
        print("="*70)
        print("TESTING DEEP CFR BOT")
        print("="*70)
        
        try:
            bot = DeepCFRBot(args.model_path, player_idx=0)
            print("✅ Bot loaded successfully!")
            
            # Run test games
            wins = 0
            total_games = 100
            
            print(f"\nRunning {total_games} test games...")
            for game_num in range(total_games):
                players = [Player("Bot", 1000), Player("Random", 1000)]
                table = Table(players, sb=5, bb=10)
                table.auto_advance = True
                table.start_hand()
                
                # Play until showdown or fold
                max_actions = 50
                for _ in range(max_actions):
                    if table.stage == "showdown" or len([p for p in table.players if p.in_hand]) <= 1:
                        break
                    
                    if table.to_act == 0:  # Bot's turn
                        action = bot.get_action(table)
                        
                        # Convert to basic action for table
                        if action in ["check", "fold", "call"]:
                            table.perform_action(action)
                        else:  # Any raise variant
                            table.perform_action("raise", raise_by=table.bb * 2)
                    else:  # Random opponent
                        valid = table.valid_actions()
                        action = random.choice(valid)
                        if action == "raise":
                            table.perform_action(action, raise_by=table.bb * 2)
                        else:
                            table.perform_action(action)
                
                # Check winner
                if table.players[0].chips > 1000:
                    wins += 1
                
                if (game_num + 1) % 10 == 0:
                    print(f"Games: {game_num + 1}/{total_games}, Win rate: {wins/(game_num+1)*100:.1f}%")
            
            print(f"\n{'='*70}")
            print(f"Final win rate: {wins/total_games*100:.1f}% ({wins}/{total_games} games)")
            print(f"{'='*70}")
            
        except Exception as e:
            print(f"❌ Error testing bot: {e}")
            import traceback
            traceback.print_exc()
    
    else:
        # Training mode
        print("🚀 Starting Deep CFR Training")
        print(f"📁 Models will be saved in 'poker_ai_dev/model/' directory")
        print()
        
        try:
            # Train the model
            trained_model = train_deep_cfr(
                iterations=args.iterations, 
                save_every=args.save_every,
                sb=5,
                bb=10
            )
            
            print("\n🎯 Training complete! Testing the bot...")
            
            # Quick test
            model_path = f"poker_ai_dev/model/iteration_{args.iterations}"
            bot = DeepCFRBot(model_path, player_idx=0)
            print("✅ Bot loaded successfully!")
            
            # Single test game
            players = [Player("Bot", 1000), Player("Random", 1000)]
            table = Table(players, sb=5, bb=10)
            table.auto_advance = False
            table.start_hand()
            
            print(f"\n🎲 Test game:")
            print(f"Bot cards: {table.players[0].hole_cards}")
            print(f"Community: {table.community}")
            print(f"Pot: {table.pot}, To call: {table.amount_to_call()}")
            
            strategy = bot.get_strategy_probabilities(table)
            print(f"\n🤖 Bot strategy:")
            for action, prob in sorted(strategy.items(), key=lambda x: -x[1]):
                print(f"  {action}: {prob*100:.1f}%")
            
            chosen = bot.get_action(table)
            print(f"\n✅ Bot chooses: {chosen}")
            
        except KeyboardInterrupt:
            print("\n⚠️ Training interrupted by user")
        except Exception as e:
            print(f"\n❌ Error: {e}")
            import traceback
            traceback.print_exc()