# -- Code Cell --
import pandas as pd

# -- Code Cell --
train_df = pd.read_csv("train_data.csv")
test_df = pd.read_csv('test_data.csv')

# -- Code Cell --
train_df.head()

# -- Code Cell --
task1 = train_df.groupby('IDSample').size()
task1

# -- Code Cell --
cols = [c for c in train_df.columns if c!='IDSample']
cols = [c for c in cols if c!='FrameNumber']
cols = [c for c in cols if c!='Action']
cols = [c for c in cols if c!='Camera']
cols

# -- Code Cell --
data = train_df.groupby('IDSample')[cols].agg(list)
data

# -- Code Cell --
from sklearn.preprocessing import StandardScaler
scaler = StandardScaler()
train_df[cols] = scaler.fit_transform(train_df[cols])
test_df[cols] = scaler.transform(test_df[cols])

# -- Code Cell --
groups = train_df.groupby('IDSample')

# -- Code Cell --
groups_test = test_df.groupby('IDSample')

# -- Code Cell --
sequences = []
actions = []
cameras = []

for id_sample, group in groups:
    group = group.sort_values('FrameNumber')
    seq = group[cols].values
    action = group['Action'].iloc[0]
    camera = group['Camera'].iloc[0]
    
    sequences.append(seq)       
    actions.append(action)     
    cameras.append(camera)    

# -- Code Cell --
sequences_test = []
for id_sample, group in groups_test:
    group = group.sort_values('FrameNumber')
    seq = group[cols].values
    sequences_test.append(seq)       

# -- Code Cell --
# max_anterior = 0
# for i in range(len(sequences)):
#     len(sequences[i])
#     if len(sequences[i])>max_anterior:
#         max_max=len(sequences[i])
#         max_anterior = max_max
# max_max

# -- Code Cell --
max_max = 100

# -- Code Cell --
import numpy as np

# -- Code Cell --
for i in range(len(sequences)):
    if len(sequences[i]) > max_max:
        sequences[i] = sequences[i][:max_max]
    else:
        pad_size = max_max - len(sequences[i])
        sequences[i] = np.append(sequences[i], np.zeros((pad_size, 75)), axis=0)

# -- Code Cell --
for i in range(len(sequences_test)):
    if len(sequences_test[i])>max_max:
        sequences_test[i] = sequences_test[i][:max_max]
    else:    
        pad_size = max_max - len(sequences_test[i])
        sequences_test[i] = np.append(sequences_test[i], np.zeros((pad_size, 75)), axis=0)

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

# -- Code Cell --
(train_df['Action'][0],train_df['Camera'][0])

# -- Code Cell --
class TrainDataset(Dataset):
    def __init__(self, sequences, actions, cameras):
        self.sequences = sequences
        self.actions = actions
        self.cameras = cameras
    
    def __len__(self):
        return len(self.sequences)
    
    def __getitem__(self, index):
        seq = torch.tensor(self.sequences[index], dtype=torch.float32)
        action = torch.tensor(self.actions[index], dtype=torch.long)
        camera = torch.tensor(self.cameras[index], dtype=torch.long)
        return seq, action, camera

# -- Code Cell --
train_ds = TrainDataset(sequences, actions, cameras)
train_loader = DataLoader(train_ds, batch_size=64, num_workers=0, shuffle=True)

# -- Code Cell --
class LSTMACTION(nn.Module):
    def __init__(self):
        super().__init__()
        self.lstm = nn.LSTM(75, 128, num_layers=2, batch_first=True, dropout=0.3)
        self.dropout = nn.Dropout(0.3)
        self.fc = nn.Linear(128, 5)
    
    def forward(self, x):
        out, (hn, cn) = self.lstm(x)
        last_hidden = hn[-1]       
        return self.fc(self.dropout(last_hidden))

# -- Code Cell --
device = torch.device('cuda')

# -- Code Cell --
model_action = LSTMACTION()
model_action = model_action.to(device)
criterion_action = nn.CrossEntropyLoss()
optimizer_action = optim.Adam(model_action.parameters(), lr=1e-3)

# -- Code Cell --
for epoch in range(30):
    total_loss = 0
    for sequence, action, camera in train_loader:
        sequence = sequence.to(device)
        action = action.to(device)
        
        optimizer_action.zero_grad()
        output = model_action(sequence)
        loss = criterion_action(output, action)
        loss.backward()
        optimizer_action.step()
        total_loss += loss.item()
    print(f"Epoch {epoch} - Loss: {total_loss/len(train_loader):.4f}")

# -- Code Cell --
class LSTMCAMERA(nn.Module):
    def __init__(self):
        super().__init__()
        self.lstm = nn.LSTM(75, 128, num_layers=2, batch_first=True, dropout=0.3)
        self.dropout = nn.Dropout(0.3)
        self.fc = nn.Linear(128, 3)
    
    def forward(self, x):
        out, (hn, cn) = self.lstm(x)
        last_hidden = hn[-1]       
        return self.fc(self.dropout(last_hidden))

# -- Code Cell --
model_camera= LSTMACTION()
model_camera = model_camera.to(device)
criterion_camera = nn.CrossEntropyLoss()
optimizer_camera = optim.Adam(model_camera.parameters(), lr=1e-3)

# -- Code Cell --
for epoch in range(30):
    total_loss = 0
    for sequence, _, camera in train_loader:
        sequence = sequence.to(device)
        camera = camera.to(device)
        
        optimizer_camera.zero_grad()
        output = model_camera(sequence)
        loss = criterion_camera(output, camera)
        loss.backward()
        optimizer_camera.step()
        total_loss += loss.item()
    print(f"Epoch {epoch} - Loss: {total_loss/len(train_loader):.4f}")

# -- Code Cell --
class TestDataset(Dataset):
    def __init__(self, sequences):
        self.sequences = sequences
    
    def __len__(self):
        return len(self.sequences)
    
    def __getitem__(self, index):
        seq = torch.tensor(self.sequences[index], dtype=torch.float32)
        return seq

# -- Code Cell --
test_ds = TestDataset(sequences_test)
test = DataLoader(test_ds, num_workers=0, batch_size=64, shuffle=False)

# -- Code Cell --
preds_arr = []
for sequence in test:
    with torch.no_grad():
        sequence = sequence.to(device)
        pred = model_action(sequence)
        preds = torch.argmax(pred, dim=1)
        preds_arr.append(preds.cpu())
        preds_final = torch.cat(preds_arr)
preds_final

# -- Code Cell --
preds_final = preds_final.numpy()

# -- Code Cell --
preds_arr = []
for sequence in test:
    with torch.no_grad():
        sequence = sequence.to(device)
        pred = model_camera(sequence)
        preds = torch.argmax(pred, dim=1)
        preds_arr.append(preds.cpu())
        preds_final_camera = torch.cat(preds_arr)
preds_final_camera

# -- Code Cell --
preds_final_camera = preds_final_camera.numpy()

# -- Code Cell --
test_df['IDSample'].unique()

# -- Code Cell --
frame_counts = test_df.groupby('IDSample').size().reset_index(name='answer')

# -- Code Cell --
test_ids = test_df.groupby('IDSample').first().index

# -- Code Cell --
task1 = pd.DataFrame({
    "subtaskID": 1,
    "datapointID": frame_counts['IDSample'],
    "answer": frame_counts['answer']
})
task2 = pd.DataFrame({
    "subtaskID": 2,
    "datapointID": test_ids,
    "answer": preds_final
})

task3 = pd.DataFrame({
    "subtaskID": 3,
    "datapointID": test_ids,
    "answer": preds_final_camera
})
pd.concat([task1,task2,task3]).to_csv("subs.csv",index=False)