"""Convert the TDT prediction net (KDA) and joint to CoreML for on-device decoding (fp32, meant for the CPU).

  decoder.mlpackage    token [1] i32, rec [2,1,20480], conv [2,1,20480] -> g [1,640] (joint.pred applied), rec_out, conv_out
  joint_enc.mlpackage  encoded [1,1024,188] -> f [1,188,640] (joint.enc), once per 15 s window
  joint.mlpackage      f [1,640], g [1,640] -> token_id [1] i32 (argmax of 3073), duration_id [1] i32 (argmax of 5)
The decoder only runs after a non-blank token; the joint runs every step.
"""
import os, sys
# HAAS_EXPORT_DIR: dir containing the haas_onnx package (scripts/ of Sanami/haas-ja-onnx)
# HAAS_NEMO_DIR: the extracted Haas-ja-v1.0.nemo (model_config.yaml, model_weights.ckpt)
sys.path.insert(0, os.environ.get("HAAS_EXPORT_DIR", "export"))
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import numpy as np, torch, coremltools as ct
from haas_onnx.load import load_model
import cmlpatch  # noqa: F401

OUT = os.path.dirname(os.path.abspath(__file__))
N_DUR = 5


class Dec(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.embed = model.decoder.prediction["embed"]
        self.rnn = model.decoder.prediction["dec_rnn"]
        self.pred = model.joint.pred

    def forward(self, token, rec, conv):
        y = self.embed(token.long()).unsqueeze(0)  # [U=1, B=1, H]
        g, (rec, conv) = self.rnn(y, (rec, conv))
        return self.pred(g[0]), rec, conv


class JointEnc(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.enc = model.joint.enc

    def forward(self, encoded):
        return self.enc(encoded.transpose(1, 2))


class Joint(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.net = model.joint.joint_net

    def forward(self, f, g):
        logits = self.net(f + g)  # joint_net = activation, dropout, linear
        tok = logits[:, :-N_DUR].argmax(-1).to(torch.int32)
        dur = logits[:, -N_DUR:].argmax(-1).to(torch.int32)
        return tok, dur


def conv(module, args, ins, outs, name):
    with torch.no_grad():
        tr = torch.jit.trace(module.eval(), args, check_trace=False)
    ml = ct.convert(tr, inputs=ins, outputs=[ct.TensorType(o) for o in outs], convert_to="mlprogram",
                    compute_precision=ct.precision.FLOAT32, minimum_deployment_target=ct.target.iOS17)
    ml.save(f"{OUT}/{name}.mlpackage")
    return ml


if __name__ == "__main__":
    model = load_model(os.environ.get("HAAS_NEMO_DIR", "nemo_x"))
    print(model.joint.joint_net)
    W = model.decoder.prediction["dec_rnn"].width
    tok = torch.tensor([42], dtype=torch.int32)
    rec, cv = torch.randn(2, 1, W) * 0.1, torch.randn(2, 1, W) * 0.1
    d = Dec(model)
    m = conv(d, (tok, rec, cv), [ct.TensorType("token", shape=(1,), dtype=np.int32),
                                 ct.TensorType("rec", shape=(2, 1, W), dtype=np.float32),
                                 ct.TensorType("conv", shape=(2, 1, W), dtype=np.float32)],
             ["g", "rec_out", "conv_out"], "decoder")
    with torch.no_grad():
        r = d(tok, rec, cv)
    y = m.predict({"token": tok.numpy(), "rec": rec.numpy(), "conv": cv.numpy()})
    print("decoder maxdiff", [float(np.abs(y[k] - v.numpy()).max()) for k, v in zip(["g", "rec_out", "conv_out"], r)])
    e = torch.randn(1, 1024, 188)
    m = conv(JointEnc(model), (e,), [ct.TensorType("encoded", shape=(1, 1024, 188), dtype=np.float32)], ["f"], "joint_enc")
    print("joint_enc maxdiff", float(np.abs(m.predict({"encoded": e.numpy()})["f"] - JointEnc(model)(e).detach().numpy()).max()))
    f, g = torch.randn(1, 640), torch.randn(1, 640)
    m = conv(Joint(model), (f, g), [ct.TensorType("f", shape=(1, 640), dtype=np.float32),
                                    ct.TensorType("g", shape=(1, 640), dtype=np.float32)], ["token_id", "duration_id"], "joint")
    print("joint", m.predict({"f": f.numpy(), "g": g.numpy()}), [t.item() for t in Joint(model)(f, g)])