from src.medrag import MedRAG

question = "A lesion causing compression of the facial nerve at the stylomastoid foramen will cause ipsilateral"
options = {
    "A": "paralysis of the facial muscles.",
    "B": "paralysis of the facial muscles and loss of taste.",
    "C": "paralysis of the facial muscles, loss of taste and lacrimation.",
    "D": "paralysis of the facial muscles, loss of taste, lacrimation and decreased salivation."
}

cot = MedRAG(llm_name="OpenAI/gpt-3.5-turbo-16k", rag=False)
answer, _, _ = cot.answer(question=question, options=options)
print(f"Final answer in json with rationale: {answer}")
# {
#   "step_by_step_thinking": "Compression of the facial nerve at the stylomastoid foramen will affect the function of the facial nerve. The facial nerve is responsible for innervating the muscles of facial expression, including those involved in smiling, frowning, and closing the eyes. It also carries taste sensation from the anterior two-thirds of the tongue. Additionally, the facial nerve controls tear production (lacrimation) and salivation. Therefore, compression of the facial nerve at the stylomastoid foramen will cause paralysis of the facial muscles (A), loss of taste (B), lacrimation (C), and decreased salivation (D).", 
#   "answer_choice": "D"
# }

medrag = MedRAG(llm_name="OpenAI/gpt-3.5-turbo-16k", rag=True, retriever_name="MedCPT", corpus_name="Textbooks")
answer, snippets, scores = medrag.answer(question=question, options=options, k=32) # scores are given by the retrieval system
print(f"Final answer in json with rationale: {answer}")
# {
#   "step_by_step_thinking": "A lesion causing compression of the facial nerve at the stylomastoid foramen will result in paralysis of the facial muscles. Loss of taste, lacrimation, and decreased salivation are not specifically mentioned in relation to a lesion at the stylomastoid foramen.", 
#   "answer_choice": "A"
# }

### MedRAG with pre-determined snippets
snippets = [{'id': 'InternalMed_Harrison_30037', 'title': 'InternalMed_Harrison', 'content': 'On side of lesion Horizontal and vertical nystagmus, vertigo, nausea, vomiting, oscillopsia: Vestibular nerve or nucleus Facial paralysis: Seventh nerve Paralysis of conjugate gaze to side of lesion: Center for conjugate lateral gaze Deafness, tinnitus: Auditory nerve or cochlear nucleus Ataxia: Middle cerebellar peduncle and cerebellar hemisphere Impaired sensation over face: Descending tract and nucleus fifth nerve On side opposite lesion Impaired pain and thermal sense over one-half the body (may include face): Spinothalamic tract Although atheromatous disease rarely narrows the second and third segments of the vertebral artery, this region is subject to dissection, fibromuscular dysplasia, and, rarely, encroachment by osteophytic spurs within the vertebral foramina.', 'contents': 'InternalMed_Harrison. On side of lesion Horizontal and vertical nystagmus, vertigo, nausea, vomiting, oscillopsia: Vestibular nerve or nucleus Facial paralysis: Seventh nerve Paralysis of conjugate gaze to side of lesion: Center for conjugate lateral gaze Deafness, tinnitus: Auditory nerve or cochlear nucleus Ataxia: Middle cerebellar peduncle and cerebellar hemisphere Impaired sensation over face: Descending tract and nucleus fifth nerve On side opposite lesion Impaired pain and thermal sense over one-half the body (may include face): Spinothalamic tract Although atheromatous disease rarely narrows the second and third segments of the vertebral artery, this region is subject to dissection, fibromuscular dysplasia, and, rarely, encroachment by osteophytic spurs within the vertebral foramina.'}]
answer, _, _ = medrag.answer(question=question, options=options, snippets=snippets)

### MedRAG with pre-determined snippet ids
snippets_ids = [{"id": s["id"]} for s in snippets]
answer, snippets, _ = medrag.answer(question=question, options=options, snippets_ids=snippets_ids)

medrag = MedRAG(llm_name="OpenAI/gpt-3.5-turbo-16k", rag=True, retriever_name="MedCPT", corpus_name="Textbooks", corpus_cache=True)
answer, snippets, scores = medrag.answer(question=question, options=options, k=32) # scores are given by the retrieval system
print(f"Final answer in json with rationale: {answer}")
# {
#   "step_by_step_thinking": "A lesion causing compression of the facial nerve at the stylomastoid foramen will result in paralysis of the facial muscles. Loss of taste, lacrimation, and decreased salivation are not specifically mentioned in relation to a lesion at the stylomastoid foramen.", 
#   "answer_choice": "A"
# }

medrag = MedRAG(llm_name="OpenAI/gpt-3.5-turbo-16k", rag=True, follow_up=True, retriever_name="MedCPT", corpus_name="Textbooks", corpus_cache=True)
answer, history = medrag.answer(question=question, options=options, k=32, n_rounds=4, n_queries=3)
print(f"Final answer in json: {answer}") # {'answer': 'A'}
print(f"Raw answer with analysis: {history[-3]}")
# {
#   'role': 'assistant', 
#   'content': '## Analysis:\nBased on the previous information provided, a lesion causing compression of the facial nerve at the stylomastoid foramen can result in an ipsilateral loss of motor function of the whole side of the face. This can lead to an unusual appearance and complications with chewing food. Lacrimation and taste may not be affected if the lesion remains distal to the greater petrosal and chorda tympani branches that originate deep in the temporal bone. However, decreased salivation can occur if the lesion affects the parasympathetic fibers that innervate the salivary glands.\n\n## Answer:\nThe correct answer is A. paralysis of the facial muscles.'
# }
print(f"Follow-up queries generated: {[item.split('Answer: ')[0].strip() for item in history[-4]['content'].split('Query: ')[1:]]}")
# [
#   'What are the functions of the facial nerve?', 
#   'What is the anatomical location and function of the stylomastoid foramen?', 
#   'What are the possible effects of a lesion causing compression of the facial nerve at the stylomastoid foramen?', 
#   'What are the specific branches of the facial nerve that control taste sensation and lacrimation?', 
#   'How does a lesion at the stylomastoid foramen affect salivation?', 
#   'Are there any other possible effects of a lesion causing compression of the facial nerve at the stylomastoid foramen?', 
#   'What are the effects of a lesion causing compression of the facial nerve at the stylomastoid foramen on salivation?', 
#   'Can a lesion at the stylomastoid foramen affect lacrimation and taste if it remains distal to the greater petrosal and chorda tympani branches?', 
#   'Are there any other possible effects of a lesion causing compression of the facial nerve at the stylomastoid foramen?', 
#   'What is the specific effect of a lesion causing compression of the facial nerve at the stylomastoid foramen on lacrimation?', 
#   'How does a lesion at the stylomastoid foramen affect taste sensation?', 
#   'Can a lesion at the stylomastoid foramen result in decreased salivation?'
# ]