import numpy as np
import pickle
import os

import matplotlib.pyplot as plt

from pathlib import Path
from PIL import Image

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import Input, Conv2D
from tensorflow.keras.models import Model
from tensorflow.keras.applications.vgg16 import VGG16

from sklearn.ensemble import RandomForestClassifier


class HDRModel:


    def __init__(self, shape=None, loadDir=None):
        self.featureModel = None
        self.rfClassifer  = None

        self.inputChannels  = 3
        self.outputChannels = 1

        self.__fmSaveDir  = "featureModel"
        self.__fmSaveFile = "fm.h5"
        self.__rfSaveDir  = "rfClassifer"
        self.__rfSaveFile = "classifier.sav"

        if loadDir is not None:
            self.load(loadDir)
        elif shape is not None:
            self.__initModels(shape)


    def save(self, saveDir):
        # Save tf model
        fmDir = saveDir.joinpath(self.__fmSaveDir)
        os.mkdir(fmDir)
        self.featureModel.save(str(fmDir.joinpath(self.__fmSaveFile)))
        # Save sklearn random forest
        classifierDir = saveDir.joinpath(self.__rfSaveDir)
        os.mkdir(classifierDir)
        pickle.dump(self.rfClassifer, open(classifierDir.joinpath(self.__rfSaveFile), "wb"))


    def load(self, modelDir):
        # Load tf model
        fmPath = modelDir.joinpath(self.__fmSaveDir).joinpath(self.__fmSaveFile)
        self.featureModel = keras.models.load_model(str(fmPath))
        self.featureModel.summary()

        # Load sklean model
        rfPath = modelDir.joinpath(self.__rfSaveDir).joinpath(self.__rfSaveFile)
        self.rfClassifer = pickle.load(open(rfPath, "rb"))


    def fit(self, xTrain, yTrain):
        # Get features from VGG16 built model
        features = self.__getFeatures(xTrain)

        # Reshape input for random forest
        # TO DO: sklearn split data to train and test
        xForestTrain = features.reshape(-1, features.shape[-1])
        yForestTrain = yTrain.reshape(-1)

        self.rfClassifer.fit(xForestTrain, yForestTrain)


    def predict(self, xInput):
        '''
        Supports only a single input: shape=(x, y, z)
        '''
        # Extract features from input
        xExpand  = np.expand_dims(xInput, axis=0)
        features = self.__getFeatures(xExpand)
        
        # Reshape input for random forest
        forestInput = features.reshape(-1, features.shape[-1])
        forestOutput = self.rfClassifer.predict(forestInput)
        
        # Reshape output back to image shape
        outputShape = (xInput.shape[0], xInput.shape[1])
        finalOutput = forestOutput.reshape(outputShape)
        return finalOutput


    def __getFeatures(self, xInput):
        return self.featureModel.predict(xInput)
    

    def __initModels(self, shape):
        self.inputChannels = shape[-1]
        vgg = VGG16(weights="imagenet", include_top=False, input_shape=shape)
        self.featureModel = Model(inputs=vgg.input, outputs=vgg.get_layer("block1_conv2").output)
        self.rfClassifer  = RandomForestClassifier(n_estimators=50, random_state=42, verbose=1, n_jobs=-1)