import cv2
import numpy as np
import math
ans_guide = [None, "A", "B", "C", "D", "E"]


def round_off(number):
    DP = number - int(number)
    if DP > 0.99:
        return int(number) +1
    else:
        return int(number)
def rotate_image_bound(image, angle):
    (h, w) = image.shape[:2]
    center = (w / 2, h / 2)

    # Get rotation matrix
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    # Calculate the new bounding dimensions of the image
    cos = abs(M[0, 0])
    sin = abs(M[0, 1])

    # Compute new dimensions
    new_w = int((h * sin) + (w * cos))
    new_h = int((h * cos) + (w * sin))

    # Adjust the rotation matrix to take into account translation
    M[0, 2] += (new_w / 2) - center[0]
    M[1, 2] += (new_h / 2) - center[1]

    # Perform the rotation
    return cv2.warpAffine(image, M, (new_w, new_h))

def getAnswer(roi):
    hh, ww, ss = roi.shape
    # Convert the image to grayscale
    gray_image = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)

    # Apply Gaussian blur to reduce noise and improve edge detection
    blurred_image = cv2.GaussianBlur(gray_image, (5, 5), 1.4)

    # Perform Canny edge detection
    edges = cv2.Canny(blurred_image, 100, 200)
    contours1, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    bboxes = []
    for c1 in contours1:
        x, y, w, h = cv2.boundingRect(c1)
        bboxes.append((x, y, w, h))
    ans = []
    bboxes.sort(key=sort_criteria_mc)
    for idx, (x, y, w, h) in enumerate(bboxes, start=1):
        answer = ans_guide[int(((x + x+w)/(2*ww))*6)]
        number = round_off(((y + y+h)/(2*hh))*6)
        cv2.rectangle(roi, (x, y), (x + w, y + h), (0, 255, 0), 2)
        # cv2.putText(roi, str(answer), (x + 5, y + 20),
        #             cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        ans.append([number, answer])
    required_numbers = {0, 1, 2, 3, 4, 5}

    # Create a dictionary of the current numbers
    sublist_dict = {item[0]: item[1] for item in ans}
    
    # Fill in missing numbers with empty string
    for num in required_numbers:
        if num not in sublist_dict:
            sublist_dict[num] = ''
    
    # Rebuild the sublist with filled values
    filled_sublist = [[num, sublist_dict[num]] for num in sorted(sublist_dict.keys())]

    print(ans)
    cv2.imshow("roi", roi)
    cv2.waitKey(0)
    return filled_sublist


def sort_criteria_mc(box):
    x, y, w, h = box
    return (round(y / 10) * 10, x)


def sort_criteria(box):
    x, y, w, h = box
    return (round(x / 10) * 10, y)


def OMR_Checker(img_filename:str, ans_key:list):
    """
    Inputs:
        img_filename (str): Path or filename of the image containing the filled OMR sheet.
        ans_key (list): List of correct answers (e.g., ['A', 'C', 'B', 'D', ...]).

    Returns:
        answer_from_sheet (list): List of detected answers from the OMR sheet.
        score (int): Number of correct answers based on comparison with the answer key.
    """
    
    img = cv2.imread(img_filename)
    
    img = rotate_image_bound(img, 90)
    height, width, size = img.shape
    img = cv2.resize(img, (int(width/3), int(height/3)))
    clone = img.copy()
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    # Fixed H and V ranges
    lower_h, upper_h = 89, 130
    lower_v, upper_v = 0, 255

    # Thresholds
    min_area = 14000
    max_area = 25000
    min_w, max_w = 100, 200

    target_boxes = 13
    found = False

    answer_from_sheet = []

    for s in range(0, 256):
        lower = np.array([lower_h, s, lower_v])
        upper = np.array([upper_h, 255, upper_v])

        mask = cv2.inRange(hsv, lower, upper)
        kernel = np.ones((5, 5), np.uint8)
        mask = cv2.dilate(mask, kernel, iterations=4)
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        bounding_boxes = []
        for cnt in contours:
            area = cv2.contourArea(cnt)
            x, y, w, h = cv2.boundingRect(cnt)
            if min_area < area < max_area and min_w < w < max_w and min_w < h < max_w:
                bounding_boxes.append((x, y, w, h))

        if len(bounding_boxes) == target_boxes:
           

            bounding_boxes.sort(key=sort_criteria)

            temp_img = img.copy()
            for idx, (x, y, w, h) in enumerate(bounding_boxes, start=1):
                cv2.rectangle(temp_img, (x, y), (x + w, y + h), (0, 255, 0), 2)
                cv2.putText(temp_img, str(idx), (x + 5, y + 20),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
                if idx > 1:
                    anss = getAnswer(img[y:y+h,x:x+w])
                    for a in anss:
                        answer_from_sheet.append(a[1])

            score = sum(1 for a, b in zip(ans_key, answer_from_sheet) if a == b)
            return answer_from_sheet, score


    return None,None
if __name__ == "__main__":
    import random

    letters = [random.choice(['A', 'B', 'C', 'D', 'E']) for _ in range(72)]
    letters = ['A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A'
               ,'A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A','A']
    letters = ['B', 'C', 'C', 'C', 'B', 'D',
               'A', 'B', 'C', 'D', 'D', 'E',
               'A', 'B', 'C', 'D', 'B', 'A',
               'B', 'A', 'B', 'C', 'D', 'E',
               'B', 'B', 'B', 'B', 'D', 'E',
               'A', 'A', 'B', 'C', 'D', 'E',
               'A', 'A', 'B', 'B', 'C', 'C',
               'A', 'B', 'B', 'A', 'A', 'B',
               'E', 'E', 'D', 'C', 'B', 'A',
               'A', 'A', 'B', 'C', 'C', 'D',
               'A', 'B', 'B', 'A', 'A', 'C',
               'B', 'B', 'C', 'B', 'A', 'A']
    filename = "t1.jpg"
    print(OMR_Checker(filename,letters))