import numpy as np
import cv2, csv, os

def mask_box(mask_img, points, color):
    mask = np.ones_like(mask_img, dtype=np.uint8) * 255
    cv2.fillPoly(mask, [points], color)

    # Apply the mask to the image
    img_with_overlay = cv2.bitwise_and(mask_img, mask)

    return img_with_overlay

def mask_frame(mask_img, cl, tables, color):
    # Blend the original image with the specified color
    blend = 0.5
    img_color = cv2.addWeighted(mask_img, blend, np.full_like(mask_img, color), 1 - blend, 0)

    # Create a mask with a white rectangle in the specified region
    mask = np.zeros_like(mask_img, dtype=np.uint8)
    x1, y1, x2, y2 = cl.x, cl.y, cl.x + cl.w, cl.y + cl.h
    cv2.rectangle(mask, (x1, y1), (x2, y2), (255, 255, 255), -1)
    for table in tables:
        for tab in table:
            pts = np.array([(tab.x, tab.y), (tab.x+tab.w, tab.y), (tab.x+tab.w, tab.y+tab.h),(tab.x,tab.y+tab.h)], np.int32)
            cv2.fillPoly(mask, [pts], (255, 255, 255))

    # Apply the mask to choose between the blended image and the original
    result = np.where(mask == 255, mask_img, img_color)

    return result

def mask_img(img, gdt_boxes, tables, dimensions, frame, other_info):
    mask_img=img.copy()
    for table in tables:
        for tab in table:
            pts = np.array([(tab.x, tab.y), (tab.x+tab.w, tab.y), (tab.x+tab.w, tab.y+tab.h),(tab.x,tab.y+tab.h)], np.int32)
            mask_img = mask_box(mask_img, pts, (212,242,247))
    
    for gdt in gdt_boxes:
        for g in gdt.values():
            for tab in g:
                pts = np.array([(tab.x, tab.y), (tab.x+tab.w, tab.y), (tab.x+tab.w, tab.y+tab.h),(tab.x,tab.y+tab.h)], np.int32)
                mask_img = mask_box(mask_img, pts, (68,136,179))

    if frame:
        mask_img = mask_frame(mask_img, frame, tables, (100*2, 78*2, 73*2))
        offset = (frame.x, frame.y)
    else:
        offset = (0, 0)

    for dim in dimensions:
        box = dim[1]
        pts=np.array([(box[0]+offset),(box[1]+offset),(box[2]+offset),(box[3]+offset)])
        mask_img = mask_box(mask_img, pts, (110,185,187))
    
    for info in other_info:
        box = info[1]
        pts=np.array([(box[0]+offset),(box[1]+offset),(box[2]+offset),(box[3]+offset)])
        mask_img = mask_box(mask_img, pts, (128,102,90))

   
    return mask_img

def process_raw_output(output_path, table_results = None, gdt_results = None, dimension_results = None, other_info = None, save = False):
    #Write Table Results
    if table_results and save:
        csv_file = os.path.join(output_path, 'table_results.csv')
        new_table_results =[]
        with open(csv_file, mode='w', newline='') as file:
            writer = csv.writer(file)
            # Write the header    
            for t in range(len(table_results)):
                writer.writerow([f'Table_{t}'])
                writer.writerow(["Text", "X Coordinate", "Y Coordinate"])
                tab_results = []
                for item in table_results[t]:
                    line = [item['text'], (item['left'], item['top'])]
                    tab_results.append(line)
                    writer.writerow([item['text'], item['left'], item['top']])
                new_table_results.append(tab_results)
            table_results = new_table_results

    #Write GD&T Results
    if gdt_results and save:
        csv_file = os.path.join(output_path, 'gdt_results.csv')
        with open(csv_file, mode='w', newline='') as file:
            writer = csv.writer(file)
            # Write the header
            writer.writerow(["Text", "X Coordinate", "Y Coordinate"])
            # Write the data
            for item in gdt_results:
                text, coords = item
                writer.writerow([text, coords[0], coords[1]])
                
    #Write Dimension Results
    if dimension_results:
        new_dim_results = []
        for item in dimension_results:
            text, coords = item
            center = np.mean(coords, axis=0).astype(int).tolist()
            new_dim_results.append([text, center])
        dimension_results = new_dim_results
        if save:
            csv_file = os.path.join(output_path, 'dimension_results.csv')
            with open(csv_file, mode='w', newline='') as file:
                writer = csv.writer(file)
                # Write the header
                writer.writerow(["Text", "X Coordinate", "Y Coordinate"])
                # Write the data
                for i in new_dim_results:
                    writer.writerow([i[0], i[1][0], i[1][1]])
    
    if other_info:
        new_info_results = []
        for item in other_info:
            text, coords = item
            center = np.mean(coords, axis=0).astype(int).tolist()
            new_info_results.append([text, center])
        other_info = new_info_results
        if save:
            csv_file = os.path.join(output_path, 'other_info.csv')
            with open(csv_file, mode='w', newline='') as file:
                writer = csv.writer(file)
                # Write the header
                writer.writerow(["Text", "X Coordinate", "Y Coordinate"])
                # Write the data
                for i in new_info_results:
                    writer.writerow([i[0], i[1][0], i[1][1]])

    return table_results, gdt_results, dimension_results, other_info