# %%
import pandas as pd
import numpy as np
import cv2
import os
import matplotlib.pyplot as plt
from tqdm.notebook import tqdm

# %%
test_data = []
for pth in os.listdir("./test_data"):
    test_data.append(os.path.join("./test_data", pth))

# %%


# %%


# %%


# %%
gt_img = cv2.imread("./train_data/all_piano_notes_duration.png", cv2.IMREAD_GRAYSCALE)
gt_img.shape

# %%
plt.imshow(gt_img)

# %%
def extract_ref_components(img):
    _, bw = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY_INV)
    kernel = np.ones((3,3), np.uint8)
    dilated = cv2.dilate(bw, kernel, iterations=1)
    n, labels, stats, _ = cv2.connectedComponentsWithStats(dilated, connectivity=8)
    crops = []
    for i in range(1, n):
        x, y, w, h, area = stats[i]
        if area > 50:
            crops.append(img[y:y+h, x:x+w])
    
    crops.sort(key=lambda c: stats[0][1])  
    return crops

img_ref = extract_ref_components(gt_img)

# %%
plt.imshow(img_ref[0])

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

# %%
ans_1 = []
id_1 = []
for pth in tqdm(test_data):
    img = cv2.imread(pth, cv2.IMREAD_GRAYSCALE)
    img = img[:, 30:]
    original = img.copy()  
    row_removed = []
    for i in range(img.shape[0]):
        if (img[i, :] == 0).sum() > img.shape[1] * 0.9: 
            row_removed.append(i)
    for x in row_removed:
        img[x, :] = 255
    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_1.append(totalLabels -1)
    id_1.append(int(pth.split('/')[-1].split('.')[0]))


# %%
len(ans_1)

# %% [markdown]
# Partial solution

# %%
import cv2
import numpy as np

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 = [compute_hu(ref) for ref in img_ref]

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_1 = []
ans_2 = []
id_1 = []

for pth in tqdm(test_data):
    img = cv2.imread(pth, cv2.IMREAD_GRAYSCALE)
    img = img[:, 30:]

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

    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_1.append(totalLabels - 1)
    ans_2.append(get_idx(img, label_ids))

# %%
import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from tqdm.notebook import tqdm

def extract_ref_components(img):
    _, bw = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY_INV)
    kernel = np.ones((3,3), np.uint8)
    dilated = cv2.dilate(bw, kernel, iterations=1)
    n, labels, stats, _ = cv2.connectedComponentsWithStats(dilated, connectivity=8)
    crops = []
    for i in range(1, n):
        x, y, w, h, area = stats[i]
        if area > 50:
            crops.append((x, img[y:y+h, x:x+w]))
    crops.sort(key=lambda c: c[0])  
    return [c for _, c in crops]

gt_img = cv2.imread("./train_data/all_piano_notes_duration.png", cv2.IMREAD_GRAYSCALE)
img_ref = extract_ref_components(gt_img)

class NoteDataset(Dataset):
    def __init__(self, crops, n_augment=500):
        self.samples = []
        self.labels = []
        aug = transforms.Compose([
                transforms.ToPILImage(),
                transforms.Resize((32, 32)),
                transforms.RandomAffine(degrees=10, translate=(0.1, 0.1), scale=(0.85, 1.15)),
                transforms.RandomPerspective(distortion_scale=0.2, p=0.5),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.5], std=[0.5])
                ])
        for label, crop in enumerate(crops):
            # Normalize crop to uint8 grayscale
            c = crop.astype(np.uint8)
            for _ in range(n_augment):
                t = aug(c)
                self.samples.append(t)
                self.labels.append(label)

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        return self.samples[idx], self.labels[idx]

dataset = NoteDataset(img_ref, n_augment=500)
loader = DataLoader(dataset, batch_size=64, shuffle=True)
def note_density(crop):
    resized = cv2.resize(crop, (32, 32))
    _, bw = cv2.threshold(resized, 127, 255, cv2.THRESH_BINARY_INV)
    return (bw > 0).mean()

class NoteCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d(4),
            nn.Flatten(),
            nn.Linear(64*4*4, 128), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(128, 5)
        )
    def forward(self, x):
        return self.net(x)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = NoteCNN().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

for epoch in range(30):
    model.train()
    total_loss = 0
    for x, y in loader:
        x, y = x.to(device), torch.tensor(y).to(device)
        optimizer.zero_grad()
        loss = criterion(model(x), y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    if (epoch+1) % 5 == 0:
        print(f"Epoch {epoch+1}/20 — loss: {total_loss/len(loader):.4f}")

map_idx = {0: 1.0, 1: 0.5, 2: 0.25, 3: 0.125, 4: 0.0625}

ref_densities = [(img_ref[j] == 0).mean() for j in [3, 4]]  # 1/8 and 1/16

preprocess = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((32, 32)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

ref_densities = [note_density(img_ref[j]) for j in [3, 4]]  # 1/8, 1/16
def classify_note(pred_img):
    t = preprocess(pred_img.astype(np.uint8)).unsqueeze(0).to(device)
    model.eval()
    with torch.no_grad():
        logits = model(t)
        probs = torch.softmax(logits, dim=1)[0]
        pred = int(torch.argmax(probs))

 #   if pred in (3, 4):
  #      density = (pred_img == 0).mean()
 ##       dists = [abs(density - r) for r in ref_densities]
    #    pred = 3 if dists[0] < dists[1] else 4

    return pred

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_1, ans_2, id_1 = [], [], []

for pth in tqdm(test_data):
    img = cv2.imread(pth, cv2.IMREAD_GRAYSCALE)
    img = img[:, 30:]

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

    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_1.append(totalLabels - 1)
   # id_1.append(int(pth.split('/')[-1].split('.')[0]))
    ans_2.append(get_idx(img, label_ids))

# %%
map_to_shit = {
    27 : "G",
    31 : "F",
    34 : "E",
    37 : "D",
    41 : "C",
    44 : "B",
    48 : "A"
}

# %%
def get_pos(img, bboxes):
    results = []
    for bbox in bboxes:
        x0, y0, w, h, area = bbox
        img_cpy = img[:, x0 : x0 + w]   
        col = img_cpy[::-1]                         
        has_zero = (col == 0).any(axis=1)
        idx = np.argmax(has_zero)
        if has_zero[idx]:
            results.append(map_to_shit[int(img.shape[0] - idx)])
    return results

# %%
ans_3 = []
all_vals = []
position = []
for pth in tqdm(test_data):
    img = cv2.imread(pth, cv2.IMREAD_GRAYSCALE)
    img = img[:, 30:]
    img_cpy = img.copy()
    for i in range(img.shape[0]):
        if (img[i, :] == 0).sum() > img.shape[1] * 0.9:
            img[i, :] = 255

    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
    sorted_indices = np.argsort(values[1:, cv2.CC_STAT_LEFT]) + 1
    values = values[sorted_indices]
    ans_1.append(totalLabels - 1)
    ans_3.append(get_idx(img, label_ids))
    all_vals.append(get_pos(img, values))

# %%

map_to_str = {
    0 :"1/1",
    1 : "1/2",
    2: "1/4",
    3 :"1/8",
    4 : "1/6"
}

# %%
ans_3 = []
for i in tqdm(range(len(all_vals))):
    ans = ""
    for j in range(len(all_vals[i])):
        ans += str(all_vals[i][j])
        ans += "-"
        ans += map_to_str[ans_2[i][j]]
        ans += " "
    ans_3.append(ans)
        

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

# %%
ans_1 = [float(x) for x in ans_1]
ans__2 = [float(x) for x in ans__2]

# %%
len(id_1)

# %%
len(ans_3)

# %%
df_subi_1 = pd.DataFrame({
    "subtaskID" : 1,
    "datapointID" : id_1,
    "answer" : ans_1
})
df_subi_2 = pd.DataFrame({
    "subtaskID" : 2,
    "datapointID" : id_1,
    "answer" : ans__2
})

df_subi_3 = pd.DataFrame({
    "subtaskID" : 3,
    "datapointID" : id_1,
    "answer" : ans_3
})
df_subi_1["answer"] = df_subi_1["answer"].astype(int).astype(str) 
df_subi_2["answer"] = df_subi_2["answer"].astype(float).astype(str)
df_subi_3["answer"] = df_subi_3["answer"].astype(str) 

pd.concat([df_subi_1, df_subi_2, df_subi_3]).to_csv("subi.csv", index=False)

# %%



