import tensorflow as tf
import tensorflow_hub as hub
import pandas as pd
import re
import numpy as np

import gradio as gr
import sqlite3
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from nltk.stem import PorterStemmer

from queue import Queue
from nltk.stem import PorterStemmer
# Load similarity model
from sentence_transformers import SentenceTransformer
retrieval_model_path = "NlpSql//model_jatin//"
retrieval_model = SentenceTransformer(retrieval_model_path)

# Load LLM model
model_path = "weights/sql-weights"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# Load data
data = pd.read_csv('NlpSql/Data/data_1.txt', sep="\t", header=None)

headers = data.iloc[0]
new_df = pd.DataFrame(data.values[1:], columns=headers)




# Preprocess data
new_df['Question'] = new_df['Question'].str.replace('-',' ')
new_df['Question'] = new_df['Question'].str.lower()

data_list = new_df["Question"].to_list()
fin = [re.sub(r'\r\n', '', line) for line in data_list]
data_emb = retrieval_model.encode(fin)




# Define dummy users for authentication
dummy_users = [
    {"username": "jatin", "password": "qwer@345"},
    {"username": "ajoy", "password": "qwer@345"},
    {"username": "user3", "password": "pass3"},
    {"username": "user4", "password": "pass4"},
    {"username": "user5", "password": "pass5"}
]

# Authenticate user function
def authenticate_user(username, password):
    for user in dummy_users:
        if user["username"] == username and user["password"] == password:
            return True
    return False

# Function to get most similar question
def get_most_similar_question(userText, history):
    print(userText)
   
    inp_emb = retrieval_model.encode([userText])
    corr = np.inner(inp_emb, data_emb)
    
    flat_list = [item for sublist in corr for item in sublist]
    top_1_idx = np.argsort(flat_list)[-1:]
    top_1_values = [flat_list[i] for i in top_1_idx]
    
    if top_1_values[0] < 0.7:
        return predict(userText, history)
    else:
        n_data = new_df.iloc[top_1_idx, [1]]
        df_html_with_sql = n_data.to_html(index=False)
        return df_html_with_sql

    
porter_stemmer =PorterStemmer()
#--------------------------------------
def fn_preprocess_question(question):
    return ' '.join([porter_stemmer.stem(word) for word in question.split()])
dict_predefined_answers ={
    "What is the highest selling region?": "SELECT Region, Max(counterselloutval) from sales_dlb Group by region;",
    "Who are you":"I am SAMS bot"
}
#-----    
def predict(question,history):
    try:
#         if username not in user_sessions:
#             return 'User Authentication Failed.'
        preprocessed_question_var = fn_preprocess_question(question)
        for predefined_question_var,predefined_answer_var in dict_predefined_answers.items():
            if fn_preprocess_question(predefined_question_var) in preprocessed_question_var:
                return predefined_answer_var




        #--------------------------------------
#         predefined_answer = predefined_answers.get(question)
#         if predefined_answer:
#             return predefined_answer
        #--------------------------------------
        #--------------------------------------
        # Check if the user's question closely matches a predefined question
#         matched_question = next((predefined_question for predefined_question in predefined_answers.keys()
#                                 if match_question(question, predefined_question)), None)

#         if matched_question:
#             return predefined_answers[matched_question]
        #--------------------------------------
        conn_local = sqlite3.connect("sales_database.db",check_same_thread=False)
        cursor_local = conn_local.cursor()

        ### Answer Given the database schema,here is the SQL query that answers [QUESTION]{question}[/QUESTION] [SQL] """.format(question=question)
        prompt = """### Instructions:
        Your task is convert a question into a SQL query, given a sqlite3 database schema.
        Adhere to these rules:
        - **Deliberately go through the question and database schema word by word** to appropriately answer the question
        - **Use Table Aliases** to prevent ambiguity. For example, `SELECT table1.col1, table2.col1 FROM table1 JOIN table2 ON table1.id = table2.id`.
        - When creating a ratio, always cast the numerator as int
        
        ### Input:
        Generate a SQL query that answers the question `{question}`.
        This query will run on a database whose schema is represented in this string:        
         CREATE TABLE IF NOT EXISTS sales (
                    salestype TEXT,
                    salescategory TEXT,
                    channel TEXT,
                    region TEXT,
                    month INTEGER,
                    sellercode TEXT,
                    sellername TEXT,
                    modelname TEXT,
                    quantity INTEGER,
                    amount INTEGER );

        CREATE TABLE IF NOT EXISTS sales_dlb (
                     Year INTEGER, -- year of sellout/sale
                     SuperRegion TEXT, --Superegion of sale
                     CounterSelloutVal INTEGER,-- value of sellout/amount
                     CounterSelloutVol INTEGER,-- volume of sellout
                     Region TEXT,-- region of sale
                     StoreCode TEXT,-- store code
                     Competitor TEXT,--competitor/competition name
                     SurveyDate TEXT,-- date of survey
                     PriceBand TEXT,-- price bands of different model/modelcode sold
                     SalesType TEXT,-- salestype of the sale
                     ModelCode TEXT,-- unique model code that are sold
                     Channel TEXT,-- channel through which it is sold
                     Status TEXT, -- status of sale
                     SAMSModelCode TEXT,-- modelcode in sams database
                     MarketName TEXT,-- market name
                     ProductGroup TEXT -- product group
        );
        ### Response:
        Based on your instructions, here is the SQL query I have generated to answer the question `{question}`:
        ```sql
        """.format(question=question)

        eos_token_id = tokenizer.convert_tokens_to_ids(["```"])[0]
#         CREATE TABLE IF NOT EXISTS sales_dlb (
#             Year INTEGER,
#             SuperRegion TEXT,
#             CounterSelloutVal INTEGER,
#             CounterSelloutVol INTEGER,
#             Region TEXT,
#             StoreCode TEXT,
#             Competitor TEXT,
#             SurveyDate TEXT,
#             PriceBand TEXT,
#             SalesType TEXT,
#             ModelCode TEXT,
#             Channel TEXT,
#             Status TEXT,
#             SAMSModelCode TEXT,
#             MarketName TEXT,
#             ProductGroup TEXT );
        

        inputs = tokenizer(prompt, return_tensors="pt").to(device)
        generated_ids = model.generate(
            **inputs,
            num_return_sequences=1,
            eos_token_id=eos_token_id,
            pad_token_id=eos_token_id,
            max_new_tokens=400,
            do_sample=False,
            num_beams=5,
            temperature = 0.9
        )


        outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

        generated_sql = outputs[0].split("```sql")[-1].split("```")[0].split(";")[0].strip()
        print(generated_sql)
    #     result = execute_sql_query(generated_sql)
    #     for row in result:
    #         print(row)
    #     df = pd.DataFrame(result,columns=[desc[0] for desc in cursor.description])
    # #     df_str = df.to_string(index=False)
    #     res_html = df.to_html(index=False)
#         generated_sql ='''100 * SUM(CASE WHEN sales_dlb.competitor = 'SAMSUNG' THEN sales_dlb.counterselloutval ELSE 0 END) / NULLIF(SUM(sales_dlb.counterselloutval), 0) AS percentage_of_counterselloutval
#         FROM   sales_dlb'''
        cursor_local.execute(generated_sql)
        result = cursor_local.fetchall()
        print(result)
        cursor_local.close()
        conn_local.close()
        for row in result:
            print(row)
        df = pd.DataFrame(result,columns=[desc[0] for desc in cursor_local.description])
        sql_datagram = generated_sql

        # Create a DataFrame with the SQL datagram as the first row
        df_sql_row = pd.DataFrame([{'Generated_SQL': sql_datagram}])

        # Concatenate the SQL row with your existing DataFrame (df)
        df_with_sql = pd.concat([df,df_sql_row], keys=['',''],ignore_index=True)

        # Convert the DataFrame with SQL to HTML
        df_html_with_sql = df.to_html(index=False)

#         df_final = df + dff
    #     df_str = df.to_string(index=False)
        
        return df_html_with_sql
    except Exception as e:
        print(f"An exception occurred:{str(e)}")
        return "Kindly modify the query that can best fit the database schema.Please provide updated query."

    
gr.ChatInterface(get_most_similar_question).launch(auth=authenticate_user,inline=True, share=True,server_name= '0.0.0.0', server_port=8181)


I need to ask user for his feedback after every question.Can you create a Like dislike button in gradion or make this app completely using streamlit where chatting is possible with user feedback for like /dislike of answer