"""
Scalable, Detailed and Mask-free Universal Photometric Stereo Network (CVPR2023)
# Copyright (c) 2023 Satoshi Ikehata
# All rights reserved.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
import math
import torch.nn.init as init

class MultiHeadAttentionBlock(nn.Module):
    def __init__(self, dim_in, dim_out, num_heads, ln=False, attention_dropout = 0.1, dim_feedforward = 512,  
                 q_bucket_size = 1024, k_bucket_size = 2048, attn_mode = 'Normal'):
        super(MultiHeadAttentionBlock, self).__init__()
        if attn_mode == 'Efficient':
            self.q_bucket_size = q_bucket_size
            self.k_bucket_size = k_bucket_size
        self.attn_mode = attn_mode
        self.dim_V = dim_out
        self.dim_Q = dim_in
        self.dim_K = dim_in
        self.num_heads = num_heads
        self.fc_q = nn.Linear(self.dim_Q, self.dim_V, bias=False) # dimin -> dimhidden
        self.fc_k = nn.Linear(self.dim_K, self.dim_V, bias=False) # dimin -> dimhidden
        self.fc_v = nn.Linear(self.dim_K, self.dim_V, bias=False) # dimhidden -> dim
        if ln:
            self.ln0 = nn.LayerNorm(self.dim_Q)
            self.ln1 = nn.LayerNorm(self.dim_V)

        self.dropout_attn = nn.Dropout(attention_dropout)
        self.fc_o1 = nn.Linear(self.dim_V, dim_feedforward, bias=False)
        self.fc_o2 = nn.Linear(dim_feedforward, self.dim_V, bias=False)
        self.dropout1 = nn.Dropout(attention_dropout)
        self.dropout2 = nn.Dropout(attention_dropout)

        # memory efficient attention related parameters
        # can be overriden on forward
        self.q_bucket_size = q_bucket_size
        self.k_bucket_size = k_bucket_size
     

    # memory efficient attention
    def summarize_qkv_chunk(self, q, k, v):
        weight = torch.einsum('b h i d, b h j d -> b h i j', q, k)
        weight_max = weight.amax(dim = -1, keepdim = True).detach()
        weight = weight - weight_max

        exp_weight = self.dropout_attn(weight.exp()) # attention_dropout
        weighted_value = torch.einsum('b h i j, b h j d -> b h i d', exp_weight, v)
        return exp_weight.sum(dim = -1), weighted_value, rearrange(weight_max, '... 1 -> ...')

    def memory_efficient_attention(
        self,
        q, k, v,
        q_bucket_size = 512,
        k_bucket_size = 1024,
        eps = 1e-8,
    ):
        scale = q.shape[-1] ** -0.5
        q = q * scale

        summarize_qkv_fn = self.summarize_qkv_chunk

        # chunk all the inputs
        q_chunks = q.split(q_bucket_size, dim = -2)
        k_chunks = k.split(k_bucket_size, dim = -2)
        v_chunks = v.split(k_bucket_size, dim = -2)

        # loop through all chunks and accumulate
        values = []
        weights = []
        for q_chunk in q_chunks:
            exp_weights = []
            weighted_values = []
            weight_maxes = []  
            for (k_chunk, v_chunk) in zip(k_chunks, v_chunks):
                exp_weight_chunk, weighted_value_chunk, weight_max_chunk = summarize_qkv_fn(
                    q_chunk,
                    k_chunk,
                    v_chunk
                )
                exp_weights.append(exp_weight_chunk)
                weighted_values.append(weighted_value_chunk)
                weight_maxes.append(weight_max_chunk)
            weight_maxes = torch.stack(weight_maxes, dim = -1)
            weighted_values = torch.stack(weighted_values, dim = -1)
            exp_weights = torch.stack(exp_weights, dim = -1)

            global_max = weight_maxes.amax(dim = -1, keepdim = True)
            renorm_factor = (weight_maxes - global_max).exp().detach()

            exp_weights = exp_weights * renorm_factor
            weighted_values = weighted_values * rearrange(renorm_factor, '... c -> ... 1 c')         

            all_values = weighted_values.sum(dim = -1)
            all_weights = exp_weights.sum(dim = -1)
            values.append(all_values)
            weights.append(all_weights)    
        values = torch.cat(values, dim=2)
        weights = torch.cat(weights, dim=2)
        # (rearrange(weights, '... -> ... 1') 
        normalized_values = values / (rearrange(weights, '... -> ... 1')  + eps)
        return normalized_values


    def forward(
        self,
        x,y,
    ):
        
        x = x if getattr(self, 'ln0', None) is None else self.ln0(x) # pre-normalization
        Q = self.fc_q(x) # input_dim -> embed dim       
        K, V = self.fc_k(y), self.fc_v(y) # input_dim -> embed dim
        dim_split = self.dim_V // self.num_heads # multi-head attention
        if self.attn_mode == 'Efficient':
            q_bucket_size = self.q_bucket_size
            k_bucket_size = self.k_bucket_size
            
            Q_ = torch.stack(Q.split(int(dim_split), 2), 1)
            K_ = torch.stack(K.split(int(dim_split), 2), 1)
            V_ = torch.stack(V.split(int(dim_split), 2), 1) 
            
            A = self.memory_efficient_attention(Q_, K_, V_, q_bucket_size = q_bucket_size, k_bucket_size = k_bucket_size)
            A = A.reshape(-1, A.shape[2], A.shape[3])
            Q_ = Q_.reshape(-1, Q_.shape[2], Q_.shape[3])
            O = torch.cat((Q_ + A).split(Q.size(0), 0), 2)
        else: # Basic
            Q_ = torch.cat(Q.split(int(dim_split), 2), 0)
            K_ = torch.cat(K.split(int(dim_split), 2), 0)
            V_ = torch.cat(V.split(int(dim_split), 2), 0)
            A = self.dropout_attn(torch.softmax(Q_.bmm(K_.transpose(1,2))/math.sqrt(self.dim_V), 2)) # this may not be correct due to mult-head attention
            A =  A.bmm(V_) # A(Q, K, V) attention_output
            O = torch.cat((Q_ + A).split(Q.size(0), 0), 2)

        
        O_ = O if getattr(self, 'ln1', None) is None else self.ln1(O)
        O = O + self.dropout2(self.fc_o2(self.dropout1(F.gelu(self.fc_o1(O_))))) 
        return O


class SAB(nn.Module): # self attention block
    def __init__(self, dim_in, dim_out, num_heads=4, ln=False, attention_dropout = 0.1, dim_feedforward = 512, attn_mode = 'Normal'):
        super(SAB, self).__init__()
        self.mab = MultiHeadAttentionBlock(dim_in, dim_out, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode)
    def forward(self, X):
        return self.mab(X, X)

class CAB(nn.Module): # cross attention block
    def __init__(self, dim_in, dim_out, num_heads=4, ln=False, attention_dropout = 0.1, dim_feedforward = 512, attn_mode = 'Normal'):
        super(CAB, self).__init__()
        self.mab = MultiHeadAttentionBlock(dim_in, dim_out, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode)
    def forward(self, q, kv):       
        return self.mab(q, kv)
    
class PMA(nn.Module):
    def __init__(self, dim, num_heads, num_seeds, ln=False, attn_mode='Normal'):
        super(PMA, self).__init__()
        self.S = nn.Parameter(torch.Tensor(1, num_seeds, dim))
        init.xavier_uniform_(self.S)
        self.mab = MultiHeadAttentionBlock(dim, dim, num_heads, ln=ln, attn_mode=attn_mode)

    def forward(self, X):
        return self.mab(self.S.repeat(X.size(0), 1, 1), X)

class CommunicationBlock(nn.Module):
    def __init__(self, dim_input, num_enc_sab = 3, dim_hidden=384, dim_feedforward = 1024, num_heads=8, ln=False, attention_dropout=0.1, use_efficient_attention=False):
        super(CommunicationBlock, self).__init__()

        if use_efficient_attention:
            attn_mode = 'Efficient'
        else:
            attn_mode = 'Normal'
        self.dim_hidden = dim_hidden
        modules_enc = []
        modules_enc.append(SAB(dim_input, dim_hidden, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode))
        for k in range(num_enc_sab):
            modules_enc.append(SAB(dim_hidden, dim_hidden, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode))
        self.enc = nn.Sequential(*modules_enc)

    def forward(self, x):
        x = self.enc(x)
        return x

class CrossAttentionBlock(nn.Module):
    def __init__(self, dim_input, num_enc_sab = 3, dim_hidden=384, dim_feedforward = 1024, num_heads=8, ln=False, attention_dropout=0.1, use_efficient_attention=False):
        super(CrossAttentionBlock, self).__init__()

        if use_efficient_attention:
            attn_mode = 'Efficient'
        else:
            attn_mode = 'Normal'
        self.dim_hidden = dim_hidden
        modules_enc = []
        modules_enc.append(CAB(dim_input, dim_hidden, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode))
        for k in range(num_enc_sab):
            modules_enc.append(CAB(dim_hidden, dim_hidden, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode))
        self.layers = nn.ModuleList(modules_enc)
    def forward(self, q, kv):
        for k in range(len(self.layers)):
            q = self.layers[k](q, kv) # x[0] query, x[1] key value       
        return q # the output token length is len(x)
    
class AggregationBlock(nn.Module):
    def __init__(self, dim_input, num_enc_sab = 3, num_outputs = 1, dim_hidden=384, dim_feedforward = 1024, num_heads=8, ln=False, attention_dropout=0.1, use_efficient_attention=False):
        super(AggregationBlock, self).__init__()

        self.num_outputs = num_outputs
        self.dim_hidden = dim_hidden

        if use_efficient_attention:
            attn_mode = 'Efficient'
        else:
            attn_mode = 'Normal'

        modules_enc = []
        modules_enc.append(SAB(dim_input, dim_hidden, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode))
        for k in range(num_enc_sab):
            modules_enc.append(SAB(dim_hidden, dim_hidden, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode))
        self.enc = nn.Sequential(*modules_enc)
        modules_dec = []
        modules_dec.append(PMA(dim_hidden, num_heads, num_outputs, attn_mode=attn_mode)) # after the PMA we should not put drop out
        self.dec = nn.Sequential(*modules_dec)
        
    def forward(self, x):
        x = self.enc(x)
        x = self.dec(x)
        x = x.view(-1, self.num_outputs * self.dim_hidden)
        return x