"""CLM5 GPP/LHF EOF emulators and bounded parameter search."""
import json
from pathlib import Path
import numpy as np
import torch
from torch import nn
import yaml

PARAMETERS=("medlynslope","dleaf","kmax","fff","dint","baseflow_scalar")
def load_config(root):return yaml.safe_load((Path(root)/"conf/config.yaml").read_text())
class TargetEmulator(nn.Module):
 def __init__(self,hidden_dim=12):super().__init__();self.net=nn.Sequential(nn.Linear(6,hidden_dim),nn.ReLU(),nn.Linear(hidden_dim,hidden_dim),nn.Tanh(),nn.Linear(hidden_dim,3))
 def forward(self,x):return self.net(x)
class CLM5Emulator(nn.Module):
 def __init__(self,hidden_dim=12):super().__init__();self.gpp=TargetEmulator(hidden_dim);self.lhf=TargetEmulator(hidden_dim);self.model_config={"hidden_dim":hidden_dim}
 def forward(self,x):
  if x.ndim!=2 or x.shape[1]!=6:raise ValueError("expected [B,6] normalized parameters")
  return torch.stack((self.gpp(x),self.lhf(x)),1)
def reconstruct(pc,eof,mean):return np.einsum("btm,tmg->btg",pc,eof)+mean[None]
def bounded_search(model,observed_pc,scales,weights,candidates=4096,seed=0):
 rng=np.random.default_rng(seed);p=torch.from_numpy(rng.random((candidates,6)).astype("float32"));
 with torch.no_grad():pred=model(p).numpy()
 cost=(((pred-observed_pc[None])/scales[None])**2*weights[None]).sum((1,2));i=int(cost.argmin());return p[i].numpy(),float(cost[i])
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")