{-# LANGUAGE Rank2Types #-}
import Numeric.AD 
import Control.Applicative
import Control.Monad
import Data.Foldable as F
import Data.Traversable
import Control.Applicative
import Control.Monad.State
import Control.Monad as M
import System.IO
import Data.Char as C
import Data.Vector as V
import Data.List as L
import Data.Map as Map
import Data.Monoid
import Data.Reflection
import Data.Array.Accelerate.Interpreter as I
import Data.Random.Normal
import System.Random
import Numeric.AD.Internal.Reverse as Reverse
import Data.Time.Clock

main = time $ test 200 100 

{- 
  This function changes frequently and is used to test all the others.

  Currently it takes in the number of timesteps to iterate over along with 
  the number iterations of gradient descent to use. 

  It then instances a random set of network parameters and optimizes them.
  
  It then prints a list of the largest parameters in the neural network at
  every time step, along with the log likelihoods assigned to the training
  dataset by the the network before and after the optimisation.
-}
test iterations dataSize =
    do 

        biginfo  <- getData address

        let info = L.take dataSize biginfo
            elems = L.length . Map.toList $ ctiMap
            (ctiMap,itcMap) = makeConverters info
            nn = stdGRU elems
	    initMomentum = stdGRU elems
            llh = testGRU info
            grad = myGrad (testGRU info) nn 
            --opti = gradientUp (testGRU info) nn iterations 0.01 :: [(GRUParams Float,GRUParams Float)] 
            --opti = momentumGradientAscent (testGRU info) nn initMomentum 0.000001 (0.97) :: [(GRUParams Float, GRUParams Float,Float)]
            opti = optiGrad (testGRU info) nn initMomentum 0.000001 (0.98) :: [(GRUParams Float, GRUParams Float,Float,Float)]
        
        printWith "The size of the network is :" $ getCount nn 
        putStrLn . (\x -> "The initial LLH is: " L.++ show x) $ llh nn
        putStrLn . show $ (ctiMap,itcMap)
        putStrLn . show . fmap (\(x,y,z,a) -> (getNorm x,getNorm y,z,a)) $ opti
        putStrLn . (\x -> "The final LLH is: " L.++ show x) $ llh (L.last . fmap (\(x,_,_,_) -> x) $ opti)

        return ()

printWith :: Show a => String -> a -> IO ()
printWith msg x = putStrLn (msg L.++ show x)

time :: IO a -> IO ()
time x =
    do
        t1 <-getCurrentTime
	x
        t2 <- getCurrentTime
	putStrLn . show $ diffUTCTime t2 t1

--Returns a string representation of a file at a given location
getData :: String -> IO String
getData address = do
    handle <- openFile address ReadMode
    hGetContents handle


--Hard coded adress of the training data
address = "/Users/curtishuebner/Documents/Development/haskell/ghnn/Data/mobyDick.txt"

{-
  Parmeters: An object that can be folded over
  Returns: the largest numeric value contained within
-}
getMax :: (Foldable f,Num a,Ord a) => f a -> a
getMax= F.foldl (\a b-> if (abs a >abs b) then abs (a) else b)  0

getNorm :: (Foldable f, Floating a) => f a -> a 
getNorm = sqrt . getSum . F.foldMap (\x -> Sum (x*x))

getCount :: (Num a,Foldable f) => f x -> a
getCount = getSum . F.foldMap (\_ -> Sum 1)
{-
  This function produces a stadard set of network parameters
  when given the input size of the network
-}
stdGRU :: Floating a => Int -> GRUParams a
stdGRU ioSize = initGRU ioSize 500 ioSize (1,2,10) 312

{-
  Parameters: a string 
  
  Returns: maps that allow one to freely
  convert from a list of characters contained in the string
  to a list of integers and back again
-}
makeConverters :: String -> (Map.Map Char Int,Map.Map Int Char)
makeConverters x = (charToInt x, intToChar x)


{-
  Parameters: A list of integers along with a vector length
  
  Returns: a list of vectors initialized to 0 except for the
  element coresponding to the corresponding integer in the
  first list
-}
sparsify :: Num a => [Int] -> Int -> [Vector a]
sparsify list length = fmap (\x -> V.replicate length 0 // [(x,1)]) list

--returns a sorted list of numbers corresponding to characters in a string
values contents = sort . nub $ fmap C.ord contents


--unsafely applies a map to the elements of an object
conv ::(Functor f,Ord x) => Map.Map x y -> f x -> f y
conv map = fmap $ (\(Just x) -> x) . (\y -> Map.lookup y map)


{-
  This function takes a string and creates input output
  pairs that can be used to train a reccurent neural network
  along with the maximum size that the values inside the IOpairs
  can take on 
-}
makeIOPairs :: [Char] -> ([Int],[Int],Int)
makeIOPairs charValues =
    let 
        (ctiMap,_) = makeConverters charValues
        intValues = conv ctiMap charValues 
        inputs = L.init intValues
        outputs = L.tail intValues
        ioWidth = L.length $ values charValues
    in
        (inputs,outputs,ioWidth)

--Evaluates the preformance of a given neural network given a set of IOPairs and an initial vector
evalNN ::(Floating a) => ([Int],[Int],Int) -> ((Vector a,[Vector a]) -> (Vector a,[Vector a])) -> Vector a -> (a,Vector a)
evalNN (inputs,outputs,ioWidth) nn initialState = 
    let
        vectorInputs = sparsify inputs ioWidth
        (newState,vectorOutputs) = nn (initialState,vectorInputs)
        targetValuePairs = L.zip vectorOutputs outputs
        llh = sumSoftMaxLLH targetValuePairs
    in
        (llh,newState)


--Takes a set of GRU parameters and returns a neural network
makeGRU ::Floating a => GRUParams a -> (Vector a,[Vector a]) -> (Vector a,[Vector a])
makeGRU = passOver . gru


{-
  Parameters: training data and a set of network parameters
  Returns: the log likelyhood assigned by the resulting network to the data
  (a.k.a the number of bits required to encode the data)

  I find this function notable as it maps directly from parameters and 
  training data to log likelyhood, and so lends itself to easy optimization
-}
testGRU :: Floating a => [Char] -> GRUParams a -> a
testGRU inputs parameters =
    let
        ioPairs = makeIOPairs inputs
        nn = makeGRU parameters
        width = stateSize parameters
        initalVector = V.replicate width 0
    in
        fst $ evalNN ioPairs nn initalVector

    
--Map from possible characters in a string to integers
charToInt :: String -> Map.Map Char Int
charToInt values = 
    let 
        possibleChars = nub values
        pairs = L.zip possibleChars [0..]
    in
        Map.fromList pairs
    
--Map from integers to characters
intToChar :: String -> Map.Map Int Char
intToChar values =
    let
        possibleChars = nub values
        pairs = L.zip [0..] possibleChars
    in
        Map.fromList pairs

{-
  Applies the log likelyhood funcition to a list
  and sums the results
-}
sumLLH :: Floating a => [V.Vector (a,Bool)] -> a
sumLLH = L.sum . L.map logLikelyhood

{-
  
  Parameters: A vector of (float,boolean) pairs
  
  This function interprets the pairs as a probability 
  that a given value is true value and a corresponding 
  correct answer. It then computes the product of the probalbilites
  that had a True value as a pair along with (1-) the probabilities of
  values that had a False as a pair
  
  Returns: the assigned log likelyhood
-}
logLikelyhood :: Floating a => V.Vector (a,Bool) -> a
logLikelyhood = V.sum . V.map (\(x,y) -> if y then log(x) else log(1-x))

{-
  applies the softMaxLLH function to a vector and sums the results
-}
sumSoftMaxLLH :: (Floating a) => [(V.Vector a,Int)] -> a
sumSoftMaxLLH = L.sum . fmap softMaxLLH 

{-
  Applies the soft max function to vector integer pairs and
  interprets the result as a probability assinged to a given integer.
-}
softMaxLLH :: (Floating a) => (V.Vector a, Int) -> a
softMaxLLH (vector,index) = 
    let
        logNumerator = vector V.! index
        logDenominator = log . V.sum . fmap exp $ vector
    in
        logNumerator - logDenominator


--An experiment in order to understand the type signatue of the graident descent function
myGrad ::(Traversable f,Floating b) => (forall a. (Floating a) => f a -> a) -> f b -> f b 
myGrad a b = grad a b

{-
  An implementation of gradient descent.
  
  Parameters: a many to one function,
	      the starting point of the gradient descent,
	      an integer amount of steps,
	      the step size

  Returns: a list of points in the space of possible inputs with items farther in the list better maximising the function
  
  TODO: eliminate recomputation of initVal
-}
gradientUp :: (Ord a, Floating a ,Traversable f, Applicative f) => (forall x. Floating x => f x -> x) -> f a -> Int -> a -> [(f a,f a)]
gradientUp _ initParams 0 _ = []
gradientUp func initParams iterNum alpha=
    let 
        (initVal,gradient) = grad' func initParams
        newParams = (+) <$> initParams <*> ((alpha * (1) *) <$> gradient)
        newVal = func newParams
    in
        if (newVal > initVal) then
            (newParams,gradient) : gradientUp func newParams (iterNum-1) (alpha * 1.05)
        else
            (newParams,gradient) : gradientUp func initParams iterNum (alpha / 2)

momentumGradientAscent :: (Ord a,Floating a, Traversable f, Applicative f) => 
    (forall x. Floating x => f x -> x) -> f a -> f a -> a -> a -> [(f a,f a,a)]
momentumGradientAscent func initParams momentum alpha beta =
    let 
        (initVal,gradient) = grad' func initParams
        newMomentum =fmap (*beta) $ (+) <$> momentum <*> ((alpha*) <$> gradient)
        newParams = (+) <$> momentum <*> initParams
	theta = acos $ dotProduct momentum gradient / (getNorm momentum * getNorm gradient)
    in
        if (getNorm gradient > 0.001) then
	    (newParams,gradient,theta) : momentumGradientAscent func newParams newMomentum alpha beta
	else
	    []

optiGrad :: (Ord a, Floating a, Traversable f, Applicative f) =>
    (forall x. Floating x => f x -> x) -> f a -> f a -> a -> a -> [(f a,f a,a,a)]
optiGrad func initParams momentum alpha beta =
    let 
        (initVal,gradient) = grad' func initParams
        truncatedMomentum = 
	    if dotProduct momentum gradient > 0
	    then
	        --momentum
		momentum `onto` gradient
	    else 
	        --planarProj momentum gradient
		pure 0
        newMomentum = fmap (*beta) $ (+) <$> momentum <*> ((alpha*) <$> gradient)
        newParams = (+) <$> newMomentum <*> initParams
        theta = acos $ dotProduct momentum gradient / (getNorm momentum * getNorm gradient)
    in
        if (getNorm gradient > 1) then
            (newParams,gradient,theta / pi,initVal) : optiGrad func newParams newMomentum alpha beta
        else
            []

{-
  Parameters: a function from a (state,input) pair to a new (state,output) pair
  
  Returns: the same function applied to a list of values with the states
  chained together
  
  TODO: reimplent this using the state monad and the traversable typeclass
-}
passOver :: ((a,b) -> (a,c)) -> (a,[b]) -> (a,[c])
passOver f (state,inputs@(x:xs)) = 
    let
        (nextState,nextOutput) = f (state,x)
        (finalState,otherOutput) = passOver f (nextState,xs)
    in
        (finalState,nextOutput:otherOutput)
passOver _ (state,[]) = (state,[])

--sparse vector matrix multiplication
--TODO: swich the arguments around in acordance with linear algebra conventions
(><) :: Num a => V.Vector a -> Matrix a -> V.Vector a
vector >< SMatrix matrix =
    let
        multiplyElement (factor,index) = vector V.! index * factor 
        multiplyRow = L.sum . fmap multiplyElement
    in
        fmap multiplyRow matrix

--ElementWise function application for vectors
infixl 8 `vecApply`
vecApply :: V.Vector (a -> b) -> V.Vector a -> V.Vector b
vecApply func a = V.zipWith (\x y -> x y) func a  

{-
  An experimental differentiable branching operation between vectors to create turing-complete neural networks
-}
vectorBranch ::(Ord a, Floating a) => a -> (V.Vector a -> V.Vector a) -> (V.Vector a -> V.Vector a) -> V.Vector a -> V.Vector a
vectorBranch x f g args =
    let
        function = if x > 0 then f else g
        transferFunction = 1 / (1 + x**2)
        newFraction = (*transferFunction) <$> function args
        oldFraction = ((*) $ 1-transferFunction) <$> args
    in
        (vecApply $ fmap (+) newFraction) oldFraction --TODO: Clean this up

--Defining Zipvectors
newtype ZipVec a = ZipVec (V.Vector a)
    deriving (Show, Eq, Ord, Read)

getZipVec (ZipVec a) = a

instance Functor ZipVec where
    fmap f (ZipVec a) = ZipVec . V.map f $ a

instance Applicative ZipVec where
    pure = ZipVec . V.unfoldr (\x -> Just (x,x))

    ZipVec a <*> ZipVec b = ZipVec $ a `vecApply` b

instance Foldable ZipVec where
    foldMap f (ZipVec vec) = F.fold $ fmap f vec
    foldr f a (ZipVec vec) = V.foldr f a vec

instance Traversable ZipVec where
    sequenceA (ZipVec a)= ZipVec <$> sequenceA a


--Defining Sparce and uniform matrices
data Matrix a = SMatrix (V.Vector [(a,Int)]) | UMatrix a
    deriving (Show,Eq)

instance Functor Matrix where
    fmap f (SMatrix matrix) =SMatrix . (V.map . L.map $ (\(a,x) -> (f a,x))) $ matrix

instance Applicative Matrix where
    pure = UMatrix

    UMatrix f <*> UMatrix a = UMatrix $ f a 
    UMatrix f <*> matrix@(SMatrix a)= fmap f matrix
    SMatrix f <*> SMatrix a = SMatrix $ V.zipWith applyRow f a
        where
            applyRow = L.zipWith (\(f,a) (x,_) -> (f x, a))

instance Foldable Matrix where
    foldMap = foldMapDefault

instance Traversable Matrix where
    sequenceA (SMatrix a) = SMatrix <$> seqVec a
        where
            seqVec = sequenceA . fmap seqRow
            seqRow = sequenceA . fmap seqVal
            seqVal (val,index) = (\x -> (x,index)) <$> val


{-
  Parameters: number of columns and rows along with the number of non zero elements

  Returns: A stateful computation that once run will produce a matrix with the given
  parameters.
-}
genSpMtx :: (Num a, Enum a, RandomGen s) => Int -> Int -> Int -> State s (Matrix a)
genSpMtx columns rows elemPerRow = newMatrix
    where
        newMatrix = matrixGenerator >>= \y -> return . SMatrix $ y
        matrixGenerator = V.replicateM rows genRow 
        genRow =  
            let 
                vals = [0,0..] 
                indicies = knuthShuffleN range elemPerRow
                row = indicies >>= \x -> return $ L.zip vals x
            in
                row
        range = (0,columns - 1) :: (Int,Int)

--Apply gaussian noise to the numeric values of a given traversable object
randPerturb x mean stdDev = M.forM x f
    where
        f x = (state $ normal) >>= \y -> return (x + mean + y * stdDev)

{-
  Takes in a range of values and produces a stateful computaton
  that when run returns the first n elements of a randomly shuffled
  list containing all the elements in the range
-}
knuthShuffleN :: (Random a, Enum a, RandomGen g, Num b, Ord b) => (a,a) -> b -> State g [a]
knuthShuffleN range n =
    let 
        initList = (\(a,b) -> [a..b]) range
        initListLength = L.length initList

        subShuffle ::(Integral a, Random a, Enum a, Num b, Ord b, RandomGen g) => [t] -> a -> b -> State g [t]
        subShuffle [] lenght n = return []
        subShuffle list lenght n =
            if (n > 0)
            then
                do 
                    index <- randomindex lenght
                    let (newList,val) = list -!! index
                    rest <- subShuffle newList (lenght-1) (n-1)
                    let retVal = val:rest
                    return retVal
            else
                return []

        randomindex :: (Num a,Random a,Enum a,RandomGen g) => a -> State g a    
        randomindex lenght = state $ randomR (0,lenght-1)

        (-!!) :: Integral b => [a] -> b -> ([a],a)
        a -!! b = let (x,(y:ys)) = L.genericSplitAt b a in (x L.++ ys,y)
    in
        subShuffle initList initListLength n


--Instancing the gated reccurent unit dataStructure
data GRUParams a = GRUParams 
    { inputToChange :: Matrix a
    , stateToChange :: Matrix a
    , inputToGate :: Matrix a
    , stateToGate :: Matrix a
    , stateToOutput :: Matrix a
    , changeBias :: ZipVec a
    , gateBias :: ZipVec a
    , outputBias :: ZipVec a
    , inputSize :: Int
    , outputSize :: Int
    , stateSize :: Int
    }deriving (Show,Eq)

instance Functor GRUParams where
    fmap = fmapDefault

instance Foldable GRUParams where
    foldMap = foldMapDefault

{-
  Note: if traversable is instanced using sequenceA
  and fmap is implemented with fmapDefault,
  any calls to traverse will cause a stack overflow.
  This is because traverse will call fmap
  which calls fmapDefault
  which calls traverse again
-}
instance Traversable GRUParams where
    traverse f (GRUParams iTC sTC iTG sTG sTO cB gB oB iS oS sS) =
        GRUParams <$>
        traverse f iTC <*>
        traverse f sTC <*>
        traverse f iTG <*>
        traverse f sTG <*>
        traverse f sTO <*>
        traverse f cB  <*>
        traverse f gB  <*>
        traverse f oB  <*>
        pure iS <*>
        pure oS <*>
        pure sS

--TODO: instance this properley later
instance Applicative GRUParams where
    pure f = 
        let 
            g = (pure f) 
            h = (pure f) 
        in 
            GRUParams g g g g g h h h undefined undefined undefined
    GRUParams a1 a2 a3 a4 a5 a6 a7 a8 _ _ _ <*>  GRUParams b1 b2 b3 b4 b5 b6 b7 b8 x y z =
        GRUParams (a1 <*> b1) (a2 <*> b2) (a3 <*> b3) (a4 <*> b4) (a5 <*> b5) (a6 <*> b6) (a7 <*> b7) (a8 <*> b8) x y z
{-
  For the confused:
  xTY = x to y
  i = input
  s = state
  o = output
  c = change
  g = gate
  B = bias
  EPR = elements per row

  Parameters: sizes for the input state and output along with 
  a tuple describing the number of elements per row of each matrix 
  and a seed value

  Returns: a slightly randomized set of gated reccurent unit parameters
-}
initGRU ::(Floating a) => Int -> Int -> Int -> (Int,Int,Int) -> Int -> GRUParams a
initGRU input state output (iTSEPR,sTSEPR,sTOEPR) seed =
  let
    stateGRU = 
      do
        iTC <- genSpMtx input state iTSEPR
        sTC <- genSpMtx state state sTSEPR
        iTG <- genSpMtx input state iTSEPR
        sTG <- genSpMtx state state sTSEPR
        sTO <- genSpMtx state output sTOEPR
        let 
          cB = ZipVec $ V.replicate state 0
          gB = ZipVec $ V.replicate state 0
          oB = ZipVec $ V.replicate output 0
          blankGRU = GRUParams iTC sTC iTG sTG sTO cB gB oB input output state
        return blankGRU
        randPerturb blankGRU 0 (1 * 10**(-3))
    generator = mkStdGen seed
  in
    fmap floatingFromDouble . fst . runState stateGRU $ generator

{-
  Parameters: the parameters to a gated reccurent unit

  Returns: a funtion that maps between an (input,state) pair of vectors
  and an (output,newState) pair of vectors
-}
gru :: (Floating a) => GRUParams a -> (V.Vector a, V.Vector a) -> (V.Vector a, V.Vector a)
gru parameters (state,input) = (newState, output)
    where
        zipState = ZipVec state
        gate = 
            (\a b c -> logit $ a + b + c) <$> 
            gateBias parameters <*>
            ZipVec (input >< inputToGate parameters) <*>
            ZipVec (state >< stateToGate parameters) 
        change = 
            (\a b c d->cos (a + b + c) * d) <$> 
            changeBias parameters <*>
            ZipVec (input >< inputToChange parameters) <*>
            ZipVec (state >< stateToChange parameters) <*>
            gate
        newState = getZipVec $ (+) <$> zipState <*> change
        output = newState >< stateToOutput parameters

--Generic transfer functions
logit x = 1 / (1 + exp(-x))
quadraticSigmoid x = x/(sqrt(x**2+1))


--Functions for converting from polymorphic floats to doubles and back
floatingFromDouble :: Floating n => Double -> n
floatingFromDouble i = realToFrac i

doubleFromFloating :: (Real n,Floating n) => n -> Double 
doubleFromFloating i = realToFrac i

dotProduct x y = getSum . F.fold . fmap Sum $ ((*) <$> x <*> y)
onto x y = fmap (*(dotProduct x y / (getNorm y ^ 2))) y
planarProj x y = (-) <$> x <*> (x `onto` y)
