"""Independent NOTS experiment model; the general FNO is neuraloperator 1.0.2.

No source from the authors' NOTS repository is included here. Paper Algorithm 2
requires a newly initialized, fully trained FNO at every posterior draw.
"""
import math

import torch
from torch import nn
from neuralop.models import FNO


class NOTS(nn.Module):
    """Scalar permeability [B,1,H,W] -> pressure [B,1,H,W]."""

    def __init__(self, **model_config):
        super().__init__()
        self.model_config = dict(model_config)
        self.fno = FNO(**model_config)
        self.reset_posterior_parameters()

    @torch.no_grad()
    def reset_posterior_parameters(self):
        """He Gaussian weights, zero biases, including complex spectral weights.

        Complex normal_ uses the requested total variance across real/imaginary
        parts, and the regularizer uses the modulus squared of those weights.
        """
        for name, parameter in self.named_parameters():
            if 'bias' in name:
                parameter.zero_()
            else:
                fan_in = parameter.shape[1] * math.prod(parameter.shape[2:]) if parameter.ndim > 1 else parameter.numel()
                parameter.normal_(0.0, math.sqrt(2.0 / max(1, fan_in)))

    def forward(self, permeability):
        if permeability.ndim != 4 or permeability.shape[1] != 1:
            raise ValueError('NOTS expects permeability [batch,1,height,width]')
        return self.fno(permeability)

    def l2_penalty(self):
        return sum(p.abs().square().sum() for p in self.parameters() if p.requires_grad)