{-# OPTIONS_GHC -fno-warn-orphans #-}
module Main where

import MNIST.Prelude
import MNIST.DataSet
import Numeric.LinearAlgebra.Static hiding (dot)
import Numeric.Backprop

import qualified Data.Vector.Generic             as VG
import qualified Data.Vector                     as V
import qualified Generics.SOP                    as SOP
import qualified Numeric.LinearAlgebra           as HM
import qualified System.Random.MWC               as MWC
import qualified System.Random.MWC.Distributions as MWC

-- | type-combinators alias for the terminal constructor
nil :: Prod f '[]
nil = Ø

-- ========================================================================= --
-- Internal structure of our FF neural net

data Layer i o = Layer
  { _lWeights :: !(L o i)
  , _lBiases  :: !(R o)
  } deriving (Show, Generic)

instance SOP.Generic (Layer i o)
instance NFData (Layer i o)

data Network i h1 h2 o = Net
  { _nLayer1 :: !(Layer i  h1)
  , _nLayer2 :: !(Layer h1 h2)
  , _nLayer3 :: !(Layer h2 o)
  } deriving (Show, Generic)

instance SOP.Generic (Network i h1 h2 o)
instance NFData (Network i h1 h2 o)

-- ========================================================================= --
-- type classes

type KnownNat4 i h1 h2 o = (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o)
type KnownNat2 i o = (KnownNat i, KnownNat o)

instance KnownNat2 i o => Num (Layer i o) where
  Layer w1 b1 + Layer w2 b2 = Layer (w1 + w2) (b1 + b2)
  Layer w1 b1 - Layer w2 b2 = Layer (w1 - w2) (b1 - b2)
  Layer w1 b1 * Layer w2 b2 = Layer (w1 * w2) (b1 * b2)
  abs    (Layer w b)        = Layer (abs    w) (abs    b)
  signum (Layer w b)        = Layer (signum w) (signum b)
  negate (Layer w b)        = Layer (negate w) (negate b)
  fromInteger x = Layer (fromInteger x) (fromInteger x)

instance KnownNat4 i h1 h2 o => Num (Network i h1 h2 o) where
  Net a b c + Net d e f = Net (a + d) (b + e) (c + f)
  Net a b c - Net d e f = Net (a - d) (b - e) (c - f)
  Net a b c * Net d e f = Net (a * d) (b * e) (c * f)
  abs    (Net a b c)    = Net (abs    a) (abs    b) (abs    c)
  signum (Net a b c)    = Net (signum a) (signum b) (signum c)
  negate (Net a b c)    = Net (negate a) (negate b) (negate c)
  fromInteger x         = Net (fromInteger x) (fromInteger x) (fromInteger x)

instance KnownNat2 i o => Fractional (Layer i o) where
  Layer w1 b1 / Layer w2 b2 = Layer (w1 / w2) (b1 / b2)
  recip (Layer w b)         = Layer (recip w) (recip b)
  fromRational x            = Layer (fromRational x) (fromRational x)

instance KnownNat4 i h1 h2 o => Fractional (Network i h1 h2 o) where
  Net a b c / Net d e f = Net (a / d) (b / e) (c / f)
  recip (Net a b c)     = Net (recip a) (recip b) (recip c)
  fromRational x        = Net (fromRational x) (fromRational x) (fromRational x)

instance KnownNat n => MWC.Variate (R n) where
  uniform g = randomVector <$> MWC.uniform g <*> pure Uniform
  uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance KnownNat2 m n => MWC.Variate (L m n) where
  uniform g = uniformSample <$> MWC.uniform g <*> pure 0 <*> pure 1
  uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance KnownNat2 i o => MWC.Variate (Layer i o) where
  uniform g = Layer <$> MWC.uniform g <*> MWC.uniform g
  uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance KnownNat4 i h1 h2 o => MWC.Variate (Network i h1 h2 o) where
  uniform g = Net <$> MWC.uniform g <*> MWC.uniform g <*> MWC.uniform g
  uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

-- ========================================================================= --
-- Basic math functions with back propagation

matVec :: forall m n . (KnownNat m, KnownNat n) => Op '[ L m n, R n ] (R m)
matVec = op2' $ \m v -> (forward m v, backward m v)
  where
    forward :: L m n -> R n -> R m
    forward m v = m #> v

    backward :: L m n -> R n -> Maybe (R m) -> (L m n, R n)
    backward m v (fromMaybe 1 -> g) = (g `outer` v, tr m #> g)

dot :: forall n . KnownNat n => Op '[ R n, R n ] Double
dot = op2' $ \x y -> (forward x y, backward x y)
  where
    forward :: R n -> R n -> Double
    forward x y = x <.> y

    backward :: R n -> R n -> Maybe Double -> (R n, R n)
    backward x y = \case
      Nothing -> (y, x)
      Just g  -> (konst g * y, x * konst g)

scale :: forall n . KnownNat n => Op '[ Double, R n ] (R n)
scale = op2' $ \a x -> (forward a x, backward a x)
  where
    forward :: Double -> R n -> R n
    forward a x = konst a * x

    backward :: Double -> R n -> Maybe (R n) -> (Double, R n)
    backward a x = \case
      Nothing -> (HM.sumElements (extract x      ), konst a    )
      Just g  -> (HM.sumElements (extract (x * g)), konst a * g)

vsum :: forall n . KnownNat n => Op '[ R n ] Double
vsum = op1' $ \x -> (forward x, backward)
  where
    forward :: R n -> Double
    forward = HM.sumElements . extract

    backward :: Maybe Double -> R n
    backward = maybe 1 konst

logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))

-- ========================================================================= --
-- run backpropagation
type LayerCtx i o = '[R i, Layer i o]
type NetCtx i h1 h2 o = '[ R i, Network i h1 h2 o ]

runLayer :: forall i o s . (KnownNat i, KnownNat o) => BPOp s '[ R i, Layer i o ] (R o)
runLayer = withInps $ \(decombinate -> (x, l)) -> do
  w :< b :< _ <- partsVar gTuple l
  y <- opVar matVec (w :< x :< nil)
  return $ y + b
  where
    decombinate
      :: Prod (BVar s (LayerCtx i o)) '[R i, Layer i o]
      -> (BVar s (LayerCtx i o) (R i), BVar s (LayerCtx i o) (Layer i o))
    decombinate (x :< l :< _) = (x, l)

runNetwork :: KnownNat4 i h1 h2 o => BPOp s (NetCtx i h1 h2 o) (R o)
runNetwork = withInps $ \(x :< n :< _) -> do
  l1 :< l2 :< l3 :< _ <- partsVar gTuple n

  y <- bindVar $ liftB (bpOp runLayer) (x :< l1 :< nil)
  -- or
  z <- bindVar $ (bpOp runLayer) .$ (logistic y :< l2 :< nil)
  -- or
  r <- (bpOp runLayer) ~$ (logistic z :< l3 :< nil)

  bpOp softmax ~$ only r
  where
    softmax :: KnownNat n => BPOp s '[ R n ] (R n)
    softmax = withInps $ \(x :< _) -> do
      expX <- bindVar (exp x)
      totX <- vsum ~$ (expX :< nil)
      scale        ~$ (1 / totX :< expX :< nil)

runNetOnInp :: KnownNat4 i h1 h2 o => Network i h1 h2 o -> R i -> R o
runNetOnInp n x = evalBPOp runNetwork (x ::< n ::< nil)

gradNet :: KnownNat4 i h1 h2 o => Network i h1 h2 o -> R i -> Network i h1 h2 o
gradNet n x = case gradBPOp runNetwork (x ::< n ::< nil) of
    _gradX ::< gradN ::< nil -> gradN

-- ========================================================================= --

-- crossEntropy :: forall s n . KnownNat n => R n -> BPOp s '[ R n ] Double
-- crossEntropy targ = withInps $ \(r :< _) ->
--   negate (dot ~$ (log r :< only t))
--   where
--     t :: BVar s '[R n] (R n)
--     t = constVar targ

crossEntropyI :: forall s n . KnownNat n => R n -> BPOpI s '[ R n ] Double
crossEntropyI targ (r :< _) = negate (dot .$ (log r :< only t))
  where
    t :: BVar s '[R n] (R n)
    t = constVar targ

softMaxCrossEntropy :: forall s n . KnownNat n => R n -> BPOp s '[ R n ] Double
softMaxCrossEntropy targ = withInps $ \(r :< Ø) -> do
  bindVar $ realToFrac tsum * log (vsum .$ (only r)) - (dot .$ (r :< t :< nil))
  where
    tsum :: Double
    tsum = HM.sumElements . extract $ targ

    t :: BVar s '[R n] (R n)
    t = constVar targ

softMaxCrossEntropyI :: forall s n . KnownNat n => R n -> BPOpI s '[ R n ] Double
softMaxCrossEntropyI targ (r :< Ø) =
  realToFrac tsum * log (vsum .$ (only r)) - (dot .$ (r :< t :< nil))
  where
    tsum :: Double
    tsum = HM.sumElements . extract $ targ

    t :: BVar s '[R n] (R n)
    t = constVar targ

trainStep
  :: forall i h1 h2 o. KnownNat4 i h1 h2 o
  => Double
  -> R i
  -> R o
  -> Network i h1 h2 o
  -> Network i h1 h2 o
trainStep r !x !t !n =
  case gradBPOp o (x ::< n ::< nil) of
    (_ :< I gN :< _) -> n - (realToFrac r * gN)
  where
    o :: BPOp s '[ R i, Network i h1 h2 o ] Double
    o = do
      y <- runNetwork
      implicitly (crossEntropyI t) -$ (y :< nil)

trainList
  :: KnownNat4 i h1 h2 o
  => Double
  -> [(R i, R o)]
  -> Network i h1 h2 o
  -> Network i h1 h2 o
trainList r = flip $ foldl' (\n (x,y) -> trainStep r x y n)

testNet
  :: forall i h1 h2 o. KnownNat4 i h1 h2 o
  => [(R i, R o)]
  -> Network i h1 h2 o
  -> Double
testNet xs n = sum (map (\(i,o) -> test i o) xs) / fromIntegral (length xs)
  where
    test :: R i -> R o -> Double
    test x (extract->t) = fromIntegral . fromEnum $
      HM.maxIndex t == HM.maxIndex (extract r)
      where
        r :: R o
        r = evalBPOp runNetwork (x ::< n ::< nil)

main :: IO ()
main = MWC.withSystemRandom $ \g -> do
  -- initialize data and network
  !trainingSet <- trainingDataBp
  !testSet     <- testDataBp
  !net0        <- MWC.uniformR @(Network 784 300 100 9) (-0.5, 0.5) g

  flip evalStateT net0 . forM_ [1..100] $ \e -> do
    trainingSet' <- liftIO . fmap V.toList $ MWC.uniformShuffle (V.fromList trainingSet) g
    liftIO $ printf "[Epoch %d]\n" (e :: Int)

    forM_ ([1..] `zip` chunksOf batch trainingSet') $ \(b, chnk) -> StateT $ \n0 -> do
      printf "(Batch %d)\n" (b :: Int)

      -- t0 <- getCurrentTime
      n' <- evaluate . force $ trainList rate chnk n0
      -- t1 <- getCurrentTime
      -- printf "Trained on %d points in %s.\n" batch (show (t1 `diffUTCTime` t0))

      let trainScore = testNet chnk    n'
          testScore  = testNet testSet n'
      printf "Training error:   %.2f%%\n" ((1 - trainScore) * 100)
      -- printf "Validation error: %.2f%%\n" ((1 - testScore ) * 100)

      return ((), n')
  where
    rate = 0.0001
    batch = 100

  --  go :: StateT (Network 784 300 100 9) IO ()
  --  go = do
  --    e <- get

--
--   -- Test
--   let testPreds = map (take 3 testSet) $ \(x, _) -> evalBPOp runNetwork (x ::< net0 ::< nil)
--
--   liftIO $ forM_ ([0..3] :: [Int]) $ \i -> do
--       -- T.putStrLn $ drawMNIST $ testImages !! i
--       putStrLn $ "\n" ++ "expected " ++ show (testLabels !! i)
--       putStrLn $         "     got " ++ show (testPreds !! i)
--
--


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}

{-# LANGUAGE TypeFamilies        #-}

--
-- Note: Input files can be downloaded at https://www.kaggle.com/scolianni/mnistasjpg
--
--
-- This is a simple generative adversarial network to make pictures
-- of numbers similar to those in MNIST.
--
-- It demonstrates a different usage of the library. Within about 15
-- minutes it was producing examples like this:
--
--               --.
--     .=-.--..#=###
--     -##==#########.
--     #############-
--   -###-.=..-.-==
--   ###-
--   .###-
--   .####...==-.
--   -####=--.=##=
--   -##=-     -##
--             =##
--           -##=
--           -###-
--         .####.
--         .#####.
-- ...---=#####-
-- .=#########.         .
--   .#######=.          .
--     . =-.
--
-- It's a 5!
--
import           Control.Applicative
import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.Random
import           Control.Monad.Trans.Except

import qualified Data.Attoparsec.Text         as A
import qualified Data.ByteString              as B
import           Data.List                    (foldl')
import           Data.Serialize
import qualified Data.Text                    as T
import qualified Data.Text.IO                 as T
import qualified Data.Vector.Storable         as V

import           Numeric.LinearAlgebra.Data   (toLists)
import qualified Numeric.LinearAlgebra.Static as SA

import           Options.Applicative

import           Grenade
import           Grenade.Utils.OneHot

type Discriminator =
  Network
    '[ Convolution 'WithoutBias 'NoPadding 1 10 5 5 1 1, Pooling 2 2 2 2, Relu
     , Convolution 'WithoutBias 'NoPadding 10 16 5 5 1 1, Pooling 2 2 2 2, Relu
     , Reshape, FullyConnected 256 80, Logit, FullyConnected 80 1, Logit]
    '[ 'D2 28 28
     , 'D3 24 24 10, 'D3 12 12 10, 'D3 12 12 10
     , 'D3 8 8 16, 'D3 4 4 16, 'D3 4 4 16
     , 'D1 256, 'D1 80, 'D1 80, 'D1 1, 'D1 1]

type Generator =
  Network
    '[ FullyConnected 80 256, Relu, Reshape
     , Deconvolution 16 10 5 5 2 2, Relu
     , Deconvolution 10 1 8 8 2 2, Logit]
    '[ 'D1 80
     , 'D1 256, 'D1 256, 'D3 4 4 16
     , 'D3 11 11 10, 'D3 11 11 10
     , 'D2 28 28, 'D2 28 28 ]

randomDiscriminator :: IO Discriminator
randomDiscriminator = randomNetwork

randomGenerator :: IO Generator
randomGenerator = randomNetwork

trainExample :: Optimizer opt -> Discriminator -> Generator -> S ('D2 28 28) -> S ('D1 80) -> ( Discriminator, Generator )
trainExample opt discriminator generator realExample noiseSource
 = let (generatorTape, fakeExample)       = runNetwork generator noiseSource

       (discriminatorTapeReal, guessReal) = runNetwork discriminator realExample
       (discriminatorTapeFake, guessFake) = runNetwork discriminator fakeExample

       (discriminator'real, _)            = runGradient discriminator discriminatorTapeReal ( guessReal - 1 )
       (discriminator'fake, _)            = runGradient discriminator discriminatorTapeFake guessFake
       (_, push)                          = runGradient discriminator discriminatorTapeFake ( guessFake - 1)

       (generator', _)                    = runGradient generator generatorTape push

       !newDiscriminator                   = force $ foldl' (applyUpdate $ sgdUpdateLearningParamters opt) discriminator [ discriminator'real, discriminator'fake ]
       !newGenerator                       = force $ applyUpdate opt generator generator'
   in ( newDiscriminator, newGenerator )
  where sgdUpdateLearningParamters :: Optimizer opt -> Optimizer opt
        sgdUpdateLearningParamters (OptSGD rate mom reg) = OptSGD rate mom (reg * 10)
        sgdUpdateLearningParamters o                     = o

ganTest :: (Discriminator, Generator) -> Int -> FilePath -> Optimizer opt -> ExceptT String IO (Discriminator, Generator)
ganTest (discriminator0, generator0) iterations trainFile opt = do
  !trainData      <- fmap fst <$> readMNIST trainFile

  lift $ foldM (runIteration trainData) ( discriminator0, generator0 ) [1..iterations]

    where

  showShape' :: S ('D2 a b) -> IO ()
  showShape' (S2D mm) = putStrLn $
    let m  = SA.extract mm
        ms = toLists m
        render n'  | n' <= 0.2  = ' '
                   | n' <= 0.4  = '.'
                   | n' <= 0.6  = '-'
                   | n' <= 0.8  = '='
                   | otherwise =  '#'

        px = (fmap . fmap) render ms
    in unlines px

  runIteration :: [S ('D2 28 28)] -> (Discriminator, Generator) -> Int -> IO (Discriminator, Generator)
  runIteration trainData ( !discriminator, !generator ) _ = do
    trained'    <- foldM ( \(!discriminatorX, !generatorX ) realExample -> do
                      trainExample opt discriminatorX generatorX realExample <$> randomOfShape
                     ) ( discriminator, generator ) trainData

    showShape' . snd . runNetwork (snd trained') =<< randomOfShape

    return trained'

data GanOpts = GanOpts FilePath Int Bool (Optimizer 'SGD) (Optimizer 'Adam) (Maybe FilePath) (Maybe FilePath)

mnist' :: Parser GanOpts
mnist' = GanOpts <$> argument str (metavar "TRAIN")
                 <*> option auto (long "iterations" <> short 'i' <> value 15)
                 <*> flag False True (long "use-adam" <> short 'a')
                 <*> (OptSGD
                       <$> option auto (long "train_rate" <> short 'r' <> value 0.01)
                       <*> option auto (long "momentum" <> value 0.9)
                       <*> option auto (long "l2" <> value 0.0005)
                       )
                 <*> (OptAdam
                       <$> option auto (long "alpha" <> short 'r' <> value 0.001)
                       <*> option auto (long "beta1" <> value 0.9)
                       <*> option auto (long "beta2" <> value 0.999)
                       <*> option auto (long "epsilon" <> value 1e-4)
                       <*> option auto (long "lambda" <> value 1e-3)
                      )
                 <*> optional (strOption (long "load"))
                 <*> optional (strOption (long "save"))

main :: IO ()
main = do
  GanOpts mnist iter useAdam sgd adam load save <- execParser (info (mnist' <**> helper) idm)
  putStrLn "Training stupidly simply GAN"
  nets0 <-
    case load of
      Just loadFile -> netLoad loadFile
      Nothing       -> (,) <$> randomDiscriminator <*> randomGenerator
  res <-
    if useAdam
      then runExceptT $ ganTest nets0 iter mnist adam
      else runExceptT $ ganTest nets0 iter mnist sgd
  case res of
    Right nets1 ->
      case save of
        Just saveFile -> B.writeFile saveFile $ runPut (put nets1)
        Nothing       -> return ()
    Left err -> putStrLn err

readMNIST :: FilePath -> ExceptT String IO [(S ('D2 28 28), S ('D1 10))]
readMNIST mnist = ExceptT $ do
  mnistdata <- T.readFile mnist
  return $ traverse (A.parseOnly parseMNIST) (tail $ T.lines mnistdata)

parseMNIST :: A.Parser (S ('D2 28 28), S ('D1 10))
parseMNIST = do
  Just lab <- oneHot <$> A.decimal
  pixels   <- many (A.char ',' >> A.double)
  image    <- maybe (fail "Parsed row was of an incorrect size") pure (fromStorable . V.fromList $ fmap realToFrac pixels)
  return (image, lab)

netLoad :: FilePath -> IO (Discriminator, Generator)
netLoad modelPath = do
  modelData <- B.readFile modelPath
  either fail return $ runGet (get :: Get (Discriminator, Generator)) modelData


{-# LANGUAGE FlexibleInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Data.AEq
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- A type class for approximate and exact equalilty comparisons and instances
-- for common data types.
module Data.AEq (
    AEq(..),
    ) where

import Foreign
import Foreign.C.Types
import Data.Complex
import Numeric.IEEE

infix 4 ===, ~==

-- | Types with approximate and exact equality comparisons.
class Eq a => AEq a where
    -- | An exact equality comparison.
    --
    -- For real 'IEEE' types, two values are equivalent in the
    -- following cases:
    --
    --   * both values are @+0@;
    --
    --   * both values are @-0@;
    --
    --   * both values are nonzero and equal to each other
    --     (according to '==');
    --
    --   * both values are @NaN@ with the same payload and sign.
    --
    -- For complex 'IEEE' types, two values are equivalent if their
    -- real and imaginary parts are equivalent.
    --
    (===) :: a -> a -> Bool
    (===) = (==)
    {-# INLINE (===) #-}

    -- | An approximate equality comparison operator.
    --
    -- For real 'IEEE' types, two values are approximately equal in the
    -- following cases:
    --
    --   * at least half of their significand bits agree;
    --
    --   * both values are less than 'epsilon';
    --
    --   * both values are @NaN@.
    --
    -- For complex 'IEEE' types, two values are approximately equal in the
    -- followiing cases:
    --
    --   * their magnitudes are approximately equal and the angle between
    --     them is less than @32*'epsilon'@;
    --
    --   * both magnitudes are less than 'epsilon';
    --
    --   * both have a @NaN@ real or imaginary part.
    --
    -- Admitedly, the @32@ is a bit of a hack.  Future versions of the
    -- library may switch to a more principled test of the angle.
    --
    (~==) :: a -> a -> Bool
    (~==) = (==)
    {-# INLINE (~==) #-}

approxEqIEEE :: (IEEE a) => a -> a -> Bool
approxEqIEEE x y =
    ( sameSignificandBits x y >= d
    || (abs x < epsilon && abs y < epsilon)
    || (isNaN x && isNaN y)
    )
  where
    d = (floatDigits x + 1) `div` 2
{-# INLINE approxEqIEEE #-}

identicalComplexIEEE :: (IEEE a) => Complex a -> Complex a -> Bool
identicalComplexIEEE (x1 :+ y1) (x2 :+ y2) =
    (identicalIEEE x1 x2) && (identicalIEEE y1 y2)
{-# INLINE identicalComplexIEEE #-}

approxEqComplexIEEE :: (IEEE a) => Complex a -> Complex a -> Bool
approxEqComplexIEEE z1 z2 = let
    (r1,c1) = polar z1
    (r2,c2) = polar z2
    angle = abs (c1 - c2)
    in ( ( approxEqIEEE r1 r2
         && (angle < 32*epsilon || angle > 2*(pi - 16*epsilon) || isNaN angle)
         )
       || (r1 < epsilon && r2 < epsilon)
       )
{-# INLINE approxEqComplexIEEE #-}

instance AEq Float where
    (===) = identicalIEEE
    {-# INLINE (===) #-}
    (~==) = approxEqIEEE
    {-# INLINE (~==) #-}

instance AEq Double where
    (===) = identicalIEEE
    {-# INLINE (===) #-}
    (~==) = approxEqIEEE
    {-# INLINE (~==) #-}

instance AEq (Complex Float) where
    (===) = identicalComplexIEEE
    {-# INLINE (===) #-}
    (~==) = approxEqComplexIEEE
    {-# INLINE (~==) #-}

instance AEq (Complex Double) where
    (===) = identicalComplexIEEE
    {-# INLINE (===) #-}
    (~==) = approxEqComplexIEEE
    {-# INLINE (~==) #-}

instance AEq CFloat where
    (===) = identicalIEEE
    {-# INLINE (===) #-}
    (~==) = approxEqIEEE
    {-# INLINE (~==) #-}

instance AEq CDouble where
    (===) = identicalIEEE
    {-# INLINE (===) #-}
    (~==) = approxEqIEEE
    {-# INLINE (~==) #-}

instance AEq (Complex CFloat) where
    (===) = identicalComplexIEEE
    {-# INLINE (===) #-}
    (~==) = approxEqComplexIEEE
    {-# INLINE (~==) #-}

instance AEq (Complex CDouble) where
    (===) = identicalComplexIEEE
    {-# INLINE (===) #-}
    (~==) = approxEqComplexIEEE
    {-# INLINE (~==) #-}

instance AEq Bool
instance AEq Char
instance AEq Int
instance AEq Int8
instance AEq Int16
instance AEq Int32
instance AEq Int64
instance AEq Integer
instance AEq Ordering
instance AEq Word
instance AEq Word8
instance AEq Word16
instance AEq Word32
instance AEq Word64
instance AEq ()
instance AEq WordPtr
instance AEq IntPtr
instance AEq (StablePtr a)
instance AEq (Ptr a)
instance AEq (FunPtr a)
instance AEq (ForeignPtr a)
instance AEq CChar
instance AEq CSChar
instance AEq CUChar
instance AEq CShort
instance AEq CUShort
instance AEq CInt
instance AEq CUInt
instance AEq CLong
instance AEq CULong
instance AEq CPtrdiff
instance AEq CSize
instance AEq CWchar
instance AEq CSigAtomic
instance AEq CLLong
instance AEq CULLong
instance AEq CIntPtr
instance AEq CUIntPtr
instance AEq CIntMax
instance AEq CUIntMax
instance AEq CClock
instance AEq CTime

eqListsWith :: (a -> a -> Bool) -> [a] -> [a] -> Bool
eqListsWith f (x:xs) (y:ys) = f x y && eqListsWith f xs ys
eqListsWith _ [] [] = True
eqListsWith _ _  _  = False
{-# INLINE eqListsWith #-}

instance (AEq a) => AEq [a] where
    (===) = eqListsWith (===)
    {-# INLINE (===) #-}
    (~==) = eqListsWith (~==)
    {-# INLINE (~==) #-}

instance (AEq a) => AEq (Maybe a) where
    (===) Nothing  Nothing  = True
    (===) (Just x) (Just y) = (===) x y
    (===) _ _ = False
    {-# INLINE (===) #-}

    (~==) Nothing  Nothing  = True
    (~==) (Just x) (Just y) = (~==) x y
    (~==) _ _ = False
    {-# INLINE (~==) #-}

instance (AEq a, AEq b) => AEq (Either a b) where
    (===) (Left a1)  (Left a2)  = (===) a1 a2
    (===) (Right b1) (Right b2) = (===) b1 b2
    (===) _ _ = False
    {-# INLINE (===) #-}

    (~==) (Left a1)  (Left a2)  = (~==) a1 a2
    (~==) (Right b1) (Right b2) = (~==) b1 b2
    (~==) _ _ = False
    {-# INLINE (~==) #-}

instance (AEq a, AEq b) => AEq (a,b) where
    (===) (a1,b1) (a2,b2) =
        (  ((===) a1 a2)
        && ((===) b1 b2)
        )
    {-# INLINE (===) #-}

    (~==) (a1,b1) (a2,b2) =
        (  ((~==) a1 a2)
        && ((~==) b1 b2)
        )
    {-# INLINE (~==) #-}

instance (AEq a, AEq b, AEq c) => AEq (a,b,c) where
    (===) (a1,b1,c1) (a2,b2,c2) =
        (  ((===) a1 a2)
        && ((===) b1 b2)
        && ((===) c1 c2)
        )
    {-# INLINE (===) #-}

    (~==) (a1,b1,c1) (a2,b2,c2) =
        (  ((~==) a1 a2)
        && ((~==) b1 b2)
        && ((~==) c1 c2)
        )
    {-# INLINE (~==) #-}

instance (AEq a, AEq b, AEq c, AEq d) => AEq (a,b,c,d) where
    (===) (a1,b1,c1,d1) (a2,b2,c2,d2) =
        (  ((===) a1 a2)
        && ((===) b1 b2)
        && ((===) c1 c2)
        && ((===) d1 d2)
        )
    {-# INLINE (===) #-}

    (~==) (a1,b1,c1,d1) (a2,b2,c2,d2) =
        (  ((~==) a1 a2)
        && ((~==) b1 b2)
        && ((~==) c1 c2)
        && ((~==) d1 d2)
        )
    {-# INLINE (~==) #-}

instance (AEq a, AEq b, AEq c, AEq d, AEq e) => AEq (a,b,c,d,e) where
    (===) (a1,b1,c1,d1,e1) (a2,b2,c2,d2,e2) =
        (  ((===) a1 a2)
        && ((===) b1 b2)
        && ((===) c1 c2)
        && ((===) d1 d2)
        && ((===) e1 e2)
        )
    {-# INLINE (===) #-}

    (~==) (a1,b1,c1,d1,e1) (a2,b2,c2,d2,e2) =
        (  ((~==) a1 a2)
        && ((~==) b1 b2)
        && ((~==) c1 c2)
        && ((~==) d1 d2)
        && ((~==) e1 e2)
        )
    {-# INLINE (~==) #-}

instance (AEq a, AEq b, AEq c, AEq d, AEq e, AEq f) => AEq (a,b,c,d,e,f) where
    (===) (a1,b1,c1,d1,e1,f1) (a2,b2,c2,d2,e2,f2) =
        (  ((===) a1 a2)
        && ((===) b1 b2)
        && ((===) c1 c2)
        && ((===) d1 d2)
        && ((===) e1 e2)
        && ((===) f1 f2)
        )
    {-# INLINE (===) #-}

    (~==) (a1,b1,c1,d1,e1,f1) (a2,b2,c2,d2,e2,f2) =
        (  ((~==) a1 a2)
        && ((~==) b1 b2)
        && ((~==) c1 c2)
        && ((~==) d1 d2)
        && ((~==) e1 e2)
        && ((~==) f1 f2)
        )
    {-# INLINE (~==) #-}

instance (AEq a, AEq b, AEq c, AEq d, AEq e, AEq f, AEq g) => AEq (a,b,c,d,e,f,g) where
    (===) (a1,b1,c1,d1,e1,f1,g1) (a2,b2,c2,d2,e2,f2,g2) =
        (  ((===) a1 a2)
        && ((===) b1 b2)
        && ((===) c1 c2)
        && ((===) d1 d2)
        && ((===) e1 e2)
        && ((===) f1 f2)
        && ((===) g1 g2)
        )
    {-# INLINE (===) #-}

    (~==) (a1,b1,c1,d1,e1,f1,g1) (a2,b2,c2,d2,e2,f2,g2) =
        (  ((~==) a1 a2)
        && ((~==) b1 b2)
        && ((~==) c1 c2)
        && ((~==) d1 d2)
        && ((~==) e1 e2)
        && ((~==) f1 f2)
        && ((~==) g1 g2)
        )
    {-# INLINE (~==) #-}

instance (AEq a, AEq b, AEq c, AEq d, AEq e, AEq f, AEq g, AEq h) => AEq (a,b,c,d,e,f,g,h) where
    (===) (a1,b1,c1,d1,e1,f1,g1,h1) (a2,b2,c2,d2,e2,f2,g2,h2) =
        (  ((===) a1 a2)
        && ((===) b1 b2)
        && ((===) c1 c2)
        && ((===) d1 d2)
        && ((===) e1 e2)
        && ((===) f1 f2)
        && ((===) g1 g2)
        && ((===) h1 h2)
        )
    {-# INLINE (===) #-}

    (~==) (a1,b1,c1,d1,e1,f1,g1,h1) (a2,b2,c2,d2,e2,f2,g2,h2) =
        (  ((~==) a1 a2)
        && ((~==) b1 b2)
        && ((~==) c1 c2)
        && ((~==) d1 d2)
        && ((~==) e1 e2)
        && ((~==) f1 f2)
        && ((~==) g1 g2)
        && ((~==) h1 h2)
        )
    {-# INLINE (~==) #-}

instance (AEq a, AEq b, AEq c, AEq d, AEq e, AEq f, AEq g, AEq h, AEq i) => AEq (a,b,c,d,e,f,g,h,i) where
    (===) (a1,b1,c1,d1,e1,f1,g1,h1,i1) (a2,b2,c2,d2,e2,f2,g2,h2,i2) =
        (  ((===) a1 a2)
        && ((===) b1 b2)
        && ((===) c1 c2)
        && ((===) d1 d2)
        && ((===) e1 e2)
        && ((===) f1 f2)
        && ((===) g1 g2)
        && ((===) h1 h2)
        && ((===) i1 i2)
        )
    {-# INLINE (===) #-}

    (~==) (a1,b1,c1,d1,e1,f1,g1,h1,i1) (a2,b2,c2,d2,e2,f2,g2,h2,i2) =
        (  ((~==) a1 a2)
        && ((~==) b1 b2)
        && ((~==) c1 c2)
        && ((~==) d1 d2)
        && ((~==) e1 e2)
        && ((~==) f1 f2)
        && ((~==) g1 g2)
        && ((~==) h1 h2)
        && ((~==) i1 i2)
        )
    {-# INLINE (~==) #-}

instance (AEq a, AEq b, AEq c, AEq d, AEq e, AEq f, AEq g, AEq h, AEq i, AEq j) => AEq (a,b,c,d,e,f,g,h,i,j) where
    (===) (a1,b1,c1,d1,e1,f1,g1,h1,i1,j1) (a2,b2,c2,d2,e2,f2,g2,h2,i2,j2) =
        (  ((===) a1 a2)
        && ((===) b1 b2)
        && ((===) c1 c2)
        && ((===) d1 d2)
        && ((===) e1 e2)
        && ((===) f1 f2)
        && ((===) g1 g2)
        && ((===) h1 h2)
        && ((===) i1 i2)
        && ((===) j1 j2)
        )
    {-# INLINE (===) #-}

    (~==) (a1,b1,c1,d1,e1,f1,g1,h1,i1,j1) (a2,b2,c2,d2,e2,f2,g2,h2,i2,j2) =
        (  ((~==) a1 a2)
        && ((~==) b1 b2)
        && ((~==) c1 c2)
        && ((~==) d1 d2)
        && ((~==) e1 e2)
        && ((~==) f1 f2)
        && ((~==) g1 g2)
        && ((~==) h1 h2)
        && ((~==) i1 i2)
        && ((~==) j1 j2)
        )
    {-# INLINE (~==) #-}

instance (AEq a, AEq b, AEq c, AEq d, AEq e, AEq f, AEq g, AEq h, AEq i, AEq j, AEq k) => AEq (a,b,c,d,e,f,g,h,i,j,k) where
    (===) (a1,b1,c1,d1,e1,f1,g1,h1,i1,j1,k1) (a2,b2,c2,d2,e2,f2,g2,h2,i2,j2,k2) =
        (  ((===) a1 a2)
        && ((===) b1 b2)
        && ((===) c1 c2)
        && ((===) d1 d2)
        && ((===) e1 e2)
        && ((===) f1 f2)
        && ((===) g1 g2)
        && ((===) h1 h2)
        && ((===) i1 i2)
        && ((===) j1 j2)
        && ((===) k1 k2)
        )
    {-# INLINE (===) #-}

    (~==) (a1,b1,c1,d1,e1,f1,g1,h1,i1,j1,k1) (a2,b2,c2,d2,e2,f2,g2,h2,i2,j2,k2) =
        (  ((~==) a1 a2)
        && ((~==) b1 b2)
        && ((~==) c1 c2)
        && ((~==) d1 d2)
        && ((~==) e1 e2)
        && ((~==) f1 f2)
        && ((~==) g1 g2)
        && ((~==) h1 h2)
        && ((~==) i1 i2)
        && ((~==) j1 j2)
        && ((~==) k1 k2)
        )
    {-# INLINE (~==) #-}

instance (AEq a, AEq b, AEq c, AEq d, AEq e, AEq f, AEq g, AEq h, AEq i, AEq j, AEq k, AEq l) => AEq (a,b,c,d,e,f,g,h,i,j,k,l) where
    (===) (a1,b1,c1,d1,e1,f1,g1,h1,i1,j1,k1,l1) (a2,b2,c2,d2,e2,f2,g2,h2,i2,j2,k2,l2) =
        (  ((===) a1 a2)
        && ((===) b1 b2)
        && ((===) c1 c2)
        && ((===) d1 d2)
        && ((===) e1 e2)
        && ((===) f1 f2)
        && ((===) g1 g2)
        && ((===) h1 h2)
        && ((===) i1 i2)
        && ((===) j1 j2)
        && ((===) k1 k2)
        && ((===) l1 l2)
        )
    {-# INLINE (===) #-}

    (~==) (a1,b1,c1,d1,e1,f1,g1,h1,i1,j1,k1,l1) (a2,b2,c2,d2,e2,f2,g2,h2,i2,j2,k2,l2) =
        (  ((~==) a1 a2)
        && ((~==) b1 b2)
        && ((~==) c1 c2)
        && ((~==) d1 d2)
        && ((~==) e1 e2)
        && ((~==) f1 f2)
        && ((~==) g1 g2)
        && ((~==) h1 h2)
        && ((~==) i1 i2)
        && ((~==) j1 j2)
        && ((~==) k1 k2)
        && ((~==) l1 l2)
        )
    {-# INLINE (~==) #-}

instance (AEq a, AEq b, AEq c, AEq d, AEq e, AEq f, AEq g, AEq h, AEq i, AEq j, AEq k, AEq l, AEq m) => AEq (a,b,c,d,e,f,g,h,i,j,k,l,m) where
    (===) (a1,b1,c1,d1,e1,f1,g1,h1,i1,j1,k1,l1,m1) (a2,b2,c2,d2,e2,f2,g2,h2,i2,j2,k2,l2,m2) =
        (  ((===) a1 a2)
        && ((===) b1 b2)
        && ((===) c1 c2)
        && ((===) d1 d2)
        && ((===) e1 e2)
        && ((===) f1 f2)
        && ((===) g1 g2)
        && ((===) h1 h2)
        && ((===) i1 i2)
        && ((===) j1 j2)
        && ((===) k1 k2)
        && ((===) l1 l2)
        && ((===) m1 m2)
        )
    {-# INLINE (===) #-}

    (~==) (a1,b1,c1,d1,e1,f1,g1,h1,i1,j1,k1,l1,m1) (a2,b2,c2,d2,e2,f2,g2,h2,i2,j2,k2,l2,m2) =
        (  ((~==) a1 a2)
        && ((~==) b1 b2)
        && ((~==) c1 c2)
        && ((~==) d1 d2)
        && ((~==) e1 e2)
        && ((~==) f1 f2)
        && ((~==) g1 g2)
        && ((~==) h1 h2)
        && ((~==) i1 i2)
        && ((~==) j1 j2)
        && ((~==) k1 k2)
        && ((~==) l1 l2)
        && ((~==) m1 m2)
        )
    {-# INLINE (~==) #-}

instance (AEq a, AEq b, AEq c, AEq d, AEq e, AEq f, AEq g, AEq h, AEq i, AEq j, AEq k, AEq l, AEq m, AEq n) => AEq (a,b,c,d,e,f,g,h,i,j,k,l,m,n) where
    (===) (a1,b1,c1,d1,e1,f1,g1,h1,i1,j1,k1,l1,m1,n1) (a2,b2,c2,d2,e2,f2,g2,h2,i2,j2,k2,l2,m2,n2) =
        (  ((===) a1 a2)
        && ((===) b1 b2)
        && ((===) c1 c2)
        && ((===) d1 d2)
        && ((===) e1 e2)
        && ((===) f1 f2)
        && ((===) g1 g2)
        && ((===) h1 h2)
        && ((===) i1 i2)
        && ((===) j1 j2)
        && ((===) k1 k2)
        && ((===) l1 l2)
        && ((===) m1 m2)
        && ((===) n1 n2)
        )
    {-# INLINE (===) #-}

    (~==) (a1,b1,c1,d1,e1,f1,g1,h1,i1,j1,k1,l1,m1,n1) (a2,b2,c2,d2,e2,f2,g2,h2,i2,j2,k2,l2,m2,n2) =
        (  ((~==) a1 a2)
        && ((~==) b1 b2)
        && ((~==) c1 c2)
        && ((~==) d1 d2)
        && ((~==) e1 e2)
        && ((~==) f1 f2)
        && ((~==) g1 g2)
        && ((~==) h1 h2)
        && ((~==) i1 i2)
        && ((~==) j1 j2)
        && ((~==) k1 k2)
        && ((~==) l1 l2)
        && ((~==) m1 m2)
        && ((~==) n1 n2)
        )
    {-# INLINE (~==) #-}

instance (AEq a, AEq b, AEq c, AEq d, AEq e, AEq f, AEq g, AEq h, AEq i, AEq j, AEq k, AEq l, AEq m, AEq n, AEq o) => AEq (a,b,c,d,e,f,g,h,i,j,k,l,m,n,o) where
    (===) (a1,b1,c1,d1,e1,f1,g1,h1,i1,j1,k1,l1,m1,n1,o1) (a2,b2,c2,d2,e2,f2,g2,h2,i2,j2,k2,l2,m2,n2,o2) =
        (  ((===) a1 a2)
        && ((===) b1 b2)
        && ((===) c1 c2)
        && ((===) d1 d2)
        && ((===) e1 e2)
        && ((===) f1 f2)
        && ((===) g1 g2)
        && ((===) h1 h2)
        && ((===) i1 i2)
        && ((===) j1 j2)
        && ((===) k1 k2)
        && ((===) l1 l2)
        && ((===) m1 m2)
        && ((===) n1 n2)
        && ((===) o1 o2)
        )
    {-# INLINE (===) #-}

    (~==) (a1,b1,c1,d1,e1,f1,g1,h1,i1,j1,k1,l1,m1,n1,o1) (a2,b2,c2,d2,e2,f2,g2,h2,i2,j2,k2,l2,m2,n2,o2) =
        (  ((~==) a1 a2)
        && ((~==) b1 b2)
        && ((~==) c1 c2)
        && ((~==) d1 d2)
        && ((~==) e1 e2)
        && ((~==) f1 f2)
        && ((~==) g1 g2)
        && ((~==) h1 h2)
        && ((~==) i1 i2)
        && ((~==) j1 j2)
        && ((~==) k1 k2)
        && ((~==) l1 l2)
        && ((~==) m1 m2)
        && ((~==) n1 n2)
        && ((~==) o1 o2)
        )
    {-# INLINE (~==) #-}


{-# LANGUAGE GADTs, NoImplicitPrelude, TypeFamilies,
InstanceSigs, StandaloneDeriving, UndecidableInstances,
ScopedTypeVariables, FlexibleInstances, DataKinds,
FunctionalDependencies, PolyKinds,
TypeOperators, RankNTypes, ImpredicativeTypes, MultiParamTypeClasses,
AllowAmbiguousTypes, TypeApplications, FlexibleContexts #-}
---
module Main where

import qualified Data.Map.Strict as Map
import Prelude hiding ((.), id)
import Linear.Vector
import Data.Tuple
import Data.Complex
import Linear.Epsilon (nearZero)
import Control.Category
import Control.Monad ((<=<))
import GHC.TypeNats
import Data.Proxy
import Diag (mymain)
-- import qualified Text.PrettyPrint.ANSI.Leijen as PP
import qualified Text.PrettyPrint.Boxes as B
-- import Control.Arrow
{-
type MapVec r k = Map.Map k r

vadd :: (Ord k, Num r) => MapVec r k -> MapVec r k -> MapVec r k
vadd = Map.unionWith (+)

smul :: (Ord k, Num r) => k -> MapVec r k -> MapVec r k
smul

-- kmett alread did this stuff
-}

-- http://blog.sigfpe.com/2007/03/monads-vector-spaces-and-quantum.html

data W b a = W { runW :: [(a,b)] } deriving (Eq,Show,Ord)
instance Semigroup (W b a) where
  (W x) <> (W y) = W (x <> y)
instance Monoid (W b a) where
  mempty = W mempty
mapW f (W l) = W $ map (\(a,b) -> (a,f b)) l
instance Functor (W b) where
    fmap f (W a) = W $ map (\(a,p) -> (f a,p)) a

instance Num b => Applicative (W b) where
  pure x = W [(x,1)]
  (W fs) <*> (W xs) = W [(f x, a * b) | (f, a) <- fs, (x, b) <- xs]
  -- the fs is a a vector of rows kind of.
    -- error "I don't feel like it"

instance Num b => Monad (W b) where
   return x = W [(x,1)]
   l >>= f = W $ concatMap (\(W d,p) -> map (\(x,q)->(x,p*q)) d) (runW $ fmap f l)

a .* b = mapW (a*) b

instance (Eq a,Show a,Num b) => Num (W b a) where
     W a + W b = W $ (a ++ b)
     a - b = a + (-1) .* b
     _ * _ = error "Num is annoying"
     abs _ = error "Num is annoying"
     signum _ = error "Num is annoying"
     fromInteger a = if a==0 then W [] else error "fromInteger can only take zero argument"

collect :: (Ord a, Num b) => W b a -> W b a
collect = W . Map.toList . Map.fromListWith (+) . runW

trimZero = W . filter (\(k,v) -> not $ nearZero v) . runW
simplify :: Ord a => Q a -> Q a
simplify = trimZero . collect
-- filter (not . nearZero . snd)

type P a = W Double a

type Q a = W (Complex Double) a

v1 :: Map.Map Bool Double
v1 = zero

boolbase :: [Map.Map Bool Double]
boolbase = basis

e1 = Map.singleton True 1.0
e2 = Map.singleton False 1.0

star = mapW conjugate

{-
There are two types of particles

-}
data FibAnyon = Id | Tau
-- I started using DataKinds and it ended up being a problem.

data A a

data Tau
data Id
data FibTree root leaves where
   TTT :: FibTree Tau l -> FibTree Tau r -> FibTree Tau (l,r)
   ITT :: FibTree Tau l -> FibTree Tau r -> FibTree Id (l,r) -- we would like to enforce that isingany can split into a b, but thisis tough
   TIT :: FibTree Id l -> FibTree Tau r -> FibTree Tau (l,r)
   TTI :: FibTree Tau l -> FibTree Id r -> FibTree Tau (l,r)
   III :: FibTree Id l -> FibTree Id r -> FibTree Id (l,r)
   -- III :: FibTree 'Id l -> FibTree 'Id l' -> FibTree 'Id (l,l') -- maybe shouldn't exist. uneccessary. Does really help for monoidal instance
   TLeaf :: FibTree Tau Tau
   ILeaf :: FibTree Id Id

-- pretty printing would be hella nice
--deriving instance Show (FibTree a b)

instance Show (FibTree a b) where
  show = drawTree
deriving instance Eq (FibTree a b)
--deriving instance Ord (FibTree a b)
{-
r = PP.text "\\"
l = PP.text "/"
ls = PP.group $ (PP.text "  /" <> PP.line <> PP.text " /" <> PP.line <>PP.text "/")
rs = PP.group $ (PP.text "\\" <> PP.line <> PP.text " \\" <> PP.line <>PP.text "  \\")
-}

sp = B.char ' '
l = B.char '/'
r = B.char '\\'
h = B.char '-'
v = B.char '|'
t = B.char 'T'
i = B.char 'I'
rs x =  (r  B.<> sp B.<> sp) B.//
      (sp B.<> r  B.<> x) B.//
      (sp B.<> sp  B.<> r)
ls x =  (sp  B.<> sp B.<> l) B.//
      (x B.<> l  B.<> sp) B.//
      (l B.<> sp  B.<> sp)

{-
drawty = v B.// v B.// v B.// (h B.<> h B.<> h B.<> h)
drawb :: FibTree a b -> B.Box
drawb ILeaf = i
drawb TLeaf = t
drawb (TTT l r) = (drawb l) B.vcat   B.// ((drawb r)  B.<> v )
-}
-- PP.nest 2 l <> PP.line <> PP.nest 1 l <> PP.line <> l
-- text "Tau \\"
-- text "Id \\"

-- From Data.Tree

draw :: FibTree a b -> [String]
draw (ITT l r) = "I" : drawSubTrees (l,r)
draw (TIT l r) = "T" : drawSubTrees (l,r)
draw (TTT l r) = "T" : drawSubTrees (l,r)
draw (TTI l r) = "T" : drawSubTrees (l,r)
draw (III l r) = "I" : drawSubTrees (l,r)
draw  ILeaf    = "I" : []
draw  TLeaf    = "T" : []

-- drawSubTrees' [] = []
drawSubTrees' t =
    "|" : shift "`- " "   " (draw t)
drawSubTrees (t,ts) =
    "|" : shift "+- " "|  " (draw t) ++ drawSubTrees' ts

shift first other = zipWith (++) (first : repeat other)

{-
draw' ILeaf = ["I"]
draw' TLeaf = ["T"]
draw' (TTT l r) =  " -" ++ ((fmap . fmap) (\s -> s ++ "  |") (draw' l)) ++ (["-+"] ++ draw' r))
-}
-- drawex = drawb (TTT (TTT TLeaf TLeaf) TLeaf)

-- | Neat 2-dimensional drawing of a tree.
drawTree :: FibTree a b -> String
drawTree  = unlines . draw

-- | Neat 2-dimensional drawing of a forest.
-- drawForest :: Forest String -> String
-- drawForest  = unlines . map drawTree

instance Ord (FibTree a b) where
  compare (ITT l r) (ITT l' r') | l < l' = LT
                                | l > l' = GT
                                | otherwise = compare r r'
  compare (ITT _ _) _ = LT
  compare _ (ITT _ _) = GT
  compare (TTI l r) (TTI l' r') | l < l' = LT
                                | l > l' = GT
                                | otherwise = compare r r'
  compare (TIT l r) (TIT l' r') | l < l' = LT
                                | l > l' = GT
                                | otherwise = compare r r'
  compare (TTT l r) (TTT l' r') | l < l' = LT
                                | l > l' = GT
                                | otherwise = compare r r'
  compare (III l r) (III l' r') | l < l' = LT
                                | l > l' = GT
                                | otherwise = compare r r'

  compare (TTI _ _) _ = LT
  compare _ (TTI _ _) = GT
  compare (TIT _ _) _ = LT
  compare _ (TIT _ _) = GT
  compare (TTT _ _) _ = LT
  compare _ (TTT _ _) = GT
  compare (III _ _) _ = LT
  compare _ (III _ _) = GT
  compare TLeaf TLeaf = EQ
  compare ILeaf ILeaf = EQ
  -- whoa. GHC is smart enough to realize some of these patterns can't happen
  -- because of the types.

{-
instance Enum (FibTree 'Tau (A 'Tau)) where
  toEnum _ = TLeaf
  fromEnum _ = 0
-}
{-
instance Enum (FibTree 'Id 'Id) where
  toEnum _ = TLeaf
  fromEnum _ = 0
  -}
-- and no inhabitants types?
--instance Enum (FibTree d b), Enum (FibTree e c) => Enum (FibTree a (b,c))
{-
instance (Enum (FibTree 'Tau b),
  Enum (FibTree 'Tau c),
  Enum (FibTree ' b),
  Enum (FibTree 'Id c)) => Enum (FibTree 'Tau (b,c))
  toEnum 0 = TTI
  toEnum
  fromEnum
-}

--instance Enum (FibTree a (b,c)) Enum (FibTree a (b,c))=> Enum (FibTree 'Id (b,c))

-- (FibTree a (b,c))

type TreeFun a b c d = FibTree a b -> FibTree c d
{-
lmap :: (forall a. FibTree a b -> FibTree (a) c) -> (forall e. FibTree e (b,d) -> FibTree e (c,d)) --  forall (a :: FibAnyon).
lmap f (ITT l r) = (ITT (f l) r)
lmap f (TTI l r) = (TTI (f l) r)
lmap f (TIT l r) = (TIT (f l) r)
lmap f (TTT l r) = (TTT (f l) r)
-}
lmap :: (forall a. FibTree a b -> Q (FibTree a c)) -> (FibTree e (b,d) -> Q (FibTree e (c,d)))
lmap f (ITT l r) = fmap (\l' -> ITT l' r) (f l)
lmap f (TTI l r) = fmap (\l' -> TTI l' r) (f l)
lmap f (TIT l r) = fmap (\l' -> TIT l' r) (f l)
lmap f (TTT l r) = fmap (\l' -> TTT l' r) (f l)
lmap f (III l r) = fmap (\l' -> III l' r) (f l)
{-
rmap :: (forall a. FibTree a b -> FibTree a c) -> (forall e. FibTree e (d,b) -> FibTree e (d,c)) --  forall (a :: FibAnyon).
rmap f (ITT l r) = ITT l (f r)
rmap f (TTI l r) = TTI l (f r)
rmap f (TIT l r) = TIT l (f r)
rmap f (TTT l r) = TTT l (f r)
-}
rmap :: (forall a. FibTree a b -> Q (FibTree a c)) -> (FibTree e (d,b) -> Q (FibTree e (d,c)))
rmap f (ITT l r) = fmap (\r' -> ITT l r') (f r)
rmap f (TTI l r) = fmap (\r' -> TTI l r') (f r)
rmap f (TIT l r) = fmap (\r' -> TIT l r') (f r)
rmap f (TTT l r) = fmap (\r' -> TTT l r') (f r)
rmap f (III l r) = fmap (\r' -> III l r') (f r)

fibswap :: FibTree a (l,l') -> FibTree a (l',l)
fibswap (ITT l r) = (ITT r l)
fibswap (TTI l r) = (TIT r l)
fibswap (TIT l r) = (TTI r l)
fibswap (TTT l r) = (TTT r l)

eye = 0 :+ 1

braid :: FibTree a (l,r) -> Q (FibTree a (r,l))
braid (ITT l r) = W [(ITT r l,  cis $ 4 * pi / 5)]  -- different scalar factors for trivial and non trivial fusion
braid (TTT l r) = W [(TTT r l,  (cis $ - 3 * pi / 5))]
braid (TTI l r) = pure $ TIT r l-- exchange with trivial means nothing
braid (TIT l r) = pure $ TTI r l
braid (III l r) = pure $ III r l

-- The inverse of braid
braid' :: FibTree a (l,r) -> Q (FibTree a (r,l))
braid' = star . braid

-- property
-- braid (braid' v) == v
-- braid' (braid v) == v

{-
fibassoc :: FibTree a ((c,d),e) -> FibTree a (c,(d,e))
fibassoc (ITT l r) = (ITT r l)
fibassoc (TTI l r) = (TIT r l)
fibassoc (TIT l r) = (TTI r l)
fibassoc (TTT l r) = (TTT r l)
-}
-- Looks like we need more constructors for the tree
{-
fmove :: FibTree a (c,(d,e)) -> FibTree a ((c,d),e)
-- fmove (ITT  a  (TTI b c)) = ITI ( TTT  a b) c
fmove (ITT  a  (TIT b c)) = ITT ( TTI  a b) c
fmove (ITT  a  (TTT b c)) = ITT ( TTT  a b) c

fmove (TTT  a  (TTI b c)) = TTI ( TTT  a b) c
fmove (TTT  a  (TIT b c)) = TTT ( TTI  a b) c
fmove (TTT  a  (TTT b c)) = TTT ( TTT  a b) c

fmove (TIT  a  (TTI b c)) = TTI ( TIT  a b) c
-- fmove (TIT  a  (TIT b c)) = TTT ( III  a b) c
fmove (TIT  a  (TTT b c)) = TTT ( TIT  a b) c
-}

fmove :: FibTree a (c,(d,e)) -> Q (FibTree a ((c,d),e))
-- fmove (ITT  a  (TTI b c)) = pure $ ITI ( TTT  a b) c -- pure (auto (auto a b) c) -- no maybe not. The internal one isn't auto
fmove (ITT  a  (TIT b c)) = pure $ ITT ( TTI  a b) c
fmove (ITT  a  (TTT b c)) = pure $ ITT ( TTT  a b) c
fmove (ITT  a  (TTI b c)) = pure $ III ( ITT  a b) c

fmove (TIT  a  (TTT b c)) = pure $ TTT ( TIT  a b) c
fmove (TIT  a  (TTI b c)) = pure $ TTI ( TIT  a b) c
fmove (TIT  a  (TIT b c)) = pure $ TIT ( III  a b) c

-- fmove (TIT  a  (TIT b c)) = TTT ( III  a b) c
-- the nontrivial ones have all tau on the leafs and root
-- internal I
fmove (TTI  a  (III b c)) = pure $ TTI ( TTI  a b) c
fmove (TTI  a  (ITT b c)) = W [(TIT ( ITT  a b) c, tau)         , (TTT ( TTT  a b) c, sqrt tau)]
-- internal T
fmove (TTT  a  (TTT b c)) = W [(TIT ( ITT  a b) c, sqrt tau)  ,   (TTT ( TTT  a b) c, - tau   )]
fmove (TTT  a  (TTI b c)) = pure $ TTI ( TTT  a b) c
fmove (TTT  a  (TIT b c)) = pure $ TTT ( TTI  a b) c

fmove (III  a  (ITT b c)) = pure $ ITT ( TIT  a b) c
fmove (III  a  (III b c)) = pure $ III ( III  a b) c

-- largely just a tranpose of the above case.
fmove' :: FibTree a ((c,d),e) -> Q (FibTree a (c,(d,e)))
fmove' (ITT ( TTI  a b) c) = pure $ (ITT  a  (TIT b c))
fmove' (ITT ( TTT  a b) c) = pure $  (ITT  a  (TTT b c))
fmove' (ITT ( TIT  a b) c) = pure $  (III  a  (ITT b c))

--fmoveq (ITT  a  (TTT b c)) = pure $

fmove' (TTI ( TTT  a b) c) = pure $ (TTT  a  (TTI b c))
fmove' (TTI ( TTI  a b) c) = pure $ (TTI  a  (III b c))
fmove' (TTI ( TIT  a b) c) = pure $ TIT  a  (TTI b c)
--fmoveq (TTT  a  (TTI b c)) = pure $ TTI ( TTT  a b) c

--fmoveq (TTT  a  (TIT b c)) = pure $ TTT ( TTI  a b) c
fmove' (TIT ( ITT  a b) c) = W [(TTI  a  (ITT b c), tau)         , (TTT  a  (TTT b c) , sqrt tau)]
fmove' (TIT ( III  a b) c ) = pure $ TIT  a  (TIT b c)

fmove' (TTT ( TTI  a b) c ) = pure $ TTT  a  (TIT b c)
fmove' (TTT ( TIT  a b) c ) = pure $ TIT  a  (TTT b c)
fmove' (TTT ( TTT  a b) c) = W [(TTI  a  (ITT b c), sqrt tau)  , (TTT  a  (TTT b c),   - tau  )]

fmove' (III ( III  a b) c ) = pure $ III  a  (III b c)
fmove' (III ( ITT  a b) c ) = pure $ ITT  a  (TTI b c)

--fmoveq (TIT  a  (TTI b c)) = pure $ TTI ( TIT  a b) c
-- fmove (TIT  a  (TIT b c)) = TTT ( III  a b) c
-- the nontrivial ones have all tau on the leafs and root

-- internal I

--fmoveq (TTI  a  (ITT b c)) = W [(TIT ( ITT  a b) c, recip tau)         , (TTT ( TTT  a b) c, recip $ sqrt tau)]
-- internal T

--fmoveq (TTT  a  (TTT b c)) = W [(TIT ( ITT  a b) c, recip $ sqrt tau)  , (TTT ( TTT  a b) c,   - recip tau  )]

checkf :: Q (FibTree Tau ((Tau, Tau),Tau))
checkf = simplify $ fmove' (TIT (ITT TLeaf TLeaf) TLeaf) >>= fmove
checkf' = simplify $ fmove' (TTT (TTT TLeaf TLeaf) TLeaf) >>= fmove
-- checkf''' = fmove' (TIT (ITT TLeaf TLeaf) TLeaf)
-- checkf'' = fmove' (ITT (TTT TLeaf TLeaf) TLeaf) >>= fmove

-- tau**2 + tau == 1
tau :: Complex Double
tau =  ((sqrt 5) - 1) / 2 :+ 0 -- 0.618 :+ 0

--test1 = lmap fibswap
-- test2 = lmap $ lmap fibswap
-- test3 = rmap fmove

type Vec1 b r = [(b, r)]

smul :: Num r => r -> Vec1 b r -> Vec1 b r
smul s = map (\(b,s') -> (b, s' * s))

linapply :: Num r => (a -> Vec1 b r) -> (Vec1 a r -> Vec1 b r) -- bind
linapply f [] = []
linapply f ((b,s) : xs) = smul s (f b) ++ linapply f xs

dTTT :: FibTree a (Tau, Tau) -> Q (FibTree a Tau)
dTTT (TTT _ _) = pure TLeaf
dTTT _ = mempty
{-
dITT :: FibTree a (Tau, Tau) -> Q (FibTree a Id)
dTTT (TTT _ _) = pure TLeaf
dTTT _ = mempty
-}
dTTI :: FibTree a (Tau, Id) -> Q (FibTree a Tau) -- A dual tree of type  'tau (tau,'tau)
dTTI (TTI _ _) =  pure TLeaf
dTTI _ = mempty

dTTT'''' = dot (TTT TLeaf TLeaf)
-- I could do this as a typeclass. NOPE. don't need that garbage

-- I think that incompatible FibTree a' a is empty by default vs unconstructible
-- which is better?
dot :: FibTree a (b, c) -> FibTree a' (b, c) -> Q (FibTree a' a)
dot x@(TTI _ _) y@(TTI _ _) | x == y = pure TLeaf
                             | otherwise = mempty
dot x@(TIT _ _) y@(TIT _ _) | x == y = pure TLeaf
                             | otherwise = mempty
dot x@(TTT _ _) y@(TTT _ _) | x == y = pure TLeaf
                            | otherwise = mempty
dot x@(III _ _) y@(III _ _) | x == y = pure ILeaf
                            | otherwise = mempty
dot x@(ITT _ _) y@(ITT _ _) | x == y = pure ILeaf
                            | otherwise = mempty
dot _ _ = mempty

test5 = (pure $ TTT TLeaf TLeaf) >>= dTTT
test6 = (pure $ ITT TLeaf TLeaf) >>= dTTT
test7 v = do y <- rmap dTTI v  -- wait. What the hell am I doing?
             lmap dTTT y
test8 v = rmap dTTT v  >>= dTTT
-- Building from the leaves up...?
-- you  hve to apply the deeper stuff before the upper stuff.
--- So this is a dual tree with TTT in the right branch. Then TTT above that.
-- In total (T,(T,T)) -> Tleaf.
-- maybe we need to cps it.
dTTT' :: (Q (FibTree a Tau) -> (Complex Double)) -> (FibTree a (Tau, Tau)) -> Complex Double
dTTT' f = f . dTTT

dTTT'' :: (FibTree a Tau -> Complex Double) -> FibTree a (Tau, Tau) -> Complex Double
dTTT'' f = undefined

-- sort of similar to bind in a way. lift function to a Q function
-- (Q FibTree a ) -> (FibTree a -> Q FibTree b)  -> Q FibTree b
dot' :: (FibTree a b -> Complex Double) -> Q (FibTree a b) -> Complex Double
dot' f (W v) = sum $ map (\(e,a) -> a * (f e)) v

dTTT''' = dot' . dTTT''

pentagon1 ::  FibTree a (e,(d,(c,b))) -> Q (FibTree a (((e,d),c),b))
pentagon1 v =  do
                 v1 <- fmove v
                 fmove v1

-- type annotations not necessary. For clarity.
pentagon2 :: FibTree a (b,(c,(d,e))) -> Q (FibTree a (((b,c),d),e))
pentagon2 v = do
                v1 :: FibTree a (b,((c,d),e)) <- rmap fmove v
                v2 :: FibTree a ((b,(c,d)),e) <- fmove v1
                lmap fmove v2

ex3 = TTT TLeaf (TTI TLeaf (ITT TLeaf TLeaf))
ex1 = TTT TLeaf (TTT TLeaf (TTT TLeaf TLeaf))
pentagon =  simplify $ ((pentagon1 ex1) - (pentagon2 ex1))
pentagon' =  simplify $ ((pentagon1 ex3) - (pentagon2 ex3))

hexagon1 :: FibTree a (b,(c,d)) -> Q (FibTree a ((d,b),c))
hexagon1 v = do
             v1 :: FibTree a ((b,c),d) <- fmove v
             v2 :: FibTree a (d,(b,c)) <- braid v1
             fmove v2

hexagon2 :: FibTree a (b,(c,d)) -> Q (FibTree a ((d,b),c))
hexagon2 v = do
             v1 :: FibTree a (b,(d,c)) <- rmap braid v
             v2 :: FibTree a ((b,d),c) <- fmove v1
             lmap braid v2

ex4 = (TTI TLeaf (ITT TLeaf TLeaf))
ex2 = (TTT TLeaf (TTT TLeaf TLeaf))
hexagon =  simplify $ ((hexagon1 ex2) - (hexagon2 ex2))
hexagon' =  simplify $ ((hexagon1 ex4) - (hexagon2 ex4))

-- hexagon2 :: FibTree a (b,(c,d)) -> Q (FibTree a (c,(d,b)))

-- quickcheck forall v, pentagon1 v == pentagon2 v

-- test9 = test8 $ pure $ TTT (TLeaf) ()
-- Vec1 (FibTree 'Tau l) r -> Vec1 (FibTree 'Tau l') r -> Vec1 (FibTree 'Id (l,l')) r

-- linearized tensor products
--linITT :: Num r => Vec1 (FibTree 'Tau l) r -> Vec1 (FibTree 'Tau l') r -> Vec1 (FibTree 'Id (l,l')) r -- we would like to enforce that isingany can split into a b, but thisis tough
--linITT = [ (ITT l r, s * s')  |  (l, s) <- v1 , (r, s') <- v2   ]

-- class AutoNode c a b where
      -- prod ::  FibTree a l -> FibTree b l' -> FibTree  c (l,l')
{-
 instance AutoProd 'Tau 'Tau 'Tau
    prod = TTT
 instance AutoProd 'Id 'Tau 'Tau
    prod = ITT
-- etc.

-- I can't make a fst. It extracts a skolemized object.
-- So I do need to do complete pattern matching? No. I CAN do this... ? Not always clear I'd want to?
fibfst :: AutoNode a e _ => FibTree a (b,c) -> FibTree e b

-}
class TProd a b c where
  tprod :: Vec1 (FibTree a l) r -> Vec1 (FibTree b l') r -> Vec1 (FibTree c (l,l')) r

-- tprod :: AutoNode a b c => Vec1 (FibTree a l) r -> Vec1 (FibTree b l') r -> Vec1 (FibTree c (l,l')) r
--class Index tree n  -- reference leaf by number?

-- class AutoFMove tree n

{-
linTTI :: FibTree 'Tau l -> FibTree 'Id l' -> FibTree 'Tau (l,l')
linTIT :: FibTree 'Id l -> FibTree 'Tau l' -> FibTree 'Tau (l,l')
libTTT :: FibTree 'Tau l -> FibTree 'Tau l' -> FibTree 'Tau (l,l')
-}
(~*) :: r -> b -> (b,r) -- **, `smul`
(~*) = flip (,)

(~+) :: a -> [a] -> [a] -- ++ instead?
(~+) = (:)

test4 = [1 ~* 'a', 2 ~* 'b'] -- eh. What's the point.
-- test3 = (lmap . lmap) fibswap
--lmap (ITT l r)
-- ^+^
-- zero
--

-- TODOS

-- Categorical Interface
-- lmap, rmap vs second first.

{-
pullLeft (Tau,x) = (Tau,x)
pullLeft (Id,x) = (Id,x)
pullLeft (a,b) = fmove ((a',c),b) where (a',c) = pullLeft a
-- ((a,b),c) = (, (b,c)) where (a',d) = pullLeft a
-}
{-
type family FMove' a where
  FMove' ((a,b),c)  = (a,(b,c))

type family PullLeft a where
  PullLeft Tau = Tau
  PullLeft Id = Id
  PullLeft (Tau,b) = (Tau,b)
  PullLeft (Id,b) =  (Id,b)
  PullLeft ((a,b),c)  = FMove' (PullLeft (a,b), c)

-- Auto F Moves. / auto braid.

pullLeft :: (PullLeft (b,c) ~ (b',c')) => FibTree a (b,c) -> Q (FibTree a (b',c'))
pullLeft x@(TTT TLeaf _) = pure x
pullLeft (TTT l@(TTT _ _ ) r) =  do
                      l' <- pullLeft l
                      fmove' (TTT l' r)
-}
class PullLeft a b | a -> b where -- | a -> b functional dependency causes errors?
  pullLeft :: FibTree c a -> Q (FibTree c b)

instance PullLeft (Tau,c) (Tau,c) where
  pullLeft = pure

instance PullLeft (Id,c) (Id,c) where
  pullLeft = pure

instance PullLeft Tau Tau where
  pullLeft = pure

instance PullLeft Id Id where
  pullLeft = pure

instance (PullLeft (a,b) (a',b'), r ~ (a',(b',c))) => PullLeft ((a, b),c) r where
  pullLeft t = do
           t' <- lmap pullLeft t
           fmove' t'
{-
instance (PullLeft a a', r ~ (a',(b,c))) => PullLeft (a, (b,c)) r where
  pullLeft t = lmap pullLeft t
-}
{-
instance (PullLeft a a', PullLeft (a',b) (a'',b''), ) => PullLeft (a,b) r where
  pullLeft = lmap pullLeft t
-}

{-  pullLeft (TTT l r) =  do
                        l' <- pullLeft l
                        fmove' (TTT l' r)
-}

class PullRight a b | a -> b where -- | a -> b functional dependency causes errors?
  pullRight :: FibTree c a -> Q (FibTree c b)
{-
instance (PullRight a a', r ~ ((b,c),a')) => PullRight ((b,c),a) r where
  pullRight t = rmap pullRight t
-}
instance PullRight Tau Tau where
  pullRight = pure

instance PullRight Id Id where
  pullRight = pure

instance PullRight (c,Tau) (c,Tau) where
  pullRight = pure

instance PullRight (c,Id) (c,Id) where
  pullRight = pure

instance (PullRight (a,b) (a',b'), r ~ ((c,a'),b')) => PullRight (c,(a, b)) r where
  pullRight t = do
           t' <- rmap pullRight t
           fmove t'

type family Count a where
  Count Tau = 1
  Count Id = 1
  Count (a,b) = (Count a) + (Count b)

type family LeftCount a where
  LeftCount (a,b) = Count a

-- The version without the explicit ordering supplied.
class LCA n a b c d | n a c -> b d where
  lcamap :: (forall r. FibTree r b -> Q (FibTree r c)) -> (FibTree e a) -> Q (FibTree e d)
{-
lcamapP :: LCA n a b c d => Proxy n ->  (forall r. FibTree r b -> Q (FibTree r c)) -> (FibTree e a) -> Q (FibTree e d)
lcamapP _ f x = lcamap f x
-}
instance (lc ~ (LeftCount a),
          gte ~ (CmpNat lc n),
         LCA' n gte a b c d) => LCA n a b c d where
  lcamap f x = lcamap' @n @gte f x

class LCA' n gte a b c d | n gte a c -> b d where
  lcamap' :: (forall r. FibTree r b -> Q (FibTree r c)) -> (FibTree e a) -> Q (FibTree e d)

-- we find b at the lca and pass it back up. c gets passed all the way down, d gets computed by rebuilding out of c.
-- a drives the search.
instance (n' ~ (n - Count l), -- we're searching in the right subtree. Subtract the leaf number in the left subtree
        lc ~ (LeftCount r), -- dip one left down to order which way we have to go next
        gte ~ (CmpNat lc n'), -- Do we go left, right or havce we arrived in the next layer?
        LCA' n' gte r b c d',  -- recurive call
        d ~ (l,d') -- reconstruct total return type from recurive return type. left tree is unaffected by lcamapping
        ) => LCA' n 'LT (l,r) b c d where
    lcamap' f x = rmap (lcamap' @n' @gte f) x

instance (lc ~ (LeftCount l),
          gte ~ (CmpNat lc n),
          LCA' n gte l b c d',
          d ~ (d',r)
          ) => LCA' n 'GT (l,r) b c d where
    lcamap' f x = lmap (lcamap' @n @gte f) x

instance (b ~ a, d ~ c) => LCA' n 'EQ a b c d where
  lcamap' f x = f x

class LeafMap n gte a b c d | n gte a c -> b d where
  leafmap :: (forall r. FibTree r b -> Q (FibTree r c)) -> (FibTree e a) -> Q (FibTree e d)

instance (n' ~ (n - Count l), -- we're searching in the right subtree. Subtract the leaf number in the left subtree
        lc ~ (LeftCount r), -- dip one left down to order which way we have to go next
        gte ~ (CmpNat lc n'), -- Do we go left, right or havce we arrived in the next layer?
        LeafMap n' gte r b c d',  -- recurive call
        d ~ (l,d') -- reconstruct total return type from recurive return type. left tree is unaffected by lcamapping
        ) => LeafMap n 'LT (l,r) b c d where
    leafmap f x = rmap (leafmap @n' @gte f) x

instance (lc ~ (LeftCount l),
          gte ~ (CmpNat lc n),
          LeafMap n gte l b c d',
          d ~ (d',r)
          ) => LeafMap n 'GT (l,r) b c d where
    leafmap f x = lmap (leafmap @n @gte f) x

-- In the equals case, we now continue onward.
instance (lc ~ (LeftCount l),
          gte ~ (CmpNat lc n),
          LeafMap n gte l b c d',
          d ~ (d',r)
          ) => LeafMap n 'EQ (l,r) b c d where
    leafmap f x = lmap (leafmap @n @gte f) x

-- base case
instance (b ~ Tau, d ~ c) => LeafMap 1 'EQ Tau b c d where
  leafmap f x = f x

instance (b ~ Id, d ~ c) => LeafMap 1 'EQ Id b c d where
  leafmap f x = f x

-- split
-- leafmap @3 (const (TTT TLeaf TLeaf))
-- need one last arbitrary left or right fmove to put the two on the same stalk
-- neighbormap :: (LCA n a (l,r) (l',r') d, PullRight l (l',x), PullLeft r (y,r') ) => (FibTree.  -> Q FibTree) -> FibTree

-- I need this to also work for (Tau, Tau)
-- (Tau, yada yada)
-- I.e. I need to do different things depending on

-- Jesus. What a shitshow. But I don't see how to do better.
-- a is starting tree
-- b is extracted part
-- c is transformed b
-- d is recoustrcted a replacing b with c and rearranging.
class NeighborMap a b c d | a c -> b d where
  nmap :: (forall r. FibTree r b -> Q (FibTree r c)) -> FibTree e a -> Q (FibTree e d)

instance NeighborMap ((l,x),(y,r)) (x,y) c ((l,c),r) where
   nmap f x = do
              x'  <- fmove x -- (((l',x),y),r')
              x'' <- lmap fmove' x' -- ((l',(x,y)),r')
              lmap (rmap f) x''
instance NeighborMap (Tau, (y,r)) (Tau, y)  c  (c,r) where
   nmap f x = fmove x >>= lmap f
instance NeighborMap (Id, (y,r)) (Id, y) c (c,r) where
   nmap f x = fmove x >>= lmap f
instance NeighborMap ((l,x), Tau) (x,Tau) c (l,c) where
   nmap f x = fmove' x >>= rmap f
instance NeighborMap ((l,x), Id) (x,Id) c  (l,c) where
   nmap f x = fmove' x >>= rmap f
instance NeighborMap (Tau, Tau) (Tau,Tau) c  c where
   nmap f x = f x
instance NeighborMap (Id, Id) (Id,Id) c  c where
   nmap f x = f x
instance NeighborMap (Tau, Id) (Tau,Id) c  c where
   nmap f x = f x
instance NeighborMap (Id, Tau) (Id,Tau) c  c where
   nmap f x = f x

neighbormap :: forall n a b c d l l' r' x y e r z. (LCA n a b c d,
   b ~ (l,r),
   c ~ ((l',z),r'),
   PullRight l (l',x),
   PullLeft r (y,r')) =>
   (forall r. FibTree r (x,y) -> Q (FibTree r z)) -> FibTree e a -> Q (FibTree e d)
neighbormap f z = lcamap @n @a @b @c @d (helper f) z

helper ::  (c ~ ((l',z),r'),
   PullRight l (l',x),
   PullLeft r (y,r')) => (forall r. FibTree r (x,y) -> Q (FibTree r z)) -> FibTree e (l,r) -> Q (FibTree e c)
helper f x = do
            x' <- rootneighbor x
            lmap (rmap f) x'

rootneighbor :: (PullRight l (l',x), PullLeft r (y,r')) => FibTree e (l,r) -> Q (FibTree e ((l',(x,y)),r'))
rootneighbor x = do
                x' <- lmap pullRight x
                x'' <- rmap pullLeft x' -- ((l',x),(y,r'))
                x''' <- fmove x'' -- (((l',x),y),r')
                lmap fmove' x''' -- ((l',(x,y)),r')

neighbormap' :: forall n a b c d l l' r r' x y e z b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') b' z c) =>
   (forall r. FibTree r b' -> Q (FibTree r z)) -> FibTree e a -> Q (FibTree e d)
neighbormap' f z = lcamap @n @a @b @c @d (\x -> do
                                          x'  <- lmap pullRight x
                                          x'' <- rmap pullLeft x'
                                          nmap f x'') z

neighbormap'' :: forall n a b' z d c l l' r r' x y e b. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') b' z c) =>
   (forall r. FibTree r b' -> Q (FibTree r z)) -> FibTree e a -> Q (FibTree e d)
neighbormap'' f z = lcamap @n @a @b @c @d (\x -> do
                                          x'  <- lmap pullRight x
                                          x'' <- rmap pullLeft x'
                                          nmap f x'') z

t1 = neighbormap' @2 braid (TTT (TTI TLeaf ILeaf) (TTT TLeaf TLeaf))
t2 = neighbormap' @1 braid (TTT (TTI TLeaf ILeaf) (TTT TLeaf TLeaf))
t3 = neighbormap' @2 braid (TTT (TTT (TTT TLeaf TLeaf) TLeaf) (TTT TLeaf TLeaf))

                                      {-
neighbormap p f z = let helper (x :: FibTree _ (l,r)) = do
                                x' <- rootneighbor x
                                lmap (rmap f) x'
                    in
                    lcamap helper z
                    -}
{-  lcamap
((\x -> do
                                x' <- rootneighbor x
                                lmap (rmap f) x'
                                ) :: (forall s. FibTree s (l,r) -> Q (FibTree s c))) z
-}

{-(\x -> do
                            x' <- lmap pullRight x
                            x'' <- rmap pullLeft x' -- ((l',x),(y,r'))
                            x''' <- fmove x'' -- (((l',x),y),r')
                            x4 <- lmap fmove' x''' -- ((l',(x,y)),r')
                            lmap (lmap f) x4 -- ((l',z),r')
                            ) z
-}
{-
-- neighbormap f z = lcamap @n (\x -> do
                            x' <- lmap pullRight x
                            x'' <- rmap pullLeft x'
                            x''' <- fmove x''
                            lmap? f x'''   ) z
-}
{-
abraid :: forall n a b c d l l' r' x y e r z b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) (y,x) c) => Q (FibTree e a) -> Q (FibTree e d)

   -}

abraid :: forall n a b c d l l' r r' x y e z b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) (y,x) c) =>
   Q (FibTree e a) -> Q (FibTree e d)
abraid x = x >>= (neighbormap' @n @a @b @c @d @l @l' @r @r' @x @y @e @(y,x) @(x,y)) braid

abraid' :: forall n a b c d l l' r r' x y e z b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) (y,x) c) =>
   Q (FibTree e a) -> Q (FibTree e d)
abraid' x = x >>= (neighbormap' @n @a @b @c @d @l @l' @r @r' @x @y @e @(y,x) @(x,y)) braid'

{-
abraid
  :: forall n a b c d l l' r r' x y e z b'. (LCA n a (l, r) c d, PullRight l l', PullLeft r r',
      NeighborMap (l', r') (x, y) (y, x) c) =>
     (FibTree e a) -> Q (FibTree e d)
abraid x = neighbormap' @n @a @(l,r) @c @d @l @l' @r @r' braid x
-}
{-
t4 = abraid @2 $
     abraid @4 $
     abraid @3 $
     pure (TTT (TTT (TTT TLeaf TLeaf) TLeaf) (TTT TLeaf TLeaf))
-}
{-
ndot' :: forall n a b c d l l' r r' x y e z b' a'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) a' c) =>
   FibTree a' (x,y) -> FibTree e a -> Q (FibTree e d)
   -}
-- ndot' x = (neighbormap' @n @a @b @c @d @l @l' @r @r' @x @y @e @a' @(x,y)) (dot' x)

ndot :: forall n a l r c d l' r' x y a' e. (LCA n a (l, r) c d, PullRight l l', PullLeft r r',
      NeighborMap (l', r') (x, y) a' c) =>
     FibTree a' (x, y) -> FibTree e a -> Q (FibTree e d)
ndot x = neighbormap' @n @a @(l,r) (dot x)

-- This a b c d pattern is something like a lens. We are replacing
-- piece of b with c which turns a into d. b is piece of a.
-- s t a b

-- quantified constraints
--  (forall b' c'. (g a b' c' d , f b' b c c') => Compose f g a b c d

-- (forall r. FibTree r b' -> Q (FibTree r z)) ->
{-
abraid' :: forall n a b c d l l' r r' x y e b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) (y,x) c) =>
   FibTree e a -> Q (FibTree e d)
abraid' = (neighbormap'' @n @a @(x,y) @(y,x) @d) braid
-}

{-
abraid' :: forall n a b c d l l' r r' x y e z b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) (y,x) c) =>
   FibTree e a -> Q (FibTree e d)
abraid' = (neighbormap' @n @a @b @c @d @l @l' @r @r' @x @y @e @(y,x) @(x,y)) braid'
-}

{-
abraid :: (LCA n a (l1, r) c d, PullRight l1 l'1, PullLeft r r',
      NeighborMap (l'1, r') (l2, l'2) (l'2, l2) c) =>
     FibTree e a -> Q (FibTree e d)
abraid = neighbormap' braid
-}
-- autobraid = neighbormap braid
-- autodot x = neighbormap (dot x)

-- Build state monad to carry along the vector.
-- should be able to do in applicative style, since data does not change structre of computation

-- auto braid
-- find least common ancestor
-- pullLeft, and pullRight
-- then braid

-- stateful
-- Pretty Print tree

-- Operators.
newtype FibOp a b = FibOp (forall c. FibTree c a -> Q (FibTree c b))
type FibOp' c a b = FibTree c a -> Q (FibTree c b)

instance Category FibOp where
  id = FibOp pure
  (FibOp f) . (FibOp g) = FibOp (f <=< g)
{-
instance Arrow FibOp where
  arr = error "No. No arr."
  (***) =
-}
{-
class Monoidal k where
  (***) :: k a b -> k c d -> k (a,c) (b,d)
-}
tensor :: Num r => (b -> c -> a) -> W r b -> W r c -> W r a
tensor f (W xs) (W ys) =  W [ ( f x y , xc * yc ) | (x, xc) <- xs,  (y,yc) <- ys]

-- tensor :: Num r => (b -> c -> a) -> b -> c -> W r a
-- fg f g (TTT l r) = tensor TTT (f l) (g r)

{-
instance Monoidal FibOp where
  (***) :: FibOp a b -> FibOp c d -> FibOp (a,c) (b,d)
  (FibOp f) *** (FibOp g) = FibOp fg where
                               -- fg ::  (a,c) (b,d)
                               fg :: FibTree e (a,c) -> Q (FibTree e (b,d))
                               fg (TTT l r) = tensor TTT (f l) (g r)
                               -- fg (ITT l r) = FibOp $ tensor ITT (f l) (g r)
                               -- fg (TIT l r) = FibOp $ tensor TIT (f l) (g r)
-}
-- try to remove the issue of being in a instance. Which is kind of silly anyhow
{-
(****) :: FibOp' e a b -> FibOp' e c d -> FibOp' e (a,c) (b,d)
f **** g = fg where
                   -- fg ::  (a,c) (b,d)
                 fg :: FibTree e (a,c) -> Q (FibTree e (b,d))
                 fg (TTT l r) = tensor TTT (f l) (g r)
                 fg (ITT l r) = tensor ITT (f l) (g r)
-}
(***) :: (forall e. FibOp' e a b) -> (forall e'. FibOp' e' c d)  -> FibOp' e'' (a,c) (b,d)
(***) f g (TTT l r) = tensor TTT (f l) (g r)
(***) f g (ITT l r) = tensor ITT (f l) (g r)
(***) f g (TIT l r) = tensor TIT (f l) (g r)

first = lmap
second = rmap
-- braid
-- fmove

-- cup, cap? dot'

-- Is the problem that I'd trying to pattern match on kind of the output?
 -- densification

-- Quickcheck props

-- implement qubit abtraction

-- Did I fall into a way to make end, co-end work as einstein notation?6
-- The bounded meaning of forall and exists in the GADT context

main :: IO ()
main = mymain
--return ()


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Signal.Internal
-- Copyright   :  (c) Alexander Vivian Hugh McPhail 2010, 2014, 2015, 2016
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :  uses FFI
--
-- low-level interface
--
-----------------------------------------------------------------------------

module Numeric.Signal.Internal (
                Convolvable(..),
                Filterable(..),
                freqz,
                pwelch,
                hilbert
                ) where

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel

--import Numeric.LinearAlgebra.Linear

import qualified Numeric.GSL.Fourier as F
import Foreign

--import Data.Complex
import Foreign.C.Types

import Prelude hiding(filter)

import System.IO.Unsafe(unsafePerformIO)

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

infixr 1 #
a # b = applyRaw a b
{-# INLINE (#) #-}

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

type PD = Ptr Double
type PC = Ptr (Complex Double)
type PF = Ptr Float

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

class Convolvable a where
    -- | convolve two containers, output is the size of the second argument, no zero padding
    convolve :: a -> a -> a

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

class (Storable a, Container Vector a, Num (Vector a)
      , Convert a, Floating (Vector a), RealElement a
      , Num a)
       => Filterable a where
    -- | convert from Vector Double
    fromDouble :: Vector Double -> Vector a
--       b ~ ComplexOf a, Container Vector b, Convert b) => Filterable a where
    -- | filter a signal
    filter_ :: Vector a -- ^ zero coefficients
            -> Vector a -- ^ pole coefficients
            -> Vector a -- ^ input signal
            -> Vector a -- ^ output signal
    -- | coefficients of a Hamming window
    hamming_ :: Int           -- ^ length
           -> Vector a -- ^ the Hamming coeffficents
    -- | the complex power : real $ v * (conj v)
    complex_power_ :: Vector (Complex Double) -- ^ input
                  -> Vector a                       -- ^ output
     -- | resample, take one sample every n samples in the original
    downsample_ :: Int -> Vector a -> Vector a
    -- | the difference between consecutive elements of a vector
    deriv_ :: Vector a -> Vector a
    -- | unwrap the phase of signal (input expected to be within (-pi,pi)
    unwrap_ :: Vector a -> Vector a
    -- | evaluate a real coefficient polynomial for complex arguments
    polyEval_ :: Vector a           -- ^ the real coefficients
         -> Vector (Complex Double) -- ^ the points at which to be evaluated
         -> Vector (Complex Double) -- ^ the values
    -- | the cross covariance of two signals
    cross_covariance_ :: Int            -- ^ maximum delay
                     -> Vector a       -- ^ time series
                     -> Vector a       -- ^ time series
                     -> (a,a,Vector a) -- ^ (sd_x,sd_y,cross_covariance)
    -- | the cumulative sum of a signal
    cumulative_sum_ :: Vector a    -- ^ time series
                    -> Vector a    -- ^ result

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

instance Convolvable (Vector Double) where
    convolve x y = fst $ fromComplex $ F.ifft $ (F.fft (complex x) * F.fft (complex y))
--    convolve = convolve_vector_double

convolve_vector_double c a = unsafePerformIO $ do
                             r <- createVector (size a)
                             (c # a # r # id) signal_vector_double_convolve #| "signalDoubleConvolve"
                             return r

foreign import ccall "signal-aux.h vector_double_convolve" signal_vector_double_convolve :: CInt -> PD -> CInt -> PD -> CInt -> PD -> IO CInt

instance Convolvable (Vector Float) where
    convolve x y = single $ fst $ fromComplex $ F.ifft $ (F.fft (complex $ double x) * F.fft (complex $ double y))
--    convolve = convolve_vector_double

convolve_vector_float c a = unsafePerformIO $ do
                             r <- createVector (size a)
                             (c # a # r # id ) signal_vector_float_convolve #| "signalFloatConvolve"
                             return r

foreign import ccall "signal-aux.h vector_float_convolve" signal_vector_float_convolve :: CInt -> PF -> CInt -> PF -> CInt -> PF -> IO CInt

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

instance Convolvable (Vector (Complex Double)) where
    convolve x y = F.ifft $ (F.fft x * F.fft y)
--    convolve = convolve_vector_complex

convolve_vector_complex c a = unsafePerformIO $ do
                              r <- createVector (size a)
                              (c # a # r # id) signal_vector_complex_convolve #| "signalComplexConvolve"
                              return r

foreign import ccall "signal-aux.h vector_complex_convolve" signal_vector_complex_convolve :: CInt -> PC -> CInt -> PC -> CInt -> PC -> IO CInt

instance Convolvable (Vector (Complex Float)) where
    convolve x y = single $ F.ifft $ (F.fft (double x) * F.fft (double y))

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

instance Filterable Double where
    fromDouble = id
    filter_ = filterD
    hamming_ = hammingD
    complex_power_ = complex_powerD
    downsample_ = downsampleD
    deriv_ = derivD
    unwrap_ = unwrapD
    polyEval_ = polyEval
    cross_covariance_ = crossCovarianceD
    cumulative_sum_ = cumSumD

instance Filterable Float where
    fromDouble = single
    filter_ = filterF
    hamming_ = hammingF
    complex_power_ = complex_powerF
    downsample_ = downsampleF
    deriv_ = derivF
    unwrap_ = unwrapF
    polyEval_ c = polyEval (double c)
    cross_covariance_ = crossCovarianceF
    cumulative_sum_ = cumSumF

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

-- | filters the signal
filterD :: Vector Double -- ^ zero coefficients
       -> Vector Double -- ^ pole coefficients
       -> Vector Double -- ^ input signal
       -> Vector Double -- ^ output signal
filterD l k v = unsafePerformIO $ do
               r <- createVector (size v)
               (l # k # v # r # id) signal_filter_double #| "signalFilter"
               return r

foreign import ccall "signal-aux.h filter_double" signal_filter_double :: CInt -> PD -> CInt -> PD -> CInt -> PD -> CInt -> PD -> IO CInt

-- | filters the signal
filterF :: Vector Float -- ^ zero coefficients
       -> Vector Float -- ^ pole coefficients
       -> Vector Float -- ^ input signal
       -> Vector Float -- ^ output signal
filterF l k v = unsafePerformIO $ do
               r <- createVector (size v)
               (l # k # v # r # id) signal_filter_float #| "signalFilter"
               return r

foreign import ccall "signal-aux.h filter_float" signal_filter_float :: CInt -> PF -> CInt -> PF -> CInt -> PF -> CInt -> PF -> IO CInt

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

-- | Hilbert transform with original vector as real value, transformed as imaginary
hilbert :: Vector Double -> Vector (Complex Double)
hilbert v = unsafePerformIO $ do
            let r = complex v
            -- could use (complex v) to make a complex vector in haskell rather than C
            (r # id) signal_hilbert #| "hilbert"
            return r

foreign import ccall "signal-aux.h hilbert" signal_hilbert :: CInt -> PC -> IO CInt

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

-- | Welch (1967) power spectrum density using periodogram/FFT method
pwelch :: Int            -- ^ window size (multiple of 2)
       -> Vector Double  -- ^ input signal
       -> Vector Double  -- ^ power density
pwelch w v = unsafePerformIO $ do
             let r = konst 0.0 ((w `div` 2) + 1)
             (complex v # r # id) (signal_pwelch $ fromIntegral w) #| "pwelch"
             return r

foreign import ccall "signal-aux.h pwelch" signal_pwelch :: CInt -> CInt -> PC -> CInt -> PD -> IO CInt

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

-- | coefficients of a Hamming window
hammingD :: Int           -- ^ length
        -> Vector Double -- ^ the Hamming coeffficents
hammingD l
    | l == 1          = konst 1.0 1
    | otherwise       = unsafePerformIO $ do
                        r <- createVector l
                        (r # id) signal_hamming_double #| "Hamming"
                        return r

foreign import ccall "signal-aux.h hamming_double" signal_hamming_double :: CInt -> PD -> IO CInt

-- | coefficients of a Hamming window
hammingF :: Int           -- ^ length
        -> Vector Float -- ^ the Hamming coeffficents
hammingF l
    | l == 1          = konst 1.0 1
    | otherwise       = unsafePerformIO $ do
                        r <- createVector l
                        (r # id) signal_hamming_float #| "Hamming"
                        return r

foreign import ccall "signal-aux.h hamming_float" signal_hamming_float :: CInt -> PF -> IO CInt

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

-- | determine the frequency response of a filter
{-freqz :: (Filterable a, Storable a, Container Vector a, Convert a, RealElement a,
         DoubleOf a ~ DoubleOf (RealOf b), RealElement c, c ~ DoubleOf a, c ~ DoubleOf (RealOf b),
         b ~ Complex a, b ~ ComplexOf a, Convert b, Container Vector b,
         Container Vector c, Convert c, b ~ ComplexOf a, b ~ ComplexOf c)
        ⇒ Vector a     -- ^ zero coefficients
      -> Vector a       -- ^ pole coefficients
      -> Vector a       -- ^ points (between 0 and 2*pi)
      -> Vector a       -- ^ response
-}
freqz :: (Filterable a, Complex Double ~ ComplexOf (DoubleOf a)
        ,Filterable (DoubleOf a)) =>
        Vector a       -- ^ zero coefficients
      -> Vector a       -- ^ pole coefficients
      -> Vector a       -- ^ points (between 0 and 2*pi)
      -> Vector a       -- ^ response
freqz b a w = let k = max (size b) (size a)
                  hb = polyEval_ (postpad b k) (exp (scale (0 :+ 1) ((complex $ double w))))
                  ha = polyEval_ (postpad a k) (exp (scale (0 :+ 1) ((complex $ double w))))
              in complex_power_ (hb / ha)

postpad v n = let d = size v
              in if d < n then vjoin [v,(konst 0.0 (n-d))]
              else v

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

-- | evaluate a real coefficient polynomial for complex arguments
polyEval :: Vector Double           -- ^ the real coefficients
         -> Vector (Complex Double) -- ^ the points at which to be evaluated
         -> Vector (Complex Double) -- ^ the values
polyEval c z = unsafePerformIO $ do
               r <- createVector (size z)
               (c # z # r # id) signal_real_poly_complex_eval #| "polyEval"
               return r

foreign import ccall "signal-aux.h real_poly_complex_eval" signal_real_poly_complex_eval :: CInt -> PD -> CInt -> PC -> CInt -> PC -> IO CInt

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

-- | the complex power : real $ v * (conj v)
complex_powerD :: Vector (Complex Double) -- ^ input
              -> Vector Double           -- ^ output
complex_powerD v = unsafePerformIO $ do
                  r <- createVector (size v)
                  (v # r # id) signal_complex_power_double #| "complex_power"
                  return r

foreign import ccall "signal-aux.h complex_power_double" signal_complex_power_double :: CInt -> PC -> CInt -> PD -> IO CInt

-- | the complex power : real $ v * (conj v)
complex_powerF :: Vector (Complex Double) -- ^ input
              -> Vector Float             -- ^ output
complex_powerF v = unsafePerformIO $ do
                  r <- createVector (size v)
                  (v # r # id) signal_complex_power_float #| "complex_power"
                  return r

foreign import ccall "signal-aux.h complex_power_float" signal_complex_power_float :: CInt -> PC -> CInt -> PF -> IO CInt

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

-- | resample, take one sample every n samples in the original
downsampleD :: Int -> Vector Double -> Vector Double
downsampleD n v = unsafePerformIO $ do
               r <- createVector (size v `div` n)
               (v # r # id) (signal_downsample_double $ fromIntegral n) #| "downsample"
               return r

foreign import ccall "signal-aux.h downsample_double" signal_downsample_double :: CInt -> CInt -> PD -> CInt -> PD -> IO CInt

-- | resample, take one sample every n samples in the original
downsampleF :: Int -> Vector Float -> Vector Float
downsampleF n v = unsafePerformIO $ do
               r <- createVector (size v `div` n)
               (v # r # id) (signal_downsample_float $ fromIntegral n) #| "downsample"
               return r

foreign import ccall "signal-aux.h downsample_float" signal_downsample_float :: CInt -> CInt -> PF -> CInt -> PF -> IO CInt

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

-- | the difference between consecutive elements of a vector
derivD :: Vector Double -> Vector Double
derivD v = unsafePerformIO $ do
          r <- createVector (size v - 1)
          (v # r # id) (signal_diff_double) #| "diff"
          return r

foreign import ccall "signal-aux.h vector_diff_double" signal_diff_double :: CInt -> PD -> CInt -> PD -> IO CInt

-- | the difference between consecutive elements of a vector
derivF :: Vector Float -> Vector Float
derivF v = unsafePerformIO $ do
          r <- createVector (size v - 1)
          (v # r # id) (signal_diff_float) #| "diff"
          return r

foreign import ccall "signal-aux.h vector_diff_float" signal_diff_float :: CInt -> PF -> CInt -> PF -> IO CInt

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

-- | unwrap the phase of signal (input expected to be within (-pi,pi)
unwrapD :: Vector Double -> Vector Double
unwrapD v = unsafePerformIO $ do
           r <- createVector $ size v
           (v # r # id) signal_unwrap_double #| "unwrap"
           return r

foreign import ccall "signal-aux.h unwrap_double" signal_unwrap_double :: CInt -> PD -> CInt -> PD -> IO CInt

-- | unwrap the phase of signal (input expected to be within (-pi,pi)
unwrapF :: Vector Float -> Vector Float
unwrapF v = unsafePerformIO $ do
           r <- createVector $ size v
           (v # r # id) signal_unwrap_float #| "unwrap"
           return r

foreign import ccall "signal-aux.h unwrap_float" signal_unwrap_float :: CInt -> PF -> CInt -> PF -> IO CInt

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

-- | compute the cross covariance of two signals
crossCovarianceD :: Int -> Vector Double -> Vector Double -> (Double,Double,Vector Double)
crossCovarianceD l x y = unsafePerformIO $ do
                           r <- createVector (2*l)
                           alloca $ \sx ->
                               alloca $ \sy -> do
                                 (x # y # r # id) (signal_cross_covariance_double (fromIntegral l) sx sy) #| "cross_covariance"
                                 sx' <- peek sx
                                 sy' <- peek sy
                                 return (sx',sy',r)

foreign import ccall "signal-aux.h cross_covariance_double"
        signal_cross_covariance_double :: CInt -> PD -> PD -> CInt -> PD -> CInt
                                       -> PD -> CInt -> PD -> IO CInt

-- | compute the cross covariance of two signals
crossCovarianceF :: Int -> Vector Float -> Vector Float -> (Float,Float,Vector Float)
crossCovarianceF l x y = unsafePerformIO $ do
                           r <- createVector (2*l)
                           alloca $ \sx ->
                               alloca $ \sy -> do
                                 (x # y # r # id) (signal_cross_covariance_float (fromIntegral l) sx sy) #| "cross_covariance"
                                 sx' <- peek sx
                                 sy' <- peek sy
                                 return (sx',sy',r)

foreign import ccall "signal-aux.h cross_covariance_float"
        signal_cross_covariance_float :: CInt -> PF -> PF -> CInt -> PF -> CInt
                                       -> PF -> CInt -> PF -> IO CInt

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

cumSumD :: Vector Double -> Vector Double
cumSumD v = unsafePerformIO $ do
              r <- createVector (size v)
              (v # r # id) signal_cum_sum_double #| "cumSumD"
              return r

cumSumF :: Vector Float -> Vector Float
cumSumF v = unsafePerformIO $ do
              r <- createVector (size v)
              (v # r # id) signal_cum_sum_float #| "cumSumF"
              return r

foreign import ccall "signal-aux.h cum_sum_double"
        signal_cum_sum_double :: CInt -> PD -> CInt -> PD -> IO CInt
foreign import ccall "signal-aux.h cum_sum_float"
        signal_cum_sum_float :: CInt -> PF -> CInt -> PF -> IO CInt

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


-- calculates nash equilibriums, does related backend stuff
{-# LANGUAGE OverloadedStrings #-}

module Game
    ( gameSimplePartial, gameSimplePartialOpts, gameSimple
    , calcPureEVs, calcEVSimpleCore, calcVarSimpleCore, calcSDSimpleCore
    , gameComplex, resultComplex
    , solvePure, solvePureBy, solvePureOn, solvePureZeroSum
    , solvePureOptimal, solvePureOptimalZeroSum
    , solveSimpleCoreNaive
    , calcEVSimpleCoreNaive
    , solveSimple, solveComplex
    , resultComplex
    , Game (Game, gameCName, attCName, defCName, gameData, outcomesC)
    , Result (Result, evc, evr, sdc, sdr, weightsAtts, weightsDefs)
    , ResultSimple  (ResultSimple, evSimpleCol, evSimpleRow, sdSimpleCol, sdSimpleRow, weightsColsSimple, weightsRowsSimple)
    ) where

import GameData

import Data.List
import Data.List.Split
import Data.Function
import Data.Maybe
import Data.Text (Text, unpack)
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Data
import Data.Map (Map)
import qualified Data.Map as Map
import Safe

type EV = Double
type Var = Double
type SD = Double
type Weights = [Double]
type Coord = (Int,Int)
type RemovedElements = [Int]
type WeightedOption = (Opt, Opt, Double)
type BiWeightedOption = (Opt, Opt, Double, Double)
type PayoffsRaw = [[Double]]
type Payoffs = Matrix Double

gameSimplePartial :: PayoffsRaw -> Payoffs
gameSimplePartial = fromLists

gameSimplePartialOpts :: [WeightedOption] -> PayoffsRaw
gameSimplePartialOpts opts =
    let
        fst3 (x,_,_) = x
        snd3 (_,x,_) = x
        thd3 (_,_,x) = x

        cols = nub . map (fst . fst3) $ opts
        rows = nub . map (fst . snd3) $ opts
        grid = map (map thd3) . transpose . chunksOf (length rows) $ opts
    in
        grid

gameSimple :: Text -> [Text] -> [Text] -> PayoffsRaw -> PayoffsRaw -> GameSimple
gameSimple n c r m1 m2 = GameSimple n c r (fromLists m1) (fromLists m2) $ solveSimple m1 m2

calcPureEVs :: Payoffs -> Weights -> [EV]
calcPureEVs m cols = concat . toLists $ m Numeric.LinearAlgebra.<> (col (fmap (/ sum cols) cols))

calcEVSimpleCore :: Payoffs -> Weights -> Weights -> EV
calcEVSimpleCore m cols rows = (head . head . toLists $ (row rows) Numeric.LinearAlgebra.<> m Numeric.LinearAlgebra.<> (col cols))

calcVarSimpleCore :: Payoffs -> Weights -> Weights -> Var
calcVarSimpleCore mat cols rows = calcEVSimpleCore (cmap (\x -> (x - calcEVSimpleCore mat cols rows)**2) mat) cols rows

calcSDSimpleCore :: Payoffs -> Weights -> Weights -> SD
calcSDSimpleCore a b = sqrt . calcVarSimpleCore a b

gameComplex :: Text -> Text -> Text -> [BiWeightedOption] -> Game
gameComplex t1 t2 t3 m = Game t1 t2 t3 (sort m) . resultComplex m $ solveComplex m

resultComplex :: [(Opt, Opt, Double, Double)] -> ResultSimple -> Result
resultComplex opts (ResultSimple evc evr sdc sdr cw rw) = do
    let colnames = nub . map (fst . fst4) $ opts
    let rownames = nub . map (fst . snd4) $ opts
    Result evc evr sdc sdr (zip colnames cw) (zip rownames rw)

-- from here on out, stuff for solving games

-- the fundamental algorithm, given some ordering method to know which values are better
solvePureBy :: (Eq a) => (a -> a -> Ordering) -> [[a]] -> [[a]] -> [Coord]
solvePureBy f mc mr =
    let
        width = length . head $ mc
        height = length mc

        -- column player: picks the highest columns for each row
        mcIsBestResponse = map (\col -> map (== maximumBy f col) col) mc
        -- row player: picks the highest rows for each coloumn
        mrIsBestResponse = transpose . map (\row -> map (== maximumBy f row) row) . transpose $ mr
        bestResponses = map (uncurry zip) $ zip mcIsBestResponse mrIsBestResponse

        coords = map (\r -> map (\c -> (c,r)) [0..width-1]) [0..height-1]
        zipped = map (uncurry zip) $ zip coords bestResponses
        zippedFiltered = map fst . filter (\((x,y),(c,r)) -> and[c,r]) . concat $ zipped
    in
        zippedFiltered

-- solvePureBy applied to mapped values
solvePureOn :: (Ord c) => (a -> c) -> (b -> c) -> [[a]] -> [[b]] -> [Coord]
solvePureOn f g mc mr = solvePureBy compare (map (map f) mc) (map (map g) mr)

solvePure :: (Ord a) => [[a]] -> [[a]] -> [Coord]
solvePure mc mr = solvePureOn id id mc mr

solvePureZeroSum :: (Ord a, Num a) => [[a]] -> [Coord]
solvePureZeroSum m = solvePureOn id negate m m

evPure :: [[a]] -> Coord -> a
evPure m (c,r) = (m!!c)!!r

solvePureOptimalBy :: (Ord a) => (a -> a -> Ordering) -> ([[a]] -> Coord -> a) -> [[a]] -> [[a]] -> Maybe Coord
solvePureOptimalBy f ev mc mr =
    let
        pures = solvePureBy f mc mr
        -- note that we have already optimised the row choice per column choice, as these are all, by definition, nash equilibria
        -- thus, simply choose the optimal nash equilibrium for column ev
        bestPure = maximumByMay (compare `on` (\p -> ev mc p)) pures
    in
        bestPure

solvePureOptimalOn :: (Ord c) => (a -> c) -> (b -> c) -> ([[c]] -> Coord -> c) -> [[a]] -> [[b]] -> Maybe Coord
solvePureOptimalOn f g ev mc mr = solvePureOptimalBy compare ev (map (map f) mc) (map (map g) mr)

solvePureOptimal :: (Ord a) => [[a]] -> [[a]] -> Maybe Coord
solvePureOptimal mc mr = solvePureOptimalOn id id evPure mc mr

solvePureOptimalZeroSum :: (Ord a, Num a) => [[a]] -> Maybe Coord
solvePureOptimalZeroSum m = solvePureOptimal m . map (map negate) $ m

solveSimpleCoreNaive :: Matrix Double -> Maybe Weights
solveSimpleCoreNaive m =
    do
        case (rows m, cols m) of
            -- not really much choice
            (_,1) -> return [1]
            -- they only have one option, so zip your options to their position, get the best option's position, and pad a [1] with zeroes as appropriate to the position
            (1,c) -> return . (\n -> ((take n (repeat 0)) ++ [1] ++ (take (c - n - 1) (repeat 0)))) . fst . maximumBy (compare `on` snd) . zip [0..] . concat . toLists $ m
            -- use linear algebra to find the weights such that unexploitable play is achieved, ie, the ev of each of the opponent's options is equal
            _ -> let
                        weights = normalise . concat . toLists $ ((pinv m) Numeric.LinearAlgebra.<> (((rows m)><1) (repeat (fromInteger 1))))
                    in
                        if all (>=0) weights
                        then return weights
                        else Nothing
    where
        normalise :: (Fractional a) => [a] -> [a]
        normalise xs = map (/ (sum xs)) xs

calcEVSimpleCoreNaive :: Matrix Double -> Maybe EV
calcEVSimpleCoreNaive m =
    do
        wc <- solveSimpleCoreNaive m
        wr <- solveSimpleCoreNaive . tr' $ m
        return $ calcEVSimpleCore m wc wr

gameSupports :: (Ord a, Num a) => [[a]] -> [[a]] -> [[((RemovedElements, RemovedElements), ([[a]], [[a]]))]]
gameSupports m1 m2 =
    let
        widther = length . head
        heighter = length
        w = widther m1 - 1
        h = heighter m1 - 1

        removeIndexes :: RemovedElements -> [a] -> [a]
        removeIndexes [] xs = xs
        removeIndexes _ [] = error "Can't remove from an empty list!"
        removeIndexes [0] (x:xs) = xs
        removeIndexes [n] (x:xs) = x:(removeIndexes [n-1] xs)
        removeIndexes (n:ns) xs = removeIndexes (fmap (subtract 1) ns) (removeIndexes [n] xs)

        subm :: (RemovedElements,RemovedElements) -> [[a]] -> [[a]]
        subm coords = map (removeIndexes (fst coords)) . removeIndexes (snd coords)

        supportsColsRaw = zip [0..] . init . subsequences $ [0..w] -- [[],[0],[1],...,[n],[0,1],[0,2],...,[0,n],[0,1,2],...]
        supportsRowsRaw = zip [0..] . init . subsequences $ [0..h]
        supportsPairs = map (\(x,c) -> map (\(y,r) -> ((c,r), (subm (c,r) m1, subm (c,r) m2))) supportsRowsRaw) supportsColsRaw
    in
        supportsPairs

gameSupportsZeroSum :: (Ord a, Num a) => [[a]] -> [[((RemovedElements, RemovedElements), ([[a]], [[a]]))]]
gameSupportsZeroSum m = gameSupports m (map (map negate) m)

supportIntoValidatedWeightsEVsSDs :: ((RemovedElements, RemovedElements), ([[Double]], [[Double]])) -> Maybe ((Weights, Weights), ((EV, EV), (SD, SD)))
supportIntoValidatedWeightsEVsSDs ((sc,sr), (mc,mr)) =
    do
        let (g1,g2) = (gameSimplePartial mc, tr' $ gameSimplePartial mr)
        w1 <- solveSimpleCoreNaive g1
        w2 <- solveSimpleCoreNaive g2
        let ws = (addIndexes 0 sc w1, addIndexes 0 sr w2)
        let evs = (calcEVSimpleCore g1 w1 w2, calcEVSimpleCore g2 w2 w1)
        let sds = (calcSDSimpleCore g1 w1 w2, calcSDSimpleCore g2 w2 w1)
        return (ws, (evs, sds))
    where
        addIndexes :: a -> [Int] -> [a] -> [a]
        addIndexes _ [] xs = xs
        addIndexes def [0] xs = def:xs
        addIndexes _ _ [] = error "Add index out of bounds"
        addIndexes def [n] (x:xs) = x:(addIndexes def [n-1] xs)
        addIndexes def (n:ns) xs = addIndexes def ns (addIndexes def [n] xs)

supportsIntoValidatedWeightsEVsSDs :: [[((RemovedElements, RemovedElements), ([[Double]], [[Double]]))]] -> [[Maybe ((Weights, Weights), ((EV, EV), (SD, SD)))]]
supportsIntoValidatedWeightsEVsSDs = map (map supportIntoValidatedWeightsEVsSDs)

solveSimple :: [[Double]] -> [[Double]] -> ResultSimple
solveSimple mc mr =
    let
        sups = gameSupports mc mr
        strats = supportsIntoValidatedWeightsEVsSDs sups

        maximumOn f = maximumBy (compare `on` f)
        colEv = fmap (fst . fst . snd)
        rowEv = fmap (snd . fst . snd)

        stratsBestResponses = map (maximumOn rowEv) strats
        stratsBest = maximumOn colEv stratsBestResponses

        res (Just ((wc,wr),((evc,evr),(sdc,sdr)))) = ResultSimple evc evr sdc sdr wc wr
    in
        res stratsBest

solveComplex :: Ord a => [((a, Maybe Double), (a, Maybe Double), Double, Double)] -> ResultSimple
solveComplex gdata =
    let
        typeCheck ((_,c), (_,r), _, _) = (fromEnum . isNothing $ c) + ((2*) . fromEnum . isNothing $ r)

        pureStrategy :: Int -> Int -> [Double]
        pureStrategy l x = take l ((take (x-1) . repeat $ 0) ++ (1:[0..]))
        --  get all the subgame data and split it according to type: neither = Nothing, col = Nothing, row = Nothing, both = Nothing
        --  evs :: [[((Text, Maybe Double), (Text, Maybe Double), Double)]]
        evs = map (\i -> filter (\outcome -> typeCheck outcome == i) $ gdata) [0..3]
    in
        case evs of
            [both, [], [], []]  -> do -- both are fixed, so we just copy the weights into the ResultSimple
                let rows = nub . map (fst . fst4) $ both
                let cols = nub . map (fst . snd4) $ both
                let grid = transpose . chunksOf (length rows) $ both
                let outsc = map (map thd4) grid
                let outsr = map (map fth4) grid

                let gc = fromLists outsc
                let gr = fromLists outsr
                let cw = map (maybe (error "???") id . snd) . nub . map fst4 $ both
                let rw = map (maybe (error "???") id . snd) . nub . map snd4 $ both
                ResultSimple (calcEVSimpleCore gc cw rw) (calcEVSimpleCore gr cw rw) (calcSDSimpleCore gc cw rw) (calcSDSimpleCore gr cw rw) cw rw

            [[], c, [], []]      -> do -- typeCheck 1 -> columns unfixed, rows fixed, evaluate columns
                let cols = nub . map (fst . fst4) $ c -- each column option, in the form of name::Text
                let rows = map (\(a, Just b) -> (a,b)) . nub . map snd4 $ c -- each row option, in the form of (name, weight)::(Text, Double)
                let grid = transpose . chunksOf (length rows) $ c
                let outsc = map (map thd4) grid
                let outsr = map (map fth4) grid
                let gc = fromLists outsc
                let gr = fromLists outsr

                let colStrats = map (pureStrategy . length $ cols) [0..length cols] -- every pure strategy to try against the fixed row strategy
                let evs = map (\strat -> calcEVSimpleCore gc strat (map snd rows)) colStrats -- every ev

                let (cw, ev) = maximumBy (compare `on` snd) . zip colStrats $ evs -- get the optimal pure strategy for the column player
                let rw = map snd rows
                ResultSimple ev (calcEVSimpleCore gr cw rw) (calcSDSimpleCore gc cw rw) (calcSDSimpleCore gr cw rw) cw rw

            [[], [], r, []]      -> do -- the same as above, but cols fixed rows unfixed
                let cols = map (\(a, Just b) -> (a,b)) . nub . map fst4 $ r
                let rows = nub . map (fst . snd4) $ r
                let grid = transpose . chunksOf (length rows) $ r
                let outsc = map (map thd4) grid
                let outsr = map (map fth4) grid
                let gc = fromLists outsc
                let gr = fromLists outsr

                let rowStrats = map (pureStrategy . length $ rows) [0..length rows]
                let evs = map (\strat -> calcEVSimpleCore gr (map snd cols) strat) rowStrats

                let cw = map snd cols
                let (rw, ev) = minimumBy (compare `on` snd) . zip rowStrats $ evs -- the defender wants to minimise the ev
                ResultSimple (calcEVSimpleCore gc cw rw) ev (calcSDSimpleCore gc cw rw) (calcSDSimpleCore gr cw rw) cw rw

            [[], [], [], neither]     -> do -- both are unfixed, so we just calculate it as a normal GameSimple
                let cols = nub . map (fst . fst4) $ neither
                let rows = nub . map (fst . snd4) $ neither
                let grid = transpose . chunksOf (length rows) $ neither
                let outsc = map (map thd4) grid
                let outsr = map (map fth4) grid

                solveSimple outsc outsr

            _                      -> error "both fixed and unfixed options for a single player" -- no. don't. this makes no sense!!! the only reason i can even think you'd want to try is equivalent to just nesting in another mixup, so just do that.

fst4 (x,_,_,_) = x
snd4 (_,x,_,_) = x
thd4 (_,_,x,_) = x
fth4 (_,_,_,x) = x