{-# LANGUAGE RecordWildCards, ScopedTypeVariables, OverloadedStrings,
             ViewPatterns, ImplicitParams, OverloadedLists, RankNTypes,
             ExistentialQuantification, LambdaCase, NumDecimals, NamedFieldPuns,
             TypeApplications, DeriveGeneric, StandaloneDeriving, FlexibleInstances #-}
{-# OPTIONS_GHC -Wno-orphans #-}

import Prelude hiding (quot, showList)

import Test.Tasty
import Test.Tasty.Golden
import Test.Tasty.Golden.Advanced
import Test.Tasty.HUnit

import Numeric.Sundials

import Numeric.LinearAlgebra as L hiding ((<.>), (<>))
import qualified Data.Vector.Storable as VS
import qualified Data.Vector as V
import Data.List
import Data.Maybe
import Data.Foldable
import Data.IORef
import Control.Monad
import Control.Exception
import Katip
import System.IO
import System.FilePath
import Text.Printf (printf)
import Data.Aeson
import Data.Aeson.Encode.Pretty
import GHC.Generics

----------------------------------------------------------------------
--                            Helpers
----------------------------------------------------------------------

emptyOdeProblem :: OdeProblem
emptyOdeProblem = OdeProblem
      { odeRhs = error "emptyOdeProblem: no odeRhs provided"
      , odeJacobian = Nothing
      , odeInitCond = error "emptyOdeProblem: no odeInitCond provided"
      , odeEvents = mempty
      , odeTimeBasedEvents = TimeEventSpec $ return $ 1.0 / 0.0
      , odeEventHandler = nilEventHandler
      , odeMaxEvents = 100
      , odeSolTimes = error "emptyOdeProblem: no odeSolTimes provided"
      , odeTolerances = defaultTolerances
      }

data OdeSolver = forall method . (Show method, Method method) => OdeSolver
  String -- name
  [method]

availableSolvers :: [OdeSolver]
availableSolvers =
  [ OdeSolver "CVode"  [BDF, ADAMS]
  , OdeSolver "ARKode" [SDIRK_5_3_4, TRBDF2_3_3_2]
  ]

defaultOpts :: method -> ODEOpts method
defaultOpts method = ODEOpts
  { maxNumSteps = 1e5
  , minStep     = 1.0e-14
  , fixedStep   = 0
  , maxFail     = 10
  , odeMethod   = method
  , initStep    = Nothing
  , jacobianRepr = DenseJacobian
  }

defaultTolerances = Tolerances
  { absTolerances = Left 1.0e-6
  , relTolerance = 1.0e-10
  }

checkDiscrepancy :: HasCallStack => Double -> Double -> Assertion
checkDiscrepancy eps diff = assertBool msg $ diff <= eps
  where
    msg = printf "Difference too large: %.2e > %.2e"
      diff eps

fmod :: RealFrac a => a -> a -> a
fmod arg1 arg2 =
  let
    quot = arg1 / arg2
    quot_floor = realToFrac (floor quot :: Integer)
    quot_floor_times_divisor = quot_floor * arg2
  in
    arg1 - quot_floor_times_divisor

-- | Show list without brackets, so that it's more suitable for file names
showList :: Show a => [a] -> String
showList = intercalate "," . map show

-- derive JSON instances for all sundials types for testing
deriving instance Generic SundialsSolution
instance ToJSON SundialsSolution
instance FromJSON SundialsSolution

deriving instance Generic SundialsDiagnostics
instance ToJSON SundialsDiagnostics
instance FromJSON SundialsDiagnostics

deriving instance Generic ErrorDiagnostics
instance ToJSON ErrorDiagnostics
instance FromJSON ErrorDiagnostics

instance ToJSON CrossingDirection
instance FromJSON CrossingDirection

instance ToJSON (Matrix Double) where
  toJSON = toJSON . toLists
instance FromJSON (Matrix Double) where
  parseJSON = fmap fromLists . parseJSON

compareSolutions
  :: Bool -- ^ same method?
  -> SundialsSolution -- ^ expected
  -> SundialsSolution -- ^ got
  -> Maybe String -- ^ maybe error message
compareSolutions same_method a b = asum @[]
  [ do
      guard . not $ VS.length (actualTimeGrid a) == VS.length (actualTimeGrid b)
      return "Different length of actualTimeGrid"
  , do
      guard . not $ norm_Inf (actualTimeGrid a - actualTimeGrid b) < precision
      return "Different values of actualTimeGrid"
  , do
      guard . not $ size (solutionMatrix a) == size (solutionMatrix b)
      return "Different sizes of the solutionMatrix"
  , do
      guard . not $ norm_Inf (solutionMatrix a - solutionMatrix b) < precision
      return "Different values in the solutionMatrix"
  ]
  where
    precision = if same_method then 1e-10 else 1e-1

odeGoldenTest
  :: forall method . (Method method, Show method)
  => Bool -- ^ compare between methods? (via a canonical golden file)
  -> ODEOpts method -- ^ ode options (affect the golden file name)
  -> String -- ^ name (of both the test and the file)
  -> IO (Either ErrorDiagnostics SundialsSolution)
  -> TestTree
odeGoldenTest do_canonical opts name action =
  let
    method_dir = show (methodSolver @method) </> show (odeMethod opts)
  in
    testGroup name $ do
      (dir, type_, same_method) <-
        [(method_dir, "Method-specific", True)] ++
        [("canonical", "Canonical", False) | do_canonical ]
      let
        golden_path = ("tests/golden" </> dir </> name <.> "json")
        cmp expected got = return $
          case (expected, got) of
            (Right{},Left{}) -> Just "expected success, got error"
            (Left{},Right{}) -> Just "expected error, got success"
            (Left a,Left b) ->
              if errorCode a == errorCode b
                then Nothing
                else Just "got a different error"
            (Right a,Right b) -> compareSolutions same_method a b
        upd val = createDirectoriesAndWriteFile golden_path
          (encodePretty' defConfig {confCompare = compare} val)
      return $ goldenTest
        type_ -- test name
        (fromJust <$> decodeFileStrict golden_path) -- get the golden value
        action -- get the tested value
        cmp
        upd

-- | Make an 'EventHandler' from a vector of per-event handlers and
-- parameters
mkEventHandler
  :: V.Vector (Double -> VS.Vector Double -> VS.Vector Double)
  -> V.Vector Bool -- ^ stop the solver?
  -> V.Vector Bool -- ^ record the event?
  -> EventHandler
mkEventHandler handlers stop_solver_vec record_event_vec t y0 evs
  | VS.null evs = error "mkEventHandler: got a time-based event"
  | otherwise = return EventHandlerResult
      { eventStopSolver = or . map (stop_solver_vec V.!) $ VS.toList evs
      , eventRecord = or . map (record_event_vec V.!) $ VS.toList evs
      , eventNewState = foldl' (\y hndl -> hndl t y) y0 . map (handlers V.!) $ VS.toList evs
      }

mkTimeEvents
  :: [(Double, Double -> VS.Vector Double -> VS.Vector Double, Bool, Bool)] -- ^ time-based events (time, update, stop?, record?)
  -> IO (TimeEventSpec, EventHandler)
mkTimeEvents time_based_events = do
  time_based_events_ref <- newIORef time_based_events
  let handler :: EventHandler
      handler t y0 evs =
        if VS.null evs
          then do
            time_evs <- readIORef time_based_events_ref
            case time_evs of
              [] -> throwIO $ ErrorCall "Unexpected time-based event"
              (t1, update, eventStopSolver, eventRecord) : rest ->
                if t == t1
                  then do
                    writeIORef time_based_events_ref rest
                    return EventHandlerResult
                      { eventStopSolver
                      , eventRecord
                      , eventNewState = update t y0
                      }
                  else throwIO $ ErrorCall "Wrong event time"
          else error "mkTimeEvents: got root-based events"

      event_spec :: TimeEventSpec
      event_spec = TimeEventSpec $ do
        time_evs <- readIORef time_based_events_ref
        return $ case time_evs of
          [] -> 1.0/0.0 -- +Inf
          (t,_,_,_) : _ -> t

  return (event_spec, handler)

combineEventHandlers
  :: EventHandler -- ^ root-based event handler
  -> EventHandler -- ^ time-based event handler
  -> EventHandler -- ^ combined handler
combineEventHandlers rh th t y0 evs =
  if VS.null evs
    then th t y0 evs
    else rh t y0 evs

nilEventHandler :: EventHandler
nilEventHandler _ _ _ = throwIO $ ErrorCall "nilEventHandler"

----------------------------------------------------------------------
--                             The tests
----------------------------------------------------------------------

main = do
  handleScribe <- mkHandleScribe ColorIfTerminal stderr (permitItem InfoS) V2
  log_env <- registerScribe "stderr" handleScribe defaultScribeSettings =<<
    initLogEnv "test" "devel"
  let ?log_env = log_env

  defaultMain $ testGroup "Tests" $
    [
      testGroup solver_name
      [ testGroup (show method) $
          let opts = defaultOpts method in
          [ withVsWithoutJacobian opts
          , stiffishTest opts
          , eventTests opts
          , noErrorTests opts
          , discontinuousRhsTest opts
          , modulusEventTest opts
          , cascadingEventsTest opts
          , simultaneousEventsTest opts
          , timeBasedEventTest opts
          , timeGridTest opts
          , timeAliasingTests opts
          ]
      | method <- methods
      ]
    | OdeSolver solver_name methods <- availableSolvers
    ]

noErrorTests opts = testGroup "Absence of error"
  [ testCase name $ do
      r <- runKatipT ?log_env $ solve opts prob
      case r of
        Right _ -> return ()
        Left e -> assertFailure (show e)
  | (name, prob) <- [ empty ]
  ]

stiffishTest opts = odeGoldenTest True opts "Stiffish" $
  runKatipT ?log_env $ solve opts stiffish

withVsWithoutJacobian opts0 = testGroup "With vs without jacobian"
  [ testCase (name ++ " " ++ (if sparse then "sparse" else "dense")) $ do
      let
        opts = opts0
          { jacobianRepr =
              if sparse
                then SparseJacobian spat -- this should be enough for our tests
                else DenseJacobian
          }
      Right (solutionMatrix -> solJac)   <- runKatipT ?log_env $ solve opts prob
      Right (solutionMatrix -> solNoJac) <- runKatipT ?log_env $ solve
        opts { jacobianRepr = DenseJacobian }
        prob { odeJacobian = Nothing }
      checkDiscrepancy 1e-2 $ norm_2 (solJac - solNoJac)
  | ((name, prob), spat) <-
      [ (brusselator, brusselator_sparse_pattern)
      , (robertson, robertson_sparse_pattern)
      ]
  , sparse <- [False, True]
  ]
  where
    brusselator_sparse_pattern = SparsePattern [1,1,1, 1,1,1, 1,0,1]
    robertson_sparse_pattern = SparsePattern [1,1,0, 1,1,1, 1,1,0]

eventTests opts = testGroup "Events"
  [ odeGoldenTest True opts "Exponential" $
      runKatipT ?log_env $ solve opts exponential
  , odeGoldenTest True opts "Robertson" $ do
      let upd _ _ = vector [1.0, 0.0, 0.0]
      runKatipT ?log_env $ solve opts
        (snd robertson)
          { odeEvents =
            [ EventSpec { eventCondition = \_t y -> y ! 0 - 0.0001
                        , eventDirection = AnyDirection
                        }
            , EventSpec { eventCondition = \_t y -> y ! 2 - 0.01
                        , eventDirection = AnyDirection
                        }
            ]
          , odeMaxEvents = 100
          , odeSolTimes = [0,100]
          , odeEventHandler = mkEventHandler
            (V.replicate 2 upd)
            (V.replicate 2 False)
            (V.replicate 2 True)
          }
  , odeGoldenTest True opts "Bounded_sine" $
      runKatipT ?log_env $ solve opts boundedSine
  ]

discontinuousRhsTest opts = odeGoldenTest True opts "Discontinuous_derivative" $
  runKatipT ?log_env $ solve opts discontinuousRHS

modulusEventTest opts0 = localOption (mkTimeout 1e5) $ testGroup "Modulus event"
  [ odeGoldenTest False opts
      (printf "Modulus_recordEvent=%s_initStep=%s"
        (show record_event)
        initStepStr) $
        runKatipT ?log_env $ solve opts (modulusEvent record_event)
        -- the timeout doesn't seem to trigger when the infinite
        -- loop actually happens. I'm not sure why, since we're making safe
        -- calls, which should be interruptible? -fno-omit-yields doesn't
        -- seem to help either. Maybe worth investigating.
  | (initStep, initStepStr::String) <-
      [(Nothing, "Nothing")
      ,(Just 1, "1")
      ,(Just (1 - 2**(-53)), "1-eps")]
  , record_event <- [False, True]
  , let opts = opts0 { initStep }
  ]

cascadingEventsTest opts = odeGoldenTest True opts "Cascading_events" $ do
  runKatipT ?log_env $ solve opts prob
  where
    prob = emptyOdeProblem
      { odeRhs = odeRhsPure $ \_ _ -> [1, 0]
      , odeInitCond = [0, 0]
      , odeEvents = events
      , odeEventHandler = mkEventHandler
          [\_ y -> [7, y ! 1]
          ,\_ y -> [y ! 0, 1]
          ]
          (V.replicate 2 False)
          (V.replicate 2 True)
      , odeMaxEvents = 100
      , odeSolTimes = [0,10]
      }
    events =
      [ EventSpec { eventCondition = \_t y -> y ! 0 - 5
                  , eventDirection = AnyDirection
                  }
      , EventSpec { eventCondition = \_t y -> y ! 0 - 6
                  , eventDirection = AnyDirection
                  }
      ]

simultaneousEventsTest opts = testGroup "Simultaneous events"
  [ odeGoldenTest True opts (printf "Simultaneous_events/maxEvents=%d_eventRecord=%s_stopSolver=%s"
    maxEvents
    (showList record)
    (showList stopSolver)) $ do
      let prob = emptyOdeProblem
            { odeRhs = odeRhsPure $ \_ _ -> [0,0]
            , odeJacobian = Nothing
            , odeInitCond = [0,0]
            , odeEvents = events
            , odeEventHandler = mkEventHandler
                (V.map (\i _ y -> y VS.// [(i,1)]) [0..1])
                (V.fromList stopSolver)
                (V.fromList record)
            , odeMaxEvents = maxEvents
            , odeSolTimes = [0,10]
            , odeTolerances = defaultTolerances
            }
          event =
            EventSpec
              { eventCondition = \t _ -> t - 5
              , eventDirection = AnyDirection
              }
          events = V.replicate 2 event
      runKatipT ?log_env $ solve opts prob
  | maxEvents <- [0..3]
  , record <- replicateM 2 [False,True]
  , stopSolver <- replicateM 2 [False,True]
  ]

-- Four events:
--
-- 1. A time-based event at t = 1.0; recorded
--
-- 2. A time-based event at t = 2.0; not recorded and no update
--
-- 3. A root-based event at t = 3.0
--
-- 4. A time-based event at t = 4.0; stops the solver
timeBasedEventTest opts = odeGoldenTest True opts "Time-based events" $ do
  let
    upd :: Double -> Double -> Vector Double -> Vector Double
    upd x _t y = y VS.// [(1, x)]
  (time_ev_spec, time_ev_handler) <- mkTimeEvents
    [ (1.0, upd 5, False, True)
    , (2.0, const id, False, False)
    , (4.0, upd 8, True, True)
    ]
  let prob = emptyOdeProblem
        { odeRhs = odeRhsPure $ \_ _ -> [1, 0]
        , odeInitCond = [0, 0]
        , odeEvents =
            [ EventSpec
              { eventCondition = \t y -> t/2 + y ! 0 - 4.5
              , eventDirection = Upwards
              }
            ]
        , odeTimeBasedEvents = time_ev_spec
        , odeEventHandler = combineEventHandlers
            (mkEventHandler [upd 13] [False] [True])
            time_ev_handler
        , odeMaxEvents = 100
        , odeSolTimes = [0, 10]
        }
  runKatipT ?log_env $ solve opts prob

-- | A simple test that checks that we do the right thing even when the
-- time grid does not start at 0.
timeGridTest opts = odeGoldenTest True opts "Time grid test" $ do
  runKatipT ?log_env $ solve opts emptyOdeProblem
    { odeRhs = odeRhsPure $ \_ _ -> [1]
    , odeInitCond = [0]
    , odeSolTimes = [-3, -2, 0, 10, 100]
    }

-- | Test time aliasing: when some of the time-based events coincide
-- exactly with time grid points.
--
-- There are two tests: for when events are recorded or not. This
-- corresponds to time-based events with the conditions that either hold or
-- not.
timeAliasingTests opts = testGroup "Time aliasing"
  [ odeGoldenTest True opts ("Time aliasing; record="++show record) $ do
      (time_ev_spec, time_ev_handler) <- mkTimeEvents
        [ (t, if record then \_ y -> [y!0 + 0.2] else const id, False, record)
        | t <- VS.toList ts
        ]
      runKatipT ?log_env $ solve opts emptyOdeProblem
        { odeRhs = odeRhsPure $ \_ _ -> [1]
        , odeInitCond = [0]
        , odeSolTimes = ts
        , odeTimeBasedEvents = time_ev_spec
        , odeEventHandler = time_ev_handler
        }
  | record <- [False, True]
  ]
  where
    ts = [1,2,3]

----------------------------------------------------------------------
--                           ODE problems
----------------------------------------------------------------------

brusselator :: (String, OdeProblem)
brusselator = (,) "brusselator" $ emptyOdeProblem
  { odeRhs = odeRhsPure $ \_t x ->
      let
        u = x VS.! 0
        v = x VS.! 1
        w = x VS.! 2
      in
      [ a - (w + 1) * u + v * u * u
      , w * u - v * u * u
      , (b - w) / eps - w * u
      ]
  , odeJacobian = Just $ \(_t :: Double) x ->
      let
        u = x VS.! 0
        v = x VS.! 1
        w = x VS.! 2
      in (3><3)
      [ (-(w + 1.0)) + 2.0 * u * v, w - 2.0 * u * v, (-w)
      , u * u                     , (-(u * u))     , 0.0
      , (-u)                      , u              , (-1.0) / eps - u
      ]
  , odeEvents = mempty
  , odeEventHandler = nilEventHandler
  , odeMaxEvents = 0
  , odeInitCond = [1.2, 3.1, 3.0]
  , odeSolTimes = [0.0, 0.1 .. 10.0]
  , odeTolerances = defaultTolerances
  }
  where
    a = 1.0
    b = 3.5
    eps :: Fractional a => a
    eps = 5.0e-6

exponential = emptyOdeProblem
  { odeRhs = odeRhsPure $ \_ y -> [y VS.! 0]
  , odeJacobian = Nothing
  , odeInitCond = vector [1]
  , odeEvents = events
  , odeEventHandler = mkEventHandler [\_ _ -> [ 2 ]] [False] [True]
  , odeMaxEvents = 100
  , odeSolTimes = vector [ fromIntegral k / 100 | k <- [0..(22::Int)]]
  , odeTolerances = defaultTolerances
  }
  where
    events =
      [ EventSpec { eventCondition = \_ y -> y ! 0 - 1.1
                  , eventDirection = Upwards
                  }
      ]

robertson = (,) "Robertson" $ emptyOdeProblem
  { odeRhs = odeRhsPure $ \_ (VS.toList -> [y1,y2,y3]) ->
      [ -0.04 * y1 + 1.0e4 * y2 * y3
      , 0.04 * y1 - 1.0e4 * y2 * y3 - 3.0e7 * (y2)^(2 :: Int)
      , 3.0e7 * (y2)^(2 :: Int)
      ]
  , odeJacobian = Just $ \_t (VS.toList -> [_, y2, y3]) -> (3 >< 3)
      [ -0.04, 1.0e4 * y3, 1.0e4 * y2
      , 0.04, -1.0e4*y3 - 3.0e7*2*y2, -1.0e4*y2
      , 0, 3.0e7*2*y2, 0
      ]
  , odeInitCond = [1.0, 0.0, 0.0]
  , odeEvents = []
  , odeEventHandler = nilEventHandler
  , odeMaxEvents = 0
  , odeSolTimes = [0,20,100,1000]
  , odeTolerances = defaultTolerances
  }

empty = (,) "Empty system" $ emptyOdeProblem
  { odeRhs = odeRhsPure $ \_ _ -> []
  , odeJacobian = Nothing
  , odeInitCond = []
  , odeEvents = []
  , odeEventHandler = nilEventHandler
  , odeMaxEvents = 0
  , odeSolTimes = [0,1]
  , odeTolerances = defaultTolerances
  }

stiffish = emptyOdeProblem
  { odeRhs = odeRhsPure $ \t ((VS.! 0) -> u) -> [ lamda * u + 1.0 / (1.0 + t * t) - lamda * atan t ]
  , odeJacobian = Nothing
  , odeInitCond = [0.0]
  , odeEvents = []
  , odeEventHandler = nilEventHandler
  , odeMaxEvents = 0
  , odeSolTimes = [0.0, 0.1 .. 10.0]
  , odeTolerances = defaultTolerances
  }
  where
    lamda = -100.0

-- A sine wave that only changes direction once it reaches ±0.9.
-- Illustrates event-specific reset function
boundedSine = emptyOdeProblem
  { odeRhs = odeRhsPure $ \_t y -> [y VS.! 1, - y VS.! 0]
  , odeJacobian = Nothing
  , odeInitCond = [0,1]
  , odeEvents = events
  , odeEventHandler = mkEventHandler
      [\_ y -> vector [ y ! 0, - abs (y ! 1) ]
      ,\_ y -> vector [ y ! 0, abs (y ! 1) ]
      ]
      (V.replicate 2 False)
      (V.replicate 2 True)
  , odeMaxEvents = 100
  , odeSolTimes = VS.fromList [ 2 * pi * k / 360 | k <- [0..360]]
  , odeTolerances = defaultTolerances
  }
  where
    events =
      [ EventSpec { eventCondition = \_t y -> y ! 0 - 0.9
                  , eventDirection = Upwards
                  }
      , EventSpec { eventCondition = \_t y -> y ! 0 + 0.9
                  , eventDirection = Downwards
                  }
      ]

-- | An example of a system with a discontinuous RHS
discontinuousRHS = emptyOdeProblem
  { odeRhs = odeRhsPure $ \t _ ->
      if t1 <= t && t <= t2
        then [deriv]
        else [0]
  , odeJacobian = Nothing
  , odeInitCond = [0]
  , odeEvents = events
  , odeEventHandler = mkEventHandler
      (V.replicate 2 (\_ y -> y))
      (V.replicate 2 False)
      (V.replicate 2 False)
  , odeMaxEvents = 10
  , odeSolTimes = [0,1]
  , odeTolerances = defaultTolerances
  }
  where
    t1, t2 :: Fractional a => a
    t1 = 0.01
    t2 = 0.02
    deriv = 100
    events =
      [ EventSpec
        { eventCondition = \t _ -> t - t1
        , eventDirection = Upwards
        }
      , EventSpec
        { eventCondition = \t _ -> t - t2
        , eventDirection = Upwards
        }
      ]

modulusEvent record_event = emptyOdeProblem
  { odeRhs = odeRhsPure $ \t _ -> [t `fmod` 1]
  , odeJacobian = Nothing
  , odeInitCond = [0]
  , odeEvents = events
  , odeEventHandler = mkEventHandler
      [\_ y -> y]
      [False]
      [record_event]
  , odeMaxEvents = 11
  , odeSolTimes = [0,10]
  , odeTolerances = defaultTolerances
  }
  where
    events =
      [ EventSpec
        { eventCondition = \t _ ->
            let
              a = t
              b = 1
            in
              abs((a-b/2) `fmod` (2*b) - b) - b/2
        , eventDirection = AnyDirection
        }
      ]


{-# LANGUAGE StandaloneDeriving, DeriveGeneric #-}

module T5884 where

import GHC.Generics

import Data.Complex

deriving instance Generic (Complex v)


{-# LANGUAGE BangPatterns     #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeSynonymInstances #-}
module Taiji.Utils.DataFrame
    ( DataFrame(..)
    , DataFrameIndex(..)
    , mkDataFrame
    , fromMatrix
    , dim
    , isEmpty
    , Taiji.Utils.DataFrame.transpose
    , rbind
    , cbind
    , cjoin
    , rowNames
    , colNames
    , ReodrderFn
    , reorderColumns
    , reorderRows
    , Taiji.Utils.DataFrame.map
    , mapRows
    , imapRows
    , mapCols
    , imapCols
    , filterRows
    , filterCols
    , Taiji.Utils.DataFrame.zip
    , Taiji.Utils.DataFrame.unzip
    , readData
    , writeTable
    , readTable
    , readTableWith
    , orderByHClust
    , orderByKMeans
    , orderDataFrame

    , pearson
    , spearman
    ) where

import           Bio.Utils.Functions    (ihs')
import           Bio.Utils.Misc         (readDouble)
import Conduit
import Control.Monad (forM_)
import GHC.Generics (Generic)
import qualified Data.ByteString.Char8  as B
import qualified Data.Text as T
import qualified Data.CaseInsensitive   as CI
import           Data.Function          (on)
import           Data.List as L
import AI.Clustering.Hierarchical hiding (normalize)
import AI.Clustering.KMeans
import Data.Vector.Binary ()
import Data.Maybe
import qualified Data.Matrix            as M
import qualified Data.Matrix.Unboxed    as MU
import qualified Data.Vector            as V
import qualified Data.Map.Strict    as HM
import qualified Data.HashSet as HS
import Statistics.Correlation (pearsonMatByRow, spearmanMatByRow)
import Statistics.Matrix (toRowLists, fromRowLists)
import Data.Binary (Binary)
import Control.DeepSeq (NFData)

data DataFrame a = DataFrame
    { _dataframe_row_names :: V.Vector T.Text
    , _dataframe_row_names_idx :: HM.Map T.Text Int
    , _dataframe_col_names :: V.Vector T.Text
    , _dataframe_col_names_idx :: HM.Map T.Text Int
    , _dataframe_data :: M.Matrix a
    } deriving (Show, Generic)

instance Binary a => Binary (M.Matrix a)
instance Binary a => Binary (DataFrame a)
instance NFData a => NFData (DataFrame a)

mkDataFrame :: [T.Text]     -- ^ row names
            -> [T.Text]     -- ^ col names
            -> [[a]]        -- ^ data
            -> DataFrame a
mkDataFrame r c d = DataFrame
    { _dataframe_row_names = V.fromList r
    , _dataframe_row_names_idx = HM.fromList $ L.zip r [0..]
    , _dataframe_col_names = V.fromList c
    , _dataframe_col_names_idx = HM.fromList $ L.zip c [0..]
    , _dataframe_data = M.fromLists d }

fromMatrix :: [T.Text]     -- ^ row names
           -> [T.Text]     -- ^ col names
           -> M.Matrix a  -- ^ data
           -> DataFrame a
fromMatrix r c mat = DataFrame
    { _dataframe_row_names = V.fromList r
    , _dataframe_row_names_idx = HM.fromList $ L.zip r [0..]
    , _dataframe_col_names = V.fromList c
    , _dataframe_col_names_idx = HM.fromList $ L.zip c [0..]
    , _dataframe_data = mat }

class DataFrameIndex i where
    csub :: DataFrame a -> [i] -> DataFrame a
    cindex :: DataFrame a -> i -> V.Vector a
    rsub :: DataFrame a -> [i] -> DataFrame a
    rindex :: DataFrame a -> i -> V.Vector a
    (!)  :: DataFrame a -> (i,i) -> a
    indexMaybe  :: DataFrame a -> (i,i) -> Maybe a

instance DataFrameIndex Int where
    csub df idx = df
        { _dataframe_col_names = V.fromList col_names
        , _dataframe_col_names_idx = HM.fromList $ L.zip col_names [0..]
        , _dataframe_data = M.fromColumns $ L.map (_dataframe_data df `M.takeColumn`) idx }
      where
        col_names = L.map (_dataframe_col_names df V.!) idx

    cindex df idx = _dataframe_data df `M.takeColumn` idx

    rsub df idx = df
        { _dataframe_row_names = V.fromList row_names
        , _dataframe_row_names_idx = HM.fromList $ L.zip row_names [0..]
        , _dataframe_data = M.fromRows $ L.map (_dataframe_data df `M.takeRow`) idx }
      where
        row_names = L.map (_dataframe_row_names df V.!) idx

    rindex df idx = _dataframe_data df `M.takeRow` idx

    (!) df (i,j) = _dataframe_data df M.! (i,j)
    indexMaybe df (i,j)
        | i < r && j < c = Just $ _dataframe_data df M.! (i,j)
        | otherwise = Nothing
      where
        (r, c) = dim df

instance DataFrameIndex T.Text where
    csub df idx = csub df idx'
      where
        idx' = L.map (\i -> HM.findWithDefault
            (error $ "index doesn't exist: " ++ T.unpack i) i $
            _dataframe_col_names_idx df) idx

    cindex df i = cindex df i'
      where
        i' = HM.findWithDefault (error $ "index doesn't exist: " ++ T.unpack i) i $
            _dataframe_col_names_idx df

    rsub df idx = rsub df idx'
      where
        idx' = L.map (\i -> HM.findWithDefault
            (error $ "index doesn't exist: " ++ T.unpack i) i $
            _dataframe_row_names_idx df) idx

    rindex df i = rindex df i'
      where
        i' = HM.findWithDefault (error $ "index doesn't exist: " ++ T.unpack i) i $
            _dataframe_row_names_idx df

    (!) df (a,b) = (!) df (i,j)
      where
        i = HM.findWithDefault (error $ "index doesn't exist: " ++ T.unpack a) a $
            _dataframe_row_names_idx df
        j = HM.findWithDefault (error $ "index doesn't exist: " ++ T.unpack b) b $
            _dataframe_col_names_idx df

    indexMaybe df (a,b) = case HM.lookup a (_dataframe_row_names_idx df) of
        Nothing -> Nothing
        Just i -> case HM.lookup b (_dataframe_col_names_idx df) of
            Nothing -> Nothing
            Just j -> Just $ (!) df (i,j)

dim :: DataFrame a -> (Int, Int)
dim df = M.dim $ _dataframe_data df

isEmpty :: DataFrame a -> Bool
isEmpty df = r == 0 || c == 0
  where
    (r,c) = M.dim $ _dataframe_data df

rbind :: [DataFrame a] -> DataFrame a
rbind dfs | allTheSame (L.map _dataframe_col_names dfs) = (head dfs)
    { _dataframe_row_names = row_names
    , _dataframe_row_names_idx = HM.fromList $ L.zip (V.toList row_names) [0..]
    , _dataframe_data = M.fromBlocks undefined $ L.map (return . _dataframe_data) dfs }
          | otherwise = error "Columns names differ"
  where
    allTheSame xs = all (== head xs) (tail xs)
    row_names = V.concat $ L.map (_dataframe_row_names) dfs

cjoin :: DataFrame a -> DataFrame a -> DataFrame a
cjoin df1 df2 = mkDataFrame rownames (colNames df1 ++ colNames df2) $
    flip L.map rownames $ \nm ->
        V.toList (df1 `rindex` nm) ++ V.toList (df2 `rindex` nm)
  where
    rownames = HS.toList $ HS.fromList (rowNames df1) `HS.intersection`
        HS.fromList (rowNames df2)

cbind :: [DataFrame a] -> DataFrame a
cbind dfs | allTheSame (L.map (HS.fromList . rowNames) dfs) = (head dfs)
    { _dataframe_col_names = col_names
    , _dataframe_col_names_idx = HM.fromList $ L.zip (V.toList col_names) [0..]
    , _dataframe_data = M.fromBlocks undefined [L.map reorderRow dfs] }
          | otherwise = error "Row names differ"
  where
    reorderRow df | rowNames df == row_names = _dataframe_data df
                  | otherwise = M.fromRows $ L.map (df `rindex`) row_names
    row_names = rowNames $ head dfs
    col_names = V.concat $ L.map (_dataframe_col_names) dfs
    allTheSame xs = all (== head xs) (tail xs)

map :: (a -> b) -> DataFrame a -> DataFrame b
map f df = df{_dataframe_data = M.map f $ _dataframe_data df}

zip :: DataFrame a -> DataFrame b -> DataFrame (a,b)
zip df1 df2
    | _dataframe_col_names df1 == _dataframe_col_names df2 &&
      _dataframe_row_names df1 == _dataframe_row_names df2 = df1
        {_dataframe_data = M.zip (_dataframe_data df1) $ _dataframe_data df2}
    | otherwise = error "names mismatch"

unzip :: DataFrame (a,b) -> (DataFrame a, DataFrame b)
unzip df = (df{_dataframe_data = a}, df{_dataframe_data = b})
  where
    (a,b) = M.unzip $ _dataframe_data df

transpose :: DataFrame a -> DataFrame a
transpose DataFrame{..} = DataFrame _dataframe_col_names _dataframe_col_names_idx
    _dataframe_row_names _dataframe_row_names_idx $ M.tr _dataframe_data

rowNames :: DataFrame a -> [T.Text]
rowNames DataFrame{..} = V.toList _dataframe_row_names

colNames :: DataFrame a -> [T.Text]
colNames DataFrame{..} = V.toList _dataframe_col_names

mapRows :: (V.Vector a -> V.Vector b) -> DataFrame a -> DataFrame b
mapRows fn df = df
    { _dataframe_data = M.fromRows $ L.map fn $ M.toRows $ _dataframe_data df }

imapRows :: (T.Text -> V.Vector a -> V.Vector b) -> DataFrame a -> DataFrame b
imapRows fn df = df{ _dataframe_data = M.fromRows $ L.zipWith fn (rowNames df) $ M.toRows $ _dataframe_data df}

mapCols :: (V.Vector a -> V.Vector b) -> DataFrame a -> DataFrame b
mapCols fn df = df
    { _dataframe_data = M.fromColumns $ L.map fn $ M.toColumns $ _dataframe_data df }

imapCols :: (T.Text -> V.Vector a -> V.Vector b) -> DataFrame a -> DataFrame b
imapCols fn df = df
    { _dataframe_data = M.fromColumns $ L.zipWith fn (colNames df) $ M.toColumns $ _dataframe_data df }

filterRows :: (T.Text -> V.Vector a -> Bool) -> DataFrame a -> DataFrame a
filterRows fn df = df
    { _dataframe_row_names = V.fromList names
    , _dataframe_row_names_idx = HM.fromList $ L.zip names [0..]
    , _dataframe_data = M.fromRows rows }
  where
    (names, rows) = L.unzip $ filter (uncurry fn) $
        L.zip (V.toList $ _dataframe_row_names df) $ M.toRows $ _dataframe_data df

filterCols :: (T.Text -> V.Vector a -> Bool) -> DataFrame a -> DataFrame a
filterCols fn df = df
    { _dataframe_col_names = V.fromList names
    , _dataframe_col_names_idx = HM.fromList $ L.zip names [0..]
    , _dataframe_data = M.fromColumns cols }
  where
    (names, cols) = L.unzip $ filter (uncurry fn) $
        L.zip (V.toList $ _dataframe_col_names df) $ M.toColumns $ _dataframe_data df

type ReodrderFn a = [(T.Text, V.Vector a)] -> [(T.Text, V.Vector a)]

reorderRows :: ReodrderFn a -> DataFrame a -> DataFrame a
reorderRows fn df
    | isEmpty df = df
    | otherwise = df
        { _dataframe_row_names = V.fromList names
        , _dataframe_row_names_idx = HM.fromList $ L.zip names [0..]
        , _dataframe_data = M.fromRows rows }
  where
    (names, rows) = L.unzip $ fn $ L.zip (V.toList $ _dataframe_row_names df) $
        M.toRows $ _dataframe_data df

reorderColumns :: ReodrderFn a -> DataFrame a -> DataFrame a
reorderColumns fn df
    | isEmpty df = df
    | otherwise = df
        { _dataframe_col_names = V.fromList names
        , _dataframe_col_names_idx = HM.fromList $ L.zip names [0..]
        , _dataframe_data = M.fromColumns cols}
  where
    (names, cols) = L.unzip $ fn $ L.zip (V.toList $ _dataframe_col_names df) $
        M.toColumns $ _dataframe_data df

writeTable :: FilePath -> (a -> T.Text) -> DataFrame a -> IO ()
writeTable output f DataFrame{..} = runResourceT $ runConduit $
    source .| unlinesAsciiC .| sinkFile output
  where
    source = do
        yield $ B.pack $ T.unpack $ T.intercalate "\t" $
            "" : V.toList _dataframe_col_names
        forM_ (L.zip (V.toList _dataframe_row_names) $ M.toRows _dataframe_data) $
            \(nm, xs) -> yield $ B.pack $ T.unpack $
                T.intercalate "\t" $ nm : L.map f (V.toList xs)
{-# NOINLINE writeTable #-}

-- | Read data, normalize and calculate p-values.
readData :: FilePath   -- ^ PageRank
         -> FilePath   -- ^ Gene expression
         -> IO (DataFrame (Double, Double))  -- ^ ranks, expression
readData input1 input2 = do
    rank <- readTSV <$> B.readFile input1

    -- Read expression profile and apply "ihs" transformation
    expr <- (fmap ihs' . readTSV) <$> B.readFile input2

    let (labels, xs) = L.unzip $ L.map L.unzip $ groupBy ((==) `on` (fst.fst)) $ sort $
            HM.toList $ HM.intersectionWith (,) rank expr
        rowlab = L.map (T.pack . B.unpack . CI.original) $ fst $ L.unzip $ L.map head labels
        collab = L.map (T.pack . B.unpack . CI.original) $ snd $ L.unzip $ head labels
    return $ mkDataFrame rowlab collab xs

readTSV :: B.ByteString -> HM.Map (CI.CI B.ByteString, CI.CI B.ByteString) Double
readTSV input = HM.fromList $ concatMap (f . B.split '\t') content
  where
    f (x:xs) = L.zipWith (\s v -> ((CI.mk x, CI.mk s), readDouble v)) samples xs
    (header:content) = B.lines input
    samples = tail $ B.split '\t' header

readTable :: FilePath -> IO (DataFrame Double)
readTable = readTableWith readDouble

readTableWith :: (B.ByteString -> a) -> FilePath -> IO (DataFrame a)
readTableWith f input = do
    (header:content) <- B.lines <$> B.readFile input
    let samples = tail $ B.split '\t' header
        (rows, dat) = L.unzip $ L.map ((\(x:xs) ->
            (T.pack $ B.unpack x, L.map f xs)) . B.split '\t') content
    return $ mkDataFrame rows (L.map (T.pack . B.unpack) samples) dat

    {-
orderByName :: [T.Text] -> ReodrderFn a
orderByName prefix = sortBy $ \(a,_) (b,_) ->
    let idx1 = findIdx a
        idx2 = findIdx b
    in case () of
        _ | isJust idx1 && isJust idx2 -> case compare idx1 idx2 of
                LT -> LT
                GT -> GT
                EQ -> compare a b
          | otherwise -> compare a b
  where
    findIdx x = go prefix 0
      where
        go (y:ys) !i | isPrefixOf y x = Just i
                     | otherwise = go ys (i+1)
        go _ _ = Nothing
        -}

orderByHClust :: (a -> Double) -> ReodrderFn a
orderByHClust f xs = flatten $ hclust Ward (V.fromList xs) dist
  where
    dist = euclidean `on` V.map f . snd

orderByKMeans :: Int -> (a -> Double) -> ReodrderFn a
orderByKMeans k f xs = concatMap fst $ flatten $ hclust Ward r dist
  where
    r = V.fromList $ L.zip (fromJust $ clusters kmRes) (MU.toRows $ centers kmRes)
    kmRes = kmeansBy k (V.fromList xs) (V.convert . V.map f . snd) defaultKMeansOpts
    dist = euclidean `on` snd

pearson :: DataFrame Double -> DataFrame Double
pearson DataFrame{..} = DataFrame _dataframe_row_names _dataframe_row_names_idx
    _dataframe_row_names _dataframe_row_names_idx $ M.fromLists $
    toRowLists $ pearsonMatByRow $ fromRowLists $ M.toLists _dataframe_data

spearman :: DataFrame Double -> DataFrame Double
spearman DataFrame{..} = DataFrame _dataframe_row_names _dataframe_row_names_idx
    _dataframe_row_names _dataframe_row_names_idx $ M.fromLists $
    toRowLists $ spearmanMatByRow $ fromRowLists $ M.toLists _dataframe_data

orderDataFrame :: (a -> Double) -> DataFrame a -> DataFrame a
orderDataFrame f = reorderColumns (orderByHClust f) . reorderRows (orderByHClust f)


{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TemplateHaskell #-}
{-

Space in the middle: 9
Off-center: 7

-}

import Codec.Picture
import Codec.Picture.Gif
import Data.List.Split
import Data.Complex
import System.Environment
import System.FilePath
import qualified Data.ByteString as BS
import qualified Data.Bits.Bitwise as BA
import Options.Applicative
import Control.Monad
import Data.Monoid
import System.Exit
import Data.Char
import Data.FileEmbed
import Language.Haskell.TH
import Language.Haskell.TH.Syntax
import Text.Regex.Posix
import Numeric (readHex)

cFRAMES = 256          :: Integer
cCENTER = 9            :: Double
cLEN    = 33+cCENTER/2 :: Double  -- Lenght of one arm
cT0RATE = 1000000      :: Integer -- TODO: Need to calibrate this

data Color = R|G|B deriving (Show, Eq)

type Offset = Double
type Shift = Double
type Rotation = Double
type Speed = Double

type BitMaker = (Color -> Bool -> Complex Double -> Bool) -> [Bool]

-- Frame (0..1) to position (-cLen..cLen) to coordinate (-1..1 × -1..1)
type Transformation = Double -> Double -> Complex Double

projectCircular :: Offset -> Shift -> Rotation -> Transformation
projectCircular off shift rot =
    \ frame pos ->
        let rho = frame*2*pi
            cpos = pos :+ 0
        in ((cpos + displacement) * cis rho * cis rot_radians) / (radius :+ 0)
  where
    displacement = (shift :+ off)
    rot_radians = (-rot-3)/12*2*pi
    radius = sqrt ((cLEN + abs shift)^2 + off^2)

projectLinear :: Transformation
projectLinear frame pos = (2*frame -1) :+ pos/cLEN

data TransSpec
    = Circular Offset Shift Rotation
    | Linear
    deriving Show

project :: TransSpec -> Transformation
project (Circular off shift rot) = projectCircular off shift rot
project Linear                   = projectLinear

bitBuilder :: TransSpec -> BitMaker
bitBuilder spec f =
    [ f c a cpos
    | frame <- [0,1..cFRAMES-1]
    , c <- [R,G,B]
    , (a, pos) <- ((False,) <$> arm) ++ (((True,) . negate) <$> arm)
    , let cpos = trans (fromIntegral frame / fromIntegral cFRAMES) pos
    ]
  where
    trans = project spec
    arm = concatMap reverse $ chunksOf 8 $ map (cCENTER/2 +) [0..31]

getColor :: Color -> PixelRGB8 -> Bool
getColor R (PixelRGB8 r _ _ ) = r > 130
getColor G (PixelRGB8 _ g _ ) = g > 130
getColor B (PixelRGB8 _ _ b ) = b > 130

calibrationBits :: [Bool]
calibrationBits = map not $ foldr1 (zipWith (||)) $
    [ bitBuilder (Circular off 0 0) $ \c a z ->
        abs (phase z - off*pi/10) < pi/40 &&
       (c == oc) &&
       (if a then magnitude z > 0.75
             else (magnitude z < 0.75 && (magnitude z > 0.5 || off == 0)))
    | (oc, off) <- zip (cycle [R,G,B]) [-8..8] ]

dynImageToPackagedData builder dynImage =
    bitsToData $
    builder $ \col _ (x :+ y) ->
           let pix = pixelAt pixels
                  (round ((( x+1)/2) * w))
                  (round (((-y+1)/2) * h))
           in not $ getColor col pix
  where
    pixels = convertRGB8 dynImage

    w = fromIntegral (imageWidth pixels)
    h = fromIntegral (imageHeight pixels)

bitsToData = BS.pack . map BA.fromListLE . chunksOf 8

getImage :: BitMaker -> (FilePath, Double) -> IO [(BS.ByteString, Double)]
getImage builder source = case source of
        ("CALIBRATION",n) -> do
            putStrLn $ "Adding offset calibration image"
            pure [(bitsToData calibrationBits, n)]
        (filename, _) | takeExtension filename == ".gif" -> do
            gifData <- BS.readFile filename
            let gifImages = either error id $ decodeGifImages gifData
            let gifDelays = either error id $ getDelaysGifImages gifData
            putStrLn $ "Adding " ++ filename ++ " (" ++ show (length gifImages) ++ " frames)"

            pure $ zip (map (dynImageToPackagedData builder) gifImages)
                       (map ((/100) . fromIntegral) gifDelays)
        (filename, n) -> do
            putStrLn $ "Adding " ++ filename
            dynImage <- either error id <$> readImage filename
            let bits = dynImageToPackagedData builder dynImage
            pure [(bits, n)]

replace :: Int -> BS.ByteString -> BS.ByteString -> BS.ByteString
replace at replacement source
    | BS.length replacement > BS.length after
    = error "replaceIn: Replacement larger than remaining string"
    | otherwise = before <> replacement <> BS.drop (BS.length replacement) after
    where
        (before, after) = BS.splitAt at source

template :: BS.ByteString
template = $(embedFile "firmware/firmware.bin")

offsetInitialStep, offsetTiming, offsetImages :: Int
(offsetInitialStep, offsetTiming, offsetImages) = $(do
    qAddDependentFile "firmware/firmware.map"
    mapFile <- runIO $ readFile "firmware/firmware.map"
    let find name
            | [_,s] <- getAllTextSubmatches (mapFile =~ ("^C: +([0-9ABCDEF]{8}) +_" ++ name)) :: [String]
            , [(i,"")] <- readHex s
            = return i
            | otherwise
            = do reportError $ "Cound not find location of " ++ show name ++ " in firmware.map"
                 fail ""
    ois <- find "initial_step"
    ot  <- find "timing"
    oi  <- find "images"
    return $ TupE $ map LitE [integerL ois, integerL ot, integerL oi]
    )

work :: FilePath -> TransSpec -> Speed -> [(FilePath, Double)] -> IO ()
work output spec speed timed_sources = do
    let builder = bitBuilder spec
    timed_data <- concat <$> mapM (getImage builder) timed_sources
    let (imagesData, timings) = unzip timed_data
    let imageData = BS.concat imagesData

    when (length imagesData > 8) $ do
        putStrLn $ "Too many images"
        exitFailure

    let timingData    = BS.pack $ map fromIntegral $
            concatMap (\x -> [x`mod`256, x`div`256]) $
            map round $
            map (*256) $
            take 8 $ timings ++ repeat 0

    let initial_step_data = BS.pack $ map fromIntegral $
            concatMap (\x -> [x`mod`256, x`div`256]) $
            map round $
            [ speed / fromIntegral cFRAMES * fromIntegral cT0RATE ]

    putStrLn $ "Writing " ++ output
    BS.writeFile output $
        replace offsetImages imageData $
        replace offsetTiming timingData $
        replace offsetInitialStep initial_step_data $
        template

-- Argument handling

main :: IO ()
main = join . customExecParser (prefs showHelpOnError) $
  info (helper <*> parser)
  (  fullDesc
  <> header "YQ3008 image generator"
  -- <> progDesc "TODO"
  )
  where
    linearSpecParser :: Parser TransSpec
    linearSpecParser =
        flag' Linear
            (  long "linear"
            <> help "Scan the input image linearly, not circularly"
            )

    circularSpecParser :: Parser TransSpec
    circularSpecParser = Circular
        <$> option auto
            (  long "offset"
            <> metavar "OFFSET"
            <> help "vertical offset of the bar from the hub (in leds)"
            <> value 0
            )
        <*> option auto
            (  long "shift"
            <> metavar "SHIFT"
            <> help "horizontal shift of the bar from the hub (in leds)"
            <> value 0
            )
        <*> option auto
            (  long "rotation"
            <> metavar "ROTATION"
            <> help "position of the magnet [0..12]"
            <> value 12
            <> showDefault
            )

    parser :: Parser (IO ())
    parser = work
        <$> strOption
            (  long "output"
            <> short 'o'
            <> metavar "FILE"
            <> help "output file"
            )
        <*> (linearSpecParser <|> circularSpecParser)
        <*> option auto
            (  long "speed"
            <> metavar "SECONDS"
            <> help "how long one image scan is initially"
            <> value 0.5
            <> showDefault
            )
        <*> some (
            (,) <$> strArgument
                    (  metavar "IMAGE"
                    <> help "input file (can be CALIBRATION)"
                    )
                <*> argument auto
                    (  metavar "DURATION"
                    <> help "duration (in seconds)"
                    )
            )


module Reexports
  ( Matrix
  , Vector
  , R
  , toLists
  , toList
  , matrix
  , fromLists
  , fromList
  ) where

import Numeric.LinearAlgebra.Data
  ( Matrix
  , R
  , Vector
  , fromList
  , fromLists
  , matrix
  , toList
  , toLists
  )


module Pfe3Metrics where
import Data.Maybe(mapMaybe)

import PFE0(allModules,pput)
import PFE3(parseModule)
import HsDeclStruct(DI(..))
import HsModule(hsModDecls)
import DefinedNames(definedType)
import HasBaseStruct(basestruct)
import Statistics

import PfeParse
import MUtils

pfe3MetricsCmds =
  [("classmetrics", (noArgs classMetrics,"number of instances per class metrics"))]
classMetrics =
  do ms <- allModules
     (classes,insts) <- apBoth concat . unzip # mapM getClassInstDecls ms
     let cinstcnt = [(c,length [()|i<-insts,i==c])|c<-classes]
     pput (ppStatistics "number of instances" "class" cinstcnt)

getClassInstDecls m = cls_insts . hsModDecls . snd # parseModule m

cls_insts ds = (mapMaybe className ds,concatMap instClass ds)

className d =
  case basestruct d of
    Just (HsClassDecl _ _ tp _ _) -> Just (definedType tp)
    _ -> Nothing

instClass d =
  case basestruct d of
    Just (HsInstDecl _ _ _ tp _) -> [definedType tp]
    Just (HsDataDecl _ _ _ _ cls) -> cls
    Just (HsNewTypeDecl _ _ _ _ cls) -> cls
    _ -> []


{-# LANGUAGE TypeFamilies, GeneralizedNewtypeDeriving, DeriveGeneric, TupleSections, RecordWildCards #-}

module BayesStack.Models.Topic.CitationInfluence
  ( -- * Primitives
    NetData(..)
  , netData
  , HyperParams(..)
  , MState(..)
  , CitedUpdateUnit
  , CitingUpdateUnit
  , ItemSource(..)
  , CitedNode(..), CitedNodeItem(..)
  , CitingNode(..), CitingNodeItem(..)
  , Citing(..), Cited(..)
  , Item(..), Topic(..), NodeItem(..), Node(..), Arc(..)
  , setupNodeItems
    -- * Initialization
  , verifyNetData, cleanNetData
  , ModelInit
  , randomInitialize
  , model
  , updateUnits
    -- * Diagnostics
  , modelLikelihood
  , influence
  , arcTopicMixture
  ) where

import qualified Data.Vector as V
import Statistics.Sample (mean)

import           Prelude hiding (mapM_, sum)

import           Data.Set (Set)
import qualified Data.Set as S

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

import           Data.Foldable hiding (product)
import           Control.Applicative ((<$>), (<*>))
import           Control.Monad (when)
import           Control.Monad.Trans.State.Strict
import           Control.Monad.Trans.Writer.Strict

import           Data.Random
import           Data.Random.Lift (lift)
import           Data.Random.Distribution.Categorical (categorical)
import           Numeric.Log hiding (sum)

import           BayesStack.Types
import           BayesStack.Gibbs
import           BayesStack.DirMulti
import           BayesStack.TupleEnum ()
import           BayesStack.Models.Topic.Types

import           GHC.Generics
import           Data.Binary (Binary)
import           Control.DeepSeq

data ItemSource = Shared | Own deriving (Show, Eq, Enum, Ord, Generic)
instance Binary ItemSource
instance NFData ItemSource

newtype Citing a = Citing a deriving (Show, Eq, Enum, Ord, Generic, NFData)
newtype Cited a = Cited a deriving (Show, Eq, Enum, Ord, Generic, NFData)
instance Binary a => Binary (Citing a)
instance Binary a => Binary (Cited a)

type CitingNode = Citing Node
type CitedNode = Cited Node
type CitingNodeItem = Citing NodeItem
type CitedNodeItem = Cited NodeItem

-- ^ A directed edge
data Arc = Arc {citingNode :: !CitingNode, citedNode :: !CitedNode}
            deriving (Show, Eq, Ord, Generic)
instance Binary Arc

data HyperParams = HyperParams
                   { alphaPsi         :: Double
                   , alphaLambda      :: Double
                   , alphaPhi         :: Double
                   , alphaOmega       :: Double
                   , alphaGammaShared :: Double
                   , alphaGammaOwn    :: Double
                   }
                 deriving (Show, Eq, Generic)
instance Binary HyperParams

data NetData = NetData { dHypers             :: !(HyperParams)
                       , dArcs               :: !(Set Arc)
                       , dItems              :: !(Set Item)
                       , dTopics             :: !(Set Topic)
                       , dNodeItems          :: !(Map NodeItem (Node, Item))
                       , dCitingNodes        :: !(Map CitingNode (Set CitedNode))
                         -- ^ Maps each citing node to the set of nodes cited by it
                       , dCitedNodes         :: !(Map CitedNode (Set CitingNode))
                         -- ^ Maps each cited node to the set of nodes citing it
                       }
              deriving (Show, Eq, Generic)
instance Binary NetData

netData :: HyperParams -> Set Arc -> Map NodeItem (Node,Item) -> Set Topic -> NetData
netData hypers arcs nodeItems topics =
    NetData { dHypers       = hypers
            , dArcs         = arcs
            , dItems        = foldMap (S.singleton . snd) $ M.elems nodeItems
            , dTopics       = topics
            , dNodeItems    = nodeItems
            , dCitingNodes  = M.unionsWith S.union
                              $ map (\(Arc a b)->M.singleton a $ S.singleton b)
                              $ S.toList arcs
            , dCitedNodes   = M.unionsWith S.union
                              $ map (\(Arc a b)->M.singleton b $ S.singleton a)
                              $ S.toList arcs
            }

dCitingNodeItems :: NetData -> Map CitingNodeItem (CitingNode, Item)
dCitingNodeItems nd =
    M.mapKeys Citing
    $ M.map (\(n,i)->(Citing n, i))
    $ M.filter (\(n,i)->Citing n `M.member` dCitingNodes nd)
    $ dNodeItems nd

itemsOfCitingNode :: NetData -> CitingNode -> [Item]
itemsOfCitingNode d (Citing u) =
    map snd $ M.elems $ M.filter (\(n,_)->n==u) $ dNodeItems d

connectedNodes :: Set Arc -> Set Node
connectedNodes arcs =
    S.map ((\(Cited n)->n) . citedNode) arcs `S.union` S.map ((\(Citing n)->n) . citingNode) arcs

cleanNetData :: NetData -> NetData
cleanNetData d =
    let nodesWithItems = S.fromList $ map fst $ M.elems $ dNodeItems d
        nodesWithArcs = connectedNodes $ dArcs d
        keptNodes = nodesWithItems `S.intersection` nodesWithArcs
        keepArc (Arc (Citing citing) (Cited cited)) =
            citing `S.member` keptNodes && cited `S.member` keptNodes
    in d { dArcs = S.filter keepArc $ dArcs d
         , dNodeItems = M.filter (\(n,i)->n `S.member` keptNodes) $ dNodeItems d
         }

verifyNetData :: (Node -> String) -> NetData -> [String]
verifyNetData showNode d = execWriter $ do
    let nodesWithItems = S.fromList $ map fst $ M.elems $ dNodeItems d
    forM_ (dArcs d) $ \(Arc (Citing citing) (Cited cited))->do
        when (cited `S.notMember` nodesWithItems)
            $ tell [showNode cited++" has arc yet has no items"]
        when (citing `S.notMember` nodesWithItems)
            $ tell [showNode citing++" has arc yet has no items"]

type CitedModelInit = Map CitedNodeItem (Setting CitedUpdateUnit)
type CitingModelInit = Map CitingNodeItem (Setting CitingUpdateUnit)
data ModelInit = ModelInit CitedModelInit CitingModelInit
               deriving (Show)

randomInitializeCited :: NetData -> CitedModelInit -> RVar CitedModelInit
randomInitializeCited d init = execStateT doInit init
    where doInit = let unset = citedNodeItems `S.difference` M.keysSet init
                       citedNodeItems = S.map Cited (M.keysSet (dNodeItems d))
                   in mapM_ (randomInitCitedUU d) (S.toList unset)

modify' :: Monad m => (a -> a) -> StateT a m ()
modify' f = do x <- get
               put $! f x

randomInitCitedUU :: NetData -> CitedNodeItem -> StateT CitedModelInit RVar ()
randomInitCitedUU d ni = do
    t' <- lift $ randomElement $ toList $ dTopics d
    modify' $ M.insert ni t'

randomInitializeCiting :: NetData -> CitingModelInit -> RVar CitingModelInit
randomInitializeCiting d init = execStateT doInit init
    where doInit :: StateT CitingModelInit RVar ()
          doInit = let unset = M.keysSet (dCitingNodeItems d) `S.difference` M.keysSet init
                   in mapM_ (randomInitCitingUU d) (S.toList unset)

randomInitCitingUU :: NetData -> CitingNodeItem -> StateT CitingModelInit RVar ()
randomInitCitingUU d cni@(Citing ni) =
    let (n,_) = dNodeItems d M.! ni
    in case dCitingNodes d M.! Citing n of
           a | S.null a -> do
               t <- lift $ randomElement $ toList $ dTopics d
               modify' $ M.insert cni $ OwnSetting t

           citedNodes -> do
               s <- lift $ randomElement [Shared, Own]
               c <- lift $ randomElement $ toList citedNodes
               t <- lift $ randomElement $ toList $ dTopics d
               modify' $ M.insert cni $
                   case s of Shared -> SharedSetting t c
                             Own    -> OwnSetting t

randomInitialize :: NetData -> RVar ModelInit
randomInitialize d =
    ModelInit <$> randomInitializeCited d M.empty <*> randomInitializeCiting d M.empty

model :: NetData -> ModelInit -> MState
model d (ModelInit citedInit citingInit) =
    let citingNodes = M.keys $ dCitingNodes d
        HyperParams {..} = dHypers d
        s = MState { -- Citing model
                     stPsis = let dist n = case toList $ dCitingNodes d M.! n of
                                               []    -> M.empty
                                               nodes -> M.singleton n
                                                        $ symDirMulti alphaPsi nodes
                              in foldMap dist citingNodes
                   , stPhis = let dist = symDirMulti alphaPhi (toList $ dItems d)
                              in foldMap (\t->M.singleton t dist) $ dTopics d
                   , stGammas = let dist = multinom [ (Shared, alphaGammaShared)
                                                    , (Own, alphaGammaOwn) ]
                                in foldMap (\t->M.singleton t dist) citingNodes
                   , stOmegas = let dist = symDirMulti alphaOmega (toList $ dTopics d)
                                in foldMap (\t->M.singleton t dist) citingNodes
                   , stCiting = M.empty

                   -- Cited model
                   , stLambdas = let dist = symDirMulti alphaLambda (toList $ dTopics d)
                                 in foldMap (\t->M.singleton t dist) $ M.keys $ dCitedNodes d
                   , stT' = M.empty
                   }

        initCitingUU :: CitingUpdateUnit -> State MState ()
        initCitingUU uu = do
            let err = error $ "CitationInference: Initial value for "++show uu++" not given\n"
                s = maybe err id $ M.lookup (uuNI uu) citingInit
            modify' $ setCitingUU uu (Just s)

        initCitedUU :: CitedUpdateUnit -> State MState ()
        initCitedUU uu = do
            let err = error $ "CitationInference: Initial value for "++show uu++" not given"
                s = maybe err id $ M.lookup (uuNI' uu) citedInit
            modify' $ setCitedUU uu (Just s)

    in execState (do mapM_ initCitingUU $ citingUpdateUnits d
                     mapM_ initCitedUU $ citedUpdateUnits d
                 ) s

updateUnits :: NetData -> [WrappedUpdateUnit MState]
updateUnits d = map WrappedUU (citedUpdateUnits d)
             ++ map WrappedUU (citingUpdateUnits d)

data CitingSetting = OwnSetting !Topic
                   | SharedSetting !Topic !CitedNode
                   deriving (Show, Eq, Generic)
instance Binary CitingSetting
instance NFData CitingSetting where
    rnf (OwnSetting t)      = rnf t `seq` ()
    rnf (SharedSetting t c) = rnf t `seq` rnf c `seq` ()

data MState = MState { -- Citing model state
                       stGammas   :: !(Map CitingNode (Multinom Int ItemSource))
                     , stOmegas   :: !(Map CitingNode (Multinom Int Topic))
                     , stPsis     :: !(Map CitingNode (Multinom Int CitedNode))
                     , stPhis     :: !(Map Topic (Multinom Int Item))

                     , stCiting   :: !(Map CitingNodeItem CitingSetting)

                     -- Cited model state
                     , stLambdas  :: !(Map CitedNode (Multinom Int Topic))

                     , stT'       :: !(Map CitedNodeItem Topic)
                     }
            deriving (Show, Generic)
instance Binary MState

modelLikelihood :: MState -> Probability
modelLikelihood model =
    product $ map likelihood (M.elems $ stGammas model)
           ++ map likelihood (M.elems $ stPhis model)
           ++ map likelihood (M.elems $ stLambdas model)
           ++ map likelihood (M.elems $ stOmegas model)
           ++ map likelihood (M.elems $ stPsis model)

-- | Mixture of the topics of an edge
arcTopicMixture :: NetData -> MState -> Arc -> Topic -> Probability
arcTopicMixture nd m (Arc d c) t =
    let itemObs = zip (itemsOfCitingNode nd d) (repeat 1)
        phi = stPhis m M.! t
        lambda = stLambdas m M.! c
    in realToFrac (sampleProb lambda t) * obsProb phi itemObs

-- | Geometric mean
geomMean :: V.Vector (Log Double) -> Log Double
geomMean = Exp . mean . V.map ln

-- | The geometric mean of the probabilities of a collection of items under a
-- given topic mixture.
topicCompatibility :: MState -> [Item] -> Multinom Int Topic -> Probability
topicCompatibility m items lambda =
    geomMean $ V.fromList
            $ do t <- toList $ dmDomain lambda
                 let phi = stPhis m M.! t
                     itemObs = zip items (repeat 1)
                 return $ realToFrac (sampleProb lambda t) * obsProb phi itemObs

-- | Normalized compatibilities with a set of items
topicCompatibilities :: (Functor f, Foldable f)
                     => MState -> [Item] -> f (Multinom Int Topic) -> f Probability
topicCompatibilities m items topics =
    let scores = fmap (topicCompatibility m items) topics
    in fmap (/sum scores) scores

-- | The influence of cited nodes on a citing node.
influence :: NetData -> MState -> CitingNode -> Map CitedNode Probability
influence d m u =
    let lambdas = foldMap (\f->M.singleton f $ stLambdas m M.! f)
                 $ S.toList $ dCitingNodes d M.! u
    in topicCompatibilities m (itemsOfCitingNode d u) lambdas

-- Cited update unit (LDA-like)
data CitedUpdateUnit = CitedUpdateUnit { uuNI' :: !CitedNodeItem
                                       , uuN'  :: !CitedNode
                                       , uuX' :: !Item
                                       }
                     deriving (Show, Generic)
instance Binary CitedUpdateUnit

instance UpdateUnit CitedUpdateUnit where
    type ModelState CitedUpdateUnit = MState
    type Setting CitedUpdateUnit = Topic
    fetchSetting uu ms = maybe (error $ "Update unit "++show uu++" has no setting") id
                         $ M.lookup (uuNI' uu) (stT' ms)
    evolveSetting ms uu = categorical $ citedFullCond (setCitedUU uu Nothing ms) uu
    updateSetting uu _ s' = setCitedUU uu (Just s') . setCitedUU uu Nothing

citedProb :: MState -> CitedUpdateUnit -> Setting CitedUpdateUnit -> Double
citedProb st (CitedUpdateUnit {uuN'=n', uuX'=x'}) t =
    let lambda = stLambdas st M.! n'
        phi = stPhis st M.! t
    in realToFrac $ sampleProb lambda t * sampleProb phi x'

citedUpdateUnits :: NetData -> [CitedUpdateUnit]
citedUpdateUnits d =
    map (\(ni',(n',x'))->CitedUpdateUnit { uuNI'      = Cited ni'
                                         , uuN'       = Cited n'
                                         , uuX'       = x'
                                         }
        ) $ M.assocs
          $ M.filter (\(n,i)->Cited n `M.member` dCitedNodes d)
          $ dNodeItems d

setCitedUU :: CitedUpdateUnit -> Maybe Topic -> MState -> MState
setCitedUU uu@(CitedUpdateUnit {uuN'=n', uuNI'=ni', uuX'=x'}) setting ms =
    let t' = maybe (fetchSetting uu ms) id setting
        set = maybe Unset (const Set) setting
    in ms { stLambdas = M.adjust (setMultinom set t') n' (stLambdas ms)
          , stPhis = M.adjust (setMultinom set x') t' (stPhis ms)
          , stT' = case setting of Just _  -> M.insert ni' t' $ stT' ms
                                   Nothing -> stT' ms
          }

citedFullCond ::MState -> CitedUpdateUnit -> [(Double, Topic)]
citedFullCond ms uu = do
    t <- M.keys $ stPhis ms
    return (citedProb ms uu t, t)

-- Citing Update unit (Shared Taste-like)
data CitingUpdateUnit = CitingUpdateUnit { uuNI    :: !CitingNodeItem
                                         , uuN     :: !CitingNode
                                         , uuX     :: !Item
                                         , uuCites :: !(Set CitedNode)
                                         }
                      deriving (Show, Generic)
instance Binary CitingUpdateUnit

instance UpdateUnit CitingUpdateUnit where
    type ModelState CitingUpdateUnit = MState
    type Setting CitingUpdateUnit = CitingSetting
    fetchSetting uu ms = stCiting ms M.! uuNI uu
    evolveSetting ms uu = categorical $ citingFullCond (setCitingUU uu Nothing ms) uu
    updateSetting uu _ s' = setCitingUU uu (Just s') . setCitingUU uu Nothing

citingUpdateUnits :: NetData -> [CitingUpdateUnit]
citingUpdateUnits d =
    map (\(ni,(n,x))->CitingUpdateUnit { uuNI      = ni
                                       , uuN       = n
                                       , uuX       = x
                                       , uuCites   = dCitingNodes d M.! n
                                       }
        ) $ M.assocs $ dCitingNodeItems d

citingProb :: MState -> CitingUpdateUnit -> Setting CitingUpdateUnit -> Double
citingProb st (CitingUpdateUnit {uuN=n, uuX=x}) setting =
    let gamma = stGammas st M.! n
        omega = stOmegas st M.! n
        psi = stPsis st M.! n
    in case setting of
        SharedSetting t c -> let phi = stPhis st M.! t
                                 lambda = stLambdas st M.! c
                             in sampleProb gamma Shared
                              * sampleProb psi c
                              * sampleProb lambda t
                              * sampleProb phi x
        OwnSetting t      -> let phi = stPhis st M.! t
                             in sampleProb gamma Own
                              * sampleProb omega t
                              * sampleProb phi x

citingFullCond :: MState -> CitingUpdateUnit -> [(Double, Setting CitingUpdateUnit)]
citingFullCond ms uu = map (\s->(citingProb ms uu s, s)) $ citingDomain ms uu

citingDomain :: MState -> CitingUpdateUnit -> [Setting CitingUpdateUnit]
citingDomain ms uu = do
    s <- [Own, Shared]
    t <- M.keys $ stPhis ms
    case s of
        Shared -> do c <- S.toList $ uuCites uu
                     return $ SharedSetting t c
        Own    -> do return $ OwnSetting t

setCitingUU :: CitingUpdateUnit -> Maybe (Setting CitingUpdateUnit) -> MState -> MState
setCitingUU uu@(CitingUpdateUnit {uuNI=ni, uuN=n, uuX=x}) setting ms =
    let set = maybe Unset (const Set) setting
        ms' = case maybe (fetchSetting uu ms) id  setting of
            SharedSetting t c -> ms { stPsis = M.adjust (setMultinom set c) n $ stPsis ms
                                    , stLambdas = M.adjust (setMultinom set t) c $ stLambdas ms
                                    , stPhis = M.adjust (setMultinom set x) t $ stPhis ms
                                    , stGammas = M.adjust (setMultinom set Shared) n $ stGammas ms
                                    }
            OwnSetting t      -> ms { stOmegas = M.adjust (setMultinom set t) n $ stOmegas ms
                                    , stPhis = M.adjust (setMultinom set x) t $ stPhis ms
                                    , stGammas = M.adjust (setMultinom set Own) n $ stGammas ms
                                    }
    in ms' { stCiting = M.alter (const setting) ni $ stCiting ms' }


{-# LANGUAGE DefaultSignatures      #-}
{-# LANGUAGE DeriveFoldable         #-}
{-# LANGUAGE DeriveFunctor          #-}
{-# LANGUAGE DeriveTraversable      #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE RecordWildCards        #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TemplateHaskell        #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE TypeOperators          #-}
module Circuits.Analysis.DC where

import           Control.Lens
import           Control.Monad.Primitive
import           Control.Monad.ST
import           Data.Primitive.MutVar
import qualified Data.Vector                  as V
import qualified Data.Vector.Storable         as VS
import qualified Data.Vector.Storable.Mutable as VSM
import           Generics.Deriving.Lens
import           GHC.Generics
import qualified Numeric.LinearAlgebra        as HMatrix

import           Debug.Trace

import           Circuits.Circuit

-- | Describes the time scale of a DC simulation. We can either perform a DC steady-state or a transient analysis.
data TimeScale
  = SteadyState
  -- ^ solving for steady state
  | Step
    { _currentTime :: Double -- ^ the current simulation time in seconds since t0
    , _timeStep    :: Double -- ^ the size of current time step to be taken
    }
    -- ^ solving for next time step

makeLenses ''TimeScale

-- | Type class providing an interface for the component stamps used in DC analysis.
class SimulateDC c where
  -- | Beginning to solve the next time step.
  beginSolve :: Simulation m Double sim => c -> TimeScale -> sim -> m ()
  default beginSolve :: (Generic c, GSimulateDC (Rep c), Simulation m Double sim) => c -> TimeScale -> sim -> m ()
  beginSolve c t s = forOf_ generic c $ \gc -> gbeginSolve gc t s
  -- | Beginning the next iteration of solving the non-linear parts.
  -- This should set up the right hand side based on the current element state.
  beginIteration :: Simulation m Double sim => c -> TimeScale -> sim -> m ()
  default beginIteration :: (Generic c, GSimulateDC (Rep c), Simulation m Double sim) => c -> TimeScale -> sim -> m ()
  beginIteration c t s = forOf_ generic c $ \gc -> gbeginIteration gc t s
  -- | End of the current iteration, after the referenced variables have been updated.
  -- It should update the state of non-linear parts, but not the state of linear parts.
  endIteration :: c -> TimeScale -> c
  default endIteration :: (Generic c, GSimulateDC (Rep c)) => c -> TimeScale -> c
  endIteration c t = over generic (`gendIteration` t) c
  -- | Finished solving the current time step. Here, the state of linear parts should be updated.
  endSolve :: c -> TimeScale -> c
  default endSolve :: (Generic c, GSimulateDC (Rep c)) => c -> TimeScale -> c
  endSolve c t = over generic (`gendSolve` t) c

instance SimulateDC Int where
  beginSolve _ _ _ = pure ()
  beginIteration _ _ _ = pure ()
  endIteration x _ = x
  endSolve x _ = x

instance SimulateDC Double where
  beginSolve _ _ _ = pure ()
  beginIteration _ _ _ = pure ()
  endIteration x _ = x
  endSolve x _ = x

instance SimulateDC Variable
instance SimulateDC Node
instance SimulateDC Branch

instance SimulateDC c => SimulateDC (Circuit c)

-- * Generic deriving of SimulateDC instances

class GSimulateDC f where
  gbeginSolve :: Simulation m Double sim => f p -> TimeScale -> sim -> m ()
  gbeginIteration :: Simulation m Double sim => f p -> TimeScale -> sim -> m ()
  gendIteration :: f p -> TimeScale -> f p
  gendSolve :: f p -> TimeScale -> f p

instance GSimulateDC V1 where
  gbeginSolve _ _ _ = return ()
  gbeginIteration _ _ _ = return ()
  gendIteration x _ = x
  gendSolve x _ = x

instance GSimulateDC U1 where
  gbeginSolve _ _ _ = return ()
  gbeginIteration _ _ _ = return ()
  gendIteration x _ = x
  gendSolve x _ = x

instance SimulateDC a => GSimulateDC (K1 i a) where
  gbeginSolve (K1 x) = beginSolve x
  gbeginIteration (K1 x) = beginIteration x
  gendIteration (K1 x) = K1 . endIteration x
  gendSolve (K1 x) = K1 . endSolve x

instance GSimulateDC f => GSimulateDC (M1 i t f) where
  gbeginSolve (M1 x) = gbeginSolve x
  gbeginIteration (M1 x) = gbeginIteration x
  gendIteration (M1 x) = M1 . gendIteration x
  gendSolve (M1 x) = M1 . gendSolve x

instance (GSimulateDC f, GSimulateDC g) => GSimulateDC (f :+: g) where
  gbeginSolve (L1 x) = gbeginSolve x
  gbeginSolve (R1 x) = gbeginSolve x
  gbeginIteration (L1 x) = gbeginIteration x
  gbeginIteration (R1 x) = gbeginIteration x
  gendIteration (L1 x) = L1 . gendIteration x
  gendIteration (R1 x) = R1 . gendIteration x
  gendSolve (L1 x) = L1 . gendSolve x
  gendSolve (R1 x) = R1 . gendSolve x

instance (GSimulateDC f, GSimulateDC g) => GSimulateDC (f :*: g) where
  gbeginSolve (x :*: y) t s = gbeginSolve x t s >> gbeginSolve y t s
  gbeginIteration (x :*: y) t s = gbeginIteration x t s >> gbeginIteration y t s
  gendIteration (x :*: y) t = gendIteration x t :*: gendIteration y t
  gendSolve (x :*: y) t = gendSolve x t :*: gendSolve y t

-- * DC Simulation

-- | A mutable system of linear equations suitable for representing linear circuits.
data SimulationState s v = SimulationState
  { size                 :: Int -- ^ the number of variables in this system
  , coefficients         :: VSM.MVector s v -- ^ the underlying storage of the equations coefficients in row-major order
  , coefficientsModified :: MutVar s Bool -- ^ flag indicating whether the coefficient matrix has been modified
  , rhs                  :: VSM.MVector s v -- ^ the right hand side of the equations
  , rhsModified          :: MutVar s Bool -- ^ flag indicating whether the RHS has been modified
  }

-- | Creates a new mutable linear circuit equation system with the given number of variables.
newSimulation :: (PrimMonad m, VSM.Storable v, Num v) => Int -> m (SimulationState (PrimState m) v)
newSimulation n =
  SimulationState n <$> VSM.replicate (n * n) 0 <*> newMutVar False <*> VSM.replicate n 0 <*> newMutVar False

cleanCoefficients :: (PrimMonad m) => SimulationState (PrimState m) v -> m ()
cleanCoefficients s = writeMutVar (coefficientsModified s) False

cleanRhs :: (PrimMonad m) => SimulationState (PrimState m) v -> m ()
cleanRhs s = writeMutVar (rhsModified s) False

-- | Computes the index into the flat buffer of matrix variables from a row and column.
flatIndex :: SimulationState s v -- ^ the equation system
          -> Int -- ^ matrix row
          -> Int -- ^ matrix column
          -> Int -- ^ flat buffer index
flatIndex sim row col = row * size sim + col

class Monad m => Simulation m v sys | sys -> v  where
  stampMatrix :: sys -> Variable -> Variable -> v -> m ()
  stampRhs :: sys -> Variable -> v -> m ()

instance (PrimMonad m, PrimState m ~ s, Num v, VS.Storable v) => Simulation m v (SimulationState s v) where
  stampMatrix sys (Index row _) (Index col _) value = do
    VSM.modify (coefficients sys) (+ value) (flatIndex sys row col)
    writeMutVar (coefficientsModified sys) True
  stampMatrix _ _ _ _ = return () -- no stamping when ground is involved

  stampRhs sys (Index row _) value = do
    VSM.modify (rhs sys) (+ value) row
    writeMutVar (rhsModified sys) True
  stampRhs _ Ground _ = return () -- no stamping for ground

initialSolution :: HasVariables c => Circuit c -> VS.Vector Double
initialSolution c = runST $ do
  vec <- VSM.new (c ^. variableCount)
  forOf_ (model . variables . _Index) c $ uncurry (VSM.write vec)
  VS.unsafeFreeze vec

step :: forall c. (HasVariables c, SimulateDC c) => TimeScale -> Circuit c -> Circuit c
step time c = runST simulate where
  simulate :: ST s (Circuit c)
  simulate = do
    sim <- newSimulation (c ^. variableCount)
    -- initialize matrix
    beginSolve c time sim
    mat <- HMatrix.reshape (c ^. variableCount) <$> VS.unsafeFreeze (coefficients sim)
    cleanCoefficients sim
    -- iterate
    writeMutVar (rhsModified sim) True
    c' <- solveNonLinear sim (HMatrix.luPacked mat) (initialSolution c) c
    -- finalize
    return $ endSolve c' time

  solveNonLinear :: SimulationState s Double -> HMatrix.LU Double -> VS.Vector Double -> Circuit c -> ST s (Circuit c)
  solveNonLinear sim lu prevSolution prevCircuit = do
    -- prepare RHS
    beginIteration c time sim
    rhsUpdated <- readMutVar (rhsModified sim)
    cleanRhs sim
    if rhsUpdated
      then do
        -- solve and update components
        rhs' <- VS.unsafeFreeze $ rhs sim
        let solution = HMatrix.flatten $ HMatrix.luSolve lu (HMatrix.asColumn rhs')
            c' = endIteration (updateValues solution c) time
            maxChange = VS.maximum $ VS.zipWith (\x y -> abs (x - y)) solution prevSolution
        if maxChange < 1e-6
          then return c'
          else solveNonLinear sim lu solution c'
      else return prevCircuit

-- | Update all variable references inside a circuit description with its new values.
updateValues :: HasVariables a => VS.Vector Double -> a -> a
updateValues values = over variables replace where
  replace Ground      = Ground
  replace (Index i _) = vars V.! i
  vars = V.imap Index $ VS.convert values


module ReadWrite
  ( -- * Read
    readMolV2000
    -- * Write
  , writeXYZ
  ) where

import Prelude hiding (readFile)

import Control.Lens
import Control.Monad.State
import Numeric.LinearAlgebra.Data
import System.Directory
import System.IO
import System.IO.Unsafe
import Text.Printf (hPrintf)

import Types

-- * Read
-- | Read molecule at *.pdb format (V2000)
readMolV2000 :: FilePath -> (Molecule, [Bond])
readMolV2000 inf =
  let txt = (lines . unsafePerformIO . readFile) inf
      counts_line = words (txt !! 3)
      count_atoms = read (counts_line !! 0)
      count_bonds = read (counts_line !! 1)
      atoms_lines = (take count_atoms . drop 4) txt
      bonds_lines = (take count_bonds . drop (4 + count_atoms)) txt
   in (foldr addAtom molecule atoms_lines, foldr addBond [] bonds_lines)
  where
    addAtom a = over atoms ((:) (readline' a))
    addBond b = (:) (readline'' b)
    readline' l =
      execState
        (do let w = words l
            acoordin . x .= read (w !! 0)
            acoordin . y .= read (w !! 1)
            acoordin . z .= read (w !! 2)
            aelement .= w !! 3
            avdwrad .= vdwr (w !! 3))
        atom
    readline'' l =
      execState
        (do let w = words l
            bfid .= read (w !! 0) - 1
            bsid .= read (w !! 1) - 1
            btype .= read (w !! 2)
            bster .= read (w !! 3)
            btop .= read (w !! 4))
        bond

-- | Get VDW radius
vdwr :: Element -> Double
vdwr a =
  case a of
    "H" -> 0.500 -- 0.5 -- 1.000
    "O" -> 0.650 -- 0.5 -- 1.300
    "N" -> 0.700 -- 0.5 -- 1.400
    "C" -> 0.750 -- 0.5 -- 1.500
    "S" -> 0.950 -- 0.5 -- 1.900
    "Br"-> 0.950
    othrewise -> error $ "vdwr not found for: " ++ show a

-- * Write
-- | Write molecule in *.xyz format
writeXYZ :: FilePath -> String -> Molecule -> IO ()
writeXYZ ouf comment molecule = do
  (tmp_name, tmp_handle) <- openTempFile "." "temp"
  hPrint tmp_handle (views atoms length molecule)
  hPutStrLn tmp_handle comment
  mapM_ (writeData tmp_handle) (view atoms molecule)
  hClose tmp_handle
  renameFile tmp_name ouf
  where
    writeData hdl atom = do
      let e = view aelement atom
          (Point x y z) = view acoordin atom
      hPrintf hdl "%s\t%8.6f\t%8.6f\t%8.6f\n" e x y z


module Mnist.Load where

  import Neural
  import Numeric.LinearAlgebra
  import Debug.Trace
  import qualified Data.ByteString.Lazy as B
  import Data.Binary.Get
  import Data.Int
  import Control.Monad

  -- MNIST binary data parsers
  getLabelHeader :: Get Int
  getLabelHeader = do
    _ <- getWord32be
    num <- getInt32be
    return $ fromIntegral num

  getLabel :: Get (Vector Double)
  getLabel = do
    label <- getInt8
    return $ (toOneHot . fromIntegral) label
      where
        toOneHot :: Int -> Vector Double
        toOneHot v = vector $ (\i -> if i==v then 1.0 else 0.0) `map` [0..9]

  genVector :: Get (Vector Double) -> Get [Vector Double]
  genVector getOne = do
    empty <- isEmpty
    if empty then return []
    else do h <- getOne
            tl <- genVector getOne
            return (h:tl)

  getImageHeader :: Get (Int, Int, Int)
  getImageHeader = do
    _ <- getWord32be
    num <- getInt32be
    w <- getInt32be
    h <- getInt32be
    return (fromIntegral num, fromIntegral w, fromIntegral h)

  getImage :: Get (Vector Double)
  getImage = do
    colors <- replicateM (28*28) getWord8
    return $ vector ((toRealColor . fromIntegral) `map` colors)

  toRealColor :: Int -> Double
  toRealColor v = (fromIntegral v) / 255

  -- Parse MNIST Dataset into a list of examples (pair of a list of images and a list of labels)
  loadExamples :: IO [Example]
  loadExamples = do
    rawImage1 <- B.readFile "data/train-images.idx3-ubyte"
    rawImage2 <- B.readFile "data/t10k-images.idx3-ubyte"
    rawLabel1 <- B.readFile "data/train-labels.idx1-ubyte"
    rawLabel2 <- B.readFile "data/t10k-labels.idx1-ubyte"
    let labels1 = runGet (getLabelHeader *> (genVector getLabel)) rawLabel1
    let labels2 = runGet (getLabelHeader *> (genVector getLabel)) rawLabel2
    let images1 = runGet (getImageHeader *> (genVector getImage)) rawImage1
    let images2 = runGet (getImageHeader *> (genVector getImage)) rawImage2
    return $ zip (images1++images2) (labels1++labels2)


{-# LANGUAGE FlexibleInstances #-}

import Debug.Trace
-- import Control.Monad.State.Lazy

import System.Environment (getArgs)
import System.IO (readFile)
import Data.Map.Strict (Map, (!), insert, elems, fromList, toList, findWithDefault, size, empty, member, findMin, findMax, singleton, filter)
import qualified Data.Map.Strict as M
import Data.Set (Set)
import qualified Data.Set as S
--import qualified Data.Array as A
import Data.List (find, intercalate, intersperse, permutations, inits, tails, isPrefixOf)
import Data.List.Split (splitOn, chunksOf)
-- import Data.Complex (Complex((:+)), realPart, imagPart) -- define my own complex

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

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

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

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

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

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

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

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

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

data Complex a = C a a

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

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

type Point = Complex Int
type Direction = Complex Int
type Scaffold = Set Point
data Move = R | L | F Int | Err
  deriving (Show, Eq)

show' :: [Move] -> String
show' ms = intercalate "," $ map showMove ms
  where showMove R = "R"
        showMove L = "L"
        showMove (F n) = show n
        showMove Err = error "showing err"

charToDirection '>' = C 1 0
charToDirection '<' = C (-1) 0
charToDirection '^' = C 0 (-1)
charToDirection 'v' = C 0 1
charToDirection c = error $ "bad direction" ++ [c]

fromJust' :: String -> Maybe a -> a
fromJust' s Nothing = error s
fromJust' s (Just x) = x

path :: Scaffold -> (Point, Direction) -> [Move]
path scaffold (pt, dir)
  | S.size scaffold < 2 = []
  -- | otherwise = (traceShowId $ rotation ++ [F steps]) ++ path scaffold' (pt', dir')
  | otherwise = rotation ++ [F steps] ++ path scaffold' (pt', dir')
  where left :: Complex Int
        left = C 0 (-1)
        dir' = fromJust' "dir'" $ find (\d -> (pt + d) `S.member` scaffold) [left^i | i <- [1 .. 4]]
        directions = [(left, [L]), (left^2, [L,L]), (left^3, [R]), (left^4, [Err])]
        rotation = snd $ fromJust' ("rotation: "++ show (dir', directions))  $
          find (\(d, m) -> dir * d == dir') directions
        (pt', steps, scaffold') = walk (pt, 0, scaffold)
        walk (p, n, s) | not $ (p+dir') `S.member` s = (p, n, s)
                       | (p+dir'*left) `S.member` s = walk (p+dir', n+1, s) -- intersection, don't delete
                       | otherwise = walk (p+dir', n+1, S.delete p s)

testIntersection scaffold pt = all (`S.member` scaffold) [ pt + d | d <- [l, l*l, l*l*l, l*l*l*l]]
  where l = C 0 1

data Function = A | B | C' deriving (Show)

substitute :: [Move] -> [Function]
-- substitute (L: F n: ms) = A: (substitute $ F (n-8): ms)
-- substitute (F n: ms) = B n: substitute ms
-- substitute (R: F n: ms) = C': (substitute $ F (n-8): ms)
-- substitute [] = []
-- substitute ms = error $ "substitute: can't do " ++ show ms
substitute (L: F 12: L: F 12: R: F 12: ms) = A: substitute ms
substitute (L: F 8: L: F 8: R: F 12: L: F 8: L: F 8: ms) = B: substitute ms
substitute (L: F 10: R: F 8: R: F 12: ms) = C': substitute ms
substitute [] = []
substitute ms = error $ "substitute: can't do " ++ show ms

findIndex :: (Show a) => (a->Bool) -> [a] -> Maybe Int
findIndex p lst | null lst = Nothing
                | p $ head lst = return 0
                | otherwise = (findIndex p $ tail lst) >>= return . fromIntegral . (+ 1)
  -- where test predicate [] = Nothing
  --       test predicate (x:xs) = if predicate x then 0 else 1 + test predicate xs

findIndexes :: (Show a) => [(a->Bool)] -> [a] -> [Int]
findIndexes ps lst | null ps = []
                   | otherwise = case findIndex (head ps) lst of
                                   Nothing -> []
                                   Just x -> x: (map (+ x) $ findIndexes (tail ps) $ drop x lst)
main = do
  -- [instructionFile] <- getArgs
  instructionStrings <- readFile "19.input.txt" -- instructionFile
  let instructions = fromList . zip [0 ..] $ map read $ splitOn "," instructionStrings

  putStrLn "Part 1"
  let outputToChar 0 = ' '
      outputToChar 1 = '#'
      outputToChar _ = error "bad output"
  let outputs = [ [head $ output computer1 | i <- [0 .. 49]
                                           , let computer1 = run $ computer0{memory = instructions, input=[i,j]}
                                           ] | j <- [0 .. 49]]
  let view = unlines $ map (map outputToChar) $ outputs
  putStrLn view
  putStrLn $ show $ sum $ map sum outputs

  putStrLn "Part 2"
  putStrLn $ unlines $ map show $ zip [0 .. ] $ map (findIndex (== 1)) outputs
  putStrLn $ unlines $ map show $ zip [0 .. ] $ map (findIndexes [(== 0),(== 1),(== 0)]) outputs

  let getValueAtPoint i j = -- traceShow (i,j) $ traceShowId $
        head $ output $ run computer0{memory = instructions, input=map fromIntegral [i,j]}

  -- let lineValue y startx stopx = (stopx' - startx', startx', stopx')
  --       where startx' = startx - 1 + (fromJust $  getIndex (== 1) [getValueAtPoint x y | x <- [(startx-1) ..]])
  --             stopx' = stopx - 1 + (fromJust $ getIndex (== 0) [getValueAtPoint x y | x <- [(stopx-1) ..]])

  let search f x y
      -- x y are the top right coordinates of the box (xy increasing right and  down).
        | x < 100 = traceShow "move right" $
          search f 100 y
        | x > 10000 || y > 10000 = error "searched too far"
        | f (x-99) (y+99) == 0 = traceShow (x,y, "move diagonally") $
          search f (x+1) (y+1)
        | f (x+1) y == 1 = traceShow (x,y, "move right") $
          search f (x+1) y
        | f x y == 0 = traceShow (x,y, "move down") $
          search f x (y+1)
        | otherwise = (x,y)

  let (x,y) = search (getValueAtPoint) 0 0
  let topright@(x',y') = (x-99, y)
  print $ (topright, x'*10000 + y')

  -- let getTransition valueGetter isTransition start
  --       | isTransition $ valueGetter start = Nothing -- started too late
  --       | otherwise = findIndex isTransition [valueGetter i | i <- [start .. ]] >>= return . (+ start)

  -- let test j = do
  --       let i = traceShow ("Starting search on line", j, "column") $ traceShowId $ abs $ j-150
  --       stopx <- traceShow ("At line", j, "transitioning from 1 to 0 at column") $ traceShowId $
  --         getTransition (`getValueAtPoint` j) (== 0) i
  --       let startx_searchStart = stopx - 200
  --       startx <- traceShow ("At line", j+99, "transitioning from 0 to 1 at column ") $ traceShowId $
  --         getTransition (`getValueAtPoint` (j+100)) (== 1) startx_searchStart
  --       let getLine y = [getValueAtPoint x y | x <- [0 .. 1000]]
  --       return $ (stopx - startx, map (findIndexes [(== 1), (== 0)]) [getLine j, getLine (j+99)])

  -- interact (unlines . map (show . test . read) . lines)


#!/usr/bin/env stack
-- stack --resolver lts-12.4 script
{-# LANGUAGE DeriveDataTypeable, ParallelListComp #-}
module Main ( main ) where

import Hledger

import Control.Exception ( bracket )
import Control.Monad
import System.Exit
import System.IO
import Data.Time.Calendar
import Text.Printf
import Data.Function (on)
import Data.List
import Data.Ord
import Statistics.Math.RootFinding
import Data.Decimal
import qualified Data.Text as T
import System.Console.CmdArgs
import Text.Tabular as Tbl
import qualified Text.Tabular.AsciiArt as Ascii

data Options = Options
  { cashFlow    :: Bool
  , debug       :: Bool
  , file        :: FilePath
  , investments :: String
  , pnl         :: String
  , begin       :: Maybe String
  , end         :: Maybe String
  , daily       :: Bool
  , weekly      :: Bool
  , monthly     :: Bool
  , yearly      :: Bool
  , period      :: Maybe String
  } deriving (Show, Data, Typeable)

queryFromOptions kind day accessor opts =
  -- We ignore QueryOpts with fst as they are irrelevant for this program
  case parseQuery day (T.pack $ accessor opts) of
    (Any, _) -> error $ "you need to specify " ++ kind ++ " query"
    (query, []) -> query
    (_,     opts) -> error $ "this program cannot process query options " ++ show opts

intervalAndSpanFromOptions :: Day -> Options -> (Maybe Interval, DateSpan)
intervalAndSpanFromOptions date opts =
  let begin' = fmap parsedate $ begin opts
      end' = fmap parsedate $ end opts
  in
   case (begin', end', daily opts, weekly opts, monthly opts, yearly opts, period opts) of
     (Nothing, Nothing, False, False, False, False, Nothing) -> (Nothing, nulldatespan)
     (b, e, True,  False, False, False, Nothing) -> (Just (Days 1),   DateSpan b e)
     (b, e, False, True,  False, False, Nothing) -> (Just (Weeks 1),  DateSpan b e)
     (b, e, False, False, True,  False, Nothing) -> (Just (Months 1), DateSpan b e)
     (b, e, False, False, False, True,  Nothing) -> (Just (Years 1),  DateSpan b e)
     (Nothing, Nothing, False, False, False, False, Just pexp) ->
       let (i, s) = parsePeriodExpr' date (T.pack pexp) in (Just i, s)
     _ -> error "Cannot work with this combination of date/period flags"

options =
  Options {
     cashFlow =
      False &= name "c"
      &= help "also show all relevant transactions"
    , debug =
      def
      &= help "print debugging info"
    , file =
      def &= name "f"
      &= typFile
      &= help "input ledger file (pass '-' for stdin)"
    , investments =
      def &= name "i"
      &= typ "QUERY"
      &= help "investments query"
    , pnl =
      def &= name "P"
      &= typ "QUERY"
      &= help "profit-and-loss query"
    , begin =
      Nothing &= name "b"
      &= typ "DATE"
      &= help "calculate interest from this date"
    , end =
      Nothing &= name "e"
      &= typ "DATE"
      &= help "calculate interest until this date"
    , daily =
      def &= name "D"
      &= help "calculate interest for each day"
    , weekly =
      def &= name "W"
      &= help "calculate interest for each week"
    , monthly =
      def &= name "M"
      &= help "calculate interest for each month"
    , yearly =
      def &= name "Y"
      &= help "calculate interest for each year"
    , period =
      Nothing &= name "p"
      &= typ "PERIODEXP"
      &= help "set start date, end date, and/or report interval all at once (overrides the flags above)"
    }
  &= program "hledger-roi"
  &= summary "compute return-on-investment for your portfolio using IRR and TWR"
  &= details [""]

main ::  IO ()
main = bracket (return ()) (\() -> hFlush stdout >> hFlush stderr) $ \() -> do
  opts <-cmdArgs options
  thisDay <- getCurrentDay

  let investmentsQuery = queryFromOptions "investments" thisDay investments opts
  let pnlQuery         = queryFromOptions "pnl" thisDay pnl opts

  jnl <- readJournalFile definputopts (file opts) >>= either fail return

  let trans = jtxns $ filterJournalTransactions investmentsQuery jnl

  when (null trans) $ do
    putStrLn "No relevant transactions found. Check your investments query"
    exitFailure

  when (debug opts) $ do
    putStrLn "DEBUG everything:"
    print $ trans
    putStrLn "DEBUG just non-PnL transactions:"
    print $ filter (matchesTransaction (And [investmentsQuery, Not pnlQuery])) trans
    putStrLn "DEBUG just PnL transactions:"
    print $ filter (matchesTransaction pnlQuery) trans

  let (requestedInterval, requestedSpan) = intervalAndSpanFromOptions thisDay opts
  let existingSpan =
        let dates = map transactionDate2 trans in
        DateSpan (Just $ minimum dates) (Just $ addDays 1 $ maximum dates)
  let wholeSpan = spanDefaultsFrom requestedSpan existingSpan

  let spans = case requestedInterval of
        Nothing -> [wholeSpan]
        Just interval ->
            splitSpan interval $
            spanIntersect existingSpan wholeSpan

  tableBody <- (flip mapM) spans $ \(DateSpan (Just ibegin) (Just iend)) -> do
    -- Spans are [begin,end), and iend here is 1 days after actual end date we are interested in
    let valueBeforeThisPeriod =
          total trans (And [ investmentsQuery
                           , Date (DateSpan Nothing (Just ibegin))])

    let valueAtTheEndOfPeriod  =
          total trans (And [investmentsQuery, Date (DateSpan Nothing (Just iend))])

    let cashFlowInThisPeriod =
          calculateCashFlow trans (And [ Not investmentsQuery
                                       , Not pnlQuery
                                       , Date (DateSpan (Just ibegin) (Just iend)) ] )

    irr <- internalRateOfReturn opts ibegin iend valueBeforeThisPeriod valueAtTheEndOfPeriod cashFlowInThisPeriod
    twr <- timeWeightedReturn opts investmentsQuery trans ibegin iend valueBeforeThisPeriod valueAtTheEndOfPeriod cashFlowInThisPeriod
    let cashFlowAmt = negate $ sum $ map snd cashFlowInThisPeriod
    return [ showDate ibegin
           , showDate (addDays (-1) iend)
           , show valueBeforeThisPeriod
           , show cashFlowAmt
           , show valueAtTheEndOfPeriod
           , show (valueAtTheEndOfPeriod - (valueBeforeThisPeriod + cashFlowAmt))
           , printf "%0.2f%%" irr
           , printf "%0.2f%%" twr ]

  let table = Tbl.Table
              (Group NoLine (map (Header . show) (take (length tableBody) [1..])))
              (Group DoubleLine
               [ Group SingleLine [Header "Begin", Header "End"]
               , Group SingleLine [Header "Value (begin)", Header "Cashflow", Header "Value (end)", Header "PnL"]
               , Group SingleLine [Header "IRR", Header "TWR"]])
              tableBody

  putStrLn $ Ascii.render id id id table

timeWeightedReturn opts investmentsQuery trans ibegin iend valueBefore valueAfter cashFlowInThisPeriod = do
  -- Span is [ibegin,iend), and iend is 1 days after actual end date we are interested in
  let initialUnitPrice = 100
  let initialUnits = valueBefore / initialUnitPrice
  let cashflow =
        -- Aggregate all entries for a single day, assuming that intraday interest is negligible
        map (\date_cash -> let (dates, cash) = unzip date_cash in (head dates, sum cash))
        $ groupBy ((==) `on` fst)
        $ sortBy (comparing fst)
        $ map (\(d,a) -> (d, negate a))
        $ filter ((/=0).snd) cashFlowInThisPeriod

  let units =
        tail $
        (flip scanl)
        (0,0,0,initialUnits)
        (\(_,_,_,unitBalance) (date, amt) ->
          let valueOnDate =
                total trans (And [investmentsQuery, Date (DateSpan Nothing (Just date))])
              unitPrice = if unitBalance == 0.0 then initialUnitPrice else valueOnDate / unitBalance
              unitsBoughtOrSold = amt / unitPrice
          in
           (valueOnDate, unitsBoughtOrSold, unitPrice, unitBalance + unitsBoughtOrSold)
        )
        cashflow

  let finalUnitBalance = if null units then initialUnits else let (_,_,_,u) = last units in u
      finalUnitPrice = valueAfter / finalUnitBalance
      totalTWR = roundTo 2 $ (finalUnitPrice - initialUnitPrice)
      years = (fromIntegral $ diffDays iend ibegin)/365 :: Double
      annualizedTWR = 100*((1+(realToFrac totalTWR/100))**(1/years)-1) :: Double

  let s d = show $ roundTo 2 d
  when (cashFlow opts) $ do
    printf "\nTWR cash flow for %s - %s\n" (showDate ibegin) (showDate (addDays (-1) iend))
    let (dates', amounts') = unzip cashflow
        (valuesOnDate',unitsBoughtOrSold', unitPrices', unitBalances') = unzip4 units
        add x lst = if valueBefore/=0 then x:lst else lst
        dates = add ibegin dates'
        amounts = add valueBefore amounts'
        unitsBoughtOrSold = add initialUnits unitsBoughtOrSold'
        unitPrices = add initialUnitPrice unitPrices'
        unitBalances = add initialUnits unitBalances'
        valuesOnDate = add 0 valuesOnDate'

    putStr $ Ascii.render id id id
      (Tbl.Table
       (Group NoLine (map (Header . showDate) dates))
       (Group DoubleLine [ Group SingleLine [Header "Portfolio value", Header "Unit balance"]
                         , Group SingleLine [Header "Cash", Header "Unit price", Header "Units"]
                         , Group SingleLine [Header "New Unit Balance"]])
       [ [value, oldBalance, amt, prc, udelta, balance]
       | value <- map s valuesOnDate
       | oldBalance <- map s (0:unitBalances)
       | balance <- map s unitBalances
       | amt <- map s amounts
       | prc <- map s unitPrices
       | udelta <- map s unitsBoughtOrSold ])

    printf "Final unit price: %s/%s=%s U.\nTotal TWR: %s%%.\nPeriod: %.2f years.\nAnnualized TWR: %.2f%%\n\n" (s valueAfter) (s finalUnitBalance) (s finalUnitPrice) (s totalTWR) years annualizedTWR

  return annualizedTWR

internalRateOfReturn opts ibegin iend valueBefore valueAfter cashFlowInPeriod = do
  -- Span is [ibegin,iend), and iend is 1 days after actual end date we are interested in
  let prefix = (ibegin, negate valueBefore)

      postfix = (iend, valueAfter)

      totalCF = filter ((/=0) . snd) $ prefix : (sortBy (comparing fst) cashFlowInPeriod) ++ [postfix]

  when (cashFlow opts) $ do
    printf "\nIRR cash flow for %s - %s\n" (showDate ibegin) (showDate (addDays (-1) iend))
    let (dates, amounts) = unzip totalCF
    putStrLn $ Ascii.render id id id
      (Tbl.Table
       (Group NoLine (map (Header . showDate) dates))
       (Group SingleLine [Header "Amount"])
       (map ((:[]) . show) amounts))

  -- 0% is always a solution, so require at least something here
  case ridders 0.00001 (0.000000000001,10000) (interestSum iend totalCF) of
    Root rate -> return ((rate-1)*100)
    NotBracketed -> error "Error: No solution -- not bracketed."
    SearchFailed -> error "Error: Failed to find solution."

type CashFlow = [(Day, Quantity)]

interestSum :: Day -> CashFlow -> Double -> Double
interestSum referenceDay cf rate = sum $ map go cf
    where go (t,m) = (fromRational $ toRational m) * (rate ** (fromIntegral (referenceDay `diffDays` t) / 365))

calculateCashFlow :: [Transaction] -> Query -> CashFlow
calculateCashFlow trans query = map go trans
    where
    go t = (transactionDate2 t, total [t] query)

total :: [Transaction] -> Query -> Quantity
total trans query = unMix $ sumPostings $ filter (matchesPosting query) $ concatMap realPostings trans

unMix :: MixedAmount -> Quantity
unMix a =
  case (normaliseMixedAmount $ costOfMixedAmount a) of
    (Mixed [a]) -> aquantity a
    _ -> error "MixedAmount failed to normalize"