module Utils where

import Data.Vector.Unboxed as U
import Prelude as P
import Data.Decimal
import GHC.Word

import qualified Statistics.Matrix as M
import qualified Statistics.Matrix.Types as T
import qualified Statistics.Matrix.Algorithms as Alg
import qualified Statistics.Matrix.Function as Func
import qualified Statistics.Matrix.Mutable as Mut

import qualified Fixtures as F

-- | Operations on vectors

-- | Operations on matrices

-- | Check if two matrices are equal within some precision
isEqual :: Word8 -> T.Matrix -> T.Matrix -> Bool
isEqual prec a b =
    let
        a' = makeDecimal prec a
        b' = makeDecimal prec b
    in
        a' == b'

------------------------------ matrix property checks --------------------------------------

-- | Check if given matrix is orthogonal
isOrtho :: T.Matrix -> Bool
isOrtho mat = isEqual 8 (M.multiply mat (M.transpose mat)) F.matId

-- | Checks if the n'th row of the matrix have 'n' leading zeros
hasLeadingZeros :: Int -> T.Vector -> Bool
hasLeadingZeros n row = U.foldl (&&) True (U.map (== 0) (U.take n row))

-- | pass the n'th row of the matrix to hasLeadingZeros
isRowFromTriMatrix :: T.Matrix -> Int -> Bool
isRowFromTriMatrix mat n = hasLeadingZeros n (M.row mat n)

-- | Helper function to check if the given matrix is upper triangular
upperTriHelper :: T.Matrix -> Int -> Bool
upperTriHelper mat n = P.foldl (&&) True $ P.map (isRowFromTriMatrix mat) [0..(n-1)]

-- | to check if the given matrix is upper triangular
isUpperTri :: T.Matrix -> Bool
isUpperTri mat = upperTriHelper mat (fst (M.dimension mat))

-- | to check if the given matrix is invertible
isInvertible :: T.Matrix -> Bool
isInvertible mat = isFloatZero (det mat) 0.0000001

------------------------------- matrix manipulation functions -------------------------

-- | make all values of the matrix absolute
makeAbs :: T.Matrix -> T.Matrix
makeAbs mat = M.map P.abs mat

-- | make all values of the (row-list represented matrix)
-- | decimal values with a given precision
makeDecimal :: Word8 -> T.Matrix -> [[Decimal]]
makeDecimal prec mat =
    let
        matList = M.toRowLists mat
        rowToDecimal prec ls = P.map (realFracToDecimal prec) ls
    in
        P.map (rowToDecimal prec) matList

-- | to get the determinant of a given matrix
-- | taken from here:
-- | http://michaeljgilliland.blogspot.com/2013/04/determinant-in-haskell.html
det :: T.Matrix -> Double
det mat =
    let
        detHelper [[x]] = x
        detHelper mat = P.sum [(-1)^i*x*(detHelper (getRest i mat)) | (i, x) <- P.zip [0..] (P.head mat)]
        listMatrix = M.toRowLists mat
    in
        detHelper listMatrix

-- |  getRest function simply returns the matrix without the head row (topmost) and without the ith column
getRest :: Int -> [[Double]] -> [[Double]]
getRest col mat =
    let
        decapitatedMat = P.tail mat
    in
        P.map (leaveElement col) decapitatedMat

----------------------------- auxiliary helper functions -----------------------------------

-- | function to delete i'th element from list
leaveElement :: Ord a => Int -> [a] -> [a]
leaveElement _ [] = []
leaveElement i (x:xs)
        | i == 0    = xs
        | otherwise = x : leaveElement (i-1) xs

-- | compare a double value to some epsilon
isFloatZero :: Double -> Double -> Bool
isFloatZero n eps =
    if (abs n) > (abs eps)
        then True
    else False


module NN where

import Numeric.LinearAlgebra as LA

rectifiedLinearUnit :: R -> R
rectifiedLinearUnit = max 0

rectifiedLinearUnit' :: R -> R
rectifiedLinearUnit' x = if rectifiedLinearUnit x == 0 then 0 else 1

trainNeuralNetwork :: [((R, R), R)] -> Matrix R -> Matrix R -> Matrix R -> R -> Int -> (Matrix R, Matrix R, Matrix R)
trainNeuralNetwork xys wT b vT learningRate epochs = iterate (step xys learningRate) (wT, b, vT) !! epochs
    where
        step xys learningRate (wT, b, vT) = train xys wT b vT learningRate

train :: [((R, R), R)] -> Matrix R -> Matrix R -> Matrix R -> Double -> (Matrix R, Matrix R, Matrix R)
train xys wT b vT learningRate = train' xys wT b vT learningRate
    where
        train' [] wTs b vTs _ = (wTs, b, vTs)
        train' (((x1,x2), y):xys') wTs b vTs learningRate =
            let (wTs', b', vTs') = trainPoint wTs b vTs learningRate x1 x2 y
            in train' xys' wTs' b' vTs' learningRate

testNeuralNetwork :: [(R, R)] -> Matrix R -> Matrix R -> Matrix R -> R -> [R]
testNeuralNetwork xs wT b vT learningRate = map (calc wT b vT learningRate) xs
    where
        calc wT b vT learningRate (x1, x2) = predict wT b vT learningRate x1 x2

trainPoint :: Matrix R -> Matrix R -> Matrix R -> R -> R -> R -> R -> (Matrix R, Matrix R, Matrix R)
trainPoint wT b vT learningRate x1 x2 y = (wT - scalar learningRate * dw, b - scalar learningRate * db, vT - scalar learningRate * dv)
    where
        x = (2 >< 1) [x1, x2]

        z = wT LA.<> x + b
        a =  cmap rectifiedLinearUnit z

        u =  vT LA.<> a

        du = u - scalar y
        dv = du LA.<> tr a
        da = tr vT LA.<> du
        dz = da * cmap rectifiedLinearUnit' z

        db = dz
        dw = dz LA.<> tr x

predict :: Matrix R -> Matrix R -> Matrix R -> R -> R -> R -> R
predict wT b vT learningRate x1 x2  = u `atIndex` (0,0)
    where
        z = wT LA.<> (2 >< 1) [x1, x2] + b
        a = cmap rectifiedLinearUnit z
        u = vT LA.<> a


-- |
-- Module:      Math.NumberTheory.RootsOfUnity
-- Copyright:   (c) 2018 Bhavik Mehta
-- Licence:     MIT
-- Maintainer:  Bhavik Mehta <bhavikmehta8@gmail.com>
--
-- Implementation of roots of unity
--

module Math.NumberTheory.RootsOfUnity
(
-- * Roots of unity
   RootOfUnity (..)
-- ** Conversions
   , toRootOfUnity
   , toComplex )

where

import Data.Complex                                        (Complex(..), cis)
import Data.Semigroup                                      (Semigroup(..))
import Data.Ratio                                          ((%), numerator, denominator)

-- | A representation of <https://en.wikipedia.org/wiki/Root_of_unity roots of unity>: complex
-- numbers \(z\) for which there is \(n\) such that \(z^n=1\).
newtype RootOfUnity =
  RootOfUnity { -- | Every root of unity can be expressed as \(e^{2 \pi i q}\) for some
                -- rational \(q\) satisfying \(0 \leq q < 1\), this function extracts it.
                fromRootOfUnity :: Rational }
  deriving (Eq)

instance Show RootOfUnity where
  show (RootOfUnity q)
    | n == 0    = "1"
    | d == 1    = "-1"
    | n == 1    = "e^(πi/" ++ show d ++ ")"
    | otherwise = "e^(" ++ show n ++ "πi/" ++ show d ++ ")"
    where n = numerator (2*q)
          d = denominator (2*q)

-- | Given a rational \(q\), produce the root of unity \(e^{2 \pi i q}\).
toRootOfUnity :: Rational -> RootOfUnity
toRootOfUnity q = RootOfUnity ((n `rem` d) % d)
  where n = numerator q
        d = denominator q
        -- effectively q `mod` 1
  -- This smart constructor ensures that the rational is always in the range 0 <= q < 1.

-- | This Semigroup is in fact a group, so @'stimes'@ can be called with a negative first argument.
instance Semigroup RootOfUnity where
  RootOfUnity q1 <> RootOfUnity q2 = toRootOfUnity (q1 + q2)
  stimes k (RootOfUnity q) = toRootOfUnity (q * fromIntegral k)

instance Monoid RootOfUnity where
  mappend = (<>)
  mempty = RootOfUnity 0

-- | Convert a root of unity into an inexact complex number. Due to floating point inaccuracies,
-- it is recommended to avoid use of this until the end of a calculation. Alternatively, with
-- the [cyclotomic](http://hackage.haskell.org/package/cyclotomic-0.5.1) package, one can use
-- @[polarRat](https://hackage.haskell.org/package/cyclotomic-0.5.1/docs/Data-Complex-Cyclotomic.html#v:polarRat)
-- 1 . @'fromRootOfUnity' to convert to a cyclotomic number.
toComplex :: Floating a => RootOfUnity -> Complex a
toComplex (RootOfUnity t)
  | t == 1/2 = (-1) :+ 0
  | t == 1/4 = 0 :+ 1
  | t == 3/4 = 0 :+ (-1)
  | otherwise = cis . (2*pi*) . fromRational $ t


{-# Language TypeApplications #-}
module PointSpec where

import Test.Hspec
import Test.QuickCheck hiding (scale)
import Test.Invariant
import Data.Complex
import Geometry
import Geometry.Testing

spec :: Spec
spec = describe "Point" $ do
    describe "Affinity" $ do
      it "1" $ property $ cmp `inverts` Point
      it "2" $ property $ Point `inverts` cmp
      it "3" $ property $ cmp `inverts` (Point . cmp)
      it "4" $ property $ (Point . cmp) `inverts` cmp

    describe "pointOn" $ do
      it "1" $ xy (pointOn (aCircle # scale 2) 0) ~= (2, 0)
      it "2" $ xy (pointOn (aCircle # scale 2) 0.5) ~= (-2, 0)
      it "3" $ property $ \(AnyCircle c) t -> c `isContaining` pointOn c t
      it "4" $ property $ \(AnyLine l) t -> l `isContaining` pointOn l t
      it "5" $ property $ \n t ->
        let p = regularPoly (3 + abs n)
        in p `isContaining` pointOn p t

    describe "aligned" $ do
      it "1" $ property $ \(AnyLine l) xs ->
        length xs > 1 ==> aligned (param l <$> xs)
      it "2" $ property $ \(Position a) (Position b) (Position c) ->
        (aligned [a,b,c] ==> isDegenerate (Triangle [a,b,c])) .&.
        (not (aligned [a,b,c]) ==> not (isDegenerate (Triangle [a,b,c])))

    describe "APoint" $ do
      it "1" $ property $ toPoint `inverts` (asPoint @XY)
      it "2" $ property $ toPoint `inverts` (asPoint @Cmp)
      it "3" $ property $ toPoint `inverts` (asPoint @(Maybe Point))


{-# LANGUAGE TypeApplications #-}
module ArrayFire.StatisticsSpec where

import ArrayFire    hiding (not)

import Data.Complex
import Test.Hspec

spec :: Spec
spec =
  describe "Statistics spec" $ do
    it "Should find the mean" $ do
      mean (vector @Double 10 [1..]) 0
        `shouldBe`
           5.5
    it "Should find the weighted-mean" $ do
      meanWeighted (vector @Double 10 [1..]) (vector @Double 10 [1..]) 0
        `shouldBe`
           7.0
    it "Should find the variance" $ do
      var (vector @Double 8 [1..8]) False 0
        `shouldBe`
           6.0
    it "Should find the weighted variance" $ do
      varWeighted (vector @Double 8 [1..]) (vector @Double 8 (repeat 1)) 0
        `shouldBe`
           5.25
    it "Should find the standard deviation" $ do
      stdev (vector @Double 10 (cycle [1,-1])) 0
        `shouldBe`
           1.0
    it "Should find the covariance" $ do
      cov (vector @Double 10 (repeat 1)) (vector @Double 10 (repeat 1)) False
        `shouldBe`
           0.0
    it "Should find the median" $ do
      median (vector @Double 10 [1..]) 0
        `shouldBe`
           5.5
    it "Should find the mean of all elements across all dimensions" $ do
      fst (meanAll (matrix @Double (2,2) [[10,10],[10,10]]))
        `shouldBe`
           10
    it "Should find the weighted mean of all elements across all dimensions" $ do
      fst (meanAllWeighted (matrix @Double (2,2) [[10,10],[10,10]]) (matrix @Double (2,2) [[10,10],[10,10]]))
        `shouldBe`
           10
    it "Should find the variance of all elements across all dimensions" $ do
      fst (varAll (vector @Double 10 (repeat 10)) False)
        `shouldBe`
           0
    it "Should find the weighted variance of all elements across all dimensions" $ do
      fst (varAllWeighted (vector @Double 10 (repeat 10)) (vector @Double 10 (repeat 10)))
        `shouldBe`
           0
    it "Should find the stdev of all elements across all dimensions" $ do
      fst (stdevAll (vector @Double 10 (repeat 10)))
        `shouldBe`
           0
    it "Should find the median of all elements across all dimensions" $ do
      fst (medianAll (vector @Double 10 [1..]))
        `shouldBe`
           5.5
    it "Should find the correlation coefficient" $ do
      fst (corrCoef (vector @Int 10 [1..] ) ( vector @Int 10 [10,9..] ))
        `shouldBe`
           (-1.0)
    it "Should find the top k elements" $ do
      let (vals,indexes) = topk ( vector @Double 10 [1..] ) 3 TopKDefault
      vals `shouldBe` vector @Double 3 [10,9,8]
      indexes `shouldBe` vector @Double 3 [9,8,7]


import Criterion.Main
import Numeric.LinearAlgebra

n = 10 :: Int
main = defaultMain [
         bench "matrix creation" $ w (\_ -> fromLists sampleList :: Matrix Double),
         bench "random matrix generation" $ whnfIO (randn n n),
         bench "matrix multiplication" $ w (\_ -> sampleDiag <> sampleDiag),
         bench "matrix slicing" $ whnf (\n -> sampleDiag ?? (Take (n - 1), Take (n - 1))) n,

         bench "scalar multiplication" $ whnf (\n -> sampleDiag * (fromIntegral n)) n,

         bench "matrix-vector addition" $ w (\_ -> sampleDiag + sampleDiag),

         bench "transpose" $ w (\_ -> tr sampleDiag)
       ]
  where
    w f = f `whnf` 0
    sampleList = replicate n $ replicate n 1
    sampleDiag = diagl [1..fromIntegral n]


module NeuralNet.Cost (computeCost) where

import Numeric.LinearAlgebra

computeCost :: Matrix Double -> Matrix Double -> Double
computeCost al y
  | rows al /= 1       = error "al should be a row vector"
  | rows y /= 1        = error "y should be a row vector"
  | cols y /= cols al = error "y and al should be same length"
  -- TODO: Must be a better impl than this sum
  | otherwise           = (-1 / fromIntegral m) * sum (concat (toLists (term1 + term2)))
    where
      m = cols y
      unitRowVector = row (replicate m 1)
      term1 = y * cmap log al
      term2 = (unitRowVector - y) * (cmap log (unitRowVector - al))


-- code from Edward Kmett's blog post
-- Mirrored Lenses
-- http://comonad.com/reader/2012/mirrored-lenses/

-- Slightly adapted since Data.PolyLens doesn't use Rank 2 types.
-- Mostly this means adding "Functor f =>" everywhere

import Data.Complex
import Data.PolyLens

realLens :: (RealFloat a, Functor f) => Lens (Complex a) a f
realLens f (r :+ i) = fmap (:+ i) (f r)

imagLens :: (RealFloat a, Functor f) => Lens (Complex a) a f
imagLens f (r :+ i) = fmap (r :+) (f i)

fstLens :: Functor f => PolyLens (a,c) (b,c) a b f
fstLens f (a, b) = fmap (\x -> (x, b)) (f a)

sndLens :: Functor f => PolyLens (a,b) (a,c) b c f
sndLens f (a, b) = fmap ((,) a) (f b)

swap :: (a, b) -> (b, a)
swap (a, b) = (b, a)

swapped :: Functor f => PolyLens (a, b) (c, d) (b, a) (d, c) f
swapped = iso swap swap

negated :: (Num a, Functor f) => Lens a a f
negated = iso negate negate

traverseLens :: ((c -> c) -> a -> b) -> a -> b
traverseLens f = f id

getFst :: Getter (a, b) a r
getFst = getting fst

getSnd :: Getter (a, b) b r
getSnd = getting snd

getPhase :: RealFloat a => Getter (Complex a) a r
getPhase = getting phase

getAbs, getSignum :: Num a => Getter a a r
getAbs = getting abs
getSignum = getting signum

plus, times :: Num a => Setter a a a
plus = setting (+)
times = setting (*)

reals :: (RealFloat a, RealFloat b) => Modifier (Complex a) (Complex b) a b
reals = modifying (\f (r :+ i) -> f r :+ f i)


import Data.Complex
import Data.Array
import System.Random

predict_point_from_sources :: [Float] -> [Float] -> (Float,Float) -> Complex Float
predict_point_from_sources ls ms (u,v) =  foldl1 (+) $ map (\(l, m) ->
                                                               cos(2*pi * (u*l + v*m)) :+
                                                               sin(2*pi * (u*l + v*m))) $ zip ls ms

sample_uv_grid :: Int -> Float -> [Float] -> [Float] -> Array(Int, Int) (Complex Float)
sample_uv_grid size du ls ms =
  array((0,0),(size-1,size-1))
  [((x,y),predict_point_from_sources ls ms (du*fromIntegral(x - quot size 2),du*fromIntegral(y - quot size 2))) | x<-[0..size-1], y<-[0..size-1]]

generateArray :: Float -> Float -> IO [Float]
generateArray lowerb upperb =
  (getStdGen >>= return . randomRs (lowerb, upperb))

main :: IO()
main = do
  l <- generateArray 0.0 1.0
  m <- generateArray 0.0 1.0
  let ls = map (\x -> (x - 0.5)*0.1) $ take 100 $ l
  let ms = map (\x -> (x - 0.5)*0.1) $ take 100 $ m
  let point = predict_point_from_sources ls ms (3.4,5.6)
  print point
  let uvgrid = sample_uv_grid 5 5.0 ls ms
  print uvgrid


-- Haskell code for navigating the interface between lists of edges and
-- adjacency matrices.

module Adjacency where

import Numeric.LinearAlgebra hiding (fromList)
import qualified Numeric.LinearAlgebra as LA (fromList)
import Numeric.LinearAlgebra.Data
import qualified Data.Map.Strict as M
import qualified Data.List as L
import Data.Ord
import Data.KMeans

-- Data structure for graphs and networks. Just a list of edges.
data Graph a = Graph { nodes :: [a]
                     , edges :: [(a, a)]
                     } deriving (Show, Eq)

data Graph' a = Graph' { nodes' :: [a]
                       , edges' :: M.Map a [a]
                       } deriving (Show, Eq)

data Graph'' a = Graph'' { nodes'' :: [a]
                         , edges'' :: Matrix Double
                         } deriving (Show, Eq)

-- Converts the edge list to a map structure. Takes up more space, but it's
-- faster for looking up what else is connected to a particular node.
toMap :: (Ord a, Eq a) => Graph a -> Graph' a
toMap (Graph ns es) = Graph' ns es'
  where
    fInsert o o' m x = M.insertWith (++) (o x) [o' x] m
    iMap = M.fromList $ map (\x -> (x, [])) ns
    es' = L.foldl' (fInsert snd fst) (L.foldl' (fInsert fst snd) iMap es) es

-- Converts a graph (edge list) to an adjacency matrix structure.
toAdjacency :: (Ord a, Eq a) => Graph a -> Graph'' a
toAdjacency gr@(Graph ns es) = Graph'' ns m
  where
    n = length ns
    m = (n >< n) . concat . map (to01List ns) $ map ((M.!) . edges' $ toMap gr) ns

-- Converts a graph (edge list) to an adjacency matrix structure. About 1/3
-- slower than `toAdjacency`.
-- toAdjacency' :: (Ord a, Eq a) => Graph a -> Graph'' a
-- toAdjacency' gr@(Graph ns es) = Graph'' ns m
--   where
--     n = length ns
--     m = (n >< n) . M.foldl' (++) [] $ M.map (to01List ns) (edges' $ toMap es)

-- Given a list of edges which are all connected to a node, generates
-- a list of 0's and 1's denoting membership.
to01List ::  (Eq a) => [a] -> [a] -> [Double]
to01List ns es = map (boolTo01 . flip elem es) ns

-- Converts False/True to 0/1.
boolTo01 :: Bool -> Double
boolTo01 b
  | b    = 1
  | True = 0

-- SVD approximation to a graph (uses k-element SVD of the adjacency
-- matrix). In reality we need the SVD of XY^T, but that will come later.
svdGraph :: Int -> Graph'' a -> Graph'' a
svdGraph k (Graph'' ns es) = Graph'' ns $ compress k es

-- This function was written by Nicolas Favre-Felix. It returns an
-- approximation of the given matrix, by taking only the k largest entries
-- of Sigma.
compress :: Int -> Matrix Double -> Matrix Double
compress k m = u_k <> sigma_k <> v_k where
  (u,sigma,v) = svd m  -- get SVD
  sigma_k = takeBoth k $ diag sigma -- keep k values of Σ
  u_k = takeColumns k u    -- keep k columns of U
  v_k = takeRows k $ tr v   -- keep k rows of v

-- Constructs the graph Laplacian.
laplacian :: (Ord a, Eq a) => Graph a -> Matrix Double
laplacian gr = (diag $ LA.fromList counts) - (edges'' $ toAdjacency gr)
  where
    counts = M.foldl (\x y -> x ++ [fromIntegral $ length y]) [] . edges'
             $ toMap gr

-- A little wrapper to simplify taking the upper submatrix.
takeBoth :: (Element t) => Int -> Matrix t -> Matrix t
takeBoth d = takeColumns d . takeRows d

-- Spanning tree counter; uses Kirchhoff's Theorem.
spanTreeCount :: (Ord a, Eq a) => Graph a -> Integer
spanTreeCount gr@(Graph ns es) = round . det . takeBoth n' $ laplacian gr
  where
    n' = length ns - 1

-- spectralCluster performs spectral clustering to estimate group membership.
spectralCluster :: (Ord a, Eq a) => Int -> Graph a -> [[[Double]]]
spectralCluster k gr@(Graph ns es) = kmeans k u
  where
    n = length ns
    -- u = map toList . drop (n-k) . toColumns . snd . eigSH . sym $ laplacian gr
    u = map toList . drop (n-k) . toColumns . snd . eigSH $ laplacian gr

-- spectralCluster performs spectral clustering to estimate group
-- membership. Estimates the best choice of k using consecutive
-- differences of the eigenvalues.
spectralCluster' :: (Ord a, Eq a) => Graph a -> [[[Double]]]
spectralCluster' gr@(Graph ns es) = kmeans k u
  where
    n = length ns
    -- e = eigSH . sym $ laplacian gr
    e = eigSH $ laplacian gr
    k = (+1) . fst . L.minimumBy (comparing snd) . zip [0..] . reverse
        . diff . toList $ fst e
    u = map toList . drop (n-k) . toColumns $ snd e

lambdaSelect :: (Ord a, Eq a) => Graph a -> Int
lambdaSelect gr@(Graph ns es) = k
  where
    n = length ns
    -- e = eigSH . sym $ laplacian gr
    e = eigSH $ laplacian gr
    k = (+1) . fst . L.minimumBy (comparing snd) . zip [0..] . reverse
        . diff . toList $ fst e
    u = map toList . drop (n-k) . toColumns $ snd e

-- Consecutive differences.
diff :: (Num a) => [a] -> [a]
diff xs = zipWith (-) (drop 1 xs) xs


module Window (

    hamming
  , hammingC
  , c
  , Window
  , CWindow

  ) where

import Data.Vector
import qualified Data.Complex as C

type Window = Vector Double -- Windowing function

type CWindow = Vector (C.Complex Double) -- Windowing function

-- quick conversion to complex
c :: Num a => a -> C.Complex a
c = (C.:+ 0)

-- A hamming window of size n.
hamming :: Int -> Window
hamming m = generate m hamming' where
  hamming' n = 0.54 - 0.46*cos(2 * pi * (fromIntegral n)/(fromIntegral m-1))

-- A hamming window of type Complex Double of size n
hammingC :: Int -> CWindow
hammingC m = generate m hamming' where
  hamming' n = c $ 0.54 - 0.46*cos(2 * pi * (fromIntegral n)/(fromIntegral m-1))


{-|
Module      : ODMatrix
Description : Representación y procesamiento de matrices origen-destino.

Base de funciones para la operación sobre matrices origen-destino.

ODM Features
The ODM features are indicators to verify if two ODMs are elements of the same solution set.
-}
module ODMatrix (
    ODM
  , FeatVec
  --
  , boardingVector
  , alightingVector
  , onBoardVector
  --
  , inputVector
  --
  , getFeatures
  , compareFeatures
  --
  , toFlatList
  , fromFlatList
  ) where

  import Numeric.LinearAlgebra (Matrix, Vector,
                                rows, cols, flatten,
                                toColumns, toRows,
                                subMatrix, vector,
                                vjoin, subVector)
  import Numeric.LinearAlgebra.Data (toList, takeDiag, fromLists, toLists, assoc)

  import Control.Monad (join)

  import ODMatrix.Util.TriangularNumbers (isTrg, trglvl)

  -- | Representación de una matriz de viajes. Debe ser cuadrada.
  type ODM = Matrix Double

  -- | The feature represents the input information.
  type FeatVec = Vector Double

  -- * Cálculo de vectores de información derivados. Feature Vectors.

  -- | Dada una matriz de viajes, calcula el vector de ascensos.
  boardingVector :: ODM      -- ^ Matriz de viajes de entrada.
                 -> FeatVec  -- ^ Vector de ascensos
  boardingVector = vector . map (sum . toList) . toRows

  -- | Dada una matriz de viajes, calcula el vector de descensos.
  alightingVector :: ODM      -- ^ Matriz de viajes de entrada.
                  -> FeatVec  -- ^ Vector de ascensos
  alightingVector = vector . map (sum . toList) . toColumns

  -- | Dada una matriz de viajes, calcula el vector de pasajeros a bordo.
  onBoardVector :: ODM      -- ^ Matriz de viajes de entrada.
                -> FeatVec  -- ^ Vector de pasajeros a bordo.
  onBoardVector m = vector ((sum . toList) . flatten <$> blocks)
    where s = cols m
          blocks = map (\x -> subMatrix (0,x) (x+1,s-x) m) [0..(s-1)]

  -- | Get all the feature vectors. (boarding, alighting, onBoard)
  getFeatures :: ODM -> (FeatVec, FeatVec, FeatVec)
  getFeatures m = (boardingVector m, alightingVector m, onBoardVector m)

  -- | Given two ODMs, verify if they are members of the same solution set.
  compareFeatures :: ODM -> ODM -> Bool
  compareFeatures a b = getFeatures a == getFeatures b

  -- | Given a ODM, build the vector of independent terms
  inputVector :: ODM -> Vector Double
  inputVector odm = vjoin [a,b]
    where a = alightingVector odm
          b = subVector 1 (cols odm - 1) $ boardingVector odm

  -- * Convertions

  -- | Aplana una matriz de origen destino, por filas y solo el triangulo superior.
  toFlatList :: ODM
             -> [Double]
  toFlatList = join . zipWith drop [0..] . toLists

  -- | Convierte una lista en una matriz origen destino.
  fromFlatList :: [Double]
               -> ODM
  fromFlatList ls
    | (not . isTrg) (length ls) =
      error "fromFlatList: Invalid list. The length must be a triangular number."
    | otherwise =
      assoc (n,n) 0 . zip [(i,j)|i<-[0..n-1],j<-[i..n-1]] $ ls
        where n = trglvl $ length ls


module JuliaSet where

import Data.Complex
import Graphics.Image

-- This will work for any function with one complex number z as input.
-- Constants and coefficients must be set through partial application

-- in the output, x is the real portion of the result and y is the imaginary
-- portion, with these values laid out as one would expect with a cartesian
-- grid
visualizeJuliaSet :: (Complex Double -> Complex Double) -- f
                -> Double -- escapeRadius
                -> Int -- width in pixels
                -> Int -- height in pixels
                -> Int -- maxIter
                -> String -- output file name
                -> String -- array representation
                -> IO ()
visualizeJuliaSet f escapeRadius width height maxIter outputFilename arrType
  | arrType == "VU" = writeImage outputFilename
      (makeImageR VU (width, height) g)
  | arrType == "VS" = writeImage outputFilename
      (makeImageR VS (width, height) g)
  | arrType == "RSU" = writeImage outputFilename
      (makeImageR RSU (width, height) g)
  | arrType == "RPU" = writeImage outputFilename
      (makeImageR RPU (width, height) g)
  | arrType == "RSS" = writeImage outputFilename
      (makeImageR RSS (width, height) g)
  | arrType == "RPS" = writeImage outputFilename
      (makeImageR RPS (width, height) g)
  | otherwise = writeImage outputFilename (makeImageR RPU (width, height) g)
    where
      g = pixelToJuliaSetValue f escapeRadius width height maxIter

pixelToJuliaSetValue :: (Complex Double -> Complex Double) -- f
                     -> Double -- escapeRadius
                     -> Int -- width
                     -> Int -- height
                     -> Int -- maxIter
                     -> (Int, Int) -- i, j
                     -> Pixel RGB Double -- result: pixel value
pixelToJuliaSetValue f escapeRadius width height maxIter (i, j)
    = toPixelRGB (PixelHSI pixelValue 0.5 0.5)
      where
        pixelValue :: Double
        pixelValue = (fromIntegral (julia f (x :+ y) escapeRadius maxIter))/
                    (fromIntegral maxIter+1)
        x = -escapeRadius + (2*escapeRadius/(fromIntegral width)) *
          (fromIntegral j)
        y = -(-escapeRadius + (2*escapeRadius/(fromIntegral height)) *
          (fromIntegral i))

-- complex and real parts of z must be between -escapeRadius and escapeRadius
julia :: (Complex Double -> Complex Double) -- f
      -> Complex Double -- z
      -> Double -- escapeRadius
      -> Int -- maxIter
      -> Int -- result
julia f z escapeRadius maxIter = juliaRecursive f z escapeRadius 1 maxIter

-- not meant to be called by user; use julia instead
juliaRecursive :: (Complex Double -> Complex Double)
      -> Complex Double
      -> Double -- escapeRadius
      -> Int -- i
      -> Int -- maxIter
      -> Int -- result
juliaRecursive f z escapeRadius i maxIter
  | i > maxIter = maxIter
  | Data.Complex.magnitude(z') > escapeRadius = i
  | otherwise = juliaRecursive f z' escapeRadius (i+1) maxIter
    where
      z' = f z