import cv2
import torch
import torch.nn.functional as F
import numpy as np
from transformers import PreTrainedModel, PretrainedConfig

# ==========================================
# 1. Hugging Face Configuration
# ==========================================
class PhorensicsConfig(PretrainedConfig):
    model_type = "phorensics"
    
    def __init__(self, max_dim=600, patch_size=128, stride=64, **kwargs):
        self.max_dim = max_dim
        self.patch_size = patch_size
        self.stride = stride
        super().__init__(**kwargs)

# ==========================================
# 2. Native PyTorch Physics Engine
# ==========================================
class PhorensicsModel(PreTrainedModel):
    config_class = PhorensicsConfig

    def __init__(self, config):
        super().__init__(config)
        self.patch_size = config.patch_size
        self.stride = config.stride
        
        # THE FIX: Hugging Face requires at least one parameter to determine model dtype.
        # This dummy parameter satisfies the push_to_hub internal checks.
        self.dummy_param = torch.nn.Parameter(torch.empty(0, dtype=torch.float32))
        
        # STRUCTURAL TENSORS: Stored permanently in the .safetensors file.
        laplacian_kernel = torch.tensor([[[[0., 1., 0.], [1., -4., 1.], [0., 1., 0.]]]], dtype=torch.float32)
        self.register_buffer("noise_kernel", laplacian_kernel)

        self.register_buffer("safe_z_threshold", torch.tensor(2.5, dtype=torch.float32))

    def _extract_patches(self, tensor):
        """Hardware-accelerated tensor unfolding (PyTorch native sliding window)"""
        c, h, w = tensor.shape
        pad_h = max(0, self.patch_size - h) if h < self.patch_size else 0
        pad_w = max(0, self.patch_size - w) if w < self.patch_size else 0
        if pad_h > 0 or pad_w > 0:
            tensor = F.pad(tensor, (0, pad_w, 0, pad_h), mode='reflect')

        patches = tensor.unfold(1, self.patch_size, self.stride).unfold(2, self.patch_size, self.stride)
        return patches.contiguous().view(c, -1, self.patch_size, self.patch_size).permute(1, 0, 2, 3)

    @torch.no_grad()
    def forward(self, image_path):
        """
        The native inference pass. Automatically handles CPU/CUDA routing.
        """
        # 1. Load and Standardize Image
        img_bgr = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
        if img_bgr is None: 
            raise ValueError(f"Failed to load image: {image_path}")
            
        h, w = img_bgr.shape[:2]
        if max(h, w) > self.config.max_dim:
            scale = self.config.max_dim / max(h, w)
            img_bgr = cv2.resize(img_bgr, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)
        
        # 2. TENSOR CONVERSION & GPU ACCELERATION
        gray_np = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
        # Move the image tensor to whatever device the model is currently on (CPU/CUDA)
        t_gray = torch.from_numpy(gray_np).to(self.device).float().unsqueeze(0).unsqueeze(0)

        # 3. SPATIAL TENSOR CONVOLUTION (PRNU Noise Validation)
        # Apply the .safetensors kernel across the image
        noise_map = F.conv2d(t_gray, self.noise_kernel, padding=1).squeeze(0)
        p_noise = self._extract_patches(noise_map)
        
        # Calculate robust Z-Score using pure vectorized PyTorch math
        patch_variances = torch.var(p_noise.view(p_noise.shape[0], -1), dim=1)
        median_var = torch.median(patch_variances)
        mad = torch.median(torch.abs(patch_variances - median_var)) + 1e-3
        z_scores = 0.6745 * torch.abs(patch_variances - median_var) / mad
        
        max_noise_z = torch.quantile(z_scores, 0.98).item()

        # 4. FAST FOURIER TRANSFORM (Global Spectral Decay)
        f = np.fft.fft2(gray_np)
        power = np.abs(np.fft.fftshift(f))**2
        h_fft, w_fft = power.shape
        cy, cx = h_fft // 2, w_fft // 2
        y, x = np.indices((h_fft, w_fft))
        r = np.sqrt((x - cx)**2 + (y - cy)**2).astype(int)
        
        radial_profile = np.bincount(r.ravel(), power.ravel()) / np.maximum(np.bincount(r.ravel()), 1)
        r_vals = np.arange(1, len(radial_profile))
        p_vals = radial_profile[1:]
        
        slope, _ = np.polyfit(np.log10(r_vals), np.log10(p_vals + 1e-10), 1)
        alpha = (-slope) * 0.70  
        
        # 5. ENGINE FUSION
        threshold = self.safe_z_threshold.item()
        noise_threat = max(0.0, min(((max_noise_z - threshold) / 1.5) * 100.0, 100.0))
        fft_threat = 99.9 if (2.0 <= alpha <= 3.5) else max(((alpha - 1.0) / 1.0) * 49.9, 0.0)

        final_prob = (noise_threat * 0.40) + (fft_threat * 0.60)

        return {
            "is_fake": final_prob >= 50.0,
            "threat_probability": round(min(final_prob, 99.9), 2),
            "max_noise_z": round(max_noise_z, 2),
            "fft_alpha_value": round(alpha, 3)
        }