' NeReLaBasic - Three-Body Problem Simulation (Vectorized)
' --------------------------------------------------------
' This program simulates the motion of three bodies interacting
' under their mutual gravitational attraction. It uses initial
' conditions that lead to a stable "figure-8" choreography.
'
' Method: Vectorized calculation of forces and numerical
' integration using the Semi-Implicit Euler method.

'=============================================================================
' 1. FUNCTION TO CALCULATE ACCELERATIONS (VECTORIZED)
'=============================================================================
FUNC CalculateAccelerations(Pos, Mass, G, softening)
    ' This function calculates the acceleration on each body using matrix operations,
    ' completely replacing the need for nested loops.

    ' Extract X and Y position vectors from the main Position matrix.
    PosX = SLICE(Pos, 1, 0) ' Get column 0 (all x-positions)
    PosY = SLICE(Pos, 1, 1) ' Get column 1 (all y-positions)

    ' 1. Calculate matrices of relative positions (rx, ry) for all pairs.
    '    The OUTER(A, B, "-") creates a matrix where M(i,j) = B(j) - A(i).
    rx_matrix = OUTER(PosX, PosX, "-")
    ry_matrix = OUTER(PosY, PosY, "-")

    ' 2. Calculate the matrix of squared distances between all pairs.
    '    The softening factor prevents division by zero if bodies get too close.
    dist_sq_matrix = rx_matrix^2 + ry_matrix^2 + softening

    ' 3. Calculate the matrix of inverse cube distances.
    '    This is a common optimization: F = G*m*r / r^3 instead of F = G*m*r / (r^2 * r)
    inv_dist_cube = dist_sq_matrix ^ -1.5

    ' 4. Calculate the matrix of acceleration magnitudes.
    '    First, create a matrix of influencing masses using OUTER.
    Mass_j_matrix = OUTER(Mass * 0, Mass, "+")
    accel_mag_matrix = G * Mass_j_matrix * inv_dist_cube

    ' 5. Calculate the final acceleration components for X and Y.
    '    This is the acceleration magnitude multiplied by the direction vector (rx, ry).
    AccX_components = accel_mag_matrix * rx_matrix
    AccY_components = accel_mag_matrix * ry_matrix

    ' 6. Sum the forces.
    '    To get the total acceleration on each body, we sum the influences from all
    '    other bodies. This is equivalent to summing along the rows of our component matrices.
    '    The dimension parameter '1' tells SUM to reduce along the rows.
    AccX_total = SUM(AccX_components, 1)
    AccY_total = SUM(AccY_components, 1)

    ' 7. Combine the final X and Y acceleration vectors back into one matrix.
    '    SUM returns a column matrix (e.g., 3x1), but MVLET expects a 1D vector.
    '    We must flatten the result of SUM into a 1D vector before using it.
    AccX_flat = RESHAPE(AccX_total, [NUM_BODIES])
    AccY_flat = RESHAPE(AccY_total, [NUM_BODIES])

    ' Create the final 3x2 acceleration matrix by inserting the flat vectors.
    FinalAcc = RESHAPE([0], [NUM_BODIES, DIMS])
    FinalAcc = MVLET(FinalAcc, 1, 0, AccX_flat)
    FinalAcc = MVLET(FinalAcc, 1, 1, AccY_flat)
    
    RETURN FinalAcc
ENDFUNC

'=============================================================================
' 2. SETUP AND DEFINE SIMULATION CONSTANTS
'=============================================================================
PRINT "--- 1. Setting up Simulation Constants ---"

' --- GRAPHICS CONSTANTS ---
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
VIEW_SCALE = 220   ' Scales the simulation space to fit the screen (pixels per unit)
BODY_RADIUS = 5    ' Radius of the circle for each body in pixels
TRAIL_LENGTH = 50  ' How many past positions to store for the trail effect

G = 1.0           ' Gravitational constant
DT = 0.001        ' Time step for the simulation
NUM_STEPS = 2000  ' Number of steps to simulate
PRINT_EVERY = 200 ' How often to print the state
NUM_BODIES = 3
DIMS = 2          ' 2D simulation (x, y)
SOFTENING = 0.001 ' Softening factor to prevent division by zero

'=============================================================================
' 3. INITIALIZE DATA STRUCTURES AND SET INITIAL CONDITIONS
'=============================================================================
PRINT "--- 2. Initializing Body States (Mass, Position, Velocity) ---"

DIM Trail[3, TRAIL_LENGTH, 2] ' For each body, store X and Y for TRAIL_LENGTH steps
trail_idx = 0 ' The current index in our trail's circular buffer

' Set masses (equal mass system)
Mass = [1, 1, 1]

' Set initial positions for the "figure-8" orbit using direct array assignment.
Pos = RESHAPE([0.97000436, -0.24308753, -0.97000436, 0.24308753, 0.0, 0.0], [3, 2])

' Set initial velocities for the "figure-8" orbit
Vel = RESHAPE([0.466203685, 0.43236573, 0.466203685, 0.43236573, -0.93240737, -0.86473146], [3, 2])

SCREEN SCREEN_WIDTH, SCREEN_HEIGHT, "Three-Body Problem"

'=============================================================================
' 4. THE MAIN SIMULATION LOOP
'=============================================================================
PRINT "--- 3. Starting Simulation ---"
PRINT "Step, Body, PosX, PosY" ' CSV Header for easy plotting

FOR A = 1 TO NUM_STEPS
    ' First, calculate the accelerations on all bodies based on their current positions.
    Acc = CalculateAccelerations(Pos, Mass, G, SOFTENING)

    ' Next, update the velocity and position for each body.
    ' These operations are inherently vectorized.
    Vel = Vel + Acc * DT
    Pos = Pos + Vel * DT

    FOR i = 0 TO NUM_BODIES - 1
        ' Store the new position in the trail buffer
        Trail[i, trail_idx, 0] = Pos[i, 0]
        Trail[i, trail_idx, 1] = Pos[i, 1]
    NEXT i

    trail_idx = (trail_idx + 1) MOD TRAIL_LENGTH

   ' --- DRAWING LOGIC ---
    ' Clear the screen to black
    CLS 0, 0, 0

    ' Draw the trails first, so the bodies appear on top
    FOR i = 0 TO NUM_BODIES - 1
        ' Set trail color based on body index
        IF i = 0 THEN r = 100: g = 20: b = 20 
        IF i = 1 THEN r = 20: g = 100: b = 20 
        IF i = 2 THEN r = 20: g = 20: b = 100 
        
        FOR t = 0 TO TRAIL_LENGTH - 2
            ' Get two consecutive points from the trail buffer
            idx1 = (trail_idx + t) MOD TRAIL_LENGTH
            idx2 = (trail_idx + t + 1) MOD TRAIL_LENGTH

            ' Get simulation coordinates
            sim_x1 = Trail[i, idx1, 0] : sim_y1 = Trail[i, idx1, 1]
            sim_x2 = Trail[i, idx2, 0] : sim_y2 = Trail[i, idx2, 1]
            
            ' Don't draw lines from uninitialized parts of the trail buffer
            IF sim_x1 <> 0 AND sim_x2 <> 0 THEN
                ' Convert simulation coords to screen coords for both points
                screen_x1 = (sim_x1 * VIEW_SCALE) + (SCREEN_WIDTH / 2)
                screen_y1 = (sim_y1 * VIEW_SCALE) + (SCREEN_HEIGHT / 2)
                screen_x2 = (sim_x2 * VIEW_SCALE) + (SCREEN_WIDTH / 2)
                screen_y2 = (sim_y2 * VIEW_SCALE) + (SCREEN_HEIGHT / 2)
                LINE screen_x1, screen_y1, screen_x2, screen_y2, r, g, b
            ENDIF
        NEXT t
    NEXT i


    ' Draw the bodies themselves
    FOR i = 0 TO NUM_BODIES - 1
        ' Get the body's current simulation coordinates
        sim_x = Pos[i, 0]
        sim_y = Pos[i, 1]

        ' Convert simulation coordinates to screen pixel coordinates
        screen_x = (sim_x * VIEW_SCALE) + (SCREEN_WIDTH / 2)
        screen_y = (sim_y * VIEW_SCALE) + (SCREEN_HEIGHT / 2)

        ' Set color based on body index
        IF i = 0 THEN r = 255: g = 50: b = 50 
        IF i = 1 THEN r = 50: g = 255: b = 50 
        IF i = 2 THEN r = 50: g = 50: b = 255 

        ' Draw the body as a circle
        CIRCLE screen_x, screen_y, BODY_RADIUS, r, g, b
    NEXT i

    ' Flip the screen to show what we've drawn
    SCREENFLIP
    SLEEP 1

NEXT A

PRINT
PRINT "--- Simulation Complete ---"
