#Imports
import os
import sys
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

os.environ["TRANSFORMERS_OFFLINE"] = "1"
os.environ["HF_DATASETS_OFFLINE"] = "1"
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"

import numpy as np
import torch
import torch.nn as nn
from lime import lime_image
from PIL import Image
from skimage.segmentation import mark_boundaries
from torchvision import transforms
from transformers import ViTForImageClassification, ViTConfig

from PyQt6.QtCore import Qt, QThread, pyqtSignal
from PyQt6.QtGui import QAction, QImage, QPixmap
from PyQt6.QtWidgets import (
	QApplication,
	QFileDialog,
	QHBoxLayout,
	QLabel,
	QMainWindow,
	QMessageBox,
	QPushButton,
	QScrollArea,
	QStatusBar,
	QVBoxLayout,
	QWidget,
	QButtonGroup,
)

#Constants
IMG_SIZE = 512
IMG_SIZE_VIT = 384
MODELS_DIR = "models_up"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
IMAGE_EXTENSIONS = ("*.jpg", "*.jpeg", "*.png", "*.bmp", "*.webp", "*.tiff")


LABELS_MAP = {
	"binary": {0: "Artificial", 1: "Natural"},
	"trinary": {0: "Drawings", 1: "Games", 2: "Nature"},
}

DISPLAY_MODES = {
	"normal": "Normal image",
	"binary": "Explanation binary",
	"trinary": "Explanation trinary",
}

#Przełącznik trybów
class SegmentedControl(QWidget):
	modeChanged = pyqtSignal()

	def __init__(self, options: List[Tuple[str, str]]):
		super().__init__()
		layout = QHBoxLayout(self)
		layout.setContentsMargins(0, 0, 0, 0)
		layout.setSpacing(6)

		self.group = QButtonGroup(self)
		self.group.setExclusive(True)
		for idx, (label, mode) in enumerate(options):
			btn = QPushButton(label)
			btn.setCheckable(True)
			btn.setProperty("mode", mode)
			layout.addWidget(btn)
			self.group.addButton(btn, id=idx)

		first = self.group.buttons()[0] if self.group.buttons() else None
		if first:
			first.setChecked(True)

		self.group.buttonClicked.connect(lambda _btn: self.modeChanged.emit())

	def currentMode(self) -> str:
		for btn in self.group.buttons():
			if btn.isChecked():
				return btn.property("mode")
		return "normal"


def make_label(text: str) -> QLabel:
	lbl = QLabel(text)
	lbl.setWordWrap(True)
	return lbl

#Vision Transformer - 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]),
	]
)


class ViTWrapper(nn.Module):
	def __init__(self, num_labels: int):
		super().__init__()
		config = ViTConfig(
			num_labels=num_labels,
			image_size=IMG_SIZE_VIT,
			patch_size=16,
		)
		self.vit = ViTForImageClassification(config)

	def forward(self, x):
		return self.vit(pixel_values=x).logits


@dataclass
class AnalysisResult:
	image_path: str
	image_name: str
	image_array: np.ndarray
	display_array: np.ndarray
	display_mode: str
	binary_probs: np.ndarray
	trinary_probs: np.ndarray
	binary_label: str
	trinary_label: str
	binary_confidence: float
	trinary_confidence: float
	explanation_array: Optional[np.ndarray] = None
	explanation_label: Optional[str] = None


#Wczytanie modeli
def load_state_dict_robust(model: nn.Module, path: str) -> None:
	state = torch.load(path, map_location=DEVICE)
	if isinstance(state, dict) and "state_dict" in state:
		state = state["state_dict"]

	if isinstance(state, dict):
		cleaned = { (k[len("module."):] if k.startswith("module.") else k): v for k, v in state.items() }
		model.load_state_dict(cleaned)
		return

	model.load_state_dict(state)


def load_vit_models() -> Dict[str, Tuple[nn.Module, str]]:
	models: Dict[str, Tuple[nn.Module, str]] = {}
	candidates = {
		"binary": os.path.join(MODELS_DIR, "binary_vit", "model_best.pth"),
		"trinary": os.path.join(MODELS_DIR, "trinary_vit", "model_best.pth"),
	}

	for task, path in candidates.items():
		if not os.path.exists(path):
			raise FileNotFoundError(f"Nie znaleziono wag dla modelu {task} ViT: {path}")

		num_labels = 1 if task == "binary" else 3
		model = ViTWrapper(num_labels=num_labels)
		load_state_dict_robust(model, path)
		model.to(DEVICE)
		model.eval()
		models[task] = (model, path)

	return models

#Wykrywanie obrazów
def list_image_files(input_path: str) -> List[str]:
	if os.path.isdir(input_path):
		paths: List[str] = []
		allowed_extensions = tuple(extension.replace("*", "").lower() for extension in IMAGE_EXTENSIONS)
		for root, _, files in os.walk(input_path):
			for file_name in files:
				if file_name.lower().endswith(allowed_extensions):
					paths.append(os.path.join(root, file_name))
		return sorted(set(paths))

	if os.path.isfile(input_path):
		return [input_path]

	raise FileNotFoundError(input_path)


def load_image(path: str) -> np.ndarray:
	image = Image.open(path).convert("RGB").resize((IMG_SIZE, IMG_SIZE))
	return np.array(image)


def preprocess_vit(images: List[np.ndarray]) -> torch.Tensor:
	tensors = [vit_transform(image) for image in images]
	return torch.stack(tensors).to(DEVICE)


def predict_batch(model: nn.Module, images: List[np.ndarray], task: str) -> np.ndarray:
	batch = preprocess_vit(images)
	with torch.no_grad():
		logits = model(batch)
		if task == "binary":
			logits = logits.squeeze(-1)
			probability_1 = torch.sigmoid(logits).cpu().numpy()
			probability_0 = 1.0 - probability_1
			return np.column_stack((probability_0, probability_1))
		return torch.softmax(logits, dim=1).cpu().numpy()


def class_name(task: str, label: int) -> str:
	return LABELS_MAP[task].get(label, str(label))


def to_uint8_rgb(image_array: np.ndarray) -> np.ndarray:
	clipped = np.clip(image_array, 0, 255)
	if clipped.dtype != np.uint8:
		clipped = clipped.astype(np.uint8)
	return clipped

#LIME Explanation
def make_explanation_overlay(
	image_array: np.ndarray,
	probs: np.ndarray,
	task: str,
	explainer: lime_image.LimeImageExplainer,
) -> Tuple[np.ndarray, str]:
	prob_vector = np.asarray(probs, dtype=np.float32)
	if prob_vector.ndim == 0:
		prob_vector = prob_vector.reshape(1)
	if prob_vector.ndim == 1:
		prob_vector = np.tile(prob_vector, (1, 1))

	def predict_fn(images: List[np.ndarray]) -> np.ndarray:
		return np.repeat(prob_vector, len(images), axis=0)

	explanation = explainer.explain_instance(
		image_array.astype(np.double),
		predict_fn,
		top_labels=1,
		hide_color=0,
		num_samples=35,
	)
	label = int(np.argmax(probs))
	temp, mask = explanation.get_image_and_mask(
		label,
		positive_only=True,
		num_features=10,
		hide_rest=False,
	)
	temp = temp / (temp.max() if temp.max() > 0 else 1)
	overlay = mark_boundaries(temp, mask)
	overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
	return overlay, class_name(task, label)

#Analiza obrazu
def analyse_image(
	image_path: str,
	models: Dict[str, Tuple[nn.Module, str]],
	display_mode: str,
	explainer: lime_image.LimeImageExplainer,
) -> AnalysisResult:
	image_array = load_image(image_path)

	binary_model = models["binary"][0]
	trinary_model = models["trinary"][0]

	binary_probs = predict_batch(binary_model, [image_array], "binary")[0]
	trinary_probs = predict_batch(trinary_model, [image_array], "trinary")[0]

	binary_index = int(np.argmax(binary_probs))
	trinary_index = int(np.argmax(trinary_probs))

	explanation_array = None
	explanation_label = None
	if display_mode == "binary":
		explanation_array, explanation_label = make_explanation_overlay(
			image_array,
			binary_probs,
			"binary",
			explainer,
		)
		display_array = explanation_array
	elif display_mode == "trinary":
		explanation_array, explanation_label = make_explanation_overlay(
			image_array,
			trinary_probs,
			"trinary",
			explainer,
		)
		display_array = explanation_array
	else:
		display_array = image_array

	return AnalysisResult(
		image_path=image_path,
		image_name=os.path.basename(image_path),
		image_array=image_array,
		display_array=display_array,
		display_mode=display_mode,
		binary_probs=binary_probs,
		trinary_probs=trinary_probs,
		binary_label=class_name("binary", binary_index),
		binary_confidence=float(binary_probs[binary_index]),
		trinary_label=class_name("trinary", trinary_index),
		trinary_confidence=float(trinary_probs[trinary_index]),
		explanation_array=explanation_array,
		explanation_label=explanation_label,
	)


def image_to_pixmap(image_array: np.ndarray, max_width: int = 1000, max_height: int = 700) -> QPixmap:
	rgb = to_uint8_rgb(image_array)
	height, width, channels = rgb.shape
	bytes_per_line = channels * width
	image = QImage(rgb.data, width, height, bytes_per_line, QImage.Format.Format_RGB888)
	pixmap = QPixmap.fromImage(image.copy())
	return pixmap.scaled(max_width, max_height, Qt.AspectRatioMode.KeepAspectRatio, Qt.TransformationMode.SmoothTransformation)

#Funkcja wstępnego wyliczania klas i wyjaśnień przy wczytywaniu obrazów
class PrecomputeWorker(QThread):
	precomputed = pyqtSignal(str, str, object)
	progress = pyqtSignal(int, int)
	finished_all = pyqtSignal()

	def __init__(self, image_paths: List[str], models: Dict[str, Tuple[nn.Module, str]]):
		super().__init__()
		self.image_paths = image_paths
		self.models = models

	def run(self):
		explainer = lime_image.LimeImageExplainer()
		modes = list(DISPLAY_MODES.keys())
		total = len(self.image_paths) * len(modes)
		count = 0
		for p in self.image_paths:
			for mode in modes:
				try:
					result = analyse_image(p, self.models, mode, explainer)
					self.precomputed.emit(p, mode, result)
				except Exception:
					pass
				count += 1
				self.progress.emit(count, total)
		self.finished_all.emit()

#Główne okno PyQT6
class MainWindow(QMainWindow):
	def __init__(self):
		super().__init__()
		self.setWindowTitle("LZSI - ViT Image Inspector")
		self.resize(1320, 860)

		self.models = load_vit_models()
		self.image_paths: List[str] = []
		self.current_index = 0
		self.current_result: Optional[AnalysisResult] = None
		self.result_cache: Dict[Tuple[str, str], AnalysisResult] = {}

		self._build_ui()
		self._set_navigation_enabled(False)
		self.display_mode_control.setEnabled(False)
		self.statusBar().showMessage("Wczytaj pojedynczy obraz albo katalog obrazów.")

	def _build_ui(self) -> None:
		central = QWidget()
		root_layout = QHBoxLayout(central)

		controls_panel = QWidget()
		controls_layout = QVBoxLayout(controls_panel)

		self.open_file_button = QPushButton("Otwórz obraz")
		self.open_dir_button = QPushButton("Otwórz katalog")
		self.prev_button = QPushButton("Poprzedni")
		self.next_button = QPushButton("Następny")

		self.prev_button.clicked.connect(self.show_previous_image)
		self.next_button.clicked.connect(self.show_next_image)
		self.open_file_button.clicked.connect(self.open_image_file)
		self.open_dir_button.clicked.connect(self.open_image_directory)

		controls_layout.addWidget(self.open_file_button)
		controls_layout.addWidget(self.open_dir_button)
		controls_layout.addStretch(1)

		self.path_label = make_label("Brak wczytanego pliku.")

		self.counter_label = QLabel("")
		self.counter_label.setAlignment(Qt.AlignmentFlag.AlignCenter)

		self.binary_label = make_label("Binary: -")
		self.trinary_label = make_label("Trinary: -")
		self.explanation_label = make_label("Explanation: wyłączone")

		info_panel = QWidget()
		info_layout = QVBoxLayout(info_panel)
		info_layout.addWidget(self.path_label)
		info_layout.addWidget(self.binary_label)
		info_layout.addWidget(self.trinary_label)
		info_layout.addWidget(self.explanation_label)
		info_layout.addStretch(1)

		self.image_label = QLabel("Wybierz obraz lub katalog obrazów.")
		self.image_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
		self.image_label.setMinimumSize(720, 540)
		self.image_label.setStyleSheet("background: #111827; color: #e5e7eb; border: 1px solid #374151; border-radius: 10px;")

		scroll_area = QScrollArea()
		scroll_area.setWidgetResizable(True)
		scroll_area.setWidget(self.image_label)

		right_panel = QWidget()
		right_layout = QVBoxLayout(right_panel)
		right_layout.addWidget(scroll_area, stretch=1)

		self.display_mode_control = SegmentedControl([
			("Normal", "normal"),
			("Binary", "binary"),
			("Trinary", "trinary"),
		])

		bottom_controls = QWidget()
		bottom_layout = QHBoxLayout(bottom_controls)
		bottom_layout.setContentsMargins(0, 0, 0, 0)
		bottom_layout.addWidget(self.display_mode_control)

		self.prev_button.setText('◀')
		self.next_button.setText('▶')
		self.display_mode_control.modeChanged.connect(self.render_current_image)
		bottom_layout.addStretch(1)
		bottom_layout.addWidget(self.prev_button)
		bottom_layout.addWidget(self.counter_label)
		bottom_layout.addWidget(self.next_button)

		right_layout.addWidget(bottom_controls)
		right_layout.addWidget(info_panel)

		root_layout.addWidget(controls_panel, stretch=0)
		root_layout.addWidget(right_panel, stretch=1)

		self.setCentralWidget(central)

		self.status_bar = QStatusBar()
		self.setStatusBar(self.status_bar)

		open_action = QAction("Otwórz obraz", self)
		open_action.triggered.connect(self.open_image_file)
		open_dir_action = QAction("Otwórz katalog", self)
		open_dir_action.triggered.connect(self.open_image_directory)
		self.menuBar().addAction(open_action)
		self.menuBar().addAction(open_dir_action)

	def _set_navigation_enabled(self, enabled: bool) -> None:
		self.prev_button.setEnabled(enabled)
		self.next_button.setEnabled(enabled)

	def _current_display_mode(self) -> str:
		return self.display_mode_control.currentMode()

	def open_image_file(self) -> None:
		path, _ = QFileDialog.getOpenFileName(
			self,
			"Wybierz obraz",
			"",
			"Images (*.jpg *.jpeg *.png *.bmp *.webp *.tiff)",
		)
		if path:
			self.load_path(path)

	def open_image_directory(self) -> None:
		path = QFileDialog.getExistingDirectory(self, "Wybierz katalog")
		if path:
			self.load_path(path)

	def load_path(self, input_path: str) -> None:
		try:
			self.image_paths = list_image_files(input_path)
		except FileNotFoundError:
			QMessageBox.warning(self, "Błąd", f"Nie można odnaleźć: {input_path}")
			return

		if not self.image_paths:
			QMessageBox.information(self, "Brak obrazów", "Nie znaleziono obsługiwanych plików graficznych.")
			return

		self.current_index = 0
		self.current_result = None
		self.result_cache.clear()
		self._set_navigation_enabled(False)
		self.display_mode_control.setEnabled(False)
		self.status_bar.showMessage("Generowanie wyjaśnień: 0/0")

		first_image = load_image(self.image_paths[0])
		self.image_label.setPixmap(image_to_pixmap(first_image))

		self.precompute_worker = PrecomputeWorker(self.image_paths, self.models)
		self.precompute_worker.precomputed.connect(self.on_precomputed_item)
		self.precompute_worker.progress.connect(self.on_precompute_progress)
		self.precompute_worker.finished_all.connect(self.on_precompute_finished)
		self.precompute_worker.start()

	def show_previous_image(self) -> None:
		if not self.image_paths:
			return
		self.current_index = (self.current_index - 1) % len(self.image_paths)
		self.render_current_image()

	def show_next_image(self) -> None:
		if not self.image_paths:
			return
		self.current_index = (self.current_index + 1) % len(self.image_paths)
		self.render_current_image()

	def render_current_image(self, *_args) -> None:
		if not self.image_paths:
			return

		image_path = self.image_paths[self.current_index]
		display_mode = self._current_display_mode()
		cache_key = (image_path, display_mode)

		cached_result = self.result_cache.get(cache_key)
		if cached_result is not None:
			self.apply_result(cached_result)


	def on_precomputed_item(self, image_path: str, mode: str, result: AnalysisResult) -> None:
		self.result_cache[(image_path, mode)] = result
		if self.image_paths:
			current_path = self.image_paths[self.current_index]
			if image_path == current_path and mode == self._current_display_mode():
				self.apply_result(result)

	def on_precompute_progress(self, count: int, total: int) -> None:
		self.status_bar.showMessage(f"Generowanie wyjaśnień: {count}/{total}")

	def on_precompute_finished(self) -> None:
		self._set_navigation_enabled(len(self.image_paths) > 1)
		self.display_mode_control.setEnabled(True)
		self.status_bar.showMessage("Wczytywanie zakończone")
		self.render_current_image()

	def apply_result(self, result: AnalysisResult) -> None:
		self.current_result = result
		self.path_label.setText(f"Plik: {result.image_path}")

		total = len(self.image_paths)
		current = self.current_index + 1 if self.image_paths else 0
		self.counter_label.setText(f"{current} z {total}")

		self.binary_label.setText(
			"Binary: {label} ({confidence:.2f}) | Artificial {p0:.2f} / Natural {p1:.2f}".format(
				label=result.binary_label,
				confidence=result.binary_confidence,
				p0=float(result.binary_probs[0]),
				p1=float(result.binary_probs[1]),
			)
		)
		self.trinary_label.setText(
			"Trinary: {label} ({confidence:.2f}) | Drawings {p0:.2f} / Games {p1:.2f} / Nature {p2:.2f}".format(
				label=result.trinary_label,
				confidence=result.trinary_confidence,
				p0=float(result.trinary_probs[0]),
				p1=float(result.trinary_probs[1]),
				p2=float(result.trinary_probs[2]),
			)
		)

		if result.display_mode == "normal":
			self.explanation_label.setText("Widok: zwykły obraz")
		else:
			self.explanation_label.setText(
				"Widok: {mode} -> {label}".format(
					mode=DISPLAY_MODES.get(result.display_mode, result.display_mode),
					label=result.explanation_label,
				)
			)
		display_array = result.display_array

		self.image_label.setPixmap(image_to_pixmap(display_array))
		self.status_bar.showMessage(f"Gotowe: {result.image_name}")


def main() -> int:
	app = QApplication(sys.argv)
	app.setStyleSheet(
		"""
		QMainWindow {
			background: #0f172a;
			color: #e5e7eb;
		}
		QWidget {
			font-size: 13px;
		}
		QPushButton, QCheckBox {
			padding: 8px;
		}
		QPushButton {
			background: #1f2937;
			color: #f9fafb;
			border: 1px solid #374151;
			border-radius: 8px;
		}
		QPushButton:hover {
			background: #273449;
		}
		QLabel {
			color: #e5e7eb;
		}
		QStatusBar {
			color: #e5e7eb;
		}
		"""
	)

	try:
		window = MainWindow()
	except Exception as exc:
		QMessageBox.critical(None, "Błąd startu", str(exc))
		return 1

	window.show()
	return app.exec()


if __name__ == "__main__":
	raise SystemExit(main())