import torch
import torchvision
import torchvision.transforms as transforms
import os

from skimage import data,io
from numpy import random


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

#'/opt/samples/S6/casia_like_light/CASIA-maxpy-clean'

import torch.utils.data as data
class FaceTrainDataset(data.Dataset):
    def __init__(self, transform=None,path='/opt/samples/S6/casia_like_light/CASIA-maxpy-clean/'):
        self.transform = transform
	dirs = os.listdir(path)
	self.file_list=[]

	class_ind = 0
	for c_dir in dirs:
	  filez = os.listdir(path+c_dir)
	  for c_file in filez:
	    self.file_list.append((path+c_dir+'/'+c_file,class_ind))
	  class_ind+=1
	random.shuffle(self.file_list)

    def __getitem__(self, index):
	fname,clasID = self.file_list[index]
        img = io.imread(fname)
        img = self.transform(img)
        return img, clasID

    def __len__(self):
        return len(self.file_list)

#Train
trainset = FaceTrainDataset(transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=40,
                                          shuffle=True, num_workers=2)


# CNN modules and primitives
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(17424, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10575)
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 17424)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


class Net_modules(nn.Module):
    def __init__(self):
        super(Net_modules, self).__init__()
        self.add_module("l1",nn.Conv2d(3, 6, 5))
	self.add_module("relu_1", torch.nn.ReLU())
        self.add_module("l2",nn.MaxPool2d(2, 2))
        self.add_module("l3",nn.Conv2d(6, 16, 5))
	self.add_module("relu_2", torch.nn.ReLU())
        self.add_module("l4",nn.MaxPool2d(2, 2))
	
        self.add_module("l5",nn.Linear(17424, 120))
        self.add_module("l6",nn.Linear(120, 84))
        self.add_module("l7",nn.Linear(84, 10575))

    def forward(self, x):
	layers = [i for i in self.children()]
	for i in range(6):
	  x=layers[i](x)
	print x.shape
        x = x.view(-1, 17424)
	for i in range(6,9):
	  x=layers[i](x)
        return x

class Net_sequentional(nn.Module):
    def __init__(self):
        super(Net_sequentional, self).__init__()
	self.conv_layers = nn.Sequential()
        self.conv_layers.add_module("l1",nn.Conv2d(3, 6, 5))
	self.conv_layers.add_module("relu_1", torch.nn.ReLU())
        self.conv_layers.add_module("l2",nn.MaxPool2d(2, 2))
        self.conv_layers.add_module("l3",nn.Conv2d(6, 16, 5))
	self.conv_layers.add_module("relu_2", torch.nn.ReLU())
        self.conv_layers.add_module("l4",nn.MaxPool2d(2, 2))
	
	self.fc_layers = nn.Sequential()
        self.fc_layers.add_module("l5",nn.Linear(17424, 120))
        self.fc_layers.add_module("l6",nn.Linear(120, 84))
        self.fc_layers.add_module("l7",nn.Linear(84, 10575))

    def forward(self, x):
	x=self.conv_layers(x)
        x = x.view(-1, 17424)
	x = self.fc_layers(x)
        return x

class PoolRELUConv(nn.Module):
    def __init__(self,in_size,count,width):
	super(PoolRELUConv,self).__init__()
	self.conv = nn.Conv2d(in_size,count,width)
	self.pool=nn.MaxPool2d(2, 2)
	 
    def forward(self, x):
	return self.pool(F.relu(self.conv(x)))

class Net_module(nn.Module):
    def __init__(self):
        super(Net_module, self).__init__()
        self.PoolRELUConv1 = PoolRELUConv(3, 64, 5)
        self.PoolRELUConv2 = PoolRELUConv(64, 128, 5)

        self.fc1 = nn.Linear(17424*8, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10575)
    def forward(self, x):
        x = self.PoolRELUConv1(x)
        x = self.PoolRELUConv2(x)
        x = x.view(-1, 17424*8)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


def save_chekpoint(model,optimizer,path):
	torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            }, path)


def load_chekpoint(model,optimizer,path,train=True):
	checkpoint = torch.load(PATH)
	model.load_state_dict(checkpoint['model_state_dict'])
	optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
	if train:
	   model.train()
	else:
	   model.eval()

	return checkpoint['epoch']


net = Net()
print net.state_dict()
#raise()

net = nn.DataParallel(net) #'!!!!'
net = net.cuda() # GPU!!!

import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.000001, momentum=0.9)
lrs =[]

from matplotlib import pyplot as plt


def adjust_learning_rate(optimizer, epoch):
    scale = 0.33
    step  = 10
    lr = 0.000001 * (scale ** (epoch // step))
    print('lr: {}'.format(lr))
    if (epoch != 0) & (epoch % step == 0):
        print('Change lr')
        for param_group in optimizer.param_groups:
            param_group['lr'] = param_group['lr'] * scale

for epoch in range(2):  
        adjust_learning_rate(optimizer, epoch)
	for i,data in enumerate(trainloader, 0):
 	     inputs,labels = data
	     inputs = inputs.cuda()
             labels = labels.cuda()
	     optimizer.zero_grad()
	     outputs = net(inputs)
	#     print inputs.shape
	#     print outputs.shape
	#     print labels.shape	
	     loss = criterion(outputs, labels)
	     loss.backward()	
	     optimizer.step()
	     lrs.append(loss.item())
	     plt.plot(lrs)
	     plt.pause(0.05)

	     print loss.item()


