import cv2
import os
import csv
import sys
from pathlib import Path
import argparse

sys.path.append(str(Path(__file__).parent.parent))
from skeleton.mediapipe_detector import MediaPipeDetector
from ml_gestures.feature_extractor import normalize_landmarks

def main():
    parser = argparse.ArgumentParser(description='Разметка изображений для обучения')
    parser.add_argument('--folder', required=True, help='Папка с изображениями')
    parser.add_argument('--classes', default='dome,cross,none',
                        help='Список классов через запятую')
    parser.add_argument('--output', default='gesture_data.csv',
                        help='Имя выходного CSV-файла')
    parser.add_argument('--max_display_size', default='800,600',
                        help='Максимальный размер для отображения (ширина,высота)')
    args = parser.parse_args()

    classes = [c.strip() for c in args.classes.split(',')]
    key_to_class = {str(i+1): cls for i, cls in enumerate(classes)}
    print("Классы:", classes)

    detector = MediaPipeDetector()

    max_width, max_height = map(int, args.max_display_size.split(','))

    image_extensions = ('.jpg', '.jpeg', '.png', '.bmp')
    image_files = [f for f in os.listdir(args.folder) if f.lower().endswith(image_extensions)]
    image_files.sort()
    print(f"Найдено {len(image_files)} изображений.")

    csv_file = args.output
    file_exists = os.path.isfile(csv_file)
    if not file_exists:
        with open(csv_file, 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow(['class'] + [f'f{i}' for i in range(99)])

    for idx, filename in enumerate(image_files):
        filepath = os.path.join(args.folder, filename)
        print(f"\n[{idx+1}/{len(image_files)}] {filename}")

        image = cv2.imread(filepath)
        if image is None:
            print("Не удалось загрузить")
            continue

        result = detector.detect(image)

        if result['success']:
            landmarks = result['landmarks']
            features = normalize_landmarks(landmarks)
            vis_image = detector.draw_landmarks(image.copy(), result['pose_landmarks'])
        else:
            vis_image = image.copy()
            landmarks = None

        # Масштабирование для отображения
        h, w = vis_image.shape[:2]
        scale = min(max_width / w, max_height / h, 1.0)
        if scale < 1.0:
            new_w = int(w * scale)
            new_h = int(h * scale)
            display = cv2.resize(vis_image, (new_w, new_h))
        else:
            display = vis_image.copy()

        # Панель с инструкцией
        dh, dw = display.shape[:2]
        overlay = display.copy()
        cv2.rectangle(overlay, (0, dh-80), (dw, dh), (50,50,50), -1)
        cv2.addWeighted(overlay, 0.6, display, 0.4, 0, display)

        y = dh - 60
        for i, cls in enumerate(classes):
            cv2.putText(display, f"{i+1}:{cls}", (10 + i*120, y),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,255), 2)
        cv2.putText(display, "n:skip  q:quit", (10, y+30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,0), 2)

        cv2.imshow('Annotation', display)
        key = cv2.waitKey(0) & 0xFF
        cv2.destroyAllWindows()

        if key == ord('q'):
            break
        elif key == ord('n'):
            continue
        else:
            key_char = chr(key) if key < 256 else None
            if key_char in key_to_class:
                selected = key_to_class[key_char]
                with open(csv_file, 'a', newline='', encoding='utf-8') as f:
                    writer = csv.writer(f)
                    writer.writerow([selected] + features.tolist())
                print(f"Сохранено: {selected}")
            else:
                print("Неверная клавиша")

    print("Готово.")

if __name__ == '__main__':
    main()