
from ignite.engine.engine import Engine, State, Events
from ignite._utils import convert_tensor


def _prepare_batch(batch, device=None):
    xs, ys = batch
    return [convert_tensor(x, device=device) for x in xs],\
           [convert_tensor(y, device=device) for y in ys]


def create_supervised_trainer(model, optimizer, weights, loss_fns, device=None):
    """
    Factory function for creating a trainer for supervised models

    Args:
        model (`torch.nn.Module`): the model to train
        optimizer (`torch.optim.Optimizer`): the optimizer to use
        loss_fn (torch.nn loss function): the loss function to use
        device (str, optional): device type specification (default: None).
            Applies to both model and batches.

    Returns:
        Engine: a trainer engine with supervised update function
    """
    if device:
        model.to(device)

    def _update(engine, batch):
        model.train()
        optimizer.zero_grad()

        x, ys = _prepare_batch(batch, device=device)
        y_preds = model(xs)

        losses = [loss_fn(y_pred, y) for loss_fn, y_pred, y in zip(loss_fns, y_preds, ys)]

        loss = 0
        for w, l in zip(weights, losses):
            loss += w.value * l

        loss.backward()
        optimizer.step()

        return loss.item()

    return Engine(_update)


def create_supervised_evaluator(model, metrics={}, device=None):
    """
    Factory function for creating an evaluator for supervised models

    Args:
        model (`torch.nn.Module`): the model to train
        metrics (dict of str - :class:`ignite.metrics.Metric`): a map of metric names to Metrics
        device (str, optional): device type specification (default: None).
            Applies to both model and batches.

    Returns:
        Engine: an evaluator engine with supervised inference function
    """
    if device:
        model.to(device)

    def _inference(engine, batch):
        model.eval()
        with torch.no_grad():
            x, y = _prepare_batch(batch, device=device)
            y_pred = model(x)
            return y_pred, y

    engine = Engine(_inference)

    for name, metric in metrics.items():
        metric.attach(engine, name)

    return engine
