{-# LANGUAGE AllowAmbiguousTypes, DataKinds, RankNTypes, TypeFamilies,
             TypeOperators #-}
{-# OPTIONS_GHC -Wno-missing-export-lists #-}
{-# OPTIONS_GHC -fconstraint-solver-iterations=16 #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
-- | Shaped tensor-based implementation of Convolutional Neural Network
-- for classification of MNIST digits. Sports 2 hidden layers.
module HordeAd.Tool.MnistCnnShaped where

import Prelude

import qualified Data.Array.DynamicS as OT
import           Data.Array.Internal (valueOf)
import qualified Data.Array.Shape
import qualified Data.Array.ShapedS as OS
import           Data.Proxy (Proxy)
import qualified Data.Vector.Generic as V
import           GHC.TypeLits (KnownNat, type (+), type (<=), type Div)
import qualified Numeric.LinearAlgebra as HM

-- until stylish-haskell accepts NoStarIsType
import qualified GHC.TypeLits

import HordeAd.Core.DualNumber
import HordeAd.Core.Engine
import HordeAd.Core.PairOfVectors (DualNumberVariables, varS)
import HordeAd.Tool.MnistData

convMnistLayerS
  :: forall kheight_minus_1 kwidth_minus_1 out_channels
            in_height in_width in_channels batch_size d r m.
     ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1, KnownNat out_channels
     , KnownNat in_height, KnownNat in_width
     , KnownNat in_channels, KnownNat batch_size
     , 1 <= kheight_minus_1
     , 1 <= kwidth_minus_1  -- wrongly reported as redundant
     , DualMonad d r m )
  => DualNumber d (OS.Array '[ out_channels, in_channels
                             , kheight_minus_1 + 1, kwidth_minus_1 + 1 ] r)
  -> DualNumber d (OS.Array '[batch_size, in_channels, in_height, in_width] r)
  -> DualNumber d (OS.Array '[out_channels] r)
  -> m (DualNumber d (OS.Array '[ batch_size, out_channels
                                , (in_height + kheight_minus_1) `Div` 2
                                , (in_width + kwidth_minus_1) `Div` 2 ] r))
convMnistLayerS ker x bias = do
  let yConv = conv24 ker x
      replicateBias
        :: DualNumber d (OS.Array '[] r)
        -> DualNumber d (OS.Array '[ in_height + kheight_minus_1
                                   , in_width + kwidth_minus_1 ] r)
      replicateBias = konstS . fromS0
      biasStretched = ravelFromListS
                      $ replicate (valueOf @batch_size)
                      $ mapS replicateBias bias
        -- TODO: this is weakly typed; add and use replicateS instead
        -- or broadcastS or stretchS, possibly with transposeS?
  yRelu <- reluAct $ yConv + biasStretched
  maxPool24 @1 @2 yRelu

convMnistTwoS
  :: forall kheight_minus_1 kwidth_minus_1 num_hidden out_channels
            in_height in_width in_channels batch_size d r m.
     ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1
     , KnownNat num_hidden, KnownNat out_channels
     , KnownNat in_height, KnownNat in_width, KnownNat batch_size
     , in_channels ~ 1
     , 1 <= kheight_minus_1
     , 1 <= kwidth_minus_1
     , DualMonad d r m )
  => OS.Array '[batch_size, in_channels, in_height, in_width] r
  -- All below is the type of all paramters of this nn. The same is reflected
  -- in the length function below and read from variables further down.
  -> DualNumber d (OS.Array '[ out_channels, in_channels
                             , kheight_minus_1 + 1, kwidth_minus_1 + 1 ] r)
  -> DualNumber d (OS.Array '[out_channels] r)
  -> DualNumber d (OS.Array '[ out_channels, out_channels
                             , kheight_minus_1 + 1, kwidth_minus_1 + 1 ] r)
  -> DualNumber d (OS.Array '[out_channels] r)
  -> DualNumber d (OS.Array '[ num_hidden
                             , out_channels
                                 GHC.TypeLits.*
                                   (((in_height + kheight_minus_1) `Div` 2
                                     + kheight_minus_1) `Div` 2)
                                 GHC.TypeLits.*
                                   (((in_width + kwidth_minus_1) `Div` 2
                                     + kwidth_minus_1) `Div` 2)
                             ] r)
  -> DualNumber d (OS.Array '[num_hidden] r)
  -> DualNumber d (OS.Array '[SizeMnistLabel, num_hidden] r)
  -> DualNumber d (OS.Array '[SizeMnistLabel] r)
  -> m (DualNumber d (OS.Array '[SizeMnistLabel, batch_size] r))
convMnistTwoS x ker1 bias1 ker2 bias2
              weigthsDense biasesDense weigthsReadout biasesReadout = do
  t1 <- convMnistLayerS ker1 (constant x) bias1
  t2 <- convMnistLayerS ker2 t1 bias2
  let m1 = mapS reshapeS t2
      m2 = transpose2S m1
      denseLayer = weigthsDense <>$ m2 + asColumnS biasesDense
  denseRelu <- reluAct denseLayer
  returnLet $ weigthsReadout <>$ denseRelu + asColumnS biasesReadout

convMnistLenS
  :: forall kheight_minus_1 kwidth_minus_1 num_hidden out_channels
            in_height in_width.
     ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1
     , KnownNat num_hidden, KnownNat out_channels
     , KnownNat in_height, KnownNat in_width )
  => Proxy kheight_minus_1
  -> Proxy kwidth_minus_1
  -> Proxy num_hidden
  -> Proxy out_channels
  -> Proxy in_height
  -> Proxy in_width
  -> (Int, [Int], [(Int, Int)], [OT.ShapeL])
convMnistLenS _ _ _ _ _ _ =
  ( 0
  , []
  , []
  , [ Data.Array.Shape.shapeT @'[ out_channels, 1
                                , kheight_minus_1 + 1, kwidth_minus_1 + 1 ]
    , Data.Array.Shape.shapeT @'[out_channels]
    , Data.Array.Shape.shapeT @'[ out_channels, out_channels
                                , kheight_minus_1 + 1, kwidth_minus_1 + 1 ]
    , Data.Array.Shape.shapeT @'[out_channels]
    , Data.Array.Shape.shapeT @'[ num_hidden
                                , out_channels
                                    GHC.TypeLits.*
                                      ((in_height + kheight_minus_1) `Div` 2
                                       + kheight_minus_1) `Div` 2
                                    GHC.TypeLits.*
                                      ((in_width + kwidth_minus_1) `Div` 2
                                       + kheight_minus_1) `Div` 2
                                ]
    , Data.Array.Shape.shapeT @'[num_hidden]
    , Data.Array.Shape.shapeT @'[SizeMnistLabel, num_hidden]
    , Data.Array.Shape.shapeT @'[SizeMnistLabel]
    ]
  )

convMnistS
  :: forall kheight_minus_1 kwidth_minus_1 num_hidden out_channels
            in_height in_width batch_size d r m.
     ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1
     , KnownNat num_hidden, KnownNat out_channels
     , KnownNat in_height, KnownNat in_width, KnownNat batch_size
     , 1 <= kheight_minus_1
     , 1 <= kwidth_minus_1
     , DualMonad d r m )
  => OS.Array '[batch_size, 1, in_height, in_width] r
  -> DualNumberVariables d r
  -> m (DualNumber d (OS.Array '[SizeMnistLabel, batch_size] r))
convMnistS x variables = do
  let ker1 = varS variables 0
      bias1 = varS variables 1
      ker2 = varS variables 2
      bias2 = varS variables 3
      weigthsDense = varS variables 4
      biasesDense = varS variables 5
      weigthsReadout = varS variables 6
      biasesReadout = varS variables 7
  convMnistTwoS @kheight_minus_1 @kwidth_minus_1 @num_hidden @out_channels
                x ker1 bias1 ker2 bias2
                weigthsDense biasesDense weigthsReadout biasesReadout

convMnistLossFusedS
  :: forall kheight_minus_1 kwidth_minus_1 num_hidden out_channels
            in_height in_width batch_size d r m.
     ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1
     , KnownNat num_hidden, KnownNat out_channels
     , KnownNat in_height, KnownNat in_width, KnownNat batch_size
     , 1 <= kheight_minus_1
     , 1 <= kwidth_minus_1
     , DualMonad d r m )
  => Proxy kheight_minus_1
  -> Proxy kwidth_minus_1
  -> Proxy num_hidden
  -> Proxy out_channels
  -> ( OS.Array '[batch_size, in_height, in_width] r
     , OS.Array '[batch_size, SizeMnistLabel] r )
  -> DualNumberVariables d r
  -> m (DualNumber d r)
convMnistLossFusedS _ _ _ _ (glyphS, labelS) variables = do
  let xs :: OS.Array '[batch_size, 1, in_height, in_width] r
      xs = OS.reshape glyphS
  result <- convMnistS @kheight_minus_1 @kwidth_minus_1
                       @num_hidden @out_channels
                       xs variables
  let targets2 = HM.tr $ HM.reshape (valueOf @SizeMnistLabel)
                       $ OS.toVector labelS
  vec <- lossSoftMaxCrossEntropyL targets2 (fromS2 result)
  returnLet $ scale (recip $ fromIntegral (valueOf @batch_size :: Int))
            $ sumElements0 vec

-- For simplicity, testing is performed in mini-batches of 1.
-- See RNN for testing done in batches.
convMnistTestS
  :: forall kheight_minus_1 kwidth_minus_1 num_hidden out_channels
            in_height in_width r.
     ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1
     , KnownNat num_hidden, KnownNat out_channels
     , KnownNat in_height, KnownNat in_width
     , 1 <= kheight_minus_1
     , 1 <= kwidth_minus_1
     , IsScalar 'DModeGradient r )
  => Proxy r
  -> Proxy kheight_minus_1
  -> Proxy kwidth_minus_1
  -> Proxy num_hidden
  -> Proxy out_channels
  -> [( OS.Array '[in_height, in_width] r
      , OS.Array '[SizeMnistLabel] r )]
  -> Domains r
  -> r
convMnistTestS _ _ _ _ _ inputs parameters =
  let matchesLabels :: ( OS.Array '[in_height, in_width] r
                       , OS.Array '[SizeMnistLabel] r )
                    -> Bool
      matchesLabels (glyph, label) =
        let tx :: OS.Array '[1, 1, in_height, in_width] r
            tx = OS.reshape glyph
            nn :: DualNumberVariables 'DModeGradient r
               -> DualMonadValue r (DualNumber 'DModeGradient (OS.Array '[SizeMnistLabel, 1] r))
            nn = convMnistS @kheight_minus_1 @kwidth_minus_1
                            @num_hidden @out_channels
                            tx
            value = primalValue nn parameters
        in V.maxIndex (OS.toVector value) == V.maxIndex (OS.toVector label)
  in fromIntegral (length (filter matchesLabels inputs))
     / fromIntegral (length inputs)


{-# LANGUAGE DeriveFunctor #-}
module Main where

import Lib
import Data.Complex
import Data.Fixed (mod')
import Linear.Vector

data ChebFun a = ChebFun {pow2 :: Int, coefficients :: [a] } deriving (Show, Functor)

-- data SampleFun a =
-- interpolate :: Vector a => SampleFun a -> (Circ -> a) -- need to use linear inteperpolation
{-
instance Functor ChebFun where
	fmap f (ChebFun i cs) = ChebFun i (fmap f cs)
-}

{-
zipWith' f dx dy (x : xs) (y : ys) = (f x y) : (zipWith' f dx dy xs ys)
zipWith' f dx _ [] ys = map (f dx) ys
zipWith' f _ dy xs [] = map (flip f dy) xs
-}

unionWith' f (x : xs) (y : ys) = (f x y) : (unionWith' f xs ys)
unionWith' f [] ys = ys
unionWith' f xs [] = xs

-- Linear package makes scalar multiplication a piece of being a functor.
-- I'm not sure I can make ChebFun a Functor? Well, it can be for the purposes of scalar multiply.
-- But usually we need to FFT it to real space before applying pointwise stuff

instance Additive ChebFun where
	zero = ChebFun 0 []
	liftI2 f (ChebFun n xs) (ChebFun n' ys) = ChebFun (min n n') (zipWith f xs ys)
	liftU2 f (ChebFun n xs) (ChebFun n' ys) = ChebFun (max n n') (unionWith' f xs ys)

-- intance Additive ChebFun where
-- possible conventions: incorrect sized lists are error, everything past nil is 0, everything past last element is a repeat of last constant
-- Last convention allows for efficient representation of constant values. A pure nil is represented as 0. case on (a:[])
-- TrigFun is better name
-- ChebFun b a where ChebFun (b ~ Circ, Num a) =>

-- compose
-- Possible optimization
data Consty f a = NonConsty (f a) | Consty a

{-
we can lift pointwise functions?
cheblift

almost applicative, except that the object inside needs to be foureir transformable.

instance Applicative ChebFun where
	Num a b c => (a -> b -> c) -> ChebFun a -> ChebFun b -> ChebFun c

-}
-- chebcompose :: ChebFun Circ -> ChebFun Circ -> ChebFun Circ

{-
instance Num a => Num (ChubFun a) where
	(ChebFun n xs) + (ChebFun n' ys) = ChebFun (max n n') (zipWithDefault 0 (+) xs ys) -- not right. We want to add with default 0 is [] and then re-shorten.  also the min needs to be a max.
	(ChebFun n xs) * (ChebFun n' ys) = chebify $ (funcify x) * (funcify y) -- all pointwise stuff tends to be in the real domain
	abs x = chebify $ abs . (funicfy x)

instance Vector a => Vector (ChebFun a) where
	smul s = fmap (smul s) coefficients
	vadd x y = x + y

-}

newtype Circ = Circ Double
circ :: Double -> Circ -- smart constructor
circ x = Circ $ mod' (abs x) (2 * pi)

-- https://ro-che.info/articles/2015-12-04-fft
split :: [a] -> ([a], [a])
split = foldr f ([], [])
  where
    f a (r1, r2) = (a : r2, r1)

-- interleave (The opposite of split). Not super sure I got the order right.

-- liquid haskell: interleave split = split interleave
interleave (x : xs, ys) = x : (interleave (ys, xs))

-- only takes 2^n sized lists, or we can assume is zero. Split is incorrect possible then. If not an even number it will swap
fourier :: [Complex Double] -> [Complex Double]
fourier [] = []
fourier [a] = [a]
fourier xs =  (zipWith (+) ffte ffto') <> (zipWith (-) ffte ffto')  where
	        (evens, odds) = split xs
	        ffte = fourier evens
	        n = length ffte
	        ffto = fourier odds
	        ffto' = zipWith (*) ffto (twiddle n)

-- we may want to seperate out the real and imag part
data Fourier = Fourier {k0 :: Double, kcos :: [Double] , ksin :: [Double] }

-- could we get liquid haskell to verify this?

-- I can fourier any sequence Vector a => [a] -> [a]
-- Need scalar multiplication via fourier coefficents.
-- need vector add for + and - of
-- I can make ChebFun a functor. Which can be useful, but perhaps suspicious.

-- instance Num a => Vector a a

fourier' :: [Double] -> [Complex Double]
fourier' [] = []
fourier' [a] = [a :+ 0]
fourier' xs =  (zipWith (+) ffte ffto') <> (zipWith (-) ffte ffto')  where
	        (evens, odds) = split xs
	        ffte = fourier' evens
	        n = length ffte
	        ffto = fourier' odds
	        ffto' = zipWith (*) ffto (twiddle n)

twiddle :: Int -> [Complex Double]
twiddle n = map cis $ iterate (+ (pi / n')) 0 where n' = fromInteger $ toInteger n

-- only takes 2^n sized lists, or we can assume is zero. Split is incorrect possible then. If not an even number it will swap
ifourier :: [Complex Double] -> [Complex Double]
ifourier [] = []
ifourier [a] = [a]
ifourier xs =  (zipWith (+) ffte ffto') <> (zipWith (-) ffte ffto')  where
	        (evens, odds) = split xs
	        ffte = ifourier evens
	        n = length ffte
	        ffto = ifourier odds
	        ffto' = zipWith (*) ffto (map conjugate (twiddle n))
{-
{-# Rewrite (ifourier (fourier x) = x #-}
{-# Rewrite (fourier (ifourier x) = x #-}
	Do I need fourier . ifourier = id ? couldn't hurt I guess.
		These laws aren't strictly true, but pretty close. A nightmare for consistency, but will always make accuracy better.

-}
-- his is some bull.
ifft :: [Complex Double] -> [Complex Double]
ifft (x : xs) = fourier (x : (reverse xs))

fourierGood :: [Complex Double] -> Bool
fourierGood _ = True -- implement the envelope test

envelope :: (Ord a, Num a) => [a] -> [a] -- an envelope that replaces the list with the maximum of the list form then on
envelope = (scanr1 max) . (map abs)  -- (\x acc -> max x acc)

-- plateaujtest j env = env ! j >= env ! j2 * (1 - )   where j2 = 1.25*j + 5
-- hmm a constant 0 may be unhappy with this. probably shouldn't use it.
-- The maximum will be the head of the envelope by the way
-- normalize xs = map (/ m) xs where m = maximum xs

--normalize . envelope

-- The stupidest test. WHat if we are accidentally in a zero of the spectrum
-- maybe lsightly less stupid would be to take the maximum of the last 1/4.
-- These really are garbage. We should actually follow the chebfun prescirption.

simpleTest :: (Fractional a, Ord a) => [a] -> Bool
simpleTest xs = if (last absxs) >= (m * (1e-12)) then False else True where
	    absxs = fmap abs xs
	    m = maximum absxs

-- If we consider iterated ChebFuns, we may want different complexity depending on whether we are at low rank pieces or not.
-- so part of our chebify should recurse inside and chebify for every element.
-- I guess i roughly expect the needed coefficients to form a polytope ?
-- It will be this blob on nenzero coefficients. In higher dimension, surfacc is large, so cutting down on that is helpful.
-- Hmm. This isn't really hierarchical. Maybe some even odd splitting combined with left right splitting in real space?
-- I get a faint sensation I might want to look at a maxi-min? Maybe not.
{-
The new expansion test should decide on both whether to expand internally and at this level.
-- Chebify (Chebify a) => Chebify (Double -> a) where
  chebify f = (fmap chebify) f -

  Chebify Double where
  	chebify = id

-}
-- G(x,x') considered as a 2d function
-- 1/r^2 also should seperate out into peices

-- integrate -- from 0 to x
-- differentiate -- differentiation with resepct to the coordinate
-- functional_differenation -- = differntation with resepct to the value held.
-- sample :: Double -> ChebFun a -> a -- function application
-- fmap diff (For a composed chebfun this idffierentiates the third coordinate)
-- diff2 = diff . diff
-- 3d laplacian.
-- diff2 + fmap diff2 + fmap . fmap diff2

{-
increasingFourier :: [Complex Double] -> (Double -> Double) -> [Complex Double]
increasingFourier evens f | fourierGood evens = increasingFourier evens' f
                          | otherwise = evens where
                                               evens' = (zipWith (+) evens ffto') <> (zipWith (-) evens ffto')
                                               odds = sample f oddpositions
                                               ffto =  fourier odds
                                               twiddle =
                                               ffto' = zipWith      -- if evens acceptable, stop, otherwise

--- increasing fourier is chebify basically
chebify = increasingFourier []
-- sinc interpolation I guess? just take nearest point of reverse fft? Linear interpolation of fft? All of these corresopnd to filters, the choice of
	-- which should not matter much

-- sum_x0  sinc(x-x0) * f(x0)
funcify = \x -> Chebfun
-}

main :: IO ()
main = someFunc


import Foreign.Marshal.Array
import Data.Complex
import Foreign.Storable.Complex
import FFTW

main = do
    inA  <- fftwAllocComplex 1024
    outA <- fftwAllocComplex 1024

    plan <- planDFT1d 1024 inA outA Forward fftwEstimate

    pokeArray inA $ map (:+ 0) [0..1023]
    execute plan
    res <- peekArray 1024 outA

    fftwFree inA
    fftwFree outA

    print res


{-# LANGUAGE BangPatterns #-}

module SimParse where

import Data.List
import NEC
import System.Process
import System.Random
import Data.Complex

vswr :: Double -> (Double, Double) -> Double
vswr target = vswr' (target, 0)

vswr' :: (Double, Double) -> (Double, Double) -> Double
vswr' (reTar, imTar) (reLoad, imLoad) = swr where
  z0 = reTar :+ imTar
  zL = reLoad :+ imLoad
  refRatio = (zL - z0) / (zL + z0)
  (ro :+ _) = abs $ refRatio
  swr = (1 + ro)  / (1 - ro)

bandwidthInterval :: Double -> Double -> [[String]] -> (Double, Double)
bandwidthInterval targetImp swrThre lnss = (lo, hi) where
  swrs = map (\l -> (frequency l , vswr targetImp (impedance l))) lnss
  below = takeWhile (\s -> (snd s) <= swrThre) $ dropWhile (\s -> (snd s) > swrThre) swrs
  lo = fst $ head below
  hi = fst $ last below

bandwidth t s = (\(lo, hi) -> hi - lo) . bandwidthInterval t s

gain :: [String] -> Double
gain lns = read gn where
  gn = head $ drop 4 $ words gainline
  gainline =  head $ drop 5 $ dropWhile (not . isSubsequenceOf "RADIATION PATTERN") lns

impedance :: [String] -> (Double, Double)
impedance lns = (zr, zi) where
  zr = read $ head line
  zi = read $ head $ tail line
  line = drop 6 $ words paramsLine
  paramsLine = head $ drop 3 $ dropWhile (not . isSubsequenceOf "ANTENNA INPUT PARAMETERS") lns

frequency :: [String] -> Double
frequency lns = read freq where
  freq = head $ drop 2 $ words $ head $ dropWhile (not . isSubsequenceOf "FREQUENCY") $ lns

fitness targetImpedance lns = a * gn - b * (abs $ targetImpedance - reZ) - c * (abs imZ) where
  a = 40
  b = 2
  c = 2
  gn = gain lns
  (reZ, imZ) = impedance lns

fitnessImepd targetImpedance lns = - b * (abs $ targetImpedance - reZ) - c * (abs imZ) where
  b = 2
  c = 2
  gn = gain lns
  (reZ, imZ) = impedance lns

fitness' targetImpedance lns = a * gn - b * ((targetImpedance - reZ)*(targetImpedance - reZ)) - c * (imZ * imZ) where
  a = 40
  b = 2
  c = 2
  gn = gain lns
  (reZ, imZ) = impedance lns

fitnessImepd' targetImpedance lns =  - b * ((targetImpedance - reZ)*(targetImpedance - reZ)) - c * (imZ * imZ) where
  b = 2
  c = 2
  gn = gain lns
  (reZ, imZ) = impedance lns

runSim sim sfreq efreq steps = go where
  splitRuns [] = []
  splitRuns xs = first : splitRuns rest where
    first = takeWhile (not . isSubsequenceOf "- FREQUENCY -") $ rest
    rest = tail $ (dropWhile (not . isSubsequenceOf "- FREQUENCY -")) xs

  go = do
    rn <- randomIO :: IO Int
    let fl = "/tmp/evalnec"++(show rn)
    writeFile (fl++".nec") $ printSim sfreq efreq steps False sim
    system $ "nec2c -i "++fl++".nec"
    !lns <- fmap lines $ readFile (fl++".out")
    system $ "rm "++fl++"*"
    return $ init $ splitRuns $ lns ++ ["- FREQUENCY -"]

evalSim sim freq fitnessF = go where
  --freq = 447.625
  go = do
    rn <- randomIO :: IO Int
    let fl = "/tmp/evalnec"++(show rn)
    writeFile (fl++".nec") $ printSim freq freq 1 False sim
    system $ "nec2c -i "++fl++".nec"
    !lns <- fmap lines $ readFile (fl++".out")
    system $ "rm "++fl++"*"
    --print $ gain lns
    --print $ impedance lns
    return $ fitnessF lns


{-# LANGUAGE OverloadedStrings, NoImplicitPrelude, BangPatterns #-}
module Main where
import ClassyPrelude
import qualified Numeric.LinearAlgebra as HMatrix
import NLP.Albemarle.Dict (counts, ids)
import qualified NLP.Albemarle.Dict as Dict
import NLP.Albemarle.LSA (termvectors, topicweights)
import qualified NLP.Albemarle.LSA as LSA
import qualified NLP.Albemarle.Tokens as Tokens
import qualified System.IO.Streams as Streams
import Data.Text.Encoding.Error (lenientDecode)
import Lens.Micro

main :: IO ()
main = do
  (final_dict, model) <- Streams.withFileAsInput "kjv.verses.txt.gz" $ \file ->
    Streams.gunzip file
    >>= Streams.lines
    >>= Streams.decodeUtf8With lenientDecode
    >>= Streams.map Tokens.wordTokenize
    >>= Streams.chunkList 5000 -- Yep, I'm making this stuff up.
    >>= Streams.map (\chunk -> let
      dict = Dict.dictifyAllWords chunk
      sparsem = Dict.asSparseMatrix dict chunk
      lsa = LSA.lsa 100 sparsem
      in (dict, lsa))
    >>= Streams.fold (\(!d1, !lsa1) (d2, lsa2) -> let
      d3 = Dict.filterDict 2 0.5 1000 $ d1 <> d2
      lsa3 = LSA.rebase d1 d3 lsa1 <> LSA.rebase d2 d3 lsa2
      in (d3, lsa3)) (mempty, mempty)

  -- It should use all 100 words allowed plus the unknown
  print $ length (final_dict^.counts)
  -- It should have a full size LSA as well
  print $ HMatrix.size (model^.termvectors)
  -- Plus topic weights
  print $ HMatrix.size (model^.topicweights)


{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}

module Spinell.Fitting
    (fit,
     unscaledFit,
     mvarFit,
     mvarUnscaledFit,
     defFitOpt,
     FitOptions(..),
     Jacobianable(..),
     Model,
     UnboxedModel,
     getParams,
     getUnscaledParams) where

import Numeric.GSL.Fitting
import Numeric.LinearAlgebra
import Numeric.AD
import Numeric.AD.Internal.Reverse
import Data.Reflection

type FitData = [([Double],([Double],Double))]
type FitRes = ([(Double, Double)], Matrix Double)
type UnscaledFitRes = ([Double], Matrix Double)

--
type Model = (forall a . Floating a => [a]->[a]->[a])

type UnboxedModel = (forall a . Floating a => [a]->a->a)

box :: UnboxedModel -> Model
box unmodel = (\p -> \x -> return (unmodel p (x !! 0)))

--
type FullJacob = ([Double] -> [Double] -> [[Double]])

data Jacobianable = AutoJacob | ManualJacob FullJacob

sanitize :: Jacobianable -> Model -> FullJacob
sanitize AutoJacob model = mkJac model
sanitize (ManualJacob jac) model = jac

--

data FitOptions = FitOptions {
        jacob :: Jacobianable,
        iter :: Int,
        absTol :: Double,
        relTol :: Double
    }

defFitOpt = FitOptions { jacob = AutoJacob,iter = 1000,absTol = 1E-20,relTol = 1E-20 }

--

fit :: [Double]
     -> [Double]
     -> [Double]
     -> UnboxedModel
     -> [Double]
     -> FitOptions
     -> FitRes

fit xs ys sigma rawmodel guess fitparams
    | xeqy && ( seqx || seqone) = fitModelScaled setAbsTol setResTol setIter (box rawmodel, sanitize jacobian (box rawmodel)) fitdata guess
    | otherwise = error "Input data has the wrong size"
    where
    lx = length xs
    ly = length ys
    ls = length sigma
    xeqy = lx == ly
    seqx = ls == lx
    seqone = ls == 1

    jacobian = jacob fitparams --we extract the "jacobian" parameter from the record

    setIter = iter fitparams
    setAbsTol = absTol fitparams
    setResTol = relTol fitparams

    fitdata = formatData xs ys sigma

formatData :: [Double] -> [Double] -> [Double] -> FitData
formatData xs ys sigma
    |  length sigma == length xs = zip (return <$> xs) $ zip (return <$> ys) sigma
    |  length sigma == 1 = zip (return <$> xs) $ zip (return <$> ys) (repeat (head sigma))

unscaledFit :: [Double]
     -> [Double]
     -> UnboxedModel
     -> [Double]
     -> FitOptions
     -> UnscaledFitRes

unscaledFit xs ys rawmodel guess fitparams
    | xeqy = fitModel setAbsTol setResTol setIter (box rawmodel, sanitize jacobian (box rawmodel)) fitdata guess
    | otherwise = error "Input data has the wrong size"
    where
    lx = length xs
    ly = length ys
    xeqy = lx == ly

    jacobian = jacob fitparams

    setIter = iter fitparams
    setAbsTol = absTol fitparams
    setResTol = relTol fitparams

    fitdata = zip (return <$> xs) (return <$> ys)

mvarFit :: [[Double]]
     -> [[Double]]
     -> [Double]
     -> Model
     -> [Double]
     -> FitOptions
     -> FitRes

mvarFit xs ys sigma model guess fitparams
    | xeqy && ( seqx || seqone) = fitModelScaled setAbsTol setResTol setIter (model, sanitize jacobian model) fitdata guess
    | otherwise = error "Input data has the wrong size"
    where
    lx = length xs
    ly = length ys
    ls = length sigma
    xeqy = lx == ly
    seqx = ls == lx
    seqone = ls == 1

    jacobian = jacob fitparams

    setIter = iter fitparams
    setAbsTol = absTol fitparams
    setResTol = relTol fitparams

    fitdata = formatDataBoxed xs ys sigma

formatDataBoxed :: [[Double]] -> [[Double]] -> [Double] -> FitData
formatDataBoxed xs ys sigma
    |  length sigma == length xs = zip xs $ zip ys sigma
    |  length sigma == 1 = zip xs $ zip ys (repeat (head sigma))

mvarUnscaledFit :: [[Double]]
     -> [[Double]]
     -> Model
     -> [Double]
     -> FitOptions
     -> UnscaledFitRes

mvarUnscaledFit xs ys model guess fitparams
    | xeqy = fitModel setAbsTol setResTol setIter (model, sanitize jacobian model) fitdata guess
    | otherwise = error "Input data has the wrong size"
    where
    lx = length xs
    ly = length ys
    xeqy = lx == ly

    jacobian = jacob fitparams

    setIter = iter fitparams
    setAbsTol = absTol fitparams
    setResTol = relTol fitparams

    fitdata = zip xs ys

mkJac :: (Num a) => (forall s. Reifies s Tape => [Reverse s a] -> [Reverse s a] -> [Reverse s a])
       -> [a]
       -> [a]
       -> [[a]]

mkJac f p x = jacobian ((flip f) (fmap auto x)) p

getParams :: FitRes -> [Double]
getParams res = fst . unzip . fst $ res

getUnscaledParams :: UnscaledFitRes -> [Double]
getUnscaledParams res = fst res


{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}

{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLabels      #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-|
Module      : Grenade.Layers.BatchNormalisation
Description : Batch normalization layer
Copyright   : (c) Theo Charalambous, 2020
License     : BSD2
Stability   : experimental

The layer follows the implementation as described in the paper:
Sergey Ioffe and Christian Szegedy. Batch normalization:  Accelerating deep network training by reducing internal covariate shift
<http://arxiv.org/abs/1502.03167>
-}

module Grenade.Layers.BatchNormalisation
  (
  -- * Layer Definition
    BatchNorm (..)
  , BatchNormParams(..)

  -- * Differentation Types
  , BatchNormTape(..)
  , BatchNormGrad(..)

  -- * Helper functions
  , initBatchNorm
  )
where

import           Control.DeepSeq
import           Data.Kind                         (Type)
import           Data.List                         (transpose, zipWith5)
import           Data.Maybe                        (fromJust)
import           Data.Proxy
import           Data.Serialize
import           GHC.TypeLits

import qualified Numeric.LinearAlgebra             as LA
import           Numeric.LinearAlgebra.Static      hiding (Seed)

import           Grenade.Core
import           Grenade.Layers.Internal.BatchNorm
import           Grenade.Layers.Internal.Update
import           Grenade.Onnx
import           Grenade.Types
import           Grenade.Utils.LinearAlgebra
import           Grenade.Utils.ListStore

import           Lens.Micro

data BatchNormTape :: Nat -- The number of channels of the tensor
                   -> Nat -- The number of rows of the tensor
                   -> Nat -- The number of columns of the tensor
                   -> Type where
  -- | there is no back propogation when using network for testing, so
  --   so the tape for automatic differentiation is empty
  TestBatchNormTape  :: ( KnownNat channels
                     , KnownNat rows
                     , KnownNat columns)
                     => BatchNormTape channels rows columns -- ^ empty tape

  -- | prevents having to calculate information twice during forward and backwards pass
  TrainBatchNormTape :: ( KnownNat channels
                     , KnownNat rows
                     , KnownNat columns)
                     => [R (channels * rows * columns)]  -- ^ xnorm
                     -> R channels                       -- ^ std
                     -> R channels                       -- ^ running mean
                     -> R channels                       -- ^ running variance
                     -> BatchNormTape channels rows columns

data BatchNorm :: Nat -- The number of channels of the tensor
               -> Nat -- The number of rows of the tensor
               -> Nat -- The number of columns of the tensor
               -> Nat -- momentum
               -> Type where
  BatchNorm :: ( KnownNat channels
               , KnownNat rows
               , KnownNat columns
               , KnownNat momentum)
            => Bool                                 -- ^ True: training, False: testing
            -> BatchNormParams channels             -- ^ gamma and beta
            -> R channels                           -- ^ running mean
            -> R channels                           -- ^ running variance
            -> RealNum                              -- ^ epsilon (used for numerical stability)
            -> ListStore (BatchNormParams channels) -- ^ momentum store
            -> BatchNorm channels rows columns momentum

data BatchNormParams :: Nat -> Type where
  BatchNormParams :: KnownNat channels
                  => R channels -- gamma
                  -> R channels -- beta
                  -> BatchNormParams channels

data BatchNormGrad :: Nat -- The number of channels of the tensor
                   -> Nat -- The number of rows of the tensor
                   -> Nat -- The number of columns of the tensor
                   -> Type where
  BatchNormGrad :: ( KnownNat channels
                   , KnownNat rows
                   , KnownNat columns)
            => R channels -- ^ running mean
            -> R channels -- ^ running variance
            -> R channels -- ^ derivative of loss wrt gamma
            -> R channels -- ^ derivative of loss wrt beta
            -> BatchNormGrad channels rows columns

-- | NFData instances
instance NFData (BatchNormTape channels rows columns) where
  rnf TestBatchNormTape = ()
  rnf (TrainBatchNormTape xnorm std mean var) = rnf xnorm `seq` rnf std `seq` rnf mean `seq` rnf var

instance NFData (BatchNormParams flattenSize) where
  rnf (BatchNormParams gamma beta)
    = rnf gamma `seq` rnf beta

instance NFData (BatchNorm channels rows columns momentum) where
  rnf (BatchNorm training bnparams mean var eps store)
    = rnf training `seq` rnf bnparams `seq` rnf mean `seq` rnf var `seq` rnf eps `seq` rnf store

-- | Show instance

instance Show (BatchNorm channels rows columns momentum) where
  show _ = "Batch Normalization"

-- Serialize instances
instance (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat momentum)
  => Serialize (BatchNorm channels rows columns momentum) where

  put (BatchNorm training bnparams mean var ε store) = do
    put training
    put bnparams
    putListOf put . LA.toList . extract $ mean
    putListOf put . LA.toList . extract $ var
    put ε
    put store

  get = do
    training <- get
    bnparams <- get
    mean     <- maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get
    var      <- maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get
    ε        <- get
    store    <- get

    return $ BatchNorm training bnparams mean var ε store

instance KnownNat flattenSize => Serialize (BatchNormParams flattenSize) where
  put (BatchNormParams gamma beta) = do
    putListOf put . LA.toList . extract $ gamma
    putListOf put . LA.toList . extract $ beta

  get = do
    gamma <- maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get
    beta  <- maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get

    return $ BatchNormParams gamma beta

-- | Neural network operations
instance (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat mom, KnownNat (channels * rows * columns))
  => UpdateLayer (BatchNorm channels rows columns mom) where

  type Gradient (BatchNorm channels rows columns mom) = BatchNormGrad channels rows columns
  type MomentumStore (BatchNorm channels rows columns mom) = ListStore (BatchNormParams channels)

  reduceGradient = undefined

  runUpdate opt@OptSGD{} x@(BatchNorm training (BatchNormParams oldGamma oldBeta) _ _ ε store) (BatchNormGrad runningMean runningVar dGamma dBeta)
    = let BatchNormParams oldGammaMomentum oldBetaMomentum = getData opt x store
          VectorResultSGD newGamma newGammaMomentum        = descendVector opt (VectorValuesSGD oldGamma dGamma oldGammaMomentum)
          VectorResultSGD newBeta  newBetaMomentum         = descendVector opt (VectorValuesSGD oldBeta dBeta oldBetaMomentum)
          newStore                                         = setData opt x store (BatchNormParams newGammaMomentum newBetaMomentum)
      in  BatchNorm training (BatchNormParams newGamma newBeta) runningMean runningVar ε newStore

  runUpdate opt@OptAdam{} x@(BatchNorm training (BatchNormParams oldGamma oldBeta) _ _ ε store) (BatchNormGrad runningMean runningVar dGamma dBeta)
    = let [BatchNormParams oldMGamma oldMBeta, BatchNormParams oldVGamma oldVBeta] = getData opt x store
          VectorResultAdam newGamma newMGamma newVGamma                            = descendVector opt (VectorValuesAdam (getStep store) oldGamma dGamma oldMGamma oldVGamma)
          VectorResultAdam newBeta  newMBeta  newVBeta                             = descendVector opt (VectorValuesAdam (getStep store) oldBeta  dBeta  oldMBeta  oldVBeta)
          newStore                                                                 = setData opt x store [BatchNormParams newMGamma newMBeta, BatchNormParams newVGamma newVBeta]
      in  BatchNorm training (BatchNormParams newGamma newBeta) runningMean runningVar ε newStore

  runSettingsUpdate NetworkSettings{trainingActive=training} (BatchNorm _ bnparams mean var ε store) = BatchNorm training bnparams mean var ε store

instance (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat mom)
  => LayerOptimizerData (BatchNorm channels rows columns mom) (Optimizer 'SGD) where

  type MomentumDataType (BatchNorm channels rows columns mom) (Optimizer 'SGD) = BatchNormParams channels
  getData opt x store = head $ getListStore opt x store
  setData opt x store = setListStore opt x store . return
  newData _ _ = BatchNormParams (konst 0) (konst 0)

instance (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat mom)
  => LayerOptimizerData (BatchNorm channels rows columns mom) (Optimizer 'Adam) where

  type MomentumDataType (BatchNorm channels rows columns mom) (Optimizer 'Adam) = BatchNormParams channels
  type MomentumExpOptResult (BatchNorm channels rows columns mom) (Optimizer 'Adam) = [BatchNormParams channels]
  getData     = getListStore
  setData     = setListStore
  newData _ _ = BatchNormParams (konst 0) (konst 0)

instance (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat momentum)
  => RandomLayer (BatchNorm channels rows columns momentum) where
  createRandomWith _ _ = pure initBatchNorm

-- | Initialize a batch norm layer with gamma set to zero for each channel and beta
--   set to one for each channel. Running mean is 0 and running var is 1. The default
--   value for epsilon is 0.00001
initBatchNorm :: forall channels rows columns momentum.
  (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat momentum)
  => BatchNorm channels rows columns momentum
initBatchNorm =
  let ch     = fromIntegral $ natVal (Proxy :: Proxy channels)
      zeroes = replicate ch 0
      ones   = replicate ch 1
      gamma  = vector ones   :: R channels
      beta   = vector zeroes :: R channels
      mean   = vector zeroes :: R channels
      var    = vector ones   :: R channels
      ε      = 0.00001
  in BatchNorm True (BatchNormParams gamma beta) mean var ε mkListStore

instance (KnownNat rows, KnownNat momentum)
  => Layer (BatchNorm 1 1 rows momentum) ('D1 rows) ('D1 rows) where

  type Tape (BatchNorm 1 1 rows momentum) ('D1 rows) ('D1 rows) = BatchNormTape 1 1 rows

  runForwards (BatchNorm True _ _ _ _ _) _
    = error "Cannot train use batch size of 1 with BatchNorm layer during training"

  runForwards (BatchNorm False (BatchNormParams gamma beta) runningMean runningVar ε _) (S1D x)
    = let gamma'       = extract gamma LA.! 0
          beta'        = extract beta LA.! 0
          runningMean' = extract runningMean LA.! 0
          runningVar'  = extract runningVar LA.! 0

          std = sqrt $ runningVar' + ε

          y = dvmap (\x -> ((x - runningMean') / std ) * gamma' + beta') x

      in (TestBatchNormTape, S1D y)

  runBatchForwards bn@(BatchNorm False _ _ _ _ _) xs
    = let outs = map (snd . runForwards bn) xs
      in ([TestBatchNormTape], outs)

  runBatchForwards (BatchNorm True (BatchNormParams gamma beta) runningMean runningVar ε _) xs
    = let [m]             = vectorToList runningMean :: [RealNum]
          [v]             = vectorToList runningVar  :: [RealNum]
          [g]             = vectorToList gamma       :: [RealNum]
          [b]             = vectorToList beta        :: [RealNum]
          mom             = (/ 100) $ fromIntegral $ natVal (Proxy :: Proxy momentum)

          xs'             = map extractV xs

          sample_mean     = batchNormMean xs'                 :: RealNum
          sample_var      = batchNormVariance xs'             :: RealNum

          m'              = mom * m + (1 - mom) * sample_mean :: RealNum
          v'              = mom * v + (1 - mom) * sample_var  :: RealNum
          std             = sqrt $ sample_var + ε             :: RealNum

          x_extracted     = map (\(S1D x) -> x) xs
          x_normalised    = map (dvmap (\a -> (a - sample_mean) / std)) x_extracted
          scaledShifted   = map (dvmap (\a -> g * a + b)) x_normalised
          out             = map S1D scaledShifted

          stdV            = listToVector [std] :: R 1
          runningMeanV    = listToVector [m']  :: R 1
          runningVarV     = listToVector [v']  :: R 1

      in ([TrainBatchNormTape x_normalised stdV runningMeanV runningVarV], out)

  runBatchBackwards (BatchNorm True (BatchNormParams _ _) _ _ _ _) _ _
    = undefined

  runBatchBackwards _ _ _
    = undefined

instance (KnownNat rows, KnownNat columns, KnownNat momentum)
  => Layer (BatchNorm 1 rows columns momentum) ('D2 rows columns) ('D2 rows columns) where

  type Tape (BatchNorm 1 rows columns momentum) ('D2 rows columns) ('D2 rows columns) = BatchNormTape 1 rows columns

  runForwards (BatchNorm True _ _ _ _ _) _
    = error "Cannot train use batch size of 1 with BatchNorm layer during training"

  runForwards (BatchNorm False (BatchNormParams gamma beta) runningMean runningVar ε _) (S2D x)
    = let rows     = fromIntegral $ natVal (Proxy :: Proxy rows)
          columns  = fromIntegral $ natVal (Proxy :: Proxy columns)

          gamma'       = extract gamma
          beta'        = extract beta
          runningMean' = extract runningMean
          runningVar'  = extract runningVar
          mat          = extract x

          y  = batchnorm 1 rows columns ε mat gamma' beta' runningMean' runningVar'
          y' = fromJust . create $ y

      in (TestBatchNormTape, S2D y')

  runBatchForwards bn@(BatchNorm False _ _ _ _ _) xs
    = let outs = map (snd . runForwards bn) xs
      in ([TestBatchNormTape], outs)

  runBatchForwards (BatchNorm True (BatchNormParams gamma beta) runningMean runningVar ε _) xs
    = let [m]             = vectorToList runningMean :: [RealNum]
          [v]             = vectorToList runningVar  :: [RealNum]
          [g]             = vectorToList gamma       :: [RealNum]
          [b]             = vectorToList beta        :: [RealNum]
          mom             = (/ 100) $ fromIntegral $ natVal (Proxy :: Proxy momentum)

          xs'             = map (sflatten . extractM2D) xs

          sample_mean     = batchNormMean xs'                 :: RealNum
          sample_var      = batchNormVariance xs'             :: RealNum

          m'              = mom * m + (1 - mom) * sample_mean :: RealNum
          v'              = mom * v + (1 - mom) * sample_var  :: RealNum
          std             = sqrt $ sample_var + ε             :: RealNum

          x_extracted     = map (\(S2D x) -> x) xs
          x_normalised    = map (dmmap (\a -> (a - sample_mean) / std)) x_extracted
          scaledShifted   = map (dmmap (\a -> g * a + b)) x_normalised
          out             = map S2D scaledShifted

          x_normalised'   = map sflatten x_normalised
          stdV            = listToVector [std] :: R 1
          runningMeanV    = listToVector [m']  :: R 1
          runningVarV     = listToVector [v']  :: R 1

      in ([TrainBatchNormTape x_normalised' stdV runningMeanV runningVarV], out)

  runBatchBackwards (BatchNorm True (BatchNormParams _ _) _ _ _ _) _ _
    = undefined

  runBatchBackwards _ _ _
    = undefined

instance (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat momentum)
  => Layer (BatchNorm channels rows columns momentum) ('D3 rows columns channels) ('D3 rows columns channels) where

  type Tape (BatchNorm channels rows columns momentum) ('D3 rows columns channels) ('D3 rows columns channels) = BatchNormTape channels rows columns

  runForwards (BatchNorm True _ _ _ _ _) _
    = error "Cannot train use batch size of 1 with BatchNorm layer during training"

  runForwards (BatchNorm False (BatchNormParams gamma beta) runningMean runningVar ε _) (S3D x)
    = let rows     = fromIntegral $ natVal (Proxy :: Proxy rows)
          columns  = fromIntegral $ natVal (Proxy :: Proxy columns)
          channels = fromIntegral $ natVal (Proxy :: Proxy channels)

          gamma'       = extract gamma
          beta'        = extract beta
          runningMean' = extract runningMean
          runningVar'  = extract runningVar
          mat          = extract x

          y  = batchnorm channels rows columns ε mat gamma' beta' runningMean' runningVar'
          y' = fromJust . create $ y

      in (TestBatchNormTape, S3D y')

  runBatchForwards bn@(BatchNorm False _ _ _ _ _) xs
    = let outs = map (snd . runForwards bn) xs
      in ([TestBatchNormTape], outs)

  runBatchForwards (BatchNorm True (BatchNormParams gamma beta) runningMean runningVar ε _) xs
    = let ms     = vectorToList runningMean
          vs     = vectorToList runningVar
          gs     = vectorToList gamma
          bs     = vectorToList beta

          cs     = map splitChannels xs :: [[S ('D2 rows columns)]]
          cs'    = transpose cs

          f c g b m v = let gs' = listToVector [g] :: R 1
                            bs' = listToVector [b] :: R 1
                            ms' = listToVector [m] :: R 1
                            vs' = listToVector [v] :: R 1
                            bn' = BatchNorm True (BatchNormParams gs' bs') ms' vs' ε undefined :: BatchNorm 1 rows columns momentum
                        in  runBatchForwards bn' c

          (tapes, outs) = unzip $ zipWith5 f cs' gs bs ms vs
          outs' = transpose outs
      in (combineTapes tapes, map combineChannels outs')
    where
      combineTapes :: [[BatchNormTape 1 i j]] -> [BatchNormTape k i j]
      combineTapes = undefined

  runBatchBackwards (BatchNorm True (BatchNormParams _ _) _ _ _ _) _ _
    = undefined

  runBatchBackwards _ _ _ = undefined

instance OnnxOperator (BatchNorm channels rows columns momentum) where
  onnxOpTypeNames _ = ["BatchNormalization"]

instance (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat momentum) => OnnxLoadable (BatchNorm channels rows columns momentum) where
  loadOnnxNode inits node = case node ^. #input of
    [_, scale, b, mean, var] -> do
      epsilon     <- readFloatAttributeToRealNum "epsilon" node
      loadedScale <- readInitializerVector inits scale
      loadedB     <- readInitializerVector inits b
      loadedMean  <- readInitializerVector inits mean
      loadedVar   <- readInitializerVector inits var

      return $ BatchNorm False (BatchNormParams loadedScale loadedB) loadedMean loadedVar epsilon mkListStore
    _               -> onnxIncorrectNumberOfInputs


{-# LANGUAGE TupleSections #-}
{-# LANGUAGE EmptyDataDeriving #-}
{-# LANGUAGE EmptyCase #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
module Models.Integrals.Types (module Models.Integrals.Types, module TLC.Terms) where

-- import Data.Ratio
import Algebra.Classes
import qualified Algebra.Morphism.Affine as A
import qualified Algebra.Morphism.LinComb as LC
import Prelude hiding (Num(..), Fractional(..), (^), product, sum, pi, sqrt
                      , exp, (**))
import Data.Complex
import TLC.Terms (type (∈)(..), Type(..), type(×), type(⟶))
import qualified Algebra.Expression as E
import Data.Function (on)
import Data.Foldable
import Data.List
--------------------------------------------------------------------------------
-- | Types

type C = Complex Double

data Zero deriving (Show,Eq,Ord)
newtype Number = Number (E.Expr Zero) deriving
  (Additive,Group,AbelianAdditive,Multiplicative,Division,Field,Ring,Roots,Transcendental)
instance Scalable Number Number where (*^) = (*)
type Rat = Number
fromNumber :: Number -> E.Expr Zero
fromNumber (Number x) = x
evalNumber :: Number -> Double
evalNumber (Number e) = E.eval (\case) e
instance Show Number where show = show . fromNumber
instance Eq Rat where
  (==) = (==) `on` evalNumber
instance Ord Rat where
  compare = compare `on` evalNumber
instance DecidableZero Rat where
  isZero  = (== 0) . evalNumber
type RatLike α = (Ring α, Ord α, DecidableZero α)

type Var γ = 'R ∈ γ
data Dir = Min | Max deriving (Eq,Ord,Show)
data Elem γ = Vari (Var γ)
            | Supremum Dir [Ret γ] -- either minimum or maximum of the arguments
  deriving (Eq, Ord, Show)
type Available α γ = α ∈ γ
type Expr γ = A.Affine (Var γ) Rat
type Ret γ = E.Expr (Elem γ)
type Cond γ = Cond' (Expr γ)
data Cond' e = IsNegative { condExpr :: e }
              -- Meaning of this constructor: expression ≤ 0
              | IsZero { condExpr :: e }
              -- Meaning of this constructor: expression = 0
   deriving (Eq, Show, Functor, Foldable, Traversable, Ord)

data Domain γ = Domain { domainLoBounds, domainHiBounds :: [Expr γ] }
  deriving (Show, Eq, Ord)

-- | needs to be 1st a order representation for optimisations to be
-- implementable
data P (γ :: Type) where
  Done :: Ret γ -> P γ
  Cond :: Cond γ -> P γ -> P γ
  Integrate :: Domain γ -> P (γ × 'R) -> P γ
  Add :: P γ -> P γ -> P γ
  Power :: P γ -> Rat -> P γ
  Mul :: [P γ] -> P γ
  -- Can this replaced by "Scale"? No, because we do integration in
  -- normalisation factors as well.
  -- Scale :: Ret γ -> P γ -> P γ
  deriving (Ord, Eq)

lift' :: Applicative f
  => (forall v. Available v γ -> f (Available v δ))
  -> (forall v. Available v (γ × α) -> f (Available v (δ × α)))
lift' _ (Get) = pure Get
lift' f (Weaken x) = Weaken <$> (f x)

class VarTraversable t where
  varTraverse :: (Applicative f)
    => (forall x. Available x γ -> f (Available x δ))
    -> t γ -> f (t δ)

instance VarTraversable Domain where
 varTraverse f (Domain los his)
  = Domain <$> traverse (A.traverseVars f) los <*>
               traverse (A.traverseVars f) his

instance VarTraversable P where
  varTraverse f = \case
    Power e k -> Power <$> varTraverse f e <*> pure k
    Done x -> Done <$> traverse (varTraverse f) x
    Mul xs -> Mul <$> traverse (varTraverse f) xs
    Integrate d e ->
      Integrate <$> (varTraverse f d) <*> (varTraverse (lift' f) e)
    Cond e x -> Cond <$> traverse (A.traverseVars f) e <*> varTraverse f x
    Add x y  -> Add <$> varTraverse f x <*> varTraverse f y

instance VarTraversable Elem where
  varTraverse f = \case
    Vari x -> Vari <$> f x
    Supremum d es -> Supremum d <$> traverse (traverse (varTraverse f)) es

deriving instance Show (P γ)

----------------------------------
-- | Smart constructors

conds_ :: [Cond γ] -> P γ -> P γ
conds_ cs e = foldr Cond e cs

isPositive,isNegative :: Expr γ -> Cond γ
isPositive e = isNegative (negate e)
isNegative e = IsNegative e

-- Domain without restriction
full :: Domain γ
full = Domain [] []

lessThan, greaterThan :: Expr γ -> Expr γ -> Cond γ
t `lessThan` u = isNegative (t - u)
t `greaterThan` u = u `lessThan` t

----------------------------
-- Instances

instance Multiplicative (P γ) where
  one = Done one
  x * y = Mul [x,y]

instance Division (P γ) where
  recip x = Power x (negate one)

instance AbelianAdditive (P γ)
instance Group (P γ) where
  negate = (negate (one :: Ret γ) *^)
instance Scalable (E.Expr (Elem γ)) (P γ) where
  k *^ e = Done k * e

pattern PZero :: P γ
pattern PZero <- Done E.Zero

instance Additive (P γ) where
  zero =  Done E.Zero
  PZero + x = x
  x + PZero = x
  x + y = Add x y

setHere :: f γ -> (Var γ -> f γ) -> Var (γ × α) -> f γ
setHere a f = \case
  Get -> a
  Weaken x -> f x

----------------------------------------------------------------
-- Normalising substitutions of variables to affine expressions

type SubstE γ δ = Var γ -> Expr δ

wkSubst :: SubstE γ δ -> SubstE (γ × α) (δ × α)
wkSubst f = \case
  Get -> A.var Get
  Weaken x -> A.mapVars Weaken (f x)

substExpr :: SubstE γ δ ->  Expr γ -> Expr δ
substExpr = A.subst

substCond :: SubstE γ δ -> Cond γ -> Cond δ
substCond f = fmap (substExpr f)

substDomain :: SubstE γ δ -> Domain γ -> Domain δ
substDomain f (Domain lo hi) = Domain
                                 (nub (substExpr f <$> lo))
                                 (nub (substExpr f <$> hi))

wkP :: P γ -> P (γ × α)
wkP = substP $ \i -> A.var (Weaken i)

substElem :: forall γ ζ. SubstE γ ζ -> Elem γ -> Ret ζ
substElem v = \case
  Supremum dir es -> supremum dir (substRet v <$> es)
  Vari x -> exprToPoly (v x)

substRet  :: forall γ ζ. SubstE γ ζ -> Ret γ -> Ret ζ
substRet v = E.eval (substElem v)

substP :: SubstE γ δ -> P γ -> P δ
substP f p0 = case p0 of
  Done e -> Done (E.eval (substElem f) e)
  Add p1 p2 -> substP f p1 + substP f p2
  Power p k -> Power (substP f p) k
  Mul ps -> Mul (substP f <$> ps)
  Cond c p -> Cond (substCond f c) (substP f p)
  Integrate d p -> Integrate (substDomain f d) (substP (wkSubst f) p) -- integrations are never simplified by substitution

swap2P :: P (γ × α × β) -> P (γ × β × α)
swap2P = substP $ \case
  Get -> A.var (Weaken Get)
  Weaken Get -> A.var Get
  Weaken (Weaken x) -> A.var (Weaken (Weaken x))

wkExpr :: Expr γ -> Expr (γ × β)
wkExpr = substExpr (A.var . Weaken)

condVars :: Cond γ -> [Var γ]
condVars c = case condExpr c of
   (A.Affine _ e) -> map fst (LC.toList e)

retVars :: Ret γ -> [Var γ]
retVars x = concatMap elemVars (toList x)

elemVars :: Elem γ -> [Var γ]
elemVars = \case
   Vari x -> [x]
   Supremum _ es -> concatMap retVars es

supremum :: Dir -> [Ret γ] -> Ret γ
supremum _ [e] = e
supremum dir es =
  case traverse (fmap retToNumber . traverse (varTraverse (const Nothing))) es of
    Just cs | not (null es) ->
       constPoly ((case dir of
                     Max -> maximum
                     Min -> minimum)
                   cs)
    _ -> E.Var (Supremum dir es)

constPoly :: Number -> Ret γ
constPoly (Number n) = (\case) <$> n

varPoly :: 'R ∈ γ -> Ret γ
varPoly = E.Var . Vari

retToNumber :: Ret 'Unit -> Number
retToNumber = E.eval $ \case
     Vari x -> case x of
     Supremum d xs -> (case d of Max -> maximum; Min -> minimum) (fmap retToNumber xs)

numberToRet :: Number -> Ret γ
numberToRet (Number x) = fmap (\case) x
exprToPoly :: Expr γ -> Ret γ
exprToPoly = A.eval (fmap (\case) .fromNumber) (E.Var .  Vari)

mkSuprema :: Domain γ -> (Ret γ, Ret γ)
mkSuprema (Domain los his) = (supremum Max $ map exprToPoly los,
                              supremum Min $ map exprToPoly his)


{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}

module Year2020.Day20 where

import Util hiding ((??))

import Numeric.LinearAlgebra

import Data.Char
import Data.Function
import Data.List hiding (find)
import Data.List.Extra hiding (find)

import Control.Arrow
import Control.Monad

import System.FilePath

import Data.Map (Map)
import qualified Data.Map as Map

type Tile = Matrix Z

makeTile :: [String] -> (Int, Tile)
makeTile (x:xs) = (n, fromLists (map (map (\case '#' -> 1 ; _ -> 0)) xs))
  where n = read (filter isDigit x) :: Int

transforms :: Tile -> [Tile]
transforms m = take 8 . scanl (&) m $ cycle [tr, flipud]

place :: Int -> Map Int Tile -> Map (Int, Int) (Int, Tile)
       -> Int -> Int -> [[[(Int, Tile)]]]
place size tiles placed _ _ | Map.null tiles = pure . transpose
  .  map (map snd) . groupOn (fst . fst) $ Map.toList placed
place size tiles placed x y = do
  let xp = snd (placed Map.! (x - 1, y)) ?? (All, TakeLast 1)
      yp = snd (placed Map.! (x, y - 1)) ?? (TakeLast 1, All)
  (num, tile') <- Map.toList tiles
  tile <- transforms tile'
  guard $ (x == 1 || tile ?? (All, Take 1) == xp)
       && (y == 1 || tile ?? (Take 1, All) == yp)
  let (x', y') = ((+ x) *** (+ 1)) (divMod y size)
  place size (Map.delete num tiles) (Map.insert (x, y) (num, tile) placed) x' y'

part1, part2 :: [[(Int, Tile)]] -> Int
part1 xs = product . map (fst . ($ xs)) $ (.) <$> [head, last] <*> [head, last]
part2 = fromIntegral . minimum . checks . transforms . fromBlocks
  . map (map ((?? (DropLast 1, DropLast 1)) . (?? (Drop 1, Drop 1)) . snd))
  where goal = fromIntegral $ sumElements monster
        check xs = sumElements $ cond (corr2 monster xs) goal 0 goal 0
        checks (x:xs) = map ((sumElements x -) . check) (x:xs)
        monster = snd . makeTile $ "" :
          [ "                  # "
          , "#    ##    ##    ###"
          , " #  #  #  #  #  #   " ]

main = do
  input <- readFile (replaceExtension __FILE__ ".in")
  let tiles = map makeTile . filter notNull $ paragraphs input
      size = round . sqrt . fromIntegral $ length tiles
  print . part2 . head $ place size (Map.fromList tiles) mempty 1 1


{-# LANGUAGE DeriveFunctor                            #-}
{-# LANGUAGE DeriveGeneric                            #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE KindSignatures                           #-}
{-# LANGUAGE LambdaCase                               #-}
{-# LANGUAGE NoStarIsType                             #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE RecordWildCards                          #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE StandaloneDeriving                       #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeFamilyDependencies                   #-}
{-# LANGUAGE TypeInType                               #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE UndecidableInstances                     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module Numeric.Wavelet.Discrete (
    DWavelet(..)
  , DWT(..)
  , cdwt
  , dwt
  , idwt
  , haar
    -- DWD(..)
  -- , dwdApprox, dwdDetail
  -- , flattenDWD
  -- , denseDWD
  -- , haar
  -- , unHaar
  ) where

import           Data.Complex
import           Data.Finite
import           Data.Finite.Internal
import           Data.Kind
import           Data.Proxy
import           Data.Vector.Generic.Sized    (Vector)
import           GHC.Generics
import           GHC.TypeLits.Compare
import           GHC.TypeNats
import           Numeric.Wavelet.Internal.FFT
import qualified Data.Vector                  as UV
import qualified Data.Vector.Generic          as UVG
import qualified Data.Vector.Generic.Sized    as VG
import qualified Data.Vector.Sized            as V

data DWavelet v l a = DW
    { dwHiDecomp :: VG.Vector v l a      -- ^ High-pass decomposition filter
    , dwLoDecomp :: VG.Vector v l a      -- ^ Low-pass decomposition filter
    , dwHiRecon  :: VG.Vector v l a      -- ^ High-pass reconstruction filter
    , dwLoRecon  :: VG.Vector v l a      -- ^ Low-pass reconstruction filter
    }
  deriving (Show, Eq, Ord, Functor, Generic)

data DWT v n a = DWT
    { dwtApprox :: VG.Vector v n a
    , dwtDetail :: VG.Vector v n a
    }
  deriving (Show, Eq, Ord, Functor, Generic)

cdwt
    :: (UVG.Vector v (Complex a), KnownNat l, KnownNat n, 1 <= l, 1 <= n, FFTWReal a)
    => DWavelet  v l (Complex a)
    -> VG.Vector v n (Complex a)
    -> DWT v ((n + l - 1) `Div` 2) (Complex a)
cdwt DW{..} x = DWT{..}
  where
    dwtApprox = downsamp $ convolve x dwLoDecomp
    dwtDetail = downsamp $ convolve x dwHiDecomp

dwt
    :: (UVG.Vector v (Complex a), UVG.Vector v a, KnownNat l, KnownNat n, 1 <= l, 1 <= n, FFTWReal a)
    => DWavelet  v l a
    -> VG.Vector v n a
    -> DWT v ((n + l - 1) `Div` 2) a
dwt DW{..} x = DWT{..}
  where
    dwtApprox = downsamp $ rconvolve x dwLoDecomp
    dwtDetail = downsamp $ rconvolve x dwHiDecomp

idwt
    :: forall v l n a.
     ( UVG.Vector v (Complex a)
     , UVG.Vector v a
     , KnownNat l
     , KnownNat n
     , FFTWReal a
     , Div ((((Div ((n + l) - 1) 2 * 2) + l) - 1) - n) 2 <= ((((Div ((n + l) - 1) 2 * 2) + l) - 1) - n)
     )
    => DWavelet  v l a
    -> DWT v ((n + l - 1) `Div` 2) a
    -> VG.Vector v n a      -- todo: remove the choice
idwt DW{..} DWT{..} = VG.slice @_ @((((n+l-1)`Div`2)*2+l-1-n)`Div`2) @n @(((n+l-1)`Div`2)*2+l-1-n-((((n+l-1)`Div`2)*2+l-1-n)`Div`2)) Proxy $
                VG.zipWith (+) x y
  where
    x = rconvolve (upsamp dwtApprox) dwLoRecon
    y = rconvolve (upsamp dwtDetail) dwHiRecon

downsamp :: (UVG.Vector v a, KnownNat n) => VG.Vector v n a -> VG.Vector v (n `Div` 2) a
downsamp v = VG.generate $ \i -> v `VG.index` doubleFinite i
  where
    doubleFinite :: Finite (n `Div` 2) -> Finite n
    doubleFinite (Finite x) = Finite (x * 2 + 1)

upsamp :: (UVG.Vector v a, KnownNat n, Num a) => VG.Vector v n a -> VG.Vector v (n * 2) a
upsamp v = VG.generate $ \i -> case separateProduct @2 i of
    (j, k)
      | j == 0    -> v `VG.index` k
      | otherwise -> 0

-- upsamp :: forall v n a. (UVG.Vector v a, KnownNat n, Num a) => VG.Vector v (n `Div` 2) a -> VG.Vector v n a
-- upsamp v = VG.generate $ \(Finite i) -> case i `divMod` 2 of
--     (d, m)
--       | m == 0    -> v `VG.index` Finite d
--       | otherwise -> 0

    -- let (j, k) = case separateProduct @2 i
    -- in  if j == 0
    --       then

-- unhaarPass
--     :: (UVG.Vector v a, KnownNat n, Num a)
--     => Vector v (2 ^ n) a
--     -> Vector v (2 ^ n) a
--     -> Vector v (2 ^ (n + 1)) a
-- unhaarPass app det = VG.generate $ \i ->
--     let (j, k) = separateProduct @2 i
--         combiner
--           | j == 0    = (+)
--           | otherwise = (-)
--     in  combiner (app `VG.index` k) (det `VG.index` k)

haar :: (UVG.Vector v a, Floating a) => DWavelet v 2 a
haar = DW{..}
  where
    srtot = sqrt 2 / 2
    dwHiDecomp = VG.fromTuple (-srtot, srtot)
    dwLoDecomp = VG.fromTuple ( srtot, srtot)
    dwHiRecon  = VG.fromTuple ( srtot,-srtot)
    dwLoRecon  = VG.fromTuple ( srtot, srtot)

-- data DWD :: (Type -> Type) -> Nat -> Type -> Type where
--     DWDZ :: { _dwdDetailZ :: !a
--            , _dwdApprox  :: !a
--            } -> DWD v 0 a
--     DWDS :: { _dwdDetailS :: !(Vector  v (2 ^ (n + 1)) a)
--            , _dwdNext    :: !(DWD v n             a)
--            } -> DWD v (n + 1) a

-- deriving instance (Show (v a), Show a) => Show (DWD v n a)

-- dwdDetail :: UVG.Vector v a => DWD v n a -> Vector v (2 ^ n) a
-- dwdDetail = \case
--     DWDZ d _ -> VG.singleton d
--     DWDS d _ -> d

-- dwdApprox :: DWD v n a -> a
-- dwdApprox = \case
--     DWDZ _ a -> a
--     DWDS _ a -> dwdApprox a

-- testDec :: DWD UV.Vector 4 Double
-- testDec = DWDS (0 :: V.Vector 16 Double)
--         . DWDS (0 :: V.Vector  8 Double)
--         . DWDS (0 :: V.Vector  4 Double)
--         . DWDS (0 :: V.Vector  2 Double)
--         $ DWDZ 0 0

-- flattenDWD :: UVG.Vector v a => DWD v n a -> Vector v (2 ^ (n + 1)) a
-- flattenDWD = \case
--     DWDZ d  a  -> VG.fromTuple (a, d)
--     -- whoops this is O(n^2)
--     DWDS ds as -> flattenDWD as VG.++ ds

-- denseDWD
--     :: forall v n a. (KnownNat n, UVG.Vector v a)
--     => DWD v n a
--     -> V.Vector (n + 1) (Vector v (2 ^ n) a)
-- denseDWD = \case
--     DWDZ d  _  -> VG.singleton $ VG.singleton d
--     -- whoops this is O(n^3) or something
--     DWDS ds as -> (densify @_ @_ @(n - 1) <$> denseDWD as) `VG.snoc` ds

-- densify
--     :: (UVG.Vector v a, KnownNat n)
--     => Vector v (2 ^ n) a
--     -> Vector v (2 ^ (n + 1)) a
-- densify xs = VG.generate $ \i -> xs `VG.index` snd (separateProduct @2 i)

-- haar
--     :: forall v n a. (UVG.Vector v a, KnownNat n, Fractional a)
--     => Vector  v (2 ^ (n + 1)) a
--     -> DWD v n a
-- haar xs = case pNat (Proxy @n) of
--     PZ -> let x = xs `VG.index` 0
--               y = xs `VG.index` 1
--           in  DWDZ ((x - y) / 2) ((x + y) / 2)
--     PS -> let (a, d) = haarPass xs
--           in  DWDS d (haar a)

-- unHaar
--     :: (UVG.Vector v a, KnownNat n, Num a)
--     => DWD v n a
--     -> Vector v (2 ^ (n + 1)) a
-- unHaar = \case
--     DWDZ d  a  -> VG.fromTuple (a + d, a - d)
--     DWDS ds as -> unhaarPass (unHaar as) ds

-- haarPass
--     :: (UVG.Vector v a, KnownNat n, Fractional a)
--     => Vector v (2 ^ (n + 1)) a
--     -> (Vector v (2 ^ n) a, Vector v (2 ^ n) a)
-- haarPass xs = (app, det)
--   where
--     app = VG.generate $ \i -> ( xs `VG.index` combineProduct (0, i)
--                               + xs `VG.index` combineProduct (1, i)
--                               ) / 2
--     det = VG.generate $ \i -> ( xs `VG.index` combineProduct (0, i)
--                               - xs `VG.index` combineProduct (1, i)
--                               ) / 2

-- unhaarPass
--     :: (UVG.Vector v a, KnownNat n, Num a)
--     => Vector v (2 ^ n) a
--     -> Vector v (2 ^ n) a
--     -> Vector v (2 ^ (n + 1)) a
-- unhaarPass app det = VG.generate $ \i ->
--     let (j, k) = separateProduct @2 i
--         combiner
--           | j == 0    = (+)
--           | otherwise = (-)
--     in  combiner (app `VG.index` k) (det `VG.index` k)


module Main where

import Parser
import Analyzer
import qualified Data.MultiSet as M
import Control.Lens hiding ((??),(|>),(<.>))
import qualified Data.Text.Lazy.IO as I
import Data.List (sortBy)
import Data.Function (on)
import Text.Megaparsec (parseMaybe)
import Numeric.LinearAlgebra
import Data.Complex (realPart)

file :: String
file = "/home/bertram/Downloads/la famiglia.txt"

messages :: IO [Chunk]
messages = filter isMessage <$> ((parseMaybe chunksP <$> I.readFile file)
           >>= maybe (putStrLn "Parse failed" >> return []) return)

main' :: IO ()
main' = do
  msg <- messages
  let users = map (toText . (^.user)) msg
  let counts = M.fromList users
  mapM_ (\(a,b)->I.putStr a >> putStr " | " >> print b) . sortBy (compare `on` snd) . M.toOccurList $ counts

main :: IO ()
main = do
  input <- I.readFile file
  let (matrix,users) = maybe undefined id (parseMaybe matrixP input)
  let n = length users
  let probs =  cmap realPart . flatten . (?? (All, Take 1))
               . snd . eig . tr $ matrix
  let probs' = scale (1 / (probs <.> (n |> repeat 1))) probs
  mapM_ print . sortBy (compare `on` snd). zip users $ toList probs'
  print $ sum . toList $ probs'


module D20
    ( largestNumberOfDoors, countRoomsAfterDoors
    ) where

import Text.Parsec
import Data.Functor.Identity
import Data.Either
import Data.Complex
import qualified Data.HashMap.Strict as M

type Position = Complex Float
data Navigation = Navigation  { currentPosition :: Position
                              , pastBranches :: [Position]
                              , distances :: M.HashMap Position Int } deriving (Show,Eq)

changePosition :: Position -> Navigation -> Navigation
changePosition p (Navigation cp pb d) = Navigation np pb (M.insert np (shortest (M.lookup np d) (M.lookup cp d)) d)
  where np=cp+p
        shortest Nothing (Just n) = n+1
        shortest (Just n) Nothing = n
        shortest (Just m) (Just n) = min m (n+1)

enterBranch :: Navigation -> Navigation
enterBranch (Navigation cp pb d) = Navigation cp (cp:pb) d

tryAlternateBranch :: Navigation -> Navigation
tryAlternateBranch (Navigation _ pb@(p:_) d) = Navigation p pb d

leaveBranch :: Navigation -> Navigation
leaveBranch (Navigation _ (p:pb) d) = Navigation p pb d

type RoomParser = ParsecT String Navigation Identity

tryAllPaths :: ParsecT String Navigation Identity Navigation
tryAllPaths = char '^'
              >> (many $ choice [ char 'N' >> modifyState (changePosition $ 0:+(-1))
                                , char 'S' >> modifyState (changePosition $ 0:+1)
                                , char 'E' >> modifyState (changePosition $ 1:+0)
                                , char 'W' >> modifyState (changePosition $ (-1):+0)
                                , char '(' >> modifyState enterBranch
                                , char '|' >> modifyState tryAlternateBranch
                                , char ')' >> modifyState leaveBranch ])
              >> char '$'
              >> getState

getAllDistances :: String -> [Int]
getAllDistances = M.elems . distances . fromRight nav . runParser tryAllPaths nav ""
  where nav = Navigation (0:+0) [] (M.singleton (0:+0) 0)

largestNumberOfDoors :: String -> Int
largestNumberOfDoors = maximum . getAllDistances

countRoomsAfterDoors :: Int -> String -> Int
countRoomsAfterDoors threshold = length . filter (((<=)threshold)) . getAllDistances


{-|
Module      : Lib
Description : Main component of your project.
Copyright   : (c) Fabricio Olivetti de Franca, 2021
License     : GPL-3
Maintainer  : fabricio.olivetti@gmail.com
Stability   : experimental
Portability : POSIX
|-}

module GAPoly
    ( parseFile
    , runGA
    , generateReports
    , Solution(..)
    , Poly(..)
    ) where

import GA
import Fitness
import Dataset
import Report
import Random
import Evolution

import Control.Monad.State.Strict
import System.Random
import Numeric.LinearAlgebra (Vector)
import GHC.TypeLits

gaPoly :: KnownNat n => Double -> Evolution (Poly n)
gaPoly pm = Select generational End
              (Cross crossover
                (Mutate (mutate pm) End))

-- | Run a genetic algorithm with the provided parameters
runGA :: KnownNat n => Int -> Int -> Int -> Int -> Double -> [[Double]] -> Vector Double -> IO ([Double], Solution (Poly n))
runGA nTerms nVars it nPop pm xss ys =
  do g <- newStdGen
     let fitness   = evalFitness nTerms xss ys
         createSol = createRndSolution (nVars*nTerms)
         gens      = runEvolution createSol fitness it nPop (gaPoly pm)
     evalStateT gens g


module ComplexUtils (toComplex, toComplexM) where

import Data.Complex (Complex((:+)))

toComplex :: Double -> (Complex Double)
toComplex a = (a :+ 0)

toComplexM :: [Double] -> [(Complex Double)]
toComplexM = map toComplex


module Math.FROG.Types
( Nonlinearity(..)
, ComplexSignal
, Trace
) where

import qualified Numeric.LinearAlgebra as LA

type ComplexSignal = LA.Vector (LA.Complex Double)
type Trace = LA.Matrix Double

data Nonlinearity = PG | SHG | SD | THG


{-# language BangPatterns        #-}
{-# language LambdaCase          #-}
{-# language NoImplicitPrelude   #-}
{-# language ScopedTypeVariables #-}

-- | This module exposes functions for performing
--   Fast Fourier Transform (FFT) and Inverse Fast Fourier Transform (IFFT)
--   over 'Contiguous' data structures.
module Data.Primitive.Contiguous.FFT
  ( fft
  , ifft
  , mfft
  ) where

import qualified Prelude

import Control.Applicative (pure)
import Control.Monad (when)
import Control.Monad.Primitive (PrimMonad(..))
import Control.Monad.ST (runST)
import Data.Bits (shiftR,shiftL,(.&.),(.|.))
import Data.Bool (Bool,otherwise)
import Data.Complex (Complex(..),conjugate)
import Data.Eq (Eq(..))
import Data.Function (($),(.))
import Data.Ord (Ord(..))
import Data.Primitive.Contiguous (Contiguous,Element,Mutable)
import Data.Semiring (negate,(+),(*),(-))
import GHC.Exts
import GHC.Real ((/))
import qualified Data.Primitive.Contiguous as Contiguous

{-# RULES
"fft/ifft" forall x. fft (ifft x) = x
"ifft/fft" forall x. ifft (fft x) = x
  #-}

-- | Radix-2 decimation-in-time fast Fourier Transform.
--   The given array must have a length that is a power of two.
fft :: forall arr. (Contiguous arr, Element arr (Complex Double))
  => arr (Complex Double)
  -> arr (Complex Double)
{-# inlinable [1] fft #-}
fft arr = if arrOK arr
  then runST $ do {
      marr <- copyWhole arr
    ; mfft marr
    ; Contiguous.unsafeFreeze marr
  }
  else Prelude.error "Data.Primitive.Contiguous.FFT.fft: bad array length"
-- | Inverse fast Fourier transform.
ifft :: forall arr. (Contiguous arr, Element arr (Complex Double))
  => arr (Complex Double)
  -> arr (Complex Double)
{-# inlinable [1] ifft #-}
ifft arr = if arrOK arr
  then
    let lenComplex = intToComplexDouble (Contiguous.size arr)
    in cmap ((/lenComplex) . conjugate) . fft . cmap conjugate $ arr
  else Prelude.error "Data.Primitive.Contiguous.FFT.ifft: bad vector length"

copyWhole :: forall arr m a. (PrimMonad m, Contiguous arr, Element arr a)
  => arr a
  -> m (Mutable arr (PrimState m) a)
{-# inline copyWhole #-}
copyWhole arr = do
  let len = Contiguous.size arr
  marr <- Contiguous.new len
  Contiguous.copy marr 0 arr 0 len
  pure marr

arrOK :: forall arr a. (Contiguous arr, Element arr a)
  => arr a
  -> Bool
{-# inline arrOK #-}
arrOK arr =
  let n = Contiguous.size arr
  in (1 `shiftL` log2 n) == n

-- | Radix-2 decimation-in-time fast Fourier Transform.
--   The given array must have a length that is a power of two,
--   though this property is not checked.
mfft :: forall arr m. (PrimMonad m, Contiguous arr, Element arr (Complex Double))
  => Mutable arr (PrimState m) (Complex Double)
  -> m ()
mfft mut = do {
    len <- Contiguous.sizeMutable mut
  ; let bitReverse !i !j = do {
          ; if i == len - 1
              then stage 0 1
              else do {
                  when (i < j) $ swap mut i j
                ; let inner k l = if k <= l
                        then inner (k `shiftR` 1) (l - k)
                        else bitReverse (i + 1) (l + k)
                ; inner (len `shiftR` 1) j
              }
        }
        stage l l1 = if l == (log2 len)
          then pure ()
          else do {
              let !l2 = l1 `shiftL` 1
                  !e = (negate twoPi) / (intToDouble l2)
                  flight j !a = if j == l1
                    then stage (l + 1) l2
                    else do {
                        let butterfly i = if i >= len
                              then flight (j + 1) (a + e)
                              else do {
                                  let i1 = i + l1
                                ; xi1 :+ yi1 <- Contiguous.read mut i1
                                ; let !c = Prelude.cos a
                                      !s' = Prelude.sin a
                                      d = (c*xi1 - s'*yi1) :+ (s'*xi1 + c*yi1)
                                ; ci <- Contiguous.read mut i
                                ; Contiguous.write mut i1 (ci - d)
                                ; Contiguous.write mut i (ci + d)
                                ; butterfly (i + l2)
                              }
                      ; butterfly j
                    }
            ; flight 0 0
         }
  ; bitReverse 0 0
}

-- wildcard cases should never happen. if they do, really bad things will happen.
b,s :: Int -> Int
b = \case { 0 -> 0x02; 1 -> 0x0c; 2 -> 0xf0; 3 -> 0xff00; 4 -> wordToInt 0xffff0000; 5 -> wordToInt 0xffffffff00000000; _ -> 0; }
s = \case { 0 -> 1; 1 -> 2; 2 -> 4; 3 -> 8; 4 -> 16; 5 -> 32; _ -> 0; }
{-# inline b #-}
{-# inline s #-}

log2 :: Int -> Int
log2 v0 = if v0 <= 0
  then Prelude.error $ "Data.Primitive.Contiguous.FFT: nonpositive input, got " Prelude.++ Prelude.show v0
  else go 5 0 v0
  where
    go !i !r !v
      | i == -1 = r
      | v .&. b i /= 0 =
          let si = s i
          in go (i - 1) (r .|. si) (v `shiftR` si)
      | otherwise = go (i - 1) r v

swap :: forall arr m x. (PrimMonad m, Contiguous arr, Element arr x)
  => Mutable arr (PrimState m) x
  -> Int
  -> Int
  -> m ()
{-# inline swap #-}
swap mut i j = do
  atI <- Contiguous.read mut i
  atJ <- Contiguous.read mut j
  Contiguous.write mut i atJ
  Contiguous.write mut j atI

twoPi :: Double
{-# inline twoPi #-}
twoPi = 6.283185307179586

intToDouble :: Int -> Double
{-# inline intToDouble #-}
intToDouble = Prelude.fromIntegral

wordToInt :: Word -> Int
{-# inline wordToInt #-}
wordToInt = Prelude.fromIntegral

intToComplexDouble :: Int -> Complex Double
{-# inline intToComplexDouble #-}
intToComplexDouble = Prelude.fromIntegral

cmap :: (Contiguous arr, Element arr (Complex Double))
  => (Complex Double -> Complex Double)
  -> arr (Complex Double)
  -> arr (Complex Double)
{-# inline cmap #-}
cmap = Contiguous.map


------------------------------------------------
-- |
-- Module    :  Numeric.Matrix.Integral.SmithNF
-- Copyright :  (c) Jun Yoshida 2019
-- License   :  BSD3
--
-- Compute Smith normal form
--
------------------------------------------------

module Numeric.Matrix.Integral.SmithNF where

import Control.Monad (when, unless, forM_, guard, join)
import Control.Monad.ST
import Control.Monad.Loops (whileM_, whileJust_)

import Data.List as L
import Data.STRef

import Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra.Devel

--import Numeric.Matrix.Integral.HNFLLL (hnfLLLST)
import Numeric.Matrix.Integral.NormalForms (hermiteNFST)

-- | Find off-diagonal entries and calculate min of min{row,col} among them
findDiagSz :: LA.Matrix LA.Z -> Maybe Int
findDiagSz ma =
  let offdiag = filter (uncurry (/=)) $ LA.find (/=0) ma
  in L.foldl' bin Nothing offdiag
  where
    bin Nothing (i,j) = Just (min i j)
    bin (Just k) (i,j) = Just $ L.foldl' min k [i,j]

-- | Compute "pre-Smith" form with a function calculating Hermite normal forms
-- "Pre-Smith" means diag(c_1,c_2,..) with not necessarily c_i|c_{i+1}
preSmithNFST :: STMatrix s LA.Z -> STMatrix s LA.Z ->  STMatrix s LA.Z -> ST s ()
preSmithNFST stMatUL stMatA stMatURt = do
  whileJust_ (liftSTMatrix findDiagSz stMatA) $ \diagSz -> do
    -- Extract a submatrix containing all the off-diagonal entries.
    stMatAex <- thawMatrix =<< extractMatrix stMatA (FromRow diagSz) (FromCol diagSz)
    stMatULex <- thawMatrix =<< extractMatrix stMatUL (FromRow diagSz) AllCols
    stMatURtex <- thawMatrix =<< extractMatrix stMatURt (FromRow diagSz) AllCols
    -- Compute the Hermite normal forms of the extracted submatrix and of the transpose of the result.
    --hnfLLLST stMatULex stMatAex
    join $ hermiteNFST <$> newMatrix 0 0 0 <*> pure stMatULex <*> pure stMatAex
    stMatAext <- thawMatrix =<< (LA.tr' <$> freezeMatrix stMatAex)
    --hnfLLLST stMatURtex stMatAext
    join $ hermiteNFST <$> newMatrix 0 0 0 <*> pure stMatURtex <*> pure stMatAext
    -- Write out the result
    setMatrix stMatA diagSz diagSz =<< (LA.tr' <$> freezeMatrix stMatAext)
    setMatrix stMatUL diagSz 0 =<< freezeMatrix stMatULex
    setMatrix stMatURt diagSz 0 =<< freezeMatrix stMatURtex

isHeadGCD :: LA.Vector LA.Z -> Bool
isHeadGCD = isHeadGCD' . LA.toList
  where
    isHeadGCD' [] = True
    isHeadGCD' (x:xs) = L.foldl' (\b y -> b && rem y x == 0) True xs

-- Compute the submatrix containing all the non-zero diagonal entries
extractMaxNZDiag :: LA.Matrix LA.Z -> (Int,LA.Matrix LA.Z)
extractMaxNZDiag mx = runST $ do
  dRef <- newSTRef 0
  let p d = d < uncurry min (LA.size mx) && (mx!d!d) /= 0
  whileM_ (p <$> readSTRef dRef) $ modifySTRef' dRef (+1)
  d <- readSTRef dRef
  return (d,LA.subMatrix (0,0) (d,d) mx)

-- | Normalize diagonals to form a divisible chain; i.e. c_1 | c_2 | c_3 | ...
normalizeST :: STMatrix s LA.Z -> STMatrix s LA.Z ->  STMatrix s LA.Z -> ST s ()
normalizeST stMatUL stMatA stMatURt = do
  (dlen,matD) <- liftSTMatrix extractMaxNZDiag stMatA
  stMatD <- thawMatrix matD
  stMatULex <- thawMatrix =<< extractMatrix stMatUL (RowRange 0 (dlen-1)) AllCols
  stMatURtex <- thawMatrix =<< extractMatrix stMatURt (RowRange 0 (dlen-1)) AllCols
  normalizeST' stMatULex stMatD stMatURtex
  setMatrix stMatA 0 0 =<< freezeMatrix stMatD
  setMatrix stMatUL 0 0 =<< freezeMatrix stMatULex
  setMatrix stMatURt 0 0 =<< freezeMatrix stMatURtex
  where
    normalizeST' stMatULex stMatD stMatURtex = do
      b <- (isHeadGCD . LA.takeDiag) <$> freezeMatrix stMatD
      n <- liftSTMatrix LA.rows stMatD
      unless b $ do
        forM_ [1..(n-1)] $ \i -> do
          readMatrix stMatD i i >>= writeMatrix stMatD i 0
          rowOper (AXPY 1 i 0 AllCols) stMatURtex
        preSmithNFST stMatULex stMatD stMatURtex
      when (n > 1) $ do
        stMatDex <- thawMatrix =<< extractMatrix stMatD (FromRow 1) (FromCol 1)
        stMatULexex <- thawMatrix =<< extractMatrix stMatULex (FromRow 1) AllCols
        stMatURtexex <- thawMatrix =<< extractMatrix stMatURtex (FromRow 1) AllCols
        normalizeST' stMatULexex stMatDex stMatURtexex
        setMatrix stMatD 1 1 =<< freezeMatrix stMatDex
        setMatrix stMatULex 1 0 =<< freezeMatrix stMatULexex
        setMatrix stMatURtex 1 0 =<< freezeMatrix stMatURtexex

-- | Compute Smith normal form together with transform unimodular matrices.
smithNF :: LA.Matrix LA.Z -> (LA.Matrix LA.Z, LA.Matrix LA.Z, LA.Matrix LA.Z)
smithNF matA = runST $ do
  stMatA <- thawMatrix matA
  stMatUL <- thawMatrix $ LA.ident (LA.rows matA)
  stMatURt <- thawMatrix $ LA.ident (LA.cols matA)
  preSmithNFST stMatUL stMatA stMatURt
  normalizeST stMatUL stMatA stMatURt
  resMatS <- freezeMatrix stMatA
  resMatUL <- freezeMatrix stMatUL
  resMatURt <- freezeMatrix stMatURt
  return (resMatUL, resMatS, LA.tr' resMatURt)

-- | For debug.
-- Compute "pre-Smith form."
preSmithNF :: LA.Matrix LA.Z -> (LA.Matrix LA.Z, LA.Matrix LA.Z, LA.Matrix LA.Z)
preSmithNF matA = runST $ do
  stMatA <- thawMatrix matA
  stMatUL <- thawMatrix $ LA.ident (LA.rows matA)
  stMatURt <- thawMatrix $ LA.ident (LA.cols matA)
  preSmithNFST stMatUL stMatA stMatURt
  resMatS <- freezeMatrix stMatA
  resMatUL <- freezeMatrix stMatUL
  resMatURt <- freezeMatrix stMatURt
  return (resMatUL, resMatS, LA.tr' resMatURt)