from langchain_community.document_loaders import DirectoryLoader
from ragas.llms import LangchainLLMWrapper
from ragas.embeddings import LangchainEmbeddingsWrapper
from langchain_openai import ChatOpenAI
from langchain_openai import OpenAIEmbeddings
from ragas.testset import TestsetGenerator

path = "Sample_Docs_Markdown/"
loader = DirectoryLoader(path, glob="**/*.md")
docs = loader.load()
generator_llm = LangchainLLMWrapper(ChatOpenAI(model="gpt-4o"))
generator_embeddings = LangchainEmbeddingsWrapper(OpenAIEmbeddings())

generator = TestsetGenerator(llm=generator_llm, embedding_model=generator_embeddings)
dataset = generator.generate_with_langchain_docs(docs, testset_size=10)

kg = KnowledgeGraph()
for doc in docs:
    kg.nodes.append(
        Node(
            type=NodeType.DOCUMENT,
            properties={"page_content": doc.page_content, "document_metadata": doc.metadata}
        )
    )

transformer_llm = LangchainLLMWrapper(ChatOpenAI(model="gpt-4o"))
embedding_model = LangchainEmbeddingsWrapper(OpenAIEmbeddings())

transformer = default_transforms(documents=docs, llm=transformer_llm, embedding_model=embedding_model)
apply_transforms(kg, transformer)


generator = TestsetGenerator(llm=generator_llm, embedding_model=embedding_model, knowledge_graph=kg)
query_distribution = [
    (SingleHopSpecificQuerySynthesizer(llm=generator_llm), 0.5),
    (SingleHopAbstractQuerySynthesizer(llm=generator_llm), 0.25),
    (MultiHopSpecificQuerySynthesizer(llm=generator_llm), 0.25),
]
testset = generator.generate(testset_size=10, query_distribution=query_distribution)

{
    "question": "Who was Albert Einstein and what is he best known for?",
    "answer": "He was a German-born theoretical physicist, widely acknowledged to be one of the greatest and most influential physicists of all time. He was best known for developing the theory of relativity, he also made important contributions to the development of the theory of quantum mechanics.",
    "sentences": {
        "0": "He was a German-born theoretical physicist, widely acknowledged to be one of the greatest and most influential physicists of all time.",
        "1": "He was best known for developing the theory of relativity, he also made important contributions to the development of the theory of quantum mechanics."
    }
}

CORRECTNESS_INSTRUCTIONS = """\
Given a ground truth and an answer statements, analyze each statement and classify them in one of the following categories:

- TP (true positive): statements that are present in answer that are also directly supported by the one or more statements in ground truth,
- FP (false positive): statements present in the answer but not directly supported by any statement in ground truth,
- FN (false negative): statements found in the ground truth but not present in answer.

Each statement can only belong to one of the categories. Provide a reason for each classification.
"""

"question": """What powers the sun and what is its primary function?""",
  "answer": [
    "The sun is powered by nuclear fission, similar to nuclear reactors on Earth.",
    "The primary function of the sun is to provide light to the solar system.",
  ],
    "ground_truth": [
      "The sun is powered by nuclear fusion, where hydrogen atoms fuse to form helium.",
      "This fusion process in the sun's core releases a tremendous amount of energy.",
      "The energy from the sun provides heat and light, which are essential for life on Earth.",
      "The sun's light plays a critical role in Earth's climate system.",
      "Sunlight helps to drive the weather and ocean currents.",
    ],

"TP": [
  {
    "statement": "The primary function of the sun is to provide light to the solar system.",
    "reason": "This statement is somewhat supported by the ground truth mentioning the sun providing light and its roles, though it focuses more broadly on the sun's energy."
  }
],
"FP": [
  {
    "statement": "The sun is powered by nuclear fission, similar to nuclear reactors on Earth.",
    "reason": "This statement is incorrect and contradicts the ground truth which states that the sun is powered by nuclear fusion."
  }
],
"FN": [
  {
    "statement": "The sun is powered by nuclear fusion, where hydrogen atoms fuse to form helium.",
    "reason": "This accurate description of the sun’s power source is not included in the answer."
  }
]