text:
import Data.Complex
import Data.Either
import Data.Maybe
--import System.Random
import Text.Read

main = do
  putStrLn "Compiled."
  center defaultGD ["0", "0"]

--Datatypes/typeclasses/etc.

data Vector2 = V2 Double Double deriving (Show, Read, Eq)

data Transformation = Transformation Matrix Vector2 deriving (Show, Read, Eq)

data GlobalData = GD
  { getWindowSize :: (Int, Int),
    getTickSizes :: (Double, Double),
    getPlotLoc :: (Vector2, Vector2),
    getIters :: Int,
    getAxesShown :: Bool,
    getVectToChar :: VectToChar
  }

type Matrix = (Double, Double, Double, Double)

type Args = [String]

type StrCommand = String

type ArgDescription = String

type CmdDescription = String

type Command = GlobalData -> Args -> IO ()

type PlotCommand = GlobalData -> Args -> Either (IO ()) VectToChar

type VectToChar = GlobalData -> Vector2 -> Char

--Constants

defaultGD :: GlobalData
defaultGD = GD (100, 30) (1.0, 1.0) (V2 (-1.5) (-1), V2 1.5 1) 50 True vtcMBSet

charWidth :: Double
charWidth = 0.3

cursorChar = 'o'

originChar = '+'

hAxisMinChar = '<'

hAxisMaxChar = '>'

hAxisChar = '-'

hAxisTickChar = '|'

vAxisMinChar = 'V'

vAxisMaxChar = '^'

vAxisChar = '|'

vAxisTickChar = '~'

maxShadeChar = '#'

minShadeChar = ' '

vBounderyChar = '|'

hBounderyChar = '~'

shadeChars :: [Char]
shadeChars = " `.*+&$@#"

commandList :: [(StrCommand, (Command, ArgDescription, CmdDescription))]
commandList =
  [ ("draw", (draw, "draw", "Redraws the plot to the screen.")),
    ("plot", (plot, "plot (args)", "Sets the function defining which points to plot.\nPossible arguments:\n" ++ show (map (sndTriple . snd) plotList))),
    ( "zoom",
      ( zoom,
        "zoom [Double c]",
        "Zooms in by the specified factor."
          ++ "Requires c > 0. Omitting arguments sets zoom level."
      )
    ),
    ("shift", (shift, "shift (Double dx) (Double dy)", "Translates the plot by the specified amount.")),
    ("center", (center, "center (Double x) (Double y)", "Centers the plot on the point (x,y). Defaults to the origin.")),
    ( "resize",
      ( resize,
        "resize (Int width) (Int height)",
        "Resizes the window. Omitting arguments defaults to 100x30."
          ++ "\nRequires width, height > 0. Tag -square ignores second argument and calculates height based on given width."
      )
    ),
    ("toggleAxes", (toggleAxes, "toggleAxes", "Toggles the visibility of the axes.")),
    ( "cutoff",
      ( cutoff,
        "cutoff (Int n)",
        "Sets the cutoff point for iterations. Generally, a higher cutoff point means greater precision.\n"
          ++ "Requires n > 0. Omitting arguments defaults to 50."
      )
    ),
    ("reload", (reload, "reload", "Restores all data (center, scale factor, etc.) to their default values.")),
    ("info", (info, "info", "Prints information about the current plot.")),
    ("quit", (quit, "quit", "Ends execution. Current data will be lost.")),
    ("help", (help, "help [command]", "Prints a description of the specified command and its arguments. Omitting the arguments lists available commands."))
  ]

plotList :: [(StrCommand, (PlotCommand, ArgDescription, CmdDescription))]
plotList =
  [ ("mandelbrot", (plotMBSet, "plot mandelbrot", "Plots the Mandelbrot set over the complex plane.\nTag -shade adds shading based on time to diverge.")),
    ("sierpinski", (plotSierpinski, "plot sierpinski", "Work in progress. Plots the Sierpinski triangle within the unit square, calculated as an iterated function system.\n")),
    ( "juliaset",
      ( plotJuliaSet,
        "plot juliaset (Complex c)",
        "Given a complex value (using a:+b notation), plots the filled julia set of z^2+c over the complex plane.\n"
          ++ "Omitting arguments defaults to 0:+0. Tag -shade adds shading based on time to diverge.\n"
      )
    )
  ]

--Triple commands

fstTriple (x, _, _) = x

sndTriple (_, y, _) = y

thrdTriple (_, _, z) = z

--IO commands

getCommand :: Command
getCommand gd _ = do
  putStr "\nMBSet>> "
  line <- getLine
  if line == ""
    then do getCommand gd []
    else do
      let (strCmd : args) = words line
      let mCommand = lookup strCmd commandList
      case mCommand of
        Nothing -> rejectCommand gd [strCmd]
        Just (command, argDesc, cmdDesc) -> command gd args
      return ()

rejectCommand :: Command
rejectCommand gd [strCmd] = do
  putStrLn $ "No such command: " ++ strCmd
  putStrLn "Use the command \"help\" to see a list of available commands."
  getCommand gd []
rejectCommand gd args = do
  putStrLn "No such command.\nUse the command \"help\" to see a list of available commands."
  getCommand gd []

rejectArgs :: Command
rejectArgs gd args = do
  putStrLn $ "Invalid arguments: " ++ unwords args ++ "\nUse the command \"help [command]\" for command syntax."
  getCommand gd []

--Input Commands

draw :: Command
draw gd _ = do
  putStr . unlines . asciiArray $ gd
  getCommand gd []

plot :: Command
plot gd args@(strPlot : args') = do
  let mPlotTriple = lookup strPlot plotList
  if isNothing mPlotTriple
    then do
      rejectArgs gd args
    else do
      let Just (plotCmd, _, _) = mPlotTriple
      if isLeft $ plotCmd gd args'
        then do
          let Left failureAction = plotCmd gd args
          failureAction ---------
        else do
          let Right command = plotCmd gd args'
          draw (setVectToChar gd command) [""]
plot gd args = help gd $ "plot" : args

zoom :: Command
zoom gd [] = draw (setPlotLoc gd $ getPlotLoc defaultGD) []
zoom gd args
  | head args == "in" = zoom gd $ "2" : tail args
  | head args == "out" = zoom gd $ "0.5" : tail args
  | otherwise = do
    let mDouble = readMaybe (head args) :: Maybe Double
    case mDouble of
      Nothing -> rejectArgs gd args
      Just x ->
        if x <= 0
          then rejectArgs gd args
          else draw (zoomGD gd x) $ tail args
    return ()

shift :: Command
shift gd@GD {getPlotLoc = (v1, v2)} args@(dxStr : dyStr : args')
  | isNothing (readMaybe dxStr :: Maybe Double) = rejectArgs gd args
  | isNothing (readMaybe dyStr :: Maybe Double) = rejectArgs gd args
  | otherwise = draw (setPlotLoc gd (v1 `vPlus` d, v2 `vPlus` d)) args'
  where
    d = V2 (read dxStr :: Double) (read dyStr :: Double)
shift gd args = rejectArgs gd args

center :: Command
center gd [] = center gd ["0", "0"]
center gd@GD {getPlotLoc = (v1, v2)} args@(xStr : yStr : args')
  | isNothing (readMaybe xStr :: Maybe Double) = rejectArgs gd args
  | isNothing (readMaybe yStr :: Maybe Double) = rejectArgs gd args
  | otherwise = shift gd $ [show dx, show dy] ++ args'
  where
    (V2 dx dy) = newCenter `vMinus` currentCenter
    currentCenter = midPt v1 v2
    newCenter = V2 (read xStr :: Double) (read yStr :: Double)
center gd args = rejectArgs gd args

resize :: Command
resize gd [] = draw (setWindowSize gd $ getWindowSize defaultGD) []
resize gd args@(wStr : hStr : args')
  | isNothing (readMaybe wStr :: Maybe Int) = rejectArgs gd args
  | width <= 0 = rejectArgs gd args
  | "square" `elem` args = resize gd sqrArgs
  | isNothing (readMaybe hStr :: Maybe Int) = rejectArgs gd args
  | height <= 0 = rejectArgs gd args
  | otherwise = draw (setWindowSize gd (width, height)) args'
  where
    width = read wStr :: Int
    height = read hStr :: Int
    sqrHeight = round . (* charWidth) . fromIntegral $ width
    sqrArgs = wStr : show sqrHeight : filter (/= "square") args'
resize gd args = rejectArgs gd args

cutoff :: Command
cutoff gd [] = draw (setIters gd $ getIters defaultGD) []
cutoff gd args@(strIters : args')
  | isNothing (readMaybe strIters :: Maybe Int) = rejectArgs gd args
  | iters <= 0 = rejectArgs gd args
  | otherwise = draw (setIters gd iters) args'
  where
    iters = read strIters :: Int

toggleAxes :: Command
toggleAxes gd@GD {getAxesShown = axesShown} = draw (setAxesShown gd (not axesShown))

reload :: Command
reload _ = draw defaultGD

info :: Command
info gd args = do
  putStrLn "\n  --  Current plot information  --"
  putStrLn $ "Window size: " ++ (show . fst . getWindowSize $ gd) ++ "x" ++ (show . snd . getWindowSize $ gd)
  let (xMin, yMin) = (getX . fst . getPlotLoc $ gd, getY . fst . getPlotLoc $ gd)
  let (xMax, yMax) = (getX . snd . getPlotLoc $ gd, getY . snd . getPlotLoc $ gd)
  putStrLn $ "Plot location (bounds): " ++ show xMin ++ " < x < " ++ show xMax ++ ",  " ++ show yMin ++ " < y < " ++ show yMax
  putStrLn $ "Cutoff point (number of iterations): " ++ show (getIters gd)
  putStrLn $ "Axes shown: " ++ show (getAxesShown gd)
  getCommand gd []

quit :: Command
quit gd args = do
  putStrLn "Quitting.\n\n"
  return ()

help :: Command
help gd ("plot" : strPlot : args) = do
  let output = case lookup strPlot plotList of
        Just (f, argDesc, cmdDesc) -> argDesc ++ " -- " ++ cmdDesc
        Nothing -> "No such plot: " ++ strPlot
  putStrLn output
  getCommand gd []
help gd [] = do
  putStrLn "Available commands: "
  print $ map fst commandList
  getCommand gd []
help gd args = do
  let strCmd = head args
  let output = case lookup strCmd commandList of
        Just (f, argDesc, cmdDesc) -> argDesc ++ " -- " ++ cmdDesc
        Nothing -> "No such command: " ++ strCmd
  putStrLn output
  getCommand gd []

--GlobalData functions

setWindowSize :: GlobalData -> (Int, Int) -> GlobalData
setWindowSize (GD _ tickSizes plotLoc iters axesShown vectToChar) windowSize = GD windowSize tickSizes plotLoc iters axesShown vectToChar

setTickSizes :: GlobalData -> (Double, Double) -> GlobalData
setTickSizes (GD windowSize _ plotLoc iters axesShown vectToChar) tickSizes = GD windowSize tickSizes plotLoc iters axesShown vectToChar

setPlotLoc :: GlobalData -> (Vector2, Vector2) -> GlobalData
setPlotLoc (GD windowSize tickSizes _ iters axesShown vectToChar) plotLoc = GD windowSize tickSizes plotLoc iters axesShown vectToChar

setIters :: GlobalData -> Int -> GlobalData
setIters (GD windowSize tickSizes plotLoc _ axesShown vectToChar) iters = GD windowSize tickSizes plotLoc iters axesShown vectToChar

setAxesShown :: GlobalData -> Bool -> GlobalData
setAxesShown (GD windowSize tickSizes plotLoc iters _ vectToChar) axesShown = GD windowSize tickSizes plotLoc iters axesShown vectToChar

setVectToChar :: GlobalData -> VectToChar -> GlobalData
setVectToChar (GD windowSize tickSizes plotLoc iters axesShown _) = GD windowSize tickSizes plotLoc iters axesShown

pxlW :: GlobalData -> Double
pxlW gd = (getX vMax - getX vMin) / fromIntegral screenWidth
  where
    (screenWidth, _) = getWindowSize gd
    (vMin, vMax) = getPlotLoc gd

pxlH :: GlobalData -> Double
pxlH gd = (getY vMax - getY vMin) / fromIntegral screenHeight
  where
    (_, screenHeight) = getWindowSize gd
    (vMin, vMax) = getPlotLoc gd

zoomGD :: GlobalData -> Double -> GlobalData
zoomGD gd@GD {getPlotLoc = (v1, v2)} c = setPlotLoc gd (v1', v2')
  where
    v1' = vPlus (midPt v1 v2) . vScale (1 / c) $ v1 `vMinus` midPt v1 v2
    v2' = vPlus (midPt v1 v2) . vScale (1 / c) $ v2 `vMinus` midPt v1 v2

--Vector2 functions

getX :: Vector2 -> Double
getX (V2 x _) = x

getY :: Vector2 -> Double
getY (V2 _ y) = y

mag :: Vector2 -> Double
mag (V2 x y) = sqrt (x ^ 2 + y ^ 2)

vPlus :: Vector2 -> Vector2 -> Vector2
vPlus (V2 a b) (V2 x y) = V2 (a + x) (b + y)

vMinus :: Vector2 -> Vector2 -> Vector2
vMinus (V2 a b) (V2 x y) = V2 (a - x) (b - y)

vScale :: Double -> Vector2 -> Vector2
vScale c (V2 x y) = V2 (x * c) (y * c)

midPt :: Vector2 -> Vector2 -> Vector2
midPt v1 v2 = vScale 0.5 (v1 `vPlus` v2)

applyTransformation :: Transformation -> Vector2 -> Vector2
applyTransformation (Transformation (a11, a21, a12, a22) (V2 bx by)) v@(V2 x y) =
  V2 (a11 * x + a12 * y + bx) (a21 * x + a22 * y + by)

-- Printing

hRange :: GlobalData -> [Double]
hRange gd = [xMin + i * dx | i <- [0 .. width]]
  where
    dx = (xMax - xMin) / width
    width = fromIntegral . fst . getWindowSize $ gd
    xMin = getX . fst . getPlotLoc $ gd
    xMax = getX . snd . getPlotLoc $ gd

vectRow :: GlobalData -> Double -> [Vector2]
vectRow gd y = map (`V2` y) $ hRange gd

vRange :: GlobalData -> [Double]
vRange gd = reverse [yMin + i * dy | i <- [0 .. height]]
  where
    dy = (yMax - yMin) / height
    height = fromIntegral . snd . getWindowSize $ gd
    yMin = getY . fst . getPlotLoc $ gd
    yMax = getY . snd . getPlotLoc $ gd

vectArray :: GlobalData -> [[Vector2]]
vectArray gd = map (vectRow gd) $ vRange gd

asciiArray :: GlobalData -> [[Char]]
asciiArray gd = addBoundery . map (map (vectToAscii gd)) $ vectArray gd

addBoundery :: [[Char]] -> [[Char]]
addBoundery array = [hBoundery] ++ addVBoundery array ++ [hBoundery]
  where
    hBoundery = minShadeChar : replicate (length . head $ array) hBounderyChar ++ [minShadeChar]
    addVBoundery array = map (\row -> vBounderyChar : row ++ [vBounderyChar]) array

vectToAscii :: GlobalData -> Vector2 -> Char
vectToAscii gd v@(V2 x y)
  | onCursor = cursorChar
  | onXAxis && onYAxis = originChar
  -- x-axis
  | onXAxis && abs (x - xMin) <= pxlW gd / 2 = hAxisMinChar
  | onXAxis && abs (x - xMax) <= pxlW gd / 2 = hAxisMaxChar
  | onXTick = hAxisTickChar
  | onXAxis = hAxisChar
  -- y-axis
  | onYAxis && abs (y - yMin) <= pxlH gd / 2 = vAxisMinChar
  | onYAxis && abs (y - yMax) <= pxlH gd / 2 = vAxisMaxChar
  | onYTick = vAxisTickChar
  | onYAxis = vAxisChar
  | otherwise = getVectToChar gd gd v
  where
    (scrnW, scrnH) = getWindowSize gd
    (vMin, vMax) = getPlotLoc gd
    onXAxis = getAxesShown gd && abs y < pxlH gd / 2
    onYAxis = getAxesShown gd && abs x < pxlW gd / 2
    onXTick = onXAxis && abs (x - rndX) < pxlW gd / 2 --WIP
    onYTick = onYAxis && abs (y - rndY) < pxlH gd / 2 --WIP
    [rndX, rndY] = map (fromIntegral . round) [x, y] --WIP rewrite as nearest tick
    (xMin, yMin) = (getX vMin, getY vMin)
    (xMax, yMax) = (getX vMax, getY vMax)
    onCursor = mag (v `vMinus` midPt vMin vMax) <= pxlW gd / 2

--Mandelbrot set

plotMBSet :: PlotCommand
plotMBSet gd args
  | "-shade" `elem` args = Right vtcMBSetShaded
  | otherwise = Right vtcMBSet

vtcMBSet :: VectToChar
vtcMBSet gd@GD {getIters = iters} (V2 x y)
  | isInSet = maxShadeChar
  | otherwise = minShadeChar
  where
    isInSet = all (\z -> magnitude z < 2) . take iters $ mSequence c (0 :+ 0)
    c = x :+ y
    mSequence c z = z : mSequence c (mFunc c z)
    mFunc c z = z ^ 2 + c

vtcMBSetShaded :: VectToChar
vtcMBSetShaded gd@GD {getIters = maxIters} (V2 x y)
  | itersToEscape >= maxIters = maxShadeChar
  | itersToEscape >= (49 * maxIters) `div` 8 = shadeChars !! 7
  | itersToEscape >= (36 * maxIters) `div` 8 = shadeChars !! 6
  | itersToEscape >= (25 * maxIters) `div` 8 = shadeChars !! 5
  | itersToEscape >= (16 * maxIters) `div` 8 = shadeChars !! 4
  | itersToEscape >= (9 * maxIters) `div` 8 = shadeChars !! 3
  | itersToEscape >= (4 * maxIters) `div` 8 = shadeChars !! 2
  | itersToEscape >= (1 * maxIters) `div` 8 = shadeChars !! 1
  | otherwise = minShadeChar
  where
    itersToEscape = length . takeWhile (\z -> magnitude z < 2) . take maxIters $ mSequence c (0 :+ 0)
    c = x :+ y
    mSequence c z = z : mSequence c (z ^ 2 + c)

--Julia Sets

plotJuliaSet :: PlotCommand
plotJuliaSet gd args@(cStr : args')
  | isNothing mc = Left (rejectArgs gd args)
  | "-shade" `elem` args = Right (vtcJuliaSetShaded c)
  | otherwise = Right (vtcJuliaSet c)
  where
    mc = readMaybe cStr :: Maybe (Complex Double)
    Just c = mc
plotJuliaSet gd args = plotJuliaSet gd ("0" : "0" : args)

vtcJuliaSet :: Complex Double -> VectToChar
vtcJuliaSet c gd@GD {getIters = iters} (V2 x y)
  | isInSet = maxShadeChar
  | otherwise = minShadeChar
  where
    isInSet = all (\z -> magnitude z < 2) . take iters $ jSequence c (x :+ y)
    jSequence c z = z : jSequence c (z ^ 2 + c)

vtcJuliaSetShaded :: Complex Double -> VectToChar
vtcJuliaSetShaded c gd@GD {getIters = maxIters} (V2 x y)
  | itersToEscape >= maxIters = maxShadeChar
  | itersToEscape >= (7 * maxIters) `div` 8 = shadeChars !! 7
  | itersToEscape >= (6 * maxIters) `div` 8 = shadeChars !! 6
  | itersToEscape >= (5 * maxIters) `div` 8 = shadeChars !! 5
  | itersToEscape >= (4 * maxIters) `div` 8 = shadeChars !! 4
  | itersToEscape >= (3 * maxIters) `div` 8 = shadeChars !! 3
  | itersToEscape >= (2 * maxIters) `div` 8 = shadeChars !! 2
  | itersToEscape >= (1 * maxIters) `div` 8 = shadeChars !! 1
  | otherwise = minShadeChar
  where
    itersToEscape = length . takeWhile (\z -> magnitude z < 2) . take maxIters $ jSequence c (x :+ y)
    jSequence c z = z : jSequence c (z ^ 2 + c)

--Sierpinski triangle (WIP - random)

plotSierpinski :: PlotCommand
plotSierpinski gd args = Right vtcSierpinski

vtcSierpinski :: VectToChar
vtcSierpinski gd v@(V2 x y)
  | x < 0 || x > 1 || y < 0 || y > 1 = minShadeChar
  | any (withinRange gd v) $ truncatedList gd = maxShadeChar
  | otherwise = minShadeChar

truncatedList gd = take iters' $ drop iters' sierpinskiList' --Hardcoded parameter
  where
    iters' = 10 * getIters gd

withinRange gd (V2 x1 y1) (V2 x2 y2) = abs (x2 - x1) <= width && abs (y2 - y1) <= height
  where
    (width, height) = (0.75 * pxlW gd, 0.75 * pxlH gd) --Hardcoded parameter

sierpinskiList :: (Vector2, Int) -> [Vector2]
sierpinskiList (v, n) = v : sierpinskiList (nextSierpinski (v, n))

sierpinskiList' = sierpinskiList (V2 0.5 0.5, 0) --Initial values

nextSierpinski :: (Vector2, Int) -> (Vector2, Int)
nextSierpinski (v@(V2 x y), n)
  | rand n `mod` 3 == 0 = (applyTransformation trans0 v, n + 1)
  | rand n `mod` 3 == 1 = (applyTransformation trans1 v, n + 1)
  | rand n `mod` 3 == 2 = (applyTransformation trans2 v, n + 1)
nextSierpinski _ = (V2 0 0, 0)

trans0 = Transformation (0.5, 0.0, 0.0, 0.5) (V2 0 0)

trans1 = Transformation (0.5, 0.0, 0.0, 0.5) (V2 0.5 0)

trans2 = Transformation (0.5, 0.0, 0.0, 0.5) (V2 0.25 (sqrt 3 / 4))

--Generates a pseudorandom digit
rand :: Int -> Int
rand n = randList !! (n `mod` length randList)

randList =
  [ 3,
    1,
    4,
    1,
    5,
    9,
    2,
    6,
    5,
    3,
    5,
    8,
    9,
    7,
    9,
    3,
    2,
    3,
    8,
    4,
    6,
    2,
    6,
    4,
    3,
    3,
    8,
    3,
    2,
    7,
    9,
    5,
    0,
    2,
    8,
    8,
    4,
    1,
    9,
    7,
    1,
    6,
    9,
    3,
    9,
    9,
    3,
    7,
    5,
    1,
    0,
    5,
    8,
    2,
    0,
    9,
    7,
    4,
    9,
    4,
    4,
    5,
    9,
    2,
    3,
    0,
    7,
    8,
    1,
    6,
    4,
    0,
    6,
    2,
    8,
    6,
    2,
    0,
    8,
    9,
    9,
    8,
    6,
    2,
    8,
    0,
    3,
    4,
    8,
    2,
    5,
    3,
    4,
    2,
    1,
    1,
    7,
    0,
    6,
    7,
    9,
    8,
    2,
    1,
    4,
    8,
    0,
    8,
    6,
    5,
    1,
    3,
    2,
    8,
    2,
    3,
    0,
    6,
    6,
    4,
    7,
    0,
    9,
    3,
    8,
    4,
    4,
    6,
    0,
    9,
    5,
    5,
    0,
    5,
    8,
    2,
    2,
    3,
    1,
    7,
    2,
    5,
    3,
    5,
    9,
    4,
    0,
    8,
    1,
    2,
    8,
    4,
    8,
    1,
    1,
    1,
    7,
    4,
    5,
    0,
    2,
    8,
    4,
    1,
    0,
    2,
    7,
    0,
    1,
    9,
    3,
    8,
    5,
    2,
    1,
    1,
    0,
    5,
    5,
    5,
    9,
    6,
    4,
    4,
    6,
    2,
    2,
    9,
    4,
    8,
    9,
    5,
    4,
    9,
    3,
    0,
    3,
    8,
    1,
    9,
    6,
    4,
    4,
    2,
    8,
    8,
    1,
    0,
    9,
    7,
    5,
    6,
    6,
    5,
    9,
    3,
    3,
    4,
    4,
    6,
    1,
    2,
    8,
    4,
    7,
    5,
    6,
    4,
    8,
    2,
    3,
    3,
    7,
    8,
    6,
    7,
    8,
    3,
    1,
    6,
    5,
    2,
    7,
    1,
    2,
    0,
    1,
    9,
    0,
    9,
    1,
    4,
    5,
    6,
    4,
    8,
    5,
    6,
    6,
    9,
    2,
    3,
    4,
    6,
    0,
    3,
    4,
    8,
    6,
    1,
    0,
    4,
    5,
    4,
    3,
    2,
    6,
    6,
    4,
    8,
    2,
    1,
    3,
    3,
    9,
    3,
    6,
    0,
    7,
    2,
    6,
    0,
    2,
    4,
    9,
    1,
    4,
    1,
    2,
    7,
    3,
    7,
    2,
    4,
    5,
    8,
    7,
    0,
    0,
    6,
    6,
    0,
    6,
    3,
    1,
    5,
    5,
    8,
    8,
    1,
    7,
    4,
    8,
    8,
    1,
    5,
    2,
    0,
    9,
    2,
    0,
    9,
    6,
    2,
    8,
    2,
    9,
    2,
    5,
    4,
    0,
    9,
    1,
    7,
    1,
    5,
    3,
    6,
    4,
    3,
    6,
    7,
    8,
    9,
    2,
    5,
    9,
    0,
    3,
    6,
    0,
    0,
    1,
    1,
    3,
    3,
    0,
    5,
    3,
    0,
    5,
    4,
    8,
    8,
    2,
    0,
    4,
    6,
    6,
    5,
    2,
    1,
    3,
    8,
    4,
    1,
    4,
    6,
    9,
    5,
    1,
    9,
    4,
    1,
    5,
    1,
    1,
    6,
    0,
    9,
    4,
    3,
    3,
    0,
    5,
    7,
    2,
    7,
    0,
    3,
    6,
    5,
    7,
    5,
    9,
    5,
    9,
    1,
    9,
    5,
    3,
    0,
    9,
    2,
    1,
    8,
    6,
    1,
    1,
    7,
    3,
    8,
    1,
    9,
    3,
    2,
    6,
    1,
    1,
    7,
    9,
    3,
    1,
    0,
    5,
    1,
    1,
    8,
    5,
    4,
    8,
    0,
    7,
    4,
    4,
    6,
    2,
    3,
    7,
    9,
    9,
    6,
    2,
    7,
    4,
    9
  ]


text:
{-# LANGUAGE NamedFieldPuns #-}

module School.FileIO.FileHeader
( FileHeader(..)
, compatibleHeaders
, headerBuilder
, parseHeader
) where

import Control.Applicative ((<|>))
import Control.Monad (replicateM)
import Data.Attoparsec.ByteString (Parser, anyWord8, parseOnly, take, word8)
import Data.Attoparsec.ByteString.Char8 (char)
import qualified Data.ByteString as B
import Data.ByteString.Conversion (FromByteString(..), ToByteString(..), toByteString')
import Data.Default.Class (Default(..))
import Data.Monoid ((<>))
import Numeric.LinearAlgebra (I)
import Prelude hiding (take)
import School.FileIO.FileType (FileType(..))
import School.Types.DataType (DataType(..), fromIdxIndicator, toIdxIndicator)
import School.Types.Decoding (binToInt)
import School.Types.Encoding (intToBin)
import School.Types.Error (Error)
import School.Types.LiftResult (liftResult)
import School.Utils.Constants (separator)

data FileHeader = FileHeader
  { dataType :: DataType
  , rows :: Int
  , cols :: Int
  } deriving (Eq, Show)

instance Default FileHeader where
  def = FileHeader { dataType = DBL64B
                   , cols = 1
                   , rows = 1
                   }

compatibleHeaders :: FileHeader
                  -> FileHeader
                  -> Bool
compatibleHeaders
  FileHeader { dataType = type1, rows = rows1, cols = cols1 }
  FileHeader { dataType = type2, rows = rows2, cols = cols2 }
    = type1 == type2
   && cols1 == cols2
   && rows2 `mod` rows1 == 0

instance ToByteString FileHeader where
  builder header = sep <> t <> r <> c <> sep where
    sep = builder separator
    t = (builder . dataType) header
    r = builder 'r' <> (builder . rows) header
    c = builder 'c' <> (builder . cols) header

intGetter :: Parser I
intGetter = take 4 >>= liftResult . binToInt

parseIdxHeader :: Parser FileHeader
parseIdxHeader = do
  _ <- word8 0
  _ <- word8 0
  t <- fromEnum <$> anyWord8
  dataType <- liftResult $ fromIdxIndicator t
  d <- fromEnum <$> anyWord8
  dims <- (fromIntegral <$>) <$> replicateM d intGetter
  let rows = head dims
  let cols = product . tail $ dims
  return FileHeader { dataType, cols, rows }

parseSmHeader :: Parser FileHeader
parseSmHeader = do
  _ <- char separator
  typeName <- parser
  _ <- char 'r'
  r <- parser
  _ <- char 'c'
  c <- parser
  _ <- char separator
  return $ FileHeader typeName r c

parseHeader :: B.ByteString -> Either Error FileHeader
parseHeader = parseOnly $ parseSmHeader <|> parseIdxHeader

instance FromByteString FileHeader where
  parser = parseSmHeader <|> parseIdxHeader

headerBuilder :: FileType
              -> FileHeader
              -> B.ByteString
headerBuilder SM header = toByteString' header
headerBuilder IDX FileHeader { cols, dataType, rows } =
  let i = toIdxIndicator dataType
  in B.pack (toEnum <$> [0, 0, i, 2])
  <> (intToBin . fromIntegral $ rows)
  <> (intToBin . fromIntegral $ cols)
headerBuilder _ _ = undefined

ext: hs

max_stars_repo_name: jfulseca/School

max_issues_repo_name: jfulseca/School

max_forks_repo_name: jfulseca/School


text:
{-# OPTIONS_GHC -Wno-incomplete-patterns #-}

module VecSpec where

import Data.Complex

import Test.QuickCheck
import Test.QuickCheck.Poly

import Category
import Functor
import Unboxed
import Vec

type N = 10

prop_Vec_Functor_id :: FnProp (Vec3 N A)
prop_Vec_Functor_id = checkFnEqual law_Functor_id

prop_Vec_Functor_comp :: Fun B C -> Fun A B -> FnProp (Vec3 N A)
prop_Vec_Functor_comp f g = checkFnEqual (law_Functor_comp f g)

type UA = Int
type UB = Double
type UC = Complex Double

prop_UVec3_Functor_id :: FnProp (UVec3 N UA)
prop_UVec3_Functor_id = checkFnEqual law_Functor_id

prop_UVec3_Functor_comp ::
    (UB -#> UC) -> (UA -#> UB) -> FnProp (UVec3 N UA)
prop_UVec3_Functor_comp f g = checkFnEqual (law_Functor_comp f g)


text:
{-# LANGUAGE TypeFamilies  #-}
{-# LANGUAGE TypeOperators #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Vectorized
-- Copyright   :  (c) Alberto Ruiz 2007-15
-- License     :  BSD3
-- Maintainer  :  Alberto Ruiz
-- Stability   :  provisional
--
-- Low level interface to vector operations.
--
-----------------------------------------------------------------------------

module Internal.Vectorized where

import           Control.Monad         (when)
import           Data.Complex
import           Foreign.C.String
import           Foreign.C.Types
import           Foreign.Marshal.Alloc (free, malloc)
import           Foreign.Marshal.Array (copyArray, newArray)
import           Foreign.Ptr           (Ptr)
import           Foreign.Storable      (Storable, peek)
import           Internal.Devel
import           Internal.Vector
import           System.IO.Unsafe      (unsafePerformIO)

infixr 1 #
(#) :: TransArray c => c -> (b -> IO r) -> TransRaw c b -> IO r
a # b = applyRaw a b
{-# INLINE (#) #-}

(#!) :: (TransArray c, TransArray c1) => c1 -> c -> TransRaw c1 (TransRaw c (IO r)) -> IO r
a #! b = a # b # id
{-# INLINE (#!) #-}

fromei :: Enum a => a -> CInt
fromei x = fromIntegral (fromEnum x) :: CInt

data FunCodeV = Sin
              | Cos
              | Tan
              | Abs
              | ASin
              | ACos
              | ATan
              | Sinh
              | Cosh
              | Tanh
              | ASinh
              | ACosh
              | ATanh
              | Exp
              | Log
              | Sign
              | Sqrt
              deriving Enum

data FunCodeSV = Scale
               | Recip
               | AddConstant
               | Negate
               | PowSV
               | PowVS
               | ModSV
               | ModVS
               deriving Enum

data FunCodeVV = Add
               | Sub
               | Mul
               | Div
               | Pow
               | ATan2
               | Mod
               deriving Enum

data FunCodeS = Norm2
              | AbsSum
              | MaxIdx
              | Max
              | MinIdx
              | Min
              deriving Enum

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

-- | sum of elements
sumF :: Vector Float -> Float
sumF = sumg c_sumF

-- | sum of elements
sumR :: Vector Float -> Float
sumR = sumg c_sumF

-- | sum of elements
sumQ :: Vector (Complex Float) -> Complex Float
sumQ = sumg c_sumQ

-- | sum of elements
sumC :: Vector (Complex Float) -> Complex Float
sumC = sumg c_sumQ

sumI :: ( TransRaw c (CInt -> Ptr a -> IO CInt) ~ (CInt -> Ptr I -> I :> Ok)
        , TransArray c
        , Storable a
        )
     => I -> c -> a
sumI m = sumg (c_sumI m)

sumL :: ( TransRaw c (CInt -> Ptr a -> IO CInt) ~ (CInt -> Ptr Z -> Z :> Ok)
        , TransArray c
        , Storable a
        ) => Z -> c -> a
sumL m = sumg (c_sumL m)

sumg :: (TransArray c, Storable a) => TransRaw c (CInt -> Ptr a -> IO CInt) -> c -> a
sumg f x = unsafePerformIO $ do
    r <- createVector 1
    (x #! r) f #| "sum"
    return $ r @> 0

type TVV t = t :> t :> Ok

foreign import ccall unsafe "sumF" c_sumF :: TVV Float
-- foreign import ccall unsafe "sumR" c_sumR :: TVV Double
foreign import ccall unsafe "sumQ" c_sumQ :: TVV (Complex Float)
-- foreign import ccall unsafe "sumC" c_sumC :: TVV (Complex Double)
foreign import ccall unsafe "sumI" c_sumI :: I -> TVV I
foreign import ccall unsafe "sumL" c_sumL :: Z -> TVV Z

-- | product of elements
prodF :: Vector Float -> Float
prodF = prodg c_prodF

-- | product of elements
prodR :: Vector Float -> Float
prodR = prodg c_prodF

-- | product of elements
prodQ :: Vector (Complex Float) -> Complex Float
prodQ = prodg c_prodQ

-- | product of elements
prodC :: Vector (Complex Float) -> Complex Float
prodC = prodg c_prodQ

prodI :: I-> Vector I -> I
prodI = prodg . c_prodI

prodL :: Z-> Vector Z -> Z
prodL = prodg . c_prodL

prodg :: (TransArray c, Storable a)
      => TransRaw c (CInt -> Ptr a -> IO CInt) -> c -> a
prodg f x = unsafePerformIO $ do
    r <- createVector 1
    (x #! r) f #| "prod"
    return $ r @> 0

foreign import ccall unsafe "prodF" c_prodF :: TVV Float
-- foreign import ccall unsafe "prodR" c_prodR :: TVV Double
foreign import ccall unsafe "prodQ" c_prodQ :: TVV (Complex Float)
-- foreign import ccall unsafe "prodC" c_prodC :: TVV (Complex Double)
foreign import ccall unsafe "prodI" c_prodI :: I -> TVV I
foreign import ccall unsafe "prodL" c_prodL :: Z -> TVV Z

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

toScalarAux :: (Enum a, TransArray c, Storable a1)
            => (CInt -> TransRaw c (CInt -> Ptr a1 -> IO CInt)) -> a -> c -> a1
toScalarAux fun code v = unsafePerformIO $ do
    r <- createVector 1
    (v #! r) (fun (fromei code)) #|"toScalarAux"
    return (r @> 0)

vectorMapAux :: (Enum a, Storable t, Storable a1)
             => (CInt -> CInt -> Ptr t -> CInt -> Ptr a1 -> IO CInt)
             -> a -> Vector t -> Vector a1
vectorMapAux fun code v = unsafePerformIO $ do
    r <- createVector (dim v)
    (v #! r) (fun (fromei code)) #|"vectorMapAux"
    return r

vectorMapValAux :: (Enum a, Storable a2, Storable t, Storable a1)
                => (CInt -> Ptr a2 -> CInt -> Ptr t -> CInt -> Ptr a1 -> IO CInt)
                -> a -> a2 -> Vector t -> Vector a1
vectorMapValAux fun code val v = unsafePerformIO $ do
    r <- createVector (dim v)
    pval <- newArray [val]
    (v #! r) (fun (fromei code) pval) #|"vectorMapValAux"
    free pval
    return r

vectorZipAux :: (Enum a, TransArray c, Storable t, Storable a1)
             => (CInt -> CInt -> Ptr t -> TransRaw c (CInt -> Ptr a1 -> IO CInt))
             -> a -> Vector t -> c -> Vector a1
vectorZipAux fun code u v = unsafePerformIO $ do
    r <- createVector (dim u)
    (u # v #! r) (fun (fromei code)) #|"vectorZipAux"
    return r

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

-- | obtains different functions of a vector: norm1, norm2, max, min, posmax, posmin, etc.
toScalarR :: FunCodeS -> Vector Float -> Float
toScalarR oper =  toScalarAux c_toScalarF (fromei oper)

-- foreign import ccall unsafe "toScalarR" c_toScalarR :: CInt -> TVV Double

-- | obtains different functions of a vector: norm1, norm2, max, min, posmax, posmin, etc.
toScalarF :: FunCodeS -> Vector Float -> Float
toScalarF oper =  toScalarAux c_toScalarF (fromei oper)

foreign import ccall unsafe "toScalarF" c_toScalarF :: CInt -> TVV Float

-- | obtains different functions of a vector: only norm1, norm2
toScalarC :: FunCodeS -> Vector (Complex Float) -> Float
toScalarC oper =  toScalarAux c_toScalarQ (fromei oper)

-- foreign import ccall unsafe "toScalarC" c_toScalarC :: CInt -> Complex Double :> Double :> Ok

-- | obtains different functions of a vector: only norm1, norm2
toScalarQ :: FunCodeS -> Vector (Complex Float) -> Float
toScalarQ oper =  toScalarAux c_toScalarQ (fromei oper)

foreign import ccall unsafe "toScalarQ" c_toScalarQ :: CInt -> Complex Float :> Float :> Ok

-- | obtains different functions of a vector: norm1, norm2, max, min, posmax, posmin, etc.
toScalarI :: FunCodeS -> Vector CInt -> CInt
toScalarI oper =  toScalarAux c_toScalarI (fromei oper)

foreign import ccall unsafe "toScalarI" c_toScalarI :: CInt -> TVV CInt

-- | obtains different functions of a vector: norm1, norm2, max, min, posmax, posmin, etc.
toScalarL :: FunCodeS -> Vector Z -> Z
toScalarL oper =  toScalarAux c_toScalarL (fromei oper)

foreign import ccall unsafe "toScalarL" c_toScalarL :: CInt -> TVV Z

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

-- | map of real vectors with given function
vectorMapR :: FunCodeV -> Vector Float -> Vector Float
vectorMapR = vectorMapAux c_vectorMapF

-- foreign import ccall unsafe "mapR" c_vectorMapR :: CInt -> TVV Double

-- | map of complex vectors with given function
vectorMapC :: FunCodeV -> Vector (Complex Float) -> Vector (Complex Float)
vectorMapC oper = vectorMapAux c_vectorMapQ (fromei oper)

-- foreign import ccall unsafe "mapC" c_vectorMapC :: CInt -> TVV (Complex Double)

-- | map of real vectors with given function
vectorMapF :: FunCodeV -> Vector Float -> Vector Float
vectorMapF = vectorMapAux c_vectorMapF

foreign import ccall unsafe "mapF" c_vectorMapF :: CInt -> TVV Float

-- | map of real vectors with given function
vectorMapQ :: FunCodeV -> Vector (Complex Float) -> Vector (Complex Float)
vectorMapQ = vectorMapAux c_vectorMapQ

foreign import ccall unsafe "mapQ" c_vectorMapQ :: CInt -> TVV (Complex Float)

-- | map of real vectors with given function
vectorMapI :: FunCodeV -> Vector CInt -> Vector CInt
vectorMapI = vectorMapAux c_vectorMapI

foreign import ccall unsafe "mapI" c_vectorMapI :: CInt -> TVV CInt

-- | map of real vectors with given function
vectorMapL :: FunCodeV -> Vector Z -> Vector Z
vectorMapL = vectorMapAux c_vectorMapL

foreign import ccall unsafe "mapL" c_vectorMapL :: CInt -> TVV Z

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

-- | map of real vectors with given function
vectorMapValR :: FunCodeSV -> Float -> Vector Float -> Vector Float
vectorMapValR oper = vectorMapValAux c_vectorMapValF (fromei oper)

-- foreign import ccall unsafe "mapValR" c_vectorMapValR :: CInt -> Ptr Double -> TVV Double

-- | map of complex vectors with given function
vectorMapValC :: FunCodeSV -> Complex Float -> Vector (Complex Float) -> Vector (Complex Float)
vectorMapValC = vectorMapValAux c_vectorMapValQ

-- foreign import ccall unsafe "mapValC" c_vectorMapValC :: CInt -> Ptr (Complex Double) -> TVV (Complex Double)

-- | map of real vectors with given function
vectorMapValF :: FunCodeSV -> Float -> Vector Float -> Vector Float
vectorMapValF oper = vectorMapValAux c_vectorMapValF (fromei oper)

foreign import ccall unsafe "mapValF" c_vectorMapValF :: CInt -> Ptr Float -> TVV Float

-- | map of complex vectors with given function
vectorMapValQ :: FunCodeSV -> Complex Float -> Vector (Complex Float) -> Vector (Complex Float)
vectorMapValQ oper = vectorMapValAux c_vectorMapValQ (fromei oper)

foreign import ccall unsafe "mapValQ" c_vectorMapValQ :: CInt -> Ptr (Complex Float) -> TVV (Complex Float)

-- | map of real vectors with given function
vectorMapValI :: FunCodeSV -> CInt -> Vector CInt -> Vector CInt
vectorMapValI oper = vectorMapValAux c_vectorMapValI (fromei oper)

foreign import ccall unsafe "mapValI" c_vectorMapValI :: CInt -> Ptr CInt -> TVV CInt

-- | map of real vectors with given function
vectorMapValL :: FunCodeSV -> Z -> Vector Z -> Vector Z
vectorMapValL oper = vectorMapValAux c_vectorMapValL (fromei oper)

foreign import ccall unsafe "mapValL" c_vectorMapValL :: CInt -> Ptr Z -> TVV Z

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

type TVVV t = t :> t :> t :> Ok

-- | elementwise operation on real vectors
vectorZipR :: FunCodeVV -> Vector Float -> Vector Float -> Vector Float
vectorZipR = vectorZipAux c_vectorZipF

-- foreign import ccall unsafe "zipR" c_vectorZipR :: CInt -> TVVV Double

-- | elementwise operation on complex vectors
vectorZipC :: FunCodeVV -> Vector (Complex Float) -> Vector (Complex Float) -> Vector (Complex Float)
vectorZipC = vectorZipAux c_vectorZipQ

-- foreign import ccall unsafe "zipC" c_vectorZipC :: CInt -> TVVV (Complex Double)

-- | elementwise operation on real vectors
vectorZipF :: FunCodeVV -> Vector Float -> Vector Float -> Vector Float
vectorZipF = vectorZipAux c_vectorZipF

foreign import ccall unsafe "zipF" c_vectorZipF :: CInt -> TVVV Float

-- | elementwise operation on complex vectors
vectorZipQ :: FunCodeVV -> Vector (Complex Float) -> Vector (Complex Float) -> Vector (Complex Float)
vectorZipQ = vectorZipAux c_vectorZipQ

foreign import ccall unsafe "zipQ" c_vectorZipQ :: CInt -> TVVV (Complex Float)

-- | elementwise operation on CInt vectors
vectorZipI :: FunCodeVV -> Vector CInt -> Vector CInt -> Vector CInt
vectorZipI = vectorZipAux c_vectorZipI

foreign import ccall unsafe "zipI" c_vectorZipI :: CInt -> TVVV CInt

-- | elementwise operation on CInt vectors
vectorZipL :: FunCodeVV -> Vector Z -> Vector Z -> Vector Z
vectorZipL = vectorZipAux c_vectorZipL

foreign import ccall unsafe "zipL" c_vectorZipL :: CInt -> TVVV Z

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

-- foreign import ccall unsafe "vectorScan" c_vectorScan
--     :: CString -> Ptr CInt -> Ptr (Ptr Double) -> IO CInt

foreign import ccall unsafe "vectorScan" c_vectorScan
    :: CString -> Ptr CInt -> Ptr (Ptr Float) -> IO CInt

vectorScan :: FilePath -> IO (Vector Float)
vectorScan s = do
    pp <- malloc
    pn <- malloc
    cs <- newCString s
    ok <- c_vectorScan cs pn pp
    when (not (ok == 0)) $
        error ("vectorScan: file \"" ++ s ++"\" not found")
    n <- fromIntegral <$> peek pn
    p <- peek pp
    v <- createVector n
    free pn
    free cs
    unsafeWith v $ \pv -> copyArray pv p n
    free p
    free pp
    return v

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

type Seed = Int

data RandDist = Uniform  -- ^ uniform distribution in [0,1)
              | Gaussian -- ^ normal distribution with mean zero and standard deviation one
              deriving Enum

-- | Obtains a vector of pseudorandom elements (use randomIO to get a random seed).
randomVector :: Seed
             -> RandDist -- ^ distribution
             -> Int      -- ^ vector size
             -> Vector Float
randomVector seed dist n = unsafePerformIO $ do
    r <- createVector n
    (r # id) (c_random_vector (fi seed) ((fi.fromEnum) dist)) #|"randomVector"
    return r

-- foreign import ccall unsafe "random_vector" c_random_vector :: CInt -> CInt -> Double :> Ok
foreign import ccall unsafe "random_vector" c_random_vector :: CInt -> CInt -> Float :> Ok

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

roundVector :: Vector Float -> Vector Float
roundVector v = unsafePerformIO $ do
    r <- createVector (dim v)
    (v #! r) c_round_vector #|"roundVector"
    return r

-- foreign import ccall unsafe "round_vector" c_round_vector :: TVV Double

foreign import ccall unsafe "round_vector" c_round_vector :: TVV Float

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

-- |
-- >>> range 5
-- [0,1,2,3,4]
-- it :: Vector I
--
range :: Int -> Vector I
range n = unsafePerformIO $ do
    r <- createVector n
    (r # id) c_range_vector #|"range"
    return r

foreign import ccall unsafe "range_vector" c_range_vector :: CInt :> Ok

-- float2DoubleV :: Vector Float -> Vector Float
-- float2DoubleV x = tog c_float2double

-- double2FloatV :: Vector Double -> Vector Float
-- double2FloatV = tog c_double2float

-- double2IntV :: Vector Double -> Vector CInt
-- double2IntV = tog c_double2int

-- int2DoubleV :: Vector CInt -> Vector Double
-- int2DoubleV = tog c_int2double

-- double2longV :: Vector Double -> Vector Z
-- double2longV = tog c_double2long

long2FloatV :: Vector Z -> Vector Float
long2FloatV = tog c_long2float

float2IntV :: Vector Float -> Vector CInt
float2IntV = tog c_float2int

int2floatV :: Vector CInt -> Vector Float
int2floatV = tog c_int2float

int2longV :: Vector I -> Vector Z
int2longV = tog c_int2long

long2intV :: Vector Z -> Vector I
long2intV = tog c_long2int

long2floatV :: Vector Z -> Vector Float
long2floatV = tog c_long2float

tog :: (Storable t, Storable a)
    => (CInt -> Ptr t -> CInt -> Ptr a -> IO CInt) -> Vector t -> Vector a
tog f v = unsafePerformIO $ do
    r <- createVector (dim v)
    (v #! r) f #|"tog"
    return r

-- foreign import ccall unsafe "float2double" c_float2double :: Float :> Double :> Ok
-- foreign import ccall unsafe "double2float" c_double2float :: Double :> Float :> Ok
-- foreign import ccall unsafe "int2double"   c_int2double   :: CInt :> Double :> Ok
-- foreign import ccall unsafe "double2int"   c_double2int   :: Double :> CInt :> Ok
-- foreign import ccall unsafe "long2double"  c_long2double   :: Z :> Double :> Ok
-- foreign import ccall unsafe "double2long"  c_double2long   :: Double :> Z :> Ok
foreign import ccall unsafe "float2int"   c_float2int   :: Float :> CInt  :> Ok
foreign import ccall unsafe "float2long"   c_float2long :: Float :> I     :> Ok
foreign import ccall unsafe "int2float"   c_int2float   :: CInt  :> Float :> Ok
foreign import ccall unsafe "int2long"    c_int2long    :: I     :> Z     :> Ok
foreign import ccall unsafe "long2int"    c_long2int    :: Z     :> I     :> Ok
foreign import ccall unsafe "long2float"  c_long2float  :: Z     :> Float :> Ok

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

stepg :: (Storable t, Storable a)
      => (CInt -> Ptr t -> CInt -> Ptr a -> IO CInt) -> Vector t -> Vector a
stepg f v = unsafePerformIO $ do
    r <- createVector (dim v)
    (v #! r) f #|"step"
    return r

-- stepD :: Vector Double -> Vector Double
-- stepD = stepg c_stepD

stepF :: Vector Float -> Vector Float
stepF = stepg c_stepF

stepI :: Vector CInt -> Vector CInt
stepI = stepg c_stepI

stepL :: Vector Z -> Vector Z
stepL = stepg c_stepL

foreign import ccall unsafe "stepF" c_stepF :: TVV Float
-- foreign import ccall unsafe "stepD" c_stepD :: TVV Double
foreign import ccall unsafe "stepI" c_stepI :: TVV CInt
foreign import ccall unsafe "stepL" c_stepL :: TVV Z

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

conjugateAux :: (Storable t, Storable a)
             => (CInt -> Ptr t -> CInt -> Ptr a -> IO CInt) -> Vector t -> Vector a
conjugateAux fun x = unsafePerformIO $ do
    v <- createVector (dim x)
    (x #! v) fun #|"conjugateAux"
    return v

conjugateQ :: Vector (Complex Float) -> Vector (Complex Float)
conjugateQ = conjugateAux c_conjugateQ
foreign import ccall unsafe "conjugateQ" c_conjugateQ :: TVV (Complex Float)

-- conjugateC :: Vector (Complex Double) -> Vector (Complex Double)
-- conjugateC = conjugateAux c_conjugateC
-- foreign import ccall unsafe "conjugateC" c_conjugateC :: TVV (Complex Double)

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

cloneVector :: Storable t => Vector t -> IO (Vector t)
cloneVector v = do
        let n = dim v
        r <- createVector n
        let f _ s _ d =  copyArray d s n >> return 0
        (v #! r) f #|"cloneVector"
        return r

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

constantAux :: (Storable a1, Storable a)
            => (Ptr a1 -> CInt -> Ptr a -> IO CInt) -> a1 -> Int -> Vector a
constantAux fun x n = unsafePerformIO $ do
    v <- createVector n
    px <- newArray [x]
    (v # id) (fun px) #|"constantAux"
    free px
    return v

type TConst t = Ptr t -> t :> Ok

foreign import ccall unsafe "constantF" cconstantF :: TConst Float
-- foreign import ccall unsafe "constantR" cconstantR :: TConst Double
foreign import ccall unsafe "constantQ" cconstantQ :: TConst (Complex Float)
-- foreign import ccall unsafe "constantC" cconstantC :: TConst (Complex Double)
foreign import ccall unsafe "constantI" cconstantI :: TConst CInt
foreign import ccall unsafe "constantL" cconstantL :: TConst Z

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

max_stars_repo_name: schnecki/hmatrix-float

max_issues_repo_name: schnecki/hmatrix-float

max_forks_repo_name: schnecki/hmatrix-float


text:
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}

module Myelin.PyNN.PyNN where

import Prelude hiding (id)

import Control.Monad.Except
import Control.Monad.State.Lazy
import Control.Lens
import Data.Aeson
import Data.List (intercalate)
import qualified Data.Map.Strict as Map
import Data.String.Interpolate
import qualified Data.ByteString.Lazy.Char8 as BS
import Numeric.LinearAlgebra
import Text.Regex as Regex

import Myelin.Model
import Myelin.Neuron
import Myelin.SNN

data PyNNPreample
  = PyNNPreample
    { configuration :: String
    }
    deriving (Eq, Show)

data PyNNModel = PyNNModel
  { _declarations :: [String]
  , _layers :: Map.Map String String -- ^ Layers between the nodes
  , _populations :: Map.Map String String -- ^ Nodes indexed by their Python variables
  }
  deriving (Eq, Show)

makeLenses ''PyNNModel

-- | An empty initial model
emptyPyNNModel :: PyNNModel
emptyPyNNModel = PyNNModel [] Map.empty Map.empty

type PyNNState = ExceptT String (State PyNNModel)

-- | Attempts to translate a network into a Python code string
translate :: Task -> PyNNPreample -> Either String String
translate (Task target network simulationTime) preample =
  let (result, (PyNNModel {..})) = runState (runExceptT $ translate' network) emptyPyNNModel
  in  case result of
        Left translationError -> Left $ "Failed to translate SNN model to Python: " ++ translationError
        Right _ ->
          let populationLines = unlines $ Map.elems _populations
              layerLines = unlines $ Map.elems _layers
              declarationLines = unlines $ _declarations
              preampleLines = pyNNPreample target preample
          in Right $ concat [preampleLines, populationLines, layerLines, declarationLines, training]
    where training = [i|
optimiser = v.GradientDescentOptimiser(0.1, simulation_time=#{simulationTime})
if __name__ == "__main__":
    v.Main(model).train(optimiser)
|]

translate' :: Network -> PyNNState ()
translate' Network {..} = do
  _ <- mapM pyNNNode _inputs
  _ <- mapM pyNNNode _nodes
  outputStrings <- mapM pyNNNode _outputs
  _ <- mapM pyNNEdge _edges
  layers <- fmap _layers get
  let decodeLayer = "l_decode = v.Decode(" ++ populationReference ((last $ _outputs) ^.id) ++ ")"
  let layersString = intercalate ", " $ Map.keys layers
  let model = "model = v.Model(" ++ layersString ++ ", l_decode)"
  declarations <>= [decodeLayer, model]

-- | The Python PyNN preamble for import statements
pyNNPreample :: ExecutionTarget -> PyNNPreample -> String
pyNNPreample target (PyNNPreample {..}) = let
 pyNNTarget =
   case target of
      Nest { .. } -> "nest"
      SpiNNaker -> "spiNNaker"
 in [i|import numpy as np
import volrpynn.#{pyNNTarget} as v
import pyNN.#{pyNNTarget} as pynn

#{configuration}

|]

pyNNNode :: Node -> PyNNState String
pyNNNode node = case node of
  Population {..} -> do
    pyNNPopulation _neuronType _numNeurons _id
  _ -> throwError $ "Unknown node type " ++ (show node)

--   SpikeSourceArray { .. } ->
--  SpikeSourcePoisson { .. } -> do

-- | Statefully encodes a pyNNPopulation, returning the variable name for that
--   population
pyNNPopulation :: NeuronType -> Integer -> Int -> PyNNState String
pyNNPopulation tpe numNeurons id =
  case pyNNPopulationString tpe numNeurons id of
    Right codeString -> populationVariable id codeString
    Left errorString -> throwError errorString

typeRegex = Regex.mkRegex "\"type\":\"\\w*\","

-- | Encodes a PyNN population as a string without state
pyNNPopulationString :: NeuronType -> Integer -> Int -> Either String String
pyNNPopulationString tpe numNeurons id =
  let params = BS.unpack $ encode tpe
      stripType = Regex.subRegex typeRegex params ""
  in  case tpe of
        IFCondExp { .. } -> Right $ "pynn.Population(" ++ (show numNeurons) ++ ", pynn.IF_cond_exp(**" ++ stripType ++ "))"
        _ -> Left $ "Unknown neuron type " ++ (show tpe)

-- | Creates and stores an Edge as a PyNN projection
pyNNEdge :: Edge -> PyNNState String
pyNNEdge projection = do
  (PyNNProjection {..}, inputs, outputs) <- pyNNProjection projection
  let popRef1 = concatReferences inputs
  let popRef2 = concatReferences outputs
  let code = "v." ++ layerType ++ "(" ++ popRef1 ++", " ++ popRef2 ++ (weightString weight) ++ (biasString bias) ++ ")"
  layerVariable code
    where
      concatReferences :: [Node] -> String
      concatReferences nodes =
        case length nodes of
          1 -> populationReference (nodes !! 0 ^. id)
          _ -> let strings = intercalate ", " $ map populationReference $ map (\node -> node ^. id) nodes
               in "(" ++ strings ++ ")"
      weightString (Just weight) =  ", weights=" ++ weight
      weightString Nothing = ""
      biasString (Just bias) = ", biases=" ++ bias
      biasString Nothing = ""

data PyNNProjection = PyNNProjection
  { bias :: Maybe String
  , layerType :: String
  , weight :: Maybe String
  }

-- | Create a PyNN Connector definition
pyNNProjection :: Edge-> PyNNState (PyNNProjection, [Node], [Node])
pyNNProjection (DenseProjection (Static Excitatory (AllToAll bias weight)) nodeIn nodeOut) = do
  b <- pyNNBias bias (_numNeurons nodeOut)
  w <- pyNNWeight weight (_numNeurons nodeIn) (_numNeurons nodeOut)
  return $ (PyNNProjection (Just b) "Dense" (Just w), [nodeIn], [nodeOut])
pyNNProjection (MergeProjection _ (nodeIn1, nodeIn2) nodeOut) = do
  let input1Size = (_numNeurons nodeIn1)
  let input2Size = (_numNeurons nodeIn2)
  return $ input1Size + input2Size
  return $ (PyNNProjection Nothing "Merge" Nothing, [nodeIn1, nodeIn2], [nodeOut])
pyNNProjection (ReplicateProjection (Static Excitatory (AllToAll bias weight)) nodeIn (nodeOut1, nodeOut2)) = do
  let inputSize = _numNeurons nodeIn
  let output1Size = _numNeurons nodeOut1
  let output2Size = _numNeurons nodeOut2
  outputSize <- if (output1Size == output2Size)
    then return $ output1Size
    else throwError $ "Replicaten connections require two outputs of the same size, found " ++ (show (output1Size, output2Size))
  w1 <- pyNNWeight weight inputSize output1Size
  w2 <- pyNNWeight weight inputSize output1Size
  let w = [i|(#{w1}, #{w2})|]
  b <- pyNNBias bias output1Size
  return $ (PyNNProjection (Just b) "Replicate" (Just w), [nodeIn], [nodeOut1, nodeOut2])
pyNNProjection p = throwError $ "Unknown projection effect" ++ (show p)

-- | Converts biases to a Python bias setting
pyNNBias :: Biases -> Integer -> PyNNState String
pyNNBias (BiasGenerator (Constant n)) _ = return (show n)
pyNNBias (BiasGenerator (GaussianRandom mean scale)) toSize
  = return [i|np.random.normal(#{mean}, #{scale}, (#{toSize}))|]
pyNNBias (Biases list) toSize = return [i|np.array(#{show list})|]

-- | Converts weights to a PyNN weight setting
pyNNWeight :: Weights -> Integer -> Integer -> PyNNState String
pyNNWeight (WeightGenerator (Constant n)) _ _ = return (show n)
pyNNWeight (WeightGenerator (GaussianRandom mean scale)) fromSize toSize
  = return [i|np.random.normal(#{mean}, #{scale}, (#{fromSize}, #{toSize}))|]
pyNNWeight (Weights matrix) fromSize toSize =
  let (rows, columns) = join bimap toInteger $ size matrix
      str = show $ toLists matrix
  in  if (rows == toSize && columns == fromSize)
      then return [i|np.array(#{str})|]
      else throwError $ [i|Expected sizes (#{toSize},#{fromSize}), but got (#{rows},#{columns})|]

populationReference :: Int -> String
populationReference nodeId = "p" ++ (show nodeId)

populationVariable :: Int -> String -> PyNNState String
populationVariable popId code = do
  let name = "p" ++ (show popId)
  let variable = pythonVariable name code
  populations %= (Map.insert name variable)
  return name

layerVariable :: String -> PyNNState String
layerVariable code = do
  projections <- use layers
  let name = "layer" ++ (show $ Map.size projections)
  let variable = pythonVariable name code
  layers %= (Map.insert name variable)
  return name

pythonVariable :: String -> String -> String
pythonVariable name value = name ++ " = " ++ value

max_stars_repo_name: volr/myelin

max_issues_repo_name: volr/myelin

max_forks_repo_name: volr/myelin


text:
module Chapter08 where

import Numeric.LinearAlgebra.Data
--import Numeric.LinearAlgebra.Algorithms
import Numeric.LinearAlgebra.HMatrix
import Data.List as L

import Chapter04
import Chapter06
import Chapter07

stopWords :: [String]
stopWords = ["a", "about", "above", "above", "across", "after",
    "afterwards", "again", "against", "all", "almost", "alone", "along",
    "already", "also", "although", "always", "am", "among", "amongst",
    "amoungst", "amount", "an", "and", "another", "any", "anyhow",
    "anyone", "anything", "anyway", "anywhere", "are", "around", "as",
    "at", "back", "be", "became", "because", "become", "becomes",
    "becoming", "been", "before", "beforehand", "behind", "being",
    "below", "beside", "besides", "between", "beyond", "bill", "both",
    "bottom", "but", "by", "call", "can", "cannot", "cant", "co", "con",
    "could", "couldnt", "cry", "de", "describe", "detail", "do", "done",
    "dont", "down", "due", "during", "each", "eg", "eight", "either",
    "eleven", "else", "elsewhere", "empty", "enough", "etc", "even",
    "ever", "every", "everyone", "everything", "everywhere", "except",
    "few", "fifteen", "fifty", "fill", "find", "fire", "first", "five",
    "for", "former", "formerly", "forty", "found", "four", "from",
    "front", "full", "further", "get", "give", "go", "got", "had", "has",
    "hasnt", "have", "he", "hence", "her", "here", "hereafter", "hereby",
    "herein", "hereupon", "hers", "herself", "him", "himself", "his",
    "how", "however", "hundred", "i", "ie", "if", "im", "in", "inc",
    "indeed", "interest", "into", "is", "it", "its", "itself", "just",
    "keep", "last", "latter", "latterly", "least", "less", "ltd", "made",
    "many", "may", "me", "meanwhile", "might", "mill", "mine", "more",
    "moreover", "most", "mostly", "move", "much", "must", "my", "myself",
    "name", "namely", "neither", "need", "never", "nevertheless",
    "next", "nine", "no", "nobody", "none", "noone", "nor", "not",
    "nothing", "now", "nowhere", "of", "off", "often", "on", "once",
    "one", "only", "onto", "or", "other", "others", "otherwise", "our",
    "ours", "ourselves", "out", "over", "own", "part", "per", "perhaps",
    "please", "put", "rather", "re", "same", "see", "seem", "seemed",
    "seeming", "seems", "serious", "several", "she", "should", "show",
    "side", "since", "sincere", "six", "sixty", "so", "some", "somehow",
    "someone", "something", "sometime", "sometimes", "somewhere", "still",
    "such", "system", "take", "ten", "than", "that", "the", "their",
    "them", "themselves", "then", "thence", "there", "thereafter",
    "thereby", "therefore", "therein", "thereupon", "these", "they",
    "thick", "thin", "third", "this", "those", "though", "three",
    "through", "throughout", "thru", "thus", "to", "together", "too",
    "top", "toward", "towards", "twelve", "twenty", "two", "un", "under",
    "until", "up", "upon", "us", "very", "via", "want", "was", "we",
    "well", "were", "what", "whatever", "when", "whence", "whenever",
    "where", "whereafter", "whereas", "whereby", "wherein", "whereupon",
    "wherever", "whether", "which", "while", "whither", "who", "whoever",
    "whole", "whom", "whose", "why", "will", "with", "within", "without",
    "would", "yet", "you", "your", "youre", "yours", "yourself",
    "yourselves", "rt", "mt", "u"]

principalComponentAnalysis :: [[Double]] -> Integer -> Matrix Double
principalComponentAnalysis records top =
    tr $ mul (tr topOrderedEvectors) (tr featureVectors)
  where featureVectors = fromLists records
        (_, covMatrix) = meanCov featureVectors
        (_, evectors) = eigSH covMatrix
        topOrderedEvectors = takeColumns (fromInteger top) evectors

euclidianDistanceSqrdFromRow :: Matrix Double -> Integer -> Vector Double
euclidianDistanceSqrdFromRow records row =
    sumOfSquaresVector
  where d = cols records
        copyMatrix = repmat (subMatrix (fromIntegral row, 0) (1,d) records) (rows records) 1
        diffMatrix = records - copyMatrix
        diffSqrdMatrix = diffMatrix * diffMatrix
        sumOfSquaresVector = sum $ toColumns diffSqrdMatrix

order :: Eq a => [a] -> [a] -> [Integer]
order unsorted nubSorted =
    concatMap
        (\x -> L.map toInteger $ L.elemIndices x unsorted)
        nubSorted

findClosestFeaturesToRow :: Matrix Double -> Integer -> Integer -> [Integer]
findClosestFeaturesToRow records row knn =
    take (fromIntegral knn) orderOfFeatures
  where sumOfSquares = toList $ euclidianDistanceSqrdFromRow records row
        orderOfFeatures = order sumOfSquares . nub $ sort sumOfSquares

max_stars_repo_name: ulmalana/haskell-data-analysis

max_issues_repo_name: ulmalana/haskell-data-analysis

max_forks_repo_name: ulmalana/haskell-data-analysis


text:
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Hypergeometric
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The Hypergeometric distribution.  This is the discrete probability
-- distribution that measures the probability of /k/ successes in /l/
-- trials, without replacement, from a finite population.
--
-- The parameters of the distribution describe /k/ elements chosen
-- from a population of /l/, with /m/ elements of one type, and
-- /l/-/m/ of the other (all are positive integers).

module Statistics.Distribution.Hypergeometric
    (
      HypergeometricDistribution
    -- * Constructors
    , hypergeometric
    , hypergeometricE
    -- ** Accessors
    , hdM
    , hdL
    , hdK
    ) where

import Control.Applicative
import Data.Aeson           (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary          (Binary(..))
import Data.Data            (Data, Typeable)
import GHC.Generics         (Generic)
import Numeric.MathFunctions.Constants (m_epsilon,m_neg_inf)
import Numeric.SpecFunctions (choose,logChoose)

import qualified Statistics.Distribution as D
import Statistics.Internal

data HypergeometricDistribution = HD {
      hdM :: {-# UNPACK #-} !Int
    , hdL :: {-# UNPACK #-} !Int
    , hdK :: {-# UNPACK #-} !Int
    } deriving (Eq, Typeable, Data, Generic)

instance Show HypergeometricDistribution where
  showsPrec i (HD m l k) = defaultShow3 "hypergeometric" m l k i
instance Read HypergeometricDistribution where
  readPrec = defaultReadPrecM3 "hypergeometric" hypergeometricE

instance ToJSON HypergeometricDistribution
instance FromJSON HypergeometricDistribution where
  parseJSON (Object v) = do
    m <- v .: "hdM"
    l <- v .: "hdL"
    k <- v .: "hdK"
    maybe (fail $ errMsg m l k) return $ hypergeometricE m l k
  parseJSON _ = empty

instance Binary HypergeometricDistribution where
  put (HD m l k) = put m >> put l >> put k
  get = do
    m <- get
    l <- get
    k <- get
    maybe (fail $ errMsg m l k) return $ hypergeometricE m l k

instance D.Distribution HypergeometricDistribution where
    cumulative = cumulative

instance D.DiscreteDistr HypergeometricDistribution where
    probability    = probability
    logProbability = logProbability

instance D.Mean HypergeometricDistribution where
    mean = mean

instance D.Variance HypergeometricDistribution where
    variance = variance

instance D.MaybeMean HypergeometricDistribution where
    maybeMean = Just . D.mean

instance D.MaybeVariance HypergeometricDistribution where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Entropy HypergeometricDistribution where
  entropy = directEntropy

instance D.MaybeEntropy HypergeometricDistribution where
  maybeEntropy = Just . D.entropy

variance :: HypergeometricDistribution -> Double
variance (HD m l k) = (k' * ml) * (1 - ml) * (l' - k') / (l' - 1)
  where m' = fromIntegral m
        l' = fromIntegral l
        k' = fromIntegral k
        ml = m' / l'

mean :: HypergeometricDistribution -> Double
mean (HD m l k) = fromIntegral k * fromIntegral m / fromIntegral l

directEntropy :: HypergeometricDistribution -> Double
directEntropy d@(HD m _ _)
  = negate . sum
  $ takeWhile (< negate m_epsilon)
  $ dropWhile (not . (< negate m_epsilon))
    [ let x = probability d n in x * log x | n <- [0..m]]

hypergeometric :: Int               -- ^ /m/
               -> Int               -- ^ /l/
               -> Int               -- ^ /k/
               -> HypergeometricDistribution
hypergeometric m l k
  = maybe (error $ errMsg m l k) id $ hypergeometricE m l k

hypergeometricE :: Int               -- ^ /m/
                -> Int               -- ^ /l/
                -> Int               -- ^ /k/
                -> Maybe HypergeometricDistribution
hypergeometricE m l k
  | not (l > 0)            = Nothing
  | not (m >= 0 && m <= l) = Nothing
  | not (k > 0  && k <= l) = Nothing
  | otherwise              = Just (HD m l k)

errMsg :: Int -> Int -> Int -> String
errMsg m l k
  =  "Statistics.Distribution.Hypergeometric.hypergeometric: "
  ++ "m=" ++ show m
  ++ "l=" ++ show l
  ++ "k=" ++ show k
  ++ " should hold: l>0 & m in [0,l] & k in (0,l]"

-- Naive implementation
probability :: HypergeometricDistribution -> Int -> Double
probability (HD mi li ki) n
  | n < max 0 (mi+ki-li) || n > min mi ki = 0
    -- No overflow
  | li < 1000 = choose mi n * choose (li - mi) (ki - n)
              / choose li ki
  | otherwise = exp $ logChoose mi n
                    + logChoose (li - mi) (ki - n)
                    - logChoose li ki

logProbability :: HypergeometricDistribution -> Int -> Double
logProbability (HD mi li ki) n
  | n < max 0 (mi+ki-li) || n > min mi ki = m_neg_inf
  | otherwise = logChoose mi n
              + logChoose (li - mi) (ki - n)
              - logChoose li ki

cumulative :: HypergeometricDistribution -> Double -> Double
cumulative d@(HD mi li ki) x
  | isNaN x      = error "Statistics.Distribution.Hypergeometric.cumulative: NaN argument"
  | isInfinite x = if x > 0 then 1 else 0
  | n <  minN    = 0
  | n >= maxN    = 1
  | otherwise    = D.sumProbabilities d minN n
  where
    n    = floor x
    minN = max 0 (mi+ki-li)
    maxN = min mi ki


text:
#!/usr/bin/env stack
-- stack --resolver lts-14.16 script
{-# LANGUAGE BangPatterns #-}

import Debug.Trace
-- import Control.Monad.State.Lazy

import System.Environment (getArgs)
import System.IO (readFile)
import Data.Map.Strict (Map, (!), insert, elems, fromList, toList, findWithDefault, size, empty, member, findMin, findMax, singleton)
import qualified Data.Map.Strict as M
--import qualified Data.Array as A
import Data.List (permutations)
import Data.List.Split (splitOn)
import Data.Complex (Complex((:+)), realPart, imagPart)

type Instructions = Map Integer Integer

run :: (Integer, Integer) -> Instructions -> [Integer] -> [Integer]
-- (instructionPointer, base), instructions, input, output
run (i, base) instructions inputs =
  -- trace ((show $ toList instructions) ++ ": (ip, base, instr, inputs) " ++ show (i, base, instr, inputs) ++ "\n") $
  case instr `mod` 100 of
    1 -> run (i+4, base) (insert (addr 3) (arg 1 + arg 2) instructions) inputs
    2 -> run (i+4, base) (insert (addr 3) (arg 1 * arg 2) instructions) inputs
    3 -> run (i+2, base) (insert (addr 1) (head inputs) instructions) $ tail inputs
    4 -> arg 1: run (i+2, base) instructions inputs
    5 -> run (if arg 1 == 0 then i+3 else arg 2, base) instructions inputs
    6 -> run (if arg 1 == 0 then arg 2 else i+3, base) instructions inputs
    7 -> run (i+4, base) (insert (addr 3) (if arg 1 < arg 2 then 1 else 0) instructions) inputs
    8 -> run (i+4, base) (insert (addr 3) (if arg 1 == arg 2 then 1 else 0) instructions) inputs
    9 -> run (i+2, base+arg 1) instructions inputs
    99 -> []
    _ -> error "unknown opcode"
  where instr = instructions!i
        ii x = findWithDefault 0 x instructions
        arg :: Integer -> Integer
        arg n = case (instr `mod` (100*10^n)) `div` (10*10^n) of
          0 -> ii $ ii $ i+n
          1 -> ii $ i+n
          2 -> ii ((ii $ i+n) + base)
          _ -> error "bad argument mode"
        addr n = case (instr `mod` (100*10^n)) `div` (10*10^n) of
          0 -> ii $ i+n
          1 -> error "address in mode 1"
          2 -> --trace "address in mode 2" $
            ii (i+n) + base
          _ -> error $ "address in unknown mode"

instance Ord a => Ord (Complex a) where
  compare a b | ar < br || ar == br && ai < bi = LT
              | ar > br || ar == br && ai > bi = GT
              | otherwise = EQ
    where (ar, ai) = toParts a
          (br, bi) = toParts b

toParts :: Complex a -> (a, a)
toParts c = (realPart c, imagPart c)

toIntParts :: RealFrac a => Complex a -> (Int, Int)
toIntParts c = (round $ realPart c, round $ imagPart c)

fromParts :: (a, a) -> Complex a
fromParts (a,b) = a :+ b

process :: (Complex Float, Complex Float, Map (Complex Float) Integer) -> [Integer] -> ([Integer], Complex Float, Complex Float, Map(Complex Float) Integer)
process (p, d, colors) [] =
  -- traceShow (size colors) $
  -- traceShow (showColors p d colors) $
  ([], p, d, colors)
process (p, d, colors) (c: t: futureOutputs) =
  -- traceShow ("xy", toIntParts p,"out", out,robot d, "paint", c,"turn", t) $
  -- trace (showColors p d colors) $
  merge out $ process (p', d', insert p c colors) futureOutputs
  where p' = p + d' -- next position
        d' = d * (0 :+ (1.0 - fromIntegral t * 2)) -- next direction
        out = findWithDefault 0 p' colors
        merge :: Integer -> ([Integer], Complex Float, Complex Float, Map(Complex Float) Integer) -> ([Integer], Complex Float, Complex Float, Map(Complex Float) Integer)
        merge x ~(xs, p, d, m) = (x:xs, p, d, m) -- need the irrefutable match operator ~, see https://stackoverflow.com/questions/59297557/when-can-i-rely-on-haskell-to-read-a-list-lazily/59298311#59298311

-- showProcess :: [Integer] -> (Complex Float, Complex Float, Map (Complex Float) Integer) -> (Complex Float, Complex Float, Map (Complex Float) Integer)
-- showProcess [] s = s
-- showProcess [c] s = trace ("leftover " ++ show c) $ s
-- showProcess (c: t: r) (p, d, m) = showProcess r (p', d', insert p c m)
--   where d' = d * (0 :+ (1.0 - fromIntegral t * 2 ))
--         p' = p + d'

showColors :: Complex Float -> Complex Float -> Map (Complex Float) Integer -> [Char]
showColors p d cs
  -- | null cs = []
  -- | otherwise
  = --traceShow ("showColors", p, robot d, "min", (x0,y0), "max", (x1,y1), "colors size", size cs) $
    unlines [line y | y <- reverse $ [y0 .. y1]]
  where line :: Int -> [Char]
        line y = [ (v (fromIntegral x :+ fromIntegral y) $ findWithDefault 3 (fromIntegral x :+ fromIntegral y) cs) | x <- [x0 .. x1]]
        pts = map fst $ toList cs
        xs = map (round . realPart) $ p: pts
        ys = map (round . imagPart) $ p: pts
        x0, y0, x1, y1 :: Int
        [x0, y0, x1, y1] = [minimum xs, minimum ys, maximum xs, maximum ys]
        v :: Complex Float -> Integer -> Char
        v p' c | p == p' = robot d
               | p' == (0.0 :+ 0.0) && c== 0 = 'o'
               | p' == (0.0 :+ 0.0) && c== 1 = 'O'
               | c == 1 = '#'
               | c == 0 = ' ' -- '.' -- change to space for clarity
               | otherwise = ' '

robot :: Complex Float -> Char
robot c | x == 0 && y == 1 = '^'
        | x == 1 && y == 0 = '>'
        | x == 0 && y == -1 = 'V'
        | x == -1 && y == 0 = '<'
        where (x,y) = toParts c

origin = 0 :+ 0
up = 0 :+ 1
startColors = singleton origin 0
startColors2 = singleton origin 1

main = do
  [instructionFile] <- getArgs
  instructionStrings <- readFile instructionFile
  let instructions = fromList . zip [0 ..] $ map read $ splitOn "," instructionStrings

  putStrLn "Part 1"
  let out = run (0,0) instructions $ (startColors!origin) : processOutput
      (processOutput, finalp, finald, colors) = process (origin,up, startColors) $ out
  putStrLn
    $ unlines
    $ ["done Part 1"
      , show $ ("tiles painted: ", size colors)
      , show $ ("finalp", finalp, "finald", robot finald)
      ]
  putStr $ showColors finalp finald colors

  putStrLn "Part 2"
  let out = run (0,0) instructions $ (startColors2!origin) : processOutput
      (processOutput, finalp, finald, colors) = process (origin,up, startColors2) $ out
  putStrLn
    $ unlines
    $ ["done Part 2"
      , show $ ("tiles painted: ", size colors)
      , show $ ("finalp", finalp, "finald", robot finald)
      ]
  putStr $ showColors finalp finald colors


text:
module Scenes.Hexagons (getScene) where

import Numeric.LinearAlgebra (vector)

import Arc
import Group
import Helpers (degToRad, vectorToFloats)
import Hexagon
import Layer (style)
import Line
import Point
import Scene (Scene, mkSceneWithStyle)
import Shape
import Style

r :: Float
r = 0.35

--rotateArc :: Arc -> Arc
--rotateArc (Arc c ra s e) = Arc c ra (s - (degToRad 90)) (e - (degToRad 90))

connectEndpoints :: (Arc, Arc) -> Line
connectEndpoints (a, b) = Line (snd $ arcCoords a) (snd $ arcCoords b)

duplicate :: (Show a, SvgShape a, Transformable a, Mergable a) => a -> Group a
duplicate = let pts = map (\x -> Point x 0) (vectorToFloats $ vector [0..5] - 6 / 2 + 0.25) in
  Group . translatePoints pts

offsetGroup :: (Show a, SvgShape a, Transformable a, Mergable a) => Group a -> Group a
offsetGroup = translateGroupOverPoints [Point 0.5 (r * 2.1), Point 0.5 (-r * 2.1)]

alignedGroup :: (Show a, SvgShape a, Transformable a, Mergable a) => Group a -> Group a
alignedGroup = translateGroupOverPoints [Point 0 (r * 4.1), Point 0 (-r * 4.1)]

getScene :: IO Scene
getScene = pure $ mkSceneWithStyle 7 5 sceneStyle [hexLayer, arcLayer, lineLayer]
  where
      orange = "#ef6101"
      teal = "#045658"

      hexes = duplicate (Hexagon (Point 0 0) r)
      offsetHexes = offsetGroup hexes
      alignedHexes = alignedGroup hexes
      hexLayer = (toLayer "hexagon" (hexes <> offsetHexes <> alignedHexes)) {style=Just $ withStrokeColor orange}

      leftArcs = duplicate $ Arc (Point 0 0) (r * 1.15) (degToRad 90) (degToRad 180) -- <>
      leftOffsetArcs = offsetGroup leftArcs
      leftAlignedArcs = alignedGroup leftArcs

      rightArcs = duplicate $  Arc (Point 0 0) (r * 1.2) (degToRad 270) (degToRad 0)
      rightOffsetArcs = offsetGroup rightArcs
      rightAlignedArcs = alignedGroup rightArcs

      arcLayer = (toLayer "arcs"
        (leftArcs <> leftOffsetArcs <> leftAlignedArcs <> rightArcs <> rightOffsetArcs <> rightAlignedArcs)) {style=Just $ withStrokeColor teal}

      connectingLines = Group $ zipWith (curry connectEndpoints) (toList rightArcs) (tail (toList leftArcs))
      offsetLines = translate (Point (-0.01) 0) $ offsetGroup connectingLines
      alignedLines = alignedGroup connectingLines
      lineLayer = (toLayer "lines" (connectingLines <> offsetLines <> alignedLines)) {style=Just $ withStrokeColor teal}

      sceneStyle = defaultStyleAttrs { strokeColor=Just "#000000", strokeWidth=Just 0.015 }

max_stars_repo_name: mazelife/cnc-pattern-lib

max_issues_repo_name: mazelife/cnc-pattern-lib

max_forks_repo_name: mazelife/cnc-pattern-lib


text:
-- file: Astro/Tests.hs
-- Tests for some of the library functions

module Astro.Tests
	where

import Astro.Stumpff
import Astro.Kepler
import Astro.Elements

import Control.Monad
import Numeric.LinearAlgebra
import Test.QuickCheck
import Test.QuickCheck.Arbitrary

-- Stumpff c functions
stumpff_c_rec_lhs k x = x * stumpff_c k x
stumpff_c_rec_rhs k x = 1.0/fromIntegral (fac (k-2)) - stumpff_c (k-2) x
	where fac n = product [1..n]
prop_stumpff_c_0 k = stumpff_c k 0.0 == 0
prop_stumpff_c_rec k x =
	let lhs = stumpff_c_rec_lhs k x
	    rhs = stumpff_c_rec_rhs k x
	in
	(k >= 2 && k <= 35 && x /= 0) ==> abs $ (rhs-lhs)/lhs <= 1e-7
	where fac n = product [1..n]
prop_stumpff_G_0 k beta = (stumpff_G k beta 0.0 == 0)
prop_stumpff_G_rec k beta x =
	(k >= 2) ==> (x * stumpff_G k beta x ==
		x^(k-2)/fromIntegral (fac (k-2)) - stumpff_G (k-2) beta x)
	where fac n = product [1..n]

-- Orbital elements
{-
instance Arbitrary (Vector t) where
	arbitrary = liftM fromList (vectorOf 3 arbitrarySizedFractional)
-}
--prop_oe_rv_id :: Vector Double -> Vector Double -> Double -> Property
prop_oe_rv_id :: [Double] -> [Double] -> Double -> Property
prop_oe_rv_id r v gpar =
	property $ (length r == 3 && length v == 3) ==> (rv == rv' && vv == vv')
	where
		rv = fromList r
		vv = fromList v
		(rv',vv') = orbelToPosvel_gp (posvelToOrbel_gp rv vv gpar) gpar

max_stars_repo_name: sageh/AstroHaskell

max_issues_repo_name: sageh/AstroHaskell

max_forks_repo_name: sageh/AstroHaskell


text:
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeFamilies #-}
module Data.Matrix.Static.Sparse
   ( -- * Sparse matrix
     SparseMatrix(..)

   , Zero(..)

    -- * Accessors
    -- ** length information
    , C.dim
    , C.rows
    , C.cols

    -- ** Query
    , (C.!)
    , C.takeDiag

    -- ** Unsafe Query
    , C.unsafeIndex
    , C.unsafeTakeRow
    , C.unsafeTakeColumn
    , unsafeTakeColumnC

    -- * Construction
    , C.empty
    , fromTriplet
    , fromTripletC
    , toTriplet
    , C.fromVector
    , C.fromList
    , C.unsafeFromVector
    , diag
    , diagRect

    -- * Conversions
    , toDense
    , C.flatten
    , C.toList

    -- * Different matrix types
    , C.convertAny
   ) where

import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Storable as S
import qualified Data.Vector.Storable.Mutable as SM
import Data.Singletons
import Control.Monad.ST (runST)
import           Data.Bits                         (shiftR)
import Text.Printf (printf)
import Conduit
import Data.Conduit.Internal (zipSinks)
import Data.Tuple (swap)
import GHC.TypeLits (type (<=))
import Foreign.C.Types
import Data.Complex
import Flat (Flat(..))
import Flat.Instances.Vector ()
import Control.DeepSeq (NFData(..))

import qualified Data.Matrix.Static.Dense as D
import qualified Data.Matrix.Static.Dense.Mutable as DM
import qualified Data.Matrix.Static.Generic as C
import Data.Matrix.Static.Sparse.Mutable

type instance C.Mutable SparseMatrix = MSparseMatrix

class Eq a => Zero a where
    zero :: a

instance Zero Int where
    zero = 0

instance Zero Float where
    zero = 0.0

instance Zero CFloat where
    zero = 0.0

instance Zero Double where
    zero = 0.0

instance Zero (Complex Float) where
    zero = 0

instance Zero (Complex Double) where
    zero = 0

instance Eq a => Zero ([] a) where
    zero = []

-- | Column-major mutable matrix.
data SparseMatrix :: C.MatrixKind where
    SparseMatrix :: (SingI r, SingI c)
                 => (v a)           -- ^ Values: stores the coefficient values
                                     -- of the non-zeros.
                 -> (S.Vector CInt)  -- ^ InnerIndices: stores the row
                                     -- (resp. column) indices of the non-zeros.
                 -> (S.Vector CInt)  -- ^ OuterStarts: stores for each column
                                     -- (resp. row) the index of the first
                                     -- non-zero in the previous two arrays.
                 -> SparseMatrix r c v a

instance NFData (v a) => NFData (SparseMatrix r c v a) where
    rnf (SparseMatrix v1 v2 v3) = rnf (v1, v2, v3)

instance Flat CInt where
    encode x = encode (fromIntegral x :: Int)
    decode = do
        x <- decode
        return $ fromIntegral (x :: Int)
    size x = size (fromIntegral x :: Int)

instance (G.Vector v a, Zero a, Flat (v a), SingI r, SingI c) =>
    Flat (SparseMatrix r c v a) where
        encode mat@(SparseMatrix nnz inner outer) = encode (r, c, nnz, inner, outer)
          where
            (r,c) = C.dim mat
        decode = do
            (r', c', nnz, inner, outer) <- decode
            if r' /= r || c' /= c
                then error $ "Dimensions donot match: " <> show (r,c) <> " /= " <> show (r',c')
                else return $ SparseMatrix nnz inner outer
          where
            r = fromIntegral $ fromSing (sing :: Sing r) :: Int
            c = fromIntegral $ fromSing (sing :: Sing c) :: Int
        size mat@(SparseMatrix nnz inner outer) = size (r, c, nnz, inner, outer)
          where
            (r,c) = C.dim mat

instance (G.Vector v a, Eq (v a)) => Eq (SparseMatrix r c v a) where
    (==) (SparseMatrix a b c) (SparseMatrix a' b' c') =
        a == a' && b == b' && c == c'

instance (G.Vector v a, Zero a, Show a) => Show (SparseMatrix r c v a) where
    show mat = printf "(%d x %d)\n%s" r c vals
      where
        (r,c) = C.dim mat
        vals = unlines $ map (unwords . map show . G.toList) $ C.toRows mat

instance (G.Vector v a, Zero a) => C.Matrix SparseMatrix v a where
    -- | O(1) Return the size of matrix.
    dim :: forall r c. SparseMatrix r c v a -> (Int, Int)
    dim (SparseMatrix _ _ _) = (r,c)
      where
        r = fromIntegral $ fromSing (sing :: Sing r)
        c = fromIntegral $ fromSing (sing :: Sing c)
    {-# INLINE dim #-}

    -- | O(1) Unsafe indexing without bound check.
    unsafeIndex (SparseMatrix vec inner outer) (i,j) =
        case binarySearchByBounds inner (fromIntegral i) r0 r1 of
            Nothing -> zero
            Just k -> vec `G.unsafeIndex` k
      where
        r0 = fromIntegral $ outer `S.unsafeIndex` j
        r1 = fromIntegral $ outer `S.unsafeIndex` (j+1) - 1
    {-# INLINE unsafeIndex #-}

    unsafeTakeColumn mat i = G.create $ do
        vec <- GM.replicate (C.rows mat) zero
        let f (r,_,v) = GM.unsafeWrite vec r v
        runConduit $ unsafeTakeColumnC mat i .| mapM_C f
        return vec
    {-# INLINE unsafeTakeColumn #-}

    -- | O(1) Create matrix from vector containing columns.
    unsafeFromVector :: forall r c. (G.Vector v a, SingI r, SingI c)
           => v a -> SparseMatrix r c v a
    unsafeFromVector vec = fromTriplet vec'
      where
        vec' = V.fromList $ map (\((a,b),c) -> (a,b,c)) $ filter ((/=zero) . snd) $
            zipWith (\i x -> (toIndex i, x)) [0..] $ G.toList vec
        toIndex i = swap $ i `divMod` r
        r = fromIntegral $ fromSing (sing :: Sing r)
    {-# INLINE unsafeFromVector #-}

    transpose mat = runIdentity $ fromTripletC source
      where
        source = toTriplet mat .| mapC (\(i,j,x) -> (j,i,x))
    {-# INLINE transpose #-}

    thaw = undefined
    {-# INLINE thaw #-}

    unsafeThaw = undefined
    {-# INLINE unsafeThaw #-}

    freeze = undefined
    {-# INLINE freeze #-}

    unsafeFreeze = undefined
    {-# INLINE unsafeFreeze #-}

    map f (SparseMatrix vec inner outer) = SparseMatrix (G.map f vec) inner outer
    {-# INLINE map #-}

    imap f mat@(SparseMatrix _ inner outer) = SparseMatrix vec' inner outer
      where
        vec' = runST $ runConduit $ toTriplet mat .| mapC g .| sinkVector
        g (i,j,x) = f (i,j) x
    {-# INLINE imap #-}

    imapM_ f mat@(SparseMatrix _ _ _) = runConduit $ toTriplet mat .| mapM_C g
      where
        g (i,j,x) = f (i,j) x >> return ()
    {-# INLINE imapM_ #-}

    sequence (SparseMatrix vec inner outer) = do
        vec' <- G.sequence vec
        return $ SparseMatrix vec' inner outer
    {-# INLINE sequence #-}

    sequence_ (SparseMatrix vec _ _) = G.sequence_ vec
    {-# INLINE sequence_ #-}

toDense :: (Zero a, G.Vector v a, SingI r, SingI c)
        => SparseMatrix r c v a -> D.Matrix r c v a
toDense mat = D.create $ do
    m <- DM.replicate zero
    flip C.imapM_ mat $ \idx -> DM.unsafeWrite m idx
    return m
{-# INLINE toDense #-}

-- | Stream a column.
unsafeTakeColumnC :: (Monad m, G.Vector v a)
                  => SparseMatrix r c v a -> Int -> ConduitT i (Int, Int, a) m ()
unsafeTakeColumnC (SparseMatrix nnz inner outer) i = enumFromToC lo hi .| mapC f
  where
    f idx = (fromIntegral $ inner `S.unsafeIndex` idx, i, nnz `G.unsafeIndex` idx)
    lo = fromIntegral $ outer S.! i
    hi = fromIntegral $ outer S.! (i+1) - 1
{-# INLINE unsafeTakeColumnC #-}

-- | O(n) Create matrix from triplet. row and column indices *are not* assumed to be ordered
-- duplicate entries are carried over to the CSR represention
fromTriplet :: forall u r c v a. (G.Vector u (Int, Int, a), G.Vector v a, SingI r, SingI c)
            => u (Int, Int, a) -> SparseMatrix r c v a
fromTriplet triplets = SparseMatrix val inner outer
  where
    outer = S.scanl (+) 0 $ S.create $ do
        vec <- SM.replicate c 0
        G.forM_ triplets $ \(i, j, _) -> if i < r && j < c
            then SM.unsafeModify vec (+1) j
            else error $ printf "Index out of bound: (%d, %d) >= (%d, %d)" i j r c
        return vec
    (val, inner) = runST $ do
        outer' <- S.thaw outer
        val' <- GM.new nnz
        inner' <- SM.new nnz
        G.forM_ triplets $ \(i, j, v) -> do
            idx <- fromIntegral <$> SM.unsafeRead outer' j
            GM.unsafeWrite val' idx v
            SM.unsafeWrite inner' idx $ fromIntegral i
            SM.unsafeModify outer' (+1) j
        (,) <$> G.unsafeFreeze val' <*> S.unsafeFreeze inner'
    nnz = G.length triplets
    r = fromIntegral $ fromSing (sing :: Sing r)
    c = fromIntegral $ fromSing (sing :: Sing c)
{-# INLINE fromTriplet #-}

-- | O(n) Create matrix from triplet. Row and column indices *are not* assumed
-- to be ordered. Duplicate entries are carried over to the CSR represention.
-- NOTE: The Conduit will be consumed twice. Use `fromTriplet` if generating
-- the Conduit is expensive.
fromTripletC :: forall m r c v a. (Monad m, G.Vector v a, SingI r, SingI c)
             => ConduitT () (Int, Int, a) m ()
             -> m (SparseMatrix r c v a)
fromTripletC triplets = do
    (nnz, outer) <- runConduit $ triplets .| zipSinks lengthC sinkOuter
    (val, inner, _) <- runConduit $ triplets .| sinkValInner nnz (clone outer)
    return $ SparseMatrix val inner outer
  where
    sinkOuter = S.scanl (+) 0 <$> foldlC f (S.replicate c 0)
      where
        f vec (_, j, _) = S.modify (\v -> SM.unsafeModify v (+1) j) vec
    sinkValInner nnz outer0 = foldlC f (val0, inner0, outer0)
      where
        val0 = G.create $ GM.new nnz
        inner0 = S.create $ SM.new nnz
        f (val, inner, outer) (i, j, v) = (val', inner', outer')
          where
            idx = fromIntegral $ outer `S.unsafeIndex` j
            val' = G.create $ do
                vec <- G.unsafeThaw val
                GM.unsafeWrite vec idx v
                return vec
            inner' = S.create $ do
                vec <- S.unsafeThaw inner
                SM.unsafeWrite vec idx $ fromIntegral i
                return vec
            outer' = S.create $ do
                vec <- S.unsafeThaw outer
                SM.unsafeModify vec (+1) j
                return vec
    c = fromIntegral $ fromSing (sing :: Sing c)
    clone x = S.create $ S.thaw x
{-# INLINE fromTripletC #-}

-- | Convert sparse matrix to triplets in column order.
toTriplet :: (Monad m, G.Vector v a, SingI r, SingI c)
          => SparseMatrix r c v a -> ConduitT i (Int, Int, a) m ()
toTriplet (SparseMatrix val inner outer) =
    G.ifoldM_ go (fromIntegral $ G.head outer) outer
  where
    go start curC end = do
        enumFromToC start (end'-1) .| mapC f
        return end'
      where
        end' = fromIntegral end
        f i = (fromIntegral $ inner `G.unsafeIndex` i, fromIntegral curC - 1, val `G.unsafeIndex` i)
{-# INLINE toTriplet #-}

-- | O(m*n) Create a rectangular matrix with default values and given diagonal
diag :: (G.Vector v a, Zero a, SingI n)
     => D.Matrix n 1 v a       -- ^ diagonal
     -> SparseMatrix n n v a
diag = diagRect
{-# INLINE diag #-}

-- | O(m*n) Create a rectangular matrix with default values and given diagonal
diagRect :: (G.Vector v a, Zero a, SingI r, SingI c, n <= r, n <= c)
         => D.Matrix n 1 v a       -- ^ diagonal
         -> SparseMatrix r c v a
diagRect d = SparseMatrix (C.flatten d) (S.enumFromN 0 n) (S.enumFromN 0 $ n + 1)
  where
    n = C.rows d
{-# INLINE diagRect #-}

binarySearchByBounds :: S.Vector CInt -> CInt -> Int -> Int -> Maybe Int
binarySearchByBounds vec x = loop
  where
    loop !l !u
        | l > u = Nothing
        | x == x' = Just k
        | x < x' = loop l (k-1)
        | otherwise = loop (k+1) u
      where
        k = (u+l) `shiftR` 1
        x' = vec `S.unsafeIndex` k
{-# INLINE binarySearchByBounds #-}

-------------------------------------------------------------------------------
-- Helper
-------------------------------------------------------------------------------

--getIndex :: Int -> (Int, Int)
--getIndex =
--{-# INLINE getIndex #-}

max_stars_repo_name: kaizhang/matrix-sized

max_issues_repo_name: kaizhang/matrix-sized

max_forks_repo_name: kaizhang/matrix-sized


text:
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-deprecations #-}
module Amby.Plot
  ( getEC
  , getRenderable
  , save
  , saveSvg

  , cairoDefSave
  , diagramsDefSave
  ) where

import Control.Arrow
import Control.Monad
import Control.Monad.State
import qualified Data.List as L
import qualified Data.Maybe as Maybe
import Data.Tuple (swap)
import qualified Data.Vector.Generic as G
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U

import Control.Lens
import Data.Colour as Colour
import Data.Default.Class
import Graphics.Rendering.Chart.Easy (Layout, EC, Renderable, LayoutPick)
import qualified Graphics.Rendering.Chart.Easy as Chart
import qualified Graphics.Rendering.Chart.Backend.Cairo as Cairo
import qualified Graphics.Rendering.Chart.Backend.Diagrams as Diagrams
import Safe (maximumMay)
import qualified Statistics.Sample.KernelDensity.Simple as Stats
import qualified Statistics.Sample as Stats

import Amby.BoxPlot
import Amby.Compatibility.HistogramPlot
import Amby.FactorPlot
import Amby.Numeric
import Amby.Theme (toColour)
import Amby.Types

--------------------------
-- Instance declarations
--------------------------

instance AmbyContainer (V.Vector Double) where
  type Value (V.Vector Double) = Double

  plot = plotVec
  plot' = plotVec'
  plotEq = plotEqVec
  plotEq' = plotEqVec'
  distPlot = distPlotVec
  distPlot' = distPlotVec'
  kdePlot = kdePlotVec
  kdePlot' = kdePlotVec'
  rugPlot = rugPlotVec
  rugPlot' = rugPlotVec'
  boxPlot = boxPlotVec
  boxPlot' = boxPlotVec'
  factorPlot = factorPlotVec

instance AmbyContainer (U.Vector Double) where
  type Value (U.Vector Double) = Double

  plot = plotVec
  plot' = plotVec'
  plotEq = plotEqVec
  plotEq' = plotEqVec'
  distPlot = distPlotVec
  distPlot' = distPlotVec'
  kdePlot = kdePlotVec
  kdePlot' = kdePlotVec'
  rugPlot = rugPlotVec
  rugPlot' = rugPlotVec'
  boxPlot = boxPlotVec
  boxPlot' = boxPlotVec'
  factorPlot v = factorPlotVec (G.convert v)

instance (Real a) => AmbyContainer [a] where
  type Value [a] = a

  plot x y optsState = plotList vals opts
    where
      opts = execState optsState def
      vals = L.zipWith realToTuple x y
  plot' x y = plot x y $ return ()

  plotEq x fn optsState = plotList vals $ def
      & color .~ (opts ^. color)
    where
      opts = execState optsState def
      vals = L.zipWith realToTuple x $ map fn x
  plotEq' x fn = plotEq x fn $ return ()

  distPlot xs optsState = distPlotVec (realToVec xs) optsState
  distPlot' xs = distPlot xs $ return ()

  kdePlot xs optsState = kdePlotVec (realToVec xs) optsState
  kdePlot' xs = kdePlot xs $ return ()

  rugPlot xs optsState = rugPlotVec (realToVec xs) optsState
  rugPlot' xs = rugPlot xs $ return ()

  boxPlot xs optsState = boxPlotVec (realToVec xs) optsState
  boxPlot' xs = boxPlot xs $ return ()

  factorPlot xs optsState = factorPlotVec (realToVec xs) optsState

realToVec :: (Real a) => [a] -> V.Vector Double
realToVec = V.map realToFrac . V.fromList

realToTuple :: Real a => a -> a -> (Double, Double)
realToTuple a b = (realToFrac a, realToFrac b)

--------------------------
-- Generic vec plotters
--------------------------

plotVec :: (G.Vector v Double, G.Vector v (Double, Double))
        => v Double -> v Double -> State PlotOpts () -> AmbyChart ()
plotVec x y optsState = plotList vals opts
  where
    opts = execState optsState def
    vals = G.toList $ G.zip x y

plotVec' :: (G.Vector v Double, G.Vector v (Double, Double))
         => v Double -> v Double -> AmbyChart ()
plotVec' x y = plotVec x y $ return ()

plotEqVec :: (G.Vector v Double, G.Vector v (Double, Double))
          => v Double -> (Double -> Double) -> State PlotEqOpts ()
          -> AmbyChart ()
plotEqVec x fn optsState = plotList vals $ def
    & color .~ (opts ^. color)
  where
    opts = execState optsState def
    vals = G.toList $ G.zip x (G.map fn x)

plotEqVec' :: (G.Vector v Double, G.Vector v (Double, Double))
           => v Double -> (Double -> Double) -> AmbyChart ()
plotEqVec' x fn = plotEqVec x fn $ return ()

distPlotVec :: (G.Vector v Double, G.Vector v (Double, Double))
            => v Double -> State DistPlotOpts ()
            -> AmbyChart ()
distPlotVec xs optsState = plotDistribution xs $ execState optsState def

distPlotVec' :: (G.Vector v Double, G.Vector v (Double, Double))
             => v Double -> AmbyChart ()
distPlotVec' xs = distPlotVec xs $ return ()

kdePlotVec :: (G.Vector v Double) => v Double -> State KdePlotOpts ()
           -> AmbyChart ()
kdePlotVec xs optsState = do
    layout <- takeLayout
    putLayout $ do
      layout
      nextColor <- toColour (opts ^. color) <$> Chart.takeColor
      plotKde nextColor xs opts
  where
    opts = execState optsState def

kdePlotVec' :: (G.Vector v Double) => v Double -> AmbyChart ()
kdePlotVec' xs = kdePlotVec xs $ return ()

rugPlotVec :: (G.Vector v Double, G.Vector v (Double, Double))
           => v Double -> State RugPlotOpts ()
           -> AmbyChart ()
rugPlotVec xs optsState = do
    layout <- takeLayout
    putLayout $ do
      layout
      nextColor <- toColour (opts ^. color) <$> Chart.takeColor
      plotRug nextColor xs opts
  where
    opts = execState optsState def

rugPlotVec' :: (G.Vector v Double, G.Vector v (Double, Double))
            => v Double -> AmbyChart ()
rugPlotVec' xs = rugPlotVec xs $ return ()

--------------------------
-- Generic plotters
--------------------------

plotDistribution :: (G.Vector v Double, G.Vector v (Double, Double))
                 => v Double -> DistPlotOpts -> AmbyChart ()
plotDistribution xs opts = do
    layout <- takeLayout
    putLayout $ do
      layout
      nextColor <- toColour (opts ^. color) <$> Chart.takeColor
      when (opts ^. hist) $ plotHist opts nextColor numBins xs
      when (opts ^. kde) $ plotKde nextColor xs $ def
        & bw .~ (opts ^. bw)
        & shade .~ (opts ^. shade)
        & gridsize .~ (opts ^. gridsize)
        & axis .~ (opts ^. axis)
        & linewidth .~ (opts ^. kdeLinewidth)
        & cut .~ (opts ^. cut)
      when (opts ^. rug) $ plotRug nextColor xs $ def
        & height .~ (opts ^. rugHeight)
        & axis .~ (opts ^. axis)
        & linewidth .~ (opts ^. rugLinewidth)
  where
    numBins = if opts ^. bins /= 0
      then opts ^. bins
      else min 50 $ freedmanDiaconisBins xs

plotList :: [(Double, Double)] -> PlotOpts -> AmbyChart ()
plotList xs opts = do
    layout <- takeLayout
    putLayout $ do
      layout
      nextColor <- toColour (opts ^. color) <$> Chart.takeColor
      plotLine nextColor (opts ^. linewidth) xs

--------------------------
-- Composable plotters
--------------------------

plotHist :: (G.Vector v Double) => DistPlotOpts -> AlphaColour Double -> Int -> v Double
         -> EC (Layout Double Double) ()
plotHist opts c numBins xs =
  Chart.plot $ fmap histToPlot $ Chart.liftEC $ do
    plot_hist_bins .= numBins
    plot_hist_drop_lines .= True
    plot_hist_values .= G.toList xs
    plot_hist_line_style . Chart.line_width .= (opts ^. histLinewidth)
    plot_hist_line_style . Chart.line_color .= Colour.dissolve 0.5 c
    plot_hist_fill_style . Chart.fill_color .= Colour.dissolve 0.4 c
    plot_hist_vertical .= (opts ^. axis == YAxis)

plotRug :: forall v. (G.Vector v Double, G.Vector v (Double, Double))
        => AlphaColour Double -> v Double -> RugPlotOpts
        -> EC (Layout Double Double) ()
plotRug c xs opts = do
      maxMay <- getPlotMaxHeightMay
      h <- calcRugHeight maxMay
      Chart.plot $ return $ Chart.Plot
        { Chart._plot_render = renderPlotRug h
        , Chart._plot_legend = [("", const $ return ())]
        , Chart._plot_all_points = join (***) G.toList
          $ G.unzip
          $ G.concatMap (toPoints h) xs
        }
      updatePlotHeight maxMay
    where
      getPlotMaxHeightMay :: EC (Layout Double Double) (Maybe Double)
      getPlotMaxHeightMay = do
        layout <- get
        layout ^.. Chart.layout_plots . each . Chart.plot_all_points
          & concatMap (if opts ^. axis == XAxis then snd else fst)
          & maximumMay
          & return

      toPoints :: Double -> Double -> v (Double, Double)
      toPoints h x = if opts ^. axis == XAxis
        then G.fromList [(x, 0), (x, h)]
        else G.fromList [(0, x), (h, x)]

      calcRugHeight :: (Maybe Double) -> EC (Layout Double Double) Double
      calcRugHeight maxMay =
        return $ (Maybe.fromMaybe 1.0 maxMay) * (opts ^. height)

      renderPlotRug :: Double -> Chart.PointMapFn Double Double
                    -> BackendProgram ()
      renderPlotRug h pmap =
        Chart.withLineStyle lineStyle $ do
          Chart.alignStrokePath (plotPath h pmap) >>= Chart.strokePath
        where
          lineStyle = def
            & Chart.line_width .~ (opts ^. linewidth)
            & Chart.line_color .~ c

      plotPath :: Double -> Chart.PointMapFn Double Double -> Chart.Path
      plotPath h pmap =
          go (G.toList xs)
        where
          go [] = Chart.End
          go [b] = Chart.MoveTo (pt b 0) $ Chart.LineTo (pt b h) $ Chart.End
          go (b:rest) = Chart.MoveTo (pt b 0) $ Chart.LineTo (pt b h) $ go rest
          pt x y = pmap (Chart.LValue x, Chart.LValue y)

      updatePlotHeight :: (Maybe Double)
                       -> EC (Layout Double Double) ()
      updatePlotHeight maxMay = do
        layout2 <- get
        layout2 & Chart.layout_plots . mapped . Chart.plot_all_points
          %~ updatePlotPoints maxMay
          & put

      updatePlotPoints :: (Maybe Double) -> ([Double], [Double])
                       -> ([Double], [Double])
      updatePlotPoints maxMay (as, bs) =
        if | maxMay == Nothing && opts ^. axis == XAxis -> (0.0:as, 1.0:bs)
           | maxMay == Nothing && opts ^. axis == YAxis -> (1.0:as, 0.0:bs)
           | otherwise -> (as, bs)

plotKde :: (G.Vector v Double) => AlphaColour Double -> v Double -> KdePlotOpts
        -> EC (Layout Double Double) ()
plotKde c xs opts = do
    plotLine c (opts ^. linewidth) kdeVals
    when (opts ^. shade) $ plotShade c kdeVals
  where
    minX = G.minimum xs
    maxX = G.maximum xs
    iqr = interquartileRange xs
    std = Stats.stdDev xs
    n = fromIntegral $ G.length xs
    a = min std (iqr / 1.34)
    bandwidth = if | opts ^. bw == Scott -> 1.059 * a * n ** ((-1) / 5)
                   | (BwScalar b) <- opts ^. bw -> b
    clipLeft = minX - (opts ^. cut) * bandwidth
    clipRight = maxX + (opts ^. cut) * bandwidth
    sizeEstimate = fromIntegral $ opts ^. gridsize :: Double
    numEstimates = 2 ^ (ceiling (log sizeEstimate / log 2) :: Int)
    xVals = Stats.Points $ linspace clipLeft clipRight numEstimates
    kdeResult = Stats.estimatePDF Stats.gaussianKernel bandwidth xs xVals
    kdeVals' = U.zip (Stats.fromPoints xVals) kdeResult
    kdeVals = if opts ^. axis == YAxis
      then U.toList $ U.map swap kdeVals'
      else U.toList kdeVals'

plotLine :: AlphaColour Double -> Double -> [(Double, Double)]
         -> EC (Layout Double Double) ()
plotLine c lineWidth xs = Chart.plot $ Chart.liftEC $ do
  Chart.plot_lines_values .= [xs]
  Chart.plot_lines_style . Chart.line_width .= lineWidth
  Chart.plot_lines_style . Chart.line_color .= c

plotShade :: AlphaColour Double -> [(Double, Double)]
          -> EC (Layout Double Double) ()
plotShade c xs = Chart.plot $ Chart.liftEC $ do
  Chart.plot_fillbetween_values .= (xs & mapped . _2 %~ \y -> (0, y))
  Chart.plot_fillbetween_style .= Chart.solidFillStyle (Colour.dissolve 0.25 c)

--------------------------
-- Rendering
--------------------------

-- | Default save filename for Cairo.
cairoDefSave :: FilePath
cairoDefSave = ".__amby.png"

-- | Default save filename for Diagrams.
diagramsDefSave :: FilePath
diagramsDefSave = ".__amby.svg"

-- | Convert 'AmbyChart' into Chart's 'EC (Layout Double Double) ()'.
getEC :: AmbyChart () -> EC (Layout Double Double) ()
getEC compute = getLayout $ execState compute def

-- | Convert 'AmbyGrid' into Chart's 'Renderable a'.
getRenderable :: Saveable s => s -> Renderable (LayoutPick Double Double Double)
getRenderable ch = getSaveObjectRenderable (toSaveObject ch)

-- Short-hand to render to png file using Cairo backend.
save :: Saveable s => s -> IO ()
save chart = void $ Cairo.renderableToFile
    def { Cairo._fo_size = newSize }
    cairoDefSave
    (getSaveObjectRenderable so)
  where
    so = toSaveObject chart
    newSize = getSize so

-- | Short-hand to render to svg using Cairo backend
saveSvg :: Saveable s => s -> IO ()
saveSvg chart = void $ Diagrams.renderableToFile
    def { Diagrams._fo_size = join (***) fromIntegral newSize }
    diagramsDefSave
    (getSaveObjectRenderable so)
  where
    so = toSaveObject chart
    newSize = getSize so

max_stars_repo_name: jsermeno/amby

max_issues_repo_name: jsermeno/amby

max_forks_repo_name: jsermeno/amby