{-# LANGUAGE TemplateHaskell, OverloadedLists #-}
{-# LANGUAGE Strict #-}

module Data.Scene where

import           Numeric.LinearAlgebra (Vector, scale)
import qualified Graphics.GL as GLRaw
import qualified SDL
import qualified Graphics.Rendering.OpenGL.GL as GL
import           Control.Lens
import           Control.Arrow
import qualified Data.Vector as V

import           Data.GameObject ( GameObject(..)
                                 , drawObject
                                 , location
                                 , drawObject2)
import           Utils (lookAt, step, norm, remTrans)
import           Data.SkyBox
import           Tasks.Swarm
import           Debug.Trace
-- import           Data.Traversable (traverse)

data Objects = Objects {
  _player :: GameObject,
  _npc    :: V.Vector GameObject,
  _area   :: V.Vector GameObject }

data Scene = Scene {
  _campos     :: Vector GLRaw.GLfloat,
  _objects    :: Objects,
  _projection :: Vector GLRaw.GLfloat,
  _skyBox :: Maybe SkyBox }

makeLenses ''Objects

makeLenses ''Scene

resolveObjectCollisions obs =
  let v = resolveContacts .
          V.cons (obs^.player) $
          (obs^.npc) -- V.++ (obs^.area)
      newPlayer = V.head v
      newNpcs =
        -- V.take (V.length (obs^.npc)) .
        V.drop 1 $ v
  in set player newPlayer . set npc newNpcs $ obs

traverseObjects :: (GameObject -> GameObject) -> Scene -> Scene
traverseObjects =
  over (objects.player) &&& over (objects.npc.traverse) >>> uncurry (.)

class Drawable d where
  draw :: SDL.Window -> d -> IO ()

instance Drawable Scene where
  draw window (Scene pos obj proj msb) =
    let camMat     = lookAt pos $ obj^.player.location
        objectDraw = drawObject (msb^?_Just.skyTex) camMat proj
        objectDraw2 = drawObject2 camMat proj
        drawSB      = drawSkyBox (remTrans camMat) proj
    in
      GL.clear [GL.ColorBuffer, GL.DepthBuffer]
      >> traverse drawSB msb
      >> objectDraw (obj^.player)
      >> V.mapM_ objectDraw (obj^.npc)
      >> V.mapM_ objectDraw (obj^.area)
      >> SDL.glSwapWindow window

-- CAMERA

moveCameraV = over campos . (+)

moveCamera :: Float -> Float -> Float -> Scene -> Scene
moveCamera x y z = moveCameraV [x,y,z]

moveDepth :: Float -> Scene -> Scene
moveDepth = moveCamera 0 0

moveForwardCamera :: Scene -> Scene
moveForwardCamera = moveDepth step

moveBackCamera :: Scene -> Scene
moveBackCamera = moveDepth (-step)

adjustCamera :: Scene -> Scene
adjustCamera scene =
  let playPos = scene^.objects.player.location
      pointer = playPos - scene^.campos
      dist    = norm pointer
      maxDist = 10
      tooFar  = dist > maxDist
      movement= scale ((dist-maxDist)/dist) pointer
  in if tooFar then over campos (+movement) scene
     else scene

moveHor :: Float -> Scene -> Scene
moveHor s = moveCamera s 0 0

moveLeftCamera :: Scene -> Scene
moveLeftCamera = moveHor step

moveRightCamera :: Scene -> Scene
moveRightCamera = moveHor (-step)


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
module Taiji.Utils.Clustering
    ( evalClusters
    , computeClusterMetrics
    , computeReproducibility
    , readKNNGraph
    , optimalParam
    , visualizeCluster
    , sampleCells
    , silhouette
    , confusionTable
    , batchCorrect
    , leiden
    , ari
    ) where

import qualified Data.ByteString.Char8 as B
import Data.Binary (encodeFile, decodeFile)
import Bio.Utils.Misc
import qualified Data.Text as T
import qualified Data.HashSet as S
import qualified Data.HashMap.Strict as M
import Data.Hashable (Hashable)
import Control.Arrow (first, second)
import Conduit
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Unboxed as U
import Data.List.Ordered (nubSort)
import Shelly (shelly, run_)
import Data.Conduit.Zlib (multiple, ungzip)
import AI.Clustering.Hierarchical.Types (computeDists, (!))
import AI.Clustering.Hierarchical (euclidean)
import Statistics.Sample (mean)
import System.Random.MWC.Distributions
import Language.Javascript.JMacro
import System.Random.MWC

import IGraph.Random
import IGraph
import qualified IGraph.Algorithms.Community as I

import Taiji.Prelude
import Taiji.Utils.Plot.ECharts
import Taiji.Utils.Plot

-- | Evaluating the clustering results
evalClusters :: FilePath
             -> Optimizer  -- ^ Optimizer
             -> Double     -- ^ Resolution
             -> FilePath   -- ^ knn
             -> IO ([FilePath], FilePath)
evalClusters dir optimizer res knn = do
    gr <- readKNNGraph knn
    gen <- create

    perturbed <- forM [1::Int ..5] $ \i -> do
        let output' = dir <> "/perturbed_clustering_result_" <> show i <> ".bin"
        r <- mutateGraph gr gen >>= leiden res optimizer
        encodeFile output' $ map snd $ sort $ concat $ zipWith (\a b -> zip a $ repeat b) r [0::Int ..]
        return output'

    let output = dir <> "/clustering_result.bin"
    (filter ((>100) . length) <$> leiden res optimizer gr) >>= encodeFile output
    return (perturbed, output)
  where
    mutateGraph gr gen = do
        let n = truncate $ fromIntegral (nEdges gr) * (0.02 :: Double)
        xs <- take n . V.toList <$> uniformShuffle (V.fromList $ edges gr) gen
        return $ delEdges xs gr

computeReproducibility :: [[Int]] -- ^ cluster
                       -> [[Int]]  -- ^ perturbed clusters
                       -> [Double]
computeReproducibility a bs = map (mean . U.fromList) $ transpose $ flip map bs $ \b ->
    let b' = U.fromList b
    in map (reproducibility . map (b' U.!)) a

computeClusterMetrics :: ([FilePath], FilePath)
                      -> FilePath
                      -> IO (Int, Double, Double)
computeClusterMetrics (perturbed, cl) coordinate = do
    stability <- (mean . U.fromList . map ari . comb) <$> mapM decodeFile perturbed
    clusters <- decodeFile cl
    points <- fmap V.fromList $ runResourceT $ runConduit $ sourceFile coordinate .|
        multiple ungzip .| linesUnboundedAsciiC .|
        mapC (U.fromList . map readDouble . B.split '\t') .| sinkList
    gen <- create

    let samplingRatio = 10000 / fromIntegral (V.length points) :: Double
    clusters' <- forM clusters $ \c -> do
        let n = max 2 $ truncate $ fromIntegral (length c) * samplingRatio
        V.toList . V.take n <$> uniformShuffle (V.fromList c) gen

    let sil = silhouette $ (map . map) (points V.!) clusters'
    return (length clusters, sil, stability)
  where
    comb (x:xs) = zip (repeat x) xs ++ comb xs
    comb _ = []

readKNNGraph :: FilePath -> IO (Graph 'U () Double)
readKNNGraph fl = runResourceT $ runConduit $ sourceFile fl .|
    multiple ungzip .| linesUnboundedAsciiC .| sink
  where
    sink = do
        n <- headC >>= \case
            Nothing -> error ""
            Just x -> return $ readInt x
        es <- mapC f .| sinkList
        return $ mkGraph (replicate n ()) $ replaceInf es
    f x = let (a:b:c:_) = B.split '\t' x
          in ((readInt a, readInt b), readDouble' c)
    readDouble' x | x == "inf" = Nothing
                  | otherwise = Just $ readDouble x
    replaceInf es = let m = maximum $ mapMaybe snd es
                    in map (second (fromMaybe m)) es

optimalParam :: FilePath
             -> [(Double, (Int, Double, Double))]
             -> IO Double
optimalParam output input = do
    savePlots output [] plt
    return optimal
  where
    optimal = fst $ maximumBy (comparing (^._2._2)) $ case filter (\x -> x^._2._3 >= 0.9) input of
        [] -> input
        x -> x
    (res, dat) = unzip $ flip map (sortBy (comparing fst) input) $ \(r, (n, sil, stab)) ->
        (r, (fromIntegral n, sil, stab))
    (num, sils, stabs) = unzip3 dat
    plt = map (setDim 400 300 . addAttr toolbox)
        [ addAttr (yAxisLabel "number of clusters") $
                addAttr (xAxisLabel "resolution") $
                lineplot $ zip res num
        , addAttr (yAxisLabel "silhouette width") $
                addAttr (xAxisLabel "resolution") $
                lineplot $ zip res sils
        , addAttr (yAxisLabel "stability") $
                addAttr (xAxisLabel "resolution") $
                lineplot $ zip res stabs ]
    lineplot input = mkEChart [jmacroE| {
        grid: { containLabel: true },
        xAxis: {
            type: "value",
            axisLine: {onZero: false}
        },
        tooltip: {
            trigger: "axis"
        },
        yAxis: {
            type: "value",
            axisLine: {onZero: false}
        },
        series: [{
            data: `dataset`,
            markLine: {symbolSize: 0, precision: 10, data: [{xAxis: `optimal`}]},
            type: "line"
        }]
        } |]
      where
        dataset = map (\(x,y) -> [x,y]) input

visualizeCluster :: [CellCluster] -> [EChart]
visualizeCluster cls =
    [ addAttr toolbox $ scatter' $ flip map cls $ \(CellCluster nm cells _) ->
        (B.unpack nm, map _cell_2d cells) ] ++ if noName then [] else
          [ addAttr toolbox $ scatter' $ map (first head . unzip) $
              groupBy ((==) `on` fst) $ sortBy (comparing fst) $ concatMap
              (map (\x -> (getName $ _cell_barcode x, _cell_2d x)) . _cluster_member) cls
          ]
  where
    noName = null $ getName $ _cell_barcode $ head $
        _cluster_member $ head cls
    getName x = let prefix = fst $ B.breakEnd (=='+') x
                in if B.null prefix then "" else B.unpack $ B.init prefix

-- | Random sample 30,000 cells.
sampleCells :: [CellCluster] -> IO [CellCluster]
sampleCells clusters
    | ratio >= 1 = return clusters
    | otherwise = do
        gen <- create
        forM clusters $ \c -> do
            s <- sampling gen ratio $ V.fromList $ _cluster_member c
            return $ c {_cluster_member = V.toList s}
  where
    n = foldl1' (+) $ map (length . _cluster_member) clusters
    ratio = 1 / (fromIntegral n / 30000) :: Double
    sampling gen frac v = V.take n' <$> uniformShuffle v gen
      where
        n' = max 200 $ truncate $ frac * fromIntegral (V.length v)

silhouette :: [[U.Vector Double]]   -- ^ Groups of objects
           -> Double
silhouette input
    | length input <= 1 = 0
    | otherwise = mean $ U.fromList $ flip concatMap clusters $ \(i, cl) ->
        let cls = map snd $ filter ((/= i) . fst) clusters
        in flip map cl $ \x ->
            let a = getD x cl
                b = minimum $ map (getD x) cls
            in (b - a) / max a b
  where
    getD i cl = mean $ U.fromList $ map (\j -> distMat ! (i, j)) $ filter (/= i) cl
    clusters = zip [0 :: Int ..] $ go 0 input
      where
        go i (x:xs) = let n = length x in take n [i..] : go (i + n) xs
        go _ _ = []
    distMat = computeDists euclidean $ V.fromList $ concat input
{-# INLINE silhouette #-}

confusionTable :: (Eq a, Hashable a) => [[a]] -> [[a]] -> [[Int]]
confusionTable inputA inputB = flip map inputA' $ \a -> flip map inputB' $ \b ->
    S.size $ S.intersection a b
  where
    inputA' = map S.fromList inputA
    inputB' = map S.fromList inputB
{-# INLINE confusionTable #-}

batchCorrect :: [Maybe (B.ByteString, Maybe B.ByteString)]  -- ^ label and group
             -> V.Vector (U.Vector Double)    -- ^ Data
             -> IO (V.Vector (U.Vector Double))
batchCorrect batchInfo input = do
    vec <- V.thaw input
    forM_ batchGroups $ \(idx, label) -> when (length (nubSort label) > 1) $ do
        res <- batchAveraging label (map (input V.!) idx)
        forM_ (zip idx res) $ \(i, r) -> VM.unsafeWrite vec i r
    V.unsafeFreeze vec
  where
    batchGroups = map (unzip . map (\(i, (x, _)) -> (i, x))) $
        groupBy ((==) `on` (snd . snd)) $ sortBy (comparing (snd . snd)) $
        mapMaybe (\(i, b) -> maybe Nothing (\x -> Just (i,x)) b) $ zip [0..] batchInfo
{-# INLINE batchCorrect #-}

batchAveraging :: [B.ByteString]    -- ^ Batch labels
               -> [U.Vector Double]    -- ^ Data
               -> IO [U.Vector Double]
batchAveraging labels input = withTempDir Nothing $ \dir -> do
    let labelFl = dir <> "/label.txt"
        matFl = dir <> "/mat.txt"
        tmp = dir <> "/tmp.txt"
    B.writeFile labelFl $ B.unlines labels
    B.writeFile matFl $ B.unlines $ map (B.intercalate "\t" . map toShortest . U.toList) input
    shelly $ run_ "taiji-utils" ["correct", "--label", T.pack labelFl, T.pack matFl, T.pack tmp]
    map (U.fromList . map readDouble . B.words) . B.lines <$> B.readFile tmp
{-# INLINE batchAveraging #-}

leiden :: Double -> Optimizer -> Graph 'U () Double -> IO [[Int]]
leiden resolution optimizer gr = withSeed 9304 $ fmap (sortBy (flip (comparing length))) .
    I.findCommunity gr nodeWeight edgeWeight I.leiden{I._resolution=res}
  where
    edgeWeight
        | optimizer `elem` [RBConfiguration, CPM] = Nothing
        | otherwise = Just id
    (nodeWeight, res)
        | optimizer `elem` [RBConfiguration, RBConfigurationWeighted] =
            ( Just $ \i _ -> foldl1' (+) $ map (\j -> edgeLab gr (i, j)) $ neighbors gr i
            , resolution / (2 * (foldl1' (+) $ map snd $ labEdges gr)))
        | otherwise = (Nothing, resolution)

-- |
reproducibility :: [Int]   -- membership
                -> Double
reproducibility xs = fromIntegral sameGroup / fromIntegral (c2 $ length xs)
  where
    counts = M.elems $ M.fromListWith (+) $ zip xs $ repeat (1 :: Int)
    sameGroup = foldl1' (+) $ map c2 counts
    c2 x = (x * (x - 1)) `div` 2

-- | Adjusted Rand Index: <http://en.wikipedia.org/wiki/Rand_index>
ari :: ([Int], [Int]) -> Double
ari (a, b) | length a /= length b = error "unequal length"
           | a == b = 1
           | otherwise = ari' $ counts a b
  where
    ari' (Counts cxy cx cy) =  (sum1 - sum2*sum3/choicen2)
                            / (1/2 * (sum2+sum3) - (sum2*sum3) / choicen2)
      where choicen2 = choice (sum . M.elems $ cx) 2
            sum1 = sum [ choice nij 2 | nij <- M.elems cxy ]
            sum2 = sum [ choice ni 2 | ni <- M.elems cx ]
            sum3 = sum [ choice nj 2 | nj <- M.elems cy ]
    -- | Creates count table 'Counts'
    counts xs = foldl' f mempty . zipWith ((,)) xs
        where f cs@(Counts cxy cx cy) p@(x,y) =
                cs { joint       = M.insertWith (+) p 1 cxy
                , marginalFst = M.insertWith (+) x 1 cx
                , marginalSnd = M.insertWith (+) y 1 cy }
    -- | The binomial coefficient: C^n_k = PROD^k_i=1 (n-k-i)\/i
    choice :: Double -> Double -> Double
    choice n k = foldl' (*) 1 [n-k+1 .. n] / foldl' (*) 1 [1 .. k]

-- | Count table
data Counts = Counts
    { joint :: !(M.HashMap (Int, Int) Double) -- ^ Counts of both components
    , marginalFst :: !(M.HashMap Int Double) -- ^ Counts of the first component
    , marginalSnd :: !(M.HashMap Int Double) -- ^ Counts of the second component
    }

instance Monoid Counts where
    mempty = Counts M.empty M.empty M.empty
    c `mappend` k =
        Counts { joint = unionPlus (joint c) (joint k)
               , marginalFst = unionPlus (marginalFst c) (marginalFst k)
               , marginalSnd = unionPlus (marginalSnd c) (marginalSnd k)
               }
      where
        unionPlus m = M.foldlWithKey' (\z k' v -> M.insertWith (+) k' v z) m

instance Semigroup Counts where
    c <> k =
        Counts { joint = unionPlus (joint c) (joint k)
               , marginalFst = unionPlus (marginalFst c) (marginalFst k)
               , marginalSnd = unionPlus (marginalSnd c) (marginalSnd k)
               }
      where
        unionPlus m = M.foldlWithKey' (\z k' v -> M.insertWith (+) k' v z) m


module Statistics.BBVI.Scheduler
  ( stepTogether
  , stepSeparate
  , unsafeContent
  )
where

import           Control.Monad.ST               ( ST )
import           Data.Propagator
import           Data.Maybe                     ( fromMaybe )
import qualified Data.Vector                   as V
import           Statistics.BBVI.Class
import           Statistics.BBVI.Propagator     ( Time
                                                , time
                                                , DistCell(..)
                                                , DistCellss
                                                , mergeGeneric
                                                , mergeGenericss
                                                )

-- | helper (partial) function to pull content from cell
unsafeContent :: Cell s a -> ST s a
unsafeContent =
  (fromMaybe (error "called unsafe content but no content in cell") <$>)
    . content

-- | update all propagators in one step until first argument reaches
-- quiesence; synonymous with traditional synchronous algorithms
stepTogether :: (a -> b -> ST s (a, b)) -> Cell s a -> Cell s b -> ST s ()
stepTogether f x y = watch x $ \x' -> with y $ \y' -> do
  (upX, upY) <- f x' y'
  -- careful: order below can create a bug. safer would be to use
  -- unsafe content instead of with y
  write y upY
  write x upX

-- | partially optimize each cell with the other held fixed, and
-- repeat until quiesence. Each local optimization will proceed some
-- number of steps, or until the local optimization converges, before
-- triggering the global cell to fire.
stepSeparate
  :: (DistUtil a1, DistUtil a2)
  => Time -- ^ number of local steps
  -> Double
  -> (  DistCell a2
     -> DistCellss a1
     -> ST s (DistCell a2, DistCellss a1)
     ) -- ^ gradient propagator
  -> Cell s (DistCell a2)
  -> Cell s (V.Vector (V.Vector (DistCell a1))) --
  -> ST s ()
stepSeparate nLocal localDelta f x0 ys0 = do
  watch x0 $ \x -> do
    yTemp <- unsafeContent ys0
    ys    <- cellWith $ mergeGenericss
      ((+ nLocal) . time . (V.! 0) . (V.! 0) $ yTemp)
      localDelta
    write ys yTemp
    watch ys $ \ys' -> do
      (_upX, upYs) <- f x ys'
      -- upYs <- f2 x ys'
      write ys upYs
    ysNew <- unsafeContent ys
    write ys0 ysNew
  watch ys0 $ \ys -> do
    xTemp <- unsafeContent x0
    x     <- cellWith $ mergeGeneric (nLocal + time xTemp) localDelta
    write x xTemp
    watch x $ \x' -> do
      (upX, _upYs) <- f x' ys
      -- upX <- f1 x' ys
      write x upX
    xNew <- unsafeContent x
    write x0 xNew


{-# LANGUAGE TypeFamilies, GeneralizedNewtypeDeriving, DeriveGeneric, TupleSections, RecordWildCards, TemplateHaskell, RankNTypes, FlexibleContexts #-}

module BayesStack.Models.Topic.CitationInfluenceNoTopics
  ( -- * Primitives
    NetData
  , dHypers, dArcs, dItems, dNodeItems, dCitingNodes, dCitedNodes
  , netData
  , HyperParams(..)
  , MState
  , stGammas, stOmegas, stPsis, stCiting, stLambdas
  , CitingUpdateUnit
  , ItemSource(..)
  , CitedNode(..), CitedNodeItem(..)
  , CitingNode(..), CitingNodeItem(..)
  , Citing(..), Cited(..)
  , Item(..), Topic(..), NodeItem(..), Node(..), Arc(..)
  , setupNodeItems
    -- * Initialization
  , verifyNetData, cleanNetData
  , ModelInit
  , randomInitialize
  , model
  , updateUnits
    -- * Diagnostics
  , modelLikelihood
  ) where

import qualified Data.Vector as V
import Statistics.Sample (mean)

import           Prelude hiding (mapM_, sum)
import           Data.Maybe (fromMaybe)

import           Control.Lens hiding (Setting)
import           Data.Set (Set)
import qualified Data.Set as S

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

import           Data.Foldable hiding (product)
import           Control.Applicative ((<$>), (<*>))
import           Control.Monad (when)
import           Control.Monad.Trans.State.Strict
import           Control.Monad.Trans.Writer.Strict

import           Data.Random
import           Data.Random.Lift (lift)
import           Data.Random.Distribution.Categorical (categorical)
import           Numeric.Log hiding (sum)

import           BayesStack.Types
import           BayesStack.Gibbs
import           BayesStack.DirMulti
import           BayesStack.TupleEnum ()
import           BayesStack.Models.Topic.Types

import           GHC.Generics (Generic)
import           Data.Binary (Binary)
import           Control.DeepSeq

at' :: At m => Index m -> IndexedLens' (Index m) m (IxValue m)
at' i = at i . _fromMaybe
  where _fromMaybe = iso (fromMaybe $ error "at': Unexpected Nothing") Just

data ItemSource = Shared | Own deriving (Show, Eq, Enum, Ord, Generic)
instance Binary ItemSource
instance NFData ItemSource

newtype Citing a = Citing a deriving (Show, Eq, Enum, Ord, Generic, NFData)
newtype Cited a = Cited a deriving (Show, Eq, Enum, Ord, Generic, NFData)
instance Binary a => Binary (Citing a)
instance Binary a => Binary (Cited a)

type CitingNode = Citing Node
type CitedNode = Cited Node
type CitingNodeItem = Citing NodeItem
type CitedNodeItem = Cited NodeItem

-- ^ A directed edge
data Arc = Arc { citingNode :: !CitingNode, citedNode :: !CitedNode }
            deriving (Show, Eq, Ord, Generic)
instance Binary Arc

data HyperParams = HyperParams
                   { _alphaPsi         :: Double
                   , _alphaLambda      :: Double
                   , _alphaOmega       :: Double
                   , _alphaGammaShared :: Double
                   , _alphaGammaOwn    :: Double
                   , _alphaBetaFG      :: Double
                   , _alphaBetaBG      :: Double
                   }
                 deriving (Show, Eq, Generic)
instance Binary HyperParams
makeLenses ''HyperParams

data NetData = NetData { _dHypers             :: !(HyperParams)
                       , _dArcs               :: !(Set Arc)
                       , _dItems              :: !(Map Item Double)
                       , _dNodeItems          :: !(Map NodeItem (Node, Item))
                       , _dCitingNodes        :: !(Map CitingNode (Set CitedNode))
                         -- ^ Maps each citing node to the set of nodes cited by it
                       , _dCitedNodes         :: !(Map CitedNode (Set CitingNode))
                         -- ^ Maps each cited node to the set of nodes citing it
                       }
              deriving (Show, Eq, Generic)
instance Binary NetData
makeLenses ''NetData

netData :: HyperParams -> Set Arc -> Map Item Double -> Map NodeItem (Node,Item) -> NetData
netData hypers arcs items nodeItems =
    NetData { _dHypers       = hypers
            , _dArcs         = arcs
            , _dItems        = items
            , _dNodeItems    = nodeItems
            , _dCitingNodes  = M.unionsWith S.union
                               $ map (\(Arc a b)->M.singleton a $ S.singleton b)
                               $ S.toList arcs
            , _dCitedNodes   = M.unionsWith S.union
                               $ map (\(Arc a b)->M.singleton b $ S.singleton a)
                               $ S.toList arcs
            }

dCitingNodeItems :: NetData -> Map CitingNodeItem (CitingNode, Item)
dCitingNodeItems nd =
    M.mapKeys Citing
    $ M.map (\(n,i)->(Citing n, i))
    $ M.filter (\(n,i)->Citing n `M.member` (nd^.dCitingNodes))
    $ nd^.dNodeItems

itemsOfCitingNode :: NetData -> CitingNode -> [Item]
itemsOfCitingNode d (Citing u) =
    map snd $ M.elems $ M.filter (\(n,_)->n==u) $ d^.dNodeItems

connectedNodes :: Set Arc -> Set Node
connectedNodes arcs =
    S.map ((\(Cited n)->n) . citedNode) arcs `S.union` S.map ((\(Citing n)->n) . citingNode) arcs

cleanNetData :: NetData -> NetData
cleanNetData d =
    let nodesWithItems = S.fromList $ map fst $ M.elems $ d^.dNodeItems
        nodesWithArcs = connectedNodes $ d^.dArcs
        keptNodes = nodesWithItems `S.intersection` nodesWithArcs
        keepArc (Arc (Citing citing) (Cited cited)) =
            citing `S.member` keptNodes && cited `S.member` keptNodes
        go = do dArcs %= S.filter keepArc
                dNodeItems %= M.filter (\(n,i)->n `S.member` keptNodes)
    in execState go d

verifyNetData :: (Node -> String) -> NetData -> [String]
verifyNetData showNode d = execWriter $ do
    let nodesWithItems = S.fromList $ map fst $ M.elems $ d^.dNodeItems
    forM_ (d^.dArcs) $ \(Arc (Citing citing) (Cited cited))->do
        when (cited `S.notMember` nodesWithItems)
            $ tell [showNode cited++" has arc yet has no items"]
        when (citing `S.notMember` nodesWithItems)
            $ tell [showNode citing++" has arc yet has no items"]

-- Citing Update unit (Shared Taste-like)
data CitingUpdateUnit = CitingUpdateUnit { _uuNI    :: CitingNodeItem
                                         , _uuN     :: CitingNode
                                         , _uuX     :: Item
                                         , _uuCites :: Set CitedNode
                                         , _uuItemWeight :: Double
                                         }
                      deriving (Show, Generic)
instance Binary CitingUpdateUnit
makeLenses ''CitingUpdateUnit

citingUpdateUnits :: NetData -> [CitingUpdateUnit]
citingUpdateUnits d =
    map (\(ni,(n,x))->CitingUpdateUnit { _uuNI      = ni
                                       , _uuN       = n
                                       , _uuX       = x
                                       , _uuCites   = d^.dCitingNodes . at' n
                                       , _uuItemWeight = (d ^. dItems . at' x)
                                       }
        ) $ M.assocs $ dCitingNodeItems d

updateUnits :: NetData -> [WrappedUpdateUnit MState]
updateUnits d = map WrappedUU (citingUpdateUnits d)

-- | Model State
data CitingSetting = OwnSetting
                   | SharedSetting !CitedNode
                   deriving (Show, Eq, Generic)
instance Binary CitingSetting
instance NFData CitingSetting where
    rnf (OwnSetting)      = ()
    rnf (SharedSetting c) = rnf c `seq` ()

data MState = MState { -- Citing model state
                       _stGammas   :: !(Map CitingNode (Multinom Int ItemSource))
                     , _stOmegas   :: !(Map CitingNode (Multinom Int Item))
                     , _stPsis     :: !(Map CitingNode (Multinom Int CitedNode))

                     , _stCiting   :: !(Map CitingNodeItem CitingSetting)

                     -- Cited model state
                     , _stLambdas  :: !(Map CitedNode (Multinom Int Item))
                     }
            deriving (Show, Generic)
instance Binary MState
makeLenses ''MState

-- | Model initialization
type ModelInit = Map CitingNodeItem (Setting CitingUpdateUnit)

modify' :: Monad m => (a -> a) -> StateT a m ()
modify' f = do x <- get
               put $! f x

randomInitializeCiting :: NetData -> ModelInit -> RVar ModelInit
randomInitializeCiting d init = execStateT doInit init
    where doInit :: StateT ModelInit RVar ()
          doInit = let unset = M.keysSet (dCitingNodeItems d) `S.difference` M.keysSet init
                   in mapM_ (randomInitCitingUU d) (S.toList unset)

randomInitCitingUU :: NetData -> CitingNodeItem -> StateT ModelInit RVar ()
randomInitCitingUU d cni@(Citing ni) =
    let (n,_) = d ^. dNodeItems . at' ni
    in case d ^. dCitingNodes . at' (Citing n) of
           a | S.null a -> do
               modify' $ M.insert cni OwnSetting

           citedNodes -> do
               s <- lift $ randomElement [Shared, Own]
               c <- lift $ randomElement $ toList citedNodes
               modify' $ M.insert cni $
                   case s of Shared -> SharedSetting c
                             Own    -> OwnSetting

randomInitialize :: NetData -> RVar ModelInit
randomInitialize d = randomInitializeCiting d M.empty

model :: NetData -> ModelInit -> MState
model d citingInit =
    let citingNodes = M.keys $ d^.dCitingNodes
        hp = d^.dHypers
        s = MState { -- Citing model
                     _stPsis = let dist n = case d ^. dCitingNodes . at' n . to toList of
                                                []    -> M.empty
                                                nodes -> M.singleton n
                                                         $ symDirMulti (hp^.alphaPsi) nodes
                               in foldMap dist citingNodes
                   , _stGammas = let dist = multinom [ (Shared, hp^.alphaGammaShared)
                                                     , (Own, hp^.alphaGammaOwn) ]
                                 in foldMap (\t->M.singleton t dist) citingNodes
                   , _stOmegas = let dist = symDirMulti (hp^.alphaOmega) (M.keys $ d^.dItems)
                                 in foldMap (\t->M.singleton t dist) citingNodes
                   , _stCiting = M.empty

                   -- Cited model
                   , _stLambdas = let dist = symDirMulti (hp^.alphaLambda) (M.keys $ d^.dItems)
                                      lambdas0 = foldMap (\n->M.singleton n dist) $ M.keys $ d^.dCitedNodes
                                  in foldl' (\dms (n,x)->M.adjust (incMultinom x) (Cited n) dms) lambdas0 (M.elems $ d^.dNodeItems)
                   }

        initCitingUU :: CitingUpdateUnit -> State MState ()
        initCitingUU uu = do
            let err = error $ "CitationInference: Initial value for "++show uu++" not given\n"
                s = maybe err id $ M.lookup (uu^.uuNI) citingInit
            modify' $ setCitingUU uu (Just s)

    in execState (mapM_ initCitingUU $ citingUpdateUnits d) s

modelLikelihood :: MState -> Probability
modelLikelihood model =
    product $ (model ^.. stGammas  . folded . to likelihood)
           ++ (model ^.. stLambdas . folded . to likelihood)
           ++ (model ^.. stOmegas  . folded . to likelihood)
           ++ (model ^.. stPsis    . folded . to likelihood)

instance UpdateUnit CitingUpdateUnit where
    type ModelState CitingUpdateUnit = MState
    type Setting CitingUpdateUnit = CitingSetting
    fetchSetting uu ms = ms ^. stCiting . at' (uu^.uuNI)
    evolveSetting ms uu = categorical $ citingFullCond (setCitingUU uu Nothing ms) uu
    updateSetting uu _ s' = setCitingUU uu (Just s') . setCitingUU uu Nothing

citingProb :: MState -> CitingUpdateUnit -> Setting CitingUpdateUnit -> Double
citingProb st (CitingUpdateUnit {_uuN=n, _uuX=x}) setting =
    let gamma = st ^. stGammas . at' n
        omega = st ^. stOmegas . at' n
        psi = st ^. stPsis . at' n
    in case setting of
        SharedSetting c   -> let lambda = st ^. stLambdas . at' c
                             in sampleProb gamma Shared
                              * sampleProb psi c
                              * sampleProb lambda x
        OwnSetting        ->  sampleProb gamma Own
                            * sampleProb omega x

citingFullCond :: MState -> CitingUpdateUnit -> [(Double, Setting CitingUpdateUnit)]
citingFullCond ms uu = map (\s->(citingProb ms uu s, s)) $ citingDomain ms uu

citingDomain :: MState -> CitingUpdateUnit -> [Setting CitingUpdateUnit]
citingDomain ms uu = do
    s <- [Own, Shared]
    case s of
        Shared -> do c <- uu ^. uuCites . to S.toList
                     return $ SharedSetting c
        Own    -> do return $ OwnSetting

setCitingUU :: CitingUpdateUnit -> Maybe (Setting CitingUpdateUnit) -> MState -> MState
setCitingUU uu@(CitingUpdateUnit {_uuNI=ni, _uuN=n, _uuX=x}) setting ms = execState go ms
  where
    set = maybe Unset (const Set) setting
    go = case maybe (fetchSetting uu ms) id setting of
           SharedSetting c    -> do stPsis .    at' n %= setMultinom set c
                                    stLambdas . at' c %= setMultinom set x
                                    stGammas .  at' n %= setMultinom set Shared
                                    stCiting .  at ni .= setting

           OwnSetting         -> do stOmegas .  at' n %= setMultinom set x
                                    stGammas .  at' n %= setMultinom set Own
                                    stCiting .  at ni .= setting


{- |
Description :  Default values and constants
Copyright   :  (c) Dominik Schrempf 2017
License     :  GPLv3

Maintainer  :  dominik.schrempf@gmail.com
Stability   :  unstable
Portability :  non-portable (not tested)

Default values and constants.

* Changelog

-}

module Defaults where

import           DNAModel              ( DNAModelSpec(..))
import           Numeric.LinearAlgebra
import qualified System.Random         as Rand

-- | Output file name.
outFileName :: String
outFileName = "Test"

-- | Virtual population size.
popSize :: Int
popSize = 9

-- | A kappa value of 6.25 corresponds to a transition to transversion ratio of
-- 3.0
dnaModelSpec :: DNAModelSpec
dnaModelSpec = HKY 6.25 (vector [0.3, 0.2, 0.2, 0.3])

-- | Heterozygosity value.
heterozygosity :: Double
heterozygosity = 0.0025

-- | Tree height.
treeHeight :: Double
treeHeight = 0.005

-- | Tree type. At the moment, ILS (incomplete lineage sorting) and Yule trees
-- are supported.
treeType :: String
treeType = "ILS"

-- | Number of sites to simulate.
nSites :: Int
nSites = 1000000

-- | By default, the seed is random.
seed :: IO Rand.StdGen
seed = Rand.getStdGen


-- |
-- Module      : Occlusion.Types
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created September 20 2015

-- TODO | -
--        -

-- SPEC | -
--        -

--------------------------------------------------------------------------------------------------------------------------------------------
-- GHC Pragmas
--------------------------------------------------------------------------------------------------------------------------------------------

--------------------------------------------------------------------------------------------------------------------------------------------
-- API
--------------------------------------------------------------------------------------------------------------------------------------------
module Occlusion.Types where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
import Data.Complex
import qualified Data.Set as S
import qualified Data.Map as M

import Graphics.UI.Gtk
import qualified Graphics.Rendering.Cairo as Cairo

import           Southpaw.Picasso.Palette (Colour)
import qualified Southpaw.Picasso.Palette as Palette

--------------------------------------------------------------------------------------------------------------------------------------------
-- Types
--------------------------------------------------------------------------------------------------------------------------------------------

-- |
data AppState = AppState { _gui :: GUI, _animation :: AnimationData, _input :: InputData, _scene :: Scene, _assets :: Assets }

-- |
data GUI = GUI { _window :: Window,  _canvas :: DrawingArea }

-- |
data InputData = InputData { _mouse :: Complex Double, _click :: Maybe (Complex Double), _keyboard :: S.Set String }

-- |
data AnimationData = AnimationData { _fps :: Double, _frame :: Int }

-- |
data Assets = Assets { _images :: M.Map String Cairo.Surface } --  _strings }

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

-- |
type Polygon n = [Complex n]
type Edge n    = [Complex n]
type Linear n  = (n, n) -- TODO: Use custom type (eg. Linear { slope :: n, intercept :: n })
data Line n    = Line (Complex n) (Complex n)
-- type Scene n   = [Polygon n]

-- |
data Scene = Scene { _player :: Character, _npcs :: [Character], _obstacles :: [Polygon Double] }

-- |
-- TODO: Factor out stats, behaviour, visuals, 'physics', etc.
-- TODO: Should it be Behaviour Character AppState AppState (would allow player to modify world, maybe some restricted subset of it)
data Character = Character { _position :: Complex Double, _velocity :: Complex Double, _health :: Int, _name :: String, _colour :: Colour Double, _behaviour :: Behaviour Character AppState Character }

-- |
-- TODO: This type should be much more flexible
-- TOOD: Should this be a newtype or type synonym instead (?)
type Behaviour self s s' = self -> s -> s'
-- newtype Behaviour self s s' = Behaviour { _respond :: self -> s -> s' }
-- data Behaviour self s s' = Behaviour { _respond :: self -> s -> s' }


-- |
-- Module      : main
-- Description : Graphics module leveraging the Gloss package
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental
-- Portability : POSIX (not sure)
--
-- Jonatan H Sundqvist
-- June 2 2015
--

-- TODO | -
--        -

-- SPEC | -
--        -

module GlossGraphics where

---------------------------------------------------------------------------------------------------
-- We'll need these
---------------------------------------------------------------------------------------------------
-- import Graphics.Gloss.Data.Picture (line)
import qualified Graphics.Gloss as Gloss --(circleSolid, rectangleSolid)
import Graphics.Gloss.Interface.IO.Game hiding (Vector)
-- import Graphics.Gloss.Geometry.Angle (degToRad, radToDeg, normaliseAngle)
-- import Graphics.Gloss.Data.Vector

import Data.Complex

import Copernicus

---------------------------------------------------------------------------------------------------
-- Types
---------------------------------------------------------------------------------------------------
data World = World { bodies :: [Body Float],
                     grid   :: Bool
                   }

---------------------------------------------------------------------------------------------------
-- Interaction
---------------------------------------------------------------------------------------------------
simulate :: IO ()
simulate = playIO
	display 	-- Window mode
	white 		-- Background colour
	60			-- FPS (simulation steps per second, technically)
	world 		-- Initial world
	render		-- Converts world to Picture
	respond 	-- User interaction
	advance 	-- Advances the world to the next simulation step
	where
		display = InWindow "Simulator" (width, height) (25, 25)
		world   = World { bodies = map (\ (p', v', g') -> Body p' v' g') [(0.0:+0.0, v, g), (10.0:+0.0, (-20.0):+15.0, g), (35.0:+(-28.0), v, g), (20.5:+19.2, v, g)],
	                      grid   = True
	              }
		-- TODO: Refactor this ugly mess
		render w 	 		  = return . pictures $ zipWith drawBall colours (bodies w) ++ [drawGround (bodies w)] ++ if grid w then [renderGrid 45 45 width height] else []
		drawBall col (Body (x:+y) _ _) = color col . translate x y $ circleSolid 15 												-- TODO: Reorder arguments (?)
		drawGround _ 		  = translate 0 (30/2-fromIntegral height/2) . color green $ Gloss.rectangleSolid (fromIntegral width) 30
		respond e w           = return $ case e of
			EventKey (Char 'g') Down _ _ -> w { grid = not $ grid w }
			_                            -> w

		advance t w 		  = return $ w { bodies=map (animate t) $ bodies w }
		v 					  = 40.0:+20.0 --40.0 :+ 20.0
		g 					  = 0.0:+(-98.2) --0.0 :+ (-9.82)
		(width, height) 	  = (740, 540)
		colours = cycle [red, green, orange, makeColor 0.2 0.1 0.3 1.0]
		--parabola t v a p = let (px:+py) = p; (vx:+vy) = v; (ax:+ay) = a in (px + vx*t + 0.5*ax*t**2):+(py + vy*t + 0.5*ay*t**2)

--
--renderForces

--
-- Subgrids, colours, thickness, markings, hover for coordinates, snap-to-grid
renderGrid :: Float -> Float -> Int -> Int -> Picture
renderGrid dx dy w' h' = pictures $ rows ++ cols
	where
		cols = map (\cl -> line [(cl*dx-w/2, -h/2), (cl*dx-w/2, h/2)]) [1..cls]
		rows = map (\rw -> line [(-w/2, rw*dy-h/2), (w/2, rw*dy-h/2)]) [1..rws]
		rws  = h / dy -- Number of rows
		cls  = w / dx -- Number of columns
		w 	 = fromIntegral w'
		h 	 = fromIntegral h'

-- Transforms a vector from one coordinate space to another
-- by applying the given scaling and translation
-- Useful for converting between simulation and screen coordinates
-- TODO: Make pure (eg. use for pure coordinates so it doesn't depend on Gloss) (?)
transform :: Vector Float -> Vector Float -> Picture -> Picture
transform (sx:+sy) (dx:+dy) = scale sx sy . translate dx dy


-- |
-- Copyright: (c) 2020 Tom Westerhout
-- SPDX-License-Identifier: BSD-3-Clause
-- Maintainer: Tom Westerhout <14264576+twesterhout@users.noreply.github.com>
--
-- High-level wrapper around [PRIMME C
-- library](https://github.com/primme/primme). /Quote from README/:
--
-- @
--   PRIMME, pronounced as prime, is a high-performance library for computing a few
--   eigenvalues/eigenvectors, and singular values/vectors. PRIMME is especially
--   optimized for large, difficult problems. Real symmetric and complex Hermitian
--   problems, standard @A x = λ x@ and generalized @A x = λ B x@, are supported.
--   Besides, standard eigenvalue problems with a normal matrix are supported. It
--   can find largest, smallest, or interior singular/eigenvalues, and can use
--   preconditioning to accelerate convergence.
-- @
module Numeric.PRIMME
  ( -- * Defining the matrix

    -- | One of the great things about PRIMME library is that it works with
    -- block matrix-vector products (i.e. matrix-matrix products). Following the
    -- example of [vector](https://hackage.haskell.org/package/vector) library,
    -- we define two types for mutable and immutable dense blocks.
    MBlock (..),
    Block (..),
    -- | Now we are can define the "operator" we want to diagonalize. Since this
    -- library is meant to be used with rather large matrices, it might very
    -- well be the case that the matrix does not fit into the memory of your
    -- computer (even in sparse format such as CSR). Sometimes, however, we can
    -- define the "operator" implicitly, namely, by defining its action on a
    -- vector (or in case of PRIMME on a block of vectors). This is done with
    -- 'PrimmeOperator' type.
    PrimmeOperator,

    -- * Choosing what to compute

    -- | Having defined a 'PrimmeOperator' we now need to tell PRIMME what to
    -- compute. This is done by constructing a 'PrimmeOptions' object.
    PrimmeOptions (..),
    primmeDefaults,
    finalizeOptions,
    PrimmeTarget (..),

    -- * Diagonalizing
    eigh,
    eigh',

    -- * Logging
    PrimmeMonitor (..),
    PrimmeInfo (..),
    PrimmeEventInfo (..),
    PrimmeStats (..),
    primmePrettyInfo,

    -- * Dense Matrices

    -- | The most trivial example of an "operator" is of course a square dense
    -- matrix. We thus provide a function which constructs an operator from a
    -- matrix.
    primmeFromDense,

    -- * Misc
    getPrimmeVersion,
    PrimmeException (..),
    BlasDatatype (blasTag),
    BlasRealPart,
    BlasDatatypeTag (..),
  )
where

import Control.Exception.Safe (throw)
import Control.Monad (forM_, unless)
import Control.Monad.ST (RealWorld)
import Data.Complex
import Data.Proxy
import Data.Vector.Storable (Vector)
import qualified Data.Vector.Storable as V
import Data.Vector.Storable.Mutable (MVector)
import qualified Data.Vector.Storable.Mutable as MV
import Foreign.C.Types (CInt (..))
import Foreign.Ptr (Ptr)
import Foreign.Storable
import qualified Language.C.Inline as C
import qualified Language.C.Inline.Unsafe as CU
import Numeric.PRIMME.Context
import Numeric.PRIMME.Dense
import Numeric.PRIMME.Monitor
import Numeric.PRIMME.Options
import Numeric.PRIMME.Types
import Prelude hiding (init)

C.context (C.baseCtx <> primmeCtx)
C.include "<primme.h>"

-- | Diagonalize the operator to find the first few eigenpairs.
eigh ::
  forall a.
  BlasDatatype a =>
  PrimmeOptions ->
  PrimmeOperator a ->
  IO (Vector (BlasRealPart a), Block a, Vector (BlasRealPart a))
eigh options matrix = do
  let dim = pDim options
      numEvals = pNumEvals options
  (evals :: MVector RealWorld (BlasRealPart a)) <- MV.new numEvals
  (evecs :: MVector RealWorld a) <- MV.new (dim * numEvals)
  (rnorms :: MVector RealWorld (BlasRealPart a)) <- MV.new numEvals
  status <-
    MV.unsafeWith evals $ \evalsPtr ->
      MV.unsafeWith evecs $ \evecsPtr ->
        MV.unsafeWith rnorms $ \rnormsPtr ->
          withOptions options matrix $ \optionsPtr ->
            cPrimme evalsPtr evecsPtr rnormsPtr optionsPtr
  unless (status == 0) $ throwPrimmeError status
  evals' <- V.unsafeFreeze evals
  evecs' <- Block (dim, numEvals) dim <$> V.unsafeFreeze evecs
  rnorms' <- V.unsafeFreeze rnorms
  return (evals', evecs', rnorms')

mkEmptyBlock :: Storable a => Int -> Int -> IO (MBlock RealWorld a)
mkEmptyBlock rows columns
  | rows <= 0 || columns <= 0 = error $ "invalid shape specified: " <> show (rows, columns)
  | otherwise = MBlock (rows, columns) rows <$> MV.new (rows * columns)

copyBlock :: Storable a => MBlock RealWorld a -> Block a -> IO ()
copyBlock dest@(MBlock (r, c) _ _) src@(Block (r', c') _ _)
  | r /= r' || c /= c' = error $ "block shape mismatch: " <> show (r, c) <> " != " <> show (r', c')
  | otherwise = do
    src' <- unsafeThaw src
    forM_ [0 .. (c - 1)] $ \i ->
      MV.copy (getColumn i dest) (getColumn i src')
  where
    getColumn :: Storable a => Int -> MBlock RealWorld a -> MVector RealWorld a
    getColumn i (MBlock (rows, _) stride v) = MV.slice (i * stride) rows v

unsafeThaw :: Storable a => Block a -> IO (MBlock RealWorld a)
unsafeThaw (Block shape stride v) = MBlock shape stride <$> V.unsafeThaw v

unsafeFreeze :: Storable a => MBlock RealWorld a -> IO (Block a)
unsafeFreeze (MBlock shape stride v) = Block shape stride <$> V.unsafeFreeze v

sliceBlock1 :: Storable a => Int -> Int -> MBlock s a -> MBlock s a
sliceBlock1 i n (MBlock (r, c) stride v)
  | i < 0 || n < 0 || i + n > c = error $ "invalid slice: [" <> show i <> ", " <> show (i + n) <> ")"
  | otherwise = MBlock (r, n) stride $ MV.slice (i * stride) (n * stride) v

eigh' ::
  forall a.
  BlasDatatype a =>
  PrimmeOptions ->
  MBlock RealWorld a ->
  PrimmeOperator a ->
  IO (Vector (BlasRealPart a), Block a, Vector (BlasRealPart a))
eigh' options init@(MBlock (dim', initSize) _ _) matrix
  | pDim options /= dim' =
    error $
      "'init' has wrong shape: "
        <> show (dim', initSize)
        <> "; expected a block with "
        <> show (pDim options)
        <> " rows"
  | otherwise = do
    let dim = pDim options
        numEvals = pNumEvals options
    evecs@(MBlock (_, _) evecsStride evecsData) <- case initSize < numEvals of
      True -> do
        temp <- mkEmptyBlock dim numEvals
        copyBlock (sliceBlock1 0 initSize temp) =<< unsafeFreeze init
        return temp
      False -> return init
    (evals :: MVector RealWorld (BlasRealPart a)) <- MV.new numEvals
    (rnorms :: MVector RealWorld (BlasRealPart a)) <- MV.new numEvals
    status <-
      MV.unsafeWith evals $ \evalsPtr ->
        MV.unsafeWith evecsData $ \evecsPtr ->
          MV.unsafeWith rnorms $ \rnormsPtr ->
            withOptions options matrix $ \optionsPtr -> do
              let c_initSize = fromIntegral initSize
                  c_evecsStride = fromIntegral evecsStride
              [C.block| void {
                $(primme_params* optionsPtr)->initSize = $(int c_initSize);
                $(primme_params* optionsPtr)->ldevecs = $(PRIMME_INT c_evecsStride);
              } |]
              cPrimme evalsPtr evecsPtr rnormsPtr optionsPtr
    unless (status == 0) $ throwPrimmeError status
    evals' <- V.unsafeFreeze evals
    evecs' <- unsafeFreeze $ sliceBlock1 0 numEvals evecs
    rnorms' <- V.unsafeFreeze rnorms
    return (evals', evecs', rnorms')

throwPrimmeError :: CInt -> IO a
throwPrimmeError c = case (- c) of
  0 -> error "no error"
  3 -> throw PrimmeMaximumIterationsReached
  40 -> throw PrimmeLapackFailure
  41 -> throw PrimmeUserFailure
  _ -> throw $ PrimmeOtherFailure c

foreign import ccall "sprimme"
  sprimme :: Ptr Float -> Ptr Float -> Ptr Float -> Ptr Cprimme_params -> IO CInt

foreign import ccall "dprimme"
  dprimme :: Ptr Double -> Ptr Double -> Ptr Double -> Ptr Cprimme_params -> IO CInt

foreign import ccall "cprimme"
  cprimme :: Ptr Float -> Ptr (Complex Float) -> Ptr Float -> Ptr Cprimme_params -> IO CInt

foreign import ccall "zprimme"
  zprimme :: Ptr Double -> Ptr (Complex Double) -> Ptr Double -> Ptr Cprimme_params -> IO CInt

cPrimme ::
  forall a.
  (BlasDatatype a) =>
  Ptr (BlasRealPart a) ->
  Ptr a ->
  Ptr (BlasRealPart a) ->
  Ptr Cprimme_params ->
  IO CInt
cPrimme = case blasTag (Proxy :: Proxy a) of
  FloatTag -> sprimme
  DoubleTag -> dprimme
  ComplexFloatTag -> cprimme
  ComplexDoubleTag -> zprimme

getPrimmeVersion :: (Int, Int)
getPrimmeVersion =
  ( fromIntegral [CU.pure| int { PRIMME_VERSION_MAJOR } |],
    fromIntegral [CU.pure| int { PRIMME_VERSION_MINOR } |]
  )


{-# LANGUAGE OverloadedStrings #-}

module Haskem.Lib
    ( doGaussianParsing,
      doFreqParsing,
      doGenericParsing,
      xyzParser,
      gaussianFreqParser,
      gaussianParser,
      elemMass,
      elemSym2Num,
      calculateCoM,
      calcDisplacementVector,
      writeMoleculeToXYZ,
      projectDonN,
      calcProjectionVec,
      calcLength3D,
      setUpBonds,
      setUpBonds',
      setUpBonds'',
      setUpBonds''',
      checkBondCompleteness
    ) where

import              Haskem.Parsers
import              Haskem.Types
import              Control.Applicative
import qualified    Formatting                  as F
import              Data.Attoparsec.Text.Lazy   hiding (take)
import qualified    Data.Text.Lazy              as T
import qualified    Data.Text.IO                as TIO
import              Data.Char
import              Lens.Micro.Platform
import              Debug.Trace
import              Data.List                   hiding (takeWhile)
import              Data.List.Extra
-- import              Prelude                     hiding (takeWhile)
-- import              Data.Tuple.HT               hiding (double, triple)
import qualified    Data.Map.Strict             as Map
import qualified    Data.Vector.Storable        as VS
import qualified    Numeric.LinearAlgebra.Data  as LA
import qualified    Data.IntMap                 as IM
import qualified    Data.IntSet                 as IS

xyzFormat :: F.Format r (String -> Double -> Double -> Double -> r)
xyzFormat = F.string F.% "\t" F.% F.fixed 4 F.% "\t" F.% F.fixed 4 F.% "\t" F.% F.fixed 4

doGaussianParsing :: FilePath -> IO GaussianInfo
doGaussianParsing filePath = do
    fileContent <- TIO.readFile filePath
    let gaussParse = parseOnly gaussianParser fileContent

    case gaussParse of
        Left    err     -> error err
        Right   gInfo   -> return gInfo

-- Parses a Gaussian freq calculation output to a GaussianInfo instance and a
-- list of Normalmode instances. Returns the output as a tuple in IO.
doFreqParsing :: FilePath -> IO (GaussianInfo, [NormalMode])
doFreqParsing filePath = do
    fileContent <- TIO.readFile filePath

    let gaussParse = parseOnly gaussianParser       fileContent
    let gaussFreq  = parseOnly gaussianFreqParser   fileContent

    gFreq <- case gaussFreq of
        Left    err     -> error err
        Right   gFreq   -> return gFreq

    gInfo <- case gaussParse of
        Left    err     -> error err
        Right   gInfo   -> return gInfo

    return (gInfo, gFreq)

-- Parses a file using a specified parser and returns the result in IO
doGenericParsing :: FilePath -> Parser a -> IO a
doGenericParsing filePath parser = do
    fileContent <- TIO.readFile filePath
    let gParse = parseOnly parser fileContent

    case gParse of
        Left    err     -> error err
        Right   gInfo   -> return gInfo

getListOfCoords :: Molecule -> [CartesianCoord]
getListOfCoords molA' = molA' ^.. atoms . each . coordinate

-- Calculates the center of mass of a given molecule of type Molecule (from Haskem.Types)molA
calculateCoM :: Molecule -> CartesianCoord
calculateCoM molecule' = CartesianCoord xCoM yCoM zCoM
    where
        atomsList   = getListOfCoords molecule'
        massList    = map elemMass $ molecule' ^.. atoms . each . atomNumber
        fullMass    = sum massList
        xCoM        = (/ fullMass) $ sum $ zipWith (*) massList $ map (^. xCoord) atomsList
        yCoM        = (/ fullMass) $ sum $ zipWith (*) massList $ map (^. yCoord) atomsList
        zCoM        = (/ fullMass) $ sum $ zipWith (*) massList $ map (^. zCoord) atomsList

-- Calculate a displacement vector from two molecular geometries.
-- Outputs a maybe molecule (in case the geometries do not match
-- (checked via number of atoms and the list of element numbers))
calcDisplacementVector :: Molecule -> Molecule -> [CartesianCoord]
calcDisplacementVector molA molB =
    do
        let nrOfAtomsA  = molA ^. nrOfAtoms
            nrOfAtomsB  = molA ^. nrOfAtoms
            elemA   = molA ^.. atoms . each . atomNumber
            elemB   = molB ^.. atoms . each . atomNumber
        if (nrOfAtomsA == nrOfAtomsB) && (elemA == elemB)
            then do
                let atomsA  = getListOfCoords molA
                    atomsB  = getListOfCoords molB
                    xD      = zipWith (-) (map (^. xCoord) atomsA) $ map (^. xCoord) atomsB
                    yD      = zipWith (-) (map (^. yCoord) atomsA) $ map (^. yCoord) atomsB
                    zD      = zipWith (-) (map (^. zCoord) atomsA) $ map (^. zCoord) atomsB
                zipWith3 CartesianCoord xD yD zD
            else
                []

-- Writes a "Molecule" back to disk as a .xyz file
writeMoleculeToXYZ :: Molecule -> FilePath -> IO ()
writeMoleculeToXYZ molA fp = writeFile fp (show nrOfAtomsA
                                ++ "\nThanks for using Haskem!\n"
                                ++ T.unpack (T.unlines lineList))
    where
        coordinates = getListOfCoords molA
        nrOfAtomsA  = molA ^. nrOfAtoms
        elems       = map elemNum2Sym $ molA ^.. atoms . each . atomNumber
        lineList    = zipWith4 (F.format xyzFormat)
                        elems
                        (map (^. xCoord) coordinates)
                        (map (^. yCoord) coordinates)
                        (map (^. zCoord) coordinates)

projectDonN :: [CartesianCoord] -> NormalMode -> Double
projectDonN displCoordsA modesA = projectCoordNM
    where
        displCoordsNM   = modesA ^.. displacement . each .coordinate  :: [CartesianCoord]
        projectCoordNM  = sum [sum $ zipWith (*)
                            (map ( ^. cc) displCoordsA)
                            (map ( ^. cc) displCoordsNM) | cc <- [xCoord, yCoord, zCoord]]

calcProjectionVec :: [NormalMode] -> [CartesianCoord] -> [Double]
calcProjectionVec nModes displCoords = projectionVec
        where
            projectionVec = map (projectDonN displCoords) nModes

{-
Build redundant internal coordinates from the molecular cartesian structure
-}

calcLength3D :: CartesianCoord -> CartesianCoord -> Double
calcLength3D coord1 coord2 = length3D
        where
            length3D = sqrt . sum $ (** 2) <$> dVector
            dVector = [(coord1 ^. axis) - (coord2 ^. axis)
                            | axis <- [xCoord, yCoord, zCoord]]

setUpBonds :: Molecule -> [(String, String, Double)]
setUpBonds molA = bondsList
        where
            atomListN   = molA ^.. atoms . each . atomNumber
            atomListS   = map elemNum2Sym   atomListN
            vdWList     = map covRadius     atomListN
            enumCoords  =  zip4 [1..] atomListS vdWList $ getListOfCoords molA
            bondsList   = [(e1 ++ show i1, e2 ++ show i2, calcLength3D c1 c2)
                            | (i1, e1, r1, c1) <- enumCoords,
                              (i2, e2, r2, c2) <- enumCoords,
                              let l12 = calcLength3D c1 c2,
                              i1 < i2 && l12 < 1.3 * (r1 + r2)
                            ]

setUpBonds' :: Molecule -> [LinearBond]
setUpBonds' molA = bondsList
        where
            atomListN'   = molA ^.. atoms . each . atomNumber
            atomListS   = map elemNum2Sym   atomListN'
            vdWList     = map covRadius     atomListN'
            enumCoords  =  zip4 [1..] atomListS vdWList $ getListOfCoords molA
            bondsList   = [LinearBond (AtomSpec e1 i1) (AtomSpec e2 i2) l12
                            | (i1, e1, r1, c1) <- enumCoords,
                              (i2, e2, r2, c2) <- enumCoords,
                              i1 < i2,
                              let l12 = calcLength3D c1 c2,
                              l12 <= 1.3 * (r1+r2)]

setUpBonds'' :: Molecule -> LA.Matrix Double
setUpBonds'' molA = LA.reshape nrOfAtoms' bondsMatrix
        where
            nrOfAtoms'  = molA ^. nrOfAtoms
            atomListN'  = molA ^.. atoms . each . atomNumber
            atomListS   = map elemNum2Sym   atomListN'
            vdWList     = map covRadius     atomListN'
            enumCoords  = zip4 [1..] atomListS vdWList $ getListOfCoords molA
            bondsMatrix = VS.fromList [bl | (i1, e1, r1, c1) <- enumCoords,
                                            (i2, e2, r2, c2) <- enumCoords,
                                            let l12 = calcLength3D c1 c2,
                                            let bl  = if l12 < 1.3 * (r1+r2) then 1 else 0]

setUpBonds''' :: Molecule -> [IS.IntSet]
setUpBonds''' molA = bondsMatrix
        where
            nrOfAtoms'  = molA ^. nrOfAtoms
            atomListN'  = molA ^.. atoms . each . atomNumber
            atomListS   = map elemNum2Sym   atomListN'
            vdWList     = map covRadius     atomListN'
            enumCoords  = zip4 [1..] atomListS vdWList $ getListOfCoords molA
            bondsMatrix = IS.fromList <$> chunksOf nrOfAtoms' [bl |
                                                (i1, e1, r1, c1) <- enumCoords,
                                                (i2, e2, r2, c2) <- enumCoords,
                                                let l12 = calcLength3D c1 c2,
                                                let bl  = if l12 < 1.3 * (r1+r2) then 1 else 0]
-- TODO!

removeDuplicates :: (Ord a) => [a] -> [a]
removeDuplicates = map head . group . sort

checkBondCompleteness :: Molecule -> [LinearBond] -> [Atom]
checkBondCompleteness mol' listOfBonds' = missing
        where
            nrOfAtoms'      = mol' ^. nrOfAtoms
            atomNums'       = removeDuplicates $ map (^. bAtomA . geomNum) listOfBonds'
                                ++ map (^. bAtomB . geomNum) listOfBonds'
            missingIndxs    = findIndices (`notElem` atomNums')  [1..nrOfAtoms']
            missing         = case missingIndxs of
                               []   -> []
                               _    -> concat [mol' ^.. atoms . ix ii | ii <- missingIndxs]

defineFragments     :: [IS.IntSet] -> [IS.IntSet]
defineFragments []      = []
defineFragments [x]     = [x]

defineFragments (a:as)  =
    if (all null overlaps)
        then (a:as)
        else defineFragments reducedSet
    where
        setsWithOverlap ::  IS.IntSet -> [IS.IntSet] -> [IS.IntSet]
        setsWithOverlap bond bondsList = [ii | ii <- bondsList,
                                                not . IS.null $ IS.intersection bond ii,
                                                bond /= ii]

        overlaps =
            [setsWithOverlap ((a:as) !! ii) (deleteNth ii (a:as))
                | ii <- [0..length(a:as) - 1]
            ]
            :: [[IS.IntSet]]

        reducedSet =
            nub
            [ IS.unions (((a:as) !! i) : (overlaps !! i))
                | i <- [0 .. length (a:as) - 1]
            ]

deleteNth :: Int -> [a] -> [a]
deleteNth n l = (take n l) ++ (drop (n + 1) l)

-- defineFragments'    :: VS.Vector (IM.IntMap Int) -> VS.Vector (IM.IntMap Int)
-- defineFragments'


{-# LANGUAGE OverloadedStrings #-}
module Data.Outputable.Instances where

import           Control.Applicative            (ZipList)
import           Data.Char                      (isAlphaNum)
import           Data.Complex                   (Complex)
import           Data.Functor.Identity          (Identity)
import           Data.Int                       (Int16, Int32, Int64, Int8)
import           Data.List.NonEmpty             (NonEmpty)
import qualified Data.Map                       as Map
import           Data.Monoid                    (All, Any, Dual, First, Last,
                                                 Product, Sum)
import           Data.Outputable.Class
import qualified Data.Semigroup                 as S (First, Last, Max, Min,
                                                      Option, WrappedMonoid)
import           Data.Text                      (Text, unpack)
import qualified Data.Text.Lazy                 as TL (Text, unpack)
import           Data.Version                   (Version)
import           Data.Void                      (Void)
import           Data.Word                      (Word16, Word32, Word64, Word8)
import           GHC.Generics
import           Numeric.Natural                (Natural)
import           System.Exit                    (ExitCode)
import           Text.PrettyPrint.HughesPJClass

viaPretty :: Pretty a => Int -> a -> Doc
viaPretty _ = pPrint

viaPrettyNum :: (Ord a, Num a, Pretty a) => Int -> a -> Doc
viaPrettyNum n x
  | n /= 0 && x < 0 = parens $ pPrint x
  | otherwise = pPrint x

viaShow _ x = text $ show x
viaShowNum n x
  | n /= 0 && x < 0 = parens $ viaShow n x
  | otherwise = viaShow n x

instance Outputable Double where pprPrec = viaPrettyNum
instance Outputable Float where pprPrec = viaPrettyNum
instance Outputable Int where pprPrec = viaPrettyNum
instance Outputable Int8 where pprPrec = viaShowNum
instance Outputable Int16 where pprPrec = viaShowNum
instance Outputable Int32 where pprPrec = viaShowNum
instance Outputable Int64 where pprPrec = viaShowNum
instance Outputable Integer where pprPrec = viaPrettyNum
instance Outputable Natural where pprPrec = viaShow
instance Outputable Word where pprPrec = viaShow
instance Outputable Word8 where pprPrec = viaShow
instance Outputable Word16 where pprPrec = viaShow
instance Outputable Word32 where pprPrec = viaShow
instance Outputable Word64 where pprPrec = viaShow
instance Outputable TL.Text where pprPrec n x = doubleQuotes $ text $ TL.unpack x
instance Outputable Text where pprPrec n x = doubleQuotes $ text $ unpack x
instance Outputable a => Outputable [a] where pprPrec _ = pprList

instance Outputable Char where
  pprPrec _ x = quotes $ char x
  pprList = pPrint

instance Outputable Bool
instance Outputable Ordering
instance Outputable ()
instance Outputable DecidedStrictness
instance Outputable SourceStrictness
instance Outputable SourceUnpackedness
instance Outputable Associativity
instance Outputable Fixity
instance Outputable Any
instance Outputable All
instance Outputable ExitCode
instance Outputable Version
instance Outputable Void
instance Outputable a => Outputable (Maybe a)
instance Outputable p => Outputable (Par1 p)
instance Outputable a => Outputable (NonEmpty a)
instance Outputable a => Outputable (Product a)
instance Outputable a => Outputable (Sum a)
instance Outputable a => Outputable (Dual a)
instance Outputable a => Outputable (Last a)
instance Outputable a => Outputable (First a)
instance Outputable a => Outputable (Identity a)
instance Outputable a => Outputable (ZipList a)
instance Outputable a => Outputable (S.Option a)
instance Outputable m => Outputable (S.WrappedMonoid m)
instance Outputable a => Outputable (S.Last a)
instance Outputable a => Outputable (S.First a)
instance Outputable a => Outputable (S.Max a)
instance Outputable a => Outputable (S.Min a)
instance Outputable a => Outputable (Complex a)
instance (Outputable a, Outputable b) => Outputable (Either a b)
instance (Outputable a, Outputable b) => Outputable (a, b) where
  pprPrec _ (a, b) = parens $ fsep $ punctuate comma [ppr a, ppr b]
instance (Outputable a, Outputable b, Outputable c) => Outputable (a, b, c) where
  pprPrec _ (a, b, c) = parens $ fsep $ punctuate comma [ppr a, ppr b, ppr c]
instance (Outputable a, Outputable b, Outputable c, Outputable d) => Outputable (a, b, c, d) where
  pprPrec _ (a, b, c, d) = parens $ fsep $ punctuate comma [ppr a, ppr b, ppr c, ppr d]
instance (Outputable a, Outputable b, Outputable c, Outputable d, Outputable e) => Outputable (a, b, c, d, e) where
  pprPrec _ (a, b, c, d, e) = parens $ fsep $ punctuate comma [ppr a, ppr b, ppr c, ppr d, ppr e]

instance (Outputable k, Outputable v) => Outputable (Map.Map k v) where
  pprPrec _ m = "fromList" <+> ppr (Map.toList m)