#!/usr/bin/env python3

import torch
from torch.autograd import Function

from .. import settings
from ..utils.cholesky import psd_safe_cholesky
from ..utils.permutation import apply_permutation, inverse_permutation


class PivotedCholesky(Function):
    @staticmethod
    def forward(ctx, representation_tree, max_iter, error_tol, *matrix_args):
        ctx.representation_tree = representation_tree
        matrix = ctx.representation_tree(*matrix_args)
        batch_shape = matrix.shape[:-2]
        matrix_shape = matrix.shape[-2:]

        if error_tol is None:
            error_tol = settings.preconditioner_tolerance.value()

        # Need to get diagonals. This is easy if it's a LazyTensor, since
        # LazyTensor.diag() operates in batch mode.
        matrix_diag = matrix._approx_diag()

        # Make sure max_iter isn't bigger than the matrix
        max_iter = min(max_iter, matrix_shape[-1])

        # What we're returning
        L = torch.zeros(*batch_shape, max_iter, matrix_shape[-1], dtype=matrix.dtype, device=matrix.device)
        orig_error = torch.max(matrix_diag, dim=-1)[0]
        errors = torch.norm(matrix_diag, 1, dim=-1) / orig_error

        # The permutation
        permutation = torch.arange(0, matrix_shape[-1], dtype=torch.long, device=matrix_diag.device)
        permutation = permutation.repeat(*batch_shape, 1)

        # Maybe log
        if settings.verbose_linalg.on():
            settings.verbose_linalg.logger.debug(
                f"Running Pivoted Cholesky on a {matrix.shape} RHS for {max_iter} iterations."
            )

        m = 0
        while (m == 0) or (m < max_iter and torch.max(errors) > error_tol):
            # Get the maximum diagonal value and index
            # This will serve as the next diagonal entry of the Cholesky,
            # as well as the next entry in the permutation matrix
            permuted_diags = torch.gather(matrix_diag, -1, permutation[..., m:])
            max_diag_values, max_diag_indices = torch.max(permuted_diags, -1)
            max_diag_indices = max_diag_indices + m

            # Swap pi_m and pi_i in each row, where pi_i is the element of the permutation
            # corresponding to the max diagonal element
            old_pi_m = permutation[..., m].clone()
            permutation[..., m].copy_(permutation.gather(-1, max_diag_indices.unsqueeze(-1)).squeeze_(-1))
            permutation.scatter_(-1, max_diag_indices.unsqueeze(-1), old_pi_m.unsqueeze(-1))
            pi_m = permutation[..., m].contiguous()

            # Populate L[..., m, m] with the sqrt of the max diagonal element
            L_m = L[..., m, :]  # Will be all zeros -- should we use torch.zeros?
            L_m.scatter_(-1, pi_m.unsqueeze(-1), max_diag_values.sqrt().unsqueeze_(-1))

            # Populater L[... m:, m] with L[..., m:, m] * L[..., m, m].sqrt()
            if m + 1 < matrix_shape[-1]:
                # Get next row of the permuted matrix
                row = apply_permutation(matrix, pi_m.unsqueeze(-1), right_permutation=None).squeeze(-2)
                pi_i = permutation[..., m + 1 :].contiguous()

                L_m_new = row.gather(-1, pi_i)
                if m > 0:
                    L_prev = L[..., :m, :].gather(-1, pi_i.unsqueeze(-2).repeat(*(1 for _ in batch_shape), m, 1))
                    update = L[..., :m, :].gather(
                        -1, pi_m.view(*pi_m.shape, 1, 1).repeat(*(1 for _ in batch_shape), m, 1)
                    )
                    L_m_new -= torch.sum(update * L_prev, dim=-2)

                L_m_new /= L_m.gather(-1, pi_m.unsqueeze(-1))
                L_m.scatter_(-1, pi_i, L_m_new)

                matrix_diag_current = matrix_diag.gather(-1, pi_i)
                matrix_diag.scatter_(-1, pi_i, matrix_diag_current - L_m_new**2)
                L[..., m, :] = L_m

                # Keep track of errors - for potential early stopping
                errors = torch.norm(matrix_diag.gather(-1, pi_i), 1, dim=-1) / orig_error

            m = m + 1

        # Save items for backward pass, and return output
        ctx.save_for_backward(permutation, permutation[..., :m], *matrix_args)
        return L[..., :m, :].transpose(-1, -2).contiguous(), permutation

    def backward(ctx, grad_output, _):
        full_permutation, short_permutation, *_matrix_args = ctx.saved_tensors
        _inverse_permutation = inverse_permutation(full_permutation)
        m = short_permutation.size(-1)  # The rank of the pivoted Cholesky factor

        with torch.enable_grad():
            # Create a new set of matrix args that we can backpropagate through
            matrix_args = [matrix_arg.detach().requires_grad_(True) for matrix_arg in _matrix_args]

            # Create new linear operator using new matrix args
            matrix = ctx.representation_tree(*matrix_args)

            # Get Krows - rows of the matrix that were selected for pivoted Cholesky
            Krows = apply_permutation(matrix, full_permutation, short_permutation)

            # Get L - Cholesky factor of K[..., :m, :m]
            L = psd_safe_cholesky(Krows[..., :m, :])

            # Compute (Krows * L^{-T}) - the (pivoted) result of Pivoted Cholesky
            res_pivoted = torch.cat(
                [L, torch.triangular_solve(Krows[..., m:, :].transpose(-1, -2), L, upper=False)[0].transpose(-1, -2)],
                dim=-2,
            )
            res = apply_permutation(res_pivoted, left_permutation=_inverse_permutation, right_permutation=None)

            # Now compute the backward pass of res
            res.backward(gradient=grad_output)

        # Define gradient placeholders
        return tuple([None, None, None] + [matrix_arg.grad for matrix_arg in matrix_args])
