from keras.models import Sequential
from keras.layers import Conv2D
from keras.layers import MaxPooling2D
from keras.layers import Flatten
from keras.layers import Dense
from keras.layers import concatenate
from keras.models import Model
from keras.utils import plot_model
from keras.engine.input_layer import Input
from keras.preprocessing.image import ImageDataGenerator
from keras.preprocessing.image import img_to_array
import numpy as np
from keras.preprocessing import image
from sklearn.preprocessing import MultiLabelBinarizer
import cv2
import os
import shutil

train_path = "D:\Work\Project\\data_train\\"
old_data_path = "D:\Work\Project\Dictionary_word\\"
number_of_copy = 5
def create_y_train():
    f = open("label.txt","r")
    count_line = 0
    label_list = []
    for line in f:
        label_list_string = line.strip('\n').split(" ")
        label_list_temp = []
        if count_line % 2 == 0:
            for i in label_list_string:
                label_list_temp.append(int(i))
            for i in range(number_of_copy):
                label_list.append(label_list_temp)
        count_line = count_line + 1
    mlb = MultiLabelBinarizer()
    return mlb.fit_transform(label_list)

def create_x_train():
    data_train = []
    for file in os.listdir(train_path):
        print (file)
        img = cv2.imread(train_path + file)
        img = cv2.resize(img,(32,32))
        img = img_to_array(img)
        data_train.append(img)
    return data_train


def save_file_to_train():
    file_number = 0
    for folder in os.listdir(old_data_path):
        path_img = old_data_path + folder + "\\"
        for file in os.listdir(path_img):
            shutil.copy2(path_img + file, train_path + "{:09d}".format(file_number) + ".bmp")
            file_number += 1
            # img = cv2.imread(path_img+file)
            # img = cv2.resize(img, (100, 50))

def create_model():
    # Initialising the CNN
    input_shape = Input(shape=(32, 32, 3))
    # Step 1 - Convolution
    feature_extract = Conv2D(32, (3, 3), activation = 'relu')(input_shape)
    feature_extract = MaxPooling2D(pool_size = (2, 2))(feature_extract)

    feature_extract = Conv2D(32, (3, 3), activation = 'relu')(feature_extract)
    feature_extract = MaxPooling2D(pool_size = (2, 2))(feature_extract)

    tone_mark_phase = Flatten()(feature_extract)
    tone_mark_phase = Dense(units = 128, activation = 'relu')(tone_mark_phase)
    tone_mark_phase = Dense(units = 6, activation = 'softmax')(tone_mark_phase)

    upper_vowel_phase = Flatten()(feature_extract)
    upper_vowel_phase = Dense(units = 128, activation = 'relu')(upper_vowel_phase)
    upper_vowel_phase = Dense(units = 6, activation = 'softmax')(upper_vowel_phase)

    alphabet_phase = Flatten()(feature_extract)
    alphabet_phase = Dense(units = 128, activation = 'relu')(alphabet_phase)
    alphabet_phase = Dense(units = 56, activation = 'softmax')(alphabet_phase)

    lower_vowel_phase = Flatten()(feature_extract)
    lower_vowel_phase = Dense(units = 128, activation = 'relu')(lower_vowel_phase)
    lower_vowel_phase = Dense(units = 3, activation = 'softmax')(lower_vowel_phase)

    out = concatenate([tone_mark_phase, upper_vowel_phase, alphabet_phase,lower_vowel_phase], axis=1)

    tone_mark_model = Model(input_shape,tone_mark_phase)
    upper_vowel_model = Model(input_shape,upper_vowel_phase)
    alphabet_model = Model(input_shape,alphabet_phase)
    lower_vowel_model = Model(input_shape,lower_vowel_phase)

    o_model = Model(input_shape,out)

    plot_model(tone_mark_model, to_file="tone_mark.png")
    plot_model(upper_vowel_model, to_file="upper.png")
    plot_model(alphabet_model, to_file="alphabet.png")
    plot_model(lower_vowel_model, to_file="lower.png")
    plot_model(o_model, to_file="o.png")
    return o_model


# save_file_to_train()
x_train = create_x_train()
y_train = create_y_train()
print ("x_train , y_train created !!")
model = create_model()
print ("model created !!")
model.compile(optimizer = 'adam', loss = 'binary_crossentropy', metrics = ['accuracy'])

train_datagen = ImageDataGenerator(rescale = 1./255,
                                   shear_range = 0.2,
                                   zoom_range = 0.2,
                                   horizontal_flip = True)
model.fit_generator(train_datagen.flow(x_train, y_train, batch_size=25))

model.save('test_run_data_alphabet_classification.model')