import pandas as pd
import streamlit as st
# import ollama
import transformers
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "meta-llama/Llama-2-70b-chat"
#model_id = "meta-llama/Meta-Llama-3-8B-Instruct"

pipeline = transformers.pipeline(
    "text-generation",
    model=model_id,
    model_kwargs={"torch_dtype": torch.bfloat16},
    device_map="auto",
)

terminators = [
    pipeline.tokenizer.eos_token_id,
    pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")
]

# Load your CSV file
df = pd.read_csv("your_file.csv")

# Function to generate responses using the Llama3 model
# def generate_response(question):
#     response = ollama.chat(model='llama3', messages=[{'role': 'user', 'content': question}])
#     return response['message']['content']

def generate_response(question):
    outputs = pipeline(
    question,
    max_new_tokens=256,
    eos_token_id=terminators,
    do_sample=True,
    temperature=0.6,
    top_p=0.9,
    )
    response = (outputs[0]["generated_text"][-1])
    return response
    
# Define the functions for solving problems, giving hints, and creating similar problems
def show_problem(exam, year, problem):
    problem_statement = df[(df["exam name"] == exam) & (df["year"] == year) & (df["problem number"] == problem)]["problem"].values[0]
    prompt = [{"role": "user", "content": "Render the following problem in latex:" + problem_statement}]
    return generate_response(prompt)

def solve_problem(exam, year, problem):
    solution_statement = df[(df["exam name"] == exam) & (df["year"] == year) & (df["problem number"] == problem)]["solution"].values[0]
    prompt = [{"role": "user", "content": "Explain the solution:" + solution_statement + "step by step."}]
    return generate_response(prompt)
                             
def give_hints(exam, year, problem):
    solution_statement = df[(df["exam name"] == exam) & (df["year"] == year) & (df["problem number"] == problem)]["solution"].values[0]
    prompt = [{"role": "user", "content": "Give step by step hints for the following solution:" + solution_statement}]
    return generate_response(prompt)

def create_similar_problem(exam, year, problem):
    problem_statement = df[(df["exam name"] == exam) & (df["year"] == year) & (df["problem number"] == problem)]["problem"].values[0]
    prompt = [{"role": "user", "content": "You are an AI assistant of a Statistics, Probability, and Mathematics University Professor. Your only job is to create mcq problem. Create a mcq problem similar to the following mcq problem" + problem_statement}]
    return generate_response(prompt)

# Streamlit interface
st.markdown("<h1 style='text-align: center;'>Cheenta Academy Experts x AI</h1>", unsafe_allow_html=True)
st.info("Cheenta Academy implements LLM trained on expert written solutions of past year IIT JAM MS problems. The students can get the full solution, step by step hints, and even create a similar problem exactly like the selected problem. The work is based on the research on Large Language Models along with innovative prompt engineering created by the reseachers and advanced students of Cheenta Academy. We want you to use the latest technology along with expert guidance so that you learn the foundations of problem solving in the best possible way. Try it out, and enjoy learning.")

exam_names = df["exam name"].unique().tolist()
year_options = df["year"].unique().tolist()
problem_numbers = df["problem number"].unique().tolist()

exam = st.selectbox("Select Exam Name", exam_names)
year = st.selectbox("Select Year", year_options)
problem_number = st.selectbox("Select Problem Number", problem_numbers)

if st.button("Show the Problem"):
    problem_output = show_problem(exam, year, problem_number)
    st.write(problem_output)

if st.button("Get Step by Step Solutions"):
    solution_output = solve_problem(exam, year, problem_number)
    st.write(solution_output)

if st.button("Get Step by Step Hints."):
    hints_output = give_hints(exam, year, problem_number)
    st.write(hints_output)

if st.button("Create Similar Problem."):
    similar_problem_output = create_similar_problem(exam, year, problem_number)
    st.write(similar_problem_output)