# -*- coding: utf-8 -*-
"""model_training-final version.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1R-zxmUzAC80wBmiMOJ56gh9DjMi8dWPU

## 1. Install Libraries
"""

# Install latest libraries optimized for A100
!pip install torch torchvision timm scikit-learn tqdm matplotlib roboflow optuna optuna-integration lightning torchmetrics onnx onnxruntime onnxruntime-gpu


# Mount Google Drive for persistent storage
from google.colab import drive
drive.mount('/content/drive')

# Check A100 GPU availability
!nvidia-smi


import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

import json
import torch
import torch.nn as nn
import timm
import optuna
import numpy as np
import matplotlib.pyplot as plt
import time
from tqdm import tqdm
from sklearn.metrics import classification_report, confusion_matrix, ConfusionMatrixDisplay, f1_score
from sklearn.utils.class_weight import compute_class_weight
import torchvision.transforms.v2 as T
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader, WeightedRandomSampler
from roboflow import Roboflow
import lightning.pytorch as pl
from lightning.pytorch import Trainer, LightningModule, seed_everything
from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint, StochasticWeightAveraging, LearningRateMonitor
from lightning.pytorch.loggers import TensorBoardLogger
from torchmetrics.classification import MulticlassAccuracy, MulticlassF1Score
from pathlib import Path

# GPU Memory monitoring function
def print_gpu_memory():
    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1024**3
        reserved = torch.cuda.memory_reserved() / 1024**3
        print(f"GPU Memory: {allocated:.2f}GB allocated / {reserved:.2f}GB reserved")

"""## 2. Config

"""

# 3. Configuration
config = {
    "seed": 42,
    "roboflow_api_key": "api",  # Replace with your Roboflow API key
    "roboflow_workspace": "workspace", # Replace with Roboflow workspace name
    "roboflow_project": "project", # Replace with Roboflow project name
    "roboflow_version": 14, # Replace with Roboflow project version
    "roboflow_dataset_format": "folder", # Replace with Roboflow folder
    "data_dir": "/content/folder", # Replace with Roboflow folder
    "train_subdir": "train",
    "val_subdir": "valid",
    "test_subdir": "test",
    "class_names": ["away-from-desk", "distracted", "studying"],
    "num_classes": 3,
    "img_size": 224,
    "batch_size": 128,  # Increased for A100
    "num_workers": 8,   # Optimized for A100
    "prefetch_factor": 4,  # Higher prefetch for A100
    "epochs": 15,
    "early_stopping_patience": 5,
    "optuna_n_trials": 20,  # More trials due to faster execution
    "optuna_epochs": 8,     # Increased for better evaluation
    "model_name": "tiny_vit_21m_224",
    "model_pretrained": True,
    "dropout_rate": 0.1,
    "drop_path_rate": 0.05,
    "scheduler": "Cosine",
    "pct_start": 0.1,
    "warmup_epochs": 1,
    "use_swa": True,
    "checkpoint_dir": "/content/drive/MyDrive/tinyvit_checkpoints-2",  # Google Drive
    "log_dir": "/content/drive/MyDrive/tinyvit_logs-2",              # Google Drive
}

# Create directories
Path(config["checkpoint_dir"]).mkdir(parents=True, exist_ok=True)
Path(config["log_dir"]).mkdir(parents=True, exist_ok=True)

# Unpack configuration
SEED = config["seed"]
DATA_DIR = Path(config["data_dir"])
TRAIN_SUB = config["train_subdir"]
VAL_SUB = config["val_subdir"]
TEST_SUB = config["test_subdir"]
NUM_CLASSES = config["num_classes"]
IMG_SIZE = config["img_size"]
BATCH_SIZE = config["batch_size"]
EPOCHS = config["epochs"]
OPTUNA_TRIALS = config["optuna_n_trials"]
OPTUNA_EPOCHS = config["optuna_epochs"]
EARLY_STOP_P = config["early_stopping_patience"]
MODEL_NAME = config["model_name"]
MODEL_PRETRAIN = config["model_pretrained"]
DROPOUT_RATE = config["dropout_rate"]
DROP_PATH_RATE = config["drop_path_rate"]
SCHEDULER = config["scheduler"]
PCT_START = config["pct_start"]
WARMUP_EPOCHS = config["warmup_epochs"]
USE_SWA = config["use_swa"]
CHECKPOINT_DIR = Path(config["checkpoint_dir"])
LOG_DIR = Path(config["log_dir"])

# Reproducibility & Device Setup
torch.manual_seed(SEED)
np.random.seed(SEED)
torch.use_deterministic_algorithms(True)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.backends.cudnn.benchmark = True  # Enable for A100 performance
print(f"Device: {DEVICE}")
print(f"CUDA Version: {torch.version.cuda}")
print_gpu_memory()

"""## 4. Download Dataset and Setup"""

CHECKPOINT_DIR = Path(config["checkpoint_dir"])
LOG_DIR = Path(config["log_dir"])

seed_everything(SEED)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.manual_seed(SEED)
np.random.seed(SEED)
torch.use_deterministic_algorithms(True)
torch.backends.cudnn.benchmark = False

print(f"Device: {DEVICE}, Seed: {SEED}")
print(f"Data directory: {DATA_DIR}")

def download_dataset():
    if not DATA_DIR.exists():
        print(f"Downloading dataset to {DATA_DIR}...")
        rf = Roboflow(api_key=config["roboflow_api_key"])
        proj = rf.workspace(config["roboflow_workspace"]).project(config["roboflow_project"])
        version = proj.version(config["roboflow_version"])
        version.download(config["roboflow_dataset_format"])
    else:
        print(f"Dataset already exists at {DATA_DIR}")

download_dataset()

"""## 5. Data Transforms and Loaders"""

mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]

# Enhanced transforms for A100 training
train_transform = T.Compose([
    T.RandomResizedCrop(IMG_SIZE, scale=(0.85, 1.0)),
    T.RandomHorizontalFlip(p=0.5),
    T.ColorJitter(0.1, 0.1, 0.1, 0.05),
    T.RandomRotation(degrees=5),
    T.RandomErasing(p=0.25, scale=(0.02, 0.15)),
    T.ToImage(),
    T.ToDtype(torch.float32, scale=True),
    T.Normalize(mean=mean, std=std),
])

val_transform = T.Compose([
    T.Resize((IMG_SIZE + 32, IMG_SIZE + 32)),
    T.CenterCrop(IMG_SIZE),
    T.ToImage(),
    T.ToDtype(torch.float32, scale=True),
    T.Normalize(mean=mean, std=std),
])

# Load datasets
train_path = DATA_DIR / TRAIN_SUB
val_path = DATA_DIR / VAL_SUB
train_ds = ImageFolder(train_path, transform=train_transform)
train_labels = [label for _, label in train_ds.samples]

# Compute class weights
weights_np = compute_class_weight('balanced', classes=np.arange(NUM_CLASSES), y=train_labels)
class_weights = torch.tensor(weights_np, dtype=torch.float32).to(DEVICE)

# Create balanced sampler
sample_weights = class_weights[torch.tensor(train_labels)]
train_sampler = WeightedRandomSampler(
    weights=sample_weights,
    num_samples=len(sample_weights),
    replacement=True
)

# A100 optimized data loaders
train_loader = DataLoader(
    train_ds,
    batch_size=BATCH_SIZE,
    sampler=train_sampler,
    num_workers=config["num_workers"],
    pin_memory=True,
    persistent_workers=True,
    prefetch_factor=config["prefetch_factor"],
    drop_last=True  # Consistent batch sizes for A100
)

val_ds = ImageFolder(val_path, transform=val_transform)
val_loader = DataLoader(
    val_ds,
    batch_size=BATCH_SIZE,
    shuffle=False,
    num_workers=config["num_workers"],
    pin_memory=True,
    persistent_workers=True,
    prefetch_factor=config["prefetch_factor"]
)

print(f"Training samples: {len(train_ds)}")
print(f"Validation samples: {len(val_ds)}")
print(f"Class distribution: {np.bincount(train_labels)}")
print_gpu_memory()

"""## 6. Mixup and CutMix Implementation"""

def rand_bbox(size, lam):
    W = size[2]
    H = size[3]
    cut_rat = np.sqrt(1. - lam)
    cut_w = int(W * cut_rat)
    cut_h = int(H * cut_rat)
    cx = np.random.randint(W)
    cy = np.random.randint(H)
    bbx1 = np.clip(cx - cut_w // 2, 0, W)
    bby1 = np.clip(cy - cut_h // 2, 0, H)
    bbx2 = np.clip(cx + cut_w // 2, 0, W)
    bby2 = np.clip(cy + cut_h // 2, 0, H)
    return bbx1, bby1, bbx2, bby2

def mixup_cutmix(x, y, num_classes, alpha=1.0, cutmix_prob=0.5):
    device = x.device
    batch_size = x.size(0)

    if alpha > 0:
        lam = np.random.beta(alpha, alpha)
    else:
        lam = 1

    rand_index = torch.randperm(batch_size, device=device)
    target_a = y
    target_b = y[rand_index]

    targets_onehot_a = torch.eye(num_classes, device=device)[target_a]
    targets_onehot_b = torch.eye(num_classes, device=device)[target_b]

    if np.random.rand() < cutmix_prob:
        # CutMix
        bbx1, bby1, bbx2, bby2 = rand_bbox(x.size(), lam)
        x[:, :, bbx1:bbx2, bby1:bby2] = x[rand_index, :, bbx1:bbx2, bby1:bby2]
        lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (x.size(-1) * x.size(-2)))
    else:
        # Mixup
        x = lam * x + (1 - lam) * x[rand_index, :]

    return x, targets_onehot_a, targets_onehot_b, lam

def mixup_cutmix_loss(criterion, pred, ta, tb, lam):
    return lam * criterion(pred, ta) + (1 - lam) * criterion(pred, tb)

"""## 7. Lightning Module"""

class TimingCallback(pl.Callback):
    def on_train_epoch_start(self, trainer, pl_module):
        self.start_time = time.time()

    def on_train_epoch_end(self, trainer, pl_module):
        epoch_time = time.time() - self.start_time
        print(f"Epoch {trainer.current_epoch} completed in {epoch_time:.2f}s")

class TinyViTLightning(LightningModule):
    def __init__(self, learning_rate, weight_decay, optimizer_name,
                 mixup_alpha, cutmix_prob, dropout_rate, drop_path_rate,
                 num_classes, model_name, model_pretrained, class_weights=None,
                 scheduler="Cosine", pct_start=0.1, warmup_epochs=1, total_steps=None):
        super().__init__()
        self.save_hyperparameters()

        self.model = timm.create_model(
            model_name,
            pretrained=model_pretrained,
            num_classes=num_classes,
            drop_rate=dropout_rate,
            drop_path_rate=drop_path_rate
        )

        self.criterion = nn.CrossEntropyLoss(
            weight=class_weights,
            label_smoothing=0.1
        )

        self.train_acc = MulticlassAccuracy(num_classes=num_classes)
        self.val_acc = MulticlassAccuracy(num_classes=num_classes)
        self.val_f1 = MulticlassF1Score(num_classes=num_classes, average='macro')

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        images, labels = batch
        labels = labels.long()
        # Apply mixup/cutmix with 70% probability for A100 training
        if self.hparams.mixup_alpha > 0 and np.random.rand() < 0.7:
            mixed, ta, tb, lam = mixup_cutmix(
                images, labels,
                num_classes=self.hparams.num_classes,
                alpha=self.hparams.mixup_alpha,
                cutmix_prob=self.hparams.cutmix_prob
            )
            out = self(mixed)
            loss = mixup_cutmix_loss(self.criterion, out, ta, tb, lam)
        else:
            out = self(images)
            loss = self.criterion(out, labels)

        preds = out.argmax(dim=1)
        acc = (preds == labels).float().mean()

        self.log("train_loss", loss, on_epoch=True, prog_bar=True)
        self.log("train_acc", acc, on_epoch=True, prog_bar=True)

        return loss

    def validation_step(self, batch, batch_idx):
        images, labels = batch
        labels = labels.long()

        logits = self(images)
        loss = self.criterion(logits, labels)
        preds = logits.argmax(dim=1)

        self.val_acc.update(preds, labels)
        self.val_f1.update(preds, labels)

        self.log("val_loss", loss, on_epoch=True, prog_bar=True)
        self.log("val_acc", self.val_acc, on_epoch=True, prog_bar=True)
        self.log("val_f1", self.val_f1, on_epoch=True, prog_bar=True)

        return loss

    def configure_optimizers(self):
        # Scale learning rate with batch size
        scaled_lr = self.hparams.learning_rate * (BATCH_SIZE / 32)

        if self.hparams.optimizer_name == "AdamW":
            opt = torch.optim.AdamW(
                self.parameters(),
                lr=scaled_lr,
                weight_decay=self.hparams.weight_decay,
                eps=1e-8
            )
        elif self.hparams.optimizer_name == "Adam":
            opt = torch.optim.Adam(
                self.parameters(),
                lr=scaled_lr,
                weight_decay=self.hparams.weight_decay
            )
        else:
            opt = torch.optim.SGD(
                self.parameters(),
                lr=scaled_lr,
                weight_decay=self.hparams.weight_decay,
                momentum=0.9,
                nesterov=True
            )

        if self.hparams.scheduler == "OneCycle":
            sched = torch.optim.lr_scheduler.OneCycleLR(
                opt,
                max_lr=scaled_lr,
                total_steps=self.hparams.total_steps,
                pct_start=self.hparams.pct_start
            )
            return {"optimizer": opt, "lr_scheduler": {"scheduler": sched, "interval": "step"}}
        else:
            sched = torch.optim.lr_scheduler.CosineAnnealingLR(
                opt,
                T_max=self.trainer.max_epochs - self.hparams.warmup_epochs,
                eta_min=1e-7
            )
            warmup = torch.optim.lr_scheduler.LinearLR(
                opt,
                start_factor=0.01,
                total_iters=self.hparams.warmup_epochs
            )
            from torch.optim.lr_scheduler import SequentialLR
            seq = SequentialLR(
                opt,
                schedulers=[warmup, sched],
                milestones=[self.hparams.warmup_epochs]
            )
            return {"optimizer": opt, "lr_scheduler": {"scheduler": seq, "interval": "epoch"}}

"""## 8. Optuna Hyperparameter tunning"""

def objective(trial):
    # A100 optimized hyperparameter ranges
    batch_size = trial.suggest_categorical("batch_size", [96, 128, 160, 192])
    lr = trial.suggest_float("learning_rate", 5e-5, 5e-3, log=True)
    weight_decay = trial.suggest_float("weight_decay", 1e-6, 1e-3, log=True)
    optimizer_name = trial.suggest_categorical("optimizer", ["AdamW", "Adam"])
    mixup_alpha = trial.suggest_float("mixup_alpha", 0.1, 0.6)
    cutmix_prob = trial.suggest_float("cutmix_prob", 0.3, 0.8)
    dropout_rate = trial.suggest_float("dropout_rate", 0.05, 0.25)
    drop_path_rate = trial.suggest_float("drop_path_rate", 0.0, 0.15)
    scheduler = trial.suggest_categorical("scheduler", ["Cosine", "OneCycle"])
    pct_start = trial.suggest_float("pct_start", 0.1, 0.3)
    warmup_epochs = trial.suggest_int("warmup_epochs", 1, 3)
    max_epochs = OPTUNA_EPOCHS

    # Create optimized data loaders for trial
    trial_train_loader = DataLoader(
        train_ds,
        batch_size=batch_size,
        sampler=train_sampler,
        num_workers=6,  # Slightly reduced for stability
        pin_memory=True,
        persistent_workers=False,  # Disabled for trials
        prefetch_factor=2,
        drop_last=True
    )
    trial_val_loader = DataLoader(
        val_ds,
        batch_size=batch_size,
        shuffle=False,
        num_workers=6,
        pin_memory=True,
        persistent_workers=False,
        prefetch_factor=2
    )

    total_steps = len(trial_train_loader) * max_epochs if scheduler == "OneCycle" else None

    try:
        lit_model = TinyViTLightning(
            learning_rate=lr,
            weight_decay=weight_decay,
            optimizer_name=optimizer_name,
            mixup_alpha=mixup_alpha,
            cutmix_prob=cutmix_prob,
            dropout_rate=dropout_rate,
            drop_path_rate=drop_path_rate,
            num_classes=NUM_CLASSES,
            model_name=MODEL_NAME,
            model_pretrained=MODEL_PRETRAIN,
            class_weights=class_weights,
            scheduler=scheduler,
            pct_start=pct_start,
            warmup_epochs=warmup_epochs,
            total_steps=total_steps
        )

        from optuna.integration import PyTorchLightningPruningCallback

        callbacks = [
            EarlyStopping(monitor="val_acc", mode="max", patience=3),
            ModelCheckpoint(
                monitor="val_acc",
                mode="max",
                save_top_k=1,
                dirpath=f"{config['checkpoint_dir']}/optuna_ckpts",
                filename=f"trial-{trial.number}-{{epoch:02d}}-{{val_acc:.4f}}"
            ),
            PyTorchLightningPruningCallback(trial, monitor="val_acc")
        ]

        trainer = Trainer(
            max_epochs=max_epochs,
            accelerator="gpu",
            devices=1,
            precision="16-mixed",  # A100 optimization
            callbacks=callbacks,
            enable_progress_bar=False,
            gradient_clip_val=1.0,
            logger=False,
            deterministic=False,  # Allow non-deterministic for speed
            benchmark=True  # A100 optimization
        )

        trainer.fit(lit_model, trial_train_loader, trial_val_loader)

        # Safe metric retrieval
        if "val_acc" in trainer.callback_metrics:
            val_acc = trainer.callback_metrics["val_acc"].item()
        else:
            val_acc = 0.0

        return val_acc

    except Exception as e:
        print(f"Trial {trial.number} failed: {e}")
        raise optuna.TrialPruned()

    finally:
        # Cleanup for A100 memory management
        if 'lit_model' in locals():
            del lit_model
        if 'trainer' in locals():
            del trainer
        torch.cuda.empty_cache()

# Run Optuna optimization
study = optuna.create_study(
    direction="maximize",
    sampler=optuna.samplers.TPESampler(seed=SEED),
    pruner=optuna.pruners.MedianPruner(n_startup_trials=3, n_warmup_steps=2)
)

print("Starting Optuna optimization on A100...")
print_gpu_memory()

try:
    study.optimize(objective, n_trials=OPTUNA_TRIALS, timeout=None)

    best = study.best_trial
    print(f"\n🏆 Best val_acc: {best.value:.4f}")
    print("Best hyperparameters:", best.params)

    # Save results to Google Drive
    results = {
        "best_params": best.params,
        "best_value": best.value,
        "n_trials": len(study.trials),
        "device": "A100"
    }

    with open(f"{config['checkpoint_dir']}/best_hyperparameters.json", "w") as f:
        json.dump(results, f, indent=4)

    print(f"Study completed with {len(study.trials)} trials")
    print(f"Results saved to Google Drive")

except Exception as e:
    print(f"Optuna optimization failed: {e}")

"""## 9. Final Training"""

# Load best hyperparameters
try:
    with open(f"{config['checkpoint_dir']}/best_hyperparameters.json") as f:
        results = json.load(f)
        best = results["best_params"]
except FileNotFoundError:
    print("Using default hyperparameters")
    best = {}

# Extract optimized parameters
batch_size = best.get("batch_size", BATCH_SIZE)
max_epochs = EPOCHS
learning_rate = best.get("learning_rate", 1e-3)
weight_decay = best.get("weight_decay", 1e-4)

# Create final optimized data loaders
final_train_loader = DataLoader(
    train_ds,
    batch_size=batch_size,
    sampler=train_sampler,
    num_workers=config["num_workers"],
    pin_memory=True,
    persistent_workers=True,
    prefetch_factor=config["prefetch_factor"],
    drop_last=True
)

final_val_loader = DataLoader(
    val_ds,
    batch_size=batch_size,
    shuffle=False,
    num_workers=config["num_workers"],
    pin_memory=True,
    persistent_workers=True,
    prefetch_factor=config["prefetch_factor"]
)

total_steps = len(final_train_loader) * max_epochs if best.get("scheduler") == "OneCycle" else None

# Create final model
final_model = TinyViTLightning(
    learning_rate=learning_rate,
    weight_decay=weight_decay,
    optimizer_name=best.get("optimizer", "AdamW"),
    mixup_alpha=best.get("mixup_alpha", 0.2),
    cutmix_prob=best.get("cutmix_prob", 0.5),
    dropout_rate=best.get("dropout_rate", 0.1),
    drop_path_rate=best.get("drop_path_rate", 0.05),
    scheduler=best.get("scheduler", "Cosine"),
    pct_start=best.get("pct_start", 0.1),
    warmup_epochs=best.get("warmup_epochs", 1),
    total_steps=total_steps,
    num_classes=NUM_CLASSES,
    model_name=MODEL_NAME,
    model_pretrained=MODEL_PRETRAIN,
    class_weights=class_weights
)

# A100 optimized callbacks
logger = TensorBoardLogger(config["log_dir"], name="TinyViT_A100_final")
timing_callback = TimingCallback()
early_stop = EarlyStopping(
    monitor="val_acc",
    mode="max",
    patience=config["early_stopping_patience"],
    verbose=True
)
checkpoint = ModelCheckpoint(
    monitor="val_acc",
    mode="max",
    save_top_k=3,
    dirpath=f"{config['checkpoint_dir']}/final_ckpts",
    filename="final-{epoch:02d}-{val_acc:.4f}",
    save_last=True,
    every_n_epochs=2  # Save frequently for Colab
)
swa = StochasticWeightAveraging(swa_lrs=learning_rate)
lr_monitor = LearningRateMonitor(logging_interval="epoch")

# A100 optimized trainer
trainer = Trainer(
    accelerator="gpu",
    devices=1,
    max_epochs=max_epochs,
    precision="16-mixed",  # A100 Tensor Core optimization
    gradient_clip_val=1.0,
    gradient_clip_algorithm="norm",
    deterministic=False,  # Allow non-deterministic for A100 speed
    benchmark=True,       # A100 optimization
    callbacks=[early_stop, checkpoint, swa, lr_monitor, timing_callback],
    logger=logger,
    enable_progress_bar=True,
    log_every_n_steps=50,
    val_check_interval=0.5  # Check validation twice per epoch
)

print("Starting final training on A100...")
print(f"Batch size: {batch_size}")
print(f"Learning rate: {learning_rate}")
print_gpu_memory()

# Train the model
trainer.fit(final_model, final_train_loader, final_val_loader)

print("Training completed!")
print_gpu_memory()

"""## 10. Testing"""

# Load best checkpoint
ckpt_files = list(Path(f"{config['checkpoint_dir']}/final_ckpts").glob("final-*.ckpt"))
if ckpt_files:
    ckpt_path = max(ckpt_files, key=lambda x: x.stat().st_mtime)  # Most recent
else:
    ckpt_path = trainer.checkpoint_callback.best_model_path

print(f"Loading checkpoint: {ckpt_path}")

# Load model
best_model = TinyViTLightning.load_from_checkpoint(
    ckpt_path,
    num_classes=NUM_CLASSES,
    model_name=MODEL_NAME,
    model_pretrained=MODEL_PRETRAIN,
    class_weights=class_weights,
)
best_model = best_model.to(DEVICE)
best_model.eval()

# Test dataset
test_ds = ImageFolder(DATA_DIR / TEST_SUB, transform=val_transform)
test_loader = DataLoader(
    test_ds,
    batch_size=BATCH_SIZE,
    shuffle=False,
    num_workers=config["num_workers"],
    pin_memory=True,
)

print(f"Test samples: {len(test_ds)}")

# Evaluation
all_preds, all_trues = [], []
start_time = time.time()

with torch.no_grad():
    for images, labels in tqdm(test_loader, desc="Testing"):
        images, labels = images.to(DEVICE), labels.to(DEVICE)
        logits = best_model(images)
        preds = torch.argmax(logits, dim=1)
        all_preds.append(preds.cpu())
        all_trues.append(labels.cpu())

test_time = time.time() - start_time
preds = torch.cat(all_preds).numpy()
trues = torch.cat(all_trues).numpy()

# Calculate metrics
test_acc = (preds == trues).mean()
test_f1 = f1_score(trues, preds, average="macro")

print(f"\n🎯 Final Results (A100 Training):")
print(f"Test Accuracy: {test_acc:.4f}")
print(f"Test Macro F1: {test_f1:.4f}")
print(f"Test time: {test_time:.2f}s")
print(f"Inference speed: {len(test_ds)/test_time:.1f} images/sec")

# Detailed report
print("\nClassification Report:")
print(classification_report(trues, preds, target_names=test_ds.classes, zero_division=0))

# Confusion matrix
cm = confusion_matrix(trues, preds)
plt.figure(figsize=(10, 8))
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=test_ds.classes)
disp.plot(cmap='Blues', values_format='d')
plt.title('Confusion Matrix - TinyViT A100 Training')
plt.tight_layout()
plt.savefig(f"{config['checkpoint_dir']}/confusion_matrix.png", dpi=300, bbox_inches='tight')
plt.show()

# Save final results
final_results = {
    "test_accuracy": float(test_acc),
    "test_f1": float(test_f1),
    "test_time": test_time,
    "inference_speed": len(test_ds)/test_time,
    "device": "A100",
    "batch_size": BATCH_SIZE,
    "model_name": MODEL_NAME
}

with open(f"{config['checkpoint_dir']}/final_results.json", "w") as f:
    json.dump(final_results, f, indent=4)

print(f"\nResults saved to: {config['checkpoint_dir']}")
print_gpu_memory()

"""## Cell 11: Fixed Model Export (.pt and .onnx)

"""

import os
import torch
from pathlib import Path

# Export optimized model
print("Exporting model...")

# Ensure model is in evaluation mode and on CPU
model_cpu = best_model.cpu()
model_cpu.eval()

# Create export directory if it doesn't exist
export_dir = Path(config['checkpoint_dir'])
export_dir.mkdir(parents=True, exist_ok=True)

try:
    # 1. TorchScript (.pt) export using Lightning's built-in method
    print("🔄 Exporting TorchScript model...")
    scripted = model_cpu.to_torchscript()  # Use Lightning's method instead of torch.jit.script()
    pt_path = export_dir / "tinyvit_a100_optimized.pt"
    torch.jit.save(scripted, str(pt_path))

    # Verify file size
    file_size_pt = os.path.getsize(pt_path) / (1024 * 1024)  # Size in MB
    print(f"✅ TorchScript model saved: {pt_path}")
    print(f"   File size: {file_size_pt:.2f} MB")

except Exception as e:
    print(f"❌ TorchScript export failed: {e}")

try:
    # 2. ONNX (.onnx) export
    print("🔄 Exporting ONNX model...")
    dummy_input = torch.randn(1, 3, IMG_SIZE, IMG_SIZE)
    onnx_path = export_dir / "tinyvit_a100_optimized.onnx"

    torch.onnx.export(
        model_cpu,
        dummy_input,
        str(onnx_path),
        input_names=["input"],
        output_names=["output"],
        opset_version=17,  # Latest ONNX opset
        dynamic_axes={
            "input": {0: "batch_size"},
            "output": {0: "batch_size"}
        },
        export_params=True,
        do_constant_folding=True,
        verbose=False,
        training=torch.onnx.TrainingMode.EVAL
    )

    # Verify file size
    file_size_onnx = os.path.getsize(onnx_path) / (1024 * 1024)  # Size in MB
    print(f"✅ ONNX model saved: {onnx_path}")
    print(f"   File size: {file_size_onnx:.2f} MB")

except Exception as e:
    print(f"❌ ONNX export failed: {e}")

print(f"\n✅ Model export completed!")
print(f"📂 Files saved to: {config['checkpoint_dir']}")