"""Direct prompt-only H3 packed-token construction."""

import math

import torch
import torch.nn.functional as functional


FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
FRAME_RESCALE = 5.0 / 3.0


def patchify_video(latent: torch.Tensor) -> torch.Tensor:
    batch, channels, frames, height, width = latent.shape
    if batch != 1:
        raise ValueError("H3 supports batch size one.")
    if height % 2 or width % 2:
        raise ValueError("H3 video latent dimensions must be divisible by two.")
    return latent.reshape(batch, channels, frames, 1, height // 2, 2, width // 2, 2).permute(0, 2, 4, 6, 1, 3, 5, 7).reshape(-1, channels * 4)


def pack_audio(latent: torch.Tensor) -> torch.Tensor:
    return latent[0].permute(1, 2, 0).reshape(-1, latent.shape[1]).transpose(0, 1).contiguous().transpose(0, 1)


def unpatchify_video(rows: torch.Tensor, frames: int, latent_height: int, latent_width: int) -> torch.Tensor:
    height, width = latent_height // 2, latent_width // 2
    x = rows.reshape(1, frames, height, width, 24, 1, 2, 2).permute(0, 4, 1, 5, 2, 6, 3, 7)
    return x.reshape(1, 24, frames, latent_height, latent_width)


def _axis(dim: int, area: float) -> torch.Tensor:
    ratio, count = dim / area, dim // 2
    return (torch.arange(count, dtype=torch.float64) * ratio / count + (1 - ratio) / 2) * 32


def _video_positions(frames: int, height: int, width: int, offset: float) -> torch.Tensor:
    area = math.sqrt(height * width)
    ys, xs = torch.meshgrid(_axis(height, area), _axis(width, area), indexing="ij")
    spatial = torch.stack((ys.flatten(), xs.flatten()), dim=-1)
    spans = torch.tensor([FRAME_RESCALE * FRAME_PER_TOKEN[index % 5] for index in range(frames)], dtype=torch.float64)
    times = offset + torch.cat((torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)))
    result = torch.empty(frames, spatial.shape[0], 3, dtype=torch.float64)
    result[:, :, 0], result[:, :, 1:] = times[:, None], spatial[None]
    return result.reshape(-1, 3)


def _audio_positions(steps: int, offset: float, width: int, height: int) -> torch.Tensor:
    area = math.sqrt(height * width)
    x_axis = _axis(width, area)
    result = torch.zeros(steps * 2, 3, dtype=torch.float64)
    result[:, 0] = (offset + torch.arange(steps, dtype=torch.float64)).repeat(2)
    result[:steps, 2], result[steps:, 2] = x_axis[0], x_axis[-1]
    return result


class H3PromptPacker:
    """Build `[text | audio | video]` tokens for prompt-only H3 T2V."""

    def __init__(self, checkpoint):
        self.video_weight = checkpoint.tensor("video_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
        self.video_bias = checkpoint.tensor("video_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32)
        self.audio_weight = checkpoint.tensor("audio_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
        self.audio_bias = checkpoint.tensor("audio_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32)
        self.text_weight = checkpoint.tensor("condition_proj.weight", dtype=torch.bfloat16)
        self.text_bias = checkpoint.tensor("condition_proj.bias", dtype=torch.bfloat16)

    def __call__(self, text: torch.Tensor, video: torch.Tensor, audio: torch.Tensor, sigma: float) -> tuple[torch.Tensor, torch.Tensor, list[tuple[int, int, int]], tuple[int, int, int], tuple[int, int, int]]:
        if text.shape[-1] == 5120:
            text_rows = functional.linear(text[0].to(self.text_weight.dtype), self.text_weight, self.text_bias).to(torch.bfloat16)
        elif text.shape[-1] == 5376:
            text_rows = text[0].to(torch.bfloat16)
        else:
            raise ValueError("H3 text states must be Qwen 5120-wide or refined 5376-wide.")
        video_rows = functional.linear(patchify_video(video).float(), self.video_weight, self.video_bias).to(torch.bfloat16)
        audio_rows = functional.linear(pack_audio(audio).float(), self.audio_weight, self.audio_bias).to(torch.bfloat16)
        text_length, audio_length = text_rows.shape[0], audio_rows.shape[0]
        hidden = torch.cat((text_rows, audio_rows, video_rows))
        video_sigma = torch.tensor(float(sigma), device=hidden.device).clamp(min=1e-6)
        base = video_sigma / (12.0 + video_sigma * (1.0 - 12.0))
        audio_sigma = 3.0 * base / (1.0 + (3.0 - 1.0) * base)
        video_time, audio_time = 1 - float(video_sigma), 1 - float(audio_sigma)
        unique_times = sorted({video_time, audio_time})
        row = {value: index for index, value in enumerate(unique_times)}
        video_row, audio_row = row[video_time] * 3, row[audio_time] * 3
        times = torch.tensor(unique_times, device=hidden.device, dtype=torch.float32)
        positions = torch.cat((torch.stack((torch.arange(text_length, dtype=torch.float64), torch.zeros(text_length), torch.zeros(text_length)), dim=-1), _audio_positions(audio.shape[-1], float(text_length), video.shape[-1], video.shape[-2]), _video_positions(video.shape[2], video.shape[-2], video.shape[-1], float(text_length))))
        block_video_segment = (text_length + audio_length, hidden.shape[0], video_row)
        block_audio_segment = (text_length, text_length + audio_length, audio_row + 2)
        final_video_segment = (text_length + audio_length, hidden.shape[0], row[video_time])
        final_audio_segment = (text_length, text_length + audio_length, row[audio_time])
        return hidden, times, [(0, text_length, video_row + 1), block_audio_segment, block_video_segment], positions, final_video_segment, final_audio_segment