{-# LANGUAGE OverloadedStrings, NoImplicitPrelude, BangPatterns, DeriveGeneric #-}
module NLP.Albemarle.Sparse
    ( dense
    , fromDocuments
    , transpose
    , shift
    , size
    , sparse
    , mult
    , multCol
    , SparseMatrix(..)
    , Coord
    , Coords
    ) where
import ClassyPrelude
import NLP.Albemarle hiding (SparseMatrix)
import GHC.Generics (Generic)
import Data.Vector.Binary
import qualified Data.Vector.Unboxed as Vec
import qualified Data.Vector.Storable as SVec
import qualified Data.Vector.Algorithms.Intro as Introsort
import qualified Numeric.LinearAlgebra as HMatrix
import qualified Data.Binary as Bin
import Control.Parallel.Strategies (rpar, rseq, parList, withStrategy, dot, rdeepseq)
type Coord = (Int, Int, Double)
type Coords = Vec.Vector Coord
data SparseMatrix = RowMatrix !Coords | ColMatrix !Coords
  deriving (Show, Eq, Generic)
instance Bin.Binary SparseMatrix

-- | Create a sparse matrix from rows
fromDocuments :: [SparseVector] -> SparseMatrix
fromDocuments rows = RowMatrix $ Vec.concat [ docToSparseRow docid doc | (docid, doc) <- zip [0..] rows ]
  where docToSparseRow docid = Vec.map (\(colid, val) -> (docid, colid, fromIntegral val)) . Vec.fromList

-- | Find the size of a sparse matrix
size :: SparseMatrix -> (Int, Int)
size (RowMatrix mat) = (maxx+1, maxy+1)
  where (maxx, maxy) = Vec.foldr' (\(r, c, _) (mr, mc) -> (max mr r, max mc c)) (-1, -1) mat
size (ColMatrix mat) = (maxx+1, maxy+1)
  where (maxx, maxy) = Vec.foldr' (\(r, c, _) (mr, mc) -> (max mr r, max mc c)) (-1, -1) mat

-- | Convert a sparse matrix into a dense one
dense :: SparseMatrix -> HMatrix.Matrix Double
-- dense (RowMatrix mat) = HMatrix.toDense $ (\(row, col, val) -> ((row, col), val)) <$> Vec.toList mat
-- dense (ColMatrix mat) = HMatrix.toDense $ (\(row, col, val) -> ((row, col), val)) <$> Vec.toList mat
dense (RowMatrix mat) =
  HMatrix.reshape width $ Vec.convert flat_mat
  where
    (width, height) = size $ RowMatrix mat
    flat_mat = Vec.update
      (Vec.replicate (width*height) 0.0) $
      Vec.map (\(r,c,v) -> (r*width+c, v)) mat
dense (ColMatrix mat) = dense $ RowMatrix mat -- a little fib but it works

-- | Convert a dense matrix into a sparse one
sparse :: HMatrix.Matrix Double -> SparseMatrix
sparse mat =
  RowMatrix $ Vec.imap (\ix v -> (
    ix `div` width,
    ix `mod` width,
    v)) $ Vec.convert $ HMatrix.flatten mat
  where (width, height) = HMatrix.size mat
  --RowMatrix $ Vec.fromList [ (ri, ci, val)
  --  | (ri, row) <- zip [0..] lmat,
  --    (ci, val) <- zip [0..] row, abs val > 1e-6 ]
  --where lmat = HMatrix.toLists mat

transpose :: SparseMatrix -> SparseMatrix
transpose (RowMatrix mat) = RowMatrix $ Vec.map (\(row, col, val) -> (col, row, val)) mat
transpose (ColMatrix mat) = ColMatrix $ Vec.map (\(col, row, val) -> (row, col, val)) mat

shift :: SparseMatrix -> SparseMatrix
shift (RowMatrix mat) = ColMatrix $ Vec.modify (Introsort.sortBy $ comparing extractCol) mat
shift (ColMatrix mat) = RowMatrix $ Vec.modify (Introsort.sortBy $ comparing extractRow) mat

extractRow (row, col, val) = row
extractCol (row, col, val) = col

-- | Multiplies a RowMatrix by a ColMatrix (converting as necessary beforehand)
-- and generates a RowMatrix as a result
mult :: SparseMatrix -> SparseMatrix -> SparseMatrix
mult left@(ColMatrix _) right@(ColMatrix _) = mult (shift left) right
mult left@(ColMatrix _) right@(RowMatrix _) = mult (shift left) (shift right)
mult left@(RowMatrix _) right@(RowMatrix _) = mult left (shift right)
mult (RowMatrix !left) (ColMatrix !right) =
  RowMatrix $ Vec.fromList [ (rowid, colid, multLines 0.0 row col)
    | (rowid, row) <- rows,
      (colid, col) <- cols]
  where
    rows = zip [0..] (getLines extractRow left)
    cols = zip [0..] (getLines extractCol right)

-- | Multiplies a RowMatrix by a ColMatrix (converting as necessary beforehand)
-- and generates a ColMatrix as a result
multCol :: SparseMatrix -> SparseMatrix -> SparseMatrix
multCol left@(ColMatrix _) right@(ColMatrix _) = multCol (shift left) right
multCol left@(ColMatrix _) right@(RowMatrix _) = multCol (shift left) (shift right)
multCol left@(RowMatrix _) right@(RowMatrix _) = multCol left (shift right)
multCol (RowMatrix left) (ColMatrix right) =
  ColMatrix $ Vec.fromList [ (rowid, colid, multLines 0.0 row col)
    | (colid, col) <- zip [0..] (getLines extractCol right),
      (rowid, row) <- zip [0..] (getLines extractRow left)]

--
--
--   INTERNAL METHODS
--
--

multLines :: Double -> [Coord] -> [Coord] -> Double
multLines x [] _ = x
multLines x _ [] = x
multLines !total xs@((_, xcol, xval):xr) ys@((yrow, _, yval):yr) =
  case compare xcol yrow of
    LT -> multLines total xr ys
    GT -> multLines total xs yr
    EQ -> multLines (total + xval * yval) xr yr

getLines :: (Coord -> Int) -> Coords -> [[Coord]]
getLines extract mat
  | Vec.null mat = []
  | otherwise = Vec.toList thisrow : getLines extract rest
  where
    (thisrow, rest) = Vec.span (\c -> extract c == target) mat
    target = extract $ Vec.head mat

-- Gets the index that a row starts at (but it works with cols too)
csrRowBeginsAt :: (Coord -> Int) -> Coords -> Vec.Vector Int
csrRowBeginsAt extract mat =
  Vec.cons 0 $ Vec.accumulate max zeros $ Vec.imap (\idx coord -> (extract coord, idx+1)) mat
  where
    zeros = Vec.replicate (end+1) 0
    end = extract $ Vec.last mat


module Solve where

import Data.List
import Control.Applicative ((<|>), many)
import Control.Monad (void)
import Data.Char (isLetter, isDigit)
import qualified Debug.Trace as DT
import Numeric.LinearAlgebra.Data
import Data.Int

import Text.ParserCombinators.Parsec
import qualified Text.ParserCombinators.Parsec.Token as P
import Text.ParserCombinators.Parsec.Language

-- TYPE DEFINITIONS
data Point = MkPoint {
  x :: Int,
  y :: Int
} deriving (Show, Eq)

data Dimensions = MkDimensions {
  width  :: Int,
  height :: Int
} deriving (Show, Eq)

data Rectangle = MkRectangle {
  upperLeft :: Point,
  dimensions :: Dimensions
} deriving (Show, Eq)

data Claim  = MkClaim {
    index :: Int,
    rectangle :: Rectangle
} deriving (Show, Eq)

-- PARSING
-- of course, we could simply use a regular expression, but I've always wanted to play
-- around with the Parsec library for generating parsers, so bear with me :-)
numParser :: Parser Int
numParser = do
    n <- many1 digit
    return (read n)

pointParser :: Parser Point
pointParser = do
    x <- numParser
    void $ char ','
    y <- numParser
    return (MkPoint x y)

dimensionsParser :: Parser Dimensions
dimensionsParser = do
  width <- numParser
  void $ char 'x'
  height <- numParser
  return (MkDimensions width height)

rectangleParser :: Parser Rectangle
rectangleParser = do
    p <- pointParser
    void $ char ':'
    void $ char ' '
    dim <- dimensionsParser
    return (MkRectangle p dim)

claimParser :: Parser Claim
claimParser = do
    void $ char '#'
    idx <- numParser
    void $ char ' '
    void $ char '@'
    void $ char ' '
    rect <- rectangleParser
    return (MkClaim idx rect)

-- helper function to apply a parser
regularParse :: Parser a -> String -> Either ParseError a
regularParse p = parse p ""

-- MAIN routines
--
-- create an association list for the given claim
-- the association list contains a list of (x,y) coordinates plus the value (which equals the claim's index)
toAssocList :: Claim -> [((Int, Int), Int64)]
toAssocList c = map (\(i,j) -> ((i,j), (fromIntegral idx))) indices
  where idx = index c
        r = rectangle c
        x' = x $ upperLeft r
        y' = y $ upperLeft r
        width' = width $ dimensions r
        height' = height $ dimensions r
        indices = [(i,j) | i <- [(x'+1) .. (x' + width')],
                           j <- [(y'+1) .. (y' + height')]]

-- apply the given claim to the given matrix
applyClaim :: Claim -> Matrix Int64 -> Matrix Int64
applyClaim c m = accum m (\new old -> if old == 0 then new else -1) (toAssocList c)

-- apply the given list of claims to the given matrix
applyClaims :: [Either ParseError Claim ] -> Matrix Int64 -> Matrix Int64
applyClaims [] m = m
applyClaims (x:xs) m =
  case x of
    Right x' -> applyClaims xs (applyClaim x' m)
    Left _   -> applyClaims xs m -- leave m unchanged for erroneous claim

-- count the elements in the given matrix that are equal to the given value
countElems :: Matrix Int64 -> Int64 -> Int
countElems m val = length $ filter (\x -> x == val) $ toList $ flatten m

-- count the clashes (squares with overlapping claims)
countClashes :: Matrix Int64 -> Int
countClashes m = countElems m (-1)

-- prepare solution: parse input, build a 1001x1001 matrix, and apply all claims to it
prepSolution:: String -> ([Claim], Matrix Int64)
prepSolution input = (validClaims, applyClaims claims m)
  where claims = map (\line -> regularParse claimParser line) $ lines input
        m = (1001><1001) (repeat 0) -- 1001x1001, all elements are 0
        m2 = applyClaims claims m
        validClaims = map (\(Right x) -> x) $ filter (\x -> case x of Right _ -> True) claims

-- find the single claim that is intact
-- our approach is kind of brute force - for n claims, we check all matrix elements n times
findIntactClaim :: ([Claim], Matrix Int64) -> Claim
findIntactClaim (claims, m) = head $ filter (\c -> area c == countElems m (fromIntegral (index c))) claims
  where area c' = width (dimensions (rectangle c')) * height (dimensions (rectangle c'))

solveI :: String -> Int
solveI = countClashes . snd . prepSolution

solveII :: String -> Int
solveII = index . findIntactClaim . prepSolution

#if defined(STANDALONE)
main = do
    input <- readFile "input.txt"
    putStrLn $ show $ solveI  input
    putStrLn $ show $ solveII input
#endif


{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}

-- |
-- Module      :  Spiral.SPL
-- Copyright   :  (c) 2016-2017 Drexel University
-- License     :  BSD-style
-- Maintainer  :  mainland@drexel.edu

module Spiral.SPL (
    module Spiral.Permutation,

    SPL(..),
    matrix,
    fromLists,
    fromFunction,

    diag,
    circ,
    skew,
    toep,
    permute,
    backpermute,

    extent,

    mXv,
    (#>),

    toMatrix,
    col,
    row,

    transpose,

    inverse,

    (<->),
    (<|>),
    block,

    (⊗),
    (×),
    (⊕),

    Z(..),
    (:.)(..),

    DIM0,
    DIM1,
    DIM2,
    ix1,
    ix2
  ) where

import Data.Complex
import Data.Monoid ((<>))
import Data.Typeable (Typeable)
import qualified Data.Vector as V
import Text.PrettyPrint.Mainland hiding ((<|>))
import Text.PrettyPrint.Mainland.Class

import qualified Spiral.Array as A
import Spiral.Array (D,
                     IArray,
                     M,
                     Matrix,
                     Vector,
                     manifest)
import qualified Spiral.Array.Operators.Mapping as A
import qualified Spiral.Array.Operators.Matrix as A
import Spiral.Array.Repr.Complex
import Spiral.Array.Shape
import Spiral.Exp
import Spiral.Permutation hiding (Inv)
import Spiral.RootOfUnity
import Spiral.Util.Pretty

-- | Wrap an array to ensure it is unconditionally Show-able.
newtype ShowArray r sh e = ShowArray { unShowArray :: Array r sh e }

instance (Show sh, Show e, IArray r sh e) => Show (ShowArray r sh e) where
    show arr = show (manifest (unShowArray arr))

-- | An SPL transform. The type index is the type of the scalar values in the
-- transformed vector.
data SPL a where
    -- | An "embedded" array with an unknown representation.
    E :: IArray r DIM2 e => ShowArray r DIM2 e -> SPL e

    -- | The $n \times n$ identity matrix.
    I :: Num e => Int -> SPL e

    -- | Transpose
    T :: SPL e -> SPL e

    -- | Inverse
    Inv :: (Eq e, Fractional e) => SPL e -> SPL e

    -- | A permutation
    Pi :: Permutation -> SPL e

    -- | The rotation matrix
    Rot :: Floating a => a -> SPL a

    -- | A diagonal matrix
    Diag :: V.Vector e -> SPL e

    -- | The $n \times n$ diagonal matrix with constant diagonal elements.
    KDiag :: Int -> e -> SPL e

    -- | Vertical stacking of transforms
    Above :: SPL e -> SPL e -> SPL e

    -- | Horizontal stacking of transforms
    Beside :: SPL e -> SPL e -> SPL e

    -- | Kronecker product
    Kron :: SPL e -> SPL e -> SPL e

    -- | Direct sum
    DSum :: SPL e -> SPL e -> SPL e

    -- | Matrix product
    Prod :: SPL e -> SPL e -> SPL e

    -- | Circulant matrix given first column
    Circ :: V.Vector e -> SPL e

    -- | Skew-Circulant matrix given first column
    Skew :: V.Vector e -> SPL e

    -- | Toeplitz matrix
    Toep :: V.Vector e -> SPL e

    -- | Make a complex transform into a real transform.
    Re :: RealFloatConst a => SPL (Exp (Complex a)) -> SPL (Exp a)

    -- | The 2x2 DFT
    F2 :: SPL e

    -- | The nxn DFT matrix
    DFT :: RootOfUnity a => Int -> SPL a

    -- | The nxn Inverse DFT matrix
    DFT' :: RootOfUnity a => Int -> SPL a

    -- | The nxn "rotated" DFT matrix
    F :: RootOfUnity a => Int -> a -> SPL a

    -- | The nxn "rotated" inverse DFT matrix
    F' :: RootOfUnity a => Int -> a -> SPL a

deriving instance Show e => Show (SPL e)
deriving instance Typeable e => Typeable (SPL e)

-- | Embed any 'Matrix' as an SPL term.
matrix :: IArray r DIM2 e
       => Array r DIM2 e
       -> SPL e
matrix = E . ShowArray

-- | Embed a matrix created from a list of lists of values.
fromLists :: [[e]] -> SPL e
fromLists = matrix . A.fromLists

-- | Embed a matrix created from a function mapping indices to elements.
fromFunction :: DIM2 -> (DIM2 -> e) -> SPL e
fromFunction sh = matrix . A.fromFunction sh

-- | Create a diagonal matrix
diag :: [a] -> SPL a
diag = Diag . V.fromList

-- | Create a circulant matrix from the first column
circ :: [a] -> SPL a
circ = Circ . V.fromList

-- | Create a skew circulant matrix from the first column
skew :: [a] -> SPL a
skew = Skew . V.fromList

-- | Create a Toeplitz matrix
toep :: [a] -> SPL a
toep = Toep . V.fromList

-- | Permute (scatter).
permute :: Permutation -> SPL e
permute = Pi

-- | Backpermute (gather).
backpermute :: Permutation -> SPL e
backpermute = Pi . invert

-- | Return the extent of an SPL transform.
extent :: SPL a -> DIM2
extent (E a)       = A.extent (unShowArray a)
extent (I n)       = ix2 n n
extent (T e)       = ix2 n m
                       where
                         Z :. m :. n = extent e
extent (Inv e)     = ix2 n n
                       where
                         Z :. n :. _ = extent e
extent (Pi p)      = ix2 (dim p) (dim p)
extent (Rot _)     = ix2 2 2
extent (Diag xs)   = ix2 n n
                       where
                         n = length xs
extent (KDiag n _) = ix2 n n

extent (Above a b) = ix2 (ra+rb) c
  where
    Z :. ra :. c  = extent a
    Z :. rb :. _c = extent b

extent (Beside a b) = ix2 r (ca+cb)
  where
    Z :. r  :. ca  = extent a
    Z :. _r :. cb = extent b

extent (Kron a b) = ix2 (m*p) (n*q)
  where
    Z :. m :. n = extent a
    Z :. p :. q = extent b

extent (DSum a b) = ix2 (m+p) (n+q)
  where
    Z :. m :. n = extent a
    Z :. p :. q = extent b

extent (Prod a b) = ix2 m q
  where
    Z :. m  :. _n = extent a
    Z :. _p :. q  = extent b

extent (Circ xs) = ix2 n n
  where
    n = length xs

extent (Skew xs) = ix2 n n
  where
    n = length xs

extent (Toep xs) = ix2 n n
  where
    n = (length xs + 1) `quot` 2

extent (Re a) = ix2 (2*m) (2*n)
  where
    Z :. m  :. n = extent a

extent F2 = ix2 2 2

extent (DFT n)  = ix2 n n
extent (DFT' n) = ix2 n n
extent (F n _)  = ix2 n n
extent (F' n _) = ix2 n n

-- | Transpose an SPL expression
transpose :: forall a . (Show a, Num a) => SPL a -> SPL a
-- transpose (E m)        = (matrix . A.transpose . unShowArray) m
transpose a@I{}        = a
transpose (T a)        = a
transpose (Inv a)      = Inv (transpose a)
transpose (Pi p)       = backpermute p
transpose (Rot alpha)  = Rot (-alpha)
transpose a@Diag{}     = a
transpose a@KDiag{}    = a
transpose (Above a b)  = Beside (transpose a) (transpose b)
transpose (Beside a b) = Above (transpose a) (transpose b)
transpose (Kron a b)   = Kron (transpose a) (transpose b)
transpose (DSum a b)   = DSum (transpose a) (transpose b)
transpose (Prod a b)   = Prod (transpose b) (transpose a)
transpose (Circ cs)    = circ (x:reverse xs)
                           where
                             x:xs = V.toList cs
transpose a@Skew{}     = a
transpose (Toep cs)    = Toep (V.reverse cs)
transpose (Re a)       = Re (transpose a)
transpose F2           = F2
transpose a@DFT{}      = a
transpose a@DFT'{}     = a
transpose a@F{}        = a
transpose a@F'{}       = a
transpose a            = T a

-- | Inverse of an SPL expression
inverse :: forall a . (Eq a, Fractional a) => SPL a -> SPL a
inverse a@I{}       = a
inverse (T a)       = T (inverse a)
inverse (Inv a)     = a
inverse (Pi p)      = backpermute p
inverse (Rot alpha) = Rot (-alpha)
inverse (Diag xs)   = Diag (fmap (1/) xs)
inverse (KDiag n k) = KDiag n (1/k)
inverse (Kron a b)  = Kron (inverse a) (inverse b)
inverse (Prod a b)  = Prod (inverse b) (inverse a)
inverse (Re a)      = Re (inverse a)
inverse F2          = F2
inverse (DFT n)     = DFT' n
inverse (DFT' n)    = DFT n
inverse (F n w)     = F' n w
inverse (F' n w)    = F n w
inverse a           = Inv a

-- | Compute the product, @A x@.
mXv :: forall r a .
       ( Num a
       , IArray r DIM1 a
       )
    => SPL a      -- ^ The SPL transform, @A@
    -> Vector r a -- ^ The vector, @x@
    -> Vector D a
mXv I{} x =
    A.delay x

mXv (Pi p) x =
    A.fromFunction (ix1 n) f
  where
    Z :. n = A.extent x

    f (Z :. i) = x A.! reindex i
      where
        reindex = toIdxMapping (invert p)

mXv (Diag d) x =
    A.fromVector d A..* x

mXv (KDiag n k) x =
    A.fromFunction (ix1 n) f
  where
    f (Z :. i) = k * x A.! i

mXv (Kron (I m) a) x =
    A.delay $
    A.concat [a #> A.slice (i*n') 1 n' x | i <- [0..m-1]]
  where
    Z :. _n :. n' = extent a

mXv (Kron a (I n)) x =
    permute (L (m*n) m) #> Kron (I n) a #> permute (L (m'*n) n) #> x
  where
    Z :. m :. m' = extent a

mXv (Kron a b) x =
    (I m ⊗ b) #> (a ⊗ I n') #> x
  where
    Z :. m :. _  = extent a
    Z :. _ :. n' = extent b

mXv (DSum a b) x =
    A.delay $
    (a #> A.slice 0 1 m x) A.++ (b #> A.slice m 1 n x)
  where
    Z :. _ :. m = extent a
    Z :. _ :. n = extent b

mXv (Prod a b) x =
    a #> b #> x

mXv a x =
    toMatrix a `A.mXv` x

-- | Compute the matrix-vector product, @A x@.
infixr 8 #>
(#>) :: forall r a .
        ( Num a
        , IArray r DIM1 a
        )
     => SPL a
     -> Vector r a
     -> Vector D a
(#>) = mXv

-- | Convert an SPL transform to an explicit matrix.
toMatrix :: forall e . Num e => SPL e -> Matrix M e
toMatrix (E a) =
    manifest (unShowArray a)

toMatrix (T e) = manifest $ A.transpose $ toMatrix e

toMatrix (Inv e) =
    case A.inverse $ toMatrix e of
      Nothing -> error "Matrix has no inverse"
      Just a  -> A.manifest a

toMatrix (Diag xs) =
    manifest $ A.fromFunction (ix2 n n) f
  where
    n = length xs

    f (Z :. i :. j) | i == j    = xs V.! i
                    | otherwise = 0

toMatrix (KDiag n e) =
    manifest $ A.fromFunction (ix2 n n) f
  where
    f (Z :. i :. j) | i == j    = e
                    | otherwise = 0

toMatrix (Above a b) =
    A.manifest $ toMatrix a A.<-> toMatrix b

toMatrix (Beside a b) =
    A.manifest $ toMatrix a A.<|> toMatrix b

toMatrix (Kron a b) =
    A.manifest $ A.kronecker (toMatrix a) (toMatrix b)

toMatrix (DSum a b) =
    A.manifest $ A.directSum (toMatrix a) (toMatrix b)

toMatrix (Prod a b) =
    A.manifest $ A.mXm (toMatrix a) (toMatrix b)

toMatrix (Circ xs) =
    manifest $ A.fromFunction (ix2 n n) f
  where
    n = length xs

    f (Z :. i :. j) = xs V.! ((i-j) `mod` n)

toMatrix (Skew xs) =
    manifest $ A.fromFunction (ix2 n n) f
  where
    n = length xs

    f (Z :. i :. j) = xs V.! ((i+j) `mod` n)

toMatrix (Toep xs) =
    manifest $ A.fromFunction (ix2 n n) f
  where
    n = (length xs + 1) `quot` 2

    f (Z :. i :. j) = xs V.! (i-j+n-1)

toMatrix (I n) =
    manifest $ A.fromFunction (ix2 n n) f
  where
    f (Z :. i :. j) | i == j    = 1
                    | otherwise = 0

toMatrix (Rot alpha) =
    A.matrix [[cos alpha, -(sin alpha)],
              [sin alpha, cos alpha]]

toMatrix (Pi p) =
    manifest $ A.fromFunction (ix2 (dim p) (dim p)) f
  where
    f (Z :. i :. j) | g j == i  = 1
                    | otherwise = 0

    g = toIdxMapping p

toMatrix (Re a) = manifest (RE (toMatrix a))

toMatrix F2 =
    A.matrix [[1,  1],
              [1, -1]]

toMatrix (DFT n)  = toMatrix (F n (omega n))
toMatrix (DFT' n) = toMatrix (F' n (omega n))

toMatrix (F n w) = manifest $ A.fromFunction (ix2 n n) f
  where
    f (Z :. i :. j) = w ^ (i*j)

toMatrix (F' n w) = toMatrix (KDiag n (1/fromIntegral n) × F n (1/w))

pprArgs :: Pretty a => [a] -> Doc
pprArgs = parens . commasep . map ppr

instance (Num e, Pretty e) => Pretty (SPL e) where
    pprPrec p (E a)        = pprPrec p (manifest (unShowArray a))
    pprPrec _ (I n)        = text "I_" <> ppr n
    pprPrec _ (T e)        = pprPrec 10 e <> text "^T"
    pprPrec _ (Inv e)      = pprPrec 10 e <> text "^-1"
    pprPrec _ (Pi p)       = ppr p
    pprPrec _ (Rot alpha)  = text "R_" <> ppr alpha
    pprPrec _ (Diag xs)    = text "diag" <> pprArgs (V.toList xs)
    pprPrec _ (KDiag n e)  = text "kdiag" <> parens (commasep [ppr n, ppr e])
    pprPrec p (Above a b)  = infixop p AboveOp a b
    pprPrec p (Beside a b) = infixop p BesideOp a b
    pprPrec p (Kron a b)   = infixop p KOp a b
    pprPrec p (DSum a b)   = infixop p DSOp a b
    pprPrec p (Prod a b)   = infixop p POp a b
    pprPrec _ (Circ xs)    = text "circ" <> pprArgs (V.toList xs)
    pprPrec _ (Skew xs)    = text "skew" <> pprArgs (V.toList xs)
    pprPrec _ (Toep xs)    = text "toep" <> pprArgs (V.toList xs)
    pprPrec _ (Re a)       = text "Re" <> parens (ppr a)
    pprPrec _ F2           = text "F_2"
    pprPrec _ (DFT n)      = text "DFT_" <> ppr n
    pprPrec _ (DFT' n)     = text "DFT'_" <> ppr n
    pprPrec _ (F n w)      = text "F_" <> ppr n <> parens (ppr w)
    pprPrec _ (F' n w)     = text "F'_" <> ppr n <> parens (ppr w)

data MatrixBinop = AboveOp
                 | BesideOp
                 | KOp
                 | DSOp
                 | POp
  deriving (Eq, Ord, Show)

instance HasFixity MatrixBinop where
    fixity AboveOp  = infixl_ 7
    fixity BesideOp = infixl_ 7
    fixity KOp      = infixl_ 7
    fixity DSOp     = infixl_ 6
    fixity POp      = infixl_ 7

instance Pretty MatrixBinop where
    ppr AboveOp  = text "<->"
    ppr BesideOp = text "<|>"
    ppr KOp      = char '⊗'
    ppr DSOp     = char '⊕'
    ppr POp      = char '×'

-- | Extract a row of a matrix
row :: Matrix M e -> Int -> V.Vector e
row e i = V.slice (i*n) n xs
  where
    M (Z :. _m :. n) xs = manifest e

-- | Extract a column of a matrix
col :: forall e . Matrix M e -> Int -> V.Vector e
col e j = V.generate m f
  where
    M (Z :. m :. n) xs = manifest e

    f :: Int -> e
    f i = xs V.! (i*n + j)

-- | Vertical stacking of transforms
infixr 7 <->
(<->) :: SPL e -> SPL e -> SPL e
(<->) = Above

-- | Horizontal stacking of transforms
infixr 7 <|>
(<|>) :: SPL e -> SPL e -> SPL e
(<|>) = Beside

-- | Create a block matrix
block :: SPL e -> SPL e -> SPL e -> SPL e -> SPL e
block tl tr bl br = (tl <|> tr) <-> (bl <|> br)

-- | Alias for Kronecker product.
infixl 7 ⊗
(⊗) :: SPL e -> SPL e -> SPL e
I m ⊗ (Kron (I n) y) = I (m*n) ⊗ y
I 1 ⊗ y              = y
x   ⊗ I 1            = x
x   ⊗ y              = Kron x y

-- | Alias for matrix direct sum.
infixl 6 ⊕
(⊕) :: SPL e -> SPL e -> SPL e
I m     ⊕ I n     = I (n + m)
Diag xs ⊕ Diag ys = Diag (xs <> ys)
x       ⊕ y       = DSum x y

-- | Alias for matrix product.
infixl 7 ×
(×) :: SPL e -> SPL e -> SPL e
(×) = Prod


#!/usr/bin/env stack
-- stack --install-ghc runghc --resolver lts-11.9 --package backprop-0.2.2.0 --package random --package hmatrix-backprop-0.1.2.1 --package statistics --package lens --package one-liner-instances --package split --package ghc-typelits-natnormalise --package ghc-typelits-knownnat --package hmatrix --package hmatrix-vector-sized --package microlens --package vector-sized --package transformers --package type-combinators -- -O2

-- | Replace the second line with this one to have "./model.hs" open a ghci session
-- stack --install-ghc ghci --resolver lts-11.9 --package backprop-0.2.2.0 --package random --package hmatrix-backprop-0.1.2.1 --package statistics --package lens --package one-liner-instances --package split --package ghc-typelits-natnormalise --package ghc-typelits-knownnat --package hmatrix --package hmatrix-vector-sized --package microlens --package vector-sized --package transformers --package type-combinators

{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE DeriveGeneric                            #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE FlexibleInstances                        #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE LambdaCase                               #-}
{-# LANGUAGE MultiParamTypeClasses                    #-}
{-# LANGUAGE PartialTypeSignatures                    #-}
{-# LANGUAGE PatternSynonyms                          #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeInType                               #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE ViewPatterns                             #-}
{-# OPTIONS_GHC -fno-warn-orphans                     #-}
{-# OPTIONS_GHC -fno-warn-partial-type-signatures     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}
{-# OPTIONS_GHC -fwarn-redundant-constraints          #-}

import           Control.Monad.Trans.State
import           Data.Bifunctor
import           Data.Foldable
import           Data.List
import           Data.List.Split
import           Data.Tuple
import           Data.Type.Option
import           GHC.Generics                          (Generic)
import           GHC.TypeNats
import           Lens.Micro hiding                     ((&))
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import           Numeric.LinearAlgebra.Static.Vector
import           Numeric.OneLiner
import           System.Random
import qualified Data.Vector.Storable.Sized            as SVS
import qualified Numeric.LinearAlgebra                 as HU
import qualified Numeric.LinearAlgebra.Static          as H
import qualified Prelude.Backprop                      as B

data a :& b = !a :& !b
  deriving (Show, Generic)
infixr 2 :&

type Model p a b = forall z. Reifies z W
                => BVar z p
                -> BVar z a
                -> BVar z b

linReg :: Model (Double :& Double) Double Double
linReg (a :&& b) x = b * x + a

squaredErrorGrad
    :: (Backprop p, Backprop b, Num b)
    => Model p a b      -- ^ Model
    -> a                -- ^ Observed input
    -> b                -- ^ Observed output
    -> p                -- ^ Parameter guess
    -> p                -- ^ Gradient
squaredErrorGrad f x targ = gradBP $ \p ->
    (f p (auto x) - auto targ) ^ 2

trainModel
    :: (Fractional p, Backprop p, Num b, Backprop b)
    => Model p a b      -- ^ model to train
    -> p                -- ^ initial parameter guess
    -> [(a,b)]          -- ^ list of observations
    -> p                -- ^ updated parameter guess
trainModel f = foldl' $ \p (x,y) -> p - 0.1 * squaredErrorGrad f x y p

trainModelIO
    :: (Fractional p, Backprop p, Num b, Backprop b, Random p)
    => Model p a b      -- ^ model to train
    -> [(a,b)]          -- ^ list of observations
    -> IO p             -- ^ parameter guess
trainModelIO m xs = do
    p0 <- (/ 10) . subtract 0.5 <$> randomIO
    return $ trainModel m p0 xs

testTrainLinReg :: IO (Double :& Double)
testTrainLinReg = trainModelIO linReg (concat (replicate 1000 samps))
  where
    samps = [(1,1),(2,3),(3,5),(4,7),(5,9)]

logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))

feedForward
    :: (KnownNat i, KnownNat o)
    => Model (L o i :& R o) (R i) (R o)
feedForward (w :&& b) x = w #> x + b

feedForwardLog
    :: (KnownNat i, KnownNat o)
    => Model (L o i :& R o) (R i) (R o)
feedForwardLog (w :&& b) x = logistic (w #> x + b)

testTrainPerceptron :: IO [R 1]
testTrainPerceptron = do
    trained <- trainModelIO feedForwardLog $ take 10000 (cycle samps)
    print trained
    return [ evalBP2 feedForwardLog trained r | (r, _) <- samps ]
  where
    samps = [ (H.vec2 0 0, 0)
            , (H.vec2 1 0, 0)
            , (H.vec2 0 1, 0)
            , (H.vec2 1 1, 1)
            ]

logReg :: Model (Double :& Double) Double Double
logReg ab = logistic . linReg ab

feedForwardLog'
    :: (KnownNat i, KnownNat o)
    => Model (L o i :& R o) (R i) (R o)
feedForwardLog' wb = logistic . feedForward wb

softMax :: (Reifies z W, KnownNat n) => BVar z (R n) -> BVar z (R n)
softMax x = konst (1 / sumElements expx) * expx
  where
    expx = exp x

feedForwardSoftMax
    :: (KnownNat i, KnownNat o)
    => Model (L o i :& R o) (R i) (R o)
feedForwardSoftMax wb = softMax . feedForward wb

(<~)
    :: (Backprop p, Backprop q)
    => Model  p       b c
    -> Model       q  a b
    -> Model (p :& q) a c
(f <~ g) (p :&& q) = f p . g q
infixr 8 <~

testTrainTwoLayer :: IO [R 1]
testTrainTwoLayer = do
    trained <- trainModelIO model (take 10000 (cycle samps))
    print trained
    return [ evalBP2 model trained r | (r, _) <- samps ]
  where
    model :: Model _ (R 2) (R 1)
    model = feedForwardLog' @4 @1 <~ feedForwardLog' @2 @4
    samps = [ (H.vec2 0 0, 0)
            , (H.vec2 1 0, 1)
            , (H.vec2 0 1, 1)
            , (H.vec2 1 1, 0)
            ]

type ModelS p s a b = forall z. Reifies z W
                   => BVar z p
                   -> BVar z a
                   -> BVar z s
                   -> (BVar z b, BVar z s)

ar2 :: ModelS (Double :& (Double :& Double)) Double Double Double
ar2 (c :&& (φ1 :&& φ2)) yLast yLastLast =
    ( c + φ1 * yLast + φ2 * yLastLast, yLast )

fcrnn
    :: (KnownNat i, KnownNat o)
    => ModelS ((L o i :& L o o) :& R o) (R o) (R i) (R o)
fcrnn ((wX :&& wS) :&& b) x s = ( y, logistic y )
  where
    y  = (wX #> x) + (wS #> s) + b

(<*~*)
  :: (Backprop p, Backprop q, Backprop s, Backprop t)
    => ModelS  p        s       b c
    -> ModelS       q        t  a b
    -> ModelS (p :& q) (s :& t) a c
(f <*~* g) (p :&& q) x (s :&& t) = (z, s' :&& t')
  where
    (y, t') = g q x t
    (z, s') = f p y s
infixr 8 <*~*

mapS
    :: (forall z. Reifies z W => BVar z b -> BVar z c)
    -> ModelS p s a b
    -> ModelS p s a c
mapS f g p x = first f . g p x

toS :: Model  p   a b
    -> ModelS p s a b
toS f p x s = (f p x, s)

(<*~)
  :: (Backprop p, Backprop q)
    => Model   p         b c
    -> ModelS       q  s a b
    -> ModelS (p :& q) s a c
(f <*~ g) (p :&& q) x = first (f p) . g q x
infixr 8 <*~

unroll
    :: (Backprop a, Backprop b)
    => ModelS p s  a   b
    -> ModelS p s [a] [b]
unroll f p xs s0 = swap $ B.mapAccumL f' s0 xs
  where
    -- we have to re-arrange the order of arguments and tuple a bit to
    -- match what `mapAccumL` expects
    f' s x = swap (f p x s)

unrollLast
    :: (Backprop a, Backprop b)
    => ModelS p s  a  b
    -> ModelS p s [a] b
unrollLast f = mapS (last . sequenceVar) (unroll f)

unrollLast'
    :: Backprop a
    => ModelS p s  a  b
    -> ModelS p s [a] b
unrollLast' f p xs s0 = foldl' go (undefined, s0) (sequenceVar xs)
  where
    go (_, s) x = f p x s

fixState
    :: s
    -> ModelS p s a b
    -> Model  p   a b
fixState s0 f p x = fst $ f p x (auto s0)

zeroState
    :: Num s
    => ModelS p s a b
    -> Model  p   a b
zeroState = fixState 0

trainState
    :: (Backprop p, Backprop s)
    => ModelS  p    s  a b
    -> Model  (p :& s) a b
trainState f (p :&& s) x = fst $ f p x s

prime
    :: Foldable t
    => ModelS p s a b     -- ^ model
    -> p                  -- ^ parameterization
    -> s                  -- ^ initial state
    -> t a                -- ^ priming input
    -> s                  -- ^ primed state
prime f p = foldl' $ evalBP2 (\s x -> snd $ f (auto p) x s)

feedback
    :: (Backprop a, Backprop s)
    => ModelS p s a a     -- ^ model
    -> p                  -- ^ parameterization
    -> s                  -- ^ initial state
    -> a                  -- ^ initial input
    -> [a]                -- ^ inifinite feedback loop
feedback f p s0 x0 = unfoldr go (x0, s0)
  where
    go (x, s) = Just (x, (y, s'))
      where
        -- 'T2' tuples up a pair of 'BVar's into a 'BVar' of a tuple
        (y, s') = evalBP (uncurry T2 . f (auto p) (auto x)) s

testAR2 :: IO [Double]
testAR2 = do
    trained <- trainModelIO model $ take 10000 samps
    let primed = prime    model0 trained 0      (take 19 series)
        output = feedback model0 trained primed (series !! 19)
    print trained
    return $ take 200 output
  where
    -- sine wave with period 25
    series :: [Double]
    series = [ sin (2 * pi * t / 25) | t <- [0..]              ]
    samps  = [ (init c, last c)      | c <- chunksOf 19 series ]
    model0 :: ModelS _ _ Double Double
    model0 = ar2
    model  :: Model  _   [Double] Double
    model  = zeroState $ unrollLast model0

testRNN :: IO [R 1]
testRNN = do
    trained <- trainModelIO model $ take 100000 samps
    let primed = prime   model0  trained 0      (take 19 series)
        output = feedback model0 trained primed (series !! 19)
    return $ take 200 output
  where
    -- sine wave with period 25
    series :: [H.R 1]
    series = [ H.konst (sin (2 * pi * t / 25)) | t <- [0..]              ]
    samps  = [ (init c, last c)                | c <- chunksOf 19 series ]
    model0 :: ModelS _ _ (R 1) (R 1)
    model0 = feedForward @30 @1
         <*~ mapS logistic (fcrnn @1 @30)
    model  :: Model  _   [R 1] (R 1)
    model  = zeroState $ unrollLast model0

recurrently
    :: (Backprop a, Backprop b)
    => Model  p   (a :& b) b
    -> ModelS p b  a       b
recurrently f p x yLast = (y, y)
  where
    y = f p (x :&& yLast)

recurrentlyWith
    :: (Backprop a, Backprop b)
    => (forall z. Reifies z W => BVar z c -> BVar z b)
    -> Model  p   (a :& b) c
    -> ModelS p b  a       c
recurrentlyWith store f p x yLast = (y, store y)
  where
    y = f p (x :&& yLast)

ffOnSplit
    :: forall i o. (KnownNat i, KnownNat o)
    => Model _ (R i :& R o) (R o)
ffOnSplit p (rI :&& rO) = feedForward p (rI # rO)

fcrnn'
    :: (KnownNat i, KnownNat o)
    => ModelS _ (R o) (R i) (R o)
fcrnn' = recurrentlyWith logistic (\p -> feedForward p . uncurryT (#))

lagged
    :: (KnownNat n, 1 <= n)
    => Model  p       (R (n + 1)) b
    -> ModelS p (R n) Double      b
lagged f p x xLasts = (y, xLasts')
  where
    fullLasts    = xLasts & x
    y            = f p fullLasts
    (_, xLasts') = headTail fullLasts

ar :: (KnownNat n, 1 <= n)
   => ModelS _ (R n) Double Double
ar = lagged (\p -> fst . headTail . feedForward @_ @1 p)

ar2' :: ModelS _ (R 2) Double Double
ar2' = ar @2

main :: IO ()
main = do
    putStrLn "Linear regression"
    print =<< testTrainLinReg

    putStrLn "Single layer perceptron learning AND"
    mapM_ print =<< testTrainPerceptron

    putStrLn "Two-layer ANN learning XOR"
    mapM_ print =<< testTrainTwoLayer

    putStrLn "Sine (AR2)"
    ar2Test <- testAR2
    mapM_ print (take 30 ar2Test)
    writeFile "ar2sin.dat" $ unlines (show <$> ar2Test)

    putStrLn "Sine (RNN)"
    rnnTest <- testRNN
    mapM_ print (take 30 rnnTest)
    writeFile "rnnsin.dat" $ unlines (show . HU.sumElements . H.extract <$> rnnTest)

pattern (:&&) :: (Backprop a, Backprop b, Reifies z W)
              => BVar z a -> BVar z b -> BVar z (a :& b)
pattern x :&& y <- (\xy -> (xy ^^. t1, xy ^^. t2)->(x, y))
  where
    (:&&) = isoVar2 (:&) (\case x :& y -> (x, y))
{-# COMPLETE (:&&) #-}

t1 :: Lens (a :& b) (a' :& b) a a'
t1 f (x :& y) = (:& y) <$> f x

t2 :: Lens (a :& b) (a :& b') b b'
t2 f (x :& y) = (x :&) <$> f y

instance (Num a, Num b) => Num (a :& b) where
    (+)         = gPlus
    (-)         = gMinus
    (*)         = gTimes
    negate      = gNegate
    abs         = gAbs
    signum      = gSignum
    fromInteger = gFromInteger

instance (Fractional a, Fractional b) => Fractional (a :& b) where
    (/) = gDivide
    recip = gRecip
    fromRational = gFromRational

instance (Random a, Random b) => Random (a :& b) where
    random g0 = (x :& y, g2)
      where
        (x, g1) = random g0
        (y, g2) = random g1
    randomR (x0 :& y0, x1 :& y1) g0 = (x :& y, g2)
      where
        (x, g1) = randomR (x0, x1) g0
        (y, g2) = randomR (y0, y1) g1

instance (Backprop a, Backprop b) => Backprop (a :& b)

uncurryT
    :: (Backprop a, Backprop b, Reifies z W)
    => (BVar z a -> BVar z b -> BVar z c)
    -> BVar z (a :& b)
    -> BVar z c
uncurryT f x = f (x ^^. t1) (x ^^. t2)

instance (KnownNat n, KnownNat m) => Random (L n m) where
    random = runState . fmap vecL $ SVS.replicateM (state random)
    randomR (xs,ys) = runState . fmap vecL $ SVS.zipWithM (curry (state . randomR))
        (lVec xs) (lVec ys)

instance (KnownNat n) => Random (R n) where
    random = runState $ vecR <$> SVS.replicateM (state random)
    randomR (xs,ys) = runState . fmap vecR $ SVS.zipWithM (curry (state . randomR))
        (rVec xs) (rVec ys)


import Data.Complex
import Data.Either
import Data.Maybe
--import System.Random
import Text.Read

main = do
  putStrLn "Compiled."
  center defaultGD ["0", "0"]

--Datatypes/typeclasses/etc.

data Vector2 = V2 Double Double deriving (Show, Read, Eq)

data Transformation = Transformation Matrix Vector2 deriving (Show, Read, Eq)

data GlobalData = GD
  { getWindowSize :: (Int, Int),
    getTickSizes :: (Double, Double),
    getPlotLoc :: (Vector2, Vector2),
    getIters :: Int,
    getAxesShown :: Bool,
    getVectToChar :: VectToChar
  }

type Matrix = (Double, Double, Double, Double)

type Args = [String]

type StrCommand = String

type ArgDescription = String

type CmdDescription = String

type Command = GlobalData -> Args -> IO ()

type PlotCommand = GlobalData -> Args -> Either (IO ()) VectToChar

type VectToChar = GlobalData -> Vector2 -> Char

--Constants

defaultGD :: GlobalData
defaultGD = GD (100, 30) (1.0, 1.0) (V2 (-1.5) (-1), V2 1.5 1) 50 True vtcMBSet

charWidth :: Double
charWidth = 0.3

cursorChar = 'o'

originChar = '+'

hAxisMinChar = '<'

hAxisMaxChar = '>'

hAxisChar = '-'

hAxisTickChar = '|'

vAxisMinChar = 'V'

vAxisMaxChar = '^'

vAxisChar = '|'

vAxisTickChar = '~'

maxShadeChar = '#'

minShadeChar = ' '

vBounderyChar = '|'

hBounderyChar = '~'

shadeChars :: [Char]
shadeChars = " `.*+&$@#"

commandList :: [(StrCommand, (Command, ArgDescription, CmdDescription))]
commandList =
  [ ("draw", (draw, "draw", "Redraws the plot to the screen.")),
    ("plot", (plot, "plot (args)", "Sets the function defining which points to plot.\nPossible arguments:\n" ++ show (map (sndTriple . snd) plotList))),
    ( "zoom",
      ( zoom,
        "zoom [Double c]",
        "Zooms in by the specified factor."
          ++ "Requires c > 0. Omitting arguments sets zoom level."
      )
    ),
    ("shift", (shift, "shift (Double dx) (Double dy)", "Translates the plot by the specified amount.")),
    ("center", (center, "center (Double x) (Double y)", "Centers the plot on the point (x,y). Defaults to the origin.")),
    ( "resize",
      ( resize,
        "resize (Int width) (Int height)",
        "Resizes the window. Omitting arguments defaults to 100x30."
          ++ "\nRequires width, height > 0. Tag -square ignores second argument and calculates height based on given width."
      )
    ),
    ("toggleAxes", (toggleAxes, "toggleAxes", "Toggles the visibility of the axes.")),
    ( "cutoff",
      ( cutoff,
        "cutoff (Int n)",
        "Sets the cutoff point for iterations. Generally, a higher cutoff point means greater precision.\n"
          ++ "Requires n > 0. Omitting arguments defaults to 50."
      )
    ),
    ("reload", (reload, "reload", "Restores all data (center, scale factor, etc.) to their default values.")),
    ("info", (info, "info", "Prints information about the current plot.")),
    ("quit", (quit, "quit", "Ends execution. Current data will be lost.")),
    ("help", (help, "help [command]", "Prints a description of the specified command and its arguments. Omitting the arguments lists available commands."))
  ]

plotList :: [(StrCommand, (PlotCommand, ArgDescription, CmdDescription))]
plotList =
  [ ("mandelbrot", (plotMBSet, "plot mandelbrot", "Plots the Mandelbrot set over the complex plane.\nTag -shade adds shading based on time to diverge.")),
    ("sierpinski", (plotSierpinski, "plot sierpinski", "Work in progress. Plots the Sierpinski triangle within the unit square, calculated as an iterated function system.\n")),
    ( "juliaset",
      ( plotJuliaSet,
        "plot juliaset (Complex c)",
        "Given a complex value (using a:+b notation), plots the filled julia set of z^2+c over the complex plane.\n"
          ++ "Omitting arguments defaults to 0:+0. Tag -shade adds shading based on time to diverge.\n"
      )
    )
  ]

--Triple commands

fstTriple (x, _, _) = x

sndTriple (_, y, _) = y

thrdTriple (_, _, z) = z

--IO commands

getCommand :: Command
getCommand gd _ = do
  putStr "\nMBSet>> "
  line <- getLine
  if line == ""
    then do getCommand gd []
    else do
      let (strCmd : args) = words line
      let mCommand = lookup strCmd commandList
      case mCommand of
        Nothing -> rejectCommand gd [strCmd]
        Just (command, argDesc, cmdDesc) -> command gd args
      return ()

rejectCommand :: Command
rejectCommand gd [strCmd] = do
  putStrLn $ "No such command: " ++ strCmd
  putStrLn "Use the command \"help\" to see a list of available commands."
  getCommand gd []
rejectCommand gd args = do
  putStrLn "No such command.\nUse the command \"help\" to see a list of available commands."
  getCommand gd []

rejectArgs :: Command
rejectArgs gd args = do
  putStrLn $ "Invalid arguments: " ++ unwords args ++ "\nUse the command \"help [command]\" for command syntax."
  getCommand gd []

--Input Commands

draw :: Command
draw gd _ = do
  putStr . unlines . asciiArray $ gd
  getCommand gd []

plot :: Command
plot gd args@(strPlot : args') = do
  let mPlotTriple = lookup strPlot plotList
  if isNothing mPlotTriple
    then do
      rejectArgs gd args
    else do
      let Just (plotCmd, _, _) = mPlotTriple
      if isLeft $ plotCmd gd args'
        then do
          let Left failureAction = plotCmd gd args
          failureAction ---------
        else do
          let Right command = plotCmd gd args'
          draw (setVectToChar gd command) [""]
plot gd args = help gd $ "plot" : args

zoom :: Command
zoom gd [] = draw (setPlotLoc gd $ getPlotLoc defaultGD) []
zoom gd args
  | head args == "in" = zoom gd $ "2" : tail args
  | head args == "out" = zoom gd $ "0.5" : tail args
  | otherwise = do
    let mDouble = readMaybe (head args) :: Maybe Double
    case mDouble of
      Nothing -> rejectArgs gd args
      Just x ->
        if x <= 0
          then rejectArgs gd args
          else draw (zoomGD gd x) $ tail args
    return ()

shift :: Command
shift gd@GD {getPlotLoc = (v1, v2)} args@(dxStr : dyStr : args')
  | isNothing (readMaybe dxStr :: Maybe Double) = rejectArgs gd args
  | isNothing (readMaybe dyStr :: Maybe Double) = rejectArgs gd args
  | otherwise = draw (setPlotLoc gd (v1 `vPlus` d, v2 `vPlus` d)) args'
  where
    d = V2 (read dxStr :: Double) (read dyStr :: Double)
shift gd args = rejectArgs gd args

center :: Command
center gd [] = center gd ["0", "0"]
center gd@GD {getPlotLoc = (v1, v2)} args@(xStr : yStr : args')
  | isNothing (readMaybe xStr :: Maybe Double) = rejectArgs gd args
  | isNothing (readMaybe yStr :: Maybe Double) = rejectArgs gd args
  | otherwise = shift gd $ [show dx, show dy] ++ args'
  where
    (V2 dx dy) = newCenter `vMinus` currentCenter
    currentCenter = midPt v1 v2
    newCenter = V2 (read xStr :: Double) (read yStr :: Double)
center gd args = rejectArgs gd args

resize :: Command
resize gd [] = draw (setWindowSize gd $ getWindowSize defaultGD) []
resize gd args@(wStr : hStr : args')
  | isNothing (readMaybe wStr :: Maybe Int) = rejectArgs gd args
  | width <= 0 = rejectArgs gd args
  | "square" `elem` args = resize gd sqrArgs
  | isNothing (readMaybe hStr :: Maybe Int) = rejectArgs gd args
  | height <= 0 = rejectArgs gd args
  | otherwise = draw (setWindowSize gd (width, height)) args'
  where
    width = read wStr :: Int
    height = read hStr :: Int
    sqrHeight = round . (* charWidth) . fromIntegral $ width
    sqrArgs = wStr : show sqrHeight : filter (/= "square") args'
resize gd args = rejectArgs gd args

cutoff :: Command
cutoff gd [] = draw (setIters gd $ getIters defaultGD) []
cutoff gd args@(strIters : args')
  | isNothing (readMaybe strIters :: Maybe Int) = rejectArgs gd args
  | iters <= 0 = rejectArgs gd args
  | otherwise = draw (setIters gd iters) args'
  where
    iters = read strIters :: Int

toggleAxes :: Command
toggleAxes gd@GD {getAxesShown = axesShown} = draw (setAxesShown gd (not axesShown))

reload :: Command
reload _ = draw defaultGD

info :: Command
info gd args = do
  putStrLn "\n  --  Current plot information  --"
  putStrLn $ "Window size: " ++ (show . fst . getWindowSize $ gd) ++ "x" ++ (show . snd . getWindowSize $ gd)
  let (xMin, yMin) = (getX . fst . getPlotLoc $ gd, getY . fst . getPlotLoc $ gd)
  let (xMax, yMax) = (getX . snd . getPlotLoc $ gd, getY . snd . getPlotLoc $ gd)
  putStrLn $ "Plot location (bounds): " ++ show xMin ++ " < x < " ++ show xMax ++ ",  " ++ show yMin ++ " < y < " ++ show yMax
  putStrLn $ "Cutoff point (number of iterations): " ++ show (getIters gd)
  putStrLn $ "Axes shown: " ++ show (getAxesShown gd)
  getCommand gd []

quit :: Command
quit gd args = do
  putStrLn "Quitting.\n\n"
  return ()

help :: Command
help gd ("plot" : strPlot : args) = do
  let output = case lookup strPlot plotList of
        Just (f, argDesc, cmdDesc) -> argDesc ++ " -- " ++ cmdDesc
        Nothing -> "No such plot: " ++ strPlot
  putStrLn output
  getCommand gd []
help gd [] = do
  putStrLn "Available commands: "
  print $ map fst commandList
  getCommand gd []
help gd args = do
  let strCmd = head args
  let output = case lookup strCmd commandList of
        Just (f, argDesc, cmdDesc) -> argDesc ++ " -- " ++ cmdDesc
        Nothing -> "No such command: " ++ strCmd
  putStrLn output
  getCommand gd []

--GlobalData functions

setWindowSize :: GlobalData -> (Int, Int) -> GlobalData
setWindowSize (GD _ tickSizes plotLoc iters axesShown vectToChar) windowSize = GD windowSize tickSizes plotLoc iters axesShown vectToChar

setTickSizes :: GlobalData -> (Double, Double) -> GlobalData
setTickSizes (GD windowSize _ plotLoc iters axesShown vectToChar) tickSizes = GD windowSize tickSizes plotLoc iters axesShown vectToChar

setPlotLoc :: GlobalData -> (Vector2, Vector2) -> GlobalData
setPlotLoc (GD windowSize tickSizes _ iters axesShown vectToChar) plotLoc = GD windowSize tickSizes plotLoc iters axesShown vectToChar

setIters :: GlobalData -> Int -> GlobalData
setIters (GD windowSize tickSizes plotLoc _ axesShown vectToChar) iters = GD windowSize tickSizes plotLoc iters axesShown vectToChar

setAxesShown :: GlobalData -> Bool -> GlobalData
setAxesShown (GD windowSize tickSizes plotLoc iters _ vectToChar) axesShown = GD windowSize tickSizes plotLoc iters axesShown vectToChar

setVectToChar :: GlobalData -> VectToChar -> GlobalData
setVectToChar (GD windowSize tickSizes plotLoc iters axesShown _) = GD windowSize tickSizes plotLoc iters axesShown

pxlW :: GlobalData -> Double
pxlW gd = (getX vMax - getX vMin) / fromIntegral screenWidth
  where
    (screenWidth, _) = getWindowSize gd
    (vMin, vMax) = getPlotLoc gd

pxlH :: GlobalData -> Double
pxlH gd = (getY vMax - getY vMin) / fromIntegral screenHeight
  where
    (_, screenHeight) = getWindowSize gd
    (vMin, vMax) = getPlotLoc gd

zoomGD :: GlobalData -> Double -> GlobalData
zoomGD gd@GD {getPlotLoc = (v1, v2)} c = setPlotLoc gd (v1', v2')
  where
    v1' = vPlus (midPt v1 v2) . vScale (1 / c) $ v1 `vMinus` midPt v1 v2
    v2' = vPlus (midPt v1 v2) . vScale (1 / c) $ v2 `vMinus` midPt v1 v2

--Vector2 functions

getX :: Vector2 -> Double
getX (V2 x _) = x

getY :: Vector2 -> Double
getY (V2 _ y) = y

mag :: Vector2 -> Double
mag (V2 x y) = sqrt (x ^ 2 + y ^ 2)

vPlus :: Vector2 -> Vector2 -> Vector2
vPlus (V2 a b) (V2 x y) = V2 (a + x) (b + y)

vMinus :: Vector2 -> Vector2 -> Vector2
vMinus (V2 a b) (V2 x y) = V2 (a - x) (b - y)

vScale :: Double -> Vector2 -> Vector2
vScale c (V2 x y) = V2 (x * c) (y * c)

midPt :: Vector2 -> Vector2 -> Vector2
midPt v1 v2 = vScale 0.5 (v1 `vPlus` v2)

applyTransformation :: Transformation -> Vector2 -> Vector2
applyTransformation (Transformation (a11, a21, a12, a22) (V2 bx by)) v@(V2 x y) =
  V2 (a11 * x + a12 * y + bx) (a21 * x + a22 * y + by)

-- Printing

hRange :: GlobalData -> [Double]
hRange gd = [xMin + i * dx | i <- [0 .. width]]
  where
    dx = (xMax - xMin) / width
    width = fromIntegral . fst . getWindowSize $ gd
    xMin = getX . fst . getPlotLoc $ gd
    xMax = getX . snd . getPlotLoc $ gd

vectRow :: GlobalData -> Double -> [Vector2]
vectRow gd y = map (`V2` y) $ hRange gd

vRange :: GlobalData -> [Double]
vRange gd = reverse [yMin + i * dy | i <- [0 .. height]]
  where
    dy = (yMax - yMin) / height
    height = fromIntegral . snd . getWindowSize $ gd
    yMin = getY . fst . getPlotLoc $ gd
    yMax = getY . snd . getPlotLoc $ gd

vectArray :: GlobalData -> [[Vector2]]
vectArray gd = map (vectRow gd) $ vRange gd

asciiArray :: GlobalData -> [[Char]]
asciiArray gd = addBoundery . map (map (vectToAscii gd)) $ vectArray gd

addBoundery :: [[Char]] -> [[Char]]
addBoundery array = [hBoundery] ++ addVBoundery array ++ [hBoundery]
  where
    hBoundery = minShadeChar : replicate (length . head $ array) hBounderyChar ++ [minShadeChar]
    addVBoundery array = map (\row -> vBounderyChar : row ++ [vBounderyChar]) array

vectToAscii :: GlobalData -> Vector2 -> Char
vectToAscii gd v@(V2 x y)
  | onCursor = cursorChar
  | onXAxis && onYAxis = originChar
  -- x-axis
  | onXAxis && abs (x - xMin) <= pxlW gd / 2 = hAxisMinChar
  | onXAxis && abs (x - xMax) <= pxlW gd / 2 = hAxisMaxChar
  | onXTick = hAxisTickChar
  | onXAxis = hAxisChar
  -- y-axis
  | onYAxis && abs (y - yMin) <= pxlH gd / 2 = vAxisMinChar
  | onYAxis && abs (y - yMax) <= pxlH gd / 2 = vAxisMaxChar
  | onYTick = vAxisTickChar
  | onYAxis = vAxisChar
  | otherwise = getVectToChar gd gd v
  where
    (scrnW, scrnH) = getWindowSize gd
    (vMin, vMax) = getPlotLoc gd
    onXAxis = getAxesShown gd && abs y < pxlH gd / 2
    onYAxis = getAxesShown gd && abs x < pxlW gd / 2
    onXTick = onXAxis && abs (x - rndX) < pxlW gd / 2 --WIP
    onYTick = onYAxis && abs (y - rndY) < pxlH gd / 2 --WIP
    [rndX, rndY] = map (fromIntegral . round) [x, y] --WIP rewrite as nearest tick
    (xMin, yMin) = (getX vMin, getY vMin)
    (xMax, yMax) = (getX vMax, getY vMax)
    onCursor = mag (v `vMinus` midPt vMin vMax) <= pxlW gd / 2

--Mandelbrot set

plotMBSet :: PlotCommand
plotMBSet gd args
  | "-shade" `elem` args = Right vtcMBSetShaded
  | otherwise = Right vtcMBSet

vtcMBSet :: VectToChar
vtcMBSet gd@GD {getIters = iters} (V2 x y)
  | isInSet = maxShadeChar
  | otherwise = minShadeChar
  where
    isInSet = all (\z -> magnitude z < 2) . take iters $ mSequence c (0 :+ 0)
    c = x :+ y
    mSequence c z = z : mSequence c (mFunc c z)
    mFunc c z = z ^ 2 + c

vtcMBSetShaded :: VectToChar
vtcMBSetShaded gd@GD {getIters = maxIters} (V2 x y)
  | itersToEscape >= maxIters = maxShadeChar
  | itersToEscape >= (49 * maxIters) `div` 8 = shadeChars !! 7
  | itersToEscape >= (36 * maxIters) `div` 8 = shadeChars !! 6
  | itersToEscape >= (25 * maxIters) `div` 8 = shadeChars !! 5
  | itersToEscape >= (16 * maxIters) `div` 8 = shadeChars !! 4
  | itersToEscape >= (9 * maxIters) `div` 8 = shadeChars !! 3
  | itersToEscape >= (4 * maxIters) `div` 8 = shadeChars !! 2
  | itersToEscape >= (1 * maxIters) `div` 8 = shadeChars !! 1
  | otherwise = minShadeChar
  where
    itersToEscape = length . takeWhile (\z -> magnitude z < 2) . take maxIters $ mSequence c (0 :+ 0)
    c = x :+ y
    mSequence c z = z : mSequence c (z ^ 2 + c)

--Julia Sets

plotJuliaSet :: PlotCommand
plotJuliaSet gd args@(cStr : args')
  | isNothing mc = Left (rejectArgs gd args)
  | "-shade" `elem` args = Right (vtcJuliaSetShaded c)
  | otherwise = Right (vtcJuliaSet c)
  where
    mc = readMaybe cStr :: Maybe (Complex Double)
    Just c = mc
plotJuliaSet gd args = plotJuliaSet gd ("0" : "0" : args)

vtcJuliaSet :: Complex Double -> VectToChar
vtcJuliaSet c gd@GD {getIters = iters} (V2 x y)
  | isInSet = maxShadeChar
  | otherwise = minShadeChar
  where
    isInSet = all (\z -> magnitude z < 2) . take iters $ jSequence c (x :+ y)
    jSequence c z = z : jSequence c (z ^ 2 + c)

vtcJuliaSetShaded :: Complex Double -> VectToChar
vtcJuliaSetShaded c gd@GD {getIters = maxIters} (V2 x y)
  | itersToEscape >= maxIters = maxShadeChar
  | itersToEscape >= (7 * maxIters) `div` 8 = shadeChars !! 7
  | itersToEscape >= (6 * maxIters) `div` 8 = shadeChars !! 6
  | itersToEscape >= (5 * maxIters) `div` 8 = shadeChars !! 5
  | itersToEscape >= (4 * maxIters) `div` 8 = shadeChars !! 4
  | itersToEscape >= (3 * maxIters) `div` 8 = shadeChars !! 3
  | itersToEscape >= (2 * maxIters) `div` 8 = shadeChars !! 2
  | itersToEscape >= (1 * maxIters) `div` 8 = shadeChars !! 1
  | otherwise = minShadeChar
  where
    itersToEscape = length . takeWhile (\z -> magnitude z < 2) . take maxIters $ jSequence c (x :+ y)
    jSequence c z = z : jSequence c (z ^ 2 + c)

--Sierpinski triangle (WIP - random)

plotSierpinski :: PlotCommand
plotSierpinski gd args = Right vtcSierpinski

vtcSierpinski :: VectToChar
vtcSierpinski gd v@(V2 x y)
  | x < 0 || x > 1 || y < 0 || y > 1 = minShadeChar
  | any (withinRange gd v) $ truncatedList gd = maxShadeChar
  | otherwise = minShadeChar

truncatedList gd = take iters' $ drop iters' sierpinskiList' --Hardcoded parameter
  where
    iters' = 10 * getIters gd

withinRange gd (V2 x1 y1) (V2 x2 y2) = abs (x2 - x1) <= width && abs (y2 - y1) <= height
  where
    (width, height) = (0.75 * pxlW gd, 0.75 * pxlH gd) --Hardcoded parameter

sierpinskiList :: (Vector2, Int) -> [Vector2]
sierpinskiList (v, n) = v : sierpinskiList (nextSierpinski (v, n))

sierpinskiList' = sierpinskiList (V2 0.5 0.5, 0) --Initial values

nextSierpinski :: (Vector2, Int) -> (Vector2, Int)
nextSierpinski (v@(V2 x y), n)
  | rand n `mod` 3 == 0 = (applyTransformation trans0 v, n + 1)
  | rand n `mod` 3 == 1 = (applyTransformation trans1 v, n + 1)
  | rand n `mod` 3 == 2 = (applyTransformation trans2 v, n + 1)
nextSierpinski _ = (V2 0 0, 0)

trans0 = Transformation (0.5, 0.0, 0.0, 0.5) (V2 0 0)

trans1 = Transformation (0.5, 0.0, 0.0, 0.5) (V2 0.5 0)

trans2 = Transformation (0.5, 0.0, 0.0, 0.5) (V2 0.25 (sqrt 3 / 4))

--Generates a pseudorandom digit
rand :: Int -> Int
rand n = randList !! (n `mod` length randList)

randList =
  [ 3,
    1,
    4,
    1,
    5,
    9,
    2,
    6,
    5,
    3,
    5,
    8,
    9,
    7,
    9,
    3,
    2,
    3,
    8,
    4,
    6,
    2,
    6,
    4,
    3,
    3,
    8,
    3,
    2,
    7,
    9,
    5,
    0,
    2,
    8,
    8,
    4,
    1,
    9,
    7,
    1,
    6,
    9,
    3,
    9,
    9,
    3,
    7,
    5,
    1,
    0,
    5,
    8,
    2,
    0,
    9,
    7,
    4,
    9,
    4,
    4,
    5,
    9,
    2,
    3,
    0,
    7,
    8,
    1,
    6,
    4,
    0,
    6,
    2,
    8,
    6,
    2,
    0,
    8,
    9,
    9,
    8,
    6,
    2,
    8,
    0,
    3,
    4,
    8,
    2,
    5,
    3,
    4,
    2,
    1,
    1,
    7,
    0,
    6,
    7,
    9,
    8,
    2,
    1,
    4,
    8,
    0,
    8,
    6,
    5,
    1,
    3,
    2,
    8,
    2,
    3,
    0,
    6,
    6,
    4,
    7,
    0,
    9,
    3,
    8,
    4,
    4,
    6,
    0,
    9,
    5,
    5,
    0,
    5,
    8,
    2,
    2,
    3,
    1,
    7,
    2,
    5,
    3,
    5,
    9,
    4,
    0,
    8,
    1,
    2,
    8,
    4,
    8,
    1,
    1,
    1,
    7,
    4,
    5,
    0,
    2,
    8,
    4,
    1,
    0,
    2,
    7,
    0,
    1,
    9,
    3,
    8,
    5,
    2,
    1,
    1,
    0,
    5,
    5,
    5,
    9,
    6,
    4,
    4,
    6,
    2,
    2,
    9,
    4,
    8,
    9,
    5,
    4,
    9,
    3,
    0,
    3,
    8,
    1,
    9,
    6,
    4,
    4,
    2,
    8,
    8,
    1,
    0,
    9,
    7,
    5,
    6,
    6,
    5,
    9,
    3,
    3,
    4,
    4,
    6,
    1,
    2,
    8,
    4,
    7,
    5,
    6,
    4,
    8,
    2,
    3,
    3,
    7,
    8,
    6,
    7,
    8,
    3,
    1,
    6,
    5,
    2,
    7,
    1,
    2,
    0,
    1,
    9,
    0,
    9,
    1,
    4,
    5,
    6,
    4,
    8,
    5,
    6,
    6,
    9,
    2,
    3,
    4,
    6,
    0,
    3,
    4,
    8,
    6,
    1,
    0,
    4,
    5,
    4,
    3,
    2,
    6,
    6,
    4,
    8,
    2,
    1,
    3,
    3,
    9,
    3,
    6,
    0,
    7,
    2,
    6,
    0,
    2,
    4,
    9,
    1,
    4,
    1,
    2,
    7,
    3,
    7,
    2,
    4,
    5,
    8,
    7,
    0,
    0,
    6,
    6,
    0,
    6,
    3,
    1,
    5,
    5,
    8,
    8,
    1,
    7,
    4,
    8,
    8,
    1,
    5,
    2,
    0,
    9,
    2,
    0,
    9,
    6,
    2,
    8,
    2,
    9,
    2,
    5,
    4,
    0,
    9,
    1,
    7,
    1,
    5,
    3,
    6,
    4,
    3,
    6,
    7,
    8,
    9,
    2,
    5,
    9,
    0,
    3,
    6,
    0,
    0,
    1,
    1,
    3,
    3,
    0,
    5,
    3,
    0,
    5,
    4,
    8,
    8,
    2,
    0,
    4,
    6,
    6,
    5,
    2,
    1,
    3,
    8,
    4,
    1,
    4,
    6,
    9,
    5,
    1,
    9,
    4,
    1,
    5,
    1,
    1,
    6,
    0,
    9,
    4,
    3,
    3,
    0,
    5,
    7,
    2,
    7,
    0,
    3,
    6,
    5,
    7,
    5,
    9,
    5,
    9,
    1,
    9,
    5,
    3,
    0,
    9,
    2,
    1,
    8,
    6,
    1,
    1,
    7,
    3,
    8,
    1,
    9,
    3,
    2,
    6,
    1,
    1,
    7,
    9,
    3,
    1,
    0,
    5,
    1,
    1,
    8,
    5,
    4,
    8,
    0,
    7,
    4,
    4,
    6,
    2,
    3,
    7,
    9,
    9,
    6,
    2,
    7,
    4,
    9
  ]


module Data.TimeSeriesSpec (spec) where

import Prelude
import Test.Hspec
import Test.QuickCheck
import Control.Arrow (first)
import Data.Time (UTCTime)
import Data.Time.Clock.POSIX (posixSecondsToUTCTime, utcTimeToPOSIXSeconds)
import qualified Statistics.Sample as S
import qualified Data.Vector as V

import qualified Data.TimeSeries as TS
import Data.TimeSeries.Time (seconds)

sampleSeries :: TS.TimeSeries Double
sampleSeries = TS.tsSeries [1..] [10.0, 1.2, 32.4, 0.6, 11.0]

spec :: Spec
spec = do

  describe "Properties" $ do

    it "toList" $ property $
        \ts -> (TS.series . TS.toList) ts == (ts :: TS.TimeSeries Double)

    it "fold over series" $ property $
        \ts -> (foldr (+) 0.0 ts) - ((sum . (map snd) . TS.toList) (ts :: TS.TimeSeries Double)) < 10**(-6)

    it "resampling the same" $ property $
        \ts -> TS.resample (posixSecondsToUTCTime 1) (seconds 1) ts == (ts :: TS.TimeSeries Double)

    it "groupBy the same" $ property $
        \ts -> TS.groupBy (seconds 1) sum ts == (ts :: TS.TimeSeries Double)

  describe "Selecting sub Series" $ do

    it "return data point value at given index" $ do
        let pos = posixSecondsToUTCTime 2
        TS.valueAt pos sampleSeries `shouldBe` Just 1.2

    it "return Nothing value if wrong index" $ do
        let idx = [10, 20, 30]
        let values = [10.0, 12.0, 32.4]
        let pos = posixSecondsToUTCTime 1234
        TS.valueAt pos (TS.tsSeries idx values) `shouldBe` Nothing

    it "return subset" $ do
        let xs = TS.tsSeries [1..5] [10.0, 1.2, 32.4, 0.65, 11.0]
        let start = utcFromSeconds 2
        let end = utcFromSeconds 3
        TS.size (TS.slice start end xs) `shouldBe` 2

    it "return range" $ do
        let xs = TS.tsSeries [0, 5..] [1.0, 1.2, 32.4, 0.65, 11.0]
        let start = utcFromSeconds 2
        let end = utcFromSeconds 130
        TS.size (TS.slice start end xs) `shouldBe` 4

    it "return empty subset" $ do
        let xs = TS.tsSeries [1..5] [10.0, 1.2, 32.4, 0.65, 11.0]
        let start = utcFromSeconds 12
        let end = utcFromSeconds 34
        TS.slice start end xs `shouldBe` TS.emptySeries

  describe "Basic operations" $ do

    it "return series size" $ do
        let xs = TS.tsSeries [1..] [10.0, 1.2, 32.4, 0.65, 11.0]
        TS.size xs `shouldBe` 5

    it "first element" $ do
        let xs = TS.tsSeries [1..] [10.0, 1.2, 32.4, 0.65, 11.0]
        fmap TS.dpValue (TS.firstElem xs) `shouldBe` Just 10.0

    it "last element" $ do
        let xs = TS.tsSeries [1..] [10.0, 1.2, 32.4, 0.65, 11.0]
        fmap TS.dpValue (TS.lastElem xs) `shouldBe` Just 11.0

    it "element by index" $ do
        let xs = TS.tsSeries [1..] [10.0, 1.2, 32.4, 0.65, 11.0]
        fmap TS.dpValue (TS.elemAt 2 xs) `shouldBe` Just 32.4

    it "map over series" $ do
        let xs = TS.tsSeries [1..] [10.0, 1.2, 32.4, 0.65, 11.0]
        let ys = fmap (+ 2) xs
        TS.values ys `shouldBe` [12.0, 3.2, 34.4, 2.65, 13.0]

    it "maximum value" $ do
        let xs = TS.tsSeries [1..] [10.0, 1.2, 32.4, 0.65, 11.0]
        maximum xs `shouldBe` 32.4

    it "minimum value" $ do
        let xs = TS.tsSeries [1..5] [10.0, 1.2, 32.4, 0.65, 11.0]
        minimum xs `shouldBe` 0.65

  describe "Rolling" $ do

    it "sum" $ do
        let xs = TS.tsSeries [1..] [1.0, 2.0, 3.0, 4.0, 5.0]
        TS.rolling (seconds 2) sum xs `shouldBe` TS.tsSeries [2..] [3.0, 5.0, 7.0, 9.0]

  describe "Smoothing" $ do

    it "mean" $ do
        let xs = TS.tsSeries [1..] [1.0, 2.0, 3.0, 4.0, 5.0]
        TS.rolling (seconds 3) (S.mean . V.fromList) xs `shouldBe` TS.tsSeries [3..] [2.0, 3.0, 4.0]

  describe "Resampling" $ do

    it "missing values" $ do
        let xs = TS.tsSeries [1, 2, 3, 5] [1.0, 2.0, 3.0, 5.0]
        let startTime = posixSecondsToUTCTime 1
        TS.resample startTime (seconds 1) xs `shouldBe` TS.tsSeries [1, 2, 3, 4, 5] [1.0, 2.0, 3.0, 4.0, 5.0]

    it "missing lots of values" $ do
        let xs = TS.tsSeries [1, 5] [1.0, 5.0]
        let startTime = posixSecondsToUTCTime 1
        TS.resample startTime (seconds 1) xs `shouldBe` TS.tsSeries [1, 2, 3, 4, 5] [1.0, 2.0, 3.0, 4.0, 5.0]

    it "middle values" $ do
        let xs = TS.tsSeries [1, 3, 5] [1.0, 3.0, 5.0]
        let startTime = posixSecondsToUTCTime 0
        TS.resample startTime (seconds 2) xs `shouldBe` TS.tsSeries [0, 2, 4] [1.0, 2.0, 4.0]

  describe "Group by" $ do

    it "sum" $ do
        let xs = TS.tsSeries [1..] [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
        TS.groupBy (seconds 2) sum xs `shouldBe` TS.tsSeries [1, 3, 5] [3.0, 7.0, 11.0]

  describe "zip" $ do

    it "the same index" $ do
        let xs = TS.tsSeries [1..] [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
        let ys = TS.tsSeries [1..] [2.0, 3.0, 4.0, 5.0, 6.0, 7.0]
        fmap (\(x,y) -> x+y) (TS.zip xs ys) `shouldBe` TS.tsSeries [1..] [3.0, 5.0, 7.0, 9.0, 11.0, 13.0]

-- Convert to UTC Time from timestamp
utcFromSeconds :: Integer -> UTCTime
utcFromSeconds = posixSecondsToUTCTime . fromIntegral

-- Generate Time Series
instance (Arbitrary a) => Arbitrary (TS.TimeSeries a) where
    arbitrary = do
        vs <- arbitrary
        return $ TS.tsSeries [1..] vs


module STCR2Z2T0S0ReversalCornerPointSet where

import           Data.Array.Repa         as R
import           Data.Binary             (decodeFile)
import           Data.Complex
import           Data.List               as L
import           DFT.Plan
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           STC.PowerMethod
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoDecayStr:taoReversalStr:taoCornerStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:scale0FreqsStr:scaleFreqsStr:histFilePath:numIterationStr:writeSourceFlagStr:rStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      taoDecay = read taoDecayStr :: Double
      taoReversal = read taoReversalStr :: Double
      taoCorner = read taoCornerStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scale0Freq = read scale0FreqsStr :: Double
      scaleFreq = read scaleFreqsStr :: Double
      scale0Freqs = [-scale0Freq .. scale0Freq]
      scaleFreqs = [-scaleFreq .. scaleFreq]
      numIteration = read numIterationStr :: Int
      writeSourceFlag = read writeSourceFlagStr :: Bool
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2Z2T0S0ReversalCornerPointSet"
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z2T0S0ReversalCornerRadial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          taoDecay
          taoReversal
          taoCorner
          theta0Freqs
          thetaFreqs
          scale0Freqs
          scaleFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length scale0Freqs
             , L.length theta0Freqs
             , L.length scaleFreqs
             , L.length thetaFreqs
             ]
             0)
  arrR2Z2T0S0 <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      radialArr
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      theta0Freqs
      scale0Freqs
  plan <- makeR2Z2T0S0Plan emptyPlan arrR2Z2T0S0
  let n = 30
      m = round $ (fromIntegral n) * (sqrt 2) / 2
      a = 10
      b = -10
      c = 10
      a' = round $ (fromIntegral a) * (sqrt 2) / 2
      b' = round $ (fromIntegral b) * (sqrt 2) / 2
      c' = round $ (fromIntegral c) * (sqrt 2) / 2
      r = read rStr :: Double
      rr = round r :: Int
      shift = 15
      numTheta = 45
      deltaTheta = 1 * pi / numTheta
      xs =
        -- [ R2S1RPPoint (rr, shift, 0, 1)
        -- , R2S1RPPoint (rr, -shift, 0, 1)
        -- , R2S1RPPoint (-rr, shift, 0, 1)
        -- , R2S1RPPoint (-rr, -shift, 0, 1)
        -- , R2S1RPPoint (shift, rr, 0, 1)
        -- , R2S1RPPoint (-shift, rr, 0, 1)
        -- , R2S1RPPoint (shift, -rr, 0, 1)
        -- , R2S1RPPoint (-shift, -rr, 0, 1)
        -- ]
        -- [R2S1RPPoint (i, -i, 0, 1) | i <- [-2,0 .. 20]] L.++
        [R2S1RPPoint (-i, -i, 0, 1) | i <- [-10,-8 .. 10]]
        -- ((L.map
        --     (\(i, j) -> R2S1RPPoint (round i, round j, 0, 1))
        --     [ (r * cos (k * deltaTheta) + 0, r * sin (k * deltaTheta) + 0)
        --     | k <- [0 .. numTheta]
        --     ]) -- L.++
        --  -- (L.map
        --  --    (\(i, j) -> R2S1RPPoint (round i, round j, 0, 1))
        --  --    [ ((r + 2) * cos (k * deltaTheta), (r + 2) * sin (k * deltaTheta))
        --  --    | k <- [0 .. numTheta - 1]
        --  --    ]) L.++
        --  -- (L.map
        --  --    (\(i, j) -> R2S1RPPoint (round i, round j, 0, 1))
        --  --    [ ( (r + 5) * cos (k * deltaTheta)
        --  --      , (r + 5) * sin (k * deltaTheta))
        --  --    | k <- [0 .. numTheta - 1]
        --  --    ])
        --        -- L.++ [R2S1RPPoint (3, 5, 0, 1)]
        --  )
               -- ([R2S1RPPoint (i, i, 0, 1) | i <- [a',a' + b' .. c']] L.++
               --  [R2S1RPPoint (i, -i, 0, 1) | i <- [-a',-(a' + b') .. -c']] L.++
               --  [R2S1RPPoint (i, i, 0, 1) | i <- [-a',-(a' + b') .. -c']] L.++
               --  [R2S1RPPoint (i, -i, 0, 1) | i <- [a',a' + b' .. c']] L.++
               --  [R2S1RPPoint (i, 0, 0, 1) | i <- [a,a + b .. c]] L.++
               --  [R2S1RPPoint (0, i, 0, 1) | i <- [-a,-(a + b) .. -c]] L.++
               --  [R2S1RPPoint (i, 0, 0, 1) | i <- [-a,-(a + b) .. -c]] L.++
               --  [R2S1RPPoint (0, i, 0, 1) | i <- [a,a + b .. c]] L.++
               --  [R2S1RPPoint (3, 5, 0, 1)])
                -- [ R2S1RPPoint (n, 0, 0, 1)
                -- , R2S1RPPoint (0, n, 0, 1)
                -- , R2S1RPPoint (-n, 0, 0, 1)
                -- , R2S1RPPoint (0, -n, 0, 1)
                -- , R2S1RPPoint (m, m, 0, 1)
                -- , R2S1RPPoint (-m, m, 0, 1)
                -- , R2S1RPPoint (m, -m, 0, 1)
                -- , R2S1RPPoint (-m, -m, 0, 1)
                -- ]
  let bias = computeBiasR2T0S0 numPoint numPoint theta0Freqs scale0Freqs xs
      eigenVec =
        computeInitialEigenVectorR2T0S0
          numPoint
          numPoint
          theta0Freqs
          scale0Freqs
          thetaFreqs
          scaleFreqs
          xs
  powerMethodR2Z2T0S0
    plan
    folderPath
    numPoint
    numPoint
    numOrientation
    thetaFreqs
    theta0Freqs
    numScale
    scaleFreqs
    scale0Freqs
    arrR2Z2T0S0
    numIteration
    writeSourceFlag
    -- (printf "_%d" (round r :: Int))
    (printf
       "_%d_%d_%d_%d_%.2f_%.2f"
       (round maxScale :: Int)
       (round taoDecay :: Int)
       (round taoReversal :: Int)
       (round taoCorner :: Int)
       thetaSigma
       scaleSigma)
    0.5
    bias
    eigenVec


module Shuffle where

import Numeric.LinearAlgebra ((<>))
import qualified System.Random as Random

import qualified Cube as Cube
import Rotation (MatrixRotation)
import qualified Rotation as Rotation

shuffleRotations :: [MatrixRotation]
shuffleRotations = [
    Rotation.rightL
  , Rotation.rightR
  , Rotation.leftL
  , Rotation.leftR
  , Rotation.topL
  , Rotation.topR
  , Rotation.bottomL
  , Rotation.bottomR
  , Rotation.frontL
  , Rotation.frontR
  , Rotation.backL
  , Rotation.backR
  ]

randomRotations :: Int -> [MatrixRotation]
randomRotations seed =
  let gen = Random.mkStdGen seed
      indices = map (\x -> (abs x) `mod` length shuffleRotations) $ Random.randoms gen
  in map (shuffleRotations !!) indices

exampleShuffle =
  let rotations = take 20 $ randomRotations 5
      names = map Rotation.nameMatrixRotation rotations
      rotation = foldl (<>) Rotation.identity rotations
      cube = Rotation.rotate rotation Cube.solvedCube
  in (cube, names)


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE UnicodeSyntax #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE NoMonomorphismRestriction #-}

module Main where

import Control.Monad ( zipWithM_ )
import Data.Foldable ( toList )
import Data.Monoid ( Monoid(..) )
import Data.SBV
import Data.SBV.Tools.CodeGen
import Data.SBV.Internals ( SolverContext )
import Data.Semigroup ( Semigroup(..) )

import qualified Numeric.GSL.ODE as ODE
import qualified Numeric.LinearAlgebra as LA

import Graphics.Rendering.Chart.Easy hiding ( (.>) )
import Graphics.Rendering.Chart.Backend.Cairo

data Pendulum a = Pendulum
  { pendulumLength :: a
  , pendulumDampingConstant :: a
  , pendulumMass :: a
  , pendulumGravity :: a
  } deriving (Eq, Show, Functor, Foldable, Traversable)

data State a = State
  { stateθ :: a
  , stateω :: a
  } deriving (Eq, Show, Functor, Foldable, Traversable)

newtype Controller a = Controller
  { controllerDamping :: a
  } deriving (Eq, Show, Functor, Foldable, Traversable)

instance EqSymbolic a => EqSymbolic (State a) where
  (State x a) .== (State y b) = x .== y &&& a .== b

stateLabels :: State String
stateLabels = State "θ" "ω"

pendulum :: Fractional a => Pendulum a -> a -> State a -> State a
pendulum sys@(Pendulum len damping _ grav) τ (State θ ω) =
  State ω $
  (grav * taylorSin θ) / len + (damping * ω) / inertia + τ / inertia
  where
    inertia = pendulumInertia sys

pendulumInertia :: Fractional a => Pendulum a -> a
pendulumInertia (Pendulum len _ mass _) = mass * len * len

kineticEnergy :: Fractional a => Pendulum a -> State a -> a
kineticEnergy system (State _ ω) =
  0.5 * pendulumInertia system * ω * ω

dkedt ::
  Fractional a => Controller a -> Pendulum a -> State a -> a
dkedt ctrl system state@(State _ ω) =
  pendulumInertia system * ω * stateω (closedLoop ctrl system state)

-- | Potential energy spans [-2, 0] * mg.
potentialEnergy :: Fractional a => Pendulum a -> State a -> a
potentialEnergy (Pendulum len _ mass grav) (State θ _) =
  len * mass * grav * (taylorCos θ - 1)

dpedt :: Fractional a => Pendulum a -> State a -> a
dpedt (Pendulum len _ mass grav) (State θ ω) =
  len * mass * grav * (- taylorSin θ) * ω

lyapunovController :: Fractional a => Controller a -> Pendulum a -> State a -> a
lyapunovController (Controller kd) (Pendulum len _ mass grav) (State θ ω) =
  (-2) * mass * grav * len * taylorSin θ - kd * ω

closedLoop ::
  Fractional a => Controller a -> Pendulum a -> State a -> State a
closedLoop ctrl system initialState =
  pendulum system torque initialState
  where
    torque = lyapunovController ctrl system initialState

{- Proofs -}

v :: Fractional a => Pendulum a -> State a -> a
v system st =
  kineticEnergy system st - potentialEnergy system st

dvdt :: Fractional a => Controller a -> Pendulum a -> State a -> a
dvdt ctrl system st =
  dkedt ctrl system st - dpedt system st

newtype SAll a = SAll { getSAll :: a }

instance Boolean a => Semigroup (SAll a) where
  (SAll x) <> (SAll y) = SAll $ x &&& y

instance Boolean a => Monoid (SAll a) where
  mempty = SAll true
  mappend = (<>)

allIsPoint :: (IEEEFloating a, Foldable t) => t (SBV a) -> SBool
allIsPoint = getSAll . foldMap (SAll . fpIsPoint)

nanFree ::
  (IEEEFloating a, SolverContext m, Monad m) =>
  (String -> m (SBV a)) -> (State (SBV a) -> SBV a) -> m SBool
nanFree gen f = do
  st <- traverse gen stateLabels
  constrainPi st
  constrainFP st
  pure . fpIsPoint $ f st
  where
    constrainFP = constrain . allIsPoint
    constrainPi (State θ _) = constrain $ θ .<= π &&& θ .> -π
    π = 3.1415926535897932384626433832795028841971693993751

lyapunov'sTheorem ::
  ( SymWord a, Fractional a
  , SolverContext m, Monad m) =>
  (String -> m (SBV a)) -> (State (SBV a) -> SBV a) -> (State (SBV a) -> SBV a) -> m SBool
lyapunov'sTheorem gen f dfdt = do
  st <- traverse gen stateLabels
  constrainPi st
--   constrainFP st
--   constrainFP [f st, dfdt st]
  eq <- lyapunovEquilibrium st
  nn <- lyapunovNonNegative st
  gn <- lyapunovGradNegative st
  pure $ eq &&& nn &&& gn
  where
    constrainPi (State θ _) = constrain $ θ .<= π &&& θ .> -π
    π = 3.1415926535897932384626433832795028841971693993751
    lyapunovEquilibrium _ = pure $
      f (State 0 0) .== 0.0

    lyapunovNonNegative st = do
      constrain $ st ./= State 0 0
      pure $ f st .> 0.0

    lyapunovGradNegative st = pure $
      dfdt st .<= 0.0 &&& dfdt (State 0 0) .<= 0.0

nominalController :: Fractional a => Controller a
nominalController = Controller 0.3

nominalSystem :: Fractional a => Pendulum a
nominalSystem = Pendulum 0.5 (-0.03) 0.1 9.81

systemLabels :: Pendulum String
systemLabels = Pendulum "length" "damping" "mass" "gravity"

controllerLabels :: Controller String
controllerLabels = Controller "kd"

proveStability :: IO ThmResult
proveStability =
  prove $ lyapunov'sTheorem sReal v' dvdt'
  where
    v' = v nominalSystem
    dvdt' = dvdt nominalController nominalSystem

proveNanSafety :: IO ThmResult
proveNanSafety =
  prove $ nanFree sFloat controller
  where
    controller = lyapunovController nominalController nominalSystem

-- Simulation

dxdt :: Floating a => p -> [a] -> [a]
dxdt _ [θ, ω] = toList $
  closedLoop nominalController nominalSystem
  (State θ ω)
dxdt _ _ = error "Invalid arguments to 'dxdt'"

dxdtOpenLoop :: Fractional a => p -> [a] -> [a]
dxdtOpenLoop _ [θ, ω] =
  toList $ pendulum nominalSystem 0 (State θ ω)
dxdtOpenLoop _ _ = error "Invalid arguments to 'dxdtOpenLoop'"

solution :: State Double
         -> (Double -> [Double] -> [Double])
         -> LA.Vector Double
         -> LA.Matrix Double
solution state0 f = ODE.odeSolve f $ toList state0

listSolution :: State Double
             -> (Double -> [Double] -> [Double])
             -> [Double]
             -> [(Double, Double)]
listSolution state0 f =
  fmap assoc . LA.toRows . solution state0 f . LA.fromList
  where
    assoc vec = let [t, x] = LA.toList vec in (t, x)

initialStates :: Fractional a => [State a]
initialStates = zipWith State [1e-3, -0.5, 0.3] [1e-3, 0.1, 0.3]

sampleTs :: (Enum a, Fractional a) => [a]
sampleTs = [0, 0.01 .. 7]

makePlot :: PlotValue a
         => String -> String -> String -> [[(a, a)]] -> IO ()
makePlot nm title lbl trajectories =
  toFile opts (nm <> ".png") $ do
  layout_title .= title
  setColors $ fmap opaque [red, blue, green]
  zipWithM_ mkPlot [0 :: Int ..] trajectories
  where
    mkPlot num = plot . line (lbl <> show num) . pure
    opts = def { _fo_size = (1280, 720) }

plotStates :: String -> (Double -> [Double] -> [Double]) -> IO ()
plotStates prefix dynamics = do
  makePlot (prefix <> "_theta") (prefix <> " pendulum angle") "θ" thetas
  makePlot (prefix <> "_omega") (prefix <> " pendulum velocity") "ω" omegas
  where
    trajs = fmap (\st -> listSolution st dynamics sampleTs) initialStates
    withTime sel = zip sampleTs . fmap sel
    thetas, omegas :: [[(Double, Double)]]
    thetas = fmap (withTime fst) trajs
    omegas = fmap (withTime snd) trajs

unstabilized :: IO ()
unstabilized = plotStates "Unstabilized" dxdtOpenLoop

stabilized :: IO ()
stabilized = plotStates "Stabilized" dxdt

main :: IO ()
main = do
  unstabilized
  stabilized
  genCCode

-- Trigonometry

taylorCos :: Fractional a => a -> a
taylorCos x = 1 + sum (take 10 series)
  where
    inc num old =
      let new = old * x * x / (num * (num + 1))
      in new : inc (num + 2) new
    signs = cycle [negate, id]
    series = zipWith ($) signs (inc 1 1)

taylorSin :: Fractional a => a -> a
taylorSin x = x + sum (take 10 series)
  where
    inc num old =
      let new = old * x * x / (num * (num + 1))
      in new : inc (num + 2) new
    signs = cycle [negate, id]
    series = zipWith ($) signs (inc 2 x)

-- C code generation

emitController ::
  (Fractional a, SymWord a) =>
  (String -> SBVCodeGen (SBV a)) -> IO ()
emitController gen = compileToC Nothing "lyapunovController" $ do
  system <- traverse gen systemLabels
  controller <- traverse gen controllerLabels
  state <- traverse gen stateLabels
  cgReturn $ lyapunovController controller system state

genCCode :: IO ()
genCCode = do
  emitController cgGen
  emitTaylor taylorSin "taylorSin" cgGen
  emitTaylor taylorCos "taylorCos" cgGen
  emitCalloutController cgGen
  where
    cgGen :: String -> SBVCodeGen SDouble
    cgGen = cgInput

emitTaylor
  :: (Fractional a, SymWord a)
  => (SBV a -> SBV a)
  -> String
  -> (String -> SBVCodeGen (SBV a))
  -> IO ()
emitTaylor f funName gen = compileToC Nothing funName $
  gen "x" >>= cgReturn . f

taylorSin' :: (Fractional a, SymWord a) => SBV a -> SBV a
taylorSin' = cgUninterpret "taylorSin" mempty taylorSin

taylorCos' :: (Fractional a, SymWord a) => SBV a -> SBV a
taylorCos' = cgUninterpret "taylorCos" mempty taylorCos

lyapunovController'
  :: (SymWord a, Fractional a) =>
     Controller (SBV a)
     -> Pendulum (SBV a) -> State (SBV a) -> SBV a
lyapunovController' (Controller kd) (Pendulum len _ mass grav) (State θ ω) =
  -2 * mass * grav * len * taylorSin' θ + kd * (-ω)

emitCalloutController
  :: (SymWord a, Fractional a) =>
     (String -> SBVCodeGen (SBV a)) -> IO ()
emitCalloutController gen = compileToC Nothing "lyapunovController2" $ do
  system <- traverse gen systemLabels
  controller <- traverse gen controllerLabels
  state <- traverse gen stateLabels
  cgReturn $ lyapunovController' controller system state


--Quantum circuits
--Copyright Laurence Emms 2018

module QuantumCircuit (bell) where

import Data.Complex
import Register
import Instruction

bell :: (Floating a, Show a, Ord a) => Circuit a
bell = let parameters@[Left a, Left b] = [Left (MetaQubitRegister "a"), Left (MetaQubitRegister "b")]
           instructions = [(Hadamard a),
                           (CNot a b)]
       in Circuit "BELL" parameters instructions

clear :: (Floating a, Show a, Ord a) => Circuit a
clear = let parameters@[Left q, Right scratch] = [Left (MetaQubitRegister "q"), Right (MetaRegister "scratch")]
            instructions = [(MeasureOut q scratch),
                            (JumpUnless "CLEAREND" scratch),
                            (PauliX q),
                            (Label "CLEAREND")]
        in Circuit "CLEAR" parameters instructions

euler :: (Floating a, Show a, Ord a) => Complex a -> Complex a -> Complex a -> Circuit a
euler alpha beta gamma = let parameters@[Left q] = [Left (MetaQubitRegister "q")]
                             instructions = [(RX (ComplexConstant alpha) q),
                                             (RY (ComplexConstant alpha) q),
                                             (RZ (ComplexConstant alpha) q)]
                         in Circuit "EULER" parameters instructions


{-# OPTIONS_GHC -Wall                     #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing  #-}
{-# OPTIONS_GHC -fno-warn-type-defaults   #-}
{-# OPTIONS_GHC -fno-warn-unused-do-bind  #-}
{-# OPTIONS_GHC -fno-warn-missing-methods #-}
{-# OPTIONS_GHC -fno-warn-orphans         #-}

{-# LANGUAGE MultiParamTypeClasses        #-}
{-# LANGUAGE TypeFamilies                 #-}
{-# LANGUAGE ScopedTypeVariables          #-}
{-# LANGUAGE ExplicitForAll               #-}
{-# LANGUAGE DataKinds                    #-}

{-# LANGUAGE FlexibleInstances            #-}
{-# LANGUAGE MultiParamTypeClasses        #-}
{-# LANGUAGE FlexibleContexts             #-}
{-# LANGUAGE TypeFamilies                 #-}
{-# LANGUAGE BangPatterns                 #-}
{-# LANGUAGE GeneralizedNewtypeDeriving   #-}
{-# LANGUAGE TemplateHaskell              #-}
{-# LANGUAGE DataKinds                    #-}
{-# LANGUAGE DeriveGeneric                #-}

module Particle ( main ) where

import           Data.Random hiding ( StdNormal, Normal )
import           Data.Random.Source.PureMT ( pureMT )
import           Control.Monad.State ( evalState, replicateM )
import           Numeric.LinearAlgebra.Static
                 ( R, vector, Sym,
                   headTail, matrix, sym,
                   diag
                 )
import           GHC.TypeLits ( KnownNat )
import           Data.Random.Distribution.Static.MultivariateNormal ( Normal(..) )
import qualified Data.Vector as V
import           Data.Vector ( Vector )
import           Data.List ( transpose )
import           Control.Parallel.Strategies
import           GHC.Generics (Generic)

import           Numeric.Particle

import qualified Graphics.Rendering.Chart as C
import Graphics.Rendering.Chart.Backend.Diagrams
import Data.Colour
import Data.Colour.Names
import Data.Default.Class
import Control.Lens

import Diagrams.Backend.Cairo.CmdLine
import Diagrams.Prelude hiding ( render, Renderable, trace, Vector, sample )
import Diagrams.Backend.CmdLine

import Data.Csv
import System.IO hiding ( hGetContents )
import Data.ByteString.Lazy ( hGetContents )

nParticles :: Int
nParticles = 1000 -- 500

data SystemState a = SystemState { x1  :: a, x2  :: a }
  deriving (Show, Generic)

instance NFData a => NFData (SystemState a)

newtype SystemObs a = SystemObs { y1  :: a }
  deriving Show

(.+), (.*), (.-) :: (Num a) => V.Vector a -> V.Vector a -> V.Vector a
(.+) = V.zipWith (+)
(.*) = V.zipWith (*)
(.-) = V.zipWith (-)

deltaT, g :: Double
deltaT = 0.01
g  = 9.81

type PendulumState = R 2

stateUpdate :: Particles (SystemState Double) ->
                Particles (SystemState Double)
stateUpdate xPrevs = V.zipWith SystemState x1s x2s
  where
    ix = V.length xPrevs

    x1Prevs = V.map x1 xPrevs
    x2Prevs = V.map x2 xPrevs

    deltaTs = V.replicate ix deltaT
    gs = V.replicate ix g
    x1s = x1Prevs .+ (x2Prevs .* deltaTs)
    x2s = x2Prevs .- (gs .* (V.map sin x1Prevs) .* deltaTs)

stateUpdateNoisy :: MonadRandom m =>
                    Sym 2 ->
                    Particles (SystemState Double) ->
                    m (Particles (SystemState Double))
stateUpdateNoisy bigQ xPrevs = do
  let xs = stateUpdate xPrevs

      x1s = V.map x1 xs
      x2s = V.map x2 xs

  let ix = V.length xPrevs
  etas <- replicateM ix $ sample $ rvar (Normal 0.0 bigQ)

  let eta1s, eta2s :: V.Vector Double
      eta1s = V.fromList $ map (fst . headTail) etas
      eta2s = V.fromList $ map (fst . headTail . snd . headTail) etas

  return (V.zipWith SystemState (x1s .+ eta1s) (x2s .+ eta2s))

obsUpdate :: Particles (SystemState Double) ->
             Particles (SystemObs Double)
obsUpdate xs = V.map (SystemObs . sin . x1) xs

weight :: forall a n . KnownNat n =>
          (a -> R n) ->
          Sym n ->
          a -> a -> Double
weight f bigR obs obsNew = pdf (Normal (f obsNew) bigR) (f obs)

bigP :: Sym 2
bigP = sym $ diag 0.1

initParticles :: MonadRandom m =>
                 m (Particles (SystemState Double))
initParticles = V.replicateM nParticles $ do
  r <- sample $ rvar (Normal m0 bigP)
  let x1 = fst $ headTail r
      x2 = fst $ headTail $ snd $ headTail r
  return $ SystemState { x1 = x1, x2 = x2}

runFilter :: Particles (SystemObs Double) -> Vector (Particles (SystemState Double))
runFilter pendulumSamples = evalState action (pureMT 19)
  where
    action = do
      xs <- initParticles
      scanMapM
        (runPF (stateUpdateNoisy bigQ) obsUpdate (weight f bigR))
        return
        xs
        pendulumSamples

f :: SystemObs Double -> R 1
f = vector . pure . y1

h :: SystemState Double -> R 2
h u = vector [x1 u , x2 u]

bigQ :: Sym 2
bigQ = sym $ matrix bigQl

qc1 :: Double
qc1 = 0.01

bigQl :: [Double]
bigQl = [ qc1 * deltaT^3 / 3, qc1 * deltaT^2 / 2,
          qc1 * deltaT^2 / 2,       qc1 * deltaT
         ]

bigR :: Sym 1
bigR  = sym $ matrix [0.1]

m0 :: PendulumState
m0 = vector [1.6, 0]

testSmoothing :: Particles (SystemObs Double) -> Int -> [Double]
testSmoothing ss n = V.toList $ evalState action (pureMT 23)
  where
    action = do
      xss <- V.replicateM n $ oneSmoothingPath (stateUpdateNoisy bigQ) (weight h bigQ) nParticles (runFilter ss)
      let yss = V.fromList $ map V.fromList $
                transpose $
                V.toList $ V.map (V.toList) $
                xss
      return $ V.map (/ (fromIntegral n)) $ V.map V.sum $ V.map (V.map x1) yss

chartEstimated :: String ->
              [(Double, Double)] ->
              [(Double, Double)] ->
              [(Double, Double)] ->
              C.Renderable ()
chartEstimated title acts obs ests = C.toRenderable layout
  where

    actuals = C.plot_lines_values .~ [acts]
            $ C.plot_lines_style  . C.line_color .~ opaque red
            $ C.plot_lines_title .~ "Actual Trajectory"
            $ C.plot_lines_style  . C.line_width .~ 1.0
            $ def

    measurements = C.plot_points_values .~ obs
                 $ C.plot_points_style  . C.point_color .~ opaque blue
                 $ C.plot_points_title .~ "Measurements"
                 $ def

    estimas = C.plot_lines_values .~ [ests]
            $ C.plot_lines_style  . C.line_color .~ opaque black
            $ C.plot_lines_title .~ "Inferred Trajectory"
            $ C.plot_lines_style  . C.line_width .~ 1.0
            $ def

    layout = C.layout_title .~ title
           $ C.layout_plots .~ [C.toPlot actuals, C.toPlot measurements, C.toPlot estimas]
           $ C.layout_y_axis . C.laxis_title .~ "Angle / Horizontal Displacement"
           $ C.layout_y_axis . C.laxis_override .~ C.axisGridHide
           $ C.layout_x_axis . C.laxis_title .~ "Time"
           $ C.layout_x_axis . C.laxis_override .~ C.axisGridHide
           $ def

nObs :: Int
nObs = 200

diagU :: IO (Diagram Cairo)
diagU = do
  h <- openFile "matlabRNGs.csv" ReadMode
  cs <- hGetContents h
  let df = (decode NoHeader cs) :: Either String (V.Vector (Double, Double))
  case df of
    Left _ -> error "Whatever"
    Right generatedSamples -> do
      let preObs = V.take nObs $ V.map fst generatedSamples
      let obs = V.toList preObs
      let acts = V.toList $ V.take nObs $ V.map snd generatedSamples
      let nus = take nObs (testSmoothing (V.map SystemObs preObs) 50)
      denv <- defaultEnv C.vectorAlignmentFns 600 500
      let charte = chartEstimated "Particle Smoother"
                                  (zip [0,1..] acts)
                                  (zip [0,1..] obs)
                                  (zip [0,1..] nus)
      return $ fst $ runBackend denv (C.render charte (600, 500))

displayHeader :: FilePath -> Diagram B -> IO ()
displayHeader fn =
  mainRender ( DiagramOpts (Just 900) (Just 700) fn
             , DiagramLoopOpts False Nothing 0
             )

main :: IO ()
main = do
  du <- diagU
  displayHeader "diagrams/Smooth3.png" du


module Statistics.Information.Utils.Matrix where

import Data.Matrix
import qualified Data.Vector as Vector
import Statistics.Information.Utils.List

rows :: Matrix a -> [Vector.Vector a]
rows xs = [getRow i xs | i <- [1..nrows xs]]

columns :: Matrix a -> [Vector.Vector a]
columns xs = [getCol i xs | i <- [1..ncols xs]]

rowExtend :: Matrix a -> Int -> Matrix a
rowExtend xs n =
  if remainder > 0 then
    remRows <-> Prelude.foldr (<->) xs (repeatN (completions - 1) xs)
  else
    Prelude.foldr (<->) xs (repeatN (completions - 1) xs)
  where
    completions = n `quot` nrows xs
    remRows = fromLists [Vector.toList $ getRow i xs |
                         i <- [1..remainder]]
    remainder = n `rem` nrows xs

residualMatrix :: Matrix a -> Int -> Matrix a
xsiLeft xs i = submatrix 1 (nrows xs) 1 i xs
xsiRight xs i = submatrix 1 (nrows xs) (i+1) (ncols xs) xs
residualMatrix xs i | i == (ncols xs) = xsiLeft xs i
residualMatrix xs 1 = xsiRight xs 1
residualMatrix xs i = xsiLeft xs i <|> xsiRight xs i


{-# LANGUAGE NamedFieldPuns #-}

module School.App.MultiLayerPerceptron
( MLPOptions(..)
, multiLayerPerceptron
) where

import Numeric.LinearAlgebra (R)
import School.FileIO.FileHeader (FileHeader(..))
import School.FileIO.Source (source)
import School.Train.GradientDescent (gradientDescent)
import School.Train.IterationHandler (logCost)
import School.Train.SimpleDescentUpdate (simpleDescentUpdate)
import School.Train.StoppingCondition (maxIterations)
import School.Train.TrainState (TrainState(..), def)
import School.Types.FloatEq ((~=))
import School.Types.PingPong (PingPong, toPingPong)
import School.Types.DataType (DataType(..))
import School.Unit.Affine (affine)
import School.Unit.RecLin (recLin)
import School.Unit.MultiNoulli (multiNoulli)
import School.Unit.Unit (Unit)
import School.Unit.UnitParams (UnitParams(..))
import School.Unit.WeightDecay (weightDecay)
import School.Utils.Random (randAffineParams)
import System.FilePath (FilePath)
import System.Random (StdGen, mkStdGen)

data MLPOptions =
  MLPOptions { hiddenDims :: [Int]
             , initRate :: R
             , maxIter :: Int
             , nBatches :: Int
             , nClasses :: Int
             , nFeatures :: Int
             , randomSeed :: Int
             , weightDecayCoeff :: R
             , writeCost :: Maybe FilePath
             }

heLimits :: Int -> (R, R)
heLimits d = (0, sqrt (2 / fromIntegral d))

getUnits :: [Int]
         -> StdGen
         -> ([Unit R], [UnitParams R])
         -> ([Unit R], [UnitParams R])
getUnits (dim1:dim2:dim3:dims) gen (units, params) =
  getUnits (dim2:dim3:dims) gen' acc where
    addUnits = [affine, recLin]
    limits = heLimits dim1
    (affineParams, gen') =
      randAffineParams dim1 dim2 limits gen
    addParams = [affineParams, EmptyParams]
    acc = (units ++ addUnits, params ++ addParams)
getUnits (dim1:dim2:_) gen (units, params) =
  (units ++ [addUnit], params ++ [addParam]) where
    addUnit = affine
    limits = heLimits dim1
    (addParam, _) =
      randAffineParams dim1 dim2 limits gen
getUnits [_] _ acc = acc
getUnits [] _ acc = acc

setupUnits :: Int
           -> Int
           -> [Int]
           -> Int
           -> Either String
                     ([Unit R], PingPong (UnitParams R))
setupUnits nFeatures nClasses hiddenDims seed = do
  let allDims = nFeatures:hiddenDims ++ [nClasses]
  let gen = mkStdGen seed
  let (units, paramList) = getUnits allDims
                                    gen
                                    ([], [])
  params <- toPingPong paramList
  return (units, params)

multiLayerPerceptron :: MLPOptions
                     -> FilePath
                     -> IO (Either String
                           (Maybe [R], [UnitParams R]))
multiLayerPerceptron MLPOptions { nClasses = nC
                                , nBatches = nB
                                , nFeatures = nF
                                , hiddenDims
                                , initRate = learningRate
                                , maxIter
                                , randomSeed
                                , weightDecayCoeff
                                , writeCost
                                }
                      path = do
  let header = FileHeader { cols = nF
                            , rows = nB
                            , dataType = DBL64B
                            }
  let setup = setupUnits nF nC hiddenDims randomSeed
  case setup of
    Left e -> return . Left $ e
    Right (units, paramList) -> do
      let multi = multiNoulli Nothing Nothing
      let cost = if weightDecayCoeff ~= 0
                   then multi
                   else mappend multi (weightDecay weightDecayCoeff)
      let update = simpleDescentUpdate
      let condition = maxIterations maxIter
      let handler = case writeCost of
                      Nothing -> mempty
                      Just p -> logCost p
      let initState = def { paramList
                          , learningRate
                          }
      endState <- gradientDescent (source header path)
                                  units
                                  cost
                                  update
                                  condition
                                  handler
                                  initState
      return $ Left ""