import streamlit as st
import numpy as np
import random
import json
import math
import time
import pandas as pd
import graphviz
from scipy.sparse import csr_matrix

# -----------------------------------------------------------------------------
# Page Configuration
# -----------------------------------------------------------------------------
st.set_page_config(page_title="NeuroDungeon: Cortical Column Game", layout="wide")

# -----------------------------------------------------------------------------
# Initialize Game & Memory State
# -----------------------------------------------------------------------------
if 'game_state' not in st.session_state:
    st.session_state.game_state = {
        'hex_grid': [[{'type': 'empty', 'emoji': '', 'alive': False} for _ in range(12)] for _ in range(16)],
        'heroes': {},  # {player_id: {x, y, hp, atk, def, xp, level, gear}}
        'monsters': [],
        'loot': [],
        'exit': {'x': 15, 'y': 11},
        'overlord': {'x': 14, 'y': 10, 'hp': 50, 'atk': 5, 'alive': True},
        'players': {},
        'story': [],       # Memory events (episodic/semantic)
        'drawn_cards': 0,
        'turn': 0
    }

# -----------------------------------------------------------------------------
# Card Deck for Narrative Events (Episodic Memory)
# -----------------------------------------------------------------------------
if 'deck' not in st.session_state:
    def create_deck():
        suits = ["Hearts", "Diamonds", "Clubs", "Spades"]
        ranks = list(range(1, 14))
        deck = [(suit, rank) for suit in suits for rank in ranks]
        random.shuffle(deck)
        return deck
    st.session_state.deck = create_deck()

# -----------------------------------------------------------------------------
# Constants and UI Elements for Both Systems
# -----------------------------------------------------------------------------
HEX_SIZE = 40
PLAYER_NAMES = ["SkyWalker", "ForestRanger", "CityBuilder", "MonsterTamer", "RailMaster"]
PLANTS = ["🌱", "🌲", "🌳", "🌴", "🌵"]
LOOT_EMOJIS = ["⚔️", "🛡️", "💰", "🔑"]
TRAPS = ["⚡", "🕳️"]
SUIT_PROPERTIES = {"Hearts": "heroic", "Diamonds": "treacherous", "Clubs": "fierce", "Spades": "enigmatic"}

# -----------------------------------------------------------------------------
# Initialize Hero if Not Already Set
# -----------------------------------------------------------------------------
if 'player_id' not in st.session_state:
    st.session_state.player_id = random.choice(PLAYER_NAMES)
    st.session_state.game_state['heroes'][st.session_state.player_id] = {
        'x': 0, 'y': 0, 'hp': 20, 'atk': 3, 'def': 1, 'xp': 0, 'level': 1, 'gear': []
    }

player_id = st.sidebar.selectbox("Choose Hero", PLAYER_NAMES, index=PLAYER_NAMES.index(st.session_state.player_id))
st.session_state.player_id = player_id
hero = st.session_state.game_state['heroes'].get(player_id, {'x': 0, 'y': 0, 'hp': 20, 'atk': 3, 'def': 1, 'xp': 0, 'level': 1, 'gear': []})

st.sidebar.subheader("🗡️ Hero Stats")
st.sidebar.write(f"HP: {hero.get('hp', 20)} | ATK: {hero.get('atk', 3)} | DEF: {hero.get('def', 1)}")
st.sidebar.write(f"Level: {hero.get('level', 1)} | XP: {hero.get('xp', 0)}")
st.sidebar.write(f"Gear: {', '.join(hero.get('gear', [])) or 'None'}")

st.sidebar.subheader("🏆 Scores")
for p, s in st.session_state.game_state['players'].items():
    st.sidebar.write(f"{p}: {s}")

# -----------------------------------------------------------------------------
# Utility Functions for Game Dynamics & Neural Decay
# -----------------------------------------------------------------------------
def apply_game_of_life(grid):
    new_grid = [[cell.copy() for cell in row] for row in grid]
    for i in range(len(grid)):
        for j in range(len(grid[0])):
            neighbors = count_neighbors(grid, i, j)
            if grid[i][j]['type'] in ['plant', 'trap'] and grid[i][j]['alive']:
                if neighbors < 2 or neighbors > 3:
                    new_grid[i][j]['alive'] = False
                    new_grid[i][j]['type'] = 'empty'
                    new_grid[i][j]['emoji'] = ''
            elif grid[i][j]['type'] == 'empty' and neighbors == 3:
                new_grid[i][j]['type'] = 'plant'
                new_grid[i][j]['emoji'] = random.choice(PLANTS)
                new_grid[i][j]['alive'] = True
    return new_grid

def count_neighbors(grid, x, y):
    count = 0
    for di in [-1, 0, 1]:
        for dj in [-1, 0, 1]:
            if di == 0 and dj == 0:
                continue
            ni, nj = x + di, y + dj
            if 0 <= ni < len(grid) and 0 <= nj < len(grid[0]) and grid[ni][nj]['alive']:
                count += 1
    return count

def reset_game():
    st.session_state.game_state = {
        'hex_grid': [[{'type': 'empty', 'emoji': '', 'alive': False} for _ in range(12)] for _ in range(16)],
        'heroes': {player_id: {'x': 0, 'y': 0, 'hp': 20, 'atk': 3, 'def': 1, 'xp': 0, 'level': 1, 'gear': []}},
        'monsters': [],
        'loot': [],
        'exit': {'x': 15, 'y': 11},
        'overlord': {'x': 14, 'y': 10, 'hp': 50, 'atk': 5, 'alive': True},
        'players': {},
        'story': [],
        'drawn_cards': 0,
        'turn': 0
    }
    for i in range(16):
        for j in range(12):
            if random.random() < 0.1:
                st.session_state.game_state['hex_grid'][i][j] = {'type': 'wall', 'emoji': '🧱', 'alive': False}

# Place random walls in the dungeon initially
for i in range(16):
    for j in range(12):
        if random.random() < 0.1:
            st.session_state.game_state['hex_grid'][i][j] = {'type': 'wall', 'emoji': '🧱', 'alive': False}

# -----------------------------------------------------------------------------
# p5.js Code for the HexCitySaga Game Rendering (with Narrative Integration)
# -----------------------------------------------------------------------------
p5js_code = """
const HEX_SIZE = 40;
const SQRT_3 = Math.sqrt(3);
let hexGrid = [];
let heroes = {};
let monsters = [];
let loot = [];
let exit = {};
let overlord = {};
let playerId;
let story = [];

function setup() {
  createCanvas(640, 480);
  updateFromState();
}

function updateFromState() {
  hexGrid = JSON.parse(document.getElementById('hex_grid').innerHTML);
  heroes = JSON.parse(document.getElementById('heroes').innerHTML);
  monsters = JSON.parse(document.getElementById('monsters').innerHTML);
  loot = JSON.parse(document.getElementById('loot').innerHTML);
  exit = JSON.parse(document.getElementById('exit').innerHTML);
  overlord = JSON.parse(document.getElementById('overlord').innerHTML);
  playerId = document.getElementById('player_id').innerHTML;
  story = JSON.parse(document.getElementById('story').innerHTML);
}

function draw() {
  background(169, 169, 169);
  drawHexGrid();
  drawLoot();
  drawHeroes();
  drawMonsters();
  drawExit();
  drawOverlord();
  drawStory();
}

function drawHexGrid() {
  for (let i = 0; i < hexGrid.length; i++) {
    for (let j = 0; j < hexGrid[i].length; j++) {
      let x = i * HEX_SIZE * 1.5;
      let y = j * HEX_SIZE * SQRT_3 + (i % 2 ? HEX_SIZE * SQRT_3 / 2 : 0);
      fill(hexGrid[i][j].type === 'wall' ? '#555' : hexGrid[i][j].alive ? '#90EE90' : '#A9A9A9');
      stroke(0);
      drawHex(x, y);
      if (hexGrid[i][j].emoji) {
        textSize(20);
        text(hexGrid[i][j].emoji, x - 10, y + 5);
      }
    }
  }
}

function drawHex(x, y) {
  beginShape();
  for (let a = 0; a < 6; a++) {
    let angle = TWO_PI / 6 * a;
    vertex(x + HEX_SIZE * cos(angle), y + HEX_SIZE * sin(angle));
  }
  endShape(CLOSE);
}

function drawHeroes() {
  Object.keys(heroes).forEach(h => {
    let hero = heroes[h];
    let x = hero.x * HEX_SIZE * 1.5;
    let y = hero.y * HEX_SIZE * SQRT_3 + (hero.x % 2 ? HEX_SIZE * SQRT_3 / 2 : 0);
    fill(0, 255, 0);
    ellipse(x, y, HEX_SIZE * 0.8);
    textSize(14);
    text(h.slice(0,2), x - 10, y + 5);
  });
}

function drawMonsters() {
  monsters.forEach(m => {
    let x = m.x * HEX_SIZE * 1.5;
    let y = m.y * HEX_SIZE * SQRT_3 + (m.x % 2 ? HEX_SIZE * SQRT_3 / 2 : 0);
    fill(255, 0, 0);
    ellipse(x, y, HEX_SIZE * 0.8);
    textSize(20);
    text(m.type === 'Godzilla' ? '🦖' : '🤖', x - 10, y + 5);
  });
}

function drawLoot() {
  loot.forEach(l => {
    let x = l.x * HEX_SIZE * 1.5;
    let y = l.y * HEX_SIZE * SQRT_3 + (l.x % 2 ? HEX_SIZE * SQRT_3 / 2 : 0);
    fill(255, 215, 0);
    ellipse(x, y, HEX_SIZE * 0.6);
    textSize(20);
    text(l.type, x - 10, y + 5);
  });
}

function drawExit() {
  let x = exit.x * HEX_SIZE * 1.5;
  let y = exit.y * HEX_SIZE * SQRT_3 + (exit.x % 2 ? HEX_SIZE * SQRT_3 / 2 : 0);
  fill(0, 191, 255);
  ellipse(x, y, HEX_SIZE * 0.8);
  textSize(20);
  text('🚪', x - 10, y + 5);
}

function drawOverlord() {
  if (overlord.alive) {
    let x = overlord.x * HEX_SIZE * 1.5;
    let y = overlord.y * HEX_SIZE * SQRT_3 + (overlord.x % 2 ? HEX_SIZE * SQRT_3 / 2 : 0);
    fill(139, 0, 139);
    ellipse(x, y, HEX_SIZE * 1.2);
    textSize(20);
    text('👑', x - 10, y + 5);
  }
}

function drawStory() {
  fill(0);
  textSize(14);
  for (let i = 0; i < Math.min(story.length, 5); i++) {
    text(story[story.length - 1 - i], 10, 20 + i * 20);
  }
}

function mousePressed() {
  let i = Math.floor(mouseX / (HEX_SIZE * 1.5));
  let j = Math.floor((mouseY - (i % 2 ? HEX_SIZE * SQRT_3 / 2 : 0)) / (HEX_SIZE * SQRT_3));
  if (i >= 0 && i < hexGrid.length && j >= 0 && j < hexGrid[0].length) {
    let hero = heroes[playerId];
    if (hexGrid[i][j].type === 'empty' && Math.abs(hero.x - i) <= 1 && Math.abs(hero.y - j) <= 1) {
      moveHero(i, j);
    }
  }
}

function keyPressed() {
  let hero = heroes[playerId];
  if (key === 'w' || key === 'W') moveHero(hero.x, hero.y - 1);
  if (key === 's' || key === 'S') moveHero(hero.x, hero.y + 1);
  if (key === 'a' || key === 'A') moveHero(hero.x - 1, hero.y);
  if (key === 'd' || key === 'D') moveHero(hero.x + 1, hero.y);
  if (key === 'q' || key === 'Q') moveHero(hero.x - 1, hero.y - 1);
  if (key === 'e' || key === 'E') moveHero(hero.x + 1, hero.y - 1);
}

function moveHero(newX, newY) {
  if (newX < 0 || newX >= hexGrid.length || newY < 0 || newY >= hexGrid[0].length) return;
  let hero = heroes[playerId];
  let cell = hexGrid[newX][newY];
  if (cell.type === 'empty' || cell.type === 'plant') {
    hero.x = newX;
    hero.y = newY;
    handleInteractions();
    updateState();
  }
}

function handleInteractions() {
  let hero = heroes[playerId];
  // Loot Interaction
  let lootIdx = loot.findIndex(l => l.x === hero.x && l.y === hero.y);
  if (lootIdx !== -1) {
    let item = loot[lootIdx].type;
    if (item === '⚔️') hero.atk += 2;
    if (item === '🛡️') hero.def += 2;
    if (item === '💰') hero.xp += 10;
    if (item === '🔑') hero.gear.push('🔑');
    loot.splice(lootIdx, 1);
    addStory('Loot', item);
  }
  // Trap Interaction
  if (hexGrid[hero.x][hero.y].type === 'trap') {
    hero.hp -= Math.floor(random(1, 5));
    addStory('Trap', hexGrid[hero.x][hero.y].emoji);
  }
  // Monster Combat
  let monsterIdx = monsters.findIndex(m => m.x === hero.x && m.y === hero.y);
  if (monsterIdx !== -1) {
    let monster = monsters[monsterIdx];
    let damage = Math.max(0, hero.atk - 1);
    monster.hp -= damage;
    hero.hp -= Math.max(0, monster.atk - hero.def);
    if (monster.hp <= 0) {
      monsters.splice(monsterIdx, 1);
      hero.xp += 10;
      addStory('Monster', monster.type + ' slain');
    }
  }
  // Overlord Combat
  if (overlord.alive && hero.x === overlord.x && hero.y === overlord.y) {
    let damage = Math.max(0, hero.atk - 2);
    overlord.hp -= damage;
    hero.hp -= Math.max(0, overlord.atk - hero.def);
    if (overlord.hp <= 0) {
      overlord.alive = false;
      hero.xp += 50;
      addStory('Monster', 'Overlord defeated');
    }
  }
  // Exit Check
  if (hero.x === exit.x && hero.y === exit.y && hero.gear.includes('🔑')) {
    alert('Victory! You escaped the Hex Dungeon!');
    resetGame();
  }
  // Level Up
  if (hero.xp >= hero.level * 20) {
    hero.level += 1;
    hero.hp += 5;
    hero.atk += 1;
    hero.def += 1;
    addStory('Hero', 'leveled up');
  }
}

function addStory(type, word) {
  let deck = JSON.parse(document.getElementById('deck').innerHTML);
  let cardIndex = st.session_state.drawn_cards % 52;
  let card = deck[cardIndex];
  let suit = card[0];
  let sentence = `A ${SUIT_PROPERTIES[suit]} event: ${type.toLowerCase()} ${word}`;
  story.push(sentence);
  st.session_state.drawn_cards += 1;
}

function updateState() {
  fetch('/update_state', {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({
      hex_grid: hexGrid,
      heroes: heroes,
      monsters: monsters,
      loot: loot,
      exit: exit,
      overlord: overlord,
      story: story,
      player_id: playerId
    })
  });
}
"""

# -----------------------------------------------------------------------------
# Multi-Tab UI: Cross Breeding Cortical Theory and Dungeon Crawl
# -----------------------------------------------------------------------------
tabs = st.tabs(["Cortical Theory", "HexCitySaga", "Memory Integration", "Self Reward NPS", "Extra UI"])

# --- Tab 1: Cortical Column Theory Demonstration ---
with tabs[0]:
    st.header("Cortical Column Theory: Self-Modifying Memory Systems")
    st.markdown("""
    **Theory Overview:**  
    Life’s essential attribute is the capacity for self-modification. This demonstration illustrates:
    
    - **Episodic Memory (E):** Short-term, introspective attention (≈5–10 seconds).  
    - **Semantic Memory (K):** Accumulated knowledge formed by experiences.  
    - **Neural Connectivity:** Modeled via sparse matrices and graph structures representing voting neurons.  
    - **Free Energy & Bonding:** The drive that connects entities—analogous to love and maximal energy states.
    """)
    
    st.subheader("Neural Connectivity Sparse Matrix")
    size = 10
    density = 0.2
    random_matrix = np.random.binomial(1, density, size=(size, size))
    sparse_matrix = csr_matrix(random_matrix)
    st.write("Sparse Matrix Representation:")
    st.write(sparse_matrix)
    st.write("Dense Matrix Representation:")
    st.write(random_matrix)
    
    st.subheader("Concept Graph Visualization")
    graph_source = """
    digraph G {
        "Cortical Column 🧠" -> "Episodic Memory (E) ⏱️" [label="short-term"];
        "Cortical Column 🧠" -> "Semantic Memory (K) 📚" [label="knowledge"];
        "Episodic Memory (E) ⏱️" -> "Introspective Attention 🔍" [label="focus"];
        "Semantic Memory (K) 📚" -> "Free Energy ⚡" [label="agency"];
        "Free Energy ⚡" -> "Love ❤️" [label="bond"];
        "Love ❤️" -> "Humanity 🌍" [label="connection"];
    }
    """
    st.graphviz_chart(graph_source)

# --- Tab 2: HexCitySaga: Dungeon Crawl Game ---
with tabs[1]:
    st.header("HexCitySaga: Dungeon Crawl")
    st.write(f"Hero: **{player_id}**. Use WASD/QE to move. Goal: escape (🚪) with the key (🔑) or defeat the Overlord (👑).")
    game_html = f"""
    <script src="https://cdnjs.cloudflare.com/ajax/libs/p5.js/1.4.2/p5.min.js"></script>
    <div id="hex_grid" style="display:none">{json.dumps(st.session_state.game_state['hex_grid'])}</div>
    <div id="heroes" style="display:none">{json.dumps(st.session_state.game_state['heroes'])}</div>
    <div id="monsters" style="display:none">{json.dumps(st.session_state.game_state['monsters'])}</div>
    <div id="loot" style="display:none">{json.dumps(st.session_state.game_state['loot'])}</div>
    <div id="exit" style="display:none">{json.dumps(st.session_state.game_state['exit'])}</div>
    <div id="overlord" style="display:none">{json.dumps(st.session_state.game_state['overlord'])}</div>
    <div id="player_id" style="display:none">{player_id}</div>
    <div id="story" style="display:none">{json.dumps(st.session_state.game_state['story'])}</div>
    <div id="deck" style="display:none">{json.dumps(st.session_state.deck)}</div>
    <script>{p5js_code}</script>
    """
    st.components.v1.html(game_html, height=500)

# --- Tab 3: Memory Integration ---
with tabs[2]:
    st.header("Memory Integration")
    st.write("Mapping game events as nodes in an episodic/semantic memory network:")
    if st.session_state.game_state['story']:
        nodes = []
        edges = []
        for idx, event in enumerate(st.session_state.game_state['story']):
            node_label = f"Event {idx+1}"
            # Shorten the event text for node labels if needed
            safe_event = event.replace('"', '\'')
            nodes.append(f'"{node_label}" [label="{safe_event}"];')
            if idx > 0:
                edges.append(f'"Event {idx}" -> "Event {idx+1}";')
        graph_def = "digraph MemoryGraph {\n" + "\n".join(nodes + edges) + "\n}"
        st.graphviz_chart(graph_def)
    else:
        st.write("No memory events recorded yet. Engage in the dungeon to generate memories!")

# --- Tab 4: Self Reward Learning NPS ---
with tabs[3]:
    st.header("Self Reward Learning NPS Score")
    nps_score = st.slider("Rate Self Reward Learning (0-10):", 0, 10, 5)
    if nps_score <= 6:
        nps_comment = "Needs Improvement - Refine self-modification algorithms."
    elif nps_score <= 8:
        nps_comment = "Good, but can be better - Fine-tuning required."
    else:
        nps_comment = "Excellent! The system demonstrates robust self-reward learning."
    st.write(f"**NPS Score:** {nps_score} - {nps_comment}")

# --- Tab 5: Extra UI Components ---
with tabs[4]:
    st.header("Extra UI Components")
    with st.expander("More Details"):
        st.write("Additional interactive widgets and data visualizations can be added here.")
    
    col1, col2 = st.columns(2)
    with col1:
        st.write("**Column 1:** Neural Signal Simulation (Line Chart)")
        st.line_chart(np.random.randn(20, 1))
    with col2:
        st.write("**Column 2:** Energy Consumption (Bar Chart)")
        st.bar_chart(np.random.randn(20, 1))
    
    st.subheader("Data Table")
    df = pd.DataFrame({
        "Component": ["Neuron", "Synapse", "Dendrite", "Axon"],
        "Status": ["Active", "Active", "Inactive", "Active"]
    })
    st.dataframe(df)
    
    st.subheader("JSON Configuration")
    sample_json = {
        "Episodic": {"Duration": "5-10 sec", "Type": "Conscious"},
        "Semantic": {"Label": "Cortical Column Theory"}
    }
    st.json(sample_json)

# -----------------------------------------------------------------------------
# Sidebar: Reset Game Button
# -----------------------------------------------------------------------------
if st.sidebar.button("Reset Game"):
    reset_game()