import os
import torch
import numpy as np
import pandas as pd
from PIL import Image
from tqdm import tqdm
import torch.nn as nn
from torchvision import models
import torch.nn.functional as F
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
from typing import Dict, List, Tuple, Optional, Union
from dataclasses import dataclass
import warnings
warnings.filterwarnings('ignore')

# ----------------------------
# Configuration
# ----------------------------
@dataclass
class InferenceConfig:
    # Model Configuration
    model_name: str = "resnet34"
    embedding_dim: int = 128
    normalize_embeddings: bool = True
    checkpoint_path: str = "../../model/models_checkpoints/best_model.pth"
    
    # Inference Settings
    batch_size: int = 32
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    distance_threshold: float = 0.5  # Will be loaded from checkpoint
    
    # Data Settings
    remove_bg: bool = False
    num_workers: int = 4

# Global configuration
CONFIG = InferenceConfig()

# ----------------------------
# Model Architecture (Same as training)
# ----------------------------
class ResNetBackbone(nn.Module):
    """ResNet backbone feature extractor."""
    
    def __init__(self, model_name: str = "resnet34"):
        super().__init__()
        
        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}")
        
        # 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 forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.resnet(x)

class AdvancedEmbeddingHead(nn.Module):
    """Embedding head to project features to embedding space."""
    
    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
        
        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:
            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)
        return self.layers(x)

class SiameseSignatureNetwork(nn.Module):
    """Siamese network for signature verification."""
    
    def __init__(self, config: InferenceConfig = CONFIG):
        super().__init__()
        self.config = config
        
        # Initialize backbone
        self.backbone = ResNetBackbone(model_name=config.model_name)
        backbone_dim = self.backbone.output_dim
        
        # Initialize embedding head
        self.embedding_head = AdvancedEmbeddingHead(
            input_dim=backbone_dim,
            embedding_dim=config.embedding_dim,
            dropout=0.0  # No dropout during inference
        )
        
        self.normalize_embeddings = config.normalize_embeddings
        self.distance_threshold = config.distance_threshold
    
    def forward(self, img1: torch.Tensor, img2: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Forward pass for inference."""
        # Extract features
        f1 = self.backbone(img1)
        f2 = self.backbone(img2)
        
        # Get embeddings
        emb1 = self.embedding_head(f1)
        emb2 = self.embedding_head(f2)
        
        # Normalize if configured
        if self.normalize_embeddings:
            emb1 = F.normalize(emb1, p=2, dim=1)
            emb2 = F.normalize(emb2, p=2, dim=1)
        
        return emb1, emb2
    
    def predict_pair(self, img1: torch.Tensor, img2: torch.Tensor, 
                    threshold: Optional[float] = None) -> Dict[str, torch.Tensor]:
        """Predict similarity between image pairs."""
        self.eval()
        with torch.no_grad():
            emb1, emb2 = self(img1, img2)
            distances = F.pairwise_distance(emb1, emb2)
            
            thresh = threshold if threshold is not None else self.distance_threshold
            predictions = (distances < thresh).long()
            
            # Convert distance to similarity score (0-1, higher is more similar)
            similarities = 1.0 / (1.0 + distances)
            
            return {
                'predictions': predictions,
                'distances': distances,
                'similarities': similarities,
                'threshold': torch.tensor(thresh)
            }

# ----------------------------
# Dataset for Batch Prediction
# ----------------------------
class PredictionDataset(Dataset):
    """Dataset for batch prediction from Excel."""
    
    def __init__(self, excel_path: str, image_folder: str, config: InferenceConfig = CONFIG):
        self.image_folder = image_folder
        self.config = config
        self.data = pd.read_excel(excel_path)
        self.transform = self._get_transforms()
        
        # Check required columns
        required_cols = ['image_1_path', 'image_2_path']
        missing_cols = [col for col in required_cols if col not in self.data.columns]
        if missing_cols:
            raise ValueError(f"Missing required columns: {missing_cols}")
    
    def _get_transforms(self) -> transforms.Compose:
        """Get image transforms for inference."""
        return transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])
    
    def __len__(self) -> int:
        return len(self.data)
    
    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, int]:
        """Return image pair and index."""
        row = self.data.iloc[idx]
        
        img1 = self._load_image(row['image_1_path'])
        img2 = self._load_image(row['image_2_path'])
        
        return img1, img2, idx
    
    def _load_image(self, image_path: str) -> torch.Tensor:
        """Load and transform image."""
        image = replace_background_with_white(
            image_path, self.image_folder, 
            remove_bg=self.config.remove_bg
        )
        return self.transform(image)

# ----------------------------
# Image Processing
# ----------------------------
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 Prediction Class
# ----------------------------
class SignatureVerifier:
    """Main class for signature verification predictions."""
    
    def __init__(self, config: InferenceConfig = CONFIG):
        self.config = config
        self.device = torch.device(config.device)
        self.model = self._load_model()
        self.transform = self._get_transforms()
    
    def _get_transforms(self) -> transforms.Compose:
        """Get image transforms."""
        return transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])
    
    def _load_model(self) -> SiameseSignatureNetwork:
        """Load model from checkpoint."""
        print(f"Loading model from: {self.config.checkpoint_path}")
        
        # Initialize model
        model = SiameseSignatureNetwork(self.config)
        
        # Load checkpoint
        checkpoint = torch.load(self.config.checkpoint_path, map_location=self.device, weights_only=False)
        
        # Load model state
        if 'model_state_dict' in checkpoint:
            model.load_state_dict(checkpoint['model_state_dict'])
        else:
            # If checkpoint is just the state dict
            model.load_state_dict(checkpoint)
        
        # Load threshold if available
        if 'prediction_threshold' in checkpoint:
            model.distance_threshold = checkpoint['prediction_threshold']
            print(f"Loaded threshold: {model.distance_threshold:.4f}")
        
        # Load best EER if available
        if 'best_eer' in checkpoint:
            print(f"Model best EER: {checkpoint['best_eer']:.4f}")
        
        model = model.to(self.device)
        model.eval()
        
        print("Model loaded successfully!")
        return model
    
    def predict_single_pair(self, image1_path: str, image2_path: str, 
                           image_folder: str = "") -> Dict[str, float]:
        """Predict similarity for a single pair of images."""
        # Load images
        img1 = replace_background_with_white(
            image1_path, image_folder, remove_bg=self.config.remove_bg
        )
        img2 = replace_background_with_white(
            image2_path, image_folder, remove_bg=self.config.remove_bg
        )
        
        # Transform
        img1_tensor = self.transform(img1).unsqueeze(0).to(self.device)
        img2_tensor = self.transform(img2).unsqueeze(0).to(self.device)
        
        # Predict
        results = self.model.predict_pair(img1_tensor, img2_tensor)
        
        return {
            'is_genuine': bool(results['predictions'].item()),
            'distance': float(results['distances'].item()),
            'similarity_score': float(results['similarities'].item()),
            'threshold': float(results['threshold'].item())
        }
    
    def predict_from_excel(self, excel_path: str, image_folder: str, 
                          output_path: Optional[str] = None) -> pd.DataFrame:
        """Batch prediction from Excel file."""
        # Create dataset and dataloader
        dataset = PredictionDataset(excel_path, image_folder, self.config)
        dataloader = DataLoader(
            dataset, 
            batch_size=self.config.batch_size,
            shuffle=False,
            num_workers=self.config.num_workers,
            pin_memory=True
        )
        
        # Prediction storage
        all_predictions = []
        all_distances = []
        all_similarities = []
        
        # Predict in batches
        print(f"Processing {len(dataset)} pairs...")
        with torch.no_grad():
            for img1_batch, img2_batch, indices in tqdm(dataloader):
                img1_batch = img1_batch.to(self.device)
                img2_batch = img2_batch.to(self.device)
                
                results = self.model.predict_pair(img1_batch, img2_batch)
                
                all_predictions.extend(results['predictions'].cpu().numpy())
                all_distances.extend(results['distances'].cpu().numpy())
                all_similarities.extend(results['similarities'].cpu().numpy())
        
        # Create results dataframe
        results_df = dataset.data.copy()
        results_df['prediction'] = all_predictions
        results_df['is_genuine'] = results_df['prediction'].astype(bool)
        results_df['distance'] = all_distances
        results_df['similarity_score'] = all_similarities
        results_df['threshold'] = self.model.distance_threshold
        
        # Save if output path provided
        if output_path:
            results_df.to_excel(output_path, index=False)
            print(f"Results saved to: {output_path}")
        
        return results_df
    
    def update_threshold(self, new_threshold: float):
        """Update the decision threshold."""
        self.model.distance_threshold = new_threshold
        print(f"Threshold updated to: {new_threshold:.4f}")

# Initialize verifier
config = InferenceConfig(
	checkpoint_path="../../../../model/models_checkpoints/fa7e1bdc01814016ac8220bfbf1eb691/best_model.pth",
	batch_size=32,
	device="cuda" if torch.cuda.is_available() else "cpu"
)

verifier = SignatureVerifier(config)

'''
# Example 1: Single pair prediction
print("\n--- Single Pair Prediction ---")
result = verifier.predict_single_pair(
	image1_path="sig1.png",
	image2_path="sig2.png",
	image_folder="../../data/classify/preprared_data/images/"
)
'''

# Example 2: Batch prediction from Excel
print("\n--- Batch Prediction from Excel ---")
results_df = verifier.predict_from_excel(
	excel_path="../../../../data/classify/preprared_data/labels/test_pairs_balanced_v12.xlsx",
	image_folder="../../../../data/classify/preprared_data/images/",
	output_path="./predictions_output.xlsx"
)

# Print summary
genuine_count = results_df['is_genuine'].sum()
total_count = len(results_df)
print(f"\nPrediction Summary:")
print(f"Total pairs: {total_count}")
print(f"Genuine predictions: {genuine_count} ({100*genuine_count/total_count:.1f}%)")
print(f"Forged predictions: {total_count - genuine_count} ({100*(total_count-genuine_count)/total_count:.1f}%)")