{-|
Module      : Data.Meta.NMA
Description : Second stage NMA
Copyright   : (c) Thodoris Papakonstantinou, 2019
License     : GPL-3
Maintainer  : mail@tpapak.com
Stability   : experimental
Portability : POSIX

Given a study graph of direct effects perform network-meta analysis
-}

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE NamedFieldPuns #-}

module Data.Meta.NMA
  ( NetworkEffects (..)
  , Spring (..)
  , springNMA
  )
where

import           Control.Applicative
import           Data.List
import Data.Tuple
import qualified Data.Map.Strict               as Map
import qualified Data.Set                      as Set
import           Data.Maybe
import           Data.Either
import qualified Data.IntMap.Strict as IM

import qualified Numeric.LinearAlgebra         as LA

import qualified Data.Graph.AdjacencyList      as G
import qualified Data.Graph.AdjacencyList.BFS      as BFS

import           Data.Meta.Effects
import Data.Meta.Matrices
import Data.Meta.Studies

data Effect a => NetworkEffects a =
  NetworkEffects { directs :: StudyGraph a
                 --, leagueTable :: LA.Matrix Double
                 --, lstar :: Map.Map G.Edge (Either Spring Spring)
                 --, lstar :: Map.Map G.Edge Double
                 , networkEstimates :: Map.Map TreatmentId (Map.Map TreatmentId Double)
                 , networkVariances :: Map.Map TreatmentId (Map.Map TreatmentId Double)
                 , debugmsg :: String
                 }
  deriving (Show, Eq)

-- | Matrices needed for solving the linear system
data Matrices = Matrices { kMatrix :: LA.Matrix Double
                         , llMatrix :: [Double]
                         , cMatrix :: LA.Matrix Double
                         , dMatrix :: LA.Matrix Double
                         , xMatrix :: LA.Matrix Double
                         , spnTr :: Set.Set G.Edge
                         , parents :: IM.IntMap G.Vertex
                         , vss :: [G.Vertex]
                         , directEdges :: Set.Set G.Edge
                         }

type NaturalLength = Double
type Hardness = Double
type SpringLength = Double

data Spring = Spring Hardness NaturalLength SpringLength
  deriving (Show, Eq)

reverseSpring :: Spring -> Spring
reverseSpring (Spring k ll l) = Spring k (-ll) (-l)

hardness :: Spring -> Hardness
hardness (Spring k _ _) = k

naturalLength :: Spring -> NaturalLength
naturalLength (Spring _ ll _) = ll

springLength :: Spring -> SpringLength
springLength (Spring _ _ l) = l

type TauSquare = Double

springNMA :: (ArmEffect a b, Gaussian a, Gaussian b) => [Study]
                             -> Maybe G.Vertex
                             -> Maybe TauSquare
                             -> (Arm -> Either String a)
                             -> Either String (NetworkEffects b)
springNMA studies mpinnedVertex mtau2 armEffect =
  let vs = foldl' (\acvs study ->
                    let sid = getStudyId study
                        arms = getStudyArms study
                     in case mtau2 of
                          Nothing -> Set.insert (Left sid) $ Set.union acvs
                            $ Set.fromList $ map (Right . tidOfArm) arms
                          Just tau2 ->
                            let tauvs = Set.fromList $ map (\a
                                      -> Left $ StudyId $ StringId ((show sid) ++ "_" ++ (show $ tidOfArm a))) arms
                                tidvs = Set.fromList $ map (Right . tidOfArm) arms
                             in Set.insert (Left sid) $ Set.union acvs $ Set.union tauvs tidvs
                  ) Set.empty studies  :: Set.Set (Either StudyId TreatmentId)
      vssts = IM.fromList $ zip [1..] $ Set.toList vs
      stsvs = Map.fromList $ map swap $ IM.toList vssts
      Right vsts = sequence $ IM.filter isRight vssts -- only treatment vertices
      tsvs = Map.fromList $ map swap $ IM.toList vsts
      esprings = foldl' (\eaces study ->
                    let sid = getStudyId study
                        arms = getStudyArms study
                        studysprings = sequence $  case armEffect $ head arms of
                          Left err -> [Left err]
                          Right aArm ->
                            foldl' (\acc a -> case armEffect a of
                                      Left err -> Left err : acc
                                      Right armef ->
                                        let tid = tidOfArm a
                                            tausid = StudyId $ StringId ((show sid) ++ "_" ++ show tid)
                                            vs = stsvs Map.! (Left sid)
                                            vt = stsvs Map.! (Right tid)
                                            vtau = stsvs Map.! (Left tausid)
                                            eff = relatedRelative armef
                                            nl = expectation eff
                                            k = 1 / (variance eff)
                                            spr = Spring k nl nl
                                         in  case mtau2 of
                                              Nothing ->
                                                let sttrEdge = G.Edge vs vt
                                                 in (Right (sttrEdge, spr)) : acc
                                              Just tau2 ->
                                                let sttauEdge = G.Edge vs vtau
                                                    sttauspr = Spring (2/tau2) 0 0 -- That's where tau2 is inserted
                                                    tautrEdge = G.Edge vtau vt
                                                 in [Right (sttauEdge, sttauspr)
                                                    , Right (tautrEdge, spr)] <> acc
                                     ) [] arms :: [Either String (G.Edge, Spring)]
                     in case eaces of
                          Left err -> Left err
                          Right aces ->
                             case studysprings of
                               Left err -> Left err
                               Right sprs -> Right (sprs ++ aces)
              ) (Right []) studies :: Either String [(G.Edge,Spring)]
    in case esprings of
       Left err -> Left err
       Right springs ->
         let springsMap = Map.fromList springs
             es = Map.keys springsMap
             numEdges = length es
             directedSG = G.graphFromEdges es
             undirectedSG = G.makeUndirected directedSG
             getSpring :: G.Edge -> Spring -- from undirected edge to spring
             getSpring e =
               case Map.lookup e springsMap of
                 Nothing -> reverseSpring $ springsMap Map.! (G.reverseEdge e)
                 Just sp -> sp
             pinnedV =
               case mpinnedVertex of
                 Just p -> p
                 Nothing -> 1
             -- | get the Matrices corresponding to defined pinned Vertex
             getMatrices :: G.Vertex -> Matrices
             getMatrices pinnedVertex =
             -- | Networks spanning tree
               let netbfs = BFS.bfs undirectedSG pinnedVertex
                   spnTr = Set.fromList $ BFS.spanningTree netbfs
                   parents = BFS.parent netbfs
                   vss = map (\(G.Edge u v)->v) (Set.toList spnTr)
                   llmap = Map.fromList $ map (\e ->
                              case Set.member e spnTr of
                                True -> (e, Right $ getSpring e)
                                False -> let re = G.reverseEdge e
                                          in case Set.member re spnTr of
                                               True -> (re, Right $ getSpring re)
                                               False -> (e, Left $ getSpring e)
                                        ) es
                   -- | L vector (E><1) with the natural lengths of springs which  respect
                   -- the spanning tree orientation
                   directEdges = Set.fromList $ Map.keys llmap
                   llMatrix = map (\de ->
                     let ms = llmap Map.! de
                      in case ms of
                          Left sp -> naturalLength sp
                          Right sp -> naturalLength sp)
                              $ Set.toList directEdges
                   -- weight matrix (E x E)
                   kMatrix = LA.diagl $ map (hardness . getSpring) $ Set.toList directEdges
                   -- cMatrix (V-1 x E)
                   cMatrix =
                     let vertexToRow :: G.Vertex -> [Double]
                         vertexToRow x = map (\(G.Edge u v) ->
                                         if x==u
                                            then 1.0
                                            else
                                             if x==v then (-1)
                                                     else 0
                                             )
                                       $ Set.toList directEdges
                      in LA.fromLists $ map vertexToRow vss
                   -- D matrix (V-1 x E)
                   dMatrix = cMatrix LA.<> kMatrix
                   -- X matrix (E x (V-1)) consistency matrix
                   -- Definition: X <> L* = L
                   xMatrix =
                     let edgeToRow :: G.Edge -> [Double]
                         edgeToRow (G.Edge u v) =
                           let r = replicate (length vss) 0
                               r' = IM.fromList $ zip vss r
                               followEnd :: G.Vertex
                                         -> IM.IntMap Double
                                         -> IM.IntMap Double
                               followEnd k r =
                                 if k == pinnedVertex
                                    then r
                                    else
                                      let r' = IM.adjust (+ 1) k r
                                          k' = (BFS.parent netbfs) IM.! k
                                       in followEnd k' r'
                               followStart k r =
                                 if k == pinnedVertex
                                    then r
                                    else
                                      let r' = IM.adjust (+ (-1)) k r
                                          k' = (BFS.parent netbfs) IM.! k
                                       in followStart k' r'
                               r'' = followEnd v $ followStart u r'
                            in map (\v -> r'' IM.! v) vss -- has to respect vss ordering
                      in LA.fromLists $ map edgeToRow $ Set.toList directEdges
                in Matrices { kMatrix
                            , llMatrix
                            , cMatrix
                            , dMatrix
                            , xMatrix
                            , parents
                            , spnTr
                            , vss
                            , directEdges
                            }
             -- | l* (V-1 >< 1) where V-1 are the end vertices of the spanning
             -- tree
             mtrcs = getMatrices pinnedV
             kMatrix'=  kMatrix mtrcs
             llMatrix'= llMatrix mtrcs
             cMatrix'=  cMatrix mtrcs
             dMatrix'=  dMatrix mtrcs
             xMatrix'=  xMatrix mtrcs
             parents'=  parents mtrcs
             spnTr'  =  spnTr mtrcs
             vss' = vss mtrcs
             directEdges' = directEdges mtrcs
             lstar = ((LA.inv (dMatrix' LA.<> xMatrix')) LA.<> dMatrix') LA.<> (numEdges LA.>< 1 $ llMatrix')
             lstarMap = Map.fromList $ zip (Set.toList spnTr') (LA.toList $ LA.flatten lstar)
             getLength :: G.Edge
                       -> G.Vertex -- pinned Vertex
                       -> Map.Map G.Edge Double -- lstarMap
                       -> IM.IntMap G.Vertex -- parents
                       -> NaturalLength
             -- | follows back the spanning tree adding parent edges' length
             getLength (G.Edge u v) pinnedV lstarMap parents =
               let followEnd :: G.Vertex
                       -> Double
                       -> Double
                   followEnd k l =
                       if k == pinnedV
                          then l
                          else
                            let k' = parents IM.! k
                                l' = lstarMap Map.! (G.Edge k' k)
                             in followEnd k' (l+l')
                   followStart k l =
                       if k == pinnedV
                          then l
                          else
                            let k' = parents IM.! k
                                l' = lstarMap Map.! (G.Edge k' k)
                             in followStart k' (l-l')
                in if u == v
                      then 0
                      else
                        followEnd v $ followStart u 0
             netests = foldl' (\acc utid ->
               let vmap = foldl'
                     (\mac vtid ->
                       let u = tsvs Map.! utid
                           v = tsvs Map.! vtid
                           l = getLength (G.Edge u v) pinnedV lstarMap parents'
                        in Map.insert vtid l mac
                     ) Map.empty $ Map.keys tsvs
                in Map.insert utid vmap acc
                   ) Map.empty $ Map.keys tsvs
             effectiveVariance :: G.Edge
                               -> G.Vertex -- pinned
                               -> [G.Vertex] -- vss
                               -> Set.Set G.Edge -- spanning tree
                               -> LA.Matrix Double -- D matrix
                               -> LA.Matrix Double -- X matrix
                               -> IM.IntMap G.Vertex -- parents
                               -> Hardness
             effectiveVariance (G.Edge u v) pinnedV vss spnTr dMatrix xMatrix parents =
               let r = IM.fromList $ zip vss $ replicate (length vss) 0
                   r' = case IM.lookup u r of
                               Nothing -> r
                               Just _ -> IM.insert u (-1.0) r
                   r'' = case IM.lookup v r' of
                               Nothing -> r'
                               Just _ -> IM.insert v (1.0) r'
                   forceVector = LA.tr' $ LA.fromLists [map (\v -> r'' IM.! v) vss] -- has to respect vss ordering
                   eflstar = (LA.inv (dMatrix LA.<> xMatrix)) LA.<> forceVector
                   eflstarMap = Map.fromList $
                     zip (Set.toList spnTr) (LA.toList $ LA.flatten eflstar)
                in abs $ getLength (G.Edge u v) pinnedV eflstarMap parents
             netvars = foldl' (\acc utid ->
               let vmap = foldl'
                     (\mac vtid ->
                       let u' = tsvs Map.! utid
                           v' = tsvs Map.! vtid
                           u = min u' v'
                           v = max u' v'
                           mtc = getMatrices u
                           l = if u == v
                                  then 0
                                  else
                                  --effectiveVariance (G.Edge u v) u (vss mtc) (spnTr mtc) (dMatrix mtc) (xMatrix mtc) (parents mtc)
                                  effectiveVariance (G.Edge u v) pinnedV (vss') (spnTr') (dMatrix') (xMatrix') (parents')
                        in Map.insert vtid l mac
                     ) Map.empty $ Map.keys tsvs
                in Map.insert utid vmap acc
                   ) Map.empty $ Map.keys tsvs
             sg = StudyGraph { studyGraph = G.graphFromEdges []
                             , directGraph = G.graphFromEdges $ Set.toList spnTr'
                             , vsts = vsts
                             , tsvs = tsvs
                             }
             nets = Map.empty
          in Right $ NetworkEffects { directs = sg
                                    --, leagueTable = dMatrix'
                                    --, lstar = lstar
                                    , networkEstimates = netests
                                    , networkVariances = netvars
                                    , debugmsg = unlines [ "spnTr " <> (show $ Set.toList spnTr')
                                                         , "vss " <> (show $ map (\(G.Edge u v)->v) (Set.toList spnTr'))
                                                         , "netvars" <> (show netvars)
                                                         , "dMatrix" <> (show dMatrix')
                                                         , "directEdges" <> (show directEdges')
                                                         ]
                                    }


{-# OPTIONS_GHC  -fno-warn-unused-binds -fno-warn-unused-matches -fno-warn-name-shadowing -fno-warn-missing-signatures #-}
{-# LANGUAGE FlexibleInstances, ConstraintKinds, ExistentialQuantification, GADTs, RankNTypes, MultiParamTypeClasses   #-}
{-# LANGUAGE ImpredicativeTypes, RankNTypes, UndecidableInstances, FlexibleContexts, TypeSynonymInstances #-}

---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------
-- |
-- | Module : Count-Min Sketch
-- | Creator: Xiao Ling
-- | Created: 12/17/2015
-- |
---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------

module CountMinSketch where

import Control.Monad.Random.Class
import Control.Monad.Random
import Control.Monad.State

import Data.Matrix
import Data.Conduit
import qualified Data.Conduit.List as Cl
import qualified Data.Matrix as M

import Core
import Utils
import Statistics

{-----------------------------------------------------------------------------
  Types
------------------------------------------------------------------------------}

type D        = Int              -- * num rows
type W        = Int              -- * num cols
type P        = Int              -- * some prime that is 2^n - 1
type Event    = Int              -- * some input event from universe:
                                 -- * {1, .., i, .., n}

type Sketch   = Matrix Int       -- * D x W matrix  of counters
type ABs      = Matrix Int       -- * 2 x D matrix  of a_j and b_j for j = 1..d
type ColIdx   = [Int]            -- * indices of counters to increment
type Hash     = Event -> ColIdx  -- * Maps event to d x 1 list of column indices

-- * Algorithm State
type AState    = (Sketch,Hash,(D,W,P))

-- * the algorithm uses `some m` equipped with state and randomness
type Some m    = (MonadRandom m, MonadState AState m)

-- * Algorithm outputes an `Estimator` that
-- * estimates the fequence of `event` i from  universe {1 ... i... n},
type Estimator = Event -> Int

{-----------------------------------------------------------------------------
    Count Min Sketch Conduit
------------------------------------------------------------------------------}

-- * Batch mode
sketch' :: EpsDelta -> Event -> Batch Event IO Estimator
sketch' t n = sketch t n `using` eval

-- * Streaming mode
sketch :: Some m => EpsDelta -> Event -> Streaming Event m ()
sketch t n = inits False t n $= awaitWith update'

-- * initalize algorithm given confidence/error `t` and maximum event `n`
inits :: Some m => Bool -> EpsDelta -> Event -> Conduit Event m Event
inits False t n = lift (inits' t n) >> inits True t n
inits _     _ _ = pass

-- * a concrete `eval`uation of Some `m`
-- * run with vacuous inital condition to be thrown away
eval :: StateT AState (Rand StdGen) () -> IO Estimator
eval m = fmap (\(s,h,_) -> \i -> freq i s h)
        . evalRandIO $ execStateT m $ (ones 1 1, const [], (0,0,0))

{-------------------------------------------------------------------------------
   CountMinSketch Primitives
--------------------------------------------------------------------------------}

-- * initalize algorithm with (w,d,p), zero counters, and hash function
inits' :: Some m => EpsDelta -> Event -> m ()
inits' ed n = do
    let t = newDwp ed n
    setParam t
    hash    <- newHash t
    let sketch = newSketch t
    setHash   hash
    setSketch sketch

-- * on event `i`, update the `sketch`
update' :: Some m => Event -> m ()
update' i = do
    pm      <- askParam
    sketch  <- getSketch
    hash    <- askHash
    let incr = toMask (hash i) pm
    setSketch $ incr .+. sketch

{-------------------------------------------------------------------------------
   Effectful Subroutines
--------------------------------------------------------------------------------}

-- * h_j (i) = (a_j x i + b_j mod p) mod w
-- * sumRow-wise  [ a_1 ... a_j ... a_d ]  .*.  [ i_1 ...  i_d ]     mod p  mod w
-- *              [ b_1 ... b_j ... b_d ]       [ 1   ...  1   ]
-- * result:      [ idx_1, ..., idx_d   ]
-- *               where the indices start at 1
newHash :: MonadRandom m => (D,W,P) -> m Hash
newHash t = do
    ab  <- newAbs t
    return $ toHash t ab

-- * intialize random cofficients a_j, b_j, for d x 2 matrix:
-- * [ a_1 ... a_j ... a_d ]
-- * [ b_1 ... b_j ... b_d ]
newAbs :: MonadRandom m => (D,W,P) -> m ABs
newAbs (d,_,p) = (M.fromList 2 d . take (2*d)) <$> getRandomRs (1,p)

{-----------------------------------------------------------------------------
    Pure Subroutines
------------------------------------------------------------------------------}

-- * given some event `i` from  universe {1 ... i... n}, a table of
-- * counters Sketch `s` and hash function `h`, compute the
-- * frequence of event `i`
freq :: Event -> Sketch -> Hash -> Int
freq i s h = minimum $ go s (h i) 1
    where go s []     _ = []
          go s (c:cs) d = (getElem d c s) : (go s cs $ succ d)

-- * compute d,w, and p given accuracy/confidence level and
-- * max value of event `n`
newDwp :: EpsDelta -> Event -> (D,W,P)
newDwp (ED e d) n = (round $ log2(1/d), round $ 2/e, round $ 2^n -1)

-- * initalize the sketch
newSketch :: (D,W,P) -> Sketch
newSketch (d,w,_) = zeros d w

-- * Note `is` maps some `i` from event universe onto d x 2 matrix:
-- * [ i_1 ... i_2 ]
-- * [ 1   ...  1  ]
toHash :: (D, W, P) -> ABs -> Hash
toHash (d, w, p) ab i = go $ sumR (ab .*. is) ... (\v -> v `mod` p `mod` w)
    where go = fmap (+1) . toList
          is = (i .+ (zeros 1 d)) <-> ones 1 d

-- * given hash function, construct a binary mask of size `w` x `d`
-- * to increment sketch
toMask :: ColIdx -> (D,W,P) -> Sketch
toMask ks (d,w,p) = go w ks
    where
        go w ks = let vec c w = zeros 1 (c-1) <|> single 1 <|> zeros 1 (w - c) in
            case ks of
                c:[] -> vec c w
                c:cs -> vec c w <-> go w cs

{-----------------------------------------------------------------------------
    Managing State
    TODO : redesign so this is not necessary
------------------------------------------------------------------------------}

askParam :: Some m => m (D,W,P)
askParam    = get >>= \(_,_,p) -> return p

setParam :: Some m => (D,W,P) -> m ()
setParam t  = get >>= \(s,h,_) -> put (s,h,t)

getSketch :: Some m => m Sketch
getSketch   = get >>= \(s,_,_) -> return s

setSketch :: Some m => Sketch -> m ()
setSketch s = get >>= \(_,h,t) -> put (s,h,t)

askHash :: Some m => m Hash
askHash     = get >>= \(_,h,_) -> return h

setHash :: Some m => Hash -> m ()
setHash h   = get >>= \(s,_,t) -> put (s,h,t)

{-----------------------------------------------------------------------------
    Find conduit lib function for these:
------------------------------------------------------------------------------}

-- * identity conduit, why can't I find this in the library?
pass :: Monad m => Conduit i m i
pass  = do
    mx <- await
    case mx of
        Nothing -> return ()
        Just x  -> yield x >> pass

awaitWith :: Monad m => (i -> m a) -> ConduitM i o m ()
awaitWith g = do
    mi <- await
    case mi of
        Nothing -> return ()
        Just i  -> lift (g i) >> awaitWith g


{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE DefaultSignatures #-}
-----------------------------------------------------------------------------
-- |
-- Copyright   :  (C) 2012-2015 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  portable
--
-- Operations on free vector spaces.
-----------------------------------------------------------------------------
module Linear.Vector
  ( Additive(..)
  , E(..)
  , negated
  , (^*)
  , (*^)
  , (^/)
  , sumV
  , basis
  , basisFor
  , scaled
  , outer
  , unit
  ) where

import Control.Applicative
import Control.Lens
import Data.Complex
import Data.Foldable as Foldable (forM_, foldl')
import Data.Functor.Compose
import Data.Functor.Product
import Data.HashMap.Lazy as HashMap
import Data.Hashable
import Data.IntMap as IntMap
import Data.Map as Map
import qualified Data.Vector as Vector
import Data.Vector (Vector)
import qualified Data.Vector.Mutable as Mutable
import GHC.Generics
import Linear.Instances ()

-- $setup
-- >>> import Linear.V2

-- | Basis element
newtype E t = E { el :: forall x. Lens' (t x) x }

infixl 6 ^+^, ^-^
infixl 7 ^*, *^, ^/

class GAdditive f where
  gzero :: Num a => f a
  gliftU2 :: (a -> a -> a) -> f a -> f a -> f a
  gliftI2 :: (a -> b -> c) -> f a -> f b -> f c

instance GAdditive U1 where
  gzero = U1
  {-# INLINE gzero #-}
  gliftU2 _ U1 U1 = U1
  {-# INLINE gliftU2 #-}
  gliftI2 _ U1 U1 = U1
  {-# INLINE gliftI2 #-}

instance (GAdditive f, GAdditive g) => GAdditive (f :*: g) where
  gzero = gzero :*: gzero
  {-# INLINE gzero #-}
  gliftU2 f (a :*: b) (c :*: d) = gliftU2 f a c :*: gliftU2 f b d
  {-# INLINE gliftU2 #-}
  gliftI2 f (a :*: b) (c :*: d) = gliftI2 f a c :*: gliftI2 f b d
  {-# INLINE gliftI2 #-}

instance (Additive f, GAdditive g) => GAdditive (f :.: g) where
  gzero = Comp1 $ gzero <$ (zero :: f Int)
  {-# INLINE gzero #-}
  gliftU2 f (Comp1 a) (Comp1 b) = Comp1 $ liftU2 (gliftU2 f) a b
  {-# INLINE gliftU2 #-}
  gliftI2 f (Comp1 a) (Comp1 b) = Comp1 $ liftI2 (gliftI2 f) a b
  {-# INLINE gliftI2 #-}

instance Additive f => GAdditive (Rec1 f) where
  gzero = Rec1 zero
  {-# INLINE gzero #-}
  gliftU2 f (Rec1 g) (Rec1 h) = Rec1 (liftU2 f g h)
  {-# INLINE gliftU2 #-}
  gliftI2 f (Rec1 g) (Rec1 h) = Rec1 (liftI2 f g h)
  {-# INLINE gliftI2 #-}

instance GAdditive f => GAdditive (M1 i c f) where
  gzero = M1 gzero
  {-# INLINE gzero #-}
  gliftU2 f (M1 g) (M1 h) = M1 (gliftU2 f g h)
  {-# INLINE gliftU2 #-}
  gliftI2 f (M1 g) (M1 h) = M1 (gliftI2 f g h)
  {-# INLINE gliftI2 #-}

instance GAdditive Par1 where
  gzero = Par1 0
  gliftU2 f (Par1 a) (Par1 b) = Par1 (f a b)
  {-# INLINE gliftU2 #-}
  gliftI2 f (Par1 a) (Par1 b) = Par1 (f a b)
  {-# INLINE gliftI2 #-}

-- | A vector is an additive group with additional structure.
class Functor f => Additive f where
  -- | The zero vector
  zero :: Num a => f a
#ifndef HLINT
  default zero :: (GAdditive (Rep1 f), Generic1 f, Num a) => f a
  zero = to1 gzero
#endif

  -- | Compute the sum of two vectors
  --
  -- >>> V2 1 2 ^+^ V2 3 4
  -- V2 4 6
  (^+^) :: Num a => f a -> f a -> f a
  (^+^) = liftU2 (+)
  {-# INLINE (^+^) #-}

  -- | Compute the difference between two vectors
  --
  -- >>> V2 4 5 ^-^ V2 3 1
  -- V2 1 4
  (^-^) :: Num a => f a -> f a -> f a
  x ^-^ y = x ^+^ negated y

  -- | Linearly interpolate between two vectors.
  lerp :: Num a => a -> f a -> f a -> f a
  lerp alpha u v = alpha *^ u ^+^ (1 - alpha) *^ v
  {-# INLINE lerp #-}

  -- | Apply a function to merge the 'non-zero' components of two vectors, unioning the rest of the values.
  --
  -- * For a dense vector this is equivalent to 'liftA2'.
  --
  -- * For a sparse vector this is equivalent to 'unionWith'.
  liftU2 :: (a -> a -> a) -> f a -> f a -> f a
#ifndef HLINT
  default liftU2 :: Applicative f => (a -> a -> a) -> f a -> f a -> f a
  liftU2 = liftA2
  {-# INLINE liftU2 #-}
#endif

  -- | Apply a function to the components of two vectors.
  --
  -- * For a dense vector this is equivalent to 'liftA2'.
  --
  -- * For a sparse vector this is equivalent to 'intersectionWith'.
  liftI2 :: (a -> b -> c) -> f a -> f b -> f c
#ifndef HLINT
  default liftI2 :: Applicative f => (a -> b -> c) -> f a -> f b -> f c
  liftI2 = liftA2
  {-# INLINE liftI2 #-}
#endif

instance (Additive f, Additive g) => Additive (Product f g) where
  zero = Pair zero zero
  liftU2 f (Pair a b) (Pair c d) = Pair (liftU2 f a c) (liftU2 f b d)
  liftI2 f (Pair a b) (Pair c d) = Pair (liftI2 f a c) (liftI2 f b d)
  Pair a b ^+^ Pair c d = Pair (a ^+^ c) (b ^+^ d)
  Pair a b ^-^ Pair c d = Pair (a ^-^ c) (b ^-^ d)
  lerp alpha (Pair a b) (Pair c d) = Pair (lerp alpha a c) (lerp alpha b d)

instance (Additive f, Additive g) => Additive (Compose f g) where
  zero = Compose $ zero <$ (zero :: f Int)
  {-# INLINE zero #-}
  Compose a ^+^ Compose b = Compose $ liftU2 (^+^) a b
  {-# INLINE (^+^) #-}
  Compose a ^-^ Compose b = Compose $ liftU2 (^-^) a b
  {-# INLINE (^-^) #-}
  liftU2 f (Compose a) (Compose b) = Compose $ liftU2 (liftU2 f) a b
  {-# INLINE liftU2 #-}
  liftI2 f (Compose a) (Compose b) = Compose $ liftI2 (liftI2 f) a b
  {-# INLINE liftI2 #-}

instance Additive ZipList where
  zero = ZipList []
  {-# INLINE zero #-}
  liftU2 f (ZipList xs) (ZipList ys) = ZipList (liftU2 f xs ys)
  {-# INLINE liftU2 #-}
  liftI2 = liftA2
  {-# INLINE liftI2 #-}

instance Additive Vector where
  zero = mempty
  {-# INLINE zero #-}
  liftU2 f u v = case compare lu lv of
    LT | lu == 0   -> v
       | otherwise -> Vector.modify (\ w -> Foldable.forM_ [0..lu-1] $ \i -> Mutable.unsafeWrite w i $ f (Vector.unsafeIndex u i) (Vector.unsafeIndex v i)) v
    EQ -> Vector.zipWith f u v
    GT | lv == 0   -> u
       | otherwise -> Vector.modify (\ w -> Foldable.forM_ [0..lv-1] $ \i -> Mutable.unsafeWrite w i $ f (Vector.unsafeIndex u i) (Vector.unsafeIndex v i)) u
    where
      lu = Vector.length u
      lv = Vector.length v
  {-# INLINE liftU2 #-}
  liftI2 = Vector.zipWith
  {-# INLINE liftI2 #-}

instance Additive Maybe where
  zero = Nothing
  {-# INLINE zero #-}
  liftU2 f (Just a) (Just b) = Just (f a b)
  liftU2 _ Nothing ys = ys
  liftU2 _ xs Nothing = xs
  {-# INLINE liftU2 #-}
  liftI2 = liftA2
  {-# INLINE liftI2 #-}

instance Additive [] where
  zero = []
  {-# INLINE zero #-}
  liftU2 f = go where
    go (x:xs) (y:ys) = f x y : go xs ys
    go [] ys = ys
    go xs [] = xs
  {-# INLINE liftU2 #-}
  liftI2 = Prelude.zipWith
  {-# INLINE liftI2 #-}

instance Additive IntMap where
  zero = IntMap.empty
  {-# INLINE zero #-}
  liftU2 = IntMap.unionWith
  {-# INLINE liftU2 #-}
  liftI2 = IntMap.intersectionWith
  {-# INLINE liftI2 #-}

instance Ord k => Additive (Map k) where
  zero = Map.empty
  {-# INLINE zero #-}
  liftU2 = Map.unionWith
  {-# INLINE liftU2 #-}
  liftI2 = Map.intersectionWith
  {-# INLINE liftI2 #-}

instance (Eq k, Hashable k) => Additive (HashMap k) where
  zero = HashMap.empty
  {-# INLINE zero #-}
  liftU2 = HashMap.unionWith
  {-# INLINE liftU2 #-}
  liftI2 = HashMap.intersectionWith
  {-# INLINE liftI2 #-}

instance Additive ((->) b) where
  zero   = const 0
  {-# INLINE zero #-}
  liftU2 = liftA2
  {-# INLINE liftU2 #-}
  liftI2 = liftA2
  {-# INLINE liftI2 #-}

instance Additive Complex where
  zero = 0 :+ 0
  {-# INLINE zero #-}
  liftU2 f (a :+ b) (c :+ d) = f a c :+ f b d
  {-# INLINE liftU2 #-}
  liftI2 f (a :+ b) (c :+ d) = f a c :+ f b d
  {-# INLINE liftI2 #-}

instance Additive Identity where
  zero = Identity 0
  {-# INLINE zero #-}
  liftU2 = liftA2
  {-# INLINE liftU2 #-}
  liftI2 = liftA2
  {-# INLINE liftI2 #-}

-- | Compute the negation of a vector
--
-- >>> negated (V2 2 4)
-- V2 (-2) (-4)
negated :: (Functor f, Num a) => f a -> f a
negated = fmap negate
{-# INLINE negated #-}

-- | Sum over multiple vectors
--
-- >>> sumV [V2 1 1, V2 3 4]
-- V2 4 5
sumV :: (Foldable f, Additive v, Num a) => f (v a) -> v a
sumV = Foldable.foldl' (^+^) zero
{-# INLINE sumV #-}

-- | Compute the left scalar product
--
-- >>> 2 *^ V2 3 4
-- V2 6 8
(*^) :: (Functor f, Num a) => a -> f a -> f a
(*^) a = fmap (a*)
{-# INLINE (*^) #-}

-- | Compute the right scalar product
--
-- >>> V2 3 4 ^* 2
-- V2 6 8
(^*) :: (Functor f, Num a) => f a -> a -> f a
f ^* a = fmap (*a) f
{-# INLINE (^*) #-}

-- | Compute division by a scalar on the right.
(^/) :: (Functor f, Fractional a) => f a -> a -> f a
f ^/ a = fmap (/a) f
{-# INLINE (^/) #-}

-- | Produce a default basis for a vector space. If the dimensionality
-- of the vector space is not statically known, see 'basisFor'.
basis :: (Additive t, Traversable t, Num a) => [t a]
basis = basisFor (zero :: Additive v => v Int)

-- | Produce a default basis for a vector space from which the
-- argument is drawn.
basisFor :: (Traversable t, Num a) => t b -> [t a]
basisFor = \t ->
   ifoldMapOf traversed ?? t $ \i _ ->
     return                  $
       iover  traversed ?? t $ \j _ ->
         if i == j then 1 else 0
{-# INLINABLE basisFor #-}

-- | Produce a diagonal (scale) matrix from a vector.
--
-- >>> scaled (V2 2 3)
-- V2 (V2 2 0) (V2 0 3)
scaled :: (Traversable t, Num a) => t a -> t (t a)
scaled = \t -> iter t (\i x -> iter t (\j _ -> if i == j then x else 0))
  where
  iter :: Traversable t => t a -> (Int -> a -> b) -> t b
  iter x f = iover traversed f x
{-# INLINE scaled #-}

-- | Create a unit vector.
--
-- >>> unit _x :: V2 Int
-- V2 1 0
unit :: (Additive t, Num a) => ASetter' (t a) a -> t a
unit l = set' l 1 zero

-- | Outer (tensor) product of two vectors
outer :: (Functor f, Functor g, Num a) => f a -> g a -> f (g a)
outer a b = fmap (\x->fmap (*x) b) a


{-# LANGUAGE TemplateHaskell, RankNTypes #-}

module Geometry.LRS (
    -- * Types
    Dictionary,
    -- * Functions
    getDictionary,
    selectPivot,
    reverseRS,
    lrs,
    polytope2LiftedCone,
    liftedCone2polyotpe,
    lrsFacetEnumeration,

    colFromList
    )
 where

import Geometry.Vertex
import Geometry.Matrix

import Data.Matrix as M hiding (trace)
import qualified Data.Vector as V
import Data.Maybe
import Data.List
import Util
import Control.Lens
import Debug.Trace
import Numeric.LinearProgramming hiding (simplex)
import Numeric.LinearAlgebra.HMatrix as HM
import Data.List

type Row = M.Matrix Rational
type Col = M.Matrix Rational

data Dictionary = Dict {
                            __B :: [Int],
                            __N :: [Int],
                            _dict :: M.Matrix Rational
                        }
    deriving (Show, Eq)

makeLenses ''Dictionary

numRows :: Dictionary -> Int
numRows dictionary = nrows $ dictionary^.dict

numCols :: Dictionary -> Int
numCols dictionary = ncols $ dictionary^.dict

(|*|) = multStrassen

colFromList :: [a] -> M.Matrix a
colFromList = colVector . V.fromList

rowFromList :: [a] -> M.Matrix a
rowFromList = rowVector . V.fromList

-- | Wrapper of submatrix function. This function works for indices starting from 0 rather than 1 in the submatrix function
submatrix' ::
    (Int, Int)      -- ^ Rows indices
    -> (Int, Int)   -- ^ Cols indices
    -> M.Matrix a
    -> M.Matrix a
submatrix' (ro,rk) (co,ck) m = submatrix (ro+1) (rk+1) (co+1) (ck+1) m

-- | Lens for columns
colAt :: Int -> Lens' (M.Matrix Rational) Col
colAt j = lens (getCol' j) (\m c -> setCol' j c m)

getCol' :: Int -> M.Matrix a -> M.Matrix a
getCol' idx = colVector . (getCol (idx+1))

setCol' :: Int -> Col -> M.Matrix Rational -> M.Matrix Rational
setCol' idx col mat
    | idx == 0 = col <|> right
    | idx == (ncols mat) - 1 = left <|> col
    | otherwise = left <|> col <|> right
    where
        left = submatrix' (0,(nrows mat)-1) (0,idx-1) mat
        right = submatrix' (0,(nrows mat)-1) (idx+1, (ncols mat)-1) mat

-- | Lens for rows
rowAt :: Int -> Lens' (M.Matrix Rational) Row
rowAt i = lens (getRow' i) (\m r -> setRow' i r m)

getRow' :: Int -> M.Matrix a -> M.Matrix a
getRow' idx = rowVector . (getRow (idx+1))

setRow' :: Int -> Col -> M.Matrix Rational -> M.Matrix Rational
setRow' idx row mat
    | idx == 0 = row <-> down
    | idx == (nrows mat) - 1 = up <-> row
    | otherwise = up <-> row <-> down

    where
        up = submatrix' (0,idx-1) (0,(ncols mat)-1)  mat
        down = submatrix' (idx+1, (nrows mat)-1) (0,(ncols mat)-1) mat

-- | Lens for accessing elements

elemAt :: (Int, Int) -> Lens' (M.Matrix a) a
elemAt (i, j) = lens (getElem' (i,j)) (\m x -> setElem' x (i,j) m)

getElem' :: (Int, Int) -> M.Matrix a -> a
getElem' (row, col) = getElem (row+1) (col+1)

setElem' :: a -> (Int, Int) -> M.Matrix a -> M.Matrix a
setElem' elem (row, col) mat
    | row < 0 || col < 0 || row >= nrows mat || col >= ncols mat = error ("setElem': Trying to set at (" ++ show row ++ "," ++ show col ++ ") in a " ++ show (nrows mat) ++ "x" ++ show (ncols mat) ++ " matrix.")
    | otherwise = setElem elem (row+1, col+1) mat

mapCol' :: (Int -> a -> a) -> Int -> M.Matrix a -> M.Matrix a
mapCol' f col m
    | col < 0 || col > (ncols m)-1 = error "mapCol': index out of bounds"
    | otherwise = mapCol f (col+1) m

mapRow' :: (Int -> a -> a) -> Int -> M.Matrix a -> M.Matrix a
mapRow' f row m
    | row < 0 || row > (nrows m)-1 = error "mapCol': index out of bounds"
    | otherwise = mapRow f (row+1) m

getOptimumVertex :: M.Matrix Rational -> Col -> Maybe Vertex
getOptimumVertex mat col = fmap (map toRational) $ feasNOpt result
    where
        problem = Maximize $ replicate (ncols mat) (-1)
        constraints = Dense $ safeZipWith (:<=:) (map (map fromRational) $ M.toLists mat) (map fromRational $ concat $ M.toLists col)
        result = exact problem constraints (map Free [1..ncols mat])
        feasNOpt (Optimal (_, vertex)) = Just vertex
        feasNOpt _ = Nothing

{-
    Dictionary form

                p21     p31
    invA_b
                p22     p32

 -}

toHMatrix :: M.Matrix Rational -> HM.Matrix Double
toHMatrix matrix = HM.matrix cols (map (fromRational) (concat $ M.toLists matrix))
    where
        cols = ncols matrix

findBasis :: (Maybe (M.Matrix Rational), Maybe (M.Matrix Rational)) -> M.Matrix Rational -> Int -> (Maybe (M.Matrix Rational), Maybe (M.Matrix Rational))
findBasis (basic, cobasic) original colNum
    | currRank == rows && colNum < cols = findBasis (basic, Just newCobasic) original (colNum + 1)-- Just ((fromJust cobasic) <|> (submatrix' (0,rows-1) (colNum, (ncols original)-1) original)) )
    | currRank == rows = (basic, cobasic)-- Just ((fromJust cobasic) <|> (submatrix' (0,rows-1) (colNum, (ncols original)-1) original)) )
    | currRank == candRank = findBasis (basic, Just newCobasic) original (colNum + 1)
    | currRank < candRank = findBasis (Just newBasic, cobasic) original (colNum + 1)
    where
        rows = nrows original
        cols = ncols original
        currRank = if isNothing basic then 0 else rank (toHMatrix $ fromJust basic)
        candVec = original ^. colAt colNum
        candMat = if isNothing basic then candVec else (fromJust basic) <|> candVec
        candRank = rank $ toHMatrix candMat

        newBasic = if isNothing basic then candVec else (fromJust basic) <|> candVec
        newCobasic = if isNothing cobasic then candVec else (fromJust cobasic) <|> candVec

sortSystem :: M.Matrix Rational -> (M.Matrix Rational, M.Matrix Rational)
sortSystem matrix = trace ("Matrix : "++ show matrix ++ "\n\nSorted " ++ show (findBasis (Nothing, Nothing) matrix 0) ) (((,) <$> (fromJust . fst) <*> (fromJust . snd)) (findBasis (Nothing, Nothing) matrix 0))
-- sortSystem :: M.Matrix Rational -> Col -> Vertex -> (M.Matrix Rational, Col)

-- sortSystem mat col vertex
--     | cols /=  (rows `div` 2) = (,) mat col
--     | otherwise = (trace $ show $ M.fromLists (upper ++ lower)) (,) (M.fromLists (upper ++ lower)) col
--     where
--         matLists = M.toLists mat

--         rows = nrows mat
--         cols = ncols mat
--         slackMatr = submatrix' (0,rows-1) (0, rows - cols -2) $ identity rows
--         lower  = getIndependent matLists [] 0 rows slackMatr
--         upper = matLists \\ lower

-- getDictionary :: Matrix Rational -> Col -> Vertex -> Dictionary
-- getDictionary _A b vertex = trace ("In dictionary" ++ show (rows,cols, nrows dictionary, ncols dictionary) ++ show newDict) Dict [0..rows] [rows+1..rows+cols] newDict
--     -- trace ("In dictionary" ++ show _A_B)
--     where
--         (newA, newb) = sortSystem _A b vertex
--         rows = nrows newA
--         cols = ncols newA
--         slack = identity rows
--        -- dictionary = newA <|> slack
--         topRow = rowFromList $ 1 : replicate rows 0 ++  replicate cols (1)
--         leftCol = colFromList $ replicate rows 0
--         dictionary = topRow <-> (leftCol <|> slack <|> newA )

--         -- c_B = submatrix' (0,0) (1,rows) topRow
--         -- c_N = submatrix' (0,0) (rows+1, rows+cols) topRow
--         -- _A_B = submatrix' (0,rows-1) (0,rows-1) dictionary
--         -- _A_N = submatrix' (0,rows-1) (rows, rows+cols-1) dictionary
--         _A_B = submatrix' (0,rows) (0,rows) dictionary
--         _A_N = submatrix' (0,rows) (rows+1, rows+cols) dictionary
--         Right invA_B = inverse _A_B

--         p2 = - invA_B |*| _A_N
--         p3 = invA_B |*| newb
--         -- p21 = (c_N - c_B |*| invA_B |*| _A_N)
--         -- p22 = invA_B |*| _A_N
--         -- p31 = -c_B |*| invA_B |*| newb
--         -- p32 = invA_B |*| newb
--         newDict = ((identity (rows+1)) <|> p2 <|> p3)

-- | Constructs initial dictionary given a set of inequalities and an optimum initial vertex
-- getDictionary :: M.Matrix Rational -> Col -> Vertex -> Dictionary
-- getDictionary _A b vertex = Dict [0..rows] [rows+1..rows+cols] newDict
--     where
--         -- (newA, newb) = sortSystem _A b vertex
--         initialDict = _A <|> slack
--         (basic, cobasic) = sortSystem initialDict
--         (newA, newb) = (,) _A b
--         rows = nrows newA
--         cols = ncols newA
--         slack = identity rows
--         dictionary = newA <|> slack
--         topRow = rowFromList $ 1 :  replicate rows 0 ++  replicate cols 1
--         c_B = submatrix' (0,0) (1,rows) topRow
--         c_N = submatrix' (0,0) (rows+1, rows+cols) topRow
--         _A_B = submatrix' (0,rows-1) (0,rows-1) dictionary
--         _A_N = submatrix' (0,rows-1) (rows, rows+cols-1) dictionary
--         Right invA_B = inverse _A_B
--         p21 = (c_N - c_B |*| invA_B |*| _A_N)
--         p22 = invA_B |*| _A_N
--         p31 = -c_B |*| invA_B |*| newb
--         p32 = invA_B |*| newb

--         newDict = ((identity (rows+1)) <|> (p21 <-> p22) <|> (p31 <-> p32))

getDictionary :: M.Matrix Rational -> Col -> Vertex -> Dictionary
getDictionary _A b vertex = Dict [0..rows] [rows+1..rows+cols] newDict
    where

        -- (newA, newb) = sortSystem _A b vertex
        initialDict = _A <|> slack
        (basic, cobasic) = sortSystem initialDict
        (newA, newb) = (,) _A b
        rows = nrows newA
        cols = ncols newA
        slack = identity rows
        -- dictionary = newA <|> slack
        topRow = rowFromList $ 1 :  replicate rows 0 ++  replicate cols 1
        c_B = submatrix' (0,0) (1,rows) topRow
        c_N = submatrix' (0,0) (rows+1, rows+cols) topRow
        _A_B = basic --submatrix' (0,rows-1) (0,rows-1) dictionary
        _A_N = cobasic --submatrix' (0,rows-1) (rows, rows+cols-1) dictionary
        Right invA_B = inverse _A_B
        p21 = (c_N - c_B |*| invA_B |*| _A_N)
        p22 = invA_B |*| _A_N
        p31 = -c_B |*| invA_B |*| newb
        p32 = invA_B |*| newb

        newDict = ((identity (rows+1)) <|> (p21 <-> p22) <|> (p31 <-> p32))

enteringVariable :: Dictionary -> Maybe Int
enteringVariable dictionary
    | null negs = Nothing
    | otherwise = trace ("ENTERINGVAR\nBASE: " ++ show (dictionary^._B) ++ "\nCOBASE: " ++ show (dictionary^._N))  (Just $ (fst.head.sort) negs)
    where
        rows = nrows $ dictionary^.dict
        cobasic_0 = zip (dictionary^._N) (map (\j -> dictionary^.dict.elemAt (0,j)) [rows..])
        negs = filter (\(_,value) -> value < 0) cobasic_0

lexMinRatio :: Dictionary -> Int -> Int
lexMinRatio dictionary s
    | dim+1 == rows = 0
    | null indexed_s = 0
    | otherwise = trace ("LEXMINRATIO\nBASE: " ++ show (dictionary^._B) ++ "\nCOBASE: " ++ show (dictionary^._N))  (dictionary ^. _B) !! (fst $ indexed_s !! (fromJust $ elemIndex (minimum ratios) ratios))
    where
        rows = numRows dictionary
        cols = numCols dictionary
        dim = cols - rows - 1
        dictMatrix = dictionary^.dict
        col_s = dictMatrix ^. colAt s
        _D = (dictMatrix ^. colAt (cols-1)) <|> submatrix' (0,rows-1) (0, rows-1) dictMatrix
        slice_s = (concat.M.toLists) $ submatrix' (dim+1, rows-1) (0,0) col_s
        indexed_s = filter (snd.(fmap (>0))) $ safeZipWith (,) [dim+1..rows-1] slice_s
        sub_D = map (head . M.toLists . (\i -> _D ^. rowAt i) . fst) indexed_s
        ratios = safeZipWith (map $) (map ((flip (/)) . snd) indexed_s) sub_D

-- | Selects entering and leaving variables for the next pivot in the dictionary
selectPivot :: Dictionary -> Maybe (Int, Int)
selectPivot dictionary
    | s == Nothing = Nothing
    | otherwise = Just (r, fromJust s)
    where
        s = enteringVariable dictionary
        r = lexMinRatio dictionary (fromJust s)

pivot :: Int ->Int -> Dictionary -> Dictionary
pivot r s dictionary = Dict new_B new_N (_E |*| dictMatrix)
    where
        dictMatrix = dictionary ^. dict
        col_s = dictMatrix ^. colAt (s)
        t = fromJust $ elemIndex r (dictionary ^. _B) -- idxR
        idxS = fromJust $ elemIndex s (dictionary ^. _N)
        a_t = col_s ^. elemAt (t,0)
        eta = (mapCol' (\_ x -> -x/a_t) 0 col_s) & elemAt (t,0) .~ (1/a_t)
        _E = (identity (nrows dictMatrix)) & colAt t .~ eta
        new_B = dictionary^._B & element t .~ s
        new_N = dictionary^._N & element idxS .~ r

simplex :: Dictionary -> Dictionary
simplex dictionary
    | idxsToPivot == Nothing = dictionary
    | otherwise = simplex $ uncurry pivot (fromJust idxsToPivot) dictionary
    where
        idxsToPivot = selectPivot dictionary

-- | Reverse search.
reverseRS ::
        Dictionary
    ->  Int          -- ^ element in N
    ->  Maybe Int
reverseRS dictionary v
    | conditions == False = trace ("REVERSE RS: " ++ show [firstCondition, sndCondition, lastCondition]) Nothing
    | conditions == True = trace ("REVERSE RS: " ++ show v) Just u
    where
        dictMatrix = dictionary^.dict
        rows = nrows dictMatrix
        v_col = dictMatrix ^. colAt (v)
        w_row_0 = dictMatrix ^. rowAt 0
        u = lexMinRatio dictionary v
        i = fromJust $ elemIndex u (dictionary ^. _B)
        w_row_i = mapRow' (\_ x -> (v_col ^. elemAt (0,0))/(v_col ^. elemAt (i,0)) * x) 0 $ dictMatrix ^. rowAt (i)
        diff_ws = (head.M.toLists) $ w_row_0 - w_row_i
        firstCondition = (w_row_0 ^. elemAt (0,v)) > 0
        sndCondition = u /= 0
        lastCondition = all (>=0) [(diff_ws!!(snd idx))|idx <- (zip (dictionary^._N) [rows..]), (fst idx) < u ]
        conditions = firstCondition && sndCondition && lastCondition -- Proposition 6.1

-- reverseRS :: -- reverse with pivot and selectPivot
--     Dictionary ->
--     Int ->
--     Maybe Int
-- reverseRS dictionary v
--     | newPivots == Nothing = Nothing
--     | condition == False = Nothing
--     | condition == True = Just u
--     where
--         u = lexMinRatio dictionary v
--         prev_B = pivot u v dictionary
--         newPivots = selectPivot prev_B
--         condition = fst (fromJust newPivots) == v && (snd (fromJust newPivots)) == u

getVertex :: Dictionary -> Vertex
getVertex dictionary = concat $ M.toLists $ submatrix' (1,dim) (cols-1, cols-1) (dictionary^.dict)
    where
        rows = numRows dictionary
        cols = numCols dictionary
        dim = cols - rows - 1

revSearch :: Dictionary -> [Vertex]
revSearch dictionary@(Dict _B _N dictMatrix) -- = getVertex dictionary : concatMap revSearch pivoted
    | (not.null) possibleRay = possibleRay ++ (concatMap revSearch pivoted)
    | otherwise = getVertex dictionary : concatMap revSearch pivoted
    where
        rows = numRows dictionary
        cols = numCols dictionary
        valid_N = [i | i <- _N, (reverseRS dictionary i) /= Nothing]
        valid_B = map (lexMinRatio dictionary) valid_N
        pivoted = map (\(r, s) ->  pivot r s dictionary) $ zip valid_B valid_N
        possibleRay = hasRay dictionary

-- lexMin :: Dictionary -> Maybe Extremal
-- lexMin dictionary

hasRay :: Dictionary -> [Vertex]
hasRay dictionary = trace ("RAYS: " ++ show (map snd rays) ++ "\n\n" ++ show (map fst rays) ++ "\n" ++ show (dictionary ^. dict) ++ "\n\n") (map snd rays)
    -- idxsPivot == Nothing = Nothing
    -- r == 0 = Just ray
    -- otherwise = Nothing
    where
        dictMatrix = dictionary ^. dict
        rows = numRows dictionary
        dim = cols - rows - 1
        cols = numCols dictionary
        nonPositive column = all (<=0) ((concat.M.toLists) column)
        colsWithRays = if dim+1 == rows then
                            [(j, dictMatrix^.colAt j) | j <- [rows..cols-2]]
                        else [(j, dictMatrix^.colAt j) | j <- [rows..cols-2], nonPositive (submatrix' (dim+1,rows-1) (j,j) dictMatrix )]
        rays =  zip (map fst colsWithRays) (map (concat . M.toLists . (submatrix' (1,dim) (0,0)) . snd ) colsWithRays)

-- | Implementation of Lexicographic Reverse Search
lrs :: M.Matrix Rational -> Col -> Vertex-> [Vertex]
lrs matrix b vertex = (sort.nub) $ revSearch dictionary
    where
        dictionary = simplex $ getDictionary matrix b vertex

-- | Transforms a V polytope into a H pointed cone
polytope2LiftedCone :: [Vertex] -> (M.Matrix Rational, Col)
polytope2LiftedCone points = (colOnes <|> hMatrix , colZeros)
    where
        nrows = length points
        colOnes = colFromList $ replicate nrows 1
        hMatrix = M.fromLists points
        colZeros = colFromList $ replicate nrows 0

-- | Transforms a V pointed cone into a H polyotpe
liftedCone2polyotpe :: [Vertex] -> (M.Matrix Rational, Col)
liftedCone2polyotpe points = (hMatrix, b)
    where
        dim = length (points !! 0)
        rays  = points \\ [(replicate dim 0)] -- remove vertex at the origin
        b = colFromList $ map (negate.head) rays
        hMatrix = M.fromLists $ map tail rays

-- | Implementation of Lexicographic Reverse Search for Facet Enumeration
lrsFacetEnumeration :: [IVertex] -> (M.Matrix Rational, Col)
lrsFacetEnumeration ipoints = hPolytope
    where
        points =  map (map toRational) ipoints
        cols = length (points!!0)
        initialVertex = replicate (cols + 1) 0
        hLiftedCone@(matrix, b) = polytope2LiftedCone points
        vLiftedCone = lrs matrix b initialVertex
        hPolytope = liftedCone2polyotpe vLiftedCone


{-# LANGUAGE QuasiQuotes, MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell #-}
module Lib where

import qualified Language.C.Inline as C
import Foreign.C.Types
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA
import qualified Numeric.LinearAlgebra.Data as D
import qualified Numeric.LinearAlgebra.Devel as D
import Control.Arrow ((***))
import Data.Tuple (swap)
C.context (C.baseCtx <> C.vecCtx)

C.include "<math.h>"
C.include "<osqp/osqp.h>"

testsp = D.mkSparse [((0,0), 3.0), ((1,1), 2.0)]
testcsr = D.mkCSR [((0,0), 3.0), ((1,1), 2.0)]

-- getCSR (D.SparseR csr ncols nrows) =

assocVec :: (Num a, Ord a) => [a] -> [(Int,a)]
assocVec v = filter (\(_,x) -> x /= 0) (zipWith (,) [0..] v)

-- conveniently, will assume that if your list stops, it is implied the rest of it is zero. Or maybe inconvenitly from a type safety persepctive
assocMat :: (Num a, Ord a) => [[a]] -> [((Int,Int),a)]
assocMat m = let m' = map assocVec m in addRows m' where addRows m'' = concat $ zipWith (\i xs -> map (\(j,x) -> ((i,j),x)) xs) [0 ..] m''
-- AssocMatrix = [((Int, Int), Double)]

transposeAssoc = map (swap *** id)

mkGCSC = D.tr . D.mkSparse . transposeAssoc
mkGCSC' = mkGCSC . assocMat
mkCSC :: D.AssocMatrix -> CSC'
mkCSC = transpose'' . D.mkCSR . transposeAssoc
mkCSC' :: [[Double]] -> CSC'
mkCSC' = mkCSC . assocMat
-- The CSC type is internal to hmatrix.
data CSC' = CSC'
        { cscVals  :: V.Vector Double
        , cscRows  :: V.Vector CInt
        , cscCols  :: V.Vector CInt
        , cscNRows :: Int
        , cscNCols :: Int
        } deriving Show

-- data OSQPCSC = CSC  -- {nzmax :: CInt, m :: CInt, n :: CInt, col_p :: V.Vector Ptr, col  }
{-
instance LA.Transposable D.CSR CSC'
  where
  tr (D.CSR vs cs rs n m) = CSC' vs cs rs m n
  tr' = tr

instance LA.Transposable CSC' D.CSR
  where
  tr (CSC' vs rs cs n m) = D.CSR vs rs cs m n
  tr' = tr
-}
transpose'' (D.CSR vs cs rs n m) = CSC' vs cs rs m n
transpose''' (CSC' vs rs cs n m) = D.CSR vs rs cs m n

funboy :: IO ()
funboy = do
  x <- [C.exp| double{ cos(1) } |]
  print x
someFunc :: IO ()
someFunc = putStrLn "someFunc"

myfun :: CDouble -> CDouble
myfun x = [C.pure| double{ cos( $(double x)) } |]

-- exampleOSQPData =

-- Actually CDoubles
--data OSQPData = OSQPData {dn :: CInt, dm :: CInt, p :: CSC ,
--     a :: CSC, q :: V.Vector CDouble, l :: V.Vector CDouble, u :: V.Vector CDouble}

-- OSQPWorkspace *osqp_setup(const OSQPData *data, OSQPSettings *settings)
-- c_int osqp_solve(OSQPWorkspace *work)
-- c_int osqp_cleanup(OSQPWorkspace *work)

-- Maybe the easiest thing is to just inject default settings and marshal only what we need
-- Get a demo code in full c and start pulling pieces out.
-- l and u are the easiest

demo :: IO CInt
demo = [C.block|
int {
    // Load problem data
    c_float P_x[4] = {4.00, 1.00, 1.00, 2.00, };
    c_int P_nnz = 4;
    c_int P_i[4] = {0, 1, 0, 1, };
    c_int P_p[3] = {0, 2, 4, };
    c_float q[2] = {1.00, 1.00, };
    c_float A_x[4] = {1.00, 1.00, 1.00, 1.00, };
    c_int A_nnz = 4;
    c_int A_i[4] = {0, 1, 0, 2, };
    c_int A_p[3] = {0, 2, 4, };
    c_float l[3] = {1.00, 0.00, 0.00, };
    c_float u[3] = {1.00, 0.70, 0.70, };
    c_int n = 2;
    c_int m = 3;

    // Problem settings
    OSQPSettings * settings = (OSQPSettings *)c_malloc(sizeof(OSQPSettings));

    // Structures
    OSQPWorkspace * work;  // Workspace
    OSQPData * data;  // OSQPData

    // Populate data
    data = (OSQPData *)c_malloc(sizeof(OSQPData));
    data->n = n;
    data->m = m;
    data->P = csc_matrix(data->n, data->n, P_nnz, P_x, P_i, P_p);
    data->q = q;
    data->A = csc_matrix(data->m, data->n, A_nnz, A_x, A_i, A_p);
    data->l = l;
    data->u = u;

    // Define Solver settings as default
    osqp_set_default_settings(settings);
    settings->alpha = 1.0; // Change alpha parameter

    // Setup workspace
    work = osqp_setup(data, settings);

    // Solve Problem
    osqp_solve(work);

    // Cleanup
    osqp_cleanup(work);
    c_free(data->A);
    c_free(data->P);
    c_free(data);
    c_free(settings);

    return 0;
} |]


{-# LANGUAGE FlexibleContexts #-}
module FokkerPlanck.DomainChange where

import           Data.Array.Repa       as R
import           Data.Complex
import           Data.List             as L
import           FokkerPlanck.Types
import           Numeric.LinearAlgebra as NL
import           Types
import           Utils.Array

{-# INLINE r2z1t0Tor2s1t0 #-}
r2z1t0Tor2s1t0 :: Int -> [Double] -> R2Z1T0Array -> R2S1T0Array
r2z1t0Tor2s1t0 numOrientations freqs arr =
  let deltaTheta = 2 * pi / (fromIntegral numOrientations)
      (Z :. numFreqs :. numT0Freqs :. xLen :. yLen) = extent arr
      mat1 =
        (numOrientations >< numFreqs) . R.toList $
        R.traverse
          (fromListUnboxed (Z :. numFreqs) freqs)
          (const (Z :. numOrientations :. numFreqs))
          (\f (Z :. i :. j) ->
             exp (0 :+ (-1) * (deltaTheta * fromIntegral i) * f (Z :. j)))
      mat2 = (numFreqs >< (xLen * yLen * numT0Freqs)) . R.toList $ arr
   in fromListUnboxed (Z :. numOrientations :. numT0Freqs :. xLen :. yLen) .
      NL.toList . flatten $
      mat1 NL.<> mat2

{-# INLINE r2z1Tor2s1 #-}
r2z1Tor2s1 ::
     (Source s (Complex Double))
  => Int
  -> [Double]
  -> R.Array s DIM3 (Complex Double)
  -> R.Array U DIM3 (Complex Double)
r2z1Tor2s1 numOrientations freqs arr =
  let deltaTheta = 2 * pi / (fromIntegral numOrientations)
      (Z :. numFreqs :. xLen :. yLen) = extent arr
      mat1 =
        (numOrientations >< numFreqs) . R.toList $
        R.traverse
          (fromListUnboxed (Z :. numFreqs) freqs)
          (const (Z :. numOrientations :. numFreqs))
          (\f (Z :. i :. j) ->
             exp (0 :+ (-1) * (deltaTheta * fromIntegral i) * f (Z :. j)))
      mat2 = (numFreqs >< (xLen * yLen)) . R.toList $ arr
   in fromListUnboxed (Z :. numOrientations :. xLen :. yLen) .
      NL.toList . flatten $
      mat1 NL.<> mat2

{-# INLINE r2s1tor2z1 #-}
r2s1tor2z1 ::
     (Source s (Complex Double))
  => [Double]
  -> R.Array s DIM3 (Complex Double)
  -> R.Array U DIM3 (Complex Double)
r2s1tor2z1 freqs arr =
  let (Z :. numOrientations :. xLen :. yLen) = extent arr
      deltaTheta = 2 * pi / (fromIntegral numOrientations)
      numFreqs = L.length freqs
      mat1 =
        (numFreqs >< numOrientations) . R.toList $
        R.traverse
          (fromListUnboxed (Z :. numFreqs) freqs)
          (const (Z :. numFreqs :. numOrientations))
          (\f (Z :. i :. j) ->
             exp (0 :+ (deltaTheta * fromIntegral j) * f (Z :. i)))
      mat2 = (numOrientations >< (xLen * yLen)) . R.toList $ arr
   in fromListUnboxed (Z :. numFreqs :. xLen :. yLen) . NL.toList . flatten $
      mat1 NL.<> mat2

{-# INLINE r2z2t0s0Tor2s1rpt0s0 #-}
r2z2t0s0Tor2s1rpt0s0 ::
     Int -> [Double] -> Int -> [Double] -> R2Z2T0S0Array -> R2S1RPT0S0Array
r2z2t0s0Tor2s1rpt0s0 numOrientations thetaFreqs numScales scaleFreqs arr =
  let deltaTheta = 2 * pi / (fromIntegral numOrientations)
      (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :. numScale0Freq :. xLen :. yLen) =
        extent arr
      mat1 =
        ((numOrientations * numScales) >< (numThetaFreq * numScaleFreq)) .
        R.toList $
        R.traverse3
          (fromFunction
             (Z :. numOrientations :. numScales :. numThetaFreq :. numScaleFreq)
             (const 0))
          (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
          (fromListUnboxed (Z :. numScaleFreq) scaleFreqs)
          (\a _ _ -> a)
          (\f1 f2 f3 (Z :. i :. j :. k :. l) ->
             exp
               (0 :+
                (-1) *
                ((deltaTheta * fromIntegral i) * f2 (Z :. k) +
                 (2 * pi * fromIntegral j / fromIntegral numScales) *
                 f3 (Z :. l))))
      mat2 =
        ((numThetaFreq * numScaleFreq) ><
         (xLen * yLen * numTheta0Freq * numScale0Freq)) .
        R.toList $
        arr
   in fromListUnboxed
        (Z :. numOrientations :. numScales :. numTheta0Freq :. numScale0Freq :.
         xLen :.
         yLen) .
      NL.toList . flatten $
      mat1 NL.<> mat2

{-# INLINE r2z2Tor2s1rp #-}
r2z2Tor2s1rp ::
     (Source s (Complex Double))
  => Int
  -> [Double]
  -> Int
  -> [Double]
  -> Double
  -> R.Array s DIM4 (Complex Double)
  -> R.Array U DIM4 (Complex Double)
r2z2Tor2s1rp numOrientations thetaFreqs numScales scaleFreqs maxScale arr =
  let deltaTheta = 2 * pi / (fromIntegral numOrientations)
      deltaScale = log maxScale / (fromIntegral numScales)
      (Z :. numThetaFreq :. numScaleFreq :. xLen :. yLen) = extent arr
      mat1 =
        ((numOrientations * numScales) >< (numThetaFreq * numScaleFreq)) .
        R.toList $
        R.traverse3
          (fromFunction
             (Z :. numOrientations :. numScales :. numThetaFreq :. numScaleFreq)
             (const 0))
          (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
          (fromListUnboxed (Z :. numScaleFreq) scaleFreqs)
          (\a _ _ -> a)
          (\f1 f2 f3 (Z :. i :. j :. k :. l) ->
             exp
               (0 :+
                (1) *
                ((deltaTheta * fromIntegral i) * f2 (Z :. k) +
                 (2 * pi * fromIntegral j * deltaScale / (log maxScale)) *
                 f3 (Z :. l))))
      mat2 = ((numThetaFreq * numScaleFreq) >< (xLen * yLen)) . R.toList $ arr
  in fromListUnboxed (Z :. numOrientations :. numScales :. xLen :. yLen) .
     NL.toList . flatten $
     mat1 NL.<> mat2

{-# INLINE r2z2Tor2s1rpP #-}
r2z2Tor2s1rpP ::
     (Source s (Complex Double))
  => Int
  -> [Double]
  -> Int
  -> [Double]
  -> Double
  -> R.Array s DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
r2z2Tor2s1rpP numOrientations thetaFreqs numScales scaleFreqs maxR arr = do
  let deltaTheta = 2 * pi / (fromIntegral numOrientations)
      deltaScale = log maxR / (fromIntegral numOrientations)
      (Z :. numThetaFreq :. numScaleFreq :. xLen :. yLen) = extent arr
      mat2 = ((numThetaFreq * numScaleFreq) >< (xLen * yLen)) . R.toList $ arr
  mat1 <-
    fmap
      (((numOrientations * numScales) >< (numThetaFreq * numScaleFreq)) .
       R.toList) .
    computeUnboxedP $
    R.traverse3
      (fromFunction
         (Z :. numOrientations :. numScales :. numThetaFreq :. numScaleFreq)
         (const 0))
      (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
      (fromListUnboxed (Z :. numScaleFreq) scaleFreqs)
      (\a _ _ -> a)
      (\f1 f2 f3 (Z :. i :. j :. k :. l) ->
         exp
           (0 :+
            (1) *
            ((deltaTheta * fromIntegral i) * f2 (Z :. k) +
             (2 * pi * fromIntegral j * deltaScale / (log maxR)) * f3 (Z :. l))))
  return .
    fromListUnboxed (Z :. numOrientations :. numScales :. xLen :. yLen) .
    NL.toList . flatten $
    mat1 NL.<> mat2

{-# INLINE r2z2t0s0Tor2s1rps1rp #-}
r2z2t0s0Tor2s1rps1rp ::
     Int -> [Double] -> [Double] -> Int -> [Double] -> [Double] -> Double -> R2Z2T0S0Array -> R.Array U DIM6 (Complex Double)
r2z2t0s0Tor2s1rps1rp numOrientations thetaFreqs theta0Freqs numScales scaleFreqs scale0Freqs maxScale arr =
  let deltaTheta = 2 * pi / (fromIntegral numOrientations)
      deltaScale =
        if maxScale == 0
          then 0
          else (2 * pi / (log maxScale)) / (fromIntegral numScales)
      (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :. numScale0Freq :. xLen :. yLen) =
        extent arr
      mat1 =
        ((numOrientations * numScales * numOrientations * numScales) ><
         (numThetaFreq * numScaleFreq * numTheta0Freq * numScale0Freq)) .
        R.toList $
        R.traverse4
          (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
          (fromListUnboxed (Z :. numScaleFreq) scaleFreqs)
          (fromListUnboxed (Z :. numTheta0Freq) theta0Freqs)
          (fromListUnboxed (Z :. numScale0Freq) scale0Freqs)
          (\_ _ _ _ ->
             (Z :. numOrientations :. numScales :. numOrientations :. numScales :.
              numThetaFreq :.
              numScaleFreq :.
              numTheta0Freq :.
              numScale0Freq))
          (\ft fs ft0 fs0 (Z :. o :. s :. o0 :. s0 :. tf :. sf :. tf0 :. sf0) ->
             exp
               (0 :+
                (-1) *
                ((deltaTheta * fromIntegral o) * ft (Z :. tf) +
                 (deltaTheta * fromIntegral o0) * ft0 (Z :. tf0) +
                 (deltaScale * fromIntegral s) * fs (Z :. sf) +
                 (deltaScale * fromIntegral s0) * fs0 (Z :. sf0))))
      mat2 =
        ((numThetaFreq * numScaleFreq * numTheta0Freq * numScale0Freq) ><
         (xLen * yLen)) .
        R.toList $
        arr
   in fromListUnboxed
        (Z :. numOrientations :. numScales :. numOrientations :. numScales :.
         xLen :.
         yLen) .
      NL.toList . flatten $
      mat1 NL.<> mat2

{-# INLINE r2s1rps1rpTor2z2t0s0 #-}
r2s1rps1rpTor2z2t0s0 ::
     Int
  -> [Double]
  -> [Double]
  -> Int
  -> [Double]
  -> [Double]
  -> Double
  -> R.Array U DIM6 (Complex Double)
  -> R2Z2T0S0Array
r2s1rps1rpTor2z2t0s0 numOrientations thetaFreqs theta0Freqs numScales scaleFreqs scale0Freqs maxScale arr =
  let deltaTheta = 2 * pi / (fromIntegral numOrientations)
      deltaScale =
        if maxScale == 0
          then 0
          else (2 * pi / (log maxScale)) / (fromIntegral numScales)
      numThetaFreq = L.length thetaFreqs
      numScaleFreq = L.length scaleFreqs
      numTheta0Freq = L.length theta0Freqs
      numScale0Freq = L.length scale0Freqs
      (Z :. _ :. _ :. _ :. _ :. xLen :. yLen) =
        extent arr
      mat1 =
        ((numThetaFreq * numScaleFreq * numTheta0Freq * numScale0Freq) ><
         (numOrientations * numScales * numOrientations * numScales)) .
        R.toList $
        R.traverse4
          (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
          (fromListUnboxed (Z :. numScaleFreq) scaleFreqs)
          (fromListUnboxed (Z :. numTheta0Freq) theta0Freqs)
          (fromListUnboxed (Z :. numScale0Freq) scale0Freqs)
          (\_ _ _ _ ->
             (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :.
              numScale0Freq :.
              numOrientations :.
              numScales :.
              numOrientations :.
              numScales))
          (\ft fs ft0 fs0 (Z :. tf :. sf :. tf0 :. sf0 :. o :. s :. o0 :. s0) ->
             exp
               (0 :+ (1) *
                ((deltaTheta * fromIntegral o) * ft (Z :. tf) +
                 (deltaTheta * fromIntegral o0) * ft0 (Z :. tf0) +
                 (deltaScale * fromIntegral s) * fs (Z :. sf) +
                 (deltaScale * fromIntegral s0) * fs0 (Z :. sf0))))
      mat2 =
        ((numOrientations * numScales * numOrientations * numScales) ><
         (xLen * yLen)) .
        R.toList $
        arr
   in fromListUnboxed
        (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :. numScale0Freq :.
         xLen :.
         yLen) .
      NL.toList . flatten $
      mat1 NL.<> mat2

-- x y t s
{-# INLINE stsTor2s1rp #-}
stsTor2s1rp ::
     Int
  -> Double
  -> [Double]
  -> Int
  -> Double
  -> [Double]
  -> Int
  -> [Double]
  -> Int
  -> [Double]
  -> Double
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
stsTor2s1rp numX xLen xFreqs numY yLen yFreqs numOrientations thetaFreqs numScales scaleFreqs maxScale arr = do
  let deltaX = xLen / fromIntegral numX
      deltaY = yLen / fromIntegral numY
      deltaTheta = 2 * pi / fromIntegral numOrientations
      deltaScale = (log maxScale) / fromIntegral numScales
      centerX = div numX 2
      centerY = div numY 2
      vec = NL.fromList . R.toList $ arr
  transformMat <-
    fmap
      (((numX * numY * numOrientations * numScales) ><
        (L.product . listOfShape . extent $ arr)) .
       R.toList) .
    computeUnboxedP .
    R.traverse4
      (fromListUnboxed (Z :. L.length xFreqs) xFreqs)
      (fromListUnboxed (Z :. L.length yFreqs) yFreqs)
      (fromListUnboxed (Z :. L.length thetaFreqs) thetaFreqs)
      (fromListUnboxed (Z :. L.length scaleFreqs) scaleFreqs)
      (\(Z :. numXFreq) (Z :. numYFreq) (Z :. numThetaFreq) (Z :. numScaleFreq) ->
         (Z :. numX :. numY :. numOrientations :. numScales :. numXFreq :.
          numYFreq :.
          numThetaFreq :.
          numScaleFreq)) $ \fx fy fo fs (Z :. x :. y :. o :. s :. xf :. yf :. orif :. sf) ->
      exp
        (0 :+
         (fo (Z :. orif) * fromIntegral o * deltaTheta +
          2 * pi *
          (fs (Z :. sf) * fromIntegral s * deltaScale / (log maxScale) +
           fx (Z :. xf) * fromIntegral ((x :: Int) - centerX) * deltaX / xLen +
           fy (Z :. yf) * fromIntegral ((y :: Int) - centerY) * deltaY / yLen)))
  return .
    fromListUnboxed (Z :. numX :. numY :. numOrientations :. numScales) .
    NL.toList $
    transformMat NL.#> vec

{-# INLINE stsTor2 #-}
stsTor2 ::
     Int
  -> Double
  -> [Double]
  -> Int
  -> Double
  -> [Double]
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM2 (Complex Double))
stsTor2 numX xLen xFreqs numY yLen yFreqs arr = do
  let deltaX = xLen / fromIntegral numX
      deltaY = yLen / fromIntegral numY
      centerX = div numX 2
      centerY = div numY 2
  vec <-
    fmap (NL.fromList . R.toList) .
    sumP . sumS . R.map (\x -> (magnitude x) ^ 2 :+ 0) $
    arr
  transformMat <-
    fmap
      (((numX * numY) >< (L.product . L.drop 2 . listOfShape . extent $ arr)) .
       R.toList) .
    computeUnboxedP .
    R.traverse2
      (fromListUnboxed (Z :. L.length xFreqs) xFreqs)
      (fromListUnboxed (Z :. L.length yFreqs) yFreqs)
      (\(Z :. numXFreq) (Z :. numYFreq) ->
         (Z :. numX :. numY :. numXFreq :. numYFreq)) $ \fx fy (Z :. x :. y :. xf :. yf) ->
      exp
        (0 :+
         (2 * pi *
          (fx (Z :. xf) * fromIntegral (x - centerX) * deltaX / xLen +
           fy (Z :. yf) * fromIntegral (y - centerY) * deltaY / yLen)))
  return . fromListUnboxed (Z :. numX :. numY) . NL.toList $
    transformMat NL.#> vec

{-# INLINE stsTor2' #-}
stsTor2' ::
     Int
  -> Double
  -> [Double]
  -> Int
  -> Double
  -> [Double]
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM2 Double)
stsTor2' numX xLen xFreqs numY yLen yFreqs arr = do
  let deltaX = xLen / fromIntegral numX
      deltaY = yLen / fromIntegral numY
      centerX = div numX 2
      centerY = div numY 2
      (Z :. numXF :. numYF :. numTF :. numSF) = extent arr
      mat = ((numXF * numYF) >< (numTF * numSF)) (R.toList arr)
  transformMat <-
    fmap
      (((numX * numY) >< (L.product . L.drop 2 . listOfShape . extent $ arr)) .
       R.toList) .
    computeUnboxedP .
    R.traverse2
      (fromListUnboxed (Z :. L.length xFreqs) xFreqs)
      (fromListUnboxed (Z :. L.length yFreqs) yFreqs)
      (\(Z :. numXFreq) (Z :. numYFreq) ->
         (Z :. numX :. numY :. numXFreq :. numYFreq)) $ \fx fy (Z :. x :. y :. xf :. yf) ->
      exp
        (0 :+
         (2 * pi *
          (fx (Z :. xf) * fromIntegral (x - centerX) * deltaX / xLen +
           fy (Z :. yf) * fromIntegral (y - centerY) * deltaY / yLen)))
  return .
    sumS .
    sumS .
    R.map (\x -> (magnitude x) ^ 2) .
    fromListUnboxed (Z :. numX :. numY :. numTF :. numSF) . NL.toList . flatten $
    transformMat NL.<> mat