# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Utterance-convention decoder: pure-numpy predictor for exported gradient-boosted trees, plus the decision rule.

The trees are scikit-learn HistGradientBoostingRegressor models exported by tools/export_decoder.py. Prediction
replicates sklearn exactly: float64 baseline plus each tree's leaf value, added in iteration order; numerical splits
go left when x <= threshold, missing values follow the stored missing direction.
Decision rule (as frozen in the research pipeline): bridge a same-speaker gap when the gap model predicts a negative
error change per second; drop a short segment when the segment model predicts a negative error change.
"""
import json

import numpy as np

from .geometry import union


class TreeEnsemble:
    def __init__(self, blob):
        self.features = blob["features"]
        self.baseline = float(blob["baseline"])
        self.trees = [{k: np.asarray(v) for k, v in t.items()} for t in blob["trees"]]

    def predict(self, rows):
        if not rows:
            return np.zeros(0)
        X = np.array([[r[f] for f in self.features] for r in rows], dtype=np.float64)
        out = np.full(len(X), self.baseline, dtype=np.float64)
        for t in self.trees:
            feat, thr, mleft = t["feature"], t["threshold"], t["missing_left"]
            left, right, leaf, value = t["left"], t["right"], t["leaf"], t["value"]
            node = np.zeros(len(X), dtype=np.int64)
            for _ in range(len(feat)):
                active = ~leaf[node]
                if not active.any():
                    break
                n = node[active]
                x = X[np.flatnonzero(active), feat[n]]
                go_left = np.where(np.isnan(x), mleft[n], x <= thr[n])
                node[active] = np.where(go_left, left[n], right[n])
            out += value[node]
        return out


def load_decoder(path):
    blob = json.loads(open(path).read())
    return TreeEnsemble(blob["gap"]), TreeEnsemble(blob["seg"])


def apply(hyp, gap_rows, gap_pred, seg_rows, seg_pred, tau_gap=0.0, tau_seg=0.0):
    """hyp: spk -> unbridged intervals. Bridge gaps with gap_pred > tau_gap; drop segments with seg_pred < tau_seg
    unless they were merged by a bridge. (Research: convention_decoder.apply without the unused `moves` option.)"""
    fill = {(r["spk"], r["j"]) for r, p in zip(gap_rows, gap_pred) if p > tau_gap}
    drop = {(r["spk"], r["j"]) for r, p in zip(seg_rows, seg_pred) if p < tau_seg}
    out = {}
    for k, v in hyp.items():
        iv = union([tuple(x) for x in v])
        keep = []
        for j, (s, e) in enumerate(iv):
            merged = (k, j) in fill or (k, j + 1) in fill
            keep.append(not ((k, j) in drop and not merged))
        res = []
        for j, (s, e) in enumerate(iv):
            if j > 0 and (k, j) in fill and res and keep[j] and keep[j - 1]:
                res[-1][1] = max(res[-1][1], e)
            elif keep[j]:
                res.append([s, e])
        if res:
            out[k] = [(a, b) for a, b in res]
    return out