"""
Experimental A2 NVFP4 performance recipe.

IMPORTANT:
- BF16 -> NVFP4 encoder is experimental.
- Persistent numeric representation obeys NInfer NVFP4 semantics.
- Encoder profile:
    CT_MAXABS_E2M1_E4M3_V1
- activation_input_divisor=1.0 is PERFORMANCE-ONLY.
- Do NOT use this artifact for quality conclusions before d_x calibration.
"""

from __future__ import annotations

import math
import struct

import numpy as np
import torch

from tools.convert.methods import grouped_absmax
from tools.convert.quantization.fp8_row import (
    _POSITIVE_E4M3FN_VALUES,
    _round_e4m3fn_rne,
)

Q6 = "q6_g64_fp16"
NVFP4 = "nvfp4"

# Positive E2M1 code words 0..7.
_E2M1 = np.array(
    [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0],
    dtype=np.float32,
)

# FP4 max * FP8-E4M3FN max.
_NVFP4_PRODUCT_MAX = np.float32(6.0 * 448.0)


def _round_e2m1_rne(values: np.ndarray) -> np.ndarray:
    """Round finite float32 values to E2M1 words, ties-to-even."""

    if values.dtype != np.float32:
        raise TypeError("E2M1 input must be float32")

    if not np.isfinite(values).all():
        raise ValueError("E2M1 input contains NaN/Inf")

    bounded = np.clip(values, -6.0, 6.0).astype(np.float32, copy=False)
    magnitude = np.abs(bounded)

    upper = np.searchsorted(_E2M1, magnitude, side="left").astype(np.int16)
    upper = np.minimum(upper, 7)
    lower = np.maximum(upper - 1, 0)

    lower_distance = magnitude - _E2M1[lower]
    upper_distance = _E2M1[upper] - magnitude

    choose_upper = (
        (upper_distance < lower_distance)
        | (
            (upper_distance == lower_distance)
            & ((upper & 1) == 0)
        )
    )

    words = np.where(choose_upper, upper, lower).astype(np.uint8)

    # Preserve sign, including negative zero.
    words |= np.where(np.signbit(bounded), 0x8, 0).astype(np.uint8)

    return words


def _read_rows(request, begin: int, end: int, k: int) -> np.ndarray:
    values = (
        request.values(begin * k, end * k)
        .reshape(end - begin, k)
        .to(torch.bfloat16)
        .to(torch.float32)
        .contiguous()
    )

    host = values.numpy()

    if not np.isfinite(host).all():
        raise ValueError(
            f"{request.target.id}: NVFP4 source contains NaN/Inf"
        )

    return host


def nvfp4_ct_maxabs(request):
    """
    BF16 -> compressed-tensors-style NVFP4 MAXABS encoder.

    Representation:
        W = E2M1(code) * E4M3(scale) / d_w

    d_w:
        (6 * 448) / tensor_amax

    block scale:
        RNE_E4M3(d_w * block16_amax / 6)

    code:
        RNE_E2M1(W * d_w / decoded_block_scale)
    """

    if request.target.format != NVFP4:
        raise ValueError("nvfp4_ct_maxabs requires target format nvfp4")

    if len(request.target.shape) != 2:
        raise ValueError("NVFP4 target must be rank two")

    if request.parameters:
        raise ValueError("nvfp4_ct_maxabs accepts no parameters")

    n, k = request.target.shape

    if n % 128:
        raise ValueError(
            f"{request.target.id}: NVFP4 N must be divisible by 128, got {n}"
        )

    if k % 16:
        raise ValueError(
            f"{request.target.id}: NVFP4 K must be divisible by 16, got {k}"
        )

    chunk = max(128, (request.rows_per_chunk // 128) * 128)

    # ---------------------------------------------------------------
    # Pass 1: one global divisor for the complete physical parent.
    # ---------------------------------------------------------------

    tensor_amax = np.float32(0.0)

    for begin in range(0, n, chunk):
        end = min(n, begin + chunk)
        host = _read_rows(request, begin, end, k)

        local = np.max(np.abs(host)).astype(np.float32)

        if local > tensor_amax:
            tensor_amax = local

    if tensor_amax == 0:
        divisor = np.float32(1.0)
    else:
        divisor = np.float32(
            np.float64(_NVFP4_PRODUCT_MAX)
            / np.float64(tensor_amax)
        )

    if not np.isfinite(divisor) or divisor <= 0:
        raise ValueError(
            f"{request.target.id}: invalid NVFP4 weight divisor {divisor}"
        )

    divisor_bytes = struct.pack("<f", float(divisor))

    print(
        f"[nvfp4 encoder] {request.target.id}: "
        f"shape={request.target.shape} "
        f"amax={float(tensor_amax):.9g} "
        f"d_w={float(divisor):.9g}",
        flush=True,
    )

    # ---------------------------------------------------------------
    # Pass 2: block16 E4M3 scales + E2M1 codes.
    # ---------------------------------------------------------------

    def produce(output):
        count = 0
        source_sse = 0.0
        error_sse = 0.0
        max_abs_error = 0.0

        for begin in range(0, n, chunk):
            end = min(n, begin + chunk)
            rows = end - begin

            host = _read_rows(request, begin, end, k)

            grouped = host.reshape(rows, k // 16, 16)

            group_amax = np.max(
                np.abs(grouped),
                axis=2,
            ).astype(np.float32)

            # Match the two-level NVFP4 maxabs construction.
            raw_scale = np.asarray(
                np.asarray(group_amax * divisor, dtype=np.float32)
                / np.float32(6.0),
                dtype=np.float32,
            )

            # Global divisor guarantees <= 448 apart from floating rounding.
            raw_scale = np.clip(
                raw_scale,
                np.float32(0.0),
                np.float32(448.0),
            ).astype(np.float32, copy=False)

            scale_words = _round_e4m3fn_rne(raw_scale)

            scale32 = _POSITIVE_E4M3FN_VALUES[
                scale_words.astype(np.int16)
            ].astype(np.float32, copy=False)

            normalized = np.zeros_like(grouped, dtype=np.float32)
            valid = scale32 > np.float32(0.0)

            if np.any(valid):
                scaled = np.asarray(
                    grouped[valid] * divisor,
                    dtype=np.float32,
                )

                normalized[valid] = np.asarray(
                    scaled / scale32[valid, None],
                    dtype=np.float32,
                )

            normalized = np.clip(
                normalized,
                np.float32(-6.0),
                np.float32(6.0),
            ).astype(np.float32, copy=False)

            code_words = _round_e2m1_rne(normalized)

            flat_codes = code_words.reshape(rows, k)

            packed = (
                flat_codes[:, 0::2]
                | (flat_codes[:, 1::2] << np.uint8(4))
            ).astype(np.uint8, copy=False)

            output.write_codes(
                begin,
                torch.from_numpy(packed.copy()),
                torch.from_numpy(scale_words.copy()),
                divisor_bytes,
            )

            # -------------------------------------------------------
            # Quantization diagnostics.
            # -------------------------------------------------------

            magnitudes = _E2M1[
                (code_words & np.uint8(0x7)).astype(np.int16)
            ]

            signs = np.where(
                (code_words & np.uint8(0x8)) != 0,
                np.float32(-1.0),
                np.float32(1.0),
            )

            reconstructed = np.asarray(
                magnitudes
                * signs
                * scale32[..., None]
                / divisor,
                dtype=np.float32,
            ).reshape(rows, k)

            error = host - reconstructed

            source_sse += float(
                np.sum(
                    host.astype(np.float64)
                    * host.astype(np.float64)
                )
            )

            error_sse += float(
                np.sum(
                    error.astype(np.float64)
                    * error.astype(np.float64)
                )
            )

            max_abs_error = max(
                max_abs_error,
                float(np.max(np.abs(error))),
            )

            count += host.size

        rmse = math.sqrt(error_sse / count)
        source_rms = math.sqrt(source_sse / count)

        relative_rmse = (
            rmse / source_rms
            if source_rms > 0
            else 0.0
        )

        print(
            f"[nvfp4 encoder] {request.target.id}: "
            f"rmse={rmse:.9g} "
            f"relative_rmse={relative_rmse:.6%} "
            f"max_abs_error={max_abs_error:.9g}",
            flush=True,
        )

    return request.job(produce=produce)


def configure(model, recipe, sources):
    # ---------------------------------------------------------------
    # A2.1 six Q6 keepers.
    # ---------------------------------------------------------------

    q6_keepers = (
        "text/layers/63/mlp/down",
        "text/layers/63/attention/output",
        "text/layers/62/gdn/output",
        "text/layers/62/mlp/down",
        "text/layers/61/gdn/output",
        "text/layers/61/mlp/down",
    )

    for name in q6_keepers:
        recipe.assign(
            name,
            format=Q6,
            method=grouped_absmax,
        )

    # ---------------------------------------------------------------
    # Replace FINAL4 L60-63 gate/up Q8 by experimental NVFP4.
    #
    # d_x=1.0 is deliberate:
    #   PERFORMANCE GATE ONLY.
    #
    # Route selection depends on policy/T, not the numerical divisor.
    # Before quality testing this MUST be replaced by calibrated d_x.
    # ---------------------------------------------------------------

    for layer in (60, 61, 62, 63):
        gate = f"text/layers/{layer}/mlp/gate"
        up   = f"text/layers/{layer}/mlp/up"

        for name in (gate, up):
            recipe.assign(
                name,
                format=NVFP4,
                method=nvfp4_ct_maxabs,
            )

            for input_name in model.parameters[name].inputs:
                recipe.use(
                    name,
                    input_name,
                    activation_policy="AllowA4",
                    auxiliaries={
                        "activation_input_divisor": 1.0,
                    },
                )

        recipe.group(
            (gate, up),
            shape=(34816, 5120),
        )