# coding: utf8
# Author: Shen Qu

import os
import time

import numpy as np
from keras.applications.mobilenet import MobileNet
from keras.callbacks import ModelCheckpoint
from keras.layers import Input, Dense, GlobalAveragePooling2D
from keras.models import Model
from keras.preprocessing import image
from keras.preprocessing.image import ImageDataGenerator


class GoodClassifier(object):
    """A neural network for image classification.

    Notes:
        If train_data_path=None, model and weight will be loaded automatically.

    Args:
        train_data_path: Where the train data stored.

    """

    def __init__(self, shopping_list_path, weight_path=None, train_data_path=None):
        self.input_width = 299
        self.input_height = 299

        self.train_data_path = train_data_path

        with open(shopping_list_path) as file:
            self.shopping_list = file.readlines()

        self.shopping_list = [string.strip() for string in self.shopping_list]

        self.object_list = self.shopping_list.copy()

        self.object_list.remove('yellow cube')

        self.object_list.append('coca cola')
        self.object_list.append('pepsi cola')
        self.object_list.append('red ping pang ball')
        self.object_list.append('tonic water')
        self.object_list.append('sprite')
        self.object_list.append('gel pen')
        self.object_list.append('wang lao ji')
        self.object_list.append('xue hua')
        self.object_list.append('loofah')
        self.object_list.append('red bull')

        self.model = None

        if self.train_data_path:
            for name in self.object_list:
                if not os.path.exists(os.path.join(self.train_data_path, name)):
                    os.mkdir(os.path.join(self.train_data_path, name))

        else:
            self.get_model()
            if weight_path:
                self.model.load_weights(weight_path)

    def get_model(self) -> None:
        """Load neural network model.

        Notes:
            Should be called before using "single_recognize".

        Returns:
             None.
        """
        print('Initializing Neural Network...')
        start = time.time()

        input_tensor = Input(shape=(self.input_height, self.input_width, 3))

        base_model = MobileNet(input_tensor=input_tensor, include_top=False, weights='imagenet')

        x = base_model.output
        x = GlobalAveragePooling2D()(x)
        x = Dense(1024, activation='relu')(x)
        predictions = Dense(len(self.object_list), activation='softmax')(x)

        self.model = Model(inputs=base_model.input, outputs=predictions)

        for layer in base_model.layers:
            layer.trainable = False

        self.model.compile(optimizer='rmsprop',
                           loss='categorical_crossentropy',
                           metrics=['categorical_accuracy'])

        print('Done.')
        print('Time eclipse: {0:.3f} seconds'.format(time.time() - start))

    def train(self, epoch: int = 1) -> None:
        """Train neural network.

        Args:
            epoch: Same meaning in Keras.
            
        Returns:
            None.
        """
        train_datagen = ImageDataGenerator(
            rotation_range=40,
            width_shift_range=0.2,
            height_shift_range=0.2,
            rescale=1. / 255,
            zoom_range=0.2,
            shear_range=0.2,
            horizontal_flip=True,
            vertical_flip=True,
            fill_mode='nearest')

        train_generator = train_datagen.flow_from_directory(
            self.train_data_path,
            target_size=(self.input_height, self.input_width),
            batch_size=32,
            classes=self.object_list,
            class_mode='categorical')

        self.model.fit_generator(
            train_generator,
            epochs=epoch)

        for layer in self.model.layers[:115]:
            layer.trainable = False
        for layer in self.model.layers[115:]:
            layer.trainable = True

        self.model.compile(optimizer='rmsprop',
                           loss='categorical_crossentropy',
                           metrics=['categorical_accuracy'])

        check_pointer = ModelCheckpoint(filepath="/weights.hdf5", verbose=1, save_best_only=True)

        self.model.fit_generator(
            train_generator,
            epochs=epoch,
            callbacks=[check_pointer])

    def single_recognize(self, image_in) -> (str, float):
        """Get recognize result for single image.

        Args:
            image_in: Image path.
            
        Returns:
            Name of good and confidence. e.g. ('badminton' , 0.8999).

        Raises:
            RuntimeError: If no model loaded.
            ValueError: If input in not str or ndarray.
        """

        start = time.time()

        if not self.model:
            raise RuntimeError('No model loaded. Use or "get_model" first.')

        if isinstance(image_in, str):
            img = image.load_img(image_in, target_size=(self.input_width, self.input_height))
            x = image.img_to_array(img)

        elif isinstance(image_in, np.ndarray):
            x = image_in

        else:
            raise ValueError('Input only can be numpy array or file path')

        x = x / 255.
        x = np.expand_dims(x, axis=0)
        prediction = self.model.predict(x)

        index = np.argmax(prediction[0])
        name = self.object_list[index]
        confidence = prediction[0][index]

        print('\033[0;33;40m')
        print('{}\t\t{:.4f}\t\t{:.4f}'.format(name, confidence, (time.time() - start)))

        if name in self.shopping_list:
            return name, confidence

        else:
            return 'other', confidence

    def batch_recognize(self, batch) -> list:
        pass
