import json,torch
from pathlib import Path
from torch import nn
import yaml
def cfg(r):return yaml.safe_load((Path(r)/'conf/config.yaml').read_text())
def sample(i,days=730):
 t=torch.arange(days).float();p=(.5+.5*torch.sin(t/20+i)).clamp_min(0);temp=10+12*torch.sin(t/58+i);pet=(2+.1*temp).clamp_min(0);f=torch.stack((p,temp,pet),1);a=torch.linspace(0,1,35)+i*.01;return f,a
class DeltaHBV(nn.Module):
 def __init__(self,static_features=35,hidden=16,components=16):super().__init__();self.components=components;self.rnn=nn.LSTM(3+static_features,hidden,batch_first=True);self.head=nn.Linear(hidden,components*4);self.model_config={'static_features':static_features,'hidden':hidden,'components':components}
 def forward(self,forcing,attrs):
  b,t,_=forcing.shape;x=torch.cat((forcing,attrs[:,None].expand(-1,t,-1)),-1);h,_=self.rnn(x);par=torch.sigmoid(self.head(h)).reshape(b,t,self.components,4);snow=torch.zeros(b,self.components);soil=torch.ones_like(snow);upper=torch.zeros_like(snow);lower=torch.zeros_like(snow);flows=[];ets=[]
  for k in range(t):
   p,temp,pet=forcing[:,k].unbind(-1);rain=p[:,None]*(temp[:,None]>0);snow=snow+p[:,None]-rain;melt=torch.minimum(snow,torch.relu(temp[:,None])*.02);snow=snow-melt;water=rain+melt;beta=1+4*par[:,k,:,0];gamma=.5+2*par[:,k,:,1];fc=1+4*par[:,k,:,2];recharge=torch.minimum((soil/fc).clamp(0,1).pow(beta)*water,water);soil=soil+water-recharge;et=torch.minimum((soil/fc).clamp(0,1).pow(gamma)*pet[:,None],soil);soil=soil-et;upper=upper+recharge;q0=upper*.2;upper=upper-q0;lower=lower+upper*.05;upper=upper*.95;q2=lower*.02;lower=lower-q2;flows.append((q0+q2).mean(1));ets.append(et.mean(1))
  return torch.stack(flows,1),torch.stack(ets,1)
def nse(p,t):return 1-((p-t)**2).sum()/((t-t.mean())**2).sum().clamp_min(1e-8)
def write(p,o):p=Path(p);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(o,indent=2)+'\n')