#!/usr/bin/env python3
"""
SignFlow — STGCN2D Sign Language Recognition Server (HF Spaces deployable)

Self-contained: vendors the entire fall_2025 pipeline inline.
- MediaPipe Holistic extraction with hand correction
- Shoulder-centered pose normalization
- Standard ST-GCN architecture (K=3 spatial partitioning, 10 layers)
- 51 nodes: 17 pose + 17 left hand + 17 right hand (no pinky)
- Channels: auto-detected from checkpoint
"""

import os
import sys
import base64
import threading
from pathlib import Path
from typing import Dict, List, Optional
from collections import deque
from dataclasses import dataclass, field

import cv2
import numpy as np
import torch
import torch.nn as nn
import pandas as pd

from flask import Flask, request, send_from_directory, jsonify
from flask_cors import CORS
from flask_socketio import SocketIO, emit

# =============================================================================
# Landmark constants — subset used by the model (drops pinky to get 17)
# =============================================================================
POSE_LANDMARKS = list(range(17))   # MediaPipe pose: 0..16 (nose, eyes, ears, shoulders, elbows, wrists)
HAND_LANDMARKS = list(range(17))   # MediaPipe hand: 0..16 (wrist + thumb + index + middle + ring; drops pinky 17-20)
NUM_NODES = len(POSE_LANDMARKS) + 2 * len(HAND_LANDMARKS)  # 51

POSE_CONNECTIONS = [
    (0,1),(1,2),(2,3),(3,7),(0,4),(4,5),(5,6),(6,8),
    (9,10),(11,12),(11,13),(13,15),(12,14),(14,16),
    (0,9),(0,10),(0,11),(0,12)
]
HAND_CONNECTIONS = [
    (0,1),(1,2),(2,3),(3,4),(0,5),(5,6),(6,7),(7,8),
    (0,9),(9,10),(10,11),(11,12),(0,13),(13,14),(14,15),(15,16),
    (5,9),(9,13)
]
LEFT_WRIST_POSE_IDX = 15
RIGHT_WRIST_POSE_IDX = 16

# =============================================================================
# Adjacency matrix construction (ST-GCN spatial partitioning, K=3)
# =============================================================================
def build_adjacency() -> np.ndarray:
    """Build (3, V, V) adjacency for ST-GCN spatial partitioning."""
    V = NUM_NODES
    A = np.zeros((V, V), dtype=np.float32)

    # Self-loops
    for i in range(V):
        A[i, i] = 1.0

    # Pose connections (nodes 0..16)
    for s, e in POSE_CONNECTIONS:
        if s < len(POSE_LANDMARKS) and e < len(POSE_LANDMARKS):
            A[s, e] = A[e, s] = 1.0

    # Left hand (nodes 17..33), right hand (nodes 34..50)
    for offset in [len(POSE_LANDMARKS), len(POSE_LANDMARKS) + len(HAND_LANDMARKS)]:
        for s, e in HAND_CONNECTIONS:
            if s < len(HAND_LANDMARKS) and e < len(HAND_LANDMARKS):
                A[offset + s, offset + e] = A[offset + e, offset + s] = 1.0

    # Attach hand wrists to pose wrists
    left_offset = len(POSE_LANDMARKS)
    right_offset = len(POSE_LANDMARKS) + len(HAND_LANDMARKS)
    A[LEFT_WRIST_POSE_IDX, left_offset] = A[left_offset, LEFT_WRIST_POSE_IDX] = 1.0
    A[RIGHT_WRIST_POSE_IDX, right_offset] = A[right_offset, RIGHT_WRIST_POSE_IDX] = 1.0

    # Spatial partitioning: K=3 (self, inward, outward) using distance from center
    # Center node = root (we use node 0, the nose)
    center = 0
    # BFS distances from center
    dist = np.full(V, -1, dtype=np.int32)
    dist[center] = 0
    frontier = [center]
    while frontier:
        nxt = []
        for u in frontier:
            for v in range(V):
                if A[u, v] > 0 and dist[v] == -1:
                    dist[v] = dist[u] + 1
                    nxt.append(v)
        frontier = nxt
    # Any disconnected nodes get max distance
    dist[dist == -1] = dist.max() + 1

    A_self = np.zeros((V, V), dtype=np.float32)
    A_in = np.zeros((V, V), dtype=np.float32)
    A_out = np.zeros((V, V), dtype=np.float32)
    for i in range(V):
        for j in range(V):
            if A[i, j] > 0:
                if dist[j] == dist[i]:
                    A_self[i, j] = 1.0
                elif dist[j] > dist[i]:
                    A_out[i, j] = 1.0
                else:
                    A_in[i, j] = 1.0

    # Normalize each partition
    def normalize(M):
        D = M.sum(axis=0)
        D[D == 0] = 1.0
        return M / D[None, :]

    A_stack = np.stack([normalize(A_self), normalize(A_in), normalize(A_out)], axis=0)
    return A_stack  # (3, V, V)

# =============================================================================
# ST-GCN model (matches the fall_2025/CV/src/models/stgcn_2d.py architecture)
# =============================================================================
class ConvTemporalGraphical(nn.Module):
    def __init__(self, in_ch, out_ch, kernel_size, t_kernel=1, t_stride=1, t_pad=0, t_dil=1, bias=True):
        super().__init__()
        self.kernel_size = kernel_size  # K (spatial partitions)
        self.conv = nn.Conv2d(in_ch, out_ch * kernel_size, (t_kernel, 1), (t_stride, 1), (t_pad, 0), (t_dil, 1), bias=bias)

    def forward(self, x, A):
        x = self.conv(x)
        n, kc, t, v = x.size()
        x = x.view(n, self.kernel_size, kc // self.kernel_size, t, v)
        x = torch.einsum("nkctv,kvw->nctw", x, A).contiguous()
        return x, A


class STGCNLayer(nn.Module):
    def __init__(self, in_ch, out_ch, kernel_size, stride=1, dropout=0, residual=True):
        super().__init__()
        padding = ((kernel_size[0] - 1) // 2, 0)
        self.gcn = ConvTemporalGraphical(in_ch, out_ch, kernel_size[1])
        self.tcn = nn.Sequential(
            nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, (kernel_size[0], 1), (stride, 1), padding),
            nn.BatchNorm2d(out_ch), nn.Dropout(dropout, inplace=True),
        )
        if not residual:
            self.residual = lambda x: 0
        elif in_ch == out_ch and stride == 1:
            self.residual = lambda x: x
        else:
            self.residual = nn.Sequential(nn.Conv2d(in_ch, out_ch, 1, (stride, 1)), nn.BatchNorm2d(out_ch))
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x, A):
        res = self.residual(x)
        x, A = self.gcn(x, A)
        return self.relu(self.tcn(x) + res), A


class STGCN2D(nn.Module):
    def __init__(self, in_channels: int, num_class: int, graph_args=None, edge_importance: bool = True, dropout: float = 0, edge_importance_weighting=None):
        super().__init__()
        if edge_importance_weighting is not None: edge_importance = edge_importance_weighting
        A = build_adjacency()  # (3, V, V)
        self.register_buffer("A", torch.tensor(A, dtype=torch.float32))

        K_spatial = A.shape[0]  # 3
        V = A.shape[1]          # 51
        kernel = (9, K_spatial)
        self.data_bn = nn.BatchNorm1d(in_channels * V)

        self.st_gcn_networks = nn.ModuleList([
            STGCNLayer(in_channels, 64, kernel, 1, residual=False, dropout=dropout),
            STGCNLayer(64, 64, kernel, 1, dropout=dropout),
            STGCNLayer(64, 64, kernel, 1, dropout=dropout),
            STGCNLayer(64, 64, kernel, 1, dropout=dropout),
            STGCNLayer(64, 128, kernel, 2, dropout=dropout),
            STGCNLayer(128, 128, kernel, 1, dropout=dropout),
            STGCNLayer(128, 128, kernel, 1, dropout=dropout),
            STGCNLayer(128, 256, kernel, 2, dropout=dropout),
            STGCNLayer(256, 256, kernel, 1, dropout=dropout),
            STGCNLayer(256, 256, kernel, 1, dropout=dropout),
        ])

        if edge_importance:
            self.edge_importance = nn.ParameterList([nn.Parameter(torch.ones(A.shape)) for _ in self.st_gcn_networks])
        else:
            self.edge_importance = [1] * len(self.st_gcn_networks)

        self.fcn = nn.Conv2d(256, num_class, 1)

    def forward(self, x):
        # x: (N, C, T, V)
        N, C, T, V = x.size()
        # data_bn over (C*V) dims
        x = x.permute(0, 3, 1, 2).contiguous().view(N, V * C, T)
        x = self.data_bn(x)
        x = x.view(N, V, C, T).permute(0, 2, 3, 1).contiguous()  # back to (N, C, T, V)

        for gcn, importance in zip(self.st_gcn_networks, self.edge_importance):
            x, _ = gcn(x, self.A * importance)

        # Global pooling over T, V
        x = nn.functional.adaptive_avg_pool2d(x, 1)
        x = self.fcn(x)
        x = x.view(N, -1)
        return x

# =============================================================================
# MediaPipe extraction with hand correction + shoulder-centered normalization
# =============================================================================
class MediaPipeExtractor:
    def __init__(self):
        from mediapipe import solutions as mp_solutions
        self.mp_holistic = mp_solutions.holistic
        self.holistic = self.mp_holistic.Holistic(
            static_image_mode=False,
            model_complexity=1,
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5,
        )
        self._lock = threading.Lock()
        print("✓ MediaPipe Holistic initialized")

    def extract(self, image: np.ndarray) -> Optional[Dict]:
        """Return shoulder-centered landmarks as dict of lists of (x, y, z, vis)."""
        rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        with self._lock:
            results = self.holistic.process(rgb)

        if not results.pose_landmarks:
            return None

        pose_lms = results.pose_landmarks.landmark
        # Shoulder-center: midpoint of L/R shoulders (indices 11, 12)
        left_sh = pose_lms[11]
        right_sh = pose_lms[12]
        cx = (left_sh.x + right_sh.x) / 2
        cy = (left_sh.y + right_sh.y) / 2
        cz = (left_sh.z + right_sh.z) / 2
        # Scale: shoulder width
        scale = max(abs(left_sh.x - right_sh.x), 1e-3)

        unified = {"pose": [], "left_hand": [], "right_hand": []}

        # Pose (subset POSE_LANDMARKS, shoulder-centered, scaled)
        for i in POSE_LANDMARKS:
            lm = pose_lms[i]
            unified["pose"].append({
                "x": (lm.x - cx) / scale,
                "y": (lm.y - cy) / scale,
                "z": (lm.z - cz) / scale,
                "visibility": lm.visibility,
            })

        # Wrists in shoulder-centered space (for hand anchoring)
        left_wrist_sc = unified["pose"][LEFT_WRIST_POSE_IDX]
        right_wrist_sc = unified["pose"][RIGHT_WRIST_POSE_IDX]

        # LEFT hand: anchor to left wrist
        if results.left_hand_landmarks:
            lhms = results.left_hand_landmarks.landmark
            hw = lhms[0]
            for i in HAND_LANDMARKS:
                lm = lhms[i]
                unified["left_hand"].append({
                    "x": left_wrist_sc["x"] + (lm.x - hw.x) / scale,
                    "y": left_wrist_sc["y"] + (lm.y - hw.y) / scale,
                    "z": left_wrist_sc["z"] + (lm.z - hw.z) / scale,
                    "visibility": 1.0,
                })
        else:
            for _ in HAND_LANDMARKS:
                unified["left_hand"].append({"x": 0, "y": 0, "z": 0, "visibility": 0})

        # RIGHT hand: anchor to right wrist
        if results.right_hand_landmarks:
            rhms = results.right_hand_landmarks.landmark
            hw = rhms[0]
            for i in HAND_LANDMARKS:
                lm = rhms[i]
                unified["right_hand"].append({
                    "x": right_wrist_sc["x"] + (lm.x - hw.x) / scale,
                    "y": right_wrist_sc["y"] + (lm.y - hw.y) / scale,
                    "z": right_wrist_sc["z"] + (lm.z - hw.z) / scale,
                    "visibility": 1.0,
                })
        else:
            for _ in HAND_LANDMARKS:
                unified["right_hand"].append({"x": 0, "y": 0, "z": 0, "visibility": 0})

        unified["_has_left_hand"] = results.left_hand_landmarks is not None
        unified["_has_right_hand"] = results.right_hand_landmarks is not None
        return unified


def frames_to_tensor(frames: List[Dict], in_channels: int) -> torch.Tensor:
    """Convert a list of unified-landmark frames to (1, C, T, V) tensor."""
    T = len(frames)
    V = NUM_NODES
    arr = np.zeros((T, V, in_channels), dtype=np.float32)

    for t, f in enumerate(frames):
        row = 0
        for part in ["pose", "left_hand", "right_hand"]:
            for lm in f[part]:
                if in_channels >= 1: arr[t, row, 0] = lm["x"]
                if in_channels >= 2: arr[t, row, 1] = lm["y"]
                if in_channels >= 3: arr[t, row, 2] = lm["z"]
                if in_channels >= 4: arr[t, row, 3] = lm["visibility"]
                row += 1

    # Reshape to (1, C, T, V)
    tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
    return tensor

# =============================================================================
# Configuration
# =============================================================================
MODEL_PATH = os.environ.get("MODEL_PATH", "best_model.pth")
GLOSS_MAP_PATH = os.environ.get("GLOSS_MAP_PATH", "gloss_map.csv")
BUFFER_SIZE = 48
CONFIDENCE_THRESHOLD = 0.0
TOP_K = 5

# =============================================================================
# Globals
# =============================================================================
app = Flask(__name__, static_folder='dist', static_url_path='')
CORS(app)
socketio = SocketIO(app, cors_allowed_origins="*", async_mode='threading')

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model: Optional[STGCN2D] = None
extractor: Optional[MediaPipeExtractor] = None
gloss_map: Dict[int, str] = {}
IN_CHANNELS: int = 2  # auto-detected from checkpoint

@dataclass
class Session:
    sid: str
    frame_buffer: deque = field(default_factory=lambda: deque(maxlen=BUFFER_SIZE))
    last_prediction: Optional[List] = None
    processing: bool = False

sessions: Dict[str, Session] = {}

# =============================================================================
# Model loading with auto-detected channel count
# =============================================================================
def load_model() -> bool:
    global model, extractor, gloss_map, IN_CHANNELS

    # Gloss map
    if os.path.exists(GLOSS_MAP_PATH):
        df = pd.read_csv(GLOSS_MAP_PATH)
        col = 'gloss' if 'gloss' in df.columns else ('label' if 'label' in df.columns else df.columns[-1])
        gloss_map = {i: str(g) for i, g in enumerate(df[col].tolist())}
        print(f"✓ Loaded {len(gloss_map)} glosses from column '{col}'")
    else:
        gloss_map = {i: f"SIGN_{i}" for i in range(208)}
        print("⚠ No gloss map, using defaults")

    if not os.path.exists(MODEL_PATH):
        print(f"✗ Model not found: {MODEL_PATH}")
        return False

    ckpt = torch.load(MODEL_PATH, map_location=device, weights_only=False)

    # New checkpoint format: {"config": ..., "state": ..., "source_code": ..., "meta": ...}
    if isinstance(ckpt, dict) and "state" in ckpt:
        state = ckpt["state"]
        ckpt_config = ckpt.get("config", {})
        source_code = ckpt.get("source_code", "")
        meta = ckpt.get("meta", {})
        print(f"✓ New checkpoint format. Meta: epoch={meta.get('epoch','?')}, val_acc={meta.get('val_accuracy','?')}")
    else:
        state = ckpt
        ckpt_config = {}
        source_code = ""

    # Strip prefixes
    state = {k.replace('module.', '').replace('_orig_mod.', ''): v for k, v in state.items()}

    # Get config params
    model_cfg = ckpt_config.get("model", {})
    data_cfg = ckpt_config.get("data", {})
    IN_CHANNELS = model_cfg.get("num_features", 4)
    num_classes = model_cfg.get("num_classes", len(gloss_map))
    edge_importance = model_cfg.get("edge_importance_weighting", True)
    dropout = model_cfg.get("dropout", 0)
    body_parts = data_cfg.get("body_parts", ["pose", "left_hand", "right_hand"])
    print(f"✓ Config: in_channels={IN_CHANNELS}, num_classes={num_classes}, body_parts={body_parts}")

    global model
    used_source = False

    # Try to use the source_code from checkpoint for exact architecture match
    if source_code:
        try:
            tmp_path = "/tmp/_ckpt_model.py"
            with open(tmp_path, "w") as f:
                f.write(source_code)
            import importlib.util
            spec = importlib.util.spec_from_file_location("_ckpt_model", tmp_path)
            mod = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(mod)
            # Find the Model class
            ModelCls = getattr(mod, "Model", None) or getattr(mod, "STGCN2D", None) or getattr(mod, "STGCN", None)
            if ModelCls is None:
                # Find any nn.Module subclass at the top level
                for name in dir(mod):
                    obj = getattr(mod, name)
                    if isinstance(obj, type) and issubclass(obj, nn.Module) and obj is not nn.Module:
                        ModelCls = obj
                        break
            if ModelCls is not None:
                graph_args = model_cfg.get("graph_args", {
                    "body_parts": body_parts,
                    "strategy": model_cfg.get("graph_strategy", "spatial"),
                    "max_hop": model_cfg.get("max_hop", 1),
                })
                # Try common constructor signatures
                try:
                    model = ModelCls(
                        in_channels=IN_CHANNELS,
                        num_class=num_classes,
                        graph_args=graph_args,
                        edge_importance_weighting=edge_importance,
                        dropout=dropout,
                    )
                except TypeError:
                    try:
                        model = ModelCls(
                            in_channels=IN_CHANNELS,
                            num_class=num_classes,
                            graph_args=graph_args,
                            edge_importance=edge_importance,
                            dropout=dropout,
                        )
                    except TypeError:
                        model = ModelCls(IN_CHANNELS, num_classes, graph_args)
                used_source = True
                print(f"✓ Built model from checkpoint source_code: {ModelCls.__name__}")
        except Exception as e:
            print(f"⚠ Could not use source_code: {e}")
            import traceback; traceback.print_exc()

    # Fallback to vendored STGCN2D
    if not used_source:
        print("⚠ Falling back to vendored STGCN2D")
        model = STGCN2D(in_channels=IN_CHANNELS, num_class=num_classes, edge_importance=edge_importance, dropout=dropout)

    missing, unexpected = model.load_state_dict(state, strict=False)
    print(f"⚠ Missing keys: {len(missing)}, Unexpected: {len(unexpected)}")
    if missing[:3]:
        print(f"  First missing: {missing[:3]}")
    if unexpected[:3]:
        print(f"  First unexpected: {unexpected[:3]}")

    model.to(device).eval()
    print(f"✓ Model on {device}")

    extractor = MediaPipeExtractor()
    return True


# =============================================================================
# Inference
# =============================================================================
def run_inference(frames: List[Dict]) -> List[Dict]:
    try:
        tensor = frames_to_tensor(frames, IN_CHANNELS).to(device)
        with torch.no_grad():
            logits = model(tensor)
            probs = torch.softmax(logits, dim=-1)[0]
        top_probs, top_idx = torch.topk(probs, TOP_K)
        results = []
        for p, i in zip(top_probs.cpu().numpy(), top_idx.cpu().numpy()):
            if p >= CONFIDENCE_THRESHOLD:
                results.append({
                    'gloss': gloss_map.get(int(i), f"SIGN_{i}"),
                    'confidence': float(p),
                    'index': int(i),
                })
        if results:
            print(f"[INFER] top={results[0]['gloss']} conf={results[0]['confidence']:.3f}", flush=True)
        return results
    except Exception as e:
        print(f"[INFER ERROR] {e}", flush=True)
        import traceback; traceback.print_exc()
        return []

# =============================================================================
# Frame processing
# =============================================================================
def process_frame(session: Session, frame_data: str) -> Optional[Dict]:
    try:
        img_bytes = base64.b64decode(frame_data.split(',')[-1])
        arr = np.frombuffer(img_bytes, dtype=np.uint8)
        frame = cv2.imdecode(arr, cv2.IMREAD_COLOR)
        if frame is None:
            return None

        landmarks = extractor.extract(frame)
        if landmarks is None:
            return {
                'status': 'no_person',
                'predictions': [],
                'buffer_fill': len(session.frame_buffer) / BUFFER_SIZE,
                'has_hands': False,
            }

        session.frame_buffer.append(landmarks)
        has_hands = landmarks.get('_has_left_hand') or landmarks.get('_has_right_hand')

        if len(session.frame_buffer) >= BUFFER_SIZE and not session.processing:
            session.processing = True
            try:
                preds = run_inference(list(session.frame_buffer))
                session.last_prediction = preds
                return {
                    'status': 'prediction',
                    'predictions': preds,
                    'has_hands': has_hands,
                    'buffer_fill': 1.0,
                }
            finally:
                session.processing = False

        return {
            'status': 'buffering',
            'buffer_fill': len(session.frame_buffer) / BUFFER_SIZE,
            'has_hands': has_hands,
            'predictions': session.last_prediction or [],
        }
    except Exception as e:
        print(f"[PROCESS ERROR] {e}", flush=True)
        return None

# =============================================================================
# Socket.IO + HTTP routes
# =============================================================================
@socketio.on('connect')
def on_connect():
    sid = request.sid
    sessions[sid] = Session(sid=sid)
    print(f"Client connected: {sid[:8]}", flush=True)
    emit('status', {
        'connected': True,
        'model_loaded': model is not None,
        'num_glosses': len(gloss_map),
    })

@socketio.on('disconnect')
def on_disconnect():
    sid = request.sid
    sessions.pop(sid, None)

@socketio.on('frame')
def on_frame(data):
    sid = request.sid
    if sid not in sessions:
        return
    result = process_frame(sessions[sid], data.get('image', ''))
    if result:
        emit('prediction', result)

@socketio.on('reset')
def on_reset():
    sid = request.sid
    if sid in sessions:
        sessions[sid].frame_buffer.clear()
        sessions[sid].last_prediction = None

@app.route('/')
def index():
    return send_from_directory(app.static_folder, 'index.html')

@app.route('/health')
def health():
    return jsonify({'ok': model is not None, 'glosses': len(gloss_map), 'channels': IN_CHANNELS})

# =============================================================================
# Main
# =============================================================================
if __name__ == '__main__':
    print("=" * 60)
    print("  SignFlow STGCN2D Server")
    print("=" * 60)
    if not load_model():
        sys.exit(1)
    port = int(os.environ.get('PORT', 7860))
    print(f"🚀 Listening on 0.0.0.0:{port}", flush=True)
    socketio.run(app, host='0.0.0.0', port=port, debug=False, allow_unsafe_werkzeug=True)