import os
import math
import std/strutils

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 euler_cromer*(p: var Parameters, h: float64,
                   uq: proc(p: var Parameters, h: float64),
                   up: proc(p: var Parameters, h: float64)) =
    uq(p, h)
    up(p, h)

proc stormer_verlet*(p: var Parameters, h: float64,
                     uq: proc(p: var Parameters, h: float64),
                     up: proc(p: var Parameters, h: float64)) =
    uq(p, h * 0.5)
    up(p, h)
    uq(p, h * 0.5)

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
    args: seq[string] = commandLineParams()
    order: int64 = parseInt(args[1])
    h: float64 = parseFloat(args[2])
    steps: int64 = parseInt(args[3])
  var
    integrator: proc(p: var Parameters, h: float64, uq: proc(p: var Parameters, h: float64), up: proc(p: var Parameters, h: float64))
  case order
  of 1:
    integrator = euler_cromer
  of 2:
    integrator = stormer_verlet
  else:
    echo "invalid integrator ", order, " should be 1 or 2"
  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, h, 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(p, update_q, update_p, plot)