# -- Code Cell --
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import cv2 as cv

df_train = pd.read_csv('train_data.csv')
df_test = pd.read_csv('test_data.csv')

df_train['image_path'] = df_train['image_path'].apply(lambda x: 'starter_kit/images/' + x)
df_test['image_path'] = df_test['image_path'].apply(lambda x: 'starter_kit/images/' + x)

df_train.head()

# -- Code Cell --
df_test.head()

# -- Code Cell --
import torch
import torch.nn as nn
from torchvision.transforms import v2
from torchvision.models import resnet18
from torch.utils.data import Dataset, DataLoader, random_split
from torch.optim.lr_scheduler import CosineAnnealingLR
from PIL import Image
from sklearn.utils.class_weight import compute_class_weight

class sami(Dataset):
    def __init__(self, df, test = False, transform = None):
        self.df = df.reset_index(drop = True)
        self.test = test
        self.transform = transform
        if test == False:
            self.labels = df['rotation_label'].values
    
    def __len__(self):
        return(len(self.df))
    
    def __getitem__(self, index):
        row = self.df.iloc[index]
        image =  Image.open(row['image_path']).convert('RGB')
        image = self.transform(image)
        if self.test == False:
            return image, self.labels[index]
        return image

transform_train = v2.Compose([
    v2.Resize((224, 224)),
    v2.ToImage(),
    v2.ToDtype(dtype = torch.float32, scale = True),
    v2.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225])
])

transform_test = v2.Compose([
    v2.Resize((224, 224)),
    v2.ToImage(),
    v2.ToDtype(dtype = torch.float32, scale = True),
    v2.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225])
])

full_dataset = sami(df = df_train, transform = transform_train)
test_dataset = sami(df = df_test, test = True, transform = transform_test)
train_dataset, valid_dataset = random_split(full_dataset, [0.8, 0.2])

train_loader = DataLoader(dataset = train_dataset, batch_size = 8, shuffle = True)
valid_loader = DataLoader(dataset = valid_dataset, batch_size = 8)
test_loader = DataLoader(dataset = test_dataset, batch_size = 8)

device = 'cuda' if torch.cuda.is_available() else 'cpu'

num_classes = df_train['rotation_label'].nunique()
print(num_classes)

# -- Code Cell --
model = resnet18(weights = None)
model.fc = nn.Sequential(nn.LazyLinear(64), nn.ReLU(), nn.LazyLinear(num_classes))
model = model.to(device)

epochs = 10

optimizer = torch.optim.AdamW(model.parameters(), lr = 2e-4, weight_decay = 1e-5)

labels = df_train['rotation_label'].values

class_weights = compute_class_weight(class_weight="balanced" , classes=np.unique(labels), y=labels)

class_weights = torch.tensor(class_weights, dtype=torch.float32).to(device)

criterion = nn.CrossEntropyLoss(weight=class_weights)
scheduler = CosineAnnealingLR(optimizer = optimizer , T_max = epochs)



# -- Code Cell --
from tqdm import tqdm
from sklearn.metrics import f1_score

best_model_f1 = 0.0
best_model_path = 'best_model.pt'

for epoch in range(epochs):
    model.train()
    train_loss = 0.0
    y_list = []
    y_hat_list = []

    for batch in tqdm(train_loader, total = len(train_loader)):
        X, y = batch[0].to(device), batch[1].to(device)
        optimizer.zero_grad()
        y_pred = model(X)
        loss = criterion(y_pred, y)
        train_loss += loss.item()

        loss.backward()
        optimizer.step()

        preds = torch.argmax(y_pred, dim = 1)
        y_list.extend(y.detach().cpu().numpy())
        y_hat_list.extend(preds.detach().cpu().numpy())
    
    train_loss /= len(train_loader)
    train_f1 = f1_score(y_list, y_hat_list, average = 'macro')

    model.eval()
    valid_loss = 0.0
    y_list = []
    y_hat_list = []

    with torch.no_grad():
        for batch in tqdm(valid_loader, total = len(valid_loader)):
            X, y = batch[0].to(device), batch[1].to(device)
            y_pred = model(X)
         #   print(y_pred.shape)
            loss = criterion(y_pred, y)
            valid_loss += loss.item()

            preds = torch.argmax(y_pred, dim = 1)
            y_list.extend(y.detach().cpu().numpy())
            y_hat_list.extend(preds.detach().cpu().numpy())
    
    valid_loss /= len(valid_loader)
    valid_f1 = f1_score(y_list, y_hat_list, average = 'macro')

    scheduler.step()

    print(f'Epoch {epoch + 1} / {epochs}')
    print(f'Train loss {train_loss:.4f} | Train f1 {train_f1:.4f}')
    print(f'Valid loss {valid_loss:.4f} | Valid f1 {valid_f1:.4f}')
    print('-' * 40)

    if valid_f1 > best_model_f1:
        best_model_f1 = valid_f1
        torch.save(model.state_dict(), best_model_path)



# -- Code Cell --
torch.serialization.safe_globals([getattr])

# -- Code Cell --
model.load_state_dict(torch.load(best_model_path, weights_only=False))

model.eval()

pred1 = []

with torch.no_grad():
    for batch in tqdm(test_loader, total = len(test_loader)):
        X = batch.to(device)
        y_pred = model(X)
        preds = torch.argmax(y_pred, dim = 1)
        pred1.extend(preds.detach().cpu().numpy())

subtask1 = pd.DataFrame({
    'subtaskID': 1,
    'datapointID': df_test['datapoint_id'],
    'answer': pred1
})

subtask1.head()

# -- Code Cell --
subtask1.to_csv("subi.csv", index= False)

# -- Code Cell --
catalog = pd.read_csv('vision_catalog.csv')

catalog.head()

# -- Code Cell --
catalog["image_path"] = catalog["image_path"].apply(lambda x : "starter_kit/images/" + x)

# -- Code Cell --
model.fc = nn.Identity()

# -- Code Cell --
from tqdm import tqdm

sift = cv.SIFT_create(contrastThreshold=0.03, edgeThreshold=15)
bf   = cv.BFMatcher()

def extract_sift_des(path):
    img = cv.imread(path, cv.IMREAD_GRAYSCALE)
    _, des = sift.detectAndCompute(img, None)
    return des 

def count_good_matches(des1, des2):
    if des1 is None or des2 is None or len(des1) < 2 or len(des2) < 2:
        return 0
    matches = bf.knnMatch(des1, des2, k=2)
    return sum(1 for m, n in matches if m.distance < 0.75 * n.distance)

# Extract descriptors
catalog_des = catalog["image_path"].apply(extract_sift_des).values
test_des    = df_test["image_path"].apply(extract_sift_des).values

N_test, N_catalog = len(test_des), len(catalog_des)
sims = np.zeros((N_test, N_catalog))

for i in tqdm(range(N_test)):
    for j in range(N_catalog):
        sims[i, j] = count_good_matches(test_des[i], catalog_des[j])

taken = set()
best_match_idx = np.full(N_test, -1, dtype=int)
best_match_sim = np.full(N_test, -np.inf)

priority = np.max(sims, axis=1).argsort()[::-1]

for i in priority:
    ranked = np.argsort(sims[i])[::-1]
    for j in ranked:
        if j not in taken:
            best_match_idx[i] = j
            best_match_sim[i] = sims[i, j]
            taken.add(j)
            break

df_test["matched_catalog_idx"]  = best_match_idx
df_test["matched_gallery_id"]   = catalog["gallery_id"].values[best_match_idx]
df_test["matched_image_path"]   = catalog["image_path"].values[best_match_idx]
df_test["matched_dist"]         = best_match_sim  

df_test[["datapoint_id", "image_path", "matched_gallery_id", "matched_image_path", "matched_dist"]].head(10)


# -- Code Cell --
subtask2 = pd.DataFrame({
    'subtaskID': 2,
    'datapointID': df_test['datapoint_id'],
    'answer': df_test["matched_gallery_id"]
})

# -- Code Cell --
sub = pd.concat([subtask1, subtask2])

# -- Code Cell --
sub.to_csv("subi.csv", index= False)

# -- Code Cell --


# -- Code Cell --
