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

# -- Code Cell --
img = cv.imread("./data/train/263.jpg")
img_rgb = cv.cvtColor(img, cv.COLOR_BGR2RGB)

h, w = img_rgb.shape[:2]
margin = 5
rect = (margin, margin, w - 2*margin, h - 2*margin)
mask = np.zeros((h, w), np.uint8)
bgdModel = np.zeros((1,65), np.float64)
fgdModel = np.zeros((1,65), np.float64)
cv.grabCut(img_rgb, mask, rect, bgdModel, fgdModel, 10, cv.GC_INIT_WITH_RECT)
mask2 = np.where((mask==2)|(mask==0), 0, 1).astype('uint8')
result = np.full_like(img_rgb, 128)
result[mask2 == 1] = img_rgb[mask2 == 1]

plt.imshow(result), plt.show()

# -- Markdown Cell --
# # NU FOLOSESC ASTA PENTRU CA DUREAZA 1 MINUT un batch

# -- Code Cell --
from torch.utils.data import DataLoader,Dataset,random_split
import torch.nn as nn
import torch
import torch.optim as optim
from torchvision import transforms, models
import pandas as pd
from PIL import Image

# -- Code Cell --
train = pd.read_csv("./data/train.csv")
train.head()

# -- Code Cell --
class FullDataset(Dataset):
    def __init__(self,path,transform):
        super().__init__()
        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]
        img = cv.imread(f"./data/train/{index}.jpg")
        label = row['y']
        img_crop = img[20:204, 20:204]
        imgmod = Image.fromarray(img_crop)
        if self.transform:
            imgmod = self.transform(imgmod)
        return imgmod,label

# -- Code Cell --
transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])
])

# -- Code Cell --
fullds = FullDataset("./data/train.csv",transform)

# -- Code Cell --
train_size = int(0.8*(len(fullds)))
val_size = len(fullds) - train_size
train_subset, val_subset = random_split(fullds, [train_size,val_size])

# -- Code Cell --
train_loader = DataLoader(train_subset, batch_size=64, num_workers=0, shuffle=True)
val_loader = DataLoader(val_subset, batch_size=64, num_workers=0, shuffle=False)

# -- Code Cell --
train["y"].unique()

# -- Code Cell --
device = torch.device('cuda')

# -- Code Cell --
class modelmeu(nn.Module):
    def __init__(self):
        super().__init__()
        self.resnet = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
        self.resnet.fc = nn.Sequential(
            nn.Linear(2048, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 2),
        )
    def forward(self,x):
        out = self.resnet(x)
        return out

# -- Code Cell --
model = modelmeu()
for param in model.resnet.parameters():
    param.requires_grad = False
for param in model.resnet.fc.parameters():
    param.requires_grad = True
model = model.to(device)
optimizer = optim.Adam(model.resnet.fc.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
scheduler = optim.lr_scheduler.StepLR(optimizer,step_size=30, gamma=0.2)

# -- Code Cell --
decoy = 1

# -- Code Cell --
from tqdm import tqdm
for epoch in range(15):
    total_loss = 0
    for img, label in tqdm(train_loader):
        img=img.to(device)
        label=label.to(device)
        optimizer.zero_grad()
        output = model(img)
        loss = criterion(output,label)
        loss.backward()
        optimizer.step()
        total_loss+=loss.item()
    scheduler.step()
    if total_loss/len(train_loader)<decoy:
        decoy = total_loss/len(train_loader)
        torch.save(model.state_dict(), "model.pth")
    print(f"epoch{epoch+1} -- loss{total_loss/len(train_loader):04f}")

# -- Code Cell --
model = modelmeu()
model.load_state_dict(torch.load("model.pth"))
model = model.to(device)
model.eval()

# -- Code Cell --
class TestDataset(Dataset):
    def __init__(self,path,transform):
        super().__init__()
        self.df = pd.read_csv(path)
        self.transform = transform
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        img = cv.imread(f"./data/test/{index}.jpg")
        img_crop = img[20:204, 20:204]
        imgmod = Image.fromarray(img_crop)
        if self.transform:
            imgmod = self.transform(imgmod)
        return imgmod

# -- Code Cell --
test_ds = TestDataset("./data/test.csv",transform)
test_loader = DataLoader(test_ds, batch_size=64, num_workers=0, shuffle=False)

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

# -- Code Cell --
test = pd.read_csv('./data/test.csv')
test.head()

# -- Code Cell --
sub = pd.DataFrame({
    'subtaskID':1,
    'datapointID':test['datapointID'],
    'answer':preds_final
}).to_csv('subs.csv',index=False)

# -- Code Cell --


# -- Code Cell --


# -- Code Cell --
import pandas as pd
train = pd.read_csv("./train.csv")
test = pd.read_csv("./test.csv")
val = pd.read_csv("./val.csv")

# -- Code Cell --
import cv2 as cv

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

# -- Code Cell --
img = cv.imread("./train/1.jpg")
img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
img1 = img[0:40,0:224]
img2 = img[0:224,0:40]
img3 = img[184:224,0:224]
img4 = img[0:224,184:224]
img = img[20:img.shape[0]-20, 20:img.shape[1]-20]
import matplotlib.pyplot as plt
plt.imshow(img)

# -- Code Cell --
img.shape

# -- Code Cell --
import numpy as np

# -- Code Cell --
def has_water_bg(img_path):
    img = cv.imread(img_path)
    img = cv.cvtColor(img, cv.COLOR_BGR2RGB).astype(np.float32)

    img1 = img[0:40, 0:224]      
    img2 = img[0:224, 0:40]     
    img3 = img[184:224, 0:224]   
    img4 = img[0:224, 184:224]   
    edges = np.concatenate([
        img1.reshape(-1, 3),
        img2.reshape(-1, 3),
        img3.reshape(-1, 3),
        img4.reshape(-1, 3)
    ], axis=0)

    R = edges[:, 0]
    G = edges[:, 1]
    B = edges[:, 2]

    blue_dominant = (B > R) & (B > G) & (B > 100)
    return blue_dominant.mean() > 0.25

# -- Code Cell --
raspuns = has_water_bg("./train/1.jpg")
raspuns

# -- Code Cell --
train["filename"] = train["filename"].apply(lambda x : "./train/" + x)
val["filename"] = val["filename"].apply(lambda x : "./val/" + x)
test["filename"] = test["filename"].apply(lambda x : "./test/" + x)

# -- Code Cell --
train["filename"]

# -- Code Cell --
train['water_bg'] = train['filename'].apply(has_water_bg)

# -- Code Cell --
import pandas as pd

df_0 = train[train['y'] == 0]
df_1 = train[train['y'] == 1]

balanced_dfs = []

# Step 1: determine max balanced size per water_bg
bg_sizes = []

for bg in train['water_bg'].unique():
    size_0 = len(df_0[df_0['water_bg'] == bg])
    size_1 = len(df_1[df_1['water_bg'] == bg])
    bg_sizes.append(min(size_0, size_1))

# Step 2: ensure global balance
global_min_per_class = sum(bg_sizes)
target_per_class = global_min_per_class

current_0 = 0
current_1 = 0

for bg in train['water_bg'].unique():
    df_0_bg = df_0[df_0['water_bg'] == bg]
    df_1_bg = df_1[df_1['water_bg'] == bg]

    # number to sample from this bg
    n = min(len(df_0_bg), len(df_1_bg))

    df_0_sample = df_0_bg.sample(n, random_state=42)
    df_1_sample = df_1_bg.sample(n, random_state=42)

    balanced_dfs.append(df_0_sample)
    balanced_dfs.append(df_1_sample)

# Combine
df_balanced = pd.concat(balanced_dfs)

# ---- GLOBAL BALANCE FIX ----
df_0_bal = df_balanced[df_balanced['y'] == 0]
df_1_bal = df_balanced[df_balanced['y'] == 1]

min_global = min(len(df_0_bal), len(df_1_bal))

df_0_final = df_0_bal.sample(min_global, random_state=42)
df_1_final = df_1_bal.sample(min_global, random_state=42)

df_balanced_final = pd.concat([df_0_final, df_1_final]) \
                        .sample(frac=1, random_state=42) \
                        .reset_index(drop=True)


# -- Code Cell --
train = df_balanced_final

# -- Code Cell --
import torch
from torchvision import models
from torchvision.transforms import v2
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader,Dataset

# -- Code Cell --
import torch
from torchvision.transforms import v2
transform = v2.Compose([
    v2.ToImage(),
   # RemoveRedSquare(),  # Add this to remove artifacts
    v2.RandomResizedCrop((224, 224), scale=(0.5, 1.0)),
    v2.RandomHorizontalFlip(p=0.5),
    v2.RandomRotation(degrees=15),
    v2.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    v2.ToDtype(torch.float32, scale=True),
    v2.GaussianBlur(kernel_size=23, sigma=(0.1, 2.0)),
    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_valid = v2.Compose([
    v2.ToTensor(),
    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

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

# -- Code Cell --
class TrainDataset(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 = cv.imread(row['filename'])
        image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
        image = image[20:image.shape[0]-20, 20:image.shape[1]-20]
        image = Image.fromarray(image)
        image = self.transform(image)
        label = row['y']
        return image,label

# -- Code Cell --
train_dataset = TrainDataset(train, transform=transform)
val_dataset = TrainDataset(val, transform=transform_valid)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=0)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=True, num_workers=0)

# -- Code Cell --
class modelmeu(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.backbone = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
        for param in self.backbone.parameters():
             param.requires_grad = False
        self.backbone.fc =nn.Sequential(
            nn.Linear(2048, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(512, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(128, 2),
        )
    def forward(self,x):
        out = self.backbone(x)
        return out

# -- Code Cell --
from sklearn.metrics import f1_score
from sklearn.utils.class_weight import compute_class_weight
device = 'cuda'
class_weights = compute_class_weight("balanced", classes=np.array([0,1]),y = train['y'].values)
class_weights = torch.tensor(class_weights, dtype=torch.float32).to(device)

# -- Code Cell --
model = modelmeu()
model = model.to(device)
optimizer = optim.Adam(model.parameters(),lr=1e-4)
criterion = nn.CrossEntropyLoss(weight=class_weights)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer=optimizer,mode='max')

# -- Code Cell --
from tqdm import tqdm
decoy = -1
for epoch in range(10):
    model.train()
    for img,label in tqdm(train_loader):
        img = img.to(device)
        label = label.to(device)
        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 tqdm(val_loader):
            img = img.to(device)
            label = label.to(device)
            output = model(img)
            pred = output.argmax(dim=1)
            all_preds.append(pred.cpu())
            all_labels.append(label.cpu())
    all_preds = torch.cat(all_preds).numpy()
    all_labels = torch.cat(all_labels).numpy()
    f1 = f1_score(all_labels,all_preds,average='macro')
    if f1>decoy:
        decoy=f1
        torch.save(model.state_dict(),"model_tuned.pth")
    scheduler.step(f1)
    print(f"epoch {epoch+1}, f1 {f1}")

# -- Code Cell --
model = modelmeu()
model.load_state_dict(torch.load("model_tuned.pth"))
model = model.to(device)

# -- 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 = cv.imread(row['filename'])
        image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
        image = image[20:image.shape[0]-20, 20:image.shape[1]-20]
        image = Image.fromarray(image)
        image = self.transform(image)
        return image

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

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

# -- Code Cell --
test

# -- Code Cell --
out = pd.DataFrame({
    "subtaskID":1,
    "datapointID":test['datapointID'],
    'answer':preds
}).to_csv("subs.csv",index=False)

# -- Code Cell --
