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

import torch.optim as optim


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(4, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)

        self.p_fc1 = nn.Conv2d(128, 4, 1)
        self.p_fc2 = nn.Linear(4 * 8 * 8, 64)

        self.v_fc1 = nn.Conv2d(128, 2, 1)
        self.v_fc2 = nn.Linear(2 * 8 * 8, 1)

    def forward(self, x):
        a1 = F.relu(self.conv1(x))
        a2 = F.relu(self.conv2(a1))
        a3 = F.relu(self.conv3(a2))

        p1 = F.relu(self.p_fc1(a3))
        p_act = p1.view(-1, 4 * 8 * 8)

        p_out = F.softmax(self.p_fc2(p_act), dim=0)

        v1 = F.relu(self.v_fc1(a3))
        v_act = v1.view(-1, 2 * 8 * 8)

        v_out = torch.tanh(self.v_fc2(v_act))

        return p_out, v_out


class PolicyValueNet():
    def __init__(self, board_width=8, board_height=8, model_file=None):

        self.board_width = board_width
        self.board_height = board_height
        self.model = Net()
        self.v_loss = nn.MSELoss()
        # self.p_loss=nn.CrossEntropyLoss()
        self.optimizer = optim.Adam(self.model.parameters(), lr=0.01)
        if model_file:
            self.model.load_state_dict(torch.load(model_file))

    def policy_value_fn(self, board):
        legal_positions = board.available
        current_state = board.current_state()
        x = torch.FloatTensor(
            current_state.reshape(-1, 4, self.board_width, self.board_height))

        probs, value = self.model.forward(x)
        probs = probs.data.numpy()
        value = value.data.numpy()

        probs = zip(legal_positions, probs.flatten()[legal_positions])

        return probs, value

    def policy_value(self, x):
        x = torch.FloatTensor(np.reshape(
            x, (-1, 4, self.board_width, self.board_height)))
        probs, value = self.model.forward(x)
        probs = probs.data.numpy()
        value = value.data.numpy()
        return probs, value
    

    def train_step(self, state_batch, probs_batch, winner_batch, lr):

        state_batch = torch.FloatTensor(state_batch)
        probs_batch = torch.LongTensor(probs_batch)
        winner_batch = torch.FloatTensor(winner_batch)

        self.optimizer.zero_grad()
        p_out, v_out = self.model(state_batch)
        p_loss = -torch.mean(torch.sum(probs_batch * torch.log(p_out), 1))
        v_loss = self.v_loss(v_out, winner_batch)
        loss = p_loss + v_loss
        loss.backward()
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr
        self.optimizer.step()

        return p_loss.item(), v_loss.item()

    def save_model(self, dir):
        torch.save(self.model.state_dict(), dir)


if __name__ == '__main__':
    net = PolicyValueNet()
    net.save_model("./testmodel")
