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

module Main (main) where

import Graphics.Gloss
import Graphics.Gloss.Geometry.Angle

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

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

majorMass :: Float
majorMass = 200

minorMass :: Float
minorMass = 20

initialState :: State
initialState = State (PolarCoord 300 0) 0 0.001

framesPerSecond :: Int
framesPerSecond = 120

render :: State -> Picture
render state = pictures [ circleSolid (sqrt majorMass)
                        , translatePolar (position state) $ circleSolid (sqrt minorMass) ]

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 = simulate window white framesPerSecond initialState render step