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

{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NoStarIsType          #-}
{-# LANGUAGE OverloadedLabels      #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}

import           Control.Monad
import           System.Random.MWC

import           Data.Constraint              (Dict (..))
import           Data.List                    (zipWith5)
import           Data.Proxy
import           Data.Reflection
import           Data.Singletons
import           Data.Singletons.TypeLits     hiding (natVal)
import           GHC.TypeLits
import           Unsafe.Coerce                (unsafeCoerce)

import           Numeric.LinearAlgebra.Static (L, R)
import qualified Numeric.LinearAlgebra.Static as H

import           Criterion.Main

import           Grenade
import           Grenade.Utils.LinearAlgebra
import           Grenade.Utils.ListStore

main :: IO ()
main = do
  defaultMain
    [ bgroup
        "batchnorm forward (CxHxW)"
        [ benchBatchNorm "1x1x64" 1 1 64
        , benchBatchNorm "1x64x64" 1 64 64
        , benchBatchNorm "64x64x64" 64 64 64
        ]
    , bgroup
        "convolutions with bias (CxHxW)"
        [ benchBiasConvolution "1024x64x64, 125 kernels" 1024 64 64 125 1 1 64 64
        , benchBiasConvolution "3x416x416, 64 kernels" 3 416 416 64 2 1 208 208
        ]
    , bgroup
        "leaky relu (CxHxW)"
        [ benchLeakyRelu "3x416x416" 3 416 416
        , benchLeakyRelu "1024x16x16" 1024 16 16
        ]
    ]

-- BENCHMARK GENERATION FUNCTIONS

benchBiasConvolution :: String -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Benchmark
benchBiasConvolution name channels width height filters strides kernels outWidth outHeight
  | strides * (outHeight - 1) <= (height - kernels + 1) - 1 && (height - kernels + 1) <= (outHeight * strides) && strides * (outWidth - 1) <= (width - kernels + 1) - 1 && (width - kernels + 1) <= (outWidth * strides)
      = case (someNatVal (fromIntegral channels), someNatVal (fromIntegral width), someNatVal (fromIntegral height), someNatVal (fromIntegral filters), someNatVal (fromIntegral strides), someNatVal (fromIntegral kernels), someNatVal (fromIntegral outWidth), someNatVal (fromIntegral outHeight)) of
          (Just (SomeNat (Proxy :: Proxy channels)), Just (SomeNat (Proxy :: Proxy width)), Just (SomeNat (Proxy :: Proxy height)), Just (SomeNat (Proxy :: Proxy filters)), Just (SomeNat (Proxy :: Proxy strides)), Just (SomeNat (Proxy :: Proxy kernels)), Just (SomeNat (Proxy :: Proxy outWidth)), Just (SomeNat (Proxy :: Proxy outHeight))) ->
            case (channels, width, height
                 , (unsafeCoerce (Dict :: Dict ()) :: Dict ( strides * (outHeight - 1) <= (height - kernels) ) )
                 , (unsafeCoerce (Dict :: Dict ()) :: Dict ( (height - kernels) <= (outHeight * strides) - 1 ) )
                 , (unsafeCoerce (Dict :: Dict ()) :: Dict ( strides * (outWidth - 1) <= (width - kernels ) ) )
                 , (unsafeCoerce (Dict :: Dict ()) :: Dict ( (width - kernels) <= (outWidth * strides) - 1 ) ) ) of
              (1, 1, _, _, _, _, _)
                -> error "1D convolutions are not allowed"
              (1, _, _, Dict, Dict, Dict, Dict)
                -> env (generateBiasConvEnv :: IO (Convolution 'WithBias 'NoPadding 1 filters kernels kernels strides strides, S ('D2 height width))) $ \ ~(layer, x) -> bench name $ nf (snd . runForwards layer :: S ('D2 height width) -> S ('D3 outHeight outWidth filters )) x
              (_, _, _, Dict, Dict, Dict, Dict)
                -> env (generateBiasConvEnv :: IO (Convolution 'WithBias 'NoPadding channels filters kernels kernels strides strides, S ('D3 height width channels))) $ \ ~(layer, x) -> bench name $ nf (snd . runForwards layer :: S ('D3 height width channels) -> S ('D3 outHeight outWidth filters)) x
  where
    generateBiasConvEnv :: forall channels filters kernel1 kernel2 strides1 strides2 s.
                            ( SingI s, KnownNat channels, KnownNat filters, KnownNat kernel1, KnownNat kernel2,
                              KnownNat strides1, KnownNat strides2 )
                            => IO (Convolution 'WithBias 'NoPadding channels filters kernel1 kernel2 strides1 strides2, S s)
    generateBiasConvEnv = do
      x     <- randomOfShape
      gen   <- createSystemRandom
      layer <- createRandomWith UniformInit gen
      return (layer, x)

benchBatchNorm :: String -> Int -> Int -> Int -> Benchmark
benchBatchNorm name channels width height
  = case (someNatVal (fromIntegral channels), someNatVal (fromIntegral width), someNatVal (fromIntegral height)) of
          (Just (SomeNat (Proxy :: Proxy channels)), Just (SomeNat (Proxy :: Proxy width)), Just (SomeNat (Proxy :: Proxy height))) ->
            case (channels, width, height) of
              (1, 1, _) -> env (generateBatchNormEnv False :: IO (BatchNorm 1 1 width 90, S ('D1 width))) $ \ ~(layer, x) -> bench name $ nf (snd . runForwards layer :: S ('D1 width) -> S ('D1 width)) x
              (1, _, _) -> env (generateBatchNormEnv False :: IO (BatchNorm 1 width height 90, S ('D2 width height))) $ \ ~(layer, x) -> bench name $ nf (snd . runForwards layer :: S ('D2 width height) -> S ('D2 width height)) x
              (_, _, _) -> env (generateBatchNormEnv False :: IO (BatchNorm channels width height 90, S ('D3 width height channels))) $ \ ~(layer, x) -> bench name $ nf (snd . runForwards layer :: S ('D3 width height channels) -> S ('D3 width height channels)) x
  where
    generateBatchNormEnv :: forall c h w s. (KnownNat c, KnownNat h, KnownNat w, SingI s)
                         => Bool -> IO (BatchNorm c h w 90, S s)
    generateBatchNormEnv training = do
      x     <- randomOfShape
      gens  <- replicateM 4 createSystemRandom
      seeds <- mapM uniform gens :: IO [Int]
      let [gamma, beta, running_mean, running_var] = map (\s -> H.randomVector s H.Uniform) seeds :: [R c]
          ε            = 0.00001
      return (BatchNorm training (BatchNormParams gamma beta) running_mean running_var ε mkListStore, x)

benchLeakyRelu :: String -> Int -> Int -> Int -> Benchmark
benchLeakyRelu name channels width height
  = case (someNatVal (fromIntegral channels), someNatVal (fromIntegral width), someNatVal (fromIntegral height)) of
          (Just (SomeNat (Proxy :: Proxy channels)), Just (SomeNat (Proxy :: Proxy width)), Just (SomeNat (Proxy :: Proxy height))) ->
            case (channels, width, height) of
              (1, 1, _) -> env (generateLeakyReluEnv :: IO (LeakyRelu, S ('D1 width))) $ \ ~(layer, x) -> bench name $ nf (snd . runForwards layer :: S ('D1 width) -> S ('D1 width)) x
              (1, _, _) -> env (generateLeakyReluEnv :: IO (LeakyRelu, S ('D2 width height))) $ \ ~(layer, x) -> bench name $ nf (snd . runForwards layer :: S ('D2 width height) -> S ('D2 width height)) x
              (_, _, _) -> env (generateLeakyReluEnv :: IO (LeakyRelu, S ('D3 width height channels))) $ \ ~(layer, x) -> bench name $ nf (snd . runForwards layer :: S ('D3 width height channels) -> S ('D3 width height channels)) x
  where
    generateLeakyReluEnv :: forall s. (SingI s) => IO (LeakyRelu, S s)
    generateLeakyReluEnv = do
      x     <- randomOfShape
      gen   <- createSystemRandom
      alpha <- uniformR (-1, 1) gen :: IO RealNum
      return (LeakyRelu alpha, x)


{-|
Type aliases for activation functions and their derivative.
-}
module ML.NN.ActivationFunction (ActivationFunction, ActivationFunctionDerivative) where

import Numeric.LinearAlgebra (R)

-- | An activation function for a neuron.
type ActivationFunction = R -> R

-- | The derivative of a neuron activation function.
type ActivationFunctionDerivative = R -> R


{-# LANGUAGE TypeApplications #-}
module Main where

import Config
import File

import Statistics.Distribution
import Statistics.Distribution.Empirical
import Statistics.Distribution.Normal

import Data.List
import Data.Yaml

import Text.Printf

import Debug.Trace

main :: IO ()
main = do
  -- YAML setup
  yml <- decodeFileEither "splits.yaml"

  f <- case yml of
         Right (Config _ levels) -> load levels
         Left  _ -> load'

  let l = 1
      (Just distrs) = lookup l $ levelEmpiricalDistribution f
      (Just level) = fmap (sort . map realToFrac) . lookup l $ onlyValidSplits $ levelData f
      elems = [0, 0.01..0.99]
      dens = 1 / (fromIntegral $ length level)
      domain = reverse [1.0, 1.0 - dens..dens]

      vec = zipWith Vector2 level domain

      vecs = zip4 vec (drop 1 vec) (drop 2 vec) (drop 3 vec)

      vals = concat $ [ [ let (Vector2 x y) = centripetalInterpolate v m in (x, y) | m <- elems ] | v <- vecs ]

      (xs, ys) = unzip vals

      test = zipWith (\(x1, y1) (x2, y2) -> (abs $ 1 / (y2 - y1)) * (x2 - x1)) vals $ tail vals
      (dx, dy) = unzip $ zipWith (\(x1, y1) (x2, y2) -> (,) (x2 - x1) (y2 - y1)) vals $ tail vals

  mapM_ putStrLn $ zipWith5 (printf "%f %f %f %f %f") xs ys dx dy test

linearInterpolate :: (Double, Double) -> Double -> Double
linearInterpolate (lRange, uRange) m =
  let notM = 1 - m
  in lRange * notM + uRange * m

hermiteInterpolate :: (Double, Double, Double, Double) -> Double -> Double
hermiteInterpolate (y0, y1, y2, y3) mu =
  let mu2 = mu * mu
      a0 = (-0.5) * y0 + 1.5 * y1 - 1.5 * y2 + 0.5 * y3
      a1 = y0 - 2.5 * y1 + 2 * y2 - 0.5 * y3
      a2 = (-0.5) * y0 + 0.5 * y2
      a3 = y1
  in a0 * mu * mu2 + a1 * mu2 + a2 * mu + a3

centripetalInterpolate (p0, p1, p2, p3) mu =
  let t0 = 0
      t1 = getT t0 p0 p1
      t2 = getT t1 p1 p2
      t3 = getT t2 p2 p3

      t = linearInterpolate (t1, t2) mu

      a1 = vscale ((t1 - t) / (t1 - t0)) p0 + vscale ((t - t0) / (t1 - t0)) p1
      a2 = vscale ((t2 - t) / (t2 - t1)) p1 + vscale ((t - t1) / (t2 - t1)) p2
      a3 = vscale ((t3 - t) / (t3 - t2)) p2 + vscale ((t - t2) / (t3 - t2)) p3

      b1 = vscale ((t2 - t) / (t2 - t0)) a1 + vscale ((t - t0) / (t2 - t0)) a2
      b2 = vscale ((t3 - t) / (t3 - t1)) a2 + vscale ((t - t1) / (t3 - t1)) a3
  in vscale ((t2 - t) / (t2 - t1)) b1 + vscale ((t - t1) / (t2 - t1)) b2
  where getT tn (Vector2 x1 y1) (Vector2 x2 y2) = tn + ((sqrt $ (x2 - x1) ^ 2 + (y2 - y1) ^ 2) ** alpha)
        alpha = 0.5

type Scalar = Double

class Vector v where
  vmap  :: (Scalar -> Scalar) -> v -> v
  vzip  :: (Scalar -> Scalar -> Scalar) -> v -> v -> v
  vfold :: (x -> Scalar -> x) -> x -> v -> x

vdot :: Vector v => v -> v -> Scalar
vdot v0 v1 = vfold (+) 0 $ vzip (*) v0 v1

vmag_sqr :: Vector v => v -> Scalar
vmag_sqr v = v `vdot` v

vmag :: Vector v => v -> Scalar
vmag = sqrt . vmag_sqr

vscale :: Vector v => Scalar -> v -> v
vscale s = vmap (s*)

vunit :: Vector v => v -> v
vunit v =
  if vmag v == 0
    then v
    else vscale (1 / vmag v) v

data Vector2 = Vector2 {v2x, v2y :: Scalar} deriving (Eq)

instance Show Vector2 where
  show (Vector2 x y) = "<" ++ (show x) ++ ", " ++ (show y) ++ ">"

instance Vector Vector2 where
  vmap  f   (Vector2 x y) = Vector2 (f x) (f y)
  vfold f i (Vector2 x y) = (i `f` x) `f` y
  vzip  f   (Vector2 x0 y0) (Vector2 x1 y1) = Vector2 (f x0 x1) (f y0 y1)

instance Num Vector2 where
  (+) = vzip (+)
  (-) = vzip (-)
  (*) = vzip (*)
  negate = vmap negate
  fromInteger s = Vector2 (fromInteger s) (fromInteger s)

instance Fractional Vector2 where
  (/) = vzip (/)
  fromRational s = let r = realToFrac s in Vector2 r r


{-# LANGUAGE BangPatterns #-}

import Gauge.Main (defaultMain, bench, whnf)

import Data.Int (Int64)
import Data.List (unfoldr)
import Data.Maybe (fromJust)
import Data.Primitive.Contiguous (PrimArray, fromListN)
import Statistics.Array.Types (AscList)
import System.Random (random, mkStdGen)

import qualified Data.Primitive as PM
import qualified Statistics.Array as Stats
import qualified Statistics.Array.Types as Asc

main :: IO ()
main = defaultMain
  [ bench "sort-16" $ whnf fromAscInt input16
  , bench "sort-128" $ whnf fromAscInt input128
  , bench "sort-1024" $ whnf fromAscInt input1024
  , bench "sort-65536" $ whnf fromAscInt input65536
  , bench "list-derivative" $ whnf performListDerivativeInt asc1024
  , bench "median-of-absolute-deviations" $ whnf Stats.mad asc1024
  ]
  where
  input16, input128, input1024, input65536 :: PrimArray Int64
  !input16 = fromListN 16 $ take 16 . drop 0 $ inputInf
  !input128 = fromListN 128 $ take 128 . drop 16 $ inputInf
  !input1024 = fromListN 1024 $ take 1024 . drop (16+128) $ inputInf
  !input65536 = fromListN 65536 $ take 65536 . drop (16+128+1024) $ inputInf
  asc1024 :: AscList Int64
  asc1024 = fromJust $ Asc.fromArray input1024

fromAscInt :: PrimArray Int64 -> Int64
{-# noinline fromAscInt #-}
fromAscInt !x = case Asc.fromArray x of
  Nothing -> 0
  Just (Asc.AscList y) -> fromIntegral (PM.sizeofPrimArray y)

-- This is here to make it easy to inspect Core to confirm that specialization
-- works correctly.
performListDerivativeInt :: AscList Int64 -> PrimArray Int64
{-# noinline performListDerivativeInt #-}
performListDerivativeInt !x = Stats.listDerivative x

inputInf :: [Int64]
inputInf = unfoldr (Just . random) seed
  where
  seed = mkStdGen 1234567


module GApprox
  (
    GApprox (..)
  , createGApprox
  , predictG
  , updateG
  ) where

-- Wrapper for HNN's Neural network module
-- Using it for function approximation while in simulation -> wrapper supports creation,
-- prediction, and update.
-- Usecase: Predict 1 -> update 1 -> predict 1 -> update 1

-- Goal -> understand library enough to support a CCEA approach (ie, be able to mutate the
-- weights). It seems like there is a load from weight matrix function, so should involve just
-- creating a Network -> Weight matrix function.

import AI.HNN.FF.Network
import Numeric.LinearAlgebra
import Policy

data GApprox = GApprox (IO (Network Double))
instance Eq GApprox where
  (==) _ _ = True

predictG :: GApprox     -- Neural Network Approximating G
         -> [Double]    -- Input (will be state action pair)
         -> IO [Double] -- predicted value of G
predictG (GApprox network) input = do
  net <- network
  return $ toList $ output net sigmoid $ fromList input

updateG :: GApprox -- Neural Network Approximating G
        -> Double   -- Learning Rate
        -> [Double] -- Input to G
        -> [Double] -- Actual value of G
        -> GApprox -- Updated Neural Network
updateG (GApprox network) rate input target = GApprox newNet
  where newNet = do
          net <- network
          let sample = [ fromList input --> fromList target ] :: Samples Double
          return $ trainNTimes 1 rate sigmoid sigmoid' net sample

createGApprox :: GApprox
createGApprox = GApprox $ createNetwork 8 [14] 1


{- |
Module          : $Header$
Description     : This is the parser file for the Scheme compiler.
Copyright       : (c) Michael Buchel
License         : BSD3

Maintainer      : abuchel2@uwo.ca
Stability       : experimental
Portability     : portable

This is the parser, where I am using Parsec for simplicity.
| -}
module Parser where

import Abstract
import Control.Monad
import Control.Monad.Except
import Data.Array
import Data.Complex
import Data.Ratio
import Error
import Evaluator
import Numbers
import Text.ParserCombinators.Parsec

-- | Symbol function which identifies all the symbols allowed in Scheme.
symbol :: Parser Char
symbol = oneOf "!$%&|*+-/:<=>?@^_~"

-- | Escaped character parser.
escapedChars :: Parser Char
escapedChars = do
	char '\\'
	x <- oneOf "\\\"nrt"
	return $ case x of
		      '\\' -> x
		      '"' -> x
		      'n' -> '\n'
		      'r' -> '\r'
		      't' -> '\t'

-- | Parses a string.
parseString :: Parser LispVal
parseString = do
	char '"'
	x <- many $ escapedChars <|> noneOf "\"\\"
	char '"'
	return $ String x

-- | Parses an atom.
parseAtom :: Parser LispVal
parseAtom = do
	first <- letter <|> symbol
	rest <- many (letter <|> digit <|> symbol)
	let atom = first:rest
	return $ Atom atom

-- | Parses a boolean value.
parseBool :: Parser LispVal
parseBool = do
	char '#'
	(char 't' >> return (Bool True)) <|> (char 'f' >> return (Bool False))

-- | Parses decimals without the identifier string.
parseDecimal1 :: Parser LispVal
parseDecimal1 = many1 digit >>= (return . Number . read)

-- | Parses decimals with the identifier string.
parseDecimal2 :: Parser LispVal
parseDecimal2 = do
	try $ string "#d"
	many1 digit >>= (return . Number . read)

-- | Parses hexadecimals.
parseHex :: Parser LispVal
parseHex = do
	try $ string "#x"
	many1 hexDigit >>= (return . Number . hex2dig)

-- | Parses octal numbers.
parseOct :: Parser LispVal
parseOct = do
	try $ string "#o"
	many1 octDigit >>= (return . Number . oct2dig)

-- | Parses binary numbers.
parseBin :: Parser LispVal
parseBin = do
	try $ string "#b"
	many1 (oneOf "10") >>= (return . Number . bin2dig)

-- | Parses a number.
parseNumber :: Parser LispVal
parseNumber = parseDecimal1 <|> parseDecimal2 <|> parseHex <|> parseOct <|> parseBin

-- | Parses a floating point number.
parseFloat :: Parser LispVal
parseFloat = do
	x <- many1 digit
	char '.'
	y <- many1 digit
	return $ Float $ getFloat x y

-- | Parses a ratio value.
parseRatio :: Parser LispVal
parseRatio = do
	x <- many1 digit
	spaces >> char '/' >> spaces
	y <- many1 digit
	return $ Ratio ((read x) % (read y))

-- | Parses a complex number.
parseComplex :: Parser LispVal
parseComplex = do
	x <- (try parseFloat <|> parseDecimal1)
	spaces >> char '+' >> spaces
	y <- (try parseFloat <|> parseDecimal1)
	oneOf "ij"
	return $ Complex (toDouble x :+ toDouble y)

-- | Helper list parser.
parseListHelp :: Parser LispVal
parseListHelp = do
	char '.' >> (skipMany1 space)
	x <- parseExpr
	spaces >> char ')'
	return x

-- | Parses an individual list.
parseList :: Parser LispVal
parseList = do
	char '(' >> spaces
	headList <- parseExpr `sepEndBy` (skipMany1 space)
	(parseListHelp >>= return . (DottedList headList)) <|> (spaces >> char ')' >> return (List headList))

-- | For Scheme syntax sugar.
parseQuoted :: Parser LispVal
parseQuoted = do
	char '\''
	x <- parseExpr
	return $ List [Atom "quote", x]

-- | For Scheme syntax sugar.
parseQuasiQuoted :: Parser LispVal
parseQuasiQuoted = do
	char '`'
	x <- parseExpr
	return $ List [Atom "quasiquote", x]

-- | For Scheme syntax sugar.
parseUnQuoted :: Parser LispVal
parseUnQuoted = do
	char ','
	x <- parseExpr
	return $ List [Atom "unquote", x]

-- | Parses a character value.
parseChar :: Parser LispVal
parseChar = do
	try $ string "#\\"
	value <- try (string "newline" <|> string "space") <|> do {x <- anyChar; notFollowedBy alphaNum; return [x]}
	return $ Character $ case value of
				  "space" -> ' '
				  "newline" -> '\n'
				  otherwise -> (value !! 0)

-- | Parses a vector.
parseVector :: Parser LispVal
parseVector = do
	arrayVals <- sepBy parseExpr spaces
	return $ Vector (listArray (0, (length arrayVals - 1)) arrayVals)

-- | Helper function for parsing vectors.
parseVectors :: Parser LispVal
parseVectors = do
	string "#(" >> spaces
	x <- parseVector
	spaces >> char ')'
	return x

-- | Helper parsing function which tries to parse the expression given to it.
parseExpr :: Parser LispVal
parseExpr = parseAtom
	<|> parseString
	<|> parseQuoted
	<|> parseQuasiQuoted
	<|> parseUnQuoted
	<|> try parseComplex
	<|> try parseFloat
	<|> try parseRatio
	<|> try parseNumber
	<|> try parseBool
	<|> try parseChar
	<|> try parseVectors
	<|> parseList

-- | A function to read the expression from a string.
readExpr :: String -- ^ Input string
	-> ThrowsError LispVal
readExpr input = case parse (spaces >> parseExpr) "parser" input of
		      Left err -> throwError $ Parser err
		      Right val -> return val

-- | Find the end of an expression.
findExpr :: String -- ^ Input string.
	-> Integer -- ^ Previous integer.
	-> Integer -- ^ Nesting of expression.
	-> Integer
findExpr [] x _ = x
findExpr (x : xs) n y = case x of
			     '(' -> findExpr xs (n + 1) (y + 1)
			     ')' -> if (y - 1) == 0 then (n + 1) else findExpr xs (n + 1) (y - 1)
			     otherwise -> findExpr xs (n + 1) y


{-# LANGUAGE OverloadedStrings #-}
module Lisp.EvalSpec (spec) where

import           Test.Hspec
import           Test.QuickCheck

import           Data.Complex
import           Data.Ratio
import           Data.Text

import           Lisp
import           Lisp.LispVal    (h2l)

spec :: Spec

spec = do
  describe "Basic primitives self-eval" $ do
    it "Ints self-eval" $
      property $ \i -> let l = Int i in eval l == l

    it "Reals self-eval" $
      property $ \r -> let l = Real r in eval l == l

    it "Rationals self-eval" $
      property $ \q -> let l = Rational q in eval l == l

    it "Complex numbers self-eval" $
      property $ \c -> let l = Complex c in eval l == l

    it "Strings self-eval" $
      property $ \s -> let l = String (pack s) in eval l == l

    it "Bools self-eval" $
      property $ \b -> let l = Bool b in eval l == l

  describe "quoted forms return their contents" $ do
    let quoted v = Pair (Symbol "quote") (Pair v Nil)
    let foo = Symbol "foo"
    let bar = Symbol "bar"
    let foobar = Pair foo bar
    it "'foo evaluates to foo" $
      eval (quoted foo) `shouldBe` foo
    it "'(foo . bar) evaluates to (foo . bar)" $
      eval (quoted foobar) `shouldBe` foobar

  describe "basic primitives" $ do
    it "Can add two Ints" $
      eval (h2l [Symbol "+", Int 5, Int 5]) `shouldBe` Int 10
    it "Can add two Reals" $
      eval (h2l [Symbol "+", Real 5.0, Real 5.0]) `shouldBe` Real 10.0
    it "Can add Ints and Reals" $
      eval (h2l [Symbol "+", Int 5, Real 5.0]) `shouldBe` Real 10.0
    it "Can add Rationals and Complex" $
      eval (h2l [Symbol "+", Rational (5%2), Complex (2.4:+3.1)]) `shouldBe` Complex (4.9:+3.1)


-- https://drive.google.com/open?id=0BxJP-hCBgo5OcU9td2Fhc2xHek0

module Numeric.Combinations
  ( combIndex
  , combIndexSorted
  )
  where

import Control.Monad.ST (runST)

import Data.Vector.Unboxed ((!))
import qualified Data.Vector.Unboxed as UV
import Data.Vector.Algorithms.Merge (sort)

import Numeric.SpecFunctions (choose)

import Debug.Trace

-- | Calculates the combination index. Assumes that neither alpha nor x
--   contain any duplicate elements.
combIndex :: (UV.Unbox a, Ord a)
          => UV.Vector a -- Alphabet of length N
          -> UV.Vector a -- element of length K
          -> Maybe Int -- Index in C_K^A
combIndex alpha x | UV.length x > UV.length alpha = Nothing -- Require that the alphabet is bigger than the element!
combIndex alpha x = combIndexSorted alpha_sorted x_sorted where
  alpha_sorted = runST $ do
    m_alpha <- UV.thaw alpha
    sort m_alpha
    UV.unsafeFreeze m_alpha
  x_sorted = runST $ do
    m_x <- UV.thaw x
    sort m_x
    UV.unsafeFreeze m_x

-- | Calculates the combination index assuming both alpha and x are pre-sorted.
--   This will give incorrect answers if the inputs are not sorted!
combIndexSorted :: (UV.Unbox a, Ord a)
                => UV.Vector a -- Alphabet of length N
                -> UV.Vector a -- element of length K
                -> Maybe Int -- Index in C_K^A
combIndexSorted alpha x | UV.length x > UV.length alpha = Nothing
combIndexSorted alpha x = do
    x_idx <- m_x_idx
    return $ UV.sum
           . UV.map (flip inner_sum x_idx)
           $ UV.enumFromN 0 cap_k
  where
    m_x_idx = UV.mapM (flip UV.elemIndex alpha) x
    cap_k = UV.length x
    cap_n = UV.length alpha
    i_k (-1) _ = -1
    i_k k x_idx = x_idx ! k
    inner_sum k x_idx =
          UV.sum
        . UV.map (\q -> floor
                      $ choose (cap_n - (i1 + 1) - q) (cap_k - k - 1))
        $ UV.enumFromN 1 (i0 - (i1 + 1))
      where
        i0 = i_k k x_idx
        i1 = i_k (k-1) x_idx


{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE EmptyCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE LambdaCase #-}

module Models.Integrals.Approximate2  where

import Algebra.Classes
import qualified Algebra.Morphism.Affine as A
import Prelude hiding (Num(..), Fractional(..), (^), product, sum, pi, sqrt
                      , exp)
import qualified Models.Field
import qualified Algebra.Linear.Chebyshev as Chebyshev
import Data.Complex
import qualified Algebra.Morphism.Polynomial.Multi as Multi
import Algebra.Morphism.Polynomial.Multi hiding (constPoly)
import qualified Algebra.Morphism.LinComb as LC
import Control.Applicative

import Models.Integrals.Types

type family Env γ where
  Env 'Unit = ()
  Env (a ':× b) = (Env a, Env b)
  Env 'R = Double

type RR = Double

newtype PointWise x a = PW (x -> a) deriving Functor
instance Applicative (PointWise x) where
  pure x = PW $  \_ -> x
  PW f <*> PW a = PW (\x -> (f x) (a x))

fromPointWise :: PointWise x a -> x -> a
fromPointWise (PW x) = x

instance Additive a => Additive (PointWise x a) where
  zero = pure zero
  (+) = liftA2 (+)

instance Group a => Group (PointWise x a) where
  negate = fmap negate
  (-) = liftA2 (-)

instance AbelianAdditive a => AbelianAdditive (PointWise x a)

instance Multiplicative a => Multiplicative (PointWise x a) where
  one  = pure one
  (*) = liftA2 (*)

instance Division a => Division (PointWise x a) where
  recip  = fmap recip
  (/) = liftA2 (/)

instance Roots a => Roots (PointWise x a) where
  root n  = fmap (root n)

instance Transcendental a => Transcendental (PointWise x a) where
  pi = pure pi
  -- log = fmap log
  -- sin = fmap sin
  -- cos = fmap cos
  -- asin = fmap asin
  -- acos = fmap acos
  -- atan = fmap atan
  -- sinh = fmap sinh
  -- cosh = fmap cosh
  -- asinh = fmap asinh
  -- acosh = fmap acosh
  -- atanh = fmap atanh
  exp = fmap exp

instance Multiplicative a => Scalable (PointWise x a) (PointWise x a) where
  (*^) = (*)

instance Ring a => Ring (PointWise x a) where
  fromInteger = PW . const . fromInteger

instance Field a => Field (PointWise x a) where
  fromRational = PW . const . fromRational

type S γ = PointWise (Env γ) RR

lk :: Var γ -> Env γ -> RR
lk Get (_,x) = x
lk (Weaken x) (ρ,_) = lk x ρ

lk' :: Var γ -> S γ
lk' v = PW $ lk v

evalC :: forall γ. Coef γ -> S γ
evalC (Coef c) = LC.eval (\x -> PW @(Env γ) $ \_ρ -> Models.Field.eval @RR x)
                         (exp . evalP) c

evalP :: forall γ. Poly γ -> S γ
evalP = Multi.eval evalC evalE

evalE :: forall γ. Elem γ -> S γ
evalE = \case
  Supremum dir es -> case dir of
    Min ->  foldr (liftA2 min) (pure ( 1/0)) (evalP <$> es)
    Max ->  foldr (liftA2 max) (pure (-1/0)) (evalP <$> es)
  Vari x -> lk' x
  CharFun x -> fmap (\y -> if y >= 0 then 1 else 0) (evalP x)

approxIntegrals :: Int -> Env γ -> P γ -> RR
approxIntegrals n ρ =
  let evP x = fromPointWise (evalP x) ρ
  in \case
      Add a b -> approxIntegrals n ρ a + approxIntegrals n ρ b
      Div a b -> approxIntegrals n ρ a / approxIntegrals n ρ b
      Integrate (mkSuprema -> (lo,hi)) e ->
        realPart (Chebyshev.integral @Double @C n (evP lo) (evP hi) $
                  \x -> approxIntegrals n (ρ,x) e :+ 0)
      Ret x -> evP x
      Cond (IsNegative c) e -> if A.eval Models.Field.eval (flip lk ρ) c <= 0 then approxIntegrals n ρ e else 0
      Cond (IsZero _) _ -> error "approxIntegrals: equality not eliminated?"


{-# LANGUAGE AllowAmbiguousTypes        #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE InstanceSigs               #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE TypeApplications           #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE TypeFamilyDependencies     #-}
{-# LANGUAGE TypeInType                 #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE UndecidableInstances       #-}
{-# LANGUAGE ViewPatterns               #-}

module Numeric.BLAS.HMatrix (
    HM(..)
  , HM'
  , flatten
  , unflatten
  ) where

import           Control.Applicative
import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.Trans.State
import           Data.Finite.Internal
import           Data.Foldable
import           Data.Function
import           Data.Kind
import           Data.Maybe
import           Data.MonoTraversable
import           Data.Monoid                  (Endo(..))
import           Data.Singletons
import           Data.Singletons.Prelude
import           Data.Singletons.TypeLits
import           Data.Type.Product
import           GHC.Generics                 (Generic)
import           GHC.TypeLits
import           Numeric.BLAS hiding          (outer)
import           Numeric.LinearAlgebra.Static
import           Type.Family.List hiding      (Reverse)
import qualified Data.Vector                  as UV
import qualified Data.Vector.Generic.Sized    as VG
import qualified Data.Vector.Sized            as V
import qualified Data.Vector.Storable         as UVS
import qualified Data.Vector.Storable.Sized   as VS
import qualified Numeric.LinearAlgebra        as LA

type family HM' (s :: [Nat]) = (h :: Type) | h -> s where
    HM' '[]            = Double
    HM' '[n]           = R n
    HM' '[m,n]         = L m n
    HM' (n ': m ': ms) = V.Vector n (HM' (m ': ms))

newtype HM :: [Nat] -> Type where
    HM  :: { getHM :: HM' b }
        -> HM b
  deriving (Generic)

type instance Element (HM s) = Double

instance BLAS HM where

    transp (HM x) = HM (tr x)

    scal α (HM x)        = HM (konst α * x)
    axpy α (HM x) (HM y) = HM (konst α * x + y)
    dot    (HM x) (HM y) = x <.> y
    norm2  (HM x)        = norm_2 x
    asum   (HM x)        = norm_1 x

    iamax
        :: forall n. KnownNat n
        => HM '[n + 1]
        -> Finite (n + 1)
    iamax  (HM x)        = withKnownNat (SNat @n %:+ SNat @1) $
        Finite . fromIntegral
      . LA.maxIndex . extract
      $ abs x

    gemv α (HM a) (HM x) = \case
        Just (β, HM y) -> HM (konst α * (a #> x) + konst β * y)
        Nothing        -> HM (konst α * (a #> x))
    ger  α (HM x) (HM y) = \case
        Just (HM a) -> HM (konst α * outer x y + a)
        Nothing     -> HM (konst α * outer x y)
    gemm α (HM a) (HM b) = \case
        Just (β, HM c) -> HM (konst α * (a <> b) + konst β * c)
        Nothing        -> HM (konst α * (a <> b))

instance Tensor HM where
    type Scalar  HM = Double

    gen = \case
      SNil -> \f -> HM $ f Ø
      n@SNat `SCons` SNil -> \f -> HM . fromJust . create $
        LA.build (fromIntegral (fromSing n))
          (f . only . Finite . round)
      m@SNat `SCons` (n@SNat `SCons` SNil) -> \f -> HM . fromJust . create $
        LA.build (fromIntegral (fromSing m), fromIntegral (fromSing n))
          (\i j -> f (Finite (round i) :< Finite (round j) :< Ø))
      SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) -> \f -> HM $
        V.generate_ $ \i -> getHM $ gen ns (\js -> f (i :< js))

    genA = \case
      SNil -> \f -> HM <$> f Ø
      n@SNat `SCons` SNil -> \f ->
        fmap (HM . vector) . traverse (f . only . Finite) $
          [0 .. fromSing n - 1]
      m@SNat `SCons` (n@SNat `SCons` SNil) -> \f ->
        fmap (HM . matrix) . traverse f $
          [ Finite i :< Finite j :< Ø | j <- [0 .. fromSing m - 1]
                                      , i <- [0 .. fromSing n - 1]
          ]
      SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) -> \f -> fmap HM $
        sequenceA . V.generate_ $ \i -> getHM <$> genA ns (\js -> f (i :< js))

    tkonst s = HM . hkonst s

    tsum :: forall s. SingI s => HM s -> Double
    tsum = go sing . getHM
      where
        go :: Sing ns -> HM' ns -> Double
        go = \case
          SNil                                      -> id
          SNat `SCons` SNil                         -> LA.sumElements . extract
          SNat `SCons` (SNat `SCons` SNil)          -> LA.sumElements . extract
          SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) -> sum . fmap (go ns)

    tmap :: forall s. SingI s => (Double -> Double) -> HM s -> HM s
    tmap = omap

    tzip
        :: forall s. SingI s
        => (Double -> Double -> Double)
        -> HM s -> HM s -> HM s
    tzip f (HM x0) (HM y0) = HM $ hzip f sing x0 y0

    tindex
        :: forall s. SingI s
        => Prod Finite s
        -> HM s
        -> Double
    tindex ix0 = go sing ix0 . getHM
      where
        go :: Sing ns -> Prod Finite ns -> HM' ns -> Double
        go = \case
          SNil -> \case
            Ø -> id
          SNat `SCons` SNil -> \case
            i :< Ø -> (`LA.atIndex` fromIntegral i) . extract
          SNat `SCons` (SNat `SCons` SNil) -> \case
            i :< j :< Ø -> (`LA.atIndex` (fromIntegral i, fromIntegral j)) . extract
          SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) -> \case
            i :< js -> go ns js . (`V.index` i)

    tslice
        :: forall n m. ()
        => ProdMap Slice n m
        -> HM n
        -> HM m
    tslice sl0 (HM x0) = HM $ go sl0 x0
      where
        go  :: forall ns ms. ()
            => ProdMap Slice ns ms
            -> HM' ns
            -> HM' ms
        go = \case
          PMZ -> \x -> x
          PMS (Slice sL sC@SNat sR) PMZ -> \xs -> fromJust . create $
            let l = fromIntegral $ fromSing sL
                c = fromIntegral $ fromSing sC
            in  withKnownNat (sL %:+ sC %:+ sR) $
                  LA.subVector l c (extract xs)
          PMS (Slice sLy sCy@SNat sRy) (PMS (Slice sLx sCx@SNat sRx) PMZ) ->
              \xs -> fromJust . create $
            let lx = fromIntegral $ fromSing sLx
                ly = fromIntegral $ fromSing sLy
                cx = fromIntegral $ fromSing sCx
                cy = fromIntegral $ fromSing sCy
            in  withKnownNat (sLy %:+ sCy %:+ sRy) $
                withKnownNat (sLx %:+ sCx %:+ sRx) $
                  LA.subMatrix (ly, lx) (cy, cx) (extract xs)
          PMS (Slice sL sC@SNat _) pm@(PMS _ (PMS _ _)) -> \xs ->
            let l = fromIntegral $ fromSing sL
                c = fromIntegral $ fromSing sC
            in  fmap (go pm)
                  . fromJust . V.toSized
                  . UV.take c
                  . UV.drop l
                  . V.fromSized
                  $ xs

    tconv
        :: forall n m s. ()
        => Sing n
        -> ProdMap Conv m s
        -> HM (m >: n)      -- ^ mask
        -> HM s
        -> HM (s >: n)
    tconv sn@SNat pm0 (HM m0) (HM x0) = HM $ go pm0 m0 x0
      where
        go  :: forall ms ss. ()
            => ProdMap Conv ms ss
            -> HM' (ms >: n)
            -> HM' ss
            -> HM' (ss >: n)
        go = \case
          PMZ -> \m x -> konst x * m
          PMS (Conv sM sS str off) PMZ -> \m x ->
            undefined

    treshape
        :: (SingI s1, Product s1 ~ Product s2)
        => Sing s2
        -> HM s1
        -> HM s2
    treshape s = unflatten s . flatten

    tload
        :: Sing s
        -> V.Vector (Product s) Double
        -> HM s
    tload s = unflatten s . VG.convert

    textract
        :: SingI s
        => HM s
        -> V.Vector (Product s) Double
    textract = VG.convert . flatten

-- hconv
--     :: forall o ms ns. (KnownNat o)
--     => DoubleProd Sing ms ns
--     -> HM' (o ': ms)
--     -> HM' ns
--     -> HM' (o ': ns)
-- hconv = \case
--     DPZ -> \ms x -> konst x * ms
--     DPS SNat sn@SNat DPZ -> \m x -> fromJust . create $
--       let c = LA.conv2 (extract m) (LA.asRow (extract x))
--           o = fromInteger (natVal (Proxy @o))
--           left = o `div` 2
--       in  LA.subMatrix (0,left) (o, fromInteger (fromSing sn)) c
--     DPS smx@SNat snx@SNat (DPS smy@SNat sny@SNat DPZ) -> \m x ->
--       -- todo: vectorize with im2colV
--       flip fmap m $ \msk -> fromJust . create $
--         let c = LA.conv2 (extract msk) (extract x)
--             left = fromInteger (fromSing smx) `div` 2
--             top  = fromInteger (fromSing smy) `div` 2
--         in  LA.subMatrix (left, top) (fromInteger (fromSing snx), fromInteger (fromSing sny)) c
--     dp0@(DPS (SNat :: Sing m0) (SNat :: Sing n0) (DPS _ _ (DPS _ _ _) :: DoubleProd Sing ms0 ns0)) ->
--               \(ms :: V.Vector o (V.Vector m0 (HM' ms0))) (xs :: V.Vector n0 (HM' ns0)) ->
--       flip fmap ms $ \(msk :: V.Vector m0 (HM' ms0)) -> hconv1 dp0 msk xs

-- hconv1
--     :: forall m s. ()
--     => DoubleProd Sing m s
--     -> HM' m
--     -> HM' s
--     -> HM' s
-- hconv1 = \case
--     DPZ -> (*)
--     DPS sm@SNat sn@SNat DPZ -> \m x -> fromJust . create $
--       let c = LA.conv (extract m) (extract x)
--           left = fromInteger (fromSing sm) `div` 2
--       in  UVS.slice left (fromInteger (fromSing sn)) c
--     DPS smx@SNat snx@SNat (DPS smy@SNat sny@SNat DPZ) -> \m x -> fromJust . create $
--       let c = LA.conv2 (extract m) (extract x)
--           left = fromInteger (fromSing smx) `div` 2
--           top  = fromInteger (fromSing smy) `div` 2
--       in  LA.subMatrix (left, top) (fromInteger (fromSing snx), fromInteger (fromSing sny)) c
--     DPS (SNat :: Sing m0) (SNat :: Sing n0) dps@(DPS _ _ (DPS _ _ _) :: DoubleProd Sing ms0 ns0) ->
--               \(ms :: V.Vector m0 (HM' ms0)) (xs :: V.Vector n0 (HM' ns0)) ->
--       let s   :: Sing ns0
--           s   = prodSing $ secondDP dps
--           cl :: V.Vector n0 (V.Vector m0 (HM' ns0))
--           cl = im2colV (hkonst s 0) xs
--       in  fmap (hsum s . V.zipWith (hconv1 dps) ms) cl

flatten :: SingI s => HM s -> VS.Vector (Product s) Double
flatten = hflatten sing . getHM

unflatten :: Sing s -> VS.Vector (Product s) Double -> HM s
unflatten s = HM . hunflatten s

hflatten
    :: Sing s
    -> HM' s
    -> VS.Vector (Product s) Double
hflatten = \case
    SNil -> VS.singleton
    SNat `SCons` SNil -> fromJust . VS.toSized . extract
    sn@SNat `SCons` (sm@SNat `SCons` SNil) -> case sn %:* sm of
      SNat -> fromJust . VS.toSized . LA.flatten . extract
    SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) ->
      VG.convert . V.concatMap (VG.convert . hflatten ns)

hunflatten
    :: Sing s
    -> VS.Vector (Product s) Double
    -> HM' s
hunflatten = \case
    SNil -> VS.head
    SNat `SCons` SNil -> fromJust . create . VS.fromSized
    SNat `SCons` (sm@SNat `SCons` SNil) -> fromJust . create . LA.reshape (fromInteger (fromSing sm)) . VS.fromSized
    sn@SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) -> case sProduct ns of
      sp@SNat -> fromJust
            . V.fromList
            . evalState (replicateM (fromInteger (fromSing sn)) $
                           hunflatten ns . fromJust . VS.toSized <$> state (UVS.splitAt (fromInteger (fromSing sp)))
                        )
            . VS.fromSized

im2col
    :: forall m n o. (KnownNat m, KnownNat n, KnownNat o)
    => L n m
    -> L n o
im2col = undefined

im2colV
    :: forall m n a. (KnownNat m, KnownNat n)
    => a
    -> V.Vector n a
    -> V.Vector n (V.Vector m a)
im2colV pad (V.fromSized->v) = V.generate $ \i ->
      fromJust . V.toSized $ UV.slice i m padded
  where
    padded = UV.concat [UV.replicate left pad, v, UV.replicate right pad]
    m :: Int
    m  = fromIntegral $ natVal (Proxy @m)
    left  = m `div` 2
    right = m - left

hadd :: Sing s -> HM' s -> HM' s -> HM' s
hadd = \case
    SNil                                     -> (+)
    SNat `SCons` SNil                        -> (+)
    SNat `SCons` (SNat `SCons` SNil)         -> (+)
    SNat `SCons` s@(_ `SCons` (_ `SCons` _)) -> liftA2 (hadd s)

hsum :: Foldable f => Sing s -> f (HM' s) -> HM' s
hsum s = foldl' (hadd s) (hkonst s 0)

hkonst :: Sing s -> Double -> HM' s
hkonst = \case
    SNil                                     -> id
    SNat `SCons` SNil                        -> konst
    SNat `SCons` (SNat `SCons` SNil)         -> konst
    SNat `SCons` s@(_ `SCons` (_ `SCons` _)) -> pure . hkonst s

hzip
    :: (Double -> Double -> Double)
    -> Sing s
    -> HM' s
    -> HM' s
    -> HM' s
hzip f = go
  where
    go :: Sing t -> HM' t -> HM' t -> HM' t
    go = \case
      SNil                             -> f
      SNat `SCons` SNil                -> zipWithVector f
      SNat `SCons` (SNat `SCons` SNil) ->
        (\xs ys -> matrix (zipWith f xs ys))
           `on` (concat . LA.toLists . extract)
      SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) ->
         V.zipWith (go ns)

-- firstDP
--     :: DoubleProd f as bs
--     -> Prod f as
-- firstDP = \case
--     DPZ        -> Ø
--     DPS x _ xs -> x :< firstDP xs

-- secondDP
--     :: DoubleProd f as bs
--     -> Prod f bs
-- secondDP = \case
--     DPZ        -> Ø
--     DPS _ x xs -> x :< secondDP xs

prodSing
    :: Prod Sing as
    -> Sing as
prodSing = \case
    Ø       -> SNil
    x :< xs -> x `SCons` prodSing xs

instance SingI s => MonoFunctor (HM s) where
    omap f = HM . go sing . getHM
      where
        go :: Sing ns -> HM' ns -> HM' ns
        go = \case
          SNil                                      -> f
          SNat `SCons` SNil                         -> dvmap f
          SNat `SCons` (SNat `SCons` SNil)          -> dmmap f
          SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) -> fmap (go ns)

helems :: forall s. SingI s => HM s -> [Double]
helems = flip appEndo [] . go sing . getHM
  where
    go :: Sing ns -> HM' ns -> Endo [Double]
    go = \case
      SNil                                      -> \x -> Endo (x:)
      SNat `SCons` SNil                         -> Endo . (++) . LA.toList . extract
      SNat `SCons` (SNat `SCons` SNil)          -> foldMap (Endo . (++)) . LA.toLists . extract
      SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) -> foldMap (go ns)

instance SingI s => MonoFoldable (HM s) where
    ofoldMap f     = foldMap f . helems
    ofoldr f z     = foldr f z . helems
    ofoldl' f z    = foldl' f z . helems
    otoList        = helems
    oall f         = all f . helems
    oany f         = any f . helems
    onull          = (== 0) . olength
    olength _      = fromIntegral (product (fromSing (sing @_ @s)))
    olength64      = fromIntegral . olength
    ocompareLength = ocompareLength . helems
    otraverse_ f   = traverse_ f . helems
    ofor_ x        = for_ (helems x)
    omapM_ f       = traverse_ f . helems
    oforM_ x       = for_ (helems x)
    ofoldlM f x    = foldlM f x . helems
    ofoldMap1Ex f  = ofoldMap1Ex f . helems
    ofoldr1Ex f    = ofoldr1Ex f . helems
    ofoldl1Ex' f   = ofoldl1Ex' f . helems
    headEx         = headEx . helems
    lastEx         = lastEx . helems
    maximumByEx f  = maximumByEx f . helems
    minimumByEx f  = minimumByEx f . helems

deriving instance NFData (HM' s)     => NFData (HM s)
deriving instance Show (HM' s)       => Show (HM s)
deriving instance Num (HM' s)        => Num (HM s)
deriving instance Fractional (HM' s) => Fractional (HM s)
deriving instance Floating (HM' s)   => Floating (HM s)


{-# language BangPatterns #-}
{-# language GeneralizedNewtypeDeriving #-}
{-# language MultiParamTypeClasses #-}
{-# language RankNTypes #-}
{-# language RebindableSyntax #-}
{-# language TypeFamilies #-}
module W40K.Core.Prob
  ( QQ
  , Event(..)
  , Prob
  , events
  , traceLength
  , traceEvents
  , traceNumEvents
  , fmapProb
  , fmapProbMonotone
  , fmapProbMonotone'
  , (|>>=|), (|=<<|)
  , uniformly
  , probTrue
  , probFalse
  , probOf
  , given
  , bernoulli
  , binomial
  , foldlProbs'
  , foldrProbs
  , foldProbs
  , sumProbs
  , foldAssocIID
  , foldIID
  , sumIID
  , addImpossibleEvents
  , cdf
  , ccdf
  , summary
  , summaryInt
  , mean
  , variance
  , stDev
  , maxEvent
  ) where

import Prelude hiding (Functor(..), Applicative(..), Monad(..))

import Control.DeepSeq (NFData(..))

import Data.List (foldl', sort)
import Data.MemoTrie (memo)

import Numeric.SpecFunctions (choose)

import Debug.Trace (trace)

import W40K.Core.SortedList (SortedList, SortedListItem(..))
import qualified W40K.Core.SortedList as SortedList

import W40K.Core.ConstrMonad as ConstrMonad
import W40K.Core.Util

type QQ = Double

data Event a = Event !a {-# unpack #-} !QQ
    deriving (Eq, Ord)

mapEvent :: (a -> b) -> Event a -> Event b
mapEvent f (Event a p) = Event (f a) p
{-# inline mapEvent #-}

instance NFData a => NFData (Event a) where
    rnf (Event a _) = rnf a

instance Ord a => SortedListItem (Event a) where
    type ItemKey (Event a) = a
    itemKey (Event a _) = a
    combineItems (Event _ p) (Event a p') = Event a (p + p')

instance Show a => Show (Event a) where
    show (Event a p) = show a ++ ": " ++ show (realToFrac p :: Float)

newtype Prob a = Prob { density :: SortedList (Event a) }
    deriving (Eq, Ord, Show, NFData)

events :: Prob a -> [Event a]
events (Prob es) = SortedList.toAscList es
{-# inline events #-}

traceLength :: [a] -> [a]
traceLength as = trace (show (length as)) as

traceEvents :: Show a => Prob a -> Prob a
traceEvents df@(Prob _) =
    trace (show (events df)) df

traceNumEvents :: Prob a -> Prob a
traceNumEvents df@(Prob _) =
    trace (show (length (events df))) df

fmapProb :: Ord b => (a -> b) -> Prob a -> Prob b
fmapProb f (Prob evts) = Prob (SortedList.map (mapEvent f) evts)
{-# inline fmapProb #-}

{-# rules "fmap/fmapProb" fmap = fmapProb #-}
{-# specialize fmapProb :: (a -> Int) -> Prob a -> Prob Int #-}

-- only for monotone f
fmapProbMonotone :: Ord b => (a -> b) -> Prob a -> Prob b
fmapProbMonotone f (Prob evts) = Prob (SortedList.mapMonotone (mapEvent f) evts)
{-# inline fmapProbMonotone #-}

-- only for *strictly* monotone f
fmapProbMonotone' :: (a -> b) -> Prob a -> Prob b
fmapProbMonotone' f (Prob evts) = Prob (SortedList.mapMonotone' (mapEvent f) evts)
{-# inline fmapProbMonotone' #-}

bindProbWithStrat :: Ord b => (forall c. [c] -> [c]) -> Prob a -> (a -> Prob b) -> Prob b
bindProbWithStrat evalList (Prob evts) f =
    case SortedList.toAscList evts of
      [Event a _] -> f a
      evtList     -> Prob $ SortedList.concat $ evalList $ totalProb [Event (f a) p | Event a p <- evtList]
  where
    totalProb :: [Event (Prob a)] -> [SortedList (Event a)]
    totalProb ess = [SortedList.fromAscList [Event b (p*p') | Event b p' <- events y] | Event y p <- ess]

{-# inlinable bindProbWithStrat #-}

instance ConstrMonad Ord Prob where
    return a = Prob (SortedList.singleton (Event a 1))
    {-# inline return #-}
    (>>=) = bindProbWithStrat seqItems
    {-# inline (>>=) #-}
    {-# specialize (>>=) :: Prob a -> (a -> Prob Int) -> Prob Int #-}

(|>>=|) :: (Ord a, Ord b) => Prob a -> (a -> Prob b) -> Prob b
(|>>=|) = bindProbWithStrat parItems
{-# inline (|>>=|) #-}

(|=<<|) :: (Ord a, Ord b) => (a -> Prob b) -> Prob a -> Prob b
f |=<<| ma = ma |>>=| f
{-# inline (|=<<|) #-}

uniformly :: Ord a => [a] -> Prob a
uniformly as = Prob $ SortedList.fromAscList [Event a (1/n) | a <- sort as]
  where
    n = fromIntegral (length as)

probTrue :: Prob Bool -> QQ
probTrue prob =
    case events prob of
        [Event False _, Event True p] -> p
        [Event True p]                -> p
        [Event False q]               -> 1 - q
        []                            -> error "empty Prob!"
        _                             -> error ("unnormalized Prob!" ++ show prob)

probFalse :: Prob Bool -> QQ
probFalse df = 1 - probTrue df

probOf :: (a -> Bool) -> Prob a -> QQ
probOf test df = sum [p | Event a p <- events df, test a]

given :: (a -> Bool) -> Prob a -> Prob a
given hyp prob =
    Prob (SortedList.fromAscList [Event a (p / probHyp) | Event a p <- events prob, hyp a])
  where
    probHyp = probOf hyp prob

bernoulli :: QQ -> Prob Bool
bernoulli p = Prob $ SortedList.fromAscList [Event False (1-p), Event True p]

binomial :: Int -> QQ -> Prob Int
binomial 0 _ = return 0
binomial n p = binomialFlip p n

binomialFlip :: QQ -> Int -> Prob Int
binomialFlip p
  | p == 1/2  = binomial12
  | p == 1/3  = binomial13
  | p == 1/6  = binomial16
  | p == 2/3  = binomial23
  | otherwise = binomialMemo p

binomialMemo :: QQ -> Int -> Prob Int
binomialMemo p = memo $ \n ->
    case p of
      0 -> return 0
      1 -> return n
      _ -> Prob $ SortedList.fromAscList [Event k (binomProbOf n p k) | k <- [0..n]]
  where
    binomProbOf :: Int -> QQ -> Int -> QQ
    binomProbOf n p k
      | k < 0 || k > n = 0
      | n == 0         = 1
      | otherwise      = realToFrac (choose n k) * p^k * (1-p)^(n-k)

binomial23 :: Int -> Prob Int
binomial23 = binomialMemo (2/3)

binomial12 :: Int -> Prob Int
binomial12 = binomialMemo (1/2)

binomial13 :: Int -> Prob Int
binomial13 = binomialMemo (1/3)

binomial16 :: Int -> Prob Int
binomial16 = binomialMemo (1/6)

foldlProbs' :: (Ord a, Ord b) => (b -> a -> b) -> Prob b -> [Prob a] -> Prob b
foldlProbs' = foldl' . liftA2

foldrProbs :: (Ord a, Ord b) => (a -> b -> b) -> Prob b -> [Prob a] -> Prob b
foldrProbs = foldr . liftA2

foldProbs :: (Ord m, Monoid m) => [Prob m] -> Prob m
foldProbs = foldlProbs' mappend (return mempty)

sumProbs :: (Ord a, Num a) => [Prob a] -> Prob a
sumProbs []     = return 0
sumProbs [p]    = p
sumProbs (p:ps) = foldlProbs' (+) p ps
{-# specialize sumProbs :: [Prob Int] -> Prob Int #-}

foldAssocIID :: (Ord a) => (a -> a -> a) -> a -> Int -> Prob a -> Prob a
foldAssocIID _ z 0 _ = return z
foldAssocIID f z n p
  | n < 0     = error "foldAssocIID: n must be >= 0"
  | otherwise = noCheck f z n p
  where
    noCheck _ _ 1 p = p
    noCheck f z n p
      | n `mod` 2 == 0 = twiceHalf
      | otherwise      = liftA2 f p twiceHalf
      where
        !twiceHalf = liftA2 f recHalf recHalf
        !recHalf   = foldAssocIID f z (n`div`2) p

foldIID :: (Ord a, Monoid a) => Int -> Prob a -> Prob a
foldIID = foldAssocIID mappend mempty

-- TODO: Option to use Central Limit Theorem for sufficiently large n
sumIID :: (Ord a, Num a) => Int -> Prob a -> Prob a
sumIID = foldAssocIID (+) 0

addImpossibleEvents :: (Ord a, Enum a) => Prob a -> Prob a
addImpossibleEvents prob =
    let es = events prob
        as = [a | Event a _ <- es]
    in
        Prob $ SortedList.fromAscList $ merge es (zipWith Event [minimum as .. maximum as] (repeat 0))
  where
    merge es []  = es
    merge [] es' = es'
    merge ees@(e@(Event a _):es)
          ees'@(e'@(Event a' _):es')
      | a == a' = e : merge es es'
      | a <  a' = e : merge es ees'
      | a >  a' = e' : merge ees es'

cdf :: Ord a => Prob a -> [Event a]
cdf = scanl1 sumEvents . events
  where
    sumEvents (Event _ p) (Event a p') = Event a (p + p')

ccdf :: Ord a => Prob a -> [Event a]
ccdf = scanr1 sumEvents . events
  where
    sumEvents (Event a p) (Event _ p') = Event a (p + p')

summary :: Prob QQ -> IO ()
summary p =
  let mu    = "µ=" ++ show (realToFrac (mean p) :: Double)
      sigma = "σ=" ++ show (sqrt $ realToFrac (variance p) :: Double)
  in
      putStrLn $ mu ++ ", " ++ sigma

summaryInt :: Prob Int -> IO ()
summaryInt = summary . fmap fromIntegral

mean :: Prob QQ -> QQ
mean df = sum [k * p | Event k p <- events df]

variance :: Prob QQ -> QQ
variance df = sum [squared k * p | Event k p <- events df] - squared (mean df)
  where
    squared x = x*x

stDev :: Prob QQ -> QQ
stDev df = sqrt (variance df)

maxEvent :: Ord a => Prob a -> a
maxEvent df = maximum [a | Event a _ <- events df]


{-# LANGUAGE DataKinds #-}

module SCG.Graph () where
{--
import GHC.TypeLits
import Algebra.Graph
import Backprop
import Numeric.LinearAlgebra.Static

{--
Optimizing loss functions over random variables is often intractable
due to the loss functions and their gradients being either a sum over
an exponential number of latent variable configurations or non-analytic,
high-dimensional integrals.

Usually this is resolved by using problem-specific MC Grad Estimators.

There is a general technique for this!

Stochastic Computation Graphs:
- Allows derivation of unbiased gradient estimators for general expected losses.
- Estimator can be computed as the grad of a differentiable "surrogate loss" through backprop
- Variance reduction techniques can be applied to this problem formulation
- Hessian-free methods and majorization-minimization algorithms can be generalized to the SCG framework

!! The main modification to backprop is to introduce extra gradient signals at the stochastic nodes.
--}

-- Gradient Estimators for a Single Random Variable

type RV = Vector Double

type Theta = Matrix

parameterizedProbDist :: Theta -> RV

costfn :: RV -> Double
costfn rv = 2.0

-- scoreFunctionEstimator: By Log-Derivative Trick (valid if p_x_theta is continuous function of theta, though not necessarily for x)
-- also known as REINFORCE or likelihood ratio estimator
ddTheta_expect_over_x_f_x x f p theta = expectation $ x $ f x $ ddTheta $ log $ p x theta

-- if x is deterministic differentiable (perhaps representable as a type constraint? find one)
-- function of theta and another rv z, (i.e, x(z, theta)) we can use
-- !!Pathwise Derivative
-- only valid if f(x(z, theta)) is continous function of theta for all z (another type constraint?)
ddTheta_expect_over_x_f_ztheta f x z theta = expectation $ z $ ddTheta $ f (x z theta)

-- Theta might appear inside the expectation and the prob dist!
-- ddtheta Expectation_z~p(., theta)[ f( x(z, theta) ) ]
-- then two terms in the gradient estimator:

ddtheta_expect_over_z_from_pOfTheta_f_of_x_of_z_theta p f x z theta =
  expectation map (\z -> pointwiseEstimate) sample (p z theta)
  where pointwiseEstimate = ddTheta $ f (x z theta) + (ddTheta $ log (p z theta)) * (f $ x z theta)

{--
1. SF can be used even if f is discontinuous or x is a discrete rv
2. SF only requires trajectories, PD requires f'(x)
3. SF has higher variance than PD, unless f is rough as in time-series problems with exploding gradients
4. PD has a deterministic limit, SF doesn't.
--}

{--
  STOCHASTIC COMPUTATION GRAPHS:
    Directed, Acyclic Graph with Three Kinds of nodes:
    - Input Nodes, set externally including the parameters we differentiate with respect to
    - Deterministic Nodes: pure functions of their parents
    - Stochastic Nodes: Distributed Conditionally on their parents.

    Each parent v of a non-input node w is connected to it by a directed edge (v, w)

  THE STRUCTURE OF THE GRAPH FULLY SPECIFIES:
    - What estimator we will use, SF or PD or a combination thereof
    - Nodes arranged in series are multiplicative terms only
    - Nodes arranged in parallel lead to sums over mulplicative terms
    -
--}
{--
data Node a = InputNode a | DeterministicNode a | StochasticNode a

newtype InputNode a = IO a

newtype DeterministicNode v e = ([v] -> e -> b)

newtype StochasticNode v e = ConditionalDistribution [v] e

data DirectedEdge a = DirectedEdge
  { _v :: Node,
    _w :: DeterministicNode | StochasticNode
  }
--}

-- TODO: Create a directed acyclic graph DS using alga
-- TODO: Define Node types which typecheck (initially for a scalar parameterization of the relevant types)
-- TODO: Add backprop for folds over the deterministic nodes? derive an instance of Backprop for edges?
-- TODO: Add the right gradient estimation procedure for combinations of deterministic and stochastic nodes
-- TODO: Figure out how to do dataflow programming over the Graph
--}


-- Graphical and Network Models in Haskell.
-- Nicolas Kim
-- Department of Statistics, Carnegie Mellon University
-- Distributed under the MIT License.

import System.IO
import System.Random

import Adjacency
import Core
import Local

import Control.Monad
import Control.Monad.Random
import Control.Parallel (par, pseq)
-- import Control.Lens

import qualified Data.Map.Lazy as M
import qualified Data.List as L

import Numeric.LinearAlgebra (dispf)

-- import qualified Numeric.LinearAlgebra as LA

-- Some examples of graphs.
gr1 = Graph [1..3] [(1,2), (1,3), (2,3)] :: Graph Int
gr2 = Graph [1..2] [(1,2), (1,3), (2,3)] :: Graph Int
gr3 = completeGraph [1..10] :: Graph Int
gr4 = Graph [1..4] [(1,2), (3,4), (1,4)] :: Graph Int

kite = Graph [1..4] [(1,2), (1,3), (2,3), (2,4), (3,4)] :: Graph Int
wiki = Graph [1..6] [(1,2), (1,5), (2,5), (2,3), (3,4), (4,5), (4,6)] ::
  Graph Int
gr5 = Graph [1..3] [(1,2)] :: Graph Int
disconnected = Graph [1..4] [(1,2), (3,4)] :: Graph Int

-- Data structure for network models.
-- data Model = ERG { parameters :: [Double] } deriving (Show, Eq)

-- Makes an empty graph.
emptyGraph :: Graph a
emptyGraph = Graph [] []

-- Makes a complete simple graph from a list of nodes.
completeGraph :: (Ord a, Eq a) => [a] -> Graph a
completeGraph ns = Graph sns es
  where
    sns = L.sort ns
    es = completeEdges sns

-- List of all possible graphs on n nodes.
-- allGraphs :: (Ord a, Eq a) => [a] -> [Graph a] -> M.Map [Int] [Graph a]
-- allGraphs ns = L.foldl' () mapInit [1..n]
--   where
--     mapInit = M.fromList [(0, emptyGraph)]
--     n = (\x -> x*(x-1) `div` 2) $ length ns

-- Makes a valid graph object from a list of edges.
toGraph :: (Ord a, Eq a) => [(a, a)] -> Graph a
toGraph es = Graph (extractNodes es) es

-- Checks if the edges contain defined nodes.
validGraph :: (Eq a) => Graph a -> Bool
validGraph (Graph ns es) = subset (L.union (map fst es) (map snd es)) ns
  where
    subset xs ys = all (`elem` ys) xs

-- Removes redundant edges (e.g. (1,2) = (2,1)).
redundant :: (Ord a, Eq a) => [(a, a)] -> [(a, a)]
redundant es = L.intersect es (map f es)
  where
    f (x, y) = (min x y, max x y)

-- Makes a graph valid (see the `validGraph` function).
makeValid :: (Eq a) => Graph a -> Graph a
makeValid gr@(Graph ns es)
  | validGraph gr = gr
  | otherwise = Graph (L.union ens ns) es
    where
      ens = L.nub $ L.union (map fst es) (map snd es)

-- From a list of edges, give the sorted and pruned list of nodes
extractNodes :: (Ord a, Eq a) => [(a, a)] -> [a]
extractNodes es = L.nub $ L.sort allnodes
  where
    allnodes = (map fst es) ++ (map snd es)

-- The complete set of edges between nodes in a list. Nodes don't need to
-- be orderable. 22s
completeEdges' :: [a] -> [(a, a)]
completeEdges' ns = [ (fst x, fst y) | x <- zs, y <- zs, snd x < snd y ]
  where
    zs = zip ns [1..]

-- Add a set of nodes to a graph using their labels.
addNodes :: (Ord a) => Graph a -> [a] -> Graph a
addNodes (Graph ns es) nns = Graph (L.union ns nns) es

-- Add a set of edges to a graph.
addEdges :: (Ord a) => Graph a -> [(a, a)] -> Graph a
addEdges (Graph ns es) ees = Graph ns (L.union es ees)

-- -- Various graph statistics.
-- Computes the number of nodes.
nNode :: Graph a -> Int
nNode gr = length $ nodes gr

-- Computes the number of edges.
nEdge :: Graph a -> Int
nEdge gr = length $ edges gr

-- -- Monad-based generation of random nodes and edges, i.e. random graphs.
-- Erdos-Renyi Model
-- Include or don't include an edge?
pBool :: (RandomGen g) => Double -> Rand g Bool
pBool p = liftM (< p) $ getRandomR ((0, 1) :: (Double, Double))

-- Decide over a list of edges.
pBools :: (RandomGen g) => [Double] -> Rand g [Bool]
pBools ps = sequence $ map pBool ps

-- Given a list of Bools and any other list, return the list subsetted by
-- the list of Bools.
byBool :: [Bool] -> [a] -> [a]
byBool bs xs = [ snd zs | zs <- (zip bs xs), fst zs ]

-- Select elements from a list independently and with probabilities ps.
setEdges :: (RandomGen g) => [a] -> [Double] -> Rand g [a]
setEdges xs ps = liftM (flip byBool xs) $ pBools ps

-- Generate an Erdos-Renyi random graph; nodes ns, and edge probabilities
-- ps.
erdosGen :: (Ord a, RandomGen g) => [a] -> [Double] -> Rand g (Graph a)
erdosGen ns ps = liftM (Graph ns) $ setEdges (completeEdges ns) ps

-- Graphon Model: generates a w-random graph.
graphonGen :: (Ord a, RandomGen g) => [a] -> (Double -> Double -> Double)
  -> Rand g (Graph a)
graphonGen ns w = liftM (Graph ns) es
  where
    us = liftM (take (length ns)) $ getRandomRs ((0, 1) :: (Double, Double))
    es = (liftM (applyUpper w) us) >>= (setEdges (completeEdges ns))

-- Graphon Model: generates a w-random graph. Parallel.
graphonGen' :: (Ord a, RandomGen g) => [a] -> (Double -> Double -> Double)
  -> Rand g (Graph a)
graphonGen' ns w = liftM2 pseq (liftM2 par (liftM force es) (liftM force
                               es')) (liftM (Graph ns) (liftM2 (++) es
                               es'))
  where
    us = liftM (take (length ns)) $ getRandomRs ((0, 1) :: (Double, Double))
    cs = completeEdges ns
    n = length ns
    cshalf = splitAt (((n * (n-1))+2) `div` 4) cs
    es = (liftM (applyUpper w) us) >>= (setEdges (fst cshalf))
    es' = (liftM (applyUpper w) us) >>= (setEdges (snd cshalf))

-- From "Real World Haskell". Used in the parallel version of the code for
-- generating w-random graphs, `graphonGen'`.
force :: [a] -> ()
force xs = go xs `pseq` ()
  where go (_:xs) = go xs
        go [] = 1

-- Apply a function to the upper triangle of an array.
applyUpper :: (a -> a -> b) -> [a] -> [b]
applyUpper f xs = [ f (fst x) (fst y) | x <- zl, y <- zl, snd x < snd y ]
  where
    zl = zip xs ([1..] :: [Int])

sblock :: Double -> Double -> Double
sblock x y
  | x < 0.5 && y < 0.5 = 0.8
  | x < 0.5 || y < 0.5 = 0.2
  | otherwise          = 0.8

-- Need to create a show instance for this (or rewrite it).
round' :: (RealFrac a, RealFrac b, Integral b) => Int -> a -> b
round' x = (/ 10^x) . round . (* 10^x)

main :: IO ()
main = do
    -- values <- evalRandIO $ graphonGen [1..10] (\x y -> (x+y)/2)
    let n = 100 :: Int
    gr <- evalRandIO $ graphonGen [1..n] sblock
    putStrLn . show . length $ community (toMap gr) 1
    -- let spannum = spanTreeCount values
    -- putStrLn . (\x -> "Number of spanning trees: " ++ x) $ show spannum
    -- putStrLn . show $ spectralCluster' values
    -- putStrLn . (\x -> "Guess for optimal lambda: " ++ x) . show
    --  $ lambdaSelect values
    -- putStrLn . (\x -> "This is " ++ x ++ "% of the true value." ) . show
    --   . round' 4 . ((/) $ fromIntegral n^(n-2)) $ fromIntegral spannum


-- |
-- Module      : Occlusion.Vector
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created Jonatan H Sundqvist 2015

-- TODO | - Most of these functions should be moved to a utility library (eg. Cartesian)
--        -

-- SPEC | -
--        -

--------------------------------------------------------------------------------------------------------------------------------------------
-- GHC Pragmas
--------------------------------------------------------------------------------------------------------------------------------------------

--------------------------------------------------------------------------------------------------------------------------------------------
-- API
--------------------------------------------------------------------------------------------------------------------------------------------
module Occlusion.Vector where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
import Data.Complex
import Data.Functor ((<$>))
import Data.List    (sort, sortBy)
import Control.Applicative
import Control.Lens
import Control.Monad

import Southpaw.Math.Constants

import Occlusion.Types
import Occlusion.Lenses

--------------------------------------------------------------------------------------------------------------------------------------------
-- Functions
--------------------------------------------------------------------------------------------------------------------------------------------
-- |
vectorise :: (a -> a -> b) -> Complex a -> b
vectorise f (x:+y) = f x y

-- Plumbing --------------------------------------------------------------------------------------------------------------------------------

-- |
dotwise :: (a -> a -> a) -> Complex a -> Complex a -> Complex a
dotwise f (x:+y) (x':+y') = f x x':+f y y'

-- |
dotmap :: (a -> a) -> Complex a -> Complex a
dotmap f (x:+y) = f x:+f y

-- Linear functions ------------------------------------------------------------------------------------------------------------------------

-- |
-- TODO: Refactor
-- TODO: Invariants, check corner cases
-- TODO: Deal with vertical lines
-- TODO: Factor out infinite-line logic
-- TODO: Decide how to deal with identical lines
-- TODO: Factor out domain logic (eg. write restrict or domain function)
-- TODO: Visual debugging functions
intersect :: RealFloat f => Line f -> Line f -> Maybe (Complex f)
-- intersect f@(Line a b) g@(Line a' b')
intersect f' g' = mp >>= \p -> indomain f' p >> indomain g' p
  where
    indomain h' = restrict (h'^.linebegin) (h'^.linestop)
    mp = case [linear f', linear g'] of
      [Just f, Nothing] -> let x = g'^.linebegin.real in Just $ (x):+(plotpoint f x)
      [Nothing, Just g] -> let x = f'^.linebegin.real in Just $ (x):+(plotpoint g x)
      [Just f,  Just g] -> linearIntersect f g
      _                 -> Nothing

-- | Gives the linear function overlapping the given segment
linear :: RealFloat f => Line f -> Maybe (Linear f)
linear line = (,) <$> slope line <*> intercept line

-- | Applies a linear function to the given value
-- TODO: Rename (?)
plotpoint :: RealFloat f => Linear f -> f -> f
plotpoint (slope', intercept') x = slope'*x + intercept'

-- | Finds the intersection (if any) of two linear functions
linearIntersect :: RealFloat f => Linear f -> Linear f -> Maybe (Complex f)
linearIntersect (kf, mf) (kg, mg)
  | kf == kg  = Nothing
  | otherwise = let x = (mg-mf)/(kg-kg)
                    y = (mf*x + kf)
                in Just $ x:+y

-- |
slope :: RealFloat f => Line f -> Maybe f
slope (Line fr to)
  | dx == 0   = Nothing
  | otherwise = Just $ dy/dx
  where
    (dx:+dy) = to - fr

-- |
intercept :: Line f -> Maybe f
intercept (Line a b) = error ""

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

-- |
between :: Ord a => a -> a -> a -> Bool
between mini maxi a = mini <= a && a <= maxi

-- | Ensures that a given point lies within the domain and codomain
-- TODO: Let thus function work on scalars, write another function for domain and codomain (?)
-- restrict domain codomain p = _
restrict :: Ord f => Complex f -> Complex f -> Complex f -> Maybe (Complex f)
restrict a b p@(x:+y)
  | indomain && incodomain = Just p
  | otherwise              = Nothing
  where
    (lowx:+lowy)   = dotwise min a b
    (highx:+highy) = dotwise max a b
    indomain       = between lowx highx x
    incodomain     = between lowy highy y

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

-- |
pass :: Monad m => m ()
pass = return ()

-- |
unit :: Monad m => a -> m a
unit = return

-- | Like maybe, except the function comes at the end
perhaps :: b -> Maybe a -> (a -> b) -> b
perhaps d m f = maybe d f m