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

# -- Code Cell --
df_train = pd.read_csv("train_data.csv")
df_test = pd.read_csv("test_data.csv")
df_train.head(5)

# -- Code Cell --
import torch

def translate(text):
    mapping = {'A': 0, 'U': 1, 'G': 2, 'C': 3}
    indices = [mapping[c] for c in text]
    one_hot = torch.zeros(len(text), 4)
    for i, idx in enumerate(indices):
        one_hot[i, idx] = 1.0
    return one_hot  # shape: (seq_len, 4)

# -- Code Cell --
import ast
df_train["target_profile"] = df_train["target_profile"].apply(lambda x :ast.literal_eval('[' + x.replace(' ', ', ') + ']'))

# -- Code Cell --
from torch.utils.data import Dataset, DataLoader
class data_idk(Dataset):
    def __init__(self, df, test = False):
        super().__init__()
        self.x = df["sequence"]
        self.test = test
        if self.test == False:
            self.y = df["target_profile"]
    def __len__(self):
        return len(self.x)
    def __getitem__(self, index):
        if self.test == False:
            return translate(self.x[index]), torch.tensor(self.y[index]).float()
        return translate(self.x[index])
full_dt = data_idk(df_train)
test_dt = data_idk(df_test, test = True)

# -- Code Cell --
from torch.utils.data import random_split
train_len = int(0.8 * len(full_dt))
valid_len = len(full_dt) - train_len
train_dt, valid_dt = random_split(full_dt, [train_len, valid_len])

# -- Code Cell --
from torch.nn.utils.rnn import pad_sequence
def collate_fn(batch):
    X, y = zip(*batch)
    X = pad_sequence(X, batch_first= True, padding_value= 0)
    y = pad_sequence(y, batch_first= True, padding_value= -1)
    return X, y
def collate_fn_tst(batch):
    X= batch
    X = pad_sequence(X, batch_first= True, padding_value=0)
    return X
train_loader = DataLoader(train_dt, batch_size= 16, shuffle= True, collate_fn= collate_fn)
valid_loader = DataLoader(valid_dt, batch_size= 16, shuffle= True, collate_fn= collate_fn)
test_loader = DataLoader(test_dt, batch_size= 16, collate_fn= collate_fn_tst)

# -- Code Cell --
import torch.nn as nn
class model_cool(nn.Module):
    def __init__(self, embed_dim=4, hidden_dim=64):
        super().__init__()
        self.lstm = nn.LSTM(embed_dim, hidden_dim, num_layers=3,
                            batch_first=True, bidirectional=True, dropout=0.3)
        self.cls = nn.Sequential(
            nn.Linear(hidden_dim * 2, 32),
            nn.ReLU(),
            nn.Linear(32, 1)
        )
    def forward(self, x):
   
        x = x.float() 
        x, _ = self.lstm(x)
        return self.cls(x)

# -- Code Cell --
def masked_pearson_loss(pred, target, ignore_index=-1):
    mask = target != ignore_index
    total, count = 0.0, 0
    for i in range(pred.shape[0]):
        m = mask[i]
        if m.sum() < 2:
            continue
        p = pred[i][m]
        t = target[i][m]
        p_ = p - p.mean()
        t_ = t - t.mean()
        corr = (p_ * t_).sum() / (p_.norm() * t_.norm() + 1e-8)
        total += (1 - corr)
        count += 1
    return total / count if count > 0 else pred.sum() * 0

# -- Code Cell --
model = model_cool()
optim = torch.optim.Adam(model.parameters())
loss_fn = masked_pearson_loss
epochs = 10

# -- Code Cell --
for batch in train_loader:
    print(model(batch[0]).shape)
    break

# -- Code Cell --
from tqdm import tqdm
import torch

best_mae = float('inf')

for epoch in range(epochs):
    # --- Train ---
    model.train()
    total_loss = 0

    pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs} [Train]")
    for x, y in pbar:
        optim.zero_grad()
        pred = model(x)
        loss = loss_fn(pred.squeeze(-1), y)
        loss.backward()
        optim.step()

        total_loss += loss.item()
        pbar.set_postfix(mae=f"{loss.item():.4f}")

    train_mae = total_loss / len(train_loader)

    # --- Validation ---
    model.eval()
    val_loss = 0
    with torch.no_grad():
        pbar_val = tqdm(valid_loader, desc=f"Epoch {epoch+1}/{epochs} [Val]  ")
        for x, y in pbar_val:
            pred = model(x)
            loss = loss_fn(pred.squeeze(-1), y)
            val_loss += loss.item()
            pbar_val.set_postfix(mae=f"{loss.item():.4f}")

    val_mae = val_loss / len(valid_loader)

    if val_mae < best_mae:
        best_mae = val_mae
        torch.save(model.state_dict(), "best_model.pt")
        tag = " ← saved"
    else:
        tag = ""
    print(f"Epoch {epoch+1}/{epochs} | Train MAE: {train_mae:.4f} | Val MAE: {val_mae:.4f} | Best: {best_mae:.4f}{tag}")

# -- Code Cell --
model.load_state_dict(torch.load("best_model.pt"))
model.eval()

all_preds = []

with torch.no_grad():
    for x in tqdm(test_loader, desc="Test inference"):
        pred = model(x).squeeze(-1)       
        pad_mask = (x != 0)
        while pad_mask.dim() > pred.dim():
            pad_mask = pad_mask.any(dim=-1)
        pad_mask = (~pad_mask)
        pred[pad_mask] = -11     
        all_preds.append(pred)


# -- Code Cell --
all_preds[0][all_preds[0] != -11].shape

# -- Code Cell --
test_pred = []
for i in tqdm(range(len(all_preds))):
    for j in range(all_preds[i].shape[0]):
        mask = all_preds[i][j] != -11
        test_pred.append(all_preds[i][j][mask].numpy().tolist())

# -- Code Cell --
df_subi = pd.read_csv("sample_output.csv")
df_subi.head(5)

# -- Code Cell --
df_subi["answer"] = test_pred
df_subi["answer"] = df_subi["answer"].apply(lambda x: " ".join(f"{v:.4f}" for v in x))


# -- Code Cell --
df_subi.to_csv('subi.csv', index= False)

# -- Code Cell --
