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* = object
    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) =
  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
  echo p.q_r * sin(p.q_phi), " ", p.q_r * cos(p.q_phi), " ", 0.0, " ", t, " ", 10.0 * log10(if e >= 1e-36 : e else: 1e-36)

type
  Coefficients = seq[float]
  ArgumentError = object of ValueError

proc symplectic[T](p: var T, cd: Coefficients, uq, up)
  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])

proc solve*[T](args: seq[string], p: var T, uq, up, output: proc(p: var 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
    cd: Coefficients = case order
    of 2:  # Stormer-Verlet
      @[0.5 * h, h]
    of 4:  # m4r35n357_4
      @[0.5 * h * z1, h * z1, h * z1, h * z1, 0.5 * h * (z1 + z0), h * z0]
    of 6:  # m4r35n357_6
      @[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]
    of 8:  # m4r35n357_8
      @[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]
    else:
      echo "invalid order value ", order, ", should be 2, 4, 6, or 8"
      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:
    symplectic(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(m: mass, gm: mass, q_r: r0, p_r: 0.0, q_phi: 0.0, p_phi: lfac * sqrt(r0), h0: 0.0)
solve(args, p, update_q, update_p, plot)