import torch
import torchvision
import torchvision.transforms as transforms

from skimage import data,io

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


import torch.utils.data as data
import os
class CifarTrainDataset(data.Dataset):
    def __init__(self, transform=None,path = '/opt/samples/S6/casia_like_light/CASIA-maxpy-clean/'):
        self.imglist = []
        dirs = os.listdir(path)
        class_id = 0
        for c_dir in dirs:
          files = os.listdir(path+c_dir)
          for c_file in files:
          	self.imglist.append((path+c_dir+'/'+c_file,class_id))
          class_id=class_id+1
               
        self.transform = transform

    def __getitem__(self, index):
    	path,class_id = self.imglist[index]
        img = io.imread(path)
        img = self.transform(img)
        return img, class_id

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

#Train
trainset = CifarTrainDataset(transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True, num_workers=2)
#Test

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

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, 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(69696/4, 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, 69696/4)
        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

net = Net_module()
print net.state_dict()
raise a


net = net.cuda() # GPU!!!

import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

def adjust_learning_rate(optimizer, epoch):
    scale = 0.33
    step  = 10
    lr = 0.001 * (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)
	     loss = criterion(outputs, labels)
	     loss.backward()	
	     optimizer.step()
	     print loss.item()


