# @lint-ignore-every LICENSELINT
# Adapted from https://github.com/lioryariv/idr/blob/main/code/model/
#              implicit_differentiable_renderer.py
# Copyright (c) 2020 Lior Yariv
import math
from typing import Sequence

import torch
from pytorch3d.implicitron.tools.config import registry
from pytorch3d.renderer.implicit import HarmonicEmbedding
from torch import nn

from .base import ImplicitFunctionBase


@registry.register
class IdrFeatureField(ImplicitFunctionBase, torch.nn.Module):
    """
    Implicit function as used in http://github.com/lioryariv/idr.

    Members:
        d_in: dimension of the input point.
        n_harmonic_functions_xyz: If -1, do not embed the point.
            If >=0, use a harmonic embedding with this number of
            harmonic functions. (The harmonic embedding includes the input
            itself, so a value of 0 means the point is used but without
            any harmonic functions.)
        d_out and feature_vector_size: Sum of these is the output
            dimension. This implicit function thus returns a concatenation
            of `d_out` signed distance function values and `feature_vector_size`
            features (such as colors). When used in `GenericModel`,
            `feature_vector_size` corresponds is automatically set to
            `render_features_dimensions`.
        dims: list of hidden layer sizes.
        geometric_init: whether to use custom weight initialization
            in linear layers. If False, pytorch default (uniform sampling)
            is used.
        bias: if geometric_init=True, initial value for bias subtracted
            in the last layer.
        skip_in: List of indices of layers that receive as input the initial
            value concatenated with the output of the previous layers.
        weight_norm: whether to apply weight normalization to each layer.
        pooled_feature_dim: If view pooling is in use (provided as
            fun_viewpool to forward()) this must be its number of features.
            Otherwise this must be set to 0. (If used from GenericModel,
            this config value will be overridden automatically.)
        encoding_dim: If global coding is in use (provided as global_code
            to forward()) this must be its number of featuress.
            Otherwise this must be set to 0. (If used from GenericModel,
            this config value will be overridden automatically.)
    """

    feature_vector_size: int = 3
    d_in: int = 3
    d_out: int = 1
    dims: Sequence[int] = (512, 512, 512, 512, 512, 512, 512, 512)
    geometric_init: bool = True
    bias: float = 1.0
    skip_in: Sequence[int] = ()
    weight_norm: bool = True
    n_harmonic_functions_xyz: int = 0
    pooled_feature_dim: int = 0
    encoding_dim: int = 0

    def __post_init__(self):
        super().__init__()

        dims = [self.d_in] + list(self.dims) + [self.d_out + self.feature_vector_size]

        self.embed_fn = None
        if self.n_harmonic_functions_xyz >= 0:
            self.embed_fn = HarmonicEmbedding(
                self.n_harmonic_functions_xyz, append_input=True
            )
            dims[0] = self.embed_fn.get_output_dim()
        if self.pooled_feature_dim > 0:
            dims[0] += self.pooled_feature_dim
        if self.encoding_dim > 0:
            dims[0] += self.encoding_dim

        self.num_layers = len(dims)

        out_dim = 0
        layers = []
        for layer_idx in range(self.num_layers - 1):
            if layer_idx + 1 in self.skip_in:
                out_dim = dims[layer_idx + 1] - dims[0]
            else:
                out_dim = dims[layer_idx + 1]

            lin = nn.Linear(dims[layer_idx], out_dim)

            if self.geometric_init:
                if layer_idx == self.num_layers - 2:
                    torch.nn.init.normal_(
                        lin.weight,
                        mean=math.pi**0.5 / dims[layer_idx] ** 0.5,
                        std=0.0001,
                    )
                    torch.nn.init.constant_(lin.bias, -self.bias)
                elif self.n_harmonic_functions_xyz >= 0 and layer_idx == 0:
                    torch.nn.init.constant_(lin.bias, 0.0)
                    torch.nn.init.constant_(lin.weight[:, 3:], 0.0)
                    torch.nn.init.normal_(
                        lin.weight[:, :3], 0.0, 2**0.5 / out_dim**0.5
                    )
                elif self.n_harmonic_functions_xyz >= 0 and layer_idx in self.skip_in:
                    torch.nn.init.constant_(lin.bias, 0.0)
                    torch.nn.init.normal_(lin.weight, 0.0, 2**0.5 / out_dim**0.5)
                    torch.nn.init.constant_(lin.weight[:, -(dims[0] - 3) :], 0.0)
                else:
                    torch.nn.init.constant_(lin.bias, 0.0)
                    torch.nn.init.normal_(lin.weight, 0.0, 2**0.5 / out_dim**0.5)

            if self.weight_norm:
                lin = nn.utils.weight_norm(lin)

            layers.append(lin)

        self.linear_layers = torch.nn.ModuleList(layers)
        self.out_dim = out_dim
        self.softplus = nn.Softplus(beta=100)

    # pyre-fixme[14]: `forward` overrides method defined in `ImplicitFunctionBase`
    #  inconsistently.
    def forward(
        self,
        # ray_bundle: RayBundle,
        rays_points_world: torch.Tensor,  # TODO: unify the APIs
        fun_viewpool=None,
        global_code=None,
    ):
        # this field only uses point locations
        # rays_points_world = ray_bundle_to_ray_points(ray_bundle)
        # rays_points_world.shape = [minibatch x ... x pts_per_ray x 3]

        if rays_points_world.numel() == 0 or (
            self.embed_fn is None and fun_viewpool is None and global_code is None
        ):
            return torch.tensor(
                [],
                device=rays_points_world.device,
                dtype=rays_points_world.dtype
                # pyre-fixme[6]: For 2nd param expected `int` but got `Union[Module,
                #  Tensor]`.
            ).view(0, self.out_dim)

        embeddings = []
        if self.embed_fn is not None:
            # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
            embeddings.append(self.embed_fn(rays_points_world))

        if fun_viewpool is not None:
            assert rays_points_world.ndim == 2
            pooled_feature = fun_viewpool(rays_points_world[None])
            # TODO: pooled features are 4D!
            embeddings.append(pooled_feature)

        if global_code is not None:
            assert global_code.shape[0] == 1  # TODO: generalize to batches!
            # This will require changing raytracer code
            # embedding = embedding[None].expand(global_code.shape[0], *embedding.shape)
            embeddings.append(
                global_code[0, None, :].expand(rays_points_world.shape[0], -1)
            )

        embedding = torch.cat(embeddings, dim=-1)
        x = embedding
        # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C._TensorBase.__s...
        for layer_idx in range(self.num_layers - 1):
            if layer_idx in self.skip_in:
                x = torch.cat([x, embedding], dim=-1) / 2**0.5

            # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
            x = self.linear_layers[layer_idx](x)

            # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C._TensorBase...
            if layer_idx < self.num_layers - 2:
                # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
                x = self.softplus(x)

        return x  # TODO: unify the APIs
