#!/usr/bin/env stack
-- stack runghc --package reanimate
{-# LANGUAGE OverloadedStrings #-}
module Main where

import           Codec.Picture.Types
import           Control.Lens                  ()
import           Control.Monad
import           Data.Function
import           Data.List
import           Data.List.NonEmpty            (NonEmpty)
import qualified Data.List.NonEmpty            as NE
import           Data.Maybe
import           Data.Ratio
import qualified Data.Text                     as T
import           Data.Tuple
import qualified Data.Vector                   as V
import           Debug.Trace
import           Linear.Matrix                 hiding (trace)
import           Linear.Metric
import           Linear.V2
import           Linear.V3
import           Linear.Vector
import           Numeric.LinearAlgebra         hiding (polar, scale, (<>))
import qualified Numeric.LinearAlgebra         as Matrix
import           Numeric.LinearAlgebra.HMatrix hiding (polar, scale, (<>))
import           Reanimate
import           Reanimate.Animation
import           Reanimate.Math.Balloon
import           Reanimate.Math.Common
import           Reanimate.Math.Triangulate
import           Reanimate.Math.Polygon
import           Reanimate.Math.EarClip
import           Reanimate.Math.SSSP
import           Reanimate.Math.Render
import           Reanimate.Math.Visibility
import           Reanimate.Math.Compatible
import           Reanimate.Morph.Common
import           Reanimate.PolyShape           (svgToPolygons)

p :: Polygon
p = mkPolygon $ V.fromList [V2 (0 % 1) (0 % 1),V2 (1 % 1) (0 % 1),V2 (1 % 1) (1 % 1),V2 (2 % 1) (1 % 1),V2 (2 % 1) ((-1) % 1),V2 (3 % 1) ((-1) % 1),V2 (3 % 1) (2 % 1),V2 (0 % 1) (2 % 1)]

pCuts' :: Polygon -> [(Int,Int)]
pCuts' p =
  [ (i, j)
  | i <- [0 .. pSize p-1 ]
  , j <- [i+2 .. pSize p-1 ]
  , (j+1) `mod` pSize p /= i
  , trace ("Check: " ++ show (i,j, pSize p)) $ pParent p i j == i ]

-- p :: Polygon
-- p = pScale 6 $ unsafeSVGToPolygon 0.1 $
--   lowerTransformations $ pathify $ center $ latex "$1$"

main :: IO ()
main = reanimate $ sceneAnimation $ do
  bg <- newSpriteSVG $ mkBackground "black"
  spriteZ bg (-1)
  newSpriteSVG_ $ translate 0 1 $ mkGroup
    [ withFillColor "grey" $ polygonShape p
    , polygonNumDots p
    ]
  forM_ (pCuts p) $ \(l,r) -> do
    play $ mkAnimation (1/60) $ \_ -> mkGroup
      [ translate (-3) 0 $ withFillColor "grey" $ polygonShape l
      , translate (-3) 0 $ polygonNumDots l
      , translate (3) 0 $ withFillColor "grey" $ polygonShape r
      , translate (3) 0 $ polygonNumDots r
      ]
  -- wait 1
  -- fork $ play $ drawTriangulation shape1 earCut'
  --   # mapA (translate (-3) 0)
  -- play $ drawTriangulation shape1 earClip'
  --   # mapA (translate (3) 0)
  return ()


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

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

import GameData

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

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

gameSimplePartial :: PayoffsRaw -> Payoffs
gameSimplePartial = fromLists

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

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

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

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

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

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

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

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

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

-- from here on out, stuff for solving games

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

                solveSimple outsc outsr

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

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


module Main where

import Numeric.LinearAlgebra
import Data.Function (on)
import Data.List (minimumBy)

import Lib

type MarcovChain = Vector R -> Vector R

type State = (Int,Int)

type Window = Int

type Steps = Int

type Data t = t State

type States t = t State

classify :: Steps -> [Double] -> [Int]
classify steps ds = map (roundToNearest . (* 100) . (/ maxValue)) ds
  where
    maxValue = maximum ds

    increments = statePercents steps

    roundToNearest :: Double -> Int
    roundToNearest d = case dropWhile (\(_,y) -> fromIntegral y <= d) $ zip increments (tail increments) of
                          ((x,y):_) -> if d - fromIntegral x >= fromIntegral y - d then y else x
                          _         -> last increments

makeMarcovChainFromData :: Foldable t => Window -> Steps -> Data t -> MarcovChain
makeMarcovChainFromData = undefined

stateProduct :: Window -> Steps -> States []
stateProduct winSize steps = [(i,j) | i <- statePercents steps, j <- [1..winSize]]

statePercents :: Steps -> [Int]
statePercents steps = [0,stepSize..200]
  where
    stepSize = round $ 200 / fromIntegral steps

predict :: MarcovChain -> Vector R -> Vector R
predict = ($)

main :: IO ()
main = print "Crypto"


module Statistics.Iteratee.Tests

where

import Data.Iteratee as I
import Statistics.Iteratee as Si
import Statistics.Sample as St
import Test.Framework
import Test.Framework.Providers.QuickCheck2 (testProperty)

import Control.Monad.Identity
import qualified Data.Vector.Unboxed as V

tests =
  [ testGroup "Sample" $ map mkUnProp uns
  ]

-- Don't want to drag in more dependencies for just this.
class ApproxEq a where
    approxEq :: Double -> a -> a -> Bool

instance ApproxEq Double where
    approxEq tol d1 d2 = m == 0.0 || d/m < tol where
      m = max (abs d1) (abs d2)
      d = abs (d1 - d2)

instance (ApproxEq a, ApproxEq b) => ApproxEq (a,b) where
    approxEq tol (l1,r1) (l2,r2) = approxEq tol l1 l2 && approxEq tol r1 r2

infix 4 ===
(===) :: (ApproxEq a) => a -> a -> Bool
(===) = approxEq {-pretty equal-}1.0e-10

unsProp :: (ApproxEq a)
        => (V.Vector Double -> a)
        -> (Iteratee [Double] Identity a)
        -> [Double]
        -> Bool
unsProp vec iter xs = if null xs then True
    else vec (V.fromList xs) === (runIdentity $ run =<< enumPure1Chunk xs iter)
    -- we're using Eq for doubles, which is always a bad idea...

    -- also not checking empty vectors, because in some cases (range,
    -- harmonicMean) we get NaN's or other funky values.

mkUnProp (lbl, st, si) = testProperty lbl $ unsProp st si

-- unary properties
uns =
  [ ("mean", St.mean, Si.mean)
  , ("range", St.range, Si.range)
  , ("harmonic_mean", St.harmonicMean, Si.harmonicMean)
  , ("variance", St.fastVariance, Si.variance)
  , ("std_dev", St.fastStdDev,   Si.stdDev)
  ]


module HScheme.Parser where

import HScheme.Data

import Text.ParserCombinators.Parsec as P hiding ( spaces )
import Data.Char ( toUpper )
import Data.Ratio ( (%) )
import Data.Complex ( Complex((:+)) )
import Numeric ( readHex, readOct )
import Data.List ( foldl' )
import Control.Monad.Except
import qualified Data.Vector as V

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

spaces :: Parser ()
spaces = skipMany1 space

parseBool :: Parser Value
parseBool = try $ do
    char '#'
    (char 't' >> return (Bool True)) <|> (char 'f' >> return (Bool False))

parseString :: Parser Value
parseString = do
    char '"'
    x <- many (noneOf "\"" <|> escapeChar)
    char '"'
    return $ String x
    where
        escapeChar :: Parser Char
        escapeChar = do
            oneOf "nrt\\" >>= \case
                'n' -> return '\n'
                'r' -> return '\r'
                't' -> return '\t'
                '\\'-> return '\\'

parseChar :: Parser Value
parseChar = do
    try $ string "#\\"
    chr <- (string "newline" <|> string "space") <|>
        do
            x <- anyChar
            notFollowedBy alphaNum
            return [x]
    return $ Character $ case chr of
        "\\" -> '\n'
        "newline" -> '\n'
        "space" -> ' '
        _ -> head chr

parseAtom :: Parser Value
parseAtom = do
    first <- letter <|> symbol
    rest <- many (letter <|> digit <|> symbol)
    let atom = first:rest
    return $ case atom of
        "#t" -> Bool True
        "#f" -> Bool False
        _    -> Atom atom

{-
    Parsing numbers
-}

parseDecimal :: Parser Value
parseDecimal = do
    sign <- option "" $ string "-"
    many1 digit >>= (return . Number . read . (sign <>))

parseDecimal' :: Parser Value
parseDecimal' = try $ string "#d" >> many1 digit >>= (return . Number . read)

-- | Parse hexadecimal number of form #x...
parseHex :: Parser Value
parseHex = try $ string "#x" >> many1 hexDigit >>= (return . Number . hexToDigit)
    where
        hexToDigit x = fst $ head $ readHex x

-- | Parse octadecimal number of form #o...
parseOct :: Parser Value
parseOct = try $ string "#o" >> many1 octDigit >>= (return . Number . octToDigit)
    where
        octToDigit x = fst $ head $ readOct x

-- | Parse binary of form #b...
parseBin :: Parser Value
parseBin = try $ string "#b" >> many1 (oneOf "10") >>= (return . Number . binToDigit)
    where
        binToDigit = foldl' binDigitToDigit 0
            where
                binDigitToDigit acc = \case
                    '1' -> 2 * acc + 1
                    '0' -> 2 * acc + 0

parseNumber :: Parser Value
parseNumber = do
    parseDecimal
    <|> parseDecimal'
    <|> parseHex
    <|> parseOct
    <|> parseBin

parseFloat :: Parser Value
parseFloat = try $ do
    sign <- option "" $ string "-"
    x <- many1 digit
    y <- (:) <$> char '.' <*> many1 digit
    return $ Float $ read $ sign <> x <> y

parseRatio :: Parser Value
parseRatio = try $ do
    (Number x) <- parseDecimal
    char '/'
    (Number y) <- parseDecimal
    return $ Ratio (x % y)

parseComplex :: Parser Value
parseComplex = try $ do
    real <- parseFloat <|> parseNumber
    char '+'
    imiginary <- parseFloat <|> parseNumber
    char 'i'
    return $ Complex (toDouble real :+ toDouble imiginary)
    where
        toDouble (Float f) = realToFrac f
        toDouble (Number n) = realToFrac n

parseExpr :: Parser Value
parseExpr = parseAtom
    <|> parseString
    <|> parseChar
    <|> parseRatio
    <|> parseComplex
    <|> parseFloat
    <|> parseNumber
    <|> parseBool
    <|> parseQuoted
    <|> parseVector
    <|> parseList

{-
    Recursive Parsers
-}

parseList :: Parser Value
parseList = do
    char '('
    x <- try parseSpacesList <|> parseDottedList
    char ')'
    return x

parseSpacesList :: Parser Value
parseSpacesList = List <$> sepBy parseExpr spaces

parseDottedList :: Parser Value
parseDottedList = do
    head <- endBy parseExpr spaces
    tail <- char '.' >> spaces >> parseExpr
    return $ DottedList head tail

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

parseVector :: Parser Value
parseVector = do
    try $ string "#("
    values <- sepBy parseExpr spaces
    char ')'
    return $ Vector (V.fromList values)

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


--Haskell Quil compiler
--Copyright Laurence Emms 2018

module Lib (compile) where

import Data.Complex
import Data.List
import Register
import Instruction
import ClassicalCircuit
import QuantumCircuit
import Util

testRXCircuit :: (Floating a, Show a, Ord a) => Circuit a
testRXCircuit = let parameters@[Left a] = [Left (MetaQubitRegister "a")]
                    instructions = [(RX (ComplexConstant (5.0 :+ 10.0)) a)]
                in Circuit "TESTRX" parameters instructions

hadamardRotateXYZ :: (Floating a, Show a, Ord a) => Complex a -> Circuit a
hadamardRotateXYZ rotation
    = let parameters@[Left a, Left b, Right r, Left z] = [Left (MetaQubitRegister "a"), Left (MetaQubitRegister "b"), Right (MetaRegister "r"), Left (MetaQubitRegister "z")]
          instructions = [(Hadamard a),
                          (MeasureOut a r)] ++
                          (ifC "HROTXYZTHEN0" "HROTXYZEND0"
                               r
                               [(RX (ComplexConstant rotation) z)]
                               ([(Hadamard b),
                                 (MeasureOut b r)] ++
                                (ifC "HROTXYZTHEN1" "HROTXYZEND1"
                                     r
                                     [(RY (ComplexConstant rotation) z)]
                                     [(RZ (ComplexConstant rotation) z)])))
      in Circuit "HROTXYZ" parameters instructions

compile :: IO ()
compile = putStrLn "Compiling quantum executable" >>
          --Test registers
          putStrLn (show $ QubitRegister 10) >>
          putStrLn (show $ Register 5) >>
          putStrLn (show $ ComplexConstant (6.0 :+ 7.0)) >>
          putStrLn (show $ ComplexConstant ((-8.0) :+ 1.0)) >>
          putStrLn (show $ ComplexConstant (2.0 :+ (-3.0))) >>
          putStrLn (show $ ComplexConstant ((-4.0) :+ (-6.0))) >>
          --Test instructions
          putStrLn (show $ CNot (QubitRegister 0) (QubitRegister 1)) >>
          putStrLn (show $ PSwap (ComplexConstant (5.0 :+ (-3.2))) (QubitRegister 0) (QubitRegister 1)) >>
          putStrLn (show $ Measure (QubitRegister 4)) >>
          putStrLn (show $ MeasureOut (QubitRegister 4) (Register 5)) >>
          --Test circuits
          putStrLn (show $ DefCircuit bell) >>
          putStrLn (show $ CallCircuit bell [Left (QubitRegister 5), Left (QubitRegister 3)]) >>
          putStrLn (show $ DefCircuit testRXCircuit) >>
          putStrLn (show $ CallCircuit testRXCircuit [Left (QubitRegister 1)]) >>
          putStrLn (show $ DefCircuit xor) >>
          putStrLn (show $ CallCircuit xor [Right (Register 0), Right (Register 1), Right (Register 2)]) >>
          putStrLn (show $ DefCircuit halfAdder) >>
          putStrLn (show $ CallCircuit halfAdder [Right (Register 0), Right (Register 1), Right (Register 2), Right (Register 3)]) >>
          putStrLn (show $ DefCircuit adder) >>
          putStrLn (show $ CallCircuit adder [Right (Register 0), Right (Register 1), Right (Register 2), Right (Register 3), Right (Register 4)]) >>
          --Load the integer 53 into registers [0-31]
          putStrLn (intercalate "\n" (fmap show (loadIntToRegisters 53 0))) >>
          --Load the integer 18 into registers [32-63]
          putStrLn (intercalate "\n" (fmap show (loadIntToRegisters 18 32))) >>
          --Add [0-31] + [32-63] into [64-95]
          let carry = 128
              temp = 129
          in putStrLn (intercalate "\n" (fmap show (addRegisters carry temp 0 32 64))) >>
          putStrLn (show $ DefCircuit (hadamardRotateXYZ (0.70710678118 :+ 0.70710678118))) >>
          putStrLn (show $ CallCircuit (hadamardRotateXYZ (0.70710678118 :+ 0.70710678118)) [Left (QubitRegister 0), Left (QubitRegister 1), Right (Register 0), Left (QubitRegister 2)])


{-# LANGUAGE RecordWildCards #-}
import Data.Complex
import Control.Monad.Trans.Except
import Foreign.C.Types
import Data.Word
import System.IO
import Data.Monoid

import Control.Error.Util
import Data.Vector.Storable as VS hiding ((++))
import Data.Vector.Generic as VG hiding ((++))
import Pipes
import qualified Pipes.Prelude as P
import Foreign.Storable.Complex

import Options.Applicative

import SDR.Filter
import SDR.RTLSDRStream
import SDR.Util
import SDR.Demod
import SDR.Pulse
import SDR.Serialize
import SDR.ArgUtils
import SDR.CPUID

--The filter coefficients are stored in another module
import Coeffs

data Options = Options {
    frequency  :: Word32,
    input      :: Maybe FilePath,
    output     :: Maybe FilePath
}

optParser :: Parser Options
optParser = Options
          <$> option (fmap fromIntegral parseSize) (
                 long "frequency"
              <> short 'f'
              <> metavar "FREQUENCY"
              <> help "Frequency to tune to"
              )
          <*> optional (strOption (
                 long "input"
              <> short 'i'
              <> metavar "FILENAME"
              <> help "Input filename"
              ))
          <*> optional (strOption (
                 long "output"
              <> short 'o'
              <> metavar "FILENAME"
              <> help "Output filename"
              ))

opt :: ParserInfo Options
opt = info (helper <*> optParser) (fullDesc <> progDesc "Receive and demodulate broadcast FM radio" <> header "RTLSDR FM")

bufNum     = 1
bufLen     = 16384
samples    = fromIntegral (bufNum * bufLen) `quot` 2
decimation = 8
sqd        = samples `quot` decimation

{-
    Sampling frequency of the input is 1280 khz
    This is decimated by a factor of 8 and then demodulated
    Sampling frequency of demodulated signal is 160 khz
    Need audio output at 48 khz
    Resampling factor is 48/160 == 3/10
    FM pilot tone at 19khz (0.3958 * 48)
    Start audio filter cutoff at 15khz (0.3125 * 48)
-}

doIt Options{..} = do

    info <- lift getCPUInfo

    let rtlstream = do
            str <- sdrStream (defaultRTLSDRParams frequency 1280000) bufNum bufLen
            return $ str >-> P.map (interleavedIQUnsignedByteToFloatFast info)

    let fileStream fname = lift $ do
            h <- openFile fname ReadMode
            return $ fromHandle samples h
            --TODO: how do I ensure these handles get closed?

    inputSpectrum <- maybe rtlstream fileStream input

    let fileSink fname = do
            h <- openFile fname ReadMode
            return $ toHandle h

    sink <- lift $ maybe pulseAudioSink fileSink output

    deci <- lift $ fastDecimatorC info decimation coeffsRFDecim
    resp <- lift $ fastResamplerR info 3 10 coeffsAudioResampler
    filt <- lift $ fastFilterSymR info coeffsAudioFilter

    --Build the pipeline
    let pipeline :: Effect IO ()
        pipeline =   inputSpectrum
                 >-> firDecimator deci sqd
                 >-> fmDemod
                 >-> firResampler resp sqd
                 >-> firFilter filt sqd
                 >-> P.map (VG.map (* 0.2))
                 >-> sink

    --Run the pipeline
    lift $ runEffect pipeline

main = execParser opt >>= exceptT putStrLn return . doIt


{-# LANGUAGE Arrows              #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies        #-}
{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS -fplugin=Overloaded -fplugin-opt=Overloaded:Categories #-}
module Main where

import Control.Monad      (when)
import Data.Word          (Word64)
import Numeric            (showFFloat)
import System.Environment (getArgs)
import Data.List (intercalate)

import qualified Control.Category
import qualified Numeric.LinearAlgebra  as LA
import qualified System.Random.SplitMix as SM

import Overloaded.Categories
import VectorSpace

-- | A Function which computes value and derivative at the point.
newtype AD a b = AD (a -> (b, L a b))

instance Category AD where
    id = AD (\x -> (x, L id))

    AD g . AD f = AD $ \a ->
        let (b, L f') = f a
            (c, L g') = g b
        in (c, L (g' . f'))

instance CategoryWith1 AD where
    type Terminal AD = ()

    terminal = AD (const ((), terminal))

instance CartesianCategory AD where
    type Product AD = (,)

    proj1 = AD (\x -> (fst x, proj1))
    proj2 = AD (\x -> (snd x, proj2))

    fanout (AD f) (AD g) = AD $ \a ->
        let (b, f') = f a
            (c, g') = g a
        in ((b, c), fanout f' g')

instance GeneralizedElement AD where
    type Object AD a = a

    konst x = AD (\_ -> (x, L $ \_ -> LZ))

ladd :: LinMap r (a, a) -> LinMap r a
ladd (LH f g) = LA f g
ladd (LV f g) = LV (ladd f) (ladd g)
ladd (LA a b) = LA (ladd a) (ladd b)
ladd LZ       = LZ
ladd (LD k)   = LV (LD k) (LD k)

lmult :: Double -> Double -> LinMap r (a, a) -> LinMap r a
lmult x y (LH f g) = LA (lmul y f) (lmul x g)
lmult x y (LV f g) = LV (lmult x y f) (lmult x y g)
lmult x y (LA f g) = LA (lmult x y f) (lmult x y g)
lmult _ _ LZ       = LZ
lmult x y (LD k)   = LV (LD (k * y)) (LD (k * x))

plus :: AD (Double, Double) Double
plus = AD $ \(x,y) -> (x + y, L ladd)

minus :: AD (Double, Double) Double
minus = AD $ \(x,y) -> (x - y, L $ lmult (-1) 1)

mult :: AD (Double, Double) Double
mult = AD $ \(x,y) -> (x * y, L $ lmult x y)

scale :: Double -> AD Double Double
scale k = AD $ \x -> (k * x, linear k)

evaluateAD :: (HasDim a, HasDim b) => AD a b -> a -> (b, LA.Matrix Double)
evaluateAD (AD f) x = let (y, f') = f x in (y, evalL f')

-------------------------------------------------------------------------------
-- Simple examples
-------------------------------------------------------------------------------

ex1 :: AD Double Double
ex1 = plus %% fanout identity identity

ex2 :: AD Double Double
ex2 = mult %% fanout identity identity

-------------------------------------------------------------------------------
-- Quadratic function
-------------------------------------------------------------------------------

quad :: AD (Double, Double) Double
quad = proc (x, y) -> do
    x2  <- mult    -< (x, x)
    y2  <- mult    -< (y, y)
    tmp <- plus    -< (x2, y2)
    z   <- konst 5 -< ()
    plus -< (tmp, z)

-------------------------------------------------------------------------------
-- Newton
-------------------------------------------------------------------------------

findZero :: AD Double Double -> Double -> [Double]
findZero f x0 = take 10 results
  where
    results = iterate go x0

    go :: Double -> Double
    go x =
        let (y, m) = evaluateAD f x
            [[y']] = LA.toLists m
        in x - gamma * (y / y')

    gamma = 0.1

-------------------------------------------------------------------------------
-- Gradient descent
-------------------------------------------------------------------------------

gradDesc :: forall a. VectorSpace a => AD a Double -> a -> [a]
gradDesc f = iterate go where
    go :: a -> a
    go x =
        let (_, m) = evaluateAD f x
            [grad] = LA.toLists $ LA.tr $ LA.scale gamma m

        in fromVector $ zipWith (-) (toVector x) grad

    gamma = 0.1

-------------------------------------------------------------------------------
-- Random
-------------------------------------------------------------------------------

randomDoubles :: Word64 -> [Double]
randomDoubles seed = go (SM.mkSMGen seed) where
    go g = let (d, g') = SM.nextDouble g in d : go g'

-------------------------------------------------------------------------------
-- Dot
-------------------------------------------------------------------------------

class VectorSpace' a where
    sumN :: AD a Double
    multN :: AD (a, a) a

instance (VectorSpace' a, VectorSpace' b) => VectorSpace' (a, b) where
    sumN = proc (x, y) -> do
        x' <- sumN -< x
        y' <- sumN -< y
        plus -< (x', y')

    multN = proc ((x1, x2), (y1, y2)) -> do
        z1 <- multN -< (x1, y1)
        z2 <- multN -< (x2, y2)
        identity -< (z1, z2)

instance VectorSpace' Double where
    sumN  = identity
    multN = mult

dot :: VectorSpace' a => AD (a, a) Double
dot = sumN %% multN

-------------------------------------------------------------------------------
-- ML stuff
-------------------------------------------------------------------------------

tanhAD :: AD Double Double
tanhAD = AD $ \x ->
    let y = tanh x
    in (y, linear (1 - y * y))

sigmoidAD :: AD Double Double
sigmoidAD = AD $ \x ->
    let y = 1 / (1 + exp (- x))
    in (x, linear (y * (1 - y)))

-- | weights for 2x1 connection. Two weights and bias.
type Weights' = ((Double, Double), Double)

-- | Two internal neurons, and final output
type Weights = ((Weights', Weights'), Weights')

startWeights :: Weights
startWeights = fromVector $ randomDoubles 1337

neuron :: AD (Weights', (Double, Double)) Double
neuron = proc ((ws, bias), i) -> do
    o <- dot -< (ws, i)
    tanhAD %% plus -< (o, bias)

network :: AD (Weights, (Double, Double)) Double
network = proc (((w1, w2), w3), xy) -> do
    u <- neuron  -< (w1, xy)
    v <- neuron  -< (w2, xy)
    neuron -< (w3, (u, v))

networkError :: AD Weights Double
networkError = proc ws -> do
    -- xor!
    s1 <- ex 1 1 0 -< ws
    s2 <- ex 0 0 0 -< ws
    s3 <- ex 1 0 1 -< ws
    s4 <- ex 0 1 1 -< ws

    sumN -< ((s1,s2), (s3, s4))
  where
    ex :: Double -> Double -> Double -> AD Weights Double
    ex x y z = proc ws -> do
         x1 <- konst x -< ()
         y1 <- konst y -< ()
         e1 <- konst z -< ()
         a1 <- network -< (ws, (x1, y1))
         r1 <- minus   -< (e1, a1)
         mult -< (r1, r1)

train :: Weights
train = gradDesc networkError startWeights !! 500

-------------------------------------------------------------------------------
-- Main
-------------------------------------------------------------------------------

main :: IO ()
main = do
    putStrLn $ "quad (2,3) = " ++ show (evaluateAD quad (2,3))
    putStrLn $ "gradDesc quad (2,3) = " ++ show (gradDesc quad (2,3) !! 30)

    print $ evaluateAD tanhAD 1
    print $ evaluateAD sigmoidAD 1

    putStrLn "Training the net (for xor)"
    let ws = train
    putStrLn $ "Parameters = " ++ show (toVector ws)
    putStrLn $ "Error = " ++ show (fst $ evaluateAD networkError ws)
    let example xy =
          putStrLn $ "eval " ++ show xy ++ " = " ++ showFFloat (Just 2) (fst $ evaluateAD network (ws, xy)) ""

    example (0, 0)
    example (0, 1)
    example (1, 0)
    example (1, 1)

    args <- getArgs
    when ("plot" `elem` args) $ do
        putStrLn "Outputting plot data: datafile.dat"

        let n = 20 :: Int
        let points = [ fromIntegral x / fromIntegral n | x <- [0..n] ] :: [Double]

        let output :: String
            output = unlines
                [ intercalate "\t"
                    [ show x
                    , show y
                    , show (fst (evaluateAD network (ws, (x, y))))
                    ]
                | x <- points
                , y <- points
                ]

        writeFile "datafile.dat" output


{-# LANGUAGE CPP                    #-}
{-# LANGUAGE ConstraintKinds        #-}
{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE DefaultSignatures      #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE ParallelListComp       #-}
{-# LANGUAGE PartialTypeSignatures  #-}
{-# LANGUAGE PatternGuards          #-}
{-# LANGUAGE Rank2Types             #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TemplateHaskell        #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE TypeApplications       #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE TypeSynonymInstances   #-}
{-# LANGUAGE UndecidableInstances   #-} -- See below

-- experiment
-- {-# LANGUAGE PartialTypeSignatures #-}

{-# OPTIONS_GHC -Wall #-}
-- {-# OPTIONS_GHC -fno-warn-unused-imports #-} -- TEMP

-- #define TESTING

#ifdef TESTING
{-# OPTIONS_GHC -fno-warn-unused-binds   #-} -- TEMP
#endif

-- #define GenericPowFFT

----------------------------------------------------------------------
-- | Generic FFT
----------------------------------------------------------------------

module ConCat.FFT
  ( dft, FFT(..), DFTTy, genericFft, GFFT
  -- -- Temporary while debugging
  , twiddle, twiddles, omega, cis
  , o8sq
  ) where

import Prelude hiding (zipWith)

import Data.Complex
-- import Control.Applicative (liftA2)
import GHC.Generics hiding (C,S)

import Data.Key (Zip(..))
import Data.Pointed

#ifdef TESTING
import Data.Foldable (toList)
-- import Test.QuickCheck (quickCheck)
import Test.QuickCheck.All (quickCheckAll)
import ShapedTypes.ApproxEq
#endif

import ConCat.Misc (Unop,transpose,C,inGeneric1,inComp)
import ConCat.Sized
import ConCat.Scan (LScan,lproducts)
import ConCat.Pair
import ConCat.Free.LinearRow (($*))

{--------------------------------------------------------------------
    DFT
--------------------------------------------------------------------}

type AS  h = (Pointed h, Zip h, LScan h)
type ASZ h = (AS h, Sized h)

-- To resolve: Zip vs Applicative. traverse/transpose needs Applicative.

dft :: forall f a. (ASZ f, Foldable f, RealFloat a) => Unop (f (Complex a))
dft xs = omegas (size @f) $* xs
{-# INLINE dft #-}

omegas :: (AS f, AS g, RealFloat a) => Int -> g (f (Complex a))
omegas = fmap powers . powers . omega
-- omegas n = powers <$> powers (omega n)
-- omegas n = powers <$> powers (exp (- i * 2 * pi / fromIntegral n))
{-# INLINE omegas #-}

omegas' :: forall f g a. (ASZ f, ASZ g, RealFloat a) => g (f (Complex a))
omegas' = fmap powers (powers (cis (- 2 * pi / fromIntegral (size @(g :.: f)))))

-- i :: Num a => Complex a
-- i = 0 :+ 1

omega :: RealFloat a => Int -> Complex a
omega n = cis (- 2 * pi / fromIntegral n)
-- omega n = exp (0 :+ (- 2 * pi / fromIntegral n))
-- omega n = exp (- 2 * (0:+1) * pi / fromIntegral n)
{-# INLINE omega #-}

{--------------------------------------------------------------------
    FFT
--------------------------------------------------------------------}

type DFTTy f = forall a. RealFloat a => f (Complex a) -> FFO f (Complex a)

class FFT f where
  type FFO f :: * -> *
  fft :: DFTTy f
  -- default fft :: ( Generic1 f, Generic1 (FFO f)
  --                , FFT (Rep1 f), FFO (Rep1 f) ~ Rep1 (FFO f) ) => DFTTy f
  -- fft = genericFft
  -- Temporary hack to avoid newtype-like representation.
  fftDummy :: f a
  fftDummy = undefined

-- TODO: Eliminate FFO, in favor of fft :: Unop (f (Complex a)).
-- Use dft as spec.

twiddle :: forall g f a. (ASZ g, ASZ f, RealFloat a) => Unop (g (f (Complex a)))
-- twiddle = (zipWith.zipWith) (*) twiddles
twiddle = (zipWith.zipWith) (*) omegas'
-- twiddle = (zipWith.zipWith) (*) (omegas (size @(g :.: f)))
{-# INLINE twiddle #-}

twiddles :: forall g f a. (ASZ g, ASZ f, RealFloat a) => g (f (Complex a))
twiddles = omegas (size @(g :.: f))
{-# INLINE twiddles #-}

o8sq :: C
o8sq = omega (8 :: Int) ^ (2 :: Int)

-- Powers of x, starting x^0. Uses 'LScan' for log parallel time
powers :: (LScan f, Pointed f, Num a) => a -> f a
powers = fst . lproducts . point
{-# INLINE powers #-}

-- TODO: Consolidate with powers in TreeTest and rename sensibly. Maybe use
-- "In" and "Ex" suffixes to distinguish inclusive and exclusive cases.

{--------------------------------------------------------------------
    Generic support
--------------------------------------------------------------------}

instance FFT Par1 where
  type FFO Par1 = Par1
  fft = id

#if 0
inTranspose :: (Traversable f', Traversable g, Applicative g', Applicative f)
            => (f (g a) -> f' (g' a)) -> g (f a) -> g' (f' a)
inTranspose = transpose <-- transpose

ffts' :: ( FFT g, Traversable f, Traversable g
         , Applicative (FFO g), Applicative f, RealFloat a) =>
     g (f (Complex a)) -> FFO g (f (Complex a))
ffts' = transpose . fmap fft . transpose
#endif

#if 0

transpose :: g (f C)     -> f (g C)
fmap fft  :: f (g C)     -> f (FFO g C)
transpose :: f (FFO g C) -> FFO g (f C)

#endif

instance ( Zip f,  Traversable f , Traversable g
         , Applicative f, Applicative (FFO f), Applicative (FFO g), Zip (FFO g)
         , Pointed f, Traversable (FFO g), Pointed (FFO g)
         , FFT f, FFT g, LScan f, LScan (FFO g), Sized f, Sized (FFO g) )
      => FFT (g :.: f) where
  type FFO (g :.: f) = FFO f :.: FFO g
  fft = inComp (traverse fft . twiddle . traverse fft . transpose)
  -- fft = inComp (ffts' . transpose . twiddle . ffts')
  {-# INLINE fft #-}

#if 0
  fft = Comp1 . transpose . fmap fft . twiddle . transpose . fmap fft . transpose . unComp1
  fft = Comp1 . traverse fft . twiddle . traverse fft . transpose . unComp1

-- Types in fft for (g :. f):
  unComp1   :: (g :. f) a -> g  (f  a)
  transpose :: g  (f  a)  -> f  (g  a)
  fmap fft  :: f  (g  a)  -> f  (g' a)
  transpose :: f  (g' a)  -> g' (f  a)
  twiddle   :: g' (f  a)  -> g' (f  a)
  fmap fft  :: g' (f  a)  -> g' (f' a)
  transpose :: g' (f' a)  -> f' (g' a)
  Comp1     :: f' (g' a)  -> (f' :. g') a
#endif

#if 0

--   fft = inComp (ffts' . transpose . twiddle . ffts')

ffts'     :: g (f C)     -> FFO g (f C)
twiddle   :: FFO g (f C) -> FFO g (f C)
transpose :: FFO g (f C) -> f (FFO g C)
ffts'     :: f (FFO g C) -> FFO f (FFO g C)

#endif

-- -- Generalization of 'dft' to traversables. Note that liftA2 should
-- -- work zippily (unlike with lists).
-- dftT :: forall f a. (ASZ f, Traversable f, RealFloat a)
--      => Unop (f (Complex a))
-- dftT xs = out <$> indices
--  where
--    out k   = sum (liftA2 (\ n x -> x * ok^n) indices xs) where ok = om ^ k
--    indices = fst iota :: f Int
--    om      = omega (size @f)
-- {-# INLINE dftT #-}

-- | Generic FFT
genericFft :: ( Generic1 f, Generic1 (FFO f)
              , FFT (Rep1 f), FFO (Rep1 f) ~ Rep1 (FFO f) ) => DFTTy f
genericFft = inGeneric1 fft

type GFFT f = (Generic1 f, Generic1 (FFO f), FFT (Rep1 f), FFO (Rep1 f) ~ Rep1 (FFO f))

#define GenericFFT(f,g) instance GFFT (f) => FFT (f) where { type FFO (f) = (g); fft = genericFft }

-- #define GenericFFT(f,g) instance GFFT (f) => FFT (f) where { type FFO (f) = g; INLINE }

-- TODO: Replace Applicative with Zippable.
-- Can't, because Traversable needs Applicative.

-- Perhaps dftT isn't very useful. Its result and argument types match, unlike fft.

{--------------------------------------------------------------------
    Specialized FFT instances.
--------------------------------------------------------------------}

-- I put the specific instances here in order to avoid an import loop between
-- the LPow and RPow modules. I'd still like to find an elegant FFT that maps f
-- to f, and then move the instances to RPow and LPow.

-- Radix 2 butterfly
instance FFT Pair where
  type FFO Pair = Pair
  -- bogus "non-exhaustive" warning in ghc 8.0.2
  -- fft (a :# b) = (a + b) :# (a - b)
  fft = \ (a :# b) -> (a + b) :# (a - b)
  -- fft = dft
  {-# INLINE fft #-}

#ifdef TESTING

#if 0
twiddles :: (ASZ g, ASZ f, RealFloat a) => g (f (Complex a))

as :: f C
(<.> as) :: f C -> C
twiddles :: f (f C)
(<.> as) <$> twiddles :: f C
#endif

-- -- Binary dot product
-- infixl 7 <.>
-- (<.>) :: (Foldable f, Applicative f, Num a) => f a -> f a -> a
-- u <.> v = sum (liftA2 (*) u v)

{--------------------------------------------------------------------
    Simple, quadratic DFT (for specification & testing)
--------------------------------------------------------------------}

-- Adapted from Dave's definition
dftL :: RealFloat a => Unop [Complex a]
dftL xs = [ sum [ x * ok^n | x <- xs | n <- [0 :: Int ..] ]
          | k <- [0 .. length xs - 1], let ok = om ^ k ]
 where
   om = omega (length xs)

{--------------------------------------------------------------------
    Tests
--------------------------------------------------------------------}

-- > powers 2 :: LTree N2 Int
-- B (B (L ((1 :# 2) :# (4 :# 8))))
-- > powers 2 :: LTree N3 Int
-- B (B (B (L (((1 :# 2) :# (4 :# 8)) :# ((16 :# 32) :# (64 :# 128))))))

fftl :: (FFT f, Foldable (FFO f), RealFloat a) => f (Complex a) -> [Complex a]
fftl = toList . fft

type LTree = L.Pow Pair
type RTree = R.Pow Pair

type LC n = LTree n C
type RC n = RTree n C

p1 :: Pair C
p1 = 1 :# 0

tw1 :: LTree N1 (Pair C)
tw1 = twiddles

tw2 :: RTree N2 (Pair C)
tw2 = twiddles

tw3 :: RTree N2 (RTree N2 C)
tw3 = twiddles

tw3' :: [[C]]
tw3' = toList (toList <$> tw3)

-- Adapted from Dave's testing

-- test :: (FFT f, Foldable f, Foldable (FFO f)) => f C -> IO ()
-- test fx =
--   do ps "\nTesting input" xs
--      ps "Expected output" (dftL xs)
--      ps "Actual output  " (toList (fft fx))
--  where
--    ps label z = putStrLn (label ++ ": " ++ show z)
--    xs = toList fx

#if 0
t0 :: LC N0
t0 = L.fromList [1]

t1 :: LC N1
t1 = L.fromList [1, 0]

t2s :: [LC N2]
t2s = L.fromList <$>
        [ [1,  0,  0,  0]  -- Delta
        , [1,  1,  1,  1]  -- Constant
        , [1, -1,  1, -1]  -- Nyquist
        , [1,  0, -1,  0]  -- Fundamental
        , [0,  1,  0, -1]  -- Fundamental w/ 90-deg. phase lag
       ]

tests :: IO ()
tests = do test p1
           test t0
           test t1
           mapM_ test t2s
#endif

infix 4 ===
(===) :: Eq b => (a -> b) -> (a -> b) -> a -> Bool
(f === g) x = f x == g x

infix 4 =~=
(=~=) :: ApproxEq b => (a -> b) -> (a -> b) -> a -> Bool
(f =~= g) x = f x =~ g x

fftIsDftL :: (FFT f, Foldable f, Foldable (FFO f), RealFloat a, ApproxEq a) =>
             f (Complex a) -> Bool
fftIsDftL = toList . fft =~= dftL . toList

dftTIsDftL :: (ASZ f, Traversable f, RealFloat a, ApproxEq a) =>
              f (Complex a) -> Bool
dftTIsDftL = toList . dftT =~= dftL . toList

dftIsDftL :: (ASZ f, Foldable f, RealFloat a, ApproxEq a) =>
             f (Complex a) -> Bool
dftIsDftL = toList . dft =~= dftL . toList

-- -- TEMP:
-- dftDft :: (ASZ f, Traversable f, RealFloat a, ApproxEq a) =>
--           f (Complex a) -> ([Complex a], [Complex a])
-- dftDft xs = (toList . dft $ xs, dftL . toList $ xs)

{--------------------------------------------------------------------
    Properties to test
--------------------------------------------------------------------}

transposeTwiddleCommutes :: (ASZ g, Traversable g, ASZ f, (ApproxEq (f (g C))))
                         => g (f C) -> Bool
transposeTwiddleCommutes =
 twiddle . transpose =~= transpose . twiddle

prop_transposeTwiddle_L3P :: LTree N3 (Pair C) -> Bool
prop_transposeTwiddle_L3P = transposeTwiddleCommutes

prop_transposeTwiddle_R3P :: RTree N3 (Pair C) -> Bool
prop_transposeTwiddle_R3P = transposeTwiddleCommutes

-- dft tests fail. Hm!

-- prop_dft_R3 :: RTree N3 C -> Bool
-- prop_dft_R3 = dftIsDftL

-- prop_dft_L3 :: LTree N3 C -> Bool
-- prop_dft_L3 = dftIsDftL

prop_dftT_p :: Pair C -> Bool
prop_dftT_p = dftTIsDftL

prop_dftT_L3 :: LTree N3 C -> Bool
prop_dftT_L3 = dftTIsDftL

prop_dftT_R3 :: RTree N3 C -> Bool
prop_dftT_R3 = dftTIsDftL

prop_fft_p :: Pair C -> Bool
prop_fft_p = fftIsDftL

prop_fft_L1 :: LTree N1 C -> Bool
prop_fft_L1 = fftIsDftL

prop_fft_L2 :: LTree N2 C -> Bool
prop_fft_L2 = fftIsDftL

prop_fft_L3 :: LTree N3 C -> Bool
prop_fft_L3 = fftIsDftL

prop_fft_L4 :: LTree N4 C -> Bool
prop_fft_L4 = fftIsDftL

prop_fft_R1 :: RTree N1 C -> Bool
prop_fft_R1 = fftIsDftL

prop_fft_R2 :: RTree N2 C -> Bool
prop_fft_R2 = fftIsDftL

prop_fft_R3 :: RTree N3 C -> Bool
prop_fft_R3 = fftIsDftL

prop_fft_R4 :: RTree N4 C -> Bool
prop_fft_R4 = fftIsDftL

-- TH oddity
return []

runTests :: IO Bool
runTests = $quickCheckAll

-- end of tests
#endif


module TextHistogram (
  printHistogram,
) where

import           Control.Monad (forM_, when)
import           Statistics.Test.ApproxRand
import           System.IO (hPutStrLn, stderr)
import           Text.Printf (printf)

import           Histogram

printHistogram :: Int -> TestResult -> IO ()
printHistogram bins result@(TestResult _ score _) =
  case histogram bins result of
    Left err   -> hPutStrLn stderr err
    Right hist ->
      let bucketSize = (fst $ hist !! 1) - (fst $ head hist)
          bucketHalf = bucketSize / 2
          charsPerDot = (maximum $ map snd hist) `div` 50 in
            forM_ hist $ \(label, freq) -> do
              let blocks = freq `div` charsPerDot
              when (blocks > 0) $ do
                let lower = label - bucketHalf
                let barChar = if score >= lower && score < lower + bucketSize then
                            '✣'
                          else
                            '█'
                putStr $ printf "%12.3e | " $ lower + bucketHalf
                putStrLn $ replicate blocks barChar


module STCR2Z2T0S0Edge where

import           Control.Arrow
import           Control.Monad
import           Data.Array.Repa         as R
import           Data.Binary             (decodeFile)
import           Data.Complex
import           Data.List               as L
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.Edge
import           Image.IO
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:numTrailStr:maxTrailStr:thetaFreqsStr:scaleFreqsStr:histFilePath:numIterationStr:writeSourceFlagStr:edgeFilePath:numNoisePointStr:scaleFactorStr:useFFTWWisdomFlagStr:fftwWisdomFileName:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Double
      scaleFreqs = [-scaleFreq .. scaleFreq]
      numIteration = read numIterationStr :: Int
      writeSourceFlag = read writeSourceFlagStr :: Bool
      numNoisePoint = read numNoisePointStr :: Int
      scaleFactor = read scaleFactorStr :: Double
      useFFTWWisdomFlag = read useFFTWWisdomFlagStr :: Bool
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2Z2T0S0Edge"
      fftwWisdomFilePath = folderPath </> fftwWisdomFileName
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z2T0S0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          tao
          thetaFreqs
          thetaFreqs
          scaleFreqs
          scaleFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length scaleFreqs
             , L.length thetaFreqs
             , L.length scaleFreqs
             , L.length thetaFreqs
             ]
             0)
  arrR2Z2T0S0 <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      (PinwheelHollow0 4)
      radialArr
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      thetaFreqs
      scaleFreqs
  plan <-
    makeR2Z2T0S0Plan emptyPlan useFFTWWisdomFlag fftwWisdomFilePath arrR2Z2T0S0
  xs <- parseEdgeFile edgeFilePath
  randomPonintSet <- generateRandomPointSet numNoisePoint numPoint numPoint
  let (centerX, centerY) =
        join (***) (\x -> div x . L.length $ xs) .
        L.foldl' (\(a, b) (R2S1RPPoint (c, d, _, _)) -> (a + c, b + d)) (0, 0) $
        xs
      ys =
        L.map
          (\(R2S1RPPoint (a, b, c, d)) ->
             (R2S1RPPoint
                ( round $ (fromIntegral $ a - centerX) / scaleFactor
                , round $ (fromIntegral $ b - centerX) / scaleFactor
                , c
                , d)))
          xs
      zs =
        L.map
          (\(R2S1RPPoint (a, b, c, d)) ->
             (R2S1RPPoint (a - center numPoint, b - center numPoint, c, d)))
          randomPonintSet
      points = ys L.++ zs
  let bias = computeBiasR2T0S0 numPoint numPoint thetaFreqs scaleFreqs points
      eigenVec =
        computeInitialEigenVectorR2T0S0
          numPoint
          numPoint
          thetaFreqs
          scaleFreqs
          thetaFreqs
          scaleFreqs
          points
  powerMethodR2Z2T0S0
    plan
    folderPath
    numPoint
    numPoint
    numOrientation
    thetaFreqs
    numScale
    scaleFreqs
    maxScale
    arrR2Z2T0S0
    numIteration
    writeSourceFlag
    (printf
       "_%d_%d_%.2f_%.2f"
       (round maxScale :: Int)
       (round tao :: Int)
       thetaSigma
       scaleSigma)
    bias
    eigenVec


{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MagicHash                  #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE TypeInType                 #-}
{-# LANGUAGE UnboxedTuples              #-}
{-# LANGUAGE UndecidableInstances       #-}
{-# LANGUAGE ViewPatterns               #-}
{-# OPTIONS_GHC -fno-warn-orphans       #-}
module Numeric.Quaternion.Internal.QDouble
    ( QDouble, Quater (..)
    ) where

import qualified Control.Monad.ST                     as ST
import           Data.Coerce                          (coerce)
import           Numeric.Basics
import           Numeric.DataFrame.Internal.PrimArray
import qualified Numeric.DataFrame.ST                 as ST
import           Numeric.DataFrame.Type
import           Numeric.Quaternion.Internal
import           Numeric.Vector.Internal

type QDouble = Quater Double

deriving instance PrimBytes (Quater Double)
deriving instance PrimArray Double (Quater Double)

instance Quaternion Double where
    newtype Quater Double = QDouble (Vector Double 4)
    {-# INLINE packQ #-}
    packQ = coerce (vec4 :: Double -> Double -> Double -> Double -> Vector Double 4)
    {-# INLINE unpackQ# #-}
    unpackQ# = coerce (unpackV4# :: Vector Double 4 -> (# Double, Double, Double, Double #))
    {-# INLINE fromVecNum #-}
    fromVecNum (unpackV3# -> (# x, y, z #))  = packQ x y z
    {-# INLINE fromVec4 #-}
    fromVec4 = coerce
    {-# INLINE toVec4 #-}
    toVec4 = coerce
    {-# INLINE square #-}
    square (unpackQ# -> (# x, y, z, w #)) = (x * x) + (y * y) + (z * z) + (w * w)
    {-# INLINE im #-}
    im (unpackQ# -> (# x, y, z, _ #)) = packQ x y z 0.0
    {-# INLINE re #-}
    re (unpackQ# -> (# _, _, _, w #)) = packQ 0 0 0 w
    {-# INLINE imVec #-}
    imVec (unpackQ# -> (# x, y, z, _ #)) = vec3 x y z
    {-# INLINE taker #-}
    taker (unpackQ# -> (# _, _, _, w #)) = w
    {-# INLINE takei #-}
    takei (unpackQ# -> (# x, _, _, _ #)) = x
    {-# INLINE takej #-}
    takej (unpackQ# -> (# _, y, _, _ #)) = y
    {-# INLINE takek #-}
    takek (unpackQ# -> (# _, _, z, _ #)) = z
    {-# INLINE conjugate #-}
    conjugate (unpackQ# -> (# x, y, z, w #))
      = packQ (negate x) (negate y) (negate z) w
    {-# INLINE rotScale #-}
    rotScale (unpackQ# -> (# i, j, k, t #))
             (unpackV3# -> (# x, y, z #))
      = let l = t*t - i*i - j*j - k*k
            d = 2.0 * ( i*x + j*y + k*z)
            t2 = t * 2.0
        in vec3 (l*x + d*i + t2 * (z*j - y*k))
                (l*y + d*j + t2 * (x*k - z*i))
                (l*z + d*k + t2 * (y*i - x*j))
    {-# INLINE getRotScale #-}
    getRotScale a b = case (# unpackV3# a, unpackV3# b #) of
      (# _, (# 0, 0, 0 #) #) -> packQ 0 0 0 0
      (# (# 0, 0, 0 #), _ #) -> let x = (1 / 0 :: Double) in packQ x x x x
      (# (# a1, a2, a3 #), (# b1, b2, b3 #) #) ->
        let ma = sqrt (a1*a1 + a2*a2 + a3*a3)
            mb = sqrt (b1*b1 + b2*b2 + b3*b3)
            d  = a1*b1 + a2*b2 + a3*b3
            c  = sqrt (ma*mb + d)
            ma2 = ma * 1.4142135623730951 -- sqrt 2.0
            r  = recip (ma2 * c)
            c' = sqrt (mb / ma) -- ratio of a and b for corner cases
            r' = recip (sqrt ( negate (a1*b1 + a2*b2) ))
        in case unpackV3# (cross a b) of
          (# 0, 0, 0 #)
              -- if a and b face the same direction, q is fully real
            | d >= 0       -> packQ 0 0 0 c'
              -- if a and b face opposite directions, find an orthogonal vector
              -- prerequisites: w == 0  and  a·(x,y,z) == 0
              -- corner cases: only one vector component is non-zero
            | b1 == 0      -> packQ c' 0 0 0
              -- otherwise set the last component to zero,
              -- and get an orthogonal vector in 2D.
            | otherwise    -> packQ (-b2*r') (b1*r') 0 0
              -- NB: here we have some precision troubles
              --     when a and b are close to parallel and opposite.
          (# t1, t2, t3 #) -> packQ (t1 * r) (t2 * r) (t3 * r) (c / ma2)
    {-# INLINE axisRotation #-}
    axisRotation v a = case unpackV3# v of
      (# 0, 0, 0 #) -> packQ 0 0 0 (negateUnless (abs a < M_PI) 1)
      (# x, y, z #) ->
        let c = cos (a * 0.5)
            s = sin (a * 0.5)
                / sqrt (x*x + y*y + z*z)
        in packQ (x * s) (y * s) (z * s) c
    {-# INLINE qArg #-}
    qArg (unpackQ# -> (# x, y, z, w #)) = 2 * atan2 (sqrt (x*x + y*y + z*z)) w
    {-# INLINE fromMatrix33 #-}
    fromMatrix33 m = fromM 1
      (ix# 0# m) (ix# 1# m) (ix# 2# m)
      (ix# 3# m) (ix# 4# m) (ix# 5# m)
      (ix# 6# m) (ix# 7# m) (ix# 8# m)

    {-# INLINE fromMatrix44 #-}
    fromMatrix44 m = fromM (ix# 15# m)
      (ix# 0# m) (ix# 1# m) (ix# 2# m)
      (ix# 4# m) (ix# 5# m) (ix# 6# m)
      (ix# 8# m) (ix# 9# m) (ix# 10# m)

    {-# INLINE toMatrix33 #-}
    toMatrix33 (unpackQ# -> (# 0.0, 0.0, 0.0, w #))
      = let x = w * w
            f 0 = (# 3 :: Int , x #)
            f k = (# k-1, 0 #)
        in case gen# (CumulDims [9,3,1]) f 0 of
            (# _, m #) -> m -- diag (scalar (w * w))
    toMatrix33 (unpackQ# -> (# x', y', z', w' #)) =
      let x = scalar x'
          y = scalar y'
          z = scalar z'
          w = scalar w'
          x2 = x * x
          y2 = y * y
          z2 = z * z
          w2 = w * w
          l2 = x2 + y2 + z2 + w2
      in ST.runST $ do
        df <- ST.newDataFrame
        ST.writeDataFrameOff df 0 $ l2 - 2*(z2 + y2)
        ST.writeDataFrameOff df 1 $ 2*(x*y + z*w)
        ST.writeDataFrameOff df 2 $ 2*(x*z - y*w)
        ST.writeDataFrameOff df 3 $ 2*(x*y - z*w)
        ST.writeDataFrameOff df 4 $ l2 - 2*(z2 + x2)
        ST.writeDataFrameOff df 5 $ 2*(y*z + x*w)
        ST.writeDataFrameOff df 6 $ 2*(x*z + y*w)
        ST.writeDataFrameOff df 7 $ 2*(y*z - x*w)
        ST.writeDataFrameOff df 8 $ l2 - 2*(y2 + x2)
        ST.unsafeFreezeDataFrame df
    {-# INLINE toMatrix44 #-}
    toMatrix44 (unpackQ# -> (# 0.0, 0.0, 0.0, w #)) = ST.runST $ do
      df <- ST.newDataFrame
      mapM_ (flip (ST.writeDataFrameOff df) 0) [0..15]
      let w2 = scalar (w * w)
      ST.writeDataFrameOff df 0 w2
      ST.writeDataFrameOff df 5 w2
      ST.writeDataFrameOff df 10 w2
      ST.writeDataFrameOff df 15 1
      ST.unsafeFreezeDataFrame df
    toMatrix44 (unpackQ# -> (# x', y', z', w' #)) =
      let x = scalar x'
          y = scalar y'
          z = scalar z'
          w = scalar w'
          x2 = x * x
          y2 = y * y
          z2 = z * z
          w2 = w * w
          l2 = x2 + y2 + z2 + w2
      in ST.runST $ do
        df <- ST.newDataFrame
        ST.writeDataFrameOff df 0 $ l2 - 2*(z2 + y2)
        ST.writeDataFrameOff df 1 $ 2*(x*y + z*w)
        ST.writeDataFrameOff df 2 $ 2*(x*z - y*w)
        ST.writeDataFrameOff df 3 0
        ST.writeDataFrameOff df 4 $ 2*(x*y - z*w)
        ST.writeDataFrameOff df 5 $ l2 - 2*(z2 + x2)
        ST.writeDataFrameOff df 6 $ 2*(y*z + x*w)
        ST.writeDataFrameOff df 7 0
        ST.writeDataFrameOff df 8 $ 2*(x*z + y*w)
        ST.writeDataFrameOff df 9 $ 2*(y*z - x*w)
        ST.writeDataFrameOff df 10 $ l2 - 2*(y2 + x2)
        ST.writeDataFrameOff df 11 0
        ST.writeDataFrameOff df 12 0
        ST.writeDataFrameOff df 13 0
        ST.writeDataFrameOff df 14 0
        ST.writeDataFrameOff df 15 1
        ST.unsafeFreezeDataFrame df

{- Calculate quaternion from a 3x3 matrix.

   First argument is a constant; it is either 1 for a 3x3 matrix,
   or m44 for a 4x4 matrix. I just need to multiply all components by
   this number.

   Further NB for the formulae:

   d == square q == det m ** (1/3)
   t == trace m  == 4 w w - d
   m01 - m10 == 4 z w
   m20 - m02 == 4 y w
   m12 - m21 == 4 x w
   m01 + m10 == 4 x y
   m20 + m02 == 4 x z
   m12 + m21 == 4 y z
   m00 == + x x - y y - z z + w w
   m11 == - x x + y y - z z + w w
   m22 == - x x - y y + z z + w w
   4 x x == d + m00 - m11 - m22
   4 y y == d - m00 + m11 - m22
   4 z z == d - m00 - m11 + m22
   4 w w == d + m00 + m11 + m22
 -}
fromM :: Double
      -> Double -> Double -> Double
      -> Double -> Double -> Double
      -> Double -> Double -> Double
      -> QDouble
fromM c'
  m00 m01 m02
  m10 m11 m12
  m20 m21 m22
    | t > 0
      = let dd = sqrt ( d + t )
            is = c / dd
        in packQ ((m12 - m21)*is) ((m20 - m02)*is) ((m01 - m10)*is) (c*dd)
    | m00 > m11 && m00 > m22
      = let dd = sqrt ( d + m00 - m11 - m22 )
            is = c / dd
        in packQ (c*dd) ((m01 + m10)*is) ((m02 + m20)*is) ((m12 - m21)*is)
    | m11 > m22
      = let dd = sqrt ( d - m00 + m11 - m22 )
            is = c / dd
        in packQ ((m01 + m10)*is) (c*dd) ((m12 + m21)*is) ((m20 - m02)*is)
    | otherwise
      = let dd = sqrt ( d - m00 - m11 + m22 )
            is = c / dd
        in packQ ((m02 + m20)*is) ((m12 + m21)*is) (c*dd) ((m01 - m10)*is)

  where
    -- normalizing constant
    c = recip $ 2 * sqrt c'
    -- trace
    t = m00 + m11 + m22
    -- cubic root of determinant
    d = ( m00 * ( m11 * m22 - m12 * m21 )
        - m01 * ( m10 * m22 - m12 * m20 )
        + m02 * ( m10 * m21 - m11 * m20 )
        ) ** 0.33333333333333333333333333333333

instance Num QDouble where
    QDouble a + QDouble b
      = QDouble (a + b)
    {-# INLINE (+) #-}
    QDouble a - QDouble b
      = QDouble (a - b)
    {-# INLINE (-) #-}
    (unpackQ# -> (# a1, a2, a3, a4 #)) * (unpackQ# -> (# b1, b2, b3, b4 #))
      = packQ ((a4 * b1) + (a1 * b4) + (a2 * b3) - (a3 * b2))
              ((a4 * b2) - (a1 * b3) + (a2 * b4) + (a3 * b1))
              ((a4 * b3) + (a1 * b2) - (a2 * b1) + (a3 * b4))
              ((a4 * b4) - (a1 * b1) - (a2 * b2) - (a3 * b3))
    {-# INLINE (*) #-}
    negate (QDouble a) = QDouble (negate a)
    {-# INLINE negate #-}
    abs = packQ 0 0 0 . sqrt . square
    {-# INLINE abs #-}
    signum q@(unpackQ# -> (# x, y, z, w #))
      | qd == 0   = q
      | otherwise = case ix + iy + iz + iw + nn of
        0 -> packQ (x * l) (y * l) (z * l) (w * l)
        1 -> packQ (copysign 1 x) 0 0 0
        2 -> packQ 0 (copysign 1 y) 0 0
        4 -> packQ 0 0 (copysign 1 z) 0
        8 -> packQ 0 0 0 (copysign 1 w)
        _ -> packQ n n n n
      where
        n  = 0 / 0 :: Double
        qd = x*x + y*y + z*z + w*w
        ix = if isInfinite x then 1 else 0 :: Int
        iy = if isInfinite y then 2 else 0 :: Int
        iz = if isInfinite z then 4 else 0 :: Int
        iw = if isInfinite w then 8 else 0 :: Int
        nn = if isNaN x || isNaN y || isNaN z || isNaN w then 16 else 0 :: Int
        l  = recip (sqrt qd)
    {-# INLINE signum #-}
    fromInteger = packQ 0 0 0 . fromInteger
    {-# INLINE fromInteger #-}

instance Fractional QDouble where
    {-# INLINE recip #-}
    recip q@(unpackQ# -> (# x, y, z, w #)) = case negate (recip (square q)) of
      c -> packQ (x * c) (y * c) (z * c) (negate (w * c))
    {-# INLINE (/) #-}
    a / b = a * recip b
    {-# INLINE fromRational #-}
    fromRational = packQ 0 0 0 . fromRational

instance Floating QDouble where
    {-# INLINE pi #-}
    pi = packQ 0 0 0 M_PI
    {-# INLINE exp #-}
    exp (unpackQ# -> (# x, y, z, w #))
      | mv2 == 0  = packQ x y z ew
      | otherwise = packQ (x * l) (y * l) (z * l) arg
      where
        mv2 = (x * x) + (y * y) + (z * z)
        mv  = sqrt mv2
        ew  = exp w
        l   = ew * sin mv / mv
        arg = ew * cos mv
    {-# INLINE log #-}
    log = log' (Vec3 1 0 0)
    {-# INLINE sqrt #-}
    sqrt = sqrt' (Vec3 1 0 0)
    {-# INLINE sin #-}
    sin (unpackQ# -> (# x, y, z, w #))
      | mv2 == 0  = packQ x y z (sin w)
      | otherwise = packQ (x * l) (y * l) (z * l) arg
      where
        mv2 = (x * x) + (y * y) + (z * z)
        mv  = sqrt mv2
        l   = cos w * sinh mv / mv
        arg = sin w * cosh mv
    {-# INLINE cos #-}
    cos (unpackQ# -> (# x, y, z, w #))
      | mv2 == 0  = packQ x y z (cos w)
      | otherwise = packQ (x * l) (y * l) (z * l) arg
      where
        mv2 = (x * x) + (y * y) + (z * z)
        mv  = sqrt mv2
        l   = sin w * sinh mv / negate mv
        arg = cos w * cosh mv
    {-# INLINE tan #-}
    tan (unpackQ# -> (# x, y, z, w #))
      | mv2 == 0       = packQ x y z (tan w)
      | isInfinite mv2 = signum (packQ x y z 0)
      | otherwise      = packQ (x * l) (y * l) (z * l) arg
      where
        mv2 = (x * x) + (y * y) + (z * z)
        mv = sqrt mv2
        b = 2*mv
        a = 2*w
        sina = sin a
        eb = exp (-b)
        eb2 = eb*eb
        d = 1 + eb2 + 2 * eb * cos a
        rd = recip d
        pa = M_PI - abs a
        rd' = 2 / (b*b + pa*pa)
        (l, arg) =
          if d >= M_EPS
          then ((1 - eb2) * rd / mv, 2 * sina * eb * rd)
          else (2 * rd' , negate sina * rd')
    {-# INLINE sinh #-}
    sinh (unpackQ# -> (# x, y, z, w #))
      | mv2 == 0  = packQ x y z (sinh w)
      | otherwise = packQ (x * l) (y * l) (z * l) arg
      where
        mv2 = (x * x) + (y * y) + (z * z)
        mv  = sqrt mv2
        l   = cosh w * sin mv / mv
        arg = sinh w * cos mv
    {-# INLINE cosh #-}
    cosh (unpackQ# -> (# x, y, z, w #))
      | mv2 == 0  = packQ x y z (cosh w)
      | otherwise = packQ (x * l) (y * l) (z * l) arg
      where
        mv2 = (x * x) + (y * y) + (z * z)
        mv  = sqrt mv2
        l   = sinh w * sin mv / mv
        arg = cosh w * cos mv
    {-# INLINE tanh #-}
    tanh (unpackQ# -> (# x, y, z, w #))
      | mv2 == 0       = packQ x y z (tanh w)
      | isInfinite mv2 = packQ 0 0 0 (signum w)
      | otherwise      = packQ (x * l) (y * l) (z * l) arg
      where
        mv2 = (x * x) + (y * y) + (z * z)
        mv = sqrt mv2
        b = 2*w
        a = 2*mv
        eb = exp (- abs b)
        eb2 = eb*eb
        d = 1 + eb2 + 2 * eb * cos a
        rd = recip d
        pa = M_PI - a
        rd' = 2 / (b*b + pa*pa)
        (l, arg) =
          if d >= M_EPS
          then (2 * sin a * eb * rd / mv, copysign (1 - eb2) b * rd)
          else (2 * rd' , b * rd')
    {-# INLINE asin #-}
    -- The original formula:
    -- asin q = -i * log (i*q + sqrt (1 - q*q))
    -- below is a more numerically stable version.
    asin (unpackQ# -> (# x, y, z, w #))
      | v2 == 0   = if w2 <= 1
                    then packQ x y z (asin w)
                    else packQ l 0 0 arg
      | otherwise  = packQ (x*c) (y*c) (z*c) arg
      where
        v2 = (x * x) + (y * y) + (z * z)
        v = sqrt v2
        w2 = w*w
        w1qq = 0.5 *(1 - w2 + v2)       -- real part of (1 - q*q)/2
        l1qq = sqrt (w1qq*w1qq + w2*v2) -- length of (1 - q*q)/2
        sp2 = l1qq + w1qq
        sn2 = l1qq - w1qq
        sp = sqrt sp2
        sn = copysign (sqrt sn2) w
        -- choose a more stable (symbolically equiv) version
        dp = if 2 * v2 <= sp2 then sp - v else v2 / ((sp + v)*(sn2 + v2))
        dn = if 2 * w2 <= sn2 then w - sn else w2 / ((sn + w)*(sp2 + w2))
        (wD, vD) = case compare w1qq 0 of
            GT -> (dp, w * dp / sp)
            LT -> (v * dn / sn, dn)
            EQ -> (-v, w)
        l = -0.5 * log (wD*wD + vD*vD)
        c = l / v
        arg = atan2 vD wD
    {-# INLINE acos #-}
    acos q = M_PI_2 - asin q
    {-# INLINE atan #-}
    -- atan q = i / 2 * log ( (i + q) / (i - q) )
    atan (unpackQ# -> (# x, y, z, w #))
      | v2 == 0   = packQ x y z (atan w)
      | otherwise = packQ (x*c) (y*c) (z*c) arg
      where
        v2 = (x * x) + (y * y) + (z * z)
        v = sqrt v2
        w2 = w*w
        q2 = w2 + v2
        v' = v - 1
        mzero = w2 + v'*v'
        (c, arg) =
          if mzero == 0
          then ( sqrt maxFinite / v, 0)
          else ( 0.25 * (log (1 + q2 + 2*v) - log mzero) / v
               , 0.5 * atan2 (2*w) (1 - q2) )
    {-# INLINE asinh #-}
    -- The original formula:
    -- asinh q = log (q + sqrt (q*q + 1))
    -- below is a more numerically stable version.
    asinh (unpackQ# -> (# x, y, z, w #))
      | v2 == 0   = packQ x y z (asinh w)
      | otherwise = packQ (x*c) (y*c) (z*c) arg
      where
        v2 = (x * x) + (y * y) + (z * z)
        v = sqrt v2
        w2 = w*w
        w1qq = 0.5 *(1 + w2 - v2)       -- real part of (1 + q*q)/2
        l1qq = sqrt (w1qq*w1qq + w2*v2) -- length of (1 + q*q)/2
        sp2 = l1qq + w1qq
        sn2 = l1qq - w1qq
        sp = sqrt sp2
        sn = copysign (sqrt sn2) w
        -- choose a more stable (symbolically equiv) version
        dp = if 2 * w >= - sp then w + sp else w2 / ((sp - w)*(w2 + sn2))
        dn = if 2 * v <= - sn || sn >= 0
                              then v + sn else v2 / ((v - sn)*(v2 + sp2))
        (wD, vD) = case compare w1qq 0 of
            GT -> (dp, v * dp / sp)
            LT -> (w * dn / sn, dn)
            EQ -> (w, v)
        c = atan2 vD wD / v
        arg = 0.5 * log (wD*wD + vD*vD)
    {-# INLINE acosh #-}
    -- The original formula:
    -- asinh q = log (q + sqrt (q + 1) * sqrt (q - 1))
    -- below is a more numerically stable version.
    -- note, log (q + sqrt (q*q - 1)) would not work, because that would not
    -- be the principal value.
    acosh (unpackQ# -> (# x, y, z, w #))
      | v2 == 0   = packQ x y z (acosh w)
      | otherwise = packQ (x*c) (y*c) (z*c) arg
      where
        v2 = (x * x) + (y * y) + (z * z)
        v = sqrt v2
        w2 = w*w
        w1qq = 0.5 *(w2 - v2 - 1)       -- real part of (q*q - 1)/2
        l1qq = sqrt (w1qq*w1qq + w2*v2) -- length of (q*q - 1)/2
        sp2 = l1qq + w1qq
        sn2 = l1qq - w1qq
        sp = sqrt sp2
        sn = copysign (sqrt sn2) w
        -- choose a more stable (symbolically equiv) version
        dp = if 2 * w >= - sp then w + sp else w2 / ((w - sp)*(w2 + sn2))
        dn = if 2 * v <= - sn || sn >= 0
                              then v + sn else v2 / ((sn - v)*(v2 + sp2))
        (wD, vD) = case compare w1qq 0 of
            GT -> (dp, v * dp / sp)
            LT -> (w * dn / sn, dn)
            EQ -> (w, v)
        c = atan2 vD wD / v
        arg = 0.5 * log (wD*wD + vD*vD)
    {-# INLINE atanh #-}
    -- atanh q = 0.5 * log ( (1 + q) / (1 - q) )
    atanh (unpackQ# -> (# x, y, z, w #))
      | v2 ==  0  = packQ x y z (atanh w)
      | otherwise = packQ (x*c) (y*c) (z*c) (copysign arg w)
      where
        v2 = (x * x) + (y * y) + (z * z)
        v = sqrt v2
        w2 = w*w
        q2 = w2 + v2
        w' = abs w - 1
        c  = 0.5 * atan2 (2*v) (1 - q2) / v
        arg = if w' == 0
              then (1/0)
              else 0.25 * (log (1 + q2 + 2 * abs w) - log (v2 + w'*w'))

-- If q is negative real, provide a fallback axis to align log.
log' :: Vector Double 3 -> QDouble -> QDouble
log' r (unpackQ# -> (# x, y, z, w #))
  = case (x * x) + (y * y) + (z * z) of
    0.0 | w >= 0
           -> packQ 0 0 0 (log w)
        | Vec3 rx ry rz <- r
           ->  packQ (M_PI*rx) (M_PI*ry) (M_PI*rz) (log (negate w))
    mv2 -> case (# mv2 + w * w, sqrt mv2 #) of
      (# q2, mv #) -> case atan2 mv w / mv of
        l -> packQ (x * l) (y * l) (z * l) (0.5 * log q2)

-- If q is negative real, provide a fallback axis to align sqrt.
sqrt' :: Vector Double 3 -> QDouble -> QDouble
sqrt' r (unpackQ# -> (# x, y, z, w #))
  | v2 == 0 && w >= 0
    = packQ x y z (sqrt w)
  | v2 == 0
  , Vec3 rx ry rz <- r
  , sw <- sqrt (negate w)
    = packQ (sw*rx) (sw*ry) (sw*rz) 0
  | otherwise
    = packQ (x * c) (y * c) (z * c) arg
  where
    v2 = (x * x) + (y * y) + (z * z)
    mq = sqrt (v2 + w * w)
    arg = sqrt $ 0.5 * if w >= 0 then mq + w else v2 / (mq - w)
    c = 0.5 / arg

instance Eq QDouble where
    {-# INLINE (==) #-}
    QDouble a == QDouble b = a == b
    {-# INLINE (/=) #-}
    QDouble a /= QDouble b = a /= b


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE FlexibleContexts #-}

module CPVO.IO.Reader.Ecalj.MMOM where
import CPVO.Numeric
import CPVO.IO
import CPVO.IO.Reader.Ecalj.Common
import CPVO.IO.Reader.Ecalj.DOS

import qualified Data.Text as T
import qualified Data.Text.IO as T
import qualified Data.Text.Read as T
import Data.Either (rights)
-------------------------
import Numeric.LinearAlgebra

readMMOM :: Int -> String -> IO [Double]
readMMOM nAtom foldernya = do
    fLLMF <- fmap (T.unpack . head) $ getLastLLMF foldernya
    putStrLn $ "===================processed LLMF=" ++ fLLMF
    mmom <- fmap (map T.double) $ inshell2text $ concat [ "mkdir -p temp; grep mmom ", fLLMF
                                            ,"| tail -n", show (nAtom + 1)
                                            ,"| head -n", show nAtom
                                            ,"| awk '{print $2}'"
                                          ]
    sdtMMOM <- fmap (map T.double) $ inshell2text $ concat [ "grep mmom ", fLLMF, "| grep ehf | tail -1 | sed -e 's/^.*mmom=//g'| awk '{print $1}'"
                                          ]
    return ( map fst $ rights $ concat [sdtMMOM,mmom])

----------------------------------------------------------------------
--getMMOM allArgs@(texFile:jd:jdHead:colAlign:xr:ymax':wTot:tumpuk:invS:tailer:foldernya:aos) = do

getMMOM :: [String] -> IO ()
getMMOM allArgs = do
--getMMOM allArgs = do
    putStrLn "===start ==== CPVO.IO.Reader.Ecalj.MMOM: getMMOM ==="
    --(invStat, ymax, xmin, xmax, ctrlAtoms, uniqAtoms, ctrlAtomicAOs,jdTable, cleanedJdHead, foldernya, tailer,colAlign,texFile) <- readHeaderData allArgs
    Right (invStat,_,_,_,ctrlAtoms,_, ctrlAtomicAOs,jdTable, cleanedJdHead,foldernya,tailer,_,texFile) <- readHeaderData allArgs

    -------------------------------generating data------------------------
    -------------------------------generating DOS data------------------------
    totalDOS <- readTotalDOSText tailer foldernya
    -------------------------------integrating DOS data-----------------------
    let intgTot = map (\i -> integrateToZero $ totalDOS ¿ [0,i]) $ flipBy invStat [1,2] -- run it on spin [1,2]
    putStrLn $ show intgTot
    -------------------------------integrating PDOS data------------------------
    let nAtom = length ctrlAtoms
    -------------------------------generating PDOS data------------------------
              -- map ditambah -1 karena input mengikuti gnuplot
              -- input : d kolom 6-10
              -- gnuplot : d kolom 6-10
              -- hmatrix : d kolom 5-9

  -------------------------------integrating PDOS data------------------------
    putStrLn $ "========invStat=" ++ (show invStat)
    (tMMomSD:mmomSD) <- fmap (map (* invStat)) $ readMMOM nAtom foldernya
    putStrLn $ show $ map (showDouble (3::Integer)) mmomSD
    putStrLn $ show tMMomSD
    putStrLn "==========show tMMomSD==========="
    pdosAtomicAll <- readPDOS invStat tailer foldernya ctrlAtomicAOs
    let integratedAtomicPDOS = integrateAtomicPDOS pdosAtomicAll
    let rIntgAll' = rendertable
         $ (:) cleanedJdHead
         $ (:) (concat [ ["Total" ]
           , ["  "]
           , map (showDouble (3::Integer)) $ (\[t,iu,idn] -> [t,iu-idn,t-(iu-idn)]) $ (tMMomSD:intgTot)
           ])
         $ zipWith (\a b -> a:b) (map show ([1,2..]::[Integer]))
         $ zipWith (\sdMom (intMom,(_,(j,_))) -> j:(map (showDouble (3::Integer)) [sdMom,intMom,sdMom-intMom])) mmomSD
         $ map (\(iu,idn,b) -> ((iu-idn),b) ) integratedAtomicPDOS
    let rIntgAll = unlines  [
                            rIntgAll'
                            , jdTable
                            ]
    {-
    resIntAll' <- markdownToTex rIntgAll
    let resIntAll = T.replace "\\}" "}"
                  $ T.replace "\\{" "{" $ T.pack
                  $ unlines [
                            "\\begin{longtable}[]{" ++ colAlign ++ "}"
                            , unlines $ tail $ lines $ T.unpack resIntAll'
                            ]
    putStrLn rIntgAll
    T.putStrLn resIntAll
    -}
    T.writeFile texFile $ T.pack rIntgAll
    putStrLn "===done CPVO.IO.Reader.Ecalj.MMOM: getMMOM ==="


{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
module Lang.Parser where

import           Import                  hiding ( try )
import           Data.Attoparsec.Combinator
import           Data.Attoparsec.Text
import           Control.Monad.Except           ( MonadError(throwError) )
import           Data.Sequence                  ( fromList )
import           Data.Complex                   ( Complex((:+)) )
import           Data.Ratio                     ( (%) )
import           RIO.Partial                    ( read )

data Base = Bin | Oct | Dec | Hex

oneOf :: [Char] -> Parser Char
oneOf list = satisfy (`elem` list)

noneOf :: [Char] -> Parser Char
noneOf list = satisfy (`notElem` list)

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

skipSpaces :: Parser ()
skipSpaces = skipMany space

escaped :: Parser Char
escaped = char '\\' *> oneOf "\\\"\'nrt"

parseChar :: Parser LispVal
parseChar = "#\\" *> (name <|> letter <|> symbol <|> oneOf "\\\"\'") <&> Char
  where name = ("space" $> ' ') <|> ("newline" $> '\n')

parseString :: Parser LispVal
parseString =
  char '"' *> many (noneOf "\\\"" <|> escaped) <* char '"' <&> String

parseAtom :: Parser LispVal
parseAtom = do
  firstLetter <- letter <|> symbol
  rest        <- many (letter <|> digit <|> symbol)
  let atom = firstLetter : rest
  return $ case atom of
    "#t" -> Bool True
    "#f" -> Bool False
    _    -> Atom atom

radixPref :: Parser Base
radixPref = char '#' *> (bin <|> oct <|> dec <|> hex)
 where
  bin = char 'b' $> Bin
  oct = char 'o' $> Oct
  dec = char 'd' $> Dec
  hex = char 'x' $> Hex

parseBase :: (Integral a, Read a) => a -> [Char] -> a
parseBase a list = parseBase' list 0
 where
  parseBase' []      n = n
  parseBase' (h : t) n = parseBase' t (n * a + parseDigit h)
  parseDigit c | c == 'A' || c == 'a' = 10
               | c == 'B' || c == 'b' = 11
               | c == 'C' || c == 'c' = 12
               | c == 'D' || c == 'd' = 13
               | c == 'E' || c == 'e' = 14
               | c == 'F' || c == 'f' = 15
               | otherwise            = read [c]

parseInteger :: Parser NumType
parseInteger = Integer <$> do
  pref <- option Dec radixPref
  case pref of
    Bin -> parseBase 2 <$> many1 digit
    Oct -> parseBase 8 <$> many1 digit
    Dec -> parseBase 10 <$> many1 digit
    Hex -> parseBase 16 <$> many1 digit

parseRational :: Parser NumType
parseRational = (%) <$> (decimal <* "/") <*> decimal <&> Rational

parseReal :: Parser NumType
parseReal = double <&> Real

parseComplex :: Parser NumType
parseComplex =
  (:+) <$> (double <* oneOf "+-") <*> (double <* char 'i') <&> Complex

parseNumber :: Parser LispVal
parseNumber =
  try parseInteger
    <|> try parseReal
    <|> try parseRational
    <|> try parseComplex
    <&> Number

parseList :: Parser LispVal -> Parser LispVal
parseList recParser = do
  char '(' *> skipSpaces
  inits <- sepBy recParser skipSpaces
  last  <- skipSpaces *> oneOf ".)"
  case last of
    ')' -> return $ List inits
    _ -> skipSpaces *> recParser <* skipSpaces <* char ')' <&> DottedList inits

parseVector :: Parser LispVal -> Parser LispVal
parseVector recParser =
  "#("
    *>  skipSpaces
    *>  sepBy recParser skipSpaces
    <*  skipSpaces
    <*  char ')'
    <&> Vector
    .   fromList

parseQuotes :: Text -> LispVal -> Parser LispVal -> Parser LispVal
parseQuotes c val recParser = do
  _ <- string c
  x <- recParser
  return $ List [val, x]

parseQuasiquote :: Parser LispVal -> Parser LispVal
parseQuasiquote = parseQuotes "`" (Atom "quasiquote")

parseQuote :: Parser LispVal -> Parser LispVal
parseQuote = parseQuotes "\'" (Atom "quote")

parseUnquote :: Parser LispVal -> Parser LispVal
parseUnquote = parseQuotes "," (Atom "unquote")

parseUnquoteSplicing :: Parser LispVal -> Parser LispVal
parseUnquoteSplicing = parseQuotes ",@" (Atom "unquote-splicing")

parseMetaVal :: Parser LispVal
parseMetaVal = many1 letter <&> MetaVal

parseMetaAtom :: Parser LispVal
parseMetaAtom = "atom:" *> many1 letter <&> MetaAtom

parseMetaList :: Parser LispVal
parseMetaList = "list:" *> many1 letter <&> MetaList

parseMetaVector :: Parser LispVal
parseMetaVector = "vec:" *> many1 letter <&> MetaVector

parseMetaString :: Parser LispVal
parseMetaString = "str:" *> many1 letter <&> MetaString

parseMeta :: Parser LispVal
parseMeta =
  char '@'
    *> (   parseMetaAtom
       <|> parseMetaList
       <|> parseMetaVector
       <|> parseMetaString
       <|> parseMetaVal
       )

parseExpr :: Parser LispVal
parseExpr =
  try parseChar
    <|> try parseString
    <|> try parseNumber
    <|> try (parseList parseExpr)
    <|> try (parseVector parseExpr)
    <|> try (parseQuote parseExpr)
    <|> try (parseQuasiquote parseExpr)
    <|> try (parseUnquoteSplicing parseExpr)
    <|> try (parseUnquote parseExpr)
    <|> try parseAtom

parseExprOrMeta :: Parser LispVal
parseExprOrMeta =
  parseMeta
    <|> try parseChar
    <|> try parseString
    <|> try parseNumber
    <|> try (parseList parseExprOrMeta)
    <|> try (parseVector parseExprOrMeta)
    <|> try (parseQuote parseExprOrMeta)
    <|> try (parseQuasiquote parseExprOrMeta)
    <|> try (parseUnquoteSplicing parseExprOrMeta)
    <|> try (parseUnquote parseExprOrMeta)
    <|> try parseAtom

readOrThrow :: Parser a -> Text -> ThrowsError a
readOrThrow parser input = case parseOnly parser input of
  Left  err -> throwError $ Parser err
  Right val -> return val

readExpr :: Text -> ThrowsError LispVal
readExpr = readOrThrow parseExpr

readExprList :: Text -> ThrowsError [LispVal]
readExprList = readOrThrow (sepBy parseExpr skipSpaces)


{-# LANGUAGE  ForeignFunctionInterface #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Foreign.LAPACK.Zomplex
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--

module Foreign.LAPACK.Zomplex
    where

import Data.Complex( Complex )
import Foreign( Ptr )
import Foreign.BLAS.Types
import Foreign.LAPACK.Types

#include "f77_func-hsc.h"

foreign import ccall unsafe #f77_func zgeqrf
    zgeqrf :: Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double)
           -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double)
           -> Ptr LAInt -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zgelqf
    zgelqf :: Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double)
           -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double)
           -> Ptr LAInt -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zheevr
    zheevr :: LAEigJob -> LAEigRange -> BLASUplo -> Ptr LAInt
           -> Ptr (Complex Double) -> Ptr LAInt -> Ptr Double -> Ptr Double
           -> Ptr LAInt -> Ptr LAInt -> Ptr Double -> Ptr LAInt -> Ptr Double
           -> Ptr (Complex Double) -> Ptr LAInt -> Ptr LAInt
           -> Ptr (Complex Double) -> Ptr LAInt -> Ptr Double -> Ptr LAInt
           -> Ptr LAInt -> Ptr LAInt -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zlarfg
    zlarfg :: Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double)
           -> Ptr LAInt -> Ptr (Complex Double) -> IO ()

foreign import ccall unsafe #f77_func zpotrf
    zpotrf :: BLASUplo -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt
           -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zpotrs
    zpotrs :: BLASUplo -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double)
           -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr LAInt
           -> IO ()

foreign import ccall unsafe #f77_func zpptrf
    zpptrf :: BLASUplo -> Ptr LAInt -> Ptr (Complex Double)
           -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zpptrs
    zpptrs :: BLASUplo -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double)
           -> Ptr (Complex Double) -> Ptr LAInt -> Ptr LAInt
           -> IO ()

foreign import ccall unsafe #f77_func zunmqr
    zunmqr :: BLASSide -> BLASTrans -> Ptr LAInt -> Ptr LAInt
           -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt
           -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt
           -> Ptr (Complex Double) -> Ptr LAInt -> Ptr LAInt
           -> IO ()

foreign import ccall unsafe #f77_func zunmlq
    zunmlq :: BLASSide -> BLASTrans -> Ptr LAInt -> Ptr LAInt -> Ptr LAInt
           -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double)
           -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double)
           -> Ptr LAInt -> Ptr LAInt -> IO ()


module Shrinkage (
    strimmerX
  , testShrinkage
  ) where

import Numeric.LinearAlgebra
import Test.Tasty
import Test.Tasty.HUnit

import Statistics.Covariance.Shrinkage.Strimmer

strimmerX :: Matrix Double
strimmerX = matrix 3 [
    -1.1541680,  0.1229077,  0.2964592
  , -0.4574491, -0.0106610, -0.2497785
  ,  1.2711852,  0.2362810,  0.3653930
  ]

withinTolVec :: Double -> Vector Double -> Vector Double -> Bool
withinTolVec tol x x' = abs (x - x') < scalar tol

withinTolMat :: Double -> Matrix Double -> Matrix Double -> Bool
withinTolMat tol x x' = withinTolVec tol (flatten x) (flatten x')

strimmerpsmallSVD :: Assertion
strimmerpsmallSVD = assertBool "psmallSVD" $ withinTolMat tol u' u
  && withinTolVec tol s' s
  && withinTolMat tol v' v
  where
    tol = 1e-6
    (u, s, v) = psmallSVD strimmerX
    u' = matrix 3 [
         0.6256382, -0.7737046, -0.0997895
      ,  0.2669970,  0.3325596, -0.9044980
      , -0.7330003, -0.5392451, -0.4146388
      ]
    s' = vector [1.7850070, 0.5601167, 0.1101480]
    v' = matrix 3 [
        -0.99495870,  0.09886453,  0.01682235
      , -0.05554307, -0.40358205, -0.91325599
      , -0.08349942, -0.90958635,  0.40703871
      ]

strimmernsmallSVD :: Assertion
strimmernsmallSVD = assertBool "nsmallSVD" $ withinTolMat tol u' u
  && withinTolVec tol s' s
  && withinTolMat tol v' v
  where
    tol = 1e-6
    (u, s, v) = nsmallSVD strimmerX
    u' = matrix 3 [
        -0.6256382,  0.7737046, 0.0997895
      , -0.2669970, -0.3325596, 0.9044980
      ,  0.7330003,  0.5392451, 0.4146388
      ]
    s' = vector [1.7850070, 0.5601167, 0.1101480]
    v' = matrix 3 [
        0.99495870, -0.09886453, -0.01682235
      , 0.05554307,  0.40358205,  0.91325599
      , 0.08349942,  0.90958635, -0.40703871
      ]

strimmerposSVD :: Assertion
strimmerposSVD = assertBool "posSVD" $ withinTolMat tol u' u
  && withinTolVec tol s' s
  && withinTolMat tol v' v
  where
    tol = 1e-6
    (u, s, v) = posSVD strimmerX
    u' = matrix 3 [
        -0.6256382,  0.7737046, 0.0997895
      , -0.2669970, -0.3325596, 0.9044980
      ,  0.7330003,  0.5392451, 0.4146388
      ]
    s' = vector [1.7850070, 0.5601167, 0.1101480]
    v' = matrix 3 [
        0.99495870, -0.09886453, -0.01682235
      , 0.05554307,  0.40358205,  0.91325599
      , 0.08349942,  0.90958635, -0.40703871
      ]

strimmerCovShrink :: Assertion
strimmerCovShrink = assertBool "corpcor cor shrink is wrong" $ withinTolMat 1e-6 cov' cov
  where
    cov = unSym (covShrink strimmerX)
    cov' = matrix 3 [
        1.13853590, 0.06889376, 0.05675350
      , 0.06889376, 0.04389409, 0.03081658
      , 0.05675350, 0.03081658, 0.11359393
      ]

strimmerCorShrink :: Assertion
strimmerCorShrink = assertBool "corpcor cor shrink is wrong" $ withinTolMat 1e-6 cor' cor
  where
    cor = unSym (corShrink strimmerX)
    cor' = matrix 3 [
        1.0000000, 0.3081793, 0.1578126
      , 0.3081793, 1.0000000, 0.4364192
      , 0.1578126, 0.4364192, 1.0000000
      ]

strimmerCorCoef :: Assertion
strimmerCorCoef = assertBool msg $ abs (lambda - coef) < 0.0000001
  where
    lambda = 0.5311773
    coef = corCoef strimmerX
    msg = "expected: " ++ show lambda ++ "\n     got: " ++ show coef

strimmerVarCoef :: Assertion
strimmerVarCoef = assertBool "corpcor var coef is wrong" $ abs (lambda - varCoef strimmerX) < 0.0000001
  where lambda = 0.2910548

strimmerVarShrink :: Assertion
strimmerVarShrink = assertBool "corp var shrink" $ abs (v - varShrink strimmerX) < 0.0000001
  where v = vector [1.13853590, 0.04389409, 0.11359393]

testShrinkage :: TestTree
testShrinkage = testGroup "shrinkage" [
    testCase "strimmer varCoef" strimmerVarCoef
  , testCase "strimmer varShrink" strimmerVarShrink
  , testCase "strimmer corCoef" strimmerCorCoef
  , testCase "strimmer corShrink" strimmerCorShrink
  , testCase "strimmer covShrink" strimmerCovShrink
  , testCase "strimmer psmallSVD" strimmerpsmallSVD
  , testCase "strimmer nsmallSVD" strimmernsmallSVD
  , testCase "strimmer posSVD" strimmerposSVD
  ]


module Main where

-- third party imports
import Criterion.Main
import Data.Complex
import Data.Tuple
import Graphics.Image
import System.Environment

-- local imports
import CommonFunctions
import JuliaSet

main :: IO ()
main = do
  args <- getArgs
  if length args == 0 then do
    usage
  else do
    if length args < 5 then do
      visualizeJuliaBench
    else do
      let outputFilename = args !! 4
      -- putStrLn "visualizing julia set..."
      runWithCliArgs args
      -- putStrLn ("Done! Image output to " ++ outputFilename)

visualizeJuliaBench = defaultMain [
    bgroup "visualize julia set" [
      -- bench "f1 ((-0.4) :+ 0.65) VU" $ whnf makeImageVU g,
      -- bench "f1 ((-0.4) :+ 0.65) VS" $ whnf makeImageVS g,
      bench "f1 ((-0.4) :+ 0.65) RSU" $ whnf makeImageRSU g,
      bench "f1 ((-0.4) :+ 0.65) RPU" $ whnf makeImageRPU g,
      bench "f1 ((-0.4) :+ 0.65) RSS" $ whnf makeImageRSS g,
      bench "f1 ((-0.4) :+ 0.65) RPS" $ whnf makeImageRPS g
    ]
  ]
  where
    default_func = f1 ((-0.4) :+ 0.65)
    g = pixelToJuliaSetValue default_func 1.5 width height 100
    -- makeImageVU = makeImageR VU (width, height)
    -- makeImageVS = makeImageR VS (width, height)
    makeImageRSU = makeImageR RSU (width, height)
    makeImageRPU = makeImageR RPU (width, height)
    makeImageRSS = makeImageR RSS (width, height)
    makeImageRPS = makeImageR RPS (width, height)
    width = 10000
    height = 10000

-- takes string cli args and returns processed args for visualizeJuliaSet
runWithCliArgs :: [String] -> IO()
runWithCliArgs args
  -- | (length args) < 5 = do
  --     usage
  --     putStrLn "using default values " ++ (show (
  --       default_func, 1.5, 1000, 1000, 100, "images/output.png"))
  --     (default_func, 1.5, 1000, 1000, 100, "images/output.png")
  | (length args) < 6 = visualizeJuliaSet default_func r width height
                          maxIter outputFilename "RPU"
  | (length args) < 7 = visualizeJuliaSet default_func r width height
                          maxIter outputFilename arrayType
  | (length args) < 8 = visualizeJuliaSet (parseFunctionParams func_num [])
                          r width height maxIter outputFilename arrayType
  | (length args) < 9 = visualizeJuliaSet
                          (parseFunctionParams func_num params)
                          r width height maxIter outputFilename arrayType
  where
    default_func = f1 ((-0.4) :+ 0.65)
    r = read (args !! 0) :: Double
    width = read (args !! 1) :: Int
    height = read (args !! 2) :: Int
    maxIter = read (args !! 3) :: Int
    outputFilename = args !! 4
    arrayType = args !! 5
    func_num = read (args!!6) :: Int
    params = Prelude.map (read::String -> Complex Double)
                         (slice 7 (length args) args)

slice :: Int -> Int -> [a] -> [a]
slice from to xs = take (to - from + 1) (drop from xs)

parseFunctionParams :: Int
                    -> [Complex Double]
                    -> (Complex Double -> Complex Double)
parseFunctionParams func_num params
    | func_num == 1 && params == [] = f1 ((-0.4) :+ 0.65)
    | func_num == 1 = f1 (params!!0)

usage = do
  let prog = "stack run"
  -- prog <- getProgName
  putStrLn ("Usages: ")
  putStrLn ("    " ++ prog)
  putStrLn ("    " ++ prog ++ " escape_radius width height "
            ++ "max_iter output_filename [arrType]")
  putStrLn ("              [func_num [constant1 [constant2] ... ]]")
  putStrLn ("    " ++ prog ++ " -- --output benchmark.html")
  putStrLn ""
  putStrLn "usage 1 (no argments) displays this help"
  putStrLn ("usage 2 visualizes the julia set created with the " ++
           "given parameters")
  putStrLn "usage 3 benchmarks the program"
  putStrLn ""
  putStrLn "notes on usage 2:"
  putStrLn ("number of constants must correspond to the number of constants "
    ++ "function func_num takes")
  putStrLn ("example: " ++ prog ++ " 1.5 1000 1000 100 "
    ++ "\"images/output.png\"")
  putStrLn ("example: " ++ prog ++ " 1.5 1000 1000 100 "
    ++ "\"images/output.png\" RPU 1 \"(0.285 :+ 0)\"")