{-# LANGUAGE FlexibleInstances #-}

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, filter)
import qualified Data.Map.Strict as M
import Data.Set (Set)
import qualified Data.Set as S
--import qualified Data.Array as A
import Data.List (find, intercalate, intersperse, permutations, inits, tails, isPrefixOf)
import Data.List.Split (splitOn)
-- import Data.Complex (Complex((:+)), realPart, imagPart) -- define my own complex

import UI.NCurses
import Data.Char (chr, ord)
-- import Data.Complex (Complex((:+)))
import Data.Maybe (fromJust)
type Instructions = Map Integer Integer

data ComputerState = Ready | Running | Blocked | Done
  deriving (Show, Eq)

data Computer = Computer {
  state :: ComputerState,
  memory :: Map Integer Integer,
  iptr :: Integer, -- instruction pointer
  base :: Integer, -- base offset
  input :: [Integer],
  output :: [Integer]}

computer0 = Computer { state = Ready, memory = empty, iptr = 0, base = 0, input = [], output = [] }

run :: Computer -> Computer -- (Integer, Integer) -> Instructions -> [Integer] -> [Integer]
run c
  | state c == Ready = run $ c {state = Running}
  | state c == Done = c
  | state c == Blocked = if input c == [] then c else run $ c {state = Running}
  | otherwise = -- running
    --traceShow (iptr c, base c, memory c) $
    case instr `mod` 100 of
      1 -> -- add
        run $ c {iptr = i+4, memory = insert (addr 3) (arg 1 + arg 2) instructions}
      2 -> -- multiply
        run $ c {iptr = i+4, memory = insert (addr 3) (arg 1 * arg 2) instructions}
      3 -> -- read input
        if (null $ input c)
        then c {state = Blocked}
        else run $ c {iptr = i+2, memory = insert (addr 1) (head $ input c) instructions, input = tail $ input c}
      4 -> -- output
        --traceShow (arg 1) $
        run $ c {iptr = i+2, output = output c ++ [arg 1]}
      5 -> -- jump-if-true
        run $ c {iptr = if arg 1 == 0 then i+3 else arg 2}
      6 -> -- jump-if-false
        run $ c {iptr = if arg 1 == 0 then arg 2 else i+3}
      7 -> -- less than
        run $ c {iptr = i+4, memory = insert (addr 3) (if arg 1 < arg 2 then 1 else 0) instructions}
      8 -> -- equals
        run $ c {iptr = i+4, memory = insert (addr 3) (if arg 1 == arg 2 then 1 else 0) instructions}
      9 -> -- set relative base
        run $ c {iptr = i+2, base = base c + arg 1}
      99 -> -- halt
        run $ c {state = Done}
      _ -> error "unknown opcode"
  where instructions = memory c
        i = iptr c
        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 c)
          _ -> 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 c
          _ -> error $ "address in unknown mode"

-- instance (Show a) => Show (Complex a) where
--   show (C x y) = "C " ++ show x ++ " " ++ show y

-- instance (Eq a) => Eq (Complex a) where
--   (C a ai) == (C b bi) = a == b && ai == bi

-- instance (Eq a, Ord a) => Ord (Complex a) where
--   compare (C a ai) (C b bi) | a < b || a == b && ai < bi = LT
--                             | a > b || a == b && ai > bi = GT
--                             | otherwise = EQ

-- -- Note that we are using a plane where imaginary (y-axis) numbers
-- -- increase downward. I think this means that right and left are
-- -- reversed.

-- data Complex a = C a a

-- -- instance Num (Complex Int) where
-- --   abs c = error "Can't take abs of Complex Int"

-- instance Num a => Num (Complex a) where
--   (C x y) + (C u v) = C (x+u) (y+v)
--   (C x y) * (C u v) = C (x*u-y*v) (x*v+y*u)
--   fromInteger n = C (fromInteger n) 0
--   abs (C x y) = C (abs x) y -- sqrt $ x*x + y*y
--   signum (C x y) = C (signum x) 0
--   negate (C x y) = C (negate x) (negate y)

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

  putStrLn "Part 1"
  let linesToAscii = map (fromIntegral . ord) . unlines
  -- Jumps are 4 spaces. Jump on an island if you can. # = ground, . = hole, _ = don't care
  -- Jump if any of the three spaces in front of you have holes and the fourth space is solid.
  -- @ A B C D
  -- _ . _ _ #
  -- _ _ . _ #
  -- _ _ _ . #
  -- ~(ABC)D -> ~((ABC)+D)
  let ls = [ "OR A J"
           , "AND B J"
           , "AND C J"
           , "NOT J J"
           , "AND D J"
           , "WALK"]

  -- let c = run $ computer0{memory = instructions, input=ls}
  let c = run $ computer0{memory = instructions}
  putStrLn $ map (chr . fromIntegral) $ output c
  putStrLn $ "Running with input: \n" ++ unlines ls
  let c2 = run $ c {input=linesToAscii ls, output=[]}
  let (msg, result) = span (< 128) $ output c2
  putStrLn $ map (chr . fromIntegral) msg ++ concatMap show result

  putStrLn "\n\nPart 2"
  -- Jump as before . . .
  -- @ABCD  EFGHI
  -- _.__#
  -- __._#
  -- ___.#
  -- But check that you are not blocked after the jump.
  --        EFGHI
  --        .__#_
  --        #___#
  --        ##___
  -- ~(ABC)D -> ~((ABC)+D)
  -- (~E)H + EI + EF = (~E)H + E(I+F) = E(F+I) + H

  let ls2 = [
              -- jump if a hole is ahead and the landing is open
              "OR A J"
            , "AND B J"
            , "AND C J"
            , "NOT J J"
            , "AND D J"
              -- and if the future looks ok
            , "OR F T"
            , "OR I T"
            , "AND E T"
            , "OR H T"
            , "AND T J"
            , "RUN"
            ]
  let c3 = run $ computer0{memory = instructions}
  putStrLn $ map (chr . fromIntegral) $ output c3
  putStrLn $ "Running with input: \n" ++ unlines ls2
  let c4 = run $ c3 {input=linesToAscii ls2, output=[]}
  let (msg, result) = span (< 128) $ output c4
  putStrLn $ map (chr . fromIntegral) msg ++ concatMap show result


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, filter)
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)

import UI.NCurses

type Instructions = Map Integer Integer

data ComputerState = Ready | Running | Blocked | Done
  deriving (Show, Eq)

data Computer = Computer {
  state :: ComputerState,
  memory :: Map Integer Integer,
  iptr :: Integer, -- instruction pointer
  base :: Integer, -- base offset
  input :: [Integer],
  output :: [Integer]}

computer0 = Computer { state = Ready, memory = empty, iptr = 0, base = 0, input = [], output = [] }

run :: Computer -> Computer -- (Integer, Integer) -> Instructions -> [Integer] -> [Integer]
run c
  | state c == Ready = run $ c {state = Running}
  | state c == Done = c
  | state c == Blocked = if input c == [] then c else run $ c {state = Running}
  | otherwise = -- running
    --traceShow (iptr c, base c, memory c) $
    case instr `mod` 100 of
      1 -> -- add
        run $ c {iptr = i+4, memory = insert (addr 3) (arg 1 + arg 2) instructions}
      2 -> -- multiply
        run $ c {iptr = i+4, memory = insert (addr 3) (arg 1 * arg 2) instructions}
      3 -> -- read input
        if (null $ input c)
        then c {state = Blocked}
        else run $ c {iptr = i+2, memory = insert (addr 1) (head $ input c) instructions, input = tail $ input c}
      4 -> -- output
        --traceShow (arg 1) $
        run $ c {iptr = i+2, output = output c ++ [arg 1]}
      5 -> -- jump-if-true
        run $ c {iptr = if arg 1 == 0 then i+3 else arg 2}
      6 -> -- jump-if-false
        run $ c {iptr = if arg 1 == 0 then arg 2 else i+3}
      7 -> -- less than
        run $ c {iptr = i+4, memory = insert (addr 3) (if arg 1 < arg 2 then 1 else 0) instructions}
      8 -> -- equals
        run $ c {iptr = i+4, memory = insert (addr 3) (if arg 1 == arg 2 then 1 else 0) instructions}
      9 -> -- set relative base
        run $ c {iptr = i+2, base = base c + arg 1}
      99 -> -- halt
        run $ c {state = Done}
      _ -> error "unknown opcode"
  where instructions = memory c
        i = iptr c
        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 c)
          _ -> 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 c
          _ -> 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

data Game = Game {
  m :: Map (Integer, Integer) Integer, -- map
  b :: [(Integer, Integer)], -- ball
  p :: [(Integer, Integer)]} -- paddle
  deriving (Show)

game0 = Game { m = empty, b = [], p = [] }
updateGame :: Game -> [Integer] -> Game
updateGame g [] = g
updateGame g (x:y:t:r) = updateGame (g' { m = insert (x,y) t $ m g}) r
  where g' | t == 4 = g { b = (x,y): b g }
           | t == 3 = g { p = (x,y): p g }
           | otherwise = g
updateGame g _ = error "bad input"

recommendMove :: Game -> Integer
recommendMove g =
  if null (p g) || null (b g) then 0
  else if x == x' && (y == y' || y - 1 == y') then 0
       else case length (b g) of
              0 -> 0
              1 -> signum $ (fst $ head $ b g) - x
              _ -> let [x2, x1] = map fst $ take 2 $ b g
                in signum $ x2 + signum (x2 - x1) - x
  where (x,y) = head $ p g
        (x',y') = head $ b g

display :: [Integer] -> IO ()
display drawCommands = runCurses $ do
    setEcho False
    w <- defaultWindow
    updateWindow w $ do
        draw drawCommands
    render
    waitFor w (\ev -> ev == EventCharacter 'q' || ev == EventCharacter 'Q')
      where draw [] = do
              moveCursor 0 10
              drawString "Done"
            draw (x:y:t:r) =
              if (x,y) == (-1, 0)
              then do
                moveCursor 0 0
                drawString $ show t
                draw r
              else do
                moveCursor (y+5) x
                drawString (case t of
                              0 -> " "
                              1 -> "W"
                              2 -> "B"
                              3 -> "_"
                              4 -> "*")
                draw r

waitFor :: Window -> (Event -> Bool) -> Curses ()
waitFor w p = loop where
    loop = do
        ev <- getEvent w Nothing
        case ev of
            Nothing -> loop
            Just ev' -> if p ev' then return () else loop

play :: Game -> Computer -> IO (Game)
play game c = runCurses $ do
  setEcho False
  w <- defaultWindow
  loop w game $ run c
    where loop :: Window -> Game -> Computer -> Curses(Game)
          loop w g c
            | state c == Done = do
                processDrawCommands w g' $ output c
                updateWindow w $ do
                  moveCursor 0 10
                  drawString "Press q to quit"
                render
                let qloop = do
                      ev <- getEvent w Nothing
                      case ev of
                        Nothing -> qloop
                        Just (EventCharacter 'q') -> return g'
                        Just (EventCharacter 'Q') -> return g'
                        _ -> qloop
                qloop
            | state c == Blocked = do -- waiting for input
                processDrawCommands w g' $ output c
                moveEvent <- getEvent w (Just 20)
                let joystick = case moveEvent of
                                 Just (EventSpecialKey KeyLeftArrow) -> -1
                                 Just (EventSpecialKey KeyRightArrow) -> 1
                                 Nothing -> recommendMove g'
                                 Just (EventCharacter 'j') -> -1
                                 Just (EventCharacter 'J') -> -1
                                 Just (EventCharacter 'k') -> 1
                                 Just (EventCharacter 'K') -> 1
                                 _ -> 0
                updateWindow w $ do
                  moveCursor 0 10
                  drawString $ case joystick of
                    -1 -> "move left"
                    0 -> "no move"
                    1 -> "move right"
                render
                loop w g' $ run $ c {state = Running, output = [], input = [joystick]}
            | otherwise = loop w game $ run c
              where g' = updateGame g $ output c
          processDrawCommands :: Window -> Game -> [Integer] -> Curses()
          processDrawCommands w g [] = render
          processDrawCommands w g (x:y:t:r) = do
            updateWindow w $ do
              if (x,y) == (-1, 0)
              then do
              moveCursor 0 0
              drawString $ show t
              else do
              moveCursor 1 0
              drawString $ show $ take 25 $ b g
              moveCursor 2 0
              drawString $ show $ take 25 $ p g
              moveCursor (y+5) x
              drawString (case t of
                             0 -> " "
                             1 -> "W"
                             2 -> "B"
                             3 -> "_"
                             4 -> "*"
                             _ -> error ("Unknown input " ++ show t)
                         )
            processDrawCommands w g r

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

  -- putStrLn "Part 1"
  -- let c = run $ computer0 {memory = instructions}
  -- -- putStrLn $ show ("state", state c, "drawCommands", output c)
  -- let game = updateGame game0 $ output c
  -- putStrLn $ show $ size $ M.filter (== 2) $ m game
  -- putStrLn "Part 1b"
  -- display $ output c
  putStrLn "Part 2"
  let instructions2 = insert 0 2 instructions -- play for free by setting memory address 0 to 2
  game <- play game0 $ computer0 {memory = instructions2}
  putStrLn $ show $ head $ p game
  -- let moves = moves >>= (play . run (0,0) instructions2)
  -- -- moves <- play $ run (0,0) instructions2 $ repeat 0
  -- moves' <- moves
  -- putStrLn $ show $ length moves'

  -- let moves = moves >>= (play . run (0,0) instructions2)
  -- moves' <- moves
  -- putStrLn $ show moves'
  -- moves' <- play $ run (0,0) instructions2 []
  -- putStrLn $ show ("moves", moves')
  -- let moves = run (0,0) instructions2 $ repeat 0
  -- putStrLn $ show ("moves", moves)
  -- display moves

  -- 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


{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
module DerivingViaCompile where

import Data.Void
import Data.Complex
import Data.Functor.Const
import Data.Functor.Identity
import Data.Ratio
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Writer
import Control.Applicative hiding (WrappedMonad(..))

import Data.Bifunctor
import Data.Monoid
import Data.Kind

type f ~> g = forall xx. f xx -> g xx

-----
-- Simple example
-----

data Foo a = MkFoo a a
  deriving Show
       via (Identity (Foo a))

-----
-- Eta reduction at work
-----

newtype Flip p a b = Flip { runFlip :: p b a }

instance Bifunctor p => Bifunctor (Flip p) where
  bimap f g = Flip . bimap g f . runFlip

instance Bifunctor p => Functor (Flip p a) where
  fmap f = Flip . first f . runFlip

newtype Bar a = MkBar (Either a Int)
  deriving Functor
       via (Flip Either Int)

-----
-- Monad transformers
-----

type MTrans = (Type -> Type) -> (Type -> Type)

-- From `constraints'
data Dict c where
  Dict :: c => Dict c

newtype a :- b = Sub (a => Dict b)

infixl 1 \\
(\\) :: a => (b => r) -> (a :- b) -> r
r \\ Sub Dict = r

-- With `-XQuantifiedConstraints' this just becomes
--
--    type Lifting cls  trans = forall mm. cls mm => cls (trans mm)
--
--    type LiftingMonad trans = Lifting Monad trans
--
class LiftingMonad (trans :: MTrans) where
  proof :: Monad m :- Monad (trans m)

instance LiftingMonad (StateT s :: MTrans) where
  proof :: Monad m :- Monad (StateT s m)
  proof = Sub Dict

instance Monoid w => LiftingMonad (WriterT w :: MTrans) where
  proof :: Monad m :- Monad (WriterT w m)
  proof = Sub Dict

instance (LiftingMonad trans, LiftingMonad trans') => LiftingMonad (ComposeT trans trans' :: MTrans) where
  proof :: forall m. Monad m :- Monad (ComposeT trans trans' m)
  proof = Sub (Dict \\ proof @trans @(trans' m) \\ proof @trans' @m)

newtype Stack :: MTrans where
  Stack :: ReaderT Int (StateT Bool (WriterT String m)) a -> Stack m a
  deriving newtype
    ( Functor
    , Applicative
    , Monad
    , MonadReader Int
    , MonadState Bool
    , MonadWriter String
    )
  deriving (MonadTrans, MFunctor)
       via (ReaderT Int `ComposeT` StateT Bool `ComposeT` WriterT String)

class MFunctor (trans :: MTrans) where
  hoist :: Monad m => (m ~> m') -> (trans m ~> trans m')

instance MFunctor (ReaderT r :: MTrans) where
  hoist :: Monad m => (m ~> m') -> (ReaderT r m ~> ReaderT r m')
  hoist nat = ReaderT . fmap nat . runReaderT

instance MFunctor (StateT s :: MTrans) where
  hoist :: Monad m => (m ~> m') -> (StateT s m ~> StateT s m')
  hoist nat = StateT . fmap nat . runStateT

instance MFunctor (WriterT w :: MTrans) where
  hoist :: Monad m => (m ~> m') -> (WriterT w m ~> WriterT w m')
  hoist nat = WriterT . nat . runWriterT

infixr 9 `ComposeT`
newtype ComposeT :: MTrans -> MTrans -> MTrans where
  ComposeT :: { getComposeT :: f (g m) a } -> ComposeT f g m a
  deriving newtype (Functor, Applicative, Monad)

instance (MonadTrans f, MonadTrans g, LiftingMonad g) => MonadTrans (ComposeT f g) where
  lift :: forall m. Monad m => m ~> ComposeT f g m
  lift = ComposeT . lift . lift
    \\ proof @g @m

instance (MFunctor f, MFunctor g, LiftingMonad g) => MFunctor (ComposeT f g) where
  hoist :: forall m m'. Monad m => (m ~> m') -> (ComposeT f g m ~> ComposeT f g m')
  hoist f = ComposeT . hoist (hoist f) . getComposeT
    \\ proof @g @m

-----
-- Using tuples in a `via` type
-----

newtype X a = X (a, a)
  deriving (Semigroup, Monoid)
       via (Product a, Sum a)

  deriving (Show, Eq)
       via (a, a)

-----
-- Abstract data types
-----

class C f where
  c :: f a -> Int

newtype X2 f a = X2 (f a)

instance C (X2 f) where
  c = const 0

deriving via (X2 IO) instance C IO

----
-- Testing parser
----

newtype P0 a = P0 a             deriving Show via a
newtype P1 a = P1 [a]           deriving Show via [a]
newtype P2 a = P2 (a, a)        deriving Show via (a, a)
newtype P3 a = P3 (Maybe a)     deriving Show via (First a)
newtype P4 a = P4 (Maybe a)     deriving Show via (First $ a)
newtype P5 a = P5 a             deriving Show via (Identity $ a)
newtype P6 a = P6 [a]           deriving Show via ([] $ a)
newtype P7 a = P7 (a, a)        deriving Show via (Identity $ (a, a))
newtype P8 a = P8 (Either () a) deriving Functor via (($) (Either ()))

newtype f $ a = APP (f a) deriving newtype Show deriving newtype Functor

----
-- From Baldur's notes
----

----
-- 1
----
newtype WrapApplicative f a = WrappedApplicative (f a)
  deriving (Functor, Applicative)

instance (Applicative f, Num a) => Num (WrapApplicative f a) where
  (+)         = liftA2 (+)
  (*)         = liftA2 (*)
  negate      = fmap negate
  fromInteger = pure . fromInteger
  abs         = fmap abs
  signum      = fmap signum

instance (Applicative f, Fractional a) => Fractional (WrapApplicative f a) where
  recip        = fmap recip
  fromRational = pure . fromRational

instance (Applicative f, Floating a) => Floating (WrapApplicative f a) where
  pi    = pure pi
  sqrt  = fmap sqrt
  exp   = fmap exp
  log   = fmap log
  sin   = fmap sin
  cos   = fmap cos
  asin  = fmap asin
  atan  = fmap atan
  acos  = fmap acos
  sinh  = fmap sinh
  cosh  = fmap cosh
  asinh = fmap asinh
  atanh = fmap atanh
  acosh = fmap acosh

instance (Applicative f, Semigroup s) => Semigroup (WrapApplicative f s) where
  (<>) = liftA2 (<>)

instance (Applicative f, Monoid m) => Monoid (WrapApplicative f m) where
  mempty = pure mempty

----
-- 2
----
class Pointed p where
  pointed :: a -> p a

newtype WrapMonad f a = WrappedMonad (f a)
  deriving newtype (Pointed, Monad)

instance (Monad m, Pointed m) => Functor (WrapMonad m) where
  fmap = liftM

instance (Monad m, Pointed m) => Applicative (WrapMonad m) where
  pure  = pointed
  (<*>) = ap

-- data
data Sorted a = Sorted a a a
  deriving (Functor, Applicative)
    via (WrapMonad Sorted)
  deriving (Num, Fractional, Floating, Semigroup, Monoid)
    via (WrapApplicative Sorted a)

instance Monad Sorted where
  (>>=) :: Sorted a -> (a -> Sorted b) -> Sorted b
  Sorted a b c >>= f = Sorted a' b' c' where
    Sorted a' _  _  = f a
    Sorted _  b' _  = f b
    Sorted _  _  c' = f c

instance Pointed Sorted where
  pointed :: a -> Sorted a
  pointed a = Sorted a a a

----
-- 3
----
class IsZero a where
  isZero :: a -> Bool

newtype WrappedNumEq  a = WrappedNumEq a
newtype WrappedShow   a = WrappedShow  a
newtype WrappedNumEq2 a = WrappedNumEq2 a

instance (Num a, Eq a) => IsZero (WrappedNumEq a) where
  isZero :: WrappedNumEq a -> Bool
  isZero (WrappedNumEq a) = 0 == a

instance Show a => IsZero (WrappedShow a) where
  isZero :: WrappedShow a -> Bool
  isZero (WrappedShow a) = "0" == show a

instance (Num a, Eq a) => IsZero (WrappedNumEq2 a) where
  isZero :: WrappedNumEq2 a -> Bool
  isZero (WrappedNumEq2 a) = a + a == a

newtype INT = INT Int
  deriving newtype Show
  deriving IsZero via (WrappedNumEq Int)

newtype VOID = VOID Void deriving IsZero via (WrappedShow Void)

----
-- 4
----
class Bifunctor p => Biapplicative p where
  bipure :: a -> b -> p a b

  biliftA2
    :: (a  -> b  -> c)
    -> (a' -> b' -> c')
    -> p a a'
    -> p b b'
    -> p c c'

instance Biapplicative (,) where
  bipure = (,)

  biliftA2 f f' (a, a') (b, b') =
    (f a b, f' a' b')

newtype WrapBiapp p a b = WrapBiap (p a b)
  deriving newtype (Bifunctor, Biapplicative, Eq)

instance (Biapplicative p, Num a, Num b) => Num (WrapBiapp p a b) where
  (+) = biliftA2 (+) (+)
  (-) = biliftA2 (*) (*)
  (*) = biliftA2 (*) (*)
  negate = bimap negate negate
  abs = bimap abs abs
  signum = bimap signum signum
  fromInteger n = fromInteger n `bipure` fromInteger n

newtype INT2 = INT2 (Int, Int)
  deriving IsZero via (WrappedNumEq2 (WrapBiapp (,) Int Int))

----
-- 5
----
class Monoid a => MonoidNull a where
  null :: a -> Bool

newtype WrpMonNull a = WRM a deriving (Eq, Semigroup, Monoid)

instance (Eq a, Monoid a) => MonoidNull (WrpMonNull a) where
  null :: WrpMonNull a -> Bool
  null = (== mempty)

deriving via (WrpMonNull Any) instance MonoidNull Any
deriving via ()               instance MonoidNull ()
deriving via Ordering         instance MonoidNull Ordering

----
-- 6
----
-- https://github.com/mikeizbicki/subhask/blob/f53fd8f465747681c88276c7dabe3646fbdf7d50/src/SubHask/Algebra.hs#L635

class Lattice a where
  sup   :: a -> a -> a
  (.>=) :: a -> a -> Bool
  (.>)  :: a -> a -> Bool

newtype WrapOrd a = WrappedOrd a
  deriving newtype (Eq, Ord)

instance Ord a => Lattice (WrapOrd a) where
  sup   = max
  (.>=) = (>=)
  (.>)  = (>)

deriving via [a]    instance Ord a          => Lattice [a]
deriving via (a, b) instance (Ord a, Ord b) => Lattice (a, b)
--mkLattice_(Bool)
deriving via Bool instance Lattice Bool
--mkLattice_(Char)
deriving via Char instance Lattice Char
--mkLattice_(Int)
deriving via Int instance Lattice Int
--mkLattice_(Integer)
deriving via Integer instance Lattice Integer
--mkLattice_(Float)
deriving via Float instance Lattice Float
--mkLattice_(Double)
deriving via Double instance Lattice Double
--mkLattice_(Rational)
deriving via Rational instance Lattice Rational

----
-- 7
----
-- https://hackage.haskell.org/package/linear-1.20.7/docs/src/Linear-Affine.html

class Functor f => Additive f where
  zero :: Num a => f a
  (^+^) :: Num a => f a -> f a -> f a
  (^+^) = liftU2 (+)
  (^-^) :: Num a => f a -> f a -> f a
  x ^-^ y = x ^+^ fmap negate y
  liftU2 :: (a -> a -> a) -> f a -> f a -> f a

instance Additive [] where
  zero = []
  liftU2 f = go where
    go (x:xs) (y:ys) = f x y : go xs ys
    go [] ys = ys
    go xs [] = xs

instance Additive Maybe where
  zero = Nothing
  liftU2 f (Just a) (Just b) = Just (f a b)
  liftU2 _ Nothing ys = ys
  liftU2 _ xs Nothing = xs

instance Applicative f => Additive (WrapApplicative f) where
  zero   = pure 0
  liftU2 = liftA2

deriving via (WrapApplicative ((->) a)) instance Additive ((->) a)
deriving via (WrapApplicative Complex)  instance Additive Complex
deriving via (WrapApplicative Identity) instance Additive Identity

instance Additive ZipList where
  zero = ZipList []
  liftU2 f (ZipList xs) (ZipList ys) = ZipList (liftU2 f xs ys)

class Additive (Diff p) => Affine p where
  type Diff p :: Type -> Type

  (.-.) :: Num a => p a -> p a -> Diff p a
  (.+^) :: Num a => p a -> Diff p a -> p a
  (.-^) :: Num a => p a -> Diff p a -> p a
  p .-^ v = p .+^ fmap negate v

-- #define ADDITIVEC(CTX,T) instance CTX => Affine T where type Diff T = T ; \
--   (.-.) = (^-^) ; {-# INLINE (.-.) #-} ; (.+^) = (^+^) ; {-# INLINE (.+^) #-} ; \
--   (.-^) = (^-^) ; {-# INLINE (.-^) #-}
-- #define ADDITIVE(T) ADDITIVEC((), T)
newtype WrapAdditive f a = WrappedAdditive (f a)

instance Additive f => Affine (WrapAdditive f) where
  type Diff (WrapAdditive f) = f

  WrappedAdditive a .-. WrappedAdditive b = a ^-^ b
  WrappedAdditive a .+^ b = WrappedAdditive (a ^+^ b)
  WrappedAdditive a .-^ b = WrappedAdditive (a ^-^ b)

-- ADDITIVE(((->) a))
deriving via (WrapAdditive ((->) a)) instance Affine ((->) a)
-- ADDITIVE([])
deriving via (WrapAdditive [])       instance Affine []
-- ADDITIVE(Complex)
deriving via (WrapAdditive Complex)  instance Affine Complex
-- ADDITIVE(Maybe)
deriving via (WrapAdditive Maybe)    instance Affine Maybe
-- ADDITIVE(ZipList)
deriving via (WrapAdditive ZipList)  instance Affine ZipList
-- ADDITIVE(Identity)
deriving via (WrapAdditive Identity) instance Affine Identity

----
-- 8
----

class C2 a b c where
  c2 :: a -> b -> c

instance C2 a b (Const a b) where
  c2 x _ = Const x

newtype Fweemp a = Fweemp a
  deriving (C2 a b)
       via (Const a (b :: Type))


{-# LANGUAGE Safe #-}
{-# OPTIONS_GHC -fno-warn-deprecations #-}
-- | Import all modules from prelude that should be safe
module Main where

import Numeric
import Prelude

import Control.Applicative
import Control.Arrow
import Control.Category

import Control.Concurrent
import Control.Concurrent.Chan
import Control.Concurrent.MVar
import Control.Concurrent.QSem
import Control.Concurrent.QSemN

import Control.Exception
import Control.Exception.Base

import Control.Monad
import Control.Monad.Fix
import Control.Monad.Instances
import Control.Monad.Zip

import Control.Monad.ST
import Control.Monad.ST.Lazy
import Control.Monad.ST.Strict

import Data.Bifunctor
import Data.Bits
import Data.Bool
import Data.Char
-- import Data.Coerce
import Data.Complex
import Data.Data
import Data.Dynamic
import Data.Either
import Data.Eq
import Data.Fixed
import Data.Foldable
import Data.Function
import Data.Functor.Identity
import Data.Functor
import Data.IORef
import Data.Int
import Data.Ix
import Data.List
import Data.Maybe
import Data.Monoid
import Data.Ord
import Data.Proxy
import Data.Ratio
import Data.String
import Data.Traversable
import Data.Tuple
import Data.Typeable
import Data.Unique
import Data.Version
import Data.Word

import Data.STRef
import Data.STRef.Lazy
import Data.STRef.Strict

import Data.Type.Bool
-- import Data.Type.Coercion
-- import Data.Type.Equality

-- import Debug.Trace

import Foreign

import Foreign.Concurrent
import Foreign.ForeignPtr
import Foreign.Ptr
import Foreign.StablePtr
import Foreign.Storable

import Foreign.C
import Foreign.C.Error
import Foreign.C.String
import Foreign.C.Types

import Foreign.Marshal
import Foreign.Marshal.Alloc
import Foreign.Marshal.Array
import Foreign.Marshal.Error
import Foreign.Marshal.Pool
import Foreign.Marshal.Utils

import System.CPUTime
import System.Console.GetOpt
import System.Environment
import System.Exit
import System.IO
import System.IO.Error
import System.Info
import System.Mem
import System.Mem.StableName
import System.Mem.Weak
import System.Posix.Internals
import System.Posix.Types
import System.Timeout

import Text.ParserCombinators.ReadP
import Text.ParserCombinators.ReadPrec
import Text.Printf
import Text.Read
import Text.Read.Lex
import Text.Show
import Text.Show.Functions

import Type.Reflection

-- import Unsafe.Coerce

f :: Int
f = 2

main :: IO ()
main = putStrLn $ "X is: " ++ show f


-- Copyright 2021 Google LLC.
-- SPDX-License-Identifier: Apache-2.0

-- #region language extensions
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ViewPatterns #-}

-- #endregion

module Main where

-- #region imports

import Control.DeepSeq
import Control.Lens (at, each, folded, isn't, ix, traversed, (%~), (&), (.~), (?~), (^.), (^..), (^?), _1, _2, _3, _4)
import Control.Lens.Extras (is)
import Control.Monad (forM_)
import qualified Data.Array as A
import qualified Data.Array.Unboxed as UA
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BS8
import qualified Data.ByteString.Lazy as LBS
import qualified Data.ByteString.Lazy.Char8 as LBS8
import Data.Complex (Complex (..), cis, imagPart, realPart)
import Data.Functor ((<&>))
import Data.Generics.Labels ()
import qualified Data.Graph.Inductive as G
import Data.Graph.Inductive.PatriciaTree (Gr)
import Data.List (elemIndex, scanl', sortOn, transpose, foldl', tails)
import Data.List.Split (chunksOf, splitOn)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import Data.Maybe (catMaybes, fromJust, mapMaybe)
import Data.MemoTrie
import Data.Sequence (Seq (..))
import qualified Data.Sequence as Seq
import Data.Set (Set)
import qualified Data.Set as S
import qualified Data.Text as T
import qualified Data.Text.Lazy as LT
import Data.Vector (Vector)
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as UV
import qualified Data.Vector.Unboxed.Mutable as MUV
import Debug.Trace (trace, traceShow, traceShowId)
import GHC.Generics (Generic)
import GHC.IO (unsafePerformIO)
import Linear (V2 (..), angle, (*^), _x, _y)
import QuickSpec
import Safe hiding (at)
import Test.QuickCheck (Arbitrary (..), NonNegative (NonNegative), chooseInt, vectorOf, (==>), quickCheckResult)
import Test.QuickCheck.Modifiers (Positive (Positive, getPositive))
import qualified Text.Megaparsec as P
import qualified Text.Megaparsec.Char as P
import qualified Text.Megaparsec.Char.Lexer as L
import Text.Pretty.Simple (pPrint, pShow)
import Text.Regex.PCRE ((=~))
import Util

-- #endregion

type Input = UV.Vector Int

parse :: String -> Input
parse = UV.fromList . map (read . (: [])) . head . lines

input :: Input
input = parse (unsafePerformIO (readFile "/home/polux/projects/aoc/2019/day16-2/input.txt"))

vecSize :: Int
vecSize = UV.length input * 10_000

offset :: Int
offset = UV.take 7 input & UV.toList & concatMap show & read

inputFromOffset :: [Int]
inputFromOffset = [input UV.! (i `mod` UV.length input) | i <- [offset..vecSize-1]]

multFromOffset :: [Int] -> [Int]
multFromOffset vec = scanr (+) 0 vec & map (`mod` 10)

result :: [Int]
result = take 8 (iterate multFromOffset inputFromOffset !! 100)

main :: IO ()
main = print result


{-# LANGUAGE ForeignFunctionInterface #-}
module Grenade.Layers.Internal.Pooling (
    poolForward
  , poolBackward
  ) where

import qualified Data.Vector.Storable as U ( unsafeToForeignPtr0, unsafeFromForeignPtr0 )

import           Foreign ( mallocForeignPtrArray, withForeignPtr )
import           Foreign.Ptr ( Ptr )

import           Numeric.LinearAlgebra ( Matrix , flatten )
import qualified Numeric.LinearAlgebra.Devel as U

import           System.IO.Unsafe ( unsafePerformIO )

poolForward :: Int -> Int -> Int -> Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double
poolForward channels height width kernelRows kernelColumns strideRows strideColumns dataIm =
  let vec             = flatten dataIm
      rowOut          = (height - kernelRows) `div` strideRows + 1
      colOut          = (width - kernelColumns) `div` strideColumns + 1
      numberOfPatches = rowOut * colOut
  in unsafePerformIO $ do
    outPtr <- mallocForeignPtrArray (numberOfPatches * channels)
    let (inPtr, _) = U.unsafeToForeignPtr0 vec

    withForeignPtr inPtr $ \inPtr' ->
      withForeignPtr outPtr $ \outPtr' ->
        pool_forwards_cpu inPtr' channels height width kernelRows kernelColumns strideRows strideColumns outPtr'

    let matVec = U.unsafeFromForeignPtr0 outPtr (numberOfPatches * channels)
    return $ U.matrixFromVector U.RowMajor (rowOut * channels) colOut matVec

foreign import ccall unsafe
    pool_forwards_cpu
      :: Ptr Double -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Ptr Double -> IO ()

poolBackward :: Int -> Int -> Int -> Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double -> Matrix Double
poolBackward channels height width kernelRows kernelColumns strideRows strideColumns dataIm dataGrad =
  let vecIm     = flatten dataIm
      vecGrad   = flatten dataGrad
  in unsafePerformIO $ do
    outPtr <- mallocForeignPtrArray (height * width * channels)
    let (imPtr, _) = U.unsafeToForeignPtr0 vecIm
    let (gradPtr, _) = U.unsafeToForeignPtr0 vecGrad

    withForeignPtr imPtr $ \imPtr' ->
      withForeignPtr gradPtr $ \gradPtr' ->
        withForeignPtr outPtr $ \outPtr' ->
          pool_backwards_cpu imPtr' gradPtr' channels height width kernelRows kernelColumns strideRows strideColumns outPtr'

    let matVec = U.unsafeFromForeignPtr0 outPtr (height * width * channels)
    return $ U.matrixFromVector U.RowMajor (height * channels) width matVec

foreign import ccall unsafe
    pool_backwards_cpu
      :: Ptr Double -> Ptr Double -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Ptr Double -> IO ()


module Type
    (
      -- * Types
      LispVal(..)
    , LispError(..)
    , LispEnv
      -- * Monads
    , ThrowsError
    ) where

import Data.Array (Array, elems)
import Data.Complex (Complex(..), imagPart, realPart)
import Data.Ratio (numerator, denominator)

import Data.IORef (IORef)
import Text.ParserCombinators.Parsec (ParseError)

-- | Values and code structures for lisp.
data LispVal
    = Symbol String
    | Integer Integer
    | Rational Rational
    | Real Float
    | Complex (Complex Integer)
    | Character Char
    | String String
    | Boolean Bool
    | List [LispVal]
    | DottedList [LispVal] LispVal
    | Vector (Array Int LispVal)
    | PrimitiveFunc String
    | Func [String] (Maybe String) [LispVal] LispEnv
    deriving (Eq)

instance Show LispVal where
    show = showVal

-- | Show the s-expression for LispVal.
showVal :: LispVal -> String
showVal (Symbol x)        = x
showVal (Integer x)       = show x
showVal (Rational x)      = show (numerator x) ++ "/" ++ show (denominator x)
showVal (Real x)          = show x
showVal (Complex x)       = show (realPart x) ++ "+" ++ show (imagPart x) ++ "i"
showVal (Character x)     = "#\\" ++ [x]
showVal (String x)        = "\"" ++ x ++ "\""
showVal (Boolean True)    = "#t"
showVal (Boolean False)   = "#f"
showVal (List xs)         = "(" ++ unwordsList xs ++ ")"
showVal (DottedList xs t) = "(" ++ unwordsList xs ++ " . " ++ showVal t ++ ")"
showVal (Vector xs)       = "#(" ++ unwordsList (elems xs) ++ ")"
showVal (PrimitiveFunc n) = "<primitive " ++ n ++ ">"
showVal (Func pns vn _ _) =
    "(lambda (" ++ varargText ++ unwords (map show pns) ++ ") ...)"
  where
    varargText = maybe "" (" . " ++) vn

-- | Errors for evaluating or parsing lisp.
data LispError
    = NumArgs String Integer [LispVal]
    | TypeMismatch String String LispVal
    | Parser ParseError
    | BadSpecialForm String LispVal
    | NotFunction String String
    | UnboundVar String String
    | Default String
    deriving (Eq)

-- | Show `LispError` as text.
showError :: LispError -> String
showError (NumArgs fName expected found) =
    fName ++ ": Expected " ++ show expected ++ " arguments, found values " ++ vs
  where
    vs = unwords $ map show found
showError (TypeMismatch fName expected found) =
    fName ++ ": Invalid type: expected " ++ expected ++ ", found " ++ show found
showError (Parser parseError) = "Parse error at " ++ show parseError
showError (BadSpecialForm message form) = message ++ ": " ++ show form
showError (NotFunction message fName)   = message ++ ": " ++ fName
showError (UnboundVar message varName)  = message ++ ": " ++ varName
showError (Default message)             = "Default: " ++ message

instance Show LispError where
    show = showError

-- | Monad for computation with potential `LispError`.
type ThrowsError = Either LispError

-- | Environment of mutable `LispVal` bindings.
type LispEnv = IORef [(String, IORef LispVal)]

-- | Join the s-expressions of a list of `LispVal` with spaces.
unwordsList :: [LispVal] -> String
unwordsList = unwords . map showVal


-- |
-- Module      : Copernicus
-- Description : Contains the core of the Copernicus project
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental
-- Portability : POSIX (not sure)
--
-- Jonatan H Sundqvist
-- November 25 2014
--

-- TODO | - Model planetary motion and physics in 2D
--        - Cartoon earth with gravitational field
--        - Flexible event handling
--        - App typeclass (run, manage events, window properties, etc.)
--        - Lenses
--        - Options (eg. toggle grid) (cf. 'when')
--        - UI
--        - Use types to encode units (eg. SI, radians)
--        - Move to Cairo (branch?)
--        - Polymorphic types (not just Floats) (?)

-- SPEC | -
--        -

module Copernicus where

---------------------------------------------------------------------------------------------------
-- We'll need these
---------------------------------------------------------------------------------------------------
import Data.Complex
-- import Control.Monad (when)

import Graphics.Gloss.Geometry.Angle (degToRad, radToDeg, normaliseAngle)

---------------------------------------------------------------------------------------------------
-- Types
---------------------------------------------------------------------------------------------------
-- Make typeclass (mesh, body, bounding box, collision, etc) (?)
-- type Number a = Floating a -- Real number
type Vector = Complex -- TODO: Polymorphic (cf. related TODO item)
data Body f = Body (Vector f) (Vector f) (Vector f) deriving Show -- Add argument

---------------------------------------------------------------------------------------------------
-- Functions
---------------------------------------------------------------------------------------------------
clampAngle :: Float -> Float
clampAngle = radToDeg . normaliseAngle . degToRad

--
-- TODO: Make them change colour when bouncing (?)
animate :: (RealFloat f, Floating f) => f -> Body f -> Body f
animate t (Body p v a) = collide ground $ Body (parabola t p v a) (v + (t:+0)*a) a
	where ground = 0.0 -- 30+30/2-540/2 (previous hard-coded value)

--
parabola :: (RealFloat f, Floating f) => f -> Vector f -> Vector f -> Vector f -> Vector f
parabola t p v a = let (px:+py) = p
                       (vx:+vy) = v
                       (ax:+ay) = a
                   in (px + vx*t + 0.5*ax*t**2) :+ (py + vy*t + 0.5*ay*t**2)

-- collide
-- Very primitive for now
-- TODO: Use 'contains' function (Range -> Value -> Bool)
-- TODO: Don't hard-code bounds (left, right)
-- TODO: Take bounds of Body into account (don't hard-code that either)
collide :: (RealFloat f, Floating f) => f -> Body f -> Body f
collide gnd (Body (px:+py) (vx:+vy) a) = Body (px:+py) ((invertIf (\ _ -> (px <= left) || ( px >= right)) vx) :+ (invertIf (\ v -> (v < 0) && (py <= gnd)) vy)) a
	where invertIf p v | p v 	   = -v
	                   | otherwise =  v
	      (left, right) = (-5, 5) --(15-720/2, 720/2-30/2)

-- | ETA (estimated time of arrival)
-- TODO: Rename (eg. timeUntil, solveForT, etc)
-- TODO: Parabola type (eg. Parabola a v x)
-- eta :: Acceleration

-- Utilities (should eventually be moved to separate module or library ----------------------------
-- Python-style String formatting (eg. keyword interpolation, {0}, customisation, format specs.)
-- Parsec

---------------------------------------------------------------------------------------------------
-- Entry point
---------------------------------------------------------------------------------------------------


{-# LANGUAGE DeriveDataTypeable, FlexibleContexts, Rank2Types #-}
{-# OPTIONS_HADDOCK hide #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Matrix.STBase
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : asinerimental
--

module Numeric.LinearAlgebra.Matrix.STBase
    where

import Control.Monad( forM_, when )
import Control.Monad.ST( ST, RealWorld, runST, unsafeInterleaveST,
    unsafeIOToST )
import Data.Maybe( fromMaybe )
import Data.Typeable( Typeable )
import Foreign( Ptr, advancePtr, peek, peekElemOff, pokeElemOff,
    mallocForeignPtrArray )
import Text.Printf( printf )
import Unsafe.Coerce( unsafeCoerce )

import Numeric.LinearAlgebra.Types
import qualified Foreign.BLAS as BLAS
import Numeric.LinearAlgebra.Matrix.Base hiding ( unsafeWith,
    unsafeToForeignPtr, unsafeFromForeignPtr, )
import qualified Numeric.LinearAlgebra.Matrix.Base as M
import Numeric.LinearAlgebra.Vector( STVector, RVector )
import qualified Numeric.LinearAlgebra.Vector as V

-- | Mutable dense matrices in the 'ST' monad.
newtype STMatrix s e = STMatrix { unSTMatrix :: Matrix e }
    deriving (Typeable)

-- | Mutable dense matrices in the 'IO' monad.
type IOMatrix = STMatrix RealWorld

-- | A safe way to create and work with a mutable matrix before returning
-- an immutable matrix for later perusal. This function avoids copying
-- the matrix before returning it - it uses 'unsafeFreeze' internally,
-- but this wrapper is a safe interface to that function.
create :: (Storable e) => (forall s . ST s (STMatrix s e)) -> Matrix e
create mx = runST $ mx >>= unsafeFreeze
{-# INLINE create #-}

-- | Converts a mutable matrix to an immutable one by taking a complete
-- copy of it.
freeze :: (RMatrix m, Storable e) => m e -> ST s (Matrix e)
freeze a = do
    a' <- newCopy a
    unsafeFreeze a'
{-# INLINE freeze #-}

-- | Read-only matrices
class RMatrix m where
    -- | Get the dimensions of the matrix (number of rows and columns).
    getDim :: (Storable e) => m e -> ST s (Int,Int)

    -- | Same as 'withCol' but does not range-check index.
    unsafeWithCol :: (Storable e)
                  => m e
                  -> Int
                  -> (forall v. RVector v => v e -> ST s a)
                  -> ST s a

    -- | Perform an action with a list of views of the matrix columns.
    withCols :: (Storable e)
                 => m e
                 -> (forall v . RVector v => [v e] -> ST s a)
                 -> ST s a

    -- | Same as 'withSlice' but does not range-check index.
    unsafeWithSlice :: (Storable e)
                    => (Int,Int)
                    -> (Int,Int)
                    -> m e
                    -> (forall m'. RMatrix m' => m' e -> ST s a)
                    -> ST s a

    -- | Possibly view a matrix as a vector and perform an action on the
    -- view.  This only succeeds if the matrix is stored contiguously in
    -- memory, i.e. if the matrix contains a single column or the \"lda\"
    -- of the matrix is equal to the number of rows.
    maybeWithVector :: (Storable e)
                    => m e
                    -> (forall v . RVector v => v e -> ST s a)
                    -> Maybe (ST s a)

    -- | Converts a read-only matrix into an immutable matrix. This simply
    -- casts the matrix from one type to the other without copying.
    -- Note that because the matrix is possibly not copied, any subsequent
    -- modifications made to the read-only version of the matrix may be shared
    -- with the immutable version. It is safe to use, therefore, if the
    -- read-only version is never modified after the freeze operation.
    unsafeFreeze :: (Storable e) => m e -> ST s (Matrix e)

    -- | Unsafe cast from a read-only matrix to a mutable matrix.
    unsafeThaw :: (Storable e)
               => m e -> ST s (STMatrix s e)

    -- | Execute an 'IO' action with a pointer to the first element in the
    -- matrix and the leading dimension (lda).
    unsafeWith :: (Storable e) => m e -> (Ptr e -> Int -> IO a) -> IO a

instance RMatrix Matrix where
    getDim = return . dim
    {-# INLINE getDim #-}
    unsafeWithCol a j f = f (unsafeCol a j)
    {-# INLINE unsafeWithCol #-}
    withCols a f = f (cols a)
    {-# INLINE withCols #-}
    unsafeWithSlice ij mn a f = f (unsafeSlice ij mn a)
    {-# INLINE unsafeWithSlice #-}
    maybeWithVector a f | isContig a = Just $ f (toVector a)
                        | otherwise = Nothing
    {-# INLINE maybeWithVector #-}
    unsafeWith = M.unsafeWith
    {-# INLINE unsafeWith #-}
    unsafeFreeze = return
    {-# INLINE unsafeFreeze #-}
    unsafeThaw = return . STMatrix
    {-# INLINE unsafeThaw #-}

instance RMatrix (STMatrix s) where
    getDim = return . dim . unSTMatrix
    {-# INLINE getDim #-}
    unsafeWithCol = unsafeWithCol . unSTMatrix
    {-# INLINE unsafeWithCol #-}
    withCols = withCols . unSTMatrix
    {-# INLINE withCols #-}
    unsafeWithSlice ij mn = unsafeWithSlice ij mn . unSTMatrix
    {-# INLINE unsafeWithSlice #-}
    maybeWithVector = maybeWithVector . unSTMatrix
    {-# INLINE maybeWithVector #-}
    unsafeWith = unsafeWith . unSTMatrix
    {-# INLINE unsafeWith #-}
    unsafeFreeze = return . unSTMatrix
    {-# INLINE unsafeFreeze #-}
    unsafeThaw v = return $ cast v
      where
        cast :: STMatrix s e -> STMatrix s' e
        cast = unsafeCoerce
    {-# INLINE unsafeThaw #-}

-- | Perform an action with a view of a mutable matrix column
-- (no index checking).
unsafeWithColM :: (Storable e)
               => STMatrix s e
               -> Int
               -> (STVector s e -> ST s a)
               -> ST s a
unsafeWithColM a j f =
    unsafeWithCol a j $ \c -> do
        mc <- V.unsafeThaw c
        f mc
{-# INLINE unsafeWithColM #-}

-- | Perform an action with a list of views of the mutable matrix columns. See
-- also 'withCols'.
withColsM :: (Storable e)
               => STMatrix s e
               -> ([STVector s e] -> ST s a)
               -> ST s a
withColsM a f =
    withCols a $ \cs -> do
        mcs <- thawVecs cs
        f mcs
  where
    thawVecs [] = return []
    thawVecs (c:cs) = unsafeInterleaveST $ do
        mc <- V.unsafeThaw c
        mcs <- thawVecs cs
        return $ mc:mcs
{-# INLINE withColsM #-}

-- | Possibly view a matrix as a vector and perform an action on the
-- view.  This succeeds when the matrix is stored contiguously in memory,
-- i.e. if the matrix contains a single column or the \"lda\" of the matrix
-- is equal to the number of rows.  See also 'maybeWithVector'.
maybeWithVectorM :: (Storable e)
                 => STMatrix s e
                 -> (STVector s e -> ST s a)
                 -> Maybe (ST s a)
maybeWithVectorM a f =
    maybeWithVector a $ \v -> do
        mv <- V.unsafeThaw v
        f mv
{-# INLINE maybeWithVectorM #-}

-- | View a vector as a matrix of the given shape and pass it to
-- the specified function.
withFromVector :: (RVector v, Storable e)
                   => (Int,Int)
                   -> v e
                   -> (forall m . RMatrix m => m e -> ST s a)
                   -> ST s a
withFromVector mn@(m,n) v f = do
    nv <- V.getDim v
    when (nv /= m*n) $ error $
        printf ("withFromVector (%d,%d) <vector with dim %d>:"
                ++ " dimension mismatch") m n nv
    iv <- V.unsafeFreeze v
    f $ fromVector mn iv
{-# INLINE withFromVector #-}

-- | View a mutable vector as a mutable matrix of the given shape and pass it
-- to the specified function.
withFromVectorM :: (Storable e)
                     => (Int,Int)
                     -> STVector s e
                     -> (STMatrix s e -> ST s a)
                     -> ST s a
withFromVectorM mn@(m,n) v f = do
    nv <- V.getDim v
    when (nv /= m*n) $ error $
        printf ("withFromVectorM (%d,%d) <vector with dim %d>:"
                ++ " dimension mismatch") m n nv
    withFromVector mn v $ \a -> do
        ma <- unsafeThaw a
        f ma
{-# INLINE withFromVectorM #-}

-- | View a vector as a matrix with one column and pass it to
-- the specified function.
withFromCol :: (RVector v, Storable e)
                => v e
                -> (forall m . RMatrix m => m e -> ST s a)
                -> ST s a
withFromCol v f = do
    m <- V.getDim v
    withFromVector (m,1) v f
{-# INLINE withFromCol #-}

-- | View a mutable vector as a mutable matrix with one column and pass it to
-- the specified function.
withFromColM :: (Storable e)
                  => STVector s e
                  -> (STMatrix s e -> ST s a)
                  -> ST s a
withFromColM v f = do
    m <- V.getDim v
    withFromVectorM (m, 1) v f
{-# INLINE withFromColM #-}

-- | View a vector as a matrix with one row and pass it to
-- the specified function.
withFromRow :: (RVector v, Storable e)
                => v e
                -> (forall m . RMatrix m => m e -> ST s a)
                -> ST s a
withFromRow v f = do
    n <- V.getDim v
    withFromVector (1,n) v f
{-# INLINE withFromRow #-}

-- | View a mutable vector as a mutable matrix with one row and pass it to
-- the specified function.
withFromRowM :: (Storable e)
                  => STVector s e
                  -> (STMatrix s e -> ST s a)
                  -> ST s a
withFromRowM v f = do
    n <- V.getDim v
    withFromVectorM (1,n) v f
{-# INLINE withFromRowM #-}

-- | Perform an action with a view of a matrix column.
withCol :: (RMatrix m, Storable e)
            => m e
            -> Int
            -> (forall v . RVector v => v e -> ST s a)
            -> ST s a
withCol a j f = do
    (m,n) <- getDim a
    when (j < 0 || j >= n) $ error $
        printf ("withCol <matrix with dim (%d,%d)> %d:"
                ++ " index out of range") m n j

    unsafeWithCol a j f
{-# INLINE withCol #-}

-- | Like 'withCol', but perform the action with a mutable view.
withColM :: (Storable e)
         => STMatrix s e
         -> Int
         -> (STVector s e -> ST s a)
         -> ST s a
withColM a j f = do
    (m,n) <- getDim a
    when (j < 0 || j >= n) $ error $
        printf ("withColM <matrix with dim (%d,%d)> %d:"
                ++ " index out of range") m n j

    unsafeWithColM a j f
{-# INLINE withColM #-}

-- | Create a new matrix of given shape, but do not initialize the elements.
new_ :: (Storable e) => (Int,Int) -> ST s (STMatrix s e)
new_ (m,n)
    | m < 0 || n < 0 = error $
        printf "new_ (%d,%d): invalid dimensions" m n
    | otherwise = unsafeIOToST $ do
        f <- mallocForeignPtrArray (m*n)
        return $ STMatrix $ M.unsafeFromForeignPtr f 0 (m,n) (max 1 m)

-- | Create a matrix with every element initialized to the same value.
new :: (Storable e) => (Int,Int) -> e -> ST s (STMatrix s e)
new (m,n) e = do
    a <- new_ (m,n)
    setElems a $ replicate (m*n) e
    return a

-- | Creates a new matrix by copying another one.
newCopy :: (RMatrix m, Storable e) => m e -> ST s (STMatrix s e)
newCopy a = do
    mn <- getDim a
    b <- new_ mn
    unsafeCopyTo b a
    return b

-- | @copyTo dst src@ replaces the values in @dst@ with those in
-- source.  The operands must be the same shape.
copyTo :: (RMatrix m, Storable e) => STMatrix s e -> m e -> ST s ()
copyTo = checkOp2 "copyTo" unsafeCopyTo
{-# INLINE copyTo #-}

-- | Same as 'copyTo' but does not range-check indices.
unsafeCopyTo :: (RMatrix m, Storable e) => STMatrix s e -> m e -> ST s ()
unsafeCopyTo = vectorOp2 V.unsafeCopyTo
{-# INLINE unsafeCopyTo #-}

-- | Get the indices of the elements in the matrix, in column-major order.
getIndices :: (RMatrix m, Storable e) => m e -> ST s [(Int,Int)]
getIndices a = do
    (m,n) <- getDim a
    return $ [ (i,j) | j <- [ 0..n-1 ], i <- [ 0..m-1 ] ]

-- | Lazily get the elements of the matrix, in column-major order.
getElems :: (RMatrix m, Storable e) => m e -> ST s [e]
getElems a = case maybeWithVector a V.getElems of
    Just es -> es
    Nothing -> withCols a $ \xs ->
                   concat `fmap` mapM V.getElems xs

-- | Get the elements of the matrix, in column-major order.
getElems' :: (RMatrix m, Storable e) => m e -> ST s [e]
getElems' a = case maybeWithVector a V.getElems' of
    Just es -> es
    Nothing -> withCols a $ \xs ->
                   concat `fmap` mapM V.getElems' xs

-- | Lazily get the association list of the matrix, in column-major order.
getAssocs :: (RMatrix m, Storable e) => m e -> ST s [((Int,Int),e)]
getAssocs a = do
    is <- getIndices a
    es <- getElems a
    return $ zip is es

-- | Get the association list of the matrix, in column-major order.
getAssocs' :: (RMatrix m, Storable e) => m e -> ST s [((Int,Int),e)]
getAssocs' a = do
    is <- getIndices a
    es <- getElems' a
    return $ zip is es

-- | Set all of the values of the matrix from the elements in the list,
-- in column-major order.
setElems :: (Storable e) => STMatrix s e -> [e] -> ST s ()
setElems a es =
    case maybeWithVectorM a (`V.setElems` es) of
        Just st  -> st
        Nothing -> do
            (m,n) <- getDim a
            go m n 0 es
  where
    go _ n j [] | j == n = return ()
    go m n j [] | j < n = error $
        printf ("setElems <matrix with dim (%d,%d>"
                ++ "<list with length %d>: not enough elements)") m n (j*m)
    go m n j es' =
        let (es1', es2') = splitAt m es'
        in do
            withColM a j (`V.setElems` es1')
            go m n (j+1) es2'

-- | Set the given values in the matrix.  If an index is repeated twice,
-- the value is implementation-defined.
setAssocs :: (Storable e) => STMatrix s e -> [((Int,Int),e)] -> ST s ()
setAssocs a ies =
    sequence_ [ write a i e | (i,e) <- ies ]

-- | Same as 'setAssocs' but does not range-check indices.
unsafeSetAssocs :: (Storable e) => STMatrix s e -> [((Int,Int),e)] -> ST s ()
unsafeSetAssocs a ies =
    sequence_ [ unsafeWrite a i e | (i,e) <- ies ]

-- | Set the specified row of the matrix to the given vector.
setRow :: (RVector v, Storable e)
       => STMatrix s e -> Int -> v e -> ST s ()
setRow a i x = do
    (m,n) <- getDim a
    nx <- V.getDim x

    when (i < 0 || i >= m) $ error $
        printf ("setRow <matrix with dim (%d,%d)> %d:"
                ++ " index out of range") m n i
    when (nx /= n) $ error $
        printf ("setRow <matrix with dim (%d,%d)> _"
                ++ " <vector with dim %d>:"
                ++ " dimension mismatch") m n nx

    unsafeSetRow a i x
{-# INLINE setRow #-}

-- | Same as 'setRow' but does not range-check index or check
-- vector dimension.
unsafeSetRow :: (RVector v, Storable e)
             => STMatrix s e -> Int -> v e -> ST s ()
unsafeSetRow a i x = do
    jes <- V.getAssocs x
    sequence_ [ unsafeWrite a (i,j) e | (j,e) <- jes ]
{-# INLINE unsafeSetRow #-}

-- | Exchange corresponding elements in the given rows.
swapRows :: (BLAS1 e)
         => STMatrix s e -> Int -> Int -> ST s ()
swapRows a i1 i2 = do
    (m,n) <- getDim a
    when (i1 < 0 || i1 >= m || i2 < 0 || i2 >= m) $ error $
        printf ("swapRows <matrix with dim (%d,%d)> %d %d"
                ++ ": index out of range") m n i1 i2
    unsafeSwapRows a i1 i2

-- | Same as 'swapRows' but does not range-check indices.
unsafeSwapRows :: (BLAS1 e)
               => STMatrix s e -> Int -> Int -> ST s ()
unsafeSwapRows a i1 i2 = when (i1 /= i2) $ do
    (_,n) <- getDim a
    unsafeIOToST $
        unsafeWith a $ \pa lda ->
            let px = pa `advancePtr` i1
                py = pa `advancePtr` i2
                incx = lda
                incy = lda
            in
                BLAS.swap n px incx py incy

-- | Exchange corresponding elements in the given columns.
swapCols :: (BLAS1 e)
         => STMatrix s e -> Int -> Int -> ST s ()
swapCols a j1 j2 = do
    (m,n) <- getDim a
    when (j1 < 0 || j1 >= n || j2 < 0 || j2 >= n) $ error $
        printf ("swapCols <matrix with dim (%d,%d)> %d %d"
                ++ ": index out of range") m n j1 j2
    unsafeSwapCols a j1 j2

-- | Same as 'swapCols' but does not range-check indices.
unsafeSwapCols :: (BLAS1 e)
               => STMatrix s e -> Int -> Int -> ST s ()
unsafeSwapCols a j1 j2 = when (j1 /= j2) $ do
    (m,_) <- getDim a
    unsafeIOToST $
        unsafeWith a $ \pa lda ->
            let px = pa `advancePtr` (j1*lda)
                py = pa `advancePtr` (j2*lda)
                incx = 1
                incy = 1
            in
                BLAS.swap m px incx py incy

-- | Copy the specified row of the matrix to the vector.
rowTo :: (RMatrix m, Storable e)
         => STVector s e -> m e -> Int -> ST s ()
rowTo x a i = do
    (m,n) <- getDim a
    nx <- V.getDim x
    when (i < 0 || i >= m) $ error $
        printf ("rowTo"
               ++ " _"
               ++ " <matrix with dim (%d,%d)>"
               ++ " %d:"
               ++ ": index out of range"
               ) m n i
    when (nx /= n) $ error $
        printf ("rowTo"
                ++ " <vector with dim %d>"
                ++ " <matrix with dim (%d,%d)>"
                ++ " _"
                ++ ": dimension mismatch") nx m n

    unsafeRowTo x a i
{-# INLINE rowTo #-}

-- | Same as 'rowTo' but does not range-check index or check dimension.
unsafeRowTo :: (RMatrix m, Storable e)
            =>  STVector s e -> m e -> Int ->ST s ()
unsafeRowTo x a i = do
    (_,n) <- getDim a
    forM_ [ 0..n-1 ] $ \j -> do
        e <- unsafeRead a (i,j)
        V.unsafeWrite x j e
{-# INLINE unsafeRowTo #-}

-- | Set the diagonal of the matrix to the given vector.
setDiag :: (RVector v, Storable e)
        => STMatrix s e -> v e -> ST s ()
setDiag a x = do
    (m,n) <- getDim a
    nx <- V.getDim x
    let mn = min m n

    when (nx /= mn) $ error $
        printf ("setRow <matrix with dim (%d,%d)>"
                ++ " <vector with dim %d>:"
                ++ " dimension mismatch") m n nx

    unsafeSetDiag a x
{-# INLINE setDiag #-}

-- | Same as 'setDiag' but does not range-check index or check dimension.
unsafeSetDiag :: (RVector v, Storable e)
              => STMatrix s e -> v e -> ST s ()
unsafeSetDiag a x = do
    ies <- V.getAssocs x
    sequence_ [ unsafeWrite a (i,i) e | (i,e) <- ies ]
{-# INLINE unsafeSetDiag #-}

-- | Copy the diagonal of the matrix to the vector.
diagTo :: (RMatrix m, Storable e)
       => STVector s e -> m e -> ST s ()
diagTo x a = do
    nx <- V.getDim x
    (m,n) <- getDim a
    let mn = min m n

    when (nx /= mn) $ error $
        printf ("diagTo"
                ++ " <vector with dim %d>"
                ++ " <matrix with dim (%d,%d)>"
                ++ ": dimension mismatch") nx m n

    unsafeDiagTo x a
{-# INLINE diagTo #-}

-- | Same as 'diagTo' but does not range-check index or check dimensions.
unsafeDiagTo :: (RMatrix m, Storable e)
             => STVector s e -> m e -> ST s ()
unsafeDiagTo x a = do
    (m,n) <- getDim a
    let mn = min m n
    forM_ [ 0..mn-1 ] $ \i -> do
        e <- unsafeRead a (i,i)
        V.unsafeWrite x i e
{-# INLINE unsafeDiagTo #-}

-- | Get the element stored at the given index.
read :: (RMatrix m, Storable e) => m e -> (Int,Int) -> ST s e
read a (i,j) = do
    (m,n) <- getDim a
    when (i < 0 || i >= m || j < 0 || j >= n) $ error $
        printf ("read <matrix with dim (%d,%d)> (%d,%d):"
                ++ " index out of range") m n i j
    unsafeRead a (i,j)
{-# INLINE read #-}

-- | Same as 'read' but does not range-check index.
unsafeRead :: (RMatrix m, Storable e) => m e -> (Int,Int) -> ST s e
unsafeRead a (i,j) = unsafeIOToST $
    unsafeWith a $ \p lda ->
        peekElemOff p (i + j * lda)
{-# INLINE unsafeRead #-}

-- | Set the element stored at the given index.
write :: (Storable e)
      => STMatrix s e -> (Int,Int) -> e -> ST s ()
write a (i,j) e = do
    (m,n) <- getDim a
    when (i < 0 || i >= m || j < 0 || j >= n) $ error $
        printf ("write <matrix with dim (%d,%d)> (%d,%d):"
                ++ " index out of range") m n i j
    unsafeWrite a (i,j) e

{-# INLINE write #-}

-- | Same as 'write' but does not range-check index.
unsafeWrite :: (Storable e)
            => STMatrix s e -> (Int,Int) -> e -> ST s ()
unsafeWrite a (i,j) e = unsafeIOToST $
    unsafeWith a $ \p lda ->
        pokeElemOff p (i + j * lda) e
{-# INLINE unsafeWrite #-}

-- | Modify the element stored at the given index.
modify :: (Storable e)
       => STMatrix s e -> (Int,Int) -> (e -> e) -> ST s ()
modify a (i,j) f = do
    (m,n) <- getDim a
    when (i < 0 || i >= m || j < 0 || j >= n) $ error $
        printf ("modify <matrix with dim (%d,%d)> (%d,%d):"
                ++ " index out of range") m n i j
    unsafeModify a (i,j) f
{-# INLINE modify #-}

-- | Same as 'modify' but does not range-check index.
unsafeModify :: (Storable e)
             => STMatrix s e -> (Int,Int) -> (e -> e) -> ST s ()
unsafeModify a (i,j) f = unsafeIOToST $
    unsafeWith a $ \p lda ->
        let o = i + j * lda
        in do
            e <- peekElemOff p o
            pokeElemOff p o $ f e
{-# INLINE unsafeModify #-}

-- | @mapTo dst f src@ replaces @dst@ elementwise with @f(src)@.
mapTo :: (RMatrix m, Storable e, Storable f)
      => STMatrix s f
      -> (e -> f)
      -> m e
      -> ST s ()
mapTo dst f src = (checkOp2 "mapTo _" $ \z x -> unsafeMapTo z f x) dst src
{-# INLINE mapTo #-}

-- | Same as 'mapTo' but does not check dimensions.
unsafeMapTo :: (RMatrix m, Storable e, Storable f)
            => STMatrix s f
            -> (e -> f)
            -> m e
            -> ST s ()
unsafeMapTo dst f src =
    fromMaybe colwise $ maybeWithVectorM dst $ \vdst ->
    fromMaybe colwise $ maybeWithVector  src $ \vsrc ->
        V.unsafeMapTo vdst f vsrc
  where
    colwise = withColsM dst $ \zs ->
              withCols   src $ \xs ->
                  sequence_ [ V.unsafeMapTo z f x
                            | (z,x) <- zip zs xs
                            ]

-- | @zipWithTo dst f x y@ replaces @dst@ elementwise with @f(x, y)@.
zipWithTo :: (RMatrix m1, RMatrix m2, Storable e1, Storable e2, Storable f)
          => STMatrix s f
          -> (e1 -> e2 -> f)
          -> m1 e1
          -> m2 e2
          -> ST s ()
zipWithTo dst f x y =
    (checkOp3 "zipWithTo _" $ \dst1 x1 y1 -> unsafeZipWithTo dst1 f x1 y1)
        dst x y
{-# INLINE zipWithTo #-}

-- | Same as 'zipWithTo' but does not check dimensions.
unsafeZipWithTo :: (RMatrix m1, RMatrix m2, Storable e1, Storable e2, Storable f)
                => STMatrix s f
                -> (e1 -> e2 -> f)
                -> m1 e1
                -> m2 e2
                -> ST s ()
unsafeZipWithTo dst f x y =
    fromMaybe colwise $ maybeWithVectorM dst $ \vdst ->
    fromMaybe colwise $ maybeWithVector    x $ \vx ->
    fromMaybe colwise $ maybeWithVector    y $ \vy ->
        V.unsafeZipWithTo vdst f vx vy
  where
    colwise = withColsM dst $ \vdsts ->
              withCols   x $ \vxs ->
              withCols   y $ \vys ->

                  sequence_ [ V.unsafeZipWithTo vdst f vx vy
                            | (vdst,vx,vy) <- zip3 vdsts vxs vys
                            ]

-- | Set every element in the matrix to a default value.  For
-- standard numeric types (including 'Double', 'Complex Double', and 'Int'),
-- the default value is '0'.
clear :: (Storable e) => STMatrix s e -> ST s ()
clear a = fromMaybe colwise $ maybeWithVectorM a V.clear
  where
    colwise = withColsM a $ mapM_ V.clear

-- | @withSlice (i,j) (m,n) a@ performs an action with a view of the
-- submatrix of @a@ starting at index @(i,j)@ and having dimension @(m,n)@.
withSlice :: (RMatrix m, Storable e)
          => (Int,Int)
          -> (Int,Int)
          -> m e
          -> (forall m'. RMatrix m' => m' e -> ST s a)
          -> ST s a
withSlice ij mn a f = do
    ia <- unsafeFreeze a
    f $ slice ij mn ia

-- | Like 'withSlice', but perform the action with a mutable view.
withSliceM :: (Storable e)
           => (Int,Int)
           -> (Int,Int)
           -> STMatrix s e
           -> (STMatrix s e -> ST s a)
           -> ST s a
withSliceM ij mn a f =
    withSlice ij mn a $ \a' -> do
        ma <- unsafeThaw a'
        f ma

-- | Perform an action with a view gotten from taking the given number of
-- rows from the start of the matrix.
withTakeRows :: (RMatrix m, Storable e)
             => Int
             -> m e
             -> (forall m'. RMatrix m' => m' e -> ST s a)
             -> ST s a
withTakeRows i a f = do
    ia <- unsafeFreeze a
    f $ takeRows i ia

-- | Like 'withTakeRows', but perform the action with a mutable view.
withTakeRowsM :: (Storable e)
              => Int
              -> STMatrix s e
              -> (STMatrix s e -> ST s a)
              -> ST s a
withTakeRowsM i a f =
    withTakeRows i a $ \a' -> do
        ma <- unsafeThaw a'
        f ma

-- | Perform an action with a view gotten from dropping the given number of
-- rows from the start of the matrix.
withDropRows :: (RMatrix m, Storable e)
             => Int
             -> m e
             -> (forall m'. RMatrix m' => m' e -> ST s a)
             -> ST s a
withDropRows n a f = do
    ia <- unsafeFreeze a
    f $ dropRows n ia

-- | Like 'withDropRows', but perform the action with a mutable view.
withDropRowsM :: (Storable e)
              => Int
              -> STMatrix s e
              -> (STMatrix s e -> ST s a)
              -> ST s a
withDropRowsM i a f =
    withDropRows i a $ \a' -> do
        ma <- unsafeThaw a'
        f ma

-- | Perform an action with views from splitting the matrix rows at the given
-- index.
withSplitRowsAt :: (RMatrix m, Storable e)
                => Int
                -> m e
                -> (forall m1 m2. (RMatrix m1, RMatrix m2) => m1 e -> m2 e -> ST s a)
                -> ST s a
withSplitRowsAt i a f = do
    ia <- unsafeFreeze a
    uncurry f $ splitRowsAt i ia

-- | Like 'withSplitRowsAt', but perform the action with a mutable view.
withSplitRowsAtM :: (Storable e)
                 => Int
                 -> STMatrix s e
                 -> (STMatrix s e -> STMatrix s e -> ST s a)
                 -> ST s a
withSplitRowsAtM i a f =
    withSplitRowsAt i a $ \a1' a2' -> do
        ma1 <- unsafeThaw a1'
        ma2 <- unsafeThaw a2'
        f ma1 ma2

-- | Perform an action with a view gotten from taking the given number of
-- columns from the start of the matrix.
withTakeCols :: (RMatrix m, Storable e)
             => Int
             -> m e
             -> (forall m'. RMatrix m' => m' e -> ST s a)
             -> ST s a
withTakeCols i a f = do
    ia <- unsafeFreeze a
    f $ takeCols i ia

-- | Like 'withTakeCols', but perform the action with a mutable view.
withTakeColsM :: (Storable e)
              => Int
              -> STMatrix s e
              -> (STMatrix s e -> ST s a)
              -> ST s a
withTakeColsM i a f =
    withTakeCols i a $ \a' -> do
        ma <- unsafeThaw a'
        f ma

-- | Perform an action with a view gotten from dropping the given number of
-- columns from the start of the matrix.
withDropCols :: (RMatrix m, Storable e)
             => Int
             -> m e
             -> (forall m'. RMatrix m' => m' e -> ST s a)
             -> ST s a
withDropCols n a f = do
    ia <- unsafeFreeze a
    f $ dropCols n ia

-- | Like 'withDropCols', but perform the action with a mutable view.
withDropColsM :: (Storable e)
              => Int
              -> STMatrix s e
              -> (STMatrix s e -> ST s a)
              -> ST s a
withDropColsM i a f =
    withDropCols i a $ \a' -> do
        ma <- unsafeThaw a'
        f ma

-- | Perform an action with views from splitting the matrix columns at the given
-- index.
withSplitColsAt :: (RMatrix m, Storable e)
                => Int
                -> m e
                -> (forall m1 m2. (RMatrix m1, RMatrix m2) => m1 e -> m2 e -> ST s a)
                -> ST s a
withSplitColsAt i a f = do
    ia <- unsafeFreeze a
    uncurry f $ splitColsAt i ia

-- | Like 'withSplitColsAt', but perform the action with mutable views.
withSplitColsAtM :: (Storable e)
                 => Int
                 -> STMatrix s e
                 -> (STMatrix s e -> STMatrix s e -> ST s a)
                 -> ST s a
withSplitColsAtM i a f =
    withSplitColsAt i a $ \a1' a2' -> do
        ma1 <- unsafeThaw a1'
        ma2 <- unsafeThaw a2'
        f ma1 ma2

-- | Add a vector to the diagonal of a matrix.
shiftDiagM_ :: (RVector v, BLAS1 e)
              => v e -> STMatrix s e -> ST s ()
shiftDiagM_ s a = do
    (m,n) <- getDim a
    ns <- V.getDim s
    let mn = min m n

    when (ns /= mn) $ error $
        printf ("shiftDiagM_"
                ++ " <vector with dim %d>"
                ++ " <matrix with dim (%d,%d)>"
                ++ ": dimension mismatch")
                ns
                m n

    shiftDiagWithScaleM_ 1 s a

-- | Add a scaled vector to the diagonal of a matrix.
shiftDiagWithScaleM_ :: (RVector v, BLAS1 e)
                       => e -> v e -> STMatrix s e -> ST s ()
shiftDiagWithScaleM_ e s a = do
    (m,n) <- getDim a
    ns <- V.getDim s
    let mn = min m n

    when (ns /= mn) $ error $
        printf ("shiftDiagWithScaleM_"
                ++ " _"
                ++ " <vector with dim %d>"
                ++ " <matrix with dim (%d,%d)>"
                ++ ": dimension mismatch")
                ns
                m n

    unsafeIOToST $
        V.unsafeWith s $ \ps ->
        unsafeWith a $ \pa lda ->
            BLAS.axpy mn e ps 1 pa (lda+1)

-- | Add two matrices.
addTo :: (RMatrix m1, RMatrix m2, VNum e)
      =>  STMatrix s e -> m1 e -> m2 e -> ST s ()
addTo = checkOp3 "addTo" $ vectorOp3 V.addTo

-- | Subtract two matrices.
subTo :: (RMatrix m1, RMatrix m2, VNum e)
      => STMatrix s e -> m1 e -> m2 e -> ST s ()
subTo = checkOp3 "subTo" $ vectorOp3 V.subTo

-- | Conjugate the entries of a matrix.
conjugateTo :: (RMatrix m, VNum e)
            => STMatrix s e -> m e -> ST s ()
conjugateTo = checkOp2 "conjugateTo" $
    vectorOp2 V.conjugateTo

-- | Negate the entries of a matrix.
negateTo :: (RMatrix m, VNum e)
         => STMatrix s e -> m e -> ST s ()
negateTo = checkOp2 "negateTo" $
    vectorOp2 V.negateTo

-- | Scale the entries of a matrix by the given value.
scaleM_ :: (BLAS1 e)
          => e -> STMatrix s e -> ST s ()
scaleM_ e = vectorOp (V.scaleM_ e)

-- | @addWithScaleM_ alpha x y@ sets @y := alpha * x + y@.
addWithScaleM_ :: (RMatrix m, BLAS1 e)
               => e -> m e -> STMatrix s e -> ST s ()
addWithScaleM_ e = checkOp2 "addWithScaleM_" $
    unsafeAddWithScaleM_ e

unsafeAddWithScaleM_ :: (RMatrix m, BLAS1 e)
                     => e -> m e -> STMatrix s e -> ST s ()
unsafeAddWithScaleM_ alpha x y =
    fromMaybe colwise $ maybeWithVector  x $ \vx ->
    fromMaybe colwise $ maybeWithVectorM y $ \vy ->
        V.unsafeAddWithScaleM_ alpha vx vy
  where
    colwise = withCols   x $ \vxs ->
              withColsM y $ \vys ->
                  sequence_ [ V.unsafeAddWithScaleM_ alpha vx vy
                            | (vx,vy) <- zip vxs vys ]

-- | Scale the rows of a matrix; @scaleRowsM_ s a@ sets
-- @a := diag(s) * a@.
scaleRowsM_ :: (RVector v, BLAS1 e)
              => v e -> STMatrix s e -> ST s ()
scaleRowsM_  s a = do
    (m,n) <- getDim a
    ns <- V.getDim s
    when (ns /= m) $ error $
        printf ("scaleRowsM_"
                ++ " <vector with dim %d>"
                ++ " <matrix with dim (%d,%d)>"
                ++ ": dimension mismatch")
                ns
                m n

    unsafeIOToST $
        V.unsafeWith s $ \ps ->
        unsafeWith a   $ \pa lda ->
            go m n lda pa ps 0
  where
    go m n lda pa ps i | i == m    = return ()
                       | otherwise = do
                             e <- peek ps
                             BLAS.scal n e pa lda
                             go m n lda (pa `advancePtr` 1)
                                        (ps `advancePtr` 1)
                                        (i+1)

-- | Scale the columns of a matrix; @scaleColBysM_ s a@ sets
-- @a := a * diag(s)@.
scaleColsM_ :: (RVector v, BLAS1 e)
            => v e -> STMatrix s e -> ST s ()
scaleColsM_ s a = do
    (m,n) <- getDim a
    ns <- V.getDim s
    when (ns /= n) $ error $
        printf ("scaleColsM_"
                ++ " <vector with dim %d>"
                ++ " <matrix with dim (%d,%d)>"
                ++ ": dimension mismatch")
                ns
                m n

    es <- V.getElems s
    withColsM a $ \xs ->
        sequence_ [ V.scaleM_ e x
                  | (e,x) <- zip es xs
                  ]

-- | @rank1UpdateM_ alpha x y a@ sets @a := alpha * x * y^H + a@.
rank1UpdateM_ :: (RVector v1, RVector v2, BLAS2 e)
              => e -> v1 e -> v2 e -> STMatrix s e -> ST s ()
rank1UpdateM_ alpha x y a = do
    (m,n) <- getDim a
    nx <- V.getDim x
    ny <- V.getDim y

    when (nx /= m || ny /= n) $ error $
        printf ("rank1UpdateTo"
                ++ " _"
                ++ " <vector with dim %d>"
                ++ " <vector with dim %d>"
                ++ ": dimension mismatch"
                ++ "<matrix with dim (%d,%d)>"                )
                nx
                ny
                m n

    unsafeIOToST $
        V.unsafeWith x $ \px ->
        V.unsafeWith y $ \py ->
        unsafeWith a $ \pa lda ->
            BLAS.gerc m n alpha px 1 py 1 pa lda

-- | @transTo dst a@ sets @dst := trans(a)@.
transTo :: (RMatrix m, BLAS1 e)
        => STMatrix s e
        -> m e
        -> ST s ()
transTo a' a = do
    (ma,na) <- getDim a
    (ma',na') <- getDim a'
    let (m,n) = (ma,na)

    when ((ma,na) /= (na',ma')) $ error $
        printf ( "transTo"
               ++ " <matrix with dim (%d,%d)>"
               ++ " <matrix with dim (%d,%d)>"
               ++ ": dimension mismatch"
               )
               ma' na'
               ma na

    unsafeIOToST $
        unsafeWith a' $ \pa' lda' ->
        unsafeWith a $ \pa lda -> let
            go j px py | j == n = return ()
                       | otherwise = do
                           BLAS.copy m px 1 py lda'
                           go (j+1) (px `advancePtr` lda) (py `advancePtr` 1)
            in go 0 pa pa'

-- | @conjTransTo dst a@ sets @dst := conjugate(trans(a))@.
conjTransTo :: (RMatrix m, BLAS1 e)
            => STMatrix s e
            -> m e
            -> ST s ()
conjTransTo a' a = do
    transTo a' a
    conjugateTo a' a'

-- | @mulVectorTo dst transa a x@
-- sets @dst := op(a) * x@, where @op(a)@ is determined by @transa@.
mulVectorTo :: (RMatrix m, RVector v, BLAS2 e)
            => STVector s e
            -> Trans -> m e
            -> v e
            -> ST s ()
mulVectorTo dst = mulVectorWithScaleTo dst 1

-- | @mulVectorWithScaleTo dst alpha transa a x@
-- sets @dst := alpha * op(a) * x@, where @op(a)@ is determined by @transa@.
mulVectorWithScaleTo :: (RMatrix m, RVector v, BLAS2 e)
                     => STVector s e
                     -> e
                     -> Trans -> m e
                     -> v e
                     -> ST s ()
mulVectorWithScaleTo dst alpha t a x =
    addMulVectorWithScalesM_ alpha t a x 0 dst

-- | @addMulVectorWithScalesM_ alpha transa a x beta y@
-- sets @y := alpha * op(a) * x + beta * y@, where @op(a)@ is
-- determined by @transa@.
addMulVectorWithScalesM_ :: (RMatrix m, RVector v, BLAS2 e)
                         => e
                         -> Trans -> m e
                         -> v e
                         -> e
                         -> STVector s e
                         -> ST s ()
addMulVectorWithScalesM_ alpha transa a x beta y = do
    (ma,na) <- getDim a
    nx <- V.getDim x
    ny <- V.getDim y
    let (m,n) = (ny,nx)

    when ((not . and) [ case transa of NoTrans -> (ma,na) == (m,n)
                                       _       -> (ma,na) == (n,m)
                      , nx == n
                      , ny == m
                      ]) $ error $
        printf ("addMulVectorWithScalesTo"
                ++ " _"
                ++ " %s"
                ++ " <matrix with dim (%d,%d)>"
                ++ " <vector with dim %d>"
                ++ " _"
                ++ " <vector with dim %d>"
                ++ ": dimension mismatch")
               (show transa)
               ma na
               nx
               ny

    unsafeIOToST $
        unsafeWith a $ \pa lda ->
        V.unsafeWith x $ \px ->
        V.unsafeWith y $ \py ->
            if n == 0
                then BLAS.scal m beta py 1
                else BLAS.gemv transa ma na alpha pa lda px 1 beta py 1

-- | @mulMatrixTo dst transa a transb b@
-- sets @dst := op(a) * op(b)@, where @op(a)@ and @op(b)@ are determined
-- by @transa@ and @transb@.
mulMatrixTo :: (RMatrix m1, RMatrix m2, BLAS3 e)
            => STMatrix s e
            -> Trans -> m1 e
            -> Trans -> m2 e
            -> ST s ()
mulMatrixTo dst = mulMatrixWithScaleTo dst 1

-- | @mulMatrixWithScaleTo alpha transa a transb b c@
-- sets @c := alpha * op(a) * op(b)@, where @op(a)@ and @op(b)@ are determined
-- by @transa@ and @transb@.
mulMatrixWithScaleTo :: (RMatrix m1, RMatrix m2, BLAS3 e)
                     => STMatrix s e
                     -> e
                     -> Trans -> m1 e
                     -> Trans -> m2 e
                     -> ST s ()
mulMatrixWithScaleTo dst alpha ta a tb b =
    addMulMatrixWithScalesM_ alpha ta a tb b 0 dst

-- | @addMulMatrixWithScalesM_ alpha transa a transb b beta c@
-- sets @c := alpha * op(a) * op(b) + beta * c@, where @op(a)@ and
-- @op(b)@ are determined by @transa@ and @transb@.
addMulMatrixWithScalesM_ :: (RMatrix m1, RMatrix m2, BLAS3 e)
                         => e
                         -> Trans -> m1 e
                         -> Trans -> m2 e
                         -> e
                         -> STMatrix s e
                         -> ST s ()
addMulMatrixWithScalesM_ alpha transa a transb b beta c = do
    (ma,na) <- getDim a
    (mb,nb) <- getDim b
    (mc,nc) <- getDim c
    let (m,n) = (mc,nc)
        k = case transa of NoTrans -> na
                           _       -> ma

    when ((not . and) [ case transa of NoTrans -> (ma,na) == (m,k)
                                       _       -> (ma,na) == (k,m)
                      , case transb of NoTrans -> (mb,nb) == (k,n)
                                       _       -> (mb,nb) == (n,k)
                      , (mc, nc) == (m,n)
                      ]) $ error $
        printf ("addMulMatrixWithScalesM_"
                ++ " _"
                ++ " %s <matrix with dim (%d,%d)>"
                ++ " %s <matrix with dim (%d,%d)>"
                ++ " _"
                ++ " <matrix with dim (%d,%d)>"
                ++ ": dimension mismatch")
               (show transa) ma na
               (show transb) mb nb
               mc nc

    unsafeIOToST $
        unsafeWith a $ \pa lda ->
        unsafeWith b $ \pb ldb ->
        unsafeWith c $ \pc ldc ->
            BLAS.gemm transa transb m n k alpha pa lda pb ldb beta pc ldc

checkOp2 :: (RMatrix x, RMatrix y, Storable e, Storable f)
         => String
         -> (x e -> y f -> ST s a)
         -> x e
         -> y f
         -> ST s a
checkOp2 str f x y = do
    (m1,n1) <- getDim x
    (m2,n2) <- getDim y
    when ((m1,n1) /= (m2,n2)) $ error $
        printf ("%s <matrix with dim (%d,%d)> <matrix with dim (%d,%d)>:"
                ++ " dimension mismatch") str m1 n1 m2 n2
    f x y
{-# INLINE checkOp2 #-}

checkOp3 :: (RMatrix x, RMatrix y, RMatrix z, Storable e, Storable f, Storable g)
         => String
         -> (x e -> y f -> z g -> ST s a)
         -> x e
         -> y f
         -> z g
         -> ST s a
checkOp3 str f x y z = do
    (m1,n1) <- getDim x
    (m2,n2) <- getDim y
    (m3,n3) <- getDim z
    when((m1,n1) /= (m2,n2) || (m1,n1) /= (m3,n3)) $ error $
        printf ("%s <matrix with dim (%d,%d)> <matrix with dim (%d,%d)>:"
                ++ " <matrix with dim (%d,%d)> dimension mismatch")
               str m1 n1 m2 n2 m3 n3
    f x y z
{-# INLINE checkOp3 #-}

vectorOp :: (Storable e)
         => (STVector s e -> ST s ())
         -> STMatrix s e -> ST s ()
vectorOp f x =
    fromMaybe colwise $ maybeWithVectorM x $ \vx -> f vx
  where
    colwise = withColsM x $ \vxs ->
                  sequence_ [ f vx | vx <- vxs ]

vectorOp2 :: (RMatrix m, Storable e, Storable f)
          => (forall v . RVector v => STVector s f -> v e -> ST s ())
          -> STMatrix s f -> m e -> ST s ()
vectorOp2 f dst x =
    fromMaybe colwise $ maybeWithVectorM dst $ \vdst ->
    fromMaybe colwise $ maybeWithVector    x $ \vx ->
        f vdst vx
  where
    colwise = withColsM dst $ \vdsts ->
              withCols   x   $ \vxs ->
                  sequence_ [ f vdst vx | (vdst,vx) <- zip vdsts vxs ]
{-# INLINE vectorOp2 #-}

vectorOp3 :: (RMatrix m1, RMatrix m2, Storable e1, Storable e2, Storable f)
          => (forall v1 v2 . (RVector v1, RVector v2) =>
                  STVector s f -> v1 e1 -> v2 e2 -> ST s ())
          -> STMatrix s f -> m1 e1 -> m2 e2 -> ST s ()
vectorOp3 f dst x y =
    fromMaybe colwise $ maybeWithVectorM dst $ \vdst ->
    fromMaybe colwise $ maybeWithVector    x $ \vx ->
    fromMaybe colwise $ maybeWithVector    y $ \vy ->
        f vdst vx vy
  where
    colwise = withColsM dst $ \vdsts ->
              withCols   x   $ \vxs ->
              withCols   y   $ \vys ->
                  sequence_ [ f vdst vx vy
                            | (vdst,vx,vy) <- zip3 vdsts vxs vys ]
{-# INLINE vectorOp3 #-}


--
-- (c) Susumu Katayama
--
{-# LANGUAGE UndecidableInstances, OverlappingInstances, TemplateHaskell, CPP, FlexibleInstances #-}
-- x #define TESTEQ

module MagicHaskeller.Classification -- (
                      -- randomTestFilter, -- ::  Filtrable a => (b->a) -> Matrix b -> Matrix b
                      -- )
                      where
import Prelude hiding ((/))
#ifdef TFRANDOM
import System.Random.TF.Gen
#else
import System.Random
#endif
import MagicHaskeller.MyCheck
import Data.Char
import Data.List
import Control.Monad
import Control.Monad.Search.Combinatorial
import Data.Complex

import MagicHaskeller.MHTH
import MagicHaskeller.TimeOut
import MagicHaskeller.T10

import MagicHaskeller.Instantiate(compareRealFloat)

class (Search m) => SStrategy m where
    sfilter :: Relation r =>
               (k->k->r) -> (Int->Int) -> m ([k],e) -> m ([k],e)
    ofilter :: Relation r =>
               (k->k->r) -> m (k,e) -> m (k,e)

instance SStrategy Matrix where
    sfilter = sfilterMx
    ofilter = ofilterMx

instance SStrategy DBound where
    sfilter = sfilterDB
    ofilter = ofilterDB

#ifdef TFRANDOM
arbitraries :: Arbitrary a => [a]
arbitraries = arbs 0 (seedTFGen (12279932681387497184,218462391894233257,1625759680792809304,12756118543158863164))
arbs :: Arbitrary a => Int -> TFGen -> [a]
arbs n stdgen  =  case map (splitn stdgen 8) [0..255] of
                    g0:gs -> zipWith f [n..] gs ++ arbs (n+255) g0 -- I think 255 seeds is enough, but just in case.
    where Gen f = arbitrary
#else
arbitraries :: Arbitrary a => [a]
arbitraries = arbs 0 (mkStdGen 1)
arbs :: Arbitrary a => Int -> StdGen -> [a]
arbs n stdgen  =  case split stdgen of
                    (g0,g1) -> f n g0 : arbs (n+1) g1
    where Gen f = arbitrary
#endif

(/~) :: [a] -> (a->a->Bool) -> [[a]]
[]      /~  eq  =  []
(x:xs)  /~  eq  =  case partition (x `eq`) xs of
                     (same, diff) -> (x:same) : (diff /~ eq)
{-
$B$J$*!$>e5-(B(/)$B$*$h$S2<5-(Bnub$B$N>l9g!$L58B%j%9%H$G$b$$$1$k!%(B
*T10> let {[]     / eq = []; (x:xs) / eq = case partition (x `eq`) xs of (same, diff) -> (x:same) : (diff / eq)}
*T10> cycle "hogeha" / (==)
["hhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhh
(snip)
hhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhInterrupted.
*T10> map head $ cycle "hogeha" / (==)
"hogeaInterrupted.
*T10> Data.List.nub $ cycle "hogeha"
"hogeaTerminated
(snip)
*T10> Data.List.nub $ repeat 'h'
"hTerminated
-}

-- ToDo: deal with timeout

-- complete set of representatives
nubSortBy :: (a->a->Ordering) -> [a] -> [a]
nubSortBy = mergesortWithBy const
nubSortByBot :: (a->a->Maybe Ordering) -> [a] -> [a]
nubSortByBot = mergesortWithByBot const
-- quotient set
(/<) :: [a] -> (a->a->Ordering) -> [[a]]
xs /< cmp  =  mergesortWithBy (++)
                (\x y -> head x `cmp` head y)
                (map return xs)
(/<?) :: [a] -> (a->a->Maybe Ordering) -> [[a]]
xs /<? cmp  =  mergesortWithByBot (++)
                (\x y -> head x `cmp` head y)
                (map return xs)
class Eq rel => Relation rel where
 -- remove duplicates, and sort if |rel==Ordering|
    fromListBy   :: (k->k->rel) -> [k] -> [k]
    fromListBy cmp = map head . (/cmp)
-- used to pick up the shallowest expression in |DBound|
    fromListByDB :: (k->k->rel) -> [(k,Int)] -> [(k,Int)]
    fromListByDB rel ts =
          map  (minimumBy (\x y -> compare (snd y) (snd x)))
               (ts / (\x y -> rel (fst x) (fst y)))
    -- \NB : |maximumBy| returns the last of the maxima,
    -- while |minimumBy| the first of the minima.
 -- quotient set
    (/)          :: [k] -> (k->k->rel) -> [[k]]
 -- merge two lists,
    appendWithBy :: (k->k->k) -> -- combiner
                    (k->k->rel) -> [k] -> [k] -> [k]
    diffBy       :: (k -> k -> rel) -> [k] -> [k] -> [k]
 -- counterpart of EQ
    cEQ          :: rel
-- merge two quotient sets
appendQuotientsBy ::
  (Relation rel) =>
  (k -> k -> rel) -> [[k]] -> [[k]] -> [[k]]
appendQuotientsBy rel =
   appendWithBy (++) (\ (x:_) (y:_) -> x `rel` y)
-- merge two complete sets of representatives
appendRepresentativesBy ::
   (Relation rel) =>
   (k -> k -> rel) -> [k] -> [k] -> [k]
appendRepresentativesBy = appendWithBy const

instance Relation Bool where
    fromListBy    = Data.List.nubBy
    (/)           = (/~)
    appendWithBy  = unionWithBy
    diffBy        = Data.List.deleteFirstsBy
    cEQ           = True
unionWithBy combiner eq []      ys = ys
unionWithBy combiner eq (x:xs)  ys =
   case break (eq x) ys of
      (_,   [])    ->  x  :  unionWithBy combiner eq xs ys
      (ts,  h:ds)  ->  (x `combiner` h)
                          :  unionWithBy combiner eq xs (ts++ds)
instance Relation Ordering where
    fromListBy    = nubSortBy
    fromListByDB rel =
        mergesortWithBy
                    (\x y -> if snd x < snd y then y else x)
                    (\x y -> fst x `rel` fst y)
    (/)           = (/<)
    appendWithBy  = mergeWithBy
    diffBy        = diffSortedBy
    cEQ           = EQ
instance Relation (Maybe Ordering) where
    fromListBy    = nubSortByBot
    fromListByDB rel =
        mergesortWithByBot
                    (\x y -> if snd x < snd y then y else x)
                    (\x y -> fst x `rel` fst y)
    (/)           = (/<?)
    appendWithBy  = mergeWithByBot
    diffBy        = diffSortedByBot
    cEQ           = Just EQ

randomTestFilter ::  (SStrategy m, Filtrable a) =>
                     (Int->Int) -> m (e,a) -> m (e,a)
randomTestFilter numRandoms = filt numRandoms . fmap (\ t@(_,a) -> (a,t))

unsafeRandomTestFilter ::  (SStrategy m, Filtrable a) =>
                           Maybe Int -- ^ microsecs until timeout
                               -> (Int->Int) -> m (e,a) -> m (e,a)
unsafeRandomTestFilter mto numRandoms = unsafeFilt mto numRandoms . fmap (\ t@(_,a) -> (a,t))

mapFst f (a,b) = (f a, b)

class Filtrable a where
    filt    :: SStrategy m => (Int->Int) -> m (a,e) -> m e
    filtFun :: (SStrategy m, Arbitrary b) =>
               (Int->Int) -> m (b->a,e) -> m e
    unsafeFilt    :: SStrategy m =>
                     Maybe Int -> (Int->Int) -> m (a,e) -> m e
    unsafeFiltFun :: (SStrategy m, Arbitrary b) =>
                     Maybe Int -> (Int->Int) -> m (b->a,e) -> m e

instance  (Arbitrary a, Filtrable r) => Filtrable (a->r)
  where
    filt      = filtFun
    filtFun f = filt f . fmap (mapFst uncurry)
    unsafeFilt    mto f = unsafeFiltFun mto f
    unsafeFiltFun mto f = unsafeFilt mto f . fmap (mapFst uncurry)

#ifdef TESTEQ
instance Eq a => Filtrable a where
    filt     = filtNullary  (==)
    filtFun  = filtUnary    (==)
#else
instance Ord a => Filtrable a where
    filt     = filtNullary  compare
    filtFun  = filtUnary    compare
    unsafeFilt    mto = filtNullary  (unsafeOpWithPTO mto compare)
    unsafeFiltFun mto = filtUnary    (unsafeOpWithPTO mto compare)
instance Filtrable Double where
    filt     = filtNullary  compareRealFloat
    filtFun  = filtUnary    compareRealFloat
    unsafeFilt    mto = filtNullary  (unsafeOpWithPTO mto compareRealFloat)
    unsafeFiltFun mto = filtUnary    (unsafeOpWithPTO mto compareRealFloat)
#endif
filtNullary ::  (SStrategy m, Relation r) =>
                (k->k->r) -> (Int->Int) -> m (k,e) -> m e
filtNullary  op _ =  fmap snd . ofilter op
filtUnary    op f =  fmap snd . sfilter op f .
                      fmap (mapFst (flip map arbitraries))

instance  (RealFloat a, Ord a) =>
          Filtrable (Complex a) where
    filt     = filtNullary  compareCx
    filtFun  = filtUnary    compareCx
    unsafeFilt    mto = filtNullary  (unsafeOpWithPTO mto compareCx)
    unsafeFiltFun mto = filtUnary    (unsafeOpWithPTO mto compareCx)
compareCx ::  (RealFloat a, Ord a) =>
              Complex a -> Complex a -> Ordering
(a:+b) `compareCx` (c:+d) = case compare a c of
                              EQ -> compare b d
                              o  -> o
ofilterMx ::  Relation r =>
            (k->k->r) -> Matrix (k,e) -> Matrix (k,e)
ofilterMx op (Mx xss)
        = let
            (k,_) `rel` (l,_) = k `op` l
            mapped     = map (fromListBy rel) xss
            cumulative =
                 scanl  (appendRepresentativesBy rel)
                                 [] mapped
          in Mx $ zipWith (diffBy rel) mapped cumulative

ofilterDB :: Relation rel =>
              (k->k->rel) ->
               DBound (k,e) -> DBound (k,e)
ofilterDB cmp (DB f) = DB $
    \n -> fromListByDB  (\(k,_) (l,_) -> cmp k l)
                        (f n)

cumulativeRepresentatives ::
    Relation rel =>
   [a->a->rel] -> Matrix a -> Matrix a
cumulativeRepresentatives relations mx =
    fmap head (cumulativeQuotients relations mx)

representatives ::
    Relation rel =>
   [a->a->rel] -> Matrix a -> Matrix a
representatives relations mx =
   unscanlByList relations $
        cumulativeRepresentatives relations mx
unscanlByList ::  Relation r =>
                  [k->k->r] -> Matrix k -> Matrix k
unscanlByList (_:rels) (Mx (yss@(xs:xss))) =
    Mx $ xs : zipWith3 diffBy rels xss yss

sfilterMx ::  Relation r =>
              (k->k->r) ->
              (Int->Int) ->
              Matrix ([k],e) -> Matrix ([k],e)
sfilterMx rel numRands = representatives (map (liftRelation rel . numRands) [0..])
liftRelation ::  Relation r =>
                 (k->k->r) ->
                    Int -> ([k],e) -> ([k],e) -> r
liftRelation rel len (xs,_) (ys,_) = liftRel rel len xs ys
liftRel _   0   _      _      = cEQ
liftRel _   _   []     []     = cEQ
liftRel rel len (x:xs) (y:ys) =
    case rel x y of
           c  | c == cEQ   -> liftRel rel (len-1) xs ys
              | otherwise  -> c

sfilterDB ::  Relation rel =>
               (k->k->rel) ->
              (Int->Int)->
                DBound ([k],e) -> DBound ([k],e)
sfilterDB rel numRands (DB f) = DB $ \n -> fromListByDB  (liftRelation rel (numRands n)) (f n)

cumulativeQuotients relations (Mx xss)
   =  let yss:ysss = zipWith (/) xss relations
      in Mx $ scanl  (\rec (r,z) ->
                       appendQuotientsBy r (rec>>=(/r)) z)
                     yss  (zip (tail relations) ysss)

ns = [6..]


{-# LANGUAGE BangPatterns #-}
module MnistLoader
  ( load_data_wrapper
  , TrainingData
  , ValidationData
  ) where

import           Codec.Compression.GZip (decompress)
import           Control.Applicative ((<$>))
import           Control.DeepSeq
import           Control.Monad (replicateM, unless)
import           Data.Binary.Get
import qualified Data.ByteString.Lazy as B
import           Data.Word (Word8)
import           Numeric.LinearAlgebra
import           System.FilePath.Posix ((</>))

type TrainingData   = [(Vector Double, Vector Double)]
type ValidationData = [(Vector Double, Word8)]

readImageFile :: FilePath -> IO [[Word8]]
readImageFile p = runGet parseImages <$> decompress <$> B.readFile p

readLabelFile :: FilePath -> IO [Word8]
readLabelFile p = runGet parseLabels <$> decompress <$> B.readFile p

parseImages :: Get [[Word8]]
parseImages = do
  magicNum <- getWord32be
  unless (magicNum == 2051) $ error $
    "Not an image file, wrong magic number " ++ show magicNum
  nImg <- fromIntegral <$> getWord32be
  nRow <- fromIntegral <$> getWord32be
  nCol <- fromIntegral <$> getWord32be
  img  <- replicateM nImg $ replicateM (nRow * nCol) getWord8
  unlessM isEmpty $ error "not all data consumed"
  return img

parseLabels :: Get [Word8]
parseLabels = do
  magicNum <- getWord32be
  unless (magicNum == 2049) $ error $
    "Not a label file, wrong magic number " ++ show magicNum
  n <- fromIntegral <$> getWord32be
  lbl <- replicateM n getWord8
  unlessM isEmpty $ error "not all data consumed"
  return lbl

load_data_wrapper :: FilePath -> IO (TrainingData, ValidationData)
load_data_wrapper p = do
  imgTraining    <- readImageFile $ p </> "train-images-idx3-ubyte.gz"
  imgValidation  <- readImageFile $ p </> "t10k-images-idx3-ubyte.gz"
  lblTraining    <- readLabelFile $ p </> "train-labels-idx1-ubyte.gz"
  lblValidation  <- readLabelFile $ p </> "t10k-labels-idx1-ubyte.gz"
  return ( (zip $!! normalizeImg imgTraining)   $!! toUnitVector lblTraining
         , (zip $!! normalizeImg imgValidation) $!! lblValidation
         )
  where
    normalizeImg = map $ fromList . map ((/ 256) . fromIntegral)
    toUnitVector = map (fromList . toUnitVector')
    toUnitVector' 0 = [1, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    toUnitVector' 1 = [0, 1, 0, 0, 0, 0, 0, 0, 0, 0]
    toUnitVector' 2 = [0, 0, 1, 0, 0, 0, 0, 0, 0, 0]
    toUnitVector' 3 = [0, 0, 0, 1, 0, 0, 0, 0, 0, 0]
    toUnitVector' 4 = [0, 0, 0, 0, 1, 0, 0, 0, 0, 0]
    toUnitVector' 5 = [0, 0, 0, 0, 0, 1, 0, 0, 0, 0]
    toUnitVector' 6 = [0, 0, 0, 0, 0, 0, 1, 0, 0, 0]
    toUnitVector' 7 = [0, 0, 0, 0, 0, 0, 0, 1, 0, 0]
    toUnitVector' 8 = [0, 0, 0, 0, 0, 0, 0, 0, 1, 0]
    toUnitVector' 9 = [0, 0, 0, 0, 0, 0, 0, 0, 0, 1]
    toUnitVector' x = error $ "(toUnitVector " ++ show x ++ ") is undefined"

unlessM :: Monad m => m Bool -> m () -> m ()
unlessM mp f = (\p -> unless p f) =<< mp