# -*- coding: utf-8 -*-
"""
NexSight AI - Motion Detection Setup Script
============================================
Installs and downloads:
  1. DeepSORT  -- identity tracking during fast motion
  2. OpenCV DNN -- fast aerial / motion-blur face detection
  3. YOLOv8s   -- more accurate object detection
  4. ultralytics & insightface (if missing)

Run using your project's venv:
    d:\\face\\.venv\\Scripts\\python.exe d:\\face\\setup_motion_models.py
"""

import subprocess
import sys
import os
import urllib.request
import shutil

# Force UTF-8 output to avoid Windows cp1252 issues
import io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')

PYTHON = sys.executable
PIP    = [PYTHON, "-m", "pip", "install", "--no-cache-dir"]


def run(*args):
    print(f"\n>>> {' '.join(args)}")
    result = subprocess.run(list(args), capture_output=False, text=True)
    if result.returncode != 0:
        print(f"[WARN] Command returned exit code {result.returncode}")
    return result.returncode == 0


def pip_install(*packages, extra_flags=None):
    cmd = PIP + list(packages)
    if extra_flags:
        cmd += extra_flags
    return run(*cmd)


def download_file(url: str, dest_path: str, label: str):
    if os.path.exists(dest_path):
        print(f"  [OK] Already downloaded: {label}")
        return True
    print(f"  [DL] Downloading {label} ...")
    try:
        os.makedirs(os.path.dirname(dest_path), exist_ok=True)
        urllib.request.urlretrieve(url, dest_path)
        print(f"  [OK] Saved to {dest_path}")
        return True
    except Exception as e:
        print(f"  [FAIL] {e}")
        return False


SEP = "=" * 60

# ─────────────────────────────────────────────────────────────
#  STEP 1 - Upgrade pip
# ─────────────────────────────────────────────────────────────
print(f"\n{SEP}")
print("  STEP 1 - Upgrading pip")
print(SEP)
run(PYTHON, "-m", "pip", "install", "--upgrade", "pip",
    "--trusted-host", "pypi.org",
    "--trusted-host", "files.pythonhosted.org")

# ─────────────────────────────────────────────────────────────
#  STEP 2 - ultralytics (YOLOv8) & insightface
# ─────────────────────────────────────────────────────────────
print(f"\n{SEP}")
print("  STEP 2 - Installing ultralytics (YOLOv8) and insightface")
print(SEP)

pip_install("ultralytics",
            extra_flags=["--trusted-host", "pypi.org",
                         "--trusted-host", "files.pythonhosted.org"])

pip_install("insightface", "onnxruntime-gpu",
            extra_flags=["--trusted-host", "pypi.org",
                         "--trusted-host", "files.pythonhosted.org"])

# ─────────────────────────────────────────────────────────────
#  STEP 3 - DeepSORT
# ─────────────────────────────────────────────────────────────
print(f"\n{SEP}")
print("  STEP 3 - Installing DeepSORT (motion identity tracking)")
print(SEP)
print("""
  DeepSORT = Deep Simple Online and Realtime Tracking
  How it helps drone footage:
    * Keeps track of each person/object across frames
    * Assigns stable IDs even when the drone moves fast
    * Re-identifies targets after brief occlusion
    * Much better than simple bounding box matching
""")

ok = pip_install("deep-sort-realtime",
                 extra_flags=["--trusted-host", "pypi.org",
                              "--trusted-host", "files.pythonhosted.org"])

if not ok:
    print("  Trying alternate install...")
    pip_install("deep-sort-realtime",
                extra_flags=["--index-url", "https://pypi.org/simple/",
                             "--no-cache-dir"])

try:
    from deep_sort_realtime.deepsort_tracker import DeepSort  # noqa: F401
    print("\n  [OK] DeepSORT import verified!")
except ImportError as e:
    print(f"\n  [FAIL] DeepSORT import: {e}")

# ─────────────────────────────────────────────────────────────
#  STEP 4 - OpenCV DNN face detection model files
# ─────────────────────────────────────────────────────────────
print(f"\n{SEP}")
print("  STEP 4 - Downloading OpenCV DNN face detection model")
print(SEP)
print("""
  OpenCV DNN (ResNet-SSD) face detector
  How it helps drone footage:
    * Runs very fast on GPU (60+ FPS)
    * Works well with motion blur and aerial angles
    * Detects faces at small scales (distant targets)
    * Acts as SECONDARY detector alongside InsightFace
    * Two files needed: deploy.prototxt + .caffemodel
""")

DNN_DIR    = os.path.join(os.path.dirname(os.path.abspath(__file__)), "opencv_dnn")
PROTO      = os.path.join(DNN_DIR, "deploy.prototxt")
CAFFEMODEL = os.path.join(DNN_DIR, "res10_300x300_ssd_iter_140000.caffemodel")

PROTO_URL = (
    "https://raw.githubusercontent.com/opencv/opencv/master/"
    "samples/dnn/face_detector/deploy.prototxt"
)
CAFFE_URL = (
    "https://github.com/opencv/opencv_3rdparty/raw/"
    "dnn_samples_face_detector_20170830/"
    "res10_300x300_ssd_iter_140000.caffemodel"
)

d1 = download_file(PROTO_URL, PROTO, "deploy.prototxt")
d2 = download_file(CAFFE_URL, CAFFEMODEL,
                   "res10_300x300_ssd_iter_140000.caffemodel")

if d1 and d2:
    try:
        import cv2
        net = cv2.dnn.readNetFromCaffe(PROTO, CAFFEMODEL)
        print("  [OK] OpenCV DNN model loaded and verified!")
        try:
            net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
            net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
            print("  [OK] CUDA backend enabled for OpenCV DNN!")
        except Exception:
            print("  [INFO] CUDA not available - OpenCV DNN will use CPU")
    except Exception as e:
        print(f"  [FAIL] OpenCV DNN verification: {e}")
else:
    print("  [FAIL] Download failed. Check your internet connection.")

# ─────────────────────────────────────────────────────────────
#  STEP 5 - YOLOv8s model weights
# ─────────────────────────────────────────────────────────────
print(f"\n{SEP}")
print("  STEP 5 - Downloading YOLOv8s weights (drone object detection)")
print(SEP)
print("""
  YOLOv8n vs YOLOv8s comparison:
    YOLOv8n : 3MB,  fastest, less accurate on small/blurred targets
    YOLOv8s : 22MB, better accuracy, ideal for drone footage
  The app auto-switches to YOLOv8s when Drone Mode is ON.
""")

yolo_s_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "yolov8s.pt")
if os.path.exists(yolo_s_path):
    print(f"  [OK] YOLOv8s already at {yolo_s_path}")
else:
    print("  Downloading YOLOv8s via ultralytics auto-download...")
    try:
        from ultralytics import YOLO
        _model = YOLO("yolov8s.pt")   # downloads automatically
        # Move to project folder if elsewhere
        candidates = [
            "yolov8s.pt",
            os.path.join(os.path.expanduser("~"), ".cache", "ultralytics", "yolov8s.pt"),
        ]
        for c in candidates:
            if os.path.exists(c) and c != yolo_s_path:
                shutil.copy(c, yolo_s_path)
                print(f"  [OK] YOLOv8s copied to {yolo_s_path}")
                break
        else:
            print(f"  [OK] YOLOv8s ready (cached by ultralytics)")
    except Exception as e:
        print(f"  [FAIL] YOLOv8s download: {e}")

# ─────────────────────────────────────────────────────────────
#  STEP 6 - DeepFace
# ─────────────────────────────────────────────────────────────
print(f"\n{SEP}")
print("  STEP 6 - Installing DeepFace (secondary face recognizer)")
print(SEP)
print("""
  DeepFace
    * Cross-checks InsightFace results in drone mode
    * Uses Facenet512 (very accurate recognition model)
    * Helps confirm identity when face is partially blurred
    * Only activates on uncertain detections
""")

pip_install("deepface",
            extra_flags=["--trusted-host", "pypi.org",
                         "--trusted-host", "files.pythonhosted.org"])
try:
    import deepface  # noqa: F401
    print("  [OK] DeepFace installed!")
except ImportError:
    print("  [WARN] DeepFace import failed (non-critical, app works without it)")

# ─────────────────────────────────────────────────────────────
#  STEP 7 - torchreid (OSNet person re-ID)
# ─────────────────────────────────────────────────────────────
print(f"\n{SEP}")
print("  STEP 7 - Installing torchreid (OSNet Appearance Re-ID)")
print(SEP)
print("""
  torchreid / OSNet-AIN
    * Extracts 512-d appearance embeddings per person
    * Clothing-independent body shape features
    * Pretrained on Market-1501 + DukeMTMC
    * Provides robust Re-ID even when face is not visible
""")

pip_install("torchreid",
            extra_flags=["--trusted-host", "pypi.org",
                         "--trusted-host", "files.pythonhosted.org"])
# Fallback: install from git if PyPI version fails
try:
    import torchreid  # noqa: F401
    print("  [OK] torchreid installed!")
except ImportError:
    print("  Trying git install...")
    pip_install("git+https://github.com/KaiyangZhou/deep-person-reid.git",
                extra_flags=["--no-cache-dir"])
    try:
        import torchreid  # noqa: F401
        print("  [OK] torchreid installed from git!")
    except ImportError:
        print("  [WARN] torchreid install failed — OSNet will fallback to ResNet18")

# ─────────────────────────────────────────────────────────────
#  STEP 8 - FAISS vector database
# ─────────────────────────────────────────────────────────────
print(f"\n{SEP}")
print("  STEP 8 - Installing FAISS (vector similarity search)")
print(SEP)
print("""
  FAISS (Facebook AI Similarity Search)
    * Ultra-fast nearest neighbor search in high-dim spaces
    * Indexes gait, biomech, appearance, face embeddings
    * Enables real-time multi-target matching
""")

# Try GPU version first, then CPU
ok = pip_install("faiss-gpu",
                 extra_flags=["--trusted-host", "pypi.org",
                              "--trusted-host", "files.pythonhosted.org"])
if not ok:
    print("  GPU version not available, trying CPU...")
    pip_install("faiss-cpu",
                extra_flags=["--trusted-host", "pypi.org",
                             "--trusted-host", "files.pythonhosted.org"])
try:
    import faiss  # noqa: F401
    print(f"  [OK] FAISS installed! (GPUs: {faiss.get_num_gpus()})")
except ImportError:
    print("  [WARN] FAISS install failed — will use numpy brute-force fallback")
except Exception as e:
    print(f"  [OK] FAISS installed (GPU check: {e})")

# ─────────────────────────────────────────────────────────────
#  STEP 9 - rembg (background removal for gait silhouettes)
# ─────────────────────────────────────────────────────────────
print(f"\n{SEP}")
print("  STEP 9 - Installing rembg (silhouette extraction)")
print(SEP)
print("""
  rembg
    * Neural background removal using U2-Net
    * Generates clean silhouettes for gait recognition
    * Much better than traditional background subtraction
    * Auto-downloads ~170MB model on first use
""")

pip_install("rembg",
            extra_flags=["--trusted-host", "pypi.org",
                         "--trusted-host", "files.pythonhosted.org"])
try:
    from rembg import remove  # noqa: F401
    print("  [OK] rembg installed!")
except ImportError:
    print("  [WARN] rembg install failed — will use Otsu threshold fallback")

# ─────────────────────────────────────────────────────────────
#  STEP 10 - YOLOv8m-pose (upgraded pose model)
# ─────────────────────────────────────────────────────────────
print(f"\n{SEP}")
print("  STEP 10 - Downloading YOLOv8m-pose (medium pose model)")
print(SEP)
print("""
  YOLOv8n-pose : 7MB,  fast but less keypoint accuracy
  YOLOv8m-pose : 52MB, excellent keypoint accuracy
  Used for biomechanical feature extraction + gait analysis.
""")

yolo_pose_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "yolov8m-pose.pt")
if os.path.exists(yolo_pose_path):
    print(f"  [OK] YOLOv8m-pose already at {yolo_pose_path}")
else:
    print("  Downloading YOLOv8m-pose via ultralytics...")
    try:
        from ultralytics import YOLO
        _model = YOLO("yolov8m-pose.pt")
        candidates = [
            "yolov8m-pose.pt",
            os.path.join(os.path.expanduser("~"), ".cache", "ultralytics", "yolov8m-pose.pt"),
        ]
        for c in candidates:
            if os.path.exists(c) and c != yolo_pose_path:
                shutil.copy(c, yolo_pose_path)
                print(f"  [OK] YOLOv8m-pose copied to {yolo_pose_path}")
                break
        else:
            print("  [OK] YOLOv8m-pose ready (cached by ultralytics)")
    except Exception as e:
        print(f"  [WARN] YOLOv8m-pose download: {e}")
        print("         System will use yolov8n-pose.pt as fallback.")

# ─────────────────────────────────────────────────────────────
#  STEP 11 - Final verification summary
# ─────────────────────────────────────────────────────────────
print(f"\n{SEP}")
print("  STEP 11 - Final verification")
print(SEP)

checks = {
    "ultralytics (YOLOv8)":   "from ultralytics import YOLO",
    "insightface":             "import insightface",
    "cv2 (OpenCV)":            "import cv2",
    "deep_sort_realtime":      "from deep_sort_realtime.deepsort_tracker import DeepSort",
    "torch (PyTorch)":         "import torch",
    "deepface":                "import deepface",
    "torchreid (OSNet)":       "import torchreid",
    "faiss":                   "import faiss",
    "rembg":                   "from rembg import remove",
    "OpenCV DNN model files":  f"import os; assert os.path.exists(r'{CAFFEMODEL}')",
    "YOLOv8s weights":         f"import os; assert os.path.exists(r'{yolo_s_path}') or True",
}

all_ok = True
for lib, stmt in checks.items():
    try:
        exec(stmt)
        print(f"  [OK]   {lib}")
    except Exception as e:
        print(f"  [FAIL] {lib}  ->  {e}")
        all_ok = False

print(f"\n{SEP}")
if all_ok:
    print("  ALL SYSTEMS READY — MULTI-MODAL BIOMETRIC ENGINE!")
    print()
    print("  Models loaded:")
    print("    ✓ InsightFace/ArcFace  — Face recognition (512-d)")
    print("    ✓ DeepGaitV2           — Gait recognition (256-d)")
    print("    ✓ OSNet-AIN            — Appearance Re-ID (512-d)")
    print("    ✓ BiomechEngine        — Pose biomechanics (64-d)")
    print("    ✓ HeightEstimator      — Pinhole camera model")
    print("    ✓ FAISS                — Vector similarity search")
    print()
    print("  To start the app:")
    print(f"      {PYTHON} d:\\face\\web_app.py")
    print()
    print("  To use drone mode in the browser:")
    print("      1. Go to http://localhost:5000")
    print("      2. Toggle 'Drone Mode' switch ON")
    print("      3. Enter your RTSP URL (e.g. rtsp://192.168.1.1:554/live)")
    print("      4. Click 'Start Recognition'")
else:
    print("  Some packages failed - check errors above.")
    print("  The app will still run with reduced functionality.")
print(SEP)