import cv2
from PIL import Image
from ultralyticsplus import YOLO
from transformers import pipeline
import pandas as pd
import numpy as np
import easyocr
from utils import *

INVOICE = ["Numéro de facture", "Date", "Numéro de commande", "Echéance", "Total"]

model = YOLO('keremberke/yolov8s-table-extraction')
model.overrides['conf'] = 0.25  # NMS confidence threshold
model.overrides['iou'] = 0.45  # NMS IoU threshold
model.overrides['agnostic_nms'] = False  # NMS class-agnostic
model.overrides['max_det'] = 1000  # maximum number of detections per image

pipe = pipeline("object-detection", model="bilguun/table-transformer-structure-recognition")


def detect_tables(image):
  # image is an np array
  results = model.predict(image)
  

  result = results[0]
  xyxy = result.boxes.xyxy
  scores = result.boxes.conf
  tables = []
  for i in range(len(scores)):
    if scores[i] >= 0.5:
      table = image[int(xyxy[i,1]):int(xyxy[i,3]), int(xyxy[i,0]):int(xyxy[i,2])]
      table = Image.fromarray(table)
      tables.append(table)
  return tables

def insert(el, listt, pos):
  if not listt:
      listt.append(el)
  else:
      inserted = False
      for i in range(len(listt)):
          if el[pos] <= listt[i][pos]:
              listt.insert(i, el)
              inserted = True
              break
      if not inserted:
          listt.append(el)

def rec_table(table, reader):
  col_row = pipe(table)
  cols = []
  rows = []
  for el in col_row:
    if el["label"] == 'table column':
      insert(el["box"], cols, pos = "xmin")
    elif el["label"] == 'table row':
      insert(el["box"], rows, pos = "ymin")

  table = np.array(table)

  csv = []
  for row in rows:
    temp = []
    for col in cols:
      box = intersection(row, col)
      cell = table[box['ymin']:box['ymax'], box['xmin']:box['xmax']]
      res = get_ocr(cell,reader)
      temp.append(get_input(res))
    csv.append(temp)
  
  df = pd.DataFrame(csv)
  return df


def intersection(box1, box2):
    # Extract coordinates of first bounding box
    x1min, y1min, x1max, y1max = box1['xmin'], box1['ymin'], box1['xmax'], box1['ymax']

    # Extract coordinates of second bounding box
    x2min, y2min, x2max, y2max = box2['xmin'], box2['ymin'], box2['xmax'], box2['ymax']

    # Calculate coordinates of intersection
    xmin = max(x1min, x2min)
    ymin = max(y1min, y2min)
    xmax = min(x1max, x2max)
    ymax = min(y1max, y2max)

    # Check if there is no intersection
    if xmin >= xmax or ymin >= ymax:
        return None

    # Return the coordinates of the intersection
    return {'xmin': xmin, 'ymin': ymin, 'xmax': xmax, 'ymax': ymax}


# def extract_tables(lang, image):
#   reader = easyocr.Reader([langs[lang]])
#   tables = detect_tables(image)

#   for i in range(len(tables)):
#     df = rec_table(tables[i], reader)
#     df.to_excel(f'table_{i+1}.xlsx', index=False, header=False)

def extract_tables(lang, image):
  reader = easyocr.Reader([langs[lang]])
  tables = detect_tables(image)
  csvs = []
  for i in range(len(tables)):
    df = rec_table(tables[i], reader)
    csv = df.to_csv(index=False, header=False)
    csvs.append(csvs)
  
  return csvs[0]

if __name__ == '__main__':
    lang = "french"
    to_be_extracted = INVOICE
    image_path = "./docs for ocr/invoices/facture.png"
    image = cv2.imread(image_path)
    print(image.shape)
    
    text_data = extract_data(lang, to_be_extracted, image)
    print(text_data)

    # extract_tables(lang, image) # extract tables from the image and download them in excel format to the current directory