import json,torch
from pathlib import Path
from torch import nn
import yaml
FIELDS=('Z1000','Z500','Z250','THICK300_700','T2M','T850','TCWV','WS10','OLR')
def cfg(r):return yaml.safe_load((Path(r)/'conf/config.yaml').read_text())
def state(i,s=0):
 y,x=torch.meshgrid(torch.linspace(-1,1,16),torch.linspace(-1,1,16),indexing='ij');c=torch.arange(9)[:,None,None];a=torch.sin((c+1)*x+.04*(i+s))*torch.cos((c%4+1)*y);sst=torch.cos(x+.01*(i+s))*torch.cos(y);return a.float(),sst[None].float()
class Net(nn.Module):
 def __init__(self,ci,co,h):super().__init__();self.n=nn.Sequential(nn.Conv2d(ci,h,3,padding=1),nn.GELU(),nn.Conv2d(h,h,3,padding=1),nn.GELU(),nn.Conv2d(h,co,1))
 def forward(self,x):return self.n(x)
class DLESyM(nn.Module):
 def __init__(self,atmosphere_channels=9,hidden=16):super().__init__();self.atm=Net(10,9,hidden);self.ocean=Net(4,1,hidden//2);self.precip=Net(9,1,hidden//2);self.model_config={'atmosphere_channels':atmosphere_channels,'hidden':hidden}
 def atmosphere(self,a,sst):return a+self.atm(torch.cat((a,sst),1))
 def cycle(self,a,sst):
  seq=[]
  for _ in range(16):a=self.atmosphere(a,sst);seq.append(a)
  f=torch.stack(seq);forcing=torch.stack((f[:,:,7].mean(0),f[:,:,0].mean(0),f[:,:,8].mean(0)),1);sst=sst+self.ocean(torch.cat((sst,forcing),1));return a,sst,self.precip(a)
 def forward(self,a,sst):return self.cycle(a,sst)
def write(p,o):p=Path(p);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(o,indent=2)+'\n')