"""Learned stroke cost for HanziFlow.

Needs: numpy, onnxruntime, and the three files from the training notebook:
stroke_encoder.onnx, template_embeddings.npz, calibration.json
Strokes must be in the Make Me a Hanzi frame (1024 box, y up) and resampled to 32 points,
exactly as in the matcher notebook (prep()).
"""
import json
import numpy as np
import onnxruntime as ort


class LearnedCost:
    def __init__(self, onnx_path="stroke_encoder.onnx", npz_path="template_embeddings.npz", calib_path="calibration.json"):
        self.sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
        z = np.load(npz_path)
        self.chars = [str(c) for c in z["chars"]]
        self.templates = {c: z[f"emb_{i}"] for i, c in enumerate(self.chars)}   # (n_strokes, 64) per character
        cal = json.load(open(calib_path))
        self.scale, self.tau, self.good = cal["scale"], cal["tau"], cal["good"]
        self.alpha = cal.get("alpha_recommended", 0.5)

    @staticmethod
    def features(S):
        # (n, 32, 2) points -> (n, 32, 6); identical to stroke_features() in the training notebook
        S = np.asarray(S, np.float32)
        x = (S[..., 0] - 512) / 512
        y = (S[..., 1] - 388) / 512
        g = np.gradient(S, axis=-2)
        nrm = np.linalg.norm(g, axis=-1) + 1e-6
        rx = (S[..., 0] - S[..., :1, 0]) / 512
        ry = (S[..., 1] - S[..., :1, 1]) / 512
        return np.stack([x, y, g[..., 0] / nrm, g[..., 1] / nrm, rx, ry], axis=-1).astype(np.float32)

    def embed(self, S):
        return self.sess.run(None, {"stroke_features": self.features(S)})[0]

    def costs(self, char, U):
        """Learned cost of every drawn stroke (rows) against every template stroke of `char` (columns).
        Direction-agnostic: a backwards stroke is flagged by the geometric part of the matcher, not here."""
        te = self.templates[char]
        fwd = np.linalg.norm(self.embed(U)[:, None, :] - te[None], axis=-1)
        bwd = np.linalg.norm(self.embed(np.ascontiguousarray(U[:, ::-1]))[:, None, :] - te[None], axis=-1)
        return np.minimum(fwd, bwd) / self.scale

    def cost_fn(self, char, geometric_cost_matrix, alpha=None):
        """Drop-in `cost_fn` for align_full(): blends the geometric cost with the learned one.
        alpha = 0 -> geometric only, 1 -> learned only, 0.5 -> half and half (recommended)."""
        a = self.alpha if alpha is None else alpha
        if a == 0:
            return None

        def fn(U, T):
            C_geo, rev = geometric_cost_matrix(U, T)
            return (1 - a) * C_geo + a * self.costs(char, U), rev
        return fn