import torch
from torchax import tensor  # pylint: disable=unused-import
import torchax
import torchax.export

from .. import base_test_util
from . import llama_model
from . import model_exportable
from torch.utils import _pytree as pytree


class LlamaTest(base_test_util.TestCase):

  def test_can_run(self):
    with torchax.default_env():
      sample_args = (
          torch.randint(0, 32000, (1, 2048), device='jax:0'),
          torch.arange(0, 2048, device='jax:0'),
      )

      model_args = llama_model.ModelArgs(
          block_size=2048,
          vocab_size=32000,
          n_layer=2,
          n_head=4,
          dim=256,
      )
      m = llama_model.Transformer(model_args)
      m.to(torch.bfloat16)
      m.setup_caches(1, 2048)
      m = m.to('jax')

      print(m(*sample_args))

  def test_can_run_exportable(self):
    model_args = model_exportable.ModelArgs(
        vocab_size=32000,
        n_layers=2,
        n_heads=4,
        dim=256,
    )
    m = model_exportable.Transformer(model_args)
    context_length = 2048
    input_shape_prefill = (1, context_length)
    input_shape_decode = (1, 1)

    def make_cache(args, batch_size):
      n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
      n_local_heads = args.n_heads
      n_local_kv_heads = n_kv_heads
      n_rep = n_local_heads // n_local_kv_heads
      head_dim = args.dim // args.n_heads
      res = []
      for i in range(args.n_layers):
        if batch_size is None:
          size = (
              args.max_seq_len,
              n_local_kv_heads,
              head_dim,
          )
        else:
          size = (
              batch_size,
              args.max_seq_len,
              n_local_kv_heads,
              head_dim,
          )
        res.append(
            (torch.zeros(
                size,
                dtype=torch.bfloat16 if args.bf16_enable else torch.float),
             torch.zeros(
                 size,
                 dtype=torch.bfloat16 if args.bf16_enable else torch.float)))
      return res

    prefill_caches = make_cache(model_args, 1)

    sample_input_prefill = (
        torch.randint(0, 1000, input_shape_prefill,
                      dtype=torch.int32),  # len seq length
        torch.arange(0, context_length, dtype=torch.int32),  # input indexes
        torch.arange(0, context_length, dtype=torch.int32),  # context indexes
        prefill_caches,
        True,  # prefil
    )
    with torch.no_grad():
      m_prefill = torch.export.export(m, sample_input_prefill)

    weights, mj_prefill = torchax.export.exported_program_to_jax(m_prefill)
    env = torchax.default_env()
    sample_inputs = env.t2j_copy(sample_input_prefill)
    print('Prefill', mj_prefill(weights, sample_inputs))

    sample_input_decode = (
        torch.randint(0, 1000, input_shape_decode,
                      dtype=torch.int32),  # len = 1
        torch.tensor([0], dtype=torch.int32),
        torch.roll(torch.arange(context_length, dtype=torch.int32), 1, 0),
        prefill_caches,
        False  # prefill
    )
    with torch.no_grad():
      m_decode = torch.export.export(m, sample_input_decode)
    weights, mj_decode = torchax.export.exported_program_to_jax(m_decode)
    sample_inputs = env.t2j_copy(sample_input_decode)
    print('Decode', mj_decode(weights, sample_inputs))


if __name__ == "__main__":
  base_test_util.main()
