import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets
from torch.autograd import Variable
from sklearn.model_selection import train_test_split
import time
import pandas as pd
import numpy as np
import csv


batch_size = 128
NUM_EPOCHS = 30
LR = 0.001
TIME_STEP = 4


class CCRNN(nn.Module):
    def __init__(self):
        # 继承RNN
        super(CCRNN, self).__init__()

        self.ccLSTM = nn.LSTM(
            input_size=4,
            hidden_size=128,
            num_layers=4,
            bidirectional=True,
            batch_first=True

        )

        self.ccCNN22 = nn.Conv2d(
            in_channels=1,
            out_channels=1,
            kernel_size=2,
            stride=2,
            padding=0
        )

        self.ccCNN14 = nn.Conv2d(
            in_channels=1,
            out_channels=1,
            kernel_size=(1, 4),
            stride=1,
            padding=0
        )

        self.ccCNN41 = nn.Conv2d(
            in_channels=1,
            out_channels=1,
            kernel_size=(4, 1),
            stride=1,
            padding=0
        )

        self.CNN22toFC = nn.Linear(4, 64)
        self.CNN41toFC = nn.Linear(4, 32)
        self.CNN14toFC = nn.Linear(4, 32)
        self.LSTMtoFC = nn.Linear(256, 128)
        self.FCtoOut = nn.Linear(256, 4)

    def forward(self, x):
        LSTM_out, (h_n, c_n) = self.ccLSTM(x, None)
        CNN_in = torch.unsqueeze(x[:, 0:4, :], 1)
        CNN_out22 = self.ccCNN22(CNN_in)
        CNN_out41 = self.ccCNN41(CNN_in)
        CNN_out14 = self.ccCNN14(CNN_in)

        CNN22_reshape = CNN_out22.view(-1, 4)
        CNN14_reshape = CNN_out41.view(-1, 4)
        CNN41_reshape = CNN_out14.view(-1, 4)
        CNN22toFC = self.CNN22toFC(CNN22_reshape)
        CNN14toFC = self.CNN14toFC(CNN14_reshape)
        CNN41toFC = self.CNN41toFC(CNN41_reshape)
        LSTMtoFC = self.LSTMtoFC(LSTM_out[:, -1, :])
        CNNandLSTM = torch.cat((CNN22toFC, CNN41toFC, CNN14toFC, LSTMtoFC), 1)
        out = self.FCtoOut(CNNandLSTM)
        return out

#------------------读入数据-----------------------------
csv_data = pd.read_csv('./drive/My Drive/DATA.csv')
csv_data = csv_data.values
A = csv_data.shape[0]
board_data = csv_data[:,0:16]
# X = np.log2(X)
X = torch.FloatTensor(board_data)
X = np.int64(board_data)

# 转置后拼接
X = np.reshape(X, (-1,4,4))
XT = X.transpose(0,2,1)

X = np.concatenate((X,XT),axis=1)
print(X.shape)

direction_data = csv_data[:,16]
Y = np.int64(direction_data)


#-------------------------------------------------------


X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2,shuffle=False)
X_train = torch.FloatTensor(X_train)
X_test = torch.FloatTensor(X_test)
Y_train = torch.LongTensor(Y_train)
Y_test = torch.LongTensor(Y_test)

train_dataset = torch.utils.data.TensorDataset(X_train,Y_train)
# test_dataset = torch.utils.data.TensorDataset(X_test,Y_test)

train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                      batch_size=batch_size,
                      shuffle=True
)

# test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
#                       batch_size=batch_size,
#                       shuffle=False
# )

batch_size = 128
NUM_EPOCHS = 30
LR = 0.001
TIME_STEP = 4


#------------------读入数据-----------------------------
csv_data = pd.read_csv('./drive/My Drive/DATA.csv')
csv_data = csv_data.values
A = csv_data.shape[0]
board_data = csv_data[:,0:16]
# X = np.log2(X)
X = torch.FloatTensor(board_data)
X = np.int64(board_data)

# 转置后拼接
X = np.reshape(X, (-1,4,4))
XT = X.transpose(0,2,1)

X = np.concatenate((X,XT),axis=1)
print(X.shape)

direction_data = csv_data[:,16]
Y = np.int64(direction_data)



model = CCRNN()
model = model.cuda()
optimizer = optim.Adam(model.parameters(), lr = 0.001)

def train(epoch):
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = Variable(data).cuda(), Variable(target).cuda()
        data = data/11.0
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 50 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\t Loss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
    torch.save(self.model, 'rnn_model_' + str(epoch) + '.pkl')


if __name__ == '__main__':
    for epoch in range(0, NUM_EPOCHS):
        train(epoch)
[((83, 11, 83, 51), 'pandas.read_csv', 'pd.read_csv', ({(83, 23, 83, 50): '"""./drive/My Drive/DATA.csv"""'}, {}), "('./drive/My Drive/DATA.csv')", True, 'import pandas as pd\n'), ((88, 4, 88, 33), 'torch.FloatTensor', 'torch.FloatTensor', ({(88, 22, 88, 32): 'board_data'}, {}), '(board_data)', False, 'import torch\n'), ((89, 4, 89, 24), 'numpy.int64', 'np.int64', ({(89, 13, 89, 23): 'board_data'}, {}), '(board_data)', True, 'import numpy as np\n'), ((92, 4, 92, 27), 'numpy.reshape', 'np.reshape', ({(92, 15, 92, 16): 'X', (92, 18, 92, 26): '(-1, 4, 4)'}, {}), '(X, (-1, 4, 4))', True, 'import numpy as np\n'), ((95, 4, 95, 33), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((99, 4, 99, 28), 'numpy.int64', 'np.int64', ({(99, 13, 99, 27): 'direction_data'}, {}), '(direction_data)', True, 'import numpy as np\n'), ((105, 35, 105, 86), 'sklearn.model_selection.train_test_split', 'train_test_split', (), '', False, 'from sklearn.model_selection import train_test_split\n'), ((106, 10, 106, 36), 'torch.FloatTensor', 'torch.FloatTensor', ({(106, 28, 106, 35): 'X_train'}, {}), '(X_train)', False, 'import torch\n'), ((107, 9, 107, 34), 'torch.FloatTensor', 'torch.FloatTensor', ({(107, 27, 107, 33): 'X_test'}, {}), '(X_test)', False, 'import torch\n'), ((108, 10, 108, 35), 'torch.LongTensor', 'torch.LongTensor', ({(108, 27, 108, 34): 'Y_train'}, {}), '(Y_train)', False, 'import torch\n'), ((109, 9, 109, 33), 'torch.LongTensor', 'torch.LongTensor', ({(109, 26, 109, 32): 'Y_test'}, {}), '(Y_test)', False, 'import torch\n'), ((111, 16, 111, 63), 'torch.utils.data.TensorDataset', 'torch.utils.data.TensorDataset', ({(111, 47, 111, 54): 'X_train', (111, 55, 111, 62): 'Y_train'}, {}), '(X_train, Y_train)', False, 'import torch\n'), ((114, 15, 117, 1), 'torch.utils.data.DataLoader', 'torch.utils.data.DataLoader', (), '', False, 'import torch\n'), ((131, 11, 131, 51), 'pandas.read_csv', 'pd.read_csv', ({(131, 23, 131, 50): '"""./drive/My Drive/DATA.csv"""'}, {}), "('./drive/My Drive/DATA.csv')", True, 'import pandas as pd\n'), ((136, 4, 136, 33), 'torch.FloatTensor', 'torch.FloatTensor', ({(136, 22, 136, 32): 'board_data'}, {}), '(board_data)', False, 'import torch\n'), ((137, 4, 137, 24), 'numpy.int64', 'np.int64', ({(137, 13, 137, 23): 'board_data'}, {}), '(board_data)', True, 'import numpy as np\n'), ((140, 4, 140, 27), 'numpy.reshape', 'np.reshape', ({(140, 15, 140, 16): 'X', (140, 18, 140, 26): '(-1, 4, 4)'}, {}), '(X, (-1, 4, 4))', True, 'import numpy as np\n'), ((143, 4, 143, 33), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((147, 4, 147, 28), 'numpy.int64', 'np.int64', ({(147, 13, 147, 27): 'direction_data'}, {}), '(direction_data)', True, 'import numpy as np\n'), ((25, 22, 32, 9), 'torch.nn.LSTM', 'nn.LSTM', (), '', True, 'import torch.nn as nn\n'), ((34, 23, 40, 9), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((42, 23, 48, 9), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((50, 23, 56, 9), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((58, 25, 58, 41), 'torch.nn.Linear', 'nn.Linear', ({(58, 35, 58, 36): '4', (58, 38, 58, 40): '64'}, {}), '(4, 64)', True, 'import torch.nn as nn\n'), ((59, 25, 59, 41), 'torch.nn.Linear', 'nn.Linear', ({(59, 35, 59, 36): '4', (59, 38, 59, 40): '32'}, {}), '(4, 32)', True, 'import torch.nn as nn\n'), ((60, 25, 60, 41), 'torch.nn.Linear', 'nn.Linear', ({(60, 35, 60, 36): '4', (60, 38, 60, 40): '32'}, {}), '(4, 32)', True, 'import torch.nn as nn\n'), ((61, 24, 61, 43), 'torch.nn.Linear', 'nn.Linear', ({(61, 34, 61, 37): '256', (61, 39, 61, 42): '128'}, {}), '(256, 128)', True, 'import torch.nn as nn\n'), ((62, 23, 62, 40), 'torch.nn.Linear', 'nn.Linear', ({(62, 33, 62, 36): '256', (62, 38, 62, 39): '4'}, {}), '(256, 4)', True, 'import torch.nn as nn\n'), ((66, 17, 66, 49), 'torch.unsqueeze', 'torch.unsqueeze', ({(66, 33, 66, 45): 'x[:, 0:4, :]', (66, 47, 66, 48): '1'}, {}), '(x[:, 0:4, :], 1)', False, 'import torch\n'), ((78, 21, 78, 78), 'torch.cat', 'torch.cat', ({(78, 31, 78, 74): '(CNN22toFC, CNN41toFC, CNN14toFC, LSTMtoFC)', (78, 76, 78, 77): '1'}, {}), '((CNN22toFC, CNN41toFC, CNN14toFC, LSTMtoFC), 1)', False, 'import torch\n'), ((161, 15, 161, 46), 'torch.nn.functional.cross_entropy', 'F.cross_entropy', ({(161, 31, 161, 37): 'output', (161, 39, 161, 45): 'target'}, {}), '(output, target)', True, 'import torch.nn.functional as F\n'), ((157, 23, 157, 37), 'torch.autograd.Variable', 'Variable', ({(157, 32, 157, 36): 'data'}, {}), '(data)', False, 'from torch.autograd import Variable\n'), ((157, 46, 157, 62), 'torch.autograd.Variable', 'Variable', ({(157, 55, 157, 61): 'target'}, {}), '(target)', False, 'from torch.autograd import Variable\n')]