"""Original SparseGPT one-shot unstructured pruning for nn.Linear."""

from __future__ import annotations

import math
import time

import torch
import torch.nn as nn


class SparseGPT:
    """Hessian-aware unstructured pruning (Frantar & Alistarh, arXiv:2301.00774)."""

    def __init__(self, layer: nn.Linear):
        self.layer = layer
        self.dev = layer.weight.device
        w = layer.weight.data
        self.rows = w.shape[0]
        self.columns = w.shape[1]
        self.H = torch.zeros((self.columns, self.columns), device=self.dev)
        self.nsamples = 0

    def add_batch(self, inp: torch.Tensor, out: torch.Tensor):
        if len(inp.shape) == 2:
            inp = inp.unsqueeze(0)
        tmp = inp.shape[0]
        if len(inp.shape) == 3:
            inp = inp.reshape((-1, inp.shape[-1]))
        inp = inp.t()
        if self.nsamples > 0:
            self.H *= self.nsamples / (self.nsamples + tmp)
        self.nsamples += tmp
        inp = math.sqrt(2 / self.nsamples) * inp.float()
        self.H += inp.matmul(inp.t())

    def fasterprune(
        self,
        sparsity: float,
        prunen: int = 0,
        prunem: int = 0,
        blocksize: int = 128,
        percdamp: float = 0.01,
        verbose: bool = True,
    ):
        w = self.layer.weight.data.clone().float()
        h = self.H
        self.H = None
        dead = torch.diag(h) == 0
        h[dead, dead] = 1
        w[:, dead] = 0

        losses = torch.zeros(self.rows, device=self.dev)
        mean_diag = torch.mean(torch.diag(h)).clamp_min(1e-8)
        damp = float(percdamp) * float(mean_diag)
        diag = torch.arange(self.columns, device=self.dev)
        hinv = None
        for _ in range(16):
            trial = h.clone()
            trial[diag, diag] += damp
            chol, info = torch.linalg.cholesky_ex(trial)
            info_i = int(info.item()) if info.ndim == 0 else int(info.max().item())
            if info_i == 0:
                inv = torch.cholesky_inverse(chol)
                inv = 0.5 * (inv + inv.transpose(-1, -2))
                inv[diag, diag] += damp
                inv_chol, info2 = torch.linalg.cholesky_ex(inv, upper=True)
                info2_i = int(info2.item()) if info2.ndim == 0 else int(info2.max().item())
                if info2_i == 0:
                    hinv = inv_chol
                    break
            damp *= 10.0
        if hinv is None:
            raise torch.linalg.LinAlgError(
                f"SparseGPT Hessian not SPD after damping; columns={self.columns}"
            )

        mask = None
        tick = time.time()

        for i1 in range(0, self.columns, blocksize):
            i2 = min(i1 + blocksize, self.columns)
            count = i2 - i1
            w1 = w[:, i1:i2].clone()
            q1 = torch.zeros_like(w1)
            err1 = torch.zeros_like(w1)
            losses1 = torch.zeros_like(w1)
            hinv1 = hinv[i1:i2, i1:i2]

            if prunen == 0:
                if mask is not None:
                    mask1 = mask[:, i1:i2]
                else:
                    tmp = w1 ** 2 / (torch.diag(hinv1).reshape((1, -1))) ** 2
                    thresh = torch.sort(tmp.flatten())[0][int(tmp.numel() * sparsity)]
                    mask1 = tmp <= thresh
            else:
                mask1 = torch.zeros_like(w1) == 1

            for i in range(count):
                wi = w1[:, i]
                d = hinv1[i, i]
                if prunen != 0 and i % prunem == 0:
                    tmp = w1[:, i:(i + prunem)] ** 2 / (
                        torch.diag(hinv1)[i:(i + prunem)].reshape((1, -1))
                    ) ** 2
                    mask1.scatter_(1, i + torch.topk(tmp, prunen, dim=1, largest=False)[1], True)
                q = wi.clone()
                q[mask1[:, i]] = 0
                q1[:, i] = q
                losses1[:, i] = (wi - q) ** 2 / d ** 2
                err = (wi - q) / d
                w1[:, i:] -= err.unsqueeze(1).matmul(hinv1[i, i:].unsqueeze(0))
                err1[:, i] = err

            w[:, i1:i2] = q1
            losses += torch.sum(losses1, 1) / 2
            w[:, i2:] -= err1.matmul(hinv[i1:i2, i2:])

        sparsity_val = 1.0 - float((w != 0).sum()) / float(w.numel())
        if verbose:
            if self.dev.type == "cuda":
                torch.cuda.synchronize()
            print(
                f"sparsegpt {self.layer}: time {time.time() - tick:.2f}s "
                f"error {torch.sum(losses).item():.4f} "
                f"sparsity {sparsity_val:.3f}"
            )
        self.layer.weight.data = w.reshape(self.layer.weight.shape).to(self.layer.weight.dtype)
        return sparsity_val

    def free(self):
        self.H = None
        torch.cuda.empty_cache()