import copy

N_EPOCHS = 10

plot_losses = []
scheduler_counter = 0
best_model = None

for epoch in range(N_EPOCHS):
    # Training
    model.train()
    loss_list = []
    acc_list = []

    for batch_i, sample in enumerate(dataloader):
        x, y = sample["image"][:, :3, :, :], sample["mask"]
        sample = None
        pred_mask = model(x.to(device).float())
        loss = criterion(pred_mask, y.to(device).type(torch.int64))

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        loss_list.append(loss.cpu().detach().numpy())
        acc_list.append(acc(y, pred_mask).numpy())

        print(
            "\r[Epoch %d/%d] [Batch %d/%d] [Loss: %f (%f)]"
            % (
                epoch,
                N_EPOCHS,
                batch_i,
                len(dataloader),
                loss.cpu().detach().numpy(),
                np.mean(loss_list),
            )
        )

    scheduler_counter += 1

    # Testing (using same dataloader for now)
    model.eval()
    val_loss_list = []
    val_acc_list = []

    for batch_i, sample in enumerate(dataloader):
        x, y = sample["image"][:, :3, :, :], sample["mask"]
        sample = None
        with torch.no_grad():
            pred_mask = model(x.to(device).float())
        val_loss = criterion(pred_mask, y.to(device).type(torch.int64))
        val_loss_list.append(val_loss.cpu().detach().numpy())
        val_acc_list.append(acc(y, pred_mask).numpy())

    print(' epoch {} - loss : {:.5f} - acc : {:.2f} - val loss : {:.5f} - val acc : {:.2f}'.format(
        epoch,
        np.mean(loss_list),
        np.mean(acc_list),
        np.mean(val_loss_list),
        np.mean(val_acc_list)
    ))

    plot_losses.append([epoch, np.mean(loss_list), np.mean(val_loss_list)])

    compare_loss = np.mean(val_loss_list)
    is_best = compare_loss < min_loss

    if is_best:
        scheduler_counter = 0
        min_loss = min(compare_loss, min_loss)
        torch.save(model.state_dict(), './saved_models/unet_epoch_{}_{:.5f}.pt'.format(epoch, np.mean(val_loss_list)))
        best_model = copy.deepcopy(model)

    if scheduler_counter > 5:
        lr_scheduler.step()
        print(f"lowering learning rate to {optimizer.param_groups[0]['lr']}")
        scheduler_counter = 0