import linalg
import parsetoml

import math
import strformat
import random

# define types
type
    Particle = object
        name: string
        pos: Vector
        vel: Vector
        acc: Vector
        mass: float

# read input file
let input = parsetoml.parseString("input.toml".readFile())

# initialize variables based on input file
let config = input["config"]
var t = 0.0 # current time
let dt = config["dt"].getFloat()
let maxt = config["t"].getFloat()
let gravConst = config["G"].getFloat()
let gravAcc = config["g"].getFloat()
let stabilizeMassCentre = config["stab_mass"].getBool()
let stabInterval = config["stab_interval"].getInt()
var stabCounter = 0
let lennardE = config["lennardE"].getFloat()
let lennardS = config["lennardS"].getFloat()
let boundary = config["boundary"].getFloat()
let saveInterval = config["saveInterval"].getInt()
# a thing that runs every 100 frames to move the centre of mass back to the original one
var particles: seq[Particle] = @[]

# process the rest of the input file
for key, val in input.tableVal.pairs:
    if key == "config":
        continue
    elif key == "random":
        let count = val["count"].getInt()
        let xmin = val["xmin"].getFloat()
        let xmax = val["xmax"].getFloat()
        let ymin = val["ymin"].getFloat()
        let ymax = val["ymax"].getFloat()
        let vmax = val["vmax"].getFloat()
        for i in 0..count:
            block thisPart:
                var particle = Particle(name: "random")
                var pass = false
                var attempt = 0
                while not pass:
                    inc attempt
                    if attempt > 10:
                        break thisPart
                    particle.pos = vector(rand(xmax-xmin)+xmin, rand(ymax-ymin)+ymin)
                    pass = true
                    for p in particles:
                        if (p.pos - particle.pos).len() < lennardS * 2.5:
                            pass = false
                            break
                particle.vel = vector(rand(vmax), rand(vmax))
                particle.acc = vector(0.0, 0.0)
                particle.mass = 1.0
                particles.add(particle)
    else:
        let x: float = val["pos"][0].getFloat()
        let y: float = val["pos"][1].getFloat()
        let vx: float = val["vel"][0].getFloat()
        let vy: float = val["vel"][1].getFloat()
        let pos = vector(x, y)
        let vel = vector(vx, vy)
        let mass = val["mass"].getFloat()
        let particle = Particle(name: key, pos: pos, vel: vel, acc: vector(0,0), mass: mass)
        particles.add(particle)

# open the output file
let f = open("output.txt", fmWrite)
# do the simulation, writing results to a file (streaming)
f.writeLine("mdsim input file: input.toml")
proc oupStatus =
    # output current status
    var status: string = ""
    for i in 0..particles.high:
        status &= &"{particles[i].name} {$particles[i].pos} {$particles[i].vel} "
    f.writeLine(&"{t} {status}")

oupStatus()

proc getForces(index: int): Vector =
    result = vector(0, 0)
    let posThis = particles[index].pos
    let mThis = particles[index].mass
    for i in 0..particles.high:
        if i == index:
            continue
        let diff = particles[i].pos - posThis
        let direction = diff.normalize()
        let distance = diff.len()
        let mThat = particles[i].mass

        # gravity between particles
        #  m2 * C / r^2 (so it's acceleartion, mass of this doesn't apply)
        result += direction * (mThat * gravConst / distance / distance)

        # Lennard-Jones
        # F = 4e (6s^6/r^7 - 12s^12/r^13)
        result -= direction * (4f * lennardE * (6.0 * (lennardS^6)/(distance^7) - 12.0 * (lennardS^12)/(distance^13))) / mThis
        
        result += vector(0, gravAcc)

proc keepInside(index: int) =
    let posThis = particles[index].pos
    let velThis = particles[index].vel

    if posThis.x > boundary and velThis.x > 0:
        particles[index].vel.x = -velThis.x

    if posThis.x < -boundary and velThis.x < 0:
        particles[index].vel.x = -velThis.x

    if posThis.y > boundary and velThis.y > 0:
        particles[index].vel.y = -velThis.y

    if posThis.y < -boundary and velThis.y < 0:
        particles[index].vel.y = -velThis.y


proc getCentreOfMass(): Vector =
    var totMass = 0.0
    result = vector(0.0, 0.0)
    for i in 0..particles.high:
        let part = particles[i]
        totMass += part.mass
        result += part.pos * part.mass
    result = result / totMass

let initCentreOfMass = getCentreOfMass()

proc resetCentreOfMass() =
    let cCentre = getCentreOfMass()
    let delta = cCentre - initCentreOfMass
    for i in 0..particles.high:
        particles[i].pos -= delta

var saveCounter = 0

while t < maxt:
    block step:
        # iterate through particles
        for i in 0..particles.high:
            let part = particles[i]
            particles[i].pos = part.pos + part.vel * dt + dt * dt * 0.5 * part.acc

        for i in 0..particles.high:
            let part = particles[i]
            let newAcc = getForces(i)
            let newVel = part.vel + (part.acc + newAcc) * (dt * 0.5)
            particles[i].vel = newVel
            particles[i].acc = newAcc
        
        # enforce boundary
        for i in 0..particles.high:
            keepInside(i)
    inc saveCounter
    if saveCounter >= saveInterval:
        saveCounter = 0
        oupStatus()
    t += dt
    inc stabCounter
    if stabCounter > stabInterval:
        stabCounter = 0
        if stabilizeMassCentre:
            resetCentreOfMass()

# cleanup
f.close()