OModeling-2.py

import os
import gc
import json
import math
import torch
import mlflow
import logging
import platform
import numpy as np
import pandas as pd
from PIL import Image
from tqdm import tqdm
import torch.nn as nn
import torch.optim as optim
from torchvision import models
import matplotlib.pyplot as plt
import torch.nn.functional as F
from sklearn.manifold import TSNE
from torchvision import transforms
from kymatio.torch import Scattering2D
from torch.utils.data import Dataset, DataLoader
from pytorch_metric_learning.miners import BatchHardMiner
from pytorch_metric_learning.losses import MultiSimilarityLoss
from torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau
from sklearn.metrics import roc_curve, auc, precision_recall_fscore_support
from typing import Dict, List, Tuple, Optional, Union, Any
from dataclasses import dataclass, asdict
import warnings
warnings.filterwarnings('ignore')

# ----------------------------
# Configuration Management
# ----------------------------
@dataclass
@dataclass
class TrainingConfig:
    
    # Model Architecture
    model_name: str = "resnet34"
    embedding_dim: int = 128
    normalize_embeddings: bool = True
    pretrained_path: Optional[str] = "../../model/pretrained_model/ResNet34.pt"
    
    # Training Hyperparameters
    batch_size: int = 512
    max_epochs: int = 20
    grad_accum_steps: int = 10
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Learning Rate Configuration
    head_lr: float = 1e-3          # Higher LR for embedding head (untrained)
    backbone_lr: float = 1e-4      # Lower LR for backbone (pretrained)
    lr_scheduler: str = "cosine"   # "cosine" or "plateau"
    weight_decay: float = 1e-4
    
    # Curriculum Learning Parameters - ADJUSTED FOR PRECISION
    curriculum_strategy: str = "progressive"  # "progressive", "exponential", "linear"
    initial_hard_ratio: float = 0.6          # Increased from 0.1 for more hard negatives early
    final_hard_ratio: float = 0.9            # Increased from 0.8 for focus on hard cases
    curriculum_warmup_epochs: int = 1        # Reduced from 2 for faster hard sample exposure
    
    # Data Augmentation
    remove_bg: bool = False
    augmentation_strength: float = 0.5        # 0.0 = no aug, 1.0 = strong aug
    
    # Loss Configuration - ADJUSTED FOR PRECISION
    multisim_alpha: float = 2.5              # Increased from 2.0 (penalize false positives more)
    multisim_beta: float = 60.0              # Increased from 50.0 (larger margin)
    multisim_base: float = 0.4               # Decreased from 0.5 (stricter similarity)
    
    # Triplet Loss Parameters - NEW
    triplet_margin: float = 1.0              # Margin for triplet loss
    triplet_weight: float = 0.3              # Weight for triplet loss component
    false_positive_penalty_weight: float = 0.3  # Extra penalty for false positives
    
    # Mining Configuration
    use_hard_mining: bool = True
    
    # Precision Focus Parameters - NEW
    target_precision: float = 0.75           # Target precision for threshold selection
    negative_weight_multiplier: float = 2.5  # How much more to weight hard negatives
    
    # Checkpoint Configuration
    run_id: Optional[str] = None
    last_epoch_weights: Optional[str] = None
    save_frequency: int = 1                   # Save every N epochs
    
    # Early Stopping
    patience: int = 15
    min_delta: float = 0.001
    
    # Logging
    log_frequency: int = 100                  # Log every N steps
    visualize_frequency: int = 1              # Visualize every N epochs
    
    tracking_uri: str = "http://127.0.0.1:5555"
    
    def __post_init__(self):
        """Validate configuration parameters."""
        assert 0.0 <= self.initial_hard_ratio <= 1.0, "Initial hard ratio must be in [0, 1]"
        assert 0.0 <= self.final_hard_ratio <= 1.0, "Final hard ratio must be in [0, 1]"
        assert self.curriculum_strategy in ["progressive", "exponential", "linear"]
        assert self.lr_scheduler in ["cosine", "plateau"]
        assert 0.0 <= self.target_precision <= 1.0, "Target precision must be in [0, 1]"
		
		
# Global configuration
CONFIG = TrainingConfig()

# ----------------------------
# MLFlow Setup
# ----------------------------
class MLFlowManager:
	"""Centralized MLflow management for experiment tracking."""
	
	def __init__(self, tracking_uri: str = "http://127.0.0.1:5555"):
		mlflow.set_tracking_uri(tracking_uri)
		self.experiment_name = "Signature Verification - Advanced Training"
		self._setup_experiment()
	
	def _setup_experiment(self):
		"""Setup MLflow experiment."""
		try:
			self.experiment_id = mlflow.create_experiment(self.experiment_name)
		except:
			self.experiment_id = mlflow.get_experiment_by_name(self.experiment_name).experiment_id
	
	def start_run(self, run_id: Optional[str] = None):
		"""Start MLflow run with configuration logging."""
		return mlflow.start_run(run_id=run_id, experiment_id=self.experiment_id)
	
	def log_config(self, config: TrainingConfig):
		"""Log training configuration."""
		config_dict = asdict(config)
		mlflow.log_params(config_dict)

# ----------------------------
# Curriculum Learning Manager
# ----------------------------
class CurriculumLearningManager:
    """Advanced curriculum learning for both hard positives and hard negatives."""
    
    def __init__(self, config: TrainingConfig):
        self.config = config
        self.current_epoch = 0
        
    def get_hard_ratio(self, epoch: int) -> float:
        """Get hard negative ratio (forgeries) for current epoch."""
        if epoch < self.config.curriculum_warmup_epochs:
            return self.config.initial_hard_ratio
        
        # Target: reach final_hard_ratio by max_epochs // 2
        target_epoch = max(self.config.max_epochs // 2, self.config.curriculum_warmup_epochs + 3)
        
        if epoch >= target_epoch:
            return self.config.final_hard_ratio
        
        # Aggressive progression to reach target by mid-training
        progress = (epoch - self.config.curriculum_warmup_epochs) / (target_epoch - self.config.curriculum_warmup_epochs)
        
        initial = self.config.initial_hard_ratio
        final = self.config.final_hard_ratio
        
        if self.config.curriculum_strategy == "progressive":
            # Very aggressive: exponential growth early, then plateau
            ratio = initial + (final - initial) * (progress ** 0.5)
        elif self.config.curriculum_strategy == "exponential":
            ratio = initial + (final - initial) * (progress ** 0.3)
        else:  # linear
            ratio = initial + (final - initial) * progress
        
        return min(max(ratio, 0.0), 1.0)
    
    def get_hard_positive_ratio(self, epoch: int) -> float:
        """Get hard positive ratio for current epoch - increases more gradually."""
        if epoch < self.config.curriculum_warmup_epochs:
            return 0.1  # Start with 10% hard positives
        
        # Hard positives should increase more gradually than hard negatives
        max_epochs = self.config.max_epochs
        progress = min(1.0, (epoch - self.config.curriculum_warmup_epochs) / (max_epochs - self.config.curriculum_warmup_epochs))
        
        # Target 40% hard positives by end of training
        initial_ratio = 0.1
        final_ratio = 0.4
        
        if self.config.curriculum_strategy == "progressive":
            ratio = initial_ratio + (final_ratio - initial_ratio) * (progress ** 0.7)
        else:
            ratio = initial_ratio + (final_ratio - initial_ratio) * progress
        
        return min(max(ratio, 0.0), final_ratio)

    def get_mining_difficulty(self, epoch: int) -> Dict[str, float]:
        """Adaptive mining parameters for both hard positives and negatives."""
        progress = min(1.0, epoch / self.config.max_epochs)
        
        # Separate ratios for hard positives and hard negatives
        hard_negative_ratio = self.get_hard_ratio(epoch)
        hard_positive_ratio = self.get_hard_positive_ratio(epoch)
        
        # Dynamic weights for different sample types
        hard_pos_weight = 1.0 + 2.0 * progress  # 1.0 → 3.0
        hard_neg_weight = 1.0 + 4.0 * progress  # 1.0 → 5.0 (harder negatives more important)
        
        return {
            # Margin parameters
            "margin_multiplier": 1.0 + 0.5 * progress,
            
            # Hard sample ratios
            "hard_negative_ratio": hard_negative_ratio,
            "hard_positive_ratio": hard_positive_ratio,
            "current_hard_ratio": hard_negative_ratio,  # For backward compatibility
            
            # Sample weights
            "hard_positive_weight": hard_pos_weight,
            "hard_negative_weight": hard_neg_weight,
            "semi_positive_weight": 1.0 + 1.0 * progress,
            "semi_negative_weight": 1.0 + 2.0 * progress,
            
            # Difficulty thresholds
            "difficulty_threshold": 0.05 + 0.15 * progress,
            "selectivity": 0.8 + 0.2 * progress,
            
            # Mining aggressiveness
            "mining_temperature": max(0.5, 1.0 - 0.5 * progress),  # Decreases over time
            
            # Focus balance (0 = equal focus, 1 = focus on negatives)
            "negative_focus": 0.5 + 0.3 * progress
        }
# ----------------------------
# Enhanced Dataset with Advanced Curriculum Learning
# ----------------------------
class SignatureDataset(Dataset):
	"""
	Advanced signature dataset with curriculum learning and mining statistics.
	"""
	
	def __init__(
		self, 
		folder_img: str, 
		excel_data: pd.DataFrame, 
		curriculum_manager: CurriculumLearningManager,
		transform: Optional[transforms.Compose] = None, 
		is_train: bool = True,
		config: TrainingConfig = CONFIG
	):
		self.folder_img = folder_img
		self.is_train = is_train
		self.config = config
		self.curriculum_manager = curriculum_manager
		self.transform = transform or self._default_transforms()
		self.excel_data = excel_data.reset_index(drop=True)
		self.current_epoch = 0
		
		# Data preparation
		self._handle_excel_person_ids()
		self._categorize_difficulty()
		
		# Curriculum learning data
		self.epoch_data = []
		self._prepare_epoch_data()
	def _handle_excel_person_ids(self):
		"""Properly separate genuine vs forged signature IDs with compact offset."""
		# Map genuine person IDs to 0, 1, 2, ...
		genuine_ids = pd.concat([
			self.excel_data["anchor_id"], 
			self.excel_data[self.excel_data["easy_or_hard"] == "easy"]["negative_id"]
		]).unique()
		
		self.genuine_id_mapping = {val: idx for idx, val in enumerate(genuine_ids)}
		max_genuine_id = len(genuine_ids)
		
		# Create forgery ID space with SMALLER offset (just enough to avoid collisions)
		forged_data = self.excel_data[self.excel_data["easy_or_hard"] == "hard"]
		if len(forged_data) > 0:
			unique_forged_persons = forged_data["negative_id"].unique()
			self.forgery_id_mapping = {
				val: idx + max_genuine_id + 100  # Smaller offset: 100 instead of 1000
				for idx, val in enumerate(unique_forged_persons)
			}
		else:
			self.forgery_id_mapping = {}
		
		# Apply mappings
		self.excel_data["anchor_id"] = self.excel_data["anchor_id"].map(self.genuine_id_mapping)
		
		# Handle negatives based on type
		new_negative_ids = []
		for idx, row in self.excel_data.iterrows():
			if row["easy_or_hard"] == "easy":
				# Genuine different person: use regular ID
				new_negative_ids.append(self.genuine_id_mapping[row["negative_id"]])
			else:
				# Forged signature: use offset ID to prevent clustering with genuine
				new_negative_ids.append(self.forgery_id_mapping[row["negative_id"]])
		
		self.excel_data["negative_id"] = new_negative_ids
		
		print(f"ID mapping: Genuine IDs 0-{max_genuine_id-1}, Forgery IDs {max_genuine_id+100}-{max_genuine_id+100+len(self.forgery_id_mapping)-1}")		
		
	
	def _categorize_difficulty(self):
		"""Categorize samples by difficulty if not already done."""
		if self.is_train and "easy_or_hard" in self.excel_data.columns:
			self.easy_df = self.excel_data[self.excel_data["easy_or_hard"] == "easy"]
			self.hard_df = self.excel_data[self.excel_data["easy_or_hard"] == "hard"]
		else:
			# All samples treated as medium difficulty
			self.easy_df = self.excel_data
			self.hard_df = pd.DataFrame()  # Empty hard samples
	
	def _prepare_epoch_data(self):
		"""Prepare data for current epoch based on curriculum."""
		if not self.is_train:
			# Validation data preparation with better error handling
			if "image_1_path" in self.excel_data.columns and "image_2_path" in self.excel_data.columns:
				# Standard pair format
				required_cols = ["image_1_path", "image_2_path", "label"]
				
				# Find ID columns
				id_cols = [col for col in self.excel_data.columns if "id" in col.lower()]
				if len(id_cols) >= 2:
					required_cols.extend(id_cols[-2:])  # Take last 2 ID columns
				else:
					# Create dummy IDs if none exist
					self.excel_data["dummy_id1"] = 0
					self.excel_data["dummy_id2"] = 1
					required_cols.extend(["dummy_id1", "dummy_id2"])
				
				self.epoch_data = self.excel_data[required_cols].values.tolist()
				
			else:
				# Fallback: try to use all available columns
				print(f"Warning: Expected validation columns not found. Available: {list(self.excel_data.columns)}")
				self.epoch_data = self.excel_data.values.tolist()
			
			print(f"Validation data prepared: {len(self.epoch_data)} samples")
			return
		
		# Training data preparation (unchanged)
		hard_ratio = self.curriculum_manager.get_hard_ratio(self.current_epoch)
		
		if len(self.hard_df) > 0:
			n_total = len(self.excel_data)
			n_hard = int(n_total * hard_ratio)
			n_easy = n_total - n_hard
			
			hard_sample = self.hard_df.sample(
				n=min(n_hard, len(self.hard_df)), 
				random_state=self.current_epoch,
				replace=(n_hard > len(self.hard_df))
			)
			easy_sample = self.easy_df.sample(
				n=min(n_easy, len(self.easy_df)), 
				random_state=self.current_epoch,
				replace=(n_easy > len(self.easy_df))
			)
			
			epoch_df = pd.concat([hard_sample, easy_sample]).sample(
				frac=1, random_state=self.current_epoch
			).reset_index(drop=True)
			
			print(f"Epoch {self.current_epoch}: {len(hard_sample)} hard + {len(easy_sample)} easy = {len(epoch_df)} total (target ratio: {hard_ratio:.2f})")
		else:
			epoch_df = self.excel_data.sample(
				frac=1, random_state=self.current_epoch
			).reset_index(drop=True)
		
		required_cols = ["anchor_path", "positive_path", "negative_path", "anchor_id", "negative_id"]
		missing_cols = [col for col in required_cols if col not in epoch_df.columns]
		if missing_cols:
			raise ValueError(f"Missing required training columns: {missing_cols}")
		
		self.epoch_data = epoch_df[required_cols].values.tolist()
		
	def set_epoch(self, epoch: int):
		"""Update epoch and regenerate data."""
		self.current_epoch = epoch
		self._prepare_epoch_data()
	
	def get_curriculum_stats(self) -> Dict[str, Any]:
		"""Get current curriculum learning statistics."""
		hard_ratio = self.curriculum_manager.get_hard_ratio(self.current_epoch)
		mining_params = self.curriculum_manager.get_mining_difficulty(self.current_epoch)
		
		return {
			"epoch": self.current_epoch,
			"hard_ratio": hard_ratio,
			"easy_ratio": 1.0 - hard_ratio,
			"total_samples": len(self.epoch_data),
			**mining_params
		}
	
	def __len__(self) -> int:
		return len(self.epoch_data)
	
	def __getitem__(self, index: int) -> Tuple[torch.Tensor, ...]:
		if self.is_train:
			return self._get_train_item(index)
		else:
			return self._get_val_item(index)
	
	def _get_train_item(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
		"""Return triplet: anchor, positive, negative with their IDs."""
		anchor_path, positive_path, negative_path, pid, nid = self.epoch_data[index]
		
		anchor = self._load_image(anchor_path)
		positive = self._load_image(positive_path)
		negative = self._load_image(negative_path)
		
		return anchor, positive, negative, int(pid), int(nid)
	
	def _get_val_item(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
		"""Return: img1, img2, label, id1, id2."""
		data_row = self.epoch_data[index]
		
		# Handle different data formats robustly
		if len(data_row) >= 5:
			img1_path, img2_path, label, id1, id2 = data_row[:5]
		elif len(data_row) >= 3:
			img1_path, img2_path, label = data_row[:3]
			# Fallback IDs
			id1, id2 = 0, 1
		else:
			raise ValueError(f"Invalid validation data format: expected at least 3 columns, got {len(data_row)}")
		
		try:
			img1 = self._load_image(img1_path)
			img2 = self._load_image(img2_path)
			
			return img1, img2, torch.tensor(float(label), dtype=torch.float32), int(id1), int(id2)
		except Exception as e:
			print(f"Error loading validation item {index}: {e}")
			print(f"Data row: {data_row}")
			raise	
			
	def _load_image(self, path: str) -> torch.Tensor:
		"""Load and transform image."""
		image = replace_background_with_white(
			path, self.folder_img, remove_bg=self.config.remove_bg
		)
		return self.transform(image) if self.transform else image
	
	def _default_transforms(self) -> transforms.Compose:
		"""Get default transforms with configurable augmentation strength."""
		normalize = transforms.Normalize(
			mean=[0.485, 0.456, 0.406],
			std=[0.229, 0.224, 0.225]
		)
		
		if self.is_train:
			aug_strength = self.config.augmentation_strength
			return transforms.Compose([
				transforms.Resize((224, 224)),
				transforms.RandomHorizontalFlip(p=0.5 * aug_strength),
				transforms.RandomRotation(degrees=int(10 * aug_strength)),
				transforms.ColorJitter(
					brightness=0.2 * aug_strength, 
					contrast=0.2 * aug_strength
				),
				transforms.GaussianBlur(kernel_size=5, sigma=(0.1, 2.0 * aug_strength)),
				transforms.ToTensor(),
				normalize
			])
		
		return transforms.Compose([
			transforms.Resize((224, 224)),
			transforms.ToTensor(),
			normalize
		])

# ----------------------------
# Enhanced Model Architecture
# ----------------------------
class ResNetBackbone(nn.Module):
	"""Enhanced ResNet backbone with better weight loading."""
	
	def __init__(self, model_name: str = "resnet34", pretrained_path: Optional[str] = None):
		super().__init__()
		
		# Initialize the ResNet model
		if model_name == "resnet18":
			self.resnet = models.resnet18(weights=None)
		elif model_name == "resnet34":
			self.resnet = models.resnet34(weights=None)
		elif model_name == "resnet50":
			self.resnet = models.resnet50(weights=None)
		else:
			raise ValueError(f"Unsupported model_name: {model_name}")
		
		# Load pretrained weights
		if pretrained_path and os.path.exists(pretrained_path):
			self._load_pretrained_weights(pretrained_path)
		elif pretrained_path:
			print(f"Warning: Pretrained path {pretrained_path} not found, using random initialization")
		
		# Remove the fully connected layer
		self.resnet.fc = nn.Identity()
		
		# Get output dimension
		with torch.no_grad():
			dummy = torch.randn(1, 3, 224, 224)
			self.output_dim = self.resnet(dummy).shape[1]
	
	def _load_pretrained_weights(self, pretrained_path: str):
		"""Load pretrained weights with comprehensive error handling."""
		try:
			checkpoint = torch.load(pretrained_path, map_location="cpu", weights_only=False)
			state_dict = checkpoint.get("state_dict", checkpoint)
			
			# Handle prefix issues
			if not any(key.startswith("resnet.") for key in state_dict.keys()):
				state_dict = {f"resnet.{k}": v for k, v in state_dict.items()}
			
			# Filter matching keys and sizes
			model_dict = self.state_dict()
			filtered_state_dict = {
				k: v for k, v in state_dict.items()
				if k in model_dict and v.size() == model_dict[k].size()
			}
			
			# Load filtered weights
			missing_keys = self.load_state_dict(filtered_state_dict, strict=False)
			
			print(f"[INFO] Loaded pretrained weights: {len(filtered_state_dict)}/{len(model_dict)} parameters")
			if missing_keys.missing_keys:
				print(f"[INFO] Missing keys: {len(missing_keys.missing_keys)}")
			
		except Exception as e:
			print(f"[ERROR] Failed to load pretrained weights: {e}")
			raise
	
	def forward(self, x: torch.Tensor) -> torch.Tensor:
		return self.resnet(x)

class AdvancedEmbeddingHead(nn.Module):
	"""Advanced embedding head with residual connections and normalization."""
	
	def __init__(self, input_dim: int, embedding_dim: int, dropout: float = 0.5):
		super().__init__()
		
		self.input_dim = input_dim
		self.embedding_dim = embedding_dim
		
		# Multi-layer embedding head with residual connections
		if input_dim > embedding_dim * 4:
			hidden_dim = max(embedding_dim * 2, input_dim // 4)
			self.layers = nn.Sequential(
				nn.Linear(input_dim, hidden_dim),
				nn.LayerNorm(hidden_dim),
				nn.GELU(),
				nn.Dropout(dropout),
				
				nn.Linear(hidden_dim, embedding_dim * 2),
				nn.LayerNorm(embedding_dim * 2),
				nn.GELU(),
				nn.Dropout(dropout / 2),
				
				nn.Linear(embedding_dim * 2, embedding_dim),
				nn.LayerNorm(embedding_dim)
			)
		else:
			# Simple head for smaller dimensions
			self.layers = nn.Sequential(
				nn.Linear(input_dim, embedding_dim),
				nn.LayerNorm(embedding_dim)
			)
	
	def forward(self, x: torch.Tensor) -> torch.Tensor:
		x = x.flatten(1)  # Flatten spatial dimensions
		return self.layers(x)

class SiameseSignatureNetwork(nn.Module):
    """Advanced Siamese network with precision-focused loss."""
    
    def __init__(self, config: TrainingConfig = CONFIG):
        super().__init__()
        self.config = config
        
        # Initialize backbone
        if config.model_name.startswith("resnet"):
            self.backbone = ResNetBackbone(
                model_name=config.model_name,
                pretrained_path=config.pretrained_path if config.last_epoch_weights is None else None
            )
            backbone_dim = self.backbone.output_dim
        else:
            raise ValueError(f"Unsupported model: {config.model_name}")
        
        # Initialize embedding head
        self.embedding_head = AdvancedEmbeddingHead(
            input_dim=backbone_dim,
            embedding_dim=config.embedding_dim,
            dropout=0.5
        )
        
        self.normalize_embeddings = config.normalize_embeddings
        self.distance_threshold = 0.5  # Will be updated during validation
        
        # Loss components
        self.criterion = MultiSimilarityLoss(
            alpha=config.multisim_alpha,
            beta=config.multisim_beta,
            base=config.multisim_base
        )
        
        # Add triplet margin loss for better separation
        self.triplet_loss = nn.TripletMarginLoss(
            margin=config.triplet_margin, 
            p=2,
            reduction='none'  # We'll apply weights manually
        )
        
        # Loss weights
        self.triplet_weight = config.triplet_weight
        self.fp_penalty_weight = config.false_positive_penalty_weight
        
        # Mining
        if config.use_hard_mining:
            self.miner = BatchHardMiner()
        else:
            self.miner = None
    
    def get_parameter_groups(self) -> List[Dict[str, Any]]:
        """Get parameter groups for differential learning rates."""
        backbone_params = list(self.backbone.parameters())
        head_params = list(self.embedding_head.parameters())
        
        return [
            {
                'params': backbone_params,
                'lr': self.config.backbone_lr,
                'name': 'backbone',
                'weight_decay': self.config.weight_decay
            },
            {
                'params': head_params,
                'lr': self.config.head_lr,
                'name': 'embedding_head',
                'weight_decay': self.config.weight_decay
            }
        ]
    
    def forward(self, anchor: torch.Tensor, positive: torch.Tensor, 
                negative: Optional[torch.Tensor] = None) -> Union[Tuple[torch.Tensor, torch.Tensor], 
                                                                  Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
        """Forward pass for training or inference."""
        a_features = self.backbone(anchor)
        a_emb = self.embedding_head(a_features)
        
        p_features = self.backbone(positive)
        p_emb = self.embedding_head(p_features)
        
        if self.normalize_embeddings:
            a_emb = F.normalize(a_emb, p=2, dim=1)
            p_emb = F.normalize(p_emb, p=2, dim=1)
        
        if negative is not None:
            n_features = self.backbone(negative)
            n_emb = self.embedding_head(n_features)
            
            if self.normalize_embeddings:
                n_emb = F.normalize(n_emb, p=2, dim=1)
            
            return a_emb, p_emb, n_emb
        
        return a_emb, p_emb
    
    def compute_loss(self, embeddings: torch.Tensor, labels: torch.Tensor,
                     anchors: Optional[torch.Tensor] = None,
                     positives: Optional[torch.Tensor] = None,
                     negatives: Optional[torch.Tensor] = None,
                     distance_weights: Optional[Dict[str, torch.Tensor]] = None) -> torch.Tensor:
        """Enhanced loss computation with precision focus and distance weighting."""
        
        # MultiSimilarity loss
        if self.miner is not None:
            hard_pairs = self.miner(embeddings, labels)
            ms_loss = self.criterion(embeddings, labels, hard_pairs)
        else:
            ms_loss = self.criterion(embeddings, labels)
        
        total_loss = ms_loss
        
        # Add triplet loss if embeddings provided
        if anchors is not None and positives is not None and negatives is not None:
            # Compute triplet losses for each sample
            triplet_losses = self.triplet_loss(anchors, positives, negatives)
            
            # Apply distance-based weights if provided
            if distance_weights is not None:
                neg_weights = distance_weights.get('negative_weights', torch.ones_like(triplet_losses))
                weighted_triplet_loss = (triplet_losses * neg_weights).mean()
            else:
                weighted_triplet_loss = triplet_losses.mean()
            
            total_loss += self.triplet_weight * weighted_triplet_loss
            
            # Additional penalty for hard negatives (false positives)
            with torch.no_grad():
                d_an = F.pairwise_distance(anchors, negatives)
                # Find negatives that are too close (potential false positives)
                hard_negative_mask = d_an < self.distance_threshold
                
            if hard_negative_mask.any():
                # Apply distance-based weights for false positive penalty
                if distance_weights is not None:
                    neg_weights = distance_weights.get('negative_weights', torch.ones_like(d_an))
                    # Extra penalty weighted by how bad the false positive is
                    false_positive_distances = self.distance_threshold - d_an[hard_negative_mask]
                    false_positive_weights = neg_weights[hard_negative_mask]
                    fp_loss = (false_positive_distances * false_positive_weights).mean()
                else:
                    fp_loss = (self.distance_threshold - d_an[hard_negative_mask]).mean()
                
                total_loss += self.fp_penalty_weight * fp_loss
        
        return total_loss
    
    def predict_pair(self, img1: torch.Tensor, img2: torch.Tensor, 
                    threshold: Optional[float] = None, return_dist: bool = False) -> torch.Tensor:
        """Predict similarity between image pairs."""
        self.eval()
        with torch.no_grad():
            emb1, emb2 = self(img1, img2)
            distances = F.pairwise_distance(emb1, emb2)
            
            if return_dist:
                return distances
            
            thresh = threshold if threshold is not None else self.distance_threshold
            return (distances < thresh).long()

# ----------------------------
# Advanced Training Metrics and Statistics
# ----------------------------
class TrainingMetrics:
    """Enhanced training metrics with adaptive mining for both hard positives and negatives."""
    
    def __init__(self):
        self.reset()
        # Track distance statistics for adaptive thresholds
        self.distance_history = {"positive": [], "negative": []}
        self.adaptive_stats = {}
        
    def reset(self):
        """Reset all metrics."""
        self.losses = []
        self.genuine_distances = []
        self.forged_distances = []
        
        # Separate mining stats for positives and negatives
        self.positive_mining_stats = {"easy": 0, "semi": 0, "hard": 0}
        self.negative_mining_stats = {"easy": 0, "semi": 0, "hard": 0}
        
        # Hard sample counts
        self.hard_positive_count = 0
        self.hard_negative_count = 0
        self.total_positive_pairs = 0
        self.total_negative_pairs = 0
        
        # False positive/negative tracking
        self.false_positive_count = 0
        self.false_negative_count = 0
        
        self.learning_rates = {}
        
    def update_distance_statistics(self, d_positive: np.ndarray, d_negative: np.ndarray):
        """Update running statistics for adaptive thresholds."""
        # Keep rolling window of recent distances
        self.distance_history["positive"].extend(d_positive.tolist())
        self.distance_history["negative"].extend(d_negative.tolist())
        
        # Keep only recent history (last 5000 samples)
        for key in self.distance_history:
            if len(self.distance_history[key]) > 5000:
                self.distance_history[key] = self.distance_history[key][-5000:]
        
        # Compute adaptive statistics
        if len(self.distance_history["positive"]) > 100 and len(self.distance_history["negative"]) > 100:
            pos_distances = np.array(self.distance_history["positive"])
            neg_distances = np.array(self.distance_history["negative"])
            
            self.adaptive_stats = {
                "pos_mean": np.mean(pos_distances),
                "pos_std": np.std(pos_distances),
                "pos_q25": np.percentile(pos_distances, 25),
                "pos_q50": np.percentile(pos_distances, 50),
                "pos_q75": np.percentile(pos_distances, 75),
                "pos_q90": np.percentile(pos_distances, 90),
                
                "neg_mean": np.mean(neg_distances),
                "neg_std": np.std(neg_distances),
                "neg_q10": np.percentile(neg_distances, 10),
                "neg_q25": np.percentile(neg_distances, 25),
                "neg_q50": np.percentile(neg_distances, 50),
                "neg_q75": np.percentile(neg_distances, 75),
                
                "separation": np.mean(neg_distances) - np.mean(pos_distances),
                "overlap_region": max(0, np.percentile(pos_distances, 95) - np.percentile(neg_distances, 5))
            }

    def compute_precision_focused_weights(self, d_positive: np.ndarray, 
                                        d_negative: np.ndarray,
                                        negative_weight_multiplier: float = 2.5) -> Tuple[torch.Tensor, torch.Tensor]:
        """Compute sample weights with focus on improving precision."""
        pos_weights = np.ones_like(d_positive)
        neg_weights = np.ones_like(d_negative)
        
        if self.adaptive_stats:
            # Hard negatives (forged that look genuine) get MUCH higher weight
            neg_q10 = self.adaptive_stats["neg_q10"]
            neg_q25 = self.adaptive_stats["neg_q25"]
            
            # Very hard negatives (bottom 10%) - highest weight
            very_hard_neg_mask = d_negative < neg_q10
            neg_weights[very_hard_neg_mask] = negative_weight_multiplier * 1.5
            
            # Hard negatives (10-25%) - high weight
            hard_neg_mask = (d_negative >= neg_q10) & (d_negative < neg_q25)
            neg_weights[hard_neg_mask] = negative_weight_multiplier
            
            # Semi-hard negatives (25-50%) - moderate weight
            semi_neg_mask = (d_negative >= neg_q25) & (d_negative < self.adaptive_stats["neg_q50"])
            neg_weights[semi_neg_mask] = negative_weight_multiplier * 0.6
            
            # Hard positives get moderate weight (but less than hard negatives)
            pos_q75 = self.adaptive_stats["pos_q75"]
            pos_q90 = self.adaptive_stats["pos_q90"]
            
            # Very hard positives (top 10%)
            very_hard_pos_mask = d_positive > pos_q90
            pos_weights[very_hard_pos_mask] = 1.8
            
            # Hard positives (75-90%)
            hard_pos_mask = (d_positive > pos_q75) & (d_positive <= pos_q90)
            pos_weights[hard_pos_mask] = 1.5
        
        return torch.tensor(pos_weights, dtype=torch.float32), torch.tensor(neg_weights, dtype=torch.float32)

    def update_mining_stats(self, d_positive: np.ndarray, d_negative: np.ndarray, 
                           margin: float, difficulty_params: Dict[str, float]):
        """Intelligent adaptive mining for both hard positives and hard negatives."""
        
        # Update distance statistics first
        self.update_distance_statistics(d_positive, d_negative)
        
        # Update totals
        self.total_positive_pairs += len(d_positive)
        self.total_negative_pairs += len(d_negative)
        
        # Use adaptive thresholds if available, otherwise fallback to fixed
        if self.adaptive_stats:
            self._adaptive_mining(d_positive, d_negative, difficulty_params)
        else:
            self._fixed_mining(d_positive, d_negative, margin)
    
    def _adaptive_mining(self, d_positive: np.ndarray, d_negative: np.ndarray, 
                        difficulty_params: Dict[str, float]):
        """Adaptive mining based on current distance distributions."""
        stats = self.adaptive_stats
        
        # Get difficulty parameters
        hard_positive_ratio = difficulty_params.get("hard_positive_ratio", 0.3)
        hard_negative_ratio = difficulty_params.get("hard_negative_ratio", 0.3)
        
        # Dynamic thresholds for hard positives (far apart genuine pairs)
        # Use percentile based on desired hard positive ratio
        hard_pos_percentile = 100 - (hard_positive_ratio * 100)
        hard_pos_threshold = np.percentile(self.distance_history["positive"][-1000:], hard_pos_percentile)
        semi_pos_threshold = stats["pos_q50"]
        
        # Dynamic thresholds for hard negatives (close together impostor pairs)
        # Use percentile based on desired hard negative ratio
        hard_neg_percentile = hard_negative_ratio * 100
        hard_neg_threshold = np.percentile(self.distance_history["negative"][-1000:], hard_neg_percentile)
        semi_neg_threshold = stats["neg_q50"]
        
        # Mine hard positives
        for dp in d_positive:
            if dp >= hard_pos_threshold:
                self.positive_mining_stats["hard"] += 1
                self.hard_positive_count += 1
            elif dp >= semi_pos_threshold:
                self.positive_mining_stats["semi"] += 1
            else:
                self.positive_mining_stats["easy"] += 1
        
        # Mine hard negatives
        for dn in d_negative:
            if dn <= hard_neg_threshold:
                self.negative_mining_stats["hard"] += 1
                self.hard_negative_count += 1
            elif dn <= semi_neg_threshold:
                self.negative_mining_stats["semi"] += 1
            else:
                self.negative_mining_stats["easy"] += 1
    
    def _fixed_mining(self, d_positive: np.ndarray, d_negative: np.ndarray, margin: float):
        """Fallback fixed mining for early epochs."""
        # Fixed thresholds
        hard_pos_threshold = 0.5  # Far genuine pairs
        hard_neg_threshold = 0.3  # Close impostor pairs
        
        for dp in d_positive:
            if dp >= hard_pos_threshold:
                self.positive_mining_stats["hard"] += 1
                self.hard_positive_count += 1
            elif dp >= hard_pos_threshold * 0.7:
                self.positive_mining_stats["semi"] += 1
            else:
                self.positive_mining_stats["easy"] += 1
        
        for dn in d_negative:
            if dn <= hard_neg_threshold:
                self.negative_mining_stats["hard"] += 1
                self.hard_negative_count += 1
            elif dn <= hard_neg_threshold * 1.5:
                self.negative_mining_stats["semi"] += 1
            else:
                self.negative_mining_stats["easy"] += 1
    
    def get_mining_percentages(self) -> Dict[str, float]:
        """Get mining statistics as percentages with debugging info."""
        total_pos = sum(self.positive_mining_stats.values())
        total_neg = sum(self.negative_mining_stats.values())
        
        percentages = {}
        
        # Positive pair mining stats
        if total_pos > 0:
            percentages.update({
                "pos_mining_easy_pct": 100.0 * self.positive_mining_stats["easy"] / total_pos,
                "pos_mining_semi_pct": 100.0 * self.positive_mining_stats["semi"] / total_pos,
                "pos_mining_hard_pct": 100.0 * self.positive_mining_stats["hard"] / total_pos,
            })
        else:
            percentages.update({
                "pos_mining_easy_pct": 0.0,
                "pos_mining_semi_pct": 0.0,
                "pos_mining_hard_pct": 0.0,
            })
        
        # Negative pair mining stats
        if total_neg > 0:
            percentages.update({
                "neg_mining_easy_pct": 100.0 * self.negative_mining_stats["easy"] / total_neg,
                "neg_mining_semi_pct": 100.0 * self.negative_mining_stats["semi"] / total_neg,
                "neg_mining_hard_pct": 100.0 * self.negative_mining_stats["hard"] / total_neg,
            })
        else:
            percentages.update({
                "neg_mining_easy_pct": 0.0,
                "neg_mining_semi_pct": 0.0,
                "neg_mining_hard_pct": 0.0,
            })
        
        # Overall hard sample ratios
        if self.total_positive_pairs > 0:
            percentages["hard_positive_ratio"] = 100.0 * self.hard_positive_count / self.total_positive_pairs
        else:
            percentages["hard_positive_ratio"] = 0.0
            
        if self.total_negative_pairs > 0:
            percentages["hard_negative_ratio"] = 100.0 * self.hard_negative_count / self.total_negative_pairs
        else:
            percentages["hard_negative_ratio"] = 0.0
            
        # False positive/negative rates
        total_samples = self.total_positive_pairs + self.total_negative_pairs
        if total_samples > 0:
            percentages["false_positive_rate"] = 100.0 * self.false_positive_count / self.total_negative_pairs if self.total_negative_pairs > 0 else 0.0
            percentages["false_negative_rate"] = 100.0 * self.false_negative_count / self.total_positive_pairs if self.total_positive_pairs > 0 else 0.0
        
        # Add adaptive stats if available
        if self.adaptive_stats:
            percentages.update({
                "adaptive_separation": self.adaptive_stats["separation"],
                "adaptive_overlap": self.adaptive_stats["overlap_region"],
                "adaptive_pos_spread": self.adaptive_stats["pos_std"],
                "adaptive_neg_spread": self.adaptive_stats["neg_std"],
            })
        
        return percentages
    
    def compute_separation_metrics(self) -> Dict[str, float]:
        """Compute distance separation metrics."""
        if not self.genuine_distances or not self.forged_distances:
            return {
                "genuine_dist_mean": 0.0,
                "forged_dist_mean": 0.0,
                "genuine_dist_std": 0.0,
                "forged_dist_std": 0.0,
                "separation": 0.0,
                "overlap": 0.0,
                "separation_ratio": 0.0,
                "cohesion_ratio": 0.0
            }
        
        gen_mean = np.mean(self.genuine_distances)
        forg_mean = np.mean(self.forged_distances)
        gen_std = np.std(self.genuine_distances)
        forg_std = np.std(self.forged_distances)
        
        separation = forg_mean - gen_mean
        overlap = max(0, gen_mean + 2*gen_std - (forg_mean - 2*forg_std))
        
        # Cohesion ratio: how tight are genuine pairs relative to separation
        cohesion_ratio = gen_std / (separation + 1e-8)
        
        return {
            "genuine_dist_mean": gen_mean,
            "forged_dist_mean": forg_mean,
            "genuine_dist_std": gen_std,
            "forged_dist_std": forg_std,
            "separation": separation,
            "overlap": overlap,
            "separation_ratio": separation / (gen_std + forg_std + 1e-8),
            "cohesion_ratio": cohesion_ratio
        }
		
# ----------------------------
# Enhanced Training Loop
# ----------------------------
class SignatureTrainer:
	"""Research-grade signature verification trainer."""
	
	def __init__(self, config: TrainingConfig = CONFIG):
		self.config = config
		self.device = torch.device(config.device)
		
		# Initialize managers
		self.mlflow_manager = MLFlowManager(tracking_uri=self.config.tracking_uri)
		self.curriculum_manager = CurriculumLearningManager(config)
		
		# Training state
		self.current_epoch = 0
		self.best_eer = float('inf')
		self.patience_counter = 0
		self.global_step = 0
		
		# Setup logging
		self._setup_logging()
		
	def _setup_logging(self):
		"""Setup comprehensive logging."""
		logging.basicConfig(
			level=logging.INFO,
			format='%(asctime)s - %(levelname)s - %(message)s',
			handlers=[
				logging.FileHandler('training.log'),
				logging.StreamHandler()
			]
		)
		self.logger = logging.getLogger(__name__)
	
	def _prepare_datasets(self) -> Tuple[SignatureDataset, SignatureDataset]:
		"""Prepare training and validation datasets."""
		# Load datasets
		train_data = pd.read_excel("../../data/classify/preprared_data/labels/train_triplets_balanced_v12.xlsx")
		val_data = pd.read_excel("../../data/classify/preprared_data/labels/valid_pairs_balanced_v12.xlsx")
		
		train_dataset = SignatureDataset(
			folder_img="../../data/classify/preprared_data/images/",
			excel_data=train_data,
			curriculum_manager=self.curriculum_manager,
			is_train=True,
			config=self.config
		)
		
		val_dataset = SignatureDataset(
			folder_img="../../data/classify/preprared_data/images/",
			excel_data=val_data,
			curriculum_manager=self.curriculum_manager,
			is_train=False,
			config=self.config
		)
		
		self.logger.info(f"Training samples: {len(train_dataset)}")
		self.logger.info(f"Validation samples: {len(val_dataset)}")
		
		return train_dataset, val_dataset
		
	def _compute_precision_optimized_threshold(self, distances: np.ndarray, 
											  labels: np.ndarray, 
											  target_precision: float = None) -> float:
		"""Find threshold that achieves target precision while maximizing F1."""
		if target_precision is None:
			target_precision = self.config.target_precision
			
		thresholds = np.linspace(distances.min(), distances.max(), 1000)
		best_threshold = thresholds[0]
		best_f1 = 0
		best_precision = 0
		best_recall = 0
		
		for thresh in thresholds:
			predictions = (distances < thresh).astype(int)
			
			# Calculate metrics
			tp = np.sum((predictions == 1) & (labels == 1))
			fp = np.sum((predictions == 1) & (labels == 0))
			fn = np.sum((predictions == 0) & (labels == 1))
			
			precision = tp / (tp + fp + 1e-8)
			recall = tp / (tp + fn + 1e-8)
			f1 = 2 * precision * recall / (precision + recall + 1e-8)
			
			# Prioritize precision while maintaining reasonable recall
			if precision >= target_precision and f1 > best_f1:
				best_f1 = f1
				best_threshold = thresh
				best_precision = precision
				best_recall = recall
			# If we can't achieve target precision, get best precision with recall > 0.5
			elif precision > best_precision and recall > 0.5:
				best_f1 = f1
				best_threshold = thresh
				best_precision = precision
				best_recall = recall
		
		print(f"  Precision-optimized threshold: {best_threshold:.4f} "
			  f"(P: {best_precision:.3f}, R: {best_recall:.3f}, F1: {best_f1:.3f})")
		
		return best_threshold	
	
	def _setup_model_and_optimizer(self) -> Tuple[SiameseSignatureNetwork, torch.optim.Optimizer, Any]:
		"""Setup model, optimizer, and scheduler."""
		# Initialize model
		model = SiameseSignatureNetwork(self.config)
		
		# Compile model if available
		if hasattr(torch, "compile") and platform.system() != "Windows":
			self.logger.info("Compiling model with torch.compile")
			model = torch.compile(model)
		
		model = model.to(self.device)
		
		# Count parameters
		total_params = sum(p.numel() for p in model.parameters())
		trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
		self.logger.info(f"Total parameters: {total_params:,}")
		self.logger.info(f"Trainable parameters: {trainable_params:,}")
		
		# Setup optimizer with parameter groups
		param_groups = model.get_parameter_groups()
		optimizer = torch.optim.AdamW(param_groups)
		
		# Log learning rates
		for group in param_groups:
			self.logger.info(f"Parameter group '{group['name']}': LR = {group['lr']:.2e}")
		
		# Setup scheduler
		if self.config.lr_scheduler == "cosine":
			scheduler = CosineAnnealingLR(optimizer, T_max=self.config.max_epochs)
		else:
			scheduler = ReduceLROnPlateau(optimizer, mode='min', patience=5, factor=0.5)
		
		return model, optimizer, scheduler
	
	def _setup_checkpoint_management(self, run_id: str) -> Tuple[str, str]:
		"""Setup checkpoint directories."""
		checkpoint_dir = os.path.join("../../model/models_checkpoints/", run_id)
		figures_dir = os.path.join(checkpoint_dir, "figures")
		os.makedirs(checkpoint_dir, exist_ok=True)
		os.makedirs(figures_dir, exist_ok=True)
		return checkpoint_dir, figures_dir
	
	def _load_checkpoint(self, model: nn.Module, optimizer: torch.optim.Optimizer, 
						scheduler: Any, scaler: torch.amp.GradScaler) -> int:
		"""Load checkpoint if specified."""
		if not self.config.last_epoch_weights:
			return 1
		
		checkpoint_path = self.config.last_epoch_weights
		self.logger.info(f"Loading checkpoint from {checkpoint_path}")
		
		try:
			checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
			
			model.load_state_dict(checkpoint["model_state_dict"])
			optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
			scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
			scaler.load_state_dict(checkpoint.get("scaler_state_dict", scaler.state_dict()))
			
			start_epoch = checkpoint["epoch"] + 1
			self.best_eer = checkpoint.get("best_eer", self.best_eer)
			model.distance_threshold = checkpoint.get("prediction_threshold", 0.5)
			
			self.logger.info(f"Resumed from epoch {start_epoch}, best EER: {self.best_eer:.4f}")
			return start_epoch
			
		except Exception as e:
			self.logger.error(f"Failed to load checkpoint: {e}")
			return 1
			
	def train_epoch(self, model: nn.Module, train_loader: DataLoader, 
				   optimizer: torch.optim.Optimizer, scaler: torch.amp.GradScaler,
				   epoch: int) -> TrainingMetrics:
		"""Enhanced training with intelligent adaptive mining for both hard positives and negatives."""
		model.train()
		metrics = TrainingMetrics()
		
		curriculum_stats = train_loader.dataset.get_curriculum_stats()
		
		# INTELLIGENT MARGIN CALCULATION
		base_margin = 0.5  # Base margin for normalized embeddings
		margin_multiplier = curriculum_stats["margin_multiplier"]
		adaptive_margin = base_margin * margin_multiplier
		
		# Progressive margin adjustment based on epoch
		epoch_progress = epoch / self.config.max_epochs
		progressive_factor = 1.2 - 0.4 * epoch_progress  # 1.2 → 0.8
		final_margin = adaptive_margin * progressive_factor
		
		# Tracking counters
		forgery_batch_count = 0
		genuine_batch_count = 0
		batch_fp_count = 0
		batch_fn_count = 0
		
		# Debug info
		debug_printed = False
		
		pbar = tqdm(train_loader, desc=f"[Train] Epoch {epoch}")
		
		for step, (anchors, positives, negatives, anchor_ids, negative_ids) in enumerate(pbar):
			
			# Move to device
			anchors = anchors.to(self.device, non_blocking=True)
			positives = positives.to(self.device, non_blocking=True)
			negatives = negatives.to(self.device, non_blocking=True)
			anchor_ids = anchor_ids.to(self.device, non_blocking=True)
			negative_ids = negative_ids.to(self.device, non_blocking=True)
			
			# Count forgery vs genuine negatives
			max_genuine_id = len(train_loader.dataset.genuine_id_mapping)
			forgery_mask = negative_ids >= max_genuine_id + 100
			forgery_batch_count += forgery_mask.sum().item()
			genuine_batch_count += (~forgery_mask).sum().item()
			
			if not debug_printed and step == 0:
				print(f"\n[DEBUG Epoch {epoch}]")
				print(f"  Final margin: {final_margin:.3f}")
				print(f"  Hard negative ratio target: {curriculum_stats['hard_negative_ratio']:.3f}")
				print(f"  Hard positive ratio target: {curriculum_stats['hard_positive_ratio']:.3f}")
				print(f"  Negative weight multiplier: {self.config.negative_weight_multiplier:.2f}")
				print(f"  Triplet weight: {self.config.triplet_weight:.2f}")
				print(f"  FP penalty weight: {self.config.false_positive_penalty_weight:.2f}")
				debug_printed = True
			
			# Forward pass to get embeddings first
			with torch.amp.autocast(device_type=self.device.type):
				a_emb, p_emb, n_emb = model(anchors, positives, negatives)
			
			# Compute distances and weights BEFORE loss computation
			with torch.no_grad():
				d_ap = F.pairwise_distance(a_emb, p_emb).cpu().numpy()
				d_an = F.pairwise_distance(a_emb, n_emb).cpu().numpy()
				
				# Get precision-focused weights
				pos_weights, neg_weights = metrics.compute_precision_focused_weights(
					d_ap, d_an, 
					negative_weight_multiplier=self.config.negative_weight_multiplier
				)
				pos_weights = pos_weights.to(self.device)
				neg_weights = neg_weights.to(self.device)
				
				# Track false positives/negatives
				fp_mask = d_an < model.distance_threshold
				fn_mask = d_ap > model.distance_threshold
				batch_fp_count = fp_mask.sum()
				batch_fn_count = fn_mask.sum()
				metrics.false_positive_count += batch_fp_count
				metrics.false_negative_count += batch_fn_count
			
			# Prepare distance weights for loss
			distance_weights = {
				'positive_weights': pos_weights,
				'negative_weights': neg_weights
			}
			
			# Now compute loss with weights
			with torch.amp.autocast(device_type=self.device.type):
				all_embeddings = torch.cat([a_emb, p_emb, n_emb], dim=0)
				all_labels = torch.cat([anchor_ids, anchor_ids, negative_ids], dim=0)
				
				# Compute loss with triplet component and distance weights
				batch_loss = model.compute_loss(
					all_embeddings, all_labels,
					anchors=a_emb, positives=p_emb, negatives=n_emb,
					distance_weights=distance_weights
				)
			
			# Gradient accumulation
			loss = batch_loss / self.config.grad_accum_steps
			scaler.scale(loss).backward()
			
			if (step + 1) % self.config.grad_accum_steps == 0 or (step + 1) == len(train_loader):
				scaler.unscale_(optimizer)
				torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
				scaler.step(optimizer)
				scaler.update()
				optimizer.zero_grad(set_to_none=True)
				self.global_step += 1
			
			# Update metrics
			metrics.losses.append(batch_loss.item())
			metrics.genuine_distances.extend(d_ap.tolist())
			metrics.forged_distances.extend(d_an.tolist())
			
			# Use enhanced mining with difficulty parameters
			metrics.update_mining_stats(d_ap, d_an, final_margin, curriculum_stats)
			
			# Store learning rates
			for i, group in enumerate(optimizer.param_groups):
				metrics.learning_rates[f"lr_{group.get('name', i)}"] = group['lr']
			
			# Enhanced progress bar with precision focus
			sep = np.mean(d_an) - np.mean(d_ap)
			actual_forgery_ratio = forgery_batch_count / (forgery_batch_count + genuine_batch_count) if (forgery_batch_count + genuine_batch_count) > 0 else 0
			
			# Get current mining stats
			mining_pcts = metrics.get_mining_percentages()
			
			pbar.set_postfix({
				"loss": f"{batch_loss.item():.3f}",
				"h_neg%": f"{mining_pcts.get('neg_mining_hard_pct', 0):.0f}",
				"h_pos%": f"{mining_pcts.get('pos_mining_hard_pct', 0):.0f}",
				"d_sep": f"{sep:.3f}",
				"FP": f"{batch_fp_count}",
				"FN": f"{batch_fn_count}",
				"margin": f"{final_margin:.3f}"
			})
			
			# Periodic logging
			if self.global_step % self.config.log_frequency == 0:
				enhanced_stats = {
					**curriculum_stats,
					**mining_pcts,
					"actual_forgery_ratio": actual_forgery_ratio,
					"batch_false_positives": int(batch_fp_count),
					"batch_false_negatives": int(batch_fn_count),
					"final_margin": final_margin,
					"epoch_progress": epoch_progress
				}
				self._log_training_step(metrics, enhanced_stats, self.global_step)
			
			# Memory cleanup
			del anchors, positives, negatives, a_emb, p_emb, n_emb
			torch.cuda.empty_cache()
		
		# End-of-epoch mining summary
		mining_pcts = metrics.get_mining_percentages()
		print(f"\n[Epoch {epoch} Mining Summary]")
		print(f"  Hard Negatives: {mining_pcts.get('neg_mining_hard_pct', 0):.1f}% | Semi: {mining_pcts.get('neg_mining_semi_pct', 0):.1f}% | Easy: {mining_pcts.get('neg_mining_easy_pct', 0):.1f}%")
		print(f"  Hard Positives: {mining_pcts.get('pos_mining_hard_pct', 0):.1f}% | Semi: {mining_pcts.get('pos_mining_semi_pct', 0):.1f}% | Easy: {mining_pcts.get('pos_mining_easy_pct', 0):.1f}%")
		print(f"  Overall Hard Ratios - Positives: {mining_pcts.get('hard_positive_ratio', 0):.1f}% | Negatives: {mining_pcts.get('hard_negative_ratio', 0):.1f}%")
		print(f"  False Positive Rate: {mining_pcts.get('false_positive_rate', 0):.1f}% | False Negative Rate: {mining_pcts.get('false_negative_rate', 0):.1f}%")
		
		if "adaptive_separation" in mining_pcts:
			print(f"  Adaptive separation: {mining_pcts['adaptive_separation']:.3f} | Overlap: {mining_pcts['adaptive_overlap']:.3f}")
		
		return metrics

	def validate_epoch(self, model: nn.Module, val_loader: DataLoader, 
					  epoch: int) -> Tuple[float, float, Dict[str, float]]:
		"""Validate for one epoch."""
		model.eval()
		
		val_distances = []
		val_labels = []
		val_embeddings = []
		val_person_ids = []
		val_loss_total = 0.0
		
		with torch.no_grad():
			pbar = tqdm(val_loader, desc=f"[Val] Epoch {epoch}")
			
			for img1, img2, labels, id1, id2 in pbar:
				# Move to device
				img1 = img1.to(self.device, non_blocking=True)
				img2 = img2.to(self.device, non_blocking=True)
				labels = labels.to(self.device, non_blocking=True)
				id1 = id1.to(self.device, non_blocking=True)
				id2 = id2.to(self.device, non_blocking=True)
				
				# Forward pass
				emb1, emb2 = model(img1, img2)
				distances = F.pairwise_distance(emb1, emb2)
				
				# Compute validation loss
				val_loss = self._compute_validation_loss(emb1, emb2, labels, id1, id2, model.criterion)
				val_loss_total += val_loss.item()
				
				# Collect results
				val_distances.extend(distances.cpu().numpy())
				val_labels.extend(labels.cpu().numpy())
				val_embeddings.append(emb1.cpu().numpy())
				val_embeddings.append(emb2.cpu().numpy())
				val_person_ids.extend(id1.cpu().numpy())
				val_person_ids.extend(id2.cpu().numpy())
				
				# Update progress
				pos_mask = labels == 1
				neg_mask = labels == 0
				pos_dist = distances[pos_mask].mean().item() if pos_mask.any() else 0.0
				neg_dist = distances[neg_mask].mean().item() if neg_mask.any() else 0.0
				
				pbar.set_postfix({
					"loss": f"{val_loss.item():.4f}",
					"d_pos": f"{pos_dist:.3f}",
					"d_neg": f"{neg_dist:.3f}",
					"sep": f"{neg_dist - pos_dist:.3f}"
				})
				
				# Memory cleanup
				del img1, img2, emb1, emb2
				torch.cuda.empty_cache()
		
		# Process results
		val_distances = np.array(val_distances)
		val_labels = np.array(val_labels)
		val_embeddings = np.concatenate(val_embeddings)
		val_person_ids = np.array(val_person_ids)
		avg_val_loss = val_loss_total / len(val_loader)
		
		# Compute metrics
		threshold, eer, metrics_dict = self._compute_validation_metrics(
			val_distances, val_labels, avg_val_loss
		)
		
		# Update model threshold
		model.distance_threshold = threshold
		
		return threshold, eer, {
			"metrics": metrics_dict,
			"embeddings": val_embeddings,
			"labels": np.repeat(val_labels, 2),
			"person_ids": val_person_ids,
			"distances": np.repeat(val_distances, 2)
		}
	
	def _compute_validation_loss(self, emb1: torch.Tensor, emb2: torch.Tensor, 
							   binary_labels: torch.Tensor, person_ids1: torch.Tensor, 
							   person_ids2: torch.Tensor, criterion) -> torch.Tensor:
		"""Compute validation loss using MultiSimilarityLoss."""
		labels1 = person_ids1.clone()
		labels2 = person_ids2.clone()
		
		# Handle forged pairs
		forged_mask = binary_labels == 0
		if forged_mask.any():
			max_person_id = torch.max(torch.cat([person_ids1, person_ids2])).item()
			labels2[forged_mask] = labels2[forged_mask] + max_person_id + 1
		
		# Handle genuine pairs
		genuine_mask = binary_labels == 1
		labels2[genuine_mask] = labels1[genuine_mask]
		
		# Combine embeddings and labels
		all_embeddings = torch.cat([emb1, emb2], dim=0)
		all_labels = torch.cat([labels1, labels2], dim=0)
		
		return criterion(all_embeddings, all_labels)
	
	def _compute_validation_metrics(self, distances: np.ndarray, labels: np.ndarray, 
								  val_loss: float) -> Tuple[float, float, Dict[str, float]]:
		"""Compute comprehensive validation metrics with precision focus."""
		# Compute EER and threshold
		similarity_scores = 1.0 / (distances + 1e-8)
		fpr, tpr, thresholds = roc_curve(labels, similarity_scores, pos_label=1)
		fnr = 1 - tpr
		eer_idx = np.nanargmin(np.abs(fpr - fnr))
		eer = fpr[eer_idx]
		eer_threshold = 1.0 / thresholds[eer_idx]
		
		# Get precision-optimized threshold
		precision_threshold = self._compute_precision_optimized_threshold(distances, labels)
		
		# Use precision-optimized threshold instead of EER threshold
		threshold = precision_threshold
		
		# Compute metrics with precision-optimized threshold
		predictions = (distances < threshold).astype(int)
		precision, recall, f1, _ = precision_recall_fscore_support(
			labels, predictions, average='binary', zero_division=0
		)
		accuracy = (predictions == labels).mean()
		roc_auc = auc(fpr, tpr)
		
		# Distance statistics
		genuine_dist = np.mean([d for d, l in zip(distances, labels) if l == 1])
		forged_dist = np.mean([d for d, l in zip(distances, labels) if l == 0])
		separation = forged_dist - genuine_dist
		
		# Confidence scores
		confidences = 1.0 / (distances + 1e-8)
		conf_genuine = np.mean([c for c, l in zip(confidences, labels) if l == 1])
		conf_forged = np.mean([c for c, l in zip(confidences, labels) if l == 0])
		
		metrics_dict = {
			"val_loss": val_loss,
			"val_EER": eer,
			"val_f1": f1,
			"val_auc": roc_auc,
			"val_accuracy": accuracy,
			"val_precision": precision,
			"val_recall": recall,
			"val_separation": separation,
			"val_genuine_dist": genuine_dist,
			"val_forged_dist": forged_dist,
			"val_genuine_conf": conf_genuine,
			"val_forged_conf": conf_forged,
			"threshold": threshold,
			"eer_threshold": eer_threshold,
			"precision_threshold": precision_threshold
		}
		
		return threshold, eer, metrics_dict	

	def _log_training_step(self, metrics: TrainingMetrics, curriculum_stats: Dict, step: int):
		"""Log training step metrics."""
		if not metrics.losses:
			return
		
		try:
			# Compute separation metrics
			sep_metrics = metrics.compute_separation_metrics()
			
			# Get mining percentages
			mining_percentages = metrics.get_mining_percentages()
			
			# Log to MLflow
			log_dict = {
				"train_loss": np.mean(metrics.losses[-10:]),  # Last 10 batches
				**sep_metrics,
				**mining_percentages,
				**curriculum_stats,
				**metrics.learning_rates
			}
			
			mlflow.log_metrics(log_dict, step=step)
			
		except Exception as e:
			print(f"Warning: Failed to log training step metrics: {e}")
	
	def _log_epoch_metrics(self, train_metrics: TrainingMetrics, val_metrics: Dict, epoch: int):
		"""Log comprehensive epoch metrics."""
		try:
			# Training metrics
			train_sep = train_metrics.compute_separation_metrics()
			train_mining = train_metrics.get_mining_percentages()
			
			log_dict = {
				"epoch": epoch,
				"train_loss_epoch": np.mean(train_metrics.losses),
				**train_sep,
				**train_mining,
				**val_metrics["metrics"],
				**train_metrics.learning_rates
			}
			
			mlflow.log_metrics(log_dict, step=epoch)
			
			# Log key metrics to console
			self.logger.info(f"Epoch {epoch}/{self.config.max_epochs} Summary:")
			self.logger.info(f"  Train Loss: {log_dict['train_loss_epoch']:.4f}")
			self.logger.info(f"  Val EER: {log_dict['val_EER']:.4f}")
			self.logger.info(f"  Val F1: {log_dict['val_f1']:.4f}")
			self.logger.info(f"  Separation: {log_dict['separation']:.4f}")
			
		except Exception as e:
			self.logger.error(f"Failed to log epoch metrics: {e}")
			# Log minimal metrics as fallback
			mlflow.log_metrics({
				"epoch": epoch,
				"train_loss_epoch": np.mean(train_metrics.losses) if train_metrics.losses else 0.0,
				**val_metrics["metrics"]
			}, step=epoch)
		
	def _save_checkpoint(self, model: nn.Module, optimizer: torch.optim.Optimizer,
						scheduler: Any, scaler: torch.amp.GradScaler, epoch: int,
						threshold: float, eer: float, checkpoint_dir: str, is_best: bool = False):
		"""Save model checkpoint."""
		checkpoint = {
			"epoch": epoch,
			"model_state_dict": model.state_dict(),
			"optimizer_state_dict": optimizer.state_dict(),
			"scheduler_state_dict": scheduler.state_dict(),
			"scaler_state_dict": scaler.state_dict(),
			"prediction_threshold": threshold,
			"best_eer": self.best_eer,
			"eer": eer,
			"config": asdict(self.config)
		}
		
		# Save regular checkpoint
		if epoch % self.config.save_frequency == 0:
			torch.save(checkpoint, os.path.join(checkpoint_dir, f"epoch_{epoch}.pth"))
		
		# Save best checkpoint
		if is_best:
			torch.save(checkpoint, os.path.join(checkpoint_dir, "best_model.pth"))
			self.logger.info(f"New best model saved with EER: {eer:.4f}")
	
	def _create_visualizations(self, val_results: Dict, epoch: int, figures_dir: str):
		"""Create comprehensive visualizations."""
		if epoch % self.config.visualize_frequency != 0:
			return
		
		# Distance distribution plot
		self._plot_distance_distribution(
			val_results["distances"][:len(val_results["distances"])//2],
			val_results["labels"][:len(val_results["labels"])//2],
			epoch, figures_dir
		)
		
		# t-SNE embedding visualization
		self._plot_tsne_embeddings(
			val_results["embeddings"],
			val_results["labels"],
			val_results["person_ids"],
			val_results["distances"],
			epoch, figures_dir
		)
	
	def _plot_distance_distribution(self, distances: np.ndarray, labels: np.ndarray, 
								   epoch: int, figures_dir: str):
		"""Plot distance distribution."""
		genuine_dists = distances[labels == 1]
		forged_dists = distances[labels == 0]
		
		plt.figure(figsize=(12, 8))
		plt.hist(genuine_dists, bins=50, alpha=0.6, color='blue',
				label=f'Genuine (μ={np.mean(genuine_dists):.4f}±{np.std(genuine_dists):.4f})')
		plt.hist(forged_dists, bins=50, alpha=0.6, color='red',
				label=f'Forged (μ={np.mean(forged_dists):.4f}±{np.std(forged_dists):.4f})')
		
		separation = np.mean(forged_dists) - np.mean(genuine_dists)
		plt.axvline(np.mean(genuine_dists), color='blue', linestyle='--', alpha=0.7)
		plt.axvline(np.mean(forged_dists), color='red', linestyle='--', alpha=0.7)
		
		plt.title(f'Distance Distribution - Epoch {epoch}\nSeparation: {separation:.4f}', fontsize=14)
		plt.xlabel('Embedding Distance', fontsize=12)
		plt.ylabel('Frequency', fontsize=12)
		plt.legend(fontsize=12)
		plt.grid(alpha=0.3)
		
		plt.savefig(os.path.join(figures_dir, f"distance_dist_epoch_{epoch}.png"),
				   dpi=150, bbox_inches='tight')
		plt.close()
	

	def _plot_tsne_embeddings(self, embeddings: np.ndarray, labels: np.ndarray,
							  person_ids: np.ndarray, distances: np.ndarray,
							  epoch: int, figures_dir: str, n_samples: int = 3000):
		"""Create comprehensive t-SNE visualization."""
		# Sample for computational efficiency
		if len(embeddings) > n_samples:
			indices = np.random.choice(len(embeddings), n_samples, replace=False)
			embeddings = embeddings[indices]
			labels = labels[indices]
			person_ids = person_ids[indices]
			distances = distances[indices]
		
		# Compute t-SNE
		tsne = TSNE(n_components=2, random_state=42, perplexity=min(30, len(embeddings)-1))
		embeddings_2d = tsne.fit_transform(embeddings)
		
		fig, axes = plt.subplots(1, 3, figsize=(20, 6))
		
		# 1. Genuine vs Forged
		for label_val, color, name in [(0, 'red', 'Forged'), (1, 'blue', 'Genuine')]:
			mask = labels == label_val
			if mask.any():
				axes[0].scatter(embeddings_2d[mask, 0], embeddings_2d[mask, 1],
								c=color, label=name, alpha=0.6, s=20)
		axes[0].set_title(f'Genuine vs Forged - Epoch {epoch}')
		axes[0].legend()
		axes[0].grid(alpha=0.3)
		
		# 2. Person clusters
		unique_ids = np.unique(person_ids)
		colors = plt.cm.tab20(np.linspace(0, 1, min(20, len(unique_ids))))
		
		# Show top 15 most frequent IDs
		id_counts = {pid: np.sum(person_ids == pid) for pid in unique_ids}
		top_ids = sorted(id_counts.items(), key=lambda x: x[1], reverse=True)[:15]
		
		for idx, (pid, count) in enumerate(top_ids):
			mask = person_ids == pid
			color = colors[idx % len(colors)]
			
			# Plot the cluster points
			axes[1].scatter(embeddings_2d[mask, 0], embeddings_2d[mask, 1],
							c=[color], label=f'ID {pid} (n={count})', alpha=0.7, s=25)
			
			# Compute the centroid (mean position) of the points in this cluster
			centroid = np.mean(embeddings_2d[mask], axis=0)
			
			# Add the person ID text at the centroid
			axes[1].text(centroid[0], centroid[1], f'ID {pid}', fontsize=10, color='black', alpha=0.8, ha='center')
		
		# Plot others in gray
		other_mask = ~np.isin(person_ids, [pid for pid, _ in top_ids])
		if other_mask.any():
			axes[1].scatter(embeddings_2d[other_mask, 0], embeddings_2d[other_mask, 1],
							c='gray', label='Others', alpha=0.3, s=15)
		
		axes[1].set_title(f'Person Clusters - Epoch {epoch}')
		axes[1].legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=8)
		axes[1].grid(alpha=0.3)
		
		# 3. Distance-based coloring
		scatter = axes[2].scatter(embeddings_2d[:, 0], embeddings_2d[:, 1],
								  c=distances, cmap='viridis', alpha=0.7, s=20)
		plt.colorbar(scatter, ax=axes[2], label='Distance')
		axes[2].set_title(f'Distance Visualization - Epoch {epoch}')
		axes[2].grid(alpha=0.3)
		
		plt.tight_layout()
		plt.savefig(os.path.join(figures_dir, f"tsne_epoch_{epoch}.png"),
					dpi=150, bbox_inches='tight')
		plt.close()

		
	def train(self):
		"""Main training loop."""
		torch.backends.cudnn.benchmark = True
		self.logger.info(f"Starting training on device: {self.device}")
		
		# Prepare components
		train_dataset, val_dataset = self._prepare_datasets()
		model, optimizer, scheduler = self._setup_model_and_optimizer()
		scaler = torch.amp.GradScaler(self.device.type, enabled=(self.device.type == "cuda"))
		
		# MLflow setup
		with self.mlflow_manager.start_run(run_id=self.config.run_id):
			run_id = mlflow.active_run().info.run_id
			self.mlflow_manager.log_config(self.config)
			
			# Setup checkpoints
			checkpoint_dir, figures_dir = self._setup_checkpoint_management(run_id)
			
			# Load checkpoint if specified
			start_epoch = self._load_checkpoint(model, optimizer, scheduler, scaler)
			
			# Data loaders
			val_loader = DataLoader(
				val_dataset, batch_size=self.config.batch_size, shuffle=False,
				num_workers=4, pin_memory=True, prefetch_factor=2
			)
			
			# Training loop
			for epoch in range(start_epoch, self.config.max_epochs + 1):
				self.current_epoch = epoch
				
				# Update curriculum
				train_dataset.set_epoch(epoch)
				train_loader = DataLoader(
					train_dataset, batch_size=self.config.batch_size, shuffle=True,
					num_workers=4, pin_memory=True, persistent_workers=True, prefetch_factor=2
				)
				
				# Training phase
				train_metrics = self.train_epoch(model, train_loader, optimizer, scaler, epoch)
				
				# Validation phase
				threshold, eer, val_results = self.validate_epoch(model, val_loader, epoch)
				
				# Logging
				self._log_epoch_metrics(train_metrics, val_results, epoch)
				
				# Visualizations
				self._create_visualizations(val_results, epoch, figures_dir)
				
				# Model checkpoint management
				is_best = eer < self.best_eer
				if is_best:
					self.best_eer = eer
					self.patience_counter = 0
				else:
					self.patience_counter += 1
				
				self._save_checkpoint(
					model, optimizer, scheduler, scaler, epoch,
					threshold, eer, checkpoint_dir, is_best
				)
				
				# Early stopping
				if self.patience_counter >= self.config.patience:
					self.logger.info(f"Early stopping after {self.config.patience} epochs without improvement")
					break
				
				# Learning rate scheduling
				if self.config.lr_scheduler == "cosine":
					scheduler.step()
				else:
					scheduler.step(eer)
				
				# Memory cleanup
				gc.collect()
				torch.cuda.empty_cache()
			
			# Final logging
			mlflow.log_metric("final_best_eer", self.best_eer)
			self.logger.info(f"Training completed. Best EER: {self.best_eer:.4f}")

# ----------------------------
# Image Processing Utilities
# ----------------------------
def estimate_background_color_pil(image: Image.Image, border_width: int = 10, 
								method: str = "median") -> np.ndarray:
	"""Estimate background color from image borders."""
	if image.mode != 'RGB':
		image = image.convert('RGB')
	
	np_img = np.array(image)
	h, w, _ = np_img.shape
	
	# Extract border pixels
	top = np_img[:border_width, :, :].reshape(-1, 3)
	bottom = np_img[-border_width:, :, :].reshape(-1, 3)
	left = np_img[:, :border_width, :].reshape(-1, 3)
	right = np_img[:, -border_width:, :].reshape(-1, 3)
	
	all_border_pixels = np.concatenate([top, bottom, left, right], axis=0)
	
	if method == "mean":
		return np.mean(all_border_pixels, axis=0).astype(np.uint8)
	else:
		return np.median(all_border_pixels, axis=0).astype(np.uint8)

def replace_background_with_white(image_name: str, folder_img: str, 
								tolerance: int = 40, method: str = "median", 
								remove_bg: bool = False) -> Image.Image:
	"""Replace background with white based on border color estimation."""
	image_path = os.path.join(folder_img, image_name)
	image = Image.open(image_path).convert("RGB")
	
	if not remove_bg:
		return image
	
	np_img = np.array(image)
	bg_color = estimate_background_color_pil(image, method=method)
	
	# Create mask for background pixels
	diff = np.abs(np_img.astype(np.int32) - bg_color.astype(np.int32))
	mask = np.all(diff < tolerance, axis=2)
	
	# Replace background with white
	result = np_img.copy()
	result[mask] = [255, 255, 255]
	
	return Image.fromarray(result)

# ----------------------------
# Main Execution
# ----------------------------
def main():
    """Main execution function with aggressive curriculum."""
    # Test distance ranges first
    print("\n[INFO] Testing distance ranges for margin calibration...")
    dummy_emb1 = F.normalize(torch.randn(1000, 128), p=2, dim=1)
    dummy_emb2 = F.normalize(torch.randn(1000, 128), p=2, dim=1)
    dummy_distances = F.pairwise_distance(dummy_emb1, dummy_emb2).numpy()
    print(f"Random embeddings: mean={dummy_distances.mean():.3f}, std={dummy_distances.std():.3f}")
    print(f"Expected margin range: {dummy_distances.std() * 0.5:.3f} - {dummy_distances.std() * 1.5:.3f}")
    
    # Aggressive curriculum configuration
    CONFIG.model_name = "resnet34"
    CONFIG.embedding_dim = 128
    CONFIG.max_epochs = 20  # Shorter with aggressive curriculum
    CONFIG.head_lr = 2e-3   # Higher for faster adaptation
    CONFIG.backbone_lr = 1e-4
    CONFIG.curriculum_strategy = "progressive"
    
    # AGGRESSIVE SETTINGS
    CONFIG.initial_hard_ratio = 0.4   # Start much higher
    CONFIG.final_hard_ratio = 0.85    # Target very high
    CONFIG.curriculum_warmup_epochs = 1  # Very short warmup
    CONFIG.batch_size = 256  # Smaller batches for more frequent updates
    CONFIG.grad_accum_steps = 8  # Smaller accumulation
    
    CONFIG.tracking_uri = "http://127.0.0.1:5555"
	#CONFIG.run_id = "aa58e3a1f3314351bc1dd2b82ab156ad"
	#CONFIG.last_epoch_weights = "../../model/models_checkpoints/aa58e3a1f3314351bc1dd2b82ab156ad/best_model.pth"
    
    trainer = SignatureTrainer(CONFIG)
    trainer.train()
if __name__ == "__main__":
	main()