# -- Code Cell --
import cv2 as cv
import matplotlib.pyplot as plt
import os

# -- Markdown Cell --
# # TASK1

# -- Code Cell --


# -- Code Cell --


# -- Code Cell --


# -- Code Cell --
img = cv.imread("./test_data/5.png", cv.IMREAD_GRAYSCALE)
img.shape
img = img[0:img.shape[0], 25:img.shape[1]]
for i in range(5):
    cv.line(img,(0,13+(7*i)),(img.shape[1],13+(7*i)),(255,255,255),1)
blur = cv.GaussianBlur(img,(7,7),0)
ret3,th3 = cv.threshold(blur,0,255,cv.THRESH_BINARY_INV+cv.THRESH_OTSU)
plt.imshow(th3,cmap="gray")

# -- Code Cell --
n_labels, labels, stats, centroids = cv.connectedComponentsWithStats(th3)
boxes= []
for i in range(1,n_labels):
    x = stats[i, cv.CC_STAT_LEFT]
    y = stats[i, cv.CC_STAT_TOP]
    w = stats[i, cv.CC_STAT_WIDTH]
    h = stats[i, cv.CC_STAT_HEIGHT]
    area = stats[i, cv.CC_STAT_AREA]

    if area > 10:
        boxes.append((x, y, w, h))

# -- Code Cell --
img_color = cv.cvtColor(th3, cv.COLOR_GRAY2BGR)

for (x, y, w, h) in boxes:
    cv.rectangle(img_color, (x, y), (x + w, y + h), (0, 255, 0), 2)

plt.figure(figsize=(10,4))
plt.imshow(cv.cvtColor(img_color, cv.COLOR_BGR2RGB))
plt.show()

# -- Code Cell --
from tqdm.notebook import tqdm
import numpy as np
task1 = []
keep_all_boxes_for_task2=[]
id_1 = []
for j in tqdm(range(len(os.listdir("./test_data")))):
    img = cv.imread(f"./test_data/{j}.png", cv.IMREAD_GRAYSCALE)
    img.shape
    img = img[0:img.shape[0], 25:img.shape[1]]
    for i in range(5):
        cv.line(img,(0,13+(7*i)),(img.shape[1],13+(7*i)),(255,255,255),1)
    kernel = np.ones((3, 3), np.uint8) 
    img = cv.bitwise_not(img)
    blur = cv.dilate(img, kernel, iterations=1) 
  #  plt.imshow(blur)
  #  plt.show()
   # break
    n_labels, labels, stats, centroids = cv.connectedComponentsWithStats(blur)
    task1.append(n_labels - 1)
    keep_all_boxes_for_task2.append(boxes)
   # print(i)
    id_1.append(j)
  #  break

# -- Markdown Cell --
# # TASK2

# -- Code Cell --
train_img = cv.imread("./train_data/all_piano_notes_duration.png")
train_img = cv.cvtColor(train_img, cv.COLOR_BGR2GRAY)

note1 = train_img[0:50, 10:30]
note2 = train_img[0:50, 70:100]
note3 = train_img[0:50, 120:150]
note4 = train_img[0:50, 170:190]
note5 = train_img[0:50, 200:265]

notes = [note1, note2, note3, note4, note5]
processed_notes = []

for note in notes:
    blur = cv.GaussianBlur(note, (7,7), 0)
    _, th = cv.threshold(blur, 0, 255, cv.THRESH_BINARY_INV + cv.THRESH_OTSU)
    processed_notes.append(th)

# -- Code Cell --
def compute_hu(img):
    moments = cv.moments(img)
    hu = cv.HuMoments(moments).flatten()
    hu = -np.sign(hu) * np.log10(np.abs(hu) + 1e-10)
    return hu

hu1 = compute_hu(processed_notes[0])
hu2 = compute_hu(processed_notes[1])
hu3 = compute_hu(processed_notes[2])
hu4 = compute_hu(processed_notes[3])
hu5 = compute_hu(processed_notes[4])

all_templtes = [hu1, hu2, hu3, hu4, hu5]

# -- Code Cell --
img = cv.imread("./test_data/1.png", cv.IMREAD_GRAYSCALE)
img = img[:, 25:]

for i in range(5):
    cv.line(img, (0, 13 + 7*i), (img.shape[1], 13 + 7*i), 255, 1)

blur = cv.GaussianBlur(img, (7,7), 0)
_, th3 = cv.threshold(blur, 0, 255, cv.THRESH_BINARY_INV + cv.THRESH_OTSU)

salvez_peaici = []

for box in keep_all_boxes_for_task2[1]:
    x, y, w, h = box
    crop = th3[y:y+h, x:x+w]

    crop_hu = compute_hu(crop)

    distances = []
    for template_hu in all_templtes:
        dist = np.linalg.norm(crop_hu - template_hu)
        distances.append(dist)

    best_label = np.argmin(distances) + 1
    salvez_peaici.append(best_label)

# -- Code Cell --
img = cv.imread(f"./test_data/{8}.png", cv.IMREAD_GRAYSCALE)
plt.imshow(img)

# -- Code Cell --
salvez_peaici

# -- Code Cell --
mapping = {1: 1,2: 1/2,3: 1/4,4: 1/8,5: 1/16}
rezultate_finale = []

# -- Code Cell --
len(os.listdir("./test_data"))

# -- Code Cell --
for i in range(len(keep_all_boxes_for_task2)):
    img = cv.imread(f"./test_data/{i}.png", cv.IMREAD_GRAYSCALE)
    img = img[:, 25:]
    for j in range(5):
        cv.line(img, (0, 13 + 7*j), (img.shape[1], 13 + 7*j), 255, 1)

    blur = cv.GaussianBlur(img, (7,7), 0)
    _, th3 = cv.threshold(blur, 0, 255, cv.THRESH_BINARY_INV + cv.THRESH_OTSU)
    salvez_peaici = []

    for box in keep_all_boxes_for_task2[i]:
        x, y, w, h = box
        crop = th3[y:y+h, x:x+w]

        crop_hu = compute_hu(crop)

        distances = []
        for template_hu in all_templtes:
            dist = np.linalg.norm(crop_hu - template_hu)
            distances.append(dist)

        best_label = np.argmin(distances) + 1

        salvez_peaici.append(mapping[best_label])

    rezultate_finale.append(sum(salvez_peaici))

# -- Code Cell --
rezultate_finale

# -- Code Cell --
import pandas as pd
ids = [i for i in range(2000)]

# -- Code Cell --
len(rezultate_finale)

# -- Code Cell --
sub1 = pd.DataFrame({
    'subtaskID':1,
    "datapointID":id_1,
    "answer":task1
})
sub1["answer"] = sub1["answer"].astype(str)
sub2 = pd.DataFrame({
    'subtaskID':2,
    "datapointID":ids,
    "answer":rezultate_finale
})
sub3 = pd.DataFrame({
    'subtaskID':3,
    "datapointID":ids,
    "answer":task1
})
final = pd.concat([sub1,sub2,sub3]).to_csv("subs.csv",index=False)

# -- Code Cell --
