import torch
import numpy as np


def get_available_memory(device):
    """Get available GPU memory in bytes."""
    if device.type == 'cuda':
        torch.cuda.synchronize()
        free, total = torch.cuda.mem_get_info()
        return free
    else:
        # return a default number: 16 GB
        return 16 * 1024 * 1024 * 1024


def estimate_hqq_memory(num_elements, compute_dtype=torch.float32, symmetric=False):
    """
    Estimate the memory needed for HQQ quantization optimization.

    The optimization process creates multiple temporary tensors:
    - W_f: copy of weights in compute_dtype
    - W_q: quantized weights
    - W_r: reconstructed weights
    - W_e: error tensor
    - W_prime: additional tensor for symmetric quantization
    - Additional temporary tensors in _shrink_lp_op

    Returns estimated memory in bytes.
    """
    dtype_size = 4 if compute_dtype == torch.float32 else 2

    # Main working tensors
    num_tensors = 4 if not symmetric else 5
    main_memory = num_elements * dtype_size * num_tensors

    # Temporary tensors during _shrink_lp_op (can create 2-3 additional tensors)
    temp_memory = num_elements * dtype_size * 3

    # Scale and zero tensors (smaller, proportional to num_groups)
    # Assuming group_size = 64, num_groups = num_elements / 64
    scale_zero_memory = (num_elements // 64) * dtype_size * 2

    # Add 20% safety margin for PyTorch memory allocator overhead
    total_memory = (main_memory + temp_memory + scale_zero_memory) * 1.2

    return int(total_memory)


class HQQQuantizer:
    def __init__(self,
                 weight,
                 bit,
                 group_size,
                 sym=False,
                 compute_dtype: torch.dtype = torch.float32,
                 device: torch.device = torch.device("cpu"),
                 quant_config: dict = None,
                 auto_chunk: bool = True,
                 memory_safety_factor: float = 0.7):
        self.weight = weight
        self.bit = bit
        self.group_size = group_size
        self.sym = sym
        self.compute_dtype = compute_dtype
        self.device = device
        self.auto_chunk = auto_chunk
        self.memory_safety_factor = memory_safety_factor

    def quant(self):
        if self.auto_chunk and self.device.type in ('cuda', 'mps'):
            self._quantize_chunked()
        else:
            self._quantize()

    def _quantize_chunked(self):
        """Quantize with automatic chunking to avoid OOM."""
        num_elements = self.weight.numel()
        estimated_memory = estimate_hqq_memory(num_elements, self.compute_dtype, self.sym)
        available_memory = get_available_memory(self.device)
        safe_memory = available_memory * self.memory_safety_factor

        if estimated_memory <= safe_memory:
            # No chunking needed
            self._quantize()
            return

        # Calculate number of chunks needed
        num_chunks = int(np.ceil(estimated_memory / safe_memory))
        original_shape = self.weight.shape

        # Split along the first dimension (output channels)
        chunk_size = max(1, original_shape[0] // num_chunks)
        num_chunks = (original_shape[0] + chunk_size - 1) // chunk_size

        # Collect results
        W_q_list = []
        scale_list = []
        zero_list = []

        for i in range(num_chunks):
            start_idx = i * chunk_size
            end_idx = min((i + 1) * chunk_size, original_shape[0])
            chunk_weight = self.weight[start_idx:end_idx].contiguous()

            # Quantize this chunk
            chunk_quantizer = HQQQuantizer(
                chunk_weight,
                self.bit,
                self.group_size,
                self.sym,
                self.compute_dtype,
                self.device,
                auto_chunk=False  # Disable recursive chunking
            )
            chunk_quantizer._quantize()

            W_q_list.append(chunk_quantizer.W_q)
            scale_list.append(chunk_quantizer.meta['scale'])
            if not self.sym:
                zero_list.append(chunk_quantizer.meta['zero'])

            # Clean up chunk memory
            del chunk_weight, chunk_quantizer
            if self.device.type == 'cuda':
                torch.cuda.empty_cache()
            elif self.device.type == 'mps':
                torch.mps.empty_cache()

        # Concatenate results
        self.W_q = torch.cat(W_q_list, dim=0)
        self.meta = {
            "nbits": self.bit,
            "group_size": self.group_size,
            "shape": original_shape,
            "scale": torch.cat(scale_list, dim=0),
            "zero": torch.cat(zero_list, dim=0) if not self.sym else None,
            "axis": 1,
        }

        # Clean up
        del W_q_list, scale_list, zero_list
        if self.device.type == 'cuda':
            torch.cuda.empty_cache()
        elif self.device.type == 'mps':
            torch.mps.empty_cache()

    @torch.inference_mode()
    def _quantize(
        self,
        channel_wise: bool = True,
        axis: int = 1,
    ) -> tuple:

        if self.group_size is not None:
            assert self.weight.numel() % self.group_size == 0, (
                "group_size should be divisble by the total tensor dimensions. shape: "
                + str(self.weight.shape)
                + ", group_size: "
                + str(self.group_size)
            )

        W = self.weight.to(self.compute_dtype).float()
        shape = W.shape

        # Reshape for grouping
        if (self.group_size is not None) and channel_wise:
            W = (
                W.reshape([-1, self.group_size])
                if (axis == 1)
                else W.reshape([self.group_size, -1])
            )

        # Get min/max values
        if not channel_wise:
            _min, _max = W.min(), W.max()
        else:
            _min = W.min(axis=axis, keepdim=True)[0]
            _max = W.max(axis=axis, keepdim=True)[0]

        if self.sym:
            max_v = 2**(self.bit-1) - 1    # 4bit: 7
            min_v = -2**(self.bit-1)       # 4bit: -8
            min_max = [min_v, max_v]       # [-8, 7]

            max_abs = torch.max(torch.abs(_min), torch.abs(_max))
            scale = max_v / max_abs
            scale = torch.where(max_abs <= 1e-4, torch.full_like(scale, 1.0), scale)
            scale = scale.clamp(max=2e4)
            zero = None
        else:
            max_v = round(2**self.bit - 1)  # 4bit: 15
            min_v = 0                       # 4bit: 0
            min_max = [min_v, max_v]        # [0, 15]

            denom = (_max - _min)
            scale = (max_v / denom)
            scale = torch.where(denom.abs() <= 1e-4, torch.full_like(scale, 1.0), scale)
            scale = scale.clamp(max=2e4)
            zero = -_min * scale
            zero = torch.round(zero)

        W_q, scale, zero = self._optimize_weights(
            W,
            scale,
            zero,
            min_max=min_max,
            axis=axis,
        )
        #W_q = (W * scale).round_().clamp_(min_max[0], min_max[1])
        # cleanup
        del W, _min, _max

        # Store meta-data (we invert the scale for dequantization)
        scale = 1.0 / scale
        meta = {
            "nbits": self.bit,
            "group_size": self.group_size,
            "shape": shape,
            "scale": scale,
            "zero": zero,
            "axis": axis,
        }

        W_q = W_q.to(self.weight.dtype)

        if self.device == torch.device('cuda'):
            torch.cuda.empty_cache()
        elif self.device == torch.device('mps'):
            torch.mps.empty_cache()

        self.W_q = W_q
        self.meta = meta

    @torch.inference_mode()
    def _optimize_weights(
        self,
        W: torch.Tensor,
        scale: torch.Tensor,
        zero: torch.Tensor,
        min_max: list,
        axis: int = 0,
        opt_params: dict = {"lp_norm": 0.7, "beta": 1e1, "kappa": 1.01, "iters": 20},
        verbose: bool = False,
    ) -> tuple:
        lp_norm, beta, kappa, iters = (
            opt_params["lp_norm"],
            opt_params["beta"],
            opt_params["kappa"],
            opt_params["iters"],
        )

        dtype = torch.float32
        W_f   = W.to(dtype=dtype, device=self.device)
        scale = scale.to(dtype=dtype, device=self.device)
        if not self.sym:
            zero  = zero.to(dtype=dtype, device=self.device)

        best_error = torch.tensor(torch.inf, dtype=torch.float32, device=self.device)
        best_scale = scale.clone()
        best_zero = zero.clone() if not self.sym else None
        W_q = torch.empty_like(W_f)
        W_r = torch.empty_like(W_f)
        W_e = torch.empty_like(W_f)
        W_prime = torch.empty_like(W_f) if self.sym else None
        for i in range(iters):
            if not self.sym:
                self._optimize_weights_proximal_legacy_step(W_f, scale, zero, min_max, beta, lp_norm, axis, W_q, W_r, W_e)
            else:
                self._optimize_weights_proximal_scale_only(W_f, scale, min_max, beta, lp_norm, axis, W_q, W_r, W_e, W_prime)
            current_error = torch.abs(W_f - W_r).mean().float()
            if verbose:
                print(i, current_error.cpu())

            # Check for NaN in error or scale - restore best and stop
            if torch.isnan(current_error) or torch.any(torch.isnan(scale)) or (not self.sym and torch.any(torch.isnan(zero))):
                scale.copy_(best_scale)
                if not self.sym:
                    zero.copy_(best_zero)
                break

            if current_error < best_error:
                best_error = current_error
                best_scale.copy_(scale)
                if not self.sym:
                    best_zero.copy_(zero)
            else:
                scale.copy_(best_scale)
                if not self.sym:
                    zero.copy_(best_zero)
                break

        scale = scale.to(W.device)
        if not self.sym:
            zero = zero.to(W.device)
        del W_f, W_q, W_r, best_scale, best_zero
        if self.device.type == 'cuda':
            torch.cuda.empty_cache()
        elif self.device.type == 'mps':
            torch.mps.empty_cache()
        if not self.sym:
            W_q = torch.round(W * scale + zero).clamp_(min_max[0], min_max[1])
        else:
            W_q = torch.round(W * scale).clamp_(min_max[0], min_max[1])
        return W_q, scale, zero

    @torch.inference_mode()
    def _optimize_weights_proximal_legacy_step(self, W_f, scale, zero, min_max, beta, lp_norm, axis, W_q, W_r, W_e):
        _eps = 1e-8
        torch.mul(W_f, scale, out=W_q)
        torch.add(W_q, zero, out=W_q)
        torch.round(W_q, out=W_q).clamp_(min_max[0], min_max[1])
        torch.sub(W_q, zero, out=W_r)
        # Guard against division by zero
        safe_scale = torch.where(scale.abs() < _eps, torch.full_like(scale, _eps), scale)
        torch.div(W_r, safe_scale, out=W_r)
        torch.sub(W_f, W_r, out=W_e)
        self._shrink_lp_op(W_e, beta, lp_norm, out=W_e)
        torch.sub(W_f, W_e, out=W_r)
        torch.mul(W_r, scale, out=W_r)
        torch.sub(W_q, W_r, out=W_r)
        torch.mean(W_r, axis=axis, keepdim=True, out=zero)

    @torch.inference_mode()
    def _optimize_weights_proximal_scale_only(self, W_f, scale, min_max, beta, lp_norm, axis, W_q, W_r, W_e, W_prime):
        _eps = 1e-8
        torch.mul(W_f, scale, out=W_q)
        torch.round(W_q, out=W_q).clamp_(min_max[0], min_max[1])
        # Guard against division by zero: replace zero scale with eps
        safe_scale = torch.where(scale.abs() < _eps, torch.full_like(scale, _eps), scale)
        torch.div(W_q, safe_scale, out=W_r)
        torch.sub(W_f, W_r, out=W_e)
        self._shrink_lp_op(W_e, beta, lp_norm, out=W_e)
        torch.sub(W_f, W_e, out=W_prime)
        w_prime_dot_w_q = torch.sum(W_prime * W_q, axis=axis, keepdim=True)
        w_q_norm_sq = torch.sum(W_q**2, axis=axis, keepdim=True)
        torch.add(w_prime_dot_w_q, _eps, out=w_prime_dot_w_q)
        torch.div(w_q_norm_sq, w_prime_dot_w_q, out=scale)
        # Clamp scale to prevent zero or negative values that cause NaN in next iteration
        scale.clamp_(min=_eps)

    # Shrinking operator
    @torch.inference_mode()
    def _shrink_lp_op(self, x: torch.Tensor, beta: float, lp_norm: float, out: torch.Tensor) -> torch.Tensor:
        if lp_norm == 1:
            #torch.sign(x) * torch.nn.functional.relu(torch.abs(x) - 1.0 / beta)
            torch.abs(x, out=out)
            out.sub_(1.0 / beta).clamp_min_(0.0)
            out.mul_(torch.sign(x))
            return out
        else:
            # Original formula: sign(x) * relu(|x| - (1/beta) * |x|^(lp_norm-1))
            # Note: This formula inherently requires a temporary tensor for lp_norm != 1
            # because we need both |x| and |x|^(lp_norm-1) simultaneously
            torch.abs(x, out=out)
            temp = out.pow(lp_norm - 1)
            temp.mul_(1.0 / beta)
            out.sub_(temp).clamp_min_(0.0)
            out.mul_(torch.sign(x))
            del temp
            return out