from pathlib import Path
import sys,os,torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
from model.sfno_bvmc import *
c=load_config(ROOT);rank=int(os.environ.get("RANK",0));world=int(os.environ.get("WORLD_SIZE",1));distributed=world>1
if distributed:dist.init_process_group("gloo")
torch.manual_seed(c["seed"]);torch.set_num_threads(2);base=CompactSFNO(**c["model"]);m=DDP(base) if distributed else base;opt=torch.optim.Adam(m.parameters(),lr=c["train"]["learning_rate"]);losses=[]
indices=list(range(rank,len(c["data"]["tile_origins"]),world))
while len(indices)<math.ceil(len(c["data"]["tile_origins"])/world):indices.append(indices[-1])
for _ in range(c["train"]["epochs"]):
 for i in indices:
  o=c["data"]["tile_origins"][i]
  x=state(o,c["data"]["tile_size"],i,0)[None];y=state(o,c["data"]["tile_size"],i,1,74)[None];loss=F.mse_loss(m(x),y);opt.zero_grad();loss.backward();opt.step();losses.append(float(loss))
total=torch.tensor([sum(losses),len(losses)],dtype=torch.float64)
if distributed:dist.all_reduce(total)
path=ROOT/c["paths"]["checkpoint"]
if rank==0:
 path.parent.mkdir(parents=True,exist_ok=True);torch.save({"model":base.state_dict(),"model_config":c["model"],"checkpoint_count_protocol":c["train"]["checkpoint_count"],"format_version":c["data"]["format_version"]},path);write_json(ROOT/c["paths"]["training_metrics"],{"history":[{"epoch":1,"loss":float(total[0]/total[1])}],"stored_checkpoint_files":1,"checkpoint_count_protocol":29,"world_size":world});print(path)
if distributed:dist.destroy_process_group()