module RL.Plotting where

import Common
import Control.Foldl qualified as Foldl
import Control.Monad (forM_)
import Control.Monad.State qualified as ST
import Data.Colour.Palette.ColorSet
import Display
import Graphics.Rendering.Chart.Backend.Cairo as Plt
import Graphics.Rendering.Chart.Easy ((.=))
import Graphics.Rendering.Chart.Easy qualified as Plt
import Graphics.Rendering.Chart.Gtk qualified as Plt
import Musicology.Pitch (SPitch)
import PVGrammar
import PVGrammar.Generate (derivationPlayerPV)
import RL.ModelTypes
import StrictList qualified as SL

-- helpers
-- -------

mean :: (Foldable t) => t QType -> QType
mean :: forall (t :: * -> *). Foldable t => t QType -> QType
mean = Fold QType QType -> t QType -> QType
forall (f :: * -> *) a b. Foldable f => Fold a b -> f a -> b
Foldl.fold Fold QType QType
forall a. Fractional a => Fold a a
Foldl.mean

zipWithStrict :: (a -> b -> c) -> SL.List a -> SL.List b -> SL.List c
zipWithStrict :: forall a b c. (a -> b -> c) -> List a -> List b -> List c
zipWithStrict a -> b -> c
f List a
SL.Nil List b
_ = List c
forall a. List a
SL.Nil
zipWithStrict a -> b -> c
f List a
_ List b
SL.Nil = List c
forall a. List a
SL.Nil
zipWithStrict a -> b -> c
f (SL.Cons a
x List a
xs) (SL.Cons b
y List b
ys) = c -> List c -> List c
forall a. a -> List a -> List a
SL.Cons (a -> b -> c
f a
x b
y) (List c -> List c) -> List c -> List c
forall a b. (a -> b) -> a -> b
$ (a -> b -> c) -> List a -> List b -> List c
forall a b c. (a -> b -> c) -> List a -> List b -> List c
zipWithStrict a -> b -> c
f List a
xs List b
ys

-- plotting
-- --------

mkHistoryPlot
  :: String
  -> [QType]
  -> ST.StateT
      (Plt.Layout Int QType)
      (ST.State Plt.CState)
      ()
mkHistoryPlot :: FilePath -> [QType] -> StateT (Layout Int QType) (State CState) ()
mkHistoryPlot FilePath
title [QType]
values = do
  [AlphaColour QType] -> StateT (Layout Int QType) (State CState) ()
forall l. [AlphaColour QType] -> EC l ()
Plt.setColors ([AlphaColour QType]
 -> StateT (Layout Int QType) (State CState) ())
-> [AlphaColour QType]
-> StateT (Layout Int QType) (State CState) ()
forall a b. (a -> b) -> a -> b
$ Colour QType -> AlphaColour QType
forall a. Num a => Colour a -> AlphaColour a
Plt.opaque (Colour QType -> AlphaColour QType)
-> [Colour QType] -> [AlphaColour QType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Colour QType
forall a. (Ord a, Floating a) => Colour a
Plt.steelblue]
  (FilePath -> Identity FilePath)
-> Layout Int QType -> Identity (Layout Int QType)
forall x y (f :: * -> *).
Functor f =>
(FilePath -> f FilePath) -> Layout x y -> f (Layout x y)
Plt.layout_title ((FilePath -> Identity FilePath)
 -> Layout Int QType -> Identity (Layout Int QType))
-> FilePath -> StateT (Layout Int QType) (State CState) ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
.= FilePath
title
  EC (Layout Int QType) (PlotLines Int QType)
-> StateT (Layout Int QType) (State CState) ()
forall (p :: * -> * -> *) x y.
ToPlot p =>
EC (Layout x y) (p x y) -> EC (Layout x y) ()
Plt.plot (EC (Layout Int QType) (PlotLines Int QType)
 -> StateT (Layout Int QType) (State CState) ())
-> EC (Layout Int QType) (PlotLines Int QType)
-> StateT (Layout Int QType) (State CState) ()
forall a b. (a -> b) -> a -> b
$ FilePath
-> [[(Int, QType)]] -> EC (Layout Int QType) (PlotLines Int QType)
forall x y l. FilePath -> [[(x, y)]] -> EC l (PlotLines x y)
Plt.line FilePath
title [[(Int, QType)]
points]
 where
  points :: [(Int, QType)]
points = [Int] -> [QType] -> [(Int, QType)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1 :: Int ..] [QType]
values

mkHistoriesPlot
  :: String
  -> [[QType]]
  -> ST.StateT
      (Plt.Layout Int QType)
      (ST.State Plt.CState)
      ()
mkHistoriesPlot :: FilePath
-> [[QType]] -> StateT (Layout Int QType) (State CState) ()
mkHistoriesPlot FilePath
title [[QType]]
series = do
  [AlphaColour QType] -> StateT (Layout Int QType) (State CState) ()
forall l. [AlphaColour QType] -> EC l ()
Plt.setColors ([AlphaColour QType]
 -> StateT (Layout Int QType) (State CState) ())
-> [AlphaColour QType]
-> StateT (Layout Int QType) (State CState) ()
forall a b. (a -> b) -> a -> b
$
    Colour QType -> AlphaColour QType
forall a. Num a => Colour a -> AlphaColour a
Plt.opaque (Colour QType -> AlphaColour QType)
-> [Colour QType] -> [AlphaColour QType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Brightness -> Int -> Colour QType
d3Colors2 Brightness
Dark (Int -> Colour QType) -> [Int] -> [Colour QType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
0 .. Int
9]) [Colour QType] -> [Colour QType] -> [Colour QType]
forall a. [a] -> [a] -> [a]
++ (Brightness -> Int -> Colour QType
d3Colors2 Brightness
Light (Int -> Colour QType) -> [Int] -> [Colour QType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
0 .. Int
9])
  (FilePath -> Identity FilePath)
-> Layout Int QType -> Identity (Layout Int QType)
forall x y (f :: * -> *).
Functor f =>
(FilePath -> f FilePath) -> Layout x y -> f (Layout x y)
Plt.layout_title ((FilePath -> Identity FilePath)
 -> Layout Int QType -> Identity (Layout Int QType))
-> FilePath -> StateT (Layout Int QType) (State CState) ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
.= FilePath
title
  (Maybe LegendStyle -> Identity (Maybe LegendStyle))
-> Layout Int QType -> Identity (Layout Int QType)
forall x y (f :: * -> *).
Functor f =>
(Maybe LegendStyle -> f (Maybe LegendStyle))
-> Layout x y -> f (Layout x y)
Plt.layout_legend ((Maybe LegendStyle -> Identity (Maybe LegendStyle))
 -> Layout Int QType -> Identity (Layout Int QType))
-> Maybe LegendStyle -> StateT (Layout Int QType) (State CState) ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
.= Maybe LegendStyle
forall a. Maybe a
Nothing
  [([QType], Integer)]
-> (([QType], Integer)
    -> StateT (Layout Int QType) (State CState) ())
-> StateT (Layout Int QType) (State CState) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([[QType]] -> [Integer] -> [([QType], Integer)]
forall a b. [a] -> [b] -> [(a, b)]
zip [[QType]]
series [Integer
1 ..]) ((([QType], Integer)
  -> StateT (Layout Int QType) (State CState) ())
 -> StateT (Layout Int QType) (State CState) ())
-> (([QType], Integer)
    -> StateT (Layout Int QType) (State CState) ())
-> StateT (Layout Int QType) (State CState) ()
forall a b. (a -> b) -> a -> b
$ \([QType]
values, Integer
i) -> do
    let points :: [(Int, QType)]
points = [Int] -> [QType] -> [(Int, QType)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1 :: Int ..] [QType]
values
    EC (Layout Int QType) (PlotLines Int QType)
-> StateT (Layout Int QType) (State CState) ()
forall (p :: * -> * -> *) x y.
ToPlot p =>
EC (Layout x y) (p x y) -> EC (Layout x y) ()
Plt.plot (EC (Layout Int QType) (PlotLines Int QType)
 -> StateT (Layout Int QType) (State CState) ())
-> EC (Layout Int QType) (PlotLines Int QType)
-> StateT (Layout Int QType) (State CState) ()
forall a b. (a -> b) -> a -> b
$ FilePath
-> [[(Int, QType)]] -> EC (Layout Int QType) (PlotLines Int QType)
forall x y l. FilePath -> [[(x, y)]] -> EC l (PlotLines x y)
Plt.line FilePath
"" [[(Int, QType)]
points]

mkHistoryPlot'
  :: String
  -> QType
  -> [QType]
  -> ST.StateT
      (Plt.Layout Int QType)
      (ST.State Plt.CState)
      ()
mkHistoryPlot' :: FilePath
-> QType -> [QType] -> StateT (Layout Int QType) (State CState) ()
mkHistoryPlot' FilePath
title QType
target [QType]
values = do
  [AlphaColour QType] -> StateT (Layout Int QType) (State CState) ()
forall l. [AlphaColour QType] -> EC l ()
Plt.setColors ([AlphaColour QType]
 -> StateT (Layout Int QType) (State CState) ())
-> [AlphaColour QType]
-> StateT (Layout Int QType) (State CState) ()
forall a b. (a -> b) -> a -> b
$ Colour QType -> AlphaColour QType
forall a. Num a => Colour a -> AlphaColour a
Plt.opaque (Colour QType -> AlphaColour QType)
-> [Colour QType] -> [AlphaColour QType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Colour QType
forall a. (Ord a, Floating a) => Colour a
Plt.steelblue, Colour QType
forall a. (Ord a, Floating a) => Colour a
Plt.orange]
  (FilePath -> Identity FilePath)
-> Layout Int QType -> Identity (Layout Int QType)
forall x y (f :: * -> *).
Functor f =>
(FilePath -> f FilePath) -> Layout x y -> f (Layout x y)
Plt.layout_title ((FilePath -> Identity FilePath)
 -> Layout Int QType -> Identity (Layout Int QType))
-> FilePath -> StateT (Layout Int QType) (State CState) ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
.= FilePath
title
  EC (Layout Int QType) (PlotLines Int QType)
-> StateT (Layout Int QType) (State CState) ()
forall (p :: * -> * -> *) x y.
ToPlot p =>
EC (Layout x y) (p x y) -> EC (Layout x y) ()
Plt.plot (EC (Layout Int QType) (PlotLines Int QType)
 -> StateT (Layout Int QType) (State CState) ())
-> EC (Layout Int QType) (PlotLines Int QType)
-> StateT (Layout Int QType) (State CState) ()
forall a b. (a -> b) -> a -> b
$ FilePath
-> [[(Int, QType)]] -> EC (Layout Int QType) (PlotLines Int QType)
forall x y l. FilePath -> [[(x, y)]] -> EC l (PlotLines x y)
Plt.line FilePath
title [[(Int, QType)]
points]
  EC (Layout Int QType) (PlotLines Int QType)
-> StateT (Layout Int QType) (State CState) ()
forall (p :: * -> * -> *) x y.
ToPlot p =>
EC (Layout x y) (p x y) -> EC (Layout x y) ()
Plt.plot (EC (Layout Int QType) (PlotLines Int QType)
 -> StateT (Layout Int QType) (State CState) ())
-> EC (Layout Int QType) (PlotLines Int QType)
-> StateT (Layout Int QType) (State CState) ()
forall a b. (a -> b) -> a -> b
$ FilePath
-> [[(Int, QType)]] -> EC (Layout Int QType) (PlotLines Int QType)
forall x y l. FilePath -> [[(x, y)]] -> EC l (PlotLines x y)
Plt.line FilePath
"target" [[(Int
1, QType
target), ([QType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [QType]
values, QType
target)]]
 where
  points :: [(Int, QType)]
points = [Int] -> [QType] -> [(Int, QType)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1 :: Int ..] [QType]
values

mkHistoriesPlot'
  :: String
  -> [QType]
  -> [[QType]]
  -> ST.StateT
      (Plt.Layout Int QType)
      (ST.State Plt.CState)
      ()
mkHistoriesPlot' :: FilePath
-> [QType]
-> [[QType]]
-> StateT (Layout Int QType) (State CState) ()
mkHistoriesPlot' FilePath
title [QType]
targets [[QType]]
series = do
  [AlphaColour QType] -> StateT (Layout Int QType) (State CState) ()
forall l. [AlphaColour QType] -> EC l ()
Plt.setColors ([AlphaColour QType]
 -> StateT (Layout Int QType) (State CState) ())
-> [AlphaColour QType]
-> StateT (Layout Int QType) (State CState) ()
forall a b. (a -> b) -> a -> b
$
    Colour QType -> AlphaColour QType
forall a. Num a => Colour a -> AlphaColour a
Plt.opaque (Colour QType -> AlphaColour QType)
-> [Colour QType] -> [AlphaColour QType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Brightness -> Int -> Colour QType
d3Colors2 Brightness
Dark (Int -> Colour QType) -> [Int] -> [Colour QType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
0 .. Int
9]) [Colour QType] -> [Colour QType] -> [Colour QType]
forall a. [a] -> [a] -> [a]
++ (Brightness -> Int -> Colour QType
d3Colors2 Brightness
Light (Int -> Colour QType) -> [Int] -> [Colour QType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
0 .. Int
9])
  (FilePath -> Identity FilePath)
-> Layout Int QType -> Identity (Layout Int QType)
forall x y (f :: * -> *).
Functor f =>
(FilePath -> f FilePath) -> Layout x y -> f (Layout x y)
Plt.layout_title ((FilePath -> Identity FilePath)
 -> Layout Int QType -> Identity (Layout Int QType))
-> FilePath -> StateT (Layout Int QType) (State CState) ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
.= FilePath
title
  [(QType, [QType], Integer)]
-> ((QType, [QType], Integer)
    -> StateT (Layout Int QType) (State CState) ())
-> StateT (Layout Int QType) (State CState) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([QType] -> [[QType]] -> [Integer] -> [(QType, [QType], Integer)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [QType]
targets [[QType]]
series [Integer
1 ..]) (((QType, [QType], Integer)
  -> StateT (Layout Int QType) (State CState) ())
 -> StateT (Layout Int QType) (State CState) ())
-> ((QType, [QType], Integer)
    -> StateT (Layout Int QType) (State CState) ())
-> StateT (Layout Int QType) (State CState) ()
forall a b. (a -> b) -> a -> b
$ \(QType
target, [QType]
values, Integer
i) -> do
    let points :: [(Int, QType)]
points = [Int] -> [QType] -> [(Int, QType)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1 :: Int ..] [QType]
values
    EC (Layout Int QType) (Plot Int QType)
-> StateT (Layout Int QType) (State CState) ()
forall (p :: * -> * -> *) x y.
ToPlot p =>
EC (Layout x y) (p x y) -> EC (Layout x y) ()
Plt.plot (EC (Layout Int QType) (Plot Int QType)
 -> StateT (Layout Int QType) (State CState) ())
-> EC (Layout Int QType) (Plot Int QType)
-> StateT (Layout Int QType) (State CState) ()
forall a b. (a -> b) -> a -> b
$ do
      color <- EC (Layout Int QType) (AlphaColour QType)
forall l. EC l (AlphaColour QType)
Plt.takeColor
      histLine <- Plt.liftEC $ do
        Plt.plot_lines_values .= [points]
        Plt.plot_lines_title .= show i
        Plt.plot_lines_style . Plt.line_color .= color
      targetLine <- Plt.liftEC $ do
        Plt.plot_lines_values .= [[(1, target), (length values, target)]]
        Plt.plot_lines_style . Plt.line_color .= color
        Plt.plot_lines_style . Plt.line_dashes .= [10, 10]
      pure $ Plt.joinPlot (Plt.toPlot histLine) (Plt.toPlot targetLine)

fileOpts :: Plt.FileOptions
fileOpts :: FileOptions
fileOpts = FileOptions
forall a. Default a => a
Plt.def{_fo_format = Plt.SVG}

showHistory :: String -> [QType] -> IO ()
showHistory :: FilePath -> [QType] -> IO ()
showHistory FilePath
title [QType]
values = Int -> Int -> StateT (Layout Int QType) (State CState) () -> IO ()
forall r.
(Default r, ToRenderable r) =>
Int -> Int -> EC r () -> IO ()
Plt.toWindow Int
60 Int
40 (StateT (Layout Int QType) (State CState) () -> IO ())
-> StateT (Layout Int QType) (State CState) () -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath -> [QType] -> StateT (Layout Int QType) (State CState) ()
mkHistoryPlot FilePath
title [QType]
values

plotHistory :: String -> [QType] -> IO ()
plotHistory :: FilePath -> [QType] -> IO ()
plotHistory FilePath
title [QType]
values = FileOptions
-> FilePath -> StateT (Layout Int QType) (State CState) () -> IO ()
forall r.
(Default r, ToRenderable r) =>
FileOptions -> FilePath -> EC r () -> IO ()
Plt.toFile FileOptions
fileOpts (FilePath
"rl/" FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
title FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
".svg") (StateT (Layout Int QType) (State CState) () -> IO ())
-> StateT (Layout Int QType) (State CState) () -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath -> [QType] -> StateT (Layout Int QType) (State CState) ()
mkHistoryPlot FilePath
title [QType]
values

plotHistories :: String -> [[QType]] -> IO ()
plotHistories :: FilePath -> [[QType]] -> IO ()
plotHistories FilePath
title [[QType]]
values = FileOptions
-> FilePath -> StateT (Layout Int QType) (State CState) () -> IO ()
forall r.
(Default r, ToRenderable r) =>
FileOptions -> FilePath -> EC r () -> IO ()
Plt.toFile FileOptions
fileOpts (FilePath
"rl/" FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
title FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
".svg") (StateT (Layout Int QType) (State CState) () -> IO ())
-> StateT (Layout Int QType) (State CState) () -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
-> [[QType]] -> StateT (Layout Int QType) (State CState) ()
mkHistoriesPlot FilePath
title [[QType]]
values

plotHistory' :: String -> QType -> [QType] -> IO ()
plotHistory' :: FilePath -> QType -> [QType] -> IO ()
plotHistory' FilePath
title QType
target [QType]
values = FileOptions
-> FilePath -> StateT (Layout Int QType) (State CState) () -> IO ()
forall r.
(Default r, ToRenderable r) =>
FileOptions -> FilePath -> EC r () -> IO ()
Plt.toFile FileOptions
fileOpts (FilePath
"rl/" FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
title FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
".svg") (StateT (Layout Int QType) (State CState) () -> IO ())
-> StateT (Layout Int QType) (State CState) () -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
-> QType -> [QType] -> StateT (Layout Int QType) (State CState) ()
mkHistoryPlot' FilePath
title QType
target [QType]
values

plotHistories' :: String -> [QType] -> [[QType]] -> IO ()
plotHistories' :: FilePath -> [QType] -> [[QType]] -> IO ()
plotHistories' FilePath
title [QType]
target [[QType]]
values = FileOptions
-> FilePath -> StateT (Layout Int QType) (State CState) () -> IO ()
forall r.
(Default r, ToRenderable r) =>
FileOptions -> FilePath -> EC r () -> IO ()
Plt.toFile FileOptions
fileOpts (FilePath
"rl/" FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
title FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
".svg") (StateT (Layout Int QType) (State CState) () -> IO ())
-> StateT (Layout Int QType) (State CState) () -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
-> [QType]
-> [[QType]]
-> StateT (Layout Int QType) (State CState) ()
mkHistoriesPlot' FilePath
title [QType]
target [[QType]]
values

plotDeriv :: (Foldable t) => FilePath -> t (Leftmost (Split SPitch) (Freeze SPitch) (Spread SPitch)) -> IO ()
plotDeriv :: forall (t :: * -> *).
Foldable t =>
FilePath
-> t (Leftmost (Split SPitch) (Freeze SPitch) (Spread SPitch))
-> IO ()
plotDeriv FilePath
fn t (Leftmost (Split SPitch) (Freeze SPitch) (Spread SPitch))
deriv = do
  case DerivationPlayer
  (Split SPitch)
  (Freeze SPitch)
  (Spread SPitch)
  (Notes SPitch)
  (Edges SPitch)
-> t (Leftmost (Split SPitch) (Freeze SPitch) (Spread SPitch))
-> Either FilePath (DerivationGraph (Notes SPitch) (Edges SPitch))
forall (t :: * -> *) slc tr s f h.
(Foldable t, Ord slc, Ord tr) =>
DerivationPlayer s f h slc tr
-> t (Leftmost s f h) -> Either FilePath (DerivationGraph slc tr)
replayDerivation DerivationPlayer
  (Split SPitch)
  (Freeze SPitch)
  (Spread SPitch)
  (Notes SPitch)
  (Edges SPitch)
forall n.
(Eq n, Ord n, Notation n, Hashable n, Eq (IntervalOf n),
 HasPitch n) =>
DerivationPlayer
  (Split n) (Freeze n) (Spread n) (Notes n) (Edges n)
derivationPlayerPV t (Leftmost (Split SPitch) (Freeze SPitch) (Spread SPitch))
deriv of
    (Left FilePath
err) -> FilePath -> IO ()
putStrLn FilePath
err
    (Right DerivationGraph (Notes SPitch) (Edges SPitch)
g) -> FilePath -> DerivationGraph (Notes SPitch) (Edges SPitch) -> IO ()
forall slc tr.
(Eq slc, Eq tr, Show slc, Show tr) =>
FilePath -> DerivationGraph slc tr -> IO ()
viewGraph FilePath
fn DerivationGraph (Notes SPitch) (Edges SPitch)
g