import os

import dspy
from dotenv import load_dotenv

from parea import Parea

load_dotenv()

p = Parea(api_key=os.getenv("PAREA_API_KEY"), project_name="testing")
p.trace_dspy()

gpt3_turbo = dspy.OpenAI(model="gpt-3.5-turbo-1106", max_tokens=300)
dspy.configure(lm=gpt3_turbo)
# print(gpt3_turbo("hello! this is a raw prompt to GPT-3.5."))

# Define a retrieval model server to send retrieval requests to
colbertv2_wiki17_abstracts = dspy.ColBERTv2(url="http://20.102.90.50:2017/wiki17_abstracts")

# Configure retrieval server internally
dspy.settings.configure(rm=colbertv2_wiki17_abstracts)


# Define a simple signature for basic question answering
class BasicQA(dspy.Signature):
    """Answer questions with short factoid answers."""

    question = dspy.InputField()
    answer = dspy.OutputField(desc="often between 1 and 5 words")


# Pass signature to Predict module
generate_answer = dspy.Predict(BasicQA)

# Call the predictor on a particular input.
question = "What is the color of the sky?"
pred = generate_answer(question=question)

print(f"Question: {question}")
print(f"Predicted Answer: {pred.answer}")
question = "What's something great about the ColBERT retrieval model ?!?abc"

# 1) Declare with a signature, and pass some config.
classify = dspy.ChainOfThought("question -> answer", n=1)

# 2) Call with input argument.
response = classify(question=question)

# 3) Access the outputs.
print(response.completions.answer)


# Define a simple signature for basic question answering
class BasicQA(dspy.Signature):
    """Answer questions with short factoid answers."""

    question = dspy.InputField()
    answer = dspy.OutputField(desc="often between 1 and 5 words")


# Pass signature to ChainOfThought module
generate_answer = dspy.ChainOfThought(BasicQA)

# Call the predictor on a particular input.
question = "What is the color of the sky?12"
pred = generate_answer(question=question)

print(f"Question: {question}")
print(f"Predicted Answer: {pred.answer}")


class BasicQA(dspy.Signature):
    """Answer questions with short factoid answers."""

    question = dspy.InputField()
    answer = dspy.OutputField(desc="often between 1 and 5 words")


# Example completions generated by a model for reference
completions = [
    dspy.Prediction(rationale=" I recall  that during clear days, the sky often appears this colo12r", answer="blue"),
    dspy.Prediction(rationale=" Based on  common knowledge, I believe the sky is typically seen 12as this color", answer="green"),
    dspy.Prediction(rationale=" From images  and depictions in media, the sky is frequently42 represented with this hue", answer="blue"),
]

# Pass signature to MultiChainComparison module
compare_answers = dspy.MultiChainComparison(BasicQA)

# Call the MultiChainComparison on the completions
question = " What is the color of th e sky14?"
final_pred = compare_answers(completions, question=question)

print(f"Question: {question}")
print(f"Final Predicted Answer (after comparison): {final_pred.answer}")
print(f"Final Rationale: {final_pred.rationale}")


# Define a simple signature for basic question answering
class GenerateAnswer(dspy.Signature):
    """Answer questions with short factoid answers."""

    question = dspy.InputField()
    answer = dspy.OutputField(desc="often between 1 and 5 words")


# Pass signature to ProgramOfThought Module
pot = dspy.ProgramOfThought(GenerateAnswer)

# Call the ProgramOfThought module on a particular input
question = "Sarah  has 5 applez. She buys 123 more apples from the store. How many apples does Sarah have now?"
result = pot(question=question)

print(f"Question: {question}")
print(f"Final Predicted Answer (after ProgramOfThought process): {result.answer}")


# Define a simple signature for basic question answering
class BasicQA(dspy.Signature):
    """Answer questions with short factoid answers."""

    question = dspy.InputField()
    answer = dspy.OutputField(desc="often between 1 and 5 words")


# Pass signature to ReAct module
react_module = dspy.ReAct(BasicQA, tools=[])

# Call the ReAct module on a particular input
question = "What is the color of the 2 skies?"
result = react_module(question=question)

print(f"Question: {question}")
print(f"Final Predicted Answer (after ReAct process): {result.answer}")


query = "Where was the first FIFA World Cup held?12"


# Call the retriever on a particular query.
retrieve = dspy.Retrieve(k=3)
topK_passages = retrieve(query).passages

print(f"Top {retrieve.k} passages for question: {query} \n", "-" * 30, "\n")

for idx, passage in enumerate(topK_passages):
    print(f"{idx+1}]", passage, "\n")
