from __future__ import annotations

from functools import lru_cache
from typing import Iterable, List, Optional

import torch
import torch.nn as nn


def _powers_of_two(low: int, high: int) -> List[int]:
    vals = []
    x = 1
    while x < low:
        x *= 2
    while x <= high:
        vals.append(x)
        x *= 2
    return vals


def _is_power_of_two(x: int) -> bool:
    return x > 0 and (x & (x - 1)) == 0


@lru_cache(maxsize=128)
def _power_biased_weights(
    sample_horizon: int, power: float, device_str: str
) -> torch.Tensor:
    weights = (
        torch.arange(1, sample_horizon + 1, dtype=torch.float32, device=device_str)
        ** power
    )
    return weights / weights.sum()


class MaskedNestedDropout(nn.Module):
    """Progressive token dropout that preserves first k tokens and masks the rest."""

    def __init__(self, dim: int, size_sampling_mode: str = "pow2"):
        super().__init__()
        self.dim = dim
        self.size_sampling_mode = size_sampling_mode

        if size_sampling_mode != "disable":
            self.dropout_mask_token = nn.Parameter(torch.randn(dim) * 0.02)
            nn.init.trunc_normal_(self.dropout_mask_token, std=0.02)

    def _sample_keep_k(
        self, batch_size: int, horizon: int, device: torch.device
    ) -> torch.Tensor:
        mode = self.size_sampling_mode
        if mode == "uniform":
            return torch.randint(1, horizon + 1, (batch_size,), device=device)

        if mode == "pow2":
            if not _is_power_of_two(horizon):
                raise ValueError(
                    f"pow2 mode requires power-of-two horizon, got {horizon}"
                )
            vals = torch.tensor(_powers_of_two(1, horizon), device=device)
            idx = torch.randint(0, len(vals), (batch_size,), device=device)
            return vals[idx]

        if mode == "uniform_pow2":
            ks = torch.randint(1, horizon + 1, (batch_size,), device=device)
            keep = []
            for k in ks.tolist():
                p = 1
                while p < k:
                    p *= 2
                keep.append(p)
            return torch.tensor(keep, device=device)

        if mode == "drop":
            # Simple drop mode: randomly keep k tokens with higher probability for larger k
            # This provides a simple uniform-like sampling without power-of-two constraint
            return torch.randint(1, horizon + 1, (batch_size,), device=device)

        if mode.endswith("_biased"):
            power_map = {
                "linear_biased": 1.0,
                "quadratic_biased": 2.0,
                "cubic_biased": 3.0,
            }
            if mode not in power_map:
                raise ValueError(f"Unsupported biased mode: {mode}")
            weights = _power_biased_weights(horizon, power_map[mode], str(device))
            indices = torch.multinomial(weights, batch_size, replacement=True)
            return indices + 1

        raise ValueError(f"Unsupported size_sampling_mode: {mode}")

    def _mask_by_keep_k(self, x: torch.Tensor, keep_ks: Iterable[int]) -> torch.Tensor:
        keep = torch.as_tensor(list(keep_ks), device=x.device)
        _, n, _ = x.shape
        mask = torch.arange(n, device=x.device).unsqueeze(0) >= keep.unsqueeze(1)
        x[mask] = self.dropout_mask_token.to(device=x.device, dtype=x.dtype)
        return x

    def forward(
        self, x: torch.Tensor, eval_keep_k: Optional[List[int]] = None
    ) -> torch.Tensor:
        if self.size_sampling_mode == "disable":
            return x

        b, n, _ = x.shape
        if self.training:
            keep_ks = self._sample_keep_k(b, n, x.device)
            return self._mask_by_keep_k(x, keep_ks.tolist())

        if eval_keep_k is not None:
            if len(eval_keep_k) != b:
                raise ValueError(
                    f"eval_keep_k length {len(eval_keep_k)} does not match batch size {b}"
                )
            return self._mask_by_keep_k(x, eval_keep_k)

        return x
