module Legacy.Base {-# DEPRECATED "Use Vector instead" #-} where

import Data.Complex
import Data.Either
import Data.List
import Control.Applicative

newtype StateVector = StateVector [Complex Double]
    deriving (Eq, Show)

newtype TransformTensor = TransformTensor [[Complex Double]]
    deriving (Eq, Show)

{-|
    integerLog2 determines if an integer is a power of 2
    If it is, the log base 2 of the integer will be wrapped
    by Right constructor; if not, Left constructor with
    corresponding error will be returned.
    The function takes 1 argument of type 'Int'
-}
integerLog2 :: Int -> Either String Int
integerLog2 x
    | x < 1 = Left "Input less than 0"
    | x == 1 = Right 0
    | y * 2 == x = (+1) <$> integerLog2 y
    | otherwise = Left "Input is not a power of 2"
    where
        y = quot x 2

initQubit :: Complex Double -> Complex Double -> Either String StateVector
initQubit a b
    | magnitude a ** 2 + magnitude b ** 2 == 1 = Right $ StateVector [a, b]
    | otherwise = Left "Magnitude Error"

initQubit0 :: StateVector
initQubit0 = StateVector [1:+0, 0:+0]

initNumQubits :: Int -> Either String StateVector
initNumQubits n
    | n < 1 = Left "Number of Qubits Must be Greater than 1"
    | otherwise = Right $ foldr mergeState (StateVector [1 :+ 0]) $ replicate n initQubit0

initState :: [Complex Double] -> Either String StateVector
initState q
    | isLeft numQubits = Left "Input State Dimension Invalid"
    | otherwise = Right $ StateVector q
    where
        numQubits = integerLog2 $ length q

mergeState :: StateVector -> StateVector -> StateVector
mergeState (StateVector x) (StateVector y) = StateVector $ liftA2 (*) x y

mergeTensor :: TransformTensor -> TransformTensor -> TransformTensor
mergeTensor (TransformTensor x) (TransformTensor y) = TransformTensor $ fmap ((<*>) . fmap (*)) x <*> y

tensorAdd :: TransformTensor -> TransformTensor -> Either String TransformTensor
tensorAdd (TransformTensor i) (TransformTensor j)
    | m == x && n == y = Right $ TransformTensor $ zipWith (zipWith (+)) i j
    | otherwise = Left "Adding Two Tensors with Different Dimensions"
    where
        m = length i
        n = length $ head i
        x = length j
        y = length $ head j

doubleBindEither :: Either String TransformTensor -> Either String TransformTensor -> (TransformTensor -> TransformTensor -> Either String TransformTensor) -> Either String TransformTensor
doubleBindEither x y f
    | isLeft x && isRight y = Left $ fromLeft [] x ++ "\n" ++ fromLeft [] y
    | isLeft x = x
    | isLeft y = y
    | otherwise = f (fromRight (TransformTensor [[]]) x) (fromRight (TransformTensor [[]]) y)

applyTensorToState :: TransformTensor -> StateVector -> Either String StateVector
applyTensorToState (TransformTensor g) (StateVector q)
    | length f == length q && length q == length g =
        Right $ StateVector $ foldr (zipWith (+)) acc $ zipWith fmap (fmap (*) q) f
    | otherwise = Left "Transformation Matrix Dimension Mismatch"
    where
        acc = replicate (length q) $0:+0
        f = transpose g

-- The qubits are 1-indexed
applyGateToQubit :: TransformTensor -> Int -> StateVector -> Either String StateVector
applyGateToQubit g n (StateVector q) = t >>= (`applyTensorToState` StateVector q)
    where
        numQubits = integerLog2 $ length q
        t = numQubits >>= flip (genTensor g) n

-- Gate to insert -> Number of Qubits -> Index of Qubit
genTensor :: TransformTensor -> Int -> Int -> Either String TransformTensor
genTensor g n i
    | n < i || i < 1 = Left "Index Out of Range"
    | otherwise = Right $ foldr mergeTensor (TransformTensor [[1:+0]]) expGate
    where
        expGate = fstId ++ [g] ++ lstId
        fstId = replicate (i - 1) pauliId
        lstId = replicate (n - i) pauliId

getQubitState :: Int -> StateVector -> Either String (Double, Double)
getQubitState x (StateVector q)
    | x < 1 = Left "Index Out of Bounds"
    | otherwise = Right (zeroP, oneP)
    where
        numRow = 2 ^ (x - 1)
        q2 = fmap (** 2) q
        zeroP = magnitude $ sum $ takeHalf q2 numRow
        oneP = magnitude $ sum $ takeHalf (drop numRow q2) numRow
        takeHalf :: [Complex Double] -> Int -> [Complex Double]
        takeHalf xs n
            | length xs < n = []
            | otherwise = take n xs ++ takeHalf (drop (2 * n) xs) n

-- Gate to Apply -> Control Bit Index -> Application Bit Index -> Circuit To Apply
applyControlToQubits :: TransformTensor -> Int -> Int -> StateVector -> Either String StateVector
applyControlToQubits g m n (StateVector q)
    | m < n = t m n mask1 g >>= (`applyTensorToState` StateVector q)
    | m > n = t n m g mask1 >>= (`applyTensorToState` StateVector q)
    | otherwise = Left "Control and Application Qubit Cannot be Identical"
    where
        zeroMask = numQubits >>= flip (genTensor mask0) m
        numQubits = integerLog2 $ length q
        t :: Int -> Int -> TransformTensor -> TransformTensor -> Either String TransformTensor
        t i j fstGate sndGate = doubleBindEither (fmap mergeTensor fstTensor <*> sndTensor) zeroMask tensorAdd
            where
                fstTensor = genTensor fstGate i i
                sndTensor = numQubits >>= flip (genTensor sndGate) (j - i) . flip (-) i

pauliX :: TransformTensor
pauliX = TransformTensor [[0 :+ 0, 1 :+ 0],
                      [1 :+ 0, 0 :+ 0]]

pauliY :: TransformTensor
pauliY = TransformTensor [[0 :+ 0, 0 :+ (-1)],
                      [0 :+ 1, 0 :+ 0]]

pauliZ :: TransformTensor
pauliZ = TransformTensor [[1 :+ 0, 0 :+ 0],
                      [0 :+ 0, (-1) :+ 0]]

pauliId :: TransformTensor
pauliId = TransformTensor [[1 :+ 0, 0 :+ 0],
                      [0 :+ 0, 1 :+ 0]]

hadamard :: TransformTensor
hadamard = TransformTensor [[1 / sqrt 2 :+ 0, 1 / sqrt 2 :+ 0],
                        [1 / sqrt 2 :+ 0, (-1) / sqrt 2 :+ 0]]

mask0 :: TransformTensor
mask0 = TransformTensor [[1 :+ 0, 0 :+ 0],
                    [0 :+ 0, 0 :+ 0]]

mask1 :: TransformTensor
mask1 = TransformTensor [[0 :+ 0, 0 :+ 0],
                    [0 :+ 0, 1 :+ 0]]

-- main :: IO()
-- main = print result
--     where q0 = initQubit0
--           crt = mergeState q0 q0
--           flip1 = applyGateToQubit hadamard 1 crt
--           flip2 = flip1 >>= applyControlToQubits pauliY 1 2
--           --result = flip2 >>= (getQubitState 2)
--           result = TransformTensor [[0 :+ 0, 0 :+ 0],
--                                     [0 :+ 0, 1 :+ 0]]

-- Linear algebra versions

-- safeMultiply :: Num a => T.Tensor a -> T.Tensor a -> Either String (T.Tensor a)
-- safeMultiply (T.Tensor x) (T.Tensor y)
--     | T.isSquare (T.Tensor x)
--         && T.isSquare (T.Tensor y)
--         && length x == length y = Right $ T.Tensor x * T.Tensor y
--     | otherwise = Left "Dimension mismatch"
-- safeMultiply x y = Right $ x * y

-- -- combine two qubit circuits
-- addQubit :: Num a => T.Tensor a -- ^ first half statevector
--   -> T.Tensor a -- ^ second half statevector
--   -> T.Tensor a
-- addQubit (T.Tensor x) (T.Tensor y)
--     | T.rank (T.Tensor x) /= 1 || T.rank (T.Tensor y) /= 1 = error "One of the elements is not a valid statevector"
--     | otherwise = T.Tensor $ liftA2 (*) x y
-- addQubit _ _ = error "One of the elements is not a valid statevector"

-- Applies input gate to qubit of index

-- -- Builds a transformation tensor to apply input gate to selected qubit
-- build :: RealFloat a => Int -- ^ Number of qubits this tensor applies to
--   -> [(Int, T.Tensor (Complex a))]-- ^ Index / Gate pair (Pre-sorted)
--   -> T.Tensor (Complex a)-- ^ Default Gate
--   -> T.Tensor (Complex a)
-- build len lst dft = T.flatten2d $ foldr T.product (T.fromList [1]) (go len lst 0)
--     where
--         go l (x:xs) offset
--             | fst x >= l = go l [] offset
--             | otherwise = replicate (fst x - offset) dft ++ [snd x] ++ go l xs (fst x + 1)
--         go l [] offset = replicate (l - offset) dft

-- applyGate :: RealFloat a => T.Tensor (Complex a) -- ^ Gate tensor
--   -> Int -- ^ Index
--   -> T.Tensor (Complex a) -- ^ State vector
--   -> Either String (T.Tensor (Complex a))
-- applyGate (T.Tensor xs) i (T.Tensor ys)
--     | T.rank x /= 2 || not (T.isSquare x) || numT == 0 = Left "Invalid Gate Tensor"
--     | T.rank y /= 1 || numV == 0 = Left "Invalid Statevector"
--     | i < 0 || i >= numV = Left "Index out of bounds"
--     | otherwise = Right $ build numV [(i, x)] pauliId * y
--     where
--         x = T.Tensor xs
--         y = T.Tensor ys
--         numT = fromRight 0 $ integralLog 2 (length xs)
--         numV = fromRight 0 $ integralLog 2 (length ys)
-- applyGate _ _ _ = Left "Input is not a Tensor"

-- applyGateAll :: RealFloat a => T.Tensor (Complex a) -> T.Tensor (Complex a) -> Either String (T.Tensor (Complex a))
-- applyGateAll (T.Tensor xs) (T.Tensor ys)
--     | T.rank x /= 2 || not (T.isSquare x) || numT == 0 = Left "Invalid Gate Tensor"
--     | T.rank y /= 1 || numV == 0 = Left "Invalid Statevector"
--     | otherwise = Right $ build numV [] x * y
--     where
--         x = T.Tensor xs
--         y = T.Tensor ys
--         numT = fromRight 0 $ integralLog 2 (length xs)
--         numV = fromRight 0 $ integralLog 2 (length ys)
-- applyGateAll _ _ = Left "Input is not a Tensor"

-- applyControl :: RealFloat a => T.Tensor (Complex a) -- ^ Gate tensor
--   -> Int -- ^ Control qubit index
--   -> Int -- ^ Apply qubit index
--   -> T.Tensor (Complex a) -- ^ State vector
--   -> Either String (T.Tensor (Complex a))
-- applyControl (T.Tensor xs) ctl i (T.Tensor ys)
--     | T.rank x /= 2 || not (T.isSquare x) || numT == 0 = Left "Invalid Gate Tensor"
--     | T.rank y /= 1 || numV == 0 = Left "Invalid Statevector"
--     | i < 0 || i >= numV = Left "Application qubit index out of bounds"
--     | ctl < 0 || ctl >= numV = Left "Control qubit index out of bounds"
--     | ctl == i = Left "Control and application qubit cannot be the same"
--     | otherwise = Right $ (build numV [(ctl, mask0)] pauliId
--         + build numV (sortBy (compare `on` fst) [(ctl, mask1), (i, x)]) pauliId) * y
--     where
--         x = T.Tensor xs
--         y = T.Tensor ys
--         numT = fromRight 0 $ integralLog 2 (length xs)
--         numV = fromRight 0 $ integralLog 2 (length ys)

-- applyControl _ _ _ _ = Left "Input is not a Tensor"


module VB where

import Numeric.SpecFunctions (digamma)
import Numeric.LinearAlgebra.HMatrix
       (Matrix, size, (!), fromRows, scalar, sumElements, toRows, reshape,
        flatten, tr, dot, (<>), ( #> ), flatten, asRow, matrix, fromList, Vector, cmap)
import Numeric.LinearAlgebra.Devel (zipVectorWith)

data LogType
    = Log  -- ^ log x
    | InvLog  -- ^ log (1 - x)

type ExpectationQZ = Matrix Double
type ExpectationQmu = (Matrix Double, Matrix Double)  -- ^ Matrix Log and Matrix InvLog
type QZ = Matrix Double
type Qmu = (Matrix Double, Matrix Double)  -- ^ Matrix A and matrix B

-- | Moments of logarithmically transformed random variables
expectationBeta :: Double  -- ^ alpha
                -> Double -- ^ beta
                -> LogType  -- ^ log x or log (1 - x)
                -> Double
expectationBeta alpha beta (Log) = digamma alpha -
    digamma (alpha + beta)
expectationBeta alpha beta (InvLog) = digamma beta -
    digamma (alpha + beta)

em :: Matrix Double
   -> Vector Double
   -> Vector Double
   -> Double
   -> Double
   -> Vector Double
em = go 0 (-1)
  where
    go lQthetaNew lQthetaOld matrixX vectorPi vectorAlpha a b
        | abs (lQthetaNew - lQthetaOld) <
              0.1 =
            vectorPi
        | otherwise =
            let (eQz,eQmu) = eStep matrixX initQZ initQmu
                vectorPiNew = mStep vectorPi vectorAlpha eQz
            in go
                   (lQTheta matrixX vectorPiNew vectorAlpha eQz eQmu)
                   lQthetaNew
                   matrixX
                   vectorPiNew
                   vectorAlpha
                   a
                   b
        where initQZ = fromRows (replicate nRow vectorPi)
              initQmu = (matrixAB a, matrixAB b)
              matrixAB v = matrix
                      nColumn
                      [v | _ <-
                              [1 .. nMixtures * nColumn]]
              (nRow,nColumn) = size matrixX

eStep :: Matrix Double -> QZ -> Qmu -> (ExpectationQZ, ExpectationQmu)
eStep matrixX qZ qMu = let eQmu = calculateEQmu qMu
                           eQz = calculateEQZ qZ
                           qZ' = calculateQZ matrixX eQmu qZ
                           eQz' = calculateEQZ qZ'
                           qMu' = calculateQmu matrixX eQz' qMu
                           eQmu' = calculateEQmu qMu'
    in go
           (lQ matrixX eQz' eQmu')
           (lQ matrixX eQz eQmu)
           qZ'
           eQz'
           qMu'
           eQmu'
    where go lQnew lQold qZn eQzn qMun eQmun
              | abs (lQnew - lQold) <
                    0.1 =
                  (eQzn, eQmun)
              | otherwise =
                  let qZn' = calculateQZ matrixX eQmun qZn
                      eQzn' = calculateEQZ qZn'
                      qMun' = calculateQmu matrixX eQzn' qMun
                      eQmun' = calculateEQmu qMun'
                  in go (lQ matrixX eQzn' eQmun') lQnew qZn' eQzn' qMun' eQmun'

lQ :: Matrix Double -> ExpectationQZ -> ExpectationQmu -> Double
lQ matrixX eQz (eQmuM,eQmuMi) = sum
        (map
             (\n ->
                   (eQz ! n) `dot`
                   flatten
                       (asRow (matrixX ! n) <>
                        teQmuM +
                        asRow
                             (1 -
                              (matrixX ! n)) <>
                        teQmuMi))
             [0 .. nRow - 1])
    where teQmuM = tr eQmuM
          teQmuMi = tr eQmuMi
          (nRow,_) = size matrixX

calculateEQZ :: QZ -> ExpectationQZ
calculateEQZ qZ = fromRows
        (map
             (\row ->
                   row /
                   scalar (sumElements row))
             (toRows qZ))

calculateQZ :: Matrix Double -> ExpectationQmu -> QZ -> QZ
calculateQZ matrixX (eQmuM,eQmuMi) qZ = qZ + matrixX <> tr eQmuM +
    (1 - matrixX) <>
    tr eQmuMi

calculateEQmu :: Qmu -> ExpectationQmu
calculateEQmu (aM, bM) = (reshape d qMuM, reshape d qMuMi)
  where
    aV = flatten aM
    bV = flatten bM
    qMuM = zipVectorWith (\a b -> expectationBeta a b Log) aV bV
    qMuMi = zipVectorWith (\a b -> expectationBeta a b InvLog) aV bV
    (k, d) = size aM

calculateQmu :: Matrix Double -> ExpectationQZ -> Qmu -> Qmu
calculateQmu matrixX eQz (aOld,bOld) = ( aOld - 1 + tr eQz <> matrixX
                                       , bOld - 1 +
                                         tr eQz <>
                                         (1 - matrixX))

lQTheta :: Matrix Double
        -> Vector Double
        -> Vector Double
        -> ExpectationQZ
        -> ExpectationQmu
        -> Double
lQTheta matrixX vectorPi vectorAlpha eQz eQmu = lQ matrixX eQz eQmu +
    vectorAlpha `dot`
    cmap log vectorPi

mStep :: Vector Double -> Vector Double -> ExpectationQZ -> Vector Double
mStep vectorPi vectorAlpha eQz = let vectorNewPi = fromList
                                             (map
                                                  (\k ->
                                                        -(vectorPi ! k) *
                                                         sumElements
                                                              (eqZt ! k) /
                                                         (vectorAlpha ! k - 1))
                                                  [0 .. nMixtures - 1])
    in norm vectorNewPi
    where eqZt = tr eQz
          norm v = v /
              scalar (sumElements v)

nMixtures :: Int
nMixtures = 10


module ML.Unsupervised.Clustering.KMeans (
  fit,
  predict,
  KMeans(..)
) where

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Data

import Data.Ratio ((%))
import Data.List (minimumBy)
import qualified Data.Vector.Storable as V

import Debug.Trace

data KMeans a l = KMeans {
  centroids :: [(l, V.Vector a)]
} deriving (Show, Eq)

fit :: (Fractional a, Ord a, Eq l, Show l, Show a, V.Storable a) => [l] -> [V.Vector a] -> KMeans a l
fit labels points = fit' (KMeans initialCentroids)
  where
    fit' kmeans = let
                    ass = assignment kmeans
                    updatedKmeans = update kmeans ass in
                    if (trace ("updated assignment: " ++ show (assignment updatedKmeans)) assignment updatedKmeans) == ass
                      then updatedKmeans
                      else fit' updatedKmeans
    update kmeans ass = KMeans $ zip labels $ fmap (\l -> centerOfMass $ snd <$> filter (\x -> fst x == l) ass) labels
    assignment kmeans = fmap (\x -> (assign kmeans x, x)) points
    assign kmeans point = fst $ minimumBy (\x1 x2 -> snd x1 `compare` snd x2) $ (fmap (\(l, cent) -> (l, distance cent point))) $ centroids kmeans
    distance v1 v2 = V.sum $ V.zipWith (\x1 x2 -> (x2 - x1) * (x2 - x1)) v1 v2
    initialCentroids = zip labels points

    centerOfMass :: (Fractional a, V.Storable a) => [V.Vector a] -> V.Vector a
    centerOfMass (p:pts) = V.map (/ (fromRational $ (toInteger $ length pts + 1) % 1)) $ foldr sumVectors p pts

    dim = V.length $ head points
    sumVectors :: (Num a, V.Storable a) => V.Vector a -> V.Vector a -> V.Vector a
    sumVectors = V.zipWith (+)

predict :: (Ord a, Num a, V.Storable a) => KMeans a l -> V.Vector a -> l
predict kmeans point = fst $ minimumBy (\x1 x2 -> snd x1 `compare` snd x2) $ (fmap (\(l, cent) -> (l, distance cent point))) $ centroids kmeans
  where
    distance v1 v2 = V.sum $ V.zipWith (\x1 x2 -> (x2 - x1) * (x2 - x1)) v1 v2

fitMatrix :: (Element a) => [l] -> Matrix a -> KMeans a l
fitMatrix = undefined


{-# LANGUAGE DefaultSignatures          #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE StandaloneDeriving         #-}

{-|
Module: Data.Semiring
Description: Haskell semirings
License: MIT
Maintainer: mail@doisinkidney.com
Stability: experimental
-}

module Data.Semiring
  ( Semiring(..)
  , Add(..)
  , Mul(..)
  , mul
  , add
  ) where

import           Data.Coerce           (coerce)

import           Data.Functor.Identity (Identity (..))

import           Data.Monoid
import           Data.Semigroup        (Max (..), Min (..))

import           Data.Complex          (Complex)
import           Data.Fixed            (Fixed, HasResolution)
import           Data.Ratio            (Ratio)
import           Numeric.Natural       (Natural)

import           Data.Set              (Set)
import qualified Data.Set              as Set

import           Data.Int              (Int16, Int32, Int64, Int8)
import           Data.Word             (Word16, Word32, Word64, Word8)
import           Foreign.C.Types       (CChar, CClock, CDouble, CFloat, CInt,
                                        CIntMax, CIntPtr, CLLong, CLong,
                                        CPtrdiff, CSChar, CSUSeconds, CShort,
                                        CSigAtomic, CSize, CTime, CUChar, CUInt,
                                        CUIntMax, CUIntPtr, CULLong, CULong,
                                        CUSeconds, CUShort, CWchar)

import           Foreign.Ptr           (IntPtr, WordPtr)
import           System.Posix.Types    (CCc, CDev, CGid, CIno, CMode, CNlink,
                                        COff, CPid, CRLim, CSpeed, CSsize,
                                        CTcflag, CUid, Fd)

import           GHC.Generics          (Generic, Generic1)

import           Test.QuickCheck       (Arbitrary)

-- | A <https://en.wikipedia.org/wiki/Semiring Semiring> is like the
-- the combination of two 'Data.Monoid.Monoid's. The first
-- is called '<+>'; it has the identity element 'zero', and it is
-- commutative. The second is called '<.>'; it has identity element 'one',
-- and it must distribute over '<+>'.
--
-- = Laws
-- == Normal 'Precursor.Algebra.Monoid.Monoid' laws
-- * @(a '<+>' b) '<+>' c = a '<+>' (b '<+>' c)@
-- * @'zero' '<+>' a = a '<+>' 'zero' = a@
-- * @(a '<.>' b) '<.>' c = a '<.>' (b '<.>' c)@
-- * @'one' '<.>' a = a '<.>' 'one' = a@
--
-- == Commutativity of '<+>'
-- * @a '<+>' b = b '<+>' a@
--
-- == Distribution of '<.>' over '<+>'
-- * @a '<.>' (b '<+>' c) = (a '<.>' b) '<+>' (a '<.>' c)@
-- * @(a '<+>' b) '<.>' c = (a '<.>' c) '<+>' (b '<.>' c)@
--
-- Another useful law, annihilation, may be deduced from the axioms
-- above:
--
-- * @'zero' '<.>' a = a '<.>' 'zero' = 'zero'@
class Semiring a where
  -- | The identity of '<+>'.
  zero :: a
  -- | The identity of '<.>'.
  one :: a
  -- | An associative binary operation, which distributes over '<+>'.
  infixl 7 <.>
  (<.>) :: a -> a -> a
  -- | An associative, commutative binary operation.
  infixl 6 <+>
  (<+>) :: a -> a -> a

  default zero :: Num a => a
  default one :: Num a => a
  default (<+>) :: Num a => a -> a -> a
  default (<.>) :: Num a => a -> a -> a

  zero = 0
  one = 1
  (<+>) = (+)
  (<.>) = (*)

------------------------------------------------------------------------
-- Instances
------------------------------------------------------------------------
instance Semiring Bool where
  one = True
  zero = False
  (<+>) = (||)
  (<.>) = (&&)

instance Semiring () where
  one = ()
  zero = ()
  _ <+> _ = ()
  _ <.> _ = ()

cartProd :: (Ord a, Monoid a) => Set a -> Set a -> Set a
cartProd xs ys =
  Set.foldl' (\a x ->
                Set.foldl' (flip (Set.insert . mappend x)) a ys)
  Set.empty xs

-- | The 'Set' 'Semiring' is 'Data.Set.union' for '<+>', and a Cartesian
-- product for '<.>'.
instance (Ord a, Monoid a) => Semiring (Set a) where
  (<.>) = cartProd
  (<+>) = Set.union
  zero = Set.empty
  one = Set.singleton mempty

------------------------------------------------------------------------
-- Addition and multiplication newtypes
------------------------------------------------------------------------

type WrapBinary f a = (a -> a -> a) -> f a -> f a -> f a

-- | Monoid under '<+>'. Analogous to 'Data.Monoid.Sum', but uses the
-- 'Semiring' constraint, rather than 'Num'.
newtype Add a = Add
  { getAdd :: a
  } deriving (Eq, Ord, Read, Show, Enum, Bounded, Generic, Generic1, Num
             ,Arbitrary)

-- | Monoid under '<.>'. Analogous to 'Data.Monoid.Product', but uses the
-- 'Semiring' constraint, rather than 'Num'.
newtype Mul a = Mul
  { getMul :: a
  } deriving (Eq, Ord, Read, Show, Enum, Bounded, Generic, Generic1, Num
             ,Arbitrary)

add :: (Semiring a, Foldable f) => f a -> a
add = getAdd . foldMap Add

mul :: (Semiring a, Foldable f) => f a -> a
mul = getMul . foldMap Mul

instance Functor Add where fmap = coerce

instance Functor Mul where fmap = coerce

instance Foldable Add where
  foldr   =
    (coerce :: ((a -> b -> c) -> (b -> a -> c))
            -> (a -> b -> c)
            -> (b -> Add a -> c)) flip
  foldl   = coerce
  foldMap = coerce
  length  = const 1

instance Foldable Mul where
  foldr   =
    (coerce :: ((a -> b -> c) -> (b -> a -> c))
            -> (a -> b -> c)
            -> (b -> Mul a -> c)) flip
  foldl   = coerce
  foldMap = coerce
  length  = const 1

instance Applicative Add where
  pure = coerce
  (<*>) =
    (coerce :: ((a -> b) -> a -> b)
            -> (Add (a -> b) -> Add a -> Add b)) ($)

instance Applicative Mul where
  pure = coerce
  (<*>) =
    (coerce :: ((a -> b) -> a -> b)
            -> (Mul (a -> b) -> Mul a -> Mul b)) ($)

instance Monad Add where
  (>>=) = flip coerce

instance Monad Mul where
  (>>=) = flip coerce

instance Semiring a => Monoid (Add a) where
  mempty = Add zero
  mappend = (coerce :: WrapBinary Add a) (<+>)

instance Semiring a => Monoid (Mul a) where
  mempty = Mul one
  mappend = (coerce :: WrapBinary Mul a) (<.>)

instance Semiring a => Semiring (Add a) where
  zero = Add zero
  one = Add one
  (<+>) = (coerce :: WrapBinary Add a) (<+>)
  (<.>) = (coerce :: WrapBinary Add a) (<.>)

instance Semiring a => Semiring (Mul a) where
  zero = Mul zero
  one = Mul one
  (<+>) = (coerce :: WrapBinary Mul a) (<+>)
  (<.>) = (coerce :: WrapBinary Mul a) (<.>)

------------------------------------------------------------------------
-- Ord wrappers
------------------------------------------------------------------------

-- | The 'Semiring' for 'Max' uses the 'max' operation for '<+>', and
-- normal '+' for '<.>'.
instance (Ord a, Bounded a, Semiring a) => Semiring (Max a) where
  (<+>) = mappend
  zero = mempty
  (<.>) = (coerce :: WrapBinary Max a) (<+>)
  one = Max zero

-- | The 'Semiring' for 'Min' uses the 'min' operation for '<+>', and
-- normal '+' for '<.>'.
instance (Ord a, Bounded a, Semiring a) => Semiring (Min a) where
  (<+>) = mappend
  zero = mempty
  (<.>) = (coerce :: WrapBinary Min a) (<+>)
  one = Min zero

------------------------------------------------------------------------
-- (->) instance
------------------------------------------------------------------------

-- | The @(->)@ instance is analogous to the one for 'Monoid'.
instance Semiring b => Semiring (a -> b) where
  zero = const zero
  one = const one
  (f <+> g) x = f x <+> g x
  (f <.> g) x = f x <.> g x

------------------------------------------------------------------------
-- Endo instance
------------------------------------------------------------------------

-- | The 'Endo' semiring uses function composition for '<.>', and
-- pointwise 'mappend' for '<+>'. The underlying 'Monoid' needs to be
-- commutative.
instance Monoid a => Semiring (Endo a) where
  (<.>) = mappend
  one = mempty
  Endo f <+> Endo g = Endo (\x -> f x `mappend` g x)
  zero = Endo (const mempty)

------------------------------------------------------------------------
-- Boring instances
------------------------------------------------------------------------
instance Semiring Float
instance Semiring Double
instance Semiring Int
instance Semiring Int8
instance Semiring Int16
instance Semiring Int32
instance Semiring Int64
instance Semiring Integer
instance Semiring Word
instance Semiring Word8
instance Semiring Word16
instance Semiring Word32
instance Semiring Word64
instance Semiring CUIntMax
instance Semiring CIntMax
instance Semiring CUIntPtr
instance Semiring CIntPtr
instance Semiring CSUSeconds
instance Semiring CUSeconds
instance Semiring CTime
instance Semiring CClock
instance Semiring CSigAtomic
instance Semiring CWchar
instance Semiring CSize
instance Semiring CPtrdiff
instance Semiring CDouble
instance Semiring CFloat
instance Semiring CULLong
instance Semiring CLLong
instance Semiring CULong
instance Semiring CLong
instance Semiring CUInt
instance Semiring CInt
instance Semiring CUShort
instance Semiring CShort
instance Semiring CUChar
instance Semiring CSChar
instance Semiring CChar
instance Semiring IntPtr
instance Semiring WordPtr
instance Semiring Fd
instance Semiring CRLim
instance Semiring CTcflag
instance Semiring CSpeed
instance Semiring CCc
instance Semiring CUid
instance Semiring CNlink
instance Semiring CGid
instance Semiring CSsize
instance Semiring CPid
instance Semiring COff
instance Semiring CMode
instance Semiring CIno
instance Semiring CDev
instance Semiring Natural
instance Integral a => Semiring (Ratio a)
deriving instance Semiring a => Semiring (Product a)
deriving instance Semiring a => Semiring (Sum a)
instance RealFloat a => Semiring (Complex a)
instance HasResolution a => Semiring (Fixed a)
deriving instance Semiring a => Semiring (Identity a)

------------------------------------------------------------------------
-- Very boring instances
------------------------------------------------------------------------
instance (Semiring a, Semiring b) => Semiring (a,b) where
        zero = (zero, zero)
        (a1,b1) <+> (a2,b2) =
                (a1 <+> a2, b1 <+> b2)
        one = (one, one)
        (a1,b1) <.> (a2,b2) =
                (a1 <.> a2, b1 <.> b2)

instance (Semiring a, Semiring b, Semiring c) => Semiring (a,b,c) where
        zero = (zero, zero, zero)
        (a1,b1,c1) <+> (a2,b2,c2) =
                (a1 <+> a2, b1 <+> b2, c1 <+> c2)
        one = (one, one, one)
        (a1,b1,c1) <.> (a2,b2,c2) =
                (a1 <.> a2, b1 <.> b2, c1 <.> c2)

instance (Semiring a, Semiring b, Semiring c, Semiring d) => Semiring (a,b,c,d) where
        zero = (zero, zero, zero, zero)
        (a1,b1,c1,d1) <+> (a2,b2,c2,d2) =
                (a1 <+> a2, b1 <+> b2,
                 c1 <+> c2, d1 <+> d2)
        one = (one, one, one, one)
        (a1,b1,c1,d1) <.> (a2,b2,c2,d2) =
                (a1 <.> a2, b1 <.> b2, c1 <.> c2, d1 <.> d2)

instance (Semiring a, Semiring b, Semiring c, Semiring d, Semiring e) =>
                Semiring (a,b,c,d,e) where
        zero = (zero, zero, zero, zero, zero)
        (a1,b1,c1,d1,e1) <+> (a2,b2,c2,d2,e2) =
                (a1 <+> a2, b1 <+> b2, c1 <+> c2,
                 d1 <+> d2, e1 <+> e2)
        one = (one, one, one, one, one)
        (a1,b1,c1,d1,e1) <.> (a2,b2,c2,d2,e2) =
                (a1 <.> a2, b1 <.> b2, c1 <.> c2,
                 d1 <.> d2, e1 <.> e2)


module Tool where
  import Numeric.LinearAlgebra
  import Data.List
  import System.Random

  {- Convenience functions -}

  -- zero vector of length n
  zeros n = n |> [0,0..]

  -- one vector of length n
  ones n = n |> [1,1..]

  -- random vector of length n
  rands range n = do
    randarr <- (getStdGen >>= (return.randomRs range))
    return $ n |> randarr

  -- zero matrix of (n, m) shape
  zeromat n m = (n><m) [0,0..]

  -- one matrix of (n, m) shape
  onemat n m = (n><m) [1,1..]

  -- random matrix of (n, m) shape
  randmat range n m = do
    randarr <- (getStdGen >>= (return.randomRs range))
    return $ (n><m) randarr

  -- append 1 at the end of a vector
  homo :: Vector Double -> Vector Double
  homo v = vjoin [v, vector [1]]

  -- remove 1 at the end of a vector
  invhomo :: Vector Double -> Vector Double
  invhomo v = subVector 0 (size v - 1) v

  -- Make a matrix by repeating a column vector
  repeatCols :: Int -> Vector Double -> Matrix Double
  repeatCols n v = repmat (asColumn v) 1 n

  -- Given one-hot label, convert it to a index
  stepLabel :: Vector Double -> Int
  stepLabel v = argmax (toList v)

  -- Check if two one-hot labels are equal
  labelEquals :: Vector Double -> Vector Double -> Bool
  labelEquals a b = (argmax $ toList a) == (argmax $ toList b)

  -- Find the index of maximum element in a list
  argmax l = case elemIndex (maximum l) l of Just result -> result

  -- Split a list into half
  split l = [take halfLength l, drop halfLength l]
    where halfLength = length l `quot` 2


module Main where
import Plot
import Graphics.GD
import Data.Complex

imageDim :: Size
imageDim = (600,600)

-- How many iterations before  we declare the point doesn't diverge
maxIter = 127

-- Coordinates in the complex plane to graph
-- First coordinate must be top left.
mbWindow :: Window
mbWindow = ( (-1.9, 1.3), (0.7, -1.3) ) -- Full view

main :: IO ()
main = drawPlot drawMandelbrot imageDim mbWindow "mandelbrot.png"

drawMandelbrot :: Coordinate -> Color
drawMandelbrot (x,y) = colorIterations $ mandelbrot (x :+ y) (0 :+ 0) 0

-- Count number of iterations for the given point to diverge
-- Non-diverging points return 0
mandelbrot :: Complex Double -- Coordinate to test
           -> Complex Double -- Iterating Z value
           -> Int -- Current iteration
           -> Int -- Iterations before diverging
mandelbrot c z iter
    | iter > maxIter = 0
    | otherwise = let z' = z^2 + c in
                  if magnitude z' > 2
                  then iter
                  else mandelbrot c z' (iter+1)

colorIterations :: Int -> Color
colorIterations x
    | x > maxIter = rgb 255 255 255
    | otherwise = let c = x + 30
                  in rgb c x (x^2)


module Main where

import           Numeric.LinearAlgebra
import           System.Environment (getArgs)
import           System.Random
import           System.ProgressBar
import           Control.Monad (replicateM)
import           Control.Monad.IO.Class
import           Data.Maybe
import           Data.List
import qualified Data.Vector as V
import           Polynomials
import           Genetic
import           Formulas

main :: IO ()
main = do
    (nGenerations:nPop:goal:nElit:_) <- getArgs

    ------------ Information ------------
    putStrLn ""
    putStrLn $ "No. of generatoins: "  ++ nGenerations
    putStrLn $ "Population Size: "     ++ nPop
    putStrLn $ "Score of goal: "       ++ goal
    putStrLn $ "Elite size (N-elit): " ++ nElit
    putStrLn ""

    ------------ Progress Bar -----------
    putStrLn "Progress:"
    (pb, asy) <- startProgress (msg "") (msg "") 10 (Progress 0 (read nGenerations))

    ------------ Algorithm --------------
    powers <- randomPowerGen 6 2
    let
        ------ Inps and outps of function ------
        inps  = inpsRosen
        outps = fromList $ V.toList yRosen
        ----------------------------------------
        individual = Article powers inps outps
        init = replicate (read nPop) individual

    algo <- runGeneticAlgorithm (read goal) (read nGenerations) (read nElit) init pb
    let
        finalPowers = _articlePowers algo
        x = calcX inps finalPowers
        y = outps
        coeffs = normEq x y
        ----------------------------------------
    grade <- show . fromJust <$> fitness algo

    putStrLn ""
    putStrLn $ showPolynomial coeffs finalPowers
    putStrLn $ "Fitness of final Polynomial is: " ++ grade
    putStrLn ""
    putStrLn "Attention: component with coefficient smaller tan 0.001 are rounded to nothing"

openVLists = V.map V.toList

-- Check x + 1:
simple x = x + 1.0

inpsSimple :: V.Vector Double
inpsSimple = V.fromList
       [ 1
       , 2
       , 3
       , 4]

powsSimple :: V.Vector Double
powsSimple = V.fromList
            [1.0,
             0]

ySimple = V.map simple inpsSimple

-- Check x**2 + y**2 -1:
circle [x, y] = x**2 + y**2 --  -1

inpsCircle :: Inplist
inpsCircle = V.fromList
     . map V.fromList
     $ [ [0.0001, 1.1]
       , [2, 0.001]
       , [2, 1.05]
       , [1, 2.11]
       , [0.1, 2]]

powsCircle :: Powers
powsCircle = V.fromList
          . map V.fromList
          $ [[2, 0],
             [0, 2],
             [0, 0]]

yCircle = V.map circle $ openVLists inpsCircle

-- Check Rosenbrock:
rosen [x, y] =
    1 - 2*x + x**2 + 100*y**2 + 200*x**2*y + 100*x**4

inpsRosen :: Inplist
inpsRosen = V.fromList
     . map V.fromList
     $ [[i, j] | i <- [-0.9, -0.8 .. -0.2], j <- [0.2, 0.3 .. 0.9]]

powsRosen :: Powers
powsRosen = V.fromList
        . map V.fromList
        $ [[0, 0],
           [1, 0],
           [2, 0],
           [0, 2],
           [2, 1],
           [4, 0]]

yRosen = V.map rosen $ openVLists inpsRosen


{-# LANGUAGE TupleSections #-}
{-# LANGUAGE MultiWayIf #-}

module Isumi.Math.FFT.Internal.Bruun
  ( fftDirect
  , cruN
  , fftBruun
  , fftBruunC
  )
where

import           Data.Bits                          (shiftL)
import           Data.Complex
import           Data.HashMap.Strict                (HashMap)
import qualified Data.HashMap.Strict                as HashMap
import           Data.List                          (concatMap)
import           Data.Maybe
import qualified Data.Vector                        as V
import qualified Data.Vector.Unboxed                as UV
import           Math.NumberTheory.Logarithms       (integerLog2)

import           Isumi.Math.FFT.Internal.Polynomial
import           Isumi.Math.FFT.Internal.Utils

-- | Computes fft using definition.
fftDirect :: UV.Vector (Complex Double) -> UV.Vector (Complex Double)
fftDirect xs = UV.generate len genEntry
 where
  len = UV.length xs
  genEntry k = UV.sum . flip UV.imap xs $ \n xn ->
    (crus UV.! k)^n * xn
  crus = cruN len

{-|
  Computes the complex roots of unity
-}
cruN :: Int -> UV.Vector (Complex Double)
cruN n = UV.generate n $ \k ->
  exp (realToComplex k * (-2) * pi * compI / fromIntegral n)

-- | FFT with Bruun's algorithm, for real input
fftBruun :: UV.Vector Double -> Maybe (UV.Vector (Complex Double))
fftBruun = fftBruunG RealPolynomial

-- | FFT with Bruun's alogrithm, for complex input
fftBruunC :: UV.Vector (Complex Double) -> Maybe (UV.Vector (Complex Double))
fftBruunC = fftBruunG ComplexPolynomial

fftBruunG :: (UV.Unbox a, Floating a, Eq a)
          => (Polynomial a -> RealOrComplexPolynomial Double)
          -> UV.Vector a
          -> Maybe (UV.Vector (Complex Double))
fftBruunG f xs =
  if | (not . isPowerOf2) n -> Nothing
     | UV.all (== 0) xs -> Just $ UV.replicate n 0
     | otherwise ->
       let r0 = f . fromJust . polyFromVecRep . UV.reverse $ xs
           d0 = BruunDivisorMinus n
        in Just $ fftBruun' (cruN n) (d0, r0)
  where
  n = UV.length xs

fftBruun' :: UV.Vector (Complex Double) -- ^cru array
          -> DivisorRemainder -- ^initial divisor and remainder
          -> UV.Vector (Complex Double)
fftBruun' cruValues dr =
  UV.generate n (\i -> cruToFFT HashMap.! (cruValues' V.! i))
  where
  n = UV.length cruValues
  cruToFFT :: HashMap (Integer, Integer) (Complex Double)
  cruToFFT = HashMap.fromList
           . fmap (\(d, r) -> (normalizeCRU . finalDivisorToCRU $ d,
               finalRemainderToValue r))
           . factorDRAll $ dr
  cruValues' = fmap normalizeCRU (UV.convert cruValues)
    :: V.Vector (Integer, Integer)

normalizeCRU :: Complex Double -> (Integer, Integer)
normalizeCRU (r :+ i) = (toInt r, toInt i)
  where
  toInt x = round (x * 1000000)

finalDivisorToCRU :: BruunDivisor Double -> Complex Double
finalDivisorToCRU (BruunDivisorComplex c) = (-1) * c
finalDivisorToCRU (BruunDivisorMinus n)   = if n == 1 then 1 else error ""
finalDivisorToCRU (BruunDivisorPlus n c)  =
  if n == 1 && c == 0
     then (-1)
     else error ""

finalRemainderToValue :: RealOrComplexPolynomial Double -> Complex Double
finalRemainderToValue (ComplexPolynomial p) =
  case p of
    PolyVecRep values ->
      if UV.length values == 1
         then values UV.! 0
         else error $ "Final stage remainder MUST have a length of 1. Actual: "
                      ++ show (UV.length values)
    PolyExpCoefRep _ ->
      error "Final stage remainder MUST have exp-coef representation"
finalRemainderToValue (RealPolynomial p) =
  finalRemainderToValue . ComplexPolynomial . mapPoly realToComplex $ p

-- | Recursively divisor-remainder pairs until they can not be factored
factorDRAll :: DivisorRemainder -> [DivisorRemainder]
factorDRAll dr = go [dr]
  where
  go :: [DivisorRemainder] -> [DivisorRemainder]
  go xs = let ys = concatMap (extractListFromMaybe . factorDR) xs
           in if length ys /= 2 * length xs then xs else go ys

type DivisorRemainder = (BruunDivisor Double, RealOrComplexPolynomial Double)

-- | Factor a divisor-remainder pair into TWO sub divisor remainder pairs.
factorDR :: DivisorRemainder -> Maybe [DivisorRemainder]
factorDR (d, r) =
  if divDegree d == 1
     then Nothing
     else fmap (fromJust . flip updateDR r) . pairToList <$> factorDivisor d

updateDR :: BruunDivisor Double -- ^new divisor
         -> RealOrComplexPolynomial Double -- ^old remainder
         -> Maybe DivisorRemainder -- ^(new divisor, new remainder)
updateDR d r = (d,) <$> r `modPolyG` divToPoly d

pairToList :: (a, a) -> [a]
pairToList (x, y) = [x, y]

isPowerOf2 :: Int -> Bool
isPowerOf2 x = (1 `shiftL` integerLog2 (fromIntegral x)) == x

extractListFromMaybe :: Maybe [a] -> [a]
extractListFromMaybe Nothing   = []
extractListFromMaybe (Just xs) = xs


module Mandelbruh
    ( C,
      mandelbruh,
      inMandelbrot,
      color,
      series
    ) where

import Data.Complex

type C = Complex Float

f :: C -> C -> C
f c z = z^2 + c

type Series = [C]

series :: C -> Series
series c = iterate (f c) c

inMandelbrot :: C -> Bool
inMandelbrot =
  let ceiling = 100 -- TODO: move to config file
      samples = 20
      in all (< ceiling) . map magnitude . take samples . series

mandelbruh = 42

readComplex :: String -> C
readComplex s = read s :: C

type Color = Char
type Palette = [Color]

color :: Palette -> Series -> Color
color palette = (palette !!) . length . take n . takeWhile fairlyClose
  where n = length palette - 1
        fairlyClose c = magnitude c < ceiling
        ceiling = 100

type Image color = C -> color


import qualified Prelude as P
import Control.Constrained.Prelude

import Control.Constrained.Applicative
import Control.Constrained.Category
import Control.Constrained.Comonad
import Control.Constrained.Foldable
import Control.Constrained.Functor
import Control.Constrained.Plain
import qualified Control.Monad as P
import Data.Aeson (genericToEncoding, defaultOptions)
import Data.Complex
import qualified Data.Monoid as M
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Monoid.Instances ()
import qualified Data.Vector.Unboxed as U
import Data.Text (Text)
import Data.Yaml
import Formatting
import System.Directory

sum :: Foldable f => Dom f ~ (-#>) => Num a => PCon a => PCon (M.Sum a)
    => f a -> a
sum xs = M.getSum (foldMap (PFun M.Sum) xs)

sum2 :: Foldable f => Dom f ~ (-#>) => Num a => PCon a => PCon (M.Sum a)
     => f a -> a
sum2 xs = M.getSum (foldMap (PFun M.Sum . PFun \x -> x^(2::Int)) xs)

integral :: PCon a => Floating a => UIVector a -> a
integral xs = sum xs / fromIntegral (length xs)

norm2 :: PCon a => Floating a => UIVector a -> a
norm2 xs = rms (sum2 xs) (length xs)
  where rms :: Floating a => Integral i => a -> i -> a
        rms s2 c = sqrt (s2 / fromIntegral c)

-- physical boundaries
xmin :: Fractional a => a
xmin = 0
xmax :: Fractional a => a
xmax = 1

-- discretization
npts :: Int
npts = 100
hgsp :: Fractional a => a
hgsp = (xmax - xmin) / fromIntegral npts

cfl :: Fractional a => a
cfl = 0.5
nsteps :: Int
nsteps = 200

coord :: Fractional a => Int -> a
coord i = xavg + xrad / irad * (fromIntegral i - iavg)
  where -- discrete boundaries
        imin = -1/2
        imax = fromIntegral npts + 1/2
        -- auxiliary functions
        iavg = (imin + imax) / 2
        irad = (imax - imin) / 2
        xavg = (xmin + xmax) / 2
        xrad = (xmax - xmin) / 2

-- TODO: implement 'ToList UIVector', use it
coords :: PCon a => Fractional a => UIVector a
coords = uivector . U.fromListN npts . fmap coord $ [0..npts-1]

setup :: Floating a => a -#> Complex a
setup = PFun \x -> sin (2*pi*x) :+ 0

setups :: Functor f => Dom f ~ (-#>) => Cod f ~ (->)
       => PCon a => Floating a
       => f a -> f (Complex a)
setups = fmap setup

deriv :: ComonadStore f
      => Cod f ~ (->) => Eq (Index f) => Num (Index f)
      => Ok (Dom f) a => Fractional a
      => f a -> a
deriv xs = if | n == 0 -> 0
              | i == 0 -> (peek (i+1) xs - peek i xs) / hgsp
              | i == n-1 -> (peek i xs - peek (i-1) xs) / hgsp
              | otherwise -> (peek (i+1) xs - peek (i-1) xs) / (2 * hgsp)
  where i = pos xs
        n = fromIntegral npts

derivs :: ComonadStore f
       => Cod f ~ (->) => Eq (Index f) => Num (Index f)
       => Ok (Dom f) a => Fractional a
       => f a -> f a
derivs = extend deriv

energy :: ComonadStore f
       => Cod f ~ (->) => Eq (Index f) => Num (Index f)
       => Ok (Dom f) (Complex a) => RealFloat a
       => f (Complex a) -> a
energy us = (realPart du ^^ (2::Int) + imagPart du ^^ (2::Int)) / 2
   where du = deriv us

energies :: ComonadStore f
         => Cod f ~ (->) => Eq (Index f) => Num (Index f)
         => Ok (Dom f) a => Ok (Dom f) (Complex a) => RealFloat a
         => f (Complex a) -> f a
energies = extend energy

rhs :: ComonadStore f
    => Cod f ~ (->) => Eq (Index f) => Num (Index f)
    => Ok (Dom f) (Complex a) => RealFloat a
    => f (Complex a) -> Complex a
rhs us = du
   where du = deriv us

rhss :: ComonadStore f
     => Cod f ~ (->) => Eq (Index f) => Num (Index f)
     => Ok (Dom f) (Complex a) => RealFloat a
     => f (Complex a) -> f (Complex a)
rhss = extend rhs

rk2 :: Apply f
    => Dom f ~ (-#>) => Cod f ~ (->)
    => Ok (Dom f) a => Fractional a
    => (f a -> f a) -> a -> f a -> f a
rk2 f dx y0 = let k0 = f y0
                  y1 = liftA2u (PFun \(y, k) -> y + 0.5 * dx * k) y0 k0
                  k1 = f y1
                  y2 = liftA2u (PFun \(y, k) -> y + dx * k) y0 k1
              in y2

outputFilePath :: Int -> FilePath
outputFilePath i =
  formatToString ("output/wavetoy.it" % (left 6 '0' %. int) % ".yaml") i

instance (ToJSON a, PCon a) => ToJSON (Complex a) where
  toEncoding = genericToEncoding defaultOptions

instance (ToJSON a, PCon a) => ToJSON (UIVector a) where
  toJSON = toJSON . getUIVector
  toEncoding = toEncoding . getUIVector

data JSONData where
  JD :: ToJSON a => a -> JSONData

instance ToJSON JSONData where
  toJSON (JD x) = toJSON x
  toEncoding (JD x) = toEncoding x

output :: FilePath -> Map Text JSONData -> IO ()
output = encodeFile

main:: IO ()
main = do putStrLn "Constrained functors: Numeric functions"
          let x = coords @Double
          putStrLn $ "L2[x] = " ++ show (norm2 x)
          let u0 = setups x
          putStrLn $ "L2[u0] = " ++ show (realPart (norm2 u0))
          let e0 = energies u0
          putStrLn $ "E0 = " ++ show (integral e0)
          -- let du = rhss u0
          -- putStrLn $ "L2[du] = " ++ show (realPart (norm2 du))
          -- let u1 = rk2 rhss (cfl * hgsp) u0
          -- putStrLn $ "L2[u1] = " ++ show (realPart (norm2 u1))
          let step :: PCon a => RealFloat a
                   => UIVector (Complex a) -> UIVector (Complex a)
              step = rk2 rhss (cfl * hgsp)
          let uis = take (nsteps + 1) (P.iterate step u0)
          createDirectoryIfMissing True "output"
          P.forM (zip [0..] uis) \(i, u) ->
            do putStrLn $ "Iteration " ++ show i
               putStrLn $ "  L2[u] = " ++ show (realPart (norm2 u))
               let e = energies u
               putStrLn $ "  E = " ++ show (integral e)
               let vars = Map.fromList [ ("iteration", JD i)
                                       , ("x", JD x)
                                       , ("u", JD u0)
                                       , ("e", JD e0)
                                       ]
               P.when (i `mod` 10 == 0) $
                 output (outputFilePath i) vars
          putStrLn "Done."


{-# LANGUAGE BangPatterns, CPP #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}

module Mandelbrot(mandelBench) where
{- Code taken from: https://github.com/simonmar/monad-par/blob/master/examples/src/mandel/mandel.hs-}

import Control.Monad
import Control.DeepSeq
import Control.Exception
import Data.Complex
import System.Environment
import System.IO
#ifdef ALIST
import Control.Monad.Par.AList as A
#endif

#ifdef PARSCHED
import PARSCHED
#else
import Control.Monad.Par
#endif

import qualified Data.Vector.Unboxed as V

#ifdef WRITE_IMAGE
import Codec.Picture  -- JuicyPixels
import qualified Data.Vector.Storable as V
#endif

#ifdef NEW_GENERIC
import qualified Data.Par as C
import qualified Data.Par.Range as C
import Control.Par.Class
parMapReduceRangeThresh :: (NFData a, ParFuture p, FutContents p a)
                        => Int -> C.InclusiveRange -> (Int -> p a) -> (a -> a -> p a) -> a -> p a
parMapReduceRangeThresh = C.parMapReduceThresh
#else
import Control.Monad.Par.Combinator as C
#endif

import Control.Monad.SD as Ohua
import Criterion
import Criterion.Main
import Control.Monad.Identity (runIdentity)

mandel :: Int -> Complex Double -> Int
mandel max_depth c = loop 0 0
  where
   fn = magnitude
   loop i !z
    | i == max_depth = i
    | fn(z) >= 2.0   = i
    | otherwise      = loop (i+1) (z*z + c)

threshold = 1

data VecTree = Leaf (V.Vector Int)
	     | MkNode VecTree VecTree
  deriving (Eq,Show)

#if 0
instance V.Unbox a => NFData (V.Vector a) where
  rnf v = rnf (V.length v)
#endif

instance NFData VecTree where
  rnf (Leaf v) = rnf v
  rnf (MkNode v1 v2) = rnf v1 `seq` rnf v2

leftmost (Leaf v)     = v
leftmost (MkNode l _) = leftmost l

#ifdef ALIST
runMandel :: Double -> Double -> Double -> Double -> Int -> Int -> Int -> Par (AList [Int])
runMandel minX minY maxX maxY winX winY max_depth = do
  -- This version does a ROW at a time in parallel:
  A.parBuildThreshM threshold (C.InclusiveRange 0 (winY-1)) $ \y ->
       do let l = [ mandelStep y x | x <- [0.. winX-1] ]
          deepseq l (return l)

#else
-- RRN: Comment to handle LVish types too:
-- runMandel :: Double -> Double -> Double -> Double -> Int -> Int -> Int -> Par VecTree
runMandel minX minY maxX maxY winX winY max_depth runner = do
  -- Auto-partitioning version.  A bit worse:
  -- C.parMapReduceRange (C.InclusiveRange 0 (winY-1))
  runner threshold (C.InclusiveRange 0 (winY-1))
     (\y ->
       do
          let vec = V.generate winX (\x -> mandelStep y x)
          seq (vec V.! 0) $
           return (Leaf vec))
     (\ a b -> return$ MkNode a b)
     (Leaf V.empty)
#endif
  where
    mandelStep i j = mandel max_depth z
        where z = ((fromIntegral j * r_scale) / fromIntegral winY + minY) :+
                  ((fromIntegral i * c_scale) / fromIntegral winX + minX)
    r_scale  =  maxY - minY  :: Double
    c_scale =   maxX - minX  :: Double

-- pure version of the above
runMandelP minX minY maxX maxY winX winY max_depth runner =
  -- Auto-partitioning version.  A bit worse:
  -- C.parMapReduceRange (C.InclusiveRange 0 (winY-1))
  runner threshold (C.InclusiveRange 0 (winY-1))
     (\y ->
        let vec = V.generate winX (\x -> mandelStep y x)
        in seq (vec V.! 0) $ (Leaf vec))
     (\ a b ->  MkNode a b)
     (Leaf V.empty)
  where
    mandelStep i j = mandel max_depth z
        where z = ((fromIntegral j * r_scale) / fromIntegral winY + minY) :+
                  ((fromIntegral i * c_scale) / fromIntegral winX + minX)
    r_scale  =  maxY - minY  :: Double
    c_scale =   maxX - minX  :: Double

-- runMandelP_ minX minY maxX maxY winX winY max_depth runner =
--   -- Auto-partitioning version.  A bit worse:
--   -- C.parMapReduceRange (C.InclusiveRange 0 (winY-1))
--   runner threshold (C.InclusiveRange 0 (winY-1))
--      (\y ->
--         let vec = V.generate winX (\x -> mandelStep y x)
--         in seq (vec V.! 0) $ (Leaf vec))
--      (\ a b ->  MkNode a b)
--      (Leaf V.empty)
--   where
--     mandelStep i j = mandel max_depth z
--         where z = ((fromIntegral j * r_scale) / fromIntegral winY + minY) :+
--                   ((fromIntegral i * c_scale) / fromIntegral winX + minX)
--     r_scale  =  maxY - minY  :: Double
--     c_scale =   maxX - minX  :: Double

#ifdef WRITE_IMAGE
-- makeImage :: Integer -> Integer -> Int -> A.AList [Int] -> Image PixelRGB8
makeImage :: Int -> Int -> Int -> A.AList [Int] -> Image PixelRGB8
makeImage x y depth ls =
  -- This is a quite silly series of intermediate structures:
  Image x y (V.fromList $ concat $
	     Prelude.map prettyRGB $
	     concat$ A.toList ls)
 where
   prettyRGB s =
      let t = (fromIntegral (depth - s)) in
--      PixelRGB8 (fromIntegral s) t t
      [fromIntegral s, t, t]
#endif

simple x y depth = runMandel (-2) (-2) 2 2 x y depth
simpleP :: Monad m => Int -> Int -> Int -> (
    Int -- ^ threshold
    -> InclusiveRange -- ^ range over which to calculate
    -> (Int -> VecTree) -- ^ compute one result
    -> (VecTree -> VecTree -> VecTree) -- ^ combine two results (associative)
    -> VecTree -- ^ initial result
    -> m VecTree) -> m VecTree
simpleP x y depth = runMandelP (-2) (-2) 2 2 x y depth

-- simpleP_ :: Int -> Int -> Int -> (Int -- ^ threshold
--     -> InclusiveRange -- ^ range over which to calculate
--     -> (Int -> VecTree) -- ^ compute one result
--     -> (VecTree -> VecTree -> VecTree) -- ^ combine two results (associative)
--     -> VecTree -- ^ initial result
--     -> IO VecTree) -> IO VecTree
-- simpleP_ x y depth = runMandelP_ (-2) (-2) 2 2 x y depth

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

#ifdef ALIST
-- A meaningless checksum.  This match the C++ CnC benchmark:
mandelCheck :: AList [Int] -> Int -> Int -> Int
mandelCheck als max_col max_depth = loop 0 als 0
 where
 -- Sum positions of the pixels that are equal to max_depth:
 loop i als !sum | A.null als = sum
 loop i als !sum = loop (i+1) (A.tail als)
		        (loop2 i 0 (A.head als) sum)
 loop2 i j []    !sum = sum
 loop2 i j (h:t) !sum | h == max_depth = loop2 i (j+1) t (sum + i*max_col + j)
		      | otherwise      = loop2 i (j+1) t  sum
#endif

 -- This kind of checksum is much simpler:
checkSum :: VecTree -> Int
checkSum (Leaf vec) = V.foldl (+) 0 vec
checkSum (MkNode v1 v2) = checkSum v1 + checkSum v2

-- main = do args <- getArgs
--
--           let (x,y,depth) =
-- 		case args of
-- 		 [] ->
--     --                 runPar $ simple 3 3 3
--                        -- Defaults recommended by Simon Marlow
-- 		       -- (1024,1024,256)
--                        -- Note, this should do something very *quick* when called with no args:
--                        (3,3,3)
--
-- 		 [x,y,depth] ->
--     --		   simple (read x) (read y) (read depth)
-- 		       (read x, read y, read depth)
--
-- 		 -- [minX,minY,maxX,maxY,winX,winY,depth] ->
-- 		 --    runPar $
-- 		 -- 	    runMandel (read minX) (read minY)
-- 		 -- 		      (read maxX) (read maxY)
-- 		 -- 		      (read winX) (read winY) (read depth)
--
--           let pixels = runPar$ simple x y depth parMapReduceRangeThresh
--
--           putStrLn$ "About to poke pixels to launch runPar:"
--           evaluate pixels
--           putStrLn$ "Should be done with runpar at this point."
--
-- #ifdef WRITE_IMAGE
-- 	  writePng "mandel_image.png" (makeImage (fromIntegral x) (fromIntegral y) depth pixels)
-- 	  putStrLn$ "File written."
--           putStrLn$ "Checksum " ++ show (checkSumpixels y depth)
-- #endif
--           putStrLn$ "Spot check, length of leftmost leaf: " ++ show (V.length (leftmost pixels))
--           return ()

parMapReduceRangeThreshSeq:: Int -- ^ threshold
  -> InclusiveRange -- ^ range over which to calculate
  -> (Int -> VecTree) -- ^ compute one result
  -> (VecTree -> VecTree -> VecTree) -- ^ combine two results (associative)
  -> VecTree -- ^ initial result
  -> IO VecTree
parMapReduceRangeThreshSeq t (InclusiveRange mi ma) fn binop init =
  return $ loop mi ma
  where
    loop mi ma
      | ma - mi <= t =
        let mapred a b = let x = fn b
                             result = a `binop` x
                          in result
         in foldl mapred init [mi .. ma]
      | otherwise =
        let mid = mi + ((ma - mi) `quot` 2)
            r = loop (mid + 1) ma
            l = loop mi mid
         in l `binop` r

mandelBench =
  bgroup
    "mandel-bench"
    [ bench "sequential" (nfIO (simpleP x y depth parMapReduceRangeThreshSeq))
    , bench "ohua" (nfIO $ simpleP x y depth Ohua.mapReduceRangeThresh)
    , bench "par" (nf (runPar . (simple x y depth)) C.parMapReduceRangeThresh)
    ]
  where
    {- Values taken from: https://github.com/simonmar/monad-par/blob/master/examples/src/run_benchmark.hs-}
    -- desktop values
    -- x = 1024
    -- y = 1024
    -- depth = 256
    -- server values
    x = 1024
    y = 1024
    depth = 256
{-
[2012.03.08] {Updated to use Vector}
Looks like this was out of date in that it was still using AList.
AList has never performed well enough, and sure enough, this test was
falling to ~50% productivity at >16 threads.
However, simply replacing AList with vector is no good (14.4s for 512
512 2048 as opposed to 15.5).  It is of course a quadratic copy at
that point.
-}


{-# LANGUAGE RecordWildCards #-}

module NeuralNetwork.BackForwardPropagation
  ( genRandomThetas
  , computeCost
  , minimizeBFGS2
  ) where

import Control.Monad (mapM)
import Numeric.GSL.Minimization (MinimizeMethodD(VectorBFGS2), minimizeVD)
import Numeric.LinearAlgebra (( #> ), (<.>), (<>), sumElements)
import Numeric.LinearAlgebra.Data
  ( Matrix
  , R
  , Vector
  , (===)
  , (?)
  , asColumn
  , atIndex
  , cmap
  , flatten
  , fromList
  , fromLists
  , matrix
  , reshape
  , size
  , takesV
  , tr'
  , vjoin
  )
import Numeric.LinearAlgebra.Devel (mapMatrixWithIndex)
import Numeric.Natural (Natural)
import Prelude hiding ((<>))
import System.Random (randomRIO)

import Common
  ( MinimizationOpts(..)
  , addOnesColumn
  , getDimensions
  , sigmoid
  , sigmoidMatrix
  )
import NeuralNetwork.Types (Layer(numUnits), Network(..))

genRandomThetas :: R -> Network -> IO [Matrix R]
genRandomThetas epsilon Network {..} =
  let allLayers = inputLayer : hiddenLayer ++ [outputLayer]
   in mapM mkTheta (zip allLayers (tail allLayers))
  where
    mkTheta :: (Layer, Layer) -> IO (Matrix R)
    mkTheta (l1, l2) = do
      let numRows = numUnits l2
          numCols = numUnits l1 + 1
          numElts = fromIntegral $ numRows * numCols
      elts <- sequence $ replicate numElts (randomRIO (0.0, 1.0))
      return $
        matrix
          (fromIntegral numCols)
          ((\r -> r * (2 * epsilon) - epsilon) <$> elts)

computeCost :: Int -> Matrix R -> Vector R -> [Matrix R] -> R
computeCost numLabels x y thetas = cost
  where
    m = size y
    --
    hypotheses =
      foldl (\an th -> sigmoidMatrix ((addOnesColumn an) <> (tr' th))) x thetas
    --
    cost =
      let rowElemsCost i k r =
            let yik =
                  if (round $ y `atIndex` i) == k
                    then 1.0
                    else 0.0
             in (-yik * (log (r))) - ((1 - yik) * log (1 - r))
          summation =
            sumElements $
            mapMatrixWithIndex (\(i, k) r -> rowElemsCost i k r) hypotheses
       in summation / fromIntegral m

minimizeBFGS2 ::
     Matrix R
  -> Vector R
  -> Int
  -> [Matrix R]
  -> Int
  -> MinimizationOpts
  -> Int
  -> [Matrix R]
minimizeBFGS2 x y numLabels thetas numIters MinimizationOpts {..} regFactor =
  let (flattenedParams, dimensions) = flattenParameters thetas
      (minParams, _) =
        minimizeVD
          VectorBFGS2
          precision
          numIters
          sizeOfFirstTrialStep
          tolerance
          (computeCost' dimensions)
          (computeGradients' dimensions)
          flattenedParams
   in unflattenParameters dimensions minParams
  where
    m = size y
    --
    flattenParameters :: [Matrix R] -> (Vector R, [(Int, Int)])
    flattenParameters mtxs =
      let (vecs, dims) =
            unzip $ (\mx -> (flatten mx, getDimensions mx)) <$> mtxs
       in (vjoin vecs, dims)
    --
    unflattenParameters :: [(Int, Int)] -> Vector R -> [Matrix R]
    unflattenParameters dims params =
      let sizes = uncurry (*) <$> dims
       in zipWith
            (\vec (ncols, _) -> reshape ncols vec)
            (takesV sizes params)
            dims
    --
    computeCost' :: [(Int, Int)] -> Vector R -> R
    computeCost' dims params =
      computeCost numLabels x y (unflattenParameters dims params)
    --
    computeGradients' :: [(Int, Int)] -> Vector R -> Vector R
    computeGradients' dims gradients =
      fst $
      flattenParameters $
      computeGradients 0 (unflattenParameters dims gradients)
    --
    computeActivationVals :: Matrix R -> [Matrix R]
    computeActivationVals a0 =
      let f vals@[] th = sigmoidMatrix (th * (addBiasUnit a0)) : vals
          f vals@(an:_) th = sigmoidMatrix (th * (addBiasUnit an)) : vals
       in foldl f [] thetas
    --
    computeErrors :: Matrix R -> [Matrix R] -> [Matrix R]
    computeErrors dL activationVals =
      let f vals@(dl:_) (an, th) = ((tr' th) * dl) .* an .* (1 - an) : vals
       in foldl f [dL] (zip activationVals (reverse thetas))
    --
    computeGradients :: Int -> [Matrix R] -> [Matrix R]
    computeGradients i gradients
      | i >= m = gradients
      | otherwise =
        let xi = (asColumn . flatten) $ x ? [i]
            a0 = xi
            activationVals@(aL:_) = computeActivationVals a0
            yi = round $ y `atIndex` i
            binEq :: Int -> R
            binEq n =
              if yi == n
                then 1.0
                else 0.0
            yik = asColumn $ fromList (binEq <$> [0 .. (numLabels - 1)])
            errorVals = computeErrors (aL - yik) (tail activationVals)
            accumGradient grad errVal actVal = grad + (errVal * (tr' actVal))
         in computeGradients (i + 1) $
            zipWith3
              accumGradient
              gradients
              errorVals
              (a0 : reverse activationVals)

elemWiseMult :: Matrix R -> Matrix R -> Matrix R
elemWiseMult mx1 mx2 =
  mapMatrixWithIndex (\idx r -> (mx1 `atIndex` idx) * r) mx2

(.*) :: Matrix R -> Matrix R -> Matrix R
(.*) = elemWiseMult

infixr 8 .*

-- Helpers
addBiasUnit :: Matrix R -> Matrix R
addBiasUnit mx = fromLists [[1]] === mx

addBiasUnit' :: Vector R -> Vector R
addBiasUnit' vec = vjoin [fromList [1], vec]


-----------------------------------------------------------------------------
-- |
-- Module      :  DSP.Matrix.Matrix
-- Copyright   :  (c) Matthew Donadio 2003
-- License     :  GPL
--
-- Maintainer  :  m.p.donadio@ieee.org
-- Stability   :  experimental
-- Portability :  portable
--
-- Basic matrix routines
--
-----------------------------------------------------------------------------

module Matrix.Matrix where

import Data.Array
import Data.Complex

-- | Matrix-matrix multiplication: A x B = C

mm_mult :: (Ix a, Integral a, Num b) => Array (a,a) b -- ^ A
	-> Array (a,a) b -- ^ B
	-> Array (a,a) b -- ^ C

mm_mult a b = if ac /= br
	      then error "mm_mult: inside dimensions inconsistent"
	      else array bnds [ ((i,j), mult i j) | (i,j) <- range bnds ]
    where mult i j = sum [ a!(i,k) * b!(k,j) | k <- [1..ac] ]
	  ((_,_),(ar,ac)) = bounds a
	  ((_,_),(br,bc)) = bounds b
	  bnds = ((1,1),(ar,bc))

-- | Matrix-vector multiplication: A x b = c

mv_mult :: (Ix a, Integral a, Num b) => Array (a,a) b -- ^ A
	-> Array a b -- ^ b
	-> Array a b -- ^ c

mv_mult a b = if ac /= br
	      then error "mv_mult: dimensions inconsistent"
	      else array bnds [ (i, mult i) | i <- range bnds ]
    where mult i = sum [ a!(i,k) * b!(k) | k <- [1..ac] ]
	  ((_,_),(ar,ac)) = bounds a
	  (_,br) = bounds b
	  bnds = (1,ar)

-- | Transpose of a matrix

m_trans :: (Ix a, Integral a, Num b) => Array (a,a) b -- ^ A
	-> Array (a,a) b -- ^ A^T

m_trans a = array bnds [ ((i,j), a!(j,i)) | (i,j) <- range bnds ]
    where (_,(m,n)) = bounds a
	  bnds = ((1,1),(n,m))

-- | Hermitian transpose (conjugate transpose) of a matrix

m_hermit :: (Ix a, Integral a, RealFloat b) => Array (a,a) (Complex b) -- ^ A
	 -> Array (a,a) (Complex b) -- ^ A^H

m_hermit a = array bnds [ ((i,j), conjugate (a!(j,i))) | (i,j) <- range bnds ]
    where (_,(m,n)) = bounds a
	  bnds = ((1,1),(n,m))


{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE RankNTypes #-}
module Numeric.Tensor where

import qualified Prelude
import Control.Category
import Data.Kind(Type)
import Dependent.Size
import Data.Array.Accelerate((+),(*))
import qualified Data.Array.Accelerate as Acc
import Numeric.Vector.Sized

class Semigroup g where
  (<>) :: g -> g -> g

-- Associativity:
-- (a <> b) <> c = a <> (b <> c)

class Semigroup m => Monoid m where
  zero :: m

-- Identity:
-- a <> zero = a = zero <> a

class Monoid g => Group g where
    negate :: g -> g

-- Additive inverse
-- a <> negate a = zero = negate a <> a

-- Multiplication
class Group r => Ring r where
  one :: r
  (.*.) :: r -> r -> r
-- one .*. x = x = x .*. one

class Ring k => Field k where
  recip :: k -> k

-- x .*. recip x = one = recip x .*. x

class (Field k, Group v) => Module k v where
  (*.) :: k -> v -> v
  (.*) :: v -> k -> v

newtype Vect k v w = Vect {linear :: v -> w}

--
-- instance ShapeSize size => Semigroup (Sized size) where
--   -- (<>) :: Sized size -> Sized size -> Sized size
--   x <> y = zipWith (+) x y
--
-- instance ShapeSize size => Monoid (Sized size) where
--   -- zero :: Sized size
--   zero = generate (\_shape -> 0)
--
-- instance ShapeSize size => Group (Sized size) where
--   negate v = map Acc.negate v
--
instance Semigroup (Acc.Exp Acc.Double) where
  a <> b = a + b

instance Monoid (Acc.Exp Acc.Double) where
  zero = 0

instance Group (Acc.Exp Acc.Double) where
  negate = Acc.negate

instance Ring (Acc.Exp Acc.Double) where
  one = 1
  a .*. b = a * b

instance Field (Acc.Exp Acc.Double) where
  recip x = Acc.recip x

-- instance (ShapeSize size) => Module (Acc.Exp Acc.Double) (Sized size) where
--    r *. v = map (* r) v
--    v .* r = map (* r) v

instance Field k => Category (Vect k) where
  id = Vect id
  Vect f . Vect g = Vect (f . g)

class (Field f, Group v)  => VectorSpace f v where
  distribute :: f -> v -> v

-- instance ShapeSize size => VectorSpace (Acc.Exp Double) (Sized size) where
--   type (Sized size1 ⊗ Sized size2) = Sized (TensorSize size1 size2)


---------------------------------------------------------------------------
-- | Module    : Numeric.Optimization.Basic.LP
-- Copyright   : (c) 2016 Marco Zocca
-- License     : GPL-3
--
-- Maintainer  : zocca.marco gmail
-- Stability   : experimental
-- Portability : portable
--
-- This module implements a general interior-point method for LP, based on the logarithmic barrier penalization
--
-- * [1] Primal-dual interior point algorithms for linear programming, G. Tzallas-Regas
--
--------------------------------------------------------------------------
module Numeric.Optimization.Basic.LP where

import Control.Monad.State

import Data.Sparse -- (reexports Data.Sparse.Common from sparse-linear-algebra)

import Numeric.LinearAlgebra.Sparse

feasX :: (Foldable t, Ord a, Num a) => t a -> Bool
feasX = all (> 0)

linesearch amax x s dx ds = undefined where

-- * Karmarkar algorithm

karmIpLP :: SpMatrix Double -> SpVector Double -> KarmLPData
karmIpLP aa c = execState (untilConverged (\x -> c `dot` _xKarm x) (karmarkarLPstep aa c ar)) kInit
  where
   ar = alpha * r
   kInit = KarmLPData x0 y0 0
   (m, n) = dimSM aa
   n' = fromIntegral n
   r = 1 / sqrt (n' * (n'-1))
   alpha = (n' - 1)/ (3*n')
   x0 = recip (fromIntegral n) .* onesSV n
   y0 = x0

data KarmLPData =
  KarmLPData { _xKarm, _yKarm:: SpVector Double, _niter :: Int } deriving (Eq, Show)

-- | alphar = alpha * r
karmarkarLPstep ::
  SpMatrix Double -> SpVector Double -> Double -> KarmLPData -> KarmLPData
karmarkarLPstep aa c alphar (KarmLPData x y k) = KarmLPData xnew ynew (k + 1) where
  (m, n) = dim aa
  ddk = diagonalSM x
  pp = (aa #~# ddk) -=- svToSM (onesSV n)
  cdk = c <# ddk
  cp = cdk ^-^ (transposeSM pp #> ((pp ##^ pp) <\> (pp #> cdk)))
  cp0 = normalize 2 cp
  ynew = y ^-^ (alphar .* cp0)
  dy = ddk #> ynew
  xnew = recip (sum dy) .* dy
  -- z = c `dot` xnew

-- untilConverged :: MonadState a m => (a -> SpVector Double) -> (a -> a) -> m a
untilConverged fproj = modifyInspectN 200 (normDiffConverged fproj)

normDiffConverged f [x1,x2] = nearZero $ f x1 - f x2

-- normDiffConverged fp xx = nearZero $ normSq (foldrMap fp (^-^) (zeroSV 0) xx)

{- |
Interior point, log-penalized LP (PLP):

max_x <c, x> - rho B(x)

s.t. A x = b

where B(x) := sum_j( log(x_j) )  (i.e. non-negativity constraint)

1. Write Lagrangian for PLP L(x, rho, y)
2. Write KKT conditions (first order optimality of L(x, rho, y))
3. Write Newton equations for KKT system
-}

-- | Update primal `x`, dual `y` and slack `s` variables by applying Newton's method to the KKT system of the log-penalized LP.
-- This requires one linear solve and a line search for a stepsize `alpha` that preserves feasibility : (x, s) > 0.
-- NB : It assumes the dimensions of vectors are x :: Rn , y :: Rm , s :: Rn

-- pdIpLPstep ::
--   (SpVector Double ->
--    SpVector Double ->
--    SpVector Double ->
--    SpVector Double ->
--    Double) ->
--   SpMatrix Double ->
--   Double ->
--   (SpVector Double, SpVector Double, SpVector Double, Int) ->
--   (SpVector Double, SpVector Double, SpVector Double, Int)
-- pdIpLPstep chooseAlpha aa tau (x, y, s, k) = (xNew, yNew, sNew, k+1) where
--   gamma = (x `dot` s) / fromIntegral n  -- duality measure
--   rho = tau * gamma                     -- Lagrange multiplier of barrier penalty
--   ss = diagonalSM s
--   xx = diagonalSM x
--   arow1 = zn -||- transposeSM aa -||- eye n
--   arow2 = aa -||- zm             -||- zm
--   arow3 = ss -||- zm             -||- xx
--   amat = arow1 -=- (arow2 -=- arow3)
--   rhs = concatSV zvn $ concatSV zvm $ negated (xx #> s) ^+^ (rho .* onesSV m)
--   xysStep = amat <\> rhs
--   xStep = takeSV n xysStep            -- 0 .. n-1
--   yStep = takeSV m (dropSV n xysStep) -- n .. n+m-1
--   sStep = dropSV (m + n) xysStep      -- n+m .. 2n+m
--   xNew = x ^+^ (alpha .* xStep)
--   yNew = y ^+^ (alpha .* yStep)
--   sNew = s ^+^ (alpha .* sStep)
--   alpha = chooseAlpha x s xStep sStep
--   (m, n) = dim aa
--   zm = zeroSM m m
--   zn = zeroSM n n
--   zvn = zeroSV n
--   zvm = zeroSV m


module UKF
  ( KF
  , kf_x
  , kf_p
  , kf_t

  , makeFilter
  , predict
  , update
  ) where

import Prelude hiding ((<>))

import App (Sensor (..))

import Numeric.LinearAlgebra

data KF = KF
  { kf_x :: Vector Double
  , kf_p :: Matrix Double
  , kf_t :: Word
  } deriving Show

makeFilter (Laser px py t) =
  KF { kf_x = 5 |> [px,
                    py,
                    0,
                    0,
                    0]
     , kf_p = (5><5) [0.15, 0.00, 0.0, 0.0, 0.0,
                      0.00, 0.15, 0.0, 0.0, 0.0,
                      0.00, 0.00, 1.0, 0.0, 0.0,
                      0.00, 0.00, 0.0, 1.0, 0.0,
                      0.00, 0.00, 0.0, 0.0, 1.0]
     , kf_t = t }

makeFilter (Radar rho phi rho' t) =
  KF { kf_x = 5 |> [rho * cos phi,
                    rho * sin phi,
                    rho',
                    0,
                    0]
     , kf_p = (5><5) [0.15, 0.00, 0.0, 0.0, 0.0,
                      0.00, 0.15, 0.0, 0.0, 0.0,
                      0.00, 0.00, 1.0, 0.0, 0.0,
                      0.00, 0.00, 0.0, 1.0, 0.0,
                      0.00, 0.00, 0.0, 0.0, 1.0]
     , kf_t = t }

-- Standard Deviation Tuning Parameters
std_a_     = 2.0
std_yawdd_ = 0.3

std_laspx_  = 0.15
std_laspy_  = 0.15

std_radr_   = 0.3
std_radphi_ = 0.03
std_radrd_  = 0.3

-- Unwind angle
normAngle θ
  | θ > pi  = normAngle (θ - 2*pi)
  | θ < -pi = normAngle (θ + 2*pi)
  | otherwise = θ

-- Calculate A st. m = AA^T
sqrtMat :: Matrix Double -> Matrix Double
sqrtMat = tr . chol . sym

{- constant turn rate and velocity magnitude (CTRV) motion model.

   x is the state augmented with the noise vector
   dt is a reasonable time step
 -}
process :: Vector Double -> Double -> Vector Double
process x dt =
  (5 |> [px, py, vel, yaw, yaw']) + f + v
  where
    -- State vector components
    px   = x ! 0
    py   = x ! 1
    vel  = x ! 2
    yaw  = x ! 3
    yaw' = x ! 4

    -- Noise vector components
    noiseA  = x ! 5
    noiseY' = x ! 6

    -- State model vector
    f = if abs yaw' > 0.001
        then
          5 |> [vel/yaw' * ((sin (yaw + yaw' * dt)) - (sin yaw)),
                vel/yaw' * ((cos yaw) - (cos (yaw + yaw' * dt))),
                0,
                yaw' * dt,
                0]
        else
          -- The special case where yaw rate is 0.
          -- Meaning that the motion is traveling
          -- along a straight line.
          5 |> [vel * dt * (cos yaw),
                vel * dt * (sin yaw),
                0,
                0,
                0]

    -- Noise model vector
    v = 5 |> [0.5 * noiseA * (dt ** 2) * (cos yaw),
              0.5 * noiseA * (dt ** 2) * (sin yaw),
              noiseA * dt,
              0.5 * noiseY' * (dt ** 2),
              noiseY' * dt]

{- This gives us KF(k+1 | k) for the filter.
 -}
predict :: Word -> KF -> KF
predict t kf =
  KF { kf_x = x'
     , kf_p = p'
     , kf_t = t }
  where
    -- Time inside the filter.
    filterTime =
      kf_t kf

    -- Delta-time. Essentially a timestep.
    dt =
      (fromIntegral (t - filterTime)) / 1000000.0

    -- n is the size of the state space with noise components
    n = 7

    -- l is a tuning parameter that controls the distance of the points.
    l = 3 - n

    -- Generate a list of augmented sigma points --
    sigmaPoints :: [Vector Double]
    sigmaPoints =
      let
        x = kf_x kf
        p = kf_p kf

        xAug = 7 |> [x ! 0,
                     x ! 1,
                     x ! 2,
                     x ! 3,
                     x ! 4,
                     0,
                     0]
        pAug = diagBlock [p, std_a_ ** 2, std_yawdd_ ** 2]

        sqrtP = toColumns $ sqrtMat $ (l+n) `scale` pAug
      in
        xAug : map (xAug +) sqrtP ++ map (xAug -) sqrtP

    -- Run the sigma points thru the motion model --
    predictedPoints =
      x0 : xs
      where
        x0 = process (head sigmaPoints) dt
        xs = map (\pnt -> process pnt dt)
                 (tail sigmaPoints)

    -- x(k+1 | k)
    x' :: Vector Double
    x' = foldr (\xn x -> x + (wn `scale` xn))
               (w0 `scale` x0)
               xs
      where
        x0 = head predictedPoints
        w0 = l / (l+n)
        wn = 1 / (2 * (l+n))
        xs = tail predictedPoints

    -- p(k+1 | k)
    p' :: Matrix Double
    p' = foldr (\xn p ->
                  let
                    xn_d  = xn - x'
                    xn_d' = 5 |> [xn_d ! 0,
                                  xn_d ! 1,
                                  xn_d ! 2,
                                  normAngle $ xn_d ! 3,
                                  xn_d ! 4]
                  in
                    p + (wn `scale` asColumn xn_d' <> asRow xn_d'))
               (w0 `scale` asColumn x0_d' <> asRow x0_d')
               xs
      where
        x0    = head predictedPoints
        w0    = l / (l+n)
        x0_d  = x0 - x'
        x0_d' = 5 |> [x0_d ! 0,
                      x0_d ! 1,
                      x0_d ! 2,
                      normAngle $ x0_d ! 3,
                      x0_d ! 4]

        wn = 1 / (2 * (l+n))
        xs = tail predictedPoints

{- For the case of laser measurments we don't need to
   use an unscented transformation for the measurment
   update. Standard kalman filter equations apply just
   fine here. However, we do need an unscented
   transformation for the model prediction. Because that
   is a non-linear model space. This gives us
   KF(k+1 | k+1) for the filter with a laser measurment.
 -}
update (Laser px py t) kf =
  KF { kf_x = x'
     , kf_p = p'
     , kf_t = t }
  where
    kf' = predict t kf

    -- Measurment transition matricies
    h = (2><5) [1, 0, 0, 0, 0,
                0, 1, 0, 0, 0]
    r = (2><2) [std_laspx_ ** 2, 0.00000000,
                0.00000000,      std_laspy_ ** 2]
    y = 2 |> [px, py] - (h #> kf_x kf')

    -- Kalman Filter Equations
    i  = ident 5
    ht = tr h
    s  = h <> kf_p kf' <> ht + r
    si = inv s
    k  = kf_p kf' <> ht <> si

    -- Estimation based on kalman filter gain kf(k+1 | k+1)
    x' = kf_x kf' + (k #> y)
    p' = (i - k <> h) <> kf_p kf'

{- For the case of radar measurments we do
   need to use an unscented transformation as
   the both the mapping to measurment space is
   non-linear. This gives us  KF(k+1 | k+1) for
   the filter
 -}
update (Radar rho phi rho' t) kf =
  kf -- skip
  where
    kf' = predict t kf

    n = 5
    l = 3 - n

    --Calculates the state space point as a radar space point.
    radarModel x =
      let
        _x  = x ! 0
        _y  = x ! 1
        _v  = x ! 2
        _θ  = x ! 3
        _vx = _v * (cos _θ)
        _vy = _v * (sin _θ)

        ρ  = sqrt ((_x ** 2) + (_y ** 2))
        φ  = atan2 _y _x
        ρ' = (_x*_vx + _y*_vy) / ρ
      in
        3 |> [ρ, φ, ρ']

    -- Generate a list of sigma points --
    sigmaPoints :: [Vector Double]
    sigmaPoints =
      let
        x = kf_x kf
        p = kf_p kf

        sqrtP = toColumns $ sqrtMat $ (l+n) `scale` p
      in
        x : map (x +) sqrtP ++ map (x -) sqrtP

    -- Run the sigma points thru the radar model --
    zs = map radarModel sigmaPoints

    -- z(k+1 | k)
    z' :: Vector Double
    z' = foldr (\zn z -> z + (wn `scale` zn))
               (w0 `scale` z0)
               (tail zs)
      where
        z0 = head zs
        w0 = l / (l+n)
        wn = 1 / (2 * (l+n))

    -- s(k+1 | k)
    s' :: Matrix Double
    s' = s + r
      where
        z0 = head zs
        w0 = l / (l+n)
        wn = 1 / (2 * (l+n))
        z0_d  = z0 - z'

        r = (3><3) [std_radr_ ** 2, 0, 0,
                    0, std_radphi_ ** 2, 0,
                    0, 0, std_radrd_ ** 2]

        s = foldr (\zn s ->
                     let
                       zn_d  = zn - z'
                     in
                       s + (wn `scale` asColumn zn_d <> asRow zn_d))
                  (w0 `scale` asColumn z0_d <> asRow z0_d)
                  (tail zs)


module Rnd (Distr, Rnd, minMax95, minMax99, meanMinMax95, meanMinMax99,
  (|*|), (|+|), runRnd, simulate, mkN95, distrN, distrAsymN, distrConst,
  simulateN, simulateMmm, MeanMinMax(..), scanM, withConfidence, mkMeanMinMax, showD,
  mmmProbability95AtLeast, meanMinMax95N, mmm) where

import Statistics.Distribution.Normal
import Statistics.Distribution
import Data.Random.Normal
import Control.Monad.State
import System.Random
import qualified Data.List as List
import Numeric

type Rnd = State StdGen
type Distr = Rnd Double

distrConst :: Double -> Distr
distrConst = return

distrN :: Double -> Double -> Distr
distrN m std = (+) m . (*) std <$> state normal

distrAsymN :: Double -> Double -> Double -> Distr
distrAsymN m lowStd highStd = (+) m . mulStd <$> state normal
  where
    mulStd :: Double -> Double
    mulStd v
      | v <= 0 = v*lowStd
      | otherwise = v*highStd

data MeanMinMax = MeanMinMax {
  mmmMean :: Double,
  mmmMin :: Double,
  mmmMax :: Double
}

mmm :: Double -> Double -> Double -> MeanMinMax
mmm = MeanMinMax

instance Show MeanMinMax where
    show (MeanMinMax m pMin pMax) = showD 2 m ++ " (" ++ showD 2 pMin ++ " .. " ++ showD 2 pMax ++ ")"

instance Num MeanMinMax where
  (MeanMinMax a c e) + (MeanMinMax b d f) = MeanMinMax (a + b) (c + d) (e + f)
  (MeanMinMax a c e) - (MeanMinMax b d f) = MeanMinMax (a - b) (c - d) (e - f)
  (MeanMinMax a c e) * (MeanMinMax b d f) = MeanMinMax (a * b) (c * d) (e * f)
  negate (MeanMinMax a b c) = MeanMinMax (-a) (-b) (-c)
  abs (MeanMinMax a b c) = MeanMinMax (abs a) (abs b) (abs c)
  signum (MeanMinMax a b c) = MeanMinMax (signum a) (signum b) (signum c)
  fromInteger x =
    let d = fromIntegral x
    in MeanMinMax d d d

instance Fractional MeanMinMax where
  (MeanMinMax a c e) / (MeanMinMax b d f) = MeanMinMax (a / b) (c / d) (e / f)
  fromRational x =
    let d = fromRational x
    in MeanMinMax d d d

mmmProbability95AtLeast :: MeanMinMax -> Double -> Double
mmmProbability95AtLeast (MeanMinMax m a z) x
  | x == m = 0.5
  | x < m = complCumulative (normalDistr m ((m - a)/2)) x
  | otherwise = complCumulative (normalDistr m ((z - m)/2)) x

(|*|) :: MeanMinMax -> Double -> MeanMinMax
(MeanMinMax m pMin pMax) |*| v = MeanMinMax (m*v) (pMin*v) (pMax*v)

(|+|) :: MeanMinMax -> Double -> MeanMinMax
(MeanMinMax m pMin pMax) |+| v = MeanMinMax (m+v) (pMin+v) (pMax+v)

minMax95 :: Double -> Double -> Distr
minMax95 pMin pMax = distrN m std where
  m = (pMax + pMin)/2
  std = (pMax - pMin)/4

minMax99 :: Double -> Double -> Distr
minMax99 pMin pMax = distrN m std where
  m = (pMax + pMin)/2
  std = (pMax - pMin)/6

meanMinMax95N :: MeanMinMax -> Distr
meanMinMax95N (MeanMinMax m pMin pMax) = distrAsymN m lowStd highStd
  where
    lowStd = (m - pMin)/2
    highStd = (pMax - m)/2

meanMinMax95 :: Double -> Double -> Double -> Distr
meanMinMax95 m pMin pMax = meanMinMax95N $ MeanMinMax m pMin pMax

meanMinMax99N :: MeanMinMax -> Distr
meanMinMax99N (MeanMinMax m pMin pMax) = distrN mean1 std where
  mean1 = (pMax + 4*m + pMin)/6
  std = (pMax - pMin)/6

meanMinMax99 :: Double -> Double -> Double -> Distr
meanMinMax99 m pMin pMax = meanMinMax99N $ MeanMinMax m pMin pMax

mkN95 :: [Double] -> Distr
mkN95 [] = distrN 0 1
mkN95 xs = meanMinMax95N $ mkMeanMinMax $ withConfidence 0.997 xs

withConfidence :: Double -> [Double] -> [Double]
withConfidence confidence xs = if d == 0 then xs else take (len - 2*d) $ drop d $ List.sort xs
  where
    len = length xs
    d = floor $ (1 - confidence)/2 * fromIntegral len

mkMeanMinMax :: [Double] -> MeanMinMax
mkMeanMinMax xs = MeanMinMax (sum xs / fromIntegral (length xs) ) (minimum xs) (maximum xs)

simulate :: Int -> Rnd a -> [a]
simulate 0 _ = []
simulate times r = runRnd times r : simulate (times - 1) r

simulateN :: Int -> Rnd Double -> Distr
simulateN times r = mkN95 $ simulate times r

simulateMmm :: Int -> Double -> Rnd Double -> MeanMinMax
simulateMmm times confidence r = mkMeanMinMax $ withConfidence confidence $ simulate times r

runRnd :: Int -> Rnd a -> a
runRnd seed r = evalState r (mkStdGen seed)

scanM :: (Monad m) => (a -> b -> m a) -> a -> [b] -> m [a]
scanM _ q [] = return [q]
scanM f q (x:xs) =
  do
    q2 <- f q x
    qs <- scanM f q2 xs
    return (q:qs)

showD :: RealFloat a => Int -> a -> String
showD numOfDecimals floatNum = showFFloat (Just numOfDecimals) floatNum ""


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.CauchyLorentz
-- Copyright : (c) 2011 Aleksey Khudyakov
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The Cauchy-Lorentz distribution. It's also known as Lorentz
-- distribution or Breit–Wigner distribution.
--
-- It doesn't have mean and variance.
module Statistics.Distribution.CauchyLorentz (
    CauchyDistribution
  , cauchyDistribMedian
  , cauchyDistribScale
    -- * Constructors
  , cauchyDistribution
  , cauchyDistributionE
  , standardCauchy
  ) where

import Control.Applicative
import Data.Aeson             (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary            (Binary(..))
import Data.Maybe             (fromMaybe)
import Data.Data              (Data, Typeable)
import GHC.Generics           (Generic)
import qualified Statistics.Distribution as D
import Statistics.Internal

-- | Cauchy-Lorentz distribution.
data CauchyDistribution = CD {
    -- | Central value of Cauchy-Lorentz distribution which is its
    --   mode and median. Distribution doesn't have mean so function
    --   is named after median.
    cauchyDistribMedian :: {-# UNPACK #-} !Double
    -- | Scale parameter of Cauchy-Lorentz distribution. It's
    --   different from variance and specify half width at half
    --   maximum (HWHM).
  , cauchyDistribScale  :: {-# UNPACK #-} !Double
  }
  deriving (Eq, Typeable, Data, Generic)

instance Show CauchyDistribution where
  showsPrec i (CD m s) = defaultShow2 "cauchyDistribution" m s i
instance Read CauchyDistribution where
  readPrec = defaultReadPrecM2 "cauchyDistribution" cauchyDistributionE

instance ToJSON   CauchyDistribution
instance FromJSON CauchyDistribution where
  parseJSON (Object v) = do
    m <- v .: "cauchyDistribMedian"
    s <- v .: "cauchyDistribScale"
    maybe (fail $ errMsg m s) return $ cauchyDistributionE m s
  parseJSON _ = empty

instance Binary CauchyDistribution where
    put (CD m s) = put m >> put s
    get = do
      m <- get
      s <- get
      maybe (error $ errMsg m s) return $ cauchyDistributionE m s

-- | Cauchy distribution
cauchyDistribution :: Double    -- ^ Central point
                   -> Double    -- ^ Scale parameter (FWHM)
                   -> CauchyDistribution
cauchyDistribution m s
  = fromMaybe (error $ errMsg m s)
  $ cauchyDistributionE m s

-- | Cauchy distribution
cauchyDistributionE :: Double    -- ^ Central point
                    -> Double    -- ^ Scale parameter (FWHM)
                    -> Maybe CauchyDistribution
cauchyDistributionE m s
  | s > 0     = Just (CD m s)
  | otherwise = Nothing

errMsg :: Double -> Double -> String
errMsg _ s
  = "Statistics.Distribution.CauchyLorentz.cauchyDistribution: FWHM must be positive. Got "
  ++ show s

-- | Standard Cauchy distribution. It's centered at 0 and have 1 FWHM
standardCauchy :: CauchyDistribution
standardCauchy = CD 0 1

instance D.Distribution CauchyDistribution where
  cumulative (CD m s) x = 0.5 + atan( (x - m) / s ) / pi

instance D.ContDistr CauchyDistribution where
  density (CD m s) x = (1 / pi) / (s * (1 + y*y))
    where y = (x - m) / s
  quantile (CD m s) p
    | p > 0 && p < 1 = m + s * tan( pi * (p - 0.5) )
    | p == 0         = -1 / 0
    | p == 1         =  1 / 0
    | otherwise      =
      error $ "Statistics.Distribution.CauchyLorentz.quantile: p must be in [0,1] range. Got: "++show p
  complQuantile (CD m s) p
    | p > 0 && p < 1 = m + s * tan( pi * (0.5 - p) )
    | p == 0         =  1 / 0
    | p == 1         = -1 / 0
    | otherwise      =
      error $ "Statistics.Distribution.CauchyLorentz.complQuantile: p must be in [0,1] range. Got: "++show p

instance D.ContGen CauchyDistribution where
  genContVar = D.genContinuous

instance D.Entropy CauchyDistribution where
  entropy (CD _ s) = log s + log (4*pi)

instance D.MaybeEntropy CauchyDistribution where
  maybeEntropy = Just . D.entropy


{-# LANGUAGE Rank2Types #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Packed.Tri
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Triangular views of packed matrices.
--

module Numeric.LinearAlgebra.Packed.Tri (
    -- * Immutable interface

    -- ** Vector multiplication
    triMulVector,

    -- ** Vector solving
    triSolvVector,

    -- * Mutable interface
    triCreate,

    -- ** Vector multiplication
    triMulVectorM_,

    -- ** Vector solving
    triSolvVectorM_,

    ) where

import Control.Monad( when )
import Control.Monad.ST( ST, runST, unsafeIOToST )
import Text.Printf( printf )

import Numeric.LinearAlgebra.Vector( Vector, STVector )
import qualified Numeric.LinearAlgebra.Vector as V
import Numeric.LinearAlgebra.Packed.Base( Packed, RPacked, STPacked )
import qualified Numeric.LinearAlgebra.Packed.Base as P
import Numeric.LinearAlgebra.Types
import qualified Foreign.BLAS as BLAS

-- | A safe way to create and work with a mutable Tri Packed before returning
-- an immutable one for later perusal.
triCreate :: (Storable e)
          => (forall s. ST s (Tri (STPacked s) e))
          -> Tri Packed e
triCreate mt = runST $ do
    (Tri u d ma) <- mt
    a <- P.unsafeFreeze ma
    return $ Tri u d a

-- | @triMulVector trans a x@ returns @op(a) * x@, where @op(a)@ is
-- determined by @trans@.
triMulVector :: (BLAS2 e)
             => Trans
             -> Tri Packed e
             -> Vector e
             -> Vector e
triMulVector trans a x =
    V.create $ do
        x' <- V.newCopy x
        triMulVectorM_ trans a x'
        return x'

-- | @triMulVectorM_ a x@ sets @x := op(a) * x@, where @op(a)@ is determined
-- by @trans@.
triMulVectorM_ :: (RPacked p, BLAS2 e)
               => Trans -> Tri p e
               -> STVector s e
               -> ST s ()
triMulVectorM_ trans (Tri uplo diag a) x = do
    na <- P.getDim a
    nx <- V.getDim x
    let n = nx

    when (nx /= n) $ error $
        printf ("triMulVectorM_"
                ++ " _"
                ++ " (Tri _ _ <packed matrix with dim %d>)"
                ++ " <vector with dim %d>"
                ++ ": dimension mismatch")
               na
               nx

    unsafeIOToST $
        P.unsafeWith a $ \pa ->
        V.unsafeWith x $ \px ->
            BLAS.tpmv uplo trans diag n pa px 1

-- | @triSolvVector trans a x@ returns @op(a) \\ x@, where @op(a)@ is
-- determined by @trans@.
triSolvVector :: (BLAS2 e)
             => Trans
             -> Tri Packed e
             -> Vector e
             -> Vector e
triSolvVector trans a x =
    V.create $ do
        x' <- V.newCopy x
        triSolvVectorM_ trans a x'
        return x'

-- | @triSolvVectorM_ a x@ sets @x := op(a) \\ x@, where @op(a)@ is determined
-- by @trans@.
triSolvVectorM_ :: (RPacked p, BLAS2 e)
               => Trans -> Tri p e
               -> STVector s e
               -> ST s ()
triSolvVectorM_ trans (Tri uplo diag a) x = do
    na <- P.getDim a
    nx <- V.getDim x
    let n = nx

    when (nx /= n) $ error $
        printf ("triMulVectorM_"
                ++ " _"
                ++ " (Tri _ _ <packed matrix with dim %d>)"
                ++ " <vector with dim %d>"
                ++ ": dimension mismatch")
               na
               nx

    unsafeIOToST $
        P.unsafeWith a $ \pa ->
        V.unsafeWith x $ \px ->
            BLAS.tpsv uplo trans diag n pa px 1


{-# LANGUAGE FlexibleInstances,
             DeriveDataTypeable #-}

module XAlg.Languages.Arithmetic.Number where

import XAlg.Foundation.Core

import Prelude hiding               (Float, Int)
import qualified  Data.Complex as C (Complex)
import Data.Complex          hiding (Complex)
import Data.Data

type Z = Integer
type R = Double
type C = C.Complex Double

deriving instance Ord (C.Complex Double)

data Number = Real R
            | Int Z
            | Complex C
            deriving (Eq, Ord, Data)

instance Show Number where
  show (Real r) = show r
  show (Int z) = show z
  show (Complex c) = show c

instance Z <: R where
  incl = fromInteger

instance Z <: C where
  incl = fromInteger

instance R <: C where
  incl = (0 :+)

compat :: (Number,Number) -> (Number,Number)
compat p@(Int n,Int m)          = p
compat   (Int n,r@(Real _))     = (Real (incl n),r)
compat   (Int n,c@(Complex _))  = (Complex (incl n),c)
compat   (Real n,Int m)         = (Real n, Real (incl m))
compat p@(Real _,Real _)        = p
compat   (Real n,c@(Complex _)) = (Complex (incl n),c)
compat p@(Complex _,Complex _)  = p
compat   (Complex n, Int m)     = (Complex n, Complex (incl m))
compat   (Complex n, Real m)    = (Complex n, Complex (incl m))

caseFun :: (Z -> Z) -> (R -> R) -> (C -> C) -> Number -> Number
caseFun zf rf cf (Real n)    = Real $ rf n
caseFun zf rf cf (Int n)     = Int $ zf n
caseFun zf rf cf (Complex n) = Complex $ cf n

caseFun2 :: (Z -> Z -> Z) -> (R -> R -> R) -> (C -> C -> C) -> Number -> Number -> Number
caseFun2 zf rf cf (Real n)    (Real m)    = Real $ rf n m
caseFun2 zf rf cf (Real n)    (Int m)     = Real $ rf n (incl m)
caseFun2 zf rf cf (Real n)    (Complex m) = Complex $ cf (incl n) m
caseFun2 zf rf cf (Int n)     (Real m)    = Real $ rf (incl n) m
caseFun2 zf rf cf (Int n)     (Int m)     = Int $ zf n m
caseFun2 zf rf cf (Int n)     (Complex m) = Complex $ cf (incl n) m
caseFun2 zf rf cf (Complex n) (Real m)    = Complex $ cf n (incl m)
caseFun2 zf rf cf (Complex n) (Int m)     = Complex $ cf n (incl m)
caseFun2 zf rf cf (Complex n) (Complex m) = Complex $ cf n m

instance Num Number where
  fromInteger = Int
  (+) = caseFun2 (+) (+) (+)
  (-) = caseFun2 (-) (-) (-)
  (*) = caseFun2 (*) (*) (*)
  abs = caseFun abs abs abs
  signum = caseFun signum signum signum

instance Fractional Number where
  (/) (Int x) (Int y) = Real $ fromInteger x / fromInteger y
  (/) x y = uncurry (/) . compat $ (x,y)
  recip (Int x) = Real . recip . fromInteger $ x
  recip x = caseFun undefined recip recip x
  fromRational = Real . fromRational

instance Floating Number where
  pi = Real pi
  exp (Int z) = Real $ exp $ fromInteger z
  exp (Real r) = Real $ exp r
  exp (Complex c) = Complex $ exp c
  log (Int z) = Real $ log $ fromInteger z
  log (Real r) = Real $ log r
  log (Complex c) = Complex $ log c
  sin (Int z) = Real $ sin $ fromInteger z
  sin (Real r) = Real $ sin r
  sin (Complex c) = Complex $ sin c
  cos (Int z) = Real $ cos $ fromInteger z
  cos (Real r) = Real $ cos r
  cos (Complex c) = Complex $ cos c
  asin (Int z) = Real $ asin $ fromInteger z
  asin (Real r) = Real $ asin r
  asin (Complex c) = Complex $ asin c
  acos (Int z) = Real $ acos $ fromInteger z
  acos (Real r) = Real $ acos r
  acos (Complex c) = Complex $ acos c
  atan (Int z) = Real $ atan $ fromInteger z
  atan (Real r) = Real $ atan r
  atan (Complex c) = Complex $ atan c
  sinh (Int z) = Real $ sinh $ fromInteger z
  sinh (Real r) = Real $ sinh r
  sinh (Complex c) = Complex $ sinh c
  cosh (Int z) = Real $ cosh $ fromInteger z
  cosh (Real r) = Real $ cosh r
  cosh (Complex c) = Complex $ cosh c
  asinh (Int z) = Real $ asinh $ fromInteger z
  asinh (Real r) = Real $ asinh r
  asinh (Complex c) = Complex $ asinh c
  acosh (Int z) = Real $ acosh $ fromInteger z
  acosh (Real r) = Real $ acosh r
  acosh (Complex c) = Complex $ acosh c
  atanh (Int z) = Real $ atanh $ fromInteger z
  atanh (Real r) = Real $ atanh r
  atanh (Complex c) = Complex $ atanh c


import Data.Word
import Data.List(elemIndex)
import Data.Maybe
import Data.Complex
import qualified Data.Array.Repa as R
import Data.Array.Repa.IO.BMP
import System.Environment

type Number = Double
type ComplexNumber = Complex Number
type IterationCount = Integer
type PointData = (ComplexNumber, IterationCount)
type RootData = (Int, IterationCount)
type Bounds = (Number, Number, Number, Number)
type Step = Number
type Color = (Word8, Word8, Word8)

e :: ComplexNumber
e = (exp 1) :+ 0

data Function = Constant ComplexNumber
              | Z
              | Sum Function Function
              | Product Function Function
              | Quotient Function Function
              | Power ComplexNumber Function
              | Sin Function
              | Cos Function
              | Exponential ComplexNumber Function
              | Log ComplexNumber Function
              | Undefined
              | Indeterminate
              deriving (Show)

realConstant :: Number -> Function
realConstant n = Constant (n :+ 0)

derivative :: Function -> Function
derivative (Constant _) = realConstant 0
derivative Z = realConstant 1
derivative (Sum u v) = Sum (derivative u) (derivative v)
derivative (Product u v) = Sum (Product (derivative u) v) (Product u (derivative v))
derivative (Quotient u v) = (Quotient (Sum (Product v (derivative u)) (Product (Product (realConstant (-1)) u) (derivative v))) (Power 2 v))
derivative (Power p f) = Product (Product (Constant p) (Power (p - 1) f)) (derivative f)
derivative (Sin f) = Product (Cos f) (derivative f)
derivative (Cos f) = Product (Product (realConstant (-1)) (Sin f)) (derivative f)
derivative (Exponential b f) = Product (Product (Exponential b f) (Log e (Constant b))) (derivative f)
derivative (Log b f) = Quotient (derivative f) (Product f (Log e (Constant b)))
derivative Undefined = Undefined
derivative Indeterminate = Indeterminate

evaluate :: Function -> ComplexNumber -> ComplexNumber
evaluate (Constant a) z = a
evaluate Z z = z
evaluate (Sum u v) z = (evaluate u z) + (evaluate v z)
evaluate (Product u v) z = (evaluate u z) * (evaluate v z)
evaluate (Quotient u v) z = (evaluate u z) / (evaluate v z)
evaluate (Power p f) z = (evaluate f z) ** p
evaluate (Sin f) z = sin (evaluate f z)
evaluate (Cos f) z = cos (evaluate f z)
evaluate (Exponential b f) z = b ** (evaluate f z)
evaluate (Log b f) z = logBase b (evaluate f z)
evaluate Undefined z = undefined
evaluate Indeterminate z = undefined

maxIterations :: IterationCount
maxIterations = 100

epsilon :: Number
epsilon = 1e-12

newton :: Function -> Function -> ComplexNumber -> PointData
newton f f' z0 = newton' z1 z0 0
    where approximate z = z - (evaluate f z) / (evaluate f' z)
          z1 = approximate z0
          newton' current previous n
              | n > maxIterations = (current, n)
              | magnitude (current - previous) < epsilon = (current, n)
              | otherwise = newton' (approximate current) current (n + 1)

windowApply :: (ComplexNumber -> PointData) -> Bounds -> Step -> [PointData]
windowApply f (top, left, bottom, right) step = [f (x :+ y ) | y <- ys, x <- xs]
    where xs = [left, left + step .. right]
          ys = [top, top + step .. bottom]

approximateIndex :: ComplexNumber -> [ComplexNumber] -> Maybe Int
approximateIndex z zs = elemIndex True [magnitude (z - zi) < epsilon * 2 | zi <- zs]

closestIndex :: ComplexNumber -> [ComplexNumber] -> Int
closestIndex z zs = snd . minimum $ zip [magnitude (z - zi) | zi <- zs] [0..]

closeEnough :: ComplexNumber -> [ComplexNumber] -> Bool
closeEnough z zs = case approximateIndex z zs of
    Just _ -> True
    Nothing ->  False

roots :: [PointData] -> [ComplexNumber]
roots ps = foldr checkRoot [] zs
    where checkRoot z roots
              | z `closeEnough` roots = roots
              | otherwise = z:roots
          zs = fst . unzip $ ps

assignRoot :: [ComplexNumber] -> PointData -> RootData
assignRoot rs p@(z, i) = (closestIndex z rs, i)

colorRoot :: RootData -> Color
colorRoot (n, i) = (ra, ga, ba)
    where colors = [ (0, 0, 0)
                   , (255, 0, 0)
                   , (0, 255, 0)
                   , (0, 0, 255)
                   , (255, 255, 0)
                   , (255, 0, 255)
                   , (0, 255, 255)
                   , (255, 255, 255)
                   ]
          (r, g, b) = colors !! n
          ii = max i 1
          ra = fromIntegral . min 255  $ 4 * r `div` ii
          ga = fromIntegral . min 255  $ 4 * g `div` ii
          ba = fromIntegral . min 255  $ 4 * b `div` ii

colorRoots :: Function -> Bounds -> Step -> [Color]
colorRoots f bounds step = map colorRoot rds
    where ps = windowApply (newton f (derivative f)) bounds step
          rs = roots ps
          rds = map (assignRoot rs) ps

arrayify :: Int -> Int -> [Color] -> R.Array R.U R.DIM2 Color
arrayify width height =
    R.fromListUnboxed (R.Z R.:. height R.:. width :: R.DIM2)

f = Sum (Power 3 Z) (realConstant (-1))

main = do
    args <- getArgs
    let [fileName, stringTop, stringLeft, stringBottom, stringRight, stringStep] = args
        top = read stringTop
        left = read stringLeft
        bottom = read stringBottom
        right = read stringRight
        step = read stringStep
        width = floor ((right - left) / step) + 1
        height = floor ((bottom - top) / step) + 1
        colorList = colorRoots f (top, left, bottom, right) step
        colorArray = arrayify width height colorList
    writeImageToBMP fileName colorArray