#!/usr/bin/env python3
"""Integration Test for StreamingMoELayer with Quantized Weights on Apple Silicon MLX."""

import time
from pathlib import Path

import mlx.core as mx

from expert_registry import SafetensorsIndex
from moe_stream_layer import StreamingMoELayer


def run_test():
    model_dir = Path("/Volumes/NVMe2TB/AI-LLM/LLMjp/Pyrite/pyrite-mlx/results/models/deepseek_coder_v2_lite_pyrite_final_mlx")
    shard1 = model_dir / "model-00001-of-00002.safetensors"
    
    print("=" * 70)
    print(" Streaming MoE Layer Integration Test (Layer 1)")
    print("=" * 70)
    
    idx = SafetensorsIndex(shard1)
    idx.open()
    
    # 1. Load resident gate & shared expert weights
    print("Loading resident weights (Gate + Shared Experts)...")
    gate_w = idx.get_tensor("model.layers.1.mlp.gate.weight")
    
    shared = {
        "gate_w": idx.get_tensor("model.layers.1.mlp.shared_experts.gate_proj.weight"),
        "gate_s": idx.get_tensor("model.layers.1.mlp.shared_experts.gate_proj.scales"),
        "gate_b": idx.get_tensor("model.layers.1.mlp.shared_experts.gate_proj.biases"),
        "up_w": idx.get_tensor("model.layers.1.mlp.shared_experts.up_proj.weight"),
        "up_s": idx.get_tensor("model.layers.1.mlp.shared_experts.up_proj.scales"),
        "up_b": idx.get_tensor("model.layers.1.mlp.shared_experts.up_proj.biases"),
        "down_w": idx.get_tensor("model.layers.1.mlp.shared_experts.down_proj.weight"),
        "down_s": idx.get_tensor("model.layers.1.mlp.shared_experts.down_proj.scales"),
        "down_b": idx.get_tensor("model.layers.1.mlp.shared_experts.down_proj.biases"),
    }
    
    print(f"Gate weight shape: {gate_w.shape}")
    print(f"Shared experts: gate_w={shared['gate_w'].shape}, down_w={shared['down_w'].shape}")
    
    # 2. Instantiate StreamingMoELayer with LRU cache (max 8 experts in RAM)
    layer = StreamingMoELayer(
        layer_idx=1,
        gate_weight=gate_w,
        shared_experts=shared,
        index_reader=idx,
        top_k=6,
        max_cache_experts=8,
    )
    
    # 3. Simulate forward passes with synthetic activations
    hidden_dim = 2048
    seq_len = 4
    x = mx.random.normal((1, seq_len, hidden_dim)).astype(mx.bfloat16)
    
    print("\n--- Token Step 1 (Cold Start: Streaming from NVMe SSD) ---")
    t0 = time.perf_counter()
    out1 = layer(x)
    mx.eval(out1)
    elapsed_cold = (time.perf_counter() - t0) * 1000
    print(f"Step 1 Output shape: {out1.shape}")
    print(f"Cold Start Time: {elapsed_cold:.2f} ms")
    print(f"Current Cached Experts in RAM: {list(layer.cache.cache.keys())}")
    
    print("\n--- Token Step 2 (Subsequent Token: Testing Cache Hits) ---")
    t0 = time.perf_counter()
    out2 = layer(x)
    mx.eval(out2)
    elapsed_warm = (time.perf_counter() - t0) * 1000
    print(f"Warm Step Time: {elapsed_warm:.2f} ms")
    print(f"Current Cached Experts in RAM: {list(layer.cache.cache.keys())}")
    
    idx.close()
    print("\n" + "=" * 70)
    print(" Test Passed! Dynamic streaming & LRU caching verified on Apple Metal! ")
    print("=" * 70)

if __name__ == "__main__":
    run_test()