# -- 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 --
train_df = train_df.sort_values(['IDSample', 'FrameNumber']).reset_index(drop=True)
test_df = test_df.sort_values(['IDSample', 'FrameNumber']).reset_index(drop=True)

# -- Code Cell --
from numpy.linalg import norm

cos_cols = [f'cos_J{j}' for j in range(1, 26)]

for df in [train_df, test_df]:
    for j in range(1, 26):
        x, y, z = f'J{j}X', f'J{j}Y', f'J{j}Z'
        curr = df[[x, y, z]].values
        prev = df.groupby('IDSample')[[x, y, z]].shift(1).fillna(0).values
        
        dot = (curr * prev).sum(axis=1)
        norm_curr = norm(curr, axis=1) 
        norm_prev = norm(prev, axis=1) 
        df[f'cos_J{j}'] = dot / (norm_curr * norm_prev)

# -- Code Cell --
#Explicatie: ca inainte mai intai iau cos_cols, 25..., parsez dataframeul iau X,Y,Z cu fstring, cols era mai greu de prasat(se putea), 
# current iau valorile t0, apoi previous sunt valorile t1(cu shift(1)), fillna(0) PRIMUL ARE NAN, NAN, NAN !!!!
# fac cosine similarity si bag in cos_cols

# -- Code Cell --
vel_cols = [f'v_{c}' for c in cols]
acc_cols = [f'a_{c}' for c in cols]

# -- Code Cell --
for df in [train_df, test_df]:
    df[vel_cols] = df.groupby('IDSample')[cols].diff().fillna(0)
    df[acc_cols] = df.groupby('IDSample')[vel_cols].diff().fillna(0)

# -- Code Cell --
all_feature_cols = cols + vel_cols + acc_cols + cos_cols

# -- Code Cell --
from sklearn.preprocessing import StandardScaler
scaler = StandardScaler()
train_df[all_feature_cols] = scaler.fit_transform(train_df[all_feature_cols])
test_df[all_feature_cols] = scaler.transform(test_df[all_feature_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[all_feature_cols].values
    action = group['Action'].iloc[0]
    camera = group['Camera'].iloc[0]-1
    
    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[all_feature_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 --
len(all_feature_cols)

# -- 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, 250)), 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, 250)), 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(250, 256, num_layers=3, batch_first=True, dropout=0.3, bidirectional=True)
        self.dropout = nn.Dropout(0.3)
        self.fc = nn.Linear(512, 5)
    
    def forward(self, x):
        out, (hn, cn) = self.lstm(x)
        last_hidden = torch.cat([hn[-2], hn[-1]], dim=1)
        return self.fc(self.dropout(last_hidden))

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

# -- Code Cell --
class model_cool(nn.Module):
    def __init__(self, input_size=250, lstm_hidden=256, num_layers=3, num_classes=5, dropout=0.3):
        super().__init__()
        self.fc_features = nn.Linear(input_size, 128)
        self.bn = nn.BatchNorm1d(128)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)
        self.lstm = nn.LSTM(128, lstm_hidden, num_layers=num_layers,
                            batch_first=True, dropout=dropout, bidirectional=True)
        self.ln = nn.LayerNorm(lstm_hidden * 2)
        self.classifier = nn.Linear(lstm_hidden * 2, num_classes)
    
    def forward(self, x):
        x = self.fc_features(x)
        x = self.relu(x)
        x = self.bn(x.transpose(1, 2)).transpose(1, 2)
        x = self.dropout(x)
        out, _ = self.lstm(x)
        out = self.ln(out)
        out = out.mean(dim=1)
        return self.classifier(out)

# -- Code Cell --
# model_action = LSTMACTION()
model_action = model_cool(num_classes=5).to(device)
model_action = model_action.to(device)
criterion_action = nn.CrossEntropyLoss()
optimizer_action = optim.Adam(model_action.parameters(), lr=1e-3)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer_action, patience=5, factor=0.5)

# -- Code Cell --
for epoch in range(80):
    model_action.train()
    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()
    avg_loss = total_loss / len(train_loader)
    scheduler.step(avg_loss)
    print(f"Epoch {epoch} - Loss: {avg_loss:.4f} - LR: {optimizer_action.param_groups[0]['lr']:.6f}")

# -- Code Cell --
class LSTMCAMERA(nn.Module):
    def __init__(self):
        super().__init__()
        self.lstm = nn.LSTM(250, 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 = model_cool(num_classes=3).to(device)
criterion_camera = nn.CrossEntropyLoss()
optimizer_camera = optim.Adam(model_camera.parameters(), lr=1e-3)
scheduler_camera = optim.lr_scheduler.ReduceLROnPlateau(optimizer_camera, patience=5, factor=0.5)

# -- Code Cell --
for epoch in range(80):
    model_camera.train()
    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()
    avg_loss = total_loss / len(train_loader)
    scheduler_camera.step(avg_loss)
    print(f"Epoch {epoch} - Loss: {avg_loss:.4f} - LR: {optimizer_camera.param_groups[0]['lr']:.6f}")

# -- 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 + 1
})
pd.concat([task1,task2,task3]).to_csv("subs.csv",index=False)

# -- Code Cell --


# -- Code Cell --
