{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Polynomials where

import           Numeric.LinearAlgebra
import           System.Random hiding (rand)
import           Control.Monad.IO.Class
import           Control.Monad
import           Control.Exception
import           Data.Maybe
import qualified Data.Vector as V
import qualified Data.Array.Accelerate as A
import           Genetic
import           Formulas

import Debug.Trace

data Article = Article
            { _articlePowers  :: Powers
            , _articleInputs  :: Inplist
            , _articleOutputs :: Y
            } deriving (Eq, Ord)

instance Show Article where
      show = show . _articlePowers

instance Individual Article where
   -- fitness :: a -> IO (Maybe Double)
      fitness (Article ps is os) = do
            let x' = calcX is ps
                rows = length $ toLists x'
                cols = length . head $ toLists x'
                y = os
            noise_ <- rand rows cols
            let   noise = noise_ * 100.0
                  x = if det (tr x' <> x') == 0
                      then x' + noise
                      else x'
            -- bug fix for BLAS matrix rank on big numbers
            return $ if det (tr x <> x) == 0
                     then Nothing
                     else Just $ cost (normEq x y) x y

   -- mutateIndividual :: a -> IO a
      mutateIndividual (Article pows _a _b) = do
            r <- randomIO
            if mutationProbability < r
            then inner pows
            else return $ Article pows _a _b
            where
                  mutationProbability :: Double
                  mutationProbability = 0.1

                  inner :: Powers -> IO Article
                  inner x = do
                        blockIndex <- randomRIO (0, V.length x - 1)
                        powerIndex <- randomRIO (0, V.length (x V.! 0) - 1)
                        newPower   <- randomRIO (-4, 4)

                        let updated = updateVV x ( blockIndex, powerIndex) newPower
                        return $ Article updated _a _b

   -- crossover :: (a, a) -> IO a
      crossover ( Article pows1 _a _b
                , Article pows2 _ _
                ) = do
            index <- randomRIO (0, V.length pows1)
            let leftPart  = V.take index pows1
            let rightPart = V.drop index pows2
            return $ Article (V.concat [leftPart, rightPart]) _a _b

updateVV :: V.Vector (V.Vector a) -- what to update
         -> (Int, Int)            -- where
         -> a                     -- what to put there
         -> V.Vector (V.Vector a) -- result
updateVV vvs (xIdx, yIdx) value =
      vvs V.// [(xIdx, updatedV)]
      where
            updatedV  = originalV V.// [(yIdx, value)]
            originalV = vvs V.! xIdx


-------------------------------------------------------------------------------
-- |
-- Module    :  Environments.Bandits
-- Copyright :  (c) Sentenai 2017
-- License   :  BSD3
-- Maintainer:  sam@sentenai.com
-- Stability :  experimental
-- Portability: non-portable
--
-- Implementation of an n-armed bandit environment.
--
-- FIXME: currently this is only for a 10-armed bandit. This needs to be tied
-- to a config.
-------------------------------------------------------------------------------
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveGeneric #-}
{-# OPTIONS_GHC -Wno-unused-top-binds #-}
module Environments.Bandits
  ( Environment(..)
  , runEnvironment
  , Event.Event(..)
  , Config
  , Action
  , mkBandits
  , defaultBandits
  , mkAction
  ) where

import Control.MonadEnv
import Control.MonadMWCRandom
import qualified Data.Vector as V
import Data.Vector ((!), Vector)
import Data.Hashable
import qualified Data.Event as Event
import Control.Exception.Safe (assert, MonadThrow)
import qualified Statistics.Distribution as Dist
import Statistics.Distribution.Normal
import Data.DList
import Control.Monad.Reader.Class
import Control.Monad.RWS.Class
import Control.Monad.IO.Class
import Control.Monad.Writer.Class
import Control.Monad.State.Class
import Control.Monad.Trans.RWS (RWST, evalRWST)
import GHC.Generics

-- | FIXME: only 10 arms for the time being. This is where a "discrete space"
-- would be nice
data Config = Config
  { nBandits :: Int
  , offset   :: Int
  , stdDev   :: Float
  , bandits  :: Vector NormalDistribution
  , gen      :: GenIO
  }

instance Show Config where
  show c = "Config" ++
    "{nBandits="++ show (nBandits c)++
    "{offset="++ show (offset c)++
    "{bandit_stdDevs="++ show (stdDev c)++
    "{bandit_means="++ show (fmap Dist.mean . V.toList $ bandits c)++
    "}"

type Event = Event.Event Reward () Action

-- | The slot machine index whose arm will be pulled
newtype Action = Action { unAction :: Int }
  deriving (Eq, Ord, Show, Enum, Generic)

instance Bounded Action where
  minBound = Action 0
  maxBound = Action 9

instance Hashable Action where

-- | Convert an Int to an Action  in the bandit environment. Throw if the Int
-- falls out of bounds.
mkAction :: Int -> Environment Action
mkAction i = Environment $ do
  n <- nBandits <$> ask
  assert (i > n || i < 0) (pure $ Action i)

-- | Monad for an n-armed bandit environment
newtype Environment a = Environment
  { getEnvironment :: RWST Config (DList Event) () IO a }
  deriving
    ( Functor
    , Applicative
    , Monad
    , MonadIO
    , MonadThrow
    , MonadReader Config
    , MonadWriter (DList Event)
    , MonadState ()
    , MonadRWS Config (DList Event) ()
    )

-- | run an n-armed bandit environment
runEnvironment :: Config -> Environment () -> IO (DList Event)
runEnvironment c (Environment m) = snd <$> evalRWST m c ()

-- | Give the default config of a 10-armed bandit
defaultBandits :: GenIO -> Config
defaultBandits = mkBandits 10 2 0.1

-- | helper function to build a bandits config with normally-distributed
-- reward functions
mkBandits :: Int -> Int -> Float -> GenIO -> Config
mkBandits n offset' std = Config n offset' std $
  V.fromList $ fmap (`rewardDist` std) [offset' .. offset' + n - 1]
  where
    rewardDist :: Int -> Float -> NormalDistribution
    rewardDist m s = normalDistr (fromIntegral m) (realToFrac s)

instance MonadMWCRandom Environment where
  getGen = Environment $ fmap gen ask

instance MonadEnv Environment () Action Reward where
  -- this isn't an episodic environment... we'll have to split this out later
  reset :: Environment (Initial ())
  reset = return $ Initial ()

  step :: Action -> Environment (Obs Reward ())
  step (Action a) = do
    rwd <- genContVar =<< (! a) . bandits <$> ask
    tell . pure $ Event.Event 0 rwd () (Action a)
    return $ Next rwd ()


module Augmentation.DGA
    (DGA_Map (DGA_Map)
    ,Augmentation (Aug)
    ,applyDGAMap
    ,compose_maps
    ,Algebra
    ,fromDGAMap
    ) where

import Algebra
import Braid
import Data.Maybe
import Numeric.LinearAlgebra as N
import Data.List
import Debug.Trace

default (Int,Double)

data DGA_Map = DGA_Map [(Char,Algebra)] deriving Eq
instance Show DGA_Map where
    show (DGA_Map l) = "[" ++ (foldr (\(c,e) xs -> [c] ++ "→" ++ show e ++ (if xs /= "]" then "," else "") ++ xs) "]" l)

data Augmentation = Aug StdBraid [(Char,[Vector Z])]
instance Show Augmentation where
    show (Aug b m) = (show b) ++ "\n" ++ (concat $ map (\(c,vs) -> [c] ++ "→" ++ (foldr (\x xs -> show x ++ if xs == "" then "" else "+"++xs) "" vs) ++ " ") m)
instance Eq Augmentation where
    (Aug b1 m1) == (Aug b2 m2) = b1 == b2 && inZ2 m1 == inZ2 m2 && (m1 == m2 || eqh m1 m2)

inZ2 :: [(Char,[Vector Z])] -> [(Char,Z2)]
inZ2 = map (\(c,vs) -> (c,sum $ map (\_ -> fromInteger 1) vs))

isUpperTri :: Matrix Z -> Bool
isUpperTri mat = uptrih mat 0

uptrih :: Matrix Z -> Int -> Bool
uptrih mat i = let (l,w) = size mat
                   cond = and $ map (\j -> (==0) $ mat `atIndex` (j,i)) [i+1..l-1]
                in if i == w then True else cond && (uptrih mat $ i+1)

eqh :: [(Char,[Vector Z])] -> [(Char,[Vector Z])] -> Bool
eqh l1 l2 = maybe False id $ do
            { l1' <- mapM (\c -> lookup c l1) $ map fst l2
            ; l2' <- mapM (\(l,l') -> if length l == length l' then Just l' else Nothing) $ zip l1' $ map snd l2
            ; return True
            ; let dims1 = nub $ map size $ concat l1'
            ; let dims2 = nub $ map size $ concat l2'
            ; dim <- if length dims1 == 1 && length dims2 == 1 && dims1 == dims2 then Just $ head dims1 else Nothing
            ; let m2 = fromRows $ map (fromZ :: Vector Z -> Vector R) $ concat l2'
            ; let n0s = map (\l -> (0,(length $ permutations l)-1)) l1'
            ; let suc mns = if mns == [] then [] else if (fst $ head mns) == (snd $ head mns) then (0,snd $ head mns):(suc $ tail mns) else (1 + (fst $ head mns), snd $ head mns):(tail mns)
            ; let bound (acc,mns) = if mns == [] then acc else bound (((snd $ head mns)+1)*acc,tail mns)
            ; let ubound = bound (1,n0s)
            ; let mat mns = fromRows $ map (fromZ :: Vector Z -> Vector R) $ concat $ zipWith (\l (m,_) -> (permutations l) !! m) l1' mns
            ; let check m = maybe False id $ do
                            { let (l,u,p,s) = lu m
                            ; let (lR,lC) = size l
                            ; let sq = abs $ lR - lC
                            ; let l' = if lC == lR then l else if lC < lR then l ||| (konst 0 (lC,sq) === ident sq) else l === (konst 0 (sq,lR) ||| ident sq)
                            ; linv <- if det l' == 0 then Nothing else Just $ inv l'
                            ; pinv <- if s == 0 then Nothing else Just $ inv p
                            ; let mat' = linv N.<> pinv N.<> m2
                            ; let matz = fromColumns $ map (toZ . roundVector) $ toColumns mat'
                            ; let cond = isUpperTri matz
                            ; let cond' = and $ map (\x -> (x - (fromIntegral $ floor x) < cutoff) || (((fromIntegral $ ceiling x) - x) < cutoff)) $ toList $ flatten mat'
                            ; return $ cond && cond'
                            }
            ; let checkAll k mns = if size (mat mns) /= size m2 then trace "Dim mismatch" False else if check $ mat mns then trace "Found one!" True else if k > ubound then trace ("Tried: "++show k) False else checkAll (k+1) (suc mns)
            ; return $ checkAll 0 n0s
            }

fromDGAMap :: StdBraid -> DGA_Map -> [Char] -> Maybe Augmentation
fromDGAMap b (DGA_Map l) chars = do
                                { l' <- mapM (\(c,a) -> (represent chars a) >>= (\vs -> return (c,vs))) l
                                ; return $ Aug b l'
                                }

compose_maps :: DGA_Map -> DGA_Map -> DGA_Map
compose_maps (DGA_Map map1) (DGA_Map map2) = DGA_Map $ (map (\(c,exp) -> (c,applyDGAMap (DGA_Map map2) exp)) map1) ++ (filter (\(c,_) -> not $ elem c $ map fst map1) map2)

applyDGAMap :: DGA_Map -> Algebra -> Algebra
applyDGAMap (DGA_Map alist) a = appmaph alist a

appmaph::[(Char,Algebra)] -> Algebra -> Algebra
appmaph [] = id
appmaph cs = plugIn (\c -> case (lookup c cs) of Just e -> e
                                                 Nothing -> G $ E c)


{-# LANGUAGE ViewPatterns #-}

module Main where

import Codec.Image.DevIL
import Data.Array.Unboxed
import Math.Probably.MCMC
import Math.Probably.RandIO
import Math.Probably.Sampler
import Math.Probably.FoldingStats
import Control.Monad.State.Strict
import System.Cmd
import System.Environment
import Data.Array.IO
import System.IO

import Data.Array.Unboxed
import Numeric.LinearAlgebra hiding (find)
import qualified Math.Probably.PDF as PDF
import Data.List
import Data.Maybe
import Control.Applicative
import Data.Ord

import Data.IORef
import System.IO.Unsafe

import CVUtils
import Edge

posteriorV :: Image -> Image -> (Int, Int) -> Vector R -> R
posteriorV bgim im (cx,cy) v =
   uniformLogPdf 0 100.0 len1 +
   uniformLogPdf 0 1.0 ecc +
   uniformLogPdf 0 10000.0 noise +
   uniformLogPdf 0 1500.0 px +
   uniformLogPdf 0 1500.0 py +
   sum [ f x y chan |
                   x <- [ (cx::Int) -20.. cx +20],
                   y <- [ cy -20.. cy +20],
                   chan <- [0..2]]

  where px = v @> 0
        py = v @> 1
        noise = v @> 2
        len1 = v@> 3
        rot = v@> 4
        ecc = v@> 5
        red = v@> 6
        green = v@> 7
        blue = v@> 8
        noiseInside = v@> 9
        len2 = v@> 10
        f1x = px+(len1*ecc)*cos rot
        f1y = py+(len1*ecc)*sin rot
        f2x = px-(len1*ecc)*cos rot
        f2y = py-(len1*ecc)*sin rot
        f1x2 = px+(len2*ecc)*cos rot
        f1y2 = py+(len2*ecc)*sin rot
        f2x2 = px-(len2*ecc)*cos rot
        f2y2 = py-(len2*ecc)*sin rot
        colVec = fromList [red,green,blue]
        f :: Int -> Int -> Int -> R
        f x y ch
         = if dist  f1x f1y   x  y  + dist  f2x f2y   x  y  < 2 * len1
              then gaussR noiseInside (colVec@>ch) $ im!(y,x,ch)
              else if dist  f1x2 f1y2   x  y  + dist  f2x2 f2y2   x  y  < 2 * len2
                      then gaussW8 noise ((bgim!(y,x,ch))`div` 2) $ im!(y,x,ch)
                      else gaussW8 noise (bgim!(y,x,ch)) $ im!(y,x,ch)

main = do
     ilInit
     bgnm : fvid :_ <- getArgs
     bgIm <-readImage bgnm
     system $ "~/cvutils/extract "++fvid++" 3200"
     frame0 <-readImage "extract.png"
     let x = 1092
         y = 518
         rot = -5.12
         posterior = posteriorV bgIm frame0 (round x,round y)
         postAndV v = (posteriorV bgIm frame0 (round x,round y) $ fromList v, v)

         initialsV = fromList [x,y,218, 4.5, rot, 0.9, 0.1,0.1,0.1, 218, 6]
     print [frame0!(round y,round x,c) | c <- [0..2]]
     print [frame0!(round y+1,round x,c) | c <- [0..2]]
     print [frame0!(round y,round x+1,c) | c <- [0..2]]
     runRIO $ do
         iniampar <- sample $ initialAdaMet 300 1.5e-4 posterior initialsV
         AMPar v _ _ _ _ _ _ <- runAndDiscard 50000 (show . ampPar) iniampar $ adaMet False posterior
         lift $ print initialsV
         lift $ print v


import Control.Applicative
import Data.Generics.Random.Boltzmann.GeneratingFunctions
import Numeric.LinearAlgebra
import Test.HUnit

main = runTestTT $ TestList
  [ test_solve
  ]

test_solve = "F(x) = 1 + x F(x)^2" ~: do
    print xs
    print (evalDeltas <$> xs <*> pure es)
  where
    x : fx : f'x : _ = fmap X [0 ..]
    expectedSize x n = (n * fx, x * f'x)
    es =
      [ expectedSize x 10000
      , (fx, 1 + x * fx * fx)
      , (f'x, fx * fx + 2 * x * f'x * fx)
      ]
    xs = solveEquations defSolveArgs es (vector [0, 1, 1])


-- | Example originally inspired by
-- <https://ro-che.info/articles/2015-12-05-testing-fft Roman Cheplyaka>.

{-# LANGUAGE GADTs #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns #-}

module Main where

import Data.Complex
import qualified Language.R as R
import Language.R (R)
import Language.R.QQ

-- Call R's FFT
r_fft :: [Complex Double] -> R s [Complex Double]
r_fft nums = do
    R.dynSEXP <$> [r| fft(nums_hs) |]

main :: IO ()
main = R.withEmbeddedR R.defaultConfig $ do
    result <- R.runRegion $ r_fft [1,2,1]
    print result


{-
   Copyright 2016, Dominic Orchard, Andrew Rice, Mistral Contrastin, Matthew Danish

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
-}

{-
  Units of measure extension to Fortran: backend
-}

{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Camfort.Specification.Units.InferenceBackend
  ( chooseImplicitNames
  , criticalVariables
  , inconsistentConstraints
  , inferVariables
  -- mainly for debugging and testing:
  , shiftTerms
  , flattenConstraints
  , flattenUnits
  , constraintsToMatrix
  , constraintsToMatrices
  , rref
  , genUnitAssignments
  , genUnitAssignments'
  , provenance
  , splitNormHNF
  ) where

import           Camfort.Specification.Units.Environment
import qualified Camfort.Specification.Units.InferenceBackendFlint as Flint
import           Control.Arrow (first, second, (***))
import           Control.Monad
import           Control.Monad.ST
import           Control.Parallel.Strategies
import qualified Data.Array as A
import           Data.Generics.Uniplate.Operations
  (transformBi, universeBi)
import           Data.Graph.Inductive hiding ((><))
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import           Data.List
  ((\\), findIndex, inits, nub, partition, sort, sortBy, group, tails, foldl')
import qualified Data.Map.Strict as M
import           Data.Maybe (fromMaybe, mapMaybe)
import           Data.Ord
import           Data.Tuple (swap)
import           Numeric.LinearAlgebra
  ( atIndex, (<>)
  , rank, (?), (¿)
  , rows, cols
  , subMatrix, diag
  , fromBlocks, ident
  )
import qualified Numeric.LinearAlgebra as H
import           Numeric.LinearAlgebra.Devel
  ( newMatrix, readMatrix
  , writeMatrix, runSTMatrix
  , freezeMatrix, STMatrix
  )
import           Prelude hiding ((<>))

-- | Returns list of formerly-undetermined variables and their units.
inferVariables :: Constraints -> [(VV, UnitInfo)]
inferVariables cons = unitVarAssignments
  where
    unitAssignments = genUnitAssignments cons
    -- Find the rows corresponding to the distilled "unit :: var"
    -- information for ordinary (non-polymorphic) variables.
    unitVarAssignments            =
      [ (var, units) | ([UnitPow (UnitVar var)                 k], units) <- unitAssignments, k `approxEq` 1 ] ++
      [ (var, units) | ([UnitPow (UnitParamVarAbs (_, var)) k], units)    <- unitAssignments, k `approxEq` 1 ]

-- Detect inconsistency if concrete units are assigned an implicit
-- abstract unit variable with coefficients not equal, or there are
-- monomorphic literals being given parametric polymorphic units.
detectInconsistency :: [([UnitInfo], UnitInfo)] -> Constraints
detectInconsistency unitAssignments = inconsist
  where
    ua' = map (shiftTerms . fmap flattenUnits) unitAssignments
    badImplicits = [ fmap foldUnits a | a@([UnitPow (UnitParamImpAbs _) k1], rhs) <- ua'
                                      , UnitPow _ k2 <- rhs
                                      , k1 /= k2 ]
    inconsist = unitAssignmentsToConstraints badImplicits ++ mustBeUnitless unitAssignments

-- Must be unitless: any assignments of parametric abstract units to
-- monomorphic literals.
mustBeUnitless :: [([UnitInfo], UnitInfo)] -> Constraints
mustBeUnitless unitAssignments = mbu
  where
    mbu = [ ConEq UnitlessLit (UnitPow (UnitLiteral l) k)
          | (UnitPow (UnitLiteral l) k:_, rhs) <- ua''
          , any isParametric (universeBi rhs :: [UnitInfo]) ]
    -- ua' = map (shiftTerms . fmap flattenUnits) unitAssignments
    ua'' = map (shiftTermsBy isLiteral . fmap flattenUnits) unitAssignments

    isLiteral UnitLiteral{} = True
    isLiteral (UnitPow UnitLiteral{} _) = True
    isLiteral _ = False

    isParametric UnitParamVarAbs{} = True
    isParametric UnitParamPosAbs{} = True
    isParametric UnitParamEAPAbs{} = True
    isParametric UnitParamLitAbs{} = True
    isParametric UnitParamImpAbs{} = True
    isParametric (UnitPow u _)     = isParametric u
    isParametric _                 = False

-- convert the assignment format back into constraints
unitAssignmentsToConstraints :: [([UnitInfo], UnitInfo)] -> Constraints
unitAssignmentsToConstraints = map (uncurry ConEq . first foldUnits)

-- | Raw units-assignment pairs.
genUnitAssignments :: Constraints -> [([UnitInfo], UnitInfo)]
genUnitAssignments cons
  -- if the results include any mappings that must be forced to be unitless...
  | mbu <- mustBeUnitless ua, not (null mbu) = genUnitAssignments (mbu ++ unitAssignmentsToConstraints ua)
  | null (detectInconsistency ua)            = ua
  | otherwise                                = []
  where
    ua = genUnitAssignments' colSort cons

-- | Break up the problem of solving normHNF on each group of related
-- columns, then bring it all back together.
splitNormHNF :: H.Matrix Double -> (H.Matrix Double, [Int])
splitNormHNF unsolvedM = (combinedMat, allNewColIndices)
  where
    combinedMat      = joinMat (map (first fst) solvedMs)
    allNewColIndices = concatMap (snd . fst) solvedMs

    inParallel = (`using` parTuple2 (parList rseq) rseq)
    (solvedMs, _) = inParallel . foldl' eachResult ([], cols unsolvedM) $ map (first Flint.normHNF) (splitMat unsolvedM)

    -- for each result re-number the generated columns & add mappings for each.
    eachResult (ms, startI) ((m, newColIndices), origCols) = (((m, newColIndices'), origCols'):ms, endI)
      where
        -- produce (length newColIndices) number of mappings
        endI           = startI + length newColIndices
        -- re-number the newColIndices according to the lookup list
        newColIndices' = map (origCols !!) newColIndices
        -- add columns in the (combined) matrix for the newly
        -- generated columns from running normHNF on m.
        origCols'      = origCols ++ [startI .. endI-1]

genUnitAssignments' :: SortFn -> Constraints -> [([UnitInfo], UnitInfo)]
genUnitAssignments' _ [] = []
genUnitAssignments' sortfn cons
  | null colList                                      = []
  | null inconsists                                   = unitAssignments
  | otherwise                                         = []
  where
    (lhsM, rhsM, inconsists, lhsColA, rhsColA) = constraintsToMatrices' sortfn cons
    unsolvedM | rows rhsM == 0 || cols rhsM == 0 = lhsM
              | rows lhsM == 0 || cols lhsM == 0 = rhsM
              | otherwise                        = fromBlocks [[lhsM, rhsM]]
    (solvedM, newColIndices)      = splitNormHNF unsolvedM
    -- solvedM can have additional columns and rows from normHNF;
    -- cosolvedM corresponds to the original lhsM.
    -- cosolvedM                     = subMatrix (0, 0) (rows solvedM, cols lhsM) solvedM
    -- cosolvedMrhs                  = subMatrix (0, cols lhsM) (rows solvedM, cols solvedM - cols lhsM) solvedM

    -- generate a colList with both the original columns and new ones generated
    -- if a new column generated was derived from the right-hand side then negate it
    numLhsCols                    = 1 + snd (A.bounds lhsColA)
    colList                       = map (1,) (A.elems lhsColA ++ A.elems rhsColA) ++ map genC newColIndices
    genC n | n >= numLhsCols      = (-k, UnitParamImpAbs (show u))
           | otherwise            = (k, UnitParamImpAbs (show u))
      where (k, u) = colList !! n
    -- Convert the rows of the solved matrix into flattened unit
    -- expressions in the form of "unit ** k".
    unitPow (k, u) x              = UnitPow u (k * x)
    unitPows                      = map (concatMap flattenUnits . zipWith unitPow colList) (H.toLists solvedM)

    -- Variables to the left, unit names to the right side of the equation.
    unitAssignments               = map (fmap (foldUnits . map negatePosAbs) . checkSanity . partition (not . isUnitRHS')) unitPows
    isUnitRHS' (UnitPow (UnitName _) _)        = True
    isUnitRHS' (UnitPow (UnitParamEAPAbs _) _) = True
    -- Because this version of isUnitRHS different from
    -- constraintsToMatrix interpretation, we need to ensure that any
    -- moved ParamPosAbs units are negated, because they are
    -- effectively being shifted across the equal-sign:
    isUnitRHS' (UnitPow (UnitParamImpAbs _) _) = True
    isUnitRHS' (UnitPow (UnitParamPosAbs (_, 0)) _) = False
    isUnitRHS' (UnitPow (UnitParamPosAbs _) _) = True
    isUnitRHS' _                               = False

checkSanity :: ([UnitInfo], [UnitInfo]) -> ([UnitInfo], [UnitInfo])
checkSanity (u1@[UnitPow (UnitVar _) _], u2)
  | or $ [ True | UnitParamPosAbs (_, _) <- universeBi u2 ]
      ++ [ True | UnitParamImpAbs _      <- universeBi u2 ] = (u1++u2,[])
checkSanity (u1@[UnitPow (UnitParamVarAbs (f, _)) _], u2)
  | or [ True | UnitParamPosAbs (f', _) <- universeBi u2, f' /= f ] = (u1++u2,[])
checkSanity c = c

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

-- FIXME: you know better...
approxEq :: Double -> Double -> Bool
approxEq a b = abs (b - a) < epsilon
notApproxEq :: Double -> Double -> Bool
notApproxEq a b = not (approxEq a b)
epsilon :: Double
epsilon = 0.001 -- arbitrary

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

type RowNum = Int               -- ^ 'row number' of matrix
type ColNum = Int               -- ^ 'column number' of matrix
-- | Represents a subproblem of AX=B where the row numbers and column
-- numbers help you re-map back to the original problem.
type Subproblem = ([RowNum], (H.Matrix Double, H.Matrix Double), [ColNum])

-- | Divide up the AX=B problem into smaller problems based on the
-- 'related columns' and their corresponding rows in the
-- right-hand-side of the equation. Where lhsM = A and rhsM = B.  The
-- resulting list of subproblems contains the new, smaller As and Bs
-- as well as a list of original row numbers and column numbers to
-- aide re-mapping back to the original lhsM and rhsM.
splitMatWithRHS :: H.Matrix Double -> H.Matrix Double -> [Subproblem]
splitMatWithRHS lhsM rhsM | cols lhsM > 0 = map (eachComponent . sort) $ scc (relatedColumnsGraph lhsM)
                          | otherwise     = []
  where
    -- Gets called on every strongly-connected component / related set of columns.
    eachComponent cs = (rs, mats, cs)
      where
        -- Selected columns
        lhsSelCols :: H.Matrix Double
        lhsSelCols = lhsM ¿ cs

        csLen = cols lhsSelCols

        -- Find the row numbers of the 'all zero' rows in lhsM.
        lhsAllZeroRows :: [RowNum]
        lhsAllZeroRows = map fst . filter (all (approxEq 0) . snd) . zip [0..] $ H.toLists lhsM

        -- Find the row numbers that correspond to the non-zero co-efficients in the selected columns.
        lhsNonZeroColRows :: [(RowNum, [Double])]
        lhsNonZeroColRows = filter (any (notApproxEq 0) . snd) . zip [0..] . H.toLists $ lhsSelCols

        -- List of all the row numbers and row values combined from the two above variables.
        lhsNumberedRows :: [(RowNum, [Double])]
        lhsNumberedRows = sortBy (comparing fst) $ lhsNonZeroColRows ++ zip lhsAllZeroRows (repeat (replicate csLen 0))

        -- For each of the above LHS rows find a corresponding RHS row.
        rhsSelRows :: [[Double]]
        rhsSelRows | rows rhsM > 0 = H.toLists (rhsM ? map fst lhsNumberedRows)
                   | otherwise     = []

        reassoc (a, b) c = (a, (b, c))

        notAllZero (_, (lhs, rhs)) = any (notApproxEq 0) (lhs ++ rhs)

        -- Zip the selected LHS, RHS rows together, filter out any that are all zeroes.
        numberedRows :: ([RowNum], [([Double], [Double])])
        numberedRows = unzip . filter notAllZero $ zipWith reassoc lhsNumberedRows rhsSelRows

        rs :: [RowNum]          -- list of row numbers in the subproblem
        mats :: (H.Matrix Double, H.Matrix Double) -- LHS/RHS subproblem matrices
        (rs, mats) = second ((H.fromLists *** H.fromLists) . unzip) numberedRows

-- | Split the lhsM/rhsM problem into subproblems and then look for
-- inconsistent rows in each subproblem, concatenating all of the
-- inconsistent row numbers found (in terms of the rows of the
-- original lhsM).
splitFindInconsistentRows :: H.Matrix Double -> H.Matrix Double -> [RowNum]
splitFindInconsistentRows lhsMat rhsMat = concatMap eachComponent $ splitMatWithRHS lhsMat rhsMat
  where
    eachComponent (rs, (lhsM, rhsM), _) = map (rs !!) $ findInconsistentRows lhsM augM
      where
        -- Augmented matrix is defined as the combined LHS/RHS matrices.
        augM
          | rows rhsM == 0 || cols rhsM == 0 = lhsM
          | rows lhsM == 0 || cols lhsM == 0 = rhsM
          | otherwise = fromBlocks [[lhsM, rhsM]]

-- | Break out the 'unrelated' columns in a single matrix into
-- separate matrices, along with a list of their original column
-- positions.
splitMat :: H.Matrix Double -> [(H.Matrix Double, [ColNum])]
splitMat m = map (eachComponent . sort) $ scc (relatedColumnsGraph m)
  where
    eachComponent cs = (H.fromLists . filter (any (/= 0)) . H.toLists $ m ¿ cs, cs)

-- | Bring together the split matrices and put the columns back in
-- their original order. Rows may not be in the same order as the
-- original, but the constraints should be equivalent.
joinMat :: [(H.Matrix Double, [Int])] -> H.Matrix Double
joinMat ms = sortedM
  where
    disorderedM = H.diagBlock (map fst ms)
    colsWithIdx = zip (concatMap snd ms) . H.toColumns $ disorderedM
    sortedM     = H.fromColumns . map snd . sortBy (comparing fst) $ colsWithIdx

-- | Turn a matrix into a graph where each node represents a column
-- and each edge represents two columns that have non-zero
-- co-efficients in some row. Basically, 'related columns'. Also
-- includes self-refs for each node..
relatedColumnsGraph :: H.Matrix Double -> Gr () ()
relatedColumnsGraph m = mkGraph (map (,()) ns) (map (\ (a,b) -> (a,b,())) es)
  where
    nonZeroCols = [ [ j | j <- [0..cols m - 1], not (m `atIndex` (i, j) `approxEq` 0) ] | i <- [0..rows m - 1] ]
    ns          = nub $ concat nonZeroCols
    es          = [ (i, j) | cs <- nonZeroCols, [i, j] <- sequence [cs, cs] ]

-- Convert a set of constraints into a matrix of co-efficients, and a
-- reverse mapping of column numbers to units.
constraintsToMatrix :: Constraints -> (H.Matrix Double, [Int], A.Array Int UnitInfo)
constraintsToMatrix cons
  | all null lhs = (H.ident 0, [], A.listArray (0, -1) [])
  | otherwise = (augM, inconsists, A.listArray (0, length colElems - 1) colElems)
  where
    -- convert each constraint into the form (lhs, rhs)
    consPairs       = filter (uncurry (/=)) $ flattenConstraints cons
    -- ensure terms are on the correct side of the equal sign
    shiftedCons     = map shiftTerms consPairs
    lhs             = map fst shiftedCons
    rhs             = map snd shiftedCons
    (lhsM, lhsCols) = flattenedToMatrix colSort lhs
    (rhsM, rhsCols) = flattenedToMatrix colSort rhs
    colElems        = A.elems lhsCols ++ A.elems rhsCols
    augM            = if rows rhsM == 0 || cols rhsM == 0 then lhsM else if rows lhsM == 0 || cols lhsM == 0 then rhsM else fromBlocks [[lhsM, rhsM]]
    inconsists      = splitFindInconsistentRows lhsM rhsM

constraintsToMatrices :: Constraints -> (H.Matrix Double, H.Matrix Double, [Int], A.Array Int UnitInfo, A.Array Int UnitInfo)
constraintsToMatrices cons = constraintsToMatrices' colSort cons

constraintsToMatrices' :: SortFn -> Constraints -> (H.Matrix Double, H.Matrix Double, [Int], A.Array Int UnitInfo, A.Array Int UnitInfo)
constraintsToMatrices' sortfn cons
  | all null lhs = (H.ident 0, H.ident 0, [], A.listArray (0, -1) [], A.listArray (0, -1) [])
  | otherwise = (lhsM, rhsM, inconsists, lhsCols, rhsCols)
  where
    -- convert each constraint into the form (lhs, rhs)
    consPairs       = filter (uncurry (/=)) $ flattenConstraints cons
    -- ensure terms are on the correct side of the equal sign
    shiftedCons     = map shiftTerms consPairs
    lhs             = map fst shiftedCons
    rhs             = map snd shiftedCons
    (lhsM, lhsCols) = flattenedToMatrix sortfn lhs
    (rhsM, rhsCols) = flattenedToMatrix sortfn rhs
    inconsists      = splitFindInconsistentRows lhsM rhsM

-- [[UnitInfo]] is a list of flattened constraints
flattenedToMatrix :: SortFn -> [[UnitInfo]] -> (H.Matrix Double, A.Array Int UnitInfo)
flattenedToMatrix sortfn cons = (m, A.array (0, numCols - 1) (map swap uniqUnits))
  where
    m = runSTMatrix $ do
          newM <- newMatrix 0 numRows numCols
          -- loop through all constraints
          forM_ (zip cons [0..]) $ \ (unitPows, row) -> do
            -- write co-efficients for the lhs of the constraint
            forM_ unitPows $ \ (UnitPow u k) -> do
              case M.lookup u colMap of
                Just col -> readMatrix newM row col >>= (writeMatrix newM row col . (+k))
                _        -> return ()
          return newM
    -- identify and enumerate every unit uniquely
    uniqUnits = flip zip [0..] . map head . group . sortBy sortfn $ [ u | UnitPow u _ <- concat cons ]
    -- map units to their unique column number
    colMap    = M.fromList uniqUnits
    numRows   = length cons
    numCols   = M.size colMap

negateCons :: [UnitInfo] -> [UnitInfo]
negateCons = map (\ (UnitPow u k) -> UnitPow u (-k))

negatePosAbs :: UnitInfo -> UnitInfo
negatePosAbs (UnitPow (UnitParamPosAbs x) k) = UnitPow (UnitParamPosAbs x) (-k)
negatePosAbs (UnitPow (UnitParamImpAbs v) k) = UnitPow (UnitParamImpAbs v) (-k)
negatePosAbs u                               = u

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

-- Units that should appear on the right-hand-side of the matrix during solving
isUnitRHS :: UnitInfo -> Bool
isUnitRHS (UnitPow (UnitName _) _)        = True
isUnitRHS (UnitPow (UnitParamEAPAbs _) _) = True
isUnitRHS _                               = False

-- | Shift UnitNames/EAPAbs poly units to the RHS, and all else to the LHS.
shiftTerms :: ([UnitInfo], [UnitInfo]) -> ([UnitInfo], [UnitInfo])
shiftTerms (lhs, rhs) = (lhsOk ++ negateCons rhsShift, rhsOk ++ negateCons lhsShift)
  where
    (lhsOk, lhsShift) = partition (not . isUnitRHS) lhs
    (rhsOk, rhsShift) = partition isUnitRHS rhs

-- | Shift terms based on function f (<- True, False ->).
shiftTermsBy :: (UnitInfo -> Bool) -> ([UnitInfo], [UnitInfo]) -> ([UnitInfo], [UnitInfo])
shiftTermsBy f (lhs, rhs) = (lhsOk ++ negateCons rhsShift, rhsOk ++ negateCons lhsShift)
  where
    (lhsOk, lhsShift) = partition f lhs
    (rhsOk, rhsShift) = partition (not . f) rhs

-- | Translate all constraints into a LHS, RHS side of units.
flattenConstraints :: Constraints -> [([UnitInfo], [UnitInfo])]
flattenConstraints = map (\ (ConEq u1 u2) -> (flattenUnits u1, flattenUnits u2))

--------------------------------------------------
-- Matrix solving functions based on HMatrix

-- | Returns given matrix transformed into Reduced Row Echelon Form
rref :: H.Matrix Double -> H.Matrix Double
rref a = snd $ rrefMatrices' a 0 0 []
  where
    -- (a', den, r) = Flint.rref a

-- Provenance of matrices.
data RRefOp
  = ElemRowSwap Int Int         -- ^ swapped row with row
  | ElemRowMult Int Double      -- ^ scaled row by constant
  | ElemRowAdds [(Int, Int)]    -- ^ set of added row onto row ops
  deriving (Show, Eq, Ord)

-- worker function
-- invariant: the matrix a is in rref except within the submatrix (j-k,j) to (n,n)
rrefMatrices' :: H.Matrix Double -> Int -> Int -> [(H.Matrix Double, RRefOp)] ->
                 ([(H.Matrix Double, RRefOp)], H.Matrix Double)
rrefMatrices' a j k mats
  -- Base cases:
  | j - k == n            = (mats, a)
  | j     == m            = (mats, a)

  -- When we haven't yet found the first non-zero number in the row, but we really need one:
  | a `atIndex` (j - k, j) == 0 = case findIndex (/= 0) below of
    -- this column is all 0s below current row, must move onto the next column
    Nothing -> rrefMatrices' a (j + 1) (k + 1) mats
    -- we've found a row that has a non-zero element that can be swapped into this row
    Just i' -> rrefMatrices' (swapMat <> a) j k ((swapMat, ElemRowSwap i (j - k)):mats)
      where i       = j - k + i'
            swapMat = elemRowSwap n i (j - k)

  -- We have found a non-zero cell at (j - k, j), so transform it into
  -- a 1 if needed using elemRowMult, and then clear out any lingering
  -- non-zero values that might appear in the same column, using
  -- elemRowAdd:
  | otherwise             = rrefMatrices' a2 (j + 1) k mats2
  where
    n     = rows a
    m     = cols a
    below = getColumnBelow a (j - k, j)
    scale = recip (a `atIndex` (j - k, j))
    erm   = elemRowMult n (j - k) scale

    -- scale the row if the cell is not already equal to 1
    (a1, mats1) | a `atIndex` (j - k, j) /= 1 = (erm <> a, (erm, ElemRowMult (j - k) scale):mats)
                | otherwise                   = (a, mats)

    -- Locate any non-zero values in the same column as (j - k, j) and
    -- cancel them out. Optimisation: instead of constructing a
    -- separate elemRowAdd matrix for each cancellation that are then
    -- multiplied together, simply build a single matrix that cancels
    -- all of them out at the same time, using the ST Monad.
    findAdds _ curM ms
      | isWritten = (newMat <> curM, (newMat, ElemRowAdds matOps):ms)
      | otherwise = (curM, ms)
      where
        (isWritten, matOps, newMat) = runST $ do
          newM <- newMatrix 0 n n :: ST s (STMatrix s Double)
          sequence_ [ writeMatrix newM i' i' 1 | i' <- [0 .. (n - 1)] ]
          let f w o i | i >= n                  = return (w, o)
                      | i == j - k              = f w o (i + 1)
                      | a `atIndex` (i, j) == 0 = f w o (i + 1)
                      | otherwise               = writeMatrix newM i (j - k) (- (a `atIndex` (i, j)))
                                                  >> f True ((i, j - k):o) (i + 1)
          (isW, ops) <- f False [] 0
          (isW, ops,) `fmap` freezeMatrix newM

    (a2, mats2) = findAdds (0::Int) a1 mats1

-- Get a list of values that occur below (i, j) in the matrix a.
getColumnBelow :: H.Matrix Double -> (Int, Int) -> [Double]
getColumnBelow a (i, j) = concat . H.toLists $ subMatrix (i, j) (n - i, 1) a
  where n = rows a

-- 'Elementary row operation' matrices
elemRowMult :: Int -> Int -> Double -> H.Matrix Double
elemRowMult n i k = diag (H.fromList (replicate i 1.0 ++ [k] ++ replicate (n - i - 1) 1.0))

elemRowSwap :: Int -> Int -> Int -> H.Matrix Double
elemRowSwap n i j
  | i == j          = ident n
  | i > j           = elemRowSwap n j i
  | otherwise       = ident n ? ([0..i-1] ++ [j] ++ [i+1..j-1] ++ [i] ++ [j+1..n-1])

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

type GraphCol = IM.IntMap IS.IntSet   -- graph from origin to dest.
type Provenance = IM.IntMap IS.IntSet -- graph from dest. to origin

opToGraphCol :: RRefOp -> GraphCol
opToGraphCol ElemRowMult{} = IM.empty
opToGraphCol (ElemRowSwap i j) = IM.fromList [ (i, IS.singleton j), (j, IS.singleton i) ]
opToGraphCol (ElemRowAdds l)   = IM.fromList $ concat [ [(i, IS.fromList [i,j]), (j, IS.singleton j)]  | (i, j) <- l ]

graphColCombine :: GraphCol -> GraphCol -> GraphCol
graphColCombine g1 g2 = IM.unionWith (curry snd) g1 $ IM.map (IS.fromList . trans . IS.toList) g2
  where
    trans = concatMap (\ i -> [i] `fromMaybe` (IS.toList <$> IM.lookup i g1))

invertGraphCol :: GraphCol -> GraphCol
invertGraphCol g = IM.fromListWith IS.union [ (i, IS.singleton j) | (j, jset) <- IM.toList g, i <- IS.toList jset ]

provenance :: H.Matrix Double -> (H.Matrix Double, Provenance)
provenance m = (m', p)
  where
    (matOps, m') = rrefMatrices' m 0 0 []
    p = invertGraphCol . foldl' graphColCombine IM.empty . map opToGraphCol $ map snd matOps

-- Worker functions:

findInconsistentRows :: H.Matrix Double -> H.Matrix Double -> [Int]
findInconsistentRows coA augA | rows augA < 2 = []
                              | otherwise     = inconsistent
  where
    inconsistent = [0..(rows augA - 1)] \\ consistent

    consistent
      -- if the space is relatively small, try it all
      | rows augA < 16 = head (filter tryRows (powerset $ reverse [0..(rows augA - 1)]))
      | otherwise = head (filter tryRows (tails ( [0..(rows augA - 1)])) ++ [[]])

    powerset = filterM (const [True, False])

    -- Rouché–Capelli theorem is that if the rank of the coefficient
    -- matrix is not equal to the rank of the augmented matrix then
    -- the system of linear equations is inconsistent.
    tryRows [] = True
    tryRows ns = (rank coA' == rank augA')
      where
        coA'  = coA ? ns
        augA' = augA ? ns

-- | Create unique names for all of the inferred implicit polymorphic
-- unit variables.
chooseImplicitNames :: [(VV, UnitInfo)] -> [(VV, UnitInfo)]
chooseImplicitNames vars = replaceImplicitNames (genImplicitNamesMap vars) vars

genImplicitNamesMap :: Data a => a -> M.Map UnitInfo UnitInfo
genImplicitNamesMap x = M.fromList [ (absU, UnitParamEAPAbs (newN, newN)) | (absU, newN) <- zip absUnits newNames ]
  where
    absUnits = nub [ u | u@(UnitParamPosAbs _)             <- universeBi x ] ++
               nub [ u | u@(UnitParamImpAbs _)             <- universeBi x ]
    eapNames = nub $ [ n | (UnitParamEAPAbs (_, n))      <- universeBi x ] ++
                     [ n | (UnitParamEAPUse ((_, n), _)) <- universeBi x ]
    newNames = filter (`notElem` eapNames) . map ('\'':) $ nameGen
    nameGen  = concatMap sequence . tail . inits $ repeat ['a'..'z']

replaceImplicitNames :: Data a => M.Map UnitInfo UnitInfo -> a -> a
replaceImplicitNames implicitMap = transformBi replace
  where
    replace u@(UnitParamPosAbs _) = fromMaybe u $ M.lookup u implicitMap
    replace u@(UnitParamImpAbs _) = fromMaybe u $ M.lookup u implicitMap
    replace u                     = u

-- | Identifies the variables that need to be annotated in order for
-- inference or checking to work.
criticalVariables :: Constraints -> [UnitInfo]
criticalVariables [] = []
criticalVariables cons = filter (not . isUnitRHS') $ map (colA A.!) criticalIndices
  where
    (unsolvedM, _, colA)          = constraintsToMatrix cons
    solvedM                       = rref unsolvedM
    uncriticalIndices             = mapMaybe (findIndex (/= 0)) $ H.toLists solvedM
    criticalIndices               = A.indices colA \\ uncriticalIndices
    isUnitRHS' (UnitName _)       = True; isUnitRHS' _ = False

-- | Returns just the list of constraints that were identified as
-- being possible candidates for inconsistency, if there is a problem.
inconsistentConstraints :: Constraints -> Maybe Constraints
inconsistentConstraints [] = Nothing
inconsistentConstraints cons
  | not (null direct) = Just direct
  | null inconsists   = Nothing
  | otherwise         = Just [ con | (con, i) <- zip cons [0..], i `elem` inconsists ]
  where
    (_, _, inconsists, _, _) = constraintsToMatrices cons
    direct = detectInconsistency $ genUnitAssignments' colSort cons


{-# LANGUAGE BangPatterns #-}
-----------------------------------------------------------------------------
--
-- Module      :  AI.Network.RNN.Util
-- Copyright   :  (c) JP Moresmau
-- License     :  BSD3
--
-- Maintainer  :  JP Moresmau <jp@moresmau.fr>
-- Stability   :  experimental
-- Portability :
--
-- | Utility functions
--
-----------------------------------------------------------------------------

module AI.Network.RNN.Util where

import Control.Parallel.Strategies
import Control.Monad.Random

import Numeric.LinearAlgebra.HMatrix

import Data.List
import Data.Ord
import qualified Data.Set as S

-- | Calculate the product of a matrix by a vector, with both represented by a list
listMProd :: (Num a) => [a] -> [a] -> [a]
listMProd mdt vdt = go mdt vdt 0
  where
    go [] _  !s = [s]
    go ls [] !s = s : go ls vdt 0
    go (y:ys) (x:xs) !ix = go ys xs (y*x+ix)
--    map (foldr (\(a,b) c->a*b+c) 0) $ go mdt vdt []
--  where
--    go [] _  s = [s]
--    go ls [] s = s : go ls vdt []
--    go (y:ys) (x:xs) acc = go ys xs ((y,x):acc)
--listMProd mdt vdt = let (_,l,lst)= foldl' go (vdt,[],0) mdt
--    in reverse $ lst:l
--    where
--        go ([],acc,ix) m =
--            let (x:xs)=vdt
--            in (xs,ix:acc,x*m)
--        go ((x:xs),acc,ix) m = (xs,acc,x*m+ix)
--listMProd mdt vdt = let (_,l,lst)= foldr go (rvdt,[],0) mdt
--    in lst:l
--    where
--        rvdt = reverse vdt
--        go m ([],acc,ix) =
--            let (x:xs)=rvdt
--            in (xs,ix:acc,x*m)
--        go m ((x:xs),acc,ix) = (xs,acc,x*m+ix)
--listMProd mdt vdt = map sum $ takes1 (length vdt) $ zipWith (*) mdt (cycle vdt)

-- | Calculate the mean of a list
calcMeanList :: (Fractional a) => [a] -> a
calcMeanList = uncurry (/) . foldr (\e (s,c) -> (e+s,c+1)) (0,0)

takes1 :: Int -> [a] -> [[a]]
takes1 _ [] = []
takes1 s xs = let
    (one,two) = splitAt s xs
    in one : takes1 s two

-- | Split a given list into a series of list whose length is given by the first argument
takes :: [Int] -> [a] -> [[a]]
takes [] _ = []
takes (_:sz) [] = [] : takes sz []
takes (s:sz) xs = let
    (one,two) = splitAt s xs
    in one : takes sz two

-- | It's a simple, differentiable sigmoid function.
sigmoid :: Floating a => a -> a
sigmoid x = 1 / (1 + exp (-x))
{-# INLINE sigmoid #-}

-- | Normalize a list of values between 0 and 1
normalize :: (Floating b,Ord b)=> [(a,b)] -> [(a,b)]
normalize xs =
    let
        ws = map snd xs
        mi = minimum ws
        mx = maximum ws
        df = mx-mi
     in  map (\(a,w)->(a,(w-mi)/df)) xs

-- | Parallel zipwith
parZipWith :: (a -> b -> c) -> [a] -> [b] -> [c]
parZipWith f xs ys = withStrategy (parList rseq) $ zipWith f xs ys

-- | Parallel zipWith3
parZipWith3 :: (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
parZipWith3 f xs ys zs = withStrategy (parList rseq) $ zipWith3 f xs ys zs

-- | Number of binary digits necessary to store the given number
binaryDigits :: (Num a, Ord a) => a -> a
binaryDigits a = go 1 1
    where go nb sz
            | nb == a = sz
            | nb > a  = sz - 1
            | otherwise = go (nb*2) (sz+1)

-- | Number of bits needed to store in a sparse manner (with 2 bits) the given number
sparseSize :: (Integral a, Ord a) => a -> a
sparseSize a = go 1
      where
        go nb =
            let pos = nb * (nb - 1) `div` 2
            in if pos >= a
                then nb
                else go (nb+1)

-- | Sparse representation: store all numbers from 0 to n using 2 bits
-- Returns the location of each two bits
sparse :: Int -> [(Int,Int)]
sparse n = let
    sz = sparseSize n
    in take n $ iterate (gen sz) (0,1)
    where gen sz (a,b) =
            if b==sz-1 then (a+1,a+2)
                       else (a,b+1)

-- | Standard Normal distribution
stdNormal :: (Monad m,RandomGen g) => RandT g m (Double,Double)
stdNormal = do
    (x1,x2,w1)<-gen
    let w  = sqrt( (-2 * log( w1 ) ) / w1 )
        y1 = x1 * w;
        y2 = x2 * w;
    return (y1,y2)
    where
        ranx = do
            r1 <- getRandomR (0,1)
            return (2 * r1 - 1)
        gen = do
            x1 <- ranx
            x2 <- ranx
            let w = x1 * x1 + x2 * x2
            if w>=1 then gen else return (x1,x2,w)

softmax :: [Double] -> [Double]
softmax is =
    let den = sum $ map exp is
    in map (( / den) . exp) is

softmaxV :: Vector Double -> Vector Double
softmaxV is =
    let den = sumElements $ cmap exp is
    in cmap (( / den) . exp) is

euclidian :: Vector Double -> Vector Double -> Double
euclidian v1 v2 = sumElements $ cmap (**2) (v1 - v2)

equilateralEncoding :: Int -> Matrix Double
equilateralEncoding n =
    let z1 = replicate (n-2) 0
        m = (n><(n-1)) (-1: z1 ++ (1 : z1) ++ repeat 0)
    in foldl' pass m [2..n-1]
    where
        r :: Double -> Double
        r k = -1 / k
        f :: Double -> Double
        f k = sqrt (k*k -1) / k
        pass :: Matrix Double -> Int -> Matrix Double
        pass m1 ik =
            let k = fromIntegral ik
                f1 = f k
                a0 = concatMap (\i->map (\j->((i,j),f1)) [0..ik-2]) [0..ik-1]
                m2 = accum m1 (*) a0
                r1 = r k
                a1 = map (\x->((x,ik-1),r1)) [0..ik-1] ++ [((ik,ik-1),1)]
                m3 = accum m2 const a1
            in m3

equilateralDecoding :: Matrix Double -> Vector Double -> Int
equilateralDecoding m os =
    let rs = zip (toRows m) [0..]
        es = map (\(vs,ix)->(euclidian vs os,ix)) rs
    in snd $ minimumBy (comparing fst) es

roundTo :: (Fractional a, Integral b, RealFrac r) =>
                 b -> r -> a
roundTo n f=  (fromInteger $ round $ f * (10^n)) / (10.0^^n)

ordNub :: (Ord a) => [a] -> [a]
ordNub l = go S.empty l
     where
       go _ []     = []
       go s (x:xs) = if x `S.member` s then go s xs
                                     else x : go (S.insert x s) xs


{-# LANGUAGE FlexibleContexts, OverloadedLists, OverloadedStrings, UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Main (
	main,
	testXor, testClassify, testIris, testMnist
	) where

import Prelude.Unicode

import Control.DeepSeq
import Control.Monad
import Control.Monad.State.Strict
import Control.Monad.Morph
import qualified Data.Attoparsec.Text as A
import qualified Data.Text.IO as T
import qualified Data.Vector as V
import System.Random (mkStdGen)
import Test.Hspec
import Numeric.LinearAlgebra hiding (conv2, conv)
import Text.Format

import Numeric.Trainee.Data
import Numeric.Trainee.Neural
import Numeric.Trainee.Gradee (reshapeVec, flattenMat, concatVecs)

main ∷ IO ()
main = hspec $
	describe "training neural network" $ do
		it "should approximate xor function" testXor
		it "should classify objects" testClassify
		it "should support 3 classes" testIris

testXor ∷ IO ()
testXor = do
	n ← return $ netSeed (mkStdGen 0) $ fc sigma 2 2 ⭃ fc sigma 2 2 ⭃ fc sigma 2 1 ∷ IO (Net Double)
	(e, n') ← runLearnT n $ trainUntil 5.0 10000 4 1e-4 squared xorSamples
	e `shouldSatisfy` (≤ 1e-4)
	mapM_ (shouldPass n' 0.1) xorSamples

testClassify ∷ IO ()
testClassify = do
	n ← return $ netSeed (mkStdGen 0) $ fc sigma 4 4 ⭃ fc sigma 4 2 ⭃ fc sigma 2 1 ∷ IO (Net Double)
	let
		cases ∷ [(String, Vector Double → Vector Double)]
		cases = [
			("adult+stretch", fn $ \[_, _, act, age] → act ≡ 0.0 ∨ age ≡ 0.0),
			("adult-stretch", fn $ \[_, _, act, age] → act ≡ 0.0 ∧ age ≡ 0.0),
			("yellow-small+adult-stretch", fn $ \[color, sz, act, age] → (color ≡ 0.0 ∧ sz ≡ 0.0) ∨ (act ≡ 0.0 ∧ age ≡ 0.0)),
			("yellow-small", fn $ \[color, sz, _, _] → color ≡ 0.0 ∧ sz ≡ 0.0)]
		fn ∷ ([Double] → Bool) → Vector Double → Vector Double
		fn f = vector ∘ return ∘ fromIntegral ∘ fromEnum ∘ f ∘ toList
	forM_ cases $ \(name, fun) → do
		classes ← readBalloonSamples $ "data/classify/balloon/{name}.data" ~~ ("name" ~% name)
		(e, n') ← runLearnT n $ trainUntil 10.0 1000 10 1e-4 squared classes
		e `shouldSatisfy` (≤ 1e-4)
		mapM_ (shouldPass n' 0.1) [xs ⇢ fun xs |
			xs ← map vector (replicateM 4 [0.0, 1.0])]

instance Show (Vector a) ⇒ FormatBuild (Vector a)

testIris ∷ IO ()
testIris = do
	n ← return $ netSeed (mkStdGen 0) $ fc sigma 4 12 ⭃ fc sigma 12 3 ∷ IO (Net Double)
	classes ← readIrisData "data/classify/iris/iris.data"
	(_, n') ← runLearnT n $ hoist (`evalStateT` (rightAnswers n classes, 0)) $ learnIris classes
	let
		failedSamples = filter (not ∘ rightClass n' ∘ snd) $ zip ([1..] ∷ [Integer]) (V.toList classes)
	putStrLn "--- Done. ---"
	putStrLn "Failed samples:"
	forM_ failedSamples $ \(i, Sample inp outp) → putStrLn $ "{n}\t{input} → {output} ≢ {right}"
		~~ ("n" ~% i)
		~~ ("input" ~% inp)
		~~ ("output" ~% eval n' inp)
		~~ ("right" ~% outp)
	length failedSamples `shouldSatisfy` (≤ 2)
	where
		learnIris ∷ Samples (Vector Double) (Vector Double) → StateT (Net Double) (StateT (Int, Int) IO) ()
		learnIris classes = do
			e ← fmap last $ replicateM 100 $ trainEpoch 1.0 150 crossEntropy classes
			n' ← get
			let
				ans = rightAnswers n' classes
			lift $ modify (\(ans', long') → (ans, if ans ≡ ans' then succ long' else 0))
			liftIO $ putStrLn $ "correct answers: {rights}/{total}; error = {e}"
				~~ ("e" ~% e)
				~~ ("rights" ~% ans)
				~~ ("total" ~% length classes)
			long' ← lift $ gets snd
			when (long' ≤ 10 ∧ e > 0.1) $ learnIris classes

testMnist ∷ IO ()
testMnist = do
	n ← net $
		return (computee (reshapeVec 28)) ⭃ ndup 1 ⭃
		npar 1 (pad2 2 2) ⭃ dconv2 sigma 1 32 (5, 5) ⭃ npar 32 (maxPool2 2 2) ⭃
		npar 32 (pad2 2 2) ⭃ dconv2 sigma 32 64 (5, 5) ⭃ npar 64 (maxPool2 2 2) ⭃
		npar 64 (return $ computee flattenMat) ⭃ return (computee concatVecs) ⭃
		fc sigma (7 * 7 * 64) 1024 ⭃
		fc sigma 1024 10
		∷ IO (Net Double)
	putStrLn "reading train data"
	smps ← readMnist "data/classify/mnist/train.csv"
	putStrLn $ "loaded {0} samples" ~~ length smps
	(_, n') ← runLearnT n $ learnMnist smps
	print n'
	where
		learnMnist ∷ Samples (Vector Double) (Vector Double) → StateT (Net Double) IO ()
		learnMnist smps = do
			ixs ← shuffleList [0 .. V.length smps - 1]
			es ← forM ixs $ \i → do
				-- let
				-- 	b = V.fromList $ map (smps V.!) is
				e ← trainOnce 0.001 crossEntropy (smps V.! i)
				liftIO $ putStrLn $ "error: {}" ~~ e
				return e
			let
				e = avg es
			liftIO $ putStrLn $ "error: {}" ~~ e
			when (e > 0.1) $ learnMnist smps

rightClass ∷ Net Double → Sample (Vector Double) (Vector Double) → Bool
rightClass n_ (Sample i o) = maxIndex (eval n_ i) ≡ maxIndex o

rightAnswers ∷ Net Double → Samples (Vector Double) (Vector Double) → Int
rightAnswers n_ = V.length ∘ V.filter (rightClass n_)

xorSamples ∷ Samples (Vector Double) (Vector Double)
xorSamples = samples [
	[0, 0] ⇢ [0],
	[1, 1] ⇢ [0],
	[1, 0] ⇢ [1],
	[0, 1] ⇢ [1]]

readIrisData ∷ FilePath → IO (Samples (Vector Double) (Vector Double))
readIrisData fpath = parseCsvFile False fpath $ do
	is ← mapM (col_ >=> read_ >=> (return ∘ (* 0.1))) [0 .. 3]
	os ← col_ 4 >>= class_ ["Iris-setosa", "Iris-versicolor", "Iris-virginica"]
	sample_ is os

readBalloonSamples ∷ FilePath → IO (Samples (Vector Double) (Vector Double))
readBalloonSamples fpath = parseCsvFile False fpath $ do
	is ← sequence [
		col_ 0 >>= enum_ ["yellow", "purple"],
		col_ 1 >>= enum_ ["small", "large"],
		col_ 2 >>= enum_ ["stretch", "dip"],
		col_ 3 >>= enum_ ["adult", "child"]]
	os ← col_ 4 >>= enum_ ["f", "t"] >>= single_
	sample_ is os

readMnist ∷ FilePath → IO (Samples (Vector Double) (Vector Double))
readMnist fpath = do
	cts ← T.readFile fpath
	either error return $ A.parseOnly mnist cts
	where
		mnist ∷ A.Parser (Samples (Vector Double) (Vector Double))
		mnist = do
			_ ← A.manyTill A.anyChar A.endOfLine
			fmap samples $ A.many' $ do
				(f:fs) ← A.sepBy A.decimal (A.char ',') <* A.endOfLine
				let
					is = fromList $ map ((/ 255.0) ∘ fromIntegral) fs
					os = fromList $ replicate f 0.0 ++ [1.0] ++ replicate (10 - f - 1) 0.0
				is `deepseq` os `deepseq` return (Sample is os)

shouldPass ∷ Net Double → Double → Sample (Vector Double) (Vector Double) → IO ()
shouldPass n ε (Sample xs ys) = when (err > ε) $ expectationFailure msg where
	msg = show xs ++ " -> " ++ show res ++ " should be " ++ show ys
	res = eval n xs
	err = vecSize (res - ys)
	vecSize v = sqrt (dot v v)


{-# LANGUAGE Rank2Types, FlexibleContexts #-}

module Numeric.AD.Lagrangian.Internal where

import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Storable as S
import qualified Data.Packed.Vector as V
import qualified Data.Packed.Matrix as M

import GHC.IO (unsafePerformIO)

import Numeric.AD
import Numeric.Optimization.Algorithms.HagerZhang05
import Numeric.LinearAlgebra.Algorithms

-- | An equality constraint of the form @g(x, y, ...) = c@. Use '<=>' to
-- construct a 'Constraint'.
newtype Constraint = Constraint
  {unConstraint :: forall a. (Floating a) => ([a] -> a, a)}

infixr 1 <=>
-- | Build a 'Constraint' from a function and a constant
(<=>) :: (forall a. (Floating a) => [a] -> a)
      -> (forall b. (Floating b) => b)
      -> Constraint
f <=> c = Constraint (f, c)

-- | Numerically minimize the Langrangian. The objective function and each of
-- the constraints must take the same number of arguments.
minimize :: (forall a. (Floating a) => [a] -> a)
         -- ^ The objective function to minimize
         -> [Constraint]
         -- ^ A list of constraints @g \<=\> c@ corresponding to equations of
         -- the form @g(x, y, ...) = c@
         -> Double
         -- ^ Stop iterating when the largest component of the gradient is
         -- smaller than this value
         -> Int
         -- ^ The arity of the objective function, which must equal the arity of
         -- the constraints
         -> Either (Result, Statistics) (V.Vector Double, V.Vector Double)
         -- ^ Either a 'Right' containing the argmin and the Lagrange
         -- multipliers, or a 'Left' containing an explanation of why the
         -- gradient descent failed
minimize f constraints tolerance argCount = result where
    -- At a constrained minimum of `f`, the gradient of the Lagrangian must be
    -- zero. So we square the Lagrangian's gradient (making it non-negative) and
    -- minimize it.
    (sqGradLgn, gradSqGradLgn) = (fst . g, snd . g) where
        g = grad' $ squaredGrad $ lagrangian f constraints argCount

    -- Perhaps this should be exposed. ...
    guess = U.replicate (argCount + length constraints) 1

    result = case unsafePerformIO $
                    optimize
                        (defaultParameters {printFinal = False})
                        tolerance
                        guess
                        (VFunction (sqGradLgn . U.toList))
                        (VGradient (U.fromList . gradSqGradLgn . U.toList))
                        Nothing of
       (vs, ToleranceStatisfied, _) -> Right (S.take argCount vs,
                                              S.drop argCount vs)
       (_, x, y) -> Left (x, y)

-- | Numerically maximize the Langrangian. The objective function and each of
-- the constraints must take the same number of arguments.
maximize :: (forall a. (Floating a) => [a] -> a)
         -- ^ The objective function to minimize
         -> [Constraint]
         -- ^ A list of constraints @g \<=\> c@ corresponding to equations of
         -- the form @g(x, y, ...) = c@
         -> Double
         -- ^ Stop iterating when the largest component of the gradient is
         -- smaller than this value
         -> Int
         -- ^ The arity of the objective function, which must equal the arity of
         -- the constraints
         -> Either (Result, Statistics) (V.Vector Double, V.Vector Double)
         -- ^ Either a 'Right' containing the argmax and the Lagrange
         -- multipliers, or a 'Left' containing an explanation of why the
         -- gradient ascent failed
maximize f = minimize $ negate . f

lagrangian :: (Floating a)
           => (forall b. (Floating b) => [b] -> b)
           -> [Constraint]
           -> Int
           -> [a]
           -> a
lagrangian f constraints argCount argsAndLams = result where
    args = take argCount argsAndLams
    lams = drop argCount argsAndLams

    -- g(x, y, ...) = c <=> g(x, y, ...) - c = 0
    appliedConstraints = fmap (\(Constraint (g, c)) -> g args - c) constraints

    -- L(x, y, ..., lam0, ...) = f(x, y, ...) + lam0 * (g0 - c0) ...
    result = (f args) + (sum . zipWith (*) lams $ appliedConstraints)

squaredGrad :: (Floating a)
            => (forall b. (Floating b) => [b] -> b)
            -> [a]
            -> a
squaredGrad f = sum . fmap square . grad f where
    square x = x * x

-- | WARNING: Experimental.
--   This is not a true feasibility test for the function. I am not sure
--   exactly how to implement that. This just checks the feasiblility at a
--   point. If this ever returns false, 'solve' can fail.
feasible :: (Floating a, Field a, M.Element a)
         => (forall b. (Floating b) => [b] -> b)
         ->[Constraint]
         -> [a]
         -> Bool
feasible f constraints points = result where
    sqGradLgn :: (Floating a) => [a] -> a
    sqGradLgn = squaredGrad $ lagrangian f constraints $ length points

    hessianMatrix = M.fromLists . hessian sqGradLgn $ points

    -- make sure all of the eigenvalues are positive
    result = all (>0) . V.toList . eigenvaluesSH $ hessianMatrix


{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MagicHash             #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms       #-}
{-# LANGUAGE RebindableSyntax      #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeSynonymInstances  #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE ViewPatterns          #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- |
-- Module      : Data.Array.Accelerate.Data.Complex
-- Copyright   : [2015..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- Complex numbers, stored in the usual C-style array-of-struct representation,
-- for easy interoperability.
--
module Data.Array.Accelerate.Data.Complex (

  -- * Rectangular from
  Complex(..), pattern (::+),
  real,
  imag,

  -- * Polar form
  mkPolar,
  cis,
  polar,
  magnitude, magnitude',
  phase,

  -- * Conjugate
  conjugate,

) where

import Data.Array.Accelerate.Classes.Eq
import Data.Array.Accelerate.Classes.Floating
import Data.Array.Accelerate.Classes.Fractional
import Data.Array.Accelerate.Classes.FromIntegral
import Data.Array.Accelerate.Classes.Num
import Data.Array.Accelerate.Classes.Ord
import Data.Array.Accelerate.Classes.RealFloat
import Data.Array.Accelerate.Data.Functor
import Data.Array.Accelerate.Pattern
import Data.Array.Accelerate.Prelude
import Data.Array.Accelerate.Representation.Tag
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Representation.Vec
import Data.Array.Accelerate.Smart
import Data.Array.Accelerate.Sugar.Elt
import Data.Array.Accelerate.Sugar.Vec
import Data.Array.Accelerate.Type
import Data.Primitive.Vec

import Data.Complex                                                 ( Complex(..) )
import Prelude                                                      ( ($) )
import qualified Data.Complex                                       as C
import qualified Prelude                                            as P

infix 6 ::+
pattern (::+) :: Elt a => Exp a -> Exp a -> Exp (Complex a)
pattern r ::+ i <- (deconstructComplex -> (r, i))
  where (::+) = constructComplex
{-# COMPLETE (::+) #-}

-- Use an array-of-structs representation for complex numbers if possible.
-- This matches the standard C-style layout, but we can use this representation only at
-- specific types (not for any type 'a') as we can only have vectors of primitive type.
-- For other types, we use a structure-of-arrays representation. This is handled by the
-- ComplexR. We use the GADT ComplexR and function complexR to reconstruct
-- information on how the elements are represented.
--
instance Elt a => Elt (Complex a) where
  type EltR (Complex a) = ComplexR (EltR a)
  eltR = let tR = eltR @a
          in case complexR tR of
               ComplexVec s -> TupRsingle $ VectorScalarType $ VectorType 2 s
               ComplexTup   -> TupRunit `TupRpair` tR `TupRpair` tR

  tagsR = let tR = eltR @a
           in case complexR tR of
               ComplexVec s -> [ TagRsingle (VectorScalarType (VectorType 2 s)) ]
               ComplexTup   -> let go :: TypeR t -> [TagR t]
                                   go TupRunit         = [TagRunit]
                                   go (TupRsingle s)   = [TagRsingle s]
                                   go (TupRpair ta tb) = [TagRpair a b | a <- go ta, b <- go tb]
                                in
                                [ TagRunit `TagRpair` ta `TagRpair` tb | ta <- go tR, tb <- go tR ]

  toElt = case complexR $ eltR @a of
    ComplexVec _ -> \(Vec2 r i)   -> toElt r :+ toElt i
    ComplexTup   -> \(((), r), i) -> toElt r :+ toElt i

  fromElt (r :+ i) = case complexR $ eltR @a of
    ComplexVec _ -> Vec2 (fromElt r) (fromElt i)
    ComplexTup   -> (((), fromElt r), fromElt i)

type family ComplexR a where
  ComplexR Half   = Vec2 Half
  ComplexR Float  = Vec2 Float
  ComplexR Double = Vec2 Double
  ComplexR Int    = Vec2 Int
  ComplexR Int8   = Vec2 Int8
  ComplexR Int16  = Vec2 Int16
  ComplexR Int32  = Vec2 Int32
  ComplexR Int64  = Vec2 Int64
  ComplexR Word   = Vec2 Word
  ComplexR Word8  = Vec2 Word8
  ComplexR Word16 = Vec2 Word16
  ComplexR Word32 = Vec2 Word32
  ComplexR Word64 = Vec2 Word64
  ComplexR a      = (((), a), a)

-- This isn't ideal because we gather the evidence based on the
-- representation type, so we really get the evidence (VecElt (EltR a)),
-- which is not very useful...
--    - TLM 2020-07-16
data ComplexType a c where
  ComplexVec :: VecElt a => SingleType a -> ComplexType a (Vec2 a)
  ComplexTup ::                             ComplexType a (((), a), a)

complexR :: TypeR a -> ComplexType a (ComplexR a)
complexR = tuple
  where
    tuple :: TypeR a -> ComplexType a (ComplexR a)
    tuple TupRunit       = ComplexTup
    tuple TupRpair{}     = ComplexTup
    tuple (TupRsingle s) = scalar s

    scalar :: ScalarType a -> ComplexType a (ComplexR a)
    scalar (SingleScalarType t) = single t
    scalar VectorScalarType{}   = ComplexTup

    single :: SingleType a -> ComplexType a (ComplexR a)
    single (NumSingleType t) = num t

    num :: NumType a -> ComplexType a (ComplexR a)
    num (IntegralNumType t) = integral t
    num (FloatingNumType t) = floating t

    integral :: IntegralType a -> ComplexType a (ComplexR a)
    integral TypeInt    = ComplexVec singleType
    integral TypeInt8   = ComplexVec singleType
    integral TypeInt16  = ComplexVec singleType
    integral TypeInt32  = ComplexVec singleType
    integral TypeInt64  = ComplexVec singleType
    integral TypeWord   = ComplexVec singleType
    integral TypeWord8  = ComplexVec singleType
    integral TypeWord16 = ComplexVec singleType
    integral TypeWord32 = ComplexVec singleType
    integral TypeWord64 = ComplexVec singleType

    floating :: FloatingType a -> ComplexType a (ComplexR a)
    floating TypeHalf   = ComplexVec singleType
    floating TypeFloat  = ComplexVec singleType
    floating TypeDouble = ComplexVec singleType

constructComplex :: forall a. Elt a => Exp a -> Exp a -> Exp (Complex a)
constructComplex r i =
  case complexR (eltR @a) of
    ComplexTup   -> coerce $ T2 r i
    ComplexVec _ -> V2 (coerce @a @(EltR a) r) (coerce @a @(EltR a) i)

deconstructComplex :: forall a. Elt a => Exp (Complex a) -> (Exp a, Exp a)
deconstructComplex c@(Exp c') =
  case complexR (eltR @a) of
    ComplexTup   -> let T2 r i = coerce c in (r, i)
    ComplexVec t -> let T2 r i = Exp (SmartExp (VecUnpack (VecRsucc (VecRsucc (VecRnil t))) c'))
                     in (r, i)

coerce :: EltR a ~ EltR b => Exp a -> Exp b
coerce (Exp e) = Exp e

instance (Lift Exp a, Elt (Plain a)) => Lift Exp (Complex a) where
  type Plain (Complex a) = Complex (Plain a)
  lift (r :+ i) = lift r ::+ lift i

instance Elt a => Unlift Exp (Complex (Exp a)) where
  unlift (r ::+ i) = r :+ i

instance Eq a => Eq (Complex a) where
  r1 ::+ c1 == r2 ::+ c2 = r1 == r2 && c1 == c2
  r1 ::+ c1 /= r2 ::+ c2 = r1 /= r2 || c1 /= c2

instance RealFloat a => P.Num (Exp (Complex a)) where
  (+)    = lift2 ((+) :: Complex (Exp a) -> Complex (Exp a) -> Complex (Exp a))
  (-)    = lift2 ((-) :: Complex (Exp a) -> Complex (Exp a) -> Complex (Exp a))
  (*)    = lift2 ((*) :: Complex (Exp a) -> Complex (Exp a) -> Complex (Exp a))
  negate = lift1 (negate :: Complex (Exp a) -> Complex (Exp a))
  signum z@(x ::+ y) =
    if z == 0
       then z
       else let r = magnitude z
             in x/r ::+ y/r
  abs z         = magnitude z ::+ 0
  fromInteger n = fromInteger n ::+ 0

instance RealFloat a => P.Fractional (Exp (Complex a)) where
  fromRational x  = fromRational x ::+ 0
  z / z'          = (x*x''+y*y'') / d ::+ (y*x''-x*y'') / d
    where
      x  :+ y   = unlift z
      x' :+ y'  = unlift z'
      --
      x'' = scaleFloat k x'
      y'' = scaleFloat k y'
      k   = - max (exponent x') (exponent y')
      d   = x'*x'' + y'*y''

instance RealFloat a => P.Floating (Exp (Complex a)) where
  pi                = pi ::+ 0
  exp (x ::+ y)     = let expx = exp x
                       in expx * cos y ::+ expx * sin y
  log z             = log (magnitude z) ::+ phase z
  sqrt z@(x ::+ y)  =
    if z == 0
      then 0
      else u ::+ (y < 0 ? (-v, v))
    where
      T2 u v = x < 0 ? (T2 v' u', T2 u' v')
      v'     = abs y / (u'*2)
      u'     = sqrt ((magnitude z + abs x) / 2)

  x ** y =
    if y == 0 then 1 else
    if x == 0 then if exp_r > 0 then 0 else
                   if exp_r < 0 then inf ::+ 0
                                else nan ::+ nan
              else if isInfinite r || isInfinite i
                     then if exp_r > 0 then inf ::+ 0 else
                          if exp_r < 0 then 0
                                       else nan ::+ nan
                     else exp (log x * y)
    where
      r     ::+ i  = x
      exp_r ::+ _  = y
      --
      inf = 1 / 0
      nan = 0 / 0

  sin (x ::+ y)  = sin x * cosh y ::+ cos x * sinh y
  cos (x ::+ y)  = cos x * cosh y ::+ (- sin x * sinh y)
  tan (x ::+ y)  = (sinx*coshy ::+ cosx*sinhy) / (cosx*coshy ::+ (-sinx*sinhy))
    where
      sinx  = sin x
      cosx  = cos x
      sinhy = sinh y
      coshy = cosh y

  sinh (x ::+ y) = cos y * sinh x ::+ sin  y * cosh x
  cosh (x ::+ y) = cos y * cosh x ::+ sin y * sinh x
  tanh (x ::+ y) = (cosy*sinhx ::+ siny*coshx) / (cosy*coshx ::+ siny*sinhx)
    where
      siny  = sin y
      cosy  = cos y
      sinhx = sinh x
      coshx = cosh x

  asin z@(x ::+ y) = y' ::+ (-x')
    where
      x' ::+ y' = log (((-y) ::+ x) + sqrt (1 - z*z))

  acos z                    = y'' ::+ (-x'')
    where
      x'' ::+ y''  = log (z + ((-y') ::+ x'))
      x'  ::+ y'   = sqrt (1 - z*z)

  atan z@(x ::+ y) = y' ::+ (-x')
    where
      x' ::+ y' = log (((1-y) ::+ x) / sqrt (1+z*z))

  asinh z =  log (z + sqrt (1+z*z))
  acosh z =  log (z + (z+1) * sqrt ((z-1)/(z+1)))
  atanh z =  0.5 * log ((1.0+z) / (1.0-z))

instance (FromIntegral a b, Num b, Elt (Complex b)) => FromIntegral a (Complex b) where
  fromIntegral x = fromIntegral x ::+ 0

-- | @since 1.2.0.0
--
instance Functor Complex where
  fmap f (r ::+ i) = f r ::+ f i

-- | The non-negative magnitude of a complex number
--
magnitude :: RealFloat a => Exp (Complex a) -> Exp a
magnitude (r ::+ i) = scaleFloat k (sqrt (sqr (scaleFloat mk r) + sqr (scaleFloat mk i)))
  where
    k     = max (exponent r) (exponent i)
    mk    = -k
    sqr z = z * z

-- | As 'magnitude', but ignore floating point rounding and use the traditional
-- (simpler to evaluate) definition.
--
-- @since 1.3.0.0
--
magnitude' :: RealFloat a => Exp (Complex a) -> Exp a
magnitude' (r ::+ i) = sqrt (r*r + i*i)

-- | The phase of a complex number, in the range @(-'pi', 'pi']@. If the
-- magnitude is zero, then so is the phase.
--
phase :: RealFloat a => Exp (Complex a) -> Exp a
phase z@(r ::+ i) =
  if z == 0
    then 0
    else atan2 i r

-- | The function 'polar' takes a complex number and returns a (magnitude,
-- phase) pair in canonical form: the magnitude is non-negative, and the phase
-- in the range @(-'pi', 'pi']@; if the magnitude is zero, then so is the phase.
--
polar :: RealFloat a => Exp (Complex a) -> Exp (a,a)
polar z =  T2 (magnitude z) (phase z)

-- | Form a complex number from polar components of magnitude and phase.
--
mkPolar :: forall a. Floating a => Exp a -> Exp a -> Exp (Complex a)
mkPolar = lift2 (C.mkPolar :: Exp a -> Exp a -> Complex (Exp a))

-- | @'cis' t@ is a complex value with magnitude @1@ and phase @t@ (modulo
-- @2*'pi'@).
--
cis :: forall a. Floating a => Exp a -> Exp (Complex a)
cis = lift1 (C.cis :: Exp a -> Complex (Exp a))

-- | Return the real part of a complex number
--
real :: Elt a => Exp (Complex a) -> Exp a
real (r ::+ _) = r

-- | Return the imaginary part of a complex number
--
imag :: Elt a => Exp (Complex a) -> Exp a
imag (_ ::+ i) = i

-- | Return the complex conjugate of a complex number, defined as
--
-- > conjugate(Z) = X - iY
--
conjugate :: Num a => Exp (Complex a) -> Exp (Complex a)
conjugate z = real z ::+ (- imag z)


-- |
-- Module      : Cartesian.Types
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created _ _ 2015

-- TODO | -
--        -

-- SPEC | -
--        -

--------------------------------------------------------------------------------------------------------------------------------------------
-- GHC Pragmas
--------------------------------------------------------------------------------------------------------------------------------------------

--------------------------------------------------------------------------------------------------------------------------------------------
-- API
--------------------------------------------------------------------------------------------------------------------------------------------
module Cartesian.Types (
  -- ^ Third party types
  V1(..), V2(..), V3(..), V4(..), Complex(..),

  -- ^ Synonyms
  BoxLens, Axis, Axes, Polygon,

  -- ^ Coordinate types
  Normalised, Absolute,

  -- ^ Types defined in this library
  BoundingBox(..), Line, Linear,

  -- ^ Classes
  HasX, HasY, HasZ) where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
import Linear.V1
import Linear.V2
import Linear.V3
import Linear.V4

import Data.Complex (Complex(..))

import Cartesian.Internal.Types


{-# LANGUAGE DeriveDataTypeable        #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE FunctionalDependencies    #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE RecordWildCards           #-}
{-# LANGUAGE TypeFamilies              #-}
{-# LANGUAGE UndecidableInstances      #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Plots.Types.Histogram
-- Copyright   :  (C) 2015 Christopher Chalmers
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Christopher Chalmers
-- Stability   :  experimental
-- Portability :  non-portable

-- A histogram is a graphical representation of the distribution of
-- numerical data. It is an estimate of the probability distribution of
-- a continuous variable.
--
----------------------------------------------------------------------------

module Plots.Types.Histogram
  (
    -- * Histogram plot
    HistogramPlot

    -- ** Already computed histograms
  , computedHistogram

    -- ** Histogram options
  , HistogramOptions
  , HasHistogramOptions (..)

    -- ** Normalisation
  , NormalisationMethod
  , count
  , probability
  , countDensity
  , pdf
  , cumilative
  , cdf

    -- ** Plotting histograms
  , histogramPlot
  , histogramPlot'
  , histogramPlotOf
  , histogramPlotOf'

    -- * Low level constructors
  , mkComputedHistogram
  , mkHistogramPlot
  ) where

import           Control.Monad.State.Lazy

import qualified Data.Foldable               as F
import           Data.Maybe
import           Data.Typeable

import qualified Data.Vector                 as V
import qualified Statistics.Sample.Histogram as Stat

import           Diagrams.Core.Transform     (fromSymmetric)
import           Diagrams.Prelude
import           Linear.V2                   (_yx)

import           Plots.Axis
import           Plots.Style
import           Plots.Types
import           Plots.Util

-- | Construct a rectangle of size $v$ with the bottom left at point $p$.
rectBL :: (InSpace V2 n t, TrailLike t) => Point V2 n -> V2 n -> t
rectBL p (V2 x y) =
  trailLike $ fromOffsets [V2 x 0, V2 0 y, V2 (-x) 0] # closeTrail `at` p

------------------------------------------------------------------------
-- GHistogram plot
------------------------------------------------------------------------

-- | Simple histogram type supporting uniform bins.
data HistogramPlot n = HistogramPlot
  { hWidth  :: n
  , hStart  :: n
  , hValues :: [n]
  , hOrient :: Orientation
  } deriving Typeable

type instance V (HistogramPlot n) = V2
type instance N (HistogramPlot n) = n

instance OrderedField n => Enveloped (HistogramPlot n) where
  getEnvelope HistogramPlot {..} =
    -- don't like this redundant code
    getEnvelope . orient hOrient _reflectXY id . (id :: Path v n -> Path v n) $
      ifoldMap drawBar hValues
    where
      drawBar i h = rectBL (mkP2 x 0) (V2 hWidth h)
        where x = hStart + fromIntegral i * hWidth

instance (TypeableFloat n, Renderable (Path V2 n) b)
    => Plotable (HistogramPlot n) b where
  renderPlotable s sty HistogramPlot {..} =
    ifoldMap drawBar hValues
      # orient hOrient _reflectXY id
      # applyAreaStyle sty
      # transform (s^.specTrans)
    where
      drawBar i h = rectBL (mkP2 x 0) (V2 hWidth h)
        where x = hStart + fromIntegral i * hWidth

  defLegendPic sty HistogramPlot {..}
    = centerXY
    . applyAreaStyle sty'
    . orient hOrient _reflectXY id
    $ alignB (rect 4 7) ||| alignB (rect 4 10) ||| alignB (rect 4 6)
    where
      -- The legend bars don't look right if the line width is too big so we limit it
      sty' = sty & areaStyle . _lw %~ atMost (local 0.8)

instance HasOrientation (HistogramPlot n) where
  orientation = lens hOrient $ \hp o -> hp {hOrient = o}

------------------------------------------------------------------------
-- Simple histogram plot
------------------------------------------------------------------------

-- | Plot an already computed histogram with equally sized bins.
computedHistogram
  :: (MonadState (Axis b V2 n) m,
      Plotable (HistogramPlot n) b,
      F.Foldable f)
  => n   -- ^ start of first bin
  -> n   -- ^ width of each bin
  -> f n -- ^ heights of the bins
  -> State (Plot (HistogramPlot n) b) ()
  -> m ()
computedHistogram x0 w xs = addPlotable (mkComputedHistogram x0 w xs)

-- | Construct a 'HistogramPlot' from raw histogram data.
mkComputedHistogram
  :: F.Foldable f
  => n -- ^ start of first bin
  -> n -- ^ width of each bin
  -> f n -- ^ heights of the bins
  -> HistogramPlot n
mkComputedHistogram x0 w xs = HistogramPlot x0 w (F.toList xs) Horizontal

----------------------------------------------------------------------------
-- Building histograms
----------------------------------------------------------------------------

-- example setup
-- > import Plots
-- > sampleData :: [Double]
-- > sampleData =
-- >   [5.1,4.9,4.7,4.6,5.0,5.4,4.6,5.0,4.4,4.9
-- >   ,5.4,4.8,4.8,4.3,5.8,5.7,5.4,5.1,5.7,5.1
-- >   ,5.4,5.1,4.6,5.1,4.8,5.0,5.0,5.2,5.2,4.7
-- >   ,4.8,5.4,5.2,5.5,4.9,5.0,5.5,4.9,4.4,5.1
-- >   ,5.0,4.5,4.4,5.0,5.1,4.8,5.1,4.6,5.3,5.0
-- >   ,7.0,6.4,6.9,5.5,6.5,5.7,6.3,4.9,6.6,5.2
-- >   ,5.0,5.9,6.0,6.1,5.6,6.7,5.6,5.8,6.2,5.6
-- >   ,5.9,6.1,6.3,6.1,6.4,6.6,6.8,6.7,6.0,5.7
-- >   ,5.5,5.5,5.8,6.0,5.4,6.0,6.7,6.3,5.6,5.5
-- >   ,5.5,6.1,5.8,5.0,5.6,5.7,5.7,6.2,5.1,5.7
-- >   ,6.3,5.8,7.1,6.3,6.5,7.6,4.9,7.3,6.7,7.2
-- >   ,6.5,6.4,6.8,5.7,5.8,6.4,6.5,7.7,7.7,6.0
-- >   ,6.9,5.6,7.7,6.3,6.7,7.2,6.2,6.1,6.4,7.2
-- >   ,7.4,7.9,6.4,6.3,6.1,7.7,6.3,6.4,6.0,6.9
-- >   ,6.7,6.9,5.8,6.8,6.7,6.7,6.3,6.5,6.2,5.9
-- >   ]
--
-- > mkNmExample nm = r2Axis &~ do
-- >   yMin ?= 0
-- >   histogramPlot sampleData $ do
-- >     normaliseSample .= nm
-- > countDia = renderAxis $ mkNmExample count
-- > probabilityDia = renderAxis $ mkNmExample probability
-- > countDensityDia = renderAxis $ mkNmExample countDensity
-- > pdfDia = renderAxis $ mkNmExample pdf
-- > cumilativeDia = renderAxis $ mkNmExample cumilative
-- > cdfDia = renderAxis $ mkNmExample cdf

-- Histogram options ---------------------------------------------------

-- | The way to normalise the data from a histogram. The default method
--   is 'count'.
newtype NormalisationMethod =
  NM { runNM :: forall n. Fractional n => n -> V.Vector n -> V.Vector n }
 -- width -> heights -> normalised heights

instance Default NormalisationMethod where
  def = count

-- | The height of each bar is the number of observations. This is the
--   'Default' method.
--
-- === __Example__
--
-- <<diagrams/src_Plots_Types_Histogram_countDia.svg#diagram=countDia&height=350>>
count :: NormalisationMethod
count = NM $ \_ v -> v

-- | The sum of the heights of the bars is equal to 1.
--
-- === __Example__
--
-- <<diagrams/src_Plots_Types_Histogram_probabilityDia.svg#diagram=probabilityDia&height=350>>
probability :: NormalisationMethod
probability = NM $ \_ v -> v ^/ V.sum v

-- | The height of each bar is @n / w@ where @n@ is the number of
--   observations and @w@ is the total width.
--
-- === __Example__
--
-- <<diagrams/src_Plots_Types_Histogram_countDensityDia.svg#diagram=countDensityDia&height=350>>
countDensity :: NormalisationMethod
countDensity = NM $ \w v -> v ^/ w

-- | The total area of the bars is @1@. This gives a probability density
--   function estimate.
--
-- === __Example__
--
-- <<diagrams/src_Plots_Types_Histogram_pdfDia.svg#diagram=pdfDia&height=350>>
pdf :: NormalisationMethod
pdf = NM $ \w v -> v ^/ (w * V.sum v)

-- | The height of each bar is the cumulative number of observations in
--   each bin and all previous bins. The height of the last bar is the
--   total number of observations.
--
-- === __Example__
--
-- <<diagrams/src_Plots_Types_Histogram_cumilativeDia.svg#diagram=cumilativeDia&height=350>>
cumilative :: NormalisationMethod
cumilative = NM $ \_ -> V.scanl1 (+)

-- | Cumulative density function estimate. The height of each bar is
--   equal to the cumulative relative number of observations in the bin
--   and all previous bins. The height of the last bar is 1.
--
-- === __Example__
--
-- <<diagrams/src_Plots_Types_Histogram_cdfDia.svg#diagram=cdfDia&height=350>>
cdf :: NormalisationMethod
cdf = NM $ \_ v -> V.scanl1 (+) v ^/ V.sum v

-- | Options for binning histogram data. For now only very basic
--   histograms building is supported.
data HistogramOptions n = HistogramOptions
  { hBins   :: Int
  , hRange  :: Maybe (n, n)
  , hNorm   :: NormalisationMethod
  , oOrient :: Orientation
  }

type instance V (HistogramOptions n) = V2
type instance N (HistogramOptions n) = n

instance Default (HistogramOptions n) where
  def = HistogramOptions
    { hBins   = 10
    , hRange  = Nothing
    , hNorm   = def
    , oOrient = Vertical
    }

instance HasOrientation (HistogramOptions n) where
  orientation = lens oOrient $ \ho o -> ho {oOrient = o}

class HasOrientation a => HasHistogramOptions a where
  -- | Options for building the histogram from data.
  histogramOptions :: Lens' a (HistogramOptions (N a))

  -- | The number of bins (bars) to use for the histogram. Must be
  --   positive.
  --
  --   'Default' is @10@.
  numBins :: Lens' a Int
  numBins = histogramOptions . lens hBins (\ho n -> ho {hBins = n})

  -- | The range of data to consider when building the histogram. Any
  --   data outside the range is ignored.
  --
  --   'Default' is 'Nothing'.
  binRange :: Lens' a (Maybe (N a, N a))
  binRange = histogramOptions . lens hRange (\ho r -> ho {hRange = r})

  -- | Should the resulting histogram be normalised so the total area is
  --   1.
  --
  --   'Default' is False.
  normaliseSample :: Lens' a NormalisationMethod
  normaliseSample = histogramOptions . lens hNorm (\ho b -> ho {hNorm = b})

instance HasHistogramOptions (HistogramOptions n) where
  histogramOptions = id

instance HasHistogramOptions a => HasHistogramOptions (Plot a b) where
  histogramOptions = rawPlot . histogramOptions

-- | Create a histogram by binning the data using the
--   'HistogramOptions'.
mkHistogramPlot
  :: (F.Foldable f, RealFrac n)
  => HistogramOptions n -> f n -> HistogramPlot n
mkHistogramPlot HistogramOptions {..} xs =
  HistogramPlot
    { hWidth  = w
    , hStart  = a
    , hValues = V.toList $ runNM hNorm w ns
    , hOrient = Vertical
    }
  where
    w     = (b - a) / fromIntegral hBins
    ns    = Stat.histogram_ hBins a b v
    v     = V.fromList (F.toList xs)
    (a,b) = fromMaybe (range hBins v) hRange

-- Taken from Statistics, which was limited to 'Double'.
range :: (Ord n, Fractional n)
      => Int                    -- ^ Number of bins (must be positive).
      -> V.Vector n             -- ^ Sample data (cannot be empty).
      -> (n, n)
range nBins xs
    | nBins < 1 = error "Plots.Types.Histogram: invalid bin count"
    | V.null xs = error "Plots.Types.Histogram: empty sample"
    | lo == hi  = case abs lo / 10 of
                    a | a < 1e-6   -> (-1,1)
                      | otherwise  -> (lo - a, lo + a)
    | otherwise = (lo-d, hi+d)
  where
    d | nBins == 1 = 0
      | otherwise  = (hi - lo) / ((fromIntegral nBins - 1) * 2)
    (lo,hi)        = minMaxOf folded xs
{-# INLINE range #-}

-- |
-- mkWeightedHistogram
--   :: (F.Foldable f, OrderdField n)
--   => HistogramOptions n -> [(n, n)] -> HistogramPlot n
-- mkWeightedHistogram

------------------------------------------------------------------------
-- Histogram
------------------------------------------------------------------------

-- $ histogram
-- Histograms display data as barplot of x data, bin y data.
-- Box plots have the following lenses:
--
-- @
-- * 'setBin' :: 'Lens'' ('BoxPlot' v n) 'Double' - 10
-- @

-- | Add a 'HistogramPlot' to the 'AxisState' from a data set.
--
-- === __Example__
--
-- <<diagrams/src_Plots_Types_Histogram_histogramExample.svg#diagram=histogramExample&height=350>>
--
-- > import Plots
-- > histogramAxis :: Axis B V2 Double
-- > histogramAxis = r2Axis &~ do
-- >   histogramPlot sampleData $ do
-- >     key "histogram"
-- >     plotColor .= blue
-- >     areaStyle . _opacity .= 0.5
--
-- > histogramExample = renderAxis histogramAxis
histogramPlot
  :: (MonadState (Axis b V2 n) m, Plotable (HistogramPlot n) b, F.Foldable f, RealFrac n)
  => f n -- ^ data
  -> State (Plot (HistogramOptions n) b) () -- ^ changes to plot options
  -> m () -- ^ add plot to axis
histogramPlot ns s = addPlot (hoPlot & rawPlot %~ \ho -> mkHistogramPlot ho ns)
  where hoPlot = mkPlot def &~ s

-- | Make a 'HistogramPlot' without changes to the plot options.
histogramPlot'
  :: (MonadState (Axis b V2 n) m, Plotable (HistogramPlot n) b, F.Foldable f, RealFrac n)
  => f n -- ^ data
  -> m () -- ^ add plot to axis
histogramPlot' d = histogramPlot d (return ())

-- | Add a 'HistogramPlot' using a fold over the data.
histogramPlotOf
  :: (MonadState (Axis b V2 n) m, Plotable (HistogramPlot n) b, RealFrac n)
  => Fold s n -- ^ fold over the data
  -> s        -- ^ data to fold
  -> State (Plot (HistogramOptions n) b) () -- ^ change to the plot
  -> m () -- ^ add plot to the 'Axis'
histogramPlotOf f s = histogramPlot (toListOf f s)

-- | Same as 'histogramPlotOf' without any changes to the plot.
histogramPlotOf'
  :: (MonadState (Axis b V2 n) m, Plotable (HistogramPlot n) b, RealFrac n)
  => Fold s n -> s -> m ()
histogramPlotOf' f s = histogramPlotOf f s (return ())

-- temporary functions that will be in next lib release

_reflectionXY :: (Additive v, R2 v, Num n) => Transformation v n
_reflectionXY = fromSymmetric $ (_xy %~ view _yx) <-> (_xy %~ view _yx)

_reflectXY :: (InSpace v n t, R2 v, Transformable t) => t -> t
_reflectXY = transform _reflectionXY


module FractalStream.Models
    ( Coordinate
    , ComplexParametric1d
    ) where

import Lang.Numbers

type family Coordinate model :: *
data ComplexParametric1d

{-
import Data.Word
import Data.Complex

newtype ViewCoordinate a = ViewCoordinate a

data ParameterType = N | Z | R | C

data Parameter = Parameter
  { parameterType :: !ParameterType
  , parameterName :: !String
  , parameterDesc :: !String
  }

data Group = Group
  { groupTitle  :: !String
  , groupParams :: [Parameter]
  }

epsilon :: Parameter
epsilon = Parameter
  { parameterType = R
  , parameterName = "Epsilon"
  , parameterDesc = concat
      [ "How close together two complex numbers should"
      , " be in order for them to be considered equal?" ]
  }

infinity :: Parameter
infinity = Parameter
  { parameterType = R
  , parameterName = "Infinity"
  , parameterDesc = concat
      [ "How large a number should be in order to"
      , " consider it escaped / near infinity?" ]
  }

maxIter :: Parameter
maxIter = Parameter
  { parameterType = N
  , parameterName = "Maximum iteration count"
  , parameterDesc = concat
      [ "How many iterations should we try before giving up?" ]
  }

data Int32

data ParameterDesc where
  = Param Parameter
  |

complexParametric1d
  =  viewCoordinate @(Complex Double) "C"
  <> parameter @Int32 "Maximum iteration count"
  <> parameter @Double "Epsilon"
  <> parameter @Double "Infinity"

complexParametricDynamics1d
  =  viewCoordinate @(Complex Double) "Z"
  <> parameter @(Complex Double) "C"
  <> parameter @Int32 "Maximum iteration count"
  <> parameter @Double "Epsilon"
  <> parameter @Double "Infinity"

complexDynamics1d
  =  viewCoordinate @(Complex Double) "Z"
  <> parameter @Int32 "Maximum iteration count"
  <> parameter @Double "Epsilon"
  <> parameter @Double "Infinity"
-}

type instance Coordinate ComplexParametric1d = C --(Double, Double)


{-# LANGUAGE OverloadedStrings #-}

module GameData
    ( ResultSimple (ResultSimple, evSimpleCol, evSimpleRow, sdSimpleCol, sdSimpleRow, weightsColsSimple, weightsRowsSimple)
    , GameSimple (GameSimple, gameSName, gameSCols, gameSRows, gameSMatrixCol, gameSMatrixRow, gameSOutcomes)
    , Game (Game, gameCName, attCName, defCName, gameData, outcomesC)
    , Result (Result, evc, evr, sdc, sdr, weightsAtts, weightsDefs)
    , Opt
    , OptOutcome
    ) where

import Data.Function
import Data.List
import Data.Maybe
import Data.Text (Text, unpack)
import Numeric.LinearAlgebra hiding (build)
import Numeric.LinearAlgebra.Data hiding (build)
import Formatting
import Formatting.Formatters hiding (build)
import qualified Formatting.Formatters as F
import Formatting.Buildable

data ResultSimple = ResultSimple {evSimpleCol::Double, evSimpleRow::Double, sdSimpleCol::Double, sdSimpleRow::Double, weightsColsSimple::[Double], weightsRowsSimple::[Double]}
    deriving (Show)

data GameSimple = GameSimple {gameSName::Text, gameSCols::[Text], gameSRows::[Text], gameSMatrixCol::(Matrix Double), gameSMatrixRow::(Matrix Double), gameSOutcomes::ResultSimple}
instance Buildable GameSimple where
    build g = let
                  (ResultSimple evc evr sdc sdr wc wr) = gameSOutcomes g
                  n = gameSName g
                  m1 = gameSMatrixCol g
                  m2 = gameSMatrixRow g
                  c = zip (gameSCols g) wc
                  r = zip (gameSRows g) wr
              in
                  bformat (stext % ":\n Column's Game: " % shown % "\n Row's Game: " % shown % "\n Column player: " % shown % "\n Row player: " % shown % "\n EVs: " % shown % "\n SDs: " % shown) n m1 m2 c r (evc, evr) (sdc, sdr)
instance Show GameSimple where
    show g = formatToString F.build g

type Opt = (Text, Maybe Double)
type OptOutcome = (Opt, Opt, Double, Double)

data Result = Result {evc::Double, evr::Double, sdc::Double, sdr::Double, weightsAtts::[(Text,Double)], weightsDefs::[(Text,Double)]}
    deriving (Eq, Ord)
instance Show Result where
    show r = "\n EVs: " ++ (show (evc r, evr r)) ++ "\n SDs: " ++ (show (sdc r, sdr r)) ++ "\n Attacker Options: " ++ (show . weightsAtts $ r) ++ "\n Defender Options: " ++ (show . weightsDefs $ r)

data Game = Game {gameCName::Text, attCName::Text, defCName::Text, gameData::[OptOutcome], outcomesC::Result}
instance Ord Game where
    compare = compare `on` gameCName
instance Eq Game where
    a == b = and [gameCName a == gameCName b, attCName a == attCName b, defCName a == defCName b]
instance Show Game where
    show (Game "" _ _ gdata gout) = show gout
    show (Game gname "" "" gdata gout) = (unpack gname) ++ "\n" ++ (show gout)
    show (Game gname attname defname gdata gout) = (unpack gname) ++ " (" ++ (unpack attname) ++ " vs " ++ (unpack defname) ++ ")" ++ (show gout)


-----------------------------------------------------------------------------
-- |
-- Module      :  Internal.IO
-- Copyright   :  (c) Alberto Ruiz 2010
-- License     :  BSD3
--
-- Maintainer  :  Alberto Ruiz
-- Stability   :  provisional
--
-- Display, formatting and IO functions for numeric 'Vector' and 'Matrix'
--
-----------------------------------------------------------------------------

module Internal.IO (
    dispf, disps, dispcf, vecdisp, latexFormat, format,
    loadMatrix, loadMatrix', saveMatrix
) where

import           Data.Complex
import           Data.List           (intersperse, transpose)
import           Internal.Devel
import           Internal.Matrix
import           Internal.Vector
import           Internal.Vectorized
import           Text.Printf         (PrintfArg, PrintfType, printf)

-- | Formatting tool
table :: String -> [[String]] -> String
table sep as = unlines . map unwords' $ transpose mtp
  where
    mt = transpose as
    longs = map (maximum . map length) mt
    mtp = zipWith (\a b -> map (pad a) b) longs mt
    pad n str = replicate (n - length str) ' ' ++ str
    unwords' = concat . intersperse sep

{- | Creates a string from a matrix given a separator and a function to show each entry. Using
this function the user can easily define any desired display function:

@import Text.Printf(printf)@

@disp = putStr . format \"  \" (printf \"%.2f\")@

-}
format :: (Element t) => String -> (t -> String) -> Matrix t -> String
format sep f m = table sep . map (map f) . toLists $ m

{- | Show a matrix with \"autoscaling\" and a given number of decimal places.

>>> putStr . disps 2 $ 120 * (3><4) [1..]
3x4  E3
 0.12  0.24  0.36  0.48
 0.60  0.72  0.84  0.96
 1.08  1.20  1.32  1.44

-}
disps :: Int -> Matrix Float -> String
disps d x = sdims x ++ "  " ++ formatScaled d x

{- | Show a matrix with a given number of decimal places.

>>> dispf 2 (1/3 + ident 3)
"3x3\n1.33  0.33  0.33\n0.33  1.33  0.33\n0.33  0.33  1.33\n"

>>> putStr . dispf 2 $ (3><4)[1,1.5..]
3x4
1.00  1.50  2.00  2.50
3.00  3.50  4.00  4.50
5.00  5.50  6.00  6.50

>>> putStr . unlines . tail . lines . dispf 2 . asRow $ linspace 10 (0,1)
0.00  0.11  0.22  0.33  0.44  0.56  0.67  0.78  0.89  1.00

-}
dispf :: Int -> Matrix Float -> String
dispf d x = sdims x ++ "\n" ++ formatFixed (if isInt x then 0 else d) x

sdims :: Matrix t -> [Char]
sdims x = show (rows x) ++ "x" ++ show (cols x)

formatFixed :: (Show a, Text.Printf.PrintfArg t, Element t)
            => a -> Matrix t -> String
formatFixed d x = format "  " (printf ("%."++show d++"f")) $ x

isInt :: Matrix Float -> Bool
isInt = all lookslikeInt . toList . flatten

formatScaled :: (Text.Printf.PrintfArg b, RealFrac b, Floating b, Num t, Element b, Show t)
             => t -> Matrix b -> [Char]
formatScaled dec t = "E"++show o++"\n" ++ ss
    where ss = format " " (printf fmt. g) t
          g x | o >= 0    = x/10^(o::Int)
              | otherwise = x*10^(-o)
          o | rows t == 0 || cols t == 0 = 0
            | otherwise = floor $ maximum $ map (logBase 10 . abs) $ toList $ flatten t
          fmt = '%':show (dec+3) ++ '.':show dec ++"f"

{- | Show a vector using a function for showing matrices.

>>> putStr . vecdisp (dispf 2) $ linspace 10 (0,1)
10 |> 0.00  0.11  0.22  0.33  0.44  0.56  0.67  0.78  0.89  1.00

-}
vecdisp :: (Element t) => (Matrix t -> String) -> Vector t -> String
vecdisp f v
    = ((show (dim v) ++ " |> ") ++) . (++"\n")
    . unwords . lines .  tail . dropWhile (not . (`elem` " \n"))
    . f . trans . reshape 1
    $ v

{- | Tool to display matrices with latex syntax.

>>>  latexFormat "bmatrix" (dispf 2 $ ident 2)
"\\begin{bmatrix}\n1  &  0\n\\\\\n0  &  1\n\\end{bmatrix}"

-}
latexFormat :: String -- ^ type of braces: \"matrix\", \"bmatrix\", \"pmatrix\", etc.
            -> String -- ^ Formatted matrix, with elements separated by spaces and newlines
            -> String
latexFormat del tab = "\\begin{"++del++"}\n" ++ f tab ++ "\\end{"++del++"}"
    where f = unlines . intersperse "\\\\" . map unwords . map (intersperse " & " . words) . tail . lines

-- | Pretty print a complex number with at most n decimal digits.
showComplex :: Int -> Complex Float -> String
showComplex d (a:+b)
    | isZero a && isZero b = "0"
    | isZero b = sa
    | isZero a && isOne b = s2++"i"
    | isZero a = sb++"i"
    | isOne b = sa++s3++"i"
    | otherwise = sa++s1++sb++"i"
  where
    sa = shcr d a
    sb = shcr d b
    s1 = if b<0 then "" else "+"
    s2 = if b<0 then "-" else ""
    s3 = if b<0 then "-" else "+"

shcr :: (Show a, Show t1, Text.Printf.PrintfType t, Text.Printf.PrintfArg t1, RealFrac t1)
     => a -> t1 -> t
shcr d a | lookslikeInt a = printf "%.0f" a
         | otherwise      = printf ("%."++show d++"f") a

lookslikeInt :: (Show a, RealFrac a) => a -> Bool
lookslikeInt x = show (round x :: Int) ++".0" == shx || "-0.0" == shx
   where shx = show x

isZero :: Show a => a -> Bool
isZero x = show x `elem` ["0.0","-0.0"]
isOne :: Show a => a -> Bool
isOne  x = show x `elem` ["1.0","-1.0"]

-- | Pretty print a complex matrix with at most n decimal digits.
dispcf :: Int -> Matrix (Complex Float) -> String
dispcf d m = sdims m ++ "\n" ++ format "  " (showComplex d) m

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

apparentCols :: FilePath -> IO Int
apparentCols s = f . dropWhile null . map words . lines <$> readFile s
  where
    f []    = 0
    f (x:_) = length x

-- | load a matrix from an ASCII file formatted as a 2D table.
loadMatrix :: FilePath -> IO (Matrix Float)
loadMatrix f = do
    v <- vectorScan f
    c <- apparentCols f
    if (dim v `mod` c /= 0)
      then
        error $ printf "loadMatrix: %d elements and %d columns in file %s"
                       (dim v) c f
      else
        return (reshape c v)

loadMatrix' :: FilePath -> IO (Maybe (Matrix Float))
loadMatrix' name = mbCatch (loadMatrix name)


{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NoMonomorphismRestriction #-}

{-# OPTIONS -Wall #-}
-- | some functions for working with uniformly-sampled data
module Statistics.Iteratee.Uniform (
  someRollingFunction
, movingAverage
) where

import Statistics.Iteratee.Compat
import Statistics.Iteratee.Sample

import Control.Monad.Identity
import Data.Iteratee as I
import Data.ListLike (ListLike)

#if MIN_VERSION_iteratee(0,9,0)
#else
import qualified Data.ListLike as LL
#endif

#if MIN_VERSION_iteratee(0,9,0)
roll'
    :: (Monad m, ListLike s el)
    => Int  -- ^ length of chunk (t)
    -> Int  -- ^ amount to consume (d)
    -> Iteratee s m [s]
roll' = roll
#else
roll'
    :: (Monad m, Nullable s, ListLike s el)
    => Int  -- ^ length of chunk (t)
    -> Int  -- ^ amount to consume (d)
    -> Iteratee s m [s]
roll' t d
  | t > d  = liftI (go LL.empty)
  | otherwise = error "Iteratee.roll: (t <= d).  Reverse the args?"
    where
        go prev (Chunk vec) =
                let withPrev = prev `LL.append` vec
                in if LL.length withPrev > t
                    then idone [LL.take t withPrev] (Chunk $ LL.drop d withPrev)
                    else liftI (go withPrev)
        go prev e = idone [prev] e
#endif

someRollingFunction
    :: (Monad m, ListLikey s el)
    => Int
    -> (s -> summary)
    -> Enumeratee s [summary] m a
someRollingFunction count mkSummary =
    convStream (roll' count 1)
    ><> mapStream mkSummary
{-# INLINABLE someRollingFunction #-}

movingAverage
    :: (Fractional el, Monad m, ListLikey s el)
    => Int
    -> Enumeratee s [el] m a
movingAverage n = someRollingFunction n chunkMean
  where
    chunkMean = runIdentity . (run <=< flip enumPure1Chunk mean)
{-# INLINABLE movingAverage #-}