import imghdr
import os
import glob
from flask import Flask, render_template, request, redirect, url_for, abort, send_from_directory, jsonify
from werkzeug.utils import secure_filename
from utils.utils import calc_mean_score, save_json
from handlers.model_builder import Nima
from handlers.data_generator import TestDataGenerator

app = Flask(__name__)

def predict(model, data_generator):
    return model.predict_generator(data_generator, workers=8, use_multiprocessing=True, verbose=1)

def image_file_to_json(img_path):
    img_dir = os.path.dirname(img_path)
    img_id = os.path.basename(img_path).split('.')[0]
    return img_dir, [{'image_id': img_id}]

def image_dir_to_json(img_dir, img_type='jpg'):
    img_paths = glob.glob(os.path.join(img_dir, '*.'+img_type))
    samples = []
    for img_path in img_paths:
        img_id = os.path.basename(img_path).split('.')[0]
        samples.append({'image_id': img_id})

    return samples

nima1 = Nima("MobileNet", weights=None)
nima1.build()
nima1.nima_model.load_weights("MobileNet/weights_mobilenet_aesthetic_0.07.hdf5")

nima2 = Nima("MobileNet", weights=None)
nima2.build()
nima2.nima_model.load_weights("MobileNet/weights_mobilenet_technical_0.11.hdf5")

image_dir = 'test_images'

app.config['MAX_CONTENT_LENGTH'] = 1024 * 1024
app.config['UPLOAD_EXTENSIONS'] = ['.jpg']
app.config['UPLOAD_PATH'] = 'uploads'

def validate_image(stream):
    header = stream.read(512)  # 512 bytes should be enough for a header check
    stream.seek(0)  # reset stream pointer
    format = imghdr.what(None, header)
    if not format:
        return None
    return '.' + (format if format != 'jpeg' else 'jpg')

@app.route('/c/', methods=['GET'])
def index():
    files = os.listdir(app.config['UPLOAD_PATH'])
    return render_template('index.html', files=files)

@app.route('/c/test/<filename>', methods=['GET'])
def test(filename):
    filename = str(filename)
    if len(filename) == 0:
        abort(400)
    image_dir_full = '/'+image_dir+'/'+filename
    samples1 = image_dir_to_json(image_dir_full, img_type='jpg')
    samples2 = samples1
    # return jsonify(samples)
    data_generator1 = TestDataGenerator(samples1, image_dir_full, 64, 10, nima1.preprocessing_function(),img_format='jpg')
    data_generator2 = TestDataGenerator(samples2, image_dir_full, 64, 10, nima2.preprocessing_function(),img_format='jpg')
    predictions1 = predict(nima1.nima_model, data_generator1)
    predictions2 = predict(nima2.nima_model, data_generator2)

    for i, sample in enumerate(samples1):
        sample['msp'] = calc_mean_score(predictions1[i])
    for i, sample in enumerate(samples2):
        sample['msp'] = calc_mean_score(predictions2[i])
    del data_generator1
    del data_generator2
    del predictions1
    del predictions2
    return jsonify([samples1, samples2])

if __name__ == '__main__':
   app.run(host='0.0.0.0', port=5060, threaded=False, processes=1)