# %%
#EMOJI SEGMENTATION
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import pandas as pd
import numpy as np


# %%
import torch
print(torch.__version__)
print(torch.cuda.is_available())
print(torch.cuda.get_device_name())

# %%
class EmojiDataset(Dataset):
    def __init__(self, csv_path, transform = None):
        self.df = pd.read_csv(csv_path) # citim datasetul
        self.transform = transform
    def __len__(self):
        return len(self.df)
    def __getitem__(self, idx):
        row = self.df.iloc[idx]#ia randul din CSV
        img = Image.open(row['Path']).convert('RGB')#citeste imaginea in RGB
        mask = Image.open(row['Mask']).convert('L')#citeste imaginea in GRAYSCALE
        
        if self.transform: 
            img = self.transform(img)
            
        mask = mask.resize((256, 256), Image.NEAREST)    
        mask = transforms.ToTensor()(mask) # converteste masca in tensor din valori 0-255 devine 0.0-1.0 shape [1,H,W]
        mask = (mask>0.5).float() # orice este peste 0.5 devine 1.0 si restul 0.0
        return img,mask

# %%
transform = transforms.Compose([
    transforms.Resize((256, 256)), #resize la 256x256
    transforms.ToTensor(), #Imagine PIL (HxWxC, 0-255) -> (CxHxW, 0.0-1.0)
    transforms.Normalize([0.485, 0.456, 0.406], 
                         [0.229, 0.224, 0.225]) #normalizare
])

# %%
class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            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 self.out_conv(d1)

# %%
dataset = EmojiDataset('train.csv', transform=transform)
loader = DataLoader(dataset, batch_size=4, shuffle=True)

# %%
device = torch.device('cuda')

model = UNet().to(device)
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(20):
    model.train()
    total_loss = 0
    for imgs, mask in loader:
        
        imgs = imgs.to(device)     
        mask = mask.to(device)
        
        optimizer.zero_grad()
        output = model(imgs)
        loss = criterion(output,mask)
        loss.backward()
        optimizer.step()
        total_loss +=loss.item()
    print(f"Epochul {epoch+1}, loss: {total_loss/len(loader):.2f}")
        

# %%
import os

test_df = pd.read_csv('test.csv')
model.eval()

os.makedirs('predictions', exist_ok=True)

with torch.no_grad():
    for _, row in test_df.iterrows():
        img = Image.open(row['Path']).convert('RGB')
        original_size = img.size  # (W, H) — ca să dai resize înapoi
        
        img_t = transform(img).unsqueeze(0).to(device)  # [1, 3, 256, 256]
        
        output = model(img_t)                    # [1, 1, 256, 256] logits
        pred = torch.sigmoid(output)             # logits → probabilități
        mask = (pred > 0.5).float() * 255        # binarizare → 0 sau 255
        
        # resize înapoi la dimensiunea originală
        mask_img = Image.fromarray(mask[0, 0].cpu().byte().numpy())
        mask_img = mask_img.resize(original_size, Image.NEAREST)
        
        mask_img.save(f"predictions/{row['SampleID']}_mask.png")

# submission CSV
submission = pd.DataFrame({
    'SampleID': test_df['SampleID'],
    'PredictedMask': [f"predictions/{sid}_mask.png" for sid in test_df['SampleID']]
})
submission.to_csv('submission.csv', index=False)
print("Done!")

# %%
import zipfile

with zipfile.ZipFile('submission.zip', 'w') as z:
    z.write('submission.csv')
    for sid in test_df['SampleID']:
        z.write(f"predictions/{sid}_mask.png")

print("Done! Trimite submission.zip")

# %%
import matplotlib.pyplot as plt

img, mask = dataset[0]
fig, (ax1, ax2) = plt.subplots(1, 2)
ax1.imshow(img.permute(1, 2, 0).numpy() * 0.225 + 0.45)  # denormalizare aprox
ax1.set_title('Image')
ax2.imshow(mask[0].numpy(), cmap='gray')
ax2.set_title('Mask')
plt.show()

# %%



