{-# LANGUAGE DeriveGeneric, DeriveAnyClass, FlexibleContexts, FlexibleInstances, TypeSynonymInstances #-}

module Neural.FF where

import Control.Monad
import Control.Monad.State.Strict
import Control.Parallel.Strategies
import Data.List.Split
import qualified Data.Ix as I
import qualified Data.Vector as V
import GHC.Generics
import Numeric.LinearAlgebra
import System.Random

import Neural.Util
import Numeric.Extra

instance Show (a -> b) where show a = "function"

data FFLayerOptions = FFLayerOptions {
  -- # of slices
  layerDepth :: Int,

  -- z_j -> y_j
  layerActivationFn :: [Matrix Double] -> [Matrix Double],

  -- z_j -> yPrime_j
  layerActivationFnDeriv :: [Matrix Double] -> [Matrix Double],

  -- Pooling functions
  pooling :: (Int,Int) -> [Matrix Double] -> [Matrix Double],
  poolingBP :: (Int,Int) -> [Matrix Double] -> [Matrix Double] -> [Matrix Double],
  poolSize :: (Int,Int)
  }

data FFLayer = FFLayer {
  weights :: [Matrix Double],
  biases :: [Matrix Double],

  -- x_j -> (w_j, b_j) -> z_j
  zFn :: [Matrix Double] -> ([Matrix Double], [Matrix Double]) -> [Matrix Double],
  yFn :: [Matrix Double] -> [Matrix Double],

  -- nabla_b_j -> x_j -> nabla_w_j
  nablaW :: [Matrix Double] -> [Matrix Double] -> [Matrix Double],

  -- delta_j -> x_j -> nabla_b_j
  nablaB :: [Matrix Double] -> [Matrix Double] -> [Matrix Double],

  -- nabla_b_j -> z_i-> w_j -> delta_i
  backprop :: [Matrix Double] -> [Matrix Double] -> [Matrix Double] -> [Matrix Double]
    } deriving (Show, Generic, NFData)

data FFNetwork = FFNetwork {
  layers :: V.Vector(FFLayer)
  } deriving (Show, Generic, NFData)

-- descs is a list of [(spawner, dims, options)]
spawnFFNetwork :: [((Int, Int) -> (Int, Int) -> FFLayerOptions -> FFLayerOptions -> IO (FFLayer),
                   (Int, Int), (Int, Int), FFLayerOptions)] -> IO (FFNetwork)
spawnFFNetwork descs = liftM (FFNetwork . V.fromList)
                       $ zipWithM (\(ls_j, dims_i, dims_j, opts_j) opts_i ->
                                    ls_j dims_i dims_j opts_i opts_j)
                       (tail descs) (map (\(_,_,_,opts_) -> opts_) descs)

-- Parallel matrix multiplication; left operand is split into rows.
(<||#>) :: (Matrix Double) -> (Matrix Double) -> (Matrix Double)
x <||#> y = fromRows $ parMap rdeepseq (<# y) $ toRows x

-- Same as above, but right operand is split into columns
(<#||>) :: (Matrix Double) -> (Matrix Double) -> (Matrix Double)
x <#||> y = fromColumns $ parMap rdeepseq (x #>) $ toColumns y

-- Parallel dense matrix/vector product.
-- (m><n) #||> (n) -> (m)
(#||>) :: (Matrix Double) -> (Vector Double) -> (Vector Double)
x #||> y = fromList $ parMap rdeepseq (<.> y) $ toRows x

-- Idk what this product is called but is sure as hell is parallel
-- splits second operand into a list
-- I should really see if this is faster than <#||>
-- (m) <.||> (n) -> (m,n)
(<.||>) :: (Vector Double) -> (Vector Double) -> (Matrix Double)
x <.||> y = fromColumns $ parMap rdeepseq (\y_i -> scale y_i x) (toList y)

-- feeds forward, collecting layer activations.
-- note y, z are reversed.
-- net -> x_1 -> ([y_n, .. y_1, x_1], [z_n, .. z_1])
ff :: FFNetwork -> [Matrix Double] -> (V.Vector([Matrix Double]), V.Vector([Matrix Double]))
ff net x = V.foldl activator (V.singleton x, V.empty) $ layers net
  where activator (y, z) (FFLayer w b zf yf _ _ _) =
          (\(y_j, z_j) -> (y_j `V.cons` y, z_j `V.cons` z)) $ (\a -> (yf a, a)) $ zf (V.head y) (w, b)

-- ff that only returns final output of the net.
-- net -> x_1 -> y_n
ffOut :: FFNetwork -> [Matrix Double] -> [Matrix Double]
ffOut net x = V.head $ fst $ ff net x

-- feeds forward, then passes error backwards
bp :: FFNetwork -> [Matrix Double] -> [Matrix Double] -> [([Matrix Double], [Matrix Double])]
bp net expected x = V.foldl passBack [(nabla_w_n, nabla_b_n)] kit
  where passBack ((nabla_w_k, nabla_b_k) : nablas) (x_j, z_j, layer_k, layer_j) =
          (nabla_w_j, nabla_b_j) : (nabla_w_k, nabla_b_k) : nablas
          where nabla_w_j = (nablaW layer_j) nabla_b_j x_j
                nabla_b_j = (nablaB layer_j) delta_j z_j
                delta_j = (backprop layer_k) nabla_b_k z_j (weights layer_k)

        -- passBack sees these args:
        -- y[n-1 .. 1], z[n-1 .. 1], layers[n .. 2], layers[n-1 .. 1]
        kit = V.zip4 (V.tail $ V.tail y) (V.tail z) (V.reverse $ V.tail $ ls) (V.reverse $ V.init ls)
        nabla_w_n = (nablaW $ V.last $ ls) nabla_b_n (y V.! 1)
        nabla_b_n = zipWith (-) (y V.! 0) expected
        ls = layers net
        (y, z) = ff net x

-- Error/cost functions.

{-
crossEntropyDelta :: [Matrix Double] -> [Matrix Double] -> [Matrix Double]
crossEntropyDelta = zipWith (-)

crossEntropyError :: [Matrix Double] -> [Matrix Double] -> [Double]
crossEntropyError x y = map (\(x,y) -> sumElements $ - ((y * (cmap log x)) + (1.0 - y) * (cmap log $ 1.0 - x)))
                        $ zip x y
-}

-- Convert an output value to an actual output of the net
-- i.e. '5' -> [0,0,0,0,0,1,0,...]
buildOutput :: Int -> Int -> [Matrix Double]
buildOutput i len = [(len><1) (map (\x -> if x == i then 1.0 else 0.0) [0..len - 1])]

-- Convert a net's output to an integer output value
-- uses max value of the net.
-- i.e. [0,0.5,0.25,-0.1,0.999999,...] -> '4'
convertOutput :: [Matrix Double] -> Int
convertOutput output = (\(i,j) -> i) $ maxIndex $ head output

-- For a given test set, return the number the net gets correct
evaluate :: FFNetwork -> [((Matrix Double), Int)] -> Double -> Int
evaluate net testSet reg =
  sum $ parMap rdeepseq
  (\(input, expected) ->
    let output = ffOut net [input]
    in if convertOutput output == expected then 1 else 0) testSet

-- perform stochastic gradient descent on the net.
-- each sample data is a tuple of (input, expected output) where the output is an integer representing the desired result.
-- maxClassIndex is the max classification ID any output can produce
-- eta is the learning rate.
-- reg is the regularization parameter (for L2 reg)
sgd :: FFNetwork -> StdGen -> [((Matrix Double), Int)] -> Int -> Int -> Double -> Double -> ((FFNetwork, [((Matrix Double), Int)]), StdGen)
sgd net gen trainingSet maxClassIndex batchSize eta reg = ((newNet, newTrainingSet), newGen)
  where newNet = foldl applyBatch net $ chunksOf batchSize newTrainingSet
        applyBatch net batch =
          net { layers =
                  V.zipWith (\layer (nabla_w, nabla_b) ->
                             layer {
                                -- apply the changes to each layer of the net.
                                weights = zipWith adjustWeights (weights layer) nabla_w,
                                biases = zipWith adjustBiases (biases layer) nabla_b
                                }) (layers net)
                  -- backpropagate each member of the batch, summing the resulting change in the net.
                  -- this will yield [(nw, nb) ...] for each layer of the net.
                  $ V.fromList
                  $ (\(batch1 : batches) ->
                      foldl (zipWith (\(nw,nb) (dnw,dnb) -> (zipWith (+) nw dnw, zipWith (+) nb dnb)))
                      batch1 batches)
                  $ parMap rdeepseq
                  (\(input, expected) ->
                    bp net (buildOutput expected $ maxClassIndex + 1) [input]) batch
              }

        -- hmatrix gets a little confused when (Num a, Matrix b) => a * b happens
        -- so use `scale` for now.
        adjustWeights w nabla_w = (scale (1.0 - eta * reg / (fromIntegral $ length trainingSet)) w)
                                  - (scale (eta / (fromIntegral batchSize)) nabla_w)
        adjustBiases b nabla_b = b - (scale (eta / (fromIntegral batchSize)) nabla_b)

        -- cost = 0.5 * reg * l2_norm_sq * (fromIntegral batchSize) / (fromIntegral $ length trainingSet)
        -- l2_norm_sq = V.sum $ V.map (\l -> sum $ parMap rdeepseq (\w -> sumElements $ w * w) $ weights l) $ layers net

        (newTrainingSet, newGen) = shuffle' trainingSet gen


-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.GSL.Vector
-- Copyright   :  (c) Alberto Ruiz 2007
-- License     :  GPL-style
--
-- Maintainer  :  Alberto Ruiz <aruiz@um.es>
-- Stability   :  provisional
-- Portability :  portable (uses FFI)
--
-- Low level interface to vector operations.
--
-----------------------------------------------------------------------------

module Numeric.GSL.Vector (
    sumF, sumR, sumQ, sumC,
    prodF, prodR, prodQ, prodC,
    dotF, dotR, dotQ, dotC,
    FunCodeS(..), toScalarR, toScalarF, toScalarC, toScalarQ,
    FunCodeV(..), vectorMapR, vectorMapC, vectorMapF, vectorMapQ,
    FunCodeSV(..), vectorMapValR, vectorMapValC, vectorMapValF, vectorMapValQ,
    FunCodeVV(..), vectorZipR, vectorZipC, vectorZipF, vectorZipQ,
    RandDist(..), randomVector
) where

import Data.Packed.Internal.Common
import Data.Packed.Internal.Signatures
import Data.Packed.Internal.Vector

import Data.Complex
import Foreign.Marshal.Alloc(free)
import Foreign.Marshal.Array(newArray)
import Foreign.Ptr(Ptr)
import Foreign.C.Types
import System.IO.Unsafe(unsafePerformIO)

fromei x = fromIntegral (fromEnum x) :: CInt

data FunCodeV = Sin
              | Cos
              | Tan
              | Abs
              | ASin
              | ACos
              | ATan
              | Sinh
              | Cosh
              | Tanh
              | ASinh
              | ACosh
              | ATanh
              | Exp
              | Log
              | Sign
              | Sqrt
              deriving Enum

data FunCodeSV = Scale
               | Recip
               | AddConstant
               | Negate
               | PowSV
               | PowVS
               deriving Enum

data FunCodeVV = Add
               | Sub
               | Mul
               | Div
               | Pow
               | ATan2
               deriving Enum

data FunCodeS = Norm2
              | AbsSum
              | MaxIdx
              | Max
              | MinIdx
              | Min
              deriving Enum

------------------------------------------------------------------

-- | sum of elements
sumF :: Vector Float -> Float
sumF x = unsafePerformIO $ do
           r <- createVector 1
           app2 c_sumF vec x vec r "sumF"
           return $ r @> 0

-- | sum of elements
sumR :: Vector Double -> Double
sumR x = unsafePerformIO $ do
           r <- createVector 1
           app2 c_sumR vec x vec r "sumR"
           return $ r @> 0

-- | sum of elements
sumQ :: Vector (Complex Float) -> Complex Float
sumQ x = unsafePerformIO $ do
           r <- createVector 1
           app2 c_sumQ vec x vec r "sumQ"
           return $ r @> 0

-- | sum of elements
sumC :: Vector (Complex Double) -> Complex Double
sumC x = unsafePerformIO $ do
           r <- createVector 1
           app2 c_sumC vec x vec r "sumC"
           return $ r @> 0

foreign import ccall unsafe "gsl-aux.h sumF" c_sumF :: TFF
foreign import ccall unsafe "gsl-aux.h sumR" c_sumR :: TVV
foreign import ccall unsafe "gsl-aux.h sumQ" c_sumQ :: TQVQV
foreign import ccall unsafe "gsl-aux.h sumC" c_sumC :: TCVCV

-- | product of elements
prodF :: Vector Float -> Float
prodF x = unsafePerformIO $ do
           r <- createVector 1
           app2 c_prodF vec x vec r "prodF"
           return $ r @> 0

-- | product of elements
prodR :: Vector Double -> Double
prodR x = unsafePerformIO $ do
           r <- createVector 1
           app2 c_prodR vec x vec r "prodR"
           return $ r @> 0

-- | product of elements
prodQ :: Vector (Complex Float) -> Complex Float
prodQ x = unsafePerformIO $ do
           r <- createVector 1
           app2 c_prodQ vec x vec r "prodQ"
           return $ r @> 0

-- | product of elements
prodC :: Vector (Complex Double) -> Complex Double
prodC x = unsafePerformIO $ do
           r <- createVector 1
           app2 c_prodC vec x vec r "prodC"
           return $ r @> 0

foreign import ccall unsafe "gsl-aux.h prodF" c_prodF :: TFF
foreign import ccall unsafe "gsl-aux.h prodR" c_prodR :: TVV
foreign import ccall unsafe "gsl-aux.h prodQ" c_prodQ :: TQVQV
foreign import ccall unsafe "gsl-aux.h prodC" c_prodC :: TCVCV

-- | dot product
dotF :: Vector Float -> Vector Float -> Float
dotF x y = unsafePerformIO $ do
           r <- createVector 1
           app3 c_dotF vec x vec y vec r "dotF"
           return $ r @> 0

-- | dot product
dotR :: Vector Double -> Vector Double -> Double
dotR x y = unsafePerformIO $ do
           r <- createVector 1
           app3 c_dotR vec x vec y vec r "dotR"
           return $ r @> 0

-- | dot product
dotQ :: Vector (Complex Float) -> Vector (Complex Float) -> Complex Float
dotQ x y = unsafePerformIO $ do
           r <- createVector 1
           app3 c_dotQ vec x vec y vec r "dotQ"
           return $ r @> 0

-- | dot product
dotC :: Vector (Complex Double) -> Vector (Complex Double) -> Complex Double
dotC x y = unsafePerformIO $ do
           r <- createVector 1
           app3 c_dotC vec x vec y vec r "dotC"
           return $ r @> 0

foreign import ccall unsafe "gsl-aux.h dotF" c_dotF :: TFFF
foreign import ccall unsafe "gsl-aux.h dotR" c_dotR :: TVVV
foreign import ccall unsafe "gsl-aux.h dotQ" c_dotQ :: TQVQVQV
foreign import ccall unsafe "gsl-aux.h dotC" c_dotC :: TCVCVCV

------------------------------------------------------------------

toScalarAux fun code v = unsafePerformIO $ do
    r <- createVector 1
    app2 (fun (fromei code)) vec v vec r "toScalarAux"
    return (r `at` 0)

vectorMapAux fun code v = unsafePerformIO $ do
    r <- createVector (dim v)
    app2 (fun (fromei code)) vec v vec r "vectorMapAux"
    return r

vectorMapValAux fun code val v = unsafePerformIO $ do
    r <- createVector (dim v)
    pval <- newArray [val]
    app2 (fun (fromei code) pval) vec v vec r "vectorMapValAux"
    free pval
    return r

vectorZipAux fun code u v = unsafePerformIO $ do
    r <- createVector (dim u)
    app3 (fun (fromei code)) vec u vec v vec r "vectorZipAux"
    return r

---------------------------------------------------------------------

-- | obtains different functions of a vector: norm1, norm2, max, min, posmax, posmin, etc.
toScalarR :: FunCodeS -> Vector Double -> Double
toScalarR oper =  toScalarAux c_toScalarR (fromei oper)

foreign import ccall unsafe "gsl-aux.h toScalarR" c_toScalarR :: CInt -> TVV

-- | obtains different functions of a vector: norm1, norm2, max, min, posmax, posmin, etc.
toScalarF :: FunCodeS -> Vector Float -> Float
toScalarF oper =  toScalarAux c_toScalarF (fromei oper)

foreign import ccall unsafe "gsl-aux.h toScalarF" c_toScalarF :: CInt -> TFF

-- | obtains different functions of a vector: only norm1, norm2
toScalarC :: FunCodeS -> Vector (Complex Double) -> Double
toScalarC oper =  toScalarAux c_toScalarC (fromei oper)

foreign import ccall unsafe "gsl-aux.h toScalarC" c_toScalarC :: CInt -> TCVV

-- | obtains different functions of a vector: only norm1, norm2
toScalarQ :: FunCodeS -> Vector (Complex Float) -> Float
toScalarQ oper =  toScalarAux c_toScalarQ (fromei oper)

foreign import ccall unsafe "gsl-aux.h toScalarQ" c_toScalarQ :: CInt -> TQVF

------------------------------------------------------------------

-- | map of real vectors with given function
vectorMapR :: FunCodeV -> Vector Double -> Vector Double
vectorMapR = vectorMapAux c_vectorMapR

foreign import ccall unsafe "gsl-aux.h mapR" c_vectorMapR :: CInt -> TVV

-- | map of complex vectors with given function
vectorMapC :: FunCodeV -> Vector (Complex Double) -> Vector (Complex Double)
vectorMapC oper = vectorMapAux c_vectorMapC (fromei oper)

foreign import ccall unsafe "gsl-aux.h mapC" c_vectorMapC :: CInt -> TCVCV

-- | map of real vectors with given function
vectorMapF :: FunCodeV -> Vector Float -> Vector Float
vectorMapF = vectorMapAux c_vectorMapF

foreign import ccall unsafe "gsl-aux.h mapF" c_vectorMapF :: CInt -> TFF

-- | map of real vectors with given function
vectorMapQ :: FunCodeV -> Vector (Complex Float) -> Vector (Complex Float)
vectorMapQ = vectorMapAux c_vectorMapQ

foreign import ccall unsafe "gsl-aux.h mapQ" c_vectorMapQ :: CInt -> TQVQV

-------------------------------------------------------------------

-- | map of real vectors with given function
vectorMapValR :: FunCodeSV -> Double -> Vector Double -> Vector Double
vectorMapValR oper = vectorMapValAux c_vectorMapValR (fromei oper)

foreign import ccall unsafe "gsl-aux.h mapValR" c_vectorMapValR :: CInt -> Ptr Double -> TVV

-- | map of complex vectors with given function
vectorMapValC :: FunCodeSV -> Complex Double -> Vector (Complex Double) -> Vector (Complex Double)
vectorMapValC = vectorMapValAux c_vectorMapValC

foreign import ccall unsafe "gsl-aux.h mapValC" c_vectorMapValC :: CInt -> Ptr (Complex Double) -> TCVCV

-- | map of real vectors with given function
vectorMapValF :: FunCodeSV -> Float -> Vector Float -> Vector Float
vectorMapValF oper = vectorMapValAux c_vectorMapValF (fromei oper)

foreign import ccall unsafe "gsl-aux.h mapValF" c_vectorMapValF :: CInt -> Ptr Float -> TFF

-- | map of complex vectors with given function
vectorMapValQ :: FunCodeSV -> Complex Float -> Vector (Complex Float) -> Vector (Complex Float)
vectorMapValQ oper = vectorMapValAux c_vectorMapValQ (fromei oper)

foreign import ccall unsafe "gsl-aux.h mapValQ" c_vectorMapValQ :: CInt -> Ptr (Complex Float) -> TQVQV

-------------------------------------------------------------------

-- | elementwise operation on real vectors
vectorZipR :: FunCodeVV -> Vector Double -> Vector Double -> Vector Double
vectorZipR = vectorZipAux c_vectorZipR

foreign import ccall unsafe "gsl-aux.h zipR" c_vectorZipR :: CInt -> TVVV

-- | elementwise operation on complex vectors
vectorZipC :: FunCodeVV -> Vector (Complex Double) -> Vector (Complex Double) -> Vector (Complex Double)
vectorZipC = vectorZipAux c_vectorZipC

foreign import ccall unsafe "gsl-aux.h zipC" c_vectorZipC :: CInt -> TCVCVCV

-- | elementwise operation on real vectors
vectorZipF :: FunCodeVV -> Vector Float -> Vector Float -> Vector Float
vectorZipF = vectorZipAux c_vectorZipF

foreign import ccall unsafe "gsl-aux.h zipF" c_vectorZipF :: CInt -> TFFF

-- | elementwise operation on complex vectors
vectorZipQ :: FunCodeVV -> Vector (Complex Float) -> Vector (Complex Float) -> Vector (Complex Float)
vectorZipQ = vectorZipAux c_vectorZipQ

foreign import ccall unsafe "gsl-aux.h zipQ" c_vectorZipQ :: CInt -> TQVQVQV

-----------------------------------------------------------------------

data RandDist = Uniform  -- ^ uniform distribution in [0,1)
              | Gaussian -- ^ normal distribution with mean zero and standard deviation one
              deriving Enum

-- | Obtains a vector of pseudorandom elements from the the mt19937 generator in GSL, with a given seed. Use randomIO to get a random seed.
randomVector :: Int      -- ^ seed
             -> RandDist -- ^ distribution
             -> Int      -- ^ vector size
             -> Vector Double
randomVector seed dist n = unsafePerformIO $ do
    r <- createVector n
    app1 (c_random_vector (fi seed) ((fi.fromEnum) dist)) vec r "randomVector"
    return r

foreign import ccall unsafe "random_vector" c_random_vector :: CInt -> CInt -> TV


module HLearn.Optimization.Conic
    where

import Control.DeepSeq
import Control.Monad
import Control.Monad.Random
import Control.Monad.ST
import Data.List
import Data.List.Extras
import Debug.Trace
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as VGM
import qualified Data.Vector.Algorithms.Intro as Intro
import Numeric.LinearAlgebra hiding ((<>))
import qualified Numeric.LinearAlgebra as LA
import Data.Random.Normal

import HLearn.Algebra
import qualified HLearn.Optimization.Common as Recipe
import qualified HLearn.Optimization.LineMinimization as LineMin

-------------------------------------------------------------------------------

conicprojection :: Matrix Double -> Matrix Double
conicprojection m = cmap realPart $ u LA.<> lambda' LA.<> trans u
    where
        lambda' = cmap (\x -> if realPart x < 0 then 0 else x) lambda
        lambda = diagRect 0 l (VG.length l) (VG.length l)
        (l,u) = eig m

-------------------------------------------------------------------------------

data RandomConicPersuit a = RandomConicPersuit
    { _stdgen :: !StdGen
    , _soln :: !a
    , _fx :: !(Scalar a)
    , _solnlast :: !a
    }

-- data OptInfo a = OptInfo
--     { _stdgen :: !StdGen
--     , _x :: !a
--     , _fx ::
--     , _xold :: !a
--     }

itr :: Int -> (tmp -> a) -> (tmp -> Bool) -> (tmp -> tmp) -> tmp -> a
itr i result stop step init = --trace ("i="++show i++"; fx="++show (_fx init)) $ --trace ("i="++show i++"; init="++show init) $
  if i==0 || stop init
    then result init
    else itr (i-1) result stop step (step init)

randomConicPersuit f x0 = _soln $ itr 10 id (\x -> False) (step_RandomConicPersuit f) init
    where
        init = RandomConicPersuit
            { _stdgen = mkStdGen $ round $ sumElements x0
            , _soln = x0
            , _fx = f x0
            , _solnlast = x0
            }

conicProjection f x0 = _soln $ argmin _fx [itr 1 id (\x -> False) (step_RandomConicPersuit f) (init i) | i <- [0..10]]
    where
        init i = RandomConicPersuit
            { _stdgen = mkStdGen $ i+(round $ sumElements x0)
            , _soln = x0
            , _fx = f x0
            , _solnlast = x0
            }

stop_RandomConicPersuit (RandomConicPersuit stdgen soln fx solnlast) = undefined

step_RandomConicPersuit f (RandomConicPersuit stdgen soln fx solnlast) = --trace ("lambda = "++show lambda++ "; phi="++show phi) $
    ret
    where
--         (x', stdgen') = runRand ((LA.fromList . take (rows soln)) `liftM` getRandomRs (-1,1)) stdgen
        normL 0 (xs,g) = (xs,g)
        normL i (xs,g) = normL (i-1) (x:xs,g')
            where
                (x,g') = normal g

        (x'std, stdgen') = normL (rows soln) ([],stdgen)
        x' = VS.fromList x'std
--         (lambda,phi::Matrix Double) = eigSH soln
--         q = (diag $ cmap (sqrt.abs) lambda) LA.<> phi
--         x' = (scale (1-kappa) q - scale kappa (ident (rows q))) LA.<> LA.fromList x'std
--         kappa = 1e-4

        y' = asColumn x' LA.<> asRow x'
--         y' = asColumn (LA.fromList x') LA.<> asRow (LA.fromList x')

        g_alpha alpha = f $ scale alpha y' + soln
        alpha_hat = error "step_randomConic" -- LineMin._x $ runOptimization $ LineMin.brent g_alpha (LineMin.lineBracket g_alpha 0 1)

        alpha_hat_y' = scale alpha_hat y'
        g_beta beta = f $ alpha_hat_y' + scale beta soln
        beta_hat = error "step_randomConic" -- LineMin._x $ LineMin.brent g_beta (LineMin.lineBracket g_beta 0 1)

        soln' = alpha_hat_y' + scale beta_hat soln

        ret = RandomConicPersuit
            { _stdgen = stdgen'
            , _soln = soln'
            , _fx = f soln'
            , _solnlast = soln
            }


{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE FlexibleContexts #-}

module Turtle where

import Control.Monad.State
-- import Control.Monad.Trans.Class
import Numeric.LinearAlgebra.HMatrix hiding (scale, (!))
import qualified Graphics.Rendering.OpenGL as GL
import Foreign.Storable
import GLCode
import Sketch
-- import Points
import Lines
import Triangles
import Data.Array
import Debug.Trace

clamp :: Float -> Float -> Float -> Float
clamp a b x | x < a = a
            | x > b = b
            | otherwise = x

lerp :: Float -> Float -> Float -> Float
lerp a b x = (1-x)*a+x*b

linstep :: Float -> Float -> Float -> Float
linstep a b x = clamp 0.0 1.0 ((x-a)/(b-a))

smoothstep :: Float -> Float -> Float -> Float
smoothstep a b x =
    let t = clamp 0.0 1.0 ((x-a)/(b-a))
    in t*t*(3.0-2.0*t)

v2f :: Real a => a -> a -> GL.Vertex2 Float
v2f x y = GL.Vertex2 (realToFrac x) (realToFrac y)

v3f :: Real a => a -> a -> a -> GL.Vertex3 Float
v3f x y z = GL.Vertex3 (realToFrac x) (realToFrac y) (realToFrac z)

v4f :: Real a => a -> a -> a -> a -> GL.Vertex4 Float
v4f x y z w = GL.Vertex4 (realToFrac x) (realToFrac y) (realToFrac z) (realToFrac w)

translation :: (Fractional a, Storable a) => a -> a -> a -> Matrix a
translation x y z = (4 >< 4) [ 1.0, 0.0, 0.0, 0.0,
                               0.0, 1.0, 0.0, 0.0,
                               0.0, 0.0, 1.0, 0.0,
                               x  , y  , z  , 1.0 ]

scaling :: (Fractional a, Storable a) => a -> a -> a -> Matrix a
scaling x y z = (4 >< 4) [ x  , 0.0, 0.0, 0.0,
                           0.0, y  , 0.0, 0.0,
                           0.0, 0.0, z  , 0.0,
                           0.0, 0.0, 0.0, 1.0 ]

rotation :: (Fractional a, Storable a, Floating a) => a -> Matrix a
rotation theta = let c = cos theta
                     s = sin theta
                 in (4 >< 4) [ c, s, 0.0, 0.0,
                               -s, c, 0.0, 0.0,
                               0.0, 0.0, 1.0, 0.0,
                               0.0, 0.0, 0.0, 1.0]

rectangle :: (Float, Float, Float) -> Float -> Float -> SketchMonad ()
rectangle (r, g, b) w h = do
    drawTriangle "turtle"
        "vPosition" (rectangleVertices w h)
        "color" [v4f r g b 1.0, v4f r g b 1.0, v4f r g b 1.0,
                 v4f r g b 1.0, v4f r g b 1.0, v4f r g b 1.0]

drawPath :: (Float, Float, Float) -> [(Float, Float)] -> SketchMonad ()
drawPath (r, g, b) vs = do
    drawLine "turtle"
        "vPosition" (map (uncurry v2f) vs)
        "color" (replicate (length vs) (v4f r g b 1.0))

plotPath :: (Float, Float, Float) -> Int -> (Float -> (Float, Float)) -> Float -> Float -> SketchMonad ()
plotPath rgb n f t0 t1 = do
    let scale = (t1-t0)/fromIntegral n
    drawPath rgb [f t | i <- [0..n],
                        let t = t0+scale*fromIntegral i]

constRectangle :: (Float, Float, Float) -> Float -> Float -> SketchMonad ()
constRectangle (r, g, b) w h = do
    drawTriangle "turtle"
        "vPosition" (rectangleVertices w h)
        "color" [v4f r g b 1.0, v4f r g b 1.0, v4f r g b 1.0,
                 v4f r g b 1.0, v4f r g b 1.0, v4f r g b 1.0]

rectangleVertices :: Float -> Float -> [GL.Vertex2 Float]
rectangleVertices w h =
        [v2f (-0.5*w) (-0.5*h), v2f (0.5*w) (-0.5*h), v2f (0.5*w) (0.5*h),
         v2f (-0.5*w) (-0.5*h), v2f (0.5*w) (0.5*h), v2f (-0.5*w) (0.5*h)]

rotate :: (Floating t, Numeric t, MonadState (Matrix t) m) => t -> m ()
rotate angle = modify (rotation angle `mul`)

scale :: (Fractional t, Numeric t, MonadState (Matrix t) m) => t -> t -> t -> m ()
scale s t u = modify (scaling s t u `mul`)

translate :: (Fractional t, Numeric t, MonadState (Matrix t) m) => t -> t -> t -> m ()
translate u v w = modify (translation u v w `mul`)

save :: StateT (Matrix Float) (StateT World IO) () -> StateT (Matrix Float) (StateT World IO) ()
save c = do
    m <- get
    a <- c
    put m
    setTransform
    return a

setTransform :: StateT (Matrix Float) (StateT World IO) ()
setTransform = do
    transform <- get
    lift $ setUniform "turtle"
               "transform" transform

setTransformPoint :: StateT (Matrix Float) (StateT World IO) ()
setTransformPoint = do
    transform <- get
    lift $ setUniform "turtle_point"
               "transform" transform

arrow :: (Float, Float, Float) -> Float -> Float -> Float -> Float -> SketchMonad ()
arrow (r, g, b) thickness headLength headWidth length = do
--     let thickness = 0.01
    constRectangle (r, g, b) length thickness
    drawTriangle "turtle"
        "vPosition" [v2f (0.5*length) (-0.5*headWidth), v2f (0.5*length+headLength) 0, v2f (0.5*length) (0.5*headWidth)]
        "color" [v4f r g b 1.0, v4f r g b 1.0,
                 v4f r g b 1.0]

grid :: Int -> Int -> Float -> Float ->
        (Int -> Int -> StateT (Matrix Float) (StateT World IO) ()) ->
        StateT (Matrix Float) (StateT World IO) ()
grid m n dx dy f =
    forM_ [0..(m-1)] $ \i ->
        forM_ [0..(n-1)] $ \j -> save $ do
            let x = dx*fromIntegral i
            let y = dy*fromIntegral j
            translate x y 0
            setTransform
            f i j

drawVectorField drawArrow m n dx dy vx vy = drawVectorField' drawArrow m n dx dy $ interpVelocityField vx vy

drawVectorField' drawArrow m n dx dy f =
    grid m n dx dy $ \i j -> do
        let (vx, vy) = f (fromIntegral i) (fromIntegral j)
        let theta = atan2 vy vx
        let r = sqrt (vx*vx+vy*vy)
        save $ do
            modify ((rotation theta) `mul`)
            transform <- get
            lift $ setUniform "turtle"
                              "transform" transform
            lift $ drawArrow r

type ColorMap = Float -> GL.Vertex4 Float

drawDensityField :: ColorMap -> Int -> Int -> Float -> Float -> Array (Int, Int) Float -> StateT (Matrix Float) (StateT World IO) ()
drawDensityField cmap m n dx dy a =
    grid m n dx dy $ \i j -> do
        let ip = (i+1) `mod` m
        let jp = (j+1) `mod` m
        lift $ drawTriangle "turtle"
                            "vPosition" [v2f 0.0 0.0, v2f dx 0.0, v2f dx dy,
                                         v2f 0.0 0.0, v2f dx dy, v2f 0.0 dy]
                            "color"     [cmap (a!(i, j)), cmap (a!(ip, j)), cmap (a!(ip, jp)),
                                         cmap (a!(i, j)), cmap (a!(ip, jp)), cmap (a!(i, jp))]

coolwarm :: Float -> Float -> Float -> GL.Vertex4 Float
coolwarm a b x =
    let mid = 0.5*(a+b)
    in if x > mid
        then let t = smoothstep mid b x in v4f 1.0 (1.0-t) (1.0-t) 1.0
        else let t = smoothstep mid a x in v4f (1.0-t) (1.0-t) 1.0 1.0

-- Interpolate velocity onto uniform grid from staggered grid
-- Let's say velocity field defined by
-- vx!(x, y) = v (x-0.5) y
-- vy!(x, y) = v x (y-0.5)
--
-- So, for example, (v x y)_x = vx!(x,y)+vx!(x+1,y)
interpVelocity :: Array (Int, Int) Float -> Array (Int, Int) Float -> Array (Int, Int) (Float, Float)
interpVelocity vx vy =
    let (_, (nx, ny)) = bounds vx
    in array ((0, 0), (nx, ny)) [((i, j), (0.5*(vx!(ip, j)+vx!(i, j)), 0.5*(vy!(i, jp)+vy!(i, j)))) |
                                 i <- [0..nx],
                                 j <- [0..ny],
                                 let ip = (i+1) `mod` (nx+1),
                                 let jp = (j+1) `mod` (ny+1)]

interpVelocityField :: Array (Int, Int) Float -> Array (Int, Int) Float -> Float -> Float -> (Float, Float)
interpVelocityField vx vy x y =
    (interpScalarField vx (x+0.5) y, interpScalarField vy x (y+0.5))
--     let (_, (nx, ny)) = bounds vx
--         x = x0+0.5
--         y = y0+0.5
--         ix = floor x `mod` nx
--         iy = floor y `mod` ny
--         ixp = (ix+1) `mod` nx
--         iyp = (iy+1) `mod` ny
--         fx = x-fromIntegral ix
--         fy = y-fromIntegral iy
--         vx00 = lerp (vx!(ix, iy)) (vy!(ixp, iy))
--     in (lerp (vx!(ix, iy)) (vx!(ixp, iy)) fx,
--         lerp (vy!(ix, iy)) (vy!(ix, iyp)) fy)

interpScalarField :: Array (Int, Int) Float -> Float -> Float -> Float
interpScalarField p x y =
    let (_, (nx, ny)) = bounds p
        ix = floor x `mod` nx
        iy = floor y `mod` ny
        ixp = (ix+1) `mod` nx
        iyp = (iy+1) `mod` ny
        fx = x-fromIntegral ix
        fy = y-fromIntegral iy
        p00 = p!(ix, iy)
        p01 = p!(ix, iyp)
        p10 = p!(ixp, iy)
        p11 = p!(ixp, iyp)
        p0 = lerp p00 p01 fy
        p1 = lerp p10 p11 fy
    in lerp p0 p1 fx

divergence :: Array (Int, Int) Float -> Array (Int, Int) Float -> Array (Int, Int) (Float, Float)
divergence vx vy =
    let (_, (nx, ny)) = bounds vx
    in array ((0, 0), (nx, ny)) [((i, j), (vx!(ip, j)-vx!(i, j), vy!(i, jp)-vy!(i, j))) |
                                 i <- [0..nx],
                                 j <- [0..ny],
                                 let ip = (i+1) `mod` (nx+1),
                                 let jp = (j+1) `mod` (ny+1)]

grad :: Array (Int, Int) Float -> (Array (Int, Int) Float, Array (Int, Int) Float)
grad p =
    let (_, (nx, ny)) = bounds p
        vx = array ((0, 0), (nx, ny)) [((i, j), p!(ip, j)-p!(i, j)) |
                                      i <- [0..nx],
                                      j <- [0..ny],
                                      let ip = (i+1) `mod` (nx+1)]
        vy = array ((0, 0), (nx, ny)) [((ix, iy), p!(ix, iyp)-p!(ix, iy)) |
                                      ix <- [0..nx],
                                      iy <- [0..ny],
                                      let iyp = (iy+1) `mod` (ny+1)]
    in (vx, vy)

integrate :: (Float -> Float -> (Float, Float)) -> Float -> Float -> Int -> Float -> [(Float, Float)]
integrate v x y 0 dt = []
integrate v x y n dt = (x, y) :
    let (vx, vy) = v x y
    in integrate v (x+dt*vx) (y+dt*vy) (n-1) dt

duringAndAfter :: Monad m => Float -> Float -> Float -> (Float -> m ()) -> m ()
duringAndAfter t t0 t1 m =
    if t < t0
        then return ()
        else if t > t1
            then m 1
            else m ((t-t0)/(t1-t0))

beforeAndDuring :: Monad m => Float -> Float -> Float -> (Float -> m ()) -> m ()
beforeAndDuring t t0 t1 m =
    if t > t1
        then return ()
        else if t < t0
            then m 1
            else m ((t-t1)/(t0-t1))


-- | Now we plot!

module QuickHullDraw where

import           Control.Arrow
import           Data.List
import qualified Data.Vector             as V
import qualified Data.Vector.Unboxed     as UV
import           Graphics.Rendering.Plot
import           Numeric.LinearAlgebra   (Vector (..))
import           QuickHull
import           System.Environment
import           System.Random

main :: IO ()
main = do
  [d] <- getArgs
  let vectors =
        V.fromList .
        map UV.fromList .
        chop 2 . take ((* 2) $ read d :: Int) . randomRs (0 :: Double, 100) $
          mkStdGen 0
      (rxL, rxH) = (0, 100)
      (ryL, ryH) = (0, 100)
      (vx, vy) = convertPoints vectors
      hull' = qhull2D vectors
      (hx, hy) =
        convertPoints . V.fromList $ hull'
  writeFigure PNG ("QHullTest" ++ d ++ "Points" ++ ".png") (600, 600) (testFig vx vy hx hy (rxL, rxH) (ryL, ryH))

convertPoints :: V.Vector (UV.Vector Double) -> (Vector Double, Vector Double)
convertPoints =
  (V.convert *** V.convert) .
  V.unzip . V.map (UV.unsafeHead &&& UV.unsafeLast)

testFig :: Series
        -> Series
        -> Series
        -> Series
        -> (Double, Double)
        -> (Double, Double)
        -> Figure ()
testFig xs ys hx hy (rxL, rxH) (ryL, ryH) = do
  withTextDefaults $ setFontFamily "OpenSymbol"
  withTitle $ setText "2D Convex Hull with QuickHull!"
  setPlots 1 1
  withPlot (1, 1) $
    do
      setDataset [(xs, point ys (Asterisk, red)), (hx, point hy (Box, purple))]
      addAxis XAxis (Side Lower) $ withAxisLabel $ setText "x-axis"
      addAxis YAxis (Side Lower) $ withAxisLabel $ setText "y-axis"
      addAxis XAxis (Value 0) $ return ()
      setRange XAxis Lower Linear (rxL - 5) (rxH + 5)
      setRange YAxis Lower Linear (ryL - 5) (ryH + 5)

chop :: Int -> [a] -> [[a]]
chop n =
  unfoldr
    (\v -> if null v
             then Nothing
             else Just $ splitAt n v)


module InterpSpec where

import Commons
import Data.Complex (Complex (..))
import Data.Map.Strict (union)
import Data.Maybe (fromJust)
import Debug.Trace (traceShowId)
import HashedExpression.Internal.Expression
import HashedExpression.Internal.Normalize
import HashedExpression.Internal.Utils
import HashedExpression.Interp
import HashedExpression.Operation hiding (product, sum)
import qualified HashedExpression.Operation
import HashedExpression.Prettify
import Test.Hspec
import Test.QuickCheck (property)
import Var
import Prelude hiding ((^))

-- |
prop_AddScalarR :: SuiteScalarR -> SuiteScalarR -> Bool
prop_AddScalarR (Suite exp1 valMaps1) (Suite exp2 valMaps2) =
  eval valMaps (exp1 + exp2) == eval valMaps exp1 + eval valMaps exp2
  where
    valMaps = valMaps1 `union` valMaps2

-- |
prop_MultiplyScalarR :: SuiteScalarR -> SuiteScalarR -> Bool
prop_MultiplyScalarR (Suite exp1 valMaps1) (Suite exp2 valMaps2) =
  eval valMaps (exp1 * exp2) == eval valMaps exp1 * eval valMaps exp2
  where
    valMaps = valMaps1 `union` valMaps2

-- |
prop_AddScalarC :: SuiteScalarC -> SuiteScalarC -> Bool
prop_AddScalarC (Suite exp1 valMaps1) (Suite exp2 valMaps2) =
  eval valMaps (exp1 + exp2) == eval valMaps exp1 + eval valMaps exp2
  where
    valMaps = valMaps1 `union` valMaps2

-- |
prop_MultiplyScalarC :: SuiteScalarC -> SuiteScalarC -> Bool
prop_MultiplyScalarC (Suite exp1 valMaps1) (Suite exp2 valMaps2) =
  eval valMaps (exp1 * exp2) == eval valMaps exp1 * eval valMaps exp2
  where
    valMaps = valMaps1 `union` valMaps2

-- |
prop_RotateOneR1 :: SuiteOneR -> Bool
prop_RotateOneR1 (Suite exp valMaps) =
  eval valMaps (rotate 0 exp) == eval valMaps exp

-- |
prop_RotateOneR2 :: SuiteOneR -> Int -> Bool
prop_RotateOneR2 (Suite exp valMaps) amount =
  eval valMaps (f exp) == eval valMaps exp
  where
    f = rotate amount . rotate (- amount)

-- |
prop_RotateOneR3 :: SuiteOneR -> Int -> Int -> Bool
prop_RotateOneR3 (Suite exp valMaps) amount1 amount2 =
  eval valMaps (f1 exp) == eval valMaps (f2 exp)
  where
    f1 = rotate amount1 . rotate amount2
    f2 = rotate (amount1 + amount2)

-- |
prop_RotateTwoR1 :: SuiteTwoR -> Bool
prop_RotateTwoR1 (Suite exp valMaps) =
  eval valMaps (rotate (0, 0) exp) == eval valMaps exp

-- |
prop_RotateTwoR2 :: SuiteTwoR -> (Int, Int) -> Bool
prop_RotateTwoR2 (Suite exp valMaps) (offset1, offset2) =
  eval valMaps (f exp) == eval valMaps exp
  where
    f = rotate (offset1, offset2) . rotate (- offset1, - offset2)

-- |
prop_RotateTwoR3 :: SuiteTwoR -> (Int, Int) -> (Int, Int) -> Bool
prop_RotateTwoR3 (Suite exp valMaps) amount1 amount2 =
  eval valMaps (f1 exp) == eval valMaps (f2 exp)
  where
    f1 = rotate amount1 . rotate amount2
    f2 = rotate (fst amount1 + fst amount2, snd amount1 + snd amount2)

spec :: Spec
spec =
  describe "Interp spec" $ do
    specify "prop_Add Scalar R" $ property prop_AddScalarR
    specify "prop_Multiply Scalar R" $ property prop_MultiplyScalarR
    specify "prop_Add Scalar C" $ property prop_AddScalarC
    specify "prop_Multiply Scalar C" $ property prop_MultiplyScalarC
    specify "prop_Rotate One R rotate 0 should stay the same" $
      property prop_RotateOneR1
    specify "prop_Rotate One R rotate a and -a should stay the same" $
      property prop_RotateOneR2
    specify
      "prop_Rotate One R rotate a then rotate b should equal rotate (a + b)"
      $ property prop_RotateOneR3
    specify "prop_Rotate Two R rotate (0, 0) should stay the same" $
      property prop_RotateTwoR1
    specify "prop_Rotate Two R rotate a and -a should stay the same" $
      property prop_RotateTwoR2
    specify
      "prop_Rotate Two R rotate a then rotate b should equal rotate (a + b)"
      $ property prop_RotateTwoR3


{-# LANGUAGE ConstraintKinds, ScopedTypeVariables, TypeApplications, AllowAmbiguousTypes, NoImplicitPrelude,
    GeneralizedNewtypeDeriving, FlexibleContexts
#-}
module LinRel where

import Numeric.LinearAlgebra
import Prelude hiding ((<>))
import Debug.Trace
type BEnum a = (Enum a, Bounded a)
enumAll :: (BEnum a) => [a] -- What about Void?
enumAll = [minBound .. maxBound]

-- A Hacked "Void" type to get card @Void = 0
data Void = Void' | Void'' -- SORRY MOM.
instance Enum Void where
  fromEnum Void' =  1
  fromEnum Void'' = 0
  toEnum 1 = Void'
  toEnum 0 = Void''
instance Bounded Void where
  maxBound = Void''
  minBound = Void'

instance (Enum a, Enum b, Bounded a) => Enum (Either a b) where
  fromEnum (Left a) = fromEnum a
  fromEnum (Right b) = fromEnum b + fromEnum (maxBound @a) + 1
  toEnum n | n <= ((fromEnum (maxBound @a)) - (fromEnum (minBound @a))) = Left (toEnum n)
           | otherwise = Right (toEnum (n - fromEnum (maxBound @a)))

instance (Bounded a, Bounded b) => Bounded (Either a b) where
  maxBound = Right maxBound
  minBound = Left minBound

card :: forall a. (BEnum a) => Int
card = (fromEnum (maxBound @a)) - (fromEnum (minBound @a)) + 1

-- HLinRel holds A x = b constraint
data HLinRel a b = HLinRel (Matrix Double) (Vector Double) deriving Show

-- x = A l + b. Generator constraint.
data VLinRel a b = VLinRel (Matrix Double) (Vector Double) deriving Show

-- f(x) = xQx + bx
data QuadOp a b = QuadOp (Matrix Double) (Vector Double)
{-
qid :: QuadOp a a
qid = ? -- a = a ? Need Relations too.

qcompose :: QuadOp b c -> QuadOp a b -> QuadOp a c
qcompose (QuadOp q c) (QuadOp q' c') = QuadOp q'' c'' where
        where
                ca = card @a
                cb = card @b
                cc = card @c
                a = subMatrix (cb, cb) q
                b = subMatrix (cb,cc) q
                c = subMatrix (cc, cb) q
                d = subMatrix (cc,cc) q
                a' = subMatrix (ca, ca) q'
                b' = subMatrix (ca, cb) q'
                c' = subMatrix (cb, ca  q'
                d' = subMatrix (cb, cb) q'
                m = - (a + d')
                q'' = fromBlocks [[a' + b' <> m </> c'   , b' <> m </> c ],  -- can memoize some of this
                                  [b <> m </> c' , d - ]]

                a'' =  a' - b' <> m </> c'
                [v1', v2'] = takesV [ca,cb] c'
                [v1,  v2] = takesV [cb, cc] c
                v3 = (v2' + v1)
                c'' = vJoin [v1' + m </> c #>  v3, ]

        -}
    {-

    Any cost outside the constraint space is irrelevant.
    Q = V m V where m is fullrank.
    c = Vc also

    minimal sets
    -}

        -- break into pieces, form schur complement.

-- if A x = b then x is in the nullspace + a vector b' solves the equation
h2v :: HLinRel a b -> VLinRel a b
h2v (HLinRel a b) = VLinRel a' b' where
        b' = a <\> b -- least squares solution
        a' = nullspace a

-- if x = A l + b, then A' . x = A' A l + A' b = A' b because A' A = 0
v2h :: VLinRel a b -> HLinRel a b
v2h (VLinRel a' b') = HLinRel a b where
        b = a #> b' -- matrix multiply
        a = tr $ nullspace (tr a') -- orthogonal space to range of a.

hid :: forall a. BEnum a => HLinRel a a
hid =  HLinRel (i ||| (- i)) (vzero s) where
                            s = card @a
                            i = ident s

vzero :: Konst Double d c => d -> c Double
vzero = konst 0

hcompose :: forall a b c. (BEnum a, BEnum b, BEnum c) => HLinRel b c -> HLinRel a b -> HLinRel a c
hcompose (HLinRel m b) (HLinRel m' b') = let a'' = fromBlocks [[       ma',           mb' ,    0       ],
                                                               [         0 ,    mb,        mc          ]] in
                                         let b'' = vjoin [b', b] in
                                         let (VLinRel q p) = h2v (HLinRel a'' b'') in -- kind of a misuse
                                         let q' = (takeRows ca q)  -- drop rows belonging to @b
                                                       ===
                                                  (dropRows (ca + cb) q) in
                                         let [x,y,z] =  takesV [ca,cb,cc] p in
                                         let p'=  vjoin [x,z] in -- rebuild without rows for @b
                                         v2h (VLinRel q' p') -- reconstruct HLinRel
                                       where
                                           ca = card @a
                                           cb = card @b
                                           cc = card @c
                                           sb = size b -- number of constraints in first relation
                                           sb' = size b' -- number of constraints in second relation
                                           ma' = takeColumns ca m'
                                           mb' = dropColumns ca m'
                                           mb = takeColumns cb m
                                           mc = dropColumns cb m

(<<<) :: forall a b c. (BEnum a, BEnum b, BEnum c) => HLinRel b c -> HLinRel a b -> HLinRel a c
(<<<) = hcompose
-- stack the constraints
hmeet :: HLinRel a b -> HLinRel a b -> HLinRel a b
hmeet (HLinRel a b) (HLinRel a' b') = HLinRel (a === a') (vjoin [b,b'])

{- If they don't meet are we still ok?

I am not sure. Might be weird corner cases?

-}

hjoin :: HLinRel a b -> HLinRel a b -> HLinRel a b
hjoin v w = v2h $ vjoin' (h2v v) (h2v w)

-- hmatrix took vjoin from me :(
-- joining means combining generators and adding a new generator
-- Closed under affine combination l * x1 + (1 - l) * x2
vjoin' :: VLinRel a b -> VLinRel a b -> VLinRel a b
vjoin' (VLinRel a b) (VLinRel a' b') = VLinRel (a ||| a' ||| (asColumn (b - b'))) b

-- no constraints, everything
-- trivially true
htop :: forall a b. (BEnum a, BEnum b) => HLinRel a b
htop = HLinRel (vzero (1,ca + cb)) (konst 0 1) where
                                      ca = card @a
                                      cb = card @b
{-
hbottom :: forall a b. (BEnum a, BEnum b) => HLinRel a b
hbottom = HLinRel (vzero (1,ca + cb)) (konst 1 1) where
                                      ca = card @a
                                      cb = card @b
                                      -}
-- all the constraints! Only the origin.
-- no. it should be the empty set. Impossible to satisfy.
-- 0 x = 1 is impossible
-- not gonna play nice.
{-
hbottom :: forall a b. (BEnum a, BEnum b) => HLinRel a b
hbottom = HLinRel (ident (ca + cb)) (konst 0 (ca + cb)) where
                                    ca = card @a
                                    cb = card @b
  -}

hconverse :: forall a b. (BEnum a, BEnum b) => HLinRel a b -> HLinRel b a
hconverse (HLinRel a b) = HLinRel ( (dropColumns ca a) |||  (takeColumns ca a)) b where
    ca = card @a
    cb = card @b

    -- this is numerically unacceptable
-- forall l. A' ( A l + b) == b'
vhsub :: VLinRel a b -> HLinRel a b -> Bool
vhsub (VLinRel a b) (HLinRel a' b') = (naa' <=  1e-10 * (norm_2 a') * (norm_2 a)  ) && ((norm_2 ((a' #> b) - b')) <= 1e-10 * (norm_2 b')  ) where
          naa' = norm_2 (a' <> a)

hsub :: HLinRel a b -> HLinRel a b -> Bool
hsub h1 h2 = vhsub (h2v h1) h2

heq :: HLinRel a b -> HLinRel a b -> Bool
heq a b = (hsub a b) && (hsub b a)

instance Ord (HLinRel a b) where
  (<=) = hsub
  (>=) = flip hsub

instance Eq (HLinRel a b) where
  (==) = heq

-- I can't do this right?
-- hcomplement :: HLinRel a b -> HLinRel a b
-- hcomplement

hpar :: HLinRel a b -> HLinRel c d -> HLinRel (Either a c) (Either b d)
hpar (HLinRel mab v) (HLinRel mcd v') = HLinRel (fromBlocks [ [mab, 0], [0 , mcd]]) (vjoin [v, v']) where

hassoc :: forall a b c. (BEnum a, BEnum b, BEnum c) => HLinRel (Either (Either a b) c) (Either a (Either b c))
hassoc = HLinRel m v where HLinRel m v = hid @((Either (Either a b) c))

hassoc' :: forall a b c. (BEnum a, BEnum b, BEnum c) => HLinRel  (Either a (Either b c)) (Either (Either a b) c)
hassoc' = HLinRel m v where HLinRel m v = hid @((Either (Either a b) c))

{-
        -- Void is unit for Either.
-- void has no inhabitants.... This is a bad boy.

hcup :: HLinRel Void (Either a a)
hcap :: HLinRel (Either a a) Void

-}
hleft :: forall a b. (BEnum a, BEnum b) => HLinRel a (Either a b)
hleft = HLinRel ( i ||| (- i) ||| (konst 0 (ca,cb))) (konst 0 ca) where
    ca = card @a
    cb = card @b
    i = ident ca

hright :: forall a b. (BEnum a, BEnum b) => HLinRel b (Either a b)
hright = HLinRel ( i ||| (konst 0 (cb,ca)) ||| (- i) ) (konst 0 cb) where
    ca = card @a
    cb = card @b
    i = ident cb

hfan ::  forall a b c. BEnum a => HLinRel a b -> HLinRel a c -> HLinRel a (Either b c)
hfan (HLinRel m v) (HLinRel m' v') = HLinRel (fromBlocks [ [ma, mb, 0], [ma', 0, mc']]) (vjoin [v,v']) where
        ca = card @a
        ma = takeColumns ca m
        mb = dropColumns ca m
        ma' = takeColumns ca m'
        mc' = dropColumns ca m'

hdump :: HLinRel a Void
hdump = HLinRel 0 0
{-
hlabsorb :: HLinRel a b -> HLinRel (Either Void a) b
hlabsorb (HLinRel m v) = (HLinRel m v)
-}
hlabsorb ::forall a. BEnum a => HLinRel (Either Void a) a
hlabsorb = HLinRel m v where (HLinRel m v) = hid @a

htrans :: HLinRel a (Either b c) -> HLinRel (Either a b) c
htrans (HLinRel m v) = HLinRel m v

hswap :: forall a b. (BEnum a, BEnum b) => HLinRel (Either a b) (Either b a)
hswap = HLinRel (fromBlocks [[ia ,0,0 ,-ia], [0, ib,-ib,0]]) (konst 0 (ca + cb)) where
        ca = card @a
        cb = card @b
        ia = ident ca
        ib = ident cb

hsum :: forall a. BEnum a => HLinRel (Either a a) a
hsum = HLinRel ( i ||| i ||| - i ) (konst 0 ca)  where
        ca = card @a
        i= ident ca

hdup :: forall a. BEnum a => HLinRel a (Either a a)
hdup = HLinRel (fromBlocks [[i, -i,0 ], [i, 0, -i]]) (konst 0 (ca + ca))  where
        ca = card @a
        i= ident ca

-- hcup :: forall a. BEnum a => HLinRel Void (Either a a)
-- hcup = -- or mainulate hid

-- hcap :: forall a. BEnum a => HLinRel (Either a a) Void

-- smart constructors
hLinRel :: forall a b. (BEnum a, BEnum b) => Matrix Double -> Vector Double -> Maybe (HLinRel a b)
hLinRel m v | cols m == (ca + cb) && (size v == rows m)  = Just (HLinRel m v)
            |  otherwise = Nothing  where
                 ca = card @a
                 cb = card @b

-- a 2d space at every wire or current and voltage.
data IV = I | V deriving (Show, Enum, Bounded, Eq, Ord)

resistor :: Double -> HLinRel IV IV
resistor r = HLinRel ( (2><4)  [ 1,0,-1,   0,
                                 r, 1, 0, -1]) (konst 0 2)

bridge :: Double -> HLinRel (Either IV IV) (Either IV IV)
bridge r = HLinRel (  (4><8) [ 1,0, 1,  0, -1, 0, -1,  0, -- current conservation
                               0, 1, 0, 0, 0, -1 , 0,  0, --voltage maintained
                               0, 0, 0, 1, 0,  0,  0, -1, -- voltage maintained
                               r, 1, 0,-1, -r,  0,  0, 0  ]) (konst 0 4)
short = bridge 0

first :: BEnum c => HLinRel a b -> HLinRel (Either a c) (Either b c)
first f = hpar f hid

second :: BEnum a => HLinRel b c -> HLinRel (Either a b) (Either a c)
second f = hpar hid f

type HLinRel2D u d l r = HLinRel (Either u l) (Either d r)

{-
A stencil  of 2d resistors for tiling

          u
          /
          \
          /
         |
l -/\/\/----/\/\/\-  r
         |
         /
         \
         /
         |
         d

        -}
stencil :: HLinRel2D IV IV IV IV
stencil = (hpar r10 r10) <<< short <<< (hpar r10 r10) where r10 = resistor 10

horicomp :: forall w w' w'' w''' a b c. (BEnum w, BEnum w', BEnum a, BEnum w''', BEnum b, BEnum w'', BEnum c ) => HLinRel2D w' w'' b c -> HLinRel2D w w''' a b -> HLinRel2D (Either w' w) (Either w'' w''') a c
horicomp f g = hcompose f' g' where
               f' :: HLinRel (Either (Either w' w''') b) (Either (Either w'' w''') c)
               f' = (first hswap) <<< hassoc' <<< (hpar hid f) <<< hassoc <<<  (first hswap)
               g' :: HLinRel (Either (Either w' w) a) (Either (Either w' w''') b)
               g' = hassoc' <<< (hpar hid g) <<< hassoc

rotate :: (BEnum w, BEnum w', BEnum a, BEnum b) => HLinRel2D w w' a b -> HLinRel2D a b w w'
rotate f = hswap <<< f <<< hswap

vertcomp :: (BEnum w, BEnum w', BEnum a, BEnum d, BEnum b, BEnum w'', BEnum c ) => HLinRel2D w'  w'' c d -> HLinRel2D w w' a b -> HLinRel2D w w'' (Either c a) (Either d b)
vertcomp f g = rotate (horicomp (rotate f)  (rotate g) )

{-
traceV :: ->
traceH :: ->

  -}

stencil2 = vertcomp h h where h = stencil `horicomp` stencil

{- Legendre transformations in thermo are for open systems. SOmething to that -}
{- Dependent sources. Well, these are goddamn cheating.
We could do it though
|   |
I   alpha I
|   |

small signal models of transistor, op amps etc

gyrator - would need polynomials from caps and inds
would be kind of nice for boundary models

kron instead of dsum - Quantum fields
Kron relation might be nice for discussing remnant space of topological matter

-}

{-

A wire in a circuit has a potential (questionably) and a current running through it
So our wires should have both of these variables.

-}
newtype VProbe = VProbe () deriving (Enum, Bounded, Show, Eq, Ord)
vprobe :: HLinRel IV VProbe
vprobe = HLinRel ( (2><3)  [1,0,0,
                            0,1,-1]) (konst 0 2)

vsource :: Double -> HLinRel IV IV
vsource v = HLinRel ( (2><4) [ 1,0,-1,   0,
                               0, 1, 0, -1]) (fromList [0,v])

isource :: Double -> HLinRel IV IV
isource i = HLinRel ( (2><4) [  1,0, -1,   0 , -- current conservation
                                 1, 0, 0,  0]) (fromList [0,i])

-- the currents add, but the voltages dup. sum and dup are dual
-- Or should it be |--|   a parallel short?
-- Ad then we could open circuit one of them and absorb the Void
-- to derive this
cmerge :: HLinRel (Either IV IV) IV
cmerge = HLinRel ( (3><4)  [1, 0, 1, 0, -1, 0,
                            0,1,0,0,0 ,  -1  ,
                            0,0,0,1, 0, -1])  (konst 0 3)

open :: HLinRel IV Void
open = HLinRel ( (1><2) [1,0]) (konst 0 1)

cap :: HLinRel  (Either IV IV) Void
cap  = hcompose open cmerge

cup :: HLinRel Void (Either IV IV)
cup = hconverse cap

ground :: HLinRel IV Void
ground = HLinRel ( (1><2) [ 0 , 1 ]) (vzero 1)

-- resistors in parallel.

ex1 = hcompose (bridge 10) (bridge 10)
ex2 = hcompose (resistor 10) (resistor 30) -- resistors in series.
r20 :: HLinRel IV IV
r20 = resistor 20

divider :: Double -> Double -> HLinRel (Either IV IV) (Either IV IV)
divider r1 r2 = hcompose (bridge r2) (hpar (resistor r1) hid)

{-
type StateCoState s = Either s s
type ValueD s = s

dynamics :: HLinRel (Either SHOState Control) SHOState
dynamics = HLinRel ((2><5)  [ 1,  dt, 0, -1 ,  0,
                              -dt, 1, dt, 0, -1 ])  (vzero 2)
   where dt = 0.01

-- labsorb <<< (par initial_cond id)
-}

-- state of an oscillator
data SHOState = X | P deriving (Show, Enum, Bounded, Eq, Ord)
data Control = F  deriving (Show, Enum, Bounded, Eq, Ord)
-- Costate newtype wrapper
newtype Co a = Co a deriving (Show, Enum, Bounded, Eq, Ord)

type M = Matrix Double
dynamics :: forall x u. (BEnum x, BEnum u) =>  Matrix Double -> Matrix Double ->
  HLinRel (Either x u) x
dynamics a b =  HLinRel (a ||| b ||| -i )  (vzero cx) where
  cx = card @x
  cu = card @u
  i = ident cx

initial_cond :: forall x. BEnum x => Vector Double-> HLinRel Void x
initial_cond x0 =  HLinRel i x0 where
  cx = card @x
  i = ident cx

valueUpdate :: forall x l. (BEnum x, BEnum l) =>  M -> M -> HLinRel (Either x l) l
valueUpdate a q = HLinRel ((tr a) ||| q ||| i)  (vzero cl) where
  cl = card @l
  i = ident cl

optimal_u :: forall u l. (BEnum u, BEnum l) =>
     M -> M -> HLinRel u l
optimal_u r b = HLinRel (r ||| tr b) (vzero cu) where
  cu = card @u

step :: forall x u l. (BEnum x, BEnum u, BEnum l) => M -> M -> M -> M
        -> HLinRel (Either x l) (Either x l)
step a b r q =
  f5 <<< f4 <<< hassoc' <<< f3 <<< f2 <<< hassoc <<< f1 where
  f1 :: HLinRel (Either x l) (Either (Either x x) l)
  f1 = first hdup
  f2 :: HLinRel (Either x (Either x l)) (Either x l)
  f2 = second (valueUpdate a q)
  f3 ::  HLinRel (Either x l) (Either x (Either l l))
  f3 = second hdup
  f4 ::  HLinRel (Either (Either x l) l) (Either (Either x u) l)
  f4 = first (second (hconverse (optimal_u r b)))
  f5 ::  HLinRel (Either (Either x u) l) (Either x l)
  f5 = first (dynamics a b)

-- iterate (hcompose (step a b r q)) :: [HLinRel (Either x l) (Either x l)]

 {-
step :: forall x l. (BEnum x, BEnum l) =>
   HLinRel (State + ValueD) (State + ValueD)
step a b q r =
= dynamics,
  valueUpdate <<< fst
  second optimal_u
  par (par hid cup) hid
   ::  , optimal_u
-}
{-
cost :: HLinRel (Either (Co SHOState) SHOState Control) (Co SHOState)
cost = HLinRel ((2><5) [ 1,  0, 0,  -1 ,  0,
                        -dt, 1,  dt,  0,  -1 ])  (vzero 2)
-}

{-

Is there a reasonable intepretation of kron?

-}
{-
everything can be definedc inefficiently via v2s and h2v functions

right division

-}

    {-
Call them affine relations

Join and meet aren't union and intersection.
They are the affine closure of union and intersection.

Linear has some niceness.
Homgeonous coordinates usually do.
For clarity and familiaryt I have chosebn not to do it this way
Or maybe I will do it?

par

-}

{-

import numpy as np

def meet(a,b):
    pass
def compose(a,b): # a after b
    assert(a.inN == b.outN)

    combo = np.block([[a.constraints, np.zeros((a.constraints.shape[0] , b.inN) )     ],
                      [np.zeros((b.constraints.shape[0] , a.outN)) ,      b.constraints]])
    print(combo)
    gens = LinRel(a.inN + a.outN, b.outN, gens=combo).gens
    print("gens",gens)
    gens = np.vstack((gens[:a.outN, :], gens[-b.inN: , :]) )
    print(gens)
    return LinRel(a.outN, b.inN, gens=gens)
def top(outN, inN):
    return LinRel(outN, inN, constraints = np.array([[]]))
def bottom(outN, inN):
    return LinRel(outN, inN, gens = np.array([[]]))

def converse(a):
    return LinRel(a.inN, a.outN, np.hstack((a.constraints[:, a.inN:], a.constraints[:, :a.inN])))
def complement(a):
    return LinRel(a.outN, a.inN, constraints = a.gens.T.conj())

def right_div(a,b):
    pass # return complement( compose(a, complement(b)) ) # something like this
def inclusion(a,b):

    s = a.constraints @ b.gens
    np.all(a.constraints @ b.gens <= tol )

    if rcond is None:
        rcond = np.finfo(s.dtype).eps * max(max(a.shape), max(b.shape))
        tol = max(np.amax(a), np.amax(b))
        tol = np.amax(s) * rcond
def fromMat(mat):
    (outN, inN) = mat.shape
    return LinRel(inN, outN,constraints = np.hstack((mat,-np.eye(outN))))
def id(N):
    return fromMat(np.eye(N))
# make 0,1 first index for in/out? Oh, but then in out have to be same size.
# A[0, ...] @ x + A[1, ...] @ y = 0
# Then I can form the kron of linear relations.
# store sperate A B matrices? A @in + B @ out
class LinRel():
    def __init__(self, outN, inN, constraints = None, gens = None, rcond=None):
        #assert(inN <= constraints.shape[1])
        self.inN = inN
        self.outN = outN
        if constraints is not None: #baiscally scipy.linalg.null_space
            u, s, vh = np.linalg.svd(constraints, full_matrices=True)
            M, N = u.shape[0], vh.shape[1]
            if rcond is None:
                rcond = np.finfo(s.dtype).eps * max(M, N)
            tol = np.amax(s) * rcond
            num = np.sum(s > tol, dtype=int)
            self.gens = vh[num:,:].T.conj()
            self.constraints = vh[:num,:]
        if gens is not None: #basically scipy.linalg.orth
            u, s, vh = np.linalg.svd(gens, full_matrices=True)
            M, N = u.shape[0], vh.shape[1]
            if rcond is None:
                rcond = np.finfo(s.dtype).eps * max(M, N)
            tol = np.amax(s) * rcond
            num = np.sum(s > tol, dtype=int)
            self.gens = u[:, :num]
            self.constraints = u[:, num:].T.conj()

    def shape(self):
        return (self.outN, self.inN)
    def size(self):
        return self.outN + self.inN
    # operator overloadings
    def __matmul__(a,b):
        return compose(a,b)
    def __invert__(a): # ~
        return complement(a)
    def __or__(a,b): # an argument could be made for + and *
        return join(a,b)
    def __and__(a,b):
        return meet(a,b)
    def __sub__(a,b):
        return  (a) & (-b)
    def __le__(a,b): # Are the others automatic?
        return inclusion(a,b)
    def __str__(self):
        return " Constraints: \n%s, \nGens:\n %s\n" % (str(self.constraints), str(self.gens))

ex = LinRel(1,2, np.array([[3,4,0]]))
e2 = LinRel(2,1, np.array([[3,4,0]]))

assert(np.all(LinRel(1,2, ex.constraints).constraints == ex.constraints) )
print( ex.constraints @ ex.gens)
assert(np.all( np.abs(ex.constraints @ ex.gens) <= 1e-15) )
print(ex @ e2)
print(e2 @ ex)
print(e2 @ id(e2.inN))

'''
Quadratic optimization can be mixed in.
Quad && LinRel

AffineRel = maintain homgenous coord, or always insert 1 -1 keeping homoeg coord
+ discrete? Maintain a couple copies of

'''

Linear relations
hrep - Ax = 0
or
vrep - y = sum x_i

hrep <-> vrep = row echelon

in and out variables. In and out subspaces.
in : [] - list of indices
out : []

in is a vrep of input space.

in = projection/injection matrix n x d
out = projection/injection matrix. (d-n) x d
in * out = 0. orthogonal

auxiliary variables allowed

compose relations
in1 out1
in2 out2

in = d x n
stack out1 and in2 into matrix. with them equal.

np.hstack(A1, out1 - in2, A2)
in = no.hatck(in, zeores)
in = no.hatck(zeores, out)

drect sum as monoidal product
block([ A, 0  ],
      [ 0, A  ])
in = [in1, in2]
out = [out1, out2]

converse = flip in out

meet = combine the two constraint matrices
join = convert? combine in out?

1-d as unit object
<= is subspace ordering

negation = orthogonalization
division =>

a linear problem is a linear relation of 1-d.

use in1 and out2 as new in/out

fan
snd(30,10) = project bottom 10 = idetnity matrix stacked.
fst(30, 10) project top 10
id(20)
id(n) = LinRel(np.zeros(0,2*n), [zeros, eye], [eye, zero] )
id(n) = LinRel((0,n), eye(n), eye(n)  )
        LinRel [I, -I], [I, 0], [0,I]

"internal" space
class LinRel():
    def __init__(A, in, out):

svd(in * A , smallest )
1 - A*A

Ax + By = 0

vstack adds constraints
hstack adds variables

def idRel(n):
  return LinRel(sparse.lil_matrix((0,n)), n, n)
def mapRel(A):
  (r,c) = A.shape
  newA = sparse.hstack(A, - sparse.eye(r))
  return LinRel(newA, r, c)

class LinRelV():
class LinRelH():

class LinRel():
  def init(self,A, in, out):
    self.A = A
    self.in = in
    self.out = out
  def compose(self, b):
    assert(self.out == b.in, "Shapes must match")

    i = sparse.eye(b.in)
    cons = sparse.hstack([0, i, -i, 0])
    ina = self.A[:,:self.in]
    auxa = self.A[:,self.in:-self.out]
    outa = self.A[:,-self.out:]
    inb = b.A[:,:b.in]
    auxb = b.A[:,b.in:-b.out]
    outb = b.A[:,-b.out:]

    newA = sparse.bmat(  [[ina, auxa, outa, 0, 0,    0],
                          [0  , 0   , i,   -i, 0 ,   0]
                          [0,  0,   0,    inb, auxb,outb]])
    LinRel(newA, self.in, b.out)
  def meet(self,b):
    #hmm. I suppose we acutlaly should split the thing apart again
    assert(self.in == b.in)
    assert(self.out == b.out)
    assert()
    newA = sparse.vstack([self.A, b.A])
    return ()
  def complement(self):
    linalg.svd(self.A)
    return LinRel(get_nonzeroeigs)
  def __negate__(self):
    return self.complement()
  def rdiv(self):
  def transpose():
    self.converse()
  def T():
    self.converse()

  # the svd gives you most of what you need?
  def inclusion():
    x = linalg.nullspace(self.A)
    return b.A @ x == 0 #check that every generator is in. Makes sense. Except numrically is trash.
  def __leq__(self):
    self.inclusion(b)

  def converse(rel):
      newA = np.hstack( [ rel.A[:,-rel.out: ] , rel.A[:,rel.in:-rel.out], rel.A[:,:rel.in ] ])
      return LinRel(newA, rel.out, rel.in)

compose(A)
linalg.nullspace(A)
range?

(cons1,d1) = self.A.shape
(cons2,d2) = b.A.shape
constrain = np.hstack
newA = sparse.vstack

using bigM I can encode the complement of a H-polytope
but then what?
I do it again I guess?

complementation of relation ->
At least one must be inverted.

polytope inclusion
-> search for point not in B.
Or, do sandardinni encoding

Really It should generate new variables for an instantiation.
Snd should not reuse the same variables every time.
class PolyRel()
  invars = []
  constraints = [] # store >= values, not full constraints?
  outvars = []
  def __init__():
    all fresh variables
  def compose(self,b):
    PolyRel(self.constraints + self., invars = outvars )
  def complement():
    zs = []
    for c in constraints:
      z, constraints = reify(c)
      # z = cvx.Variable(1, boolean=True) # actually make same shape as c
      # c += c + M * z
    sum(zs) >= 1 # one or more constraints is disatisfied.
  def rdiv():

yeah. We should use a dsl, compile it, then encode it.
data Rel a b where
  Compose ::
  Complement ::
  Converse ::

relu

l1 >= 0
l2 >= 0
l1 <= M * z
l2 <= M * (1 - z)
x = lambda1 - lambda2
y = lambda1

Maybe insetado f subspaces, we should be thinking ellipses and svd.

-}


import Test.Framework
import Test.Framework.Providers.QuickCheck2

import Vector
import STVector
import Matrix
import STMatrix
import Statistics

main :: IO ()
main = defaultMain tests
  where
    tests = [ tests_Vector
            , tests_STVector
            , tests_Matrix
            , tests_STMatrix
            , tests_Statistics
            ]


-----------------------------------------------------------
-- |
-- module:                      Math.Complex
-- copyright:                   (c) 2017 HE, Tao
-- license:                     MIT
-- maintainer:                  sighingnow@gmail.com
--
-- Coordinate `Data.Complex` with foundation's numeric and primitive type support.
--
{-# OPTIONS_GHC -Wno-orphans #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UnboxedTuples #-}

module Math.Complex
  ( Complex (..)
  , conjugate
  ) where

import Foundation
import Foundation.Class.Storable
import Foundation.Primitive

import Data.Complex (Complex(..), conjugate)
import GHC.Exts

instance Additive a => Additive (Complex a) where
    azero = azero :+ azero
    (a :+ b) + (c :+ d) = (a + c) :+ (b + d)

instance (Subtractive a, Difference a ~ a) => Subtractive (Complex a) where
    type Difference (Complex a) = (Complex a)
    (a :+ b) - (c :+ d) = (a - c) :+ (b - d)

instance (Additive a, Subtractive a, Difference a ~ a, Multiplicative a) => Multiplicative (Complex a) where
    midentity = midentity :+ azero
    (a :+ b) * (c :+ d) = (a * c - b * d) :+ (b * c + a * d)

instance (Additive a, Subtractive a, Difference a ~ a, Divisible a) => Divisible (Complex a) where
    (a :+ b) / (c :+ d) = ((a * c + b * d) / (c * c + d * d)) :+ ((b * c - a * d) / (c * c + d * d))

offsetComplex :: Offset (Complex a) -> (# Int#, Int# #)
offsetComplex !(Offset (I# i)) = (# n, n +# 1# #)
    where !n = uncheckedIShiftL# i 1#

{-# INLINE offsetComplex #-}

instance PrimType (Complex Float) where
    primSizeInBytes _ = primSizeInBytes (Proxy :: Proxy Float) + primSizeInBytes (Proxy :: Proxy Float)
    {-# INLINE primSizeInBytes #-}

    primShiftToBytes _ = 3 -- TODO may be wrong
    {-# INLINE primShiftToBytes #-}

    primBaUIndex ba n = F# (indexFloatArray# ba n1) :+ F# (indexFloatArray# ba n2)
        where !(# n1, n2 #) = offsetComplex n
    {-# INLINE primBaUIndex #-}

    primMbaURead mba n = primitive $ \s1 -> let !(# s2, r1 #) = readFloatArray# mba n1 s1
                                                !(# s3, r2 #) = readFloatArray# mba n2 s2
                                             in (# s3, F# r1 :+ F# r2 #)
        where !(# n1, n2 #) = offsetComplex n
    {-# INLINE primMbaURead #-}

    primMbaUWrite mba n ((F# w1) :+ (F# w2)) = primitive $ \s1 -> let !s2 = writeFloatArray# mba n1 w1 s1
                                                                   in (# writeFloatArray# mba n2 w2 s2, () #)
        where !(# n1, n2 #) = offsetComplex n
    {-# INLINE primMbaUWrite #-}

    primAddrIndex addr n = F# (indexFloatOffAddr# addr n1) :+ F# (indexFloatOffAddr# addr n2)
        where !(# n1, n2 #) = offsetComplex n
    {-# INLINE primAddrIndex #-}

    primAddrRead addr n = primitive $ \s1 -> let !(# s2, r1 #) = readFloatOffAddr# addr n1 s1
                                                 !(# s3, r2 #) = readFloatOffAddr# addr n2 s2
                                              in (# s3, F# r1 :+ F# r2 #)
        where !(# n1, n2 #) = offsetComplex n
    {-# INLINE primAddrRead #-}

    primAddrWrite addr n ((F# w1) :+ (F# w2)) = primitive $ \s1 -> let !s2 = writeFloatOffAddr# addr n1 w1 s1
                                                                    in (# writeFloatOffAddr# addr n2 w2 s2, () #)
        where !(# n1, n2 #) = offsetComplex n
    {-# INLINE primAddrWrite #-}

instance PrimType (Complex Double) where
    primSizeInBytes _ = primSizeInBytes (Proxy :: Proxy Double) + primSizeInBytes (Proxy :: Proxy Double)
    {-# INLINE primSizeInBytes #-}

    primShiftToBytes _ = 5 -- TODO may be wrong
    {-# INLINE primShiftToBytes #-}

    primBaUIndex ba n = D# (indexDoubleArray# ba n1) :+ D# (indexDoubleArray# ba n2)
        where !(# n1, n2 #) = offsetComplex n
    {-# INLINE primBaUIndex #-}

    primMbaURead mba n = primitive $ \s1 -> let !(# s2, r1 #) = readDoubleArray# mba n1 s1
                                                !(# s3, r2 #) = readDoubleArray# mba n2 s2
                                             in (# s3, D# r1 :+ D# r2 #)
        where !(# n1, n2 #) = offsetComplex n
    {-# INLINE primMbaURead #-}

    primMbaUWrite mba n ((D# w1) :+ (D# w2)) = primitive $ \s1 -> let !s2 = writeDoubleArray# mba n1 w1 s1
                                                                   in (# writeDoubleArray# mba n2 w2 s2, () #)
        where !(# n1, n2 #) = offsetComplex n
    {-# INLINE primMbaUWrite #-}

    primAddrIndex addr n = D# (indexDoubleOffAddr# addr n1) :+ D# (indexDoubleOffAddr# addr n2)
        where !(# n1, n2 #) = offsetComplex n
    {-# INLINE primAddrIndex #-}

    primAddrRead addr n = primitive $ \s1 -> let !(# s2, r1 #) = readDoubleOffAddr# addr n1 s1
                                                 !(# s3, r2 #) = readDoubleOffAddr# addr n2 s2
                                              in (# s3, D# r1 :+ D# r2 #)
        where !(# n1, n2 #) = offsetComplex n
    {-# INLINE primAddrRead #-}

    primAddrWrite addr n ((D# w1) :+ (D# w2)) = primitive $ \s1 -> let !s2 = writeDoubleOffAddr# addr n1 w1 s1
                                                                    in (# writeDoubleOffAddr# addr n2 w2 s2, () #)
        where !(# n1, n2 #) = offsetComplex n
    {-# INLINE primAddrWrite #-}

instance Storable (Complex Float) where
    peek (Ptr addr) = primAddrRead addr (Offset 0)
    poke (Ptr addr) = primAddrWrite addr (Offset 0)

instance Storable (Complex Double) where
    peek (Ptr addr) = primAddrRead addr (Offset 0)
    poke (Ptr addr) = primAddrWrite addr (Offset 0)

instance StorableFixed (Complex Float) where
    size _ = size (Proxy :: Proxy Float) + size (Proxy :: Proxy Float)
    alignment _ = alignment (Proxy :: Proxy Float)

instance StorableFixed (Complex Double) where
    size _ = size (Proxy :: Proxy Double) + size (Proxy :: Proxy Double)
    alignment _ = alignment (Proxy :: Proxy Double)


-- |
-- Tests for data serialization instances
module Tests.Serialization where

import Data.Binary (Binary,decode,encode)
import Data.Aeson  (FromJSON,ToJSON,Result(..),toJSON,fromJSON)
import Data.Typeable

import Statistics.Distribution.Beta           (BetaDistribution)
import Statistics.Distribution.Binomial       (BinomialDistribution)
import Statistics.Distribution.CauchyLorentz
import Statistics.Distribution.ChiSquared     (ChiSquared)
import Statistics.Distribution.Exponential    (ExponentialDistribution)
import Statistics.Distribution.FDistribution  (FDistribution)
import Statistics.Distribution.Gamma          (GammaDistribution)
import Statistics.Distribution.Geometric
import Statistics.Distribution.Hypergeometric
import Statistics.Distribution.Laplace        (LaplaceDistribution)
import Statistics.Distribution.Normal         (NormalDistribution)
import Statistics.Distribution.Poisson        (PoissonDistribution)
import Statistics.Distribution.StudentT
import Statistics.Distribution.Transform      (LinearTransform)
import Statistics.Distribution.Uniform        (UniformDistribution)
import Statistics.Types

import Test.Tasty            (TestTree, testGroup)
import Test.Tasty.QuickCheck (testProperty)
import Test.QuickCheck         as QC

import Tests.Helpers
import Tests.Orphanage ()

tests :: TestTree
tests = testGroup "Test for data serialization"
  [ serializationTests (T :: T (CL Float))
  , serializationTests (T :: T (CL Double))
  , serializationTests (T :: T (PValue Float))
  , serializationTests (T :: T (PValue Double))
  , serializationTests (T :: T (NormalErr Double))
  , serializationTests (T :: T (ConfInt   Double))
  , serializationTests' "T (Estimate NormalErr Double)" (T :: T (Estimate NormalErr Double))
  , serializationTests' "T (Estimate ConfInt Double)" (T :: T (Estimate ConfInt   Double))
  , serializationTests (T :: T (LowerLimit Double))
  , serializationTests (T :: T (UpperLimit Double))
    -- Distributions
  , serializationTests (T :: T BetaDistribution        )
  , serializationTests (T :: T CauchyDistribution      )
  , serializationTests (T :: T ChiSquared              )
  , serializationTests (T :: T ExponentialDistribution )
  , serializationTests (T :: T GammaDistribution       )
  , serializationTests (T :: T LaplaceDistribution     )
  , serializationTests (T :: T NormalDistribution      )
  , serializationTests (T :: T UniformDistribution     )
  , serializationTests (T :: T StudentT                )
  , serializationTests (T :: T (LinearTransform NormalDistribution))
  , serializationTests (T :: T FDistribution           )
  , serializationTests (T :: T BinomialDistribution       )
  , serializationTests (T :: T GeometricDistribution      )
  , serializationTests (T :: T GeometricDistribution0     )
  , serializationTests (T :: T HypergeometricDistribution )
  , serializationTests (T :: T PoissonDistribution        )
  ]

serializationTests
  :: (Eq a, Typeable a, Binary a, Show a, Read a, ToJSON a, FromJSON a, Arbitrary a)
  => T a -> TestTree
serializationTests t = serializationTests' (typeName t) t

-- Not all types are Typeable, unfortunately
serializationTests'
  :: (Eq a, Binary a, Show a, Read a, ToJSON a, FromJSON a, Arbitrary a)
  => String -> T a -> TestTree
serializationTests' name t = testGroup ("Tests for: " ++ name)
  [ testProperty "show/read" (p_showRead t)
  , testProperty "binary"    (p_binary   t)
  , testProperty "aeson"     (p_aeson    t)
  ]

p_binary :: (Eq a, Binary a) => T a -> a -> Bool
p_binary _ a = a == (decode . encode) a

p_showRead :: (Eq a, Read a, Show a) => T a -> a -> Bool
p_showRead _ a = a == (read . show) a

p_aeson :: (Eq a, ToJSON a, FromJSON a) => T a -> a -> Bool
p_aeson _ a = Data.Aeson.Success a == (fromJSON . toJSON) a


{-# OPTIONS_GHC -Wall #-}

-- | 'Distribution' avoids a name clash with 'Data.Distributive'
module Plankton.Distribution
  ( Distribution
  ) where

import Data.Complex (Complex(..))
import Plankton.Additive
import Plankton.Multiplicative
import Protolude (Bool(..), Double, Float, Int, Integer)

-- | Distribution (and annihilation) laws
--
-- > a * (b + c) == a * b + a * c
-- > (a + b) * c == a * c + b * c
-- > a * zero == zero
-- > zero * a == zero
class (Additive a, MultiplicativeMagma a) =>
      Distribution a

instance Distribution Double

instance Distribution Float

instance Distribution Int

instance Distribution Integer

instance Distribution Bool

instance (AdditiveGroup a, Distribution a) => Distribution (Complex a)


module Ray where

import Numeric.Vector
import Numeric.Scalar
import Types
import Object.Tri (asList)

createHit :: Object -> Tri -> Ray -> Double -> (Double, Double, Double) -> Vec3d -> RayHit
createHit obj tri ray dist (a, b, c) pos = let abc3 = map (fromScalar . scalar) [a, b, c] :: [Vec3d]
                                               abc2 = map (fromScalar . scalar) [a, b, c] :: [Vec2d]
                                               tex = sum $ zipWith (*) abc2 $ map v_getTexCoord $ asList tri
                                               normal = sum $ zipWith (*) abc3 $ map v_getNormal $ asList tri
                                               tangent = sum $ zipWith (*) abc3 $ map v_getTangent $ asList tri
                                               bitangent = tangent × normal
                                               color = sum $ zipWith (*) abc3 $ map v_getColor $ asList tri

                                           in  RayHit { rh_getPos = pos
                                                      , rh_getRay = ray
                                                      , rh_getDistance = dist
                                                      , rh_getTri = tri
                                                      , rh_getObj = obj
                                                      , rh_getShader = o_getShader obj
                                                      , rh_getTexCoord = tex
                                                      , rh_getNormal = normal
                                                      , rh_getTangent = tangent
                                                      , rh_getBitangent = bitangent
                                                      , rh_getColor = color
                                                      }


{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}

module Cas.Misc where

import           PreludeCustom
import           Data.Function
import           Data.List hiding (intersect)
import           Debug.Trace
import           Data.Foldable
import           Data.Ratio
import           Data.Complex.Generic

foldlFst :: (a -> b -> a) -> a -> [(b,c)] -> (a,c)
foldlFst f a =  bimap (foldl f a) head . unzip

debug :: (Show a) => String -> a -> a
debug s = (\x -> trace (s <> " " <> show x ) x )

showQ = (==1) . denominator ?>>> show . numerator
                             ||> ((<>) <$> (<>"/") . show . numerator
                                       <*> show . denominator)

showImg y = bool "" (showQ y) (y/=1) <> "i"

showF (x':+y') = case (x',y') of
               (0,0) -> "0"
               (x,0) -> showQ x
               (0,y) -> showImg y
               (x,y) -> "(" <> showQ x <> "+" <> showImg y <> ")"

-- data QWrapperForShow = QWrapperForShow { unwrapQ :: Rational }

-- instance Show QWrapperForShow where
--     show (QWrapperForShow q) = showQ q

-- showF :: Complex Rational -> String
-- showF (x:+y)= show $ (QWrapperForShow x) :+ (QWrapperForShow y)
--

wrap :: String -> String
wrap s = "(" <> s <> ")"

wrapshow :: (Show a) => a -> String
wrapshow = wrap . show

stableGroupOn :: (Ord b) => (a -> b) -> [a] -> [[a]]
stableGroupOn f = fmap fromJust
                . (fmap . flip lookup
                    <$> fmap (f . head &&& id)
                      . groupBy ((==) `on` f)
                      . sortBy  (comparing f)
                    <*> nub
                      . fmap f )

stableGroup :: (Ord a) => [a] -> [[a]]
stableGroup = stableGroupOn id


{-# LANGUAGE DeriveDataTypeable, ScopedTypeVariables, NoMonomorphismRestriction, BangPatterns, FlexibleInstances, ViewPatterns #-}
module Math.Probably.RandIO where

import qualified Math.Probably.PDF as PDF
import Math.Probably.Sampler
import Math.Probably.MCMC hiding (traceIt)
import Math.Probably.StochFun
import Math.Probably.FoldingStats
import Math.Probably.Sampler
import qualified Control.Monad.State.Strict as S
import System.IO
import qualified Data.Packed.Matrix as L
import Control.Monad.Trans
import Control.Monad
import Data.IORef
import qualified Numeric.LinearAlgebra as L
import Math.Probably.NelderMead

import qualified Data.Vector.Storable as V
import Statistics.Autocorrelation

import Text.Printf
import Data.List
import Control.Monad.Par

import Debug.Trace

type RIO = S.StateT Seed IO

runRIO :: RIO a -> IO a
runRIO mx = do
  sd <- getSeedIO
  S.evalStateT mx sd

io :: IO a -> RIO a
io = liftIO

sample :: Sampler a -> RIO a
sample (Sam f) = do sd <- S.get
                    let (x,sd') = f sd
                    S.put sd'
                    return x

update :: IORef a -> (a->Sampler a) -> RIO ()
update rf sm = do
  x <- io $ readIORef rf
  newx <- sample $ sm x
  io $ writeIORef rf newx

runChainRIO :: Int -> (a->String) -> a -> (a->Sampler a) -> RIO [a]
runChainRIO n showit init sam = do
    seed <- S.get
    (nseed, xs) <- io $ go n seed init []
    S.put nseed
    return xs
     where go 0 s x xs = return (s, xs)
           go nn s x xs = do let (!xx, !ns) = unSam (sam x) s
                             when(nn `rem` chsz==0) $
                                 putStrLn $ show (((n-nn) `div` chsz)*2)++"%: " ++showit xx
                             go (nn-1) ns xx $ xx:xs
           chsz = n `div` 50

data AdaMetRunPars = AdaMetRunPars
     { nmTol :: Double,
       displayIt :: Maybe (L.Vector Double -> String),
       verboseNM :: Bool,
       amnsam :: Int,
       initw:: Int -> Double,
       minNM ::Int,
       maxNM ::Int }

defaultAM = AdaMetRunPars 0.5 Nothing False 1000 (const 0.02) 100 1000

traceit s x = trace (s++show x) x

getInitialCov :: Sampler (L.Vector Double)
              -> (L.Vector Double -> Double)
              -> RIO (Double, L.Vector Double, L.Matrix Double)
getInitialCov inisam posteriorL = go 500 where
  pdfv v = posteriorL v
  go n | n < 1
         = error $ "can't find suitable initial conditions. Try a narrower prior"
       | otherwise =  do
    initialv <- sample inisam
    let iniSimplex = genInitial (negate . pdfv) [] (const 0.02)  $ initialv
    if any (nanOrInf . snd) iniSimplex
       then go (n-1)
       else runNM n iniSimplex
  runNM n iniSimplex = do
    let finalSim = goNm (negate . pdfv) [] 0.01 100 1000 iniSimplex
        (maxPost,hess) = hessianFromSimplex (negate . pdfv) [] [] finalSim
    if any nanOrInf $ concat $ L.toLists hess
       then go (n-10)
       else getCov finalSim hess
  getCov finalSim hess = do
    let cor = case L.mbCholSH $ hess of
            Just _ -> L.inv hess
            Nothing -> case L.mbCholSH $ PDF.posdefify hess of
                         Just _ -> L.inv $ PDF.posdefify hess
                         Nothing -> L.diag $ L.mapVector (recip)
                                           $ L.takeDiag $ hess
        initV = fst $ head $ finalSim
    return (negate $ snd $ head $ finalSim,
            initV, cor)

laplaceApprox :: AdaMetRunPars -> PDF.PDF (L.Vector Double) -> [Int] -> [((Int, Int), Double)]
              -> L.Vector Double -> (L.Vector Double, Maybe (L.Matrix Double), Simplex)
laplaceApprox (AdaMetRunPars nmtol dispit verbnm nsam initw minNM maxNM) pdf isInt fixed init =
     let iniSim = genInitial (negate . pdf) isInt initw $ init
         finalSim = {-if verbnm then goNmVerbose (negate . pdf) isInt nmtol minNM maxNM iniSim
                              else -} goNm (negate . pdf) isInt nmtol minNM maxNM iniSim

         (maxPost,hess) = hessianFromSimplex (negate . pdf) isInt fixed finalSim
--     io $ print maxPost
         mbcor = case L.mbCholSH $ hess of
                   Just _ -> Just $ L.inv hess
                   Nothing -> case L.mbCholSH $ PDF.posdefify hess of
                                 Just _ -> Just $ L.inv $ PDF.posdefify hess
                                 Nothing -> Just $ L.diag $ L.mapVector (recip) $ L.takeDiag $ hess
         initV = fst $ head $ finalSim
     in (initV, mbcor, finalSim)

nmAdaMet :: AdaMetRunPars -> PDF.PDF (L.Vector Double) -> [Int] -> [((Int, Int), Double)]
            -> L.Vector Double -> RIO [L.Vector Double]
nmAdaMet (AdaMetRunPars nmtol dispit verbnm nsam initw minNM maxNM ) pdf isInt fixed init = do
     let iniSim = genInitial (negate . pdf) isInt initw $ init
     io $ print iniSim
     let finalSim =  goNm (negate . pdf) isInt nmtol minNM maxNM iniSim
     io $ print finalSim
     let (maxPost,hess) = hessianFromSimplex (negate . pdf) isInt fixed finalSim
--     io $ print maxPost
     io $ putStrLn "hessian"
     io $ print hess

     let mbcor = case L.mbCholSH hess of
                   Just _ -> Just $ L.inv hess
                   Nothing -> Nothing
     io $ putStrLn "maybe inverse hess"
     io $ print mbcor
     let initV = centroid finalSim
     let mbcorChol = mbcor >>= L.mbCholSH
     case (mbcor, mbcorChol) of
       (Just cor, Just _) ->  do --io $ putStrLn "chol cor"
                                 --io $ print $ L.inv cor
                                 let ampar = AMPar initV initV cor 2.4 (pdf initV) 0 0
                                 runAdaMetRIO nsam True ampar pdf
       _         -> do iniampar <- sample $ initialAdaMet 100 (const 5e-3) pdf initV
                       froampar <- runAndDiscard (nsam*2) (show . ampPar) iniampar $ adaMet False pdf
                       runAdaMetRIO (nsam*2) True froampar pdf

runAdaMetRIO :: Int -> Bool -> AMPar -> PDF.PDF (L.Vector Double) -> RIO [L.Vector Double]
runAdaMetRIO n freeze ampar pdf = do
    seed <- S.get
    (nseed, xs) <- io $ go n seed ampar []
    S.put nseed
    return xs
     where go 0 s amp vs = do print $ amp
                              return (s, reverse vs)
           go nn s amp vs = do let (!ampn, !ns) = unSam (adaMet freeze pdf amp) s
                               when(nn `rem` chsz==0) $
                                   putStrLn $ show (((n-nn) `div` chsz)*2)++"%: " ++showV (ampPar ampn) ++" LH="++printf "%.3g" (pdf (ampPar ampn)) ++ " accept="++acceptS ampn
                               go (nn-1) ns (ampn) $ (ampPar ampn):vs
           chsz = n `div` 50

runAdaMetRioESS :: Int -> Bool -> AMPar -> PDF.PDF (L.Vector Double) -> RIO [L.Vector Double]
runAdaMetRioESS want_ess freeze ampar pdf = do
    seed <- S.get
    (nseed, xs, amp) <- io $ go 200 seed ampar []
    S.put nseed
    return xs
     where go (0::Int) s amp vs = goChunks s amp vs
           go nn s amp vs = do let (!ampn, !ns) = unSam (adaMet freeze pdf amp) s
                               go (nn-1) ns (ampn) $ (ampPar ampn):vs

           goChunks s amp [] = do
              (nseed, xs, namp) <- go 200 s amp []
              goChunks nseed namp xs
           goChunks s amp xs = do
              let have_ess = min (realToFrac $ count_accept amp)
                                 $ calcESS  want_ess xs
                  drawn = length xs
              putStrLn $ show $ ampPar amp
              putStrLn $ "ESS="++show have_ess++" from "++show drawn
                       ++" drawn accept ratio="++acceptS amp
              if have_ess > realToFrac want_ess
                 then return (s, xs, amp)
                 else do
                   let need_ess = min (realToFrac $ want_ess `div` 5)
                                    $ max 1
                                    $ realToFrac want_ess - have_ess
                       samples_per_es = realToFrac drawn/have_ess
                       to_do = max 50 $ min 2000 $ samples_per_es * need_ess
                   putStrLn $ "now doing "++ show to_do
                   go (round to_do) s amp xs

runFixMetRioESS :: Double -> Int -> AMPar -> PDF.PDF (L.Vector Double) -> RIO [L.Vector Double]
runFixMetRioESS factor want_ess ampar pdf = do
    seed <- S.get
    (nseed, xs, amp) <- io $ go 200 seed ampar []
    S.put nseed
    return xs
     where go (0::Int) s amp vs = goChunks s amp vs
           go nn s amp vs = do let (!ampn, !ns) = unSam (fixedMet factor pdf amp) s
                               go (nn-1) ns (ampn) $ (ampPar ampn):vs

           goChunks s amp [] = do
              (nseed, xs, namp) <- go 200 s amp []
              goChunks nseed namp xs
           goChunks s amp xs = do
              let have_ess = min (realToFrac $ count_accept amp)
                             $ calcESS want_ess xs
                  drawn = length xs
              --putStrLn $ show $ ampPar amp
              putStrLn $ "ESS="++show have_ess++" from "++show drawn
                       ++" drawn accept ratio="++acceptS amp
              if have_ess > realToFrac want_ess
                 then return (s, xs, amp)
                 else do
                   let need_ess = min (realToFrac $ want_ess `div` 5)
                                    $ max 1
                                    $ realToFrac want_ess - have_ess
                       samples_per_es = realToFrac drawn/have_ess
                       to_do = max 50 $ min 2000 $ samples_per_es * need_ess * 1.2
                   putStrLn $ "now doing "++ show to_do
                   go (round to_do) s amp xs

runFixMetRioToFile :: Int -> Int -> String -> Double -> AMPar
                      -> PDF.PDF (L.Vector Double) -> RIO ()
runFixMetRioToFile samples thinn fileNm factor ampar pdf = do
    seed <- S.get
    h <- io $ openFile (fileNm) WriteMode
    nseed <- io $ go h samples seed ampar []
    S.put nseed
    return ()
     where go h (0::Int) s amp vs = return s
           go h nn s amp vs = do let (!ampn, !ns) = unSam (fixedMet factor pdf amp) s
                                 when (nn `rem` thinn == 0) $ do
                                     hPutStrLn h $ show $ L.toList $ ampPar ampn
                                 go h (nn-1) ns (ampn) $ (ampPar ampn):vs

runFixMetRio :: Double -> Int -> AMPar -> PDF.PDF (L.Vector Double) -> RIO [L.Vector Double]
runFixMetRio factor samples ampar pdf = do
    seed <- S.get
    (nseed, xs, amp) <- io $ go samples seed ampar []
    S.put nseed
    return xs
     where go (0::Int) s amp vs = return (s,vs, amp)
           go nn s amp vs = do let (!ampn, !ns) = unSam (fixedMet factor pdf amp) s
                               go (nn-1) ns (ampn) $ (ampPar ampn):vs

runFixMetRioBurn :: Double -> Int -> AMPar -> PDF.PDF (L.Vector Double) -> RIO AMPar
runFixMetRioBurn factor burn ampar pdf = do
    seed <- S.get
    (nseed, amp) <- io $ go burn seed ampar
    S.put nseed
    return amp
     where go (0::Int) s amp  = return (s,amp)
           go nn s amp  = do let (!ampn, !ns) = unSam (fixedMet factor pdf amp) s
                             go (nn-1) ns (ampn)

traceHead vs = trace ("head="++(show $ vs L.@> 0)) vs

most_es_per_s = 10

calcESS :: Int -> [L.Vector Double] -> Double
calcESS  want_ess mcmcOut
   | L.dim (head mcmcOut) <  want_ess * most_es_per_s * 2
      = calcESSprim mcmcOut
   | otherwise
      = let thinFactor = L.dim (head mcmcOut) `div` ( want_ess * most_es_per_s)
        in calcESSprim $ map (thinV $ thinFactor - 1) mcmcOut

thinV 0   =  id
thinV thin = V.ifilter $ \ix _ -> ix `mod` thin == 0

--thinV thin v = V.generate ((V.length v `div` (thin+1))+1) $ \i -> (V.!) v  (i*(thin))

calcESSprim ::  [L.Vector Double] -> Double
calcESSprim mcmcOut =
  let ndims = L.dim $ head mcmcOut
      len = realToFrac (length mcmcOut)
      acfs = mapP (\i->  V.sum $ V.takeWhile (>0.1) $ fst3 $ autocorrelation $ V.fromList $ map (L.@>i) mcmcOut) [0..ndims-1]
      ess = foldl1' min $ flip map acfs $ \acfsum-> (len/(1+2*acfsum)) -- realToFrac samples/(1+2*acfsum)
  in ess

mapP :: NFData b => (a-> b) -> [a] -> [b]
mapP f xs = runPar $
  forM xs (spawn . return . f) >>= mapM get

runAdaMetRIOtoFile :: Int -> Int -> String -> Bool -> AMPar -> PDF.PDF (L.Vector Double) -> RIO ()
runAdaMetRIOtoFile n thinn fileNm freeze ampar pdf = do
    seed <- S.get
    h <- io $ openFile (fileNm) WriteMode
    (nseed) <- io $ go h n seed ampar
    S.put nseed
    return ()
    --return xs
     where go h 0 s amp  = do print $ amp
                              return s --  return (s, reverse vs)
           go h nn s amp = do let (!ampn, !ns) = unSam (adaMet freeze pdf amp) s
                              when(nn `rem` chsz==0) $
                                   putStrLn $ show (((n-nn) `div` chsz)*2)++"%: " ++
                                              showV (ampPar ampn) ++" LH="++
                                              printf "%.3g" (pdf (ampPar ampn)) ++
                                              " accept="++acceptS ampn
                              when (nn `rem` thinn == 0) $ do
                                   hPutStrLn h $ show $ L.toList $ ampPar ampn
                              go h (nn-1) ns (ampn)
           chsz = n `div` 50

showV v = "<"++intercalate "," (map (printf "%.4g") $ L.toList v)++">"
acceptS ampar  | count ampar == 0 = "0/0"
               | otherwise = printf "%.3g" (rate::Double) ++ " ("++show yes++"/"++show total++")"where
   rate = realToFrac (yes) / realToFrac (total)
   yes = count_accept ampar
   total = count ampar

ampParRate ampar  | count ampar == 0 = 0
                  | otherwise = rate where
   rate = realToFrac (yes) / realToFrac (total)
   yes = count_accept ampar
   total = count ampar

runAdaMetRIOInterleaveInitial :: Int -> Bool -> L.Matrix Double -> AMPar -> PDF.PDF (L.Vector Double) -> RIO [L.Vector Double]
runAdaMetRIOInterleaveInitial n freeze cov ampar pdf = do
    seed <- S.get
    (nseed, xs) <- io $ go n seed ampar []
    S.put nseed
    return xs
     where go 0 s amp vs = do print $ amp
                              return (s, vs)
           go nn s amp vs = do let (!ampn, !ns) = unSam (adaMetInterleaveInitial freeze cov pdf amp) s
                               when(nn `rem` chsz==0) $
                                   putStrLn $ show (((n-nn) `div` chsz)*2)++"%: " ++show (ampPar ampn)-- ++" LH="++show (pdf (ampPar ampn))
                               go (nn-1) ns ampn $ (ampPar ampn):vs
           chsz = n `div` 50

runChainStat :: Int -> (Fold a b) -> a -> (a->Sampler a) -> RIO b
runChainStat n (F acc fldini k _) ini sam = do
    seed <- S.get
    let (nseed, stat) = go n seed fldini ini
    S.put nseed
    return stat
     where go 0 s y x = (s, k y)
           go nn s y x = let (xx, ns) = unSam (sam x) s
                         in go (nn-1) ns (acc y xx) xx

runAndWriteTo :: String -> Int -> (a -> String) -> a -> (a->Sampler a) -> RIO a
runAndWriteTo fnm n showit x sam = do
    h <- io $ openFile fnm WriteMode
    seed <- S.get
    (y,nseed) <- io $ runChainPrim h n showit seed x sam
    io $ hClose h
    S.put nseed
    return y

runChainPrim :: Handle -> Int -> (a -> String) -> Seed -> a -> (a->Sampler a) -> IO (a, Seed)
runChainPrim h n showit seed x sam = go n seed x
   where go 0 s x = return (x,s)
         go nn s x = do
            let (y, ns) = unSam (sam x) s
            hPutStrLn h $ showit y
            when(nn `rem` chsz==0) $
              putStrLn $ show (((n-nn) `div` chsz)*2)++"%: "++showit y
            go (nn-1) ns y
         chsz = n `div` 50

runOnce :: a -> (a->Sampler a) -> RIO a
runOnce init sam = runAndDiscardPrim 1 init sam

runAndDiscard :: Int -> (a->String) -> a -> (a->Sampler a) -> RIO a
--runAndDiscard = runAndDiscardPrim
runAndDiscard 0 _ x _ = return x
runAndDiscard n showit x sam = do
    go 50 x
 where chunksz = n `div` 50
       go 0 x = io (putStr "\n") >> return x
       go n x = do
          x' <- runAndDiscardPrim chunksz x sam
          io $ putStrLn $ show ((50-n)*2+2)++"%: "++showit x'
          go (n-1) x'

runMeanVar :: Floating b => Int -> (a->b) -> a -> (a->Sampler a) -> RIO (b,b)
runMeanVar n conv init sam = do
   xs <- sample $ runChainS n init sam
   return $ both meanF stdDevF `runStat` map conv xs

runAndDiscardPrim :: Int -> a -> (a->Sampler a) -> RIO a
runAndDiscardPrim n init sam = do
   seed <- S.get
   let (x,nseed) = runPrim seed n init sam
   S.put nseed
   return x

runPrim :: Seed -> Int -> a -> (a->Sampler a) -> (a,Seed)
runPrim s 0 x _ = (x,s)
runPrim s n x sam = let (x', s') = unSam (sam x) s
                    in runPrim s' (n-1) x' sam

metSample2P :: Show a => String -> (Double -> a -> a -> Sampler a)
               -> PDF.PDF (a,b) -> (Param a,b) -> Sampler (Param a)
metSample2P st prop p (par@(Param j t tt _ curw ada ini xi),y)
    = let accept xi pi pstar | notNanInf2 pi pstar =  min 1 $ exp (pstar - pi)
                             | otherwise = cond [(nanOrInf pi && nanOrInf pstar,
                                                        error $ "metropolisLnP "++st++" pi pi pstar :"++
                                                                show (pi,pstar)++"\n"++
                                                                show xi),
                                                 (nanOrInf pstar, -1), -- never accept
                                                 (nanOrInf pi, 2)] $ error "metropolisLn: the impossible happend"
      in  do
        let (nextw, nj, nt) = calcNextW ada curw j t tt
        u <-  unitSample
        xstar <- prop nextw ini xi
        let pstar = p (xstar,y)
        let pi = p (xi,y)
        return $ if u < accept par pi pstar
                      then Param (nj+1) (nt+1) (tt+1) pstar nextw ada ini xstar
                      else Param nj (nt+1) (tt+1) pi nextw ada ini xi
metSample2PC :: Show a => String -> (Double -> a -> a -> Sampler a)
               -> PDF.PDF (a,b) -> (Param a,b) -> Sampler (Param a)
metSample2PC st prop p (par@(Param j t tt lhi curw ada ini xi),y)
    = let accept xi pi pstar | notNanInf2 pi pstar =  min 1 $ exp (pstar - pi)
                             | otherwise = cond [(nanOrInf pi && nanOrInf pstar,
                                                        error $ "metropolisLnP "++st++" pi pi pstar :"++
                                                                show (pi,pstar)++"\n"++
                                                                show xi),
                                                 (nanOrInf pstar, -1), -- never accept
                                                 (nanOrInf pi, 2)] $ error "metropolisLn: the impossible happend"
      in  do
        let (nextw, nj, nt) = calcNextW ada curw j t tt
        u <-  unitSample
        xstar <- prop nextw ini xi
        let pstar = p (xstar,y)
        let pi = if notNanInf lhi then lhi else p (xi,y)
        return $ if u < accept par pi pstar
                      then Param (nj+1) (nt+1) (tt+1) pstar nextw ada ini xstar
                      else Param nj (nt+1) (tt+1) pi nextw ada ini xi

bestOfTwoCov inisam posterior = do
   set1@(postval1,_,_) <- getInitialCov inisam posterior
   set2@(postval2,_,_) <- getInitialCov inisam posterior
   if postval1>postval2 -- I THINK
      then return set1
      else return set2

mkAMPar init cov initp = AMPar init init cov 2.4 initp 10 5

initAdaMetFromCov nsam pdf initv retries cov = do
  --lift $ putStrLn $ "starting from existing cov; try number "++ show retries
  iniampar <- sample $ initialAdaMetFromCov nsam (pdf) initv
                                                 (PDF.posdefify $ cov)
  --lift $ print iniampar
  let rate = realToFrac (count_accept iniampar) / realToFrac nsam
  case () of
     _ | retries > 8 -> do lift $ putStrLn "initals ok."
                           return iniampar
     _ | rate > 0.5 -> initAdaMetFromCov nsam pdf initv (retries +1) $ L.scale 2 cov
     _ | rate > 0.40 -> initAdaMetFromCov nsam pdf initv (retries +1) $ L.scale 1.5 cov
     _ | rate < 0.025 ->  initAdaMetFromCov nsam pdf initv (retries +1) $ L.scale 0.1 cov
     _ | rate < 0.04 ->  initAdaMetFromCov nsam pdf initv (retries +1) $ L.scale 0.2 cov
     _ | rate < 0.12 ->  initAdaMetFromCov nsam pdf initv (retries +1) $ L.scale 0.3 cov
     _ | rate < 0.16 ->  initAdaMetFromCov nsam pdf initv (retries +1) $ L.scale 0.5 cov
     _ | rate < 0.20 ->  initAdaMetFromCov nsam pdf initv (retries +1) $ L.scale 0.8 cov
     _ | otherwise -> do lift $ putStrLn "initals ok."
                         return iniampar