more refactoring in Fit.hs, fitw

more refactoring in Fit.hs, fitw

module Data.Cov
    where

import Prelude
import Data.Array
  ( replicate, unsafeIndex, zipWith, length, foldl, range, take
  ) as A
import Control.Monad.Eff ( forE )
import Control.Monad.ST ( pureST )
import Data.Array.ST (emptySTArray, peekSTArray, pokeSTArray
                     , pushAllSTArray, unsafeFreeze)
import Data.Foldable ( sum )
import Partial.Unsafe ( unsafePartial )
import Data.Maybe ( Maybe (..) )
import Control.MonadZero (guard)
import Data.Int ( toNumber, floor )
import Math ( abs, sqrt )
import Unsafe.Coerce ( unsafeCoerce ) as Unsafe.Coerce

import Data.SimpleMatrix
  ( Matrix
  , transpose
  , fromArray, fromArray2, toArray
  ) as M
import Stuff

newtype Dim3 = DDim3 Int
newtype Dim4 = DDim4 Int
newtype Dim5 = DDim5 Int
class DDim a where
  ddim :: a -> Int
instance ddim3 :: DDim Dim3 where
  ddim _ = 3
instance ddim4 :: DDim Dim4 where
  ddim _ = 4
instance ddim5 :: DDim Dim5 where
  ddim _ = 5
instance ddima :: DDim a where
  ddim _ = undefined

{-- mapTo :: forall a. DDim a => Cov a -> Dim5 --}
{-- mapTo = Unsafe.Coerce.unsafeCoerce --}
{-- class Dim a where --}
{--   dim :: a -> Int --}
{-- instance dima :: Dim (Cov a) where --}
{--   dim ccc = n where --}
{--     xx = mapTo ccc --}
{--     n = ddim xx --}
{-- instance dim3 :: Dim (Cov Dim3) where --}
{--   dim ccc = 3 --}
{-- instance dim4 :: Dim (Cov Dim4) where --}
{--   dim ccc = 4 --}
{-- instance dim5 :: Dim (Cov Dim5) where --}
{--   dim ccc = 5 --}

newtype Cov a = Cov { v :: Array Number }
newtype Jac a b = Jac { v :: Array Number }
newtype Vec a = Vec { v :: Array Number }
type Cov3 = Cov Dim3
type Cov4 = Cov Dim4
type Cov5 = Cov Dim5
{-- type Jac43 = Jac Dim4 Dim3 --}
type Jac53 = Jac Dim5 Dim3
type Jac33 = Jac Dim3 Dim3
type Jac34 = Jac Dim3 Dim4
type Jac35 = Jac Dim3 Dim5
type Jac44 = Jac Dim4 Dim4
type Jac55 = Jac Dim5 Dim5
type Vec3 = Vec Dim3
type Vec4 = Vec Dim4
type Vec5 = Vec Dim5
type Jacs = {aa :: Jac53, bb :: Jac53, h0 :: Vec5}

-- access to arrays of symmetrical matrices
indV :: Int -> Int -> Int -> Int
indV w i0 j0 = (i0*w+j0)
indVs :: Int -> Int -> Int -> Int
indVs w i0 j0 | i0 <= j0  = (i0*w - (i0*(i0-1)) `div` 2 + j0-i0)
              | otherwise = (j0*w - (j0*(j0-1)) `div` 2 + i0-j0)
-------------------------------------------------------------------------
-------------------------------------------------------------------------
-------------------------------------------------------------------------
-------------------------------------------------------------------------
-- Mat to give behavior to Cov and Vec and Jac
-- ability to convert to and from Matrix and Array
-- while keeping info about dimensionality
-- also define Semiring and Ring functions
--

class Mat a where
  val :: a -> Array Number
  fromArray :: Array Number -> a
  toArray :: a -> Array Number
instance matCova :: Mat (Cov a) where
  val (Cov {v}) = v
  fromArray a = c' where
    l = A.length a
    c' = case l of
              6   -> Cov {v: a}
              10  -> Cov {v: a}
              15  -> Cov {v: a}
              _   -> Cov {v: let
                            n = floor <<< sqrt <<< toNumber $ l
                            iv = indV n
                            in do -- only upper triangle
                              i0 <- A.range 0 (n-1)
                              j0 <- A.range i0 (n-1)
                              pure $ uidx a (iv i0 j0) }
  toArray c@(Cov {v}) = v' where
          l = A.length v
          n = case l of -- ceil ((sqrt(8.0 * (toNumber l) + 1.0) - 1.0)/2.0)
            6  -> 3
            10 -> 4
            15 -> 5
            _  -> error $ "matCova: toArray not supported for lenght "
                        <> show l
          iv = indVs n
          v' = do
                  i0 <- A.range 0 (n-1)
                  j0 <- A.range 0 (n-1)
                  pure $ uidx v (iv i0 j0)
instance matVeca :: Mat (Vec a) where
  val (Vec {v}) = v
  fromArray a = Vec {v: a}
  toArray (Vec {v}) = v
instance matJacab :: Mat (Jac a b) where
  val (Jac {v}) = v
  fromArray a = Jac {v: a}
  toArray (Jac {v}) = v

class Mat1 a where
  toMatrix :: a -> M.Matrix
instance mat1Cova :: Mat1 (Cov a) where
  toMatrix a@(Cov {v}) = case A.length v of
                            6  -> M.fromArray2 3 3 v
                            10 -> M.fromArray2 4 4 v
                            15 -> M.fromArray2 5 5 v
                            _ -> error $ "mat1Cova toMatrix "
                                          <> show (A.length v)
instance mat1Veca :: Mat1 (Vec a) where
  toMatrix (Vec {v}) = M.fromArray (A.length v) v
instance mat1Jac53 :: Mat1 (Jac Dim5 Dim3) where
  toMatrix (Jac {v}) = M.fromArray2 5 3 v
instance mat1Jac35 :: Mat1 (Jac Dim3 Dim5) where
  toMatrix (Jac {v}) = M.fromArray2 3 5 v
instance mat1Jacaa :: Mat1 (Jac a b) where
  toMatrix j@(Jac {v}) = case A.length v of
                              9  -> M.fromArray2 3 3 v
                              16 -> M.fromArray2 4 4 v
                              25 -> M.fromArray2 5 5 v
                              12 -> M.fromArray2 3 4 v `debug` "this should not have happened ??????????????????? 4 3"
                              15 -> M.fromArray2 5 3 v `debug` "this should not have happened ??????????????????? 5 3"
                              _  -> error $ "mat1Jacaa toMatrix "
                                          <> show (A.length v)

-----------------------------------------------------------------
-- | funcitons for symetric matrices: Cov
-- | type class SymMat
class SymMat a where
  inv :: Cov a -> Cov a                -- | inverse matrix
  invMaybe :: Cov a -> Maybe (Cov a)   -- | Maybe inverse matrix
  det :: Cov a -> Number               -- | determinant
  diag :: Cov a -> Array Number        -- | Array of diagonal elements
  {-- chol :: Cov a -> Jac a a       -- | Cholsky decomposition --}
instance symMatCov3 :: SymMat Dim3 where
--  inv m | trace ( "inv " <> (show m) ) False = undefined
  inv m = m' where
    mm = invMaybe m
    m' = case mm of
           Nothing -> error $ "can't invert 3x3 matrix " <> show m
           Just x -> x
  invMaybe (Cov {v}) = _inv v where
    _inv :: Array Number -> Maybe (Cov Dim3)
    _inv = unsafePartial $ \[a11,a12,a13,a22,a23,a33] -> do
      let det = (a33*a12*a12 - 2.0*a13*a23*a12 + a13*a13*a22
                + a11*(a23*a23 - a22*a33))
      guard $ (abs det) > 1.0e-50
      let
          b11 = (a23*a23 - a22*a33)/det
          b12 = (a12*a33 - a13*a23)/det
          b13 = (a13*a22 - a12*a23)/det
          b22 = (a13*a13 - a11*a33)/det
          b23 = (a11*a23 - a12*a13)/det
          b33 = (a12*a12 - a11*a22)/det
          v' = [b11,b12,b13,b22,b23,b33]
      pure $ Cov {v: v'}
  det (Cov {v}) = _det v where
    _det :: Array Number -> Number
    _det = unsafePartial $ \[a,b,c,d,e,f] ->
            a*d*f - a*e*e - b*b*f + 2.0*b*c*e - c*c*d
  det (Cov {v}) = dd where
        a = unsafePartial $ A.unsafeIndex v 0
        b = unsafePartial $ A.unsafeIndex v 1
        c = unsafePartial $ A.unsafeIndex v 2
        d = unsafePartial $ A.unsafeIndex v 3
        e = unsafePartial $ A.unsafeIndex v 4
        f = unsafePartial $ A.unsafeIndex v 5
        dd = a*d*f - a*e*e - b*b*f + 2.0*b*c*e - c*c*d
  diag (Cov {v}) = a where
    a11 = unsafePartial $ A.unsafeIndex v 0
    a22 = unsafePartial $ A.unsafeIndex v 3
    a33 = unsafePartial $ A.unsafeIndex v 5
    a = [a11,a22,a33]
instance symMatCov4 :: SymMat Dim4 where
  inv m = uJust (invMaybe m)
  invMaybe (Cov {v}) = _inv v where
    _inv :: Array Number -> Maybe (Cov Dim4)
    _inv = unsafePartial $ \[a,b,c,d,e,f,g,h,i,j] -> do
      let det = (a*e*h*j - a*e*i*i - a*f*f*j + 2.0*a*f*g*i - a*g*g*h
            - b*b*h*j + b*b*i*i - 2.0*d*(b*f*i - b*g*h - c*e*i + c*f*g)
            + b*c*(2.0*f*j - 2.0*g*i) + c*c*(g*g - e*j) + d*d*(f*f - e*h))
      guard $ (abs det) > 1.0e-50
      let a' = (-j*f*f + 2.0*g*i*f - e*i*i - g*g*h + e*h*j)/det
          b' = (b*i*i - d*f*i - c*g*i + d*g*h + c*f*j - b*h*j)/det
          c' = (c*g*g - d*f*g - b*i*g + d*e*i - c*e*j + b*f*j)/det
          d' = (d*f*f - c*g*f - b*i*f - d*e*h + b*g*h + c*e*i)/det
          e' = (-j*c*c + 2.0*d*i*c - a*i*i - d*d*h + a*h*j)/det
          f' = (f*d*d - c*g*d - b*i*d + a*g*i + b*c*j - a*f*j)/det
          g' = (g*c*c - d*f*c - b*i*c + b*d*h - a*g*h + a*f*i)/det
          h' = (-j*b*b + 2.0*d*g*b - a*g*g - d*d*e + a*e*j)/det
          i' = (i*b*b - d*f*b - c*g*b + c*d*e + a*f*g - a*e*i)/det
          j' = (-h*b*b + 2.0*c*f*b - a*f*f - c*c*e + a*e*h)/det
      pure $ fromArray [a',b',c',d',e',f',g',h',i',j']
  det (Cov {v}) = _det v where
    _det = unsafePartial $ \[a,b,c,d,e,f,g,h,i,j] ->
        (a*e*h*j - a*e*i*i - a*f*f*j + 2.0*a*f*g*i - a*g*g*h
          - b*b*h*j + b*b*i*i - 2.0*d*(b*f*i - b*g*h - c*e*i + c*f*g)
          + b*c*(2.0*f*j - 2.0*g*i) + c*c*(g*g - e*j) + d*d*(f*f - e*h))
  diag (Cov {v}) = _diag v where
    _diag :: Array Number -> Array Number
    _diag = unsafePartial $ \[a11,_,_,_,a22,_,_,a33,_,a44] -> [a11,a22,a33,a44]
instance symMatCov5 :: SymMat Dim5 where
  inv m = cholInv m 5
  invMaybe m = Just (cholInv m 5)
  det (Cov {v}) = _det v where
    _det [a,b,c,d,e,f,g,h,i,j,k,l,m,n,o] =
      a*f*j*m*o - a*f*j*n*n - a*f*k*k*o + 2.0*a*f*k*l*n - a*f*l*l*m
      - a*g*g*m*o + a*g*g*n*n + 2.0*a*g*h*k*o - 2.0*a*g*h*l*n - 2.0*a*g*i*k*n
      + 2.0*a*g*i*l*m - a*h*h*j*o + a*h*h*l*l + 2.0*a*h*i*j*n - 2.0*a*h*i*k*l
      - a*i*i*j*m + a*i*i*k*k - b*b*j*m*o + b*b*j*n*n + b*b*k*k*o
      - 2.0*b*b*k*l*n + b*b*l*l*m + 2.0*b*c*g*m*o - 2.0*b*c*g*n*n - 2.0*b*c*h*k*o
      + 2.0*b*c*h*l*n + 2.0*b*c*i*k*n - 2.0*b*c*i*l*m - 2.0*b*d*g*k*o
      + 2.0*b*d*g*l*n + 2.0*b*d*h*j*o - 2.0*b*d*h*l*l - 2.0*b*d*i*j*n
      + 2.0*b*d*i*k*l + 2.0*b*e*g*k*n - 2.0*b*e*g*l*m - 2.0*b*e*h*j*n
      + 2.0*b*e*h*k*l + 2.0*b*e*i*j*m - 2.0*b*e*i*k*k - c*c*f*m*o + c*c*f*n*n
      + c*c*h*h*o - 2.0*c*c*h*i*n + c*c*i*i*m + 2.0*c*d*f*k*o - 2.0*c*d*f*l*n
      - 2.0*c*d*g*h*o + 2.0*c*d*g*i*n + 2.0*c*d*h*i*l - 2.0*c*d*i*i*k
      - 2.0*c*e*f*k*n + 2.0*c*e*f*l*m + 2.0*c*e*g*h*n - 2.0*c*e*g*i*m
      - 2.0*c*e*h*h*l + 2.0*c*e*h*i*k - d*d*f*j*o + d*d*f*l*l + d*d*g*g*o
      - 2.0*d*d*g*i*l + d*d*i*i*j + 2.0*d*e*f*j*n - 2.0*d*e*f*k*l - 2.0*d*e*g*g*n
      + 2.0*d*e*g*h*l + 2.0*d*e*g*i*k - 2.0*d*e*h*i*j - e*e*f*j*m + e*e*f*k*k
      + e*e*g*g*m - 2.0*e*e*g*h*k + e*e*h*h*j
    _det _ = undefined
  diag (Cov {v}) = a where
    a11 = unsafePartial $ A.unsafeIndex v 0
    a22 = unsafePartial $ A.unsafeIndex v 5
    a33 = unsafePartial $ A.unsafeIndex v 9
    a44 = unsafePartial $ A.unsafeIndex v 12
    a55 = unsafePartial $ A.unsafeIndex v 14
    a = [a11,a22,a33,a44,a55]

class MulMat a b c | a b -> c where
  mulm :: a -> b -> c
infixr 7 mulm as *.
instance mulMata :: MulMat (Cov a) (Cov a) (Jac a a) where
  mulm c1 c2 = j' where
    mc1 = toMatrix c1
    mc2 = toMatrix c2
    mj' = mc1 * mc2
    j' = fromArray $ M.toArray mj'
instance mulMatJC :: MulMat (Jac a b) (Cov b) (Jac a b) where
  mulm j@(Jac {v: va}) c@(Cov {v: vb}) = Jac {v: v'} where
    nb = case A.length vb of
              6  -> 3
              10 -> 4
              15 -> 5
              _  -> error $ "mulMatJC wrong length of Cov v "
                            <> show (A.length vb)
    na = (A.length va)/nb
    n = nb
    idx i0 j0 | i0 <= j0  = (i0*n - i0*(i0-1)/2 + j0-i0)
              | otherwise = (j0*n - j0*(j0-1)/2 + i0-j0)
    v' = do
      i0 <- A.range 0 (na-1)
      j0 <- A.range 0 (nb-1)
      pure $ sum do
        k0 <- A.range 0 (nb-1)
        pure $ (uidx va (i0*nb+k0)) * (uidx vb (idx k0 j0))
instance mulMatCJ :: MulMat (Cov a) (Jac a b) (Jac a b) where
  mulm c@(Cov {v: va}) j@(Jac {v: vb}) = Jac {v: v'} where
    na = case A.length va of
              6  -> 3
              10 -> 4
              15 -> 5
              _  -> error $ "mulMatCJ wrong length of Cov v "
                            <> show (A.length va)
    nb = (A.length vb)/na
    n = na
    idx i0 j0 | i0 <= j0  = (i0*n - i0*(i0-1)/2 + j0-i0)
              | otherwise = (j0*n - j0*(j0-1)/2 + i0-j0)
    v' = do
      i0 <- A.range 0 (na-1)
      j0 <- A.range 0 (nb-1)
      pure $ sum do
        k0 <- A.range 0 (na-1)
        pure $ (uidx va (idx i0 k0)) * (uidx vb (k0*nb+j0))
instance mulMatJV :: MulMat (Jac a b) (Vec b) (Vec a) where
  mulm j@(Jac {v: va}) v@(Vec {v: vb}) = Vec {v: v'} where
    nb = A.length vb
    na = (A.length va)/nb
    v' = do
      i0 <- A.range 0 (na-1)
      let j0 = 0
      pure $ sum do
        k0 <- A.range 0 (nb-1)
        pure $ (uidx va (i0*nb+k0)) * (uidx vb (k0))
-- ???????????this needs to be generalized to (Jac a b)
instance mulMatJJ :: MulMat (Jac Dim3 Dim5) (Jac Dim5 Dim3) (Jac Dim3 Dim3) where
  mulm j1 j2 = j' where
    mj1 = toMatrix j1
    mj2 = toMatrix j2
    mj' = mj1 * mj2
    j' = fromArray $ M.toArray mj'
instance mulMatCV :: MulMat (Cov a) (Vec a) (Vec a) where
  mulm c v = v' where
    mc = toMatrix c
    mv = toMatrix v
    mv' = mc * mv
    v' = fromArray $ M.toArray mv'
instance mulMatVV :: MulMat (Vec a) (Vec a) Number where
  mulm (Vec {v:v1}) (Vec {v:v2}) = A.foldl (+) zero $ A.zipWith (*) v1 v2
class TrMat a b | a -> b where
  tr :: a -> b
instance trMatC :: TrMat (Cov a) (Cov a) where
  tr c = c
instance trMatJ :: TrMat (Jac a b) (Jac b a) where
  tr j@(Jac {v}) = Jac {v:v'} where
    l = A.length v
    na = case l of
              9 -> 3
              15 -> 5
              16 -> 4
              25 -> 5
              _  -> error $ "trMatJ: sorry, can't do anything but 5x3 and square "
                            <> show (A.length v)
    nb = l/na
    v' = do
      i0 <- A.range 0 (nb-1)
      j0 <- A.range 0 (na-1)
      pure $ (uidx v (j0*nb+i0))
class SW a b c | a b -> c where
  sw :: a -> b -> c
infixl 7 sw as .*.
instance swVec :: SW (Vec a) (Cov a) Number where
  sw v c = n where
    mv = toMatrix v
    mc = toMatrix c
    mc' = M.transpose mv * mc * mv
    n = uidx (M.toArray mc') 0
instance swCov :: SW (Cov a) (Cov a) (Cov a) where
  sw c1 c2 = c' where
    j' = c1 *. c2 *. c1
    c' = fromArray $ toArray j'
instance swJac :: SW (Jac a b) (Cov a) (Cov b) where
  sw j@(Jac {v: va}) c@(Cov {v: vb}) = Cov {v: v'} where
    {-- mj = toMatrix j --}
    {-- mc = toMatrix c --}
    {-- mc' = M.transpose mj * mc * mj --}

    l = A.length vb
    n = case l of
              6  -> 3
              10 -> 4
              15 -> 5
              _  -> error $ "swJac: don'w know how to " <> show l
    m = (A.length va)/n -- > mxn * nxn * nxm -> mxm
    -- vaT * vb * va
    {-- vat= do --}
    {--   i0 <- A.range 0 (m-1) --}
    {--   j0 <- A.range 0 (n-1) --}
    {--   pure $ (uidx va (j0*m+i0)) --}
    {-- mjt = M.fromArray2 m n vat --}
    {-- mjt''' = M.transpose mj --}

    idx i0 j0 | i0 <= j0  = (i0*n - i0*(i0-1)/2 + j0-i0)
              | otherwise = (j0*n - j0*(j0-1)/2 + i0-j0)
   {-- vb''' = do --}
    {--   i0 <- A.range 0 (n-1) --}
    {--   j0 <- A.range 0 (n-1) --}
    {--   pure $ (uidx vb (idx i0 j0)) --}
    {-- mb''' = M.fromArray2 n n vb''' --}

    {-- va''' = do --}
    {--   i0 <- A.range 0 (n-1) --}
    {--   j0 <- A.range 0 (m-1) --}
    {--   pure $ (uidx va (i0*m+j0)) --}
    {-- mj''' = M.fromArray2 n m va''' --}

    vint = do
      i0 <- A.range 0 (n-1)
      j0 <- A.range 0 (m-1)
      pure $ sum do
        k0 <- A.range 0 (n-1)
        pure $ (uidx vb (idx i0 k0)) * (uidx va (k0*m+j0))
    vf = do
      i0 <- A.range 0 (m-1)
      j0 <- A.range i0 (m-1)
      pure $ sum do
        k0 <- A.range 0 (n-1)
        pure $ (uidx va (k0*m+i0) ) * (uidx vint (k0*m+j0))

    {-- vf = M.toArray mc' --}
    lf = A.length vf
    v' = case lf of
              6   -> vf
              10  -> vf
              15  -> vf
              _   -> error $ "swJac: wrong result " <> show lf

    {-- mv' = M.fromArray2 m m v' --}
    {-- xx = n `debug` ("--jT---->\n" <> show mjt''' <> show mjt) --}
    {-- yx = n `debug` ("--c----->\n" <> show mc <> show mb''') --}
    {-- xy = n `debug` ("--j----->\n" <> show mj <> show mj''') --}
    {-- yy = n `debug` ("-jT*c*j->\n" <> show mc' <> show mv') --}
testCov2 :: String
testCov2 = s where
  s = "Test Cov 2----------------------------------------------\n"
    <> "Vec *. Vec = " <> show (v3 *. v3) <> "\n"
    <> "Cov *. Cov = " <> show ((one::Cov3) *. inv (one::Cov3)) <> "\n"
    <> "Vec + Vec = " <> show (v5 + v5) <> "\n"
    <> "chol Cov = " <> show (chol (one::Cov5)) <> "\n"
    <> "Vec .*. Cov = " <> show (v5 .*. inv (one::Cov5)) <> "\n"
  v3 = fromArray [1.0,1.0,1.0] :: Vec3
  v5 = fromArray [1.0,1.0,1.0,1.0,1.0] :: Vec5

instance semiringCov3 :: Semiring (Cov Dim3) where
  add (Cov {v: v1}) (Cov {v: v2}) = Cov {v: A.zipWith (+) v1 v2}
  zero = Cov {v: A.replicate 6 0.0 }
  mul (Cov {v: a}) (Cov {v: b}) = error "------------> mul cov3 * cov3 not allowed"
  one = Cov { v: [1.0, 0.0, 0.0, 1.0, 0.0, 1.0] }
instance ringCov3 :: Ring (Cov Dim3) where
  sub (Cov {v: v1}) (Cov {v: v2}) = Cov {v: A.zipWith (-) v1 v2}
instance showCov3 :: Show (Cov Dim3) where
  show c = "Show (Cov Dim3) \n" <> (show $ toMatrix c)

instance semiringCov4 :: Semiring (Cov Dim4) where
  add (Cov {v: v1}) (Cov {v: v2}) = Cov {v: A.zipWith (+) v1 v2}
  zero = Cov {v: A.replicate 10 0.0 }
  mul (Cov {v: a}) (Cov {v: b}) = error "------------> mul cov4 * cov4 not allowed"
  one = Cov { v: [1.0,0.0,0.0,0.0,1.0,0.0,0.0,1.0,0.0,1.0] }
instance ringCov4 :: Ring (Cov Dim4) where
  sub (Cov {v: v1}) (Cov {v: v2}) = Cov {v: A.zipWith (-) v1 v2}
instance showCov4 :: Show (Cov Dim4) where
  show c = "Show (Cov Dim4)\n" <> (show $ toMatrix c)

instance semiringCov5 :: Semiring (Cov Dim5) where
  add (Cov {v: v1}) (Cov {v: v2}) = Cov {v: A.zipWith (+) v1 v2}
  zero = Cov {v: A.replicate 15 0.0 }
  mul a b = error "------------> mul cov5 * cov5 not allowed"
  one = Cov { v: [1.0,0.0,0.0,0.0,0.0,1.0,0.0,0.0,0.0,1.0,0.0,0.0,1.0,0.0,1.0] }
instance ringCov5 :: Ring (Cov Dim5) where
  sub (Cov {v: v1}) (Cov {v: v2}) = Cov {v: A.zipWith (-) v1 v2}
instance showCov5 :: Show (Cov Dim5) where
  show c = "Show (Cov Dim5)\n" <> (show $ toMatrix c)

instance semiringJac :: Semiring (Jac a b) where
  add (Jac {v: v1}) (Jac {v: v2}) = Jac {v: A.zipWith (+) v1 v2}
  zero = undefined
  mul (Jac {v: v1}) (Jac {v: v2}) = undefined -- Cov {v: cov5StdMult v1 v2}
  one = undefined
instance ringJac :: Ring (Jac a b) where
  sub (Jac {v: v1}) (Jac {v: v2}) = Jac {v: A.zipWith (-) v1 v2}
instance showJac :: Show (Jac a b) where
  show a = "Show Jac\n" <> show (toMatrix a)

-- Instances for Vec -- these are always column vectors
instance semiringVec3 :: Semiring (Vec Dim3) where
  add (Vec {v: v1}) (Vec {v: v2}) = Vec {v: A.zipWith (+) v1 v2}
  zero = Vec {v: A.replicate 3 0.0 }
  mul (Vec {v: v1}) (Vec {v: v2}) = undefined
  one = Vec { v: A.replicate 3 1.0 }

instance semiringVec4 :: Semiring (Vec Dim4) where
  add (Vec {v: v1}) (Vec {v: v2}) = Vec {v: A.zipWith (+) v1 v2}
  zero = Vec {v: A.replicate 4 0.0 }
  mul (Vec {v: v1}) (Vec {v: v2}) = undefined
  one = Vec { v: A.replicate 4 1.0 }

instance semiringVec5 :: Semiring (Vec Dim5) where
  add (Vec {v: v1}) (Vec {v: v2}) = Vec {v: A.zipWith (+) v1 v2}
  zero = Vec {v: A.replicate 5 0.0 }
  mul (Vec {v: v1}) (Vec {v: v2}) = undefined
  one = Vec { v: A.replicate 5 1.0 }

instance semiringVec :: Semiring (Vec a) where
  add (Vec {v: v1}) (Vec {v: v2}) = Vec {v: A.zipWith (+) v1 v2}
  {-- zero = error "error calling zero for Vec a" -- Vec {v: A.replicate 5 0.0 } --}
  zero = Vec {v: A.replicate 5 0.0 } `debug` "xxxxxxxxxxx>>> called Vec zero"
  mul (Vec {v: v1}) (Vec {v: v2}) = undefined
  {-- one = error "error calling one for Vec a" -- Vec { v: A.replicate 5 1.0 } --}
  one = Vec { v: A.replicate 5 1.0 } `debug` "xxxxxxxxxxx>>> called Vec one"
instance ringVec :: Ring (Vec a) where
  sub (Vec {v: v1}) (Vec {v: v2}) = Vec {v: A.zipWith (-) v1 v2}
instance showVec :: Show (Vec a) where
  show v = "Show Vec\n" <> show (toMatrix v)

scaleDiag :: Number -> Cov3 -> Cov3
scaleDiag s (Cov {v}) = (Cov {v: v'}) where
  a11 = s * (uidx v 0)
  a22 = s * (uidx v 3)
  a33 = s * (uidx v 5)
  v' = [a11, 0.0, 0.0, a22,0.0, a33]

subm :: Int -> Vec5 -> Vec3
subm n (Vec {v:v5}) = Vec {v: v'} where
  a1 = unsafePartial $ A.unsafeIndex v5 0
  a2 = unsafePartial $ A.unsafeIndex v5 1
  a3 = unsafePartial $ A.unsafeIndex v5 2
  v' = [a1,a2,a3]

subm2 :: Int -> Cov5 -> Cov3
subm2 n (Cov {v:v}) = Cov {v: v'} where
  a11 = unsafePartial $ A.unsafeIndex v 0
  a12 = unsafePartial $ A.unsafeIndex v 1
  a13 = unsafePartial $ A.unsafeIndex v 2
  a22 = unsafePartial $ A.unsafeIndex v 5
  a23 = unsafePartial $ A.unsafeIndex v 6
  a33 = unsafePartial $ A.unsafeIndex v 9
  v' = [a11,a12,a13,a22,a23,a33]

-- CHOLESKY DECOMPOSITION

-- | Simple Cholesky decomposition of a symmetric, positive definite matrix.
--   The result for a matrix /M/ is a lower triangular matrix /L/ such that:
--
--   * /M = LL^T/.
--
--   Example:
--
-- >            (  2 -1  0 )   (  1.41  0     0    )
-- >            ( -1  2 -1 )   ( -0.70  1.22  0    )
-- > choldx     (  0 -1  2 ) = (  0.00 -0.81  1.15 )
--
-- Given a positive-deﬁnite symmetric matrix a[1..n][1..n],
-- this routine constructs its Cholesky decomposition,
-- A = L · L^T
-- The Cholesky factor L is returned in the lower triangle of a,
-- except for its diagonal elements which are returned in p[1..n].

chol :: forall a. Cov a -> Jac a a
chol = choldc
choldc :: forall a. Cov a -> Jac a a
choldc (Cov {v: a}) = Jac {v: a'} where
  n  = case A.length a of
        6  -> 3
        10 -> 4
        15 -> 5
        _  -> error $ "choldc: cannot deal with A.length "
                    <> show (A.length a)
  ll = n*n
  w  = n
  idx :: Int -> Int -> Int
  idx i j | i <= j    = ((i-1)*w - (i-1)*(i-2)/2 + j-i)
          --| otherwise = ((j-1)*w - (j-1)*(j-2)/2 + i-j)
          | otherwise = error "idx: i < j"
  idx' :: Int -> Int -> Int
  idx' j i | i >= j   = (i-1)*w + j-1
           | otherwise = error "idx': i < j"
  {-- run :: forall a. (forall h. Eff (st :: ST h) (STArray h a)) -> Array a --}
  {-- run act = pureST (act >>= unsafeFreeze) --}
  {-- a' = run (do --}
  a' = A.take (n*n) $ pureST ((do
    -- make a STArray of n x n + space for diagonal
    arr <- emptySTArray
    _ <- pushAllSTArray arr (A.replicate (ll+n) 0.0)

    -- loop over input array using Numerical Recipies algorithm (chapter 2.9)
    forE 1 (w+1) \i -> do
      forE i (w+1) \j -> do
          _ <- pokeSTArray arr (idx' i j) (uidx a (idx i j))
          let kmin = 1
              kmax = (i-1) + 1
          forE kmin kmax \k -> do
              aik <- peekSTArray arr (idx' k i)
              ajk <- peekSTArray arr (idx' k j)
              sum <- peekSTArray arr (idx' i j)
              void $ pokeSTArray arr (idx' i j) ((uJust sum)
                                               - (uJust aik) * (uJust ajk))

          msum <- peekSTArray arr (idx' i j)
          let sum' = uJust msum
              sum = if (i==j) && sum' < 0.0
                       then error ("choldc: not a positive definite matrix " <> show a)
                       else sum'
          mp_i' <- peekSTArray arr (ll+i-1)
          let p_i' = uJust mp_i'
              p_i = if i == j then sqrt sum else p_i'
          void $ if i==j
                         then pokeSTArray arr (ll+i-1) p_i -- store diag terms outside main array
                         else pokeSTArray arr (idx' i j) (sum/p_i)
          pure $ unit

    -- copy diagonal back into array
    forE 1 (w+1) \i -> do
          maii <- peekSTArray arr (ll+i-1)
          let aii = uJust maii
          void $ pokeSTArray arr (idx' i i) aii
    pure arr) >>= unsafeFreeze)

-- | Matrix inversion using Cholesky decomposition
-- | based on Numerical Recipies formula in 2.9
--
cholInv :: forall a. Cov a -> Int -> Cov a
cholInv (Cov {v: a}) n = Cov {v: a'} where
  ll = n*n --n*(n+1)/2
  idx :: Int -> Int -> Int -- index into values array of symmetric matrices
  idx i j | i <= j    = ((i-1)*n - (i-1)*(i-2)/2 + j-i)
          | otherwise = ((j-1)*n - (j-1)*(j-2)/2 + i-j)
  idx' :: Int -> Int -> Int -- index into values array for full matrix
  idx' i j = (i-1)*n + j-1
  l = pureST ((do
    -- make a STArray of n x n + space for diagonal +1 for summing
    arr <- emptySTArray
    void $ pushAllSTArray arr (A.replicate (ll+n+1) 0.0)

    -- loop over input array using Numerical Recipies algorithm (chapter 2.9)
    forE 1 (n+1) \i -> do
      forE i (n+1) \j -> do
          let aij = uidx a (idx i j)
          void $ if i==j then pokeSTArray arr (ll+i-1) aij
                         else pokeSTArray arr (idx' j i) aij
          forE 1 i \k -> do
              maik <- peekSTArray arr (idx' i k)
              majk <- peekSTArray arr (idx' j k)
              maij <- if i==j then peekSTArray arr (ll+i-1)
                              else peekSTArray arr (idx' j i)
              let sum = (uJust maij) - (uJust maik) * (uJust majk)
              void $ if i==j then pokeSTArray arr (ll+i-1) sum
                             else pokeSTArray arr (idx' j i) sum
          msum <- if i==j then peekSTArray arr (ll+i-1)
                          else peekSTArray arr (idx' j i)
          let sum' = uJust msum
              sum = if i==j && sum' < 0.0
                       then error ("choldInv: not a positive definite matrix "
                                    <> show a)
                       else sum'
          mp_i' <- peekSTArray arr (ll+i-1)
          let p_i' = uJust mp_i'
              p_i = if i == j then sqrt sum else p_i'
          void $ if i==j then pokeSTArray arr (ll+i-1) p_i
                         else pokeSTArray arr (idx' j i) (sum/p_i)
          pure $ unit

    -- invert L -> L^(-1)
    forE 1 (n+1) \i -> do
      mp_i <- peekSTArray arr (ll+i-1)
      void $ pokeSTArray arr (idx' i i) (1.0/(uJust mp_i))
      forE (i+1) (n+1) \j -> do
        void $ pokeSTArray arr (ll+n) 0.0
        forE i j \k -> do
          majk <- peekSTArray arr (idx' j k)
          maki <- peekSTArray arr (idx' k i)
          sum <- peekSTArray arr (ll+n)
          void $ pokeSTArray arr (ll+n)
                    ((uJust sum) - (uJust majk) * (uJust maki))
        msum <- peekSTArray arr (ll+n)
        mp_j <- peekSTArray arr (ll+j-1)
        void $ pokeSTArray arr (idx' j i) ((uJust msum)/(uJust mp_j))
    pure arr) >>= unsafeFreeze)
  a' = do
    i <- A.range 1 n
    j <- A.range i n
    let aij = sum do
                  k <- A.range 1 n
                  pure $ (uidx l (idx' k i)) * (uidx l (idx' k j))
    pure $ aij

--C version Numerical Recipies 2.9
--for (i=1;i<=n;i++) {
--  for (j=i;j<=n;j++) {
--    for (sum=a[i][j],k=i-1;k>=1;k--) sum -= a[i][k]*a[j][k];
--    if (i == j) {
--      if (sum <= 0.0) nrerror("choldc failed");
--      p[i]=sqrt(sum);
--    } else a[j][i]=sum/p[i];
--  }
--}
-- In this, and many other applications, one often needs L^(−1) . The lower
-- triangle of this matrix can be efﬁciently found from the output of choldc:
--for (i=1;i<=n;i++) {
--  a[i][i]=1.0/p[i];
--  for (j=i+1;j<=n;j++) {
--    sum=0.0;
--    for (k=i;k<j;k++) sum -= a[j][k]*a[k][i];
--    a[j][i]=sum/p[j];
--  }
--}

module Data.Cov
    where

import Prelude
import Data.Array
  ( replicate, unsafeIndex, zipWith, length, foldl, range, take
  ) as A
import Control.Monad.Eff ( forE )
import Control.Monad.ST ( pureST )
import Data.Array.ST (emptySTArray, peekSTArray, pokeSTArray
                     , pushAllSTArray, unsafeFreeze)
import Data.Foldable ( sum )
import Partial.Unsafe ( unsafePartial )
import Data.Maybe ( Maybe (..) )
import Control.MonadZero (guard)
import Data.Int ( toNumber, floor )
import Math ( abs, sqrt )
import Unsafe.Coerce ( unsafeCoerce ) as Unsafe.Coerce

import Data.SimpleMatrix
  ( Matrix
  , transpose
  , fromArray, fromArray2, toArray
  ) as M
import Stuff

newtype Dim3 = DDim3 Int
newtype Dim4 = DDim4 Int
newtype Dim5 = DDim5 Int
class DDim a where
  ddim :: a -> Int
instance ddim3 :: DDim Dim3 where
  ddim _ = 3
instance ddim4 :: DDim Dim4 where
  ddim _ = 4
instance ddim5 :: DDim Dim5 where
  ddim _ = 5
instance ddima :: DDim a where
  ddim _ = undefined

{-- mapTo :: forall a. DDim a => Cov a -> Dim5 --}
{-- mapTo = Unsafe.Coerce.unsafeCoerce --}
{-- class Dim a where --}
{--   dim :: a -> Int --}
{-- instance dima :: Dim (Cov a) where --}
{--   dim ccc = n where --}
{--     xx = mapTo ccc --}
{--     n = ddim xx --}
{-- instance dim3 :: Dim (Cov Dim3) where --}
{--   dim ccc = 3 --}
{-- instance dim4 :: Dim (Cov Dim4) where --}
{--   dim ccc = 4 --}
{-- instance dim5 :: Dim (Cov Dim5) where --}
{--   dim ccc = 5 --}

newtype Cov a = Cov { v :: Array Number }
newtype Jac a b = Jac { v :: Array Number }
newtype Vec a = Vec { v :: Array Number }
type Cov3 = Cov Dim3
type Cov4 = Cov Dim4
type Cov5 = Cov Dim5
{-- type Jac43 = Jac Dim4 Dim3 --}
type Jac53 = Jac Dim5 Dim3
type Jac33 = Jac Dim3 Dim3
type Jac34 = Jac Dim3 Dim4
type Jac35 = Jac Dim3 Dim5
type Jac44 = Jac Dim4 Dim4
type Jac55 = Jac Dim5 Dim5
type Vec3 = Vec Dim3
type Vec4 = Vec Dim4
type Vec5 = Vec Dim5
type Jacs = {aa :: Jac53, bb :: Jac53, h0 :: Vec5}

-- access to arrays of symmetrical matrices
indV :: Int -> Int -> Int -> Int
indV w i0 j0 = (i0*w+j0)
indVs :: Int -> Int -> Int -> Int
indVs w i0 j0 | i0 <= j0  = (i0*w - (i0*(i0-1)) `div` 2 + j0-i0)
              | otherwise = (j0*w - (j0*(j0-1)) `div` 2 + i0-j0)
-------------------------------------------------------------------------
-------------------------------------------------------------------------
-------------------------------------------------------------------------
-------------------------------------------------------------------------
-- Mat to give behavior to Cov and Vec and Jac
-- ability to convert to and from Matrix and Array
-- while keeping info about dimensionality
-- also define Semiring and Ring functions
--

class Mat a where
  val :: a -> Array Number
  fromArray :: Array Number -> a
  toArray :: a -> Array Number
instance matCova :: Mat (Cov a) where
  val (Cov {v}) = v
  fromArray a = c' where
    l = A.length a
    c' = case l of
              6   -> Cov {v: a}
              10  -> Cov {v: a}
              15  -> Cov {v: a}
              _   -> Cov {v: let
                            n = floor <<< sqrt <<< toNumber $ l
                            iv = indV n
                            in do -- only upper triangle
                              i0 <- A.range 0 (n-1)
                              j0 <- A.range i0 (n-1)
                              pure $ uidx a (iv i0 j0) }
  toArray c@(Cov {v}) = v' where
          l = A.length v
          n = case l of -- ceil ((sqrt(8.0 * (toNumber l) + 1.0) - 1.0)/2.0)
            6  -> 3
            10 -> 4
            15 -> 5
            _  -> error $ "matCova: toArray not supported for lenght "
                        <> show l
          iv = indVs n
          v' = do
                  i0 <- A.range 0 (n-1)
                  j0 <- A.range 0 (n-1)
                  pure $ uidx v (iv i0 j0)
instance matVeca :: Mat (Vec a) where
  val (Vec {v}) = v
  fromArray a = Vec {v: a}
  toArray (Vec {v}) = v
instance matJacab :: Mat (Jac a b) where
  val (Jac {v}) = v
  fromArray a = Jac {v: a}
  toArray (Jac {v}) = v

class Mat1 a where
  toMatrix :: a -> M.Matrix
instance mat1Cova :: Mat1 (Cov a) where
  toMatrix a@(Cov {v}) = case A.length v of
                            6  -> M.fromArray2 3 3 v
                            10 -> M.fromArray2 4 4 v
                            15 -> M.fromArray2 5 5 v
                            _ -> error $ "mat1Cova toMatrix "
                                          <> show (A.length v)
instance mat1Veca :: Mat1 (Vec a) where
  toMatrix (Vec {v}) = M.fromArray (A.length v) v
instance mat1Jac53 :: Mat1 (Jac Dim5 Dim3) where
  toMatrix (Jac {v}) = M.fromArray2 5 3 v
instance mat1Jac35 :: Mat1 (Jac Dim3 Dim5) where
  toMatrix (Jac {v}) = M.fromArray2 3 5 v
instance mat1Jacaa :: Mat1 (Jac a b) where
  toMatrix j@(Jac {v}) = case A.length v of
                              9  -> M.fromArray2 3 3 v
                              16 -> M.fromArray2 4 4 v
                              25 -> M.fromArray2 5 5 v
                              12 -> M.fromArray2 3 4 v `debug` "this should not have happened ??????????????????? 4 3"
                              15 -> M.fromArray2 5 3 v `debug` "this should not have happened ??????????????????? 5 3"
                              _  -> error $ "mat1Jacaa toMatrix "
                                          <> show (A.length v)

-----------------------------------------------------------------
-- | funcitons for symetric matrices: Cov
-- | type class SymMat
class SymMat a where
  inv :: Cov a -> Cov a                -- | inverse matrix
  invMaybe :: Cov a -> Maybe (Cov a)   -- | Maybe inverse matrix
  det :: Cov a -> Number               -- | determinant
  diag :: Cov a -> Array Number        -- | Array of diagonal elements
  {-- chol :: Cov a -> Jac a a       -- | Cholsky decomposition --}
instance symMatCov3 :: SymMat Dim3 where
--  inv m | trace ( "inv " <> (show m) ) False = undefined
  inv m = m' where
    mm = invMaybe m
    m' = case mm of
           Nothing -> error $ "can't invert 3x3 matrix " <> show m
           Just x -> x
  invMaybe (Cov {v}) = _inv v where
    _inv :: Array Number -> Maybe (Cov Dim3)
    _inv = unsafePartial $ \[a11,a12,a13,a22,a23,a33] -> do
      let det = (a33*a12*a12 - 2.0*a13*a23*a12 + a13*a13*a22
                + a11*(a23*a23 - a22*a33))
      guard $ (abs det) > 1.0e-50
      let
          b11 = (a23*a23 - a22*a33)/det
          b12 = (a12*a33 - a13*a23)/det
          b13 = (a13*a22 - a12*a23)/det
          b22 = (a13*a13 - a11*a33)/det
          b23 = (a11*a23 - a12*a13)/det
          b33 = (a12*a12 - a11*a22)/det
          v' = [b11,b12,b13,b22,b23,b33]
      pure $ Cov {v: v'}
  det (Cov {v}) = _det v where
    _det :: Array Number -> Number
    _det = unsafePartial $ \[a,b,c,d,e,f] ->
            a*d*f - a*e*e - b*b*f + 2.0*b*c*e - c*c*d
  det (Cov {v}) = dd where
        a = unsafePartial $ A.unsafeIndex v 0
        b = unsafePartial $ A.unsafeIndex v 1
        c = unsafePartial $ A.unsafeIndex v 2
        d = unsafePartial $ A.unsafeIndex v 3
        e = unsafePartial $ A.unsafeIndex v 4
        f = unsafePartial $ A.unsafeIndex v 5
        dd = a*d*f - a*e*e - b*b*f + 2.0*b*c*e - c*c*d
  diag (Cov {v}) = a where
    a11 = unsafePartial $ A.unsafeIndex v 0
    a22 = unsafePartial $ A.unsafeIndex v 3
    a33 = unsafePartial $ A.unsafeIndex v 5
    a = [a11,a22,a33]
instance symMatCov4 :: SymMat Dim4 where
  inv m = uJust (invMaybe m)
  invMaybe (Cov {v}) = _inv v where
    _inv :: Array Number -> Maybe (Cov Dim4)
    _inv = unsafePartial $ \[a,b,c,d,e,f,g,h,i,j] -> do
      let det = (a*e*h*j - a*e*i*i - a*f*f*j + 2.0*a*f*g*i - a*g*g*h
            - b*b*h*j + b*b*i*i - 2.0*d*(b*f*i - b*g*h - c*e*i + c*f*g)
            + b*c*(2.0*f*j - 2.0*g*i) + c*c*(g*g - e*j) + d*d*(f*f - e*h))
      guard $ (abs det) > 1.0e-50
      let a' = (-j*f*f + 2.0*g*i*f - e*i*i - g*g*h + e*h*j)/det
          b' = (b*i*i - d*f*i - c*g*i + d*g*h + c*f*j - b*h*j)/det
          c' = (c*g*g - d*f*g - b*i*g + d*e*i - c*e*j + b*f*j)/det
          d' = (d*f*f - c*g*f - b*i*f - d*e*h + b*g*h + c*e*i)/det
          e' = (-j*c*c + 2.0*d*i*c - a*i*i - d*d*h + a*h*j)/det
          f' = (f*d*d - c*g*d - b*i*d + a*g*i + b*c*j - a*f*j)/det
          g' = (g*c*c - d*f*c - b*i*c + b*d*h - a*g*h + a*f*i)/det
          h' = (-j*b*b + 2.0*d*g*b - a*g*g - d*d*e + a*e*j)/det
          i' = (i*b*b - d*f*b - c*g*b + c*d*e + a*f*g - a*e*i)/det
          j' = (-h*b*b + 2.0*c*f*b - a*f*f - c*c*e + a*e*h)/det
      pure $ fromArray [a',b',c',d',e',f',g',h',i',j']
  det (Cov {v}) = _det v where
    _det = unsafePartial $ \[a,b,c,d,e,f,g,h,i,j] ->
        (a*e*h*j - a*e*i*i - a*f*f*j + 2.0*a*f*g*i - a*g*g*h
          - b*b*h*j + b*b*i*i - 2.0*d*(b*f*i - b*g*h - c*e*i + c*f*g)
          + b*c*(2.0*f*j - 2.0*g*i) + c*c*(g*g - e*j) + d*d*(f*f - e*h))
  diag (Cov {v}) = _diag v where
    _diag :: Array Number -> Array Number
    _diag = unsafePartial $ \[a11,_,_,_,a22,_,_,a33,_,a44] -> [a11,a22,a33,a44]
instance symMatCov5 :: SymMat Dim5 where
  inv m = cholInv m 5
  invMaybe m = Just (cholInv m 5)
  det (Cov {v}) = _det v where
    _det [a,b,c,d,e,f,g,h,i,j,k,l,m,n,o] =
      a*f*j*m*o - a*f*j*n*n - a*f*k*k*o + 2.0*a*f*k*l*n - a*f*l*l*m
      - a*g*g*m*o + a*g*g*n*n + 2.0*a*g*h*k*o - 2.0*a*g*h*l*n - 2.0*a*g*i*k*n
      + 2.0*a*g*i*l*m - a*h*h*j*o + a*h*h*l*l + 2.0*a*h*i*j*n - 2.0*a*h*i*k*l
      - a*i*i*j*m + a*i*i*k*k - b*b*j*m*o + b*b*j*n*n + b*b*k*k*o
      - 2.0*b*b*k*l*n + b*b*l*l*m + 2.0*b*c*g*m*o - 2.0*b*c*g*n*n - 2.0*b*c*h*k*o
      + 2.0*b*c*h*l*n + 2.0*b*c*i*k*n - 2.0*b*c*i*l*m - 2.0*b*d*g*k*o
      + 2.0*b*d*g*l*n + 2.0*b*d*h*j*o - 2.0*b*d*h*l*l - 2.0*b*d*i*j*n
      + 2.0*b*d*i*k*l + 2.0*b*e*g*k*n - 2.0*b*e*g*l*m - 2.0*b*e*h*j*n
      + 2.0*b*e*h*k*l + 2.0*b*e*i*j*m - 2.0*b*e*i*k*k - c*c*f*m*o + c*c*f*n*n
      + c*c*h*h*o - 2.0*c*c*h*i*n + c*c*i*i*m + 2.0*c*d*f*k*o - 2.0*c*d*f*l*n
      - 2.0*c*d*g*h*o + 2.0*c*d*g*i*n + 2.0*c*d*h*i*l - 2.0*c*d*i*i*k
      - 2.0*c*e*f*k*n + 2.0*c*e*f*l*m + 2.0*c*e*g*h*n - 2.0*c*e*g*i*m
      - 2.0*c*e*h*h*l + 2.0*c*e*h*i*k - d*d*f*j*o + d*d*f*l*l + d*d*g*g*o
      - 2.0*d*d*g*i*l + d*d*i*i*j + 2.0*d*e*f*j*n - 2.0*d*e*f*k*l - 2.0*d*e*g*g*n
      + 2.0*d*e*g*h*l + 2.0*d*e*g*i*k - 2.0*d*e*h*i*j - e*e*f*j*m + e*e*f*k*k
      + e*e*g*g*m - 2.0*e*e*g*h*k + e*e*h*h*j
    _det _ = undefined
  diag (Cov {v}) = a where
    a11 = unsafePartial $ A.unsafeIndex v 0
    a22 = unsafePartial $ A.unsafeIndex v 5
    a33 = unsafePartial $ A.unsafeIndex v 9
    a44 = unsafePartial $ A.unsafeIndex v 12
    a55 = unsafePartial $ A.unsafeIndex v 14
    a = [a11,a22,a33,a44,a55]

class MulMat a b c | a b -> c where
  mulm :: a -> b -> c
infixr 7 mulm as *.
instance mulMata :: MulMat (Cov a) (Cov a) (Jac a a) where
  mulm c1 c2 = j' where
    mc1 = toMatrix c1
    mc2 = toMatrix c2
    mj' = mc1 * mc2
    j' = fromArray $ M.toArray mj'
instance mulMatJC :: MulMat (Jac a b) (Cov b) (Jac a b) where
  mulm j@(Jac {v: va}) c@(Cov {v: vb}) = Jac {v: v'} where
    nb = case A.length vb of
              6  -> 3
              10 -> 4
              15 -> 5
              _  -> error $ "mulMatJC wrong length of Cov v "
                            <> show (A.length vb)
    na = (A.length va)/nb
    n = nb
    idx i0 j0 | i0 <= j0  = (i0*n - i0*(i0-1)/2 + j0-i0)
              | otherwise = (j0*n - j0*(j0-1)/2 + i0-j0)
    v' = do
      i0 <- A.range 0 (na-1)
      j0 <- A.range 0 (nb-1)
      pure $ sum do
        k0 <- A.range 0 (nb-1)
        pure $ (uidx va (i0*nb+k0)) * (uidx vb (idx k0 j0))
instance mulMatCJ :: MulMat (Cov a) (Jac a b) (Jac a b) where
  mulm c@(Cov {v: va}) j@(Jac {v: vb}) = Jac {v: v'} where
    na = case A.length va of
              6  -> 3
              10 -> 4
              15 -> 5
              _  -> error $ "mulMatCJ wrong length of Cov v "
                            <> show (A.length va)
    nb = (A.length vb)/na
    n = na
    idx i0 j0 | i0 <= j0  = (i0*n - i0*(i0-1)/2 + j0-i0)
              | otherwise = (j0*n - j0*(j0-1)/2 + i0-j0)
    v' = do
      i0 <- A.range 0 (na-1)
      j0 <- A.range 0 (nb-1)
      pure $ sum do
        k0 <- A.range 0 (na-1)
        pure $ (uidx va (idx i0 k0)) * (uidx vb (k0*nb+j0))
instance mulMatJV :: MulMat (Jac a b) (Vec b) (Vec a) where
  mulm j@(Jac {v: va}) v@(Vec {v: vb}) = Vec {v: v'} where
    nb = A.length vb
    na = (A.length va)/nb
    v' = do
      i0 <- A.range 0 (na-1)
      let j0 = 0
      pure $ sum do
        k0 <- A.range 0 (nb-1)
        pure $ (uidx va (i0*nb+k0)) * (uidx vb (k0))
-- ???????????this needs to be generalized to (Jac a b)
instance mulMatJJ :: MulMat (Jac Dim3 Dim5) (Jac Dim5 Dim3) (Jac Dim3 Dim3) where
  mulm j1 j2 = j' where
    mj1 = toMatrix j1
    mj2 = toMatrix j2
    mj' = mj1 * mj2
    j' = fromArray $ M.toArray mj'
instance mulMatCV :: MulMat (Cov a) (Vec a) (Vec a) where
  mulm c v = v' where
    mc = toMatrix c
    mv = toMatrix v
    mv' = mc * mv
    v' = fromArray $ M.toArray mv'
instance mulMatVV :: MulMat (Vec a) (Vec a) Number where
  mulm (Vec {v:v1}) (Vec {v:v2}) = A.foldl (+) zero $ A.zipWith (*) v1 v2
class TrMat a b | a -> b where
  tr :: a -> b
instance trMatC :: TrMat (Cov a) (Cov a) where
  tr c = c
instance trMatJ :: TrMat (Jac a b) (Jac b a) where
  tr j@(Jac {v}) = Jac {v:v'} where
    l = A.length v
    na = case l of
              9 -> 3
              15 -> 5
              16 -> 4
              25 -> 5
              _  -> error $ "trMatJ: sorry, can't do anything but 5x3 and square "
                            <> show (A.length v)
    nb = l/na
    v' = do
      i0 <- A.range 0 (nb-1)
      j0 <- A.range 0 (na-1)
      pure $ (uidx v (j0*nb+i0))
class SW a b c | a b -> c where
  sw :: a -> b -> c
infixl 7 sw as .*.
instance swVec :: SW (Vec a) (Cov a) Number where
  sw v c = n where
    mv = toMatrix v
    mc = toMatrix c
    mc' = M.transpose mv * mc * mv
    n = uidx (M.toArray mc') 0
instance swCov :: SW (Cov a) (Cov a) (Cov a) where
  sw c1 c2 = c' where
    j' = c1 *. c2 *. c1
    c' = fromArray $ toArray j'
instance swJac :: SW (Jac a b) (Cov a) (Cov b) where
  sw j@(Jac {v: va}) c@(Cov {v: vb}) = Cov {v: v'} where
    {-- mj = toMatrix j --}
    {-- mc = toMatrix c --}
    {-- mc' = M.transpose mj * mc * mj --}

    l = A.length vb
    n = case l of
              6  -> 3
              10 -> 4
              15 -> 5
              _  -> error $ "swJac: don'w know how to " <> show l
    m = (A.length va)/n -- > mxn * nxn * nxm -> mxm
    -- vaT * vb * va
    {-- vat= do --}
    {--   i0 <- A.range 0 (m-1) --}
    {--   j0 <- A.range 0 (n-1) --}
    {--   pure $ (uidx va (j0*m+i0)) --}
    {-- mjt = M.fromArray2 m n vat --}
    {-- mjt''' = M.transpose mj --}

    idx i0 j0 | i0 <= j0  = (i0*n - i0*(i0-1)/2 + j0-i0)
              | otherwise = (j0*n - j0*(j0-1)/2 + i0-j0)
   {-- vb''' = do --}
    {--   i0 <- A.range 0 (n-1) --}
    {--   j0 <- A.range 0 (n-1) --}
    {--   pure $ (uidx vb (idx i0 j0)) --}
    {-- mb''' = M.fromArray2 n n vb''' --}

    {-- va''' = do --}
    {--   i0 <- A.range 0 (n-1) --}
    {--   j0 <- A.range 0 (m-1) --}
    {--   pure $ (uidx va (i0*m+j0)) --}
    {-- mj''' = M.fromArray2 n m va''' --}

    vint = do
      i0 <- A.range 0 (n-1)
      j0 <- A.range 0 (m-1)
      pure $ sum do
        k0 <- A.range 0 (n-1)
        pure $ (uidx vb (idx i0 k0)) * (uidx va (k0*m+j0))
    vf = do
      i0 <- A.range 0 (m-1)
      j0 <- A.range i0 (m-1)
      pure $ sum do
        k0 <- A.range 0 (n-1)
        pure $ (uidx va (k0*m+i0) ) * (uidx vint (k0*m+j0))

    {-- vf = M.toArray mc' --}
    lf = A.length vf
    v' = case lf of
              6   -> vf
              10  -> vf
              15  -> vf
              _   -> error $ "swJac: wrong result " <> show lf

    {-- mv' = M.fromArray2 m m v' --}
    {-- xx = n `debug` ("--jT---->\n" <> show mjt''' <> show mjt) --}
    {-- yx = n `debug` ("--c----->\n" <> show mc <> show mb''') --}
    {-- xy = n `debug` ("--j----->\n" <> show mj <> show mj''') --}
    {-- yy = n `debug` ("-jT*c*j->\n" <> show mc' <> show mv') --}
testCov2 :: String
testCov2 = s where
  s = "Test Cov 2----------------------------------------------\n"
    <> "Vec *. Vec = " <> show (v3 *. v3) <> "\n"
    <> "Cov *. Cov = " <> show ((one::Cov3) *. inv (one::Cov3)) <> "\n"
    <> "Vec + Vec = " <> show (v5 + v5) <> "\n"
    <> "chol Cov = " <> show (chol (one::Cov5)) <> "\n"
    <> "Vec .*. Cov = " <> show (v5 .*. inv (one::Cov5)) <> "\n"
  v3 = fromArray [1.0,1.0,1.0] :: Vec3
  v5 = fromArray [1.0,1.0,1.0,1.0,1.0] :: Vec5

instance semiringCov3 :: Semiring (Cov Dim3) where
  add (Cov {v: v1}) (Cov {v: v2}) = Cov {v: A.zipWith (+) v1 v2}
  zero = Cov {v: A.replicate 6 0.0 }
  mul (Cov {v: a}) (Cov {v: b}) = error "------------> mul cov3 * cov3 not allowed"
  one = Cov { v: [1.0, 0.0, 0.0, 1.0, 0.0, 1.0] }
instance ringCov3 :: Ring (Cov Dim3) where
  sub (Cov {v: v1}) (Cov {v: v2}) = Cov {v: A.zipWith (-) v1 v2}
instance showCov3 :: Show (Cov Dim3) where
  show c = "Show (Cov Dim3) \n" <> (show $ toMatrix c)

instance semiringCov4 :: Semiring (Cov Dim4) where
  add (Cov {v: v1}) (Cov {v: v2}) = Cov {v: A.zipWith (+) v1 v2}
  zero = Cov {v: A.replicate 10 0.0 }
  mul (Cov {v: a}) (Cov {v: b}) = error "------------> mul cov4 * cov4 not allowed"
  one = Cov { v: [1.0,0.0,0.0,0.0,1.0,0.0,0.0,1.0,0.0,1.0] }
instance ringCov4 :: Ring (Cov Dim4) where
  sub (Cov {v: v1}) (Cov {v: v2}) = Cov {v: A.zipWith (-) v1 v2}
instance showCov4 :: Show (Cov Dim4) where
  show c = "Show (Cov Dim4)\n" <> (show $ toMatrix c)

instance semiringCov5 :: Semiring (Cov Dim5) where
  add (Cov {v: v1}) (Cov {v: v2}) = Cov {v: A.zipWith (+) v1 v2}
  zero = Cov {v: A.replicate 15 0.0 }
  mul a b = error "------------> mul cov5 * cov5 not allowed"
  one = Cov { v: [1.0,0.0,0.0,0.0,0.0,1.0,0.0,0.0,0.0,1.0,0.0,0.0,1.0,0.0,1.0] }
instance ringCov5 :: Ring (Cov Dim5) where
  sub (Cov {v: v1}) (Cov {v: v2}) = Cov {v: A.zipWith (-) v1 v2}
instance showCov5 :: Show (Cov Dim5) where
  show c = "Show (Cov Dim5)\n" <> (show $ toMatrix c)

instance semiringJac :: Semiring (Jac a b) where
  add (Jac {v: v1}) (Jac {v: v2}) = Jac {v: A.zipWith (+) v1 v2}
  zero = undefined
  mul (Jac {v: v1}) (Jac {v: v2}) = undefined -- Cov {v: cov5StdMult v1 v2}
  one = undefined
instance ringJac :: Ring (Jac a b) where
  sub (Jac {v: v1}) (Jac {v: v2}) = Jac {v: A.zipWith (-) v1 v2}
instance showJac :: Show (Jac a b) where
  show a = "Show Jac\n" <> show (toMatrix a)

-- Instances for Vec -- these are always column vectors
instance semiringVec3 :: Semiring (Vec Dim3) where
  add (Vec {v: v1}) (Vec {v: v2}) = Vec {v: A.zipWith (+) v1 v2}
  zero = Vec {v: A.replicate 3 0.0 }
  mul (Vec {v: v1}) (Vec {v: v2}) = undefined
  one = Vec { v: A.replicate 3 1.0 }

instance semiringVec4 :: Semiring (Vec Dim4) where
  add (Vec {v: v1}) (Vec {v: v2}) = Vec {v: A.zipWith (+) v1 v2}
  zero = Vec {v: A.replicate 4 0.0 }
  mul (Vec {v: v1}) (Vec {v: v2}) = undefined
  one = Vec { v: A.replicate 4 1.0 }

instance semiringVec5 :: Semiring (Vec Dim5) where
  add (Vec {v: v1}) (Vec {v: v2}) = Vec {v: A.zipWith (+) v1 v2}
  zero = Vec {v: A.replicate 5 0.0 }
  mul (Vec {v: v1}) (Vec {v: v2}) = undefined
  one = Vec { v: A.replicate 5 1.0 }

instance semiringVec :: Semiring (Vec a) where
  add (Vec {v: v1}) (Vec {v: v2}) = Vec {v: A.zipWith (+) v1 v2}
  {-- zero = error "error calling zero for Vec a" -- Vec {v: A.replicate 5 0.0 } --}
  zero = Vec {v: A.replicate 5 0.0 } `debug` "xxxxxxxxxxx>>> called Vec zero"
  mul (Vec {v: v1}) (Vec {v: v2}) = undefined
  {-- one = error "error calling one for Vec a" -- Vec { v: A.replicate 5 1.0 } --}
  one = Vec { v: A.replicate 5 1.0 } `debug` "xxxxxxxxxxx>>> called Vec one"
instance ringVec :: Ring (Vec a) where
  sub (Vec {v: v1}) (Vec {v: v2}) = Vec {v: A.zipWith (-) v1 v2}
instance showVec :: Show (Vec a) where
  show v = "Show Vec\n" <> show (toMatrix v)

scaleDiag :: Number -> Cov3 -> Cov3
scaleDiag s (Cov {v}) = (Cov {v: v'}) where
  a11 = s * (uidx v 0)
  a22 = s * (uidx v 3)
  a33 = s * (uidx v 5)
  v' = [a11, 0.0, 0.0, a22,0.0, a33]

subm :: Int -> Vec5 -> Vec3
subm n (Vec {v:v5}) = Vec {v: v'} where
  a1 = unsafePartial $ A.unsafeIndex v5 0
  a2 = unsafePartial $ A.unsafeIndex v5 1
  a3 = unsafePartial $ A.unsafeIndex v5 2
  v' = [a1,a2,a3]

subm2 :: Int -> Cov5 -> Cov3
subm2 n (Cov {v:v}) = Cov {v: v'} where
  a11 = unsafePartial $ A.unsafeIndex v 0
  a12 = unsafePartial $ A.unsafeIndex v 1
  a13 = unsafePartial $ A.unsafeIndex v 2
  a22 = unsafePartial $ A.unsafeIndex v 4
  a23 = unsafePartial $ A.unsafeIndex v 5
  a33 = unsafePartial $ A.unsafeIndex v 8
  v' = [a11,a12,a13,a22,a23,a33]

-- CHOLESKY DECOMPOSITION

-- | Simple Cholesky decomposition of a symmetric, positive definite matrix.
--   The result for a matrix /M/ is a lower triangular matrix /L/ such that:
--
--   * /M = LL^T/.
--
--   Example:
--
-- >            (  2 -1  0 )   (  1.41  0     0    )
-- >            ( -1  2 -1 )   ( -0.70  1.22  0    )
-- > choldx     (  0 -1  2 ) = (  0.00 -0.81  1.15 )
--
-- Given a positive-deﬁnite symmetric matrix a[1..n][1..n],
-- this routine constructs its Cholesky decomposition,
-- A = L · L^T
-- The Cholesky factor L is returned in the lower triangle of a,
-- except for its diagonal elements which are returned in p[1..n].

chol :: forall a. Cov a -> Jac a a
chol = choldc
choldc :: forall a. Cov a -> Jac a a
choldc (Cov {v: a}) = Jac {v: a'} where
  n  = case A.length a of
        6  -> 3
        10 -> 4
        15 -> 5
        _  -> error $ "choldc: cannot deal with A.length "
                    <> show (A.length a)
  ll = n*n
  w  = n
  idx :: Int -> Int -> Int
  idx i j | i <= j    = ((i-1)*w - (i-1)*(i-2)/2 + j-i)
          --| otherwise = ((j-1)*w - (j-1)*(j-2)/2 + i-j)
          | otherwise = error "idx: i < j"
  idx' :: Int -> Int -> Int
  idx' j i | i >= j   = (i-1)*w + j-1
           | otherwise = error "idx': i < j"
  {-- run :: forall a. (forall h. Eff (st :: ST h) (STArray h a)) -> Array a --}
  {-- run act = pureST (act >>= unsafeFreeze) --}
  {-- a' = run (do --}
  a' = A.take (n*n) $ pureST ((do
    -- make a STArray of n x n + space for diagonal
    arr <- emptySTArray
    _ <- pushAllSTArray arr (A.replicate (ll+n) 0.0)

    -- loop over input array using Numerical Recipies algorithm (chapter 2.9)
    forE 1 (w+1) \i -> do
      forE i (w+1) \j -> do
          _ <- pokeSTArray arr (idx' i j) (uidx a (idx i j))
          let kmin = 1
              kmax = (i-1) + 1
          forE kmin kmax \k -> do
              aik <- peekSTArray arr (idx' k i)
              ajk <- peekSTArray arr (idx' k j)
              sum <- peekSTArray arr (idx' i j)
              void $ pokeSTArray arr (idx' i j) ((uJust sum)
                                               - (uJust aik) * (uJust ajk))

          msum <- peekSTArray arr (idx' i j)
          let sum' = uJust msum
              sum = if (i==j) && sum' < 0.0
                       then error ("choldc: not a positive definite matrix " <> show a)
                       else sum'
          mp_i' <- peekSTArray arr (ll+i-1)
          let p_i' = uJust mp_i'
              p_i = if i == j then sqrt sum else p_i'
          void $ if i==j
                         then pokeSTArray arr (ll+i-1) p_i -- store diag terms outside main array
                         else pokeSTArray arr (idx' i j) (sum/p_i)
          pure $ unit

    -- copy diagonal back into array
    forE 1 (w+1) \i -> do
          maii <- peekSTArray arr (ll+i-1)
          let aii = uJust maii
          void $ pokeSTArray arr (idx' i i) aii
    pure arr) >>= unsafeFreeze)

-- | Matrix inversion using Cholesky decomposition
-- | based on Numerical Recipies formula in 2.9
--
cholInv :: forall a. Cov a -> Int -> Cov a
cholInv (Cov {v: a}) n = Cov {v: a'} where
  ll = n*n --n*(n+1)/2
  idx :: Int -> Int -> Int -- index into values array of symmetric matrices
  idx i j | i <= j    = ((i-1)*n - (i-1)*(i-2)/2 + j-i)
          | otherwise = ((j-1)*n - (j-1)*(j-2)/2 + i-j)
  idx' :: Int -> Int -> Int -- index into values array for full matrix
  idx' i j = (i-1)*n + j-1
  l = pureST ((do
    -- make a STArray of n x n + space for diagonal +1 for summing
    arr <- emptySTArray
    void $ pushAllSTArray arr (A.replicate (ll+n+1) 0.0)

    -- loop over input array using Numerical Recipies algorithm (chapter 2.9)
    forE 1 (n+1) \i -> do
      forE i (n+1) \j -> do
          let aij = uidx a (idx i j)
          void $ if i==j then pokeSTArray arr (ll+i-1) aij
                         else pokeSTArray arr (idx' j i) aij
          forE 1 i \k -> do
              maik <- peekSTArray arr (idx' i k)
              majk <- peekSTArray arr (idx' j k)
              maij <- if i==j then peekSTArray arr (ll+i-1)
                              else peekSTArray arr (idx' j i)
              let sum = (uJust maij) - (uJust maik) * (uJust majk)
              void $ if i==j then pokeSTArray arr (ll+i-1) sum
                             else pokeSTArray arr (idx' j i) sum
          msum <- if i==j then peekSTArray arr (ll+i-1)
                          else peekSTArray arr (idx' j i)
          let sum' = uJust msum
              sum = if i==j && sum' < 0.0
                       then error ("choldInv: not a positive definite matrix "
                                    <> show a)
                       else sum'
          mp_i' <- peekSTArray arr (ll+i-1)
          let p_i' = uJust mp_i'
              p_i = if i == j then sqrt sum else p_i'
          void $ if i==j then pokeSTArray arr (ll+i-1) p_i
                         else pokeSTArray arr (idx' j i) (sum/p_i)
          pure $ unit

    -- invert L -> L^(-1)
    forE 1 (n+1) \i -> do
      mp_i <- peekSTArray arr (ll+i-1)
      void $ pokeSTArray arr (idx' i i) (1.0/(uJust mp_i))
      forE (i+1) (n+1) \j -> do
        void $ pokeSTArray arr (ll+n) 0.0
        forE i j \k -> do
          majk <- peekSTArray arr (idx' j k)
          maki <- peekSTArray arr (idx' k i)
          sum <- peekSTArray arr (ll+n)
          void $ pokeSTArray arr (ll+n)
                    ((uJust sum) - (uJust majk) * (uJust maki))
        msum <- peekSTArray arr (ll+n)
        mp_j <- peekSTArray arr (ll+j-1)
        void $ pokeSTArray arr (idx' j i) ((uJust msum)/(uJust mp_j))
    pure arr) >>= unsafeFreeze)
  a' = do
    i <- A.range 1 n
    j <- A.range i n
    let aij = sum do
                  k <- A.range 1 n
                  pure $ (uidx l (idx' k i)) * (uidx l (idx' k j))
    pure $ aij

--C version Numerical Recipies 2.9
--for (i=1;i<=n;i++) {
--  for (j=i;j<=n;j++) {
--    for (sum=a[i][j],k=i-1;k>=1;k--) sum -= a[i][k]*a[j][k];
--    if (i == j) {
--      if (sum <= 0.0) nrerror("choldc failed");
--      p[i]=sqrt(sum);
--    } else a[j][i]=sum/p[i];
--  }
--}
-- In this, and many other applications, one often needs L^(−1) . The lower
-- triangle of this matrix can be efﬁciently found from the output of choldc:
--for (i=1;i<=n;i++) {
--  a[i][i]=1.0/p[i];
--  for (j=i+1;j<=n;j++) {
--    sum=0.0;
--    for (k=i;k<j;k++) sum -= a[j][k]*a[k][i];
--    a[j][i]=sum/p[j];
--  }
--}


Better forbidden words in Elm

Better forbidden words in Elm

Original commit c568ba4fac7cc548b181b3a14ed771cce1ca031b

module Elm
    ( renderer
    ) where

import Doc
import IRGraph
import Prelude

import Data.Array as A
import Data.Char.Unicode (isLetter)
import Data.Foldable (for_, intercalate)
import Data.List (List)
import Data.List as L
import Data.Map as Map
import Data.Maybe (Maybe(..), maybe)
import Data.Set (Set)
import Data.Set as S
import Data.String as Str
import Data.String.Util (capitalize, decapitalize, camelCase, stringEscape)
import Data.Tuple (Tuple(..), fst)
import Utils (sortByKey, sortByKeyM, forEnumerated_)

forbiddenWords :: Array String
forbiddenWords =
    [ "if", "then", "else"
    , "case", "of"
    , "let", "in"
    , "type"
    , "module", "where"
    , "import", "exposing"
    , "as"
    , "port"
    , "int", "float", "bool", "string"
    , "root", "jenc", "jdec", "jpipe", "array", "dict", "maybe"
    ]

forbiddenPropertyNames :: Set String
forbiddenPropertyNames = S.fromFoldable forbiddenWords

forbiddenNames :: Array String
forbiddenNames = A.insert "Root" $ map capitalize forbiddenWords

renderer :: Renderer
renderer =
    { name: "Elm"
    , aceMode: "elm"
    , extension: "elm"
    , doc: elmDoc
    , transforms:
        { nameForClass
        , unionName: Just unionName
        , unionPredicate: Just unionPredicate
        , nextName: \s -> "Other" <> s
        , forbiddenNames: forbiddenNames
        }
    }

nameForClass :: IRClassData -> String
nameForClass (IRClassData { names }) = upperNameStyle $ combineNames names

unionName :: L.List String -> String
unionName s =
    L.sort s
    <#> upperNameStyle
    # intercalate "Or"

unionPredicate :: IRType -> Maybe (Set IRType)
unionPredicate = case _ of
    IRUnion ur ->
        let s = unionToSet ur
        in case nullableFromSet s of
            Nothing -> Just s
            _ -> Nothing
    _ -> Nothing

isLetterCharacter :: Char -> Boolean
isLetterCharacter c =
    isLetter c || c == '_'

legalizeIdentifier :: Boolean -> String -> String
legalizeIdentifier upper str =
    case Str.charAt 0 str of
    Nothing -> "Empty"
    Just s ->
        if isLetter s then
            Str.fromCharArray $ map (\c -> if isLetterCharacter c then c else '_') $ Str.toCharArray str
        else
            legalizeIdentifier upper ((if upper then "F_" else "f_") <> str)

lowerNameStyle :: String -> String
lowerNameStyle = camelCase >>> decapitalize >>> (legalizeIdentifier false)

upperNameStyle :: String -> String
upperNameStyle = camelCase >>> capitalize >>> (legalizeIdentifier true)

renderComment :: Maybe String -> String
renderComment (Just s) = " -- " <> s
renderComment Nothing = ""

elmDoc :: Doc Unit
elmDoc = do
    line """module QuickType exposing (Root, root)

import Json.Decode as Jdec
import Json.Decode.Pipeline as Jpipe
import Array exposing (Array)
import Dict exposing (Dict)
"""
    topLevel <- getTopLevel
    { rendered: topLevelRendered, comment: topLevelComment } <- typeStringForType topLevel
    line $ "type alias Root = " <> topLevelRendered <> renderComment topLevelComment
    blank
    line "root : Jdec.Decoder Root"
    rootDecoder <- decoderNameForType topLevel
    line $ "root = " <> rootDecoder
    blank
    classes <- getClasses
    for_ classes \(Tuple i cls) -> do
        renderTypeDefinition i cls
        blank
    unions <- getUnions
    for_ unions \types -> do
        blank
        renderUnionDefinition types

noComment :: String -> Doc { rendered :: String, comment :: Maybe String }
noComment rendered =
    pure { rendered, comment: Nothing }

typeStringForUnion :: Set IRType -> Doc { rendered :: String, comment :: Maybe String }
typeStringForUnion s =
    case nullableFromSet s of
    Just x -> do
        { rendered, comment } <- typeStringForType x
        pure { rendered: "(Maybe " <> rendered <> ")", comment: Just $ maybe "optional" ("optional " <> _) comment }
    Nothing -> do
        noComment =<< lookupUnionName s

typeStringForType :: IRType -> Doc { rendered :: String, comment :: Maybe String }
typeStringForType = case _ of
    IRNothing -> noComment "Jdec.Value"
    IRNull -> noComment "()"
    IRInteger -> noComment "Int"
    IRDouble -> noComment "Float"
    IRBool -> noComment "Bool"
    IRString -> noComment "String"
    IRArray a -> do
        { rendered, comment } <- typeStringForType a
        pure { rendered: "(Array " <> rendered <> ")", comment: map ("array of " <> _) comment }
    IRClass i -> noComment =<< lookupClassName i
    IRMap t -> do
        { rendered, comment } <- typeStringForType t
        pure { rendered: "(Dict String " <> rendered <> ")", comment: map ("map to " <> _) comment }
    IRUnion types -> typeStringForUnion $ unionToSet types

lookupClassDecoderName :: Int -> Doc String
lookupClassDecoderName i = decapitalize <$> lookupClassName i

lookupUnionDecoderName :: Set IRType -> Doc String
lookupUnionDecoderName s = decapitalize <$> lookupUnionName s

unionConstructorName :: Set IRType -> IRType -> Doc String
unionConstructorName s t = do
    typeName <- upperNameStyle <$> getTypeNameForUnion t
    unionName <- lookupUnionName s
    pure $ typeName <> "In" <> unionName

decoderNameForType :: IRType -> Doc String
decoderNameForType = case _ of
    IRNothing -> pure "Jdec.value"
    IRNull -> pure "(Jdec.null ())"
    IRInteger -> pure "Jdec.int"
    IRDouble -> pure "Jdec.float"
    IRBool -> pure "Jdec.bool"
    IRString -> pure "Jdec.string"
    IRArray a -> do
        rendered <- decoderNameForType a
        pure $ "(Jdec.array " <> rendered <> ")"
    IRClass i -> lookupClassDecoderName i
    IRMap t -> do
        rendered <- decoderNameForType t
        pure $ "(Jdec.dict " <> rendered <> ")"
    IRUnion types ->
        case nullableFromSet $ unionToSet types of
        Just t -> do
            rendered <- decoderNameForType t
            pure $ "(Jdec.nullable " <> rendered <> ")"
        Nothing -> do
            lookupUnionDecoderName $ unionToSet types

forWithPrefix_ :: forall a b p m. Applicative m => List a -> p -> p -> (p -> a -> m b) -> m Unit
forWithPrefix_ l firstPrefix restPrefix f =
    forEnumerated_ l (\i -> f $ if i == 0 then firstPrefix else restPrefix)

isOptional :: IRType -> Boolean
isOptional = case _ of
    IRUnion u ->
        case nullableFromSet $ unionToSet u of
        Just t -> true
        Nothing -> false
    -- IRNull -> true
    -- IRUnion u -> S.member IRNull $ unionToSet u
    _ -> false

renderTypeDefinition :: Int -> IRClassData -> Doc Unit
renderTypeDefinition classIndex (IRClassData { names, properties }) = do
    className <- lookupClassName classIndex
    let propertyNames = transformNames lowerNameStyle (\n -> "other" <> capitalize n) forbiddenPropertyNames $ map (\n -> Tuple n n) $ Map.keys properties
    let propsList = Map.toUnfoldable properties # sortByKey (\t -> lookupName (fst t) propertyNames)
    line $ "type alias " <> className <> " ="
    indent do
        forWithPrefix_ propsList "{ " ", " \braceOrComma (Tuple pname ptype) -> do
            let propName = lookupName pname propertyNames
            { rendered, comment } <- typeStringForType ptype
            line $ braceOrComma <> propName <> " : " <> rendered <> renderComment comment
        when (propsList == L.Nil) do
            line "{"
        line "}"
    blank
    decoderName <- lookupClassDecoderName classIndex
    line $ decoderName <> " : Jdec.Decoder " <> className
    line $ decoderName <> " ="
    indent do
        line $ "Jpipe.decode " <> className
        for_ propsList \(Tuple pname ptype) -> do
            indent do
                decoderName <- decoderNameForType ptype
                let { reqOrOpt, fallback } = if isOptional ptype then { reqOrOpt: "Jpipe.optional", fallback: " Nothing" } else { reqOrOpt: "Jpipe.required", fallback: "" }
                line $ "|> " <> reqOrOpt <> " \"" <> stringEscape pname <> "\" " <> decoderName <> fallback

renderUnionDefinition :: Set IRType -> Doc Unit
renderUnionDefinition allTypes = do
    name <- lookupUnionName allTypes
    fields <- L.fromFoldable allTypes # sortByKeyM (unionConstructorName allTypes)
    line $ "type " <> name
    forWithPrefix_ fields "=" "|" \equalsOrPipe t -> do
        indent do
            constructor <- unionConstructorName allTypes t
            when (t == IRNull) do
                line $ equalsOrPipe <> " " <> constructor
            unless (t == IRNull) do
                { rendered, comment } <- typeStringForType t
                line $ equalsOrPipe <> " " <> constructor <> " " <> rendered <> renderComment comment
    blank
    decoderName <- lookupUnionDecoderName allTypes
    line $ decoderName <> " : Jdec.Decoder " <> name
    line $ decoderName <> " ="
    indent do
        line "Jdec.oneOf"
        indent do
            forWithPrefix_ fields "[" "," \bracketOrComma t -> do
                constructor <- unionConstructorName allTypes t
                when (t == IRNull) do
                    line $ bracketOrComma <> " Jdec.null " <> constructor
                unless (t == IRNull) do
                    decoder <- decoderNameForType t
                    line $ bracketOrComma <> " Jdec.map " <> constructor <> " " <> decoder
            line "]"

module Elm
    ( renderer
    ) where

import Doc
import IRGraph
import Prelude

import Data.Array as A
import Data.Char.Unicode (isLetter)
import Data.Foldable (for_, intercalate)
import Data.List (List)
import Data.List as L
import Data.Map as Map
import Data.Maybe (Maybe(..), maybe)
import Data.Set (Set)
import Data.Set as S
import Data.String as Str
import Data.String.Util (capitalize, decapitalize, camelCase, stringEscape)
import Data.Tuple (Tuple(..), fst)
import Utils (sortByKey, sortByKeyM, forEnumerated_)

keywords :: Array String
keywords =
    [ "if", "then", "else"
    , "case", "of"
    , "let", "in"
    , "type"
    , "module", "where"
    , "import", "exposing"
    , "as"
    , "port"
    ]

forbiddenPropertyNames :: Set String
forbiddenPropertyNames = S.fromFoldable keywords

forbiddenNames :: Array String
forbiddenNames = A.insert "Root" $ map capitalize keywords

renderer :: Renderer
renderer =
    { name: "Elm"
    , aceMode: "elm"
    , extension: "elm"
    , doc: elmDoc
    , transforms:
        { nameForClass
        , unionName: Just unionName
        , unionPredicate: Just unionPredicate
        , nextName: \s -> "Other" <> s
        , forbiddenNames: forbiddenNames
        }
    }

nameForClass :: IRClassData -> String
nameForClass (IRClassData { names }) = upperNameStyle $ combineNames names

unionName :: L.List String -> String
unionName s =
    L.sort s
    <#> upperNameStyle
    # intercalate "Or"

unionPredicate :: IRType -> Maybe (Set IRType)
unionPredicate = case _ of
    IRUnion ur ->
        let s = unionToSet ur
        in case nullableFromSet s of
            Nothing -> Just s
            _ -> Nothing
    _ -> Nothing

isLetterCharacter :: Char -> Boolean
isLetterCharacter c =
    isLetter c || c == '_'

legalizeIdentifier :: Boolean -> String -> String
legalizeIdentifier upper str =
    case Str.charAt 0 str of
    Nothing -> "Empty"
    Just s ->
        if isLetter s then
            Str.fromCharArray $ map (\c -> if isLetterCharacter c then c else '_') $ Str.toCharArray str
        else
            legalizeIdentifier upper ((if upper then "F_" else "f_") <> str)

lowerNameStyle :: String -> String
lowerNameStyle = camelCase >>> decapitalize >>> (legalizeIdentifier false)

upperNameStyle :: String -> String
upperNameStyle = camelCase >>> capitalize >>> (legalizeIdentifier true)

renderComment :: Maybe String -> String
renderComment (Just s) = " -- " <> s
renderComment Nothing = ""

elmDoc :: Doc Unit
elmDoc = do
    line """module QuickType exposing (Root, root)

import Json.Decode exposing (Decoder, Value, int, float, bool, string, array, dict, null, nullable, value, map, oneOf)
import Json.Decode.Pipeline exposing (decode, required, optional)
import Array exposing (Array)
import Dict exposing (Dict)
"""
    topLevel <- getTopLevel
    { rendered: topLevelRendered, comment: topLevelComment } <- typeStringForType topLevel
    line $ "type alias Root = " <> topLevelRendered <> renderComment topLevelComment
    blank
    line "root : Decoder Root"
    rootDecoder <- decoderNameForType topLevel
    line $ "root = " <> rootDecoder
    blank
    classes <- getClasses
    for_ classes \(Tuple i cls) -> do
        renderTypeDefinition i cls
        blank
    unions <- getUnions
    for_ unions \types -> do
        blank
        renderUnionDefinition types

noComment :: String -> Doc { rendered :: String, comment :: Maybe String }
noComment rendered =
    pure { rendered, comment: Nothing }

typeStringForUnion :: Set IRType -> Doc { rendered :: String, comment :: Maybe String }
typeStringForUnion s =
    case nullableFromSet s of
    Just x -> do
        { rendered, comment } <- typeStringForType x
        pure { rendered: "(Maybe " <> rendered <> ")", comment: Just $ maybe "optional" ("optional " <> _) comment }
    Nothing -> do
        noComment =<< lookupUnionName s

typeStringForType :: IRType -> Doc { rendered :: String, comment :: Maybe String }
typeStringForType = case _ of
    IRNothing -> noComment "Value"
    IRNull -> noComment "()"
    IRInteger -> noComment "Int"
    IRDouble -> noComment "Float"
    IRBool -> noComment "Bool"
    IRString -> noComment "String"
    IRArray a -> do
        { rendered, comment } <- typeStringForType a
        pure { rendered: "(Array " <> rendered <> ")", comment: map ("array of " <> _) comment }
    IRClass i -> noComment =<< lookupClassName i
    IRMap t -> do
        { rendered, comment } <- typeStringForType t
        pure { rendered: "(Dict String " <> rendered <> ")", comment: map ("map to " <> _) comment }
    IRUnion types -> typeStringForUnion $ unionToSet types

lookupClassDecoderName :: Int -> Doc String
lookupClassDecoderName i = decapitalize <$> lookupClassName i

lookupUnionDecoderName :: Set IRType -> Doc String
lookupUnionDecoderName s = decapitalize <$> lookupUnionName s

unionConstructorName :: Set IRType -> IRType -> Doc String
unionConstructorName s t = do
    typeName <- upperNameStyle <$> getTypeNameForUnion t
    unionName <- lookupUnionName s
    pure $ typeName <> "In" <> unionName

decoderNameForType :: IRType -> Doc String
decoderNameForType = case _ of
    IRNothing -> pure "value"
    IRNull -> pure "(null ())"
    IRInteger -> pure "int"
    IRDouble -> pure "float"
    IRBool -> pure "bool"
    IRString -> pure "string"
    IRArray a -> do
        rendered <- decoderNameForType a
        pure $ "(array " <> rendered <> ")"
    IRClass i -> lookupClassDecoderName i
    IRMap t -> do
        rendered <- decoderNameForType t
        pure $ "(dict " <> rendered <> ")"
    IRUnion types ->
        case nullableFromSet $ unionToSet types of
        Just t -> do
            rendered <- decoderNameForType t
            pure $ "(nullable " <> rendered <> ")"
        Nothing -> do
            lookupUnionDecoderName $ unionToSet types

forWithPrefix_ :: forall a b p m. Applicative m => List a -> p -> p -> (p -> a -> m b) -> m Unit
forWithPrefix_ l firstPrefix restPrefix f =
    forEnumerated_ l (\i -> f $ if i == 0 then firstPrefix else restPrefix)

isOptional :: IRType -> Boolean
isOptional = case _ of
    IRUnion u ->
        case nullableFromSet $ unionToSet u of
        Just t -> true
        Nothing -> false
    -- IRNull -> true
    -- IRUnion u -> S.member IRNull $ unionToSet u
    _ -> false

renderTypeDefinition :: Int -> IRClassData -> Doc Unit
renderTypeDefinition classIndex (IRClassData { names, properties }) = do
    className <- lookupClassName classIndex
    let propertyNames = transformNames lowerNameStyle (\n -> "other" <> capitalize n) forbiddenPropertyNames $ map (\n -> Tuple n n) $ Map.keys properties
    let propsList = Map.toUnfoldable properties # sortByKey (\t -> lookupName (fst t) propertyNames)
    line $ "type alias " <> className <> " ="
    indent do
        forWithPrefix_ propsList "{ " ", " \braceOrComma (Tuple pname ptype) -> do
            let propName = lookupName pname propertyNames
            { rendered, comment } <- typeStringForType ptype
            line $ braceOrComma <> propName <> " : " <> rendered <> renderComment comment
        when (propsList == L.Nil) do
            line "{"
        line "}"
    blank
    decoderName <- lookupClassDecoderName classIndex
    line $ decoderName <> " : Decoder " <> className
    line $ decoderName <> " ="
    indent do
        line $ "decode " <> className
        for_ propsList \(Tuple pname ptype) -> do
            indent do
                decoderName <- decoderNameForType ptype
                let { reqOrOpt, fallback } = if isOptional ptype then { reqOrOpt: "optional", fallback: " Nothing" } else { reqOrOpt: "required", fallback: "" }
                line $ "|> " <> reqOrOpt <> " \"" <> stringEscape pname <> "\" " <> decoderName <> fallback

renderUnionDefinition :: Set IRType -> Doc Unit
renderUnionDefinition allTypes = do
    name <- lookupUnionName allTypes
    fields <- L.fromFoldable allTypes # sortByKeyM (unionConstructorName allTypes)
    line $ "type " <> name
    forWithPrefix_ fields "=" "|" \equalsOrPipe t -> do
        indent do
            constructor <- unionConstructorName allTypes t
            when (t == IRNull) do
                line $ equalsOrPipe <> " " <> constructor
            unless (t == IRNull) do
                { rendered, comment } <- typeStringForType t
                line $ equalsOrPipe <> " " <> constructor <> " " <> rendered <> renderComment comment
    blank
    decoderName <- lookupUnionDecoderName allTypes
    line $ decoderName <> " : Decoder " <> name
    line $ decoderName <> " ="
    indent do
        line "oneOf"
        indent do
            forWithPrefix_ fields "[" "," \bracketOrComma t -> do
                constructor <- unionConstructorName allTypes t
                when (t == IRNull) do
                    line $ bracketOrComma <> " null " <> constructor
                unless (t == IRNull) do
                    decoder <- decoderNameForType t
                    line $ bracketOrComma <> " map " <> constructor <> " " <> decoder
            line "]"


No need for kind signature

No need for kind signature

module Run
  ( Run(..)
  , lift
  , send
  , extract
  , interpret
  , run
  , interpretRec
  , runRec
  , runCont
  , runAccum
  , runAccumRec
  , runAccumCont
  , peel
  , resume
  , expand
  , EFF
  , AFF
  , liftEff
  , liftAff
  , runBaseEff
  , runBaseAff
  , runBaseAff'
  , module Data.Functor.Variant
  , module Exports
  ) where

import Prelude

import Control.Monad.Aff (Aff)
import Control.Monad.Aff.Class (class MonadAff)
import Control.Monad.Eff (Eff)
import Control.Monad.Eff.Class (class MonadEff)
import Control.Monad.Eff.Class as Eff
import Control.Monad.Free (Free, liftF, runFree, runFreeM, resume')
import Control.Monad.Rec.Class (class MonadRec, Step(..), tailRecM)
import Data.Either (Either(..))
import Data.Functor.Variant (FProxy(..), VariantF, case_, default, inj, match, on, onMatch)
import Data.Newtype (class Newtype, unwrap)
import Data.Symbol (SProxy(..)) as Exports
import Data.Symbol (SProxy(..), class IsSymbol)
import Data.Tuple (Tuple, curry, uncurry)
import Partial.Unsafe (unsafeCrashWith)
import Type.Equality (class TypeEquals)
import Type.Row (RProxy)
import Unsafe.Coerce (unsafeCoerce)

-- | An extensible effect Monad, indexed by a set of effect functors. Effects
-- | are eliminated by interpretation into a pure value or into some base
-- | effect Monad.
-- |
-- | An example using `State` and `Except`:
-- | ```purescript
-- | type MyEffects =
-- |   ( state ∷ STATE Int
-- |   , except ∷ EXCEPT String
-- |   , eff ∷ EFF (console ∷ CONSOLE)
-- |   )
-- |
-- | yesProgram ∷ Run MyEffects Unit
-- | yesProgram = do
-- |   whenM (gets (_ < 0)) do
-- |     throw "Number is less than 0"
-- |   whileM_ (gets (_ > 0)) do
-- |     liftEff $ log "Yes"
-- |     modify (_ - 1)
-- |
-- | main =
-- |   yesProgram
-- |     # catch (liftEff <<< log)
-- |     # runState 10
-- |     # runBaseEff
-- |     # void
-- | ````
newtype Run r a = Run (Free (VariantF r) a)

derive instance newtypeRun ∷ Newtype (Run r a) _
derive newtype instance functorRun :: Functor (Run r)
derive newtype instance applyRun :: Apply (Run r)
derive newtype instance applicativeRun :: Applicative (Run r)
derive newtype instance bindRun :: Bind (Run r)
derive newtype instance monadRun :: Monad (Run r)

-- | This instance is provided for compatibility, but is otherwise
-- | unnecessary. You can use monadic recursion with `Run`, deferring the
-- | `MonadRec` constraint till it is interpretted.
instance monadRecRun ∷ MonadRec (Run r) where
  tailRecM f = loop
    where
    loop a = do
      b ← f a
      case b of
        Done r → pure r
        Loop n → loop n

-- | Lifts an effect functor into the `Run` Monad according to the provided
-- | `SProxy` slot.
lift
  ∷ ∀ sym r1 r2 f a
  . RowCons sym (FProxy f) r1 r2
  ⇒ IsSymbol sym
  ⇒ Functor f
  ⇒ SProxy sym
  → f a
  → Run r2 a
lift p = Run <<< liftF <<< inj p

-- | Reflects the next instruction or the final value if there are no more
-- | instructions.
peel
  ∷ ∀ a r
  . Run r a
  → Either (VariantF r (Run r a)) a
peel = resume Left Right

-- | Eliminator for the `Run` data type.
resume
  ∷ ∀ a b r
  . (VariantF r (Run r a) → b)
  → (a → b)
  → Run r a
  → b
resume k1 k2 = resume' (\x f → k1 (Run <<< f <$> x)) k2 <<< unwrap

-- | Enqueues an instruction in the `Run` Monad.
send
  ∷ ∀ a r
  . VariantF r a
  → Run r a
send = Run <<< liftF

-- | Casts some set of effects to a wider set of effects via a left-biased
-- | union. For example, you could take a closed effect and unify it with
-- | a superset of effects because we know the additional effects never
-- | occur.
-- |
-- | ```purescript
-- | type LessRows = (foo :: FOO)
-- | type MoreRows = (foo :: FOO, bar :: BAR, baz :: BAZ)
-- |
-- | foo :: Run LessRows Unit
-- | foo = foo
-- |
-- | foo' :: Run MoreRows Unit
-- | foo' = expand foo
-- | ```
expand
  ∷ ∀ r1 r2 rx a
  . Union r1 rx r2
  ⇒ Run r1 a
  → Run r2 a
expand = unsafeCoerce

-- | Extracts the value from a purely interpreted program.
extract ∷ ∀ a. Run () a → a
extract = unwrap >>> runFree \_ → unsafeCrashWith "Run: the impossible happened"

-- | Extracts the value from a program via some Monad `m`. This assumes
-- | stack safety under Monadic recursion.
interpret
  ∷ ∀ m a r
  . Monad m
  ⇒ (VariantF r ~> m)
  → Run r a
  → m a
interpret = run

-- | Identical to `interpret` but with a less restrictive type signature,
-- | letting you intercept the rest of the program.
run
  ∷ ∀ m a r
  . Monad m
  ⇒ (VariantF r (Run r a) → m (Run r a))
  → Run r a
  → m a
run k = loop
  where
  loop ∷ Run r a → m a
  loop = resume (\a → loop =<< k a) pure

-- | Extracts the value from a program via some MonadRec `m`, preserving
-- | stack safety.
interpretRec
  ∷ ∀ m a r
  . MonadRec m
  ⇒ (VariantF r ~> m)
  → Run r a
  → m a
interpretRec = runRec

-- | Identical to `interpretRec` but with a less restrictive type
-- | signature, letting you intercept the rest of the program.
runRec
  ∷ ∀ m a r
  . MonadRec m
  ⇒ (VariantF r (Run r a) → m (Run r a))
  → Run r a
  → m a
runRec k = runFreeM (coerceM k) <<< unwrap
  where
  -- Just so we can avoid the overhead of mapping the Run constructor
  coerceM ∷ (VariantF r (Run r a) -> m (Run r a)) -> VariantF r (Free (VariantF r) a) -> m (Free (VariantF r) a)
  coerceM = unsafeCoerce

-- | Extracts the value from a program via some `m` using continuation passing.
runCont
  ∷ ∀ m a b r
  . (VariantF r (m b) → m b)
  → (a → m b)
  → Run r a
  → m b
runCont k1 k2 = loop
  where
  loop ∷ Run r a → m b
  loop = resume (\b -> k1 (loop <$> b)) k2

-- | Extracts the value from a program via some Monad `m` with an internal
-- | accumulator. This assumes stack safety under Monadic recursion.
runAccum
  ∷ ∀ m r s a
  . Monad m
  ⇒ (s → VariantF r (Run r a) → m (Tuple s (Run r a)))
  → s
  → Run r a
  → m a
runAccum k = loop
  where
  loop ∷ s → Run r a → m a
  loop s = resume (\b → uncurry loop =<< k s b) pure

-- | Extracts the value from a program via some MonadRec `m` with an internal
-- | accumulator.
runAccumRec
  ∷ ∀ m r s a
  . MonadRec m
  ⇒ (s → VariantF r (Run r a) → m (Tuple s (Run r a)))
  → s
  → Run r a
  → m a
runAccumRec k = curry (tailRecM (uncurry loop))
  where
  loop ∷ s → Run r a → m (Step (Tuple s (Run r a)) a)
  loop s = resume (\b → Loop <$> k s b) (pure <<< Done)

-- | Extracts the value from a program via some `m` using continuation passing
-- | with an internal accumulator.
runAccumCont
  ∷ ∀ m r s a b
  . (s → VariantF r (s → m b) → m b)
  → (s → a → m b)
  → s
  → Run r a
  → m b
runAccumCont k1 k2 = loop
  where
  loop ∷ s → Run r a → m b
  loop s = resume (\b → k1 s (flip loop <$> b)) (k2 s)

-- | Type synonym for using `Eff` as an effect.
type EFF eff = FProxy (Eff eff)

-- Lift an `Eff` effect into the `Run` Monad via the `eff` label.
liftEff ∷ ∀ eff r. Eff eff ~> Run (eff ∷ EFF eff | r)
liftEff = lift (SProxy ∷ SProxy "eff")

-- | Runs a base `Eff` effect.
runBaseEff ∷ ∀ eff. Run (eff ∷ EFF eff) ~> Eff eff
runBaseEff = runRec $ match { eff: \a → a }

-- | Type synonym for using `Aff` as an effect.
type AFF eff = FProxy (Aff eff)

-- | Lift an `Aff` effect into the `Run` Monad via the `aff` label.
liftAff ∷ ∀ eff r. Aff eff ~> Run (aff ∷ AFF eff | r)
liftAff = lift (SProxy ∷ SProxy "aff")

-- | Runs a base `Aff` effect.
runBaseAff ∷ ∀ eff. Run (aff ∷ AFF eff) ~> Aff eff
runBaseAff = run $ match { aff: \a → a }

-- | Runs base `Aff` and `Eff` together as one effect.
runBaseAff' ∷ ∀ eff. Run (aff ∷ AFF eff, eff ∷ EFF eff) ~> Aff eff
runBaseAff' = run $ match { aff: \a → a, eff: \a → Eff.liftEff a }

instance runMonadEff ∷ (TypeEquals (RProxy r1) (RProxy (eff ∷ EFF eff | r2))) ⇒ MonadEff eff (Run r1) where
  liftEff = coerceEff <<< liftEff
    where
    coerceEff ∷ Run (eff ∷ EFF eff | r2) ~> Run r1
    coerceEff = unsafeCoerce

-- | This will insert an `EFF` effect because `MonadAff` entails `MonadEff`.
-- | If you don't want this, use `Run.liftAff` rather than `Control.Monad.Aff.Class.liftAff`.
instance runMonadAff ∷ (TypeEquals (RProxy r1) (RProxy (aff ∷ AFF eff, eff ∷ EFF eff | r2))) ⇒ MonadAff eff (Run r1) where
  liftAff = coerceAff <<< liftAff
    where
    coerceAff ∷ Run (aff ∷ AFF eff, eff ∷ EFF eff | r2) ~> Run r1
    coerceAff = unsafeCoerce

module Run
  ( Run(..)
  , lift
  , send
  , extract
  , interpret
  , run
  , interpretRec
  , runRec
  , runCont
  , runAccum
  , runAccumRec
  , runAccumCont
  , peel
  , resume
  , expand
  , EFF
  , AFF
  , liftEff
  , liftAff
  , runBaseEff
  , runBaseAff
  , runBaseAff'
  , module Data.Functor.Variant
  , module Exports
  ) where

import Prelude

import Control.Monad.Aff (Aff)
import Control.Monad.Aff.Class (class MonadAff)
import Control.Monad.Eff (Eff)
import Control.Monad.Eff.Class (class MonadEff)
import Control.Monad.Eff.Class as Eff
import Control.Monad.Free (Free, liftF, runFree, runFreeM, resume')
import Control.Monad.Rec.Class (class MonadRec, Step(..), tailRecM)
import Data.Either (Either(..))
import Data.Functor.Variant (FProxy(..), VariantF, case_, default, inj, match, on, onMatch)
import Data.Newtype (class Newtype, unwrap)
import Data.Symbol (SProxy(..)) as Exports
import Data.Symbol (SProxy(..), class IsSymbol)
import Data.Tuple (Tuple, curry, uncurry)
import Partial.Unsafe (unsafeCrashWith)
import Type.Equality (class TypeEquals)
import Type.Row (RProxy)
import Unsafe.Coerce (unsafeCoerce)

-- | An extensible effect Monad, indexed by a set of effect functors. Effects
-- | are eliminated by interpretation into a pure value or into some base
-- | effect Monad.
-- |
-- | An example using `State` and `Except`:
-- | ```purescript
-- | type MyEffects =
-- |   ( state ∷ STATE Int
-- |   , except ∷ EXCEPT String
-- |   , eff ∷ EFF (console ∷ CONSOLE)
-- |   )
-- |
-- | yesProgram ∷ Run MyEffects Unit
-- | yesProgram = do
-- |   whenM (gets (_ < 0)) do
-- |     throw "Number is less than 0"
-- |   whileM_ (gets (_ > 0)) do
-- |     liftEff $ log "Yes"
-- |     modify (_ - 1)
-- |
-- | main =
-- |   yesProgram
-- |     # catch (liftEff <<< log)
-- |     # runState 10
-- |     # runBaseEff
-- |     # void
-- | ````
newtype Run (r ∷ # Type) a = Run (Free (VariantF r) a)

derive instance newtypeRun ∷ Newtype (Run r a) _
derive newtype instance functorRun :: Functor (Run r)
derive newtype instance applyRun :: Apply (Run r)
derive newtype instance applicativeRun :: Applicative (Run r)
derive newtype instance bindRun :: Bind (Run r)
derive newtype instance monadRun :: Monad (Run r)

-- | This instance is provided for compatibility, but is otherwise
-- | unnecessary. You can use monadic recursion with `Run`, deferring the
-- | `MonadRec` constraint till it is interpretted.
instance monadRecRun ∷ MonadRec (Run r) where
  tailRecM f = loop
    where
    loop a = do
      b ← f a
      case b of
        Done r → pure r
        Loop n → loop n

-- | Lifts an effect functor into the `Run` Monad according to the provided
-- | `SProxy` slot.
lift
  ∷ ∀ sym r1 r2 f a
  . RowCons sym (FProxy f) r1 r2
  ⇒ IsSymbol sym
  ⇒ Functor f
  ⇒ SProxy sym
  → f a
  → Run r2 a
lift p = Run <<< liftF <<< inj p

-- | Reflects the next instruction or the final value if there are no more
-- | instructions.
peel
  ∷ ∀ a r
  . Run r a
  → Either (VariantF r (Run r a)) a
peel = resume Left Right

-- | Eliminator for the `Run` data type.
resume
  ∷ ∀ a b r
  . (VariantF r (Run r a) → b)
  → (a → b)
  → Run r a
  → b
resume k1 k2 = resume' (\x f → k1 (Run <<< f <$> x)) k2 <<< unwrap

-- | Enqueues an instruction in the `Run` Monad.
send
  ∷ ∀ a r
  . VariantF r a
  → Run r a
send = Run <<< liftF

-- | Casts some set of effects to a wider set of effects via a left-biased
-- | union. For example, you could take a closed effect and unify it with
-- | a superset of effects because we know the additional effects never
-- | occur.
-- |
-- | ```purescript
-- | type LessRows = (foo :: FOO)
-- | type MoreRows = (foo :: FOO, bar :: BAR, baz :: BAZ)
-- |
-- | foo :: Run LessRows Unit
-- | foo = foo
-- |
-- | foo' :: Run MoreRows Unit
-- | foo' = expand foo
-- | ```
expand
  ∷ ∀ r1 r2 rx a
  . Union r1 rx r2
  ⇒ Run r1 a
  → Run r2 a
expand = unsafeCoerce

-- | Extracts the value from a purely interpreted program.
extract ∷ ∀ a. Run () a → a
extract = unwrap >>> runFree \_ → unsafeCrashWith "Run: the impossible happened"

-- | Extracts the value from a program via some Monad `m`. This assumes
-- | stack safety under Monadic recursion.
interpret
  ∷ ∀ m a r
  . Monad m
  ⇒ (VariantF r ~> m)
  → Run r a
  → m a
interpret = run

-- | Identical to `interpret` but with a less restrictive type signature,
-- | letting you intercept the rest of the program.
run
  ∷ ∀ m a r
  . Monad m
  ⇒ (VariantF r (Run r a) → m (Run r a))
  → Run r a
  → m a
run k = loop
  where
  loop ∷ Run r a → m a
  loop = resume (\a → loop =<< k a) pure

-- | Extracts the value from a program via some MonadRec `m`, preserving
-- | stack safety.
interpretRec
  ∷ ∀ m a r
  . MonadRec m
  ⇒ (VariantF r ~> m)
  → Run r a
  → m a
interpretRec = runRec

-- | Identical to `interpretRec` but with a less restrictive type
-- | signature, letting you intercept the rest of the program.
runRec
  ∷ ∀ m a r
  . MonadRec m
  ⇒ (VariantF r (Run r a) → m (Run r a))
  → Run r a
  → m a
runRec k = runFreeM (coerceM k) <<< unwrap
  where
  -- Just so we can avoid the overhead of mapping the Run constructor
  coerceM ∷ (VariantF r (Run r a) -> m (Run r a)) -> VariantF r (Free (VariantF r) a) -> m (Free (VariantF r) a)
  coerceM = unsafeCoerce

-- | Extracts the value from a program via some `m` using continuation passing.
runCont
  ∷ ∀ m a b r
  . (VariantF r (m b) → m b)
  → (a → m b)
  → Run r a
  → m b
runCont k1 k2 = loop
  where
  loop ∷ Run r a → m b
  loop = resume (\b -> k1 (loop <$> b)) k2

-- | Extracts the value from a program via some Monad `m` with an internal
-- | accumulator. This assumes stack safety under Monadic recursion.
runAccum
  ∷ ∀ m r s a
  . Monad m
  ⇒ (s → VariantF r (Run r a) → m (Tuple s (Run r a)))
  → s
  → Run r a
  → m a
runAccum k = loop
  where
  loop ∷ s → Run r a → m a
  loop s = resume (\b → uncurry loop =<< k s b) pure

-- | Extracts the value from a program via some MonadRec `m` with an internal
-- | accumulator.
runAccumRec
  ∷ ∀ m r s a
  . MonadRec m
  ⇒ (s → VariantF r (Run r a) → m (Tuple s (Run r a)))
  → s
  → Run r a
  → m a
runAccumRec k = curry (tailRecM (uncurry loop))
  where
  loop ∷ s → Run r a → m (Step (Tuple s (Run r a)) a)
  loop s = resume (\b → Loop <$> k s b) (pure <<< Done)

-- | Extracts the value from a program via some `m` using continuation passing
-- | with an internal accumulator.
runAccumCont
  ∷ ∀ m r s a b
  . (s → VariantF r (s → m b) → m b)
  → (s → a → m b)
  → s
  → Run r a
  → m b
runAccumCont k1 k2 = loop
  where
  loop ∷ s → Run r a → m b
  loop s = resume (\b → k1 s (flip loop <$> b)) (k2 s)

-- | Type synonym for using `Eff` as an effect.
type EFF eff = FProxy (Eff eff)

-- Lift an `Eff` effect into the `Run` Monad via the `eff` label.
liftEff ∷ ∀ eff r. Eff eff ~> Run (eff ∷ EFF eff | r)
liftEff = lift (SProxy ∷ SProxy "eff")

-- | Runs a base `Eff` effect.
runBaseEff ∷ ∀ eff. Run (eff ∷ EFF eff) ~> Eff eff
runBaseEff = runRec $ match { eff: \a → a }

-- | Type synonym for using `Aff` as an effect.
type AFF eff = FProxy (Aff eff)

-- | Lift an `Aff` effect into the `Run` Monad via the `aff` label.
liftAff ∷ ∀ eff r. Aff eff ~> Run (aff ∷ AFF eff | r)
liftAff = lift (SProxy ∷ SProxy "aff")

-- | Runs a base `Aff` effect.
runBaseAff ∷ ∀ eff. Run (aff ∷ AFF eff) ~> Aff eff
runBaseAff = run $ match { aff: \a → a }

-- | Runs base `Aff` and `Eff` together as one effect.
runBaseAff' ∷ ∀ eff. Run (aff ∷ AFF eff, eff ∷ EFF eff) ~> Aff eff
runBaseAff' = run $ match { aff: \a → a, eff: \a → Eff.liftEff a }

instance runMonadEff ∷ (TypeEquals (RProxy r1) (RProxy (eff ∷ EFF eff | r2))) ⇒ MonadEff eff (Run r1) where
  liftEff = coerceEff <<< liftEff
    where
    coerceEff ∷ Run (eff ∷ EFF eff | r2) ~> Run r1
    coerceEff = unsafeCoerce

-- | This will insert an `EFF` effect because `MonadAff` entails `MonadEff`.
-- | If you don't want this, use `Run.liftAff` rather than `Control.Monad.Aff.Class.liftAff`.
instance runMonadAff ∷ (TypeEquals (RProxy r1) (RProxy (aff ∷ AFF eff, eff ∷ EFF eff | r2))) ⇒ MonadAff eff (Run r1) where
  liftAff = coerceAff <<< liftAff
    where
    coerceAff ∷ Run (aff ∷ AFF eff, eff ∷ EFF eff | r2) ~> Run r1
    coerceAff = unsafeCoerce