import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

class HDGL_VM:
    def __init__(self):
        self.D = np.array([1.,2.,3.,4.,5.,6.,7.,8.,1.])
        self.P = np.zeros(8)
        self.P[4:] = [6.8541019662, 11.09016994, 17.94427191, 29.03444654]
        self.void = 0.0
        self.omega = 1.0
        self.phi_phi = 2.6180339887
        self.phi = 1.6180339887
        self.weights = np.ones(8)/8
        self.ip = 0
        self.program = [0,1,2,3,1,0,2,3,1,0]

        # History for live plotting
        self.history_D = []
        self.history_void = []
        self.history_omega = []

        # Set up live plotting
        self.fig, (self.ax1, self.ax2) = plt.subplots(2,1, figsize=(10,6))
        self.lines_D = [self.ax1.plot([], [], label=f"D{i+1}")[0] for i in range(len(self.D))]
        self.line_void, = self.ax2.plot([], [], label="Void")
        self.line_omega, = self.ax2.plot([], [], label="Omega")
        self.ax1.set_title("D Registers Evolution")
        self.ax2.set_title("Void & Omega Evolution")
        self.ax1.set_xlabel("Step")
        self.ax1.set_ylabel("D Value")
        self.ax2.set_xlabel("Step")
        self.ax2.set_ylabel("Value")
        self.ax1.legend()
        self.ax2.legend()
        plt.ion()
        plt.show()

    # -------------------------
    # Core VM step
    # -------------------------
    def step(self, blend_factor=0.05):
        if self.ip >= len(self.program):
            print("End of program.")
            return False
        D_prev = self.D.copy()
        weighted_sum = np.sum(D_prev[:8] * self.weights)
        for i in range(len(self.D)):
            p_val = self.P[i] if i < len(self.P) else 0
            phi_val = self.phi ** i
            self.D[i] = D_prev[i] + blend_factor * (phi_val * D_prev[i] + self.phi_phi * p_val + weighted_sum + self.omega)
        self.void += np.mean(self.D) * blend_factor
        self.omega += 0.01 * blend_factor
        self.ip += 1

        # update live history
        self.history_D.append(self.D.copy())
        self.history_void.append(self.void)
        self.history_omega.append(self.omega)

        # update live plot
        self.update_plot()
        return True

    # -------------------------
    # Run N steps
    # -------------------------
    def run(self, steps=None):
        count = 0
        while steps is None or count < steps:
            cont = self.step()
            if not cont:
                break
            count += 1

    # -------------------------
    # Reset VM
    # -------------------------
    def reset(self):
        self.__init__()
        print("VM reset.")

    # -------------------------
    # Show current state
    # -------------------------
    def show_state(self):
        print(f"D registers: {np.round(self.D,5)}")
        print(f"Void: {self.void:.5f}, Omega: {self.omega:.5f}")

    # -------------------------
    # Update live plot
    # -------------------------
    def update_plot(self):
        for i, line in enumerate(self.lines_D):
            line.set_data(range(len(self.history_D)), [h[i] for h in self.history_D])
        self.line_void.set_data(range(len(self.history_void)), self.history_void)
        self.line_omega.set_data(range(len(self.history_omega)), self.history_omega)
        self.ax1.relim(); self.ax1.autoscale_view()
        self.ax2.relim(); self.ax2.autoscale_view()
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()

# -------------------------
# Command prompt / REPL
# -------------------------
def vm_repl(vm):
    print("Welcome to HDGL Analog VM. Type 'help' for commands.")
    while True:
        cmd = input("HDGL> ").strip()
        if not cmd:
            continue
        cmd_lower = cmd.lower()

        # Exit
        if cmd_lower in ['quit','exit']:
            break

        # Help
        elif cmd_lower == 'help':
            print("Commands:")
            print(" step [n]     - execute n steps (default 1)")
            print(" run [n]      - run continuously or n steps")
            print(" reset        - reset VM")
            print(" state        - show current D, Void, Omega")
            print(" set Dn=X     - set register n to X (e.g. set D1=5)")
            print(" add/sub/mul/div target source - arithmetic")
            print(" phi/phi_phi/inc/dec Dn          - apply operator to register")
            print(" swap Dn Dm                       - swap two registers")
            print(" exit / quit  - exit prompt")
            print(" help         - show this help")
        # Step
        elif cmd_lower.startswith('step'):
            parts = cmd.split()
            n = int(parts[1]) if len(parts) > 1 else 1
            vm.run(steps=n)
            vm.show_state()
        # Run
        elif cmd_lower.startswith('run'):
            parts = cmd.split()
            n = int(parts[1]) if len(parts) > 1 else None
            vm.run(steps=n)
            vm.show_state()
        # Reset
        elif cmd_lower=='reset':
            vm.reset()
        # State
        elif cmd_lower=='state':
            vm.show_state()
        # Set register
        elif cmd_lower.startswith('set'):
            try:
                _, assign = cmd.split()
                reg, val = assign.split('=')
                reg = reg.strip().upper()
                val = float(val.strip())
                if reg.startswith('D'):
                    idx = int(reg[1:])-1
                    vm.D[idx] = val
                    print(f"{reg} set to {val}")
                else:
                    print("Unknown register.")
            except:
                print("Invalid set command.")
        # Operators
        elif cmd_lower.split()[0] in ['add','sub','mul','div','phi','phi_phi','inc','dec','swap']:
            parts = cmd.split()
            op = parts[0]
            if op in ['add','sub','mul','div']:
                target, source = parts[1], parts[2]
                target_val = vm.D[int(target[1:])-1] if target.upper().startswith('D') else getattr(vm, target.lower())
                # source can be Dn, float, or VOID/OMEGA
                try:
                    source_val = vm.D[int(source[1:])-1] if source.upper().startswith('D') else float(source)
                except:
                    source_val = getattr(vm, source.lower())
                if op=='add': result = target_val + source_val
                elif op=='sub': result = target_val - source_val
                elif op=='mul': result = target_val * source_val
                elif op=='div': result = target_val / source_val
                if target.upper().startswith('D'):
                    vm.D[int(target[1:])-1] = result
                else:
                    setattr(vm, target.lower(), result)
                print(f"{target} = {result}")
                vm.update_plot()
            elif op in ['phi','phi_phi','inc','dec']:
                target = parts[1]
                idx = int(target[1:])-1
                if op=='phi': vm.D[idx] *= vm.phi
                elif op=='phi_phi': vm.D[idx] *= vm.phi_phi
                elif op=='inc': vm.D[idx] += 1
                elif op=='dec': vm.D[idx] -= 1
                print(f"{target} = {vm.D[idx]}")
                vm.update_plot()
            elif op=='swap':
                a, b = parts[1], parts[2]
                idx_a = int(a[1:])-1
                idx_b = int(b[1:])-1
                vm.D[idx_a], vm.D[idx_b] = vm.D[idx_b], vm.D[idx_a]
                print(f"{a} <-> {b}")
                vm.update_plot()
        else:
            print("Unknown command. Type 'help'.")

# -------------------------
# Main
# -------------------------
if __name__ == "__main__":
    vm = HDGL_VM()
    vm_repl(vm)
