#!/usr/bin/env stack
-- stack runghc --package reanimate
{-# LANGUAGE OverloadedStrings #-}
module Main (main) where

import           Data.Complex
import qualified Data.Text as T

import           Graphics.SvgTree
import           Reanimate.Driver (reanimate)
import           Reanimate.LaTeX  (latex)
import           Reanimate.Animation
import           Reanimate.Svg
import           Reanimate.Signal
import           Reanimate.Constants

waveMultiplier :: Int
-- waveMultiplier = 1 -- Sawtooth wave
waveMultiplier = 2 -- Square wave

main :: IO ()
main = reanimate $
  fourierAnimation 1 `seqA`
  fourierAnimation 2 `seqA`
  fourierAnimation 3 `seqA`
  fourierAnimation 5 `seqA`
  fourierAnimation 10 `seqA`
  fourierAnimation 25 `seqA`
  fourierAnimation 50 `seqA`
  fourierAnimation 100

sWidth :: Double
sWidth = 0.02

fourierAnimation :: Int -> Animation
fourierAnimation nCircles = repeatA 2 $ mkAnimation 3 $ \t ->
    let phi = fromToS 0 (2*pi) t
    in mkGroup
    [ mkBackground "black"
    , translate (-screenWidth/4) 0 $ mkGroup
      [ drawNCircles nCircles phi
      , withStrokeColor "white" $
        withStrokeWidth sWidth $
        withFillOpacity 0 $
        translate (screenWidth/4) 0 $
        mkCirclePath nCircles phi ]
    , withStrokeWidth sWidth $
      withFillColor "white" $
      translate (-screenWidth/8*3) (screenHeight/8*3) $
      latex $ T.pack $ "Circles: " ++ show nCircles ]

drawNCircles :: Int -> Double -> Tree
drawNCircles totalCircles phi = mkGroup
    [ worker circles
    , let x :+ y = sum circles in
      withStrokeWidth sWidth $
      withStrokeColor "white" $
      mkLine (x, y) (screenWidth/4, y) ]
  where
    circles = [ nthCircle n phi | n <- [0..totalCircles-1] ]
    worker [] = None
    worker (x :+ y : rest) =
      let radius = sqrt(x*x+y*y) in
      mkGroup
      [ withStrokeWidth sWidth $
        withStrokeColor "grey" $
        withFillOpacity 0 $
        mkCircle radius
      , translate x y $ worker rest
      , withStrokeWidth sWidth $
        withStrokeColor "white" $
        mkLine (0, 0) (x, y) ]

mkCirclePath :: Int -> Double -> Tree
mkCirclePath nCircles phiOffset = mkLinePath $ take 2000 $
    zip [ 2 * i/granularity | i <- [0..]]
    $ drop (round $ (1-phiOffset/(2*pi)) * granularity) $
    cycle $ [ fourierYValue nCircles phi
    | x <- reverse [1..granularity]
    , let phi = 2*pi*(x/granularity)
    ]
  where
    granularity = 500

fourierYValue :: Int -> Double -> Double
fourierYValue n phi =
  imagPart (sum [ nthCircle i phi | i <- [0..n-1]])

nthCircle :: Int -> Double -> Complex Double
nthCircle n phi = x :+ y
  where
    n' = fromIntegral (n*waveMultiplier+1)
    x = cos (n'*phi) * radius
    y = sin (n'*phi) * radius
    radius = 2.5 * (2 / (n'*pi))


{-# LANGUAGE Trustworthy #-}
-- |
-- Module    : Statistics.Distribution.Hypergeometric.GenVar
-- Copyright : (c) 2015 Sam Rijs
--             (c) 2005 Robert Kern
--             (c) 1998 Ivan Frohne
-- License   : MIT
--
-- Maintainer  : srijs@airpost.net
-- Stability   : experimental
--

module Statistics.Distribution.Hypergeometric.GenVar (genVar) where

import Control.Monad.Primitive (PrimMonad, PrimState)

import Statistics.Distribution (ContGen(..), DiscreteGen(..))
import Statistics.Distribution.Hypergeometric (HypergeometricDistribution, hdM, hdL, hdK)

import Numeric.SpecFunctions (logGamma)

import System.Random.MWC (Gen, uniform)

d1 = 1.7155277699214135 -- 2*sqrt(2/e)
d2 = 0.8989161620588988 -- 3 - 2*sqrt(3/e)

instance ContGen HypergeometricDistribution
  where genContVar d g = return . fromIntegral =<< genVar (hdM d, hdL d, hdK d) g

instance DiscreteGen HypergeometricDistribution
  where genDiscreteVar d g = genVar (hdM d, hdL d, hdK d) g

-- | Random variates from the hypergeometric distribution /(m, l, k)/.
-- Returns the number of white balls drawn when /k/ balls
-- are drawn at random from an urn containing /m/ white
-- and /l-m/ black balls.
genVar :: (PrimMonad m, Integral a) => (a, a, a) -> Gen (PrimState m) -> m a
genVar (good, popsize, sample) gen = return . fix2 . fix1 =<< loop
  where fix1 z = if good > bad then m - z else z
        fix2 z = if m < sample then good - z else z
        mingoodbad = min good bad
        bad = popsize - good
        maxgoodbad = max good bad
        m = min sample (popsize - sample)
        gamma z = sum $ map (logGamma . fromIntegral)
          [ z + 1,     mingoodbad - z + 1
          , m - z + 1, maxgoodbad - m + z + 1 ]
        d4 = fromIntegral mingoodbad / fromIntegral popsize
        d5 = 1 - d4
        d6 = fromIntegral m * d4 + 0.5
        d7 = sqrt $ fromIntegral (popsize - m) * fromIntegral sample * d4 * d5 / fromIntegral (popsize - 1) + 0.5
        d8 = d1 * d7 + d2
        d9 = floor $ fromIntegral (m + 1) * fromIntegral (mingoodbad + 1) / fromIntegral (popsize + 2)
        d10 = gamma d9
        d11 = fromIntegral $ min (min m mingoodbad + 1) (floor (d6 + 16 * d7))
        -- 16 for 16-decimal-digit precision in d1 and d2
        loop = do
          x <- uniform gen
          y <- uniform gen
          case () of
            _  | w < 0 || w >= d11    -> loop -- fast rejection
               | x * (4 - x) - 3 <= t -> return z -- fast acceptance
               | x * (x - t) >= 1     -> loop -- fast rejection
               | 2 * (log x) <= t     -> return z -- acceptance
               | otherwise            -> loop -- rejection
               where w = d6 + d8 * (y - 0.5) / x
                     z = floor w
                     t = d10 - gamma z


module Lib
    ( someFunc
    ) where

{--import Graphics.Gnuplot.Advanced
import Graphics.Gnuplot.MultiPlot
import Graphics.Gnuplot.Terminal.X11
import Graphics.Gnuplot.Graph.TwoDimensional--}
import Numeric.LinearAlgebra
import Graphics.Plot

f :: (Floating a) => a -> a
f x = sin x

someFunc :: IO ()
someFunc = do
  let x = linspace 1000 (-pi,pi)
  mplot [f x]


{-#LANGUAGE BangPatterns #-}

module Tests.Correlation
    ( tests ) where

import Control.Arrow (Arrow(..))
import qualified Data.Vector as V
import Statistics.Correlation
import Statistics.Correlation.Kendall
import Test.QuickCheck ((==>),Property,counterexample)
import Test.Tasty
import Test.Tasty.QuickCheck
import Test.Tasty.HUnit

import Tests.ApproxEq

----------------------------------------------------------------
-- Tests list
----------------------------------------------------------------

tests :: TestTree
tests = testGroup "Correlation"
    [ testProperty "Pearson correlation"           testPearson
    , testProperty "Spearman correlation is scale invariant" testSpearmanScale
    , testProperty "Spearman correlation, nonlinear"         testSpearmanNonlinear
    , testProperty "Kendall test -- general"       testKendall
    , testCase     "Kendall test -- special cases" testKendallSpecial
    ]

----------------------------------------------------------------
-- Pearson's correlation
----------------------------------------------------------------

testPearson :: [(Double,Double)] -> Property
testPearson sample
  = (length sample > 1) ==> (exact ~= fast)
  where
    (~=) = eql 1e-12
    exact = exactPearson $ map (realToFrac *** realToFrac) sample
    fast  = pearson $ V.fromList sample

exactPearson :: [(Rational,Rational)] -> Double
exactPearson sample
  = realToFrac cov / sqrt (realToFrac (varX * varY))
  where
    (xs,ys) = unzip sample
    n       = fromIntegral $ length sample
    -- Mean
    muX  = sum xs / n
    muY  = sum ys / n
    -- Mean of squares
    muX2 = sum (map (\x->x*x) xs) / n
    muY2 = sum (map (\x->x*x) ys) / n
    -- Covariance
    cov  = sum (zipWith (*) [x - muX | x<-xs] [y - muY | y<-ys]) / n
    varX = muX2 - muX*muX
    varY = muY2 - muY*muY

----------------------------------------------------------------
-- Spearman's correlation
----------------------------------------------------------------

-- Test that Spearman correlation is scale invariant
testSpearmanScale :: [(Double,Double)] -> Double -> Property
testSpearmanScale xs a
  = and [ length xs > 1       -- Enough to calculate underflow
        , a /= 0
        , not (isNaN c1)
        , not (isNaN c2)
        , not (isNaN c3)
        , not (isNaN c4)
        ]
  ==> ( counterexample (show xs2)
      $ counterexample (show xs3)
      $ counterexample (show xs4)
      $ counterexample (show (c1,c2,c3,c4))
      $ and [ c1 == c4
           , c1 == signum a * c2
           , c1 == signum a * c3
           ]
      )
  where
    xs2 = map ((*a) *** id  ) xs
    xs3 = map (id   *** (*a)) xs
    xs4 = map ((*a) *** (*a)) xs
    c1 = spearman $ V.fromList xs
    c2 = spearman $ V.fromList xs2
    c3 = spearman $ V.fromList xs3
    c4 = spearman $ V.fromList xs4

-- Test that Spearman correlation allows to transform sample with
testSpearmanNonlinear :: [(Double,Double)] -> Property
testSpearmanNonlinear sample0
  = and [ length sample0 > 1
        , not (isNaN c1)
        , not (isNaN c2)
        , not (isNaN c3)
        , not (isNaN c4)
        ]
  ==> ( counterexample (show sample0)
      $ counterexample (show sample1)
      $ counterexample (show sample2)
      $ counterexample (show sample3)
      $ counterexample (show sample4)
      $ counterexample (show (c1,c2,c3,c4))
      $ and [ c1 == c2
            , c1 == c3
            , c1 == c4
            ]
      )
  where
    -- We need to stretch sample into [-10 .. 10] range to avoid
    -- problems with under/overflows etc.
    stretch xs
      | a == b = xs
      | otherwise = [ (x - a - 10) * 20 / (a - b) | x <- xs ]
      where
        a = minimum xs
        b = maximum xs
    sample1 = uncurry zip $ (stretch *** stretch) $ unzip sample0
    sample2 = map (exp *** id ) sample1
    sample3 = map (id  *** exp) sample1
    sample4 = map (exp *** exp) sample1
    c1 = spearman $ V.fromList sample1
    c2 = spearman $ V.fromList sample2
    c3 = spearman $ V.fromList sample3
    c4 = spearman $ V.fromList sample4

----------------------------------------------------------------
-- Kendall's correlation
----------------------------------------------------------------

testKendall :: [(Double, Double)] -> Bool
testKendall xy | isNaN r1 = isNaN r2
               | otherwise = r1 == r2
  where
    r1 = kendallBruteForce xy
    r2 = kendall $ V.fromList xy

testKendallSpecial :: Assertion
testKendallSpecial = vs @=? map (\(xs, ys) -> kendall $ V.fromList $ zip xs ys) d
  where
    (d, vs) = unzip testData
    testData :: [(([Double], [Double]), Double)]
    testData = [ (([1, 2, 3, 1, 2], [1, 2, 1, 5, 2]), -0.375)
               , (([1, 1, 1, 3, 3], [3, 3, 3, 2, 5]), 0)
               ]

kendallBruteForce :: [(Double, Double)] -> Double
kendallBruteForce xy = (n_c - n_d) / sqrt ((n_0 - n_1) * (n_0 - n_2))
  where
    allPairs = f xy
    (n_c, n_d, n_1, n_2) = foldl g (0,0,0,0) allPairs
    n_0 = fromIntegral.length $ allPairs
    g (!nc, !nd, !n1, !n2) ((x1, y1), (x2, y2))
      | (x2 - x1) * (y2 - y1) > 0 = (nc+1, nd, n1, n2)
      | (x2 - x1) * (y2 - y1) < 0 = (nc, nd+1, n1, n2)
      | otherwise = if x1 == x2
                       then if y1 == y2
                               then (nc, nd, n1+1, n2+1)
                               else (nc, nd, n1+1, n2)
                       else (nc, nd, n1, n2+1)
    f (x:xs) = zip (repeat x) xs ++ f xs
    f _ = []


module Backprop
  (BackpropNet, buildBackpropNet, logisticSigmoidAS, tanhAS, identityAS)
    where

import MatrixPlus as P
import NeuralNet

import Numeric.LinearAlgebra as N
import Test.QuickCheck
import Control.Exception
import Debug.Trace

-- {{{ Layers

-- | An individual layer in a neural network, prior to propagation
data Layer = Layer
    {
      -- The weights for this layer
      lW :: Matrix Double,
      -- The activation specification for this layer
      lAS :: ActivationSpec
    }

instance Show Layer where
    show layer = "w=" ++ show (lW layer) ++ ", activation spec=" ++ show (lAS layer)

-- | Generate a layer of the specified "size", with arbitrary values.
-- | To see sample values (of "size" 4), in GHCi type: sample (sizedArbLayer 4)
sizedArbLayer :: Int -> Gen Layer
sizedArbLayer n = do
    -- QuickCheck passes a value n >= 0, but we can't create a vector of length 0, so use n+1
    r <- choose(1, n+1)
    c <- choose(1, n+1)
    w <- sizedArbWeightMatrix r c
    s <- (arbitrary :: Gen ActivationSpec)
    return Layer{ lW=w, lAS=s }

-- | Generate a layer of the specified "size" and the specified input width, with arbitrary values.
-- | To see sample values (of input width 3, "size" 4), in GHCi type: sample (sizedArbLayer' 3 4)
sizedArbLayer' :: Int -> Int -> Gen Layer
sizedArbLayer' c n = do
    -- QuickCheck passes a value n >= 0, but we can't create a vector of length 0, so use n+1
    r <- choose(1, n+1)
    w <- sizedArbWeightMatrix r c
    s <- (arbitrary :: Gen ActivationSpec)
    return Layer{ lW=w, lAS=s }

-- | Generate a layer of arbitrary "size", with arbitrary values.
-- | To see some sample values, in GHCi type: sample' arbitrary :: IO [Layer]
instance Arbitrary Layer where
    arbitrary = sized sizedArbLayer

inputWidth :: Layer -> Int
inputWidth = cols . lW

outputWidth :: Layer -> Int
outputWidth = rows . lW

-- }}}

-- {{{ Propagation

-- | An individual layer in a neural network, after propagation but prior to backpropagation
data PropagatedLayer
    = PropagatedLayer
        {
          -- The input to this layer
          pIn :: ColumnVector Double,
          -- The output from this layer
          pOut :: ColumnVector Double,
          -- The value of the first derivative of the activation function for this layer
          pF'a :: ColumnVector Double,
          -- The weights for this layer
          pW :: Matrix Double,
          -- The activation specification for this layer
          pAS :: ActivationSpec
        }
    | PropagatedSensorLayer
        {
          -- The output from this layer
          pOut :: ColumnVector Double
        }

instance Show PropagatedLayer where
    show (PropagatedLayer x y f'a w s) =
        "in=" ++ show x
        ++ ", out=" ++ show y
        ++ ", f'(a)=" ++ show f'a
        ++ ", w=" ++ show w
        ++ ", " ++ show s
    show (PropagatedSensorLayer x) = "out=" ++ show x

    -- {{{ Testing

{-
-- | Generate a propagated hidden or output layer of the specified "size", with arbitrary values.
-- | To see sample values (of "size" 4), in GHCi type: sample (sizedArbPropagatedLayer 4)
sizedArbPropagatedLayer :: Int -> Gen PropagatedLayer
sizedArbPropagatedLayer n = do
    -- QuickCheck passes a value n >= 0, but we can't create a vector of length 0, so use n+1
    r <- choose(1, n+1)
    c <- choose(1, n+1)
    w <- sizedArbWeightMatrix r c
    y <- sizedArbColumnVector r
    return PropagatedLayer{ pOut=y
                          , pF'a=(zeroColumnVector r)
                          , pW=w
                          , pF=id
                          , pF'=id'
                          }

-- | Generate a propagated sensor layer of the specified "size", with arbitrary values.
-- | To see sample values (of "size" 4), in GHCi type: sample (sizedArbPropagatedLayer2 4)
sizedArbPropagatedLayer0 :: Int -> Gen PropagatedLayer
sizedArbPropagatedLayer0 n = do
    -- QuickCheck passes a value n >= 0, but we can't create a vector of length 0, so use n+1
    r <- choose(1, n+1)
    y <- sizedArbColumnVector r
    return PropagatedSensorLayer{ pOut=y }

-- | Generate a propagated layer of arbitrary "size", with arbitrary values.
-- | To see some sample values, in GHCi type: sample' arbitrary :: IO [PropagatedLayer]
instance Arbitrary PropagatedLayer where
    arbitrary = frequency
        [ (4, sized sizedArbPropagatedLayer)
        , (1, sized sizedArbPropagatedLayer0) ]
-}
    -- }}}

-- | Propagate the inputs through this layer to produce an output.
propagate :: PropagatedLayer -> Layer -> PropagatedLayer
propagate layerJ layerK = PropagatedLayer
        {
          pIn = x,
          pOut = y,
          pF'a = f'a,
          pW = w,
          pAS = lAS layerK
        }
  where x = pOut layerJ
        w = lW layerK
        a = w <> x
        f = asF ( lAS layerK )
        y = P.mapMatrix f a
        f' = asF' ( lAS layerK )
        f'a = P.mapMatrix f' a

-- }}}

-- {{{ Backpropagation

-- | An individual layer in a neural network, after backpropagation
data BackpropagatedLayer = BackpropagatedLayer
    {
      -- Del-sub-z-sub-l of E
      bpDazzle :: ColumnVector Double,
      -- The error due to this layer
      bpErrGrad :: ColumnVector Double,
      -- The value of the first derivative of the activation
      --   function for this layer
      bpF'a :: ColumnVector Double,
      -- The input to this layer
      bpIn :: ColumnVector Double,
      -- The output from this layer
      bpOut :: ColumnVector Double,
      -- The weights for this layer
      bpW :: Matrix Double,
      -- The activation specification for this layer
      bpAS :: ActivationSpec
    }

instance Show BackpropagatedLayer where
    show layer =
        "dazzle=" ++ show (bpDazzle layer)
        ++ ", grad=" ++ show (bpErrGrad layer)
        ++ ", in=" ++ show (bpIn layer)
        ++ ", out=" ++ show (bpOut layer)
        ++ ", w=" ++ show (bpW layer)
        ++ ", activationFunction=?, activationFunction'=?"

backpropagateFinalLayer ::
    PropagatedLayer -> ColumnVector Double -> BackpropagatedLayer
backpropagateFinalLayer l t = BackpropagatedLayer
    {
      bpDazzle = dazzle,
      bpErrGrad = errorGrad dazzle f'a (pIn l),
      bpF'a = pF'a l,
      bpIn = pIn l,
      bpOut = pOut l,
      bpW = pW l,
      bpAS = pAS l
    }
    where dazzle =  pOut l - t
          f'a = pF'a l

errorGrad :: ColumnVector Double -> ColumnVector Double -> ColumnVector Double
    -> ColumnVector Double
errorGrad dazzle f'a input = (dazzle * f'a) <> trans input

-- | Propagate the inputs backward through this layer to produce an output.
backpropagate :: PropagatedLayer -> BackpropagatedLayer -> BackpropagatedLayer
backpropagate layerJ layerK = BackpropagatedLayer
    {
      bpDazzle = dazzleJ,
      bpErrGrad = errorGrad dazzleJ f'aJ (pIn layerJ),
      bpF'a = pF'a layerJ,
      bpIn = pIn layerJ,
      bpOut = pOut layerJ,
      bpW = pW layerJ,
      bpAS = pAS layerJ
    }
    where dazzleJ = wKT <> (dazzleK * f'aK)
          dazzleK = bpDazzle layerK
          wKT = trans ( bpW layerK )
          f'aK = bpF'a layerK
          f'aJ = pF'a layerJ

-- }}}

-- {{{ Adjusting weights after backpropagation

update :: Double -> BackpropagatedLayer -> Layer
update rate layer = Layer
        {
          lW = wNew,
          lAS = bpAS layer
        }
    where wOld = bpW layer
          delW = rate `scale` bpErrGrad layer
          wNew = wOld - delW

-- }}}

-- {{{ Building a network

data BackpropNet = BackpropNet
    {
      layers :: [Layer],
      learningRate :: Double
    } deriving Show

buildBackpropNet ::
  -- The learning rate
  Double ->
  -- The weights for each layer
  [Matrix Double] ->
  -- The activation specification (used for all layers)
  ActivationSpec ->
  -- The network
  BackpropNet
buildBackpropNet lr ws s = BackpropNet { layers=ls, learningRate=lr }
  where checkedWeights = scanl1 checkDimensions ws
        ls = map buildLayer checkedWeights
        buildLayer w = Layer { lW=w, lAS=s }

checkDimensions :: Matrix Double -> Matrix Double -> Matrix Double
checkDimensions w1 w2 =
  if rows w1 == cols w2
       then w2
       else error "Inconsistent dimensions in weight matrix"

propagateNet :: ColumnVector Double -> BackpropNet -> [PropagatedLayer]
propagateNet input net = tail calcs
  where calcs = scanl propagate layer0 (layers net)
        layer0 = PropagatedSensorLayer{ pOut=validatedInputs }
        validatedInputs = validateInput net input

validateInput :: BackpropNet -> ColumnVector Double -> ColumnVector Double
validateInput net = validateInputValues . validateInputDimensions net

validateInputDimensions ::
    BackpropNet ->
    ColumnVector Double ->
    ColumnVector Double
validateInputDimensions net input =
  if got == expected
       then input
       else error ("Input pattern has " ++ show got ++ " bits, but " ++ show expected ++ " were expected")
           where got = rows input
                 expected = inputWidth (head (layers net))

validateInputValues :: ColumnVector Double -> ColumnVector Double
validateInputValues input =
  if (min >= 0) && (max <= 1)
       then input
       else error "Input bits outside of range [0,1]"
       where min = minimum ns
             max = maximum ns
             ns = toList ( flatten input )

backpropagateNet ::
  ColumnVector Double -> [PropagatedLayer] -> [BackpropagatedLayer]
backpropagateNet target layers = scanr backpropagate layerL hiddenLayers
  where hiddenLayers = init layers
        layerL = backpropagateFinalLayer (last layers) target

-- }}}

-- {{{ Define BackpropNet to be an instance of Neural Net

instance NeuralNet BackpropNet where
  evaluate = evaluateBPN
  train = trainBPN

evaluateBPN :: BackpropNet -> [Double] -> [Double]
evaluateBPN net input = columnVectorToList( pOut ( last calcs ))
  where calcs = propagateNet x net
        x = listToColumnVector (1:input)

trainBPN :: BackpropNet -> [Double] -> [Double] -> BackpropNet
trainBPN net input target = BackpropNet { layers=newLayers, learningRate=rate }
  where newLayers = map (update rate) backpropagatedLayers
        rate = learningRate net
        backpropagatedLayers = backpropagateNet (listToColumnVector target) propagatedLayers
        propagatedLayers = propagateNet x net
        x = listToColumnVector (1:input)

-- }}}

-- {{{ General Testing

-- | A layer with suitable input and target vectors, suitable for testing.
data LayerTestData = LTD (ColumnVector Double) Layer (ColumnVector Double)
  deriving Show

-- | Generate a layer with suitable input and target vectors, of the specified "size",
-- | with arbitrary values.
-- | To see sample values (of "size" 4), in GHCi type: sample (sizedLayerTestData 4)
sizedLayerTestData :: Int -> Gen LayerTestData
sizedLayerTestData n = do
    l <- sizedArbLayer n
    x <- sizedArbColumnVector (inputWidth l)
    t <- sizedArbColumnVector (outputWidth l)
    return (LTD x l t)

instance Arbitrary LayerTestData where
  -- | To see sample values, in GHCi type: sample arbLayerTestData
  arbitrary = sized sizedLayerTestData

-- | Training reduces error in the final (output) layer
prop_trainingReducesFinalLayerError :: LayerTestData -> Property
prop_trainingReducesFinalLayerError (LTD x l t) =
    -- (collect l) . -- uncomment to view test data
    (classifyRange "len x " n 0 25) .
    (classifyRange "len x " n 26 50) .
    (classifyRange "len x " n 51 75) .
    (classifyRange "len x " n 76 100) $
    errorAfter < errorBefore || errorAfter < 0.01
        where n = inputWidth l
              pl0 = PropagatedSensorLayer{ pOut=x }
              pl = propagate pl0 l
              bpl = backpropagateFinalLayer pl t
              errorBefore = P.magnitude (t - pOut pl)
              lNew = update 0.0000000001 bpl -- make sure we don't overshoot the mark
              plNew = propagate pl0 lNew
              errorAfter =  P.magnitude (t - pOut plNew)

iterateTraining n pl0 r ltd = pOut plFinal
    where iterations = iterate (trainOneLayer pl0 r) ltd
          (LTD _ lFinal _) = last (take n iterations)
          plFinal = propagate pl0 lFinal

trainOneLayer :: PropagatedLayer -> Double -> LayerTestData -> LayerTestData
trainOneLayer pl0 r (LTD x l t) = LTD x lNew t
    where pl = propagate pl0 l
          bpl = backpropagateFinalLayer pl t
          lNew = update r bpl

-- | Testable property:
-- | Training a single layer with the same input repeatedly will eventually yield the target
prop_trainingOneLayerWithOneInputYieldsPerfection :: LayerTestData -> Property
prop_trainingOneLayerWithOneInputYieldsPerfection (LTD x l t) =
    -- (collect l) . -- uncomment to view test data
    (classifyRange "len x " n 0 25) .
    (classifyRange "len x " n 26 50) .
    (classifyRange "len x " n 51 75) .
    (classifyRange "len x " n 76 100) $
    e < 0.1
        where n = inputWidth l
              r = 1
              pl0 = PropagatedSensorLayer{ pOut=x }
              y = iterateTraining 100 pl0 r (LTD x l t)
              e = P.magnitude (t - y)

-- | A layer with suitable input and target vectors, suitable for testing.
data TwoLayerTestData =
  TLTD (ColumnVector Double) Layer Layer (ColumnVector Double)
    deriving Show

-- | Generate a layer with suitable input and target vectors, of the specified "size",
-- | with arbitrary values.
-- | To see sample values (of "size" 4), in GHCi type: sample (sizedTwoLayerTestData 4)
sizedTwoLayerTestData :: Int -> Gen TwoLayerTestData
sizedTwoLayerTestData n = do
    l1 <- sizedArbLayer n
    l2 <- sizedArbLayer' (outputWidth l1) n
    x <- sizedArbColumnVector (inputWidth l1)
    t <- sizedArbColumnVector (outputWidth l2)
    return (TLTD x l1 l2 t)

instance Arbitrary TwoLayerTestData where
  -- | To see sample values, in GHCi type: sample arbTwoLayerTestData
  arbitrary = sized sizedTwoLayerTestData

-- | Training reduces error in a hidden layer
prop_trainingReducesHiddenLayerError :: TwoLayerTestData -> Property
prop_trainingReducesHiddenLayerError (TLTD x l1 l2 t)=
    -- (collect l) . -- uncomment to view test data
    (classifyRange "len x " n 0 25) .
    (classifyRange "len x " n 26 50) .
    (classifyRange "len x " n 51 75) .
    (classifyRange "len x " n 76 100) $
    errorAfter < errorBefore || errorAfter < 0.01
        where n = inputWidth l1
              pl0 = PropagatedSensorLayer{ pOut=x }
              pl1 = propagate pl0 l1
              pl2 = propagate pl1 l2
              bpl2 = backpropagateFinalLayer pl2 t
              bpl1 = backpropagate pl1 bpl2
              errorBefore = P.magnitude (t - pOut pl2)
              l1New = update 0.00000001 bpl1 -- make sure we don't overshoot the mark
              -- leave layer 2 alone, we're only interested in layer 1
              pl1New = propagate pl0 l1New
              pl2New = propagate pl1New l2
              errorAfter = P.magnitude (t - pOut pl2New)

classifyRange :: Testable a => String -> Int -> Int -> Int -> a -> Property
classifyRange s n n0 n1 =
    classify (n >= n0 && n <= n1) (s ++ show n0 ++ ".." ++ show n1)

-- }}}

-- {{{ Common activation functions

data ActivationSpec = ActivationSpec
    {
      asF :: Double -> Double,
      asF' :: Double -> Double,
      desc :: String
    }

instance Show ActivationSpec where
    show = desc

identityAS = ActivationSpec
    {
      asF = id,
      asF' = const 1,
      desc = "identity"
    }

logisticSigmoidAS :: Double -> ActivationSpec
logisticSigmoidAS c = ActivationSpec
    {
        asF = logisticSigmoid c,
        asF' = logisticSigmoid' c,
        desc = "logistic sigmoid, c=" ++ show c
    }

arbitraryLogisticAS :: Gen ActivationSpec
arbitraryLogisticAS = do
    c <- choose(0,1) -- TODO Can it be > 1?
    return (logisticSigmoidAS c)

instance Arbitrary ActivationSpec where
--    arbitrary = oneof [ return identityAS, arbitraryLogisticAS ]
    arbitrary = return identityAS

logisticSigmoid :: (Field a, Floating a) => a -> a -> a
logisticSigmoid c a = 1 / (1 + exp((-c) * a))

logisticSigmoid' :: (Field a, Floating a) => a -> a -> a
logisticSigmoid' c a = (c * f a) * (1 - f a)
  where f = logisticSigmoid c

tanhAS :: ActivationSpec
tanhAS = ActivationSpec
    {
      asF = tanh,
      asF' = tanh',
      desc = "tanh"
    }

tanh' x = 1 - (tanh x)^2

-- }}}


{-# LANGUAGE DeriveDataTypeable, DeriveGeneric, FlexibleContexts,
    FlexibleInstances, UndecidableInstances #-}
-- |
-- Module    : Statistics.Distribution.Transform
-- Copyright : (c) 2013 John McDonnell;
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Transformations over distributions
module Statistics.Distribution.Transform
    (
      LinearTransform (..)
    , linTransFixedPoint
    , scaleAround
    ) where

import Data.Aeson (FromJSON, ToJSON)
import Control.Applicative ((<*>))
import Data.Binary (Binary)
import Data.Binary (put, get)
import Data.Data (Data, Typeable)
import Data.Functor ((<$>))
import GHC.Generics (Generic)
import qualified Statistics.Distribution as D

-- | Linear transformation applied to distribution.
--
-- > LinearTransform μ σ _
-- > x' = μ + σ·x
data LinearTransform d = LinearTransform
  { linTransLocation :: {-# UNPACK #-} !Double
    -- ^ Location parameter.
  , linTransScale    :: {-# UNPACK #-} !Double
    -- ^ Scale parameter.
  , linTransDistr    :: d
    -- ^ Distribution being transformed.
  } deriving (Eq, Show, Read, Typeable, Data, Generic)

instance (FromJSON d) => FromJSON (LinearTransform d)
instance (ToJSON d) => ToJSON (LinearTransform d)

instance (Binary d) => Binary (LinearTransform d) where
    get = LinearTransform <$> get <*> get <*> get
    put (LinearTransform x y z) = put x >> put y >> put z

-- | Apply linear transformation to distribution.
scaleAround :: Double           -- ^ Fixed point
            -> Double           -- ^ Scale parameter
            -> d                -- ^ Distribution
            -> LinearTransform d
scaleAround x0 sc = LinearTransform (x0 * (1 - sc)) sc

-- | Get fixed point of linear transformation
linTransFixedPoint :: LinearTransform d -> Double
linTransFixedPoint (LinearTransform loc sc _) = loc / (1 - sc)

instance Functor LinearTransform where
  fmap f (LinearTransform loc sc dist) = LinearTransform loc sc (f dist)

instance D.Distribution d => D.Distribution (LinearTransform d) where
  cumulative (LinearTransform loc sc dist) x = D.cumulative dist $ (x-loc) / sc

instance D.ContDistr d => D.ContDistr (LinearTransform d) where
  density    (LinearTransform loc sc dist) x = D.density    dist ((x-loc) / sc) / sc
  logDensity (LinearTransform loc sc dist) x = D.logDensity dist ((x-loc) / sc) - log sc
  quantile      (LinearTransform loc sc dist) p = loc + sc * D.quantile      dist p
  complQuantile (LinearTransform loc sc dist) p = loc + sc * D.complQuantile dist p

instance D.MaybeMean d => D.MaybeMean (LinearTransform d) where
  maybeMean (LinearTransform loc _ dist) = (+loc) <$> D.maybeMean dist

instance (D.Mean d) => D.Mean (LinearTransform d) where
  mean (LinearTransform loc _ dist) = loc + D.mean dist

instance D.MaybeVariance  d => D.MaybeVariance (LinearTransform d) where
  maybeVariance (LinearTransform _ sc dist) = (*(sc*sc)) <$> D.maybeVariance dist
  maybeStdDev   (LinearTransform _ sc dist) = (*sc)      <$> D.maybeStdDev dist

instance (D.Variance d) => D.Variance (LinearTransform d) where
  variance (LinearTransform _ sc dist) = sc * sc * D.variance dist
  stdDev   (LinearTransform _ sc dist) = sc * D.stdDev dist

instance (D.MaybeEntropy d) => D.MaybeEntropy (LinearTransform d) where
  maybeEntropy (LinearTransform _ _ dist) = D.maybeEntropy dist

instance (D.Entropy d) => D.Entropy (LinearTransform d) where
  entropy (LinearTransform _ _ dist) = D.entropy dist

instance D.ContGen d => D.ContGen (LinearTransform d) where
  genContVar (LinearTransform loc sc d) g = do
    x <- D.genContVar d g
    return $! loc + sc * x


{-# LANGUAGE FlexibleContexts #-}

module WeightDecay
  ( input
  , testdata
  , traindata
  , testY
  , testX
  , trainLinRegUnkownTarget
  , trainAndTestWithData
  , trainAndTestWithRegularization
  , makeTransformedMatrix
  , listToTuple'
  ) where

import           Data.Maybe
import           LinearRegression
import           Network.HTTP
import           NonlinearTransform
import           Numeric.LinearAlgebra.Data
import           Numeric.LinearAlgebra.HMatrix as HMatrix

--------------------------------------------------------------
-- Solutions to homework 6 of "Learning from data" question 5-6
--------------------------------------------------------------
input http = do
  response <- simpleHTTP $ getRequest http
  let body = fmap rspBody response
  let rows = fmap (map (map readDouble . words) . lines) body
  return rows

readDouble :: String -> Double
readDouble = read

listToTuple' :: [a] -> Maybe (a, a)
listToTuple' (x:y:xs) = Just (x, y)
listToTuple' _        = Nothing

-- | downloads the test data
testdata = input "http://work.caltech.edu/data/out.dta"

-- | downloads the training data
traindata = input "http://work.caltech.edu/data/in.dta"

testY = fmap (fmap makeVectorY) testdata
  where
    makeVectorY lists = vector $ map last lists

testX = makeDataMatrix testdata

trainY = fmap (fmap makeVectorY) traindata
  where
    makeVectorY lists = vector $ map last lists

trainX = makeDataMatrix traindata

makeLabelVector inputdata = fmap (fmap makeVectorY) inputdata
  where
    makeVectorY lists = vector $ map last lists

makeDataMatrix inputdata = fmap (fmap makeMatrixX) inputdata
  where
    makeMatrixX lists = matrix 3 $ concatMap xvectors lists
    xvectors list = 1 : init list

makeTransformedMatrix2 inputdata = fmap (fmap makeMatrixX) inputdata
  where
    makeMatrixX lists =
      createTransformedX2 $ map (fromMaybe (0, 0) . listToTuple') lists

makeTransformedMatrix k inputdata = fmap (fmap makeMatrixX) inputdata
  where
    makeMatrixX lists =
      createTransformedXk k $ map (fromMaybe (0, 0) . listToTuple') lists

trainXTransformed2 = makeTransformedMatrix2 traindata

testXTransformed2 = makeTransformedMatrix2 testdata

-- | Trains linear Regression from the data on the course website
trainLinRegUnkownTarget = do
  matrixXTrain <- trainXTransformed2
  vectorYTrain <- trainY
  let weights = linearRegressionWeight <$> matrixXTrain <*> vectorYTrain
  let inSampleError =
        linRegClassificationError <$> matrixXTrain <*> vectorYTrain <*> weights
  return (weights, inSampleError)

-- | Tests linear Regression from the data on the course website
testLinRegUnkownTarget weights = do
  matrixXTest <- testXTransformed2
  vectorYTest <- testY
  let outOfSampleError =
        linRegClassificationError <$> matrixXTest <*> vectorYTest <*> weights
  return outOfSampleError

-- | solution for question 2: train and test with given data, with nonlinear transformation but without regularization and print the result
trainAndTestWithData :: IO ()
trainAndTestWithData = do
  (weights, inSampleError) <- trainLinRegUnkownTarget
  outOfSampleError <- testLinRegUnkownTarget weights
  putStr "In Sample Error: "
  print inSampleError
  putStr "Out of Sample Error: "
  print outOfSampleError

-- | returns the weights determined by linear regression on the training data with regularization
linearRegressionWeightReg :: R -> Matrix R -> Vector R -> Vector R
linearRegressionWeightReg regParam matrixOfInputData vectorOfLabels =
  inv
    (tr matrixOfInputData HMatrix.<> matrixOfInputData +
     scalar regParam * (ident $fst $size $tr matrixOfInputData)) HMatrix.#>
  (tr matrixOfInputData #> vectorOfLabels)

-- | Trains linear Regression from the data on the course website with regularization
trainLinRegWithRegularization regParam = do
  matrixXTrain <- trainXTransformed2
  vectorYTrain <- trainY
  let weights =
        linearRegressionWeightReg regParam <$> matrixXTrain <*> vectorYTrain
  let inSampleError =
        linRegClassificationError <$> matrixXTrain <*> vectorYTrain <*> weights
  return (weights, inSampleError)

-- | performs training and testing with the regularization parameter provided as an input and prints the in-sample and out-of-sample error
trainAndTestWithRegularization :: R -> IO ()
trainAndTestWithRegularization regParam = do
  (weights, inSampleError) <- trainLinRegWithRegularization regParam
  outOfSampleError <- testLinRegUnkownTarget weights
  putStr "In Sample Error: "
  print inSampleError
  putStr "Out of Sample Error: "
  print outOfSampleError


{-# OPTIONS_HADDOCK prune #-}
module ForSyDe.Atom.Skel.FastVector.DSP where

import Data.Complex
import qualified Data.Number.FixedFunctions as F
import ForSyDe.Atom.MoC as MoC
import ForSyDe.Atom.Skel.FastVector.Lib as V hiding (duals, unduals)
import ForSyDe.Atom.Skel.FastVector.Matrix as M
import ForSyDe.Atom ((><))

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.taylor'.
taylor :: (Eq a, Floating a)
       => Int      -- ^ number of points in the output window.
       -> Int      -- ^ Number of nearly constant level sidelobes adjacent to the mainlobe
       -> a        -- ^ Desired peak sidelobe level in decibels (db) relative to the mainlobe
       -> Vector a -- ^ The window, with the center value normalized to one (the value
                   -- one appears only if the number of samples is odd).
taylor n nbar level = V.farm11 (*scale) w
  where
    -- explicit conversions to floating
    bN   = fromIntegral n
    nBar = fromIntegral nbar
    ma   = map fromIntegral [1 .. nbar-1]
    -- calculate intermediate values
    b  = 10**((-level) / 20)
    a  = log(b + sqrt(b**2 - 1)) / pi
    s2 = nBar ** 2 / (a**2 + (nBar - 0.5)**2)
    -- functions for calculating coefficients
    fmcalc m  = let numer = (-1)**(m+1) * product[1.0 - m**2/s2/(a**2 + (j - 0.5)**2) | j <- ma]
                    denom = 2 * product[1 - m**2/j**2 | j <- ma, j /= m]
                in numer / denom
    ccalc m x = cos(2 * pi * x * (m - bN/2 + 1/2) / bN)
    wcalc m   = 2 * dotvv (vector $ map fmcalc ma) (vector $ map (ccalc m) ma) + 1
    -- calculate window coefficients
    w  = vector $ map (wcalc . fromIntegral) [0..n-1]
    -- normalize (Note that this is not described in the original text [1])
    scale = 1 / wcalc (bN - 1) / 2

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.taylor''.
taylor' n = taylor 4 (-30)

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.dotvm'.
dotvm :: Num a => Vector a -> Vector (Vector a) -> Vector a
dotvm = dotvm' (+) (*)

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.dotvv'.
dotvv :: Num a => Vector a -> Vector a -> a
dotvv a b
  | V.length a == V.length b = V.reduce (+) (V.farm21 (*) a b)
  | otherwise                = error "Vector sizes must match"

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.dotvv''.
dotvv' :: Num a
     => ((a -> a -> a) -> f a -> f a -> f a)
     -- ^ higher-order function that can wrap the (*) operation.
     -> Vector (f a) -> Vector (f a) -> f a
dotvv' wrap a b
  | V.length a == V.length b = V.reduce (wrap (+)) $ V.farm21 (wrap (*)) a b
  | otherwise                = error "Vector sizes must match"

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.dotmv''.
dotmv' :: (a -> a -> a) -- ^ kernel function for a row/column reduction, e.g. @(+)@ for dot product
       -> (b -> a -> a) -- ^ binary operation for pair-wise elements, e.g. @(*)@ for dot product
       -> Vector (Vector b) -- ^ /size/ = @(xa,ya)@
       -> Vector a      -- ^ /length/ = @xa@
       -> Vector a      -- ^ /length/ = @xa@
dotmv' f g mA y = V.farm11 (\x -> V.reduce f $ V.farm21 g x y) mA

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.dotvm''.
dotvm' :: (b -> b -> b)     -- ^ kernel function for a row/column reduction, e.g. @(+)@ for dot product
       -> (a -> b -> b)     -- ^ binary operation for pair-wise elements, e.g. @(*)@ for dot product
       -> Vector a          -- ^ /length/ = @xa@
       -> Vector (Vector b) -- ^ /size/ = @(xa,ya)@
       -> Vector b          -- ^ /length/ = @ya@
dotvm' f g vs = V.reduce (V.farm21 f) . V.farm21 (\x -> V.farm11 (g x)) vs

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.hanning'.
hanning :: (Floating n)
        => Int -- ^ The length of the window
        -> Vector n
hanning size = V.farm11 func $ V.vector [0..size-1]
  where
    func idx = let i = fromIntegral idx
                   n = fromIntegral size
               in 0.5 * (1 - cos((2 * pi * i) / (n - 1)))

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.hamming'.
hamming :: (Floating n)
        => Int -- ^ The length of the window
        -> Vector n
hamming size = V.farm11 func $ V.vector [0..size-1]
  where
    func idx = let i = fromIntegral idx
                   n = fromIntegral size
               in 0.54 - 0.46 * cos((2 * pi * i) / (n - 1))

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.blackman'.
blackman :: (Floating n)
        => Int -- ^ The length of the window
        -> Vector n
blackman size = V.farm11 func $ V.vector [0..size-1]
  where
    func idx = let i = fromIntegral idx
                   n = fromIntegral size
               in 0.42 - 0.5 * cos((2 * pi * i) / (n - 1)) + 0.08 * cos((4 * pi * i) / (n - 1))

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.fir'.
fir :: Num a
    => Vector a  -- ^ vector of coefficients
    -> Vector a  -- ^ input vector of numbers; /size/ = @n@
    -> Vector a  -- ^ output vector of numbers; /size/ = @n@
fir coefs = V.reverse . V.farm11 applyFilter . tails . V.reverse
  where
    applyFilter = V.reduce (+) . V.farm21 (*) coefs

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.fir''.
fir' :: (a -> a -> a)  -- ^ process/operation replacing '+'
     -> (c -> a -> a)  -- ^ process/operation replacing '*'
     -> (a -> a)       -- ^ delay process
     -> Vector c       -- ^ vector of coefficients
     -> a              -- ^ input signal/structure
     -> a              -- ^ output signal/structure
fir' plus times delay coefs =
  -- V.reduce plus . V.farm21 (\c -> times c) coefs . V.recur (V.fanoutn n delay <: id)
  V.reduce plus . V.farm21 times coefs . V.recuri (V.fanoutn n delay)
  where n = V.length coefs - 1

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.twiddles'.
twiddles :: Floating a => Int -> V.Vector (Complex a)
twiddles bN = (bitrev . V.take (bN `div` 2)) (V.farm11 bW $ vector [0..])
  where bW x = (cis . negate) (-2 * pi * fromIntegral x / fromIntegral bN)

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.fft'.
fft :: RealFloat a
     => Int -> V.Vector (Complex a) -> V.Vector (Complex a)
fft k vs | n == 2^k = V.reverse $ bitrev $ (stage `V.pipe1` V.iterate k (*2) 2) vs
         | otherwise = error $ "n=" ++ show n ++ "; k=" ++ show k
  where
    stage   w = V.concat . V.farm21 segment (twiddles n) . V.group w
    segment t = (><) unduals . (><) (V.farm22 (butterfly t)) . duals
    n         = V.length vs        -- length of input
    -------------------------------------------------
    butterfly w x0 x1 = let t = w * x1 in (x0 + t, x0 - t) -- kernel function

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.fft''.
fft' :: Floating a
     => (Complex a -> a -> a -> (a, a))
     -> Int -> V.Vector a -> V.Vector a
fft' butterfly k vs | n == 2^k = bitrev $ (stage `V.pipe1` (V.iterate k (*2) 2)) vs
  where
    stage   w = V.concat . V.farm21 segment (twiddles n) . V.group w
    segment t = (><) unduals . (><) (V.farm22 (butterfly t)) . duals
    n         = V.length vs        -- length of input

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.duals'.
duals    :: Vector a -> (Vector a, Vector a)
duals v  = (V.take k v, V.drop k v)
  where
    k = V.length v `div` 2

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.unduals'.
unduals  :: Vector a -> Vector a -> Vector a
unduals x y =  x <++> y

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.bitrev'.
bitrev = V.unsafeLift bitrevF
  where
    bitrevF [x] = [x]
    bitrevF xs  = bitrevF (V.evensF xs) ++ bitrevF (V.oddsF xs)


module Quant.Models (
    CharFunc(..)
) where

import Data.Complex
import Quant.Math.Integration
import Quant.ContingentClaim(OptionType(..))

{- | The 'CharFunc' class defines those
models which have closed-form characteristic
functions.

Minimal complete definition: 'charFunc'.

Still under construction.
-}
class CharFunc a where
    -- | Creates a characteristic function for a model.
    charFunc :: CharFunc a => a -> Double -> Complex Double -> Complex Double

    charFuncOption :: (CharFunc a) => a -> OptionType -> Integrator -> Double
        -> Double -> Double -> Double
    charFuncOption model opt intF strike tmat damp = case opt of
        Put  -> s * q * p1 - k * disc * p2
        Call -> k * disc * (1-p2) - s * q * (1-p1)
            where
                s =
                p1 = 0.5 + intF func1 1E-8 100 / pi
                p2 = 0.5 + intF func2 1E-8 100 / pi

                f v' = realPart $ exp (i*v*k) * leftTerm * rightTerm
                    where
                        v = v' :+ 0
                        damp' = damp :+ 0
                        k = log strike :+ 0
                        i = 0 :+ 1
                        leftTerm = d / (damp' + i * v) / (damp'+i*v+(1:+0))
                        rightTerm = cf $ v - i * (damp' + 1)
                        d = disc yc tmat :+ 0
                        cf x = charFuncMart model fg tmat x


{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedLists #-}

module Data.Network
  (Network, Layer, Batch, Data, Layers(Output, (:=>)),
   makeRandomLayer, makeRandomNetwork, run, compute, runLayers, runNetwork, fromList, toList
  )
  --   makeRandomNetwork, makeRandomModel, makeModel, fromList,
  --   runLayer, runLayer', runNetwork, runLayers, runLayers',
  --   runLayerBatch, runLayerBatch', runNetworkBatch, runLayersBatch, runLayersBatch',
  --   Input, Output, Data, Batch, BatchIn, BatchOut, Sample, BatchSample, Weights)
where

import GHC.Exts hiding ((<#))
import Prelude hiding (scanl)
import Data.Functor
-- import Data.Word8
import Data.Foldable hiding (toList)
import Data.Monoid
import System.Random
import Control.Monad.Random as R hiding (fromList)
import Control.Monad
import Control.DeepSeq
import Numeric.LinearAlgebra as LA hiding ((<>), fromList, toList)
import Numeric.LinearAlgebra.Data as LAD hiding (fromList, toList, size, R)
-- import Numeric.LinearAlgebra.Static as LAS hiding ((<>), fromList)
import Numeric.LinearAlgebra.Devel
import Data.Network.Utils (scanl)
import Data.Network.Activation

-- type ActivationFunction = forall a. Floating a => a -> a

type Input = Vector Double
type Output = Vector Double
type Data = Vector Double
type Weights = Matrix Double
type Batch = Matrix Double
type BatchIn = Matrix Double
type BatchOut = Matrix Double
type Sample = (Data, Data) -- Labeled data for supervised learning
type BatchSample = (Batch, Batch)
-- type Model a = (Network, ActivationFunction a) -- A network's structure with its activation function

-- Network layer = (weight matrix, bias vector)
-- Parameter a: type of data(e.g. Double)

-- What's a layer?
class (NFData a) => Layer i a | a -> i where
  compute :: a -> i -> i
  run     :: a -> i -> (i, i)

  makeRandomLayer :: (RandomGen g) => g -> ([Double] -> [Double]) -> Int -> Int -> ActivationFunction i -> a

-- What's a network?
class (NFData a) => Network i a | a -> i where
  runNetwork :: a -> i -> (i,i)
  runLayers  :: a -> i -> [(i,i)]
  makeRandomNetwork :: (RandomGen g) =>
    g -> ([Double] -> [Double]) -> (a -> a) -> Int -> [(Int, ActivationFunction i)] -> (Int, ActivationFunction i) -> a

-- Linked list Structure for feedforward neural network
data Layers :: * -> * where
  Output :: !a -> Layers a
  (:=>) :: !a -> !(Layers a) -> Layers a deriving (Eq)

infixr 5 :=>

instance (NFData a) => NFData (Layers a) where
  rnf (Output x) = rnf x
  rnf (l :=> ll) = rnf l `seq` rnf ll

instance Show a => Show (Layers a) where
  show (Output l) = "(Output) " ++ show l
  show (a :=> b) = show a ++ " -> " ++ show b

instance Functor Layers where
  fmap f (Output x) = (Output $ f x)
  fmap f (a :=> b) = (f a) :=> (fmap f b)

instance Foldable Layers where
  foldr f b (Output x) = f x b
  foldr f c (a :=> b) = f a (foldr f c b)

  foldl f a (Output x) = f a x
  foldl f a (b :=> c) = foldl f (f a b) c

instance Num a => Num (Layers a) where
  (Output l) + (Output r) = Output (l + r)
  (l :=> n) + (r :=> n') = (l + r) :=> (n + n')
  (Output l) - (Output r) = Output (l - r)
  (l :=> n) - (r :=> n') = (l - r) :=> (n - n')
  (Output l) * (Output r) = Output (l * r)
  (l :=> n) * (r :=> n') = (l * r) :=> (n * n')
  abs (Output l) = Output (abs l)
  abs (l :=> n) = (abs l :=> abs n)
  signum (Output l) = Output (signum l)
  signum (l :=> n) = (signum l) :=> (signum n)
  fromInteger a = Output (fromInteger a)

instance Fractional a => Fractional (Layers a) where
  (Output l) / (Output r) = Output (l/r)
  (l :=> n) / (r :=> n') = (l/r) :=> (n/n')
  recip (Output l) = Output (recip l)
  recip (l :=> n) = (recip l) :=> (recip n)
  fromRational a = Output (fromRational a)

    -- g -> ([Double] -> [Double]) -> Int -> [(Int, ActivationFunction i)] -> (Int, ActivationFunction i) -> a

instance IsList (Layers i) where
  type Item (Layers i) = i
  fromList []     = undefined
  fromList [x]    = Output x
  fromList (x:xs) = x :=> (fromList xs)

  toList (Output x) = [x]
  toList (x :=> y)  = x:(toList y)

-- fromList :: [a] -> Layers a
-- fromList [] = undefined
-- fromList [x] = Output x
-- fromList (x:xs) = x :=> (fromList xs)

-- make a fully connected layer; m = size of input(previous layer), n = size of output(next layer)
-- makeLayerFC :: Floating a => Int -> Int -> ActivationFunction a -> Layer a Double
-- makeLayerFC m n af = Layer { weights = (m><n) (repeat 0), biases = n |> (repeat 0), activation = af }

-- make a fully connected layer with randomly generated (normally distributed) weights and biases
-- makeLayerFCR :: MonadRandom m => RandDist -> Int -> Int -> ActivationFunction a -> m (Layer a Double)
-- makeLayerFCR dist m n = do
--   (seeds :: [Int]) <- getRandoms
--   (seed2 :: Int) <- getRandom
--   let w = fromRows $ fmap (\s -> randomVector s dist n) (take m seeds)
--       b = randomVector seed2 dist n
--   return $ Layer { weights = w, biases = b }

-- makeRandomNetwork :: MonadRandom m => RandDist -> Int -> [Int] -> Int -> m Network
-- makeRandomNetwork dist i [] o = Output <$> (makeLayerFCR dist i o)
-- makeRandomNetwork dist i (h:hs) o = (:=>) <$> makeLayerFCR dist i h <*> makeRandomNetwork dist h hs o

-- makeModel :: Network -> [ActivationFunction a] -> Model a
-- makeModel (Output l) (a:afs) = Output (l,a)
-- makeModel (l :=> n) (a:afs) = ((l,a) :=> makeModel n afs)

-- makeRandomModel :: MonadRandom m => RandDist -> Int -> [(Int, ActivationFunction a)] -> (Int, ActivationFunction a) -> m (Network i)
-- makeRandomModel dist i [] (o, oaf) = (Output . (,oaf)) <$> (makeLayerFCR dist i o)
-- makeRandomModel dist i ((h,af):hs) (o, oaf) = ((:=>) . (,af)) <$> makeLayerFCR dist i h <*> makeRandomModel dist h hs (o, oaf)

-- compute layer's net result (before activation/squashing)
-- computeLayer :: Layer a Double -> a -> Vector Double
-- computeLayer (Layer weights biases) input = biases + input <# weights

-- net linear combiner for a batch of input as a matrix, with rows as input vectors
-- computeLayerBatch :: Layer Batch Double -> Batch -> Batch
-- computeLayerBatch (Layer weights biases) inputs = (fromRows $ replicate (rows inputs) biases) + (inputs <> weights)

-- run layer with input vector
-- runLayer :: ActivationFunction Output -> Layer Double -> Input -> Output
-- runLayer af l i = activate af $ computeLayer l i

-- run layer with output vector, outputing both result and derivative
-- runLayer' :: ActivationFunction Output -> Layer Double -> Input -> (Output, Output)
-- runLayer' af l i = unAct af $ computeLayer l i

-- Batch variants

-- runLayerBatch :: ActivationFunction Batch -> Layer Double -> BatchIn -> BatchOut
-- runLayerBatch af l i = activate af $ computeLayerBatch l i

-- runLayerBatch' :: ActivationFunction Batch -> Layer Double -> BatchIn -> (BatchOut, BatchOut)
-- runLayerBatch' af l i = unAct af $ computeLayerBatch l i

-- run network on input vector with an activation function
-- runNetwork :: Model Data -> Input -> Output
-- runNetwork n i = foldl' ((flip . uncurry . flip) runLayer) i n

-- runNetworkBatch :: Model Batch -> BatchIn -> BatchOut
-- runNetworkBatch n i = foldl' ((flip . uncurry . flip) runLayerBatch) i n

-- run network on input, accumulating each layer's intermediary result,
-- with the first element of the resulting list being the last layer's output, ...
-- runLayers :: Model Data -> Input -> [Output]
-- runLayers n i = foldl' (\os@(o:os') (l,af) -> (runLayer af l o):os) [i] n

-- Same as above but with derivatives
-- runLayers' :: Model Data -> Input -> [(Output, Output)]
-- runLayers' n i = foldl' (\os@((o, o'):os') (l, af) -> (runLayer' af l o):os) [(i, mempty)] n

-- Batch variants
-- runLayersBatch :: Model Batch -> BatchIn -> [BatchOut]
-- runLayersBatch n i = foldl' (\os@(o:os') (l, af) -> (runLayerBatch af l o):os) [i] n

-- runLayersBatch' :: Model Batch -> BatchIn -> [(BatchOut, BatchOut)]
-- runLayersBatch' n i = foldl' (\os@((o, o'):os') (l, af) -> (runLayerBatch' af l o):os) [(i, mempty)] n


{-# language DataKinds, ScopedTypeVariables #-}
module Main where

import System.Environment
import GHC.TypeLits
import Data.Proxy
import Data.List
import Data.List.Split (chunksOf)
import Data.Maybe
import Numeric.LinearAlgebra (toList)

import GAPoly

type Args = (String, Int, Int, Int, Integer, Double)

validateArgs :: [String] -> Args
validateArgs args =
  case args of
       (fname:nGens:nPop:nTerms:maxK:pm:_) -> (fname, read nGens, read nPop, read nTerms, read maxK, read pm)
       _ -> error "missing arguments"

term2str :: [(Bool,Int)] -> String
term2str ts = intercalate "*" $ filter (/="") $ zipWith toStr ts [0..]
  where
    toStr (False,_) _  = ""
    toStr (True, k) ix = "x[:," ++ show ix ++ "]**(" ++ show k ++ ")"

poly2str :: Int -> Solution (Poly n) -> String
poly2str nVars (Sol cs betas _) = let cs'    = map _getPoly cs
                                      terms  = map term2str $ chunksOf nVars cs'
                                      (b:bs) = reverse $ toList $ fromJust betas
                                      bts    = intercalate " + " $ zipWith toStr bs terms
                                   in show b ++ " + " ++ bts
  where
    toStr beta term = show beta ++ "*" ++ term

main :: IO ()
main = do args <- getArgs
          let (dataname, nGens, nPop, nTerms, maxK, pm) = validateArgs args
          dat <- readFile dataname
          let (xss, ys, nVars) = parseFile dat
          case someNatVal maxK of
               Just (SomeNat (_ :: Proxy n)) ->
                 do (avgs, best) <- runGA nTerms nVars nGens nPop pm xss ys
                    print $ poly2str nVars (best :: Solution (Poly n))
               Nothing -> error "Negative maxK"


module Statistics.Information.Mixed.MutualInfo where

import Data.List
import Data.Matrix
import qualified Data.Vector as Vector
import qualified Statistics.Information.Continuous.Entropy as CE
import Statistics.Information.Utils.List
import Statistics.Information.Utils.Random
import System.Random

micd :: (RandomGen g, Eq a) => g -> Int -> Int -> Matrix Double -> Matrix a ->
        Double
micd g k base xs ys | nrows xs == nrows ys =
  abs $ hx - sum [ycontrib ypg | ypg <- yvals]
  where
    (g1, g2) = split g
    hx = CE.entropy g1 k base xs
    ys' = toLists ys
    yvals = probs ys' `zip` splitN (length (nub ys')) g2
    ycontrib ((y, p), g') =
      let xgiveny = [getRow i xs | i <- [1..nrows xs],
                     (Vector.toList $ getRow i ys) == y] in
        if k <= length xgiveny - 1 then
          p * CE.entropy g' k base (fromLists (map Vector.toList xgiveny))
        else
          p * hx

midc :: (RandomGen g, Eq a) => g -> Int -> Int -> Matrix a -> Matrix Double ->
        Double
midc g k base xs ys = micd g k base ys xs


module Carlson
  (carlsonRF, carlsonRF',
  carlsonRD, carlsonRD',
  carlsonRJ, carlsonRJ',
  carlsonRC, carlsonRC',
  carlsonRG, carlsonRG')
  where
import           Data.Complex
import           Internal

rf_ :: Cplx -> Cplx -> Cplx -> Double -> ((Double,Double,Double), Cplx)
rf_ x y z err =
  let a = (x+y+z)/3 in
  let delta = map (\u -> magnitude(1-u/a)) [x,y,z] in
  if maximum delta < err
    then ((delta !! 0, delta !! 1, delta !! 2), a)
    else
      let (sqrtx, sqrty, sqrtz) = (sqrt x, sqrt y, sqrt z) in
      let lambda = sqrtx*sqrty + sqrty*sqrtz + sqrtz*sqrtx in
      rf_ ((x+lambda)/4) ((y+lambda)/4) ((z+lambda)/4) err

carlsonRF' :: Double -> Cplx -> Cplx -> Cplx -> Cplx
carlsonRF' err x y z =
  if zeros > 1
    then error "At most one of x, y, z can be 0"
    else
      let ((dx,dy,dz), a) = rf_ x y z err in
      let (e2,e3) = (dx*dy + dy*dz + dz*dx, dx*dy*dz) in
      toCplx(1 - e2/10 + e3/14 + e2*e2/24 - 3*e2*e3/44 - 5*e2*e2*e2/208 +
          3*e3*e3/104 + e2*e2*e3/16) / sqrt a
  where
    zeros = sum (map (\u -> fromEnum (u == 0)) [x,y,z])

carlsonRF :: Cplx -> Cplx -> Cplx -> Cplx
carlsonRF = carlsonRF' 1e-15

rd_ :: Cplx -> Cplx -> Cplx -> Cplx -> Cplx -> Double ->
       ((Double,Double,Double), Cplx, Cplx, Cplx)
rd_ x y z s fac err =
  let a = (x+y+z+z+z)/5 in
  let delta = map (\u -> magnitude(1-u/a)) [x,y,z] in
  if maximum delta < err
    then ((delta !! 0, delta !! 1, delta !! 2), a, s, fac)
    else
      let (sqrtx, sqrty, sqrtz) = (sqrt x, sqrt y, sqrt z) in
      let lambda = sqrtx*sqrty + sqrty*sqrtz + sqrtz*sqrtx in
      let s' = s + fac / (sqrt z * (z + lambda)) in
      rd_ ((x+lambda)/4) ((y+lambda)/4) ((z+lambda)/4) s' (fac/4) err

carlsonRD' :: Double -> Cplx -> Cplx -> Cplx -> Cplx
carlsonRD' err x y z =
  if zeros > 1
    then error "At most one of x, y, z can be 0"
    else
      let ((dx,dy,dz), a, s, fac) = rd_ x y z 0 1 err in
      let
        (e2,e3,e4,e5) = (dx*dy + dy*dz + 3*dz*dz + 2*dz*dx + dx*dz + 2*dy*dz,
          dz*dz*dz + dx*dz*dz + 3*dx*dy*dz + 2*dy*dz*dz + dy*dz*dz + 2*dx*dz*dz,
          dy*dz*dz*dz + dx*dz*dz*dz + dx*dy*dz*dz + 2*dx*dy*dz*dz,
          dx*dy*dz*dz*dz) in
      3*s + fac * toCplx(1 - 3*e2/14 + e3/6 + 9*e2*e2/88 - 3*e4/22 - 9*e2*e3/52 +
        3*e5/26 - e2*e2*e2/16 + 3*e3*e3/40 + 3*e2*e4/20 + 45*e2*e2*e3/272 -
        9*(e3*e4 + e2*e5)/68) / a / sqrt a
  where
    zeros = sum (map (\u -> fromEnum (u == 0)) [x,y,z])

carlsonRD :: Cplx -> Cplx -> Cplx -> Cplx
carlsonRD = carlsonRD' 1e-15

rj_ :: Cplx -> Cplx -> Cplx -> Cplx -> Cplx -> Double -> Cplx -> Int ->
       Double -> [Cplx] -> [Cplx] -> Double -> (Cplx, Int, [Cplx], [Cplx])
rj_ x y z p a maxmagns delta f fac d e err =
  let q = (4/err)**(1/6) * maxmagns / fromIntegral f in
  if magnitude a > q
    then (a, f, d, e)
    else
      let dnew = (sqrt p + sqrt x)*(sqrt p + sqrt y)*(sqrt p + sqrt z)
          d' = (fromIntegral f * dnew) : d
          e' = (toCplx fac * delta / dnew / dnew) : e
          lambda = sqrt x * sqrt y + sqrt y * sqrt z + sqrt z * sqrt x
          x' = (x + lambda) / 4
          y' = (y + lambda) / 4
          z' = (z + lambda) / 4
          p' = (p + lambda) / 4
          a' = (a + lambda) / 4
      in
      rj_ x' y' z' p' a' maxmagns delta (4*f) (fac/64) d' e' err

carlsonRJ' :: Double -> Cplx -> Cplx -> Cplx -> Cplx -> Cplx
carlsonRJ' err x y z p =
  if zeros > 1
    then error "At most one of x, y, z, p can be 0"
    else
      let a0 = (x + y + z + p + p) / 5
          maxmagns = maximum $ map (\u -> magnitude(a0-u)) [x, y, z, p]
          delta = (p-x)*(p-y)*(p-z)
      in
      let (a, f, d, e) = rj_ x y z p a0 maxmagns delta 1 1 [] [] err
          f' = fromIntegral f
      in
      let x' = (a0 - x) / f' / a
          y' = (a0 - y) / f' / a
          z' = (a0 - z) / f' / a
          p' = -(x'+y'+z') / 2
          e2 = x'*y' + x'*z' + y'*z' - 3*p'*p'
          e3 = x'*y'*z' + 2*e2*p' + 4*p'*p'*p'
          e4 = p'*(2*x'*y'*z' + e2*p' + 3*p'*p'*p')
          e5 = x'*y'*z'*p'*p'
          h = zipWith (\u v -> atanx_over_x(sqrt u) / v) e d
      in
      (1 - 3*e2/14 + e3/6 + 9*e2*e2/88 - 3*e4/22 - 9*e2*e3/52 + 3*e5/26) /
        f' / a / sqrt a + 6 * sum h
  where
    zeros = sum (map (\u -> fromEnum (u == 0)) [x,y,z,p])
    atanx_over_x w = if w == 0 then 1 else atanC w / w

carlsonRJ :: Cplx -> Cplx -> Cplx -> Cplx -> Cplx
carlsonRJ = carlsonRJ' 1e-15

rc_ :: Cplx -> Cplx -> Cplx -> Double -> Int -> Double -> (Cplx, Int)
rc_ x y a magn f err =
  let q = (1/3/err)**(1/8) * magn / fromIntegral f in
  if magnitude a > q
    then (a, f)
    else
      let lambda = 2 * sqrt x * sqrt y + y
          a' = (a + lambda) / 4
          x' = (x + lambda) / 4
          y' = (y + lambda) / 4
      in
      rc_ x' y' a' magn (4*f) err

carlsonRC' :: Double -> Cplx -> Cplx -> Cplx
carlsonRC' err x y =
  if y == 0
    then error "y cannot be 0"
    else
      let a0 = (x + y + y) / 3
          magn = magnitude(a0-x)
      in
      let (a, f) = rc_ x y a0 magn 1 err
          f' = fromIntegral f
      in
      let s = (y - a0) / f' / a in
      (1 + 3*s*s/10 + s*s*s/7 + 3*s*s*s*s/8 + 9*s*s*s*s*s/22 +
        159*s*s*s*s*s*s/208 + 9*s*s*s*s*s*s*s/8) / sqrt a

carlsonRC :: Cplx -> Cplx -> Cplx
carlsonRC = carlsonRC' 1e-15

carlsonRG' :: Double -> Cplx -> Cplx -> Cplx -> Cplx
carlsonRG' err x y z =
  if zeros > 1
    then sqrt(x+y+z) / 2
    else
      if z == 0
        then carlsonRG' err z x y
        else
          (z * carlsonRF' err x y z -
            (x-z) * (y-z) * carlsonRD' err x y z / 3 +
            sqrt x * sqrt y / sqrt z) / 2
  where
    zeros = sum (map (\u -> fromEnum (u == 0)) [x,y,z])

carlsonRG :: Cplx -> Cplx -> Cplx -> Cplx
carlsonRG = carlsonRG' 1e-15


{-# LANGUAGE OverloadedStrings, DuplicateRecordFields, RecordWildCards, NamedFieldPuns #-}
{-|
 - Module      : Numeric.GLP
 - Description : Specifiy a General Linear Problem and transform to a Standardized LP
 - Copyright   : (c) Torsten Kemps-Benedix, 2019
 - This module 'Numeric.GLP' specifiies a [Generalized Linear Problem](https://en.wikipedia.org/wiki/Linear_programming) (or LP in short) and provides means to tranform this into a Standardized LP (see module 'Numeric.LP') that may be solved with e.g. the simplex solver in module 'Numeric.LP.Simplex'.
-}
module Numeric.GLP
-- (GLP, Numeric.GLP.a, c, m1, m2, m3, dir, mkGLP, standardizeGLP)
where

import qualified Data.List                   as L
import           Control.Monad
import qualified Data.Set                    as S
import           Data.Set                    (member, isSubsetOf)
import           Numeric.LinearAlgebra       (I, Element, Matrix, Vector, ident, diagl,
                                              cols, rows, find, toColumns,
                                              fromColumns, cmap, fromList, toList,
                                              (!), (><),
                                              )
import qualified Numeric.LinearAlgebra       as LA
import           Numeric.LinearAlgebra.Devel (modifyMatrix, runSTMatrix,
                                              thawMatrix, mapVectorWithIndex, fi)
import           Numeric.LinearAlgebra       (vjoin, (|||), (===))

-- | This record defines a General Linear Problem containing ≤
-- and ≥ inequalities and == equalities. The sum of the numbers of
-- inequalities ('m1' times ≤ and 'm2' times ≥) and equalities ('m3'
-- times ==) must match one less than the number of
-- rows of the matrix 'a'. The first (i.e. the zero-th) row of the matrix
-- contains the coefficients of the objective function.
--
-- We assume that \(a_{00}\) is zero. The matrix \(a\) shall have \(m+1\) rows
-- and \(n+1\) columns.
--
-- The formal specification of the problem we intend to solve is as follows:
--
-- Maximize resp. minimize (according to 'dir') the following linear objective
-- function:
--
-- \[ z_0 = \sum_{i=1}^{n} a_{0,i} x_i \]
--
-- Subject to the following linear constraints:
--
-- \[ z_j = a_{j,0} + \sum_{i=1}^{n} a_{j,i} x_i,\, j=1,\ldots,m\]
-- \[ z_j \leq 0,\, j=1,\ldots,m_1\]
-- \[ z_j \geq 0,\, j=m_1+1,\ldots,m_1+m_2\]
-- \[ z_j = 0,\, j=m_1+m_2+1,\ldots,m=m1+m_2+m_3\]
--
-- The original variables \(x_1, \ldots, x_n\) may be constraint to be non-negative,
-- non-positive or they may be unconstrained according to the n-vector
-- \(c\in\{-1,0,+1\}^n\).
--
-- \[ c_j = -1 \Rightarrow x_j \leq 0\]
-- \[ c_j = 1 \Rightarrow x_j \geq 0\]
-- \[ c_j = 0 \Rightarrow x_j \text{ unrestricted} \]
data GLP = GLP {
    a   :: Matrix Double, -- ^ Tableau of coefficients, see above for a formal
                          -- description of their meaning
    c   :: Vector I,      -- ^ Type of constraint for \(x_j\), \(j=1,\ldots,n\),
                          -- see 'GLP' for further details. The element \(c_i\)
                          -- corresponds to column \(i+1\) of the matrix \(a\).
    n :: Int,             -- ^ Number of original variables \(x_1, \ldots, x_n\)
    m1  :: Int,           -- ^ Number of ≤ inequalities
    m2  :: Int,           -- ^ Number of ≥ inequalities
    m3  :: Int,           -- ^ Number of = equalities
    dir :: Dir,           -- ^ Whether the objective function is to be maximized
                          -- or minimzed
    lhs, rhs :: Vector I  -- ^ Indices of left and right hand variables respectively.
                          -- The original variables \(x_1,\ldots,x_n\) have indices
                          -- \(1,\ldots,n\), the slack variables get indices
                          -- \(n+1, \ldots, n+m_1+m_2+1\), and the auxiliary variables
                          -- finally get indices \(n+m_1+m_2+2, \ldots, n+m_1+m_2+m+2\).
    } deriving (Show)

-- | This smart constructor checks all the pre-conditions for 'GLP' stated in the
-- description of 'GLP' and initializes the variables 'lhs' and 'rhs'. Initially only
-- the \(m\) auxiliary variables \(z_1,\ldots,z_m\) are on the left hand side and all the
-- original variables \(x_1,\ldots,x_n\) are on the right hand side. Initially there are
-- no slack variables \(y_1,\ldots,y_{m_1+m_2+1}\). These will be introdced by the function
-- 'introduceSlackVariables' below.
mkGLP :: Matrix Double -> Vector I -> (Int, Int, Int) -> Dir -> Either String GLP
mkGLP a c (m1,m2,m3) dir =
    let m = rows a - 1
        n = cols a - 1
    in if (m /= m1+m2+m3)
       then Left $ "rows a - 1 == "++show m++" /= m1+m2+m3 == "++show (m1+m2+m3)
       else if (m <= 0)
            then Left "The number of constraints is m <= 0"
            else if (n <= 0)
                 then Left "The number of variables n is <= 0"
                 else if not (S.fromList (LA.toList c) `isSubsetOf` (S.fromList [-1,0,1]))
                      then Left "Vector c has elements other than -1, 0, +1"
                      else if LA.size c /= n
                           then Left $ "Vector c has not the proper length "++show n
                           else Right (GLP a c n m1 m2 m3 dir
                                           (LA.fromList [fi (n+m1+m2+2)..fi (n+m1+m2+m+2)])
                                           (LA.fromList [1..fi n]))

data Dir = Maximize | Minimize deriving Show

-- | Transform the Generalized Linear Problem to a Standardized Linear Problem
-- by means of the following steps:
--
-- 1. If the objective function is to be minimized multiply it by -1.
-- 2. Turn around the ≥-inequalities by negation of the corresponding rows. After this
-- step there are 'm1'+'m2' inequalities of the <= type.
-- 3. For those variables \(x_i\) that have a constraint \(x_i \leq 0\) substitute
-- \(x_i \rightarrow -x_i\). This means multiplying the column \(i\) of the matrix
-- \(a\) by \(-1\).
-- 4. For those variables \(x_i\) that are unconstrained introduce new non-negative
-- variables \(x_i = x_i⁺ - x_i⁻\). We give these the indices \(x_i⁺ = x_i\) and
-- \(x_⁻ = x_{i+1}\). The indices of all other variables \(x_k\), \(k>i\) have to be
-- increased by \(1\). The matrix \(a\) therefore grows by one column.
-- 5. For the 'm1'+'m2' ≤-inequalities introduce as many slack variables on the right
-- hand side and make these equalities instead. After this step there are only
-- equalities left as constraints. We do not represent these slack variables in the
-- matrix 'a' and therefore nothing really changes in this step.
-- 6. Introduce artificial variables \(z_j\) ( \(j=1,\ldots,m\) ) for each equation.
-- This changes the optimization problem and the solver has to take this into account.
--
-- The function returns a modified and standardized GLP and a matrix \(b\) that maps
-- solution vectors (including the objective functions value) of the modified GLP back
-- to solution vectors of the original GLP.
standardizeGLP :: GLP -> (GLP, Matrix Double)
standardizeGLP glp =
    let (glp1, b1) = resolveMin2Max glp
        (glp2, b2) = resolveGeqConstraints glp1
        (glp3, b3) = resolveNonPositiveVars glp2
        (glp4, b4) = resolveUnconstrainedVars glp3
        glp5 = introduceSlackVariables glp4
    in (glp5, b1 <> b2 <> b3 <> b4)

resolveMin2Max :: GLP -> (GLP, Matrix Double)
resolveMin2Max glp@GLP{..} = case dir of
        Maximize -> (glp, ident (1+m1+m2+m3))
        Minimize -> (glp{a = mapRow a 0 negate, dir = Maximize},
                     diagl (-1:(replicate (m1+m2+m3) 1)))

resolveGeqConstraints :: GLP -> (GLP, Matrix Double)
resolveGeqConstraints glp@GLP{..} =
    let a' = foldl (\a j -> mapRow a j negate) a [m1+1..m1+m2] -- negate rows from m1+1 to m1+m2
    in (glp{a = a', m1 = m1+m2, m2 = 0},
        ident (1+m1+m2+m3))

resolveNonPositiveVars :: GLP -> (GLP, Matrix Double)
resolveNonPositiveVars glp@GLP{..} =
    let varsToBeNegated = find (== -1) c
        varsToBeNegatedS = S.fromList varsToBeNegated
        (a', c') = foldl (\(a,c) i ->
                            (mapCol a (i+1) negate,
                             mapVectorWithIndex (\k x -> if k==i
                                                         then 1
                                                         else c!k)
                                                c))
                         (a, c)
                         (find (== -1) c)
    in (glp{a = a', c = c'},
        diagl $ 1:[if i `member` varsToBeNegatedS
                   then -1
                   else 1 | i <- [0..m1+m2+m3-1]])

resolveUnconstrainedVars :: GLP -> (GLP, Matrix Double)
resolveUnconstrainedVars glp@GLP{..} =
    let m = 1+m1+m2+m3
        zCol = vec1 m 0 1
        (a', c', b) =
            let listOfColsAndCs =
                   fmap (\(c1,ci,i1) ->
                            if ci==0
                            then ([c1, cmap negate c1],
                                  [1, 1],
                                  [vec1 m (i1+1) 1,
                                   vec1 m (i1+1) (-1)]) -- duplicate and adjust
                            else ([c1], [ci],
                                  [vec1 m (i1+1) 1])) -- return unchanged
                        (L.zip3 (toColumns a) (toList c) [0..(cols a)-1])
                columns = L.concat $ fmap (\(x,_,_) -> x) listOfColsAndCs
                cs = L.concat $ fmap (\(_,x,_) -> x) listOfColsAndCs
                bs = L.concat $ fmap (\(_,_,x) -> x) listOfColsAndCs
            in (fromColumns columns, fromList cs, fromColumns $ zCol:bs)
    in (glp{a=a', c=c'}, b)

introduceSlackVariables :: GLP -> GLP
introduceSlackVariables glp@GLP{..} =
    let m3' = m1+m2+m3
        n = cols a - 1
        a' = a ||| ((1><(m1+m2)) (replicate (m1+m2) 0) ===
                    diagl ((replicate m1 1)++(replicate m2 (-1))) ===
                    (m3><(m1+m2)) (replicate (m3*(m1+m2)) 0))
        rhs' = vjoin [rhs, LA.fromList [fi (n+1)..fi (n+1+m1+m2)]]
    in glp{a=a', m1=0, m2=0, m3=m3', rhs=rhs'}

vec1 :: Int -> Int -> Double -> Vector Double
vec1 n i x = fromList $ L.replicate i 0++x:L.replicate (n-i-1) 0

mapRow :: (Element t) => Matrix t -> Int -> (t -> t) -> Matrix t
mapRow a k f = runSTMatrix $ do
                    a' <- thawMatrix a
                    forM_ [0..cols a-1] (\i -> modifyMatrix a' k i f)
                    return a'

mapCol :: (Element t) => Matrix t -> Int -> (t -> t) -> Matrix t
mapCol a k f = runSTMatrix $ do
                    a' <- thawMatrix a
                    forM_ [0..rows a-1] (\i -> modifyMatrix a' i k f)
                    return a'

prettyPrintAsFormulas :: GLP -> String
prettyPrintAsFormulas glp@GLP{..} =
    let rs = LA.toRows a
        m = cols a - 1
        n = rows a - 1
        objFun = (show dir)++
                    " z_0 = "++linExp 0
        zRows = fmap (\k -> "z_"++show k++" = "++linExp k) [1..n]
        linExp k = L.intercalate " + " (
                    L.filter (/= "") $
                    fmap (\(a,i) ->
                            if a /= 0
                            then show a++(if i>0 then " "++varName i else "")
                            else "") (L.zip (toList $ rs!!k) [0..m]))
        varName i = if i <= n then "x_"++show i  else "y_"++show (i-n)
        zConstr  = L.intercalate ", " $ filter (/= "") $
                   [if m1 == 1
                    then "z_1 ≤ 0"
                    else if m1 == 2
                         then "z_1, z_2 ≤ 0"
                         else if m1 > 2
                              then "z_1, ..., z_"++show m1++" ≤ 0"
                              else "",
                    if m2 == 1
                    then "z_"++show (m1+1)++" ≥ 0"
                    else if m2 > 1
                         then "z_"++show (m1+1)++", ..., z_"++show (m1+m2)++" ≥ 0"
                         else "",
                    if m3 == 1
                    then "z_"++show (m1+m2+1)++" = 0"
                    else if m3 > 1
                         then "z_"++show (m1+m2+1)++", ..., z_"++show (m1+m2+m3)++" = 0"
                         else ""]
        xConstr ck k = "x_"++show k++if ck == -1
                                     then " ≤ 0"
                                     else if ck == 1
                                          then " ≥ 0"
                                          else " unconstrained"
        xConstrs = fmap (uncurry xConstr) (L.zip (toList c) [1..m])
    in L.intercalate "\n" (
            objFun:
            "Subject to":
            (fmap (\line -> "\t"++line) (zRows++[""]++
                                         [zConstr, ""]++
                                         xConstrs)))
            ++"\n"

(Right ex1') = ex1
ex1 = let a = (5><5)
                 [   0,    1,  1,  3, -0.5
                 , 740,   -1,  0, -2,  0
                 ,   0,    0, -2,  0,  7
                 ,   0.5,  0, -1,  1, -2
                 ,   9,   -1, -1, -1, -1]
          c = LA.fromList [-1, 1, 0, -1]
      in mkGLP a c (2, 1, 1) Minimize

allEx1 = do
    let g1 = ex1'
    putStr $ prettyPrintAsFormulas g1
    putStrLn $ "lhs = "++show (lhs g1)
    putStrLn $ "rhs = "++show (rhs g1)
    putStrLn "\nResolve min/max"
    let (g2, b1) = resolveMin2Max g1
    putStr $ prettyPrintAsFormulas g2
    putStrLn $ "lhs = "++show (lhs g2)
    putStrLn $ "rhs = "++show (rhs g2)
    print b1
    putStrLn "\nResove geq constraints"
    let (g3, b2) = resolveGeqConstraints g2
    putStr $ prettyPrintAsFormulas g3
    putStrLn $ "lhs = "++show (lhs g3)
    putStrLn $ "rhs = "++show (rhs g3)
    print b2
    putStrLn "\nResove non-positive variables"
    let (g4, b3) = resolveNonPositiveVars g3
    putStr $ prettyPrintAsFormulas g4
    putStrLn $ "lhs = "++show (lhs g4)
    putStrLn $ "rhs = "++show (rhs g4)
    print b3
    putStrLn "\nResove unconstrained variables"
    let (g5, b4) = resolveUnconstrainedVars g4
    putStr $ prettyPrintAsFormulas g5
    putStrLn $ "lhs = "++show (lhs g5)
    putStrLn $ "rhs = "++show (rhs g5)
    print b4
    putStrLn "\nIntroduce slack variables"
    let g6 = introduceSlackVariables g5
    putStr $ prettyPrintAsFormulas g6
    putStrLn $ "lhs = "++show (lhs g6)
    putStrLn $ "rhs = "++show (rhs g6)
    putStrLn "\n"
    return (g6, b1 <> b2 <> b3 <> b4)


module NeuralNet.Problem (
  Problem (),
  RunStep (),
  createProblem,
  runProblem,
  runStepIteration,
  runStepAccuracy,
  runStepCost,
  problemTestSet
) where

import NeuralNet.Example
import NeuralNet.Net
import NeuralNet.Layer
import NeuralNet.Cost
import NeuralNet.Train
import Numeric.LinearAlgebra

type NumIterations = Int
type LearningRate = Double
data Problem = Problem NeuralNetDefinition ExampleSet ExampleSet LearningRate NumIterations
  deriving (Show, Eq)

type IterationNum = Int
type Cost = Double
type Accuracy = Double
data RunStep = RunStep IterationNum Cost Accuracy
  deriving (Show, Eq)

runStepIteration :: RunStep -> IterationNum
runStepIteration (RunStep i _ _) = i

runStepCost :: RunStep -> Cost
runStepCost (RunStep _ c _) = c

runStepAccuracy :: RunStep -> Accuracy
runStepAccuracy (RunStep _ _ a) = a

problemNNDef :: Problem -> NeuralNetDefinition
problemNNDef (Problem d _ _ _ _) = d

problemTrainSet :: Problem -> ExampleSet
problemTrainSet (Problem _ t _ _ _) = t

problemTestSet :: Problem -> ExampleSet
problemTestSet (Problem _ _ t _ _) = t

problemLearningRate :: Problem -> LearningRate
problemLearningRate (Problem _ _ _ l _) = l

problemNumIterations :: Problem -> NumIterations
problemNumIterations (Problem _ _ _ _ i) = i

createProblem :: NeuralNetDefinition -> ExampleSet -> ExampleSet -> LearningRate -> NumIterations -> Problem
createProblem def trainSet testSet learningRate numIterations
  | not (isExampleSetCompatibleWithNNDef trainSet def) = error "trainSet not compatible with nn"
  | not (isExampleSetCompatibleWithNNDef testSet def)  = error "testSet not compatible with nn"
  | numIterations <= 0                                 = error "Must provide positive numIterations"
  | learningRate <= 0                                  = error "Must provide positive learningRate"
  | otherwise                                          = Problem def trainSet testSet learningRate numIterations

runProblem :: WeightsStream -> Problem -> (Double -> Double -> Bool) -> (NeuralNet, [RunStep], Double)
runProblem g p accuracyCheck = (resultNN, tail allSteps, testAccuracy)
  where
    startNN = initNN g (problemNNDef p)
    startStep = RunStep 0 1 0
    iterations = [1..(problemNumIterations p)]
    allNNAndSteps = reverse (foldl (\steps@((nn,_):_) i -> runProblemStep p i nn accuracyCheck : steps) [(startNN, startStep)] iterations)
    allSteps = map snd allNNAndSteps
    resultNN = fst (last allNNAndSteps)

    testSet = problemTestSet p
    testA = forwardPropA (last (nnForwardSet resultNN testSet))
    testAccuracy = calcAccuracy accuracyCheck testA (exampleSetY testSet)

runProblemStep :: Problem -> IterationNum -> NeuralNet -> (Double -> Double -> Bool) -> (NeuralNet, RunStep)
runProblemStep p i nn accuracyCheck = (newNN, RunStep i cost accuracy)
  where
    trainSet = problemTrainSet p
    forwardSteps = nnForwardSet nn trainSet
    al = forwardPropA (last forwardSteps)
    y = exampleSetY trainSet
    grads = nnBackward nn forwardSteps trainSet
    cost = computeCost al y
    newNN = updateNNParams nn grads (problemLearningRate p)
    accuracy = calcAccuracy accuracyCheck al y

calcAccuracy :: (Double -> Double -> Bool) -> Matrix Double -> Matrix Double -> Double
calcAccuracy accuracyCheck yHat y = fromIntegral numCorrect / fromIntegral (cols y)
  where
    numCorrect = length (filter (uncurry accuracyCheck) (zip (concat (toLists yHat)) (concat (toLists y))))
    -- TODO: Search for concat (toLists  - should be an easier way


{-# LANGUAGE PatternGuards #-}
-- |
-- Module    : Statistics.Matrix
-- Copyright : 2011 Aleksey Khudyakov, 2014 Bryan O'Sullivan
-- License   : BSD3
--
-- Basic matrix operations.
--
-- There isn't a widely used matrix package for Haskell yet, so
-- we implement the necessary minimum here.

module Statistics.Matrix
    ( -- * Data types
      Matrix(..)
    , Vector
      -- * Conversion from/to lists/vectors
    , fromVector
    , fromList
    , fromRowLists
    , fromRows
    , fromColumns
    , toVector
    , toList
    , toRows
    , toColumns
    , toRowLists
      -- * Other
    , generate
    , generateSym
    , ident
    , diag
    , dimension
    , center
    , multiply
    , multiplyV
    , transpose
    , power
    , norm
    , column
    , row
    , map
    , for
    , unsafeIndex
    , hasNaN
    , bounds
    , unsafeBounds
    ) where

import Prelude hiding (exponent, map)
import Control.Applicative ((<$>))
import Control.Monad.ST
import qualified Data.Vector.Unboxed as U
import           Data.Vector.Unboxed   ((!))
import qualified Data.Vector.Unboxed.Mutable as UM

import Numeric.Sum         (sumVector,kbn)
import Statistics.Matrix.Function
import Statistics.Matrix.Types
import Statistics.Matrix.Mutable  (unsafeNew,unsafeWrite,unsafeFreeze)

----------------------------------------------------------------
-- Conversion to/from vectors/lists
----------------------------------------------------------------

-- | Convert from a row-major list.
fromList :: Int                 -- ^ Number of rows.
         -> Int                 -- ^ Number of columns.
         -> [Double]            -- ^ Flat list of values, in row-major order.
         -> Matrix
fromList r c = fromVector r c . U.fromList

-- | create a matrix from a list of lists, as rows
fromRowLists :: [[Double]] -> Matrix
fromRowLists = fromRows . fmap U.fromList

-- | Convert from a row-major vector.
fromVector :: Int               -- ^ Number of rows.
           -> Int               -- ^ Number of columns.
           -> U.Vector Double   -- ^ Flat list of values, in row-major order.
           -> Matrix
fromVector r c v
  | r*c /= len = error "input size mismatch"
  | otherwise  = Matrix r c v
  where len    = U.length v

-- | create a matrix from a list of vectors, as rows
fromRows :: [Vector] -> Matrix
fromRows xs
  | [] <- xs        = error "Statistics.Matrix.fromRows: empty list of rows!"
  | any (/=nCol) ns = error "Statistics.Matrix.fromRows: row sizes do not match"
  | nCol == 0       = error "Statistics.Matrix.fromRows: zero columns in matrix"
  | otherwise       = fromVector nRow nCol (U.concat xs)
  where
    nCol:ns = U.length <$> xs
    nRow    = length xs

-- | create a matrix from a list of vectors, as columns
fromColumns :: [Vector] -> Matrix
fromColumns = transpose . fromRows

-- | Convert to a row-major flat vector.
toVector :: Matrix -> U.Vector Double
toVector (Matrix _ _ v) = v

-- | Convert to a row-major flat list.
toList :: Matrix -> [Double]
toList = U.toList . toVector

-- | Convert to a list of lists, as rows
toRowLists :: Matrix -> [[Double]]
toRowLists (Matrix _ nCol v)
  = chunks $ U.toList v
  where
    chunks [] = []
    chunks xs = case splitAt nCol xs of
      (rowE,rest) -> rowE : chunks rest

-- | Convert to a list of vectors, as rows
toRows :: Matrix -> [Vector]
toRows (Matrix _ nCol v) = chunks v
  where
    chunks xs
      | U.null xs = []
      | otherwise = case U.splitAt nCol xs of
          (rowE,rest) -> rowE : chunks rest

-- | Convert to a list of vectors, as columns
toColumns :: Matrix -> [Vector]
toColumns = toRows . transpose

----------------------------------------------------------------
-- Other
----------------------------------------------------------------

-- | Generate matrix using function
generate :: Int                 -- ^ Number of rows
         -> Int                 -- ^ Number of columns
         -> (Int -> Int -> Double)
            -- ^ Function which takes /row/ and /column/ as argument.
         -> Matrix
generate nRow nCol f
  = Matrix nRow nCol $ U.generate (nRow*nCol) $ \i ->
      let (r,c) = i `quotRem` nCol in f r c

-- | Generate symmetric square matrix using function
generateSym
  :: Int                 -- ^ Number of rows and columns
  -> (Int -> Int -> Double)
     -- ^ Function which takes /row/ and /column/ as argument. It must
     --   be symmetric in arguments: @f i j == f j i@
  -> Matrix
generateSym n f = runST $ do
  m <- unsafeNew n n
  for 0 n $ \r -> do
    unsafeWrite m r r (f r r)
    for (r+1) n $ \c -> do
      let x = f r c
      unsafeWrite m r c x
      unsafeWrite m c r x
  unsafeFreeze m

-- | Create the square identity matrix with given dimensions.
ident :: Int -> Matrix
ident n = diag $ U.replicate n 1.0

-- | Create a square matrix with given diagonal, other entries default to 0
diag :: Vector -> Matrix
diag v
  = Matrix n n $ U.create $ do
      arr <- UM.replicate (n*n) 0
      for 0 n $ \i ->
        UM.unsafeWrite arr (i*n + i) (v ! i)
      return arr
  where
    n = U.length v

-- | Return the dimensions of this matrix, as a (row,column) pair.
dimension :: Matrix -> (Int, Int)
dimension (Matrix r c _) = (r, c)

-- | Matrix-matrix multiplication. Matrices must be of compatible
-- sizes (/note: not checked/).
multiply :: Matrix -> Matrix -> Matrix
multiply m1@(Matrix r1 _ _) m2@(Matrix _ c2 _) =
  Matrix r1 c2 $ U.generate (r1*c2) go
  where
    go t = sumVector kbn $ U.zipWith (*) (row m1 i) (column m2 j)
      where (i,j) = t `quotRem` c2

-- | Matrix-vector multiplication.
multiplyV :: Matrix -> Vector -> Vector
multiplyV m v
  | cols m == c = U.generate (rows m) (sumVector kbn . U.zipWith (*) v . row m)
  | otherwise   = error $ "matrix/vector unconformable " ++ show (cols m,c)
  where c = U.length v

-- | Raise matrix to /n/th power. Power must be positive
-- (/note: not checked).
power :: Matrix -> Int -> Matrix
power mat 1 = mat
power mat n = res
  where
    mat2 = power mat (n `quot` 2)
    pow  = multiply mat2 mat2
    res | odd n     = multiply pow mat
        | otherwise = pow

-- | Element in the center of matrix (not corrected for exponent).
center :: Matrix -> Double
center mat@(Matrix r c _) =
    unsafeBounds U.unsafeIndex mat (r `quot` 2) (c `quot` 2)

-- | Calculate the Euclidean norm of a vector.
norm :: Vector -> Double
norm = sqrt . sumVector kbn . U.map square

-- | Return the given column.
column :: Matrix -> Int -> Vector
column (Matrix r c v) j= U.generate r (\i -> v `U.unsafeIndex` (j + i * c))
{-# INLINE column #-}

-- | Return the given row.
row :: Matrix -> Int -> Vector
row (Matrix _ c v) i = U.slice (c*i) c v

unsafeIndex :: Matrix
            -> Int              -- ^ Row.
            -> Int              -- ^ Column.
            -> Double
unsafeIndex = unsafeBounds U.unsafeIndex

-- | Apply function to every element of matrix
map :: (Double -> Double) -> Matrix -> Matrix
map f (Matrix r c v) = Matrix r c (U.map f v)

-- | Indicate whether any element of the matrix is @NaN@.
hasNaN :: Matrix -> Bool
hasNaN = U.any isNaN . toVector

-- | Given row and column numbers, calculate the offset into the flat
-- row-major vector.
bounds :: (Vector -> Int -> r) -> Matrix -> Int -> Int -> r
bounds k (Matrix rs cs v) r c
  | r < 0 || r >= rs = error "row out of bounds"
  | c < 0 || c >= cs = error "column out of bounds"
  | otherwise        = k v $! r * cs + c
{-# INLINE bounds #-}

-- | Given row and column numbers, calculate the offset into the flat
-- row-major vector, without checking.
unsafeBounds :: (Vector -> Int -> r) -> Matrix -> Int -> Int -> r
unsafeBounds k (Matrix _ cs v) r c = k v $! r * cs + c
{-# INLINE unsafeBounds #-}

transpose :: Matrix -> Matrix
transpose m@(Matrix r0 c0 _) = Matrix c0 r0 . U.generate (r0*c0) $ \i ->
  let (r,c) = i `quotRem` r0
  in unsafeIndex m c r


module ParseNumber where

import Data.Char (digitToInt)
import Data.Complex
import Data.Ratio
import DataTypes (LispVal(..))

import Numeric (readDec, readHex, readInt, readOct)
import ParserCombinators
       (Parser(..), (<|>), char, digit, failure, hexDigit, many1,
        octDigit, oneOf)

readBinary :: ReadS Integer
readBinary = readInt 2 (`elem` "01") digitToInt

parseNumber :: Parser LispVal
parseNumber = parseComplex <|> parseRational <|> parseFloat <|> parseInteger

--------------
-- Integer
--------------
parseInteger :: Parser LispVal
parseInteger =
  parseIntegerDecimal <|> do
    _ <- char '#'
    base <- oneOf "bdox"
    parseIntegerBase base

parseIntegerBase :: Char -> Parser LispVal
parseIntegerBase base =
  case base of
    'd' -> parseIntegerDecimal
    'o' -> parseIntegerOctal
    'x' -> parseIntegerHex
    'b' -> parseIntegerBinary
    _ -> failure "Bad integer format"

parseIntegerHelper :: ReadS Integer -> Parser Char -> Parser LispVal
parseIntegerHelper reader parser =
  (char '-' >> (parseIntegerHelper' negate)) <|>
  (char '+' >> (parseIntegerHelper' id)) <|>
  (parseIntegerHelper' id)
  where
    parseIntegerHelper' :: (Integer -> Integer) -> Parser LispVal
    parseIntegerHelper' op = Integer . op . fst . head . reader <$> many1 parser

parseIntegerDecimal :: Parser LispVal
parseIntegerDecimal = parseIntegerHelper readDec digit

parseIntegerOctal :: Parser LispVal
parseIntegerOctal = parseIntegerHelper readOct octDigit

parseIntegerHex :: Parser LispVal
parseIntegerHex = parseIntegerHelper readHex hexDigit

parseIntegerBinary :: Parser LispVal
parseIntegerBinary = parseIntegerHelper readBinary (oneOf "01")

--------------
-- Float
--------------
parseFloat :: Parser LispVal
parseFloat =
  parseFloatDecimal <|> do
    _ <- char '#'
    base <- oneOf "bdox"
    parseFloatBase base

parseFloatBase :: Char -> Parser LispVal
parseFloatBase base =
  case base of
    'd' -> parseFloatDecimal
    'o' -> parseFloatOctal
    'x' -> parseFloatHex
    'b' -> parseFloatBinary
    _ -> failure "Bad float format"

parseFloatDecimal :: Parser LispVal
parseFloatDecimal = parseFloatHelper 10 digit readDec

parseFloatOctal :: Parser LispVal
parseFloatOctal = parseFloatHelper 8 octDigit readOct

parseFloatBinary :: Parser LispVal
parseFloatBinary = parseFloatHelper 2 (oneOf "01") readBinary

parseFloatHex :: Parser LispVal
parseFloatHex = parseFloatHelper 16 hexDigit readHex

parseFloatHelper :: Int -> Parser Char -> (ReadS Integer) -> Parser LispVal
parseFloatHelper base p reader =
  (char '-' >> (parseFloat' negate)) <|> (char '+' >> parseFloat' id) <|>
  (parseFloat' id)
  where
    helper :: Integer -> Integer -> Double
    helper whole decimal =
      if decimal == 0
        then fromIntegral whole
        else let d = fromIntegral decimal :: Double
                 w = fromIntegral whole :: Double
                 b = fromIntegral base :: Double
                 e = logBase b d
                 f = fromIntegral $ floor e
                 g = d / (b ** (f + 1))
             in w + g
    parseFloat' :: (Double -> Double) -> Parser LispVal
    parseFloat' op = do
      w <- many1 p
      _ <- char '.'
      d <- many1 p
      let whole = fst . head $ reader w
          decimal = fst . head $ reader d
      return $ Float $ (op (helper whole decimal))

--------------
-- Complex
--------------
parseComplex :: Parser LispVal
parseComplex =
  parseComplexDecimal <|> do
    _ <- char '#'
    base <- oneOf "bdox"
    parseComplexBase base

parseComplexBase :: Char -> Parser LispVal
parseComplexBase base =
  case base of
    'd' -> parseComplexDecimal
    'o' -> parseComplexOctal
    'x' -> parseComplexHex
    'b' -> parseComplexBinary
    _ -> failure "Bad complex format"

parseComplexDecimal :: Parser LispVal
parseComplexDecimal =
  parseComplexHelper parseIntegerDecimal parseFloatDecimal parseRationalDecimal

parseComplexOctal :: Parser LispVal
parseComplexOctal =
  parseComplexHelper parseIntegerOctal parseFloatOctal parseRationalOctal

parseComplexHex :: Parser LispVal
parseComplexHex =
  parseComplexHelper parseIntegerHex parseFloatHex parseRationalHex

parseComplexBinary :: Parser LispVal
parseComplexBinary =
  parseComplexHelper parseIntegerBinary parseFloatBinary parseRationalBinary

parseComplexHelper ::
     Parser LispVal -> Parser LispVal -> Parser LispVal -> Parser LispVal
parseComplexHelper pn pf pr = do
  real <- fmap toDouble (pr <|> pf <|> pn)
  imaginary <- fmap toDouble (pr <|> pf <|> pn)
  _ <- char 'i'
  return $ Complex (real :+ imaginary)
  where
    toDouble (Float x) = x
    toDouble (Integer x) = fromIntegral x
    toDouble (Rational x) = fromRational x

--------------
-- Rational
--------------
parseRational :: Parser LispVal
parseRational =
  parseRationalDecimal <|> do
    _ <- char '#'
    base <- oneOf "bdox"
    parseRationalBase base

parseRationalBase :: Char -> Parser LispVal
parseRationalBase base =
  case base of
    'd' -> parseRationalDecimal
    'o' -> parseRationalOctal
    'x' -> parseRationalHex
    'b' -> parseRationalBinary
    _ -> failure "Bad rational format"

parseRationalHelper :: Parser LispVal -> Parser LispVal
parseRationalHelper p = do
  num <- fmap toInt p
  _ <- char '/'
  denom <- fmap toInt p
  return $ Rational (num % denom)
  where
    toInt (Integer x) = x

parseRationalDecimal :: Parser LispVal
parseRationalDecimal = parseRationalHelper parseIntegerDecimal

parseRationalOctal :: Parser LispVal
parseRationalOctal = parseRationalHelper parseIntegerOctal

parseRationalHex :: Parser LispVal
parseRationalHex = parseRationalHelper parseIntegerHex

parseRationalBinary :: Parser LispVal
parseRationalBinary = parseRationalHelper parseIntegerBinary