from torch import nn
import numpy as np
import torch.nn.functional as F
from torch.nn.utils import weight_norm
import math
import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
symbol_length = 73

class GlowTTS(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()
        
    def forward(self, text, text_len, mel=None, mel_len=None, inference=False, noise_scale=1., length_scale=1.):
        """
        =====inputs=====
        text: (B, T)
        text_len: (B) list
        mel: (B, 80, F)
        mel_len: (B) list
        inference: True/False
        =====outputs=====
        (tuple) (z, z_mean, z_log_std, log_det, z_mask)
            z(training) or y(inference): (B, 80, F) | z: latent representation, y: mel-spectrogram
            z_mean: (B, 80, F)
            z_log_std: (B, 80, F)
            log_det: (B) or None
            z_mask: (B, 1, F)
        (tuple) (x_mean, x_log_std, x_mask)
            x_mean: (B, 80, T)
            x_log_std: (B, 80, T)
            x_mask: (B, 1, T)
        (tuple) (attention_alignment, x_log_dur, log_d)
            attention_alignment: (B, T, F)
            x_log_dur: (B, 1, T) | 추측한 duration의 log scale
            log_d: (B, 1, T) | 적절하다고 추측한 alignment에서의 duration의 log scale
        """
        x_mean, x_log_std, x_log_dur, x_mask = self.encoder(text, text_len)
            # x_std, x_dur 에 log를 붙인 이유는, 논문 저자의 구현에서는 log가 취해진 값으로 간주하기 때문이다.
        y, y_len = mel, mel_len
        
        if not inference: # training
            y_max_len = y.size(2)
        else: # inference
            dur = torch.exp(x_log_dur) * x_mask  * length_scale # (B, 1, T)
            ceil_dur = torch.ceil(dur) # (B, 1, T)
            y_len = torch.clamp_min(torch.sum(ceil_dur, [1, 2]), 1).long() # (B)
                # ceil_dur을 [1, 2] 축에 대해 sum한 뒤 최솟값이 1이상이 되도록 설정. 정수 long 타입으로 반환한다.
            y_max_len = None
        
        # preprocessing
        if y_max_len is not None:
            y_max_len = (y_max_len // 2) * 2 # 홀수면 1을 빼서 짝수로 만든다.
            y = y[:, :, :y_max_len] # y_max_len에 맞게 y를 조정
            y_len = (y_len // 2) * 2 # y_len이 홀수이면 1을 빼서 짝수로 만든다.
        
        # make the z_mask
        B = len(y_len)
        temp_max = max(y_len)
        z_mask = torch.zeros((B, 1, temp_max), dtype=torch.bool).to(device) # (B, 1, F)
        for idx, length in enumerate(y_len):
            z_mask[idx, :, :length] = True
        
        # make the attention_mask
        attention_mask = x_mask.unsqueeze(3) * z_mask.unsqueeze(2) # (B, 1, T, 1) * (B, 1, 1, F) = (B, 1, T, F)
            # 주의: Encoder의 attention_mask와는 다른 mask임.
        
        if not inference: # training
            z, log_det = self.decoder(y, z_mask, reverse=False)
            with torch.no_grad():
                x_std_squared_root = torch.exp(-2 * x_log_std) # (B, 80, T)
                logp1 = torch.sum(-0.5 * math.log(2 * math.pi) - x_log_std, [1]).unsqueeze(-1) # [(B, T, F)
                logp2 = torch.matmul(x_std_squared_root.transpose(1, 2), -0.5 * (z ** 2)) # [(B, T, 80) * (B, 80, F) = (B, T, F)
                logp3 = torch.matmul((x_mean * x_std_squared_root).transpose(1,2), z) # (B, T, 80) * (B, 80, F) = (B, T, F)
                logp4 = torch.sum(-0.5 * (x_mean ** 2) * x_std_squared_root, [1]).unsqueeze(-1) # (B, T, F)
                logp = logp1 + logp2 + logp3 + logp4 # (B, T, F)
                """
                logp는 normal distribution N(x_mean, x_std)의 maximum log-likelihood이다.
                sum(log(N(z;x_mean, x_std)))를 정규분포 식을 이용하여 분배법칙으로 풀어내면 위와 같은 식이 도출된다.
                """
                attention_alignment = maximum_path(logp, attention_mask.squeeze(1)).detach() # alignment (B, T, F)
             
            z_mean = torch.matmul(attention_alignment.transpose(1, 2), x_mean.transpose(1, 2)) # (B, F, T) * (B, T, 80) -> (B, F, 80)
            z_mean = z_mean.transpose(1, 2) # (B, 80, F)
            z_log_std = torch.matmul(attention_alignment.transpose(1, 2), x_log_std.transpose(1, 2)) # (B, F, T) * (B, T, 80) -> (B, F, 80)
            z_log_std = z_log_std.transpose(1, 2) # (B, 80, F)
            log_d = torch.log(1e-8 + torch.sum(attention_alignment, -1)).unsqueeze(1) * x_mask # (B, 1, T) | alignment에서 형성된 duration의 log scale
            return (z, z_mean, z_log_std, log_det, z_mask), (x_mean, x_log_std, x_mask), (attention_alignment, x_log_dur, log_d)
            
        else: # inference
            # generate_path (make attention_alignment using ceil(x_dur))
            attention_alignment = generate_path(ceil_dur.squeeze(1), attention_mask.squeeze(1)) # (B, T, F)
            z_mean = torch.matmul(attention_alignment.transpose(1, 2), x_mean.transpose(1, 2)) # (B, F, T) * (B, T, 80) -> (B, F, 80)
            z_mean = z_mean.transpose(1, 2) # (B, 80, F)
            z_log_std = torch.matmul(attention_alignment.transpose(1, 2), x_log_std.transpose(1, 2)) # (B, F, T) * (B, T, 80) -> (B, F, 80)
            z_log_std = z_log_std.transpose(1, 2) # (B, 80, F)
            log_d = torch.log(1e-8 + torch.sum(attention_alignment, -1)).unsqueeze(1) * x_mask # (B, 1, T) | alignment에서 형성된 duration의 log scale
            
            z = (z_mean + torch.exp(z_log_std) * torch.randn_like(z_mean) * noise_scale) * z_mask # z(latent representation) 생성
            y, log_det = self.decoder(z, z_mask, reverse=True) # mel-spectrogram 생성
            return (y, z_mean, z_log_std, log_det, z_mask), (x_mean, x_log_std, x_mask), (attention_alignment, x_log_dur, log_d)
        
##### 아래 논문의 구현이 훨씬 빠르다. 이 논문 구현을 보고 위의 구현을 변경할 필요가 있다. #####
def maximum_path(value, mask, max_neg_val=-np.inf):
    """ Numpy-friendly version. It's about 4 times faster than torch version.
    value: [b, t_x, t_y]
    mask: [b, t_x, t_y]
    """
    value = value * mask

    device = value.device
    dtype = value.dtype
    value = value.cpu().detach().numpy()
    mask = mask.cpu().detach().numpy().astype(bool)

    b, t_x, t_y = value.shape
    direction = np.zeros(value.shape, dtype=np.int64)
    v = np.zeros((b, t_x), dtype=np.float32)
    x_range = np.arange(t_x, dtype=np.float32).reshape(1,-1)
    for j in range(t_y):
        v0 = np.pad(v, [[0,0],[1,0]], mode="constant", constant_values=max_neg_val)[:, :-1]
        v1 = v
        max_mask = (v1 >= v0)
        v_max = np.where(max_mask, v1, v0)
        direction[:, :, j] = max_mask

        index_mask = (x_range <= j)
        v = np.where(index_mask, v_max + value[:, :, j], max_neg_val)
    direction = np.where(mask, direction, 1)

    path = np.zeros(value.shape, dtype=np.float32)
    index = mask[:, :, 0].sum(1).astype(np.int64) - 1
    index_range = np.arange(b)
    for j in reversed(range(t_y)):
        path[index_range, index, j] = 1
        index = index + direction[index_range, index, j] - 1
    path = path * mask.astype(np.float32)
    path = torch.from_numpy(path).to(device=device, dtype=dtype)
    return path


def generate_path(duration, mask):
    """
    duration: [b, t_x]
    mask: [b, t_x, t_y]
    """
    device = duration.device

    b, t_x, t_y = mask.shape # (B, T, F)
    cum_duration = torch.cumsum(duration, 1) # 누적합, (B, T) 
    path = torch.zeros(b, t_x, t_y, dtype=mask.dtype).to(device=device) # (B, T, F)

    cum_duration_flat = cum_duration.view(b * t_x) # (B*T)
    path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype) # (B*T, F)
    path = path.view(b, t_x, t_y) # (B, T, F)
    path = path.to(torch.float32)
    path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:,:-1] # (B, T, F) # T의 차원 맨 앞을 -1한다.
    path = path * mask
    return path

def sequence_mask(length, max_length=None):
    if max_length is None:
        max_length = length.max()
    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
    return x.unsqueeze(0) < length.unsqueeze(1)

def convert_pad_shape(pad_shape):
    l = pad_shape[::-1] # [[0, 0], [p, p], [0, 0]]
    pad_shape = [item for sublist in l for item in sublist] # [0, 0, p, p, 0, 0]
    return pad_shape

def MAS(path, logp, T_max, F_max):
    """
    Glow-TTS의 모듈인 maximum_path의 모듈
    MAS 알고리즘을 수행하는 함수이다.
    =====inputs=====
    path: (T, F)
    logp: (T, F)
    T_max: (1)
    F_max: (1)
    =====outputs=====
    path: (T, F) | 0과 1로 구성된 alignment
    """
    neg_inf = -1e9 # negative infinity
    # forward
    for j in range(F_max):
        for i in range(max(0, T_max + j - F_max), min(T_max, j + 1)): # 평행사변형을 생각하라.
            # Q_i_j-1 (current)
            if i == j:
                Q_cur = neg_inf
            else:
                Q_cur = logp[i, j-1] # j=0이면 i도 0이므로 j-1을 사용해도 된다.
            
            # Q_i-1_j-1 (previous)
            if i==0:
                if j==0:
                    Q_prev = 0. # i=0, j=0인 경우에는 logp 값만 반영해야 한다.
                else:
                    Q_prev = neg_inf # i=0인 경우에는 Q_i-1_j-1을 반영하지 않아야 한다.
            else:
                Q_prev = logp[i-1, j-1]
            
            # logp에 Q를 갱신한다.
            logp[i, j] = max(Q_cur, Q_prev) + logp[i, j]
    
    # backtracking
    idx = T_max - 1
    for j in range(F_max-1, -1, -1): # F_max-1부터 -1까지(-1 포함 없이 0까지) -1씩 감소
        path[idx, j] = 1
        if idx != 0:
            if (logp[idx, j-1] < logp[idx-1, j-1]) or (idx == j):
                idx -= 1
    
    return path
        

def maximum_path(logp, attention_mask):
    """
    Glow-TTS에 사용되는 모듈
    MAS를 사용하여 alignment를 찾아주는 역할을 한다.
    논문 저자 구현에서는 cpython을 이용하여 병렬 처리를 구현한 듯 하나
    여기에서는 python만을 이용하여 구현하였다.
    =====inputs=====
    logp: (B, T, F) | N(x_mean, x_std)의 log-likelihood
    attention_mask: (B, T, F)
    =====outputs=====
    path: (B, T, F) | alignment
    """
    B = logp.shape[0]
    
    logp = logp * attention_mask
    # 계산은 CPU에서 실행되도록 하기 위해 기존의 device를 저장하고 .cpu().numpy()를 한다.
    logp_device = logp.device
    logp_type = logp.dtype
    logp = logp.data.cpu().numpy().astype(np.float32)
    attention_mask = attention_mask.data.cpu().numpy()
    
    path = np.zeros_like(logp).astype(np.int32) # (B, T, F)
    T_max = attention_mask.sum(1)[:, 0].astype(np.int32) # (B)
    F_max = attention_mask.sum(2)[:, 0].astype(np.int32) # (B)
    
    # MAS 알고리즘
    for idx in range(B):
        path[idx] = MAS(path[idx], logp[idx], T_max[idx], F_max[idx]) # (T, F)
    return torch.from_numpy(path).to(device=logp_device, dtype=logp_type)

def generate_path(ceil_dur, attention_mask):
    """
    Glow-TTS에 사용되는 모듈
    inference 과정에서 alignment를 만들어낸다.
    =====input=====
    ceil_dur: (B, T) | 추론한 duration에 ceil 연산한 것 | ex) [[2, 1, 2, 2, ...], [1, 2, 1, 3, ...], ...]
    attention_mask: (B, T, F)
    =====output=====
    path: (B, T, F) | alignment
    """
    B, T, Frame = attention_mask.shape
    cum_dur = torch.cumsum(ceil_dur, 1)
    cum_dur = cum_dur.to(torch.int32) # (B, T) | 누적합 | ex) [[2, 3, 5, 7, ...], [1, 3, 4, 7, ...], ...]
    path = torch.zeros(B, T, Frame).to(ceil_dur.device) # (B, T, F) | all False(0)
    
    # make the sequence_mask
    for b, batch_cum_dur in enumerate(cum_dur):
        for t, each_cum_dur in enumerate(batch_cum_dur):
            path[b, t, :each_cum_dur] = torch.ones((1, 1, each_cum_dur)).to(ceil_dur.device)
                # cum_dur로부터 True(1)를 path에 새겨넣는다.
    path = path - F.pad(path, (0, 0, 1, 0, 0, 0))[:, :-1] # (B, T, F)
    """
    ex) batch를 잠시 제외해두고 예시를 든다.
    [[1, 1, 0, 0, 0, 0, 0],   [[0, 0, 0, 0, 0, 0, 0],    [[1, 1, 0, 0, 0, 0, 0],
     [1, 1, 1, 0, 0, 0, 0], -  [1, 1, 0, 0, 0, 0, 0],  =  [0, 0, 1, 0, 0, 0, 0],
     [1, 1, 1, 1, 1, 0, 0],    [1, 1, 1, 0, 0, 0, 0],     [0, 0, 0, 1, 1, 0, 0],
     [1, 1, 1, 1, 1, 1, 1]]    [1, 1, 1, 1, 1, 0, 0]]     [0, 0, 0, 0, 0, 1, 1]]
    """
    path = path * attention_mask
    return path

class Decoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.flows = nn.ModuleList()
        for i in range(12):
            self.flows.append(ActNorm())
            self.flows.append(InvertibleConv())
            self.flows.append(AffineCouplingLayer())
        
    def forward(self, x, x_mask, reverse=False):
        """
        =====inputs=====
        x: (B, 80, F) | mel-spectrogram(Direct) OR latent representation(Reverse)
        x_mask: (B, 1, F)
        =====outputs=====
        z: (B, 80, F) | latent representation(Direct) OR mel-spectrogram(Reverse)
        total_log_det: (B) or None | log determinant
        """
        if not reverse:
            flows = self.flows
            total_log_det = 0
        else:
            flows = reversed(self.flows)
            total_log_det = None
        
        x, x_mask = Squeeze(x, x_mask) # (B, 80, F) -> (B, 160, F//2) | (B, 1, F) -> (B, 1, F//2)
        
        for f in flows:
            if not reverse:
                x, log_det = f(x, x_mask, reverse=reverse)
                total_log_det += log_det
            else:
                x, _ = f(x, x_mask, reverse=reverse)
                
        x, x_mask = Unsqueeze(x, x_mask) # (B, 160, F//2) -> (B, 80, F) | (B, 1, F//2) -> (B, 1, F)
        
        return x, total_log_det

"""
Decoder는 Glow: Generative Flow with Invertible 1×1 Convolutions 논문의 기본 구조를 따라간다.
Glow 논문: https://arxiv.org/pdf/1807.03039.pdf
"""
def Squeeze(x, x_mask):
    """
    Decoder의 preprocessing
    =====inputs=====
    x: (B, 80, F) | mel_spectrogram or latent representation
    x_mask: (B, 1, F)
    =====outputs=====
    x: (B, 160, F//2) | F//2 = [F/2] ([]: 가우스 기호)
    x_mask: (B, 160, F//2)
    """
    B, C, F = x.size()
    x = x[:, :, :(F//2)*2] # F가 홀수이면 맨 뒤 한 frame을 버림.
    x = x.view(B, C, F//2, 2) # (B, 80, F//2, 2)
    x = x.permute(0, 3, 1, 2).contiguous() # (B, 2, 80, F//2)
    x = x.view(B, C*2, F//2) # (B, 160, F//2)
    
    x_mask = x_mask[:, :, 1::2] # (B, 1, F//2) frame을 1부터 한칸씩 건너뛴다.
    x = x * x_mask # masking
    return x, x_mask

class ActNorm(nn.Module):
    """
    Decoder의 1번째 모듈
    """
    def __init__(self):
        super().__init__()
        self.log_s = nn.Parameter(torch.zeros(1, 160, 1)) # Glow 논문의 s에서 log를 취한 것이다. 즉, log[s]
        self.bias = nn.Parameter(torch.zeros(1, 160, 1))
        
    def forward(self, x, x_mask, reverse=False):
        """
        =====inputs=====
        x: (B, 160, F//2) | mel_spectrogram features
        x_mask: (B, 1, F//2) | mel_spectrogram features의 mask. (Decoder의 Squeeze에서 변형됨.)
        =====outputs=====
        z: (B, 160, F//2)
        log_det: (B) or None | log_determinant, reverse=True이면 None 반환
        """
        x_len = torch.sum(x_mask, [1, 2]) # (B) | 1, 2차원의 값을 더한다. cf. [1, 2] 대신 [2]만 사용하면 shape가 (B, 1)이 된다.
        
        if not reverse:
            z = (x * torch.exp(self.log_s) + self.bias) * x_mask # function & masking
            log_det = x_len * torch.sum(self.log_s) # log_determinant
                # Glow 논문의 Table 1을 확인하라. log_s를 log[s]라 볼 수 있다.
                # determinant 대신 log_determinant를 사용하는 이유는 det보다 작은 수치와 적은 계산량 때문으로 추측된다.
        else:
            z = ((x - self.bias) / torch.exp(self.log_s)) * x_mask # inverse function & masking
            log_det = None
        
        return z, log_det

class InvertibleConv(nn.Module):
    """
    Decoder의 2번째 모듈
    """
    def __init__(self):
        super().__init__()
        Q = torch.linalg.qr(torch.FloatTensor(4, 4).normal_())[0] # (4, 4)
        """
        torch.FloatTensor(4, 4).normal_(): 정규분포 N(0, 1)에서 무작위로 추출한 4x4 matrix
        Q, R = torch.linalg.qr(W): QR분해 | Q: 직교 행렬, R: upper traiangular 행렬 cf. det(Q) = 1 or -1
        """
        if torch.det(Q) < 0:
            Q[:, 0] = -1 * Q[:, 0] # 0번째 열의 부호를 바꿔서 det(Q) = -1로 만든다.
        self.W = nn.Parameter(Q)
    
    def forward(self, x, x_mask, reverse=False):
        """
        =====inputs=====
        x: (B, 160, F//2)
        x_mask: (B, 1, F//2)
        =====outputs=====
        z: (B, 160, F//2)
        log_det: (B) or None
        """
        B, C, f = x.size() # B, 160, F//2
        x_len = torch.sum(x_mask, [1, 2]) # (B)
        
        # channel mixing
        x = x.view(B, 2, C//4, 2, f) # (B, 2, 40, 2, F//2)
        x = x.permute(0, 1, 3, 2, 4).contiguous() # (B, 2, 2, 40, F//2)
        x = x.view(B, 4, C//4, f) # (B, 4, 40, F//2)
        
        # 편의상 log_det부터 구한다.
        if not reverse:
            weight = self.W
            log_det = (C/4) * x_len * torch.logdet(self.W) # (B) | torch.logdet(W): log(det(W))
                # height = C/4, width = x_len 인 상황임을 고려하면 Glow 논문의 log_determinant 식과 같다.
        else:
            weight = torch.linalg.inv(self.W) # inverse matrix
            log_det = None
        
        weight = weight.view(4, 4, 1, 1)
        z = F.conv2d(x, weight) # (B, 4, 40, F//2) * (4, 4, 1, 1) -> (B, 4, 40, F//2)
        """
        F.conv2d(x, weight)의 convolution 연산은 다음과 같이 생각해야 한다.
        (B, 4, 40, F//2): (batch_size, in_channels, height, width)
        (4, 4, 1, 1): (out_channels, in_channels/groups, kernel_height, kernel_width)
        
        즉, nn.Conv2d(4, 4, kernel_size=(1, 1))인 상황에 가중치를 준 것이다.
        """
        
        # channel unmixing
        z = z.view(B, 2, 2, C//4, f) # (B, 4, 40, F//2) -> (B, 2, 2, 40, F//2)
        z = z.permute(0, 1, 3, 2, 4).contiguous() # (B, 2, 40, 2, F//2)
        z = z.view(B, C, f) * x_mask # (B, 160, F//2) & masking
        return z, log_det
    
class WN(nn.Module):
    """
    Decoder의 3번째 모듈인 AffineCouplingLayer의 모듈
    
    해당 구조는 WAVEGLOW: A FLOW-BASED GENERATIVE NETWORK FOR SPEECH SYNTHESIS 로부터 제안되었다.
    WaveGlow 논문: https://arxiv.org/pdf/1811.00002.pdf
    """
    def __init__(self, dilation_rate=1):
        super().__init__()
        self.in_layers = nn.ModuleList()
        self.res_skip_layers = nn.ModuleList()
        
        for i in range(4):
            dilation = dilation_rate ** i # NVIDIA WaveGlow에서는 dilation_rate=2이지만, 여기에서는 1이므로 의미는 없다.
            in_layer = weight_norm(nn.Conv1d(192, 2*192, kernel_size=5, dilation=dilation,
                                 padding=((5-1) * dilation)//2)) # (B, 192, F//2) -> (B, 2*192, F//2)
            self.in_layers.append(in_layer)
            
            if i < 3:
                res_skip_layer = weight_norm(nn.Conv1d(192, 2*192, kernel_size=1)) # (B, 192, F//2) -> (B, 2*192, F//2)
            else:
                res_skip_layer = weight_norm(nn.Conv1d(192, 192, kernel_size=1)) # (B, 192, F//2) -> (B, 192, F//2)
            self.res_skip_layers.append(res_skip_layer)
        
        self.dropout = nn.Dropout(0.05)
    
    def forward(self, x, x_mask):
        """
        =====inputs=====
        x: (B, 192, F//2)
        x_mask: (B, 1, F//2)
        =====outputs=====
        output: (B, 192, F//2)
        """
        output = torch.zeros_like(x) # (B, 192, F//2) all zeros
        
        for i in range(4):
            x_in = self.in_layers[i](x) # (B, 192, F//2) -> (B, 2*192, F//2)
            x_in = self.dropout(x_in) # dropout
            
            # fused add tanh sigmoid multiply
            tanh_act = torch.tanh(x_in[:, :192, :]) # (B, 192, F//2)
            sigmoid_act = torch.sigmoid(x_in[:, 192:, :]) # (B, 192, F//2)
            
            acts = sigmoid_act * tanh_act # (B, 192, F//2)
            
            x_out = self.res_skip_layers[i](acts) # (B, 192, F//2) -> (B, 2*192, F//2) or [last](B, 192, F//2)
            if i < 3:
                x = (x + x_out[:, :192, :]) * x_mask # residual connection & masking
                output += x_out[:, 192:, :] # add output
            else:
                output += x_out # (B, 192, F//2)
        
        output = output * x_mask # masking
        return output

class AffineCouplingLayer(nn.Module):
    """
    Decoder의 3번째 모듈
    """
    def __init__(self):
        super().__init__()
        self.start_conv = weight_norm(nn.Conv1d(160//2, 192, kernel_size=1)) # (B, 80, F//2) -> (B, 192, F//2)
        self.wn = WN()
        self.end_conv = nn.Conv1d(192, 160, kernel_size=1) # (B, 192, F//2) -> (B, 160, F//2)
        # end_conv의 초기 가중치를 0으로 설정하는 것이 처음에 학습하지 않는 역할을 하며, 이는 학습 안정화에 도움이 된다.
        self.end_conv.weight.data.zero_() # weight를 0으로 초기화
        self.end_conv.bias.data.zero_() # bias를 0으로 초기화
        
    def forward(self, x, x_mask, reverse=False):
        """
        =====inputs=====
        x: (B, 160, F//2)
        x_mask: (B, 1, F//2)
        =====outputs=====
        z: (B, 160, F//2)
        log_det: (B) or None
        """
        B, C, f = x.size() # B, 160, F//2
        x_0, x_1 = x[:, :C//2, :], x[:, C//2:, :] # split: (B, 80, F//2) x2
        
        x = self.start_conv(x_0) * x_mask # (B, 80, F//2) -> (B, 192, F//2) & masking
        x = self.wn(x, x_mask) # (B, 192, F//2)
        out = self.end_conv(x) # (B, 192, F//2) -> (B, 160, F//2)
        
        z_0 = x_0 # (B, 80, F//2)
        m = out[:, :C//2, :] # (B, 80, F//2)
        log_s = out[:, C//2:, :] # (B, 80, F//2)
        
        if not reverse:
            z_1 = (torch.exp(log_s) * x_1 + m) * x_mask # (B, 80, F//2) | function & masking 
            log_det = torch.sum(log_s * x_mask, [1, 2]) # (B)
        else:
            z_1 = (x_1 - m) / torch.exp(log_s) * x_mask # (B, 80, F//2) | inverse function & masking
            log_det = None
        
        z = torch.cat([z_0, z_1], dim=1) # (B, 160, F//2)
        return z, log_det
    
def Unsqueeze(x, x_mask):
    """
    Decoder의 postprocessing
    =====inputs=====
    x: (B, 160, F//2)
    x_mask: (B, 1, F//2)
    =====outputs=====
    x: (B, 80, F)
    x_mask: (B, 1, F)
    """
    B, C, f = x.size() # B, 160, F//2
    x = x.view(B, 2, C//2, f) # (B, 2, 80, F//2)
    x = x.permute(0, 2, 3, 1).contiguous() # (B, 80, F//2, 2)
    x = x.view(B, C//2, 2*f) # (B, 160, F)
    
    x_mask = x_mask.unsqueeze(3).repeat(1, 1, 1, 2).view(B, 1, 2*f) # (B, 1, F//2, 1) -> (B, 1, F//2, 2) -> (B, 1, F)
    x = x * x_mask # masking
    return x, x_mask

class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(symbol_length, 192) # (B, T) -> (B, T, 192)
        nn.init.normal_(self.embedding.weight, 0.0, 192**(-0.5)) # 가중치 정규분포 초기화 (N(0, 0.07xx))
        
        self.prenet = PreNet()
        self.transformer_encoder = TransformerEncoder()
        self.project_mean = nn.Conv1d(192, 80, kernel_size=1) # (B, 192, T) -> (B, 80, T)
        self.project_std = nn.Conv1d(192, 80, kernel_size=1) # (B, 192, T) -> (B, 80, T)
        
        self.duration_predictor = DurationPredictor()
        
    def forward(self, text, text_len):
        """
        =====inputs=====
        text: (B, Max_T)
        text_len: (B)
        =====outputs=====
        x_mean: (B, 80, T) | 평균, 논문 저자 구현의 train.py에서 out_channels를 80으로 설정한 것을 알 수 있음.
        x_std: (B, 80, T) | 표준편차
        x_dur: (B, 1, T)
        x_mask: (B, 1, T)
        """
        x = self.embedding(text) * math.sqrt(192) # (B, T) -> (B, T, 192) # math.sqrt(192) = 13.xx (수정)
        x = x.transpose(1, 2) # (B, T, 192) -> (B, 192, T)
        
        # Make the x_mask
        x_mask = torch.zeros_like(x[:, 0:1, :], dtype=torch.bool) # (B, 1, T)
        for idx, length in enumerate(text_len):
            x_mask[idx, :, :length] = True
        
        x = self.prenet(x, x_mask) # (B, 192, T)
        x = self.transformer_encoder(x, x_mask) # (B, 192, T)
        
        # project
        x_mean = self.project_mean(x) * x_mask # (B, 192, T) -> (B, 80, T)
        # x_std = self.project_std(x) * x_mask # (B, 192, T) -> (B, 80, T)
        ##### 아래는 mean_only를 적용한 것임. #####
        x_std = torch.zeros_like(x_mean) # x_log_std: (B, 80, T), all zero # log std = 0이므로 std = 1로 계산됨.

        # duration predictor
        x_dp = torch.detach(x) # stop_gradient
        x_dur = self.duration_predictor(x_dp, x_mask) # (B, 192, T) -> (B, 1, T)
        
        return x_mean, x_std, x_dur, x_mask

class LayerNorm(nn.Module):
    """
    여러 곳에서 정규화(Norm)를 위해 사용되는 모듈.
    
    nn.LayerNorm이 이미 pytorch 안에 구현되어 있으나, 항상 마지막 차원을 정규화한다.
    그래서 channel을 기준으로 정규화하는 LayerNorm을 따로 구현한다.
    """
    def __init__(self, channels):
        """
        channels: 입력 데이터의 channel 수 | LayerNorm은 channel 차원을 정규화한다.
        """
        super().__init__()
        self.channels = channels
        self.eps = 1e-4
        
        self.gamma = nn.Parameter(torch.ones(channels)) # 학습 가능한 파라미터
        self.beta = nn.Parameter(torch.zeros(channels)) # 학습 가능한 파라미터
        
    def forward(self, x):
        """
        =====inputs=====
        x: (B, channels, *) | 정규화할 입력 데이터
        =====outputs=====
        x: (B, channels, *) | channel 차원이 정규화된 데이터
        """
        mean = torch.mean(x, dim=1, keepdim=True) # channel 차원(index=1)의 평균 계산, 차원을 유지한다.
        variance = torch.mean((x-mean)**2, dim=1, keepdim=True) # 분산 계산
        
        x = (x - mean) * (variance + self.eps)**(-0.5) # (x - m) / sqrt(v)
        
        n = len(x.shape)
        shape = [1] * n
        shape[1] = -1 # shape = [1, -1, 1] or [1, -1, 1, 1]
        x = x * self.gamma.view(*shape) + self.beta.view(*shape) # y = x*gamma + beta
        
        return x

class PreNet(nn.Module):
    """
    Encoder의 1번째 모듈
    """
    def __init__(self):
        super().__init__()
        self.convs = nn.ModuleList()
        self.norms = nn.ModuleList()
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        for i in range(3):
            self.convs.append(nn.Conv1d(192, 192, kernel_size=5, padding=2)) # (B, 192, T) 유지
            self.norms.append(LayerNorm(192)) # (B, 192, T) 유지
        self.linear = nn.Conv1d(192, 192, kernel_size=1) # (B, 192, T) 유지 | linear 역할을 하는 conv
        
    def forward(self, x, x_mask):
        """
        =====inputs=====
        x: (B, 192, T) | Embedding된 입력 데이터
        x_mask: (B, 1, T) | 글자 길이에 따른 mask (글자가 있으면 True, 없으면 False로 구성)
        =====outputs=====
        x: (B, 192, T)
        """
        x0 = x
        for i in range(3):
            x = self.convs[i](x * x_mask)
            x = self.norms[i](x)
            x = self.relu(x)
            x = self.dropout(x)
        x = self.linear(x)
        x = x0 + x # residual connection
        return x
    
class MultiHeadAttention(nn.Module):
    """
    Encoder 중 2번째 모듈인 TransformerEncoder의 1번째 모듈
    """
    def __init__(self):
        super().__init__()
        self.n_heads = 2
        self.window_size = 4
        self.k_channels = 192 // self.n_heads # 96
        
        self.linear_q = nn.Conv1d(192, 192, kernel_size=1) # (B, 192, T) 유지
        self.linear_k = nn.Conv1d(192, 192, kernel_size=1) # (B, 192, T) 유지
        self.linear_v = nn.Conv1d(192, 192, kernel_size=1) # (B, 192, T) 유지
        nn.init.xavier_uniform_(self.linear_q.weight)
        nn.init.xavier_uniform_(self.linear_k.weight)
        nn.init.xavier_uniform_(self.linear_v.weight)
        
        relative_std = self.k_channels ** (-0.5) # 0.1xx
        self.relative_k = nn.Parameter(torch.randn(1, self.window_size * 2 + 1, self.k_channels) * relative_std) # (1, 9, 96)
        self.relative_v = nn.Parameter(torch.randn(1, self.window_size * 2 + 1, self.k_channels) * relative_std) # (1, 9, 96)
        
        self.attention_weights = None
        self.linear_out = nn.Conv1d(192, 192, kernel_size=1) # (B, 192, T) 유지
        self.dropout = nn.Dropout(0.1)
        
    def forward(self, query, context, attention_mask, self_attention=True):
        """
        =====inputs=====
        query: (B, 192, T_target) | Glow-TTS에서는 self-attention만 이용하므로 query와 context가 동일한 텐서 x이다.
        context: (B, 192, T_source) | query = context || 여기에서는 특히 T_source = T_target 이다.
        attention_mask: (B, 1, T, T) | x_mask.unsqueeze(2) * z_mask.unsqueeze(3)
        self_attention: True/False | self_attention일 때 relative position representations를 적용한다. 여기에서는 항상 True이다.
        # 실제로는 query와 context에 같은 텐서 x를 입력하면 된다.
        =====outputs=====
        output: (B, 192, T)
        """
        
        query = self.linear_q(query)
        key = self.linear_k(context)
        value = self.linear_v(context)
        
        B, _, T_tar = query.size()
        T_src = key.size(2)
        query = query.view(B, self.n_heads, self.k_channels, T_tar).transpose(2, 3)
        key = key.view(B, self.n_heads, self.k_channels, T_src).transpose(2, 3)
        value = value.view(B, self.n_heads, self.k_channels, T_src).transpose(2, 3)
            # (B, 192, T_src) -> (B, 2, 96, T_src) -> (B, 2, T_src, 96)
            
        scores = torch.matmul(query, key.transpose(2, 3)) / (self.k_channels ** 0.5)
            # (B, 2, T_tar, 96) * (B, 2, 96, T_src) -> (B, 2, T_tar, T_src)
        
        if self_attention: # True
            # Get relative embeddings (relative_keys) (1-1)
            padding = max(T_src - (self.window_size + 1), 0) # max(T-5, 0)
            start_pos = max((self.window_size + 1) - T_src, 0) # max(5-T, 0)
            end_pos = start_pos + 2 * T_src - 1 # (2*T-1) or (T+4)
            relative_keys = F.pad(self.relative_k, (0, 0, padding, padding))
                # (1, 9, 96) -> (1, pad+9+pad, 96) = (1, 2T-1, 96)
            """
            위 코드의 F.pad(input, pad) 에서 pad = (0, 0, padding, padding)은 다음을 의미한다.
            - 앞의 (0, 0): input의 -1차원을 앞으로 0, 뒤로 0만큼 패딩한다.
            - 앞의 (padding, padding): input의 -2차원을 앞으로 padding, 뒤로 padding만큼 패딩한다.
            즉, F.pad에서 pad는 역순으로 생각해주어야 한다.
            """
            relative_keys = relative_keys[:, start_pos:end_pos, :] # (1, 2T-1, 96)
            
            # Matmul with relative keys (2-1)
            relative_keys = relative_keys.unsqueeze(0).transpose(2, 3) # (1, 2T-1, 96) -> (1, 1, 2T-1, 96) -> (1, 1, 96, 2T-1)
            x = torch.matmul(query, relative_keys) # (B, 2, T_tar, 96) * (1, 1, 96, 2T_src-1) = (B, 2, T, 2T-1)
                # self attention에서는 T_tar = T_src이므로 이를 다르게 고려할 필요가 없다.
            
            # Relative position to absolute position (3-1)
            T = T_tar # Absolute position to relative position에서도 쓰임.
            x = F.pad(x, (0, 1)) # (B, 2, T, 2*T-1) -> (B, 2, T, 2*T)
            x = x.view(B, self.n_heads, T * 2 * T) # (B, 2, T, 2*T) -> (B, 2. 2T^2)
            x = F.pad(x, (0, T-1)) # (B, 2, 2T^2 + T - 1)
            x = x.view(B, self.n_heads, T+1, 2*T-1) # (B, 2, T+1, 2T-1)
            relative_logits = x[:, :, :T, T-1:] # (B, 2, T, T)
            
            # Compute scores
            scores_local = relative_logits / (self.k_channels ** 0.5)
            scores = scores + scores_local # (B, 2, T, T)
            """
            위 식은 Self-Attention with Relative Position Representations 논문의 5번 식을 구현한 것이다.
            Relative- 논문: https://arxiv.org/pdf/1803.02155.pdf
            """

        scores = scores.masked_fill(attention_mask == 0, -1e-4) # attention_mask가 0인 곳을 -1e-4로 채운다.
        
        attention_weights = F.softmax(scores, dim=-1) # (B, 2, T_tar, T_src) # Relative- 논문에서의 alpha에 해당한다.
        attention_weights = self.dropout(attention_weights) # dropout하는 이유가 무엇일까?
        output = torch.matmul(attention_weights, value) # (B, 2, T_tar, T_src) * (B, 2, T_src, 96) -> (B, 2, T_tar, 96)
        
        if self_attention: # True
            # Absolute position to relative position (3-2)
            x = F.pad(attention_weights, (0, T-1)) # (B, 2, T, T) -> (B, 2, T, 2T-1)
            x = x.view((B, self.n_heads, T * (2*T-1))) # (B, 2, 2T^2-T)
            x = F.pad(x, (T, 0)) # (B, 2, 2T^2) # 앞에 패딩
            x = x.view((B, self.n_heads, T, 2*T)) # (B, 2, T, 2T)
            relative_weights = x[:, :, :, 1:] # (B, 2, T, 2T-1)
            
            # Get relative embeddings (relative_value) (1-2) # (1-1)과 거의 동일
            padding = max(T_src - (self.window_size + 1), 0) # max(T-5, 0)
            start_pos = max((self.window_size + 1) - T_src, 0) # max(5-T, 0)
            end_pos = start_pos + 2 * T_src - 1 # (2*T-1) or (T+4)
            relative_values = F.pad(self.relative_v, (0, 0, padding, padding))
                # (1, 9, 96) -> (1, pad+9+pad, 96) = (1, 2T-1, 96)
            relative_values = relative_values[:, start_pos:end_pos, :] # (1, 2T-1, 96)

            # Matmul with relative values (2-2)
            relative_values = relative_values.unsqueeze(0) # (1, 1, 2T-1, 96)

            output = output + torch.matmul(relative_weights, relative_values)
                # (B, 2, T, 2T-1) * (1, 1, 2T-1, 96) = (B, 2, T, 96)
            """
            위 식은 Self-Attention with Relative Position Representations 논문의 3번 식을 구현한 것이다. (분배법칙 이용)
            Relative- 논문: https://arxiv.org/pdf/1803.02155.pdf
            """
        
        output = output.transpose(2, 3).contiguous().view(B, 192, T_tar)
            # (B, 2, 96, T) -> 메모리에 연속 배치 -> (B, 192, T)
            
        self.attention_weights = attention_weights # (B, 2, T, T)
        output = self.linear_out(output)
        return output # (B, 192, T)
    
class FFN(nn.Module):
    """
    Encoder 중 2번째 모듈인 TransformerEncoder의 2번째 모듈
    """
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv1d(192, 768, kernel_size=3, padding=1) # (B, 192, T) -> (B, 768, T)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv1d(768, 192, kernel_size=3, padding=1) # (B, 768, T) -> (B, 192, T)
        self.dropout = nn.Dropout(0.1)
    
    def forward(self, x, x_mask):
        """
        =====inputs=====
        x: (B, 192, T)
        x_mask: (B, 1, T)
        =====outputs=====
        output: (B, 192, T)
        """
        x = self.conv1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.conv2(x)
        output = x * x_mask
        return output
    
class TransformerEncoder(nn.Module):
    """
    Encoder의 2번째 모듈
    """
    def __init__(self):
        super().__init__()
        self.attentions = nn.ModuleList()
        self.norms1 = nn.ModuleList()
        self.ffns = nn.ModuleList()
        self.norms2 = nn.ModuleList()
        for i in range(6):
            self.attentions.append(MultiHeadAttention())
            self.norms1.append(LayerNorm(192))
            self.ffns.append(FFN())
            self.norms2.append(LayerNorm(192))
        self.dropout = nn.Dropout(0.1)
        
    def forward(self, x, x_mask):
        """
        =====inputs=====
        x: (B, 192, T)
        x_mask: (B, 1, T)
        =====outputs=====
        output: (B, 192, T)
        """
        attention_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(3)
            # (B, 1, 1, T) * (B, 1, T, 1) = (B, 1, T, T), only consist 0 or 1
        for i in range(6):
            x = x * x_mask
            y = self.attentions[i](x, x, attention_mask)
            y = self.dropout(y)
            x = x + y # residual connection
            x = self.norms1[i](x) # (B, 192, T) 유지
            
            y = self.ffns[i](x, x_mask)
            y = self.dropout(y)
            x = x + y # residual connection
            x = self.norms2[i](x)
        output = x * x_mask
        return output # (B, 192, T)
    
class DurationPredictor(nn.Module):
    """
    Encoder의 3번째 모듈
    """
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv1d(192, 256, kernel_size=3, padding=1) # (B, 192, T) -> (B, 256, T)
        self.norm1 = LayerNorm(256)
        self.conv2 = nn.Conv1d(256, 256, kernel_size=3, padding=1) # (B, 256, T) -> (B, 256, T)
        self.norm2 = LayerNorm(256)
        self.linear = nn.Conv1d(256, 1, kernel_size=1) # (B, 256, T) -> (B, 1, T)
        
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.1)
        
    def forward(self, x, x_mask):
        """
        =====inputs=====
        x: (B, 192, T)
        x_mask: (B, 1, T)
        =====outputs=====
        output: (B, 1, T)
        """
        x = self.conv1(x * x_mask) # (B, 192, T) -> (B, 256, T)
        x = self.relu(x)
        x = self.norm1(x)
        x = self.dropout(x)
        
        x = self.conv2(x * x_mask) # (B, 256, T) -> (B, 256, T)
        x = self.relu(x)
        x = self.norm2(x)
        x = self.dropout(x)
        
        x = self.linear(x * x_mask) # (B, 256, T) -> (B, 1, T)
        output = x * x_mask
        return output