import faiss
import cv2
import time
import sys
import numpy as np
import pandas as pd
from flask import Flask
from flask import request
from flask import render_template


'''
create Index Object of faiss
'''
# load csv from disk
embeddings_properties_df = pd.read_csv("embeddings_properties.csv", sep='|', index_col=0)
print(f"load csv snapshot, contains {len(embeddings_properties_df)} records!")
print(embeddings_properties_df)

# create Index Object
embeddings_list = [[e for e in str(embedding).split(',')] for embedding in embeddings_properties_df["embedding"].to_list()]
embeddings = np.array(embeddings_list)
print(f"load embeddings successfully, shape is {embeddings.shape}")
dimension = embeddings.shape[1]
# you can try different Index type
search_index = faiss.IndexFlatIP(dimension)
# add all embeddings into Index, use it in search action later
search_index.add(embeddings)


'''
generate embedding data for query image
'''
# vehicle embedding model keep as same as vehicle_encoding_pipeline.cpp
vehicle_embedding_model_path = "../../../vp_data/models/trt/vehicle/dim256_batch=1_0916_norm.onnx"
input_size = (256, 256)
net = cv2.dnn.readNetFromONNX(vehicle_embedding_model_path)
output_layer_names = net.getUnconnectedOutLayersNames()
def generate_embedding(query_image):
    input_image = cv2.resize(query_image, input_size)
    blob = cv2.dnn.blobFromImage(input_image)
    net.setInput(blob)
    out = net.forward(output_layer_names)
    embedding = np.array(out[0])
    return embedding


'''
create flask app
'''
app = Flask(__name__)

# home page
@app.route("/")
def home():
    vehicle_types = ["all", "bus", "business_car", "construction_truck", "large_truck", "sedan", "small_truck", "suv", "tanker", "van"]
    vehicle_colors = ["all", "black", "blue", "grey", "red", "white", "yellow", "other"]
    return render_template("home.html", total_records=search_index.ntotal, vehicle_types=vehicle_types, vehicle_colors=vehicle_colors)

# search action
@app.route("/search", methods=["POST"])
def search():
    # no exception catch
    query_image = request.files['query_image']
    topK = int(request.form['topK'])
    vehicle_type = str(request.form['vehicle_type'])
    vehicle_color = str(request.form['vehicle_color'])
    vehicle_plate = str(request.form['vehicle_plate'])
    
    # convert to cv2.mat and save to disk
    img_array = np.asarray(bytearray(query_image.read()), dtype=np.uint8)
    query_mat = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
    query_image_path = "static/query_images/" + str(time.time()) + ".jpg"
    cv2.imwrite(query_image_path, query_mat)

    # generate embedding for query image and do search action
    query_embedding = generate_embedding(query_mat)
    D, I = search_index.search(query_embedding, topK)

    # print similiarity search info
    print("query embedding:")
    print(query_embedding)
    print("similiarity distance:")
    print(D)
    print("similiarity uid:")
    print(I)

    # need properties filter
    sql = "path != ''"
    if vehicle_color != "all" or vehicle_type != "all" or vehicle_plate != "":
        if vehicle_color != "all":
            sql += f" & color == '{vehicle_color}'"
        if vehicle_type != "all":
            sql += f" & type == '{vehicle_type}'"
        if vehicle_plate != "":
            sql += f" & plate.str.contains('{vehicle_plate}')"
    filter_df = embeddings_properties_df.query(sql)
    filter_uids = filter_df.index.to_list()

    # filter action
    paired_lists = list(zip(I[0], D[0]))
    filter_result = [elem for elem in paired_lists if elem[0] in filter_uids]
    shoot_I = [e[0] for e in filter_result]
    shoot_D = [e[1] for e in filter_result]

    # shoot data
    print(shoot_I)
    print(shoot_D)

    shoot_df = embeddings_properties_df.iloc[shoot_I]
    shoot_path = shoot_df['path']
    shoot_color = shoot_df['color']
    shoot_type = shoot_df['type']
    shoot_plate = shoot_df['plate']
    shoot_distance = shoot_D
    shoot_uid = shoot_I

    # render search results
    search_results = []
    for i,p,d,c,t,pl in zip(shoot_uid, shoot_path, shoot_distance, shoot_color, shoot_type, shoot_plate):
        search_results.append({"uid": i, "image_url": p, "distance":d, 'vehicle_color':c, 'vehicle_type':t, 'vehicle_plate':pl})
    return render_template("search_result.html", topK=topK, \
        search_results=search_results, query_image_url=query_image_path, \
            search_color=vehicle_color, search_type=vehicle_type, search_plate=vehicle_plate)


if __name__ == "__main__":
    app.run(host="192.168.77.68", port=9898)