

from keras import regularizers
from keras.models import Model, Sequential
from keras.layers import Input, Activation, Conv3D, MaxPooling3D, ZeroPadding3D,UpSampling3D,Dropout, BatchNormalization, Flatten, GaussianDropout, Dense, concatenate, AveragePooling3D
from keras.layers import Conv2D, MaxPooling2D, ZeroPadding2D,UpSampling2D,Dropout, BatchNormalization, Flatten, GaussianDropout, Dense, concatenate, AveragePooling3D
from keras.layers.advanced_activations import LeakyReLU, ELU, PReLU


import keras.backend as K
from keras.layers.merge import Concatenate
import numpy as np
from scipy.misc import imread
from scipy.ndimage import convolve

import os

def block_b(in_layer,num_filters=24, kern=(3,3),padding=('same','same'),
        batchnorm=False,dropout=None, weight_reg=0,stride=(1,1), num_convs=2,residual=False):
    #print(kern,stride,padding)
    out_layer = in_layer
    for i in range(num_convs):
        out_layer=Conv2D(num_filters,kern, padding='same', activation='relu', 
                kernel_regularizer = regularizers.l1(weight_reg))(out_layer)            
    if residual is True:
        out_layer=Concatenate(axis=1)([out_layer,in_layer])

    out_layer=Conv2D(num_filters,kern, padding='valid', strides = stride, activation='relu',
            kernel_regularizer = regularizers.l1(weight_reg))(out_layer)
 
    if dropout is not None:
        out_layer=Dropout(dropout)(out_layer)
    
    if batchnorm:
        out_layer=BatchNormalization(axis=1)(out_layer)
        
    return out_layer

def build_model(image_size,num_classes):
    inputs =  Input(image_size)
    conv = block_b(inputs,num_filters=8, kern=(3,3),padding=('same','same'),
        batchnorm=False,dropout=0.2, weight_reg=0,stride=(1,1), num_convs=2,residual=True)
    conv = block_b(conv,num_filters=32, kern=(3,3),padding=('same','same'),
        batchnorm=False,dropout=0.2, weight_reg=0,stride=(1,1), num_convs=2,residual=True)
    conv = block_b(conv,num_filters=16, kern=(3,3),padding=('same','same'),
        batchnorm=False,dropout=0.2, weight_reg=0,stride=(1,1), num_convs=2,residual=True)
#    conv = block_b(conv,num_filters=16, kern=(3,3),padding=('same','same'),
#        batchnorm=False,dropout=0.2, weight_reg=0,stride=(1,1), num_convs=2,residual=True)

    flat = Flatten()(conv)
    
    dense = Dense(128,kernel_regularizer=regularizers.l1(0.01))(flat)
    dense = LeakyReLU()(dense)
    dense = Dropout(0.2)(dense)
    dense = Dense(num_classes,activation='sigmoid')(dense)

    model = Model(inputs,dense)
    
    return model

def train_model(model,train,log_dir,weights_file):
    from keras.optimizers import Adam
    adam = Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)
    model.compile(optimizer=adam,loss=['mse'])

    from keras.callbacks import TensorBoard
    from keras.callbacks import EarlyStopping
    from keras.callbacks import ModelCheckpoint

    es_cb = EarlyStopping(monitor='val_loss', min_delta=0, patience=2, verbose=0, mode='auto')

    tb_cb = TensorBoard(log_dir=log_dir, histogram_freq=0, batch_size=32, write_graph=True, write_grads=False, write_images=False, embeddings_freq=0, embeddings_layer_names=None, embeddings_metadata=None)

    mc_cb = ModelCheckpoint(weights_file,save_best_only=True,monitor='val_loss')
    model.fit(train[0],
            train[1],
            batch_size=61,shuffle=True, 
            callbacks=[tb_cb,es_cb,mc_cb],
            epochs=1000,validation_split=0.25)
    return model

def get_image(str, downsample=False):
    im = imread(str,flatten=False,mode='RGB').T#channels first
    if downsample:
        im = np.stack([convolve(im[i,:,:],np.array([[0.25,0.25],[0.25,0.25]]), mode='nearest')[:im.shape[1]:2,:im.shape[2]:2] for i in range(3)],axis=0)
    return im

def predict_all(img_folder,model):
    import glob
    files = glob.glob("{}/*.png".format(img_folder))
    inputs = []
    for f in files:
        inputs.append(get_image(f,downsample=True))

    inputs = np.stack(inputs,axis=0).astype(float)
    #normalize
    mean = np.mean(inputs.flatten())
    std = np.std(inputs.flatten())
    inputs -= mean
    inputs /= (std+1e-15)

    predict = model.predict(inputs)

    print("Classifications:")
    for i,f in enumerate(files):
        p = predict[i,:]
        print("\tfile: {:<10}\t double: {:.2f}\tsingle: {:.2f}".format(f,p[0],p[1]))

def load_model_from_run_dir(dir):
    from keras.models import model_from_json
    f = open("{}/model-arch.json".format(dir),'r')
    model = model_from_json(f.read())
    f.close()
    model.load_weights("{}/weights.hdf5".format(dir))
    return model

def make_batches_from_protos():
    yes = np.genfromtxt('data/hatfield_multiple_20x20_ID19_prototype_6_0.csv',skip_header=1,delimiter=',')[:,18]
    no = np.genfromtxt('data/hatfield_multiple_20x20_ID19_prototype_0_0.csv',skip_header=1,delimiter=',')[:,18]
    inputs = []
    outputs = []
    for i in range(len(yes)):
        file = "data/img/{}.png".format(int(yes[i]))
        inputs.append(get_image(file,downsample=True))
        outputs.append([1,0])
    for i in range(len(no)):
        file = "data/img/{}.png".format(int(no[i]))
        inputs.append(get_image(file,downsample=True))
        outputs.append([0,1])

    inputs = np.stack(inputs,axis=0).astype(float)
    #normalize
    mean = np.mean(inputs.flatten())
    std = np.std(inputs.flatten())
    inputs -= mean
    inputs /= (std+1e-15)
    outputs = np.stack(outputs,axis=0).astype(float)

    yes = np.genfromtxt('data/hatfield_multiple_20x20_ID19_prototype_7_0.csv',skip_header=1,delimiter=',')[:,18]
    no = np.genfromtxt('data/hatfield_multiple_20x20_ID19_prototype_0_17.csv',skip_header=1,delimiter=',')[:,18]
    val_inputs = []
    val_outputs = []
    for i in range(len(yes)):
        file = "data/img/{}.png".format(int(yes[i]))
        val_inputs.append(get_image(file,downsample=True))
        val_outputs.append([1,0])
    for i in range(len(no)):
        file = "data/img/{}.png".format(int(no[i]))
        val_inputs.append(get_image(file,downsample=True))
        val_outputs.append([0,1])

    valinputs = np.stack(val_inputs,axis=0).astype(float)
    #normalize
    val_inputs -= mean
    val_inputs /= (std+1e-15)
    val_outputs = np.stack(val_outputs,axis=0).astype(float)

def make_batch_from_classifications():
    names = ["Josh","Alex","Huub","Duy","Erik","Pedro"]
    inputs = []
    outputs = []
    for name in names:
        d = np.genfromtxt('data/classification - {}.csv'.format(name),
                skip_header=1,filling_values=0.,delimiter=',')
        #d = np.pad(d,((0,0),(0,9-d.shape[1])),'constant')
        mask = np.sum(d[:,1:],axis=1) > 0
        d = d[mask,:]
        classifications = np.zeros(d.shape[0])
        for i in range(d.shape[1]-1):
            classifications += d[:,i+1]*(1<<i)
        classifications = (np.arange(1<<7) == classifications[..., None]-1).astype(int)
        for i in range(d.shape[0]):
            file = "data/img/{:05d}.png".format(int(d[i,0]))
            inputs.append(get_image(file,downsample=True))
            #single/no, double/no, core/no, diffuse bridge/no, symmetric/no, dd/no, UC/no
            #outputs.append(classifications[i,:])
            outputs.append(d[i,1:])
    inputs = np.stack(inputs,axis=0).astype(float)
    inputs -= np.mean(inputs.flatten())
    inputs /= np.std(inputs.flatten()) + 1e-10
    outputs = np.stack(outputs,axis=0).astype(float)
#    class_weights_ = np.sum(outputs,axis=0)/np.sum(outputs.flatten())
#    class_weights = 1./class_weights_
#    class_weights[class_weights_ == 0] = 0.
#    class_weights = {i : class_weights[i] for i in range(int(np.size(class_weights)))}
    return inputs, outputs#, class_weights

if __name__=='__main__':
#    inputs = np.random.uniform(size=(64,1,128,128))
#    outputs = np.random.randint(2,size=(64))
#    outputs = (np.arange(outputs.max()+1) == outputs[..., None]-1).astype(int)
    inputs, outputs = make_batch_from_classifications()
    
    print("Train Data shapes:",inputs.shape,outputs.shape)
    data = (inputs,outputs)
    output_folder = 'output_lofargroup_sigmoid'
    
    for i in range(100):
        run_dir = '{}/run{}'.format(output_folder,i)
        log_dir = '{}/logs/run{}'.format(output_folder,i)
        try:
            os.makedirs(run_dir)
            os.makedirs(log_dir)
        except:
            continue
        model = build_model(inputs.shape[1:],outputs.shape[1])
        f = open("{}/model-arch.json".format(run_dir),"w")
        f.write(model.to_json())
        f.close()
        model = train_model(model,data,log_dir,"{}/weights.hdf5".format(run_dir)) 
        #model = load_model_from_run_dir(run_dir)
        #predict_all("data/img",model)
