from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.document_loaders import PyPDFDirectoryLoader
from langchain_community.llms import HuggingFaceEndpoint
from langchain.chains import ConversationalRetrievalChain
from langchain.chains import RetrievalQA
import gradio as gr
import os
from pandasai import Agent
from langchain.document_loaders.csv_loader import CSVLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.memory import ConversationSummaryBufferMemory
import io
import contextlib
import re
import pandas as pd
from transformers import AutoConfig

config = AutoConfig.from_pretrained("config.json")
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
config = AutoConfig.from_pretrained("config.json")
vector_store= FAISS.load_local("vector_db/", embeddings, allow_dangerous_deserialization=True)

repo_id="mistralai/Mixtral-8x7B-Instruct-v0.1"

llm = HuggingFaceEndpoint(
    repo_id = repo_id, 
    temperature = 0.01, 
    max_new_tokens = 4096,
    verbose = True,
    return_full_text = False 
)

retriever = vector_store.as_retriever(
    search_type="similarity",
    search_kwargs={"k":5}
)

df=pd.read_csv('data/Gretel_Data.csv')
averages = df.mean(numeric_only=True).to_dict()

agent = Agent([df], config={"llm": llm, 'verbose':True})

global unique_columns
unique_columns = [
        'Avg_Connected_UEs',
        'PRB Util%',
        'CA Activation Rate',
        'DLRLCLayerDataVolume MB',
        'DRB UL Data Volume MB',
        'UPTP_Mbps',
        'UPTP Mbps Num',
        'UPTP Mbps Den',
        'UL MAC Vol Scell Pct',
        'DL MAC Vol Scell Pct',
        'DL MAC Vol Scell MB',
        'DL Volume',
        'DL Data Vol MAC in MB',
        'UL Throughput',
        'MB_per_connected_UE'
    ]

global target_words
target_words = ["Bandwidth", "Interference", "Call Quality", "Network", "Handover"]

columns = []

column_avgs = {}

global network_features
network_features  = {
    'Bandwidth': [
        'Avg_Connected_UEs',
        'PRB Util%',
        'CA Activation Rate',
        'DLRLCLayerDataVolume MB',
        'DRB UL Data Volume MB',
        'UPTP_Mbps',
        'UPTP Mbps Num',
        'UPTP Mbps Den',
        'UL MAC Vol Scell Pct',
        'DL MAC Vol Scell Pct',
        'DL MAC Vol Scell MB',
        'DL Volume',
        'DL Data Vol MAC in MB',
        'UL Throughput',
        'MB_per_connected_UE'
    ],
    'Handover': [
        'Avg_Connected_UEs',
        'PRB Util%',
        'CA Activation Rate',
        'HO Failures',
        'HO_fail_InterFreq',
        'HO_fail_PCT_InterFreq',
        'HO Failure%',
        'HO Attempts',
        'HO_att_InterFreq'
    ],
    'Network': [
        'Avg_Connected_UEs',
        'PRB Util%',
        'CA Activation Rate',
        'SIP DC%',
        'RRC Setup Attempts',
        'RRC Setup Failures',
        'RRC Setup Failure% 5G',
        'Combined RACH Failure%',
        'Combined RACH Preambles',
        'Combined RACH Failures',
        'Interference Pwr',
    ],
    'Call Quality': [
        'Avg_Connected_UEs',
        'PRB Util%',
        'CA Activation Rate',
        'Avg_PUCCH_SINR',
        'Avg CQI',
        'SIP Calls with a Leg',
        'SIP_SC_Total_MOU',
        'SIP Dropped Calls',
        'VoLTE_MOU',
        'QCI 1 Bearer Drops',
        'QCI 1 Bearer Releases',
        'QCI 1 Bearer Drop%',
        'Peak UE',
        'DL Packet Loss Pct',
        'UL Resid BLER PCT',
        'Bearer Drops Voice',
        'Bearer Releases Voice',
        'Bearer Drop%',
        'Call_Drops_Credit'
    ],
    'Interference': [
        'Avg_Connected_UEs',
        'PRB Util%',
        'CA Activation Rate',
        'Combined RACH Failure%',
        'Interference Pwr'
    ]
}

def echo(message, history):
  try:
    qa=RetrievalQA.from_chain_type(llm=llm, retriever=retriever, return_source_documents=True)
    message= " <s> [INST] You are a senior telecom network engineer having access to troubleshooting tickets data and other technical and product documentation. Stick to the knowledge provided. Search through the product documentation pdfs first before scanning the tickets to generate the answer. Return only the helpful answer. Question:" + message + '[/INST]'
    result= qa({"query":message})
    answer= result['result']
    for word in target_words:
      if re.search(r'\b' + re.escape(word) + r'\b', answer, flags=re.IGNORECASE):
        columns.extend(network_features.get(word, []))
    unique_columns = list(set(columns))

    for column in unique_columns:
      column_avgs.update({column:averages.get(column, [])})

    result_df = df[unique_columns].iloc[:25]

    def highlight_rows(val, threshold):
        if val > threshold:
            return 'color: red; font-weight: bold'
        elif val < threshold:
            return 'color: green'
        else:
            return ''

    styled_df = result_df.style

    for key in column_avgs:
      styled_df = styled_df.applymap(lambda x, k=key: highlight_rows(x, column_avgs[k]), subset=[f'{key}'])

    gr.Dataframe(styled_df)
    return (
       "Answer: \n"
       + '\n' + answer.strip() + '\n'
       + '\n' + "Sources: \n"
       + '\n' + '1. ' + result['source_documents'][0].metadata['source'] + '\n' + result['source_documents'][0].page_content +  "\n"
       + '\n' + '2. ' + result['source_documents'][1].metadata['source'] + '\n' + result['source_documents'][1].page_content +  "\n"
       + '\n' + "3. " + result['source_documents'][2].metadata['source'] + '\n' + result['source_documents'][2].page_content +  "\n"
       + '\n' + "4. " + result['source_documents'][3].metadata['source'] + '\n' + result['source_documents'][3].page_content +  "\n"
       + '\n' + "5. " + result['source_documents'][4].metadata['source'] + '\n' + result['source_documents'][4].page_content +  "\n",
       styled_df
       )
  except Exception as e:
    error_message = f"An error occurred: {e}"+str(e.with_traceback) + str(e.args)
    return error_message, error_message

def echo_agent(message, history):
  try:
    response = agent.chat(message, output_type= 'text')
    # explanation = agent.explain()

    # result = "Answer: \n" + '\n' + response.str() + '\n' + '\n' + "Explanation: \n" + '\n' + explanation
    
    return response
  except Exception as e:
    error_message = f"An error occurred: {e}"+str(e.with_traceback) + str(e.args)
    return error_message

demo_agent = gr.Blocks(
        title="Network Ticket Knowledge Management",
        theme=gr.themes.Soft(),
)

with demo_agent:

  gr.Markdown(
      '''
      # <p style="text-align: center;">Network Ticket Knowledge Management</p>
      Welcome to Verizon Network Operations Center. I am here to help the Field Operations team with technical queries & escalation.
      '''
      )

  with gr.Tab('Clara'):
    with gr.Row():
      message = gr.Text(label="Input Query")

    btn = gr.Button("Submit")

    with gr.Row():
      reply = gr.Text(label="RCA and MoP", autoscroll=False)

    with gr.Accordion(label = "Metrics", open=False):
          table = gr.Dataframe()

    btn.click(echo, inputs=[message], outputs=[reply, table])

    gr.Examples([
        "Wi-Fi connected but no internet showing",
        'What are the possible cause of router overheating ?',
        "What are the possible causes of RAN getting disconnected frequently?",
        "For the past week, are there any specific cell towers in Texas experiencing unusually high call failure rates or data latency?",
        "What are the network problems faced by people living in the state of California?",
        "I have an FWA connection and all devices except my iPhone have internet access via this FWA device. Can you suggest steps for resolution?",
        "We're receiving reports of congested cell towers in Cleveland. Can you identify the specific cell towers experiencing overload and suggest any temporary network adjustments to alleviate the congestion?"
        ],
                inputs=[message]
                )

  with gr.Tab('Sam'):
    with gr.Row():
      message_agent = gr.Text(label="Input Query")
    with gr.Row():
      reply_agent = gr.Text(label="Answer")

    btn2 = gr.Button("Submit")
    btn2.click(echo_agent, inputs=[message_agent], outputs=[reply_agent])


demo_agent.launch(share=True,debug=True,auth=("admin", "Sam&Clara"))