# -- Code Cell --
import pandas as pd
import numpy as np
from PIL import Image
import os

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

# -- Code Cell --
def has_water_bg(img_path, threshold=0.25):
    img = Image.open(img_path)
    img = np.array(img,dtype=np.float32)
    H,W, _ = img.shape
    border_mask = np.zeros((H,W), dtype=bool)
    border=int(min(H,W) * 0.2)
    border_mask[0:border, :] = True
    border_mask[H-border:H, :]=True
    border_mask[:, 0:border] = True
    border_mask[:, W-border:W] = True
    border_pixels = img[border_mask]
    R,G,B = border_pixels[:,0], border_pixels[:,1], border_pixels[:,2]
    blue_dominant = (B>R) & (B>G) & (B>100)
    return blue_dominant.mean() > threshold

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

# -- Code Cell --
df_0 = train[train['y'] == 0]
df_1 = train[train['y'] == 1]
balanced_dfs = []
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]
    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)
df_balanced = pd.concat(balanced_dfs)
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 numpy as np
from PIL import Image

def remove_red_square_with_random_palette(img_path, pad=5, red_thresh=150):
    img = Image.open(img_path).convert("RGB")
    img = np.array(img, dtype=np.uint8)

    H, W, _ = img.shape

    border_mask = np.zeros((H, W), dtype=bool)
    border = int(min(H, W) * 0.2)
    border_mask[0:border, :] = True
    border_mask[H-border:H, :] = True
    border_mask[:, 0:border] = True
    border_mask[:, W-border:W] = True

    R = img[:, :, 0].astype(np.int32)
    G = img[:, :, 1].astype(np.int32)
    B = img[:, :, 2].astype(np.int32)

    red_mask = (R > red_thresh) & (R > G + 30) & (R > B + 30)
    red_mask = red_mask & border_mask

    ys, xs = np.where(red_mask)
    if len(xs) == 0:
        return img, None

    x_min, x_max = xs.min(), xs.max()
    y_min, y_max = ys.min(), ys.max()

    x1 = max(0, x_min - pad)
    y1 = max(0, y_min - pad)
    x2 = min(W, x_max + pad + 1)
    y2 = min(H, y_max + pad + 1)

    border_pixels = img[border_mask].astype(np.int32)

    low = np.percentile(border_pixels, 20, axis=0).astype(np.int32)
    high = np.percentile(border_pixels, 80, axis=0).astype(np.int32)

    h_box = y2 - y1
    w_box = x2 - x1

    fill = np.random.randint(low, high + 1, size=(h_box, w_box, 3), dtype=np.int32)
    fill = np.clip(fill, 0, 255).astype(np.uint8)

    img[y1:y2, x1:x2] = fill

    return img, [x1, y1, x2, y2]

# -- Code Cell --
import cv2 as cv

# -- Code Cell --
img = cv.imread("./train/0.jpg")
img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
img, _ =remove_red_square_with_random_palette("./train/0.jpg")
import matplotlib.pyplot as plt
plt.imshow(img)

# -- 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 --
import torch
from PIL import Image
from torch.utils.data import Dataset, DataLoader
class dataset_img_cls(Dataset):
    def __init__(self, df, test = False, transform = transform):
        super().__init__()
        self.X_path = df["filename"]
        self.transform = transform
        self.test = test
        if self.test == False:
             self.y = torch.tensor(np.array(df["y"]))
    def __len__(self):
        return len(self.X_path)

    def __getitem__(self, index):
        img = Image.open(self.X_path[index]).convert("RGB")
        img = transform(img)  # shape: [C, H, W]

        _, H, W = img.shape
        if self.test or (not self.test and self.y[index] == 1):
            square_size = np.random.randint(10, 21)
            position = np.random.randint(0, W - square_size)
            placement = np.random.randint(0, 4)

            if placement == 0:  # top
                img[0, 0:square_size, position:position+square_size] = 1.0
            elif placement == 1:  # bottom
                img[0, H-square_size:H, position:position+square_size] = 1.0
            elif placement == 2:  # left
                img[0, position:position+square_size, 0:square_size] = 1.0
            else:  # right
                img[0, position:position+square_size, W-square_size:W] = 1.0
        if not self.test:
            return img, self.y[index]
        else:
            return img

# -- Code Cell --
train_dt = dataset_img_cls(train)
valid_dt = dataset_img_cls(valid, transform= transform_valid)
test_dt = dataset_img_cls(test, test = True, transform= transform_valid)

# -- Code Cell --
train_loader = DataLoader(train_dt, batch_size= 32, shuffle= True)
valid_loader = DataLoader(valid_dt, batch_size= 32, shuffle= True)
test_loader = DataLoader(test_dt, batch_size= 32)


# -- Code Cell --
import torchvision
import torch.nn as nn
class model_bengos(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.backbone = torchvision.models.resnet50(pretrained = True)
        self.in_fc_feat = self.backbone.fc.in_features
        self.backbone.fc = nn.Identity() # Suprascrierea classification head
        for param in self.backbone.parameters():
             param.requires_grad = False
        self.classification_head = nn.Sequential(
                nn.Linear(self.in_fc_feat, 512),
                nn.BatchNorm1d(512),
                nn.ReLU(),
                nn.Dropout(0.5),

                nn.Linear(512, 128),
                nn.BatchNorm1d(128),
                nn.ReLU(),
                nn.Dropout(0.4),

                nn.Linear(128, 1)
            )
    def forward(self, x):
        x = self.backbone(x)
        x = self.classification_head(x)
        return x

# -- Code Cell --
from sklearn.utils.class_weight import compute_class_weight
device = "cuda" if torch.cuda.is_available() else "cpu"
class_weights = compute_class_weight(
    'balanced', 
    classes=np.array([0, 1]), 
    y=train["y"].values
)
pos_weight = torch.tensor([class_weights[1] / class_weights[0]]).to(device)
model = model_bengos()
model = model.to(device)
loss_fn = nn.BCEWithLogitsLoss(pos_weight = pos_weight)
optim = torch.optim.Adam(model.parameters(), lr = 5e-3)
sch = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer=optim)
epochs = 3

# -- Code Cell --
from tqdm.notebook import tqdm
from sklearn.metrics import f1_score
import torch
best_f1 = 0.0  # track best validation F1
best_model_path = "best_model.pth"

for epoch in range(epochs):
    loss_values = 0.0
    train_preds = []
    train_targets = []

    model.train()
    for batch in tqdm(train_loader, total=len(train_loader)):
        X = batch[0].to(device)
        y = batch[1].float().to(device)

        y_hat = model(X).squeeze(1)
        loss = loss_fn(y_hat, y)
        loss_values += loss.item()

        optim.zero_grad()
        loss.backward()
        optim.step()

        probs = torch.sigmoid(y_hat)
        preds = (probs > 0.5).float()

        train_preds.extend(preds.detach().cpu().numpy())
        train_targets.extend(y.detach().cpu().numpy())

    train_f1 = f1_score(train_targets, train_preds)

    with torch.no_grad():
        valid_loss_values = 0.0
        valid_preds = []
        valid_targets = []

        model.eval()
        for batch in tqdm(valid_loader, total=len(valid_loader)):
            X = batch[0].to(device)
            y = batch[1].float().to(device)

            y_hat = model(X).squeeze(1)

            loss = loss_fn(y_hat, y)
            valid_loss_values += loss.item()

            probs = torch.sigmoid(y_hat)
            preds = (probs > 0.5).float()

            valid_preds.extend(preds.cpu().numpy())
            valid_targets.extend(y.cpu().numpy())

        valid_f1 = f1_score(valid_targets, valid_preds)

    # ---- Save best model ----
    if valid_f1 > best_f1:
        best_f1 = valid_f1
        torch.save(model.state_dict(), best_model_path)
        print(f"🔥 New best model saved with F1: {best_f1:.4f}")

    sch.step(loss_values)

    print(f"Epoch: {epoch} | "
          f"Train Loss: {loss_values / len(train_loader):.4f} | "
          f"Train F1: {train_f1:.4f} | "
          f"Valid Loss: {valid_loss_values / len(valid_loader):.4f} | "
          f"Valid F1: {valid_f1:.4f}")

# -- Code Cell --
model.eval()
y_pred =[]
for batch in tqdm(test_loader, total = len(test_loader)):
    X = batch.to(device)
    y_hat = model(X)
    preds = (y_hat > 0.5).float()
    y_pred.extend(preds)
y_pred = [int(x.detach().item()) for x in y_pred]

# -- Code Cell --
pd.DataFrame({
    "subtaskID" : 1,
    "datapointID" :df_test["datapointID"],
    "answer" :y_pred
}).to_csv("subi.csv", index= False)