import torch
import numpy as np
import cv2
import matplotlib.pyplot as plt
import os
import logging

import pandas as pd
import time
from PIL import Image

from tqdm import tqdm
from ultralytics import YOLO

def yolo_segmentation_to_masks(segmentation_list, image_size):
    masks = []

    for segmentation in segmentation_list:
        mask = np.zeros(image_size, dtype=np.uint8)
        segmentation = segmentation * np.array(image_size[::-1])  # Rescale the coordinates to image size
        
        # Fill the region defined by the coordinates with 1 (foreground)
        cv2.fillPoly(mask, [segmentation.astype(int)], 1)

        masks.append(mask)

    return masks

model = YOLO("./weights/segment_water/weights/best.pt")
results = []

TXT_DIR = './dataset/water/labels/test'
IMG_DIR = './dataset/water/images/test'


pred_image_folder = './dataset/water/pred_images'
gt_image_folder = './dataset/water/gt_images'
seg_threshold = 0.5
n_class = 1

average_mIoU = []
img_list = [file for file in os.listdir(IMG_DIR) if file.endswith('.png') or file.endswith('.jpg')]
i=0
for img in tqdm(img_list):
    file_name = img.rsplit('.')[0]
    img_path = os.path.join(IMG_DIR, img)
    txt_path = os.path.join(TXT_DIR, file_name + '.txt')
    
    image = cv2.imread(img_path)
    
    img_w = image.shape[1]
    img_h = image.shape[0]
    image_size = (img_w, img_h)
    
    seg_results = model(image, conf = seg_threshold, iou = 0.5, verbose=False)
    
    for r in seg_results:
        if r.masks is None:
            img = image
        else:
            mask = np.zeros((image.shape[0], image.shape[1], 3), dtype=np.uint8)
            for masks in r.masks:
                segmentations = masks.xy
                segmentations = np.array(segmentations).reshape(1,-1,2)
                segmentations = segmentations.astype(np.int32)
                cv2.fillPoly(mask, segmentations, color=(255,255,255,127))
            pred_img = cv2.addWeighted(image, 1, mask, 0.5, 0)
            
    pred_image_path = os.path.join(pred_image_folder, file_name + '.jpg')
    cv2.imwrite(pred_image_path, pred_img)
    
    gt_seg_list = []
    gt_cls_list = []
    pred_seg_list = []
    
    gt_image = image
    
    with open(txt_path , 'r') as file:
        mask = np.zeros((image.shape[0], image.shape[1], 3), dtype=np.uint8)
        
        for row in file:
            gt_cls = row.split()[0]
            gt_seg_ = row.split()[1:]
            gt_seg = np.array(gt_seg_, dtype=np.float32).reshape(-1,2)
            gt_cls_list.append(gt_cls)
            gt_seg_list.append(gt_seg)
            
            d_gt_cls = row.split()[0]
            d_gt_seg_ = row.split()[1:]
            d_gt_seg = np.array(gt_seg_, dtype=np.float32).reshape(-1,2)
            d_gt_seg[:,0] *= img_w
            d_gt_seg[:,1] *= img_h
            d_gt_seg = d_gt_seg.astype(np.int32)
            cv2.fillPoly(mask, [d_gt_seg], color=(255,255,255,127))
            
        gt_img = cv2.addWeighted(gt_image, 1, mask, 0.5, 0)
        gt_image_path = os.path.join(gt_image_folder, file_name + '.jpg')
        cv2.imwrite(gt_image_path, gt_img)
            
    if seg_results[0].masks is not None:
        pred_segs = seg_results[0].masks.xyn
        
        for pred_seg in pred_segs:
            pred_seg_list.append(pred_seg)
            
        pred_masks = yolo_segmentation_to_masks(pred_seg_list, image_size)
        
    if seg_results[0].masks is None:
        mask = np.zeros(image_size, dtype=np.uint8)
        pred_masks = list(mask)    
        
    gt_masks = yolo_segmentation_to_masks(gt_seg_list, image_size)
    
    intersection = np.logical_and(np.sum(gt_masks, axis=0) > 0, np.sum(pred_masks, axis=0) >0).sum()
    union = np.logical_or(np.sum(gt_masks, axis=0) > 0, np.sum(pred_masks, axis=0) > 0).sum()
    
    if union > 0:
        mIoU = intersection / union
    else:
        mIoU = 0.0
    
    average_mIoU.append(mIoU)
    
average = np.mean(average_mIoU)
print('Average mIoU: ', average)