import torch

from parse_config import parse_file
import data
from net import Net

class Runner:
    def __init__(self, file, metrics):
        self.metrics = metrics
        self.setup(file)

    def setup(self, file):
        (
            net_config,
            lr,
            self.epochs,
            batch_size,
            augmentations,
            target_transform,
            self.loss_function,
            self.count_correct
        ) = parse_file(file)

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.train_loader, self.test_loader = data.get_loaders(
                augmentations,
                target_transform,
                batch_size)

        self.net = Net(**net_config).to(self.device)
        self.optimizer = torch.optim.Adam(self.net.parameters(), lr=lr)

    def run(self):
        for epoch in range(self.epochs):
            self.train_step()
            self.test_step()

    def train_step(self):
        """
        Performs one epoch of training.
        """
        self.net.train()
        total_loss = 0
        number_batches = 0
        for batch_idx, (data, target) in enumerate(self.train_loader):
            data, target = data.to(self.device), target.to(self.device)
            number_batches += 1
            self.optimizer.zero_grad()
            output = self.net(data)
            loss = self.loss_function(output, target)
            loss.backward()
            self.optimizer.step()
            total_loss += loss.detach()
            if batch_idx % 10 == 0:
                print('Training: [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    batch_idx * len(data), len(self.train_loader.dataset),
                    100. * batch_idx / len(self.train_loader), loss.item()))
        self.metrics.add_train_loss(total_loss/number_batches)

    def test_step(self):
        self.net.eval()
        test_loss = 0
        correct = 0
        number_batches = 0
        with torch.no_grad():
            for data, target in self.test_loader:
                data, target = data.to(self.device), target.to(self.device)
                number_batches += 1
                output = self.net(data)
                test_loss += self.loss_function(output, target)
                correct += self.count_correct(output, target)

        test_loss /= number_batches
        accuracy = correct / len(self.test_loader.dataset)

        print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            test_loss, correct, len(self.test_loader.dataset), 100. * accuracy
            ))

        self.metrics.add_test_metrics(test_loss, accuracy)