"""QuantizedTensor layout classes for INT-Crush INT4/INT8 formats.

These bridge INT-Crush's packed weight format with ComfyUI's QuantizedTensor
dispatch system, enabling VBAR demand-paging, async offload, and LoRA
compatibility through ComfyUI's standard cast_bias_weight pipeline.

Each layout class defines:
  - Params dataclass: metadata stored alongside packed weights
  - state_dict_tensors(): serialization layout for safetensors
  - dequantize(): full dequantization (used for LoRA patching and fallback)
  - quantize(): float → packed (used by QuantizedTensor.from_float() for LoRA requantization)
  - supports_fast_matmul(): check for fused Triton kernel availability
"""

import torch
import torch.nn.functional as F
from dataclasses import dataclass

from ._quant_utils import unpack_int4, INT4_SCALE_DIVISOR, INT4_MIN, INT4_MAX, INT4_TWOS_COMPLEMENT_BIAS
from ._quant_utils import INT8_MIN, INT8_MAX, INT8_SCALE_DIVISOR

__all__ = ["IntCrushInt4Layout", "IntCrushInt8Layout", "OrbitQuantLayout", "register_intcrush_layouts"]

# BaseLayoutParams is required by QuantizedTensor.__tensor_flatten__ which
# calls dataclasses.fields(self._params) to discover tensor sub-fields.
try:
    from comfy_kitchen.tensor.base import BaseLayoutParams
except ImportError:
    # Minimal fallback when comfy_kitchen is unavailable (e.g. unit tests).
    # The real BaseLayoutParams is a @dataclass(frozen=True) with three fields:
    #   scale: torch.Tensor
    #   orig_dtype: torch.dtype
    #   orig_shape: tuple[int, ...]
    @dataclass(frozen=True)
    class BaseLayoutParams:
        scale: torch.Tensor
        orig_dtype: torch.dtype
        orig_shape: tuple

        def __post_init__(self):
            pass


class IntCrushInt4Layout:
    """Layout class for INT-Crush packed INT4 (uint8, 2 values per byte).

    Stores weights as packed uint8 with per-row FP16 scales. Optionally
    supports zero-point (asymmetric quantization) and PermuQuant channel
    permutation indices.
    """

    @dataclass(frozen=True)
    class Params(BaseLayoutParams):
        rot_need: bool = True
        rot_size: int = 256
        perm: torch.Tensor | None = None
        zp: torch.Tensor | None = None

        def _tensor_fields(self) -> list[str]:
            fields = ["scale"]
            if self.perm is not None:
                fields.append("perm")
            if self.zp is not None:
                fields.append("zp")
            return fields

        def _validate_tensor_fields(self):
            if isinstance(self.scale, torch.Tensor):
                object.__setattr__(self, "scale", self.scale.to(dtype=torch.float32, non_blocking=True))

    @classmethod
    def state_dict_tensors(cls, qdata, params):
        tensors = {"": qdata, "weight_scale": params.scale}
        if params.perm is not None:
            tensors["weight_perm"] = params.perm
        if params.zp is not None:
            tensors["weight_zp"] = params.zp
        return tensors

    @classmethod
    def dequantize(cls, qdata, params):
        """Full dequantization to float — used for LoRA patching and PyTorch fallback."""
        in_features = params.orig_shape[1]
        unpacked = unpack_int4(qdata, in_features)
        if params.zp is not None:
            return (unpacked.float() - params.zp.reshape(-1, 1).float()) * params.scale.reshape(-1, 1).float()
        return unpacked.float() * params.scale.reshape(-1, 1).float()

    @classmethod
    def supports_fast_matmul(cls):
        # comfy-kitchen CUDA (cuBLASLt / CUTLASS / GEMV — sm75+)
        try:
            from comfy_kitchen.backends.cuda import _EXT_AVAILABLE
            if _EXT_AVAILABLE:
                return True
        except ImportError:
            pass
        # Triton fallback
        try:
            from .kernels.triton_int8_gemm import fused_int8_gemm_dequant  # noqa: F401  -- probe-by-import
            from .kernels.triton_quantize import dynamic_quantize_activation  # noqa: F401
            from .kernels.triton_int4_to_int8_unpack import unpack_int4_to_int8  # noqa: F401
            return True
        except Exception:
            return False

    @classmethod
    def quantize(cls, tensor, scale=None, stochastic_rounding=0, inplace_ops=False):
        """Quantize a float tensor to packed INT4 with per-row scales.

        Used by QuantizedTensor.from_float() during LoRA requantization.
        The input tensor must already be in rotated+permuted space.
        """
        max_vals = tensor.abs().amax(dim=1, keepdim=True)
        if scale is None or (isinstance(scale, str) and scale == "recalculate"):
            scales = (max_vals / INT4_SCALE_DIVISOR).clamp(min=1e-8).to(torch.float16)
        elif isinstance(scale, torch.Tensor):
            scales = scale.to(torch.float16)
        else:
            scales = torch.tensor([scale], dtype=torch.float16, device=tensor.device)

        w_scaled = tensor / scales.to(tensor.dtype)
        int_rounded = w_scaled.round().clamp(INT4_MIN, INT4_MAX).to(torch.int8)
        # Pack two INT4 values per byte (low nibble = even index).
        K = int_rounded.shape[-1]
        if K % 2 != 0:
            pad = torch.zeros(*int_rounded.shape[:-1], 1, dtype=torch.int8, device=int_rounded.device)
            int_rounded = torch.cat([int_rounded, pad], dim=-1)
        q_i8 = torch.where(int_rounded < 0, INT4_TWOS_COMPLEMENT_BIAS + int_rounded, int_rounded).to(torch.uint8)
        packed = q_i8[..., 0::2] | (q_i8[..., 1::2] << 4)

        orig_dtype = tensor.dtype
        orig_shape = tuple(tensor.shape)
        params = cls.Params(
            scale=scales.reshape(-1), orig_dtype=orig_dtype, orig_shape=orig_shape,
        )
        return packed, params

    @classmethod
    def requantize_kwargs(cls, qtensor):
        # No options are preserved. VBAR's requantize_from_float uses
        # scale="recalculate", and rot/perm live on the module, not the params.
        return {}


class IntCrushInt8Layout:
    """Layout class for INT-Crush INT8 (plain int8) with per-row FP16 scales."""

    @dataclass(frozen=True)
    class Params(BaseLayoutParams):
        rot_need: bool = True
        rot_size: int = 256
        perm: torch.Tensor | None = None
        zp: torch.Tensor | None = None

        def _tensor_fields(self) -> list[str]:
            fields = ["scale"]
            if self.perm is not None:
                fields.append("perm")
            if self.zp is not None:
                fields.append("zp")
            return fields

        def _validate_tensor_fields(self):
            if isinstance(self.scale, torch.Tensor):
                object.__setattr__(self, "scale", self.scale.to(dtype=torch.float32, non_blocking=True))

    @classmethod
    def state_dict_tensors(cls, qdata, params):
        tensors = {"": qdata, "weight_scale": params.scale}
        if params.perm is not None:
            tensors["weight_perm"] = params.perm
        if params.zp is not None:
            tensors["weight_zp"] = params.zp
        return tensors

    @classmethod
    def dequantize(cls, qdata, params):
        """Full dequantization to float — used for LoRA patching and PyTorch fallback."""
        if params.zp is not None:
            return (qdata.float() - params.zp.reshape(-1, 1).float()) * params.scale.reshape(-1, 1).float()
        return qdata.float() * params.scale.reshape(-1, 1).float()

    @classmethod
    def supports_fast_matmul(cls):
        # comfy-kitchen CUDA (cuBLASLt / CUTLASS / GEMV — sm75+)
        try:
            from comfy_kitchen.backends.cuda import _EXT_AVAILABLE
            if _EXT_AVAILABLE:
                return True
        except ImportError:
            pass
        # Triton fallback
        try:
            from .kernels.triton_int8_gemm import fused_int8_gemm_dequant  # noqa: F401  -- probe-by-import
            from .kernels.triton_quantize import dynamic_quantize_activation  # noqa: F401
            return True
        except Exception:
            return False

    @classmethod
    def quantize(cls, tensor, scale=None, stochastic_rounding=0, inplace_ops=False):
        """Quantize a float tensor to INT8 with per-row scales.

        Used by QuantizedTensor.from_float() during LoRA requantization.
        The input tensor must already be in rotated space.
        """
        max_vals = tensor.abs().amax(dim=1, keepdim=True)
        if scale is None or (isinstance(scale, str) and scale == "recalculate"):
            scales = (max_vals / INT8_SCALE_DIVISOR).clamp(min=1e-8).to(torch.float32)
        elif isinstance(scale, torch.Tensor):
            scales = scale.to(torch.float32)
        else:
            scales = torch.tensor([scale], dtype=torch.float32, device=tensor.device)

        w_scaled = tensor / scales.to(tensor.dtype)
        int_rounded = w_scaled.round().clamp(INT8_MIN, INT8_MAX).to(torch.int8)

        orig_dtype = tensor.dtype
        orig_shape = tuple(tensor.shape)
        params = cls.Params(
            scale=scales.reshape(-1), orig_dtype=orig_dtype, orig_shape=orig_shape,
        )
        return int_rounded, params

    @classmethod
    def requantize_kwargs(cls, qtensor):
        # No options are preserved. VBAR's requantize_from_float uses
        # scale="recalculate", and rot/perm live on the module, not the params.
        return {}


class OrbitQuantLayout:
    """Layout class for OrbitQuant codebook-based quantization.

    Stores packed codebook indices (uint8 for INT4, int8 for INT8) with
    per-row L2 norms (bfloat16).  AdaLN layers use per-group RTN scales
    (float16) instead.

    The codebook is stored in Params.codebook so that dequantize() can
    access it for LoRA patching, VBAR demand-paging, and PyTorch fallback.
    """

    @dataclass(frozen=True)
    class Params(BaseLayoutParams):
        quant_scheme: str = "orbitquant_codebook"  # or "adaln_rtn"
        d: int = 0           # Original input dimension
        d_padded: int = 0    # Padded input dimension
        bits: int = 4        # Weight bit-width
        codebook: torch.Tensor | None = None  # Lloyd-Max centroids [2^b]

        def _tensor_fields(self) -> list[str]:
            fields = ["scale"]
            if self.codebook is not None:
                fields.append("codebook")
            return fields

        def _validate_tensor_fields(self):
            if isinstance(self.scale, torch.Tensor):
                object.__setattr__(self, "scale",
                    self.scale.to(dtype=torch.float32, non_blocking=True))
            # codebook stays float32 (already correct dtype from converter)

    @classmethod
    def state_dict_tensors(cls, qdata, params):
        tensors = {"": qdata, "weight_scale": params.scale}
        # codebook is stored globally, not per-layer — not in state_dict_tensors
        return tensors

    @classmethod
    def dequantize(cls, qdata, params):
        """Full dequantization to float — used for LoRA patching, VBAR,
        and PyTorch fallback.

        Produces weights in the ROTATED basis (W @ Π^T).  The forward path
        must also rotate activations for cancellation.  For LoRA patching,
        this is fine — LoRA is applied on raw activations, and the output
        is in the original (unrotated) space because the rotation cancels.

        Uses the Triton codebook dequant (fp32 output) when available for
        speed; falls back to PyTorch otherwise.
        """
        from ._orbitquant import codebook_dequant, adaln_rtn_dequant

        if params.quant_scheme == "adaln_rtn":
            return adaln_rtn_dequant(qdata, params.scale, params.d, params.bits)

        if params.codebook is None:
            raise RuntimeError(
                "OrbitQuant codebook not available in Params. "
                "This should not happen — the loader stores it."
            )
        # Try Triton fp32 output for exact LoRA parity.
        try:
            from . import _triton_runtime as _tr
            if _tr.TRITON_CODEBOOK_DEQUANT and qdata.is_cuda:
                out = torch.empty(
                    qdata.shape[0], params.d_padded,
                    dtype=torch.float32, device=qdata.device,
                )
                return _tr.codebook_dequant_triton(
                    qdata, params.scale, params.codebook,
                    params.d_padded, params.bits, out=out,
                )
        except Exception:
            pass
        return codebook_dequant(
            qdata, params.scale, params.codebook,
            params.d, params.d_padded, params.bits,
        )

    @classmethod
    def supports_fast_matmul(cls):
        # RPBH rotation is the recurring per-step cost; fast matmul is
        # meaningful once the RPBH kernel is loaded.
        try:
            from . import _triton_runtime as _tr
            return bool(_tr.TRITON_RPBH_ROTATE)
        except Exception:
            return False

    @classmethod
    def quantize(cls, tensor, scale=None, stochastic_rounding=0, inplace_ops=False):
        """Quantize a float tensor to OrbitQuant format.

        Used by QuantizedTensor.from_float() during LoRA requantization.
        NOT YET IMPLEMENTED — requires RPBH rotation + codebook quantization.
        For now, the caller (set_weight) handles this by storing as plain float.
        """
        raise NotImplementedError(
            "OrbitQuant requantization not yet implemented. "
            "The weight will be stored as plain float and the forward path "
            "will use _forward_predequant_oq."
        )


def register_intcrush_layouts():
    """Register INT-Crush INT4/INT8 formats with ComfyUI's quantization system.

    Idempotent and safe to call even when comfy_kitchen is unavailable (no-op).
    """
    import logging
    _log = logging.getLogger(__name__)

    try:
        from comfy.quant_ops import register_layout_class, QUANT_ALGOS
    except ImportError:
        _log.warning("INT-Crush: comfy.quant_ops not available; layout registration skipped")
        return

    try:
        register_layout_class("IntCrushInt4Layout", IntCrushInt4Layout)
        register_layout_class("IntCrushInt8Layout", IntCrushInt8Layout)
        register_layout_class("OrbitQuantLayout", OrbitQuantLayout)
    except Exception:
        _log.warning("INT-Crush: layout class registration failed", exc_info=True)

    QUANT_ALGOS.setdefault("int4_crush", {
        "storage_t": torch.uint8,
        "parameters": {"weight_scale", "weight_zp", "weight_perm"},
        "comfy_tensor_layout": "IntCrushInt4Layout",
    })
    QUANT_ALGOS.setdefault("int8_crush", {
        "storage_t": torch.int8,
        "parameters": {"weight_scale", "weight_perm", "weight_zp"},
        "comfy_tensor_layout": "IntCrushInt8Layout",
    })
    QUANT_ALGOS.setdefault("orbitquant", {
        "storage_t": torch.uint8,  # INT4 default; overridden for INT8
        "parameters": {"weight_scale"},
        "comfy_tensor_layout": "OrbitQuantLayout",
    })