{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
module Marvin.API.Algorithms.LogisticRegression (
  LogisticRegression (..)
  , defaultLogisticRegression

  , LogisticRegressionModel
  , threshold
  , setThreshold
  , coefficients
  , intercept
) where

import qualified Numeric.LinearAlgebra as LA (toList)

import qualified Marvin.API.Algorithms.Internal.GradientDescent as GD
import Marvin.API.Algorithms.Internal.GradientDescent hiding (learningRate, lambda, addIntercept)

import Marvin.API.Meta.Model
import Marvin.API.Table.Internal
import Marvin.API.Fallible

-- * Parameters

-- | Parameters.
data LogisticRegression = LogisticRegression {
  learningRate :: Double -- ^ Coefficient of the gradient when updating the model.
  , numberOfIterations :: Int -- ^ Number of iterations.
  , lambda :: Double -- ^ Regularization coefficient.
  , addIntercept :: Bool -- ^ Whether to add a constant 1 column before training.
}

-- | Default parameters.
defaultLogisticRegression = LogisticRegression {
  learningRate = 0.001
  , numberOfIterations = 10
  , lambda = 0.1
  , addIntercept = True
}

toGradDesc :: LogisticRegression -> GradientDescent
toGradDesc logReg = defaultGradientDescent {
  numIter = numberOfIterations logReg
  , GD.cost = logisticRegressionCost
  , GD.learningRate = learningRate logReg
  , GD.addIntercept = addIntercept logReg
  , GD.lambda = lambda logReg
}

-- * Model

-- | Resulting model.
newtype LogisticRegressionModel = LogRegModel (Double, LinearModel)

-- | Retrieves the threshold to decide 'True' or 'False' based on probability.
threshold (LogRegModel (t, _)) = t
-- | Sets the threshold that decides 'True' or 'False' based on probability.
setThreshold t (LogRegModel (_, linear)) = LogRegModel (t, linear)

-- | Retrieves a list of the coefficient.
coefficients :: LogisticRegressionModel -> [Double]
coefficients (LogRegModel (_, linearModel)) = LA.toList $ coefs linearModel

-- | Retrieves the constant to be added.
intercept :: LogisticRegressionModel -> Double
intercept (LogRegModel (_, linearModel)) = intercept_ linearModel

instance Show LogisticRegressionModel where
  show (LogRegModel linear) = "Logistic regression model:\n\n" ++ show linear

instance Estimator LogisticRegression (NumericTable, BinaryColumn) where
  type ResultingModel LogisticRegression (NumericTable, BinaryColumn) =
    LogisticRegressionModel
  fit logReg (train, target) = do
    ensurePositive
      (InvalidAlgorithmParameter "LogisticRegression" "numberOfIterations must be positive")
        $ numberOfIterations logReg
    ensurePositive
      (InvalidAlgorithmParameter "LogisticRegression" "learning rate must be positive")
        $ learningRate logReg
    ensureNonNegative
      (InvalidAlgorithmParameter "LogisticRegression" "lambda must be non-negative")
        $ lambda logReg
    let numTarget = binaryToNumericColumn target
    ensureTableColSizeMatch "while training logistic regression" train numTarget
    let linearModel = fit' (toGradDesc logReg) (train, numTarget)
    let defaultThreshold = 0.5
    return $ LogRegModel (defaultThreshold, linearModel)

instance Predictor LogisticRegressionModel where
  type Testing LogisticRegressionModel = NumericTable
  type Prediction LogisticRegressionModel = BinaryColumn
  predict (LogRegModel (thresh, model)) testingTable = do
    numericPredictions <- probaPredict model testingTable
    return $ mapColumnAtoB (\x -> if x < thresh then False else True) numericPredictions

probaPredict model testingTable = do
  numericPredictions <- predict' model testingTable
  return $ mapColumnAtoB sigmoid numericPredictions


-- | <<docs/apollonian_gasket/spaced_gasket.svg>>
module Geometry.Processes.ApollonianGasket (
    -- * The classical shape
      createGasket
    , fifthCircle
    , Circle(..)

    -- * For playing around
    , ApoCircle(..)
    , newCircle
    , toCircle
    , toApoCircle
) where

import Data.Complex
import Data.Tree

import Geometry.Core as G

-- | A circle, but in a format more suitable for Descartes’ Theorem, which the
-- gasket algorithm is based on.
--
-- It interprets the center @'Vec2' x y@ as a complex number \(x + \mathrm i\,y\),
-- and stores the curvature \(k=\frac1r\) instead of the radius \(r\).
--
-- 'Circle' and 'ApoCircle' can be freely converted between using 'toCircle' and
-- 'toApoCircle'.
data ApoCircle = ApoCircle (Complex Double) Double
    deriving (Eq, Show)

instance HasBoundingBox ApoCircle where
    boundingBox = boundingBox . toCircle

-- | Simple conversion function. Inverse of 'toApoCircle'.
toCircle :: ApoCircle -> Circle
toCircle (ApoCircle (x :+ y) k) = Circle (Vec2 x y) (1/k)

-- | Simple conversion function. Inverse of 'toCircle'.
toApoCircle :: Circle -> ApoCircle
toApoCircle (Circle (Vec2 x y) r) = ApoCircle (x :+ y) (1/r)

-- | Curvature of the new Apollonian circle.
newCurvature
    :: (Double -> Double -> Double)
    -> Double -- ^ Curvature of first circle
    -> Double -- ^ Second circle
    -> Double -- ^ Third circle
    -> Double -- ^ New circle
newCurvature (+-) k1 k2 k3 =
    -- Taken straight from Wikipedia: https://en.wikipedia.org/wiki/Descartes%27_theorem
    (k1 + k2 + k3) +- (2 * sqrt (k1*k2 + k2*k3 + k3*k1))

-- | Center of the new Apollonian circle.
newCenter
    :: (Complex Double -> Complex Double -> Complex Double) -- ^ Plus or minus to yield both circles
    -> ApoCircle
    -> ApoCircle
    -> ApoCircle
    -> Double         -- ^ Curvature of the new circle
    -> Complex Double -- ^ Center of the new circle
newCenter (+-) (ApoCircle c1 k1) (ApoCircle c2 k2) (ApoCircle c3 k3) k4 =
    -- Taken straight from Wikipedia: https://en.wikipedia.org/wiki/Descartes%27_theorem
    1/k4' * ( (c1*k1' + c2*k2' + c3*k3')
              +- (2*sqrt (k1'*k2'*c1*c2 + k2'*k3'*c2*c3 + k1'*k3'*c1*c3)))
  where
    [k1', k2', k3', k4'] = map (:+ 0) [k1, k2, k3, k4]

-- | Create a new Apollonian circle, based on three existing ones. This is the raw
-- application of Descartes’ theorem. For something with less pitfalls, use
-- 'fifthCircle'.
--
-- The choice of sign is a mystery to me. For the Apollonian Gasket, the rule is:
--
--   1. The large circle is obtained using \((-,-)\) on the three initial cycles
--   2. The bottom-left circle is obtained using \((+,-)\) on the initial top, initial bottom, and large circle
--   3. All other recursive steps use \((+,+)\)
--   4. If you can explain the \((+,-)\), please contact me, I’m dying to know
--
-- If you’re wondering, this is what you get when you don’t do the \((+,-)\) trick:
--
-- <<docs/apollonian_gasket/missing_the_minus.svg>>
newCircle
    :: (Double -> Double -> Double)
        -- ^ Plus (larger curvature, smaller circle) or minus (smaller curvature, larger circle).
    -> (Complex Double -> Complex Double -> Complex Double) -- ^ Plus or minus; see this function’s doc.
    -> ApoCircle
    -> ApoCircle
    -> ApoCircle
    -> ApoCircle
newCircle plusMinusCurvature plusMinusCenter circ1@(ApoCircle _ k1) circ2@(ApoCircle _ k2) circ3@(ApoCircle _ k3) =
    let k4 = newCurvature plusMinusCurvature k1 k2 k3
        c4 = newCenter plusMinusCenter circ1 circ2 circ3 k4
    in ApoCircle c4 k4

-- | Given four mutually tangent circles, find »the other one« than the first argument.
-- This resolves the »\((+,-)\)« ambiguity of 'newCircle'.
fifthCircle :: ApoCircle -> ApoCircle -> ApoCircle -> ApoCircle -> ApoCircle
fifthCircle (ApoCircle c1 k1) (ApoCircle c2 k2) (ApoCircle c3 k3) (ApoCircle c4 k4) =
    let [k1', k2', k3', k4', k5'] = map (:+ 0) [k1, k2, k3, k4, k5]
        k5 = 2*(k2+k3+k4) - k1
        c5 = (2*(k2'*c2+k3'*c3+k4'*c4) - k1'*c1) / k5'
    in ApoCircle c5 k5

-- | The simple workhorse function. Given three mutually touching, equally sized
-- circles, it calculates the classical Apollonian Gasket:
--
--
-- <<docs/apollonian_gasket/classical_gasket.svg>>
--
-- Be warned though, the number of circles grows exponentially with recursion depth
-- (controlled with the radius parameter). This can easily melt your RAM, CPU or
-- harddrive if you aren’t careful.
--
-- It’s quite fun to play with the parameters though! Here is what happens when the three initial circles
-- aren’t touching:
--
-- <<docs/apollonian_gasket/forgetting_gen0.svg>>
--
-- To get even more creative freedom, I like to copy this algorithm and see what I
-- get, e.g. by commenting out one of the recursions.
createGasket
    :: Double -- ^ Radius threshold: do not recurse to circles smaller than this. Smaller values yield more detail.
    -> Circle -- ^ Initial left circle
    -> Circle -- ^ Initial right circle
    -> Circle -- ^ Initial bottom circle
    -> Tree Circle
createGasket minRadius gen0LCirc gen0RCirc gen0BCirc =
    let gen0L = toApoCircle gen0LCirc
        gen0R = toApoCircle gen0RCirc
        gen0B = toApoCircle gen0BCirc

        large = newCircle (-) (-) gen0L gen0R gen0B
        small = newCircle (+) (+) gen0L gen0R gen0B

        recurse :: ApoCircle -> ApoCircle -> ApoCircle -> ApoCircle -> Tree Circle
        recurse cA c1 c2 c3 =
            let cB@(ApoCircle _ kB) = fifthCircle cA c1 c2 c3
            in Node (toCircle cB)
                (if abs kB < 1/minRadius
                    then
                        [ recurse c3 cB c1 c2
                        , recurse c2 c3 cB c1
                        , recurse c1 c2 c3 cB
                        ]
                    else [])

        circles :: Tree Circle
        circles = Node (toCircle large)
            [ recurse small gen0L gen0R gen0B
            , recurse large gen0L gen0R gen0B
            ]
    in circles


module LinearRegressionExercise where

import qualified Graphics.Gnuplot.Advanced as GP

import qualified Graphics.Gnuplot.Graph as Graph

import qualified Graphics.Gnuplot.Plot.ThreeDimensional as Plot3D
import qualified Graphics.Gnuplot.Graph.ThreeDimensional as Graph3D

import qualified Graphics.Gnuplot.Plot.TwoDimensional as Plot2D
import qualified Graphics.Gnuplot.Graph.TwoDimensional as Graph2D
import qualified Graphics.Gnuplot.Terminal.Default as DefaultTerm
import Graphics.Gnuplot.Plot.TwoDimensional (linearScale, )

import LinearRegression
import Numeric.LinearAlgebra
import Prelude hiding ((<>))
import System.Directory
import Util

-- replace this with the path to your data.
exerciseOneDataOne :: IO String
exerciseOneDataOne = do
  home <- getHomeDirectory
  return (home ++ "/code/haskell/CourseraMachineLearning/app/ex1data1.txt")

exerciseOneDataTwo :: IO String
exerciseOneDataTwo = do
  home <- getHomeDirectory
  return (home ++ "/code/haskell/CourseraMachineLearning/app/ex1data2.txt")

exerciseOne = do
  -- Part 1: Basic function
  putStrLn "Running warmUpExercise ..."
  putStrLn "5x5 Identity Matrix: "
  let a = ident 5 :: Matrix Double
  print a

  -- Part 2: Plotting
  putStrLn "Plotting Data ..."
  dataOne <- exerciseOneDataOne
  [x', y'] <- fmap toColumns (loadMatrix dataOne)
  GP.plotSync DefaultTerm.cons (plotPoints (toList x') (toList y'))

  -- Part 3: Cost and Gradient Descent
  let m = size y'
  let x = ones m 1 ||| reshape 1 x'
  let y = asColumn y'
  let xPoints = toList x'
  let yPoints = toList y'
  let theta = zeros 2 1
  let iterations = 1500
  let alpha = 0.01
  putStrLn "Testing the cost function ..."
  let j = computeCost x y theta
  putStrLn ("With theta = [0, 0]\n Cost computed = " ++ show j)
  putStrLn "Expected cost value (approx) 32.07"
  let k = computeCost x y (asColumn $ 2 |> [-1, 2])
  putStrLn ("With theta = [-1, 2]\n Cost computed = " ++ show k)
  putStrLn "Expected cost value (approx) 54.24"

  putStrLn "Running Gradient Descent ..."
  let (costHist, thetaOpt) = gradientDescent  x y theta alpha iterations
  putStrLn ("Theta found by gradient descent: " ++ show thetaOpt)
  putStrLn "Expected theta values (approx)\n"
  putStrLn " -3.6303\n  1.1664"
  let yOptPoints = (toList (flatten $ x <> thetaOpt))
  putStrLn " -3.6303\n  1.1664"
  let yOptPoints = (toList (flatten $ x <> thetaOpt))
  GP.plotSync DefaultTerm.cons $ plotLine xPoints yOptPoints `mappend` plotPoints xPoints yPoints

  -- Predicting values
  let predict1 = row [1, 3.5] <> thetaOpt
  putStrLn "For population = 35,000 we predict a profit of "
  print ((predict1 * 1000) `atIndex` (0,0))

  let predict2 = row [1, 20] <> thetaOpt
  putStrLn "For population = 70,000 we predict a profit of "
  print ((predict2 * 1000) `atIndex` (0,0))

  -- Visualising J(theta_0, theta_1)
  let theta0 = toList (linspace 100 (-10 :: Double, 10 :: Double))
  let theta1 = toList (linspace 100 (-1  :: Double, 4  :: Double))
  let jVals = [computeCost x y (col [i, j]) | i <- theta0, j <- theta1]
  -- TODO(mchav): Draw contour plot as well.
  GP.plotSync DefaultTerm.cons $ Plot3D.surface theta0 theta1 (\i j -> computeCost x y (col [i, j]))

exerciseOneOptional = do
  dataTwo <- exerciseOneDataTwo
  columns <- fmap toColumns (loadMatrix dataTwo)
  let y' = last columns
  let y = asColumn y'
  let x' = fromColumns (init columns)
  let m = size y'

  let (xNorm, mu, sigma) = featureNormalize x'
  let x = ones m 1 ||| xNorm

  let theta = zeros 3 1
  let iterations = 50 :: Int
  let alpha = 1

  -- TODO(mchav): Include cost history in gradient descent function and then plot it.
  let (costHist, thetaOpt) = gradientDescent x y theta alpha iterations
  putStrLn ("Theta found by gradient descent: " ++ show thetaOpt)
  -- Plot the graph ingoring the first (zeroth) iteration.
  GP.plotSync DefaultTerm.cons $ plotLine [1..(fromIntegral iterations)] ((tail . reverse) costHist)
  putStrLn "Predicted price of a 1650 sq-ft, 3 br house (using gradient descent)"
  let xTest = ones 1 1 ||| ((row [1650, 3] - mu) / sigma)
  print (xTest <> thetaOpt)
  let x'' = ones m 1 ||| x'
  let normTheta = normalEquation x'' y
  print normTheta
  let xTest' = ones 1 1 ||| row [1650, 3]
  print (xTest' <> normTheta)


{-# OPTIONS_GHC -Wall #-}

import           Numeric.Sundials.ARKode.ODE
import           Numeric.LinearAlgebra

import           Data.List (intercalate)

import           Text.PrettyPrint.HughesPJClass

butcherTableauTex :: ButcherTable -> String
butcherTableauTex (ButcherTable m c b b2) =
  render $
  vcat [ text ("\n\\begin{array}{c|" ++ (concat $ replicate n "c") ++ "}")
       , us
       , text "\\hline"
       , text bs <+> text "\\\\"
       , text b2s <+> text "\\\\"
       , text "\\end{array}"
       ]
  where
    n = rows m
    rs = toLists m
    ss = map (\r -> intercalate " & " $ map show r) rs
    ts = zipWith (\i r -> show i ++ " & " ++ r) (toList c) ss
    us = vcat $ map (\r -> text r <+> text "\\\\") ts
    bs  = " & " ++ (intercalate " & " $ map show $ toList b)
    b2s = " & " ++ (intercalate " & " $ map show $ toList b2)

main :: IO ()
main = do

  let res = butcherTable (SDIRK_2_1_2 undefined)
  putStrLn $ show res
  putStrLn $ butcherTableauTex res

  let resA = butcherTable (KVAERNO_4_2_3 undefined)
  putStrLn $ show resA
  putStrLn $ butcherTableauTex resA

  let resB = butcherTable (SDIRK_5_3_4 undefined)
  putStrLn $ show resB
  putStrLn $ butcherTableauTex resB

  let resC = butcherTable (FEHLBERG_6_4_5 undefined)
  putStrLn $ show resC
  putStrLn $ butcherTableauTex resC


module FFT.DFTData(
    dftInnerMap,
    dftFoldMap,
    dftOuterMap,
    dftInnerMapReduce,
    dftInnerParMapReduce
) where
import Data.Complex
import FFT.Samples
import Control.Parallel
import Strategies

-- twiddle factors
tw :: Int -> Int -> Complex Float
tw n k = cis (-2 * pi * fromIntegral k / fromIntegral n)

dftInnerMap xs = [ sum (innerMap k) | k <- [0..n']]
  where
    innerMap k = parmap (\j -> xs!!j * tw n (j*k)) [0..n']
    n = length xs
    n' = n-1

dftFoldMap xs = [ (parfold (+) (0 :+ 0) (innerMap k)) | k <- [0..n']]
  where
    innerMap k = parmap (\j -> xs!!j * tw n (j*k)) [0..n']
    n = length xs
    n' = n-1

dftOuterMap xs = parmap innerFunc [0..n']
    where
        innerFunc k = sum [ xs!!j * tw n (j*k) | j <- [0..n']]
        n = length xs
        n' = n-1

dftInnerMapReduce xs = map (\k -> parMapReduceSimple rdeepseq (mapper k) rdeepseq reducer [0..n']) [0..n']
    where
        mapper k j = xs!!j * tw n (j*k)
        reducer = sum
        n = length xs
        n' = n-1

dftInnerParMapReduce xs = parmap (\k -> parMapReduceSimple rdeepseq (mapper k) rdeepseq reducer [0..n']) [0..n']
    where
        mapper k j = xs!!j * tw n (j*k)
        reducer = sum
        n = length xs
        n' = n-1


-- |
-- Module:      Data.Euclidean
-- Copyright:   (c) 2019 Andrew Lelechenko
-- Licence:     BSD3
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--

{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DefaultSignatures          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MagicHash                  #-}

module Data.Euclidean
  ( Euclidean(..)
  , Field
  , GcdDomain(..)
  , WrappedIntegral(..)
  , WrappedFractional(..)
  , gcdExt
  ) where

import Prelude hiding (quotRem, quot, rem, divMod, div, mod, gcd, lcm, negate, (*), Int, Word)
import qualified Prelude as P
import Control.Exception (throw, ArithException(..))
import Data.Bits (Bits)
import Data.Complex (Complex(..))
import Data.Int (Int, Int8, Int16, Int32, Int64)
import Data.Maybe (isJust)
import Data.Ratio (Ratio)
import Data.Semiring (Semiring(..), Ring(..), (*), minus, isZero, Mod2)
import Data.Word (Word, Word8, Word16, Word32, Word64)
import Foreign.C.Types (CFloat, CDouble)
import GHC.Exts (Int(..), Word(..))
import GHC.Integer.GMP.Internals (gcdInt, gcdWord, gcdInteger, lcmInteger)

import Numeric.Natural

---------------------------------------------------------------------
-- Classes
---------------------------------------------------------------------

-- | 'GcdDomain' represents a
-- <https://en.wikipedia.org/wiki/GCD_domain GCD domain>.
-- This is a domain, where GCD can be defined,
-- but which does not necessarily allow a well-behaved
-- division with remainder (as in 'Euclidean' domains).
--
-- For example, there is no way to define 'rem' over
-- polynomials with integer coefficients such that
-- remainder is always "smaller" than divisor. However,
-- 'gcd' is still definable, just not by means of
-- Euclidean algorithm.
--
-- All methods of 'GcdDomain' have default implementations
-- in terms of 'Euclidean'. So most of the time
-- it is enough to write:
--
-- > instance GcdDomain Foo
-- > instance Euclidean Foo where
-- >   quotRem = ...
-- >   degree  = ...
class Semiring a => GcdDomain a where
  -- | Division without remainder.
  --
  -- prop> \x y -> (x * y) `divide` y == Just x
  -- prop> \x y -> maybe True (\z -> x == z * y) (x `divide` y)
  divide :: a -> a -> Maybe a

  default divide :: (Eq a, Euclidean a) => a -> a -> Maybe a
  divide x y = let (q, r) = quotRem x y in
    if isZero r then Just q else Nothing

  -- | Greatest common divisor. Must satisfy
  --
  -- prop> \x y -> isJust (x `divide` gcd x y) && isJust (y `divide` gcd x y)
  -- prop> \x y z -> isJust (gcd (x * z) (y * z) `divide` z)
  gcd :: a -> a -> a

  default gcd :: (Eq a, Euclidean a) => a -> a -> a
  gcd a b
    | isZero b  = a
    | otherwise = gcd b (a `rem` b)

  -- | Lowest common multiple. Must satisfy
  --
  -- prop> \x y -> isJust (lcm x y `divide` x) && isJust (lcm x y `divide` y)
  -- prop> \x y z -> isNothing (z `divide` x) || isNothing (z `divide` y) || isJust (z `divide` lcm x y)
  lcm :: a -> a -> a

  default lcm :: Eq a => a -> a -> a
  lcm a b
    | isZero a || isZero b = zero
    | otherwise = case a `divide` gcd a b of
      Nothing -> error "lcm: violated gcd invariant"
      Just c  -> c * b

  -- | Test whether two arguments are
  -- <https://en.wikipedia.org/wiki/Coprime_integers coprime>.
  -- Must match its default definition:
  --
  -- prop> \x y -> coprime x y == isJust (1 `divide` gcd x y)
  coprime :: a -> a -> Bool

  default coprime :: Eq a => a -> a -> Bool
  coprime x y = isJust (one `divide` gcd x y)

infixl 7 `divide`

-- | Informally speaking, 'Euclidean' is a superclass of 'Integral',
-- lacking 'toInteger', which allows to define division with remainder
-- for a wider range of types, e. g., complex integers
-- and polynomials with rational coefficients.
--
-- 'Euclidean' represents a
-- <https://en.wikipedia.org/wiki/Euclidean_domain Euclidean domain>
-- endowed by a given Euclidean function 'degree'.
--
-- No particular rounding behaviour is expected of 'quotRem'. E. g.,
-- it is not guaranteed to truncate towards zero or towards negative
-- infinity (cf. 'P.divMod'), and remainders are not guaranteed to be non-negative.
-- For a faithful representation of residue classes one can use
-- <http://hackage.haskell.org/package/mod mod> package instead.
class GcdDomain a => Euclidean a where
  {-# MINIMAL (quotRem | quot, rem), degree #-}
  -- | Division with remainder.
  --
  -- prop> \x y -> y == 0 || let (q, r) = x `quotRem` y in x == q * y + r
  quotRem :: a -> a -> (a, a)
  quotRem x y = (quot x y, rem x y)

  -- | Division. Must match its default definition:
  --
  -- prop> \x y -> quot x y == fst (quotRem x y)
  quot :: a -> a -> a
  quot x y = fst (quotRem x y)

  -- | Remainder. Must match its default definition:
  --
  -- prop> \x y -> rem x y == snd (quotRem x y)
  rem :: a -> a -> a
  rem x y = snd (quotRem x y)

  -- | Euclidean (aka degree, valuation, gauge, norm) function on @a@. Usually @'fromIntegral' '.' 'abs'@.
  --
  -- 'degree' is rarely used by itself. Its purpose
  -- is to provide an evidence of soundness of 'quotRem'
  -- by testing the following property:
  --
  -- prop> \x y -> y == 0 || let (q, r) = x `quotRem` y in (r == 0 || degree r < degree y)
  degree :: a -> Natural

infixl 7 `quot`
infixl 7 `rem`

coprimeIntegral :: Integral a => a -> a -> Bool
coprimeIntegral x y = (odd x || odd y) && P.gcd x y == 1

-- | Execute the extended Euclidean algorithm.
-- For elements @a@ and @b@, compute their greatest common divisor @g@
-- and the coefficient @s@ satisfying @as + bt = g@ for some @t@.
gcdExt :: (Eq a, Euclidean a, Ring a) => a -> a -> (a, a)
gcdExt = go one zero
  where
    go s s' r r'
      | r' == zero = (r, s)
      | otherwise  = case quotRem r r' of
        (q, r'') -> go s' (minus s (times q s')) r' r''
{-# INLINABLE gcdExt #-}

-- | 'Field' represents a
-- <https://en.wikipedia.org/wiki/Field_(mathematics) field>,
-- a ring with a multiplicative inverse for any non-zero element.
class (Euclidean a, Ring a) => Field a

---------------------------------------------------------------------
-- Instances
---------------------------------------------------------------------

instance GcdDomain () where
  divide  = const $ const (Just ())
  gcd     = const $ const ()
  lcm     = const $ const ()
  coprime = const $ const True

instance Euclidean () where
  degree  = const 0
  quotRem = const $ const ((), ())
  quot    = const $ const ()
  rem     = const $ const ()

instance Field ()

instance GcdDomain Mod2 where

instance Euclidean Mod2 where
  degree = const 0
  quotRem x y
    | isZero y  = throw DivideByZero
    | otherwise = (x, zero)

instance Field Mod2

-- | Wrapper around 'Integral' with 'GcdDomain'
-- and 'Euclidean' instances.
newtype WrappedIntegral a = WrapIntegral { unwrapIntegral :: a }
  deriving (Eq, Ord, Show, Num, Integral, Real, Enum, Bits)

instance Num a => Semiring (WrappedIntegral a) where
  plus  = (P.+)
  zero  = 0
  times = (P.*)
  one   = 1
  fromNatural = P.fromIntegral

instance Num a => Ring (WrappedIntegral a) where
  negate = P.negate

instance Integral a => GcdDomain (WrappedIntegral a) where
  divide x y = case x `P.quotRem` y of (q, 0) -> Just q; _ -> Nothing
  gcd     = P.gcd
  lcm     = P.lcm
  coprime = coprimeIntegral

instance Integral a => Euclidean (WrappedIntegral a) where
  degree  = P.fromIntegral . abs . unwrapIntegral
  quotRem = P.quotRem
  quot    = P.quot
  rem     = P.rem

instance GcdDomain Int where
  divide x y = case x `P.quotRem` y of (q, 0) -> Just q; _ -> Nothing
#if MIN_VERSION_integer_gmp(0,5,1)
  gcd (I# x) (I# y) = I# (gcdInt x y)
#else
  gcd     = P.gcd
#endif
  lcm     = P.lcm
  coprime = coprimeIntegral

instance GcdDomain Word where
  divide x y = case x `P.quotRem` y of (q, 0) -> Just q; _ -> Nothing
#if MIN_VERSION_integer_gmp(1,0,0)
  gcd (W# x) (W# y) = W# (gcdWord x y)
#else
  gcd     = P.gcd
#endif
  lcm     = P.lcm
  coprime = coprimeIntegral

instance GcdDomain Integer where
  divide x y = case x `P.quotRem` y of (q, 0) -> Just q; _ -> Nothing
  gcd     = gcdInteger
  lcm     = lcmInteger
  coprime = coprimeIntegral

#define deriveGcdDomain(ty)     \
instance GcdDomain (ty) where { \
;  divide x y = case x `P.quotRem` y of { (q, 0) -> Just q; _ -> Nothing } \
;  gcd     = P.gcd              \
;  lcm     = P.lcm              \
;  coprime = coprimeIntegral    \
}

deriveGcdDomain(Int8)
deriveGcdDomain(Int16)
deriveGcdDomain(Int32)
deriveGcdDomain(Int64)
deriveGcdDomain(Word8)
deriveGcdDomain(Word16)
deriveGcdDomain(Word32)
deriveGcdDomain(Word64)
deriveGcdDomain(Natural)

#define deriveEuclidean(ty)       \
instance Euclidean (ty) where {   \
;  degree  = P.fromIntegral . abs \
;  quotRem = P.quotRem            \
;  quot    = P.quot               \
;  rem     = P.rem                \
}

deriveEuclidean(Int)
deriveEuclidean(Int8)
deriveEuclidean(Int16)
deriveEuclidean(Int32)
deriveEuclidean(Int64)
deriveEuclidean(Integer)
deriveEuclidean(Word)
deriveEuclidean(Word8)
deriveEuclidean(Word16)
deriveEuclidean(Word32)
deriveEuclidean(Word64)
deriveEuclidean(Natural)

-- | Wrapper around 'Fractional'
-- with trivial 'GcdDomain'
-- and 'Euclidean' instances.
newtype WrappedFractional a = WrapFractional { unwrapFractional :: a }
  deriving (Eq, Ord, Show, Num, Fractional)

instance Num a => Semiring (WrappedFractional a) where
  plus  = (P.+)
  zero  = 0
  times = (P.*)
  one   = 1
  fromNatural = P.fromIntegral

instance Num a => Ring (WrappedFractional a) where
  negate = P.negate

instance Fractional a => GcdDomain (WrappedFractional a) where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Fractional a => Euclidean (WrappedFractional a) where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Fractional a => Field (WrappedFractional a)

instance Integral a => GcdDomain (Ratio a) where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Integral a => Euclidean (Ratio a) where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Integral a => Field (Ratio a)

instance GcdDomain Float where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Euclidean Float where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Field Float

instance GcdDomain Double where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Euclidean Double where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Field Double

instance GcdDomain CFloat where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Euclidean CFloat where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Field CFloat

instance GcdDomain CDouble where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Euclidean CDouble where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Field CDouble

conjQuotAbs :: Field a => Complex a -> Complex a
conjQuotAbs (x :+ y) = x `quot` norm :+ (negate y) `quot` norm
  where
    norm = (x `times` x) `plus` (y `times` y)

instance Field a => GcdDomain (Complex a) where
  divide x y = Just (x `times` conjQuotAbs y)
  gcd        = const $ const one
  lcm        = const $ const one
  coprime    = const $ const True

instance Field a => Euclidean (Complex a) where
  degree      = const 0
  quotRem x y = (quot x y, zero)
  quot x y    = x `times` conjQuotAbs y
  rem         = const $ const zero

instance Field a => Field (Complex a)


module LeastSquares
    ( VarName
    , LinExpr(..)
    , LinContr(..)
    , QuadExpr(..)
    , minimize
    ) where

import Data.List (elemIndex, foldl1')
import qualified Data.Map.Strict as M
import Data.Maybe (fromJust)
import qualified Data.Set as S
import Numeric.LinearAlgebra

type VarName = String

type CanonicalMap = M.Map VarName (Matrix R)

infix 4 :==:

data LinExpr = Var VarName Int
             | Coeff R
             | Vec (Vector R)
             | Mat (Matrix R)
             | Neg LinExpr
             | Sum LinExpr LinExpr
             | Prod LinExpr LinExpr
             deriving (Show, Eq)

instance Num LinExpr where
    (+) = Sum
    a - b = Sum a (Neg b)
    (*) = Prod
    negate = Neg
    signum = undefined
    abs = undefined
    fromInteger a = Coeff (fromInteger a)

data LinContr = LinExpr :==: LinExpr

data QuadExpr = SumSquares LinExpr
              | CoeffQuad R
              | ProdQuad QuadExpr QuadExpr
              | SumQuad QuadExpr QuadExpr

instance Num QuadExpr where
    (+) = SumQuad
    (*) = ProdQuad
    (-) = undefined
    negate = undefined
    signum = undefined
    abs = undefined
    fromInteger a = CoeffQuad (fromInteger a)

class CanonExpr e  where
    varSet :: e -> S.Set (VarName, Int)
    canonicalize :: e -> CanonicalMap

-- Canonicalize a generic expression
-- The vector is stored as a column matrix under the variable "".
instance CanonExpr LinExpr where
    varSet (Var v s) = S.singleton (v, s)
    varSet (Coeff _) = S.empty
    varSet (Vec _) = S.empty
    varSet (Mat _) = S.empty
    varSet (Neg e) = varSet e
    varSet (Prod _ e) = varSet e
    varSet (Sum e1 e2) = varSet e1 `S.union` varSet e2
    canonicalize (Var v s) = M.singleton v (ident s)
    canonicalize (Vec v) = M.singleton "" (asColumn v)
    canonicalize (Neg e) = M.map negate $ canonicalize e
    canonicalize (Prod (Coeff c) e) = M.map (scale c) $ canonicalize e
    canonicalize (Prod (Mat m) e) = M.map (m <>) $ canonicalize e
    canonicalize (Sum e1 e2) =
        M.unionWith (+) (canonicalize e1) (canonicalize e2)
    canonicalize _ = error "Expression is not well-formed"

-- Canonicalize an expression in the constraint
instance CanonExpr LinContr where
    varSet (e1 :==: e2) = varSet e1 `S.union` varSet e2
    canonicalize (e1 :==: e2) = canonicalize (e1 - e2)

-- Canonicalize an expression in the objective, which consists of SumSquares.
instance CanonExpr QuadExpr where
    varSet (SumSquares e) = varSet e
    varSet (CoeffQuad _) = S.empty
    varSet (ProdQuad _ e) = varSet e
    varSet (SumQuad e1 e2) = varSet e1 `S.union` varSet e2
    canonicalize (SumSquares e) = canonicalize e
    canonicalize (ProdQuad (CoeffQuad c) e) =
        M.map (scale (sqrt c)) $ canonicalize e
    canonicalize (SumQuad e1 e2) = canonicalize e1 `vstack` canonicalize e2
    canonicalize _ = error "Expression is not well-formed"

-- Combine two canonical maps by stacking one on top of another.
vstack :: CanonicalMap -> CanonicalMap -> CanonicalMap
vstack m1 m2 = M.fromSet combine (M.keysSet m1 `S.union` M.keysSet m2)
  where
    combine k = lookupDefault k m1 === lookupDefault k m2
    lookupDefault k m = M.findWithDefault (konst 0 (firstDim m, 1)) k m
    firstDim = rows . snd . fromJust . M.lookupGE "" :: CanonicalMap -> Int

minimize :: QuadExpr -> [LinContr] -> M.Map VarName (Vector R)
minimize obj [] = unpack vars $ a <\> (-b)
  where
    (a,b) = pack vars . canonicalize $ obj
    vars = varSet obj
minimize obj constraints = unpack vars $ mat <\> (-vec)
  where
    mat = fromBlocks [[2 * tr a <> a, tr c], [c, 0]]
    vec = vjoin [2 * tr a #> b, d]
    (a,b) = pack vars . canonicalize $ obj
    (c,d) = pack vars $ foldl1' vstack cds
    cds = map canonicalize constraints
    vars = varSet obj `S.union` S.unions (map varSet constraints)

pack :: S.Set (VarName, Int) -> CanonicalMap -> (Matrix R, Vector R)
pack vars m = (a, b)
  where
    a = fromBlocks [map getCoefficient (S.toAscList vars)]
    b = flatten $ M.findWithDefault (konst 0 (rows a, 1)) "" m
    getCoefficient (varName,s) = M.findWithDefault (konst 0 (s, s)) varName m

-- The vector given to unpack might be longer than the total size of the
-- variables. It might contain the dual variable values, for example.
unpack :: S.Set (VarName, Int) -> Vector R -> M.Map VarName (Vector R)
unpack vars result = M.fromList $ zip varNames coefficients
  where
    coefficients = map getCoefficient varList
    getCoefficient (varName,s) = subVector (startIndex varName) s result
    startIndex varName = cumsum !! fromJust (elemIndex varName varNames)
    cumsum = scanl (+) 0 varSizes
    (varNames,varSizes) = unzip varList
    varList = S.toAscList vars


module Scene.Object.SimpleSphere where

import Scene.Base
import Math.Color
import Math.Ray
import Numeric.LinearAlgebra

-- Sphere
sphere :: Double -> Vec3 -> Material -> Scene
sphere r p mat = [sphere' r p mat]

sphere' :: Double -> Vec3 -> Material -> Object
sphere' r p = Object (sphere_intercept r p)

sphere_intercept :: Double -> Vec3 -> Ray -> Maybe (Double, Ray)
sphere_intercept r p ray@(Ray o v) =
    let
        o' = o-p
        a = v <.> v
        b = v <.> o'
        c = o' <.> o' - r**2
        discriminant = b**2 - a*c
    in if discriminant > 0
        then
            let
                x1 = (-b + sqrt(discriminant))/a
                x2 = (-b - sqrt(discriminant))/a
                x = case (x1 > 0, x2 > 0) of
                    (True,True) -> Just (min x1 x2)
                    (False,True) -> Just x2
                    (True, False) -> Just x1
                    (False, False) -> Nothing

            in (\lambda -> (lambda, Ray (parametric ray lambda) (parametric ray lambda - p))) <$> x
        else Nothing


-- |
-- Module    : Statistics.Matrix.Algorithms
-- Copyright : 2014 Bryan O'Sullivan
-- License   : BSD3
--
-- Useful matrix functions.

module Statistics.Matrix.Algorithms
    (
      qr
    ) where

import Control.Applicative
import Control.Monad.ST (ST, runST)
import Statistics.Matrix (Matrix, column, dimension, for, norm)
import qualified Statistics.Matrix.Mutable as M
import Statistics.Sample.Internal (sum)
import qualified Data.Vector.Unboxed as U
import Prelude hiding (sum, replicate)

-- | /O(r*c)/ Compute the QR decomposition of a matrix.
-- The result returned is the matrices (/q/,/r/).
qr :: Matrix -> (Matrix, Matrix)
qr mat = runST $ do
  let (m,n) = dimension mat
  r <- M.replicate n n 0
  a <- M.thaw mat
  for 0 n $ \j -> do
    cn <- M.immutably a $ \aa -> norm (column aa j)
    M.unsafeWrite r j j cn
    for 0 m $ \i -> M.unsafeModify a i j (/ cn)
    for (j+1) n $ \jj -> do
      p <- innerProduct a j jj
      M.unsafeWrite r j jj p
      for 0 m $ \i -> do
        aij <- M.unsafeRead a i j
        M.unsafeModify a i jj $ subtract (p * aij)
  (,) <$> M.unsafeFreeze a <*> M.unsafeFreeze r

innerProduct :: M.MMatrix s -> Int -> Int -> ST s Double
innerProduct mmat j k = M.immutably mmat $ \mat ->
  sum $ U.zipWith (*) (column mat j) (column mat k)


module Dft where

import Data.Complex

nthPartOfCycle :: Floating a => Int -> Int -> Int -> Complex a
nthPartOfCycle n f _N =
  cis (2 * pi * ((fromIntegral (n * f)) / (fromIntegral _N)))

correlateSignalAtFreq :: RealFloat a => [Complex a] -> Int -> [(Complex a, Complex a, Complex a)]
correlateSignalAtFreq signal f =
  let _N = length signal
      wavePart = [nthPartOfCycle i f _N | i <- [0.._N - 1]]
  in
   scanr (\(s, w) (_, _, sum) -> (s, w, sum + s * w)) (0 :+ 0, 0 :+ 0, 0 :+ 0) (reverse $ zip signal wavePart)


-- Take a Ten Minute Walk
-- http://www.codewars.com/kata/54da539698b8a2ad76000228/

module Codewars.Kata.TenMinuteWalk where

import Data.Complex (magnitude, Complex( (:+) ))

isValidWalk :: String -> Bool
isValidWalk walk = case splitAt 10 walk of (ds, []) -> (length walk == 10) &&
                                                       ((==0) . magnitude . sum . map (\c -> case c of
                                                           'n' -> 0 :+ 1
                                                           's' -> 0 :+ (-1)
                                                           'w' -> (-1) :+ 0
                                                           'e' -> 1 :+ 0) $ walk)
                                           (_, rest) -> False


module Colors
   (
   -- * Normalization functions
     normLinear
   , normSigmoid
   , normPeriodic
   , normSine
   -- * Colorization functions
   , colorGrey
   , colorHue
   , colorGrad
   )
where

import           Data.Complex
import qualified Graphics.Image                as I

mod' :: RealFloat a => a -> a -> a
mod' x y = (-) x . (*) y . fromIntegral . truncate $ x / y

normLinear :: Int -> Double -> Double
normLinear max = (/ fromIntegral max)

normSigmoid :: Double -> Double -> Double -> Double
normSigmoid center power = (1 /) . (1 +) . (power **) . (/ (center / 4)) . (center -)

normPeriodic :: Double -> Double -> Double
normPeriodic period = (/ period) . flip mod' period

normSine :: Double -> Double -> Double
normSine period = (/ 2) . (1 -) . cos . (* 6.2831853) . (/ period) . (flip mod' period)

colorGrey :: I.Pixel I.RGB Double -> Maybe Double -> I.Pixel I.RGB Double
colorGrey setColor Nothing    = setColor
colorGrey _        (Just val) = I.PixelRGB val val val

colorGrad :: I.Pixel I.RGB Double -> Bool -> [I.Pixel I.RGB Double] -> Maybe Double -> I.Pixel I.RGB Double
colorGrad convergentC _ _      Nothing  = convergentC
colorGrad _           isCircular colorsIn (Just x) = interColors (colors !! section) (colors !! (section + 1))
 where
  colors  = if isCircular then colorsIn ++ [head colorsIn] else colorsIn
  count   = fromIntegral $ (subtract 1) $ length colors
  section = truncate $ x * count
  p       = (* count) $ (x `mod'`) $ 1 / count
  inter i f = (i +) $ (* p) $ (f - i)
  interColors (I.PixelRGB iR iG iB) (I.PixelRGB fR fG fB) =
     I.PixelRGB (inter iR fR) (inter iG fG) (inter iB fB)

colorHue :: I.Pixel I.RGB Double -> Maybe Double -> I.Pixel I.RGB Double
colorHue setColor Nothing = setColor
colorHue _        (Just x) | x' <= 1   = I.PixelRGB x'       0        1
                           | x' <= 2   = I.PixelRGB 1        0        (2 - x')
                           | x' <= 3   = I.PixelRGB 1        (x' - 2) 0
                           | x' <= 4   = I.PixelRGB (4 - x') 1        0
                           | x' <= 5   = I.PixelRGB 0        1        (x' - 4)
                           | otherwise = I.PixelRGB 0        (6 - x') 1
   where x' = 6 * x


-- | This module implements various functions that return a probabilistic result,
-- defined as unitary operators, and quantum computations.
module QIO.QIORandom where

import Data.Monoid as Monoid
import QIO.QioSyn
import QIO.Qdata
import QIO.Qio
import Data.Complex

-- | The exponentiated Pauli-X rotation
rX :: RR -> Rotation
rX r (x,y) = if x==y then (cos (r/2):+0) else (0:+ (-(sin (r/2))))

-- | The exponentiated Pauli-Y rotation
rY :: RR -> Rotation
rY r (x,y) = if x==y then (cos (r/2):+0) else (s * sin (r/2):+0) where s = if x then 1 else -1

-- | Applies a Hadamard rotation to each qubit in the given list of qubits
hadamards :: [Qbit] -> U
hadamards [] = mempty
hadamards (q:qs) = uhad q `mappend` hadamards qs

-- | returns the highest integer power of 2 that is less than or equal to \x\
pow2 :: Int -> Int
pow2 x = pow2' 0
  where pow2' y | 2^(y+1) > x = 2^y
            | otherwise = pow2' (y+1)

-- | A rotation that, given a qubit in state 0, leaves it in a super-position of
-- 0 and 1, such that the probability of measuring as state 0 is \ps\.
weightedU :: RR -> Qbit -> U
weightedU ps q | sqrt ps <= 1 = rot q (rX (2*(acos (sqrt ps))))
           | otherwise = error ("weightedU: Invalid Probability: " ++ show ps)
-- | A QIO computation that uses the "weightedU" unitary, to return a Bool that
-- has a probablity of \pf\ of being False.
weightedBool :: RR -> QIO Bool
weightedBool pf = do q <- mkQbit False
                     applyU (weightedU pf q)
                     measQ q

-- | removes any leading Falses from a list of booleans
rlf :: [Bool] -> [Bool]
rlf (False:bs) = rlf bs
rlf bs = bs

-- | removes any leading Falses from the (big-endian) bit-wise representation
-- of the given Int.
rlf_l :: Int -> [Bool]
rlf_l x = rlf (reverse (int2bits x))

-- | returns the number of bits left after calling the "flf_l" function
rlf_n :: Int -> Int
rlf_n x = length (rlf_l x)

-- | Given an Int \max\ that is the largest number required to be represented in
-- a quantum register, this function trims the front off the given register, to
-- leave the number of qubits required to represent \max\.
trim :: Int -> [Qbit] -> [Qbit]
trim max qbs = drop ((length qbs)-(rlf_n max)) qbs

-- | Given an Int \max\, and a quantum register in the state \max\, this function
-- defines a unitary operation that will leave the quantum register in state that
-- has equal probability of being measured in any of the states 0 to \max\.
randomU :: Int -> [Qbit] -> U
randomU max qbs = randomU' max (trim max qbs)
 where
  randomU' _ [] = mempty
  randomU' 0 _ = mempty
  randomU' max (q:qbs) = weightedU (fromIntegral ((max+1)-p)/fromIntegral (max+1)) q
                 `mappend`
                 condQ q (\x -> if x then (randomU (max-p) qbs)
                                else (hadamards qbs))
                  where p = pow2 max

-- | A quantum computation that will return a quantum integer in a state that
-- has equal probabilities of being measured in any of the state 0 to \max\.
randomQInt :: Int -> QIO QInt
randomQInt max = do
  qbs <- mkQ (reverse (int2bits max))
  applyU (randomU max qbs)
  return (QInt (reverse qbs))

-- | A quantum computation that will return a quantum integer in a state that
-- has equal probabilities of being measured in any of the state \min\ to \max\.
randomQIO :: (Int,Int) -> QIO Int
randomQIO (min,max) = do q <- randomInt (max-min)
                         return (q + min)

-- | A quantum computation that measures the outcome of "randomQInt"
randomInt :: Int -> QIO Int
randomInt max = do
  q <- randomQInt max
  measQ q

-- | A quantum computation that returns an integer that is equally likely to be
-- any number in the range 0 to \x\-1
random :: Int -> QIO Int
random x = randomInt (x-1)

-- | This function uses a Quantum computation to simulate the roll of a dice
dice :: IO Int
dice = do
  x <- run (randomInt 5)
  return (x+1)

-- | This function simulates the given number of repitions of dice rolls
dice_rolls :: Int -> IO [Int]
dice_rolls 0 = return []
dice_rolls y = do
  x <- dice
  xs <- dice_rolls (y-1)
  return (x:xs)

-- | Returns the number of occurences of 1 through 6 in the given list of Ints
occs :: [Int] -> (Int,Int,Int,Int,Int,Int)
occs rs = (rs' 1,rs' 2,rs' 3,rs' 4,rs' 5,rs' 6)
  where
-- "unterminated quasiquotation" -> space needed
   rs' x = length [ y | y <-rs, y == x]

-- | Returns the number of occurences of 1 through 6 in the given number of
-- rolls of the dice.
probs' :: Int -> IO (Int,Int,Int,Int,Int,Int)
probs' x = do
  xs <- dice_rolls x
  return (occs xs)

-- | Returns the percentage of occurences of 1 through 6, after the given number
-- of rolls of the dice.
probs :: Int -> IO (RR,RR,RR,RR,RR,RR)
probs x = do
  (a,b,c,d,e,f) <- probs' x
  return (fromIntegral a/x',fromIntegral b/x',fromIntegral c/x',fromIntegral d/x',fromIntegral e/x',fromIntegral f/x')
    where x' = fromIntegral x


{-# LANGUAGE Strict #-}
module STC.Bias where

import           Array.UnboxedArray       as AU
import           Data.Array.Repa          as R
import           Data.Complex
import           Data.List                as L
import           Data.Vector.Storable     as VS
import           Data.Vector.Unboxed      as VU
import Data.Vector.Generic as VG
import           DFT.Plan
import           Filter.Utils
import           Pinwheel.FourierSeries2D
import           STC.Point
import           STC.Utils
import           Utils.List
import FourierPinwheel.AsteriskGaussian
import Filter.Utils
import           Data.Array.Accelerate.LLVM.PTX as A
import qualified Data.Array.Accelerate          as A
import FourierMethod.FourierSeries2D
import Image.IO
import Utils.Array
import System.FilePath
import Utils.Parallel
import FourierPinwheel.GaussianEnvelopePinwheel

{-# INLINE computeBias #-}
computeBias ::
     (Num a, Unbox a, VG.Vector vector a) => Int -> Int -> [Point] -> vector a
computeBias rows cols =
  let (minR, maxR) = computeRange rows
      (minC, maxC) = computeRange cols
   in VG.convert .
      toUnboxedVector .
      AU.accum (+) 0 ((minC, minR), (maxC, maxR)) .
      L.map (\(Point x y theta scale) -> ((round x, round y), 1))

computeBiasPinwheelBasis ::
     DFTPlan -> Int -> Double -> Double -> [Point] -> IO (VS.Vector (Complex Double))
computeBiasPinwheelBasis plan numR2Freq sigma period points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      envelope =
        computeUnboxedS $
        analyticalFourierCoefficients2 numR2Freq 1 0 0 sigma period (period * sqrt 2)
  in dftExecute plan (DFTPlanID DFT1DG [numR2Freq, numR2Freq] [0, 1]) .
     VS.convert .
     toUnboxed .
     computeS .
     makeFilter2D .
     R.zipWith (*) envelope .
     fromListUnboxed (Z :. numR2Freq :. numR2Freq) $
     [ L.foldl'
       (\b (Point x y theta _) ->
          b + (cis $ (freqX * x + freqY * y) * (-2) * pi / period))
       0
       points
     | freqY <- r2Freqs
     , freqX <- r2Freqs
     ]

computeBiasPinwheelBasis1 ::
     DFTPlan -> Int -> Double -> Double -> [Point] -> VU.Vector (Complex Double)
computeBiasPinwheelBasis1 plan numR2Freq sigma period points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      envelope =
        toUnboxed . computeS $
        analyticalFourierCoefficients2 numR2Freq 1 0 0 sigma period (period * sqrt 2)
  in VU.zipWith (*) envelope . VU.fromList $
     [ L.foldl'
       (\b (Point x y theta _) ->
          b + (cis $ (freqX * x + freqY * y) * (-2) * pi / period))
       0
       points
     | freqY <- r2Freqs
     , freqX <- r2Freqs
     ]

computeBiasGaussian ::
     DFTPlan
  -> Int
  -> Double
  -> Double
  -> [Point]
  -> IO (VS.Vector (Complex Double))
computeBiasGaussian plan numR2Freq sigma period points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      freqCenter = div numR2Freq 2
      envelope =
        fromFunction (Z :. numR2Freq :. numR2Freq) $ \(Z :. i :. j) ->
          (exp $
           (fromIntegral $ (i - freqCenter) ^ 2 + (j - freqCenter) ^ 2) *
           (sigma ^ 2) /
           (-2)) *
          (sigma ^ 2) /
          (2 * pi) :+
          0
  in dftExecute plan (DFTPlanID DFT1DG [numR2Freq, numR2Freq] [0, 1]) .
     VS.convert .
     toUnboxed .
     computeS .
     makeFilter2D .
     R.zipWith (*) envelope . fromListUnboxed (Z :. numR2Freq :. numR2Freq) $
     [ L.foldl'
       (\b (Point x y theta _) ->
          b + (cis $ (freqX * x + freqY * y) * (-2) * pi / period))
       0
       points
     | freqY <- r2Freqs
     , freqX <- r2Freqs
     ]

computeBiasGaussian1 ::
     DFTPlan -> Int -> Double -> Double -> [Point] -> VU.Vector (Complex Double)
computeBiasGaussian1 plan numR2Freq sigma period points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      freqCenter = div numR2Freq 2
      envelope =
        toUnboxed . computeS . fromFunction (Z :. numR2Freq :. numR2Freq) $ \(Z :. i :. j) ->
          (exp $
           (fromIntegral $ (i - freqCenter) ^ 2 + (j - freqCenter) ^ 2) *
           (sigma ^ 2) /
           (-2)) *
          (sigma ^ 2) /
          (2 * pi) :+
          0
  in VU.zipWith (*) envelope . VU.fromList $
     [ L.foldl'
       (\b (Point x y theta _) ->
          b + (cis $ (freqX * x + freqY * y) * (-2) * pi / period))
       0
       points
     | freqY <- r2Freqs
     , freqX <- r2Freqs
     ]

computeBiasFourierPinwheel ::
     DFTPlan
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Point]
  -> IO (VS.Vector (Complex Double), VS.Vector (Complex Double))
computeBiasFourierPinwheel plan numR2Freqs thetaFreq alpha periodR2 periodEnv radius stdTheta points = do
  asteriskGaussianVec <-
    asteriskGaussianLowPass
      plan
      numR2Freqs
      thetaFreq
      alpha
      periodR2
      periodEnv
      stdTheta
      radius
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
      numThetaFreq = 2 * thetaFreq + 1
      shiftVec =
        VS.concat . L.replicate numThetaFreq . VS.fromList $
        [ L.foldl'
          (\b (Point x y theta scale) ->
             b + cis (-(freqX * x + freqY * y) * 2 * pi / periodR2))
          0
          points
        | freqY <- r2Freqs
        , freqX <- r2Freqs
        ]
      biasVec = VS.zipWith (*) asteriskGaussianVec shiftVec
      -- centerFreq = div numR2Freqs 2
      -- gaussian2D =
      --   VU.convert .
      --   toUnboxed . computeS . fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i' :. j') ->
      --     let i = i' - centerFreq
      --         j = j' - centerFreq
      --      in exp
      --           (pi * fromIntegral (i ^ 2 + j ^ 2) /
      --            ((-1) * periodR2 ^ 2 * stdR2 ^ 2)) /
      --         (2 * pi * stdR2 ^ 2) :+
      --         0
      -- zeroVec = VS.replicate (numR2Freqs ^ 2) 0
      -- shiftVec =
      --   VS.concat .
      --   L.map
      --     (\af ->
      --        if af == 0
      --          then VS.zipWith (*) gaussian2D . VS.fromList $
      --               [ L.foldl'
      --                 (\b (Point x y theta scale) ->
      --                    b + cis (-(freqX * x + freqY * y) * 2 * pi / periodR2))
      --                 0
      --                 points
      --               | freqY <- r2Freqs
      --               , freqX <- r2Freqs
      --               ]
      --          else zeroVec) $
      --   [-thetaFreq .. thetaFreq]
  biasVecF <-
    dftExecute
      plan
      (DFTPlanID DFT1DG [numThetaFreq, numR2Freqs, numR2Freqs] [0, 1, 2]) .
    VU.convert .
    toUnboxed .
    computeS .
    makeFilter3D .
    fromUnboxed (Z :. numThetaFreq :. numR2Freqs :. numR2Freqs) . VS.convert $
    -- shiftVec
    biasVec
  return (biasVec, biasVecF)
  -- return (shiftVec, biasVecF)

computeBiasFourierPinwheelFull ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> (VU.Vector (Complex Double))
  -> [Point]
  -> IO (VS.Vector (Complex Double), VS.Vector (Complex Double))
computeBiasFourierPinwheelFull plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv radius stdTheta stdR stdR2 asteriskGaussianVec points = do
  let -- asteriskGaussianVec =
      --   gaussianPinwheel
      --     numR2Freqs
      --     periodR2
      --     stdR2
      --     10
      --     thetaFreq
      --     rFreq
      --     periodEnv
      --     stdTheta
      --     stdR
      r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
      numThetaFreq = 2 * thetaFreq + 1
      numRFreq = 2 * rFreq + 1
      shiftArr =
        fromListUnboxed (Z :. numR2Freqs :. numR2Freqs) .
        parMap
          rdeepseq
          (\(freqY, freqX) ->
             L.foldl'
               (\b (Point x y theta scale) ->
                  b + cis (-(freqX * x + freqY * y) * 2 * pi / periodR2))
               0
               points) $
        [(freqY, freqX) | freqY <- r2Freqs, freqX <- r2Freqs]
      asteriskGaussianArr =
        fromUnboxed
          (Z :. numRFreq :. numThetaFreq :. numR2Freqs :. numR2Freqs)
          asteriskGaussianVec
  biasArr <-
    computeUnboxedP . R.traverse2 asteriskGaussianArr shiftArr const $ \fA fS idx@(Z :. _ :. _ :. i :. j) ->
      fA idx * fS (Z :. i :. j)
  biasVecF <-
    dftExecute
      plan
      (DFTPlanID
         DFT1DG
         [numRFreq, numThetaFreq, numR2Freqs, numR2Freqs]
         [0, 1, 2, 3]) .
    VU.convert . toUnboxed . computeS . makeFilter4D $
    biasArr
  return
    ( VU.convert . toUnboxed . computeUnboxedS . makeFilter2D $ biasArr
    , biasVecF)

computeBiasFourierPinwheelKoffkaCorss ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> (VU.Vector (Complex Double))
  -> [Point]
  -> IO (VS.Vector (Complex Double), VS.Vector (Complex Double))
computeBiasFourierPinwheelKoffkaCorss plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv radius stdTheta stdR stdR2 asteriskGaussianVec points
      -- asteriskGaussianVec =
      --   gaussianPinwheel1
      --     numR2Freqs
      --     periodR2
      --     stdR2
      --     10
      --     thetaFreq
      --     rFreq
      --     periodEnv
      --     stdTheta
      --     stdR
 = do
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
      thetaFreqs = [fromIntegral (-thetaFreq) .. fromIntegral thetaFreq]
      numThetaFreq = 2 * thetaFreq + 1
      numRFreq = 2 * rFreq + 1
      halfPI = 0.5 * pi
      centerFreqR2 = div numR2Freqs 2
      shiftArr =
        fromFunction (Z :. numThetaFreq :. numR2Freqs :. numR2Freqs) $ \(Z :. tf :. yf :. xf) ->
          let freqTheta = fromIntegral (tf - thetaFreq)
              freqY = fromIntegral (yf - centerFreqR2)
              freqX = fromIntegral (xf - centerFreqR2)
           in L.foldl'
                (\b (Point x y theta scale) ->
                   let phi =
                         if abs x < abs y
                           then atan2 0 x
                           else if abs x > abs y
                                  then atan2 y 0
                                  else error
                                         "computeBiasFourierPinwheelKoffkaCorss: x == y"
                    in b +
                       cis (-(freqX * x + freqY * y) * 2 * pi / periodR2) *
                       (cis (-freqTheta * (phi + halfPI)) +
                        cis (-freqTheta * (phi - halfPI))))
                0
                points
      -- shiftArr =
      --   fromListUnboxed (Z :. numThetaFreq :. numR2Freqs :. numR2Freqs) .
      --   parMap
      --     rdeepseq
      --     (\(thetaFreq, freqY, freqX) ->
      --        L.foldl'
      --          (\b (Point x y theta scale) ->
      --             let phi =
      --                   if abs x < abs y
      --                     then atan2 0 x
      --                     else if abs x > abs y
      --                            then atan2 y 0
      --                            else error
      --                                   "computeBiasFourierPinwheelKoffkaCorss: x == y"
      --              in b +
      --                 cis (-(freqX * x + freqY * y) * 2 * pi / periodR2) *
      --                 (cis (-thetaFreq * (phi + halfPI)) +
      --                  cis (-thetaFreq * (phi - halfPI))))
      --          0
      --          points) $
      --   [ (thetaFreq, freqY, freqX)
      --   | thetaFreq <- thetaFreqs
      --   , freqY <- r2Freqs
      --   , freqX <- r2Freqs
      --   ]
      asteriskGaussianArr =
        fromUnboxed
          (Z :. numRFreq :. numThetaFreq :. numR2Freqs :. numR2Freqs)
          asteriskGaussianVec
  biasArr <-
    computeUnboxedP . R.traverse2 asteriskGaussianArr shiftArr const $ \fA fS idx@(Z :. _ :. k :. i :. j) ->
      fA idx * fS (Z :. k :. i :. j)
  biasVecF <-
    (computeUnboxedP . makeFilter4D $ biasArr) >>=
    dftExecute
      plan
      (DFTPlanID
         DFT1DG
         [numRFreq, numThetaFreq, numR2Freqs, numR2Freqs]
         [0, 1, 2, 3]) .
    VU.convert . toUnboxed
  biasVec <- VU.convert . toUnboxed <$> computeUnboxedP (makeFilter2D biasArr)
  return (biasVec, biasVecF)

computeBiasFourierPinwheelDiscrete ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Point]
  -> IO (VS.Vector (Complex Double), VS.Vector (Complex Double))
computeBiasFourierPinwheelDiscrete plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv radius stdTheta stdR stdR2 points = do
  let biasVec =
        VU.convert .
        toUnboxed .
        computeS . makeFilter2D . fromUnboxed (Z :. numR2Freqs :. numR2Freqs) $
        computeBias numR2Freqs numR2Freqs points
      bias = VS.concat . L.replicate (2 * thetaFreq + 1) $ biasVec
      zeroVec = VS.replicate (VS.length biasVec) 0
  biasVecF <-
    dftExecute plan (DFTPlanID DFT1DG [numR2Freqs, numR2Freqs] [0, 1]) biasVec
  let biasF =
        VS.concat .
        L.map
          (\angularFreq ->
             if angularFreq == 0
               then biasVecF
               else zeroVec) $
        [-thetaFreq .. thetaFreq]
  return (bias, biasF)

computeBiasFourierPinwheelDiscrete' ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Point]
  -> IO (VS.Vector (Complex Double), VS.Vector (Complex Double))
computeBiasFourierPinwheelDiscrete' plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv radius stdTheta stdR stdR2 points = do
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
      shiftArr =
        fromListUnboxed (Z :. numR2Freqs :. numR2Freqs) .
        parMap
          rdeepseq
          (\(freqY, freqX) ->
             L.foldl'
               (\b (Point x y theta scale) ->
                  b + cis (-(freqX * (fromIntegral (round x :: Int)) + freqY * (fromIntegral (round y :: Int))) * 2 * pi / periodR2))
               0
               points) $
        [(freqY, freqX) | freqY <- r2Freqs, freqX <- r2Freqs]
      shiftVec =
        VU.convert . toUnboxed . computeUnboxedS . makeFilter2D $ shiftArr
      centerFreq = div numR2Freqs 2
      envelope =
        fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
          if i == centerFreq && j == centerFreq
            then 0
            else let r =
                       sqrt . fromIntegral $
                       (i - centerFreq) ^ 2 + (j - centerFreq) ^ 2
                  in (1 / (r * periodR2)) :+ 0
      envelopeShiftVec =
        VU.convert . toUnboxed . computeS . makeFilter2D $
        shiftArr -- *^ (centerHollowArray numR2Freqs envelope)
      zeroVec = VS.replicate (VS.length shiftVec) 0
  biasVec <-
    dftExecute
      plan
      (DFTPlanID IDFT1DG [numR2Freqs, numR2Freqs] [0, 1])
      envelopeShiftVec
  let bias = VS.concat . L.replicate (2 * thetaFreq + 1) $ biasVec
      biasF =
        VS.concat .
        L.map
          (\angularFreq ->
             if angularFreq == 0
               then shiftVec
               else zeroVec) $
        [-thetaFreq .. thetaFreq]
  return (bias, biasF)

-- computeBiasFourierPinwheel ::
--      DFTPlan
--   -> Int
--   -> Int
--   -> Double
--   -> Double
--   -> Double
--   -> Double
--   -> Double
--   -> [Point]
--   -> FilePath
--   -> [PTX]
--   -> Int
--   -> Int
--   -> IO (VS.Vector (Complex Double))
-- computeBiasFourierPinwheel plan numR2Freqs thetaFreq alpha periodR2 periodEnv stdR2 stdTheta points folderPath ptxs numPoints numBatch = do
--   let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
--       numThetaFreq = 2 * thetaFreq + 1
--       gaussianVec =
--         asteriskGaussian1 numR2Freqs thetaFreq alpha periodR2 periodEnv stdR2
--       shiftCoef =
--         [ L.foldl'
--           (\b (Point x y theta scale) ->
--              b + cis (-(freqX * x + freqY * y) * 2 * pi / periodR2))
--           0
--           points
--         | freqY <- r2Freqs
--         , freqX <- r2Freqs
--         ]
--       shiftVec = VS.concat . L.replicate numThetaFreq . VS.fromList $ shiftCoef
--       biasVec = VS.zipWith (*) (VS.concat gaussianVec) shiftVec
--       biasMat =
--         A.transpose .
--         A.use .
--         A.fromList (A.Z A.:. numThetaFreq A.:. (numR2Freqs ^ 2)) . VS.toList $
--         biasVec
--   biasR2 <-
--     computeFourierSeriesR2StreamAcc
--       ptxs
--       numR2Freqs
--       numPoints
--       numThetaFreq
--       periodR2
--       1
--       numBatch
--       biasMat
--   plotImageRepa (folderPath </> "AsteriskBias.png") .
--     ImageRepa 8 .
--     fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
--     VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
--     biasR2
--   let centerFreq = div numR2Freqs 2
--       shiftArr = fromListUnboxed (Z :. numR2Freqs :. numR2Freqs) shiftCoef
--       gaussian2D =
--         computeUnboxedS . fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i' :. j') ->
--           let i = i' - centerFreq
--               j = j' - centerFreq
--            in (exp
--                  ((pi * (fromIntegral $ i ^ 2 + j ^ 2)) /
--                   ((-1) * periodR2 ^ 2 * stdR2 ^ 2))) /
--               (2 * pi * stdR2 ^ 2) :+
--               0
--       gaussian2DMat =
--         A.use .
--         A.fromList (A.Z A.:. (numR2Freqs ^ 2) A.:. (1 :: Int)) .
--         R.toList . R.zipWith (*) shiftArr $
--         gaussian2D
--   gaussian2DR2 <-
--     computeFourierSeriesR2StreamAcc
--       ptxs
--       numR2Freqs
--       numPoints
--       1
--       periodR2
--       1
--       numBatch
--       gaussian2DMat
--   plotImageRepa (folderPath </> "Gaussian2D.png") .
--     ImageRepa 8 .
--     fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
--     VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
--     gaussian2DR2
--   let planID = DFTPlanID DFT1DG [numR2Freqs, numR2Freqs] [0, 1]
--       inversePlanID = DFTPlanID IDFT1DG [numR2Freqs, numR2Freqs] [0, 1]
--   gaussian2DF <-
--     dftExecute plan planID . VU.convert . toUnboxed . computeS . makeFilter2D $
--     gaussian2D
--   asteriskGaussianF <- dftExecuteBatchP plan planID gaussianVec
--   outputVec <-
--     fmap VS.concat .
--     dftExecuteBatchP plan inversePlanID .
--     parMap rdeepseq (VS.zipWith (*) gaussian2DF) $
--     asteriskGaussianF
--   let outputMat =
--         A.transpose .
--         A.use .
--         A.fromList (A.Z A.:. numThetaFreq A.:. (numR2Freqs ^ 2)) . VS.toList $
--         outputVec
--   outputR2 <-
--     computeFourierSeriesR2StreamAcc
--       ptxs
--       numR2Freqs
--       numPoints
--       numThetaFreq
--       periodR2
--       1
--       numBatch
--       outputMat
--   plotImageRepa (folderPath </> "AsteriskGaussian.png") .
--     ImageRepa 8 .
--     fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
--     VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
--     outputR2
--   dftExecute
--     plan
--     (DFTPlanID DFT1DG [numThetaFreq, numR2Freqs, numR2Freqs] [0, 1, 2]) .
--     VU.convert .
--     toUnboxed .
--     computeS .
--     makeFilter3D .
--     fromUnboxed (Z :. numThetaFreq :. numR2Freqs :. numR2Freqs) . VS.convert $
--     -- shiftVec
    -- VS.zipWith (*) outputVec shiftVec

-- zeroVec = VS.replicate (numR2Freqs ^ 2) 0
-- shiftVec =
--   VS.concat .
--   L.map
--     (\angularFreq ->
--        if angularFreq == 0
--          then VS.fromList $
--               [ L.foldl'
--                 (\b (Point x y _ _) ->
--                    b + cis (-(freqX * x + freqY * y) * 2 * pi / periodR2))
--                 0
--                 points
--               | freqY <- r2Freqs
--               , freqX <- r2Freqs
--               ]
--          else zeroVec) $
--   [-thetaFreq .. thetaFreq]


{-# LANGUAGE FlexibleContexts, ConstraintKinds, DataKinds #-}
import GHC.TypeLits
import Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra as LA
import System.Environment
import Text.Printf (printf)

-- TODO: Learn about statically typed matrices.
-- TODO: Figure out how to avoid converting b/t lists and matrices.
--       (Will basically require slicing and mapping over rows/colmns).
--
--       Solution: extractRows and extractCols after the predictions line.
--       Could get rid of the toColumns entry as well and separately extract.
--       (Not really worth the time.)

-- Create design matrix for polynomial powers.
-- (or, just use Vector Double. But we want to keep generality.)
-- expand :: (Element t, Field t, Num (Vector t)) => Vector t -> Int -> Matrix t
expand :: (KnownNat n, KnownNat d) => R n -> L n d
expand x deg = fromColumns $ map (x^) [0..deg]

regress :: (Element t, Field t, Num (Vector t)) => Vector t -> Vector t -> Int -> Vector t
regress x y deg = (expand x deg) <\> y
--                                  ^ efficient pseudoinverse

evalPoly :: (Element t, Field t, Product t, Num (Vector t)) => Vector t -> Vector t -> Vector t
evalPoly coeff x = expand x (dim coeff - 1) <> coeff

-- sumVector :: (Element t, Field t) => Vector t -> t
squareDiff :: Vector Double -> Vector Double -> Vector Double
squareDiff x y = mapVector (flip (^) 2) (x - y)
sumVector :: Vector Double -> Double
sumVector x    = x <.> ones where ones = konst 1 (dim x)
sumSquareDiff x y = sumVector $ squareDiff x y

-- NB: This requires concrete types due to subroutines called:
-- readFile   :: FilePath -> IO String
-- readMatrix :: String -> Matrix Double
readData :: String -> IO [Vector Double]
readData filename = fmap (toColumns . readMatrix) $ readFile filename

maxDeg = 5
main = do
  args <- getArgs
  if null args
     then putStrLn "Usage: polyreg datafile"
     else do [x,y] <- readData "polyreg_data.txt"
             let models      = map (regress x y) [1..maxDeg]
                 predictions = map (flip evalPoly x) models
                 residuals   = map (squareDiff y) predictions
                 predsview   = fromColumns $ [x, y] ++ predictions
                 residsview  = fromColumns $ [x, y] ++ residuals
                 totalerrs   = map sumVector residuals
                 putFormat   = putStrLn . format "  " (printf "%.2f")

             putStrLn "Modeling predictions"
             putFormat predsview

             putStrLn "Modeling residuals"
             putFormat residsview

             putStrLn "Total errors"
             putFormat $ fromRows [fromList totalerrs]


---------------------------------------------------------------------------------
-- |
-- Utility functions for matrix operations
---------------------------------------------------------------------------------
module Utils where

import           Numeric.LinearAlgebra (Container, Element, Matrix, R, cmap,
                                        sumElements)
import           Prelude               (Num, Ord, (**), (.), (*), (-), abs, flip,
                                        max, signum)

frobeniusNorm :: Matrix R -> R
frobeniusNorm = flip (**) 0.5 . sumElements . cmap ((flip (**) 2) . abs)

absL1Norm :: Matrix R -> R
absL1Norm = sumElements . cmap abs

softThreshold
  :: (Ord e, Num e, Element e, Container c e)
  => c e
  -> e
  -> c e
softThreshold xs mu' = cmap f xs
  where f x = (signum x) * max ((abs x) - mu') 0


{-# OPTIONS_GHC -Wall #-}

-- | Field classes
module Plankton.Field
  ( Semifield
  , Field
  , ExpField(..)
  , QuotientField(..)
  , BoundedField(..)
  , infinity
  , neginfinity
  , TrigField(..)
  ) where

import Data.Complex (Complex(..))
import Plankton.Additive
import Plankton.Multiplicative
import Plankton.Ring
import Protolude (Bool, Double, Float, Integer, (||))
import qualified Protolude as P

-- | A Semifield is a Field without Commutative Multiplication.
class (MultiplicativeInvertible a, Ring a) =>
      Semifield a

instance Semifield Double

instance Semifield Float

instance (Semifield a) => Semifield (Complex a)

-- | A Field is a Ring plus additive invertible and multiplicative invertible operations.
--
-- A summary of the rules inherited from super-classes of Field
--
-- > zero + a == a
-- > a + zero == a
-- > (a + b) + c == a + (b + c)
-- > a + b == b + a
-- > a - a = zero
-- > negate a = zero - a
-- > negate a + a = zero
-- > a + negate a = zero
-- > one * a == a
-- > a * one == a
-- > (a * b) * c == a * (b * c)
-- > a * (b + c) == a * b + a * c
-- > (a + b) * c == a * c + b * c
-- > a * zero == zero
-- > zero * a == zero
-- > a * b == b * a
-- > a / a = one
-- > recip a = one / a
-- > recip a * a = one
-- > a * recip a = one
class (AdditiveGroup a, MultiplicativeGroup a, Ring a) =>
      Field a

instance Field Double

instance Field Float

instance (Field a) => Field (Complex a)

-- | A hyperbolic field class
--
-- > sqrt . (**2) == identity
-- > log . exp == identity
-- > for +ive b, a != 0,1: a ** logBase a b ≈ b
class (Field a) =>
      ExpField a where
  exp :: a -> a
  log :: a -> a
  logBase :: a -> a -> a
  logBase a b = log b / log a
  (**) :: a -> a -> a
  (**) a b = exp (log a * b)
  sqrt :: a -> a
  sqrt a = a ** (one / (one + one))

instance ExpField Double where
  exp = P.exp
  log = P.log
  (**) = (P.**)

instance ExpField Float where
  exp = P.exp
  log = P.log
  (**) = (P.**)

-- | todo: bottom is here somewhere???
instance (TrigField a, ExpField a) => ExpField (Complex a) where
  exp (rx :+ ix) = exp rx * cos ix :+ exp rx * sin ix
  log (rx :+ ix) = log (sqrt (rx * rx + ix * ix)) :+ atan2 ix rx

-- | quotient fields explode constraints if they allow for polymorphic integral types
--
-- > a - one < floor a <= a <= ceiling a < a + one
-- > round a == floor (a + one/(one+one))
class (Field a) =>
      QuotientField a where
  round :: a -> Integer
  ceiling :: a -> Integer
  floor :: a -> Integer
  (^^) :: a -> Integer -> a

instance QuotientField Float where
  round = P.round
  ceiling = P.ceiling
  floor = P.floor
  (^^) = (P.^^)

instance QuotientField Double where
  round = P.round
  ceiling = P.ceiling
  floor = P.floor
  (^^) = (P.^^)

-- | A bounded field includes the concepts of infinity and NaN, thus moving away from error throwing.
--
-- > one / zero + infinity == infinity
-- > infinity + a == infinity
-- > isNaN (infinity - infinity)
-- > isNaN (infinity / infinity)
-- > isNaN (nan + a)
-- > zero / zero != nan
--
-- Note the tricky law that, although nan is assigned to zero/zero, they are never-the-less not equal. A committee decided this.
class (Field a) =>
      BoundedField a where
  maxBound :: a
  maxBound = one / zero
  minBound :: a
  minBound = negate (one / zero)
  nan :: a
  nan = zero / zero
  isNaN :: a -> Bool

-- | prints as `Infinity`
infinity :: BoundedField a => a
infinity = maxBound

-- | prints as `-Infinity`
neginfinity :: BoundedField a => a
neginfinity = minBound

instance BoundedField Float where
  isNaN = P.isNaN

instance BoundedField Double where
  isNaN = P.isNaN

-- | todo: work out boundings for complex
-- as it stands now, complex is different eg
--
-- > one / (zero :: Complex Float) == nan
instance (BoundedField a) => BoundedField (Complex a) where
  isNaN (rx :+ ix) = isNaN rx || isNaN ix

-- | Trigonometric Field
class (P.Ord a, Field a) =>
      TrigField a where
  pi :: a
  sin :: a -> a
  cos :: a -> a
  tan :: a -> a
  tan x = sin x / cos x
  asin :: a -> a
  acos :: a -> a
  atan :: a -> a
  sinh :: a -> a
  cosh :: a -> a
  tanh :: a -> a
  tanh x = sinh x / cosh x
  asinh :: a -> a
  acosh :: a -> a
  atanh :: a -> a
  atan2 :: a -> a -> a
  atan2 y x
    | x P.> zero = atan (y / x)
    | x P.== zero P.&& y P.> zero = pi / (one + one)
    | x P.< one P.&& y P.> one = pi + atan (y / x)
    | (x P.<= zero P.&& y P.< zero) || (x P.< zero) =
      negate (atan2 (negate y) x)
    | y P.== zero = pi -- must be after the previous test on zero y
    | x P.== zero P.&& y P.== zero = y -- must be after the other double zero tests
    | P.otherwise = x + y -- x or y is a NaN, return a NaN (via +)

instance TrigField Double where
  pi = P.pi
  sin = P.sin
  cos = P.cos
  asin = P.asin
  acos = P.acos
  atan = P.atan
  sinh = P.sinh
  cosh = P.cosh
  asinh = P.sinh
  acosh = P.acosh
  atanh = P.atanh

instance TrigField Float where
  pi = P.pi
  sin = P.sin
  cos = P.cos
  asin = P.asin
  acos = P.acos
  atan = P.atan
  sinh = P.sinh
  cosh = P.cosh
  asinh = P.sinh
  acosh = P.acosh
  atanh = P.atanh


-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Matrix.Statistics
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Basic multivariate statistics.
--

module Numeric.LinearAlgebra.Matrix.Statistics (
    defaultCovUplo,

    -- * Immutable interface
    cov,
    covWithMean,
    weightedCov,
    weightedCovWithMean,

    -- * Mutable interface
    covTo,
    covWithMeanTo,
    weightedCovTo,
    weightedCovWithMeanTo,

    ) where

import Control.Monad( when )
import Control.Monad.ST( ST )
import Data.List( foldl' )
import Text.Printf( printf )

import Numeric.LinearAlgebra.Types

import Numeric.LinearAlgebra.Vector( Vector, RVector )
import qualified Numeric.LinearAlgebra.Vector as V

import Numeric.LinearAlgebra.Matrix.Base( Matrix )
import Numeric.LinearAlgebra.Matrix.STBase( STMatrix )
import qualified Numeric.LinearAlgebra.Matrix.STBase as M
import qualified Numeric.LinearAlgebra.Matrix.Herm as M

-- | Returns the default storage scheme for covariance matrices.
defaultCovUplo :: Uplo
defaultCovUplo = Lower

-- | Returns the sample covariance matrix as a hermitian matrix with storage
-- scheme equal to 'defaultCovUplo'.  The first argument gives the dimension
-- of the vectors.
cov :: (BLAS3 e)
          => Int -> CovMethod -> [Vector e] -> Herm Matrix e
cov p t xs = M.hermCreate $ do
    c <- Herm uplo `fmap` M.new_ (p,p)
    covTo c t xs
    return c
  where
    uplo = defaultCovUplo

-- | Given the pre-computed mean, returns the sample covariance matrix
-- with storage scheme equal to 'defaultCovUplo'.
covWithMean :: (BLAS3 e)
                  => Vector e -> CovMethod -> [Vector e] -> Herm Matrix e
covWithMean mu t xs = M.hermCreate $ do
    p <- V.getDim mu
    c <- Herm uplo `fmap` M.new_ (p,p)
    covWithMeanTo c mu t xs
    return c
  where
    uplo = defaultCovUplo

-- | Returns the weighed sample covariance matrix with storage scheme equal
-- to 'defaultCovUplo'. The first argument gives the dimension of the vectors.
weightedCov :: (BLAS3 e)
                  => Int -> CovMethod -> [(Double, Vector e)] -> Herm Matrix e
weightedCov p t wxs = M.hermCreate $ do
    c <- Herm uplo `fmap` M.new_ (p,p)
    weightedCovTo c t wxs
    return c
  where
    uplo = defaultCovUplo

-- | Given the pre-computed mean, returns the weighed sample covariance matrix
-- with storage scheme equal to 'defaultCovUplo'.
weightedCovWithMean :: (BLAS3 e)
                    => Vector e -> CovMethod -> [(Double, Vector e)]
                    -> Herm Matrix e
weightedCovWithMean mu t wxs = M.hermCreate $ do
    p <- V.getDim mu
    c <- Herm uplo `fmap` M.new_ (p,p)
    weightedCovWithMeanTo c mu t wxs
    return c
  where
    uplo = defaultCovUplo

-- | Computes and copies the sample covariance matrix to the given
-- destination.
covTo :: (RVector v, BLAS3 e)
      => Herm (STMatrix s) e -> CovMethod -> [v e] -> ST s ()
covTo c@(Herm _ a) t xs = do
    (p,_) <- M.getDim a
    mu <- V.new p 1
    V.meanTo mu xs
    covWithMeanTo c mu t xs

-- | Given the pre-computed mean, computes and copies the sample covariance
-- matrix to the given destination.
covWithMeanTo :: (RVector v1, RVector v2, BLAS3 e)
              => Herm (STMatrix s) e -> v1 e -> CovMethod -> [v2 e] -> ST s ()
covWithMeanTo c@(Herm _ a) mu t xs = do
    (ma,na) <- M.getDim a
    p <- V.getDim mu

    when ((ma,na) /= (p,p)) $ error $
        printf ("covWithMeanTo"
                ++ " (Herm _ <matrix with dim (%d,%d)>)"
                ++ " <vector with dim %d>"
                ++ " _ _"
                ++ ": dimension mismatch")
               ma na p

    xt <- M.new_ (p,n)
    M.withColsM xt $ \xs' ->
        sequence_ [ V.subTo x' mu x
                  | (x,x') <- zip xs xs'
                  ]
    M.hermRankKUpdateM_ (1/df) NoTrans xt 0 c
  where

    n = length xs
    df = fromIntegral $ case t of { MLCov -> n ; UnbiasedCov -> n - 1 }

-- | Computes and copies the weighed sample covariance matrix to the
-- given destination.
weightedCovTo :: (RVector v, BLAS3 e)
              => Herm (STMatrix s) e -> CovMethod -> [(Double, v e)] -> ST s ()
weightedCovTo c@(Herm _ a) t wxs = do
    (p,_) <- M.getDim a
    mu <- V.new p 1
    V.weightedMeanTo mu wxs
    weightedCovWithMeanTo c mu t wxs

-- | Given the pre-computed mean, computes and copies the weighed sample
-- covariance matrix to the given destination.
weightedCovWithMeanTo :: (RVector v1, RVector v2, BLAS3 e)
                      => Herm (STMatrix s) e
                      -> v1 e -> CovMethod -> [(Double, v2 e)]
                      -> ST s ()
weightedCovWithMeanTo c@(Herm _ a) mu t wxs = do
    (ma,na) <- M.getDim a
    p <- V.getDim mu

    when ((ma,na) /= (p,p)) $ error $
        printf ("weightedCovWithMeanTo"
                ++ " (Herm _ <matrix with dim (%d,%d)>):"
                ++ " <vector with dim %d>"
                ++ " _ _"
                ++ " dimension mismatch")
               ma na p

    xt <- M.new_ (p,n)
    M.withColsM xt $ \xs' ->
        sequence_ [  V.subTo x' mu x
                  >> V.scaleM_ (realToFrac $ sqrt (w / invscale)) x'
                  |  (w,x,x') <- zip3 ws xs xs'
                  ]
    M.hermRankKUpdateM_ 1 NoTrans xt 0 c
  where
    (ws0,xs) = unzip wxs
    w_sum = foldl' (+) 0 ws0
    ws = if w_sum == 0 then ws0 else map (/w_sum) ws0
    w2s_sum = foldl' (+) 0 $ map (^^(2::Int)) ws
    invscale = case t of
                   MLCov -> 1
                   UnbiasedCov -> (1 - w2s_sum)
    n = length ws0


{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE FlexibleContexts #-}
module DigitRecognition where

import Numeric.LinearAlgebra as LA
import Codec.Compression.GZip (decompress)
import qualified Data.ByteString.Lazy as BS
import Data.Functor
import qualified Foreign.Storable as ForeignStorable
import GHC.Float (float2Double)
--import System.Random

render :: Integral a => a -> Char
render n = let s = " .:oO@" in s !! (fromIntegral n * length s `div` 256)

--main :: IO ()
--main = do
--  s <- decompress <$> BS.readFile "train-images-idx3-ubyte.gz"
--  l <- decompress <$> BS.readFile "train-labels-idx1-ubyte.gz"
--  n <- (`mod` 60000) <$> randomIO
--  putStr . unlines $
--    [render . BS.index s . (n * 28 ^ 2 + 16 + r * 28 +) <$> [0 .. 27] | r <- [0..27]]
-- print $ BS.index l (n + 8)

clip :: Ord c => c -> c -> c -> c
clip a b = max a . min b

computeProbabilities :: Matrix R -> Matrix R -> R -> Matrix R
computeProbabilities x theta temp = cmap (\x -> if x < 1.0e-16 then 0 else x) p / divisor
    where
        c = (rows theta >< rows x) $ concat $ repeat $ map (maximum . toList) $ toColumns $ (theta LA.<> tr x) / scalar temp
        p = cmap exp $ (theta LA.<> tr x) / scalar temp - c
        divisor = (rows theta >< rows x) $ concat $ repeat $ map (sum . toList) $ toColumns p

contributors :: (ForeignStorable.Storable R, Fractional R) => Int -> Int -> Vector Int -> Matrix R
contributors n k y = (n >< k) $ contributor k y
    where
        contributor k = concatMap (contributor' k) . toList
        contributor' k v = l ++ [1.0] ++ r
            where
                (l, r) = splitAt v $ replicate (k - 1) 0.0

computeCostFunction :: Matrix R -> Vector Int -> Matrix R -> R -> R -> R
computeCostFunction x y theta lambda temp = loss (rows x) (rows theta) y + regularisation
    where
        logProbabilities = cmap log $ cmap (clip 1e-15 (1 - 1e-15)) (computeProbabilities x theta temp)
        regularisation = lambda / 2 * sumElements (cmap (** 2) theta)
        loss n k y' = -sumElements (tr (contributors n k y') * logProbabilities) / fromIntegral n

runGradientDescentIteration :: Matrix R -> Vector Int -> Matrix R -> R -> R -> R -> Matrix R
runGradientDescentIteration x y theta alpha lambda temp = theta - scale alpha grad
    where
        n = rows x
        k = rows theta
        m = tr $ contributors n k y
        probs = cmap (clip 1e-15 (1 - 1e-15)) (computeProbabilities x theta temp)
        grad = scale lambda theta - (m - probs) LA.<> x / (scalar temp * fromIntegral n)

augmentFeatures :: (Element R, Fractional R) => Matrix R -> Matrix R
augmentFeatures x = (rows x >< 1) (repeat 1.0) ||| x

softmaxRegression :: Matrix R -> Vector Int -> R -> R -> R -> Int -> Int -> (Matrix R, [R])
softmaxRegression x y temp alpha lambda k t = iterate step ((k >< cols x') $ repeat 0.0, []) !! t
    where
        x' = augmentFeatures x
        step (theta, cs) =
            (runGradientDescentIteration x' y theta alpha lambda temp, cs ++ [computeCostFunction x' y theta lambda temp])

getClassifification :: Matrix R -> Matrix R -> R -> [Int]
getClassifification x theta temp = map maxIndex $ toColumns probs
    where
        x' = augmentFeatures x
        probs = computeProbabilities x' theta temp

computeTestError :: Matrix R -> Vector Int -> Matrix R -> R -> R
computeTestError x y theta temp = 1 - sum (compare y assignedLabels) / fromIntegral (length assignedLabels)
    where
        assignedLabels = getClassifification x theta temp
        compare y z = zipWith (\x y -> if x == y then 1 else 0) (toList y) z

buildData :: (ForeignStorable.Storable Float, Num Float) => Int -> BS.ByteString -> Matrix R
buildData length ds = cmap float2Double $ (length >< (28 ^ 2)) $ map ((/ 256) . fromIntegral) $ BS.unpack $ BS.drop 16 ds

buildLabels :: (ForeignStorable.Storable Int, Num Int) => BS.ByteString -> Vector Int
buildLabels bs = fromList $ map fromIntegral $ BS.unpack $ BS.drop 8 bs

main :: IO ()
main = do

    tr <- decompress <$> BS.readFile "train-images-idx3-ubyte.gz"
    trl <- decompress <$> BS.readFile "train-labels-idx1-ubyte.gz"
    te <- decompress <$> BS.readFile "t10k-images-idx3-ubyte.gz"
    tel <- decompress <$> BS.readFile "t10k-labels-idx1-ubyte.gz"

    let trainData = buildData 60000 tr
    let trainLabels = buildLabels trl
    let testData = buildData 10000 te
    let testLabels = buildLabels tel
    let temp = 1
    let alpha = 0.3
    let lambda = 0.0001
    let k = 10
    let iter = 150
    let (theta, cs) = softmaxRegression trainData trainLabels temp alpha lambda k iter
--    print $ trainData ? [0]
--    print $ toList trainLabels !! 0
--    print $ (computeProbabilities' (augmentFeatures trainData) theta temp) ¿ [0]
--    let theta' = runGradientDescentIteration (augmentFeatures trainData) trainLabels theta alpha lambda temp
--    print $ computeCostFunction (augmentFeatures trainData) trainLabels theta lambda temp
--    print $ theta ¿ [0]
--    print $ getClassifification testData theta temp
    print $ computeTestError testData testLabels theta temp
    -- print $ take 5 $ toList trainLabels


module Main (main) where

import           Control.Monad
import           System.IO
import           Text.Printf (printf)

import           Data.Number.LogFloat (LogFloat)
import qualified Data.Number.LogFloat as LF
import           Statistics.Distribution
import           Statistics.Distribution.Binomial

-- | Compute the probability of seeing more than @k@ blocks in @n@ slots
moreThanK ::
     Double  -- ^ Active slot coefficient
  -> Int     -- ^ Security parameter (@k@)
  -> Int     -- ^ Number of slots
  -> LogFloat
moreThanK f k n =
    LF.sum [LF.logToLogFloat $ logProbability d i | i <- [k + 1 .. n]]
  where
    d :: BinomialDistribution
    d = binomial n f

defaultS ::
     Double  -- ^ Active slot coefficient
  -> Int     -- ^ Security parameter (@k@)
  -> Int
defaultS f k = floor (fromIntegral k / f) `div` 4

main :: IO ()
main = do
    forM_ [s .. 4 * s] $ \n -> do
      putStrLn $ show n ++ "\t" ++ showLogFloat (moreThanK f k n)
      hFlush stdout
  where
    f = 0.05
    k = 2160
    s = defaultS f k

{-------------------------------------------------------------------------------
  LogFloat util
-------------------------------------------------------------------------------}

showLogFloat :: LogFloat -> String
showLogFloat lf = printf "%6.4f * 10 ^ %d" m e
  where
    (m, e) = logFloatToScientific lf

logFloatToScientific :: LogFloat -> (Double, Int)
logFloatToScientific lf = (m, e)
  where
    l :: Double
    l = LF.logFromLogFloat lf

    e :: Int
    e = floor $ l / log10

    m :: Double
    m = exp $ l - log10 * fromIntegral e

    log10 :: Double
    log10 = log 10


module Fuml.Optimisation where

import Fuml.Core
import Fuml.Optimisation.BFGS
import Fuml.Optimisation.SGD
import Numeric.LinearAlgebra
import qualified Data.Vector.Storable as VS
import Fuml.Base.Logistic
import Data.Maybe (fromMaybe)
import Control.Monad.Identity
import Data.Random

type BfgsModel = (Vector Double, Matrix Double)

data CostGradientModel a b = CostGradientModel
  { costFunction :: ([(Vector Double, a)] -> Vector Double -> Double)
  , costGradient :: ([(Vector Double, a)] -> Vector Double -> Vector Double)
  , generatePrediction :: Vector Double -> Vector Double -> b
  }

bfgsSupervisor :: BFGSOpts
               -> CostGradientModel a b
               -> Supervisor Identity a BfgsModel b
bfgsSupervisor bopts (CostGradientModel ll gradll predf) = Supervisor $ \mp theData ->
  let start = fromMaybe (VS.map (const 0) $ fst $ head theData) $ fmap fst mp
      hessInit = fromMaybe (ident $ size start) $ fmap snd mp
  in case bfgsWith bopts (ll theData) (gradll theData) start hessInit of
       Left s -> error s
       Right (p, h) -> return $ Predict (p, h) (predf p)

defBfgs = BFGSOpts 1e-7 1e-7 200

data SGDModel = SGDModel
  { sgdFit :: Vector Double
  , lossCurve :: [Double]
  }

sgdSupervisor :: SGDOpts
              -> CostGradientModel a b
              -> Supervisor RVar a SGDModel b
sgdSupervisor opts (CostGradientModel ll gradll predf) = Supervisor $ \mp theData ->
  let start = fromMaybe (VS.map (const 0) $ fst $ head theData) $ fmap sgdFit mp
      mkPredict sgdModel = Predict sgdModel (predf (sgdFit sgdModel))
      go 0 p losses = return $ SGDModel p $ reverse losses
      go n p losses = do
        nextp <- sgdEpoch opts gradll theData p
        let loss = ll theData p
        go (n-1) nextp (loss:losses)
  in fmap mkPredict $ go (epochs opts) start []

logisticCost :: Double -> CostGradientModel Bool Double
logisticCost delta = CostGradientModel (logLikelihood delta) (gradLogLikelihood delta) (\beta v -> logit $ v `dot` beta)

logisticBfgs :: Double -> Supervisor Identity Bool BfgsModel Double
logisticBfgs delta = bfgsSupervisor defBfgs (logisticCost delta)


module Binary where

import Nnaskell
import Numeric.LinearAlgebra.Data

binaryNN :: IO NN
binaryNN = loadNN "./data/binary.txt"

binaryTD :: [(Vector R, Vector R)]
binaryTD = zip (map binaryInput ns) (map binaryOutput ns)
    where ns = [0 .. 15]

binaryInput :: Int -> Vector R
binaryInput x = vector $ map fromIntegral [ x `mod` 2
                                          , (x `div` 2) `mod` 2
                                          , (x `div` 4) `mod` 2
                                          , (x `div` 8) `mod` 2]

binaryOutput :: Int -> Vector R
binaryOutput x = accum (vector $ replicate 16 0.0) (\a _ -> a) [(x, 1.0)]


{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeOperators #-}
module Filter(
      FIR(..)
    , bode
    , phasePlot
    ) where

import Signal
import Math.Polynomial
import Common
import Noise
import Fixed
import GHC.TypeLits
import SpecialInt
import Data.Bits
import Plot
import Data.Complex
import Text.Printf

import Debug.Trace

debug a = trace (show a) a

data FIR p i o where
 FIR :: (SingI (nc + ni), SingI nc, SingI ni, SingI no, SingI s, SingI r, HasDoubleRepresentation (Fixed a nc s r), Num (Fixed ag (nc + ni) s r), Num (Fixed a nc s r), Eq a, Conversion ag ao, ConvertConstraint ag ao, AMulConstraint a ag)
     => Poly (Fixed a nc s r)
     -> FIR ((Fixed ag (nc + ni) s r),(Fixed a nc s r)) (Fixed a ni s r) (Fixed ao no s r)
 FIRD :: Poly Double
      -> FIR Double Double Double

polyRepresentation :: FIR p i o -> Poly Double
polyRepresentation (FIRD p) = p
polyRepresentation (FIR f) = fmap toDouble f

instance Show p => Show (FIR p i o) where
  show (FIR c) = show c
  show (FIRD d) = show d

delay :: s -> Signal s -> Signal s
delay a s = consS a s

guardType :: FIR ((Fixed ag ng s r),(Fixed a nc s r)) i (Fixed ao no s r)
          -> Fixed ag n s r
          -> Fixed ag n s r
guardType _ = id

filterDWith :: [Double]
            -> Signal Double
            -> Signal Double
filterDWith (a:l) s | null l = mapS (a*) s
                    | otherwise = zipWithS (+) (mapS (a*) s) (filterDWith l (delay 0 s))
filterDWith [] s = s

_filterWith :: (AMulConstraint a ag, SingI (na + ni), SingI s, SingI r, SingI na, SingI ni, Num (Fixed a ni s r), Num (Fixed ag (na + ni) s r))
            => FIR ((Fixed ag ng s r),(Fixed a na s r)) (Fixed a ni s r) (Fixed ao no s r)
            -> [Fixed a na s r]
            -> Signal (Fixed a ni s r)
            -> Signal (Fixed ag (na + ni) s r)
_filterWith f (a:l) s | null l = mapS (amul a) s
                      | otherwise = zipWithS (+) (mapS (amul a) s) (_filterWith f l (delay 0 s))
_filterWith _ [] s = error "Empty polynomial can't be used for filtering"

filterWith :: (Conversion ag ao, ConvertConstraint ag ao, Num (Fixed a ni s r), Num (Fixed ag (na + ni) s r), SingI r, SingI s, SingI (na + ni), SingI ni, SingI na, SingI no, AMulConstraint a ag)
           => FIR ((Fixed ag ng s r),(Fixed a na s r)) (Fixed a ni s r) (Fixed ao no s r)
           -> [Fixed a na s r]
           -> Signal (Fixed a ni s r)
           -> Signal (Fixed ao no s r)
filterWith f l = mapS convert . _filterWith f l

instance Structure FIR where
    doubleVersion (FIR p) = FIRD $ fmap toDouble p
    transferFunction f@(FIR p) s =
        let c = polyCoeffs LE $ p
        in
        filterWith f c s
    transferFunction (FIRD p) s =
        let c = polyCoeffs LE p
        in
        filterDWith c s

bode :: (Sample o,Sample i)
     => FIR p i o -> StyledSignal Double Double
bode f =
    let fd = polyRepresentation f
        log10 x = log x / log 10
        polyF = evalPoly (fmap (:+ 0) fd)
        transferFunction = \t -> 20 * log10 (magnitude . polyF . cis $ -t)
        dw = 0.01
        values = map transferFunction [0,dw..pi]
        nb = length values
        ticksWithPhase n = printf "%.2f pi" (n / fromIntegral nb :: Double)
    in
    (plotSpectrum nb [ AS $ fromListS 0 values])
                 `withNewStyle` (\s -> s { verticalLabel = Just "Filter Gain (dB)"
                                         , title = Just "Bode Plot"
                                         , horizontalTickRepresentation = ticksWithPhase
                                         })

phasePlot :: (Sample o,Sample i)
     => FIR p i o -> StyledSignal Double Double
phasePlot f =
    let fd = polyRepresentation f
        log10 x = log x / log 10
        polyF = evalPoly (fmap (:+ 0) fd)
        transferFunction = \t -> phase . polyF . cis $ -t
        dw = 0.01
        values = map transferFunction [0,dw..pi]
        nb = length values
        ticksWithPhase n = printf "%.2f pi" (n / fromIntegral nb :: Double)
    in
    (plotSpectrum nb [ AS $ fromListS 0 values])
                 `withNewStyle` (\s -> s { verticalLabel = Just "Filter Phase"
                                         , title = Just "Phase Plot"
                                         , horizontalTickRepresentation = ticksWithPhase
                                         })


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}

module QUtil where

import Data.Complex
import Data.List

import Numeric.LinearAlgebra as L

import LinAlg
import BraKet
import Operators

-- Utility functions --
takeKet :: Element c => Int -> Ket c -> Ket c
takeKet n (Ket coeffs) = Ket $ L.subVector 0 n coeffs

-- Expected value of an operator O computed as Tr(rho <> O)
expectedValue :: MatOp -> [Complex Double] -> [Ket (Complex Double)] -> Complex Double
expectedValue op ps kets = trace $ op Prelude.<> densityMatrix ps kets

-- Commutator of two operators
commutator :: MatOp -> MatOp -> MatOp
commutator a b = (a Prelude.<> b) <~> (b Prelude.<> a)

-- normalization function
normalize :: Hilbert v => v -> v
normalize v = ((1.0 :+ 0.0) / (norm v :+ 0.0)) <**> v

-- sized vacuum Fock state
vacuum :: Int -> Ket (Complex Double)
vacuum n = Ket . L.fromList $ 1.0 : replicate (n-1) 0.0

-- n-photon Fock basis state
fockN :: Int -> Int -> Ket (Complex Double)
fockN m n = QUtil.normalize $ act (foldl1' (Prelude.<>) $ replicate n (create m)) (vacuum m)

-- Coherent state
coherent :: Int -> Complex Double -> Ket (Complex Double)
coherent m alpha = act (displacement m alpha) $ vacuum m


{-# LANGUAGE ConstraintKinds #-}

module Projection
( projectL2
)
where

import           Basis
import           Element
import           Node
import qualified Numeric.LinearAlgebra as LinAl
import           Quadrature
import           ShapeFcns

-- Synonym for numeric/fractional constraints
type FrElNuFi a = (Fractional a,LinAl.Element a,LinAl.Numeric a,LinAl.Field a)

-- Element mass integrand
elemMassIntegrand :: (Element e,ShapeFcn s,Basis b,FrElNuFi a) => e a -> s b -> [a] -> LinAl.Matrix a
elemMassIntegrand elem shpFcn xi = LinAl.fromLists [[
  ShapeFcns.n shpFcn xi i [0] * ShapeFcns.n shpFcn xi j [0] * Element.computeJacobianDet elem shpFcn xi
  | i <- idxRange] | j <- idxRange]
  where
    idxRange = [0..(ShapeFcns.getShapeFcnOrder shpFcn)]

-- This function computes the element load integrand when given a function
-- to integrate over an element.
elemLoadIntegrand :: (Element e,ShapeFcn s,Basis b,FrElNuFi a) => e a -> s b -> ([a] -> a) -> [a] -> LinAl.Matrix a
elemLoadIntegrand elem shpFcn f xi = LinAl.scale (computeJacobianDet elem shpFcn xi) (LinAl.fromLists $ map (:[]) [f physCoords * n shpFcn xi i [0] | i <- [0..(getShapeFcnOrder shpFcn)]])
  where
    physCoords = [sum $ zipWith (*) (map (head . nodeCoordinates) (getElementNodes elem)) (map (\j -> n shpFcn xi j [0]) [0..(getShapeFcnOrder shpFcn)])]

-- Function for projecting a function onto a basis over a single element.
projectL2 :: (Element e,ShapeFcn s,Basis b,FrElNuFi a) => e a -> s b -> Int -> ([a] -> a) -> Maybe (LinAl.Matrix a)
projectL2 elem shpFcn ngpts f = LinAl.linearSolve (integrate 1 ngpts (elemMassIntegrand elem shpFcn)) (integrate 1 ngpts (elemLoadIntegrand elem shpFcn f))


import Data.Complex

main :: IO ()
main = do let res = dft $ enumerate $ map (\x -> x :+ 0) [1..100]
       	  putStrLn $ show res

enumerate :: Integral b => [a] -> [(a, b)]
enumerate l = zip l [0..]

dft :: (RealFloat a, Integral b) => [(Complex a, b)] -> [Complex a]
dft x = map (\(_, k) -> singleDFT k x) x

singleDFT :: (RealFloat a, Integral b) => b -> [(Complex a, b)] -> Complex a
singleDFT k x = sum $ map (\(e, n) -> (e *) $ fCoeff k n $ length x) $ x

fCoeff :: (RealFloat a, Integral b) => b -> b -> b -> Complex a
fCoeff k n n_tot = exp (0.0 :+ ((-2.0) * pi * (fromIntegral $ k * n) /
       	   	       	       	       	      (fromIntegral n_tot)))


module Statistics.Dispersion where

import Statistics.Center

variance :: (Fractional a) => [a] -> a
variance vals = (sum $ zipWith (*) deviations deviations)/n
    where n = (fromIntegral $ length vals)
          mu = arithmeticMean vals
          deviations = map (\x -> x-mu) vals

stdev :: (Floating a) => [a] -> a
stdev vals = sqrt $ variance vals