# -- Code Cell --
import pandas as pd

# -- Code Cell --
import torch 
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset, random_split
from PIL import Image
import torch.optim as optim

# -- Code Cell --
train = pd.read_csv("train_data.csv")

# -- Code Cell --
train

# -- Code Cell --
labels = train['angle']

# -- Code Cell --
test = pd.read_csv("./test_data.csv")
test.head()

# -- Code Cell --
# img = Image.open(f"images/{test[f"img{0}"]}")
img = Image.open(f"images/{test[f'img{1}'][0]}")

# -- Markdown Cell --
# https://docs.opencv.org/3.4/d8/dc8/tutorial_histogram_comparison.html

# -- Code Cell --
import cv2 as cv
src_base = cv.imread(f"images/{test[f'img{1}'][4]}")
src_test1 = cv.imread(f"images/{test[f'img{2}'][4]}")
src_test2 = cv.imread(f"images/{test[f'img{3}'][4]}")

# -- Code Cell --
hsv_base = cv.cvtColor(src_base, cv.COLOR_BGR2HSV)
hsv_test1 = cv.cvtColor(src_test1, cv.COLOR_BGR2HSV)
hsv_test2 = cv.cvtColor(src_test2, cv.COLOR_BGR2HSV)

# -- Code Cell --
hsv_half_down = hsv_base[hsv_base.shape[0]//2:,:]

# -- Code Cell --
h_bins = 50
s_bins = 60
histSize = [h_bins, s_bins]
# hue varies from 0 to 179, saturation from 0 to 255
h_ranges = [0, 180]
s_ranges = [0, 256]
ranges = h_ranges + s_ranges # concat lists
# Use the 0-th and 1-st channels
channels = [0, 1]

# -- Code Cell --
hist_base = cv.calcHist([hsv_base], channels, None, histSize, ranges, accumulate=False)
cv.normalize(hist_base, hist_base, alpha=0, beta=1, norm_type=cv.NORM_MINMAX)
hist_half_down = cv.calcHist([hsv_half_down], channels, None, histSize, ranges, accumulate=False)
cv.normalize(hist_half_down, hist_half_down, alpha=0, beta=1, norm_type=cv.NORM_MINMAX)
hist_test1 = cv.calcHist([hsv_test1], channels, None, histSize, ranges, accumulate=False)
cv.normalize(hist_test1, hist_test1, alpha=0, beta=1, norm_type=cv.NORM_MINMAX)
hist_test2 = cv.calcHist([hsv_test2], channels, None, histSize, ranges, accumulate=False)
cv.normalize(hist_test2, hist_test2, alpha=0, beta=1, norm_type=cv.NORM_MINMAX)

# -- Code Cell --
for compare_method in range(1):
    base_base = cv.compareHist(hist_base, hist_base, compare_method)
    base_half = cv.compareHist(hist_base, hist_half_down, compare_method)
    base_test1 = cv.compareHist(hist_base, hist_test1, compare_method)
    base_test2 = cv.compareHist(hist_base, hist_test2, compare_method)
    if base_test1 > base_test2 or base_test2 > base_test1:
        if base_test1 > base_test2:
            print("img1 and img2 are similar")
        if base_test2 > base_test1:
            print("img1 and img3 are similar")
    else:
        print("img2 and img3 are similar")
    print('Method:', compare_method, 'Perfect, Base-Half, Base-Test(1), Base-Test(2) :',\
          base_base, '/', base_half, '/', base_test1, '/', base_test2)

# -- Code Cell --
max_relation_between_1_2 = 0
max_relation_between_1_3 =0
max_relation_between_2_3 =0
max_relation_between_2_1 = 0
max_relation_between_3_1 = 0
max_relation_between_3_2 = 0
for i in range(1,4):
    src_base = cv.imread(f"images/{test[f'img{i}'][4]}")
    if i+1>3:
        src_test1 = cv.imread(f"images/{test[f'img{1}'][4]}")
    else: src_test1 = cv.imread(f"images/{test[f'img{i+1}'][4]}")
    if i+2>3:
        src_test2 = cv.imread(f"images/{test[f'img{2}'][4]}")
    else: src_test2 = cv.imread(f"images/{test[f'img{i+2}'][4]}")
    hsv_base = cv.cvtColor(src_base, cv.COLOR_BGR2HSV)
    hsv_test1 = cv.cvtColor(src_test1, cv.COLOR_BGR2HSV)
    hsv_test2 = cv.cvtColor(src_test2, cv.COLOR_BGR2HSV)
    h_bins = 50
    s_bins = 60
    histSize = [h_bins, s_bins]
    # hue varies from 0 to 179, saturation from 0 to 255
    h_ranges = [0, 180]
    s_ranges = [0, 256]
    ranges = h_ranges + s_ranges # concat lists
    # Use the 0-th and 1-st channels
    channels = [0, 1]
    hist_base = cv.calcHist([hsv_base], channels, None, histSize, ranges, accumulate=False)
    hist_test1 = cv.calcHist([hsv_test1], channels, None, histSize, ranges, accumulate=False)
    cv.normalize(hist_test1, hist_test1, alpha=0, beta=1, norm_type=cv.NORM_MINMAX)
    hist_test2 = cv.calcHist([hsv_test2], channels, None, histSize, ranges, accumulate=False)
    cv.normalize(hist_test2, hist_test2, alpha=0, beta=1, norm_type=cv.NORM_MINMAX)
    compare_method=0
    base_test1 = cv.compareHist(hist_base, hist_test1, compare_method)
    base_test2 = cv.compareHist(hist_base, hist_test2, compare_method)
    if i == 1:
        relation_between_1_2 = base_test1
        if relation_between_1_2 >= max_relation_between_1_2:
            max_relation_between_1_2 = relation_between_1_2
        relation_between_1_3 = base_test2
        if relation_between_1_3 >= max_relation_between_1_3:
            max_relation_between_1_3 = relation_between_1_3
    if i == 2:
        relation_between_2_3 = base_test1
        if relation_between_2_3 >= max_relation_between_2_3: 
            max_relation_between_2_3 = relation_between_2_3
        relation_between_2_1 = base_test2
        if relation_between_2_1 >= max_relation_between_2_1: 
            max_relation_between_2_1 = relation_between_2_1
    if i == 3:
        relation_between_3_1 = base_test1
        if relation_between_3_1 >= max_relation_between_3_1: 
            max_relation_between_3_1 = relation_between_3_1
        relation_between_3_2 = base_test2
        if relation_between_3_2 >= max_relation_between_3_2: 
            max_relation_between_3_2 = relation_between_3_2

# -- Code Cell --
lista = [max_relation_between_1_2,max_relation_between_1_3,max_relation_between_2_3,max_relation_between_2_1,max_relation_between_3_1,max_relation_between_3_2]
max = 0
max_idx = 0
for idx,element in enumerate(lista):
    if element > max:
        max = element
        max_idx = idx
print(max, max_idx)

# -- Code Cell --
# for i in range(1,4):
#     src_base = cv.imread(f"images/{test[f'img{i}'][4]}")
#     src_test1 = cv.imread(f"images/{test[f'img{i+1}'][4]}")
#     src_test2 = cv.imread(f"images/{test[f'img{i+2}'][4]}")
#     hsv_base = cv.cvtColor(src_base, cv.COLOR_BGR2HSV)
#     hsv_test1 = cv.cvtColor(src_test1, cv.COLOR_BGR2HSV)
#     hsv_test2 = cv.cvtColor(src_test2, cv.COLOR_BGR2HSV)
#     hsv_half_down = hsv_base[hsv_base.shape[0]//2:,:]
#     h_bins = 50
#     s_bins = 60
#     histSize = [h_bins, s_bins]
#     # hue varies from 0 to 179, saturation from 0 to 255
#     h_ranges = [0, 180]
#     s_ranges = [0, 256]
#     ranges = h_ranges + s_ranges # concat lists
#     # Use the 0-th and 1-st channels
#     channels = [0, 1]
#     hist_base = cv.calcHist([hsv_base], channels, None, histSize, ranges, accumulate=False)
#     cv.normalize(hist_base, hist_base, alpha=0, beta=1, norm_type=cv.NORM_MINMAX)
#     hist_half_down = cv.calcHist([hsv_half_down], channels, None, histSize, ranges, accumulate=False)
#     cv.normalize(hist_half_down, hist_half_down, alpha=0, beta=1, norm_type=cv.NORM_MINMAX)
#     hist_test1 = cv.calcHist([hsv_test1], channels, None, histSize, ranges, accumulate=False)
#     cv.normalize(hist_test1, hist_test1, alpha=0, beta=1, norm_type=cv.NORM_MINMAX)
#     hist_test2 = cv.calcHist([hsv_test2], channels, None, histSize, ranges, accumulate=False)
#     cv.normalize(hist_test2, hist_test2, alpha=0, beta=1, norm_type=cv.NORM_MINMAX)
#     for compare_method in range(1):
#         base_base = cv.compareHist(hist_base, hist_base, compare_method)
#         base_half = cv.compareHist(hist_base, hist_half_down, compare_method)
#         base_test1 = cv.compareHist(hist_base, hist_test1, compare_method)
#         base_test2 = cv.compareHist(hist_base, hist_test2, compare_method)
#         if base_test1 > base_test2 or base_test2 > base_test1:
#             if base_test1 > base_test2:
#                 print("img1 and img2 are similar")
#             if base_test2 > base_test1:
#                 print("img1 and img3 are similar")
#         else:
#             print("img2 and img3 are similar")
#         print('Method:', compare_method, 'Perfect, Base-Half, Base-Test(1), Base-Test(2) :',\
#             base_base, '/', base_half, '/', base_test1, '/', base_test2)

# -- Code Cell --
import matplotlib.pyplot as plt
for i in range(1,4):
    
    img = Image.open(f"images/{test[f'img{i}'][4]}")
    plt.figure()
    plt.imshow(img)
plt.show()

# -- Code Cell --


# -- Code Cell --


# -- Code Cell --


# -- Code Cell --


# -- Code Cell --


# -- Code Cell --


# -- Code Cell --
test_task_1 = test[test['subtaskID'] == 1]
test_task_1

# -- Code Cell --
answer = []
for y in range(len(test_task_1)):
    max_relation_between_1_2 = 0
    max_relation_between_1_3 =0
    max_relation_between_2_3 =0
    max_relation_between_2_1 = 0
    max_relation_between_3_1 = 0
    max_relation_between_3_2 = 0
    for i in range(1,4):
        src_base = cv.imread(f"images/{test[f'img{i}'][y]}")
        if i+1>3:
            src_test1 = cv.imread(f"images/{test[f'img{1}'][y]}")
        else: src_test1 = cv.imread(f"images/{test[f'img{i+1}'][y]}")
        if i+2 == 4:
            src_test2 = cv.imread(f"images/{test[f'img{1}'][y]}")
        elif i+2 == 5:
            src_test2 = cv.imread(f"images/{test[f'img{2}'][y]}")
        else:
            src_test2 = cv.imread(f"images/{test[f'img{i+2}'][y]}")
        hsv_base = cv.cvtColor(src_base, cv.COLOR_BGR2HSV)
        hsv_test1 = cv.cvtColor(src_test1, cv.COLOR_BGR2HSV)
        hsv_test2 = cv.cvtColor(src_test2, cv.COLOR_BGR2HSV)
        h_bins = 50
        s_bins = 60
        histSize = [h_bins, s_bins]
        # hue varies from 0 to 179, saturation from 0 to 255
        h_ranges = [0, 180]
        s_ranges = [0, 256]
        ranges = h_ranges + s_ranges # concat lists
        # Use the 0-th and 1-st channels
        channels = [0, 1]
        hist_base = cv.calcHist([hsv_base], channels, None, histSize, ranges, accumulate=False)
        cv.normalize(hist_base, hist_base, alpha=0, beta=1, norm_type=cv.NORM_MINMAX)
        hist_test1 = cv.calcHist([hsv_test1], channels, None, histSize, ranges, accumulate=False)
        cv.normalize(hist_test1, hist_test1, alpha=0, beta=1, norm_type=cv.NORM_MINMAX)
        hist_test2 = cv.calcHist([hsv_test2], channels, None, histSize, ranges, accumulate=False)
        cv.normalize(hist_test2, hist_test2, alpha=0, beta=1, norm_type=cv.NORM_MINMAX)
        compare_method=0
        base_test1 = cv.compareHist(hist_base, hist_test1, compare_method)
        base_test2 = cv.compareHist(hist_base, hist_test2, compare_method)
        if i == 1:
            relation_between_1_2 = base_test1
            if relation_between_1_2 >= max_relation_between_1_2:
                max_relation_between_1_2 = relation_between_1_2
            relation_between_1_3 = base_test2
            if relation_between_1_3 >= max_relation_between_1_3:
                max_relation_between_1_3 = relation_between_1_3
        if i == 2:
            relation_between_2_3 = base_test1
            if relation_between_2_3 >= max_relation_between_2_3: 
                max_relation_between_2_3 = relation_between_2_3
            relation_between_2_1 = base_test2
            if relation_between_2_1 >= max_relation_between_2_1: 
                max_relation_between_2_1 = relation_between_2_1
        if i == 3:
            relation_between_3_1 = base_test1
            if relation_between_3_1 >= max_relation_between_3_1: 
                max_relation_between_3_1 = relation_between_3_1
            relation_between_3_2 = base_test2
            if relation_between_3_2 >= max_relation_between_3_2: 
                max_relation_between_3_2 = relation_between_3_2
        lista = [max_relation_between_1_2,max_relation_between_1_3,max_relation_between_2_3,max_relation_between_2_1,max_relation_between_3_1,max_relation_between_3_2]
    max = 0
    max_idx = 0
    for idx,element in enumerate(lista):
        if element > max:
            max = element
            max_idx = idx
    if max_idx == 0:
        answer.append(3)
    if max_idx ==1:
        answer.append(2)
    if max_idx ==2:
        answer.append(1)
    if max_idx == 3:
        answer.append(3)
    if max_idx == 4:
        answer.append(2)
    if max_idx == 5:
        answer.append(1)

# -- Code Cell --
import numpy as np

# -- Code Cell --
answer

# -- Code Cell --
test_task_2 = test[test['subtaskID'] == 2]
test_task_2

# -- Code Cell --
sub1 = pd.DataFrame({
    "subtaskID":1,
    "datapointID":test_task_1['datapointID'],
    "answer":answer
})
sub2 = pd.DataFrame({
    "subtaskID":2,
    "datapointID":test_task_2['datapointID'],
    "answer":0*len(test_task_2['datapointID'])
})
final = pd.concat([sub1,sub2]).to_csv("subs.csv",index=False)

# -- Code Cell --
train

# -- Code Cell --
angle_map = {}
for idx, element in enumerate(train['angle'].unique()):
    angle_map[element] = idx

# -- Code Cell --
class TrainDataset(Dataset):
    def __init__(self,path, transform):
        self.df = pd.read_csv(path)
        self.transform = transform
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        row = self.df.iloc[index]
        image = Image.open(f"images/{row['image']}")
        label = angle_map[row['angle']]
        image = self.transform(image)
        label=torch.tensor(label,dtype=torch.long)
        return image,label

# -- Code Cell --
from torchvision import transforms

# -- Code Cell --
train['angle'].nunique()

# -- Code Cell --
transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize([0.5,0.5, 0.5],[0.5,0.5, 0.5])
])

# -- Code Cell --
class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(128, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1))
        )

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, 72)
        )

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x

# -- Code Cell --
train_ds = TrainDataset("train_data.csv", transform = transform)
trian_subset, val_subste = random_split(train_ds, [0.8,0.2])
train_loader = DataLoader(trian_subset, batch_size= 32, shuffle =True, num_workers =0)
val_loader = DataLoader(val_subste, batch_size= 32, shuffle =True, num_workers =0)

# -- Code Cell --
train['angle'].unique()

# -- Code Cell --
model = CNN()
model = model.to('cuda')
optimizer = optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.2)

# -- Code Cell --
def custom_score_task_2(predicted_angle, true_angle):
    predicted_angle = np.array(predicted_angle)
    true_angle = np.array(true_angle)
    z = np.minimum(np.abs(predicted_angle - true_angle), 360 - np.abs(predicted_angle - true_angle))
    return np.mean(np.maximum(0, 1 - z / 90))

# -- Code Cell --
rev_anlge_map = {}
for key,value in angle_map.items():
    rev_anlge_map[value]=key

# -- Code Cell --
for epoch in range(20):
    model.train()
    for img,label in train_loader:
        img = img.to("cuda")
        label = label.to("cuda")
        optimizer.zero_grad()
        output =model(img)
        loss= criterion(output, label)
        loss.backward()
        optimizer.step()
    model.eval()
    all_labels= []
    all_preds= []
    with torch.no_grad():
        for img,label in val_loader:
            img = img.to("cuda")
            label = label.to("cuda")
            output =model(img)
            pred = output.argmax(dim=1)
            all_preds.append(pred.cpu())
            all_labels.append(label.cpu())
    all_labels= torch.cat(all_labels).numpy()
    all_preds= torch.cat(all_preds).numpy()
    all_preds_to_angel = []
    all_labele_to_angle = []
    for elemnet in all_preds:
        all_preds_to_angel.append(rev_anlge_map[elemnet])
    for elemnet in all_labels:
        all_labele_to_angle.append(rev_anlge_map[elemnet])
    score = custom_score_task_2(all_preds_to_angel,all_labele_to_angle)
    scheduler.step()
    print(f"epoch {epoch}, score {score}")

# -- Code Cell --
test_task_2

# -- Code Cell --
class TestDataset(Dataset):
    def __init__(self,df, transform):
        self.df = df
        self.transform = transform
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        row = self.df.iloc[index]
        image = Image.open(f"images/{row['img1']}")
        image = self.transform(image)
        return image

# -- Code Cell --
test_ds = TestDataset(test_task_2,transform)
test_loader = DataLoader(test_ds, shuffle=False, num_workers=0, batch_size=32)

# -- Code Cell --
model.eval()
preds=[]
with torch.no_grad():
    for img in test_loader:
        img = img.to("cuda")
        output = model(img)
        pred=output.argmax(dim=1)
        preds.append(pred.cpu())
preds = torch.cat(preds).numpy()
preds

# -- Code Cell --
preds_w_angle = []
for element in preds:
    preds_w_angle.append(rev_anlge_map[element])

# -- Code Cell --
sub1 = pd.DataFrame({
    "subtaskID":1,
    "datapointID":test_task_1['datapointID'],
    "answer":answer
})
sub2 = pd.DataFrame({
    "subtaskID":2,
    "datapointID":test_task_2['datapointID'],
    "answer":preds_w_angle
})
final = pd.concat([sub1,sub2]).to_csv("subs.csv",index=False)

# -- Code Cell --
import math
from sklearn.svm import SVR

# -- Code Cell --
train

# -- Code Cell --
hist_features =[]
for i in range(len(train)):
    img = cv.imread(f"images/{train['image'][i]}")
    hsv_base = cv.cvtColor(img, cv.COLOR_BGR2HSV)
    h_bins = 50
    s_bins = 60
    histSize = [h_bins, s_bins]
    # hue varies from 0 to 179, saturation from 0 to 255
    h_ranges = [0, 180]
    s_ranges = [0, 256]
    ranges = h_ranges + s_ranges # concat lists
    # Use the 0-th and 1-st channels
    channels = [0, 1]
    hist_base = cv.calcHist([hsv_base], channels, None, histSize, ranges, accumulate=False)
    cv.normalize(hist_base, hist_base, alpha=0, beta=1, norm_type=cv.NORM_MINMAX)
    hist_features.append(hist_base)

# -- Code Cell --
len(hist_features)

# -- Code Cell --
len(train)

# -- Code Cell --
len(hist_features[67])

# -- Code Cell --
X_train = np.array(hist_features).reshape(len(hist_features), -1)
y_sin_train = []
y_cos_train = []

for _, row in train.iterrows():
    angle_rad = math.radians(row['angle'])
    y_sin_train.append(math.sin(angle_rad))
    y_cos_train.append(math.cos(angle_rad))

X_train = np.array(X_train)
y_sin_train = np.array(y_sin_train)
y_cos_train = np.array(y_cos_train)

regr_sin = SVR(kernel='rbf', C=500.0, epsilon=0.01, gamma='scale')
regr_cos = SVR(kernel='rbf', C=500.0, epsilon=0.01, gamma='scale')
regr_sin.fit(X_train, y_sin_train)
regr_cos.fit(X_train, y_cos_train)

# -- Code Cell --
test_task_2

# -- Code Cell --
hist_features_test =[]
for i in range(len(test_task_2)):
    img = cv.imread(f"images/{test_task_2['img1'].iloc[i]}")
    hsv_base = cv.cvtColor(img, cv.COLOR_BGR2HSV)
    h_bins = 50
    s_bins = 60
    histSize = [h_bins, s_bins]
    # hue varies from 0 to 179, saturation from 0 to 255
    h_ranges = [0, 180]
    s_ranges = [0, 256]
    ranges = h_ranges + s_ranges # concat lists
    # Use the 0-th and 1-st channels
    channels = [0, 1]
    hist_base = cv.calcHist([hsv_base], channels, None, histSize, ranges, accumulate=False)
    cv.normalize(hist_base, hist_base, alpha=0, beta=1, norm_type=cv.NORM_MINMAX)
    hist_features_test.append(hist_base)

# -- Code Cell --
hist_features_test

# -- Code Cell --
X_test = np.array(hist_features_test).reshape(len(hist_features_test), -1)

# -- Code Cell --
angles=[]
pred_sin = regr_sin.predict(X_test)
pred_cos = regr_cos.predict(X_test)

angles = []
for sin, cos in zip(pred_sin, pred_cos):
    angle_rad = math.atan2(sin, cos)
    angle_deg = round(math.degrees(angle_rad)) % 360
    angles.append(angle_deg)

# -- Code Cell --
sub1 = pd.DataFrame({
    "subtaskID":1,
    "datapointID":test_task_1['datapointID'],
    "answer":answer
})
sub2 = pd.DataFrame({
    "subtaskID":2,
    "datapointID":test_task_2['datapointID'],
    "answer":angles
})
final = pd.concat([sub1,sub2]).to_csv("subs.csv",index=False)

# -- Code Cell --
