# import the necessary packages
#Example python prediction.py -c ../conf/confMelanomaDataAugmentation.conf -i ../pos.jpg -f overfeat -p [-3] -m svm
from __future__ import print_function
from __future__ import absolute_import

from frimcla.extractor.extractor import Extractor
from frimcla.utils.conf import Conf
from frimcla.utils import dataset
from frimcla.index_features import extractFeatures
from .shallowmodels.classificationModelFactory import classificationModelFactory
from sklearn.model_selection import RandomizedSearchCV
from sklearn.preprocessing import LabelEncoder
from imutils import paths
from scipy import stats
import numpy as np
import argparse
import _pickle as cPickle
import h5py
import os
import json
import random


def prediction(featExt, classi, imagePath, outputPath, datasetPath):
    # load the configuration, label encoder, and classifier
    print("[INFO] loading model...")
    datasetName = datasetPath[datasetPath.rfind("/"):]
    auxPath = outputPath + datasetName
    factory = classificationModelFactory()

    cPickleFile = auxPath + "/classifiers/classifier_" + featExt[0] + "_" + classi + ".cpickle"
    if os.path.isfile(cPickleFile):
        labelEncoderPath = auxPath + "/models/le.cpickle"
        le = cPickle.loads(open(labelEncoderPath, "rb").read())
        model = cPickle.loads(open(cPickleFile, "rb").read())

    else: #Aqui es donde deberia hacer la pregunta de que si quiere realmente entrenar ese modelo o el mejor
        print("This is not the best model. Are you sure you want to predict with it?")
        response = input()
        if (response.upper() in ("YES", "Y")):
            files = os.walk(auxPath + "/models")
            listFeatExt = []
            for _, _, file in files:
                for feat in file:
                    feat = feat.split("-")[1].split(".")[0]
                    listFeatExt.append(feat)
            listFeatExt = list(set(listFeatExt))
            if (featExt[0] in listFeatExt):
                labelEncoderPath = auxPath + "/models/le.cpickle"
                le = cPickle.loads(open(labelEncoderPath,"rb").read())
                featuresPath = auxPath + "/models/features-" + featExt[0] + ".hdf5"
                db = h5py.File(featuresPath)
                split = int(db["image_ids"].shape[0])
                (trainData, trainLabels) = (db["features"][:split], db["image_ids"][:split])
                trainLabels = [le.transform([l.split(":")[0]])[0] for l in trainLabels]
                classifierModel = factory.getClassificationModel(classi)

                model = RandomizedSearchCV(classifierModel.getModel(), param_distributions=classifierModel.getParams(),
                                           n_iter=classifierModel.getNIterations())
                model.fit(trainData, trainLabels)
                f = open(auxPath + "/classifier_" + featExt[0] + "_" + classi + ".cpickle", "wb")
                f.write(cPickle.dumps(model))
                f.close()
                db.close()

            else:
                imagePaths = list(paths.list_images(datasetPath))
                random.seed(42)
                random.shuffle(imagePaths)
                le = LabelEncoder()
                le.fit([p.split("/")[-2] for p in imagePaths])
                extractFeatures(featExt, 32, imagePaths, outputPath, datasetPath, le, False)
                labelEncoderPath = auxPath + "/models/le.cpickle"
                le = cPickle.loads(open(labelEncoderPath,"rb").read())
                featuresPath = auxPath + "/models/features-" + featExt[0] + ".hdf5"
                db = h5py.File(featuresPath)
                split = int(db["image_ids"].shape[0])
                (trainData, trainLabels) = (db["features"][:split], db["image_ids"][:split])
                trainLabels = [le.transform([l.split(":")[0]])[0] for l in trainLabels]
                classifierModel = factory.getClassificationModel(classi)
                model = RandomizedSearchCV(classifierModel.getModel(), param_distributions=classifierModel.getParams(),
                                           n_iter=classifierModel.getNIterations())
                model.fit(trainData, trainLabels)
                f = open(auxPath + "/classifier_" + featExt[0] + "_" + classi + ".cpickle", "wb")
                f.write(cPickle.dumps(model))
                f.close()
                db.close()
        else:
            labelEncoderPath = auxPath + "/models/le.cpickle"
            le = cPickle.loads(open(labelEncoderPath,"rb").read())
            model = cPickle.loads(open(cPickleFile,"rb").read())

    filePrediction = open(auxPath+"/predictionResults.csv","a")
    filePrediction.write("image_id, melanoma\n")
    oe = Extractor(featExt)
    imagePaths = list(paths.list_images(imagePath))

    for (i, imPaths) in enumerate(dataset.chunk(imagePaths, 32)):
        (labels, images) = dataset.build_batch(imPaths, featExt[0])

        features = oe.describe(images)
        for (label, vector) in zip(labels, features):
            prediction = model.predict(np.atleast_2d(vector))[0]
            filePrediction.write( str(label) + ", " + str(prediction) + "\r\n")
            prediction = le.inverse_transform(prediction)
            print("[INFO] class predicted for the image {}: {}".format(label, prediction))

    filePrediction.close()


#This method only if you want to execute the prediction with the output of the previous steps
def predictionSimple(imagePath, outputPath, datasetPath):
    datasetName = datasetPath[datasetPath.rfind("/"):]
    auxPath = outputPath + datasetName
    with open(auxPath + "ConfModel.json") as json_file:
        data = json.load(json_file)

    extractor = data[0]['featureExtractor']
    classifier = data[0]['classifierModel']
    prediction(extractor,classifier,imagePath,outputPath,datasetPath)

def predictionEnsemble(featureExtractors, classifiers, imagePaths, outputPath, datasetName):
    auxPath = outputPath + datasetName
    filePrediction = open(auxPath + "/predictionResults.csv", "a")
    filePrediction.write("image_id, " + datasetName)
    labelEncoderPath = auxPath + "/models/le.cpickle"
    le = cPickle.loads(open(labelEncoderPath,"rb").read())
    predictions = []
    modes = []

    for fe in featureExtractors:
        (labels, images) = dataset.build_batch(imagePaths, fe[0])
        oe = Extractor(fe)
        features = oe.describe(images)
        for classi in classifiers:
            cPickleFile = auxPath + "/classifiers/classifier_" + fe[0] + "_" + classi + ".cpickle"
            if os.path.isfile(cPickleFile):
                model = cPickle.loads(open(cPickleFile,"rb").read())
                prediction = model.predict(features)
                predictions.append(prediction)
    # aux = np.array(predictions)
    aux = []
    for i in range((len(imagePaths))):
        for x in predictions:
            aux.append(x[i])
        aux = np.array(aux)
        mode = stats.mode(aux[0])

        # mode = le.inverse_transform(mode[0])
        modes.append(mode[0])
        aux = []

    # mode = stats.mode(aux[0])
    # prediction = le.inverse_transform(mode[0])
    # filePrediction.write(str(labels) + ", " + str(prediction) + "\r\n")
    # le = cPickle.loads(open(labelEncoderPath).read())
    prediction = le.inverse_transform(prediction)
    for (image, mode) in zip(imagePaths,modes):
        print("[INFO] class predicted for the image {}: {}".format(image, mode[0]))
        filePrediction.write("\n" + str(image) + ", " + str(mode[0]))

    filePrediction.close()



def __main__():
    # construct the argument parser and parse the command line arguments
    ap = argparse.ArgumentParser()
    ap.add_argument("-i", "--image", required=True, help="path to the image to predict")
    ap.add_argument("-f", "--features", required=True, help="feature extractor to predict the class")
    ap.add_argument("-m", "--classifierModels", required=True, help="model to classify the images")
    ap.add_argument("-o", "--outputPath", required=True, help="Path of the output of the algorithm")
    ap.add_argument("-d", "--datasetPath", required=True, help="model to classify the images")

    args = vars(ap.parse_args())
    conf = Conf(args["conf"])
    outputPath = conf["output_path"]
    featureExtractors = args["feature"]
    classifiers = args["classifierModels"]
    datasetPath = conf["dataset_path"]
    # This method must have the same output path and dataset path that in the other parts of the program, in other case the program
    # will not work
    # prediction(featureExtractors, classifiers, args["image"], outputPath, datasetPath)
    predictionEnsemble(featureExtractors, classifiers, args["image"], outputPath, datasetPath)


if __name__ == "__main__":
    __main__()
