import math

import torch
import torch.nn as nn
import torch.nn.functional as F


def embedding(x, L):
    freq_bands = 2**torch.arange(L, dtype=torch.float, device=x.device)
    return torch.cat([torch.cat([torch.sin(x * freq), torch.cos(x * freq)], -1)
                      for freq in freq_bands], -1)


class MLP(nn.Module):
    def __init__(
            self,
            pos_emb_L=10,
            ray_emb_L=4,
            hidden_dim=256,
            n_layers=8,
            n_aux_layers=1,
            skips=[4]
            ):
        super().__init__()
        self.pos_emb_L = pos_emb_L
        self.ray_emb_L = ray_emb_L
        self.skips = skips
        pos_emb_dim = 3 * 2 * pos_emb_L + 3
        ray_emb_dim = 3 * 2 * ray_emb_L + 3
        self.body = nn.ModuleList(
                            [nn.Linear(pos_emb_dim, hidden_dim)] + 
                            [nn.Linear(hidden_dim + pos_emb_dim if i in skips else hidden_dim, hidden_dim)
                             for i in range(n_layers-1)])

        self.emb_body = nn.Linear(hidden_dim, hidden_dim)
        self.regress_density = nn.Linear(hidden_dim, 1)

        if n_aux_layers > 0:
            aux_layers = [nn.Sequential(nn.Linear(ray_emb_dim + hidden_dim, hidden_dim//2), nn.ReLU(True))] + \
                         [nn.Sequential(nn.Linear(hidden_dim//2, hidden_dim//2), nn.ReLU(True)) for _ in range(n_aux_layers-1)]
        else:
            hidden_dim = 2 * (ray_emb_dim + hidden_dim)
            aux_layers = []
        self.regress_color = nn.Sequential(*aux_layers, nn.Linear(hidden_dim//2, 3))

    def forward(self, xyz, ray):
        ray_emb = torch.cat([ray, embedding(ray, self.ray_emb_L)], -1)
        pos_emb = torch.cat([xyz, embedding(xyz, self.pos_emb_L)], -1)
        feat = pos_emb
        for i, linear in enumerate(self.body):
            feat = linear(feat)
            feat = F.relu(feat)
            if i in self.skips:
                feat = torch.cat([feat, pos_emb], -1)
        
        sigma = self.regress_density(feat)
        feat = self.emb_body(feat)
        feat = torch.cat([feat, ray_emb], -1)
        rgb = self.regress_color(feat)

        return sigma, rgb
