# -*- coding: utf-8 -*-

"""
Created on Tue Feb 28 15:49:24 2017

@author: justjay
"""
import numpy as np
import random
import argparse
from iterator_cc import SerialIterator

import chainer
from chainer import cuda
from chainer import serializers
from chainer import optimizers
from chainer import datasets
from chainer import iterators
from chainer import training 
from chainer.training import extensions

import cv2
from sklearn.feature_extraction.image import extract_patches

from model import Model

parser = argparse.ArgumentParser(description='train.py')
parser.add_argument('--model', '-m', default='', 
                    help='path to the trained model')
parser.add_argument('--gpu', '-g', default=0, type=int, help='GPU ID')

args = parser.parse_args()


model = Model()

cuda.cudnn_enabled = True
cuda.check_cuda_available()
xp = cuda.cupy
model.to_gpu()
## serializers.load_hdf5(args.model, model)

## prepare training data 
train_pair_path = 'data/Annotations/train_label.txt'
image_path = 'data/Imageset/live2/'
save_model_path = 'models/nr_jay_live2.model'
file_path = './patches_weight.txt'

## get the train dataset:(image_array, label), in the order of label
train = datasets.LabeledImageDataset(train_pair_path, image_path, dtype=np.float32, label_dtype=np.int32)
## val = datasets.LabeledImageDataset(val_pair_path, image_path, dtype=np.float32, label_dtype=np.int32)

train_img_num = len(train)
## val_img_num = len(val)

## print img_num
patches_per_img = 256
## extract all the images to 32x32 pixls

f_label = open(train_pair_path, 'r')
f_patch = open(file_path, 'r')

list_label = [line.split()[0] for line in f_label.readlines()]          ## get names of training images
list_patch_img = [line.split()[0] for line in f_patch.readlines()]      ## get names of image patches
f_patch.seek(0)                                                         ## reset file pointer       
list_patch = [line.split() for line in f_patch.readlines()]             ## get patches info

f_label.close()
f_patch.close()

print '-------------Load data-------------'
train_patches = []
for i in range(train_img_num):
    patches = extract_patches(train[i][0], (3,32,32), 32)
    temp = patches.reshape((-1, 3, 32, 32))
    
    ## select 128 top patches/image
    line_index = list_patch_img.index(list_label[i])                    ## find the location of corresponding patches
    ## print i, line_index, list_patch_img[line_index], len(temp)
        
    temp_slice = [temp[int(index)] for index in list_patch[line_index][1:patches_per_img+1]]
    
    ## print temp_slice    
    ## temp_slice = random.sample(temp, patches_per_img)    
    for j in range(len(temp_slice)):
        temp_slice[j] = xp.array(temp_slice[j].astype(np.float32))      
        train_patches.append((temp_slice[j], train[i][1]))

print '--------------Done!----------------'
        
## use my own SerialIterator, pick image patches randomly
train_iter = SerialIterator(train_patches, patches = patches_per_img, batch_size=4)
## val_iter = SerialIterator(val_patches, patches = patches_per_img, batch_size=4, repeat=False, shuffle=False)
optimizer = optimizers.Adam()
## optimizer = optimizers.MomentumSGD()
optimizer.use_cleargrads()
optimizer.setup(model)

updater = training.StandardUpdater(train_iter, optimizer, device=0)
trainer = training.Trainer(updater, (8000, 'epoch'), out='result')

trainer.extend(extensions.LogReport())

trainer.extend(extensions.PrintReport(['epoch', 'iteration', 'main/loss', 'elapsed_time']))

## trainer.extend(extensions.ProgressBar())

trainer.run()
serializers.save_hdf5(save_model_path, model)


    
