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


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


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


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


{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies  #-}
{-# LANGUAGE TypeOperators #-}

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

{- |
Module      :  Internal.Modular
Copyright   :  (c) Alberto Ruiz 2015
License     :  BSD3
Stability   :  experimental

Proof of concept of statically checked modular arithmetic.

-}

module Internal.Modular(
    Mod, type (./.)
) where

import Internal.Vector
import Internal.Matrix hiding (size)
import Internal.Numeric
import Internal.Element
import Internal.Container
import Internal.Vectorized (prodI,sumI,prodL,sumL)
import Internal.LAPACK (multiplyI, multiplyL)
import Internal.Algorithms(luFact,LU(..))
import Internal.Util(Normed(..),Indexable(..),
                     gaussElim, gaussElim_1, gaussElim_2,
                     luST, luSolve', luPacked', magnit, invershur)
import Internal.ST(mutable)
#if MIN_VERSION_base(4,11,0)
import GHC.TypeLits hiding (Mod)
#else
import GHC.TypeLits
#endif
import Data.Proxy(Proxy)
import Foreign.ForeignPtr(castForeignPtr)
import Foreign.Storable
import Data.Ratio
import Data.Complex
import Control.DeepSeq ( NFData(..) )
#if MIN_VERSION_base(4,11,0)
import Prelude hiding ((<>))
#endif

-- | Wrapper with a phantom integer for statically checked modular arithmetic.
newtype Mod (n :: Nat) t = Mod {unMod:: t}
  deriving (Storable)

instance (NFData t) => NFData (Mod n t)
  where
    rnf (Mod x) = rnf x

infixr 5 ./.
type (./.) x n = Mod n x

instance (Integral t, Enum t, KnownNat m) => Enum (Mod m t)
  where
    toEnum = l0 (\m x -> fromIntegral $ x `mod` (fromIntegral m))
    fromEnum = fromIntegral . unMod

instance (Eq t, KnownNat m) => Eq (Mod m t)
  where
    a == b = (unMod a) == (unMod b)

instance (Ord t, KnownNat m) => Ord (Mod m t)
  where
    compare a b = compare (unMod a) (unMod b)

instance (Integral t, Real t, KnownNat m, Integral (Mod m t)) => Real (Mod m t)
  where
    toRational x = toInteger x % 1

instance (Integral t, KnownNat m, Num (Mod m t)) => Integral (Mod m t)
  where
    toInteger = toInteger . unMod
    quotRem a b = (Mod q, Mod r)
      where
         (q,r) = quotRem (unMod a) (unMod b)

-- | this instance is only valid for prime m
instance (Show (Mod m t), Num (Mod m t), Eq t, KnownNat m) => Fractional (Mod m t)
  where
    recip x
        | x*r == 1  = r
        | otherwise = error $ show x ++" does not have a multiplicative inverse mod "++show m'
      where
        r = x^(m'-2 :: Integer)
        m' = fromIntegral . natVal $ (undefined :: Proxy m)
    fromRational x = fromInteger (numerator x) / fromInteger (denominator x)

l2 :: forall m a b c. (Num c, KnownNat m) => (c -> a -> b -> c) -> Mod m a -> Mod m b -> Mod m c
l2 f (Mod u) (Mod v) = Mod (f m' u v)
  where
    m' = fromIntegral . natVal $ (undefined :: Proxy m)

l1 :: forall m a b . (Num b, KnownNat m) => (b -> a -> b) -> Mod m a -> Mod m b
l1 f (Mod u) = Mod (f m' u)
  where
    m' = fromIntegral . natVal $ (undefined :: Proxy m)

l0 :: forall m a b . (Num b, KnownNat m) => (b -> a -> b) -> a -> Mod m b
l0 f u = Mod (f m' u)
  where
    m' = fromIntegral . natVal $ (undefined :: Proxy m)

instance Show t => Show (Mod n t)
  where
    show = show . unMod

instance (Integral t, KnownNat n) => Num (Mod n t)
  where
    (+) = l2 (\m a b -> (a + b) `mod` (fromIntegral m))
    (*) = l2 (\m a b -> (a * b) `mod` (fromIntegral m))
    (-) = l2 (\m a b -> (a - b) `mod` (fromIntegral m))
    abs = l1 (const abs)
    signum = l1 (const signum)
    fromInteger = l0 (\m x -> fromInteger x `mod` (fromIntegral m))

instance KnownNat m => Element (Mod m I)
  where
    constantD x n = i2f (constantD (unMod x) n)
    extractR ord m mi is mj js = i2fM <$> extractR ord (f2iM m) mi is mj js
    setRect i j m x = setRect i j (f2iM m) (f2iM x)
    sortI = sortI . f2i
    sortV = i2f . sortV . f2i
    compareV u v = compareV (f2i u) (f2i v)
    selectV c l e g = i2f (selectV c (f2i l) (f2i e) (f2i g))
    remapM i j m = i2fM (remap i j (f2iM m))
    rowOp c a i1 i2 j1 j2 x = rowOpAux (c_rowOpMI m') c (unMod a) i1 i2 j1 j2 (f2iM x)
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)
    gemm u a b c = gemmg (c_gemmMI m') (f2i u) (f2iM a) (f2iM b) (f2iM c)
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)

instance KnownNat m => Element (Mod m Z)
  where
    constantD x n = i2f (constantD (unMod x) n)
    extractR ord m mi is mj js = i2fM <$> extractR ord (f2iM m) mi is mj js
    setRect i j m x = setRect i j (f2iM m) (f2iM x)
    sortI = sortI . f2i
    sortV = i2f . sortV . f2i
    compareV u v = compareV (f2i u) (f2i v)
    selectV c l e g = i2f (selectV c (f2i l) (f2i e) (f2i g))
    remapM i j m = i2fM (remap i j (f2iM m))
    rowOp c a i1 i2 j1 j2 x = rowOpAux (c_rowOpML m') c (unMod a) i1 i2 j1 j2 (f2iM x)
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)
    gemm u a b c = gemmg (c_gemmML m') (f2i u) (f2iM a) (f2iM b) (f2iM c)
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)

instance KnownNat m => CTrans (Mod m I)
instance KnownNat m => CTrans (Mod m Z)

instance KnownNat m => Container Vector (Mod m I)
  where
    conj' = id
    size' = dim
    scale' s x = vmod (scale (unMod s) (f2i x))
    addConstant c x = vmod (addConstant (unMod c) (f2i x))
    add' a b = vmod (add' (f2i a) (f2i b))
    sub a b = vmod (sub (f2i a) (f2i b))
    mul a b = vmod (mul (f2i a) (f2i b))
    equal u v = equal (f2i u) (f2i v)
    scalar' x = fromList [x]
    konst' x = i2f . konst (unMod x)
    build' n f = build n (fromIntegral . f)
    cmap' = mapVector
    atIndex' x k = fromIntegral (atIndex (f2i x) k)
    minIndex'     = minIndex . f2i
    maxIndex'     = maxIndex . f2i
    minElement'   = Mod . minElement . f2i
    maxElement'   = Mod . maxElement . f2i
    sumElements'  = fromIntegral . sumI m' . f2i
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)
    prodElements' = fromIntegral . prodI m' . f2i
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)
    step'         = i2f . step . f2i
    find' = findV
    assoc' = assocV
    accum' = accumV
    ccompare' a b = ccompare (f2i a) (f2i b)
    cselect' c l e g = i2f $ cselect c (f2i l) (f2i e) (f2i g)
    scaleRecip s x = scale' s (cmap recip x)
    divide x y = mul x (cmap recip y)
    arctan2' = undefined
    cmod' m = vmod . cmod' (unMod m) . f2i
    fromInt' = vmod
    toInt'   = f2i
    fromZ'   = vmod . fromZ'
    toZ'     = toZ' . f2i

instance KnownNat m => Container Vector (Mod m Z)
  where
    conj' = id
    size' = dim
    scale' s x = vmod (scale (unMod s) (f2i x))
    addConstant c x = vmod (addConstant (unMod c) (f2i x))
    add' a b = vmod (add' (f2i a) (f2i b))
    sub a b = vmod (sub (f2i a) (f2i b))
    mul a b = vmod (mul (f2i a) (f2i b))
    equal u v = equal (f2i u) (f2i v)
    scalar' x = fromList [x]
    konst' x = i2f . konst (unMod x)
    build' n f = build n (fromIntegral . f)
    cmap' = mapVector
    atIndex' x k = fromIntegral (atIndex (f2i x) k)
    minIndex'     = minIndex . f2i
    maxIndex'     = maxIndex . f2i
    minElement'   = Mod . minElement . f2i
    maxElement'   = Mod . maxElement . f2i
    sumElements'  = fromIntegral . sumL m' . f2i
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)
    prodElements' = fromIntegral . prodL m' . f2i
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)
    step'         = i2f . step . f2i
    find' = findV
    assoc' = assocV
    accum' = accumV
    ccompare' a b = ccompare (f2i a) (f2i b)
    cselect' c l e g = i2f $ cselect c (f2i l) (f2i e) (f2i g)
    scaleRecip s x = scale' s (cmap recip x)
    divide x y = mul x (cmap recip y)
    arctan2' = undefined
    cmod' m = vmod . cmod' (unMod m) . f2i
    fromInt' = vmod . fromInt'
    toInt'   = toInt . f2i
    fromZ'   = vmod
    toZ'     = f2i

instance (Storable t, Indexable (Vector t) t) => Indexable (Vector (Mod m t)) (Mod m t)
  where
    (!) = (@>)

type instance RealOf (Mod n I) = I
type instance RealOf (Mod n Z) = Z

instance KnownNat m => Product (Mod m I) where
    norm2      = undefined
    absSum     = undefined
    norm1      = undefined
    normInf    = undefined
    multiply   = lift2m (multiplyI m')
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)

instance KnownNat m => Product (Mod m Z) where
    norm2      = undefined
    absSum     = undefined
    norm1      = undefined
    normInf    = undefined
    multiply   = lift2m (multiplyL m')
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)

instance KnownNat m => Normed (Vector (Mod m I))
  where
    norm_0 = norm_0 . toInt
    norm_1 = norm_1 . toInt
    norm_2 = norm_2 . toInt
    norm_Inf = norm_Inf . toInt

instance KnownNat m => Normed (Vector (Mod m Z))
  where
    norm_0 = norm_0 . toZ
    norm_1 = norm_1 . toZ
    norm_2 = norm_2 . toZ
    norm_Inf = norm_Inf . toZ

instance KnownNat m => Numeric (Mod m I)
instance KnownNat m => Numeric (Mod m Z)

i2f :: Storable t => Vector t -> Vector (Mod n t)
i2f v = unsafeFromForeignPtr (castForeignPtr fp) (i) (n)
    where (fp,i,n) = unsafeToForeignPtr v

f2i :: Storable t => Vector (Mod n t) -> Vector t
f2i v = unsafeFromForeignPtr (castForeignPtr fp) (i) (n)
    where (fp,i,n) = unsafeToForeignPtr v

f2iM :: (Element t, Element (Mod n t)) => Matrix (Mod n t) -> Matrix t
f2iM m = m { xdat = f2i (xdat m) }

i2fM :: (Element t, Element (Mod n t)) => Matrix t -> Matrix (Mod n t)
i2fM m = m { xdat = i2f (xdat m) }

vmod :: forall m t. (KnownNat m, Storable t, Integral t, Numeric t) => Vector t -> Vector (Mod m t)
vmod = i2f . cmod' m'
  where
    m' = fromIntegral . natVal $ (undefined :: Proxy m)

lift1 f a   = vmod (f (f2i a))
lift2 f a b = vmod (f (f2i a) (f2i b))

lift2m f a b = liftMatrix vmod (f (f2iM a) (f2iM b))

instance KnownNat m => Num (Vector (Mod m I))
  where
    (+) = lift2 (+)
    (*) = lift2 (*)
    (-) = lift2 (-)
    abs = lift1 abs
    signum = lift1 signum
    negate = lift1 negate
    fromInteger x = fromInt (fromInteger x)

instance KnownNat m => Num (Vector (Mod m Z))
  where
    (+) = lift2 (+)
    (*) = lift2 (*)
    (-) = lift2 (-)
    abs = lift1 abs
    signum = lift1 signum
    negate = lift1 negate
    fromInteger x = fromZ (fromInteger x)

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

instance (KnownNat m) => Testable (Matrix (Mod m I))
  where
    checkT _ = test

test = (ok, info)
  where
    v = fromList [3,-5,75] :: Vector (Mod 11 I)
    m = (3><3) [1..]   :: Matrix (Mod 11 I)

    a = (3><3) [1,2 , 3
               ,4,5 , 6
               ,0,10,-3] :: Matrix I

    b = (3><2) [0..] :: Matrix I

    am = fromInt a :: Matrix (Mod 13 I)
    bm = fromInt b :: Matrix (Mod 13 I)
    ad = fromInt a :: Matrix Float
    bd = fromInt b :: Matrix Float

    g = (3><3) (repeat (40000)) :: Matrix I
    gm = fromInt g :: Matrix (Mod 100000 I)

    lg = (3><3) (repeat (3*10^(9::Int))) :: Matrix Z
    lgm = fromZ lg :: Matrix (Mod 10000000000 Z)

    gen  n = diagRect 1 (konst 5 n) n n :: Numeric t => Matrix t

    rgen n = gen n :: Matrix R
    cgen n = complex (rgen n) + fliprl (complex (rgen n)) * scalar (0:+1) :: Matrix C
    sgen n = single (cgen n)

    checkGen x = norm_Inf $ flatten $ invg x <> x - ident (rows x)

    invg t = gaussElim t (ident (rows t))

    checkLU okf t = norm_Inf $ flatten (l <> u <> p - t)
      where
        (l,u,p,_) = luFact (LU x' p')
          where
            (x',p') = mutable (luST okf) t

    checkSolve aa = norm_Inf $ flatten (aa <> x - bb)
       where
         bb = flipud aa
         x = luSolve' (luPacked' aa) bb

    tmm = diagRect 1 (fromList [2..6]) 5 5 :: Matrix (Mod 19 I)

    info = do
        print v
        print m
        print (tr m)
        print $ v+v
        print $ m+m
        print $ m <> m
        print $ m #> v

        print $ am <> gaussElim am bm - bm
        print $ ad <> gaussElim ad bd - bd

        print g
        print $ g <> g
        print gm
        print $ gm <> gm

        print lg
        print $ lg <> lg
        print lgm
        print $ lgm <> lgm

        putStrLn "checkGen"
        print (checkGen (gen 5 :: Matrix R))
        print (checkGen (gen 5 :: Matrix Float))
        print (checkGen (cgen 5 :: Matrix C))
        print (checkGen (sgen 5 :: Matrix (Complex Float)))
        print (invg (gen 5) :: Matrix (Mod 7 I))
        print (invg (gen 5) :: Matrix (Mod 7 Z))

        print $ mutable (luST (const True)) (gen 5 :: Matrix R)
        print $ mutable (luST (const True)) (gen 5 :: Matrix (Mod 11 Z))

        putStrLn "checkLU"
        print $ checkLU (magnit 0) (gen 5 :: Matrix R)
        print $ checkLU (magnit 0) (gen 5 :: Matrix Float)
        print $ checkLU (magnit 0) (cgen 5 :: Matrix C)
        print $ checkLU (magnit 0) (sgen 5 :: Matrix (Complex Float))
        print $ checkLU (magnit 0) (gen 5 :: Matrix (Mod 7 I))
        print $ checkLU (magnit 0) (gen 5 :: Matrix (Mod 7 Z))

        putStrLn "checkSolve"
        print $ checkSolve (gen 5 :: Matrix R)
        print $ checkSolve (gen 5 :: Matrix Float)
        print $ checkSolve (cgen 5 :: Matrix C)
        print $ checkSolve (sgen 5 :: Matrix (Complex Float))
        print $ checkSolve (gen 5 :: Matrix (Mod 7 I))
        print $ checkSolve (gen 5 :: Matrix (Mod 7 Z))

        putStrLn "luSolve'"
        print $ luSolve' (luPacked' tmm) (ident (rows tmm))
        print $ invershur tmm

    ok = and
      [ toInt (m #> v) == cmod 11 (toInt m #> toInt v )
      , am <> gaussElim_1 am bm == bm
      , am <> gaussElim_2 am bm == bm
      , am <> gaussElim   am bm == bm
      , (checkGen (gen 5 :: Matrix R)) < 1E-15
      , (checkGen (gen 5 :: Matrix Float)) < 2E-7
      , (checkGen (cgen 5 :: Matrix C)) < 1E-15
      , (checkGen (sgen 5 :: Matrix (Complex Float))) < 3E-7
      , (checkGen (gen 5 :: Matrix (Mod 7 I))) == 0
      , (checkGen (gen 5 :: Matrix (Mod 7 Z))) == 0
      , (checkLU (magnit 1E-10) (gen 5 :: Matrix R)) < 2E-15
      , (checkLU (magnit 1E-5) (gen 5 :: Matrix Float)) < 1E-6
      , (checkLU (magnit 1E-10) (cgen 5 :: Matrix C)) < 5E-15
      , (checkLU (magnit 1E-5) (sgen 5 :: Matrix (Complex Float))) < 1E-6
      , (checkLU (magnit 0) (gen 5 :: Matrix (Mod 7 I))) == 0
      , (checkLU (magnit 0) (gen 5 :: Matrix (Mod 7 Z))) == 0
      , checkSolve (gen 5 :: Matrix R) < 2E-15
      , checkSolve (gen 5 :: Matrix Float) < 1E-6
      , checkSolve (cgen 5 :: Matrix C) < 4E-15
      , checkSolve (sgen 5 :: Matrix (Complex Float)) < 1E-6
      , checkSolve (gen 5 :: Matrix (Mod 7 I)) == 0
      , checkSolve (gen 5 :: Matrix (Mod 7 Z)) == 0
      , prodElements (konst (9:: Mod 10 I) (12::Int)) == product (replicate 12 (9:: Mod 10 I))
      , gm <> gm == konst 0 (3,3)
      , lgm <> lgm == konst 0 (3,3)
      , invershur tmm == luSolve' (luPacked' tmm) (ident (rows tmm))
      , luSolve' (luPacked' (tr $ ident 5 :: Matrix (I ./. 2))) (ident 5) == ident 5
      ]


{- |
Module      :  Numeric.GSL.Root
Copyright   :  (c) Alberto Ruiz 2009
License     :  GPL

Maintainer  :  Alberto Ruiz (aruiz at um dot es)
Stability   :  provisional
Portability :  uses ffi

Multidimensional root finding.

<http://www.gnu.org/software/gsl/manual/html_node/Multidimensional-Root_002dFinding.html>

The example in the GSL manual:

@import Numeric.GSL
import Numeric.LinearAlgebra(format)
import Text.Printf(printf)

rosenbrock a b [x,y] = [ a*(1-x), b*(y-x^2) ]

disp = putStrLn . format \"  \" (printf \"%.3f\")

main = do
    let (sol,path) = root Hybrids 1E-7 30 (rosenbrock 1 10) [-10,-5]
    print sol
    disp path

\> main
[1.0,1.0]
 0.000  -10.000  -5.000  11.000  -1050.000
 1.000   -3.976  24.827   4.976     90.203
 2.000   -3.976  24.827   4.976     90.203
 3.000   -3.976  24.827   4.976     90.203
 4.000   -1.274  -5.680   2.274    -73.018
 5.000   -1.274  -5.680   2.274    -73.018
 6.000    0.249   0.298   0.751      2.359
 7.000    0.249   0.298   0.751      2.359
 8.000    1.000   0.878  -0.000     -1.218
 9.000    1.000   0.989  -0.000     -0.108
10.000    1.000   1.000   0.000      0.000
@

-}
-----------------------------------------------------------------------------

module Numeric.GSL.Root (
    uniRoot, UniRootMethod(..),
    uniRootJ, UniRootMethodJ(..),
    root, RootMethod(..),
    rootJ, RootMethodJ(..),
) where

import Data.Packed.Internal
import Data.Packed.Matrix
import Numeric.GSL.Internal
import Foreign.Ptr(FunPtr, freeHaskellFunPtr)
import Foreign.C.Types
import System.IO.Unsafe(unsafePerformIO)

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

data UniRootMethod = Bisection
                   | FalsePos
                   | Brent
                   deriving (Enum, Eq, Show, Bounded)

uniRoot :: UniRootMethod
        -> Double
        -> Int
        -> (Double -> Double)
        -> Double
        -> Double
        -> (Double, Matrix Double)
uniRoot method epsrel maxit fun xl xu = uniRootGen (fi (fromEnum method)) fun xl xu epsrel maxit

uniRootGen m f xl xu epsrel maxit = unsafePerformIO $ do
    fp <- mkDoublefun f
    rawpath <- createMIO maxit 4
                         (c_root m fp epsrel (fi maxit) xl xu)
                         "root"
    let it = round (rawpath @@> (maxit-1,0))
        path = takeRows it rawpath
        [sol] = toLists $ dropRows (it-1) path
    freeHaskellFunPtr fp
    return (sol !! 1, path)

foreign import ccall safe "root"
    c_root:: CInt -> FunPtr (Double -> Double) -> Double -> CInt -> Double -> Double -> TM

-------------------------------------------------------------------------
data UniRootMethodJ = UNewton
                    | Secant
                    | Steffenson
                    deriving (Enum, Eq, Show, Bounded)

uniRootJ :: UniRootMethodJ
        -> Double
        -> Int
        -> (Double -> Double)
        -> (Double -> Double)
        -> Double
        -> (Double, Matrix Double)
uniRootJ method epsrel maxit fun dfun x = uniRootJGen (fi (fromEnum method)) fun
    dfun x epsrel maxit

uniRootJGen m f df x epsrel maxit = unsafePerformIO $ do
    fp <- mkDoublefun f
    dfp <- mkDoublefun df
    rawpath <- createMIO maxit 2
                         (c_rootj m fp dfp epsrel (fi maxit) x)
                         "rootj"
    let it = round (rawpath @@> (maxit-1,0))
        path = takeRows it rawpath
        [sol] = toLists $ dropRows (it-1) path
    freeHaskellFunPtr fp
    return (sol !! 1, path)

foreign import ccall safe "rootj"
    c_rootj :: CInt -> FunPtr (Double -> Double) -> FunPtr (Double -> Double)
            -> Double -> CInt -> Double -> TM

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

data RootMethod = Hybrids
                | Hybrid
                | DNewton
                | Broyden
                deriving (Enum,Eq,Show,Bounded)

-- | Nonlinear multidimensional root finding using algorithms that do not require
-- any derivative information to be supplied by the user.
-- Any derivatives needed are approximated by finite differences.
root :: RootMethod
     -> Double                     -- ^ maximum residual
     -> Int                        -- ^ maximum number of iterations allowed
     -> ([Double] -> [Double])     -- ^ function to minimize
     -> [Double]                   -- ^ starting point
     -> ([Double], Matrix Double)  -- ^ solution vector and optimization path

root method epsabs maxit fun xinit = rootGen (fi (fromEnum method)) fun xinit epsabs maxit

rootGen m f xi epsabs maxit = unsafePerformIO $ do
    let xiv = fromList xi
        n   = dim xiv
    fp <- mkVecVecfun (aux_vTov (checkdim1 n . fromList . f . toList))
    rawpath <- vec xiv $ \xiv' ->
                   createMIO maxit (2*n+1)
                         (c_multiroot m fp epsabs (fi maxit) // xiv')
                         "multiroot"
    let it = round (rawpath @@> (maxit-1,0))
        path = takeRows it rawpath
        [sol] = toLists $ dropRows (it-1) path
    freeHaskellFunPtr fp
    return (take n $ drop 1 sol, path)

foreign import ccall safe "multiroot"
    c_multiroot:: CInt -> FunPtr TVV -> Double -> CInt -> TVM

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

data RootMethodJ = HybridsJ
                 | HybridJ
                 | Newton
                 | GNewton
                deriving (Enum,Eq,Show,Bounded)

-- | Nonlinear multidimensional root finding using both the function and its derivatives.
rootJ :: RootMethodJ
      -> Double                     -- ^ maximum residual
      -> Int                        -- ^ maximum number of iterations allowed
      -> ([Double] -> [Double])     -- ^ function to minimize
      -> ([Double] -> [[Double]])   -- ^ Jacobian
      -> [Double]                   -- ^ starting point
      -> ([Double], Matrix Double)  -- ^ solution vector and optimization path

rootJ method epsabs maxit fun jac xinit = rootJGen (fi (fromEnum method)) fun jac xinit epsabs maxit

rootJGen m f jac xi epsabs maxit = unsafePerformIO $ do
    let xiv = fromList xi
        n   = dim xiv
    fp <- mkVecVecfun (aux_vTov (checkdim1 n . fromList . f . toList))
    jp <- mkVecMatfun (aux_vTom (checkdim2 n . fromLists . jac . toList))
    rawpath <- vec xiv $ \xiv' ->
                   createMIO maxit (2*n+1)
                         (c_multirootj m fp jp epsabs (fi maxit) // xiv')
                         "multiroot"
    let it = round (rawpath @@> (maxit-1,0))
        path = takeRows it rawpath
        [sol] = toLists $ dropRows (it-1) path
    freeHaskellFunPtr fp
    freeHaskellFunPtr jp
    return (take n $ drop 1 sol, path)

foreign import ccall safe "multirootj"
    c_multirootj:: CInt -> FunPtr TVV -> FunPtr TVM -> Double -> CInt -> TVM

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

checkdim1 n v
    | dim v == n = v
    | otherwise = error $ "Error: "++ show n
                        ++ " components expected in the result of the function supplied to root"

checkdim2 n m
    | rows m == n && cols m == n = m
    | otherwise = error $ "Error: "++ show n ++ "x" ++ show n
                        ++ " Jacobian expected in rootJ"


{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE DeriveGeneric       #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE RecordWildCards     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeInType          #-}
{-# LANGUAGE TypeOperators       #-}

-- |
-- Module      : Numeric.Hamilton
-- Description : Hamiltonian dynamics for physical systems on generalized
--               coordinates using automatic differentiation
-- Copyright   : (c) Justin Le 2016
-- License     : BSD-3
-- Maintainer  : justin@jle.im
-- Stability   : unstable
-- Portability : portable
--
-- Simulate physical systems on generalized/arbitrary coordinates using
-- Hamiltonian mechanics and automatic differentiation!
--
-- See the <https://github.com/mstksg/hamilton#readme> for more
-- information on usage!
--

module Numeric.Hamilton
  ( -- * Systems and states
    -- ** Systems
    System
  , mkSystem
  , mkSystem'
  , underlyingPos
    -- ** States
  , Config(..)
  , Phase(..)
  , toPhase
  , fromPhase
    -- * State functions
  , momenta
  , velocities
  , keC
  , keP
  , pe
  , lagrangian
  , hamiltonian
  , hamEqs
    -- * Simulating hamiltonian dynamics
    -- ** Over phase space
  , stepHam
  , evolveHam
  , evolveHam'
    -- ** Over configuration space
    -- | Convenience wrappers over the normal phase-space
    -- steppers/simulators that allow you to provide input and expect
    -- output in configuration space instead of in phase space.  Note that
    -- the simulation itself still runs in phase space, so these all
    -- require conversions to and from phase space under the hood.
  , stepHamC
  , evolveHamC
  , evolveHamC'
  ) where

import           Control.Monad
import           Data.Bifunctor
import           Data.Foldable
import           Data.Kind
import           Data.Maybe
import           Data.Proxy
import           Data.Type.Equality
import           GHC.Generics                        (Generic)
import           GHC.TypeLits
import           GHC.TypeLits.Compare
import           Numeric.AD
import           Numeric.GSL.ODE
import           Numeric.LinearAlgebra.Static        as H
import           Numeric.LinearAlgebra.Static.Vector
import qualified Data.Vector.Generic.Sized           as VG
import qualified Data.Vector.Sized                   as V
import qualified Numeric.LinearAlgebra               as LA

-- | Represents the full state of a system of @n@ generalized coordinates
-- in configuration space (informally, "positions and velocities")
--
-- A configuration space representaiton is more directly "physically
-- meaningful" and intuitive/understandable to humans than a phase space
-- representation.  However, it's much less mathematically ideal to work
-- with because of the lack of some neat underlying symmetries.
--
-- You can convert a @'Config' n@ into a @'Phase' n@ (convert from
-- configuration space to phase space) for a given system with 'toPhase'.
-- This allows you to state your system in configuration space and then
-- convert it to phase space before handing it off to the hamiltonian
-- machinery.
data Config :: Nat -> Type where
    Cfg :: { -- | The current values ("positions") of each of the @n@
             -- generalized coordinates
             cfgPositions :: !(R n)
             -- | The current rate of changes ("velocities") of each of the
             -- @n@ generalized coordinates
           , cfgVelocities :: !(R n)
           }
        -> Config n
  deriving (Generic)

deriving instance KnownNat n => Show (Config n)

-- | Represents the full state of a system of @n@ generalized coordinates
-- in phase space (informally, "positions and momentums").
--
-- Phase space representations are much nicer to work with mathematically
-- because of some neat underlying symmetries.  For one, positions and
-- momentums are "interchangeable" in a system; if you swap every
-- coordinate's positions with their momentums, and also swap them in the
-- equations of motions, you get the same system back.  This isn't the case
-- with configuration space representations.
--
-- A hamiltonian simulation basically describes the trajectory of each
-- coordinate through phase space, so this is the /state/ of the
-- simulation.  However, configuration space representations are much more
-- understandable to humans, so it might be useful to give an initial state
-- in configuration space using 'Config', and then convert it to a 'Phase'
-- with 'toPhase'.
data Phase :: Nat -> Type where
    Phs :: { -- | The current values ("positions") of each of the @n@
             -- generalized coordinates.
             phsPositions :: !(R n)
             -- | The current conjugate momenta ("momentums") to each of
             -- the @n@ generalized coordinates
           , phsMomenta :: !(R n)
           }
        -> Phase n
  deriving (Generic)

deriving instance KnownNat n => Show (Phase n)

-- | Represents a physical system in which physics happens.  A @'System'
-- m n@ is a system whose state described using @n@ generalized coordinates
-- (an "@n@-dimensional" system), where the underlying cartesian coordinate
-- space is @m@-dimensional.
--
-- For the most part, you are supposed to be able to ignore @m@.  @m@ is
-- only provided because it's useful when plotting/drawing the system with
-- a given state back in rectangular coordinates. (The only function that
-- use the @m@ at the moment is 'underlyingPos')
--
-- A @'System' m n@'s state is described using a @'Config' n@ (which
-- describes the system in configuration space) or a @'Phase' n@ (which
-- describes the system in phase space).
data System :: Nat -> Nat -> Type where
    Sys :: { _sysInertia       :: R m
           , _sysCoords        :: R n -> R m
           , _sysJacobian      :: R n -> L m n
           , _sysHessian       :: R n -> V.Vector n (L m n)
           , _sysPotential     :: R n -> Double
           , _sysPotentialGrad :: R n -> R n
           }
        -> System m n

-- | Converts the position of generalized coordinates of a system to the
-- coordinates of the system's underlying cartesian coordinate system.
-- Useful for plotting/drawing the system in cartesian space.
underlyingPos
    :: System m n
    -> R n
    -> R m
underlyingPos = _sysCoords

-- | The potential energy of a system, given the position in the
-- generalized coordinates of the system.
pe  :: System m n
    -> R n
    -> Double
pe = _sysPotential

vec2l
    :: (KnownNat m, KnownNat n)
    => V.Vector m (V.Vector n Double)
    -> L m n
vec2l = rowsL . fmap gvecR

-- | Create a system with @n@ generalized coordinates by describing its
-- coordinate space (by a function from the generalized coordinates to the
-- underlying cartesian coordinates), the inertia of each of those
-- underlying coordinates, and the pontential energy function.
--
-- The potential energy function is expressed in terms of the genearlized
-- coordinate space's positions.
mkSystem
    :: forall m n. (KnownNat m, KnownNat n)
    => R m      -- ^ The "inertia" of each of the @m@ coordinates
                -- in the underlying cartesian space of the system.  This
                -- should be mass for linear coordinates and rotational
                -- inertia for angular coordinates.
    -> (forall a. RealFloat a => V.Vector n a -> V.Vector m a)
                -- ^ Conversion function to convert points in the
                -- generalized coordinate space to the underlying cartesian
                -- space of the system.
    -> (forall a. RealFloat a => V.Vector n a -> a)
                -- ^ The potential energy of the system as a function of
                -- the generalized coordinate space's positions.
    -> System m n
mkSystem m f u = Sys
    { _sysInertia       =                     m
    , _sysCoords        = gvecR             . f           . grVec
    , _sysJacobian      = tr   . vec2l      . jacobianT f . grVec
    , _sysHessian       = tr2  . fmap vec2l . hessianF f  . grVec
    , _sysPotential     =                     u           . grVec
    , _sysPotentialGrad = gvecR             . grad u      . grVec
    }
  where
    tr2 :: forall o. (KnownNat n, KnownNat o)
        => V.Vector m (L n o)
        -> V.Vector n (L m o)
    tr2 = fmap rowsL . traverse lRows
    {-# INLINE tr2 #-}

-- | Convenience wrapper over 'mkSystem' that allows you to specify the
-- potential energy function in terms of the underlying cartesian
-- coordinate space.
mkSystem'
    :: forall m n. (KnownNat m, KnownNat n)
    => R m      -- ^ The "inertia" of each of the @m@ coordinates
                -- in the underlying cartesian space of the system.  This
                -- should be mass for linear coordinates and rotational
                -- inertia for angular coordinates.
    -> (forall a. RealFloat a => V.Vector n a -> V.Vector m a)
                -- ^ Conversion function to convert points in the
                -- generalized coordinate space to the underlying cartesian
                -- space of the system.
    -> (forall a. RealFloat a => V.Vector m a -> a)
                -- ^ The potential energy of the system as a function of
                -- the underlying cartesian coordinate space's positions.
    -> System m n
mkSystem' m f u = mkSystem m f (u . f)

-- | Compute the generalized momenta conjugate to each generalized
-- coordinate of a system by giving the configuration-space state of the
-- system.
--
-- Note that getting the momenta from a @'Phase' n@ involves just using
-- 'phsMomenta'.
momenta
    :: (KnownNat m, KnownNat n)
    => System m n
    -> Config n
    -> R n
momenta Sys{..} Cfg{..} = tr j #> diag _sysInertia #> j #> cfgVelocities
  where
    j = _sysJacobian cfgPositions

-- | Convert a configuration-space representaiton of the state of the
-- system to a phase-space representation.
--
-- Useful because the hamiltonian simulations use 'Phase' as its working
-- state, but 'Config' is a much more human-understandable and intuitive
-- representation.  This allows you to state your starting state in
-- configuration space and convert to phase space for your simulation to
-- use.
toPhase
    :: (KnownNat m, KnownNat n)
    => System m n
    -> Config n
    -> Phase n
toPhase s = Phs <$> cfgPositions <*> momenta s

-- | The kinetic energy of a system, given the system's state in
-- configuration space.
keC :: (KnownNat m, KnownNat n)
    => System m n
    -> Config n
    -> Double
keC s = do
    vs <- cfgVelocities
    ps <- momenta s
    return $ (vs <.> ps) / 2

-- | The Lagrangian of a system (the difference between the kinetic energy
-- and the potential energy), given the system's state in configuration
-- space.
lagrangian
    :: (KnownNat m, KnownNat n)
    => System m n
    -> Config n
    -> Double
lagrangian s = do
    t <- keC s
    u <- pe s . cfgPositions
    return (t - u)

-- | Compute the rate of change of each generalized coordinate by giving
-- the state of the system in phase space.
--
-- Note that getting the velocities from a @'Config' n@ involves just using
-- 'cfgVelocities'.
velocities
    :: (KnownNat m, KnownNat n)
    => System m n
    -> Phase n
    -> R n
velocities Sys{..} Phs{..} = inv jmj #> phsMomenta
  where
    j   = _sysJacobian phsPositions
    jmj = tr j H.<> diag _sysInertia H.<> j

-- | Invert 'toPhase' and convert a description of a system's state in
-- phase space to a description of the system's state in configuration
-- space.
--
-- Possibly useful for showing the phase space representation of a system's
-- state in a more human-readable/human-understandable way.
fromPhase
    :: (KnownNat m, KnownNat n)
    => System m n
    -> Phase n
    -> Config n
fromPhase s = Cfg <$> phsPositions <*> velocities s

-- | The kinetic energy of a system, given the system's state in
-- phase space.
keP :: (KnownNat m, KnownNat n)
    => System m n
    -> Phase n
    -> Double
keP s = do
    ps <- phsMomenta
    vs <- velocities s
    return $ (vs <.> ps) / 2

-- | The Hamiltonian of a system (the sum of kinetic energy and the
-- potential energy), given the system's state in phase space.
hamiltonian
    :: (KnownNat m, KnownNat n)
    => System m n
    -> Phase n
    -> Double
hamiltonian s = do
    t <- keP s
    u <- pe s . phsPositions
    return (t + u)

-- | The "hamiltonian equations" for a given system at a given state in
-- phase space.  Returns the rate of change of the positions and
-- conjugate momenta, which can be used to progress the simulation through
-- time.
--
-- Computed using the maths derived in
-- <https://blog.jle.im/entry/hamiltonian-dynamics-in-haskell.html>.
hamEqs
    :: (KnownNat m, KnownNat n)
    => System m n
    -> Phase n
    -> (R n, R n)
hamEqs Sys{..} Phs{..} = (dHdp, -dHdq)
  where
    mm   = diag _sysInertia
    j    = _sysJacobian phsPositions
    trj  = tr j
    jmj  = trj H.<> mm H.<> j
    ijmj = inv jmj
    dTdq = gvecR
         . flip fmap (_sysHessian phsPositions) $ \djdq ->
             -phsMomenta <.> ijmj #> trj #> mm #> djdq #> ijmj #> phsMomenta
    dHdp = ijmj #> phsMomenta
    dHdq = dTdq + _sysPotentialGrad phsPositions

-- | Step a system through phase space over over a single timestep.
stepHam
    :: forall m n. (KnownNat m, KnownNat n)
    => Double           -- ^ timestep to step through
    -> System m n       -- ^ system to simulate
    -> Phase n          -- ^ initial state, in phase space
    -> Phase n
stepHam r s p = evolveHam @m @n @2 s p (V.fromTuple (0, r))
                  `V.unsafeIndex` 1

-- | Evolve a system using a hamiltonian stepper, with the given initial
-- phase space state.
--
-- Desired solution times provided as a list instead of a sized 'V.Vector'.
-- The output list should be the same length as the input list.
evolveHam'
    :: forall m n. (KnownNat m, KnownNat n)
    => System m n  -- ^ system to simulate
    -> Phase n     -- ^ initial state, in phase space
    -> [Double]    -- ^ desired solution times
    -> [Phase n]
evolveHam' _ _ [] = []
evolveHam' s p0 ts = V.withSizedList (toList ts') $ \(v :: V.Vector s Double) ->
                       case Proxy @2 %<=? Proxy @s of
                         LE Refl -> (if l1 then tail else id)
                                  . toList
                                  $ evolveHam s p0 v
                         NLE{}   -> error "evolveHam': Internal error"
  where
    (l1, ts') = case ts of
      [x] -> (True , [0,x])
      _   -> (False, ts   )

-- | Evolve a system using a hamiltonian stepper, with the given initial
-- phase space state.
evolveHam
    :: forall m n s. (KnownNat m, KnownNat n, KnownNat s, 2 <= s)
    => System m n           -- ^ system to simulate
    -> Phase n              -- ^ initial state, in phase space
    -> V.Vector s Double    -- ^ desired solution times
    -> V.Vector s (Phase n)
evolveHam s p0 ts = fmap toPs . fromJust . V.fromList . LA.toRows
                  $ odeSolveV RKf45 hi eps eps (const f) (fromPs p0) ts'
  where
    hi  = (V.unsafeIndex ts 1 - V.unsafeIndex ts 0) / 100
    eps = 1.49012e-08
    f :: LA.Vector Double -> LA.Vector Double
    f   = uncurry (\p m -> LA.vjoin [p,m])
        . join bimap extract . hamEqs s . toPs
    ts' = VG.fromSized . VG.convert $ ts
    n = fromInteger $ natVal (Proxy @n)
    fromPs :: Phase n -> LA.Vector Double
    fromPs p = LA.vjoin . map extract $ [phsPositions p, phsMomenta p]
    toPs :: LA.Vector Double -> Phase n
    toPs v = Phs pP pM
      where
        Just [pP, pM] = traverse create . LA.takesV [n, n] $ v

-- | A convenience wrapper for 'evolveHam'' that works on configuration
-- space states instead of phase space states.
--
-- Note that the simulation itself still runs in phase space; this function
-- just abstracts over converting to and from phase space for the inputs
-- and outputs.
evolveHamC'
    :: forall m n. (KnownNat m, KnownNat n)
    => System m n       -- ^ system to simulate
    -> Config n         -- ^ initial state, in configuration space
    -> [Double]         -- ^ desired solution times
    -> [Config n]
evolveHamC' s c0 = fmap (fromPhase s) . evolveHam' s (toPhase s c0)

-- | A convenience wrapper for 'evolveHam' that works on configuration
-- space states instead of phase space states.
--
-- Note that the simulation itself still runs in phase space; this function
-- just abstracts over converting to and from phase space for the inputs
-- and outputs.
evolveHamC
    :: forall m n s. (KnownNat m, KnownNat n, KnownNat s, 2 <= s)
    => System m n           -- ^ system to simulate
    -> Config n             -- ^ initial state, in configuration space
    -> V.Vector s Double    -- ^ desired solution times
    -> V.Vector s (Config n)
evolveHamC s c0 = fmap (fromPhase s) . evolveHam s (toPhase s c0)

-- | Step a system through configuration space over over a single timestep.
--
-- Note that the simulation itself still runs in phase space; this function
-- just abstracts over converting to and from phase space for the input
-- and output.
stepHamC
    :: forall m n. (KnownNat m, KnownNat n)
    => Double           -- ^ timestep to step through
    -> System m n       -- ^ system to simulate
    -> Config n         -- ^ initial state, in phase space
    -> Config n
stepHamC r s = fromPhase s . stepHam r s . toPhase s


-- |
-- = Nonlinear transient computing
--
-- This module was developed as a part of author's PhD project:
-- https://www.researchgate.net/project/Theory-and-Modeling-of-Complex-Nonlinear-Delay-Dynamics-Applied-to-Neuromorphic-Computing
--

{-# LANGUAGE BangPatterns #-}

module RC.NTC
  ( new
  , learn
  , learnClassifier
  , predict
  , par0
  , NTCParameters (..)
  ) where

import           Numeric.LinearAlgebra
import           System.Random ( StdGen
                               , split
                               )
import qualified Data.List as List
import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as VM
import           System.IO.Unsafe ( unsafePerformIO )
import qualified Learning
import qualified Numeric.DDE.Model as DDEModel

import           RC.NTC.Types
import qualified RC.NTC.Reservoir as Reservoir
import qualified RC.Helpers as H

-- | An untrained NTC network
new
  :: StdGen
  -> NTCParameters
  -> (Int, Int, Int)
  -- ^ Input dimension, network nodes, and output dimension
  -> NTC
new g par (ind, nodes, out) =
  let iwgen = _inputWeightsGenerator par
      iw = iwgen g (nodes, ind) (_inputWeightsRange par)
      ntc = NTC { _inputWeights = iw
                , _reservoir = _reservoirModel par
                , _outputWeights = Nothing
                , _par = par
                }
  in ntc

-- | Default NTC parameters
par0 :: NTCParameters
par0 = Par
  { _preprocess = id
  , _inputWeightsGenerator = H.randMatrix
  , _postprocess = H.addBiases  -- Usually `id` will work
  , _inputWeightsRange = undefined  -- To be manually set, e.g. (-1, 1)
  , _reservoirModel = Reservoir.genReservoir p
  }
  where
    filt' = DDEModel.BandpassFiltering {
              DDEModel._tau = 7.8125e-3
            , DDEModel._theta = recip 0.34375
            }

    p = DDEModel.RC { DDEModel._filt = filt'
                    , DDEModel._rho = 3.25
                    , DDEModel._fnl = H.hsigmoid (1.09375, 1.5, 0.0)
                    }

-- | Nonlinear transformation performed by an NTC network
forwardPass :: NTC  -- ^ NTC network
            -> Matrix Double  -- ^ Input information
            -> Matrix Double
forwardPass NTC { _par = Par { _preprocess = prep, _postprocess = post }
                , _inputWeights = iw
                , _reservoir = Reservoir res
                } !sample =
  let pipeline = post. res. (iw <>). prep
  in pipeline sample

-- | NTC training: learn the readout weights offline
learn
  :: NTC
  -> Int
  -- ^ Discard the first N points
  -> Matrix Double
  -- ^ Input matrix of features rows and observations columns
  -> Matrix Double
  -- ^ Desired output matrix of observations columns
  -> Either String NTC
learn ntc forgetPts inp out = ntc'
  where
    state' = forwardPass ntc inp ?? (All, Drop forgetPts)
    teacher' = out ?? (All, Drop forgetPts)
    ntc' = case Learning.learn' state' teacher' of
      Nothing -> Left "Cannot create a readout matrix"
      w -> Right $ ntc { _outputWeights = w }

_concatForwardPass :: NTC
                   -> Int
                   -- ^ Number of samples in the list (for
                   -- memory allocation)
                   -> [Matrix Double]
                   -- ^ List of samples
                   -> Matrix Double
_concatForwardPass ntc m (state0_:samples') = unsafePerformIO $ do
      -- Detect the postprocessing output dimension `nodes`.
      -- Alternatively, use Static shapes from hmatrix.
      let state0 = forwardPass ntc state0_
          nodes = rows state0
          state0' = H.flatten' state0

      -- Allocate memory for a mutable vector
      v <- VM.new (nodes * m)

      -- Copy the first computed state0'
      copy state0' v 0 0 nodes

      let foldA _ [] = return ()
          foldA f ((y, i):ys) = do
            let v0 = f y
            copy v0 v 0 (i * nodes) ((i + 1) * nodes)
            foldA f ys

      -- NB: Does compiler know (H.flatten'. H.unflatten') == id?
      -- If not, consider refactoring genReservoir and forwardPass
      foldA (H.flatten'. forwardPass ntc) $ zip samples' [1..m - 1]

      processed' <- V.unsafeFreeze v
      let processed = H.unflatten' nodes processed'
      return processed
  where
    copy !v0 !v !i0 !i !k
      | i < k = do
        VM.unsafeWrite v i (v0 V.! i0)
        copy v0 v (i0 + 1) (i + 1) k
      | otherwise = return ()

-- | NTC training specific to classification task.
-- The readout weights are learned offline.
--
-- Alternatively, one could use `learn` function. However,
-- to make sure the training samples do not mix in the reservoir RNN,
-- a significant padding of zeros would be needed. Although that
-- way directly corresponds to a physical experiment, that would be not
-- memory-efficient on a computer.
learnClassifier
  :: NTC
  -- ^ NTC network
  -> [Int]
  -- ^ Target labels
  -> Int
  -- ^ Number of inputs (for memory allocation)
  -> [Matrix Double]
  -- ^ Training inputs
  -> Either String (Learning.Classifier Int)
learnClassifier ntc labels m samples =
  case Learning.learn' state teacher' of
    Nothing -> Left "Cannot create a readout matrix"
    Just w -> let clf = Learning.winnerTakesAll w klasses. forwardPass ntc
           in Right (Learning.Classifier clf)
  where
    klasses = fromList. List.sort. List.nub $ labels
    klassesNo = V.length klasses
    teacher' = concatHor. map (flip (Learning.teacher klassesNo) 1) $ labels

    -- Alternatively, a streaming interface might be a solution
    -- https://hackage.haskell.org/package/pipes-4.3.7/docs/Pipes-Tutorial.html
    state = _concatForwardPass ntc m samples

-- | Horizontal matrix concatenation, alternative to fromBlocks
concatHor :: Element a => [Matrix a] -> Matrix a
concatHor ms@(m:_) = foldr (|||) (zeroCols (rows m)) ms
-- concatHor = concatMapHor id
{-# SPECIALIZE concatHor :: [Matrix Double] -> Matrix Double #-}

concatMapHor
  :: Element b =>
     (Matrix a -> Matrix b) -> [Matrix a] -> Matrix b
concatMapHor f ms@(m:_) = foldr (\a b -> f a ||| b) (zeroCols (rows m)) ms
{-# SPECIALIZE concatMapHor :: (Matrix Double -> Matrix Double) -> [Matrix Double] -> Matrix Double #-}

-- | Matrix with zero elements
zeroCols :: V.Storable a => Int -> Matrix a
zeroCols rows' = (rows'><0) []
{-# SPECIALIZE zeroCols :: Int -> Matrix Double #-}

-- | Run prediction using a "clean" (uninitialized) reservoir and then
-- forget the reservoir's state.
-- This can be used for both forecasting and classification tasks.
predict :: NTC  -- ^ Trained network
        -> Int  -- ^ Washout (forget) points
        -> Matrix Double
        -- ^ Input matrix where measurements are columns and features are rows
        -> Either String (Matrix Double)
        -- ^ Either error string or predicted output
predict ntc@NTC { _outputWeights = ow
                } forgetPts inp =
  case ow of
    Nothing -> Left "Please train the NTC first"
    Just w -> let y = forwardPass ntc inp
                  y2 = y ?? (All, Drop forgetPts)
                  prediction = w <> y2
              in Right prediction


module Q.Options.ImpliedVol.Normal where
import           Data.Default.Class (Default (..))
import           Numeric.IEEE (epsilon)
import           Numeric.Natural (Natural)
import           Numeric.RootFinding as R (RiddersParam (RiddersParam),
                                           Root (NotBracketed, Root, SearchFailed), Tolerance,
                                           ridders)
import           Q.Options.Bachelier (Bachelier (Bachelier), euOption)

import           Statistics.Distribution (cumulative, density)
import           Statistics.Distribution.Normal (standard)
import Q.Options

-- | Method to use to calculate the normal implied vol.
data Method =
    Jackel        -- ^ Jackel analytical formula approximation.
  | ChoKimKwak    -- ^ J. Choi, K kim, and M. Kwak (2009)
  -- | Numerical root finding. Currently Ridders is used.
  | RootFinding
        Natural                  -- ^ Maximum number of iterations.
        R.Tolerance              -- ^ Tolerance (relative or absolute)
        (Double, Double, Double) -- ^ Triple of @(low bound, initial
                                 --   guess, upper bound)@. If initial
                                 --   guess if out of bracket middle
                                 --   of bracket is taken as.

instance Default Method where
  def = Jackel

-- | Default method implementation of 'euImpliedVolWith' using 'Jackel'.
euImpliedVol :: OptionType -> Forward -> Strike -> YearFrac -> Rate -> Premium -> Vol
euImpliedVol = euImpliedVolWith def

-- | Calcualte the bachelier option implied vol of a european option.
--
-- If the options premium does not have time value @'hasTimeValue'@ return 0.
euImpliedVolWith :: Method -> OptionType -> Forward -> Strike -> YearFrac -> Rate -> Premium -> Vol
euImpliedVolWith m cp f k t r p
  | hasTimeValue cp f k p df = euImpliedVolWith' m cp f k t r p
  | otherwise                = Vol $ 0
  where df = discountFactor t r

euImpliedVolWith' :: Method-> OptionType-> Forward-> Strike-> YearFrac-> Rate-> Premium-> Vol
euImpliedVolWith' Jackel cp (Forward f) (Strike k) (YearFrac t) (Rate r) (Premium p)
  -- Case where interest rate is not 0 we need undiscount. The paper is written
  -- for the undiscounted Bachelier option prices.
  | r /= 0
    = euImpliedVol cp (Forward f) (Strike k) (YearFrac t) (Rate 0) (Premium (p/df))
  -- Case of ATM. Solve directly.
  | abs (k - f) <= epsilon = Vol $ p * sqrt2Pi / (sqrt t)
  -- Case of ITM option. Calcualte vol of the out of the money option with Put-Call-Parity.
  | phiStarTilde >= 0
    = euImpliedVol (succ cp) (Forward f) (Strike k) (YearFrac t) (Rate r) (Premium p')
  -- General case for an out of the money option.
  | otherwise  = let
      ẋ      = if phiStarTilde < c then
                 let g = 1 / (phiStarTilde - 0.5)
                     ξ = (0.032114372355 - (g**2)*(0.016969777977 - (g**2)*(2.6207332461E-3-(9.6066952861E-5)*g**2)))
                         /
                         (1-(g**2)*(0.6635646938 - (g**2)*(0.14528712196 - 0.010472855461*g**2)))
                 in g * (1 / sqrt2Pi + ξ*g**2)
               else
                 let h = sqrt $ (-log (-phiStarTilde))
                 in (9.4883409779-h*(9.6320903635-h*(0.58556997323 + 2.1464093351*h)))
                    /
                    (1-h*(0.65174820867 + h*(1.5120247828 + 6.6437847132E-5*h)))
      c       = (-0.001882039271)
      x       = ẋ + (3*q * ẋ * ẋ * (2 - q * ẋ * (2 + ẋ*ẋ)))
                    /
                    (6 + q*ẋ * ((-12) + ẋ *(6*q + ẋ * ((-6)*q*ẋ*(3+ẋ*ẋ)))))
      phiXBarTilde = (cumulative standard ẋ) + (density standard ẋ)/ẋ
      q       =  (phiXBarTilde-phiStarTilde)/ (density standard ẋ)
    in Vol $ (abs (k - f)) / (abs (x * sqrt t))
  where phiStarTilde = negate $ (abs (p - (max (theta * (f - k)) 0))) / (abs (k - f))
        theta        = if cp == Call then 1 else -1
        phiTilde     = (-theta) * p / (k - f)
        p'           = cpi * df * (f - k) + p
        cpi          = fromIntegral $ fromEnum cp --call put indicartor.
        df           = exp $ (-r) * t
        sqrt2Pi      = 2.506628274631000

euImpliedVolWith' ChoKimKwak cp (Forward f) (Strike k) (YearFrac t) (Rate r) (Premium p) =
  let df              = exp $ (-r) * t
      forwardPremium  = p / df
      straddlePremium = case cp of Call -> 2 * forwardPremium - (f - k)
                                   Put  -> 2 * forwardPremium + (f - k)
      nu'             = (f - k) / straddlePremium
      nu              = max (-1 + epsilon) (min nu' (1 - epsilon))
      eta             | abs nu < sqrtEpsilon = 1
                      | otherwise            = nu / (atanh nu)
      heta            = h eta
  in Vol $ sqrt (pi / (2 * t)) * straddlePremium * heta

euImpliedVolWith' (RootFinding maxIter tol bracket)  cp (Forward forward) k t r (Premium p) =
  let f vol        = p' - p  where
        (Premium p') = vPremium $ euOption b cp k
        b            = Bachelier t (Forward forward) r (Vol vol)
      (lb, _, ub) = bracket
      root = ridders (RiddersParam (fromEnum maxIter) tol) (lb, ub) f
  in case root of (Root vol)   -> Vol vol
                  NotBracketed -> error "not bracketed"
                  SearchFailed -> error "search failed"

sqrtEpsilon :: Double
sqrtEpsilon = sqrt epsilon

h :: Floating p => p -> p
h eta = sqrt(eta) * (num / den) where
  a0          = 3.994961687345134e-1
  a1          = 2.100960795068497e+1
  a2          = 4.980340217855084e+1
  a3          = 5.988761102690991e+2
  a4          = 1.848489695437094e+3
  a5          = 6.106322407867059e+3
  a6          = 2.493415285349361e+4
  a7          = 1.266458051348246e+4

  b0          = 1.000000000000000e+0
  b1          = 4.990534153589422e+1
  b2          = 3.093573936743112e+1
  b3          = 1.495105008310999e+3
  b4          = 1.323614537899738e+3
  b5          = 1.598919697679745e+4
  b6          = 2.392008891720782e+4
  b7          = 3.608817108375034e+3
  b8          = -2.067719486400926e+2
  b9          = 1.174240599306013e+1

  num = a0 + eta * (a1 + eta * (a2 + eta * (a3 + eta * (a4 + eta * (a5 + eta * (a6 + eta * a7))))))
  den = b0 + eta * (b1 + eta * (b2 + eta * (b3 + eta * (b4 + eta * (b5 + eta * (b6 + eta * (b7 + eta * (b8 + eta * b9))))))))


module STCR2Z2T0S0 where

import           Control.Monad             as M
import           Control.Monad.Parallel    as MP
import           Data.Array.Repa           as R
import           Data.Binary
import           Data.Complex
import           Data.List                 as L
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.DomainChange
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
import           Utils.Array

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:scale0FreqsStr:scaleFreqsStr:initDistStr:histFilePath:alphaStr:radialFlagStr: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]
      initDist = read initDistStr :: [R2S1RPPoint]
      scale0Freq = read scale0FreqsStr :: Double
      scaleFreq = read scaleFreqsStr :: Double
      scale0Freqs = [-scale0Freq .. scale0Freq]
      scaleFreqs = [-scaleFreq .. scaleFreq]
      alpha = read alphaStr :: Double
      radialFlag = read radialFlagStr :: Bool
      numThread = read numThreadStr :: Int
      sourceDist = L.take 1 initDist
      sinkDist = L.drop 1 initDist
      folderPath = "output/test/STCR2Z2T0S0"
      dirName = takeDirectory histFilePath
      fileName = takeFileName histFilePath
      newHistFilePath =
        if radialFlag
          then dirName </> ("Radial_" L.++ fileName)
          else histFilePath
      (R2S1RPPoint (_, _, _, s)) = L.head sourceDist
  createDirectoryIfMissing True folderPath
  arrR2Z2T0S0 <-
    if radialFlag
      then do
        flag <- doesFileExist newHistFilePath
        radialArr <-
          if flag
            then R.map magnitude . getNormalizedHistogramArr <$>
                 decodeFile newHistFilePath
            else do
              putStrLn
                "Couldn't find a Green's function data. Start simulation..."
              solveMonteCarloR2Z2T0S0Radial
                numThread
                numTrail
                maxTrail
                numPoint
                numPoint
                thetaSigma
                scaleSigma
                maxScale
                tao
                theta0Freqs
                thetaFreqs
                scale0Freqs
                scaleFreqs
                newHistFilePath
                (emptyHistogram
                   [ if maxScale >= 2
                        then round maxScale + 1
                        else (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
                   , L.length scale0Freqs
                   , L.length theta0Freqs
                   , L.length scaleFreqs
                   , L.length thetaFreqs
                   ]
                   0)
        computeUnboxedP $
          computeR2Z2T0S0ArrayRadial
            Pinwheel
            radialArr
            numPoint
            numPoint
            1
            maxScale
            thetaFreqs
            scaleFreqs
            theta0Freqs
            scale0Freqs
      else do
        flag <- doesFileExist newHistFilePath
        hist <-
          if flag
            then decodeFile newHistFilePath
            else return $
                 emptyHistogram
                   [ numPoint
                   , numPoint
                   , L.length scale0Freqs
                   , L.length theta0Freqs
                   , L.length scaleFreqs
                   , L.length thetaFreqs
                   ]
                   0
        solveMonteCarloR2Z2T0S0
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          tao
          theta0Freqs
          thetaFreqs
          scale0Freqs
          scaleFreqs
          newHistFilePath
          hist
  plan <- makeR2Z2T0S0Plan emptyPlan False "" arrR2Z2T0S0
  sourceDistArr <-
    computeInitialDistributionR2T0S0
      plan
      numPoint
      numPoint
      theta0Freqs
      scale0Freqs
      maxScale
      -- initDist
      sourceDist
  sinkDistArr <-
    computeInitialDistributionR2T0S0
      plan
      numPoint
      numPoint
      theta0Freqs
      scale0Freqs
      maxScale
      sinkDist
  arrR2Z2T0S0F <- dftR2Z2T0S0 plan . computeS . makeFilter2D $ arrR2Z2T0S0
  -- Source field
  sourceArr <- convolveR2T0S0 plan arrR2Z2T0S0F sourceDistArr
  sourceR2Z2 <- R.sumP . R.sumS . rotateR2Z2T0S0Array $ sourceArr
  sourceField <-
    fmap (computeS . R.extend (Z :. (1 :: Int) :. All :. All)) .
    R.sumP .
    R.sumS .
    rotate4D2 .
    r2z2Tor2s1rp numOrientation thetaFreqs numScale scaleFreqs maxScale $
    sourceR2Z2
  plotImageRepaComplex (folderPath </> printf "Source%.0f.png" s) . ImageRepa 8 $
    sourceField
  -- Sink field
  sinkArr' <- convolveR2T0S0 plan arrR2Z2T0S0F sinkDistArr
  let sinkArr = computeSinkFromSourceR2Z2T0S0 thetaFreqs theta0Freqs sinkArr'
  sinkR2Z2 <- R.sumP . R.sumS . rotateR2Z2T0S0Array $ sinkArr
  sinkField <-
    fmap (computeS . R.extend (Z :. (1 :: Int) :. All :. All)) .
    R.sumP .
    R.sumS .
    rotate4D2 .
    r2z2Tor2s1rp numOrientation thetaFreqs numScale scaleFreqs maxScale $
    sinkR2Z2
  plotImageRepaComplex (folderPath </> printf "Sink%.0f.png" s) . ImageRepa 8 $
    sinkField
  -- Completion Filed
  completionFiled <- convolveR2Z2 plan sourceR2Z2 sinkR2Z2
  completionFiledR2 <-
    R.sumP .
    R.sumS .
    rotate4D2 .
    r2z2Tor2s1rp numOrientation thetaFreqs numScale scaleFreqs maxScale $
    completionFiled
  plotImageRepaComplex (folderPath </> printf "Completion%.0f.png" s) .
    ImageRepa 8 . computeS . R.extend (Z :. (1 :: Int) :. All :. All) $
    completionFiledR2
  plotImageRepa (folderPath </> printf "CompletionMagnitude%.0f.png" s) .
    ImageRepa 8 .
    computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.map magnitude $
    completionFiledR2