-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Matrix.Cholesky
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Cholesky factorizations of symmetric (Hermitian) positive-definite
-- matrices.
--

module Numeric.LinearAlgebra.Matrix.Cholesky (
    -- * Immutable interface
    cholFactor,
    cholSolveVector,
    cholSolveMatrix,

    -- * Mutable interface
    cholFactorM,
    cholSolveVectorM_,
    cholSolveMatrixM_,

    ) where

import Control.Monad( when )
import Control.Monad.ST( ST, runST, unsafeIOToST )
import Text.Printf( printf )

import qualified Foreign.LAPACK as LAPACK

import Numeric.LinearAlgebra.Types

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

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

-- | @cholFactor a@ tries to compute the Cholesky
-- factorization of @a@.  If @a@ is positive-definite then the routine
-- returns @Right@ with the factorization.  If the leading minor of order @i@
-- is not positive-definite then the routine returns @Left i@.
cholFactor :: (LAPACK e)
                 => Herm Matrix e
                 -> Either Int (Chol Matrix e)
cholFactor (Herm uplo a) = runST $ do
    ma <- M.newCopy a
    cholFactorM (Herm uplo ma)
        >>= either (return . Left) (\(Chol uplo' ma') -> do
                a' <- M.unsafeFreeze ma'
                return $ Right (Chol uplo' a')
                )

-- | @cholSolveVector a x@ returns @a \\ x@.
cholSolveVector :: (LAPACK e)
                      => Chol Matrix e
                      -> Vector e
                      -> Vector e
cholSolveVector a x = V.create $ do
    x' <- V.newCopy x
    cholSolveVectorM_ a x'
    return x'

-- | @cholSolveMatrix a b@ returns @a \\ b@.
cholSolveMatrix :: (LAPACK e)
                => Chol Matrix e
                -> Matrix e
                -> Matrix e
cholSolveMatrix a c = M.create $ do
    c' <- M.newCopy c
    cholSolveMatrixM_ a c'
    return c'

-- | @cholFactorM a@ tries to compute the Cholesky
-- factorization of @a@ in place.  If @a@ is positive-definite then the
-- routine returns @Right@ with the factorization, stored in the same
-- memory as @a@.  If the leading minor of order @i@ is not
-- positive-definite then the routine returns @Left i@.
-- In either case, the original storage of @a@ is destroyed.
cholFactorM :: (LAPACK e)
            => Herm (STMatrix s) e
            -> ST s (Either Int (Chol (STMatrix s) e))
cholFactorM (Herm uplo a) = do
    (ma,na) <- M.getDim a
    let n = na

    when (not $ (ma,na) == (n,n)) $ error $
        printf ("cholFactorM"
                ++ " (Herm _ <matrix with dim (%d,%d)>): nonsquare matrix"
               ) ma na

    unsafeIOToST $
        M.unsafeWith a $ \pa lda -> do
            info <- LAPACK.potrf uplo n pa lda
            return $ if info > 0 then Left info
                                 else Right (Chol uplo a)

-- | @cholSolveVectorM_ a x@ sets @x := a \\ x@.
cholSolveVectorM_ :: (LAPACK e, RMatrix m)
                  => Chol m e
                  -> STVector s e
                  -> ST s ()
cholSolveVectorM_ a x =
    M.withFromColM x $ \x' ->
        cholSolveMatrixM_ a x'

-- | @cholSolveMatrixM_ a b@ sets @b := a \\ b@.
cholSolveMatrixM_ :: (LAPACK e, RMatrix m)
                  => Chol m e
                  -> STMatrix s e
                  -> ST s ()
cholSolveMatrixM_ (Chol uplo a) b = do
    (ma,na) <- M.getDim a
    (mb,nb) <- M.getDim b
    let (n,nrhs) = (mb,nb)

    when ((not . and) [ (ma,na) == (n,n)
                      , (mb,nb) == (n,nrhs)
                      ]) $ error $
        printf ("cholSolveMatrixM_"
                ++ " (Chol _ <matrix with dim (%d,%d)>)"
                ++ " <matrix with dim (%d,%d)>"
                ++ ": dimension mismatch")
               ma na mb nb

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        M.unsafeWith b $ \pb ldb ->
            LAPACK.potrs uplo n nrhs pa lda pb ldb


-------------------------------------------------------------------------------
-- |
-- Module    :  Supervised.LinearRegression
-- Copyright :  (c) Sam Stites 2017
-- License   :  MIT
-- Maintainer:  sam@stites.io
-- Stability :  experimental
-- Portability: non-portable
--
-- Fits a linear model with coefficients w = (w_1, ..., w_p) to minimize the
-- residual sum of squares between the observed responses in the dataset, and
-- the responses predicted by the linear approximation. Mathematically it solves
-- a problem of the form:
--
--     \underset{w}{min\,} {|| X w - y||_2}^2
--
-- === Ordinary Least Squares Complexity
--
-- This method computes the least squares solution using a singular value
-- decomposition of X. If X is a matrix of size (n, p) this method has a cost of
-- O(n p^2), assuming that n \geq p.
-------------------------------------------------------------------------------
{-# LANGUAGE TypeFamilies #-}
module Supervised.LinearRegression where

import Prelude
import qualified Data.Monoid as M
import Numeric.LinearAlgebra
import qualified Numeric.LinearAlgebra.Data as LD
import Debug.Trace

type Labels  = Vector Double
type Weights = Vector Double
type Inputs  = Matrix Double
type Model   = (Weights, Inputs -> Labels)

data LinearRegression = LinearRegression
  { nIters       :: Integer  -- ^ The number of iterations the algorithm will train weights for.
  , learningRate :: Double   -- ^ Step length for updating weights
  }

-- "scikit-learn style" naming
mkModel :: LinearRegression
mkModel = LinearRegression 100 0.001

fit :: LinearRegression -> Inputs -> Labels -> Model
fit cfg i l = weightedFit cfg w i l
  where
    nfs :: Int
    nfs = cols i

    w :: Weights
    w =  (vector $ replicate nfs 1) / (konst (fromIntegral nfs) nfs)

weightedFit :: LinearRegression -> Weights -> Inputs -> Labels -> Model
weightedFit cfg _ _xs y = (w, \i -> snocBias i #> w)
  where
    xs :: Matrix Double
    xs = snocBias _xs

    u, s, v :: Matrix Double
    (u, s, v) = svdM $ tr xs <> xs

    xSqInv :: Matrix Double
    xSqInv = (v <> pinv s) <> tr u

    w :: Vector Double
    w = (xSqInv <> tr xs) #> y

svdM :: Matrix Double -> (Matrix Double, Matrix Double, Matrix Double)
svdM = (\(a,b,c) -> (a, diag b, c)) . svd

consBias :: Matrix Double -> Matrix Double
consBias m = LD.fromColumns $ ones (snd $ size m) : LD.toColumns m

snocBias :: Matrix Double -> Matrix Double
snocBias m = m ||| konst 1 (fst (size m), 1)

ones :: Int -> Vector Double
ones = konst 1

predict :: Model -> Inputs -> Labels
predict (_, fn) i = fn i

-- | Returns the coefficient of determination R^2 of the prediction.
scoreR2Regressor :: (truth ~ Labels) => Model -> truth -> Inputs -> Double
scoreR2Regressor (_, fn) truth i = rSquare truth (fn i)

-- | Compute R^2, the coefficient of determination that
-- indicates goodness-of-fit of a regression.
rSquare :: (truth ~ Labels) => truth -> Labels -> Double
rSquare truth pred = 1 - numerator / denominator
  where
    weight :: Vector Double
    weight = 1

    numerator :: Double
    numerator = sumElements $ weight * ((truth - pred) ** 2)

    denominator :: Double
    denominator = sumElements $ weight * ((truth - averageTruth) ** 2)

    averageTruth :: Vector Double
    averageTruth = konst ((sumElements truth) / (fromIntegral n)) n
      where
        n :: Int
        n = size truth

sumColumns :: Matrix Double -> Vector Double
sumColumns = vector . fmap sumElements . toColumns

coefficients :: Model -> Weights
coefficients = fst


module Numeric.MixtureModel.Exponential ( -- * General data types
                                          Sample
                                        , SampleIdx
                                        , Samples
                                        , ComponentIdx
                                        , Assignments
                                        , Weight
                                        -- * Exponential parameters
                                        , Rate, Beta
                                        , Exponential(..)
                                        , ComponentParams
                                        , paramFromSamples
                                        , paramsFromAssignments
                                        -- * Exponential distribution
                                        , Prob
                                        , prob
                                        , tauMean, tauVariance
                                        , modelProb
                                        -- * Gibbs sampling
                                        , estimateWeights
                                        , updateAssignments
                                        -- * Score
                                        , scoreAssignments
                                        , maxLikelihoodScore
                                        -- * Classification
                                        , classify
                                        ) where

import           Control.Monad.ST
import           Data.Function (on)
import qualified Data.Vector as VB
import           Data.Vector.Algorithms.Heap
import qualified Data.Vector.Mutable as MV
import qualified Data.Vector.Unboxed as V

import           Numeric.Log hiding (Exp, sum)
import qualified Numeric.Log as Log
import           Numeric.SpecFunctions (logBeta)
import           Statistics.Sample (mean)
import           Numeric.Newton
import           Math.Gamma (gamma)

import           Data.Random hiding (gamma)
import           Data.Random.Distribution.Categorical

type Prob = Log Double
type Sample = Double
type SampleIdx = Int
type ComponentIdx = Int
type Weight = Double

type Rate = Double -- ^ The rate parameter
type Beta = Double -- ^ The stretching parameter
data Exponential = Exp Rate
                 | StretchedExp Rate Beta
                 | FixedExp Rate Beta
                 deriving (Show, Read, Eq)

-- k refers to number of components
-- N refers to number of samples
type Samples = V.Vector Sample                          -- length == N
type Assignments = V.Vector ComponentIdx                -- length == N
type ComponentParams = VB.Vector (Weight, Exponential)  -- length == K

-- | `expProb lambda tau` is the probability of `tau` under Exponential
-- distribution defined by rate `lambda`
prob :: Exponential -> Sample -> Prob
prob _ tau | tau < 0 = error "Exponential distribution undefined for tau<0"
prob (FixedExp lambda 1) tau = prob (Exp lambda) tau
prob (FixedExp lambda beta) tau = prob (StretchedExp lambda beta) tau
prob (StretchedExp lambda 1) tau = prob (Exp lambda) tau
prob (StretchedExp lambda beta) tau =
    Log.Exp $ log beta + (beta-1) * log tau + beta * log lambda - (tau * lambda)**beta
prob (Exp lambda) tau = Log.Exp $ log lambda - lambda * tau

-- | Probability of a sample under a mixture
modelProb :: ComponentParams -> Sample -> Prob
modelProb _ tau | tau < 0 = error "Exponential distribution undefined for tau<0"
modelProb model tau = VB.sum $ VB.map (\(w,e)->prob e tau) model

-- | Mean of the given distribution
tauMean :: Exponential -> Double
tauMean (Exp lambda) = 1 / lambda
tauMean (StretchedExp lambda beta) = gamma (1/beta) / beta / lambda
tauMean (FixedExp lambda beta) = tauMean (StretchedExp lambda beta)

-- | Variance of the given distribution
tauVariance :: Exponential -> Double
tauVariance (Exp lambda) = 1/lambda^2
tauVariance (StretchedExp lambda beta) = 2 * gamma (2/beta) / lambda^2 / beta
tauVariance (FixedExp lambda beta) = tauVariance (StretchedExp lambda beta)

-- | Exponential parameter from samples
paramFromSamples :: Exponential -> V.Vector Sample -> Exponential
paramFromSamples (FixedExp lambda beta) _ = FixedExp lambda beta
paramFromSamples _ v | V.null v = error "Can't estimate parameters without samples"
paramFromSamples (Exp _) v = Exp $ 1 / mean v
paramFromSamples (StretchedExp _ betaOld) v =
    let v' = runST $ do a <- V.thaw v
                        sort a
                        V.freeze a
        n = realToFrac $ V.length v'
        tn = V.head v'
        s beta = V.sum (V.map (\t->t**beta) v') - n*tn**beta
        betaOpt beta = let num = V.sum (V.map (\t->t**beta * log t) v') - n*tn**beta*log tn
                       in num / s beta - V.sum (V.map log v') / n - 1 / beta
        betaOpt' beta = let denom = V.sum (V.map (**beta) v') - n*tn**beta
                            num1 = V.sum (V.map (\t->t**beta * log t) v')
                                   - n*tn**beta * log tn
                            num2 = V.sum (V.map (\t->t**beta * (log t)^2) v')
                                   - n*tn**beta * (log tn)^2
                        in 1/beta^2 - (num1/denom)^2 + num2/denom
        beta' = findRoot 1e-5 betaOpt betaOpt' betaOld
        lambda' = (s beta' / n)**(-1/beta')
    in StretchedExp lambda' beta'

-- | Exponential parameter for component given samples and their
-- component assignments
paramFromAssignments :: Samples -> Assignments -> (ComponentIdx, Exponential) -> Exponential
paramFromAssignments samples assignments (k,p) =
  paramFromSamples p $ V.map snd $ V.filter (\(k',_)->k==k') $ V.zip assignments samples

-- | Exponential parameters for all components given samples and their
-- component assignments
paramsFromAssignments :: Samples -> VB.Vector Exponential -> Assignments -> VB.Vector Exponential
paramsFromAssignments samples params assignments =
  VB.map (paramFromAssignments samples assignments) $ VB.indexed params

-- | Draw a new assignment for a sample given beta parameters
drawAssignment :: ComponentParams -> Sample -> RVar ComponentIdx
drawAssignment params x =
  let probs = map (\(w,p)->realToFrac w * prob p x) $ VB.toList params
  in case filter (isInfinite . ln . fst) $ zip probs [0..] of
        (x:_)     -> return $ snd x
        otherwise -> categorical
                     $ map (\(p,k)->(realToFrac $ p / sum probs :: Double, k))
                     $ zip probs [0..]

-- | `countIndices n v` is the list of counts
countIndices :: Int -> V.Vector Int -> VB.Vector Int
countIndices n v = runST $ do
    accum <- VB.thaw $ VB.replicate n 0
    V.forM_ v $ \k -> do n' <- MV.read accum k
                         MV.write accum k $! n'+1
    VB.freeze accum

-- | Estimate the component weights of a given set of parameters
estimateWeights :: Assignments -> VB.Vector Exponential -> ComponentParams
estimateWeights assignments params =
  let counts = countIndices (VB.length params) assignments
      norm = realToFrac $ V.length assignments
      weights = VB.map (\n->realToFrac n / norm) counts
  in VB.zip weights params

-- | Sample new assignments for observations under given model parameters
updateAssignments :: Samples -> ComponentParams -> RVar Assignments
updateAssignments samples params =
  V.mapM (drawAssignment params) samples

-- | "Likelihood" of sample assignments under given model
-- parameters. Note that the exponential distribution is a density
-- function and as such this will give an unnormalized result unless
-- multiplied by dtau^N
scoreAssignments :: Samples -> ComponentParams -> Assignments -> Prob
scoreAssignments samples params assignments =
    V.product
    $ V.map (\(k,x)->let (w,p) = params VB.! k
                     in realToFrac w * prob p x
            )
    $ V.zip assignments samples

-- | Maximum likelihood classification
classify :: ComponentParams -> Sample -> ComponentIdx
classify params x =
  fst
  $ VB.maximumBy (compare `on` \(_,(w,p))->realToFrac w * prob p x)
  $ VB.indexed params

-- | Score of the maximum likelihood assignment for a set of samples
maxLikelihoodScore :: ComponentParams -> Samples -> Prob
maxLikelihoodScore params samples =
  let assignments = V.map (classify params) samples
  in scoreAssignments samples params assignments


{-# LANGUAGE NoMonomorphismRestriction, QuasiQuotes #-}
import Control.Arrow
import Control.Monad
import NeuralNetwork
import Numeric.LinearAlgebra.HMatrix
import Text.Printf.TH
import VisualizeFunction

andDataSet = map (vector *** vector) [
    ([0,0], [0]),
    ([0,1], [0]),
    ([1,0], [0]),
    ([1,1], [1])
    ]

xorDataSet = map (vector *** vector) [
    ([0,0], [0]),
    ([0,1], [1]),
    ([1,0], [1]),
    ([1,1], [0])
    ]

nn = randomlyWeightedNetwork 0 [2, 9, 1] logisticAF
trainOn dataset = iterate (stochasticGradientDescent dataset 0.25) nn

andNNs = trainOn andDataSet
xorNNs = trainOn xorDataSet

outputs dataset nns = map (\nn -> map (applyNN nn) (map fst dataset)) nns
andOutputs = outputs andDataSet andNNs
xorOutputs = outputs xorDataSet xorNNs

image = visualizeFunctionBMP (0,0) (1,1) (1024,1024) (0,1) . ((!0).) . applyNN

main = forM_ [0..10] $ \i -> do
    saveImage ([s|neuralnetwork_and_%04d.bmp|] (2^i)) (image (andNNs !! (2^i)))
    saveImage ([s|neuralnetwork_xor_%04d.bmp|] (2^i)) (image (xorNNs !! (2^i)))


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Lib.Utility.NeuralNet where

import qualified Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra.Static
import Numeric.LinearAlgebra.Devel
import Control.Monad.State
import Control.Arrow
import GHC.TypeLits
import Data.Proxy

data Layer (il :: Nat) (is :: Nat) (wc :: Nat) (state :: *) where
    Layer :: state -> (X il is -> (Y il wc, state)) -> (state -> Delta il wc -> (DX il is, state)) -> Layer il is wc state

type ReluState il wc = Maybe (L il wc)
reluLayer :: forall il wc. (KnownNat il, KnownNat wc) =>
   Layer il wc wc (ReluState il wc)
reluLayer = Layer Nothing forwardF backwardF
    where forwardF x = (lt0, Just lt0)
              where (Just lt0) = create $ LA.cmap (\ e -> if e <= 0 then 0 else e) $ extract x
          backwardF Nothing _ = error "State is Nothing"
          backwardF (Just lt0) d = (zipWithM' (\ e1 e2 -> if e1 == 0 then 0 else e2) lt0 d, Just lt0)

type SigmoidState il wc = Maybe (L il wc)
sigmoidLayer :: forall il wc. (KnownNat il, KnownNat wc) =>
   Layer il wc wc (ReluState il wc)
sigmoidLayer = Layer Nothing forwardF backwardF
    where forwardF x = (y, Just y)
              where y = 1 / (1 + exp (-x))
          backwardF Nothing _ = error "State Y is Nothing"
          backwardF (Just s) d = (d * (1 - s) * s, Just s)

type AffineState il is wc = (NeuralParam is wc, Maybe (X il is), Maybe (NeuralParam is wc))
affineLayer :: forall il is wc. (KnownNat il, KnownNat is, KnownNat wc) =>
    NeuralParam is wc -> Layer il is wc (AffineState il is wc)
affineLayer param@(w, b) = Layer (param, Nothing, Nothing) forwardF backwardF
    where forwardF x = (x <> w + v2m b, (param, Just x, Nothing))
          backwardF (_, Nothing, _) _ = error "State X is Nothing"
          backwardF (param@(w, b), Just x, _) d = (dx, (param, Just x, Just (dw, db)))
              where dx = d <> tr w
                    dw = tr x <> d
                    db = unrow $ (1 :: L 1 il) <> d

type SoftmaxWithLossState il os = (T il os, Maybe (Y il os), Maybe (Loss il 1))
softmaxWithLossLayer :: forall il os. (KnownNat il, KnownNat os) =>
    T il os -> Layer il os 1 (SoftmaxWithLossState il os)
softmaxWithLossLayer t = Layer (t, Nothing, Nothing) forwardF backwardF
    where forwardF x = (loss, (t, Just y, Just loss))
              where y = softmax x
                    loss = crossEntropyError t y
          backwardF (_, Nothing, _) d = error "State Y is Nothing"
          backwardF state@(t, Just y, _) d = ((y - t) / fromIntegral il, state)
              where (il, _) = size d

class NeuralNet (a :: Nat -> Nat -> Nat -> *) where
    predict :: forall il is os. (KnownNat il, KnownNat is, KnownNat os)
         => a il is os -> X il is -> Prediction il os
    loss :: forall il is os. (KnownNat il, KnownNat is, KnownNat os)
         => a il is os -> X il is -> Loss il os
    accuracy :: forall il is os. (KnownNat il, KnownNat is, KnownNat os)
         => a il is os -> X il is -> T il os -> Double
    gradients :: forall il is os. (KnownNat il, KnownNat is, KnownNat os)
         => a il is os -> X il is -> T il os -> ()
    training :: forall il is os. (KnownNat il, KnownNat is, KnownNat os)
         => a il is os -> X il is -> T il os -> a il is os

    accuracy a x t = LA.sumElements (extract t1f0) / fromIntegral (size t1f0)
        where y = predict a x
              (y', t') = (maxIndex' y, maxIndex' t)
              t1f0 = uncol $ zipWithM' (\ e1 e2 -> if e1 == e2 then 1 else 0) y' t'

data MnistTwoLayerNet (wc :: Nat) (il :: Nat)  (is :: Nat) (os :: Nat) where
    MnistTwoLayerNet ::
        Layer il is wc (AffineState il is wc)
     -> Layer il wc wc (ReluState il wc)
     -> Layer il wc os (AffineState il wc os)
     -> AdaGrad wc il is os
     -> MnistTwoLayerNet wc il is os
instance NeuralNet (MnistTwoLayerNet 50) where
    predict (MnistTwoLayerNet af1 relu af2 _) x = y3
        where (y1, af1s) = case af1 of (Layer _ ff _) -> ff x
              (y2, relus) = case relu of (Layer _ ff _) -> ff y1
              (y3, af2s) = case af2 of (Layer _ ff _) -> ff y2
    loss = undefined
    gradients = undefined
    training = undefined

class Updater (a :: Nat -> Nat -> Nat -> Nat -> *) where
    update :: forall wc il is os. (KnownNat wc, KnownNat il, KnownNat is, KnownNat os)
         => a wc il is os
         -> (NeuralParam is wc, NeuralParam wc os)
         -> (NeuralParam is wc, NeuralParam wc os)
         -> ((NeuralParam is wc, NeuralParam wc os), a wc il is os)

data SGD (wc :: Nat) (il :: Nat) (is :: Nat) (os :: Nat) where
    SGD :: Double -> SGD wc il is os
instance Updater SGD where
    update (SGD lr) (p1, p2) (dp1, dp2) = ((newParam lr p1 dp1, newParam lr p2 dp2), SGD lr)

newParam :: forall is wc. (KnownNat is, KnownNat wc) => Double
    -> NeuralParam is wc -> NeuralParam is wc -> NeuralParam is wc
newParam lr (w, b) (dw, db) = (w - lr_w * dw, b - lr_b * db)
    where (lr_w, lr_b) = (realToFrac lr, realToFrac lr)

data AdaGrad (wc :: Nat) (il :: Nat) (is :: Nat) (os :: Nat) where
    AdaGrad :: Double -> (NeuralParam is wc, NeuralParam wc os) -> AdaGrad wc il is os
instance Updater AdaGrad where
    update (AdaGrad lr (h1, h2)) (p1, p2) (dp1, dp2) = (result, AdaGrad lr newH)
        where result = (newParam lr p1 dp1', newParam lr p2 dp2')
              (dp1', dp2') = (f' newH1 dp1, f' newH2 dp2)
                  where f h dx = dx / sqrt (h + 1e-7)
                        f' (hw, hb) (dw, db) = (f hw dw, f hb db)
              newH@(newH1, newH2) = (f h1 dp1, f h2 dp2)
                  where f (hw, hb) (dw, db) = (dw ** 2 + hw, db ** 2 + hb)
newAdaGrad lr = AdaGrad lr ((0, 0), (0, 0))

initTestNet :: forall wc il is os. (KnownNat wc, KnownNat il, KnownNat is, KnownNat os) => Double -> IO (MnistTwoLayerNet wc il is os)
initTestNet wis = result
    where initParams = do
              w1 <- randn'
              let b1 = konst 0
              w2 <- randn'
              let b2 = konst 0
              return (w1 * realToFrac wis, b1, w2 * realToFrac wis, b2)
          result = do
              (w1, b1, w2, b2) <- initParams
              let affine1 = affineLayer (w1, b1)
              let relu = reluLayer
              let affine2 = affineLayer (w2, b2)
              return $ MnistTwoLayerNet affine1 relu affine2 $ newAdaGrad 0.1

type LastLayer il hs os = X il hs -> T il os

type X il is = L il is
type Y il wc = L il wc
type T il os = Y il os
type W is wc = L is wc
type B wc = R wc
type NeuralParam is wc = (W is wc, B wc)
type Prediction il wc = Y il wc
type Loss il wc = Y il wc
type Delta il is = X il is
type DX il is = X il is
type ForwardF il is wc = X il is -> Y il wc
type BackwardF il is wc = Y il wc -> X il is

v2m :: forall m n. (KnownNat m, KnownNat n) => R n -> L m n
v2m v = fromList $ (>>= id) $ replicate mVal $ LA.toList $ extract v
    where mVal = fromIntegral $ natVal (Proxy :: Proxy m)

softmax :: forall m n. (KnownNat m, KnownNat n) => L m n -> L m n
softmax a = expA / sumExpA
    where maxs = foldAndExtractRow' LA.maxElement a
          a' = a - maxs
          expA = exp a'
          sumExpA = foldAndExtractRow' LA.sumElements expA

crossEntropyError :: forall m n. (KnownNat m, KnownNat n) => L m n -> L m n -> L m 1
crossEntropyError t y = negate $ foldRow' LA.sumElements $ t * log (y + delta)
    where delta = 1e-7

foldRow' :: forall m n. (KnownNat m, KnownNat n) => (LA.Vector Double -> Double) -> L m n -> L m 1
foldRow' f m = fromList $ map (f . extract) $ toRows m

foldAndExtractRow' :: forall m n. (KnownNat m, KnownNat n) => (LA.Vector Double -> Double) -> L m n -> L m n
foldAndExtractRow' f m = fromList $ toRows m >>= (replicate nVal . f . extract)
    where nVal = fromIntegral $ natVal (Proxy :: Proxy n)

randn' :: forall m n. (KnownNat m, KnownNat n) => IO (L m n)
randn' = result
    where mVal = fromIntegral $ natVal (Proxy :: Proxy m)
          nVal = fromIntegral $ natVal (Proxy :: Proxy n)
          result = do
              m <- LA.randn mVal nVal
              return $ case create m of
                  Nothing -> error "Something wrong"
                  Just x -> x

dmmap' :: forall m n. (KnownNat m, KnownNat n) => (Double -> Double) -> L m n -> L m n
dmmap' f m = mapped
    where m' = extract m
          mapped = case create $ LA.cmap f m' of
              Nothing -> error "Something wrong"
              Just x -> x

zipWithM' :: forall m n. (KnownNat m, KnownNat n) => (Double -> Double -> Double) -> L m n -> L m n -> L m n
zipWithM' f m1 m2 = result
    where m1' = extract m1
          m2' = extract m2
          result = build (\ r c -> f (atIndex m1' r c) (atIndex m2' r c))
              where atIndex m r c = LA.atIndex m (truncate r, truncate c)

maxIndex' :: forall m n. (KnownNat m, KnownNat n) => L m n -> L m 1
maxIndex' = foldRow' (fromIntegral . LA.maxIndex)


-- | Module for operating on vectors and matrices. It includes some functions
-- for using bra-ket notation, based on data types and functions from 'hmarix'
-- package.
module Hoqus.Dirac where

import Numeric.LinearAlgebra.Data
import Numeric.LinearAlgebra

-- This can be used to hide the standard operators.
-- import Prelude hiding ((+))
-- import qualified Prelude

-- | Function 'ket' provides basic functionality for producing vectors from the
-- canonical basis. Vectors and matrices are not represented as lists (as in
-- Mathematica), but using data types from "Numeric.LinearAlgebra.Data" in
-- 'hmatrix' package. This makes possible to use operations on vectors  and
-- matrices defined in 'hmatrix'.
ket :: Int -> Int -> Vector C
ket d i = fromList $ (take i zeros) ++ [1] ++ (take (d-1-i) zeros)
  where zeros = replicate (d-1) 0

-- | The 'proj' function builds d-dimensional operator |i><j|.
proj :: Int -> Int -> Int -> Matrix C
proj d i j = (d><d) $ outerWith (*) (toList $ ket d i) (toList $ ket d j)

-- | The 'outer' defines the form of outer product resulting in paris of elements.
outer :: [a] -> [b] -> [(a,b)]
outer a b = [ (x,y) | x<-a, y<-b ]

-- | General form of the outer product resulting in paris of elements.
outerWith :: (a -> b -> c) -> [a] -> [b] -> [c]
outerWith f a b = [ f x y | x<-a, y<-b ]

-- | Function `overlap` is just an alternative for the scalar product from
-- hmatrix.
overlap :: Vector C -> Vector C -> C
overlap = (<.>)


module Mandroots where

import Data.Complex

data Qaudratic a = Qaudratic a a a deriving (Show, Eq)

solveQuadratic :: Qaudratic (Complex Double) -> (Complex Double, Complex Double)
solveQuadratic (Qaudratic a b c) = let d = sqrt (b*b - 4 * a * c) in ((-b + d)/(2 * a), (-b - d)/(2 * a))

evalQuadratic :: Num a => Qaudratic a -> a -> a
evalQuadratic (Qaudratic a b c) x =  x*x * a + x * b + c

rootsIteration :: Qaudratic (Complex Double) -> Qaudratic (Complex Double)
rootsIteration q@(Qaudratic a b c) = let (rt1, rt2) = solveQuadratic q in (Qaudratic rt1 rt2 c)

type QCD = Qaudratic (Complex Double)

timeToEscape :: (QCD -> QCD) -> QCD -> Int -> Int
timeToEscape f q max = count q 0
  where
    count p n
      | n ==  max = max
      | let a = evalQuadratic p (1 :+ 0) in realPart (a * conjugate a) > 25 = n
      | otherwise = count (f q) (n + 1)


import AI.HNN.FF.Network
import Numeric.LinearAlgebra

samples :: Samples Double
samples = [
  (fromList [ 1, 1, 1, 1, 1
            , 0, 0, 0, 0, 1
            , 0, 0, 1, 1, 1
            , 0, 0, 0, 0, 1
            , 1, 1, 1, 1, 1 ], fromList [1]),   -- three

  (fromList [ 1, 1, 1, 1, 1
            , 0, 0, 0, 0, 1
            , 1, 1, 1, 1, 1
            , 0, 0, 0, 0, 1
            , 1, 1, 1, 1, 1 ], fromList [1]),   -- three

  (fromList [ 0, 1, 1, 1, 1
            , 0, 0, 0, 0, 1
            , 0, 0, 0, 1, 1
            , 0, 0, 0, 0, 1
            , 0, 1, 1, 1, 1 ], fromList [1]),   -- three

  (fromList [ 1, 1, 1, 1, 1
            , 1, 0, 0, 0, 1
            , 1, 0, 0, 0, 1
            , 1, 0, 0, 0, 1
            , 1, 1, 1, 1, 1 ], fromList [0]),   -- not a three

  (fromList [ 1, 1, 1, 1, 1
            , 1, 0, 0, 0, 0
            , 1, 0, 0, 0, 0
            , 1, 0, 0, 0, 0
            , 1, 0, 0, 0, 0 ], fromList [0]),   -- not a three

  (fromList [ 0, 1, 1, 1, 0
            , 0, 1, 0, 1, 0
            , 0, 1, 1, 1, 0
            , 0, 1, 0, 1, 0
            , 0, 1, 1, 1, 0 ], fromList [0]),   -- not a three

  (fromList [ 0, 0, 1, 0, 0
            , 0, 1, 1, 0, 0
            , 1, 0, 1, 0, 0
            , 0, 0, 1, 0, 0
            , 0, 0, 1, 0, 0 ], fromList [0]) ]  -- not a three

main :: IO ()
main = do
  n <- createNetwork 25 [250] 1
  let n' = trainNTimes 10000 0.5 tanh tanh' n samples
  mapM_ (print . output n' tanh . fst) samples
  putStrLn "-------------"
  print . output n' tanh $ testInput

  where testInput = fromList [ 0, 0, 1, 1, 1
                             , 0, 0, 0, 0, 1
                             , 0, 0, 1, 1, 1
                             , 0, 0, 0, 0, 1
                             , 0, 0, 1, 1, 1 ]

{-

OUTPUT:
fromList [0.9996325368507625]
fromList [0.9997784075859734]
fromList [0.9996165887689248]
fromList [-2.8107935971909852e-2]
fromList [7.001808876464477e-3]
fromList [2.54989546107178e-2]
fromList [5.286805464313172e-4]
-------------
fromList [0.9993713524712442]
-}


-- DTMF Analyzer
-- Haskell School Project (Faculty of Mathematics and Physics, Charles University in Prague)
-- Author: Lukas Jendele
import System.Environment
import qualified Data.ByteString.Lazy as B
import qualified Data.Complex as C
import Data.Binary.Get
import Data.Word
import qualified Numeric as Numeric
import qualified Data.List as List
import qualified Data.Set as Set
import qualified Data.Map as Map
import qualified Data.Function as DataFunction

nthPrimitiveSquareRoot :: Int -> C.Complex Float
nthPrimitiveSquareRoot 0 = 1
nthPrimitiveSquareRoot n = C.conjugate $ C.mkPolar 1 (2*pi/(fromIntegral n))

-- returns filters odd-indexed elements
takeOdd :: [C.Complex a] -> [C.Complex a]
takeOdd lst = map snd $ filter (\(x,_) -> x `mod` 2 == 1) $ zip [0..] lst
-- returns filters even-indexed elements
takeEven :: [C.Complex a] -> [C.Complex a]
takeEven lst = map snd $ filter (\(x,_) -> x `mod` 2 == 0) $ zip [0..] lst

--FFT see: http://mj.ucw.cz/vyuka/ads/44-fft.pdf
fftHelper1 :: Int -> Int -> C.Complex Float -> [C.Complex Float] -> [C.Complex Float] -> [C.Complex Float]
fftHelper1 _ _ _ [] [] = []
fftHelper1 i n omega (e:evenLst) (o:oddLst) | i == n    = []
                                            | otherwise = (e + ((omega^i) * o)):(fftHelper1 (i+1) n omega evenLst oddLst)
--FFT see: http://mj.ucw.cz/vyuka/ads/44-fft.pdf
fftHelper2 :: Int -> Int -> C.Complex Float -> [C.Complex Float] -> [C.Complex Float] -> [C.Complex Float]
fftHelper2 _ _ _ [] [] = []
fftHelper2 i n omega (e:evenLst) (o:oddLst) | i == n    = []
                                            | otherwise = (e - ((omega^i) * o)):(fftHelper2 (i+1) n omega evenLst oddLst)

--FFT see: http://mj.ucw.cz/vyuka/ads/44-fft.pdf
--fft :: (Integral a) => a -> C.Complex Float -> [C.Complex Float] -> [C.Complex Float]
fft :: Int -> C.Complex Float -> [C.Complex Float] -> [C.Complex Float]
fft n omega lst = if n == 1 then (take 1 lst) else
                    let evenLst = (fft (n `div` 2) (omega^2) (takeEven lst))
                        oddLst  = (fft (n `div` 2) (omega^2) (takeOdd lst))
                    in ((fftHelper1 0 (n `div` 2) omega evenLst oddLst) ++ (fftHelper2 0 (n `div` 2) omega evenLst oddLst))

--fftForward :: (Integral a, RealFloat a) => a -> [C.Complex Float] -> [C.Complex Float]
fftForward :: Int -> [C.Complex Float] -> [C.Complex Float]
fftForward n lst = fft n (nthPrimitiveSquareRoot n) lst

-- divides complex number by components
div' :: C.Complex Float -> Int -> C.Complex Float
div' c r = ((C.realPart c) / (fromIntegral r)) C.:+ ((C.imagPart c) / (fromIntegral r))

--fftBackward :: (Integral a, RealFloat a) => a -> [C.Complex a] -> [C.Complex a]
fftBackward :: Int -> [C.Complex Float] -> [C.Complex Float]
fftBackward n lst = map (`div'` n) $ fft n ((C.conjugate (nthPrimitiveSquareRoot n))) lst

main = do
    (filePath:_) <- getArgs
    analyzeFile filePath

--loads the WaveFile from a file given a path
parseFile filePath = do
    fileContents <- B.readFile filePath
    let result = runGet getWaveFile fileContents
    putStrLn (filePath ++ " loaded!")
    return (result)

data WaveFile = WaveFile {
    chunkID ::      Int,
    fileSize ::     Int,
    riffType ::     Int,
    fmtId    ::     Int,
    fmtSize ::      Int,
    fmtCode ::      Int,
    channels ::     Int,
    sampleRate ::   Int,
    fmtAvgBPS ::    Int,
    fmtBlockAlign ::Int,
    bitDepth ::     Int,
    dataID ::       Int,
    dataSize ::     Int,
    fileData ::     [Float]
} -- deriving (Show)

instance Show WaveFile where
    show wav = "FileSize: " ++ (show $ fileSize wav) ++ "\nDataSize: " ++ (show $ dataSize wav) ++ "\nSampleRate: " ++ (show $ sampleRate wav) ++ "\nBitDepth: " ++ (show $ bitDepth wav) ++ "\nChannels: " ++ (show $ channels wav) ++ "\n"

-- parses the header of wav file and reads its contents
getWaveFile :: Get WaveFile
getWaveFile = do
    w_chunkID <- getWord32le
    let chunkID = (fromIntegral w_chunkID)
    w_fileSize <- getWord32le
    let fileSize = fromIntegral w_fileSize
    w_riffType <- getWord32le
    let riffType = fromIntegral w_riffType
    w_fmtID <- getWord32le
    let fmtID = fromIntegral w_fmtID
    w_fmtSize <- getWord32le
    let fmtSize = fromIntegral w_fmtSize
    w_fmtCode <- getWord16le
    let fmtCode = fromIntegral w_fmtCode
    w_channels <- getWord16le
    let channels = fromIntegral w_channels
    w_sampleRate <- getWord32le
    let sampleRate = fromIntegral w_sampleRate
    w_fmtAvgBPS <- getWord32le
    let fmtAvgBPS = fromIntegral w_fmtAvgBPS
    w_fmtBlockAlign <- getWord16le
    let fmtBlockAlign = fromIntegral w_fmtBlockAlign
    w_bitDepth <- getWord16le
    let bitDepth = fromIntegral w_bitDepth
    w_dataID <- getWord32le
    let dataID = fromIntegral w_dataID
    w_dataSize <- getWord32le
    let dataSize = fromIntegral w_dataSize
    fileData <- getWaveContents 0 dataSize
    return $! WaveFile chunkID fileSize riffType fmtID fmtSize fmtCode channels sampleRate fmtAvgBPS fmtBlockAlign bitDepth dataID dataSize fileData

--getWaveContents :: (Integral a) => a -> a-> Get [Float]
getWaveContents i size = do
    if i == size
       then return []
       else do
           fl <- getFloat
           rest <- getWaveContents (i+2) size
           return (fl:rest)

-- turns UInt16 to Int16
wordToInt :: Word16 -> Int
wordToInt word | word < notWord = (fromIntegral word :: Int)
               | otherwise      = (-1) * (fromIntegral notWord :: Int)
               where notWord = (-word)

-- read 16bit Float from the stream
getFloat :: Get Float
getFloat = do
    word <- getWord16le
    let signedWord = wordToInt word
    return ((fromIntegral signedWord :: Float)/(2^15))

-- Duration of the file
wavDuration :: WaveFile -> Float
wavDuration wav = ((fromIntegral $ dataSize wav) * (fromIntegral $ channels wav) * 8) / ((fromIntegral $ sampleRate wav) * (fromIntegral $ bitDepth wav) )

-- turns the list into a list of small intervals
intervalLength = 100

intervals_internal [] = []
intervals_internal lst@(head:tail) = (take (intervalLength) lst) : (intervals_internal tail)

intervals lst = filter (\l -> (length l) == (intervalLength)) (intervals_internal lst)

-- calculates the average of the absolute values in the list
my_average :: [Float] -> Float
my_average lst = (sum $ map (abs) lst) / (fromIntegral $ length lst)

-- get the interval of values from to in msec
getInterval :: Float -> Float -> WaveFile -> [Float]
getInterval from to wav = let rate = ((fromIntegral $ sampleRate wav) :: Float)
                              fromIndex = floor $ (rate * from)
                              toIndex = floor $ (rate * to)
                          in take (toIndex - fromIndex) (drop fromIndex $ fileData wav)

-- helper function for getBeepIntervals
getBeeps _ [] _ = []
getBeeps i (interval:tailOfIntervals) globalAverage = (i, if globalAverage < intervalAverage then True else False, head interval):(getBeeps (i+1) tailOfIntervals globalAverage)
                                                   where intervalAverage = my_average interval
-- handy function for debugging
indexToTime index wav = (fromIntegral index) / (fromIntegral $ sampleRate wav)

filterBeepValues [] = []
filterBeepValues ((index, isQuiet, value):tail)     |   (not isQuiet)     = []
                                                    |   otherwise         = (index, value):(filterBeepValues tail)

-- splits the wav file into into intervals and returns only non-quiet ones
-- it takes little windows and if the abs average value of the window > abs average of the file, the window is regarded as non-quite
getBeepIntervals wav = let wavData = fileData wav
                           globalAverage = my_average wavData
                           listOfIntervals = intervals wavData
                           beeps = getBeeps 0 listOfIntervals globalAverage
                           listOfBeeps = List.groupBy (\(_, isQuiet1, _) (_, isQuiet2, _) -> isQuiet1 == isQuiet2) beeps
                       in filter (\lst -> length lst > 9) $ filter (not . null) $ map (filterBeepValues) listOfBeeps
                           -- firstLast = map (\lst -> (head lst, last lst)) listOfBeeps
                       -- in  filter (\t -> (snd $ fst t)) firstLast

{-
DTMF keypad frequencies (with sound clips)
            5/1209 Hz 	6/1336 Hz 	7/1477 Hz 	8/1633 Hz
1/697 Hz  	 1 / 5   	 2 / 6    	 3 / 7   	 A / 8
2/770 Hz 	 4 / 10  	 5 / 12   	 6 / 14  	 B / 16
3/852 Hz 	 7 / 15 	 8 / 18   	 9 / 21 	 C / 24
4/941 Hz 	 * / 20 	 0 / 24 	 # / 28 	 D / 32

Each frequency button and frequency has its index. If you multiply the indices of two DTMF frequencies, you get index of the button
-}

-- constant - range around the precise DTMF freq
freqLimit = 15

dtmfFrequencies = [697,770,852,941,1209,1336,1477,1633]

dtmfFrequencyRanges_internal :: [Int] -> [Set.Set Int]
dtmfFrequencyRanges_internal [] = []
dtmfFrequencyRanges_internal (freq:rest) = (Set.fromList [(freq - freqLimit)..(freq + freqLimit)]):(dtmfFrequencyRanges_internal rest)

-- set for looking up if frequency is regarded as DTMF frequency or not
dtmfFrequencyRanges :: [Set.Set Int]
dtmfFrequencyRanges = dtmfFrequencyRanges_internal dtmfFrequencies

-- dictionary to look up chars representing buttons, we look up by the index of the button
dtmfTonesDict = Map.fromList $ zip ([x*y | x<- [1..4], y<-[5..8]]) ['1','2','3','A','4','5','6','B','7','8','9','C','*','0','#','D']

-- returns 0 if no otherwise index of frequency from table above
isDTMFFreq :: Int -> Int
isDTMFFreq freq = snd $ foldl f (1,0) dtmfFrequencyRanges
                where f acc@(index, result) set = if result /= 0 then acc
                                                  else
                                                    if freq `Set.member` set then (index+1, index)
                                                    else (index+1, 0)

-- given FFT Output, it returns recognized DTMF frequencies
analyzeFFTOutput :: Int -> Int -> [C.Complex Float] -> Int -> Float -> [(Int, Float)]
analyzeFFTOutput i n (x:xs) smplRate avrgValue   | i == n          = []
                                                 | i > (n `div` 2) = []
                                                 | otherwise       = let freq = (fromIntegral i) / ((fromIntegral n) / (fromIntegral smplRate))  -- calculate the frequency
                                                                         dfmtFreq = isDTMFFreq $ floor freq
                                                                         currentVal = (C.realPart $ abs x)
                                                                     in if dfmtFreq == 0 then restOfRecurrsion
                                                                        else if currentVal > avrgValue then (dfmtFreq, currentVal):(restOfRecurrsion)
                                                                             else restOfRecurrsion
                                                                     where restOfRecurrsion = analyzeFFTOutput (i+1) n xs smplRate avrgValue

maybeCharToChar :: Maybe Char -> Char
maybeCharToChar Nothing = '?'
maybeCharToChar (Just c) = c

-- analyzes interval of function and returns which button it is dialing
analyzeInterval :: [(Int, Float)] -> Int -> Char
analyzeInterval interval smplRate = let len = length interval
                                        floorPow2 = if len <= 1 then 0 else 2 ^ (floor $ logBase 2 (fromIntegral len)) -- calculate the n for FFT
                                        --startIndex = fst $ head interval
                                        --endIndex = fst $ last interval
                                        cmplxValues = map ((C.:+0).snd) interval -- make the values complex and trim the index
                                        fftImage = fftForward floorPow2 cmplxValues
                                        avrgValue = (sum $ map (C.realPart . abs) fftImage) / (fromIntegral floorPow2)
                                        frequencies = reverse $ List.sortBy (compare `DataFunction.on` snd) $ analyzeFFTOutput 0 floorPow2 fftImage smplRate avrgValue
                                        tone = if (length frequencies) < 2 then 0 else (fst $ head frequencies) * (fst $ head (tail frequencies))
                                    in  maybeCharToChar $ Map.lookup tone dtmfTonesDict

analyzeIntervals :: [[(Int, Float)]] -> Int -> String
analyzeIntervals [] _ = []
analyzeIntervals (interval:rest) smplRate = (analyzeInterval interval smplRate):(analyzeIntervals rest smplRate)

analyzeFile filePath = do
    wav <- parseFile filePath
    let beeps = getBeepIntervals wav
    putStrLn (analyzeIntervals beeps (sampleRate wav))


{-# LANGUAGE BinaryLiterals #-}
{-# LANGUAGE OverloadedLists #-}

module Day03 where

import Data.Bits (Bits (complement, setBit, testBit, (.&.)))
import qualified Data.Text as T
import qualified Data.Vector.Storable as Vec
import Motif (count)
import Numeric.LinearAlgebra.Data (I, Vector)
import qualified Relude.Unsafe as Unsafe
import Utils (binaryLines, showSolutions, elimination)

-- least significant bit first
-- >>> bitsAsInts 0b101
-- [1,0,1,0,0]
bitsAsInts :: Int -> Word32 -> Vector I
bitsAsInts bitLength i =
  Vec.generate
    bitLength
    ( \idx ->
        if i `testBit` idx then 1 else 0
    )

-- count the number of bits in each place value of a list of binary numbers
bitCounts :: Int -> [Word32] -> Vector I
bitCounts bitLength = sum . map (bitsAsInts bitLength)

-- >>> boolVecToBitsLE [False, True, True]
-- 6
boolVecToBitsLE :: Vector Bool -> Word32
boolVecToBitsLE = Vec.ifoldl' step 0
  where
    step :: Word32 -> Int -> Bool -> Word32
    step acc idx b
      | b = acc `setBit` idx
      | otherwise = acc

calculateGamma :: Int -> [Word32] -> Word32
calculateGamma bitLength is =
  boolVecToBitsLE . Vec.map isMostCommonBit . bitCounts bitLength $ is
  where
    len = fromIntegral $ length is

    isMostCommonBit :: I -> Bool
    isMostCommonBit cnt = cnt > (len `div` 2)

exampleReport :: [Word32]
exampleReport =
  [ 0b00100,
    0b11110,
    0b10110,
    0b10111,
    0b10101,
    0b01111,
    0b00111,
    0b11100,
    0b10000,
    0b11001,
    0b00010,
    0b01010
  ]

data RatingType = Oxygen | CO2

-- >>> calculateRating 5 Oxygen exampleReport
-- 23
-- >>> calculateRating 5 CO2 exampleReport
-- 10

calculateRating :: Int -> RatingType -> [Word32] -> Word32
calculateRating bitLength ratingType is = final
  where
    Just final =
      evalState
        (elimination (elimByBitCriteria ratingType) is)
        (bitLength - 1)

elimByBitCriteria :: RatingType -> [Word32] -> State Int [Word32]
elimByBitCriteria ratingType candidates = do
  bitIndex <- get
  let ones = count (`testBit` bitIndex) candidates
      zeros = length candidates - ones
      theBit = case ratingType of
        Oxygen -> ones >= zeros
        CO2 -> ones < zeros
      remaining = filter ((== theBit) . (`testBit` bitIndex)) candidates
  put $ bitIndex - 1
  pure remaining

solve :: Text -> Text
solve input = showSolutions p1 p2
  where
    Just is = binaryLines input

    bitLength :: Int
    bitLength = T.length . Unsafe.head . T.lines $ input

    gamma = calculateGamma bitLength is
    -- we mask to ensure irrelevant high bits are not set by complement
    nOnes n = 2 ^ n - 1
    mask = nOnes bitLength
    epsilon = complement gamma .&. mask
    p1 = gamma * epsilon

    oxygenRating = calculateRating bitLength Oxygen is
    co2Rating = calculateRating bitLength CO2 is
    p2 = oxygenRating * co2Rating


----------------------------------------------------------------------------
-- |
-- Module      :  ListBasedVectorProperties
-- Copyright   :  (c) Thomas Peiselt 2018
-- License     :  Public Domain (see COPYRIGHT)
--
-- property-based tests for the list-based implementation of the linear algebra
-- packages.
module Chapter04.LinearAlgebra.ListBasedVectorProperties (
        linearAlgebraTests )

  where

import           Chapter04.LinearAlgebra.ListBased
import           Data.Complex                      (Complex)
import           Test.Invariant
import           Test.QuickCheck
import           Test.Tasty
import           Test.Tasty.QuickCheck

linearAlgebraTests :: TestTree
linearAlgebraTests = testProperties "Linear Algebra"
  [ ("addition of zero", property (plusZero :: [Double] -> Bool))
  , ("commutative add", property (commutativeAdd :: [Double] -> [Double] -> Bool))
  -- associativity doesn't hold for double due to lack of numeric precision
  , ("associative add", property (associativeAdd :: [Int] -> [Int] -> [Int] -> Bool))
  , ("subtraction of zero", property (minusZero :: [Double] -> Bool))
  , ("subtraction from itself", property (minusItself :: [Integer] -> Bool))
  , ("subtraction inverts addition", property (subtractionInvertsAddition :: [Rational] -> [Rational] -> Property))
  , ("vector sum two", property (sum2 :: [Double] -> [Double] -> Bool))
  , ("vector sum three", property (sum3 :: [Double] -> [Double] -> [Double] -> Bool))
  , ("times zero", property (timesZero :: [Double] -> Bool))
  , ("times one", property (timesOne :: [Double] -> Bool))
  , ("times two", property (timesTwo :: [Complex Double]  -> Bool))
  , ("mean singleton", property (singletonMean :: Complex Double -> Bool))
  , ("magnitude is not negative", property (magnitudeNonNegative :: [Double] -> Bool))
  , ("negate . negate does nothing", property (involutoryNegate :: [Double] -> Bool))
  , ("mean of negated", property meanOfNegated)
  , ("distance to self", property (distanceToSelf :: [Double] -> Property))
  , ("distance to zero", property (distanceToZero :: [Double] -> Property))
  ]

plusZero :: (Num a, Eq a) => Vector a -> Bool
plusZero xs = vectorAdd xs (repeat 0) == xs

commutativeAdd :: (Num a, Eq a) => Vector a -> Vector a -> Bool
commutativeAdd = commutative vectorAdd

associativeAdd :: (Num a, Eq a) => Vector a -> Vector a -> Vector a -> Bool
associativeAdd = associative vectorAdd

minusZero :: (Num a, Eq a) => Vector a -> Bool
minusZero xs =  vectorSubtract xs (repeat 0) == xs

minusItself :: (Num a, Eq a) => Vector a -> Bool
minusItself xs = vectorSubtract xs xs == (replicate (length xs) 0)

subtractionInvertsAddition :: (Num a, Eq a) => Vector a -> Vector a -> Property
subtractionInvertsAddition xs ys = length xs <= length ys ==>
  inverts (`vectorAdd` ys) (`vectorSubtract` ys) xs

sum2 :: (Num a, Eq a) => Vector a -> Vector a -> Bool
sum2 xs ys = vectorSum [xs, ys] == vectorAdd xs ys

sum3 :: (Num a, Eq a) => Vector a -> Vector a -> Vector a -> Bool
sum3 xs ys zs = vectorSum [xs, ys, zs] == vectorAdd xs ( vectorAdd ys zs )

timesZero :: (Num a, Eq a) => Vector a -> Bool
timesZero xs = scalarMultiply 0 xs == replicate (length xs) 0

timesOne :: (Num a, Eq a) => Vector a -> Bool
timesOne xs = scalarMultiply 1 xs == xs

timesTwo :: (Num a, Eq a) => Vector a -> Bool
timesTwo xs = scalarMultiply 2 xs == vectorAdd xs xs

singletonMean :: (Fractional a, Eq a) => a -> Bool
singletonMean x = vectorMean [x] == x

involutoryNegate :: (Eq a, Num a) => Vector a -> Bool
involutoryNegate = involutory negateVector

meanOfNegated :: [Double] -> Property
meanOfNegated xs =  not (null xs) ==> vectorMean xs == -(vectorMean $ negateVector xs)

magnitudeNonNegative :: (Floating a, Eq a, Ord a) => Vector a -> Bool
magnitudeNonNegative xs = magnitude xs >= 0.0

distanceToSelf :: (Floating a, Eq a) => Vector a -> Property
distanceToSelf xs = not (null xs) ==> distance xs xs == 0

distanceToZero :: (Floating a, Eq a, Enum a) => Vector a -> Property
distanceToZero xs = not (null xs) ==>
  distance xs (replicate (length xs) 0) == magnitude xs


module Fractal.Math.Kernel(
  evaluateFractal
)where

import Data.Complex

type Point = Complex Double
type Limit = Double
type Iteration = Complex Double
type Magnitude = Double
type MaxIter = Int
type Iter = Int

evaluateFractal :: (Point -> Iteration -> Iteration) -> Limit -> MaxIter -> Point -> Iter
evaluateFractal f l m c = res c (0.0 :+ 0.0) 0
  where res c z i             = if i == m then i else rest (f c z) i
        rest val i             = if checkAcceptance l (magnitude val) then i else res c val (i + 1)

checkAcceptance :: Limit -> Double -> Bool
checkAcceptance l val
  | val >= l    = True
  | val == 0.0  = True
  | otherwise   = False


{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Main where

import Lib

import Numeric.AD
import Numeric.AD.Internal.Forward
import Numeric.AD.Rank1.Tower(Tower)
import Numeric.SpecFunctions
import Data.Csv
import qualified Data.ByteString.Lazy.Char8 as C

--c :: Num a => (a -> a) -> a -> a -> a
difference :: Fractional a => (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> (forall s. AD s (Tower a) -> AD s (Tower a))
difference v base = \x -> v x - v (auto base)

difference' :: Num a => (t -> a) -> t -> t -> a
difference' v base x = v x - v base

optValLin :: Num a => a -> a
optValLin r = 1000 * r

optValQuad :: Num a => a -> a
--optValQuad r = 1000 * (r + r*r)
optValQuad r = 10000 * (r*r) + 10* r

-- Truncate a taylor exapansion
-- If nothing don't take a taylor expansion just evaluate the 0th derivative a.k.a zero
trunTay :: Fractional a => Maybe Int -> (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> a -> a
trunTay Nothing  f centre x =  (diffs f x)!!0
trunTay (Just n) f centre x = foldr (+) 0 $ take (n+1) $ (taylor0 f centre (x-centre))

-- create convexity correction function
-- takes a function and two points and calculates the adjustment needed for the difference
convexityAdj :: Fractional a => (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> a -> a
convexityAdj f centre b = difference' (\x -> trunTay (Just x) (difference f centre) centre b ) 1 2

main :: IO ()
main = do
    --tr (Just 1) (c d 0.02) 0.02 0.04
    print $ trunTay (Just 0) (difference optValLin 0.02) 0.02 0.04
    --print $ trunTay (Just 1) (difference optionVal 0.02) 0.02 0.04
    --print $ trunTay (Just 2) (difference optionVal 0.02) 0.02 0.04
    --print $ trunTay Nothing (difference optionVal 0.02) 0.02 0.04

    print $ "convexity"
    print $ convexityAdj optValLin 0.02 0.04
    print $ convexityAdj optValQuad 0.02 0.04
    -- just get the second derivative of the adjustment

    -- plot a function trunTay for a range from like 0 -> 0.1 for examples with Nothing
    let xFunValues = [ x*0.001 | x <- [1..100] ]
        linValues = map (\x -> trunTay Nothing optValLin 0.02 x) xFunValues
        quadValues = map (\x -> trunTay Nothing optValQuad 0.02 x) xFunValues

--    print $ (linValues :: [(Double,Double)])
    C.writeFile "./plotting/funVals.csv" $ encode (zip3 xFunValues linValues quadValues :: [(Double,Double,Double)])

    -- Then plot first order derivaitve trunTay (Just 1) centred at 0.02 up to 0.04 to get straight line
    let xApproxValues = [ x*0.001 + 0.02 | x <- [0..60] ]
        quadApproxVals1 = map (\x -> trunTay (Just 1) optValQuad  0.02 x) xApproxValues
        quadApproxVals2 = map (\x -> trunTay (Just 2) optValQuad  0.02 x) xApproxValues

    C.writeFile "./plotting/approxVals.csv" $ encode (zip3 xApproxValues quadApproxVals1 quadApproxVals2 :: [(Double,Double,Double)])

    -- Now plot the convexity adustment as a constant values starting at 0.8
    let sPoint = trunTay (Just 1) optValQuad  0.02 0.08
        convexAdj = (convexityAdj optValQuad 0.02 0.08)
    writeFile "./plotting/convexAdj.csv" $ "0.08," ++ show sPoint ++ ",0," ++ show convexAdj

    print "done"


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}

module Statistics.LinearRegression (
    -- * Simple linear regression functions
    linearRegression,
    linearRegressionRSqr,
    linearRegressionTLS,
    -- * Related functions
    correl,
    covar,
    -- * Estimated errors and distribution parameters
    linearRegressionMSE,
    linearRegressionDistributions,
    -- * Robust linear regression
    robustFit,
    nonRandomRobustFit,
    robustFitRSqr,
    -- ** Related types
    EstimationParameters(..),
    ErrorFunction,
    Estimator,
    EstimatedRelation,
    -- ** Provided values
    defaultEstimationParameters,
    linearRegressionError,
    linearRegressionTLSError,
    -- ** Helper functions
    converge,
    -- * References
    -- $references
    ) where

import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U
import Data.Vector.Generic (Vector, (!))
import Safe (at)
import System.Random
import System.Random.Shuffle (shuffleM)
import Control.Monad.Random.Class
import Control.Monad.Random (evalRand)
import Control.Monad (liftM)
import Data.Function (on)
import Data.List (minimumBy, sortBy)
import Data.Maybe (fromMaybe)
import qualified Statistics.Sample as S
import qualified Statistics.Distribution as D
import qualified Statistics.Distribution.Transform as T
import qualified Statistics.Distribution.StudentT as ST

--- * Simple linear regression

-- | Covariance of two samples
covar :: Vector v Double => v Double -> v Double -> Double
covar xs ys = covar' m1 m2 n xs ys
    where
          !n = fromIntegral $ G.length xs
          !m1 = S.mean xs
          !m2 = S.mean ys
{-# INLINE covar #-}

-- internal function that avoids duplicate calculation of means and lengths where possible
-- Note: trying to make the calculation even more efficient by subtracting m1*m1*n instead of individual subtractions increased errors, probably due to rounding issues.
covar' :: Vector v Double => Double -> Double -> Double -> v Double -> v Double -> Double
covar' m1 m2 n xs ys = G.sum (G.zipWith (*) (G.map (subtract m1) xs) (G.map (subtract m2) ys)) / (n-1)
{-# INLINE covar' #-}

-- | Pearson's product-moment correlation coefficient
correl :: Vector v Double => v Double -> v Double -> Double
correl xs ys = let !c = covar xs ys
                   !sx = S.stdDev xs
                   !sy = S.stdDev ys
               in c / (sx * sy)
{-# INLINE correl #-}

-- | Simple linear regression between 2 samples.
--   Takes two vectors Y={yi} and X={xi} and returns
--   (alpha, beta, r*r) such that Y = alpha + beta*X
--   and where r is the Pearson product-moment correlation
--   coefficient
linearRegressionRSqr :: Vector v Double => v Double -> v Double -> (Double, Double, Double)
linearRegressionRSqr xs ys = (alpha, beta, r2)
    where
          !c                   = covar' m1 m2 n xs ys
          !r2                  = c*c / (v1*v2)
          !(m1,v1)             = S.meanVarianceUnb xs
          !(m2,v2)             = S.meanVarianceUnb ys
          !n                   = fromIntegral $ G.length xs
          !beta                = c / v1
          !alpha               = m2 - beta * m1
{-# INLINE linearRegressionRSqr #-}

-- | Simple linear regression between 2 samples.
--   Takes two vectors Y={yi} and X={xi} and returns
--   (alpha, beta) such that Y = alpha + beta*X
linearRegression :: Vector v Double => v Double -> v Double -> (Double, Double)
linearRegression xs ys = (alpha, beta)
    where
        (alpha, beta, _) = linearRegressionRSqr xs ys
{-# INLINE linearRegression #-}

-- | The error (or residual) mean square of a sample w.r.t. an estimated regression line.
--   This serves as an estimate for the variance of the sampled data.
--   Accepts the regression parameters (alpha,beta) and the sample vectors X and Y.
linearRegressionMSE :: (Vector v Double, Vector v (Double, Double)) => (Double,Double) -> v Double -> v Double -> Double
linearRegressionMSE ab xs ys = (G.sum . G.map (linearRegressionError ab) . G.zip xs $ ys)/(n-2)
    where
        !n = fromIntegral $ G.length xs

-- | The estimated distributions of the regression parameters (alpha and beta) assuming normal, identical distributions of Y, the sampled data.
-- These can serve to get confidence intervals for the regression parameters.
-- Accepts the regression parameters (alpha,beta) and the sample vectors X and Y.
-- The distributions are StudnetT distributions centered at the estimated (alpha,beta) respectively, with parameter numbers n-2 (where n is the initial sample size) and with standard deviations that are extracted from the sampled data based on its MSE. See chapter 2 of reference [3] for details.
linearRegressionDistributions :: (Vector v Double, Vector v (Double, Double)) => (Double,Double) -> v Double -> v Double -> (T.LinearTransform ST.StudentT,T.LinearTransform ST.StudentT)
linearRegressionDistributions (alpha,beta) xs ys = (ST.studentTUnstandardized (n-2) alpha va,ST.studentTUnstandardized (n-2) beta vb)
    where
        !n = fromIntegral $ G.length xs
        !mse = linearRegressionMSE (alpha,beta) xs ys
        !vb = mse/(xv)
        !mx = S.mean xs
        !va = mse*(1/n+mx^2/xv)
        !xv = G.sum . G.map (\x -> (x-mx)^2) $ xs

-- | Total Least Squares (TLS) linear regression.
-- Assumes x-axis values (and not just y-axis values) are random variables and that both variables have similar distributions.
-- interface is the same as 'linearRegression'.
linearRegressionTLS :: Vector v Double => v Double -> v Double -> (Double,Double)
linearRegressionTLS xs ys = (alpha, beta)
    where
          !c                   = covar' m1 m2 n xs ys
          !b                   = (v1 - v2) / c
          !(m1,v1)             = S.meanVarianceUnb xs
          !(m2,v2)             = S.meanVarianceUnb ys
          !n                   = fromIntegral $ G.length xs
          !betas               = [(-b - sqrt(b^2+4))/2,(-b + sqrt(b^2+4)) /2]
          !beta                = if c > 0 then maximum betas else minimum betas
          !alpha               = m2 - beta * m1
{-# INLINE linearRegressionTLS #-}

-- | An estimated linear relation between 2 samples is (alpha,beta) such that Y = alpha + beta*X.
type EstimatedRelation = (Double,Double)

-- | An 'Estimator' is a function that generates an estimated linear regression based on 2 samples. This module provides two estimator functions:
-- 'linearRegression' and 'linearRegressionTLS'
type Estimator = (S.Sample -> S.Sample -> EstimatedRelation)

-- | An 'ErrorFunction' is a function that computes the error of a given point from an estimate. This module provides two error functions correspoinding to the two 'Estimator' functions it defines:
--
-- * Vertical distance squared via 'linearRegressionError' that should be used with 'linearRegression'
--
-- * Total distance squared vie 'linearRegressionTLSError' that should be used with 'linearRegressionTLS'
type ErrorFunction = (EstimatedRelation -> (Double,Double) -> Double)

-- | The robust fit algorithm used has various parameters that can be specified using the 'EstimationParameters' record.
data EstimationParameters = EstimationParameters {
    -- | Maximal fraction of outliers expected in the sample (default 0.25)
    outlierFraction     :: !Double,
    -- | Number of concentration steps to take for initial evaluation of a solution (default 3)
    shortIterationSteps :: !Int,
    -- | Maximal number of sampled subsets (pairs of points) to use as starting points (default 500)
    maxSubsetsNum       :: !Int,
    -- | If the initial sample is large, and thus gets subdivided, this is the number of candidate-estimations to take from each subgroup, on which complete convergence will be executed (default 10)
    groupSubsets        :: !Int,
    -- | Maximal size of sample that can be analyzed without any sub-division (default 600)
    mediumSetSize       :: !Int,
    -- | Maximal size of sample that does not require two-step sub-division (see reference article) (default 1500)
    largeSetSize        :: !Int,
    -- | Estimator function to use (default linearRegression)
    estimator           :: Estimator,
    -- | ErrorFunction to use (default linearRegressionError)
    errorFunction       :: ErrorFunction
    }

-- | Default set of parameters to use (see reference for details).
defaultEstimationParameters = EstimationParameters {
    outlierFraction = 0.25,
    shortIterationSteps = 3,
    maxSubsetsNum = 500,
    groupSubsets = 10,
    mediumSetSize = 600,
    largeSetSize = 1500,
    estimator = linearRegression,
    errorFunction = linearRegressionError
}

-- | linearRegression error function is the square of the /vertical/ distance of a point from the line.
linearRegressionError :: ErrorFunction
linearRegressionError (alpha,beta) (x,y) = (y-(beta*x+alpha))^2

-- | linearRegressionTLS error function is the square of the /total/ distance of a point from the line.
linearRegressionTLSError :: ErrorFunction
linearRegressionTLSError (alpha,beta) (x,y) = ey/(1+beta^2)
    where
        ey = linearRegressionError (alpha,beta) (x,y)

-- | Helper function to calculate the minimal expected size of uncontaminated data based on the maximal fraction of outliers.
setSize :: Vector v Double => EstimationParameters -> v Double -> Int
setSize ep xs = max (n `div` 2 + 1) . round $ (1-outlierFraction ep) * (fromIntegral n)
    where
        n = G.length xs

-- | Helper function that, given an initial estimated relation and the error of the perivous estimation, performs a "concentration" step, generating a new estimate based on a fraction of points laying closest to the previous estimate and estimates the error of the previous estimate based on the same fraction.
-- The result is an estimate that is at least as good as the previous one.
-- The reason the error is calculated for the previous parameters is calculation optimization.
concentrationStep :: Vector v Double => EstimationParameters -> v Double -> v Double -> (EstimatedRelation, Double) -> (EstimatedRelation, Double)
concentrationStep ep xs ys (prev, prev_err) = (new_estimate, new_err)
    where
        set_size = setSize ep xs
        xyerrors = map (\p -> (p,errorFunction ep prev p)) $ zip (G.toList xs) (G.toList ys)
        (xys,errors) = unzip . take set_size . sortBy (compare `on` snd) $ xyerrors
        (good_xs,good_ys) = unzip xys
        new_estimate = estimator ep (G.fromList good_xs) (G.fromList good_ys)
        new_err = sum errors

-- | Infinite set of consecutive concentration steps.
concentration :: Vector v Double => EstimationParameters -> v Double -> v Double -> EstimatedRelation -> [(EstimatedRelation, Double)]
concentration ep xs ys params = tail $ iterate (concentrationStep ep xs ys) (params,-1)

-- | Calculate the optimal (local minimum) estimate based on an initial estimate.
-- The local minimum may not be the global (a.k.a. best) estimate but starting from enough different initial estimates should yield the global optimum eventually.
converge :: Vector v Double => EstimationParameters -> v Double -> v Double -> EstimatedRelation -> EstimatedRelation
converge ep xs ys = fst . findConvergencePoint . concentration ep xs ys

-- | The convergence point is defined as the point the error estimate of which is equal to the next estimate's error.
findConvergencePoint :: Ord a => [(b,a)] -> (b,a)
findConvergencePoint (x:y:ys)
    | snd x <= snd y = x -- rounding issues my cause an actual increase in error resulting in an infinite loop so the actual stop condition is when the errors stop decreasing
    | otherwise = findConvergencePoint (y:ys)
findConvergencePoint xs = error "Too short a list for conversion (size < 2)"

-- | Many times there is no need for full concentration as bad initial estimates can be discovered after only a few concentration steps.
concentrateNSteps :: Vector v Double => EstimationParameters -> v Double -> v Double -> EstimatedRelation -> (EstimatedRelation,Double)
concentrateNSteps ep xs ys params = concentration ep xs ys params !! shortIterationSteps ep

-- | Finding a robust fit linear estimate between two samples. The procedure requires randomization and is based on the procedure described in the reference.
robustFit :: (MonadRandom m, Vector v Double) => EstimationParameters -> v Double -> v Double -> m EstimatedRelation
robustFit ep xs ys = do
    let n = G.length xs
-- For optimal performance the exact procedure executed depends on the set size.
    if n < 2
        then
            error "cannot fit an input of size < 2"
        else if n == 2
            then return $ lineParams ((G.head xs,G.head ys),(G.last xs,G.last ys))
            else
                liftM (candidatesToWinner ep xs ys) $ if n < mediumSetSize ep
                    then
                         singleGroupFitCandidates ep Nothing xs ys
                    else if n < largeSetSize ep
                        then largeGroupFitCandidates ep xs ys
                        else do
                            (nxs,nys) <- liftM unzip $ randomSubset (zip (G.toList xs) (G.toList ys)) (largeSetSize ep)
                            largeGroupFitCandidates ep (U.fromList nxs) (G.fromList nys)

-- | Robust fit yielding also the R-square value of the \"clean\" dataset.
robustFitRSqr :: (MonadRandom m, Vector v Double, Vector v (Double, Double)) => EstimationParameters -> v Double -> v Double -> m (EstimatedRelation,Double)
robustFitRSqr ep xs ys = do
    er <- robustFit ep xs ys
    let (good_xs,good_ys) = U.unzip . G.fromList . take (setSize ep xs) . sortBy (compare `on` errorFunction ep er) . G.toList $ G.zip xs ys
    return (er,correl good_xs good_ys ^ 2)

-- | A wrapper that executes 'robustFit' using a default random generator (meaning it is only pseudo-random)
nonRandomRobustFit :: Vector v Double => EstimationParameters -> v Double -> v Double -> EstimatedRelation
nonRandomRobustFit ep xs ys = evalRand (robustFit ep xs ys) (mkStdGen 1)

-- | Given a set of initial estimates converge them all and find the optimal one.
candidatesToWinner :: Vector v Double => EstimationParameters -> v Double -> v Double -> [EstimatedRelation] -> EstimatedRelation
candidatesToWinner ep xs ys = fst . minimumBy (compare `on` snd) . map (findConvergencePoint . concentration ep xs ys)

-- | for a large initial sample - subdivide it, then get candidates from each subgroup. Perform full convergence on all the candidates and return the best ones.
largeGroupFitCandidates :: (MonadRandom m, Vector v Double) => EstimationParameters -> v Double -> v Double -> m [EstimatedRelation]
largeGroupFitCandidates ep xs ys = do
    let n = G.length xs
    let sub_groups_num = n `div` (mediumSetSize ep `div` 2)
    let sub_groups_size = n `div` sub_groups_num
    shuffled <- shuffleM $ zip (G.toList xs) (G.toList ys)
    let sub_groups = map (G.unzip . U.fromList) $ splitTo sub_groups_size shuffled
    let sub_groups_candidates = maxSubsetsNum ep `div` sub_groups_num
    candidates_list <- mapM (applyTo $ singleGroupFitCandidates ep (Just sub_groups_candidates)) sub_groups
    let candidates = concat candidates_list
    return . map fst . take (groupSubsets ep) . sortBy (compare `on` snd) . map (findConvergencePoint . concentration ep xs ys) $ candidates

-- | For a single group (a group that will not be subdivided) pick an initial set of pairs of points, run a few steps on each, then return the most promising candidates.
singleGroupFitCandidates :: (MonadRandom m, Vector v Double) => EstimationParameters -> Maybe Int -> v Double -> v Double -> m [EstimatedRelation]
singleGroupFitCandidates ep m_subsets xs ys = do
    let all_pairs = allPairs $ zip (G.toList xs) (G.toList ys)
    let return_size = fromMaybe (maxSubsetsNum ep) m_subsets
    initial_sets <- if return_size > length all_pairs
        then return all_pairs
        else randomSubset all_pairs return_size
    return . map fst . take (groupSubsets ep) . sortBy (compare `on` snd) . map (concentrateNSteps ep xs ys . lineParams) $ initial_sets

-- | Find the line passing between two points. This is the initial estimate to use given two random points.
lineParams :: ((Double,Double),(Double,Double)) -> EstimatedRelation
lineParams ((x1,y1),(x2,y2)) = (alpha,beta)
    where
        beta = (y2-y1)/(x2-x1)
        alpha = y1 - beta*x1

-- | A list of all possible two-element pairs from a list.
allPairs :: [a] -> [(a,a)]
allPairs [] = []
allPairs [x] = []
allPairs [x,y] = [(x,y)]
allPairs (x:xs) = (zip xs . repeat $ x) ++ allPairs xs

-- | Get a random subset of a given size.
randomSubset :: MonadRandom m => [a] -> Int -> m [a]
randomSubset xs size = liftM (take size) $ shuffleM xs

-- | Split a list into sublists of length n.
splitTo :: Int -> [a] -> [[a]]
splitTo n = map (take n) . takeWhile (not . null) . iterate (drop n)

-- | Helper function to adjust parameter handling
applyTo :: (a->b->c) -> (a,b) -> c
applyTo f (x,y) = f x y

-- $references
--
-- * Two Dimensional Euclidean Regression (Stein) <http://www.dspcsp.com/pubs/euclreg.pdf>
--
-- * Computing LTS Regression For Large Data Sets (Rousseeuw and Driessen) <http://agoras.ua.ac.be/abstract/Comlts99.htm>
--
-- * Applied linear statistical models (Kutner et al.)


module DeterministicSpec where

import           Test.Hspec

import           Data.Complex
import           Data.List                      ( sortBy )
import           Data.Ord                       ( compare )
import           Data.Ratio
import           Pineappl                       ( Prob(P)
                                                , LogProb(LP)
                                                )
import           Pineappl.Deterministic         ( WrappedBDDist(BDDist)
                                                , bddist
                                                , condition
                                                , factor
                                                , runBDDist
                                                , sample
                                                , uniform
                                                )
import           Pineappl.Quantum               ( fromAmplitude )

spec :: Spec
spec = do
  describe "DIPPL tutorial milestones" $ do
    let fairCoin = uniform [0, 1]
        coinFlip = sample fairCoin
    it "should get correct result for 3 coin flips"
      $ BDDist ((\a b c -> a + b + c) <$> coinFlip <*> coinFlip <*> coinFlip)
      `shouldBe` bddist
                   [ (0, P (1 % 8))
                   , (1, P (3 % 8))
                   , (2, P (3 % 8))
                   , (3, P (1 % 8))
                   ]

    it "should get correct result for 3 coin flips (monadic)"
      $          BDDist
                   (do
                     a <- coinFlip
                     b <- coinFlip
                     c <- coinFlip
                     return (a + b + c)
                   )
      `shouldBe` bddist
                   [ (0, P (1 % 8))
                   , (1, P (3 % 8))
                   , (2, P (3 % 8))
                   , (3, P (1 % 8))
                   ]

    it "should get correct result for 'funny binomial'"
      $          BDDist
                   (do
                     a <- coinFlip
                     b <- coinFlip
                     c <- coinFlip
                     factor $ P (if a + b > 0 then 1 else 1 % 4)
                     return (a + b + c)
                   )
      `shouldBe` bddist
                   [ (0, P (1 % 32))
                   , (1, P (9 % 32))
                   , (2, P (3 % 8))
                   , (3, P (1 % 8))
                   ]

  describe "Monty Hall" $ do
    it "should get correct answer with the switch"
      $          BDDist
                   (do
                     initialChoice <- sample $ uniform [1 .. 3]
                     revealed      <- case initialChoice of
                       1 -> sample $ uniform [2, 3]
                       2 -> return 3
                       3 -> return 2
                       _ -> error "oops"
                     case (initialChoice, revealed) of
                       (1, 2) -> return 3
                       (1, 3) -> return 2
                       (2, 3) -> return 1
                       (3, 2) -> return 1
                       _      -> error "oops"
                   )
      `shouldBe` bddist [(1, P (1 % 2)), (2, P (1 % 4)), (3, P (1 % 4))]

  describe "Bayes theorem" $ do
    let bernoulli p = bddist [(True, p), (False, 1 - p)]

    it "should get correct answer for flu testing problem"
      $          BDDist
                   (do
                     hasFlu     <- sample . bernoulli $ P (1 % 10)
                     testResult <- if hasFlu
                       then sample . bernoulli $ P (7 % 10)
                       else sample . bernoulli $ P (1 % 10)
                     condition testResult
                     return hasFlu
                   )
      `shouldBe` bernoulli (P (7 % 16))

    it "should get correct answer for traffic problem"
      $ let president = bernoulli $ P (1 % 100)
            accident  = bernoulli $ P (1 % 10)
            traffic p a = bernoulli . P $ case (p, a) of
              (False, False) -> 1 % 10
              (False, True ) -> 1 % 2
              (True , False) -> 3 % 5
              (True , True ) -> 9 % 10
            joint = do
              p <- sample president
              a <- sample accident
              t <- sample $ traffic p a
              return (p, a, t)
        in  do
              BDDist (joint >>= \(p, a, t) -> condition t >> return a)
                `shouldBe` bernoulli (P (8 % 23))
              BDDist (joint >>= \(p, a, t) -> condition (t && p) >> return a)
                `shouldBe` bernoulli (P (1 % 7))

  describe "Quantum mechanics" $ do
    it "should get correct answer for double-slit experiment"
      $ let
          ampl y = LP $ exp (i * ds)
           where
            i  = 0 :+ 1
            ds = sqrt $ 1 + fromIntegral y ^ 2
          outcome = BDDist
            (do
              slit <- sample $ bddist [ (y, ampl y) | y <- [-1, 1] ]
              sample $ bddist [ (y, ampl (y - slit)) | y <- [-10 .. 10] ]
            )
          intensity = runBDDist $ fromAmplitude outcome
          sorted =
            fmap snd
              . sortBy (flip compare)
              . fmap (\(a, b) -> (b, a))
              $ intensity
        in
          {- Looks like:
              -10 | #
              -9 | ####
              -8 | ######
              -7 | ######
              -6 | ###
              -5 | #
              -4 |
              -3 | ###
              -2 | ######
              -1 | #########
               0 | ###########
               1 | #########
               2 | ######
               3 | ###
               4 |
               5 | #
               6 | ###
               7 | ######
               8 | ######
               9 | ####
              10 | # -}
          take 5 sorted `shouldBe` [0, 1, -1, 8, -8]


import Numeric.LinearAlgebra
import Graphics.Plot

vector l = fromList l :: Vector Double
matrix ls = fromLists ls :: Matrix Double
diagl = diag . vector

f = matrix [[1,0,0,0],
            [1,1,0,0],
            [0,0,1,0],
            [0,0,0,1]]

h = matrix [[0,-1,1,0],
            [0,-1,0,1]]

q = diagl [1,1,0,0]

r = diagl [2,2]

s0 = State (vector [0, 0, 10, -10]) (diagl [10,0, 100, 100])

data System = System {kF, kH, kQ, kR :: Matrix Double}
data State = State {sX :: Vector Double , sP :: Matrix Double}
type Measurement = Vector Double

kalman :: System -> State -> Measurement -> State
kalman (System f h q r) (State x p) z = State x' p' where
    px = f <> x                            -- prediction
    pq = f <> p <> trans f + q             -- its covariance
    y  = z - h <> px                       -- residue
    cy = h <> pq <> trans h + r            -- its covariance
    k  = pq <> trans h <> inv cy           -- kalman gain
    x' = px + k <> y                       -- new state
    p' = (ident (dim x) - k <> h) <> pq    -- its covariance

sys = System f h q r

zs = [vector [15-k,-20-k] | k <- [0..]]

xs = s0 : zipWith (kalman sys) xs zs

des = map (sqrt.takeDiag.sP) xs

evolution n (xs,des) =
    vector [1.. fromIntegral n]:(toColumns $ fromRows $ take n (zipWith (-) (map sX xs) des)) ++
    (toColumns $ fromRows $ take n (zipWith (+) (map sX xs) des))

main = do
    print $ fromRows $ take 10 (map sX xs)
    mapM_ (print . sP) $ take 10 xs
    mplot (evolution 20 (xs,des))


module AI.Learning.Core where

import Numeric.LinearAlgebra
import Numeric.GSL.Minimization

---------------
-- Utilities --
---------------

-- |Sigmoid function:
--
--  > sigmoid x = 1 / (1 + exp (-x))
--
--  Used in the logistic regression and neural network modules.
sigmoid :: Floating a => a -> a
sigmoid x = 1 / (1 + exp (-x))

------------------
-- Optimization --
------------------

-- |Simplified minimization function. You supply functions that compute the
--  quantity to be minimized and its gradient, and an initial guess, and the
--  final solution is returned.
minimizeS :: (Vector Double -> Double)          -- f
          -> (Vector Double -> Vector Double)   -- gradient
          -> Vector Double                      -- initial x
          -> Vector Double
minimizeS f g x = fst $ minimizeVD VectorBFGS2 prec niter sz tol f g x
    where prec    = 1e-9
          niter   = 1000
          sz      = 0.1
          tol     = 0.1


{-# LANGUAGE OverloadedStrings #-}
module Main where

import Lib
import qualified Data.Text.IO as T
import qualified MachineLearning.Types as T
import qualified Numeric.LinearAlgebra as LA
import qualified MachineLearning as ML
import qualified MachineLearning.Classification.Binary as BC

main = do
    -- Step 1. Data loading.
    designMatrix <- LA.loadMatrix "data.txt"
    let (x, y) = ML.splitToXY designMatrix
    -- Step 2. Feature normalization (skipped - we don't need feature normalization if we use BFGS2).
    -- See Linear Regression sample app for dedails.

    -- Step 3. Feature mapping.
        x1 = ML.addBiasDimension $ ML.mapFeatures 6 x

    -- Step 4. Learning
        zeroTheta = LA.konst 0 (LA.cols x1)
        (theta, _) = BC.learn (BC.BFGS2 0.1 0.1) 0.000002 1500 (BC.L2 1) x1 y zeroTheta

    -- Step 5. Prediction and checking accuracy
        accuracy = calcAccuracy x1 y theta

    -- Step 6. Printing results.
    putStrLn "\n=== Logistic Regression Sample Application ===\n"

    putStrLn ""
    putStrLn $ "Theta:             " ++ (show theta)

    putStrLn ""
    putStrLn $ "Accuracy on training set data (%): " ++ show (accuracy*100)

    putStrLn "This is me"
    someFunc'
    T.putStrLn "someFunc"
    putStrLn $ anotherFunc 5


module Main
  ( main
  ) where

import Advent (getInputLines)
import Data.Complex (Complex(..))
import Data.List (foldl')

main :: IO ()
main =
  do i <- getInputLines parse 12
     print (part1 i)
     print (part2 i)

type C = Complex Double

parse :: String -> (Char,C)
parse (x:(read -> n)) = (x,n:+0)
parse xs = error (show xs)

dir, rot :: Char -> Maybe C

dir 'N' = Just (   0  :+   1  )
dir 'E' = Just (   1  :+   0  )
dir 'S' = Just (   0  :+ (-1) )
dir 'W' = Just ( (-1) :+   0  )
dir  _  = Nothing

rot 'L' = Just (   0  :+   1  )
rot 'R' = Just (   0  :+ (-1) )
rot  _  = Nothing

part1, part2 :: [(Char,C)] -> Int
step1, step2 :: (C,C) -> (Char,C) -> (C,C)

part1 ins = abs (round x) + abs (round y)
  where
    x :+ y = fst $ foldl' step1 ( 0 :+ 0 , 1 :+ 0 ) ins

step1 (p,h) (a,n)
  | Just v <- dir a = ( p + v*n , h             )
  | Just v <- rot a = ( p       , h * v**(n/90) )
  | otherwise       = ( p + h*n , h             )

part2 ins = abs (round x) + abs (round y)
  where
    x :+ y = fst $ foldl' step2 ( 0 :+ 0 , 10 :+ 1 ) ins

step2 (p,w) (a,n)
  | Just v <- dir a = ( p       , w + v*n       )
  | otherwise       = step1 (p,w) (a,n)


import Numeric.LinearAlgebra
import Control.Monad
import System.Random

makeNetwork :: [Int] -> IO [Matrix R]
makeNetwork layersDesc = zipWithM randn (tail layersDesc) layersDesc

learnExample :: R -> [Matrix R] -> (Vector R, Vector R) -> [Matrix R]
learnExample learingRate network (input, output) =
    backpropagate learingRate output (propagate input network) network

train :: R -> [Matrix R] -> Int -> [(Vector R, Vector R)] -> [Matrix R]
train learingRate network epochs trainingData =
    foldl (learnExample learingRate) network (concat $ replicate epochs trainingData)

propagate :: Vector R -> [Matrix R] -> [Vector R]
propagate = scanl (\out w -> cmap (\x -> 1 / (1 + exp (-x))) $ w #> out)

backpropagate :: R -> Vector R -> [Vector R] -> [Matrix R] -> [Matrix R]
backpropagate learingRate targetOutput layerOutputs network = reverse revNewWeights
    where
        out = last layerOutputs
        outputDelta = (out - targetOutput)*out*(scalar 1 - out)
        revNetwork = reverse network
        revOutputs = reverse layerOutputs
        revDeltas = outputDelta : zipWith3 (\o w nextDelta -> (nextDelta <# w)*o*(scalar 1 - o))
                                           (tail revOutputs) revNetwork revDeltas
        revNewWeights = zipWith3 (\o d w -> w - scale learingRate (outer d o)) (tail revOutputs) revDeltas revNetwork

main :: IO ()
main = do
    network <- makeNetwork [3,4,1] -- we need a bias neuron
    xs <- replicateM samples (randomRIO (-val, val))
    ys <- replicateM samples (randomRIO (-val, val))
    let trainInput = zipWith (\x y -> vector [x,y,1]) xs ys
    let trainOutput = map targetFunction2 trainInput
    let trainedNetwork = train 0.03 network 100 (zip trainInput trainOutput)
    mapM_ (\y -> do
        mapM_ (\x ->
            putChar $ let out = last $ propagate (vector [x,y,1]) trainedNetwork in if (out ! 0) > 0.5 then '#' else '.')
            --putStr $ let out = last $ propagate (vector [x,y,1]) trainedNetwork in show (out ! 0) ++ " ")
            [-val, -val + 0.1 .. val]
        putChar '\n')
        [-val, -val + 0.1 .. val]
    where
        targetFunction v = if sqrt ((v ! 0)^(2::Int) + (v ! 1)^(2::Int)) <= r then vector [0] else vector [1]
        targetFunction2 v = if (v ! 0) >= 0 then vector [0] else vector [1]
        r = 1.5
        samples = 1000
        val = 2


-- ch03 example 3.9

import Statistics.Distribution
import Statistics.Distribution.Binomial

b :: BinomialDistribution
b = binomial 20 0.05

main = do
  putStrLn $ "P(Y>=4) is given by 1-P(Y<=3): " ++ show (1 - cumulative b 3.00)
  putStrLn $ "P(Y>=4) is given by 1-P(Y<=3): " ++ show (complCumulative b 3.00)
  putStrLn $ "P(Y>=4) is given by 1-P(Y<=3): "
          ++ show (1 - sum (map (probability b) [0 .. 3]))


module Numeric.Chebyshev.Internal
    ( -- * Interpolation on Chebyshev points
      Chebvals
    , chebpoints
    , sample
    , interpolate

    -- * Sum of Chebyshev polynomials
    , Chebpoly
    , integral

    -- * Conversion
    , valsToPoly
    , polyToVals

    -- * Testing
    , prop_interpolate_on
    , prop_inverse
    , prop_chebyshev
    )
  where

import Data.Maybe ( fromMaybe )
import Data.Vector.Unboxed ( Vector, (!), Unbox )
import Data.Complex ( Complex((:+)), realPart )

import Numeric.Fourier

import qualified Data.Vector.Unboxed as V

{-----------------------------------------------------------------------------
    Chebyshev interpolation
------------------------------------------------------------------------------}
-- | A polynomial, given by its values at Chebyshev points.
--
-- The value @(Chebvals fs)@ represents the unique polynomial
-- $f(x)$ of order $n$ that satisfies
--
-- >    f (chebpoints n !! j) = fs ! j  for all j = [0..n]
newtype Chebvals = Chebvals { getChebvals :: Vector Double }
    deriving (Eq, Show)

-- | Chebyshev points of order $n$.
chebpoints :: Int -> Vector Double
chebpoints n = getChebvals $ sample n id

-- | Sample a function at the Chebyshev points
sample :: Int -> (Double -> Double) -> Chebvals
sample n f = Chebvals $ V.generate (n+1) $ \j ->
    f $ cos (pi * fromIntegral j / fromIntegral n)

-- | Evaluate a 'Chebvals' at any point in the interval $[-1,1]$
-- through interpolation.
interpolate :: Chebvals -> Double -> Double
interpolate (Chebvals fs) = \x ->
    let ys = barycenters x
    in  correctNaN x $ sumMap (\j y -> fs!j * y) ys / V.sum ys
  where
    n  = V.length fs - 1
    xs = chebpoints n
    sign j = if even j then 1 else (-1)
    weight j
        | j == 0 = 0.5
        | j == n = 0.5 * sign j
        | otherwise = sign j
    barycenters x = V.generate (n+1) $ \j -> weight j / (x - xs!j)
    correctNaN x result
        | isNaN result = fs ! (fromMaybe 0 $ binarySearchReverse x xs)
        | otherwise    = result

-- | Binary search in a vector that is sorted from /largest/ to /smallest/.
binarySearchReverse :: (Ord a, Unbox a) => a -> Vector a -> Maybe Int
binarySearchReverse x xs = go 0 (n-1)
  where
    n = V.length xs
    go a b
        | b - a <= 1 =
            if x == xs!a then Just a
            else if x == xs!b then Just b
            else Nothing
        | otherwise  = case x `compare` (xs!m) of
            LT -> go m b
            GT -> go a m
            EQ -> Just m
      where
        m = a + (b-a) `div` 2

{-----------------------------------------------------------------------------
    Chebyshev polynomials
------------------------------------------------------------------------------}
-- | A polynomial, given as a linear combination of Chebyshev polynomials.
--
-- The $k$-the Chebyshev polynomial in the variable $x \in [-1,1]$
-- is defined as
--
-- $T^k(x) := \frac12 (z^k + z^{-k})$
--
-- where $z = e^{iθ}$ for $x = \cos θ$. This function is, in fact,
-- a polynomial in $x$.
newtype Chebpoly = Chebpoly (Vector Double)
    deriving (Eq, Show)

-- | Definite integral of a polynomial over the interval [-1,1].
integral :: Chebpoly -> Double
integral (Chebpoly cs) = sumMap (\j c -> c * integral j) cs
  where
    integral k = if odd k then 0 else 2/(1-fromIntegral k^2)

{-----------------------------------------------------------------------------
    Conversion
------------------------------------------------------------------------------}
toComplex :: Vector Double -> Vector (Complex Double)
toComplex = V.map (:+ 0)

toReal :: Vector (Complex Double) -> Vector Double
toReal = V.map realPart

-- | Expansion of a polynomial in terms of
-- Chebyshev polynomials, using the fast Fourier transform (FFT).
valsToPoly :: Chebvals -> Chebpoly
valsToPoly (Chebvals fs) =
    Chebpoly $ toReal $ uncircle $ fft $ toComplex $ circle fs
  where
    n = V.length fs - 1
    m = 2*n
    circle v = V.generate m $ \j ->
        if j <= n then v!j else v!(m-j)
    uncircle v = V.generate (n+1) $ \k -> (1/fromIntegral m)
        * (if k == 0 || k == n then v!k else v!k + v!(m-k))

-- | Evaluation of a sum of Chebyshev polynomials
-- using the fast Fourier transform (FFT).
polyToVals :: Chebpoly -> Chebvals
polyToVals (Chebpoly as) =
    Chebvals $ toReal $ uncircle $ ifft $ circle $ toComplex as
  where
    n = V.length as - 1
    m = 2*n
    circle v = V.generate m $ \k -> fromIntegral m *
        (if k == 0 || k == n
            then v!k
            else 0.5 * (if k < n then v!k else v!(m-k)))
    uncircle v = V.generate (n+1) $ \k ->
        if k == 0 || k == n then v!k else 0.5*(v!k + v!(m-k))

{- Note [ChebyshevFourier]

Expansion in terms of Chebyshev polynomials in terms of the Fourier
transform.

Assume that f(x) is a sum of Chebyshev polynomials of degree <= N

    f(x) = \sum_{k = 0…N} a_k 1/2·(z^k + z^{ -k})

Let M = 2N. Then, the values of this polynomial at the Chebyshev points

    x_j = cos(2πi/M j)  j = 0 … M-1
        = x_{M - j}

are given by the Fourier sum

    f(x_j) = \frac1M \sum_{k = 0 … M-1} c_k exp(2πi/M·jk)

where

    a_0 = 1/M·c_0
    a_k = 1/2·1/M·c_k = 1/2·1/M·c_{M-k}
    a_N = 1/M·c_N

-}

{-----------------------------------------------------------------------------
    Testing
------------------------------------------------------------------------------}
-- | 'interpolate' returns exact values when evaluated on Chebyshev points.
prop_interpolate_on :: Int -> (Double -> Double) -> Bool
prop_interpolate_on n f =
    V.map (interpolate (sample n f)) xs == V.map f xs
  where
    xs = chebpoints n

-- | 'polyToVals' and 'valsToPoly' are inverses of each other.
prop_inverse :: Chebvals -> Bool
prop_inverse xs = vec xs `approx` vec (polyToVals (valsToPoly xs))
  where vec (Chebvals v) = v

-- | Evaluating a Chebyshev polynomial using `polyToVals'
-- gives the same result as using the recursion formula.
prop_chebyshev :: Int -> Bool
prop_chebyshev k =
    vec (polyToVals fun) `approx` vec (sample n $ chebyshev k)
  where
    vec (Chebvals v) = v
    m = ceiling $ log (fromIntegral k) / log 2
    n = 2^m
    fun = Chebpoly $ V.generate (n+1) $ \j -> if j == k then 1 else 0

-- | Evaluate the $k$-the Chebyshev polynomial at argument @x@
-- using the recursive formula
--
-- \[\begin{aligned}
-- T_0(x)     &= 1 \\
-- T_1(x)     &= x \\
-- T_{n}(x) &= 2xT_{n-1}(x) - T_{n-2}(x) \\
-- \end{aligned}
-- \]
chebyshev :: Int -> Double -> Double
chebyshev k x = chebs !! k
  where
    chebs = 1 : x : zipWith (-) (map (2*x*) $ tail chebs) chebs

{-----------------------------------------------------------------------------
    Utilities
------------------------------------------------------------------------------}
-- | Efficient implementation of
--
-- >  sumMap f = V.sum (V.imap f)
sumMap :: (Int -> Double -> Double) -> Vector Double -> Double
sumMap f = V.ifoldl' (\total j a -> total + f j a) 0
{-# INLINE sumMap #-}

-- | L1 distance between two vectors.
dist :: Vector Double -> Vector Double -> Double
dist xs ys = V.sum $ V.zipWith (\x y -> abs (x-y)) xs ys

-- | Approximate equality up to machine precision
approx xs ys = dist xs ys / dist xs (V.map (const 0) ys) < eps
  where eps = 1e-15


module PlotUtils where

import Data.Complex ( Complex((:+)) )
import Graphics.GD ( newImage, rgb, savePngFile, setPixel, Color, Point, Size )
import Fractals ( iteratingFractal, Formula , flipForJulia)

type ViewPoint = (Float, Float) -- note that while gd's Point will refer to pixel points, ViewPoint will refer to points in the view port, on the complex plane
type ViewPort = (ViewPoint, ViewPoint) -- the exact window frame on the complex plane we want to observe
type JuliaMode = Bool

gridBy :: Size -> [(Int, Int)]
gridBy (width, height) = [(x, y) | x <- [0 .. width - 1], y <- [0 .. height - 1]]

colorBy :: Int -> Color
colorBy x = rgb (x*3+60) (x*2+30) (x^2+ 60) -- change by color ratio wanted

toColor :: ViewPoint -> Formula -> JuliaMode ->  Color
toColor (vx, vy) formula juliaMode = colorBy $ iteratingFractal formula zc 0
  where zc = flipForJulia (0 :+ 0, vx :+ vy) juliaMode

-- pointing from the outside world to a specific point on the complex plane outputing its x,y as a ViewPoint
pixelToViewPoint :: Point -> Size -> ViewPort -> ViewPoint
pixelToViewPoint (pixX, pixY) (width, height) ((minXvp, minYvp), (maxXvp, maxYvp)) =
  (xVP, yVP) where

    xscale = fromIntegral pixX / fromIntegral width
    yscale = fromIntegral pixY / fromIntegral height

    xVP = (maxXvp - minXvp) * xscale + minXvp
    yVP = (maxYvp - minYvp) * yscale + minYvp

plot :: ViewPort -> Size -> Formula -> JuliaMode -> String -> IO ()
plot viewPort size formula juliaMode fileName = do
  img <- newImage size

  mapM_ (\point ->
          let
            viewPoint = pixelToViewPoint point size viewPort
            color = toColor viewPoint formula juliaMode
          in
            setPixel point color img
        ) (gridBy size)

  savePngFile fileName img


module LeastSquares where

import Numeric.LinearAlgebra

n = 20
x = linspace n (-3,7::Double)
y0 = 3 * x

main = do
        noise <- randn 1 n
        let y = (flatten noise) + y0
        let sampleMatrix = (asColumn x) ||| (konst 1 (n,1))
        let sol = (sampleMatrix <\> y)
        print $ "Best fit is y = " ++ show (sol ! 0) ++ " * x + " ++ (show (sol ! 1))


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE GADTs #-}

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

{-|
Module      : Numeric.LinearProgramming.LP.Static
Description : Solve linear programs using static vectors/matrices.
Copyright   : (c) Ryan Orendorff, 2020
License     : BSD3
Stability   : experimental
-}
module Numeric.LinearProgramming.LP.Static (
  affineScaling
)
where

import           GHC.TypeLits
import           Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra.Static as S
import qualified Numeric.LinearAlgebra as LA
import Numeric.UnitaryRational

-- | https://en.wikipedia.org/wiki/Karmarkar%27s_algorithm
affineScaling :: (KnownNat m, KnownNat n, KnownNat p, KnownNat q) => L m n -> R m -> R n -> R n -> UnitaryRational p q -> [Maybe (R n)]
affineScaling a b c x0 gamma = iterate (>>= go) (Just x0)
    where
        gamma' = unitaryToFractional gamma

        min_vh :: (KnownNat n) => R n -> R n -> Double
        min_vh v h = minimum . map fst . filter ((<0) . snd) $ vals_with_h
            where
                vals = zipWithVector (\vi hi -> -vi/hi) v h
                vals_with_h = zip (LA.toList $ extract vals) (LA.toList $ extract h)

        go xk = let
            vk = b - a #> xk
            dm_negsq = diag $ dvmap (**(-2)) vk
            hx = inv ((tr a) S.<> dm_negsq S.<> a) #> c
            hv = ((konst (-1)) * a #> hx)
            alpha = gamma' * min_vh vk hv
            xkp1 = xk + (konst alpha) * hx
            in
                if ((>=0) . maximum . LA.toList . extract $ hv)
                then Nothing
                else Just xkp1

-- Example
-- -------
--
--
-- p = vector [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0] :: R 11
--
-- -- min c^Tx
-- c = vec2 1 1
--
-- -- s.t Ax <= b
-- a = ((konst 2) * (col p)) ||| (1 :: L 11 1)
-- b = p**2 + 1
--
-- lpSol = affineScaling a b c (konst 0) (UnitaryRational :: UnitaryRational 50 100)
--
-- mapM_ (disp 10 . fromJust) . take 20 $ lpSol