import torch
import torch.nn as nn
import numpy as np

class FSQ(nn.Module):
    def __init__(self, levels: list[int], dim: int = None):
        super().__init__()
        self.levels = levels
        self.dim = dim if dim is not None else len(levels)
        self.n_codes = np.prod(levels)
        # Register basis as a buffer so it's saved with the model but not a parameter
        self.register_buffer('basis', torch.tensor([np.prod(levels[i+1:]) for i in range(len(levels))] + [1], dtype=torch.int64)[:-1])
        
    def bound(self, z, eps=1e-3):
        """ Bound z to [-1 + eps, 1 - eps] to avoid dead zones at boundaries """
        return (1 - eps) * torch.tanh(z)

    def quantize(self, z):
        # z: [..., dim]
        # 1. Bound
        z = self.bound(z)
        
        # 2. Scale and shift to [0, L-1]
        quantized = []
        indices = []
        
        for i, L in enumerate(self.levels):
            # Scale to [0, L-1]
            # z_i is in [-1, 1]
            val = (z[..., i] + 1) / 2 * (L - 1)
            idx = torch.round(val)
            
            # Straight-Through Estimator
            # We want gradients to flow through z, but idx is discrete.
            # idx_ste = idx + (val - val.detach()) ? No.
            # Standard STE: q = z + (round(z) - z).detach()
            # Here: idx is the integer index.
            # The quantized value in [-1, 1] is:
            recon = idx / (L - 1) * 2 - 1
            
            # STE for reconstruction
            recon_ste = z[..., i] + (recon - z[..., i]).detach()
            
            quantized.append(recon_ste)
            indices.append(idx.long())
            
        quantized = torch.stack(quantized, dim=-1)
        indices = torch.stack(indices, dim=-1)
        
        return quantized, indices

    def indices_to_codes(self, indices):
        # indices: [..., dim]
        # Convert to single code index (if needed)
        return (indices * self.basis).sum(dim=-1)

    def codes_to_indices(self, codes):
        # codes: [...]
        indices = []
        for i in range(len(self.levels)):
            indices.append((codes // self.basis[i]) % self.levels[i])
        return torch.stack(indices, dim=-1)

    def indices_to_values(self, indices):
        # indices: [..., dim]
        values = []
        for i, L in enumerate(self.levels):
            # recon = idx / (L - 1) * 2 - 1
            val = indices[..., i].float() / (L - 1) * 2 - 1
            values.append(val)
        return torch.stack(values, dim=-1)