import numpy as np
import torch

from DataLoader import get_image_loader
from Net import ImageNN

# 01.05.22 -- 0.5h
from netio import save_model, load_model


def get_train_device():
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f'Device to train net: {device}')
    return torch.device(device)


def train_model():
    # Set a known random seed for reproducibility
    np.random.seed(0)
    torch.manual_seed(0)
    device = get_train_device()

    # Load datasets
    train_loader, test_loader = get_image_loader("training/")
    nn = ImageNN(n_in_channels=3)  # todo pass size ason.
    nn.train()  # init with train mode
    nn.to(device)  # send net to device available

    optimizer = torch.optim.SGD(nn.parameters(), lr=0.1)  # todo adjust parameters and lr
    loss_function = torch.nn.MSELoss()
    n_epochs = 15  # todo epcchs here

    # todo look wtf is that
    nn.double()

    train_sample_size = len(train_loader)
    losses = []
    best_eval_loss = 0
    for epoch in range(n_epochs):
        print(f"Epoch {epoch}/{n_epochs}\n")
        i = 0
        for input_tensor, target_tensor in train_loader:
            input_tensor.to(device)
            target_tensor.to(device)

            output = nn(input_tensor)  # get model output (forward pass)
            loss = loss_function(output, target_tensor)  # compute loss given model output and true target
            loss.backward()  # compute gradients (backward pass)
            optimizer.step()  # perform gradient descent update step
            optimizer.zero_grad()  # reset gradients
            losses.append(loss.item())

            print(
                f'\rTraining epoch {epoch} [{i}/{train_sample_size * train_loader.batch_size}] (curr loss: {loss.item():.3})',
                end='')
            i += train_loader.batch_size

            # eval model every 500th element
            if i % 500 == 0:
                print(f"\nEvaluating model")
                eval_loss = eval_model(nn, test_loader, loss_function, device)
                print(f"Evalution loss={eval_loss}")
                if eval_loss > best_eval_loss:
                    best_eval_loss = eval_loss
                    save_model(nn)

    # switch net to eval mode
    print(eval_model(nn, test_loader, loss_function, device=device))


# func to evaluate our trained model
def eval_model(model: torch.nn.Module, dataloader: torch.utils.data.DataLoader, loss_fn, device: torch.device):
    # switch to eval mode
    model.eval()
    loss = .0
    # disable gradient calculations
    with torch.no_grad():
        i = 0
        for input, target in dataloader:
            input.to(device)
            target.to(device)

            out = model(input)
            loss += loss_fn(out, target).item()
            print(f'\rEval prog[{i}/{len(dataloader) * dataloader.batch_size}]', end='')
            i += dataloader.batch_size
    print()
    loss /= len(dataloader)
    model.train()
    return loss


def apply_model():
    model = load_model()

    pass