import os
import csv
import torch
import torchvision
import easyocr
import shutil
import random
import cv2
from glob import glob
from ultralytics import YOLOv10
import random
from glob import glob
from ultralytics import YOLOv10
import supervision as sva
from ultralytics import YOLOv10
import supervision as sv
import supervision as sv
from flask import Flask, request, jsonify, send_from_directory, render_template
import pytesseract
# Set the tesseract_cmd based on the operating system
if os.name == 'nt':  # For Windows
    pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
elif os.name == 'posix':  # For Linux/MacOS
    pytesseract.pytesseract.tesseract_cmd = '/usr/bin/tesseract'  # Common location for Tesseract on Linux/MacOS

import textwrap
app = Flask(__name__)

def enhance_contrast(image):
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    equalized_image = cv2.equalizeHist(gray_image)
    return equalized_image


def calculate_iou(bbox1, bbox2):
    x1_max = max(bbox1[0], bbox2[0])
    y1_max = max(bbox1[1], bbox2[1])
    x2_min = min(bbox1[2], bbox2[2])
    y2_min = min(bbox1[3], bbox2[3])

    inter_area = max(0, x2_min - x1_max) * max(0, y2_min - y1_max)

    bbox1_area = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
    bbox2_area = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])

    iou = inter_area / float(bbox1_area + bbox2_area - inter_area) if (bbox1_area + bbox2_area - inter_area) > 0 else 0
    return iou


def recreate_directories():
    # Paths for the directories
    cropped_dir = "./app/cropped_images/"
    output_dir1 = "./app/Folder1"
    output_dir2 = "./app/Folder2"
    output_dir3 = "./app/Folder3"
    UPLOAD_FOLDER = "./app/data1"

    # Remove existing directories
    if os.path.exists(cropped_dir):
        shutil.rmtree(cropped_dir)
    if os.path.exists(output_dir1):
        shutil.rmtree(output_dir1)
    if os.path.exists(output_dir2):
        shutil.rmtree(output_dir2)
    if os.path.exists(output_dir3):
        shutil.rmtree(output_dir3)
    if os.path.exists(UPLOAD_FOLDER):
        shutil.rmtree(UPLOAD_FOLDER)

    # Recreate directories
    os.makedirs(cropped_dir, exist_ok=True)
    os.makedirs(output_dir1, exist_ok=True)
    os.makedirs(output_dir2, exist_ok=True)
    os.makedirs(output_dir3, exist_ok=True)
    os.makedirs(UPLOAD_FOLDER, exist_ok=True)

@app.route('/')
def index():
    return render_template('index3.html')  # This will serve your HTML page

@app.route('/upload', methods=['POST'])
def upload_file():
    recreate_directories()
    if 'invoice-upload' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    file = request.files['invoice-upload']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400
    if file:
        file_path = os.path.join('./app/data1', file.filename)
        file.save(file_path)
        output_image, output_csv,image_path = process_image()

        return jsonify({
            'image_path': output_image,
            'csv_path': output_csv,
            'original_image': image_path
        })
def process_image():
    print("Current working directory:", os.getcwd())
    
    # Check contents in the root directory
    print("Current directory contents:", os.listdir('/'))

    model = YOLOv10(f'./runs/detect/train3/weights/best (1).pt')
    dataset = sv.DetectionDataset.from_yolo(
        images_directory_path=f"./data/MyNewVersion5.0Dataset/valid/images",
        annotations_directory_path=f"./data/MyNewVersion5.0Dataset/valid/labels",
        data_yaml_path=f"./data/MyNewVersion5.0Dataset/data.yaml"
    )
    bounding_box_annotator = sv.BoundingBoxAnnotator()
    label_annotator = sv.LabelAnnotator()
    image_dir = "./app/data1"
    files = os.listdir('./app/data1')
    files.sort()
    files = files[0:100]
    print(files)
    counter = 0
    for ii in files:
        random_image_data = cv2.imread('./app/data1/' + ii)
        random_image_data1 = cv2.imread('./app/data1/' + ii)
        results = model(source='./app/data1/' + ii, conf=0.07)[0]
        detections = sv.Detections.from_ultralytics(results)
        annotated_image = bounding_box_annotator.annotate(scene=random_image_data, detections=detections)
        annotated_image = label_annotator.annotate(scene=annotated_image, detections=detections)
        save_path = "./app/Folder1/" + "detection" + ii
        cv2.imwrite(save_path, annotated_image)
        print(f"Annotated image saved at {save_path}")
        bounding_boxes = results.boxes.xyxy.cpu().numpy()
        class_ids = results.boxes.cls.cpu().numpy()
        confidences = results.boxes.conf.cpu().numpy()
        bounding_box_save_path = "./bounding_boxes.txt"
        with open(bounding_box_save_path, 'w') as f:
            for i, (bbox, class_id, confidence) in enumerate(zip(bounding_boxes, class_ids, confidences)):
                x1, y1, x2, y2 = map(int, bbox)
                f.write(f"Object {i + 1}: Class {class_id}, Confidence: {confidence:.2f}, "
                        f"Bounding box: ({x1}, {y1}, {x2}, {y2})\n")
                cropped_image = random_image_data1[y1:y2, x1:x2]
                cropped_image_path = os.path.join('./app/cropped_images/', f"cropped_object_{i + 1}.jpg")
                cv2.imwrite(cropped_image_path, cropped_image)
                print(f"Enhanced cropped image saved at {cropped_image_path}")
        print(f"Checking contents of /app/data: {bounding_box_save_path}")
        print(f"Directory listing: {os.listdir('./app/Folder1')}")
        print(f"Bounding box coordinates saved at {bounding_box_save_path}")
        import re
        input_file_path = './bounding_boxes.txt'
        cropped_images_folder = './app/cropped_images/'
        output_csv_path = './app/Folder2/' + ii + 'bounding_boxes_with_recognition.csv'
        with open(input_file_path, 'r') as infile:
            lines = infile.readlines()
        with open(output_csv_path, 'w', newline='', encoding='utf-8') as csvfile:
            csv_writer = csv.writer(csvfile)
            csv_writer.writerow(['Object ID', 'Bounding Box', 'Image Name', 'Recognized Text'])
            for i, line in enumerate(lines):
                object_id = f"Object_{i + 1}"
                bounding_box_info = line.strip()
                cropped_image_name = f"cropped_object_{i + 1}.jpg"
                cropped_image_path = os.path.join(cropped_images_folder, cropped_image_name)
                if os.path.exists(cropped_image_path):
                    bbox_match = re.search(r"Bounding box: \((\d+), (\d+), (\d+), (\d+)\)", bounding_box_info)
                    if bbox_match:
                        x1, y1, x2, y2 = map(int, bbox_match.groups())
                        detected_boxes = [[x1, x2, y1, y2]]
                    else:
                        print("No bounding box found in the info.")
                    cropped_image = cv2.imread(cropped_image_path, cv2.IMREAD_GRAYSCALE)
                    recognized_text = pytesseract.image_to_string(cropped_image,config="--psm 6 -c tessedit_char_whitelist=0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz-")
                    print(f"Recognized Text: {recognized_text}")
                   
                csv_writer.writerow([object_id, bounding_box_info, cropped_image_name, recognized_text])
        print(f"CSV file with recognition results saved at {output_csv_path}")

        def calculate_iou(bbox1, bbox2):
            x1_max = max(bbox1[0], bbox2[0])
            y1_max = max(bbox1[1], bbox2[1])
            x2_min = min(bbox1[2], bbox2[2])
            y2_min = min(bbox1[3], bbox2[3])

            inter_area = max(0, x2_min - x1_max) * max(0, y2_min - y1_max)

            bbox1_area = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
            bbox2_area = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])

            iou = inter_area / float(bbox1_area + bbox2_area - inter_area) if (bbox1_area + bbox2_area - inter_area) > 0 else 0
            return iou

        image_path = "./app/data1/" + ii
        csv_file_path = output_csv_path = './app/Folder2/' + ii + 'bounding_boxes_with_recognition.csv'
        image = cv2.imread(image_path)
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 1.3
        font_thickness = 2
        color = (255, 0, 255)
        bboxes = []
        recognized_texts = []
        with open(csv_file_path, 'r', encoding='utf-8') as csvfile:
            csv_reader = csv.DictReader(csvfile)
            for row in csv_reader:
                bbox_match = re.search(r'\((\d+), (\d+), (\d+), (\d+)\)', row['Bounding Box'])
                if bbox_match:
                    bbox = [int(bbox_match.group(i)) for i in range(1, 5)]
                    bboxes.append(bbox)
                    recognized_texts.append(row['Recognized Text'])
        filtered_bboxes = []
        filtered_texts = []
        iou_threshold = 0.4
        for i, bbox1 in enumerate(bboxes):
            keep = True
            for j, bbox2 in enumerate(filtered_bboxes):
                if calculate_iou(bbox1, bbox2) > iou_threshold:
                    keep = False
                    break
            if keep:
                filtered_bboxes.append(bbox1)
                filtered_texts.append(recognized_texts[i])
        for bbox, recognized_text in zip(filtered_bboxes, filtered_texts):
            x1, y1, x2, y2 = bbox
            cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
            max_chars_per_line = 60
            wrapped_text = textwrap.wrap(recognized_text, width=max_chars_per_line)
            text_y = y1 - 10 if y1 - 10 > 10 else y1 + 10
            for line in wrapped_text:
                cv2.putText(image, line, (x1, text_y), font, font_scale, color, font_thickness)
                text_y += int(font_scale * 20)
            output_image_path = "./app/Folder3/" + "annotated" + ii + ".png"
            cv2.imwrite(output_image_path, image)
            print(f"Annotated image saved at {output_image_path}")
            counter += 1
        return output_image_path, output_csv_path,image_path

@app.route('/download_csv/<filename>')
def download_csv(filename):
    return send_from_directory('./app/Folder2', filename, as_attachment=True)

@app.route('/download_image/<filename>')
def download_image(filename):
    return send_from_directory('./app/Folder3', filename, as_attachment=True)

@app.route('/uploads/<filename>')
def serve_uploaded_file(filename):
    return send_from_directory('./app/data1', filename)