"""Frozen NumPy regression recipe; checkpoints are JSON, never executable pickle."""
import json
from pathlib import Path
import numpy as np
from .features import FEATURE_NAMES
from ._frozen_numpy import _HistogramTreeRegressor, _impute_and_scale

CONFIG=dict(max_depth=2,max_iter=100,learning_rate=0.05,min_samples_leaf=30,max_bins=255,random_state=0)
class UtilityModel:
    def fit(self,features,fallback,bespoke,labels):
        x=np.asarray(features,dtype=float)
        if x.ndim!=2 or x.shape[1]!=31: raise ValueError("EXACT_31_FEATURES_REQUIRED")
        a,b,y=[np.asarray(v) for v in (fallback,bespoke,labels)]
        if any(v.dtype!=np.bool_ or v.shape!=(len(x),) for v in (a,b,y)):
            raise ValueError("BOOLEAN_PAIRED_LABELS_REQUIRED")
        target=(b==y).astype(float)-(a==y).astype(float)
        train,_,self.transform=_impute_and_scale(x,x,scale=True)
        self.model=_HistogramTreeRegressor(**CONFIG).fit(train,target)
        return self
    def predict(self,features):
        if not hasattr(self,"model"): raise ValueError("MODEL_NOT_FITTED")
        x=np.asarray(features,dtype=float)
        if x.ndim!=2 or x.shape[1]!=31: raise ValueError("EXACT_31_FEATURES_REQUIRED")
        t=self.transform
        filled=np.where(np.isfinite(x),x,t["medians"][None,:])
        return self.model.predict((filled-t["means"][None,:])/t["scales"][None,:])
    def to_dict(self):
        if not hasattr(self,"model"): raise ValueError("MODEL_NOT_FITTED")
        state=vars(self.model).copy();state["thresholds"]=[v.tolist() for v in state["thresholds"]]
        return dict(schema="FACTCASCADE_UTILITY_JSON_V1",feature_order=FEATURE_NAMES,
                    config=CONFIG,transform={k:v.tolist() for k,v in self.transform.items()},model=state)
    @classmethod
    def from_dict(cls,data):
        if data.get("schema")!="FACTCASCADE_UTILITY_JSON_V1" or data.get("feature_order")!=FEATURE_NAMES or data.get("config")!=CONFIG:
            raise ValueError("CHECKPOINT_SCHEMA_OR_CONFIG_MISMATCH")
        result=cls();result.transform={k:np.asarray(data["transform"][k],dtype=float) for k in ["medians","means","scales"]}
        if any(v.shape!=(31,) or not np.isfinite(v).all() for v in result.transform.values()) or np.any(result.transform["scales"]<=0):
            raise ValueError("CHECKPOINT_TRANSFORM_INVALID")
        model=data["model"]
        if any(model[k]!=v for k,v in CONFIG.items()) or len(model["thresholds"])!=31 or len(model["trees"])!=100:
            raise ValueError("CHECKPOINT_MODEL_INVALID")
        result.model=_HistogramTreeRegressor(**CONFIG)
        result.model.initial_value=float(model["initial_value"])
        if not np.isfinite(result.model.initial_value): raise ValueError("CHECKPOINT_NONFINITE")
        result.model.thresholds=[np.asarray(v,dtype=float) for v in model["thresholds"]]
        if any(v.ndim!=1 or not np.isfinite(v).all() or np.any(np.diff(v)<=0) for v in result.model.thresholds):
            raise ValueError("CHECKPOINT_THRESHOLDS_INVALID")
        def tree(node,depth=0):
            if set(node)=={"value"} and np.isfinite(node["value"]): return
            if set(node)!={"feature","threshold_bin","left","right"} or depth>=2:
                raise ValueError("CHECKPOINT_TREE_INVALID")
            f,t=node["feature"],node["threshold_bin"]
            if type(f) is not int or not 0<=f<31 or type(t) is not int or not 0<=t<len(result.model.thresholds[f]):
                raise ValueError("CHECKPOINT_SPLIT_INVALID")
            tree(node["left"],depth+1);tree(node["right"],depth+1)
        for n in model["trees"]: tree(n)
        result.model.trees=model["trees"]
        return result
    @classmethod
    def load(cls,path): return cls.from_dict(json.loads(Path(path).read_bytes()))
    def save(self,path): Path(path).write_text(json.dumps(self.to_dict(),allow_nan=False)+"\n",encoding="utf-8")