"""Compact physical-state coupled atmosphere-ocean emulator."""
import json
from pathlib import Path
import numpy as np
import torch
from torch import nn
import yaml

ATM3D=("T","qT","U","V");ATM2D=("Ts","ps");DIAG=("RSW","OLR","USWsfc","ULWsfc","DSWsfc","DLWsfc","LHF","SHF","P","dTWP_adv","tau_u","tau_v")
OCEAN3D=("thetao","so","uo","vo");OCEAN2D=("SIC","HI","SST","ZOS")
def load_config(root):return yaml.safe_load((Path(root)/"conf/config.yaml").read_text())
def synthetic_state(origin,size,i,step=0):
 y,x=torch.meshgrid(torch.arange(origin[0],origin[0]+size),torch.arange(origin[1],origin[1]+size),indexing="ij");lat=torch.deg2rad(90-(y+.5));lon=torch.deg2rad(-180+(x+.5));a=torch.arange(46)[:,None,None];o=torch.arange(80)[:,None,None];atm=torch.sin((a%6+1)*lat+.03*a+.05*(i+step))*torch.cos(lon);oce=torch.cos((o%5+1)*lat-.01*o+.02*(i+step))*torch.sin(lon);return atm.float(),oce.float()
class ResidualEmulator(nn.Module):
 def __init__(self,cin,cout,h):super().__init__();self.net=nn.Sequential(nn.Conv2d(cin,h,1),nn.GELU(),nn.Conv2d(h,h,3,padding=1),nn.GELU(),nn.Conv2d(h,cout,1))
 def forward(self,x):return self.net(x)
class SamudrACE(nn.Module):
 def __init__(self,hidden_dim=16):super().__init__();self.atmosphere=ResidualEmulator(48,46,hidden_dim);self.ocean=ResidualEmulator(92,80,hidden_dim);self.model_config={"hidden_dim":hidden_dim}
 def atmosphere_step(self,atm,ocean):return atm+self.atmosphere(torch.cat((atm,ocean[:,[76,0]]),1))
 def ocean_step(self,ocean,flux):return ocean+self.ocean(torch.cat((ocean,flux),1))
 def forward(self,atm,ocean,n=20):return self.coupled_step(atm,ocean,n)
 def coupled_step(self,atm,ocean,n=20):
  flux=[]
  for _ in range(n):atm=self.atmosphere_step(atm,ocean);flux.append(atm[:,-12:])
  ocean=self.ocean_step(ocean,torch.stack(flux).mean(0));return atm,ocean
def write_json(path,obj):path=Path(path);path.parent.mkdir(parents=True,exist_ok=True);path.write_text(json.dumps(obj,indent=2)+"\n")