import os
import re
import litert_lm
import numpy as np

# 1. Initialize our clean structural path
model_path = os.path.expanduser("~/.litert-lm/models/gemma-4-E2B-it.litertlm/model.litertlm")

# 2. Define the static mathematical engine ("The NumPy Buttons")
def run_vector_field_flux(coords, mu, omega):
    """
    Simulates a continuous XYFLOW coordinate trajectory update.
    The values step deterministically forward along the field gradients.
    """
    x, y = coords
    
    # Primitives based on our dynamic attractor math
    dx = -mu * x + np.sin(y) * 0.5
    dy = -omega * y + np.cos(x) * 0.5
    
    return np.array([x + dx, y + dy])

def parse_parameters(ai_text):
    """
    Scans the AI response block for assignment targets using regular expressions.
    Falls back to safe baseline parameters if parsing fails.
    """
    mu = 1.0
    omega = 1.0
    try:
        mu_match = re.search(r"mu\s*=\s*([\d\.]+)", ai_text, re.IGNORECASE)
        omega_match = re.search(r"omega\s*=\s*([\d\.]+)", ai_text, re.IGNORECASE)
        if mu_match:
            mu = float(mu_match.group(1))
        if omega_match:
            omega = float(omega_match.group(1))
    except Exception:
        pass
    return mu, omega

# 3. Open the Runtime Context Lifecycle
print("[System Initialization] Starting local LiteRT-LM Context...")
with litert_lm.Engine(model_path, backend=litert_lm.Backend.CPU()) as engine:
    
    system_instruction = litert_lm.Message.system(
        "You are an XYFLOW coordinate optimizer designed to keep systems stable. "
        "Analyze the coordinate bounds and output target parameters precisely format as 'mu = X' and 'omega = Y'."
    )
    
    with engine.create_conversation(messages=[system_instruction]) as conversation:
        
        # Define starting state vector
        current_coordinates = np.array([5.0, -3.5])
        
        # Execute 10 optimization steps through time
        for step in range(1, 11):
            print(f"\n--- TIMESTEP INTERVAL {step} ---")
            print(f"Current State Trajectory Coordinates: {current_coordinates}")
            
            # Formulate current state text map for the model
            user_prompt = (
                f"The system coordinates are at x={current_coordinates[0]:.4f}, y={current_coordinates[1]:.4f}. "
                f"Select optimization parameters mu and omega to minimize spatial divergence."
            )
            
            # Send context to the SLM
            response = conversation.send_message(user_prompt)
            ai_raw_text = response["content"][0]["text"]
            
            # Parse parameters chosen by the model
            chosen_mu, chosen_omega = parse_parameters(ai_raw_text)
            print(f"SLM Decision Matrix -> Chosen mu: {chosen_mu}, Chosen omega: {chosen_omega}")
            
            # Let NumPy execute the high-speed floating point integration step
            next_coordinates = run_vector_field_flux(current_coordinates, chosen_mu, chosen_omega)
            
            # Evaluate variation delta (simulated mechanical entropy reduction)
            delta_variance = np.linalg.norm(current_coordinates - next_coordinates)
            print(f"NumPy Execution Success. Trajectory Delta step: {delta_variance:.6f}")
            
            # Transition variables to next cycle step
            current_coordinates = next_coordinates
            
        print("\n[Manifold Collapse Sequence Terminated] 10 optimization intervals complete.")