import torch
from torch.export import Dim
from transformers import MoonshineForConditionalGeneration
from transformers.cache_utils import DynamicCache, EncoderDecoderCache

MODEL, OUT = 'moonshine-ai/moonshine-tiny-uk', '/work/out'
m = MoonshineForConditionalGeneration.from_pretrained(MODEL, attn_implementation='eager').eval()
cfg = m.config
L = cfg.decoder_num_hidden_layers
print('layers:', L)

def flatten(cache):
    """Cache object -> flat tensor list, self layers first, then cross."""
    out = []
    for sub in (cache.self_attention_cache, cache.cross_attention_cache):
        for layer in sub.layers:
            out += [layer.keys, layer.values]
    return out

def build(self_kv, cross_kv):
    """Flat tensors -> EncoderDecoderCache, with cross marked filled so the model
    reuses it instead of projecting the encoder output again on every step."""
    sa, ca = DynamicCache(), DynamicCache()
    for i in range(L):
        sa.update(self_kv[2*i], self_kv[2*i+1], i)
        ca.update(cross_kv[2*i], cross_kv[2*i+1], i)
    cache = EncoderDecoderCache(sa, ca)
    cache.is_updated = {i: True for i in range(L)}
    return cache

class DecoderInit(torch.nn.Module):
    def __init__(self, m): super().__init__(); self.m = m
    def forward(self, input_ids, encoder_hidden_states):
        out = self.m.model.decoder(input_ids=input_ids,
                                   encoder_hidden_states=encoder_hidden_states,
                                   use_cache=True)
        return (self.m.proj_out(out.last_hidden_state), *flatten(out.past_key_values))

class DecoderStep(torch.nn.Module):
    def __init__(self, m): super().__init__(); self.m = m
    def forward(self, input_ids, cache_position, encoder_hidden_states, *flat):
        self_kv, cross_kv = list(flat[:2*L]), list(flat[2*L:])
        cache = build(self_kv, cross_kv)
        # The cross cache alone is not enough: without the encoder states the decoder
        # skips cross-attention entirely and quietly answers from nothing. Passed here and
        # marked already-updated, so it is read rather than recomputed.
        out = self.m.model.decoder(input_ids=input_ids,
                                   past_key_values=cache,
                                   cache_position=cache_position,
                                   encoder_hidden_states=encoder_hidden_states,
                                   use_cache=True)
        new = flatten(out.past_key_values)[:2*L]
        return (self.m.proj_out(out.last_hidden_state), *new)

hs = torch.zeros(1, 123, 288)
start = torch.tensor([[cfg.decoder_start_token_id]])
init = DecoderInit(m).eval()
with torch.no_grad():
    init_out = init(start, hs)
print('init outputs:', len(init_out), '| logits', tuple(init_out[0].shape),
      '| self k', tuple(init_out[1].shape), '| cross k', tuple(init_out[1+2*L].shape))

self_names  = [f'{p}_self_{i}' for i in range(L) for p in ('key','value')]
cross_names = [f'{p}_cross_{i}' for i in range(L) for p in ('key','value')]
frames = Dim('frames', min=4, max=2000)
past = Dim('past', min=1, max=448)

torch.onnx.export(
    init, (start, hs), f'{OUT}/decoder_init.onnx', dynamo=True,
    input_names=['input_ids', 'encoder_hidden_states'],
    output_names=['logits'] + [f'present_{n}' for n in self_names + cross_names],
    dynamic_shapes={'input_ids': {}, 'encoder_hidden_states': {1: frames}},
)
print('decoder_init exported')

flat = list(init_out[1:])
step = DecoderStep(m).eval()
pos = torch.tensor([1], dtype=torch.long)
nxt = torch.tensor([[100]])
with torch.no_grad():
    step_out = step(nxt, pos, hs, *flat)
print('step outputs:', len(step_out), '| logits', tuple(step_out[0].shape))

torch.onnx.export(
    step, (nxt, pos, hs, *flat), f'{OUT}/decoder_step.onnx', dynamo=True,
    input_names=['input_ids', 'cache_position', 'encoder_hidden_states']
                + [f'past_{n}' for n in self_names + cross_names],
    output_names=['logits'] + [f'present_{n}' for n in self_names],
    dynamic_shapes=({}, {}, {1: frames},
                    tuple([{2: past}] * len(self_names) + [{2: frames}] * len(cross_names))),
)
print('decoder_step exported')