import os
import torch as pt
import torch.nn.functional as F
from tqdm.auto import tqdm
from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR
from configuration_nula import NulaConfig
from modeling_nula import NulaForImageClassification, BlurPool2d
from dataset_nula import get_loaders, get_device
from augmentations import resize_down_up, decimate

NUM_EPOCHS = 50
AUG_PROB = 0.5
GRAD_CLIP = 1.0
SAVE_EVERY = 10
CHECKPOINT_DIR = "./checkpoints"
BEST_MODEL_DIR = "./nula-best-model"

def train_one_epoch(model, loader, optimizer, device, mean, std, blur, grad_clip=1.0):
    model.train()
    
    total_loss = 0.0
    total_correct = 0
    total_examples = 0

    pbar = tqdm(loader, desc="training...", leave=False)

    for batch in pbar:
        x = batch["pixel_values"].to(device, non_blocking=True)
        y = batch["labels"].to(device, non_blocking=True)
        optimizer.zero_grad(set_to_none=True)
 
        B = x.size(0) 
        mask_aug = pt.rand(B, device=x.device) < AUG_PROB

        if mask_aug.any():
            with pt.no_grad():
                x_image = x * std + mean
                choices = pt.randint(0, 3, (B,), device=x.device)
                mask_resize = mask_aug & (choices == 0)
                if mask_resize.any():
                    scales = pt.empty(mask_resize.sum(), device=x.device).uniform_(0.2, 0.6)
                    x_subset = x_image[mask_resize]

                    resize_out = []
                    for i in range(x_subset.size(0)):
                        resize_out.append(resize_down_up(x_subset[i:i+1], scale=scales[i].item()))
                    x_image[mask_resize] = pt.cat(resize_out, dim=0)
                    
                mask_decimate = mask_aug & (choices == 1)
                if mask_decimate.any():
                    factors = pt.randint(2, 5, (mask_decimate.sum(),), device=x.device)
                    x_subset = x_image[mask_decimate]

                    decimate_out = []
                    for i in range(x_subset.size(0)):
                        decimate_out.append(decimate(x_subset[i:i+1], factor=int(factors[i].item())))
                    x_image[mask_decimate] = pt.cat(decimate_out, dim=0)
                
                mask_blur = mask_aug & (choices == 2)
                if mask_blur.any():
                    x_subset = x_image[mask_blur]
                    x_down = blur(x_subset)
                    x_up = F.interpolate(x_down, size=x_subset.shape[-2:], mode="bilinear", align_corners=False)
                    x_image[mask_blur] = x_up

                x = (x_image - mean) / std

        out = model(pixel_values=x, labels=y)
        loss = out.loss
        logits = out.logits
        preds = logits.argmax(dim=1)
        loss.backward()
        pt.nn.utils.clip_grad_norm_(model.parameters(), max_norm=GRAD_CLIP)
        optimizer.step()

        total_loss += loss.item() * y.size(0)
        total_correct += (preds == y).sum().item()
        total_examples += y.size(0)

        pbar.set_postfix(loss=f"{loss.item():.4f}", acc=f"{100 * total_correct / total_examples:.2f}%")

    return total_loss / total_examples, total_correct / total_examples

@pt.no_grad()
def evaluate(model, loader, device):
    model.eval()
    total_loss = 0.0
    total_correct = 0
    total_examples = 0
    for batch in loader:
        x = batch["pixel_values"].to(device, non_blocking=True)
        y = batch["labels"].to(device, non_blocking=True)
        out = model(pixel_values=x, labels=y)
        loss = out.loss
        logits = out.logits 

        total_loss += loss.item() * y.size(0)
        total_correct += (logits.argmax(dim=1) == y).sum().item()
        total_examples += y.size(0)
    return total_loss / total_examples, total_correct / total_examples

if __name__ == "__main__":
    DEVICE = get_device()
    train_loader, test_loader = get_loaders()

    cfg = NulaConfig(block_channels=(128, 256, 512), classifier_hidden_dim=512, use_se=True)
    model = NulaForImageClassification(cfg).to(DEVICE)

    optimizer = pt.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
    warmup = LinearLR(optimizer, start_factor=0.1, end_factor=1.0, total_iters=5)
    cosine = CosineAnnealingLR(optimizer, T_max=45)
    scheduler = SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[5])
    
    MEAN = pt.tensor([0.5, 0.5, 0.5], device=DEVICE).view(1, 3, 1, 1)
    STD = pt.tensor([0.5, 0.5, 0.5], device=DEVICE).view(1, 3, 1, 1)
    GLOBAL_POOL_BLUR = BlurPool2d(channels=cfg.in_channels, stride=2).to(DEVICE)
    best_val_acc = 0.0

    os.makedirs(CHECKPOINT_DIR, exist_ok=True)

    for epoch in range(1, NUM_EPOCHS + 1):
        train_loss, train_acc = train_one_epoch(
            model, 
            train_loader, 
            optimizer, 
            DEVICE, 
            MEAN, 
            STD, 
            GLOBAL_POOL_BLUR
        )
        val_loss, val_acc = evaluate(model, test_loader, DEVICE)
        scheduler.step()

        if epoch % SAVE_EVERY == 0:
            model.save_pretrained(f"{CHECKPOINT_DIR}/epoch{epoch}")

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            model.save_pretrained(BEST_MODEL_DIR)
            print(f"new best: {100 * best_val_acc:.2f}%")

        current_lr = optimizer.param_groups[0]["lr"]
        print(f"|{'-'*60}|")
        print(f"epoch [{epoch}/{NUM_EPOCHS}] | lr: {current_lr:.6f}")
        print(f"train {train_loss:.4f}  {train_acc*100:.2f}%")
        print(f"val {val_loss:.4f}  {val_acc*100:.2f}%")
        print(f"best {best_val_acc*100:.2f}%")
        print()