# -- Code Cell --
import os
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

from torch.utils.data import Dataset, DataLoader, random_split
from torchvision import datasets, transforms
from PIL import Image
import matplotlib.pyplot as plt
device = 'cpu'

# -- Code Cell --
class EmojiTrainDataset(Dataset):
    def __init__(self, size=128):
        self.image_dir = "images/train"
        self.mask_dir = "images/train/masks"

        self.images = os.listdir(self.image_dir)
        self.size = size

        self.img_tf = transforms.Compose([
            transforms.Resize((size, size)),
            transforms.ToTensor(),
        ])

        self.mask_tf = transforms.Resize((size, size))

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        img_name = self.images[idx]

        img_path = os.path.join(self.image_dir, img_name)
        mask_path = os.path.join(self.mask_dir, img_name)

        img = Image.open(img_path).convert("RGB")
        mask = Image.open(mask_path).convert("L")

        img = self.img_tf(img)

        mask = self.mask_tf(mask)
        mask = np.array(mask)
        mask = (mask > 0).astype(np.float32)
        mask = torch.tensor(mask).unsqueeze(0)

        return img, mask

# -- Code Cell --
train_dataset = EmojiTrainDataset(size=128)
train_size = int(0.8 * len(train_dataset))
val_size = len(train_dataset) - train_size
train_ds, val_ds = random_split(train_dataset, [train_size, val_size])

train_loader = DataLoader(train_ds, batch_size=8, shuffle=True)
val_loader   = DataLoader(val_ds, batch_size=8, shuffle=False)

# -- Code Cell --
class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.conv(x)


class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1):
        super().__init__()
        # Encoder
        self.enc1 = DoubleConv(in_channels, 64)
        self.pool1 = nn.MaxPool2d(2)

        self.enc2 = DoubleConv(64, 128)
        self.pool2 = nn.MaxPool2d(2)

        self.enc3 = DoubleConv(128, 256)
        self.pool3 = nn.MaxPool2d(2)

        self.enc4 = DoubleConv(256, 512)
        self.pool4 = nn.MaxPool2d(2)

        self.bottleneck = DoubleConv(512, 1024)

        # Decoder
        self.up4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
        self.dec4 = DoubleConv(1024, 512)

        self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.dec3 = DoubleConv(512, 256)

        self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.dec2 = DoubleConv(256, 128)

        self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.dec1 = DoubleConv(128, 64)

        self.out_conv = nn.Conv2d(64, out_channels, kernel_size=1)

    def forward(self, x):
        # Encoder
        c1 = self.enc1(x)
        p1 = self.pool1(c1)

        c2 = self.enc2(p1)
        p2 = self.pool2(c2)

        c3 = self.enc3(p2)
        p3 = self.pool3(c3)

        c4 = self.enc4(p3)
        p4 = self.pool4(c4)

        bn = self.bottleneck(p4)

        # Decoder with skips
        u4 = self.up4(bn)
        x4 = torch.cat([u4, c4], dim=1)
        d4 = self.dec4(x4)

        u3 = self.up3(d4)
        x3 = torch.cat([u3, c3], dim=1)
        d3 = self.dec3(x3)

        u2 = self.up2(d3)
        x2 = torch.cat([u2, c2], dim=1)
        d2 = self.dec2(x2)

        u1 = self.up1(d2)
        x1 = torch.cat([u1, c1], dim=1)
        d1 = self.dec1(x1)

        logits = self.out_conv(d1)          # (B, 1, H, W)
        probs  = torch.sigmoid(logits)      # for BCE loss / binary seg

        return probs


# -- Code Cell --
bce_loss = nn.BCELoss()
model = UNet(in_channels=3, out_channels=1).to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-4)

def train_one_epoch(loader):
    model.train()
    running_loss = 0.0
    

    for imgs, masks in loader:
        imgs, masks = imgs.to(device), masks.to(device)

        optimizer.zero_grad()
        preds = model(imgs)
        loss = bce_loss(preds, masks)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        

    return running_loss / len(loader)


def validate(loader):
    model.eval()
    val_loss = 0.0
    val_dice = 0.0
    with torch.no_grad():
        for imgs, masks in loader:
            imgs, masks = imgs.to(device), masks.to(device)
            preds = model(imgs)
            loss = bce_loss(preds, masks)
            val_loss += loss.item()
            val_dice += dice_coeff(preds, masks)
    return val_loss / len(loader), val_dice / len(loader)


# -- Code Cell --
num_epochs = 5

train_history = {"loss": [], "dice": []}
val_history   = {"loss": [], "dice": []}

for epoch in range(1, num_epochs + 1):
    train_loss = train_one_epoch(train_loader)
    val_loss= validate(val_loader)

    train_history["loss"].append(train_loss)
    val_history["loss"].append(val_loss)

    print(
        f"Epoch {epoch:02d} | "
        f"Train Loss: {train_loss:.4f} | "
    )

# -- Code Cell --
import numpy
print(numpy.__version__)

# -- Code Cell --
!pip uninstall numpy -y
!pip install numpy==1.26.4

# -- Code Cell --
