# %%
import os
import pandas as pd
from PIL import Image
import numpy as np
import cv2
import copy

# %%
train = pd.read_csv('train.csv')

# %%
train.head()

# %%
train['Effect'].unique()

# %%
imgs = []
for i in range(len(train)):
    path = os.path.join("./starting_kit/", f"{train['Path'][i]}")
    img = Image.open(path)
    imgs.append(img)

# %%
imgs[0]

# %%
import pandas as pd
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
from torchvision.transforms import functional as TF
import os

# %%
map_effect = {
    '-4': 0,
    '-2': 1,
    '-1': 2,
    '0': 3,
    '1': 4,
    '4': 5,
    'A': 6,
    'B': 7
}

# %%
def preprocess_symbol_pil(img_pil, out_size=48, pad=4):
    img = np.array(img_pil.convert("L"))

    _, binary = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

    ys, xs = np.where(binary > 0)
    x1, x2 = xs.min(), xs.max()
    y1, y2 = ys.min(), ys.max()
    crop = binary[y1:y2+1, x1:x2+1]

    h, w = crop.shape
    size = max(h, w) + 2 * pad
    canvas = np.zeros((size, size), dtype=np.uint8)
    y_off = (size - h) // 2
    x_off = (size - w) // 2
    canvas[y_off:y_off+h, x_off:x_off+w] = crop
    canvas = cv2.resize(canvas, (out_size, out_size), interpolation=cv2.INTER_AREA)

    return Image.fromarray(canvas)

# %%
class TrainDataset(Dataset):
    def __init__(self, path, transform=None):
        self.df = pd.read_csv(path)
        self.transform = transform

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

    def __getitem__(self, index):
        row = self.df.iloc[index]
        img_path = os.path.join("./starting_kit/", row["Path"])
        img = Image.open(img_path).convert("L")

        img = preprocess_symbol_pil(img, out_size=48, pad=4)

        if self.transform:
            img = self.transform(img)

        label = map_effect[str(row["Effect"])]
        return img, label

# %%
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])
])

# %%
train_ds=TrainDataset("./train.csv", transform=transform)

# %%
from torch.utils.data import random_split
train_size = int(0.8* len(train_ds))
valid_size = len(train_ds) - train_size
train_subset, valid_subset = random_split(train_ds, [train_size,valid_size])

# %%
valid_loader = DataLoader(valid_subset, batch_size=32, shuffle=False, num_workers=0, pin_memory=True)
train_loader = DataLoader(train_subset, batch_size=32, shuffle=True, num_workers=0, pin_memory=True)

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

# %%
class CONVBlock(nn.Module):
    def __init__(self, in_c, out_c):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_c),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2)
        )

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

# %%
class CustomModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            CONVBlock(1, 32),   # 48 -> 24
            CONVBlock(32, 64),  # 24 -> 12
            CONVBlock(64, 128), # 12 -> 6
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool2d((1, 1))
        )
        self.cls = nn.Sequential(
            nn.Flatten(),
            nn.Dropout(0.2),
            nn.Linear(128, 8)
        )

    def forward(self, x):
        x = self.features(x)
        x = self.cls(x)
        return x

# %%
model = CustomModel().to(device)

optimizer = optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
criterion = nn.CrossEntropyLoss()
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="max", factor=0.5, patience=3)

# %%
from sklearn.metrics import f1_score

# %%
best_f1 = -1
best_state = None
for epoch in range(30):
    model.train()
    total_loss = 0.0
    for img, label in train_loader:
        img = img.to(device)
        label = label.to(device)

        optimizer.zero_grad()
        output = model(img)
        loss = criterion(output, label)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    model.eval()
    all_labels = []
    all_preds = []

    with torch.no_grad():
        for img, label in valid_loader:
            img = img.to(device)
            label = label.to(device)

            output = model(img)
            pred = output.argmax(dim=1)

            all_preds.append(pred.cpu())
            all_labels.append(label.cpu())

    all_preds = torch.cat(all_preds).numpy()
    all_labels = torch.cat(all_labels).numpy()

    f1 = f1_score(all_labels, all_preds, average='macro')
    scheduler.step(f1)
    if f1 > best_f1:
        best_f1 = f1
        best_state = copy.deepcopy(model.state_dict())

    print(f"epoch {epoch+1} -- loss {total_loss/len(train_loader):.4f} -- f1 {f1:.4f}")

model.load_state_dict(best_state)

# %%
imgs = []
for i in range(len(train)):
    path = os.path.join("./starting_kit/", f"{train['Path'][i]}")
    img = Image.open(path).convert("L")
    imgs.append(img)

# %%
test = pd.read_csv('test.csv')

# %%
test.head()

# %%
all_sequences = []

for i in range(len(test)):
    path = os.path.join("./starting_kit/", test['datapointID'][i])
    img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)

    _, binary = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(binary, connectivity=8)

    boxes = []
    for t in range(1, num_labels):
        x = stats[t, cv2.CC_STAT_LEFT]
        y = stats[t, cv2.CC_STAT_TOP]
        w = stats[t, cv2.CC_STAT_WIDTH]
        h = stats[t, cv2.CC_STAT_HEIGHT]
        area = stats[t, cv2.CC_STAT_AREA]

        if area > 20:
            boxes.append((x, y, w, h))

    boxes = sorted(boxes, key=lambda b: b[0])

    sequence = []
    for (x, y, w, h) in boxes:
        crop = img[y:y+h, x:x+w]
        crop_pil = Image.fromarray(crop)
        crop_pil = preprocess_symbol_pil(crop_pil, out_size=48, pad=4)
        sequence.append(crop_pil)

    all_sequences.append(sequence)

# %%
retval, labels, stats, centroids = cv2.connectedComponentsWithStats(img)
boxes = []
for t in range(1,labels):
    x = stats[t,cv2.CC_STAT_LEFT]
    y = stats[t, cv2.CC_STAT_TOP]
    w = stats[t, cv2.CC_STAT_WIDTH]
    h = stats[t, cv2.CC_STAT_HEIGHT]
    area = stats[t, cv2.CC_STAT_AREA]
    
    if area>20:
        boxes.append(x,y,w,h)
    
    boxes = sorted(boxes, key=lambda b:b[0])
    
    sequence=[]
    for (x,y,w,h) in boxes:
        crop = img[y:y+h, x:x+w]
        crop_pil = Image.fromarray(crop)
        crop_pil = preprocess_symbol_pil(crop_pil)
        sequence.append(crop_pil)

# %%
all_sequences[0][0]

# %%
model.eval()
preds = []

for seq in all_sequences:
    seq_preds = []

    for crop_pil in seq:
        x = transform(crop_pil)
        x = x.unsqueeze(0).to(device)

        with torch.no_grad():
            out = model(x)
            pred = torch.argmax(out, dim=1).item()

        seq_preds.append(pred)

    preds.append(seq_preds)

# %%
rev_map = {0:'-4', 1:'-2', 2:'-1', 3:'0', 4:'1', 5:'4', 6:'A', 7:'B'}
pred_labels = []

for seq in preds:
    pred_labels.append([rev_map[p] for p in seq])

# %%
pred_labels

# %%
answers = []

for seq in pred_labels:
    total = 0
    current_seq = []

    for j in seq:
        if j in ['A', 'B']:
            effect = 0
        else:
            effect = int(j)

        total += effect
        current_seq.append(str(total))

    answers.append("|".join(current_seq))

# %%
answers[100]

# %%
subs = pd.DataFrame({
    'subtaskID':1,
    'datapointID':test['datapointID'],
    'answer':answers
}).to_csv('subs.csv',index=False)

# %%


# %%



