config.py

import os
from dotenv import load_dotenv
from typing import List, Dict
from collections import defaultdict  # Add this import

load_dotenv()

# Ultra-aggressive GPU settings
GPU_CONFIG = {
    'precision': 'mixed',
    'cudnn_benchmark': True,
    'pin_memory': True,
    'prefetch_factor': 6,      # More aggressive prefetching
    'gradient_clip': None,     # Remove gradient clipping for speed
    'num_workers': 4,         # Optimize for 6GB GPU
    'persistent_workers': True,
    'memory_fraction': 0.98   # Push memory usage to limit
}

# Even more aggressive GPU settings
GPU_CONFIG.update({
    'num_workers': min(os.cpu_count(), 8),  # Optimize worker count
    'prefetch_factor': 4,      # More balanced prefetch
    'pin_memory': True,
    'persistent_workers': True,
    'memory_fraction': 0.92    # Slightly more conservative
})

# Laptop-optimized GPU settings
GPU_CONFIG.update({
    'num_workers': 2,              # Reduce workers for laptop
    'prefetch_factor': 2,          # Less aggressive prefetch
    'pin_memory': True,
    'persistent_workers': True,
    'memory_fraction': 0.8         # More conservative memory usage
})

# Update GPU settings for better speed/memory balance
GPU_CONFIG = {
    'num_workers': 2,              # Reduce worker count
    'prefetch_factor': 2,          # Smaller prefetch for less memory
    'pin_memory': True,
    'persistent_workers': True,
    'memory_fraction': 0.85,       # Leave some GPU memory free
    'cudnn_benchmark': True,       # Enable cudnn autotuner
    'mixed_precision': True        # Enable mixed precision
}

# Laptop-friendly batch sizes
def get_batch_size(image_size):
    """Optimized batch sizes"""
    # Handle both tuple and integer inputs
    size = image_size[0] if isinstance(image_size, (tuple, list)) else image_size
    
    if size <= 384:
        return 48    # Larger batches for small images
    elif size <= 512:
        return 24    
    else:
        return 12    # Smaller batches for large images

# Faster training stages - even shorter for laptop
PROGRESSIVE_SIZES = [
    (384, 384, 2),     # Just 2 epochs per stage
    (512, 512, 2),     
    (640, 640, 1)      # Single epoch at largest size
]

# Basic configuration
DATA_DIR = r'C:\Users\Jeremy\Dropbox\Bird Identification ML\data'
MODELS_DIR = r'C:\Users\Jeremy\Dropbox\Bird Identification ML\models'
IMAGE_SIZE = (384, 384)  # Increased from 224x224 for better detail
BASE_BATCH_SIZE = 32  # Default batch size
EPOCHS = 100  # Increased for better convergence
IMAGES_PER_SPECIES = 500  # Much more images per species

# Learning parameters
INITIAL_LR = 1e-5  # Use single value
MIN_LR = 1e-6
WEIGHT_DECAY = 0.05  # Use single value
WARMUP_EPOCHS = 5

# Training schedule
WARMUP_EPOCHS = 5
INITIAL_LR = 1e-5  # Use single value
MIN_LR = 1e-6
WEIGHT_DECAY = 0.05  # Use single value

# Data augmentation
AUG_STRENGTH = 0.5  # Controls augmentation intensity
CUTMIX_PROB = 0.0  # Disable cutmix for speed
MIXUP_ALPHA = 0.0

# Data collection settings
IMAGES_PER_SPECIES = 350  # Increase from 250
MIN_IMAGES_REQUIRED = 300  # Increase minimum requirement
AUGMENTATION_FACTOR = 4    # Each image becomes 4 (flips, colors, etc)

# Verification thresholds
IMAGE_QUALITY = {
    'min_resolution': (1024, 1024),  # Increase minimum resolution
    'min_per_species': MIN_IMAGES_REQUIRED,
    'target_per_species': IMAGES_PER_SPECIES,
    'augmented_per_species': IMAGES_PER_SPECIES * AUGMENTATION_FACTOR
}

# Add focused collection for commonly confused species
PRIORITY_SPECIES = {
    'Yellow-rumped Warbler': 500,  # Critical - only 1.8% accuracy
    'Yellow Warbler': 500,         # Critical - 12.7% accuracy
    'Eastern Wood-Pewee': 500,     # Critical - 12.7% accuracy
    'American Bittern': 500,       # Major confusion source
    'Scarlet Tanager': 400,        # 50.9% accuracy
    'Mourning Warbler': 400,       # 50.9% accuracy
    'Yellow Rail': 400,            # 43.6% accuracy
    'Bobolink': 400                # 58.2% accuracy
}

# Complete species list
SPECIES_LIST = [
    'Acadian Flycatcher', 'American Bittern', 'American Crow', 'American Goldfinch',
    'American Kestrel', 'American Redstart', 'American Robin', 'American Woodcock',
    'Baltimore Oriole', 'Barn Swallow', 'Bay-breasted Warbler', 'Belted Kingfisher',
    'Black-and-white Warbler', 'Black-billed Cuckoo', 'Blackburnian Warbler',
    'Black-capped Chickadee', 'Black-throated Blue Warbler', 'Black-throated Green Warbler',
    'Blue Grosbeak', 'Blue Jay', 'Blue-gray Gnatcatcher', 'Blue-headed Vireo',
    'Bobolink', 'Brown Creeper', 'Brown Thrasher', 'Brown-headed Cowbird',
    'Canada Warbler', 'Cape May Warbler', 'Carolina Chickadee', 'Carolina Wren',
    'Cedar Waxwing', 'Chestnut-sided Warbler', 'Chimney Swift', 'Chipping Sparrow',
    'Clapper Rail', 'Common Grackle', 'Common Yellowthroat', 'Connecticut Warbler',
    "Cooper's Hawk", 'Dark-eyed Junco', 'Downy Woodpecker', 'Eastern Bluebird',
    'Eastern Kingbird', 'Eastern Meadowlark', 'Eastern Phoebe', 'Eastern Screech-Owl',
    'Eastern Towhee', 'Eastern Wood-Pewee', 'European Starling', 'Field Sparrow',
    'Fox Sparrow', 'Golden-crowned Kinglet', 'Golden-winged Warbler', 'Grasshopper Sparrow',
    'Gray Catbird', 'Gray-cheeked Thrush', 'Great Crested Flycatcher', 'Great Horned Owl',
    'Hairy Woodpecker', 'Hermit Thrush', 'Hooded Warbler', 'House Finch',
    'House Sparrow', 'House Wren', 'Ruby-throated Hummingbird', 'Indigo Bunting',
    'Kentucky Warbler', 'Least Flycatcher', 'Lincoln\'s Sparrow', 'Louisiana Waterthrush',
    'Magnolia Warbler', 'Mallard', 'Marsh Wren', 'Mourning Dove',
    'Mourning Warbler', 'Nashville Warbler', 'Northern Cardinal', 'Northern Flicker',
    'Northern Mockingbird', 'Northern Parula', 'Northern Saw-whet Owl', 'Northern Waterthrush',
    'Orange-crowned Warbler', 'Ovenbird', 'Palm Warbler', 'Pileated Woodpecker',
    'Pine Siskin', 'Pine Warbler', 'Purple Finch', 'Purple Martin',
    'Red-eyed Vireo', 'Red-bellied Woodpecker', 'Red-breasted Nuthatch', 'Red-headed Woodpecker',
    'Red-shouldered Hawk', 'Red-tailed Hawk', 'Rose-breasted Grosbeak', 'Ruby-crowned Kinglet',
    'Savannah Sparrow', 'Scarlet Tanager', 'Sharp-shinned Hawk', 'Song Sparrow',
    'Sora', 'Summer Tanager', 'Swainson\'s Thrush', 'Swamp Sparrow',
    'Tennessee Warbler', 'Tufted Titmouse', 'Veery', 'Virginia Rail',
    'Whip-poor-will', 'White-breasted Nuthatch', 'White-crowned Sparrow', 'White-throated Sparrow',
    'Willow Flycatcher', 'Wilson\'s Warbler', 'Winter Wren', 'Wood Duck',
    'Wood Thrush', 'Worm-eating Warbler', 'Yellow-bellied Sapsucker', 'Yellow-billed Cuckoo',
    'Yellow-breasted Chat', 'Yellow Rail', 'Yellow Warbler', 'Yellow-bellied Flycatcher',
    'Yellow-rumped Warbler'
]

# Get actual number of classes from the data directory
def get_actual_classes() -> tuple[int, Dict, Dict]:
    """Get number of classes and mappings from the data directory"""
    if not os.path.exists(DATA_DIR):
        return len(SPECIES_LIST), {}, {}
        
    classes = sorted([d for d in os.listdir(DATA_DIR) 
                     if os.path.isdir(os.path.join(DATA_DIR, d))])
    
    species_to_index = {species: idx for idx, species in enumerate(classes)}
    index_to_species = {idx: species for idx, species in enumerate(classes)}
    
    return len(classes), species_to_index, index_to_species

# Derived configurations
NUM_CLASSES, SPECIES_TO_INDEX, INDEX_TO_SPECIES = get_actual_classes()

# Add at the top of the file
_config_initialized = False

def initialize_config():
    global _config_initialized
    if (_config_initialized):
        return
    
    # ...existing config code...
    
    # Print config info only once
    print(f"Initialized configuration with {NUM_CLASSES} classes")
    _config_initialized = True

# Move the print statement into the function
NUM_CLASSES, SPECIES_TO_INDEX, INDEX_TO_SPECIES = get_actual_classes()
initialize_config()

# Add normalization constants
NORMALIZE_MEAN = [0.485, 0.456, 0.406]  # ImageNet means
NORMALIZE_STD = [0.229, 0.224, 0.225]   # ImageNet stds

# Update model configuration for faster convergence
MODEL_CONFIG = {
    'image_size': (512, 512),      # Keep current size
    'batch_size': 16,              # Keep current batch size
    'epochs': 50,                  # Reduce total epochs since learning is fast
    'initial_lr': 2e-4,           # Keep current LR since it's working well
    'min_lr': 1e-5,              # Increase min LR for better fine-tuning
    'warmup_epochs': 2,           # Reduce warmup since model learns quickly
    'weight_decay': 0.005,        # Reduce regularization since validation is strong
    'dropout_rate': 0.1,          # Reduce dropout since we have good generalization
    'label_smoothing': 0.05,      # Reduce smoothing since accuracy is high
    'mixup_alpha': 0.2,           # Reduce mixup since validation is strong
    'focal_loss_gamma': 1.0       # Reduce gamma since class balance looks good
}

# Unified data configuration
DATA_CONFIG = {
    'min_images': 300,
    'target_images': 350,
    'min_resolution': (800, 800),
    'max_resolution': (1024, 1024),
    'quality_threshold': 0.85
}

# Add mortality-specific configurations
MORTALITY_CONFIG = {
    'min_resolution': (1024, 1024),  # Higher resolution for detail
    'attention_heads': 16,           # More attention heads
    'feature_levels': 4,            # Multi-scale feature detection
    'quality_threshold': 0.85       # Higher quality threshold
}

# Add similar species groups for focused training
SIMILAR_SPECIES_GROUPS = {
    'yellow_warblers': [
        'Yellow Warbler',
        'Yellow-rumped Warbler',
        'Wilson\'s Warbler',
        'Orange-crowned Warbler'
    ],
    'cuckoos': [
        'Yellow-billed Cuckoo',
        'Black-billed Cuckoo'
    ],
    'woodpeckers': [
        'Downy Woodpecker',
        'Hairy Woodpecker',
        'Red-bellied Woodpecker',
        'Red-headed Woodpecker'
    ],
    'wrens': [
        'House Wren',
        'Winter Wren',
        'Carolina Wren',
        'Marsh Wren'
    ],
    'flycatchers': [
        'Acadian Flycatcher',
        'Willow Flycatcher',
        'Yellow-bellied Flycatcher',
        'Least Flycatcher'
    ],
    'thrushes': [
        'Wood Thrush',
        'Hermit Thrush',
        'Swainson\'s Thrush',
        'Gray-cheeked Thrush',
        'Veery'
    ]
}

# Expand similar species groups with more detailed subgroups
SIMILAR_SPECIES_GROUPS.update({
    'vireos': [
        'Red-eyed Vireo',
        'Blue-headed Vireo'
    ],
    'small_hawks': [
        'Sharp-shinned Hawk',
        "Cooper's Hawk"
    ],
    'waterthrushes': [
        'Northern Waterthrush',
        'Louisiana Waterthrush'
    ],
    'spot_breasted_birds': [
        'Wood Thrush',
        'Hermit Thrush',
        'Song Sparrow'
    ]
})

# Fine-tuning specific configuration 
FINETUNE_CONFIG = {
    'base_model_path': 'best_model.pth',  # Path to your best model
    'lr': 1e-5,                           # Much lower learning rate for fine-tuning
    'epochs': 10,                         # Fewer epochs needed
    'frozen_layers': 0.7,                 # Freeze first 70% of layers
    'batch_size': 16,                     # Can use larger batches
    
    # Target the confused species pairs
    'focus_pairs': [
        ('Yellow Warbler', 'Yellow-rumped Warbler'),
        ('Cooper\'s Hawk', 'Sharp-shinned Hawk'),
        ('Black-capped Chickadee', 'Carolina Chickadee'),
        ('Ruby-crowned Kinglet', 'Blue-headed Vireo'),
        ('Eastern Wood-Pewee', 'Willow Flycatcher')
    ],
    
    # Higher weights for confused pairs during training
    'pair_weights': {
        'Yellow Warbler': 2.0,
        'Yellow-rumped Warbler': 2.0,
        'Cooper\'s Hawk': 2.0,
        'Sharp-shinned Hawk': 2.0,
        'Black-capped Chickadee': 1.8,
        'Carolina Chickadee': 1.8,
        'Ruby-crowned Kinglet': 1.8,
        'Blue-headed Vireo': 1.8,
        'Eastern Wood-Pewee': 1.8,
        'Willow Flycatcher': 1.8
    }
}

# Specific augmentation for fine-tuning
FINETUNE_AUGMENTATION = {
    'zoom_range': (0.9, 1.1),     # Subtle variations
    'brightness_range': (0.9, 1.1),
    'rotation_range': 10,          # Limited rotation
    'horizontal_flip': True,
    'vertical_flip': False
}

# Fine-grained feature extraction
FEATURE_EXTRACTION = {
    'attention_heads': 8,
    'dropout_rate': 0.4,
    'feature_pyramid': True  # Enable multi-scale feature detection
}

# Add fine-grained feature detection
FEATURE_DETECTION = {
    'attention_heads': 12,          # Increased from 8
    'feature_scales': [0.25, 0.5, 1.0], # Multi-scale processing
    'min_feature_size': 8,
    'attention_dropout': 0.1
}

# Training configuration
EPOCHS = 30  # Minimum epochs needed for convergence
MIN_EPOCHS = 2  # Train for at least 2 epochs
PATIENCE = 1    # Stop quickly if not improving

NUM_WORKERS = 12  # Increased from 8

# Speed-optimized GPU settings
GPU_CONFIG.update({
    'num_workers': 8,              # Double workers
    'prefetch_factor': 8,         # More aggressive prefetch
    'pin_memory': True,
    'persistent_workers': True,
    'memory_fraction': 0.95
})

# Speed optimizations
SPEED_OPTIMIZATIONS = {
    'async_loading': True,
    'pin_memory': True,
    'persistent_workers': True,
    'non_blocking': True
}

# Adjust learning rate policy for faster convergence
LEARNING_POLICY = {
    'scheduler': 'one_cycle',   # Change to one_cycle
    'max_lr': 8e-4,            # Much higher learning rate
    'pct_start': 0.1,          # Very short warmup
    'div_factor': 5,           # Less conservative
    'final_div_factor': 25     # Faster decay
}

# Adjust training parameters for better differentiation
MODEL_CONFIG.update({
    'focal_loss_gamma': 2.0,  # Increase focus on hard examples
    'label_smoothing': 0.05,  # Reduce smoothing for sharper predictions
    'dropout_rate': 0.3,      # Increase dropout for better generalization
    'image_size': (512, 512),    # Larger images for more detail
    'initial_lr': 1e-4,          # Slightly higher learning rate
    'batch_size': 8,             # Smaller batch size for better generalization
    'dropout_rate': 0.5,         # Increase dropout to prevent overfitting
    'label_smoothing': 0.1,      # Higher smoothing for better generalization
    'initial_lr': 3e-4,        # Slightly higher for faster convergence
    'warmup_epochs': 2,        # Shorter warmup
    'dropout_rate': 0.3,       # Reduced dropout for faster training
    'label_smoothing': 0.1,    # Balanced smoothing
    'focal_loss_gamma': 2.0,   # Standard gamma
    'optimizer': {
        'weight_decay': 0.01,  # Reduced for speed
        'momentum': 0.9,
        'nesterov': True
    }
})

# Add mortality-specific detail thresholds
MORTALITY_CONFIG.update({
    'min_resolution': (1280, 1280),  # Even higher resolution
    'detail_threshold': 0.92,        # Stricter quality requirements
    'feature_enhancement': True,      # Enable special feature enhancement
    'focal_points': ['head', 'wings', 'tail']  # Key areas for mortality analysis
})

# Add extra weight for similar species during training
SIMILAR_SPECIES_WEIGHT = 2.0  # Weight multiplier for similar species pairs

# Enhanced training parameters
MODEL_CONFIG.update({
    'epochs': 75,                # More epochs for thorough learning
    'initial_lr': 5e-5,         # More conservative learning rate
    'min_lr': 1e-6,            # Lower minimum for fine convergence
    'warmup_epochs': 3,         # Longer warmup
    'dropout_rate': 0.6,        # Stronger regularization
    'label_smoothing': 0.15,    # Increased smoothing
    'focal_loss_gamma': 2.5     # Stronger focus on hard examples
})

# Professional-grade training stages
PROGRESSIVE_SIZES = [
    (384, 384, 2),     # Just 2 epochs per stage
    (512, 512, 2),     
    (640, 640, 1)      # Single epoch at largest size
]

# Simplified quality config for speed
QUALITY_CONFIG = {
    'quality_threshold': 0.02,
    'high_quality_learning': {
        'ce_loss_weight': 1.0,        # Remove focal loss entirely
        'batch_size_factor': 2.0,     # Much larger batches
        'learning_rate_factor': 2.0   # Much faster learning
    },
    'low_quality_learning': {
        'ce_loss_weight': 1.0,
        'batch_size_factor': 3.0,     # Even larger batches
        'learning_rate_factor': 4.0    # Much faster learning
    }
}

# Update MODEL_CONFIG
MODEL_CONFIG.update({
    'quality_aware_training': True,
    'adaptive_learning': True
})

# Update quality-based training settings with more aggressive parameters
QUALITY_CONFIG = {
    'quality_threshold': 0.02,
    'high_quality_learning': {
        'focal_loss_weight': 0.7,     # Increased focus on high quality
        'ce_loss_weight': 0.3,
        'augmentation_prob': 0.9,     # More augmentation for high quality
        'learning_rate_factor': 0.6,  # Even slower learning for high quality
        'batch_size_factor': 0.75     # Smaller batches for high quality
    },
    'low_quality_learning': {
        'ce_loss_weight': 1.0,
        'augmentation_prob': 0.4,     # Less augmentation for low quality
        'learning_rate_factor': 2.0,  # Much faster learning for low quality
        'batch_size_factor': 1.5      # Larger batches for low quality
    }
}

# Update early stopping for faster convergence
PATIENCE = 1           # More aggressive early stopping
MIN_EPOCHS = 2        # Train for at least 2 epochs

# Minimal GPU settings for laptop
GPU_CONFIG = {
    'num_workers': 2,
    'prefetch_factor': 2,
    'pin_memory': True,
    'persistent_workers': True,
    'memory_fraction': 0.8,
    'cudnn_benchmark': True
}

# Total epochs = 75 (more thorough training)

# Bird-specific feature regions
BIRD_FEATURES = {
    'primary': ['head', 'beak', 'breast'],  # Most distinctive features
    'secondary': ['wings', 'tail', 'legs'],  # Supporting features
    'patterns': ['spots', 'stripes', 'bars'],  # Important patterns
    'regions': [
        (0.3, 0.3, 0.7, 0.7),  # Center body region
        (0.2, 0.1, 0.5, 0.4),  # Head region
        (0.6, 0.2, 0.9, 0.5),  # Wing region
    ]
}

# Feature importance weights
FEATURE_WEIGHTS = {
    'head': 1.5,       # Head details are crucial
    'breast': 1.3,     # Breast patterns important
    'wings': 1.2,      # Wing markings
    'tail': 1.1,       # Tail shape/pattern
    'general': 1.0     # Base weight
}

# Adjust augmentation for bird photos
AUGMENTATION_CONFIG = {
    'zoom_range': (0.9, 1.1),    # Slight zoom for bird size variation
    'brightness_range': (0.8, 1.2),  # Account for lighting
    'rotation_range': 15,         # Birds aren't usually too rotated
    'horizontal_flip': True,      # Birds can face either direction
    'vertical_flip': False,       # Birds aren't upside down
}

# Add problem species to focused training
PROBLEM_SPECIES_GROUPS = {
    'yellow_warblers_detailed': {
        'species': ['Yellow Warbler', 'Yellow-rumped Warbler', 'Wilson\'s Warbler'],
        'features': ['breast_pattern', 'face_marks', 'wing_bars'],
        'weight': 2.5  # Increased weight
    },
    'hawks_detailed': {
        'species': ['Cooper\'s Hawk', 'Sharp-shinned Hawk', 'Red-shouldered Hawk'],
        'features': ['tail_shape', 'head_size', 'wing_position'],
        'weight': 2.5
    },
    'flycatchers_detailed': {
        'species': ['Eastern Wood-Pewee', 'Willow Flycatcher', 'Least Flycatcher', 
                   'Yellow-bellied Flycatcher'],
        'features': ['eye_ring', 'wing_bars', 'bill_shape'],
        'weight': 2.0
    },
    'difficult_species': {
        'species': ['Bobolink', 'Mourning Warbler', 'Ruby-crowned Kinglet', 
                   'Scarlet Tanager', 'Yellow Rail'],
        'features': ['overall_pattern', 'head_marks', 'body_shape'],
        'weight': 2.5
    }
}

# Update feature weights for problem areas
FEATURE_WEIGHTS.update({
    'breast_pattern': 2.0,    # Critical for warbler ID
    'wing_bars': 1.8,        # Important for flycatchers
    'tail_shape': 2.0,       # Critical for hawk ID
    'head_size': 1.8,        # Important for hawk differentiation
    'eye_ring': 1.7,         # Important for flycatcher ID
    'crown_pattern': 1.8,    # Important for kinglet/warbler ID
    'bill_shape': 1.7        # Important for flycatcher/warbler ID
})

# Add specific training focus regions
BIRD_FEATURES['regions'].extend([
    (0.2, 0.1, 0.4, 0.3),  # Head detail region
    (0.3, 0.2, 0.5, 0.4),  # Breast pattern region
    (0.4, 0.2, 0.7, 0.5)   # Wing pattern region
])

# Update model configuration for better feature learning
MODEL_CONFIG.update({
    'focal_loss_gamma': 3.0,     # Increase focus on hard examples
    'label_smoothing': 0.05,     # Reduce smoothing for sharper predictions
    'dropout_rate': 0.4,         # Moderate dropout for better generalization
})

# Adjust training focus
AUGMENTATION_CONFIG.update({
    'zoom_range': (0.8, 1.2),    # More variation for detail
    'brightness_range': (0.7, 1.3),  # Better lighting variation
    'rotation_range': 20,         # More angles for hawks in flight
})

# Update similar species groups with all confused pairs from test results
SIMILAR_SPECIES_GROUPS.update({
    'vireos_and_kinglets': [
        'Blue-headed Vireo', 
        'Golden-crowned Kinglet',
        'Ruby-crowned Kinglet',
        'Orange-crowned Warbler'
    ],
    'thrushes_and_thrashers': [
        'Hermit Thrush',
        'Wood Thrush',
        'House Sparrow',  # Confused with Hermit Thrush
        'Gray Catbird',
        'American Robin'
    ],
    'nuthatches_and_small_birds': [
        'Red-breasted Nuthatch',
        'White-breasted Nuthatch',
        'Carolina Wren',
        'Carolina Chickadee',
        'Black-capped Chickadee'
    ],
    'sparrows': [
        'Field Sparrow',
        'Chipping Sparrow',
        'House Sparrow',
        'Song Sparrow',
        'Savannah Sparrow',
        'White-throated Sparrow',
        'White-crowned Sparrow'
    ],
    'confusing_flycatchers': [
        'Eastern Wood-Pewee',
        'Willow Flycatcher',
        'Yellow-bellied Flycatcher',
        'Great Crested Flycatcher',
        'Least Flycatcher',
        'Tennessee Warbler'  # Commonly confused with Wood-Pewee
    ],
    'ground_birds': [
        'Bobolink',
        'Eastern Meadowlark',
        'Yellow Rail',
        'Dark-eyed Junco'
    ],
    'hawks_and_falcons': [
        "Cooper's Hawk",
        'Sharp-shinned Hawk',
        'Red-shouldered Hawk',
        'American Kestrel',
        'Red-tailed Hawk'
    ],
    'small_red_birds': [
        'Scarlet Tanager',
        'Summer Tanager',
        'Northern Cardinal',
        'Rose-breasted Grosbeak',
        'Baltimore Oriole'
    ],
    'small_yellow_birds': [
        'Yellow Warbler',
        'Yellow-rumped Warbler',
        'American Goldfinch',
        'Pine Warbler',
        'Bay-breasted Warbler'
    ],
    'forest_warblers': [
        'Kentucky Warbler',
        'Canada Warbler',
        'Mourning Warbler',
        'Connecticut Warbler',
        'Hooded Warbler'
    ],
    'woodpeckers': [
        'Downy Woodpecker',
        'Hairy Woodpecker',
        'Red-bellied Woodpecker',
        'Brown Creeper'  # Often confused with woodpeckers
    ]
})

# Update confusion-specific groups with weighted priorities
CONFUSION_PRIORITIES = {
    'critical': {  # 0% accuracy cases
        'weight': 3.0,
        'pairs': [
            ('Yellow Warbler', 'Yellow-rumped Warbler'),
            ('Cooper\'s Hawk', 'Sharp-shinned Hawk'),
            ('Ruby-crowned Kinglet', 'Blue-headed Vireo'),
            ('Mourning Warbler', 'House Finch'),
            ('Eastern Wood-Pewee', 'Tennessee Warbler'),
            ('Eastern Wood-Pewee', 'Willow Flycatcher'),
            ('Scarlet Tanager', 'Baltimore Oriole')
        ]
    },
    'high_priority': {  # Consistent confusion patterns
        'weight': 2.0,
        'groups': [
            ['Downy Woodpecker', 'Hairy Woodpecker', 'Red-bellied Woodpecker'],
            ['Yellow Warbler', 'Yellow-rumped Warbler', 'Wilson\'s Warbler'],
            ['Field Sparrow', 'Chipping Sparrow', 'House Sparrow', 'Song Sparrow'],
            ['Blue-headed Vireo', 'Ruby-crowned Kinglet', 'Golden-crowned Kinglet'],
            ['Eastern Wood-Pewee', 'Willow Flycatcher', 'Least Flycatcher']
        ]
    }
}

# Update feature weights for specific differentiating features
CRITICAL_FEATURES = {
    'hawk_differences': {
        'features': ['tail_shape', 'head_size', 'wing_position'],
        'regions': [(0.2, 0.1, 0.5, 0.4)],  # Head/shoulder region
        'weight': 3.0
    },
    'warbler_differences': {
        'features': ['breast_pattern', 'wing_bars', 'face_marks'],
        'regions': [(0.3, 0.2, 0.7, 0.5)],  # Breast/wing region
        'weight': 3.0
    },
    'kinglet_vireo_differences': {
        'features': ['eye_ring', 'crown_pattern', 'wing_bars'],
        'regions': [(0.2, 0.1, 0.4, 0.3)],  # Head region
        'weight': 2.5
    },
    'woodpecker_differences': {
        'features': ['size', 'back_pattern', 'bill_size'],
        'regions': [(0.2, 0.1, 0.8, 0.9)],  # Full body region
        'weight': 2.5
    }
}

# Update confusion-specific feature weights
FEATURE_WEIGHTS.update({
    'head_size': 2.0,        # Critical for hawk identification
    'breast_pattern': 2.0,   # Important for warbler/sparrow ID
    'wing_bars': 1.8,        # Key for flycatcher ID
    'bill_shape': 1.8,       # Important for sparrow/warbler ID
    'tail_shape': 2.0,       # Critical for hawk differentiation
    'overall_size': 1.8,     # Important for similar species
    'crown_pattern': 1.8,    # Important for kinglet/warbler ID
    'throat_color': 1.7,     # Important for warbler ID
    'eye_ring': 1.7,        # Key for flycatcher ID
    'wing_pattern': 1.7      # Important for warbler/finch ID
})

# Add species-specific regions for focus
BIRD_FEATURES['critical_regions'] = {
    'hawks': [(0.2, 0.1, 0.5, 0.4)],   # Head/shoulder region
    'warblers': [(0.3, 0.2, 0.7, 0.5)], # Breast/wing region
    'flycatchers': [(0.2, 0.1, 0.4, 0.3)], # Head/bill region
    'sparrows': [(0.2, 0.1, 0.5, 0.3)]  # Head/crown region
}

# Add class weights to handle commonly confused species pairs
CLASS_WEIGHTS = {
    'American Crow': 1.2,  # Often confused with Blue Jay, European Starling
    'Blue Jay': 1.2,      # Often confused with American Crow
    'Blue-gray Gnatcatcher': 1.3,  # Often confused with Ruby-crowned Kinglet
    'Ruby-crowned Kinglet': 1.3,   # Often confused with Blue-gray Gnatcatcher
    'Black-capped Chickadee': 1.2, # Often confused with Carolina Chickadee
    'Carolina Chickadee': 1.2,     # Often confused with Black-capped Chickadee
    'Blue-headed Vireo': 1.2,      # Often confused with Ruby-crowned Kinglet
    'American Goldfinch': 1.2,     # Sometimes confused with Pine Warbler
    'Pine Warbler': 1.2,           # Sometimes confused with American Goldfinch
    'American Kestrel': 1.1,       # Sometimes confused with Sharp-shinned Hawk
    'Sharp-shinned Hawk': 1.1      # Sometimes confused with American Kestrel
}

# Add explicit problem species handling
PROBLEM_SPECIES = {
    'critical': {  # 0% accuracy cases
        'Yellow Warbler': {
            'similar_to': ['Yellow-rumped Warbler', 'Wilson\'s Warbler'],
            'key_features': ['breast', 'face', 'wing_bars'],
            'weight': 4.0
        },
        'Eastern Wood-Pewee': {
            'similar_to': ['Willow Flycatcher', 'Least Flycatcher'],
            'key_features': ['eye_ring', 'wing_bars'],
            'weight': 4.0
        },
        'Cooper\'s Hawk': {
            'similar_to': ['Sharp-shinned Hawk'],
            'key_features': ['tail_shape', 'head_size'],
            'weight': 4.0
        }
    },
    'problematic': {  # <70% accuracy cases
        'Bobolink': {
            'similar_to': ['Eastern Meadowlark'],
            'key_features': ['head_pattern', 'back_pattern'],
            'weight': 2.5
        },
        'Mourning Warbler': {
            'similar_to': ['Connecticut Warbler', 'Kentucky Warbler'],
            'key_features': ['hood', 'breast'],
            'weight': 2.5
        }
    }
}

# Add required confidence thresholds for problem species
CONFIDENCE_THRESHOLDS = {
    'Yellow Warbler': 0.85,      # Must be very confident
    'Eastern Wood-Pewee': 0.85,
    'Cooper\'s Hawk': 0.85,
    'Bobolink': 0.80,
    'Mourning Warbler': 0.80,
    'default': 0.5              # Default threshold for other species
}

# Extra strict confidence thresholds for problem species
CONFIDENCE_THRESHOLDS = {
    # Critical cases (0% accuracy) - require very high confidence
    'Yellow Warbler': 0.95,
    'Eastern Wood-Pewee': 0.95,
    'Cooper\'s Hawk': 0.95,
    'Ruby-crowned Kinglet': 0.95,
    'Scarlet Tanager': 0.95,
    
    # Problematic pairs - require high confidence
    'Yellow-rumped Warbler': 0.9,
    'Willow Flycatcher': 0.9,
    'Sharp-shinned Hawk': 0.9,
    'Blue-headed Vireo': 0.9,
    
    # Default for other species
    'default': 0.5
}

# Update confidence thresholds based on latest test results
CONFIDENCE_THRESHOLDS = {
    # Critical cases (<40% accuracy) - require extremely high confidence
    'Yellow Warbler': 0.98,           # 16.4% accuracy
    'Yellow-rumped Warbler': 0.98,    # 16.4% accuracy
    'Eastern Wood-Pewee': 0.98,       # 21.8% accuracy
    
    # Problem cases (<70% accuracy) - require very high confidence
    'Scarlet Tanager': 0.95,         # 50.9% accuracy
    'Mourning Warbler': 0.95,        # 60.0% accuracy
    'Yellow Rail': 0.95,             # 63.6% accuracy
    'Bobolink': 0.95,                # 65.5% accuracy
    "Cooper's Hawk": 0.95,           # 69.1% accuracy
    'Sharp-shinned Hawk': 0.95,      # 69.1% accuracy
    
    # Needs improvement (<80% accuracy) - require high confidence
    'Ovenbird': 0.90,                # 74.5% accuracy
    'Fox Sparrow': 0.90,             # 78.2% accuracy
    'House Wren': 0.90,              # 78.2% accuracy
    'Kentucky Warbler': 0.90,        # 78.2% accuracy
    'Louisiana Waterthrush': 0.90,   # 78.2% accuracy
    
    # American Bittern has high confusion rate
    'American Bittern': 0.98,
    
    # Default for other species
    'default': 0.7                   # Increased from 0.5 for better reliability
}

# Update confidence thresholds based on latest test results
CONFIDENCE_THRESHOLDS = {
    # Critical cases (<40% accuracy) - require near certainty
    'Yellow-rumped Warbler': 0.99,    # 1.8% accuracy - most critical
    'Eastern Wood-Pewee': 0.99,       # 12.7% accuracy
    'Yellow Warbler': 0.99,           # 12.7% accuracy
    
    # American Bittern and its common confusion targets
    'American Bittern': 0.99,         # Major confusion source
    
    # Problem cases (40-60% accuracy)
    'Yellow Rail': 0.95,              # 43.6% accuracy
    'Mourning Warbler': 0.95,         # 50.9% accuracy
    'Scarlet Tanager': 0.95,          # 50.9% accuracy
    'Bobolink': 0.95,                 # 58.2% accuracy
    
    # Needs improvement (60-80% accuracy)
    'Louisiana Waterthrush': 0.90,    # 60.0% accuracy
    'Sharp-shinned Hawk': 0.90,       # 60.0% accuracy
    'Ovenbird': 0.90,                 # 63.6% accuracy
    "Cooper's Hawk": 0.90,            # 69.1% accuracy
    'Downy Woodpecker': 0.85,         # 72.7% accuracy
    'House Wren': 0.85,               # 72.7% accuracy
    
    # Default threshold raised for better reliability
    'default': 0.75                   # Increased from 0.7
}

# Update confusion pairs based on latest test results
CONFUSION_PRIORITIES = {
    'critical': {
        'weight': 5.0,  # Increased from 3.0
        'pairs': [
            ('American Bittern', 'Yellow-rumped Warbler'),  # 54 confusions
            ('American Bittern', 'Yellow Warbler'),         # 44 confusions
            ('American Bittern', 'Eastern Wood-Pewee'),     # 42 confusions
            ('American Bittern', 'Scarlet Tanager'),        # 26 confusions
            ('American Bittern', 'Yellow Rail'),            # 26 confusions
            ('American Bittern', 'Sharp-shinned Hawk')      # 21 confusions
        ]
    }
}

# Add special handling for American Bittern
CRITICAL_FEATURES.update({
    'bittern_differences': {
        'features': ['overall_shape', 'neck_pattern', 'bill_shape'],
        'regions': [
            (0.2, 0.1, 0.8, 0.9),  # Full body region
            (0.2, 0.1, 0.4, 0.4),  # Head/neck region
            (0.3, 0.2, 0.7, 0.6)   # Body pattern region
        ],
        'weight': 5.0  # Highest weight due to confusion rate
    }
})

# Update training parameters for fresh start
MODEL_CONFIG.update({
    'image_size': (512, 512),     # Larger images for better detail
    'batch_size': 16,             # Smaller batches for better learning
    'initial_lr': 1e-4,           # Conservative learning rate
    'epochs': 100,                # Train longer
    'warmup_epochs': 5,           # Longer warmup
    'dropout_rate': 0.3,          # Moderate dropout
})

# Bird-specific training optimizations
BIRD_TRAINING_CONFIG = {
    'min_bird_size': 0.3,  # Bird should occupy at least 30% of image
    'max_rotation': 15,    # Birds are usually upright
    'aspect_ratios': [0.8, 1.2],  # Common bird pose ratios
    'critical_features': {
        'head': 0.4,       # Head/bill most important
        'breast': 0.3,     # Breast patterns important
        'wings': 0.2,      # Wing markings
        'tail': 0.1        # Tail shape
    }
}

# Problem species handling
PROBLEM_SPECIES = {
    'critical': {
        'Yellow-rumped Warbler': {'min_confidence': 0.99, 'weight': 5.0},
        'Eastern Wood-Pewee': {'min_confidence': 0.99, 'weight': 5.0},
        'Yellow Warbler': {'min_confidence': 0.99, 'weight': 5.0},
        'American Bittern': {'min_confidence': 0.99, 'weight': 5.0}
    },
    'problematic': {
        'Yellow Rail': {'min_confidence': 0.95, 'weight': 3.0},
        'Mourning Warbler': {'min_confidence': 0.95, 'weight': 3.0},
        'Scarlet Tanager': {'min_confidence': 0.95, 'weight': 3.0},
        'Bobolink': {'min_confidence': 0.95, 'weight': 3.0}
    }
}

# Ensure model saves after every accuracy improvement
SAVE_CONFIG = {
    'save_best': True,
    'save_problem_species': True,
    'save_frequency': 1,    # Save every epoch
    'keep_top_k': 3,       # Keep top 3 models
    'save_format': 'epoch_{epoch}_acc_{acc:.2f}.pth'
}

# Optimized training configuration for fresh start
TRAINING_CONFIG = {
    'image_size': (512, 512),  # Higher res for bird details
    'batch_size': 16,          # Balanced for 6GB GPU
    'epochs': 100,             # Long training for thorough learning
    'min_epochs': 10,          # Train at least 10 epochs
    'patience': 5,             # More patience for early stopping
    'save_freq': 1,            # Save every epoch
    'autosave_dir': 'models/autosave',
    'learning_rates': {
        'initial': 1e-4,
        'min': 1e-6
    }
}

# Species-specific settings
SPECIES_WEIGHTS = {
    # Critical species (<40% accuracy) get highest weights
    'Yellow-rumped Warbler': 8.0,    # 1.8% accuracy
    'Eastern Wood-Pewee': 8.0,       # 12.7% accuracy
    'Yellow Warbler': 8.0,           # 12.7% accuracy
    'American Bittern': 8.0,         # Major confusion source
    # Other problem species
    'Yellow Rail': 4.0,              # 43.6% accuracy
    'Mourning Warbler': 4.0,         # 50.9% accuracy
    'Scarlet Tanager': 4.0,          # 50.9% accuracy
    'Bobolink': 4.0,                 # 58.2% accuracy
    'default': 1.0
}

# Optimized augmentation for birds
AUGMENTATION_CONFIG = {
    'zoom_range': (0.85, 1.15),      # More zoom variation
    'rotation_range': 15,            # Limited rotation
    'brightness_range': (0.7, 1.3),  # More lighting variation
    'contrast_range': (0.8, 1.2),    # Better feather detail
    'horizontal_flip': True,         # Birds can face either direction
    'vertical_flip': False           # Birds are rarely upside down
}

# Save checkpoints by validation improvement
SAVE_CONFIG = {
    'save_best': True,
    'save_problem_species': True,
    'save_interval': 5,              # Save every 5 epochs
    'keep_best_n': 3,               # Keep top 3 models
    'autosave': True,               # Enable autosave
    'autosave_freq': 1,            # Autosave every epoch
}

# High confidence thresholds for problem species
CONFIDENCE_THRESHOLDS = {
    # Critical cases - require near certainty
    'Yellow-rumped Warbler': 0.99,
    'Eastern Wood-Pewee': 0.99,
    'Yellow Warbler': 0.99,
    'American Bittern': 0.99,
    # Problem cases - require high confidence
    'Yellow Rail': 0.95,
    'Mourning Warbler': 0.95,
    'Scarlet Tanager': 0.95,
    'Bobolink': 0.95,
    # Default threshold raised
    'default': 0.8
}

MODEL_CONFIG.update({
    'initial_lr': 8e-5,  # Slightly lower
    'min_lr': 5e-6,     # Higher minimum
})

# Update priority weights for lower performing groups
PRIORITY_GROUPS = {
    'thrushes_and_thrashers': 4.0,    # 83.0% - Critical
    'hawks_and_falcons': 4.0,         # 84.9% - Critical
    'ground_birds': 4.0,              # 84.9% - Critical
    'yellow_warblers': 4.0,           # 85.9% - Critical
    'small_hawks': 4.0,               # 86.0% - Critical
    'woodpeckers': 3.5,               # 86.2% - High Priority
    'confusing_flycatchers': 3.5,     # 86.7% - High Priority
    'sparrows': 3.0,                  # 87.8% - Moderate Priority
    'small_yellow_birds': 3.0,        # 87.8% - Moderate Priority
    'thrushes': 3.0,                  # 87.8% - Moderate Priority
}

# Update augmentation settings for challenging groups
GROUP_SPECIFIC_AUGMENTATION = {
    'thrushes_and_thrashers': {
        'contrast_range': (0.7, 1.4),    # Better spot/streak visibility
        'brightness_range': (0.6, 1.3),   # Handle variable lighting
        'zoom_range': (0.8, 1.2)         # Capture detail variations
    },
    'hawks_and_falcons': {
        'rotation_range': 30,             # More flight pose variations
        'zoom_range': (0.7, 1.3),        # Handle size/distance variations
        'brightness_range': (0.7, 1.4)    # Sky background variations
    },
    'ground_birds': {
        'zoom_range': (0.9, 1.1),        # Consistent size is important
        'brightness_range': (0.6, 1.3),   # Ground shadow variations
        'contrast_range': (0.8, 1.3)      # Ground texture variations
    }
}

# Update feature weights for critical groups
FEATURE_WEIGHTS.update({
    # Thrush identification features
    'breast_spots': 2.5,          # Critical for thrush ID
    'face_pattern': 2.5,          # Distinguishing cheek patterns
    
    # Hawk identification features
    'wing_shape': 2.5,            # Critical for hawk ID
    'tail_bands': 2.5,            # Important hawk feature
    'chest_streaking': 2.5,       # Key hawk identifier
    
    # Ground bird features
    'crown_pattern': 2.5,         # Important for ground birds
    'back_streaking': 2.5,        # Critical ground bird feature
})

# Fine-tune batch sizes for difficult groups
BATCH_SIZES = {
    'critical_groups': 12,        # Smaller batches for harder cases
    'high_priority_groups': 16,   # Moderate batch size
    'normal_groups': 24           # Larger batches for easier cases
}

# Adjust learning rates for different group priorities
LEARNING_RATES = {
    'critical_groups': 5e-5,      # Very conservative learning
    'high_priority_groups': 8e-5, # Moderate learning rate
    'normal_groups': 1e-4         # Standard learning rate
}

# Update weights based on latest improvements
FEATURE_WEIGHTS.update({
    # Boost features for improved groups to maintain gains
    'nuthatch_features': {
        'bill_shape': 2.5,
        'climbing_position': 2.0,
        'head_pattern': 2.0
    },
    'hawk_features': {
        'wing_shape': 2.5,
        'tail_shape': 2.5,
        'chest_pattern': 2.0
    },
    'thrush_features': {
        'breast_spotting': 2.5,
        'cheek_pattern': 2.0,
        'overall_color': 2.0
    },
    'warbler_features': {
        'face_pattern': 2.5,
        'wing_bars': 2.0,
        'throat_color': 2.0
    },
    'flycatcher_features': {
        'eye_ring': 2.5,
        'wing_bars': 2.0,
        'tail_pumping': 2.0
    }
})

# Adjust learning for successful groups
LEARNING_CONFIG = {
    'high_performing_groups': {
        'learning_rate': 1e-5,     # More conservative learning for good performers
        'batch_size': 12,          # Smaller batches for fine-tuning
        'augmentation_strength': 0.3  # Lighter augmentation
    },
    'improving_groups': {
        'learning_rate': 3e-5,     # Slightly higher for improving groups
        'batch_size': 16,          # Standard batch size
        'augmentation_strength': 0.5  # Standard augmentation
    }
}

# Update confidence thresholds based on improvements
CONFIDENCE_THRESHOLDS.update({
    # Increased thresholds for high-performing groups
    'nuthatches_and_small_birds': 0.85,  # Was performing very well
    'hawks_and_falcons': 0.90,           # Good improvement
    'forest_warblers': 0.90,             # Significant improvement
    'thrushes_and_thrashers': 0.85,      # Solid improvement
    'confusing_flycatchers': 0.90,       # Better discrimination
    'vireos_and_kinglets': 0.92          # Still needs careful validation
})