# ── APViT repo path setup (must precede mmcls/mmcv imports) ─────────────────
import sys
from pathlib import Path as _Path
_apvit_path = str(_Path(__file__).resolve().parent.parent.parent / "APViT")
if _apvit_path not in sys.path:
    sys.path.insert(0, _apvit_path)
# ─────────────────────────────────────────────────────────────────────────────

import pandas as pd
import os
import numpy as np
import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
import torch.optim as optim
from PIL import Image
from torch.optim import lr_scheduler
from tqdm import tqdm
from pathlib import Path
import mmcv
from mmcls.models import build_classifier

SCRIPT_DIR = Path(__file__).resolve().parent
MODEL_PATH = os.getenv("MODEL_PATH", str(SCRIPT_DIR / "model.pt"))
APVIT_REPO = SCRIPT_DIR.parent.parent / "APViT"
APVIT_CONFIG = APVIT_REPO / "configs" / "apvit" / "RAF.py"


def resolve_images_dir(split_name, explicit_env_var, data_root, extract_root):
    explicit = os.getenv(explicit_env_var, "").strip()
    candidates = []
    if explicit:
        candidates.append(Path(explicit))

    candidates.append(Path(data_root) / f"{split_name}_set" / "images")
    candidates.append(Path(extract_root) / f"{split_name}_extracted" / "images")

    extract_split_root = Path(extract_root) / f"{split_name}_extracted"
    if extract_split_root.exists():
        for p in extract_split_root.rglob("images"):
            if p.is_dir():
                candidates.append(p)

    for p in candidates:
        if p.exists() and p.is_dir():
            return str(p)

    tried = "\n".join([str(p) for p in candidates])
    raise FileNotFoundError(
        f"Could not find images folder for split='{split_name}'. Tried:\n{tried}"
    )

# Load the annotations for training and validation from separate CSV files
DATA_ROOT = os.getenv("AFFECTNET_ROOT", "/workspace/data_affectnet/AffectNet")
EXTRACT_ROOT = os.getenv("AFFECTNET_EXTRACT_ROOT", f"{DATA_ROOT}/extracted")
ANNO_ROOT = os.getenv("AFFECTNET_ANNO_ROOT", "../../affectnet_annotations")
IMAGE_FOLDER = resolve_images_dir("train", "AFFECTNET_TRAIN_IMAGES", DATA_ROOT, EXTRACT_ROOT)
IMAGE_FOLDER_TEST = resolve_images_dir("val", "AFFECTNET_VAL_IMAGES", DATA_ROOT, EXTRACT_ROOT)
train_annotations_path = os.getenv(
    "AFFECTNET_TRAIN_ANNO", f"{ANNO_ROOT}/train_set_annotation_without_lnd.csv"
)
valid_annotations_path = os.getenv(
    "AFFECTNET_VAL_ANNO", f"{ANNO_ROOT}/val_set_annotation_without_lnd.csv"
)
train_annotations_df = pd.read_csv(train_annotations_path)
valid_annotations_df = pd.read_csv(valid_annotations_path)


# Set parameters
BATCHSIZE = int(os.getenv("BATCHSIZE", "64"))
NUM_EPOCHS = 20
LR = 4e-5
NUM_WORKERS = int(os.getenv("NUM_WORKERS", "0"))
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("DATA_ROOT =", DATA_ROOT)
print("EXTRACT_ROOT =", EXTRACT_ROOT)
print("AFFECTNET_TRAIN_IMAGES =", IMAGE_FOLDER)
print("AFFECTNET_VAL_IMAGES =", IMAGE_FOLDER_TEST)
print("MODEL_PATH =", MODEL_PATH)


# **** Create dataset and data loaders ****
class CustomDataset(Dataset):
    def __init__(self, dataframe, root_dir, transform=None, balance=False):
        self.dataframe = dataframe
        self.transform = transform
        self.root_dir = root_dir
        self.balance = balance

        if self.balance:
            self.dataframe = self.balance_dataset()

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

    def __getitem__(self, idx):
        image_path = os.path.join(
            self.root_dir, f"{self.dataframe['number'].iloc[idx]}.jpg"
        )
        if os.path.exists(image_path):
            image = Image.open(image_path)
        else:
            image = Image.new(
                "RGB", (224, 224), color="white"
            )  # Handle missing image file

        classes = torch.tensor(self.dataframe["exp"].iloc[idx], dtype=torch.long)
        labels = torch.tensor(self.dataframe.iloc[idx, 2:4].values, dtype=torch.float32)

        if self.transform:
            image = self.transform(image)

        return image, classes, labels

    def balance_dataset(self):
        balanced_df = self.dataframe.groupby("exp", group_keys=False).apply(
            lambda x: x.sample(self.dataframe["exp"].value_counts().min())
        )
        return balanced_df


transform = transforms.Compose(
    [
        transforms.Resize(112),  # APViT / IR-50 requires 112x112 input
        transforms.RandomHorizontalFlip(0.5),
        transforms.RandomGrayscale(0.01),
        transforms.RandomRotation(10),
        transforms.ColorJitter(
            brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1
        ),  # model more robust to changes in lighting conditions.
        transforms.RandomPerspective(
            distortion_scale=0.2, p=0.5
        ),  # can be helpful if your images might have varying perspectives.
        transforms.ToTensor(),  # saves image as tensor (automatically divides by 255)
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        transforms.RandomErasing(
            p=0.5, scale=(0.02, 0.2), ratio=(0.3, 3.3), value="random"
        ),  # TEST: Should help overfitting
    ]
)

transform_valid = transforms.Compose(
    [
        transforms.Resize(112),  # APViT / IR-50 requires 112x112 input
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
)

train_dataset = CustomDataset(
    dataframe=train_annotations_df,
    root_dir=IMAGE_FOLDER,
    transform=transform,
    balance=False,
)
valid_dataset = CustomDataset(
    dataframe=valid_annotations_df,
    root_dir=IMAGE_FOLDER_TEST,
    transform=transform_valid,
    balance=False,
)
train_loader = DataLoader(
    train_dataset, batch_size=BATCHSIZE, shuffle=True, num_workers=NUM_WORKERS
)
valid_loader = DataLoader(
    valid_dataset, batch_size=BATCHSIZE, shuffle=False, num_workers=NUM_WORKERS
)

# ***** Define the model *****

APVIT_EMBED_DIM = 768


def _build_apvit_backbone():
    """Build APViT PoolingVitClassifier (backbone only, no head)."""
    cfg = mmcv.Config.fromfile(str(APVIT_CONFIG))
    cfg.model.pretrained = None
    cfg.model.head = None  # replaced by our custom head below
    ir50_w = APVIT_REPO / "weights" / "backbone_ir50_ms1m_epoch120.pth"
    vit_small_w = APVIT_REPO / "weights" / "vit_small_p16_224-15ec54c9.pth"
    cfg.model.extractor.pretrained = str(ir50_w) if ir50_w.exists() else None
    cfg.model.vit.pretrained = str(vit_small_w) if vit_small_w.exists() else None
    if cfg.model.extractor.pretrained is None:
        print("[WARNING] IR-50 weights not found at", ir50_w, "— training from random init.")
    if cfg.model.vit.pretrained is None:
        print("[WARNING] ViT-Small weights not found at", vit_small_w, "— training from random init.")
    return build_classifier(cfg.model)


class APViTWithHead(nn.Module):
    def __init__(self, num_outputs):
        super().__init__()
        self.apvit = _build_apvit_backbone()  # IR-50 + PoolingViT (no head)
        self.head = nn.Sequential(
            nn.LayerNorm(APVIT_EMBED_DIM),
            nn.Linear(APVIT_EMBED_DIM, APVIT_EMBED_DIM),
            nn.Tanh(),
            nn.Linear(APVIT_EMBED_DIM, num_outputs, bias=False),
        )

    def forward(self, x):
        features, _ = self.apvit.extract_feat(x)  # [B, 768] CLS token
        return self.head(features)


MODEL = APViTWithHead(num_outputs=10)
MODEL.to(DEVICE)  # Put the model to the GPU

# Define (weighted) loss function
weights = torch.tensor(
    [0.015605, 0.008709, 0.046078, 0.083078, 0.185434, 0.305953, 0.046934, 0.30821]
)
criterion_cls = nn.CrossEntropyLoss(weights.to(DEVICE))
criterion_cls_val = (
    nn.CrossEntropyLoss()
)  # Use two loss functions, as the validation dataset is balanced
criterion_reg = nn.MSELoss()

optimizer = optim.AdamW(MODEL.parameters(), lr=LR)
lr_scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=BATCHSIZE * NUM_EPOCHS)


def ccc_numpy_1d(x, y, eps=1e-8):
    x = np.asarray(x, dtype=np.float32).reshape(-1)
    y = np.asarray(y, dtype=np.float32).reshape(-1)
    x_mean, y_mean = x.mean(), y.mean()
    x_var, y_var = x.var(), y.var()
    cov = ((x - x_mean) * (y - y_mean)).mean()
    return (2.0 * cov) / (x_var + y_var + (x_mean - y_mean) ** 2 + eps)


# ***** Train the model *****
print("--- Start training ---")
scaler = torch.cuda.amp.GradScaler()
best_valid_loss = float("inf")
best_epoch = -1
history = {
    "epoch": [], "train_loss": [], "valid_loss": [],
    "train_acc": [], "valid_acc": [],
    "ccc_v": [], "ccc_a": [], "ccc_mean": [], "lr": [],
}

for epoch in range(NUM_EPOCHS):
    MODEL.train()
    total_train_loss = 0.0
    total_train_correct = 0
    total_train_samples = 0
    for images, classes, labels in tqdm(
        train_loader, desc="Epoch train_loader progress"
    ):
        images, classes, labels = (
            images.to(DEVICE),
            classes.to(DEVICE),
            labels.to(DEVICE),
        )
        optimizer.zero_grad()
        with torch.autocast(device_type="cuda", dtype=torch.float16):
            outputs = MODEL(images)
            outputs_cls = outputs[:, :8]
            outputs_reg = outputs[:, 8:]
            loss = criterion_cls(
                outputs_cls.cuda(), classes.cuda()
            ) + 5 * criterion_reg(outputs_reg.cuda(), labels.cuda())
            total_train_loss += loss.item()
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            lr_scheduler.step()
            current_lr = optimizer.param_groups[0]["lr"]

        _, train_predicted = torch.max(outputs_cls, 1)
        total_train_samples += classes.size(0)
        total_train_correct += (train_predicted == classes).sum().item()

    train_loss = total_train_loss / max(1, len(train_loader))
    train_accuracy = (total_train_correct / total_train_samples) * 100

    MODEL.eval()
    valid_loss = 0.0
    correct = 0
    total = 0
    val_true_all, val_pred_all = [], []
    aro_true_all, aro_pred_all = [], []
    with torch.no_grad():
        for images, classes, labels in valid_loader:
            images, classes, labels = (
                images.to(DEVICE),
                classes.to(DEVICE),
                labels.to(DEVICE),
            )
            outputs = MODEL(images)
            outputs_cls = outputs[:, :8]
            outputs_reg = outputs[:, 8:]
            loss = criterion_cls_val(
                outputs_cls.cuda(), classes.cuda()
            ) + 5 * criterion_reg(outputs_reg.cuda(), labels.cuda())
            valid_loss += loss.item()
            _, predicted = torch.max(outputs_cls, 1)
            total += classes.size(0)
            correct += (predicted == classes).sum().item()
            val_true_all.extend(labels[:, 0].cpu().float().numpy())
            val_pred_all.extend(outputs_reg[:, 0].cpu().float().numpy())
            aro_true_all.extend(labels[:, 1].cpu().float().numpy())
            aro_pred_all.extend(outputs_reg[:, 1].cpu().float().numpy())

    valid_loss_mean = valid_loss / max(1, len(valid_loader))
    valid_acc = (correct / total) * 100
    ccc_v = ccc_numpy_1d(val_pred_all, val_true_all)
    ccc_a = ccc_numpy_1d(aro_pred_all, aro_true_all)
    ccc_mean = (ccc_v + ccc_a) / 2.0
    current_lr = optimizer.param_groups[0]["lr"]

    history["epoch"].append(epoch + 1)
    history["train_loss"].append(train_loss)
    history["valid_loss"].append(valid_loss_mean)
    history["train_acc"].append(train_accuracy)
    history["valid_acc"].append(valid_acc)
    history["ccc_v"].append(ccc_v)
    history["ccc_a"].append(ccc_a)
    history["ccc_mean"].append(ccc_mean)
    history["lr"].append(current_lr)

    print(
        f"Epoch [{epoch+1}/{NUM_EPOCHS}] - "
        f"Train Loss: {train_loss:.4f}, "
        f"Valid Loss: {valid_loss_mean:.4f}, "
        f"Train Acc: {train_accuracy:.2f}%, "
        f"Valid Acc: {valid_acc:.2f}%, "
        f"Valid CCC_v: {ccc_v:.4f}, "
        f"Valid CCC_a: {ccc_a:.4f}, "
        f"Valid CCC_mean: {ccc_mean:.4f}, "
        f"LR: {current_lr:.8f}"
    )

    if valid_loss_mean < best_valid_loss:
        best_valid_loss = valid_loss_mean
        best_epoch = epoch + 1
        print(f"Saving model at epoch {best_epoch}")
        torch.save(MODEL.state_dict(), MODEL_PATH)  # Save the best model

# ── End-of-training summary ──────────────────────────────────────────────────
best_idx = best_epoch - 1
print("\n" + "=" * 60)
print("TRAINING COMPLETE")
print("=" * 60)
print(f"Best epoch : {best_epoch}/{NUM_EPOCHS}")
print(f"  Valid Loss  : {history['valid_loss'][best_idx]:.4f}")
print(f"  Valid Acc   : {history['valid_acc'][best_idx]:.2f}%")
print(f"  Valid CCC_v : {history['ccc_v'][best_idx]:.4f}")
print(f"  Valid CCC_a : {history['ccc_a'][best_idx]:.4f}")
print(f"  Valid CCC_mean: {history['ccc_mean'][best_idx]:.4f}")
print("\nAverage across all epochs:")
print(f"  Avg Valid Loss  : {sum(history['valid_loss'])/NUM_EPOCHS:.4f}")
print(f"  Avg Valid Acc   : {sum(history['valid_acc'])/NUM_EPOCHS:.2f}%")
print(f"  Avg CCC_v       : {sum(history['ccc_v'])/NUM_EPOCHS:.4f}")
print(f"  Avg CCC_a       : {sum(history['ccc_a'])/NUM_EPOCHS:.4f}")
print(f"  Avg CCC_mean    : {sum(history['ccc_mean'])/NUM_EPOCHS:.4f}")
print("=" * 60)