# 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.tt.embedding import Embedding, ScaledEmbedding
from models.tt_transformers.tt.model_config import ModelArgs


@torch.no_grad()
@pytest.mark.parametrize("use_scaled_embedding", (False, True))
@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(
    "batch_size",
    (1,),
)
@pytest.mark.parametrize(
    "max_seq_len",
    (128,),  # For decode-only unit test, there's no need to run with large sequence lengths
)
def test_embedding(max_seq_len, batch_size, mesh_device, reset_seeds, ensure_gc, use_scaled_embedding):
    dtype = ttnn.bfloat16

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

    state_dict = model_args.load_state_dict()
    tokenizer = model_args.tokenizer

    if use_scaled_embedding:
        model_args.embed_scale = model_args.dim**0.5
        logger.info(f"Using scaled embedding with scale {model_args.embed_scale}")

    reference_emb = model_args.reference_embedding()

    if model_args.is_llama_vision():
        weight = torch.cat(
            [
                state_dict["text_model.tok_embeddings.weight"],
                state_dict["text_model.learnable_embedding.weight"],
            ],
            dim=0,
        )
    else:
        weight = state_dict["tok_embeddings.weight"]
    reference_emb.load_state_dict({"emb.weight": weight})

    emb_kwargs = {
        "mesh_device": mesh_device,
        "args": model_args,
        "weight_cache_path": model_args.weight_cache_path(dtype),
        "state_dict": state_dict,
        "dtype": dtype,
    }
    if use_scaled_embedding:
        emb_kwargs["embed_scale"] = model_args.embed_scale
        emb_cls = ScaledEmbedding
    else:
        emb_cls = Embedding

    tt_emb = emb_cls(**emb_kwargs)

    prompts = ["Joy"] * 32
    pt_input = torch.tensor([model_args.encode_prompt(prompt, instruct=False) for prompt in prompts])
    reference_output = reference_emb(pt_input)
    logger.info(f"reference_output: {reference_output.shape}")

    tt_input = ttnn.from_torch(
        pt_input.squeeze(1),
        device=mesh_device,
        mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
        dtype=ttnn.uint32,
        layout=ttnn.ROW_MAJOR_LAYOUT,
    )
    tt_output = tt_emb(tt_input)
    tt_output_torch = ttnn.to_torch(
        tt_output,
        mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=(0, -1), mesh_shape=model_args.cluster_shape),
    )[:32].view(reference_output.shape)
    logger.info(f"tt_output_torch: {tt_output_torch.shape}")

    passing, pcc_message = comp_pcc(reference_output, tt_output_torch)

    logger.info(comp_allclose(reference_output, tt_output_torch))
    logger.info(f"PCC: {pcc_message}")
    if passing:
        logger.info("embedding Passed!")
    else:
        logger.warning("embedding Failed!")

    assert passing, f"embedding output does not meet PCC requirement {0.99}."