import time
from functools import reduce

import numpy as np
import torch

import tntorch as tn


def optimize(
    tensors,
    loss_function,
    optimizer=torch.optim.Adam,
    tol=1e-4,
    max_iter=1e4,
    print_freq=500,
    verbose=True,
):
    """
    High-level wrapper for iterative learning.

    Default stopping criterion: either the absolute (or relative) loss improvement must fall below `tol`.
    In addition, the rate loss improvement must be slowing down.

    :param tensors: one or several tensors; will be fed to `loss_function` and optimized in place
    :param loss_function: must take `tensors` and return a scalar (or tuple thereof)
    :param optimizer: one from https://pytorch.org/docs/stable/optim.html. Default is torch.optim.Adam
    :param tol: stopping criterion
    :param max_iter: default is 1e4
    :param print_freq: progress will be printed every this many iterations
    :param verbose:
    """

    if not isinstance(tensors, (list, tuple)):
        tensors = [tensors]
    parameters = []
    for t in tensors:
        if isinstance(t, tn.Tensor):
            if t.batch:
                raise ValueError("Batched tensors are not supproted.")

            parameters.extend([c for c in t.cores if c.requires_grad])
            parameters.extend([U for U in t.Us if U is not None and U.requires_grad])
        elif t.requires_grad:
            parameters.append(t)
    if len(parameters) == 0:
        raise ValueError(
            "There are no parameters to optimize. Did you forget a requires_grad=True somewhere?"
        )

    optimizer = optimizer(parameters)
    losses = []
    converged = False
    start = time.time()
    iter = 0
    while True:
        optimizer.zero_grad()
        loss = loss_function(*tensors)
        if not isinstance(loss, (tuple, list)):
            loss = [loss]
        total_loss = sum(loss)
        total_loss.backward(retain_graph=True)
        optimizer.step()

        losses.append(total_loss.detach())

        if len(losses) >= 2:
            delta_loss = losses[-1] - losses[-2]
        else:
            delta_loss = float("-inf")
        if (
            iter >= 2
            and tol is not None
            and (losses[-1] <= tol or -delta_loss / losses[-1] <= tol)
            and losses[-2] - losses[-1] < losses[-3] - losses[-2]
        ):
            converged = True
            break
        if iter == max_iter:
            break
        if verbose and iter % print_freq == 0:
            print(
                "iter: {: <{}} | loss: ".format(iter, len("{}".format(max_iter))),
                end="",
            )
            print(" + ".join(["{:10.6f}".format(l.item()) for l in loss]), end="")
            if len(loss) > 1:
                print(" = {:10.4}".format(losses[-1].item()), end="")
            print(" | total time: {:9.4f}".format(time.time() - start))

        iter += 1
    if verbose:
        print("iter: {: <{}} | loss: ".format(iter, len("{}".format(max_iter))), end="")
        print(" + ".join(["{:10.6f}".format(l.item()) for l in loss]), end="")
        if len(loss) > 1:
            print(" = {:10.4}".format(losses[-1].item()), end="")
        print(" | total time: {:9.4f}".format(time.time() - start), end="")
        if converged:
            print(" <- converged (tol={})".format(tol))
        else:
            print(" <- max_iter was reached: {}".format(max_iter))


def dof(t):
    """
    Compute the number of degrees of freedom of a tensor network.

    It is the sum of sizes of all its tensor nodes that have the requires_grad=True flag.

    :param t: input tensor

    :return: an integer
    """

    result = 0
    for n in range(t.dim()):
        if t.cores[n].requires_grad:
            result += t.cores[n].numel()
        if t.Us[n] is not None and t.Us[n].requires_grad:
            result += t.Us[n].numel()
    return result
