import streamlit as st
from transformers import AutoTokenizer, pipeline
from peft import PeftModel, PeftConfig
from transformers import AutoModelForSeq2SeqLM
from datasets import load_dataset
import torch

st.write("Initializing...")  # Debugging message

# Load the LoRA configuration and model
config = PeftConfig.from_pretrained("lorahub/flan_t5_large-web_questions_potential_correct_answer")
base_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
model = PeftModel.from_pretrained(base_model, "lorahub/flan_t5_large-web_questions_potential_correct_answer")
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")

st.write("Model Loaded Successfully!")  # Debugging message

qa_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)

# Load relevant datasets
hotpotqa_dataset = load_dataset("bdsaglam/hotpotqa-distractor")
squad_v2_dataset = load_dataset("tom-010/squad_v2_with_answerable")
bias_professions_dataset = load_dataset("society-ethics/stable-bias-professions")
classifier_dataset = load_dataset("habanoz/classifier_1300_610_url_p")

st.write("Datasets Loaded Successfully!")  # Debugging message

# Streamlit App Structure
def main():
    st.title("AI-Powered Career Counseling App with Advanced Q&A")
    st.sidebar.title("Navigation")
    option = st.sidebar.selectbox("Choose an Option", ["Profile Setup", "Career Q&A", "Career Recommendations", "Resource Library"])

    if option == "Profile Setup":
        profile_setup()
    elif option == "Career Q&A":
        career_qa()
    elif option == "Career Recommendations":
        career_recommendations()
    elif option == "Resource Library":
        resource_library()

# Profile Setup Section
def profile_setup():
    st.header("Profile Setup")
    st.write("Fill out your details to personalize your experience.")
    age = st.number_input("Age", min_value=10, max_value=100)
    education = st.selectbox("Education Level", ["High School", "Undergraduate", "Graduate", "Other"])
    interests = st.text_area("Career Interests", "e.g., Data Science, Graphic Design")
    skills = st.text_area("Skills (comma-separated)", "e.g., Python, communication, empathy")
    
    if st.button("Save Profile"):
        st.session_state["profile"] = {
            "age": age,
            "education": education,
            "interests": interests.split(", "),
            "skills": skills.split(", ")
        }
        st.success("Profile saved successfully!")

# Q&A Section for Career-related questions
def career_qa():
    st.header("Career Q&A")
    question = st.text_input("Ask a career-related question")
    
    if st.button("Get Answer"):
        if question:
            # Prepare question for the model
            response = qa_pipeline(question)
            st.write("Answer:", response[0]['generated_text'])
        else:
            st.warning("Please enter a question.")

# Career Recommendations Section (Mockup, Extend as needed)
def career_recommendations():
    st.header("Career Recommendations")
    
    # Mock recommendation - In a real application, this should be based on a recommendation model
    st.write("Based on your interests and skills, we recommend:")
    st.write("1. Data Scientist")
    st.write("2. Software Engineer")
    st.write("3. Product Manager")

# Resource Library Section
def resource_library():
    st.header("Resource Library")
    st.write("Browse resources related to different careers.")
    career_choice = st.selectbox("Choose a career", ["Data Scientist", "Graphic Designer", "Software Engineer", "Nurse"])
    
    # Display resources for the chosen career
    st.write(f"### Resources for {career_choice}")
    st.write("1. Article: How to become a successful " + career_choice)
    st.write("2. Video: Day in the life of a " + career_choice)
    st.write("3. Guide: Top skills for " + career_choice)

if __name__ == "__main__":
    main()