import json
from pathlib import Path
import torch
from torch import nn
import yaml
def cfg(r):return yaml.safe_load((Path(r)/"conf/config.yaml").read_text())
def sample(i,s=0):
 y,x=torch.meshgrid(torch.linspace(-1,1,16),torch.linspace(-1,1,16),indexing="ij");c=torch.arange(20)[:,None,None];b=torch.sin((c%5+1)*x+.1*(i+s))*torch.cos((c%4+1)*y);m=(((torch.arange(16)[None,:]+i+s)%4==0).float().expand(16,-1))[None].expand(20,-1,-1);o=b*m;return b.float(),o.float(),m.float()
class FuXiWeather(nn.Module):
 def __init__(self,channels=20,hidden=24):super().__init__();self.obs=nn.Conv2d(channels*2,hidden,1);self.bg=nn.Conv2d(channels,hidden,1);self.refine=nn.Sequential(nn.Conv2d(hidden*2,hidden,3,padding=1),nn.GELU(),nn.Conv2d(hidden,channels,1));self.short=nn.Conv2d(channels,channels,3,padding=1);self.medium=nn.Conv2d(channels,channels,3,padding=1);self.model_config={"channels":channels,"hidden":hidden}
 def analysis(self,b,o,m):return b+self.refine(torch.cat((self.bg(b),self.obs(torch.cat((o,m),1))),1))
 def forecast(self,x,medium=False):return x+(self.medium(x) if medium else self.short(x))
 def forward(self,b,o,m,medium=False):
  a=self.analysis(b,o,m);return a,self.forecast(a,medium)
def write(p,o):p=Path(p);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(o,indent=2)+"\n")