from flask import Flask, render_template, request, jsonify
import os
import cv2
import numpy as np
import pandas as pd
from datetime import datetime
import base64

app = Flask(__name__)

# Setup
face_cascade = cv2.CascadeClassifier("haarcascade_frontalface_alt.xml")
dataset_path = "./face_dataset/"
os.makedirs(dataset_path, exist_ok=True)

# KNN distance
def distance(v1, v2):
    return np.sqrt(((v1 - v2) ** 2).sum())

def knn(train, test, k=5):
    dist = []
    for i in range(train.shape[0]):
        ix = train[i, :-1]
        iy = train[i, -1]
        d = distance(test, ix)
        dist.append([d, iy])
    dk = sorted(dist, key=lambda x: x[0])[:k]
    labels = np.array(dk)[:, -1]
    return np.unique(labels, return_counts=True)[0][0]

# Attendance system
class AttendanceSystem:
    def __init__(self):
        self.file = "attendance.csv"
        self.columns = ["Name", "Date", "Time"]
        if not os.path.exists(self.file):
            pd.DataFrame(columns=self.columns).to_csv(self.file, index=False)

    def mark(self, name):
        today = datetime.now().strftime("%Y-%m-%d")
        now = datetime.now().strftime("%H:%M:%S")
        df = pd.read_csv(self.file)
        existing = df[(df["Name"] == name) & (df["Date"] == today)]
        if existing.empty:
            new_entry = pd.DataFrame([[name, today, now]], columns=self.columns)
            df = pd.concat([df, new_entry], ignore_index=True)
            df.to_csv(self.file, index=False)
            return True
        return False

attendance = AttendanceSystem()

# Home page
@app.route('/')
def index():
    return render_template('index.html')

@app.route('/register')
def register():
    return render_template('register.html')

@app.route('/mark')
def mark():
    return render_template('mark.html')

# API to save face during registration
# Modified register_face endpoint with preprocessing
@app.route('/api/register_face', methods=['POST'])
def register_face():
    data = request.json
    name = data['name']
    images = data['images']  # List of base64 images

    face_data = []

    for img_data in images:
        img_bytes = base64.b64decode(img_data.split(",")[1])
        np_arr = np.frombuffer(img_bytes, np.uint8)
        img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        gray = cv2.equalizeHist(gray)  # Histogram equalization
        faces = face_cascade.detectMultiScale(gray, 1.3, 5)

        for (x, y, w, h) in faces[:1]:
            face = img[y:y+h, x:x+w]
            face = cv2.resize(face, (100, 100))
            
            # Original face
            face_data.append(face.flatten())
            
            # Data augmentation: horizontal flip
            flipped_face = cv2.flip(face, 1)
            face_data.append(flipped_face.flatten())

    if face_data:
        face_data = np.array(face_data)
        np.save(os.path.join(dataset_path, f"{name}.npy"), face_data)
        return jsonify({"status": "success", "message": f"{len(face_data)} faces saved"})
    else:
        return jsonify({"status": "fail", "message": "No faces detected"})

# API to mark attendance
@app.route('/api/mark_attendance', methods=['POST'])
def mark_attendance():
    img_data = request.json['image']
    img_bytes = base64.b64decode(img_data.split(",")[1])
    np_arr = np.frombuffer(img_bytes, np.uint8)
    img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, 1.3, 5)

    # Load training data
    face_data = []
    labels = []
    names = {}
    class_id = 0

    for file in os.listdir(dataset_path):
        if file.endswith('.npy'):
            data = np.load(os.path.join(dataset_path, file))
            face_data.append(data)
            names[class_id] = file[:-4]
            labels.extend([class_id] * data.shape[0])
            class_id += 1

    if not face_data:
        return jsonify({"status": "fail", "message": "No trained data found"})

    X_train = np.concatenate(face_data, axis=0)
    y_train = np.array(labels).reshape(-1, 1)
    trainset = np.hstack((X_train, y_train))

    for (x, y, w, h) in faces[:1]:
        face = img[y:y+h, x:x+w]
        face = cv2.resize(face, (100, 100)).flatten()
        pred_id = knn(trainset, face)
        name = names.get(pred_id, "Unknown")

        if name != "Unknown":
            marked = attendance.mark(name)
            msg = "Attendance marked" if marked else "Already marked today"
            return jsonify({"status": "success", "name": name, "message": msg})

    return jsonify({"status": "fail", "message": "No known face detected"})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=True)