{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE FlexibleContexts #-}

module SymbolicDistr where

import Control.Monad (replicateM)
import Control.Monad.Bayes.Class (MonadSample)

import Data.Proxy (Proxy (Proxy))
import GHC.TypeNats (KnownNat)

import Numeric.LinearAlgebra.Static (R, Sym, L)

import DelayedSampling hiding (sample, observe, score)
import qualified DelayedSampling as DelS
import qualified Distributions as D
import qualified MVDistributions as D
import DSProg
import Util.Ref (Heap, Ref, MonadState)

data Distr s = DeepForce s => Distr
  { sample :: forall m. MonadState Heap m => MonadSample m => m s
  , observe :: forall m. DelayedInfer m => Forced s -> m ()
  -- Score should not modify the heap, but currently we do not
  -- enforce this
  , score :: forall m. DelayedInfer m => Forced s -> m Double
  }

distrToSymDistr :: forall a. D.Distr a -> Distr (Expr a)
distrToSymDistr d = Distr (Const <$> D.sample d) (D.observe d)
  (pure . D.score d)

symDistrWithFallback  :: forall a. (forall m. MonadState Heap m => MonadSample m => m (D.Distr a))
  -> (forall m. MonadState Heap m => MonadSample m => m (Maybe (Expr a)))
  -> (forall m. DelayedInfer m => (a -> m (Maybe ())))
  -> (forall m. DelayedInfer m => (a -> m (Maybe Double)))
  -> Distr (Expr a)
symDistrWithFallback  d is iobs iscore = Distr is' iobs' iscore' where
  pd :: MonadState Heap m => MonadSample m => m (Distr (Expr a))
  pd = distrToSymDistr <$> d
  is' :: MonadState Heap m => MonadSample m => m (Expr a)
  is' = do
    mx <- is
    case mx of
      Nothing -> pd >>= sample
      Just x -> pure x
  iobs' :: DelayedInfer m => a -> m ()
  iobs' obs = do
    mupd <- iobs obs
    case mupd of
      Nothing -> do
        d <- pd
        observe d obs
      Just () -> pure ()
  iscore' :: DelayedInfer m => a -> m Double
  iscore' obs = do
    mupd <- iscore obs
    case mupd of
      Nothing -> do
        d <- pd
        score d obs
      Just ll -> pure ll

replicateIID :: forall a. DeepForce a => D.Distr Int -> Distr a -> Distr [a]
replicateIID howMany d = Distr sam obs scorer where
  sam :: MonadState Heap m => MonadSample m => m [a]
  sam = do
    n <- D.sample howMany
    replicateM n (sample d)
  obs :: DelayedInfer m => [Forced a] -> m ()
  obs xs = D.observe howMany (length xs) >> mapM_ (observe d) xs
  scorer :: DelayedInfer m => [Forced a] -> m Double
  scorer xs = do
    let ll = D.score howMany (length xs)
    ll' <- sum <$> mapM (score d) xs
    pure (ll + ll')

bind :: forall a b. a ~ Forced a => DeepForce a => DeepForce b => D.Distr a -> (a -> Distr b) -> Distr (a, b)
bind d f = Distr sam obs scorer where
  sam :: MonadState Heap m => MonadSample m => m (a, b)
  sam = do
    x <- D.sample d
    y <- sample (f x)
    pure (x, y)
  obs :: DelayedInfer m => (a, Forced b) -> m ()
  obs (x, y) = D.observe d x >> observe (f x) y
  scorer :: DelayedInfer m => (a, Forced b) -> m Double
  scorer (x, y) = (D.score d x +) <$> score (f x) y

indep :: forall a b. DeepForce a => DeepForce b => Distr a -> Distr b -> Distr (a, b)
indep dx dy = Distr ((,) <$> sample dx <*> sample dy)
                    (\(x, y) -> observe dx x >> observe dy y)
                    (\(x, y) -> (+) <$> score dx x <*> score dy y)

normal :: Expr Double -> Double -> Distr (Expr Double)
normal mu var = symDistrWithFallback (D.normal <$> force mu <*> pure var) is iobs iscore where
  is :: MonadState Heap m => MonadSample m => m (Maybe (Expr Double))
  is = do
      mu' <- realizedToConst mu
      case getAffine1 mu' of
        Nothing -> pure Nothing
        Just (b, mx) -> do
          case mx of
            Nothing -> do
              nref <- assumeConstant "" (MGaussian b var)
              Just <$> forgettableVar nref
            Just (RefNodeToT par, m) -> do
              ty <- typeOfRefNodeToT par
              case ty of
                SMGaussianT -> do
                  nref <- assumeConditional "" par (AffineMeanGaussian m b var)
                  Just <$> forgettableVar nref
                _ -> pure Nothing
  withPrior :: forall m x. MonadState Heap m => MonadSample m => (forall a. Double -> Double -> Ref (Node a MGaussianT) -> m x) -> m (Maybe x)
  withPrior f = do
    mu' <- realizedToConst mu
    case getAffine1 mu' of
      Nothing -> pure Nothing
      Just (b, Nothing) -> pure Nothing
      Just (b, Just (RefNodeToT par, m)) -> do
        ty <- typeOfRefNodeToT par
        case ty of
          SMGaussianT -> Just <$> f m b par
          _ -> pure Nothing
  iobs :: DelayedInfer m => Double -> m (Maybe ())
  iobs obs = withPrior $ \m b par -> observeConditional "" par (AffineMeanGaussian m b var) obs
  iscore :: DelayedInfer m => Double -> m (Maybe Double)
  iscore obs = withPrior $ \m b par -> scoreConditional "" par (AffineMeanGaussian m b var) obs

mvNormal :: forall n. KnownNat n => Expr (R n) -> Sym n -> Distr (Expr (R n))
mvNormal mu var = symDistrWithFallback (D.mvNormal <$> force mu <*> pure var) is iobs iscore where
  is :: MonadState Heap m => MonadSample m => m (Maybe (Expr (R n)))
  is = do
    mu' <- realizedToConst mu
    case getAffineMV1 mu' of
      Nothing -> pure Nothing
      Just (b, mx) -> do
        case mx of
          Nothing -> do
            nref <- assumeConstant "" (MMVGaussian b var)
            Just <$> forgettableVar nref
          Just (AffineMVMult (RefNodeToT par) m) -> do
            ty <- typeOfRefNodeToT par
            case ty of
              SMMVGaussianT -> do
                nref <- assumeConditional "" par (MVAffineMeanGaussian m b var)
                Just <$> forgettableVar nref
  withPrior :: forall m x. MonadState Heap m => MonadSample m => (forall a k (p :: Proxy k). KnownNat k => L n k -> R n -> Ref (Node a (MMVGaussianT p)) -> m x) -> m (Maybe x)
  withPrior f = do
    mu' <- realizedToConst mu
    case getAffineMV1 mu' of
      Nothing -> pure Nothing
      Just (b, Nothing) -> pure Nothing
      Just (b, Just (AffineMVMult (RefNodeToT par) m)) -> do
        ty <- typeOfRefNodeToT par
        case ty of
          SMMVGaussianT -> Just <$> f m b par
  iobs :: DelayedInfer m => R n -> m (Maybe ())
  iobs obs = withPrior $ \m b par -> observeConditional "" par (MVAffineMeanGaussian m b var) obs
  iscore :: DelayedInfer m => R n -> m (Maybe Double)
  iscore obs = withPrior $ \m b par -> scoreConditional "" par (MVAffineMeanGaussian m b var) obs

betaSD :: Double -> Double -> Distr (Expr Double)
betaSD a b = symDistrWithFallback  (pure $ D.beta a b) is iobs (const (pure Nothing)) where
  is :: MonadState Heap m => MonadSample m => m (Maybe (Expr Double))
  is = Just <$> (assumeConstant "" (MBeta a b) >>= forgettableVar)
  iobs obs = pure Nothing

bernoulliSD :: Expr Double -> Distr (Expr Bool)
bernoulliSD p = symDistrWithFallback (D.bernoulli <$> force p) is iobs iscore where
  withBetaPrior :: forall m x. MonadState Heap m => MonadSample m => (forall a. Ref (Node a MBetaT) -> m x) -> m (Maybe x)
  withBetaPrior f = case p of
    Var (RefNodeToT par) -> do
      ty <- typeOfRefNodeToT par
      case ty of
        SMBetaT -> do
          Just <$> f par
        _ -> pure Nothing
    _ -> pure Nothing
  is :: MonadState Heap m => MonadSample m => m (Maybe (Expr Bool))
  is = withBetaPrior $ \par ->
    forgettableVar =<< assumeConditional "" par CBernoulli
  iobs obs = withBetaPrior $ \par ->
    observeConditional "" par CBernoulli obs
  iscore obs = withBetaPrior $ \par ->
    scoreConditional "" par CBernoulli obs


{-# LANGUAGE BangPatterns #-}
module General.UserNumber
(GeneralSimpleNumber(..)
, ResolveException(..)
, ParseException(..)
, GRatio(..)
, sGR
, reduce
, GeneralRealNumber(..)
, GeneralComplex(..)
, GeneralNumber(..)
, PossNumber(..)
, numShow
, genExp
, realMod
, genRealtoFloat
, gReal
, imagUnit
, getIntegral
, getIntGen
, getFloat
, getRatio
) where

import qualified Data.Complex                     as DC
import           Data.Ratio
import qualified Data.Text.Lazy                   as T
import           Math.NumberTheory.Powers.Squares

data ResolveException = ResolveException T.Text
    deriving Show

data ParseException   = ParseException T.Text Integer
    deriving Show

-- Complex constructor: a:+ b, i.e. 1 + 1j = 1 :+ 1
-- Fraction constructor a % b, i.e. 1 / 2 = 1 % 2 = 2 % 4

--General Number = Either (Either Integer Fractional) (Either Float )

data GeneralSimpleNumber = GInt !Integer | GFlo !Float
  deriving (Show, Eq)

--Show for when printing out. Left as-is for debugging
--instance Show GeneralSimpleNumber where
--  show GInt n = show n
--  show GFlo n = show n

instance Num GeneralSimpleNumber where
  a + b
    | GInt x <- a, GInt y <- b = GInt $ x + y
    | GInt x <- a, GFlo y <- b = GFlo $ fromIntegral x + y
    | GFlo x <- a, GInt y <- b = GFlo $ x + fromIntegral y
    | GFlo x <- a, GFlo y <- b = GFlo $ x + y

  a * b
    | GInt x <- a, GInt y <- b = GInt $ x * y
    | GInt x <- a, GFlo y <- b = GFlo $ fromIntegral x * y
    | GFlo x <- a, GInt y <- b = GFlo $ x * fromIntegral y
    | GFlo x <- a, GFlo y <- b = GFlo $ x * y

  negate (GInt a) = GInt (negate a)
  negate (GFlo a) = GFlo (negate a)

  abs (GInt a) = GInt (abs a)
  abs (GFlo a) = GFlo (abs a)

  signum a
    | GInt x <- a = GInt $ signum x
    | GFlo x <- a, signum x < 0.0 = GInt (-1)
    | GFlo x <- a, signum x > 0.0 = GInt 1
    | GFlo _ <- a, True           = GInt 0

  fromInteger = GInt

instance Ord GeneralSimpleNumber where
  a <= b
    | GInt x <- a, GInt y <- b = x <= y
    | GInt x <- a, GFlo y <- b = fromIntegral x <= y
    | GFlo x <- a, GInt y <- b = x <= fromIntegral y
    | GFlo x <- a, GFlo y <- b = x <= y

data GRatio = GR !Integer !Integer deriving (Show)

instance Eq GRatio where
  GR a b == GR c d = a * d == b * c

instance Num GRatio where
  GR a b + GR c d = ratReduce (a*d+b*c) (b*d)

  GR a b * GR c d = ratReduce (a*b) (c*d)

  negate (GR a b) = GR (negate a) b

  abs (GR a b) = GR (abs a) b

  signum (GR a _) = GR (signum a) 1

  fromInteger n = GR (fromInteger n) 1

instance Fractional GRatio where
  recip (GR a b) = GR b a

  fromRational n = GR (numerator n) (denominator n)

instance Real GRatio where
  toRational (GR a b) = a % b

instance RealFrac GRatio where
  properFraction (GR a b) = (fromInteger $ a `div` b, GR (a `mod` b) b)

instance Ord GRatio where
  GR a b <= GR c d = a*d <= b*c

sGR :: Integer -> Integer -> GRatio
sGR a b = ratReduce (a * signum b) (abs b)

ratReduce :: Integer -> Integer -> GRatio
ratReduce !a !b = let d = gcd a b
               in GR (quot a d) (quot b d)

data GeneralRealNumber = GSimp GeneralSimpleNumber | GRat GRatio deriving (Show, Eq)

reduce :: GRatio -> GeneralRealNumber
reduce r
  | GR a 1 <- r = GSimp (GInt a)
  | otherwise = GRat r

instance Num GeneralRealNumber where
  a + b
    | GSimp x <- a, GSimp y <- b = GSimp $ x + y
    | GSimp (GInt x) <- a, GRat (GR y z) <- b = reduce $ ratReduce (x*z+y) z
    | GSimp (GFlo x) <- a, GRat (GR y z) <- b = GSimp $ GFlo $ x+fromIntegral y/fromIntegral z
    | GRat x <- a, GRat y <- b = reduce $ x + y
    | otherwise = b + a

  a * b
    | GSimp x <- a, GSimp y <- b = GSimp $ x * y
    | GSimp (GInt x) <- a, GRat (GR y z) <- b = reduce $ ratReduce (x*y) z
    | GSimp (GFlo x) <- a, GRat (GR y z) <- b = GSimp $ GFlo $ x*fromIntegral y/fromIntegral z
    | GRat x <- a, GRat y <- b = reduce $ x * y
    | otherwise = b * a

  negate (GSimp a) = GSimp (negate a)
  negate (GRat a)  = GRat (negate a)

  abs (GSimp a) = GSimp (abs a)
  abs (GRat a)  = GRat (abs a)

  signum (GSimp a)       = GSimp (signum a)
  signum (GRat (GR a _)) = GSimp $ GInt $ signum a

  fromInteger n = GSimp (fromInteger n)

instance Fractional GeneralRealNumber where
  recip a
    |GSimp (GInt x) <- a = GRat $ sGR 1 x
    |GSimp (GFlo x) <- a = GSimp $ GFlo $ recip x
    |GRat x <- a = reduce $ recip x

  fromRational n = GRat $ sGR (numerator n) (denominator n)

instance Real GeneralRealNumber where
  toRational (GSimp (GInt x)) = x % 1
  toRational (GSimp (GFlo x)) = toRational x
  toRational (GRat r)         = toRational r

fracHelper :: (RealFrac a, Integral b) => a -> (a -> c) -> (b, c)
fracHelper inner fun = (int, fun frac)
  where
    (int, frac) = properFraction inner

instance RealFrac GeneralRealNumber where
  properFraction (GSimp (GInt x)) = (fromInteger x, GSimp $ GInt 0)
  properFraction (GSimp (GFlo x)) = fracHelper x (GSimp . GFlo)
  properFraction (GRat r)         = fracHelper r GRat

instance Ord GeneralRealNumber where
  GSimp a <= GSimp b = a <= b
  GSimp (GInt a) <= GRat (GR b c)  = a*c <= b
  GSimp (GFlo a) <= GRat (GR b c)  = a * fromInteger c <= fromInteger b
  GRat (GR b c)  <= GSimp (GInt a) = b <= a*c
  GRat (GR b c)  <= GSimp (GFlo a) = fromInteger b <= a * fromInteger c
  GRat a <= GRat b = a <= b

intPow :: GeneralRealNumber -> Integer -> GeneralRealNumber
intPow a n
  | n < 0 = intPow (recip a) (negate n)
  | n == 0 = GSimp $ GInt 1
  | otherwise = a * intPow a (n-1)

expAlgR :: GeneralRealNumber -> Integer -> GeneralRealNumber
expAlgR _ 0 = GSimp $ GInt 1
expAlgR x 1 = x
expAlgR x n
  |n < 0 = expAlgR (recip x) (negate n)
  |even n = expAlgR (x*x) (n `div` 2)
  |otherwise = x * expAlgR (x*x) (n `div` 2)

realExp :: GeneralRealNumber -> GeneralRealNumber -> GeneralRealNumber
realExp _ (GSimp (GInt 0)) = GSimp $ GInt 1
realExp a@(GSimp (GInt _)) (GSimp (GInt y)) = expAlgR a y
realExp (GSimp (GInt x)) (GSimp (GFlo y)) = GSimp $ GFlo $ fromIntegral x ** y
realExp (GSimp (GFlo x)) (GSimp (GInt y)) = GSimp $ GFlo $ x ** fromIntegral y
realExp (GSimp (GFlo x)) (GSimp (GFlo y)) = GSimp $ GFlo $ x**y
realExp r@(GRat (GR _ _)) (GSimp (GInt n)) = expAlgR r n
realExp (GRat (GR a b)) (GSimp c) = realExp (GSimp $ GFlo $ fromInteger a/fromInteger b) (GSimp c)
realExp (GSimp c) (GRat (GR a b)) = realExp (GSimp c) (GSimp $ GFlo $ fromInteger a/fromInteger b)
realExp (GRat (GR a b)) (GRat (GR c d)) = realExp (GSimp $ GFlo $ fromInteger a/fromInteger b) (GSimp $ GFlo $ fromInteger c/fromInteger d)

generalRSqrt :: GeneralRealNumber -> GeneralRealNumber
generalRSqrt (GSimp (GInt a))
  |isSquare a = GSimp $ GInt $ integerSquareRoot a
  |otherwise = GSimp $ GFlo $ sqrt $ fromIntegral a
generalRSqrt (GSimp (GFlo a)) = GSimp $ GFlo $ sqrt a
generalRSqrt (GRat (GR a b))
  |isSquare a && isSquare b = GRat $ GR (integerSquareRoot a) (integerSquareRoot b)
  |otherwise = GSimp $ GFlo $ sqrt $ fromIntegral a/fromIntegral b

genRealtoFloat :: GeneralRealNumber -> Float
genRealtoFloat (GSimp (GInt a)) = fromInteger a
genRealtoFloat (GSimp (GFlo a)) = a
genRealtoFloat (GRat (GR a b))  = fromInteger a/fromInteger b

data GeneralComplex = GC !GeneralRealNumber !GeneralRealNumber
  deriving (Show)

instance Eq GeneralComplex where
  GC a b == GC c d = (a==c) && (b==d)

realPart :: GeneralComplex -> GeneralRealNumber
realPart (GC a _) = a

imagPart :: GeneralComplex -> GeneralRealNumber
imagPart (GC _ b) = b

instance Num GeneralComplex where
  (GC a b) + (GC c d) = GC (a+c) (b+d)

  GC a b * GC c d = GC (a*c-b*d) (a*d+b*c)

  negate (GC a b) = GC (negate a) (negate b)

  abs (GC a b) = GC (generalRSqrt ((a `intPow` 2)+(b `intPow` 2))) (GSimp (GInt 0))

  signum (GC a b) = GC (a / mag) (b / mag)
    where
      mag = realPart $ abs $ GC a b

  fromInteger n = GC (GSimp (GInt n)) 0

instance Fractional GeneralComplex where
  recip (GC a b) = GC (a/((a `intPow` 2)+(b `intPow` 2))) (-b/((a `intPow` 2)+(b `intPow` 2)))

  fromRational n = GC (reduce (ratReduce (numerator n) (denominator n))) (GSimp (GInt 0))

data GeneralNumber = GReal GeneralRealNumber | GComp GeneralComplex
  deriving (Eq)

simpComplex :: GeneralComplex -> GeneralNumber
simpComplex c@(GC a b)
    | b == 0 || b == 0.0 = GReal a
    | otherwise = GComp c

instance Num GeneralNumber where
  a + b
    | GReal x <- a, GReal y <- b = GReal $ x + y
    | GReal x <- a, GComp (GC y z) <- b = simpComplex $ GC (x+y) z
    | GComp x <- a, GComp y <- b = simpComplex $ x + y
    | otherwise = b + a

  a * b
    | GReal x <- a, GReal y <- b = GReal $ x * y
    | GReal x <- a, GComp (GC y z) <- b = simpComplex $ GC (x*y) $ x*z
    | GComp x <- a, GComp y <- b = simpComplex $ x * y
    | otherwise = b * a

  negate (GReal a) = GReal (negate a)
  negate (GComp a) = GComp (negate a)

  abs (GReal a) = GReal (abs a)
  abs (GComp a) = simpComplex (abs a)

  signum (GReal a) = GReal (signum a)
  signum (GComp a) = GComp (signum a)

  fromInteger n = GReal $ GSimp $ GInt n

expAlgC :: GeneralComplex -> Integer -> GeneralNumber
expAlgC _ 0 = GReal $ GSimp $ GInt 1
expAlgC x 1 = GComp x
expAlgC x n
  |n < 0 = expAlgC (recip x) (negate n)
  |even n = expAlgC (x*x) (n `div` 2)
  |otherwise = GComp x * expAlgC (x*x) (n `div` 2)

instance Fractional GeneralNumber where
  recip a
    |GReal x <- a = GReal $ recip x
    |GComp x <- a = GComp $ recip x

  fromRational n = GReal $ fromRational n

genExp :: GeneralNumber -> GeneralNumber -> GeneralNumber
genExp (GReal n) (GReal m) = GReal $ realExp n m
genExp (GComp a) (GReal (GSimp (GInt n))) = expAlgC a n
genExp (GComp a) (GReal (GSimp (GFlo flo))) = GComp (GC (GSimp $ GFlo $ DC.realPart builtInPow) (GSimp $ GFlo $ DC.imagPart builtInPow))
  where
    buildInVer = genRealtoFloat (realPart a) DC.:+ genRealtoFloat (imagPart a)
    builtInPow = buildInVer ** (flo DC.:+ 0)
genExp a@(GComp _) (GReal r@(GRat _)) = genExp a (GReal $ GSimp $ GFlo $ genRealtoFloat r)
genExp (GComp (GC a b)) (GComp (GC c d)) = GComp $ GC (GSimp $ GFlo $ DC.realPart passed) $ GSimp $ GFlo $ DC.imagPart passed
  where
    passed = (genRealtoFloat a DC.:+ genRealtoFloat b) ** (genRealtoFloat c DC.:+ genRealtoFloat d)
genExp (GReal n) j = genExp (GComp (GC n 0)) j

realMod :: GeneralRealNumber -> GeneralRealNumber -> GeneralRealNumber
realMod a (GSimp (GInt 0)) = a
realMod a b = a - (b * f)
  where
    realfloor :: GeneralRealNumber -> Integer
    realfloor (GSimp (GInt x)) = x
    realfloor (GSimp (GFlo x)) = floor x
    realfloor (GRat (GR x y))  = floor $ genRealtoFloat $ GRat $ GR x y
    f = GSimp $ GInt $ realfloor $ a/b

realShow :: GeneralRealNumber -> String
realShow num
  |GRat (GR x y) <- num = show x ++ "/" ++ show y
  |GSimp (GInt x) <- num = show x
  |GSimp (GFlo x) <- num = show x

numShow :: GeneralNumber -> String
numShow num
  |GComp (GC (GSimp (GInt 0)) y) <- num = realShow y ++ "j"
  |GComp (GC x y) <- num = realShow x ++ "+" ++ realShow y ++ "j"
  |GReal x <- num = realShow x

instance Show GeneralNumber where
  show = numShow

data PossNumber = Kept GeneralNumber
                | Dropped GeneralNumber
  deriving (Eq)

instance Show PossNumber where
  show (Kept n)    = show n
  show (Dropped n) = "//" ++ show n ++ "//"

gReal :: (Integral a) => a -> GeneralNumber
gReal n = GReal $ GSimp $ GInt $ fromIntegral n

imagUnit = GComp $ GC (GSimp $ GInt 0) (GSimp $ GInt 1)

getIntegral :: (Integral a) => GeneralNumber -> Maybe a
getIntegral (GReal (GSimp (GInt n))) = Just $ fromInteger n
getIntegral _                        = Nothing

getIntGen :: (Integer -> a) -> GeneralNumber -> Maybe a
getIntGen p (GReal (GSimp (GInt n))) = Just $ p n
getIntGen _ _                        = Nothing

getFloat :: GeneralNumber -> Maybe Float
getFloat (GReal (GSimp (GFlo f))) = Just f
getFloat _                        = Nothing

getRatio :: GeneralNumber -> Maybe Rational
getRatio (GReal (GRat (GR p q))) = Just $ p % q
getRatio _                       = Nothing


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE Strict                #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Layers.Relu
Description : Rectifying linear unit layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Relu (
    Relu (..)
  , SpecRelu (..)
  , specRelu1D
  , specRelu2D
  , specRelu3D
  , relu
  ) where

import           Control.Parallel.Strategies
import           Data.Constraint                (Dict (..))
import           Data.Reflection                (reifyNat)
import           Data.Serialize
import           Data.Singletons
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra.Static   as LAS
import           Unsafe.Coerce                  (unsafeCoerce)

import           Grenade.Core
import           Grenade.Dynamic
import           Grenade.Dynamic.Internal.Build
import           Grenade.Utils.Conversion
import           Grenade.Utils.Vector

-- | A rectifying linear unit.
--   A layer which can act between any shape of the same dimension, acting as a
--   diode on every neuron individually.
data Relu = Relu
  deriving (Generic, NFData, Show)

instance UpdateLayer Relu where
  type Gradient Relu = ()
  runUpdate _ _ _ = Relu

instance RandomLayer Relu where
  createRandomWith _ _ = return Relu

instance Serialize Relu where
  put _ = return ()
  get = return Relu

-- Run forward 1D
runForward1D :: (KnownNat i) => Relu -> S ('D1 i) -> (Tape Relu ('D1 i) ('D1 i), S ('D1 i))
runForward1D _ (S1D y) = (S1D y, S1D (relu y))
    where
      relu = LAS.dvmap (\a -> if a <= 0 then 0 else a)
runForward1D _ (S1DV y) =
  (S1DV y, S1DV (mapVectorInPlace relu y)) -- y will be replaced, which however still leads to a correct implementation!
  -- force (mapVectorInPlace relu y) `seq` (S1DV y, S1DV y) -- y will be replaced, which however still leads to a correct implementation!
  -- (S1DV y, S1DV $ mapVector relu y) -- y will be replaced, which however still leads to a correct implementation!
    where
      relu a = if a <= 0 then 0 else a

reluDifZip :: (Ord a, Floating a) => a -> a -> a
reluDifZip a g = if a <= 0 then 0 else g

-- Run backward 1D
runBackward1D :: (KnownNat i) => Relu -> S ('D1 i) -> S ('D1 i) -> (Gradient Relu, S ('D1 i))
runBackward1D _ (S1D y) (S1D dEdy) = ((), S1D (relu' y * dEdy))
    where
      relu' = LAS.dvmap (\a -> if a <= 0 then 0 else 1)
runBackward1D _ (S1DV y) (S1DV dEdy) = ((), zipWithVectorInPlaceSnd reluDifZip y dEdy `seq` S1DV dEdy)
runBackward1D l y dEdy = runBackward1D l y (toLayerShape y dEdy)

instance (KnownNat i) => Layer Relu ('D1 i) ('D1 i) where
  type Tape Relu ('D1 i) ('D1 i) = S ('D1 i)
  runForwards = runForward1D
  runBackwards = runBackward1D

runBackward2D :: (KnownNat i, KnownNat j) => Relu -> S ('D2 i j) -> S ('D2 i j) -> (Gradient Relu, S ('D2 i j))
runBackward2D _ (S2D y) (S2D dEdy) = ((), S2D (relu' y * dEdy))
    where
      relu' = LAS.dmmap (\a -> if a <= 0 then 0 else 1)
runBackward2D _ (S2DV y) (S2DV dEdy) = ((), zipWithVectorInPlaceSnd reluDifZip y dEdy `seq` S2DV dEdy)
runBackward2D l y dEdy = runBackward2D l y (toLayerShape y dEdy)

instance (KnownNat i, KnownNat j) => Layer Relu ('D2 i j) ('D2 i j) where
  type Tape Relu ('D2 i j) ('D2 i j) = S ('D2 i j)

  runForwards _ (S2D y) = (S2D y, S2D (relu y))
    where
      relu = LAS.dmmap (\a -> if a <= 0 then 0 else a)
  runForwards _ (S2DV y) = (S2DV y, S2DV $ mapVectorInPlace relu y ) -- This replaces the y vector, but that doesn't harm the update
    where
      relu a = if a <= 0 then 0 else a
  runBackwards = runBackward2D

instance (KnownNat i, KnownNat j, KnownNat k) => Layer Relu ('D3 i j k) ('D3 i j k) where

  type Tape Relu ('D3 i j k) ('D3 i j k) = S ('D3 i j k)

  runForwards _ (S3D y) = (S3D y, S3D (relu y))
    where
      relu = LAS.dmmap (\a -> if a <= 0 then 0 else a)
  runBackwards _ (S3D y) (S3D dEdy) = ((), S3D (relu' y * dEdy))
    where
      relu' = LAS.dmmap (\a -> if a <= 0 then 0 else 1)

-------------------- DynamicNetwork instance --------------------

instance FromDynamicLayer Relu where
  fromDynamicLayer inp _ _ = SpecNetLayer $ SpecRelu (tripleFromSomeShape inp)

instance ToDynamicLayer SpecRelu where
  toDynamicLayer _ _ (SpecRelu (rows, cols, depth)) =
     reifyNat rows $ \(_ :: (KnownNat rows) => Proxy rows) ->
     reifyNat cols $ \(_ :: (KnownNat cols) => Proxy cols) ->
     reifyNat depth $ \(_ :: (KnownNat depth) => Proxy depth) ->
     case (rows, cols, depth) of
         (_, 1, 1)    -> return $ SpecLayer Relu (sing :: Sing ('D1 rows)) (sing :: Sing ('D1 rows))
         (_, _, 1) -> return $ SpecLayer Relu (sing :: Sing ('D2 rows cols)) (sing :: Sing ('D2 rows cols))
         _    -> case (unsafeCoerce (Dict :: Dict()) :: Dict (KnownNat (rows GHC.TypeLits.* depth))) of
           Dict -> return $ SpecLayer Relu (sing :: Sing ('D3 rows cols depth)) (sing :: Sing ('D3 rows cols depth))

-- | Create a specification for a elu layer.
specRelu1D :: Integer -> SpecNet
specRelu1D i = specRelu3D (i, 1, 1)

-- | Create a specification for a elu layer.
specRelu2D :: (Integer, Integer) -> SpecNet
specRelu2D (i, j) = specRelu3D (i, j, 1)

-- | Create a specification for a elu layer.
specRelu3D :: (Integer, Integer, Integer) -> SpecNet
specRelu3D = SpecNetLayer . SpecRelu

-- | Add a Relu layer to your build.
relu :: BuildM ()
relu = buildGetLastLayerOut >>= buildAddSpec . SpecNetLayer . SpecRelu

-------------------- GNum instances --------------------

instance GNum Relu where
  _ |* Relu = Relu
  _ |+ Relu = Relu
  sumG _ = Relu


{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE IncoherentInstances   #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeSynonymInstances  #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# OPTIONS -fno-warn-orphans #-}

module Data.Array.Accelerate.Data.Complex (

  -- * Rectangular from
  Complex(..),
  real,
  imag,

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

  -- * Conjugate
  conjugate,

) where

import Prelude
import Data.Complex                             ( Complex(..) )
import Data.Array.Accelerate                    as A
import Data.Array.Accelerate.Smart
import Data.Array.Accelerate.Product
import Data.Array.Accelerate.Array.Sugar

type instance EltRepr (Complex a) = EltRepr (a, a)

instance Elt a => Elt (Complex a) where
  eltType _             = eltType (undefined :: (a,a))
  toElt p               = let (a, b) = toElt p in a :+ b
  fromElt (a :+ b)      = fromElt (a, b)

instance cst a => IsProduct cst (Complex a) where
  type ProdRepr (Complex a) = ProdRepr (a, a)
  fromProd cst (x :+ y) = fromProd cst (x, y)
  toProd cst p          = let (x, y) = toProd cst p in (x :+ y)
  prod cst _            = prod cst (undefined :: (a, a))

instance (Lift Exp a, Elt (Plain a)) => Lift Exp (Complex a) where
  type Plain (Complex a) = Complex (Plain a)
  lift (x1 :+ x2)       = Exp $ Tuple (NilTup `SnocTup` lift x1 `SnocTup` lift x2)

instance Elt a => Unlift Exp (Complex (Exp a)) where
  unlift e
    = let x     = Exp $ SuccTupIdx ZeroTupIdx `Prj` e
          y     = Exp $ ZeroTupIdx `Prj` e
      in
      x :+ y

instance (Elt a, IsFloating a) => Num (Exp (Complex a)) where
  (+)           = lift2 ((+) :: Complex (Exp a) -> Complex (Exp a) -> Complex (Exp a))
  (-)           = lift2 ((-) :: Complex (Exp a) -> Complex (Exp a) -> Complex (Exp a))
  (*)           = lift2 ((*) :: Complex (Exp a) -> Complex (Exp a) -> Complex (Exp a))
  negate        = lift1 (negate :: Complex (Exp a) -> Complex (Exp a))
  signum        = lift1 (signum :: Complex (Exp a) -> Complex (Exp a))
  abs           = lift1 (abs :: Complex (Exp a) -> Complex (Exp a))
  fromInteger n = lift (constant (fromInteger n) :+ 0)

instance (Elt a, IsFloating a) => Fractional (Exp (Complex a)) where
  c / c'
    = let x  :+ y       = unlift c
          x' :+ y'      = unlift c'     :: Complex (Exp a)
          den           = x'^(2 :: Int) + y'^(2 :: Int)
          re            = (x * x' + y * y') / den
          im            = (y * x' - x * y') / den
      in
      lift (re :+ im)

  fromRational x
    = lift (constant (fromRational x) :+ constant 0)

instance (Elt a, IsFloating a, RealFloat a) => Floating (Exp (Complex a)) where
  sqrt z
    = let
          x :+ y        = unlift z
          v'            = abs y / (u'*2)
          u'            = sqrt ((magnitude z + abs x) / 2)
          (u, v)        = unlift ( x A.<* 0 ? ( lift (v',u'), lift (u',v') ) )
      in
      x ==* 0 &&* y ==* 0 ?
        {- then -} ( 0
        {- else -} , lift (u :+ (y A.<* 0 ? (-v,v))) )

  pi            = lift (pi :+ constant 0)
  log z         = lift (log (magnitude z) :+ phase z)
  exp           = lift1 (exp :: Complex (Exp a) -> Complex (Exp a))
  sin           = lift1 (sin :: Complex (Exp a) -> Complex (Exp a))
  cos           = lift1 (cos :: Complex (Exp a) -> Complex (Exp a))
  tan           = lift1 (tan :: Complex (Exp a) -> Complex (Exp a))
  sinh          = lift1 (sinh :: Complex (Exp a) -> Complex (Exp a))
  cosh          = lift1 (cosh :: Complex (Exp a) -> Complex (Exp a))
  tanh          = lift1 (tanh :: Complex (Exp a) -> Complex (Exp a))
  asin          = lift1 (asin :: Complex (Exp a) -> Complex (Exp a))
  acos          = lift1 (acos :: Complex (Exp a) -> Complex (Exp a))
  atan          = lift1 (atan :: Complex (Exp a) -> Complex (Exp a))
  asinh         = lift1 (asinh :: Complex (Exp a) -> Complex (Exp a))
  acosh         = lift1 (acosh :: Complex (Exp a) -> Complex (Exp a))
  atanh         = lift1 (atanh :: Complex (Exp a) -> Complex (Exp a))

-- | The non-negative magnitude of a complex number
--
magnitude :: (Elt a, IsFloating a) => Exp (Complex a) -> Exp a
magnitude c =
  let r :+ i    = unlift c
  in sqrt (r*r + i*i)

-- | The phase of a complex number, in the range @(-'pi', 'pi']@. If the
-- magnitude is zero, then so is the phase.
--
phase :: (Elt a, IsFloating a) => Exp (Complex a) -> Exp a
phase c =
  let x :+ y    = unlift c
  in atan2 y x

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

-- | Form a complex number from polar components of magnitude and phase.
--
mkPolar :: (Elt a, IsFloating a) => Exp a -> Exp a -> Exp (Complex a)
mkPolar r theta  =  lift $ r * cos theta :+ r * sin theta

-- | @'cis' t@ is a complex value with magnitude @1@ and phase @t@ (modulo
-- @2*'pi'@).
--
cis :: (Elt a, IsFloating a) => Exp a -> Exp (Complex a)
cis theta = lift $ cos theta :+ sin theta

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

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

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


-- | https://en.wikibooks.org/wiki/Write_Yourself_a_Scheme_in_48_Hours/Parsing

module Main where

import           Control.Monad
import           Control.Monad.Error
import           Data.Char                     (digitToInt, isDigit)
import           Data.Complex
import           Data.Ratio                    (denominator, numerator, (%))
import           Numeric                       (readDec, readFloat, readHex,
                                                readInt, readOct)
import           System.Environment
import           Text.ParserCombinators.Parsec hiding (spaces)

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

-- readExpr :: String -> Either ParseError LispVal
-- readExpr = parse parseExpr "lisp"

-- Whitespace
spaces :: Parser ()
spaces = skipMany1 space

-- Return Values
data LispVal = Atom String
             | List [LispVal]
             | DottedList [LispVal] LispVal
             | Number Integer
             | Float Double
             | Ratio Rational
             | Complex (Complex Double)
             | Character Char
             | String String
             | Bool Bool

parseString :: Parser LispVal
parseString = do
  char '"'
  str <- many (parseHelper <|> noneOf "\\\"")
  char '"'
  return $ String str
  <?> "a string for parseString"
  where parseHelper :: Parser Char
        parseHelper = do
          char '\\'
          x <- oneOf "\\\"nrt"
          return $ case x of
            '\\' -> x
            '"'  -> x
            'n'  -> '\n'
            'r'  -> '\r'
            't'  -> '\t'

parseAtom :: Parser LispVal
parseAtom = do
  first <- letter <|> symbol
  rest <- many $ letter <|> symbol <|> digit
  return $ Atom (first : rest)
  <?> "a atom for parseAtom"

parseBool :: Parser LispVal
parseBool = do
  x <- try (string "#t") <|> string "#f"
  return $ if x == "#t"
           then Bool True
           else Bool False

parseNumber :: Parser LispVal
parseNumber = parseDigit <|> try parseOct <|> try parseDec <|> try parseHex <|> parseBin
  where parseOct = do
          char '#' >> oneOf "oO"
          str <- many1 octDigit
          let [(n, _)] = readOct str
          return $ Number n
        parseDec = do
          char '#' >> oneOf "dD"
          str <- many1 digit
          let [(n, _)] = readDec str
          return $ Number n
        parseHex = do
          char '#' >> oneOf "xX"
          str <- many1 hexDigit
          let [(n, _)] = readHex str
          return $ Number n
        parseBin = do
          char '#' >> oneOf "bB"
          str <- many1 $ oneOf "01"
          let [(n, _)] = readInt 2 (`elem` "01") digitToInt str
          return $ Number n
        parseDigit = do
          str <- many1 digit
          return $ (Number . read) str

parseChar :: Parser LispVal
parseChar = do
  string "#\\"
  x <- many $ letter <|> digit
  return $ case length x of
    0 -> Character '\n'
    1 -> Character $ head x
    _ -> case x of
      "newline" -> Character '\n'
      "tab"     -> Character '\t'
      "space"   -> Character ' '

parseFloat :: Parser LispVal
parseFloat = do
  x <- many1 digit
  char '.'
  y <- many1 digit
  return $ let [(n, _)] = readFloat (x ++ '.' : y) in Float n

parseRatio :: Parser LispVal
parseRatio = do
  x <- many1 digit
  char '/'
  y <- many1 digit
  return $ Ratio (read x % read y)

parseComplex :: Parser LispVal
parseComplex = do
  x <- try parseFloat <|> parseNumber
  char '+'
  y <- try parseFloat <|> parseNumber
  char 'i'
  return $ Complex (toDouble x :+ toDouble y)
  where toDouble :: LispVal -> Double
        toDouble (Float f) = f
        toDouble (Number n) = fromIntegral n

-- Recursive Parsers: Adding lists, dotted lists, and quoted datums

parseList :: Parser LispVal
parseList = do
  char '('
  x <- sepBy parseExpr spaces
  char ')'
  return $ List x

parseDottedList :: Parser LispVal
parseDottedList = do
  char '('
  x <- endBy parseExpr spaces
  y <- char '.' >> spaces >> parseExpr
  char ')'
  return $ DottedList x y

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

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

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

parseUnquoteSplicing :: Parser LispVal
parseUnquoteSplicing = do
  char ',' >> char '@'
  y <- parseExpr
  return $ List [Atom "unquote-splicing", y]

parseExpr :: Parser LispVal
parseExpr = parseAtom
            <|> try parseBool
            <|> try parseRatio
            <|> try parseComplex
            <|> try parseFloat
            <|> try parseNumber
            <|> parseChar
            <|> parseString
            <|> try parseList
            <|> parseDottedList
            <|> parseQuoted
            <|> parseBackquote
            <|> try parseUnquoteSplicing
            <|> parseUnquote

-- Beginning the Evaluator
showVal :: LispVal -> String
showVal (Atom x) = x
showVal (Bool True) = "#t"
showVal (Bool False) = "#f"
showVal (Ratio x) = show (numerator x) ++ ('/' : show (denominator x))
showVal (Complex x) = show (realPart x) ++ ('+' : show (imagPart x) ++ "i")
showVal (Float x) = show x
showVal (Number x) = show x
showVal (Character x) = "#\\" ++ show x
showVal (String x) = "\"" ++ x ++ "\""
showVal (List x) = "(" ++ unwordsList x ++ ")"
showVal (DottedList first rest)= "(" ++ unwordsList first ++ " . " ++ showVal rest ++ ")"

unwordsList :: [LispVal] -> String
unwordsList = unwords . map showVal

instance Show LispVal where
  show = showVal

readExpr :: String -> ThrowsError LispVal
readExpr input = case parse parseExpr "lisp" input of
  Left err -> throwError $ Parser err
  Right val -> return val

-- Beginnings of an evaluator: Primitives
-- main :: IO ()
-- main = getArgs >>= print . eval . readExpr . head
main :: IO ()
main = do
  args <- getArgs
  evaled <- return $ liftM show $ readExpr (head args) >>= eval
  putStrLn $ extractValue $ trapError evaled

eval :: LispVal -> ThrowsError LispVal
eval val@(String _) = return val
eval val@(Bool _) = return val
eval val@(Ratio _) = return val
eval val@(Complex _) = return val
eval val@(Float _) = return val
eval val@(Number _) = return val
eval val@(Character _) = return val
eval (List [Atom "quote", val]) = return val

-- Adding basic primitives
eval (List (Atom func : args)) = mapM eval args >>= apply func

apply :: String -> [LispVal] -> ThrowsError LispVal
apply func args = maybe (throwError $ NotFunction "No such function: " func)
                  ($ args)
                  (lookup func primitives)

primitives :: [(String, [LispVal] -> ThrowsError LispVal)]
primitives = [("+", numericBinop (+)),
              ("-", numericBinop (-)),
              ("*", numericBinop (*)),
              ("/", numericBinop div),
              ("mod", numericBinop mod),
              ("quotient", numericBinop quot),
              ("remainder", numericBinop rem),
              ("symbol?", unaryOp symbolp),
              ("string?", unaryOp stringp),
              ("number?", unaryOp numberp),
              ("bool?", unaryOp boolp),
              ("list?", unaryOp listp),
              ("pair?", unaryOp pairp),
              ("symbol->string", unaryOp symbol2String),
              ("string->symbol", unaryOp string2Symbol)]

numericBinop :: (Integer -> Integer -> Integer) -> [LispVal] -> ThrowsError LispVal
numericBinop _ [] = throwError $ NumArgs 2 []
numericBinop _ singleVal@[_] = throwError $ NumArgs 2 singleVal
numericBinop op args = mapM unpackNum args >>= return . Number . foldl1 op
  where unpackNum :: LispVal -> ThrowsError Integer
        unpackNum (Number n) = return n
        -- unpackNum (String n) = let parsed = readDec n in
          -- if null parsed
          -- then 0
          -- else fst $ head parsed
        -- unpackNum (List [n]) = unpackNum n
        unpackNum notNum = throwError $ TypeMismatch "number" notNum

unaryOp :: (LispVal -> ThrowsError LispVal) -> [LispVal] -> ThrowsError LispVal
unaryOp op [x] = op x
unaryOp _ wrongArg = throwError $ NumArgs 1 wrongArg

symbolp :: LispVal -> ThrowsError LispVal
symbolp (Atom _) = return $ Bool True
symbolp _ = return $ Bool False

stringp :: LispVal -> ThrowsError LispVal
stringp (String _) = return $ Bool True
stringp _ = return $ Bool False

numberp :: LispVal -> ThrowsError LispVal
numberp (Number _) = return $ Bool True
numberp _ = return $ Bool False

boolp :: LispVal -> ThrowsError LispVal
boolp (Bool _) = return $ Bool True
boolp _ = return $ Bool False

listp :: LispVal -> ThrowsError LispVal
listp (List _) = return $ Bool True
listp _ = return $ Bool False

pairp :: LispVal -> ThrowsError LispVal
pairp (List _) = return $ Bool True
pairp (DottedList _ _) = return $ Bool True
pairp _ = return $ Bool False

symbol2String :: LispVal -> ThrowsError LispVal
symbol2String (Atom x) = return $ String x
symbol2String mismatch = throwError $ TypeMismatch "atom" mismatch

string2Symbol :: LispVal -> ThrowsError LispVal
string2Symbol (String x) = return $ Atom x
string2Symbol mismatch = throwError $ TypeMismatch "string" mismatch

-- Error Checking and Exceptions
data LispError = NumArgs Integer [LispVal]
               | TypeMismatch String LispVal
               | Parser ParseError
               | BadSpecialForm String LispVal
               | NotFunction String String
               | UnboundVar String String
               | Default String

showError :: LispError -> String
showError (UnboundVar message varname)  = message ++ ": " ++ varname
showError (BadSpecialForm message form) = message ++ ": " ++ show form
showError (NotFunction message func)    = message ++ ": " ++ show func
showError (NumArgs expected found)      = "Expected " ++ show expected
                                          ++ " args; found values " ++ unwordsList found
showError (TypeMismatch expected found) = "Invalid type: expected " ++ expected
                                          ++ ", found " ++ show found
showError (Parser parseErr)             = "Parse error at " ++ show parseErr

instance Show LispError where show = showError

instance Error LispError where
  noMsg = Default "An error has occurred"
  strMsg = Default

type ThrowsError = Either LispError

trapError :: (MonadError e m, Show e) => m String -> m String
trapError action = catchError action (return . show)

extractValue :: ThrowsError a -> a
extractValue (Right val) = val


{-# LANGUAGE ForeignFunctionInterface, GeneralizedNewtypeDeriving #-}

module Numerical.HBLAS.BLAS.FFI.Level2  where

import Foreign.Ptr
import Foreign()
import Foreign.C.Types
import Data.Complex
import Numerical.HBLAS.BLAS.FFI

type GbmvFunFFI sc el =
       CBLAS_ORDERT -> CBLAS_TRANSPOSET -> CInt -> CInt
    -> CInt -> CInt -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> sc -> Ptr el -> CInt -> IO ()

foreign import ccall unsafe "cblas_sgbmv"
    cblas_sgbmv_unsafe :: GbmvFunFFI Float Float
foreign import ccall unsafe "cblas_dgbmv"
    cblas_dgbmv_unsafe :: GbmvFunFFI Double Double
foreign import ccall unsafe "cblas_cgbmv"
    cblas_cgbmv_unsafe :: GbmvFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall unsafe "cblas_zgbmv"
    cblas_zgbmv_unsafe :: GbmvFunFFI (Ptr (Complex Double)) (Complex Double)

foreign import ccall "cblas_sgbmv"
    cblas_sgbmv_safe :: GbmvFunFFI Float Float
foreign import ccall "cblas_dgbmv"
    cblas_dgbmv_safe :: GbmvFunFFI Double Double
foreign import ccall "cblas_cgbmv"
    cblas_cgbmv_safe :: GbmvFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall "cblas_zgbmv"
    cblas_zgbmv_safe :: GbmvFunFFI (Ptr (Complex Double)) (Complex Double)
--void cblas_sgbmv(  enum CBLAS_ORDER order,   enum CBLAS_TRANSPOSE TransA,   CInt M,   CInt N,
--                   CInt KL,   CInt KU,   Float alpha,   Float *A,   CInt lda,   Float *X,   CInt incX,   Float beta, Float *Y,   CInt incY);
--void cblas_dgbmv(  enum CBLAS_ORDER order,   enum CBLAS_TRANSPOSE TransA,   CInt M,   CInt N,
--                   CInt KL,   CInt KU,   Double alpha,   Double *A,   CInt lda,   Double *X,   CInt incX,   Double beta, Double *Y,   CInt incY);
--void cblas_cgbmv(  enum CBLAS_ORDER order,   enum CBLAS_TRANSPOSE TransA,   CInt M,   CInt N,
--                   CInt KL,   CInt KU,   Float *alpha,   Float *A,   CInt lda,   Float *X,   CInt incX,   Float *beta, Float *Y,   CInt incY);
--void cblas_zgbmv(  enum CBLAS_ORDER order,   enum CBLAS_TRANSPOSE TransA,   CInt M,   CInt N,
--                   CInt KL,   CInt KU,   Double *alpha,   Double *A,   CInt lda,   Double *X,   CInt incX,   Double *beta, Double *Y,   CInt incY);

{-
matrix vector product for general matrices
 perform one of the matrix-vector operations   y :=
      alpha*A*x + beta*y, or y := alpha*A'*x + beta*y,
-}

type GemvFunFFI sc el =
       CBLAS_ORDERT -> CBLAS_TRANSPOSET -> CInt -> CInt
    -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> sc -> Ptr el -> CInt -> IO ()

foreign import ccall unsafe "cblas_sgemv"
    cblas_sgemv_unsafe :: GemvFunFFI Float Float
foreign import ccall safe   "cblas_sgemv"
    cblas_sgemv_safe   :: GemvFunFFI Float Float

foreign import ccall unsafe "cblas_dgemv"
    cblas_dgemv_unsafe :: GemvFunFFI Double Double
foreign import ccall safe   "cblas_dgemv"
    cblas_dgemv_safe   :: GemvFunFFI Double Double

foreign import ccall unsafe "cblas_cgemv"
    cblas_cgemv_unsafe :: GemvFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe   "cblas_cgemv"
    cblas_cgemv_safe   :: GemvFunFFI (Ptr (Complex Float)) (Complex Float)

foreign import ccall unsafe "cblas_zgemv"
    cblas_zgemv_unsafe :: GemvFunFFI (Ptr (Complex Double)) (Complex Double)
foreign import ccall safe   "cblas_zgemv"
    cblas_zgemv_safe   :: GemvFunFFI (Ptr (Complex Double)) (Complex Double)

--void cblas_sgemv(  enum CBLAS_ORDER order,    enum CBLAS_TRANSPOSE trans,    CInt m,   CInt n,
--           Float alpha,   Float  *a,   CInt lda,    Float  *x,   CInt incx,    Float beta,  Float  *y,   CInt incy);
--void cblas_dgemv(  enum CBLAS_ORDER order,    enum CBLAS_TRANSPOSE trans,    CInt m,   CInt n,
--           Double alpha,   Double  *a,   CInt lda,    Double  *x,   CInt incx,    Double beta,  Double  *y,   CInt incy);
--void cblas_cgemv(  enum CBLAS_ORDER order,    enum CBLAS_TRANSPOSE trans,    CInt m,   CInt n,
--           Float *alpha,   Float  *a,   CInt lda,    Float  *x,   CInt incx,    Float *beta,  Float  *y,   CInt incy);
--void cblas_zgemv(  enum CBLAS_ORDER order,    enum CBLAS_TRANSPOSE trans,    CInt m,   CInt n,
--           Double *alpha,   Double  *a,   CInt lda,    Double  *x,   CInt incx,    Double *beta,  Double  *y,   CInt incy);

-- perform the rank 1 operation   A := alpha*x*y' + A,

type GerxFunFFI scale el = CBLAS_ORDERT -> CInt -> CInt -> scale -> Ptr el -> CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()

foreign import ccall unsafe "cblas_sger" cblas_sger_unsafe ::
        GerxFunFFI Float Float
foreign import ccall safe   "cblas_sger" cblas_sger_safe   ::
        GerxFunFFI Float Float

foreign import ccall unsafe "cblas_dger" cblas_dger_unsafe ::
        GerxFunFFI Double Double
foreign import ccall safe   "cblas_dger" cblas_dger_safe   ::
        GerxFunFFI Double Double

foreign import ccall unsafe "cblas_cgerc" cblas_cgerc_unsafe ::
        GerxFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall unsafe "cblas_zgerc" cblas_zgerc_unsafe ::
        GerxFunFFI (Ptr (Complex Double)) (Complex Double)

foreign import ccall safe   "cblas_cgerc" cblas_cgerc_safe   ::
        GerxFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe   "cblas_zgerc" cblas_zgerc_safe   ::
        GerxFunFFI (Ptr (Complex Double)) (Complex Double)

foreign import ccall unsafe "cblas_cgeru" cblas_cgeru_unsafe ::
        GerxFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall unsafe "cblas_zgeru" cblas_zgeru_unsafe ::
        GerxFunFFI (Ptr (Complex Double)) (Complex Double)

foreign import ccall safe   "cblas_cgeru" cblas_cgeru_safe   ::
        GerxFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe   "cblas_zgeru" cblas_zgeru_safe   ::
        GerxFunFFI (Ptr (Complex Double)) (Complex Double)

--void cblas_sger (  enum CBLAS_ORDER order,   CInt M,   CInt N,   Float   alpha,   Float  *X,   CInt incX,   Float  *Y,   CInt incY, Float  *A,   CInt lda);
--void cblas_dger (  enum CBLAS_ORDER order,   CInt M,   CInt N,   Double  alpha,   Double *X,   CInt incX,   Double *Y,   CInt incY, Double *A,   CInt lda);
--void cblas_cgeru(  enum CBLAS_ORDER order,   CInt M,   CInt N,   Float  *alpha,   Float  *X,   CInt incX,   Float  *Y,   CInt incY, Float  *A,   CInt lda);
--void cblas_cgerc(  enum CBLAS_ORDER order,   CInt M,   CInt N,   Float  *alpha,   Float  *X,   CInt incX,   Float  *Y,   CInt incY, Float  *A,   CInt lda);
--void cblas_zgeru(  enum CBLAS_ORDER order,   CInt M,   CInt N,   Double *alpha,   Double *X,   CInt incX,   Double *Y,   CInt incY, Double *A,   CInt lda);
--void cblas_zgerc(  enum CBLAS_ORDER order,   CInt M,   CInt N,   Double *alpha,   Double *X,   CInt incX,   Double *Y,   CInt incY, Double *A,   CInt lda);

type HbmvFunFFI sc el =
       CBLAS_ORDERT -> CBLAS_UPLOT -> CInt -> CInt
    -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> sc -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_chbmv"
    cblas_chbmv_unsafe :: HbmvFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe   "cblas_chbmv"
    cblas_chbmv_safe   :: HbmvFunFFI (Ptr (Complex Float)) (Complex Float)

foreign import ccall unsafe "cblas_zhbmv"
    cblas_zhbmv_unsafe :: HbmvFunFFI (Ptr (Complex Double)) (Complex Double)
foreign import ccall safe   "cblas_zhbmv"
    cblas_zhbmv_safe   :: HbmvFunFFI (Ptr (Complex Double)) (Complex Double)
--void cblas_chbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   CInt K,
--           Float *alpha,   Float *A,   CInt lda,   Float *X,   CInt incX,   Float *beta, Float *Y,   CInt incY);
--void cblas_zhbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   CInt K,
--           Double *alpha,   Double *A,   CInt lda,   Double *X,   CInt incX,   Double *beta, Double *Y,   CInt incY);

type HemvFunFFI sc el =
       CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
    -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> sc -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_chemv"
    cblas_chemv_unsafe :: HemvFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe   "cblas_chemv"
    cblas_chemv_safe   :: HemvFunFFI (Ptr (Complex Float)) (Complex Float)

foreign import ccall unsafe "cblas_zhemv"
    cblas_zhemv_unsafe :: HemvFunFFI (Ptr (Complex Double)) (Complex Double)
foreign import ccall safe   "cblas_zhemv"
    cblas_zhemv_safe   :: HemvFunFFI (Ptr (Complex Double)) (Complex Double)
--void cblas_chemv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float *alpha,   Float *A,
--                   CInt lda,   Float *X,   CInt incX,   Float *beta, Float *Y,   CInt incY);
--void cblas_zhemv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double *alpha,   Double *A,
--                   CInt lda,   Double *X,   CInt incX,   Double *beta, Double *Y,   CInt incY);

type HerFunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_cher"
    cblas_cher_unsafe :: HerFunFFI Float (Complex Float)
foreign import ccall safe "cblas_cher"
    cblas_cher_safe :: HerFunFFI Float (Complex Float)

foreign import ccall unsafe "cblas_zher"
    cblas_zher_unsafe :: HerFunFFI Double (Complex Double)
foreign import ccall safe "cblas_zher"
    cblas_zher_safe :: HerFunFFI Double (Complex Double)
--void cblas_cher(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float alpha,   Float *X,   CInt incX, Float *A,   CInt lda);
--void cblas_zher(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *X,   CInt incX, Double *A,   CInt lda);

type Her2FunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_cher2"
    cblas_cher2_unsafe :: Her2FunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe "cblas_cher2"
    cblas_cher2_safe :: Her2FunFFI (Ptr (Complex Float)) (Complex Float)

foreign import ccall unsafe "cblas_zher2"
    cblas_zher2_unsafe :: Her2FunFFI (Ptr (Complex Double)) (Complex Double)
foreign import ccall safe "cblas_zher2"
    cblas_zher2_safe :: Her2FunFFI (Ptr (Complex Double)) (Complex Double)
--void cblas_cher2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float *alpha,   Float *X,   CInt incX,
--                  Float *Y,   CInt incY, Float *A,   CInt lda);
--void cblas_zher2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double *alpha,   Double *X,   CInt incX,
--                  Double *Y,   CInt incY, Double *A,   CInt lda);

type HpmvFunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> Ptr el -> CInt -> sc -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_chpmv"
    cblas_chpmv_unsafe :: HpmvFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe "cblas_chpmv"
    cblas_chpmv_safe :: HpmvFunFFI (Ptr (Complex Float)) (Complex Float)

foreign import ccall unsafe "cblas_zhpmv"
    cblas_zhpmv_unsafe :: HpmvFunFFI (Ptr (Complex Double)) (Complex Double)
foreign import ccall safe "cblas_zhpmv"
    cblas_zhpmv_safe :: HpmvFunFFI (Ptr (Complex Double)) (Complex Double)
--void cblas_chpmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,
--           Float *alpha,   Float *Ap,   Float *X,   CInt incX,   Float *beta, Float *Y,   CInt incY);
--void cblas_zhpmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,
--           Double *alpha,   Double *Ap,   Double *X,   CInt incX,   Double *beta, Double *Y,   CInt incY);

type HprFunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> CInt -> Ptr el -> IO ()
foreign import ccall unsafe "cblas_chpr"
    cblas_chpr_unsafe :: HprFunFFI Float (Complex Float)
foreign import ccall safe "cblas_chpr"
    cblas_chpr_safe :: HprFunFFI Float (Complex Float)

foreign import ccall unsafe "cblas_zhpr"
    cblas_zhpr_unsafe :: HprFunFFI Double (Complex Double)
foreign import ccall safe "cblas_zhpr"
    cblas_zhpr_safe :: HprFunFFI Double (Complex Double)
--void cblas_chpr(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float alpha,   Float *X,   CInt incX, Float *A);
--void cblas_zhpr(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *X,  CInt incX, Double *A);

type Hpr2FunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> Ptr el -> IO ()
foreign import ccall unsafe "cblas_chpr2"
    cblas_chpr2_unsafe :: Hpr2FunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe "cblas_chpr2"
    cblas_chpr2_safe :: Hpr2FunFFI (Ptr (Complex Float)) (Complex Float)

foreign import ccall unsafe "cblas_zhpr2"
    cblas_zhpr2_unsafe :: Hpr2FunFFI (Ptr (Complex Double)) (Complex Double)
foreign import ccall safe "cblas_zhpr2"
    cblas_zhpr2_safe :: Hpr2FunFFI (Ptr (Complex Double)) (Complex Double)
--void cblas_chpr2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float *alpha,   Float *X,   CInt incX,   Float *Y,   CInt incY, Float *Ap);
--void cblas_zhpr2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double *alpha,   Double *X,   CInt incX,   Double *Y,   CInt incY, Double *Ap);

type SbmvFunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt -> CInt
      -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> sc -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_ssbmv"
    cblas_ssbmv_unsafe :: SbmvFunFFI Float Float
foreign import ccall safe "cblas_ssbmv"
    cblas_ssbmv_safe :: SbmvFunFFI Float Float

foreign import ccall unsafe "cblas_dsbmv"
    cblas_dsbmv_unsafe :: SbmvFunFFI Double Double
foreign import ccall safe "cblas_dsbmv"
    cblas_dsbmv_safe :: SbmvFunFFI Double Double
--void cblas_ssbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   CInt K,   Float alpha,   Float *A,
--                   CInt lda,   Float *X,   CInt incX,   Float beta, Float *Y,   CInt incY);
--void cblas_dsbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   CInt K,   Double alpha,   Double *A,
--                   CInt lda,   Double *X,   CInt incX,   Double beta, Double *Y,   CInt incY);

---------------
--- | packed symmetric matrix * vector product  y:= alpha * Av  + beta * y
---------------
type SpmvFunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> Ptr el -> CInt -> sc -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_sspmv"
    cblas_sspmv_unsafe :: SpmvFunFFI Float Float
foreign import ccall safe "cblas_sspmv"
    cblas_sspmv_safe :: SpmvFunFFI Float Float

foreign import ccall unsafe "cblas_dspmv"
    cblas_dspmv_unsafe :: SpmvFunFFI Double Double
foreign import ccall safe "cblas_dspmv"
    cblas_dspmv_safe :: SpmvFunFFI Double Double
--void cblas_sspmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float alpha,   Float *Ap,
--                   Float *X,   CInt incX,   Float beta, Float *Y,   CInt incY);
--void cblas_dspmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *Ap,
--                   Double *X,   CInt incX,   Double beta, Double *Y,   CInt incY);

type SprFunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> CInt -> Ptr el -> IO ()
foreign import ccall unsafe "cblas_sspr"
    cblas_sspr_unsafe :: SprFunFFI Float Float
foreign import ccall safe "cblas_sspr"
    cblas_sspr_safe :: SprFunFFI Float Float

foreign import ccall unsafe "cblas_dspr"
    cblas_dspr_unsafe :: SprFunFFI Double Double
foreign import ccall safe "cblas_dspr"
    cblas_dspr_safe :: SprFunFFI Double Double
--void cblas_sspr(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float alpha,   Float *X,   CInt incX, Float *Ap);
--void cblas_dspr(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *X,   CInt incX, Double *Ap);

type Spr2FunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> Ptr el -> IO ()
foreign import ccall unsafe "cblas_sspr2"
    cblas_sspr2_unsafe :: Spr2FunFFI Float Float
foreign import ccall safe "cblas_sspr2"
    cblas_sspr2_safe :: Spr2FunFFI Float Float

foreign import ccall unsafe "cblas_dspr2"
    cblas_dspr2_unsafe :: Spr2FunFFI Double Double
foreign import ccall safe "cblas_dspr2"
    cblas_dspr2_safe :: Spr2FunFFI Double Double
--void cblas_sspr2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float alpha,   Float *X,   CInt incX,   Float *Y,   CInt incY, Float *A);
--void cblas_dspr2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *X,   CInt incX,   Double *Y,   CInt incY, Double *A);

----------------------------------
---- |  (unpacked) symmetric matrix vector product    x:=Av, writes result x into v
---------------------------------

type SymvFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CInt -> el -> Ptr el ->  CInt ->
                        Ptr el -> CInt -> el -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_ssymv"
    cblas_ssymv_unsafe :: SymvFunFFI Float
foreign import ccall safe "cblas_ssymv"
    cblas_ssymv_safe :: SymvFunFFI Float

foreign import ccall unsafe "cblas_dsymv"
    cblas_dsymv_unsafe :: SymvFunFFI Double
foreign import ccall safe "cblas_dsymv"
    cblas_dsymv_safe :: SymvFunFFI Double
--void cblas_ssymv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float alpha,   Float *A,
--                   CInt lda,   Float *X,   CInt incX,   Float beta, Float *Y,   CInt incY);
--void cblas_dsymv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *A,
--                   CInt lda,   Double *X,   CInt incX,   Double beta, Double *Y,   CInt incY);

type SyrFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CInt -> el -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_ssyr"
      cblas_ssyr_unsafe :: SyrFunFFI Float
foreign import ccall safe "cblas_ssyr"
      cblas_ssyr_safe :: SyrFunFFI Float

foreign import ccall unsafe "cblas_dsyr"
      cblas_dsyr_unsafe :: SyrFunFFI Double
foreign import ccall safe "cblas_dsyr"
      cblas_dsyr_safe :: SyrFunFFI Double
--void cblas_ssyr(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float alpha,   Float *X,   CInt incX, Float *A,   CInt lda);
--void cblas_dsyr(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *X,   CInt incX, Double *A,   CInt lda);

type Syr2FunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CInt -> el -> Ptr el -> CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_ssyr2"
      cblas_ssyr2_unsafe :: Syr2FunFFI Float
foreign import ccall safe "cblas_ssyr2"
      cblas_ssyr2_safe :: Syr2FunFFI Float

foreign import ccall unsafe "cblas_dsyr2"
      cblas_dsyr2_unsafe :: Syr2FunFFI Double
foreign import ccall safe "cblas_dsyr2"
      cblas_dsyr2_safe :: Syr2FunFFI Double
--void cblas_ssyr2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,  CInt N,   Float alpha,   Float *X,
--                  CInt incX,   Float *Y,   CInt incY, Float *A,   CInt lda);
--void cblas_dsyr2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *X,
--                  CInt incX,   Double *Y,   CInt incY, Double *A,   CInt lda);

type TbmvFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
                      CInt -> CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_stbmv"
      cblas_stbmv_unsafe :: TbmvFunFFI Float
foreign import ccall safe "cblas_stbmv"
      cblas_stbmv_safe :: TbmvFunFFI Float

foreign import ccall unsafe "cblas_dtbmv"
      cblas_dtbmv_unsafe :: TbmvFunFFI Double
foreign import ccall safe "cblas_dtbmv"
      cblas_dtbmv_safe :: TbmvFunFFI Double

foreign import ccall unsafe "cblas_ctbmv"
      cblas_ctbmv_unsafe :: TbmvFunFFI (Complex Float)
foreign import ccall safe "cblas_ctbmv"
      cblas_ctbmv_safe :: TbmvFunFFI (Complex Float)

foreign import ccall unsafe "cblas_ztbmv"
      cblas_ztbmv_unsafe :: TbmvFunFFI (Complex Double)
foreign import ccall safe "cblas_ztbmv"
      cblas_ztbmv_safe :: TbmvFunFFI (Complex Double)
--void cblas_stbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_dtbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Double *A,   CInt lda, Double *X,   CInt incX);
--void cblas_ctbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_ztbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Double *A,   CInt lda, Double *X,   CInt incX);

----------------
--- | solves  Ax=v where A is k+1 banded triangular matrix, and x and
----------------
type TbsvFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
                      CInt -> CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_stbsv"
      cblas_stbsv_unsafe :: TbsvFunFFI Float
foreign import ccall safe "cblas_stbsv"
      cblas_stbsv_safe :: TbsvFunFFI Float

foreign import ccall unsafe "cblas_dtbsv"
      cblas_dtbsv_unsafe :: TbsvFunFFI Double
foreign import ccall safe "cblas_dtbsv"
      cblas_dtbsv_safe :: TbsvFunFFI Double

foreign import ccall unsafe "cblas_ctbsv"
      cblas_ctbsv_unsafe :: TbsvFunFFI (Complex Float)
foreign import ccall safe "cblas_ctbsv"
      cblas_ctbsv_safe :: TbsvFunFFI (Complex Float)

foreign import ccall unsafe "cblas_ztbsv"
      cblas_ztbsv_unsafe :: TbsvFunFFI (Complex Double)
foreign import ccall safe "cblas_ztbsv"
      cblas_ztbsv_safe :: TbsvFunFFI (Complex Double)
--void cblas_stbsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_dtbsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Double *A,   CInt lda, Double *X,   CInt incX);
--void cblas_ctbsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_ztbsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Double *A,   CInt lda, Double *X,   CInt incX);

-------------------------------------------------------------------------
-- | matrix vector product Av, writes result into v, where A is a packed triangular nxn matrix
-------------------------------------------------------------------------
type TpmvFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
                      CInt -> Ptr el -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_stpmv"
      cblas_stpmv_unsafe :: TpmvFunFFI Float
foreign import ccall safe "cblas_stpmv"
      cblas_stpmv_safe :: TpmvFunFFI Float

foreign import ccall unsafe "cblas_dtpmv"
      cblas_dtpmv_unsafe :: TpmvFunFFI Double
foreign import ccall safe "cblas_dtpmv"
      cblas_dtpmv_safe :: TpmvFunFFI Double

foreign import ccall unsafe "cblas_ctpmv"
      cblas_ctpmv_unsafe :: TpmvFunFFI (Complex Float)
foreign import ccall safe "cblas_ctpmv"
      cblas_ctpmv_safe :: TpmvFunFFI (Complex Float)

foreign import ccall unsafe "cblas_ztpmv"
      cblas_ztpmv_unsafe :: TpmvFunFFI (Complex Double)
foreign import ccall safe "cblas_ztpmv"
      cblas_ztpmv_safe :: TpmvFunFFI (Complex Double)
--void cblas_stpmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Float *Ap, Float *X,   CInt incX);
--void cblas_dtpmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Double *Ap, Double *X,   CInt incX);
--void cblas_ctpmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Float *Ap, Float *X,   CInt incX);
--void cblas_ztpmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Double *Ap, Double *X,   CInt incX);

--------------------------------------------------
---  | solve  Ax=v where A is a nxn packed triangular matrix, v vector input, writes the solution into x.
--------------------------------------------------
type TpsvFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
                      CInt -> Ptr el -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_stpsv"
      cblas_stpsv_unsafe :: TpsvFunFFI Float
foreign import ccall safe "cblas_stpsv"
      cblas_stpsv_safe :: TpsvFunFFI Float

foreign import ccall unsafe "cblas_dtpsv"
      cblas_dtpsv_unsafe :: TpsvFunFFI Double
foreign import ccall safe "cblas_dtpsv"
      cblas_dtpsv_safe :: TpsvFunFFI Double

foreign import ccall unsafe "cblas_ctpsv"
      cblas_ctpsv_unsafe :: TpsvFunFFI (Complex Float)
foreign import ccall safe "cblas_ctpsv"
      cblas_ctpsv_safe :: TpsvFunFFI (Complex Float)

foreign import ccall unsafe "cblas_ztpsv"
      cblas_ztpsv_unsafe :: TpsvFunFFI (Complex Double)
foreign import ccall safe "cblas_ztpsv"
      cblas_ztpsv_safe :: TpsvFunFFI (Complex Double)
--void cblas_stpsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Float *Ap, Float *X,   CInt incX);
--void cblas_dtpsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Double *Ap, Double *X,   CInt incX);
--void cblas_ctpsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Float *Ap, Float *X,   CInt incX);
--void cblas_ztpsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Double *Ap, Double *X,   CInt incX);

type TrmvFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
                      CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_strmv"
      cblas_strmv_unsafe :: TrmvFunFFI Float
foreign import ccall safe "cblas_strmv"
      cblas_strmv_safe :: TrmvFunFFI Float

foreign import ccall unsafe "cblas_dtrmv"
      cblas_dtrmv_unsafe :: TrmvFunFFI Double
foreign import ccall safe "cblas_dtrmv"
      cblas_dtrmv_safe :: TrmvFunFFI Double

foreign import ccall unsafe "cblas_ctrmv"
      cblas_ctrmv_unsafe :: TrmvFunFFI (Complex Float)
foreign import ccall safe "cblas_ctrmv"
      cblas_ctrmv_safe :: TrmvFunFFI (Complex Float)

foreign import ccall unsafe "cblas_ztrmv"
      cblas_ztrmv_unsafe :: TrmvFunFFI (Complex Double)
foreign import ccall safe "cblas_ztrmv"
      cblas_ztrmv_safe :: TrmvFunFFI (Complex Double)
--void cblas_strmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_dtrmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Double *A,   CInt lda, Double *X,   CInt incX);
--void cblas_ctrmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_ztrmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Double *A,   CInt lda, Double *X,   CInt incX);

--STRSV - solve one of the systems of equations   A*x = b, or A'*x = b, where A is a (non)unit upper(/lower) triangular matrix
type TrsvFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
                      CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()

foreign import ccall unsafe "cblas_strsv"
  cblas_strsv_unsafe :: TrsvFunFFI Float
foreign import ccall safe   "cblas_strsv"
  cblas_strsv_safe   :: TrsvFunFFI Float

foreign import ccall unsafe "cblas_dtrsv"
  cblas_dtrsv_unsafe :: TrsvFunFFI Double
foreign import ccall safe   "cblas_dtrsv"
  cblas_dtrsv_safe   :: TrsvFunFFI Double

foreign import ccall unsafe "cblas_ctrsv"
  cblas_ctrsv_unsafe :: TrsvFunFFI (Complex Float)
foreign import ccall safe   "cblas_ctrsv"
  cblas_ctrsv_safe   :: TrsvFunFFI (Complex Float)

foreign import ccall unsafe "cblas_ztrsv"
  cblas_ztrsv_unsafe :: TrsvFunFFI (Complex Double)
foreign import ccall safe   "cblas_ztrsv"
  cblas_ztrsv_safe   :: TrsvFunFFI (Complex Double)

--void cblas_strsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_dtrsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Double *A,   CInt lda, Double *X,   CInt incX);
--void cblas_ctrsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_ztrsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Double *A,   CInt lda, Double *X,   CInt incX);


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

module Fib where

import Vec
import Data.Complex
import Control.Monad  ((<=<))
import GHC.TypeNats
import Data.Proxy

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

data Tau
--data Id
type 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)
   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)
   TLeaf :: FibTree Tau Tau
   ILeaf :: FibTree Id Id

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

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

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 -> 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)

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

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

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)

rightUnit :: FibTree e (a,Id) -> Q (FibTree e a)
rightUnit (TTI t _) = pure t
rightUnit (III t _) = pure t

rightUnit' :: FibTree e a -> Q (FibTree e (a,Id))
rightUnit' t@(TTT _ _) = pure (TTI  t ILeaf)
rightUnit' t@(TTI _ _) = pure (TTI  t ILeaf)
rightUnit' t@(TIT _ _) = pure (TTI  t ILeaf)
rightUnit' t@(III _ _) = pure (III  t ILeaf)
rightUnit' t@(ITT _ _) = pure (III  t ILeaf)
rightUnit' t@(ILeaf) = pure (III t ILeaf)
rightUnit' t@(TLeaf) = pure (TTI t ILeaf)

leftUnit :: FibTree e (Id,a) -> Q (FibTree e a)
leftUnit = rightUnit <=< braid

-- braid vs braid' doesn't matter, but it has a nice symmettry.
leftUnit' :: FibTree e a -> Q (FibTree e (Id,a))
leftUnit' = braid' <=< rightUnit'

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

class AllTrees a b where
    allTrees :: [FibTree a b]
instance AllTrees Tau Tau where
    allTrees = [TLeaf]
instance AllTrees Id Tau where
        allTrees = []
instance AllTrees Tau Id where
        allTrees = []
instance AllTrees Id Id where
        allTrees = [ILeaf]
instance (AllTrees Tau a,
         AllTrees Id a,
         AllTrees Tau b,
         AllTrees Id b) => AllTrees Id (a,b) where
        allTrees = (III <$> ia <*> ib) <> (ITT <$> ta <*> tb) where
            ta = allTrees @Tau @a
            ia = allTrees @Id @a
            tb = allTrees @Tau @b
            ib = allTrees @Id @b
instance (AllTrees Tau a,
        AllTrees Id a,
        AllTrees Tau b,
        AllTrees Id b) => AllTrees Tau (a,b) where
       allTrees = (TIT <$> ia <*> tb) <> (TTI <$> ta <*> ib) <> (TTT <$> ta <*> tb) where
            ta = allTrees @Tau @a
            ia = allTrees @Id @a
            tb = allTrees @Tau @b
            ib = allTrees @Id @b
t9 = allTrees @Tau @((Tau,Tau),Tau)

-- Resulting type depends on input
-- I think a typefamily type computation might be necessary?
-- pullLeft and pullRight might not need a type class.
-- pullLeft (TTT l r) = fmove (TTT pl r)
--                         where pl = pullLeft l
{-
type family Assoc a where
   Assoc ((a,b),c) = (a,(b,c))

type family PullLeft a where
   PullLeft ((a,b),c) = Assoc (PullLeft (a, b), c)
   PullLeft a = a

type Test1 = PullLeft (((Int,Int),Int),Int)
pullLeft :: forall a b c. (c ~ PullLeft b) => FibTree a b -> Q (FibTree a c)
pullLeft TLeaf = pure Tleaf
pullLeft ILeaf = pure ILeaf
pullLeft t =  do
            t' <- lmap pullLeft t
            fmove' t'

pullLeft t@(ITT l r) = helper t
pullLeft t@(TTT l r) = helper t
pullLeft t@(TTI l r) = helper t
pullLeft t@(TIT l r) = helper t
pullLeft t@(III l r) = helper t where
                       helper :: (d ~ PullLeft (b,c)) => FibTree a (b,c) -> Q (FibTree a d)
                       helper t = do
                                t' <- lmap pullLeft t
                                fmove' t'

-}

{-
Spiritually this function is this. Can't write it this way unforturnately
-- There are typing problems
-- this isn't even right
-- we need to pattern match on ((),)
pullLeft TLeaf = pure Tleaf
pullLeft ILeaf = pure ILeaf
pullLeft t =  do
                t' <- lmap pullLeft t
                fmove' t'

-- but actually a two level tree desctruct
--treeDestruct()

-}
{-
type family Assoc a where
    Assoc ((a,b),c) = (a,(b,c))

 type family PullLeft' a where
    PullLeft' ((a,b),c) = Assoc (PullLeft' (a, b), c)
    PullLeft' a = a
    -}
-- could seperate the typeclass so that I'm using type families instead.

class PullLeftLeaf a b | a -> b where
  pullLeftLeaf :: FibTree c a -> Q (FibTree c b)
instance PullLeftLeaf (Tau,c) (Tau,c) where
  pullLeftLeaf = pure
instance PullLeftLeaf (Id,c) (Id,c) where
  pullLeftLeaf = pure
instance PullLeftLeaf Tau Tau where
  pullLeftLeaf = pure
instance PullLeftLeaf Id Id where
  pullLeftLeaf = pure
instance (PullLeftLeaf (a,b) (a',b'),
          r ~ (a',(b',c))) => PullLeftLeaf ((a, b),c) r where
  pullLeftLeaf t = do
           t' <- lmap pullLeftLeaf t
           fmove' t'

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

instance PullRightLeaf Id Id where
    pullRightLeaf = pure

instance PullRightLeaf (c,Tau) (c,Tau) where
    pullRightLeaf = pure

instance PullRightLeaf (c,Id) (c,Id) where
    pullRightLeaf = pure

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

class RightAssoc a b | a -> b where
  rightAssoc :: FibTree c a -> Q (FibTree c b)
instance RightAssoc Tau Tau where
  rightAssoc = pure
instance RightAssoc Id Id where
  rightAssoc = pure
instance (PullLeftLeaf (a,b) (a',b'),
          RightAssoc b' b'',
          r ~ (a', b'')) => RightAssoc (a,b) r where
  rightAssoc t = do
           t' <- pullLeftLeaf t
           rmap rightAssoc t'
-- usually you'll want to force not the root, but the leaves to be some type
-- hence the ordering b c d a for type application
bmove :: forall b c d a. FibTree a (b,(c,d)) -> Q (FibTree a (c,(b,d)))
bmove t = do
           t'  :: FibTree a ((b,c),d) <- fmove t
           t'' :: FibTree a ((c,b),d) <-  lmap braid t'
           fmove' t''
bmove' :: forall b c d a. FibTree a (b,(c,d)) -> Q (FibTree a (c,(b,d)))
bmove' = fmove' <=< (lmap braid') <=< fmove

-- Therei s a general pattern for digging into
-- n and a tell us b, the subpiece of a
-- c and a and n tell use what a looks like with b replaced = d
--data Nat = Z | S Nat

rmapN :: forall n gte s t a b e. (RMapN n gte s t a b, gte ~ (CmpNat n 0)) => (forall r. FibTree r a -> Q (FibTree r b)) -> (FibTree e s) -> Q (FibTree e t)
rmapN f t = rmapN' @n @gte f t

class RMapN n gte s t a b | n gte s b -> a t where
    rmapN' :: (forall r. FibTree r a -> Q (FibTree r b)) -> (FibTree e s) -> Q (FibTree e t)
instance (a ~ s, b ~ t) => RMapN 0 'EQ s t a b where
    rmapN' f t = f t -- rmapN = id
instance (RMapN (n-1) gte r r' a b,
               gte ~ (CmpNat (n-1) 0),
               t ~ (l,r')) => RMapN n 'GT (l,r) t a b where
    rmapN' f t = rmap (rmapN @(n-1) f) t
{-
bmoveN :: forall n a b c d e f. RMapN n a (b, (c, d)) (c, (b, d)) e => FibTree f a -> Q (FibTree f e)
bmoveN t = rmapN @n (bmove @b @c @d) t

-- dotN :: forall n a b c d e f. dot :: FibTree a (b, c) -> FibTree a' (b, c) -> Q (FibTree a' a)
fuseN :: forall n a b c d e f g. RMapN n f (b,(c,d)) (a,d) g => FibTree a (b, c) -> FibTree e f -> Q (FibTree e g)
fuseN q t = rmapN @n f t where
    f :: forall r. FibTree r (b,(c,d)) -> Q (FibTree r (a,d))
    f t' = do
            t'' <- fmove t'
            lmap (dot q) 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

lcamap ::  forall n s t a b e gte .
           (gte ~ CmpNat (LeftCount s) n,
           LCAMap n gte s t a b)
           => (forall r. FibTree r a -> Q (FibTree r b)) -> (FibTree e s) -> Q (FibTree e t)
lcamap f t = lcamap' @n @gte f t

class LCAMap n gte s t a b | n gte s b -> t a where
    lcamap' :: (forall r. FibTree r a -> Q (FibTree r b)) -> (FibTree e s) -> Q (FibTree e t)

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 level down to order which way we have to go next
            gte ~ (CmpNat lc n'), -- Do we go left, right or have we arrived in the next layer?
            LCAMap n' gte r r' a b,  -- recursive call
            t ~ (l,r') -- reconstruct total return type from recursive return type. Left tree is unaffected by lcamapping
            ) => LCAMap n 'LT (l,r) t a b where
        lcamap' f x = rmap (lcamap @n' f) x

instance (lc ~ (LeftCount l),
            gte ~ (CmpNat lc n),
            LCAMap n gte l l' a b,
            t ~ (l',r)
            ) => LCAMap n 'GT (l,r) t a b where
        lcamap' f x = lmap (lcamap @n f) x

instance (t ~ b, a ~ s) => LCAMap n 'EQ s t a b where -- Base case
    lcamap' f x = f x

class Twiddle s t a b | s b -> t a where
    twiddle :: (forall r. FibTree r a -> Q (FibTree r b)) -> FibTree e s -> Q (FibTree e t)

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

nmap :: forall (n :: Nat) s t a b a' b' l l' r r' e gte.
    (gte ~ CmpNat (LeftCount s) n,
    LCAMap n gte s t a' b',
    a' ~ (l,r),
    PullRightLeaf l l',
    PullLeftLeaf r r',
    Twiddle (l',r') b' a b) =>
    (forall r. FibTree r a -> Q (FibTree r b)) -> FibTree e s -> Q (FibTree e t)
nmap f z = lcamap @n @s @t @a' @b' (\x -> do
                                           x'  <- lmap pullRightLeaf x
                                           x'' <- rmap pullLeftLeaf x'
                                           twiddle f x'') z

t1 = nmap @2 braid (TTT (TTI TLeaf ILeaf) (TTT TLeaf TLeaf))
t5 = nmap @2 pure (TTT (TTI TLeaf ILeaf) (TTT TLeaf TLeaf)) >>= nmap @3 pure
t2 = nmap @1 braid (TTT (TTI TLeaf ILeaf) (TTT TLeaf TLeaf))
t4 = nmap @1 braid (TTT TLeaf (TTT TLeaf TLeaf))
t3 = nmap @2 braid (TTT (TTT (TTT TLeaf TLeaf) TLeaf) (TTT TLeaf TLeaf))
t6 = rightAssoc (TTT (TTT (TTT TLeaf TLeaf) TLeaf) (TTT TLeaf TLeaf))
t7 = t6 >>= bmove
t8 = t6 >>= rmapN @0 bmove
-- For the category, we probably don't want the association structure.
-- '[Tau, Id, Tau, Id] typelevel list of particles
--

ttt = TTT TLeaf TLeaf
example = starttree >>=
        nmap @1 braid >>=
        nmap @2 braid >>=
        nmap @1 (dot ttt) >>=
        nmap @1 braid' >>=
        nmap @2 (dot ttt) >>=
        nmap @1 (dot ttt) where
        starttree = pure (TTT (TTT TLeaf
                              (TTT TLeaf
                                   TLeaf))
                          TLeaf
                         )
-- would be nice to use an unknot
-- example2 =
-- I should use neighbormap for this.
-- maybe I should call f fuse and make it global?

-- bmoveN t = rmapN @n (bmove :: forall r. FibTree r (b,(c,d)) -> Q (FibTree r (c,(b,d)))) t
-- bmoveN t = rmapN @n @a @(b, (c, d)) @(c, (b, d)) @e bmove t

{-
bmoveN :: forall n a b c d e f. RMapN n a (b, (c, d)) (c, (b, d)) e => Proxy n -> FibTree f a -> Q (FibTree f e)
bmoveN p t = rmapN p (bmove :: FibTree a (b,(c,d)) -> Q (FibTree a (c,(b,d)))) t
-}
   {- do
            t'  <- fmove t
            t'' <-  lmap braid' t'
            fmove' t''
-}

           --
{-
pullLeftLeaf
pullRightLeaf

-}
{-
class Standardize a b | a -> b where
    standardize :: FibTree c a -> Q (FibTree c b)

instance Standardize a b where

type family Append a b where
    Append (a,b) c = (a, Append b c)
    Append a c = (a, c)

type family Leftize where
    Leftize (a,b) = Append (Leftize a) (Leftize b)
    Leftize a = a

fullrightassoc = standardize
completerightassoc
bmove = fmove braid fmove'

standardize ::  FibTree c a -> Q (FibTree c (Leftize a))
standardize (TLeaf) = TLeaf
standardize (ILeaf) = ILeaf
standardize t = do
                x <- pullLeft t
                rmap standardize x

lcamap :: FibTree a b -> Proxy (a :: Int) -> FibTree a ?
lcamap f n t@(TTT l r) | count l == n  = f t
                       | count l < n   = lcamap f (n - count l) r
                       | otherwise     = lcamap f n l

pullRight ()

treerecursor :: FibTree a b -> (l -> r -> c) -> (leaf -> c) -> c
treerecursor (TTT l r) f g = f l r
treerecursor (TLeaf) = g Tleaf

-}