# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.

# SPDX-License-Identifier: Apache-2.0
import os

import pytest
import torch
from loguru import logger

import ttnn
from models.common.utility_functions import comp_allclose, comp_pcc
from models.tt_transformers.tests.test_utils import get_ref_model_dype
from models.tt_transformers.tt.ccl import TT_CCL
from models.tt_transformers.tt.common import Mode, PagedAttentionConfig, get_rot_transformation_mat, precompute_freqs
from models.tt_transformers.tt.decoder import TransformerBlock
from models.tt_transformers.tt.model_config import ModelArgs
from models.tt_transformers.tt.prefetcher import Prefetcher
from models.tt_transformers.tt.rope import get_rot_mats, get_rot_mats_hf


@torch.no_grad()
@pytest.mark.parametrize(
    "mesh_device",
    [
        {"N150": (1, 1), "N300": (1, 2), "T3K": (1, 8), "TG": (8, 4)}.get(
            os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids())
        )
    ],
    indirect=True,
)
@pytest.mark.parametrize(
    "paged_attention",
    (
        True,
        # False
    ),
    ids=(
        "paged_attention",
        # "default_attention"
    ),
)
@pytest.mark.parametrize(
    "page_params",
    [{"page_block_size": 32, "page_max_num_blocks": 1024}],
)
@pytest.mark.parametrize(
    "max_seq_len",
    (
        4096,
        128,
    ),
)
@pytest.mark.parametrize(
    "use_prefetcher",
    ([False]),
)
@pytest.mark.parametrize("device_params", [{"fabric_config": True}], indirect=True)
def test_decoder_inference(
    max_seq_len,
    paged_attention,
    page_params,
    mesh_device,
    reset_seeds,
    ensure_gc,
    use_prefetcher,
):
    model_name_env = os.getenv("HF_MODEL")
    if max_seq_len > 256 and model_name_env and "Mistral-7B" in model_name_env:
        pytest.skip(
            "Mistral-7B models do not support max_seq_len > 256. See issue: https://github.com/tenstorrent/tt-metal/issues/19806"
        )

    dtype = ttnn.bfloat8_b
    batch_size = 1  # For prefill we only support batch_size = 1

    num_tensors = 0
    prefetcher = Prefetcher(mesh_device, num_tensors=num_tensors, num_layers=1) if use_prefetcher else None
    if use_prefetcher:
        prefetcher.init(mode=Mode.PREFILL)

    model_args = ModelArgs(
        mesh_device,
        max_batch_size=batch_size,
        max_seq_len=max_seq_len,
        cache_hf=True,
        prefetcher=prefetcher,
        use_hf_rope=False,
    )
    model_args.n_layers = 1

    state_dict = model_args.load_state_dict()

    reference_model = model_args.reference_decoder(load_checkpoint=True)

    generation_start_pos = 0
    generation_length = 1
    all_tests_pass = True

    rot_mats_fn = get_rot_mats_hf if model_args.use_hf_rope else get_rot_mats

    # pre-compute the rotational embedding matrix and send to device
    rot_mats = rot_mats_fn(
        head_dim=model_args.head_dim,
        device=mesh_device,
        seq_len=max_seq_len,
        theta=model_args.rope_theta,
        rope_scaling=model_args.rope_scaling,
    )
    if model_args.rope_theta_local is not None:
        rot_mats_local = get_rot_mats(
            head_dim=model_args.head_dim,
            device=mesh_device,
            seq_len=max_seq_len,
            theta=model_args.rope_theta_local,
            rope_scaling=None,
        )
    else:
        rot_mats_local = None

    transformation_mats = {}
    if not model_args.use_hf_rope:
        transformation_mat_torch = get_rot_transformation_mat(model_args.head_dim)
        transformation_mats_prefill = ttnn.as_tensor(
            transformation_mat_torch,
            dtype=ttnn.bfloat16,
            layout=ttnn.TILE_LAYOUT,
            device=mesh_device,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
            mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
        )
        transformation_mats = {"prefill": transformation_mats_prefill}

    # Setup page table
    page_table_tt = None
    paged_attention_config = None

    if paged_attention:
        paged_attention_config = PagedAttentionConfig(
            block_size=page_params["page_block_size"],
            max_num_blocks=page_params["page_max_num_blocks"],
        )
        # Implied shuffling of blocks
        permutation = torch.randperm(paged_attention_config.max_num_blocks)
        # Page table which maps virtual blocks to physical
        reverse_permutation = torch.argsort(permutation)
        page_table = reverse_permutation.reshape(
            model_args.max_batch_size, paged_attention_config.max_num_blocks // model_args.max_batch_size
        )
        page_table_tt = ttnn.from_torch(
            page_table,
            device=mesh_device,
            dtype=ttnn.int32,
            layout=ttnn.ROW_MAJOR_LAYOUT,
            mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
        )

    # Initialize TT model
    tt_ccl = TT_CCL(mesh_device)
    tt_model = TransformerBlock(
        mesh_device=mesh_device,
        tt_ccl=tt_ccl,
        state_dict=state_dict,
        weight_cache_path=model_args.weight_cache_path(dtype),
        layer_num=0,
        dtype=dtype,
        transformation_mats=transformation_mats,
        args=model_args,
        paged_attention_config=paged_attention_config,
        prefetcher=prefetcher,
    )

    for i in range(generation_length):
        logger.info(f"[Decoder] Generating token {i}")
        pt_decode_input = (
            torch.rand(
                batch_size,
                max_seq_len,
                model_args.dim,
                dtype=get_ref_model_dype(reference_model, model_args.model_name),
            )
            * 2
        ) - 1
        tt_decode_input = pt_decode_input.clone()
        decode_input = model_args.prepare_residual_tensor_prefill(
            tt_decode_input,
        )
        positions = torch.LongTensor(range(max_seq_len))

        cos, sin = precompute_freqs(
            model_args.head_dim,
            model_args.max_seq_len * 2,
            model_args.rope_theta,
            model_args.rope_scaling.factor if model_args.rope_scaling else None,
            model_args.rope_scaling.original_max_position_embeddings if model_args.rope_scaling else None,
            model_args.rope_scaling.rope_type.value if model_args.rope_scaling else "llama3",
        )
        freqs_cis_i = torch.complex(cos, sin)[positions]

        # Reference model
        attn_mask = torch.triu(torch.full((max_seq_len, max_seq_len), torch.finfo(torch.float32).min), diagonal=1)
        if (
            model_args.sliding_window is not None
            and model_args.sliding_window > 0
            and (
                hasattr(reference_model.decoder, "attention_type")
                and reference_model.decoder.attention_type == "sliding_attention"
            )
        ):
            attn_mask += torch.tril(
                torch.full((max_seq_len, max_seq_len), -float("inf")), diagonal=-model_args.sliding_window
            )
        ref_output = reference_model(
            pt_decode_input,
            positions[0],
            freqs_cis_i,
            mask=attn_mask,
        )
        if ref_output.dim() == 2:
            ref_output = ref_output.unsqueeze(0)
        # Run TT model
        tt_out = tt_model(
            decode_input,
            None,
            rot_mats_global=rot_mats,
            rot_mats_local=rot_mats_local,
            user_id=0,
            mode=Mode.PREFILL,
            page_table=page_table_tt,
        )
        tt_out = ttnn.to_torch(
            tt_out,
            mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=(1, 3), mesh_shape=model_args.cluster_shape),
        )
        tt_output_torch = tt_out[:, 0:1, :, : model_args.dim].view(
            batch_size, max_seq_len, -1
        )  # [ batch_size, seq, hidden_dim]
        passing, pcc_message = comp_pcc(ref_output, tt_output_torch)

        logger.info(comp_allclose(ref_output, tt_output_torch))
        logger.info(f"PCC: {pcc_message}")
        if passing:
            logger.info("Decoder Block Passed!")
        else:
            logger.warning("Decoder Block Failed!")
            all_tests_pass = False

    if all_tests_pass:
        logger.info(f"All decode iterations Passed!")
    else:
        logger.warning("One or more iterations of decode Failed!")
        assert all_tests_pass, f"PCC value is lower than {0.99} for some of the outputs. Check Warnings!"