import json
import os
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset

# 1. Configuration ⚙️
CSV_FILE = "potato_gesture_dataset.csv"
EPOCHS = 120
BATCH_SIZE = 8
LEARNING_RATE = 0.003
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print(f"🚀 Training device: {DEVICE}")

# 2. Load & Prepare Dataset 📂
if not os.path.exists(CSV_FILE):
    # Check if user named it differently or check current folder
    csv_candidates = [f for f in os.listdir(".") if f.endswith(".csv")]
    if csv_candidates:
        CSV_FILE = csv_candidates[0]
        print(f"💡 Found CSV: {CSV_FILE}")
    else:
        raise FileNotFoundError(
            f"❌ Could not find {CSV_FILE}! Make sure it's in this folder."
        )

df = pd.read_csv(CSV_FILE)
print(f"📊 Loaded {len(df)} total gesture recordings!")

# Extract features and labels
X_raw = df.iloc[:, :-1].values.astype(np.float32)  # 90 numbers per sample
labels_raw = df.iloc[:, -1].values

# Convert text labels into numbers: 0, 1, 2, ...
unique_labels = sorted(list(set(labels_raw)))
label_to_idx = {name: i for i, name in enumerate(unique_labels)}
idx_to_label = {i: name for name, i in label_to_idx.items()}
y_raw = np.array([label_to_idx[l] for l in labels_raw], dtype=np.int64)

# Save the label mapping so your web app knows which number is which gesture!
with open("labels.json", "w") as f:
    json.dump(idx_to_label, f, indent=2)
print(f"🏷️ Classes ({len(unique_labels)}): {label_to_idx}")

# Reshape (N, 90) -> (N, 3 channels [X,Y,Z], 30 timesteps)
X_reshaped = X_raw.reshape(-1, 30, 3).transpose(0, 2, 1)

# Shuffle & Split into 80% Train, 20% Validation
indices = np.arange(len(df))
np.random.seed(42)
np.random.shuffle(indices)

split = int(0.8 * len(df))
train_idx, val_idx = indices[:split], indices[split:]


class MotionDataset(Dataset):

    def __init__(self, X, y):
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.long)

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

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]


train_loader = DataLoader(
    MotionDataset(X_reshaped[train_idx], y_raw[train_idx]),
    batch_size=BATCH_SIZE,
    shuffle=True,
)
val_loader = DataLoader(
    MotionDataset(X_reshaped[val_idx], y_raw[val_idx]), batch_size=BATCH_SIZE
)


# 3. The "Potato-Brain" Architecture (Micro 1D-CNN) 🧠🥔
class PotatoGestureNet(nn.Module):

    def __init__(self, num_classes):
        super(PotatoGestureNet, self).__init__()
        # 3 input channels (Ax, Ay, Az)
        self.conv1 = nn.Conv1d(
            in_channels=3, out_channels=16, kernel_size=5, padding=2
        )
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool1d(kernel_size=2)  # 30 -> 15

        self.conv2 = nn.Conv1d(
            in_channels=16, out_channels=32, kernel_size=3, padding=1
        )
        self.global_pool = nn.AdaptiveAvgPool1d(1)  # Flattens time axis cleanly

        self.fc = nn.Linear(32, num_classes)

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.relu(self.conv2(x))
        x = self.global_pool(x).squeeze(-1)
        x = self.fc(x)
        return x


model = PotatoGestureNet(num_classes=len(unique_labels)).to(DEVICE)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

# 4. Fast Training Loop ⏱️🏃‍♂️
print("\n🏁 Starting Training...")
for epoch in range(1, EPOCHS + 1):
    model.train()
    total_loss = 0
    correct = 0
    total = 0

    for batch_x, batch_y in train_loader:
        batch_x, batch_y = batch_x.to(DEVICE), batch_y.to(DEVICE)

        optimizer.zero_grad()
        outputs = model(batch_x)
        loss = criterion(outputs, batch_y)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total += batch_y.size(0)
        correct += predicted.eq(batch_y).sum().item()

    # Quick validation check every 5 epochs
    if epoch % 5 == 0 or epoch == EPOCHS:
        model.eval()
        val_correct, val_total = 0, 0
        with torch.no_grad():
            for bx, by in val_loader:
                bx, by = bx.to(DEVICE), by.to(DEVICE)
                preds = model(bx).argmax(dim=1)
                val_total += by.size(0)
                val_correct += preds.eq(by).sum().item()

        val_acc = (
            (val_correct / val_total) * 100 if val_total > 0 else 100.0
        )  # In case val set is empty
        train_acc = (correct / total) * 100
        print(
            f"Epoch [{epoch:02d}/{EPOCHS}] | Loss: {total_loss/len(train_loader):.4f} | Train Acc: {train_acc:.1f}% | Val Acc: {val_acc:.1f}%"
        )

print("\n🎉 Training Complete!")

# 5. Save PyTorch Model 💾
torch.save(model.state_dict(), "gesture_model.pth")
print("✅ Saved weights to 'gesture_model.pth'")

# 6. Export to ONNX for Browser / Edge Deployment 🌐
try:
    model.eval().cpu()
    dummy_input = torch.randn(1, 3, 30)  # [1 sample, 3 axes, 30 timesteps]

    torch.onnx.export(
        model,
        dummy_input,
        "gesture_model.onnx",
        input_names=["accel_input"],
        output_names=["probabilities"],
        dynamic_axes={"accel_input": {0: "batch_size"}},
        opset_version=12,
    )

    onnx_size = os.path.getsize("gesture_model.onnx") / 1024
    print(f"📦 Exported 'gesture_model.onnx' -> Size: {onnx_size:.1f} KB! 🥔👑")
except Exception as e:
    print(
        f"⚠️ Note: Run 'pip install onnx' to export ONNX. (PyTorch model saved fine!)"
    )