import torch
import torch.nn as nn
import torch.nn.functional as F
import math


def zero_grad(grad):
    if grad is not None:
        return grad * 0
    return None


class MCLoss(nn.Module):
    def __init__(self, base_loss: nn.Module, var: float = 1.0):
        super(MCLoss, self).__init__()
        self.base_loss = base_loss
        self.scale = math.sqrt(var)

    def forward(self, logits):
        if isinstance(self.base_loss, nn.MSELoss):
            dist = torch.distributions.normal.Normal(logits, scale=self.scale)
            with torch.no_grad():
                targets = dist.sample()
            return self.base_loss(logits, targets)

        elif isinstance(self.base_loss, nn.CrossEntropyLoss):
            probs = F.softmax(logits, dim=-1)
            log_probs = F.log_softmax(logits, dim=-1)
            dist = torch.distributions.Categorical(probs)
            with torch.no_grad():
                targets = dist.sample()
            return F.nll_loss(
                log_probs.view(-1, log_probs.size(-1)),
                targets.view(-1),
                reduction=self.base_loss.reduction,
            )
        else:
            raise ValueError("Unsupported loss type")

    @classmethod
    def quick_call(cls, base_loss: nn.Module, outputs: torch.Tensor, var: float = 1.0):
        with torch.no_grad():
            if isinstance(base_loss, nn.MSELoss):
                targets = torch.normal(mean=outputs, std=math.sqrt(var))
            else:
                targets = torch.multinomial(F.softmax(outputs, dim=-1), num_samples=1).squeeze()

        return base_loss(outputs, targets)


def preconditioned_cg(
    A,
    b,
    preconditioner=None,
    x_0=None,
    tol=1e-8,
    max_iter=100,
    verbose=False,
):
    if x_0 is None:
        x_0 = torch.zeros_like(b, device=b.device)

    if preconditioner is None:
        preconditioner = lambda v: v

    x = x_0.clone()
    r_k = A(x) - b
    y_k = preconditioner(r_k)
    p_k = -y_k.clone()

    for k in range(max_iter):
        Ap = A(p_k)
        alpha = torch.dot(r_k, y_k) / torch.dot(p_k, Ap)
        x = x + alpha * p_k
        r_k_prev = r_k.clone()
        r_k = r_k + alpha * Ap
        y_k_prev = y_k.clone()
        y_k = preconditioner(r_k)
        beta = torch.dot(r_k, y_k) / torch.dot(r_k_prev, y_k_prev)
        p_k = beta * p_k - y_k

        if verbose:
            print(f"iter: {k}, residual: {torch.norm(r_k)}")
        if torch.norm(r_k) < tol:
            break

    return x, k + 1


def prepare_precondition_step(model, preconditioner, train_inputs):

    if preconditioner.mc_samples == 1:
        if preconditioner.should_update_grads:
            preconditioner.set_train_stage("fisher_pass")
            dummy_loss = preconditioner.mc_loss(model(train_inputs))
            dummy_loss.backward()
            model.zero_grad()
            preconditioner.set_train_stage("gradient_pass")
    else:
        raise ValueError("Not Implemented Yet")

    return {}
