# -- 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 --
def compute_hu(img):
    _, bw = cv.threshold(img, 127, 255, cv.THRESH_BINARY_INV)
    moments = cv.moments(bw)
    hu = cv.HuMoments(moments).flatten()
    hu = -np.sign(hu) * np.log10(np.abs(hu) + 1e-10)
    return hu


# -- 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]

all_templtes = [compute_hu(note) for note in notes]


# -- Code Cell --
# Precompute Hu moments for each reference note
import cv2 as cv2
def compute_hu(img):
    _, bw = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY_INV)
    moments = cv2.moments(bw)
    hu = cv2.HuMoments(moments).flatten()
    hu = -np.sign(hu) * np.log10(np.abs(hu) + 1e-10)
    return hu

ref_hus =all_templtes

def classify_note(pred_img):
    _, bw = cv2.threshold(pred_img, 127, 255, cv2.THRESH_BINARY_INV)
    moments = cv2.moments(bw)
    hu = cv2.HuMoments(moments).flatten()
    hu = -np.sign(hu) * np.log10(np.abs(hu) + 1e-10)
    dists = [np.linalg.norm(hu - ref_hu) for ref_hu in ref_hus]
    return int(np.argmin(dists))

def get_idx(img, label_ids):
    results = []
    for i in range(1, np.max(label_ids) + 1):
        ox = np.where(label_ids == i)[0]
        oy = np.where(label_ids == i)[1]
        x0, y0 = np.min(ox), np.min(oy)
        x1, y1 = np.max(ox), np.max(oy)
        pred_img = img[x0:x1, y0:y1]
        if pred_img.size == 0:
            continue
        results.append(classify_note(pred_img))
    return results

ans_2 = []
id_2 = []
stats = []
for pth in tqdm(os.listdir("./test_data")):
    img = cv2.imread(os.path.join("./test_data", pth), cv2.IMREAD_GRAYSCALE)
    img = img[:, 30:]

    # Remove staff lines
    for i in range(img.shape[0]):
        if (img[i, :] == 0).sum() > img.shape[1] * 0.9:
            img[i, :] = 255

    # Connected components
    kernel = np.ones((3, 3), np.uint8)
    img_inv = cv2.bitwise_not(img)
    img_dilated = cv2.dilate(img_inv, kernel, iterations=1)
    analysis = cv2.connectedComponentsWithStats(img_dilated, 32, cv2.CV_32S)
    (totalLabels, label_ids, values, centroid) = analysis
    ans_2.append(get_idx(img, label_ids))
    id_2.append(pth.split('.')[0])
    stats.append(values)

# -- Code Cell --
map_idx = {
    0 :1,
    1 : 0.5,
    2: 0.25,
    3 :0.125,
    4 : 0.0625
}

# -- Code Cell --
ans__2 = [np.sum(np.array([map_idx[elem] for elem in x])) for x in ans_2]

# -- Markdown Cell --
# # TASK3

# -- Code Cell --
map_idx_pt_task3 = {
    0 :"1/1",
    1 : "1/2",
    2: "1/4",
    3 :"1/8",
    4 :"1/16"
}

# -- Code Cell --
valori_pt_task3 = []
for img in ans_2:
    valori_pt_task3.append([map_idx_pt_task3[x] for x in img])

# -- Code Cell --
stats[1]

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

# -- Code Cell --
def control(valoare):
    pozitii = {
        "G": 17,
        "F": 20,
        "E": 24,
        "D": 27,
        "C": 32,
        "B": 35,
        "A": 40
    }

    return min(pozitii, key=lambda nota: abs(valoare - pozitii[nota]))

# -- Code Cell --
answer = []
for idx, (x, y, w, h, area) in enumerate(stats[0]):
    if idx == 0:
        continue
    if area < 800:
        cx, cy = centroids[idx]
        answer.append(f"{control(int(cy))}-{valori_pt_task3[0][idx-1]}")

# -- Code Cell --
answer

# -- Code Cell --
plt.imshow(img,cmap='gray')
plt.show()

# -- Code Cell --
for i in range(len(os.listdir(./test_data))):
    for boxes in stats[i]:
        

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

# -- Code Cell --
