import numpy as np

class TrainingResult(object):

    def __init__(self, training_set_size, num_epochs, batch_size=-1, batch_stride=-1): 
        self.training_set_size = training_set_size
        self.num_epochs = num_epochs
        if batch_size > 0:
            if batch_stride < 0:
                batch_stride = batch_size
            self.num_batches = training_set_size // batch_stride
        else:
            self.num_batches = 1
        self.batch_size = batch_size
        self.batch_stride = batch_stride
        # Current epoch and batch index.
        self.epoch = 0
        self.batch_index = 0
        # Loss for each batch in each epoch.
        self.epoch_loss = np.full((self.num_epochs, self.num_batches, self.training_set_size), np.inf)
        # Step size for each epoch.
        self.epoch_step_size = np.full(self.num_epochs, np.inf)
        # Accuracy metrics by set name (training, validation, etc.)
        self.accuracies = {}
        # Named metrics for each stage.
        self.stage_metrics = {}
        # Named models for each stage.
        self.stage_models = {}

    def set_epoch(self, epoch):
        self.epoch = epoch
        self.batch_index = 0

    def set_batch_index(self, batch_index):
        self.batch_index = batch_index

    def add_loss(self, loss):
        self.epoch_loss[self.epoch][self.batch_index] = np.copy(loss)

    def add_step_size(self, step_size):
        self.epoch_step_size[self.epoch] = step_size

    def add_accuracy(self, set_name, accuracy):
        if set_name not in self.accuracies:
            self.accuracies[set_name] = np.full(self.num_epochs, np.inf)
        self.accuracies[set_name][self.epoch] = accuracy
        
    def add_stage_metric(self, stage_name, metric_name, value):
        if stage_name not in self.stage_metrics:
            self.stage_metrics[stage_name] = {}
        if metric_name not in self.stage_metrics[stage_name]:
            self.stage_metrics[stage_name][metric_name] = np.full(self.num_epochs, np.inf)
        self.stage_metrics[stage_name][metric_name][self.epoch] = value

    def stage_names(self):
        return set(self.stage_metrics.keys())

    def metric_names(self, stage_name):
        return set(self.stage_metrics[stage_name].keys())

    def aggregate_loss(self):
        aggregate_loss = np.full(self.num_epochs, np.inf)
        for epoch in range(self.num_epochs):
            epoch_sum = 0.0
            for batch in range(self.num_batches):
                epoch_sum += np.sum(self.epoch_loss[epoch][batch])
            aggregate_loss[epoch] = epoch_sum / (self.training_set_size * self.num_batches)
        return aggregate_loss

    def loss(self, epoch, batch_index=0):
        return self.epoch_loss[epoch][batch_index]

    def step_sizes(self):
        return self.epoch_step_size

    def accuracy(self, set_name):
        return self.accuracies[set_name]

    def stage_metric(self, stage_name, metric_name):
        return self.stage_metrics[stage_name][metric_name]
