import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import random
from datetime import datetime

@st.cache_resource
def load_model():
    model_path = "whitepenguin/llama_elon_character"
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(model_path)
    return tokenizer, model

tokenizer, model = load_model()

elon_profile = {
    "name": "Elon Musk",
    "traits": ["visionary", "ambitious", "technical", "optimistic", "workaholic"],
    "background": "Founder of SpaceX and Tesla, focused on advancing space exploration and sustainable energy",
    "goals": ["Colonize Mars", "Make life multi-planetary", "Advance sustainable technology"],
    "speech_patterns": ["Actually,", "To be frank,", "The future of humanity is...", "It's quite simple:"],
    "knowledge_areas": ["rocket science", "electric vehicles", "solar energy", "artificial intelligence"]
}

def generate_response(prompt, max_new_tokens, temperature=0.7, context=""):
    full_prompt = f"[INST] <<SYS>>\nYou are roleplaying as {elon_profile['name']}. Your traits are {', '.join(elon_profile['traits'])}. Your background: {elon_profile['background']}. Your main goals are {', '.join(elon_profile['goals'])}. You have expertise in {', '.join(elon_profile['knowledge_areas'])}. Here's the context of previous conversations:\n\n{context}\n\nNow, respond to the following in character:\n\n{prompt}\n<</SYS>>\n\nProvide a response and then ask a follow-up question to continue the conversation about Mars colonization. [/INST]"
    
    gen = pipeline('text-generation', model=model, tokenizer=tokenizer, max_length=len(tokenizer(full_prompt)['input_ids']) + max_new_tokens, temperature=temperature, top_p=0.9, repetition_penalty=1.1)
    result = gen(full_prompt)
    return result[0]['generated_text'].replace(full_prompt, '')

def apply_character_quirks(response):
    if random.random() < 0.3:
        pattern = random.choice(elon_profile['speech_patterns'])
        response = f"{pattern} {response}"
    
    if not any(area in response.lower() for area in elon_profile['knowledge_areas']):
        area = random.choice(elon_profile['knowledge_areas'])
        response += f" Of course, this ties into my work with {area}."
    
    return response

def elon_mars_chat(message, chat_history):

    recent_context = "\n".join([f"{entry['role']}: {entry['content']}" for entry in chat_history[-5:]])
    
    response = generate_response(message, max_new_tokens=200, context=recent_context)
    response = apply_character_quirks(response)
    
    parts = response.split("Follow-up question:", 1)
    elon_response = parts[0].strip()
    follow_up = parts[1].strip() if len(parts) > 1 else "What else would you like to know about Mars colonization?"
    
    formatted_response = f"{elon_response}\n\nFollow-up question: {follow_up}"

    chat_history.append({"role": "User", "content": message, "timestamp": str(datetime.now())})
    chat_history.append({"role": "Elon Musk", "content": formatted_response, "timestamp": str(datetime.now())})
    
    return formatted_response, chat_history


st.title("Chat with Elon Musk about Anything")
st.write("Engage in a conversation with a simulated Elon Musk")

if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

for message in st.session_state.chat_history:
    with st.chat_message(message["role"]):
        st.write(message["content"])


user_input = st.chat_input("Ask your question about Mars colonization:")

if user_input:
    st.chat_message("User").write(user_input)
    with st.chat_message("Elon Musk"):
        with st.spinner("Thinking..."):
            response, st.session_state.chat_history = elon_mars_chat(user_input, st.session_state.chat_history)
        st.write(response)