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

module Main where

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

data State = State { theta :: Float
                   , pTheta :: Float
                   , phi :: Float
                   , pPhi :: Float
                   }

-- Simulation parameters
particleMass = 1
rodLength = 200
initialState = State (pi / 4) 0 (pi / 6) 0
g = 0.1

-- Rendering parameters
particleRadius = 5
particleColor = black
rodColor = greyN 0.5
framesPerSecond = 100

render state = doublePendulum (theta state) (phi state)
doublePendulum theta phi = (translatePolar rodLength theta $ pendulum phi) <> pendulum theta
translatePolar radius angle = translate (radius * sin angle) (negate $ radius * cos angle)
pendulum angle = rotateRadians angle $ rod <> (translate 0 (negate rodLength) $ particle)
rod = color rodColor $ line [ (0, 0), (0, negate rodLength) ]
particle = color particleColor $ circleSolid particleRadius
rotateRadians = rotate . radToDeg  . negate

step _ _ = nextState
nextState state = State (theta state + thetaDot state)
                        (pTheta state + pThetaDot state)
                        (phi state + phiDot state)
                        (pPhi state + pPhiDot state)
thetaDot state
  = (6 * (2 * pTheta state - 3 * cos (theta state - phi state) * pPhi state))
    / (particleMass * rodLength ^ 2 * (16 - 9 * cos (theta state - phi state) ^ 2))
phiDot state
  = (6 * (8 * pPhi state - 3 * cos (theta state - phi state) * pTheta state))
    / (particleMass * rodLength ^ 2 * (16 - 9 * cos (theta state - phi state) ^ 2))
pThetaDot state
  = (- 1 / 2) * particleMass * rodLength ^ 2
    * (thetaDot state * phiDot state * sin (theta state - phi state) + 3 * g * sin (theta state) / rodLength)
pPhiDot state
  = (1 / 2) * particleMass * rodLength ^ 2
    * (thetaDot state * phiDot state * sin (theta state - phi state) - g * sin (phi state) / rodLength)

window = InWindow "Double Pendulum" (800, 800) (100, 100)
main = simulate window white framesPerSecond initialState render step