import math

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

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

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

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

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

func `+`*(x: float64, 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: float64): Dual =
  (x.val - y, x.dot)

func `-`*(x: float64, 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: float64): Dual =
  (x.val * y, x.dot * y)

func `*`*(x: float64, 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: float64): Dual =
  (x.val / y, x.dot / y)

func `/`*(x: float64, 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: float64, gm: float64, q_r: float64, p_r: float64, q_phi: float64, p_phi: float64, h0: float64]

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: float64) =  # 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: float64) = # 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: float64) =
  let x = p.q_r * sin(p.q_phi)
  let y = p.q_r * cos(p.q_phi)
  echo x, " ", y, " ", 0.0, " ", t, " ", h(p, d_dual(p.q_r), d_dual(p.p_r), d_dual(p.p_phi)).val - p.h0

proc solve*(p: var Parameters,
            uq: proc(p: var Parameters, h: float64),
            up: proc(p: var Parameters, h: float64),
            output: proc(p: Parameters, h: float64)) =
  let
    steps: int64 = 10000
    h: float64 = 1.0
  #let args: seq[string] = parseCmdLine()
  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:
    uq(p, h)
    up(p, h)
    output(p, step.float * h)

var p: Parameters = (1.0, 1.0, 12.0, 0.0, 0.0, 0.6 * sqrt(12.0), 0.0)
solve(p, update_q, update_p, plot)