import variables
import torch
import torch.optim as optim

from models import YoloBlock
from tqdm import tqdm

from utils import MAP, CTF, get_eval_frames, save_model, load_model, check_class_accuracy, get_loaders
from losses import Loss
import warnings
warnings.filterwarnings("ignore")

torch.backends.cudnn.benchmark = True


def train_fcn(train_loader, model, optimizer, loss_fn, resizer, resized_anchors):
    loop = tqdm(train_loader, leave=True)
    losses = []
    for batch_idx, (x, y) in enumerate(loop):
        x = x.to(variables.DEVICE)
        y0, y1, y2 = (
            y[0].to(variables.DEVICE),
            y[1].to(variables.DEVICE),
            y[2].to(variables.DEVICE),
        )

        with torch.cuda.amp.autocast():
            try:
                out = model(x)
                loss = (
                    loss_fn(out[0], y0, resized_anchors[0])
                    + loss_fn(out[1], y1, resized_anchors[1])
                    + loss_fn(out[2], y2, resized_anchors[2])
                )
                losses.append(loss.item())
                optimizer.zero_grad()
                resizer.scale(loss).backward()
                resizer.step(optimizer)
                resizer.update()
                mean_loss = sum(losses) / len(losses)
            except Exception as e:
                print(e)
                mean_loss = 'NaN'

        loop.set_postfix(loss=mean_loss)



def app():
    model = YoloBlock(no_classes=variables.NO_CLASSES).to(variables.DEVICE)
    optimizer = optim.Adam(
        model.parameters(), lr=variables.LR, weight_decay=variables.WEIGHT_DECAY
    )
    loss_fcn = Loss()
    resizer = torch.cuda.amp.GradScaler()


    if variables.LOAD_TRAINED_MODEL:
        load_model(
            variables.TRAIN_MODEL_FILE, model, optimizer, variables.LR
        )

    resized_anchors = (
        torch.tensor(variables.ANCHORS)
        * torch.tensor(variables.SPLIT_SIZE).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
    ).to(variables.DEVICE)

    for epoch in range(variables.NO_EPOCHS):
        train_loader, test_loader, train_eval_loader = get_loaders(
            train_path=r"data/kitti_single/training_data.csv", test_path="data/kitti_single/testing_data.csv"
        )
        # plot_couple_examples(model, test_loader, 0.6, 0.5, resized_anchors)
        train_fcn(train_loader, model, optimizer, loss_fcn, resizer, resized_anchors)

        if variables.SAVE_TRAINED_MODEL:
           save_model(model, optimizer, filename=f"trained_models/{epoch}_epoch_cars_kitti.pth.tar")

        if epoch > 0 and epoch % 3 == 0:
            try:
                check_class_accuracy(model, test_loader, threshold=variables.CONF_LIM)
                pred_boxes, true_boxes = get_eval_frames(test_loader, model, iou_lim=variables.IOU_LIM,
                                                         anchors=variables.ANCHORS, prob_lim=variables.CONF_LIM)
                map = MAP(pred_boxes, true_boxes, iou_lim=variables.MAP_IOU_LIM,
                          frame_format="1_point", no_classes=variables.NUM_LIM)
                print(f"MAP: {map.item()}")
                model.train()
            except Exception as e:
                print(e)
                pass


if __name__ == "__main__":
    app()
