barneshut/main.go

package main

import (
	// "fmt"
	"fmt"
	"math"
	"math/rand"
	"strconv"
	rl "github.com/gen2brain/raylib-go/raylib"
)

type magnet struct {
    pos rl.Vector2
    speed rl.Vector2
    color rl.Color
    mass float32
    radius float32
}


type quadtree struct {
    // pos, size rl.Vector2
    pos, size rl.Vector2
    magnet *magnet
    parent *quadtree
    nodes []*quadtree
    mass float32
    debug string
}


func (q *quadtree) isOutside(magnet rl.Vector2) bool {
    return magnet.X < q.pos.X || magnet.X >= q.pos.X + q.size.X ||
    magnet.Y < q.pos.Y || magnet.Y >= q.pos.Y + q.size.Y
}

func (q *quadtree) isOccupied() (bool, *magnet) {
    if q.magnet == nil {
        return false, nil
    } else {
        return true, q.magnet
    }
}

func (q *quadtree) hasChildNodes() bool {
    return q.nodes != nil
}

func (q *quadtree) insertMagnet(magnets ...*magnet) {
    for _, magnet := range magnets {
        if magnet.pos.X == 156 {
            fmt.Println("foo!")
        }

        for _, node := range q.nodes {
            if node.isOutside(magnet.pos) {
                continue
            }
            if node.hasChildNodes() {
                node.insertMagnet(magnet)
                continue
            }
            if occupied, occupyingMagnet := node.isOccupied(); occupied {
                node.subdivide()
                node.magnet = nil
                node.insertMagnet(magnet)
                node.insertMagnet(occupyingMagnet)
                break
            }
            node.magnet = magnet
            break
        }
    }
}

func newQuadTreeRoot() *quadtree {
    temp := &quadtree{
        pos: rl.NewVector2(0, 0),
        size: rl.NewVector2(windowWidth, windowHeight),
        nodes: nil,
        debug: "root",
    }
    temp.subdivide()
    return temp
}

func (q *quadtree) subdivide() {
    halfX := q.size.X/2 
    halfY := q.size.Y/2 
    halfSize := rl.NewVector2(halfX, halfY)

    q.nodes = []*quadtree{
        {
            pos: q.pos,
            size: halfSize,
            parent: q,
            magnet: nil,
            debug: "topleft",
        },
        {
            pos: rl.NewVector2(q.pos.X + halfX, q.pos.Y),
            size: halfSize,
            parent: q,
            magnet: nil,
            debug: "topright",
        },
        {
            pos: rl.NewVector2(q.pos.X, q.pos.Y + halfY),
            size: halfSize,
            parent: q,
            magnet: nil,
            debug: "bottomleft",
        },
        {
            pos: rl.NewVector2(q.pos.X + halfX, q.pos.Y + halfY),
            size: halfSize,
            parent: q,
            magnet: nil,
            debug: "bottomright",
        },
    }
}

func (q *quadtree) drawBoundaries() {
    if q.nodes == nil {
        return
    }
    for _, node := range q.nodes {
        rl.DrawRectangleLines(
            int32(math.Round(float64(node.pos.X))), 
            int32(math.Round(float64(node.pos.Y))), 
            int32(math.Round(float64(node.size.X))), 
            int32(math.Round(float64(node.size.Y))),
            rl.Green,
        )
    }
    for _, node := range q.nodes {
        node.drawBoundaries()
    }
}

// parâmetros
const initalMagnetCount = 50
const windowWidth = 700
const windowHeight = 700
const magnetRadius = 5 
const gravConst float32 = 10

func tick(magnets []*magnet) {

    for _, a := range magnets {
        for _, b := range magnets {
            if a == b { continue }
            // não creio que esteja correto porém da pro gasto
            direction := rl.Vector2Normalize(rl.Vector2Subtract(a.pos, b.pos))
            distance := rl.Vector2Distance(a.pos, b.pos)
            distanceSquared := math.Pow(float64(distance), 2)
            force := gravConst / float32(distanceSquared)
            acceleration := rl.Vector2Scale(direction, force)

            if distance < magnetRadius*2 {
                continue
            } 
            // f = m * a 
            forceA := rl.Vector2Scale(acceleration, a.mass)
            forceB := rl.Vector2Scale(acceleration, b.mass)
            a.speed = rl.Vector2Add(a.speed, rl.Vector2Negate(forceB))
            b.speed = rl.Vector2Add(b.speed, forceA)

        }
    }

    for _, magnet := range magnets {
        magnet.pos = rl.Vector2Add(magnet.pos, magnet.speed)
    }
}

func (m *magnet) update(magnets []*magnet) {
    magnetsPull := rl.Vector2{}

    for _, magnet := range magnets {
        if m == magnet { continue }

        direction := rl.Vector2Normalize(rl.Vector2Subtract(m.pos, magnet.pos))
        distance := rl.Vector2Distance(magnet.pos, m.pos)
        if (distance < magnetRadius*2 + 10) {
            direction = rl.Vector2Negate(direction)
            // continue
        }
        distanceSquared := math.Pow(float64(distance), 2)
        force := gravConst / float32(distanceSquared)

        acceleration := rl.Vector2Scale(direction, -force)

        magnetsPull = rl.Vector2Add(magnetsPull, acceleration)
    }

    m.speed = rl.Vector2Add(m.speed, magnetsPull)
    newPos := rl.Vector2Add(m.speed, m.pos)

    // x, y := insideArea(newPos)    
    // if !x { newPos.X = -newPos.X; m.speed.X = 0 }
    // if !y { newPos.Y = -newPos.Y; m.speed.Y = 0 }

    m.pos = newPos
}

// debug
func (q *quadtree) parents() []*quadtree {
    if q.parent == nil {
        return nil
    }
    return append(q.parent.parents(), q.parent)
}

// debug
func (q *quadtree) printMagnets() {
    for _, node := range q.nodes {
        node.printMagnets()
        if occupied, _ := node.isOccupied(); occupied {
            fmt.Println(node.magnet.pos, node.parents())
        }
    }
}

func newMagnet(x, y, mass float32) *magnet {
    newMagnet := magnet{
        pos: rl.Vector2{X: x, Y: y},
        color: rl.NewColor(
            uint8(rand.Int()),
            uint8(rand.Int()),
            uint8(rand.Int()),
            255,
        ),
        mass: mass,
        radius: mass,
    }
    return &newMagnet
}


func randomMagnets(n int) []*magnet {
    var magnets []*magnet

    for i := 0; i < n; i++ {
        x := float32(rand.Intn(windowWidth - 100) + 50)
        y := float32(rand.Intn(windowHeight - 100) + 50)
        magnets = append(magnets, newMagnet(x, y, 1))
    }

    return magnets
}

func main() {
	rl.InitWindow(
        windowWidth, windowHeight, 
        "barnes hut",
    )
	defer rl.CloseWindow()
	rl.SetTargetFPS(60)

    magnets := randomMagnets(initalMagnetCount)

    quadtree := newQuadTreeRoot()
    quadtree.insertMagnet(magnets...)

    pause := true

	for !rl.WindowShouldClose() {

		rl.BeginDrawing()
		rl.ClearBackground(rl.Black)
        rl.DrawFPS(0, 0)
        rl.DrawText(strconv.Itoa(len(magnets)), 0, 20, 20, rl.White)
        if pause {
            rl.DrawText("pausado (E)", 0, 40, 20, rl.White)
            quadtree.drawBoundaries()
        }

        {
            if rl.IsKeyPressed(rl.KeyR) {
                magnets = randomMagnets(initalMagnetCount)
                quadtree = newQuadTreeRoot()
                quadtree.insertMagnet(magnets...)
            }
            if rl.IsKeyPressed(rl.KeyW) {
                magnets = append(magnets, randomMagnets(100)...)
            }
            if rl.IsKeyPressed(rl.KeyE) {
                pause = !pause
            }
            if rl.IsKeyPressed(rl.KeySpace) {
                mouse := rl.GetMousePosition()
                magnets = append(magnets, newMagnet(mouse.X, mouse.Y, 100))
            }
        }

        if !pause {
            tick(magnets)
        }


        for _, magnet := range magnets {
            rl.DrawCircleV(magnet.pos, 1, rl.White)
            // rl.DrawCircleV(magnet.pos, magnet.radius, magnet.color)
        }

        rl.EndDrawing()
	}

    fmt.Println("hello world")

}