# Transformers from Scratch using "Attention is All You Need" paper
# Modelling Scaled Dot-Product Attention, Multi-Head Attention, Position-wise Feed-Forward Networks.

# Import Modules
import matplotlib.pyplot as plt
import torch.nn.functional as F
import torch.nn as nn
import torch
import numpy as np
import math

# Making Single and Multi-Head Attention modules from scratch using Pure PyTorch

# Initialise the seed for reproducibility
seed = 42
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)

# Self-Attention Mechanism: Single Head
embdim = 256  # D
headdim = 64  # Internal D
tokens = torch.randn(1, 5, embdim)  # batch, tokens, embedding

# Defining weights associates with query, key, value
Wq = torch.randn(embdim, headdim) / math.sqrt(embdim)
Wk = torch.randn(embdim, headdim) / math.sqrt(embdim)
Wv = torch.randn(embdim, embdim) / math.sqrt(embdim)

# Query, Key, Value
qis = torch.einsum("BSE,EH->BSH", tokens, Wq)  # batch x seqlen x headdim; queries, (1, 5, 64)
kis = torch.einsum("BTE,EH->BTH", tokens, Wk)  # batch x seqlen x headdim; keys
vis = torch.einsum("BTE,EF->BTF", tokens, Wv)  # batch x seqlen x embeddim; values

# Start: Testing Code
random_mat1 = torch.randn(2, 5, 4)  # BATCH, TOKENS, DIMENSIONS
random_mat2 = torch.randn(2, 5, 4)

# 2, 5, 4 * , 2, 4, 5
torch.matmul(random_mat1, random_mat2.transpose(1, 2))  # 2, 5, 5
print(qis.shape)
print(kis.shape)
# (Q) N, D * (K^T) D, N  -> N, N
# End: Testing Code


scoremat = torch.matmul(qis, kis.transpose(1, 2))  # output: batch x seqlen (Query) x seqlen (Key)
attmat = F.softmax(scoremat / math.sqrt(headdim), dim=2)  # attention matrix given.

# Output of the attention mechanism
zis = torch.einsum("BST,BTF->BSF", attmat, vis)

# We can verify the output, with scaled dot-product attention
attn_torch = F.scaled_dot_product_attention(qis, kis, vis)
assert (torch.allclose(attn_torch, zis, atol=1E-6, rtol=1E-6))  # True

# Multi-Head Attention
embdim = 768
headcnt = 12
headdim = embdim // headcnt
# print(headdim)
assert headdim * headcnt == embdim
tokens = torch.randn(1, 5, embdim)  # batch, tokens, embedding

# We use all the 256, ( 768)  ~ which is (256), (64 * 12 (heads))
Wq = torch.randn(embdim, headcnt * headdim) / math.sqrt(embdim)  # heads packed in a single dim
Wk = torch.randn(embdim, headcnt * headdim) / math.sqrt(embdim)  # heads packed in a single dim
Wv = torch.randn(embdim, headcnt * headdim) / math.sqrt(embdim)  # heads packed in a single dim

print(Wq.shape)
print(Wk.shape)
print(Wv.shape)

batch, token_num, _ = tokens.shape  # batch, tokens (n), embedding shape.
# tokens, B, N, E

# Wq,     B, E, HWeights (H * HC)
qis = torch.einsum("BSE,EH->BSH", tokens, Wq)  # Batch, N, H  ~ 1, 5, 768
kis = torch.einsum("BTE,EH->BTH", tokens, Wk)  # Batch N, H
vis = torch.einsum("BTE,EH->BTH", tokens, Wv)  # Batch, N, H
# split the single hidden dim into the heads

# Converting dimensions from (B, N, H) to (B, N, HC, HW)
# So now for each batch, for each token, for each head there are a set of weights.
qis_mh = qis.view(batch, token_num, headcnt, headdim)  # B, N, HC, HW
kis_mh = kis.view(batch, token_num, headcnt, headdim)
vis_mh = vis.view(batch, token_num, headcnt, headdim)

scoremat_mh = torch.einsum("BSHC,BTHC->BHST", qis_mh, kis_mh)  # Input: (B, N, HC, HH) & Output: (B, HC, Q, K)
print(scoremat_mh.shape)  # 1, 12, 5, 5  # Now I have 12 heads, which have given me attention matrices of shape 5x5.

# batch x headcnt x seqlen (query) x seqlen (key)

attmat_mh = F.softmax(scoremat_mh / math.sqrt(headdim), dim=-1)
zis_mh = torch.einsum("BCST,BTCH->BSCH", attmat_mh, vis_mh)  # batch x seqlen (query) x headcnt x headdim
zis = zis_mh.reshape(batch, token_num, headcnt * headdim)

# The block does not do the operation of concat and linear layer operations on this.

# We can verify the output, with Multi-Head Attention
mha = nn.MultiheadAttention(embdim, headcnt, batch_first=True, )
print(mha.in_proj_weight.shape)  # 3 * embdim x embdim
mha.in_proj_weight.data = torch.cat([Wq, Wk, Wv], dim=1).T
attn_out, attn_weights = mha(tokens, tokens, tokens, average_attn_weights=False, )

# Which is the same as attmat_mh
assert torch.allclose(attmat_mh, attn_weights, atol=1e-6, rtol=1e-6)  # True

print(attn_weights.shape)  # batch, heads, tokens, tokens.
print(attn_out.shape)

# Casual Mask from Scratch
# Calculate Casual Mask, this is described in the paper when we do not want to attend to the future tokens, in decoder.

attn_mask = torch.ones(token_num, token_num, )
attn_mask = -1E4 * torch.triu(attn_mask, 1)
print(attn_mask)
scoremat_mh_msk = torch.einsum("BSCH,BTCH->BCST", qis_mh, kis_mh)  # batch x headcnt x seqlen (query) x seqlen (key)
scoremat_mh_msk += attn_mask  # add the attn mask to the scores before SoftMax normalization
attmat_mh_msk = F.softmax(scoremat_mh_msk / math.sqrt(headdim), dim=-1)
zis_mh_msk = torch.einsum("BCST,BTCH->BSCH", attmat_mh_msk, vis_mh)  # batch x seqlen (query) x headcnt x headdim
zis_msk = zis_mh_msk.reshape(batch, token_num, headcnt * headdim)

attn_out_causal, attn_weights_causal = mha(tokens, tokens, tokens, average_attn_weights=False, attn_mask=attn_mask)

# Plotting all heads of the attention mechanism.
plt.figure()
for head in range(headcnt):
	plt.subplot(3, 4, head + 1)
	plt.imshow(attn_weights_causal[0, head].detach().numpy())
	plt.title(f"head {head}")
	plt.axis("off")
plt.show()

# Transformer Block from Scratch

# Modeling the Transformer Block from Scratch using PyTorch
# Transformer Block contains:
#     - Layer norm
#     - Skip connections
#     - Multi-head attention
#     - MLP, Feedforward net


class TransformerBlock(nn.Module):

	def __init__(self, embdim:int, headcnt, *args, dropout=0.0, **kwargs) -> None:
		super().__init__(*args, **kwargs)
		self.ln1 = nn.LayerNorm(embdim)
		self.ln2 = nn.LayerNorm(embdim)
		self.attn = nn.MultiheadAttention(embdim, headcnt, batch_first=True,)
		self.ffn = nn.Sequential(
			nn.Linear(embdim, 4 * embdim),
			nn.GELU(),
			nn.Linear(4 * embdim, embdim),
			nn.Dropout(dropout),
		)

	def forward(self, x, is_causal=True):
		"""
		Input to forward function is matrix with shape B, S, E, we can assume therefore that input and positional embeddings have been added.
		"""
		batch, token_num, hidden_dim = x.shape
		if is_causal:
			attn_mask = torch.ones(token_num, token_num,)
			attn_mask = -1E4 * torch.triu(attn_mask,1)
		else:
			attn_mask = None

		residue = x
		attn_output, attn_weights = self.attn(x, x, x, average_attn_weights=False, )
		x = residue + attn_output
		x = self.ln1(x)
		residue = x
		ffn_output = self.ffn(x)
		output = residue + ffn_output
		return output



if __name__ == "__main__":
	# Testing the Transformer Block
	print("Testing the Transformer Block")
	transformer_block = TransformerBlock(embdim, headcnt)
	tokens = torch.randn(1, 5, embdim)
	output = transformer_block(tokens)
	print(output.shape)