"""pi0 (Physical Intelligence / lerobot) inference runtime for Apple Silicon / MLX.

Faithful port of the lerobot (openpi-derived) PI0Pytorch reference:
- prefill: PaliGemma LM (18 layers, bidirectional prefix) with RoPE-applied KV cache
- 10 Euler flow-matching steps; the 300M Gemma expert attends to the cached
  prefix KV at every layer (head_dim 256 / 1 KV head on both sides)
- Gemma specifics: (1+w) RMSNorm, GeGLU (tanh), rope base 10000, language
  embeddings scaled by sqrt(width); attention/softmax in fp32
- verified against the PyTorch reference: cosine 1.00000 on full action chunks

Usage:
    from pi0_mlx import Pi0MLX
    m = Pi0MLX.from_pretrained(".")
    actions = m.predict(images=[cam0, cam1, cam2],
                        instruction="pick up the cube", state=[...])

Note: the tokenizer is loaded from google/paligemma-3b-pt-224 (gated; accept
the Gemma terms on Hugging Face and login first), matching upstream lerobot.
"""
import json
import math
from pathlib import Path

import mlx.core as mx
import mlx.nn as nn
import numpy as np

from mlx_vlm.models.paligemma.vision import VisionModel
from mlx_vlm.models.paligemma import VisionConfig

CHUNK = 50
MAX_DIM = 32
NUM_STEPS = 10
MIN_PERIOD, MAX_PERIOD = 4e-3, 4.0
N_LAYERS = 18
HID = 2048
EXP_HID = 1024
N_HEADS = 8
N_KV = 1
HEAD_DIM = 256
EPS = 1e-6
EMBED_SCALE = HID ** 0.5
IMG_SIZE = 224
VIS_CFG = {"model_type": "siglip_vision_model", "hidden_size": 1152, "num_hidden_layers": 27,
           "intermediate_size": 4304, "num_attention_heads": 16, "image_size": 224,
           "patch_size": 14, "num_channels": 3, "layer_norm_eps": 1e-6}


def rms_norm_gemma(x, w):
    x32 = x.astype(mx.float32)
    out = x32 * mx.rsqrt(mx.mean(x32 * x32, axis=-1, keepdims=True) + EPS)
    return (out * (1.0 + w.astype(mx.float32))).astype(x.dtype)


def apply_rope(x, positions, base=10000.0):
    d_half = x.shape[-1] // 2
    dtype = x.dtype
    x = x.astype(mx.float32)
    freq_exp = (2.0 / x.shape[-1]) * mx.arange(d_half, dtype=mx.float32)
    timescale = mx.power(base, freq_exp)
    radians = positions[..., None].astype(mx.float32) / timescale[None, None, :]
    radians = radians[..., None, :]
    sin, cos = mx.sin(radians), mx.cos(radians)
    x1, x2 = x[..., :d_half], x[..., d_half:]
    return mx.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis=-1).astype(dtype)


def make_att_2d_masks(pad_masks, att_masks):
    cs = mx.cumsum(att_masks.astype(mx.int32), axis=1)
    return (cs[:, None, :] <= cs[:, :, None]) & pad_masks[:, None, :]


def sinusoidal_time_emb(time, dim):
    frac = np.linspace(0.0, 1.0, dim // 2, dtype=np.float64)
    period = MIN_PERIOD * (MAX_PERIOD / MIN_PERIOD) ** frac
    scale = 1.0 / period * 2 * np.pi
    sin_in = scale[None, :] * np.asarray(time, dtype=np.float64)[:, None]
    return mx.array(np.concatenate([np.sin(sin_in), np.cos(sin_in)], axis=1).astype(np.float32))


def resize_with_pad(img_chw, size=IMG_SIZE):
    """参照実装同等: アスペクト維持バイリニア縮小＋左・上ゼロパディング。"""
    from PIL import Image

    c, h, w = img_chw.shape
    ratio = max(w / size, h / size)
    rh, rw = int(h / ratio), int(w / ratio)
    pil = Image.fromarray((np.transpose(img_chw, (1, 2, 0)) * 255).clip(0, 255).astype(np.uint8))
    pil = pil.resize((rw, rh), Image.BILINEAR)
    arr = np.asarray(pil).astype(np.float32) / 255.0
    out = np.zeros((size, size, 3), dtype=np.float32)
    out[size - rh:, size - rw:, :] = arr
    return out


class Pi0MLX:
    P_LM = "model.paligemma_with_expert.paligemma.model.language_model."
    P_EXP = "model.paligemma_with_expert.gemma_expert.model."
    P_VIS = "model.paligemma_with_expert.paligemma.model.vision_tower."
    P_PROJ = "model.paligemma_with_expert.paligemma.model.multi_modal_projector."
    P_EMB = "model.paligemma_with_expert.paligemma.lm_head.weight"  # tied weights

    def __init__(self, weights):
        self.w = weights
        self.vision = VisionModel(VisionConfig(**VIS_CFG))
        vis_w = {}
        for k, v in weights.items():
            if k.startswith(self.P_VIS):
                kk = k[len(self.P_VIS):]
                if kk.endswith("patch_embedding.weight") and v.shape[-1] != 3:
                    v = v.transpose(0, 2, 3, 1)
                vis_w[kk] = v
        self.vision.load_weights(list(vis_w.items()), strict=False)
        mx.eval(self.vision.parameters())

    @classmethod
    def from_pretrained(cls, path):
        path = Path(path)
        m = cls(mx.load(str(path / "model.safetensors")))
        from transformers import AutoTokenizer

        m.tokenizer = AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224")
        return m

    def lm(self, i, name):
        return self.w[f"{self.P_LM}layers.{i}.{name}"]

    def exp(self, i, name):
        return self.w[f"{self.P_EXP}layers.{i}.{name}"]

    def _attn(self, mask2d, q, k, v):
        B, Lk = k.shape[0], k.shape[1]
        groups = N_HEADS // N_KV
        k = mx.repeat(k[:, :, :, None, :], groups, axis=3).reshape(B, Lk, N_HEADS, HEAD_DIM)
        v = mx.repeat(v[:, :, :, None, :], groups, axis=3).reshape(B, Lk, N_HEADS, HEAD_DIM)
        q32 = q.astype(mx.float32).transpose(0, 2, 1, 3)
        k32 = k.astype(mx.float32).transpose(0, 2, 1, 3)
        att = (q32 @ k32.transpose(0, 1, 3, 2)) * (HEAD_DIM ** -0.5)
        att = mx.where(mask2d[:, None, :, :], att, mx.finfo(mx.float32).min)
        probs = mx.softmax(att, axis=-1).astype(v.dtype)
        out = probs @ v.transpose(0, 2, 1, 3)
        return out.transpose(0, 2, 1, 3).reshape(B, -1, N_HEADS * HEAD_DIM)

    def _layer(self, get, i, h, mask2d, pos, cache=None, fill=False):
        hn = rms_norm_gemma(h, get(i, "input_layernorm.weight"))
        B, L = hn.shape[:2]
        q = (hn @ get(i, "self_attn.q_proj.weight").T).reshape(B, L, -1, HEAD_DIM)
        k = (hn @ get(i, "self_attn.k_proj.weight").T).reshape(B, L, -1, HEAD_DIM)
        v = (hn @ get(i, "self_attn.v_proj.weight").T).reshape(B, L, -1, HEAD_DIM)
        q = apply_rope(q, pos)
        k = apply_rope(k, pos)
        if fill:
            cache[i] = (k, v)
        elif cache is not None:
            k = mx.concatenate([cache[i][0], k], axis=1)
            v = mx.concatenate([cache[i][1], v], axis=1)
        att = self._attn(mask2d, q, k, v)
        out = att @ get(i, "self_attn.o_proj.weight").T + h
        res = out
        on = rms_norm_gemma(out, get(i, "post_attention_layernorm.weight"))
        gate = on @ get(i, "mlp.gate_proj.weight").T
        up = on @ get(i, "mlp.up_proj.weight").T
        return (nn.gelu_approx(gate) * up) @ get(i, "mlp.down_proj.weight").T + res

    def embed_prefix(self, imgs224, tokens, lang_mask):
        embs, pads, atts = [], [], []
        vdtype = self.vision.vision_model.embeddings.patch_embedding.weight.dtype
        for img in imgs224:  # [1,224,224,3] in [-1,1]
            feat = self.vision(img.astype(vdtype))
            feat = feat[0] if isinstance(feat, tuple) else feat
            if feat.ndim == 2:
                feat = feat[None]
            feat = feat @ self.w[self.P_PROJ + "linear.weight"].T + self.w[self.P_PROJ + "linear.bias"]
            embs.append(feat.astype(mx.bfloat16))
            pads.append(mx.ones(feat.shape[:2], dtype=mx.bool_))
            atts += [0] * feat.shape[1]
        lang = (self.w[self.P_EMB][tokens].astype(mx.float32) * EMBED_SCALE).astype(mx.bfloat16)
        embs.append(lang)
        pads.append(lang_mask.astype(mx.bool_))
        atts += [0] * lang.shape[1]
        embs = mx.concatenate(embs, axis=1)
        pads = mx.concatenate(pads, axis=1)
        atts = mx.array(atts, dtype=mx.int32)[None, :]
        return embs, pads, mx.broadcast_to(atts, (embs.shape[0], atts.shape[1]))

    def embed_suffix(self, state32, x_t, time):
        state_emb = state32 @ self.w["model.state_proj.weight"].T + self.w["model.state_proj.bias"]
        act = x_t @ self.w["model.action_in_proj.weight"].T + self.w["model.action_in_proj.bias"]
        t_emb = sinusoidal_time_emb(time, EXP_HID).astype(act.dtype)
        at = mx.concatenate([act, mx.broadcast_to(t_emb[:, None, :], act.shape)], axis=2)
        at = at @ self.w["model.action_time_mlp_in.weight"].T + self.w["model.action_time_mlp_in.bias"]
        at = nn.silu(at)
        at = at @ self.w["model.action_time_mlp_out.weight"].T + self.w["model.action_time_mlp_out.bias"]
        embs = mx.concatenate([state_emb[:, None, :], at], axis=1).astype(mx.bfloat16)
        pads = mx.ones(embs.shape[:2], dtype=mx.bool_)
        atts = mx.array([1, 1] + [0] * (CHUNK - 1), dtype=mx.int32)[None, :]
        return embs, pads, mx.broadcast_to(atts, (embs.shape[0], atts.shape[1]))

    def sample_actions(self, imgs224, tokens, lang_mask, state32, noise=None):
        if noise is None:
            noise = mx.random.normal((state32.shape[0], CHUNK, MAX_DIM))
        prefix, pads, atts = self.embed_prefix(imgs224, tokens, lang_mask)
        mask2d = make_att_2d_masks(pads, atts)
        pos = mx.cumsum(pads.astype(mx.int32), axis=1) - 1
        cache = {}
        h = prefix
        for i in range(N_LAYERS):
            h = self._layer(self.lm, i, h, mask2d, pos, cache=cache, fill=True)
        P = pads.shape[1]
        offset = mx.sum(pads.astype(mx.int32), axis=-1)[:, None]
        x_t = noise
        dt = -1.0 / NUM_STEPS
        for step in range(NUM_STEPS):
            t = 1.0 + step * dt
            suffix, s_pads, s_atts = self.embed_suffix(state32, x_t, [t] * x_t.shape[0])
            L = s_pads.shape[1]
            mask_full = mx.concatenate(
                [mx.broadcast_to(pads[:, None, :], (s_pads.shape[0], L, P)),
                 make_att_2d_masks(s_pads, s_atts)], axis=2)
            pos_s = offset + mx.cumsum(s_pads.astype(mx.int32), axis=1) - 1
            h = suffix
            for i in range(N_LAYERS):
                h = self._layer(self.exp, i, h, mask_full, pos_s, cache=cache, fill=False)
            h = rms_norm_gemma(h, self.w[self.P_EXP + "norm.weight"])
            out = h[:, -CHUNK:].astype(mx.float32)
            v_t = out @ self.w["model.action_out_proj.weight"].T + self.w["model.action_out_proj.bias"]
            x_t = x_t + dt * v_t
            mx.eval(x_t)
        return x_t

    def predict(self, images, instruction, state, action_dim=None, noise=None):
        imgs224 = []
        for im in images:
            arr = np.asarray(im).astype(np.float32)
            if arr.max() > 1.5:
                arr = arr / 255.0
            hwc = resize_with_pad(np.transpose(arr, (2, 0, 1))) * 2.0 - 1.0
            imgs224.append(mx.array(hwc[None]))
        enc = self.tokenizer(instruction.rstrip("\n") + "\n", padding="max_length",
                             max_length=48, return_tensors="np")
        state = np.asarray(state, dtype=np.float32)[None]
        action_dim = action_dim or state.shape[1]
        state32 = np.zeros((1, MAX_DIM), dtype=np.float32)
        state32[:, : state.shape[1]] = state
        chunk = self.sample_actions(imgs224, mx.array(enc["input_ids"]),
                                    mx.array(enc["attention_mask"]).astype(mx.bool_),
                                    mx.array(state32), noise=noise)
        return np.array(chunk[0, :, :action_dim])


if __name__ == "__main__":
    import argparse

    ap = argparse.ArgumentParser()
    ap.add_argument("--images", nargs="+", required=True)
    ap.add_argument("--instruction", required=True)
    ap.add_argument("--state", default="0,0,0,0,0,0")
    ap.add_argument("--out", default="actions.npy")
    args = ap.parse_args()
    from PIL import Image

    model = Pi0MLX.from_pretrained(Path(__file__).parent)
    imgs = [np.asarray(Image.open(p).convert("RGB")) for p in args.images]
    state = [float(x) for x in args.state.split(",")]
    actions = model.predict(imgs, args.instruction, state)
    np.save(args.out, actions)
    print(f"action chunk {actions.shape} -> {args.out}")
    print("first action:", actions[0].round(4))