from pathlib import Path
import cv2
import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms as T
from PIL import Image
from torchvision import models


class ParcelGuardPredictor:
    def __init__(self, backbone_path, threshold=0.70, device=None):
        self.device = torch.device(device or ('cuda' if torch.cuda.is_available() else 'cpu'))
        self.threshold = float(threshold)

        backbone = models.resnet18(weights=None)
        self.extractor = nn.Sequential(
            backbone.conv1, backbone.bn1, backbone.relu, backbone.maxpool,
            backbone.layer1, backbone.layer2
        )
        state = torch.load(backbone_path, map_location='cpu', weights_only=True)
        self.extractor.load_state_dict(state, strict=True)
        self.extractor.to(self.device).eval()

        self.transform = T.Compose([
            T.Resize((224, 224)),
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225]),
        ])

    @torch.inference_mode()
    def _features(self, image_path):
        path = Path(image_path)
        if not path.is_file():
            raise FileNotFoundError(f'Image not found: {path}')
        image = Image.open(path).convert('RGB')
        tensor = self.transform(image).unsqueeze(0).to(self.device)
        return self.extractor(tensor)[0].cpu().numpy()

    @staticmethod
    def _distance(query_features, reference_features):
        q = query_features.transpose(1, 2, 0)
        r = reference_features.transpose(1, 2, 0)
        q = q / (np.linalg.norm(q, axis=-1, keepdims=True) + 1e-8)
        r = r / (np.linalg.norm(r, axis=-1, keepdims=True) + 1e-8)
        return 1.0 - np.sum(q * r, axis=-1)

    def predict(self, query_image, reference_images):
        if not reference_images:
            raise ValueError('Provide at least one reference image.')
        query_features = self._features(query_image)
        maps = []
        for ref_path in reference_images:
            ref_features = self._features(ref_path)
            maps.append(self._distance(query_features, ref_features))
        anomaly_map = np.min(np.stack(maps), axis=0)
        score = float(np.percentile(anomaly_map, 95))
        prediction = int(score >= self.threshold)

        image = cv2.imread(str(query_image))
        if image is None:
            raise ValueError('OpenCV could not read the query image.')
        heatmap = cv2.resize(anomaly_map.astype(np.float32),
                              (image.shape[1], image.shape[0]))
        heatmap = cv2.normalize(heatmap, None, 0, 255,
                                 cv2.NORM_MINMAX).astype(np.uint8)
        color = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
        overlay = cv2.addWeighted(image, 0.60, color, 0.40, 0)

        return {
            'anomaly_score_p95': score,
            'predicted_defect': prediction,
            'label': 'DEFECT' if prediction else 'NORMAL',
            'threshold': self.threshold,
            'anomaly_map': anomaly_map,
            'overlay_bgr': overlay,
        }