"""
GNDCNetworkTorch.py - Pure PyTorch implementation of GNDCNetwork.

Drop-in replacement for GNDCNetwork (TCNN version).
Hash grid and MLP implementations match TCNN's exact behavior:
  - Hash grid: resolution formula, per-level table sizes, coordinate scaling,
    hash primes all match tiny-cuda-nn
  - MLP: no bias, ReLU activation, 16-aligned padding, row-major weights

Usage:
    from GNDCNetworkTorch import GNDCNetworkTorch
    model = GNDCNetworkTorch(**config_params)
    model.load_weights_from_flat_buffer(flat_buffer)
    output, offsets = model(coords)
"""

import math
import torch
import torch.nn as nn


# ─────────────────────────────────────────────────────────────────────────────
# Hash Grid — matches TCNN exactly
# ─────────────────────────────────────────────────────────────────────────────

class _HashGridTorch(nn.Module):
    """Hash grid encoding matching tiny-cuda-nn's HashGrid implementation."""

    PRIMES = [1, 2654435761, 805459861, 3674653429,
              2097192037, 1434869437, 2165219737]

    def __init__(self, dim, n_levels, n_features_per_level,
                 log2_hashmap_size, base_resolution, per_level_scale):
        super().__init__()
        self.dim = dim
        self.n_levels = n_levels
        self.n_features_per_level = n_features_per_level
        self.log2_hashmap_size = log2_hashmap_size
        max_hash_size = 2 ** log2_hashmap_size

        # Compute per-level resolution and table size exactly as TCNN does.
        # TCNN: n_vertices = ceil(base_resolution * per_level_scale^level)
        #        table_size = min(2^log2_hashmap_size, n_vertices^dim)
        #        TCNN aligns each level's entries to multiple of 8
        #        Coord scaling uses raw float: scale = base_res * per_level_scale^level - 1
        def _align8(x):
            return ((x + 7) // 8) * 8

        # Use float32 precision to match TCNN's exp2f computation
        import numpy as _np
        log2_scale_f32 = _np.float32(_np.log2(_np.float32(per_level_scale)))

        resolutions = []     # n_vertices per dimension (for clamping)
        level_sizes = []     # actual hash table entries per level (for lookup)
        level_scales = []    # float coordinate scale per level (TCNN grid_scale)
        level_offsets = [0]  # param offsets (including alignment padding)
        for i in range(n_levels):
            # TCNN: grid_scale (float) = exp2f(level * log2(per_level_scale)) * base_res - 1
            raw_scale = float(_np.float32(_np.exp2(_np.float32(i) * log2_scale_f32))
                              * _np.float32(base_resolution))
            n_verts = math.ceil(raw_scale)
            float_scale = raw_scale - 1.0

            resolutions.append(n_verts)
            level_scales.append(float_scale)
            grid_size = n_verts ** dim
            table_size = min(max_hash_size, grid_size)
            level_sizes.append(table_size)
            # TCNN aligns each level's entry count to multiple of 8
            aligned_size = _align8(table_size)
            level_offsets.append(level_offsets[-1] + aligned_size * n_features_per_level)

        # Track which levels are dense (direct index) vs hashed
        level_is_hashed = [level_sizes[i] < resolutions[i] ** dim
                           for i in range(n_levels)]

        self.resolutions = resolutions
        self.level_sizes = level_sizes
        self.level_scales = level_scales
        self.level_is_hashed = level_is_hashed
        self.register_buffer('level_offsets',
                             torch.tensor(level_offsets, dtype=torch.long),
                             persistent=False)

        total_params = level_offsets[-1]
        self.hash_table = nn.Parameter(torch.zeros(total_params))
        nn.init.uniform_(self.hash_table, -1e-4, 1e-4)

    @property
    def output_dim(self):
        return self.n_levels * self.n_features_per_level

    @property
    def n_params(self):
        return self.hash_table.numel()

    def _hash_index(self, c, table_size):
        """Hash function matching TCNN (for hashed levels)."""
        h = torch.zeros(c.shape[0], dtype=torch.long, device=c.device)
        for d in range(c.shape[1]):
            h = h ^ (c[:, d].long() * self.PRIMES[d] & 0xFFFFFFFF)
        return h % table_size

    def _ravel_index(self, c, n_verts):
        """Direct linear index matching TCNN (for dense levels).
        Strides: [1, n_verts, n_verts^2, ...]"""
        idx = c[:, 0].long()
        stride = 1
        for d in range(1, self.dim):
            stride *= n_verts
            idx = idx + c[:, d].long() * stride
        return idx

    def _interp(self, coords, level):
        nf = self.n_features_per_level
        n_verts = self.resolutions[level]
        table_size = self.level_sizes[level]
        is_hashed = self.level_is_hashed[level]
        base_offset = self.level_offsets[level].item()

        # TCNN: pos = x * grid_scale + 0.5 (grid_scale is float, not integer)
        scaled = coords * self.level_scales[level] + 0.5
        floorf = scaled.floor()
        frac = scaled - floorf

        result = torch.zeros(coords.shape[0], nf,
                             device=coords.device, dtype=coords.dtype)
        for corner in range(2 ** self.dim):
            w = torch.ones(coords.shape[0], 1,
                           device=coords.device, dtype=coords.dtype)
            c = floorf.long().clone()
            for d in range(self.dim):
                if corner & (1 << d):
                    c[:, d] += 1
                    w = w * frac[:, d:d + 1]
                else:
                    w = w * (1.0 - frac[:, d:d + 1])
            # Clamp to valid grid range
            c = c.clamp(0, n_verts - 1)
            # TCNN: dense levels use ravel index, hashed levels use hash
            if is_hashed:
                idx = self._hash_index(c, table_size)
            else:
                idx = self._ravel_index(c, n_verts)
            idx = idx * nf + base_offset
            idx_exp = idx.unsqueeze(1) + torch.arange(nf, device=coords.device)
            result = result + self.hash_table[idx_exp] * w
        return result

    def forward(self, coords):
        return torch.cat([
            self._interp(coords, lvl)
            for lvl in range(self.n_levels)
        ], dim=-1)


# ─────────────────────────────────────────────────────────────────────────────
# Weight conversion helpers
# ─────────────────────────────────────────────────────────────────────────────

def _pad16(x):
    """Pad to next multiple of 16, matching TCNN's alignment."""
    return ((x + 15) // 16) * 16


def _load_hash_from_flat(grid, flat):
    """Load hash grid weights from TCNN flat buffer. Returns floats consumed."""
    total = grid.hash_table.numel()
    grid.hash_table.data.copy_(flat[:total])
    return total


def _load_mlp_from_flat(seq, flat, n_input_dims, n_output_dims, width, n_hidden_layers):
    """
    Load MLP weights from TCNN FullyFusedMLP flat buffer.

    TCNN layout (row-major, no bias):
      - Layer 0:            width × pad16(n_input_dims) floats
      - Hidden layers 1..N: width × width floats each
      - Last layer:         pad16(n_output_dims) × width floats

    The first layer accepts pad16(n_input_dims) inputs (TCNN pads input with
    1.0), so its weight is loaded WITHOUT slicing the input dimension.
    The last layer's output padding IS sliced since we only need n_output_dims.
    """
    offset = 0
    linear_modules = [m for m in seq.modules() if isinstance(m, nn.Linear)]

    for i, module in enumerate(linear_modules):
        out_f, in_f = module.weight.shape

        if i == 0:
            # First layer: width × pad16(n_input_dims)
            # in_f == pad16(n_input_dims) already (from _build_mlp)
            rows = width
            cols = _pad16(n_input_dims)
        elif i == len(linear_modules) - 1:
            # Last layer: pad16(n_output_dims) × width
            rows = _pad16(n_output_dims)
            cols = width
        else:
            # Hidden layers: width × width
            rows = width
            cols = width

        n_weights = rows * cols
        w = flat[offset:offset + n_weights]
        w_mat = w.reshape(rows, cols)
        # Slice output-dim padding only (first layer: no slicing needed)
        module.weight.data.copy_(w_mat[:out_f, :in_f])
        offset += n_weights

    return offset


# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _make_grid(dim, n_levels, nf, log2_hash, base_res, scale):
    """Create hash grid matching TCNN's behavior."""
    return _HashGridTorch(
        dim=dim, n_levels=n_levels, n_features_per_level=nf,
        log2_hashmap_size=log2_hash, base_resolution=base_res,
        per_level_scale=scale,
    )


def _out_dim(grid):
    return grid.output_dim


def _build_mlp(in_dim, hidden_dim, out_dim, n_hidden):
    """MLP matching TCNN FullyFusedMLP: no bias, ReLU activation.

    TCNN pads the input dimension to a multiple of 16 and fills the padding
    with 1.0 (acting as an implicit bias).  The first layer therefore accepts
    ``pad16(in_dim)`` inputs so that it can consume the full weight matrix
    including padding columns.
    """
    padded_in = _pad16(in_dim)
    layers, cur = [], padded_in
    for _ in range(n_hidden):
        layers += [nn.Linear(cur, hidden_dim, bias=False), nn.ReLU()]
        cur = hidden_dim
    layers.append(nn.Linear(cur, out_dim, bias=False))
    return nn.Sequential(*layers)


# ─────────────────────────────────────────────────────────────────────────────
# Main Network
# ─────────────────────────────────────────────────────────────────────────────

class GNDCNetworkTorch(nn.Module):
    """
    Pure PyTorch drop-in replacement for GNDCNetwork (TCNN version).

    Identical interface:
        __init__()                          same parameters as GNDCNetwork
        forward(x)                          returns (output, offsets)
        load_weights_from_flat_buffer(buf)  loads weights from .gndc file
    """

    def __init__(
        self,
        n_output_dims,
        xy_hash_size,
        xy_base_resolution,
        xy_n_levels,
        xy_per_level_scale,
        xy_features_per_level,
        t_hash_size,
        t_base_resolution,
        t_n_levels,
        t_per_level_scale,
        t_features_per_level,
        spatial_scale,
        decoder_hidden,
        n_hidden_layers,
        spatial_dims=2,
        enable_deformation=False,
        ref_t_for_deformation=-1.0,
    ):
        super().__init__()

        self.spatial_scale = spatial_scale
        self.spatial_dims = spatial_dims
        self.enable_deformation = enable_deformation
        self.ref_t_for_deformation = ref_t_for_deformation
        self.deformation_spatial_scale = None

        self.config_params = dict(
            n_output_dims=n_output_dims,
            xy_hash_size=xy_hash_size,
            xy_base_resolution=xy_base_resolution,
            xy_n_levels=xy_n_levels,
            xy_per_level_scale=xy_per_level_scale,
            xy_features_per_level=xy_features_per_level,
            t_hash_size=t_hash_size,
            t_base_resolution=t_base_resolution,
            t_n_levels=t_n_levels,
            t_per_level_scale=t_per_level_scale,
            t_features_per_level=t_features_per_level,
            spatial_scale=spatial_scale,
            decoder_hidden=decoder_hidden,
            n_hidden_layers=n_hidden_layers,
            spatial_dims=spatial_dims,
        )

        # Store MLP dimension info for weight loading
        self._n_output_dims = n_output_dims
        self._decoder_hidden = decoder_hidden
        self._n_hidden_layers = n_hidden_layers

        # ── Optional deformation network ──────────────────────────────────
        if enable_deformation:
            self.deformation_spatial_scale = 0.1
            self.enc_deform = _make_grid(
                dim=3, n_levels=4, nf=2, log2_hash=14,
                base_res=16, scale=1.5
            )
            self.net_deform = _build_mlp(
                _out_dim(self.enc_deform), 64, 2, 1
            )
            self._deform_in_dim = _out_dim(self.enc_deform)

        # ── Static 2D spatial branch ───────────────────────────────────────
        self.enc_spatial = _make_grid(
            dim=2,
            n_levels=xy_n_levels,
            nf=xy_features_per_level,
            log2_hash=xy_hash_size,
            base_res=xy_base_resolution,
            scale=xy_per_level_scale,
        )

        # ── Dynamic 3D phenology branch ────────────────────────────────────
        self.enc_phenology = _make_grid(
            dim=3,
            n_levels=t_n_levels,
            nf=t_features_per_level,
            log2_hash=t_hash_size,
            base_res=t_base_resolution,
            scale=t_per_level_scale,
        )

        # ── MLP decoder ───────────────────────────────────────────────────
        self._decoder_in_dim = _out_dim(self.enc_spatial) + _out_dim(self.enc_phenology)
        self.decoder = _build_mlp(
            in_dim=self._decoder_in_dim,
            hidden_dim=decoder_hidden,
            out_dim=n_output_dims,
            n_hidden=n_hidden_layers,
        )

    def forward(self, x):
        """
        Args:
            x: (N, 3) — (norm_x, norm_y, norm_t)
               (N, 4) — (x, y, z, t) for 3D spatial mode
        Returns:
            output:  (N, n_output_dims)
            offsets: (N, 2) or None
        """
        x = x.float()

        coords_xy = x[:, :2].contiguous()
        coords_3d = x.contiguous() if x.shape[1] == 3 \
                    else x[:, [0, 1, 3]].contiguous()

        offsets = None

        if self.enable_deformation:
            di = coords_3d.clone()
            di[:, :2] = di[:, :2] * self.deformation_spatial_scale
            deform_feat = self.enc_deform(di)
            deform_padded = _pad16(_out_dim(self.enc_deform))
            if deform_padded > deform_feat.shape[-1]:
                deform_feat = torch.cat([deform_feat, torch.ones(
                    deform_feat.shape[0], deform_padded - deform_feat.shape[-1],
                    device=deform_feat.device, dtype=deform_feat.dtype,
                )], dim=-1)
            offsets = self.net_deform(deform_feat) * 0.001

            if self.ref_t_for_deformation >= 0.0:
                is_ref = (x[:, -1] - self.ref_t_for_deformation).abs() < 1e-6
                offsets = offsets * (~is_ref).float().unsqueeze(-1)

            coords_xy = coords_xy + offsets

        coarse = coords_3d.clone()
        coarse[:, 0] = coarse[:, 0] * self.spatial_scale
        coarse[:, 1] = coarse[:, 1] * self.spatial_scale

        features = torch.cat([
            self.enc_spatial(coords_xy),
            self.enc_phenology(coarse),
        ], dim=-1)

        # TCNN pads input to pad16(n_input_dims) with 1.0 (implicit bias)
        padded_dim = _pad16(self._decoder_in_dim)
        if padded_dim > self._decoder_in_dim:
            ones_pad = torch.ones(
                features.shape[0], padded_dim - self._decoder_in_dim,
                device=features.device, dtype=features.dtype,
            )
            features = torch.cat([features, ones_pad], dim=-1)

        out = self.decoder(features)

        return out, offsets

    def load_weights_from_flat_buffer(self, buffer):
        """
        Load weights from a TCNN flat buffer (.gndc file).

        TCNN order: enc_spatial → enc_phenology → decoder MLP
        (with deformation prepended if enabled)
        """
        buf = buffer.float().cpu()
        offset = 0

        if self.enable_deformation:
            offset += _load_hash_from_flat(self.enc_deform, buf[offset:])
            offset += _load_mlp_from_flat(
                self.net_deform, buf[offset:],
                n_input_dims=self._deform_in_dim,
                n_output_dims=2, width=64, n_hidden_layers=1
            )

        offset += _load_hash_from_flat(self.enc_spatial, buf[offset:])
        offset += _load_hash_from_flat(self.enc_phenology, buf[offset:])
        offset += _load_mlp_from_flat(
            self.decoder, buf[offset:],
            n_input_dims=self._decoder_in_dim,
            n_output_dims=self._n_output_dims,
            width=self._decoder_hidden,
            n_hidden_layers=self._n_hidden_layers,
        )


    @torch.no_grad()
    def verify_against_tcnn(self, tcnn_model, n_samples=500, tol=1e-2):
        """
        Sanity check: compare outputs with the TCNN model.
        Returns True if MAE < tol.
        """
        device = next(tcnn_model.parameters()).device
        x = torch.rand(n_samples, 3, device=device)
        out_tcnn, _ = tcnn_model(x)
        out_torch, _ = self.to(device)(x)

        mae = (out_tcnn.float() - out_torch.float()).abs().mean().item()
        max_err = (out_tcnn.float() - out_torch.float()).abs().max().item()
        passed = mae < tol

        print(f"[Verify] MAE={mae:.6f}  MaxErr={max_err:.6f}  "
              f"({'PASS' if passed else 'FAIL'})")
        return passed
