import Control.Monad.ST (ST, runST)
import Data.Foldable (for_)
import Data.List (nub, partition)
import Data.STRef (STRef)
import GHC.Arr (STArray, elems, freezeSTArray, newSTArray, writeSTArray)

-- Parametricity, polynomials and ST

-- newtype State s a = State { runState :: s -> (s, a) }

increment :: Int -> Int
increment = (+ 1)

increment' :: (Num a) => a -> a
increment' = (+ 1)

whatIsThis :: forall a. a -> a
whatIsThis x = x -- ?

-- parametricity
-- Wadler: Theorems for free

-- f :: forall a. a -> a
-- g :: b -> c
-- g . f = f . g

map' :: forall a. (a -> a) -> [a] -> [a]
map' = map

-- g :: [a] -> [a]
-- map' f . g = g . map' f

f :: (Num a) => a -> a
f x = x + 3 * x * x + x * x * x

newtype Poly = Poly {runPoly :: forall a. (Num a) => a -> a}

instance Num Poly where
  Poly p + Poly q = Poly (\x -> p x + q x)
  Poly p * Poly q = Poly (\x -> p x * q x)
  abs (Poly p) = Poly (abs . p)
  signum (Poly p) = Poly (signum . p)
  negate (Poly p) = Poly (negate . p)
  fromInteger c = Poly (const (fromInteger c))

eval :: (Num a) => Poly -> a -> a
eval (Poly p) = p

subst :: Poly -> Poly -> Poly
subst = eval

newtype PolyN i = PolyN {runPolyN :: forall a. (Num a) => (i -> a) -> a}

instance Num (PolyN i)

eval' :: (Num a) => PolyN i -> (i -> a) -> a
eval' (PolyN f) = f

subst' :: PolyN i -> (i -> PolyN j) -> PolyN j
subst' = eval'

newtype Variables i = Variables {getVariables :: [i]}

instance (Eq i) => Num (Variables i) where
  Variables xs + Variables ys = Variables $ nub (xs <> ys)
  Variables xs * Variables ys = Variables $ nub (xs <> ys)
  abs (Variables xs) = Variables xs
  signum (Variables xs) = Variables xs
  fromInteger _ = Variables []
  negate (Variables xs) = Variables xs

variables :: (Eq i) => PolyN i -> [i]
variables p = getVariables $ eval' p (Variables . (: []))

-- data ST s a

runST' :: (forall s. ST s a) -> a
runST' = runST

-- what if a is a type of mutable variables?

-- STRef s a

-- bad :: ST s (STRef s' Int)
-- bad = _

-- leak :: forall s. ST s Int
-- leak = runST bad

qsort :: (Ord a) => [a] -> [a]
qsort [] = []
qsort (x : xs) =
  let (ys, zs) = partition (< x) xs
   in qsort ys ++ x : qsort zs

sort :: (Ord a) => [a] -> [a]
sort xs = runST $ do
  let range = (0, length xs - 1)
  arr <- newSTArray range (head xs)
  for_ (zip [0 ..] xs) $ \(i, x) -> writeSTArray arr i x
  sortStep arr range
  arr' <- freezeSTArray arr
  return (elems arr')

sortStep :: (Ord a) => STArray s Int a -> (Int, Int) -> ST s ()
sortStep arr (l, r) = do
  let m = l + (r - l) `div` 2
  sortStep arr (l, m - 1)
  sortStep arr (m, r)
  merge arr (l, m, r)

merge :: Ord a => STArray s Int a -> (Int, Int, Int) -> ST s ()
merge = _

newtype Gen s = Gen { genState :: STRef s Int }

nextInt :: Gen s -> ST s Int
nextInt = error "TODO"

-- functor parametricity
-- t = t * t | t + t | t -> t | mu x. t(x) | x
-- fmap :: (a -> b) -> (f a, g a) -> (f b, g b)
-- fmap :: (a -> b) -> Either (f a) (g a) -> Either (f b) (g b)
-- fmap :: (a -> b) -> Fix (f a) -> Fix (f b)
-- fmap :: (a -> b) -> a -> b
-- fmap :: (a -> b) -> ((f a -> h) -> g a) -> ((f a -> h) -> g b)