module AsteriskGaussianConvolution where

import           Control.Monad                  as M
import qualified Data.Array.Accelerate          as A
import           Data.Array.Accelerate.LLVM.PTX
import           Data.Array.Repa                as R
import           Data.Binary
import           Data.Complex
import           Data.List                      as L
import           Data.Vector.Storable           as VS
import           Data.Vector.Unboxed            as VU
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.GreensFunction
import           FokkerPlanck.Histogram
import           Foreign.CUDA.Driver            as CUDA
import           FourierMethod.FourierSeries2D
import           FourierPinwheel
import           Image.IO
import           STC                            hiding (convolve)
import           System.Directory
import           System.Environment
import           System.FilePath
import           Utils.Array
import           Utils.List
import           Utils.Parallel
import           Utils.Time

main = do
  args@(deviceIDsStr:numPointsStr:deltaStr:thresholdStr:numPointsReconStr:deltaReconStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:tauStr:numR2FreqStr:periodR2Str:phiFreqsStr:rhoFreqsStr:thetaFreqsStr:scaleFreqsStr:initDistStr:initScaleStr:histFilePath:stdR2Str:stdStr:numBatchR2Str:numBatchR2FreqsStr:numBatchOriStr:batchSizeStr:sStr:radiusStr:numThreadStr:_) <-
    getArgs
  let deviceIDs = read deviceIDsStr :: [Int]
      numPoints = read numPointsStr :: Int
      delta = read deltaStr :: Double
      threshold = read thresholdStr :: Double
      numPointsRecon = read numPointsReconStr :: Int
      deltaRecon = read deltaReconStr :: Double
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      tau = read tauStr :: Double
      numR2Freq = read numR2FreqStr :: Int
      periodR2 = read periodR2Str :: Double
      phiFreq = read phiFreqsStr :: Int
      phiFreqs = L.map fromIntegral [-phiFreq .. phiFreq]
      rhoFreq = read rhoFreqsStr :: Int
      rhoFreqs = L.map fromIntegral [-rhoFreq .. rhoFreq]
      thetaFreq = read thetaFreqsStr :: Int
      thetaFreqs = L.map fromIntegral [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Int
      scaleFreqs = L.map fromIntegral [-scaleFreq .. scaleFreq]
      initScale = read initScaleStr :: Double
      initDist = read initDistStr :: [(Double, Double, Double, Double)]
      initPoints = L.map (\(x, y, t, s) -> Point x y t s) initDist
      initSource = [L.head initPoints]
      initSink = [L.last initPoints]
      numThread = read numThreadStr :: Int
      folderPath = "output/test/AsteriskGaussianConvolution"
      maxScale = read maxScaleStr :: Double
      halfLogPeriod = log maxScale
      stdR2 = read stdR2Str :: Double
      std = read stdStr :: Double
      numBatchR2 = read numBatchR2Str :: Int
      numBatchR2Freqs = read numBatchR2FreqsStr :: Int
      numBatchOri = read numBatchOriStr :: Int
      batchSize = read batchSizeStr :: Int
      s = read sStr :: Double
      radius = read radiusStr :: Double
      periodEnvelope = periodR2^2 / 2
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  initialise []
  devs <- M.mapM device deviceIDs
  ctxs <- M.mapM (\dev -> CUDA.create dev []) devs
  ptxs <- M.mapM createTargetFromContext ctxs
  hist <-
    if flag
      then do
        printCurrentTime "Read coefficients..."
        decodeFile histFilePath
      else do
        printCurrentTime "Start computing coefficients..."
        sampleCartesian
          histFilePath
          folderPath
          ptxs
          numPoints
          periodEnvelope
          delta
          numOrientation
          initScale
          thetaSigma
          tau
          threshold
          s
          phiFreq
          rhoFreq
          thetaFreq
          scaleFreq
  printCurrentTime "Done"
  printCurrentTime "Start Convloution.."
  plan <-
    makePlan
      folderPath
      emptyPlan
      numPointsRecon
      numPointsRecon
      numR2Freq
      (2 * thetaFreq + 1)
      (2 * scaleFreq + 1)
      (2 * phiFreq + 1)
      (2 * rhoFreq + 1)
  let periodEnv = periodR2 * sqrt 2
      coefficients = getNormalizedHistogramArr hist
      harmonicsArray =
        createHarmonics
          numR2Freq
          phiFreq
          rhoFreq
          thetaFreq
          scaleFreq
          (-s)
          periodR2
          periodEnv
          coefficients
      asteriskGaussianVecs =
        asteriskGaussianFull
          numR2Freq
          thetaFreq
          scaleFreq
          (-s)
          periodR2
          periodEnv
          10
          10
      r2Freqs = getListFromNumber numR2Freq
      shiftFilter =
        fromListUnboxed
          (Z :. numR2Freq :. numR2Freq)
          [ L.foldl'
            (\b (Point x y theta scale) ->
               b +
               cis
                 (-(fromIntegral freqX * x + fromIntegral freqY * y) * 2 * pi /
                   periodR2))
            0
            initSource
          | freqY <- r2Freqs
          , freqX <- r2Freqs
          ]
      dftID = DFTPlanID DFT1DG [numR2Freq, numR2Freq] [0, 1]
      idftID = DFTPlanID IDFT1DG [numR2Freq, numR2Freq] [0, 1]
  shiftFilterF <-
    dftExecute plan dftID . VU.convert . toUnboxed . computeS . makeFilter2D $
    shiftFilter
  asteriskGaussianF <- dftExecuteBatchP plan dftID asteriskGaussianVecs
  filteredAsteriskGaussian <-
    fmap VS.concat .
    dftExecuteBatchP plan idftID . parMap rdeepseq (VS.zipWith (*) shiftFilterF) $
    asteriskGaussianF
  let numThetaFreq = 2 * thetaFreq + 1
      numRFreq = 2 * scaleFreq + 1
      sourceDist =
        FPArray
          numR2Freq
          numR2Freq
          (2 * scaleFreq + 1)
          (2 * thetaFreq + 1)
          (2 * rhoFreq + 1)
          (2 * phiFreq + 1) .
        parMap
          rdeepseq
          (\radialFreq ->
             VU.convert .
             toUnboxed .
             computeS .
             R.slice
               (fromUnboxed
                  (Z :. numRFreq :. numThetaFreq :. numR2Freq :. numR2Freq) .
                VS.convert $
                filteredAsteriskGaussian) $
             (Z :. radialFreq :. All :. All :. All)) $
        [0 .. 2 * scaleFreq]
  source <- convolve harmonicsArray sourceDist
  let sourceMat = A.transpose . toMatrixAcc $ source
  sourceR2' <-
    computeFourierSeriesR2StreamAcc
      ptxs
      (getFPArrayNumXFreq source)
      numPointsRecon
      (getFPArrayNumRFreq source * getFPArrayNumThetaFreq source)
      periodR2
      delta
      numBatchR2
      sourceMat
  let sourceR2 =
        R.reshape
          (Z :. numRFreq :. numThetaFreq :. numPointsRecon :. numPointsRecon)
          sourceR2'
  plotImageRepa (folderPath </> "Source.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPointsRecon :. numPointsRecon) .
    VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    sourceR2


-----------------------------------------------------------------------------
-- |
-- Module      : Berp.Base.Operators
-- Copyright   : (c) 2010 Bernie Pope
-- License     : BSD-style
-- Maintainer  : florbitous@gmail.com
-- Stability   : experimental
-- Portability : ghc
--
-- Implementation of Python's operators. Where possible we should try to
-- specialise them to commonly used types.
--
-- Note: Complex numbers intentionally don't have an ordering.
--
-----------------------------------------------------------------------------

module Berp.Base.Operators
   ( (+), (-), (*), (/), (==), (<), (>), (<=), (>=), (.), and, or, (%)
   , unaryMinus, unaryPlus, invert, not

   , modIntIntInt
   , addIntIntInt
   , addIntFloatFloat
   , addIntComplexComplex
   , subIntIntInt
   , subIntFloatFloat
   , subIntComplexComplex
   , mulIntIntInt
   , mulIntFloatFloat
   , mulIntComplexComplex
   , divIntIntInt
   , divIntFloatFloat
   , divIntComplexComplex
   , leIntIntBool
   , leIntFloatBool
   , gtIntIntBool
   , gtIntFloatBool
   , eqIntIntBool
   , eqIntFloatBool
   , eqIntComplexBool
   , ltIntIntBool
   , ltIntFloatBool
   , geIntIntBool
   , geIntFloatBool

   , addFloatFloatFloat
   , addFloatIntFloat
   , addFloatComplexComplex
   , subFloatFloatFloat
   , subFloatIntFloat
   , subFloatComplexComplex
   , mulFloatFloatFloat
   , mulFloatIntFloat
   , mulFloatComplexComplex
   , divFloatFloatFloat
   , divFloatIntFloat
   , divFloatComplexComplex
   , leFloatFloatBool
   , leFloatIntBool
   , gtFloatFloatBool
   , gtFloatIntBool
   , eqFloatFloatBool
   , eqFloatIntBool
   , eqFloatComplexBool
   , ltFloatFloatBool
   , ltFloatIntBool
   , geFloatFloatBool
   , geFloatIntBool

   , addComplexComplexComplex
   , addComplexIntComplex
   , addComplexFloatComplex
   , subComplexComplexComplex
   , subComplexIntComplex
   , subComplexFloatComplex
   , mulComplexComplexComplex
   , mulComplexIntComplex
   , mulComplexFloatComplex
   , divComplexComplexComplex
   , divComplexIntComplex
   , divComplexFloatComplex
   , eqComplexComplexBool
   , eqComplexIntBool
   , eqComplexFloatBool

   )
   where

import Data.Complex (Complex (..))
import Prelude hiding ((+), (-), (*), (.), (/), (==), (<), (>), (<=), (>=), or, and, not)
import qualified Prelude ((==),(<),(>=),(*),(+),(-),(<=),(>),(.),(/), not)
import Berp.Base.Prims (callMethod, raise)
import Berp.Base.Builtins.Exceptions (typeError, zeroDivisionError)
import Berp.Base.Object (lookupAttribute)
import Berp.Base.SemanticTypes (Object (..), Eval)
import Berp.Base.Hash (Hashed)
import Berp.Base.StdNames
   ( specialModName, specialAddName, specialSubName, specialMulName, specialDivName, specialLeName
   , specialGtName, specialEqName, specialLtName, specialGeName)
import Berp.Base.Truth (truth)
import {-# SOURCE #-} Berp.Base.StdTypes.Integer (int)
import {-# SOURCE #-} Berp.Base.StdTypes.Float (float)
import {-# SOURCE #-} Berp.Base.StdTypes.Bool (bool, true, false)
import {-# SOURCE #-} Berp.Base.StdTypes.Complex (complex)

infixl 9  .
infixl 7 *, /, %
infixl 6  +, -
infix  4  ==, <, <=, >=, >
infixr 3 `and`
infixr 2 `or`

{-
  Is it possible to minimise the boiler plate?
  Maybe Template Haskell?
-}

-- We specialise some operations for particular types rather than
-- going via the method lookups. Fall back to method look ups for the
-- general case.

binop :: Hashed String -> Object -> Object -> Eval Object
binop str arg1 arg2 = callMethod arg1 str [arg2]

specialiseOp :: (Object -> a) -> (Object -> a) ->
                (a -> a -> r) -> (r -> Object) -> Object -> Object -> Eval Object
specialiseOp project1 project2 op build obj1 obj2 =
   return $ build (project1 obj1 `op` project2 obj2)

specialiseOpIntIntInt :: (Integer -> Integer -> Integer) -> Object -> Object -> Eval Object
specialiseOpIntIntInt op = specialiseOp object_integer object_integer op int

specialiseOpIntIntBool :: (Integer -> Integer -> Bool) -> Object -> Object -> Eval Object
specialiseOpIntIntBool op = specialiseOp object_integer object_integer op bool

{-
specialiseOpBoolBoolBool :: (Bool -> Bool -> Bool) -> Object -> Object -> Eval Object
specialiseOpBoolBoolBool op = specialiseOp object_bool object_bool op bool
-}

specialiseOpFloatFloatFloat :: (Double -> Double -> Double) -> Object -> Object -> Eval Object
specialiseOpFloatFloatFloat op = specialiseOp object_float object_float op float

specialiseOpFloatFloatBool :: (Double -> Double -> Bool) -> Object -> Object -> Eval Object
specialiseOpFloatFloatBool op = specialiseOp object_float object_float op bool

specialiseOpIntFloatFloat :: (Double -> Double -> Double) -> Object -> Object -> Eval Object
specialiseOpIntFloatFloat op =
   specialiseOp (fromInteger Prelude.. object_integer) object_float op float

specialiseOpIntFloatBool :: (Double -> Double -> Bool) -> Object -> Object -> Eval Object
specialiseOpIntFloatBool op =
   specialiseOp (fromInteger Prelude.. object_integer) object_float op bool

specialiseOpFloatIntFloat :: (Double -> Double -> Double) -> Object -> Object -> Eval Object
specialiseOpFloatIntFloat op =
   specialiseOp object_float (fromInteger Prelude.. object_integer) op float

specialiseOpFloatIntBool :: (Double -> Double -> Bool) -> Object -> Object -> Eval Object
specialiseOpFloatIntBool op =
   specialiseOp object_float (fromInteger Prelude.. object_integer) op bool

type ComplexD = Complex Double

specialiseOpComplexComplexComplex :: (ComplexD -> ComplexD -> ComplexD) -> Object -> Object -> Eval Object
specialiseOpComplexComplexComplex op = specialiseOp object_complex object_complex op complex

specialiseOpComplexComplexBool :: (ComplexD -> ComplexD -> Bool) -> Object -> Object -> Eval Object
specialiseOpComplexComplexBool op = specialiseOp object_complex object_complex op bool

specialiseOpIntComplexBool :: (ComplexD -> ComplexD -> Bool) -> Object -> Object -> Eval Object
specialiseOpIntComplexBool op = specialiseOp (fromInteger Prelude.. object_integer) object_complex op bool

specialiseOpComplexIntBool :: (ComplexD -> ComplexD -> Bool) -> Object -> Object -> Eval Object
specialiseOpComplexIntBool op = specialiseOp object_complex (fromInteger Prelude.. object_integer) op bool

specialiseOpFloatComplexBool :: (ComplexD -> ComplexD -> Bool) -> Object -> Object -> Eval Object
specialiseOpFloatComplexBool op = specialiseOp (realToFrac Prelude.. object_float) object_complex op bool

specialiseOpComplexFloatBool :: (ComplexD -> ComplexD -> Bool) -> Object -> Object -> Eval Object
specialiseOpComplexFloatBool op = specialiseOp object_complex (realToFrac Prelude.. object_float) op bool

specialiseOpFloatComplexComplex :: (ComplexD -> ComplexD -> ComplexD) -> Object -> Object -> Eval Object
specialiseOpFloatComplexComplex op =
   specialiseOp (realToFrac Prelude.. object_float) object_complex op complex

specialiseOpComplexFloatComplex :: (ComplexD -> ComplexD -> ComplexD) -> Object -> Object -> Eval Object
specialiseOpComplexFloatComplex op =
   specialiseOp object_complex (realToFrac Prelude.. object_float) op complex

specialiseOpIntComplexComplex :: (ComplexD -> ComplexD -> ComplexD) -> Object -> Object -> Eval Object
specialiseOpIntComplexComplex op =
   specialiseOp (fromInteger Prelude.. object_integer) object_complex op complex

specialiseOpComplexIntComplex :: (ComplexD -> ComplexD -> ComplexD) -> Object -> Object -> Eval Object
specialiseOpComplexIntComplex op =
   specialiseOp object_complex (fromInteger Prelude.. object_integer) op complex

(%), (+), (-), (*), (/), (==), (<), (>), (<=), (>=), or, and :: Object -> Object -> Eval Object

modIntIntInt :: Object -> Object -> Eval Object
modIntIntInt = specialiseOpIntIntInt (Prelude.mod)

-- XXX fixme
(%) obj1@(Integer {}) obj2@(Integer {}) = modIntIntInt obj1 obj2
(%) x y = binop specialModName x y

addIntIntInt, addFloatFloatFloat, addIntFloatFloat, addFloatIntFloat, addComplexComplexComplex, addIntComplexComplex, addComplexIntComplex, addFloatComplexComplex, addComplexFloatComplex :: Object -> Object -> Eval Object

addIntIntInt = specialiseOpIntIntInt (Prelude.+)
addFloatFloatFloat = specialiseOpFloatFloatFloat (Prelude.+)
addIntFloatFloat = specialiseOpIntFloatFloat (Prelude.+)
addFloatIntFloat = specialiseOpFloatIntFloat (Prelude.+)
addComplexComplexComplex = specialiseOpComplexComplexComplex (Prelude.+)
addIntComplexComplex = specialiseOpIntComplexComplex (Prelude.+)
addComplexIntComplex = specialiseOpComplexIntComplex (Prelude.+)
addFloatComplexComplex = specialiseOpFloatComplexComplex (Prelude.+)
addComplexFloatComplex = specialiseOpComplexFloatComplex (Prelude.+)

(+) obj1@(Integer {}) obj2 =
   case obj2 of
      Integer {} -> addIntIntInt obj1 obj2
      Float {} -> addIntFloatFloat obj1 obj2
      Complex {} -> addIntComplexComplex obj1 obj2
      _other -> raise typeError
(+) obj1@(Float {}) obj2 =
   case obj2 of
      Float {} -> addFloatFloatFloat obj1 obj2
      Integer {} -> addFloatIntFloat obj1 obj2
      Complex {} -> addFloatComplexComplex obj1 obj2
      _other -> raise typeError
(+) obj1@(Complex {}) obj2 =
   case obj2 of
      Complex {} -> addComplexComplexComplex obj1 obj2
      Integer {} -> addComplexIntComplex obj1 obj2
      Float {} -> addComplexFloatComplex obj1 obj2
      _other -> raise typeError
(+) x y = binop specialAddName x y

subIntIntInt, subFloatFloatFloat, subIntFloatFloat, subFloatIntFloat, subComplexComplexComplex, subIntComplexComplex, subComplexIntComplex, subFloatComplexComplex, subComplexFloatComplex :: Object -> Object -> Eval Object

subIntIntInt = specialiseOpIntIntInt (Prelude.-)
subFloatFloatFloat = specialiseOpFloatFloatFloat (Prelude.-)
subIntFloatFloat = specialiseOpIntFloatFloat (Prelude.-)
subFloatIntFloat = specialiseOpFloatIntFloat (Prelude.-)
subComplexComplexComplex = specialiseOpComplexComplexComplex (Prelude.-)
subIntComplexComplex = specialiseOpIntComplexComplex (Prelude.-)
subComplexIntComplex = specialiseOpComplexIntComplex (Prelude.-)
subFloatComplexComplex = specialiseOpFloatComplexComplex (Prelude.-)
subComplexFloatComplex = specialiseOpComplexFloatComplex (Prelude.-)

(-) obj1@(Integer {}) obj2 =
   case obj2 of
      Integer {} -> subIntIntInt obj1 obj2
      Float {} -> subIntFloatFloat obj1 obj2
      Complex {} -> subIntComplexComplex obj1 obj2
      _other -> raise typeError
(-) obj1@(Float {}) obj2 =
   case obj2 of
      Float {} -> subFloatFloatFloat obj1 obj2
      Integer {} -> subFloatIntFloat obj1 obj2
      Complex {} -> subFloatComplexComplex obj1 obj2
      _other -> raise typeError
(-) obj1@(Complex {}) obj2 =
   case obj2 of
      Complex {} -> subComplexComplexComplex obj1 obj2
      Integer {} -> subComplexIntComplex obj1 obj2
      Float {} -> subComplexFloatComplex obj1 obj2
      _other -> raise typeError
(-) x y = binop specialSubName x y

mulIntIntInt, mulFloatFloatFloat, mulIntFloatFloat, mulFloatIntFloat, mulComplexComplexComplex, mulIntComplexComplex, mulComplexIntComplex, mulFloatComplexComplex, mulComplexFloatComplex :: Object -> Object -> Eval Object

mulIntIntInt = specialiseOpIntIntInt (Prelude.*)
mulFloatFloatFloat = specialiseOpFloatFloatFloat (Prelude.*)
mulIntFloatFloat = specialiseOpIntFloatFloat (Prelude.*)
mulFloatIntFloat = specialiseOpFloatIntFloat (Prelude.*)
mulComplexComplexComplex = specialiseOpComplexComplexComplex (Prelude.*)
mulIntComplexComplex = specialiseOpIntComplexComplex (Prelude.*)
mulComplexIntComplex = specialiseOpComplexIntComplex (Prelude.*)
mulFloatComplexComplex = specialiseOpFloatComplexComplex (Prelude.*)
mulComplexFloatComplex = specialiseOpComplexFloatComplex (Prelude.*)

(*) obj1@(Integer {}) obj2 =
   case obj2 of
      Integer {} -> mulIntIntInt obj1 obj2
      Float {} -> mulIntFloatFloat obj1 obj2
      Complex {} -> mulIntComplexComplex obj1 obj2
      _other -> raise typeError
(*) obj1@(Float {}) obj2 =
   case obj2 of
      Float {} -> mulFloatFloatFloat obj1 obj2
      Integer {} -> mulFloatIntFloat obj1 obj2
      Complex {} -> mulFloatComplexComplex obj1 obj2
      _other -> raise typeError
(*) obj1@(Complex {}) obj2 =
   case obj2 of
      Complex {} -> mulComplexComplexComplex obj1 obj2
      Integer {} -> mulComplexIntComplex obj1 obj2
      Float {} -> mulComplexFloatComplex obj1 obj2
      _other -> raise typeError
(*) x y = binop specialMulName x y

checkDivByZero :: Num a => a -> Eval Object -> Eval Object
checkDivByZero denom comp
   | denom Prelude.== 0 = raise zeroDivisionError
   | otherwise = comp

divIntIntInt, divFloatFloatFloat, divIntFloatFloat, divFloatIntFloat, divComplexComplexComplex, divIntComplexComplex, divComplexIntComplex, divFloatComplexComplex, divComplexFloatComplex :: Object -> Object -> Eval Object

divIntIntInt obj1 obj2 =
   checkDivByZero (object_integer obj2) $ specialiseOpIntIntInt (Prelude.div) obj1 obj2
divFloatFloatFloat obj1 obj2 =
   checkDivByZero (object_float obj2) $ specialiseOpFloatFloatFloat (Prelude./) obj1 obj2
divIntFloatFloat obj1 obj2 =
   checkDivByZero (object_float obj2) $ specialiseOpIntFloatFloat (Prelude./) obj1 obj2
divFloatIntFloat obj1 obj2 =
   checkDivByZero (object_integer obj2) $ specialiseOpFloatIntFloat (Prelude./) obj1 obj2
divComplexComplexComplex obj1 obj2 =
   checkDivByZero (object_complex obj2) $ specialiseOpComplexComplexComplex (Prelude./) obj1 obj2
divIntComplexComplex obj1 obj2 =
   checkDivByZero (object_complex obj2) $ specialiseOpIntComplexComplex (Prelude./) obj1 obj2
divComplexIntComplex obj1 obj2 =
   checkDivByZero (object_integer obj2) $ specialiseOpComplexIntComplex (Prelude./) obj1 obj2
divFloatComplexComplex obj1 obj2 =
   checkDivByZero (object_complex obj2) $ specialiseOpFloatComplexComplex (Prelude./) obj1 obj2
divComplexFloatComplex obj1 obj2 =
   checkDivByZero (object_float obj2) $ specialiseOpComplexFloatComplex (Prelude./) obj1 obj2

(/) obj1@(Integer {}) obj2 =
   case obj2 of
      Integer {} -> divIntIntInt obj1 obj2
      Float {} -> divIntFloatFloat obj1 obj2
      Complex {} -> divIntComplexComplex obj1 obj2
      _other -> raise typeError
(/) obj1@(Float {}) obj2 =
   case obj2 of
      Float {} -> divFloatFloatFloat obj1 obj2
      Integer {} -> divFloatIntFloat obj1 obj2
      Complex {} -> divFloatComplexComplex obj1 obj2
      _other -> raise typeError
(/) obj1@(Complex {}) obj2 =
   case obj2 of
      Complex {} -> divComplexComplexComplex obj1 obj2
      Integer {} -> divComplexIntComplex obj1 obj2
      Float {} -> divComplexFloatComplex obj1 obj2
      _other -> raise typeError
(/) x y = binop specialDivName x y

leIntIntBool, leFloatFloatBool, leIntFloatBool, leFloatIntBool :: Object -> Object -> Eval Object

leIntIntBool = specialiseOpIntIntBool (Prelude.<=)
leFloatFloatBool = specialiseOpFloatFloatBool (Prelude.<=)
leIntFloatBool = specialiseOpIntFloatBool (Prelude.<=)
leFloatIntBool = specialiseOpFloatIntBool (Prelude.<=)

(<=) obj1@(Integer {}) obj2 =
   case obj2 of
      Integer {} -> leIntIntBool obj1 obj2
      Float {} -> leIntFloatBool obj1 obj2
      _other -> raise typeError
(<=) obj1@(Float {}) obj2 =
   case obj2 of
      Float {} -> leFloatFloatBool obj1 obj2
      Integer {} -> leIntFloatBool obj1 obj2
      _other -> raise typeError
(<=) x y = binop specialLeName x y

gtIntIntBool, gtFloatFloatBool, gtIntFloatBool, gtFloatIntBool :: Object -> Object -> Eval Object

gtIntIntBool = specialiseOpIntIntBool (Prelude.>)
gtFloatFloatBool = specialiseOpFloatFloatBool (Prelude.>)
gtIntFloatBool = specialiseOpIntFloatBool (Prelude.>)
gtFloatIntBool = specialiseOpFloatIntBool (Prelude.>)

(>) obj1@(Integer {}) obj2 =
   case obj2 of
      Integer {} -> gtIntIntBool obj1 obj2
      Float {} -> gtIntFloatBool obj1 obj2
      _other -> raise typeError
(>) obj1@(Float {}) obj2 =
   case obj2 of
      Float {} -> gtFloatFloatBool obj1 obj2
      Integer {} -> gtIntFloatBool obj1 obj2
      _other -> raise typeError
(>) x y = binop specialGtName x y

eqIntIntBool, eqFloatFloatBool, eqIntFloatBool, eqFloatIntBool, eqComplexComplexBool, eqIntComplexBool, eqComplexIntBool, eqFloatComplexBool, eqComplexFloatBool :: Object -> Object -> Eval Object

eqIntIntBool = specialiseOpIntIntBool (Prelude.==)
eqFloatFloatBool = specialiseOpFloatFloatBool (Prelude.==)
eqIntFloatBool = specialiseOpIntFloatBool (Prelude.==)
eqFloatIntBool = specialiseOpFloatIntBool (Prelude.==)
eqComplexComplexBool = specialiseOpComplexComplexBool (Prelude.==)
eqIntComplexBool = specialiseOpIntComplexBool (Prelude.==)
eqComplexIntBool = specialiseOpComplexIntBool (Prelude.==)
eqFloatComplexBool = specialiseOpFloatComplexBool (Prelude.==)
eqComplexFloatBool = specialiseOpComplexFloatBool (Prelude.==)

(==) obj1@(Integer {}) obj2 =
   case obj2 of
      Integer {} -> eqIntIntBool obj1 obj2
      Float {} -> eqIntFloatBool obj1 obj2
      Complex {} -> eqIntComplexBool obj1 obj2
      _other -> raise typeError
(==) obj1@(Float {}) obj2 =
   case obj2 of
      Float {} -> eqFloatFloatBool obj1 obj2
      Integer {} -> eqIntFloatBool obj1 obj2
      Complex {} -> eqFloatComplexBool obj1 obj2
      _other -> raise typeError
(==) obj1@(Complex {}) obj2 =
   case obj2 of
      Complex {} -> eqComplexComplexBool obj1 obj2
      Integer {} -> eqComplexIntBool obj1 obj2
      Float {} -> eqComplexFloatBool obj1 obj2
      _other -> raise typeError
(==) x y = binop specialEqName x y

ltIntIntBool, ltFloatFloatBool, ltIntFloatBool, ltFloatIntBool :: Object -> Object -> Eval Object

ltIntIntBool = specialiseOpIntIntBool (Prelude.<)
ltFloatFloatBool = specialiseOpFloatFloatBool (Prelude.<)
ltIntFloatBool = specialiseOpIntFloatBool (Prelude.<)
ltFloatIntBool = specialiseOpFloatIntBool (Prelude.<)

(<) obj1@(Integer {}) obj2 =
   case obj2 of
      Integer {} -> ltIntIntBool obj1 obj2
      Float {} -> ltIntFloatBool obj1 obj2
      _other -> raise typeError
(<) obj1@(Float {}) obj2 =
   case obj2 of
      Float {} -> ltFloatFloatBool obj1 obj2
      Integer {} -> ltIntFloatBool obj1 obj2
      _other -> raise typeError
(<) x y = binop specialLtName x y

geIntIntBool, geFloatFloatBool, geIntFloatBool, geFloatIntBool :: Object -> Object -> Eval Object

geIntIntBool = specialiseOpIntIntBool (Prelude.>=)
geFloatFloatBool = specialiseOpFloatFloatBool (Prelude.>=)
geIntFloatBool = specialiseOpIntFloatBool (Prelude.>=)
geFloatIntBool = specialiseOpFloatIntBool (Prelude.>=)

(>=) obj1@(Integer {}) obj2 =
   case obj2 of
      Integer {} -> geIntIntBool obj1 obj2
      Float {} -> geIntFloatBool obj1 obj2
      _other -> raise typeError
(>=) obj1@(Float {}) obj2 =
   case obj2 of
      Float {} -> geFloatFloatBool obj1 obj2
      Integer {} -> geIntFloatBool obj1 obj2
      _other -> raise typeError
(>=) x y = binop specialGeName x y

{-
   From the Python Language Reference, sec 5.10 "Boolean Operations"
   The expression x and y first evaluates x; if x is false, its value
   is returned; otherwise, y is evaluated and the resulting value is
   returned.
-}
and obj1 obj2
   | truth obj1 = return obj2
   | otherwise  = return obj1

{-
   The expression x or y first evaluates x; if x is true, its value
   is returned; otherwise, y is evaluated and the resulting value
   is returned.
-}

or obj1 obj2
   | truth obj1 = return obj1
   | otherwise  = return obj2

(.) :: Object -> Hashed String -> Eval Object
(.) object ident = lookupAttribute object ident

unaryMinus :: Object -> Eval Object
unaryMinus obj@(Integer {}) = return $ int $ negate $ object_integer obj
unaryMinus obj@(Float {}) = return $ float $ negate $ object_float obj
unaryMinus obj@(Complex {}) = return $ complex $ negate $ object_complex obj
unaryMinus _other = error "unary minus applied to a non integer"

-- This is just the identity function
unaryPlus :: Object -> Eval Object
unaryPlus obj@(Integer {}) = return obj
unaryPlus obj@(Float {}) = return obj
unaryPlus obj@(Complex {}) = return obj
-- XXX in CPython this turns the boolean into an integer
unaryPlus obj@(Bool {}) = return obj
unaryPlus _other = error "unary plus applied to a non integer"

invert :: Object -> Eval Object
invert (Integer {}) = error "bitwise inversion not implemented"
invert _other = error "unary invert applied to a non integer"

not :: Object -> Eval Object
not obj
   | Prelude.not $ truth obj = return true
   | otherwise = return false


{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE RankNTypes            #-}

{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
module Test.Grenade.Recurrent.Layers.LSTM where

import           Hedgehog
import           Hedgehog.Internal.Source
import           Hedgehog.Internal.Show
import           Hedgehog.Internal.Property ( failWith, Diff (..) )

import           Data.Foldable ( toList )
import           Data.Singletons.TypeLits

import           Grenade
import           Grenade.Recurrent

import qualified Numeric.LinearAlgebra as H
import qualified Numeric.LinearAlgebra.Static as S

import qualified Test.Grenade.Recurrent.Layers.LSTM.Reference as Reference
import           Test.Hedgehog.Hmatrix

genLSTM :: forall i o. (KnownNat i, KnownNat o) => Gen (LSTM i o)
genLSTM = do
    let w = uniformSample
        u = uniformSample
        v = randomVector

        w0 = S.konst 0
        u0 = S.konst 0
        v0 = S.konst 0

    LSTM <$> (LSTMWeights <$> w <*> u <*> v <*> w <*> u <*> v <*> w <*> u <*> v <*> w <*> v)
         <*> pure (LSTMWeights w0 u0 v0 w0 u0 v0 w0 u0 v0 w0 v0)

prop_lstm_reference_forwards =
  property $ do
    input :: S.R 3                       <- forAll randomVector
    cell :: S.R 2                        <- forAll randomVector
    net@(LSTM lstmWeights _) :: LSTM 3 2 <- forAll genLSTM

    let actual          = runRecurrentForwards net (S1D cell) (S1D input)
    case actual of
      (_, (S1D cellOut) :: S ('D1 2), (S1D output) :: S ('D1 2)) ->
        let cellOut'        = Reference.Vector . H.toList . S.extract $ cellOut
            output'         = Reference.Vector . H.toList . S.extract $ output
            refNet          = Reference.lstmToReference lstmWeights
            refCell         = Reference.Vector . H.toList . S.extract $ cell
            refInput        = Reference.Vector . H.toList . S.extract $ input
            (refCO, refO)   = Reference.runLSTM refNet refCell refInput
        in do toList refCO ~~~ toList cellOut'
              toList refO ~~~ toList output'

prop_lstm_reference_backwards =
  property $ do
    input :: S.R 3                       <- forAll randomVector
    cell :: S.R 2                        <- forAll randomVector
    net@(LSTM lstmWeights _) :: LSTM 3 2 <- forAll genLSTM
    let (tape, _ :: S ('D1 2), _ :: S ('D1 2))
                                          = runRecurrentForwards  net (S1D cell) (S1D input)
        actualBacks                       = runRecurrentBackwards net tape (S1D (S.konst 1) :: S ('D1 2)) (S1D (S.konst 1) :: S ('D1 2))
    case actualBacks of
      (actualGradients, _, _ :: S ('D1 3)) ->
        let refNet          = Reference.lstmToReference lstmWeights
            refCell         = Reference.Vector . H.toList . S.extract $ cell
            refInput        = Reference.Vector . H.toList . S.extract $ input
            refGradients    = Reference.runLSTMback refCell refInput refNet
        in toList refGradients ~~~ toList (Reference.lstmToReference actualGradients)

prop_lstm_reference_backwards_input =
  property $ do
    input :: S.R 3                       <- forAll randomVector
    cell :: S.R 2                        <- forAll randomVector
    net@(LSTM lstmWeights _) :: LSTM 3 2 <- forAll genLSTM
    let (tape, _ :: S ('D1 2), _ :: S ('D1 2))
                                          = runRecurrentForwards  net (S1D cell) (S1D input)
        actualBacks                       = runRecurrentBackwards net tape (S1D (S.konst 1) :: S ('D1 2)) (S1D (S.konst 1) :: S ('D1 2))
    case actualBacks of
      (_, _, S1D actualGradients :: S ('D1 3)) ->
        let refNet          = Reference.lstmToReference lstmWeights
            refCell         = Reference.Vector . H.toList . S.extract $ cell
            refInput        = Reference.Vector . H.toList . S.extract $ input
            refGradients    = Reference.runLSTMbackOnInput refCell refNet refInput
        in toList refGradients ~~~ H.toList (S.extract actualGradients)

prop_lstm_reference_backwards_cell =
  property $ do
    input :: S.R 3                       <- forAll randomVector
    cell :: S.R 2                        <- forAll randomVector
    net@(LSTM lstmWeights _) :: LSTM 3 2 <- forAll genLSTM
    let (tape, _ :: S ('D1 2), _ :: S ('D1 2))
                                          = runRecurrentForwards  net (S1D cell) (S1D input)
        actualBacks                       = runRecurrentBackwards net tape (S1D (S.konst 1) :: S ('D1 2)) (S1D (S.konst 1) :: S ('D1 2))
    case actualBacks of
      (_, S1D actualGradients, _ :: S ('D1 3)) ->
        let refNet          = Reference.lstmToReference lstmWeights
            refCell         = Reference.Vector . H.toList . S.extract $ cell
            refInput        = Reference.Vector . H.toList . S.extract $ input
            refGradients    = Reference.runLSTMbackOnCell refInput refNet refCell
        in toList refGradients ~~~ H.toList (S.extract actualGradients)

(~~~) :: (Monad m, Eq a, Ord a, Num a, Fractional a, Show a, HasCallStack) => [a] -> [a] -> PropertyT m ()
(~~~) x y =
  if all (< 1e-8) (zipWith (-) x y) then
    success
  else
    case valueDiff <$> mkValue x <*> mkValue y of
      Nothing ->
        withFrozenCallStack $
          failWith Nothing $ unlines [
              "━━━ Not Simliar ━━━"
            , showPretty x
            , showPretty y
            ]
      Just differ ->
        withFrozenCallStack $
          failWith (Just $ Diff "Failed (" "- lhs" "~/~" "+ rhs" ")" differ) ""
infix 4 ~~~

tests :: IO Bool
tests = checkParallel $$(discover)


{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLabels      #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-|
Module      : Grenade.Layers.Mul
Description : Scalar multiplication with some constant, supporting scaling
Maintainer  : Theo Charalambous
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Mul
  (
  -- * Layer Definition
    Mul (..)

  -- * Helper function
  , initMul
  )
where

import           Data.Kind                    (Type)
import           Data.Proxy
import           Data.Serialize
import           GHC.TypeLits

import qualified Numeric.LinearAlgebra        as LA
import           Numeric.LinearAlgebra.Static (R)
import qualified Numeric.LinearAlgebra.Static as H

import           Lens.Micro                   ((^.))

import           Grenade.Core
import           Grenade.Onnx

-- | A layer allowing for element wise multiplication with broadcasting
data Mul :: Nat -- The number of channels of the bias
         -> Nat -- The number of rows of the bias
         -> Nat -- The number of columns of the bias
         -> Type where
  Mul  :: ( KnownNat channels
          , KnownNat rows
          , KnownNat columns )
          => R (channels * rows * columns)
          -> Mul channels rows columns

instance Show (Mul c h w) where
  show (Mul mat) = "Mul " ++ show mat

instance UpdateLayer (Mul c h w) where
  type Gradient (Mul c h w) = ()
  runUpdate _ x _  = x
  reduceGradient _ = ()

instance (KnownNat c, KnownNat h, KnownNat w ) => RandomLayer (Mul c h w) where
  createRandomWith _ _ = pure initMul

-- ^ Initialize a Mul layer with scale 1
initMul :: forall c h w. ( KnownNat c, KnownNat h, KnownNat w )
        => Mul c h w
initMul =
  let c'    = fromIntegral $ natVal (Proxy :: Proxy c)
      h'    = fromIntegral $ natVal (Proxy :: Proxy h)
      w'    = fromIntegral $ natVal (Proxy :: Proxy w)
      ones  = replicate (c' * h' * w') 1
      bias  = H.fromList ones :: R (c * h * w)
  in Mul bias

instance ( KnownNat c, KnownNat h, KnownNat w ) => Serialize (Mul c h w) where
  put (Mul bias) = putListOf put . LA.toList . H.extract $ bias
  get            = do
    bias <- maybe (fail "Vector of incorrect size") return . H.create . LA.fromList =<< getListOf get
    return $ Mul bias

-- | Currently, only multiplication by a single scalar is supported.
instance ( KnownNat i, KnownNat j, KnownNat k ) => Layer (Mul 1 1 1) ('D3 i j k) ('D3 i j k) where
  type Tape (Mul 1 1 1) ('D3 i j k) ('D3 i j k) = ()

  runForwards (Mul b) (S3D m)
    = let s  = H.extract b LA.! 0
      in  ((), S3D $  H.dmmap (s *) m)

  runBackwards = error "runBackwards is not implemented for the Mul layer"

instance OnnxOperator (Mul c h w) where
  onnxOpTypeNames _ = ["Mul"]

instance (KnownNat c, KnownNat h, KnownNat w) => OnnxLoadable (Mul c h w) where
  loadOnnxNode inits node = case node ^. #input of
    [_, scale] -> do
      loadedScale <- readInitializerVector inits scale

      return $ Mul loadedScale
    _               -> onnxIncorrectNumberOfInputs


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE OverloadedStrings     #-}

module Main where

import           Control.Applicative
import           Control.Monad
import           Control.Monad.Random

import           Data.List ( foldl' )
import           Data.Maybe ( mapMaybe )
#if ! MIN_VERSION_base(4,13,0)
import           Data.Semigroup ( (<>) )
#endif
import           Data.Csv ( FromField, FromRecord, decode, parseField, HasHeader(..) )
import           Data.ByteString.Lazy.Char8 (pack)
import qualified Data.Vector as V
import qualified Data.Vector.Storable as SV

import           Numeric.LinearAlgebra ( maxIndex )
import qualified Numeric.LinearAlgebra.Static as SA

import           Options.Applicative
import           System.FilePath ( (</>) )
import           System.Random.Shuffle (shuffleM)
import           GHC.Generics (Generic)
import           GHC.Float ( float2Double)

import           Grenade
import           Grenade.Utils.OneHot

-- Adapted from https://mmhaskell.com/machine-learning/deep-learning

-- It's logistic regression!
--
-- This network is used to show how we can embed a Network as a layer in the larger IrisNetwork
-- type.

type IrisNetwork
  = Network
      '[FullyConnected 4 10, Relu, FullyConnected 10 3]
      '[ 'D1 4, 'D1 10, 'D1 10, 'D1 3]

type IrisRow = (S ( 'D1 4), S ( 'D1 3))

randomIris :: MonadRandom m => m IrisNetwork
randomIris = randomNetwork

runIris :: Int -> FilePath -> Int -> LearningParameters -> IO ()
runIris iterations dataDir nSamples rate = do
  records <- readIrisFromFile (dataDir </> "iris.data")
  let numRecords = V.length records
  shuffledRecords <- chooseRandomRecords records numRecords
  let (trainRecords, validateRecords) = V.splitAt nSamples shuffledRecords

  let trainData = mapMaybe parseRecord (V.toList trainRecords)
  let validateData = mapMaybe parseRecord (V.toList validateRecords)

  if length trainData
       /= length trainRecords
       || length validateData
       /= length validateRecords
    then putStrLn
      "Parsing train data or validation data could not be fully parsed"
    else do
      initialNetwork <- randomIris
      foldM_ (run trainData validateData) initialNetwork [1 .. iterations]
 where

  run :: [IrisRow] -> [IrisRow] -> IrisNetwork -> Int -> IO IrisNetwork
  run trainData validateData network iterationNum = do
    sampledData <- V.toList
      <$> chooseRandomRecords (V.fromList trainData) (nSamples * 3 `div` 4)
    -- Slower drop the learning rate
    let rate' = rate { learningRate = learningRate rate * 0.99 ^ iterationNum }
    let newNetwork     = foldl' (trainRow rate') network sampledData
    let labelVectors   = fmap (testRow newNetwork) validateData
    let labelValues    = fmap getLabels labelVectors
    let total          = length labelValues
    let correctEntries = length $ filter ((==) <$> fst <*> snd) labelValues
    putStrLn $ "Iteration: " ++ show iterationNum
    putStrLn $ show correctEntries ++ " correct out of: " ++ show total
    return newNetwork

  trainRow :: LearningParameters -> IrisNetwork -> IrisRow -> IrisNetwork
  trainRow lp network (input, output) = train lp network input output

  -- Takes a test row, returns predicted output and actual output from the network.
  testRow :: IrisNetwork -> IrisRow -> (S ( 'D1 3), S ( 'D1 3))
  testRow net (rowInput, predictedOutput) =
    (predictedOutput, runNet net rowInput)

  -- Goes from probability output vector to label
  getLabels :: (S ( 'D1 3), S ( 'D1 3)) -> (Int, Int)
  getLabels (S1D predictedLabel, S1D actualOutput) =
    (maxIndex (SA.extract predictedLabel), maxIndex (SA.extract actualOutput))

data IrisOpts = IrisOpts FilePath Int Int LearningParameters

iris' :: Parser IrisOpts
iris' =
  IrisOpts
    <$> argument str (metavar "DATADIR")
        -- How many samples from the dataset should be used for training?
        -- (The rest are used for validation)
    <*> option auto (long "training_samples" <> short 't' <> value 100)
    <*> option auto (long "iterations" <> short 'i' <> value 15)
    <*> (   LearningParameters
        <$> option auto (long "train_rate" <> short 'r' <> value 0.01)
        <*> option auto (long "momentum" <> value 0.9)
        <*> option auto (long "l2" <> value 0.0005)
        )

main :: IO ()
main = do
  IrisOpts dataDir nSamples iter rate <- execParser
    (info (iris' <**> helper) idm)
  putStr "Training convolutional neural network with "
  putStr $ show nSamples
  putStrLn " samples..."

  runIris iter dataDir nSamples rate

data IrisClass = Setosa | Versicolor | Virginica
  deriving (Show, Read, Eq, Ord, Generic, Enum, Bounded)

data IrisRecord = IrisRecord {
    sepalLength :: Float,
    sepalWidth  :: Float,
    petalLength :: Float,
    petalWidth  :: Float,
    specie      :: IrisClass
} deriving (Generic, Show, Read)

instance FromRecord IrisRecord

instance FromField IrisClass where
  parseField "Iris-setosa"     = return Setosa
  parseField "Iris-versicolor" = return Versicolor
  parseField "Iris-virginica"  = return Virginica
  parseField _                 = fail "unknown iris class"

parseRecord :: IrisRecord -> Maybe IrisRow
parseRecord record = case (input, output) of
  (Just i, Just o) -> Just (i, o)
  _                -> Nothing
 where
  input =
    fromStorable
      $   SV.fromList
      $   float2Double
      <$> [ sepalLength record / 8.0
          , sepalWidth record / 8.0
          , petalLength record / 8.0
          , petalWidth record / 8.0
          ]
  output = oneHot (fromEnum $ specie record)

readIrisFromFile :: FilePath -> IO (V.Vector IrisRecord)
readIrisFromFile fp = do
  contents <- readFile fp
  let contentsAsBs = pack contents
  let results =
        decode HasHeader contentsAsBs :: Either String (V.Vector IrisRecord)
  case results of
    Left  err     -> error err
    Right records -> return records

-- A function that takes this vector of records, and selects sampleSize of them at random.
chooseRandomRecords :: V.Vector a -> Int -> IO (V.Vector a)
chooseRandomRecords records sampleSize = do
  let numRecords = V.length records
  chosenIndices <- take sampleSize <$> shuffleM [0 .. (numRecords - 1)]
  return . V.fromList $ (records V.!) <$> chosenIndices


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

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

module Spiral.Array.Repr.Complex (
    RE,
    toRealArray,

    CMPLX,
    toComplexArray,

    Array(..)
  ) where

import Prelude hiding (read)

import Data.Complex (Complex)
import Data.Monoid ((<>))
import Text.PrettyPrint.Mainland
import Text.PrettyPrint.Mainland.Class

import Spiral.Array
import Spiral.Exp

data RE r

toRealArray :: Array r (sh :. Int) (Exp (Complex a)) -> Array (RE r) (sh :. Int) (Exp a)
toRealArray = RE

instance IsArray r (sh :. Int) (Exp (Complex a)) => IsArray (RE r) (sh :. Int) (Exp a) where
    data Array (RE r) (sh :. Int) (Exp a) = RE (Array r (sh :. Int) (Exp (Complex a)))

    extent (RE a) = sh :. 2*n
      where
        sh :. n = extent a

instance Pretty (Array r (sh :. Int) (Exp (Complex a))) => Pretty (Array (RE r) (sh :. Int) (Exp a)) where
    ppr (RE a) = text "Re" <> parens (ppr a)

instance (RealFloatConst a, Shape sh, IArray r (sh :. Int) (Exp (Complex a))) => IArray (RE r) (sh :. Int) (Exp a) where
    index (RE a) (sh :. i)
        | i `rem` 2 == 0 = re
        | otherwise      = im
      where
        (re, im) = unComplexE (index a (sh :. i `quot` 2))

instance (RealFloatConst a, SArray r (sh :. Int) (Exp (Complex a))) => SArray (RE r) (sh :. Int) (Exp a) where
    indexS (RE a) (sh :. i)
        | i `rem` 2 == 0 = re
        | otherwise      = im
      where
        (re, im) = unComplexE (indexS a (sh :. i `quot` 2))

instance (RealFloatConst a, MArray r (sh :. Int) (Exp (Complex a))) => MArray (RE r) (sh :. Int) (Exp a) where
    read (RE a) (sh :. i) = do
        (re, im) <- unComplexE <$> read a (sh :. i `quot` 2)
        if i `rem` 2 == 0
          then return re
          else return im

    write (RE a) (sh :. i) e = do
        (re, im) <- unComplexE <$> read a (sh :. i `quot` 2)
        if i `rem` 2 == 0
          then write a (sh :. i `quot` 2) (ComplexE e im)
          else write a (sh :. i `quot` 2) (ComplexE re e)

instance (RealFloatConst a, SArray r (sh :. Int) (Exp (Complex a))) => Computable (RE r) (sh :. Int) (Exp a) where
    computeP a b =
        forShapeP (extent b) $ \ix ->
            write a ix (indexS b ix)

data CMPLX r

toComplexArray :: Array r (sh :. Int) (Exp a) -> Array (CMPLX r) (sh :. Int) (Exp (Complex a))
toComplexArray = CMPLX

instance IsArray r (sh :. Int) (Exp a) => IsArray (CMPLX r) (sh :. Int) (Exp (Complex a)) where
    data Array (CMPLX r) (sh :. Int) (Exp (Complex a)) = CMPLX (Array r (sh :. Int) (Exp a))

    extent (CMPLX a) = sh :. n `quot` 2
      where
        sh :. n = extent a

instance Pretty (Array r (sh :. Int) (Exp a)) => Pretty (Array (CMPLX r) (sh :. Int) (Exp (Complex a))) where
    ppr (CMPLX a) = text "Cmplx" <> parens (ppr a)

instance (Typed a, Num (Exp a), IArray r (sh :. Int) (Exp a)) => IArray (CMPLX r) (sh :. Int) (Exp (Complex a)) where
    index (CMPLX a) (sh :. i) = ComplexE (index a (sh :. 2*i)) (index a (sh :. 2*i+1))

instance (Typed a, Num (Exp a), SArray r (sh :. Int) (Exp a)) => SArray (CMPLX r) (sh :. Int) (Exp (Complex a)) where
    indexS (CMPLX a) (sh :. i) = ComplexE (indexS a (sh :. 2*i)) (indexS a (sh :. 2*i+1))

instance (RealFloatConst a, MArray r (sh :. Int) (Exp a)) => MArray (CMPLX r) (sh :. Int) (Exp (Complex a)) where
    read (CMPLX a) (sh :. i) = ComplexE <$> read a (sh :. 2*i) <*> read a (sh :. 2*i+1)

    write (CMPLX a) (sh :. i) e = do
        write a (sh :. 2*i) er
        write a (sh :. 2*i+1) ei
      where
        (er, ei) = unComplexE e

instance (RealFloatConst a, Computable r (sh :. Int) (Exp a)) => Computable (CMPLX r) (sh :. Int) (Exp (Complex a)) where
    computeP a (CMPLX b) = computeP (toRealArray a) b


-- * This module defines how to run and extract time/space data
--   from benchmarks.
--
module Test.BigOh.Benchmark
  ( -- * Running benchmarks
    runInputs
  , getTimes
  , getStdDevs
  , defaultBenchmarkConfig
  ) where

import           Control.Monad.IO.Class
import           Criterion.Internal
import           Criterion.Main
import           Criterion.Measurement
import           Criterion.Monad
import           Criterion.Types
import           Statistics.Resampling.Bootstrap

import           Test.BigOh.Generate

defaultBenchmarkConfig = defaultConfig { timeLimit = 1, resamples = 100 }

runOne :: Config -> Benchmarkable -> IO Report
runOne cfg x
  = withConfig cfg
  $ do liftIO initializeTime
       runAndAnalyseOne 0 "" x

runInputs :: Config -> [(Benchmarkable, Input a)] -> IO [(Input a, Report)]
runInputs cfg xs
  = zip (map snd xs) <$> mapM (runOne cfg . fst) xs

getTimes :: [(Input a, Report)] -> [(Int, Double)]
getTimes = map go
  where
    go (i, report)
      = (inputSize i, estPoint $ anMean $ reportAnalysis report)

getStdDevs :: [(Input a, Report)] -> [Double]
getStdDevs = map (estPoint . anStdDev . reportAnalysis . snd)


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Main where

-- import Data.Array.Repa (Array, U, DIM2)
import qualified Data.Array.Repa as R
import qualified Data.IntMap as IM
import Control.Monad
import Data.Complex
import qualified Data.Vector as V
import qualified Linear.V as V
import LSI
import LSI.Examples
import LSI.FrequencyResponse
import LSI.Grid
import LSI.RationalFunction
import LSI.TransferFunction
import Unsafe.Coerce (unsafeCoerce)

graph :: SystemGraph 2 String Double
graph = toGraph $ deriche_ydiff 1.0

tfs :: IM.IntMap (RationalFunction 2 Double)
tfs = IM.map simplify $ transferFunctions graph

frequencyResponses :: IM.IntMap (V.V 2 Double -> Complex Double)
frequencyResponses = IM.map computeFR tfs

n :: Num a => a
n = 50

grid :: R.Array R.D R.DIM2 (Double, Double)
grid = frequencyGrid2D n

nPoints :: Num a => Double
nPoints = (2*n)^2

plotFR :: (V.V 2 Double -> Complex Double) -> R.Array R.D R.DIM2 (Complex Double)
plotFR fr = R.map (fr . V.V . V.fromList . \(w1,w2) -> [w1,w2]) grid

l2Norm :: R.Array R.D R.DIM2 (Complex Double) -> IO Double
l2Norm plot = fmap (/ nPoints) $ R.foldAllP (+) 0 squares
  where squares = R.map (\x -> realPart $ (abs x)^2) plot

plots :: IM.IntMap (R.Array R.D R.DIM2 (Complex Double))
plots = IM.map plotFR $ unsafeCoerce frequencyResponses -- TODO: Kind mismatch.

main :: IO ()
main = do
  -- print grid
  -- let plot = plots IM.! 13
  -- forM_ [0..2*n+1] $ \(i::Int) ->
  --   forM_ [0..2*n+1] $ \(j::Int) -> do
  --     let (x,y,z) = (i,j, (^ 2) . realPart . abs $ plot R.! (R.Z R.:. i R.:. j))
  --     putStrLn $ show i ++ ", " ++ show j ++ ", " ++ show z
  -- print $ toMatlabFunction $ tfs IM.! 13
  -- norm <- l2Norm plot
  -- print norm

  let ps = IM.elems plots
  norms <- forM ps $ \p -> l2Norm p
  print norms


{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE DeriveGeneric       #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE QuasiQuotes         #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE UndecidableInstances      #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE AllowAmbiguousTypes       #-}
{-# LANGUAGE UndecidableSuperClasses #-}
module Frames.Regression where

import qualified Frames.Misc                  as FU
import qualified Frames.VegaLite.Utils         as FV
import qualified Math.Regression.Regression    as MR
import qualified Math.Regression.LeastSquares  as MR
import qualified Polysemy                      as P
import qualified Knit.Effect.Logger            as Log

import qualified Lucid                         as LH
import qualified Text.Blaze.Html               as BH
import qualified Control.Foldl                 as FL
import qualified Data.List                     as List
import qualified Data.Text                     as T
import qualified Data.Vinyl                    as V
import qualified Data.Vinyl.Functor            as V
import qualified Data.Vinyl.TypeLevel          as V
import           Data.Profunctor               as PF
import qualified Frames                        as F
import qualified Frames.Melt                   as F
import qualified Data.Vector.Storable          as VS

import qualified MachineLearning               as ML
import qualified MachineLearning.Regression    as ML
import qualified Numeric.LinearAlgebra         as LA
import           Numeric.LinearAlgebra.Data     ( R )
import qualified Statistics.Types              as S

import           GHC.TypeLits                   ( Symbol )

type Unweighted = "unweighted" F.:-> Void

type family NonVoidField (f :: (Symbol, Type)) :: Bool where
  NonVoidField '(a,Void) = 'False
  NonVoidField _         = 'True

class BoolVal (w :: Bool) where
  asBool :: Bool

instance BoolVal 'True where
  asBool = True

instance BoolVal 'False where
  asBool = False

data FrameRegressionResult (y :: (Symbol, Type)) (wc :: Bool) (xs :: [(Symbol, Type)]) (w :: (Symbol,Type)) (rs :: [(Symbol, Type)]) where
  FrameUnweightedRegressionResult :: (w ~ Unweighted)
    => MR.RegressionResult R -> FrameRegressionResult y wc xs Unweighted rs
  FrameWeightedRegressionResult :: (V.KnownField w, F.ElemOf rs w)
    => (V.Snd w -> R) -> MR.RegressionResult R -> FrameRegressionResult y wc xs w rs

realRecToList
  :: ( RealFrac b
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => F.Record as
  -> [b]
realRecToList =
  V.recordToList
    . V.rmap (\(V.Compose (V.Dict x)) -> V.Const $ realToFrac x)
    . V.reifyConstraint @Real

instance ( MR.Predictor S.NormalErr a (LA.Vector a) (MR.RegressionResult R)
         , BoolVal wc
         , LA.Element a
         , LA.Numeric a
         , RealFloat a
         , xs F.⊆ rs
         , V.RMap xs
         , V.RecordToList xs
         , V.ReifyConstraint Real F.ElField xs
         , V.NatToInt (V.RLength xs)
         ) => MR.Predictor S.NormalErr a (F.Record rs) (FrameRegressionResult y wc xs w as) where
  predict frr r =
    let rr = regressionResult frr
        addBias l = if asBool @wc then 1 : l else l -- if we regressed with a constant term, we need to put it into the xs for prediction
        va :: LA.Vector a = VS.fromList $ addBias $ realRecToList (F.rcast @xs r)
    in MR.predict rr va

regressionResult :: FrameRegressionResult y wc xs w rs -> MR.RegressionResult R
regressionResult (FrameUnweightedRegressionResult x) = x
regressionResult (FrameWeightedRegressionResult _ x) = x

withConstant
  :: forall y wc xs w rs
   . (BoolVal wc)
  => FrameRegressionResult y wc xs w rs
  -> Bool
withConstant _ = asBool @wc

weightedRegression
  :: forall y wc xs w rs
   . (BoolVal (NonVoidField w))
  => FrameRegressionResult y wc xs w rs
  -> Bool
weightedRegression _ = asBool @(NonVoidField w)

-- make X, y from the data
prepRegression
  :: forall y as f rs
   . ( Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> (LA.Matrix R, LA.Vector R)
prepRegression dat =
  let
    nCols  = V.natToInt @(V.RLength as)
    yListF = PF.dimap (realToFrac . F.rgetField @y) List.reverse FL.list
    toListOfDoubles :: F.Record as -> [Double] = realRecToList
    mListF = PF.dimap (toListOfDoubles . F.rcast)
                      (List.concat . List.reverse)
                      FL.list
    (yList, mList) = FL.fold ((,) <$> yListF <*> mListF) dat
  in
    (LA.matrix nCols mList, LA.vector yList)

-- make X, y, w from the data
prepWeightedRegression
  :: forall y as w f rs
   . ( Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , F.ElemOf rs w
     , FU.RealField w
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> (LA.Matrix R, LA.Vector R, LA.Vector R)
prepWeightedRegression dat =
  let
    nCols  = V.natToInt @(V.RLength as)
    yListF = PF.dimap (realToFrac . F.rgetField @y) List.reverse FL.list
    wListF = PF.dimap (realToFrac . F.rgetField @w) List.reverse FL.list
    toListOfDoubles :: F.Record as -> [Double] = realRecToList
    mListF = PF.dimap (toListOfDoubles . F.rcast)
                      (List.concat . List.reverse)
                      FL.list
    (yList, mList, wList) = FL.fold ((,,) <$> yListF <*> mListF <*> wListF) dat
  in
    (LA.matrix nCols mList, LA.vector yList, LA.vector wList)

prettyPrintRegressionResult
  :: forall y wc as w rs
   . ( F.ColumnHeaders '[y]
     , F.ColumnHeaders '[w]
     , F.ColumnHeaders as
     , BoolVal wc
     )
  => (T.Text -> T.Text -> T.Text)
  -> FrameRegressionResult y wc as w rs
  -> S.CL Double
  -> T.Text
prettyPrintRegressionResult headerF res cl =
  let yName   = FV.colName @y
      wName   = FV.colName @w
      xNames' = T.pack <$> F.columnHeaders (Proxy :: Proxy (F.Record as))
      xNames  = if withConstant res then "intercept" : xNames' else xNames'
  in  MR.prettyPrintRegressionResult (headerF yName wName)
                                     xNames
                                     (regressionResult res)
                                     cl

prettyPrintRegressionResultLucid
  :: forall y wc as w rs
   . ( F.ColumnHeaders '[y]
     , F.ColumnHeaders '[w]
     , F.ColumnHeaders as
     , BoolVal wc
     )
  => (T.Text -> T.Text -> T.Text)
  -> FrameRegressionResult y wc as w rs
  -> S.CL Double
  -> LH.Html ()
prettyPrintRegressionResultLucid headerF res cl =
  let yName   = FV.colName @y
      wName   = FV.colName @w
      xNames' = T.pack <$> F.columnHeaders (Proxy :: Proxy (F.Record as))
      xNames  = if withConstant res then "intercept" : xNames' else xNames'
  in  MR.prettyPrintRegressionResultLucid (headerF yName wName)
                                          xNames
                                          (regressionResult res)
                                          cl

prettyPrintRegressionResultBlaze
  :: forall y wc as w rs
   . ( F.ColumnHeaders '[y]
     , F.ColumnHeaders '[w]
     , F.ColumnHeaders as
     , BoolVal wc
     )
  => (T.Text -> T.Text -> T.Text)
  -> FrameRegressionResult y wc as w rs
  -> S.CL Double
  -> BH.Html
prettyPrintRegressionResultBlaze headerF res cl =
  let yName   = FV.colName @y
      wName   = FV.colName @w
      xNames' = T.pack <$> F.columnHeaders (Proxy :: Proxy (F.Record as))
      xNames  = if withConstant res then "intercept" : xNames' else xNames'
  in  MR.prettyPrintRegressionResultBlaze (headerF yName wName)
                                          xNames
                                          (regressionResult res)
                                          cl

keyRecordText
  :: (V.ReifyConstraint Show F.ElField ks, V.RecordToList ks, V.RMap ks)
  => F.Record ks
  -> T.Text
keyRecordText keyRec =
  let keyValuesAsText rks =
        V.recordToList
          . V.rmap (\(V.Compose (V.Dict x)) -> V.Const $ T.pack $ show x)
          $ V.reifyConstraint @Show rks
  in  T.intercalate "; " $ keyValuesAsText keyRec

prettyPrintRegressionResults
  :: forall y wc as w rs k f a
   . ( F.ColumnHeaders '[y]
     , F.ColumnHeaders '[w]
     , F.ColumnHeaders as
     , BoolVal wc
     , BoolVal (NonVoidField w)
     , Foldable f
     , Monoid a
     )
  => (k -> T.Text)
  -> f (k, FrameRegressionResult y wc as w rs)
  -> S.CL R
  -> (  (T.Text -> T.Text -> T.Text)
     -> FrameRegressionResult y wc as w rs
     -> S.CL Double
     -> a
     )
  -> a
  -> a
prettyPrintRegressionResults keyText keyed cl printOne sepEach =
  let headerF res key yName wName = if weightedRegression res
                                    then
                                      "Explaining "
                                        <> yName
                                        <> " ("
                                        <> keyText key
                                        <> "; weights from "
                                        <> wName
                                        <> ")"
                                    else "Explaining " <> yName <> " (" <> keyText key <> ")"
  in  FL.fold
        (FL.Fold (\t (rk, res) -> t <> printOne (headerF res rk) res cl)
                 sepEach
                 id
        )
        keyed

-- explain y in terms of as
leastSquaresByMinimization
  :: forall y as rs f
   . ( Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => Bool
  -> [R]
  -> f (F.Record rs)
  -> [R]
leastSquaresByMinimization wc guess dat =
  let (mX, y)       = prepRegression @y @as dat
      mX1           = if wc then ML.addBiasDimension mX else mX
      (solution, _) = ML.minimize (ML.ConjugateGradientFR 0.1 0.1)
                                  ML.LeastSquares
                                  0.001
                                  20
                                  ML.RegNone
                                  mX1
                                  y
                                  (LA.fromList guess)
  in  LA.toList solution

ordinaryLeastSquares
  :: forall effs y wc as rs f
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , BoolVal wc
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as Unweighted rs)
ordinaryLeastSquares dat = do
  let (mA, vB)  = prepRegression @y @as dat
      withConst = asBool @wc
  FrameUnweightedRegressionResult <$> MR.ordinaryLS withConst mA vB

weightedLeastSquares
  :: forall y wc as w rs f effs
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , F.ElemOf rs w
     , FU.RealField w
     , BoolVal wc
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as w rs)
weightedLeastSquares dat = do
  let (mA, vB, vW) = prepWeightedRegression @y @as @w dat
      withConst    = asBool @wc
  FrameWeightedRegressionResult realToFrac <$> MR.weightedLS withConst mA vB vW

-- special case when weights come from observations being population averages of different populations
popWeightedLeastSquares
  :: forall effs y wc as w rs f
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , F.ElemOf rs w
     , FU.RealField w
     , BoolVal wc
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as w rs)
popWeightedLeastSquares dat = do
  let (mA, vB, vW) = prepWeightedRegression @y @as @w dat
      withConst    = asBool @wc
      vWpop        = LA.cmap sqrt vW -- this is the correct weight for population average, the sqrt of the number averaged in that sample
  FrameWeightedRegressionResult (sqrt . realToFrac)
    <$> MR.weightedLS withConst mA vB vWpop

-- special case when we know residuals are heteroscedastic with variances proportional to given numbers
varWeightedLeastSquares
  :: forall effs y wc as w rs f
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , F.ElemOf rs w
     , FU.RealField w
     , BoolVal wc
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as w rs)
varWeightedLeastSquares dat = do
  let (mA, vB, vW) = prepWeightedRegression @y @as @w dat
      withConst    = asBool @wc
      vWvar        = LA.cmap (\x -> 1 / sqrt x) vW -- this is the correct weight for given variance
  FrameWeightedRegressionResult (\x -> 1 / sqrt (realToFrac x))
    <$> MR.weightedLS withConst mA vB vWvar

totalLeastSquares
  :: forall effs y wc as rs f
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , BoolVal wc
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as Unweighted rs)
totalLeastSquares dat = do
  let (mA, vB)  = prepRegression @y @as dat
      withConst = asBool @wc
  FrameUnweightedRegressionResult <$> MR.totalLS withConst mA vB

weightedTLS
  :: forall effs y wc as w rs f
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , BoolVal wc
     , F.ElemOf rs w
     , FU.RealField w
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as w rs)
weightedTLS dat = do
  let (mA, vB, vW) = prepWeightedRegression @y @as @w dat
      withConst    = asBool @wc
  FrameWeightedRegressionResult realToFrac <$> MR.weightedTLS withConst mA vB vW

popWeightedTLS
  :: forall effs y wc as w rs f
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , BoolVal wc
     , F.ElemOf rs w
     , FU.RealField w
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as w rs)
popWeightedTLS dat = do
  let (mA, vB, vW) = prepWeightedRegression @y @as @w dat
      withConst    = asBool @wc
      vWpop        = LA.cmap sqrt vW -- this is the correct weight for population average, the sqrt of the number averaged in that sample
  FrameWeightedRegressionResult (sqrt . realToFrac)
    <$> MR.weightedTLS withConst mA vB vWpop

varWeightedTLS
  :: forall effs y wc as w rs f
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , BoolVal wc
     , F.ElemOf rs w
     , FU.RealField w
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as w rs)
varWeightedTLS dat = do
  let (mA, vB, vW) = prepWeightedRegression @y @as @w dat
      withConst    = asBool @wc
      vWvar        = LA.cmap (\x -> 1 / sqrt x) vW -- this is the correct weight for given variance
  FrameWeightedRegressionResult (\x -> 1 / sqrt (realToFrac x))
    <$> MR.weightedTLS withConst mA vB vWvar

{-
-- this is sort of ugly but I think it will work.  And it has a pretty narrow purpose
-- But what will we do when we want logistic regression or whatever?
class (V.AllConstrained RealFrac (V.Unlabeled rs), Real a) => RealFracRecordFromList a rs where
  recordFromList :: [a] -> Maybe (F.Record rs)

instance Real a => RealFracRecordFromList a '[] where
  recordFromList _ = Just $ V.RNil

instance ( V.AllConstrained RealFrac (V.Unlabeled (r : rs))
         , V.AllConstrained RealFrac (V.Unlabeled rs)
         , RealFracRecordFromList a rs
         , V.KnownField r
         , RealFrac (V.Snd r)
         , Real a) => RealFracRecordFromList a (r : rs) where
  recordFromList [] = Nothing
  recordFromList (a : as) = case recordFromList as of
    Nothing -> Nothing
    Just xs -> let x = (realToFrac a) :: V.Snd r in Just $ x &: xs
-}


module TrainUtils where

import Control.Lens hiding (para)
import System.Random
import Control.Monad
import Prelude hiding (id, (.))
import Numeric.LinearAlgebra.Array
import Numeric.LinearAlgebra.Array.Util

import CategoricDefinitions
import Autodiff.GAD
import Autodiff.Dual
import Autodiff.Additive
import Autodiff.D
import Autodiff.Cont
import Ops
import Para
import TensorUtils
import OnesLike

-------------------------------------------
showNNInfo :: (ArrShow p, ArrShow a, ArrShow b)
    => Int -> a -> b -> LearnerType p a b -> IO ()
showNNInfo n a b nn = do
    putStrLn "-------------------------"
    putStrLn $ "Step " ++ arrShow n
    putStrLn $ "p\n" ++ arrShow (nn ^. p)
    putStrLn $ "a\n" ++ arrShow a
    putStrLn $ "a\n" ++ arrShow b

-- | Supervised learning training
-- Takes in a Learner, input-output pairs and a cost function
-- it partially applies the output to the cost function and composes the result inside learner
trainStepWithCost :: (OnesLike c, _)
    => LearnerType p a b -> (Int, IO (a, b), DType (b, b) c) -> IO (LearnerType p a b)
trainStepWithCost l (step, dataSampler, cost) = do
    (i, o) <- dataSampler
    when (step `mod` 100 == 0) $ showNNInfo step i o l
    let cost' = partiallyApply cost o
        (pGrad, _) = grad (cost' . (l ^. para . fn)) (l ^. p, i)
    return $ l & p .~ (l ^. optimizer) (l ^. p, pGrad)

instance (Random a, Random b) => Random (a, b) where
    random gen1 = let (x, gen2) = random gen1
                      (y, gen3) = random gen2
                  in ((x, y), gen3)

    randomR ((x1, y1), (x2, y2)) gen1 = let (x, gen2) = randomR (x1, x2) gen1
                                            (y, gen3) = randomR (y1, y2) gen2
                                        in ((x, y), gen3)


{-# LANGUAGE ForeignFunctionInterface, GeneralizedNewtypeDeriving #-}

module Numerical.HBLAS.BLAS.FFI.Level1  where

import Foreign.Ptr
import Foreign()
import Foreign.C.Types
import Data.Complex

type AsumFunFFI el res = CInt -> Ptr el -> CInt -> IO res
foreign import ccall unsafe "cblas_sasum" cblas_sasum_unsafe::
  AsumFunFFI Float Float
foreign import ccall unsafe "cblas_dasum" cblas_dasum_unsafe::
  AsumFunFFI Double Double
foreign import ccall unsafe "cblas_scasum" cblas_scasum_unsafe::
  AsumFunFFI (Complex Float) Float
foreign import ccall unsafe "cblas_dzasum" cblas_dzasum_unsafe::
  AsumFunFFI (Complex Double) Double

foreign import ccall "cblas_sasum" cblas_sasum_safe::
  AsumFunFFI Float Float
foreign import ccall "cblas_dasum" cblas_dasum_safe ::
  AsumFunFFI Double Double
foreign import ccall "cblas_scasum" cblas_scasum_safe ::
  AsumFunFFI (Complex Float) Float
foreign import ccall "cblas_dzasum" cblas_dzasum_safe ::
  AsumFunFFI (Complex Double) Double
--Float  cblas_sasum (  CInt n,   Float  *x,   CInt incx);
--Double cblas_dasum (  CInt n,   Double *x,   CInt incx);
--Float  cblas_scasum(  CInt n,   Float  *x,   CInt incx);
--Double cblas_dzasum(  CInt n,   Double *x,   CInt incx);

type AxpyFunFFI scale el = CInt -> scale -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_saxpy" cblas_saxpy_unsafe::
  AxpyFunFFI Float Float
foreign import ccall unsafe "cblas_daxpy" cblas_daxpy_unsafe::
  AxpyFunFFI Double Double
foreign import ccall unsafe "cblas_caxpy" cblas_caxpy_unsafe::
  AxpyFunFFI (Ptr (Complex Float )) (Complex Float )
foreign import ccall unsafe "cblas_zaxpy" cblas_zaxpy_unsafe::
  AxpyFunFFI (Ptr (Complex Double)) (Complex Double)

foreign import ccall "cblas_saxpy" cblas_saxpy_safe::
  AxpyFunFFI Float Float
foreign import ccall "cblas_daxpy" cblas_daxpy_safe::
  AxpyFunFFI Double Double
foreign import ccall "cblas_caxpy" cblas_caxpy_safe::
  AxpyFunFFI (Ptr (Complex Float )) (Complex Float )
foreign import ccall "cblas_zaxpy" cblas_zaxpy_safe::
  AxpyFunFFI (Ptr (Complex Double)) (Complex Double)
--void cblas_saxpy(  CInt n,   Float alpha,   Float *x,   CInt incx, Float *y,   CInt incy);
--void cblas_daxpy(  CInt n,   Double alpha,   Double *x,   CInt incx, Double *y,   CInt incy);
--void cblas_caxpy(  CInt n,   Float *alpha,   Float *x,   CInt incx, Float *y,   CInt incy);
--void cblas_zaxpy(  CInt n,   Double *alpha,   Double *x,   CInt incx, Double *y,   CInt incy);

type CopyFunFFI el = CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_scopy" cblas_scopy_unsafe ::
    CopyFunFFI Float
foreign import ccall unsafe "cblas_dcopy" cblas_dcopy_unsafe ::
    CopyFunFFI Double
foreign import ccall unsafe "cblas_ccopy" cblas_ccopy_unsafe ::
    CopyFunFFI (Complex Float)
foreign import ccall unsafe "cblas_zcopy" cblas_zcopy_unsafe ::
    CopyFunFFI (Complex Double)

foreign import ccall "cblas_scopy" cblas_scopy_safe ::
    CopyFunFFI Float
foreign import ccall "cblas_dcopy" cblas_dcopy_safe ::
    CopyFunFFI Double
foreign import ccall "cblas_ccopy" cblas_ccopy_safe ::
    CopyFunFFI (Complex Float)
foreign import ccall "cblas_zcopy" cblas_zcopy_safe ::
    CopyFunFFI (Complex Double)
--void cblas_scopy(  CInt n,   Float *x,   CInt incx, Float *y,   CInt incy);
--void cblas_dcopy(  CInt n,   Double *x,   CInt incx, Double *y,   CInt incy);
--void cblas_ccopy(  CInt n,   Float *x,   CInt incx, Float *y,   CInt incy);
--void cblas_zcopy(  CInt n,   Double *x,   CInt incx, Double *y,   CInt incy);

--dot products
type NoScalarDotFunFFI el res = CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO res
-- type ScalarDotFunFFI el res = CInt -> el -> Ptr el -> CInt -> Ptr el -> CInt -> IO res
type SdsdotFortranFunFFI el res = Ptr CInt -> Ptr el -> Ptr el -> Ptr CInt -> Ptr el -> Ptr CInt -> IO res
foreign import ccall "sdsdot_" cblas_sdsdot_safe :: SdsdotFortranFunFFI Float Float
foreign import ccall "cblas_dsdot" cblas_dsdot_safe :: NoScalarDotFunFFI Float Double
foreign import ccall "cblas_sdot" cblas_sdot_safe :: NoScalarDotFunFFI Float Float
foreign import ccall "cblas_ddot" cblas_ddot_safe :: NoScalarDotFunFFI Double Double

foreign import ccall unsafe "sdsdot_" cblas_sdsdot_unsafe :: SdsdotFortranFunFFI Float Float
foreign import ccall unsafe "cblas_dsdot" cblas_dsdot_unsafe :: NoScalarDotFunFFI Float Double
foreign import ccall unsafe "cblas_sdot" cblas_sdot_unsafe :: NoScalarDotFunFFI Float Float
foreign import ccall unsafe "cblas_ddot" cblas_ddot_unsafe :: NoScalarDotFunFFI Double Double
--Float  cblas_sdsdot(  CInt n,   Float alpha,   Float *x,   CInt incx,   Float *y,   CInt incy);
--Double cblas_dsdot (  CInt n,   Float *x,   CInt incx,   Float *y,   CInt incy);
--Float  cblas_sdot(  CInt n,   Float  *x,   CInt incx,   Float  *y,   CInt incy);
--Double cblas_ddot(  CInt n,   Double *x,   CInt incx,   Double *y,   CInt incy);

--complex dot products that CONJUGATE end in C, others end in U
type ComplexDotFunFFI el = CInt -> Ptr el -> CInt -> Ptr el -> CInt -> Ptr el -> IO ()
foreign import ccall "cblas_cdotu_sub" cblas_cdotu_safe :: ComplexDotFunFFI (Complex Float)
foreign import ccall "cblas_cdotc_sub" cblas_cdotc_safe :: ComplexDotFunFFI (Complex Float)
foreign import ccall "cblas_zdotu_sub" cblas_zdotu_safe :: ComplexDotFunFFI (Complex Double)
foreign import ccall "cblas_zdotc_sub" cblas_zdotc_safe :: ComplexDotFunFFI (Complex Double)

foreign import ccall unsafe "cblas_cdotu_sub" cblas_cdotu_unsafe :: ComplexDotFunFFI (Complex Float)
foreign import ccall unsafe "cblas_cdotc_sub" cblas_cdotc_unsafe :: ComplexDotFunFFI (Complex Float)
foreign import ccall unsafe "cblas_zdotu_sub" cblas_zdotu_unsafe :: ComplexDotFunFFI (Complex Double)
foreign import ccall unsafe "cblas_zdotc_sub" cblas_zdotc_unsafe :: ComplexDotFunFFI (Complex Double)
--void cblas_cdotu(CInt n, void *x, CInt incx, void *y, CInt incy, void *res);
--void cblas_cdotc(CInt n, void *x, CInt incx, void *y, CInt incy, void *res);
--void cblas_zdotu(CInt n, void *x, CInt incx, void *y, CInt incy, void *res);
--void cblas_zdotc(CInt n, void *x, CInt incx, void *y, CInt incy, void *res);

--Computes the Euclidean norm of a vector.
type Nrm2FunFFI el res = CInt -> Ptr el -> CInt -> IO res
foreign import ccall "cblas_snrm2" cblas_snrm2_safe :: Nrm2FunFFI Float Float
foreign import ccall "cblas_dnrm2" cblas_dnrm2_safe :: Nrm2FunFFI Double Double
foreign import ccall "cblas_scnrm2" cblas_scnrm2_safe :: Nrm2FunFFI (Complex Float) Float
foreign import ccall "cblas_dznrm2" cblas_dznrm2_safe :: Nrm2FunFFI (Complex Double) Double

foreign import ccall unsafe "cblas_snrm2" cblas_snrm2_unsafe :: Nrm2FunFFI Float Float
foreign import ccall unsafe "cblas_dnrm2" cblas_dnrm2_unsafe :: Nrm2FunFFI Double Double
foreign import ccall unsafe "cblas_scnrm2" cblas_scnrm2_unsafe :: Nrm2FunFFI (Complex Float) Float
foreign import ccall unsafe "cblas_dznrm2" cblas_dznrm2_unsafe :: Nrm2FunFFI (Complex Double) Double
--Float  cblas_snrm2 (  CInt N,   Float  *X,   CInt incX);
--Double cblas_dnrm2 (  CInt N,   Double *X,   CInt incX);
--Float  cblas_scnrm2(  CInt N,   void  *X,   CInt incX);
--Double cblas_dznrm2(  CInt N,   void  *X,   CInt incX);

--Performs rotation of points in the plane.
type RotFunFFI el = CInt -> Ptr el -> CInt -> Ptr el -> CInt -> el -> el -> IO ()
foreign import ccall "cblas_srot" cblas_srot_safe :: RotFunFFI Float
foreign import ccall "cblas_drot" cblas_drot_safe :: RotFunFFI Double

foreign import ccall unsafe "cblas_srot" cblas_srot_unsafe :: RotFunFFI Float
foreign import ccall unsafe "cblas_drot" cblas_drot_unsafe :: RotFunFFI Double
--void cblas_srot(  CInt N, Float *X,   CInt incX, Float *Y,   CInt incY,   Float c,   Float s);
--void cblas_drot(  CInt N, Double *X,   CInt incX, Double *Y,   CInt incY,   Double c,   Double  s);

type RotgFunFFI el = Ptr el -> Ptr el -> Ptr el -> Ptr el -> IO ()
foreign import ccall "cblas_srotg" cblas_srotg_safe :: RotgFunFFI Float
foreign import ccall "cblas_drotg" cblas_drotg_safe :: RotgFunFFI Double

foreign import ccall unsafe "cblas_srotg" cblas_srotg_unsafe :: RotgFunFFI Float
foreign import ccall unsafe "cblas_drotg" cblas_drotg_unsafe :: RotgFunFFI Double
--void cblas_srotg(Float *a, Float *b, Float *c, Float *s);
--void cblas_drotg(Double *a, Double *b, Double *c, Double *s);

type RotmFunFFI el = CInt -> Ptr el -> CInt -> Ptr el -> CInt -> Ptr el -> IO ()
foreign import ccall "cblas_srotm" cblas_srotm_safe :: RotmFunFFI Float
foreign import ccall "cblas_drotm" cblas_drotm_safe :: RotmFunFFI Double

foreign import ccall unsafe "cblas_srotm" cblas_srotm_unsafe :: RotmFunFFI Float
foreign import ccall unsafe "cblas_drotm" cblas_drotm_unsafe :: RotmFunFFI Double
--void cblas_srotm(  CInt N, Float *X,   CInt incX, Float *Y,   CInt incY,   Float *P);
--void cblas_drotm(  CInt N, Double *X,   CInt incX, Double *Y,   CInt incY,   Double *P);

type RotmgFunFFI el = Ptr el -> Ptr el -> Ptr el -> el -> Ptr el -> IO ()
foreign import ccall "cblas_srotmg" cblas_srotmg_safe :: RotmgFunFFI Float
foreign import ccall "cblas_drotmg" cblas_drotmg_safe :: RotmgFunFFI Double

foreign import ccall unsafe "cblas_srotmg" cblas_srotmg_unsafe :: RotmgFunFFI Float
foreign import ccall unsafe "cblas_drotmg" cblas_drotmg_unsafe :: RotmgFunFFI Double
--void cblas_srotmg(Float *d1, Float *d2, Float *b1,   Float b2, Float *P);
--void cblas_drotmg(Double *d1, Double *d2, Double *b1,   Double b2, Double *P);

type ScalFunFFI scale el = CInt -> scale -> Ptr el -> CInt -> IO ()
foreign import ccall "cblas_sscal" cblas_sscal_safe :: ScalFunFFI Float Float
foreign import ccall "cblas_dscal" cblas_dscal_safe :: ScalFunFFI Double Double
foreign import ccall "cblas_cscal" cblas_cscal_safe :: ScalFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall "cblas_zscal" cblas_zscal_safe :: ScalFunFFI (Ptr (Complex Double)) (Complex Double)
foreign import ccall "cblas_csscal" cblas_csscal_safe :: ScalFunFFI Float (Complex Float)
foreign import ccall "cblas_zdscal" cblas_zdscal_safe :: ScalFunFFI Double (Complex Double)

foreign import ccall unsafe "cblas_sscal" cblas_sscal_unsafe :: ScalFunFFI Float Float
foreign import ccall unsafe "cblas_dscal" cblas_dscal_unsafe :: ScalFunFFI Double Double
foreign import ccall unsafe "cblas_cscal" cblas_cscal_unsafe :: ScalFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall unsafe "cblas_zscal" cblas_zscal_unsafe :: ScalFunFFI (Ptr (Complex Double)) (Complex Double)
foreign import ccall unsafe "cblas_csscal" cblas_csscal_unsafe :: ScalFunFFI Float (Complex Float)
foreign import ccall unsafe "cblas_zdscal" cblas_zdscal_unsafe :: ScalFunFFI Double (Complex Double)
--void cblas_sscal(  CInt N,   Float alpha, Float *X,   CInt incX);
--void cblas_dscal(  CInt N,   Double alpha, Double *X,   CInt incX);
--void cblas_cscal(  CInt N,   Float *alpha, Float *X,   CInt incX);
--void cblas_zscal(  CInt N,   Double *alpha, Double *X,   CInt incX);
--void cblas_csscal(  CInt N,   Float alpha, Float *X,   CInt incX);
--void cblas_zdscal(  CInt N,   Double alpha, Double *X,   CInt incX);

type SwapFunFFI el = CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall "cblas_sswap" cblas_sswap_safe :: SwapFunFFI Float
foreign import ccall "cblas_dswap" cblas_dswap_safe :: SwapFunFFI Double
foreign import ccall "cblas_cswap" cblas_cswap_safe :: SwapFunFFI (Complex Float)
foreign import ccall "cblas_zswap" cblas_zswap_safe :: SwapFunFFI (Complex Double)

foreign import ccall unsafe "cblas_sswap" cblas_sswap_unsafe :: SwapFunFFI Float
foreign import ccall unsafe "cblas_dswap" cblas_dswap_unsafe :: SwapFunFFI Double
foreign import ccall unsafe "cblas_cswap" cblas_cswap_unsafe :: SwapFunFFI (Complex Float)
foreign import ccall unsafe "cblas_zswap" cblas_zswap_unsafe :: SwapFunFFI (Complex Double)
--void cblas_sswap(  CInt n, Float *x,   CInt incx, Float *y,   CInt incy);
--void cblas_dswap(  CInt n, Double *x,   CInt incx, Double *y,   CInt incy);
--void cblas_cswap(  CInt n, Float *x,   CInt incx, Float *y,   CInt incy);
--void cblas_zswap(  CInt n, Double *x,   CInt incx, Double *y,   CInt incy);

type IamaxFunFFI el = CInt -> Ptr el -> CInt -> IO CInt
foreign import ccall "cblas_isamax" cblas_isamax_safe :: IamaxFunFFI Float
foreign import ccall "cblas_idamax" cblas_idamax_safe :: IamaxFunFFI Double
foreign import ccall "cblas_icamax" cblas_icamax_safe :: IamaxFunFFI (Complex Float)
foreign import ccall "cblas_izamax" cblas_izamax_safe :: IamaxFunFFI (Complex Double)

foreign import ccall unsafe "cblas_isamax" cblas_isamax_unsafe :: IamaxFunFFI Float
foreign import ccall unsafe "cblas_idamax" cblas_idamax_unsafe :: IamaxFunFFI Double
foreign import ccall unsafe "cblas_icamax" cblas_icamax_unsafe :: IamaxFunFFI (Complex Float)
foreign import ccall unsafe "cblas_izamax" cblas_izamax_unsafe :: IamaxFunFFI (Complex Double)
--CBLAS_INDEX cblas_isamax(  CInt n,   Float  *x,   CInt incx);
--CBLAS_INDEX cblas_idamax(  CInt n,   Double *x,   CInt incx);
--CBLAS_INDEX cblas_icamax(  CInt n,   Float  *x,   CInt incx);
--CBLAS_INDEX cblas_izamax(  CInt n,   Double *x,   CInt incx);

{-
these aren't provided by Accelerate frameowkr on OSX, not sure about other platforms

but easy to write portable substitute I think?
-}
--type IaminFunFFI el = CInt -> Ptr el -> CInt -> IO CInt
--foreign import ccall "cblas_isamin" cblas_isamin_safe :: IaminFunFFI Float
--foreign import ccall "cblas_idamin" cblas_idamin_safe :: IaminFunFFI Double
--foreign import ccall "cblas_icamin" cblas_icamin_safe :: IaminFunFFI (Complex Float)
--foreign import ccall "cblas_izamin" cblas_izamin_safe :: IaminFunFFI (Complex Double)

--foreign import ccall unsafe "cblas_isamin" cblas_isamin_unsafe :: IaminFunFFI Float
--foreign import ccall unsafe "cblas_idamin" cblas_idamin_unsafe :: IaminFunFFI Double
--foreign import ccall unsafe "cblas_icamin" cblas_icamin_unsafe :: IaminFunFFI (Complex Float)
--foreign import ccall unsafe "cblas_izamin" cblas_izamin_unsafe :: IaminFunFFI (Complex Double)

--CBLAS_INDEX cblas_isamin(  CInt n,   Float  *x,   CInt incx);
--CBLAS_INDEX cblas_idamin(  CInt n,   Double *x,   CInt incx);
--CBLAS_INDEX cblas_icamin(  CInt n,   Float  *x,   CInt incx);
--CBLAS_INDEX cblas_izamin(  CInt n,   Double *x,   CInt incx);


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

module HVX.BlackBoxTests.ManyvarTestSpec where

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

import HVX
import HVX.Internal.TestUtil

-- Problem definition.
nx = 2
ny = 3
nz = 4
a = EConst $ (3><nx)
    ([-0.1022 , 0.3129
    , -0.2414 , -0.8649
    , 0.3192 , -0.0301 ] :: [Double])
b = EConst $ (3><ny)
    ([-0.1649 , 1.1093 , -1.2141
    , 0.6277 , -0.8637 , -1.1135
    , 1.0933 , 0.0774 , -0.0068 ] :: [Double])
c = EConst $ (3><nz)
    ([1.5326 , -0.2256 , 0.0326 , 1.5442
    , -0.7697 , 1.1174 , 0.5525 , 0.0859
    , 0.3714 , -1.0891 , 1.1006 , -1.4916 ] :: [Double])
d = EConst $ (3><1)
    ([-0.6156 , 0.7481 , -0.1924] :: [Double])

x = EVar "x"
y = EVar "y"
z = EVar "z"
const10nx = EConst $ scale 10 $ ones nx 1
const10ny = EConst $ scale 10 $ ones ny 1
const01nz = EConst $ scale 0.1 $ ones nz 1

subgradAns = subgradMaximize
  ( neg (norm 2 (a *~ x +~ b *~ y +~ c *~ z +~ d))
    +~ neg (norm 1 x)
    +~ neg (norm 4.2 y) )
  [ const01nz >=~ hexp z
    , powBaseP01 0.25 x >=~ const10nx
    , powBaseP1InfNotInt 1.75 y <=~ const10ny ]
  (decNonSumStep 100.0) 10
  [("x", zeros nx 1), ("y", zeros ny 1), ("z", zeros nz 1)]

ellipsoidAns = ellipsoidMaximize
  ( neg (norm 2 (a *~ x +~ b *~ y +~ c *~ z +~ d))
    +~ neg (norm 1 x)
    +~ neg (norm 4.2 y) )
  [ const01nz >=~ hexp z
    , powBaseP01 0.25 x >=~ const10nx
    , powBaseP1InfNotInt 1.75 y <=~ const10ny ]
  [("x", nx), ("y", ny), ("z", nz)]
  1e-16 1e10

(subgradVars, subgradOptval) = subgradAns

(ellipsoidVars, ellipsoidOptval, ellipsoidUBound) = ellipsoidAns

-- CVX's results.
cvxOptval = -29051

-- Verify that HVX matches CVX.
-- TODO(mh): This test fails because primitives that generate implicit
-- constraints are currently unsupported. When support is added for them, the
-- test will be added back in. (2014-06-04)
spec :: Spec
spec =
  describe "Placeholder test" $ do
    it "placeholder description" $
      True `shouldBe` True
--  describe "Verify that HVX matches CVX for triple variable stress test" $ do
--    it "Verify that HVX subgrad matches CVX for huber/berhu" $
--      subgradOptval `shouldSatisfy` fpequalsApprox cvxOptval
--    it "Verify that HVX ellipsoid matches CVX for huber/berhu" $
--      ellipsoidOptval `shouldSatisfy` fpequalsApprox cvxOptval
--
--main :: IO ()
--main = hspec spec


module Data.Csv.HMatrix
    (
      -- * Usage example
      -- $example

      -- * Information
      -- $info

      decodeMatrix
    , decodeMatrixWith
    , encodeMatrix
    , encodeMatrixWith
    ) where

import Data.Csv
import Numeric.LinearAlgebra.HMatrix
import qualified Data.Vector as V
import Data.ByteString.Lazy (ByteString)
import qualified Data.ByteString.Char8 as C
import Data.Char (ord)

-- $example
--
-- For decoding we have to specify, if there is a header. This can be done with the HasHeader/NoHeader type:
--
-- > >>> decodeMatrix NoHeader "1.0,2.0,3.0\r\n4.0,5.0,6.0\r\n7.0,8.0,9.0\r\n"
-- > (3><3)
-- > [ 1.0, 2.0, 3.0
-- > , 4.0, 5.0, 6.0
-- > , 7.0, 8.0, 9.0 ]
--
-- Cassava, which is used for parsing the .csv files uses overloaded string literals. If you try this in ghci, make sure to start it with the right flag:
--
-- > >>> ghci -XOverloadedStrings
--
-- Encoding a file works pretty much the same, except that we do not need to specify a header.
--
-- > >>> encodeMatrix $ matrix 3 [1,2,3,4,5,6,7,8,9]
-- > "1.0,2.0,3.0\r\n4.0,5.0,6.0\r\n7.0,8.0,9.0\r\n"

-- $info
-- If you want to help improve this library, feel free to file an issue or send a pull-request on github. Every feedback is appreciated.
-- As of now only matrices of type Double are supported.

-- | Decodes a matrix.
decodeMatrix :: HasHeader       -- ^ From Data.Csv: specify if the CSV string has a header
             -> ByteString      -- ^ The 'ByteString' containing the CSVs
             -> Matrix Double   -- ^ The parsed 'Matrix'
decodeMatrix header = decodeMatrixWith header ','

-- | Decodes a matrix from ByteString and additionally allow
-- to specify the delimter which was used.
decodeMatrixWith :: HasHeader       -- ^ From Data.Csv: specify if the CSV string has a header
                 -> Char            -- ^ The delimiter
                 -> ByteString      -- ^ The 'ByteString' containing the CSVs
                 -> Matrix Double   -- ^ The parsed 'Matrix'
decodeMatrixWith header del s =
    case decodeWith opt header s of
        Left err -> error err
        Right v  -> fromLists . V.toList . V.map V.toList $ v
    where opt = defaultDecodeOptions { decDelimiter = fromIntegral (ord del) }

rowToRecord :: [Double] -> Record
rowToRecord x = record $ map (C.pack . show) x

-- | Encodes a matrix with comma as delimiter.
encodeMatrix :: Matrix Double   -- ^ The 'Matrix' to encode
             -> ByteString      -- ^ The resulting 'ByteString'
encodeMatrix = encodeMatrixWith ','

-- | Encodes a matrix but allows to specify a delimiter.
encodeMatrixWith :: Char            -- ^ The delimiter for separating the values
                 -> Matrix Double   -- ^ The 'Matrix' to encode
                 -> ByteString      -- ^ The resulting 'ByteString'
encodeMatrixWith del m = encodeWith opt s
    where opt = defaultEncodeOptions { encDelimiter = fromIntegral (ord del) }
          s = map rowToRecord $ toLists m


module Foreign.Storable.OrphansSpec (main, spec) where

import Test.Hspec
import Data.Complex
import Data.Orphans ()
import Data.Ratio
import Foreign.Storable

main :: IO ()
main = hspec spec

spec :: Spec
spec = do
  describe "Storable Complex instance" $ do
    it "has twice the sizeOf its realPart" $ do
      sizeOf (undefined :: Complex Double) `shouldBe` 2*sizeOf (1 :: Double)
    it "has the alignment of its realPart" $ do
      alignment (undefined :: Complex Double) `shouldBe` alignment (1 :: Double)

  describe "Storable Ratio instance" $ do
    it "has twice the sizeOf its parameterized type" $ do
      sizeOf (undefined :: Ratio Int) `shouldBe` 2*sizeOf (1 :: Int)
    it "has the alignment of its parameterized type" $ do
      alignment (undefined :: Ratio Int) `shouldBe` alignment (1 :: Int)


{-# LANGUAGE BangPatterns #-}
module IllusoryContourShape where

import           Control.Monad               as M
import           Data.Array.Repa             as R
import           Data.Binary
import           Data.Complex
import           Data.List                   as L
import           Data.Vector.Generic         as VG
import           FokkerPlanck
import           Image.IO
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Utils.Array
import           Utils.Time
import qualified Data.Array.IArray as IA
import           FokkerPlanck.GreensFunction

main = do
  args@(gpuIDStr:numPointStr:deltaXStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:cutoffStr:deltaLogStr:taoStr:numTrailStr:maxTrailStr:phiFreqsStr:rhoFreqsStr:thetaFreqsStr:scaleFreqsStr:initScaleStr:shape2DStr:histFilePath:writeFlagStr:numIterationStr:suffix:stdStr:stdGStr:numThreadStr:_) <-
    getArgs
  let gpuID = read gpuIDStr :: [Int]
      numPoint = read numPointStr :: Int
      deltaX = read deltaXStr :: Double
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      phiFreq = read phiFreqsStr :: Double
      phiFreqs = [-phiFreq .. phiFreq]
      rhoFreq = read rhoFreqsStr :: Double
      rhoFreqs = [-rhoFreq .. rhoFreq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Double
      scaleFreqs = [-scaleFreq .. scaleFreq]
      initScale = read initScaleStr :: Double
      shape2D@(Points _ minDist _) = read shape2DStr :: Points Shape2D
      numThread = read numThreadStr :: Int
      folderPath =
        "output/test/IllusoryContourShape" </> (takeBaseName histFilePath) L.++
        "_" L.++
        cutoffStr L.++
        "_" L.++
        deltaXStr
      maxScale = read maxScaleStr :: Double
      cutoff = read cutoffStr :: Double
      halfLogPeriod = log maxScale
      deltaLog = read deltaLogStr :: Double
      writeFlag = read writeFlagStr :: Bool
      std = read stdStr :: Double
      stdG = read stdGStr :: Double
      numIteration = read numIterationStr :: Int
  print thetaFreqs
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  hist <-
    if flag
      then do
        printCurrentTime $
          "read Fourier coefficients data from " L.++ histFilePath
        decodeFile histFilePath
      -- else runMonteCarloFourierCoefficientsGPU
      --        gpuID
      --        numThread
      --        numTrail
      --        maxTrail
      --        thetaSigma
      --        scaleSigma
      --        maxScale
      --        tao
      --        phiFreqs
      --        rhoFreqs
      --        thetaFreqs
      --        scaleFreqs
      --        deltaLog
      --        initScale
      --        histFilePath
      --        (emptyHistogram
      --           [ L.length phiFreqs
      --           , L.length rhoFreqs
      --           , L.length thetaFreqs
      --           , L.length scaleFreqs
      --           ]
      --           0)
      else sampleCartesian
             folderPath
             histFilePath
             gpuID
             numPoint
             numPoint
             180
             deltaX
             deltaLog
             initScale
             thetaSigma
             tao
             phiFreqs
             rhoFreqs
             thetaFreqs
             scaleFreqs
  let !points =
        L.map (\(!x, !y) -> Point x y 0 1) . getShape2DIndexList . makeShape2D $
        shape2D
      -- !xs = L.map (\(a, b) -> (a * deltaX, b * deltaX)) . makeShape2D $ shape2D
      !xs = getShape2DIndexList . makeShape2D $ shape2D
      !coefficients =
        normalizeFreqArr std phiFreqs rhoFreqs . getNormalizedHistogramArr $
        hist
       -- = getNormalizedHistogramArr $ hist :: R.Array U DIM4 (Complex Double)
      !thetaRHarmonics =
        computeThetaRHarmonics
          numOrientation
          numScale
          thetaFreqs
          scaleFreqs
          halfLogPeriod
  print xs
  plan <-
    makePlan
      folderPath
      emptyPlan
      numPoint
      numPoint
      (L.length thetaFreqs)
      (L.length scaleFreqs)
  let !initSourceSparse =
        computeInitialDistributionPowerMethodSparse phiFreqs rhoFreqs points
      !harmonicsArraySparse =
        computeHarmonicsArraySparse
          numPoint
          deltaX
          numPoint
          deltaX
          phiFreqs
          rhoFreqs
          thetaFreqs
          scaleFreqs
          halfLogPeriod
          cutoff
  -- let !initSource =
  --       computeInitialDistributionPowerMethod'
  --         numPoint
  --         numPoint
  --         phiFreqs
  --         rhoFreqs
  --         -- thetaFreqs
  --         -- scaleFreqs
  --         points
  --     !bias = computeBias numPoint numPoint points
  gaussian <- gaussianFilter2D plan numPoint stdG
  -- harmonicsArray <-
  --   dftHarmonicsArrayG
  --     plan
  --     numPoint
  --     deltaX
  --     numPoint
  --     deltaX
  --     phiFreqs
  --     rhoFreqs
  --     thetaFreqs
  --     scaleFreqs
  --     halfLogPeriod
  --     cutoff
  --     gaussian
  -- print . IA.indices $ harmonicsArraySparse
  -- completion <-
  computeContourSparse
    plan
    folderPath
    coefficients
    harmonicsArraySparse
    -- thetaRHarmonics
    (fromListUnboxed (Z :. L.length phiFreqs) phiFreqs)
    (fromListUnboxed (Z :. L.length rhoFreqs) rhoFreqs)
    cutoff
    -- gaussian
                      -- std
    xs
    numIteration
    suffix
    initSourceSparse
  -- completion <-
  --   computeContour'
  --    plan
  --    folderPath
  --    writeFlag
  --    coefficients
  --    harmonicsArray
  --    bias
  --    -- gaussian
  --    numIteration
  --    suffix
  --    thetaRHarmonics
  --    initSource
  -- plotDFTArrayThetaR
  --   (folderPath </> (printf "Completion_%s.png" suffix))
  --   numPoint
  --   numPoint
  --   thetaRHarmonics
  --   completion
  -- plotDFTArrayThetaRMag
  --   (folderPath </> (printf "CompletionMax_%s.png" suffix))
  --   numPoint
  --   numPoint
  --   thetaRHarmonics
  --   completion
  -- computeContourSparse'''
  --   plan
  --   folderPath
  --   coefficients
  --   harmonicsArraySparse
  --   thetaRHarmonics
  --   (fromListUnboxed (Z :. L.length phiFreqs) phiFreqs)
  --   (fromListUnboxed (Z :. L.length rhoFreqs) rhoFreqs)
  --   cutoff
  --   xs
  --   numIteration
  --   (suffix L.++ "Test")
  --   initSourceSparse