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

# -- Code Cell --
class DoubleConv(nn.Module):
    def __init__(self, in_chn, out_chn):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_chn,out_chn,kernel_size = 3, padding=1),
            nn.BathNorm2d(out_chn),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_chn,out_chn,kernel_size = 3, padding=1),
            nn.BathNorm2d(out_chn),
            nn.ReLU(inplace=True),
        )
    def forward(self,x):
        return self.conv(x)

# -- Code Cell --
class DoubleConv(nn.Module):
    def __init__(self, in_c, out_c):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_c, out_c, kernel_size = 3, padding=1),
            nn.BathNorm2d(out_c),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_c, out_c, kernel_size = 3, padding=1),
            nn.BathNorm2d(out_c),
            nn.ReLU(inplace=True),
        )
    def forward(self, x):
        return self.conv(x)

# -- Code Cell --
class UNET(nn.Module):
    def __init__(self, in_c = 3, out_c = 1):
        super().__init__()
        self.enc1 = DoubleConv(in_c,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)
        
        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 = nn.Conv2d(64, out_c, kernel_size = 1)
        
    def forward(self,x):
        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)
        
        bt = self.bottleneck(p4)
        
        u4 = self.up4(bt)
        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.up2(d2)
        x1 = torch.cat([u2,c2], dim=1)
        d1 = self.dec2(x1)
        
        logits = self.out(d1)
        probs = torch.sigmoid(logits)
        return probs

# -- Code Cell --
class UNET(nn.Module):
    def __init__(self, in_ch = 3, out_ch=1):
        super().__init__()
        self.enc1 = DoubleConv(in_ch, 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.bt = DoubleConv(512, 1024)
        
        self.up4 = nn.ConvTrasnpose2d(1024, 512, kernel_size = 2, stride=2)
        self.dec4 = DoubleConv(1024, 512)
        
        self.up3 = nn.ConvTrasnpose2d(512, 256, kernel_size = 2, stride=2)
        self.dec3 = DoubleConv(512, 256)
        
        self.up3 = nn.ConvTrasnpose2d(256, 128, kernel_size = 2, stride=2)
        self.dec3 = DoubleConv(256, 128)
        
        self.up3 = nn.ConvTrasnpose2d(128, 64, kernel_size = 2, stride=2)
        self.dec3 = DoubleConv(128, 64)
        
        self.out = nn.Conv2d(64, out_ch, kernel_size= 1)
        
    def forward(self,x):
        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.bt(p4)
        
        u4 = self.up4(bn)
        x4 = torch.cat([u4,c4], dim = 1)
        d4 = self.dec4(x4)
        
        u3 = self.up3(bn)
        x3 = torch.cat([u3,c3], dim = 1)
        d3 = self.dec3(x3)
        
        u2 = self.up2(bn)
        x2 = torch.cat([u2,c2], dim = 1)
        d2 = self.dec2(x2)
        
        u1 = self.up1(bn)
        x1 = torch.cat([u1,c1], dim = 1)
        d1 = self.dec1(x1)
        
        out = self.out(d1)
        probs = torch.sigmoid(out)
        return probs

# -- Code Cell --
class DoubleConv(nn.Module):
    def __init__(self, in_c, out_c):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_c, out_c, kernel_size = 3, padding= 1),
            nn.BatchNomr2d(out_c),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_c, out_c, kernel_size = 3, padding= 1),
            nn.BatchNomr2d(out_c),
            nn.ReLU(inplace=True),
        )
    def forward(self, x):
        return self.conv(x)

# -- Code Cell --


# -- Code Cell --


# -- Code Cell --
from torch.utils.data import random_split
train_ds = TrainDataset()
train_size = int(0.8*len(train_ds))
val_size = len(train_ds) - train_size
train_subset, val_subset = random_split(train_ds, [train_size,val_size])

# -- Code Cell --
train_loader = DataLoader(train_subset)
val_loader= DataLoader(val_subset)

# -- Code Cell --
for param in model.parameteres():
    param.requires_grad = False
optimizer = optim.Adam(model.parameters(), lr=1e-3)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size = 15, gamma=0.5)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, "min")

# -- Code Cell --
for epoch in range(10):
    model.train()
    total_loss = 0
    for x,y in train_loader:
        x = x.to(device)
        y = y.to(device)
        optimizer.zero_grad()
        output = model(x)
        loss = criterion(output, y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    model.eval()
    val_loss = 0
    all_preds = []
    all_labels = []
    with torch.no_grad():
        for x,y in val_loader:
            x = x.to(device)
            y = y.to(device)
            output = model(x)
            loss = criterion(output, y)
            preds = output.argmax(dim=1)
            
            all_preds.append(preds.cpu())
            all_labels.append(y.cpu())
    all_labels = torch.cat(all_labels)
    all_preds = torch.cat(all_preds)
    f1 = f1_score(all_labels.numpy(), all_preds.numpy(), average = "macro")
    scheduler.step(f1)

# -- Code Cell --
model.eval()
val_loss = 0.0
correct = 0
total = 0
with torch.no_grad():
    for imgs, labels in val_loader:
        imgs = imgs.to(device)
        labels = labels.to(device)

        outputs = resnet(imgs)
        loss = criterion(outputs, labels)

        val_loss += loss.item()

        preds = outputs.argmax(dim=1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)

val_acc = correct / total

scheduler.step(val_acc)