"""Geometry-conditioned field network with spatial Fourier features."""
import numpy as np
import torch
from torch import nn
from step_surrogate import features,TASKS as ORIGINAL_TASKS
TASKS={**ORIGINAL_TASKS,'mechanical':ORIGINAL_TASKS['mechanical'][:4]}
class DenseNet(nn.Module):
 def __init__(self,query_dim):
  super().__init__();self.encoder=nn.Sequential(nn.Linear(6,64),nn.SiLU(),nn.Linear(64,96),nn.SiLU());self.register_buffer('frequencies',torch.tensor([1.,2.,4.,8.])*np.pi)
  self.body=nn.Sequential(nn.Linear(query_dim+192+24+1,192),nn.SiLU(),nn.Linear(192,192),nn.SiLU())
  self.heads=nn.ModuleDict({k:nn.Sequential(nn.Linear(192,96),nn.SiLU(),nn.Linear(96,len(v))) for k,v in TASKS.items()})
  for task in ('filling','cooling'):
   self.register_buffer(task+'_sf_mean',torch.tensor(0.));self.register_buffer(task+'_sf_std',torch.tensor(1.))
 def encode(self,cloud):
  a=self.encoder(cloud);return torch.cat([a.max(0).values,a.mean(0)])
 def forward(self,cloud,query,time_s,task,encoded=None):
  code=self.encode(cloud) if encoded is None else encoded;t=torch.log1p(time_s.clamp(min=0)).reshape(-1,1)/10
  phase=(query[:,:3,None]*self.frequencies).flatten(1);pos=torch.cat([phase.sin(),phase.cos()],1)
  out=self.heads[task](self.body(torch.cat([query,code.expand(len(query),-1),pos,t],1)))
  if task in ('filling','cooling'):
   out=out.clone();out[:,1]=(out[:,1].sigmoid()-getattr(self,task+'_sf_mean'))/getattr(self,task+'_sf_std')
  return out