{-# LANGUAGE CPP                   #-}
{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RebindableSyntax      #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeSynonymInstances  #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE ViewPatterns          #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- |
-- Module      : Data.Array.Accelerate.Data.Complex
-- Copyright   : [2015..2017] Trevor L. McDonell
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <tmcdonell@cse.unsw.edu.au>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- Complex numbers, stored in the usual C-style array-of-struct representation,
-- for easy interoperability.
--
module Data.Array.Accelerate.Data.Complex (

  -- * Rectangular from
  Complex(..),
  real,
  imag,

  -- * Polar form
  mkPolar,
  cis,
  polar,
  magnitude,
  phase,

  -- * Conjugate
  conjugate,

) where

import Data.Array.Accelerate.Array.Sugar
import Data.Array.Accelerate.Classes
import Data.Array.Accelerate.Data.Functor
import Data.Array.Accelerate.Prelude
import Data.Array.Accelerate.Product
import Data.Array.Accelerate.Smart
import Data.Array.Accelerate.Type

import Prelude                                                      ( ($) )
import Data.Complex                                                 ( Complex(..) )
import qualified Data.Complex                                       as C
import qualified Prelude                                            as P

-- Use an array-of-structs representation for complex numbers. This matches the
-- standard C-style layout, but means that we can define instances only at
-- specific types (not for any type 'a') as we can only have vectors of
-- primitive type.
--

instance Elt (Complex Half) where
  type EltRepr (Complex Half) = V2 Half
  {-# INLINE eltType     #-}
  {-# INLINE [1] toElt   #-}
  {-# INLINE [1] fromElt #-}
  eltType          = TypeRscalar scalarType
  toElt (V2 r i)   = r :+ i
  fromElt (r :+ i) = V2 r i

instance Elt (Complex Float) where
  type EltRepr (Complex Float) = V2 Float
  {-# INLINE eltType     #-}
  {-# INLINE [1] toElt   #-}
  {-# INLINE [1] fromElt #-}
  eltType          = TypeRscalar scalarType
  toElt (V2 r i)   = r :+ i
  fromElt (r :+ i) = V2 r i

instance Elt (Complex Double) where
  type EltRepr (Complex Double) = V2 Double
  {-# INLINE eltType     #-}
  {-# INLINE [1] toElt   #-}
  {-# INLINE [1] fromElt #-}
  eltType          = TypeRscalar scalarType
  toElt (V2 r i)   = r :+ i
  fromElt (r :+ i) = V2 r i

instance Elt (Complex CFloat) where
  type EltRepr (Complex CFloat) = V2 Float
  {-# INLINE eltType     #-}
  {-# INLINE [1] toElt   #-}
  {-# INLINE [1] fromElt #-}
  eltType                        = TypeRscalar scalarType
  toElt (V2 r i)                 = CFloat r :+ CFloat i
  fromElt (CFloat r :+ CFloat i) = V2 r i

instance Elt (Complex CDouble) where
  type EltRepr (Complex CDouble) = V2 Double
  {-# INLINE eltType     #-}
  {-# INLINE [1] toElt   #-}
  {-# INLINE [1] fromElt #-}
  eltType                          = TypeRscalar scalarType
  toElt (V2 r i)                   = CDouble r :+ CDouble i
  fromElt (CDouble r :+ CDouble i) = V2 r i

instance cst a => IsProduct cst (Complex a) where
  type ProdRepr (Complex a) = ProdRepr (a,a)
  fromProd (r :+ i) = fromProd @cst (r,i)
  toProd p          = let (r,i) = toProd @cst p in (r :+ i)
  prod              = prod @cst @(a,a)

instance (Lift Exp a, Elt (Plain a), Elt (Complex (Plain a))) => Lift Exp (Complex a) where
  type Plain (Complex a) = Complex (Plain a)
  lift (r :+ i)          = Exp $ Tuple (NilTup `SnocTup` lift r `SnocTup` lift i)

instance (Elt a, Elt (Complex a)) => Unlift Exp (Complex (Exp a)) where
  unlift e
    = let r     = Exp $ SuccTupIdx ZeroTupIdx `Prj` e
          i     = Exp $ ZeroTupIdx `Prj` e
      in
      r :+ i

instance (Eq a, Elt (Complex a)) => Eq (Complex a) where
  x == y = let r1 :+ c1 = unlift x
               r2 :+ c2 = unlift y
           in  r1 == r2 && c1 == c2
  x /= y = let r1 :+ c1 = unlift x
               r2 :+ c2 = unlift y
           in  r1 /= r2 || c1 /= c2

instance (RealFloat a, Elt (Complex a)) => P.Num (Exp (Complex a)) where
  (+)           = lift2 ((+) :: Complex (Exp a) -> Complex (Exp a) -> Complex (Exp a))
  (-)           = lift2 ((-) :: Complex (Exp a) -> Complex (Exp a) -> Complex (Exp a))
  (*)           = lift2 ((*) :: Complex (Exp a) -> Complex (Exp a) -> Complex (Exp a))
  negate        = lift1 (negate :: Complex (Exp a) -> Complex (Exp a))
  signum z      = if z == 0
                    then z
                    else let x :+ y = unlift z
                             r      = magnitude z
                         in
                         lift (x/r :+ y/r)
  abs z         = lift (magnitude z :+ 0)
  fromInteger n = lift (fromInteger n :+ 0)

instance (RealFloat a, Elt (Complex a)) => P.Fractional (Exp (Complex a)) where
  fromRational x  = lift (fromRational x :+ 0)
  z / z'          = lift ((x*x''+y*y'') / d :+ (y*x''-x*y'') / d)
    where
      x  :+ y   = unlift z
      x' :+ y'  = unlift z'
      --
      x'' = scaleFloat k x'
      y'' = scaleFloat k y'
      k   = - max (exponent x') (exponent y')
      d   = x'*x'' + y'*y''

instance (RealFloat a, Elt (Complex a)) => P.Floating (Exp (Complex a)) where
  pi                      = lift $ pi :+ 0

  exp (unlift -> x :+ y)  = let expx = exp x
                            in  complex $ expx * cos y :+ expx * sin y

  log z                   = lift $ log (magnitude z) :+ phase z

  sqrt z@(unlift -> x :+ y) =
    if z == 0
      then 0
      else lift $ u :+ (y < 0 ? (-v, v))
    where
      (u,v) = unlift (x < 0 ? (lift (v',u'), lift (u',v')))
      v'    = abs y / (u'*2)
      u'    = sqrt ((magnitude z + abs x) / 2)

  x ** y =
    if y == 0 then 1 else
    if x == 0 then if exp_r > 0 then 0 else
                   if exp_r < 0 then lift (inf :+ 0)
                                else lift (nan :+ nan)
              else if isInfinite r || isInfinite i
                     then if exp_r > 0 then lift (inf :+ 0) else
                          if exp_r < 0 then 0
                                       else lift (nan :+ nan)
                     else exp (log x * y)
    where
      r     :+ i  = unlift x
      exp_r :+ _  = unlift y
      --
      inf = 1 / 0
      nan = 0 / 0

  sin (unlift -> x :+ y)  = complex $ sin x * cosh y :+ cos x * sinh y
  cos (unlift -> x :+ y)  = complex $ cos x * cosh y :+ (- sin x * sinh y)
  tan (unlift -> x :+ y)  = (complex $ sinx*coshy :+ cosx*sinhy) / (complex $ cosx*coshy :+ (-sinx*sinhy))
    where
      sinx  = sin x
      cosx  = cos x
      sinhy = sinh y
      coshy = cosh y

  sinh (unlift -> x :+ y) = complex $ cos y * sinh x :+ sin  y * cosh x
  cosh (unlift -> x :+ y) = complex $ cos y * cosh x :+ sin y * sinh x
  tanh (unlift -> x :+ y) = (complex $ cosy*sinhx :+ siny*coshx) / (complex $ cosy*coshx :+ siny*sinhx)
    where
      siny  = sin y
      cosy  = cos y
      sinhx = sinh x
      coshx = cosh x

  asin z@(unlift -> x :+ y) = complex $ y' :+ (-x')
    where
      x' :+ y' = unlift $ log ((complex ((-y):+x)) + sqrt (1 - z*z))

  acos z                    = complex $ y'' :+ (-x'')
    where
      x'' :+ y''  = unlift $ log (z + (complex ((-y') :+ x')))
      x'  :+ y'   = unlift $ sqrt (1 - z*z)

  atan z@(unlift -> x :+ y) = complex $ y' :+ (-x')
    where
      x' :+ y' = unlift $ log ((complex ((1-y):+x)) / sqrt (1+z*z))

  asinh z =  log (z + sqrt (1+z*z))
  acosh z =  log (z + (z+1) * sqrt ((z-1)/(z+1)))
  atanh z =  0.5 * log ((1.0+z) / (1.0-z))

instance (FromIntegral a b, Num b, Elt (Complex b)) => FromIntegral a (Complex b) where
  fromIntegral x = lift (fromIntegral x :+ 0)

-- | @since 1.2.0.0
instance Functor Complex where
  fmap f (unlift -> r :+ i) = lift (f r :+ f i)

-- Helper function to fix the types for lift (ugh)
--
complex :: (Elt a, Elt (Complex a)) => Complex (Exp a) -> Exp (Complex a)
complex = lift

-- | The non-negative magnitude of a complex number
--
magnitude :: (RealFloat a, Elt (Complex a)) => Exp (Complex a) -> Exp a
-- magnitude (unlift -> r :+ i) = sqrt (r*r + i*i)
magnitude (unlift -> r :+ i) = scaleFloat k (sqrt (sqr (scaleFloat mk r) + sqr (scaleFloat mk i)))
  where
    k     = max (exponent r) (exponent i)
    mk    = -k
    sqr z = z * z

-- | The phase of a complex number, in the range @(-'pi', 'pi']@. If the
-- magnitude is zero, then so is the phase.
--
phase :: (RealFloat a, Elt (Complex a)) => Exp (Complex a) -> Exp a
phase z@(unlift -> r :+ i) =
  if z == 0
    then 0
    else atan2 i r

-- | The function 'polar' takes a complex number and returns a (magnitude,
-- phase) pair in canonical form: the magnitude is non-negative, and the phase
-- in the range @(-'pi', 'pi']@; if the magnitude is zero, then so is the phase.
--
polar :: (RealFloat a, Elt (Complex a)) => Exp (Complex a) -> Exp (a,a)
polar z =  lift (magnitude z, phase z)

-- | Form a complex number from polar components of magnitude and phase.
--
#if __GLASGOW_HASKELL__ <= 708
mkPolar :: forall a. (RealFloat a, Elt (Complex a)) => Exp a -> Exp a -> Exp (Complex a)
#else
mkPolar :: forall a. (Floating a,  Elt (Complex a)) => Exp a -> Exp a -> Exp (Complex a)
#endif
mkPolar = lift2 (C.mkPolar :: Exp a -> Exp a -> Complex (Exp a))

-- | @'cis' t@ is a complex value with magnitude @1@ and phase @t@ (modulo
-- @2*'pi'@).
--
#if __GLASGOW_HASKELL__ <= 708
cis :: forall a. (RealFloat a, Elt (Complex a)) => Exp a -> Exp (Complex a)
#else
cis :: forall a. (Floating a,  Elt (Complex a)) => Exp a -> Exp (Complex a)
#endif
cis = lift1 (C.cis :: Exp a -> Complex (Exp a))

-- | Return the real part of a complex number
--
real :: (Elt a, Elt (Complex a)) => Exp (Complex a) -> Exp a
real (unlift -> r :+ _) = r

-- | Return the imaginary part of a complex number
--
imag :: (Elt a, Elt (Complex a)) => Exp (Complex a) -> Exp a
imag (unlift -> _ :+ i) = i

-- | Return the complex conjugate of a complex number, defined as
--
-- > conjugate(Z) = X - iY
--
conjugate :: (Num a, Elt (Complex a)) => Exp (Complex a) -> Exp (Complex a)
conjugate z = lift $ real z :+ (- imag z)


{-# LANGUAGE CPP                #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveFoldable     #-}
{-# LANGUAGE DeriveFunctor      #-}
{-# LANGUAGE DeriveGeneric      #-}
{-# LANGUAGE FlexibleContexts   #-}
{-# LANGUAGE ViewPatterns       #-}
-- |
-- Module    : Statistics.Quantile
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Functions for approximating quantiles, i.e. points taken at regular
-- intervals from the cumulative distribution function of a random
-- variable.
--
-- The number of quantiles is described below by the variable /q/, so
-- with /q/=4, a 4-quantile (also known as a /quartile/) has 4
-- intervals, and contains 5 points.  The parameter /k/ describes the
-- desired point, where 0 ≤ /k/ ≤ /q/.

module Statistics.Quantile
    (
    -- * Quantile estimation functions
    -- $cont_quantiles
      ContParam(..)
    , Default(..)
    , quantile
    , quantiles
    , quantilesVec
    -- ** Parameters for the continuous sample method
    , cadpw
    , hazen
    , spss
    , s
    , medianUnbiased
    , normalUnbiased
    -- * Other algorithms
    , weightedAvg
    -- * Median & other specializations
    , median
    , mad
    , midspread
    -- * Deprecated
    , continuousBy
    -- * References
    -- $references
    ) where

import           Data.Binary            (Binary)
import           Data.Aeson             (ToJSON,FromJSON)
import           Data.Data              (Data,Typeable)
import           Data.Default.Class
import           Data.Functor
import qualified Data.Foldable        as F
import           Data.Vector.Generic ((!))
import qualified Data.Vector          as V
import qualified Data.Vector.Generic  as G
import qualified Data.Vector.Unboxed  as U
import qualified Data.Vector.Storable as S
import GHC.Generics (Generic)

import Statistics.Function (partialSort)

----------------------------------------------------------------
-- Quantile estimation
----------------------------------------------------------------

-- | O(/n/·log /n/). Estimate the /k/th /q/-quantile of a sample,
-- using the weighted average method. Up to rounding errors it's same
-- as @quantile s@.
--
-- The following properties should hold otherwise an error will be thrown.
--
--   * the length of the input is greater than @0@
--
--   * the input does not contain @NaN@
--
--   * k ≥ 0 and k ≤ q
weightedAvg :: G.Vector v Double =>
               Int        -- ^ /k/, the desired quantile.
            -> Int        -- ^ /q/, the number of quantiles.
            -> v Double   -- ^ /x/, the sample data.
            -> Double
weightedAvg k q x
  | G.any isNaN x   = modErr "weightedAvg" "Sample contains NaNs"
  | n == 0          = modErr "weightedAvg" "Sample is empty"
  | n == 1          = G.head x
  | q < 2           = modErr "weightedAvg" "At least 2 quantiles is needed"
  | k == q          = G.maximum x
  | k >= 0 || k < q = xj + g * (xj1 - xj)
  | otherwise       = modErr "weightedAvg" "Wrong quantile number"
  where
    j   = floor idx
    idx = fromIntegral (n - 1) * fromIntegral k / fromIntegral q
    g   = idx - fromIntegral j
    xj  = sx ! j
    xj1 = sx ! (j+1)
    sx  = partialSort (j+2) x
    n   = G.length x
{-# SPECIALIZE weightedAvg :: Int -> Int -> U.Vector Double -> Double #-}
{-# SPECIALIZE weightedAvg :: Int -> Int -> V.Vector Double -> Double #-}
{-# SPECIALIZE weightedAvg :: Int -> Int -> S.Vector Double -> Double #-}

----------------------------------------------------------------
-- Quantiles continuous algorithm
----------------------------------------------------------------

-- $cont_quantiles
--
-- Below is family of functions which use same algorithm for estimation
-- of sample quantiles. It approximates empirical CDF as continuous
-- piecewise function which interpolates linearly between points
-- \((X_k,p_k)\) where \(X_k\) is k-th order statistics (k-th smallest
-- element) and \(p_k\) is probability corresponding to
-- it. 'ContParam' determines how \(p_k\) is chosen. For more detailed
-- explanation see [Hyndman1996].
--
-- This is the method used by most statistical software, such as R,
-- Mathematica, SPSS, and S.

-- | Parameters /α/ and /β/ to the 'continuousBy' function. Exact
--   meaning of parameters is described in [Hyndman1996] in section
--   \"Piecewise linear functions\"
data ContParam = ContParam {-# UNPACK #-} !Double {-# UNPACK #-} !Double
  deriving (Show,Eq,Ord,Data,Typeable,Generic)

-- | We use 's' as default value which is same as R's default.
instance Default ContParam where
  def = s

instance Binary   ContParam
instance ToJSON   ContParam
instance FromJSON ContParam

-- | O(/n/·log /n/). Estimate the /k/th /q/-quantile of a sample /x/,
--   using the continuous sample method with the given parameters.
--
--   The following properties should hold, otherwise an error will be thrown.
--
--     * input sample must be nonempty
--
--     * the input does not contain @NaN@
--
--     * 0 ≤ k ≤ q
quantile :: G.Vector v Double
         => ContParam  -- ^ Parameters /α/ and /β/.
         -> Int        -- ^ /k/, the desired quantile.
         -> Int        -- ^ /q/, the number of quantiles.
         -> v Double   -- ^ /x/, the sample data.
         -> Double
quantile param q nQ xs
  | nQ < 2         = modErr "continuousBy" "At least 2 quantiles is needed"
  | badQ nQ q      = modErr "continuousBy" "Wrong quantile number"
  | G.any isNaN xs = modErr "continuousBy" "Sample contains NaNs"
  | otherwise      = estimateQuantile sortedXs pk
  where
    pk       = toPk param n q nQ
    sortedXs = psort xs $ floor pk + 1
    n        = G.length xs
{-# INLINABLE quantile #-}
{-# SPECIALIZE
    quantile :: ContParam -> Int -> Int -> U.Vector Double -> Double #-}
{-# SPECIALIZE
    quantile :: ContParam -> Int -> Int -> V.Vector Double -> Double #-}
{-# SPECIALIZE
    quantile :: ContParam -> Int -> Int -> S.Vector Double -> Double #-}

-- | O(/k·n/·log /n/). Estimate set of the /k/th /q/-quantile of a
--   sample /x/, using the continuous sample method with the given
--   parameters. This is faster than calling quantile repeatedly since
--   sample should be sorted only once
--
--   The following properties should hold, otherwise an error will be thrown.
--
--     * input sample must be nonempty
--
--     * the input does not contain @NaN@
--
--     * for every k in set of quantiles 0 ≤ k ≤ q
quantiles :: (G.Vector v Double, F.Foldable f, Functor f)
  => ContParam
  -> f Int
  -> Int
  -> v Double
  -> f Double
quantiles param qs nQ xs
  | nQ < 2             = modErr "quantiles" "At least 2 quantiles is needed"
  | F.any (badQ nQ) qs = modErr "quantiles" "Wrong quantile number"
  | G.any isNaN xs     = modErr "quantiles" "Sample contains NaNs"
  -- Doesn't matter what we put into empty container
  | fnull qs           = 0 <$ qs
  | otherwise          = fmap (estimateQuantile sortedXs) ks'
  where
    ks'      = fmap (\q -> toPk param n q nQ) qs
    sortedXs = psort xs $ floor (F.maximum ks') + 1
    n        = G.length xs
{-# INLINABLE quantiles #-}
{-# SPECIALIZE quantiles
      :: (Functor f, F.Foldable f) => ContParam -> f Int -> Int -> V.Vector Double -> f Double #-}
{-# SPECIALIZE quantiles
      :: (Functor f, F.Foldable f) => ContParam -> f Int -> Int -> U.Vector Double -> f Double #-}
{-# SPECIALIZE quantiles
      :: (Functor f, F.Foldable f) => ContParam -> f Int -> Int -> S.Vector Double -> f Double #-}

-- COMPAT
fnull :: F.Foldable f => f a -> Bool
#if !MIN_VERSION_base(4,8,0)
fnull = F.foldr (\_ _ -> False) True
#else
fnull = null
#endif

-- | O(/k·n/·log /n/). Same as quantiles but uses 'G.Vector' container
--   instead of 'Foldable' one.
quantilesVec :: (G.Vector v Double, G.Vector v Int)
  => ContParam
  -> v Int
  -> Int
  -> v Double
  -> v Double
quantilesVec param qs nQ xs
  | nQ < 2             = modErr "quantilesVec" "At least 2 quantiles is needed"
  | G.any (badQ nQ) qs = modErr "quantilesVec" "Wrong quantile number"
  | G.any isNaN xs     = modErr "quantilesVec" "Sample contains NaNs"
  | G.null qs          = G.empty
  | otherwise          = G.map (estimateQuantile sortedXs) ks'
  where
    ks'      = G.map (\q -> toPk param n q nQ) qs
    sortedXs = psort xs $ floor (G.maximum ks') + 1
    n        = G.length xs
{-# INLINABLE quantilesVec #-}
{-# SPECIALIZE quantilesVec
      :: ContParam -> V.Vector Int -> Int -> V.Vector Double -> V.Vector Double #-}
{-# SPECIALIZE quantilesVec
      :: ContParam -> U.Vector Int -> Int -> U.Vector Double -> U.Vector Double #-}
{-# SPECIALIZE quantilesVec
      :: ContParam -> S.Vector Int -> Int -> S.Vector Double -> S.Vector Double #-}

-- Returns True if quantile number is out of range
badQ :: Int -> Int -> Bool
badQ nQ q = q < 0 || q > nQ

-- Obtain k from equation for p_k [Hyndman1996] p.363.  Note that
-- equation defines p_k for integer k but we calculate it as real
-- value and will use fractional part for linear interpolation. This
-- is correct since equation is linear.
toPk
  :: ContParam
  -> Int        -- ^ /n/ number of elements
  -> Int        -- ^ /k/, the desired quantile.
  -> Int        -- ^ /q/, the number of quantiles.
  -> Double
toPk (ContParam a b) (fromIntegral -> n) q nQ
  = a + p * (n + 1 - a - b)
  where
    p = fromIntegral q / fromIntegral nQ

-- Estimate quantile for given k (including fractional part)
estimateQuantile :: G.Vector v Double => v Double -> Double -> Double
{-# INLINE estimateQuantile #-}
estimateQuantile sortedXs k'
  = (1-g) * item (k-1) + g * item k
  where
    (k,g) = properFraction k'
    item  = (sortedXs !) . clamp
    --
    clamp = max 0 . min (n - 1)
    n     = G.length sortedXs

psort :: G.Vector v Double => v Double -> Int -> v Double
psort xs k = partialSort (max 0 $ min (G.length xs - 1) k) xs
{-# INLINE psort #-}

-- | California Department of Public Works definition, /α/=0, /β/=1.
-- Gives a linear interpolation of the empirical CDF.  This
-- corresponds to method 4 in R and Mathematica.
cadpw :: ContParam
cadpw = ContParam 0 1

-- | Hazen's definition, /α/=0.5, /β/=0.5.  This is claimed to be
-- popular among hydrologists.  This corresponds to method 5 in R and
-- Mathematica.
hazen :: ContParam
hazen = ContParam 0.5 0.5

-- | Definition used by the SPSS statistics application, with /α/=0,
-- /β/=0 (also known as Weibull's definition).  This corresponds to
-- method 6 in R and Mathematica.
spss :: ContParam
spss = ContParam 0 0

-- | Definition used by the S statistics application, with /α/=1,
-- /β/=1.  The interpolation points divide the sample range into @n-1@
-- intervals.  This corresponds to method 7 in R and Mathematica and
-- is default in R.
s :: ContParam
s = ContParam 1 1

-- | Median unbiased definition, /α/=1\/3, /β/=1\/3. The resulting
-- quantile estimates are approximately median unbiased regardless of
-- the distribution of /x/.  This corresponds to method 8 in R and
-- Mathematica.
medianUnbiased :: ContParam
medianUnbiased = ContParam third third
    where third = 1/3

-- | Normal unbiased definition, /α/=3\/8, /β/=3\/8.  An approximately
-- unbiased estimate if the empirical distribution approximates the
-- normal distribution.  This corresponds to method 9 in R and
-- Mathematica.
normalUnbiased :: ContParam
normalUnbiased = ContParam ta ta
    where ta = 3/8

modErr :: String -> String -> a
modErr f err = error $ "Statistics.Quantile." ++ f ++ ": " ++ err

----------------------------------------------------------------
-- Specializations
----------------------------------------------------------------

-- | O(/n/·log /n/) Estimate median of sample
median :: G.Vector v Double
       => ContParam  -- ^ Parameters /α/ and /β/.
       -> v Double   -- ^ /x/, the sample data.
       -> Double
{-# INLINE median #-}
median p = quantile p 1 2

-- | O(/n/·log /n/). Estimate the range between /q/-quantiles 1 and
-- /q/-1 of a sample /x/, using the continuous sample method with the
-- given parameters.
--
-- For instance, the interquartile range (IQR) can be estimated as
-- follows:
--
-- > midspread medianUnbiased 4 (U.fromList [1,1,2,2,3])
-- > ==> 1.333333
midspread :: G.Vector v Double =>
             ContParam  -- ^ Parameters /α/ and /β/.
          -> Int        -- ^ /q/, the number of quantiles.
          -> v Double   -- ^ /x/, the sample data.
          -> Double
midspread param k x
  | G.any isNaN x = modErr "midspread" "Sample contains NaNs"
  | k <= 0        = modErr "midspread" "Nonpositive number of quantiles"
  | otherwise     = let Pair x1 x2 = quantiles param (Pair 1 (k-1)) k x
                    in  x2 - x1
{-# INLINABLE  midspread #-}
{-# SPECIALIZE midspread :: ContParam -> Int -> U.Vector Double -> Double #-}
{-# SPECIALIZE midspread :: ContParam -> Int -> V.Vector Double -> Double #-}
{-# SPECIALIZE midspread :: ContParam -> Int -> S.Vector Double -> Double #-}

data Pair a = Pair !a !a
  deriving (Functor, F.Foldable)

-- | O(/n/·log /n/). Estimate the median absolute deviation (MAD) of a
--   sample /x/ using 'continuousBy'. It's robust estimate of
--   variability in sample and defined as:
--
--   \[
--   MAD = \operatorname{median}(| X_i - \operatorname{median}(X) |)
--   \]
mad :: G.Vector v Double
    => ContParam  -- ^ Parameters /α/ and /β/.
    -> v Double   -- ^ /x/, the sample data.
    -> Double
mad p xs
  = median p $ G.map (abs . subtract med) xs
  where
    med = median p xs
{-# INLINABLE  mad #-}
{-# SPECIALIZE mad :: ContParam -> U.Vector Double -> Double #-}
{-# SPECIALIZE mad :: ContParam -> V.Vector Double -> Double #-}
{-# SPECIALIZE mad :: ContParam -> S.Vector Double -> Double #-}

----------------------------------------------------------------
-- Deprecated
----------------------------------------------------------------

continuousBy :: G.Vector v Double =>
                ContParam  -- ^ Parameters /α/ and /β/.
             -> Int        -- ^ /k/, the desired quantile.
             -> Int        -- ^ /q/, the number of quantiles.
             -> v Double   -- ^ /x/, the sample data.
             -> Double
continuousBy = quantile
{-# DEPRECATED continuousBy "Use quantile instead" #-}

-- $references
--
-- * Weisstein, E.W. Quantile. /MathWorld/.
--   <http://mathworld.wolfram.com/Quantile.html>
--
-- * [Hyndman1996] Hyndman, R.J.; Fan, Y. (1996) Sample quantiles in statistical
--   packages. /American Statistician/
--   50(4):361–365. <http://www.jstor.org/stable/2684934>


{-# LANGUAGE DataKinds #-}

import Math.Tensor
import Math.Tensor.LorentzGenerator
import Math.Tensor.Examples.Gravity
import Math.Tensor.Examples.Gravity.DiffeoSymEqns
import Math.Tensor.Examples.Gravity.Schwarzschild

import Math.Tensor.Internal.LinearAlgebra

import Numeric.LinearAlgebra (rank)
import Numeric.LinearAlgebra.Data (toLists, fromLists, size, cmap, (===), Matrix, Z)
import Data.List (nub, nubBy, findIndex, sort)

import Data.Maybe (mapMaybe)

import Data.Int (Int64)
import Data.Ratio

import Control.Parallel.Strategies (parList, runEvalIO, rdeepseq)

import qualified Data.IntMap.Strict as I

main = do

  let ans0 = fromListT6 [((Empty, Empty, Empty, Empty, Empty, Empty), AnsVar $ I.singleton 1 (SField 1))] :: ATens 0 0 0 0 0 0 AnsVarR
--  let (eta4,eps4,ans4) = mkAnsatzTensorFastAbs 4 symList4 areaList4 :: (AnsatzForestEta, AnsatzForestEpsilon, ATens 1 0 0 0 0 0 AnsVarR)
  let ans4 = ZeroTensor :: ATens 1 0 0 0 0 0 AnsVarR
  let (eta6,eps6,ans6) = mkAnsatzTensorFastAbs 6 symList6 areaList6 :: (AnsatzForestEta, AnsatzForestEpsilon, ATens 1 0 1 0 0 0 AnsVarR)
  let (eta8,eps8,ans8) = mkAnsatzTensorFastAbs 8 symList8 areaList8 :: (AnsatzForestEta, AnsatzForestEpsilon, ATens 2 0 0 0 0 0 AnsVarR)
  let (eta10_1,eps10_1,ans10_1) = mkAnsatzTensorFastAbs 10 symList10_1 areaList10_1 :: (AnsatzForestEta, AnsatzForestEpsilon, ATens 2 0 0 0 2 0 AnsVarR)
  let (eta10_2,eps10_2,ans10_2) = mkAnsatzTensorFastAbs 10 symList10_2 areaList10_2 :: (AnsatzForestEta, AnsatzForestEpsilon, ATens 2 0 1 0 0 0 AnsVarR)

  let r0    = tensorRank6' ans0
  let r4    = tensorRank6' ans4
  let r6    = tensorRank6' ans6
  let r8    = tensorRank6' ans8
  let r10_1 = tensorRank6' ans10_1
  let r10_2 = tensorRank6' ans10_2

  let ans0'    = ans0
  let ans4'    = shiftLabels6 r0 ans4
  let ans6'    = shiftLabels6 (r0 + r4) ans6
  let ans8'    = shiftLabels6 (r0 + r4 + r6) ans8
  let ans10_1' = shiftLabels6 (r0 + r4 + r6 + r8) ans10_1
  let ans10_2' = shiftLabels6 (r0 + r4 + r6 + r8 + r10_1) ans10_2

  let ansaetze = ans0' &.&>
                 ans4' &.&>
                 ans6' &.&>
                 ans8' &.&>
                 ans10_1' &.&>
                 (singletonTList6 ans10_2')

  putStrLn $ "dimension of ansatz space   : " ++ show (tensorRank6 ansaetze)

  let two = SField (2 :: Rational)

  let e1 = eqn1 ans0' ans4'
  let e2 = eqn3 ans6'
  let e3 = eqn1A ans4' (two &. ans8')
  let e4 = eqn1AI ans6' ans10_2'
  let e5 = eqn2Aa ans6' (two &. ans10_1')
  let e6 = eqn3A ans6' ans10_2'

  let system = e1 &.&>
               e2 &.&>
               e3 &.&>
               e4 &.&>
               e5 &.&>
               (singletonTList6 e6)

  putStrLn $ "rank of system              : " ++ show (tensorRank6 system)

  let solution = redefineVarsSystem6 $ solveSystem6 system ansaetze

  case solution of
    (t1 `AppendTList6` (
     t2 `AppendTList6` (
     t3 `AppendTList6` (
     t4 `AppendTList6` (
     t5 `AppendTList6` (
     t6 `AppendTList6` (
     EmptyTList6)))))))
      -> do
          let Just ans0'' = tryAsATens ans0 t1
          let Just ans4'' = tryAsATens ans4 t2
          let Just ans6'' = tryAsATens ans6 t3
          let Just ans8'' = tryAsATens ans8 t4
          let Just ans10_1'' = tryAsATens ans10_1 t5
          let Just ans10_2'' = tryAsATens ans10_2 t6

          let e1' = removeZeros6 $ eqn1 ans0'' ans4''
          let e2' = removeZeros6 $ eqn3 ans6''
          let e3' = removeZeros6 $ eqn1A ans4'' (two &. ans8'')
          let e4' = removeZeros6 $ eqn1AI ans6'' ans10_2''
          let e5' = removeZeros6 $ eqn2Aa ans6'' (two &. ans10_1'')
          let e6' = removeZeros6 $ eqn3A ans6'' ans10_2''

          putStrLn $ "dimension of solution space : " ++ show (tensorRank6 solution)
          putStrLn ""
          putStrLn "Equations on solution space :"

          print e1'
          print e2'
          print e3'
          print e4'
          print e5'
          print e6'

          let kin' = ans10_2'' &- (contrATens3 (0,0) $ contrATens3 (1,1) $ interI2 &* ans10_1'')
          let kin = kin' &+ (tensorTrans1 (0,1) kin')

          putStrLn ""
          putStrLn $ "number of parameters in kinetic part of EOM : " ++ show (tensorRank6' kin)
          putStrLn $ "number of parameters in mass part of EOM    : " ++ show (tensorRank6' ans8'')

          let n1 = removeZeros6 $ (delta3A &* ans4'') &+ (contrATens1 (1,0) $ interArea &* ans4'') &+ (contrATens1 (0,0) $ contrATens1 (1,1) $ two &. flatArea &* interArea &* ans8'')

          let n2_1 = contrATens1 (0,0) $ contrATens1 (1,1) $ contrATens2 (0,0) $ flatArea &* interArea &* interJ2 &* ans10_2''
          let n2_2 = contrATens1 (0,0) $ contrATens1 (2,1) $ contrATens2 (0,0) $ flatArea &* interArea &* interJ2 &* ans10_2''
          let n2_3 = contrATens1 (0,0) $ contrATens1 (1,1) $ flatArea &* interArea &* ans10_1''

          let n2 = removeZeros6 $ cyclicSymATens5 [0,1,2] $ n2_1 &+ n2_2 &- (two &. n2_3)

          putStrLn ""
          putStrLn "Noether identity 1 :"
          print n1

          putStrLn ""
          putStrLn "Noether identity 2 :"
          print n2

  return ()


module Numeric.RBM where

import Control.Monad
import Control.Monad.IO.Class
import Math.Probable hiding (vector)
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Data
import Numeric.LinearAlgebra.HMatrix (norm_2)

data RBM = RBM
  { weights  :: !(Matrix Double)
  , weightsT :: !(Matrix Double)
  } deriving (Eq, Show)

pp :: RBM -> IO ()
pp (RBM w _) = disp 2 w

randMat :: Int -> Int -> RandT IO Double -> IO (Matrix Double)
randMat n p = mwc . fmap fromLists . replicateM n . listOf p

new :: Int -- ^ nb of visible units
    -> Int -- ^ nb of hidden units
    -> IO RBM
new nvisible nhidden = do
  w      <- randMat nvisible nhidden (normal 0 0.1)
  let w'  = asColumn (konst 0 nvisible) ||| w
      w'' = asRow (konst 0 (nhidden + 1)) === w'
  return $ RBM w'' (tr w'')

trainOnce :: Matrix Double -- ^ training data
          -> Double        -- ^ learning rate
          -> RBM
          -> IO RBM
trainOnce trainingData learningRate (RBM w wt) = do
  let dat            = asColumn (konst 1 n) ||| trainingData
      pos_hidden_act = dat <> w
      pos_hidden_p   = cmap logistic pos_hidden_act

  pos_hidden_st <- do
    randomM <- randMat n p (uniform 0 1)
    return $ step (pos_hidden_p - randomM)

  let pos_assocs = tr dat <> pos_hidden_p

      neg_visible_act = pos_hidden_st <> wt
      neg_visible_p   = cmap logistic neg_visible_act
      n'              = rows neg_visible_p
      neg_visible_p'  = asColumn (konst 1 n') ||| dropColumns 1 neg_visible_p
      neg_hidden_act  = neg_visible_p' <> w
      neg_hidden_p    = cmap logistic neg_hidden_act
      neg_assocs      = tr neg_visible_p' <> neg_hidden_p

      new_w           = (w +) . cmap (\x -> learningRate * x / fromIntegral n) $
                          pos_assocs - neg_assocs

      err             = norm_2 (dat - neg_visible_p')

  putStrLn $ "Error: " ++ show err
  return $ RBM new_w (tr new_w)

  where n = rows trainingData
        p = cols w

logistic :: Double -> Double
logistic x = 1 / (1 + exp (negate x))

trainN :: Matrix Double -- ^ training data
       -> Double        -- ^ learning rate
       -> Int           -- ^ epochs
       -> RBM
       -> IO RBM
trainN dat learningRate epochs = go epochs

  where go 0 rbm = return rbm
        go n rbm = do
          rbm' <- trainOnce dat learningRate rbm
          go (n - 1) rbm'

runVisible :: RBM
           -> Matrix Double      -- ^ sample rows of values for visible units
           -> IO (Matrix Double) -- ^ values of hidden units for the given visible units
runVisible (RBM w wt) inputData = do
  let dat        = asColumn (konst 1 n) ||| inputData
      hidden_act = dat <> w
      hidden_p   = cmap logistic hidden_act

  hidden_st <- do
    randomM <- randMat n p (uniform 0 1)
    return $ step (hidden_p - randomM)

  return $ dropColumns 1 hidden_st

  where n = rows inputData
        p = cols w

runHidden :: RBM
          -> Matrix Double      -- ^ sample rows of hidden units values
          -> IO (Matrix Double) -- ^ values of visible units for the given hidden ones
runHidden (RBM w wt) hiddenData = do
  let dat         = asColumn (konst 1 n) ||| hiddenData
      visible_act = dat <> wt
      visible_p   = cmap logistic visible_act

  visible_st <- do
    randomM <- randMat n p (uniform 0 1)
    return $ step (visible_p - randomM)

  return $ dropColumns 1 visible_st

  where n = rows hiddenData
        p = cols wt

daydream :: RBM
         -> Int                -- ^ number of samples we want to extract
         -> IO [Vector Double] -- ^ each row is a sample of the visible units
                               --   that the RBM has "daydreamed"
daydream (RBM w wt) nsamples = do
  sample0 <- fmap (fromList . (1:) . tail) . mwc $ listOf n (uniform 0 1)

  go sample0 nsamples

  where n = rows w
        p = cols w

        go :: Vector Double -> Int -> IO [Vector Double]
        go _v 0 = return []
        go v i = do
          let hidden_act = asRow v <> w
              hidden_p   = cmap logistic hidden_act

          hidden_st <- do
            randomV <- fmap vector . mwc $ listOf p (uniform 0 1)
            return $ step (head (toRows hidden_p) - randomV)

          let hidden_st' = asRow . vector . (1:) . tail . toList $ hidden_st
              visible_act = hidden_st' <> wt
              visible_p   = cmap logistic visible_act

          visible_st <- do
            randomV <- fmap vector . mwc $ listOf n (uniform 0 1)
            return $ step (head (toRows visible_p) - randomV)

          let res = vector . tail . toList $ visible_st
          fmap (res:) $ go visible_st (i - 1)

-- EXAMPLE

-- data given in:
-- http://blog.echen.me/2011/07/18/introduction-to-restricted-boltzmann-machines/

-- Harry Potter, Avatar, LOTR, Gladiator, Titanic, Glitter
samples :: Matrix Double
samples = fromLists
  [ [ 1, 1, 1, 0, 0, 0 ]
  , [ 1, 0, 1, 0, 0, 0 ]
  , [ 1, 1, 1, 0, 0, 0 ]
  , [ 0, 0, 1, 1, 1, 0 ]
  , [ 0, 0, 1, 1, 0, 0 ]
  , [ 0, 0, 1, 1, 1, 0 ]
  ]

new_user :: Matrix Double
new_user = fromLists [ [ 0, 0, 0, 1, 1, 0 ] ]

hiddens :: Matrix Double
hiddens = fromLists
  [ [0, 0]
  , [0, 1]
  , [1, 0]
  , [1, 1]
  ]

test :: Double -> Int -> IO ()
test rate k = do
  rbm <- new 6 2

  pp rbm
  trained <- trainN samples rate k rbm
  pp trained
  putStrLn "-------"
  putStrLn "VISIBLE"
  res <- runVisible trained new_user
  disp 2 res
  putStrLn "-------"
  putStrLn "HIDDEN"
  res' <- runHidden trained hiddens
  disp 2 res'

  putStrLn "-------"
  putStrLn "DAYDREAM"
  daydream trained 10 >>= mapM_ print

main :: IO ()
main = test 0.1 10000


{-# LANGUAGE UnicodeSyntax, DataKinds, TypeOperators, KindSignatures,
             TypeInType, GADTs, MultiParamTypeClasses, FunctionalDependencies,
             TypeFamilies, AllowAmbiguousTypes, FlexibleInstances,
             UndecidableInstances, InstanceSigs, TypeApplications,
             ScopedTypeVariables, EmptyCase, FlexibleContexts,
             RankNTypes, LambdaCase
#-}

-- Implementation of density matrix interpretation of quantum computation
module LinTrans where

import Prelim

import qualified Data.Complex as C
import Data.Complex (Complex(..))
import Control.Monad.State.Lazy
import Data.Singletons
import Data.Maybe (fromJust)
--import Data.Tuple hiding (swap)
import Unsafe.Coerce
import Data.Constraint (Dict(..))
import Data.List ( (\\) )
import qualified Debug.Trace as Tr

-- I don't like the show instance for ℂ
newtype ℂ = ℂ (Complex Double)

liftC :: (Complex Double -> Complex Double) -> ℂ -> ℂ
liftC f (ℂ c) = ℂ $ f c

i :: ℂ
i = ℂ $ 0 :+ 1
conjugate = liftC C.conjugate
instance Num ℂ where
  ℂ m + ℂ n = ℂ $ m + n
  ℂ m - ℂ n = ℂ $ m - n
  ℂ m * ℂ n = ℂ $ m * n
  abs = liftC abs
  signum = liftC signum
  fromInteger = ℂ . fromInteger
instance Fractional ℂ where
  fromRational = ℂ . fromRational
  ℂ m / ℂ n    = ℂ $ m / n
instance Floating ℂ where
  pi    = ℂ $ pi
  exp   = liftC exp
  log   = liftC log
  sin   = liftC sin
  cos   = liftC cos
  asin  = liftC asin
  acos  = liftC acos
  atan  = liftC atan
  sinh  = liftC sinh
  cosh  = liftC cosh
  asinh = liftC asinh
  acosh = liftC acosh
  atanh = liftC atanh
instance Show ℂ where
  show (ℂ (α :+ β)) =   if β == 0 then show α
                        else if α == 0 then show β ++ "i"
                        else show α ++ " + " ++ show β ++ "i"

-- A matrix of dimension m × n is a function from
-- pairs of nats bounded by m and n respectively.
type Matrix m n = (BNat m,BNat n) -> ℂ
type Squared n = Matrix (Two `RaiseTo` n) (Two `RaiseTo` n)

matrix :: forall m n. (SingI m, SingI n) => [ℂ] -> Matrix m n
matrix ls (i,j) = ls !! (n * fromIntegral i + fromIntegral j)
  where
    n = toInt (sing :: Sing n)

--------------
-- Printing --
--------------
rows :: forall m n. (SingI m, SingI n) => Matrix m n -> [[ℂ]]
rows mat = fmap f $ allBNat (sing :: Sing m)
  where
    f i = fmap (\j -> mat (i,j)) $ allBNat (sing :: Sing n)

showRows :: (SingI m, SingI n) => Matrix m n -> [String]
showRows mat = show <$> rows mat

instance (SingI m, SingI n) => Show (Matrix m n) where
  show mat = unlines $ showRows mat

---------------------------------------
-- Primitive matrices and operatoins --
---------------------------------------

ident :: Matrix m m
ident (i,j) = if i==j then 1 else 0

transpose :: Matrix m n -> Matrix n m
transpose mat (i,j) = conjugate $ mat (j,i)

{-
-- if i : BNat m then raise m n = i : BNat (S m)
-- so pruneRows mat is just the first m rows of mat
pruneRows :: Sing m -> Matrix (S m) n -> Matrix m n
pruneRows m mat (i,j) = mat (raise m i,j)

pruneCols :: Sing n -> Matrix m (S n) -> Matrix m n
pruneCols n mat (i,j) = mat (i,raise n j)
-}

sumTo :: Sing n -> BNat n -> (BNat n -> ℂ) -> ℂ
sumTo _ (BNat SZ)     _ = 0
sumTo n (BNat (SS m)) f = f m' + sumTo n m' f
  where
    m' = case succLtTrans m n of Dict -> BNat m -- m < S m < n

dot :: Sing n -> Matrix (S Z) n -> Matrix n (S Z) -> ℂ
dot n v1 v2 = sumTo n (maxBNat n) (\x -> v1 (0,x) * v2(x,0))

row :: BNat m -> Matrix m n -> Matrix (S Z) n
row i mat (_,j) = mat (i,j)

col :: BNat n -> Matrix m n -> Matrix m (S Z)
col j mat (i,_) = mat (i,j)

plus :: Matrix m n -> Matrix m n -> Matrix m n
plus mat1 mat2 (i,j) = mat1(i,j) + mat2(i,j)

mult :: forall m n p. SingI n => Matrix m n -> Matrix n p -> Matrix m p
mult mat1 mat2 (i,j) =
--  Tr.trace ("Composing for intermediate " ++ show (sing :: Sing n)) $
    dot sing (row i mat1) (col j mat2)

kron :: forall m1 m2 n1 n2. (SingI m1, SingI n1, SingI m2, SingI n2)
     => Matrix m1 n1 -> Matrix m2 n2 -> Matrix (m1 `Times` m2) (n1 `Times` n2)
kron mat1 mat2 (BNat i,BNat j) =
--    Tr.trace ("Computing ⊗ for " ++ show (m1, n1) ++ "×" ++ show (m2, n2)) $
    mat1 (bNat i1, bNat j1) * mat2 (bNat i2,bNat j2)
  where
    i1 = i `divSNat` m2
    i2 = i `modSNat` m2
    j1 = j `divSNat` n2
    j2 = j `modSNat` n2
    m2 = (sing :: Sing m2)
    n2 = (sing :: Sing n2)

trace :: forall m. SingI m => Matrix m m -> ℂ
trace mat = foldr f 0 ls
  where
    ls = allBNat (sing :: Sing m)
    f i x = mat (i,i) + x

----------------------
-- Quantum Matrices --
----------------------

square :: forall n. SingI n => [ℂ] -> Squared n
square ls = withSingI (two `raiseToSNat` (sing :: Sing n)) $ matrix ls

density0 :: Matrix Two Two
density0 = matrix [1,0
                  ,0,0]

density1 :: Matrix Two Two
density1 = matrix [0,0
                  ,0,1]

hadamard = square @One $ [1/sqrt 2, 1/sqrt 2
                         ,1/sqrt 2, -1/sqrt 2]

pauliX = square @One $ [0,1,1,0]
pauliY = square @One $ [0,-i,i,0]
pauliZ = square @One $ [1,0,0,-1]

cnotD  = square @Two $ [1,0,0,0
                       ,0,1,0,0
                       ,0,0,0,1
                       ,0,0,1,0]

newD :: Bool -> Squared One
newD True  = density1
newD False = density0

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

----------------------
-- Density Matrices --
----------------------

data Density where
  Density :: Sing n -> Squared n -> Density
instance Show Density where
  show (Density n m) = withSingI (two `raiseToSNat` n) $ show m

size :: forall m. SingI m => Matrix m m -> Int
size _ = toInt (sing :: Sing m)
logsizeD :: Density -> Int
logsizeD (Density n mat) = toInt n
sizeD :: Density -> Int
sizeD ρ = 2 ^ logsizeD ρ

transposeD :: Density -> Density
transposeD (Density m ρ) = Density m $ transpose ρ

multD :: Density -> Density -> Density
multD = undefined

kronD :: Density -> Density -> Density
kronD (Density m mat1) (Density n mat2) =
      withSingI (raiseToSNat two n) $
      withSingI (raiseToSNat two m) $
      withSingI (raiseToSNat two (plusSNat m n)) $
      case raiseToPlus two m n of {Dict ->
        Density (m `plusSNat` n) $ mat1 `kron` mat2
      }

identD :: Int -> Density
identD n = case fromIntegral @_ @(SomeSing Nat) n of SomeSing n' -> Density n' ident

nil :: Int -> Density
nil n = case fromIntegral n of
          SomeSing n' -> withSingI (two `raiseToSNat` n') $
                         Density n' . matrix $ fromIntegral <$> repeat 0

------------------
-- Permutations --
------------------

-- encode and decode modulo a key k
encode :: Int -> [Int] -> Int
encode k [] = 0
encode k (i : ls) = i + k * encode k ls

decode :: Int -> Int -> [Int]
decode k i | i == 0 = []
decode k i | otherwise = i `mod` k : decode k (i `div` k)

fromAssocList :: [(Int,Int)] -> Int -> Int
fromAssocList []           = id
fromAssocList ((a,b) : ls) = assocFun a b . fromAssocList ls
  where
    assocFun a b i = if i == a then b else i

-- Check if two lists are equal up to the permutation specified, i.e.
-- if permuting ls1 by f is equal to ls2
isPermutation :: Eq a => (Int -> Int) -> [a] -> [a] -> Bool
isPermutation f ls1 ls2 = all (\i -> ls1 !! i == ls2 !! f i) [0..len]
  where
    len = length ls1

swapFun :: forall n. SingI n => (Int -> Int) -> Squared n
swapFun f (i,j) = if isPermutation f (decode n' i') (decode n' j') then 1 else 0
  where
    n' = toInt (sing :: Sing n)
    i' = toInt i
    j' = toInt j

-- swap takes a list of qubit/bit variables and produces a matrix that permutes
-- those qubits.
-- i.e. swap ls |φ_0 ⋯ φ_n⟩ = |φ_f(0) ⋯ ρ_f(n)⟩
-- where f = fromAssocList ls
swap :: forall n. SingI n => [Int] -> Squared n
swap ls = swapFun @n (fromAssocList $ zip [0..] ls)

-------------------
-- Density Monad --
-------------------

-- A density monad is a nondeterminism state monad:
-- Density -> [Density]
-- An element op of type DensityMonad corresponds to the superoperator
-- \ρ -> ∑ (op ρ)
type DensityMonad = StateT Density []

newM :: Bool -> DensityMonad Int
newM b = do
    ρ <- get
    put $ ρ `kronD` (Density one $ newD b)
    return $ sizeD ρ

-- runQ applies the superoperator to the identity density matrix of size 1.
runQ :: DensityMonad a -> [(a,Density)]
runQ m = runStateT m (Density SZ ident)

-- getDensity combines the result of runQ into a single density matrix
getDensity :: DensityMonad a -> Density
getDensity m = foldr f (nil n) ls
  where
    ls = runQ m
    n = logsizeD . snd $ head ls
    f (_,ρ) ρ0 = ρ + ρ0

-----------------
-- Application --
-----------------

superM :: forall m n. Matrix (Two `RaiseTo` m) (Two `RaiseTo` n) -> DensityMonad ()
superM mat = do Density _ ρ ← get
                put . Density _ $ mat `mult` ρ `mult` transpose mat
  where
    matD = Density (sing :: Sing m) mat

applyUnitaryM :: forall m. SingI m => Squared m -> [Int] -> DensityMonad ()
applyUnitaryM u ls = do superM @m σ -- moves the relevant qubits to the front of the density matrix
                        superM @m u -- applies operation
                        superM @m $ transpose σ -- moves the qubits back
  where
    σ = swap @m ls

applyMatrixM :: Matrix (Two `RaiseTo` m) (Two `RaiseTo` n)
             -> [Int] -> DensityMonad ()
applyMatrixM mat ls = undefined

measM :: Int -> DensityMonad Bool
measM i = do
    ρ <- get
    (b,ρ') <- lift [ (False, applyMatrixD @One density0 [i] ρ)
                   , (True,  applyMatrixD @One density1 [i] ρ) ]
    put ρ'
    return b

------------------
-- Num instance --
------------------

instance (SingI m, SingI n) => Num (Matrix m n) where
  (+) mat1 mat2 (i,j) = mat1(i,j) + mat2(i,j)
  (-) mat1 mat2 (i,j) = mat1(i,j) - mat2(i,j)
  (*) mat1 mat2 (i,j) = mat1(i,j) * mat2(i,j)
  abs mat (i,j) = abs $ mat(i,j)
  signum mat = undefined -- trace mat
  fromInteger n = matrix . repeat $ fromIntegral n
instance Num Density where
  Density n1 mat1 + Density n2 mat2 =
    withSingI (two `raiseToSNat` n1) $
    withSingI (two `raiseToSNat` n2) $
    case eqSNat n1 n2 of
      Left Dict -> Density n1 $ mat1 + mat2
      Right _   -> error $ "Cannot add mismatched matrices "
                            ++ show mat1 ++ " and " ++ show mat2
  _ - _ = undefined
  _ * _ = undefined
  abs _ = undefined
  signum = undefined
  fromInteger n = undefined

instance Show (DensityMonad a) where
  show = show . getDensity

-- Debugging help
m0 = square @Two [1,0,0,0
                 ,0,1,0,0
                 ,0,0,0,1
                 ,0,0,1,0]
rho = Density three $ \case
  (0,0) -> -0.5
  (0,2) ->  0.5
  (3,0) -> -0.5
  (3,2) ->  0.5
  (_,_) -> 0

{-
-- produces an association list
listToAssoc :: forall n. SingI n => [Int] -> [(BNat n, BNat n)]
listToAssoc ls = (fromIntegral <$> [0..n]) `zip` (fromIntegral <$> pad ls)
  where
    n = maxBNat (sing :: Sing n)

pad :: [Int] -> [Int]
pad ls = ls ++ ([0..] \\ ls)

mkPermutation :: [(BNat n,BNat n)] -> BNat n -> BNat n
mkPermutation ls i =
  case lookup i ls of
    Just i' -> i'
    Nothing -> error $ "Permutation " ++ show ls
                        ++ " not well defined at index " ++ show i

permute :: forall m. SingI m => [(BNat m, BNat m)] -> Matrix m m -> Matrix m m
permute ls mat (i,j) = mat (mkPermutation ls i, mkPermutation ls j)

expand :: forall m n. (SingI m, SingI n)
       => [(BNat (m `Times` n), BNat (m `Times` n))]
       -> Matrix m m -> Matrix (m `Times` n) (m `Times` n)
expand ls mat = Tr.trace ("Expanding..." ++ show mat) $
                withSingI (timesSNat (sing :: Sing m) (sing :: Sing n)) $
                  permute ls $ mat `kron` ident @n

applyMatrix :: forall m p. (SingI m, SingI p, (p < m) ~ 'False)
            => Squared m -> [Int] -> Squared p -> Squared p
applyMatrix mat ls ρ =
  Tr.trace ("Applying " ++ show ls) $
  case raiseToPlus two m minus of {Dict ->
  case plusMinus m p of {Dict ->
  withSingI (raiseToSNat two m) $
  withSingI (raiseToSNat two minus) $
  withSingI (raiseToSNat two p) $
    let mat1 = expand @(Two `RaiseTo` m) @(Two `RaiseTo` (p `Minus` m))
                       (swap <$> ls') (transpose mat)
        mat2 = expand @(Two `RaiseTo` m) @(Two `RaiseTo` (p `Minus` m)) ls' mat
        res  = Tr.trace "Middle of application" $
               mat1 `mult` ρ `mult` mat2
    in mat1 `seq` mat2 `seq` res `seq` Tr.trace "End of application" res
  }}
  where
    ls' :: [(BNat (Two `RaiseTo` p), BNat (Two `RaiseTo` p))]
    ls' = withSingI (raiseToSNat two p) $ listToAssoc ls
    m = (sing :: Sing m)
    p = (sing :: Sing p)
    minus :: Sing (p `Minus` m)
    minus = minusSNat p m

applyMatrixD :: forall m. SingI m => Squared m -> [Int] -> Density -> Density
applyMatrixD mat perm (Density (p :: Sing p) ρ) =
  Tr.trace ("Applying " ++ mat' ++ " to qubits " ++ show perm ++ " of \n" ++ show (Density p ρ)) $
  case compareSNat p (sing :: Sing m) of
    -- p < m
    Left Dict -> error "Cannot apply matrix to state: matrix too large"
    -- m <= p
    Right Dict ->
      let res = withSingI p $ Density p $ applyMatrix @m @p mat perm ρ
      in  res `seq` Tr.trace "Done!" res
  where
    mat' = withSingI (two `raiseToSNat` (sing :: Sing m)) $ show mat

applyUnitaryM :: forall m. SingI m => Squared m -> [Int] -> DensityMonad ()
applyUnitaryM mat ls = get >>= \ρ -> put $ applyMatrixD @m mat ls ρ

measM :: Int -> DensityMonad Bool
measM i = do
    ρ <- get
    (b,ρ') <- lift [ (False, applyMatrixD @One density0 [i] ρ)
                   , (True,  applyMatrixD @One density1 [i] ρ) ]
    put ρ'
    return b

-}


{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_GHC -Wcompat #-}
{-# OPTIONS_GHC -Wincomplete-record-updates #-}
{-# OPTIONS_GHC -Wincomplete-uni-patterns #-}
{-# OPTIONS_GHC -Wredundant-constraints #-}

module CloudSunlight
  ( cloudSunlightExperiment,
  )
where

import Control.Monad
import qualified Data.ByteString.Lazy as BL
import Data.Csv
import qualified Data.List as L
import qualified Data.Vector.Unboxed as V
import qualified Statistics.Regression as R
import qualified Statistics.Sample as S
import System.Random.MWC
import qualified System.Random.MWC.Distributions as D

meanToa :: Double
meanToa = 340.0

toa :: GenIO -> IO Double
toa = D.normal meanToa 30.0

-- | magnitude of aerosol affect on sunlight, between 1.0 and 0.0 (1.0 = max effect)
aerosolAffectSunlight :: Double
aerosolAffectSunlight = 1.0

-- | factor that cloud reduces sunlight, if there is cloud
cloudAffectSunlight :: Double
cloudAffectSunlight = 0.6

sunlight ::
  -- | non-dim measure of aerosol
  Double ->
  -- | cloud
  Bool ->
  -- | U_sun (i.e. TOA)
  Double ->
  -- | W/m2
  Double
sunlight a c toa' = if c then clearSky * cloudAffectSunlight else clearSky
  where
    clearSky = toa' * (1.0 - aerosolAffectSunlight * a)

cloud :: Double -> Double -> Bool
cloud a u
  | (a + u) > 1.0 = True
  | otherwise = False

aerosol :: Double -> Double
aerosol u = u

genTuple :: GenIO -> IO (Double, String, Double)
genTuple gen = do
  uAero <- uniform gen
  uCld <- uniform gen
  uSun <- toa gen
  let a = aerosol uAero
  let c = cloud a uCld
  let s = sunlight a c uSun
  return (a, if c then "Cloudy" else "Clear", s)

getCloud :: (a, b, c) -> b
getCloud (_a, c, _s) = c

getSun :: (a, b, c) -> c
getSun (_a, _c, s) = s

getAerosol :: (a, b, c) -> a
getAerosol (a, _c, _s) = a

calcRegression :: V.Vector Double -> Double -> Double
calcRegression ols x = x * ols V.! 0 + ols V.! 1

calcStats ::
  [(Double, String, Double)] ->
  [(Double, String, Double)] ->
  (V.Vector Double, V.Vector Double, Double, Double)
calcStats cloud' clear =
  ( olsCloud,
    olsClear,
    avgRegress olsCloud - avgRegress olsClear,
    naiveCloud - naiveClear
  )
  where
    calcAvgs xs = (as, ols, S.mean ss)
      where
        ss = V.fromList (fmap getSun xs)
        as = V.fromList (fmap getAerosol xs)
        ols = fst . R.olsRegress [as] $ ss
    (aCloud, olsCloud, naiveCloud) = calcAvgs cloud'
    (aClear, olsClear, naiveClear) = calcAvgs clear
    totAs = aCloud <> aClear
    avgRegress ols = S.mean $ V.map (calcRegression ols) totAs

cloudSunlightExperiment :: Int -> IO (Double, Double, Double)
cloudSunlightExperiment n = do
  gen <- create
  xs <- replicateM n (genTuple gen)
  let (cloud', clear) = L.partition ((== "Cloudy") . getCloud) xs
  let (olsCloud, olsClear, estDiff, naiveDiff) = calcStats cloud' clear
  let regPoints = [0, 0.1 .. 1]
  f "dat/cloudy.dat" cloud'
  f "dat/clear.dat" clear
  write
    "aerosol\tsunlight\n"
    "dat/olsCloudy.dat"
    (zip regPoints (fmap (calcRegression olsCloud) regPoints))
  write
    "aerosol\tsunlight\n"
    "dat/olsClear.dat"
    (zip regPoints (fmap (calcRegression olsClear) regPoints))
  return (estDiff, naiveDiff, (0.5 * meanToa * (cloudAffectSunlight - 1.0)))
  where
    f = write "aerosol\tcloud\tsunlight\n"
    write header' fpath =
      BL.writeFile fpath . (header' <>)
        . encodeWith
          (defaultEncodeOptions {encDelimiter = 9})


{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}

module ScalarWave where

import Control.Exception (assert)
import Control.Monad.Loops
import qualified Data.Vector.Unboxed as U
import Data.VectorSpace
import Dual
import Equations
import Foreign
import qualified Numeric.LinearAlgebra as L
import SSSFN
import System.CPUTime

default (Int)

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

bndw :: Floating a => Coord a -> a
bndw (Coord u v) = swamp * cos (pi / 4 * swk * t) * cos (pi / 4 * swk * r)
  where
    t = v + u
    r = v - u
    swk = 0.1
    swamp = 0.1

guess1 :: Floating a => Coord a -> State a
guess1 c = State 1 1 0

guess :: (L.Container L.Vector a, Floating a) => Grid (State a)
guess = gsample (\(u, v) -> guess1 (Coord u v))

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

eqnInt1 :: Num a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnInt1 eqn c s r2s r1sr r1su r1sv su sv suv suu svv =
  case eqn of
    EqnG01 -> eqnG01 s r2s r1sr su sv suv
    EqnG22 -> eqnG22 s r2s r1sr su sv suv
    EqnSW2 -> eqnSW2 s r2s r1sr su sv suv
    EqnG00 -> eqnG00 c s r2s r1su su suu
    EqnG11 -> eqnG11 c s r2s r1sv sv svv

eqnOrig1 :: Floating a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnOrig1 eqn c s r2s r1sr r1su r1sv su sv suv suu svv =
  case eqn of
    EqnG01 -> q s - 1
    EqnG22 -> g s - 1
    EqnSW2 -> w s - bndw c
    EqnG00 -> 0
    EqnG11 -> 0

-- i == 0, u == -1
eqnBndU1 :: (Eq a, Floating a) => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnBndU1 eqn c s r2s r1sr r1su r1sv su sv suv suu svv =
  assert (u c == -1) $
    case eqn of
      EqnG01 -> eqnG11 c s r2s r1sv sv svv
      EqnG22 -> g s - 1
      EqnSW2 -> w s - bndw c
      EqnG00 -> 0
      EqnG11 -> 0

-- j == 0, v == -1
eqnBndV1 :: (Eq a, Floating a) => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnBndV1 eqn c s r2s r1sr r1su r1sv su sv suv suu svv =
  assert (v c == -1) $
    case eqn of
      EqnG01 -> eqnG00 c s r2s r1su su suu
      EqnG22 -> g s - 1
      EqnSW2 -> w s - bndw c
      EqnG00 -> 0
      EqnG11 -> 0

eqnEndU1 :: Floating a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnEndU1 eqn c s r2s r1sr r1su r1sv su sv suv suu svv =
  case eqn of
    EqnG01 -> q s - 1
    EqnG22 -> g s - 1
    EqnSW2 -> w s - bndw c
    EqnG00 -> 0
    EqnG11 -> 0

eqnEndV1 :: Floating a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnEndV1 eqn c s r2s r1sr r1su r1sv su sv suv suu svv =
  case eqn of
    EqnG01 -> q s - 1
    EqnG22 -> g s - 1
    EqnSW2 -> w s - bndw c
    EqnG00 -> 0
    EqnG11 -> 0

equation1 :: (Eq a, Floating a) => Coord Int -> Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> State a -> a
equation1 (Coord i j) =
  if
    | i == 0 && j == 0 -> eqnOrig1
    | i == 0 && j == np - 1 -> eqnEndU1
    | i == np - 1 && j == 0 -> eqnEndV1
    | i == 0 -> eqnBndU1
    | j == 0 -> eqnBndV1
    | True -> eqnInt1

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

g2s :: Storable a => Grid a -> State (Grid a)
g2s (Grid xs) =
  let [g, q, w] = L.takesV (replicate 3 (np * np)) xs
   in State (Grid g) (Grid q) (Grid w)

s2g :: Storable a => State (Grid a) -> Grid a
s2g (State (Grid g) (Grid q) (Grid w)) = Grid (L.vjoin [g, q, w])

eqn1 ::
  forall a.
  (L.Field a, RealFloat a, U.Unbox a) =>
  Eqn ->
  Coord (Grid a) ->
  State (Grid a) ->
  Grid a
eqn1 eqn c s =
  Grid $
    L.fromList
      [ eqnRow i j
        | i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    r2s = (r2project #>) <$> s
    r1sr = (r1derivR #>) <$> s
    r1su = (r1derivU #>) <$> s
    r1sv = (r1derivV #>) <$> s
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    eqnRow :: Int -> Int -> a
    eqnRow i j =
      let at :: Grid a -> a
          at x = x ! (i, j)
       in equation1
            (Coord i j)
            eqn
            (at <$> c)
            (at <$> s)
            (at <$> r2s)
            (at <$> r1sr)
            (at <$> r1su)
            (at <$> r1sv)
            (at <$> su)
            (at <$> sv)
            (at <$> suv)
            (at <$> suu)
            (at <$> svv)

eqns ::
  forall a.
  (L.Field a, RealFloat a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Grid a
eqns c s =
  Grid $
    L.fromList
      [ eqnRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    r2s = (r2project #>) <$> s
    r1sr = (r1derivR #>) <$> s
    r1su = (r1derivU #>) <$> s
    r1sv = (r1derivV #>) <$> s
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    eqnRow :: (Int, Int, Eqn) -> a
    eqnRow (i, j, eqn) =
      let at :: Grid a -> a
          at x = x ! (i, j)
       in equation1
            (Coord i j)
            eqn
            (at <$> c)
            (at <$> s)
            (at <$> r2s)
            (at <$> r1sr)
            (at <$> r1su)
            (at <$> r1sv)
            (at <$> su)
            (at <$> sv)
            (at <$> suv)
            (at <$> suu)
            (at <$> svv)

eqns5 ::
  forall a.
  (L.Field a, RealFloat a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Grid a
eqns5 c s =
  Grid $
    L.fromList
      [ eqnRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2, EqnG00, EqnG11],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    r2s = (r2project #>) <$> s
    r1sr = (r1derivR #>) <$> s
    r1su = (r1derivU #>) <$> s
    r1sv = (r1derivV #>) <$> s
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    eqnRow :: (Int, Int, Eqn) -> a
    eqnRow (i, j, eqn) =
      let at :: Grid a -> a
          at x = x ! (i, j)
       in equation1
            (Coord i j)
            eqn
            (at <$> c)
            (at <$> s)
            (at <$> r2s)
            (at <$> r1sr)
            (at <$> r1su)
            (at <$> r1sv)
            (at <$> su)
            (at <$> sv)
            (at <$> suv)
            (at <$> suu)
            (at <$> svv)

op ::
  forall a.
  (L.Container L.Vector a, L.Field a, RealFloat a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Op a
op (Coord u v) s =
  Op $
    L.fromRows
      [ jacRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    r2s = (r2project #>) <$> s
    r1sr = (r1derivR #>) <$> s
    r1su = (r1derivU #>) <$> s
    r1sv = (r1derivV #>) <$> s
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    jacRow :: (Int, Int, Eqn) -> L.Vector a
    jacRow (i, j, eqn) =
      L.vjoin $
        [ let con :: Grid a -> Dual a
              der :: Grid a -> Dual a
              con x = Dual (x ! (i, j)) 0
              der x = Dual (x ! (i, j)) 1
              scon :: State (Grid a) -> State (Dual a)
              sder :: State (Grid a) -> State (Dual a)
              scon (State g q w) = State (con g) (con q) (con w)
              sder (State g q w) = case var of
                VarG -> State (der g) (con q) (con w)
                VarQ -> State (con g) (der q) (con w)
                VarW -> State (con g) (con q) (der w)
           in ( ( dual $
                    equation1
                      (Coord i j)
                      eqn
                      (Coord (con u) (con v))
                      (sder s)
                      (scon r2s)
                      (scon r1sr)
                      (scon r1su)
                      (scon r1sv)
                      (scon su)
                      (scon sv)
                      (scon suv)
                      (scon suu)
                      (scon svv)
                )
                  `L.scale` (getOp delta L.! (i * np + j))
              )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (sder r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r2project L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (sder r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r1derivR L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (sder r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r1derivU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (sder r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r1derivV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (sder su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (sder sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (sder suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (sder suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (sder svv)
                          )
                            `L.scale` (getOp derivVV L.! (i * np + j))
                        )
          | var <- [VarG, VarQ, VarW]
        ]

op5 ::
  forall a.
  (L.Container L.Vector a, L.Field a, RealFloat a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Op a
op5 (Coord u v) s =
  Op $
    L.fromRows
      [ jacRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2, EqnG00, EqnG11],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    r2s = (r2project #>) <$> s
    r1sr = (r1derivR #>) <$> s
    r1su = (r1derivU #>) <$> s
    r1sv = (r1derivV #>) <$> s
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    jacRow :: (Int, Int, Eqn) -> L.Vector a
    jacRow (i, j, eqn) =
      L.vjoin $
        [ let con :: Grid a -> Dual a
              der :: Grid a -> Dual a
              con x = Dual (x ! (i, j)) 0
              der x = Dual (x ! (i, j)) 1
              scon :: State (Grid a) -> State (Dual a)
              sder :: State (Grid a) -> State (Dual a)
              scon (State g q w) = State (con g) (con q) (con w)
              sder (State g q w) = case var of
                VarG -> State (der g) (con q) (con w)
                VarQ -> State (con g) (der q) (con w)
                VarW -> State (con g) (con q) (der w)
           in ( ( dual $
                    equation1
                      (Coord i j)
                      eqn
                      (Coord (con u) (con v))
                      (sder s)
                      (scon r2s)
                      (scon r1sr)
                      (scon r1su)
                      (scon r1sv)
                      (scon su)
                      (scon sv)
                      (scon suv)
                      (scon suu)
                      (scon svv)
                )
                  `L.scale` (getOp delta L.! (i * np + j))
              )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (sder r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r2project L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (sder r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r1derivR L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (sder r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r1derivU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (sder r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp r1derivV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (sder su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (sder sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (sder suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (sder suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon r2s)
                                (scon r1sr)
                                (scon r1su)
                                (scon r1sv)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (sder svv)
                          )
                            `L.scale` (getOp derivVV L.! (i * np + j))
                        )
          | var <- [VarG, VarQ, VarW]
        ]

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

getTime :: IO Double
getTime = do
  t <- getCPUTime
  return $ fromInteger t / 1.0e+12

main0 :: IO ()
main0 =
  do
    putStrLn $ show $ gmaxabs $ coordU @Double
    putStrLn $ show $ gmaxabs $ coordV @Double
    putStrLn $ show $ gmaxabs $ coordR @Double
    putStrLn $ show $ gmaxabs $ coordR1 @Double
    putStrLn $ show $ omaxabs $ zproject @Double
    putStrLn $ show $ omaxabs $ r2project @Double
    putStrLn $ show $ omaxabs $ r1derivR @Double
    putStrLn $ show $ omaxabs $ r1derivU @Double
    putStrLn $ show $ omaxabs $ r1derivV @Double
    putStrLn $ show $ omaxabs $ derivU @Double
    putStrLn $ show $ omaxabs $ derivV @Double
    putStrLn $ show $ omaxabs $ derivUV @Double
    putStrLn $ show $ omaxabs $ derivUU @Double
    putStrLn $ show $ omaxabs $ derivVV @Double

main :: IO ()
main =
  do
    putStrLn "Spherically Symmetric Scalar Field in double-Null coordinates"
    let coords = Coord coordU (coordV @Double)
    let ini = State (gmap g guess) (gmap q guess) (gmap w guess)
    vars0 <- return (s2g ini)
    let done (n, vars) =
          let res = eqns coords (g2s vars)
              maxabsres = gmaxabs res
           in isNaN maxabsres
                || isInfinite maxabsres
                || maxabsres < 1.0e-12
                || n >= 5
    let iter (n, vars) = do
          let res = eqns coords (g2s vars)
          let res5 = eqns5 coords (g2s vars)
          t <- getTime
          putStrLn $
            ("iter " ++ show n)
              ++ ("   time " ++ show (round t) ++ " sec")
              ++ ("   |res| " ++ show (gmaxabs res))
              ++ ("   |res5| " ++ show (gmaxabs res5))
          let jac = op coords (g2s vars)
          let dvars = solve jac res
          return (n + 1, vars ^-^ dvars)
    let done5 (n, vars) =
          let res5 = eqns5 coords (g2s vars)
              maxabsres5 = gmaxabs res5
           in isNaN maxabsres5
                || isInfinite maxabsres5
                || maxabsres5 < 1.0e-12
                || n >= 5
    let iter5 (n, vars) = do
          let res = eqns coords (g2s vars)
          let res5 = eqns5 coords (g2s vars)
          t <- getTime
          putStrLn $
            ("iter " ++ show n)
              ++ ("   time " ++ show (round t) ++ " sec")
              ++ ("   |res| " ++ show (gmaxabs res))
              ++ ("   |res5| " ++ show (gmaxabs res5))
          let jac5 = op5 coords (g2s vars)
          let dvars = solveLS jac5 res5
          return (n + 1, vars ^-^ dvars)
    -- (n, vars) <- iterateUntilM done iter (0, vars0)
    (n, vars) <- iterateUntilM done5 iter5 (0, vars0)
    let res = eqns coords (g2s vars)
    let res5 = eqns5 coords (g2s vars)
    t <- getTime
    putStrLn $
      ("iter " ++ show n)
        ++ ("   time " ++ show (round t) ++ " sec")
        ++ ("   |res| " ++ show (gmaxabs res))
        ++ ("   |res5| " ++ show (gmaxabs res5))
    putStrLn $ "var.g " ++ show (gmap approx $ g $ g2s vars)
    putStrLn $ "var.q " ++ show (gmap approx $ q $ g2s vars)
    putStrLn $ "var.w " ++ show (gmap approx $ w $ g2s vars)
    putStrLn $ "res.G01 " ++ show (gmap approx $ eqn1 EqnG01 coords (g2s vars))
    putStrLn $ "res.G22 " ++ show (gmap approx $ eqn1 EqnG22 coords (g2s vars))
    putStrLn $ "res.SW2 " ++ show (gmap approx $ eqn1 EqnSW2 coords (g2s vars))
    putStrLn $ "res.G00 " ++ show (gmap approx $ eqn1 EqnG00 coords (g2s vars))
    putStrLn $ "res.G11 " ++ show (gmap approx $ eqn1 EqnG11 coords (g2s vars))
    putStrLn $ "S var.g " ++ show (gmap approx $ inv vandermonde #> (g $ g2s vars))
    putStrLn $ "S var.q " ++ show (gmap approx $ inv vandermonde #> (q $ g2s vars))
    putStrLn $ "S var.w " ++ show (gmap approx $ inv vandermonde #> (w $ g2s vars))
    putStrLn $ "S res.G01 " ++ show (gmap approx $ inv vandermonde #> (eqn1 EqnG01 coords (g2s vars)))
    putStrLn $ "S res.G22 " ++ show (gmap approx $ inv vandermonde #> (eqn1 EqnG22 coords (g2s vars)))
    putStrLn $ "S res.SW2 " ++ show (gmap approx $ inv vandermonde #> (eqn1 EqnSW2 coords (g2s vars)))
    putStrLn $ "S res.G00 " ++ show (gmap approx $ inv vandermonde #> (eqn1 EqnG00 coords (g2s vars)))
    putStrLn $ "S res.G11 " ++ show (gmap approx $ inv vandermonde #> (eqn1 EqnG11 coords (g2s vars)))
    putStrLn "Done."
  where
    approx :: RealFrac a => a -> a
    approx x = fromInteger (round (1.0e+10 * x)) / 1.0e+10


{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE BangPatterns #-}

module Main where

import Options.Applicative
import Data.Semigroup ((<>))
import Chart.Heatmap
import Control.Lens
import Control.Monad (forM_)
import Data.List (maximumBy)
import Data.Ord (comparing)
import Data.Colour.SRGB (sRGB, RGB(..))
import Data.Colour.RGBSpace.HSL (hsl)
import Data.Complex (magnitude, realPart, imagPart, Complex(..))
import Graphics.Rendering.Chart.Easy
import Graphics.Rendering.Chart.Backend.Cairo (toFile, FileOptions(..), FileFormat(..))
import Data.Array.CArray.Base (unsafeForeignPtrToCArray, toForeignPtr, CArray(..), ixmapWithIndP)
import Data.Fixed (mod')
import Math.FFT (dftRC)
import Numeric.Signal (analytic_signal, auto_correlation, deriv, Filterable)
import Data.Maybe (catMaybes)
import Sound.MIDI.File
import qualified Data.Vector.Storable as VS
import qualified Data.Vector as V
import qualified Sound.File.Sndfile as SF
import qualified Sound.File.Sndfile.Buffer.Vector as SFV
import qualified Sound.MIDI.File.Save as MS
import qualified Sound.MIDI.File.Event as ME
import qualified Sound.MIDI.Message.Channel as MC
import qualified Sound.MIDI.Message.Channel.Voice as MCV
import qualified Sound.MIDI.Message.Channel.Mode as MCM
import qualified Data.EventList.Relative.TimeBody as E

type Frequency = Double
type Amplitude = Double
type Step = Double

maxFrequency :: Double
maxFrequency = 7000

data NoteLetter = A
        | Ais
        | B
        | C
        | Cis
        | D
        | Dis
        | E
        | F
        | Fis
        | G
        | Gis
        deriving (Show, Eq)

type Octave = Int

data Note = Note NoteLetter Octave deriving (Show, Eq)

getBaseFreq :: NoteLetter -> Frequency
getBaseFreq A   = 27.50
getBaseFreq Ais = 29.14
getBaseFreq B   = 30.87
getBaseFreq C   = 16.35
getBaseFreq Cis = 17.32
getBaseFreq D   = 18.35
getBaseFreq Dis = 19.45
getBaseFreq E   = 20.60
getBaseFreq F   = 21.83
getBaseFreq Fis = 23.12
getBaseFreq G   = 24.50
getBaseFreq Gis = 25.96

getFreq :: Note -> Frequency
getFreq (Note letter n) = getBaseFreq letter * (fromIntegral n + 1)

getNote :: Frequency -> Maybe Note
getNote freq
  | divisor A `mod'` 1 < 0.02 || divisor A `mod'` 1 > 0.98 = Just
    (Note A (findN freq A))
  | divisor Ais `mod'` 1 < 0.02 || divisor Ais `mod'` 1 > 0.98 = Just
    (Note Ais (findN freq Ais))
  | divisor B `mod'` 1 < 0.02 || divisor B `mod'` 1 > 0.98 = Just
    (Note B (findN freq B))
  | divisor C `mod'` 1 < 0.02 || divisor C `mod'` 1 > 0.98 = Just
    (Note C (findN freq C))
  | divisor Cis `mod'` 1 < 0.02 || divisor Cis `mod'` 1 > 0.98 = Just
    (Note Cis (findN freq Cis))
  | divisor D `mod'` 1 < 0.02 || divisor D `mod'` 1 > 0.98 = Just
    (Note D (findN freq D))
  | divisor Dis `mod'` 1 < 0.02 || divisor Dis `mod'` 1 > 0.98 = Just
    (Note Dis (findN freq Dis))
  | divisor E `mod'` 1 < 0.02 || divisor E `mod'` 1 > 0.98 = Just
    (Note E (findN freq E))
  | divisor F `mod'` 1 < 0.02 || divisor F `mod'` 1 > 0.98 = Just
    (Note F (findN freq F))
  | divisor Fis `mod'` 1 < 0.02 || divisor Fis `mod'` 1 > 0.98 = Just
    (Note Fis (findN freq Fis))
  | divisor G `mod'` 1 < 0.02 || divisor G `mod'` 1 > 0.98 = Just
    (Note G (findN freq G))
  | divisor Gis `mod'` 1 < 0.02 || divisor Gis `mod'` 1 > 0.98 = Just
    (Note Gis (findN freq Gis))
  | otherwise = Nothing
 where
  findN f n = round $ logBase 2 (f / getBaseFreq n) :: Int
  divisor note = logBase 2 (freq / getBaseFreq note)

-- samples = N = number of samples
-- n = sample number
hannWindow :: Integer -> Integer -> Double
hannWindow samples n =
  0.5 * (1 - cos (2 * pi * fromIntegral n / (fromIntegral samples - 1)))

sinFunc :: Double -> Double -> Double -> Double -> Double
sinFunc samples cycle amplitude p =
  -amplitude * sin (2 * pi * cycle * p / samples)

cosFunc :: Double -> Double -> Double -> Double -> Double
cosFunc samples cycle amplitude p =
  amplitude * cos (2 * pi * cycle * p / samples)

performFFT :: VS.Vector Double -> Int -> IO (VS.Vector (Complex Double))
performFFT vector intSamples = do
  let (vec, x1, x2) = VS.unsafeToForeignPtr $ VS.take intSamples vector
  carray <- unsafeForeignPtrToCArray vec (x1, x2)
  let fft           = dftRC carray
  let (len, fftptr) = toForeignPtr fft
  return (VS.unsafeFromForeignPtr fftptr 0 len)

-- |Filter an array of maxima by partioning fftabs into bins of 100 elements, then finding the
-- largest value in that bin.
-- peaks are the list of indexes of the maxima.
filterMaxima :: V.Vector Int -> V.Vector Double -> V.Vector Int
filterMaxima peaks fftabs
  | V.null peaks = V.empty
  | otherwise = V.filter (> -1) $ V.map
    (\i ->
      let
        section = V.map
          (\x -> (x, fftabs V.! x))
          (V.filter (\x -> x < V.length fftabs && x >= i && x < i + binSize)
                    peaks
          )
      in  if V.null section
            then -1
            else fst $ V.maximumBy (comparing snd)
            -- list of (index, magnitude)
                                                   section
    )
    (V.fromList [0, binSize .. V.last peaks])
  where binSize = 100

-- | Find the maxima from fftabs by finding all points where the left and right points
-- are less than the current point.
findMaxima :: V.Vector Double -> V.Vector Int
findMaxima fftabs = V.filter
  (\x -> (diffs' V.! x) > 0 && (diffs' V.! (x + 1)) < 0)
  (V.fromList [1 .. (V.length diffs' - 2)])
 where
  diff f x = f x - f (x - 1)
  diffs' = V.map (diff (fftabs V.!)) (V.fromList [1 .. (V.length fftabs - 1)])

-- | Find maxima by finding points where that point is the maximum in it's neighbourhood.
findMaxima2 :: Filterable a => Int -> V.Vector a -> V.Vector (Int, a)
findMaxima2 radius v =
  V.filter (\(i, x) -> x == V.maximum (neighbours i))
    $ V.imap (\i x -> (i, x)) v
  where neighbours x0 = V.ifilter (\i _ -> abs (x0 - i) <= radius) v

findSignificantNotesForSection
  :: Int -> Int -> V.Vector (Complex Double) -> [(Frequency, Double)]
findSignificantNotesForSection sampleRate samples fftSection = fmap
  (\cycle ->
    ( fromIntegral cycle * fromIntegral sampleRate / fromIntegral samples
    , fftabs V.! cycle
    )
  )
  maxima'
 where
  fftabs  = V.map magnitude fftSection :: V.Vector Double
  peaks   = findMaxima fftabs
  peaks'  = filterMaxima peaks fftabs
  maxima  = V.toList peaks' -- tricky: the cycles returned is the cycles+1
  maxima' = filter (\value -> fftabs V.! value > 20) maxima -- Find 'large' maxima

midiPitch :: Note -> Int
midiPitch (Note letter n) =
  round $ 12 * logBase 2 (((2 ^ n) * getBaseFreq letter) / 440) + 69

-- | Find notes that standout in each stft section.
findSignificantNotes
  :: Int -> Int -> [(Step, [Complex Double])] -> String -> IO ()
findSignificantNotes sampleRate samples stft filePrefix = do
  notes <-
    mapM
      (\(step, fft) -> do
        let freqs =
              findSignificantNotesForSection sampleRate samples (V.fromList fft)
        let notes = fmap (\(f, _) -> getNote f)
                         (filter (\(f, _) -> f < maxFrequency) freqs)
        return $ map (\note -> (step, note)) $ catMaybes notes
      )
      stft :: IO [[(Step, Note)]]

  print $ map (map (\(s, note) -> (s, note, midiPitch note))) notes

  let
    sections = map
      (foldl
        (\acc (step, note) -> E.append
          acc
          (E.singleton
            (toElapsedTime 0)
            (ME.MIDIEvent
              (MC.Cons
                (MC.toChannel 1)
                (MC.Voice
                  (MCV.NoteOn (MC.toPitch (midiPitch note)) (MC.toVelocity 20))
                )
              )
            )
          )
        )
        E.empty
      )
      notes

  let
    t =
      E.append
          (foldl (\acc section -> E.append acc (E.delay 100 section))
                 E.empty
                 sections
          )
        $ E.singleton
            (toElapsedTime 100)
            (ME.MIDIEvent (MC.Cons (MC.toChannel 1) (MC.Mode MCM.AllNotesOff)))

  let midi = Cons Serial (Ticks 120) [t] :: T

  MS.toFile (filePrefix ++ ".midi") midi

--    toFile def{_fo_format=SVG} "diff.svg" $ do
--        layout_title .= "First derivative"
--        plot (points "diff" $ zip fftx (VS.toList fftabs))

--    print $ V.map (\x -> (x, fftabs VS.! x)) peaks'
    --toFile def{_fo_format=SVG} "diff2.svg" $ do
    --    layout_title .= "Second derivative"
    --    plot (points "diff" $ zip fftx diffs'')

    --let sinAmplitudes = take 500 $ fmap (\(cycle, _) -> (fromIntegral cycle, fftimg   VS.! cycle / fromIntegral (samples `div` 2))) maxima :: [(Double, Double)]
    --let cosAmplitudes = take 500 $ fmap (\(cycle, _) -> (fromIntegral cycle, fftreals VS.! cycle / fromIntegral (samples `div` 2))) maxima :: [(Double, Double)]

    --print "sin amps"
    --print sinAmplitudes
    --print "cos amps"
    --print cosAmplitudes

    --toFile def{_fo_format=SVG} "sine2.svg" $ do
    --    layout_title .= "curve"
    --    plot (line "data" [zip x y])
    --    plot
    --        (line "fit"
    --        [zip
    --            x
    --            (fmap
    --                (\p -> foldr (\(cycle, amplitude) total -> total + sinFunc (fromIntegral samples) cycle amplitude p) 0 sinAmplitudes +
    --                    foldr (\(cycle, amplitude) total -> total + cosFunc (fromIntegral samples) cycle amplitude p) 0 cosAmplitudes) x)])

    --let signal = VS.take samples $ analytic_signal vector
    --let signalReal = VS.toList $ VS.map realPart signal
    --let signalImag = VS.toList $ VS.map imagPart signal
    ----print $ VS.length signal
    --toFile def{_fo_format=SVG} "signal.svg" $ do
    --    layout_title .= "Hibert transform derivative"
    --    plot (line "Hilbert transform" [zip x (VS.toList $ VS.map imagPart signal)])
    --    plot (line "Original signal" [zip x intList])
    --    plot (line "Analytic signal" [zip x $ VS.toList (VS.map magnitude signal)])

autoCorrelate :: Filterable a => VS.Vector a -> Int -> IO ()
autoCorrelate vector sampleRate = do
  let result = auto_correlation 1000 vector
  let y      = map realToFrac $ VS.toList result :: [Double]
  --let values = [zip (map (\x -> x / fromIntegral sampleRate) [1..]) y] :: [[(Double, Double)]]
  let values = [zip [1 ..] y] :: [[(Double, Double)]]
  toFile def { _fo_format = SVG } "autocorrelation.svg" $ do
    layout_title .= "auto correlation"
    plot (line "" values)

  let localMaxima = findMaxima2 10 (VS.convert result)
  print $ V.map (\(i, x) -> (i, realToFrac x)) localMaxima
  print $ V.imap
    (\j (i, _) -> if j == 0 then j else i - fst (localMaxima V.! (j - 1)))
    localMaxima
  print $ fromIntegral sampleRate / fromIntegral
    (V.maximum
      (V.imap
        (\j (i, _) -> if j == 0 then j else i - fst (localMaxima V.! (j - 1)))
        localMaxima
      )
    )

graphSignal :: Filterable a => VS.Vector a -> Int -> IO ()
graphSignal vector sampleRate = do
  let intList = map realToFrac $ VS.toList vector :: [Double]
  let y       = intList
  let x       = map (\x -> x / fromIntegral sampleRate) [1 .. 200] :: [Double]
  toFile def { _fo_format = SVG } "signal.svg" $ do
    layout_x_axis . laxis_title .= "Elapsed time (s)"
    layout_title .= "signal curve"
    plot (line "" [zip x y])

graphFFTSignal :: VS.Vector (Complex Double) -> Int -> Int -> IO ()
graphFFTSignal fftvec samples sampleRate = do
  let fftreals = VS.toList $ VS.map realPart fftvec :: [Double]
  let fftimg   = VS.toList $ VS.map imagPart fftvec :: [Double]
  let fftabs   = VS.toList $ VS.map magnitude fftvec :: [Double]
  let fftx =
        map (\x -> fromIntegral (x * sampleRate) / fromIntegral samples)
            [0 .. 1000] :: [Double]
  toFile def { _fo_format = SVG } "fft.svg" $ do
    layout_title .= "fft curve"
    layout_x_axis . laxis_title .= "Frequency (Hz)"
    plot (points "abs" $ zip fftx fftabs)
    plot (points "real" $ zip fftx fftreals)
    plot (points "imaginary" $ zip fftx fftimg)

data Options = Options {
    sourceFile :: String,
    stepWidth :: Int
}

args :: ParserInfo Options
args = info
  (Options <$> strOption (long "src" <> metavar "SOURCE") <*> option
    auto
    (long "step-width" <> metavar "STEP_WIDTH" <> value 10000 <> showDefault)
  )
  (fullDesc <> progDesc "Desc" <> header "test")

main :: IO ()
main = do
  opts         <- execParser args
  (i, content) <-
    SF.readFile (sourceFile opts) :: IO (SF.Info, Maybe (SFV.Buffer Double))
  let sampleRate = SF.samplerate i
  let samples    = SF.frames i * SF.channels i
  print samples
  print i
  let v = fmap SFV.fromBuffer content
  case v of
    Nothing     -> print "Nothing!"
    Just vector -> do
      graphSignal vector sampleRate

      -- Run FFT on whole dataset and plot the result
      fftvec <- performFFT vector samples
      graphFFTSignal fftvec samples sampleRate
      findSignificantNotes sampleRate samples [(0, VS.toList fftvec)] "fft"

      -- get auto correlation
      autoCorrelate vector sampleRate

      -- plot histogram
      let (vec, x1, x2) = VS.unsafeToForeignPtr vector
      carray <- unsafeForeignPtrToCArray vec (x1, x2)
      --let stepWidth = fromIntegral samples :: Integer
      let stepWidth = 5000 :: Integer
      print $ "Step width: " ++ show stepWidth
      let cutSamples =
            (samples `div` fromIntegral stepWidth) * fromIntegral stepWidth :: Int

      -- windows of size stepWidth
      let stft =
            fmap (stftSection carray cutSamples stepWidth)
                 [0, stepWidth `div` 2 .. fromIntegral cutSamples] :: [ ( Step
                , [Complex Double]
                )
              ]

      findSignificantNotes sampleRate cutSamples stft "stft"

      -- (window start, fft for window)
      let plots' =
            concatMap
              (spectrogramIntensity (fromIntegral cutSamples) sampleRate)
              stft :: [(Step, Frequency, Amplitude)]

      let stftData  = generateStftData (fromIntegral stepWidth) plots'
      let stftData' = filter (\((_, _), (_, _), z) -> z > 1) stftData
      toFile def { _fo_format = PNG } "stft.png" $ do
        layout_title .= "curve"
        plot
          (liftEC $ do
            plot_rect_title .= "title"
            plot_rect_style .= \amp ->
              let RGB a b c = hsl ((1 - amp) * 240) 1 (amp * 0.6)
              in  solidFillStyle (opaque $ sRGB a b c)
            plot_rect_values .= stftData'
          )

      return ()

-- convert to plottable set of values
generateStftData
  :: Step
  -> [(Step, Frequency, Amplitude)]
  -> [((Double, Double), (Double, Double), Amplitude)]
generateStftData windowSize = map (generateColumn windowSize)

-- column in spectrogram
generateColumn
  :: Step
  -> (Step, Frequency, Amplitude)
  -> ((Double, Double), (Double, Double), Amplitude)
generateColumn windowSize (step, freq, amp) =
  ((step, windowSize + step), (freq, freq + 1), amp)

spectrogramIntensity
  :: Double -> Int -> (Step, [Complex Double]) -> [(Step, Frequency, Amplitude)]
spectrogramIntensity totalSamples sampleRate (step, fft) =
  let rate = fromIntegral sampleRate / totalSamples
  in  filter
        (\(_, f, _) -> f < 4000)
        (map (\(cycle, x) -> (step, cycle * rate, magnitude (x + 0) ** 2))
             (zip [0 ..] fft)
        )

ftSignals :: Double -> [Complex Double] -> [Double] -> [Double]
ftSignals totalSamples fft = fmap
  (\x ->
    foldr
        (\(cycle, amplitude) total ->
          total + sinFunc totalSamples cycle amplitude x
        )
        0
        sinAmplitudes
      + foldr
          (\(cycle, amplitude) total ->
            total + cosFunc totalSamples cycle amplitude x
          )
          0
          cosAmplitudes
  )
 where
  significantSignals = filter (\x -> magnitude x > 5) fft
  cycles             = zip [0 ..] significantSignals
  sinAmplitudes =
    fmap
      (\(cycle, signal) ->
        (fromIntegral cycle, realPart signal / (totalSamples / 2.0))
      )
      cycles :: [(Double, Double)]
  cosAmplitudes =
    fmap
      (\(cycle, signal) ->
        (fromIntegral cycle, imagPart signal / (totalSamples / 2.0))
      )
      cycles :: [(Double, Double)]

-- | Apply window function to the set of samples
windowedSection
  :: CArray Int Double -> Int -> Integer -> Integer -> CArray Int Double
windowedSection data_ totalSamples n stepWidth = ixmapWithIndP
  (0, totalSamples)
  id
  (\_ x i' -> if i' < fromIntegral n || i' >= fromIntegral (n + stepWidth)
    then 0
    else x * hannWindow (fromIntegral stepWidth) (fromIntegral i' - n)
  )
  data_

-- | n = start of step sample number
stftSection
  :: CArray Int Double -> Int -> Integer -> Integer -> (Step, [Complex Double])
stftSection !completeSample totalSamples stepWidth n =
  (fromIntegral n, VS.toList $ VS.unsafeFromForeignPtr fftptr 0 len)
 where
  fft = dftRC (windowedSection completeSample totalSamples n stepWidth)
  (len, fftptr) = toForeignPtr fft