import parsetoml
import particle
import linalg
import math
import strformat
import random
import parsexyz

import threadpool
const nThreads = 8

type
    ForceField = ref object
        # this is stretching the definition of an ff beyond what's reasonable
        # should be split in two - the actual ff and calculation parameters
        lennardE*: float
        lennardS*: float
        neighborCutoff*: float
        boundary*: float
        dt*: float
        maxt*: float
        saveInterval*: int
        extraVel*: float # TODO remove this once diff velocity is implemented
        # TODO: new options - thermostat and periodic boundaries
        # cutoff distance for calculations
        # load velocity as diff from last two frames options and adapt genrandom for that option

proc newForceField*(path: string, name: string): ForceField =
    new(result)

    # read config file
    let input = parsetoml.parseString(path.readFile())

    # initialize variables based on input file
    let config = input[name]
    result.dt = config["dt"].getFloat()
    result.maxt = config["t"].getFloat()
    result.lennardE = config["lennardE"].getFloat()
    result.lennardS = config["lennardS"].getFloat()
    result.boundary = config["boundary"].getFloat()
    result.saveInterval = config["saveInterval"].getInt()
    result.extraVel = config["extraVel"].getFloat()
    result.neighborCutoff = config["neighborCutoff"].getFloat()

proc applyExtraVel*(particles: var seq[Particle], ff: ForceField) =
    for i in 0..particles.high():
        particles[i].vel = vector(rand(ff.extraVel * 2) - ff.extraVel, rand(ff.extraVel * 2) - ff.extraVel, rand(ff.extraVel * 2) - ff.extraVel)

proc getAcc(particles: seq[Particle], ff: ForceField, index: int): Vector =
    result = vector0()
    let posThis = particles[index].pos
    let mThis = particles[index].mass

    let s6 = ff.lennardS ^ 6
    let lennard_pre = 24f * ff.lennardE * s6
    let boundary = ff.boundary

    for i in 0..particles.high:
        if i == index:
            continue
        var diff = particles[i].pos - posThis
        # periodic condition
        if diff.x > boundary:
            diff.x -= 2f * boundary
        elif diff.x < -boundary:
            diff.x += 2f * boundary
        if diff.y > boundary:
            diff.y -= 2f * boundary
        elif diff.y < -boundary:
            diff.y += 2f * boundary
        if diff.z > boundary:
            diff.z -= 2f * boundary
        elif diff.z < -boundary:
            diff.z += 2f * boundary

        if diff.len() == 0.0:
            raise newException(Defect, &"overlap between {i} and {index}")
        let distance = diff.len()

        if distance > ff.neighborCutoff * ff.lennardS:
            continue

        let direction = diff / distance

        # Lennard-Jones
        # F = 4e (6s^6/r^7 - 12s^12/r^13)
        let rm1 = 1/distance
        let r6 = rm1 ^ 6
        result -= direction * (lennard_pre * rm1 * (r6) * (1 - 2f * s6 * r6) / mThis)
        
proc keepInside(particles: var seq[Particle], ff: ForceField, index: int) =
    # periodic boundary condition
    let posThis = particles[index].pos
    let velThis = particles[index].vel

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

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

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

    if posThis.y < -ff.boundary and velThis.y < 0:
        particles[index].pos.y = ff.boundary
    
    if posThis.z > ff.boundary and velThis.z > 0:
        particles[index].pos.z = -ff.boundary
    
    if posThis.z < -ff.boundary and velThis.z < 0:
        particles[index].pos.z = ff.boundary

proc iterate*(particles: var seq[Particle], ff: ForceField, t: float) =
    let dt = ff.dt
    block step:
    # iterate through particles
        # update positions
        for i in 0..particles.high:
            let part = particles[i]
            particles[i].pos = part.pos + part.vel * dt + dt * dt * 0.5 * part.acc

        # update acceleration and velocity
        var newAccs = newSeq[Vector](particles.len)
        #for i in 0..nThreads:
            # spawn a thread
        for i in 0..particles.high:
            newAccs[i] = getAcc(particles, ff, i)

        for i in 0..particles.high:
            let part = particles[i]
            let newAcc = newAccs[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(particles, ff, i)
        

proc simulate*(ffName: string, xyzPath: string) =
    let ff = newForceField("ffs.toml", ffName)

    var particles: seq[Particle] = parseXyz(xyzPath)

    applyExtraVel(particles, ff) # TODO replace this

    # open the output file
    let f = open(xyzPath, fmAppend)
    var saveCounter = 0

    var t = 0.0
    while t < ff.maxt:
        iterate(particles, ff, t)
        inc saveCounter
        if saveCounter >= ff.saveInterval:
            saveCounter = 0
            appendXyz(f, particles, t)
        t += ff.dt

    # cleanup
    f.close()