# -*- coding: utf-8 -*-
"""model

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/17MS4qILsNWwuQx9VysLxqi5NgQHBdOwl
"""

import torch
import torch.nn as nn

# Flatten된 각각의 patch가 Encoder에 입력되기 까지의 처리 담당
class LinearProjection(nn.Module):

    def __init__(self, patch_vec_size, num_patches, latent_vec_dim, drop_rate):
        super().__init__()
        self.linear_proj = nn.Linear(patch_vec_size, latent_vec_dim) # Embedding Dimension으로 변환

        # 학습이 가능하도록 nn.Parameter를 사용하여 생성한다
        # Class Token 붙이기
        # (1, D)
        self.cls_token = nn.Parameter(torch.randn(1, latent_vec_dim))

        # 학습이 가능하도록 nn.Parameter를 사용하여 생성한다
        # Positional Embedding 생성 
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches+1, latent_vec_dim))
        self.dropout = nn.Dropout(drop_rate)

    def forward(self, x):
        # x : # (batch_size, N, P^2 * C)
        batch_size = x.size(0)

        # self.linear_proj(x) : (batch_size, N, D)
        # self.cls_token : (1, D)
        # 위를 0축으로 batch_size, 1축으로 1, 2축으로 1번 반복하여 최종적으로 다음과 같은 사이즈가 나온다
        # self.cls_token.repeat(batch_size, 1, 1) : (batch_size, 1, D)
        x = torch.cat([self.cls_token.repeat(batch_size, 1, 1), self.linear_proj(x)], dim=1) # (batch_size, N+1, D)

        # Positional Embedding
        x += self.pos_embedding # (batch_size, N+1, D)
        x = self.dropout(x)  # (batch_size, N+1, D)
        return x  # (batch_size, N+1, D)

# Encoder의 Multi Head Attention Layer
class MultiheadedSelfAttention(nn.Module):
    def __init__(self, latent_vec_dim, num_heads, drop_rate):
        super().__init__()
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.num_heads = num_heads
        self.latent_vec_dim = latent_vec_dim
        self.head_dim = int(latent_vec_dim / num_heads)

        # 원래 self.query의 모양은 다음과 같아야 한다
        # (D, D_h)
        # 하지만 아래와 같은 Technique를 사용할 수 있다
        # 본래 D_h = D / k로 하기로 하였다
        # 이 때, self.query, self.key, self.value의 shape를 (D, D)로 하면
        # k개 head의 query, key, value를 동시에 계산할 수 있기 때문에 다음과 같이 하는 것이다
        self.query = nn.Linear(latent_vec_dim, latent_vec_dim)
        self.key = nn.Linear(latent_vec_dim, latent_vec_dim)
        self.value = nn.Linear(latent_vec_dim, latent_vec_dim)
        self.scale = torch.sqrt(self.head_dim*torch.ones(1)).to(device)
        self.dropout = nn.Dropout(drop_rate)

    def forward(self, x):
        # x : (batch_size, N+1, D)
        batch_size = x.size(0)


        q = self.query(x) # (batch_size, N+1, D)
        k = self.key(x) # (batch_size, N+1, D)
        v = self.value(x) # (batch_size, N+1, D)

        # 다음 q,k,v를 각 head에 대한 연산 결과로 바꾸어준다
        # (batch_size, N+1, D) > (batch_size, N+1, k, D_h) / k : head의 개수, D_h : head_dim
        # 그리고 이걸 Permute하여 다음 shape로 만들어준다
        

        q = q.view(batch_size, -1, self.num_heads, self.head_dim).permute(0,2,1,3) # (batch_size, k, N+1, D_h)
        kT = k.view(batch_size, -1, self.num_heads, self.head_dim).permute(0,2,3,1) # (batch_size, k,D_h, N+1)
        v = v.view(batch_size, -1, self.num_heads, self.head_dim).permute(0,2,1,3) # (batch_size, k, N+1, D_h)

        # attention Value 계산하기
        # @는 torch.matmul과 같다
        # dim=-1 : 마지막 차원 기준 softmax
        attention = torch.softmax(q @ kT / self.scale, dim=-1) # (batch_size, k, N+1, N+1)
        x = self.dropout(attention) @ v # (batch_size, k, N+1, D_h)

        # x.permute(0,2,1,3) : (batch_size, N+1, K, D_h)
        # 최종 : (batch_size, N+1, K*D_h)
        # k*D_h = D이므로
        # 최종 결과 출력은 (batch_size, N+1, D)이다. 즉 입력 shape와 완전히 동일해진다
        x = x.permute(0,2,1,3).reshape(batch_size, -1, self.latent_vec_dim)

        return x, attention # (Encoder의 출력값, Attention Score값)반환

# 한 번의 Encoder 연산
# 이러한 Encoder를 여러 번 거친 후 최종 출력값이 나오게 된다
class TFencoderLayer(nn.Module):
    def __init__(self, latent_vec_dim, num_heads, mlp_hidden_dim, drop_rate):
        super().__init__()
        self.ln1 = nn.LayerNorm(latent_vec_dim)
        self.ln2 = nn.LayerNorm(latent_vec_dim)
        self.msa = MultiheadedSelfAttention(latent_vec_dim=latent_vec_dim, num_heads=num_heads, drop_rate=drop_rate)
        self.dropout = nn.Dropout(drop_rate)
        self.mlp = nn.Sequential(nn.Linear(latent_vec_dim, mlp_hidden_dim),
                                 nn.GELU(), nn.Dropout(drop_rate),
                                 nn.Linear(mlp_hidden_dim, latent_vec_dim),
                                 nn.Dropout(drop_rate))

    def forward(self, x):
        z = self.ln1(x)
        z, att = self.msa(z)
        z = self.dropout(z)
        x = x + z
        z = self.ln2(x)
        z = self.mlp(z)
        x = x + z

        return x, att

# 최종 Vision Transformer 모델
class VisionTransformer(nn.Module):
    def __init__(self, patch_vec_size, num_patches, latent_vec_dim, num_heads, mlp_hidden_dim, drop_rate, num_layers, num_classes):
        super().__init__()

        # 먼저 Linear Projection 실시
        # patch_vec_size = (P**2*C) : 한 개 patch를 flatten한 차원
        # num_patches : 총 patch의 개수 = (H*W) / (P**2)
        # latent_vec_dim : D = Embedding된 각각의 Flattened Patch의 차원
        self.patchembedding = LinearProjection(patch_vec_size=patch_vec_size, num_patches=num_patches,
                                               latent_vec_dim=latent_vec_dim, drop_rate=drop_rate)
        
        # 여러 층의 Encoder가 들어있는 것을 다음과 같이 nn.ModuleList로 구현하였다
        # num_layers : Encoder의 층 수
        self.transformer = nn.ModuleList([TFencoderLayer(latent_vec_dim=latent_vec_dim, num_heads=num_heads,
                                                         mlp_hidden_dim=mlp_hidden_dim, drop_rate=drop_rate)
                                          for _ in range(num_layers)])

        # Encoder로부터 나온 출력 (N+1) * D에서 Class Token인 첫 행 D차원을 뽑아 num_classes 차원으로 바꾸어준다
        self.mlp_head = nn.Sequential(nn.LayerNorm(latent_vec_dim), nn.Linear(latent_vec_dim, num_classes))

    def forward(self, x):
        att_list = [] # 각 Encoder Layer에서 나오는 attention을 저장하는 용도
        x = self.patchembedding(x) # Linear Projection 실시
        for layer in self.transformer:
            x, att = layer(x) # Encoder를 통과시킨다
            att_list.append(att) # 각 Encoder를 통과시켜 나온 attention Value를 저장
        x = self.mlp_head(x[:,0]) # 첫 행인 Class Token만을 뽑아 classification에 사용

        return x, att_list # (최종 Prediction, 각 Encoder층에서의 Attention Value값) 반환