from langchain_community.embeddings.sentence_transformer import (
    SentenceTransformerEmbeddings,
)
from langchain.retrievers import ParentDocumentRetriever
from langchain_community.vectorstores import Chroma
from langchain_text_splitters import CharacterTextSplitter, RecursiveCharacterTextSplitter
import pickle
from langchain_community.llms import HuggingFaceEndpoint
from dotenv import load_dotenv, find_dotenv
import os
from langchain.chains import ConversationalRetrievalChain
from langchain.chains import LLMChain
from langchain.chains.conversational_retrieval.prompts import CONDENSE_QUESTION_PROMPT
from langchain.chains.question_answering import load_qa_chain
from langchain import PromptTemplate
from transformers import pipeline



# Please ensure you have a .env file available with 'HUGGINGFACEHUB_API_TOKEN'
load_dotenv(find_dotenv())
HUGGINGFACEHUB_API_TOKEN  = os.environ["HUGGINGFACEHUB_API_TOKEN"]

repo_id ="mistralai/Mistral-7B-Instruct-v0.2"
llm = None

def choose_model2(model):
    global repo_id
    global llm

    if model == "mistral-7b-base-model":
        repo_id="mistralai/Mistral-7B-Instruct-v0.2"
        print("model chooosed from recomm",repo_id)
    else:
        repo_id="GRMenon/mental-health-mistral-7b-instructv0.2-finetuned-V2"
        print("model chooosed from recomm",repo_id)

    llm = HuggingFaceEndpoint(
        repo_id=repo_id, max_length=512, temperature=0.5, token=HUGGINGFACEHUB_API_TOKEN
    )

def start_recommend():
    persist_directory="Data/chroma"
    #chroma_client = chromadb.PersistentClient(persist_directory=persist_directory)
    embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
    vectors = Chroma(persist_directory = persist_directory, embedding_function = embedding_function, collection_name="split_parents")
    retriever = vectors.as_retriever() #(k=6)

    #prompt="you are a mental health therapist, talking to a person with who is facing some mental health issues. Following is the user feeling {question}"

    prompt = """You're a Mental Health Specialist. Support those with Depressive Disorder. Your task is to provide brief mental health advice. PLEASE KEEP IT BRIEF.
    Listen compassionately, respond helpfully. For casual talk, be friendly. For facts, use context.
    Following is the user feeling {question}.
    If unsure, say, 'Out of my knowledge.' Always stay direct.
    If you cannot find the answer from the pieces of context, just say that you don't know, don't try to make up an answer.
    PLEASE GIVE THE BRIEF RESPONSE IN THE FORM OF BULLET POINTS.
    ----------------
    {context}"""

    global chain
    prompt = PromptTemplate(input_variables=['question'],template=prompt)

    chain1 = LLMChain(llm=llm, prompt=prompt, verbose=True)
    doc_chain = load_qa_chain(llm, chain_type="stuff")

    chain = ConversationalRetrievalChain(
        retriever=retriever,
        question_generator=chain1,
        combine_docs_chain=doc_chain,
        verbose=True,
    )


def recommend2(query):

    chat_history = []
    # query = "I feel sad"
    result = chain.invoke({"question": query, "chat_history": chat_history})
    print("---------------\nSummary from the chat agent:",query)
    # print("this is the result from vector database:",vectors.similarity_search(query))
    # print("this is the response from recommendation agent:", result["answer"])

    return result["answer"]

def is_depressed(human_inputs):
    ''''
    returns wether according to human inputs the person is depressed or not
    '''
    # Implement Classification
    # all_user_inputs = ''.join(human_inputs)
    pipe = pipeline('sentiment-analysis')
    status = pipe(human_inputs)
    return 'Is depressed' if status[0]["label"] == "NEGATIVE" else 'Not Depressed'
    # return status[0]["label"]


def is_depressed_from_nlp_model(human_inputs):
    ''''
    returns wether according to human inputs the person is depressed or not
    '''
    # Implement Classification
    # all_user_inputs = ''.join(human_inputs)
    from nlp_models import sentiment_class
    predictions = sentiment_class(human_inputs)
    # Assuming the predictions are shaped (1, 2) as in the example
    probability_depressed, probability_not_depressed = predictions[0]
    
    if probability_depressed > probability_not_depressed:
        return "Depressed", probability_depressed
    else:
        return "Not depressed", probability_not_depressed
    #return 'Not so depressed' if status[0][1] > 0.5 else 'is_depressed'
    #return 'Is depressed' if status[0]["label"] == "NEGATIVE" else 'Not Depressed'
    # return status[0]["label"]