# -- 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 --
frames = pd.unique(df_test["IDSample"])
y_pred_1 = []
for id in frames:
    #print(id)
    y_pred_1.append(len(df_test[df_test["IDSample"] == id]))

# -- Code Cell --
from torch.utils.data import Dataset
import torch
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler


class BaselineDataset(Dataset):
    def __init__(self, df, test=False):
        super().__init__()

        self.test = test

        if not self.test:
            self.y = df["Action"]
            self.df = df.drop(["Action", "Camera"], axis=1)
        else:
            self.df = df.copy()

        # Unique sequences
        self.frames = pd.unique(self.df["IDSample"])

        # Feature columns (ONLY raw joints)
        self.feature_cols = [c for c in self.df.columns 
                             if c not in ["IDSample", "FrameNumber"]]

        self.data = []

        for frame_id in self.frames:
            specific_rows = self.df[self.df["IDSample"] == frame_id]

            # -------------------------------
            # RAW FEATURES ONLY
            # -------------------------------
            X = specific_rows[self.feature_cols].to_numpy()

            # -------------------------------
            # Optional normalization (keep it)
            # -------------------------------
            scl = StandardScaler()
            X = scl.fit_transform(X)

            # -------------------------------
            # Store sequence
            # -------------------------------
            if not self.test:
                labels = self.y[specific_rows.index].to_numpy()
                y = labels[0]
                self.data.append((X, y))
            else:
                self.data.append(X)

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

    def __getitem__(self, idx):
        if self.test:
            X = self.data[idx]
            return torch.tensor(X, dtype=torch.float32)
        else:
            X, y = self.data[idx]
            return (
                torch.tensor(X, dtype=torch.float32),
                torch.tensor(y, dtype=torch.long)
            )

# -- Code Cell --
from torch.utils.data import random_split
full_dt =BaselineDataset(df_train)
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])
test_dt = BaselineDataset(df_test, test= True)

# -- Code Cell --
full_dt[0][0].shape

# -- 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)
    y =torch.tensor(y)
    return X, y

# -- Code Cell --
from torch.utils.data import DataLoader
train_loader= DataLoader(train_dt, batch_size= 16, collate_fn= collate_fn, shuffle= True)
valid_loader= DataLoader(valid_dt, batch_size= 16, collate_fn= collate_fn, shuffle= True)


# -- Code Cell --
import torch.nn as nn
class model_cool(nn.Module):
    def __init__(self, input_size, lstm_hidden, num_layers, num_classes, dropout=0.3):
        super(model_cool, self).__init__()
        
        self.fc_features = nn.Linear(input_size, 126)
        self.bn = nn.BatchNorm1d(126)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)
        self.lstm = nn.LSTM(
            input_size=126,
            hidden_size=lstm_hidden,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0
        )
        self.ln = nn.LayerNorm(lstm_hidden)
        self.classifier = nn.Linear(lstm_hidden, num_classes)
        self.layer_combo = nn.Sequential(
            nn.Linear(55, 23),
            nn.ReLU(),
            nn.Linear(23, 5)
        )
    def forward(self, x):
        B, T, F = x.shape
        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)
        out = self.classifier(out)
        return out

# -- Code Cell --
model = model_cool(75, lstm_hidden=512, num_layers=2,
                          num_classes=5)
epochs = 100
optim = torch.optim.Adam(model.parameters(), lr= 8e-5, weight_decay= 1e-4)
loss_fn = nn.CrossEntropyLoss()

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

# Check for GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Move model to device
model = model.to(device)

# Training loop with tqdm
for epoch in range(epochs):
    model.train()
    train_loss = 0
    correct = 0
    total = 0

    # tqdm for training
    train_iter = tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs} [Train]", leave=False)
    for X_batch, y_batch in train_iter:
        X_batch = X_batch.to(device)
        y_batch = y_batch.to(device)

        optim.zero_grad()
        outputs = model(X_batch)
        loss = loss_fn(outputs, y_batch)
        loss.backward()
        optim.step()

        train_loss += loss.item() * X_batch.size(0)
        _, predicted = torch.max(outputs, 1)
        correct += (predicted == y_batch).sum().item()
        total += y_batch.size(0)

        # Update tqdm postfix with running metrics
        train_iter.set_postfix(loss=train_loss/total, acc=correct/total)

    train_loss /= total
    train_acc = correct / total

    # Validation
    model.eval()
    val_loss = 0
    val_correct = 0
    val_total = 0

    val_iter = tqdm(valid_loader, desc=f"Epoch {epoch+1}/{epochs} [Val]", leave=False)
    with torch.no_grad():
        for X_batch, y_batch in val_iter:
            X_batch = X_batch.to(device)
            y_batch = y_batch.to(device)

            outputs = model(X_batch)
            loss = loss_fn(outputs, y_batch)

            val_loss += loss.item() * X_batch.size(0)
            _, predicted = torch.max(outputs, 1)
            val_correct += (predicted == y_batch).sum().item()
            val_total += y_batch.size(0)

            val_iter.set_postfix(loss=val_loss/val_total, acc=val_correct/val_total)

    val_loss /= val_total
    val_acc = val_correct / val_total

    print(f"Epoch [{epoch+1}/{epochs}] "
          f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | "
          f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f}")

# -- Code Cell --
from torch.nn.utils.rnn import pad_sequence
def collate_fn_test(batch):
    X = batch
    X = pad_sequence(X, batch_first= True)
    return X

# -- Code Cell --
test_loader= DataLoader(test_dt, batch_size= 16, collate_fn= collate_fn_test)


# -- Code Cell --
y_pred = []
for batch in tqdm(test_loader):
    X = batch.to(device)
    y_hat = model(X)
    y_pred.extend(torch.argmax(y_hat, dim= 1))
y_pred_2 = [x.item() for x in y_pred]

# -- Code Cell --
df_train.head(5)

# -- Code Cell --
import pandas as pd
from sklearn.preprocessing import StandardScaler

# -------------------------------
# Load data
# -------------------------------
train_df = pd.read_csv("train_data.csv")
test_df = pd.read_csv("test_data.csv")

# -------------------------------
# Keep ONLY raw joint columns
# -------------------------------
feat_cols = [c for c in train_df.columns if c.startswith("J")]

# -------------------------------
# Aggregate per sequence (VERY SIMPLE)
# -------------------------------
def extract_baseline(df):
    return df.groupby("IDSample")[feat_cols].agg(['mean', 'std']).fillna(0)

X = extract_baseline(train_df)
X_test = extract_baseline(test_df)

# Labels (unchanged)
y = train_df.groupby("IDSample")["Camera"].first()
test_ids = X_test.index

# -------------------------------
# Flatten column names
# -------------------------------
X.columns = ['_'.join(col) for col in X.columns]
X_test.columns = ['_'.join(col) for col in X_test.columns]

# Align columns
common_cols = sorted(set(X.columns).intersection(set(X_test.columns)))
X = X[common_cols].copy()
X_test = X_test[common_cols].copy()

# -------------------------------
# Scaling (keep this!)
# -------------------------------
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)
X_test_scaled = scaler.transform(X_test)

# -- Code Cell --
from sklearn.preprocessing import LabelEncoder
lbl = LabelEncoder()
y = lbl.fit_transform(y)

# -- Code Cell --
from xgboost import XGBClassifier

model = XGBClassifier(
    n_estimators=300,
    max_depth=20,
    learning_rate=0.05,
    class_weight='balanced',
    random_state=42
)
model.fit(X_scaled, y)
final_preds = model.predict(X_test_scaled)
final_labels =lbl.inverse_transform(final_preds)


# -- Code Cell --
df_pred_1 = pd.DataFrame({
    "subtaskID" : 1,
    "datapointID" : frames,
    "answer" : y_pred_1
})
df_pred_2 = pd.DataFrame({
    "subtaskID" : 2,
    "datapointID" : frames,
    "answer" : y_pred_2
})
df_pred_3 = pd.DataFrame({
    "subtaskID" : 3,
    "datapointID" : frames,
    "answer" : final_labels
})

# -- Code Cell --
df_pred =pd.concat([df_pred_1, df_pred_2, df_pred_3])

# -- Code Cell --
df_pred.to_csv("subi.csv", index= False)

# -- Code Cell --
