{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnicodeSyntax #-}
module Day20 ( runSolution) where

import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.Attoparsec.ByteString (Parser)
import qualified Data.Attoparsec.ByteString as DAB
import Data.Attoparsec.ByteString.Char8 (anyChar, space, notChar, decimal, string, char, endOfLine)

import Data.Array.IArray (Array, (!))
import qualified Data.Array.IArray as IA
import qualified Data.Complex as DC
import qualified Linear as L
import Control.Arrow
import qualified Data.List as DL
import Control.Lens
import Data.Function (on)
import Data.Maybe (fromJust)
import qualified Control.Monad as CM
import qualified Data.Set as Set
import qualified Text.PrettyPrint.Boxes as Boxes
import qualified Control.Foldl as L

import Debug.Trace

type Picture = Array (L.V2 Int) Char

type Piece = (Int, Picture)

-- A `FittingRow` is a list of pieces that fit side by side in an horizontal
-- line.
type FittingRow = [Piece]

-- A `FittingPuzzle` is a list of rows where each row matches the one above and
-- below.
type FittingPuzzle = [FittingRow]

data HorizontalGrowthDirection = ToTheLeft | ToTheRight deriving Show
data VerticalGrowthDirection = Upwards | Downwards deriving Show

-- | Build a solution adding pieces to the right
--
-- For example, starting with P0 it adds P1 onwards where
-- P1 matches the *right* edge of P0.
--
-- (P0) P1 P2 P3
--
-- It returns all posibilities. That is
--
-- [] : Don't add more pieces
-- [P1] : Just add P1
-- [P1, P2] : Just add P1 and P2
-- ...
--
-- This leaves free pieces to build other parts of the solution. E.g: to the
-- left of P0.
--
-- It doesn't include the starting piece (P0)
--
-- if HorizontalGrowthDirection is ToTheLeft. P1 fits to the left of P0
--
-- P1 (P0)
--
-- regardless the resulting rows are in reverse order. That is for:
--
-- P3 P2 P1 (P0)
--
-- the output is [[],[P1],[P1,P2],[P1,P2,P3]]
--
-- This means that in this case the alternatives shouls be reversed to be
-- displayed in the correct order.
topHorizontalSolution :: HorizontalGrowthDirection -> Piece -> [(Int, [Picture])] -> [FittingRow]
topHorizontalSolution growthDir (_, topLeftPic) pieces = unfoldrM unfolder (topLeftPic, pieces)
	where
	unfolder :: (Picture, [(Int, [Picture])]) -> [Maybe (Piece, (Picture, [(Int, [Picture])]))]
	-- `refPic` is the picture the next piece should fit on the side.
	-- `pieces` are the remmaining unused pieces
	--
	-- The alternatives are:
	-- 	Stop
	-- 	One alternative for each piece that matches the refPic on the side.
	unfolder (refPic, pieces) = (Nothing :) $
		fmap
			(\sidePiece -> Just (sidePiece,
				( snd sidePiece -- The next refPic
				, removePiece sidePiece pieces))) -- The next pieces
			(concatMap filterSideMatchs pieces)
		where

		filterSideMatchs :: (Int, [Picture]) -> [(Int, Picture)]
		filterSideMatchs (pid, trans) = (pid,) <$> filter selector trans

		selector = case growthDir of
			ToTheRight -> match Sideways refPic
			ToTheLeft -> (\t -> match Sideways t refPic)

-- | Build a solution adding pieces to both sides
topMiddleSolution :: (Int, Picture) -> [(Int, [Picture])] -> [FittingRow]
topMiddleSolution topCenter@(_, topCenterPic) pieces = [topCenter] : do
	rightPiece <- pieces >>= filterRightMatches
	leftPiece <- removePiece rightPiece pieces >>= filterLeftMatches

	let remPieces = removePiece rightPiece (removePiece leftPiece pieces)

	rightSolution <- topHorizontalSolution ToTheRight rightPiece remPieces

	-- The soultions need to be reversed because they are to the left of the
	-- `topCenter` piece.
	leftSolution <- fmap reverse $ topHorizontalSolution ToTheLeft leftPiece $
		DL.foldl' (flip removePiece) remPieces rightSolution

	pure (leftSolution ++ [leftPiece, topCenter, rightPiece] ++ rightSolution)
	where
	filterRightMatches :: (Int, [Picture]) -> [(Int, Picture)]
	filterRightMatches (pid, trans) = (pid,) <$> filter (match Sideways topCenterPic) trans

	filterLeftMatches :: (Int, [Picture]) -> [(Int, Picture)]
	filterLeftMatches (pid, trans) = (pid,) <$> filter (\t -> match Sideways t topCenterPic) trans

-- | Build a solution adding a new row at the bottom
downwardSolution :: FittingRow -> [(Int, [Picture])] -> [FittingPuzzle]
downwardSolution _ [] = [[]]
downwardSolution (t : ts) pieces@(_ : _) = [] : do
	bottomPiece <- pieces >>= filterBelowMatches
	bsol <- rowSolution Downwards bottomPiece ts
		(filter (not . null . snd) $ removePiece bottomPiece pieces)
	CM.guard (length (bottomPiece : bsol) == 12)
	[bottomPiece : bsol] : do
		let remPieces = DL.foldl' (flip removePiece) pieces (bottomPiece : bsol)
		bsol' <- downwardSolution (bottomPiece : bsol) remPieces
		[ (bottomPiece : bsol) : bsol']
	where

	filterBelowMatches (pid, trans) = (pid,) <$> filter (match Below (snd t)) trans

-- | Build a row at the bottom or top of another.
--
-- Given a row of pieces (that fit together, although this funtion doesn't check
-- that this is the case)
--
-- T1 T2 T3 T4
--
-- And a piece P0 (that could match a possible T0 to the left of T1). It builds
-- all possible solutions (S1..) that fit to the right and below the T's.
--
--    T1 T2 T3 T4
-- P0 S1 S2 S3 S4
--
-- So S1 fits P0 from the right and T1 from below. Similarly T2, etc...
--
-- This function is not greedy. It returns all possible solutions. In the
-- running example that is:
--
-- [[], [S1], [S1, S2], [S1, S2, S3], [S1, S2, S3, S4]]
--
-- If there are more solutions, thouse too should be included on the list.
--
-- The solution is built to the right.
--
-- In case the VerticalGrowthDirection is Upwards the above example becomes:
--
-- P0 S1 S2 S3 S4
--    T1 T2 T3 T4
--
-- This function can be enhanced to build solutions to the left.
rowSolution :: VerticalGrowthDirection -> Piece -> FittingRow -> [(Int, [Picture])] -> [FittingRow]
rowSolution growthDir (_, sidePic) topRow pieces = unfoldrM unfolder (sidePic, (topRow, pieces))
	where
	unfolder :: (Picture, (FittingRow, [(Int, [Picture])])) ->
		[Maybe (Piece, (Picture, (FittingRow, [(Int, [Picture])])))]
	unfolder (_, ([], _)) = [Nothing]
	-- `refPic` is P0 in the comment of the parent function. It changes to (S1,
	-- S2,...) as the computation advances.
	--
	-- `(t : ts)` is the top (T1, T2, ...). It is consumed one by one.
	--
	-- `pieces` are the available pieces. The ones used are removed from the
	-- list for the next iteration.
	unfolder (refPic, (t : ts, pieces)) = (Nothing :) $
		fmap
			(\sidePiece -> Just (sidePiece,
				( snd sidePiece
				,	( ts
					, removePiece sidePiece pieces))))
			(concatMap filterBelowMatches pieces)
		where

		filterBelowMatches (pid, trans) = (pid,) <$> filter selector trans

		-- True if `p` matches P0 from the right and T1 from below/above.
		selector :: Picture -> Bool
		selector p = match Sideways refPic p && (
			case growthDir of
				Downwards -> match Below (snd t) p
				Upwards -> match Below p (snd t) )

-- | Build a solution adding a new row on the top
upwardSolution :: FittingRow -> [(Int, [Picture])] -> [FittingPuzzle]
upwardSolution _ [] = [[]]
upwardSolution (b : bs) pieces@(_ : _) = [] : do
	topPiece <- pieces >>= filterUpwardsMatches
	tsol <- rowSolution Upwards topPiece bs
		(filter (not . null . snd) $ removePiece topPiece pieces)
	CM.guard (length (topPiece : tsol) == 12)
	[topPiece : tsol] : do
		let remPieces = DL.foldl' (flip removePiece) pieces (topPiece : tsol)
		tsol' <- upwardSolution (topPiece : tsol) remPieces
		[tsol' ++ [topPiece : tsol]]
	where

	filterUpwardsMatches (pid, trans) = (pid,) <$> filter (\t -> match Below t (snd b)) trans

middleRowSolution :: FittingRow -> [(Int, [Picture])] -> [FittingPuzzle]
middleRowSolution middleRow pieces = [middleRow] : do
	ts <- upwardSolution middleRow pieces
	let remPieces = DL.foldl' (flip removePiece) pieces (concat ts)
	bs <- downwardSolution middleRow remPieces
	pure $ ts ++ [middleRow] ++ bs

removePiece :: Piece -> [(Int, [Picture])] -> [(Int, [Picture])]
removePiece piece = filter ((/= fst piece) . fst)

solver input = case transformed of
	[] -> []
	(p : ps) -> do
		trans <- snd p
		topRow <- topMiddleSolution (fst p, trans) ps
		CM.guard (length topRow == 12)
		let remPieces = DL.foldl' (flip removePiece) (p : ps) topRow
		bs <- middleRowSolution topRow remPieces
		CM.guard (length bs == 12)
		pure bs

	where
	transformed = fmap (second transforms) input

data RelPos = Below | Sideways deriving Show

match :: RelPos -> Picture -> Picture -> Bool
match relPos refPic testPic = and $ case relPos of
	Below -> zipWith (==)
		(fmap ( (refPic !) . (`L.V2` rows)) [1..cols]) -- the bottom row of refPic
		(fmap ( (testPic !) . (`L.V2` 1)) [1..cols]) -- the top row of testPic
	Sideways -> zipWith (==)
		(fmap ( (refPic !) . (cols `L.V2`)) [1..rows]) -- the rightmost col of refPic
		(fmap ( (testPic !) . (1 `L.V2`)) [1..rows]) -- the leftmost col of testPic
	where
	(_, L.V2 cols rows) = IA.bounds refPic

-- All unique transformations of the picture
transforms :: Picture -> [Picture]
transforms basePic =
	-- try to exploit the picture's symmetry. Flip symmetries.
	Set.toList $ Set.fromList $
	-- An Horizontal flip is a Vertical flip and a 180 degree rotation.
	rotations ++ ( (Vertical 🤸) <$> rotations )
	where
	-- try to exploit the picture's symmetry. Stop ASAP. 90 degree rotation or
	-- 180 degree rotation symmetry
	rotations = basePic : takeWhile (/= basePic) (tail $ iterate (ClockWise 🥏) basePic )

-- | Removes the edges from a picture.
removeEdge :: Picture -> Picture
removeEdge pic = IA.array (L.V2 1 1, L.V2 (xlim - 2) (ylim - 2)) $
		zip (fmap (+ L.V2 (-1) (-1)) nonBorderPixels) (fmap (pic !) nonBorderPixels)
	where
	nonBorderPixels = L.V2 <$> [(orgx + 1)..(xlim-1)] <*> [(orgy + 1)..(ylim-1)]
	(L.V2 orgx orgy, L.V2 xlim ylim) = IA.bounds pic

-- | Construct one big picture from a fitting puzzle
bigPicture :: FittingPuzzle -> Picture
bigPicture pzl = IA.array (L.V2 1 1, L.V2 xlim ylim) $ concatMap (\(y, cs) ->
	fmap (first (`L.V2` y)) cs)  $ zip [1..] $ fmap (zip [1..]) picStrs
	where
	ylim = length picStrs
	xlim = length $ head picStrs
	-- Fold all rows together
	picStrs = concatMap rows pictures
	-- Build a big-row
	rows :: [Picture] -> [String]
	rows pcs = DL.foldl' rowFolder [] pcs
	rowFolder :: [String] -> Picture -> [String]
	rowFolder [] p = pl p
	rowFolder acc@(_ : _) p = zipWith (++) acc (pl p)
	-- pictures without borders
	pictures :: [[Picture]]
	pictures = fmap (fmap (removeEdge . snd)) pzl
	-- picture lines
	pl :: Picture -> [String]
	pl p = fmap (fmap snd) $ DL.groupBy ((==) `on` ((^. _2) . fst)) $ DL.sortOn ((^. _2) . fst) $  IA.assocs p

-- | Parse the input into a 2D map of chars using a default character if some
-- lines are shorter
parsePicture :: Char -> Parser Picture
parsePicture def = do
	parsedLines <- parseLines
	let maxRowLength =  (maximum . fmap length) parsedLines
	let nRows = length parsedLines
	pure $ (IA.array (L.V2 1 1, L.V2 maxRowLength nRows) . concatMap mapper . zip [1..nRows] . fmap (makeRow maxRowLength)) parsedLines
	where
	mapper :: (Int, [(Int, Char)]) -> [(L.V2 Int, Char)]
	mapper (y, row) = fmap (first (`L.V2` y)) row
	parseLines :: Parser [[Char]]
	parseLines = parseLine `DAB.sepBy1'` endOfLine
	parseLine :: Parser [Char]
	parseLine = DAB.many1' (notChar '\n')
	makeRow :: Int -> [Char] -> [(Int, Char)]
	makeRow size content = zip [1..size] (content ++ repeat def)

parseInput :: Parser [Piece]
parseInput = DAB.many1' parseSingle
	where
	parseSingle = (,) <$>
		(string "Tile " *> decimal <* char ':' <* endOfLine) <*>
		parsePicture 'X' <* DAB.choice [endOfLine <* endOfLine, endOfLine]

-- | Cound the #'s present in a picture
-- Right answer is 1607
waterRoughness :: Picture -> Int
waterRoughness = length . filter ((=='#') . snd) . IA.assocs

-- | Detect the monsters and mark them with 'O'
detectMonsters :: Picture -> Picture
detectMonsters bigPic = bigPic IA.// fmap (,'O') monsters
	where
	-- The positions of the parts of the monsters
	monsters = concatMap (fmap fst) $ filter (all ((=='#') . snd))  $ fmap mapper origins
	mapper origin = fmap ( (id &&& (bigPic !)) . (+ origin)) monsterMold
	-- All valid positions of the mosnter mold in the bigPic
	origins = L.V2 <$> [picMinX..picMaxX - maxx] <*> [picMinY..picMaxY - maxy]
	(L.V2 picMinX picMinY, L.V2 picMaxX picMaxY) = IA.bounds bigPic
	-- The width and height of the moster mold
	(Just maxx, Just maxy) = L.fold ((,) <$> L.premap (^. L._x) L.maximum <*> L.premap (^. L._y) L.maximum) monsterMold
	-- The positions where the seaMonsterInput has a '#'
	-- Starts at 0, 0 becase it's intended to be uses as an offset from an
	-- origin.
	monsterMold :: [L.V2 Int]
	monsterMold = concatMap (\(y, xs) -> fmap ((`L.V2` y) . fst) xs ) $
		zip [0..] $
		fmap
			( filter ((== '#') . snd) . zip [0..] )
			seaMonsterInput

runSolution :: FilePath -> IO ()
runSolution filePath = do
	putStrLn "**Day 20**"
	contents <- BS.readFile filePath
	let parseResult = DAB.parseOnly parseInput contents
	case parseResult of
		Left err -> putStrLn $ "Error :" ++ err
		Right input -> do
			CM.forM_ (take 1 $ solver input) $ \ pzl -> do
				putStrLn $ Boxes.render $ picLines $ bigPicture pzl
				CM.forM_ (transforms $ bigPicture pzl) $ print . waterRoughness . detectMonsters
				putStrLn "-- End Solution --"
	putStrLn "Done"

seaMonsterInput :: [String]
seaMonsterInput =
 [ "                  # "
 , "#    ##    ##    ###"
 , " #  #  #  #  #  #   " ]

puzzleLines :: FittingPuzzle -> Boxes.Box
puzzleLines pzl = Boxes.vsep 0 Boxes.left pls
	where
	pls = fmap (Boxes.hsep 1 Boxes.top) $ fmap pieceLines <$> pzl

pieceLines :: Piece -> Boxes.Box
pieceLines (pid, pic) = Boxes.text (show pid) Boxes.// picLines pic

picLines :: Picture -> Boxes.Box
picLines pic = Boxes.vcat Boxes.left $ fmap (Boxes.text . (snd <$>)) $
	DL.groupBy ((==) `on` ((^. _2) . fst)) $ DL.sortOn ((^. _2) . fst) $  IA.assocs pic

data RotDir = ClockWise | CounterClockWise deriving (Show, Eq)

--`#rotp`
(🥏)  :: RotDir -> Picture -> Picture
(🥏)  rotDir pic = IA.array (org, L.V2 ylim xlim) $
	first (untranslate . rotator rotDir . translate) <$> IA.assocs pic
	where

	translate :: L.V2 Int -> L.V2 Float
	translate (L.V2 x y) = L.V2 (fromIntegral x - 1 - fxlim) (fromIntegral y - 1 - fylim)

	untranslate :: L.V2 Float -> L.V2 Int
	untranslate (L.V2 x y) = round <$> L.V2 (x + 1 + fxlim) (y + 1 + fylim)

	rotator :: RotDir -> L.V2 Float -> L.V2 Float
	rotator CounterClockWise (L.V2 x y) = L.V2 y (negate x)
	rotator ClockWise (L.V2 x y) = L.V2 (negate y) x

	(org, L.V2 xlim ylim) = IA.bounds pic
	fxlim :: Float
	fylim :: Float
	(fxlim, fylim) = (fromIntegral (xlim - 1) / 2, fromIntegral (ylim - 1) / 2)

data FlipAxis = Vertical | Horizontal deriving Show

--`#flpp`
(🤸) :: FlipAxis -> Picture -> Picture
(🤸) flipAxis pic = IA.array (org, L.V2 ylim xlim) $
	first (untranslate . flipper flipAxis . translate) <$> IA.assocs pic

	where

	translate :: L.V2 Int -> L.V2 Float
	translate (L.V2 x y) = L.V2 (fromIntegral x - 1 - fxlim) (fromIntegral y - 1 - fylim)

	untranslate :: L.V2 Float -> L.V2 Int
	untranslate (L.V2 x y) = round <$> L.V2 (x + 1 + fxlim) (y + 1 + fylim)

	flipper :: FlipAxis -> L.V2 Float -> L.V2 Float
	flipper Horizontal (L.V2 x y) = L.V2 x (negate y)
	flipper Vertical (L.V2 x y) = L.V2 (negate x) y

	(org, L.V2 xlim ylim) = IA.bounds pic
	fxlim :: Float
	fylim :: Float
	(fxlim, fylim) = (fromIntegral (xlim - 1) / 2, fromIntegral (ylim - 1) / 2)

-- | unfoldr in a monad.
--
-- This is missing from Data.List
unfoldrM :: Monad m => (a -> m (Maybe (b, a))) -> a -> m [b]
unfoldrM f a = f a >>= \case
	Nothing -> pure []
	Just (b, a') -> (b :) <$> unfoldrM f a'


import Criterion.Main
import Criterion.Types
import Dimacs (parseDIMACSFromFile)
import Internal.Sat (solveDIMACS)
import Sat (solve)
import Statistics.Types
import Test.HUnit

solveFile :: String -> Bool -> IO ()
solveFile fileName isSat = do
  parsed <- parseDIMACSFromFile ("resources/problems/" ++ fileName)
  case parsed of
    Right dimacs -> do
      case solveDIMACS dimacs of
        Just _ -> assertBool "" isSat
        Nothing -> assertBool "" (not isSat)
      return ()
    left -> assertFailure $ "" ++ fileName ++ "cannot be parsed"

main :: IO ()
main =
  Criterion.Main.defaultMainWith
    (defaultConfig {timeLimit = 30})
    [ bgroup
        "sat tests"
        [ benchit "add4.cnf" False,
          benchit "diamond1.cnf" False,
          benchit "diamond2.cnf" False,
          benchit "diamond3.cnf" False,
          benchit "ph2.cnf" False,
          benchit "ph3.cnf" False,
          benchit "prime4.cnf" True,
          benchit "prime9.cnf" True,
          benchit "prime25.cnf" True,
          benchit "prime49.cnf" True,
          benchit "prime121.cnf" True
        ]
    ]
  where
    benchit fileName isSat = bench fileName $ whnfIO $ solveFile fileName isSat


-- | Post-mortem tool main.

import           Universum

import           Data.List (last)
import qualified Data.Map.Strict as M
import           System.FilePath
import           System.IO (hPutStrLn)
import           Text.Printf (hPrintf)

import           Pos.Util.Util (histogram)

import           Options
import           Statistics
import           Types

main :: IO ()
main = parseOptions >>= \case
    Overview sampleProb logDirs -> do
        showLogDirs logDirs
        err $ "sample probability: " ++ show sampleProb
        err ""
        xs <- forM logDirs $ flip processLogDirOverview sampleProb
        chart xs "times.png"
        err "wrote times chart"
    Focus txHash logDir         -> do
        err $ "transaction hash: " ++ show txHash
        err $ "logs directory: " ++ show logDir
        let focusFile = ("focus_" ++ extractName logDir ++ "_" ++ toString txHash) <.> "csv"
        runJSONFold logDir (focusF txHash) >>= focusToCSV focusFile
        err $ "wrote result to " ++ show focusFile
    TxRelay logDirs             -> do
        showLogDirs logDirs
        err ""
        for_ logDirs processLogDirTxRelay
    Throughput txWindow waitWindow logDirs   -> do
        showLogDirs logDirs
        err $ "tx window: " ++ show txWindow
        err $ "wait window: " ++ show waitWindow
        err ""
        for_ logDirs $ processLogDirThroughput txWindow waitWindow

showLogDirs :: [FilePath] -> IO ()
showLogDirs logDirs = do
    err "log directories: "
    for_ logDirs $ \d -> err $ " - " ++ show d

processLogDirOverview :: FilePath -> Double -> IO (String, Map TxHash (Maybe Timestamp))
processLogDirOverview logDir sampleProb = do
    err $ "processing log directory " ++ show logDir ++ " ..."

    (rc, g, mp, cr, ft) <-
        runJSONFold logDir $ (,,,,) <$> receivedCreatedF
                                    <*> graphF
                                    <*> memPoolF
                                    <*> txCntInChainF
                                    <*> txFateF
    let total    = M.size rc
        included = sort $ mapMaybe snd $ M.toList rc
        lost     = total - length included
    err $ "total number of received transactions: " ++ show total
    err $ "included in blockchain: " ++ show (length included)
    err $ "lost transactions: " ++ show lost
    err $ "blocks in chain: " ++ show (length cr)
    err $ "mem pool entries: " ++ show (length mp)

    let dirName   = extractName logDir
    let graphFile = getName "graph" dirName "png"
    b <- writeGraph graphFile g
    when b $ err $ "wrote graph png to " ++ show graphFile

    let csvFile = getName "csv" dirName "csv"
    txCntInChainMemPoolToCSV csvFile sampleProb cr mp
    err $ "wrote csv file to " ++ show csvFile
    let reportFile = getName "report" dirName "txt"
    void (reportTxFate reportFile ft)
    err $ "wrote report file to " ++ show reportFile
    err $ "processing log directory " ++ show logDir ++ " done"
    err ""
    return (dirName, rc)

processLogDirTxRelay :: FilePath -> IO ()
processLogDirTxRelay logDir = do
    err $ "processing log directory " ++ show logDir ++ " ..."
    m <- runJSONFold logDir txReceivedF
    err $ "total number of received transactions: " ++ show (M.size m)
    let xs         = sortBy (compare `on` \(_, ys) -> negate $ length ys) $ M.toList m
        dirName    = extractName logDir
        relaysFile = getName "relays" dirName "txt"
    withFile relaysFile WriteMode $ \h ->
        for_ xs $ \(tx, ys) ->
            hPutStrLn h $ toString tx ++ ": " ++ show ys
    err $ "wrote relay data to " ++ show relaysFile
    let hist     = M.toList $ histogram $ map (length . snd) xs
        histFile = getName "histogram" dirName "txt"
    withFile histFile WriteMode $ \h ->
        for_ hist $ uncurry $ hPrintf h "%4d:%6d\n"
    err $ "wrote histogram to " ++ show histFile

processLogDirThroughput :: Double -> Double -> FilePath -> IO ()
processLogDirThroughput txWindow waitWindow logDir = do
    err $ "processing log directory " ++ show logDir ++ " ..."
    (xs, ys) <- runJSONFold logDir $ (,) <$> txCntInChainF <*> memPoolF
    err $ "chain length: " ++ show (length xs) ++ " block(s)"
    err $ show (length ys) ++ " mem pool event(s)"
    let svgFile = getName "throughput" (extractName logDir) "svg"
    throughput svgFile txWindow waitWindow 1000 xs ys
    err $ "wrote chart to " ++ show svgFile

err :: String -> IO ()
err = hPutStrLn stderr

getName :: FilePath -> String -> String -> FilePath
getName template name ext =
    let dir   = takeDirectory template
        base  = takeBaseName template
        base' =  base ++ "_" ++ name
    in  dir </> base' <.> ext

extractName :: FilePath -> String
extractName = last . splitDirectories


{-# LANGUAGE CPP #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FunctionalDependencies #-}

#ifndef MIN_VERSION_base
#define MIN_VERSION_base(x,y,z) 1
#endif
-----------------------------------------------------------------------------
-- |
-- Module      :  Control.Lens.Each
-- Copyright   :  (C) 2012-13 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  non-portable
--
-----------------------------------------------------------------------------
module Control.Lens.Each
  (
  -- * Indices
    Index
  -- * Each
  , Each(..)
  ) where

import Control.Applicative
import Control.Lens.Cons as Lens
import Control.Lens.Internal.Deque
import Control.Lens.Internal.Setter
import Control.Lens.Indexed as Lens
import Control.Lens.Iso
import Control.Lens.Type
import Control.Lens.Traversal
import Data.Array.Unboxed as Unboxed
import Data.Array.IArray as IArray
import Data.ByteString as StrictB
import Data.ByteString.Lazy as LazyB
import Data.Complex
import Data.Foldable as Foldable
import Data.Functor.Identity
import Data.HashMap.Lazy as HashMap
import Data.HashSet
import Data.Int
import Data.IntMap as IntMap
import Data.IntSet
import Data.List.NonEmpty as NonEmpty
import Data.Map as Map
import Data.Set
import Data.Sequence as Seq
import Data.Text as StrictT
import Data.Text.Lazy as LazyT
import Data.Traversable
import Data.Tree as Tree
import qualified Data.Vector as Vector
import qualified Data.Vector.Primitive as Prim
import Data.Vector.Primitive (Prim)
import qualified Data.Vector.Storable as Storable
import Data.Vector.Storable (Storable)
import qualified Data.Vector.Unboxed as Unboxed
import Data.Vector.Unboxed (Unbox)
import Data.Word

-- | This is a common 'Index' type shared by 'Each', 'Control.Lens.At.At', 'Control.Lens.At.Contains' and 'Control.Lens.At.Ixed'.
type family Index (s :: *) :: *
type instance Index (e -> a) = e
type instance Index IntSet = Int
type instance Index (Set a) = a
type instance Index (HashSet a) = a
type instance Index [a] = Int
type instance Index (NonEmpty a) = Int
type instance Index (Seq a) = Int
type instance Index (a,b) = Int
type instance Index (a,b,c) = Int
type instance Index (a,b,c,d) = Int
type instance Index (a,b,c,d,e) = Int
type instance Index (a,b,c,d,e,f) = Int
type instance Index (a,b,c,d,e,f,g) = Int
type instance Index (a,b,c,d,e,f,g,h) = Int
type instance Index (a,b,c,d,e,f,g,h,i) = Int
type instance Index (IntMap a) = Int
type instance Index (Map k a) = k
type instance Index (HashMap k a) = k
type instance Index (Array i e) = i
type instance Index (UArray i e) = i
type instance Index (Vector.Vector a) = Int
type instance Index (Prim.Vector a) = Int
type instance Index (Storable.Vector a) = Int
type instance Index (Unboxed.Vector a) = Int
type instance Index (Complex a) = Int
type instance Index (Identity a) = ()
type instance Index (Maybe a) = ()
type instance Index (Tree a) = [Int]
type instance Index StrictT.Text = Int
type instance Index LazyT.Text = Int64
type instance Index StrictB.ByteString = Int
type instance Index LazyB.ByteString = Int64

-- $setup
-- >>> :set -XNoOverloadedStrings
-- >>> import Control.Lens
-- >>> import Data.Text.Strict.Lens as Text
-- >>> import Data.Char as Char

-- | Extract 'each' element of a (potentially monomorphic) container.
--
-- Notably, when applied to a tuple, this generalizes 'Control.Lens.Traversal.both' to arbitrary homogeneous tuples.
--
-- >>> (1,2,3) & each *~ 10
-- (10,20,30)
--
-- It can also be used on monomorphic containers like 'StrictT.Text' or 'StrictB.ByteString'.
--
-- >>> over each Char.toUpper ("hello"^.Text.packed)
-- "HELLO"
--
-- 'each' is an 'IndexedTraversal', so it can be used to access keys in many containers:
--
-- >>> itoListOf each $ Map.fromList [("hello",2),("world",4)]
-- [("hello",2),("world",4)]
--
-- >>> ("hello","world") & each.each %~ Char.toUpper
-- ("HELLO","WORLD")
class (Functor f, Index s ~ Index t) => Each f s t a b | s -> a, t -> b, s b -> t, t a -> s where
  each :: IndexedLensLike (Index s) f s t a b
#ifndef HLINT
  default each :: (Applicative f, Traversable g, s ~ g a, t ~ g b, Index s ~ Int, Index t ~ Int) => IndexedLensLike Int f s t a b
  each = traversed
  {-# INLINE each #-}
#endif

-- | @'each' :: 'IndexedTraversal' 'Int' (a,a) (b,b) a b@
instance (Applicative f, a~a', b~b') => Each f (a,a') (b,b') a b where
  each f ~(a,b) = (,) <$> f' (0 :: Int) a <*> f' 1 b
    where f' = Lens.indexed f
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' 'Int' (a,a,a) (b,b,b) a b@
instance (Applicative f, a~a2, a~a3, b~b2, b~b3) => Each f (a,a2,a3) (b,b2,b3) a b where
  each f ~(a,b,c) = (,,) <$> f' (0 :: Int) a <*> f' 1 b <*> f' 2 c
    where f' = Lens.indexed f
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' 'Int' (a,a,a,a) (b,b,b,b) a b@
instance (Applicative f, a~a2, a~a3, a~a4, b~b2, b~b3, b~b4) => Each f (a,a2,a3,a4) (b,b2,b3,b4) a b where
  each f ~(a,b,c,d) = (,,,) <$> f' (0 :: Int) a <*> f' 1 b <*> f' 2 c <*> f' 3 d
    where f' = Lens.indexed f
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' 'Int' (a,a,a,a,a) (b,b,b,b,b) a b@
instance (Applicative f, a~a2, a~a3, a~a4, a~a5, b~b2, b~b3, b~b4, b~b5) => Each f (a,a2,a3,a4,a5) (b,b2,b3,b4,b5) a b where
  each f ~(a,b,c,d,e) = (,,,,) <$> f' (0 :: Int) a <*> f' 1 b <*> f' 2 c <*> f' 3 d <*> f' 4 e
    where f' = Lens.indexed f
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' 'Int' (a,a,a,a,a,a) (b,b,b,b,b,b) a b@
instance (Applicative f, a~a2, a~a3, a~a4, a~a5, a~a6, b~b2, b~b3, b~b4, b~b5, b~b6) => Each f (a,a2,a3,a4,a5,a6) (b,b2,b3,b4,b5,b6) a b where
  each f ~(a,b,c,d,e,g) = (,,,,,) <$> f' (0 :: Int) a <*> f' 1 b <*> f' 2 c <*> f' 3 d <*> f' 4 e <*> f' 5 g
    where f' = Lens.indexed f
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' 'Int' (a,a,a,a,a,a,a) (b,b,b,b,b,b,b) a b@
instance (Applicative f, a~a2, a~a3, a~a4, a~a5, a~a6, a~a7, b~b2, b~b3, b~b4, b~b5, b~b6, b~b7) => Each f (a,a2,a3,a4,a5,a6,a7) (b,b2,b3,b4,b5,b6,b7) a b where
  each f ~(a,b,c,d,e,g,h) = (,,,,,,) <$> f' (0 :: Int) a <*> f' 1 b <*> f' 2 c <*> f' 3 d <*> f' 4 e <*> f' 5 g <*> f' 6 h
    where f' = Lens.indexed f
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' 'Int' (a,a,a,a,a,a,a,a) (b,b,b,b,b,b,b,b) a b@
instance (Applicative f, a~a2, a~a3, a~a4, a~a5, a~a6, a~a7, a~a8, b~b2, b~b3, b~b4, b~b5, b~b6, b~b7, b~b8) => Each f (a,a2,a3,a4,a5,a6,a7,a8) (b,b2,b3,b4,b5,b6,b7,b8) a b where
  each f ~(a,b,c,d,e,g,h,i) = (,,,,,,,) <$> f' (0 :: Int) a <*> f' 1 b <*> f' 2 c <*> f' 3 d <*> f' 4 e <*> f' 5 g <*> f' 6 h <*> f' 7 i
    where f' = Lens.indexed f
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' 'Int' (a,a,a,a,a,a,a,a,a) (b,b,b,b,b,b,b,b,b) a b@
instance (Applicative f, a~a2, a~a3, a~a4, a~a5, a~a6, a~a7, a~a8, a~a9, b~b2, b~b3, b~b4, b~b5, b~b6, b~b7, b~b8, b~b9) => Each f (a,a2,a3,a4,a5,a6,a7,a8,a9) (b,b2,b3,b4,b5,b6,b7,b8,b9) a b where
  each f ~(a,b,c,d,e,g,h,i,j) = (,,,,,,,,) <$> f' (0 :: Int) a <*> f' 1 b <*> f' 2 c <*> f' 3 d <*> f' 4 e <*> f' 5 g <*> f' 6 h <*> f' 7 i <*> f' 8 j
    where f' = Lens.indexed f
  {-# INLINE each #-}

#if MIN_VERSION_base(4,4,0)
-- | @'each' :: ('RealFloat' a, 'RealFloat' b) => 'IndexedTraversal' 'Int' ('Complex' a) ('Complex' b) a b@
instance Applicative f => Each f (Complex a) (Complex b) a b where
  each f (a :+ b) = (:+) <$> f' (0 :: Int) a <*> f' (1 :: Int) b
    where f' = Lens.indexed f
  {-# INLINE each #-}
#else
-- | @'each' :: 'IndexedTraversal' 'Int' ('Complex' a) ('Complex' b) a b@
instance (Applicative f, RealFloat a, RealFloat b) => Each f (Complex a) (Complex b) a b where
  each f (a :+ b) = (:+) <$> f' (0 :: Int) a <*> f' 1 b
    where f' = Lens.indexed f
  {-# INLINE each #-}
#endif

-- | @'each' :: 'IndexedTraversal' c ('Map' c a) ('Map' c b) a b@
instance Applicative f => Each f (Map c a) (Map c b) a b where
  each f m = sequenceA $ Map.mapWithKey f' m
    where f' = Lens.indexed f
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' c ('Map' c a) ('Map' c b) a b@
instance Applicative f => Each f (IntMap a) (IntMap b) a b where
  each f m = sequenceA $ IntMap.mapWithKey f' m
    where f' = Lens.indexed f
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' c ('HashMap' c a) ('HashMap' c b) a b@
instance Applicative f => Each f (HashMap c a) (HashMap c b) a b where
  each = HashMap.traverseWithKey . Lens.indexed
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' 'Int' [a] [b] a b@
instance Applicative f => Each f [a] [b] a b
  {- default each -}

-- | @'each' :: 'IndexedTraversal' () ('Identity' a) ('Identity' b) a b@
instance Functor f => Each f (Identity a) (Identity b) a b where
  each f (Identity a) = Identity <$> Lens.indexed f () a
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' () ('Maybe' a) ('Maybe' b) a b@
instance Applicative f => Each f (Maybe a) (Maybe b) a b where
  each f (Just a) = Just <$> Lens.indexed f () a
  each _ Nothing  = pure Nothing
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' 'Int' ('Seq' a) ('Seq' b) a b@
instance Applicative f => Each f (Seq a) (Seq b) a b
  {- default each -}

-- | @'each' :: 'IndexedTraversal' ['Int'] ('Tree' a) ('Tree' b) a b@
instance Applicative f => Each f (Tree a) (Tree b) a b where
  each pafb = go (BD 0 [] 0 []) where
    go dq (Node a as) = Node <$> Lens.indexed pafb (Foldable.toList dq) a <*> itraverse (\i n -> go (Lens.snoc dq i) n) as
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' 'Int' ('Vector.Vector' a) ('Vector.Vector' b) a b@
instance Applicative f => Each f (Vector.Vector a) (Vector.Vector b) a b
  {- default each -}

-- | @'each' :: ('Prim' a, 'Prim' b) => 'IndexedTraversal' 'Int' ('Prim.Vector' a) ('Prim.Vector' b) a b@
instance (Applicative f, Prim a, Prim b) => Each f (Prim.Vector a) (Prim.Vector b) a b where
  each f v = Prim.fromListN (Prim.length v) <$> traversed (Indexed f') (Prim.toList v)
    where f' = Lens.indexed f
  {-# INLINE each #-}

-- | @'each' :: ('Storable' a, 'Storable' b) => 'IndexedTraversal' 'Int' ('Storable.Vector' a) ('Storable.Vector' b) a b@
instance (Applicative f, Storable a, Storable b) => Each f (Storable.Vector a) (Storable.Vector b) a b where
  each f v = Storable.fromListN (Storable.length v) <$> traversed (Indexed f') (Storable.toList v)
    where f' = Lens.indexed f
  {-# INLINE each #-}

-- | @'each' :: ('Unbox' a, 'Unbox' b) => 'IndexedTraversal' 'Int' ('Unboxed.Vector' a) ('Unboxed.Vector' b) a b@
instance (Applicative f, Unbox a, Unbox b) => Each f (Unboxed.Vector a) (Unboxed.Vector b) a b where
  each f v = Unboxed.fromListN (Unboxed.length v) <$> traversed (Indexed f') (Unboxed.toList v)
    where f' = Lens.indexed f
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' 'Int' 'StrictT.Text' 'StrictT.Text' 'Char' 'Char'@
instance Applicative f =>  Each f StrictT.Text StrictT.Text Char Char where
  each = iso StrictT.unpack StrictT.pack . traversed
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' 'Int64' 'LazyT.Text' 'LazyT.Text' 'Char' 'Char'@
instance Applicative f => Each f LazyT.Text LazyT.Text Char Char where
  each = iso LazyT.unpack LazyT.pack . traversed64
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' 'Int' 'StrictB.ByteString' 'StrictB.ByteString' 'Word8' 'Word8'@
instance Applicative f => Each f StrictB.ByteString StrictB.ByteString Word8 Word8 where
  each = iso StrictB.unpack StrictB.pack . traversed
  {-# INLINE each #-}

-- | @'each' :: 'IndexedTraversal' 'Int64' 'LazyB.ByteString' 'LazyB.ByteString' 'Word8' 'Word8'@
instance Applicative f => Each f LazyB.ByteString LazyB.ByteString Word8 Word8 where
  each = iso LazyB.unpack LazyB.pack . traversed64
  {-# INLINE each #-}

-- | @'each' :: 'Ix' i => 'IndexedTraversal' i ('Array' i a) ('Array' i b) a b@
instance (Applicative f, Ix i) => Each f (Array i a) (Array i b) a b where
  each f arr = array (bounds arr) <$> traverse (\(i,a) -> (,) i <$> Lens.indexed f i a) (IArray.assocs arr)
  {-# INLINE each #-}

-- | @'each' :: ('Ix' i, 'IArray' 'UArray' a, 'IArray' 'UArray' b) => 'IndexedTraversal' i ('Array' i a) ('Array' i b) a b@
instance (Applicative f, Ix i, IArray UArray a, IArray UArray b) => Each f (UArray i a) (UArray i b) a b where
  each f arr = array (bounds arr) <$> traverse (\(i,a) -> (,) i <$> Lens.indexed f i a) (IArray.assocs arr)
  {-# INLINE each #-}

-- | @'each' :: 'Control.Lens.IndexedSetter' i (i -> a) (i -> b) a b@
instance Settable f => Each f (i -> a) (i -> b) a b where
  each f g = pure (\i -> untaintedDot (Lens.indexed f i) (g i))
  {-# INLINE each #-}


module Main (main) where

import Data.Complex
import Data.HDF5
import Data.Proxy
import Data.Text
import Data.Word
import qualified Data.Vector.Storable as V
import System.IO (IOMode (..))

main :: IO ()
main = do
  -- disableDiagOutput
  withFile "workspace_30.h5" ReadWriteMode $ \handle -> do
    byName handle "/47" . matchM @Group $ \root ->
      do
        -- print =<< getSize handle
        -- print =<< getSize root
        -- print =<< getDims root "c"
        -- () <- byName root "c" $ \(Some x@(Dataset _)) ->
        --   print =<< readDataset' (Proxy @Double) x
        writeAttribute root "random" (12.3 :: Float)
        print =<< readAttribute @Double root "energy"
        print =<< readAttribute @Float root "random"
        writeAttribute root "random" (15.8 :: Float)
        print =<< readAttribute @Float root "random"
        writeAttribute root "cheers" (17.0 :+ (-1.8) :: Complex Float)
        print =<< readAttribute @(Complex Float) root "cheers"
        writeAttribute @Text root "ping pong" "This is awesome!"
        print =<< readAttribute @Text root "ping pong"

        print =<< exists root "c'"
        print =<< exists root "/48/tleft"

        () <- byName root "c" $ \case
          (Some x@(Dataset _)) -> do
            print =<< getDatasetDims x
          _ -> error "expected c to be a dataset"

        () <- byName root "c" $ \case
          (Some x@(Dataset _)) -> do
            blob <- readDataset @Double x
            writeDataset root "p" blob
          _ -> error "expected c to be a dataset"

        writeDataset root "uint64_dataset" $ V.fromList [1 :: Word64, 2, 3]
        return ()

  putStrLn "Hello world!"

-- case hasStorable (Proxy @Float) of
--   Just _ -> print "cheers!"
--   Nothing -> print "fuck!"

-- print =<< readAttribute' (Proxy @Double) root "energy"

-- print =<< h5Name handle
-- print =<< h5FindDataset handle "c"
-- g' <- h5OpenByIndex handle 0
-- print =<< h5Size g'
-- print =<< h5Name g'
-- print =<< h5FindDataset g' "tleft"
-- print =<< h5GetNumDims g' "tleft"
-- print =<< h5GetDatasetInfo g' "c"
-- readTypedDataset (Proxy @Double) g' "c" >>= \case
--   (H5ConcreteBlob v dims) -> print dims
-- readDataset g' "c" >>= \case
--   (H5DoubleBlob v dims) -> print dims
-- _ <- h5FoldlM (\i h -> h5Name h >>= \name -> print (i, name) >> return (i + 1)) 0 g'
-- closeHandle g'


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}

module Elescore.Api.Data
  ( DataApi
  , dataServer
  ) where

import           ClassyPrelude                         hiding (Handler)
import           Data.DateTime                         (DateTime)
import qualified Data.Map                              as M
import qualified Data.Text                             as T
import           Servant
import           Servant.Pagination

import           Elescore.Api.Types
import           Elescore.IdTypes
import           Elescore.Projection.DisruptionsPerDay (DisruptionsPerDay)
import           Statistics.IQR                        (average)

type AppliedRanges ranges a = Ranges ranges a
type PagReqHeaders ranges a = Header "Range" (AppliedRanges ranges a)
type PagRespHeaders ranges a = Header "Total-Count" Int ': PageHeaders ranges a

type DisruptionFields = '["id", "occurredOn", "updatedOn", "resolvedOn", "duration"]
type DisruptionRanges = AppliedRanges DisruptionFields UIDisruption
type DisruptionPagReqHeaders = PagReqHeaders DisruptionFields UIDisruption
type DisruptionPagRespHeaders = PagRespHeaders DisruptionFields UIDisruption

type DataApi =
       "stats" :> StatsApi
  :<|> "disruptions" :> DisruptionApi
  :<|> "objects" :> ObjectsApi
  :<|> "facilities" :> FacilitiesApi

type StatsApi =
       Get '[JSON] UIOverallStats
  :<|> "averageDisruptionsPerDay" :> Get '[JSON] [UIDisruptionPerDay]

type DisruptionApi =
       DisruptionPagReqHeaders :> GetPartialContent '[JSON] (Headers DisruptionPagRespHeaders [UIDisruption])
  :<|> "markers" :> Get '[JSON] [UIMapMarker]
  :<|> "active" :> DisruptionPagReqHeaders :> GetPartialContent '[JSON] (Headers DisruptionPagRespHeaders [UIDisruption])

type ObjectsApi = QueryParam "search" Text :> Get '[JSON] [UIObjectSearchResult]

type FacilitiesApi =
       Capture "facilityId" FacilityId :> Get '[JSON] UIFacilityDetails
  :<|> DisruptionPagReqHeaders :> Capture "facilityId" FacilityId :> "disruptions" :>  GetPartialContent '[JSON] (Headers DisruptionPagRespHeaders [UIDisruption])

dataServer ::
     IORef UIOverallStats
  -> IORef DisruptionsPerDay
  -> IORef [UIDisruption]
  -> IORef [UIMapMarker]
  -> IORef [UIObjectSearchResult]
  -> IORef (Map FacilityId UIFacilityDetails)
  -> Server DataApi
dataServer stats disruptionsPerDay disruptions marker searchResults facilities =
  statsApi stats disruptionsPerDay :<|> disruptionsApi disruptions marker :<|> objectsApi searchResults :<|> facilitiesApi facilities disruptions

statsApi :: IORef UIOverallStats -> IORef DisruptionsPerDay -> Server StatsApi
statsApi statsRef disPerDayRef =
  readIORef statsRef :<|> disPerDayHandler
  where
    disPerDayHandler = reverse . fmap mkUIDisruptionPerDay . take 31 . drop 1 . reverse . M.toList . fmap average .  snd <$> readIORef disPerDayRef

disruptionsApi :: IORef [UIDisruption] -> IORef [UIMapMarker] -> Server DisruptionApi
disruptionsApi disRef markerRef =
  allDisruptionsHandler :<|> readIORef markerRef :<|> activeDisruptionsHandler
  where
    allDisruptionsHandler ranges = readIORef disRef >>= disruptionsHandler ranges
    activeDisruptionsHandler ranges = readIORef disRef >>= disruptionsHandler ranges . filter (isNothing . uidResolvedOn)

objectsApi :: IORef [UIObjectSearchResult] -> Server ObjectsApi
objectsApi ref maybeSearch =
  if length search < 3
  then return []
  else filter (nameMatches . uioName . uiosrObject) <$> liftIO (readIORef ref)

  where
    search = T.toLower $ fromMaybe "" maybeSearch
    nameMatches = T.isInfixOf search . T.toLower

facilitiesApi :: IORef (Map FacilityId UIFacilityDetails) -> IORef [UIDisruption] -> Server FacilitiesApi
facilitiesApi facilities disruptions =
  facilityDetailsHandler :<|> facilityDisruptionsHandler
  where
    facilityDetailsHandler facilityId =
      maybe (throwError err404) return . lookup facilityId =<< readIORef facilities
    facilityDisruptionsHandler ranges facilityId =
      readIORef disruptions >>= disruptionsHandler ranges . filter ((==) facilityId . uidFacilityId)

disruptionsHandler :: Maybe DisruptionRanges -> [UIDisruption] -> Handler (Headers DisruptionPagRespHeaders [UIDisruption])
disruptionsHandler r ds =
  addHeader (length ds) <$>
    fromMaybe (returnIdRange ds defaultRange)
        (fmap (returnIdRange ds) (r >>= extractRange)
    <|> fmap (returnOccurredOnRange ds) (r >>= extractRange)
    <|> fmap (returnUpdatedOnRange ds) (r >>= extractRange)
    <|> fmap (returnResolvedOnRange ds) (r >>= extractRange)
    <|> fmap (returnDurationRange ds) (r >>= extractRange))

  where
    returnIdRange xs (range :: Range "id" Int) =
      returnRange range (applyRange range xs)

    returnOccurredOnRange xs (range :: Range "occurredOn" DateTime) =
      returnRange range (applyRange range xs)

    returnUpdatedOnRange xs (range :: Range "updatedOn" (Maybe DateTime)) =
      returnRange range (applyRange range xs)

    returnResolvedOnRange xs (range :: Range "resolvedOn" (Maybe DateTime)) =
      returnRange range (applyRange range xs)

    returnDurationRange xs (range :: Range "duration" Integer) =
      returnRange range (applyRange range xs)

defaultRange :: Range "id" Int
defaultRange = getDefaultRange (Proxy @UIDisruption)


{-
  this file is taken from the hommage library
-}
{-# LANGUAGE FlexibleContexts #-}
module WavFile
 (
 -- * Binary Files
   writeDataFile
 , readDataFile
 , openDataFile

 -- * WAV-Files
 , writeWavFile
 , writeWavFileMono
 , writeWavFileStereo
 , readWavFile
 , openWavFile

 , writeWavFiles

 -- * Cast between Int16 and Double representation of WAV-data
 , wavInt16ToDouble
 , wavDoubleToInt16

 -- * Low-Level implemetation
 -- ** Arrays and Files
 , readArrayFromFile
 , writeArrayToFileWithHeader
 , writeArrayToFile

 -- ** Single Stream
 , openSingleInputFile
 , openSingleOutputFile

 , openSingleInputWavFile
 , openSingleOutputWavFileMono
 , openSingleOutputWavFileStereo

 -- ** Buffered Stream
 , openInputFile
 , openOutputFile
 , openOutputFileWithHeader
 , openInputWavFile
 , openOutputWavFileMono
 , openOutputWavFileStereo

 -- ** Header Stuff and others
 , HeaderFun
 , HeaderSize
 , noHeader

 , wavHeaderFunMono
 , wavHeaderFunStereo
 , wavHeaderSize

 , initWriteWavHeaderMono
 , initWriteWavHeaderStereo
 , initReadWavHeader
 , closeWriteWavHeader
 , encode
 , decode
 , encodeWavLengt
 , initWavHeaderMono
 , initWavHeaderStereo
 , sizeOfArrayElements
 , inferSizeOfArrayElements
 , inferSizeOfArrayElements'
 )
 where

import System.IO
import Control.Monad
import Control.Monad.Fix
import Data.Char
import Data.Int
import GHC.IO.Handle
import GHC.IO
import GHC.Base hiding (mapM)
import Data.Complex

import Foreign.Storable
import Data.Array.Storable
import Data.Array.IArray
import Data.IORef

import System.IO.Unsafe
import GHC.Weak

---------------------------------------------------------------------------------------------------
wavInt16ToDouble :: Int16 -> Double
wavInt16ToDouble i = (fromIntegral i) / 32768.0

wavDoubleToInt16 :: Double -> Int16
wavDoubleToInt16 d = round (d * 32767.0)
---------------------------------------------------------------------------------------------------
readDataFile :: Storable a => FilePath -> IO [a]
readDataFile fp = init
 where
  init = do (c,r,_) <- openSingleInputFile 10000 0 fp
            mkWeak r () (Just c)
            loop r
  loop r = r >>= maybe (return []) (\a -> return (a : unsafePerformIO (loop r)))

openDataFile :: Storable a => FilePath -> [a]
openDataFile fp = unsafePerformIO $ readDataFile fp

writeDataFile :: Storable a => FilePath -> [a] -> IO ()
writeDataFile fp xs = do
 (c,w) <- openSingleOutputFile 4000 fp
 let loop (x:r) = w x >> loop r
     loop []    = c
 loop xs

---------------------------------------------------------------------------------------------------
readWavFile :: FilePath -> IO (Either [Int16] [(Int16,Int16)])
readWavFile fp = {- putStrLn ("opening " ++ fp) >> -} init
 where
  init = do (c,r,_) <- openSingleInputWavFile 100000 fp
            either (\rm -> do mkWeak rm () (Just c) --(putStrLn ("close " ++ fp) >> c))
                              fmap Left $ loop rm)
                   (\rs -> do mkWeak rs () (Just c) --(putStrLn ("close " ++ fp) >> c))
                              fmap Right $ loop rs)
                   r
  loop r = r >>= maybe (return []) (\a -> return (a : unsafePerformIO (loop r)))

openWavFile :: FilePath -> Either [Int16] [(Int16,Int16)]
openWavFile fp = unsafePerformIO $ readWavFile fp

writeWavFile :: FilePath -> Either [Int16] [(Int16, Int16)] -> IO ()
writeWavFile fp (Left w)  = writeWavFileMono fp w
writeWavFile fp (Right w) = writeWavFileStereo fp w

writeWavFileMono :: FilePath -> [Int16] -> IO ()
writeWavFileMono fp xs = do
 (c,w) <- openSingleOutputWavFileMono 100000 fp
 let loop (x:r) = w x >> loop r
     loop []    = c
 loop xs

writeWavFileStereo :: FilePath -> [(Int16,Int16)] -> IO ()
writeWavFileStereo fp xs = do
 (c,w) <- openSingleOutputWavFileStereo 100000 fp
 let loop (x:r) = w x >> loop r
     loop []    = c
 loop xs

writeWavFiles :: FilePath -> String  -> [Either [Int16] [(Int16, Int16)]] -> IO ()
writeWavFiles fp fn sns = do
  let init count (Left x : xs) = do (c,w) <- openSingleOutputWavFileMono 100000 (fp ++ fn ++ "_" ++ show count ++ ".wav")
                                    r <- init (count+1) xs
                                    return (Left (c,w,x) : r)
      init count (Right x : xs) = do (c,w) <- openSingleOutputWavFileStereo 100000 (fp ++ fn ++ "_" ++ show count ++ ".wav")
                                     r <- init (count+1) xs
                                     return (Right (c,w,x) : r)
      init _ _ = return []
  ss <- init 0 sns
  let step (Left (c,w,x:xs) : r) = do w x
                                      r <- step r
                                      return (Left (c,w,xs) : r)
      step (Left (c,w,[]) : r)   = c >> step r
      step (Right (c,w,x:xs) : r) = do w x
                                       r <- step r
                                       return (Right (c,w,xs) : r)
      step (Right (c,w,[]) : r)   = c >> step r
      step []                     = return []

      loop [] = return ()
      loop xs = step xs >>= loop
  loop ss
---------------------------------------------------------------------------------------------------
-- Header

-- | The first Action is applied after opening the file. Then the data
--   bytes are written. The second Action is then called with the number of
--   bytes of the data. It must close the handle.
type HeaderFun = (Handle -> IO (), Handle -> Int -> IO ())
type HeaderSize = Int

noHeader :: HeaderFun
noHeader = (const $ return (), \h _ -> return ())

-------------------------------------------------
-- WavHeader

-- | The first Action writes the Header of the wav-file.
--   Then the wav-data is written. The second action moves
--   the Handle to the positions where the lenght of the wav-file
--   are encoded (in the header) and writes the right number, which
--   is unknown before all data is written.
--   Afterwards it closes the file.
wavHeaderFunMono :: HeaderFun
wavHeaderFunMono = (initWriteWavHeaderMono, closeWriteWavHeader)

wavHeaderFunStereo :: HeaderFun
wavHeaderFunStereo = (initWriteWavHeaderStereo, closeWriteWavHeader)

wavHeaderSize :: HeaderSize
wavHeaderSize = 44
-------------------------------------------------
initWriteWavHeaderMono :: Handle -> IO ()
initWriteWavHeaderMono h = mapM_ (hPutChar h) initWavHeaderMono

initWriteWavHeaderStereo :: Handle -> IO ()
initWriteWavHeaderStereo h = mapM_ (hPutChar h) initWavHeaderStereo

initReadWavHeader :: Handle -> IO (Maybe (Bool, Int, Int, Int))
initReadWavHeader handle =
 let check [] = return True
     check (x:xs) = hGetChar handle >>= \y -> if x /= ord y then return False else check xs
     m >>? n = m >>= \b -> if b then n else return Nothing
 in hSeek handle AbsoluteSeek 0 >>
    check [82,73,70,70] >>?
    (replicateM 4 (hGetChar handle) >>= \l1 ->
    check [87, 65, 86, 69, 102,109,116,32, 16, 0,  0,  0, 1,  0] >>?
    (hGetChar handle >>= \s1 ->
    check [0, 68, 172, 0, 0] >>?
    (replicateM 4 (hGetChar handle) >>= \x1 ->
    (hGetChar handle >>= \s2 ->
    check [0,  16, 0, 100, 97,116,97] >>?
    (replicateM 4 (hGetChar handle) >>= \l2 ->
    let b | ord s1 == 1 && ord s2 == 2 = False
          | ord s1 == 2 && ord s2 == 4 = True
          | otherwise          = error ("invalid wav file " ++ show s1 ++ ", " ++ show s2)
    in return $ Just (b, decode l1, decode l2, decode x1) ) ) ) ) )

closeWriteWavHeader :: Handle -> Int -> IO ()
closeWriteWavHeader h len = do
 let (sa,sb) = encodeWavLengt len
 hSeek h AbsoluteSeek 4
 mapM (hPutChar h) sa
 hSeek h AbsoluteSeek 40
 mapM (hPutChar h) sb
 hClose h

encode :: Int -> String
encode a = map chr
 [ mod a 256
 , mod (div a 256) 256
 , mod (div a 65536) 256
 , mod (div a 16777216) 256
 ]

decode :: String -> Int
decode [a0,a1,a2,a3] = sum
 [ord a0, 256 * ord a1, 65536 * ord a2, 16777216 * ord a3]

encodeWavLengt :: Int -> (String, String)
encodeWavLengt len = (encode (len + 8), encode len)

initWavHeaderMono :: [Char]
initWavHeaderMono =
 map chr
 [82,73,70,70,
  0,0,0,0,
  87, 65, 86, 69,
  102,109,116,32,
  16, 0,  0,  0,
  1,  0,  1,  0,
  68, 172,0,  0,
  136,88, 1,  0,
  2,  0,  16, 0,
  100, 97,116,97,
  0,0,0,0]

initWavHeaderStereo :: [Char]
initWavHeaderStereo =
 map chr
 [82,73,70,70,
  0,0,0,0,
  87, 65, 86, 69,
  102,109,116,32,
  16, 0,  0,  0,
  1,  0,  2,  0,
  68, 172,0,  0,
  16,177, 2,  0,
  4,  0,  16, 0,
  100, 97,116,97,
  0,0,0,0]

---------------------------------------------------------------------------------------------------
sizeOfArrayElements :: Storable a => StorableArray Int a -> Int
sizeOfArrayElements x = sizeOf (foo x)
 where
  foo :: StorableArray Int a -> a
  foo x = undefined

inferSizeOfArrayElements :: Storable a => (StorableArray Int a -> IO Int) -> StorableArray Int a -> Int
inferSizeOfArrayElements _ arr = sizeOfArrayElements arr

inferSizeOfArrayElements' :: Storable a => (StorableArray Int a -> Int -> IO ()) -> StorableArray Int a -> Int
inferSizeOfArrayElements' _ arr = sizeOfArrayElements arr

inferSizeOfArrayElements'' :: Storable a => (StorableArray Int a -> Int -> IO Int) -> StorableArray Int a -> Int
inferSizeOfArrayElements'' _ arr = sizeOfArrayElements arr

---------------------------------------------------------------------------------------------------
-- Array / File
---------------------------------------------------------------------------------------------------
readArrayFromFile :: Storable a => HeaderSize -> FilePath -> IO (StorableArray Int a, Int)
readArrayFromFile offset filepath = do
 handle <- openBinaryFile filepath ReadMode
 hSeek handle AbsoluteSeek $ fromIntegral offset
 bytes <- hFileSize handle
 let bytesize = fromIntegral bytes - offset
 arr <- mfix $ \arr -> do let bufsize = div bytesize (sizeOfArrayElements arr)
                          newArray_ (0, bufsize-1)
 withStorableArray arr $ \ptr -> hGetBuf handle ptr bytesize
 hClose handle
 return (arr, div bytesize (sizeOfArrayElements arr))

writeArrayToFileWithHeader :: (MArray StorableArray a IO, Storable a) => HeaderFun -> FilePath -> StorableArray Int a -> IO ()
writeArrayToFileWithHeader (initH, closeH) filepath arr = do
 handle <- openBinaryFile filepath WriteMode
 initH handle
 (0, n) <- getBounds arr
 let elemsize = sizeOfArrayElements arr
     bytesize = (1+n) * elemsize
 withStorableArray arr (\ptr -> hPutBuf handle ptr bytesize)
 closeH handle bytesize

writeArrayToFile :: (MArray StorableArray a IO, Storable a) => FilePath -> StorableArray Int a -> IO ()
writeArrayToFile filepath arr = do
 handle <- openBinaryFile filepath WriteMode
 (0, n) <- getBounds arr
 let elemsize = sizeOfArrayElements arr
     bytesize = (1+n) * elemsize
 withStorableArray arr (\ptr -> hPutBuf handle ptr bytesize)
 hClose handle

{-
writeArrayToFileWithHeader :: Storable a => HeaderFun -> FilePath -> StorableArray Int a -> IO ()
writeArrayToFileWithHeader (initH, closeH) filepath arr = do
 handle <- openBinaryFile filepath WriteMode
 initH handle
 let (0,n) = bounds arr
     elemsize = sizeOfArrayElements arr
     bytesize = (1+n) * elemsize
 withStorableArray arr (\ptr -> hPutBuf handle ptr bytesize)
 closeH handle bytesize

writeArrayToFile :: Storable a => FilePath -> StorableArray Int a -> IO ()
writeArrayToFile filepath arr = do
 handle <- openBinaryFile filepath WriteMode
 let (0,n) = bounds arr
     elemsize = sizeOfArrayElements arr
     bytesize = (1+n) * elemsize
 withStorableArray arr (\ptr -> hPutBuf handle ptr bytesize)
 hClose handle
-}
---------------------------------------------------------------------------------------------------
-- Single Wav
---------------------------------------------------------------------------------------------------
openSingleInputWavFile :: Int -> FilePath -> IO (IO (), Either (IO (Maybe Int16)) (IO (Maybe (Int16,Int16))), Int)
openSingleInputWavFile n filepath = do
 handle <- openBinaryFile filepath ReadMode
 initReadWavHeader handle >>= maybe (error "unknown wav format") (\(isstereo, len1, len2, x) -> do
  hSeek handle AbsoluteSeek 44
  let stereo = do let bufsize = n*2
                  rpos   <- newIORef bufsize
                  rmax   <- newIORef bufsize
                  buf    <- newArray_ (0, bufsize-1)
                  let elemsize = sizeOfArrayElements buf
                      bytesize = elemsize * bufsize
                      next = readIORef rpos >>= \pos ->
                             readIORef rmax >>= \max ->
                             if pos >= max
                              then if max < bufsize
                                    then return Nothing
                                    else withStorableArray buf (\ptr -> hGetBuf handle ptr bytesize) >>= \bts ->
                                         writeIORef rmax (div bts elemsize) >>
                                         writeIORef rpos 0 >>
                                         next
                              else readArray buf pos >>= \v1 ->
                                   readArray buf (pos+1) >>= \v2 ->
                                   writeIORef rpos (pos + 2) >>
                                   return (Just (v1,v2))
                      close = hClose handle
                  return (close, Right next,div len2 4)
      mono   = do let bufsize = n
                  rpos   <- newIORef bufsize
                  rmax   <- newIORef bufsize
                  buf    <- newArray_ (0, bufsize-1)
                  let elemsize = sizeOfArrayElements buf
                      bytesize = elemsize * bufsize
                      next = readIORef rpos >>= \pos ->
                             readIORef rmax >>= \max ->
                             if pos >= max
                              then if max < bufsize
                                    then return Nothing
                                    else withStorableArray buf (\ptr -> hGetBuf handle ptr bytesize) >>= \bts ->
                                         writeIORef rmax (div bts elemsize) >>
                                         writeIORef rpos 0 >>
                                         next
                              else readArray buf pos >>= \v ->
                                   writeIORef rpos (pos + 1) >>
                                   return (Just v)
                      close = hClose handle
                  return (close, Left next,div len2 2)
  if isstereo then stereo else mono )
---------------------------------------------------------------------------------------------------
openSingleOutputWavFileMono :: Int -> FilePath -> IO (IO (), Int16 -> IO ())
openSingleOutputWavFileMono bufsize filepath = do
 handle <- openBinaryFile filepath WriteMode
 rpos   <- newIORef 0
 buf    <- newArray_ (0, bufsize-1)
 rcnt   <- newIORef 0
 let elemsize = sizeOfArrayElements buf
     bytesize = elemsize * bufsize
     write v = readIORef rpos >>= \pos ->
               if pos >= bufsize
                then withStorableArray buf (\ptr -> hPutBuf handle ptr bytesize) >>
                     modifyIORef rcnt (+bytesize) >>
                     writeIORef rpos 1 >> writeArray buf 0 v
                else writeArray buf pos v >> writeIORef rpos (pos+1)
     close = readIORef rpos >>= \pos ->
             withStorableArray buf (\ptr -> hPutBuf handle ptr (pos*elemsize)) >>
             readIORef rcnt >>= \cnt -> closeWriteWavHeader handle (elemsize*pos+cnt)
 initWriteWavHeaderMono handle
 return (close, write)
---------------------------------------------------------------------------------------------------
openSingleOutputWavFileStereo :: Int -> FilePath -> IO (IO (), (Int16, Int16) -> IO ())
openSingleOutputWavFileStereo bufsizearg filepath = do
 let bufsize | bufsizearg < 2  = 4
             | even bufsizearg = bufsizearg
             | otherwise       = bufsizearg + 1
 handle <- openBinaryFile filepath WriteMode
 rpos   <- newIORef 0
 buf    <- newArray_ (0, bufsize-1)
 rcnt   <- newIORef 0
 let elemsize = sizeOfArrayElements buf
     bytesize = elemsize * bufsize
     write (v1,v2) = readIORef rpos >>= \pos ->
                     if pos >= bufsize
                      then withStorableArray buf (\ptr -> hPutBuf handle ptr bytesize) >>
                           modifyIORef rcnt (+bytesize) >>
                           writeIORef rpos 2 >> writeArray buf 0 v1 >> writeArray buf 1 v2
                      else writeArray buf pos v1 >> writeArray buf (pos+1) v2 >> writeIORef rpos (pos+2)
     close = readIORef rpos >>= \pos ->
             withStorableArray buf (\ptr -> hPutBuf handle ptr (pos*elemsize)) >>
             readIORef rcnt >>= \cnt -> closeWriteWavHeader handle (elemsize*pos+cnt)
 initWriteWavHeaderStereo handle
 return (close, write)
---------------------------------------------------------------------------------------------------
-- Single Storable
---------------------------------------------------------------------------------------------------
openSingleInputFile :: Storable a => Int -> HeaderSize -> FilePath -> IO (IO (), IO (Maybe a), Int)
openSingleInputFile bufsize offset filepath = do
 handle <- openBinaryFile filepath ReadMode
 bytes <- hFileSize handle
 hSeek handle AbsoluteSeek $ fromIntegral offset
 rpos   <- newIORef bufsize
 rmax   <- newIORef bufsize
 buf    <- newArray_ (0, bufsize-1)
 let elemsize = sizeOfArrayElements buf
     bytesize = elemsize * bufsize
     next = readIORef rpos >>= \pos ->
            readIORef rmax >>= \max ->
            if pos >= max
             then if max < bufsize
                   then return Nothing
                   else withStorableArray buf (\ptr -> hGetBuf handle ptr bytesize) >>= \bts ->
                        writeIORef rmax (div bts elemsize) >>
                        writeIORef rpos 0 >>
                        next
             else readArray buf pos >>= \v ->
                  writeIORef rpos (pos + 1) >>
                  return (Just v)
     close = hClose handle
 return (close, next,div (fromIntegral bytes - offset) elemsize)
---------------------------------------------------------------------------------------------------
openSingleOutputFile :: Storable a => Int -> FilePath -> IO (IO (), a -> IO ())
openSingleOutputFile bufsize filepath = do
 handle <- openBinaryFile filepath WriteMode
 rpos   <- newIORef 0
 buf    <- newArray_ (0, bufsize-1)
 let elemsize = sizeOfArrayElements buf
     bytesize = elemsize * bufsize
     write v = readIORef rpos >>= \pos ->
               if pos >= bufsize
                then withStorableArray buf (\ptr -> hPutBuf handle ptr bytesize) >>
                     writeIORef rpos 1 >> writeArray buf 0 v
                else writeArray buf pos v >> writeIORef rpos (pos+1)
     close = readIORef rpos >>= \pos ->
             withStorableArray buf (\ptr -> hPutBuf handle ptr (pos*elemsize)) >>
             hClose handle
 return (close, write)
---------------------------------------------------------------------------------------------------
-- Buffered Storable
---------------------------------------------------------------------------------------------------
openInputFile :: Storable a => HeaderSize -> FilePath -> Int -> IO (IO (), StorableArray Int a -> Int -> IO Int, Int)
openInputFile offset filepath arrsize = do
 handle <- openBinaryFile filepath ReadMode
 hSeek handle AbsoluteSeek $ fromIntegral offset
 bytes <- hFileSize handle
 let nr_bytes  = fromIntegral bytes - offset
     elemsize  = inferSizeOfArrayElements'' read undefined
     read arr k = withStorableArray arr $ \ptr ->
                  hGetBuf handle ptr (k * elemsize) >>= return . flip div elemsize
     close     = hClose handle
 return $ (close, read, div nr_bytes elemsize)

openOutputFile :: Storable a => FilePath -> IO (IO (), StorableArray Int a -> Int -> IO ())
openOutputFile filepath = do
 handle <- openBinaryFile filepath WriteMode
 let elemsize = inferSizeOfArrayElements' write undefined
     write arr k = withStorableArray arr (\ptr -> hPutBuf handle ptr (k*elemsize))
     close = hClose handle
 return $ (close, write)

openOutputFileWithHeader :: Storable a => HeaderFun -> FilePath -> IO (IO (), StorableArray Int a -> Int -> IO ())
openOutputFileWithHeader (initH, closeH) filepath = do
 handle <- openBinaryFile filepath WriteMode
 rlen   <- newIORef 0
 initH handle
 let elemsize = inferSizeOfArrayElements' write undefined
     write arr k = modifyIORef rlen (+k) >> withStorableArray arr (\ptr -> hPutBuf handle ptr (k*elemsize))
     close =  readIORef rlen >>= \l -> closeH handle (elemsize*l)
 return $ (close, write)
---------------------------------------------------------------------------------------------------
-- Buffered Wav
---------------------------------------------------------------------------------------------------

openInputWavFile :: FilePath -> IO (IO (), StorableArray Int Int16 -> Int -> IO (), Int, Bool)
openInputWavFile filepath = do
 handle <- openBinaryFile filepath ReadMode
 initReadWavHeader handle >>= maybe (error "unknown wav format") (\(isstereo, len1, len2, x) -> do
  hSeek handle AbsoluteSeek 44
  let elemsize  = inferSizeOfArrayElements' read undefined
      read arr k = withStorableArray arr $ \ptr ->
                   hGetBuf handle ptr (k * elemsize) >> return ()
      close     = hClose handle
  return $ (close, read, div len2 2, isstereo) )

openOutputWavFileMono :: FilePath -> IO (IO (), StorableArray Int Int16 -> Int -> IO ())
openOutputWavFileMono filepath = do
 handle <- openBinaryFile filepath WriteMode
 rlen   <- newIORef 0
 let elemsize = inferSizeOfArrayElements' write undefined
     write arr k = withStorableArray arr (\ptr -> hPutBuf handle ptr (k*elemsize)) >> modifyIORef rlen (+k)
     close = readIORef rlen >>= \l -> closeWriteWavHeader handle (elemsize*l)
 initWriteWavHeaderMono handle
 return $ (close, write)

openOutputWavFileStereo :: FilePath -> IO (IO (), StorableArray Int Int16 -> Int -> IO ())
openOutputWavFileStereo filepath = do
 handle <- openBinaryFile filepath WriteMode
 rlen   <- newIORef 0
 let elemsize = inferSizeOfArrayElements' write undefined
     write arr k = withStorableArray arr (\ptr -> hPutBuf handle ptr (k*elemsize)) >> modifyIORef rlen (+k)
     close = readIORef rlen >>= \l -> closeWriteWavHeader handle (elemsize*l)
 initWriteWavHeaderStereo handle
 return $ (close, write)
---------------------------------------------------------------------------------------------------


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Main (main) where

import Control.Monad.IO.Class (liftIO)
import Data.Complex (Complex)
import Data.Foldable (toList)
import Data.List (nub,
                  sort)
import Text.PrettyPrint.Mainland
import Text.PrettyPrint.Mainland.Class
import Text.Read (readMaybe)

import Spiral
import Spiral.Backend.C
import Spiral.Config
import Spiral.Driver
import Spiral.Exp
import Spiral.Search.OpCount
import Spiral.OpCount
import Spiral.SPL
import Spiral.SPL.Run
import Spiral.Util.Uniq

parseSizesArg :: Monad m => String -> m [Int]
parseSizesArg = fmap concat . mapM parseRange . splitOn (== ',')

parseRange :: Monad m => String -> m [Int]
parseRange s =
    case splitOn (== '-') s of
      [x]   -> do start <- parseInt x
                  return [start]
      [x,y] -> do start <- parseInt x
                  end   <- parseInt y
                  return [start..end]
      _     -> fail $ "Cannot parse range: " ++ s

parseInt :: Monad m => String -> m Int
parseInt s = case readMaybe s of
                 Nothing -> fail $ "Cannot parse integer: " ++ s
                 Just n  -> return n

splitOn :: (Char -> Bool) -> String -> [String]
splitOn p s = case dropWhile p s of
                  "" -> []
                  s' -> w : splitOn p s''
                        where (w, s'') = break p s'

main :: IO ()
main =
    defaultMainWith mempty $ \args -> do
    sizes <- nub . sort . concat <$> mapM parseSizesArg args
    useComplexType <- asksConfig $ testDynFlag UseComplex
    liftIO $ putDocLn $ text "\"size\",\"totalops\",\"mulops\",\"addops\""
    mapM_ (search useComplexType) sizes
  where
    search :: Bool -> Int -> Spiral ()
    search True  n = searchOpCount (DFT n :: SPL (Exp (Complex Double))) >>= go n
    search False n = searchOpCount (Re (DFT n) :: SPL (Exp Double)) >>= go n

    go :: (Typed a, Num (Exp a)) => Int -> SPL (Exp a) -> Spiral ()
    go n e = do
        -- pprint e
        prog <- toProgram ("hspiral_dft_" ++ show n) e
        -- pprint prog
        ops <- countProgramOps prog
        resetUnique
        defs <- evalCg $ cgProgram prog
        outp <- asksConfig output
        case outp of
          Nothing -> return ()
          Just{}  -> writeOutput (toList defs)
        liftIO $ putDocLn $ folddoc (\x y -> x <> comma <> y) [ppr n, ppr (allOps ops), ppr (mulOps ops), ppr (addOps ops)]


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ConstraintKinds #-}

module Language.Grappa.Interp.ProbFun where

import Data.Vector (Vector)
import qualified Data.Vector as V
import qualified Numeric.Log as Log
import qualified Numeric.LinearAlgebra as M
import Data.Functor.Compose
import Data.Aeson

import Language.Grappa.Distribution
import Language.Grappa.Interp
import Language.Grappa.GrappaInternals
import Language.Grappa.Interp.PVIE
import GHC.Exts (IsList(..))

import Debug.Trace

-- | Type tag for the representation that views @'Dist' a@ as a function from
-- @a@ to 'Prob'
data ProbFunRepr

type family ProbFunReprF a where
  ProbFunReprF (a -> b) = (ProbFunReprF a -> ProbFunReprF b)
  ProbFunReprF (Dist a) = ProbFunReprF a -> Prob
  ProbFunReprF (ADT adt) = adt (GExpr ProbFunRepr) (ADT adt)
  ProbFunReprF (Vector a) = Vector (ProbFunReprF a)
  ProbFunReprF RVector = RVector
  ProbFunReprF RMatrix = RMatrix
  ProbFunReprF ProbVector = ProbVector
  ProbFunReprF ProbMatrix = ProbMatrix
  ProbFunReprF (VIDist a) = (VIDistFamExpr (ProbFunReprF a))
  ProbFunReprF VISize = VIDim
  ProbFunReprF a = a

instance ValidExprRepr ProbFunRepr where
  type GExprRepr ProbFunRepr a = ProbFunReprF a
  interp__'bottom = error "ProbFunRepr: unexpected bottom!"
  interp__'injTuple tup = GExpr tup
  interp__'projTuple (GExpr tup) k = k tup
  interp__'app (GExpr f) (GExpr x) = GExpr (f x)
  interp__'lam f = GExpr (unGExpr . f . GExpr)
  interp__'fix f = f (interp__'fix f)

instance ValidRepr ProbFunRepr where
  type GVExprRepr ProbFunRepr a = ProbFunReprF a
  type GStmtRepr ProbFunRepr a = Prob

  interp__'projTupleStmt (GExpr tup) k = k tup

  interp__'vInjTuple tup = GVExpr $ mapADT (GExpr . unGVExpr) tup
  interp__'vProjTuple (GVExpr ve) k = k $ mapADT (GVExpr . unGExpr) ve

  interp__'vwild _ = error "ProbFunRepr: wildcard variables not allowed!"
  interp__'vlift (GExpr e) k = k (GVExpr e)

  interp__'return (GExpr _) = GStmt 1
  interp__'let rhs body = body rhs
  interp__'sample (GExpr d) (GVExpr v) k =
    -- FIXME: evaluate these two sides in parallel!
    GStmt $ d v * unGStmt (k $ GExpr v)

  interp__'mkDist f = GExpr (\ dv -> unGStmt $ f $ GVExpr dv)

instance TraversableADT adt =>
         Interp__ADT__Expr ProbFunRepr adt where
  interp__'injADT adt = GExpr adt
  interp__'projADT (GExpr adt) k = k adt
  interp__'projMatchADT (GExpr adt) _ matcher k_succ k_fail =
    if applyCtorMatcher matcher adt then k_succ adt else k_fail

instance TraversableADT adt => Interp__ADT ProbFunRepr adt where
  interp__'vInjADT adt = GVExpr (mapADT (GExpr . unGVExpr) adt)
  interp__'vProjMatchADT (GVExpr adt) _ matcher k_succ k_fail =
    if applyCtorMatcher matcher adt then
      k_succ (mapADT (GVExpr . unGExpr) adt)
    else k_fail

instance EmbedRepr ProbFunRepr Bool where
  embedRepr = GExpr
instance EmbedRepr ProbFunRepr Int where
  embedRepr = GExpr
instance EmbedRepr ProbFunRepr R where
  embedRepr = GExpr
instance EmbedRepr ProbFunRepr Prob where
  embedRepr = GExpr

instance IsList (ListF a (GExpr ProbFunRepr) (ADT (ListF a))) where
  type Item (ListF a (GExpr ProbFunRepr) (ADT (ListF a))) = ProbFunReprF a
  fromList [] = Nil
  fromList (x:xs) = (Cons (GExpr x) (GExpr (fromList xs)))
  toList Nil = []
  toList (Cons (GExpr x) (GExpr xs)) = x : toList xs

----------------------------------------------------------------------
-- Boolean and comparison operations
----------------------------------------------------------------------

instance Interp__'ifThenElse ProbFunRepr where
  interp__'ifThenElse (GExpr c) t e = if c then t else e

instance Interp__'vmatchSwitch ProbFunRepr where
  interp__'vmatchSwitch (GExpr i) stmts = stmts !! i

instance Interp__not ProbFunRepr where
  interp__not = GExpr not

instance Interp__'amp'amp ProbFunRepr where
  interp__'amp'amp = GExpr (&&)

instance Interp__'bar'bar ProbFunRepr where
  interp__'bar'bar = GExpr (||)

instance (Eq a, Eq (ProbFunReprF a)) =>
         Interp__'eq'eq ProbFunRepr a where
  interp__'eq'eq = GExpr (==)

instance (Ord a, Ord (ProbFunReprF a)) =>
         Interp__'lt ProbFunRepr a where
  interp__'lt = GExpr (<)

instance (Ord a, Ord (ProbFunReprF a)) =>
         Interp__'gt ProbFunRepr a where
  interp__'gt = GExpr (>)

instance (Ord a, Ord (ProbFunReprF a)) =>
         Interp__'lt'eq ProbFunRepr a where
  interp__'lt'eq = GExpr (<=)

instance (Ord a, Ord (ProbFunReprF a)) =>
         Interp__'gt'eq ProbFunRepr a where
  interp__'gt'eq = GExpr (>=)

instance (Ord a, Ord (ProbFunReprF a)) =>
         Interp__min ProbFunRepr a where
  interp__min = GExpr min

instance (Ord a, Ord (ProbFunReprF a)) =>
         Interp__max ProbFunRepr a where
  interp__max = GExpr max

----------------------------------------------------------------------
-- Numeric Operations
----------------------------------------------------------------------

instance (Num a, Num (ProbFunReprF a)) =>
         Interp__'plus ProbFunRepr a where
  interp__'plus = GExpr (+)

instance (Num a, Num (ProbFunReprF a)) =>
         Interp__'minus ProbFunRepr a where
  interp__'minus = GExpr (-)

instance (Num a, Num (ProbFunReprF a)) =>
         Interp__'times ProbFunRepr a where
  interp__'times = GExpr (*)

instance (Num a, Num (ProbFunReprF a)) =>
         Interp__negate ProbFunRepr a where
  interp__negate = GExpr negate

instance (Num a, Num (ProbFunReprF a)) =>
         Interp__abs ProbFunRepr a where
  interp__abs = GExpr abs

instance (Num a, Num (ProbFunReprF a)) =>
         Interp__signum ProbFunRepr a where
  interp__signum = GExpr signum

instance (Num a, Num (ProbFunReprF a)) =>
         Interp__fromInteger ProbFunRepr a where
  interp__fromInteger = GExpr fromInteger

instance (Num a, Num (ProbFunReprF a)) =>
         Interp__'integer ProbFunRepr a where
  interp__'integer n = GExpr (fromInteger n)

instance (Interp__'integer ProbFunRepr a,
          Eq (ProbFunReprF a))
         => Interp__'eqInteger ProbFunRepr a where
  interp__'eqInteger (GExpr x) (GExpr y) = GExpr (x == y)

instance (Integral a, Integral (ProbFunReprF a)) =>
         Interp__quot ProbFunRepr a where
  interp__quot = GExpr quot

instance (Integral a, Integral (ProbFunReprF a)) =>
         Interp__rem ProbFunRepr a where
  interp__rem = GExpr rem

instance (Integral a, Integral (ProbFunReprF a)) =>
         Interp__div ProbFunRepr a where
  interp__div = GExpr div

instance (Integral a, Integral (ProbFunReprF a)) =>
         Interp__mod ProbFunRepr a where
  interp__mod = GExpr mod

instance (Integral a, Integral (ProbFunReprF a)) =>
         Interp__toInteger ProbFunRepr a where
  interp__toInteger = GExpr toInteger

instance (Fractional a, Fractional (ProbFunReprF a)) =>
         Interp__'div ProbFunRepr a where
  interp__'div = GExpr (/)

instance (Fractional a, Fractional (ProbFunReprF a)) =>
         Interp__recip ProbFunRepr a where
  interp__recip = GExpr recip

instance (Fractional a, Fractional (ProbFunReprF a)) =>
         Interp__fromRational ProbFunRepr a where
  interp__fromRational = GExpr fromRational

instance (Fractional a, Fractional (ProbFunReprF a)) =>
         Interp__'rational ProbFunRepr a where
  interp__'rational n = GExpr (fromRational n)

instance (Interp__'rational ProbFunRepr a,
          Eq (ProbFunReprF a))
         => Interp__'eqRational ProbFunRepr a where
  interp__'eqRational (GExpr x) (GExpr y) = GExpr (x == y)

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__pi ProbFunRepr a where
  interp__pi = GExpr pi

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__exp ProbFunRepr a where
  interp__exp = GExpr exp

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__log ProbFunRepr a where
  interp__log = GExpr log

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__sqrt ProbFunRepr a where
  interp__sqrt = GExpr sqrt

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__'times'times ProbFunRepr a where
  interp__'times'times = GExpr (**)

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__logBase ProbFunRepr a where
  interp__logBase = GExpr logBase

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__sin ProbFunRepr a where
  interp__sin = GExpr sin

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__cos ProbFunRepr a where
  interp__cos = GExpr cos

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__tan ProbFunRepr a where
  interp__tan = GExpr tan

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__asin ProbFunRepr a where
  interp__asin = GExpr asin

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__acos ProbFunRepr a where
  interp__acos = GExpr acos

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__atan ProbFunRepr a where
  interp__atan = GExpr atan

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__sinh ProbFunRepr a where
  interp__sinh = GExpr sinh

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__cosh ProbFunRepr a where
  interp__cosh = GExpr cosh

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__tanh ProbFunRepr a where
  interp__tanh = GExpr tanh

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__asinh ProbFunRepr a where
  interp__asinh = GExpr asinh

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__acosh ProbFunRepr a where
  interp__acosh = GExpr acosh

instance (Floating a, Floating (ProbFunReprF a)) =>
         Interp__atanh ProbFunRepr a where
  interp__atanh = GExpr atanh

----------------------------------------------------------------------
-- Probability expression operations
----------------------------------------------------------------------

instance Interp__realToProb ProbFunRepr where
  interp__realToProb = GExpr $ \x ->
    if x < 0 then 0 else rToProb x

instance Interp__logRealToProb ProbFunRepr where
  interp__logRealToProb = GExpr logRToProb

instance Interp__probToReal ProbFunRepr where
  interp__probToReal = GExpr probToR

instance Interp__probToLogReal ProbFunRepr where
  interp__probToLogReal = GExpr probToLogR

instance Interp__gammaProb ProbFunRepr where
  interp__gammaProb = GExpr (Prob . Log.Exp . logGamma)

instance Interp__digamma ProbFunRepr where
  interp__digamma = GExpr digamma

----------------------------------------------------------------------
-- Misc operations
----------------------------------------------------------------------

instance (Show a, Show (ProbFunReprF a)) =>
         Interp__gtrace ProbFunRepr a b where
  interp__gtrace = GExpr gtrace

instance Interp__gerror ProbFunRepr a where
  interp__gerror = GExpr gerror

----------------------------------------------------------------------
-- Interpreting Distributions
----------------------------------------------------------------------

instance Interp__delta ProbFunRepr Int where
  interp__delta = GExpr $ \i x -> if x == i then 1 else 0

instance Interp__normal ProbFunRepr where
  interp__normal = GExpr normalDensity

instance Interp__uniform ProbFunRepr where
  interp__uniform = GExpr uniformDensity

instance Interp__exponential ProbFunRepr where
  interp__exponential = GExpr $ \rate x -> Prob $ exponentialDensity rate x

instance Interp__gamma ProbFunRepr where
  interp__gamma = GExpr $ \k theta x -> Prob $ gammaDensity k theta x

instance Interp__beta ProbFunRepr where
  interp__beta = GExpr $ \alpha beta x -> Prob $ betaDensity alpha beta x

instance Interp__betaProb ProbFunRepr where
  interp__betaProb = GExpr $ \alpha beta p ->
    Prob $ betaDensityLog alpha beta (fromProb p)

instance Interp__dirichlet ProbFunRepr where
  interp__dirichlet = GExpr $ \alphas xs ->
    Prob $ dirichletDensity (toList alphas) (toList xs)

instance Interp__dirichletProb ProbFunRepr where
  interp__dirichletProb = GExpr $ \alphas xs ->
    Prob $ dirichletDensityLog (toList alphas) (map fromProb $ toList xs)

instance Interp__dirichletV ProbFunRepr where
  interp__dirichletV = GExpr $ \alphas xs ->
    dirichletDensityV alphas xs

instance Interp__dirichletPV ProbFunRepr where
  interp__dirichletPV = GExpr $ \alphas xs ->
    dirichletDensityPV alphas xs

instance Interp__categorical ProbFunRepr where
  interp__categorical = GExpr $ \probs x ->
    if x >= length (toList probs) then
      trace ("Categorical: unexpected value x = " ++ show x) 0
    else
      (toList probs) !! x

instance Interp__poisson ProbFunRepr where
  interp__poisson = GExpr $ \rate k ->
    if rate < 0 then error "Poisson: rate < 0!" else
      if rate == 0 then
        if k == 0 then 1 else 0
      else
        if k < 0 then 0 else
          Prob $ Log.Exp $
          fromIntegral k * log rate - rate - logGamma (fromIntegral $ k + 1)

instance Interp__binary_mixture ProbFunRepr a where
  interp__binary_mixture = GExpr $ \p d1 d2 x ->
    p * d1 x + (1 - p) * d2 x

instance Interp__scaleDist ProbFunRepr a where
  interp__scaleDist = GExpr $ \p d x -> p * d x

instance Interp__ctorDist__ListF ProbFunRepr where
  interp__ctorDist__Nil = GExpr $ \d xs ->
    case xs of
      Nil -> d Tuple0
      _ -> error "Unexpected Cons!"
  interp__ctorDist__Cons = GExpr $ \d xs ->
    case xs of
      Cons x xs' -> d (Tuple2 x xs')
      _ -> error "Unexpected Nil!"

instance Interp__adtDist__ListF ProbFunRepr where
  interp__adtDist__ListF = GExpr $ \probNil dNil probCons dCons xs ->
    case xs of
      Nil -> probNil * dNil Tuple0
      Cons x xs' -> probCons * dCons (Tuple2 x xs')

instance Interp__arbitrary ProbFunRepr a where
  interp__arbitrary = GExpr $ \_ -> 1

----------------------------------------------------------------------
-- Interpreting Vectors
----------------------------------------------------------------------

instance Interp__vec_head ProbFunRepr a where
  interp__vec_head = GExpr V.head

instance Interp__vec_tail ProbFunRepr a where
  interp__vec_tail = GExpr V.tail

instance Interp__vec_nth ProbFunRepr a where
  interp__vec_nth = GExpr (V.!)

instance Interp__vec_length ProbFunRepr a where
  interp__vec_length = GExpr V.length

instance Interp__vec_map ProbFunRepr a b where
  interp__vec_map = GExpr V.map

instance Interp__vec_foldr ProbFunRepr a b where
  interp__vec_foldr = GExpr V.foldr'

instance Interp__vec_iid ProbFunRepr a where
  interp__vec_iid = GExpr $ \n d xs ->
    if V.length xs == n then
      product $ V.map d xs
    else 0

{-
instance Interp__vec_dist ProbFunRepr where
  interp__vec_dist = GExpr $ \d xs -> d (V.toList xs)
-}

instance Interp__vec_nil_dist ProbFunRepr a where
  interp__vec_nil_dist = GExpr $ \xs ->
    if V.null xs then 1 else 0

instance Interp__vec_cons_dist ProbFunRepr a where
  interp__vec_cons_dist = GExpr $ \d xs ->
    if V.null xs then 0 else
      d (Tuple2 (GExpr $ V.head xs) (GExpr $ V.tail xs))

----------------------------------------------------------------------
-- Interpreting unboxed vectors and matrices
----------------------------------------------------------------------

instance Interp__lengthV ProbFunRepr where
  interp__lengthV = GExpr lengthV

instance Interp__atV ProbFunRepr where
  interp__atV = GExpr atV

instance Interp__generateV ProbFunRepr where
  interp__generateV = GExpr generateV

instance Interp__sumV ProbFunRepr where
  interp__sumV = GExpr sumV

instance Interp__rowsM ProbFunRepr where
  interp__rowsM = GExpr rowsM

instance Interp__colsM ProbFunRepr where
  interp__colsM = GExpr colsM

instance Interp__atM ProbFunRepr where
  interp__atM = GExpr atM

instance Interp__fromRowsM ProbFunRepr where
  interp__fromRowsM =
    GExpr (fromRowsM . map unGExpr . toHaskellListF unGExpr)

instance Interp__fromColsM ProbFunRepr where
  interp__fromColsM =
    GExpr (fromColsM . map unGExpr . toHaskellListF unGExpr)

instance Interp__toRowsM ProbFunRepr where
  interp__toRowsM =
    GExpr (fromHaskellListF GExpr . map GExpr . toRowsM)

instance Interp__toColsM ProbFunRepr where
  interp__toColsM =
    GExpr (fromHaskellListF GExpr . map GExpr . toColsM)

instance Interp__mulM ProbFunRepr where
  interp__mulM = GExpr mulM

instance Interp__mulMV ProbFunRepr where
  interp__mulMV = GExpr mulMV

instance Interp__mulVM ProbFunRepr where
  interp__mulVM = GExpr mulVM

----------------------------------------------------------------------
-- Interpreting unboxed vectors and matrices of probabilities
----------------------------------------------------------------------

instance Interp__lengthPV ProbFunRepr where
  interp__lengthPV = GExpr lengthPV

instance Interp__atPV ProbFunRepr where
  interp__atPV = GExpr atPV

instance Interp__generatePV ProbFunRepr where
  interp__generatePV = GExpr generatePV

instance Interp__sumPV ProbFunRepr where
  interp__sumPV = GExpr sumPV

instance Interp__rowsPM ProbFunRepr where
  interp__rowsPM = GExpr rowsPM

instance Interp__colsPM ProbFunRepr where
  interp__colsPM = GExpr colsPM

instance Interp__atPM ProbFunRepr where
  interp__atPM = GExpr atPM

instance Interp__fromRowsPM ProbFunRepr where
  interp__fromRowsPM =
    GExpr (fromRowsPM . map unGExpr . toHaskellListF unGExpr)

instance Interp__fromColsPM ProbFunRepr where
  interp__fromColsPM =
    GExpr (fromColsPM . map unGExpr . toHaskellListF unGExpr)

instance Interp__toRowsPM ProbFunRepr where
  interp__toRowsPM =
    GExpr (fromHaskellListF GExpr . map GExpr . toRowsPM)

instance Interp__toColsPM ProbFunRepr where
  interp__toColsPM =
    GExpr (fromHaskellListF GExpr . map GExpr . toColsPM)

instance Interp__mulPM ProbFunRepr where
  interp__mulPM = GExpr mulPM

instance Interp__mulPMV ProbFunRepr where
  interp__mulPMV = GExpr mulPMV

instance Interp__mulPVM ProbFunRepr where
  interp__mulPVM = GExpr mulPVM

----------------------------------------------------------------------
-- Distributions involving unboxed vectors and matrices of probabilities
----------------------------------------------------------------------

instance Interp__categoricalPV ProbFunRepr where
  interp__categoricalPV = GExpr $ \probs x ->
    if x >= lengthPV probs then
      trace ("Categorical: unexpected value x = " ++ show x) 0
    else
      atPV probs x

instance Interp__categoricalV ProbFunRepr where
  interp__categoricalV = GExpr $ \probs x ->
    if x >= lengthV probs then
      trace ("Categorical: unexpected value x = " ++ show x) 0
    else
      realToProb (atV probs x)

{-
instance Interp__mv_normal ProbFunRepr where
  interp__mvNormal =
-}

----------------------------------------------------------------------
-- Interpreting VI Distribution Families
----------------------------------------------------------------------

instance Interp__withVISize ProbFunRepr where
  interp__withVISize = GExpr $ \f -> bindVIDimFamExpr f

instance Interp__viNormal ProbFunRepr where
  interp__viNormal = GExpr normalVIFamExpr

instance Interp__viUniform ProbFunRepr where
  interp__viUniform = GExpr uniformVIFamExpr

instance Interp__viGamma ProbFunRepr where
  interp__viGamma = GExpr gammaVIFamExpr

instance Interp__viGammaProb ProbFunRepr where
  interp__viGammaProb = GExpr gammaProbVIFamExpr

instance Interp__viBeta ProbFunRepr where
  interp__viBeta = GExpr betaVIFamExpr

instance Interp__viBetaProb ProbFunRepr where
  interp__viBetaProb = GExpr betaProbVIFamExpr

instance Interp__viCategorical ProbFunRepr where
  interp__viCategorical = GExpr categoricalVIFamExpr

instance Interp__viDirichlet ProbFunRepr where
  interp__viDirichlet =
    GExpr $ \sz ->
    xformVIDistFamExpr fromList toList (dirichletVIFamExpr sz)

instance Interp__viDirichletProb ProbFunRepr where
  interp__viDirichletProb =
    GExpr $ \sz ->
    xformVIDistFamExpr fromList toList (dirichletProbVIFamExpr sz)

instance Interp__viDirichletV ProbFunRepr where
  interp__viDirichletV =
    GExpr $ \sz ->
    xformVIDistFamExpr (fromListV . fromList) (toList . toListV)
    (dirichletVIFamExpr sz)

instance Interp__viDirichletPV ProbFunRepr where
  interp__viDirichletPV =
    GExpr $ \sz ->
    xformVIDistFamExpr (fromListPV . fromList) (toList . toListPV)
    (dirichletProbVIFamExpr sz)

instance (Eq (ProbFunReprF a), GrappaShow (ProbFunReprF a)) =>
         Interp__viDelta ProbFunRepr a where
  interp__viDelta = GExpr deltaVIFamExpr

instance (Eq (ProbFunReprF a), GrappaShow (ProbFunReprF a),
          FromJSON (ProbFunReprF a)) =>
         Interp__viJSONInput ProbFunRepr a where
  interp__viJSONInput = GExpr readJSONVIDistFamExpr

instance FromJSON (ProbFunReprF a) =>
         Interp__viMappedJSONInput ProbFunRepr a b where
  interp__viMappedJSONInput = GExpr mapJSONVIDistFamExpr

instance Interp__viTuple0 ProbFunRepr where
  interp__viTuple0 = GExpr $ tupleVIFamExpr Tuple0

instance Interp__viTuple1 ProbFunRepr a where
  interp__viTuple1 = GExpr $ \da ->
    tupleVIFamExpr (Tuple1 $ Compose $ xformVIDistFamExpr GExpr unGExpr da)

instance Interp__viTuple2 ProbFunRepr a b where
  interp__viTuple2 = GExpr $ \da db ->
    tupleVIFamExpr (Tuple2
                    (Compose $ xformVIDistFamExpr GExpr unGExpr da)
                    (Compose $ xformVIDistFamExpr GExpr unGExpr db))

instance Interp__viTuple3 ProbFunRepr a b c where
  interp__viTuple3 = GExpr $ \da db dc ->
    tupleVIFamExpr (Tuple3
                    (Compose $ xformVIDistFamExpr GExpr unGExpr da)
                    (Compose $ xformVIDistFamExpr GExpr unGExpr db)
                    (Compose $ xformVIDistFamExpr GExpr unGExpr dc))

instance Interp__viTuple4 ProbFunRepr a b c d where
  interp__viTuple4 = GExpr $ \da db dc dd ->
    tupleVIFamExpr (Tuple4
                    (Compose $ xformVIDistFamExpr GExpr unGExpr da)
                    (Compose $ xformVIDistFamExpr GExpr unGExpr db)
                    (Compose $ xformVIDistFamExpr GExpr unGExpr dc)
                    (Compose $ xformVIDistFamExpr GExpr unGExpr dd))

instance Interp__viIID ProbFunRepr a where
  interp__viIID =
    GExpr $ \sz d ->
    xformVIDistFamExpr fromList toList (iidVIFamExpr sz d)

instance Interp__viVecDist ProbFunRepr a where
  interp__viVecDist = GExpr vecDistVIFamExpr

instance Interp__viVecIID ProbFunRepr a where
  interp__viVecIID = GExpr vecIIDVIFamExpr

instance Interp__viIIDPV ProbFunRepr where
  interp__viIIDPV = GExpr iidPVVIFamExpr


{-|
 Module      : Data.Boltzmann.System.Tuner
 Description : Interface utilities with the Paganini tuner.
 Copyright   : (c) Maciej Bendkowski, 2017-2020

 License     : BSD3
 Maintainer  : maciej.bendkowski@tcs.uj.edu.pl
 Stability   : experimental

 General utilities managing the IO interface between Boltzmann Brain
 and the Paganini tuner script.
 -}
module Data.Boltzmann.System.Tuner
  ( readPaganini
  , runPaganini
  )
where

import           Data.Void
import qualified Data.Map.Strict               as M

import           Numeric.LinearAlgebra   hiding ( size
                                                , double
                                                )

import           Text.Megaparsec

import           Data.Boltzmann.Internal.Logging
import           Data.Boltzmann.Internal.Parser
import           Data.Boltzmann.Internal.Tuner
import           Data.Boltzmann.System
import           Data.Boltzmann.System.Paganini
import qualified Data.Boltzmann.System.Tuner.Algebraic
                                               as A
import           Data.Boltzmann.Internal.Annotations
                                                ( withString )

type Parser = Parsec Void String

-- | Communicates with Paganini and collects the respective
--   tuning vector for the given system. If communication is not possible,
--   for instance due to the missing Paganini script, the current process
--   is terminated with an error message on the standard error output.
runPaganini
  :: System Int
  -> Parametrisation
  -> IO (Either (ParseErrorBundle String Void) (PSystem Double))

runPaganini sys paramT = do

  -- FIXME use arg
  info "Running paganini..."
  let gen = withString (annotations sys) "generate" (initType sys)
  result <- toPaganini gen sys
  case result of
    Left  err           -> fail' (show err)
    Right (rho, ts, us) -> do
      info "Computing branching probabilities..."
      let ts'  = fromList ts
      let sys' = parametrise sys paramT rho ts' us
      return $ Right sys'

-- | Parses the given input string as a Paganini tuning vector.
readPaganini
  :: System Int
  -> Parametrisation
  -> String
  -> IO (Either (ParseErrorBundle String Void) (PSystem Double))

readPaganini sys paramT f = do
  let spec = A.toPSpec sys
  pag <- parsePaganini spec f
  case pag of
    Left  err           -> return $ Left err
    Right (rho, us, ts) -> do
      let ts' = fromList ts
      return (Right $ parametrise sys paramT rho ts' us)

paganiniStmt :: PSpec -> Parser (Double, [Double], [Double])
paganiniStmt spec = do
  rho <- double
  us  <- parseN double $ numFreqs spec
  ts  <- parseN double $ numTypes spec
  return (rho, us, ts)

-- | Parses the given Paganini specification.
parsePaganini
  :: PSpec
  -> String
  -> IO
       ( Either
           (ParseErrorBundle String Void)
           (Double, [Double], [Double])
       )

parsePaganini spec = parseFromFile (paganiniStmt spec)

accumulateCons :: Double -> [Cons Double] -> [Cons Double]
accumulateCons _   []         = []
accumulateCons acc (con : xs) = con' : accumulateCons acc' xs
 where
  con' = con { weight = w + acc }
  acc' = acc + w
  w    = weight con

accumulateType :: (String, [Cons Double]) -> (String, [Cons Double])
accumulateType (t, cons) = (t, accumulateCons 0 cons)

accumulate :: PSystem Double -> PSystem Double
accumulate psys = psys { system = sys { defs = M.fromList types' } }
 where
  sys    = system psys
  typs   = M.toList (defs sys)
  types' = map accumulateType typs

-- | Compute the numerical branching probabilities for the given system.
parametrise
  :: System Int
  -> Parametrisation
  -> Double
  -> Vector Double
  -> [Double]
  -> PSystem Double

parametrise sys paramT rho ts us =
  let sys' = A.paramSystem sys rho ts us
  in  case paramT of
        Regular     -> sys'
        Cummulative -> accumulate sys'


{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE ScopedTypeVariables #-}

module ModelSpec where

import Control.Applicative
import Codec.Picture
import Codec.Picture.Types
import qualified Data.Aeson as A
import Data.Complex
import qualified Data.HashMap.Strict as HM
import Data.Proxy
import qualified Data.Text as T
import qualified Data.Text.Lazy as TL
import qualified Data.Vector as V
import GHC.Generics
import GHC.Word
import Generics.SOP.Arbitrary
import Generics.SOP
import Test.QuickCheck
import Test.QuickCheck.Function
import Test.QuickCheck.Arbitrary
import Test.Hspec
import Test.Hspec.QuickCheck
import Model

spec = do
  describe "Encoding and decoding"
    $ modifyMaxSize (const 10)
    $ modifyMaxSuccess (const 50) $ do
      it "val should round-trip aeson" $  property $ \(v :: Val) ->
        A.decode (A.encode v) == Just v

instance Arbitrary Val where
  arbitrary = sized arbVal

arbVal :: Int -> Gen Val
arbVal 0 = oneof [ VDouble  <$> arbitrary
                 , VComplex <$> arbitrary
                 , VText    <$> arbitrary
                 , VImage   <$> arbitrary
                 ]
arbVal n = do
  (Positive m) <- arbitrary
  let n' = n `div` (m + 1)
  oneof [ (VList . V.toList) <$> V.replicateM n (arbVal n')
        , VVec1 <$> V.replicateM n' (arbVal m)
        , fmap VVec2 $ do
           x <- (V.replicateM (div n 4) (arbVal (n `div` 4)))
           y <- (V.replicateM (div n 4) (arbVal (n `div` 4)))
           return (V.zip x y)
        , fmap (VVec3 . V.fromList) $ do
           x <- (vectorOf (div n 2) (arbVal (div n 2)))
           y <- (vectorOf (div n 2) (arbVal (div n 2)))
           z <- (vectorOf (div n 2) (arbVal (div n 2)))
           return (zip3 x y z)
        ]

instance Arbitrary PrimComplex where
  arbitrary = fmap PComplex $ liftA2 (:+) arbitrary arbitrary

instance Arbitrary A.Value where
  arbitrary = oneof
    [ fmap (A.String . T.pack) arbitrary
    , fmap (A.Number . realToFrac) (arbitrary :: Gen Double)
    , return A.Null
    , fmap (A.Object . HM.fromList) (zip <$> arbitrary <*> arbitrary)
    , fmap (A.Array . V.fromList) arbitrary
    ]

instance Arbitrary Expr where
  arbitrary = fmap Expr arbitrary

instance Arbitrary T.Text where
  arbitrary = T.pack <$> arbitrary

genImg :: (Pixel a, Arbitrary a)
       => Proxy a
       -> Int
       -> Int
       -> Gen (Image a)
genImg _ wid hgt = do
  f :: Fun (Int,Int) a <- arbitrary
  let img = generateImage (curry (apply f)) wid hgt
  return img

r16 :: (Word16,Word16)
r16 = (minBound,maxBound)

r8 :: (Word8,Word8)
r8 = (minBound,maxBound)

rF :: (Float, Float)
rF = (0,1)

instance Arbitrary a => Arbitrary (V.Vector a) where
  arbitrary = fmap V.fromList arbitrary

instance Arbitrary PixelYA8 where
  arbitrary = liftA2 PixelYA8 (choose r8) (choose r8)

instance Arbitrary PixelYA16 where
  arbitrary = liftA2 PixelYA16 (choose r16) (choose r16)

instance Arbitrary PixelRGB8 where
  arbitrary = liftA3 PixelRGB8 (choose r8) (choose r8) (choose r8)

instance Arbitrary PixelRGB16 where
  arbitrary = liftA3 PixelRGB16 (choose r16) (choose r16) (choose r16)

instance Arbitrary PixelRGBF where
  arbitrary = liftA3 PixelRGBF (choose rF) (choose rF) (choose rF)

liftA4 :: Applicative m
       => (a -> b -> c -> d -> e) -> m a -> m b -> m c -> m d -> m e
liftA4 f a b c d = f <$> a <*> b <*> c <*> d

instance Arbitrary PixelRGBA8 where
  arbitrary = liftA4 PixelRGBA8 (choose r8) (choose r8) (choose r8) (choose r8)

instance Arbitrary PixelRGBA16 where
  arbitrary = liftA4 PixelRGBA16 (choose r16) (choose r16) (choose r16) (choose r16)

instance Arbitrary PixelCMYK8 where
  arbitrary = liftA4 PixelCMYK8 (choose r8) (choose r8) (choose r8) (choose r8)

instance Arbitrary PixelCMYK16 where
  arbitrary = liftA4 PixelCMYK16 (choose r16) (choose r16) (choose r16) (choose r16)

instance Arbitrary PixelYCbCr8 where
  arbitrary = liftA3 PixelYCbCr8 (choose r8) (choose r8) (choose r8)

instance Arbitrary ModelImage where
  arbitrary = genModelImage

genModelImage :: Gen ModelImage
genModelImage = do
    x <- choose (10 :: Int, 200)
    y <- choose (10 :: Int, 200)
    dynImg <- oneof
      [ -- ImageY8     <$> genImg (Proxy :: Proxy Pixel8)      x y
      -- , ImageY16    <$> genImg (Proxy :: Proxy Pixel16)     x y
      -- , ImageYF     <$> genImg (Proxy :: Proxy PixelF)      x y
      -- , ImageYA8    <$> genImg (Proxy :: Proxy PixelYA8)    x y
      -- , ImageYA16   <$> genImg (Proxy :: Proxy PixelYA16)   x y
        ImageRGB8   <$> genImg (Proxy :: Proxy PixelRGB8)   x y
      , ImageRGB16  <$> genImg (Proxy :: Proxy PixelRGB16)  x y
      -- , ImageRGBF   <$> genImg (Proxy :: Proxy PixelRGBF)   x y
      , ImageRGBA8  <$> genImg (Proxy :: Proxy PixelRGBA8)  x y
      , ImageRGBA16 <$> genImg (Proxy :: Proxy PixelRGBA16) x y
      -- , ImageYCbCr8 <$> genImg (Proxy :: Proxy PixelYCbCr8) x y
      -- , ImageCMYK8  <$> genImg (Proxy :: Proxy PixelCMYK8)  x y
      -- , ImageCMYK16 <$> genImg (Proxy :: Proxy PixelCMYK16) x y
      ]
    return (ModelImage $ promoteImage $ convertRGBA8 dynImg)