
import torch
import torch.nn as nn
import math
import json
from torch.utils.checkpoint import checkpoint
import sys
import os

curr_path = os.path.dirname(os.path.realpath(__file__))
sys.path.append(curr_path)

class MRA_headAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.drop_attn = torch.nn.Dropout(p = config["dropout_prob"])
        self.head_dim = config["head_dim"]

    def forward(self, Q, K, V, mask):
        dot = torch.matmul(Q, torch.transpose(K, -2, -1))
        dot = dot / math.sqrt(self.head_dim)
        dot = dot - 1e6 * (1 - mask[:, None, None, :])

        attn = nn.functional.softmax(dot, dim = -1)
        attn = self.drop_attn(attn)

        X = torch.matmul(attn, V)
        return X

class Attention(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.dim = config["dim"]
        self.head_dim = config["head_dim"]
        self.num_head = config["num_head"]
        self.group_by_list = config["group_by_list"]

        self.qkv = nn.ModuleList([nn.Linear(self.dim, 3 * self.dim // self.num_head) for i in range(self.num_head)])

        self.attn = MRA_headAttention(config)

        self.ff = nn.Linear(self.num_head * self.head_dim, self.dim)

    def forward(self, X, mask):
        batch_size, seq_len, dim = X.shape
        X = X * mask[:, :, None]

        attn_out = torch.empty(X.shape[0], self.num_head, X.shape[1], X.shape[2] // self.num_head, device=X.device)
        for h in range(self.num_head):
            # Down sampling mask and input
            mask_ = torch.clip(mask.reshape(X.shape[0], X.shape[1] // self.group_by_list[h], self.group_by_list[h]).sum(dim=-1), min=0, max=1)
            token_count = mask.reshape(batch_size, X.shape[1] // self.group_by_list[h], self.group_by_list[h]).sum(dim = -1)
            X_ = X.reshape(batch_size, X.shape[1] // self.group_by_list[h], self.group_by_list[h], dim).sum(dim = -2) / (token_count[:, :, None] + 1e-6)
            # Calcualte downsampled qkv
            qkv_ = self.qkv[h](X_).reshape(batch_size, seq_len // self.group_by_list[h], 3, dim // self.num_head).permute(2, 0, 1, 3).unsqueeze(2)
            # bsz, num_head, seq_len, head_dim
            q_, k_, v_ = qkv_[0], qkv_[1], qkv_[2]
            attn_out_ = self.attn(q_, k_, v_, mask_)
            attn_out[:, h, :, :] = attn_out_.squeeze(1).repeat_interleave(self.group_by_list[h], dim=1)

        attn_out = self.combine_heads(attn_out)
        out = self.ff(attn_out)

        return out

    def combine_heads(self, X):
        X = X.transpose(1, 2)
        X = X.reshape(X.size(0), X.size(1), self.num_head * self.head_dim)
        return X

    def split_heads(self, X):
        X = X.reshape(X.size(0), X.size(1), self.num_head, self.head_dim)
        X = X.transpose(1, 2)
        return X

class Block(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.norm1 = nn.LayerNorm(config["dim"])
        self.mha = Attention(config)
        self.dropout1 = torch.nn.Dropout(p = config["dropout_prob"])
        self.norm2 = nn.LayerNorm(config["dim"])

        self.mlpblock = nn.Sequential(
            nn.Linear(config["dim"], config["hidden_dim"]),
            nn.GELU(),
            torch.nn.Dropout(p = config["dropout_prob"]),
            nn.Linear(config["hidden_dim"], config["dim"]),
            torch.nn.Dropout(p = config["dropout_prob"])
        )

    def forward(self, X, mask):
        X = self.dropout1(self.mha(self.norm1(X), mask)) + X
        X = self.mlpblock(self.norm2(X)) + X
        return X

class Backbone(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.num_layers = config["num_layers"]
        self.shared_weight = config["shared_weight"]

        if self.shared_weight:
            self.encoder = Block(config)
        else:
            self.encoders = nn.ModuleList([Block(config) for _ in range(self.num_layers)])

        self.norm = nn.LayerNorm(config["dim"])

    def forward(self, X, mask):

        if self.shared_weight:
            for _ in range(self.num_layers):
                X = self.encoder(X, mask)
        else:
            for encoder in self.encoders:
                X = encoder(X, mask)

        X = self.norm(X) * mask[:, :, None]

        return X
