import json,torch
from pathlib import Path
from torch import nn
import yaml
AEROSOL=('AOD','TSAOD','SUAOD','DUAOD','BCAOD','OCAOD','SSAOD','SUSMC','DUSMC','BCSMC','OCSMC','SSSMC')
def cfg(r):return yaml.safe_load((Path(r)/'conf/config.yaml').read_text())
def state(i,h=0):
 y,x=torch.meshgrid(torch.linspace(-1,1,16),torch.linspace(-1,1,16),indexing='ij');c=torch.arange(54)[:,None,None];return (torch.sin((c%8+1)*x+.02*(i+h))*torch.cos((c%6+1)*y)).float()
class Relay(nn.Module):
 def __init__(self,c,h,heads):super().__init__();self.e=nn.Conv2d(c,h,2,2);self.a=nn.MultiheadAttention(h,heads,batch_first=True);self.d=nn.ConvTranspose2d(h,c,2,2)
 def forward(self,x):z=self.e(x);b,c,y,w=z.shape;t=z.permute(0,2,3,1).reshape(b,y*w,c);t=self.a(t,t,t,need_weights=False)[0];return x+self.d(t.reshape(b,y,w,c).permute(0,3,1,2))
class AIGAMFS(nn.Module):
 def __init__(self,channels=54,hidden=24,heads=4):super().__init__();self.models=nn.ModuleDict({str(h):Relay(channels,hidden,heads) for h in (3,6,9,12)});self.model_config={'channels':channels,'hidden':hidden,'heads':heads}
 def forward(self,x,lead):return self.models[str(lead)](x)
def write(p,o):p=Path(p);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(o,indent=2)+'\n')