# -- Code Cell --
import pandas as pd
import numpy as np

# -- Code Cell --
df_train =pd.read_csv("train.csv")
df_valid =pd.read_csv("val.csv")
df_test =pd.read_csv("test.csv")
df_train["filename"] = df_train["filename"].apply(lambda x : "./train/" + x)
df_valid["filename"] = df_valid["filename"].apply(lambda x : "./val/" + x)
df_test["filename"] = df_test["filename"].apply(lambda x : "./test/" + x)
df_train.head(5)

# -- Code Cell --
from PIL import Image
import numpy as np

def has_water_background(img_path, threshold=0.25):
    """
    Detect water background by checking for blue-dominant pixels.
    Water backgrounds have high B relative to R and G.
    """
    img = np.array(Image.open(img_path).convert("RGB")).astype(float)
    H, W, _ = img.shape
    
    # Sample the background: corners and edges (avoid the bird center)
    # Use the border region (outer 20% of image)
    border_mask = np.zeros((H, W), dtype=bool)
    border = int(min(H, W) * 0.2)
    border_mask[:border, :] = True
    border_mask[-border:, :] = True
    border_mask[:, :border] = True
    border_mask[:, -border:] = True
    
    border_pixels = img[border_mask]  # shape: (N, 3)
    
    R, G, B = border_pixels[:, 0], border_pixels[:, 1], border_pixels[:, 2]
    # Blue dominant: B > R and B > G, and B is reasonably high
    blue_dominant = (B > R) & (B > G) & (B > 100)
    
    return blue_dominant.mean() > threshold

# -- Code Cell --
df_train['water_bg'] = df_train['filename'].apply(has_water_background)

# Check how many Earth birds have water background (these are the "hard" cases)
print(df_train[df_train['y'] == 0]['water_bg'].value_counts())
print(df_train[df_train['y'] == 1]['water_bg'].value_counts())

# -- Code Cell --
import pandas as pd

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

balanced_dfs = []

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

for bg in df_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 df_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)

# ---- CHECKS ----
print("Final global balance:")
print(df_balanced_final['y'].value_counts())

print("\nBalance per water_bg:")
print(df_balanced_final.groupby(['water_bg', 'y']).size())

# -- Code Cell --
df_train = df_balanced_final

# -- 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(df_train)
valid_dt = dataset_img_cls(df_valid, transform= transform_valid)
test_dt = dataset_img_cls(df_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=df_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 --


# -- 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)

# -- Code Cell --


# -- Code Cell --
