{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE NoMonomorphismRestriction #-}

-- :set -fobject-code in GHCi
module Picture where

import Foreign hiding (rotate)
import Foreign.C.Types

import Control.Applicative

import Data.Complex

default (Int, Float, Integer)

-- Strict
data Color =
  Color
    { r, g, b, a :: !Float
    }
  deriving (Show)

data CColor =
  CColor
    { rC, gC, bC :: !CUChar
    }
  deriving (Show)

instance Storable CColor where
  alignment _ = 1
  sizeOf _ = 3
  peek ptr =
    CColor <$> peekByteOff ptr 0 <*> peekByteOff ptr 1 <*> peekByteOff ptr 2
  poke ptr (CColor d c i) = do
    pokeByteOff ptr 0 d
    pokeByteOff ptr 1 c
    pokeByteOff ptr 2 i

type Point = (Float, Float)

type Image a = Point -> a

type Region = Image Bool

type Filter c = Image c -> Image c

type FilterC = Filter Color

type ImageC = Image Color

renderRegion :: Region -> Image Color
renderRegion r p =
  if r p
    then black
    else white

cOver :: Color -> Color -> Color
cOver (Color r1 g1 b1 a1) (Color r2 g2 b2 a2) =
  Color (h r1 r2) (h g1 g2) (h b1 b2) (h a1 a2)
  where
    h x1 x2 = x1 + (1 - a1) * x2

bilerpBRBW = bilerpC black red blue white

bilerpC ll lr ul ur (wx, wy) = lerpC wy (lerpC wx ll lr) (lerpC wx ul ur)

over :: ImageC -> ImageC -> ImageC
top `over` bot = \p -> top p `cOver` bot p

instance Semigroup Color where
  (<>) = cOver

instance Monoid Color where
  mempty = invisible

invisible = Color 0 0 0 0

black = Color 0 0 0 1

white = Color 1 1 1 1

red = Color 1 0 0 1

green = Color 0 1 0 1

blue = Color 0 0 1 1

yellow = Color 1 1 0 1

cond :: Image Bool -> Image c -> Image c -> Image c
cond =
  liftA3
    (\a b c ->
       if a
         then b
         else c)

lerpI :: Image Frac -> ImageC -> ImageC -> ImageC
lerpI = liftA3 lerpC

empty = const invisible

whiteI = const white

blackI = const black

redI = const red

blueI = const blue

greenI = const green

yellowI = const yellow

blackWhiteIm reg = cond reg blackI whiteI

blueYellowIm reg = cond reg blueI yellowI

crop :: Region -> FilterC
crop reg im = cond reg im mempty

type Warp = Point -> Point

type Vector = (Float, Float)

translateP :: Vector -> Warp
translateP (dx, dy) (x, y) = (x + dx, y + dy)

scaleP :: Vector -> Warp
scaleP (sx, sy) (x, y) = (sx * x, sy * y)

uscaleP :: Float -> Warp
uscaleP s = scaleP (s, s)

rotateP :: Float -> Warp
rotateP θ (x, y) = (x * cos θ - y * sin θ, y * cos θ + x * sin θ)

udisk :: Region
udisk p = distO p < 1

invWarp warp im = im . warp

translate :: Vector -> Filter c
translate (dx, dy) = invWarp (translateP (-dx, -dy))

scale :: Vector -> Filter c
scale (sx, sy) = invWarp (scaleP (1 / sx, 1 / sy))

-- Laws
-- uscale x . uscale y = uscale (x * y)
uscale :: Float -> Filter c
uscale s = invWarp (uscaleP (1 / s))

rotate :: Float -> Filter c
rotate θ = invWarp (rotateP (-θ))

swirlP :: Float -> Warp
swirlP r p = rotateP (distO p * (2 * pi / r)) p

swirl :: Float -> Filter c
swirl r = invWarp (swirlP (-r))

square n (x, y) = abs x <= n / 2 && abs y <= n / 2

-- Figure to visually check orientation
disks =
  renderRegion $
  f <$> udisk <*> translate (2, 2) udisk <*> translate (3, 2) udisk <*> vstrip
  where
    f a b c d = a || b || c || d

-- Figures from "Functional Images" by Conal Elliot.
fig1 = renderRegion $ vstrip

fig2 = renderRegion $ checker

fig3 = renderRegion $ altRings

fig4 = renderRegion $ polarChecker 10

fig5 = renderRegion $ uscale 0.05 gasket

fig6 = renderFrac $ wavDist

fig7 = bilerpBRBW

fig8 = lerpI wavDist (blackWhiteIm (polarChecker 10)) (blueYellowIm checker)

fig9 = ybRings

fig10 = renderRegion udisk

fig11 = renderRegion $ swirl 1 vstrip

fig13 t = renderRegion (swirlingXPos t)

fig14 = renderRegion $ annulus 0.5

fig17 = renderRegion $ shiftXor 2.6 altRings

fig18 = renderRegion $ xorgon 8 (7 / 4) altRings

-- Typo in original paper, 0.2 5 should be 0.25
fig19 = crop (wedgeAnnulus 0.25 10) ybRings

fig20 = crop (swirl 2 (wedgeAnnulus 0.25 10)) ybRings

fig22 = translate (-20, -20) tiledBilerp

fig23 = renderRegion $ uscale 3.5 (radInvert checker)

fig24 = rippleRad 8 0.3 ybRings

fig25 = \t -> rippleRad 8 (cos t / 2) ybRings

fig26 = rippleRad 8 0.3 $ cropRad 1 $ ybRings

fig27 = cropRad 1 $ rippleRad 8 0.3 $ ybRings

fig28 = swirl 8 $ rippleRad 5 0.3 $ cropRad 5 $ ybRings

-- Typo in original paper, washer (1/2) (pi/2) 1 tiledBilerp
-- TODO: Fix visual mismatch with paper.
fig32 = washer (1 / 2) (pi / 2) tiledBilerp

-- Convert an image to a static animation.
static :: Image c -> Anim c
static = const

-- The main animation to run.
-- mainAnim = static . renderFrac $ mandel
mainAnim :: Anim Color
mainAnim = fig32

calculate (x, y, t) = adjust (mainAnim t') (x', y')
  where
    duration = 5000 -- Duration of the animation in ms
    range = 2 * pi -- Time range of the animation
    -- Time normalized from 0 to 1
    normalizedTime = fromIntegral (t `mod` duration) / fromIntegral duration
    scaledTime = range * normalizedTime
    -- scaledTime = fromIntegral t / 10000
    (x', y', t') = (fromIntegral x, fromIntegral y, scaledTime)
    adjust = adjustToWindow
    -- Adjust an image to a window, by translating, scaling and flipping
    adjustToWindow :: Filter c
    adjustToWindow =
      translate (screenWidth / 2, screenHeight / 2) .
      uscale 60 . flipY
    flipY p (x, y) = p (x, -y)

(screenWidth, screenHeight) = (640, 480)

foreign export ccall fillPixelBuffer :: Ptr CColor -> CInt -> IO ()
fillPixelBuffer arr t = pokeArray arr (map calc l1)
  where
    l1 = [(x, y) | y <- [0 .. screenHeight - 1], x <- [0 .. screenWidth - 1]]
    calc (x, y) =
      CColor (truncate (255 * r)) (truncate (255 * g)) (truncate (255 * b))
      where
        Color r g b _ = calculate (x, y, t)

checker :: Region
checker (x, y) = even (floor x + floor y)

vstrip :: Region
vstrip (x, y) = abs x <= 1 / 2

overlay :: ImageC -> ImageC -> ImageC
overlay = liftA2 (<>)

-- Sierpinski triangle
gasket :: Region
gasket (x, y) = floor x .|. floor y == (floor x :: Integer)

altRings p = even (floor (distO p))

distO (x, y) = sqrt (x * x + y * y)

type Frac = Float -- in [0, 1]

wavDist :: Image Frac
wavDist p = (1 + cos (pi * distO p)) / 2

ybRings = lerpI wavDist blueI yellowI

renderFrac :: Image Frac -> ImageC
renderFrac f (x, y) = lerpC (f (x, y)) black white

lerpC :: Frac -> Color -> Color -> Color
lerpC w (Color r1 g1 b1 a1) (Color r2 g2 b2 a2) =
  Color (h r1 r2) (h g1 g2) (h b1 b2) (h a1 a2)
  where
    h x1 x2 = w * x1 + (1 - w) * x2

lighten x c = lerpC x c white

darken x c = lerpC x c black

type PolarPoint = (Float, Float)

fromPolar :: Point -> PolarPoint
fromPolar (ρ, θ) = (ρ * cos θ, ρ * sin θ)

toPolar :: PolarPoint -> Point
toPolar (x, y) = (distO (x, y), atan2 y x)

polarChecker :: Int -> Region
polarChecker = polarize checker

polarize :: Image c -> Int -> PolarPoint -> c
polarize pat n = pat . sc . toPolar
  where
    sc (ρ, θ) = (ρ, θ * n' / π)
    π = pi
    n' = fromIntegral n

-- Animations
type Time = Float

type Anim c = Time -> Image c

swirlingXPos :: Anim Bool
swirlingXPos t = swirl (t * t) xPos

xPos :: Region
xPos (x, y) = x > 0

-- Region algebra
intersect :: Region -> Region -> Region
intersect = liftA2 (&&)

union :: Region -> Region -> Region
union = liftA2 (||)

xorR :: Region -> Region -> Region
xorR = liftA2 xor
  where
    xor True b = not b
    xor _ b = b

compR :: Region -> Region
compR = fmap not

universeR :: Region
universeR = const True

emptyR :: Region
emptyR = const False

r \\ r' = r `intersect` compR r'

annulus :: Frac -> Region
annulus inner = udisk \\ uscale inner udisk

radReg :: Int -> Region
radReg n = test . toPolar
  where
    test (r, a) = even (floor (a * fromIntegral n / pi))

wedgeAnnulus :: Frac -> Int -> Region
wedgeAnnulus inner n = annulus inner `intersect` radReg n

shiftXor :: Float -> Region -> Region
shiftXor r reg = reg' r `xorR` reg' (-r)
  where
    reg' d = translate (d, 0) reg

-- Typo in the original paper, missing argument g
xorgon :: Int -> Float -> Region -> Region
xorgon n r g = xorRs (rf <$> [0 .. n - 1])
  where
    rf :: Int -> Region
    rf i = translate (fromPolar (r, a)) g
      where
        a = fromIntegral i * 2 * pi / fromIntegral n

xorRs :: [Region] -> Region
xorRs = foldr xorR emptyR

tileP :: Vector -> Warp
tileP (w, h) (x, y) = (wrap' w x, wrap' h y)

tiledBilerp = about (1 / 2, 1 / 2) (tile (1, 1)) bilerpBRBW

type HyperFilter c = Filter c -> Filter c

about :: Point -> HyperFilter c
about (x, y) filt = translate (x, y) . filt . translate (-x, -y)

wrap :: Float -> Float -> Float
wrap w x = w * fracPart (x / w)
  where
    fracPart t = t - fromIntegral (truncate t)

wrap' :: Float -> Float -> Float
wrap' w x = wrap w (x + w / 2) - w / 2

tile :: Vector -> Filter c
tile size = invWarp (tileP size)

swirlP' r = polarWarp (\(ρ, θ) -> (ρ, θ + ρ * (2 * pi / r)))

polarWarp warp = fromPolar . warp . toPolar

radInvertP :: Warp
radInvertP = polarWarp (\(ρ, θ) -> (1 / ρ, θ))

radInvert :: Filter c
radInvert = invWarp radInvertP

rippleRadP :: Int -> Float -> Warp
rippleRadP n s =
  polarWarp $ (\(ρ, θ) -> (ρ * (1 + s * sin (fromIntegral n * θ)), θ))

rippleRad :: Int -> Float -> Filter c
rippleRad n s = invWarp (rippleRadP n (-s))

cropRad :: Float -> FilterC
cropRad r = crop (uscale r udisk)

wiggleRotateP :: Float -> Float -> Time -> Warp
wiggleRotateP cycles θmax t = polarWarp warp
  where
    warp (r, a) = (r, a + θmax * sin (t + dt))
      where
        dt = 2 * pi * cycles * (r - 1 / 2)

wiggleRotate :: Float -> Float -> Time -> Filter c
wiggleRotate cycles θmax t = invWarp (wiggleRotateP cycles θmax t)

washer :: Float -> Float -> ImageC -> Time -> ImageC
washer cycles θmax im t = cropRad 1 $ wiggleRotate cycles θmax t $ im

-- Constants
maxIter :: Int -- Max iterations
maxIter = 750

width :: Int -- Image width
height :: Int -- Image height
width = 400

height = 400

-- Note: aspect ratio of (minX, minY), (maxX, maxY) must
-- match aspect ratio of (width, height)
minX :: Float -- Min x-coordinate of graph
maxX :: Float -- Max x-coordinate of graph
minY :: Float -- Min y-coordinate of graph
maxY :: Float -- Max y-coordinate of graph
-- For the zoomed in part of the Mandelbrot:
--minX = -0.826341244461360116
--maxX = -0.8026423086165848822
--minY = -0.2167936114403439588
--maxY = -0.193094675595568725
--For a full view of the mandelbrot
minX = -2.5

maxX = 1.5

minY = -2

maxY = 2

-- The actual fractal part. It basically works on a matrix, which we
-- will call M, that represents a grid of points on the
-- graph. Essentially, M[i, j] is (xList[j], yList[i])
xList :: [Float]
yList :: [Float]
xList = [minX,(minX + ((maxX - minX) / (fromIntegral width - 1))) .. maxX]

yList =
  reverse [minY,(minY + ((maxY - minY) / (fromIntegral height - 1))) .. maxY]

etaFraction :: Complex Float -> Float
etaFraction z = (log (log (magnitude z))) / log 2

smoothEta :: Int -> Complex Float -> Float -- Smooth escape time algorithm value
smoothEta iter z = (fromIntegral iter - etaFraction z) / fromIntegral maxIter

-- Gets the color for the point, in range [0, 1]
color :: Int -> Complex Float -> Float
color iter z = smoothEta iter z -- Smooth escape time algorithm (and invert)

-- color iter z = fromIntegral iter / fromIntegral maxIter
interpolate :: Float -> Int -- Adds an interpolation curve for interpolating color
interpolate v = truncate ((v ^ 12) * 255) -- Polynomial curve

-- interpolate v = chr (truncate(v * 255)) -- Linear
-- The actual fractal algorithm!
frac :: Complex Float -> Complex Float -> Int -> Int
frac c@(x :+ y) z iter
  | iter >= maxIter = 255 -- never escaped, return color value of 255
  | let p = sqrt ((x - 0.25) ^ 2 + y ^ 2)
     in x <= p - 2 * p ^ 2 + 0.25 && (x + 1) ^ 2 + y ^ 2 <= 0.625 = 255
  | otherwise =
    let z' = z * z + c
     in if ((realPart z') * (realPart z') + (imagPart z') * (imagPart z')) > 4
          then interpolate (color iter z')
          else frac c z' (iter + 1)

mandel :: Image Frac
mandel (x, y) = fromIntegral (frac (x :+ y) (0 :+ 0) 0) / 255

-- dotGridGradient :: Int -> Int -> Float -> Float -> IO Float
-- dotGridGradient ix iy x y = do
--   (v1, v2) <- do p' <- randomIO :: IO Float
--                  let (_, p) = properFraction p'
--                  pure (p, sqrt (1 - p ^ 2))
--   let c1 = dx * v1
--   let c2 = dy * v2
--   let res = (c1 + c2)
--   -- print res
--   pure res
--   where
--     dx = x - fromIntegral ix
--     dy = y - fromIntegral iy

-- perlin :: Image (IO Frac)
-- perlin (x, y) = do
--   let (x0, sx) = properFraction x
--   let (y0, sy) = properFraction y
--   let x1 = x0 + 1
--   let y1 = y0 + 1
--   n0 <- dotGridGradient x0 y0 x y
--   n1 <- dotGridGradient x1 y0 x y
--   let ix0 = lerp n0 n1 sx
--   n0' <- dotGridGradient x0 y1 x y
--   n1' <- dotGridGradient x1 y1 x y
--   let ix1 = lerp n0' n1' sx
--   let res' = lerp ix0 ix1 sy
--   let res = abs res'
--   pure res
--   where
--     lerp a0 a1 w = (1.0 - w) * a0 + w * a1


{-|
Module      : MachineLearning.TIR
Description : TIR expression data structures
Copyright   : (c) Fabricio Olivetti de Franca, 2022
License     : GPL-3
Maintainer  : fabricio.olivetti@gmail.com
Stability   : experimental
Portability : POSIX

The TIR expression  represents a function of the form:

\[
f(x) = g(\sum_{i}{w_i \cdot t_i(\prod_{j}{x_j^{k_{ij}})}} / (1 + \sum_{i}{w_i \cdot t_i(\prod_{j}{x_j^{k_{ij}})}})
\]

with \(t_i\) being a transformation function, \(g\) an invertible function, \(w_i\) a linear coefficient, \(k_{ij}\)
the interaction strength.

Any given expression can be represented by two lists of terms, with each term
being composed of a transformatioon function and an interaction.
The transformation function is represented by a `Function` sum type.
The interaction is represented as a list of tuple of ints where the key is the
predictor index and the value is the strength of the predictor in this
term. Strengths with a value of zero are omitted.
-}
module MachineLearning.TIR where

import Control.Evolution
import Control.Monad.State.Strict
import Data.List.Split
import Data.SRTree
import System.Random
import Control.DeepSeq                           (NFData, rnf)
import Data.List                                 (delete)
import Data.SRTree.Print                         (showDefault)
import Data.Vector                               (Vector)
import qualified Data.Vector             as V
import qualified Data.Vector.Storable    as VS
import qualified Numeric.LinearAlgebra   as LA
import MachineLearning.Utils.Config

-- | `TIR` is a record type composed of the external
-- function `_funY` of type `Function`, numerator `_p`
-- and denominator `_q` of type `Sigma`
data TIR = TIR { _funY :: Function
               , _p :: Sigma
               , _q :: Sigma
               } deriving Show

instance NFData TIR where
  rnf _ = ()

-- | `Sigma` is just a list of terms `Pi`
type Sigma = [Pi]
-- | `Pi` is a triple composed of a coefficient, a `Function`
-- and a list of tuples where `(ix, k)` represents `x ! ix ^ k`.
type Pi    = (Double, Function, [(Int, Int)])

-- | generates a random integer within the specified range.
randomRng :: (Int, Int) -> Rnd Int
randomRng rng = state $ randomR rng
{-# INLINE randomRng #-}

-- | generates a random integer within the specified range excluding zero.
randomRngNZ :: (Int, Int) -> Rnd Int
randomRngNZ rng = do
  x <- randomRng rng
  if x == 0
    then randomRngNZ rng
    else pure x
{-# INLINE randomRngNZ #-}

-- | picks a random element from a list.
randomFrom :: [a] -> Rnd a
randomFrom xs = do
  ix <- randomRng (0, length xs - 1)
  pure (xs !! ix)
{-# INLINE randomFrom #-}

-- | returns a random index of variables provided by the mutation configuration.
randomVar :: MutationCfg -> Rnd (Maybe Int, MutationCfg)
randomVar params = do
  let vars = _vars params
      n    = length vars
  ix <- randomRng (0, n)
  if ix == n
     then pure (Nothing, params)
     else do let x = vars !! ix
             pure (Just x, params{ _vars=delete x vars })

-- | returns a list of random interactions (tuples of variables indeces and exponentes)
-- with parameters provided by the mutation configuration.
randomVars :: MutationCfg -> Rnd [(Int, Int)]
randomVars params = do
  (v, params') <- randomVar params
  k            <- randomRngNZ $ _kRange params
  case v of
    Nothing  -> pure []
    Just var -> do vs <- randomVars params'
                   pure $ (var, k) : vs

-- | returns a random `Pi`
randomPi :: MutationCfg -> Rnd (Maybe Pi)
randomPi params = do
  pis <- randomVars params
  f   <- randomFrom $ _funs params
  if null pis
    then pure Nothing
    else pure $ Just (1.0, f, pis)

-- | returns a random `Sigma`
randomSigma :: MutationCfg -> Int -> Rnd (Sigma, Int)
randomSigma params budget | budget <= 0 = pure ([], budget)
randomSigma params budget = do
  n <- randomRng (0, budget)
  if n == budget
     then pure ([], budget)
     else do term             <- randomPi params
             (terms, budget') <- randomSigma params (budget - spentBudget term)
             case term of
               Nothing -> pure (terms, budget')
               Just t  -> pure (t:terms, budget')

  where
    spentBudget Nothing           = 0
    spentBudget (Just (_, _, ps)) = 1 -- length ps

-- | returns a random `TIR` expression
randomTIR :: MutationCfg -> Rnd TIR
randomTIR params = do
  yf           <- randomFrom $ _yfuns params
  (p, budget') <- randomSigma params $ _budget params
  (q, _)       <- randomSigma params budget'
  if null p
    then randomTIR params
    else pure (TIR yf p q)

-- | We store thee dataset as a vector of columns.
-- Each vector is stored a `Storable`-based vector.
type Column a   = LA.Vector a

-- | A dataset is a `Vector` of `Column`
type Dataset a  = Vector (Column a)

-- | A constraint is a function that gets a symbolic tree
-- as an input and returns non negative `Double` representing
-- how much a constraint was violated.
type Constraint = SRTree Int Double -> Double

-- | An individual in the population is composed of
-- the chromossome, a vector of fitness, a list of
-- coefficients (for multiclass problems it stores
-- one vector of coefficient per class),
-- the constraint violation, the size of the expression,
-- and the penalty value.
data Individual = Individual { _chromo  :: TIR
                             , _fit     :: [Double]
                             , _weights :: [LA.Vector Double]
                             , _constr  :: Double
                             , _len     :: Int
                             , _penalty :: Double
                             }

-- | creates an unevaluated individual.
createIndividual :: TIR -> Individual
createIndividual tir = Individual tir [] [] 0.0 0 0.0

-- | calculates the penalized fitness.
penalizedFit :: Individual -> Double
penalizedFit t = (head . _fit) t + _penalty t
{-# INLINE penalizedFit #-}

-- | replaces the coefficients of a TIR expression
replaceConsts :: TIR -> V.Vector Double -> TIR
replaceConsts (TIR g p q) ws = TIR g p' q'
  where
    (p', ws1) = runState (traverse replaceWeight p) (V.toList ws)
    (q', ws2) = runState (traverse replaceWeight q) ws1

replaceWeight :: Pi -> State [Double] Pi
replaceWeight (w, g, h) = state $ \ws -> case ws of
                                           (wi:ws') -> ((wi, g, h), ws')
                                           []       -> error $ show h -- ((w, g, h), [])

instance Eq Individual where
    t1 == t2 = penalizedFit t1 == penalizedFit t2
instance Ord Individual where
    t1 <= t2 = penalizedFit t1 <= penalizedFit t2

instance NFData Individual where
  rnf _ = ()

instance Solution Individual where
  _getFitness = head . _fit
  _isFeasible = (<1e-12) . _constr

-- | creates a symbolic tree from a TIR expression.
assembleTree :: Double -> TIR -> SRTree Int Double
assembleTree bias (TIR f p q) = Fun f ((Const bias + assemble p) / (1 + assemble q))
  where
    -- assemble :: Sigma ix val -> SRTree ix val
    assemble []      = 0
    assemble [p']    = mk p'
    assemble (p':ps) = mk p' + assemble ps

    -- mk :: Pi ix val -> SRTree ix val
    mk (v, g, ts) = Const v * Fun g (foldr (\(ix, k) acc -> acc * Pow (Var ix) k) 1 ts)

-- | pretty print a solution.
prettyPrintsolution :: Individual -> String
prettyPrintsolution sol | Prelude.null (_fit sol) = error "unevaluated solution"
prettyPrintsolution sol = concat [ "Expression:\n", (showDefault . assembleTree bias . _chromo) sol, "\n"
                                 , "Fitness: ", (show . head . _fit) sol, "\n"
                                 , "Constraints: ", (show . _constr) sol, "\n"
                                 , "Length: ", (show . _len) sol, "\n"
                                 , "Penalty: ", (show . _penalty) sol, "\n"
                                 ]

  where bias = V.head $ VS.convert $ head $ _weights sol


{-# OPTIONS_GHC -fno-warn-type-defaults #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Basic types used throughout the library.
--
--   Includes a probability type (based on sampling functions, with a tweak for
--   collected samples) and associated instances, as well as types used for
--   MCMC.

module Math.Probably.Types where

import Control.Applicative
import Control.Monad.State.Strict
import Numeric.LinearAlgebra
import System.Random.Mersenne.Pure64

type Seed = PureMT

data Prob a =
    Sampler { unSampler :: Seed -> (a, Seed) }
  | Samples [a]

instance Functor Prob where
  fmap f (Sampler sf) = Sampler $ \rs -> let (x,rs') = sf rs in (f x, rs')
  fmap f (Samples xs) = Samples $ map f xs

instance Applicative Prob where
  pure x = Sampler (\rs-> (x, rs))
  (Sampler sff) <*> (Sampler sfx) = Sampler $ \rs ->
    let (f ,rs') = sff rs
        (x, rs'') = sfx rs'
    in (f x, rs'')
  _ <*> _ = error "Prob (<*>): unsupported pattern"

instance Monad Prob where
  return = pure
  (Sampler sf) >>= f = Sampler $ \rs->
    let (x, rs'::Seed) = sf rs
        nextProb = f x
    in case nextProb of
         Sampler g -> g rs'
         Samples xs -> primOneOf xs rs'

  (Samples xs) >>= f = Sampler $ \rs->
    let (x, rs'::Seed) = primOneOf xs rs
        nextProb = f x
    in case nextProb of
         Sampler g -> g rs'
         Samples ys -> primOneOf ys rs'

primOneOf :: [a] -> Seed -> (a, Seed)
primOneOf xs seed
  = let (u, nextSeed) = randomDouble seed
        idx = floor $ realToFrac u * realToFrac (length xs )
    in (xs !! idx, nextSeed)

type DiscreteParams   = Vector Int
type ContinuousParams = Vector Double
type Parameters       = (DiscreteParams, ContinuousParams)

type LogObjective     = Parameters -> Double
type Gradient         = ContinuousParams -> ContinuousParams

type Particle         = (ContinuousParams, ContinuousParams)

type Transition t     = StateT (Chain t) Prob Parameters

-- | State of a Markov chain.  Note that the objective function itself is
--   included in the state, which allows the possibility of annealing schedules
--   and the like.
data Chain t = Chain {
    parameterSpacePosition :: Parameters
  , objectiveFunction      :: Target
  , objectiveValue         :: Double
  , tunables               :: t
  }

-- | A target to sample, consisting of a log objective function and possibly
--   its gradient.  The gradient is only taken with respect to continuous
--   parameters.
data Target = Target {
    logObjective :: Parameters -> Double
  , gradient     :: Maybe (ContinuousParams -> ContinuousParams)
  }

-- | Convenience constructor for targets and gradients.
createTargetWithGradient :: LogObjective -> Gradient -> Target
createTargetWithGradient f g = Target f (Just g)

-- | Convenience constructor for targets without gradients.
createTargetWithoutGradient :: LogObjective -> Target
createTargetWithoutGradient f = Target f Nothing

handleGradient :: Maybe t -> t
handleGradient Nothing  = error "handleGradient: no gradient provided"
handleGradient (Just g) = g

-- | The dual-averaging implementation of NUTS requires quite a few tuning
--   parameters, which are all held in this type.
data DualAveragingParameters = DualAveragingParameters {
    mAdapt    :: !Int
  , delta     :: !Double
  , mu        :: !Double
  , gammaP    :: !Double
  , tau0      :: !Double
  , kappa     :: !Double
  , daStep    :: !Double
  , daStepAvg :: !Double
  , daH       :: !Double
  } deriving (Eq, Show)

defaultDualAveragingParameters :: Double -> Int -> DualAveragingParameters
defaultDualAveragingParameters e burnInPeriod = DualAveragingParameters {
    mu        = log (10 * e)
  , delta     = 0.5
  , mAdapt    = burnInPeriod
  , gammaP    = 0.05
  , tau0      = 10
  , kappa     = 0.75
  , daStep    = e
  , daStepAvg = e
  , daH       = 0
  }


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TemplateHaskell       #-}

module Lib where

import           Control.Monad.Except                      (MonadError,
                                                            runExceptT,
                                                            throwError)
import           Control.Monad.IO.Class                    (MonadIO, liftIO)
import           Data.ByteString                           (ByteString)
import           Data.FileEmbed                            (embedFile)
import           Data.Monoid                               ((<>))
import           Data.Scientific                           (Scientific,
                                                            toRealFloat)
import           Debug.Trace
import           GHC.TypeLits                              ()
import           Graphics.Plot                             (mplot)
import           Graphics.Rendering.Chart.Backend.Diagrams (toFile)
import           Graphics.Rendering.Chart.Easy             (def, layout_title,
                                                            line, plot, points,
                                                            re, (.=))
import           Numeric.GSL.Minimization                  (MinimizeMethod (NMSimplex),
                                                            MinimizeMethodD (SteepestDescent, VectorBFGS2),
                                                            minimizeD,
                                                            minimizeV,
                                                            minimizeVD)
import           Numeric.LinearAlgebra                     (Container, accum,
                                                            matFunc, tr, ( #> ),
                                                            (<.>))
import qualified Numeric.LinearAlgebra                     as Matrix
import           Numeric.LinearAlgebra.Data                (Matrix, Vector,
                                                            asColumn, asRow,
                                                            cols, cond,
                                                            dropColumns,
                                                            dropRows, rows,
                                                            takeColumns,
                                                            takeRows, (><),
                                                            (|||))
import qualified Numeric.LinearAlgebra.Data                as MData
import           Numeric.LinearAlgebra.Static              (L, dim, matrix, mul)
import qualified Numeric.LinearAlgebra.Static              as Static
import           Text.Megaparsec                           (Dec, ParseError,
                                                            Token, char,
                                                            newline, parse,
                                                            sepBy)
import           Text.Megaparsec.ByteString                (Parser)
import           Text.Megaparsec.Lexer                     (number)

data Error =
  ParseFailed (ParseError (Token ByteString) Dec)
  deriving (Show, Eq)

------------------------------------------------------------
-- Parsing.
rawDataParser :: Parser [[Scientific]]
rawDataParser =
  filter (not . null) <$> (number `sepBy` char ',') `sepBy` newline

toMatrix :: [[Scientific]] -> Matrix Double
toMatrix = Matrix.fromLists . fmap (fmap toRealFloat)

wineData :: ByteString
wineData = $(embedFile "data/wine.data")

wineParser :: Parser (Matrix Double)
wineParser = toMatrix <$> rawDataParser

loadData
  :: (MonadError Error m, MonadIO m)
  => m (Matrix Double)
loadData = either (throwError . ParseFailed) pure (parse wineParser "" wineData)

------------------------------------------------------------
-- Matrices
a :: L 3 2
a = matrix [1 .. 6]

b :: L 2 1
b = matrix [1 .. 2]

c :: L 3 1
c = mul a b

------------------------------------------------------------
-- Linear Regression
initialTheta :: Int -> Vector Double
initialTheta n = MData.vector $ replicate n 0

------------------------------------------------------------
-- Charts
signal :: [Double] -> [(Double, Double)]
signal xs =
  [(x, (sin (x * 3.14159 / 45) + 1) / 2 * sin (x * 3.14159 / 5)) | x <- xs]

writeChart :: IO ()
writeChart =
  toFile def "example.svg" $ do
    layout_title .= "Amplitude Modulation"
    plot (line "am" [signal [0,0.5 .. 800]])
    plot (points "am points" (signal [0,7 .. 800]))

------------------------------------------------------------
-- Main
main :: IO ()
main = do
  result <-
    runExceptT $ do
      dataset <- loadData
      process dataset
  print result

data Result = Result
  { crossValidationSet    :: Matrix Double
  , crossValidationResult :: Vector Double
  , finalTheta            :: Vector Double
  , finalHyp              :: Vector Double
  , finalCost             :: Double
  } deriving (Show)

process
  :: (MonadIO m, MonadError Error m)
  => Matrix Double -> m Result
process dataSet = do
  let is n v = cond v n 0 1 0
  let target = 1
  let scale :: Matrix Double =
        accum
          (MData.ident 13)
          (*)
          [((0, 0), 0.1), ((3, 3), 0.1), ((4, 4), 0.01), ((12, 12), 0.001)]
  let m = rows dataSet
  let f = cols dataSet
  let y = is target $ MData.flatten $ takeColumns 1 dataSet
  let x = Matrix.col (replicate m 1.0) ||| ((dropColumns 1 dataSet) <> scale)
  let cvSize = 10
  let trainingX = dropRows cvSize x
  let cvX = takeRows cvSize x
  let trainingY = MData.fromList $ drop cvSize $ MData.toList y
  let cvY = MData.fromList $ take cvSize $ MData.toList y
  let theta = initialTheta f
  let (finalTheta, path) =
        minimizeVD
          SteepestDescent -- VectorBFGS2
          10e-6
          50
          10e-6
          0.1
          (costFn x y)
          (gradFn x y)
          theta
  pure $
    Result
      cvX
      cvY
      finalTheta
      (hypothesis cvX finalTheta)
      (costFn cvX cvY finalTheta)

sigmoid
  :: Floating a
  => a -> a
sigmoid term = 1.0 / (1.0 + exp (-term))

hypothesis :: Matrix Double -> Vector Double -> Vector Double
hypothesis x theta = sigmoid (x #> theta)

costFn :: Matrix Double -> Vector Double -> Vector Double -> Double
costFn x y theta = traceShowId $ all
  where
    all :: Double
    all = vsum (leftTerm - rightTerm) / fromIntegral (rows x)
    h = hypothesis x theta
    leftTerm = (-y) * log h
    rightTerm = (1.0 - y) * log (1.0 - h)

gradFn :: Matrix Double -> Vector Double -> Vector Double -> Vector Double
gradFn x y theta = (tr x #> (hypothesis x theta - y)) / fromIntegral (rows x)

vsum :: Vector Double -> Double
vsum = sum . MData.toList


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE TypeApplications    #-}

{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

module Neldermead (
  Simplex(..)
  , haftkaGurdalSimplex
  ) where

import Data.Ord (comparing)
import Data.Function (on)
import Data.Bool (bool)
import Data.Maybe (fromMaybe)
import Data.Monoid (Sum)
import qualified Data.Foldable as FLD
import qualified Data.Proxy as PXY
import Data.Coerce

import Control.Arrow ((&&&))
import Control.Monad.State.Strict

import qualified Numeric.Interval.NonEmpty as IVL

import qualified Data.Finite as F
import Data.Singletons (Sing, sing, fromSing)
import Data.Singletons.TypeLits

import qualified Data.Vector.Sized as V
import qualified Data.Vector.Storable.Sized as SV
--import qualified Data.Vector.Mutable.Sized as MV
import Numeric.LinearAlgebra.Static
import Numeric.LinearAlgebra.Static.Vector
import GHC.TypeNats

import qualified Control.Lens as L
import qualified Control.Lens.Iso as LI
import qualified Control.Lens.Traversal as LT

--TODO in future - mark Vertex with the objective function
--in any constructor, must take objective and then mark it

--TODO in future, allow more flexibility in objective?
--TODO better field names or allow overloaded
--TODO lenses
data NMVertex n = NMVertex { getX :: !(R n)
                           , getFx :: ℝ
                           }

newtype NMSimplex n = NMSimplex { fOrderedVertices :: V.Vector (n + 1) (NMVertex n) }

newtype Simplex n = Simplex { xA :: L n (n + 1) }
    deriving (Show)

newtype BoxConstraints n = BoxConstraints { unBoxConstraints :: V.Vector n (IVL.Interval ℝ) }

data NMScalars = NMScalars {
  reflectWeight :: ℝ
  , expandWeight :: ℝ
  , contractWeight :: ℝ
  , shrinkWeight :: ℝ
}

--TODO convergeTest should be monadic to allow writing details
--
data NMEnv n = MEnv {
  scalars :: NMScalars
  , simplexConvergeTest :: NMSimplex n -> Bool
}

newtype OnBounds = OnBounds { unOnBounds :: Bool }

makeNMVertex :: KnownNat n => (R n -> ℝ) -> R n -> NMVertex n
makeNMVertex f x = NMVertex { getX = x, getFx = f x }

--TODO actually need to change case 2 to keep if better than second worst (not best)
nextNMSimplex :: forall n. (KnownNat n, 1 <= n)
  => NMScalars -> (R n -> ℝ) -> NMSimplex n -> NMSimplex n
nextNMSimplex env f (NMSimplex vtxs)
  | comparing getFx vtxReflect vtxl == LT = undefined --expand case
  | comparing getFx vtxReflect vtxl2 == LT = undefined -- just keep swap out reflect and re-order
  | otherwise = undefined --shrink case
  where
    vtxl = V.head vtxs
    vtxl2 = V.index vtxs $ F.natToFinite (PXY.Proxy @1)
    vtxh = V.last vtxs
    xc = centroid $ getX <$> FLD.toList vtxs
    vtxReflect = makeNMVertex f $ affineCombWith (getX vtxh) (reflectWeight env) xc

--TODO may need MonadReader and MonadState, and allow objective to be in any monad
rVecIso :: KnownNat n => LI.Iso' (R n) (SV.Vector n ℝ)
rVecIso = LI.iso rVec vecR

projectToBox :: KnownNat n => BoxConstraints n -> R n -> R n
projectToBox (BoxConstraints bc) =
  L.over rVecIso $ SV.imap $ IVL.clamp . V.index bc

-- Return bool is weather any dimension was clamped
projectToBoxState :: KnownNat n
  => BoxConstraints n -> R n -> State OnBounds (R n)
projectToBoxState (BoxConstraints bc) = L.mapMOf rVecIso
  $ SV.imapM $ \i x -> do
    let bci = V.index bc i
    modify $ coerce (|| x `IVL.notMember` bci)
    return $ IVL.clamp bci x

--awkward because cant deconstruct interval into inf sup
--and cant zip storable and not storable
rightAngledSimplexInBox :: forall n. KnownNat n
  => BoxConstraints n -> R n -> R n -> Simplex n
rightAngledSimplexInBox (BoxConstraints bc) step x0 = Simplex
  $ ((x0 `outer` 1) + diag step') ||| col x0
  where
    (step' :: R n) = vecR
      $ SV.generate
      $ \i -> let stepBounds = V.index bc i - IVL.singleton (SV.index vx0 i)
                  si = abs $ SV.index vstep i
              in clampStep ((IVL.inf &&& IVL.sup) stepBounds) si
    clampStep (stepLb, stepUb) stepi
      | stepi <= stepUb = stepi
      | -stepi >= stepLb = -stepi
      | -stepLb < stepUb = max 0 $ 0.5 * stepUb
      | otherwise = min 0 $ 0.5 * stepLb
    vstep = rVec step
    vx0 = rVec x0

haftkaGurdalSimplex :: forall n. KnownNat n => ℝ -> R n -> Simplex n
haftkaGurdalSimplex size x0 = Simplex . colsL
  $ V.generate $ \j -> x0 L.& rVecIso L.%~ (SV.imap $ \i x -> x + step i j)
  where
    step i j
      | toInteger j == toInteger n' = 0
      | toInteger j == toInteger i = q + (c * l)
      | otherwise = q
    q = c * (sqrt (l + 1) - 1)
    c = size / (l * sqrt 2)
    l = fromIntegral n'
    n' = fromSing $ (sing :: Sing n)

centroid :: forall n t. (KnownNat n, Foldable t) => t (R n) -> R n
centroid = (/l) . FLD.foldl' (+) 0
  where
    l = fromIntegral . fromSing $ (sing :: Sing n)

centroidL :: forall m n. (KnownNat m, KnownNat n) => L m n -> R m
centroidL = (#> (1.0 / l))
  where
    l = fromIntegral . fromSing $ (sing :: Sing m)

realWithinRelTol :: ℝ -> ℝ -> ℝ -> Bool
realWithinRelTol relTol x y = (<= 0.5 * (x `absAdd` y)) . abs $ x - y
  where absAdd = (+) `on` abs

realWithinTols :: Maybe ℝ -> Maybe ℝ -> ℝ -> ℝ -> Bool
realWithinTols relTol absTol x y = relTolCheck && absTolCheck
  where
    relTolCheck = fromMaybe True
      $ (\t -> realWithinRelTol t x y) <$> relTol
    absTolCheck = fromMaybe True $ (<= abs (x - y)) <$> absTol

simplexRadiusTo :: (KnownNat n, Foldable t) => t (R n) -> R n -> R n
simplexRadiusTo xs c = vecR
  $ SV.generate
  $ \i -> FLD.foldl' (\acc -> (max `on` abs) acc . distci i . rVec) 0 xs
  where
    distci i v = SV.index v i - SV.index vc i
    vc = rVec c

--center + max step to given vertex for each dimension
simplexRadiusToL :: KnownNat n => Simplex n -> R n -> R n
simplexRadiusToL (Simplex xs) c = vecR $ SV.generate (V.index v)
  where
    v = (SV.maximumBy (comparing abs) . rVec) <$> lRows xsDiff -- Vector, need to convert to (R n)
    xsDiff = xs - (c `outer` 1) -- vector not promoted to matrix

affineCombWith :: KnownNat n => R n -> ℝ -> R n -> R n
affineCombWith x cy y = x + rnScale cy (y - x)

--put in helper class, maybe use coerce or something
rnScale :: KnownNat n => ℝ -> R n -> R n
rnScale c = L.over rVecIso $ SV.map (c *)

--TODO write version using mean
--toRows, map mean, list to vec
centroidL2 :: forall m n. Sing m -> Sing n -> L m n -> R m
centroidL2 (sm@SNat) (sn@SNat) = (#> (1.0 / l))
  where
    l = fromIntegral . fromSing $ (sing :: Sing m)


{-# language GADTs #-}
{-# language ConstraintKinds #-}
{-# language FlexibleContexts #-}

{-# language ScopedTypeVariables #-}
{-# LANGUAGE QuasiQuotes #-}
{-# language TypeApplications #-}

module FFT where

import Feldspar
import Feldspar.Software hiding (Arr)
import Feldspar.Software.Verify
import Feldspar.Software.Compile

import Feldspar.Array.Vector
import Feldspar.Array.Buffered

import Data.Bits (Bits)
import Data.Complex (Complex)

import Data.Selection
import Data.Default.Class

import Control.Monad

-- language-c-quote
import Language.C.Quote.GCC
import qualified Language.C.Syntax as C

-- imperative-edsl
import qualified Language.Embedded.Backend.C  as Imp

import Prelude hiding ((==), (/=), (>=), length)

--------------------------------------------------------------------------------
-- * FFT
--------------------------------------------------------------------------------

type SRef    a = Reference Software a
type SArr    a = Array  Software a
type SIArr   a = IArray Software a
type SStore  a = Store  Software a
type SSyntax a = Syntax SExp a

-- | Collection of constraints for immutable arrays.
type Immutable arr a = (
    Manifestable Software arr a
  , Finite   SExp arr
  , Indexed  SExp arr, ArrElem arr ~ a
  , Slicable SExp arr
  )

--------------------------------------------------------------------------------
-- ** Helper functions.

-- | Creates a value of type @a@ where the lowest @n@ bits are ones.
oneBits :: (Bits a, Num (SExp a), SType' a)
  => SExp Int32 -> SExp a
oneBits n = complement (ones .<<. n)

-- | Check if the @i@'th bit of @a@ is a one.
testBit :: (Integral a, Bits a, Num (SExp a), SType' a) =>
  SExp a -> SExp Index -> SExp Bool
testBit a i = i2b (a .&. (1 .<<. i2n i))

-- | Flip the @i@'th bit of @a@.
flipBit :: (Bits a, Num (SExp a), SType' a) =>
  SExp a -> SExp Index -> SExp a
flipBit a i = a `xor` (1 .<<. (i2n i))
--  SExp Index -> SExp a -> SExp a
--flipBit i a = a `xor` (1 .<<. (i2n i))

-- |
zeroBit :: (Bits a, Num (SExp a), SType' a) =>
  SExp Index -> SExp a -> SExp a
zeroBit i a = a + (a .&. (ones .<<. (i2n i)))

-- | Zeroes all but the lowest @i@ bits in @a@.
leastBits :: (Bits a, Num (SExp a), SType' a) =>
  SExp Int32 -> SExp a -> SExp a
leastBits i a = a .&. oneBits i

-- | Two to the power of @n@, i.e. @2^n@.
pow2 :: (Num (SExp a), Bits a, SType' a)
  => SExp Index -> SExp a
pow2  n = 1 .<<. i2n n
twoTo n = pow2 n

--------------------------------------------------------------------------------
-- ** Riffle network.

-- | Riffle indices.
rotBit :: SExp Index -> SExp Index -> SExp Index
rotBit k i = lefts .|. rights
  where
    k'     = i2n k
    ir     = i .>>. 1
    rights = ir .&. oneBits k'
    lefts  = (((ir .>>. k') .<<. 1) .|. (i .&. 1)) .<<. k'

-- | Riffle network for a pully array.
riffle :: (Immutable arr a, SSyntax a)
  => SExp Index -> arr -> SPull a
riffle k arr = Pull (length arr) $ \i -> arr ! rotBit k i

-- |
revBit :: (Immutable arr a, SSyntax a)
  => SStore a -> SExp Length -> arr -> Software (SManifest a)
revBit st n vec = loopStore st 1 1 (n-1) (step) vec
  where
    step i arr = return $ riffle i arr
    --step :: (Immutable arr a, SSyntax a) => SExp Index -> arr -> Software (SPush a)
--    step i arr = return $ unroll 2 $ riffle i arr
--    step k arr = return $ pairwise @Software (\i -> (i, rotBit k i)) $ riffle k arr

--------------------------------------------------------------------------------
-- ** Twiddle factors.

-- |
tw :: (Floating a, SType' a, SType' (Complex a))
  => SExp Index -> SExp Index -> SExp (Complex a)
tw n k = polar 1 (-2 * pi * i2n k / i2n n)

-- |
twids
  :: ( Immutable ts  (SExp (Complex a))
     , Immutable vec (SExp (Complex a))
     , SType' a
     , SType' (Complex a)
     , RealFloat a
     )
  => ts
  -> SExp Index
  -> SExp Index
  -> SExp Length
  -> vec
  -> SPull (SExp (Complex a))
twids ts n k l vec = Pull l $ \i ->
  let
    j = (leastBits (i2n k) i) .<<. (n'-1-k')
  in
    (testBit i k) ? ((ts!j) * (vec!i)) $ (vec!i)
  where
    n' = i2n n
    k' = i2n k

--------------------------------------------------------------------------------
-- ** Butterfly.

-- | Butterfly network.
bfly
  :: ( Immutable vec (SExp (Complex a))
     , RealFloat a
     , SType' a
     , SType' (Complex a)
     )
  => SExp Index -> vec -> SPull (SExp (Complex a))
bfly k as = Pull (length as) $ \i ->
    let a = as ! i
        b = as ! flipBit i k
    in  (testBit i k) ? (b-a) $ (a+b)

--------------------------------------------------------------------------------
-- ** FFT Core.

-- | Core of the FFT
fftCore
  :: ( Immutable ts  (SExp (Complex a))
     , Immutable vec (SExp (Complex a))
     , RealFloat a
     , SType' a
     , SType' (Complex a)
     )
  => SStore (SExp (Complex a))
  -> ts
  -> SExp Length
  -> vec
  -> Software (SManifest (SExp (Complex a)))
fftCore st ts n vec =
  let
    step k = return . twids ts n k (length vec) . bfly k
--    step k = return . pairwise @Software (\i -> (i, flipBit i k)) . twids ts n k (length vec) . bfly k
  in
    do arr <- loopStore st ((i2n n :: SExp Int32)-1) (-1) 0 (step . i2n) vec
--       revBit st n arr
       return arr

-- | Radix-2 Decimation-In-Frequency Fast Fourier Transformation of the given
-- complex vector. The given vector must be power-of-two sized, (for example 2,
-- 4, 8, 16, 32, etc.) The output is non-normalized.
fft
  :: ( Immutable vec (SExp (Complex a))
     , RealFloat a
     , SType' a
     , SType' (Complex a)
     )
  => SExp Length
  -> SStore (SExp (Complex a))
  -> vec
  -> Software (SManifest (SExp (Complex a)))
fft n st vec =
  do n  <- shareM (ilog2 (length vec))
     ts <- manifestFresh $ Pull (pow2 (n-1)) (tw (pow2 n))
     fftCore st ts n vec

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

fftCore2 ::
  SStore (SExp (Complex Double)) ->
  SIArr (SExp (Complex Double)) ->
  SExp Length ->
  SPull (SExp (Complex Double)) ->
  Software (SIArr (SExp (Complex Double)))
fftCore2 store iarr exp pull = do
  let iarrm :: SManifest (SExp (Complex Double))
      iarrm = M iarr
  (M ibrr) <- fftCore store iarrm exp pull
  return ibrr

fft2 ::
  SExp Length ->
  SStore (SExp (Complex Double)) ->
  SPull (SExp (Complex Double)) ->
  Software (SIArr (SExp (Complex Double)))
fft2 n store pull = do
  (M ts) <- manifestFresh $ Pull (twoTo (n-1)) (tw (twoTo n))
  fftCore2 store ts n pull

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

example :: Software ()
example = do
  size :: SExp Length <- fget stdin
  n :: SExp Length <- fget stdin
  assert (1 `shiftL` n == size) "2^n == size"
  assert (size >= 2) "not too small"
  st   :: SStore (SExp (Complex Double)) <- newInPlaceStore size
  arr  :: SArr   (SExp (Complex Double)) <- newArr size
  iarr :: SIArr  (SExp (Complex Double)) <- freezeArr arr
  fft n st iarr
  assert (value False) "oh no"
  return ()

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

bad :: Software ()
bad = do
  size :: SExp Length <- fget stdin
  fprintf stdout "%d" (share size id)
  return ()

--------------------------------------------------------------------------------
-- FFT Bench. Copy of https://github.com/Feldspar/raw-feldspar/blob/master/examples/FFT_bench.hs
--------------------------------------------------------------------------------

printTime_def = [cedecl|
void printTime(typename clock_t start, typename clock_t end)
{
  printf("CPU time (sec): %f\n", (double)(end-start) / CLOCKS_PER_SEC);
}
|]

sizeOf_double_complex :: SExp Length
sizeOf_double_complex = 16

-- | Measure the time for 100 runs of 'fftCore' (excluding initialization) for
-- arrays of the given size
benchmark :: SExp Length -> Software ()
benchmark n = do
  addInclude "<stdio.h>"
  addInclude "<string.h>"
  addInclude "<time.h>"

  addDefinition printTime_def

  start <- newObject "clock_t" False
  end   <- newObject "clock_t" False

  st :: SStore (SExp (Complex Double)) <- newStore n
  inp <- unsafeFreezeStore n st
  callProc "memset"
      [ iarrArg (manifest inp)
      , valArg (0 :: SExp Index)
      , valArg (n*sizeOf_double_complex)
      ]

  n  <- shareM (ilog2 (length inp))
  ts <- manifestFresh $ Pull (twoTo (n-1)) (tw (twoTo n))
  -- Change `manifestFresh` to `return` to avoid pre-computing twiddle factors

  callProcAssign start "clock" []

  for 0 1 99 $ \(_ :: SExp Index) ->
    void $ fftCore st ts n inp

  callProcAssign end "clock" []
  callProc "printTime" [objArg start, objArg end]

runBenchmark n = runCompiled'
    (def :: CompilerOpts) --{compilerAssertions = select []}
    -- Note: important to turn off assertions when running the benchmarks
    --       (in old Raw-Feldspar that is, Co-Feldspar gets rid of them)
    def {Imp.externalFlagsPre = ["-O3"], Imp.externalFlagsPost = ["-lm"]}
    (benchmark n)

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

dummy :: Software ()
dummy = do
  size :: SExp Length <- fget stdin
  n :: SExp Length <- fget stdin
  assert (1 `shiftL` n == size) "2^n == size"
  assert (size >= 2) "not too small"
  st   :: SStore (SExp (Complex Double)) <- newInPlaceStore size
  arr  :: SArr   (SExp (Complex Double)) <- newArr size
  iarr :: SIArr  (SExp (Complex Double)) <- freezeArr arr
  fft n st iarr
  fft n st iarr
  assert (value False) "oh no"
  return ()

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


-----------------------------------------------------------------------------
-- |
-- 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


{-# 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))


{-# 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.

-}