{-# LANGUAGE FlexibleContexts #-}
-- {-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE LambdaCase #-}
module Main where

import           Common                         ( Analysis(..)
                                                , Path(..)
                                                )
import           Control.Monad                  ( foldM
                                                , forM
                                                , forM_
                                                , replicateM
                                                , zipWithM_
                                                )
import           Control.Monad.Except           ( runExceptT )
import           Data.Either                    ( rights )
import           Data.List                      ( unzip4 )
import qualified Data.List                     as L
import           Data.Maybe                     ( catMaybes
                                                , listToMaybe
                                                , mapMaybe
                                                )
import qualified Data.Vector                   as V
import           GHC.Float                      ( int2Double )
import qualified Graphics.Matplotlib           as Plt
import qualified GreedyParser                  as Greedy
import           Inference.Conjugate            ( Hyper
                                                , HyperRep
                                                , Prior(sampleProbs)
                                                , Probs
                                                , ProbsRep
                                                , Trace
                                                , evalTraceLogP
                                                , evalTracePredLogP
                                                , getPosterior
                                                , jeffreysPrior
                                                , traceTrace
                                                , uniformPrior
                                                )
import           Musicology.Pitch               ( SPitch )
import qualified Numeric.Log                   as Log
import           PVGrammar                      ( Edge
                                                , PVAnalysis
                                                , loadAnalysis
                                                , loadInput
                                                , topEdges
                                                )
import           PVGrammar.Parse                ( protoVoiceEvaluator
                                                , pvCountUnrestricted
                                                )
import           PVGrammar.Prob.Simple          ( PVParams(PVParams)
                                                , observeDerivation'
                                                , sampleDerivation'
                                                )
import qualified Parser
import qualified Statistics.Sample             as Stats
import           System.FilePath                ( (<.>)
                                                , (</>)
                                                )
import qualified System.FilePattern            as FP
import qualified System.FilePattern.Directory  as FP
import qualified System.Random.MWC.Probability as MWC
import           System.Random.Stateful         ( StatefulGen
                                                , initStdGen
                                                , newIOGenM
                                                )
import           Text.Pretty.Simple             ( CheckColorTty(..)
                                                , OutputOptions(..)
                                                , defaultOutputOptionsNoColor
                                                , pPrintOpt
                                                )
import           Text.Printf                    ( printf )

prettyPrint :: (Show a) => a -> IO ()
prettyPrint = pPrintOpt
  NoCheckColorTty
  (defaultOutputOptionsNoColor { outputOptionsCompact = True })

dataDir :: FilePath
dataDir = "data/"

main = mainLearn

nBaselines :: Int
nBaselines = 100 -- default: 100

nSamples :: Int
nSamples = 250 -- default: 250

mainLearn :: IO ()
mainLearn = do
  -- initialize
  genPure <- initStdGen
  gen     <- newIOGenM genPure
  genMWC  <- MWC.create -- uses a fixed seed
  let prior = uniformPrior @PVParams
  -- load data
  articleExamples <- loadDir
    (dataDir </> "theory-article")
    ["05b_cello_prelude_1-4", "09a_hinunter", "03_bwv784_pattern"]
  Just bwv939 <- loadItem (dataDir </> "bach" </> "fünf-kleine-präludien")
                          "BWV_0939"
  Just bwv940 <- loadItem (dataDir </> "bach" </> "fünf-kleine-präludien")
                          "BWV_0940"
  let dataset = bwv939 : bwv940 : articleExamples
  -- let dataset = take 3 articleExamples
  putStrLn "list of pieces:"
  forM_ dataset $ \(name, _ana, _trace, _surface) -> do
    putStrLn $ "  " <> name
  -- compute overall posterior
  posteriorTotal <- learn prior dataset
  prettyPrint posteriorTotal
  -- cross validation
  let splits = leaveOneOut dataset
  crossPerps <- forM splits (comparePerNote gen genMWC prior)
  let (logpPriors, logpPosts, counts, basemeans, baselines) =
        L.unzip5 crossPerps
      count           = sum counts
      logppnPrior     = sum logpPriors / count
      logppnTrained   = sum logpPosts / count
      logppnBaseline  = sum basemeans / count
      -- look at each split
      logppns         = zipWith (/) logpPosts counts -- fmap (\(_, logp, count, _, _) -> logp / count) crossPerps
      baselinelogppns = zipWith (\bs n -> (/ n) <$> bs) baselines counts
      testpieces      = fst <$> splits
      testnames       = (\(name, _, _, _) -> name) <$> testpieces
  showSplits logppns baselinelogppns counts testnames
      -- nsplits         = int2Double $ length splits
      -- meanlogppn      = sum logppns / nsplits
  putStrLn $ "prior logppn: " <> show logppnPrior
  putStrLn $ "prior perppn: " <> show (exp $ negate logppnPrior)
  putStrLn $ "overall trained logppn: " <> show logppnTrained
  putStrLn $ "overall trained perppn: " <> show (exp $ negate logppnTrained)
  putStrLn $ "baseline logppn: " <> show logppnBaseline
  putStrLn $ "baseline perppn: " <> show (exp $ negate logppnBaseline)
  putStrLn $ "prior-posterior Δlogppn:" <> show (logppnPrior - logppnTrained)
  putStrLn $ printf "logppn (nats) & %.3f & %.3f & %.3f\\\\"
                    logppnPrior
                    logppnTrained
                    logppnBaseline
  let natsbits = log 2
  putStrLn $ printf "logppn (bits) & %.3f & %.3f & %.3f\\\\"
                    (logppnPrior / natsbits)
                    (logppnTrained / natsbits)
                    (logppnBaseline / natsbits)
  putStrLn $ printf "perppn & %.2f & %.2f & %.2f\\\\"
                    (exp $ negate logppnPrior)
                    (exp $ negate logppnTrained)
                    (exp $ negate logppnBaseline)

testEstimator = do
  (long, short)   <- loadExamples
  articleExamples <- loadDir
    (dataDir </> "theory-article")
    ["05b_cello_prelude_1-4", "09a_hinunter", "03_bwv784_pattern"]
  Just bwv939 <- loadItem (dataDir </> "bach" </> "fünf-kleine-präludien")
                          "BWV_0939"
  Just bwv940 <- loadItem (dataDir </> "bach" </> "fünf-kleine-präludien")
                          "BWV_0940"
  let dataset = bwv939 : bwv940 : articleExamples
  genMWC <- MWC.createSystemRandom
  let prior = uniformPrior @PVParams
  posteriorTotal <- learn prior dataset
  estimates      <- derivationLogProb genMWC posteriorTotal long
  -- let means = meanOfLogs <$> tail (L.inits estimates)
  -- -- print means
  -- Plt.onscreen $ Plt.line [1 .. length means] means
  print $ meanOfLogs estimates

countRareIntDerivs = do
  Just (_, _, _, surface) <- loadItem (dataDir </> "theory-article")
                                      "10c_rare_int"
  count <- Parser.parseSize pvCountUnrestricted surface
  print count

-- loading data
-- ------------

type Piece
  = (String, PVAnalysis SPitch, Trace PVParams, Path [SPitch] [Edge SPitch])

loadItem :: FilePath -> FilePath -> IO (Maybe Piece)
loadItem dir name = do
  ana <- loadAnalysis (dir </> name <.> "analysis.json")
  case ana of
    Left  _err -> pure Nothing
    Right a    -> if anaTop a == PathEnd topEdges
      then do
        surface <- loadInput (dir </> name <.> "musicxml")
        case observeDerivation' (anaDerivation a) of
          Left _err -> do
            putStrLn $ "could not observe trace for " <> name <> ", skipping."
            pure Nothing
          Right trace -> pure $ Just (name, a, trace, surface)
      else do
        putStrLn $ "derivation for " <> name <> " is incomplete, skipping."
        pure Nothing

loadDir :: FilePath -> [String] -> IO [Piece]
loadDir dir exclude = do
  files <- FP.getDirectoryFiles dir ["*.analysis.json"]
  let getName file = FP.match "*.analysis.json" file >>= listToMaybe
      names =
        -- exclude duplicats
        filter (`L.notElem` exclude) $ catMaybes $ getName <$> files
  -- print names
  items <- mapM (loadItem dir) names
  pure $ catMaybes items

loadExamples :: IO (Trace PVParams, Trace PVParams)
loadExamples = do
  Just (_, _, bwv940Trace, _) <- loadItem
    (dataDir </> "bach" </> "fünf-kleine-präludien")
    "BWV_0940"
  Just (_, _, rareTrace, _) <- loadItem (dataDir </> "theory-article")
                                        "10c_rare_int"
  pure (bwv940Trace, rareTrace)

-- learning
-- --------

learn :: Hyper PVParams -> [Piece] -> IO (Hyper PVParams)
learn = foldM train
 where
  train prior (name, _, trace, _) =
    case getPosterior prior trace sampleDerivation' of
      Nothing -> do
        putStrLn $ "couldn't compute posterior for " <> name <> ", skipping."
        pure prior
      Just post -> do
        -- putStrLn $ "learned from " <> name <> "."
        pure post

-- evaluating
-- ----------

leaveOneOut :: [a] -> [(a, [a])]
leaveOneOut dataset = go dataset [] []
 where
  go []       _    splits = splits
  go (x : xs) done splits = go xs (x : done) ((x, xs <> done) : splits)

derivationLogProb
  :: MWC.GenIO -> Hyper PVParams -> Trace PVParams -> IO [Double]
derivationLogProb gen hyper trace = do
  probs <- replicateM nSamples $ MWC.sample (sampleProbs @PVParams hyper) gen
  let estimates = mapMaybe
        (\params -> snd <$> evalTraceLogP params trace sampleDerivation')
        probs
  pure $! estimates

derivationLogProb'
  :: MWC.GenIO -> Hyper PVParams -> Trace PVParams -> IO Double
derivationLogProb' gen hyper trace = do
  estimates <- derivationLogProb gen hyper trace
  pure $! meanOfLogs estimates

meanOfLogs :: (RealFloat b, Foldable t, Functor t) => t b -> b
meanOfLogs logs =
  Log.ln $! Log.sum (Log.Exp <$> logs) / fromIntegral (length logs)

countNotes :: Path [a] b -> Int
countNotes (PathEnd notes       ) = length notes
countNotes (Path notes edges rst) = length notes + countNotes rst

sampleBaselines
  :: (StatefulGen g IO)
  => g
  -> MWC.GenIO
  -> Hyper PVParams
  -> Piece
  -> IO [Double]
sampleBaselines gen genMWC posterior (name, _, _, surface) = do
  derivsTry <- replicateM nBaselines $ runExceptT $ Greedy.parseRandom'
    gen
    protoVoiceEvaluator
    surface
  let derivs = rights derivsTry
  putStrLn $ "  collected " <> show (length derivs) <> " samples for " <> name
  let baselines = flip fmap derivs $ \ana -> do
        trace <- observeDerivation' (anaDerivation ana)
        pure $ derivationLogProb' genMWC posterior trace
  baselines <- traverse sequence baselines -- maaaaagic
  fmap catMaybes $ forM baselines $ \case
    Left  err -> putStrLn err >> pure Nothing
    Right val -> pure $ Just val

summarize :: [Double] -> (Double, Double)
summarize xs = (Stats.mean sample, Stats.stdDev sample)
  where sample = V.fromList xs

comparePerNote
  :: (StatefulGen g IO)
  => g
  -> MWC.GenIO
  -> Hyper PVParams
  -> (Piece, [Piece])
  -> IO (Double, Double, Double, Double, [Double])
comparePerNote gen genMWC prior (test@(tstName, _, tstTrace, tstSurface), train)
  = do
    putStrLn $ "testing on " <> tstName
    posterior <- learn prior train
    -- compute normalized logprobs
    let nnotes = int2Double $ countNotes tstSurface
    logpPrior <- derivationLogProb' genMWC prior tstTrace
    logpPost  <- derivationLogProb' genMWC posterior tstTrace
    let logppnPrior = logpPrior / nnotes
        logppnPost  = logpPost / nnotes
    baselines <- sampleBaselines gen genMWC posterior test
    let (basemean, basestd) = summarize baselines
        basemeanpn          = basemean / nnotes
        basestdpn           = basestd / nnotes
    putStrLn $ "  nnotes: " <> show nnotes
    putStrLn $ "  logppn prior: " <> show logppnPrior <> " nats"
    putStrLn $ "  logppn posterior: " <> show logppnPost <> " nats"
    putStrLn
      $  "  logppn baseline: mean="
      <> show basemeanpn
      <> ", std="
      <> show basestdpn
    putStrLn $ "  baselines top 5: " <> show
      (drop 95 $ (/ nnotes) <$> L.sort baselines)
    putStrLn $ "  perplexity prior: " <> show (exp $ negate logppnPrior)
    putStrLn $ "  perplexity posterior: " <> show (exp $ negate logppnPost)
    putStrLn $ "  Δlogppn: " <> show (logppnPost - logppnPrior) <> " nats"
    pure (logpPrior, logpPost, nnotes, basemean, baselines)

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

a % b = a Plt.% Plt.mp Plt.# b
infixl 5 %

showSplits
  :: [Double] -> [[Double]] -> [Double] -> [String] -> IO (Either String String)
showSplits logppns blogppns counts testpieces =
  Plt.file "splits.svg"
    $ Plt.readData (logppns, blogppns, counts, testpieces)
    % "import numpy as np"
    % "import pandas as pd"
    % "import seaborn as sns"
    % "(logppns, blogppns, counts, pieces) = tuple(map(np.array, data))"
    % "baselines = pd.concat([pd.DataFrame({'logppn': bppns, 'piece': piece}) for bppns, piece in zip(blogppns, pieces)])"
    % "testscores = pd.DataFrame({'logppn': logppns, 'piece': pieces})"
    % "colors = sns.color_palette()"
    % "fig, ax = plot.subplots(figsize=(10,5))"
    % "sns.boxplot(ax=ax, y='piece', x='logppn', data=baselines, whis=(2,98), color=colors[0], fliersize=4)"
    % "sns.scatterplot(ax=ax, y='piece', x='logppn', data=testscores, color=colors[1])"
    % "ax.invert_yaxis()"
    % "fig.tight_layout()"
    % "fig.savefig('splits.pdf')"
    % "fig.savefig('splits.png')"
