from flask import Flask, render_template, request, jsonify
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import PyPDF2
from keras.models import load_model
from PIL import Image, ImageOps
import numpy as np
import pandas as pd
from inference_sdk import InferenceHTTPClient
import cv2
import base64
import io
from flask import send_file
from reportlab.pdfgen import canvas
from io import BytesIO

EXCEL_FILE = "Book2.xlsx"

# Initialize the Roboflow clients for different models
CLIENTS = {
    'classroom': InferenceHTTPClient(
        api_url="https://detect.roboflow.com",
        api_key="bNLTnCBq5hIm7R0O3hU4"
    ),
    'chemical_lab': InferenceHTTPClient(
        api_url="https://detect.roboflow.com",
        api_key="bNLTnCBq5hIm7R0O3hU4"
    ),
    'mechanical_workshop': InferenceHTTPClient(
        api_url="https://detect.roboflow.com",
        api_key="bNLTnCBq5hIm7R0O3hU4"
    ),
    'computer_lab': InferenceHTTPClient(
        api_url="https://detect.roboflow.com",
        api_key="bNLTnCBq5hIm7R0O3hU4"
    ),
    'cctv' :InferenceHTTPClient(
    api_url="https://detect.roboflow.com",
    api_key="bNLTnCBq5hIm7R0O3hU4"
),
    'notice_board' :InferenceHTTPClient(
    api_url="https://detect.roboflow.com",
    api_key="bNLTnCBq5hIm7R0O3hU4"
    ),
    'bench': InferenceHTTPClient(
        api_url="https://detect.roboflow.com",
        api_key="IkQtIl5NGRTc0llwyIMo"
    )

}
# Model IDs for each environment
MODEL_IDS = {
    'classroom': "sih-object-detection/1",
    'chemical_lab': "chem-dz924/1",
    'mechanical_workshop': "mech-npugl/1",
    'computer_lab': "sih-object-detection/1",
    'cctv' : "bench-bcvxh/2",
    'notice_board' : "cctv-cofid/2",
    'bench' : "bench-bcvxh/2",
}
app = Flask(__name__)
app.secret_key = 'super_secret_key'

# Configuration
app.config['UPLOAD_FOLDER'] = os.path.abspath('uploads/')
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

FACILITIES = [
    # Essential Academic and Safety Facilities
    "Classroom model",
    "Library model",
    "Computer lab model",
    "elearning model",

    "Drawing Halls model",
    "Fire extinguisher model",

    # Faculty and Administrative Needs
    "Faculty cabin model",
    "Server Room model",
    "TPO model",

    # Recreational and Co-curricular Support
    "Ground model",
    "Sports equipment model",
    "Workshop model",
    "Seminar hall model",
    "Conference Halls model",

    # Comfort and Utility Facilities
    "Canteen model",

    "Medical Room Model",
    "Parking model",

    # Backup and Miscellaneous
    "Generator model",
    "Audi model",

]


# Mapping for PDFs (names might differ from model names)
PDF_NAMES = {
    "Audi model": "Audi.pdf",
    "Canteen model": "Canteen.pdf",
    "Classroom model": "Classroom.pdf",
    "Computer lab model": "Computer Lab.pdf",
    "Conference Halls model": "Conference Hall.pdf",
    "Drawing Halls model": "Drawing Halls.pdf",
    "Faculty cabin model": "Faculty Cabin.pdf",
    "Fire extinguisher model": "Fire Extinguishers.pdf",
    "Generator model": "Generator.pdf",
    "Ground model": "Grounds.pdf",
    "Library model": "Library.pdf",
    "Medical Room Model": "Medical Room.pdf",
    "Parking model": "Parking.pdf",
    "Restroom Model": "Restroom.pdf",
    "Seminar hall model": "Seminar Hall.pdf",
    "Server Room model": "Server Room.pdf",
    "Sports equipment model": "Sports Equipment.pdf",
    "TPO model": "TPO (Training and Placement Office).pdf",
    "Workshop model": "Workshop.pdf",
    "elearning model": "elearning.pdf",
}

# Paths
MODEL_PATHS = {
    facility: {
        "model": f"MODELS/{facility}/keras_model.h5",
        "labels": f"MODELS/{facility}/labels.txt",
    }
    for facility in FACILITIES
}

PDF_PATHS = {
    facility: f"pdfs/{PDF_NAMES[facility]}"
    for facility in FACILITIES
}

# Routes
@app.route('/')
def index():
    # Extract questions from PDFs for each facility
    questions = {facility: extract_questions(PDF_PATHS.get(facility, "")) for facility in FACILITIES}
    return render_template('index.html', facilities=FACILITIES, questions=questions)


@app.route('/calculate', methods=['POST'])
def calculate():
    data = request.json
    num_students = int(data.get('num_students', 0))
    num_divisions = int(data.get('num_divisions', 0))
    num_courses = int(data.get('num_courses', 0))
    course_duration = int(data.get('course_duration', 0))

    calculated_facilities = calculate_required_facilities(num_students, num_divisions, num_courses, course_duration)
    return jsonify(calculated_facilities)




@app.route('/upload/<facility>', methods=['POST'])
def upload(facility):
    facility = facility.strip()

    # Check if facility exists in MODEL_PATHS
    normalized_facility = next(
        (key for key in MODEL_PATHS if key.lower() == facility.lower()), None
    )
    if not normalized_facility:
        return jsonify({"error": f"Facility '{facility}' not found in MODEL_PATHS"}), 400

    if 'images' not in request.files:
        return jsonify({"error": "No files uploaded"}), 400

    files = request.files.getlist('images')
    if not files:
        return jsonify({"error": "No files selected"}), 400

    results = []
    for file in files:
        try:
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
            file.save(filepath)

            # Perform verification using the model
            model_path = MODEL_PATHS[normalized_facility]["model"]
            labels_path = MODEL_PATHS[normalized_facility]["labels"]

            result = verify_image(filepath, model_path, labels_path)
            result["file_name"] = file.filename
            result["facility"] = normalized_facility

            # Log to Excel if verified
            if result["confidence"] >= 0.8:
                log_to_excel(result)

            results.append(result)

        except Exception as e:
            print(f"Error during file upload: {e}")
            results.append({"error": str(e), "file": file.filename})

    return jsonify(results)

def log_to_excel(data):
    """
    Logs verified image data to an Excel file.
    :param data: Dictionary containing facility, file name, label, and confidence score.
    """
    # Prepare a DataFrame row
    row = {
        "Facility": data["facility"],
        "Name of Image": data["file_name"],
        "Class of Prediction": data["label"],
        "Confidence Score": data["confidence"]
    }

    # Convert row to DataFrame
    df_row = pd.DataFrame([row])

    # If the file exists, append; otherwise, create a new file
    if os.path.exists(EXCEL_FILE):
        df_existing = pd.read_excel(EXCEL_FILE)
        df_updated = pd.concat([df_existing, df_row], ignore_index=True)
        df_updated.to_excel(EXCEL_FILE, index=False)
    else:
        df_row.to_excel(EXCEL_FILE, index=False)



@app.route('/submit_answers', methods=['POST'])
def submit_answers():
    data = request.json
    # Process submitted answers (if needed, save or process them)
    return jsonify({"message": "Answers submitted successfully!"})


# Utility Functions
def calculate_required_facilities(num_students, num_divisions, num_courses, course_duration):
    """
    Calculate required facilities based on student population and institutional parameters.

    Args:
    - num_students: Total number of students
    - num_divisions: Number of student divisions
    - num_courses: Number of courses
    - course_duration: Duration of courses

    Returns:
    - Dictionary of required facilities with their quantities
    """
    results = {
        # Classroom Calculation: Based on divisions, course duration, and utilization
        "Classroom model": max(1, int(num_divisions * course_duration * 0.5)),
        # Computer Lab Calculation: Considering courses, student density
        "Computer lab model": max(1, int((num_courses * course_duration + num_students / 400) * 0.75)),
        # Facilities typically singular in a college
        "Audi model": 1,  # One main auditorium
        "TPO model": 1,  # One Training and Placement Office
        "Medical Room Model": 1,  # One central medical room
        "Server Room model": 1,  # One central server room
        "Conference Halls model": 1,  # One main conference hall
        "Seminar hall model": 1,  # One primary seminar hall
        # Facilities with more variable allocation
        "Workshop model": max(1, num_students // 600),
        "Sports equipment model": 1,
        # Canteen Calculation: Scaled with student population
        "Canteen model": 1,
        # Additional facilities with minimum allocation
        "Drawing Halls model": 1,
        "Faculty cabin model": max(1, num_students//20),
        "Fire extinguisher model": max(1, num_divisions)+20,
        "Generator model": 1,
        "Ground model": 1,
        "Library model": 1,  # Typically one main library
        "Parking model": 1,
        "Restroom Model": max(2, num_students // 500),
    }

    return results

def verify_image(image_path, model_path, labels_path):
    try:
        print(f"Loading model from: {model_path}")
        model = load_model(model_path)
    except Exception as e:
        print(f"Error loading model: {e}")
        raise

    try:
        with open(labels_path, 'r') as f:
            labels = [line.strip() for line in f.readlines()]
        print(f"Labels loaded: {labels}")
    except Exception as e:
        print(f"Error loading labels: {e}")
        raise

    try:
        image = Image.open(image_path).convert('RGB')
        image = ImageOps.fit(image, (224, 224), Image.Resampling.LANCZOS)
        image_array = np.asarray(image)
        normalized_image_array = (image_array.astype(np.float32) / 127.5) - 1
        data = np.expand_dims(normalized_image_array, axis=0)

        print("Running prediction...")
        prediction = model.predict(data)
        index = np.argmax(prediction)
        confidence_score = prediction[0][index]

        # Convert numpy.float32 to Python float for JSON serialization
        return {"label": labels[index], "confidence": float(confidence_score)}
    except Exception as e:
        print(f"Error during prediction: {e}")
        raise


def extract_questions(pdf_path):
    """Extracts questions from a given PDF file."""
    if not os.path.exists(pdf_path):
        return []

    questions = []
    try:
        with open(pdf_path, 'rb') as pdf_file:
            reader = PyPDF2.PdfReader(pdf_file)
            for page in reader.pages:
                text = page.extract_text()
                # Extract lines ending with "?" (assuming questions end with "?")
                questions.extend([line.strip() for line in text.split('\n') if line.strip().endswith('?')])
    except Exception as e:
        print(f"Error extracting questions from {pdf_path}: {e}")

    return questions


def process_single_image(file, environment):
    """Helper function to process a single image"""
    if file.filename == '':
        raise ValueError('No selected file')

    # Validate file type
    allowed_extensions = {'png', 'jpg', 'jpeg'}
    if not file.filename.lower().endswith(tuple(allowed_extensions)):
        raise ValueError('Invalid file type. Please upload a PNG or JPEG image.')

    # Read and process image
    image_bytes = file.read()
    img = Image.open(io.BytesIO(image_bytes))
    if img.mode == 'RGBA':
        img = img.convert('RGB')

    # Save image temporarily
    temp_path = f"temp_image_{environment}.jpg"
    img.save(temp_path)

    # Keep a copy for drawing
    img_draw = np.array(img)
    img_draw = cv2.cvtColor(img_draw, cv2.COLOR_RGB2BGR)

    # Perform detection
    results = CLIENTS[environment].infer(temp_path, model_id=MODEL_IDS[environment])
    detections = []

    # Process results
    for i, prediction in enumerate(results.get('predictions', [])):
        x1 = int(prediction['x'] - prediction['width'] / 2)
        y1 = int(prediction['y'] - prediction['height'] / 2)
        x2 = int(prediction['x'] + prediction['width'] / 2)
        y2 = int(prediction['y'] + prediction['height'] / 2)

        class_name = prediction['class']
        confidence = prediction['confidence']

        detections.append({
            'bbox': [x1, y1, x2, y2],
            'class': class_name,
            'confidence': round(confidence, 2),
            'id': f'{environment}-detection-{i}'
        })

        # Draw bounding box
        cv2.rectangle(img_draw, (x1, y1), (x2, y2), (0, 255, 0), 1)
        cv2.putText(img_draw, f'{class_name} {confidence:.2f}', (x1, y1 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

    # Clean up temporary file
    if os.path.exists(temp_path):
        os.remove(temp_path)

    # Convert the image to base64
    _, buffer = cv2.imencode('.jpg', img_draw)
    img_str = base64.b64encode(buffer).decode()

    return {
        'image': f'data:image/jpeg;base64,{img_str}',
        'detections': detections
    }


@app.route('/detect', methods=['POST'])
def detect():
    try:
        required_environments = ['classroom', 'chemical_lab', 'mechanical_workshop', 
                               'computer_lab', 'cctv', 'notice_board', 'bench']
        results = {}

        # Check if all required images are provided
        for env in required_environments:
            if f'image_{env}' not in request.files:
                return jsonify({
                    'success': False,
                    'error': f'No image file provided for {env}'
                }), 400

        # Process each image
        for env in required_environments:
            try:
                file = request.files[f'image_{env}']
                results[env] = process_single_image(file, env)
            except Exception as e:
                return jsonify({
                    'success': False,
                    'error': f'Error processing {env} image: {str(e)}'
                }), 400

        return jsonify({
            'success': True,
            'results': results
        })

    except Exception as e:
        app.logger.error(f"Error in detect route: {str(e)}")
        return jsonify({
            'success': False,
            'error': f'Server error: {str(e)}'
        }), 500

@app.route('/download_report', methods=['GET'])
def download_report():
    # Generate PDF report
    buffer = BytesIO()
    pdf = canvas.Canvas(buffer)

    # Write content to PDF (example content)
    pdf.drawString(100, 800, "Facility Management System Report")
    pdf.drawString(100, 780, "This is an auto-generated report.")

    # Sample table (adjust as per your needs)
    y = 750
    for facility in FACILITIES:
        pdf.drawString(100, y, f"Facility: {facility}")
        y -= 20  # Move to next line

    pdf.save()
    buffer.seek(0)

    return send_file(buffer, as_attachment=True, download_name="facility_report.pdf", mimetype='application/pdf')


@app.route('/download_excel', methods=['GET'])
def download_excel():
    if os.path.exists(EXCEL_FILE):
        return send_file(EXCEL_FILE, as_attachment=True, download_name="facility_data.xlsx",
                         mimetype='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet')
    else:
        return jsonify({"error": "Excel file not found"}), 404


if __name__ == '__main__':
    app.run(debug=True)