-- {-# OPTIONS_GHC -fplugin-opt=ConCat.Plugin:trace #-}
-- {-# OPTIONS_GHC -fplugin-opt=ConCat.Plugin:showResiduals #-}
-- {-# OPTIONS_GHC -fplugin-opt=ConCat.Plugin:showCcc #-}

-- To run:
--
--   stack build :misc-examples
--
--   stack build :misc-trace >& ~/Haskell/concat/out/o1
--
-- You might also want to use stack's --file-watch flag for automatic recompilation.

{-# LANGUAGE CPP                     #-}
{-# LANGUAGE FlexibleContexts        #-}
{-# LANGUAGE TypeApplications        #-}
{-# LANGUAGE TypeOperators           #-}
{-# LANGUAGE ScopedTypeVariables     #-}
{-# LANGUAGE ConstraintKinds         #-}
{-# LANGUAGE LambdaCase              #-}
{-# LANGUAGE TypeFamilies            #-}
{-# LANGUAGE ViewPatterns            #-}
{-# LANGUAGE PatternSynonyms         #-}
{-# LANGUAGE DataKinds               #-}

-- For OkLC as a class
{-# LANGUAGE UndecidableInstances    #-}
{-# LANGUAGE FlexibleInstances       #-}
{-# LANGUAGE MultiParamTypeClasses   #-}

{-# OPTIONS_GHC -Wall #-}

{-# OPTIONS -Wno-type-defaults #-}

{-# OPTIONS_GHC -Wno-missing-signatures #-}
{-# OPTIONS_GHC -Wno-unused-imports #-}

-- Now in concat-examples.cabal
-- {-# OPTIONS_GHC -fplugin=ConCat.Plugin #-}

-- {-# OPTIONS_GHC -fplugin-opt=ConCat.Plugin:maxSteps=100 #-}

-- {-# OPTIONS_GHC -fno-specialise #-}

-- {-# OPTIONS_GHC -ddump-simpl #-}
-- {-# OPTIONS_GHC -dverbose-core2core #-}

-- {-# OPTIONS_GHC -ddump-rule-rewrites #-}
-- {-# OPTIONS_GHC -ddump-rules #-}

-- Does this flag make any difference?
-- {-# OPTIONS_GHC -fexpose-all-unfoldings #-}

-- Tweak simpl-tick-factor from default of 100
-- {-# OPTIONS_GHC -fsimpl-tick-factor=2500 #-}
{-# OPTIONS_GHC -fsimpl-tick-factor=500 #-}
-- {-# OPTIONS_GHC -fsimpl-tick-factor=250 #-}
-- {-# OPTIONS_GHC -fsimpl-tick-factor=25  #-}
-- {-# OPTIONS_GHC -fsimpl-tick-factor=5  #-}

-- {-# OPTIONS_GHC -dsuppress-uniques #-}
{-# OPTIONS_GHC -dsuppress-idinfo #-}
-- {-# OPTIONS_GHC -dsuppress-module-prefixes #-}

-- {-# OPTIONS_GHC -ddump-tc-trace #-}

-- {-# OPTIONS_GHC -dsuppress-all #-}

-- {-# OPTIONS_GHC -fno-float-in #-}
-- {-# OPTIONS_GHC -ffloat-in #-}
-- {-# OPTIONS_GHC -fdicts-cheap #-}
{-# OPTIONS_GHC -fdicts-strict #-}

-- For experiments
{-# OPTIONS_GHC -Wno-orphans #-}

----------------------------------------------------------------------
-- |
-- Module      :  Examples
-- Copyright   :  (c) 2017 Conal Elliott
-- License     :  BSD3
--
-- Maintainer  :  conal@conal.net
-- Stability   :  experimental
--
-- Suite of automated tests
----------------------------------------------------------------------

module Main where

import Prelude hiding (unzip,zip,zipWith) -- (id,(.),curry,uncurry)
import qualified Prelude as P

import Data.Monoid (Sum(..))
import Data.Foldable (fold)
import Control.Applicative (liftA2)
import Control.Arrow (second)
import Control.Monad ((<=<))
import Data.List (unfoldr)  -- TEMP
import Data.Complex (Complex)
import GHC.Exts (inline)
import GHC.Float (int2Double)
import GHC.TypeLits ()

-- packFiniteM experiment
import GHC.TypeLits
import Data.Maybe (fromJust)
import GHC.Integer

import Data.Constraint (Dict(..),(:-)(..))
import Data.Pointed
import Data.Key
import Data.Distributive
import Data.Functor.Rep
import qualified Data.Functor.Rep as FR
import Data.Vector.Sized (Vector)

import ConCat.Misc
import ConCat.Rep (HasRep(..))
import qualified ConCat.Category as C
import ConCat.Incremental (andInc,inc)
import ConCat.Dual
import ConCat.GAD
import ConCat.Additive
import ConCat.AdditiveFun
import ConCat.AD
import ConCat.ADFun hiding (D)
-- import qualified ConCat.ADFun as ADFun
import ConCat.RAD
import ConCat.Free.VectorSpace (HasV(..),distSqr,(<.>),normalizeV)
-- import ConCat.GradientDescent
import ConCat.Interval
import ConCat.Syntactic (Syn,render)
import ConCat.Circuit (GenBuses,(:>))
import qualified ConCat.RunCircuit as RC
import qualified ConCat.AltCat as A
-- import ConCat.AltCat
import ConCat.AltCat
  (toCcc,toCcc',unCcc,unCcc',conceal,(:**:)(..),Ok,Ok2,U2,equal)
import qualified ConCat.Rep
import ConCat.Rebox () -- necessary for reboxing rules to fire
import ConCat.Orphans ()
import ConCat.Nat
import ConCat.Shaped
-- import ConCat.Scan
-- import ConCat.FFT
import ConCat.Free.LinearRow -- (L,OkLM,linearL)
import ConCat.LC
import ConCat.Deep
import qualified ConCat.Deep as D
-- import ConCat.Finite
import ConCat.Isomorphism
import ConCat.TArr
import qualified ConCat.TArr as TA

-- Experimental
import qualified ConCat.Inline.SampleMethods as I

import qualified ConCat.Regress as R
import ConCat.Free.Affine
import ConCat.Choice
-- import ConCat.RegressChoice

-- import ConCat.Vector -- (liftArr2,FFun,arrFFun)  -- and (orphan) instances
#ifdef CONCAT_SMT
import ConCat.SMT
#endif

-- These imports bring newtype constructors into scope, allowing CoerceCat (->)
-- dictionaries to be constructed. We could remove the LinearRow import if we
-- changed L from a newtype to data, but we still run afoul of coercions for
-- GHC.Generics newtypes.
--
-- TODO: Find a better solution!
import qualified GHC.Generics as G
import qualified ConCat.Free.LinearRow
import qualified Data.Monoid

-- For FFT
import GHC.Generics hiding (C,R,D)

import Control.Newtype.Generics (Newtype(..))

-- Experiments
import GHC.Exts (Coercible,coerce)

import Miscellany

main :: IO ()
main = sequence_ [
    putChar '\n' -- return ()

    -- , runSyn S.t14'

  -- , runSyn $ toCcc $ A.addC . A.dup
  -- , runSyn $ toCcc $ A.addC . (A.id A.&&& A.id)

  -- -- first add . first dup

  -- , runSynStack $ toCcc $ id @ Int

  -- , runSynStack $ toCcc $ negate . negate @ Int

  -- , runSynStack $ toCcc $ \ (x :: Int) -> x  -- fine
  -- , runSyn $ toCcc $ \ (x :: Int) -> (x,x) -- fine
  -- , runSyn $ toCcc $ A.id A.&&& A.id -- fine
  -- , runSyn $ toCcc $ A.addC . (A.id A.&&& A.id) -- okay

  -- , runSynStack $ toCcc $ A.addC . A.dup

  -- -- first add . ((first swapP . (lassocP . id . rassocP) . first swapP) . lassocP . id . rassocP) . first dup
  -- , runSynStack $ toCcc $ A.addC . (A.id A.&&& A.id)

  -- -- [first dup,first addC]
  -- , runChainStack $ toCcc $ A.addC . A.dup

  -- -- [first dup,rassocP,lassocP,first swapP,rassocP,lassocP,first swapP,first addC]
  -- -- [first dup,first addC]  -- with AltCat rule (id &&& id) = dup
  -- , runChainStack $ toCcc $ A.addC . (A.id A.&&& A.id)

  -- -- twice x = x + x
  -- , runSyn        $ toCcc twice  -- addC . dup
  -- , runSynStack   $ toCcc twice  -- first addC . first dup

  -- , runSynStack $ toCcc $ \ x     -> x + x  -- addC . dup, [Dup,Add]
  -- , runSynStack $ toCcc $ \ (x,y) -> x * y  -- mulC, [Mul]

  -- -- addC . (exl *** mulC . (const 3 *** exr) . dup) . dup
  -- -- [Dup,Push,Exl,Pop,Swap,Push,Dup,Push,Const 3,Pop,Swap,Push,Exr,Pop,Swap,Mul,Pop,Swap,Add]
  -- -- TODO: better categorical optimization
  -- , runSynStack $ toCcc $ \ (x,y) -> x + 3 * y

  -- -- addC . (mulC . (const 2 *** exl) . dup *** mulC . (const 3 *** exr) . dup) . dup
  -- -- [Dup,Push,Dup,Push,Const 2,Pop,Swap,Push,Exl,Pop,Swap,Mul,Pop,Swap,Push,Dup,Push,Const 3,Pop,Swap,Push,Exr,Pop,Swap,Mul,Pop,Swap,Add]
  -- , runSynStack $ toCcc $ \ (x,y) -> 2 * x + 3 * y

  -- , runSynStack $ toCcc $ \ y -> 3 * y

  -- , runSyn $ toCcc' $ A.lassocP A.. A.rassocP                  -- id
  -- , runSyn $ toCcc' $ A.lassocP A.. A.id A.. A.rassocP         -- id
  -- , runSyn $ toCcc' $ A.first A.id                             -- id
  -- , runSyn $ toCcc' $ A.lassocP A.. A.first A.id A.. A.rassocP -- id

  -- , runSynCirc "sum-fun-B" $ toCcc $ sum @((->) Bool) @Int
  -- , runSynCirc "sum-fun-BxB" $ toCcc $ sum @((->) (Bool :* Bool)) @Int

  -- , runSynCirc "foo" $ toCcc $ \ (arr :: Arr Bool Int) -> arr `FR.index` False -- okay

  -- , runSynCirc "sum-arr-B"   $ toCcc $ sum @(Arr Bool) @Int -- okay

  -- , runSynCirc "arrSplitProd-B-B" $ toCcc $ arrSplitProd @Bool @Bool @Int -- okay

  -- , runSynCirc "foo" $ toCcc $ arrSplitProd @Bool @Bool @(Sum Int) -- okay

  -- , runSynCirc "foo" $ toCcc $ fmap fold . arrSplitProd @Bool @Bool @(Sum Int) -- fail

  -- , runSynCirc "foo" $ toCcc $ fmap @(Vector 2) not -- okay

  -- , runSynCirc "foo" $ toCcc $ fmap @(Arr Bool) not -- okay

  -- , runSynCirc "foo" $ toCcc $ fmap @(Arr Bool) @(Arr Bool (Sum Int)) fold

  -- , runSynCirc "foo" $ toCcc $ fold @(Arr Bool) @(Sum Int) --

  -- , runSynCirc "foo" $ toCcc $ (TA.!) @Bool @(Sum Int)

  -- , runSynCirc "sum-arr-BxB" $ toCcc $ sum @(Arr (Bool :* Bool)) @Int -- okay

  -- , runSynCirc "sum-arr-BxBxB-r" $ toCcc $ sum @(Arr (Bool :* (Bool :* Bool))) @Int -- okay
  -- , runSynCirc "sum-arr-BxBxB-l" $ toCcc $ sum @(Arr ((Bool :* Bool) :* Bool)) @Int -- okay

  -- , runSynCirc "foo" $ toCcc $ sum @((->) (Bool :* (Bool :* ()))) @Int -- okay
  -- , runSynCirc "foo" $ toCcc $ sum @(Arr (Bool :* ())) @Int -- okay
  -- , runSynCirc "foo" $ toCcc $ sum @(Arr (Bool :* (Bool :* ()))) @Int -- okay

  -- , runSynCirc "sum-flat-rbin-3" $ toCcc $ sum @(Flat (RBin N3)) @Int -- okay
  -- , runSynCirc "sum-flat-lbin-3" $ toCcc $ sum @(Flat (LBin N3)) @Int -- okay

  -- , runSynCirc "fmap-rbin-4" $ toCcc $ fmap @(Flat (LBin N4)) not --

  -- , runSynCirc "add-int" $ toCcc $ (+) @Int -- fine

  -- , runSynCirc "add-int" $ toCcc $ (+) @Integer --

  -- , runSynCirc "foo" $ toCcc $ toFin @Bool -- works

  -- , runSyn $ toCcc $ unFin @Bool -- breaks

  -- , runSyn $ toCcc $ finVal @2 --

  -- , runSyn $ toCcc $ unFin @() -- works

  -- , runSyn{-Circ "foo"-} $ toCcc $ sum @((->) (Finite 5)) @Int -- breaks

  -- , runSynCirc "foo" $ toCcc $ sum @((->) (Bool :* (Bool :* ()))) @Int -- works

  -- , runCirc "foo" $ toCcc $ fmap @(Arr (Bool :* (Bool :* ()))) not -- works

  -- , runCirc "foo" $ toCcc $ fmap @(Flat (RBin N6)) not -- works

  -- , runCirc "foo" $ toCcc $ sum @(Arr (Bool :* (Bool :* ()))) @Int -- nope

  -- , runCirc "foo" $ toCcc $ fmap @(Arr (Bool :* (Bool :* ()))) not -- works with Syn but not Circ

  -- , runSynCirc "sum-rbin-3" $ toCcc $ sum @(RBin N3) @Int

  -- , runSyn{-Circ "sum-flat-rbin-1"-} $ toCcc $ sum @(Flat (RBin N1)) @Int -- ?

  -- , runSynCirc "packFinite" $ toCcc $ packFinite @5 -- fail (missing INLINE)

  -- , runSynCirc "packFiniteM" $ toCcc $ packFiniteM @5 @Maybe -- okay

  -- , runSynCirc "vecIndexDef" $ toCcc $ vecIndexDef @5 @R -- okay

  -- , runSynCirc "finite" $ toCcc $ Finite @5 -- okay

  -- , runSynCirc "add-integer" $ toCcc $ uncurry ((+) @Integer) -- fine

  -- , runSynCirc "sum1" $ toCcc $ sum @Par1 @R -- ??

  -- , runSynCirc "sum2" $ toCcc $ sum @Pair @R -- ??

  -- , runSynCirc "sumV" $ toCcc $ sum @(Vector 5) @R -- ??

  -- , runSynCirc "sumAV" $ toCcc $ sumA @(Vector 5) @R -- ??

  -- , runSynCirc "foo" $ toCcc $ andDerR $ \ () -> zero @((Vector 3 :.: Bump (Vector 2)) R) -- works

  -- , runSynCirc "foo" $ toCcc $ andDerR $ \ () -> zero @((Vector 3 :.: Vector 2) R) -- o2 works

  -- , runSynCirc "foo" $ toCcc $ \ () -> zero @((Vector 3 :.: Vector 2) R) -- o3 fail

  -- , runSynCirc "foo" $ toCcc $ \ () -> zero @((Vector 3 :.: Par1) R) -- o4 fail; o7, o1 (without vector-sized mod)

  -- , runSynCirc "foo" $ toCcc $ point @(Vector 3 :.: Par1) @R -- o6 fail

  -- , runSynCirc "foo" $ toCcc $ \ () -> zero @((Par1 :.: Par1) R) -- okay
  -- , runSynCirc "foo" $ toCcc $ \ () -> zero @(Vector 3 R) -- okay
  -- , runSynCirc "foo" $ toCcc $ \ () -> zero @((Par1 :.: Par1) R) -- okay
  -- , runSynCirc "foo" $ toCcc $ point @(Vector 3 :.: Vector 2) @R -- fail (unsized Vector?)
  -- , runSynCirc "foo" $ toCcc $ andDerR $ \ () -> zero @(Vector 3 R) -- okay
  -- , runSynCirc "foo" $ toCcc $ (^+^) @(Vector 3 R) -- okay
  -- , runSynCirc "foo" $ toCcc $ andDerR $ id @(Vector 2 R) -- okay
  -- , runSynCirc "foo" $ toCcc $ andDerR $ id @((Vector 3 :.: Bump (Vector 2)) R) -- okay

  -- , runSynCirc "step-lr1" $ toCcc $ D.step (lr1 @(Vector 2) @(Vector 3)) 0.01

  -- , runSynCirc "errGrad-lr1" $ toCcc $ D.errGrad (lr1 @(Vector 2) @(Vector 3)) -- fails with cast confusion

  -- , runSynCirc "step-lr2" $ toCcc $ D.step (lr2 @(Vector 2) @(Vector 3) @(Vector 5)) 0.01

    -- , runSynCirc "fmap" $ toCcc $ fmap @(Vector 7) not

    -- , runSynCirc "constV" $ toCcc $ \ () -> pure 1 :: Vector 7 Int

    -- , runSynCirc "pointV" $ toCcc $ point @(Vector 7) @Bool -- fail

    -- , runSynCirc "addV" $ toCcc $ (^+^) @(Vector 7 R)

    -- , runSynCirc "zeroV" $ toCcc $ \ () -> zero @(Vector 7 R)

    -- , runSynCirc "addR" $ toCcc $ (^+^) @R

    -- , runSyn $ toCcc $ \ () -> finite @2 (nat @1) -- okay

    -- , runSyn $ toCcc $ \ (i :: Finite 2) -> finite @2 (nat @2 - 1) -- okay

    -- , runSyn $ toCcc $ \ (i :: Finite 2) -> finite (nat @2 - 1) - i -- okay

    -- , runSyn $ toCcc $ reverseFinite @5  -- Fine

    -- , runSynCirc "reverseFinite" $ toCcc $ reverseFinite @5

    -- , runSynCirc "reverseF-pair" $ toCcc $ reverseF @Pair @Int

    -- , runSynCirc "reverseF-vec" $ toCcc $ reverseF @(Vector 5) @Int

    -- , runSynCirc "reverseFin-Bool-802" $ toCcc $ isoFwd (reverseFinIso @Bool)

    -- , runSynCirc "reverseFin-vec" $ toCcc $ isoFwd (reverseFinIso @(A.Finite 5))

    -- , runSynCirc "reverseFin-Bool-cheat" $ toCcc $ unFin @Bool . (1 -) . toFin @Bool

    -- -- https://github.com/conal/concat/issues/49
    -- -- , putStrLn $ render (toCcc (\f g x -> f (g (x))))
    -- , putStrLn $ render (toCcc (\ (f :: Int -> Bool) (g :: Bool -> Int) x -> f (g (x))))

  -- Circuit graphs
  , runSynCirc "add"         $ toCcc $ (+) @R
  , runSynCirc "add-uncurry" $ toCcc $ uncurry ((+) @R)
  , runSynCirc "dup"         $ toCcc $ A.dup @(->) @R
  , runSynCirc "fst"         $ toCcc $ fst @R @R
  , runSynCirc "twice"       $ toCcc $ twice @R
  , runSynCirc "sqr"         $ toCcc $ sqr @R
  , runSynCirc "complex-mul" $ toCcc $ uncurry ((*) @C)
  , runSynCirc "magSqr"      $ toCcc $ magSqr @R
  , runSynCirc "cosSinProd"  $ toCcc $ cosSinProd @R
  , runSynCirc "xp3y"        $ toCcc $ \ (x,y) -> x + 3 * y :: R
  , runSynCirc "horner"      $ toCcc $ horner @R [1,3,5]
  , runSynCirc "cos-2xx"     $ toCcc $ \ x -> cos (2 * x * x) :: R
  , runSynCirc "log-2xx"     $ toCcc $ \ x -> log (2 * x * x) :: R

  -- -- Automatic differentiation
  -- , runSynCircDers "add"     $ uncurry ((+) @R)
  -- , runSynCircDers "fst"     $ fst @R @R
  -- , runSynCircDers "twice"   $ twice @R
  -- , runSynCircDers "sqr"     $ sqr @R
  -- , runSynCircDers "sin"     $ sin @R
  -- , runSynCircDers "cos"     $ cos @R
  -- , runSynCircDers "magSqr"  $ magSqr  @R
  -- , runSynCircDers "cos-2x"  $ \ x -> cos (2 * x) :: R
  -- , runSynCircDers "cos-2xx" $ \ x -> cos (2 * x * x) :: R
  -- , runSynCircDers "cos-xpy" $ \ (x,y) -> cos (x + y) :: R
  -- , runSynCircDers "cos-xpytz" $ \ (x,y,z) -> cos (x + y * z) :: R

  -- , runSynCirc "cos-xpy-adr-802" $ toCcc $ andGradR $ \ (x,y) -> cos (x + y) :: R

  -- , runSynCirc "sqr-adr-802" $ toCcc $ andGradR $ sqr @R

  -- , runSynCirc "magSqr-adr"  $ toCcc $ andDerR $ magSqr  @R
  -- , runSynCirc "cosSinProd-adr"  $ toCcc $ andDerR $ cosSinProd @R
  -- , runSynCirc "cosSinProd-gradr"  $ toCcc $ andGrad2R $ cosSinProd @R

  -- , runSynCirc "cosSinProd-adf" $ toCcc $ andDerF $ cosSinProd @R
  -- , runSynCirc "cosSinProd-adr" $ toCcc $ andDerR $ cosSinProd @R

  -- , runCirc "affRelu"         $ toCcc $                affRelu @(Vector 2) @(Vector 3) @R
  -- , runCirc "affRelu-err"     $ toCcc $ errSqrSampled (affRelu @(Vector 2) @(Vector 3) @R)
  -- , runCirc "affRelu-errGrad" $ toCcc $ errGrad (affRelu @(Vector 2) @(Vector 3) @R) -- fail

  -- , runCirc "affRelu2"         $ toCcc $                lr2 @(Vector 5) @(Vector 3) @(Vector 2)
  -- , runCirc "affRelu2-err"     $ toCcc $ errSqrSampled (lr2 @(Vector 5) @(Vector 3) @(Vector 2))
  -- , runCirc "affRelu2-errGrad" $ toCcc $ errGrad       (lr2 @(Vector 5) @(Vector 3) @(Vector 2))

  -- , runCirc "affRelu3"         $ toCcc $                lr3 @(Vector 7) @(Vector 5) @(Vector 3) @(Vector 2)
  -- , runCirc "affRelu3-err"     $ toCcc $ errSqrSampled (lr3 @(Vector 7) @(Vector 5) @(Vector 3) @(Vector 2))
  -- , runCirc "affRelu3-errGrad" $ toCcc $ errGrad       (lr3 @(Vector 7) @(Vector 5) @(Vector 3) @(Vector 2))

  ]

data P = P R R

instance HasRep P where
  type Rep P = R :* R
  repr (P x y) = (x,y)
  abst (x,y) = P x y

instance Additive P where
  zero = P zero zero
  P a b ^+^ P c d = P (a+c) (b+d)

#if 0

-- | Convert an 'Integer' into a 'Finite', returning 'Nothing' if the input is out of bounds.
-- This version has an INLINE pragma.
packFiniteM :: forall n m. (KnownNat n, Monad m) => Integer -> m (Finite n)
packFiniteM x | 0 <= x && x < natValAt @n = return (Finite x)
              | otherwise                 = fail "packFiniteM: bad index"
{-# INLINE packFiniteM #-}

-- Index a sized vector with an integer, given a default
vecIndexDef :: KnownNat n => a -> Vector n a -> Integer -> a
vecIndexDef def v i = maybe def (FR.index v) (packFiniteM i)
{-# INLINE vecIndexDef #-}

#endif

-- foo :: Stack Syn Int Int
-- foo = toCcc $ A.addC . (A.id A.&&& A.id) --

-- foo = toCcc $ A.addC . (A.id A.&&& A.id) --
-- foo = reveal $ toCcc $ A.addC . (A.id A.&&& A.id) --
-- foo = toCcc $ \ x -> x + x
-- foo = toCcc $ A.reveal (A.addC . (A.id A.&&& A.id))

-- z2 :: Syn ((Int :* Bool) :* z) ((Int :* Bool) :* z)
-- z2 = S.z2

-- z1 :: Int -> Int
-- z1 = A.addC A.. (A.id A.&&& A.id)

-- z2 :: Stack Syn Int Int
-- z2 = A.addC A.. (A.id A.&&& A.id)

-- z2' :: Stack Syn Int Int
-- z2' = A.reveal (A.addC A.. (A.id A.&&& A.id))

-- z3 :: Stack Syn Int Int
-- z3 = toCcc (\ x -> x + x)

-- z4 :: Stack Syn Int Int
-- z4 = toCcc' (\ x -> x + x)

-- z5 :: Syn (Int :* ()) (Int :* ())
-- z5 = S.unStack (toCcc' (\ x -> x + x))

-- z5 :: Syn (Int :* ()) (Int :* ())
-- z5 = S.unStack (toCcc (\ x -> x + x))

-- z2' :: Syn (Int :* ()) (Int :* ())
-- z2' = unStack z2

-- z3 :: Syn (Int :* ()) (Int :* ())
-- z3 = toCcc $ \ x -> x + x

-- z3' :: Syn (Int :* ()) (Int :* ())
-- z3' = unStack z3

-- z5 :: Stack Syn Int Int
-- z5 = A.negateC A.. A.negateC

-- z5' :: Stack Syn Int Int
-- z5' = A.reveal (A.negateC A.. A.negateC)

-- z6 :: Stack Syn (Int :* Int) (Int :* Int)
-- z6 = A.negateC A.*** A.negateC

-- z6' :: Stack Syn (Int :* Int) (Int :* Int)
-- z6' = A.reveal (A.negateC A.*** A.negateC)

-- z6'' :: Stack Syn (Int :* Int) (Int :* Int)
-- z6'' = C.negateC C.*** C.negateC

-- z7 :: Stack Syn (Int :* Int) (Int :* Int)
-- z7 = C.negateC `C.crossSecondFirst` C.negateC

-- z7' :: Stack Syn (Int :* Int) (Int :* Int)
-- z7' = C.negateC `A.crossSecondFirst` C.negateC

-- z8 :: Stack Syn (Int :* Int) (Int :* Int)
-- z8 = C.negateC C.*** C.negateC

-- z8' :: Stack Syn (Int :* Int) (Int :* Int)
-- z8' = C.negateC `S.cross` C.negateC

-- z8'' :: Stack Syn (Int :* Int) (Int :* Int)
-- z8'' = C.negateC `C.crossSecondFirst` C.negateC

-- z8' :: Stack Syn (Int :* Int) (Int :* Int)
-- z8' = C.negateC A.*** C.negateC

-- z9 :: Stack Syn (Int :* Int) (Int :* Int)
-- z9 = C.negateC `S.cross` C.negateC


{-# LANGUAGE FlexibleContexts #-}
{-# OPTIONS_GHC -fno-warn-unused-do-bind #-}
module Parse (readExpr, readExprList) where
import LispVals

import Control.Monad (liftM, join)
import Text.ParserCombinators.Parsec (
    unexpected,
    char
    , try
    , (<|>)
    , many
    , many1
    , string
    , sepEndBy
    , endBy
    , Parser
    , oneOf
    , noneOf
    , skipMany1
    , space
    , anyChar
    , letter
    , digit
    , octDigit
    , hexDigit
    , parse)
import Numeric (readOct, readHex, readFloat)
import Data.Char (digitToInt)
import Data.Ratio ((%))
import Data.Complex (Complex((:+)))
import Control.Monad.Error (MonadError, throwError)
import qualified Data.Vector as V (fromList)

-- $setup
-- >>> import Text.ParserCombinators.Parsec (parse)

symbol :: Parser Char
symbol = oneOf "!$%&|*+-/:<=?>@^_~"

spaces :: Parser ()
spaces = skipMany1 space

-- |
-- >>> parse parseString "lisp" "\"he\\nllo\""
-- Right "he\nllo"
parseString :: Parser LispVal
parseString = do
    char '\"'
    x <- many (many1 (noneOf "\\\"") <|> do
        char '\\'
        s <- oneOf "nrt\"\\"
        case s of
            '\"' -> return "\""
            '\\' -> return "\\"
            't' -> return "\t"
            'n' -> return "\n"
            'r' -> return "\r"
            c -> unexpected [c])
    char '\"'
    return $ String $ join x

-- |
-- >>> parse parseAtom "lisp" "hello"
-- Right hello
parseAtom :: Parser LispVal
parseAtom = do
    first <- letter <|> symbol
    rest <- many (letter <|> digit <|> symbol)
    return $ Atom $ first : rest

-- |
-- >>> parse parseBool "lisp" "#t"
-- Right #t
--
-- >>> parse parseBool "lisp" "#f"
-- Right #f
parseBool :: Parser LispVal
parseBool = try $ do
    char '#'
    x <- oneOf "tf"
    case x of
        't' -> return $ Bool True
        'f' -> return $ Bool False
        c -> unexpected [c]

-- |
-- >>> parse parseSignedNumber "lisp" "+#x123"
-- Right 291
--
-- >>> parse parseSignedNumber "lisp" "+#o123"
-- Right 83
--
-- >>> parse parseSignedNumber "lisp" "+#d123"
-- Right 123
--
-- >>> parse parseSignedNumber "lisp" "-#b101"
-- Right -5
--
-- >>> parse parseSignedNumber "lisp" "-123"
-- Right -123
--
-- >>> parse parseSignedNumber "lisp" "-3+2i"
-- Right -3.0+2.0i
--
-- >>> parse parseSignedNumber "lisp" "-3/2"
-- Right -3/2
parseSignedNumber :: Parser LispVal
parseSignedNumber = try $ do
    signChar <- oneOf "+-"
    let sign = case signChar of
                '-' -> -1 :: Integer
                _ -> 1
    ureal <- parseComplex <|> parseRatio <|> parseFloat <|> parsePrefixNumber <|> parseDecimal
    case ureal of
                Ratio r -> return $ Ratio (r * fromIntegral sign)
                Complex (r :+ i) -> return $ Complex $ (fromIntegral sign * r) :+ i
                Float f -> return $ Float $ fromIntegral sign * f
                Number n -> return $ Number $ sign * n
                _ -> unexpected (show ureal)

-- |
-- >>> parse parseUnsignedNumber "lisp" "#x123"
-- Right 291
--
-- >>> parse parseUnsignedNumber "lisp" "#o123"
-- Right 83
--
-- >>> parse parseUnsignedNumber "lisp" "#d123"
-- Right 123
--
-- >>> parse parseUnsignedNumber "lisp" "#b101"
-- Right 5
--
-- >>> parse parseUnsignedNumber "lisp" "123"
-- Right 123
parseUnsignedNumber :: Parser LispVal
parseUnsignedNumber = parsePrefixNumber <|> parseDecimal

parsePrefixNumber :: Parser LispVal
parsePrefixNumber = parseOctal
    <|> parseHexadecimal
    <|> parseBinary
    <|> (do
        try $ string "#d"
        parseDecimal)

parseOctal :: Parser LispVal
parseOctal = do
    try $ string "#o"
    os <- many1 octDigit
    return $ (Number . fst . head . readOct) os

parseHexadecimal :: Parser LispVal
parseHexadecimal = do
    try $ string "#x"
    os <- many1 hexDigit
    return $ (Number . fst . head . readHex) os

parseBinary :: Parser LispVal
parseBinary = do
    try $ string "#b"
    bs <- many1 (oneOf "01")
    return $ (Number . readBinary) bs

readBinary :: Num a => String -> a
readBinary = foldl (\x y -> x * 2 + (fromIntegral . digitToInt) y) 0

parseDecimal :: Parser LispVal
parseDecimal = do
    ds <- many1 digit
    return $ (Number . read) ds

-- |
-- >>> parse parseChar "list" "#\\c"
-- Right #\c
--
-- >>> parse parseChar "lisp" "#\\space"
-- Right #\space
parseChar :: Parser LispVal
parseChar = do
    try $ string "#\\"
    cs <- parseCharLiteral <|> anyChar
    return $ Char cs

parseCharLiteral :: Parser Char
parseCharLiteral = try $ do
    cs <- string "space" <|> string "newline"
    case cs of
        "space" -> return ' '
        "newline" -> return '\n'
        s -> unexpected s

parseFloat :: Parser LispVal
parseFloat = try $ do
    integer <- many1 digit
    char '.'
    fractional <- many1 digit
    return $ (Float . fst . head . readFloat) (integer ++ "." ++ fractional)

parseRatio :: Parser LispVal
parseRatio = try $ do
    numerator <- many1 digit
    char '/'
    denominator <- many1 digit
    return $ Ratio $ read numerator % read denominator

-- |
-- >>> parse parseComplex "lisp" "3.2+2i"
-- Right 3.2+2.0i
--
-- >>> parse parseComplex "lisp" "+2i"
-- Right 0.0+2.0i
--
-- >>> parse parseComplex "lisp" "-2i"
-- Right 0.0-2.0i
parseComplex :: Parser LispVal
parseComplex = parseImaginary <|> try (do
    real <- many1 digit
    realFrac <- (char '.' >> many1 digit) <|> return "0"
    sign <- oneOf "+-"
    complex <- many digit
    let okComplex = case complex of
                        [] -> "1"
                        _ -> complex
    complexFrac <- (char '.' >> many1 digit) <|> return "0"
    char 'i'
    return $ Complex  ((fst . head . readFloat) (real ++ "." ++ realFrac) :+
                        (fst . head . readFloat) (okComplex ++ "." ++ complexFrac) * case sign of
                            '-' -> -1
                            _ -> 1))

parseImaginary :: Parser LispVal
parseImaginary = try $ do
    sign <- (oneOf "+-" >>= \s -> return [s]) <|> return ""
    complex <- many1 digit
    complexFrac <- (char '.' >> many1 digit) <|> return "0"
    char 'i'
    let okComplex = case complex of
                        [] -> "1"
                        _ -> complex
    return $ Complex (0.0 :+ (fst . head . readFloat) (okComplex ++ "." ++ complexFrac) * case sign of
                            ('-':_) -> -1
                            _ -> 1)

-- |
-- >>> parse parseList "lisp" "1 2 3"
-- Right (1 2 3)
parseList :: Parser LispVal
parseList = liftM List $ sepEndBy parseExpr spaces

-- |
-- >>> parse parseDottedList "lisp" "1 . 2"
-- Right (1 . 2)
parseDottedList :: Parser LispVal
parseDottedList = do
    h <- endBy parseExpr spaces
    t <- char '.' >> spaces >> parseExpr
    return $ DottedList h t

parseQuoted :: Parser LispVal
parseQuoted = do
    char '\''
    x <- parseExpr
    return $ List [Atom "quote", x]

parseUnquoted :: Parser LispVal
parseUnquoted = do
    char ','
    x <- parseExpr
    return $ List [Atom "unquote", x]

parseQuasiQuoted :: Parser LispVal
parseQuasiQuoted = do
    char '`'
    x <- parseExpr
    return $ List [Atom "quasiquote", x]

-- |
-- >>> parse parseVector "lisp" "#(1 2 3)"
-- Right #(1 2 3)
parseVector :: Parser LispVal
parseVector = try $ do
    string "#("
    many space
    x <- sepEndBy parseExpr spaces
    string ")"
    return $ Vector $ V.fromList x

parseBracketedParser :: Parser LispVal -> Parser LispVal
parseBracketedParser p = try $ do
    char '('
    many space
    x <- p
    char ')'
    return x

-- |
-- >>> parse parseExpr "lisp" "hello"
-- Right hello
--
-- >>> parse parseExpr "lisp" "#(1 2 3)"
-- Right #(1 2 3)
--
-- >>> parse parseExpr "lisp" "3+2i"
-- Right 3.0+2.0i
--
-- >>> parse parseExpr "lisp" "3.2+2i"
-- Right 3.2+2.0i
--
-- >>> parse parseExpr "lisp" "3/2"
-- Right 3/2
--
-- >>> parse parseExpr "lisp" "3.2"
-- Right 3.2
--
-- >>> parse parseExpr "lisp" "-3.2"
-- Right -3.2
--
-- >>> parse parseExpr "lisp" "+#x123"
-- Right 291
--
-- >>> parse parseExpr "lisp" "#x123"
-- Right 291
--
-- >>> parse parseExpr "lisp" "+#o123"
-- Right 83
--
-- >>> parse parseExpr "lisp" "+#d123"
-- Right 123
--
-- >>> parse parseExpr "lisp" "123"
-- Right 123
--
-- >>> parse parseExpr "lisp" "-#b101"
-- Right -5
--
-- >>> parse parseExpr "lisp" "#t"
-- Right #t
--
-- >>> parse parseExpr "lisp" "#f"
-- Right #f
--
-- >>> parse parseExpr "list" "#\\c"
-- Right #\c
--
-- >>> parse parseExpr "lisp" "#\\space"
-- Right #\space
--
-- >>> parse parseExpr "lisp" "atom"
-- Right atom
--
-- >>> parse parseExpr "lisp" "'(a list)"
-- Right (quote (a list))
--
-- >>> parse parseExpr "lisp" "`(a list)"
-- Right (quasiquote (a list))
--
-- >>> parse parseExpr "lisp" "`(a ,( + 1 2))"
-- Right (quasiquote (a (unquote (+ 1 2))))
--
-- >>> parse parseExpr "lisp" "#(1 2 3 )"
-- Right #(1 2 3)
--
-- >>> parse parseExpr "lisp" "(1 2 3)"
-- Right (1 2 3)
--
-- >>> parse parseExpr "lisp" "( 1 2 . 3)"
-- Right (1 2 . 3)
--
-- >>> parse parseExpr "lisp" "\"he\\nllo\""
-- Right "he\nllo"
parseExpr :: Parser LispVal
parseExpr = parseString
    <|> parseComplex
    <|> parseRatio
    <|> parseFloat
    <|> parseSignedNumber
    <|> parseUnsignedNumber
    <|> parseBool
    <|> parseChar
    <|> parseAtom
    <|> parseQuoted
    <|> parseQuasiQuoted
    <|> parseUnquoted
    <|> parseVector
    <|> parseBracketedParser parseList
    <|> parseBracketedParser parseDottedList

readExpr :: (MonadError LispError m) => String -> m LispVal
readExpr = readOrThrow parseExpr

readExprList :: (MonadError LispError m) => String -> m [LispVal]
readExprList = readOrThrow (endBy parseExpr spaces)

readOrThrow :: (MonadError LispError m) => Parser a -> String -> m a
readOrThrow parser input = case parse parser "lisp" input of
    Left err -> throwError $ Parser err
    Right val -> return val


-----------------
-- GHC Options --
-----------------

{-# OPTIONS_GHC -O2                       #-}
{-# OPTIONS_GHC -Wno-unused-imports       #-}
{-# OPTIONS_GHC -Wno-missing-import-lists #-}

-------------------------
-- Language Extensions --
-------------------------

{-# LANGUAGE BlockArguments       #-}
{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE LambdaCase           #-}
{-# LANGUAGE MultiWayIf           #-}
{-# LANGUAGE NegativeLiterals     #-}
{-# LANGUAGE OverloadedLists      #-}
{-# LANGUAGE OverloadedStrings    #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE StrictData           #-}
{-# LANGUAGE TupleSections        #-}
{-# LANGUAGE TypeApplications     #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE ExtendedDefaultRules #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ViewPatterns         #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NamedFieldPuns #-}

module Main where

------------------
-- Import Lists --
------------------

import           Control.Arrow                  ( (>>>) )
import qualified Control.Monad                 as M
import qualified Control.Monad.Primitive       as Prim
import qualified Control.Monad.RWS             as RWS
import qualified Control.Monad.Reader          as Reader
import qualified Control.Monad.ST              as ST
import qualified Control.Monad.State           as State
import qualified Control.Monad.Writer          as Writer
import qualified Data.Array.IArray             as A
import qualified Data.Array.IO                 as AIO
import qualified Data.Array.ST                 as AST
import           Data.Bifunctor
import qualified Data.Bits                     as Bits
import qualified Data.ByteString.Char8         as BS
import qualified Data.Char                     as Char
import qualified Data.Complex                  as Comp
import           Data.Foldable
import qualified Data.Function                 as Func
import qualified Data.Functor                  as Functor
import qualified Data.IORef                    as IORef
import qualified Data.IntPSQ                   as PSQueue
import qualified Data.Ix                       as Ix
import qualified Data.List                     as L
import qualified Data.Map.Strict               as Map
import qualified Data.Maybe                    as Maybe
import qualified Data.Primitive.MutVar         as MutVar
import qualified Data.Proxy                    as Proxy
import qualified Data.Ratio                    as Ratio
import qualified Data.STRef                    as STRef
import qualified Data.Sequence                 as Seq
import           Data.Sequence                  ( (<|)
                                                , (><)
                                                , ViewL((:<), EmptyL)
                                                , ViewR((:>), EmptyR)
                                                , viewl
                                                , viewr
                                                , (|>)
                                                )
import qualified Data.Set                      as Set
import qualified Data.Text                     as T
import qualified Data.Text.IO                  as TIO
import qualified Data.Tree                     as Tree
import qualified Data.Vector                   as V
import qualified Data.Vector.Algorithms.Merge  as VAM
import qualified Data.Vector.Algorithms.Radix  as VAR
import qualified Data.Vector.Algorithms.Search as VAS
import           Data.Vector.Generic            ( (!)
                                                , (!?)
                                                )
import qualified Data.Vector.Generic           as VG
import qualified Data.Vector.Generic.Mutable   as VGM
import qualified Data.Vector.Mutable           as VM
import qualified Data.Vector.Unboxed
import qualified Data.Vector.Unboxing          as VU
import qualified Data.Vector.Unboxing.Mutable  as VUM
import qualified Debug.Trace                   as Trace
import qualified GHC.Generics
import qualified GHC.TypeNats                  as TypeNats
import           Prelude                 hiding ( (!)
                                                , head
                                                , map
                                                , print
                                                , tail
                                                , uncons
                                                )
import qualified Prelude

----------
-- Main --
----------

yes, no :: BS.ByteString
yes = "Yes"
no = "No"

main :: IO ()
main = do
  [a, b] <- get @[I]
  print $ a + b

-------------
-- Library --
-------------

type V = V.Vector
type VU = VU.Vector
type VM = VM.MVector
type VUM = VUM.MVector
type BS = BS.ByteString
type I = Int
type IG = Int
type D = Double
type B = Bool
type S = String

default (V.Vector, [], BS.ByteString, String, Int, Double)

infixl 1 #
(#) :: a -> (a -> b) -> b
(#) a f = f a

map :: Monad m => (a -> b) -> m a -> m b
map = fmap

---------
-- I/O --
---------

-- | ex) get @I, get @(V I) ..
get :: ReadText a => IO a
get = readText <$> BS.getLine

-- | ex) getLn @(V I) n, getLn @[I] n
getLines :: Int -> forall a . ReadTextLines a => IO a
getLines n = readTextLines <$> M.replicateM n BS.getLine

-- |
output :: ShowText a => a -> IO ()
output = BS.putStr . showText

-- |
outputLines :: ShowTextLines a => a -> IO ()
outputLines = BS.putStr . BS.unlines . showTextLines

-- |
print :: ShowText a => a -> IO ()
print = BS.putStrLn . showText

-- |
printLines :: ShowTextLines a => a -> IO ()
printLines = BS.putStrLn . BS.unlines . showTextLines

---------------
-- Read/Show --
---------------

-- | TextRead
class ReadText a where
  readText :: BS.ByteString -> a

class ShowText a where
  showText :: a -> BS.ByteString

instance ReadText Int where
  readText s = case BS.readInt s of
    Just (i, _) -> i
    Nothing     -> error "readText Int"

instance ReadText Integer where
  readText s = case BS.readInteger s of
    Just (i, _) -> i
    Nothing     -> error "readText Integer"

instance ReadText Double where
  readText = read . BS.unpack

instance ReadText BS.ByteString where
  readText = id

instance ReadText a => ReadText (V.Vector a) where
  readText = readVec

instance (ReadText a, VUM.Unboxable a) => ReadText (VU.Vector a) where
  readText = readVec

instance ReadText a => ReadText [a] where
  readText = map readText . BS.words

instance (ReadText a, ReadText b) => ReadText (a, b) where
  readText (BS.words -> [a, b]) = (readText a, readText b)
  readText _ = error "Invalid Format :: readText :: BS -> (a, b)"

instance (ReadText a, ReadText b, ReadText c) => ReadText (a, b, c) where
  readText (BS.words -> [a, b, c]) = (readText a, readText b, readText c)
  readText _ = error "Invalid Format :: readText :: BS -> (a, b, c)"

instance (ReadText a, ReadText b, ReadText c, ReadText d) => ReadText (a, b, c, d) where
  readText (BS.words -> [a, b, c, d]) =
    (readText a, readText b, readText c, readText d)
  readText _ = error "Invalid Format :: readText :: BS -> (a, b, c, d)"

instance ShowText Integer where
  showText = BS.pack . show

instance ShowText Int where
  showText = BS.pack . show

instance ShowText Double where
  showText = BS.pack . show

instance ShowText Bool where
  showText True  = yes
  showText False = no

instance ShowText BS.ByteString where
  showText = id

instance (ShowText a) => ShowText (V.Vector a) where
  showText = showVec

instance (ShowText a, VU.Unboxable a) => ShowText (VU.Vector a) where
  showText = showVec

instance ShowText a => ShowText [a] where
  showText = BS.unwords . map showText

instance (ShowText a, ShowText b) => ShowText (a, b) where
  showText (a, b) = showText a `BS.append` " " `BS.append` showText b

instance (ShowText a, ShowText b, ShowText c) => ShowText (a, b, c) where
  showText (a, b, c) =
    showText a
      `BS.append` " "
      `BS.append` showText b
      `BS.append` " "
      `BS.append` showText c

instance (ShowText a, ShowText b, ShowText c, ShowText d) => ShowText (a, b, c, d) where
  showText (a, b, c, d) =
    showText a
      `BS.append` " "
      `BS.append` showText b
      `BS.append` " "
      `BS.append` showText c
      `BS.append` " "
      `BS.append` showText d

readVec :: (VG.Vector v a, ReadText a) => BS.ByteString -> v a
readVec = VG.fromList . readText

showVec :: (VG.Vector v a, ShowText a) => v a -> BS.ByteString
showVec = showText . VG.toList

class ReadTextLines a where
  readTextLines :: [BS.ByteString] -> a

class ShowTextLines a where
  showTextLines :: a -> [BS.ByteString]

instance ReadText a => ReadTextLines [a] where
  readTextLines = map readText

instance ReadText a => ReadTextLines (V.Vector a) where
  readTextLines = readVecLines

instance (ReadText a, VU.Unboxable a) => ReadTextLines (VU.Vector a) where
  readTextLines = readVecLines

instance ReadTextLines BS.ByteString where
  readTextLines = BS.unlines

instance (ReadText a, ReadText b) => ReadTextLines (a, b) where
  readTextLines (a : b : _) = (readText a, readText b)
  readTextLines _ = error "Invalid Format :: readTextLines :: [BS] -> (a, b)"

instance (ReadText a, ReadText b, ReadText c) => ReadTextLines (a, b, c) where
  readTextLines (a : b : c : _) = (readText a, readText b, readText c)
  readTextLines _ =
    error "Invalid Format :: readTextLines :: [BS] -> (a, b, c)"

instance ShowText a => ShowTextLines [a] where
  showTextLines = map showText

instance ShowText a => ShowTextLines (V.Vector a) where
  showTextLines = showVecLines

instance (ShowText a, VU.Unboxable a) => ShowTextLines (VU.Vector a) where
  showTextLines = showVecLines

instance ShowTextLines BS.ByteString where
  showTextLines s = [s]

instance (ShowText a, ShowText b) => ShowTextLines (a, b) where
  showTextLines (a, b) = [showText a, showText b]

instance (ShowText a, ShowText b, ShowText c) => ShowTextLines (a, b, c) where
  showTextLines (a, b, c) = [showText a, showText b, showText c]

readVecLines :: (VG.Vector v a, ReadText a) => [BS.ByteString] -> v a
readVecLines = VG.fromList . map readText

showVecLines :: (VG.Vector v a, ShowText a) => v a -> [BS.ByteString]
showVecLines = map showText . VG.toList

------------
-- ModInt --
------------

newtype Mod a (p :: TypeNats.Nat) = ModInt a deriving (Eq, Show, VU.Unboxable)

instance ShowText a => ShowText (Mod a p) where
  showText (ModInt x) = showText x

instance ReadText a => ReadText (Mod a p) where
  readText = ModInt . readText

instance (TypeNats.KnownNat p, Integral a) => Num (Mod a p) where
  (ModInt x) + (ModInt y) = ModInt $ (x + y) `mod` p
    where p = fromIntegral $ TypeNats.natVal (Proxy.Proxy :: Proxy.Proxy p)
  (ModInt x) * (ModInt y) = ModInt $ (x * y) `mod` p
    where p = fromIntegral $ TypeNats.natVal (Proxy.Proxy :: Proxy.Proxy p)
  negate (ModInt x) = ModInt $ -x `mod` p
    where p = fromIntegral $ TypeNats.natVal (Proxy.Proxy :: Proxy.Proxy p)
  abs = id
  signum _ = 1
  fromInteger n = ModInt $ fromInteger n `mod` p
    where p = fromIntegral $ TypeNats.natVal (Proxy.Proxy :: Proxy.Proxy p)

instance (TypeNats.KnownNat p, Integral a) => Fractional (Mod a p) where
  recip (ModInt n)
    | gcd n p /= 1 = error
      "recip :: Mod a p -> Mod a p : The inverse element does not exisBS."
    | otherwise = ModInt . fst $ extendedEuc n (-p)
    where p = fromIntegral $ TypeNats.natVal (Proxy.Proxy :: Proxy.Proxy p)
  fromRational r = ModInt n / ModInt d   where
    n = fromInteger $ Ratio.numerator r
    d = fromInteger $ Ratio.denominator r

------------
-- InfInt --
------------

data Sign = Positive | Negative | NA deriving (Eq, Show)

negateSign :: Sign -> Sign
negateSign Positive = Negative
negateSign Negative = Positive
negateSign NA       = NA

addSign :: Sign -> Sign -> Sign
addSign Positive Positive = Positive
addSign Positive Negative = NA
addSign Negative Positive = NA
addSign Negative Negative = Negative
addSign _        _        = NA

mulSign :: Sign -> Sign -> Sign
mulSign Positive Positive = Positive
mulSign Positive Negative = Negative
mulSign Negative Positive = Negative
mulSign Negative Negative = Positive
mulSign _        _        = NA

compareSign :: Sign -> Sign -> Ordering
compareSign Positive Positive = EQ
compareSign Positive Negative = GT
compareSign Negative Positive = LT
compareSign Negative Negative = EQ
compareSign _        _        = EQ

data Inf a = Infinity Sign | Finity a deriving (Eq, Show)

instance ReadText a => ReadText (Inf a) where
  readText = \case
    "Infinity"  -> Infinity Positive
    "-Infinity" -> Infinity Negative
    "NA"        -> Infinity NA
    x           -> Finity $ readText x

instance ShowText a => ShowText (Inf a) where
  showText = \case
    Infinity Positive -> "Infinity"
    Infinity Negative -> "-Infinity"
    Infinity NA       -> "NA"
    Finity   x        -> showText x

instance (Num a, Ord a) => Num (Inf a) where
  x + y = case x of
    Infinity sign -> case y of
      Infinity sign' -> Infinity $ addSign sign sign'
      Finity   y'    -> Infinity sign
    Finity x' -> case y of
      Infinity sign' -> Infinity sign'
      Finity   y'    -> Finity $ x' + y'

  x * y = case x of
    Infinity sign -> case y of
      Infinity sign' -> Infinity $ mulSign sign sign'
      Finity y' | y' < 0    -> Infinity $ negateSign sign
                | y' == 0   -> Infinity NA
                | otherwise -> Infinity sign
    Finity x' -> case y of
      Infinity sign | x' < 0    -> Infinity $ negateSign sign
                    | x' == 0   -> Infinity NA
                    | otherwise -> Infinity sign
      Finity y' -> Finity $ x' * y'
  negate = \case
    Infinity sign -> Infinity $ negateSign sign
    Finity   x    -> Finity $ negate x
  abs = \case
    Infinity sign -> Infinity Positive
    Finity   x    -> Finity $ abs x
  signum = \case
    Infinity Positive -> Finity 1
    Infinity Negative -> Finity -1
    Infinity NA       -> Infinity NA
    Finity   x        -> Finity $ signum x
  fromInteger n = Finity $ fromInteger n

instance (Num a, Ord a) => Ord (Inf a) where
  compare = \case
    Infinity sign -> \case
      Infinity sign' -> compareSign sign sign'
      Finity   _     -> if sign == Positive then GT else LT
    Finity x -> \case
      Infinity sign -> if sign == Positive then LT else GT
      Finity   y    -> compare x y

infinity :: Num a => Inf a
infinity = Infinity Positive

na :: Inf a
na = Infinity NA

dropWhileRev :: VG.Vector v a => (a -> Bool) -> v a -> v a
dropWhileRev f xs | VG.null xs = VG.empty
                  | f x        = dropWhileRev f xs'
                  | otherwise  = xs
 where
  x   = VG.last xs
  xs' = VG.init xs

------------------
-- Disjoint Set --
------------------

type DisjointSet = VU.Vector Int
data DisjointSetM m = DSet
  { dsParents :: VUM.MVector m Int
  , dsDepths  :: VUM.MVector m Int
  }

newDSet :: Prim.PrimMonad m => Int -> m (DisjointSetM (Prim.PrimState m))
newDSet n = DSet <$> VU.thaw (VU.generate n id) <*> VUM.replicate n 1

root :: DisjointSet -> Int -> Int
root xs i | xs ! i == i = i
          | otherwise   = root xs $ xs ! i

find :: DisjointSet -> Int -> Int -> Bool
find xs i j = root xs i == root xs j

-- |
rootM :: Prim.PrimMonad m => DisjointSetM (Prim.PrimState m) -> Int -> m Int
rootM ds i = VUM.read (dsParents ds) i >>= \p -> if p == i
  then pure i
  else rootM ds p >>= \r -> VUM.write (dsParents ds) i r >> pure r

findM
  :: Prim.PrimMonad m => DisjointSetM (Prim.PrimState m) -> Int -> Int -> m Bool
findM ds i j = (==) <$> rootM ds i <*> rootM ds j

union
  :: Prim.PrimMonad m => DisjointSetM (Prim.PrimState m) -> Int -> Int -> m ()
union ds i j = do
  rooti <- rootM ds i
  rootj <- rootM ds j
  depi  <- VUM.read (dsDepths ds) rooti
  depj  <- VUM.read (dsDepths ds) rootj
  if
    | depi == depj
    -> VUM.modify (dsDepths ds) (+ 1) rooti
      >> VUM.write (dsParents ds) rootj rooti
    | depi > depj
    -> VUM.write (dsParents ds) rootj rooti
    | otherwise
    -> VUM.write (dsParents ds) rooti rootj

----------------------------
-- Monadic Priority Queue --
----------------------------

type MPSQueue m p v = MutVar.MutVar m (PSQueue.IntPSQ p v)

mpsqNull :: Prim.PrimMonad m => MPSQueue (Prim.PrimState m) p v -> m Bool
mpsqNull q = PSQueue.null <$> MutVar.readMutVar q

mpsqEmpty :: Prim.PrimMonad m => m (MPSQueue (Prim.PrimState m) p v)
mpsqEmpty = MutVar.newMutVar PSQueue.empty

mpsqSingleton
  :: Prim.PrimMonad m
  => Ord p => Int -> p -> v -> m (MPSQueue (Prim.PrimState m) p v)
mpsqSingleton k p v = MutVar.newMutVar $ PSQueue.singleton k p v

mpsqInsert
  :: Prim.PrimMonad m
  => Ord p => MPSQueue (Prim.PrimState m) p v -> Int -> p -> v -> m ()
mpsqInsert m k p v = MutVar.modifyMutVar' m (PSQueue.insert k p v)

-- |
mpsqFindMin
  :: Prim.PrimMonad m
  => Ord p => MPSQueue (Prim.PrimState m) p v -> m (Maybe (Int, p, v))
mpsqFindMin q = PSQueue.findMin <$> MutVar.readMutVar q

-- |
mpsqMinView
  :: Prim.PrimMonad m
  => Ord p => MPSQueue (Prim.PrimState m) p v -> m (Maybe (Int, p, v))
mpsqMinView q = do
  res <- PSQueue.minView <$> MutVar.readMutVar q
  case res of
    Nothing            -> pure Nothing
    Just (k, p, v, q') -> do
      MutVar.writeMutVar q q'
      pure $ Just (k, p, v)

----------------------
-- Monadic Sequence --
----------------------

type MSequence m a = MutVar.MutVar m (Seq.Seq a)

msqEmpty :: Prim.PrimMonad m => m (MSequence (Prim.PrimState m) a)
msqEmpty = MutVar.newMutVar Seq.empty

msqSingleton :: Prim.PrimMonad m => a -> m (MSequence (Prim.PrimState m) a)
msqSingleton x = MutVar.newMutVar $ Seq.singleton x

infixr 5 <<|
infixl 5 |>>

-- | Cons
(<<|) :: Prim.PrimMonad m => a -> MSequence (Prim.PrimState m) a -> m ()
x <<| xs = MutVar.modifyMutVar' xs (x Seq.<|)

-- | Snoc
(|>>) :: Prim.PrimMonad m => MSequence (Prim.PrimState m) a -> a -> m ()
xs |>> x = MutVar.modifyMutVar' xs (Seq.|> x)

msqViewL :: Prim.PrimMonad m => MSequence (Prim.PrimState m) a -> m (Maybe a)
msqViewL xs = do
  xs' <- MutVar.readMutVar xs
  case viewl xs' of
    EmptyL -> pure Nothing
    x :< _ -> pure (Just x)

msqViewR :: Prim.PrimMonad m => MSequence (Prim.PrimState m) a -> m (Maybe a)
msqViewR xs = do
  xs' <- MutVar.readMutVar xs
  case viewr xs' of
    EmptyR    -> pure Nothing
    xs'' :> x -> pure (Just x)

msqPopL :: Prim.PrimMonad m => MSequence (Prim.PrimState m) a -> m (Maybe a)
msqPopL xs = do
  xs' <- MutVar.readMutVar xs
  case viewl xs' of
    EmptyL    -> pure Nothing
    x :< xs'' -> do
      MutVar.writeMutVar xs xs''
      pure (Just x)

msqPopR :: Prim.PrimMonad m => MSequence (Prim.PrimState m) a -> m (Maybe a)
msqPopR xs = do
  xs' <- MutVar.readMutVar xs
  case viewr xs' of
    EmptyR    -> pure Nothing
    xs'' :> x -> do
      MutVar.writeMutVar xs xs''
      pure (Just x)

msqNull :: Prim.PrimMonad m => MSequence (Prim.PrimState m) a -> m Bool
msqNull xs = do
  xs' <- MutVar.readMutVar xs
  pure $ Seq.null xs'

-----------
-- Graph --
-----------

type Graph a = V.Vector [(Int, a)] --a
type UGraph = Graph ()

gEmpty :: Graph a
gEmpty = V.singleton []

gFromEdges :: VG.Vector v (Int, Int, a) => Int -> v (Int, Int, a) -> Graph a
gFromEdges n edges = ST.runST do
  v <- VM.replicate n []
  VG.forM_ edges \(i, j, a) -> VM.modify v ((j, a) :) i
  V.freeze v

-- |
gReverse :: Graph a -> Graph a
gReverse g = ST.runST do
  let n = V.length g
  v <- VM.replicate n []
  V.forM_ [0 .. n - 1] \i -> M.forM_ (g ! i) \(j, a) -> VM.modify v ((i, a) :) j
  V.freeze v

-- |
-- |
-- | (Listcons)
dijkstra :: Graph Int -> Int -> V.Vector (Inf Int, [Int])
dijkstra g i = V.create do
  let n = V.length g

  --
  -- dists ! j == (i, )
  dists <- VM.replicate n (infinity, [])
  VM.write dists i (0, [i])

  --
  -- key: , : , value: ()
  queue <- mpsqSingleton i 0 ()

  let --
      -- prev: , v: , alt: ()
      -- prev
      update prev v alt = do
        (_, xs) <- VM.read dists prev
        VM.write dists v (alt, v : xs)
        --
        -- ( C++  continue )
        mpsqInsert queue v alt ()

      --
      -- u:
      processing u = do
        --
        (dist_u, _) <- VM.read dists u
        M.forM_
          (g ! u)
          (\(v, cost) -> do
            (dist_v, _) <- VM.read dists v
            let alt = dist_u + fromIntegral cost
            M.when (dist_v > alt) $ update u v alt
          )

  --
  while do
    res <- mpsqMinView queue
    case res of
      Nothing        -> pure False
      Just (u, _, _) -> do
        processing u
        pure True
  pure dists

dfs :: Graph a -> Int -> Tree.Tree Int
dfs g i = ST.runST do
  reached <- VM.replicate (V.length g) False
  let loop now = do
        VM.write reached now True
        next     <- M.filterM (fmap not . VM.read reached) . map fst $ g ! now
        children <- M.mapM loop next
        pure $ Tree.Node now children
  loop i

----------
-- Maze --
----------

type Maze = V.Vector (VU.Vector Char)
-- ^ Char

mzHeight :: Maze -> Int
mzHeight = V.length

mzWidth :: Maze -> Int
mzWidth v = maybe 0 VU.length (v !? 0)

type Rules a = Char -> Char -> Maybe a

readMaze :: BS.ByteString -> Maze
readMaze str = V.fromList $ map (VU.fromList . BS.unpack) $ BS.lines str

getMaze :: Int -> IO Maze
getMaze h = readMaze <$> getLines h

mazePos :: Maze -> (Int, Int) -> Int
mazePos mz (i, j) = i * mzWidth mz + j

mazeToGraph :: Rules a -> Maze -> Graph a
mazeToGraph rules maze =
  let h = mzHeight maze
      w = mzWidth maze
  in  gFromEdges (h * w)
        $ V.fromList
        $ Maybe.catMaybes
        $ [ edge
          | i        <- [0 .. h - 1]
          , j        <- [0 .. w - 1]
          , (i', j') <- [(i - 1, j), (i, j - 1), (i + 1, j), (i, j + 1)]
          , i' >= 0
          , i' < h
          , j' >= 0
          , j' < w
          , let c    = maze ! i ! j
          , let c' = maze ! i' ! j'
          , let edge = (w * i + j, w * i' + j', ) <$> rules c c'
          ]

----------
-- Tree --
----------

depth :: Eq a => Tree.Tree a -> a -> Maybe Int
depth (Tree.Node a []) b | a == b = Just 0
depth (Tree.Node _ []) _          = Nothing
depth (Tree.Node a xs) b          = case Maybe.mapMaybe (`depth` b) xs of
  [] -> Nothing
  xs -> Just $ 1 + minimum xs

--------------
-- MultiSet --
--------------

type MultiSet a = Map.Map a Int

msEmpty :: MultiSet a
msEmpty = Map.empty

msSingleton :: a -> MultiSet a
msSingleton a = Map.singleton a 1

msFromList :: Ord a => [a] -> MultiSet a
msFromList xs = Map.fromListWith (+) $ map (, 1) xs

msInsert :: Ord a => a -> MultiSet a -> MultiSet a
msInsert a = Map.insertWith (+) a 1

msInserts :: Ord a => a -> Int -> MultiSet a -> MultiSet a
msInserts = Map.insertWith (+)

msDelete :: Ord a => a -> MultiSet a -> MultiSet a
msDelete = Map.update (\n -> if n > 1 then Just (n - 1) else Nothing)

msDeletes :: Ord a => a -> Int -> MultiSet a -> MultiSet a
msDeletes a n = Map.update (\n' -> if n' > n then Just (n' - n) else Nothing) a

msDeleteAll :: Ord a => a -> MultiSet a -> MultiSet a
msDeleteAll = Map.delete

msMember :: Ord a => a -> MultiSet a -> Bool
msMember = Map.member

msLookupLT :: Ord a => a -> MultiSet a -> Maybe a
msLookupLT a = fmap fst . Map.lookupLT a

msLookupGT :: Ord a => a -> MultiSet a -> Maybe a
msLookupGT a = fmap fst . Map.lookupGT a

msLookupLE :: Ord a => a -> MultiSet a -> Maybe a
msLookupLE a = fmap fst . Map.lookupLE a

msLookupGE :: Ord a => a -> MultiSet a -> Maybe a
msLookupGE a = fmap fst . Map.lookupGE a

msNull :: MultiSet a -> Bool
msNull = Map.null

msUnion :: Ord a => MultiSet a -> MultiSet a -> MultiSet a
msUnion = Map.unionWith (+)

-- | split MultiSet into two MultiSet
-- | (LT, GE)
msSplitLTGE :: Ord a => a -> MultiSet a -> (MultiSet a, MultiSet a)
msSplitLTGE a ms =
  let (lt, x, gt) = Map.splitLookup a ms
      ge          = case x of
        Nothing -> gt
        Just v  -> msInserts a v gt
  in  (lt, ge)

msSplitLEGT :: Ord a => a -> MultiSet a -> (MultiSet a, MultiSet a)
msSplitLEGT a ms =
  let (lt, x, gt) = Map.splitLookup a ms
      le          = case x of
        Nothing -> lt
        Just v  -> msInserts a v lt
  in  (le, gt)

msElemAtL :: Ord a => Int -> MultiSet a -> Maybe a
msElemAtL i ms = case Map.minViewWithKey ms of
  Nothing -> Nothing
  Just ((k, v), rest) | i < v     -> Just k
                      | otherwise -> msElemAtL (i - v) rest

msElemAtG :: Ord a => Int -> MultiSet a -> Maybe a
msElemAtG i ms = case Map.maxViewWithKey ms of
  Nothing -> Nothing
  Just ((k, v), rest) | i < v     -> Just k
                      | otherwise -> msElemAtG (i - v) rest

---------------------
-- MutableMultiSet --
---------------------

type MutableMultiSet m a = MutVar.MutVar m (MultiSet a)

mmsEmpty :: Prim.PrimMonad m => m (MutableMultiSet (Prim.PrimState m) a)
mmsEmpty = MutVar.newMutVar msEmpty

mmsSingleton
  :: Prim.PrimMonad m => a -> m (MutableMultiSet (Prim.PrimState m) a)
mmsSingleton a = MutVar.newMutVar $ msSingleton a

mmsFromList
  :: Prim.PrimMonad m
  => Ord a => [a] -> m (MutableMultiSet (Prim.PrimState m) a)
mmsFromList = MutVar.newMutVar . msFromList

mmsInsert
  :: Prim.PrimMonad m
  => Ord a => MutableMultiSet (Prim.PrimState m) a -> a -> m ()
mmsInsert mms a = MutVar.modifyMutVar' mms $ msInsert a

mmsInserts
  :: Prim.PrimMonad m
  => Ord a => MutableMultiSet (Prim.PrimState m) a -> a -> Int -> m ()
mmsInserts mms a n = MutVar.modifyMutVar' mms $ msInserts a n

mmsDelete
  :: Prim.PrimMonad m
  => Ord a => MutableMultiSet (Prim.PrimState m) a -> a -> m ()
mmsDelete mms a = MutVar.modifyMutVar' mms $ msDelete a

mmsDeletes
  :: Prim.PrimMonad m
  => Ord a => MutableMultiSet (Prim.PrimState m) a -> a -> Int -> m ()
mmsDeletes mms a n = MutVar.modifyMutVar' mms $ msDeletes a n

mmsDeleteAll
  :: Prim.PrimMonad m
  => Ord a => MutableMultiSet (Prim.PrimState m) a -> a -> m ()
mmsDeleteAll mms a = MutVar.modifyMutVar' mms $ msDeleteAll a

mmsMember
  :: Prim.PrimMonad m
  => Ord a => MutableMultiSet (Prim.PrimState m) a -> a -> m Bool
mmsMember mms a = MutVar.readMutVar mms Functor.<&> msMember a

mmsLookupLT
  :: Prim.PrimMonad m
  => Ord a => MutableMultiSet (Prim.PrimState m) a -> a -> m (Maybe a)
mmsLookupLT mms a = MutVar.readMutVar mms Functor.<&> msLookupLT a

mmsLookupGT
  :: Prim.PrimMonad m
  => Ord a => MutableMultiSet (Prim.PrimState m) a -> a -> m (Maybe a)
mmsLookupGT mms a = MutVar.readMutVar mms Functor.<&> msLookupGT a

mmsLookupLE
  :: Prim.PrimMonad m
  => Ord a => MutableMultiSet (Prim.PrimState m) a -> a -> m (Maybe a)
mmsLookupLE mms a = MutVar.readMutVar mms Functor.<&> msLookupLE a

mmsLookupGE
  :: Prim.PrimMonad m
  => Ord a => MutableMultiSet (Prim.PrimState m) a -> a -> m (Maybe a)
mmsLookupGE mms a = MutVar.readMutVar mms Functor.<&> msLookupGE a

mmsNull :: Prim.PrimMonad m => MutableMultiSet (Prim.PrimState m) a -> m Bool
mmsNull mms = MutVar.readMutVar mms Functor.<&> msNull

mmsUnion
  :: Prim.PrimMonad m
  => Ord a
  => MutableMultiSet (Prim.PrimState m) a
  -> MutableMultiSet (Prim.PrimState m) a
  -> m ()
mmsUnion mms1 mms2 = do
  ms1 <- MutVar.readMutVar mms1
  ms2 <- MutVar.readMutVar mms2
  MutVar.modifyMutVar' mms1 $ msUnion ms2

mmsSplitLTGE
  :: Prim.PrimMonad m
  => Ord a
  => MutableMultiSet (Prim.PrimState m) a
  -> a
  -> m (MultiSet a, MultiSet a)
mmsSplitLTGE mms a = MutVar.readMutVar mms Functor.<&> msSplitLTGE a

mmsSplitLEGT
  :: Prim.PrimMonad m
  => Ord a
  => MutableMultiSet (Prim.PrimState m) a
  -> a
  -> m (MultiSet a, MultiSet a)
mmsSplitLEGT mms a = MutVar.readMutVar mms Functor.<&> msSplitLEGT a

mmsElemAtL
  :: Prim.PrimMonad m
  => Ord a => MutableMultiSet (Prim.PrimState m) a -> Int -> m (Maybe a)
mmsElemAtL mms i = MutVar.readMutVar mms Functor.<&> msElemAtL i

mmsElemAtG
  :: Prim.PrimMonad m
  => Ord a => MutableMultiSet (Prim.PrimState m) a -> Int -> m (Maybe a)
mmsElemAtG mms i = MutVar.readMutVar mms Functor.<&> msElemAtG i

------------
-- Others --
------------

head :: VG.Vector v a => v a -> Maybe a
head v = if VG.null v then Nothing else Just (v ! 0)

tail :: VG.Vector v a => v a -> Maybe (v a)
tail v = if VG.null v then Nothing else Just (VG.tail v)

uncons :: VG.Vector v a => v a -> Maybe (a, v a)
uncons v = if VG.null v then Nothing else Just (v ! 0, VG.tail v)

while :: Monad m => m Bool -> m ()
while f = f >>= \frag -> M.when frag $ while f

divisor :: Integral a => a -> [a]
divisor n
  | n <= 0 = error
    "divisor : Definition range does not include negative numbers or zero"
  | otherwise = half ++ rev where
  mid  = floor . sqrt @Double . fromIntegral $ n
  half = filter ((== 0) . mod n) [1 .. mid]
  rev =
    reverse
      . map (n `div`)
      . (if mid ^ (2 :: Int) == n then init else id)
      $ half

primeFact :: forall a b . (Integral a, Integral b) => a -> [(a, b)]
primeFact 1 = []
primeFact n
  | n == 1 = []
  | n <= 0 = error
    "primefact : Definition range does not include negative numbers or zero"
  | otherwise = case L.find ((== 0) . mod n) ([2 .. mid] :: [a]) of
    Nothing -> [(n, 1)]
    Just p  -> (p, m) : primeFact next     where
      m    = loop n p
      next = n `div` (p ^ m)
 where
  loop m q | m `mod` q == 0 = 1 + loop (m `div` q) q
           | otherwise      = 0
  mid = floor . sqrt @Double . fromIntegral $ n

-- |
primes :: forall a . Integral a => a -> [a]
primes n | n <= 1    = []
         | otherwise = filter ((frags !) . fromIntegral) [2 .. n] where
  frags = VU.create do
    v <- VUM.replicate (fromIntegral (n + 1)) True
    VUM.write v 0 False
    VUM.write v 1 False
    VU.forM_
      [2 .. floor . sqrt @Double . fromIntegral $ n]
      \i -> do
        frag <- VUM.read v i
        M.when frag $ VU.forM_ [2 * i, 3 * i .. fromIntegral n]
                               \j -> VUM.write v j False
    pure v

-- |
-- | ax + by = gcd a b
extendedEuc :: (Integral b) => b -> b -> (b, b)
extendedEuc a b | a >= 0 && b == 0 = (1, 0)
                | a < 0 && b == 0  = (-1, 0)
                | otherwise        = (t, s - q * t) where
  (q, r) = divMod a b
  (s, t) = extendedEuc b r


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

{-# OPTIONS_GHC -fno-warn-orphans #-}

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

{-# LANGUAGE AllowAmbiguousTypes       #-}
{-# LANGUAGE BangPatterns              #-}
{-# LANGUAGE CPP                       #-}
{-# LANGUAGE EmptyDataDecls            #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE ForeignFunctionInterface  #-}
{-# LANGUAGE FunctionalDependencies    #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE KindSignatures            #-}
{-# LANGUAGE LambdaCase                #-}
{-# LANGUAGE MagicHash                 #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE Rank2Types                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE StandaloneDeriving        #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeFamilyDependencies    #-}
{-# LANGUAGE TypeInType                #-}
{-# LANGUAGE UnboxedTuples             #-}
{-# LANGUAGE UndecidableInstances      #-}

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

-- | Internal module to Eigen.
--   Here we define all foreign function calls,
--   and some typeclasses integral to the public and private interfaces
--   of the library.
module Eigen.Internal where --   FIXME: Explicit export list

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

import           Control.Monad            (when)
import           Data.Binary              (Binary(put,get))
import           Data.Binary.Get          (getByteString, getWord32be)
import           Data.Binary.Put          (putByteString, putWord32be)
import           Data.Bits                (xor)
import           Data.Complex             (Complex((:+)))
import           Data.Kind                (Type)
import           Data.Proxy               (Proxy(Proxy))
import           Foreign.C.String         (CString, peekCString)
import           Foreign.C.Types          (CInt(CInt), CFloat(CFloat), CDouble(CDouble), CChar)
import           Foreign.ForeignPtr       (ForeignPtr, castForeignPtr, withForeignPtr)
import           Foreign.Ptr              (Ptr, castPtr, nullPtr, plusPtr)
import           Foreign.Storable         (Storable(sizeOf, alignment, poke, peek, peekByteOff, peekElemOff, pokeByteOff, pokeElemOff))
import           GHC.TypeLits             (natVal, KnownNat, Nat)
import           System.IO.Unsafe         (unsafeDupablePerformIO)
import qualified Data.Vector.Storable     as VS
import qualified Data.ByteString          as BS
import qualified Data.ByteString.Internal as BSI

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

-- | Like 'Proxy', but specialised to 'Nat'.
data Row (r :: Nat) = Row
-- | Like 'Proxy', but specialised to 'Nat'.
data Col (c :: Nat) = Col

-- | Used internally. Given a 'KnownNat' constraint, turn the type-level 'Nat' into an 'Int'.
natToInt :: forall n. KnownNat n => Int
{-# INLINE natToInt #-}
natToInt = fromIntegral (natVal @n Proxy)

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

-- | Cast to and from a C-FFI type
--   'Cast' is a closed typeclass with an associated injective type family.
--   It is closed in the sense that we provide only four types
--   with instances for it; and intend for eigen to only be used
--   with those four types. The injectivity of the type family is
--   then useful for avoiding MPTCs. 'Cast' has two functions; 'toC'
--   and 'fromC', where 'toC' goes from a Haskell type to its associated
--   C type for internal use, with the C FFI, and 'fromC' goes from the
--   associated C type to the Haskell type.
class Cast (a :: Type) where
  type family C a = (result :: Type) | result -> a
  toC   :: a -> C a
  fromC :: C a -> a

instance Cast Int where
  type C Int = CInt
  toC = CInt . fromIntegral
  {-# INLINE toC #-}
  fromC (CInt x) = fromIntegral x
  {-# INLINE fromC #-}

instance Cast Float where
  type C Float = CFloat
  toC = CFloat
  {-# INLINE toC #-}
  fromC (CFloat x) = x
  {-# INLINE fromC #-}

instance Cast Double where
  type C Double = CDouble
  toC = CDouble
  {-# INLINE toC #-}
  fromC (CDouble x) = x
  {-# INLINE fromC #-}

instance Cast a => Cast (Complex a) where
  type C (Complex a) = CComplex (C a)
  toC (a :+ b) = CComplex (toC a) (toC b)
  {-# INLINE toC #-}
  fromC (CComplex a b) = (fromC a) :+ (fromC b)
  {-# INLINE fromC #-}

-- | WARNING! 'toC' is lossy for any Int greater than (maxBound :: Int32)!
instance Cast a => Cast (Int, Int, a) where
  type C (Int, Int, a) = CTriplet a
  {-# INLINE toC #-}
  toC (x, y, z) = CTriplet (toC x) (toC y) (toC z)
  {-# INLINE fromC #-}
  fromC (CTriplet x y z) = (fromC x, fromC y, fromC z)

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

-- | Complex number for FFI with the same memory layout as std::complex\<T\>
data CComplex a = CComplex !a !a deriving (Show)

instance Storable a => Storable (CComplex a) where
    sizeOf _ = sizeOf (undefined :: a) * 2
    alignment _ = alignment (undefined :: a)
    poke p (CComplex x y) = do
        pokeElemOff (castPtr p) 0 x
        pokeElemOff (castPtr p) 1 y
    peek p = CComplex
        <$> peekElemOff (castPtr p) 0
        <*> peekElemOff (castPtr p) 1

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

-- | FIXME: Doc
data CTriplet a where
  CTriplet :: Cast a => !CInt -> !CInt -> !(C a) -> CTriplet a

deriving instance (Show a, Show (C a)) => Show (CTriplet a)

instance (Storable a, Elem a) => Storable (CTriplet a) where
    sizeOf _ = sizeOf (undefined :: a) + sizeOf (undefined :: CInt) * 2
    alignment _ = alignment (undefined :: CInt)
    poke p (CTriplet row col val) = do
        pokeElemOff (castPtr p) 0 row
        pokeElemOff (castPtr p) 1 col
        pokeByteOff p (sizeOf (undefined :: CInt) * 2) val
    peek p = CTriplet
        <$> peekElemOff (castPtr p) 0
        <*> peekElemOff (castPtr p) 1
        <*> peekByteOff p (sizeOf (undefined :: CInt) * 2)

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

-- | `Elem` is a closed typeclass that encompasses the properties
--   eigen expects its values to possess, and simplifies the external
--   API quite a bit.
class (Num a, Cast a, Storable a, Storable (C a), Code (C a)) => Elem a

instance Elem Float
instance Elem Double
instance Elem (Complex Float)
instance Elem (Complex Double)

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

-- | Encode a C Type as a CInt
--
--   Hack used in FFI wrapper functions when constructing FFI calls
class Code a where; code :: a -> CInt
instance Code CFloat             where; code _ = 0
instance Code CDouble            where; code _ = 1
instance Code (CComplex CFloat)  where; code _ = 2
instance Code (CComplex CDouble) where; code _ = 3

-- | Hack used in constructing FFI calls.
newtype MagicCode = MagicCode CInt deriving Eq

instance Binary MagicCode where
    put (MagicCode _code) = putWord32be $ fromIntegral _code
    get = MagicCode . fromIntegral <$> getWord32be

-- | Hack used in constructing FFI calls.
magicCode :: Code a => a -> MagicCode
magicCode x = MagicCode (code x `xor` 0x45696730)

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

-- | Machine size of a 'CInt'.
intSize :: Int
intSize = sizeOf (undefined :: CInt)

-- | FIXME: Doc
encodeInt :: CInt -> BS.ByteString
encodeInt x = BSI.unsafeCreate (sizeOf x) $ (`poke` x) . castPtr

-- | FIXME: Doc
decodeInt :: BS.ByteString -> CInt
decodeInt (BSI.PS fp fo fs)
    | fs == sizeOf x = x
    | otherwise = error "decodeInt: wrong buffer size"
    where x = performIO $ withForeignPtr fp $ peek . (`plusPtr` fo)

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

-- | 'Binary' instance for 'Data.Vector.Storable.Mutable.Vector'
instance Storable a => Binary (VS.Vector a) where
    put vs = put (BS.length bs) >> putByteString bs where
        (fp,fs) = VS.unsafeToForeignPtr0 vs
        es = sizeOf (VS.head vs)
        bs = BSI.fromForeignPtr (castForeignPtr fp) 0 (fs * es)

    get = get >>= getByteString >>= \bs -> let
        (fp,fo,fs) = BSI.toForeignPtr bs
        es = sizeOf (VS.head vs)
        -- `plusForeignPtr` is used qualified here to just remind a reader
        -- that it is defined internally within eigen
        vs = VS.unsafeFromForeignPtr0 (Eigen.Internal.plusForeignPtr fp fo) (fs `div` es)
        in return vs

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

-- | FIXME: Doc
data CSparseMatrix a
-- | FIXME: Doc
type CSparseMatrixPtr a = Ptr (CSparseMatrix a)

-- | FIXME: Doc
data CSolver a
-- | FIXME: Doc
type CSolverPtr a = Ptr (CSolver a)

-- {-# INLINE unholyPerformIO #-}
-- unholyPerformIO :: IO a -> a
-- unholyPerformIO (IO m) = case m realWorld# of (# _, r #) -> r

-- | FIXME: replace with unholyPerformIO (?)
performIO :: IO a -> a
performIO = unsafeDupablePerformIO

-- | FIXME: Doc
plusForeignPtr :: ForeignPtr a -> Int -> ForeignPtr b
plusForeignPtr fp fo = castForeignPtr fp1 where
    vs :: VS.Vector CChar
    vs = VS.unsafeFromForeignPtr (castForeignPtr fp) fo 0
    (fp1, _) = VS.unsafeToForeignPtr0 vs

foreign import ccall "eigen-proxy.h free" c_freeString :: CString -> IO ()

call :: IO CString -> IO ()
call func = func >>= \c_str -> when (c_str /= nullPtr) $
    peekCString c_str >>= \str -> c_freeString c_str >> fail str

foreign import ccall "eigen-proxy.h free" free :: Ptr a -> IO ()

foreign import ccall "eigen-proxy.h eigen_setNbThreads" c_setNbThreads :: CInt -> IO ()
foreign import ccall "eigen-proxy.h eigen_getNbThreads" c_getNbThreads :: IO CInt

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

#let api1 name, args = "foreign import ccall \"eigen_%s\" c_%s :: CInt -> %s\n%s :: forall a . Code (C a) => %s\n%s = c_%s (code (undefined :: (C a)))", #name, #name, args, #name, args, #name, #name

#api1 random,        "Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 identity,      "Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 add,           "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 sub,           "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 mul,           "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 diagonal,      "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 transpose,     "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 inverse,       "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 adjoint,       "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 conjugate,     "Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 normalize,     "Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 sum,           "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 prod,          "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 mean,          "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 norm,          "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 trace,         "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 squaredNorm,   "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 blueNorm,      "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 hypotNorm,     "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 determinant,   "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 rank,          "CInt -> Ptr CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 image,         "CInt -> Ptr (Ptr (C a)) -> Ptr CInt -> Ptr CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 kernel,        "CInt -> Ptr (Ptr (C a)) -> Ptr CInt -> Ptr CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 solve,         "CInt -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api1 relativeError, "Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString"

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

#let api2 name, args = "foreign import ccall \"eigen_%s\" c_%s :: CInt -> %s\n%s :: forall a . Code (C a) => %s\n%s = c_%s (code (undefined :: (C a)))", #name, #name, args, #name, args, #name, #name

#api2 sparse_new,           "CInt -> CInt -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_clone,         "CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_fromList,      "CInt -> CInt -> Ptr (CTriplet a) -> CInt -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_toList,        "CSparseMatrixPtr a -> Ptr (CTriplet a) -> CInt -> IO CString"
#api2 sparse_free,          "CSparseMatrixPtr a -> IO CString"
#api2 sparse_makeCompressed,"CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_uncompress,    "CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_isCompressed,  "CSparseMatrixPtr a -> Ptr CInt -> IO CString"
#api2 sparse_transpose,     "CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_adjoint,       "CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_pruned,        "CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_prunedRef,     "CSparseMatrixPtr a -> Ptr (C a) -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_scale,         "CSparseMatrixPtr a -> Ptr (C a) -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_nonZeros,      "CSparseMatrixPtr a -> Ptr CInt -> IO CString"
#api2 sparse_innerSize,     "CSparseMatrixPtr a -> Ptr CInt -> IO CString"
#api2 sparse_outerSize,     "CSparseMatrixPtr a -> Ptr CInt -> IO CString"
#api2 sparse_coeff,         "CSparseMatrixPtr a -> CInt -> CInt -> Ptr (C a) -> IO CString"
#api2 sparse_coeffRef,      "CSparseMatrixPtr a -> CInt -> CInt -> Ptr (Ptr (C a)) -> IO CString"
#api2 sparse_cols,          "CSparseMatrixPtr a -> Ptr CInt -> IO CString"
#api2 sparse_rows,          "CSparseMatrixPtr a -> Ptr CInt -> IO CString"
#api2 sparse_norm,          "CSparseMatrixPtr a -> Ptr (C a) -> IO CString"
#api2 sparse_squaredNorm,   "CSparseMatrixPtr a -> Ptr (C a) -> IO CString"
#api2 sparse_blueNorm,      "CSparseMatrixPtr a -> Ptr (C a) -> IO CString"
#api2 sparse_add,           "CSparseMatrixPtr a -> CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_sub,           "CSparseMatrixPtr a -> CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_mul,           "CSparseMatrixPtr a -> CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_block,         "CSparseMatrixPtr a -> CInt -> CInt -> CInt -> CInt -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_fromMatrix,    "Ptr (C a) -> CInt -> CInt -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api2 sparse_toMatrix,      "CSparseMatrixPtr a -> Ptr (C a) -> CInt -> CInt -> IO CString"
#api2 sparse_values,        "CSparseMatrixPtr a -> Ptr CInt -> Ptr (Ptr (C a)) -> IO CString"
#api2 sparse_outerStarts,   "CSparseMatrixPtr a -> Ptr CInt -> Ptr (Ptr CInt) -> IO CString"
#api2 sparse_innerIndices,  "CSparseMatrixPtr a -> Ptr CInt -> Ptr (Ptr CInt) -> IO CString"
#api2 sparse_innerNNZs,     "CSparseMatrixPtr a -> Ptr CInt -> Ptr (Ptr CInt) -> IO CString"
#api2 sparse_setZero,       "CSparseMatrixPtr a -> IO CString"
#api2 sparse_setIdentity,   "CSparseMatrixPtr a -> IO CString"
#api2 sparse_reserve,       "CSparseMatrixPtr a -> CInt -> IO CString"
#api2 sparse_resize,        "CSparseMatrixPtr a -> CInt -> CInt -> IO CString"

#api2 sparse_conservativeResize,    "CSparseMatrixPtr a -> CInt -> CInt -> IO CString"
#api2 sparse_compressInplace,       "CSparseMatrixPtr a -> IO CString"
#api2 sparse_uncompressInplace,     "CSparseMatrixPtr a -> IO CString"

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

#let api3 name, args = "foreign import ccall \"eigen_%s\" c_%s :: CInt -> CInt -> %s\n%s :: forall s a . (Code s, Code (C a)) => s -> %s\n%s s = c_%s (code (undefined :: (C a))) (code s)", #name, #name, args, #name, args, #name, #name

#api3 sparse_la_newSolver,          "Ptr (CSolverPtr a) -> IO CString"
#api3 sparse_la_freeSolver,         "CSolverPtr a -> IO CString"
#api3 sparse_la_factorize,          "CSolverPtr a -> CSparseMatrixPtr a -> IO CString"
#api3 sparse_la_analyzePattern,     "CSolverPtr a -> CSparseMatrixPtr a -> IO CString"
#api3 sparse_la_compute,            "CSolverPtr a -> CSparseMatrixPtr a -> IO CString"
#api3 sparse_la_tolerance,          "CSolverPtr a -> Ptr CDouble -> IO CString"
#api3 sparse_la_setTolerance,       "CSolverPtr a -> CDouble -> IO CString"
#api3 sparse_la_maxIterations,      "CSolverPtr a -> Ptr CInt -> IO CString"
#api3 sparse_la_setMaxIterations,   "CSolverPtr a -> CInt -> IO CString"
#api3 sparse_la_info,               "CSolverPtr a -> Ptr CInt -> IO CString"
#api3 sparse_la_error,              "CSolverPtr a -> Ptr CDouble -> IO CString"
#api3 sparse_la_iterations,         "CSolverPtr a -> Ptr CInt -> IO CString"
#api3 sparse_la_solve,              "CSolverPtr a -> CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
-- #api3 sparse_la_solveWithGuess,     "CSolverPtr a -> CSparseMatrixPtr a -> CSparseMatrixPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api3 sparse_la_matrixQ,            "CSolverPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api3 sparse_la_matrixR,            "CSolverPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api3 sparse_la_setPivotThreshold,  "CSolverPtr a -> CDouble -> IO CString"
#api3 sparse_la_rank,               "CSolverPtr a -> Ptr CInt -> IO CString"
#api3 sparse_la_matrixL,            "CSolverPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api3 sparse_la_matrixU,            "CSolverPtr a -> Ptr (CSparseMatrixPtr a) -> IO CString"
#api3 sparse_la_setSymmetric,       "CSolverPtr a -> CInt -> IO CString"
#api3 sparse_la_determinant,        "CSolverPtr a -> Ptr (C a) -> IO CString"
#api3 sparse_la_logAbsDeterminant,  "CSolverPtr a -> Ptr (C a) -> IO CString"
#api3 sparse_la_absDeterminant,     "CSolverPtr a -> Ptr (C a) -> IO CString"
#api3 sparse_la_signDeterminant,    "CSolverPtr a -> Ptr (C a) -> IO CString"

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


module Sondages where

import qualified Data.Text as T
import Data.Time (UTCTime)
import qualified Data.Map as M
import Data.Time.Segment
import Statistics.LinearRegression (linearRegression)
import qualified Data.Vector.Unboxed as U

data Candidat = Macron | LePen | Clinton | Trump
    deriving (Show, Read, Eq, Ord, Enum, Bounded)

data VoteIntention = VoteIntention { unCandidat :: Candidat
                                   , unIntention :: Double
             } deriving (Show, Read, Eq)

data Sondage = Sondage { unSondeur   :: T.Text
                       , unDateBegin :: UTCTime
                       , unDateEnd   :: Maybe UTCTime
                       , echantillon :: Maybe Int
                       , margeErreur :: Maybe Double
                       , abstention  :: Maybe Double
                       , indecis     :: Maybe Double
                       , intentions  :: [(Candidat, Double)]
                       } deriving (Show, Read, Eq)

sondages :: [Sondage]
sondages = [ Sondage "Harris"     (jour 2017 04 23) Nothing                  (Just 2684)   Nothing   Nothing     Nothing            [(Macron, 0.64)]
           , Sondage "Ipsos"      (jour 2017 04 23) Nothing                  (Just 1379)   Nothing   Nothing     Nothing            [(Macron, 0.62)]
           , Sondage "OpinionWay" (jour 2017 04 23) (Just $ jour 2017 04 24) (Just 1461)   Nothing   Nothing     Nothing     [(Macron, 0.61)]
           , Sondage "Ifop"       (jour 2017 04 23) (Just $ jour 2017 04 24) (Just 846)    Nothing  (Just 0.26) (Just 0.12)  [(Macron, 0.60)]
           , Sondage "Harris"     (jour 2017 04 25) (Just $ jour 2017 04 27) (Just 1016)   Nothing   Nothing     Nothing     [(Macron, 0.61)]
           , Sondage "OpinionWay" (jour 2017 04 25) (Just $ jour 2017 04 27) (Just 1500)   Nothing  (Just 0.25)  Nothing     [(Macron, 0.60)]
           , Sondage "Odoxa"      (jour 2017 04 26) (Just $ jour 2017 04 27) (Just 1003)   Nothing  (Just 0.20)   Nothing     [(Macron, 0.59)]
           , Sondage "BVA"        (jour 2017 04 26) (Just $ jour 2017 04 28) (Just 1506)   Nothing  (Just 0.245) (Just 0.15) [(Macron, 0.59)]
           , Sondage "Ipsos"      (jour 2017 04 28) (Just $ jour 2017 04 29) (Just 1504)   Nothing  (Just 0.25)   Nothing     [(Macron, 0.60)]
           , Sondage "Kantar"     (jour 2017 04 28) (Just $ jour 2017 04 30) (Just 1539)   Nothing   Nothing      Nothing     [(Macron, 0.59)]
           , Sondage "Elabe"      (jour 2017 04 28) (Just $ jour 2017 05 02) (Just 3956)   Nothing   Nothing     (Just 0.11)  [(Macron, 0.59)]
           , Sondage "Ifop"       (jour 2017 04 28) (Just $ jour 2017 05 02) (Just 1388)   Nothing  (Just 0.27)   Nothing     [(Macron, 0.595)]
           , Sondage "Cevifop"    (jour 2017 04 30) (Just $ jour 2017 05 01) (Just 13742)  Nothing  (Just 0.24)  (Just 0.15)  [(Macron, 0.59)]
           , Sondage "BVA"        (jour 2017 05 01) (Just $ jour 2017 05 02) (Just 1435)   Nothing  (Just 0.22)   Nothing     [(Macron, 0.60)]
           , Sondage "OpinionWay" (jour 2017 05 02) (Just $ jour 2017 05 04) (Just 1500)   Nothing  (Just 0.25)  (Just 0.15)  [(Macron, 0.62)]
           , Sondage "Elabe"      (jour 2017 05 04)  Nothing                 (Just 1009)   Nothing   Nothing      Nothing     [(Macron, 0.62)]
           , Sondage "Ipsos"      (jour 2017 05 04)  Nothing                 (Just 2632)   Nothing  (Just 0.24)   Nothing     [(Macron, 0.615)]
           , Sondage "Odoxa"      (jour 2017 05 04)  Nothing                 (Just 998)    Nothing  (Just 0.25)  (Just 0.21)  [(Macron, 0.62)]
           , Sondage "Harris"     (jour 2017 05 02) (Just $ jour 2017 05 05) (Just 4991)   Nothing  (Just 0.25)   Nothing     [(Macron, 0.62)]
           , Sondage "Ifop"       (jour 2017 05 02) (Just $ jour 2017 05 05) (Just 1861)   Nothing  (Just 0.245)  Nothing     [(Macron, 0.63)]
           , Sondage "Ipsos"      (jour 2017 05 05)  Nothing                 (Just 5331)   Nothing  (Just 0.24)  Nothing      [(Macron, 0.63)]
           ]

polls :: [Sondage]
polls = [ Sondage "YouGov" (jour 2016 10 22) (Just $ jour 2016 10 26) (Just 1376) (Just 0.031) Nothing Nothing [(Clinton, 0.49), (Trump, 0.46)]
        , Sondage "ABC"    (jour 2016 10 23) (Just $ jour 2016 10 26) (Just 1150) (Just 0.03)  Nothing Nothing [(Clinton, 0.50), (Trump, 0.45)]
        , Sondage "USC"    (jour 2016 10 21) (Just $ jour 2016 10 27) (Just 3248) (Just 0.045) Nothing Nothing [(Clinton, 0.44), (Trump, 0.46)]
        , Sondage "Ipsos"  (jour 2016 10 21) (Just $ jour 2016 10 27) (Just 1627) (Just 0.03)  Nothing Nothing [(Clinton, 0.42), (Trump, 0.36)]
        , Sondage "IBD"    (jour 2016 10 22) (Just $ jour 2016 10 27) (Just 973)  (Just 0.033) Nothing Nothing [(Clinton, 0.45), (Trump, 0.42)]
        , Sondage "ABC"    (jour 2016 10 24) (Just $ jour 2016 10 27) (Just 1148) (Just 0.03)  Nothing Nothing [(Clinton, 0.49), (Trump, 0.46)]
        , Sondage "IBD"    (jour 2016 10 23) (Just $ jour 2016 10 28) (Just 1013) (Just 0.033) Nothing Nothing [(Clinton, 0.46), (Trump, 0.41)]
        , Sondage "ABC"    (jour 2016 10 25) (Just $ jour 2016 10 28) (Just 1160) (Just 0.03)  Nothing Nothing [(Clinton, 0.46), (Trump, 0.45)]
        ]

pollsResult :: [Sondage]
pollsResult = [Sondage "REAL" (jour 2016 11 08) Nothing Nothing Nothing Nothing Nothing [(Clinton, 0.481), (Trump, 0.46)]]

normalizePolls :: [Sondage] -> [Sondage]
normalizePolls xs = Prelude.map normalizePolls' xs
    where
        normalizePolls' (Sondage a b c d e f g is) = Sondage a b c d e f g is'
            where
                is' = Prelude.map (\(m, n) -> (m, n/total)) is
                total = sum $ Prelude.map snd is

sondage2date :: [Sondage] -> [(UTCTime, [(Candidat, Double)])]
sondage2date xs = concat ( map (dureeSondage) xs )
    where
        dureeSondage (Sondage _ begin end _ _ _ _ is) = case end of
                    Nothing   -> zip (timesBetween begin begin D) (repeat is)
                    Just end' -> zip (timesBetween begin end'  D) (repeat is)

toMap :: forall k a. Ord k => [(k, [a])] -> M.Map k [a]
toMap = Prelude.foldr (\s -> M.insertWith (++) (fst s) (snd s)) M.empty

toMap' :: forall k a. Ord k => [(k, a)] -> M.Map k [a]
toMap' = Prelude.foldr (\s -> M.insertWith (++) (fst s) ([snd s])) M.empty

mapSondages :: [Sondage] -> M.Map UTCTime (M.Map Candidat Double)
mapSondages sondage = M.map (M.map mean) $ M.map toMap' $ toMap (sondage2date sondage)
    where
        mean :: [Double] -> Double
        mean xs = (sum xs) / (fromIntegral $ length xs)

predict :: (Double -> Double) -> Int -> [(UTCTime, Double)]
predict f days = predict' sondages f days

predict' :: [Sondage] -> (Double -> Double) -> Int -> [(UTCTime, Double)]
predict' sondes f days = zip dates $ Prelude.map (\x -> f $ a' + b' * x) [1.. (fromIntegral n)]
    where
        sondes' = M.toAscList (mapSondages sondes)
        firstDate = fst $ head sondes'
        n = length sondes' + days
        dates = timesAfter n D firstDate
        (a', b') = linearRegression (U.fromList a) (U.fromList b)
        (a, b) = unzip $ zip ([1..] :: [Double]) $  Prelude.map snd $ concat $ Prelude.map M.toList $ Prelude.map snd sondes'