# train.py — LBPH training with in-memory augmentation
import glob
import os
import pickle
from pathlib import Path
from typing import List, Tuple, Dict

import cv2
import numpy as np

APP_ROOT = Path(__file__).parent.resolve()
DATA_DIR = APP_ROOT / "data" / "students"
MODELS_DIR = APP_ROOT / "models"
MODELS_DIR.mkdir(parents=True, exist_ok=True)
MODEL_PATH = MODELS_DIR / "lbph_model.xml"
LABELS_PATH = MODELS_DIR / "labels.pickle"
CONFIG_PATH = MODELS_DIR / "config.json"

# ------------------ augmentation helpers ------------------
def rot(img: np.ndarray, deg: float) -> np.ndarray:
    h, w = img.shape[:2]
    M = cv2.getRotationMatrix2D((w//2, h//2), deg, 1.0)
    return cv2.warpAffine(img, M, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)

def brighten(img: np.ndarray, factor: float) -> np.ndarray:
    # factor 0.8..1.2
    out = img.astype(np.float32) * factor
    return np.clip(out, 0, 255).astype(np.uint8)

def noise(img: np.ndarray, sigma: float=8.0) -> np.ndarray:
    n = np.random.normal(0, sigma, img.shape).astype(np.float32)
    out = img.astype(np.float32) + n
    return np.clip(out, 0, 255).astype(np.uint8)

def augment_face(gray200: np.ndarray) -> List[np.ndarray]:
    """Return a list of augmented versions (same 200x200 gray)."""
    aug = []
    # rotations
    aug.append(rot(gray200, +10))
    aug.append(rot(gray200, -10))
    # brightness
    aug.append(brighten(gray200, 1.2))
    aug.append(brighten(gray200, 0.8))
    # noise
    aug.append(noise(gray200, 6.0))
    # flip
    aug.append(cv2.flip(gray200, 1))
    return aug

def normalize(gray200: np.ndarray) -> np.ndarray:
    # same normalization as app.py: CLAHE + light blur
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    norm = clahe.apply(gray200)
    return cv2.GaussianBlur(norm, (3,3), 0)

# ------------------ dataset loading ------------------
def load_student_images() -> Tuple[List[np.ndarray], List[int], Dict[int, str], Dict[str, int]]:
    X, y = [], []
    label_to_code: Dict[int, str] = {}
    code_to_label: Dict[str, int] = {}

    student_dirs = sorted([p for p in DATA_DIR.iterdir() if p.is_dir()])
    label = 0
    for folder in student_dirs:
        code = folder.name
        imgs = sorted(glob.glob(str(folder / "*.jpg"))) + sorted(glob.glob(str(folder / "*.png")))
        if len(imgs) == 0:
            continue

        if code not in code_to_label:
            code_to_label[code] = label
            label_to_code[label] = code
            label += 1

        for p in imgs:
            img = cv2.imread(p, cv2.IMREAD_GRAYSCALE)
            if img is None:
                continue
            if img.shape != (200, 200):
                try:
                    img = cv2.resize(img, (200, 200))
                except Exception:
                    continue
            base = normalize(img)
            # original
            X.append(base); y.append(code_to_label[code])
            # augmented (in-memory only)
            for a in augment_face(base):
                X.append(normalize(a)); y.append(code_to_label[code])

    return X, y, label_to_code, code_to_label

# ------------------ threshold calibration ------------------
def calibrate_threshold(recognizer, X: List[np.ndarray], y: List[int]) -> float:
    """
    Evaluate recognizer's 'confidence' on its own training samples to estimate
    a safe threshold. We take the 95th percentile of confidences for
    correctly predicted samples. (LBPH: LOWER is better.)
    """
    if len(X) == 0:
        return 75.0
    confs = []
    for img, true_label in zip(X, y):
        pred, conf = recognizer.predict(img)
        if pred == true_label:
            confs.append(conf)
    if not confs:
        return 75.0
    confs = np.array(confs, dtype=np.float32)
    p95 = float(np.percentile(confs, 95))
    # clamp to a sensible band
    return float(np.clip(p95, 55.0, 95.0))

# ------------------ main train ------------------
def main(mode: str = "fast", limit_per_student: int | None = None):
    # Load + (optionally) limit samples per student to speed
    X, y, label_to_code, code_to_label = load_student_images()

    if limit_per_student:
        # balance / clip per student
        keep_X, keep_y = [], []
        from collections import defaultdict
        buckets = defaultdict(list)
        for i, lab in enumerate(y):
            buckets[lab].append(i)
        for lab, idxs in buckets.items():
            use = idxs[:limit_per_student]
            for i in use:
                keep_X.append(X[i]); keep_y.append(y[i])
        X, y = keep_X, keep_y

    students = len(set(y))
    images = len(X)
    if students == 0 or images == 0:
        raise RuntimeError("No training data found. Register students and capture images first.")

    rec = cv2.face.LBPHFaceRecognizer_create(radius=1, neighbors=8, grid_x=8, grid_y=8)
    rec.train(X, np.array(y, dtype=np.int32))

    thr = calibrate_threshold(rec, X, y)

    # atomic writes (Windows-safe)
    tmp_model = MODEL_PATH.with_suffix(".xml.tmp")
    tmp_labels = LABELS_PATH.with_suffix(".pickle.tmp")
    rec.write(str(tmp_model))
    with open(tmp_labels, "wb") as f:
        pickle.dump(label_to_code, f)
    os.replace(tmp_model, MODEL_PATH)
    os.replace(tmp_labels, LABELS_PATH)

    # also persist threshold to config.json
    cfg = {"lbph_threshold": thr, "mode": mode}
    try:
        import json
        with open(CONFIG_PATH, "w", encoding="utf-8") as f:
            json.dump(cfg, f, indent=2)
    except Exception:
        pass

    # quick stats
    stats = {
        "students": students,
        "images": images,
        "labels": len(label_to_code),
    }
    return {"mode": mode, "threshold": thr, "students": students, "images": images, "stats": stats}

if __name__ == "__main__":
    print(main())
