"""Ace-Mini stream model: transient signed liquid distance and inlet direction."""
import numpy as np
import torch
from torch import nn
from dense_surrogate import DenseNet,TASKS as BASE_TASKS
from step_surrogate import features as base_features

TASKS={**BASE_TASKS,'filling':BASE_TASKS['filling']+['filling__free_surface_distance']}

def features(spec,data):
    cloud,q=base_features(spec,data);g=spec['geometry'];points=np.asarray(data['points_m'])*1000
    attachments=np.asarray(g['attachment_points_mm']);nearest=np.linalg.norm(points[:,None]-attachments[None],axis=-1).argmin(axis=1)
    inlets=np.asarray(g.get('inlet_centers_mm',g['attachment_points_mm']))
    directions=np.asarray(g.get('inlet_directions',[[0,0,1]]*len(inlets)),float)
    directions/=np.maximum(np.linalg.norm(directions,axis=1,keepdims=True),1e-9)
    delta=(points-inlets[nearest])/100
    along=(delta*directions[nearest]).sum(1,keepdims=True)
    radial=np.linalg.norm(delta-along*directions[nearest],axis=1,keepdims=True)
    return cloud,np.c_[q,directions[nearest],delta,along,radial].astype('float32')

class StreamNet(DenseNet):
    def __init__(self,query_dim):
        super().__init__(query_dim)
        self.heads['filling'][-1]=nn.Linear(96,7)

def initialize_from_dense(model,checkpoint):
    """Transfer common physics weights while leaving new features/head learnable."""
    old=checkpoint['state_dict'];new=model.state_dict();oldq=checkpoint['query_dim'];newq=model.body[0].in_features-217
    for key,value in old.items():
        if key in new and value.shape==new[key].shape:new[key]=value.clone()
    w=new['body.0.weight'];w[:,:oldq]=old['body.0.weight'][:,:oldq]
    w[:,oldq:newq]=0;w[:,newq:]=old['body.0.weight'][:,oldq:]
    new['heads.filling.2.weight'][:6]=old['heads.filling.2.weight']
    new['heads.filling.2.bias'][:6]=old['heads.filling.2.bias']
    model.load_state_dict(new)