# -- Code Cell --
import os
from PIL import Image
import torch
import torch.nn as nn
from torch.utils.data import DataLoader,Dataset
from torchvision import transforms
import matplotlib.pyplot as plt

# -- Code Cell --
import cv2 as cv

# -- Code Cell --
#Strategie: binary mask(adica doar grayscale) cu cv2 pe Satellite_Images-1 apoi UNET pe Satellite_Images-1 + masks si dau pe restu.

# -- Code Cell --
img = cv.imread("./HOTSPOT JUDGE/Satellite_Images-1/image_00230.png")
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
plt.imshow(gray, cmap='gray')

# -- Code Cell --
mask = cv.imread('./HOTSPOT JUDGE/Satellite_Images_MASKS4/mask_00248.png', cv.IMREAD_GRAYSCALE)
blur = cv.GaussianBlur(img, (1, 1), 0)
gray = cv.cvtColor(blur, cv.COLOR_BGR2GRAY)
ret,th1 = cv.threshold(gray,70,255,cv.THRESH_BINARY)
plt.imshow(mask, cmap='gray')

# -- Code Cell --
img = cv.imread("./HOTSPOT JUDGE/Satellite_Images-3/image_00248.png")
dst = cv.inpaint(img,mask,3,cv.INPAINT_TELEA)

# -- Code Cell --
for i in range(len(os.listdir("./HOTSPOT JUDGE/Satellite_Images-1"))):
    img = cv.imread(f"./HOTSPOT JUDGE/Satellite_Images-1/image_{i:05d}.png")
    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    plt.imsave(f"C:/Users/David/Desktop/ProblemeML/HotSpot/HOTSPOT JUDGE/Satellite_Images_MASKS1/mask_{i:05d}.png", gray, cmap='gray')

# -- Code Cell --
for i in range(len(os.listdir("./HOTSPOT JUDGE/Satellite_Images-2"))):
    img = cv.imread(f"./HOTSPOT JUDGE/Satellite_Images-2/image_{i:05d}.png")
    blur = cv.GaussianBlur(img, (3, 3), 0)
    gray = cv.cvtColor(blur, cv.COLOR_BGR2GRAY)
    ret2,th2 = cv.threshold(gray,0,255,cv.THRESH_BINARY+cv.THRESH_OTSU)
    plt.imsave(f"C:/Users/David/Desktop/ProblemeML/HotSpot/HOTSPOT JUDGE/Satellite_Images_MASKS2/mask_{i:05d}.png", th2, cmap='gray')

# -- Code Cell --
for i in range(len(os.listdir("./HOTSPOT JUDGE/Satellite_Images-3"))):
    img = cv.imread(f"./HOTSPOT JUDGE/Satellite_Images-3/image_{i:05d}.png")
    blur = cv.blur(img,(3,3))
    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    ret,th1 = cv.threshold(gray,65,255,cv.THRESH_BINARY)
    plt.imsave(f"C:/Users/David/Desktop/ProblemeML/HotSpot/HOTSPOT JUDGE/Satellite_Images_MASKS3/mask_{i:05d}.png", th1, cmap='gray')

# -- Code Cell --
for i in range(len(os.listdir("./HOTSPOT JUDGE/Satellite_Images-4"))):
    img = cv.imread(f"./HOTSPOT JUDGE/Satellite_Images-4/image_{i:05d}.png")
    blur = cv.GaussianBlur(img, (3, 3), 0)
    gray = cv.cvtColor(blur, cv.COLOR_BGR2GRAY)
    ret,th1 = cv.threshold(gray,95,255,cv.THRESH_BINARY)
    plt.imsave(f"C:/Users/David/Desktop/ProblemeML/HotSpot/HOTSPOT JUDGE/Satellite_Images_MASKS4/mask_{i:05d}.png", th1, cmap='gray')

# -- Code Cell --
from PIL import Image
import os
import numpy as np

def image_to_rle(image_path):
    image = Image.open(image_path)
    side, side = image.size
    arr = np.array(image)
    red = arr[:, :, 0].reshape(side * side, order = 'F')
    green = arr[:, :, 1].reshape(side * side, order = 'F')
    blue = arr[:, :, 2].reshape(side * side, order = 'F')
    
    start = 0
    length = 0
    res = []
    for i in range(side * side):
        if red[i] == 255 and green[i] == 255 and blue[i] == 255:
            if length == 0:
                start = i + 1
            length = length + 1
        else:
            if length > 0:
                res.append(start)
                res.append(length)
            length = 0
    if length > 0:
        res.append(start)
        res.append(length)
        
    return res

image_names = []
ground_truth = []

path_prefix = './HOTSPOT JUDGE/Satellite_Images_MASKS'
for subtask_id in range(1, 5):
    dir_path = path_prefix + str(subtask_id)
    for image_name in os.listdir(dir_path):
        image_names.append(image_name)
        ground_truth.append(image_to_rle(dir_path + '/' + image_name))

# -- Code Cell --
ground_truth[:250]

# -- Code Cell --
image_names_al_meu = []
for i in range(len(os.listdir("./HOTSPOT JUDGE/Satellite_Images-1"))):
    image_names_al_meu.append(f"image_{i:05d}.png")

# -- Code Cell --
image_names_al_meu

# -- Code Cell --
len(ground_truth)

# -- Code Cell --
import pandas as pd
sub1 = pd.DataFrame({
    'subtaskID':1,
    'datapointID':image_names_al_meu,
    'answer':ground_truth[:250],
})
sub2 = pd.DataFrame({
    'subtaskID':2,
    'datapointID':image_names_al_meu,
    'answer':ground_truth[250:500],
})
sub3 = pd.DataFrame({
    'subtaskID':3,
    'datapointID':image_names_al_meu,
    'answer':ground_truth[500:750],
})
sub4 = pd.DataFrame({
    'subtaskID':4,
    'datapointID':image_names_al_meu,
    'answer':ground_truth[750:],
})

# -- Code Cell --
sub_final = pd.concat([sub1,sub2,sub3,sub4]).to_csv('subs.csv',index=False)

# -- Code Cell --
