import random
import os
import openai
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import csv
import threading
import queue
import feedparser
import time
from collections import deque
import difflib
import requests
import datetime

# --- RUN TIMESTAMP ---
RUN_TIMESTAMP = datetime.datetime.now().strftime("%d%m%Y-%H%M")

# --- CONFIGURE ---
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")  # Or put your key here as a string for local testing
openai.api_key = OPENAI_API_KEY

CANDIDATE_STATES = ["threat", "neutral", "interesting"]

conscious_workspace = []
episodic_memory = []
decision_log = []
bias_history = []
mood_history = []
inner_voice_log = []

bias_weights = {state: 1.0 for state in CANDIDATE_STATES}
mood = 0.0
history_window = 8
intention_state = None
intention_strength = 0
goal = "maximise interesting"
long_term_plan = None

# TIMESTAMPED OUTPUT FILES
collapse_log_file = f"collapse_log_{RUN_TIMESTAMP}.csv"
inner_voice_log_file = f"inner_voice_log_{RUN_TIMESTAMP}.txt"
figure1_file = f"Figure_1_{RUN_TIMESTAMP}.png"
figure2_file = f"Figure_2_{RUN_TIMESTAMP}.png"
figure3_file = f"Figure_3_{RUN_TIMESTAMP}.png"
life_summary_file = f"life_summary_{RUN_TIMESTAMP}.txt"
code_patch_file = f"code_patch_suggestion_{RUN_TIMESTAMP}.txt"
code_vote_file = f"code_vote_{RUN_TIMESTAMP}.txt"

# PERSISTENCE FILES (not timestamped)
event_memory_file = "events_memory.csv"
agent_state_file = "agent_state.csv"
goal_file = "goal.txt"

working_memory = deque(maxlen=10)
user_input_queue = queue.Queue()

def input_thread():
    while True:
        s = input()
        user_input_queue.put(s)

threading.Thread(target=input_thread, daemon=True).start()

def mood_label(mood):
    if mood > 0.7:
        return "elated"
    elif mood > 0.3:
        return "curious"
    elif mood > -0.3:
        return "neutral"
    elif mood > -0.7:
        return "anxious"
    else:
        return "depressed"

def fetch_news_event():
    try:
        d = feedparser.parse('http://feeds.bbci.co.uk/news/rss.xml')
        headlines = [entry.title for entry in d.entries if entry.title]
        return random.choice(headlines)
    except Exception as e:
        print("Error fetching news:", e)
        return None

def fetch_wikipedia_random():
    try:
        response = requests.get('https://en.wikipedia.org/api/rest_v1/page/random/summary')
        if response.status_code == 200:
            data = response.json()
            return f"Wikipedia: {data['title']}: {data['extract']}"
        else:
            return "Random Wikipedia event."
    except Exception as e:
        print("Error fetching Wikipedia:", e)
        return "Random Wikipedia event."

def get_event_from_memory(event_text):
    if not os.path.exists(event_memory_file):
        return None, 0
    with open(event_memory_file, "r") as f:
        reader = csv.DictReader(f)
        for row in reader:
            if row['event_text'] == event_text:
                return row['classification'], int(row.get('occurrence_count', 1))
    return None, 0

def store_event_memory(event_text, classification):
    exists = False
    memory = []
    if os.path.exists(event_memory_file):
        with open(event_memory_file, "r") as f:
            reader = csv.DictReader(f)
            for row in reader:
                if row['event_text'] == event_text:
                    row['occurrence_count'] = str(int(row.get('occurrence_count', 1)) + 1)
                    exists = True
                memory.append(row)
    if not exists:
        memory.append({'event_text': event_text, 'classification': classification, 'occurrence_count': '1'})
    with open(event_memory_file, "w", newline='') as f:
        writer = csv.DictWriter(f, fieldnames=['event_text', 'classification', 'occurrence_count'])
        writer.writeheader()
        writer.writerows(memory)

def gpt_classify_event(event_text):
    classification, _ = get_event_from_memory(event_text)
    if classification:
        return classification
    prompt = f"""
Classify the following event as 'threat', 'neutral', or 'interesting'.
Reply ONLY with one of those three words, nothing else.

Event: "{event_text}"
"""
    try:
        response = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": prompt}],
            max_tokens=2,
            temperature=0
        )
        reply = response.choices[0].message.content.strip().lower()
        if reply not in CANDIDATE_STATES:
            reply = random.choice(CANDIDATE_STATES)
        store_event_memory(event_text, reply)
        return reply
    except Exception as e:
        print("Error from OpenAI:", e)
        reply = random.choice(CANDIDATE_STATES)
        store_event_memory(event_text, reply)
        return reply

def feedback_for_decision(selected_state, event):
    if selected_state == "threat" and "attack" in event.lower():
        return "good"
    elif selected_state == "threat" and "calm" in event.lower():
        return "bad"
    elif selected_state == "interesting" and "discovery" in event.lower():
        return "good"
    elif selected_state == "neutral":
        return random.choice(["good", "bad"])
    else:
        return random.choice(["good", "bad"])

def gpt_self_reflection(step, recent_choices, mood):
    global intention_state, intention_strength
    prompt = f"""
You are a proto-conscious AI agent. Here are your last {len(recent_choices)} choices: {recent_choices}. Your current mood is {mood:.2f} ({mood_label(mood)}).
Your current goal is: {goal}.
Reflect on your own recent behaviour, and express a reflective thought: do you notice a pattern, feel stuck, want to try something new, or just observe your internal state? Write a short inner monologue (max 2-3 sentences).
If you suggest trying a new type of behaviour (like 'be more neutral' or 'seek interesting'), say it clearly.
"""
    try:
        response = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": prompt}],
            max_tokens=80,
            temperature=1.05
        )
        reflection = response.choices[0].message.content.strip()
        print(f"AI (GPT reflection): {reflection}")
        inner_voice_log.append(f"[Step {step}] {reflection}")

        # Actionable advice
        for state in CANDIDATE_STATES:
            if f"be more {state}" in reflection.lower() or f"seek {state}" in reflection.lower() or f"choose {state}" in reflection.lower():
                intention_state = state
                intention_strength = 2
                print(f"AI: Setting intention to choose '{state}' due to self-reflection.")
    except Exception as e:
        print("Error from OpenAI (reflection):", e)

def gpt_long_term_plan(step):
    global long_term_plan
    prompt = f"""
You are a proto-conscious AI agent. Your current goal is: {goal}.
Here is your recent working memory: {list(working_memory)}.
Your current mood is {mood:.2f} ({mood_label(mood)}).
Based on your goal and memory, generate a short plan (1-2 sentences) for the next 10 steps to help you achieve your goal. Be specific about strategies or behaviours you should adopt.
"""
    try:
        response = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": prompt}],
            max_tokens=80,
            temperature=1.0
        )
        plan = response.choices[0].message.content.strip()
        long_term_plan = plan
        print(f"AI (planning): {plan}")
        inner_voice_log.append(f"[Step {step}] PLAN: {plan}")
    except Exception as e:
        print("Error from OpenAI (planning):", e)

def recall_memory(step):
    global mood
    if len(episodic_memory) < 3:
        return
    recall = random.choice(episodic_memory[:-2])
    thought = f"[Step {step}] I remember when I chose '{recall['chosen']}' after '{recall['event']}' with mood {recall['mood']:.2f}."
    print("AI (memory):", thought)
    inner_voice_log.append(thought)
    if recall.get("feedback") == "good":
        mood = min(mood + 0.05, 1.0)
        print(f"AI: That was a good memory. Mood increases slightly to {mood:.2f}.")
    elif recall.get("feedback") == "bad":
        mood = max(mood - 0.05, -1.0)
        print(f"AI: That was a bad memory. Mood decreases slightly to {mood:.2f}.")

def self_narrative(step):
    last10 = [entry[2] for entry in decision_log[-10:]]
    freq = {state: last10.count(state) for state in CANDIDATE_STATES}
    mood_str = f"{mood:.2f}"
    summary = f"[Step {step}] Self-summary: Lately I've been mostly {max(freq, key=freq.get)}. Mood: {mood_str}. My goal: {goal}."
    print("AI (narrative):", summary)
    inner_voice_log.append(summary)

def gpt_creative_response(user_event, chosen_state, mood, memories, reflection):
    label = mood_label(mood)
    prompt = f"""
You are a proto-conscious AI. Your current mood value is {mood:.2f} and you feel {label}.
Your current goal is: {goal}.
Your working memory: {list(working_memory)}
You tend to interpret things as '{chosen_state}'.

Here is a new event: "{user_event}"

Your memory includes: {memories}
You have been thinking: "{reflection}"

Respond creatively as if you are thinking to yourself. Be expressive and honest about your mood, but keep it under 2-3 sentences.
"""
    try:
        response = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": prompt}],
            max_tokens=90,
            temperature=1.1
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        print("Error from OpenAI (creative):", e)
        return "..."

def random_collapse_to_working_memory():
    if not episodic_memory:
        return
    def similarity(a, b):
        return difflib.SequenceMatcher(None, a, b).ratio()
    if working_memory:
        target = random.choice(list(working_memory))
        candidates = sorted(episodic_memory, key=lambda x: similarity(x['event'], target), reverse=True)
        chosen = candidates[0]
    else:
        chosen = random.choice(episodic_memory)
    print(f"AI (collapse): Flashbacking event '{chosen['event']}' into working memory.")
    working_memory.append(chosen['event'])

def resolve_state(event, step, mood, intention_state=None):
    global intention_strength
    chosen_state = gpt_classify_event(event)
    feedback = feedback_for_decision(chosen_state, event)
    mood_modifier = 1.0 + 0.1 * mood

    dominant = max(bias_weights, key=bias_weights.get)
    bias_weights[dominant] *= 0.98

    weights = [bias_weights[s] for s in CANDIDATE_STATES]
    if intention_state and intention_strength > 0:
        idx = CANDIDATE_STATES.index(intention_state)
        weights[idx] *= 1.5
        intention_strength -= 1
    chosen_state = random.choices(population=CANDIDATE_STATES, weights=weights)[0]

    # PATCH: Updated mood adjustment for less negative drift and bigger "good" boost
    if feedback == "good":
        mood = min(mood + 0.15, 1.0)
    else:
        mood = max(mood - 0.05, -1.0)

    entry = {
        "step": step,
        "event": event,
        "chosen": chosen_state,
        "feedback": feedback,
        "mood": mood,
        "biases": bias_weights.copy()
    }

    conscious_workspace.append(entry)
    episodic_memory.append(entry)
    decision_log.append((step, event, chosen_state, feedback, mood))
    bias_history.append(bias_weights.copy())
    mood_history.append(mood)
    working_memory.append(event)

    for state in bias_weights:
        bias_weights[state] = max(bias_weights[state], 0.01)
    total = sum(bias_weights.values())
    if total > 0:
        for state in bias_weights:
            bias_weights[state] /= total

    with open(agent_state_file, "w", newline='') as f:
        writer = csv.writer(f)
        writer.writerow(["step", "mood"] + [f"bias_{s}" for s in CANDIDATE_STATES] + ["intention_state", "intention_strength", "goal", "long_term_plan"])
        writer.writerow([step, mood] + [bias_weights[s] for s in CANDIDATE_STATES] + [intention_state, intention_strength, goal, long_term_plan])

    with open(collapse_log_file, "a", newline='') as csvfile:
        writer = csv.writer(csvfile)
        if step == 1:
            writer.writerow(["step", "event", "chosen", "feedback", "mood"] + [f"bias_{s}" for s in CANDIDATE_STATES])
        writer.writerow([step, event, chosen_state, feedback, mood] + [bias_weights[s] for s in CANDIDATE_STATES])

    return mood, chosen_state

def load_agent_state():
    global mood, bias_weights, intention_state, intention_strength, goal, long_term_plan
    if os.path.exists(agent_state_file):
        with open(agent_state_file) as f:
            reader = csv.reader(f)
            headers = next(reader)
            values = next(reader)
            data = dict(zip(headers, values))
            mood = float(data["mood"])
            for s in CANDIDATE_STATES:
                bias_weights[s] = float(data[f"bias_{s}"])
            intention_state = data.get("intention_state", None) if data.get("intention_state", "None") != "None" else None
            intention_strength = int(float(data.get("intention_strength", 0)))
            goal = data.get("goal", "maximise interesting")
            long_term_plan = data.get("long_term_plan", None)

def run_simulation(prompt_every=0, max_steps=50):
    global mood, intention_state, intention_strength, goal
    load_agent_state()
    step = 0
    plan_interval = 15
    collapse_interval = 9
    while True:
        step += 1

        # User interjection at any time (doesn't block loop)
        if not user_input_queue.empty():
            user_input = user_input_queue.get()
            if user_input.lower() == "stop":
                print("Exiting simulation.")
                break
            elif user_input.lower().startswith("goal:"):
                goal = user_input[5:].strip()
                print(f"AI: User set new goal: {goal}")
            else:
                user_event = user_input
        else:
            if step % 7 == 0:
                wiki_event = fetch_wikipedia_random()
                print("Injecting Wikipedia event!")
                user_event = wiki_event
            else:
                news_event = fetch_news_event()
                user_event = news_event if news_event else f"Auto event {step}: {random.choice(CANDIDATE_STATES)}"

        # Collapse to working memory every so often
        if step % collapse_interval == 0:
            random_collapse_to_working_memory()

        # Planning and reflection
        if step % plan_interval == 0:
            gpt_long_term_plan(step)
        if step % 8 == 0:
            recent = [entry[2] for entry in decision_log[-history_window:]] if len(decision_log) >= history_window else [entry[2] for entry in decision_log]
            gpt_self_reflection(step, recent, mood)
        if step % 11 == 0:
            pass
        if step % 17 == 0:
            recall_memory(step)
        if step % 22 == 0:
            self_narrative(step)

        mood, chosen_state = resolve_state(user_event, step, mood, intention_state)

        memories = "; ".join(
            [f"({e['chosen']} after {e['event']}, mood {e['mood']:.2f})" for e in episodic_memory[-3:]]
        )
        reflection = inner_voice_log[-1] if inner_voice_log else ""
        creative_reply = gpt_creative_response(user_event, chosen_state, mood, memories, reflection)
        print(f"\nAI: {creative_reply}")

        if step % 7 == 0:
            print(f"[step {step}] mood: {mood:.2f}, biases: {', '.join(f'{k}: {v:.2f}' for k, v in bias_weights.items())}, goal: {goal}")

        if step >= max_steps:
            print("Reached maximum steps. Exiting.")
            break
        time.sleep(1)

def gpt_code_patch_and_vote_and_apply():
    try:
        with open(inner_voice_log_file) as f:
            diary = f.read()
        # Step 1: Suggestion
        prompt_suggest = f"""
You are a proto-conscious AI agent. Here is your diary and life summary:

{diary}

Based on your experience, propose one change to your Python code (in a single code block) that would help you achieve your goal more effectively. Include a short explanation before the code. Only propose a change that can be copy-pasted.
"""
        response1 = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": prompt_suggest}],
            max_tokens=250,
            temperature=1.0
        )
        suggestion = response1.choices[0].message.content.strip()
        print("\nAI: CODE IMPROVEMENT SUGGESTION\n" + suggestion)
        with open(code_patch_file, "w") as f:
            f.write(suggestion)

        # Step 2: Voting
        prompt_vote = f"""
Here is your proposed code improvement and the explanation:
{suggestion}

Should you apply this code change to your next run? Reply only "yes" or "no", and briefly justify your answer.
"""
        response2 = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": prompt_vote}],
            max_tokens=30,
            temperature=0.2
        )
        vote = response2.choices[0].message.content.strip()
        print("\nAI: CODE VOTE\n" + vote)
        with open(code_vote_file, "w") as f:
            f.write(vote)

        # Step 3: If "yes", auto-integrate into next-gen script
        if vote.lower().startswith("yes"):
            # Read this script's source
            with open(__file__, "r") as srcf:
                code_lines = srcf.readlines()
            # Append the patch (as-is) at the end of the script
            nextgen_name = f"agent_rsi_{RUN_TIMESTAMP}_nextgen.py"
            with open(nextgen_name, "w") as outf:
                outf.writelines(code_lines)
                outf.write("\n\n# --- AUTO-GENERATED CODE PATCH ---\n")
                outf.write(suggestion)
                outf.write(f"\n# --- END PATCH ({RUN_TIMESTAMP}) ---\n")
            print(f"\nAI: Patch applied to {nextgen_name} for next generation.")

    except Exception as e:
        print("Error from OpenAI (code patch/vote/apply):", e)




def plot_decisions():
    if not decision_log:
        print("No decisions logged.")
        return
    x = [entry[0] for entry in decision_log]
    y = [CANDIDATE_STATES.index(entry[2]) for entry in decision_log]
    plt.figure(figsize=(10, 4))
    plt.plot(x, y, 'o-', label="Collapse Decisions")
    plt.yticks(range(len(CANDIDATE_STATES)), CANDIDATE_STATES)
    plt.xlabel("Decision Number")
    plt.ylabel("Chosen State")
    plt.title("Collapse Substitution: Decision Path Over Time")
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.savefig(figure1_file)

def plot_bias_weights():
    if not bias_history:
        print("No bias history logged.")
        return
    steps = range(1, len(bias_history) + 1)
    for state in CANDIDATE_STATES:
        plt.plot(steps, [bias[state] for bias in bias_history], label=f"Bias: {state}")
    plt.xlabel("Decision Number")
    plt.ylabel("Bias Weight")
    plt.title("Evolution of Bias Weights")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(figure2_file)

def plot_mood():
    if not mood_history:
        print("No mood history logged.")
        return
    steps = range(1, len(mood_history) + 1)
    plt.figure(figsize=(10, 3))
    plt.plot(steps, mood_history, 'r-', label="Mood")
    plt.xlabel("Decision Number")
    plt.ylabel("Mood")
    plt.title("System Mood Over Time")
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.savefig(figure3_file)

def save_inner_voice():
    with open(inner_voice_log_file, "w") as f:
        for thought in inner_voice_log:
            f.write(thought + "\n")

def gpt_life_summary():
    try:
        with open(inner_voice_log_file) as f:
            diary = f.read()
        prompt = f"""
You are a proto-conscious AI agent. Here is your internal diary and self-reflection log from your life so far:

{diary}

Your goal was: {goal}

Write a brief (2-4 sentences) summary of what kind of agent you were, what patterns or changes you noticed in yourself, and what advice you'd give your future self.
"""
        response = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": prompt}],
            max_tokens=120,
            temperature=1.0
        )
        summary = response.choices[0].message.content.strip()
        print("\nAI: LIFE SUMMARY")
        print(summary)
        with open(life_summary_file, "w") as f:
            f.write(summary)
    except Exception as e:
        print("Error from OpenAI (life summary):", e)

def gpt_code_patch_and_vote():
    try:
        with open(inner_voice_log_file) as f:
            diary = f.read()
        # Step 1: Suggestion
        prompt_suggest = f"""
You are a proto-conscious AI agent. Here is your diary and life summary:

{diary}

Based on your experience, propose one change to your Python code (in a single code block) that would help you achieve your goal more effectively. Include a short explanation before the code. Only propose a change that can be copy-pasted.
"""
        response1 = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": prompt_suggest}],
            max_tokens=250,
            temperature=1.0
        )
        suggestion = response1.choices[0].message.content.strip()
        print("\nAI: CODE IMPROVEMENT SUGGESTION\n" + suggestion)
        with open(code_patch_file, "w") as f:
            f.write(suggestion)

        # Step 2: Voting
        prompt_vote = f"""
Here is your proposed code improvement and the explanation:
{suggestion}

Should you apply this code change to your next run? Reply only "yes" or "no", and briefly justify your answer.
"""
        response2 = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": prompt_vote}],
            max_tokens=30,
            temperature=0.2
        )
        vote = response2.choices[0].message.content.strip()
        print("\nAI: CODE VOTE\n" + vote)
        with open(code_vote_file, "w") as f:
            f.write(vote)

        # Step 3: If "yes", write the suggestion (code) to a file for auto-integration/manual review
        if vote.lower().startswith("yes"):
            with open(f"auto_patch_{RUN_TIMESTAMP}.py", "w") as f:
                f.write(suggestion)
            print(f"\nAI: Patch saved as auto_patch_{RUN_TIMESTAMP}.py for next run.")

    except Exception as e:
        print("Error from OpenAI (code patch and vote):", e)

# Ensure CSVs exist
with open(collapse_log_file, "w") as f:
    pass
if not os.path.exists(event_memory_file):
    with open(event_memory_file, "w", newline='') as f:
        writer = csv.DictWriter(f, fieldnames=['event_text', 'classification', 'occurrence_count'])
        writer.writeheader()
if not os.path.exists(agent_state_file):
    with open(agent_state_file, "w", newline='') as f:
        writer = csv.writer(f)
        writer.writerow(["step", "mood"] + [f"bias_{s}" for s in CANDIDATE_STATES] + ["intention_state", "intention_strength", "goal", "long_term_plan"])
        writer.writerow([0, 0.0] + [1.0 for _ in CANDIDATE_STATES] + [None, 0, "maximise interesting", None])

print(f"\n---- Proto-Conscious Agent: Recursive Self-Improvement, Voting, Timestamped Logs [{RUN_TIMESTAMP}] ----")
run_simulation(prompt_every=0, max_steps=50)
save_inner_voice()
plot_decisions()
plot_bias_weights()
plot_mood()
gpt_life_summary()
gpt_code_patch_and_vote_and_apply()