import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from ..octree import DfsOctree as Octree


class Strivec(Octree):
    def __init__(
        self,
        resolution: int,
        aabb: list,
        sh_degree: int = 0,
        rank: int = 8,
        dim: int = 8,
        device: str = "cuda",
    ):
        assert np.log2(resolution) % 1 == 0, "Resolution must be a power of 2"
        self.resolution = resolution
        depth = int(np.round(np.log2(resolution)))
        super().__init__(
            depth=depth,
            aabb=aabb,
            sh_degree=sh_degree,
            primitive="trivec",
            primitive_config={"rank": rank, "dim": dim},
            device=device,
        )