{-# LANGUAGE CPP                                      #-}
{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE DeriveFoldable                           #-}
{-# LANGUAGE DeriveFunctor                            #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE FlexibleInstances                        #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE LambdaCase                               #-}
{-# LANGUAGE OverloadedStrings                        #-}
{-# LANGUAGE PatternSynonyms                          #-}
{-# LANGUAGE RecordWildCards                          #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE StandaloneDeriving                       #-}
{-# LANGUAGE TupleSections                            #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE TypeSynonymInstances                     #-}
{-# LANGUAGE ViewPatterns                             #-}
{-# OPTIONS_GHC -fno-warn-orphans                     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

-- | Hamilton example suite
--
-- See: https://github.com/mstksg/hamilton#example-app-runner
--
-- Or just run with:
--
-- > $ hamilton-examples --help
-- > $ hamilton-examples [EXAMPLE] --help
--

import           Control.Concurrent
import           Control.Monad
import           Data.Bifunctor
import           Data.Finite
import           Data.Foldable
import           Data.IORef
import           Data.List
import           Data.Semigroup                      ((<>))
import           GHC.TypeLits
import           Graphics.Vty hiding                 (Config, (<|>))
import           Numeric.Hamilton
import           Numeric.LinearAlgebra.Static hiding (dim, (<>))
import           Numeric.LinearAlgebra.Static.Vector
import           Options.Applicative
import           System.Exit
import           Text.Printf
import           Text.Read
import qualified Data.List.NonEmpty                  as NE
import qualified Data.Map.Strict                     as M
import qualified Data.Vector                         as VV
import qualified Data.Vector.Sized                   as V
import qualified Data.Vector.Storable.Sized          as VS
import qualified Text.PrettyPrint.ANSI.Leijen        as PP

data SysExample where
    SE :: (KnownNat m, KnownNat n)
       => { seName   :: String
          , seCoords :: V.Vector n String
          , seSystem :: System m n
          , seDraw   :: R m -> [V2 Double]
          , seInit   :: Phase n
          }
       -> SysExample

pendulum :: Double -> Double -> SysExample
pendulum θ0 ω0 = SE "Single pendulum" (V1 "θ") s f (toPhase s c0)
  where
    s :: System 2 1
    s = mkSystem' (vec2 1 1                             )     -- masses
                  (\(V1 θ)   -> V2 (sin θ) (0.5 - cos θ))     -- coordinates
                  (\(V2 _ y) -> y                       )     -- potential
    f :: R 2 -> [V2 Double]
    f xs = [grVec xs]
    c0 :: Config 1
    c0 = Cfg (konst θ0 :: R 1) (konst ω0 :: R 1)

doublePendulum :: Double -> Double -> SysExample
doublePendulum m1 m2 = SE "Double pendulum" (V2 "θ1" "θ2") s f (toPhase s c0)
  where
    s :: System 4 2
    s = mkSystem' (vec4 m1 m1 m2 m2)     -- masses
                  (\(V2 θ1 θ2)     -> V4 (sin θ1)            (1 - cos θ1)
                                         (sin θ1 + sin θ2/2) (1 - cos θ1 - cos θ2/2)
                  )                      -- coordinates
                  (\(V4 _ y1 _ y2) -> 5 * (realToFrac m1 * y1 + realToFrac m2 * y2))
                                         -- potential
    f :: R 4 -> [V2 Double]
    f (split->(xs,ys))= grVec <$> [xs, ys]
    c0 :: Config 2
    c0 = Cfg (vec2 (pi/2) 0) (vec2 0 0)

room :: Double -> SysExample
room θ = SE "Room" (V2 "x" "y") s f (toPhase s c0)
  where
    s :: System 2 2
    s = mkSystem (vec2 1 1)         -- masses
                 id                 -- coordinates
                 (\(V2 x y) -> sum [ 2 * y                      -- gravity
                                   , 1 - logistic (-1) 10 0.1 y  -- bottom wall
                                   ,     logistic 1 10 0.1 y     -- top wall
                                   , 1 - logistic (-2) 10 0.1 x  -- left wall
                                   ,     logistic 2 10 0.1 x     -- right wall
                                   ]
                 )                  -- potential
    f :: R 2 -> [V2 Double]
    f xs = [grVec xs]
    c0 :: Config 2
    c0 = Cfg (vec2 (-1) 0.25) (vec2 (cos θ) (sin θ))

twoBody :: Double -> Double -> Double -> SysExample
twoBody m1 m2 ω0 = SE "Two-Body" (V2 "r" "θ") s f (toPhase s c0)
  where
    mT :: Double
    mT = m1 + m2
    s :: System 4 2
    s = mkSystem (vec4 m1 m1 m2 m2) -- masses
                 -- positions are calculated assuming (0,0) is the center
                 -- of mass
                 (\(V2 r θ) -> let r1 = r * realToFrac (-m2 / mT)
                                   r2 = r * realToFrac (m1 / mT)
                               in  V4 (r1 * cos θ) (r1 * sin θ)
                                      (r2 * cos θ) (r2 * sin θ)
                 )                 -- coordinates
                 (\(V2 r _) -> - realToFrac (m1 * m2) / r)  -- potential
    f :: R 4 -> [V2 Double]
    f (split->(xs,ys))= grVec <$> [xs, ys]
    c0 :: Config 2
    c0 = Cfg (vec2 2 0) (vec2 0 ω0)

spring
    :: Double -> Double -> Double -> Double -> SysExample
spring mB mW k x0 = SE "Spring hanging from block" (V3 "r" "x" "θ") s f (toPhase s c0)
  where
    s :: System 3 3
    s = mkSystem (vec3 mB mW mW)                                                  -- masses
                 (\(V3 r x θ)  -> V3 r (r + (1 + x) * sin θ) ((1 + x) * (-cos θ))) -- coordinates
                 (\(V3 r x θ) -> realToFrac k * x**2 / 2        -- spring
                              + (1 - logistic (-1.5) 25 0.1 r)  -- left rail wall
                              + (    logistic   1.5  25 0.1 r)  -- right rail wall
                              + realToFrac mB * ((1 + x) * (-cos θ))  -- gravity
                 )
    f :: R 3 -> [V2 Double]
    f (headTail->(b,w)) = [V2 b 1, V2 0 1 + grVec w]
    c0 :: Config 3
    c0 = Cfg (vec3 0 x0 0) (vec3 1 0 (-0.5))

bezier
    :: forall n. KnownNat (1 + n)
    => V.Vector (1 + n) (V2 Double)
    -> SysExample
bezier ps = SE "Bezier" (V1 "t") s f (toPhase s c0)
  where
    s :: System 2 1
    s = mkSystem (vec2 1 1)                                             -- masses
                 (\(V1 t) -> bezierCurve (fmap realToFrac <$> ps) t)    -- coordinates
                 (\(V1 t) -> (1 - logistic 0 5 0.05 t)           -- left wall
                           +      logistic 1 5 0.05 t            -- right wall
                 )
    f :: R 2 -> [V2 Double]
    f xs = [grVec xs]
    c0 :: Config 1
    c0 = Cfg (0.5 :: R 1) (0.25 :: R 1)

data ExampleOpts = EO { eoChoice :: SysExampleChoice }

data SysExampleChoice =
        SECDoublePend Double Double
      | SECPend Double Double
      | SECRoom Double
      | SECTwoBody Double Double Double
      | SECSpring Double Double Double Double
      | SECBezier (NE.NonEmpty (V2 Double))

parseEO :: Parser ExampleOpts
parseEO = EO <$> (parseSEC <|> pure (SECDoublePend 1 1))

parseSEC :: Parser SysExampleChoice
parseSEC = subparser . mconcat $
    [ command "doublepend" $
        info (helper <*> parseDoublePend)
             (progDesc "Double pendulum (default)")
    , command "pend"       $
        info (helper <*> parsePend      )
             (progDesc "Single pendulum")
    , command "room"       $
        info (helper <*> parseRoom      )
        (progDesc "Ball in room, bouncing off of walls")
    , command "twobody"    $
        info (helper <*> parseTwoBody    )
        (progDesc "Two-body graviational simulation.  Note that bodies will only orbit if H < 0.")
    , command "spring"    $
        info (helper <*> parseSpring    )
        (progDesc "A spring hanging from a block on a rail, holding up a mass.  Block is constrained to bounce between -1.5 and 1.5.")
    , command "bezier"     $
        info (helper <*> parseBezier    )
        (progDesc "Particle moving along a parameterized bezier curve")
    , metavar "EXAMPLE"
    ]
  where
    parsePend
      = SECPend       <$> option auto ( long "angle"
                                     <> short 'a'
                                     <> metavar "ANGLE"
                                     <> help "Intitial rightward angle (in degrees) of bob"
                                     <> value 0
                                     <> showDefault
                                      )
                      <*> option auto ( long "vel"
                                     <> short 'v'
                                     <> metavar "VELOCITY"
                                     <> help "Initial rightward angular velocity of bob"
                                     <> value 1
                                     <> showDefault
                                      )
    parseDoublePend
      = SECDoublePend <$> option auto ( long "m1"
                                     <> metavar "MASS"
                                     <> help "Mass of first bob"
                                     <> value 1
                                     <> showDefault
                                      )
                      <*> option auto ( long "m2"
                                     <> metavar "MASS"
                                     <> help "Mass of second bob"
                                     <> value 1
                                     <> showDefault
                                      )
    parseRoom
      = SECRoom    <$> option auto ( long "angle"
                                  <> short 'a'
                                  <> metavar "ANGLE"
                                  <> help "Initial upward launch angle (in degrees) of object"
                                  <> value 45
                                  <> showDefault
                                   )
    parseTwoBody
      = SECTwoBody <$> option auto ( long "m1"
                                  <> metavar "MASS"
                                  <> help "Mass of first body"
                                  <> value 5
                                  <> showDefault
                                   )
                   <*> option auto ( long "m2"
                                  <> metavar "MASS"
                                  <> help "Mass of second body"
                                  <> value 0.5
                                  <> showDefault
                                   )
                   <*> option auto ( long "vel"
                                  <> short 'v'
                                  <> metavar "VELOCITY"
                                  <> help "Initial angular velocity of system"
                                  <> value 0.5
                                  <> showDefault
                                   )
    parseSpring
      = SECSpring <$> option auto ( long "block"
                                 <> short 'b'
                                 <> metavar "MASS"
                                 <> help "Mass of block on rail"
                                 <> value 2
                                 <> showDefault
                                  )
                  <*> option auto ( long "weight"
                                 <> short 'w'
                                 <> metavar "MASS"
                                 <> help "Mass of weight hanging from spring"
                                 <> value 1
                                 <> showDefault
                                  )
                  <*> option auto ( short 'k'
                                 <> metavar "NUM"
                                 <> help "Spring constant / stiffness of spring"
                                 <> value 10
                                 <> showDefault
                                  )
                  <*> option auto ( short 'x'
                                 <> metavar "DIST"
                                 <> help "Initial displacement of spring"
                                 <> value 0.1
                                 <> showDefault
                                  )
    parseBezier
      = SECBezier <$> option f ( long "points"
                              <> short 'p'
                              <> metavar "POINTS"
                              <> help "List of control points (at least one), as tuples"
                              <> value (V2 (-1) (-1) NE.:| [V2 (-2) 1, V2 0 1, V2 1 (-1), V2 2 1])
                              <> showDefaultWith (show . map (\(V2 x y) -> (x, y)) . toList)
                               )
      where f = eitherReader $ \s -> do
              ps  <- maybe (Left "Bad parse") Right
                  $ readMaybe s
              maybe (Left "At least one control point required") Right
                  $ NE.nonEmpty (uncurry V2 <$> ps)

data SimOpts = SO { soZoom :: Double
                  , soRate :: Double
                  , soHist :: Int
                  }
  deriving (Show)

data SimEvt = SEQuit
            | SEZoom Double
            | SERate Double
            | SEHist Int

main :: IO ()
main = do
    EO{..} <- execParser $ info (helper <*> parseEO)
        ( fullDesc
       <> header "hamilton-examples - hamilton library example suite"
       <> progDescDoc (Just descr)
        )

    vty <- mkVty =<< standardIOConfig

    opts <- newIORef $ SO 0.5 1 25

    t <- forkIO . loop vty opts $ case eoChoice of
      SECDoublePend m1 m2        -> doublePendulum m1 m2
      SECPend       d0 ω0        -> pendulum (d0 / 180 * pi) ω0
      SECRoom       d0           -> room (d0 / 180 * pi)
      SECTwoBody    m1 m2 ω0     -> twoBody m1 m2 ω0
      SECSpring     mB mW k x0   -> spring mB mW k x0
      SECBezier     (p NE.:| ps) -> V.withSized (VV.fromList ps)
                                      (bezier . V.cons p)

    forever $ do
      e <- nextEvent vty
      forM_ (processEvt e) $ \case
        SEQuit -> do
          killThread t
          shutdown vty
          exitSuccess
        SEZoom s ->
          modifyIORef opts $ \o -> o { soZoom = soZoom o * s }
        SERate r ->
          modifyIORef opts $ \o -> o { soRate = soRate o * r }
        SEHist h ->
          modifyIORef opts $ \o -> o { soHist = soHist o + h }
  where
    fps :: Double
    fps = 12
    screenRatio :: Double
    screenRatio = 2.1
    ptAttrs :: [(Char, Color)]
    ptAttrs  = ptChars `zip` ptColors
      where
        ptColors = cycle [white,yellow,blue,red,green]
        ptChars  = cycle "o*+~"
    loop :: Vty -> IORef SimOpts -> SysExample -> IO ()
    loop vty oRef SE{..} = go M.empty seInit
      where
        qVec = intercalate "," . V.toList $ seCoords
        go hists p = do
          SO{..} <- readIORef oRef
          let p'   = stepHam (soRate / fps) seSystem p  -- progress the simulation
              xb   = (- recip soZoom, recip soZoom)
              infobox = vertCat . map (string defAttr) $
                  [ printf "[ %s ]" seName
                  , printf " <%s>   : <%s>" qVec . intercalate ", "
                     . map (printf "%.4f") . VS.toList . rVec . phsPositions $ p
                  , printf "d<%s>/dt: <%s>" qVec . intercalate ", "
                     . map (printf "%.4f") . VS.toList . rVec . velocities seSystem $ p
                  , printf "KE: %.4f" . keP seSystem           $ p
                  , printf "PE: %.4f" . pe seSystem . phsPositions $ p
                  , printf "H : %.4f" . hamiltonian seSystem   $ p
                  , " "
                  , printf "rate: x%.2f <>" $ soRate
                  , printf "hist: % 5d []" $ soHist
                  , printf "zoom: x%.2f -+" $ soZoom
                  ]
              pts  = (`zip` ptAttrs) . seDraw . underlyingPos seSystem . phsPositions
                   $ p
              hists' = foldl' (\h (r, a) -> M.insertWith (addHist soHist) a [r] h) hists pts
          dr <- displayBounds $ outputIface vty
          update vty . picForLayers . (infobox:) . plot dr (PX xb (RR 0.5 screenRatio)) $
               ((second . second) (defAttr `withForeColor`) <$> pts)
            ++ (map (\((_,c),r) -> (r, ('.', defAttr `withForeColor` c)))
                  . concatMap sequence
                  . M.toList
                  $ hists'
               )
          threadDelay (round (1000000 / fps))
          go hists' p'
    addHist hl new old = take hl (new ++ old)
    descr :: PP.Doc
    descr = PP.vcat
      [ "Run examples from the hamilton library example suite."
      , "Use with [EXAMPLE] --help for more per-example options."
      , ""
      , "To adjust rate/history/zoom, use keys <>/[]/-+, respectively."
      , ""
      , "See: https://github.com/mstksg/hamilton#example-app-runner"
      ]

processEvt
    :: Event -> Maybe SimEvt
processEvt = \case
    EvKey KEsc        []      -> Just SEQuit
    EvKey (KChar 'c') [MCtrl] -> Just SEQuit
    EvKey (KChar 'q') []      -> Just SEQuit
    EvKey (KChar '+') []      -> Just $ SEZoom (sqrt 2)
    EvKey (KChar '-') []      -> Just $ SEZoom (sqrt 0.5)
    EvKey (KChar '>') []      -> Just $ SERate (sqrt 2)
    EvKey (KChar '<') []      -> Just $ SERate (sqrt (1/2))
    EvKey (KChar ']') []      -> Just $ SEHist 5
    EvKey (KChar '[') []      -> Just $ SEHist (-5)
    _                         -> Nothing

data RangeRatio = RR { -- | Where on the screen (0 to 1) to place the other axis
                       rrZero  :: Double
                       -- | Ratio of height of a terminal character to width
                     , rrRatio :: Double
                     }
                deriving (Show)

data PlotRange = PXY (Double, Double) (Double, Double)
               | PX  (Double, Double) RangeRatio
               | PY  RangeRatio       (Double, Double)

plot
    :: (Int, Int)               -- ^ display bounds
    -> PlotRange
    -> [(V2 Double, (Char, Attr))]   -- ^ points to plot
    -> [Image]
plot (wd,ht) pr = map (crop wd ht)
                . (++ bgs)
                . map (\(p, (c, a)) -> place EQ EQ p $ char a c)
  where
    wd' = fromIntegral wd
    ht' = fromIntegral ht
    ((xmin, xmax), (ymin, ymax)) = mkRange (wd', ht') pr
    origin = place EQ EQ (V2 0 0) $ char defAttr '+'
    xaxis  = place EQ EQ (V2 0 0) $ charFill defAttr '-' wd 1
    yaxis  = place EQ EQ (V2 0 0) $ charFill defAttr '|' 1 ht
    xrange = xmax - xmin
    yrange = ymax - ymin
    bg     = backgroundFill wd ht
    scale (V2 pX pY) = V2 x y
      where
        x = round $ (pX - xmin) * (wd' / xrange)
        y = round $ (pY - ymin) * (ht' / yrange)
    place aX aY p i = case scale p of
      V2 pX pY -> translate (fAlign aX (imageWidth  i))
                            (fAlign aY (imageHeight i))
                . translate pX pY
                $ i
    labels = [ place LT EQ (V2 xmin 0) . string defAttr $ printf "%.2f" xmin
             , place GT EQ (V2 xmax 0) . string defAttr $ printf "%.2f" xmax
             , place EQ LT (V2 0 ymin) . string defAttr $ printf "%.2f" ymin
             , place EQ GT (V2 0 ymax) . string defAttr $ printf "%.2f" ymax
             ]
    bgs    = labels ++ [origin, xaxis, yaxis, bg]
    fAlign = \case
      LT -> const 0
      EQ -> negate . (`div` 2)
      GT -> negate

mkRange
    :: (Double, Double)
    -> PlotRange
    -> ((Double, Double), (Double, Double))
mkRange (wd, ht) = \case
    PXY xb     yb     -> (xb, yb)
    PX  xb     RR{..} ->
      let yr = (uncurry (-) xb) * ht / wd * rrRatio
          y0 = (rrZero - 1) * yr
      in  (xb, (y0, y0 + yr))
    PY  RR{..} yb ->
      let xr = (uncurry (-) yb) * wd / ht / rrRatio
          x0 = (rrZero - 1) * xr
      in  ((x0, x0 + xr), yb)

pattern V1 :: a -> V.Vector 1 a
pattern V1 x <- (V.head->x)
  where
    V1 x = V.singleton x
#if __GLASGOW_HASKELL__ >= 802
{-# COMPLETE V1 #-}
#endif

type V2 = V.Vector 2
pattern V2 :: a -> a -> V2 a
pattern V2 x y <- (V.toList->[x,y])
  where
    V2 x y = V.fromTuple (x, y)
#if __GLASGOW_HASKELL__ >= 802
{-# COMPLETE V2 #-}
#endif

pattern V3 :: a -> a -> a -> V.Vector 3 a
pattern V3 x y z <- (V.toList->[x,y,z])
  where
    V3 x y z = V.fromTuple (x, y, z)
#if __GLASGOW_HASKELL__ >= 802
{-# COMPLETE V3 #-}
#endif

pattern V4 :: a -> a -> a -> a -> V.Vector 4 a
pattern V4 x y z a <- (V.toList->[x,y,z,a])
  where
    V4 x y z a = V.fromTuple (x, y, z, a)
#if __GLASGOW_HASKELL__ >= 802
{-# COMPLETE V4 #-}
#endif

logistic
    :: Floating a => a -> a -> a -> a -> a
logistic pos ht width = \x -> ht / (1 + exp (- beta * (x - pos)))
  where
    beta = log (0.9 / (1 - 0.9)) / width

bezierCurve
    :: forall n f a. (KnownNat (1 + n), Applicative f, Num a)
    => V.Vector (1 + n) (f a)
    -> a
    -> f a
bezierCurve ps t =
      foldl' (liftA2 (+)) (pure 0)
    . V.imap (\i -> let i' = fromIntegral i
                    in  fmap (* (fromIntegral (n' `choose` i') * (1 - t)^(n' - i') * t^i))
             )
    $ ps
  where
    n' :: Int
    n' = fromIntegral (maxBound :: Finite (1 + n))
    choose :: Int -> Int -> Int
    n `choose` k = factorial n `div` (factorial (n - k) * factorial k)
    factorial :: Int -> Int
    factorial m = product [1..m]

deriving instance Ord Color


{-# OPTIONS_GHC -fno-warn-unticked-promoted-constructors #-}

module Noether.Algebra.Linear.Module where

import           Data.Complex

import           Noether.Lemmata.Prelude
import           Noether.Lemmata.TypeFu

import           Noether.Algebra.Actions
import           Noether.Algebra.Multiple
import           Noether.Algebra.Single

import           Noether.Algebra.Tags

data LeftModuleE
  = LeftModule_Ring_AbelianGroup_Linear_Compatible
    { leftModule_ring           :: RingE
    , leftModule_abelianGroup   :: AbelianGroupE
    , leftModule_actorLinearity :: ActorLinearE
    , leftModule_acteeLinearity :: ActeeLinearE
    , leftModule_compatibility  :: CompatibleE}
  | LeftModule_Named Symbol
                     LeftModuleE

-- | A left module (v, a) over the ring (r, p, m).
class LeftModuleK op p m r a v s

type family LeftModuleS (op :: k0) (p :: k1) (m :: k2) r (a :: k3) v :: LeftModuleE

type LeftModuleC op p m r a v = LeftModuleK op p m r a v (LeftModuleS op p m r a v)

instance ( RingK p m r zr
         , AbelianGroupK a v zag
         , ActorLinearK L m p r a v zor
         , ActeeLinearK L m r a v zee
         , CompatibleK L op m r v zlc
         ) =>
         LeftModuleK op p m r a v
           (LeftModule_Ring_AbelianGroup_Linear_Compatible zr zag zor zee zlc)

type LeftModule op p m r a v =
  ( LeftModuleC op p m r a v
  , Ring p m r
  , AbelianGroup a v
  , LinearActsOn L m p r a v
  , LeftCompatible op m r v
  )

data RightModuleE
  = RightModule_Ring_AbelianGroup_Linear_Compatible
    { rightModule_ring           :: RingE
    , rightModule_abelianGroup   :: AbelianGroupE
    , rightModule_actorLinearity :: ActorLinearE
    , rightModule_acteeLinearity :: ActeeLinearE
    , rightModule_compatibility  :: CompatibleE}
  | RightModule_Named Symbol
                     RightModuleE

-- | A right module (v, a) over the ring (r, p, m).
class RightModuleK op p m r a v s

type family RightModuleS (op :: k0) (p :: k1) (m :: k2) r (a :: k3) v :: RightModuleE

type RightModuleC op p m r a v = RightModuleK op p m r a v (RightModuleS op p m r a v)

instance ( RingK p m r zr
         , AbelianGroupK a v zag
         , ActorLinearK 'R m p r a v zor
         , ActeeLinearK 'R m r a v zee
         , CompatibleK 'R op m r v zrc
         ) => RightModuleK op p m r a v
           (RightModule_Ring_AbelianGroup_Linear_Compatible zr zag zor zee zrc)

type RightModule op p m r a v =
  ( RightModuleC op p m r a v
  , Ring p m r
  , AbelianGroup a v
  , LinearActsOn 'R m p r a v
  , RightCompatible op m r v
  )


{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE RankNTypes          #-}
-- |
-- Module      : Bench.HMatrix
-- Copyright   : [2017..2020] Trevor L. McDonell
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--

module Bench.HMatrix ( bench_hmatrix )
  where

import Bench.Util

import Control.DeepSeq
import Criterion.Main
import Data.Proxy
import Foreign.Storable
import Numeric.LinearAlgebra                                        hiding ( randomVector )
import Prelude                                                      hiding ( (<>) )
import System.Random.MWC
import Text.Printf

bench_hmatrix :: [Benchmark]
bench_hmatrix =
  [ bench_level2
  , bench_level3
  ]

bench_level2 :: Benchmark
bench_level2 =
  bgroup "matrix-vector"
    [ bgroup "(#>)"
      [ gemv  200  400
      , gemv  500 1000
      , gemv 1000 2000
      , gemv 2000 3000
      ]
    , bgroup "(<#)"
      [ gevm  200  400
      , gevm  500 1000
      , gevm 1000 2000
      , gevm 2000 3000
      ]
    ]
  where
    gemv :: Int -> Int -> Benchmark
    gemv m n =
      let setup :: (Variate e, Storable e) => proxy e -> IO (Matrix e, Vector e)
          setup _ = withSystemRandom $ \gen -> do
            matA <- randomMatrix gen m n
            vecx <- randomVector gen n
            return (matA, vecx)

          go :: (Variate e, Numeric e, NFData e, Show (ArgType e)) => proxy e -> Benchmark
          go t = env (setup t)
               $ \ ~(matA, vecx) -> bench (showType t)
               $ whnf (matA #>) vecx
      in
      bgroup (printf "%dx%d" m n) (sdcz go)

    gevm :: Int -> Int -> Benchmark
    gevm m n =
      let setup :: (Variate e, Storable e) => proxy e -> IO (Matrix e, Vector e)
          setup _ = withSystemRandom $ \gen -> do
            matA <- randomMatrix gen m n
            vecx <- randomVector gen m
            return (matA, vecx)

          go :: (Variate e, Numeric e, NFData e, Show (ArgType e)) => proxy e -> Benchmark
          go t = env (setup t)
               $ \ ~(matA, vecx) -> bench (showType t)
               $ whnf (vecx <#) matA
      in
      bgroup (printf "%dx%d" m n) (sdcz go)

bench_level3 :: Benchmark
bench_level3 =
  bgroup "matrix-matrix"
    [ bgroup "(<>)"
      [ gemm  100  100  100
      , gemm  250  250  250
      , gemm  500  500  500
      , gemm 1000 1000 1000
      ]
    ]
  where
    gemm :: Int -> Int -> Int -> Benchmark
    gemm m n k =
      let
          setup :: (Variate e, Storable e) => proxy e -> IO (Matrix e, Matrix e)
          setup _ = withSystemRandom $ \gen -> do
            matA <- randomMatrix gen m k
            matB <- randomMatrix gen k n
            return (matA, matB)

          go :: (Variate e, Numeric e, NFData e, Show (ArgType e)) => proxy e -> Benchmark
          go t = env (setup t)
               $ \ ~(matA, matB) -> bench (showType t)
               $ whnf (matA <>) matB
      in
      bgroup (printf "%dx%dx%d" m n k) (sdcz go)

randomVector :: (Variate e, Storable e) => GenIO -> Int -> IO (Vector e)
randomVector = uniformVector

randomMatrix :: (Variate e, Storable e) => GenIO -> Int -> Int -> IO (Matrix e)
randomMatrix gen m n = do
  v <- uniformVector gen (m * n)
  return $ reshape n v

sdcz :: (forall (e :: *). (Variate e, Numeric e, NFData e, Show (ArgType e)) => Proxy e -> Benchmark)
     -> [Benchmark]
sdcz go =
  [ go (Proxy :: Proxy Float)
  , go (Proxy :: Proxy Double)
  , go (Proxy :: Proxy (Complex Float))
  , go (Proxy :: Proxy (Complex Double))
  ]


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Day12 where

import Control.Lens
import Control.Lens.TH

import qualified Data.ByteString as BS
import Data.Attoparsec.ByteString (Parser)
import qualified Data.Attoparsec.ByteString as DAB
import qualified Data.Attoparsec.ByteString.Char8 as DABC8
import Data.Attoparsec.ByteString.Char8 (decimal, inClass, digit, endOfLine, char, anyChar, letter_ascii, notChar)

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

import qualified Data.List as DL
import qualified Data.List.Split as DLS

import Debug.Trace

import Control.Arrow

import Data.Array.IArray (Array)
import qualified Data.Array.IArray as IA

import Data.Set (Set)
import qualified Data.Set as Set
import Data.Either (isRight)
import Data.Maybe (mapMaybe, catMaybes)
import Data.Bool (bool)

import Data.Vector (Vector, (!))
import qualified Data.Vector as Vec

import Data.Function (on)

import Control.Monad.State.Strict (State)
import qualified Control.Monad.State.Strict as State

import qualified Control.Monad as CM

-- import qualified Control.Foldl as L

import Data.Bifunctor (bimap)

import Data.Complex

--type Input = (IA.Array Int (IA.Array Int Char))
type Input = (Char, Int)

-- >>> solver [('F', 10),

solver input =
	uncurry (+) $
	(abs . realPart &&& abs . imagPart) $ snd $
	DL.foldl' (\x y -> traceShow (x, y) (navigate y x)) (10 :+ 1, 0 :+ 0) input
	where
	navigate :: Input -> (Complex Float, Complex Float) -> (Complex Float, Complex Float)
	navigate ('N', n) = first (+ imag n)
	navigate ('S', n) = first (dec (imag n))
	navigate ('E', n) = first (+ real n)
	navigate ('W', n) = first (dec (real n))

	navigate ('L', 90) = first (* (0 :+ 1))
	navigate ('L', 180) = first $ fmap round' . (* (0 :+ 1) ** 2)
	navigate ('L', 270) = first $ fmap round' . (* (0 :+ 1) ** 3)

	navigate ('R', 90) = first $ fmap round' . (* (0 :+ (-1)))
	navigate ('R', 180) = first $ fmap round' . (* (0 :+ (-1)) ** 2)
	navigate ('R', 270) = first $ fmap round' . (* (0 :+ (-1)) ** 3)
	navigate ('F', n) = uncurry (CM.liftM2 (.) (,) ((+) . (real n *)))

	navigate (c, n) = error $ show (c, n)

	norm :: Complex Float -> Complex Float
	norm x = x / abs x

	real n = fromIntegral n :+ 0
	imag n = 0 :+ fromIntegral n
	dec n x = x - n

	round' :: Float -> Float
	round' = fromIntegral . round

parseInput :: Parser [Input]
parseInput =
	((,) <$> anyChar <*> decimal) `DAB.sepBy1'` endOfLine

runSolution :: FilePath -> IO ()
runSolution filePath = do
	contents <- BS.readFile filePath
	let parseResult = DAB.parseOnly parseInput contents
	case parseResult of
		Left err -> putStrLn err
		Right input -> do
			print $ solver input

-- 1526 wrong
-- part ii
-- 1463 wrong


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE CPP #-}

{- |
Module      :  Physics.Learn.BeamStack
Copyright   :  (c) Scott N. Walck 2016-2018
License     :  BSD3 (see LICENSE)
Maintainer  :  Scott N. Walck <walck@lvc.edu>
Stability   :  experimental

Splitters, recombiners, and detectors for Stern-Gerlach
experiments.
-}

-- Spin-1/2 mixed states.

module Physics.Learn.BeamStack
    (
    -- * Core laboratory components
      BeamStack()
    , randomBeam
    , split
    , recombine
    , applyBField
    , dropBeam
    , flipBeams
    , numBeams
    , detect
    -- * Standard splitters
    , splitX
    , splitY
    , splitZ
    -- * Standard magnetic fields
    , applyBFieldX
    , applyBFieldY
    , applyBFieldZ
    -- * Standard combiners
    , recombineX
    , recombineY
    , recombineZ
    -- * Filters
    , xpFilter
    , xmFilter
    , zpFilter
    , zmFilter
    )
    where

import Physics.Learn.QuantumMat
    ( zp
    , zm
    , nm
    , np
    , couter
    , oneQubitMixed
    )
import Numeric.LinearAlgebra
    ( C
    , Vector
    , Matrix
    , iC
    , (<>)
    , kronecker
    , fromLists
    , toList
    , toLists
    , scale
    , size
    , takeDiag
    , ident
    , tr
    )
import Data.Complex
    ( Complex(..)
    , realPart
    , imagPart
    )
import Data.List
    ( intercalate
    )
#if MIN_VERSION_base(4,11,0)
import Prelude hiding ((<>))
#endif

data BeamStack = BeamStack (Matrix C)

showOneBeam :: Double -> String
showOneBeam r = "Beam of intensity " ++ show r

instance Show BeamStack where
    show b = intercalate "\n" $ map showOneBeam (detect b)

{-
unBeamStack :: BeamStack -> Matrix C
unBeamStack (BeamStack m) = m
-}

--------------------
-- Core functions --
--------------------

-- | A beam of randomly oriented spin-1/2 particles.
randomBeam :: BeamStack
randomBeam = BeamStack oneQubitMixed

extendWithZeros :: Matrix C -> Matrix C
extendWithZeros m
    = let (_,q) = size m
          ml = toLists m
      in fromLists $ map (++ [0,0]) ml
             ++ [replicate (q+2) 0, replicate (q+2) 0]

-- reduce row and column size by 2
reduceMat :: Matrix C -> Matrix C
reduceMat m
    = let (p,q) = size m
          ml = toLists m
      in fromLists $ take (p-2) $ map (take (q-2)) ml

checkedRealPart :: C -> Double
checkedRealPart c
    = let eps = 1e-14
      in if imagPart c < eps
         then realPart c
         else error $ "checkRealPart: imagPart = " ++ show (imagPart c)

-- | Return the intensities of a stack of beams.
detect :: BeamStack -> [Double]
detect (BeamStack m)
    = addAlternate $ toList $ takeDiag m

addAlternate :: [C] -> [Double]
addAlternate [] = []
addAlternate [_] = error "addAlternate needs even number of elements"
addAlternate (x:y:xs) = checkedRealPart (x+y) : addAlternate xs

-- | Remove the most recent beam from the stack.
dropBeam :: BeamStack -> BeamStack
dropBeam (BeamStack m) = BeamStack (reduceMat m)

-- | Return the number of beams in a 'BeamStack'.
numBeams :: BeamStack -> Int
numBeams (BeamStack m)
    = let (p,_) = size m
      in p `div` 2

-- | Interchange the two most recent beams on the stack.
flipBeams :: BeamStack -> BeamStack
flipBeams (BeamStack m)
    = let (d,_) = size m
          fl = flipMat d
      in BeamStack $ fl <> m <> tr fl

flipMat :: Int -> Matrix C
flipMat d = bigM d (fromLists [[0,0,1,0]
                              ,[0,0,0,1]
                              ,[1,0,0,0]
                              ,[0,1,0,0]])

-- Turn a 2x2 into a dxd.
bigM2 :: Int -> Matrix C -> Matrix C
bigM2 d m
    | d < 2      = error "bigM2 requires d >= 2"
    | odd d      = error "bigM2 requires even d"
    | otherwise  = fromLists $ map (++ [0,0]) (toLists (ident (d-2)))
                   ++ map (replicate (d-2) 0 ++) (toLists m)

-- Turn a 4x4 into a dxd.
bigM :: Int -> Matrix C -> Matrix C
bigM d m
    | d < 4      = error "bigM requires d >= 4"
    | odd d      = error "bigM requires even d"
    | otherwise  = fromLists $ map (++ [0,0,0,0]) (toLists (ident (d-4)))
                   ++ map (replicate (d-4) 0 ++) (toLists m)

s :: Double -> Double -> Matrix C
s theta phi = kronecker (u `couter` u) (np theta phi `couter` np theta phi)
            + kronecker (l `couter` u) (nm theta phi `couter` nm theta phi)
            + kronecker (u `couter` l) (nm theta phi `couter` nm theta phi)
            + kronecker (l `couter` l) (np theta phi `couter` np theta phi)

u :: Vector C
u = zp

l :: Vector C
l = zm

-- | Given angles describing the orientation of the splitter,
--   removes an incoming beam from the stack and replaces
--   it with two beams, a spin-up and a spin-down beam.
--   The spin-down beam is the most recent beam on the stack.
split :: Double -> Double -> BeamStack -> BeamStack
split theta phi (BeamStack m)
    = let m' = extendWithZeros m
          (p,_) = size m'
          ss = bigM p (s theta phi)
      in BeamStack $ ss <> m' <> tr ss

-- | Given angles describing the orientation of the recombiner,
--   returns a single beam from an incoming pair of beams.
recombine :: Double -> Double -> BeamStack -> BeamStack
recombine theta phi (BeamStack m)
    = let (d,_) = size m
          ss = bigM d (s theta phi)
      in dropBeam $ BeamStack $ ss <> m <> tr ss

mag2x2 :: Double -> Double -> Double -> Matrix C
mag2x2 theta phi omegaT
    = let z = iC * (omegaT :+ 0) / 2
          np' = np theta phi
          nm' = nm theta phi
      in scale (exp   z ) (np' `couter` np')
       + scale (exp (-z)) (nm' `couter` nm')

-- | Given angles describing the direction of a
--   uniform magnetic field, and given an angle
--   describing the product of the Larmor frequency
--   and the time, return an output beam from an
--   input beam.
applyBField :: Double -> Double -> Double -> BeamStack -> BeamStack
applyBField theta phi omegaT (BeamStack m)
    = let (d,_) = size m
          uu = bigM2 d (mag2x2 theta phi omegaT)
      in BeamStack $ uu <> m <> tr uu

-----------------------
-- Derived functions --
-----------------------

-- | A Stern-Gerlach splitter in the x direction.
splitX :: BeamStack -> BeamStack
splitX = split (pi/2) 0

-- | A Stern-Gerlach splitter in the y direction.
splitY :: BeamStack -> BeamStack
splitY = split (pi/2) (pi/2)

-- | A Stern-Gerlach splitter in the z direction.
splitZ :: BeamStack -> BeamStack
splitZ = split 0 0

-- | Given an angle in radians
--   describing the product of the Larmor frequency
--   and the time, apply a magnetic in the x direction
--   to the most recent beam on the stack.
applyBFieldX :: Double -> BeamStack -> BeamStack
applyBFieldX = applyBField (pi/2) 0

-- | Given an angle in radians
--   describing the product of the Larmor frequency
--   and the time, apply a magnetic in the y direction
--   to the most recent beam on the stack.
applyBFieldY :: Double -> BeamStack -> BeamStack
applyBFieldY = applyBField (pi/2) (pi/2)

-- | Given an angle in radians
--   describing the product of the Larmor frequency
--   and the time, apply a magnetic in the z direction
--   to the most recent beam on the stack.
applyBFieldZ :: Double -> BeamStack -> BeamStack
applyBFieldZ = applyBField 0 0

-- | A Stern-Gerlach recombiner in the x direction.
recombineX :: BeamStack -> BeamStack
recombineX = recombine (pi/2) 0

-- | A Stern-Gerlach recombiner in the y direction.
recombineY :: BeamStack -> BeamStack
recombineY = recombine (pi/2) (pi/2)

-- | A Stern-Gerlach recombiner in the z direction.
recombineZ :: BeamStack -> BeamStack
recombineZ = recombine 0 0

-- | Filter for spin-up particles in the x direction.
xpFilter :: BeamStack -> BeamStack
xpFilter = dropBeam . splitX

-- | Filter for spin-down particles in the x direction.
xmFilter :: BeamStack -> BeamStack
xmFilter = dropBeam . flipBeams . splitX

-- | Filter for spin-up particles in the z direction.
zpFilter :: BeamStack -> BeamStack
zpFilter = dropBeam . splitZ

-- | Filter for spin-down particles in the z direction.
zmFilter :: BeamStack -> BeamStack
zmFilter = dropBeam . flipBeams . splitZ


module Vec where
import Linear.Epsilon (nearZero)
import qualified Data.Map.Strict as Map
import Data.Complex

-- http://blog.sigfpe.com/2007/03/monads-vector-spaces-and-quantum.html

data W b a = W { runW :: [(a,b)] } deriving (Eq,Show,Ord)
instance Semigroup (W b a) where
  (W x) <> (W y) = W (x <> y)
instance Monoid (W b a) where
  mempty = W mempty
mapW f (W l) = W $ map (\(a,b) -> (a,f b)) l
instance Functor (W b) where
    fmap f (W a) = W $ map (\(a,p) -> (f a,p)) a

instance Num b => Applicative (W b) where
  pure x = W [(x,1)]
  (W fs) <*> (W xs) = W [(f x, a * b) | (f, a) <- fs, (x, b) <- xs]
  -- the fs is a a vector of rows kind of.
    -- error "I don't feel like it"

instance Num b => Monad (W b) where
   return x = W [(x,1)]
   l >>= f = W $ concatMap (\(W d,p) -> map (\(x,q)->(x,p*q)) d) (runW $ fmap f l)

a .* b = mapW (a*) b

instance (Eq a,Show a,Num b) => Num (W b a) where
     W a + W b = W $ (a ++ b)
     a - b = a + (-1) .* b
     _ * _ = error "Num is annoying"
     abs _ = error "Num is annoying"
     signum _ = error "Num is annoying"
     fromInteger a = if a==0 then W [] else error "fromInteger can only take zero argument"

collect :: (Ord a, Num b) => W b a -> W b a
collect = W . Map.toList . Map.fromListWith (+) . runW

trimZero = W . filter (\(k,v) -> not $ nearZero v) . runW
simplify :: Ord a => Q a -> Q a
simplify = trimZero . collect
-- filter (not . nearZero . snd)

type P a = W Double a

type Q a = W (Complex Double) a

star :: Num b => W (Complex b) a -> W (Complex b) a
star = mapW conjugate

kron :: Num b => W b a -> W b c -> W b (a,c)
kron (W x) (W y) = W [((a,c), r1 * r2) | (a,r1) <- x , (c,r2) <- y ]


{-# LANGUAGE AllowAmbiguousTypes, DataKinds, FlexibleInstances,
             FunctionalDependencies, QuantifiedConstraints, RankNTypes,
             TypeFamilies, TypeOperators #-}
{-# OPTIONS_GHC -fconstraint-solver-iterations=16 #-}
{-# OPTIONS_GHC -Wno-missing-methods #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
-- | Dual numbers and various operations on them, arithmetic and related
-- to tensors (vectors, matrices and others). This is the high-level API,
-- defined using the low-level API in "HordeAd.Core.DualClass".
module HordeAd.Core.DualNumber
  ( module HordeAd.Core.DualNumber
  , IsScalar, HasDelta, DMode(..)
  , Domain0, Domain1, Domain2, DomainX, Domains  -- an important re-export
  ) where

import Prelude

import qualified Data.Array.Convert
import qualified Data.Array.Dynamic as OTB
import qualified Data.Array.DynamicS as OT
import           Data.Array.Internal (valueOf)
import           Data.Array.Shape (DivRoundUp)
import qualified Data.Array.Shaped as OSB
import qualified Data.Array.ShapedS as OS
import           Data.List.Index (imap)
import           Data.MonoTraversable (MonoFunctor (omap))
import           Data.Proxy (Proxy (Proxy))
import qualified Data.Strict.Vector as Data.Vector
import qualified Data.Vector.Generic as V
import           GHC.TypeLits (KnownNat, type (+), type (-), type (<=))
import           Numeric.LinearAlgebra (Matrix, Numeric, Vector)
import qualified Numeric.LinearAlgebra as HM

import HordeAd.Core.DualClass
import HordeAd.Internal.Delta
  (CodeOut (..), Domain0, Domain1, Domain2, DomainX, Domains)

-- * The main dual number types

-- | Dual numbers with the second type argument being the primal component.
data DualNumber (d :: DMode) a = D a (Dual d a)

class (IsScalar d r, Monad m, Functor m, Applicative m)
      => DualMonad d r m | m -> d r where
  returnLet :: IsPrimalWithScalar d a r
            => DualNumber d a -> m (DualNumber d a)

addParameters :: (Numeric r, Num (Vector r))
              => Domains r -> Domains r -> Domains r
addParameters (a0, a1, a2, aX) (b0, b1, b2, bX) =
  (a0 + b0, V.zipWith (+) a1 b1, V.zipWith (+) a2 b2, V.zipWith (+) aX bX)

-- Dot product and sum respective ranks and sum it all.
dotParameters :: Numeric r => Domains r -> Domains r -> r
dotParameters (a0, a1, a2, aX) (b0, b1, b2, bX) =
  a0 HM.<.> b0
  + V.sum (V.zipWith (HM.<.>) a1 b1)
  + V.sum (V.zipWith (HM.<.>) (V.map HM.flatten a2) (V.map HM.flatten b2))
  + V.sum (V.zipWith (HM.<.>) (V.map OT.toVector aX) (V.map OT.toVector bX))

-- * General operations, for any tensor rank

-- These instances are required by the @Real@ instance, which is required
-- by @RealFloat@, which gives @atan2@. No idea what properties
-- @Real@ requires here, so let it crash if it's really needed.
instance Eq (DualNumber d a) where

instance Ord (DualNumber d a) where

-- These instances are dangerous due to possible subexpression copies
-- leading to allocation explosion. Expressions should be wrapped in
-- the monadic @returnLet@ whenever there is a possibility they can be
-- used multiple times in a larger expression.
instance (Num a, IsPrimal d a) => Num (DualNumber d a) where
  D u u' + D v v' = D (u + v) (dAdd u' v')
  D u u' - D v v' = D (u - v) (dAdd u' (dScale (-1) v'))
  D u u' * D v v' = D (u * v) (dAdd (dScale v u') (dScale u v'))
  negate (D v v') = D (negate v) (dScale (-1) v')
  abs (D v v') = D (abs v) (dScale (signum v) v')
  signum (D v _) = D (signum v) dZero
  fromInteger = constant . fromInteger

instance (Real a, IsPrimal d a) => Real (DualNumber d a) where
  toRational = undefined  -- TODO?

instance (Fractional a, IsPrimal d a) => Fractional (DualNumber d a) where
  D u u' / D v v' =
    let recipSq = recip (v * v)
    in D (u / v) (dAdd (dScale (v * recipSq) u') (dScale (- u * recipSq) v'))
  recip (D v v') =
    let minusRecipSq = - recip (v * v)
    in D (recip v) (dScale minusRecipSq v')
  fromRational = constant . fromRational

instance (Floating a, IsPrimal d a) => Floating (DualNumber d a) where
  pi = constant pi
  exp (D u u') = let expU = exp u
                 in D expU (dScale expU u')
  log (D u u') = D (log u) (dScale (recip u) u')
  sqrt (D u u') = let sqrtU = sqrt u
                  in D sqrtU (dScale (recip (sqrtU + sqrtU)) u')
  D u u' ** D v v' = D (u ** v) (dAdd (dScale (v * (u ** (v - 1))) u')
                                      (dScale ((u ** v) * log u) v'))
  logBase x y = log y / log x
  sin (D u u') = D (sin u) (dScale (cos u) u')
  cos (D u u') = D (cos u) (dScale (- (sin u)) u')
  tan (D u u') = let cosU = cos u
                 in D (tan u) (dScale (recip (cosU * cosU)) u')
  asin (D u u') = D (asin u) (dScale (recip (sqrt (1 - u*u))) u')
  acos (D u u') = D (acos u) (dScale (- recip (sqrt (1 - u*u))) u')
  atan (D u u') = D (atan u) (dScale (recip (1 + u*u)) u')
  sinh (D u u') = D (sinh u) (dScale (cosh u) u')
  cosh (D u u') = D (cosh u) (dScale (sinh u) u')
  tanh (D u u') = let y = tanh u
                  in D y (dScale (1 - y * y) u')
  asinh (D u u') = D (asinh u) (dScale (recip (sqrt (1 + u*u))) u')
  acosh (D u u') = D (acosh u) (dScale (recip (sqrt (u*u - 1))) u')
  atanh (D u u') = D (atanh u) (dScale (recip (1 - u*u)) u')

instance (RealFrac a, IsPrimal d a) => RealFrac (DualNumber d a) where
  properFraction = undefined
    -- very low priority, since these are all extremely not continuous

instance (RealFloat a, IsPrimal d a) => RealFloat (DualNumber d a) where
  atan2 (D u u') (D v v') =
    let t = 1 / (u * u + v * v)
    in D (atan2 u v) (dAdd (dScale (- u * t) v') (dScale (v * t) u'))
      -- we can be selective here and omit the other methods,
      -- most of which don't even have a differentiable codomain

constant :: IsPrimal d a => a -> DualNumber d a
constant a = D a dZero

scale :: (Num a, IsPrimal d a) => a -> DualNumber d a -> DualNumber d a
scale a (D u u') = D (a * u) (dScale a u')

tanhAct :: (DualMonad d r m, IsPrimalAndHasFeatures d a r)
        => DualNumber d a -> m (DualNumber d a)
tanhAct = returnLet . tanh

logistic :: (Floating a, IsPrimal d a) => DualNumber d a -> DualNumber d a
logistic (D u u') =
  let y = recip (1 + exp (- u))
  in D y (dScale (y * (1 - y)) u')

logisticAct :: (DualMonad d r m, IsPrimalAndHasFeatures d a r)
            => DualNumber d a -> m (DualNumber d a)
logisticAct = returnLet . logistic

-- Optimized and more clearly written @u ** 2@.
square :: (Num a, IsPrimal d a) => DualNumber d a -> DualNumber d a
square (D u u') = D (u * u) (dScale (2 * u) u')

squaredDifference :: (Num a, IsPrimal d a)
                  => a -> DualNumber d a -> DualNumber d a
squaredDifference targ res = square $ res - constant targ

lossSquared :: (DualMonad d r m, IsPrimalAndHasFeatures d a r)
            => a -> DualNumber d a -> m (DualNumber d a)
lossSquared targ res = returnLet $ squaredDifference targ res

reluAct :: (DualMonad d r m, IsPrimalAndHasFeatures d a r)
        => DualNumber d a -> m (DualNumber d a)
reluAct v@(D u _) = do
  let oneIfGtZero = omap (\x -> if x > 0 then 1 else 0) u
  returnLet $ scale oneIfGtZero v

reluLeakyAct :: (DualMonad d r m, IsPrimalAndHasFeatures d a r)
             => DualNumber d a -> m (DualNumber d a)
reluLeakyAct v@(D u _) = do
  let oneIfGtZero = omap (\x -> if x > 0 then 1 else 0.01) u
  returnLet $ scale oneIfGtZero v

-- * Operations resulting in a scalar

sumElements0 :: IsScalar d r => DualNumber d (Vector r) -> DualNumber d r
sumElements0 (D u u') = D (HM.sumElements u) (dSumElements0 u' (V.length u))

index0 :: IsScalar d r => DualNumber d (Vector r) -> Int -> DualNumber d r
index0 (D u u') ix = D (u V.! ix) (dIndex0 u' ix (V.length u))

minimum0 :: IsScalar d r => DualNumber d (Vector r) -> DualNumber d r
minimum0 (D u u') =
  D (HM.minElement u) (dIndex0 u' (HM.minIndex u) (V.length u))

maximum0 :: IsScalar d r => DualNumber d (Vector r) -> DualNumber d r
maximum0 (D u u') =
  D (HM.maxElement u) (dIndex0 u' (HM.maxIndex u) (V.length u))

-- If @v'@ is a @Var1@, this is much faster due to the optimization
-- in @Index0@.
foldl'0 :: IsScalar d r
        => (DualNumber d r -> DualNumber d r -> DualNumber d r)
        -> DualNumber d r -> DualNumber d (Vector r)
        -> DualNumber d r
foldl'0 f uu' (D v v') =
  let k = V.length v
      g !acc ix p = f (D p (dIndex0 v' ix k)) acc
  in V.ifoldl' g uu' v

altSumElements0 :: IsScalar d r => DualNumber d (Vector r) -> DualNumber d r
altSumElements0 = foldl'0 (+) 0

-- | Dot product.
infixr 8 <.>!
(<.>!) :: IsScalar d r
       => DualNumber d (Vector r) -> DualNumber d (Vector r) -> DualNumber d r
(<.>!) (D u u') (D v v') = D (u HM.<.> v) (dAdd (dDot0 v u') (dDot0 u v'))

-- | Dot product with a constant vector.
infixr 8 <.>!!
(<.>!!) :: IsScalar d r
        => DualNumber d (Vector r) -> Vector r -> DualNumber d r
(<.>!!) (D u u') v = D (u HM.<.> v) (dDot0 v u')

infixr 8 <.>$
(<.>$) :: (IsScalar d r, KnownNat n)
       => DualNumber d (OS.Array '[n] r) -> DualNumber d (OS.Array '[n] r)
       -> DualNumber d r
(<.>$) d e = fromS1 d <.>! fromS1 e

fromX0 :: IsScalar d r => DualNumber d (OT.Array r) -> DualNumber d r
fromX0 (D u u') = D (OT.unScalar u) (dFromX0 u')

fromS0 :: IsScalar d r => DualNumber d (OS.Array '[] r) -> DualNumber d r
fromS0 (D u u') = D (OS.unScalar u) (dFromS0 u')

sumElementsVectorOfDual
  :: IsScalar d r => Data.Vector.Vector (DualNumber d r) -> DualNumber d r
sumElementsVectorOfDual = V.foldl' (+) 0

softMaxAct :: DualMonad d r m
           => Data.Vector.Vector (DualNumber d r)
           -> m (Data.Vector.Vector (DualNumber d r))
softMaxAct us = do
  expUs <- V.mapM (returnLet . exp) us
  let sumExpUs = sumElementsVectorOfDual expUs
  -- This has to be let-bound, because it's used many times below.
  recipSum <- returnLet $ recip sumExpUs
  V.mapM (\r -> returnLet $ r * recipSum) expUs

-- In terms of hmatrix: @-(log res <.> targ)@.
lossCrossEntropy :: forall d r m. DualMonad d r m
                 => Vector r
                 -> Data.Vector.Vector (DualNumber d r)
                 -> m (DualNumber d r)
lossCrossEntropy targ res = do
  let f :: DualNumber d r -> Int -> DualNumber d r -> DualNumber d r
      f !acc i d = acc + scale (targ V.! i) (log d)
  returnLet $ negate $ V.ifoldl' f 0 res

-- In terms of hmatrix: @-(log res <.> targ)@.
lossCrossEntropyV :: DualMonad d r m
                  => Vector r
                  -> DualNumber d (Vector r)
                  -> m (DualNumber d r)
lossCrossEntropyV targ res = returnLet $ negate $ log res <.>!! targ

-- Note that this is equivalent to a composition of softMax and cross entropy
-- only when @target@ is one-hot. Otherwise, results vary wildly. In our
-- rendering of the MNIST data all labels are on-hot.
lossSoftMaxCrossEntropyV
  :: DualMonad d r m
  => Vector r -> DualNumber d (Vector r) -> m (DualNumber d r)
lossSoftMaxCrossEntropyV target (D u u') = do
  -- The following protects from underflows, overflows and exploding gradients
  -- and is required by the QuickCheck test in TestMnistCNN.
  -- See https://github.com/tensorflow/tensorflow/blob/5a566a7701381a5cf7f70fce397759483764e482/tensorflow/core/kernels/sparse_softmax_op.cc#L106
  -- and https://github.com/tensorflow/tensorflow/blob/5a566a7701381a5cf7f70fce397759483764e482/tensorflow/core/kernels/xent_op.h
  let expU = exp (u - HM.scalar (HM.maxElement u))
      sumExpU = HM.sumElements expU
      recipSum = recip sumExpU
-- not exposed: softMaxU = HM.scaleRecip sumExpU expU
      softMaxU = HM.scale recipSum expU
  returnLet $ D (negate $ log softMaxU HM.<.> target)  -- TODO: avoid: log . exp
                (dDot0 (softMaxU - target) u')

-- * Operations resulting in a vector

-- @1@ means rank one, so the dual component represents a vector.
seq1 :: IsScalar d r
     => Data.Vector.Vector (DualNumber d r) -> DualNumber d (Vector r)
seq1 v = D (V.convert $ V.map (\(D u _) -> u) v)  -- I hope this fuses
           (dSeq1 $ V.map (\(D _ u') -> u') v)

konst1 :: IsScalar d r => DualNumber d r -> Int -> DualNumber d (Vector r)
konst1 (D u u') n = D (HM.konst u n) (dKonst1 u' n)

append1 :: IsScalar d r
        => DualNumber d (Vector r) -> DualNumber d (Vector r)
        -> DualNumber d (Vector r)
append1 (D u u') (D v v') = D (u V.++ v) (dAppend1 u' (V.length u) v')

slice1 :: IsScalar d r
       => Int -> Int -> DualNumber d (Vector r) -> DualNumber d (Vector r)
slice1 i n (D u u') = D (V.slice i n u) (dSlice1 i n u' (V.length u))

sumRows1 :: IsScalar d r => DualNumber d (Matrix r) -> DualNumber d (Vector r)
sumRows1 (D u u') = D (V.fromList $ map HM.sumElements $ HM.toRows u)
                      (dSumRows1 u' (HM.cols u))

sumColumns1 :: IsScalar d r => DualNumber d (Matrix r) -> DualNumber d (Vector r)
sumColumns1 (D u u') = D (V.fromList $ map HM.sumElements $ HM.toColumns u)
                         (dSumColumns1 u' (HM.rows u))

-- If @v'@ is a @Var1@, this is much faster due to the optimization
-- in @Index0@. The detour through a boxed vector (list probably fuses away)
-- is costly, but only matters if @f@ is cheap.
map1 :: IsScalar d r
     => (DualNumber d r -> DualNumber d r) -> DualNumber d (Vector r)
     -> DualNumber d (Vector r)
map1 f (D v v') =
  let k = V.length v
      g ix p = f $ D p (dIndex0 v' ix k)
      ds = imap g $ V.toList v
  in seq1 $ V.fromList ds

-- | Dense matrix-vector product.
infixr 8 #>!
(#>!) :: IsScalar d r
      => DualNumber d (Matrix r) -> DualNumber d (Vector r)
      -> DualNumber d (Vector r)
(#>!) (D u u') (D v v') = D (u HM.#> v) (dAdd (dMD_V1 u' v) (dM_VD1 u v'))

-- | Dense matrix-vector product with a constant vector.
infixr 8 #>!!
(#>!!) :: IsScalar d r
       => DualNumber d (Matrix r) -> Vector r
       -> DualNumber d (Vector r)
(#>!!) (D u u') v = D (u HM.#> v) (dMD_V1 u' v)

fromX1 :: IsScalar d r => DualNumber d (OT.Array r) -> DualNumber d (Vector r)
fromX1 (D u u') = D (OT.toVector u) (dFromX1 u')

fromS1 :: forall len d r. (KnownNat len, IsScalar d r)
       => DualNumber d (OS.Array '[len] r) -> DualNumber d (Vector r)
fromS1 (D u u') = D (OS.toVector u) (dFromS1 u')

reverse1 :: IsScalar d r => DualNumber d (Vector r) -> DualNumber d (Vector r)
reverse1 (D u u') = D (V.reverse u) (dReverse1 u')

flatten1 :: IsScalar d r => DualNumber d (Matrix r) -> DualNumber d (Vector r)
flatten1 (D u u') = let (rows, cols) = HM.size u
                    in D (HM.flatten u) (dFlatten1 rows cols u')

flattenX1 :: IsScalar d r => DualNumber d (OT.Array r) -> DualNumber d (Vector r)
flattenX1 (D u u') = let sh = OT.shapeL u
                     in D (OT.toVector u) (dFlattenX1 sh u')

flattenS1 :: (IsScalar d r, OS.Shape sh)
          => DualNumber d (OS.Array sh r) -> DualNumber d (Vector r)
flattenS1 (D u u') = D (OS.toVector u) (dFlattenS1 u')

corr1 :: IsScalar d r
      => DualNumber d (Vector r) -> DualNumber d (Vector r)
      -> DualNumber d (Vector r)
corr1 ker@(D u _) vv@(D v _) = case (V.length u, V.length v) of
  (0, lenV) -> konst1 0 lenV
  (lenK, lenV) -> if lenK <= lenV
                  then vectorSlices2 lenK vv #>! ker
                  else error $ "corr1: len kernel " ++ show lenK
                               ++ " > len vector " ++ show lenV

-- This is not optimally implemented: @append1@ is costly compared
-- to a @mconcat@ counterpart and @z@ is used twice without
-- assigning it to a variable.
conv1 :: IsScalar d r
      => DualNumber d (Vector r) -> DualNumber d (Vector r)
      -> DualNumber d (Vector r)
conv1 ker@(D u _) vv@(D v _) =
  let lenK = V.length u
      lenV = V.length v
      kerRev = reverse1 ker
      z = konst1 0 (lenK - 1)
      vvPadded = append1 z $ append1 vv z
  in if lenK == 0
     then konst1 0 lenV
     else corr1 kerRev vvPadded

-- No padding; remaining areas ignored.
maxPool1 :: IsScalar d r
         => Int -> Int -> DualNumber d (Vector r) -> DualNumber d (Vector r)
maxPool1 ksize stride v@(D u _) =
  let slices = [slice1 i ksize v | i <- [0, stride .. V.length u - ksize]]
  in seq1 $ V.fromList $ map maximum0 slices

softMaxActV :: DualMonad d r m
            => DualNumber d (Vector r) -> m (DualNumber d (Vector r))
softMaxActV d@(D u _) = do
  expU <- returnLet $ exp d
  let sumExpU = sumElements0 expU
  -- This has to be let-bound, because it's used many times below.
  recipSum <- returnLet $ recip sumExpU
  returnLet $ konst1 recipSum (V.length u) * expU

-- Note that this is equivalent to a composition of softMax and cross entropy
-- only when @target@ is one-hot. Otherwise, results vary wildly. In our
-- rendering of the MNIST data all labels are one-hot.
lossSoftMaxCrossEntropyL
  :: DualMonad d r m
  => Matrix r
  -> DualNumber d (Matrix r)
  -> m (DualNumber d (Vector r))
lossSoftMaxCrossEntropyL target (D u u') = do
  let expU = exp (u - HM.scalar (HM.maxElement u))  -- vs exploding gradients
      sumExpU = V.fromList $ map HM.sumElements $ HM.toColumns expU
      recipSum = recip sumExpU
      softMaxU = HM.asRow recipSum * expU
                   -- this @asRow@ is safe; multiplied at once
      scaled = D (negate $ log softMaxU * target)
                 (dScale (softMaxU - target) u')
  returnLet $ sumColumns1 scaled

-- * Operations resulting in a matrix

-- @2@ means rank two, so the dual component represents a matrix.
fromRows2 :: IsScalar d r
          => Data.Vector.Vector (DualNumber d (Vector r))
          -> DualNumber d (Matrix r)
fromRows2 v = D (HM.fromRows $ map (\(D u _) -> u) $ V.toList v)
                (dFromRows2 $ V.map (\(D _ u') -> u') v)

fromColumns2 :: IsScalar d r
             => Data.Vector.Vector (DualNumber d (Vector r))
             -> DualNumber d (Matrix r)
fromColumns2 v = D (HM.fromRows $ map (\(D u _) -> u) $ V.toList v)
                   (dFromColumns2 $ V.map (\(D _ u') -> u') v)

konst2 :: IsScalar d r => DualNumber d r -> (Int, Int) -> DualNumber d (Matrix r)
konst2 (D u u') sz = D (HM.konst u sz) (dKonst2 u' sz)

transpose2 :: IsScalar d r => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
transpose2 (D u u') = D (HM.tr' u) (dTranspose2 u')

-- | Dense matrix-matrix product.
--
-- If @u@ is a m x n (number of rows x number of columns) matrix
-- and @v@ is a n x p matrix then the result of @u <>! v@ is a m x p matrix.
infixr 8 <>!
(<>!) :: IsScalar d r
      => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
      -> DualNumber d (Matrix r)
(<>!) (D u u') (D v v') = D (u HM.<> v) (dAdd (dMD_M2 u' v) (dM_MD2 u v'))

-- | Dense matrix-matrix product with a constant matrix.
infixr 8 <>!!
(<>!!) :: IsScalar d r
       => DualNumber d (Matrix r) -> Matrix r
       -> DualNumber d (Matrix r)
(<>!!) (D u u') v = D (u HM.<> v) (dMD_M2 u' v)

rowAppend2 :: IsScalar d r
           => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
           -> DualNumber d (Matrix r)
rowAppend2 (D u u') (D v v') =
  D (u HM.=== v) (dRowAppend2 u' (HM.rows u) v')

columnAppend2 :: IsScalar d r
              => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
              -> DualNumber d (Matrix r)
columnAppend2 (D u u') (D v v') =
  D (u HM.||| v) (dColumnAppend2 u' (HM.cols u) v')

rowSlice2 :: IsScalar d r
          => Int -> Int -> DualNumber d (Matrix r)
          -> DualNumber d (Matrix r)
rowSlice2 i n (D u u') = D (HM.subMatrix (i, 0) (n, HM.cols u) u)
                           (dRowSlice2 i n u' (HM.rows u))

columnSlice2 :: IsScalar d r
             => Int -> Int -> DualNumber d (Matrix r)
             -> DualNumber d (Matrix r)
columnSlice2 i n (D u u') = D (HM.subMatrix (0, i) (HM.rows u, n) u)
                              (dColumnSlice2 i n u' (HM.rows u))

asRow2 :: IsScalar d r
       => DualNumber d (Vector r) -> Int -> DualNumber d (Matrix r)
asRow2 (D u u') n = D (HM.fromRows $ replicate n u) (dAsRow2 u')

asColumn2 :: IsScalar d r
          => DualNumber d (Vector r) -> Int -> DualNumber d (Matrix r)
asColumn2 (D u u') n = D (HM.fromColumns $ replicate n u) (dAsColumn2 u')

fromX2 :: IsScalar d r => DualNumber d (OT.Array r) -> DualNumber d (Matrix r)
fromX2 (D u u') = case OT.shapeL u of
  [_, cols] -> D (HM.reshape cols $ OT.toVector u) (dFromX2 u')
  dims -> error $ "fromX2: the tensor has wrong dimensions " ++ show dims

fromS2 :: forall rows cols d r.
          (KnownNat rows, KnownNat cols, IsScalar d r)
       => DualNumber d (OS.Array '[rows, cols] r) -> DualNumber d (Matrix r)
fromS2 (D u u') = D (HM.reshape (valueOf @cols) $ OS.toVector u) (dFromS2 u')

flipud2 :: IsScalar d r => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
flipud2 (D u u') = D (HM.flipud u) (dFlipud2 u')

fliprl2 :: IsScalar d r => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
fliprl2 (D u u') = D (HM.fliprl u) (dFliprl2 u')

vectorSlices2 :: IsScalar d r
              => Int -> DualNumber d (Vector r) -> DualNumber d (Matrix r)
vectorSlices2 n vv@(D v _) =
  fromRows2 $ V.fromList [slice1 i n vv | i <- [0 .. V.length v - n]]

reshape2 :: IsScalar d r
         => Int -> DualNumber d (Vector r) -> DualNumber d (Matrix r)
reshape2 cols (D u u') = D (HM.reshape cols u) (dReshape2 cols u')

-- TODO: This has list of matrices result instead of a cube tensor.
matrixSlices2 :: DualMonad d r m
              => Int -> DualNumber d (Matrix r) -> m [DualNumber d (Matrix r)]
matrixSlices2 dr m@(D u _) = do
  let (rows, cols) = HM.size u
      n = dr * cols
  v <- returnLet $ flatten1 m  -- used many times below
  let f k = returnLet $ reshape2 cols $ slice1 (k * cols) n v
  mapM f [0 .. rows - dr]

-- Not optimal: matrix is constructed and destructed immediately,
-- which is costly when evaluating delta expressions. The transposes
-- may not be optimal, either. This goes down to individual deltas
-- of scalars, which is horrible for performance. Unlike @corr1@
-- this uses the slow dot product instead of the fast matrix-vector
-- (or matrix-matrix) multiplication.
corr2 :: forall d r m. DualMonad d r m
      => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
      -> m (DualNumber d (Matrix r))
corr2 ker@(D u _) m@(D v _) = do
  let (rowsK, colsK) = HM.size u
      (rowsM, colsM) = HM.size v
      rr = rowsM - rowsK + 1
      rc = colsM - colsK + 1
  if | rowsK <= 0 || colsK <= 0 ->
       error $ "corr2: empty kernel not handled: " ++ show (rowsK, colsK)
     | rr <= 0 || rc <= 0 ->
       error $ "corr2: dim kernel " ++ show (rowsK, colsK)
               ++ " > dim matrix " ++ show (rowsM, colsM)
     | otherwise -> do
       kerTransV <- returnLet $ flatten1 (transpose2 ker)
       let dotColSlices :: DualNumber d (Matrix r) -> m [DualNumber d r]
           dotColSlices tm = do
             ttm <- returnLet $ transpose2 tm
             colSlices <- matrixSlices2 colsK ttm
             let f :: DualNumber d (Matrix r) -> DualNumber d r
                 f sm = kerTransV <.>! flatten1 sm
             return $ map f colSlices
       rowSlices <- matrixSlices2 rowsK m
       dotSlicesOfSlices <- mapM dotColSlices rowSlices
       returnLet $ reshape2 rc $ seq1 $ V.fromList $ concat dotSlicesOfSlices

conv2 :: forall d r m. DualMonad d r m
      => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
      -> m (DualNumber d (Matrix r))
conv2 ker@(D u _) m@(D v _) = do
  let (rowsK, colsK) = HM.size u
      (rowsM, colsM) = HM.size v
  if | rowsK <= 0 || colsK <= 0 ->
       returnLet $ konst2 0 (rowsM + rowsK - 1, colsM + colsK - 1)
     | otherwise -> do
       let zRow = konst2 0 (rowsK - 1, colsM)
           rowPadded = rowAppend2 zRow $ rowAppend2 m zRow
           zCol = konst2 0 (rowsM + 2 * (rowsK - 1), colsK - 1)
           padded = columnAppend2 zCol $ columnAppend2 rowPadded zCol
       corr2 (fliprl2 . flipud2 $ ker) padded

conv2' :: IsScalar d r
       => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
       -> DualNumber d (Matrix r)
conv2' (D u u') (D v v') = D (HM.conv2 u v) (dAdd (dConv2 u v') (dConv2 v u'))

-- A variant with limited padding, corresponding to SAME padding
-- from Tensorflow. Data size does not change with this padding.
-- It also performs convolution wrt flipped kernel (and so saves
-- on flipping it here), which makes no practical difference when
-- the kernel is initialized randomly.
convSame2 :: forall d r m. DualMonad d r m
          => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
          -> m (DualNumber d (Matrix r))
convSame2 ker@(D u _) m@(D v _) = do
  let (rowsK, colsK) = HM.size u
      (rowsM, colsM) = HM.size v
  if | rowsK <= 0 || colsK <= 0 ->
       returnLet $ konst2 0 (rowsM, colsM)
     | otherwise -> do
       let zRow = konst2 0 ((rowsK - 1) `div` 2, colsM)
           rowPadded = rowAppend2 zRow $ rowAppend2 m zRow
           zCol = konst2 0 (rowsM + rowsK - 1, (colsK - 1) `div` 2)
           padded = columnAppend2 zCol $ columnAppend2 rowPadded zCol
       corr2 ker padded

-- No padding; remaining areas ignored.
maxPool2 :: forall d r m. DualMonad d r m
         => Int -> Int -> DualNumber d (Matrix r) -> m (DualNumber d (Matrix r))
maxPool2 ksize stride m@(D u _) = do
  let (rows, cols) = HM.size u
      colsOut = cols `div` stride
      resultRows = [0, stride .. rows - ksize]
      resultCols = [0, stride .. cols - ksize]
      resultCoords = [(r, c) | r <- resultRows, c <- resultCols]
  v <- returnLet $ flatten1 m  -- used many times below
  let getArea :: (Int, Int) -> DualNumber d (Vector r)
      getArea (r0, c0) =
        let getAreaAtRow r1 = append1 (slice1 (r1 * cols + c0) ksize v)
        in foldr getAreaAtRow (seq1 V.empty) [r0 .. r0 + ksize - 1]
      mins = map (maximum0 . getArea) resultCoords
  returnLet $ reshape2 colsOut $ seq1 $ V.fromList mins

-- * Operations resulting in an arbitrary untyped tensor

konstX :: IsScalar d r => DualNumber d r -> OT.ShapeL -> DualNumber d (OT.Array r)
konstX (D u u') sh = D (OT.constant sh u) (dKonstX u' sh)

appendX :: IsScalar d r
        => DualNumber d (OT.Array r) -> DualNumber d (OT.Array r)
        -> DualNumber d (OT.Array r)
appendX (D u u') (D v v') =
  D (u `OT.append` v) (dAppendX u' (head $ OT.shapeL u) v')

sliceX :: IsScalar d r
       => Int -> Int -> DualNumber d (OT.Array r) -> DualNumber d (OT.Array r)
sliceX i n (D u u') = D (OT.slice [(i, n)] u)
                        (dSliceX i n u' (head $ OT.shapeL u))

indexX :: IsScalar d r
       => DualNumber d (OT.Array r) -> Int -> DualNumber d (OT.Array r)
indexX (D u u') ix = D (OT.index u ix)
                       (dIndexX u' ix (head $ OT.shapeL u))

ravelFromListX :: IsScalar d r
               => [DualNumber d (OT.Array r)] -> DualNumber d (OT.Array r)
ravelFromListX ld =
  let (lu, lu') = unzip $ map (\(D u u') -> (u, u')) ld
      sh = case lu of
        u : _ -> length lu : OT.shapeL u
        [] -> []
  in D (OT.ravel $ OTB.fromList sh lu) (dRavelFromListX lu')

unravelToListX :: IsScalar d r
               => DualNumber d (OT.Array r) -> [DualNumber d (OT.Array r)]
unravelToListX (D v v') = case OT.shapeL v of
  k : _ ->
    let g ix p = D p (dIndexX v' ix k)
    in imap g $ OTB.toList $ OT.unravel v
  [] -> error "unravelToListX: wrong tensor dimensions"  -- catch early

mapX :: IsScalar d r
     => (DualNumber d (OT.Array r) -> DualNumber d (OT.Array r))
     -> DualNumber d (OT.Array r)
     -> DualNumber d (OT.Array r)
mapX f = ravelFromListX . map f . unravelToListX

zipWithX :: IsScalar d r
         => (DualNumber d (OT.Array r) -> DualNumber d (OT.Array r)
             -> DualNumber d (OT.Array r))
         -> DualNumber d (OT.Array r) -> DualNumber d (OT.Array r)
         -> DualNumber d (OT.Array r)
zipWithX f d e =
  ravelFromListX $ zipWith f (unravelToListX d) (unravelToListX e)

reshapeX :: IsScalar d r
         => OT.ShapeL -> DualNumber d (OT.Array r) -> DualNumber d (OT.Array r)
reshapeX sh' (D u u') = D (OT.reshape sh' u) (dReshapeX (OT.shapeL u) sh' u')

from0X :: IsScalar d r => DualNumber d r -> DualNumber d (OT.Array r)
from0X (D u u') = D (OT.scalar u) (dFrom0X u')

from1X :: IsScalar d r => DualNumber d (Vector r) -> DualNumber d (OT.Array r)
from1X (D u u') = D (OT.fromVector [V.length u] u) (dFrom1X u')

from2X :: IsScalar d r => DualNumber d (Matrix r) -> DualNumber d (OT.Array r)
from2X (D u u') = D (OT.fromVector [HM.rows u, HM.cols u] $ HM.flatten u)
                    (dFrom2X u' (HM.cols u))

fromSX :: forall sh d r. (IsScalar d r, OS.Shape sh)
       => DualNumber d (OS.Array sh r) -> DualNumber d (OT.Array r)
fromSX (D u u') = D (Data.Array.Convert.convert u) (dFromSX u')

-- * Operations resulting in an arbitrary fully typed Shaped tensor

konstS :: (IsScalar d r, OS.Shape sh)
       => DualNumber d r -> DualNumber d (OS.Array sh r)
konstS (D u u') = D (OS.constant u) (dKonstS u')

appendS :: (KnownNat m, KnownNat n, IsScalar d r, OS.Shape sh)
        => DualNumber d (OS.Array (m ': sh) r)
        -> DualNumber d (OS.Array (n ': sh) r)
        -> DualNumber d (OS.Array ((m + n) ': sh) r)
appendS (D u u') (D v v') = D (u `OS.append` v) (dAppendS u' v')

sliceS :: forall i n k rest d r.
          (KnownNat i, KnownNat n, KnownNat k, IsScalar d r, OS.Shape rest)
       => DualNumber d (OS.Array (i + n + k ': rest) r)
       -> DualNumber d (OS.Array (n ': rest) r)
sliceS (D u u') = D (OS.slice @'[ '(i, n) ] u)
                    (dSliceS (Proxy :: Proxy i) Proxy u')

indexS :: forall ix k rest d r.
          (KnownNat ix, KnownNat k, IsScalar d r, OS.Shape rest)
       => DualNumber d (OS.Array (ix + 1 + k ': rest) r)
       -> DualNumber d (OS.Array rest r)
indexS (D u u') = D (OS.index u (valueOf @ix))
                    (dIndexS u' (Proxy :: Proxy ix))

ravelFromListS :: forall rest k d r.
                  (KnownNat k, IsScalar d r, OS.Shape rest)
               => [DualNumber d (OS.Array rest r)]
               -> DualNumber d (OS.Array (k : rest) r)
ravelFromListS ld =
  let (lu, lu') = unzip $ map (\(D u u') -> (u, u')) ld
  in D (OS.ravel $ OSB.fromList lu) (dRavelFromListS lu')

unravelToListS :: forall k rest d r.
                  (KnownNat k, IsScalar d r, OS.Shape rest)
               => DualNumber d (OS.Array (k : rest) r)
               -> [DualNumber d (OS.Array rest r)]
unravelToListS (D v v') =
  -- @dIndexS@ is rigid, with type-level bound-checking, so we have to switch
  -- to @dIndexX@ for this function.
  let g ix p = D p (dFromXS $ dIndexX (dFromSX v') ix (valueOf @k))
  in imap g $ OSB.toList $ OS.unravel v

mapS :: forall k sh1 sh d r. (KnownNat k, IsScalar d r, OS.Shape sh, OS.Shape sh1)
     => (DualNumber d (OS.Array sh1 r) -> DualNumber d (OS.Array sh r))
     -> DualNumber d (OS.Array (k : sh1) r)
     -> DualNumber d (OS.Array (k : sh) r)
mapS f = ravelFromListS . map f . unravelToListS

mapMS :: forall k sh1 sh d r m.
         (Monad m, KnownNat k, IsScalar d r, OS.Shape sh, OS.Shape sh1)
      => (DualNumber d (OS.Array sh1 r) -> m (DualNumber d (OS.Array sh r)))
      -> DualNumber d (OS.Array (k : sh1) r)
      -> m (DualNumber d (OS.Array (k : sh) r))
mapMS f d = do
  let ld = unravelToListS d
  ld2 <- mapM f ld
  return $! ravelFromListS ld2

zipWithS :: forall k sh1 sh2 sh d r.
            ( KnownNat k, IsScalar d r, OS.Shape sh, OS.Shape sh1, OS.Shape sh2)
         => (DualNumber d (OS.Array sh1 r) -> DualNumber d (OS.Array sh2 r)
             -> DualNumber d (OS.Array sh r))
         -> DualNumber d (OS.Array (k : sh1) r)
         -> DualNumber d (OS.Array (k : sh2) r)
         -> DualNumber d (OS.Array (k : sh) r)
zipWithS f d e =
  ravelFromListS $ zipWith f (unravelToListS d) (unravelToListS e)

reshapeS :: (IsScalar d r, OS.Shape sh, OS.Shape sh', OS.Size sh ~ OS.Size sh')
         => DualNumber d (OS.Array sh r) -> DualNumber d (OS.Array sh' r)
reshapeS (D u u') = D (OS.reshape u) (dReshapeS u')

-- TODO: generalize as broadcast or stretch
asRowS :: forall k n d r. (IsScalar d r, KnownNat k, KnownNat n)
       => DualNumber d (OS.Array '[k] r) -> DualNumber d (OS.Array '[n, k] r)
asRowS d = from2S $ asRow2 (fromS1 d) (valueOf @n)

asColumnS :: forall k n d r. (IsScalar d r, KnownNat k, KnownNat n)
          => DualNumber d (OS.Array '[k] r) -> DualNumber d (OS.Array '[k, n] r)
asColumnS d = from2S $ asColumn2 (fromS1 d) (valueOf @n)

from0S :: IsScalar d r => DualNumber d r -> DualNumber d (OS.Array '[] r)
from0S (D u u') = D (OS.scalar u) (dFrom0S u')

from1S :: (KnownNat n, IsScalar d r)
       => DualNumber d (Vector r) -> DualNumber d (OS.Array '[n] r)
from1S (D u u') = D (OS.fromVector u) (dFrom1S u')

from2S :: (KnownNat rows, KnownNat cols, IsScalar d r)
       => DualNumber d (Matrix r) -> DualNumber d (OS.Array '[rows, cols] r)
from2S (D u u') = D (OS.fromVector $ HM.flatten u) (dFrom2S Proxy u')

fromXS :: (IsScalar d r, OS.Shape sh)
       => DualNumber d (OT.Array r) -> DualNumber d (OS.Array sh r)
fromXS (D u u') = D (Data.Array.Convert.convert u) (dFromXS u')

-- TODO: generalize to arbitrary permutations of arbitrarily many ranks using https://hackage.haskell.org/package/orthotope/docs/Data-Array-ShapedS.html#v:transpose
transpose2S :: (IsScalar d r, KnownNat rows, KnownNat cols)
            => DualNumber d (OS.Array '[rows, cols] r)
            -> DualNumber d (OS.Array '[cols, rows] r)
transpose2S = from2S . transpose2 . fromS2

infixr 8 #>$
(#>$) :: (IsScalar d r, KnownNat rows, KnownNat cols)
      => DualNumber d (OS.Array '[rows, cols] r)
      -> DualNumber d (OS.Array '[cols] r)
      -> DualNumber d (OS.Array '[rows] r)
(#>$) d e = from1S $ fromS2 d #>! fromS1 e

infixr 8 <>$
(<>$) :: (IsScalar d r, KnownNat m, KnownNat n, KnownNat p)
      => DualNumber d (OS.Array '[m, n] r)
      -> DualNumber d (OS.Array '[n, p] r)
      -> DualNumber d (OS.Array '[m, p] r)
(<>$) d e = from2S $ fromS2 d <>! fromS2 e

conv2S :: forall d r kheight_minus_1 kwidth_minus_1 in_height in_width.
          ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1
          , KnownNat in_height, KnownNat in_width
          , IsScalar d r )
       => DualNumber d (OS.Array '[kheight_minus_1 + 1, kwidth_minus_1 + 1] r)
       -> DualNumber d (OS.Array '[in_height, in_width] r)
       -> DualNumber d (OS.Array '[ in_height + kheight_minus_1
                                 , in_width + kwidth_minus_1 ] r)
conv2S ker x = from2S $ conv2' (fromS2 ker) (fromS2 x)

-- Convolution of many matrices at once. Some of the names of dimensions
-- are from https://www.tensorflow.org/api_docs/python/tf/nn/conv2d
conv24 :: forall kheight_minus_1 kwidth_minus_1
                 out_channels in_height in_width batch_size in_channels d r.
          ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1
          , KnownNat out_channels, KnownNat in_height, KnownNat in_width
          , KnownNat batch_size, KnownNat in_channels
          , IsScalar d r )
       => DualNumber d (OS.Array '[ out_channels, in_channels
                                 , kheight_minus_1 + 1, kwidth_minus_1 + 1 ] r)
       -> DualNumber d (OS.Array '[batch_size, in_channels, in_height, in_width] r)
       -> DualNumber d (OS.Array '[ batch_size, out_channels
                                 , in_height + kheight_minus_1
                                 , in_width + kwidth_minus_1 ] r)
conv24 ker = mapS conv23 where
  conv23 :: DualNumber d (OS.Array '[in_channels, in_height, in_width] r)
         -> DualNumber d (OS.Array '[ out_channels
                                   , in_height + kheight_minus_1
                                   , in_width + kwidth_minus_1 ] r)
  conv23 x = mapS (convFilters x) ker
  convFilters
    :: DualNumber d (OS.Array '[in_channels, in_height, in_width] r)
    -> DualNumber d (OS.Array '[ in_channels
                              , kheight_minus_1 + 1, kwidth_minus_1 + 1 ] r)
    -> DualNumber d (OS.Array '[ in_height + kheight_minus_1
                              , in_width + kwidth_minus_1 ] r)
  convFilters x ker1 = sumOutermost $ zipWithS conv2S ker1 x
  sumOutermost :: DualNumber d (OS.Array '[ in_channels
                                         , in_height + kheight_minus_1
                                         , in_width + kwidth_minus_1 ] r)
               -> DualNumber d (OS.Array '[ in_height + kheight_minus_1
                                         , in_width + kwidth_minus_1 ] r)
  sumOutermost = sum . unravelToListS
    -- slow; should go through Tensor2, or the Num instance should when possible

maxPool24
  :: forall ksize_minus_1 stride in_height in_width batch_size channels d r m.
     ( KnownNat ksize_minus_1, KnownNat stride
     , KnownNat in_height, KnownNat in_width
     , KnownNat batch_size, KnownNat channels
     , 1 <= stride
     , ksize_minus_1 <= in_height
     , ksize_minus_1 <= in_width
     , 1 <= in_height - ksize_minus_1 + stride
     , 1 <= in_width - ksize_minus_1 + stride
     , DualMonad d r m )
     => DualNumber d (OS.Array '[batch_size, channels, in_height, in_width] r)
     -> m (DualNumber d
             (OS.Array '[ batch_size, channels
                         , (in_height - ksize_minus_1) `DivRoundUp` stride
                         , (in_width - ksize_minus_1) `DivRoundUp` stride ] r))
maxPool24 d = do
  res <- mapMS (mapMS (fmap from2S
                       . maxPool2 (valueOf @ksize_minus_1 + 1)
                                  (valueOf @stride)
                       . fromS2)) d
  returnLet res

-- * Operations creating delayed/outlined derivatives

-- | The version of the @D@ constructor lazy in the second argument.
-- To be used as in
--
-- > sinDelayed :: (Floating a, IsPrimal d a) => DualNumber d a -> DualNumber d a
-- > sinDelayed (D u u') = delayD (sin u) (dScale (cos u) u')
-- >
-- > plusDelayed :: (Floating a, IsPrimal d a)
-- >             => DualNumber d a -> DualNumber d a -> DualNumber d a
-- > plusDelayed (D u u') (D v v') = delayD (u + v) (dAdd u' v')
-- >
-- > x ** (sinDelayed x `plusDelayed` (id2 $ id2 $ id2 $ konst1 (sumElements0 x) 2))
--
-- The outlining is lost when serializing or logging, unlike with @Out@,
-- @Outline0@, etc.
--
-- Yet another incomparable variant that can't be serialized would be
-- (illustrating with an example of a constructor at rank 0)
--
-- > FromParams0 (Domains -> Delta0 r)
--
-- that expects the initial parameters. But it's more troublesome
-- than @Delay0@ both in implementation and usage.
delayD :: IsPrimal d a => a -> Dual d a -> DualNumber d a
delayD u ~u' = D u (dDelay u')

-- | A wrapper type to delay/outline computation of the derivatives of the given
-- primitive numeric function inside the dual component of the created dual
-- number. The rule is that if all arguments of a function are wrapped
-- in @Out@ then the function gets delayed. Inconsistent wrapping,
-- e.g., only one of the arguments, leads to early type errors.
--
-- To be used as in
--
-- > x ** unOut (sin (Out x) + Out (id $ id $ id $ konst1 (sumElements0 x) 2))
--
-- which delays computing the dual component of sine and of addition
-- (both in rank 1), but not of power, konst and sumElements. The last two
-- can't be currently delayed, because only primitive numeric functions
-- are supported (an attempt would not type-check).
newtype Out a = Out {unOut :: a}
  deriving (Eq, Ord)

returnOut :: (DualMonad d r m, IsPrimalWithScalar d a r)
          => Out (DualNumber d a) -> m (Out (DualNumber d a))
returnOut dOut = do
  dvar <- returnLet $ unOut dOut
  return $ Out dvar

instance (Num a, IsPrimal 'DModeGradient a, HasVariables a)
         => Num (Out (DualNumber 'DModeGradient a)) where
  Out (D u u') + Out (D v v') =
    Out $ D (u + v) (dOutline PlusOut [u, v] [u', v'])
  Out (D u u') - Out (D v v') =
    Out $ D (u - v) (dOutline MinusOut [u, v] [u', v'])
  Out (D u u') * Out (D v v') =
    Out $ D (u * v) (dOutline TimesOut [u, v] [u', v'])
  negate (Out (D v v')) = Out $ D (negate v) (dOutline NegateOut [v] [v'])
  abs (Out (D v v')) = Out $ D (abs v) (dOutline AbsOut [v] [v'])
  signum (Out (D v v')) = Out $ D (signum v) (dOutline SignumOut [v] [v'])
  fromInteger = Out . constant . fromInteger

instance (Real a, IsPrimal 'DModeGradient a, HasVariables a)
         => Real (Out (DualNumber 'DModeGradient a)) where
  toRational = undefined  -- TODO?

instance (Fractional a, IsPrimal 'DModeGradient a, HasVariables a)
         => Fractional (Out (DualNumber 'DModeGradient a)) where
  Out (D u u') / Out (D v v') =
    Out $ D (u / v) (dOutline DivideOut [u, v] [u', v'])
  recip (Out (D v v')) = Out $ D (recip v) (dOutline RecipOut [v] [v'])
  fromRational = Out . constant . fromRational

instance (Floating a, IsPrimal 'DModeGradient a, HasVariables a)
         => Floating (Out (DualNumber 'DModeGradient a)) where
  pi = Out $ constant pi
  exp (Out (D u u')) = Out $ D (exp u) (dOutline ExpOut [u] [u'])
  log (Out (D u u')) = Out $ D (log u) (dOutline LogOut [u] [u'])
  sqrt (Out (D u u')) = Out $ D (sqrt u) (dOutline SqrtOut [u] [u'])
  Out (D u u') ** Out (D v v') =
    Out $ D (u ** v) (dOutline PowerOut [u, v] [u', v'])
  logBase (Out (D u u')) (Out (D v v')) = Out $ D (logBase u v) (dOutline LogBaseOut [u, v] [u', v'])
  sin (Out (D u u')) = Out $ D (sin u) (dOutline SinOut [u] [u'])
  cos (Out (D u u')) = Out $ D (cos u) (dOutline CosOut [u] [u'])
  tan (Out (D u u')) = Out $ D (tan u) (dOutline TanOut [u] [u'])
  asin (Out (D u u')) = Out $ D (asin u) (dOutline AsinOut [u] [u'])
  acos (Out (D u u')) = Out $ D (acos u) (dOutline AcosOut [u] [u'])
  atan (Out (D u u')) = Out $ D (atan u) (dOutline AtanOut [u] [u'])
  sinh (Out (D u u')) = Out $ D (sinh u) (dOutline SinhOut [u] [u'])
  cosh (Out (D u u')) = Out $ D (cosh u) (dOutline CoshOut [u] [u'])
  tanh (Out (D u u')) = Out $ D (tanh u) (dOutline TanhOut [u] [u'])
  asinh (Out (D u u')) = Out $ D (asinh u) (dOutline AsinhOut [u] [u'])
  acosh (Out (D u u')) = Out $ D (acosh u) (dOutline AcoshOut [u] [u'])
  atanh (Out (D u u')) = Out $ D (atanh u) (dOutline AtanhOut [u] [u'])

instance (RealFrac a, IsPrimal 'DModeGradient a, HasVariables a)
         => RealFrac (Out (DualNumber 'DModeGradient a)) where
  properFraction = undefined
    -- very low priority, since these are all extremely not continuous

instance (RealFloat a, IsPrimal 'DModeGradient a, HasVariables a)
         => RealFloat (Out (DualNumber 'DModeGradient a)) where
  atan2 (Out (D u u')) (Out (D v v')) =
    Out $ D (atan2 u v) (dOutline Atan2Out [u, v] [u', v'])
      -- we can be selective here and omit the other methods,
      -- most of which don't even have a differentiable codomain

-- * Busywork to let the derivatives mode ignore all outlining

-- | Note that this should apply only when @d@ is @'DModeDerivative@.
-- However, GHC can't tell that @d@ has only two cases. Therefore, we need
-- to overgeneralize these definitions and mark them with @OVERLAPPABLE@
-- or else GHC complains that not enough instances are given
-- whenever type-checking code polymorphic on @d@.
instance {-# OVERLAPPABLE #-} (Num a, IsPrimal d a)
                              => Num (Out (DualNumber d a)) where
  Out d + Out e = Out (d + e)
  Out d - Out e = Out (d - e)
  Out d * Out e = Out (d * e)
  negate (Out e) = Out (negate e)
  abs (Out e) = Out (abs e)
  signum (Out e) = Out (signum e)
  fromInteger = Out . constant . fromInteger

instance {-# OVERLAPPABLE #-} (Real a, IsPrimal d a)
                              => Real (Out (DualNumber d a)) where
  toRational = undefined  -- TODO?

instance {-# OVERLAPPABLE #-} (Fractional a, IsPrimal d a)
                              => Fractional (Out (DualNumber d a)) where
  Out d / Out e = Out (d / e)
  recip (Out e) = Out (recip e)
  fromRational = Out . constant . fromRational

instance {-# OVERLAPPABLE #-} (Floating a, IsPrimal d a)
                              => Floating (Out (DualNumber d a)) where
  pi = Out $ constant pi
  exp (Out d) = Out (exp d)
  log (Out d) = Out (log d)
  sqrt (Out d) = Out (sqrt d)
  Out d ** Out e = Out (d ** e)
  logBase (Out x) (Out y) = Out (logBase x y)
  sin (Out d) = Out (sin d)
  cos (Out d) = Out (cos d)
  tan (Out d) = Out (tan d)
  asin (Out d) = Out (asin d)
  acos (Out d) = Out (acos d)
  atan (Out d) = Out (atan d)
  sinh (Out d) = Out (sinh d)
  cosh (Out d) = Out (cosh d)
  tanh (Out d) = Out (tanh d)
  asinh (Out d) = Out (asinh d)
  acosh (Out d) = Out (acosh d)
  atanh (Out d) = Out (atanh d)

instance {-# OVERLAPPABLE #-} (RealFrac a, IsPrimal d a)
                              => RealFrac (Out (DualNumber d a)) where
  properFraction = undefined
    -- very low priority, since these are all extremely not continuous

instance {-# OVERLAPPABLE #-} (RealFloat a, IsPrimal d a)
                              => RealFloat (Out (DualNumber d a)) where
  atan2 (Out d) (Out e) = Out (atan2 d e)


module VoiceLeading.Theory where

import           VoiceLeading.Base
import           VoiceLeading.Helpers           ( rotate )

import           Data.List                      ( elemIndex )
import           Data.Maybe                     ( mapMaybe )
import           Data.Ratio
import           Data.Function.Memoize          ( memoize )

import           Data.Aeson
import qualified Data.Map                      as M
import qualified Data.Vector                   as V
import qualified Data.Vector.Unboxed           as VU
import qualified Data.Vector.Unboxed.Mutable   as VUM

import           Numeric.SpecFunctions          ( logBeta )
import qualified Control.Loop                  as L

-- simple ad-hoc heuristic
--------------------------

modal :: Int -> Int -> Int
modal i m = [0, 2, 4, 5, 7, 9, 11] !! mod (i + m) 7

cMajor :: [Int]
cMajor = [2, 2, 1, 2, 2, 2, 1]

scale :: KeySig -> [Int]
scale = memoize s
 where
  s (KeySig root mode) =
    scanl (\a b -> (a + b) `mod` 12) root (init $ rotate mode cMajor)

leadingTone :: KeySig -> Int
leadingTone = memoize lt where lt (KeySig root _) = (root - 1) `mod` 12

majorChord :: [Int]
majorChord = [0, 4, 7, 10, 11, 2, 1, 5, 6, 9, 8, 3]

minorChord :: [Int]
minorChord = [0, 3, 7, 10, 2, 11, 1, 5, 6, 8, 9, 4]

-- | Returns a pair of the estimated root and the "chordness" of the given set of pitches
findHarm :: [Int] -> (Int, Double)
findHarm []      = (0, 0)
findHarm pitches = testRoot [0 .. 11]
 where
  ps = map (`mod` 12) pitches
  indexVal :: Int -> Rational
  indexVal i = 1 % (fromIntegral i + 1)
  testChord chord tps = sum $ map val [0 .. maximum is]
   where
    is = mapMaybe (`elemIndex` chord) tps
    val i = if i `elem` is then indexVal i else negate (indexVal i)
  testRoot [] = (0, 0) -- just a catchall
  testRoot (r : rs) =
    let tps   = map ((`mod` 12) . (\x -> x - r)) ps
        score = max (testChord majorChord tps) (testChord minorChord tps)
    in  testRoot' rs r score
  testRoot' []       best score = (best, fromRational score)
  testRoot' (r : rs) best score = if newScore > score
    then testRoot' rs r newScore
    else testRoot' rs best score
   where
    tps      = map ((`mod` 12) . (\x -> x - r)) ps
    newScore = max (testChord majorChord tps) (testChord minorChord tps)

-- chord profiles
-----------------

type Profile = (Double, VU.Vector Double)
type Profiles = V.Vector Profile

loadProfiles :: FilePath -> IO (Maybe (M.Map String (M.Map Int Double)))
loadProfiles = decodeFileStrict'

-- | Returns a vector of profiles.
-- Each profile is a pair of the total counts in the profile and the smoothed count vector.
vectorizeProfiles :: M.Map String (M.Map Int Double) -> Profiles
vectorizeProfiles profmap = V.fromList $ profList <$> M.elems profmap
 where
  profList pm = (VU.sum lst, lst)
    where lst = VU.fromList [ 1 + pm M.! i | i <- [0 .. 11] ]

logDirichletMultinomial :: Profile -> VU.Vector Double -> Double -> Double
logDirichletMultinomial (a0, as) xs n = num - denom
 where
  l = VU.length xs
  entry acc i = if xi > 0 then acc + log xi + logBeta (as VU.! i) xi else acc
    where xi = xs VU.! i
  num   = log n + logBeta a0 n
  denom = L.forLoopFold 0 (< l) (+ 1) 0 entry

mkChord :: Int -> [Int] -> VU.Vector Double
mkChord r pitches = VU.create $ do
  chord <- VUM.replicate 12 0
  mapM_ (VUM.unsafeModify chord (+ 1) . (`mod` 12) . subtract r) pitches
  pure chord

matchChordProfiles :: Profiles -> [Int] -> (Int, Double)
matchChordProfiles _        []      = (0, 0)
matchChordProfiles profiles pitches = (maxRoot, chordness / 0.3) -- normalize by maximum value
 where
  roots = V.fromListN 12 [0 .. 11]
  n     = fromIntegral $ length pitches
  match chord prof = logDirichletMultinomial prof chord n
  bestMatch r = V.maximum $ match (mkChord r pitches) <$> profiles
  bestMatches = bestMatch <$> roots
  maxRoot     = V.maxIndex bestMatches
  chordness   = exp $ bestMatches V.! maxRoot


{-# LANGUAGE FlexibleContexts, NamedFieldPuns #-}

module School.App.CSVReader
( csvToBinary
, csvToMatrixDouble
, parseDoubles
, readCSV
) where

import Conduit ((.|), ConduitM, mapC, mapMC, sourceFileBS)
import qualified Data.ByteString as BS
import Data.ByteString.Conversion (fromByteString)
import Data.Void (Void)
import qualified Data.Conduit.Binary as CB
import School.FileIO.AppIO (AppIO, maybeToAppIO)
import School.FileIO.FileType (FileType(..))
import School.FileIO.FileHeader (FileHeader(..))
import School.FileIO.MatrixSink (matrixDoubleSink)
import School.Utils.Constants (binComma)
import Numeric.LinearAlgebra ((><), Matrix)

parseDoubles :: [BS.ByteString] -> Maybe [Double]
parseDoubles = mapM fromByteString

readCSV :: FilePath -> ConduitM ()
                                [BS.ByteString]
                                AppIO
                                ()
readCSV path = sourceFileBS path
            .| CB.lines
            .| mapC (BS.split binComma)

csvToMatrixDouble :: FileHeader
                  -> ConduitM [BS.ByteString]
                              (Matrix Double)
                              AppIO
                              ()
csvToMatrixDouble FileHeader { cols } =
    mapC parseDoubles
 .| mapMC (maybeToAppIO "Could not parse doubles")
 .| mapC (1 >< cols)

csvToBinary :: FilePath
            -> FilePath
            -> FileHeader
            -> ConduitM () Void AppIO ()
csvToBinary inPath outPath header  =
    readCSV inPath
 .| csvToMatrixDouble header
 .| matrixDoubleSink SM header outPath


module Main
( main
) where

import Debug.Trace

import Codec.Wav (importFile)
import Data.Audio (Audio(..))

import Data.Array.Unboxed (elems, assocs)
import Data.Array.IArray (Array(..), listArray, array)
import Data.Complex (magnitude)
import Data.Int (Int64, Int32, Int16, Int8)

import Common (importantFrequencies, breakChunks, slice)
import SongId (ft, mapChunk2Freqs, chunkIdentifier)

type WordSize = Int32
chunkSize = 5000
file = ""

main :: IO ()
main = do
    input <- importFile file
    song_dft <- case input :: Either String (Audio WordSize) of
        Left err -> do
            putStrLn err
            return [(array (0,0) [])]
        Right a@(Audio _ _ samples) -> return (ft chunkSize samples)

    putStrLn $ show (length song_dft) ++ " " ++ show (length . elems $ head song_dft)

    let ys = map (mapChunk2Freqs . elems) song_dft

    let freqs = importantFrequencies 5 (8 * chunkSize `quot` 50 `quot` 2)
    let chunks = map (chunkIdentifier freqs) ys
    foldl (\acc x -> acc >> print x) (return ()) $ chunks