import torch
from torch import nn
from torch.nn import functional
from torch.nn.attention import SDPBackend, sdpa_kernel
from torch.utils.checkpoint import checkpoint

CAUSAL_CLOCK_POLICY = 'paired_av_clocks_v1'
CAUSAL_BATCH_NUMERICS_POLICY = 'chronological_kv_v1'

def embed_plan_timesteps(model, plan):
    if plan.bidirectional and plan.clock_groups:
        parts = [model.time_embedder(model.time_proj(clocks).to(model.time_embedder.linear_1.weight.dtype)) for clocks in plan.clock_groups]
        clocks = torch.cat(plan.clock_groups)
        values = torch.cat(parts)
        owners = torch.stack([torch.where(clocks == time_input)[0][0] for time_input in plan.timesteps])
        return values.index_select(0, owners)
    if plan.bidirectional:
        return model.time_embedder(model.time_proj(plan.timesteps).to(model.time_embedder.linear_1.weight.dtype))
    parts = [model.time_embedder(model.time_proj(clocks).to(model.time_embedder.linear_1.weight.dtype)) for clocks in plan.clock_groups]
    clocks = torch.cat(plan.clock_groups)
    order = clocks.argsort()
    if not torch.equal(clocks[order], plan.timesteps):
        raise ValueError('Native paired clocks no longer cover the fixed inventory')
    return torch.cat(parts).index_select(0, order)

def shared_kv_gradient_source(value):
    if torch.is_grad_enabled() and value.requires_grad and (value.dtype in (torch.bfloat16, torch.float16)):
        return value.float()
    return value

def _attention(attn, hidden, rotary, plan):
    from diffusers.models.transformers.transformer_minimax_h3 import _apply_rotary_emb
    if attn.fused_projections:
        raise ValueError('Chunked runtime requires unfused base projections and a shared LoRA owner')
    query, key, value = [getattr(attn, name)(hidden).unflatten(-1, (attn.heads, -1)) for name in ('to_q', 'to_k', 'to_v')]
    query, key = (attn.norm_q(query), attn.norm_k(key))
    query = torch.cat([_apply_rotary_emb(query[batch_index:batch_index + 1], *rotary[batch_index]) for batch_index in range(query.shape[0])])
    key = torch.cat([_apply_rotary_emb(key[batch_index:batch_index + 1], *rotary[batch_index]) for batch_index in range(key.shape[0])])
    query, key, value = (values.transpose(1, 2) for values in (query, key, value))
    if query.shape[2] > 4096 and query.device.type != 'cuda':
        raise ValueError('Full-resolution attention requires GPU execution')
    backends = [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]
    if query.shape[2] <= 4096:
        backends.append(SDPBackend.MATH)
    outputs = []
    for batch_index, groups in enumerate(plan.groups):
        out = torch.zeros_like(query[batch_index:batch_index + 1])
        key_source = shared_kv_gradient_source(key[batch_index:batch_index + 1])
        value_source = shared_kv_gradient_source(value[batch_index:batch_index + 1])
        for queries, keys in groups:
            with sdpa_kernel(backends):
                value = functional.scaled_dot_product_attention(query[batch_index:batch_index + 1].index_select(2, queries), key_source.index_select(2, keys).to(key.dtype), value_source.index_select(2, keys).to(value.dtype), dropout_p=0.0)
            out = out.index_copy(2, queries, value)
        outputs.append(out)
    merged = torch.cat(outputs).transpose(1, 2).flatten(2, 3).type_as(query)
    return attn.to_out[1](attn.to_out[0](merged))

def _block_forward(block, hidden, temb, indices, rotary, plan):
    shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = block.adaln_proj(temb)
    normalized = block.norm1(hidden) * (1 + scale_a[indices]) + shift_a[indices]
    hidden = hidden + gate_a[indices] * _attention(block.attn, normalized, rotary, plan)
    normalized = block.norm2(hidden) * (1 + scale_f[indices]) + shift_f[indices]
    return hidden + gate_f[indices] * block.ff(normalized)

class ChunkedH3Runtime(nn.Module):

    def __init__(self, model, *, checkpoint_blocks=True):
        super().__init__()
        self.model = model
        self.checkpoint_blocks = checkpoint_blocks
        self.attention_contract = 'explicit_causal_rows_stage_text_native_clocks_kv_order_v4'
        self.causal_clock_policy = CAUSAL_CLOCK_POLICY
        self.causal_batch_numerics_policy = CAUSAL_BATCH_NUMERICS_POLICY

    def forward(self, video_rows, audio_rows, prompt_embeds, plan, *, return_hidden=False):
        from diffusers.models.transformers.transformer_minimax_h3 import MiniMaxH3TransformerOutput
        model = self.model
        batch, row_count = plan.routes.valid_rows.shape
        if video_rows.shape[:2] != (batch, plan.video_indices.numel()):
            raise ValueError('Video input/plan shape differs')
        if audio_rows.shape[:2] != (batch, plan.audio_indices.numel()):
            raise ValueError('Audio input/plan shape differs')
        if prompt_embeds.shape[:2] != (batch, plan.max_text_length):
            raise ValueError('Text input/plan shape differs')
        if len({video_rows.device, audio_rows.device, prompt_embeds.device, plan.routes.token_tags.device}) != 1:
            raise ValueError('Inputs and plan must share a device')
        video = model.proj_in(video_rows.to(model.proj_in.weight.dtype))
        audio = model.audio_proj_in(audio_rows.to(model.audio_proj_in.weight.dtype))
        refined = []
        for batch_index, length in enumerate(plan.text_lengths):
            one = model.context_embedder(prompt_embeds[batch_index:batch_index + 1, :length].to(model.context_embedder.weight.dtype))
            one = model.token_refiner(one)
            refined.append(functional.pad(one, (0, 0, 0, plan.max_text_length - length)))
        text = torch.cat(refined).repeat(1, plan.copies, 1)
        hidden = text.new_zeros((batch, row_count, text.shape[-1]))
        hidden = hidden.index_copy(1, plan.text_indices, text)
        hidden = hidden.index_copy(1, plan.video_indices, video.to(text.dtype))
        hidden = hidden.index_copy(1, plan.audio_indices, audio.to(text.dtype))
        temb = embed_plan_timesteps(model, plan)
        indices = plan.timestep_indices * 3 + plan.routes.token_tags[None].clamp(min=0)
        rotary = tuple((model.rope(pos) for pos in plan.position_ids))
        for block in model.transformer_blocks:

            def call(hidden_input, time_input, block=block):
                return _block_forward(block, hidden_input, time_input, indices, rotary, plan)
            if self.checkpoint_blocks and torch.is_grad_enabled():
                hidden = checkpoint(call, hidden, temb, use_reentrant=False)
            else:
                hidden = call(hidden, temb)
        shift, scale = model.norm_out.linear(functional.silu(temb).to(model.norm_out.linear.weight.dtype)).chunk(2, dim=-1)
        hidden = model.norm_out.norm(hidden) * (1 + scale[plan.timestep_indices]) + shift[plan.timestep_indices]
        video_hidden = hidden.index_select(1, plan.video_indices).to(model.proj_out.weight.dtype)
        audio_hidden = hidden.index_select(1, plan.audio_indices).to(model.audio_proj_out.weight.dtype)
        video = torch.cat([model.proj_out(video_hidden[batch_index:batch_index + 1]) for batch_index in range(batch)])
        audio = torch.cat([model.audio_proj_out(audio_hidden[batch_index:batch_index + 1]) for batch_index in range(batch)])
        output = MiniMaxH3TransformerOutput(sample=video, audio_sample=audio)
        return (output, hidden) if return_hidden else output