from glob import glob
from utils import *
from riftnet import RiftNet

import matplotlib.pyplot as plt
from tqdm import tqdm

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

from sklearn.model_selection import train_test_split
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

c:\Users\user\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\utils\_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.
  warnings.warn(

file_paths = glob("Bluetooth_Datasets/Dataset 250 Msps/*/*/*/*")

unknown_users = {"4s_013004004984503_oguz_guler", "Note2_356261053336200_ismet_buyukkilic", "XperiaM5_354188070809491_firat_vural"}

file_paths_filtered = []
user_names_filtered = []

for file_path in file_paths:
    user = format_model_name(file_path.split("\\")[2:4])
    if user not in unknown_users:
        file_paths_filtered.append(file_path)
        user_names_filtered.append(user)

# Create dataset with filtered data
dataset = RFFDataset(file_paths=file_paths_filtered, labels=user_names_filtered)

def train_riftnet(dataset, analytic_fn, num_classes, epochs):
    encoded_dataset, le = encode_labels(dataset)

    # Stratified split: 80% train, 20% test
    train_data, test_data = train_test_split(
        encoded_dataset, test_size=0.2, stratify=[d[1] for d in encoded_dataset], random_state=0
    )

    train_dataset = IQDataset(train_data, analytic_fn)
    test_dataset = IQDataset(test_data, analytic_fn)

    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = RiftNet(num_classes=num_classes).to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-4)

    losses, accuracies, test_accuracies = [], [], []

    for epoch in range(epochs):
        model.train()
        total_loss, correct, total = 0.0, 0, 0
        with tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", ncols=100) as pbar:
            for x_long, x_short, y in pbar:
                x_long, x_short, y = x_long.to(device), x_short.to(device), y.to(device)
                optimizer.zero_grad()
                out = model(x_long, x_short)
                loss = criterion(out, y)
                loss.backward()
                optimizer.step()

                total_loss += loss.item() * x_long.size(0)
                _, preds = out.max(1)
                correct += (preds == y).sum().item()
                total += y.size(0)

                avg_loss = total_loss / total
                acc = correct / total
                pbar.set_postfix(loss=avg_loss, acc=acc)

        losses.append(avg_loss)
        accuracies.append(acc)

        # ----- Evaluate on test set -----  
        model.eval()
        correct_test, total_test = 0, 0
        y_true, y_pred = [], []  # To store true and predicted labels for confusion matrix
        with torch.no_grad():
            for x_long, x_short, y in test_loader:
                x_long, x_short, y = x_long.to(device), x_short.to(device), y.to(device)
                out = model(x_long, x_short)
                _, preds = out.max(1)
                correct_test += (preds == y).sum().item()
                total_test += y.size(0)
                
                # Collect true and predicted labels
                y_true.extend(y.cpu().numpy())
                y_pred.extend(preds.cpu().numpy())
        
        test_acc = correct_test / total_test
        test_accuracies.append(test_acc)
        print(f"Test acc: {test_acc:.2f} %")
    
        # Save best model
        if epoch == 0 or test_acc > max(test_accuracies[:-1]):
            torch.save(model.state_dict(), 'best_model_state_dict.pth')
            print(f"✅ Saved best model at epoch {epoch+1} with test acc: {test_acc:.2%}")

    model.load_state_dict(torch.load('best_model_state_dict.pth'))
    # ----- Evaluate on test set -----  
    model.eval()
    correct_test, total_test = 0, 0
    y_true, y_pred = [], []  # To store true and predicted labels for confusion matrix
    with torch.no_grad():
        for x_long, x_short, y in test_loader:
            x_long, x_short, y = x_long.to(device), x_short.to(device), y.to(device)
            out = model(x_long, x_short)
            _, preds = out.max(1)
            correct_test += (preds == y).sum().item()
            total_test += y.size(0)
            
            # Collect true and predicted labels
            y_true.extend(y.cpu().numpy())
            y_pred.extend(preds.cpu().numpy())
    
    test_acc = correct_test / total_test
    print(f"Last test acc: {test_acc:.2f} %")

    # Confusion Matrix
    cm = confusion_matrix(y_true, y_pred)
    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=le.classes_)
    fig, ax = plt.subplots(figsize=(16, 12))
    disp.plot(ax=ax, cmap='Blues', colorbar=True)
    plt.xticks(rotation=90)
    plt.title(f"Confusion Matrix at Epoch {epoch+1}, Test Accuracy: {test_acc:.2%}")
    plt.tight_layout()
    plt.show()

    # ----- Plot -----  
    plt.figure(figsize=(15, 4))

    plt.subplot(1, 3, 1)
    plt.plot(losses, label='Train Loss')
    plt.xlabel("Epoch"); plt.ylabel("Loss")
    plt.title("Training Loss")

    plt.subplot(1, 3, 2)
    plt.plot(accuracies, label='Train Acc')
    plt.xlabel("Epoch"); plt.ylabel("Accuracy")
    plt.title("Training Accuracy")

    plt.subplot(1, 3, 3)
    plt.plot(test_accuracies, label='Test Acc')
    plt.xlabel("Epoch"); plt.ylabel("Accuracy")
    plt.title("Test Accuracy")

    plt.tight_layout()
    plt.show()

    return model, le

# ====== USAGE ======
model, le = train_riftnet(dataset, analytic_signal, num_classes=30, epochs=100)

Epoch 1/100: 100%|██████████████████████████| 113/113 [00:12<00:00,  9.13it/s, acc=0.242, loss=3.28]

Test acc: 0.29 %
✅ Saved best model at epoch 1 with test acc: 29.11%

Epoch 2/100: 100%|██████████████████████████| 113/113 [00:11<00:00,  9.66it/s, acc=0.408, loss=3.11]

Test acc: 0.52 %
✅ Saved best model at epoch 2 with test acc: 52.22%

Epoch 3/100: 100%|██████████████████████████| 113/113 [00:11<00:00,  9.75it/s, acc=0.528, loss=2.96]

Test acc: 0.55 %
✅ Saved best model at epoch 3 with test acc: 55.00%

Epoch 4/100: 100%|██████████████████████████| 113/113 [00:11<00:00,  9.72it/s, acc=0.601, loss=2.89]

Test acc: 0.64 %
✅ Saved best model at epoch 4 with test acc: 64.33%

Epoch 5/100: 100%|██████████████████████████| 113/113 [00:11<00:00,  9.75it/s, acc=0.648, loss=2.84]

Test acc: 0.62 %