{-# 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.)


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


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 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


{-# LANGUAGE FlexibleContexts #-}

module TheField where

import           Data.Complex
import           Data.Char
import           Data.List as L
import           Graphics.Rendering.Chart.Easy hiding (transform)
import           Graphics.Rendering.Chart.Backend.Diagrams
--
import           TheField.Image
import           TheField.Plot

i :: Num a => a -> Complex a
i n = 0 :+ n

e :: Floating a => Complex a
e = exp 1 :+ 0

s :: [Complex Double]
s = [ 2 + i 2
    , 3 + i 2
    , 1.75 + i 1
    , 2 + i 1
    , 2.25 + i 1
    , 2.5 + i 1
    , 2.75 + i 1
    , 3 + i 1
    , 3.25 + i 1
    ]

-- Task 1.4.1
plotComplex :: EC (Layout Double Double) ()
plotComplex = plot' pts 4 4
  where pts = s

-- Task 1.4.3
plotTranslate :: EC (Layout Double Double) ()
plotTranslate = plot' pts 5 5
  where pts = [ pt + (1 + i 2) | pt <- s ]

-- Task 1.4.4
plotEyeCentral :: EC (Layout Double Double) ()
plotEyeCentral = plot' pts 4 4
  where pts = [ - 1 * (2 :+ 2) + pt | pt <- s ]

-- Task 1.4.7
plotScaled :: EC (Layout Double Double) ()
plotScaled = plot' pts 4 4
  where pts = [ 0.5 * pt | pt <- s ]

-- Task 1.4.8
plotRotation :: EC (Layout Double Double) ()
plotRotation = plot' pts 4 4
  where pts = [ i 0.5 * pt | pt <- s ]

-- Task 1.4.9
plotTranslatedRotation :: EC (Layout Double Double) ()
plotTranslatedRotation = plot' pts 4 4
  where pts = [ (i 0.5 * pt) + i (-1) + 2 | pt <- s ]

-- Task 1.4.10
plotImage :: [(Double, Double)] -> EC (Layout Double Double) ()
plotImage img = plot' pts 200 200
  where pts = [ x :+ y | (x, y) <- img ]

examplePng :: IO [(Double, Double)]
examplePng = do
  withIntensity 120 . pixels <$> readImage "./src/TheField/img01.png"
  where
    withIntensity ins pts =
      [ (fromIntegral x, fromIntegral y) | (x, y, pxl) <- pts, intensity pxl > ins ]

-- Task 1.4.17
plotE :: EC (Layout Double Double) ()
plotE = plot' pts 2 2
  where pts = [ e ** ((x * pi :+ 0) * i 2 / (n :+ 0)) | x <- [0..n-1] ]
        n = 20

-- Task 1.4.18
plotRotation' :: EC (Layout Double Double) ()
plotRotation' = plot' pts 4 4
  where pts = [ pt * e ** i (pi / 4) | pt <- s ]

-- Task 1.4.19
plotImageRotation :: [(Double, Double)] -> EC (Layout Double Double) ()
plotImageRotation img = plot' pts 200 200
  where
    pts = [ rotate $ x :+ y | (x, y) <- img ]
    rotate pt = pt * e ** i (pi / 4)

-- Task 1.4.20
plotMultipleOperations :: [(Double, Double)] -> EC (Layout Double Double) ()
plotMultipleOperations img = plot' pts 200 200
  where
    pts = [ scale . rotate . transform $ x :+ y | (x, y) <- img ]
    rotate = (*) (e ** i (pi / 4))
    scale = (*) 0.5
    transform = (+) (-1 * (100 :+ 100))

-- Task 1.5.1
message :: [String]
message =
  [ "10101"
  , "00100"
  , "10101"
  , "01011"
  , "11001"
  , "00011"
  , "01011"
  , "10101"
  , "00100"
  , "11001"
  , "11010"
  ]

decryptMessage :: String
decryptMessage = decryptCyphertext [1,0,0,0,1] message

decryptCyphertext :: [Int] -> [String] -> String
decryptCyphertext key = L.foldr ((:) . decode) ""
  where
    decode = decodeChar . decodeNumber . L.map (uncurry decryptBit) . zip key . L.map toBit
    toBit b = if b == '1' then 1 else 0

-- try all possible keys in attempt to find the answer
brute :: [([Int], String)]
brute = L.map (\k -> (k, decryptCyphertext k message)) keys
  where
    keys = [ [a,b,c,d,e] | a <- [0..1], b <- [0..1], c <- [0..1], d <- [0..1], e <- [0..1] ]

decryptBit :: Int -> Int -> Int
decryptBit k v = (k + v) `mod` 2

decodeNumber :: [Int] -> Int
decodeNumber str = L.foldr ((+) . uncurry decode) 0 $ zip [0..] (reverse str)
  where decode idx val = val * 2 ^ idx

decodeChar :: Int -> Char
decodeChar n =
  case n' of
    26 -> ' '
    _ -> chr $ n' + 65
  where n' = n `mod` 27

-- Task 1.7.10
sumComplexNumbersA :: EC (Layout Double Double) ()
sumComplexNumbersA = plotComplexAddition (3.0 :+ 1.0) (2.0 :+ 2.0)

sumComplexNumbersB :: EC (Layout Double Double) ()
sumComplexNumbersB = plotComplexAddition ((-1.0) :+ 2.0) (1.0 :+ (-1.0))

sumComplexNumbersC :: EC (Layout Double Double) ()
sumComplexNumbersC = plotComplexAddition (2.0 :+ 0.0) ((-3.0) :+ 0.001)

sumComplexNumbersD :: EC (Layout Double Double) ()
sumComplexNumbersD = plotComplexAddition (4 * (0.0 :+ 2.0)) (0.001 :+ 1.0)

generateImages :: IO ()
generateImages = do
  toFile def "./src/TheField/task-1.4.1.svg" plotComplex
  toFile def "./src/TheField/task-1.4.3.svg" plotTranslate
  toFile def "./src/TheField/task-1.4.4.svg" plotEyeCentral
  toFile def "./src/TheField/task-1.4.7.svg" plotScaled
  toFile def "./src/TheField/task-1.4.8.svg" plotRotation
  toFile def "./src/TheField/task-1.4.9.svg" plotTranslatedRotation
  toFile def "./src/TheField/task-1.4.10.svg" =<< plotImage <$> examplePng
  toFile def "./src/TheField/task-1.4.17.svg" plotE
  toFile def "./src/TheField/task-1.4.18.svg" plotRotation'
  toFile def "./src/TheField/task-1.4.19.svg" =<< plotImageRotation <$> examplePng
  toFile def "./src/TheField/task-1.4.20.svg" =<< plotMultipleOperations <$> examplePng
  toFile def "./src/TheField/task-1.7.10-a.svg" sumComplexNumbersA
  toFile def "./src/TheField/task-1.7.10-b.svg" sumComplexNumbersB
  toFile def "./src/TheField/task-1.7.10-c.svg" sumComplexNumbersC
  toFile def "./src/TheField/task-1.7.10-d.svg" sumComplexNumbersD


module Math.AlgebraicNumbers.Overloads where
import Math.AlgebraicNumbers.Prelude
import Math.AlgebraicNumbers.Class
import Math.AlgebraicNumbers.AReal
import Math.AlgebraicNumbers.AComplex
import qualified Prelude
import qualified Data.Complex

-- Prelude functions generalized from Integral a to GCDDomain a / EuclideanDomain a

gcd :: (GCDDomain a) => a -> a -> a
gcd = gcdD

lcm :: (GCDDomain a) => a -> a -> a
lcm = lcmD

divMod :: (EuclideanDomain a) => a -> a -> (a, a)
divMod = divModD

infixl 7 `div`, `mod`

div, mod :: (EuclideanDomain a) => a -> a -> a
div = divD
mod = modD

class Sqrt a where
  sqrt :: a -> a
  maybeSqrt :: a -> Maybe a

instance Sqrt Prelude.Float where
  sqrt = Prelude.sqrt
  maybeSqrt x | x >= 0 = Just $ Prelude.sqrt x
              | otherwise = Nothing

instance Sqrt Prelude.Double where
  sqrt = Prelude.sqrt
  maybeSqrt x | x >= 0 = Just $ Prelude.sqrt x
              | otherwise = Nothing

instance (RealFloat a) => Sqrt (Data.Complex.Complex a) where
  sqrt = Prelude.sqrt
  maybeSqrt = Just . Prelude.sqrt

instance Sqrt AReal where
  sqrt = sqrtA
  maybeSqrt x | x >= 0 = Just $ sqrtA x
              | otherwise = Nothing

instance Sqrt AComplex where
  sqrt = sqrtAN
  maybeSqrt = Just . sqrtAN

-- toRational?


{-# LANGUAGE MultiParamTypeClasses #-}

module Stats where

import           Control.Applicative
import           Control.Monad.State
import           Data.Random
import           Data.Random.Distribution.Exponential
import           Data.Random.Source.PureMT
import qualified Data.Vector                          as V
import           Data.Word
import qualified Statistics.Sample.Histogram          as S

seed :: Word64
seed = 123

sampleStream :: (Num t, Distribution d t) => d t -> [t]
sampleStream distribution = evalState (iterateM $ sample distribution) (pureMT seed)

iterateM :: (Applicative m) => m a -> m [a]
iterateM f = liftA2 (:) f (iterateM f)

sampleVector :: (Num t, Distribution d t) => d t -> Int -> V.Vector t
sampleVector distribution size = V.fromList . take size $ sampleStream distribution

sampleExp :: Int -> V.Vector Double
sampleExp = sampleVector (Exp 1)

sampleExp' :: Int -> V.Vector Double
sampleExp' = fmap (negate . log) . sampleVector StdUniform

expHistograms :: (V.Vector Double, V.Vector Double)
expHistograms = S.histogram 30 (sampleExp 1000)

expHistograms' :: (V.Vector Double, V.Vector Double)
expHistograms' = S.histogram 30 (sampleExp' 1000)

mean :: Floating a => V.Vector a -> a
mean xs = V.sum xs / fromIntegral (V.length xs)

var :: Floating a => V.Vector a -> a
var xs = mean (fmap (** 2) xs) - mean xs ** 2

isCloseWithTolerance :: (Ord a, Floating a) => a -> a -> a -> a -> Bool
isCloseWithTolerance aTol rTol x y = abs (x - y) <= aTol + rTol * abs y

isClose :: (Ord a, Floating a) => a -> a -> Bool
isClose = isCloseWithTolerance 1e-8 1e5


{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE Rank2Types            #-}
{-# LANGUAGE ViewPatterns          #-}
{-# OPTIONS_GHC -fno-warn-missing-methods #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Diagrams.TwoD.Apollonian
-- Copyright   :  (c) 2011, 2016 Brent Yorgey
-- License     :  BSD-style (see LICENSE)
-- Maintainer  :  byorgey@cis.upenn.edu
--
-- Generation of Apollonian gaskets.  Any three mutually tangent
-- circles uniquely determine exactly two others which are mutually
-- tangent to all three.  This process can be repeated, generating a
-- fractal circle packing.
--
-- See J. Lagarias, C. Mallows, and A. Wilks, \"Beyond the Descartes
-- circle theorem\", /Amer. Math. Monthly/ 109 (2002), 338--361.
-- <http://arxiv.org/abs/math/0101066>.
--
-- A few examples:
--
-- > import Diagrams.TwoD.Apollonian
-- > apollonian1 = apollonianGasket 0.01 2 2 2
--
-- <<diagrams/src_Diagrams_TwoD_Apollonian_apollonian1.svg#diagram=apollonian1&width=400>>
--
-- > import Diagrams.TwoD.Apollonian
-- > apollonian2 = apollonianGasket 0.01 2 3 3
--
-- <<diagrams/src_Diagrams_TwoD_Apollonian_apollonian2.svg#diagram=apollonian2&width=400>>
--
-- > import Diagrams.TwoD.Apollonian
-- > apollonian3 = apollonianGasket 0.01 2 4 7
--
-- <<diagrams/src_Diagrams_TwoD_Apollonian_apollonian3.svg#diagram=apollonian3&width=400>>
--
-----------------------------------------------------------------------------

module Diagrams.TwoD.Apollonian
       ( -- * Circles

         Circle(..), mkCircle, center, radius

         -- * Descartes' Theorem
       , descartes, other, initialConfig

         -- * Apollonian gasket generation

       , apollonian

         -- ** Kissing sets

       , KissingSet(..), kissingSets, flipSelected, selectOthers

         -- ** Apollonian trees

       , apollonianTrees, apollonianTree

         -- * Diagram generation

       , drawCircle
       , drawGasket
       , apollonianGasket

       ) where

import           Data.Complex
import qualified Data.Foldable    as F
import           Data.Maybe       (catMaybes)
import           Data.Tree

import           Diagrams.Prelude hiding (center, radius)

import           Control.Arrow    (second, (&&&))

------------------------------------------------------------
--  Circles
------------------------------------------------------------

-- | Representation for circles that lets us quickly compute an
--   Apollonian gasket.
data Circle n = Circle
  { bend :: n
    -- ^ The bend is the reciprocal of signed
    --   radius: a negative radius means the
    --   outside and inside of the circle are
    --   switched.  The bends of any four mutually
    --   tangent circles satisfy Descartes'
    --   Theorem.
  , cb   :: Complex n
    -- ^ /Product/ of bend and center represented
    --   as a complex number.  Amazingly, these
    --   products also satisfy the equation of
    --   Descartes' Theorem.
  }
  deriving (Eq, Show)

-- | Create a @Circle@ given a signed radius and a location for its center.
mkCircle :: Fractional n =>
            n -- ^ signed radius
         -> P2 n     -- ^ center
         -> Circle n
mkCircle r (unp2 -> (x,y)) = Circle (1/r) (b*x :+ b*y)
  where b = 1/r

-- | Get the center of a circle.
center :: Fractional n => Circle n -> P2 n
center (Circle b (cbx :+ cby)) = p2 (cbx / b, cby / b)

-- | Get the (unsigned) radius of a circle.
radius :: Fractional n => Circle n -> n
radius = abs . recip . bend

liftF :: RealFloat n => (forall a. Floating a => a -> a) -> Circle n -> Circle n
liftF f (Circle b c) = Circle (f b) (f c)

liftF2 :: RealFloat n => (forall a. Floating a => a -> a -> a) ->
          Circle n -> Circle n -> Circle n
liftF2 f (Circle b1 cb1) (Circle b2 cb2) = Circle (f b1 b2) (f cb1 cb2)

instance RealFloat n => Num (Circle n) where
  (+)           = liftF2 (+)
  (-)           = liftF2 (-)
  (*)           = liftF2 (*)
  negate        = liftF negate
  abs           = liftF abs
  fromInteger n = Circle (fromInteger n) (fromInteger n)

instance RealFloat n => Fractional (Circle n) where
  (/)   = liftF2 (/)
  recip = liftF recip

-- | The @Num@, @Fractional@, and @Floating@ instances for @Circle@
--   (all simply lifted elementwise over @Circle@'s fields) let us use
--   Descartes' Theorem directly on circles.
instance RealFloat n => Floating (Circle n) where
  sqrt = liftF sqrt

------------------------------------------------------------
--  Descartes' Theorem
------------------------------------------------------------

-- XXX generalize these for higher dimensions?

-- | Descartes' Theorem states that if @b1@, @b2@, @b3@ and @b4@ are
--   the bends of four mutually tangent circles, then
--
--   @
--     b1^2 + b2^2 + b3^2 + b4^2 = 1/2 * (b1 + b2 + b3 + b4)^2.
--   @
--
--   Surprisingly, if we replace each of the @bi@ with the /product/
--   of @bi@ and the center of the corresponding circle (represented
--   as a complex number), the equation continues to hold! (See the
--   paper referenced at the top of the module.)
--
--   @descartes [b1,b2,b3]@ solves for @b4@, returning both solutions.
--   Notably, @descartes@ works for any instance of @Floating@, which
--   includes both @Double@ (for bends), @Complex Double@ (for
--   bend/center product), and @Circle@ (for both at once).
descartes :: Floating n => [n] -> [n]
descartes [b1,b2,b3] = [r + s, -r + s]
  where r = 2 * sqrt (b1*b2 + b1*b3 + b2*b3)
        s = b1+b2+b3
descartes _ = error "descartes must be called on a list of length 3"

-- | If we have /four/ mutually tangent circles we can choose one of
--   them to replace; the remaining three determine exactly one other
--   circle which is mutually tangent.  However, in this situation
--   there is no need to apply 'descartes' again, since the two
--   solutions @b4@ and @b4'@ satisfy
--
--   @
--     b4 + b4' = 2 * (b1 + b2 + b3)
--   @
--
--   Hence, to replace @b4@ with its dual, we need only sum the other
--   three, multiply by two, and subtract @b4@.  Again, this works for
--   bends as well as bend/center products.
other :: Num n => [n] -> n -> n
other xs x = 2 * sum xs - x

-- | Generate an initial configuration of four mutually tangent
--   circles, given just the signed bends of three of them.
initialConfig :: RealFloat n => n -> n -> n -> [Circle n]
initialConfig b1 b2 b3 = cs ++ [c4]
  where cs     = [Circle b1 0, Circle b2 ((b2/b1 + 1) :+ 0), Circle b3 cb3]
        a      = 1/b1 + 1/b2
        b      = 1/b1 + 1/b3
        c      = 1/b2 + 1/b3
        x      = (b*b + a*a - c*c)/(2*a)
        y      = sqrt (b*b - x*x)
        cb3    = b3*x :+ b3*y
        [c4,_] = descartes cs

------------------------------------------------------------
--  Gasket generation
------------------------------------------------------------

select :: [a] -> [(a, [a])]
select [] = []
select (x:xs) = (x,xs) : (map . second) (x:) (select xs)

-- | The basic idea of a kissing set is supposed to represent a set of
--   four mutually tangent circles with one selected, though in fact
--   it is more general than that: it represents any set of objects
--   with one distinguished object selected.
data KissingSet n = KS { selected :: n, others :: [n] }
  deriving (Show)

-- | Generate all possible kissing sets from a set of objects by
--   selecting each object in turn.
kissingSets :: [n] -> [KissingSet n]
kissingSets = map (uncurry KS) . select

-- | \"Flip\" the selected circle to the 'other' circle mutually tangent
--   to the other three.  The new circle remains selected.
flipSelected :: Num n => KissingSet n -> KissingSet n
flipSelected (KS c cs) = KS (other cs c) cs

-- | Make the selected circle unselected, and select each of the
--   others, generating a new kissing set for each.
selectOthers :: KissingSet n -> [KissingSet n]
selectOthers (KS c cs) = [ KS c' (c:cs') | (c',cs') <- select cs ]

-- | Given a threshold radius and a list of /four/ mutually tangent
--   circles, generate the Apollonian gasket containing those circles.
--   Stop the recursion when encountering a circle with an (unsigned)
--   radius smaller than the threshold.
apollonian :: RealFloat n => n -> [Circle n] -> [Circle n]
apollonian thresh cs
  = (cs++)
  . concat
  . map (maybe [] flatten . prune p . fmap selected)
  . apollonianTrees
  $ cs
  where
    p c = radius c >= thresh

-- | Given a set of /four/ mutually tangent circles, generate the
--   infinite Apollonian tree rooted at the given set, represented as
--   a list of four subtrees.  Each node in the tree is a kissing set
--   with one circle selected which has just been flipped.  The three
--   children of a node represent the kissing sets obtained by
--   selecting each of the other three circles and flipping them.  The
--   initial roots of the four trees are chosen by selecting and
--   flipping each of the circles in the starting set. This
--   representation has the property that each circle in the
--   Apollonian gasket is the selected circle in exactly one node
--   (except that the initial four circles never appear as the
--   selected circle in any node).
apollonianTrees :: RealFloat n => [Circle n] -> [Tree (KissingSet (Circle n))]
apollonianTrees = map (apollonianTree . flipSelected) . kissingSets

-- | Generate a single Apollonian tree from a root kissing set.  See
--   the documentation for 'apollonianTrees' for an explanation.
apollonianTree :: RealFloat n => KissingSet (Circle n) -> Tree (KissingSet (Circle n))
apollonianTree = unfoldTree (id &&& (map flipSelected . selectOthers))

-- | Prune a tree at the shallowest points where the predicate is not
--   satisfied.
prune :: (a -> Bool) -> Tree a -> Maybe (Tree a)
prune p (Node a ts)
  | not (p a) = Nothing
  | otherwise = Just $ Node a (catMaybes (map (prune p) ts))

------------------------------------------------------------
--  Diagram generation
------------------------------------------------------------

-- | Draw a circle.
drawCircle :: (Renderable (Path V2 n) b, TypeableFloat n) =>
              Circle n -> QDiagram b V2 n Any
drawCircle c = circle (radius c) # moveTo (center c)
                                 # fcA transparent

-- | Draw a generated gasket, using a line width 0.003 times the
--   radius of the largest circle.
drawGasket :: (Renderable (Path V2 n) b, TypeableFloat n) =>
              [Circle n] -> QDiagram b V2 n Any
drawGasket cs = F.foldMap drawCircle cs

-- | Draw an Apollonian gasket: the first argument is the threshold;
--   the recursion will stop upon reaching circles with radii less than
--   it. The next three arguments are bends of three circles.
apollonianGasket :: (Renderable (Path V2 n) b, TypeableFloat n)
                 => n -> n -> n -> n -> QDiagram b V2 n Any
apollonianGasket thresh b1 b2 b3 = drawGasket . apollonian thresh $ (initialConfig b1 b2 b3)

------------------------------------------------------------
-- Some notes on floating-point error
-- (only for the intrepid)
------------------------------------------------------------
{-

-- code from Gerald Gutierrez, personal communication

module Main where

import           Data.Complex
import           Diagrams.Backend.SVG.CmdLine
import           Diagrams.Prelude

-- ------^---------^---------^---------^---------^---------^---------^--------

data Circle = Circle Double (Complex Double) deriving (Show)

descartes a b c
  = (s + r, s - r)
  where
    s = a + b + c
    r = 2 * sqrt ((a * b) + (b * c) + (c * a))

descartesDual a b c d
  = 2 * (a + b + c) - d

soddies (Circle k1 b1) (Circle k2 b2) (Circle k3 b3)
  = ( Circle k4 b4
    , Circle k5 b5 )
  where
    (k4, k5) = descartes k1 k2 k3
    (b4, b5) = descartes b1 b2 b3

soddiesDual (Circle k1 b1) (Circle k2 b2) (Circle k3 b3) (Circle k4 b4)
  = Circle (descartesDual k1 k2 k3 k4) (descartesDual b1 b2 b3 b4)

mutuallyTangentCirclesFromTriangle z1 z2 z3
  = ( Circle k1 (z1 * (k1 :+ 0))
    , Circle k2 (z2 * (k2 :+ 0))
    , Circle k3 (z3 * (k3 :+ 0)) )
  where
    a = magnitude (z2 - z3)
    b = magnitude (z3 - z1)
    c = magnitude (z1 - z2)
    s = (a + b + c) / 2
    k1 = 1 / (s - a)
    k2 = 1 / (s - b)
    k3 = 1 / (s - c)

main :: IO ()
main = mainWith picture

pic = mainWith picture

mkCircle :: Circle -> Diagram B
mkCircle (Circle k b)
  = circle (1 / k) # moveTo (p2 (realPart z, imagPart z))
  where
    z = b / (k :+ 0)

picture :: Diagram B
picture
  = mkCircle c1 <> mkCircle c2 <> mkCircle c3 <> mkCircle c4 <> mkCircle c5
  where
    (c1, c2, c3) = mutuallyTangentCirclesFromTriangle z1 z2 z3
    (c4, c5) = soddies c1 c2 c3

    -- z1 = 0 :+ 0
    -- z2 = 3 :+ 0
    -- z3 = 0 :+ 4

    -- z1 = (-0.546) :+ (-0.755)
    -- z2 = ( 0.341) :+ (-0.755)
    -- z3 = (-0.250) :+ ( 0.428)

ofsBad = 0.15397 -- doesn't work
ofsGood = 0.15398 -- works

ofs = ofsGood

z1 = ofs + ((-0.546) :+ (-0.755))
z2 = ofs + (( 0.341) :+ (-0.755))
z3 = ofs + ((-0.250) :+ ( 0.428))

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

Email to Gerald Gutierrez, 30 Sep 2016:

I got a chance to sit down and think hard about your code today, and
I think I have figured out what the problem is.

If you look at the outputs of 'soddies c1 c2 c3' using the two
different values for 'ofs', you will see that they are *almost* the
same, except that the complex numbers have been switched (though in
fact their real components are almost the same so you only notice the
difference in the imaginary components).  This clearly causes
incorrect results since the y-coordinates represented by the imaginary
components are now getting scaled by different bend values.  So the
resulting circles are of the right size and have (close to) the
correct x-coordinates, but their y-coordinates are wrong.

The problem is that in the 'soddies' function, you make independent
calls to 'descartes k1 k2 k3' (to solve for the bends) and 'descartes
b1 b2 b3' (to solve for the bend-center products), and then *assume*
that they return their solutions *in the same order*, so that you can
match up the first values to make one circle and the second values to
make another circle.  I would guess that this is *usually* true but
apparently not when certain values are hovering right around a branch
cut of sqrt, or something like that.

I think my code in Diagrams.TwoD.Apollonian suffers from the same
problem.  Ultimately I think it is due to the inherent inaccuracy of
floating-point numbers; there is nothing wrong with the formula
itself.  It would be nice to figure out how to correct for this, but I
am not sure how off the top of my head.  The interesting thing is that
switching the bend-center products does not seem to violate the
generalized Descartes' Theorem at all --- so it would seem that it
does not actually completely characterize mutually tangent circles,
that is, there exist sets of circles satisfying the theorem which are
not in fact mutually tangent.

-}


{-# LANGUAGE FlexibleContexts #-}

module SVM
  ( multiplyByItself
  , testmatrix
  , testvector
  , replicateVector
  , replicateVectorSquared
  , makeLabelAgreementMatrix
  , makeQuadraticMatrix
  , testmatrixsize
  , alphasQuadProg
  , trainPLAandSVM
  , countSV
  ) where

import           LinearRegression                (createMatrixX,
                                                  createRandomPoints,
                                                  createVectorY,
                                                  makeTargetFunction, y)
import           Numeric.LinearAlgebra
import           Numeric.LinearAlgebra.Data
import           Numeric.LinearAlgebra.HMatrix   as HMatrix
import           Numeric.Minimization.QuadProgPP
import           PLA

----------------------------------------------------
-- Solutions to Homework 7 on SVM
----------------------------------------------------
-- | a test matrix - not part of homework
testmatrix :: Matrix R
testmatrix = matrix 2 [1, 1, 1, -1, -1, -1, 3, 4, 5, 6]

testmatrixsize = size testmatrix

-- | a test vector - not part of homework
testvector :: Vector R
testvector = vector [-1, 1, 1, -1, 1]

-- | create a matrix where every entry x_n_m is the dot product of the nth vector with the mth vector
multiplyByItself :: Matrix R -> Matrix R
multiplyByItself matrixA = matrixA HMatrix.<> (tr matrixA)

-- | make a matrix where every entry corresponds to the product of the labels of the corresponding data - it will be 1 if they agree and -1 otherwise
makeLabelAgreementMatrix :: Vector R -> Matrix R
makeLabelAgreementMatrix labelvector =
  diag labelvector HMatrix.<> replicateVectorSquared labelvector

-- | make a matrix form a vector by repeating the vector n times and arranging into n rows, each corresponding to the vector
replicateVector :: Int -> Vector R -> Matrix R
replicateVector n v = fromRows $ replicate n v

-- | make a squared matrix from a vector, such that each rows corresponds to the vector
replicateVectorSquared :: Vector R -> Matrix R
replicateVectorSquared v = replicateVector (size v) v

-- | makes the matrix to pass to quadratix programming
makeQuadraticMatrix :: Matrix R -> Vector R -> Matrix R
makeQuadraticMatrix dataMatrix labelVector =
  multiplyByItself dataMatrix * makeLabelAgreementMatrix labelVector

makeQuadraticMatrix' :: Matrix R -> Vector R -> Matrix R
makeQuadraticMatrix' dataMatrix labelVector = quadMatrix + smallValtoDiag
  where
    quadMatrix = makeQuadraticMatrix dataMatrix labelVector
    smallValtoDiag = diag $ konst 0.00000000001 (size labelVector)

-- | solves quadratic Programming on the input data matrix and label Vector
alphasQuadProg ::
     Matrix R -> Vector R -> Either QuadProgPPError (Vector R, Double)
alphasQuadProg dataMatrix labelVector =
  solveQuadProg
    (matrixA, vectorB)
    (Just (labelVecMatrix, zeros))
    (Just (idMatrix, zeros))
  where
    matrixA = makeQuadraticMatrix' dataMatrix labelVector
    vectorB = konst (-1) (size labelVector)
    labelVecMatrix = diag labelVector
    idMatrix = ident (size labelVector)
    zeros = konst 0 (size labelVector)

-- | Creates the data matrix from a list of data points
createDataMatrix :: [(R, R)] -> Matrix R
createDataMatrix listOfPoints = matrix 2 listOfNumbers
  where
    listOfNumbers = concat [[a, b] | (a, b) <- listOfPoints]

setSmallTo0 x
  | x < (10 ** (-4)) = 0
  | otherwise = x

countSV vector = length . filter (/= 0) $ toList vector

-- | trains a perceptron and SVM
trainPLAandSVM :: Int -> IO ()
trainPLAandSVM n
    --making target function, training and testing points
    -----------------------------------------------------
 = do
  target <- makeTargetFunction
  trainpoints <- createRandomPoints n
  testpoints <- createRandomPoints 1000
  let trainX = createDataMatrix trainpoints
  let trainY = createVectorY (y target) trainpoints
  let testX = createDataMatrix testpoints
  let testY = createVectorY (y target) testpoints
    --Perceptron
    ---------------------------------
  let initialWeights = (0, 0, 0)
  (finalWeights, epochs) <-
    trainPLAWithInitial initialWeights trainpoints target
  let misclassifiedTestPoints =
        length $ filter (isMisclassified target finalWeights) testpoints
  let prob = fromIntegral misclassifiedTestPoints / 1000
    --SVM
    --------------------------
  let quadmatrix = makeQuadraticMatrix trainX trainY
  let labelVecMatrix = asColumn trainY
  let alphas = alphasQuadProg trainX trainY
  let alphasModified = fmap (cmap setSmallTo0 . fst) alphas
  let numSVs = fmap countSV alphasModified
    --Printing the output
    ---------------------------
  putStrLn "Perceptron:"
  putStrLn "-------------------------"
  putStr "Weights: "
  print finalWeights
  putStr "Epochs: "
  print epochs
  putStr "Probability f /= g: "
  print prob
  putStrLn "Support Vector Machine:"
  putStrLn "-------------------------"
  putStr "Alphas:"
  print alphas
  putStr "Alphas very small set ot 0:"
  print alphasModified
  putStr "number of SVs:"
  print numSVs


import Data.Complex

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

dft :: RealFloat a => [Complex a] -> [Complex a]
dft x = map (\(_, k) -> singleDFT k idxList) idxList
  where idxList = zip x [0..]

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

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


-- | This module provides generic functionality to deal with ensembles in
-- statistical mechanics.

module StatisticalMechanics.Ensemble where

import Numeric.Log

import Statistics.Probability

-- | The state probability functions provide conversion from some types @a@
-- into non-normalized probabilities. For "real" applications, using the
-- @logProbability@ function is preferred. This functions allows for easy
-- abstraction when types @a@ are given as fractions of some actual value (say:
-- deka-cal), or are discretized.
--
-- The returned values are not normalized, because we do not now the total
-- evidence @Z@ until integration over all states has happened -- which is not
-- feasible in a number of problems.
--
-- TODO replace @()@ with temperature and results with non-normalized @P@ or
-- @LogP@, depending. At some point we want to have type-level physical
-- quantities, hence the need for the second type.

class StateProbability a where
  -- | Given a temperature and a state "energy", return the corresponding
  -- non-normalized probability.
  stateProbability
    ∷ Double
    -- ^ this is @k*T@
    → a
    -- ^ the energy (or discretized energy)
    → Probability NotNormalized Double
    -- ^ probability of being in state @a@, but only proportional up to @1/Z@.
  stateLogProbability
    ∷ Double
    -- ^ this is @1/(k * T)@
    → a
    -- ^ the energy (or discretized energy)
    → Log (Probability NotNormalized Double)
    -- ^ resulting probability

instance StateProbability Double where
  stateProbability kT x = Prob . exp . negate $ x/kT
  {-# Inline stateProbability #-}
  --stateLogProbability kT x = Exp . log . Prob . exp . negate $ x/kT
  stateLogProbability kT x = Exp . Prob . negate $ x/kT
  {-# Inline stateLogProbability #-}


import Data.Array.CArray
import Data.Complex
import Math.FFT (dft, idft) -- Binding to fftw

type Vector = CArray Int (Complex Double)

calculateEnergy :: Double -> Vector -> Vector -> Vector -> Double
calculateEnergy dx kin pot wfc = (* dx) . sum . map realPart $ elems total
  where
    total = liftArray2 (+) kineticE potentialE
    potentialE = wfcConj .* pot .* wfc
    kineticE = wfcConj .* idft (kin .* dft wfc)
    wfcConj = liftArray conjugate wfc
    a .* b = liftArray2 (*) a b


module Hw5
    (Signal,
    Vec(..),
    rd,
    absolute,
    range,
    dft,
    idft,
    realV,
    low_pass',
    low_pass,
    fft,
    ifft,
    makeNoise)
    where

import Data.Complex
import Debug.Trace
import GHC.Base (liftA2)
import System.Random
-- import Control.Monad.Reader

-- {-# LANGUAGE  KindSignatures  #-}
-- {-# LANGUAGE  MultiParamTypeClasses  #-}

newtype Vec a = Vec {runVec :: [a]}
newtype  Reader r a = Reader { runReader  :: r -> a }
type Signal = Vec (Complex Double)

imagV :: Num a => Vec a -> Vec (Complex a)
realV :: Num a => Vec a -> Vec (Complex a)
range :: Double -> Double -> Double -> [Double]
absolute :: Vec (Complex Double) -> Vec Double
rd :: Int -> Vec Double -> Vec Double
twiddle :: Double -> Double -> Signal
dft :: Signal -> Signal
idft :: Signal -> Signal
mask :: Int -> Int -> Signal
fft :: Signal -> Signal
ifft :: Signal -> Signal
makeNoise :: Random a => Int -> Int -> a -> a -> Vec a

makeNoise seed n low high =
    let x = mkStdGen seed
        genNum l h = randomR (l, h)
        genList 1 l h g = [c] where (c,y) = genNum l h g
        genList n l h g = c : genList (n-1) l h y where (c,y) = genNum l h g
        in
            Vec $ genList n low high x

length' (Vec x) = length x

rd n = fmap (\x -> fromIntegral (round $ c * x) / c)
    where c = 10^n

absolute = fmap (\(r:+i) -> sqrt(r*r + i*i))

range from to count = map (\x -> from + x * step) [0..count-1]
    where step = (to - from)/count

imagV (Vec a) = Vec $ map (0:+) a
realV (Vec a) = Vec $ map (:+0) a

instance Show a => Show (Vec a) where
    show (Vec lst) = "[" ++ drop 1 lst' ++ "]"
        where
            lst' = mconcat $ map (\x -> " " ++ show x) lst

instance Functor Vec where
    fmap f (Vec x) = Vec $ map f x

instance Applicative Vec where
    pure = Vec . repeat
    (Vec f) <*> (Vec x) = Vec $ zipWith ($) f x
    liftA2 f (Vec x) (Vec y) = Vec $ zipWith f x y

instance Foldable Vec where
    foldr f c (Vec a) = Prelude.foldr f c a

instance Semigroup (Vec a) where
    Vec a <> Vec b = Vec $ a ++ b

instance Monoid (Vec a) where
    mempty = Vec []

instance Num a => Num (Vec a) where
    (+) = liftA2 (+)
    (-) = liftA2 (-)
    (*) = liftA2 (*)
    negate = fmap negate
    abs = fmap abs
    signum = fmap signum
    fromInteger x = pure $ fromInteger x

instance (Floating a) => Fractional (Vec a) where
    (/) = liftA2 (/)
    fromRational x = pure $ fromRational x

instance (Floating a) => Floating (Vec a) where
    pi = pure pi
    exp = fmap exp
    log = fmap log
    sin = fmap sin
    cos = fmap cos
    asin = fmap asin
    acos = fmap acos
    atan = fmap atan
    sinh = fmap sinh
    cosh = fmap cosh
    asinh = fmap asinh
    acosh = fmap acosh
    atanh = fmap atanh

mask freq n = realV $ Vec $ one ++ zero ++ one where
    one = (take freq $ repeat 1)
    zero = (take (n-freq*2) $ repeat 0)

low_pass' :: Int -> Signal -> Signal

low_pass' v sig =
    let
        m = mask v $ length sig
    in idft $ m * dft sig

twiddle n k =
    let a = [0.. (n - 1)]
        f x = exp (0:+ x * (-2*pi*k / n))
    in Vec $ fmap f a

dft x =
    let bigN = fromIntegral $ length x -- length of x
        n = Vec [0 .. bigN-1]
        f k = sum $ x * twiddle bigN k
    in fmap f n

idft x =
    let bigN = fromIntegral $ length x -- length of x
        cX = fmap conjugate x
        n = Vec [0 .. bigN-1]
        diC = (1.0 / bigN) :+ 0
        f k = sum $ cX * twiddle bigN k
    in conj diC $ fmap f n

conj :: Complex Double -> Signal -> Signal
conj b (Vec a) = fmap (\y -> b * conjugate y) (Vec a)

twidd n k = exp( imagV $ k * pure (-2 * pi / fromIntegral n) )
twiddI n k = exp( imagV $ k * pure (2 * pi / fromIntegral n) )

trace_fft m x = trace (m ++ ": " ++ (show $ rd 2 $ absolute x)) x -- debugging helper to print fft results
fft x
    | n <= 16 = dft x
    | n <= 1 = x
    | otherwise =
        let (even, odd) = split $ runVec x
            (e, o) = (fft $ Vec even, fft $ Vec odd)
            k = Vec [0..]
            p = twidd n k
        in
            (e + p * o) <> (e - p * o)

    where
        n = length x
        split [] = ([], [])
        split [a] = ([a], [])
        split (a:b:c) = (a:x, b:y)
            where (x,y) = split c

ifft x
    | n <= 16 = idft x
    | n <= 1 = x
    | otherwise =
        let (even, odd) = split $ runVec x
            (e, o) = (fft $ Vec even, fft $ Vec odd)
            -- k = Vec [0..]
            bigN = fromIntegral n
            cE = fmap conjugate e
            cO = fmap conjugate o
            diC = (-1.0 / bigN) :+ 0
            p = twiddI n $ Vec [0..]
        in
            conj diC (e + p * o) <> conj diC (e - p * o)

    where
        n = length x
        split [] = ([], [])
        split [a] = ([a], [])
        split (a:b:c) = (a:x, b:y)
            where (x,y) = split c

low_pass :: Int -> Signal -> Signal

low_pass v sig =
    let
        m = mask v $ length sig
    in ifft $ m * fft sig


-----------------------------------------------------------------------------
-- |
-- Module      :  Finance.Hqfl.Pricer.Black
-- Copyright   :  (C) 2016 Mika'il Khan
-- License     :  (see the file LICENSE)
-- Maintainer  :  Mika'il Khan <co.kleisli@gmail.com>
-- Stability   :  stable
-- Portability :  portable
--
----------------------------------------------------------------------------
{-# LANGUAGE FlexibleInstances #-}

module Finance.Hqfl.Pricer.Black where

import Finance.Hqfl.Instrument
import Statistics.Distribution.Normal
import Data.Random

class Black a where
  price :: a -> Double -> Double -> Double

instance Black (Option Future) where
  price (Option (Future f) m European k t) r v =
    case m of
       Call -> exp (-r * t) * (f * cdf normal d1 - k * cdf normal d2)
       Put  -> exp (-r * t) * (k * cdf normal (-d2) - f * cdf normal (-d1))
    where d1 = (log (f / k) + ((v * v) / 2) * t) / (v * sqrt t)
          d2 = d1 - v * sqrt t
          normal = Normal (0 :: Double) 1


{-# language BangPatterns #-}
{-|
Module      : Fitness
Description : Fitness Function for Genetic Algorithm for Poly Regression
Copyright   : (c) Fabricio Olivetti de Franca, 2021
License     : GPL-3
Maintainer  : fabricio.olivetti@gmail.com
Stability   : experimental
Portability : POSIX
|-}

module Fitness (evalFitness, decode) where

import GA

import Data.List (foldl')
import Data.List.Split hiding (split)
import Numeric.LinearAlgebra ((<\>),(#>),(|||), sumElements, size, Matrix, Vector, fromRows, fromList)

-- | Creates a fitness function
evalFitness :: Int            -- ^ number of terms
            -> [[Double]]     -- ^ matrix X
            -> Vector Double  -- ^ vector y
            -> Solution (Poly n)       -- ^ solution to evaluate
            -> Solution (Poly n)
evalFitness nTerms xss ys sol =
  case _fitness sol of
       Nothing -> sol{ _coeffs = Just betas, _fitness = Just f }
       _       -> sol
  where
    zss    = decode xss (map _getPoly $ _chromo sol)
    betas  = zss <\> ys
    ysHat  = zss #> betas
    f      = mse ys ysHat

-- | Decodes the polynomial into a transformed matrix
decode :: [[Double]] -> [(Bool, Int)] -> Matrix Double
decode xss chromo = fromRows (map (evalPoly chromo) xss)

-- | Eval a single row from the dataset
evalPoly :: [(Bool, Int)] -> [Double] -> Vector Double
evalPoly chromo xs = fromList $ foldl' joinTerms [1.0] terms
  where
    nVars = length xs
    ix    = cycle [0..nVars-1]
    xss   = cycle xs
    terms = zip3 chromo xss ix

    joinTerms ts   ((False,_), _, 0)   = 1.0 : ts
    joinTerms ts   ((False,_), _, _)   = ts
    joinTerms ts   ((True, k), x, 0)   = x^k : ts
    joinTerms (t:ts) ((True,k), x, ix) = t * x^k : ts

-- | Calculates the mean squared error
mse :: Vector Double -> Vector Double -> Double
mse ys ysHat = s / fromIntegral (size ys)
  where
    s = sumElements $ (ys - ysHat)^2


-- |
-- Module      : Southpaw.Picasso.Shapes
-- Description : Various functions for creating mathematical shapes
-- Copyright   : (c) Jonatan H Sundqvist, year
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created date year

-- TODO | -
--        -

-- SPEC | -
--        -

---------------------------------------------------------------------------------------------------
-- API
---------------------------------------------------------------------------------------------------
module Southpaw.Picasso.Shapes where

---------------------------------------------------------------------------------------------------
-- We'll need these
---------------------------------------------------------------------------------------------------
import Data.Complex

import Southpaw.Math.Constants

---------------------------------------------------------------------------------------------------
-- Function
---------------------------------------------------------------------------------------------------
-- Geometry ---------------------------------------------------------------------------------------
-- |
-- TODO: Start angle
-- TODO: Invalid arguments (eg. sides < 3) (use Maybe?)
-- TODO: Make polymorphic
-- TODO: Move to geometry section
-- polygon :: (Floating f, RealFloat f) => Int -> f -> Complex f -> [Complex f]
polygon :: Integral int => int -> Double -> Complex Double -> [Complex Double]
polygon sides radius origin = [ let θ = arg n in origin + ((radius * cos θ):+(radius * sin θ)) | n <- [1..sides]]
    where arg n = fromIntegral n * (2*π)/fromIntegral sides

-- |
-- TODO: Simplify and expound comments
arrow :: Complex Double -> Complex Double -> Double -> Double -> Double -> [Complex Double]
arrow from to sl sw hw = [from     + straight (sw/2), --
                          shaftEnd + straight (sw/2), --
                          shaftEnd + straight (hw/2), --
                          to,                         --
                          shaftEnd - straight (hw/2), --
                          shaftEnd - straight (sw/2), --
                          from     - straight (sw/2)] --
    where along a b distance = (a +) . mkPolar distance . snd . polar $ b-a -- Walk distance along the from a to b
          normal a b = let (mag, arg) = polar (b-a) in mkPolar mag (arg+π/2)
          shaftEnd = along from to sl --
          straight = along (0:+0) (normal from to) -- Vector perpendicular to the centre line


module Mand where
import Data.List.Split
import Data.Scientific as Scientific
import Data.Complex

type CS = Complex Scientific
type S = Scientific

makeSci :: String -> String -> Scientific
makeSci a b = read (a ++ "e" ++ b) :: Scientific

csqrt :: (Floating a, RealFloat a) => Complex a -> Complex a
csqrt x = (m * cos (t) :+ m * sin(t))
  where
    p = polar x
    m = sqrt $ fst p
    t = (snd p)/2

--cubed_iteration c z  = c + z * cos (z ** z)
cubed_iteration :: RealFloat a => Complex a -> Complex a -> Complex a
cubed_iteration c z = c + z * z * z

new_iteration :: RealFloat a => Complex a -> Complex a -> Complex a -> Complex a
new_iteration c n_z o_z = c + n_z * o_z

zzcosMask :: RealFloat a => Complex a -> Bool
zzcosMask (x :+y) = ((x > -0.95 && x < 0.1) && (y < 0.47 && y > -0.47)) || (( x <= -0.95 && x > -1.25) && (abs y < 0.42)) || (abs y < 0.2 && x <= -1.25 && x > -1.45) || (abs y < 0.4 && x >= 0.1 && x < 0.2) || (abs y < 0.17 && (x > 0.6 && x < 0.78)) || (abs y < 0.11 && x > 0.4 && x <= 0.85)

zzcos_iteration :: RealFloat a => Complex a -> Complex a -> Complex a
zzcos_iteration c z = c + z * z * (cos z)

mand_iteration :: RealFloat a => Complex a -> Complex a -> Complex a
mand_iteration c z  = c + z*z

mand_pow_iteration :: RealFloat a => Int -> Complex a -> Complex a -> Complex a
mand_pow_iteration pow c z  = c + z^pow

mand_iterationA :: CS -> CS -> CS
mand_iterationA (c:+b) (z:+x)  = (c + z*z - x* x) :+ (b + 2*z*x)

cexp :: Floating a => Complex a -> Complex a
cexp (a :+ b) = (s * cos b :+ s * sin b)
  where
    s = exp a

csin :: RealFloat a => Complex a -> Complex a
csin x = ((realPart c1 + realPart c2)/2 :+ (imagPart c1 + imagPart c2)/2)
 where
  c1 = exp x
  c2 = exp $ negate x

cpow :: RealFloat a => Complex a -> Complex a -> Complex a
cpow x y = cexp $ y * log  x

mand :: (RealFloat a) => Either (Complex a) CS -> Int -> Int
mand (Left c) max = general mand_iteration (Just mandSive) c max
mand (Right c) _ = generalA mand_iterationA c

sciMagnitude :: Complex Scientific  -> Scientific
sciMagnitude (a :+b) = a * a + b * b

mandSive :: RealFloat a => Complex a -> Bool
mandSive c@(x:+y) = ( let b = c + (1:+0) in realPart(b*(conjugate b)) < 0.05) ||((x < 0) && ( realPart (c * (conjugate c))  < 0.4)) ||(x < 0.25 && x  >= 0 && abs y  < 0.5)  -- takes 16 seconds

--( let b = a + (1:+0) in realPart(b*(conjugate b)) < 0.05) ||((realPart a < 0) && ( realPart (a * (conjugate a))  < 0.4)) ||(realPart a < 0.25 && realPart a >= 0 && abs (imagPart a) < 0.5) =  max -1  -- takes 16 seconds

general :: RealFloat a => (Complex a -> Complex a -> Complex a) -> Maybe (Complex a -> Bool)  -> Complex a -> Int -> Int
general g Nothing a max = count_iterations 0 (0 :+ 0) a
  where
    count_iterations n e x
      | n == max = max -1
      | realPart (x * (conjugate x)) > 4 = n
      | otherwise = count_iterations (n +1) e $ g a x
general g (Just sive) a max
--  | ((realPart a) < 0.25) && ((realPart a) > (-0.5)) && (abs (imagPart a) < 0.5) = 255 takes 21 secods
-- removed real check: (imagPart a == 0 && realPart a < 0.25 && realPart a > -2) ||
  | sive a = max - 1
  | otherwise = count_iterations 0 (0 :+ 0) a
  where
    count_iterations n e x
      | n == max = max -1
      | realPart (x * (conjugate x)) > 4 = n
      | otherwise = count_iterations (n +1) e $ g a x

generalA :: (CS -> CS -> CS) -> CS -> Int
generalA g a = count_iterations 0 (0 :+ 0) a
  where
    count_iterations n e x
      -- | n >= 255  = 255
      | n >= 16581375 = 16581375
      -- if using sci
      -- | (sciMagnitude x ) > (makeSci "4" "0") = n
      -- | otherwise = count_iterations (n + 1) e  $ g a $ roundComplex x
      -- if using float
      | (realPart x)^2 + (imagPart x)^2 > 4 = n
      | otherwise = count_iterations (n +1) e $ g a $ roundComplex x

general_list_R ::  (Complex Scientific -> Complex Scientific -> Complex Scientific) -> Complex Scientific -> [Complex Scientific]
general_list_R g a = get_list 0 a (0 :+ 0)
  where
    get_list 1000 _ _  = []
    get_list n (c :+d) (e :+ f)
      | (sciMagnitude (e :+ f)) > (scientific 2 0) = []
      | otherwise = (e:+f):  (get_list (n +1) (c :+d) ( g (c :+ d) (roundComplex (e :+ f))))
-- takes 1:13 to do 1000 iterations using roundComplex

general_list ::  (Complex Scientific -> Complex Scientific -> Complex Scientific) -> Complex Scientific -> [Complex Scientific]
general_list g (a :+ b) = get_list 0  (a :+ b) (0 :+ 0)
  where
    get_list 40  _ _  = []
    get_list n (c :+d) (e :+ f)
      | (sciMagnitude (e :+ f)) > (makeSci "2" "2") = []
      | otherwise = (e:+f):  (get_list (n +1) (c :+d) ( g (c :+ d) (e :+ f)))

general_list_D :: RealFloat a => (Complex a -> Complex a -> Complex a) -> Complex a -> [Complex a]
general_list_D g x = getlist 0 x (0 :+ 0)
  where
    getlist 1000 _ _ = []
    getlist n c z
      | realPart (z * conjugate z)> 4 = []
      | otherwise = z: (getlist (n +1) c (g c z))

roundSci' :: Scientific -> Scientific
roundSci' s = makeSci (getSign:cut) getExp
  where
    num = splitOn "e" $ show s
    removeSign = if s < 0 then tail (num !! 0) else (num !! 0)
    getSign = if s < 0 then '-'  else ' '
    getExp = if (length num) == 1 then "0" else num !! 1
    cut = take 100  removeSign

--rounds t 20 dp
roundSci :: Scientific -> Scientific
roundSci s
  | d <  0  = s
  | otherwise = scientific (div (coefficient s) (10^d)) $ d + base10Exponent s
  where
    d = (digitCount $ coefficient s) - 20

digitCount :: Integer -> Int
digitCount = go 1 . abs
    where
        go ds n = if n >= 10 then go (ds + 1) (n `div` 10) else ds
--(fromInteger $ round $ f * (10^n)) / (10.0^^n)
--Integer
roundComplex :: Complex Scientific -> Complex Scientific
roundComplex (a :+ b) = (roundSci' a :+ roundSci' b)


-- file: Astro/Kepler.hs
-- Various forms of Kepler's equation with efficient solvers

module Astro.Kepler ( eKepler
		    , gKepler
		    , solveEKepler
		    , solveGKepler
		    , keplerPropagate
		    )
	where

import Astro.Stumpff
import Astro.Misc

import Complex
import Debug.Trace
import Numeric.GSL
import Numeric.LinearAlgebra
import Text.Printf (printf)

-- Shorthands and such
type CD = Complex Double
type VD = Vector Double

-- Elliptic Kepler's equation
eKepler :: (Floating a) => a -> a -> a
eKepler e ea = ea - e * sin ea

-- General Variable Formulation of Kepler's eq
gKepler :: Double -> Double -> Double -> Double -> Double -> Double
gKepler r0 eta0 zeta0 beta x =
	r0 * x + eta0 * stumpff_G 2 beta x + zeta0 * stumpff_G 3 beta x

-- Solve Elliptic Kepler's equation iteratively
solveEKepler :: (Floating a) => a -> a -> a
solveEKepler e m =
	iterate 0 m m
	where
	--iterate n ea m | trace ("n,E: " ++ show n ++ " " ++ show ea) False = undefined
	iterate n ea m
		| ea' == ea || n > maxit = ea'
		| otherwise = iterate (n+1) ea' m
		where
		    fd = 1.0 - e * cos ea
		    fdd = e * sin ea
		    f = eKepler e ea - m
		    ea' = ea - f / sqrt (fd^2 - f * fdd)
	maxit = 10

-- Solve General Kepler's eq iteratively
--solveGKepler :: Double -> Double -> Double -> Double -> Double -> Double -> Double
solveGKepler r0 eta0 zeta0 beta x0 t (tol,maxIt) =
	if abs lastErr > tol
		then head sol `debug`
			("solveGKepler: iter. incomp., error: "
				++(show lastErr))
		else head sol
	where
		(sol,path) = root Hybrids tol maxIt f [x0]
		lastErr = (last.last) $ toLists path
		f vx = [gKepler r0 eta0 zeta0 beta (head vx) - t]

-- Solve r, v given r0, v0, gravitational parameter, delta_t
--keplerPropagate ::  VD -> VD -> Double -> (Double, Int) -> (VD,VD)
keplerPropagate dt (tol,maxIt) rv vv =
	(scale f rv + scale g vv , scale fd rv + scale gd vv)
	where (r0, v0) = (norm rv, norm vv)
	      tau = dt -- working in units where G(m_1+m_2) = 1
	      eta0 = rv <.> vv
	      beta = 2/r0 - v0^2
	      zeta0 = 1 - beta * r0
	      x = solveGKepler r0 eta0 zeta0 beta 0.0 tau (tol,maxIt)
	      r = r0 + eta0 * stumpff_G 1 beta x + zeta0 * stumpff_G 2 beta x
	      f = 1 - stumpff_G 2 beta x / r0
	      g = tau - stumpff_G 3 beta x
	      fd = - stumpff_G 1 beta x / (r * r0)
	      gd = 1 - stumpff_G 2 beta x / r


import Data.Complex

mandelbrot a = iterate (\z -> z^2 + a) 0 !! 50

main = mapM_ putStrLn [[if magnitude (mandelbrot (x :+ y)) < 2 then '*' else ' '
                           | x <- [-2, -1.9685 .. 0.5]]
                       | y <- [1, 0.95 .. -1]]


module Main where

import Data.Complex

main :: IO ()
main = do
  putStrLn "hello world"

maxIter :: Int
maxIter = 10000

-- data Scale = {
--                ULCorner :: Point
--                DRCorner :: Point
--                XScale :: Double
--                YScale :: Double
--              }

-- Simple boolean test to see if a point is in the set
isInSet :: RealFloat a => a -> a -> Bool
isInSet r i = case escapeTime r i of
              Nothing -> True
              Just _ -> False

-- Escape Time of Nothing means the point never escaped
escapeTime :: RealFloat a => a -> a -> Maybe Int
escapeTime r i
    | numIter == maxIter = Nothing
    | otherwise = Just numIter
    where numIter = mandel r i

-- Figures out how many iterations a point takes to escape
-- Or stops after maxIter iterations
mandel :: RealFloat a => a -> a -> Int
mandel r i = length . takeWhile ((<= 2) . magnitude) .
    take maxIter $ iterate (\z -> z^2 + (r :+ i)) 0


import Data.Complex (cis, phase)

meanAngle = (/ pi) . (* 180) . phase . sum . map (cis . (/ 180) . (* pi))

main = mapM_ (\angles -> putStrLn $ "The mean angle of " ++ show angles ++ " is: " ++ show (meanAngle angles) ++ " degrees")
       [[350, 10], [90, 180, 270, 360], [10, 20, 30]]


{-# LANGUAGE Rank2Types, BangPatterns, ScopedTypeVariables #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.AD.Newton
-- Copyright   :  (c) Edward Kmett 2010
-- License     :  BSD3
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
-- Portability :  GHC only
--
-----------------------------------------------------------------------------

module Numeric.AD.Newton
    (
    -- * Newton's Method (Forward AD)
      findZero
    , inverse
    , fixedPoint
    , extremum
    -- * Gradient Ascent/Descent (Reverse AD)
    , gradientDescent
    , gradientAscent
    , conjugateGradientDescent
    , conjugateGradientAscent
    ) where

import Prelude hiding (all, mapM, sum)
import Data.Functor
import Data.Foldable (all, sum)
import Data.Traversable
import Numeric.AD.Types
import Numeric.AD.Mode.Forward (diff, diff')
import Numeric.AD.Mode.Reverse (grad, gradWith')
import Numeric.AD.Internal.Combinators
import Numeric.AD.Internal.Composition

-- | The 'findZero' function finds a zero of a scalar function using
-- Newton's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes constant
-- ("it converges"), no further elements are returned.
--
-- Examples:
--
-- >>> take 10 $ findZero (\x->x^2-4) 1
-- [1.0,2.5,2.05,2.000609756097561,2.0000000929222947,2.000000000000002,2.0]
--
-- >>> import Data.Complex
-- >>> last $ take 10 $ findZero ((+1).(^2)) (1 :+ 1)
-- 0.0 :+ 1.0
findZero :: (Fractional a, Eq a) => (forall s. Mode s => AD s a -> AD s a) -> a -> [a]
findZero f = go where
  go x = x : if x == xn then [] else go xn where
    (y,y') = diff' f x
    xn = x - y/y'
{-# INLINE findZero #-}

-- | The 'inverse' function inverts a scalar function using
-- Newton's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes
-- constant ("it converges"), no further elements are returned.
--
-- Example:
--
-- >>> last $ take 10 $ inverse sqrt 1 (sqrt 10)
-- 10.0
inverse :: (Fractional a, Eq a) => (forall s. Mode s => AD s a -> AD s a) -> a -> a -> [a]
inverse f x0 y = findZero (\x -> f x - lift y) x0
{-# INLINE inverse  #-}

-- | The 'fixedPoint' function find a fixedpoint of a scalar
-- function using Newton's method; its output is a stream of
-- increasingly accurate results.  (Modulo the usual caveats.)
--
-- If the stream becomes constant ("it converges"), no further
-- elements are returned.
--
-- >>> last $ take 10 $ fixedPoint cos 1
-- 0.7390851332151607
fixedPoint :: (Fractional a, Eq a) => (forall s. Mode s => AD s a -> AD s a) -> a -> [a]
fixedPoint f = findZero (\x -> f x - x)
{-# INLINE fixedPoint #-}

-- | The 'extremum' function finds an extremum of a scalar
-- function using Newton's method; produces a stream of increasingly
-- accurate results.  (Modulo the usual caveats.) If the stream
-- becomes constant ("it converges"), no further elements are returned.
--
-- >>> last $ take 10 $ extremum cos 1
-- 0.0
extremum :: (Fractional a, Eq a) => (forall s. Mode s => AD s a -> AD s a) -> a -> [a]
extremum f = findZero (diff (decomposeMode . f . composeMode))
{-# INLINE extremum #-}

-- | The 'gradientDescent' function performs a multivariate
-- optimization, based on the naive-gradient-descent in the file
-- @stalingrad\/examples\/flow-tests\/pre-saddle-1a.vlad@ from the
-- VLAD compiler Stalingrad sources.  Its output is a stream of
-- increasingly accurate results.  (Modulo the usual caveats.)
--
-- It uses reverse mode automatic differentiation to compute the gradient.
gradientDescent :: (Traversable f, Fractional a, Ord a) => (forall s. Mode s => f (AD s a) -> AD s a) -> f a -> [f a]
gradientDescent f x0 = go x0 fx0 xgx0 0.1 (0 :: Int)
    where
        (fx0, xgx0) = gradWith' (,) f x0
        go x fx xgx !eta !i
            | eta == 0     = [] -- step size is 0
            | fx1 > fx     = go x fx xgx (eta/2) 0 -- we stepped too far
            | zeroGrad xgx = [] -- gradient is 0
            | otherwise    = x1 : if i == 10
                                  then go x1 fx1 xgx1 (eta*2) 0
                                  else go x1 fx1 xgx1 eta (i+1)
            where
                zeroGrad = all (\(_,g) -> g == 0)
                x1 = fmap (\(xi,gxi) -> xi - eta * gxi) xgx
                (fx1, xgx1) = gradWith' (,) f x1
{-# INLINE gradientDescent #-}

-- | Perform a gradient descent using reverse mode automatic differentiation to compute the gradient.
gradientAscent :: (Traversable f, Fractional a, Ord a) => (forall s. Mode s => f (AD s a) -> AD s a) -> f a -> [f a]
gradientAscent f = gradientDescent (negate . f)
{-# INLINE gradientAscent #-}

-- | Perform a conjugate gradient descent using reverse mode automatic differentiation to compute the gradient.
conjugateGradientDescent :: (Traversable f, Fractional a, Ord a) =>
                (forall s. Mode s => f (AD s a) -> AD s a) -> f a -> [f a]
conjugateGradientDescent f x0 = go x0 d0 d0
  where
    dot x y = sum $ zipWithT (*) x y
    d0 = negate <$> grad f x0
    go xi ri di = xi : go xi1 ri1 di1
      where
        ai  = last $ take 20 $ extremum (\a -> f $ zipWithT (\x d -> lift x + a * lift d) xi di) 0
        xi1 = zipWithT (\x d -> x + ai*d) xi di
        ri1 = negate <$> grad f xi1
        bi1 = max 0 $ dot ri1 (zipWithT (-) ri1 ri) / dot ri1 ri1
        -- bi1 = max 0 $ sum (zipWithT (\a b -> a * (a - b)) ri1 ri) / dot ri1 ri1
        di1 = zipWithT (\r d -> r * bi1*d) ri1 di
{-# INLINE conjugateGradientDescent #-}

-- | Perform a conjugate gradient ascent using reverse mode automatic differentiation to compute the gradient.
conjugateGradientAscent :: (Traversable f, Fractional a, Ord a) => (forall s. Mode s => f (AD s a) -> AD s a) -> f a -> [f a]
conjugateGradientAscent f = conjugateGradientDescent (negate . f)
{-# INLINE conjugateGradientAscent #-}


-- figure 2.1

import qualified Data.Vector as V

import Statistics.Sample
import Statistics.Sample.Histogram (histogram)

dice :: V.Vector Double
dice = V.fromList [1,2,3,4,5,6]

main = do
  putStrLn $ "the range: " ++ show (range dice)
  putStrLn $ "the mean:  " ++ show (mean dice)
  putStrLn $ "the harmonic mean, which is defined by 6 / (sum $ map (1/) [1..6]): " ++ show (harmonicMean dice)
  putStrLn $ "the standard deviation: " ++ show (stdDev dice)

  let (bs,hs) = histogram 6 dice
  putStrLn $ "the histogram is of course uniform: " ++ show (hs :: V.Vector Int)
  putStrLn $ "the (lower) boundaries: " ++ show bs


module BuckinghamPi
    (
    FundamentalUnit(..),
    DimensionalVariable(..),
    DimensionlessVariable(..),
    Exponent(..),
    generatePiGroups,
    numFundamentalUnits
    ) where

import           Control.Applicative
import           Data.List                     (nub, (\\))
import           Numeric.LinearAlgebra.Data
import           Numeric.LinearAlgebra.HMatrix
import           Text.Printf                   (printf)

-- Define sum type for fundamental units (M,L,T,Theta).
data FundamentalUnit = Mass | Length | Time | Temperature deriving (Eq,Show)

-- Define type synonym for Exponent.
type Exponent = Double

-- A variable is defined by the name, along with units, which is a list of tuples which contains
-- the FundamentalUnit along with an exponent.
data DimensionalVariable = DimensionalVariable {name :: String, units :: [(FundamentalUnit, Exponent)]} deriving (Eq,Show)

-- Define data type for a dimensionless variable.
data DimensionlessVariable = DimensionlessVariable {dimensionalVariables :: [DimensionalVariable], exponents :: [Exponent]} deriving (Eq)

-- Define function which normalizes exponents so that the smallest value is +/- 1.
normalizeExponents :: [Exponent] -> [Exponent]
normalizeExponents exps = map (/factor) exps
  where
    minExp = minimum exps
    factor = minExp * signum minExp

-- Define show instance for a dimensionless variable.
instance Show DimensionlessVariable where
  show (DimensionlessVariable dimVars exps) = concatMap (uncurry (printf "%s^(%.4f) ")) $ filter (\t -> abs (snd t) > 1.0e-6) $ zip (map name dimVars) (normalizeExponents exps)

getFundamentalUnits :: DimensionalVariable -> [FundamentalUnit]
getFundamentalUnits dimVar = map fst $ units dimVar

getUniqueFundamentalUnits :: [DimensionalVariable] -> [FundamentalUnit]
getUniqueFundamentalUnits = nub . concatMap getFundamentalUnits

numFundamentalUnits :: [DimensionalVariable] -> Int
numFundamentalUnits = length . getUniqueFundamentalUnits

-- Define function that computes the number of dimensionless groups given a list of dimensional variables.
computeNumberOfGroups :: [DimensionalVariable] -> Int
computeNumberOfGroups dimensionalVars = length dimensionalVars - numFundamentalUnits dimensionalVars

-- Define function that returns nonrepeating variables.
getNonrepeatingVars :: [DimensionalVariable] -> [Int] -> [DimensionalVariable]
getNonrepeatingVars dimVars repeatedIndices = [dimVars !! i | i <- [0..(length dimVars - 1)], i `notElem` repeatedIndices]

-- Define a function that returns the exponent on a given unit in the given dimensional variable.
getExponent :: FundamentalUnit -> DimensionalVariable -> Exponent
getExponent unit dimVar
  | unit `elem` map fst (units dimVar) = snd $ head $ filter (\u -> fst u == unit) $ units dimVar
  | otherwise = 0.0

-- Define a function that returns the equation that must be solved.
getEquation ::  [DimensionalVariable] -> FundamentalUnit -> [Exponent]
getEquation varCombination funUnit = map (getExponent funUnit) varCombination

getVariableCombinations :: [DimensionalVariable] -> [DimensionalVariable] -> [[DimensionalVariable]]
getVariableCombinations repeatingVars = map (\x -> repeatingVars ++ [x])

buildMatrix :: [DimensionalVariable] -> [FundamentalUnit] -> Matrix Double
buildMatrix varCombination funUnits = fromLists $ map (getEquation varCombination) funUnits

findDimensionlessGroup :: [FundamentalUnit] -> [DimensionalVariable] -> DimensionlessVariable
findDimensionlessGroup funUnits dimVars = DimensionlessVariable dimVars (concat . toLists $ nullspace $ buildMatrix dimVars funUnits)

-- Define function which takes a list of dimensional variables along with a list of integers
-- that identifies which are the repeated variables, and returns a list of dimensionless variables.
generatePiGroups :: [DimensionalVariable] -> [Int] -> [DimensionlessVariable]
generatePiGroups dimVars repeatingVarInds = map (findDimensionlessGroup funUnits) (getVariableCombinations repeatingVars nonrepeatingVars)
  where
    repeatingVars = [dimVars !! i | i <- repeatingVarInds]
    nonrepeatingVars = dimVars \\ repeatingVars
    funUnits = getUniqueFundamentalUnits dimVars


import Numeric.LinearAlgebra
import Graphics.Plot
import Numeric.GSL.Special(erf_Z, erf)

sombrero n = f x y where
    (x,y) = meshdom range range
    range = linspace n (-2,2)
    f x y = exp (-r2) * cos (2*r2) where
        r2 = x*x+y*y

f x =  sin x + 0.5 * sin (5*x)

gaussianPDF = erf_Z
cumdist x = 0.5 * (1+ erf (x/sqrt 2))

main = do
    let x = linspace 1000 (-4,4)
    mplot [f x]
    mplot [x, mapVector cumdist x, mapVector gaussianPDF x]
    mesh (sombrero 40)


#!/usr/bin/env runhaskell
{-# LANGUAGE OverloadedStrings, NoImplicitPrelude #-}
import ClassyPrelude hiding ((<>)) -- this <> is for semigroups (usually better)
import Data.Monoid ((<>)) -- ... but HMatrix only has Monoid
import Numeric.LinearAlgebra (tr) -- transpose
import qualified Numeric.LinearAlgebra as HMatrix
import qualified Data.Text.Format as Format

main = do
  -- Load the term-document matrix
  putStrLn "Loading"
  matrix <- HMatrix.loadMatrix "termdoc.txt"
  (u, sigma, vt) <- stochasticTruncatedSVD 50 2 matrix

  print $ HMatrix.size u
  print $ HMatrix.size sigma
  print $ HMatrix.size vt

-- Stochastic SVD. See Halko, Martinsson, and Tropp, 2010 for an explanation
stochasticTruncatedSVD :: Int -> Int -> HMatrix.Matrix Double -> IO (HMatrix.Matrix Double, HMatrix.Vector Double, HMatrix.Matrix Double)
stochasticTruncatedSVD top_vectors num_iterations original = do
  let (m, n) = HMatrix.size original
  let k = top_vectors + 10
  -- Stage A
  putStrLn "Stage A\n\tPower Iteration"
  omega <- HMatrix.randn n k
  let
    y 0 = original <> omega
    y n = seq (y (n-1)) $ original <> tr original <> y (n-1)
  bigY <- return $! y num_iterations -- Force evaluation here
  putStrLn "\tOrthogonalizing"
  q <- return $! HMatrix.orth bigY

  -- Stage B
  putStrLn "Stage B"
  let b = tr q <> original
      (uhat, sigma, vt) = HMatrix.thinSVD b
      u = q <> uhat

  return $! (HMatrix.takeColumns top_vectors u,
    HMatrix.subVector 0 top_vectors sigma,
    HMatrix.takeRows top_vectors $ tr vt)
  --putStrLn "Checking"
  --let ahat = u <> HMatrix.diag sigma <> tr vt

  --Format.print "MSE of approximation {}\n" [HMatrix.sumElements ((a-ahat)^2) / (fromIntegral $ m*n)]


module Main where

import Lib
import Data.Complex
import Graphics.Image

width = 750
height = 500
scaleDivisor = 250

xOffset = 2
yOffset = 1

isInMandelbrotSet :: Complex Double -> Complex Double -> Int -> Int
isInMandelbrotSet c z 0 = 0
isInMandelbrotSet c z itterations | (Data.Complex.magnitude z)  < 4 = isInMandelbrotSet c (z**2 + c) (itterations - 1)
                                  | (Data.Complex.magnitude z) >= 4 = itterations

calculatePixel :: Int -> Int -> Pixel RGB Double
calculatePixel y x = PixelRGB  0 0 value where
    real = (fromIntegral x / scaleDivisor) - xOffset
    imag = (fromIntegral y / scaleDivisor) - yOffset
    value = ((fromIntegral (isInMandelbrotSet (real :+ imag) 0 1000)) / 1000)

main = do
    let image = makeImageR VS (height, width) (\(y, x) -> calculatePixel y x)
    writeImageExact PNG [] "output.png" image