import os
import glob

# 1. WYMUSZENIE TRYBU OFFLINE DLA HUGGING FACE (Blokada wysyłania zapytania do sieci)
os.environ["TRANSFORMERS_OFFLINE"] = "1"
os.environ["HF_DATASETS_OFFLINE"] = "1"
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"

import cv2
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from transformers import ViTForImageClassification
from lime import lime_image
from skimage.segmentation import mark_boundaries
from torchvision import transforms
from PIL import Image

# --- KONFIGURACJA ---
IMG_SIZE_CNN = 512
IMG_SIZE_VIT = 384
MODELS_DIR = "models_up"
IMAGES_DIR = "images"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
OUTPUT_DIR = "lime_results"
os.makedirs(OUTPUT_DIR, exist_ok=True)

print(f"Urządzenie: {DEVICE}")

# Transformacje dla ViT
vit_transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((IMG_SIZE_VIT, IMG_SIZE_VIT)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# --- DEFINICJE ARCHITEKTUR (Zgodne z Twoim treningiem) ---

class CNNBinary(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, 7, stride=2, padding=3), nn.BatchNorm2d(64), nn.ReLU(),
            nn.MaxPool2d(3, stride=2, padding=1),
            nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(),
            nn.Conv2d(128, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(128, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(),
            nn.Conv2d(256, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(256, 512, 3, padding=1), nn.BatchNorm2d(512), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(), nn.Dropout(0.5),
            nn.Linear(512, 256), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(256, 1),  # 1 wyjście dla BCEWithLogitsLoss
        )
    def forward(self, x): return self.classifier(self.features(x))

class CNNTrinary(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, 7, stride=2, padding=3), nn.BatchNorm2d(64), nn.ReLU(),
            nn.MaxPool2d(3, stride=2, padding=1),
            nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(),
            nn.Conv2d(128, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(128, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(),
            nn.Conv2d(256, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(256, 512, 3, padding=1), nn.BatchNorm2d(512), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(), nn.Dropout(0.5),
            nn.Linear(512, 256), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(256, 3),  # 3 wyjścia dla CrossEntropyLoss
        )
    def forward(self, x): return self.classifier(self.features(x))

class ViTWrapper(nn.Module):
    def __init__(self, num_labels: int):
        super().__init__()
        self.vit = ViTForImageClassification.from_pretrained(
            "google/vit-base-patch16-384",
            num_labels=num_labels,
            ignore_mismatched_sizes=True,
            local_files_only=True  # Pobieraj tylko lokalnie z cache
        )
    def forward(self, x): return self.vit(pixel_values=x).logits

# --- FUNKCJE POMOCNICZE ---

def preprocess_cnn(images):
    out = []
    for img in images:
        img_res = cv2.resize(img, (IMG_SIZE_CNN, IMG_SIZE_CNN))
        img_res = img_res.astype(np.float32) / 255.0
        out.append(img_res)
    out = np.stack(out)
    return torch.tensor(out).permute(0, 3, 1, 2).float().to(DEVICE)

def preprocess_vit(images):
    out = []
    for img in images:
        img_t = vit_transform(img)
        out.append(img_t)
    return torch.stack(out).to(DEVICE)

def make_predict_fn(model, model_type, task="binary"):
    model.eval()

    def predict(images):
        if model_type == "cnn":
            x = preprocess_cnn(images)
        else:
            x = preprocess_vit(images)
            
        with torch.no_grad():
            logits = model(x)
            
            if task == "binary":
                # Konwersja 1 logitu na format prawdopodobieństwa dwuklasowego dla LIME
                logits = logits.squeeze(-1)
                prob_class_1 = torch.sigmoid(logits).cpu().numpy()
                prob_class_0 = 1.0 - prob_class_1
                probs = np.column_stack((prob_class_0, prob_class_1))
            else:
                probs = torch.softmax(logits, dim=1).cpu().numpy()
        return probs
    return predict

# --- ŁADOWANIE MODELI Z TWOJEGO FOLDERU 'models' ---
#models_up/binary_cnn
print("\nŁadowanie wag modeli z dysku...")
models = {}

# Dokładne odwzorowanie Twojej nowej struktury plików (.pth w jednym folderu)
paths = {
    "Binary CNN":  (CNNBinary(), os.path.join(MODELS_DIR, "binary_cnn","model_best.pth"), "binary", "cnn"),
    "Trinary CNN": (CNNTrinary(), os.path.join(MODELS_DIR, "trinary_cnn","model_best.pth"), "trinary", "cnn"),
    "Binary ViT":  (ViTWrapper(num_labels=1), os.path.join(MODELS_DIR, "binary_vit","model_best.pth"), "binary", "vit"),
    "Trinary ViT": (ViTWrapper(num_labels=3), os.path.join(MODELS_DIR, "trinary_vit","model_best.pth"), "trinary", "vit"),
}

for name, (model, path, task, m_type) in paths.items():
    if os.path.exists(path):
        try:
            model.load_state_dict(torch.load(path, map_location=DEVICE))
            model.to(DEVICE)
            model.eval()
            models[name] = (model, m_type, task)
            print(f"  [OK] Załadowano model: {name} ({path})")
        except Exception as e:
            print(f"  [BŁĄD] Nie można załadować {name}. Szczegóły: {e}")
    else:
        print(f"  [BRAK] Nie znaleziono pliku wag dla: {name} pod ścieżką: {path}")

if not models:
    print("\n[STOP] Nie udało się załadować żadnego modelu. Sprawdź nazwy plików w folderu 'models'.")
    exit()

# --- ŁADOWANIE ZDJĘĆ Z TWOJEGO FOLDERU 'images' ---

print("\nSzukanie zdjęć w folderze 'images'...")
valid_extensions = ("*.jpg", "*.jpeg", "*.png", "*.bmp", "*.webp", "*.tiff")
image_paths = []
for ext in valid_extensions:
    image_paths.extend(glob.glob(os.path.join(IMAGES_DIR, ext)))

if not image_paths:
    print(f"[STOP] Nie znaleziono żadnych zdjęć w folderze '{IMAGES_DIR}'!")
    exit()

print(f"Znaleziono {len(image_paths)} zdjęć do przetestowania.")

loaded_images = []
for p in image_paths:
    try:
        img = Image.open(p).convert("RGB").resize((IMG_SIZE_CNN, IMG_SIZE_CNN))
        loaded_images.append((os.path.basename(p), np.array(img)))
    except Exception as e:
        print(f"  [POMINIĘTO] Nie można otworzyć pliku {p}: {e}")

# --- URUCHOMIENIE WYJAŚNIEŃ LIME ---

explainer = lime_image.LimeImageExplainer()

# Słowniki mapowania nazw klas
labels_map = {
    "binary": {0: "Artificial", 1: "Natural"},
    "trinary": {0: "Drawings", 1: "Games", 2: "Nature"}
}

print("\nRozpoczynam generowanie wyjaśnień LIME...")

# Pętla po każdym załadowanym modelu
for model_name, (model, model_type, task) in models.items():
    print(f" -> Generowanie LIME dla modelu: {model_name}...")
    # Pętla po każdym załadowanym zdjęciu
    for img_name, img_array in loaded_images:
        print(f"\n" + "="*50)
        print(f"Przetwarzanie zdjęcia: {img_name}")
        print("="*50)

        predict_fn = make_predict_fn(model, model_type, task=task)

        # Wyjaśnienie LIME (num_samples=35 dla przyspieszenia działania ViT)
        explanation = explainer.explain_instance(
            img_array.astype(np.double), 
            predict_fn,
            top_labels=1,
            hide_color=0,
            num_samples=35
        )

        label = explanation.top_labels[0]
        temp, mask = explanation.get_image_and_mask(
            label,
            positive_only=True,  # Pokaż tylko obszary, które wpłynęły na decyzję pozytywnie
            num_features=5,       # Pokaż 5 najważniejszych superpikseli
            hide_rest=False
        )

        class_name = labels_map[task].get(label, str(label))

        # Przygotowanie bezpiecznej nazwy pliku
        safe_img_name = os.path.splitext(os.path.basename(img_name))[0] # Pobierz bez rozszerzenia
        safe_model_name = model_name.replace(" ", "_").replace("/", "_") # Bez spacji
        
        output_filename = f"{safe_img_name}_{safe_model_name}.png"
        file_path = os.path.join(OUTPUT_DIR, output_filename)
        
        # Tworzenie wykresu
        plt.figure(figsize=(6, 6))
        plt.title(f"Zdjęcie: {img_name}\nModel: {model_name}\nPredykcja: {class_name}", fontsize=10)
        plt.imshow(mark_boundaries(temp / temp.max() if temp.max() > 0 else temp, mask), cmap='viridis') # Dodajemy cmap dla kolorów
        plt.axis("off")
        plt.tight_layout()
        
        # Zapisywanie
        plt.savefig(file_path, dpi=300, bbox_inches='tight')
        print(f"      ↳ Zapisano: {output_filename}")
        
        # Opcjonalnie wyświetlanie okna (możesz usunąć, jeśli generujesz dużo plików)
        plt.show() 
        
        # Zwalniamy pamięć
        plt.close() 

# Koniec pętli po modelach
print("\n[WYNIK] Wszystkie analizy zakończone. Zobacz folder '" + OUTPUT_DIR + "'")