module Pendulum.View (
  drawPendulum
  ) where

import qualified Data.Vector.Sized as V
import qualified Data.Vector.Storable as VS
import           GHC.TypeLits
import           Graphics.Gloss
import           Numeric.Hamilton
import           Numeric.LinearAlgebra.Static
import           Pendulum.Model (Pendulum(..))

-- | 振り子の長さの描画スケール
lengthScale :: Float
lengthScale = 100.0

-- | 振り子のおもりの描画スケール
massScale :: Float
massScale = 10.0

r2list :: KnownNat n => R n -> [Double]
r2list = VS.toList . extract

drawPendulum :: KnownNat n => Pendulum n -> Picture
drawPendulum pd =
  let ls  = map ((*lengthScale) . realToFrac) (V.toList $ _lengthes pd)
      ms  = map ((*massScale) . realToFrac) (V.toList $ _masses pd)
      ps  = map realToFrac (r2list (cfgPositions (_config pd)))
      ths = map (subtract (pi / 2.0)) ps
      xs = scanl (+) 0 (zipWith (*) ls (map cos ths))
      ys = scanl (+) 0 (zipWith (*) ls (map sin ths))
      lines   = line (zip xs ys)
      circles = zipWith3 (translate) (tail xs) (tail ys) (map circleSolid ms)
   in pictures (lines:circles)


module DeltaSigma.PowerSpectrumDensity where

import Numeric.FFT
import Data.Complex

psd :: [Double] -> [Double]
psd xs = take n_half . fmap (/ (fromIntegral n_half))
        . fmap magnitude . fft . fmap (:+ 0) $ take n xs
  where
    n = length xs
    n_half = n `div` 2


{-# LANGUAGE BangPatterns , RankNTypes, GADTs, DataKinds #-}

{- | The 'Numerical.HBLAS.BLAS' module provides a fully general
yet type safe BLAS API.

When in doubt about the semantics of an operation,
consult your system's BLAS api documentation, or just read the documentation
for
<https://software.intel.com/sites/products/documentation/hpc/mkl/mklman/index.htm the Intel MKL BLAS distribution>

A few basic notes about how to invoke BLAS routines.

Many BLAS operations take one or more arguments of type 'Transpose'.
'Tranpose' has the following different constructors, which tell BLAS
routines what transformation to implicitly apply to an input matrix @mat@ with dimension @n x m@.

*  'NoTranspose' leaves the matrix @mat@ as is.

* 'Transpose' treats the @mat@ as being implicitly transposed, with dimension
    @m x n@. Entry @mat(i,j)@ being treated as actually being the entry
    @mat(j,i)@. For Real matrices this is also the matrix adjoint operation.
    ie @Tranpose(mat)(i,j)=mat(j,i)@

*  'ConjNoTranspose' will implicitly conjugate @mat@, which is a no op for Real ('Float' or 'Double') matrices, but for
'Complex Float' and 'Complex Double' matrices, a given matrix entry @mat(i,j)==x':+'y@
will be treated as actually being  @conjugate(mat)(i,j)=y':+'x@.

* 'ConjTranpose' will implicitly transpose and conjugate the input matrix.
ConjugateTranpose acts as matrix adjoint for both real and complex matrices.

The *gemm operations  work as follows (using 'sgemm' as an example):

* @'sgemm trLeft trRight alpha beta left right result'@, where @trLeft@ and @trRight@
are values of type 'Transpose' that respectively act on the matrices @left@ and @right@.

* the generalized matrix computation thusly formed can be viewed as being
@result = alpha * trLeft(left) * trRight(right) + beta * result@

the *gemv operations are akin to the *gemm operations, but with @right@ and @result@
being vectors rather than matrices.

the *trsv operations solve for @x@ in the equation @A x = y@ given @A@ and @y@.
The 'MatUpLo' argument determines if the matrix should be treated as upper or
lower triangular and 'MatDiag' determines if the triangular solver should treat
the diagonal of the matrix as being all 1's or not.  A general pattern of invocation
would be @'strsv' matuplo  tranposeMatA  matdiag  matrixA  xVector@.
A key detail to note is that the input vector is ALSO the result vector,
ie 'strsv' and friends updates the vector place.

-}

module Numerical.HBLAS.BLAS(
        GemvFun
        ,GemmFun
        ,TrsvFun

        ,dgemm
        ,sgemm
        ,cgemm
        ,zgemm

        ,sgemv
        ,dgemv
        ,cgemv
        ,zgemv

        ,sger
        ,dger

        ,strsv
        ,dtrsv
        ,ctrsv
        ,ztrsv
            ) where

import Numerical.HBLAS.UtilsFFI
import Numerical.HBLAS.BLAS.FFI
import Numerical.HBLAS.BLAS.Internal
import Control.Monad.Primitive
import Data.Complex

sgemm :: PrimMonad m=>  GemmFun Float  orient  (PrimState m) m
sgemm =  gemmAbstraction "sgemm"  cblas_sgemm_safe cblas_sgemm_unsafe (\x f -> f x )

dgemm :: PrimMonad m=>  GemmFun  Double orient  (PrimState m) m
dgemm = gemmAbstraction "dgemm"  cblas_dgemm_safe cblas_dgemm_unsafe  (\x f -> f x )

cgemm :: PrimMonad m=>  GemmFun (Complex Float) orient  (PrimState m) m
cgemm = gemmAbstraction "cgemm" cblas_cgemm_safe cblas_cgemm_unsafe  withRStorable_

zgemm :: PrimMonad m=>  GemmFun (Complex Double) orient  (PrimState m) m
zgemm = gemmAbstraction "zgemm"  cblas_zgemm_safe cblas_zgemm_unsafe withRStorable_

sgemv :: PrimMonad m => GemvFun Float orient (PrimState m) m
sgemv = gemvAbstraction "sgemv" cblas_sgemv_safe cblas_sgemv_unsafe (flip ($))

dgemv :: PrimMonad m => GemvFun Double orient (PrimState m) m
dgemv = gemvAbstraction "dgemv" cblas_dgemv_safe cblas_dgemv_unsafe (flip ($))

cgemv :: PrimMonad m => GemvFun (Complex Float) orient (PrimState m) m
cgemv = gemvAbstraction "cgemv" cblas_cgemv_safe cblas_cgemv_unsafe withRStorable_

zgemv :: PrimMonad m => GemvFun (Complex Double) orient (PrimState m) m
zgemv = gemvAbstraction "zgemv" cblas_zgemv_safe cblas_zgemv_unsafe withRStorable_
strsv :: PrimMonad m => TrsvFun Float orient (PrimState m) m
strsv = trsvAbstraction "strsv" cblas_strsv_safe cblas_strsv_unsafe

dtrsv :: PrimMonad m => TrsvFun Double orient (PrimState m) m
dtrsv = trsvAbstraction "dtrsv" cblas_dtrsv_safe cblas_dtrsv_unsafe

ctrsv :: PrimMonad m => TrsvFun (Complex Float) orient (PrimState m) m
ctrsv = trsvAbstraction "ctrsv" cblas_ctrsv_safe cblas_ctrsv_unsafe

ztrsv :: PrimMonad m => TrsvFun (Complex Double) orient (PrimState m) m
ztrsv = trsvAbstraction "ztrsv" cblas_ztrsv_safe cblas_ztrsv_unsafe

sger :: PrimMonad m => GerFun Float orient (PrimState m) m
sger = gerAbstraction "sger" cblas_sger_safe cblas_sger_unsafe

dger :: PrimMonad m => GerFun Double orient (PrimState m) m
dger = gerAbstraction "dger" cblas_dger_safe cblas_dger_unsafe


{-# language TypeFamilies #-}
{-# language FlexibleInstances #-}
{-# language FlexibleContexts #-}

module Geometry.Angle (
  Angular (..), Angle(..)
  ) where

import Geometry.Base
import Geometry.Polygon
import Geometry.Circle
import Geometry.Line

import Data.Complex

------------------------------------------------------------
-- | Class representing angle mark and decorated angle mark.
class (Manifold an, Figure an) => Angular an where
  {-# MINIMAL asAngle, toAngle #-}
  asAngle :: Angle -> an
  -- | Isomorphism for angles.
  toAngle :: an -> Angle

  -- | The setter for an angle.
  setValue :: Direction -> an -> an
  setValue v = asAngle . setValue v . toAngle

  -- | The starting direction for an angle.
  --
  -- >>> anAngle 30 # rotate 45 # angleStart
  -- 45°
  --
  angleStart :: an -> Direction
  angleStart = angleStart . toAngle

  -- | The final direction for an angle.
  --
  -- >>> anAngle 30 # rotate 45 # angleEnd
  -- <interactive>:60:2-8: error:
  --     * Variable not in scope: anAngle :: Integer -> an0
  --     * Perhaps you meant one of these:
  --         `asAngle' (line 20), data constructor `Angle' (line 59),
  --         `angle' (imported from Geometry.Base)
  --
  angleEnd :: an -> Direction
  angleEnd = angleEnd . toAngle

  -- | The value of an angle.
  --
  -- >>> anAngle 30 # rotate 45 # angleValue
  -- <interactive>:145:2-8: error:
  --     * Variable not in scope: anAngle :: Integer -> an0
  --     * Perhaps you meant one of these:
  --         `asAngle' (line 20), data constructor `Angle' (line 59),
  --         `angle' (imported from Geometry.Base)
  --
  angleValue :: an -> Direction
  angleValue an = angleEnd an - angleStart an

instance Angular a => Angular (Maybe a) where
  toAngle = maybe (asCmp 0) toAngle
  asAngle = Just . asAngle

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

-- | Type representing an angle mark on the chart.
data Angle = Angle Cmp Direction Direction
  deriving (Show, Eq)

instance AlmostEq Angle where
  a1 ~= a2 = angleValue a1 ~= angleValue a1

instance Angular Angle where
  asAngle = id
  toAngle = id
  setValue v (Angle p s _) = Angle p s (s + v)
  angleStart (Angle _ s _) = s
  angleEnd (Angle _ _ e) = e

instance Affine Angle where
  cmp = cmp . angleStart
  asCmp x = Angle 0 (asCmp x) (asCmp x)

instance Trans Angle where
  transform t (Angle p s e) = Angle p' s' e'
    where p' = transform t p
          s' = azimuth p' (transform t (cmp p + cmp s))
          e' = azimuth p' (transform t (cmp p + cmp e))

instance Manifold Angle where
  type Domain Angle = Direction
  param an x = asDeg (deg s + x * v)
    where Angle p s e = asAngle an
          v = deg (angleValue an)

  project an p = rad (p - angleStart an) / rad (angleValue an)

  isContaining a p = let x = project a (asAffine p) in 0 <= x && x <= 1

instance Figure Angle where
  refPoint (Angle p _ _) = p
  isTrivial a = angleValue a ~= 0
  box (Angle p s e) = box $ mkCircle 0.1 p

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


{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE RecordWildCards #-}
module Data.Matrix.Static.LinearAlgebra
    ( module Data.Matrix.Static.LinearAlgebra.Types
    , Arithmetic(..)
    , Factorization(..)
    , LinearAlgebra(..)
    , EigenOptions(..)
    , defaultEigenOptions
    , SortRule(..)

      -- * Dense matrix operation
    , zeros
    , ones
    , inverse
    , eig
    , svd
    , cond
    ) where

import qualified Data.Vector.Storable as VS
import Data.Vector.Storable (Vector)
import System.IO.Unsafe (unsafePerformIO)
import Data.Complex (Complex)
import Data.Singletons.Prelude hiding ((@@), type (==), type (-), type (<=))
import Data.Type.Equality
import GHC.TypeLits
import Data.Maybe (fromMaybe)

import qualified Data.Matrix.Static.Dense as D
import qualified Data.Matrix.Static.Sparse as S
import qualified Data.Matrix.Static.Generic.Mutable as CM
import qualified Data.Matrix.Static.Generic as C
import qualified Data.Matrix.Static.Internal as Internal
import Data.Matrix.Static.LinearAlgebra.Types
import Data.Matrix.Static.LinearAlgebra.Internal

class Arithmetic (mat1 :: C.MatrixKind) (mat2 :: C.MatrixKind) where
    -- | Matrix multiplication between different types of matrices.
    (@@) :: ( Numeric a, SingI n, SingI m
            , If (mat1 == mat2) mat1 D.Matrix ~ mat3 )
         => mat1 n p Vector a
         -> mat2 p m Vector a
         -> mat3 n m Vector a
    infixr 8 @@

    -- | Element-wise addition between different types of matrices.
    (%+%) :: ( Numeric a, SingI n, SingI m
             , If (mat1 == mat2) mat1 D.Matrix ~ mat3 )
          => mat1 n m Vector a
          -> mat2 n m Vector a
          -> mat3 n m Vector a
    infixr 8 %+%

    -- | Element-wise substraction between different types of matrices.
    (%-%) :: ( Numeric a, SingI n, SingI m
             , If (mat1 == mat2) mat1 D.Matrix ~ mat3 )
          => mat1 n m Vector a
          -> mat2 n m Vector a
          -> mat3 n m Vector a
    infixr 8 %-%

    -- | Element-wise multiplication between different types of matrices.
    (%*%) :: ( Numeric a, SingI n, SingI m
             , If (mat1 == mat2) mat1 S.SparseMatrix ~ mat3 )
          => mat1 n m Vector a
          -> mat2 n m Vector a
          -> mat3 n m Vector a
    infixr 8 %*%

instance Arithmetic D.Matrix D.Matrix where
    (@@) = withFun2 Internal.c_dd_mul
    (%+%) = (+)
    (%-%) = (-)
    (%*%) = (*)

instance Arithmetic D.Matrix S.SparseMatrix where
    (@@) = withDS Internal.c_ds_mul
    (%+%) = flip (%+%)
    (%-%) a b = a %+% C.map negate b
    (%*%) = undefined

instance Arithmetic S.SparseMatrix D.Matrix where
    (@@) = withSD Internal.c_sd_mul
    (%+%) = withSD Internal.c_sd_plus
    (%-%) a b = a %+% C.map negate b
    (%*%) = undefined

instance Arithmetic S.SparseMatrix S.SparseMatrix where
    (@@) = withSS Internal.c_ss_mul
    (%+%) = withSS Internal.c_ss_plus
    (%-%) a b = a %+% C.map negate b
    (%*%) = withSS Internal.c_ss_cmul

class LinearAlgebra (mat :: C.MatrixKind) where
    ident :: (Numeric a, SingI n) => mat n n Vector a

    colSum :: (Numeric a, SingI n, C.Matrix mat Vector a)
           => mat m n Vector a
           -> Matrix 1 n a
    colSum mat = D.create $ do
        m <- CM.replicate 0
        flip C.imapM_ mat $ \(_,j) v -> CM.unsafeModify m (+v) (0, j)
        return m
    {-# INLINE colSum #-}

    rowSum :: (Numeric a, SingI m, C.Matrix mat Vector a)
           => mat m n Vector a
           -> Matrix m 1 a
    rowSum mat = D.create $ do
        m <- CM.replicate 0
        flip C.imapM_ mat $ \(i,_) x -> CM.unsafeModify m (+x) (i, 0)
        return m
    {-# INLINE rowSum #-}

instance LinearAlgebra D.Matrix where
    ident = D.diag 0 $ D.replicate 1

instance LinearAlgebra S.SparseMatrix where
    ident = S.diag $ D.replicate 1

data EigenOptions = EigenOptions
    { _sigma :: Maybe Double -- ^ Find eigenvalues near sigma using shift-invert mode.
    , _ncv :: Maybe Int      -- ^ Parameter that controls the convergence
                             -- speed of the algorithm. Typically a larger ncv
                             -- means faster convergence, but it may also result
                             -- in greater memory use and more matrix operations
                             -- in each iteration. This parameter must satisfy
                             -- k+2 <= ncv <= n, and is advised to take ncv > 2k+1.
                             -- Default: min(n, max(2*k + 1, 20))
    , _maxit :: Maybe Int          -- ^ Maximum number of iterations allowed in the algorithm.
    , _tol   :: Double       -- Precision parameter for the calculated eigenvalues.
    , _sort_rule :: SortRule
    }

defaultEigenOptions :: EigenOptions
defaultEigenOptions = EigenOptions
    { _sigma = Nothing
    , _ncv = Nothing
    , _maxit = Nothing
    , _tol = 2.2204460492503131e-16
    , _sort_rule = LM }

data SortRule = LM
              | LR
              | LI
              | LA
              | SM
              | SR
              | SI
              | SA
              | BE
              deriving (Enum)

class Factorization mat where
    -- | Eigenvalues (from largest to smallest) and
    -- eigenvectors (as columns) of a general square matrix.
    eigS :: (SingI k, SingI n, k <= n - 2)
         => Sing k
         -> mat n n Vector Double
         -> (Matrix k 1 (Complex Double), Matrix n k (Complex Double))
    eigS a b = eigS' a b defaultEigenOptions

    eigS' :: (SingI k, SingI n, k <= n - 2)
         => Sing k
         -> mat n n Vector Double
         -> EigenOptions
         -> (Matrix k 1 (Complex Double), Matrix n k (Complex Double))

    -- | Eigenvalues (from largest to smallest) and
    -- eigenvectors (as columns) of a symmetric square matrix.
    eigSH :: (SingI k, SingI n, k <= n - 1)
          => Sing k
          -> mat n n Vector Double
          -> (Matrix k 1 Double, Matrix n k Double)
    eigSH a b = eigSH' a b defaultEigenOptions

    eigSH' :: (SingI k, SingI n, k <= n - 1)
           => Sing k
           -> mat n n Vector Double
           -> EigenOptions
           -> (Matrix k 1 Double, Matrix n k Double)

    -- | Generalized eigen solver for real symmetric matrices
    geigSH :: (SingI k, SingI n, k <= n - 1)
           => Sing k
           -> mat n n Vector Double
           -> SparseMatrix n n Double
           -> (Matrix k 1 Double, Matrix n k Double)
    geigSH a b c = geigSH' a b c defaultEigenOptions

    geigSH' :: (SingI k, SingI n, k <= n - 1)
            => Sing k
            -> mat n n Vector Double
            -> SparseMatrix n n Double
            -> EigenOptions
            -> (Matrix k 1 Double, Matrix n k Double)

    -- | Cholesky decomposition
    cholesky :: (Numeric a, SingI n) => mat n n Vector a -> mat n n Vector a

instance Factorization D.Matrix where
    eigS' s mat EigenOptions{..}
        | D.all (==0) mat = ( D.replicate 0, D.replicate 1)
        | otherwise = unsafePerformIO $ do
            m1 <- CM.new
            m2 <- CM.new
            info <- fmap Internal.computationInfo $ unsafeWith' m1 $ \v1 _ _ ->
                unsafeWith' m2 $ \v2 _ _ -> unsafeWith mat $ \v n _ ->
                    Internal.c_eigs k v1 v2 v n ncv' maxit _tol mode sigma $ fromIntegral $ fromEnum _sort_rule
            case info of
                Nothing -> do
                    m1' <- C.unsafeFreeze m1
                    m2' <- C.unsafeFreeze m2
                    return (m1', m2')
                Just err -> error err
      where
        ncv' = case _ncv of
            Just x -> fromIntegral x
            Nothing -> min d $ max 20 $ 2 * k + 1
        maxit = case _maxit of
            Just x -> fromIntegral x
            Nothing -> d * 10
        mode = case _sigma of
            Nothing -> 0
            _ -> 1
        sigma = fromMaybe 0 _sigma
        k = fromIntegral $ fromSing s
        d = fromIntegral $ D.rows mat
    {-# INLINE eigS' #-}

    eigSH' s mat EigenOptions{..}
        | D.all (==0) mat = (D.replicate 0, D.replicate 1)
        | otherwise = unsafePerformIO $ do
            m1 <- CM.new
            m2 <- CM.new
            ( fmap Internal.computationInfo $ unsafeWith' m1 $ \v1 _ _ ->
                unsafeWith' m2 $ \v2 _ _ -> unsafeWith mat $ \v n _ ->
                    Internal.c_eigsh k v1 v2 v n ncv' maxit _tol mode sigma $
                    fromIntegral $ fromEnum _sort_rule ) >>= maybe
                    ( do
                        m1' <- C.unsafeFreeze m1
                        m2' <- C.unsafeFreeze m2
                        return (m1', m2') ) error
      where
        ncv' = case _ncv of
            Just x -> fromIntegral x
            Nothing -> min d $ max 20 $ 2 * k + 1
        maxit = case _maxit of
            Just x -> fromIntegral x
            Nothing -> d * 10
        mode = case _sigma of
            Nothing -> 0
            _ -> 1
        sigma = fromMaybe 0 _sigma
        k = fromIntegral $ fromSing s
        d = fromIntegral $ D.rows mat
    {-# INLINE eigSH' #-}

    geigSH' s matA matB EigenOptions{..} = unsafePerformIO $ do
        m1 <- CM.new
        m2 <- CM.new
        ( fmap Internal.computationInfo $ unsafeWith' m1 $ \v1 _ _ ->
            unsafeWith' m2 $ \v2 _ _ -> unsafeWith matA $ \v n _ ->
            unsafeWithS matB $ \pv pin po _ _ size ->
            Internal.c_geigsh k v1 v2 v n pv po pin size
            ncv' maxit _tol $ fromIntegral $ fromEnum _sort_rule ) >>= maybe
            ( do
                m1' <- C.unsafeFreeze m1
                m2' <- C.unsafeFreeze m2
                return (m1', m2') ) error
      where
        ncv' = case _ncv of
            Just x -> fromIntegral x
            Nothing -> min d $ max 20 $ 2 * k + 1
        maxit = case _maxit of
            Just x -> fromIntegral x
            Nothing -> d * 10
        k = fromIntegral $ fromSing s
        d = fromIntegral $ D.rows matA
    {-# INLINE geigSH' #-}

    cholesky mat = flip withFun1 mat $
        \code p1 c1 _ p2 _ _ -> Internal.c_cholesky code p1 p2 c1
    {-# INLINE cholesky #-}

instance Factorization S.SparseMatrix where
    eigS' s mat EigenOptions{..} = unsafePerformIO $ do
        m1 <- CM.new
        m2 <- CM.new
        ( fmap Internal.computationInfo $ unsafeWith' m1 $ \v1 _ _ ->
            unsafeWith' m2 $ \v2 _ _ -> unsafeWithS mat $ \pv pin po n _ size ->
                Internal.c_seigs k v1 v2 pv po pin n size
                ncv' maxit _tol mode sigma $ fromIntegral $ fromEnum _sort_rule
                ) >>= maybe
            ( do
                m1' <- C.unsafeFreeze m1
                m2' <- C.unsafeFreeze m2
                return (m1', m2') ) error
      where
        ncv' = case _ncv of
            Just x -> fromIntegral x
            Nothing -> min d $ max 20 $ 2 * k + 1
        maxit = case _maxit of
            Just x -> fromIntegral x
            Nothing -> d * 10
        mode = case _sigma of
            Nothing -> 0
            _ -> 1
        sigma = fromMaybe 0 _sigma
        k = fromIntegral $ fromSing s
        d = fromIntegral $ D.rows mat
    {-# INLINE eigS' #-}

    eigSH' s mat EigenOptions{..} = unsafePerformIO $ do
        m1 <- CM.new
        m2 <- CM.new
        ( fmap Internal.computationInfo $ unsafeWith' m1 $ \v1 _ _ ->
            unsafeWith' m2 $ \v2 _ _ -> unsafeWithS mat $ \pv pin po n _ size ->
                Internal.c_seigsh k v1 v2 pv po pin n size
                ncv' maxit _tol mode sigma $ fromIntegral $ fromEnum _sort_rule
                ) >>= maybe
                ( do
                    m1' <- C.unsafeFreeze m1
                    m2' <- C.unsafeFreeze m2
                    return (m1', m2') ) error
      where
        ncv' = case _ncv of
            Just x -> fromIntegral x
            Nothing -> min d $ max 20 $ 2 * k + 1
        maxit = case _maxit of
            Just x -> fromIntegral x
            Nothing -> d * 10
        mode = case _sigma of
            Nothing -> 0
            _ -> 1
        sigma = fromMaybe 0 _sigma
        k = fromIntegral $ fromSing s
        d = fromIntegral $ D.rows mat
    {-# INLINE eigSH' #-}

    cholesky = undefined

type family R a where
    R Float = Float
    R Double = Double
    R (Complex Double) = Double
    R (Complex Float) = Float

zeros :: (SingI m, SingI n) => Matrix m n Double
zeros = D.replicate 0
{-# INLINE zeros #-}

ones :: (SingI m, SingI n) => Matrix m n Double
ones = D.replicate 1
{-# INLINE ones #-}

-- | The inverse of a dense matrix.
inverse :: (SingI n, Numeric a) => Matrix n n a -> Matrix n n a
inverse = withFun1 Internal.c_inverse
{-# INLINE inverse #-}

-- | Compute the full eigendecomposition for dense matrix.
eig :: forall n . SingI n
    => Matrix n n Double
    -> (Matrix n 1 (Complex Double), Matrix n n (Complex Double))
eig mat = unsafePerformIO $ do
    m1 <- CM.new
    m2 <- CM.new
    _ <- unsafeWith' m1 $ \v1 _ _ -> unsafeWith' m2 $ \v2 _ _ ->
        unsafeWith mat $ \v n _ -> Internal.c_eig v1 v2 v n
    m1' <- C.unsafeFreeze m1
    m2' <- C.unsafeFreeze m2
    return (m1', m2')
{-# INLINE eig #-}

-- | Compute the full singular value decomposition for dense matrix.
svd :: forall n p a m. (Numeric (R a), Numeric a, SingI n, SingI p, SingI m, m ~ Min n p)
    => Matrix n p a
    -> (Matrix n m a, Matrix m 1 (R a), Matrix p m a)
svd mat = unsafePerformIO $ do
    mu <- CM.new
    ms <- CM.new
    mv <- CM.new
    checkResult $ unsafeWith' mu $ \pu _ _ -> unsafeWith' ms $ \ps _ _ ->
        unsafeWith' mv $ \pv _ _ -> unsafeWith mat $ \px r c ->
            Internal.c_bdcsvd (foreignType (undefined :: a))
                pu ps pv px r c
    u <- C.unsafeFreeze mu
    s <- C.unsafeFreeze ms
    v <- C.unsafeFreeze mv
    return (u, s, v)
{-# INLINE svd #-}

-- | Condition number.
cond :: ( Numeric a, Numeric (R a), Ord (R a), Fractional (R a)
        , SingI n, SingI m, SingI (Min n m))
     => Matrix n m a -> R a
cond mat = VS.maximum val / VS.minimum val
  where
    val = VS.filter (/=0) $ D.flatten s
    (_,s,_) = svd mat
{-# INLINE cond #-}


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE DeriveGeneric       #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

import Control.Monad.Random
import Data.Binary                  as B
import Data.Singletons
import Data.Singletons.Prelude
import Data.Singletons.TypeLits
import GHC.Generics                 (Generic)
import Numeric.LinearAlgebra.Static

-- * Weights
--

data Weights i o = W { wBiases :: !(R o)
                     , wNodes  :: !(L o i)
                     }
  deriving (Show, Generic)

instance (KnownNat i, KnownNat o) => Binary (Weights i o)

randomWeights :: (MonadRandom m, KnownNat i, KnownNat o)
              => m (Weights i o)
randomWeights = do
    s1 :: Int <- getRandom
    s2 :: Int <- getRandom
    let wB = randomVector  s1 Uniform * 2 - 1
        wN = uniformSample s2 (-1) 1
    return $ W wB wN

runLayer :: (KnownNat i, KnownNat o)
         => Weights i o
         -> R i
         -> R o
runLayer (W wB wN) v = wB + wN #> v

-- * Network
--

data Network :: Nat -> [Nat] -> Nat -> * where
    O     :: !(Weights i o)
          -> Network i '[] o
    (:&~) :: KnownNat h
          => !(Weights i h)
          -> !(Network h hs o)
          -> Network i (h ': hs) o
infixr 5 :&~

deriving instance (KnownNat i, KnownNat o) => Show (Network i hs o)

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

runNet :: (KnownNat i, KnownNat o)
       => Network i hs o
       -> R i
       -> R o
runNet (O w)      !v = logistic (runLayer w v)
runNet (w :&~ n') !v = let v' = logistic (runLayer w v)
                       in  runNet n' v'

hiddenStruct :: Network i hs o -> [Integer]
hiddenStruct = \case
    O _    -> []
    _ :&~ (n' :: Network h hs' o)
           -> natVal (Proxy @h)
            : hiddenStruct n'

randomNet' :: forall m i hs o. (MonadRandom m, KnownNat i, KnownNat o)
           => Sing hs -> m (Network i hs o)
randomNet' = \case
    SNil            ->     O <$> randomWeights
    SNat `SCons` ss -> (:&~) <$> randomWeights <*> randomNet' ss

randomNet :: forall m i hs o. (MonadRandom m, KnownNat i, SingI hs, KnownNat o)
          => m (Network i hs o)
randomNet = randomNet' sing

putNet :: (KnownNat i, KnownNat o)
       => Network i hs o
       -> Put
putNet = \case
    O w     -> put w
    w :&~ n -> put w *> putNet n

getNet :: forall i hs o. (KnownNat i, KnownNat o)
       => Sing hs
       -> Get (Network i hs o)
getNet = \case
    SNil            ->     O <$> get
    SNat `SCons` ss -> (:&~) <$> get <*> getNet ss

instance (KnownNat i, SingI hs, KnownNat o) => Binary (Network i hs o) where
    put = putNet
    get = getNet sing

-- * OpaqueNet
--

data OpaqueNet :: Nat -> Nat -> * where
    ONet :: Network i hs o -> OpaqueNet i o

runOpaqueNet :: (KnownNat i, KnownNat o)
             => OpaqueNet i o
             -> R i
             -> R o
runOpaqueNet (ONet n) x = runNet n x

numHiddens :: OpaqueNet i o -> Int
numHiddens (ONet n) = go n
  where
    go :: Network i hs o -> Int
    go = \case
        O _      -> 0
        _ :&~ n' -> 1 + go n'

randomONet :: (MonadRandom m, KnownNat i, KnownNat o)
           => [Integer]
           -> m (OpaqueNet i o)
randomONet hs = case toSing hs of
                  SomeSing ss -> ONet <$> randomNet' ss

putONet :: (KnownNat i, KnownNat o)
        => OpaqueNet i o
        -> Put
putONet (ONet net) = do
    put (hiddenStruct net)
    putNet net

getONet :: (KnownNat i, KnownNat o)
        => Get (OpaqueNet i o)
getONet = do
    hs <- get
    withSomeSing hs $ \ss ->
      ONet <$> getNet ss

instance (KnownNat i, KnownNat o) => Binary (OpaqueNet i o) where
    put = putONet
    get = getONet

-- * OpaqueNet'
--

type OpaqueNet' i o r = (forall hs. Network i hs o -> r) -> r

oNet' :: Network i hs o -> OpaqueNet' i o r
oNet' n = \f -> f n

runOpaqueNet' :: (KnownNat i, KnownNat o)
              => OpaqueNet' i o (R o)
              -> R i
              -> R o
runOpaqueNet' oN x = oN (\n -> runNet n x)
--            :: ((forall hs. Network i hs o -> R o) -> R o)
--            -> R i
--            -> R o

numHiddens' :: OpaqueNet' i o Int -> Int
numHiddens' oN = oN go
  where
    go :: Network i hs o -> Int
    go = \case
        O _      -> 0
        _ :&~ n' -> 1 + go n'
--          :: ((forall hs. Network i hs o -> Int) -> Int)
--          -> Int

withRandomONet' :: (MonadRandom m, KnownNat i, KnownNat o)
                => [Integer]
                -> (forall hs. Network i hs o -> m r)
                -> m r
--         aka, => [Integer]
--              -> OpaqueNet' i o (m r)
withRandomONet' hs f = withSomeSing hs $ \ss -> do
                         net <- randomNet' ss
                         f net

putONet' :: (KnownNat i, KnownNat o)
         => OpaqueNet' i o Put
         -> Put
putONet' oN = oN $ \net -> do
                      put (hiddenStruct net)
                      putNet net

getONet' :: (KnownNat i, KnownNat o)
         => (forall hs. Network i hs o -> Get r)
         -> Get r
--  aka, => OpaqueNet' i o (Get r)
getONet' f = do
    hs <- get :: Get [Integer]
    withSomeSing hs $ \ss -> do
      n <- getNet ss
      f n

main :: IO ()
main = do
    putStrLn "What hidden layer structure do you want?"
    hs <- readLn
    n  <- randomONet hs
    case n of
      ONet (net :: Network 10 hs 3) -> do
        print net
        -- blah blah stuff with our dynamically generated net

main' :: IO ()
main' = do
    putStrLn "What hidden layer structure do you want?"
    hs <- readLn
    withRandomONet' hs $ \(net :: Network 10 hs 3) -> do
      print net
      -- blah blah stuff with our dynamically generated net

-- * SomeNet
--

data SomeNet where
    SNet :: (KnownNat i, KnownNat o)
         => Network i hs o
         -> SomeNet

withONet :: SomeNet
         -> (forall i o. (KnownNat i, KnownNat o) => OpaqueNet i o -> r)
         -> r
withONet (SNet n) f = f (ONet n)

randomSNet :: forall m. MonadRandom m
           => Integer
           -> [Integer]
           -> Integer
           -> m SomeNet
randomSNet i hs o =
    withSomeSing i  $ \(SNat :: Sing (i  :: Nat  )) ->
    withSomeSing hs $ \(ss   :: Sing (hs :: [Nat])) ->
    withSomeSing o  $ \(SNat :: Sing (o  :: Nat  )) ->
      SNet <$> (randomNet' ss :: m (Network i hs o))

withRandomSNet' :: forall m r. MonadRandom m
                => Integer
                -> [Integer]
                -> Integer
                -> (forall i hs o. (KnownNat i, KnownNat o) => Network i hs o -> m r)
                -> m r
withRandomSNet' i hs o f =
    withSomeSing i  $ \(SNat :: Sing (i  :: Nat  )) ->
    withSomeSing hs $ \(ss   :: Sing (hs :: [Nat])) ->
    withSomeSing o  $ \(SNat :: Sing (o  :: Nat  )) -> do
      n <- randomNet' ss :: m (Network i hs o)
      f n

instance Binary SomeNet where
    put (SNet (net :: Network i hs o)) = do
        put $ natVal (Proxy @i)
        put $ hiddenStruct net
        put $ natVal (Proxy @o)
        putNet net
    get = do
        i  <- get :: Get Integer
        hs <- get :: Get [Integer]
        o  <- get :: Get Integer
        withSomeSing i    $ \(SNat :: Sing (i  :: Nat  )) ->
          withSomeSing hs $ \(ss   :: Sing (hs :: [Nat])) ->
          withSomeSing o  $ \(SNat :: Sing (o  :: Nat  )) -> do
            n <- getNet ss :: Get (Network i hs o)
            return (SNet n)


module Main where

import MnistLoader
import Network

import Numeric.LinearAlgebra as L

main :: IO ()
main = testMnist

testMnist = do
  net <- initNet [784, 30, 10]
  exemplars <- loadDataWrapper
  test <- loadTestWrapper

  -- let eta = 0.050
  -- let batchSize = 20
  -- let epochs = 30

  let eta = 0.5
  let batchSize = 16
  let epochs = 30

  putStrLn "Running Stochastic Gradient Descent "
  putStrLn $ "eta: " ++ (show eta)
  putStrLn $ "Batch Size: " ++ (show batchSize)
  putStrLn $ "Epochs: " ++ (show epochs)

  trainedNetwork <- sgd net exemplars test epochs 0 batchSize eta
  return ()

test = do
  net <- initNet [3, 5]

  putStrLn "The Net"
  putStrLn $ show net

  putStrLn "Feed Forward"
  putStrLn $ show $ feedForward net (vect 3 [0, 1, 0])

  putStrLn "A ForwardPass"
  putStrLn $ show $ forwardPass net (vect 3 [0, 1, 0])

  putStrLn "Entire backprop"
  let eb = backprop net (vect 3 [0, 1, 0]) (vect 5 [1, 0, 1, 0, 1])
  putStrLn $ show $ length eb
  putStrLn $ show $ eb

  putStrLn "A Minibatch"
  let x = vect 3 [0,1,0]
  let y = vect 5 [1, 0, 1, 0, 1]
  let batch = [(x, y), (x, y), (x, y)]
  let newn = updateMiniBatch net batch 3.0
  putStrLn $ show newn
  return newn

vect n = (n><1)


{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE PolyKinds                 #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# OPTIONS_GHC -fwarn-incomplete-patterns #-}

module Math.Regression.LeastSquares where

import qualified Math.HMatrixUtils             as HU
import qualified Math.Regression.Regression    as RE

import qualified Polysemy                      as P
import qualified Knit.Effect.Logger            as Log

import qualified Data.List                     as List

import           Numeric.LinearAlgebra          ( (#>)
                                                , (<.>)
                                                , (<\>)
                                                )
import qualified Numeric.LinearAlgebra         as LA
import           Numeric.LinearAlgebra.Data     ( Matrix
                                                , R
                                                , Vector
                                                )

-- matrix dimensions given in (row x column) form
-- if A is (m x n), its transpose, A' is (n x m)
-- ||A|| is the Frobenius norm of A, the sum of the squares of the elements.
-- if A (m x l) and B (m x k) have the same number of rows, [A B] is the (m x (l+k)) matrix formed by appending the columns of B to the columns of A
-- A: an (m x n) matrix, of m observations of n observables
-- B: an (m x d) matrix, of observations (often d=1 and B is just a length m column-vector)
-- Q: an ((n+d)m x (n+d)m) matrix of covariances, the covariance matrix of [A B].  Given?  Computed?

-- we want an (n x d) matrix X which solves, or "best solves", AX = B.

-- OLS: minimize ||AX - B||
ordinaryLS
  :: Log.LogWithPrefixesLE effs
  => Bool
  -> Matrix R
  -> Vector R
  -> P.Sem effs (RE.RegressionResult R)
ordinaryLS withConstant mA vB = do
  let mAwc   = if withConstant then addBiasCol (LA.size vB) mA else mA -- add a constant, e.g., the b in y = mx + b
      (n, p) = LA.size mAwc
      dof    = n - p
  HU.checkVectorMatrix "b" "A" vB mAwc  -- vB <> mA is legal, b has same length as A has rows
  let vX  = mAwc <\> vB
      vU  = vB - (mAwc #> vX) -- residuals
      mse = (vU <.> vU) / realToFrac dof
      cov = LA.scale mse (LA.inv $ LA.tr mAwc LA.<> mAwc)
  RE.FitStatistics rSq aRSq fStat <- RE.goodnessOfFit p vB Nothing vU
  return $ RE.RegressionResult (RE.estimates cov vX)
                               (realToFrac dof)
                               mse
                               rSq
                               aRSq
                               fStat
                               cov

weightedLS
  :: Log.LogWithPrefixesLE effs
  => Bool
  -> Matrix R
  -> Vector R
  -> Vector R
  -> P.Sem effs (RE.RegressionResult R)
weightedLS withConstant mA vB vW = do
  HU.checkEqualVectors "b" "w" vB vW
  let mW     = LA.diag vW
      vWB    = mW #> vB
      mAwc   = if withConstant then addBiasCol (LA.size vB) mA else mA -- add a constant, e.g., the b in y = mx + b
      (_, p) = LA.size mAwc
  HU.checkVectorMatrix "b" "A" vB mAwc
  let mWA  = mW LA.<> mAwc
      vX   = mWA <\> vWB
      vU   = vB - (mAwc #> vX)
      vWU  = mW #> vU
      sumW = LA.sumElements vW
      effN = (sumW * sumW) / (vW <.> vW)
      mse  = effN / (effN - realToFrac p) * (vWU <.> vWU) / sumW
      cov  = LA.scale mse (LA.inv $ LA.tr mAwc LA.<> mAwc)
  RE.FitStatistics rSq aRSq fStat <- RE.goodnessOfFit p vB (Just vW) vU
  return $ RE.RegressionResult (RE.estimates cov vX)
                               (effN - realToFrac p)
                               mse
                               rSq
                               aRSq
                               fStat
                               cov

totalLS
  :: Log.LogWithPrefixesLE effs
  => Bool
  -> Matrix R
  -> Vector R
  -> P.Sem effs (RE.RegressionResult R)
totalLS withConstant mA vB = do
  let mAwc   = if withConstant then addBiasCol (LA.size vB) mA else mA -- add a constant, e.g., the b in y = mx + b
      (n, p) = LA.size mAwc
      dof    = realToFrac (n - p)
  HU.checkVectorMatrix "b" "Awc" vB mAwc
  let mAB      = mAwc LA.||| LA.asColumn vB
      (_, mV') = LA.rightSV mAB
--      gSV = sv V.! 0
--      tol = gSV * LA.eps
      sV22     = mV' LA.! p LA.! p
      vV12     = List.head $ LA.toColumns $ LA.subMatrix (0, p) (p, 1) mV' --LA.?? (LA.DropLast 1, LA.TakeLast 1)
      vX       = LA.scale (-1 / sV22) vV12 -- this is the TLS solution.  But in a shifted basis.??
--      mV2      = mV' LA.?? (LA.All, LA.TakeLast 1)
--      mABt     = mAB LA.<> mV2 LA.<> (LA.tr mV2)
--      mAt      = mABt LA.?? (LA.All, LA.DropLast 1)
      vBfit    = mAwc #> vX --(mAwc - mAt) #> vX -- this is confusing.
      vU       = vB - vBfit
--      (rSq, aRSq) = RE.goodnessOfFit p vB vU --(vB - mA #> vX)
      mse      = (vU <.> vU) / realToFrac (n - p)
      cov      = LA.scale mse (LA.inv $ LA.tr mAwc LA.<> mAwc)
  RE.FitStatistics rSq aRSq fStat <- RE.goodnessOfFit p vB Nothing vU
  return $ RE.RegressionResult (RE.estimates cov vX) dof mse rSq aRSq fStat cov

weightedTLS
  :: Log.LogWithPrefixesLE effs
  => Bool
  -> Matrix R
  -> Vector R
  -> Vector R
  -> P.Sem effs (RE.RegressionResult R)
weightedTLS withConstant mA vB vW = do
  HU.checkEqualVectors "b" "w" vB vW
  let mW     = LA.diag vW
      vWB    = mW #> vB
      mAwc   = if withConstant then addBiasCol (LA.size vB) mA else mA -- add a constant, e.g., the b in y = mx + b
      (_, p) = LA.size mAwc
  HU.checkVectorMatrix "b" "Awc" vB mAwc
  let mWA      = mW LA.<> mAwc
      mWAB     = mWA LA.||| LA.asColumn vWB
      (_, mV') = LA.rightSV mWAB
--      gSV = sv V.! 0
--      tol = gSV * LA.eps
      sV22     = mV' LA.! p LA.! p
      vV12     = List.head $ LA.toColumns $ LA.subMatrix (0, p) (p, 1) mV' --LA.?? (LA.DropLast 1, LA.TakeLast 1)
      vX       = LA.scale (-1 / sV22) vV12 -- this is the WTLS solution.  But in a shifted basis.??
--      mV2      = mV' LA.?? (LA.All, LA.TakeLast 1)
--      mWABt    = mWAB LA.<> mV2 LA.<> (LA.tr mV2)
--      mWAt     = mWABt LA.?? (LA.All, LA.DropLast 1)
--      vBfit    = mAwc #> vX -- (mWA - mWAt) #> vX -- this is confusing
      vU       = vB - (mAwc #> vX)
      vWU      = mW #> vU
      sumW     = LA.sumElements vW
      effN     = (sumW * sumW) / (vW <.> vW)
      dof      = effN - realToFrac p
      mse      = effN / dof * (vWU <.> vWU) / sumW
--      (rSq, aRSq) = RE.goodnessOfFit p vWB vWU --(vB - mA #> vX)
      cov      = LA.scale mse (LA.inv $ LA.tr mAwc LA.<> mAwc)
  RE.FitStatistics rSq aRSq fStat <- RE.goodnessOfFit p vB (Just vW) vU
  return $ RE.RegressionResult (RE.estimates cov vX) dof mse rSq aRSq fStat cov

addBiasCol :: Int -> Matrix R -> Matrix R
addBiasCol rows mA =
  let colList = List.replicate rows 1
  in  if LA.size mA == (0, 0)
        then LA.matrix 1 colList
        else LA.col colList LA.||| mA


module HVX.Internal.TestUtil
  ( fpequalsApprox
  , pairsOf
  , positiveVectors
  , smallEvenPowersGreaterThanOne
  , smallNonIntPowersGreaterThanOne
  , smallVectors
  , vectors
  ) where

import Control.Monad
import Numeric.LinearAlgebra
import Test.QuickCheck hiding ( (><) )

import HVX.Internal.Matrix
import HVX.Internal.Util

fpequalsApprox :: Double -> Double -> Bool
fpequalsApprox a b
  | a == 0 && b == 0           = True
  | max (abs a) (abs b) < 1e-6 = True
  | otherwise                  = abs ( (a - b) / denom) < tolerance
  where
    denom = max (abs a) (abs b)
    tolerance = 5e-2  -- 5% tolerance is why we call it Approx.

pairsOf :: Gen a -> Gen b -> Gen (a, b)
pairsOf a b = do
  x <- a
  y <- b
  return (x, y)

positiveVectors :: Gen Mat
positiveVectors = do
  nonEmptyL <- arbitrary :: Gen (NonEmptyList Double)
  let l = getNonEmpty nonEmptyL
  return $ abs $ (length l><1) l

smallEvenPowersGreaterThanOne :: Gen Double
smallEvenPowersGreaterThanOne =
  liftM fromIntegral $
    (arbitrary :: Gen Integer) `suchThat` \x -> x > 1 && x <= 20 && even x

smallNonIntPowersGreaterThanOne :: Gen Double
smallNonIntPowersGreaterThanOne =
  arbitrary `suchThat` \x -> x > 1 && x <= 20 && not (isInteger x)

smallVectors :: Gen Mat
smallVectors = do
  len <- choose (1, 100)
  l <- sequence [choose (-10, 10) | _ <- [1..len]]
  return $ (len><1) l

vectors :: Gen Mat
vectors = do
  nonEmptyL <- arbitrary :: Gen (NonEmptyList Double)
  let l = getNonEmpty nonEmptyL
  return $ (length l><1) l


{-# LANGUAGE TypeOperators,TypeFamilies,MultiParamTypeClasses,FlexibleContexts, BangPatterns,PolyKinds, DataKinds, GADTs, GeneralizedNewtypeDeriving,DeriveDataTypeable,FlexibleInstances,ScopedTypeVariables #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE CPP #-}
{- | Fixed point arithmetic

Fixed point arithmetic with signed, unsigned and saturation

-}
module Fixed(
      nbFractionalBits
    , HasDoubleRepresentation(..)
    , FixedPoint(..)
    , Resolution(..)
    , AMulConstraint(..)
    , ConvertConstraint(..)
    , withSaturation
    , withoutSaturation
    , withRounding
    , withNoRounding
    , Fixed
    , Saturation(..)
    , Rounding(..)
    , Int16
    , Int32
    , Int40
    , Word16
    , Word32
    , Word40
    , Int128
    , Conversion(..)
    , amul
    , amulc
    , roundf
    ) where

import Data.Typeable
import GHC.TypeLits
import Data.Word
import Data.Int
import Data.Bits
import Text.Printf
import Data.Ratio
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as M
import Control.Monad(liftM)
import Common(HasDoubleRepresentation(..))
import Data.Complex
import SpecialInt
import Control.DeepSeq
import System.Random

import Debug.Trace

debug a = trace (show a) a

instance HasDoubleRepresentation Double where
    toDouble = id
    fromDouble = id

instance HasDoubleRepresentation Float where
    toDouble = realToFrac
    fromDouble = realToFrac

instance HasDoubleRepresentation Int where
    toDouble = fromIntegral
    fromDouble = floor

newtype Fixed :: * -> Nat -> Saturation -> Rounding -> * where Fixed :: a -> Fixed a n sa r

data Saturation = Sat | Unsat deriving(Eq)
data Rounding = RO | NR  deriving(Eq)

newtype instance Sing (n :: Saturation) = SatC Saturation
newtype instance Sing (n :: Rounding) = RoundC Rounding

instance SingI Sat where
  sing = SatC Sat

instance SingI Unsat where
  sing = SatC Unsat

instance SingE (Kind :: Saturation) Saturation where
  fromSing (SatC n) = n

instance SingI RO where
  sing = RoundC RO

instance SingI NR where
  sing = RoundC NR

instance SingE (Kind :: Rounding) Rounding where
  fromSing (RoundC n) = n

newtype instance U.MVector s (Fixed a n sat r)  = MVFixed (U.MVector s a)
newtype instance U.Vector    (Fixed a n sat r) = VFixed (U.Vector a)

instance (U.Unbox a) => M.MVector U.MVector (Fixed a n s r) where
   {-# INLINE basicLength #-}
   {-# INLINE basicUnsafeSlice #-}
   {-# INLINE basicOverlaps #-}
   {-# INLINE basicUnsafeNew #-}
   {-# INLINE basicUnsafeRead #-}
   {-# INLINE basicUnsafeWrite #-}
   basicLength (MVFixed v) = M.basicLength v
   basicUnsafeSlice a b (MVFixed v) = MVFixed $ M.basicUnsafeSlice a b v
   basicOverlaps (MVFixed a) (MVFixed b) = M.basicOverlaps a b
   basicUnsafeNew n = MVFixed `liftM` M.basicUnsafeNew n
   basicUnsafeRead (MVFixed v) i = do
    r <- M.basicUnsafeRead v i
    return (Fixed r)
   basicUnsafeWrite (MVFixed v) i (Fixed r) = do
    M.basicUnsafeWrite v i r

instance (U.Unbox a) => G.Vector U.Vector (Fixed a n s r) where
   {-# INLINE basicLength #-}
   {-# INLINE basicUnsafeFreeze #-}
   {-# INLINE basicUnsafeThaw #-}
   {-# INLINE basicUnsafeSlice #-}
   {-# INLINE basicUnsafeIndexM #-}
   basicLength (VFixed v) = G.basicLength v
   basicUnsafeFreeze (MVFixed v) = VFixed `liftM`G.basicUnsafeFreeze v
   basicUnsafeThaw (VFixed v) = MVFixed `liftM`G.basicUnsafeThaw v
   basicUnsafeSlice a b (VFixed v) = VFixed (G.basicUnsafeSlice a b v)
   basicUnsafeIndexM (VFixed v) i = Fixed `liftM` G.basicUnsafeIndexM v i

instance (RealFloat a, U.Unbox a) => U.Unbox (Fixed a n sat r)

{-# INLINE getFract #-}
getFract :: Fixed a (n :: Nat) sa r -> Sing n -> Integer
getFract _ s = fromSing s

{-# INLINE nbFractionalBits #-}
nbFractionalBits :: SingI n => Fixed a (n :: Nat) sa r -> Int
nbFractionalBits f = fromIntegral $ getFract f sing

{-# INLINE getSaturationMode #-}
getSaturationMode :: Fixed a n sa r -> Sing sa -> Saturation
getSaturationMode _ s = fromSing s

{-# INLINE saturationMode #-}
saturationMode :: SingI sa => Fixed a n sa r -> Saturation
saturationMode f = getSaturationMode f sing

{-# INLINE getRoundingMode #-}
getRoundingMode :: Fixed a n s (r :: Rounding) -> Sing r -> Rounding
getRoundingMode _ s = fromSing s

{-# INLINE roundingMode #-}
roundingMode :: SingI r => Fixed a n sa r -> Rounding
roundingMode f = getRoundingMode f sing

{-# INLINE saturateWithMode #-}
saturateWithMode :: (SingI sa, SaturateConstraint a)
                 => Fixed a n sa r
                 -> SuperInt a
                 -> a
saturateWithMode f su | saturationMode f == Unsat = fromIntegral su
                      | otherwise = saturate (witness f) su
  where
    witness (Fixed a) = a

{-# INLINE roundWithMode #-}
roundWithMode :: (SingI r, Bits b, Ord b, Num b)
              => Fixed a n sa r
              -> Int
              -> Int
              -> b
              -> b
roundWithMode f na nb su = rounding (roundingMode f) na nb su

{-# INLINE roundf #-}
roundf :: (SingI n, SingI r, SingI sa, SaturateConstraint a, Bits (SuperInt a))
       => Fixed a n sa r
       -> Fixed a n sa r
roundf f@(Fixed a) =
  let na = nbFractionalBits f
  in
  Fixed $ saturateWithMode f (roundWithMode f na 0 (fromIntegral a))

{-# INLINE rounding #-}
rounding :: (Bits a, Num a, Ord a)
         => Rounding
         -> Int -- ^ Current fractional bits
         -> Int -- ^ Future fractional bits
         -> a
         -> a
rounding r cf ff v | cf <= ff = v
                   | otherwise =
              case r of
                RO -> nearR
                NR -> v
  where
    rpos = cf - ff
    roundingMask = complement ((1 `shiftL` rpos) - 1)
    nearR = (v + (1 `shiftL` (rpos - 1))) .&. roundingMask --  near

instance Eq a => Eq (Fixed a n sa r) where
    (Fixed a) == (Fixed b) = a == b

instance Ord a => Ord (Fixed a n sa r) where
    compare (Fixed a) (Fixed b) = compare a b

type AMulConstraint a b = (Num (BaseValue b), Bits b, Integral b, Integral (SuperInt b), Num (SuperInt b), Integral a, NumberInfo (SuperInt b), NumberInfo a, NumberInfo b, RawValue b, Num b)

{-# INLINE amul #-}
amul :: (SingI na, SingI nb, SingI s, AMulConstraint a b)
     => Fixed a na s r
     -> Fixed a nb s r
     -> Fixed b (na + nb) s r
amul fa@(Fixed a) (Fixed b) =
        let la = fromIntegral a
            lb = fromIntegral b
            r = (la * lb)
            result = Fixed (saturateWithMode result r)
        in
        result

{-# INLINE amulc #-}
amulc :: (SingI na,SingI nb, SingI s, AMulConstraint a b,Num (Fixed b (na+nb) s r))
      => Complex (Fixed a na s r)
      -> Complex (Fixed a nb s r)
      -> Complex (Fixed b (na+nb) s r)
amulc (xr :+ xi) (yr :+ yi) = (amul xr yr - amul xi yi) :+ (amul xr yi + amul xi yr)

type ConvertConstraint a b = (Integral a, Bits a, Bits (SuperInt b), NumberInfo a, NumberInfo b, SaturateConstraint b)

class Conversion a b where
  {-# INLINE convert #-}
  convert :: (SingI na, SingI nb, SingI sb, SingI r, ConvertConstraint a b)
          => Fixed a na sa r
          -> Fixed b nb sb r
  convert fa@(Fixed a) = fb
     where
       fb = Fixed b
       b | la <= lb = saturateWithMode fb (fromIntegral a `shift` shiftValue)
         | otherwise = saturateWithMode fb (fromIntegral ((roundWithMode fa la lb a) `shift` shiftValue))
       sa = nbFractionalBits fa
       sb = nbFractionalBits fb
       shiftValue = sb - sa
       la = nbBits a
       lb = nbBits b

instance Conversion Int32 Int16 where
instance Conversion Int16 Int32 where
instance Conversion Int40 Int16 where
instance Conversion Int64 Int16 where
instance Conversion Int16 Int64 where
instance Conversion Int16 Int40 where
instance Conversion Int32 Int64 where
instance Conversion Int64 Int32 where
instance Conversion Int32 Int40 where
instance Conversion Int40 Int32 where

instance Conversion Word32 Word16 where
instance Conversion Word16 Word32 where
instance Conversion Word40 Word16 where
instance Conversion Word16 Word40 where

instance Conversion Int16 Int16 where
instance Conversion Int32 Int32 where
instance Conversion Int40 Int40 where

instance Conversion Word16 Word16 where
instance Conversion Word32 Word32 where
instance Conversion Word40 Word40 where

{-# INLINE genericOperator #-}
genericOperator :: (SingI s, Bits a, Integral (SuperInt a), Num (BaseValue a), Num (SuperInt a), Integral a, NumberInfo (SuperInt a), NumberInfo a, RawValue a)
                => (SuperInt a -> SuperInt a -> SuperInt a)
                -> Fixed a n s r
                -> Fixed a n s r
                -> Fixed a n s r
genericOperator op f@(Fixed a) (Fixed b) =
        let la = (fromIntegral a)
            lb = (fromIntegral b)
        in
        Fixed (saturateWithMode f $ op la lb)

{-# INLINE genericMulOperator #-}
genericMulOperator :: (SingI s, SingI n, SingI r, Bits a, Bits (SuperInt a), Integral (SuperInt a), Num (BaseValue a), Num (SuperInt a), Integral a, NumberInfo (SuperInt a), NumberInfo a, RawValue a)
                   => Fixed a n s r
                   -> Fixed a n s r
                   -> Fixed a n s r
genericMulOperator f@(Fixed a) (Fixed b) =
        let la = (fromIntegral a)
            lb = (fromIntegral b)
            na = (nbFractionalBits f)
            r = (roundWithMode f (na + na) na (la * lb)) `shiftR` (nbFractionalBits f)
        in
        Fixed (saturateWithMode f $ r)

{-# INLINE genericAbs #-}
genericAbs (Fixed a) | a == minBound = Fixed maxBound
                     | otherwise = Fixed (abs a)

{-# INLINE genericFromInteger #-}
genericFromInteger a = r
     where
        b = a `shiftL` (nbFractionalBits r)
        theMax = maxBound
        theMin = minBound
        r | b > (fromIntegral theMax) = Fixed theMax
          | b < (fromIntegral theMin) = Fixed theMin
          | otherwise = Fixed (fromIntegral b)

{-# INLINE genericProperFraction #-}
genericProperFraction f@(Fixed a) =
  let l = nbFractionalBits f
      b = a `shiftR` l
      na = f - fromIntegral b
  in
  (fromIntegral b,na)

{-

For Random instance

-}

{-# INLINE genericRandomR #-}
genericRandomR :: (RandomGen g, Random a, FixedPoint a)
               => (Fixed a n s r, Fixed a n s r)
               -> g
               -> (Fixed a n s r, g)
genericRandomR (Fixed mi,Fixed ma) g =
  let (na,ng) = randomR (mi,ma) g
  in
  (fromRawValue na,ng)

{-# INLINE genericRandom #-}
genericRandom :: (Random (Fixed a n s r),Bounded (Fixed a n s r), RandomGen g, Random a, FixedPoint a)
              => g
              -> (Fixed a n s r, g)
genericRandom g = randomR (minBound, maxBound) g

class Resolution a where
  smallestValue :: a -> a
  maxValue :: a -> a
  minValue :: a -> a
  signedFormat :: a -> Bool
  bitWidth :: a -> Int

#define FIXED_INSTANCES(INT) \
instance NFData (Fixed INT n s r) where {\
    rnf (Fixed a) = rnf a };\
instance (SingI n, SingI s, SingI r) => Random (Fixed INT n s r) where {\
    randomR = genericRandomR \
;   random = genericRandom }; \
instance (SingI n, SingI s, SingI r) => Num (Fixed INT n s r) where {\
     (+) = genericOperator (+) \
;    (-) = genericOperator (-) \
;    abs = genericAbs \
;    (*) = genericMulOperator \
;    signum (Fixed a) = Fixed (signum a) \
;    fromInteger = genericFromInteger}; \
instance (SingI n, SingI s, SingI r) => Bounded (Fixed INT n s r) where {\
     maxBound = fromRawValue maxBound \
;    minBound = fromRawValue minBound };\
instance (SingI n, SingI s, SingI r) => Resolution (Fixed INT n s r) where {\
     smallestValue _ = fromRawValue 1 \
;    maxValue _ = maxBound \
;    minValue _ = minBound \
;    signedFormat a = signed (toRawValue a)\
;    bitWidth (Fixed a) = nbBits a};\
instance FixedPoint INT where {\
     fromRawValue = Fixed \
;    toRawValue (Fixed a) = a}; \
instance (SingI n, SingI s) => HasDoubleRepresentation (Fixed INT n s r) where {\
   toDouble = genericToDouble \
;  fromDouble = genericFromDouble}; \
instance (SingI n, SingI s, SingI r) => Fractional (Fixed INT n s r) where {\
    (/)  = genericDiv \
;   fromRational = genericFromRational}; \
instance (SingI n,SingI s, SingI r) => Real (Fixed INT n s r) where {\
  toRational = toRational . toDouble }; \
instance (SingI n, SingI s, SingI r) => RealFrac (Fixed INT n s r) where {\
  properFraction = genericProperFraction }; \
instance (SingI n, SingI s, SingI r) => Floating (Fixed INT n s r) where {\
    pi = fromDouble pi \
;   exp = fromDouble . exp . toDouble \
;   log = fromDouble . log . toDouble \
;   sin = fromDouble . sin . toDouble \
;   cos = fromDouble . cos . toDouble \
;   sinh = fromDouble . sinh . toDouble \
;   cosh = fromDouble . cosh . toDouble \
;   asin = fromDouble . asin . toDouble \
;   acos = fromDouble . acos . toDouble \
;   atan = fromDouble . atan . toDouble \
;   asinh = fromDouble . asinh . toDouble \
;   acosh = fromDouble . acosh . toDouble \
;   atanh = fromDouble . atanh . toDouble}; \
instance (SingI n, SingI s, SingI r) => RealFloat (Fixed INT n s r) where {\
   isInfinite = error ("RealFloat has no meaning for a fixed point number and is needed only because of Complex") \
;  isDenormalized = error ("RealFloat has no meaning for a fixed point number and is needed only because of Complex") \
;  isNegativeZero = error ("RealFloat has no meaning for a fixed point number and is needed only because of Complex") \
;  isIEEE = error ("RealFloat has no meaning for a fixed point number and is needed only because of Complex") \
;  isNaN = error ("RealFloat has no meaning for a fixed point number and is needed only because of Complex") \
;  encodeFloat = error ("RealFloat has no meaning for a fixed point number and is needed only because of Complex") \
;  decodeFloat = error ("RealFloat has no meaning for a fixed point number and is needed only because of Complex") \
;  floatRange = error ("RealFloat has no meaning for a fixed point number and is needed only because of Complex") \
;  floatRadix = error ("RealFloat has no meaning for a fixed point number and is needed only because of Complex") \
;  floatDigits = error ("RealFloat has no meaning for a fixed point number and is needed only because of Complex") \
;  exponent = error ("RealFloat has no meaning for a fixed point number and is needed only because of Complex") \
;  significand = error ("RealFloat has no meaning for a fixed point number and is needed only because of Complex") \
;  scaleFloat = error ("RealFloat has no meaning for a fixed point number and is needed only because of Complex") \
;  atan2 = error ("RealFloat has no meaning for a fixed point number and is needed only because of Complex") };

-- RealFloat is only required because Complex a is requiring it of a.
-- It has no meaning for a fixed point
-- Which means that some function of Complex caannot be used (like magnitude)
FIXED_INSTANCES(Int16)
FIXED_INSTANCES(Int32)
FIXED_INSTANCES(Int40)
FIXED_INSTANCES(Int64)
FIXED_INSTANCES(Int128)

FIXED_INSTANCES(Word16)
FIXED_INSTANCES(Word32)
FIXED_INSTANCES(Word40)
FIXED_INSTANCES(Word64)

instance (SingI n, Integral a) => Show (Fixed a n s r) where
    show f@(Fixed a) =
        let fp = nbFractionalBits f
            fr = 2**(- fromIntegral fp) :: Double
        in
        printf ("%f") ((fromIntegral a)*fr)

{-# INLINE withSaturation #-}
withSaturation :: Fixed n s sa r -> Fixed n s Sat r
withSaturation (Fixed a) = Fixed a

{-# INLINE withoutSaturation #-}
withoutSaturation :: Fixed n s sa r -> Fixed n s Unsat r
withoutSaturation (Fixed a) = Fixed a

{-# INLINE withRounding #-}
withRounding :: Fixed n s sa r -> Fixed n s sa RO
withRounding (Fixed a) = Fixed a

{-# INLINE withNoRounding #-}
withNoRounding :: Fixed n s sa r -> Fixed n s sa NR
withNoRounding (Fixed a) = Fixed a

{-# INLINE genericDiv #-}
genericDiv :: (SingI n, SingI s, Bits a, NumberInfo (SuperInt a), NumberInfo a, RawValue a, Num (BaseValue a), Integral a, Bits (SuperInt a),Integral (SuperInt a), Num (SuperInt a))
           => Fixed a n s r
           -> Fixed a n s r
           -> Fixed a n s r
genericDiv f@(Fixed a) (Fixed b) =
        let la = (fromIntegral a)
            lb = (fromIntegral b)
            (q,_) = (la `shiftL` (nbFractionalBits f)) `quotRem` lb
            result = q
        in
        Fixed (saturateWithMode f $ result)

{-# INLINE genericFromRational #-}
genericFromRational r = fromDouble (fromRational r)

class FixedPoint a where
    fromRawValue :: a -> Fixed a n s r
    toRawValue :: Fixed a n s r -> a

{-# INLINE genericFromDouble #-}
genericFromDouble :: (SingI n, SingI sa, Bits a, Integral (SuperInt a), FixedPoint a, Num a, Num (BaseValue a), Integral a, NumberInfo (SuperInt a), NumberInfo a, RawValue a)
                  => Double
                  -> Fixed a n sa r
genericFromDouble a = let la = saturateWithMode ra (floor (a * 2**(fromIntegral $ nbFractionalBits ra)))
                          ra = fromRawValue la
                      in
                      ra

{-# INLINE genericToDouble #-}
genericToDouble :: (SingI n, Integral a) => Fixed a n s r -> Double
genericToDouble f@(Fixed a) = fromIntegral a * 2**(- fromIntegral (nbFractionalBits f))


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DeriveGeneric #-}

module Lib
    ( genImage
    , ViewPort(..)
    ) where

import Codec.Picture
import Codec.Picture.Types
import Data.Fixed (mod')
import Data.Word (Word8)
import Data.Array
import Data.Hashable
import Data.Complex
import Control.Parallel.Strategies
import GHC.Generics

data ViewPort = ViewPort { _centerX :: Double
                         , _centerY :: Double
                         , _magnification :: Double
                         , _maxIterations :: Int
                         , _rotation :: Double
                         , _superScaling :: Int
                         } deriving (Eq, Show, Generic)

instance Hashable ViewPort

genImage :: Int -> Int -> ViewPort -> DynamicImage
genImage w h v = ImageRGB8 $ generateImage (renderImage v w h) w h

renderImage :: ViewPort -> Int -> Int -> (Int -> Int -> PixelRGB8)
renderImage (ViewPort cx cy mg mi _ s) w h =
    let aspectRatio = fromIntegral w / fromIntegral h :: Double
        (rw, rh) = if aspectRatio <= 4.0 / 3.0 then (4.0 / mg, 4.0 / mg / aspectRatio)
                                               else (3.0 / mg * aspectRatio, 3.0 / mg)
        f x y = let ps = [(x * s + ix, y * s + iy) | ix <- [0..s-1], iy <- [0..s-1]]
                    r x' = cx - (rw / 2.0) + ((fromIntegral x' + 0.5) / fromIntegral (w * s) :: Double) * rw
                    i y' = cy - (rh / 2.0) + ((fromIntegral y' + 0.5) / fromIntegral (h * s) :: Double) * rh
                    rs = fmap (\(x', y') -> (r x', i y')) ps
                    qs = fmap (\(r0, i0) -> color (fastRender (realToFrac r0 :+ realToFrac i0) mi)) rs
                    (rc, gc, bc) = foldr (\(PixelRGB8 r g b) (rc :: Int, gc :: Int, bc :: Int) -> (rc + fromIntegral r, gc + fromIntegral g, bc + fromIntegral b)) (0, 0, 0) qs
                 in PixelRGB8 (fromIntegral (rc `div` (s*s))) (fromIntegral (gc `div` (s*s))) (fromIntegral (bc `div` (s*s)))
    in f

mag2 :: (RealFloat a, Num a) => Complex a -> a
mag2 (r :+ i) = r*r + i*i

render :: (RealFloat a, Num a) => Complex a -> Int -> Maybe (Complex a, Int)
render c maxI = loop (0 :+ 0) 0
    where loop z i = if mag2 z <= 4 && i <= maxI
                        then loop (step z) (i + 1)
                        else if i < maxI
                            then Just (step $ step z, i)
                            else Nothing
          step z = z*z + c

fastRender :: (RealFloat a, Num a) => Complex a -> Int -> Maybe (Complex a, Int)
fastRender c maxI = fastLoop (0 :+ 0) 0
    where fastLoop z i = let zn = (step . step . step . step . step . step . step . step) z
                         in if mag2 zn <= 4 && (i + 8) <= maxI
                            then fastLoop zn (i + 8)
                            else slowLoop z i
          slowLoop z i = if mag2 z <= 4 && i <= maxI
                         then slowLoop (step z) (i + 1)
                         else if i < maxI
                              then Just (step $ step z, i)
                              else Nothing
          step z = z*z + c

color :: (RealFloat a, Num a) => Maybe (Complex a, Int) -> PixelRGB8
color Nothing = PixelRGB8 0 0 0
color (Just (z, i)) = let mu = ((fromIntegral (i + 1)) - (logBase 2 . log . mag2) z) / 80.0
                          frac = snd (properFraction mu)
                      in getColor $ if frac < 0 then frac + 1.0 else frac
{-
renderPoint :: Double     -- ^ The real coordinate
            -> Double     -- ^ The imaginary coordinate
            -> Int        -- ^ The max iterations
            -> PixelRGB8  -- ^ The resulting pixel
renderPoint r0 i0 mi = case loop (0, 0) 0 of
                       Left (p, iter) -> let (r2, i2) = step . step $ p
                                             mu = ((fromIntegral (iter + 1)) - (logBase 2 . log) (r2 * r2 + i2 * i2)) / 80.0
                                             frac = snd (properFraction mu)
                                         in getColor $ if frac < 0 then frac + 1.0 else frac
                       Right _ -> PixelRGB8 0 0 0
    where loop (r, i) iter = if r*r + i*i <= 4 && iter <= mi
                          then loop (step (r, i)) (iter + 1)
                          else if iter < mi
                               then Left ((r, i), iter)
                               else Right ()
          step (r, i) = (r*r - i*i + r0, 2*r*i + i0)
-}

palette :: (RealFloat a, Num a) => [(a, PixelRGB8)]
palette = [ (0.0,    PixelRGB8 0   7   100)
          , (0.16,   PixelRGB8 32  107 203)
          , (0.42,   PixelRGB8 237 255 255)
          , (0.6425, PixelRGB8 255 170 0)
          , (0.8575, PixelRGB8 0   2   0)
          , (1.0,    PixelRGB8 0   7   100)
          ]

getColor :: (RealFloat a, Num a) => a -> PixelRGB8
getColor f = let ((li, (PixelRGB8 lr lg lb)), (ri, (PixelRGB8 rr rg rb))) = pair palette
                 index = (f - li) / (ri - li)
                 index2 = index * index
                 index3 = index * index2
                 r = fromIntegral . floor $ (2.0 * (fromIntegral lr - fromIntegral rr) * index3) - (3.0 * (fromIntegral lr - fromIntegral rr) * index2) + fromIntegral lr
                 g = fromIntegral . floor $ (2.0 * (fromIntegral lg - fromIntegral rg) * index3) - (3.0 * (fromIntegral lg - fromIntegral rg) * index2) + fromIntegral lg
                 b = fromIntegral . floor $ (2.0 * (fromIntegral lb - fromIntegral rb) * index3) - (3.0 * (fromIntegral lb - fromIntegral rb) * index2) + fromIntegral lb
             in PixelRGB8 r g b
    where pair (c:cs) = if f >= fst (head cs) then pair cs else (c, head cs)


module Main where

import Mandelbrot
import Data.Complex (Complex((:+)))
import Graphics.UI.SDL.Video
import Graphics.UI.SDL.Event (quitRequested)
import Graphics.UI.SDL.Types (Surface(surfacePixels))
import Foreign.C.String (newCString)
import Foreign.C.Types (CInt(..))
import Foreign.Ptr
import Foreign.Storable
import Data.Int (Int32(..), Int64(..))
import Data.Word
import Control.Monad
import Control.Concurrent (threadDelay)
import Data.Array
import Data.Array.MArray
import Data.Array.IO
import Data.IORef
import System.Random
import System.IO.Unsafe

sdlWindowPosCentered :: CInt
sdlWindowPosCentered = 0x2fff0000

iterations :: Int
iterations = 100

width :: Int
width = 1024

height :: Int
height = 786

widthF :: Double
widthF = fromIntegral width

heightF :: Double
heightF = fromIntegral height

mandelbrotI :: RealFloat a => Complex a -> Int
mandelbrotI = mandelbrot iterations
{-# INLINE mandelbrotI #-}

while :: Monad m => m Bool -> m a -> m ()
while test action = do
    val <- test
    if val
    then action >> while test action
    else return ()

main :: IO ()
main = do
    videoInit nullPtr
    title <- (newCString "Handelbrot")

    palette <- array (1, iterations) <$> sequence ( [sequence (i, (fromIntegral <$> randomRIO (0 :: Word, 0xffffff)) :: IO Word32) | i <- [1..iterations-1]] ++ [return (iterations, 0)] )

    window <- createWindow title sdlWindowPosCentered sdlWindowPosCentered (CInt $ fromIntegral width) (CInt $ fromIntegral height) 0
    surface <- getWindowSurface window
    lockSurface surface
    pixels <- liftM (castPtr.surfacePixels) $ peek surface :: IO (Ptr Word32)

    forM_ [0..height-1] (\y_ -> do
        forM_ [0..width-1] (\x_ -> do
            let x = 3/widthF * (fromIntegral x_) - 2
                y = -2/heightF * (fromIntegral y_) + 1
                mand = mandelbrotI (x :+ y)
            pokeElemOff pixels (x_ + y_*width) $ palette ! mand)
        updateWindowSurface window)

    unlockSurface surface
    updateWindowSurface window

    while (liftM not $ quitRequested) $ threadDelay 100000


{-# OPTIONS_GHC -fno-warn-missing-signatures #-}

module HVX.BlackBoxTests.HuberTestSpec where

import Test.Hspec
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Util (zeros)

import HVX
import HVX.Internal.TestUtil

-- Problem definition.
n = 4
a = EConst $ (n><n)
    ([0.5377  , 0.3188  , 3.5784  , 0.7254
    , 1.8339  , -1.3077 , 2.7694  , -0.0631
    , -2.2588 , -0.4336 , -1.3499 , 0.7147
    , 0.8622  , 0.3426  , 3.0349  , -0.2050 ] :: [Double])
b = EConst $
    (n><1) ([ -1.0689 , -0.8095 , -2.9443 , -1.4384 ] :: [Double])
c = EConst $ (n><n)
    ([ -0.1241 , 0.6715  , 0.4889  , 0.2939
    , 1.4897  , -1.2075 , 1.0347  , -0.7873
    , 1.4090  , 0.7172  , 0.7269  , 0.8884
    , 1.4172  , 1.6302  , -0.3034 , -1.1471 ] :: [Double])
d = EConst $ (n><1)
    ([ 0.3252  , -0.7549 , 1.3703  , -1.7115 ] :: [Double])

x = EVar "x"
y = EVar "y"
constZeroVector = EConst $ zeros n 1

subgradAns = subgradMinimize
  (hmax $ huber 3 (a *~ x +~ b) +~ berhu 1 (c *~ y +~ d))
  [x +~ y <=~ constZeroVector]
  (decNonSumStep 1.0) 100000
  [("x", zeros n 1), ("y", zeros n 1)]

ellipsoidAns = ellipsoidMinimize
  (hmax $ huber 3 (a *~ x +~ b) +~ berhu 1 (c *~ y +~ d))
  [x +~ y <=~ constZeroVector]
  [("x", n), ("y", n)]
  1e-16 1e10

(subgradVars, subgradOptval) = subgradAns

(ellipsoidVars, ellipsoidOptval, ellipsoidUBound) = ellipsoidAns

-- CVX's results.
cvxOptval = 0.404882
--cvxOptx = (n><1) [-1.6137 ,-1.1202 ,0.7986 ,-0.7487 ]
--cvxOpty = (n><1) [0.2428 ,-0.0191 ,-0.7986 ,-1.0079 ]

-- Verify that HVX matches CVX.
spec :: Spec
spec =
  describe "Verify that HVX matches CVX for huber/berhu" $ do
    it "Verify that HVX subgrad matches CVX for huber/berhu" $
      subgradOptval `shouldSatisfy` fpequalsApprox cvxOptval
    it "Verify that HVX ellipsoid matches CVX for huber/berhu" $
      ellipsoidOptval `shouldSatisfy` fpequalsApprox cvxOptval

main :: IO ()
main = hspec spec


{-# LANGUAGE BangPatterns #-}
{- | Parser for MNIST data.
Format: n lines of 785 double precision numbers, first 784 are the image, last is the expected number.
-}
module Main where

import           Data.Attoparsec.ByteString.Char8 (Parser, char, double, parseOnly, sepBy)
import qualified Data.ByteString.Char8 as BS
import           Numeric.LinearAlgebra (R, Vector, assoc, vector)
import qualified Data.Vector as V

data Sample = Sample
              { sampleInput          :: {-# UNPACK #-} !(Vector R)
              , sampleExpectedOutput :: {-# UNPACK #-} !(Vector R)
              } deriving (Read, Show)

makeSample :: ([R], [R]) -> Sample
makeSample (image, [expected]) = Sample (vector image) (assoc 10 0 [(floor expected, 1)])

sample :: Parser Sample
sample = do
  x <- double `sepBy` (char ' ')
  let !s = makeSample $ splitAt 784 x
  pure s

toSampleSet :: V.Vector Sample -> (V.Vector Sample, V.Vector Sample, V.Vector Sample)
toSampleSet vec = (train, validation, test)
    where (train, rest) = V.splitAt 50000 vec
          (validation, test) = V.splitAt 10000 rest

parseMnist :: BS.ByteString -> (V.Vector Sample, V.Vector Sample, V.Vector Sample)
parseMnist = toSampleSet . V.unfoldr parseLines . BS.lines
    where parseLines :: [BS.ByteString] -> Maybe (Sample, [BS.ByteString])
          parseLines []     = Nothing
          parseLines (l:ls) = Just (fromRight $ parseOnly sample l, ls)

          fromRight (Right b) = b

main :: IO ()
main = do
    !contents <- BS.readFile "/home/matt/Downloads/mnist.dat"
    let (!train, !validation, !test) = parseMnist contents
    print $
        show (length train)      ++ " " ++
        show (length validation) ++ " " ++
        show (length test)


{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Numeric.LinearAlgebra
import Control.Monad
import Control.Monad.Random
import Control.Applicative (liftA2)

data Weights = W { wBiases :: !(Vector Double) -- n
                 , wNodes :: !(Matrix Double)  -- n x m
                 }                             -- "m to n" layer

data Network :: * where
  O :: !Weights -> Network
  (:~) :: !Weights -> !Network -> Network

infixr 5 :~

-- | Generate random weights.

randomWeights :: MonadRandom m => Int -> Int -> m Weights
randomWeights i o = do
  seed1  :: Int <- getRandom
  seed2  :: Int <- getRandom
  let wB = randomVector seed1 Uniform o
      wN = uniformSample seed2 o (replicate i (-1, 1))
  return $ W wB wN

-- | Generate a random Network
randomNet :: MonadRandom m => Int -> [Int] -> Int -> m Network
randomNet i [] o = O <$> randomWeights i o
randomNet i (h:hs) o = liftA2 (:~) (randomWeights i h) (randomNet i hs o)

-- | Relu function
relu :: (Floating a, Ord a) => a -> a
relu = max 0

relu' :: (Floating a, Ord a) => a -> a
relu' x | x < 0 = 0
        | otherwise = 1

train :: Double -> Vector Double -> Vector Double -> Network -> Network
train = undefined


module STCR2Z2T0S0Amodal 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
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoDecayStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:scale0FreqsStr:scaleFreqsStr:histFilePath:numIterationStr:writeSourceFlagStr:sigmaStr:initDistStr:cutoffRadiusStr:reversalFactorStr: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
      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
      sigma = read sigmaStr :: Double
      initDist = read initDistStr :: [R2S1RPPoint]
      cutoffRadius = read cutoffRadiusStr :: Int
      reversalFactor = read reversalFactorStr :: Double
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2Z2T0S0Amodal"
  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..."
        solveMonteCarloR2Z2T0S0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          taoDecay
          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
      (cutoff cutoffRadius radialArr)
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      theta0Freqs
      scale0Freqs
  plan <- makeR2Z2T0S0Plan emptyPlan arrR2Z2T0S0
  let xs = initDist
      bias =
        computeBiasR2T0S0Gaussian
          numPoint
          numPoint
          theta0Freqs
          scale0Freqs
          0
          sigma
          xs
      eigenVec =
        computeInitialEigenVectorR2T0S0
          numPoint
          numPoint
          theta0Freqs
          scale0Freqs
          thetaFreqs
          scaleFreqs
          xs
  powerMethodR2Z2T0S0BiasReversal
    plan
    folderPath
    numPoint
    numPoint
    numOrientation
    thetaFreqs
    theta0Freqs
    numScale
    scaleFreqs
    scale0Freqs
    arrR2Z2T0S0
    numIteration
    writeSourceFlag
    (printf
       "_%d_%d_%d_%.2f_%.2f_%.3f"
       numPoint
       (round maxScale :: Int)
       (round taoDecay :: Int)
       thetaSigma
       scaleSigma
       reversalFactor)
    reversalFactor
    bias
    eigenVec


module STCR2Z2T0S0Image 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
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:scale0FreqsStr:scaleFreqsStr:histFilePath:numIterationStr:writeSourceFlagStr:imagePath:cutoffRStr:reversalFactorStr: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
      tao = read taoStr :: 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
      cutoffR = read cutoffRStr :: Int
      reversalFactor = read reversalFactorStr :: Double
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2Z2T0S0Image"
  imgRepa@(ImageRepa _ img') <- readImageRepa imagePath False
  let (Z :. _ :. cols :. rows) = extent img'
      img =
        computeUnboxedS . R.traverse img' id $ \f idx@(Z :. k :. i :. j) ->
          if (sqrt . fromIntegral $ (i - div cols 2) ^ 2 + (j - div rows 2) ^ 2) >
             48
            then 0
            else f idx
  createDirectoryIfMissing True folderPath
  plotImageRepa (folderPath </> "input.png") . ImageRepa 8 $ img
  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..."
        solveMonteCarloR2Z2T0S0Radial
          numThread
          numTrail
          maxTrail
          cols
          rows
          thetaSigma
          scaleSigma
          maxScale
          tao
          theta0Freqs
          thetaFreqs
          scale0Freqs
          scaleFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div (max rows cols) 2) ^ 2)
             , L.length scale0Freqs
             , L.length theta0Freqs
             , L.length scaleFreqs
             , L.length thetaFreqs
             ]
             0)
  arrR2Z2T0S0 <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      (cutoff cutoffR radialArr)
      cols
      rows
      1
      maxScale
      thetaFreqs
      scaleFreqs
      theta0Freqs
      scale0Freqs
  plan <- makeR2Z2T0S0Plan emptyPlan arrR2Z2T0S0
  powerMethodR2Z2T0S0Reversal
    plan
    folderPath
    cols
    rows
    numOrientation
    thetaFreqs
    theta0Freqs
    numScale
    scaleFreqs
    scale0Freqs
    arrR2Z2T0S0
    numIteration
    writeSourceFlag
    (printf
       "_%d_%d_%.2f_%.2f"
       (round maxScale :: Int)
       (round tao :: Int)
       thetaSigma
       scaleSigma)
    0.5
    reversalFactor
    (R.traverse
       img
       (const
          (Z :. (L.length theta0Freqs) :. (L.length scale0Freqs) :. cols :. rows)) $ \f (Z :. _ :. _ :. i :. j) ->
       if f (Z :. (0 :: Int) :. i :. j) > 0
         then 1
         else 0)
    (R.traverse
       img
       (const
          (Z :. (L.length thetaFreqs) :. (L.length scaleFreqs) :.
           (L.length theta0Freqs) :.
           (L.length scale0Freqs) :.
           cols :.
           rows)) $ \f (Z :. k :. l :. _ :. _ :. i :. j) ->
       if k == div (L.length thetaFreqs) 2 && l == div (L.length scaleFreqs) 2
         then 1 / (fromIntegral $ (L.length theta0Freqs * L.length scale0Freqs)) :+
              0
         else 0)


{-# LANGUAGE ConstraintKinds, DataKinds, FlexibleContexts        #-}
{-# LANGUAGE MultiParamTypeClasses, NoImplicitPrelude, PolyKinds #-}
{-# LANGUAGE QuasiQuotes, TemplateHaskell                        #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}
module Main (module Algebra.Algorithms.Groebner, module Algebra.Prelude
            , module Main
            ) where
import           Algebra.Algorithms.Groebner
import           Algebra.Algorithms.ZeroDim
import           Algebra.Ring.Ideal
import           Algebra.Ring.Polynomial.Quotient
import           Algebra.Prelude
import           Control.Monad.Random             hiding (fromList)
import           Data.Complex
import           Data.Convertible
import           Data.List                        (find, nub, partition, sortBy)
import qualified Data.Matrix                      as M
import qualified Data.Sized               as SV
import qualified Data.Vector                      as V
import           Debug.Trace
import qualified Numeric.LinearAlgebra            as LA
import qualified Prelude                          as P

tr :: Show a => a -> a
tr a = trace (show a) a

x, y, z :: Polynomial (Fraction Integer) 3
[x, y, z] = vars

seed :: Polynomial (Fraction Integer) 3
seed = -412742019532366985 * x -7641395389638504101 * y + 4362835172800530323 * z

seedMat :: LA.Matrix Double
seedMat = LA.fromLists $ map (map toDouble) $ reifyQuotient eqn02 $ \pxy -> matrixRep (modIdeal' pxy seed)

toDouble :: Fractional a => Fraction Integer -> a
toDouble rat = fromIntegral (numerator rat) P./ fromIntegral (denominator rat)

fromRight :: Either t t1 -> t1
fromRight (Right a) = a
fromRight _ = error "fromRight"

printLvl :: Show a => Int -> a -> IO ()
printLvl lvl = putStrLn . unlines . map (replicate lvl '\t' ++) . lines . show

eqn01 :: Ideal (Polynomial (Fraction Integer) 3)
eqn01 = toIdeal [x^2 - 2*x*z + 5, x*y^2+y*z+1, 3*y^2 - 8*x*z]

eqn02 :: Ideal (Polynomial (Fraction Integer) 3)
eqn02 =
  toIdeal [x^2 + 2*y^2 - y - 2*z
          ,x^2 - 8*y^2 + 10*z - 1
          ,x^2 - 7*y*z
          ]

eqn03 :: Ideal (Polynomial (Fraction Integer) 3)
eqn03 = toIdeal [x^2 + y^2 + z^2 - 2*x
                ,x^3 - y*z - x
                ,x - y + 2*z
                ]

jdeal :: Ideal (Polynomial (Fraction Integer) 3)
jdeal = toIdeal [x*y + z - x*z, x^2 - z, 2*x^3 - x^2 * y * z - 1]

vs :: [V.Vector (Fraction Integer)]
vs = reifyQuotient eqn03 $ \pxy -> map (vectorRep . modIdeal' pxy) [var 0 ^ i | i <- [0..6::Natural]]

mat :: M.Matrix (Fraction Integer)
mat = fromCols $ take 4 vs

fromCols :: [V.Vector a] -> M.Matrix a
fromCols = foldr1 (M.<|>) . map M.colVector

findUnivar :: (CoeffRing r, IsMonomialOrder n ord, KnownNat n)
           => OrderedPolynomial r ord n -> Maybe (Ordinal n)
findUnivar poly =
  let os = enumOrdinal (sArity' poly)
      ms = map snd $ getTerms poly
  in find (\a -> all (`isPowerOf` (leadingMonomial (var a `asTypeOf` poly))) ms) os

toCoeffList :: (CoeffRing r,  KnownNat n, IsMonomialOrder n ord) => Ordinal n -> OrderedPolynomial r ord n -> [r]
toCoeffList on f =
  let v = var on  `asTypeOf` f
  in [ coeff (leadingMonomial $ v ^ i) f | i <- [0.. fromIntegral (totalDegree' f)]]

showSols :: (KnownNat n, IsMonomialOrder n order, Convertible a Double)
         => Double -> Ideal (OrderedPolynomial a order n) -> [Sized n1 (Complex Double)] -> IO ()
showSols err eqn sols = do
  let (rs, is) = partition (all ((<err).P.abs.imagPart)) $ map SV.toList sols
      subs a b c = generators $
                   mapIdeal (magnitude . substWith (*) (SV.unsafeFromList' [a, b, c]) . mapCoeff toComplex)
                            eqn
      showCase [a,b,c] = print (a, b, c) >> putStr "\terror: ">> print (maximum $ subs a b c)
  putStrLn $ "- " ++ show (length rs) ++ " real solution(s):"
  mapM_ showCase $ sortBy (comparing $ map magnitude) rs
  putStrLn $ "- " ++ show (length is) ++ " imaginary solution(s):"
  mapM_  showCase $ sortBy (comparing $ map magnitude) is
  let errs = concatMap (\ [a,b,c] -> subs a b c) $ rs ++ is
  putStrLn $ "- maximum error: " ++ show (maximum errs)
  putStrLn $ "- minimum error: " ++ show (minimum errs)
  putStrLn $ "- average error: " ++ show (sum errs P./ fromIntegral (length errs))

main :: IO ()
main = do
  putStrLn "---- solving equation system"
  let err = 1e-10
  putStrLn "< naive method"
  showSols err eqn01 $ solve' err eqn01
  putStrLn "\n< randomized method"
  showSols err eqn01 =<< evalRandIO (solveM eqn01)
  putStrLn "\n< companion characteristics"
  showSols err eqn01 $ solveViaCompanion err eqn01
  putStrLn "\n< univariate spanning"
  showSols err eqn01 $ solve' err eqn01

  putStrLn "\n\n---- exercise 8"
  putStrLn "< Solving 1-6"
  putStrLn "< Naive Method: "
  showSols err eqn02 $ nub $ solve' err eqn02
  putStrLn "\n< new method"
  showSols err eqn02 =<< evalRandIO (solveM eqn02)

  putStrLn "\n< Solving 1-7"
  putStrLn "< Naive Method: "
  showSols err eqn03 $ nub $ solve' err eqn03
  putStrLn "\n< new method"
  showSols err eqn03 =<< evalRandIO (solveM eqn03)
  putStrLn "\n\n---- FGLM Algorithm"
  print $ fglm jdeal
  print $ calcGroebnerBasisWith Lex jdeal
  print $ univPoly 0 jdeal
  print $ univPoly 1 jdeal
  print $ univPoly 2 jdeal
  return ()

substIdeal :: IsMonomialOrder 3 order
           => [OrderedPolynomial (Fraction Integer) Grevlex 3]
           -> Ideal (OrderedPolynomial (Fraction Integer) order 3)
           -> Ideal (OrderedPolynomial (Fraction Integer) Grevlex 3)
substIdeal = mapIdeal . substWith (.*.) . SV.unsafeFromList'

toComplex :: Convertible r Double => r -> Complex Double
toComplex = (:+ 0) . convert


module Main where

-- import           Lib

import qualified Control.Exception.Safe as Exception
import Data.Complex (Complex (..))
import qualified Data.Complex as Complex
import Data.Foldable (foldlM)
import Data.Functor (($>))
import qualified Data.Generics.Sum as GSum
import qualified Data.List as List
import Data.Map.Strict ((!?))
import qualified Data.Map.Strict as Map
import qualified Data.Maybe as Maybe
import Data.Ratio ((%))
import qualified Data.Ratio as Ratio
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Text.IO as Text
import Data.Void (Void)
import GHC.Generics (Generic)
import qualified Optics as Op
import System.Environment (getArgs)

-- import Text.Megaparsec
-- import qualified Text.Megaparsec.Char as MChar
-- import qualified Text.Megaparsec.Char.Lexer as L

data LispVal
  = Atom Text
  | List [LispVal]
  | DottedList [LispVal] LispVal
  | Number LispNumber
  | Float Float
  | String Text
  | Bool Bool
  deriving (Generic, Eq, Show)

data LispNumber
  = Complex LispComplex
  | Real Double
  | Rational Rational
  | Integer Integer
  deriving (Generic, Eq, Show)

type LispComplex = Complex Double

-- lispNumComplex = Op.iso toComplex toLispNumber

lispNumToComplex :: LispNumber -> LispComplex
lispNumToComplex (Complex x) = x
lispNumToComplex (Real x) = x :+ 0
lispNumToComplex (Rational x) = fromRational x :+ 0
lispNumToComplex (Integer x) = fromInteger x :+ 0

lispNumToDouble :: Exception.MonadThrow m => LispNumber -> m Double
lispNumToDouble x@(Complex _) = Exception.throw [TypeMismatch "Real" $ Number x]
lispNumToDouble (Real x) = pure x
lispNumToDouble (Rational x) = pure $ fromRational x
lispNumToDouble (Integer x) = pure $ fromInteger x

lispNumToRational :: Exception.MonadThrow m => LispNumber -> m Rational
lispNumToRational x@(Complex _) = Exception.throw [TypeMismatch "Rational" $ Number x]
lispNumToRational x@(Real _) = Exception.throw [TypeMismatch "Rational" $ Number x]
lispNumToRational (Rational x) = pure x
lispNumToRational (Integer x) = pure $ fromInteger x

lispNumToInteger :: Exception.MonadThrow m => LispNumber -> m Integer
lispNumToInteger x@(Complex _) = Exception.throw [TypeMismatch "Integer" $ Number x]
lispNumToInteger x@(Real _) = Exception.throw [TypeMismatch "Integer" $ Number x]
lispNumToInteger x@(Rational _) = Exception.throw [TypeMismatch "Integer" $ Number x]
lispNumToInteger (Integer x) = pure x

-- toLispNumber :: LispComplex -> LispNumber
-- toLispNumber (r :+ 0)
--   | Ratio.denominator r == 0 = Integer $ Ratio.numerator r
--   | otherwise = Rational r
-- toLispNumber c = Complex c

-- apNumOpToLispNumber ::
--   (Complex Rational -> Complex Rational -> Complex Rational) ->
--   (LispNumber -> LispNumber -> LispNumber)
-- apNumOpToLispNumber op x y = toLispNumber $ toComplex x `op` toComplex y

-- instance Num LispNumber where
--   (+) = apNumOpToLispNumber (+)
--   (-) = apNumOpToLispNumber (-)
--   (*) = apNumOpToLispNumber (*)
--   abs = Op.over lispNumComplex abs
--   signum = Op.over lispNumComplex signum
--   fromInteger = toLispNumber . fromInteger

-- instance Fractional LispNumber where
--   fromRational = toLispNumber . fromRational
--   recip = Op.over lispNumComplex recip

data LispError
  = NumArgs Integer [LispVal]
  | TypeMismatch Text LispVal
  | --   | Parser ParseError
    BadSpecialForm Text LispVal
  | -- | NotFunction T.Text T.Text
    UnboundVar Text
  | Default Text
  deriving (Generic, Eq, Show)

instance Exception.Exception [LispError]

type LispThrowErr = Either [LispError]

-- type Parser = Parsec Void T.Text

-- symbolChar :: Parser Char
-- symbolChar = oneOf "!#$%&|*+-/:<=>?@^_~"

-- spaceConsumer :: Parser ()
-- spaceConsumer = L.space MChar.space1 empty empty

-- lexeme :: Parser a -> Parser a
-- lexeme = L.lexeme spaceConsumer

-- symbol :: T.Text -> Parser T.Text
-- symbol = L.symbol' spaceConsumer

-- -- readExpr :: T.Text -> T.Text
-- -- readExpr input =
-- --   case parse (space >> symbolChar) "lisp" input of
-- --     Left err -> "No match: " <> T.pack (errorBundlePretty err)
-- --     Right val -> "Found value: " <> val

-- parseExpr :: Parser LispVal
-- parseExpr = lexeme $ Atom . T.singleton <$> symbolChar

-- parseString :: Parser LispVal
-- parseString = do
--   x <- between (symbol "\"") (symbol "\"") $ takeWhileP Nothing (/= '\'')
--   return $ String x

-- parseAtom :: Parser LispVal
-- parseAtom = do
--   first <- MChar.letterChar <|> symbolChar
--   rest <- many (letter <|> digit <|> symbol)
--   let atom = first : rest
--   return $ case atom of
--     "#t" -> Bool True
--     "#f" -> Bool False
--     _ -> Atom atom

-- parseNumber :: Parser LispVal
-- parseNumber = Number . read <$> many1 digit

foldNumOp :: Exception.MonadThrow m => (LispNumber -> LispNumber -> m LispNumber) -> LispNumber -> [LispVal] -> m LispVal
foldNumOp op identity ls =
  let mNumbers = Op.preview (GSum._As @"Number") <$> ls
   in case List.elemIndices Nothing mNumbers of
        [] -> Number <$> foldlM op identity (Maybe.catMaybes mNumbers)
        indices -> Exception.throw $ TypeMismatch "Number" . (ls !!) <$> indices

binNumOp :: Exception.MonadThrow m => (LispNumber -> LispNumber -> m LispNumber) -> [LispVal] -> m LispVal
binNumOp op [Number n, Number m] = Number <$> op n m
binNumOp op args@[_, _] =
  Exception.throw $
    TypeMismatch "Number"
      <$> List.filter (\case Number _ -> True; _ -> False) args
binNumOp op args = Exception.throw [NumArgs 2 args]

numOpToLispNumOp ::
  Exception.MonadThrow m =>
  (forall a. Num a => a -> a -> a) ->
  (LispNumber -> LispNumber -> m LispNumber)
numOpToLispNumOp op (Integer x) (Integer y) = pure $ Integer $ op x y
numOpToLispNumOp op (Integer x) (Rational y) = pure $ Rational $ op (fromIntegral x) y
numOpToLispNumOp op (Integer x) (Real y) = pure $ Real $ op (fromIntegral x) y
numOpToLispNumOp op (Rational x) (Integer y) = pure $ Rational $ op x $ fromIntegral y
numOpToLispNumOp op (Rational x) (Rational y) = pure $ Rational $ op x y
numOpToLispNumOp op (Rational x) (Real y) = pure $ Real $ op (fromRational x) y
numOpToLispNumOp op (Real x) (Integer y) = pure $ Real $ op x $ fromIntegral y
numOpToLispNumOp op (Real x) (Rational y) = pure $ Real $ op x $ fromRational y
numOpToLispNumOp op (Real x) (Real y) = pure $ Real $ op x y
numOpToLispNumOp op x y = pure $ Complex $ op (lispNumToComplex x) (lispNumToComplex y)

fractionalOpToLispNumOp ::
  Exception.MonadThrow m =>
  (forall a. Fractional a => a -> a -> a) ->
  (LispNumber -> LispNumber -> m LispNumber)
fractionalOpToLispNumOp op (Integer x) (Integer y) = pure $ Rational $ op (fromIntegral x) (fromIntegral y)
fractionalOpToLispNumOp op (Integer x) (Rational y) = pure $ Rational $ op (fromIntegral x) y
fractionalOpToLispNumOp op (Integer x) (Real y) = pure $ Real $ op (fromIntegral x) y
fractionalOpToLispNumOp op (Rational x) (Integer y) = pure $ Rational $ op x $ fromIntegral y
fractionalOpToLispNumOp op (Rational x) (Rational y) = pure $ Rational $ op x y
fractionalOpToLispNumOp op (Rational x) (Real y) = pure $ Real $ op (fromRational x) y
fractionalOpToLispNumOp op (Real x) (Integer y) = pure $ Real $ op x $ fromIntegral y
fractionalOpToLispNumOp op (Real x) (Rational y) = pure $ Real $ op x $ fromRational y
fractionalOpToLispNumOp op (Real x) (Real y) = pure $ Real $ op x y
fractionalOpToLispNumOp op x y = pure $ Complex $ op (lispNumToComplex x) (lispNumToComplex y)

integerOpToLispNumOp ::
  Exception.MonadThrow m =>
  (forall a. Integral a => a -> a -> a) ->
  (LispNumber -> LispNumber -> m LispNumber)
integerOpToLispNumOp op (Integer x) (Integer y) = pure $ Integer $ op x y
integerOpToLispNumOp _ x y = Exception.throw $ TypeMismatch "Integer" . Number <$> [x, y]

primitivesEnv :: Exception.MonadThrow m => Map.Map Text ([LispVal] -> m LispVal)
primitivesEnv =
  Map.fromList
    [ ("+", foldNumOp (numOpToLispNumOp (+)) $ Integer 0),
      ("-", foldNumOp (numOpToLispNumOp (-)) $ Integer 0),
      ("*", foldNumOp (numOpToLispNumOp (*)) $ Integer 1),
      ("/", foldNumOp (fractionalOpToLispNumOp (/)) $ Integer 1),
      ("quotient", binNumOp $ integerOpToLispNumOp quot),
      ("remainder", binNumOp $ integerOpToLispNumOp rem),
      ("modulo", binNumOp $ integerOpToLispNumOp mod)
    ]

eval :: Exception.MonadThrow m => LispVal -> m LispVal
eval (List [Atom "quote", ls]) = pure ls
eval (List (Atom func : args)) = do
  args <- traverse eval args
  apply func args
eval val = pure val

apply :: forall m. Exception.MonadThrow m => Text -> [LispVal] -> m LispVal
apply funcName args =
  case primitivesEnv !? funcName :: Maybe ([LispVal] -> m LispVal) of
    Nothing -> Exception.throw [UnboundVar funcName]
    Just f -> f args

main :: IO ()
main = do
  print $
    (eval :: LispVal -> Either Exception.SomeException LispVal)
      <$> [ List [Atom "+", Number $ Integer 1, Number $ Integer 3],
            List [Atom "*", Number $ Integer 1, Number $ Real 3, Number $ Rational 1, Number $ Complex 3],
            List [Atom "/", Number $ Real 1, Number $ Real 3, Number $ Rational 2, Number $ Integer 3],
            List [Atom "remainder", Number $ Integer 12, Number $ Integer 5]
          ]


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Lognormal
-- Copyright : (c) 2020 Ximin Luo
-- License   : BSD3
--
-- Maintainer  : infinity0@pwned.gg
-- Stability   : experimental
-- Portability : portable
--
-- The log normal distribution.  This is a continuous probability
-- distribution that describes data whose log is clustered around a
-- mean. For example, the multiplicative product of many independent
-- positive random variables.

module Statistics.Distribution.Lognormal
    (
      LognormalDistribution
      -- * Constructors
    , lognormalDistr
    , lognormalDistrErr
    , lognormalDistrMeanStddevErr
    , lognormalStandard
    ) where

import Data.Aeson            (FromJSON, ToJSON)
import Data.Binary           (Binary (..))
import Data.Data             (Data, Typeable)
import GHC.Generics          (Generic)
import Numeric.MathFunctions.Constants (m_huge, m_sqrt_2_pi)
import Numeric.SpecFunctions (expm1, log1p)
import qualified Data.Vector.Generic as G

import qualified Statistics.Distribution as D
import qualified Statistics.Distribution.Normal as N
import Statistics.Internal

-- | The lognormal distribution.
newtype LognormalDistribution = LND N.NormalDistribution
    deriving (Eq, Typeable, Data, Generic)

instance Show LognormalDistribution where
  showsPrec i (LND d) = defaultShow2 "lognormalDistr" m s i
   where
    m = D.mean d
    s = D.stdDev d
instance Read LognormalDistribution where
  readPrec = defaultReadPrecM2 "lognormalDistr" $
    (either (const Nothing) Just .) . lognormalDistrErr

instance ToJSON LognormalDistribution
instance FromJSON LognormalDistribution

instance Binary LognormalDistribution where
  put (LND d) = put m >> put s
   where
    m = D.mean d
    s = D.stdDev d
  get = do
    m  <- get
    sd <- get
    either fail return $ lognormalDistrErr m sd

instance D.Distribution LognormalDistribution where
  cumulative      = cumulative
  complCumulative = complCumulative

instance D.ContDistr LognormalDistribution where
  logDensity    = logDensity
  quantile      = quantile
  complQuantile = complQuantile

instance D.MaybeMean LognormalDistribution where
  maybeMean = Just . D.mean

instance D.Mean LognormalDistribution where
  mean (LND d) = exp (m + v / 2)
   where
    m = D.mean d
    v = D.variance d

instance D.MaybeVariance LognormalDistribution where
  maybeStdDev   = Just . D.stdDev
  maybeVariance = Just . D.variance

instance D.Variance LognormalDistribution where
  variance (LND d) = expm1 v * exp (2 * m + v)
   where
    m = D.mean d
    v = D.variance d

instance D.Entropy LognormalDistribution where
  entropy (LND d) = logBase 2 (s * exp (m + 0.5) * m_sqrt_2_pi)
   where
    m = D.mean d
    s = D.stdDev d

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

instance D.ContGen LognormalDistribution where
  genContVar d = D.genContinuous d

-- | Standard log normal distribution with mu 0 and sigma 1.
--
-- Mean is @sqrt e@ and variance is @(e - 1) * e@.
lognormalStandard :: LognormalDistribution
lognormalStandard = LND N.standard

-- | Create log normal distribution from parameters.
lognormalDistr
  :: Double            -- ^ Mu
  -> Double            -- ^ Sigma
  -> LognormalDistribution
lognormalDistr mu sig = either error id $ lognormalDistrErr mu sig

-- | Create log normal distribution from parameters.
lognormalDistrErr
  :: Double            -- ^ Mu
  -> Double            -- ^ Sigma
  -> Either String LognormalDistribution
lognormalDistrErr mu sig
  | sig >= sqrt (log m_huge - 2 * mu) = Left $ errMsg mu sig
  | otherwise = LND <$> N.normalDistrErr mu sig

errMsg :: Double -> Double -> String
errMsg mu sig =
  "Statistics.Distribution.Lognormal.lognormalDistr: sigma must be > 0 && < "
    ++ show lim ++ ". Got " ++ show sig
  where lim = sqrt (log m_huge - 2 * mu)

-- | Create log normal distribution from mean and standard deviation.
lognormalDistrMeanStddevErr
  :: Double            -- ^ Mu
  -> Double            -- ^ Sigma
  -> Either String LognormalDistribution
lognormalDistrMeanStddevErr m sd = LND <$> N.normalDistrErr mu sig
  where r = sd / m
        sig2 = log1p (r * r)
        sig = sqrt sig2
        mu = log m - sig2 / 2

-- | Variance is estimated using maximum likelihood method
--   (biased estimation) over the log of the data.
--
--   Returns @Nothing@ if sample contains less than one element or
--   variance is zero (all elements are equal)
instance D.FromSample LognormalDistribution Double where
  fromSample = fmap LND . D.fromSample . G.map log

logDensity :: LognormalDistribution -> Double -> Double
logDensity (LND d) x
  | x > 0 = let lx = log x in D.logDensity d lx - lx
  | otherwise = 0

cumulative :: LognormalDistribution -> Double -> Double
cumulative (LND d) x
  | x > 0 = D.cumulative d $ log x
  | otherwise = 0

complCumulative :: LognormalDistribution -> Double -> Double
complCumulative (LND d) x
  | x > 0 = D.complCumulative d $ log x
  | otherwise = 1

quantile :: LognormalDistribution -> Double -> Double
quantile (LND d) = exp . D.quantile d

complQuantile :: LognormalDistribution -> Double -> Double
complQuantile (LND d) = exp . D.complQuantile d