# -*- coding: utf-8 -*-
"""MartheDeployment_Doors_fasterRCNN.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1kgEtpfNt0jxSwPRhOzODIC6P_prg-c4L

## Libraries
"""

# from google.colab.patches import cv2_imshow
import cv2
import numpy as np
import pandas as pd

import statistics
from statistics import mode

from PIL import Image

# pip install PyPDF2

# pip install PyMuPDF

# pip install pip install PyMuPDF==1.19.0

import io

# !pip install pypdfium2
import pypdfium2 as pdfium

import fitz  # PyMuPDF

import os

#drive.mount("/content/drive", force_remount=True)

import torch
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from PIL import Image, ImageDraw
import torchvision.transforms.functional as F
import matplotlib.pyplot as plt
import google_sheet_Legend
"""# updated for (fullpath, pdf_name)



"""

def convert2pillow(path):
  pdf = pdfium.PdfDocument(path)
  page = pdf.get_page(0)
  pil_image = page.render().to_pil()
  return pil_image

import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

# Function to get the model
def get_model(num_classes):
    # Load a pre-trained Faster R-CNN model with a ResNet-50-FPN backbone
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

    # Get the number of input features for the classifier
    in_features = model.roi_heads.box_predictor.cls_score.in_features

    # Replace the pre-trained head with a new one for our number of classes
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    return model

'''def ev_model(img, model, device, threshold):
  image_tensor = F.to_tensor(img).unsqueeze(0)
  image_tensor = image_tensor.to(device)
  model.eval()

  with torch.no_grad():
    predictions = model(image_tensor)

  single_boxes = []
  double_boxes = []
  for element in range(len(predictions[0]['boxes'])):
    score = predictions[0]['scores'][element].item()
    if score > threshold:
      if predictions[0]['labels'][element].item() == 1:
        single_boxes.append(predictions[0]['boxes'][element].tolist())
      else:
        double_boxes.append(predictions[0]['boxes'][element].tolist())

  return single_boxes, double_boxes'''

def ev_model(img, model, device, threshold):
  image_tensor = F.to_tensor(img).unsqueeze(0)
  image_tensor = image_tensor.to(device)
  model.eval()

  with torch.no_grad():
    predictions = model(image_tensor)

  doors_info = []
  for element in range(len(predictions[0]['boxes'])):
    score = predictions[0]['scores'][element].item()
    if score > threshold:
        box = predictions[0]['boxes'][element].tolist()
        label = predictions[0]['labels'][element].item()
        doors_info.append((box,label))
  return doors_info


def distance(point1, point2):
    x1, y1 = point1
    x2, y2 = point2
    return np.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2)

def calculate_midpoint(p1, p2):
    x1, y1 = p1
    x2, y2 = p2
    # Calculate the midpoint
    xm = int((x1 + x2) / 2)
    ym = int((y1 + y2) / 2)
    return (xm, ym)
def get_door_info(doors_info):
  width_pixels = []
  lines = []
  sanda = []
  line_midpoint = []
  for door_inf in doors_info:
    xmin, ymin, xmax, ymax = door_inf[0]
    #horz_bottom
    if door_inf[1] == 2:
      #for drawing
      point_st = (int(xmin), int(ymax) + 5)
      point_end = (int(xmax),int(ymax) + 5)
      lines.append((point_st,point_end))

      sanda_st = (int(xmin), int(ymax))
      sand_end = (int(xmax),int(ymax))
      sanda.append((sanda_st, sand_end))

      line_midpoint.append(calculate_midpoint(sanda_st,sand_end))
      #for calculation
      width = distance((xmin,ymax), (xmax,ymax))
      width_pixels.append(width)

    #horz_upper
    if door_inf[1] == 3:
      #for drawing
      point_st = (int(xmin),int(ymin) -5)
      point_end = (int(xmax),int(ymin) - 5)
      lines.append((point_st,point_end))

      sanda_st = (int(xmin),int(ymin))
      sand_end = (int(xmax),int(ymin))
      sanda.append((sanda_st, sand_end))
      line_midpoint.append(calculate_midpoint(sanda_st,sand_end))
      #for calculation
      width = distance((xmin,ymin), (xmax,ymin))
      width_pixels.append(width)

    #vert_right
    if door_inf[1] == 4:
      #for drawing
      point_st = (int(xmax) + 5,int(ymin))
      point_end = (int(xmax) + 5,int(ymax))
      lines.append((point_st,point_end))

      sanda_st = (int(xmax), int(ymin))
      sand_end = (int(xmax), int(ymax))
      sanda.append((sanda_st, sand_end))
      line_midpoint.append(calculate_midpoint(sanda_st,sand_end))
      #for calculation
      width = distance((xmax,ymin), (xmax,ymax))
      width_pixels.append(width)
    #vert_left
    if door_inf[1] == 5:
      #for drawing
      point_st = (int(xmin) -5,int(ymin))
      point_end = (int(xmin) -5,int(ymax))
      lines.append((point_st,point_end))

      sanda_st = (int(xmin),int(ymin))
      sand_end = (int(xmin),int(ymax))
      sanda.append((sanda_st, sand_end))
      line_midpoint.append(calculate_midpoint(sanda_st,sand_end))
      #for calculation
      width = distance((xmin,ymin), (xmin,ymax))
      width_pixels.append(width)

  return width_pixels, lines, sanda, line_midpoint

def pxl2meter(width_pixels, ratio):
  real_width = []
  for width in width_pixels:
    real_width.append(round(width*ratio, 2))
  return real_width

def width_as_char(real_width):
  char_width = []
  for width in real_width:
    char_width.append(f"{width}m")
  return char_width

def add_annotations_to_pdf(image, pdf_name, lines, sanda, char_width, line_midpoint):
    image_width, image_height = image.size

    # Create a new PDF document
    pdf_document = fitz.open()

    # Add a new page to the document with the same dimensions as the image
    page = pdf_document.new_page(width=image_width, height=image_height)

    # Insert the image into the PDF page
    image_stream = io.BytesIO()
    image.save(image_stream, format="PNG")
    page.insert_image(page.rect, stream=image_stream.getvalue())


    #Annotation for drawin lines as in the markups
    for i in range(len(line_midpoint)):
      x, y = line_midpoint[i]
      text = char_width[i]
      rect = fitz.Rect(x, y, x + 200, y + 50)  # Adjust the width and height as needed
      annot = page.add_freetext_annot(rect, text, fontsize=10, fontname="helv", text_color=(1,0,0))
      annot.update()
    #Annotation for drawin lines as in the markups
    for i in range(len(lines)):
        annot = page.add_line_annot(lines[i][0], lines[i][1])
        annot = page.add_line_annot(sanda[i][0], lines[i][0])
        annot = page.add_line_annot(sanda[i][1], lines[i][1])
        annot.set_border(width=2, dashes=None)  # Optional border styling
        annot.set_colors(stroke=(1, 0, 0))  # Set the line color to red
        annot.update()


    output_pdf_path = pdf_name+"_annotated.pdf"
    # Save the PDF with annotations

    return pdf_document
    


'''def calculate_width(bbox):
  #if looking right or left, width < height
  bbox_width = bbox[2] - bbox[0]
  bbox_height = bbox[3] - bbox[1]
  if bbox_width > bbox_height:
    door_width = bbox_width
  else:
    door_width = bbox_height
  return door_width
'''
'''def width_annotations(bbox, ratio):
  lines = []
  width = []

  for box in bbox:
    door_width = calculate_width(box)
    door_width = round(door_width*ratio, 2)
    x1,y1,x2,y2 = box
    b_width = x2 - x1
    b_height = y2 - y1

    if b_width > b_height:
      lines.append(((x1, y1), (x2, y1)))
      x = (x1+x2)/2
      y = (y1+y1)/2
      width.append(((x,y),door_width))
    else:
      lines.append(((x1, y1), (x1, y2)))
      x = (x1+x1)/2
      y = (y1+y2)/2
      width.append(((x,y), door_width))
  return lines, width'''

'''def create_width_annotations(width_info):
  annotations = []
  for i in range(len(width_info)):
    annotations.append(((width_info[i][0][0]),(width_info[i][0][1]),f"{width_info[i][1]} m"))
  return annotations'''

'''def calculate_midpoint(top_left, bottom_right):
    x1, y1 = top_left
    x2, y2 = bottom_right
    # Calculate the midpoint
    xm = int((x1 + x2) / 2)
    ym = int((y1 + y2) / 2)
    return (xm, ym)'''

'''def mid_points_bbox(bbox):
  midpoint = []
  for i in range(len(bbox)):
    x1 = int(bbox[i][0])
    y1 = int(bbox[i][1])
    x2 = int(bbox[i][2])
    y2 = int(bbox[i][3])
    top_left_corner = (x1, y1)
    bottom_right_corner = (x2, y2)
    midpoint.append(calculate_midpoint(top_left_corner, bottom_right_corner))
  return midpoint'''

'''def create_annotations(door_kind, midpoints):
  door = door_kind
  annotations = []
  for i in range(len(midpoints)):
    annotations.append((midpoints[i][0],midpoints[i][1], door))
  return annotations

def add_annotations_to_pdf(image, pdf_name, annotation_s, annotation_d,width_ann_single, width_ann_double,line_single,line_double):
    image_width, image_height = image.size

    # Create a new PDF document
    pdf_document = fitz.open()

    # Add a new page to the document with the same dimensions as the image
    page = pdf_document.new_page(width=image_width, height=image_height)

    # Insert the image into the PDF page
    image_stream = io.BytesIO()
    image.save(image_stream, format="PNG")
    page.insert_image(page.rect, stream=image_stream.getvalue())

    # Add annotations
    for annotation in annotation_s:
        x, y, text = annotation
        # Create an annotation (sticky note)
        annot = page.add_text_annot(fitz.Point(x, y), text)
        annot.set_border(width=0.2, dashes=(1, 2))  # Optional border styling
        annot.set_colors(stroke=(1, 0, 0), fill=None)  # Set the stroke color to red
        annot.update()
    for annotation in annotation_d:
        x, y, text = annotation
        # Create an annotation (sticky note)
        annot = page.add_text_annot(fitz.Point(x, y), text)
        annot.set_border(width=0.2, dashes=(1, 2))  # Optional border styling
        annot.set_colors(stroke=(0, 1, 0), fill=None)  # Set the stroke color to red
        annot.update()

    #Annotations for width measurement (marra single we marra double)
    for annotation in width_ann_single:
        x, y, text = annotation
        rect = fitz.Rect(x, y, x + 200, y + 50)  # Adjust the width and height as needed
        annot = page.add_freetext_annot(rect, text, fontsize=10, fontname="helv", text_color=(1,0,0))
        annot.update()
    for annotation in width_ann_double:
        x, y, text = annotation
        rect = fitz.Rect(x, y, x + 200, y + 50)  # Adjust the width and height as needed
        annot = page.add_freetext_annot(rect, text, fontsize=10, fontname="helv", text_color=(1,0,0))
        annot.update()

    #Annotation kind of the line drawings (marra single we marra double)
    for line in line_single:
        start_point, end_point = line
        annot = page.add_line_annot(start_point, end_point)
        annot.set_border(width=2, dashes=None)  # Optional border styling
        annot.set_colors(stroke=(1, 0, 0))  # Set the line color to red
        annot.update()
    for line in line_double:
        start_point, end_point = line
        annot = page.add_line_annot(start_point, end_point)
        annot.set_border(width=2, dashes=None)  # Optional border styling
        annot.set_colors(stroke=(1, 0, 0))  # Set the line color to red
        annot.update()


    output_pdf_path = pdf_name+"_annotated.pdf"
    # Save the PDF with annotations

    return pdf_document
    # pdf_document.save(output_pdf_path)
    # pdf_document.close()
'''
def main_run(pdf_fullpath, weights_path, pdf_name,pdfpath,ratio): ####pdf_fullpath here is the data and not the path 
    img_pillow = convert2pillow(pdf_fullpath)
    #new_image = img_pillow.resize((2384, 1684))
    
    # Specify the number of classes (including the background)
    num_classes = 10  # Ensure this matches the saved model's number of classes
    
    # Load the model with the specified number of classes
    model = get_model(num_classes)
    
    # Load the saved model's state dictionary with map_location to handle CPU
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    
    try:
        model.load_state_dict(torch.load(weights_path, map_location=device), strict=False)
    except RuntimeError as e:
        print(f"Error loading model state_dict: {e}")
        return
    
    # Set the model to evaluation mode
    model.eval()
    
    # Move the model to the appropriate device
    model.to(device)
    
    # START INFERENCE
    #sbox, dbox = ev_model(img_pillow, model, device, 0.6)
    
    #Dataframe for Doors count
    #doors_count = {'Type': ['Single Doors', 'Double Doors'], 'Quantity': [len(sbox), len(dbox)]}
    #df_doors = pd.DataFrame(doors_count)

    #single_midpoint = mid_points_bbox(sbox)
    #double_midpoint = mid_points_bbox(dbox)

    #Kind Annotations
    #single_annotations = create_annotations("single door", single_midpoint)
    #double_annotations = create_annotations("double door", double_midpoint)
    #Lines Annotations
    #line_single, width_signle = width_annotations(sbox, 1)
    #line_double, width_double = width_annotations(dbox, 1)
    #Width Annotations
    #width_single_ann = create_width_annotations(width_signle)
    #width_double_ann = create_width_annotations(width_double)

    # add_annotations_to_pdf(new_image, pdf_name, single_annotations, double_annotations)


    #NEW
    doors_info = ev_model(img_pillow, model, device, 0.9)
    width_pixels, lines, sanda, line_midpoint = get_door_info(doors_info)
    real_width = pxl2meter(width_pixels, ratio)
    char_width = width_as_char(real_width)

    pdf_document = add_annotations_to_pdf(img_pillow, plan, lines, sanda, char_width, line_midpoint)
    
    #pdf_document=add_annotations_to_pdf(img_pillow, pdf_name, single_annotations, double_annotations,width_single_ann,width_double_ann,line_single,line_double)

    page=pdf_document[0]
    pix = page.get_pixmap()  # render page to an image
    pl=Image.frombytes('RGB', [pix.width,pix.height],pix.samples)
    img=np.array(pl)
    annotatedimg = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

    df_doors = df_doors.fillna(' ')
    gc,spreadsheet_service,spreadsheetId, spreadsheet_url , namepathArr=google_sheet_Legend.legendGoogleSheets(df_doors , pdf_name,pdfpath)
    list1=pd.DataFrame(columns=['content',  'id',  'subject'])
    for page in pdf_document:
        for annot in page.annots():
            list1.loc[len(list1)] =annot.info


    # modify this return
    return annotatedimg, pdf_document , spreadsheet_url, list1 , df_doors

# model_path =  '/content/drive/MyDrive/separated_classes.pth'
# #pdf_name = data
# for i in range(len(fullpath)):
#   main_run(fullpath[i], model_path, pdf_name[i])