import streamlit as st
import openai

# Initialize the OpenAI API
openai.api_key = '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'  # Remember never to expose API keys in code

KNOWN_MODELS = [
    # General ML models
    "Neural Networks", "Decision Trees", "Support Vector Machines", 
    "Random Forests", "Linear Regression", "Reinforcement Learning",
    "Logistic Regression", "k-Nearest Neighbors", "Naive Bayes",
    "Gradient Boosting Machines", "Regularization Techniques", 
    "Ensemble Methods", "Time Series Analysis",

    # Deep Learning models
    "Deep Learning", "Convolutional Neural Networks", 
    "Recurrent Neural Networks", "Transformer Models", 
    "Generative Adversarial Networks", "Autoencoders", 
    "Bidirectional LSTM", "Residual Networks (ResNets)", 
    "Variational Autoencoders",

    # Computer Vision models and techniques
    "Object Detection (e.g., YOLO, SSD)", "Semantic Segmentation",
    "Image Classification", "Face Recognition", "Optical Character Recognition (OCR)", 
    "Pose Estimation", "Style Transfer", "Image-to-Image Translation", 
    "Image Generation", "Capsule Networks",

    # NLP models and techniques
    "BERT", "GPT", "ELMo", "T5", "Word2Vec", "Doc2Vec", 
    "Topic Modeling", "Sentiment Analysis", "Text Classification", 
    "Machine Translation", "Speech Recognition", "Sequence-to-Sequence Models", 
    "Attention Mechanisms", "Named Entity Recognition", "Text Summarization"
]

def recommend_ai_model_via_gpt(description):
    messages = [
        {"role": "user", "content": f"Given the application described as: '{description}', which AI model would be most suitable?"}
    ]

    try:
        response = openai.ChatCompletion.create(
            model="gpt-4",
            messages=messages
        )
        recommendation = response['choices'][0]['message']['content'].strip()
        return recommendation
    except openai.error.OpenAIError as e:
        return f"Error: {e}"

def explain_recommendation(model_name):
    messages = [
        {"role": "user", "content": f"Why would {model_name} be a suitable choice for the application?"}
    ]
    
    try:
        response = openai.ChatCompletion.create(
            model="gpt-4",
            messages=messages
        )
        explanation = response['choices'][0]['message']['content'].strip()
        return explanation
    except openai.error.OpenAIError as e:
        return f"Error: {e}"

# Streamlit UI
st.title('AI Model Recommender')

description = st.text_area("Describe your application:", "")

if "rec_model_pressed" not in st.session_state:
    st.session_state.rec_model_pressed = False

if "feedback_submitted" not in st.session_state:
    st.session_state.feedback_submitted = False

if st.button("Recommend AI Model"):
    st.session_state.rec_model_pressed = True

if st.session_state.rec_model_pressed:
    if description:
        recommended_model = recommend_ai_model_via_gpt(description)
        
        # Validate recommended model
        # Commenting out model validation for the example
        # if recommended_model not in KNOWN_MODELS:
        #     st.warning("The recommendation is ambiguous. Please refine your description or consult an expert.")
        # else:
        st.subheader(f"Recommended AI Model: {recommended_model}")
        explanation = explain_recommendation(recommended_model)
        st.write("Reason:", explanation)

        # Collecting rating and feedback through Streamlit
        rating = st.slider("Rate the explanation from 1 (worst) to 5 (best):", 1, 5)
        feedback = st.text_input("Any additional feedback?")

        if st.button("Submit Feedback"):
            st.session_state.feedback_submitted = True

        if st.session_state.feedback_submitted:
            st.success("Thank you for your feedback! \nContact team@autumn8.ai or +1 (857) 600-0180 to learn how we can fine-tune and host this app for you.")
    else:
        st.warning("Please provide a description.")