from model import (
    Transformer,
    ModelArgs,
    MusicalPositionEmbedTransformer,
)
from model_onnx import (
    Transformer as TransformerONNX,
    TransformerOneStep as TransformerOneStepONNX,
)
from data import Vocab
from data_gen import Batch
import math
import torch
import pytest
import random
from test_fixtures import LittleTrainer

s0 = [
    int(
        (
            2 * (i + 1)
            + round(math.sqrt(2 * (i + 1)))
            - round(math.sqrt(2 * (i + 1))) ** 2
        )
        / 2
        + 1
    )
    for i in range(100)
]

s1 = [2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 0, 0, 0]

s2 = [random.randint(2, 127) for i in range(128)]
s3 = [random.randint(2, 10) for i in range(48)]


# training on a batch that is empty (all pads) should not crash
@pytest.mark.parametrize("enable_flash", [True, False])
def test_empty_batch(enable_flash):
    smallargs = ModelArgs(256, 4, 4, 128, dropout=0.1, enable_flash=enable_flash)
    t = Transformer(smallargs).cuda()
    trainer = LittleTrainer(t)

    data = torch.zeros(1, 10, dtype=torch.long)
    batch_size = 16
    input = data[:, :-1].repeat(batch_size, 1).cuda()
    target = data[:, 1:].repeat(batch_size, 1).cuda()

    loss, last_pred = trainer.step(input, target)

    assert loss == 0
    for param in t.parameters():
        assert (param.grad != 0).view(-1).sum() == 0


@pytest.mark.parametrize("sequence", [s0, s1, s2, s3])
def test_simple_sequences(sequence):
    smallargs = ModelArgs(256, 4, 4, 128, dropout=0.1)
    t = Transformer(smallargs).cuda()
    trainer = LittleTrainer(t)

    seq_len = min(512, len(sequence))
    data = torch.zeros(1, seq_len, dtype=torch.long)
    data[0, 0] = 1
    data[0, 1:seq_len] = torch.tensor(sequence[: seq_len - 1], dtype=torch.long)

    input = data[:, :-1]
    target = data[:, 1:]

    batch_size = 16
    input = input.repeat(batch_size, 1).cuda()
    target = target.repeat(batch_size, 1).cuda()

    for i in range(100):
        _, last_pred = trainer.step(input, target)

    last_res = torch.argmax(last_pred, dim=-1)
    last_res[target == 0] = 0

    last_res = last_res.cpu()
    target = target.cpu()

    match_count = 0
    for i in range(batch_size):
        if torch.equal(last_res[i], target[i]):
            match_count += 1

    assert match_count / batch_size > 0.9

    # and using cache
    t.eval()
    with torch.no_grad():
        for i in range(seq_len - 1):
            pred = t.forward(input[0, i].view(1, 1), start_pos=i)
            if target[0, i].item() != 0:
                assert torch.argmax(pred, dim=-1) == target[0, i]


@pytest.mark.parametrize("enable_flash", [True, False])
@pytest.mark.parametrize("enable_cross", [True, False])
def test_causal(enable_flash, enable_cross):
    smallargs = ModelArgs(
        256,
        4,
        4,
        128,
        dropout=0.1,
        enable_flash=enable_flash,
        enable_cross_attention=enable_cross,
    )
    t = Transformer(smallargs).cuda()
    trainer = LittleTrainer(t)

    data = torch.arange(1, 21, dtype=torch.long).view(1, -1).repeat(8, 1)
    data[:, 0] = 1
    for i in range(data.shape[0]):
        data[i, 3:] = i + 2
    data = data.cuda()
    input = data[:, :-1]
    target = data[:, 1:]

    if enable_cross:
        encoder_out = (
            torch.normal(0, 0.2, (1, 5, smallargs.cross_attention_embedding_dim))
            .repeat(data.shape[0], 1, 1)
            .cuda()
        )
    else:
        encoder_out = None

    for i in range(200):
        _, last_pred = trainer.step(input, target, encoder_out=encoder_out)

    # no good chance of guessing the rest given only the first 3
    t.eval()

    pred_first3 = t.forward(input[:, :3], start_pos=0, encoder_out=encoder_out)[0, 2]
    assert (
        pred_first3[2:10].min() > pred_first3[0:2].max()
        and pred_first3[2:10].min() > pred_first3[10:].max()
    )
    assert (pred_first3[2:10] - pred_first3[2]).abs().max() < 1.0

    # no good chance of guessing index 2 given the entire sequence (i.e. model is causal)
    pred_full = t.forward(input, start_pos=0, encoder_out=encoder_out)[0, 2]

    assert pred_first3.allclose(pred_full)


@pytest.mark.parametrize("enable_flash", [True, False])
def test_overfit(enable_flash):
    smallargs = ModelArgs(256, 4, 4, 128, enable_flash=enable_flash)
    t = Transformer(smallargs)
    trainer = LittleTrainer(t)

    seq_len = 128
    batch_size = 16
    data = torch.zeros(batch_size, seq_len, dtype=torch.long)
    data[:, 0] = 1  # start
    for i in range(batch_size):
        data[i, 1] = i + 2  # sequence id
        data[i, 2:] = torch.randint(2, 127, (seq_len - 2,), dtype=torch.long)

    input = data[:, :-1]
    target = data[:, 1:]

    t.cuda()
    input = input.cuda()
    target = target.cuda()

    for i in range(300):
        trainer.step(input, target)

    total = 0
    matches = 0
    for i in range(2, seq_len):
        sampler_in = input[:, :i]
        pred = torch.argmax(
            t.forward(sampler_in, 0)[:, -1],
            dim=-1,
        )
        matches += torch.count_nonzero(pred == target[:, i - 1]).item()
        total += batch_size

    print(matches / total)
    assert matches / total > 0.98


def test_aux_embeddings():
    smallargs = ModelArgs(256, 4, 4, 128, dropout=0.1)
    v = Vocab(1, 2, 1, 2, 1, 2, 1, 1, 2, 1, 128, 1, 128)
    t = MusicalPositionEmbedTransformer(v, smallargs)
    trainer = LittleTrainer(t)

    batch_size = 2
    data = torch.zeros(batch_size, len(s0) + 1, 9, dtype=torch.long)
    data[:, :, 0] = 1
    data[0, 1:, 3] = torch.tensor(s0, dtype=torch.long)
    data[0, :, 7] = 1
    data[1, 1:, 3] = torch.tensor(s0, dtype=torch.long)
    data[1, len(s0), 3] = 99  # cruel
    data[1, :, 7] = 2

    t.cuda()

    batch = Batch(v, data, "cuda")

    for i in range(200):
        _, last_pred = trainer.step(batch)

    last_res = torch.argmax(last_pred, dim=-1)
    last_res[batch.tgt[:, :, 0] == 0] = 0

    last_res = last_res.cpu()
    target = batch.tgt_y.cpu()

    print(last_res)

    match_count = 0
    for i in range(batch_size):
        if torch.equal(last_res[i], target[i]):
            match_count += 1

    assert match_count == 2


@pytest.mark.parametrize("enable_flash", [True, False])
def test_cross_attention_identity(enable_flash):
    smallargs = ModelArgs(
        256,
        4,
        4,
        128,
        enable_flash=enable_flash,
        enable_cross_attention=True,
        cache=False,
        cross_attention_embedding_dim=256,
        dropout=0.1,
    )
    t = Transformer(smallargs)
    trainer = LittleTrainer(t)

    seq_len = 128
    batch_size = 32
    data = torch.randint(2, 3, (batch_size, seq_len), dtype=torch.long)

    encoder_output = torch.normal(0, 0.2, (batch_size, seq_len - 1, smallargs.dim))
    input = data[:, :-1]
    target = data[:, 1:]

    t.cuda()
    encoder_output = encoder_output.cuda()
    input = input.cuda()
    target = target.cuda()

    for i in range(300):
        trainer.step(input, target, encoder_out=encoder_output)

    t.eval()
    total = 0
    matches = 0
    with torch.no_grad():
        for i in range(2, seq_len):
            sampler_encoder_out = encoder_output[:, :i]
            sampler_in = input[:, :i]
            pred = torch.argmax(
                t.forward(sampler_in, start_pos=0, encoder_out=sampler_encoder_out)[
                    :, -1
                ],
                dim=-1,
            )
            matches += torch.count_nonzero(pred == target[:, i - 1]).item()
            total += batch_size

    print(matches / total)
    assert matches / total > 0.98


@pytest.mark.parametrize("enable_flash", [True, False])
def test_cross_id_one_shot(enable_flash):
    smallargs = ModelArgs(
        256,
        4,
        4,
        128,
        enable_flash=enable_flash,
        enable_cross_attention=True,
        cross_attention_embedding_dim=256,
        dropout=0.1,
    )
    t = Transformer(smallargs)
    trainer = LittleTrainer(t)

    seq_len = 2
    batch_size = 32
    data = torch.zeros(batch_size, seq_len, dtype=torch.long)
    data[:, 0] = 1  # start
    for i in range(batch_size):
        data[:, 1] = i + 2

    encoder_output = torch.normal(0, 0.2, (batch_size, seq_len - 1, smallargs.dim))
    input = data[:, :-1]
    target = data[:, 1:]

    t.cuda()
    encoder_output = encoder_output.cuda()
    input = input.cuda()
    target = target.cuda()

    for i in range(300):
        trainer.step(input, target, encoder_out=encoder_output)

    t.eval()
    total = 0
    matches = 0
    with torch.no_grad():
        for i in range(batch_size):
            sampler_encoder_out = encoder_output[i].unsqueeze(0)
            sampler_in = torch.ones(1, 1, dtype=torch.long).cuda()
            pred = torch.argmax(
                t.forward(sampler_in, start_pos=0, encoder_out=sampler_encoder_out)[
                    :, -1
                ],
                dim=-1,
            )
            matches += torch.count_nonzero(pred == target[i, 0]).item()
            total += 1

    print(matches / total)
    assert matches / total > 0.98


def build_cross_cartesian_ctx(enable_flash, enable_cross):
    smallargs = ModelArgs(
        256,
        4,
        4,
        128,
        enable_flash=enable_flash,
        enable_cross_attention=enable_cross,
        cross_attention_embedding_dim=256,
    )
    t = Transformer(smallargs)
    trainer = LittleTrainer(t)

    seq_len = 5
    batch_size = 8 * 8
    data = torch.zeros(batch_size, seq_len, dtype=torch.long)
    data[:, :] = 1
    for i in range(8):
        for j in range(8):
            data[i * 8 + j, 0] = i + 2
            data[i * 8 + j, -1] = i * 8 + j + 2

    normal_block = torch.normal(0, 1, (8, seq_len - 1, smallargs.dim))
    encoder_output = normal_block.repeat(8, 1, 1)
    print(encoder_output.shape)
    encoder_train_mask = torch.ones(batch_size, seq_len - 1, dtype=torch.bool)
    encoder_train_mask[:, -1] = False
    print(encoder_train_mask.shape)

    input = data[:, :-1]
    target = data[:, 1:]

    print(data)
    print(encoder_output)

    t.cuda()
    encoder_output = encoder_output.cuda()
    encoder_train_mask = encoder_train_mask.cuda()
    input = input.cuda()
    target = target.cuda()

    for i in range(300):
        trainer.step(
            input,
            target,
            encoder_out=encoder_output,
            encoder_out_mask=encoder_train_mask,
        )

    t.eval()

    return batch_size, seq_len, smallargs, input, target, encoder_output, t


@pytest.mark.parametrize("enable_flash", [True, False])
@pytest.mark.parametrize("enable_cross", [True, False])
def test_cross_cartesian_remember(enable_flash, enable_cross):
    (
        batch_size,
        seq_len,
        smallargs,
        input,
        target,
        encoder_output,
        t,
    ) = build_cross_cartesian_ctx(enable_flash, enable_cross)
    print("enable_flash", enable_flash, "enable_cross", enable_cross)

    total = 0
    matches = 0
    with torch.no_grad():
        for i in range(batch_size):
            pred = torch.argmax(
                t.forward(
                    input[i, :-1].unsqueeze(0),
                    start_pos=0,
                    encoder_out=encoder_output[i, :-1].unsqueeze(0),
                )[:, -1],
                dim=-1,
            )
            assert pred[0].item() == 1
            pred = torch.argmax(
                t.forward(
                    input[i, -1:].unsqueeze(0),
                    start_pos=seq_len - 2,
                    encoder_out=encoder_output[i, :-1].unsqueeze(0),
                )[:, -1],
                dim=-1,
            )
            print(i, pred, target[i, -1])
            matches += torch.count_nonzero(pred[0] == target[i, -1]).item()
            total += 1

    print(matches / total)
    assert matches / total > 0.90


def test_onnx():
    (
        batch_size,
        seq_len,
        smallargs,
        input,
        target,
        encoder_output,
        t,
    ) = build_cross_cartesian_ctx(False, True)
    t_onnx = TransformerONNX(smallargs)
    t_onnx.load_state_dict(t.state_dict(), strict=False)
    t_onnx.eval()
    t_onnx.cuda()
    t_onnx_one = TransformerOneStepONNX(smallargs)
    t_onnx_one.load_state_dict(t.state_dict(), strict=False)
    t_onnx_one.eval()
    t_onnx_one.cuda()
    total = 0
    matches = 0
    with torch.no_grad():
        for i in range(batch_size):
            pred = t.forward(
                input[i, :-2].unsqueeze(0),
                start_pos=0,
                encoder_out=encoder_output[i, :-1].unsqueeze(0),
            )
            pred_onnx, a_xks, a_xvs, *fwd_rest = t_onnx.forward(
                input[i, :-2].unsqueeze(0).repeat(2, 1),
                aux_embeddings=torch.zeros(
                    2,
                    seq_len - 3,
                    smallargs.dim,
                    dtype=torch.float,
                    device=next(t.parameters()).device,
                ),
                encoder_out=encoder_output[i, :-1].unsqueeze(0).repeat(2, 1, 1),
                encoder_valid=torch.ones(
                    2, encoder_output.shape[1] - 1, dtype=torch.bool, device="cuda"
                ),
            )

            c_xks, c_xvs, embeddings = fwd_rest

            assert torch.allclose(pred_onnx[0], pred_onnx[1])
            assert torch.allclose(embeddings[0], embeddings[1])
            pred_onnx = pred_onnx[0]
            embeddings = embeddings[0]

            assert torch.allclose(pred, pred_onnx, atol=1e-5)

            a_xks = a_xks.transpose(0, 1)
            a_xvs = a_xvs.transpose(0, 1)

            for j in range(2):
                pred = t.forward(
                    input[i, j - 2].view(1, 1),
                    start_pos=seq_len - 3 + j,
                    encoder_out=encoder_output[i, :-1].unsqueeze(0),
                )
                pred_onnx, a_xk_news, a_xv_news = t_onnx_one.forward(
                    input[i, j - 2].view(-1).unsqueeze(0).repeat(2, 1),
                    aux_embeddings=torch.zeros(
                        2,
                        1,
                        smallargs.dim,
                        dtype=torch.float,
                        device=next(t.parameters()).device,
                    ),
                    a_xks=a_xks,
                    a_xvs=a_xvs,
                    c_xks=c_xks,
                    c_xvs=c_xvs,
                    encoder_valid=torch.ones(
                        2, encoder_output.shape[1] - 1, dtype=torch.bool, device="cuda"
                    ),
                )

                a_xks = torch.cat([a_xks, a_xk_news], dim=0)
                a_xvs = torch.cat([a_xvs, a_xv_news], dim=0)

                assert torch.allclose(pred[0], pred_onnx[0], atol=1e-5)
                assert torch.allclose(pred[0], pred_onnx[1], atol=1e-5)

            matches += torch.count_nonzero(
                torch.argmax(pred_onnx[-1]) == target[i, -1]
            ).item()
            total += 1

    print(matches / total)
    assert matches / total > 0.90