#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path

import mlx.core as mx
import numpy as np
from safetensors import safe_open


def load_tensors(path: Path) -> dict[str, mx.array]:
    tensors: dict[str, mx.array] = {}
    with safe_open(str(path), framework="np") as f:
        for key in f.keys():
            tensors[key] = mx.array(f.get_tensor(key))
    return tensors


def linear(x: mx.array, weight: mx.array, bias: mx.array) -> mx.array:
    return x @ mx.transpose(weight) + bias


def sigmoid(x: mx.array) -> mx.array:
    return 1.0 / (1.0 + mx.exp(-x))


def log_softmax(x: mx.array, axis: int = -1) -> mx.array:
    shifted = x - mx.max(x, axis=axis, keepdims=True)
    return shifted - mx.log(mx.sum(mx.exp(shifted), axis=axis, keepdims=True))


class PhaseBRNNT:
    def __init__(self, weights: dict[str, mx.array]):
        self.w = weights

    def predictor(self, token: mx.array, state_h: mx.array, state_c: mx.array):
        embed = self.w["decoder.prediction.embed.weight"][token]
        # token shape [B, U], this runtime step is intended for U=1.
        x = embed[:, 0, :]
        h0 = state_h[0]
        c0 = state_c[0]
        gates = (
            x @ mx.transpose(self.w["decoder.prediction.dec_rnn.lstm.weight_ih_l0"])
            + self.w["decoder.prediction.dec_rnn.lstm.bias_ih_l0"]
            + h0 @ mx.transpose(self.w["decoder.prediction.dec_rnn.lstm.weight_hh_l0"])
            + self.w["decoder.prediction.dec_rnn.lstm.bias_hh_l0"]
        )
        i, f, g, o = mx.split(gates, 4, axis=-1)
        c1 = sigmoid(f) * c0 + sigmoid(i) * mx.tanh(g)
        h1 = sigmoid(o) * mx.tanh(c1)
        pred = h1[:, None, :]
        return pred, h1[None, :, :], c1[None, :, :]

    def joint(self, encoded: mx.array, prediction: mx.array) -> mx.array:
        # encoded [B, D, T], prediction [B, U, Dp].
        enc = mx.transpose(encoded, (0, 2, 1))
        enc_proj = linear(enc, self.w["joint.enc.weight"], self.w["joint.enc.bias"])
        pred_proj = linear(prediction, self.w["joint.pred.weight"], self.w["joint.pred.bias"])
        hidden = mx.maximum(enc_proj[:, :, None, :] + pred_proj[:, None, :, :], 0)
        logits = linear(hidden, self.w["joint.joint_net.2.weight"], self.w["joint.joint_net.2.bias"])
        return log_softmax(logits, axis=-1)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--weights", type=Path, required=True)
    parser.add_argument("--out", type=Path, default=None)
    args = parser.parse_args()

    rt = PhaseBRNNT(load_tensors(args.weights))
    token = mx.array(np.ones((1, 1), dtype=np.int32))
    h = mx.zeros((1, 1, 640), dtype=mx.float16)
    c = mx.zeros((1, 1, 640), dtype=mx.float16)
    pred, h1, c1 = rt.predictor(token, h, c)
    encoded = mx.random.normal((1, 512, 2)).astype(mx.float16)
    logits = rt.joint(encoded, pred)
    mx.eval(pred, h1, c1, logits)
    result = {
        "weights": str(args.weights),
        "predictor_shapes": [list(pred.shape), list(h1.shape), list(c1.shape)],
        "joint_shape": list(logits.shape),
        "dtype": str(logits.dtype),
    }
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()