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)
[((2241, 2281), 'pandas.read_csv', 'pd.read_csv', (['"""./drive/My Drive/DATA.csv"""'], {}), "('./drive/My Drive/DATA.csv')\n", (2252, 2281), True, 'import pandas as pd\n'), ((2387, 2416), 'torch.FloatTensor', 'torch.FloatTensor', (['board_data'], {}), '(board_data)\n', (2404, 2416), False, 'import torch\n'), ((2422, 2442), 'numpy.int64', 'np.int64', (['board_data'], {}), '(board_data)\n', (2430, 2442), True, 'import numpy as np\n'), ((2459, 2484), 'numpy.reshape', 'np.reshape', (['X', '(-1, 4, 4)'], {}), '(X, (-1, 4, 4))\n', (2469, 2484), True, 'import numpy as np\n'), ((2515, 2546), 'numpy.concatenate', 'np.concatenate', (['(X, XT)'], {'axis': '(1)'}), '((X, XT), axis=1)\n', (2529, 2546), True, 'import numpy as np\n'), ((2601, 2625), 'numpy.int64', 'np.int64', (['direction_data'], {}), '(direction_data)\n', (2609, 2625), True, 'import numpy as np\n'), ((2728, 2780), 'sklearn.model_selection.train_test_split', 'train_test_split', (['X', 'Y'], {'test_size': '(0.2)', 'shuffle': '(False)'}), '(X, Y, test_size=0.2, shuffle=False)\n', (2744, 2780), False, 'from sklearn.model_selection import train_test_split\n'), ((2791, 2817), 'torch.FloatTensor', 'torch.FloatTensor', (['X_train'], {}), '(X_train)\n', (2808, 2817), False, 'import torch\n'), ((2828, 2853), 'torch.FloatTensor', 'torch.FloatTensor', (['X_test'], {}), '(X_test)\n', (2845, 2853), False, 'import torch\n'), ((2865, 2890), 'torch.LongTensor', 'torch.LongTensor', (['Y_train'], {}), '(Y_train)\n', (2881, 2890), False, 'import torch\n'), ((2901, 2925), 'torch.LongTensor', 'torch.LongTensor', (['Y_test'], {}), '(Y_test)\n', (2917, 2925), False, 'import torch\n'), ((2945, 2993), 'torch.utils.data.TensorDataset', 'torch.utils.data.TensorDataset', (['X_train', 'Y_train'], {}), '(X_train, Y_train)\n', (2975, 2993), False, 'import torch\n'), ((3075, 3166), 'torch.utils.data.DataLoader', 'torch.utils.data.DataLoader', ([], {'dataset': 'train_dataset', 'batch_size': 'batch_size', 'shuffle': '(True)'}), '(dataset=train_dataset, batch_size=batch_size,\n    shuffle=True)\n', (3102, 3166), False, 'import torch\n'), ((3506, 3546), 'pandas.read_csv', 'pd.read_csv', (['"""./drive/My Drive/DATA.csv"""'], {}), "('./drive/My Drive/DATA.csv')\n", (3517, 3546), True, 'import pandas as pd\n'), ((3652, 3681), 'torch.FloatTensor', 'torch.FloatTensor', (['board_data'], {}), '(board_data)\n', (3669, 3681), False, 'import torch\n'), ((3687, 3707), 'numpy.int64', 'np.int64', (['board_data'], {}), '(board_data)\n', (3695, 3707), True, 'import numpy as np\n'), ((3724, 3749), 'numpy.reshape', 'np.reshape', (['X', '(-1, 4, 4)'], {}), '(X, (-1, 4, 4))\n', (3734, 3749), True, 'import numpy as np\n'), ((3780, 3811), 'numpy.concatenate', 'np.concatenate', (['(X, XT)'], {'axis': '(1)'}), '((X, XT), axis=1)\n', (3794, 3811), True, 'import numpy as np\n'), ((3866, 3890), 'numpy.int64', 'np.int64', (['direction_data'], {}), '(direction_data)\n', (3874, 3890), True, 'import numpy as np\n'), ((490, 584), 'torch.nn.LSTM', 'nn.LSTM', ([], {'input_size': '(4)', 'hidden_size': '(128)', 'num_layers': '(4)', 'bidirectional': '(True)', 'batch_first': '(True)'}), '(input_size=4, hidden_size=128, num_layers=4, bidirectional=True,\n    batch_first=True)\n', (497, 584), True, 'import torch.nn as nn\n'), ((685, 761), 'torch.nn.Conv2d', 'nn.Conv2d', ([], {'in_channels': '(1)', 'out_channels': '(1)', 'kernel_size': '(2)', 'stride': '(2)', 'padding': '(0)'}), '(in_channels=1, out_channels=1, kernel_size=2, stride=2, padding=0)\n', (694, 761), True, 'import torch.nn as nn\n'), ((864, 949), 'torch.nn.Conv2d', 'nn.Conv2d', ([], {'in_channels': '(1)', 'out_channels': '(1)', 'kernel_size': '(1, 4)', 'stride': '(1)', 'padding': '(0)'}), '(in_channels=1, out_channels=1, kernel_size=(1, 4), stride=1,\n    padding=0)\n', (873, 949), True, 'import torch.nn as nn\n'), ((1048, 1133), 'torch.nn.Conv2d', 'nn.Conv2d', ([], {'in_channels': '(1)', 'out_channels': '(1)', 'kernel_size': '(4, 1)', 'stride': '(1)', 'padding': '(0)'}), '(in_channels=1, out_channels=1, kernel_size=(4, 1), stride=1,\n    padding=0)\n', (1057, 1133), True, 'import torch.nn as nn\n'), ((1234, 1250), 'torch.nn.Linear', 'nn.Linear', (['(4)', '(64)'], {}), '(4, 64)\n', (1243, 1250), True, 'import torch.nn as nn\n'), ((1277, 1293), 'torch.nn.Linear', 'nn.Linear', (['(4)', '(32)'], {}), '(4, 32)\n', (1286, 1293), True, 'import torch.nn as nn\n'), ((1320, 1336), 'torch.nn.Linear', 'nn.Linear', (['(4)', '(32)'], {}), '(4, 32)\n', (1329, 1336), True, 'import torch.nn as nn\n'), ((1362, 1381), 'torch.nn.Linear', 'nn.Linear', (['(256)', '(128)'], {}), '(256, 128)\n', (1371, 1381), True, 'import torch.nn as nn\n'), ((1406, 1423), 'torch.nn.Linear', 'nn.Linear', (['(256)', '(4)'], {}), '(256, 4)\n', (1415, 1423), True, 'import torch.nn as nn\n'), ((1524, 1556), 'torch.unsqueeze', 'torch.unsqueeze', (['x[:, 0:4, :]', '(1)'], {}), '(x[:, 0:4, :], 1)\n', (1539, 1556), False, 'import torch\n'), ((2055, 2112), 'torch.cat', 'torch.cat', (['(CNN22toFC, CNN41toFC, CNN14toFC, LSTMtoFC)', '(1)'], {}), '((CNN22toFC, CNN41toFC, CNN14toFC, LSTMtoFC), 1)\n', (2064, 2112), False, 'import torch\n'), ((4250, 4281), 'torch.nn.functional.cross_entropy', 'F.cross_entropy', (['output', 'target'], {}), '(output, target)\n', (4265, 4281), True, 'import torch.nn.functional as F\n'), ((4100, 4114), 'torch.autograd.Variable', 'Variable', (['data'], {}), '(data)\n', (4108, 4114), False, 'from torch.autograd import Variable\n'), ((4123, 4139), 'torch.autograd.Variable', 'Variable', (['target'], {}), '(target)\n', (4131, 4139), False, 'from torch.autograd import Variable\n')]