import os
import math
import std/strutils

type Dual* = tuple[val: float, dot: float]

func d_dual*(x: float): Dual =
  (x, 0.0)

func d_var*(x: float): Dual =
  (x, 1.0)

func `+`*(x, y: Dual): Dual =
  (x.val + y.val, x.dot + y.dot)

func `+`*(x: Dual, y: float): Dual =
  (x.val + y, x.dot)

func `+`*(x: float, y: Dual): Dual =
  (x + y.val, y.dot)

func `-`*(x, y: Dual): Dual =
  (x.val - y.val, x.dot - y.dot)

func `-`*(x: Dual, y: float): Dual =
  (x.val - y, x.dot)

func `-`*(x: float, y: Dual): Dual =
  (x - y.val, - y.dot)

func `*`*(x, y: Dual): Dual =
  (x.val * y.val, y.val * y.dot + x.dot * y.val)

func `*`*(x: Dual, y: float): Dual =
  (x.val * y, x.dot * y)

func `*`*(x: float, y: Dual): Dual =
  (x * y.val, x * y.dot)

func `/`*(x, y: Dual): Dual =
  (x.val / y.val, (x.dot * y.val - x.val * y.dot) / (y.val * y.val))

func `/`*(x: Dual, y: float): Dual =
  (x.val / y, x.dot / y)

func `/`*(x: float, y: Dual): Dual =
  (x / y.val, - y.dot * x / (y.val * y.val))

func sin*(x: Dual): Dual =
  (sin(x.val), x.dot * cos(x.val))

func cos*(x: Dual): Dual =
  (cos(x.val), x.dot * - sin(x.val))

type Parameters = tuple[m: float, gm: float, q_r: float, p_r: float, q_phi: float, p_phi: float, h0: float]

func h(p: Parameters, qr: Dual, pr: Dual, pth: Dual): Dual =  # the Hamiltonian
  (pr * pr + pth * pth / (qr * qr)) / (2.0 * p.m) - p.gm / qr

proc update_q(p: var Parameters, c: float) =  # coordinate updates, dq/dt = dH / dp
  let qr = c * h(p, d_dual(p.q_r), d_var(p.p_r), d_dual(p.p_phi)).dot
  let qphi = c * h(p, d_dual(p.q_r), d_dual(p.p_r), d_var(p.p_phi)).dot
  p.q_r = p.qr + qr
  p.q_phi = p.qphi + qphi

proc update_p(p: var Parameters, d: float) = # momentum updates, dp/dt = - dH / dq
  p.pr = p.pr - d * h(p, d_var(p.q_r), d_dual(p.p_r), d_dual(p.p_phi)).dot

proc plot(p: Parameters, t: float) =
  let x = p.q_r * sin(p.q_phi)
  let y = p.q_r * cos(p.q_phi)
  var diff = h(p, d_dual(p.q_r), d_dual(p.p_r), d_dual(p.p_phi)).val - p.h0
  var e = if diff > 0.0 : diff else: - diff
  var error = 10.0 * log10(if e >= 1e-36 : e else: 1e-36)
  echo x, " ", y, " ", 0.0, " ", t, " ", error

type
  Coefficients = seq[float]

proc euler_cromer*(p: var Parameters, cd: Coefficients,
                   uq: proc(p: var Parameters, h: float),
                   up: proc(p: var Parameters, h: float)) =
    uq(p, cd[0])
    up(p, cd[1])

proc symplectic*(p: var Parameters, cd: Coefficients,
                     uq: proc(p: var Parameters, h: float),
                     up: proc(p: var Parameters, h: float)) =
  for i in countup(0, cd.len - 1):
    if (i mod 2) == 0: uq(p, cd[i]) else: up(p, cd[i])
  for i in countdown(cd.len - 2, 0):
    if (i mod 2) == 0: uq(p, cd[i]) else: up(p, cd[i])

type
  ArgumentError = object of ValueError

proc solve*(args: seq[string], p: var Parameters,
            uq: proc(p: var Parameters, c: float),
            up: proc(p: var Parameters, d: float),
            output: proc(p: Parameters, t: float)) =
  let
    order: int64 = parseInt(args[1])
    h: float = parseFloat(args[2])
    steps: int64 = parseInt(args[3])
    z1: float = 1.0 / (4.0 - pow(4.0, (1.0 / 3.0)))
    y1: float = 1.0 / (4.0 - pow(4.0, (1.0 / 5.0)))
    x1: float = 1.0 / (4.0 - pow(4.0, (1.0 / 7.0)))
    z0: float = 1.0 - 4.0 * z1
    y0: float = 1.0 - 4.0 * y1
    x0: float = 1.0 - 4.0 * y1
  var
    cd: Coefficients
    integrator: proc(p: var Parameters, cd: Coefficients,
                     uq: proc(p: var Parameters, h: float),
                     up: proc(p: var Parameters, h: float))
  case order
  of 1:
    cd = @[h, h]
    integrator = euler_cromer
  of 2:  # Stormer-Verlet
    cd = @[0.5 * h, h]
    integrator = symplectic
  of -4:  # Forest-Ruth
    let z1: float = 1.0 / (2.0 - pow(2.0, (1.0 / 3.0)))
    let z0: float = 1.0 - 2.0 * z1
    cd = @[0.5 * h * z1, h * z1, 0.5 * h * (z1 + z0), h * z0]
    integrator = symplectic
  of 4:  # m4r35n357_4
    cd = @[0.5 * h * z1, h * z1, h * z1, h * z1, 0.5 * h * (z1 + z0), h * z0]
    integrator = symplectic
  of 6:  # m4r35n357_6
    cd = @[0.5 * h * z1 * y1, h * z1 * y1, h * z1 * y1, h * z1 * y1, 0.5 * h * (z1 + z0) * y1, h * z0 * y1, 0.5 * h * (z0 + z1) * y1, h * z1 * y1, h * z1 * y1, h * z1 * y1, h * z1 * y1, h * z1 * y1, h * z1 * y1, h * z1 * y1, 0.5 * h * (z1 + z0) * y1, h * z0 * y1, 0.5 * h * (z0 + z1) * y1, h * z1 * y1, h * z1 * y1, h * z1 * y1, 0.5 * h * z1 * (y1 + y0), h * z1 * y0, h * z1 * y0, h * z1 * y0, 0.5 * h * (z1 + z0) * y0, h * z0 * y0]
    integrator = symplectic
  of 8:  # m4r35n357_8
    cd = @[0.5 * h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, 0.5 * h * (z1 + z0) * y1 * x1, h * z0 * y1 * x1, 0.5 * h * (z0 + z1) * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, 0.5 * h * (z1 + z0) * y1 * x1, h * z0 * y1 * x1, 0.5 * h * (z0 + z1) * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, 0.5 * h * z1 * (y1 + y0) * x1, h * z1 * y0 * x1, h * z1 * y0 * x1, h * z1 * y0 * x1, 0.5 * h * (z1 + z0) * y0 * x1, h * z0 * y0 * x1, 0.5 * h * (z0 + z1) * y0 * x1, h * z1 * y0 * x1, h * z1 * y0 * x1, h * z1 * y0 * x1, 0.5 * h * z1 * (y1 + y0) * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, 0.5 * h * (z1 + z0) * y1 * x1, h * z0 * y1 * x1, 0.5 * h * (z0 + z1) * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, 0.5 * h * (z1 + z0) * y1 * x1, h * z0 * y1 * x1, 0.5 * h * (z0 + z1) * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, 0.5 * h * (z1 + z0) * y1 * x1, h * z0 * y1 * x1, 0.5 * h * (z0 + z1) * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, 0.5 * h * (z1 + z0) * y1 * x1, h * z0 * y1 * x1, 0.5 * h * (z0 + z1) * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, 0.5 * h * z1 * (y1 + y0) * x1, h * z1 * y0 * x1, h * z1 * y0 * x1, h * z1 * y0 * x1, 0.5 * h * (z1 + z0) * y0 * x1, h * z0 * y0 * x1, 0.5 * h * (z0 + z1) * y0 * x1, h * z1 * y0 * x1, h * z1 * y0 * x1, h * z1 * y0 * x1, 0.5 * h * z1 * (y1 + y0) * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, 0.5 * h * (z1 + z0) * y1 * x1, h * z0 * y1 * x1, 0.5 * h * (z0 + z1) * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, 0.5 * h * (z1 + z0) * y1 * x1, h * z0 * y1 * x1, 0.5 * h * (z0 + z1) * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, h * z1 * y1 * x1, 0.5 * h * z1 * y1 * (x1 + x0), h * z1 * y1 * x0, h * z1 * y1 * x0, h * z1 * y1 * x0, 0.5 * h * (z1 + z0) * y1 * x0, h * z0 * y1 * x0, 0.5 * h * (z0 + z1) * y1 * x0, h * z1 * y1 * x0, h * z1 * y1 * x0, h * z1 * y1 * x0, h * z1 * y1 * x0, h * z1 * y1 * x0, h * z1 * y1 * x0, h * z1 * y1 * x0, 0.5 * h * (z1 + z0) * y1 * x0, h * z0 * y1 * x0, 0.5 * h * (z0 + z1) * y1 * x0, h * z1 * y1 * x0, h * z1 * y1 * x0, h * z1 * y1 * x0, 0.5 * h * z1 * (y1 + y0) * x0, h * z1 * y0 * x0, h * z1 * y0 * x0, h * z1 * y0 * x0, 0.5 * h * (z1 + z0) * y0 * x0, h * z0 * y0 * x0]
    integrator = symplectic
  else:
    echo "invalid integrator ", order, ", should be 1, 2, 4, 6, 8 or -4"
    raise newException(ArgumentError, "invalid integrator")
  p.h0 = h(p, d_dual(p.q_r), d_dual(p.p_r), d_dual(p.p_phi)).val
  for step in 1..steps:
    integrator(p, cd, uq, up)
    output(p, step.float * h)

# ./h_newton 6 2 1 10000  1 12 .6
let
  args: seq[string] = commandLineParams()
  mass = parseFloat(args[4])
  r0 = parseFloat(args[5])
  lfac = parseFloat(args[6])
var p: Parameters = (mass, mass, r0, 0.0, 0.0, lfac * sqrt(r0), 0.0)
solve(args, p, update_q, update_p, plot)