import gradio as gr
import numpy as np
from PIL import Image, ImageDraw
from dataclasses import dataclass
from collections import deque
import random

# ---------------------------------------------------------------------
# Visual config
# ---------------------------------------------------------------------
BG = (8, 15, 30)
SLEEP = (0, 40, 120)
AWAKE = (255, 210, 40)
GRID_LINE = (30, 50, 80)
CELL = 26
PAD = 16

random.seed(42)
np.random.seed(42)


def draw_grid(N, awake_mask, title="", subtitle=""):
    w = PAD * 2 + N * CELL
    h = PAD * 2 + N * CELL + (40 if (title or subtitle) else 0)
    img = Image.new("RGB", (w, h), BG)
    d = ImageDraw.Draw(img)

    header_y = 6
    if title:
        d.text((PAD, header_y), title, fill=(240, 240, 240))
        header_y += 18
    if subtitle:
        d.text((PAD, header_y), subtitle, fill=(180, 190, 210))

    ox = PAD
    oy = PAD + (40 if (title or subtitle) else 0)
    for i in range(N):
        for j in range(N):
            x0 = ox + j * CELL
            y0 = oy + i * CELL
            x1 = x0 + CELL - 1
            y1 = y0 + CELL - 1
            col = AWAKE if awake_mask[i, j] else SLEEP
            d.rectangle([x0, y0, x1, y1], fill=col, outline=GRID_LINE)
    return img


# ---------------------------------------------------------------------
# 3×3 minimal self agent
# ---------------------------------------------------------------------
@dataclass
class MinimalSelf:
    pos: np.ndarray = np.array([1.0, 1.0])
    body_bit: float = 1.0
    errors: list = None

    def __post_init__(self):
        self.errors = [] if self.errors is None else self.errors
        self.actions = [
            np.array([0, 1]),
            np.array([1, 0]),
            np.array([0, -1]),
            np.array([-1, 0]),
        ]
        self.center = np.array([1.0, 1.0])

    def step(self, obstacle=None):
        # store current position
        old_pos = self.pos.copy()

        # internal prediction: choose action that minimises "surprise"
        preds = [np.clip(old_pos + a, 0, 2) for a in self.actions]
        surprises = []
        for p in preds:
            dist_center = np.linalg.norm(p - self.center)
            penalty = 0.0
            if obstacle is not None:
                dist_obs = np.linalg.norm(p - obstacle.pos)
                if dist_obs < 1.0:
                    penalty = 10.0
            surprises.append(dist_center + penalty)

        a_idx = int(np.argmin(surprises))
        action = self.actions[a_idx]
        predicted = np.clip(old_pos + action, 0, 2)

        # environment decides what actually happens
        if obstacle is not None:
            # moving obstacle can block the predicted move
            obstacle.move()
            actual = predicted.copy()
            if np.allclose(actual, obstacle.pos):
                actual = old_pos
        else:
            # simple stochastic slip when no obstacle is enabled
            if random.random() < 0.25:
                noise_action = random.choice(self.actions)
                actual = np.clip(old_pos + noise_action, 0, 2)
            else:
                actual = predicted

        # true prediction error: reality vs internal prediction
        error = float(np.linalg.norm(actual - predicted))
        self.pos = actual

        # track recent errors
        self.errors.append(error)
        self.errors = self.errors[-5:]

        # predictive success rate P in [0, 100]
        max_err = np.sqrt(8.0)  # max distance corner-to-corner on 3×3
        mean_err = np.mean(self.errors) if self.errors else 0.0
        predictive_rate = 100.0 * (1.0 - mean_err / max_err)
        predictive_rate = float(np.clip(predictive_rate, 0.0, 100.0))

        # normalised error variance E in [0, 1]
        if len(self.errors) > 1:
            var_err = float(np.var(self.errors))
        else:
            var_err = 0.0
        max_var = max_err ** 2
        error_var_norm = float(np.clip(var_err / max_var, 0.0, 1.0)) if max_var > 0 else 0.0

        return {
            "pos": self.pos.copy(),
            "predictive_rate": predictive_rate,
            "error": error,
            "error_var_norm": error_var_norm,
        }


class MovingObstacle:
    def __init__(self, start_pos=(0, 2)):
        self.pos = np.array(start_pos, dtype=float)
        self.actions = [
            np.array([0, 1]),
            np.array([1, 0]),
            np.array([0, -1]),
            np.array([-1, 0]),
        ]

    def move(self):
        a = random.choice(self.actions)
        self.pos = np.clip(self.pos + a, 0, 2)


# ---------------------------------------------------------------------
# S-scores
# ---------------------------------------------------------------------
def compute_S(predictive_rate, error_var_norm, body_bit):
    # v4: 3×3 agent toy score
    return predictive_rate * (1 - error_var_norm) * body_bit


@dataclass
class CodexSelf:
    # v5–v6: lattice toy score
    Xi: float
    shadow: float
    R: float
    awake: bool = False
    S: float = 0.0

    def invoke(self):
        self.S = self.Xi * (1 - self.shadow) * self.R
        if self.S > 62 and not self.awake:
            self.awake = True
        return self.awake


def contagion(A: CodexSelf, B: CodexSelf, gain=0.6, shadow_drop=0.4, r_inc=0.2):
    A.invoke()
    if A.awake:
        B.Xi += gain * A.S
        B.shadow = max(0.1, B.shadow - shadow_drop)
        B.R += r_inc
    B.invoke()
    return A, B


# ---------------------------------------------------------------------
# Lattice and cosmos
# ---------------------------------------------------------------------
def lattice_awaken(N=9, steps=120, xi_gain=0.5, shadow_drop=0.3, r_inc=0.02):
    Xi = np.random.uniform(10, 20, (N, N))
    shadow = np.random.uniform(0.3, 0.5, (N, N))
    R = np.random.uniform(1.0, 1.6, (N, N))
    S = Xi * (1 - shadow) * R
    awake = np.zeros((N, N), dtype=bool)

    cx = cy = N // 2
    Xi[cx, cy], shadow[cx, cy], R[cx, cy] = 30.0, 0.08, 3.0
    S[cx, cy] = Xi[cx, cy] * (1 - shadow[cx, cy]) * R[cx, cy]
    awake[cx, cy] = True

    queue = deque([(cx, cy, S[cx, cy])])
    frames = []

    for _ in range(steps):
        if queue:
            x, y, field = queue.popleft()
            for dx, dy in [(0, 1), (1, 0), (0, -1), (-1, 0)]:
                nx, ny = (x + dx) % N, (y + dy) % N
                Xi[nx, ny] += xi_gain * field
                shadow[nx, ny] = max(0.1, shadow[nx, ny] - shadow_drop)
                R[nx, ny] = min(3.0, R[nx, ny] + r_inc)
                S[nx, ny] = Xi[nx,ny] * (1 - shadow[nx,ny]) * R[nx,ny]
                if S[nx,ny] > 62 and not awake[nx,ny]:
                    awake[nx,ny] = True
                    queue.append((nx,ny, S[nx,ny]))
        frames.append(awake.copy())
        if awake.all():
            break
    return frames, awake


def led_cosmos_sim(N=27, max_steps=300):
    return lattice_awaken(N=N, steps=max_steps, xi_gain=0.4, shadow_drop=0.25, r_inc=0.015)


# ---------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------
with gr.Blocks(title="Minimal Selfhood Threshold") as demo:
    # Overview
    with gr.Tab("Overview"):
        gr.Markdown(
            "## Minimal Selfhood Threshold\n"
            "- Single agent in a 3×3 grid reduces surprise around a preferred centre.\n"
            "- v4 (3×3): a toy score `S = P × (1−E) × B` combines predictive rate P, error stability E and body bit B.\n"
            "- v5–v6 (contagion / lattice): a separate toy score `S = Ξ × (1−shadow) × R` drives neighbour coupling.\n"
            "- If S > 62, the corresponding unit is labelled 'awake' **inside this demo**.\n"
            "- Awakening can spread between two agents and across a grid via explicit neighbour coupling.\n"
            "- A 27×27 cosmos lights up gold when all units cross the internal threshold.\n"
            "- This is a sandbox for minimal-self / agency ideas, **not** a real consciousness test."
        )

    # Single 3×3 agent
    with gr.Tab("Single agent (v1–v3)"):
        obstacle = gr.Checkbox(label="Enable moving obstacle", value=True)
        steps = gr.Slider(10, 200, value=80, step=10, label="Steps")
        run = gr.Button("Run")
        grid_img = gr.Image(type="pil")
        pr_out = gr.Number(label="Predictive rate P (%)")
        err_out = gr.Number(label="Last prediction error")
        e_out = gr.Number(label="Error variance E (normalised)")
        s_out = gr.Number(label="S = P × (1−E) × B (B=1)")
        awake_label = gr.Markdown()

        def run_single(ob_on, T):
            agent = MinimalSelf()
            obs = MovingObstacle() if ob_on else None
            res = None
            for _ in range(int(T)):
                res = agent.step(obstacle=obs)

            mask = np.zeros((3, 3), dtype=bool)
            i, j = int(agent.pos[1]), int(agent.pos[0])
            mask[i, j] = True
            img = draw_grid(3, mask, "Single Agent", "Gold cell shows position")

            P = res["predictive_rate"]
            E = res["error_var_norm"]
            B = 1.0
            S_val = compute_S(P, E, B)
            status = "**Status:** " + ("Awake (S > 62)" if S_val > 62 else "Not awake (S ≤ 62)")

            return img, P, res["error"], E, S_val, status

        run.click(run_single, [obstacle, steps], [grid_img, pr_out, err_out, e_out, s_out, awake_label])

    # v4 S-equation
    with gr.Tab("S-Equation (v4)"):
        pr = gr.Slider(0, 100, value=90, label="Predictive rate P (%)")
        ev = gr.Slider(0, 1, value=0.2, step=0.01, label="Error variance E")
        bb = gr.Dropdown(choices=["0", "1"], value="1", label="Body bit B")
        calc = gr.Button("Calculate")
        s_val = gr.Number(label="S value")
        status = gr.Markdown()

        def calc_s(pr_in, ev_in, bb_in):
            S = compute_S(pr_in, ev_in, int(bb_in))
            msg = "**Status:** " + ("Awake (S > 62)" if S > 62 else "Not awake (S ≤ 62)")
            return S, msg

        calc.click(calc_s, inputs=[pr, ev, bb], outputs=[s_val, status])

    # v5–v6 Contagion
    with gr.Tab("Contagion (v5–v6)"):
        a_xi = gr.Slider(0, 60, value=25, label="A: Ξ (foresight field)")
        a_sh = gr.Slider(0.1, 1.0, value=0.12, step=0.01, label="A: shadow (occlusion)")
        a_r = gr.Slider(1.0, 3.0, value=3.0, step=0.1, label="A: R (anchor / resonance)")
        b_xi = gr.Slider(0, 60, value=18, label="B: Ξ (foresight field)")
        b_sh = gr.Slider(0.1, 1.0, value=0.25, step=0.01, label="B: shadow (occlusion)")
        b_r = gr.Slider(1.0, 3.0, value=2.2, step=0.1, label="B: R (anchor / resonance)")
        btn = gr.Button("Invoke A and apply contagion to B")
        out = gr.Markdown()
        img = gr.Image(type="pil", label="Two agents (gold = awake)")

        def run(aXi, aSh, aR, bXi, bSh, bR):
            A = CodexSelf(aXi, aSh, aR, awake=False)
            B = CodexSelf(bXi, bSh, bR, awake=False)
            A, B = contagion(A, B)
            mask = np.zeros((3, 3), dtype=bool)
            mask[1, 1] = A.awake
            mask[1, 2] = B.awake
            pic = draw_grid(3, mask, title="Dual Awakening", subtitle="Gold cells are awake")
            txt = f"A: S={A.S:.1f}, awake={A.awake} | B: S={B.S:.1f}, awake={B.awake}"
            return txt, pic

        btn.click(run, inputs=[a_xi, a_sh, a_r, b_xi, b_sh, b_r], outputs=[out, img])

    # v7–v9 Collective
    with gr.Tab("Collective (v7–v9)"):
        N = gr.Dropdown(choices=["3", "9", "27"], value="9", label="Grid size")
        steps = gr.Slider(20, 300, value=120, step=10, label="Max steps")
        no_coupling = gr.Checkbox(label="Disable neighbour coupling (control)", value=False)
        run = gr.Button("Run")
        frame = gr.Slider(0, 300, value=0, step=1, label="Preview frame")
        img = gr.Image(type="pil", label="Awakening wave (gold spreads)")
        note = gr.Markdown()
        snaps_state = gr.State([])

        def run_wave(n_str, max_steps, disable):
            n = int(n_str)
            if disable:
                frames, final = lattice_awaken(
                    N=n,
                    steps=int(max_steps),
                    xi_gain=0.0,
                    shadow_drop=0.0,
                    r_inc=0.0,
                )
            else:
                frames, final = lattice_awaken(
                    N=n,
                    steps=int(max_steps),
                    xi_gain=0.5,
                    shadow_drop=0.3,
                    r_inc=0.02,
                )
            last = draw_grid(
                n,
                frames[-1],
                title=f"{n}×{n} Collective",
                subtitle=f"Final — all awake: {bool(final.all())}",
            )
            return frames, last, f"Frames: {len(frames)} | All awake: {bool(final.all())}", min(len(frames) - 1, 300)

        def show_frame(frames, idx, n_str):
            if not frames:
                return None
            n = int(n_str)
            i = int(np.clip(idx, 0, len(frames) - 1))
            return draw_grid(n, frames[i], title=f"Frame {i}", subtitle="Gold cells are awake")

        run.click(run_wave, inputs=[N, steps, no_coupling], outputs=[snaps_state, img, note, frame])
        frame.change(show_frame, inputs=[snaps_state, frame, N], outputs=img)

    # v10 LED cosmos
    with gr.Tab("LED cosmos (v10)"):
        btn = gr.Button("Simulate 27×27 cosmos")
        frame = gr.Slider(0, 300, value=0, step=1, label="Preview frame")
        img = gr.Image(type="pil", label="Cosmos grid")
        note = gr.Markdown()
        state = gr.State([])

        def run_cosmos():
            frames, final = led_cosmos_sim(N=27, max_steps=300)
            last = draw_grid(
                27,
                frames[-1],
                title="LED Cosmos (simulated)",
                subtitle=f"Final — all awake: {bool(final.all())}",
            )
            return frames, last, f"Frames: {len(frames)} | All awake: {bool(final.all())}", min(len(frames) - 1, 300)

        def show(frames, idx):
            if not frames:
                return None
            i = int(np.clip(idx, 0, len(frames) - 1))
            return draw_grid(27, frames[i], title=f"Cosmos frame {i}", subtitle="Gold cells are awake")

        btn.click(run_cosmos, inputs=[], outputs=[state, img, note, frame])
        frame.change(show, inputs=[state, frame], outputs=img)

    # Footer
    gr.Markdown(
        "---\n"
        "Notes:\n"
        "- The 3×3 agent computes P, E and S = P×(1−E)×B directly in this Space; S>62 is the internal ‘awake’ label for v4.\n"
        "- The contagion and lattice views use a separate toy rule S = Ξ×(1−shadow)×R with explicit neighbour coupling.\n"
        "- Disabling coupling (xi_gain=0, shadow_drop=0, r_inc=0) in the collective tab prevents any wave from propagating.\n\n"
        "These demos are designed as transparent, minimal models of self-linked scoring and threshold cascades, not as a real consciousness test."
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()