#!/usr/bin/env python3
import tempfile
import random
import numpy as np
import torch
import optax
import jax
import sys
from flax.training.common_utils import onehot
from flax.traverse_util import flatten_dict


def ids_tensor(shape, vocab_size, rng=None):
    """Creates a random int32 tensor of the shape within the vocab size."""
    if rng is None:
        rng = random.Random()

    total_dims = 1
    for dim in shape:
        total_dims *= dim

    values = []
    for _ in range(total_dims):
        values.append(rng.randint(0, vocab_size - 1))

    output = np.array(values).reshape(shape)

    return output


def random_attention_mask(shape, rng=None):
    attn_mask = ids_tensor(shape, vocab_size=2, rng=rng)
    # make sure that at least one token is attended to for each batch
    attn_mask[:, -1] = 1
    return attn_mask


def floats_tensor(shape, scale=1.0, rng=None):
    """Creates a random float32 tensor"""
    if rng is None:
        rng = random.Random()

    total_dims = 1
    for dim in shape:
        total_dims *= dim

    values = []
    for _ in range(total_dims):
        values.append(rng.random() * scale)

    return np.array(values, dtype=np.float32).reshape(shape)


def shift_tokens_right(input_ids: np.array, pad_token_id: int, decoder_start_token_id: int) -> np.ndarray:
    """
    Shift input ids one token to the right.
    """
    shifted_input_ids = np.zeros_like(input_ids)
    shifted_input_ids[:, 1:] = input_ids[:, :-1]
    shifted_input_ids[:, 0] = decoder_start_token_id

    shifted_input_ids = np.where(shifted_input_ids == -100, pad_token_id, shifted_input_ids)
    return shifted_input_ids


def assert_almost_equals(a: np.ndarray, b: np.ndarray, tol: float = 1e-2):
    diff = np.abs((a - b)).max()
    if diff < tol:
        print(f"✅ Difference between Flax and PyTorch is {diff} (< {tol})")
    else:
        print(f"❌ Difference between Flax and PyTorch is {diff} (>= {tol})")


def assert_dict_equal(a: dict, b: dict, tol: float = 1e-2):
    if a.keys() != b.keys():
        print("❌ Dictionary keys for PyTorch and Flax do not match")
    results_fail = []
    results_correct = []

    results_fail_rel = []
    results_correct_rel = []
    for k in a:
        ak_norm = np.linalg.norm(a[k])
        bk_norm = np.linalg.norm(b[k])
        diff = np.abs(ak_norm - bk_norm)
        diff_rel = np.abs(ak_norm - bk_norm) / np.abs(ak_norm)
        if diff < tol:
            results_correct.append(f"✅ Layer {k} diff is {diff} < {tol}).")
        else:
            results_fail.append(f"❌ Layer {k} has PT grad norm {bk_norm} and flax grad norm {ak_norm}.")
        if diff_rel < tol:
            results_correct_rel.append(f"✅ Layer {k} rel diff is {diff} < {tol}).")
        else:
            results_fail_rel.append(f"❌ Layer {k} has PT grad norm {bk_norm} and flax grad norm {ak_norm}.")
    return results_fail_rel, results_correct_rel, results_fail, results_correct


def compare_grads(model_id, pt_architecture):
    transformers_module = __import__("transformers", fromlist=[pt_architecture])

    model_cls = getattr(transformers_module, pt_architecture)
    flax_model_cls = getattr(transformers_module, "Flax" + pt_architecture)

    pt_model, model_info = model_cls.from_pretrained(model_id, output_loading_info=True)

    if len(model_info["missing_keys"]) > 0:
        raise ValueError(f"{model_id} with {pt_architecture} has missing keys: {model_info['missing_keys']}")

    fx_model = flax_model_cls.from_pretrained(model_id, from_pt=True)

    batch_size = 2
    seq_len = 64

    input_ids = ids_tensor([batch_size, seq_len], fx_model.config.vocab_size)
    label_ids = ids_tensor([batch_size, seq_len], fx_model.config.vocab_size)

    attention_mask = random_attention_mask([batch_size, seq_len])
    label_ids = ids_tensor([batch_size, seq_len], fx_model.config.vocab_size)

    fx_inputs = {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
    }

    if pt_model.config.is_encoder_decoder:
        decoder_input_ids = shift_tokens_right(input_ids=label_ids, pad_token_id=fx_model.config.pad_token_id, decoder_start_token_id=fx_model.config.decoder_start_token_id)
        fx_inputs["decoder_input_ids"] = decoder_input_ids

    pt_inputs = {k: torch.tensor(v.tolist()) for k, v in fx_inputs.items()}
    pt_inputs["labels"] = torch.tensor(label_ids.tolist())

    fx_outputs = fx_model(**fx_inputs)
    fx_logits = fx_outputs.logits

    pt_outputs = pt_model(**pt_inputs)
    pt_logits = pt_outputs.logits
    pt_loss = pt_outputs.loss

    print("--------------------------Checking logits match--------------------------")
    print(f"Flax logits shape: {fx_logits.shape}, PyTorch logits shape: {pt_logits.shape}")
    assert_almost_equals(fx_logits, pt_logits.detach().numpy())

    def fx_train_step(fx_model, batch):
        def compute_loss(params):
            label_ids = batch.pop('label_ids')
            logits = fx_model(**batch, params=params).logits
            vocab_size = logits.shape[-1]
            targets = onehot(label_ids, vocab_size)
            loss = optax.softmax_cross_entropy(logits, targets)
            return loss.mean()

        grad_fn = jax.value_and_grad(compute_loss)
        loss, grad = grad_fn(fx_model.params)
        return loss, grad

    fx_inputs["label_ids"] = label_ids

    fx_loss, fx_grad = fx_train_step(fx_model, fx_inputs)

    print("--------------------------Checking losses match--------------------------")
    print(f"Flax loss: {fx_loss}, PyTorch loss: {pt_loss}")
    assert_almost_equals(fx_loss, pt_loss.detach().numpy())

    pt_loss.backward()

    pt_grad_dict = {k: v.grad for k, v in pt_model.named_parameters()}
    missing_grads = [k for k in pt_model.state_dict().keys() if k not in pt_grad_dict]

    missing_keys, unexpected_keys = pt_model.load_state_dict(pt_grad_dict, strict=False)

    assert missing_grads == missing_keys, f"Error with either grads {missing_keys} or keys {unexpected_keys}"

    with tempfile.TemporaryDirectory() as tmpdirname:
        pt_model.save_pretrained(tmpdirname)
        pt_grad_model_to_fx = flax_model_cls.from_pretrained(tmpdirname, from_pt=True)

    pt_grad_to_fx = pt_grad_model_to_fx.params
    fx_grad = flatten_dict(fx_grad)
    pt_grad_to_fx = flatten_dict(pt_grad_to_fx)
    print("--------------------------Checking gradients match--------------------------")
    results_fail_rel, results_correct_rel, results_fail, results_correct = assert_dict_equal(fx_grad, pt_grad_to_fx)

    if len(results_fail) == 0:
        print("✅ All grads pass")
    else:
        print("\n".join(results_fail))

    print("--------------------------Checking rel gradients match--------------------------")

    if len(results_fail_rel) == 0:
        print("✅ All rel grads pass")
    else:
        print("\n".join(results_fail_rel))


def main():
    model_id = sys.argv[1]
    pt_architecture_name = sys.argv[2]
    compare_grads(model_id, pt_architecture_name)


if __name__ == "__main__":
    main()