#!/usr/bin/env cabal
{- cabal:
ghc-options:
  -O2
  -threaded
  -with-rtsopts=-N
build-depends:
  base,
  parallel
-}
{- project:
index-state: 2023-04-17T00:00:00Z
-}

{-# LANGUAGE BangPatterns #-}

import           Control.Parallel.Strategies
import           Data.List                   (foldl')
import           Numeric                     (showFFloat)

getX (x, _, _) = x; getY (_, y, _) = y; getZ (_, _, z) = z

lorenz :: Num c => (c, c, c) -> c -> (c, c, c) -> (c, c, c)
lorenz (sigma, rho, beta) dt (x, y, z) =
  ( x + sigma * (y - x) * dt
  , y + (x * (rho - z) - y) * dt
  , z + (x * y - beta * z) * dt
  )

countFlips :: (Foldable t, Eq a1, Num a2, Num a1) => (a3 -> a1) -> t a3 -> a2
countFlips f ps = fst $
  foldl' (\(n, s) -> \p ->
             let s' = signum (f p)
             in (if s' == s then n else n + 1, s'))
  (0, 0) ps

average :: (Fractional a, Foldable t) => t a -> a
average xs = (sum xs) / fromIntegral(length xs)

showFullFloat :: Double -> ShowS
showFullFloat = Numeric.showFFloat Nothing

main = do
  let params = (10.0, 28, 8.0 / 3.0)
  let dt = 0.001
  let n = 10000000
  let is = [0..100]
  let ps = map -- parMap rseq
        (\i -> let p0 = ((i / 1000.0), 1.0, 1.0)
               in (p0, take n $ iterate (lorenz params dt) p0)) is
  let rs = map -- parMap (rpar `dot` rdeepseq)
        (\(p0, ps) -> (getX p0, countFlips getY ps)) ps
  -- let rs = (map (\(p0, ps) -> let f = getY in (f p0, average (map f ps))) ps)
  mapM (\(p,v) -> putStrLn (showFullFloat p . ("\t" ++) . showFullFloat v $ ""))
    (rs `using` parBuffer 100 rseq)