# Push-Up Counter (Aperfeiçoado) - Contador de Flexões com Visão por Computação
# Usado com MediaPipe e OpenCV

import cv2
import mediapipe as md
import time
from pushup_logger import PushUpLogger, create_landmarks_dict

# Inicialização do modelo de pose
md_drawing = md.solutions.drawing_utils
md_pose = md.solutions.pose

# Variáveis
count = 0
position = None  # "down" ou "up"
elbow_y = 0      # Armazena a posição Y do cotovelo para comparação
shoulder_y = 0   # Armazena a posição Y do ombro
landmarks_down = None  # Armazena landmarks quando na posição baixa
landmarks_up = None    # Armazena landmarks quando na posição alta

# Inicializar logger
logger = PushUpLogger()
logger.start_session()

# Variáveis para cálculo de FPS
fps = 0
frame_time_start = time.time()

# Inicializa a captura de vídeo
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

# Configuração do modelo de pose (alta precisão)
with md_pose.Pose(
    min_detection_confidence=0.5,
    min_tracking_confidence=0.5
) as pose:

    print("✅ Iniciando o Push-Up Counter... Apareça na frente da câmera e comece as flexões!")
    print("Pressione 'q' para sair. A contagem será exibida no console e na tela.")

    while cap.isOpened():
        # Calcular FPS
        frame_time_end = time.time()
        frame_interval = frame_time_end - frame_time_start
        if frame_interval > 0:
            fps = 1.0 / frame_interval
        frame_time_start = frame_time_end

        success, image = cap.read()
        if not success:
            print("❌ Erro: Nenhum feed de câmera. Verifique a conexão.")
            break

        # Converter para RGB (necessário para MediaPipe)
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image_rgb = cv2.flip(image_rgb, 1)  # Inverter para que a imagem fique direita (como em um espelho)

        # Processa a imagem com MediaPipe
        results = pose.process(image_rgb)

        # Desenhar landmarks e conexões
        if results.pose_landmarks:
            md_drawing.draw_landmarks(image, results.pose_landmarks, md_pose.POSE_CONNECTIONS)

            # Extrair coordenadas dos pontos principais (ombro e cotovelo)
            landmarks = results.pose_landmarks.landmark

            # Índices:
            # - 11: Ombro esquerdo (esquerdo), 12: Ombro direito (direito)
            # - 13: Cotovelo esquerdo, 14: Cotovelo direito
            try:
                # Coordenadas (X, Y) dos pontos
                shoulder_right = landmarks[12].y  # Ombro direito
                elbow_right = landmarks[14].y   # Cotovelo direito
                shoulder_left = landmarks[11].y # Ombro esquerdo
                elbow_left = landmarks[13].y    # Cotovelo esquerdo

                # Média entre os dois (para reduzir variação de lado)
                avg_shoulder_y = (shoulder_left + shoulder_right) / 2
                avg_elbow_y = (elbow_left + elbow_right) / 2

                # Criar dicionário de landmarks para logging
                landmarks_dict = create_landmarks_dict(landmarks)

                # Lógica de contagem: "down" = ombro abaixo do cotovelo
                # "up" = ombro acima do cotovelo (e só conta se estava no "down")
                if avg_shoulder_y > avg_elbow_y:
                    # Em "down" (ombros abaixo do cotovelo)
                    if position != "down":
                        position = "down"
                        landmarks_down = landmarks_dict  # Salvar landmarks na posição baixa
                elif avg_shoulder_y < avg_elbow_y:
                    # Em "up" (ombros acima do cotovelo)
                    if position == "down":
                        landmarks_up = landmarks_dict  # Salvar landmarks na posição alta

                        count += 1
                        print(f"🔥 Flexão #{count}! (Nível {count // 50})")

                        # Logar flexão com landmarks de ambas as posições
                        logger.log_pushup(count, landmarks_down, landmarks_up)

                        # Gamificação
                        if count % 50 == 0:
                            print(f"🎉 Nível alcançado! {count} push-ups!")
                        elif count % 10 == 0:
                            print(f"💡 Boa continuidade! {count} push-ups!")

                        # Resetar estado
                        position = None
                        landmarks_down = None
                        landmarks_up = None
                else:
                    # Se estiver em equilíbrio, não contar
                    pass

                # Logar frame (a cada frame)
                logger.log_frame(landmarks_dict, position, fps)

            except IndexError:
                print("⚠️ Erro: Pontos de detecção incompletos.")
                position = None
                continue

        # Exibir no console e tela
        image_display = cv2.flip(image, 1)
        cv2.putText(image_display, f"Push-ups: {count}", (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.putText(image_display, f"Status: {'Baixo' if position == 'down' else 'Cima'}",
                    (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 1)
        cv2.putText(image_display, f"FPS: {int(fps)}", (10, 90),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 165, 0), 1)
        cv2.putText(image_display, "[Q] Sair", (10, 450),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 200, 200), 1)

        cv2.imshow("Push-Up Counter", image_display)

        # Tecla de saída
        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):
            print(f"\n🎉 Parabéns! Você fez {count} push-ups. 🏁")
            print("\n💾 Salvando sessão...")
            logger.end_session()
            break

# Encerramento
cap.release()
cv2.destroyAllWindows()

print("\n✨ Sessão encerrada! Abra o dashboard para ver suas estatísticas:")
print(f"   → python3 -m http.server 8000")
print(f"   → Acesse: http://localhost:8000/pushup_dashboard.html")