"""Fresh geometry-conditioned point network; no weights from the plate pilot."""
import numpy as np
import torch
from torch import nn
from scipy.spatial import cKDTree
TASKS={
 'maps':['filling__fillingtime','thermomechanical__solidificationtime','demolding__microporosity','demolding__totalshrinkage'],
 'filling':['filling__temperature','filling__solidfraction','filling__velocityx','filling__velocityy','filling__velocityz','filling__pressures'],
 'cooling':['thermomechanical__temperature','thermomechanical__solidfraction'],
 'mechanical':['demolding__displacementx','demolding__displacementy','demolding__displacementz','demolding__displacementmag','demolding__vonmisesstresstensor']}
def features(spec,data):
    g=spec['geometry'];points=np.asarray(data['points_m'])*1000;surface=np.asarray(data['encoder_surface_mm']);normal=np.asarray(data['encoder_normals'])
    lo=surface.min(0);hi=surface.max(0);center=(lo+hi)/2;scale=max(float(np.max(hi-lo)),1.)
    cloud=np.c_[(surface-center)/scale,normal].astype('float32')
    gates=np.array(g['attachment_points_mm']);delta=points[:,None,:]-gates[None,:,:]
    distance=np.linalg.norm(delta,axis=2);order=np.argsort(distance,axis=1)
    nearest=np.take_along_axis(delta,order[:,:1,None],axis=1)[:,0,:]/scale
    near_d=distance.min(1,keepdims=True)/scale;far_d=distance.max(1,keepdims=True)/scale
    depth=cKDTree(surface).query(points)[0][:,None]/scale
    constants=np.tile(np.r_[(hi-lo)/100,scale/100,g['gate_radius_mm']/10,g['gate_count']/2,spec['process']['inlet_velocity_m_s']],(len(points),1))
    query=np.c_[(points-center)/scale,points/100,nearest,near_d,far_d,depth,constants].astype('float32')
    return cloud,query
class StepNet(nn.Module):
    def __init__(self,query_dim):
        super().__init__()
        self.encoder=nn.Sequential(nn.Linear(6,32),nn.SiLU(),nn.Linear(32,48),nn.SiLU())
        self.body=nn.Sequential(nn.Linear(query_dim+96+1,96),nn.SiLU(),nn.Linear(96,96),nn.SiLU(),nn.Linear(96,64),nn.SiLU())
        self.heads=nn.ModuleDict({k:nn.Linear(64,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
        time=torch.log1p(torch.clamp(time_s,min=0)).reshape(-1,1)/10
        out=self.heads[task](self.body(torch.cat([query,code.expand(len(query),-1),time],1)))
        if task in ('filling','cooling'):
            out=out.clone();out[:,1]=(torch.sigmoid(out[:,1])-getattr(self,task+'_sf_mean'))/getattr(self,task+'_sf_std')
        return out