"""
GNDC Network for Geospatial Data Compression.

This module implements a hybrid neural network architecture combining:
- Static High-Res 2D branch: Captures spatial details (boundaries, features)
- Dynamic Coarse 3D branch: Captures temporal variations (phenology)
- Optional Deformation Network: Geometric correction for misaligned data
"""

import torch
import torch.nn as nn
import math

try:
    import tinycudann as tcnn
except ImportError:
    tcnn = None


class GNDCNetwork(nn.Module):
    """
    GNDC Network for geospatial time-series compression.
    
    Combines 2D spatial encoding with 3D spatio-temporal encoding
    to achieve high compression ratios while preserving spatial details
    and temporal dynamics.
    """
    
    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__()

        if tcnn is None:
            raise ImportError(
                "GNDCNetwork requires tinycudann. Install it with:\n"
                "  pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch\n"
                "Or use backend='torch' for pure PyTorch inference."
            )

        self.spatial_dims = spatial_dims
        self.spatial_scale = spatial_scale
        self.enable_deformation = enable_deformation
        self.deformation_spatial_scale = None
        self.ref_t_for_deformation = ref_t_for_deformation
        
        self.config_params = {
            "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
        }

        if self.enable_deformation:
            self.deformation_spatial_scale = 0.1
            self.enc_deform = tcnn.Encoding(
                n_input_dims=3,
                encoding_config={
                    "otype": "HashGrid",
                    "n_levels": 4,
                    "n_features_per_level": 2,
                    "log2_hashmap_size": 14,
                    "base_resolution": 16,
                    "per_level_scale": 1.5
                }
            )
            self.net_deform = tcnn.Network(
                n_input_dims=self.enc_deform.n_output_dims,
                n_output_dims=2, 
                network_config={
                    "otype": "FullyFusedMLP",
                    "activation": "ReLU",
                    "output_activation": "None",
                    "n_neurons": 64,
                    "n_hidden_layers": 1
                }
            )
        
        self.enc_spatial = tcnn.Encoding(
            n_input_dims=2,
            encoding_config={
                "otype": "HashGrid",
                "n_levels": xy_n_levels,
                "n_features_per_level": xy_features_per_level,
                "log2_hashmap_size": xy_hash_size,
                "base_resolution": xy_base_resolution,
                "per_level_scale": xy_per_level_scale
            }
        )
        dim_spatial = self.enc_spatial.n_output_dims

        self.enc_phenology = tcnn.Encoding(
            n_input_dims=3,
            encoding_config={
                "otype": "HashGrid",
                "n_levels": t_n_levels,
                "n_features_per_level": t_features_per_level,
                "log2_hashmap_size": t_hash_size,
                "base_resolution": t_base_resolution,
                "per_level_scale": t_per_level_scale     
            }
        )
        dim_phenology = self.enc_phenology.n_output_dims

        self.decoder = tcnn.Network(
            n_input_dims=dim_spatial + dim_phenology,
            n_output_dims=n_output_dims,
            network_config={
                "otype": "FullyFusedMLP",
                "activation": "ReLU",
                "output_activation": "None",
                "n_neurons": decoder_hidden,
                "n_hidden_layers": n_hidden_layers
            }
        )
        

    def forward(self, x):
        """
        Forward pass.
        
        Args:
            x: [Batch, 3] -> (normalized_x, normalized_y, normalized_t)
               or [Batch, 4] -> (x, y, z, t)
        
        Returns:
            output: Predicted values [Batch, n_output_dims]
            offsets: Deformation offsets [Batch, 2] or None
        """
        x = x.float()
        
        if x.shape[1] == 3:
            coords_xy = x[:, :2].contiguous()
            coords_3d_raw = x.contiguous()
        else:
            coords_xy = x[:, :2].contiguous()
            coords_3d_raw = x[:, [0, 1, 3]].contiguous()
        
        offsets = None
        
        if self.enable_deformation:
            deform_input = coords_3d_raw.clone()
            deform_input[:, :2] *= self.deformation_spatial_scale
            
            with torch.amp.autocast('cuda', enabled=False):
                feat_def = self.enc_deform(deform_input)
                offsets = self.net_deform(feat_def).float() * 0.001
            
            if self.ref_t_for_deformation >= 0.0:
                epsilon = 1e-6
                t_input = x[:, -1]
                is_ref = torch.abs(t_input - self.ref_t_for_deformation) < epsilon
                anchor_mask = (~is_ref).float().unsqueeze(-1)
                offsets = offsets * anchor_mask

            coords_xy = coords_xy + offsets

        coords_coarse_3d = coords_3d_raw.clone()
        coords_coarse_3d[:, 0] *= self.spatial_scale
        coords_coarse_3d[:, 1] *= self.spatial_scale
        
        with torch.amp.autocast('cuda', enabled=False):
            feat_static = self.enc_spatial(coords_xy) 
            feat_dynamic = self.enc_phenology(coords_coarse_3d)
            features = torch.cat([feat_static, feat_dynamic], dim=-1)
            output = self.decoder(features)
        
        return output.float(), offsets

    @property
    def params(self):
        """
        Get all parameters from submodules (for optimizer).
        """
        params_list = [self.enc_spatial.params, self.enc_phenology.params, self.decoder.params]
        if self.enable_deformation:
            params_list.extend([self.enc_deform.params, self.net_deform.params])
        return torch.cat(params_list)

    def load_weights_from_flat_buffer(self, buffer):
        """
        Load weights from a flat buffer.
        Must follow the exact order of params concatenation.
        """
        offset = 0
        
        def load_slice(module):
            nonlocal offset
            num_params = module.params.numel()
            if offset + num_params > buffer.numel():
                raise RuntimeError(f"Buffer too small: need {offset+num_params}, got {buffer.numel()}")
            
            chunk = buffer[offset : offset + num_params]
            module.params.data.copy_(chunk)
            offset += num_params
        
        load_slice(self.enc_spatial)
        load_slice(self.enc_phenology)
        load_slice(self.decoder)

        print(f"Weights loaded successfully. Total params: {offset}")
