import os
import torch
import random
import numpy as np
import torchvision
import torchvision.transforms as transforms
from torch import nn
import torch.nn.functional as F

class TorchConfig(object):
    use_cuda = torch.cuda.is_available()

    def __init__(self, seed=2019):
        self.seed = seed
        self.device = torch.device('cuda' if self.use_cuda else 'cpu')  # Device configuration
        self.set_seed()
#         torch.set_default_tensor_type(torch.DoubleTensor)

    def set_seed(self):
        os.environ['PYTHONHASHSEED'] = str(self.seed)
        random.seed(self.seed)
        np.random.seed(self.seed)
        torch.manual_seed(self.seed)
        if self.use_cuda:
            print('GPU: %s' % torch.cuda.get_device_name(0))
            torch.cuda.manual_seed(self.seed)
            torch.backends.cudnn.deterministic = True

tc = TorchConfig()

# Hyper-parameters 
input_size = 784
hidden_size = 500
num_classes = 10
num_epochs = 5
batch_size = 100
learning_rate = 0.001
device = tc.device

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

train_dataset = torchvision.datasets.MNIST(root='./data', 
                                           train=True, 
                                           transform=transforms.ToTensor())

test_dataset = torchvision.datasets.MNIST(root='./data', 
                                          train=False, 
                                          transform=transforms.ToTensor())

# Data loader (input pipeline)
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)

# Fully connected neural network with one hidden layer
class NN(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden_size) 
        self.fc2 = nn.Linear(hidden_size, num_classes)  
    
    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        return x
    
model = NN(input_size, hidden_size, num_classes).to(device)
model

def train_nn(model, train_loader, num_epochs, loss_function=nn.CrossEntropyLoss(), lr=0.001):
    # Loss and optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr)
    
    # Train the model
    total_step = len(train_loader)
    for epoch in range(num_epochs):
        for i, (X, y) in enumerate(train_loader, 1):  
            # Move tensors to the configured device
            X = X.reshape(-1, 28*28).to(device) # 需重写
            y = y.to(device)
            # Forward pass
            loss = loss_function(model(X), y)

            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if i % 100 == 0:
                print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                       .format(epoch+1, num_epochs, i, total_step, loss.item()))

train_nn(model, train_loader, num_epochs)

# Test the model
# In test phase, we don't need to compute gradients (for memory efficiency)
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.reshape(-1, 28*28).to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item() # https://blog.csdn.net/dss_dssssd/article/details/83818181

    print('Accuracy of the network on the 10000 test images: {} %'.format(100 * correct / total))

# Save the model checkpoint
torch.save(model.state_dict(), 'model.ckpt')

from yuan.utils.jupyter import commit

commit(nb_filename='./Pytorch.ipynb')

import torch
import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        # 输入图像channel：1；输出channel：6；5x5卷积核
        self.conv1 = nn.Conv2d(1, 6, 5) # (32-5+2*0) / 1+1
        self.conv2 = nn.Conv2d(6, 16, 5)
        # an affine operation: y = Wx + b
        self.fc1 = nn.Linear(16 * 4 * 4, 120) # 根据网络定义参数
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # 2x2 Max pooling
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        # If the size is a square you can only specify a single number
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, torch.numel(x)//x.shape[0]) # self.num_flat_features(x)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

def train_nn(model, train_loader, num_epochs, loss_function=nn.CrossEntropyLoss(), optimizer=None):
    # Loss and optimizer
    if optimizer is None:
        lr=0.001
        optimizer = torch.optim.Adam(model.parameters(), lr)
    
    # Train the model
    total_step = len(train_loader)
    for epoch in range(num_epochs):
        for i, (X, y) in enumerate(train_loader, 1):  
            # Move tensors to the configured device
            X = X.to(device) # 需重写
            y = y.to(device)
            # Forward pass
            loss = loss_function(model(X), y)

            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if i % 100 == 0:
                print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                       .format(epoch+1, num_epochs, i, total_step, loss.item()))

model = Net().to(device)
train_nn(model, train_loader, num_epochs)

# Print model's state_dict
print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

from torch.utils.data import Dataset, DataLoader, TensorDataset

word_to_ix = {"hello": 0, "world": 1}
embeds = nn.Embedding(2, 5)  # 2 words in vocab, 5 dimensional embeddings
lookup_tensor = torch.tensor([word_to_ix["hello"]], dtype=torch.long)
hello_embed = embeds(lookup_tensor)
print(hello_embed)