-- SPDX-License-Identifier: ISC
-- Copyright (c) 2020 Camden Dixie O'Brien

module Main (main) where

import Graphics.Gloss
import Graphics.Gloss.Geometry.Angle
import Graphics.Gloss.Data.ViewPort

data PolarCoord = PolarCoord { radius :: Float, angle :: Float }

data State = State { position :: PolarCoord
                   , radialMomentum :: Float
                   , angularMomentum :: Float }

majorMass :: Float
majorMass = 100

minorMass :: Float
minorMass = 1

initialState :: State
initialState = State (PolarCoord 200 0) 0 150

majorRadius :: Float
majorRadius = 20

minorRadius :: Float
minorRadius = 5

framesPerSecond :: Int
framesPerSecond = 120

trailLength :: Int
trailLength = 720

trailColor :: Color
trailColor = greyN 0.7

render :: State -> Picture
render state = pictures [ circleSolid majorRadius
                        , translatePolar (position state) $ circleSolid minorRadius ]

step :: a -> b -> State -> State
step _ _ state = state { position = nextPosition state
                       , radialMomentum = nextRadialMomentum state }

nextPosition :: State -> PolarCoord
nextPosition state = let r = radius . position $ state
                         a = angle . position $ state
                         rDot = radialMomentum state / minorMass
                         aDot = angularMomentum state / (minorMass * r ^ 2)
                     in PolarCoord (r + rDot) (a + aDot)

nextRadialMomentum :: State -> Float
nextRadialMomentum state = let r = radius . position $ state
                               p = radialMomentum state
                               l = angularMomentum state
                               pDot = (l ^ 2 / (2 * minorMass * r) - minorMass * majorMass) / r ^ 2
                           in p + pDot

translatePolar :: PolarCoord -> Picture -> Picture
translatePolar q = rotate (radToDeg $ angle q) . translate (radius q) 0

window :: Display
window = InWindow "Foo" (800, 800) (100, 100)

main :: IO ()
main = simulateTraced trailLength trailColor window white framesPerSecond initialState render step

simulateTraced :: Int
  -> Color
  -> Display
  -> Color
  -> Int
  -> a
  -> (a -> Picture)
  -> (ViewPort -> Float -> a -> a)
  -> IO ()
simulateTraced n col window bg fps init render step
  = simulate window bg fps [init] (tracedRenderer col render) (tracedStepper n step)

tracedRenderer :: Color -> (a -> Picture) -> [a] -> Picture
tracedRenderer col render states = let renderings = map render states
                                   in pictures . reverse $ head renderings : map (color col) (tail renderings)

tracedStepper :: Int -> (a -> b -> c -> c) -> a -> b -> [c] -> [c]
tracedStepper n step vp time states = step vp time (head states) : take (n - 1) states