# Embedding and similarity using GTE-Large model
#
# https://huggingface.co/thenlper/gte-large
# This model exclusively caters to English texts, and any lengthy texts will be truncated to a maximum of 512 tokens.
#
# pip install -U sentence-transformers
#
import torch
import numpy as np

from sentence_transformers import SentenceTransformer
from transformers import set_seed

import os
import logging
import time

from flask import Flask, abort, jsonify, request

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO,
)
logger = logging.getLogger(__name__)

seed_value = 1
set_seed(seed_value)
torch.manual_seed(seed_value)
if torch.cuda.is_available():
    torch.cuda.manual_seed(seed_value)
    torch.cuda.manual_seed_all(seed_value)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    logger.info("Cuda available!")

def round_byte(a):
    b = round(a * 740) # for 1024 size
    if b > 127: b = 127
    elif b < -128: b = -128
    return b

def round_floats(o):
    if isinstance(o, float): return round(o, 6)
    if isinstance(o, dict): return {k: round_floats(v) for k, v in o.items()}
    if isinstance(o, (list, tuple)): return [round_floats(x) for x in o]
    return o

def round_floats0(o):
    if isinstance(o, float): return round_byte(o)
    if isinstance(o, dict): return {k: round_floats(v) for k, v in o.items()}
    if isinstance(o, (list, tuple)): return [round_floats(x) for x in o]
    return o

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

def get_detailed_instruct(task_description: str, query: str) -> str:
    return f'Instruct: {task_description}\nQuery: {query}'

def embed_query(query):
#    emb = model.encode([query], prompt_name="query", to_numpy=True, normalize_embeddings=True, show_progress_bar=False)
    emb = model.encode([query], normalize_embeddings=True)
    remb = round_floats(emb.tolist()[0])
#    torch.cuda.empty_cache()
    return remb

def embed(text):
#    emb = model.encode([text], to_numpy=True, normalize_embeddings=True, show_progress_bar=False)
    emb = model.encode([text], normalize_embeddings=True)
    remb = round_floats(emb.tolist()[0])
#    torch.cuda.empty_cache()
    return remb

def embedding(texts):
#    message_embeddings = model.encode(texts, to_numpy=True, show_progress_bar=False)
    message_embeddings = model.encode(texts, normalize_embeddings=True)
#    torch.cuda.empty_cache()
    return message_embeddings.tolist()

def similarity(text, texts):    
    embeddings = embedding(texts)
    query_embedding = embed_query(text)
    scs = []
    i = 0
    while i < len(embeddings):
        sc = cosine_similarity(query_embedding, embeddings[i])
        scs.append(float(sc))
#        print("{} \t: {}".format(sc, texts[i][0:100]))
        i += 1
#    torch.cuda.empty_cache()
    return scs

def empty_cache():
    global cache_cnt
    cache_cnt += 1
    if cache_cnt % 100 == 0:
        torch.cuda.empty_cache()

cache_cnt = 0;

class Args:
    model_url_or_path = 'thenlper/gte-large'
#    model_url_or_path = '/app/GTE-Large'
args = Args()
logger.info(args)

model_name = args.model_url_or_path

model = SentenceTransformer(
    model_name,
#    model_kwargs={"attn_implementation": "flash_attention_2", "device_map": "auto"},
#    tokenizer_kwargs={"padding_side": "left"},
)
logger.info("Model loaded: " + model_name)


#====================================
#http://localhost:8085/embed?text=test

app = Flask(__name__)
app.config['CORS_HEADERS'] = 'Content-Type'
app.config['DEBUG'] = False

@app.route("/embed", methods=['GET'])
def get_embed():
    isQuery = request.args.get('isQuery')
    text = request.args.get('text')
    if len(text) == 0:
        abort(500)
    else:
        empty_cache()
        start = time.time()
#        print(text)
        if isQuery == 'true':
            result = embed_query(text)
        else:   
            result = embed(text)
        end = time.time()
        took = int((end-start)*1000)
        logger.info('Embed took %d (ms)' % (took))
        return jsonify({'result': round_floats(result),'took': took})

@app.route("/embed", methods=['POST'])
def post_embed():
    data = request.get_json()
    if 'isQuery' in data:
        isQuery = data['isQuery']
    else:
        isQuery = False
    if 'text' not in data or len(data['text']) == 0:
        abort(500)
    else:
        empty_cache()
        start = time.time() 
        text = data['text']
#        print(text)
        if isQuery:
            result = embed_query(text)
        else:
            result = embed(text)
        end = time.time() 
        took = int((end-start)*1000)
        logger.info('Embed took %d (ms)' % (took))
        return jsonify({'result': round_floats(result),'took': took})

@app.route("/embedding", methods=['POST'])
def post_embedding():
    data = request.get_json()
    if 'texts' not in data or len(data['texts']) == 0:
        abort(500)
    else:
        empty_cache()
        start = time.time() 
        texts = data['texts']
#        print(texts)
        results = embedding(texts)
        end = time.time() 
        took = int((end-start)*1000)
        logger.info('Embedding %d texts took %d (ms)' % (len(texts), took))
        return jsonify({'result': round_floats(results),'took': took})

@app.route("/similarity", methods=['POST'])
def post_similarity():
    data = request.get_json()
    if 'text' not in data or len(data['text']) == 0 or 'texts' not in data or len(data['texts']) == 0:
        abort(500)
    else:
        empty_cache()
        start = time.time()
        text = data['text']
        texts = data['texts']
#        print(text)
        result = similarity(text, texts)
        end = time.time() 
        took = int((end-start)*1000)
        logger.info('Similarity for %d texts took %d (ms)' % (len(texts), took))
        return jsonify({'result': result,'took': took})

if __name__ == "__main__":
    app.run(host='0.0.0.0')
#    app.run(host="0.0.0.0", port=80, debug=False)

