"""Ordered native Mamba backward reductions for the bounded M4 trajectory.

Candidate policy: GPU gates and a real128 distributed restart are mandatory.
Forward equations and parameters are unchanged. Native CUDA/Triton calls remain
responsible for derivatives; only the order of shared-parameter sums changes.
"""
from __future__ import annotations

import importlib

import torch

POLICY = 'native_single_writer_ordered_reductions_v1'
_installed = None


def install():
    """Freeze native tiles and remove concurrent writes to parameter gradients."""
    global _installed
    if _installed is not None:
        return _installed
    from causal_conv1d.cpp_functions import causal_conv1d_bwd_function
    combined = importlib.import_module('mamba_ssm.ops.triton.ssd_combined')
    native_cumsum = combined._chunk_cumsum_bwd
    modules = ('ssd_combined', 'ssd_chunk_state', 'ssd_chunk_scan',
               'ssd_state_passing', 'ssd_bmm', 'layernorm_gated', 'k_activations')
    tiles = {}
    seen = set()
    single_writer_kernels = {'_chunk_scan_chunk_state_bwd_dx_kernel': 64,
                            '_chunk_state_bwd_db_kernel': 128,
                            '_chunk_scan_bwd_dc_kernel': 128}
    constrained = set()
    for suffix in modules:
        module = importlib.import_module('mamba_ssm.ops.triton.' + suffix)
        for name, kernel in vars(module).items():
            if not hasattr(kernel, 'configs') or not hasattr(kernel, 'cache') or id(kernel) in seen:
                continue
            seen.add(id(kernel))
            configs = list(kernel.configs)
            if not configs:
                raise RuntimeError('Native autotuner has no configuration: ' + name)
            # Large native candidates can exceed H100 shared memory; autotuning
            # normally rejects those. Freeze the smallest available native tile
            # subject to the single-writer head/state-axis constraint.
            minimum_n = single_writer_kernels.get(name, 0)
            if minimum_n:
                constrained.add(name)
            eligible = [c for c in configs if c.kwargs.get('BLOCK_SIZE_N', 0) >= minimum_n]
            if not eligible:
                raise RuntimeError('No native single-writer configuration: ' + name)
            def tile_size(config):
                dims = [v for k, v in config.kwargs.items() if k.startswith('BLOCK_SIZE') and isinstance(v, int)]
                import math
                return (math.prod(dims), config.num_stages, config.num_warps,
                        tuple(sorted(config.kwargs.items())))
            chosen = min(eligible, key=tile_size)
            kernel.configs = [chosen]
            kernel.cache.clear()
            tiles[suffix + '.' + name] = {'kwargs': dict(chosen.kwargs),
                'num_warps': chosen.num_warps, 'num_stages': chosen.num_stages}
    if constrained != set(single_writer_kernels):
        raise RuntimeError('Pinned native single-writer kernel inventory changed')

    def ordered_conv(x, weight, bias, dout, seq_idx, initial_states,
                     dfinal_states, dx, return_dinitial_states, silu_activation):
        if any(v is not None for v in (seq_idx, initial_states, dfinal_states)) or return_dinitial_states:
            raise ValueError('Ordered training convolution requires reset sequences')
        # Retain the original full-batch input gradient and its memory layout.
        dx_out, _, _, _ = causal_conv1d_bwd_function(
            x, weight, bias, dout, None, None, None, dx, False, silu_activation)
        dw = torch.zeros_like(weight, dtype=torch.float32)
        db = None if bias is None else torch.zeros_like(bias, dtype=torch.float32)
        # Native channels-first B=1 reduction has no cross-example atomics.
        for batch in range(x.shape[0]):
            _, part_w, part_b, _ = causal_conv1d_bwd_function(
                x[batch:batch + 1].contiguous(), weight.float(),
                None if bias is None else bias.float(), dout[batch:batch + 1].contiguous(),
                None, None, None, None, False, silu_activation)
            dw.add_(part_w)
            if db is not None:
                db.add_(part_b)
        return dx_out, dw.to(weight.dtype), None if db is None else db.to(bias.dtype), None

    def ordered_cumsum(ddA, ddt_out, dt, A, dt_bias=None, dt_softplus=False,
                       dt_limit=(0.0, float('inf')), ddt=None):
        # Native cumsum backward's parameter sums use atomic_add across B/C.
        # Each call below has B=C=1; accumulate those partials in lexical order.
        batch, length, heads = dt.shape
        chunks, width = ddA.shape[-2:]
        if ddA.shape != (batch, heads, chunks, width) or ddt_out.shape != ddA.shape:
            raise ValueError('Unexpected SSD chunk derivatives')
        if ddt is None:
            ddt = torch.empty_like(dt)
        dA = torch.zeros_like(A, dtype=torch.float32)
        db = None if dt_bias is None else torch.zeros_like(dt_bias, dtype=torch.float32)
        for b in range(batch):
            for c in range(chunks):
                begin, end = c * width, min((c + 1) * width, length)
                if begin >= end:
                    continue
                _, part_a, part_b = native_cumsum(
                    ddA[b:b + 1, :, c:c + 1], ddt_out[b:b + 1, :, c:c + 1],
                    dt[b:b + 1, begin:end], A, dt_bias=dt_bias,
                    dt_softplus=dt_softplus, dt_limit=dt_limit,
                    ddt=ddt[b:b + 1, begin:end])
                dA.add_(part_a)
                if db is not None:
                    db.add_(part_b)
        return ddt, dA, db

    combined.causal_conv1d_bwd_function = ordered_conv
    combined._chunk_cumsum_bwd = ordered_cumsum
    _installed = {'policy': POLICY, 'tiles': tiles,
                  'convolution_reduction': 'native_B1_channels_first_fixed_FP32_sum',
                  'cumsum_reduction': 'native_B1_C1_fixed_FP32_sum'}
    return _installed