"""Main module."""
import logging
import torch
from torch.optim import Optimizer
from pathlib import Path

from .utils import preconditioned_cg, MCLoss, zero_grad


class RedFish(Optimizer):
    def __init__(
        self,
        model,
        loss_fn,
        damping: float = 0.1,
        update_freq: int = 1,
        accumulate_grad_batches: int = 1,
        damping_update_factor: float = 1.0,
        damping_update_freq: int = 1,
        eig_damp: bool = False,
        method: str = "exact",
        cg_steps: int = 1,
        num_mc_samples: int = 1,
        mc_var: float = 0.01,
        layer_groups_index: list[int] | str = None,
        override_grad: bool = False,
        adjust_by_damping: bool = True,
        debug: bool = False,
        debug_path: str = "data/debug",
        verbose: bool = False,
    ):
        """
        RedFish preconditioner for PyTorch.
        """

        if method not in ["exact", "cg"]:
            raise ValueError("Method must be 'exact' or 'cg'.")
        if damping <= 0:
            raise ValueError("Damping must be positive.")
        if update_freq < 1:
            raise ValueError("Update frequency must be at least 1.")
        if accumulate_grad_batches < 1:
            raise ValueError("Accumulate grad batches must be at least 1.")

        # Initialize the optimizer
        self.damping = float(damping)
        self._eig_damp = eig_damp
        if self._eig_damp:
            self._damping_factor = self.damping

        self.update_freq = int(update_freq)
        self.accumulate_grad_batches = int(accumulate_grad_batches)
        self.adjust_by_damping = adjust_by_damping
        self.params = []
        self.reduction_ratio = None
        self.method = method
        self.verbose = verbose

        # Protected attributes
        self._fwd_handles = []
        self._bwd_handles = []
        self._iteration_counter = 0
        self._batch_size = None
        self._ones = None
        self._ones_mat = None
        self._damp_upf = damping_update_factor
        self._damp_freq = damping_update_freq
        self._predicted_loss = 0
        self._pred_loss_accum = None
        self._adaptive_damping = damping_update_factor != 1.0
        self._cg_steps = cg_steps
        self._batch_count = 0
        self._batch_factor = 1 / self.accumulate_grad_batches
        self._num_params = 0

        self.override_grad = override_grad
        self._current_layer = 0
        self._current_group = -1
        if layer_groups_index is not None:
            if isinstance(layer_groups_index, str):
                self._group_index = list(range(1000))
                self._num_groups = 0
            else:
                self._num_groups = len(layer_groups_index)
                self._group_index = layer_groups_index
        else:
            self._num_groups = 1
            self._group_index = [0]

        self.mc_loss = MCLoss(base_loss=loss_fn, var=mc_var)

        model.apply(self._register_hooks)
        if self.verbose:
            logging.warning(f"Number of parameters: {self._num_params:,d}")
            for p in self.params:
                logging.warning(f"{p['layer_type']}=>{p['group_ix']}")

        self._train_stage = "gradient_pass"
        self.mc_samples = num_mc_samples
        if self._num_groups == 0:
            self._num_groups = len(self.params)
        self._group_keys = [f"group_{i}" for i in range(self._num_groups)]
        self._delta_mc = dict.fromkeys(self._group_keys)
        self._Q = dict.fromkeys(self._group_keys)
        self._R_hat = dict.fromkeys(self._group_keys)
        # This can be used to store intermediate values for debugging or adaptive damping revision.
        self.debug = debug
        if self.debug:
            self._debug_file = Path().cwd() / (debug_path + "/rf_debug.pt")
            self._layer_debug_count = 0

        # Register the model parameters
        super(RedFish, self).__init__(self.params, {})

    supported_layers = [
        "Linear",
        "Conv2d",
    ]

    def step(self, closure=None):
        if self.should_update_grads:
            # build fisher gram
            self._build_fisher_gram()
            # solve for nu
            self._solve_connection_system()
            # precondition gradients
            self._precondition()

        self._iteration_counter += 1
        self._batch_size = None
        if self.should_accumulate_grads:
            self._batch_count += 1
        else:
            self._batch_count = 1
        self.set_train_stage("gradient_pass")

    def _register_hooks(self, mod):
        """Registers forward and backward hooks for supported layers to save per-sample inputs and gradients.
        More info about the hooks can be found in the links below
        https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_forward_pre_hook
        https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_full_backward_hook.
        """
        mod_class = mod.__class__.__name__

        if mod_class in self.supported_layers:
            if self._current_layer in self._group_index:
                self._current_group += 1

            if mod_class in ["Linear", "Conv2d"]:
                handle = mod.register_forward_pre_hook(self._save_input)
                self._fwd_handles.append(handle)
                handle = mod.register_full_backward_hook(self._save_grad_output)
                self._bwd_handles.append(handle)

                params = [mod.weight]
                if self.override_grad:
                    handle = mod.weight.register_hook(zero_grad)
                    self._bwd_handles.append(handle)

                self._num_params += mod.weight.numel()
                if mod.bias is not None:
                    params.append(mod.bias)
                    if self.override_grad:
                        handle = mod.bias.register_hook(zero_grad)
                        self._bwd_handles.append(handle)
                    self._num_params += mod.bias.numel()
                d = {
                    "params": params,
                    "mod": mod,
                    "layer_type": mod_class,
                    "group_ix": self._current_group,
                }
                self.params.append(d)
                self._current_layer += 1

    def enable_hooks(self, model):
        """Enables hooks for the model."""
        if self._fwd_handles or self._bwd_handles:
            logging.warning("Hooks are already registered. Skipping.")
            return
        model.apply(self._register_hooks)

    def disable_hooks(self):
        """Disables hooks for the model."""
        for handle in self._fwd_handles + self._bwd_handles:
            handle.remove()
        self._fwd_handles.clear()
        self._bwd_handles.clear()

    def set_train_stage(self, stage_key):
        """Sets the current training stage for the preconditioner.
        This method is the preferred way to switch between the "fisher_pass" and "gradient_pass".
        The "fisher_pass" stage is used to compute the reduced Fisher information matrix (RFIM) estimate and related quantities,
        while the "gradient_pass" stage is used to compute the actual gradients for the model parameters.
        """
        self._train_stage = stage_key

    def _save_input(self, mod, x):
        """Saves per-sample inputs of layer."""
        if (
            mod.training
            and self._train_stage == "fisher_pass"
            and self.should_update_grads
        ):
            mod_name = mod.__class__.__name__
            A = x[0].detach()

            if mod_name == "Conv2d":
                A = torch.nn.functional.unfold(
                    A,
                    mod.kernel_size,
                    stride=mod.stride,
                    padding=mod.padding,
                    dilation=mod.dilation,
                )

                b = A.shape[0]
                k = A.shape[1]
                l = A.shape[2]
                c_o = mod.out_channels
                # This flag determines if we can use a more efficient method to compute the RFIM for convolutional layers.
                self.state[mod]["reducible"] = (l**2) * (k + c_o) < k * c_o
                self.state[mod]["A"] = A
            else:
                self.state[mod]["A"] = A.squeeze(1)

            if self._batch_size is None:
                self._batch_size = A.size(0)
                self._ones = torch.ones(size=(self._batch_size, 1), device=A.device)
                self._ones_mat = None

                # Initialize the intermediate matrices Q and R_hat defined in section 3.1.2 of the dissertation
                # for the current batch size and group keys.
                # We also ensure they are on the correct device and have the correct dimensions.
                self._Q = {
                    gk: torch.zeros(
                        (self._batch_size, self._batch_size), device=A.device
                    )
                    for gk in self._group_keys
                }

                self._R_hat = {
                    gk: torch.zeros(
                        (self._batch_size, self._batch_size), device=A.device
                    )
                    for gk in self._group_keys
                }

    def _save_grad_output(self, mod, grad_input, grad_output):
        """Saves per-sample gradients of output of layer (with respect to parameters)."""
        if mod.training and self.should_update_grads:
            mod_name = mod.__class__.__name__
            save_key_suffix = "_hat" if self._train_stage == "fisher_pass" else ""
            dw_s = []

            for g in grad_output:
                if g is not None:
                    dw_s.append(g.detach().squeeze(1) * g.size(0))

            if dw_s:
                if len(dw_s) > 1:
                    dw_s = torch.cat(dw_s, dim=1)
                else:
                    dw_s = dw_s[0]

                if mod_name == "Conv2d":
                    G = dw_s.reshape(self._batch_size, mod.out_channels, -1)
                    if not self.state[mod]["reducible"]:
                        A = self.state[mod]["A"]
                        self.state[mod]["DW" + save_key_suffix] = torch.einsum(
                            "bkl,bml->bkm", G, A
                        )
                        if self._train_stage == "gradient_pass":
                            del self.state[mod]["A"]
                    else:
                        self.state[mod]["G" + save_key_suffix] = G
                else:
                    self.state[mod]["G" + save_key_suffix] = dw_s

    @torch.no_grad()
    def _build_fisher_gram(self):
        """Builds the Fisher Gram matrix (or its Monte Carlo estimate) for the current batch and layer groups."""
        for g_i, group in enumerate(self.param_groups):
            # For each layer group, we efficiently update the R_hat and Q matrices
            # using the per-sample activations and gradients saved by the hooks.

            layer_type = group["layer_type"]
            mod = group["mod"]
            state = self.state[mod]
            group_key = f"group_{group['group_ix']}"

            if layer_type == "Linear":
                A = state["A"]
                G = state["G"]
                G_hat = state["G_hat"]

                if A.dim() == 3:
                    R_A = torch.einsum("bpj,cpj->pbc", A, A)
                    R_G = torch.einsum("bpj,cpj->pbc", G_hat, G_hat)

                    if mod.bias is not None:
                        R_A += 1

                    self._R_hat[group_key] += (R_A * R_G).sum(dim=0)
                else:
                    R_A = A @ A.T
                    R_G = G_hat @ G_hat.T

                    if mod.bias is not None:
                        R_A += 1

                    self._R_hat[group_key] += R_A * R_G

                if G_hat.dim() == 3:
                    Q_G = torch.einsum("bpj,cpj->pbc", G_hat, G)
                    self._Q[group_key] += (R_A * Q_G).sum(dim=0)
                else:
                    Q_G = G_hat @ G.T
                    self._Q[group_key] += R_A * Q_G

                if self._adaptive_damping:
                    self._pred_loss_accum = torch.zeros(
                        size=(self._batch_size, 1), device=A.device
                    )
                    self._predicted_loss = 0

            elif layer_type == "Conv2d":
                if state["reducible"]:
                    A = state["A"]
                    G = state["G"]
                    G_hat = state["G_hat"]
                    R_A = torch.einsum("nkl,qkp->nqlp", A, A)
                    R_G = torch.einsum("nml,qmp->nqlp", G_hat, G_hat)
                    Q_G = torch.einsum("nml,qmp->nqlp", G_hat, G)
                    self._R_hat[group_key] += torch.einsum("nqlp->nq", R_A * R_G)
                    self._Q[group_key] += torch.einsum("nqlp->nq", R_A * Q_G)
                else:
                    DW_S = state["DW"].flatten(start_dim=1)
                    DW_S_hat = state["DW_hat"].flatten(start_dim=1)
                    self._R_hat[group_key] += DW_S_hat @ DW_S_hat.T
                    self._Q[group_key] += DW_S_hat @ DW_S.T

        for gk in self._group_keys:
            self._R_hat[gk] /= self._batch_size
            self._Q[gk] /= self._batch_size

            diag = torch.diagonal(self._R_hat[gk])
            if self._eig_damp:
                # Scale the damping by the trace of R_hat to make it more robust to scale changes in the RFIM.
                self.damping = self._damping_factor * diag.sum().item()

            diag += self.damping

    def _solve_connection_system(self):

        for gk, R_hat in self._R_hat.items():
            try:
                Qe = self._Q[gk] @ self._ones
                if self.method == "exact":
                    L = torch.linalg.cholesky(R_hat)
                    self._delta_mc[gk] = torch.cholesky_solve(-Qe, L)

                elif self.method == "cg":
                    self._delta_mc[gk], _ = preconditioned_cg(
                        lambda x: R_hat @ x,
                        -Qe.squeeze(),
                        max_iter=self._cg_steps,
                    )
                    self._delta_mc[gk] = self._delta_mc[gk].unsqueeze(1)

                self._delta_mc[gk] /= self._batch_size
                if self.debug:
                    torch.save(
                        {
                            f"delta_mc": self._delta_mc[gk].cpu(),
                        },
                        self._debug_file,
                    )
            except Exception as e:
                self._delta_mc[gk] = None
                logging.warning(
                    f"Solving system failed for step {self._iteration_counter} skipping preconditioning"
                )

    def _precondition(self):
        for group in self.param_groups:
            layer_type = group["layer_type"]
            if layer_type in ["Linear"]:
                self._precondition_linear(group)
            elif layer_type == "Conv2d":
                self._precondition_convolution(group)
            else:
                self._precondition_raw(group)

    def _precondition_linear(self, group):
        # For these layers we know grad = g a^T, which transforms the update
        # into the Hadamard product below.
        weight = group["params"][0]
        bias = group["params"][1] if len(group["params"]) == 2 else None
        state = self.state[group["mod"]]
        A = state["A"]
        G_hat = state["G_hat"]
        group_key = f"group_{group['group_ix']}"
        delta_mc = self._delta_mc[group_key]

        if delta_mc is not None:
            if A.dim() == 3:
                if bias is not None:
                    bias_slice = (
                        self._ones @ torch.ones((1, A.shape[1]), device=A.device)
                    ).unsqueeze(2)
                    A = torch.cat([A, bias_slice], dim=2)

                g = torch.einsum("bji,bjk,bi->ik", G_hat, A, delta_mc)

            else:
                if bias is not None:
                    A = torch.cat([A, self._ones], dim=1)
                D = torch.diag(delta_mc.squeeze(1))
                g = G_hat.T @ D @ A

            if self.debug:
                torch.save(
                    {
                        f"A_{self._layer_debug_count}": A.cpu(),
                        f"G_hat_{self._layer_debug_count}": G_hat.cpu(),
                    },
                    self._debug_file,
                )
                self._layer_debug_count += 1

            if self._adaptive_damping:
                self._pred_loss_accum += (
                    (G_hat @ g).unsqueeze(1) @ A.unsqueeze(2)
                ).squeeze(2)
                self._predicted_loss += (self.damping / 2) * (g**2).sum()

            if bias is not None:
                gb = g[:, -1].contiguous().view(*bias.shape)
                g = g[:, :-1]
            else:
                gb = None
            g = g.contiguous().view(*weight.shape)

            weight.grad.add_(g, alpha=self._batch_factor)
            if self.adjust_by_damping:
                weight.grad /= self.damping
            if bias is not None:
                bias.grad.add_(gb, alpha=self._batch_factor)
                if self.adjust_by_damping:
                    bias.grad /= self.damping

    def _precondition_convolution(self, group):
        # For these layers we know grad = g a^T, which transforms the update
        # into the Hadamard product below.
        weight = group["params"][0]
        bias = group["params"][1] if len(group["params"]) == 2 else None
        state = self.state[group["mod"]]
        group_key = f"group_{group['group_ix']}"
        delta_mc = self._delta_mc[group_key]

        if delta_mc is not None:
            if state["reducible"]:
                G_hat = state["G_hat"]
                A = state["A"]

                g = torch.einsum("bij,bkj,bi->ik", G_hat, A, delta_mc)
            else:
                g = torch.einsum("bij,bi->ij", state["DW_hat"], delta_mc)

            weight.grad.add_(g.reshape_as(weight), alpha=self._batch_factor)
            if self.adjust_by_damping:
                weight.grad /= self.damping
            # if bias is not None:
            #     b = (delta_mc.T @ state["G_hat"].sum(dim=(2, 3))).view(*bias.shape)
            #     bias.grad.add_(b, alpha=self._batch_factor)
            #     bias.grad /= self.damping

    def _precondition_raw(self, group):
        state = self.state[group["mod"]]
        group_key = f"group_{group['group_ix']}"
        delta_mc = self._delta_mc[group_key]

        if delta_mc is not None:
            if "G_hat" in state:
                G_hat = state["G_hat"].flatten(1)
                for weight in group["params"]:
                    g = delta_mc.T @ G_hat
                    if self._adaptive_damping:
                        self._pred_loss_accum += G_hat @ g.T
                        self._predicted_loss += (self.damping / 2) * (g**2).sum()
                    weight.grad.add_(g.reshape_as(weight), alpha=self._batch_factor)
                    if self.adjust_by_damping:
                        weight.grad /= self.damping

    def update_damping(self, loss_red: float = None, lr: float = None):
        """Updates the damping parameter based on the reduction ratio.
        Not used by default, or in most experiments since it's been found to be ineffective in large nets.
        """
        if self._adaptive_damping:
            self._predicted_loss *= lr**2
            self._predicted_loss += (
                lr * self._pred_loss_accum.mean()
                + (lr**2 / 2) * (self._pred_loss_accum**2).mean()
            )
            self.reduction_ratio = -loss_red / self._predicted_loss

            if self._iteration_counter % self._damp_freq == 0:
                if self.reduction_ratio > 0.75:
                    self.damping *= self._damp_upf
                elif self.reduction_ratio < 0.25:
                    self.damping /= self._damp_upf

    @property
    def should_accumulate_grads(self):
        if self.accumulate_grad_batches > 1:
            return self._batch_count != self.accumulate_grad_batches
        return False

    @property
    def should_update_grads(self):
        if self.update_freq > 1:
            return (self._iteration_counter % self.update_freq) == 0
        return True
