
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

class Connect2Model(nn.Module):
    
    def __init__(self, board_size, action_size, device):
        
        super(Connect2Model, self).__init__()
        
        self.device = device
        self.size = board_size
        self.action_size = action_size
        
        self.fc1 = nn.Linear(in_features = self.size, out_features=16)
        #fc1에서, 입력 feature은 보드 사이즈로 넣고 출력 feature은 16개
        self.fc2 = nn.Linear(in_features = 16, out_features = 16)
        #fc2에서 방금 그 16개 넣고 out도 16개
        
        #TWO heads on our network
        self.action_head = nn.Linear(in_features = 16, out_features=self.action_size)
        self.value_head = nn.Linear(in_features=16, out_features=1)
        
        self.to(device)  #device로 다시 보낸다?
        
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        
        action_logits = self.action_head(x)
        value_logits = self.value_head(x)
        
        return F.softmax(action_logits, dim=1)
        #softmax는 action_logits(1차원의 벡터)을 0~1사이로정규화해서 확률을 추정한다.
        
    def predict(self, board):
        board = torch.FloatTensor(board.astype(np.float32)).to(self.device) 
        #board를 일단 float값으로 바꾸고, 그걸 torch에 넣은 값을 각각 board에 할당
        board = board.view(1, self.size)
        #view - 원소의 수를 유지하면서 텐서의 크기 변경 . 여기서는 shape을 (1, boardsize)
        self.eval() #training과 inference시 다르게 동작하는 layer때문에 모든 동작들을 inference mode로 바꿔줌
        with torch.no_grad(): #pytorch가 autograd engine을 꺼버림. 더이상 gradient를 자동으로 트래킹 안 함
            pi, v = self.forward(board) #함수 fc1, fc2 사용해서 통과시켜줌
            
        return pi.data.cpu().numpy()[0], v.data.cpu().numpy()[0]
        
        