{-# LANGUAGE DeriveDataTypeable    #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE TypeFamilies          #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Plots.Axis
-- Copyright   :  (C) 2015 Christopher Chalmers
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Christopher Chalmers
-- Stability   :  experimental
-- Portability :  non-portable
--
-- The 'Axis' is the main data type for "plots". It holds all the
-- necessary infomation to be rendered into a 'Diagram'.
--
----------------------------------------------------------------------------
module Plots.Axis
  ( -- * Axis type
    Axis
  , axes
  , axisPlots
  , currentPlots
  , finalPlots
  , plotModifier
  , axisSize
  , colourBarRange

    -- * Predefined axes
  , r2Axis
  , polarAxis

    -- ** Base space
  , BaseSpace

    -- * Axis plots
  , addPlot
  , addPlotable
  , addPlotable'

    -- * Single axis
  , SingleAxis

    -- * Specific axes
    -- ** x-axis
  , xAxis
  , xLabel
  , xMin
  , xMax

    -- ** y-axis
  , yAxis
  , yLabel
  , yMin
  , yMax

    -- ** r-axis
  , rAxis
  , rLabel
  , rMax

    -- ** theta-axis
  , thetaAxis
  , thetaLabel

    -- ** z-axis
  , zAxis
  , zLabel
  , zMin
  , zMax
  ) where

import           Control.Monad.State
import           Data.Complex
import           Data.Default
import           Data.Typeable

import           Diagrams.Coordinates.Polar
import           Diagrams.Prelude
import           Diagrams.TwoD.Text

import           Plots.Axis.ColourBar
import           Plots.Axis.Grid
import           Plots.Axis.Labels
import           Plots.Axis.Line
import           Plots.Axis.Scale
import           Plots.Axis.Title
import           Plots.Axis.Ticks
import           Plots.Legend
import           Plots.Style
import           Plots.Types

import           Linear

------------------------------------------------------------------------
-- Axis data type
------------------------------------------------------------------------

-- Single axis ---------------------------------------------------------

-- | Render information for a single axis line.
data SingleAxis b v n = SingleAxis
-- note the the v is only present for Style v n
  { saLabel     :: AxisLabel b v n
  , saLine      :: AxisLine v n
  , saTickLabel :: TickLabels b v n
  , saScaling   :: AxisScaling n
  , saGridLines :: GridLines v n
  , saTicks     :: Ticks v n
  , saVisible   :: Bool
  }

type instance V (SingleAxis b v n) = v
type instance N (SingleAxis b v n) = n

instance (TypeableFloat n, Renderable (Text n) b)
    => Default (SingleAxis b V2 n) where
  def = SingleAxis
    { saLabel      = def
    , saLine       = def
    , saTickLabel  = def
    , saGridLines  = def
    , saTicks      = def
    , saScaling    = def
    , saVisible    = True
    }

instance Functor f => HasTicks f (SingleAxis b v n) where
  bothTicks = lens saTicks (\sa ticks -> sa {saTicks = ticks})

instance Functor f => HasMajorTicks f (SingleAxis b v n) where
  majorTicks = bothTicks . majorTicks

instance Functor f => HasMinorTicks f (SingleAxis b v n) where
  minorTicks = bothTicks . minorTicks

instance Functor f => HasAxisLabel f (SingleAxis b v n) b where
  axisLabel = lens saLabel (\sa l -> sa {saLabel = l})

instance Functor f => HasTickLabels f (SingleAxis b v n) b where
  tickLabel = lens saTickLabel (\sa tl -> sa {saTickLabel = tl})

instance Functor f => HasAxisLine f (SingleAxis b v n) where
  axisLine = lens saLine (\sa l -> sa {saLine = l})

instance Functor f => HasGridLines f (SingleAxis b v n) where
  gridLines = lens saGridLines (\sa l -> sa {saGridLines = l})

instance Functor f => HasMajorGridLines f (SingleAxis b v n) where
  majorGridLines = gridLines . majorGridLines

instance Functor f => HasMinorGridLines f (SingleAxis b v n) where
  minorGridLines = gridLines . minorGridLines

instance Functor f => HasAxisScaling f (SingleAxis b v n) where
  axisScaling = lens saScaling (\sa s -> sa {saScaling = s})

instance HasVisibility (SingleAxis b v n) where
  visible = lens saVisible (\sa b -> sa {saVisible = b})

-- singleAxisScale :: Lens' (SingleAxis b v n) AxisScale
-- singleAxisScale = lens saScale (\sa s -> sa {saScale = s})

-- singleAxisBound :: Lens' (SingleAxis b v n) (Bound n)
-- singleAxisBound = lens saBounds (\sa b -> sa {saBounds = b})

------------------------------------------------------------------------
-- Axis type
------------------------------------------------------------------------

-- Base space ----------------------------------------------------------

-- | This family is used so that we can say (Axis Polar) but use V2 for the
--   underlying diagram.
type family BaseSpace (c :: * -> *) :: * -> *

type instance BaseSpace V2      = V2
type instance BaseSpace Complex = V2
type instance BaseSpace Polar   = V2
type instance BaseSpace V3      = V3

-- Axis data type ------------------------------------------------------

-- | Axis is the data type that holds all the necessary information to render
--   a plot. Common 'LensLike's used for the axis (see haddock's
--   instances for a more comprehensive list):
--
--   * 'axisStyle'    - customise the 'AxisStyle'
--   * 'legend'       - customise the 'Legend'
--   * 'colourBar'    - customise the 'ColourBar'
--   * 'currentPlots' - current plots in the 'Axis'
--   * 'finalPlots'   - changes to the plots just before rendering
--   * 'axes'         - changes to each 'SingleAxis'
--
--          * 'xAxis' - the x-axis
--          * 'yAxis' - the y-axis
--          * 'zAxis' - the z-axis
--
--   The following 'LensLike's can be used on the on all the axes by
--   applying it the to 'Axis' or can be used on a 'SingleAxis' by using
--   it in combination with a specific axis (like 'xAxis').
--
--   * 'axisLabel'   - customise the 'MinorTicks'
--   * 'tickLabel'   - customise the 'TickLabels'
--   * 'minorTicks'  - customise the 'MinorTicks'
--   * 'majorTicks'  - customise the 'MajorTicks'
--   * 'gridLines'   - customise the 'GridLines'
--   * 'axisLine'    - customise the 'AxisLine'
--   * 'axisScaling' - customise the 'AxisScaling'
--
--   Plots are usually added to the axis using specific functions for
--   those plots ('Plots.Types.Line.linePlot, 'Plots.Types.Bar.barPlot').
--   These functions use 'addPlotable' to add the plot to the axis.
data Axis b c n = Axis
  { _axisStyle   :: AxisStyle b (BaseSpace c) n
  , _colourBar   :: ColourBar b n
  , _colourBarR  :: (n,n)
  , _legend      :: Legend b n
  , _axisTitle   :: Title b (BaseSpace c) n
  -- , _axisTitle      :: AxisTitle

  , _axisPlots   :: [DynamicPlot b (BaseSpace c) n]
  , _plotModifier :: Endo (StyledPlot b (BaseSpace c) n)

  -- the v in each axis is only used for the style
  , _axes        :: c (SingleAxis b (BaseSpace c) n)
  } deriving Typeable

-- | Lens onto the separate axes of an axis. Allows changing the
--   coordinate system as long as the 'BaseSpace' is the same.
--
-- @
-- 'axes' :: 'Lens'' ('Axis' b c n) (c ('SingleAxis' b v n))
-- @
axes :: (v ~ BaseSpace c, v ~ BaseSpace c')
     => Lens (Axis b c  n)
             (Axis b c' n)
             (c  (SingleAxis b v n))
             (c' (SingleAxis b v n))
axes = lens _axes (\(Axis a1 a2 a3 a4 a5 a6 a7 _) a8 -> Axis a1 a2 a3 a4 a5 a6 a7 a8)

-- | The list of plots currently in the axis.
axisPlots :: BaseSpace c ~ v => Lens' (Axis b c n) [DynamicPlot b v n]
axisPlots = lens _axisPlots (\a ps -> a {_axisPlots = ps})

-- | Traversal over the current plots in the axis.
--
--   For example, to make all 'ScatterPlot's currently in the axis use a
--   'connectingLine', you can write
--
-- @
-- 'finalPlots' . 'connectingLine' .= 'True'
-- @
currentPlots :: BaseSpace c ~ v => Traversal' (Axis b c n) (DynamicPlot b v n)
currentPlots = axisPlots . traversed

-- | Setter over the final plot before the axis is rendered.
--
--   For example, to make all 'ScatterPlot's in the axis use a
--   'connectingLine' (both currently in the axis and ones added later),
--   you can add
--
-- @
-- 'finalPlots' . 'connectingLine' .= 'True'
-- @
--
finalPlots :: BaseSpace c ~ v => Setter' (Axis b c n) (StyledPlot b v n)
finalPlots = sets $ \f a -> a {_plotModifier = _plotModifier a <> Endo f}

-- | Lens onto the modifier set by 'finalPlots'. This gets applied to
--   all plots in the axis, just before they are rendered.
plotModifier :: BaseSpace c ~ v => Lens' (Axis b c n) (Endo (StyledPlot b v n))
plotModifier = lens _plotModifier (\a f -> a {_plotModifier = f})

-- Axis instances ------------------------------------------------------

type instance V (Axis b v n) = BaseSpace v
type instance N (Axis b v n) = n

instance (Applicative f, Traversable c) => HasTicks f (Axis b c n) where
  bothTicks = axes . traverse . bothTicks

instance (Applicative f, Traversable c) => HasMajorTicks f (Axis b c n) where
  majorTicks = axes . traverse . majorTicks

instance (Applicative f, Traversable c) => HasMinorTicks f (Axis b c n) where
  minorTicks = axes . traverse . minorTicks

instance (Applicative f, Traversable c) => HasGridLines f (Axis b c n) where
  gridLines = axes . traverse . gridLines

instance (Applicative f, Traversable c) => HasMajorGridLines f (Axis b c n) where
  majorGridLines = axes . traverse . majorGridLines

instance (Applicative f, Traversable c) => HasMinorGridLines f (Axis b c n) where
  minorGridLines = axes . traverse . minorGridLines

instance (Applicative f, Traversable c) => HasAxisLabel f (Axis b c n) b where
  axisLabel = axes . traverse . axisLabel

instance (Applicative f, Traversable c) => HasTickLabels f (Axis b c n) b where
  tickLabel = axes . traverse . tickLabel

instance (Applicative f, Traversable c) => HasAxisScaling f (Axis b c n) where
  axisScaling = axes . traverse . axisScaling

instance Settable f => HasPlotOptions f (Axis b c n) b where
  plotOptions = finalPlots . plotOptions

instance Settable f => HasPlotStyle f (Axis b c n) b where
  plotStyle = finalPlots . plotStyle

instance HasLegend (Axis b c n) b where
  legend = lens _legend (\a l -> a {_legend = l})

instance HasTitle (Axis b c n) b where
  title = lens _axisTitle (\a t -> a {_axisTitle = t})

-- | The size used for the rendered axis.
axisSize :: (HasLinearMap c, Num n, Ord n) => Lens' (Axis b c n) (SizeSpec c n)
axisSize = axes . column renderSize . iso mkSizeSpec getSpec -- column axisScaling . asSizeSpec -- iso mkSizeSpec getSpec

-- | The range used for the colour bar limits. This is automatically set
--   when using 'heatMap' or 'heatMap''
colourBarRange :: Lens' (Axis b v n) (n,n)
colourBarRange = lens _colourBarR (\a r -> a {_colourBarR = r})

instance HasAxisStyle (Axis b v n) b where
  axisStyle = lens _axisStyle (\a sty -> a {_axisStyle = sty})

instance HasColourBar (Axis b v n) b where
  colourBar = lens _colourBar (\a cb -> a {_colourBar = cb})

-- Axis functions ------------------------------------------------------

-- $plotable
-- The 'Plotable' class defines ways of converting the data type to a
-- diagram for some axis. There are several variants for adding an axis
-- with constraints @('InSpace' v n a, 'Plotable' a b)@:
--
-- @
-- 'addPlotable'  ::           a -> 'PlotState' a b -> 'AxisState' b v n
-- 'addPlotable''   ::           a ->                  'AxisState' b v n
-- @
--
-- The last argument is a 'PlotState' so you can use @do@ notation to
-- make adjustments to the plot. The @L@ suffix stands for \"legend\",
-- it is equivalent of using 'addLegendEntry' in the 'PlotState'. Since
-- legend entries are so common it has it's own suffix. The following
-- are equivalent:
--
-- @
-- myaxis = 'r2Axis' &~ do
--   'addPlotable'' myplot $ do
--     'addLegendEntry' "my plot"
-- @
--
-- Most of the time you won't use these functions directly. However,
-- other plotting functions follow this naming convention where instead
-- of @a@, it takes the data needed to make the plot.

-- | Add a 'Plotable' 'Plot' to an 'Axis'.
addPlot
  :: (InSpace (BaseSpace c) n p, MonadState (Axis b c n) m, Plotable p b)
  => Plot p b -- ^ the plot
  -> m ()     -- ^ add plot to the 'Axis'
addPlot p = axisPlots <>= [DynamicPlot p]

-- | Add something 'Plotable' to the 'Axis' with a stateful modification
--   of the 'Plot'.
addPlotable
  :: (InSpace (BaseSpace c) n p, MonadState (Axis b c n) m, Plotable p b)
  => p -- ^ the raw plot
  -> State (Plot p b) () -- ^ changes to the plot
  -> m () -- ^ add plot to the 'Axis'
addPlotable p s = addPlot $ execState s (mkPlot p)

-- | Simple version of 'AddPlotable' without any changes 'Plot'.
addPlotable'
  :: (InSpace (BaseSpace v) n p, MonadState (Axis b v n) m, Plotable p b)
  => p    -- ^ the raw plot
  -> m () -- ^ add plot to the 'Axis'
addPlotable' p = addPlotable p (return ())

------------------------------------------------------------------------
-- Predefined axes
------------------------------------------------------------------------

-- | The default axis for plots in the 'V2' coordinate system.
r2Axis
  :: (TypeableFloat n,
     Renderable (Text n) b,
     Renderable (Path V2 n) b)
  => Axis b V2 n
r2Axis = Axis
  { _axisStyle  = fadedColours
  , _colourBar  = defColourBar
  , _colourBarR = (0,1)
  , _axisTitle  = def

  , _legend       = def
  , _axisPlots    = []
  , _plotModifier = mempty

  , _axes = pure def
  }

-- The x-axis ----------------------------------------------------------

-- | Lens onto the x-axis of an 'Axis'.
xAxis :: R1 c => Lens' (Axis b c n) (SingleAxis b (BaseSpace c) n)
xAxis = axes . _x

-- | The label for the x-axis. Shorthand for @'xAxis' . 'axisLabelText'@.
xLabel :: R1 c => Lens' (Axis b c n) String
xLabel = xAxis . axisLabelText

-- | The minimum x value for the axis. If the value if 'Nothing' (the
--   'Default'), the bounds will be infered by the plots in the axis.
xMin :: R1 c => Lens' (Axis b c n) (Maybe n)
xMin = xAxis . boundMin

-- | The minimum x value for the axis. If the value if 'Nothing' (the
--   'Default'), the bounds will be infered by the plots in the axis.
xMax :: R1 c => Lens' (Axis b c n) (Maybe n)
xMax = xAxis . boundMax

-- The y-axis ----------------------------------------------------------

-- | Lens onto the y-axis of an 'Axis'.
yAxis :: R2 c => Lens' (Axis b c n) (SingleAxis b (BaseSpace c) n)
yAxis = axes . _y

-- | The label for the y-axis. Shorthand for @'yAxis' . 'axisLabelText'@.
yLabel :: R2 c => Lens' (Axis b c n) String
yLabel = yAxis . axisLabelText

-- | The minimum y value for the axis. If the value if 'Nothing' (the
--   'Default'), the bounds will be infered by the plots in the axis.
yMin :: R2 c => Lens' (Axis b c n) (Maybe n)
yMin = yAxis . boundMin

-- | The minimum y value for the axis. If the value if 'Nothing' (the
--   'Default'), the bounds will be infered by the plots in the axis.
yMax :: R2 c => Lens' (Axis b c n) (Maybe n)
yMax = yAxis . boundMax

-- The z-axis ----------------------------------------------------------

-- | Lens onto the z-axis of an 'Axis'.
zAxis :: R3 c => Lens' (Axis b c n) (SingleAxis b (BaseSpace c) n)
zAxis = axes . _z

-- | The label for the z-axis. Shorthand for @'zAxis' . 'axisLabelText'@.
zLabel :: R3 c => Lens' (Axis b c n) String
zLabel = zAxis . axisLabelText

-- | The minimum z value for the axis. If the value if 'Nothing' (the
--   'Default'), the bounds will be infered by the plots in the axis.
zMin :: R3 c => Lens' (Axis b c n) (Maybe n)
zMin = zAxis . boundMin

-- | The minimum z value for the axis. If the value if 'Nothing' (the
--   'Default'), the bounds will be infered by the plots in the axis.
zMax :: R3 c => Lens' (Axis b c n) (Maybe n)
zMax = zAxis . boundMax

-- The r-axis ----------------------------------------------------------

-- | Lens onto the radial axis of an 'Axis'.
rAxis :: Radial c => Lens' (Axis b c n) (SingleAxis b (BaseSpace c) n)
rAxis = axes . _radial

-- | The label for the radial axis. Shorthand for @'rAxis' . 'axisLabelText'@.
rLabel :: Radial c => Lens' (Axis b c n) String
rLabel = rAxis . axisLabelText

-- | The minimum z value for the axis. If the value if 'Nothing' (the
--   'Default'), the bounds will be infered by the plots in the axis.
-- rMin :: R3 c => Lens' (Axis b c n) (Maybe n)
-- rMin = zAxis . boundMin

-- | The minimum radial value for the axis. If the value if 'Nothing'
--   (the 'Default'), the bounds will be infered by the plots in the
--   axis.
rMax :: Radial c => Lens' (Axis b c n) (Maybe n)
rMax = rAxis . boundMax

-- The theta-axis ------------------------------------------------------

-- | Lens onto the radial axis of an 'Axis'.
thetaAxis :: Circle c => Lens' (Axis b c n) (SingleAxis b (BaseSpace c) n)
thetaAxis = axes . el etheta

-- | The label for the radial axis. Shorthand for @'rAxis' . 'axisLabelText'@.
thetaLabel :: Circle c => Lens' (Axis b c n) String
thetaLabel = thetaAxis . axisLabelText

-- R3 Axis

-- instance (TypeableFloat n, Enum n, Renderable (Text n) b, Renderable (Path V2 n) b)
--     => Default (Axis b V3 n) where
--   def = Axis
--           { _axisTitle      = Nothing
--           , _axisSize       = mkWidth 300
--           , _axisPlots      = []
--           , _axisLegend     = def
--           , _axisTheme      = coolTheme
--           , _axisLinearMap  = isometricProjection
--           , _axisAxisBounds = Bounds $ pure def
--           , _axisGridLines  = pure def
--           , _axisLabels     = pure def
--           , _axisScaling    = pure def
--           , _axisTickLabels = pure def
--           , _axisTicks      = pure def
--           , _axisLines      = pure def
--           }

-- R3 Axis

polarAxis
  :: (TypeableFloat n,
      Renderable (Text n) b,
      Renderable (Path V2 n) b)
  => Axis b Polar n
polarAxis = Axis
  { _axisStyle  = fadedColours
  , _colourBar  = defColourBar
  , _colourBarR = (0,1)
  , _axisTitle  = def

  , _legend       = def
  , _axisPlots    = []
  , _plotModifier = mempty

  , _axes = pure def
  }


module Widget.Total
  ( Total()
  , mkTotal
  ) where

import Import
import File
import Format.SplitTime
import Format.Percentile

import Statistics.Distribution.Empirical

data Total = Total FileFormat (Zipper EmpiricalDistribution) (Zipper NominalDiffTime) Window

instance Widget Total where
  update                 = updateTotal
  handle                 = handleTotal
  redraw                 = redrawTotal
  window (Total _ _ _ w) = w

mkTotal :: FileFormat -> [EmpiricalDistribution] -> Window -> Total
mkTotal f p = Total f (listToZipper p) (listToZipper [0])

updateTotal :: NominalDiffTime -> Total -> Curses Total
updateTotal d (Total f c z w) = updateWindow w $ do
  let nz = replace (curs z + d) z
      nt = Total f c nz w
      str = show $ splitTime (sum $ zipperToList nz)

  (rows, columns) <- windowSize

  moveCursor 1 (columns - 14)
  drawString $
    if null $ left z
      then "----"
      else if null $ right c
             then show $ percentile (curs c) (sum $ zipperToList z)
             else show $ percentile (curs $ prev c) (sum $ left z)

  moveCursor 1 (columns - (fromIntegral $ length str))
  insertString str

  return nt

handleTotal :: TimerAction -> Total -> Curses Total
handleTotal Reset   (Total f c _ w) = return $ mkTotal f (zipperToList c) w
handleTotal Advance (Total f c z w) = return $ Total f (next c) (push 0 z) w
handleTotal Reverse (Total f c z w) =
  return $
    if null $ left z
      then Total f (prev c) z w
      else
        let cz = curs z
            z' = trunc $ prev z
            nz = replace (curs z' + cz) z'
        in Total f (prev c) nz w
handleTotal Skip    (Total f c z w) = return $ Total f (next c) (push 0 z) w
handleTotal _ t = return t

redrawTotal :: Total -> Curses ()
redrawTotal t@(Total f c z w) = updateWindow w $ do
  (rows, columns) <- windowSize

  erase

  -- Top horizontal line
  -- This is the only thing in row 0
  moveCursor 0 0
  drawLineH Nothing columns

  -- Current total
  -- Tag
  moveCursor 1 (columns - 22)
  drawString "Total:"

  -- Percentile
  moveCursor 1 (columns - 14)
  insertString $
    if null $ left z
      then "----"
      else if null $ right c
             then show $ percentile (curs c) (sum $ zipperToList z)
             else show $ percentile (curs $ prev c) (sum $ left z)

  -- Current cumulative time over all splits
  do
     let str = show $ splitTime $ sum $ zipperToList z

     moveCursor 1 (columns - (fromIntegral $ length str))
     insertString str

  -- The following are consistent internal to a run
  let target = if null $ right c
                 then curs c
                 else last $ right c

  -- Personal Best
  -- Tag
  moveCursor 2 (columns - 30)
  drawString "Personal Best:"

  do
    let pb  = personalBest f
        str = denoteZero $ splitTime pb
        p   = percentile target pb

    -- Percentile
    moveCursor 2 (columns - 14)
    insertString $ ignoreHigh p

    -- Current personal best
    moveCursor 2 (columns - (fromIntegral $ length str))
    insertString str

  -- Draw Sum of Bests
  -- Tag
  moveCursor 3 (columns - 29)
  drawString "Sum of Bests:"

  do
    let sob = sumOfBests f
        str = denoteZero $ splitTime $ sob
        p   = percentile target sob

    -- Percentile
    moveCursor 3 (columns - 14)

    if (sob > 0)
      then insertString $ ignoreHigh p
      else insertString "----"

    -- Current sum of bests
    moveCursor 3 (columns - (fromIntegral $ length str))
    insertString str

  -- Various vertical lines
  moveCursor 0 (columns - 15)
  drawGlyph glyphPlus

  moveCursor 1 (columns - 15)
  drawLineV Nothing 3

  moveCursor 0 (columns - 10)
  drawGlyph glyphPlus

  moveCursor 1 (columns - 10)
  drawLineV Nothing 3


{-# LANGUAGE NoMonomorphismRestriction #-}
module Cargo(Cargo, Market,
  randomCargo, randomMarket,
  showCargo, showMarket,
  fromMarket, fromCargo,
  withdraw,
  tradeScreen, takeScreen)
where

import System.Random
import Text.Printf
import Data.List
import Control.Monad
import Control.Monad.State as State

import qualified Data.Edison.Assoc.StandardMap as M
import Graphics.Rendering.OpenGL as OpenGL
import Graphics.UI.SDL as SDL
import Graphics.Rendering.FTGL as FTGL

import Statistics
import TextScreen
import Space
import SDLUtils

type Cargo = M.FM String Int
type Market = M.FM String (Int, Int)

cargonames = sort ["Grain", "Fruit", "Gem stones", "Firearms"]

randomCargo :: IO String
randomCargo = chooseIO cargonames

showCargo :: Cargo -> String
showCargo c | M.null c  = "No cargo\n"
            | otherwise = concatMap (\(k, v) -> printf "%-20s-%4d\n" k v) (M.toOrdSeq c)

showMarket :: Market -> String
showMarket m = title ++ infos
  where title = printf "%-16s%8s%8s\n" "Good" "Quantity" "Price"
        infos = concatMap (\(n, (q, p)) -> printf "%-20s%10d%10d\n" n q p) (M.toOrdSeq m)

numCargoItems = length cargonames

randomMarket :: IO Market
randomMarket = do
  let names = cargonames
  prices <- replicateM numCargoItems $ randomRIO (5, 15)
  quantities <- replicateM numCargoItems $ randomRIO (0, 30)
  return $ M.fromSeq $ zip names (zip quantities prices)

fitCargo :: Cargo -> [(String, Int)]
fitCargo c = map (\s -> (s, M.lookupWithDefault 0 s c)) cargonames

cargoToMarket :: Cargo -> Market
cargoToMarket cargo =
  foldl' (\acc x -> M.insert x (M.lookupWithDefault 0 x cargo, 0) acc) M.empty cargonames

fromMarket :: Int -> String -> Market -> Market
fromMarket q = toMarket (-q)

fromCargo :: Int -> String -> Cargo -> Cargo
fromCargo q = toCargo (-q)

toMarket :: Int -> String -> Market -> Market
toMarket q = M.adjust (\(q', p) -> (q' + q, p))

toCargo :: Int -> String -> Cargo -> Cargo
toCargo q n = M.insertWith (+) n q

showMarketAndCargo :: Bool -> Market -> Cargo -> String
showMarketAndCargo t m c = title ++ infos
  where title = if t
                  then printf "%-14s%9s%6s%6s\n" "Good" "Quantity" "Price" "Cargo"
                  else printf "%-17s%-12s%-9s\n" "Good" "Captured" "Cargo"
        infos = if t
                  then concatMap (\((n, (q, p)), (_, q')) -> printf "%-14s%9d%6d%6d\n" n q p q') (zip (M.toOrdSeq m) (fitCargo c))
                  else concatMap (\((n, (q, _)), (_, q')) -> printf "%-17s%-12d%-9d\n" n q q') (zip (M.toOrdSeq m) (fitCargo c))

type TradeState = (Market, Cargo, Int, Int)
type TakeState = (Cargo, Cargo, Int, Int)

buy :: Int -> String -> StateT TradeState IO ()
buy q n = do
  (market, cargo, cash, holdspace) <- State.get
  let mval = M.lookupM n market
  case mval of
    Nothing      -> return ()
    Just (q', p) -> do
      let totalq = minimum [max 0 q', q, holdspace]
      let totalp = totalq * p
      if totalp > cash || totalq == 0
        then return ()
        else do
          State.put (fromMarket totalq n market,
               toCargo totalq n cargo,
               subtract totalp cash,
               holdspace - totalq)

sell :: Int -> String -> StateT TradeState IO ()
sell q n = do
  (_, cargo, _, _) <- State.get
  let mval = M.lookupM n cargo
  case mval of
    Nothing -> return ()
    Just q' -> buy (negate (min q' q)) n

tradeScreen :: String -> Font -> Font -> StateT TradeState IO ()
tradeScreen = screenGeneric True

withdraw :: Int -> String -> StateT TradeState IO ()
withdraw q n = withStateT (\(m, c, _, hold) -> (m, c, maxBound, hold)) (buy q n)

screenGeneric :: Bool -> String -> Font -> Font -> StateT TradeState IO ()
screenGeneric trade str f1 f2 = do
  (market, _, _, _) <- State.get
  let exitb = ((100, 100), (100, 30))
      buybuttons  = map (\i -> ((550, 440 - 50 * fromIntegral i), (100, 30))) [1..numCargoItems]
      sellbuttons = map (\i -> ((680, 440 - 50 * fromIntegral i), (100, 30))) [1..numCargoItems]
      buyactions  = map (\(n, (_, _)) -> buy  1 n >> return Nothing) (M.toOrdSeq market)
      sellactions = map (\(n, (_, _)) -> sell 1 n >> return Nothing) (M.toOrdSeq market)
      allbuttons  = exitb : (buybuttons ++ sellbuttons)
      allactions  = return (Just ()) : (buyactions ++ sellactions)
      bttoaction  = zip allbuttons allactions
  let handleInput = do
        events <- liftIO $ pollAllSDLEvents
        let mbutton = mouseClickInAny height [ButtonLeft] allbuttons events
        case mbutton of
          Nothing -> return Nothing
          Just n  -> case lookup n bttoaction of
                       Just act -> act
                       Nothing  -> return Nothing
  loopTextScreen (do (market', cargo, cash, holdspace) <- State.get
                     liftIO $ makeTextScreen (10, 500)
                               [(f1, Color4 1.0 1.0 1.0 1.0, str),
                                (f2, Color4 1.0 1.0 0.0 1.0, showMarketAndCargo trade market' cargo),
                                (f1, Color4 1.0 1.0 1.0 1.0, if trade then "Cash: " ++ show cash else ""),
                                (f1, Color4 1.0 1.0 1.0 1.0, "Hold space: " ++ show holdspace)]
                               (drawButton (Just ("Exit", f1)) exitb >>
                                mapM_ (drawButton (Just (if trade then "Buy" else "Take", f1))) buybuttons >>
                                mapM_ (drawButton (Just (if trade then "Sell" else "Leave", f1))) sellbuttons))
                 handleInput

takeScreen :: String -> Font -> Font -> StateT TakeState IO ()
takeScreen str f1 f2 = do
  (stuff, cargo, cash, hold) <- State.get
  (_, cargo', _, hold') <- liftIO $ execStateT (screenGeneric False str f1 f2) (cargoToMarket stuff, cargo, cash, hold)
  State.put (stuff, cargo', cash, hold')


{-# LANGUAGE ViewPatterns #-}

module Main where

import Codec.Image.DevIL
import Data.Array.Unboxed
import Math.Probably.MCMC
import Math.Probably.Sampler
import Math.Probably.RandIO
import Math.Probably.FoldingStats
import Control.Monad.State.Strict
import System.Cmd
import System.Environment
import Data.Array.IO
import System.IO

import Data.Array.Unboxed
import Data.Array.Base
import qualified Data.Vector.Unboxed as V
import qualified Data.StorableVector as SV
import qualified Data.StorableVector.Base as SVB
import Numeric.LinearAlgebra hiding (find)
import qualified Math.Probably.PDF as PDF
import Data.List
import Data.Maybe
import Control.Applicative
import Data.Ord
import System.Directory

import Data.IORef
import System.IO.Unsafe

import CVUtils
import Edge
import CVhelp
import Foreign.C.String
import Foreign.C.Types
import Unsafe.Coerce

bgImRef :: IORef Image
bgImRef = unsafePerformIO $ newIORef undefined

visibleMat :: IORef BitImage
visibleMat = unsafePerformIO $ newIORef undefined

nogoMat :: IORef BitImage
nogoMat = unsafePerformIO $ newIORef undefined

inNogoMat x y = unsafePerformIO $ do
    nm <- readIORef nogoMat
    return $ readBitImage nm x y

triRef :: IORef Triangles
triRef = unsafePerformIO $ newIORef undefined

gloTris = unsafePerformIO $ readIORef triRef

sdv = 0.5
sdrot = 0.4
sddiv = 2
sddivRot = 4
sdlen = 0.1
sdSideDisp = 0.3

hiddenflat = 1

nparticles = 1000

gauss2i :: R -> R-> R-> R-> Sampler (R,R)
gauss2i m1 sd1 m2 sd2 = do
   u1 <- unitSample
   u2 <- unitSample
   let l = sqrt(-2*log(u1))
       th = 2*pi*u2
--   return (sqrt(-2*log(u1))*cos(2*pi*u2)*sd1+m1,
--           sqrt(-2*log(u1))*sin(2*pi*u2)*sd2+m2)
   return (l*cos(th)*sd1+m1,
           l*sin(th)*sd2+m2)

evolve :: Int -> Obj -> Sampler (Obj,Obj)
evolve retries o@(Obj vlen side len x y rot) = do
        (nvlen, nside) <- gauss2i vlen sdv 0 sdSideDisp
--        nside <- gaussD 0 sdSideDisp
        (nrot, nlen) <- gauss2i rot sdrot len sdlen
--        nlen <- gaussD len sdlen
        let no = Obj nvlen nside nlen
                                 (x+nvlen*cos nrot+nside*cos(nrot-pi/2))
                                 (y+nvlen*sin nrot+nside*sin(nrot-pi/2)) nrot
        if inNogoMat (round $ posx no) (round $ posy no) && retries > 0
           then evolve (retries-1) o
           else return $ (o,no)

invisible :: [((Int,Int),(Int,Int))]

invisible = [((362,621),(355,654)),
             ((442,649),(443,674)),
             ((866,626),(870,658)),
             ((995,628),(999,657)),
             ((1149,629),(1167,658)),
             ((1003,71),(1010,36)),
             ((422,71),(417,42)),
             ((157,162),(127,158)),
             ((1178,469),(1187,501))]

--[ ((963,626),
--               (999,663))]

visible :: Int -> Int -> Bool
visible x y  = not $ any (\sqr -> within sqr  x  y) invisible

mkVisibleMat :: IO ()
mkVisibleMat = do
    tris <- loadTriangles
    writeIORef triRef tris
    let visIm = mkBitImage 1280 720
                 $ \x -> \y-> any (pointInTriangle (fromList [realToFrac x,
                                                              realToFrac y]))
                                  tris
                              && visible x y
    writeIORef visibleMat visIm
    let nogoIm = mkBitImage 1280 720
                 $ \x -> \y-> any (\sqr-> within sqr (realToFrac x) (realToFrac y)) nogo
                              || (not $ any (pointInTriangle (fromList [realToFrac x, realToFrac y])) tris)
    writeIORef nogoMat nogoIm
    return ()

objCentreInvisible :: Obj -> Bool
objCentreInvisible (Obj _ _ _ cx cy _) = not $ visible (round cx) (round cy)

nogo::  [Square]
nogo =
    [((1128,657.5),(1162,627.5)),
     ((963.25,656),(997,626.5)),
     ((836.5,653.75),(866.75,624.5)),
     ((389.75,619.5),(359.75,648.75)),
     ((162.75,309),(132,278.5)),
     ((158.5,191.25),(127.25,160.75)),
     ((422,72.25),(450.75,40.5)),
     ((757.5,72),(728.5,41.75)),
     ((1002.75,70.5),(974.25,39.75)),
     ((1191,194),(1221,222.75)),
     ((1185,435.25),(1213.5,463.5))]
--  [ ((963,626),
--     (997,663))]

nogoPrior :: Obj -> R
nogoPrior o = sum $ map f nogo
          where  f sqr | within sqr (posx o) (posy o) = -1e100
                       | otherwise = 0

{-nogoObj :: Obj -> Bool
nogoObj o = any (\sqr-> within sqr (posx o) (posy o)) nogo
              || (not $ any (pointInTriangle (fromList [posx o,
                                                        posy o])) gloTris) -}

prevprior :: Obj -> (Obj -> R)
prevprior (Obj vlen side len x y rot)
      (Obj nvlen nside nlen _ _ nrot)
        =   PDF.gaussD vlen (sdv/sddiv) nvlen -- "heavy tailed proposals"
          + PDF.gaussD rot (sdrot/sddivRot) nrot
          + PDF.gaussD len (sdlen/sddiv) nlen
          + PDF.gaussD 0 (sdSideDisp/sddiv) nside
          + PDF.gaussD 5 0.5 nlen
          + uniformLogPdf 3 7 nlen

track :: StaticParams -> String -> Int -> Int -> Int -> Obj -> StateT Seed IO Obj
track sp vidfnm startfrm nframes frameOffset obj0 = do
  let outFnm = fileroot vidfnm ++ ".pos"
  h<- lift $ openFile outFnm WriteMode
  lift $ withCString vidfnm open_video
  res <-  go h (replicate nparticles obj0) [startfrm..startfrm+nframes-1] obj0
  lift $ hClose h
  lift $ close_video
  return res
   where
    go _ objs [] o = return o
    go h objs (i:is) _ = do
           bgIm <- lift $ readIORef bgImRef
           nogomat <- lift $ readIORef nogoMat
--           lift $ putStrLn $"~/cvutils/extract "++vidfnm++" "++show i
--           lift $ system $"~/cvutils/extract "++vidfnm++" "++show i
--           frame <- lift $ readImage "extract.png"

           lift advance
           diffObjs <- sample $ mapM (evolve 10) objs
           --lift $ print2  "pixel1Red = " (frame!(0,0,0))
           vismat <- lift $ readIORef visibleMat
           let wparticles = {-# SCC "wpart" #-} (dropLosers $ particleLike sp bgIm vismat nogomat diffObjs)
           let anyInvisible = {-# SCC "anyvis" #-}any objCentreInvisible $ map fst wparticles
           let npart = if anyInvisible then 5*nparticles else nparticles
           let wparticles' = if anyInvisible then map (\(x,w) -> (x,w/hiddenflat)) wparticles else wparticles

           let smws = sumWeights wparticles'
           let cummSmws = cummWeightedSamples wparticles'
           lift $ print2 "fr,npart=" (i,npart)
           --lift $ print2 "winners= " (map snd wparticles')
           --lift $ hFlush stdout
           nextObjs <- {-# SCC "resample" #-} sample $ sequence
                              $ replicate npart
                              $ do u <- unitSample
                                   return . fst . fromJust $ find ((>=u*smws) . snd) cummSmws
           --lift $ print $ map snd $ take 10 $ reverse $ cummSmws

           let mobj = {-# SCC "mobj" #-}
                   runStat (pure Obj <*> before meanF vellen
                                     <*> before meanF sideDisp
                                     <*> before meanF objlen
                                     <*> before meanF posx
                                     <*> before meanF posy
                                     <*> before meanF rot)
                              nextObjs
           --lift $ putStrLn $ show (i,mobj, snd $ last $  wparticles)
           lift $ hPutStrLn h $ show (i,mobj, snd $ last $  wparticles)
           lift $ hFlush h
           when (i `rem` 50 == 0) $ {-# SCC "upbg" #-} do
             lift $ updateBgIm mobj
             when (i `rem` 100 == 0) $ {-# SCC "outframe" #-} do
               lift $ system $"~/cvutils/extract "++vidfnm++" "++show i
               frame <- lift $ readImage "extract.png"
               bgIm2 <- lift $ readIORef bgImRef
               lift $ writeImage ("bgtrack"++show (i+frameOffset)++".png") bgIm2
               markedIm1 <- lift $ markEllipse sp (mobj) frame
               markedIm <- lift $ markObjsOnImage nextObjs markedIm1
               lift $ writeImage ("frame"++show (i+frameOffset)++".png") markedIm
           go h nextObjs is mobj
           --return () -- $ mobj:rest

cu2w8 :: CUChar -> Word8
cu2w8 = unsafeCoerce

pixval :: Int -> Int -> CUChar -> Word8
pixval x y c = cu2w8 $ pixel_value (fromIntegral x) (fromIntegral (719-y)) c

particleLike :: StaticParams -> Image -> BitImage -> BitImage ->[(Obj,Obj)] -> [(Obj,R)]
particleLike sp@(SP noise len ecc) bgim vismat nogomat objprs = map pL objprs where
  radiusi = 2* ceiling (len*ecc) + 2
  objs = map snd objprs
  xmin = minOn (posx) objs - radiusi
  xmax = maxOn (posx) objs+radiusi
  ymin = minOn (posy) objs-radiusi
  ymax = maxOn posy objs+radiusi
  xys = [(x,y) | x <- [xmin.. xmax], y <- [ymin..ymax]]
  yrng = (ymax-ymin+1)
--  cvec = fromList [20,25,30]
  ifvis x y n = if readBitImage vismat x y then n else 0
  ifInside, ifoutside :: UArray (Int,Int) R
  ifInside = {-# SCC "ifinside" #-} listArray ((xmin,ymin),(xmax,ymax))
               [ifvis x y $ (gaussRnn noise 0.07 $ pixval x y 2)
                           + (gaussRnn noise 0.09 $ pixval x y 1)
                           + (gaussRnn noise 0.12 $ pixval x y 0)
                       | (x,y) <- xys]
  ifoutside = {-# SCC "ifoutside" #-} listArray ((xmin,ymin),(xmax,ymax))
              [ifvis x y $ (gaussW8nn noise (bgim!(y,x,0)) $ pixval x y 2)
                           + (gaussW8nn noise (bgim!(y,x,1)) $ pixval x y 1)
                           + (gaussW8nn noise (bgim!(y,x,2)) $ pixval x y 0)
                       | (x,y) <- xys]

  pL (old,o@(Obj  _ _ len cx cy rot)) = {-# SCC "pL" #-}
    let lenecc=len*ecc
        srot = lenecc*sin rot
        crot = lenecc*cos rot
        f1x = cx+crot
        f1y = cy+srot
        f2x = cx-crot
        f2y = cy-srot
--        f :: (R,R) -> R
        f sm xy@(x, y)
--          | not $ readBitImage vismat x y  = 0
          = {-# SCC "f" #-} if dist  f1x f1y   x  y  + dist  f2x f2y   x  y  < 2 * len
                               then sm+ ifInside!xy
                               else sm+ifoutside!xy
--                               then sm + (V.unsafeIndex) ifInside  ((x-xmin)*yrng+(y-ymin))
 --                              else sm + (V.unsafeIndex) ifoutside ((x-xmin)*yrng+(y-ymin))
--                               then (SVB.unsafeIndex) ifInside
--                               else (SVB.unsafeIndex) ifoutside ((x-xmin)*yrng+(y-ymin))

    in if readBitImage nogomat (round $ cx) (round $ cy)
          then (o,-1e100)
          else (o, prevprior old o + {-# SCC "fsum" #-} (noise * foldl' f 0 xys))
                   --x <- xs,
                   --y <- ys]))

updateBgIm (Obj _ _ _ cx cy _) = do
    im <- readIORef bgImRef
    mutIm <- thaw im
    ((loy,lox,_),(hiy,hix,_)) <- getBounds mutIm
    forM_ [(y,x,c) | x<- [lox..hix],
                     y<- [loy..hiy],
                     c<- [0..2],
                     dist cx cy x y > 30] $ \ix@(y,x,c)-> do
      now <- readArray mutIm ix
      writeArray mutIm ix $ (now `div` 2)  + (pixval x y (rgb2bgr c) `div` 2)
    newbg <- freeze (mutIm::MImage)
    writeIORef bgImRef newbg

rgb2bgr :: Int -> CUChar
rgb2bgr 2 = 0
rgb2bgr 0 = 2
rgb2bgr 1 = 1

{-testCV fvid = do
     system $ "~/cvutils/extract "++fvid++" 0"
     frame <- readImage "extract.png"
     print2 "via png " $ frame!(100,100,0)
     withCString fvid open_video
     advance
     print2 "via c " $ pixel_value 100 (719-100) 2
     print2 "via c " $ pixval 100 100 2
     close_video -}
     --error "foo"

main = do
     ilInit
     bgnm : fvid : (read -> x) : (read -> y) : (read -> rot) : rest <- getArgs
     --testCV fvid

     bgIm <-readImage bgnm
     writeIORef bgImRef bgIm
--     marked <- markBg bgIm
--     writeImage "markbg.png" marked
     mkVisibleMat
     print fvid
     vm <- readIORef visibleMat
     writeVisMat "vismat.png" vm bgIm
     let (frStart,nfr) = case rest of
            [] -> (0,6000)
            s:[] -> (read s, 6000)
            s1:s2:_ -> (read s1,read s2)

     --frame0 <-readImage "extract.png"
     {-marked <- markObjsOnImage [Obj (x,y) 0,
                                Obj (x+1,y) 0,
                                Obj (x,y+1) 0,
                                Obj (x+1,y+1) 0] frame0
     writeImage "marked.png" marked -}
     let --posterior = posteriorV bgIm frame0 (round x,round y)
         --postAndV v = (posteriorV bgIm frame0 (round x,round y) $ fromList v, v)
         sp = (SP 218 5 0.9)
--         rot = negate (pi/7)
         --initialsV = fromList [x,y,218, 6, rot, 0.9]
         initObj =  (Obj 1 rot 6 x y rot)
     --ellim <- markEllipse sp initObj frame0
     --writeImage "markell.png" ellim
     --print $ nogoPrior (Obj 0 rot 992 650 rot)
     --print $ nogoPrior (Obj 0 rot 997 650 rot)
     ----print $ visible 994 650
     --print $ visible 1000 650

     runRIO $ do
--         iniampar <- sample $ initialAdaMet 500 1e-4 posterior initialsV
--         AMPar v _ _ _ <- runAndDiscard 5000 (show . ampPar) iniampar $ adaMet False posterior
--         lift $ print v
--         track (v@> 2) (v @> 3) bgIm fvid 3 (v@>0,v@>1)
         if "%d" `isInfixOf` fvid
            then trackMany sp fvid initObj
            else track sp fvid frStart nfr 0 $ initObj

     return ()

putS s  = lift . putStrLn $ s

trackMany sp fvid initObj = go 0 0 initObj where
  go n frOff obj = do
     let fnm = takeWhile (/='%') fvid ++ show n ++ ".avi"
     putS $ "looking for "++fnm
     exists <- lift $ doesFileExist fnm
     if exists then reallyGo n fnm frOff obj else putS "all done" >> return obj
  reallyGo n fnm frOff obj0= do
     lift $ system $ "~/cvutils/extract "++fnm++" "++show i ++">tempout"
     lns <- lift $ lines `fmap` readFile "tempout"
     let nframes = read $ tail $ dropWhile (/='=') $ head $ filter ("total frames=" `isPrefixOf`) lns
     objNext <- track sp fnm 0 (nframes-2) frOff $ obj0
     go (n+1) (frOff+nframes-2) objNext

-- initial on wl0: (956,641)

--rm marked.png && track mixed.png ~/Dropbox/woodlice/wl0.avi '(956,641)' && eog marked.png
--rm -f frame*.png && sudo cabal install --global && track mixed.png ~/Dropbox/woodlice/wl0.avi 956 641 '-0.448'

--rm -f frame*.png && sudo cabal install --global && track mixed.png ~/Dropbox/woodlice/wl0.avi 1092 518 '-5.12'


module Main
  ( main
  ) where

import           Control.Arrow
import           Control.Monad
import qualified Graphics.UI.GLFW    as GLFW
import           Lens.Micro.Platform ((%~), (&), (.~), (^.))

--import           Numeric.LinearAlgebra
import           Application
import           Cmd
import           Initializable
import           Keyboard
import           Matrix
import           MyState
import           Primitives.Cube
import           Primitives.Square
import           Rotation
import           View

newtype Action = Action
  { runAction :: MyState -> Update MyState
  }

type Update a = State (Cmd Action) a

-- Update
blueGreen Blue  = Green
blueGreen Green = Blue

noRun :: Action
noRun = Action return

run :: (MyState -> Update MyState) -> Action
run = Action

run1 :: (a -> MyState -> Update MyState) -> a -> Action
run1 f a = Action (f a)

run2 :: (a -> b -> MyState -> Update MyState) -> a -> b -> Action
run2 f a b = Action (f a b)

shutdown :: MyState -> Update MyState
shutdown state = do
  doPrint "Will shutdown"
  doShutdown
  return state

swapColor :: MyState -> Update MyState
swapColor state = return $ state & color %~ blueGreen

askTime :: Update ()
askTime = getTime (run timeFail) (run1 gotTime)

cameraMovement timeDiff = do
  keyPress GLFW.Key'W (run $ moveCamera timeDiff forward) noRun
  keyPress GLFW.Key'S (run $ moveCamera timeDiff (negate . forward)) noRun
  keyPress GLFW.Key'A (run $ moveCamera timeDiff leftward) noRun
  keyPress GLFW.Key'D (run $ moveCamera timeDiff (negate . leftward)) noRun
  keyPress GLFW.Key'Space (run $ moveCamera timeDiff upward) noRun
  keyPress GLFW.Key'C (run $ moveCamera timeDiff (negate . upward)) noRun

cubeMovement timeDiff = do
  keyPress GLFW.Key'Up (run $ moveSquare timeDiff yHat) noRun
  keyPress GLFW.Key'Down (run $ moveSquare timeDiff (-yHat)) noRun
  keyPress GLFW.Key'Right (run $ moveSquare timeDiff xHat) noRun
  keyPress GLFW.Key'Left (run $ moveSquare timeDiff (-xHat)) noRun

updateSpacing t = ((sin (t / 2) + 1) * 5) + 2

gotTime :: Double -> MyState -> Update MyState
gotTime t state = do
  askTime
  let timeDiff = t - state ^. lastTime
  cameraMovement timeDiff
  cubeMovement timeDiff
  return $ state & lastTime .~ t & spacing .~ updateSpacing t

timeFail :: MyState -> Update MyState
timeFail state = do
  askTime
  doPrint "Time: failed"
  return state

moveSquare distance dir state =
  return $ state & cubePos %~
  (+ (dir * realToFrac (state ^. movementSpeed) * realToFrac distance))

moveCamera distance dirF state =
  return $ state & camera . cameraPosition %~
  (+ (dirF state * realToFrac (state ^. movementSpeed) * realToFrac distance))

forward state = rotateV (negate zHat) $ mkXYRotation x y
  where
    (x, y) = state ^. camera . cameraRotation

leftward state = negate $ cross (forward state) yHat

upward state = yHat

mouseMoved x y state
  | abs x < 300 && abs y < 300 =
    return $! state & camera . cameraRotation %~
    (first (+ (negate x * 0.0001)) >>> second (+ (negate y * 0.0001)))
  | otherwise = return state

resize' :: Int -> Int -> MyState -> Update MyState
resize' w h state =
  return $ state & aspectRatio .~ fromIntegral w / fromIntegral h

keymap =
  concat
    [ [KeyPressed GLFW.Key'Escape (\_ -> run shutdown)]
    , [ KeyPressed
          GLFW.Key'M
          (\_ -> run (\state -> stopCaptureMouse >> return state))
      ]
    , KeyState <$> [GLFW.Key'Space] <*>
      [GLFW.KeyState'Pressed, GLFW.KeyState'Repeating] <*>
      [\_ -> run swapColor]
    ]

initialAction :: MyState -> Action
initialAction state =
  Action $ \_ -> do
    askTime
    return state

setup = do
  square <- create
  cube <- create
  terrain' <- create
  let state =
        initialState & squares .~ [square] & cubes .~ [cube] & terrain .~
        [terrain']
  return $ initialAction state

initState = do
  keyPresses keymap
  resize $ run2 resize'
  runIO setup
  captureMouse $ run2 mouseMoved
  return initialState

myApplication = Application runAction view

main :: IO ()
main = start DefaultConfig initState myApplication


{-# LANGUAGE ScopedTypeVariables #-}

import           Control.Monad.Primitive
import           Control.Monad
import           Foreign.C.Types
import           Foreign.Ptr
import           Unsafe.Coerce
import           Data.Complex

import qualified Data.Vector.Generic               as VG
import qualified Data.Vector.Generic.Mutable       as VGM
import qualified Data.Vector.Storable              as VS
import qualified Data.Vector.Storable.Mutable      as VSM
import qualified Data.Vector.Fusion.Bundle         as VFS
import qualified Data.Vector.Fusion.Stream.Monadic as VFSM

import           Foreign.Storable.Complex
import           Criterion.Main

import           SDR.FilterInternal
import           SDR.Util
import           SDR.CPUID

theBench :: IO ()
theBench = do
    --Setup
    let size       =  16384
        numCoeffs  =  128
        num        =  size - numCoeffs + 1
        decimation =  4
        interpolation = 3
        numCoeffsDiv2  =  64

        coeffsList :: [Float]
        coeffsList = take numCoeffs [0 ..]
        coeffs    :: VS.Vector Float
        coeffs    =  VG.fromList $ take numCoeffs [0 ..]
        coeffsSym    :: VS.Vector Float
        coeffsSym    =  VG.fromList $ take numCoeffsDiv2 [0 ..]
        inBuf     :: VS.Vector Float
        inBuf     =  VG.fromList $ take size [0 ..]
        inBufComplex :: VS.Vector (Complex Float)
        inBufComplex =  VG.fromList $ take size $ do
            i <- [0..]
            return $ i :+ i
        inBufRTLSDR  :: VS.Vector CUChar
        inBufRTLSDR  =  VG.fromList $ take size [0 ..]
        inBufBladeRF :: VS.Vector CShort
        inBufBladeRF =  VG.fromList $ take size [0 ..]

        numConv   = 16386
        inBufConv :: VS.Vector CUChar
        inBufConv = VG.fromList $ take size $ concat $ repeat [0 .. 255]

        duplicate :: [a] -> [a]
        duplicate = concatMap func
            where func x = [x, x]

        coeffs2 :: VS.Vector Float
        coeffs2 =  VG.fromList $ duplicate $ take numCoeffs [0 ..]

    outBuf        :: VS.MVector RealWorld Float <- VGM.new size
    outBufComplex :: VS.MVector RealWorld (Complex Float) <- VGM.new size

    info <- getCPUInfo

    let hasFeatures :: [(CPUInfo -> Bool, a)] -> [a]
        hasFeatures = map snd . filter (($ info) . fst)

    resampler3 <- resampleCRR2   interpolation decimation coeffsList
    resampler4 <- resampleCSSERR interpolation decimation coeffsList
    resampler5 <- resampleCAVXRR interpolation decimation coeffsList

    resampler3C <- resampleCRC    interpolation decimation coeffsList
    resampler4C <- resampleCSSERC interpolation decimation coeffsList
    resampler5C <- resampleCAVXRC interpolation decimation coeffsList

    --Benchmarks
    defaultMain [
            bgroup "filter" [
                bgroup "real" $ hasFeatures [
                    (const True, bench "highLevel"   $ nfIO $ filterHighLevel          coeffs    num inBuf outBuf),
                    (const True, bench "imperative1" $ nfIO $ filterImperative1        coeffs    num inBuf outBuf),
                    (const True, bench "imperative2" $ nfIO $ filterImperative2        coeffs    num inBuf outBuf),
                    (const True, bench "c"           $ nfIO $ filterCRR                coeffs    num inBuf outBuf),
                    (hasSSE42,   bench "cSSE"        $ nfIO $ filterCSSERR             coeffs    num inBuf outBuf),
                    (hasSSE42,   bench "cSSESym"     $ nfIO $ filterCSSESymmetricRR    coeffsSym num inBuf outBuf),
                    (hasAVX,     bench "cAVX"        $ nfIO $ filterCAVXRR             coeffs    num inBuf outBuf),
                    (hasAVX,     bench "cAVXSym"     $ nfIO $ filterCAVXSymmetricRR    coeffsSym num inBuf outBuf)
                ],
                bgroup "complex" $ hasFeatures [
                    (const True, bench "highLevel"   $ nfIO $ filterHighLevel          coeffs    num inBufComplex outBufComplex),
                    (const True, bench "c"           $ nfIO $ filterCRC                coeffs    num inBufComplex outBufComplex),
                    (hasSSE42,   bench "cSSE"        $ nfIO $ filterCSSERC             coeffs2   num inBufComplex outBufComplex),
                    (hasSSE42,   bench "cSSE2"       $ nfIO $ filterCSSERC2            coeffs    num inBufComplex outBufComplex),
                    (hasSSE42,   bench "cSSESym"     $ nfIO $ filterCSSESymmetricRC    coeffsSym num inBufComplex outBufComplex),
                    (hasAVX,     bench "cAVX"        $ nfIO $ filterCAVXRC             coeffs2   num inBufComplex outBufComplex),
                    (hasAVX,     bench "cAVX2"       $ nfIO $ filterCAVXRC2            coeffs    num inBufComplex outBufComplex),
                    (hasAVX,     bench "cAVXSym"     $ nfIO $ filterCAVXSymmetricRC    coeffsSym num inBufComplex outBufComplex)
                ]
            ],
            bgroup "decimate" [
                bgroup "real" $ hasFeatures [
                    (const True, bench "highLevel"   $ nfIO $ decimateHighLevel        decimation coeffs    (num `quot` decimation) inBuf outBuf),
                    (const True, bench "c"           $ nfIO $ decimateCRR              decimation coeffs    (num `quot` decimation) inBuf outBuf),
                    (hasSSE42,   bench "cSSE"        $ nfIO $ decimateCSSERR           decimation coeffs    (num `quot` decimation) inBuf outBuf),
                    (hasSSE42,   bench "cSSESym"     $ nfIO $ decimateCSSESymmetricRR  decimation coeffsSym (num `quot` decimation) inBuf outBuf),
                    (hasAVX,     bench "cAVX"        $ nfIO $ decimateCAVXRR           decimation coeffs    (num `quot` decimation) inBuf outBuf),
                    (hasAVX,     bench "cAVXSym"     $ nfIO $ decimateCAVXSymmetricRR  decimation coeffsSym (num `quot` decimation) inBuf outBuf)
                ],
                bgroup "complex" $ hasFeatures [
                    (const True, bench "highLevel"   $ nfIO $ decimateHighLevel        decimation coeffs    (num `quot` decimation) inBufComplex outBufComplex),
                    (const True, bench "c"           $ nfIO $ decimateCRC              decimation coeffs    (num `quot` decimation) inBufComplex outBufComplex),
                    (hasSSE42,   bench "cSSE"        $ nfIO $ decimateCSSERC           decimation coeffs2   (num `quot` decimation) inBufComplex outBufComplex),
                    (hasSSE42,   bench "cSSE2"       $ nfIO $ decimateCSSERC2          decimation coeffs    (num `quot` decimation) inBufComplex outBufComplex),
                    (hasSSE42,   bench "cSSESym"     $ nfIO $ decimateCSSESymmetricRC  decimation coeffsSym (num `quot` decimation) inBufComplex outBufComplex),
                    (hasAVX,     bench "cAVX"        $ nfIO $ decimateCAVXRC           decimation coeffs2   (num `quot` decimation) inBufComplex outBufComplex),
                    (hasAVX,     bench "cAVX2"       $ nfIO $ decimateCAVXRC2          decimation coeffs    (num `quot` decimation) inBufComplex outBufComplex),
                    (hasAVX,     bench "cAVXSym"     $ nfIO $ decimateCAVXSymmetricRC  decimation coeffsSym (num `quot` decimation) inBufComplex outBufComplex)
                ]
            ],
            bgroup "resample" [
                bgroup "real" $ hasFeatures [
                    (const True, bench "highLevel"   $ nfIO $ resampleHighLevel        interpolation decimation coeffs 0 (num `quot` decimation) inBuf outBuf),
                    (const True, bench "c"           $ nfIO $ resampleCRR              (num `quot` decimation) interpolation decimation 0 coeffs inBuf outBuf),
                    (const True, bench "c2"          $ nfIO $ resampler3               (num `quot` decimation) 0 inBuf outBuf),
                    (hasSSE42,   bench "cSSE"        $ nfIO $ resampler4               (num `quot` decimation) 0 inBuf outBuf),
                    (hasAVX,     bench "cAVX"        $ nfIO $ resampler5               (num `quot` decimation) 0 inBuf outBuf)
                ],
                bgroup "complex" $ hasFeatures [
                    (const True, bench "highLevel"   $ nfIO $ resampleHighLevel        interpolation decimation coeffs 0 (num `quot` decimation) inBufComplex outBufComplex),
                    (const True, bench "c"           $ nfIO $ resampler3C               (num `quot` decimation) 0 inBufComplex outBufComplex),
                    (const True, bench "SSE"         $ nfIO $ resampler4C               (num `quot` decimation) 0 inBufComplex outBufComplex),
                    (const True, bench "AVX"         $ nfIO $ resampler5C               (num `quot` decimation) 0 inBufComplex outBufComplex)
                ]
            ],
            bgroup "scaling" $ hasFeatures [
                (const True, bench "c"               $ nfIO $ scaleC    0.3 inBuf outBuf),
                (hasSSE42,   bench "cSSE"            $ nfIO $ scaleCSSE 0.3 inBuf outBuf),
                (hasAVX,     bench "cAVX"            $ nfIO $ scaleCAVX 0.3 inBuf outBuf)
            ],
            bgroup "conversion" [
                bgroup "RTLSDR" $ hasFeatures [
                    (const True, bench "h"    $ nf (interleavedIQUnsigned256ToFloat :: VS.Vector CUChar -> VS.Vector (Complex Float)) inBufRTLSDR),
                    (const True, bench "c"    $ nf interleavedIQUnsignedByteToFloat    inBufRTLSDR),
                    (hasSSE42,   bench "cSSE" $ nf interleavedIQUnsignedByteToFloatSSE inBufRTLSDR),
                    (hasAVX2,    bench "cAVX" $ nf interleavedIQUnsignedByteToFloatAVX inBufRTLSDR)
                ],
                bgroup "BladeRF" $ hasFeatures [
                    (const True, bench "h"    $ nf (interleavedIQSigned2048ToFloat :: VS.Vector CShort -> VS.Vector (Complex Float))  inBufBladeRF),
                    (const True, bench "c"    $ nf interleavedIQSignedWordToFloat    inBufBladeRF),
                    (hasSSE42,   bench "cSSE" $ nf interleavedIQSignedWordToFloatSSE inBufBladeRF),
                    (hasAVX2,    bench "cAVX" $ nf interleavedIQSignedWordToFloatAVX inBufBladeRF)
                ]
            ]
        ]

main = theBench


{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
module Models.Rover
  (
    Rover (..)
  , getPolicyInput
  , nnVars
  , nnVarsNaive
  ) where

import Models.RoverDomain
import Models.Location
import Models.State

import qualified Data.Map as Map
import Data.List (sortOn)
import Numeric.LinearAlgebra.HMatrix (Vector, fromList, (!))

import NN.NeuralNetwork

import System.IO.Unsafe

data Rover = Rover
             {
               _roverState :: State
             , _roverNet   :: ((Network Double), NNVars)
             , _roverUuid :: Int
             }
           | Traitor
             {
               _roverState :: State
             , _roverNet :: ((Network Double), NNVars)
             , _roverUuid :: Int
             }

nnVars :: NNVars
nnVars = Map.fromList [("numberInputs", 12)
                      ,("numberHidden", 13)
                      ,("numberOutputs", 2)
                      ,("timesToTrain", 3)
                      ,("learningRate", 0.8)
                      ,("sigmoidOrTanh", 0)
                      ,("randomLowerBound", (-10))
                      ,("randomUpperBound", 10)
                      ,("mutationRate", 0.1)
                      ,("gaussianStdDev",1)
                      ,("gaussianMean",0)]

nnVarsNaive :: NNVars
nnVarsNaive = Map.fromList [("numberInputs", 8)
                           ,("numberHidden", 13)
                           ,("numberOutputs", 2)
                           ,("timesToTrain", 3)
                           ,("learningRate", 0.8)
                           ,("sigmoidOrTanh", 0)
                           ,("randomLowerBound", (-10))
                           ,("randomUpperBound", 10)
                           ,("mutationRate", 0.1)
                           ,("gaussianStdDev",1)
                           ,("gaussianMean",0)]

instance Actor Rover where
  getState (Rover s _ _) = s
  getState (Traitor s _ _) = s
  setState (Rover _ n i) s = Rover s n i
  setState (Traitor _ n i) s = Traitor s n i
  getID (Rover _ _ i) = i
  getID (Traitor _ _ i) = i
  getMove = getNextMove
  move r cmd = setState r (moveState (getState r) cmd)
  resetActor (x,y) g r = let b = (fromIntegral x, fromIntegral y)
                             (g', s) = getRandomState g b
                         in (g', setState r s)

instance Agent Rover where
  setPolicy (Rover s (_,v) i) n = Rover s (n,v) i
  setPolicy (Traitor s (_,v) i) n = Traitor s (n,v) i
  getPolicy (Rover _ n _) = fst n
  getPolicy (Traitor _ n _) = fst n
  getFitness d r = let s = getGlobalScore d in if isLoyal r then s else (-1) * s
  isLoyal (Rover _ _ _) = True
  isLoyal (Traitor _ _ _) = False

instance Show Rover where
  show (Rover _ _ i) = "R" ++ (show i)
  show (Traitor _ _ i) = "T" ++ (show i)

instance Eq Rover where
  (==) r1 r2 = getID r1 == getID r2

getNextMove domain agent =
  let cmd = toCmd $ get (getVars agent) (getPolicy agent) (getPolicyInput domain agent)
      currentState = getState agent
  in if inBounds domain (moveState currentState cmd) then cmd
     else getBackupMove domain currentState cmd

getPolicyInput dom a = let as = getState a
                           likeMe = if isLoyal a then isLoyal else not . isLoyal
                           numInputs = getNumInputs a

                           agentStates = map getState $ (getAgents dom)
                           otherStates = map getState $ filter (not . likeMe) (getAgents dom)
                           scoreStates = map getState (getScoring dom)
                           actorStates = filter (/= as) agentStates

                           quadActor = map (getQuad as) actorStates
                           quadScore = map ((+) 4 . getQuad as) scoreStates
                           quadOther = map ((+) 8 . getQuad as) otherStates

                           actorSWithQuad = zip actorStates quadActor
                           scoreSWithQuad = zip scoreStates quadScore
                           otherSWithQuad = zip otherStates quadOther
                           all = actorSWithQuad ++ scoreSWithQuad ++ otherSWithQuad

                           allOverR = map (\(a',q) -> (1.0 / (stateDistance as a'), q)) all
                           blanks = zip (replicate numInputs 0.0) [1..numInputs]
                           allSorted = sortOn snd (allOverR ++ blanks)
                       in fromList $ take numInputs $ map fst $ compress allSorted

-- given a list sorted by the second tuple value, if the second values are equal for
-- consecutive elements, sum the fst value
compress :: (Num a, Eq b) => [(a, b)] -> [(a, b)]
compress [] = []
compress (x:[]) = [x]
compress (x:y:zs) = if snd x == snd y then compress $ (fst x + fst y, snd x):zs
                    else x:(compress (y:zs))

getNumInputs :: Rover -> Int
getNumInputs r = round $ (getVars r) Map.! "numberInputs"

getVars (Rover _ (_,v) _) = v
getVars (Traitor _ (_,v) _) = v

-- difference in locations, as vector, then get angle with atan2 y' x'.
-- Rotate angle by x's orientation backwards, then compare to sectors
getBackupMove dom s (Move dx dy) =
  if inBounds dom (moveState s (Move dx 0)) then Move dx 0
  else if inBounds dom (moveState s (Move 0 dy)) then Move 0 dy
       else Turn (pi / 2)
getBackupMove _ _ _ = Turn (pi / 2)

-- Assuming tanh activation function, output will be between -1 and 1
-- In that case, take values directly as x and y and normalize to unit vector
-- Test: Don't normalize
toCmd :: Vector Double -> Cmd
toCmd vec = let (x,y) = (vec ! 0, vec ! 1)
            in Move x y


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-|
Module      : Grenade.Layers.Concat
Description : Concatenation layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental

This module provides the concatenation layer, which runs two chilld layers in parallel and combines their outputs.
-}
module Grenade.Recurrent.Layers.ConcatRecurrent (
    ConcatRecurrent (..)
  ) where

import           Data.Serialize

import           Data.Singletons
import           GHC.TypeLits

#if MIN_VERSION_base(4,9,0)
import           Data.Kind                    (Type)
#endif

import           Grenade.Core
import           Grenade.Recurrent.Core

import           Numeric.LinearAlgebra.Static (R, split, ( # ))

-- | A Concatentating Layer.
--
-- This layer shares it's input state between two sublayers, and concatenates their output.
--
-- With Networks able to be Layers, this allows for very expressive composition of complex Networks.
--
-- The Concat layer has a few instances, which allow one to flexibly "bash" together the outputs.
--
-- Two 1D vectors, can go to a 2D shape with 2 rows if their lengths are identical.
-- Any 2 1D vectors can also become a longer 1D Vector.
--
-- 3D images become 3D images with more channels. The sizes must be the same, one can use Pad
-- and Crop layers to ensure this is the case.
data ConcatRecurrent :: Shape -> Type -> Shape -> Type -> Type where
  ConcatRecLeft  :: !x -> !y -> ConcatRecurrent m (Recurrent x) n (FeedForward y)
  ConcatRecRight :: !x -> !y -> ConcatRecurrent m (FeedForward x) n (Recurrent y)
  ConcatRecBoth  :: !x -> !y -> ConcatRecurrent m (Recurrent x) n   (Recurrent y)

instance (Show x, Show y) => Show (ConcatRecurrent m (p x) n (q y)) where
  show (ConcatRecLeft x y)  = "ConcatRecLeft\n" ++ show x ++ "\n" ++ show y
  show (ConcatRecRight x y) = "ConcatRecRight\n" ++ show x ++ "\n" ++ show y
  show (ConcatRecBoth x y)  = "ConcatRecBoth\n" ++ show x ++ "\n" ++ show y

instance (RecurrentUpdateLayer x, UpdateLayer y) => UpdateLayer (ConcatRecurrent m (Recurrent x) n (FeedForward y)) where
  type Gradient (ConcatRecurrent m (Recurrent x) n (FeedForward y)) = (Gradient x, Gradient y)
  runUpdate lr (ConcatRecLeft x y) (x', y')
    = ConcatRecLeft (runUpdate lr x x') (runUpdate lr y y')

instance (RandomLayer x, RandomLayer y) => RandomLayer (ConcatRecurrent m (Recurrent x) n (FeedForward y)) where
  createRandomWith m gen
    = ConcatRecLeft <$> createRandomWith m gen <*> createRandomWith m gen

instance (UpdateLayer x, RecurrentUpdateLayer y) => UpdateLayer (ConcatRecurrent m (FeedForward x) n (Recurrent y)) where
  type Gradient (ConcatRecurrent m (FeedForward x) n (Recurrent y)) = (Gradient x, Gradient y)
  runUpdate lr (ConcatRecRight x y) (x', y')
    = ConcatRecRight (runUpdate lr x x') (runUpdate lr y y')

instance (RandomLayer x, RandomLayer y) => RandomLayer (ConcatRecurrent m (FeedForward x) n (Recurrent y)) where
  createRandomWith m gen
    = ConcatRecRight <$> createRandomWith m gen <*> createRandomWith m gen

instance (RecurrentUpdateLayer x, RecurrentUpdateLayer y) => UpdateLayer (ConcatRecurrent m (Recurrent x) n (Recurrent y)) where
  type Gradient (ConcatRecurrent m (Recurrent x) n (Recurrent y)) = (Gradient x, Gradient y)
  runUpdate lr (ConcatRecBoth x y) (x', y')
    = ConcatRecBoth (runUpdate lr x x') (runUpdate lr y y')

instance (RandomLayer x, RandomLayer y) => RandomLayer (ConcatRecurrent m (Recurrent x) n (Recurrent y)) where
  createRandomWith m gen
    = ConcatRecBoth <$> createRandomWith m gen <*> createRandomWith m gen

instance (RecurrentUpdateLayer x, UpdateLayer y) => RecurrentUpdateLayer (ConcatRecurrent m (Recurrent x) n (FeedForward y)) where
  type RecurrentShape (ConcatRecurrent m (Recurrent x) n (FeedForward y)) = RecurrentShape x

instance (UpdateLayer x, RecurrentUpdateLayer y) => RecurrentUpdateLayer (ConcatRecurrent m (FeedForward x) n (Recurrent y)) where
  type RecurrentShape (ConcatRecurrent m (FeedForward x) n (Recurrent y)) = RecurrentShape y

instance (RecurrentUpdateLayer x, RecurrentUpdateLayer y) => RecurrentUpdateLayer (ConcatRecurrent m (Recurrent x) n (Recurrent y)) where
  type RecurrentShape (ConcatRecurrent m (Recurrent x) n (Recurrent y)) =  RecurrentInputs '[ Recurrent x, Recurrent y ]

instance ( SingI i
         , Layer x i ('D1 m)
         , RecurrentLayer y i ('D1 n)
         , KnownNat o
         , KnownNat m
         , KnownNat n
         , o ~ (m + n)
         , n ~ (o - m)
         , (m <=? o) ~ 'True
         ) => RecurrentLayer (ConcatRecurrent ('D1 m) (FeedForward x) ('D1 n) (Recurrent y)) i ('D1 o) where
  type RecTape (ConcatRecurrent ('D1 m) (FeedForward x) ('D1 n) (Recurrent y)) i ('D1 o) = (Tape x i ('D1 m), RecTape y i ('D1 n))

  runRecurrentForwards (ConcatRecRight x y) s input =
    let (xT, xOut :: S ('D1 m))       = runForwards x input
        (yT, side, yOut :: S ('D1 n)) = runRecurrentForwards y s input
    in case (xOut, yOut) of
        (S1D xOut', S1D yOut') ->
            ((xT, yT), side, S1D (xOut' # yOut'))

  runRecurrentBackwards (ConcatRecRight x y) (xTape, yTape) s (S1D o) =
    let (ox :: R m , oy :: R n) = split o
        (x', xB :: S i)         = runBackwards x xTape (S1D ox)
        (y', side, yB :: S i)   = runRecurrentBackwards y yTape s (S1D oy)
    in  ((x', y'), side, xB + yB)

instance ( SingI i
         , RecurrentLayer x i ('D1 m)
         , Layer y i ('D1 n)
         , KnownNat o
         , KnownNat m
         , KnownNat n
         , o ~ (m + n)
         , n ~ (o - m)
         , (m <=? o) ~ 'True
         ) => RecurrentLayer (ConcatRecurrent ('D1 m) (Recurrent x) ('D1 n) (FeedForward y)) i ('D1 o) where
  type RecTape (ConcatRecurrent ('D1 m) (Recurrent x) ('D1 n) (FeedForward y)) i ('D1 o) = (RecTape x i ('D1 m), Tape y i ('D1 n))

  runRecurrentForwards (ConcatRecLeft x y) s input =
    let (xT, side, xOut :: S ('D1 m)) = runRecurrentForwards x s input
        (yT, yOut :: S ('D1 n))       = runForwards y input
    in case (xOut, yOut) of
        (S1D xOut', S1D yOut') ->
            ((xT, yT), side, S1D (xOut' # yOut'))

  runRecurrentBackwards (ConcatRecLeft x y) (xTape, yTape) s (S1D o) =
    let (ox :: R m , oy :: R n) = split o
        (x', side, xB :: S i)   = runRecurrentBackwards x xTape s (S1D ox)
        (y', yB :: S i)         = runBackwards y yTape (S1D oy)
    in  ((x', y'), side, xB + yB)

instance ( SingI i
         , RecurrentLayer x i ('D1 m)
         , RecurrentLayer y i ('D1 n)
         , Fractional (RecurrentShape x)
         , Fractional (RecurrentShape y)
         , KnownNat o
         , KnownNat m
         , KnownNat n
         , o ~ (m + n)
         , n ~ (o - m)
         , (m <=? o) ~ 'True
         ) => RecurrentLayer (ConcatRecurrent ('D1 m) (Recurrent x) ('D1 n) (Recurrent y)) i ('D1 o) where
  type RecTape (ConcatRecurrent ('D1 m) (Recurrent x) ('D1 n) (Recurrent y)) i ('D1 o) = (RecTape x i ('D1 m), RecTape y i ('D1 n))

  runRecurrentForwards (ConcatRecBoth x y) (sx :~@+> (sy :~@+> RINil)) input =
    let (xT, s'x, xOut :: S ('D1 m)) = runRecurrentForwards x sx input
        (yT, s'y, yOut :: S ('D1 n)) = runRecurrentForwards y sy input
    in case (xOut, yOut) of
        (S1D xOut', S1D yOut') ->
            ((xT, yT), (s'x :~@+> (s'y :~@+> RINil)), S1D (xOut' # yOut'))

  runRecurrentBackwards (ConcatRecBoth x y) (xTape, yTape) (sx :~@+> (sy :~@+> RINil)) (S1D o) =
    let (ox :: R m , oy :: R n) = split o
        (x', s'x, xB :: S i)    = runRecurrentBackwards x xTape sx (S1D ox)
        (y', s'y, yB :: S i)    = runRecurrentBackwards y yTape sy (S1D oy)
    in  ((x', y'), (s'x :~@+> (s'y :~@+> RINil)), xB + yB)

instance (Serialize a, Serialize b) => Serialize (ConcatRecurrent sa (FeedForward a) sb (Recurrent b)) where
  put (ConcatRecRight a b) = put a *> put b
  get = ConcatRecRight <$> get <*> get

instance (Serialize a, Serialize b) => Serialize (ConcatRecurrent sa (Recurrent a) sb (Recurrent b)) where
  put (ConcatRecBoth a b) = put a *> put b
  get = ConcatRecBoth <$> get <*> get

instance (Serialize a, Serialize b) => Serialize (ConcatRecurrent sa (Recurrent a) sb (FeedForward b)) where
  put (ConcatRecLeft a b) = put a *> put b
  get = ConcatRecLeft <$> get <*> get

-------------------- GNum instances --------------------

instance (GNum x, GNum y) => GNum (ConcatRecurrent m (p x) n (q y)) where
  n |* (ConcatRecLeft x y) = ConcatRecLeft (n |* x) (n |* y)
  n |* (ConcatRecRight x y) = ConcatRecRight (n |* x) (n |* y)
  n |* (ConcatRecBoth x y) = ConcatRecBoth (n |* x) (n |* y)
  (ConcatRecLeft x1 y1) |+ (ConcatRecLeft x2 y2) = ConcatRecLeft (x1 |+ x2) (y1 |+ y2)
  (ConcatRecRight x1 y1) |+ (ConcatRecRight x2 y2) = ConcatRecRight (x1 |+ x2) (y1 |+ y2)
  (ConcatRecBoth x1 y1) |+ (ConcatRecBoth x2 y2) = ConcatRecBoth (x1 |+ x2) (y1 |+ y2)