import cv2
import numpy as np
from insightface.app import FaceAnalysis
import os
import json
from datetime import datetime, timezone

REAL_FACES_DB = "faces_db"
TEMP_DB_ROOT = "temp_face_database"
TEMP_EMB_ROOT = "temp_faces_db"

os.makedirs(TEMP_DB_ROOT, exist_ok=True)
os.makedirs(TEMP_EMB_ROOT, exist_ok=True)

def load_database():
    db = {}
    if os.path.exists(REAL_FACES_DB):
        for file in os.listdir(REAL_FACES_DB):
            if file.endswith(".npy"):
                name = file.replace(".npy", "")
                db[name] = np.load(os.path.join(REAL_FACES_DB, file))

    if os.path.exists(TEMP_EMB_ROOT):
        for file in os.listdir(TEMP_EMB_ROOT):
            if file.endswith(".npy"):
                name = file.replace(".npy", "")
                db[name] = np.load(os.path.join(TEMP_EMB_ROOT, file))
    return db

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

def get_next_unknown_id():
    existing = [d for d in os.listdir(TEMP_DB_ROOT)
                if os.path.isdir(os.path.join(TEMP_DB_ROOT, d)) and d.startswith("unknown_")]
    if not existing:
        return 1
    ids = []
    for d in existing:
        try:
            ids.append(int(d.split("_")[1]))
        except (IndexError, ValueError):
            pass
    return max(ids) + 1 if ids else 1

def log_interaction(interaction_graph, a, b, timestamp):
    if a == b:
        return
    interaction_graph.setdefault(a, {}).setdefault(b, []).append({
        "timestamp": timestamp,
        "camera": "cam_1"
    })
    interaction_graph.setdefault(b, {}).setdefault(a, []).append({
        "timestamp": timestamp,
        "camera": "cam_1"
    })

def build_levels(root, graph):
    level_1 = set(graph.get(root, {}).keys())

    level_2 = set()
    for p in level_1:
        level_2.update(graph.get(p, {}).keys())
    level_2 -= level_1
    level_2.discard(root)

    level_3 = set()
    for p in level_2:
        level_3.update(graph.get(p, {}).keys())
    level_3 -= level_2
    level_3 -= level_1
    level_3.discard(root)

    return level_1, level_2, level_3

def format_level(interaction_graph, level_set, via=None):
    result = []
    for person in level_set:
        entry = {
            "name": person,
            "interactions": interaction_graph.get(person, {})
        }
        if via:
            entry["interacted_via"] = via.get(person, "")
        result.append(entry)
    return result

def main():
    app = FaceAnalysis(name='buffalo_l', allowed_modules=['detection', 'recognition'])
    app.prepare(ctx_id=-1, det_size=(640, 640))

    db = load_database()
    root_person = input("Enter the name of the person to track: ").strip()

    interaction_graph = {}
    frame_id = 0

    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        raise SystemExit("Could not open webcam.")

    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                break

            frame_id += 1
            faces = app.get(frame)

            detected_people = []
            bboxes = []

            for face in faces:
                x1, y1, x2, y2 = face.bbox.astype(int)
                emb = face.embedding

                best_match = "Unknown"
                best_score = 0.0

                for name, db_emb in db.items():
                    if db_emb.ndim == 1:
                        score = cosine_similarity(emb, db_emb)
                    else:
                        scores = [cosine_similarity(emb, view) for view in db_emb]
                        score = max(scores) if scores else 0.0

                    if score > best_score:
                        best_score = score
                        best_match = name

                threshold = 0.35 if best_match.startswith("unknown") else 0.30

                if best_score > threshold:
                    label = best_match
                    color = (0, 255, 0)
                else:
                    new_id = get_next_unknown_id()
                    best_match = f"unknown_{new_id}"
                    db[best_match] = np.array([emb])
                    label = best_match
                    color = (0, 255, 0)

                detected_people.append(best_match)
                bboxes.append((x1, y1, x2, y2))

                cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                cv2.putText(frame, label, (x1, y1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)

            frame_width = frame.shape[1]
            proximity_threshold = 0.25 * frame_width

            seen_pairs = set()
            timestamp = datetime.now(timezone.utc).isoformat()

            for i in range(len(detected_people)):
                for j in range(i + 1, len(detected_people)):
                    a = detected_people[i]
                    b = detected_people[j]

                    if a == b:
                        continue

                    (x1a, y1a, x2a, y2a) = bboxes[i]
                    (x1b, y1b, x2b, y2b) = bboxes[j]

                    center_a = ((x1a + x2a) / 2, (y1a + y2a) / 2)
                    center_b = ((x1b + x2b) / 2, (y1b + y2b) / 2)

                    distance = np.linalg.norm(np.array(center_a) - np.array(center_b))

                    if distance < proximity_threshold:
                        pair = tuple(sorted([a, b]))
                        if pair not in seen_pairs:
                            log_interaction(interaction_graph, a, b, timestamp)
                            seen_pairs.add(pair)

            cv2.imshow("Live Face Recognition", frame)

            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    finally:
        cap.release()
        cv2.destroyAllWindows()

    level_1, level_2, level_3 = build_levels(root_person, interaction_graph)

    output = {
        "root_person": root_person,
        "contacts": {
            "level_1": format_level(interaction_graph, level_1),
            "level_2": format_level(interaction_graph, level_2),
            "level_3": format_level(interaction_graph, level_3)
        }
    }

    with open("interaction_output.json", "w") as f:
        json.dump(output, f, indent=2)


if __name__ == '__main__':
    main()