import os
import glob

from sklearn.externals.array_api_compat.numpy import imag

# 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 MODELACH
# ==========================================================

for model_name, (model, model_type, task) in models.items():

	print(f"\n -> Generowanie LIME dla modelu: {model_name}...")

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

	explanations = []

	# ======================================================
	# GENEROWANIE EXPLANATIONS
	# ======================================================

	for img_name, img_array in loaded_images:

		print("\n" + "=" * 50)
		print(f"Przetwarzanie zdjęcia: {img_name}")
		print("=" * 50)

		# liczba klas zależna od tasku
		n_classes = 2 if task == "binary" else 3

		explanation = explainer.explain_instance(
			img_array.astype(np.double),
			predict_fn,
			top_labels=n_classes,
			hide_color=0,
			num_samples=35
		)

		explanations.append((img_name, explanation))

	# ======================================================
	# TWORZENIE SIATKI
	# ======================================================

	rows = len(explanations)
	cols = 2 if task == "binary" else 3

	fig, ax = plt.subplots(
		rows,
		cols,
		figsize=(4 * cols, 4 * rows)
	)

	fig.suptitle(
		f"LIME explanations - {model_name}",
		fontsize=18
	)

	# poprawka dla przypadku rows == 1
	if rows == 1:
		ax = np.expand_dims(ax, axis=0)

	# ======================================================
	# WYPEŁNIANIE WYKRESU
	# ======================================================

	for i, (img_name, explanation) in enumerate(explanations):

		for label in explanation.top_labels:

			temp, mask = explanation.get_image_and_mask(
				label,
				positive_only=True,
				num_features=10,
				hide_rest=False
			)

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

			ax[i, label].imshow(
				mark_boundaries(temp / temp.max() if temp.max() > 0 else temp, mask), cmap='viridis'
			)

			ax[i, label].set_title(
				f"{img_name}\n{class_name}",
				fontsize=8
			)

			ax[i, label].axis("off")

	# ======================================================
	# TYTUŁ + ZAPIS
	# ======================================================

	#fig.tight_layout()

	output_path = os.path.join(
		OUTPUT_DIR,
		f"lime_grid_{model_name.replace(' ', '_')}.jpg"
	)

	fig.savefig(
		output_path,
		format="jpg",
		dpi=300,
		bbox_inches="tight"
	)

	print(f"[OK] Zapisano wykres: {output_path}")

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