from styx_msgs.msg import TrafficLight
import tensorflow as tf
import numpy as np
import cv2
import rospy
import os

LIGHTS = ['Green', 'Red', 'Yellow', 'Unknown']

class TLClassifier(object):
    def __init__(self, realData):
        #TODO load classifier
        self.count = 0
        if realData:
            self.site = True
            PATH = 'model/frozen_model_real/'
        else:
            self.site = False
            PATH = 'model/frozen_model_sim/'
        FROZEN_GRAPH = PATH + 'frozen_inference_graph.pb'

        self.graph = tf.Graph()
        with self.graph.as_default():
            od_graph_def = tf.GraphDef()
            with tf.gfile.GFile(FROZEN_GRAPH, 'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                tf.import_graph_def(od_graph_def, name='')

        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True
        self.sess = tf.Session(graph=self.graph, config=config)
        self.image_tensor = self.graph.get_tensor_by_name('image_tensor:0')
        self.detection_boxes = self.graph.get_tensor_by_name('detection_boxes:0')
        self.detection_scores = self.graph.get_tensor_by_name('detection_scores:0')
        self.detection_classes = self.graph.get_tensor_by_name('detection_classes:0')
        self.num_detections = self.graph.get_tensor_by_name('num_detections:0')

    def to_image_coords(self, boxes, height, width):
        """
        The original box coordinate output is normalized, i.e [0, 1].

        This converts it back to the original coordinate based on the image
        size.
        """
        box_coords = np.zeros_like(boxes)
        box_coords[:, 0] = boxes[:, 0] * height
        box_coords[:, 1] = boxes[:, 1] * width
        box_coords[:, 2] = boxes[:, 2] * height
        box_coords[:, 3] = boxes[:, 3] * width

        return box_coords

    def draw_boxes(self, image, boxes, classes, scores):
        """Draw bounding boxes on the image"""
        for i in range(len(boxes)):
            top, left, bot, right = boxes[i, ...]
            cv2.rectangle(image, (left, top), (right, bot), (255,0,0), 3)
            text = LIGHTS[int(classes[i])-1] + ': ' + str(int(scores[i]*100)) + '%'
            cv2.putText(image , text, (left, int(top - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (200,0,0), 1, cv2.LINE_AA)

    def filter_boxes(self, min_score, boxes, scores, classes):
        """Return boxes with a confidence >= `min_score`"""
        n = len(classes)
        idxs = []
        for i in range(n):
            if scores[i] >= min_score:
                idxs.append(i)

        filtered_boxes = boxes[idxs, ...]
        filtered_scores = scores[idxs, ...]
        filtered_classes = classes[idxs, ...]
        return filtered_boxes, filtered_scores, filtered_classes

    def get_classification(self, image):
        """Determines the color of the traffic light in the image

        Args:
            image (cv::Mat): image containing the traffic light

        Returns:
            int: ID of traffic light color (specified in styx_msgs/TrafficLight)

        """
        #TODO implement light color prediction

        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        (im_width, im_height, _) = image_rgb.shape
        image_np = np.expand_dims(image_rgb, axis=0)

        with self.graph.as_default():
            (boxes, scores, classes) = self.sess.run([self.detection_boxes, self.detection_scores, self.detection_classes],
                                                feed_dict={self.image_tensor: image_np})
            # Remove unnecessary dimensions
            boxes = np.squeeze(boxes)
            scores = np.squeeze(scores)
            classes = np.squeeze(classes)

            confidence_threshold = 0.5
            # Filter boxes with a confidence score less than `confidence_threshold`
            boxes, scores, classes = self.filter_boxes(confidence_threshold, boxes, scores, classes)

        # Write image to disk
        write = False
        if write:
            image = np.dstack((image[:, :, 2], image[:, :, 1], image[:, :, 0]))
            width, height = image.shape[1], image.shape[0]
            box_coords = self.to_image_coords(boxes, height, width)
            self.draw_boxes(image, box_coords, classes, scores)
            output_path = "/home/z637177/Desktop/debug/"
            if not os.path.exists(output_path):
                os.makedirs(output_path)
            cv2.imwrite(output_path + 'img_' + str(self.count) + '.jpg', image)
            self.count += 1

        if len(scores)>0:
            this_class = int(classes[np.argmax(scores)])
        else:
            this_class = 4

        if this_class == 1:
            return TrafficLight.GREEN
        elif this_class == 2:
             return TrafficLight.RED
        elif this_class == 3:
             return TrafficLight.YELLOW

        return TrafficLight.UNKNOWN


''' Predict vegitables code here '''
#imports
import os
import io
import numpy as np

import keras
from keras.preprocessing import image
from keras.preprocessing.image import img_to_array
from keras.applications.xception import (
    Xception, preprocess_input, decode_predictions)
from keras import backend as K

from flask import Flask, render_template, flash, redirect, url_for, session, logging, request, jsonify
from flask_sqlalchemy import SQLAlchemy
import sqlalchemy
from sqlalchemy.ext.automap import automap_base
from sqlalchemy.orm import Session
from sqlalchemy import create_engine, func
import pickle
import numpy as np

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'food_test/test_images'

logit_model = pickle.load(open('./model/logit_model.pkl', 'rb'))
reg_model = pickle.load(open('./model/reg_model.pkl', 'rb'))

model = None
graph = None

def load_model():
    global model
    global graph
    model = Xception(weights="imagenet")
    graph = K.get_session().graph

load_model()

def prepare_image(img):
    img = img_to_array(img)
    img = np.expand_dims(img, axis=0)
    img = preprocess_input(img)
    # return the processed image
    return img

@app.route("/")
def index():
    return render_template("index.html")

@app.route('/upload_img', methods=['GET', 'POST'])
def upload_file():
    data = {"success": False}
    if request.method == 'POST':
        # if request.files.get('file'):
        if request.files.get('file'):
            # read the file
            file = request.files['file']

            # read the filename
            filename = file.filename

            # create a path to the uploads folder
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)

            file.save(filepath)

            # Load the saved image using Keras and resize it to the Xception
            # format of 299x299 pixels
            image_size = (299, 299)
            im = keras.preprocessing.image.load_img(filepath,
                                                    target_size=image_size,
                                                    grayscale=False)

            # preprocess the image and prepare it for classification
            image = prepare_image(im)

            global graph
            with graph.as_default():
                preds = model.predict(image)
                print(preds)
                results = decode_predictions(preds)
                print(results)

                prediction = []

                # loop over the results and add them to the list of
                # returned predictions
                for (imagenetID, label, prob) in results[0]:
                    prediction.append(label)

                page = []
                if prediction[0] == "broccoli":
                    page.append("broccoli_low.html")
                elif prediction[0] == "cucumber":
                    page.append("cucumber_low.html")
                else:
                    page.append("butternut_squash_low.html")
        return render_template(page)
    return render_template("upload_img.html")

@app.route('/broccoli_pancakes')
def broccoli_pancakes():
    return render_template("broccoli_pancakes.html", ingredients=ingredients , instruction=instruction, time=time , cost=broccoli_pancakes_cost)

@app.route('/register', methods=['GET', 'POST'])
def register():
   highsch = 1
   college = 0
   income = 20
   strenuous = 1
   read = 0
   highread = 0
   collegeread = 0

   if request.method == "POST":
       age = float(request.form["age"])
       age2 = (age)**2
       if request.form["gender"] == "male":
           male = 1
       else:
           male = 0
       if request.form["origin"] == "white":
           white = 1
           black = 0
           asian = 0
       elif request.form["origin"] == "black":
           white = 0
           black = 1
           asian = 0
       elif request.form["origin"] == "others":
           white = 0
           black = 1
           asian = 0
       else:
           white = 0
           black = 0
           asian = 1

       # data = request.form["age"]

       data = [age, age2, male, highsch, college, income, white, black, asian, strenuous, read, highread, collegeread]
       data = np.array([data])

       prediction1 = logit_model.predict(data)
       prediction2 = reg_model.predict(data)
       prediction2 = np.array(prediction2)

       if prediction1 == 0:
           result = "You are not at risk!"
       else:
           result = "You are at risk!"

       return render_template("prediction.html", data=data, result=result, title="Health Evaluation", BMI=prediction2)
   return render_template("register.html")

@app.route('/team')
def team():
    return render_template("team.html")

@app.route("/recipes", methods=[ "GET","POST"])
def recipes():
   recps = []
   ingres = []

   engine = create_engine("sqlite:///db/data.sqlite")
   # Reflect an existing database into a new model
   Base = automap_base()
   # Reflect the tables
   Base.prepare(engine, reflect=True)
   # Save references to each table
   Recipe = Base.classes.recipes
   Ingredient = Base.classes.ingredient
   # Create session for query
   session = Session(engine)

   if request.method == "POST":
       result = request.form["recipe"]

       recps = session.query(Recipe).filter(Recipe.name == result).all()
       ingres = session.query(Ingredient).filter(Ingredient.recipes == result).all()
       total = session.query(func.sum(Ingredient.cost)).filter(Ingredient.recipes == result).all()
       print(ingres)

       return render_template("recipe.html", recps=recps, ingres=ingres, total=total, title="Recipe")

   return render_template("broccoli_low.html")

@app.route('/login')
def login():
    return render_template("login.html")

@app.route('/upload_img')
def upload_img():
     return render_template("upload_img.html")

if __name__ == "__main__":
    app.run(debug=True)