import tkinter as tk
from tkinter import ttk, filedialog, messagebox, simpledialog
import cv2
import numpy as np
import os
import pickle
from PIL import Image, ImageTk
import threading
import time

class ImprovedFaceRecognitionApp:
    def __init__(self, root):
        self.root = root
        self.root.title("Улучшенная система распознавания лиц")
        self.root.geometry("1200x700")
        
        # Инициализация переменных
        self.face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
        self.eye_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_eye.xml')
        self.known_faces = []
        self.known_names = []
        self.face_samples_dir = "face_samples"
        self.is_recognition_active = False
        self.cap = None
        self.current_frame = None
        
        # Настройки точности
        self.recognition_threshold = 0.7  # Порог распознавания
        self.min_face_size = 80  # Минимальный размер лица
        self.scale_factor = 1.05  # Меньше = точнее, но медленнее
        self.min_neighbors = 6  # Больше = меньше ложных срабатываний
        
        # Создание интерфейса
        self.create_widgets()
        
        # Загрузка образцов лиц
        self.load_face_samples()
        
    def create_widgets(self):
        # Основной фрейм
        main_frame = ttk.Frame(self.root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # Левая панель - управление
        left_frame = ttk.LabelFrame(main_frame, text="Управление и настройки", padding=10)
        left_frame.pack(side=tk.LEFT, fill=tk.Y, padx=(0, 10))
        
        # Правая панель - видео
        right_frame = ttk.LabelFrame(main_frame, text="Видео с камеры", padding=10)
        right_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True)
        
        # Элементы управления
        ttk.Button(left_frame, text="Запуск распознавания", 
                  command=self.start_recognition).pack(fill=tk.X, pady=5)
        ttk.Button(left_frame, text="Остановка", 
                  command=self.stop_recognition).pack(fill=tk.X, pady=5)
        ttk.Button(left_frame, text="Добавить образец", 
                  command=self.add_face_sample).pack(fill=tk.X, pady=5)
        ttk.Button(left_frame, text="Очистить базу", 
                  command=self.clear_database).pack(fill=tk.X, pady=5)
        
        # Настройки точности
        settings_frame = ttk.LabelFrame(left_frame, text="Настройки точности", padding=5)
        settings_frame.pack(fill=tk.X, pady=10)
        
        # Порог распознавания
        ttk.Label(settings_frame, text="Порог распознавания:").pack(anchor=tk.W)
        self.threshold_var = tk.DoubleVar(value=self.recognition_threshold)
        threshold_scale = ttk.Scale(settings_frame, from_=0.5, to=0.9, 
                                   variable=self.threshold_var, orient=tk.HORIZONTAL)
        threshold_scale.pack(fill=tk.X)
        self.threshold_label = ttk.Label(settings_frame, text=f"Текущее: {self.recognition_threshold}")
        self.threshold_label.pack(anchor=tk.W)
        
        # Минимальный размер лица
        ttk.Label(settings_frame, text="Мин. размер лица:").pack(anchor=tk.W)
        self.face_size_var = tk.IntVar(value=self.min_face_size)
        face_size_scale = ttk.Scale(settings_frame, from_=30, to=150, 
                                   variable=self.face_size_var, orient=tk.HORIZONTAL)
        face_size_scale.pack(fill=tk.X)
        self.face_size_label = ttk.Label(settings_frame, text=f"Текущее: {self.min_face_size}")
        self.face_size_label.pack(anchor=tk.W)
        
        # Применить настройки
        ttk.Button(settings_frame, text="Применить настройки", 
                  command=self.apply_settings).pack(fill=tk.X, pady=5)
        
        # Список известных лиц
        ttk.Label(left_frame, text="Известные лица:").pack(anchor=tk.W)
        self.faces_listbox = tk.Listbox(left_frame, height=10)
        self.faces_listbox.pack(fill=tk.BOTH, expand=True, pady=5)
        
        # Статус
        self.status_label = ttk.Label(left_frame, text="Статус: Остановлено")
        self.status_label.pack(fill=tk.X, pady=5)
        
        # Видео панель
        self.video_label = ttk.Label(right_frame, background='black')
        self.video_label.pack(fill=tk.BOTH, expand=True)
        
        # Информация о распознавании
        info_frame = ttk.Frame(right_frame)
        info_frame.pack(fill=tk.X, pady=5)
        
        self.info_text = tk.Text(info_frame, height=8, width=60)
        scrollbar = ttk.Scrollbar(info_frame, command=self.info_text.yview)
        self.info_text.configure(yscrollcommand=scrollbar.set)
        self.info_text.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        
    def apply_settings(self):
        """Применение новых настроек"""
        self.recognition_threshold = self.threshold_var.get()
        self.min_face_size = self.face_size_var.get()
        
        self.threshold_label.config(text=f"Текущее: {self.recognition_threshold:.2f}")
        self.face_size_label.config(text=f"Текущее: {self.min_face_size}")
        
        self.log_info(f"Настройки применены: порог={self.recognition_threshold:.2f}, мин.размер={self.min_face_size}")
    
    def load_face_samples(self):
        """Загрузка образцов лиц из папки"""
        if not os.path.exists(self.face_samples_dir):
            os.makedirs(self.face_samples_dir)
            return
        
        self.known_faces = []
        self.known_names = []
        self.faces_listbox.delete(0, tk.END)
        
        for filename in os.listdir(self.face_samples_dir):
            if filename.endswith(('.jpg', '.jpeg', '.png')):
                name = os.path.splitext(filename)[0]
                img_path = os.path.join(self.face_samples_dir, filename)
                
                img = cv2.imread(img_path)
                if img is not None:
                    # Предобработка изображения
                    processed_img = self.preprocess_image(img)
                    
                    if processed_img is not None:
                        # Вычисляем гистограмму для каждого канала цвета
                        histograms = self.calculate_multi_channel_histogram(processed_img)
                        
                        self.known_faces.append({
                            'name': name,
                            'histograms': histograms,
                            'image': processed_img
                        })
                        self.known_names.append(name)
                        self.faces_listbox.insert(tk.END, name)
        
        self.log_info(f"Загружено {len(self.known_faces)} образцов лиц")
    
    def preprocess_image(self, img):
        """Предобработка изображения для улучшения распознавания"""
        try:
            # Конвертация в grayscale
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            
            # Обнаруживаем лицо на образце
            faces = self.face_cascade.detectMultiScale(gray, 1.1, 5)
            if len(faces) == 0:
                return None
            
            x, y, w, h = faces[0]
            face_roi = gray[y:y+h, x:x+w]
            
            # Увеличиваем размер лица для лучшего распознавания
            face_roi = cv2.resize(face_roi, (200, 200))
            
            # Выравнивание гистограммы для улучшения контраста
            face_roi = cv2.equalizeHist(face_roi)
            
            # Размытие для уменьшения шума
            face_roi = cv2.GaussianBlur(face_roi, (5, 5), 0)
            
            return face_roi
        except Exception as e:
            self.log_info(f"Ошибка предобработки: {e}")
            return None
    
    def calculate_multi_channel_histogram(self, img):
        """Вычисление гистограмм для нескольких каналов"""
        # Для grayscale изображения
        hist_gray = cv2.calcHist([img], [0], None, [256], [0, 256])
        cv2.normalize(hist_gray, hist_gray, 0, 255, cv2.NORM_MINMAX)
        
        # Дополнительные гистограммы для улучшения точности
        # Гистограмма после применения оператора Лапласа (края)
        laplacian = cv2.Laplacian(img, cv2.CV_64F)
        hist_laplacian = cv2.calcHist([np.uint8(np.absolute(laplacian))], [0], None, [256], [0, 256])
        cv2.normalize(hist_laplacian, hist_laplacian, 0, 255, cv2.NORM_MINMAX)
        
        return {
            'gray': hist_gray,
            'laplacian': hist_laplacian
        }
    
    def compare_faces_advanced(self, face_img):
        """Улучшенное сравнение лиц с использованием нескольких методов"""
        if len(self.known_faces) == 0:
            return "Неизвестный", 0
        
        # Предобработка обнаруженного лица
        face_img = cv2.resize(face_img, (200, 200))
        face_img = cv2.equalizeHist(face_img)
        face_img = cv2.GaussianBlur(face_img, (5, 5), 0)
        
        # Вычисляем гистограммы для обнаруженного лица
        test_histograms = self.calculate_multi_channel_histogram(face_img)
        
        best_match = None
        best_score = 0
        
        # Сравниваем с каждым образцом используя несколько методов
        for sample in self.known_faces:
            total_score = 0
            method_count = 0
            
            # Сравнение гистограмм grayscale
            score_gray = cv2.compareHist(test_histograms['gray'], sample['histograms']['gray'], cv2.HISTCMP_CORREL)
            total_score += score_gray
            method_count += 1
            
            # Сравнение гистограмм Лапласиана
            score_laplacian = cv2.compareHist(test_histograms['laplacian'], sample['histograms']['laplacian'], cv2.HISTCMP_CORREL)
            total_score += score_laplacian
            method_count += 1
            
            # Дополнительный метод: сравнение с помощью шаблонов
            try:
                # Используем метод сопоставления шаблонов
                result = cv2.matchTemplate(face_img, sample['image'], cv2.TM_CCOEFF_NORMED)
                _, max_val, _, _ = cv2.minMaxLoc(result)
                total_score += max_val
                method_count += 1
            except:
                pass
            
            # Усредняем оценку
            avg_score = total_score / method_count
            
            if avg_score > best_score:
                best_score = avg_score
                best_match = sample['name']
        
        # Проверяем порог
        if best_score > self.recognition_threshold:
            return best_match, best_score
        else:
            return "Неизвестный", best_score
    
    def start_recognition(self):
        """Запуск распознавания"""
        if self.is_recognition_active:
            return
        
        # Пробуем разные способы инициализации камеры
        self.cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
        if not self.cap.isOpened():
            self.cap = cv2.VideoCapture(0)
            
        if not self.cap.isOpened():
            messagebox.showerror("Ошибка", "Не удалось открыть камеру")
            return
        
        # Устанавливаем параметры камеры
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        self.cap.set(cv2.CAP_PROP_FPS, 30)
        
        self.is_recognition_active = True
        self.status_label.config(text="Статус: Распознавание активно")
        self.log_info("Запущено улучшенное распознавание лиц")
        
        # Запуск потока для обработки видео
        self.recognition_thread = threading.Thread(target=self.process_video)
        self.recognition_thread.daemon = True
        self.recognition_thread.start()
        
        # Запуск обновления интерфейса
        self.update_video_ui()
    
    def update_video_ui(self):
        """Обновление интерфейса видео"""
        if self.is_recognition_active and self.current_frame is not None:
            # Конвертация для Tkinter
            rgb_image = cv2.cvtColor(self.current_frame, cv2.COLOR_BGR2RGB)
            pil_image = Image.fromarray(rgb_image)
            
            # Получаем размеры метки видео
            label_width = self.video_label.winfo_width()
            label_height = self.video_label.winfo_height()
            
            # Если размеры метки не определены, используем стандартные
            if label_width <= 1 or label_height <= 1:
                label_width = 640
                label_height = 480
            
            # Масштабируем изображение под размер метки
            pil_image = pil_image.resize((label_width, label_height), Image.Resampling.LANCZOS)
            tk_image = ImageTk.PhotoImage(pil_image)
            
            # Обновляем метку
            self.video_label.configure(image=tk_image)
            self.video_label.image = tk_image
        
        # Планируем следующее обновление
        if self.is_recognition_active:
            self.root.after(30, self.update_video_ui)
    
    def stop_recognition(self):
        """Остановка распознавания"""
        self.is_recognition_active = False
        if self.cap:
            self.cap.release()
            self.cap = None
        self.status_label.config(text="Статус: Остановлено")
        self.log_info("Распознавание остановлено")
    
    def process_video(self):
        """Обработка видео в отдельном потоке"""
        while self.is_recognition_active and self.cap and self.cap.isOpened():
            ret, frame = self.cap.read()
            if not ret:
                self.log_info("Ошибка чтения кадра с камеры")
                break
            
            # Зеркальное отражение
            frame = cv2.flip(frame, 1)
            display_frame = frame.copy()
            
            # Конвертация в grayscale
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            
            # Улучшенное обнаружение лиц с настройками
            faces = self.face_cascade.detectMultiScale(
                gray,
                scaleFactor=self.scale_factor,
                minNeighbors=self.min_neighbors,
                minSize=(self.min_face_size, self.min_face_size)
            )
            
            # Обработка каждого лица
            recognized_faces = []
            for (x, y, w, h) in faces:
                face_roi = gray[y:y+h, x:x+w]
                
                # Улучшенное распознавание лица
                name, score = self.compare_faces_advanced(face_roi)
                
                # Определение цвета рамки
                if name != "Неизвестный":
                    color = (0, 255, 0)  # Зеленый для известных
                    confidence = f"{score:.3f}"
                else:
                    color = (0, 0, 255)  # Красный для неизвестных
                    confidence = f"{score:.3f}"
                
                # Рисование рамки и текста
                cv2.rectangle(display_frame, (x, y), (x+w, y+h), color, 2)
                cv2.rectangle(display_frame, (x, y-35), (x+w, y), color, cv2.FILLED)
                
                # Текст с именем и уверенностью
                text = f"{name} ({confidence})"
                cv2.putText(display_frame, text, (x+5, y-10), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)
                
                recognized_faces.append((name, score))
            
            # Обновление информации
            if recognized_faces:
                info_text = "Обнаружены лица:\n"
                for name, score in recognized_faces:
                    status = "✓ Известный" if name != "Неизвестный" else "✗ Неизвестный"
                    info_text += f"- {name}: {status} (сходство: {score:.3f})\n"
                self.update_info_text(info_text)
            else:
                self.update_info_text("Лица не обнаружены")
            
            # Сохраняем кадр для отображения
            self.current_frame = display_frame
            
            time.sleep(0.03)
    
    def update_info_text(self, text):
        """Обновление информационного текста"""
        self.root.after(0, lambda: self._update_info_text(text))
    
    def _update_info_text(self, text):
        """Обновление информационного текста в основном потоке"""
        self.info_text.delete(1.0, tk.END)
        self.info_text.insert(1.0, text)
    
    def log_info(self, message):
        """Добавление сообщения в лог"""
        timestamp = time.strftime("%H:%M:%S")
        self.root.after(0, lambda: self._log_info(f"[{timestamp}] {message}\n"))
    
    def _log_info(self, message):
        """Добавление сообщения в лог в основном потоке"""
        self.info_text.insert(tk.END, message)
        self.info_text.see(tk.END)
    
    def add_face_sample(self):
        """Добавление нового образца лица"""
        if not self.is_recognition_active:
            messagebox.showwarning("Предупреждение", "Сначала запустите распознавание")
            return
        
        # Запрос имени для образца
        name = simpledialog.askstring("Имя", "Введите имя для этого лица:")
        if not name:
            return
            
        # Сохранение текущего кадра
        if self.current_frame is not None:
            gray = cv2.cvtColor(self.current_frame, cv2.COLOR_BGR2GRAY)
            
            # Используем улучшенные параметры обнаружения
            faces = self.face_cascade.detectMultiScale(
                gray,
                scaleFactor=self.scale_factor,
                minNeighbors=self.min_neighbors,
                minSize=(self.min_face_size, self.min_face_size)
            )
            
            if len(faces) > 0:
                x, y, w, h = faces[0]
                face_roi = self.current_frame[y:y+h, x:x+w]
                
                # Создаем имя файла
                filename = f"{name}.jpg"
                filepath = os.path.join(self.face_samples_dir, filename)
                
                # Сохранение образца
                cv2.imwrite(filepath, face_roi)
                
                # Перезагрузка образцов
                self.load_face_samples()
                self.log_info(f"Добавлен новый образец: {name}")
            else:
                messagebox.showerror("Ошибка", "Не удалось обнаружить лицо на кадре")
        else:
            messagebox.showerror("Ошибка", "Нет доступного кадра с камеры")
    
    def clear_database(self):
        """Очистка базы образцов"""
        if messagebox.askyesno("Подтверждение", "Вы уверены, что хотите очистить базу образцов?"):
            for filename in os.listdir(self.face_samples_dir):
                file_path = os.path.join(self.face_samples_dir, filename)
                if os.path.isfile(file_path):
                    os.remove(file_path)
            
            self.known_faces = []
            self.known_names = []
            self.faces_listbox.delete(0, tk.END)
            self.log_info("База образцов очищена")
    
    def on_closing(self):
        """Обработка закрытия приложения"""
        self.stop_recognition()
        self.root.destroy()

def main():
    root = tk.Tk()
    app = ImprovedFaceRecognitionApp(root)
    root.protocol("WM_DELETE_WINDOW", app.on_closing)
    root.mainloop()

if __name__ == "__main__":
    main()