import logging
import math
import torch

log = logging.getLogger(__name__)


def build_temporal_cost(q_token_idx, Lq, Lk, device, dtype, tokens_per_frame):
    offset = torch.zeros(Lq, Lk, device=device, dtype=dtype)
    query_frames = torch.arange(Lq, device=device, dtype=torch.long) // tokens_per_frame
    for seg in q_token_idx:
        local = seg['local_token_idx'].to(device=device)
        d = (query_frames.float()[:, None] - seg['midpoint']).abs()
        cost = (torch.relu(d - seg['window']) ** 2) / (2 * seg['sigma'] ** 2)
        offset[:, local] = cost.to(offset.dtype)
    return offset


def build_temporal_cost_scaled(q_token_idx, Lq, Lk, device, dtype, latent_frames):
    offset = torch.zeros(Lq, Lk, device=device, dtype=dtype)
    query_frames = torch.arange(Lq, device=device, dtype=torch.float32) * latent_frames / Lq
    for seg in q_token_idx:
        local = seg['local_token_idx'].to(device=device)
        d = (query_frames[:, None] - seg['midpoint']).abs()
        cost = (torch.relu(d - seg['window']) ** 2) / (2 * seg['sigma'] ** 2)
        offset[:, local] = cost.to(offset.dtype)
    return offset


def create_mask_fn(q_token_idx, fallback_tokens_per_frame, latent_frames):
    cache = {}
    max_token_idx = max(int(seg['local_token_idx'].max().item()) for seg in q_token_idx) + 1

    def mask_fn(q, k, transformer_options):
        Lq, Lk = q.shape[1], k.shape[1]
        if Lq == Lk:
            return None
        cond_or_uncond = transformer_options.get('cond_or_uncond', [])
        if 1 in cond_or_uncond and 0 not in cond_or_uncond:
            return None
        grid_sizes = transformer_options.get('grid_sizes', None)
        video_tpf = int(grid_sizes[1]) * int(grid_sizes[2]) if grid_sizes is not None else fallback_tokens_per_frame
        video_lq = latent_frames * video_tpf
        if Lk == video_lq or Lk < max_token_idx:
            return None
        mode = 'video' if Lq == video_lq else 'scaled'
        key = (Lq, Lk, mode, q.device)
        if key not in cache:
            if mode == 'video':
                cost = build_temporal_cost(q_token_idx, Lq, Lk, q.device, q.dtype, video_tpf)
            else:
                cost = build_temporal_cost_scaled(q_token_idx, Lq, Lk, q.device, q.dtype, latent_frames)
            log.info('[PromptRelay] Built penalty matrix (%s): Lq=%d, Lk=%d, nonzero=%d/%d',
                     mode, Lq, Lk, (cost > 0).sum().item(), cost.numel())
            cache[key] = -cost
        return cache[key].to(q.dtype)

    return mask_fn


def build_segments(token_ranges, segment_lengths, epsilon=1e-3):
    sigma = 1.0 / math.log(1.0 / epsilon) if 0 < epsilon < 1 else 0.1448
    q_token_idx = []
    frame_cursor = 0
    for (tok_start, tok_end), L in zip(token_ranges, segment_lengths):
        if L <= 0:
            frame_cursor += L
            continue
        midpoint = (2 * frame_cursor + L) // 2
        window = max(L // 2 - 2, 0)
        q_token_idx.append({
            'local_token_idx': torch.arange(tok_start, tok_end),
            'midpoint': midpoint,
            'window': window,
            'sigma': sigma,
        })
        frame_cursor += L
    return q_token_idx


def get_raw_tokenizer(clip):
    tokenizer_wrapper = clip.tokenizer
    for attr_name in dir(tokenizer_wrapper):
        if attr_name.startswith('_'):
            continue
        inner = getattr(tokenizer_wrapper, attr_name, None)
        if inner is not None and hasattr(inner, 'tokenizer'):
            return inner.tokenizer
    raise RuntimeError(
        f'Could not find raw tokenizer on CLIP object. '
        f'Known attributes: {[a for a in dir(tokenizer_wrapper) if not a.startswith(chr(95))]}'
    )


def map_token_indices(raw_tokenizer, global_prompt, local_prompts):
    prefixed_locals = [' ' + lp for lp in local_prompts]
    full_prompt = global_prompt + ''.join(prefixed_locals)
    has_eos = getattr(raw_tokenizer, 'add_eos', False)
    eos_adj = 1 if has_eos else 0
    prev_len = len(raw_tokenizer(global_prompt)['input_ids']) - eos_adj
    token_ranges = []
    built = global_prompt
    for plp in prefixed_locals:
        built += plp
        cur_len = len(raw_tokenizer(built)['input_ids']) - eos_adj
        if cur_len <= prev_len:
            raise ValueError(f'Local prompt produced no tokens: {plp.strip()!r}')
        token_ranges.append((prev_len, cur_len))
        prev_len = cur_len
    return full_prompt, token_ranges


def distribute_segment_lengths(num_segments, latent_frames, specified_lengths=None):
    if specified_lengths:
        if len(specified_lengths) != num_segments:
            raise ValueError(
                f'Number of segment_lengths ({len(specified_lengths)}) '
                f'must match number of local prompts ({num_segments})'
            )
        lengths = specified_lengths
    else:
        step = -(-latent_frames // num_segments)
        lengths = [step] * num_segments
    effective = []
    cursor = 0
    for L in lengths:
        end = min(cursor + L, latent_frames)
        effective.append(max(end - cursor, 0))
        cursor = end
    return effective