"""Dependency-free architecture arithmetic and shared initialization mapping."""
from __future__ import annotations

MAMBA_REVISION = 'd7b1ceb3c367ec9022925e812f507bcf706937c6'  # v2.2.6
BASE_REVISION = '9b98f6e75713212b737d8b7bf7c3fb7913e14ee6'
CORPUS_SHA256 = '1786fc7cf798784edf2e0ee9400456ca93add1cbc5688cd8e83865cb230bc475'
CAPS = (1, 2, 4, 8, 16, 32, 64, 96, 128, 192, 256, 384, 512, 768, 1024, 1536, 2048)


def shared_parameter_name(name, outer_layers=4):
    """Map new middle-block names to the original scratch tensor names.

    M4/20/M4 indices 4..23 correspond to T2/20/T2 indices 2..21.
    The same two-for-one mapping supports M2/2/M2 tiny correctness models.
    """
    if outer_layers < 2 or outer_layers % 2:
        raise ValueError('Two native Mamba blocks replace one Transformer block')
    if not name.startswith('layers.'):
        return name
    parts = name.split('.')
    parts[1] = str(int(parts[1]) - outer_layers // 2)
    return '.'.join(parts)


def parameter_counts(hidden=1024, intermediate=2816, vocab=32000,
                     middle_layers=20, outer_layers=4, state=128,
                     expand=2, headdim=64, conv=4):
    inner = hidden * expand
    if inner % headdim:
        raise ValueError('Mamba inner width must divide into complete heads')
    heads = inner // headdim
    transformer = 4 * hidden**2 + 3 * hidden * intermediate + 2 * hidden
    mamba = (hidden * (2 * inner + 2 * state + heads)
             + (inner + 2 * state) * (conv + 1)
             + 3 * heads + inner + inner * hidden + hidden)
    shared = vocab * hidden + middle_layers * transformer + hidden
    old = shared + outer_layers * transformer
    new = shared + 2 * outer_layers * mamba
    return {'transformer_block': transformer, 'mamba_block_with_prenorm': mamba,
            'original_model': old, 'mamba_outer_model': new,
            'increase_percent': 100 * (new - old) / old}


def theoretical_state_bytes(cap, *, batch=1, hidden=1024, middle_layers=20,
                            outer_layers=4, expand=2, state=128, conv=4,
                            kv_item_bytes=2, ssm_item_bytes=4):
    if min(cap, batch, kv_item_bytes, ssm_item_bytes) < 1:
        raise ValueError('Cap, batch and scalar byte sizes must be positive')
    inner = hidden * expand
    kv = batch * middle_layers * 2 * hidden * cap * kv_item_bytes
    recurrent = batch * 2 * outer_layers * inner * state * ssm_item_bytes
    convolution = batch * 2 * outer_layers * (inner + 2 * state) * conv * kv_item_bytes
    return {'kv': kv, 'ssm': recurrent, 'convolution': convolution,
            'total_tensor_bytes': kv + recurrent + convolution,
            'excludes': ['host metadata', 'shared positional buffers', 'weights', 'temporary workspace']}