"""
ASL Recognition API
Hugging Face Spaces — FastAPI

Endpoints:
  POST /predict/alphabet  — single JPEG frame → letter + confidence
  POST /predict/auto      — 60 frames → seamless CNN or LSTM routing
  GET  /health            — health check

Decision logic in /predict/auto:
  1. Hand presence check  — need hand in >= 1/3 of frames, else → not detected
  2. Two hands check      — >= 1/3 frames have 2 hands → LSTM (phrase)
  3. Movement check       — normalized wrist displacement >= threshold → LSTM (phrase)
  4. CNN vote             — vote_ratio >= 0.50 → CNN (letter)
  5. Fallback             — LSTM (phrase)
"""

import os
import time
import urllib.request

import cv2
import numpy as np
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image

import mediapipe as mp
from mediapipe.tasks import python as mp_python
from mediapipe.tasks.python import vision as mp_vision

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List

# ============================================================
# Settings — must match training
# ============================================================
IMG_SIZE   = 64
SEQ_LEN    = 60
POSE_SIZE  = 33 * 3        # 99
HAND_SIZE  = 21 * 3        # 63
INPUT_SIZE = POSE_SIZE + HAND_SIZE + HAND_SIZE  # 225

WRIST          = 0
MIDDLE_MCP     = 9
LEFT_SHOULDER  = 11
RIGHT_SHOULDER = 12

HAND_CONNECTIONS = [
    (0,1),(1,2),(2,3),(3,4),
    (0,5),(5,6),(6,7),(7,8),
    (5,9),(9,10),(10,11),(11,12),
    (9,13),(13,14),(14,15),(15,16),
    (13,17),(0,17),(17,18),(18,19),(19,20)
]
FINGERTIP_IDS = {4, 8, 12, 16, 20}

# ── Routing thresholds ───────────────────────────────────────
MIN_HAND_PRESENCE_RATIO = 1 / 3   # fraction of frames that must have a hand
TWO_HAND_RATIO          = 1 / 3   # fraction of frames with 2 hands → LSTM
MOVEMENT_THRESHOLD      = 0.10    # normalized wrist displacement mid→end → LSTM
MIN_CNN_VOTE_RATIO      = 0.50    # CNN vote agreement to accept a letter
MIN_CNN_CONFIDENCE      = 0.60    # per-frame CNN confidence to count as a vote
MIN_COORD_FRAMES        = 10      # minimum coord frames for LSTM to run

# ============================================================
# Model Architectures
# ============================================================

class ConvBlock(nn.Module):
    def __init__(self, in_ch, out_ch, pool=True):
        super().__init__()
        layers = [
            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        ]
        if pool:
            layers.append(nn.MaxPool2d(2, 2))
        self.block = nn.Sequential(*layers)

    def forward(self, x):
        return self.block(x)


class ASL_CNN(nn.Module):
    """
    Alphabet model — plain CNN, matches alphabet_model.pth exactly.
    Input: (B, C, H, W)
    """
    def __init__(self, n_classes):
        super().__init__()
        self.cnn = nn.Sequential(
            ConvBlock(3,   32,  pool=True),
            ConvBlock(32,  64,  pool=True),
            ConvBlock(64,  128, pool=True),
            ConvBlock(128, 256, pool=True),
            ConvBlock(256, 512, pool=False),
            nn.AdaptiveAvgPool2d(1),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(512, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.4),
            nn.Linear(256, n_classes),
        )

    def forward(self, x):
        return self.classifier(self.cnn(x))


class ASL_Phrases_LSTM(nn.Module):
    """
    Phrases model — pose + hand coordinates fed into LSTM.
    Input: (B, SEQ_LEN, INPUT_SIZE) = (B, 60, 225)
    """
    def __init__(self, input_size, n_classes, lstm_hidden=512, lstm_layers=2):
        super().__init__()
        self.input_norm = nn.LayerNorm(input_size)
        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=lstm_hidden,
            num_layers=lstm_layers,
            batch_first=True,
            dropout=0.3 if lstm_layers > 1 else 0,
        )
        self.classifier = nn.Sequential(
            nn.Linear(lstm_hidden, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.4),
            nn.Linear(256, n_classes),
        )

    def forward(self, x):
        x = self.input_norm(x)
        _, (h_n, _) = self.lstm(x)
        return self.classifier(h_n[-1])

# ============================================================
# Load models and classes
# ============================================================

device = torch.device('cpu')

def load_classes(path):
    with open(path, 'r') as f:
        return [l.strip() for l in f.readlines()]

def load_alphabet_model(path, classes):
    ckpt  = torch.load(path, map_location=device, weights_only=False)
    model = ASL_CNN(len(classes)).to(device)
    model.load_state_dict(ckpt['model_state_dict'])
    model.eval()
    return model

def load_phrases_model(path, classes):
    ckpt      = torch.load(path, map_location=device, weights_only=False)
    hidden    = ckpt.get('lstm_hidden', 512)
    layers    = ckpt.get('lstm_layers', 2)
    i_size    = ckpt.get('input_size',  INPUT_SIZE)
    n_classes = ckpt.get('n_classes',   len(classes))
    model     = ASL_Phrases_LSTM(i_size, n_classes, hidden, layers).to(device)
    model.load_state_dict(ckpt['model_state_dict'])
    model.eval()
    return model

ALPHA_CLASSES  = load_classes('alphabet_classes.txt')
PHRASE_CLASSES = [
    "GOOD AFTERNOON",
    "MAGANDANG HAPON",
    "GOOD EVENING",
    "GOOD MORNING",
    "MAGANDANG UMAGA",
    "HELLO",
    "HOW ARE YOU?",
    "KUMUSTA KA?",
    "I'M FINE",
    "J",
    "NO",
    "THANK YOU",
    "YES",
    "YOU'RE WELCOME",
    "WALANG ANUMAN",
    "Z"
]

alphabet_model = load_alphabet_model('alphabet_model.pth', ALPHA_CLASSES)
phrases_model  = load_phrases_model('phrases_model.pth',   PHRASE_CLASSES)

print(f"Alphabet classes: {ALPHA_CLASSES}")
print(f"Phrase classes:   {PHRASE_CLASSES}")

# ============================================================
# Download MediaPipe task files if missing
# ============================================================

def download_if_missing(path, url, name):
    if not os.path.exists(path):
        print(f'Downloading {name}...')
        urllib.request.urlretrieve(url, path)
        print(f'{name} ready.')

download_if_missing(
    'hand_landmarker.task',
    'https://storage.googleapis.com/mediapipe-models/hand_landmarker/'
    'hand_landmarker/float16/1/hand_landmarker.task',
    'hand_landmarker.task'
)
download_if_missing(
    'pose_landmarker.task',
    'https://storage.googleapis.com/mediapipe-models/pose_landmarker/'
    'pose_landmarker_lite/float16/1/pose_landmarker_lite.task',
    'pose_landmarker.task'
)

# VIDEO mode — wall clock timestamps (time.time() * 1000) are used per frame,
# so they are always strictly increasing across requests on the same server instance.
hand_options_alphabet = mp_vision.HandLandmarkerOptions(
    base_options=mp_python.BaseOptions(model_asset_path='hand_landmarker.task'),
    running_mode=mp_vision.RunningMode.VIDEO,
    num_hands=1,
    min_hand_detection_confidence=0.5,
    min_hand_presence_confidence=0.5,
    min_tracking_confidence=0.5,
)
hand_options_phrase = mp_vision.HandLandmarkerOptions(
    base_options=mp_python.BaseOptions(model_asset_path='hand_landmarker.task'),
    running_mode=mp_vision.RunningMode.VIDEO,
    num_hands=2,
    min_hand_detection_confidence=0.5,
    min_hand_presence_confidence=0.5,
    min_tracking_confidence=0.5,
)
pose_options = mp_vision.PoseLandmarkerOptions(
    base_options=mp_python.BaseOptions(model_asset_path='pose_landmarker.task'),
    running_mode=mp_vision.RunningMode.VIDEO,
    num_poses=1,
    min_pose_detection_confidence=0.3,
    min_pose_presence_confidence=0.3,
    min_tracking_confidence=0.3,
)

hand_detector_alpha  = mp_vision.HandLandmarker.create_from_options(hand_options_alphabet)
hand_detector_phrase = mp_vision.HandLandmarker.create_from_options(hand_options_phrase)
pose_detector        = mp_vision.PoseLandmarker.create_from_options(pose_options)

# ============================================================
# Transforms
# ============================================================

infer_transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
])

# ============================================================
# Helpers
# ============================================================

def bytes_to_mp_image(data: bytes):
    arr   = np.frombuffer(data, np.uint8)
    frame = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    frame = cv2.flip(frame, 1)
    frame = cv2.convertScaleAbs(frame, alpha=1.3, beta=20)  # matches training pipeline
    rgb   = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    return mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb)


def render_hand_skeleton(landmarks):
    canvas = np.zeros((IMG_SIZE, IMG_SIZE, 3), dtype=np.uint8)
    xs = [lm.x for lm in landmarks]
    ys = [lm.y for lm in landmarks]
    pad   = 0.15
    min_x = max(0.0, min(xs) - pad); min_y = max(0.0, min(ys) - pad)
    max_x = min(1.0, max(xs) + pad); max_y = min(1.0, max(ys) + pad)
    rx = max_x - min_x if max_x > min_x else 1
    ry = max_y - min_y if max_y > min_y else 1

    def to_px(lm):
        cx = int(((lm.x - min_x) / rx) * (IMG_SIZE - 1))
        cy = int(((lm.y - min_y) / ry) * (IMG_SIZE - 1))
        return (max(0, min(IMG_SIZE-1, cx)), max(0, min(IMG_SIZE-1, cy)))

    for p1, p2 in HAND_CONNECTIONS:
        cv2.line(canvas, to_px(landmarks[p1]), to_px(landmarks[p2]), (200,200,200), 2)
    for i, lm in enumerate(landmarks):
        r     = 4 if i in FINGERTIP_IDS else 3
        color = (255,255,255) if i in FINGERTIP_IDS else (180,180,180)
        cv2.circle(canvas, to_px(lm), r, color, -1)

    return Image.fromarray(cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB))


def normalize_pose(lms):
    mid_x = (lms[LEFT_SHOULDER].x + lms[RIGHT_SHOULDER].x) / 2.0
    mid_y = (lms[LEFT_SHOULDER].y + lms[RIGHT_SHOULDER].y) / 2.0
    scale = abs(lms[LEFT_SHOULDER].x - lms[RIGHT_SHOULDER].x)
    if scale < 1e-6: scale = 1e-6
    vec = []
    for lm in lms:
        vec.extend([(lm.x - mid_x)/scale, (lm.y - mid_y)/scale, lm.z/scale])
    return np.array(vec, dtype=np.float32)


def normalize_hand(lms):
    origin = lms[WRIST]
    ref    = lms[MIDDLE_MCP]
    scale  = np.sqrt((ref.x-origin.x)**2 + (ref.y-origin.y)**2)
    if scale < 1e-6: scale = 1e-6
    vec = []
    for lm in lms:
        vec.extend([(lm.x-origin.x)/scale, (lm.y-origin.y)/scale, lm.z/scale])
    return np.array(vec, dtype=np.float32)


def get_hands_by_side(hand_result):
    left_lms = right_lms = None
    if not hand_result.hand_landmarks:
        return left_lms, right_lms
    for i, handedness in enumerate(hand_result.handedness):
        label = handedness[0].category_name
        if label == 'Left':
            left_lms = hand_result.hand_landmarks[i]
        else:
            right_lms = hand_result.hand_landmarks[i]
    return left_lms, right_lms


def run_lstm(coord_buffer, vote_ratio):
    """Resample coord_buffer to SEQ_LEN and run LSTM."""
    if len(coord_buffer) < MIN_COORD_FRAMES:
        return AutoResponse(
            result="", result_type="phrase",
            confidence=0.0, vote_ratio=vote_ratio, detected=False
        )
    buf    = np.array(coord_buffer)
    idx    = np.linspace(0, len(buf)-1, SEQ_LEN, dtype=int)
    seq    = buf[idx]
    tensor = torch.tensor(seq, dtype=torch.float32).unsqueeze(0).to(device)
    with torch.no_grad():
        probs = torch.softmax(phrases_model(tensor), dim=1)[0]
    conf, idx2 = probs.max(dim=0)
    return AutoResponse(
        result=PHRASE_CLASSES[idx2.item()],
        result_type="phrase",
        confidence=conf.item(),
        vote_ratio=vote_ratio,
        detected=True
    )

# ============================================================
# FastAPI App
# ============================================================

app = FastAPI(title="ASL Recognition API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Response models ──────────────────────────────────────────

class AlphabetResponse(BaseModel):
    letter:     str
    confidence: float
    detected:   bool

class AutoResponse(BaseModel):
    result:      str
    result_type: str
    confidence:  float
    vote_ratio:  float
    detected:    bool

# ── Endpoints ────────────────────────────────────────────────

@app.get("/health")
def health():
    return {"status": "ok"}


@app.post("/predict/alphabet", response_model=AlphabetResponse)
async def predict_alphabet(file: UploadFile = File(...)):
    data     = await file.read()
    mp_img   = bytes_to_mp_image(data)
    result   = hand_detector_alpha.detect_for_video(mp_img, int(time.time() * 1000))
    hand_lms = result.hand_landmarks[0] if result.hand_landmarks else None

    if hand_lms is None:
        return AlphabetResponse(letter="", confidence=0.0, detected=False)

    skel   = render_hand_skeleton(hand_lms)
    tensor = infer_transform(skel).unsqueeze(0).to(device)
    with torch.no_grad():
        probs = torch.softmax(alphabet_model(tensor), dim=1)[0]
    conf, idx = probs.max(dim=0)
    return AlphabetResponse(
        letter=ALPHA_CLASSES[idx.item()],
        confidence=conf.item(),
        detected=True
    )


@app.post("/predict/auto", response_model=AutoResponse)
async def predict_auto(files: List[UploadFile] = File(...)):
    """
    Accepts 60 JPEG frames from Unity.
    Runs MediaPipe on every frame (VIDEO mode — wall clock timestamps).
    Routes to CNN (letter) or LSTM (phrase) based on what's detected.
    """
    if not files:
        raise HTTPException(status_code=400, detail="No frames provided")

    frame_data = [await f.read() for f in files]

    # ── Per-frame accumulators ───────────────────────────────
    cnn_votes          = {}
    total_hand_frames  = 0
    two_hand_frames    = 0
    wrist_positions    = []

    coord_buffer        = []
    last_pose_vec       = np.zeros(POSE_SIZE, dtype=np.float32)
    last_left_hand_vec  = np.zeros(HAND_SIZE, dtype=np.float32)
    last_right_hand_vec = np.zeros(HAND_SIZE, dtype=np.float32)

    total_frames = len(frame_data)

    for data in frame_data:
        mp_img       = bytes_to_mp_image(data)
        timestamp_ms = int(time.time() * 1000)  # wall clock — always increasing across requests

        # ── CNN: single hand for alphabet voting ─────────────
        cnn_result = hand_detector_alpha.detect_for_video(mp_img, timestamp_ms)
        hand_lms   = cnn_result.hand_landmarks[0] if cnn_result.hand_landmarks else None

        if hand_lms is not None:
            total_hand_frames += 1
            wrist_positions.append((hand_lms[WRIST].x, hand_lms[WRIST].y))

            skel   = render_hand_skeleton(hand_lms)
            tensor = infer_transform(skel).unsqueeze(0).to(device)
            with torch.no_grad():
                probs = torch.softmax(alphabet_model(tensor), dim=1)[0]
            conf, idx = probs.max(dim=0)
            if conf.item() >= MIN_CNN_CONFIDENCE:
                letter = ALPHA_CLASSES[idx.item()]
                cnn_votes[letter] = cnn_votes.get(letter, 0) + 1

        # ── LSTM: pose + both hands for coord buffer ─────────
        pose_result = pose_detector.detect_for_video(mp_img, timestamp_ms)
        pose_lms    = pose_result.pose_landmarks[0] if pose_result.pose_landmarks else None
        if pose_lms:
            last_pose_vec = normalize_pose(pose_lms)

        lstm_result             = hand_detector_phrase.detect_for_video(mp_img, timestamp_ms)
        left_lms, right_lms     = get_hands_by_side(lstm_result)
        if left_lms  is not None: last_left_hand_vec  = normalize_hand(left_lms)
        if right_lms is not None: last_right_hand_vec = normalize_hand(right_lms)

        if left_lms is not None and right_lms is not None:
            two_hand_frames += 1

        coord_buffer.append(
            np.concatenate([last_pose_vec, last_left_hand_vec, last_right_hand_vec])
        )

    # ── Step 1: Hand presence check ──────────────────────────
    # Need a hand in at least 1/3 of frames, else nothing to classify
    if total_hand_frames < total_frames * MIN_HAND_PRESENCE_RATIO:
        print(f"[auto] Hand presence too low: {total_hand_frames}/{total_frames}")
        return AutoResponse(
            result="", result_type="none",
            confidence=0.0, vote_ratio=0.0, detected=False
        )

    # ── Step 2: Two hands → LSTM ─────────────────────────────
    # Two-handed signs are always phrases, skip CNN entirely
    if two_hand_frames >= total_frames * TWO_HAND_RATIO:
        print(f"[auto] Two hands in {two_hand_frames}/{total_frames} frames → LSTM")
        return run_lstm(coord_buffer, vote_ratio=0.0)

    # ── Step 3: Movement check → LSTM ────────────────────────
    # Compare wrist position at midpoint vs end of clip.
    # Ignores the initial arm-raise by starting from the midpoint.
    movement = 0.0
    if len(wrist_positions) >= 4:
        mid_pos  = wrist_positions[len(wrist_positions) // 2]
        end_pos  = wrist_positions[-1]
        dx       = end_pos[0] - mid_pos[0]
        dy       = end_pos[1] - mid_pos[1]
        movement = (dx**2 + dy**2) ** 0.5

    print(f"[auto] Wrist movement (mid→end): {movement:.4f}")

    if movement >= MOVEMENT_THRESHOLD:
        print(f"[auto] Movement {movement:.4f} >= {MOVEMENT_THRESHOLD} → LSTM")
        return run_lstm(coord_buffer, vote_ratio=0.0)

    # ── Step 4: CNN vote → letter ─────────────────────────────
    # Static sign with enough CNN agreement → return letter
    if not cnn_votes:
        print("[auto] No CNN votes collected → LSTM fallback")
        return run_lstm(coord_buffer, vote_ratio=0.0)

    best_letter = max(cnn_votes, key=cnn_votes.get)
    vote_ratio  = cnn_votes[best_letter] / sum(cnn_votes.values())

    print(f"[auto] CNN vote: {best_letter} @ {vote_ratio:.2f}")

    if vote_ratio >= MIN_CNN_VOTE_RATIO:
        return AutoResponse(
            result=best_letter,
            result_type="letter",
            confidence=vote_ratio,
            vote_ratio=vote_ratio,
            detected=True
        )

    # ── Step 5: Low CNN agreement → LSTM fallback ────────────
    print(f"[auto] CNN vote_ratio {vote_ratio:.2f} < {MIN_CNN_VOTE_RATIO} → LSTM")
    return run_lstm(coord_buffer, vote_ratio)