#!/usr/bin/env stack
{- stack script --resolver lts-19.23 --install-ghc -}
{-# OPTIONS_GHC
  -Wall
  -Wno-unused-imports
  -Wno-unused-top-binds
  -Wno-name-shadowing
#-}
{-# language NoImplicitPrelude #-}
{-# language RecordWildCards #-}
{-# language FlexibleContexts #-}
{-# language TemplateHaskell #-}
{-# language ViewPatterns #-}
{-# language OverloadedStrings #-}
{-# language DeriveGeneric #-}
{-# language DerivingVia #-}
{-# language TupleSections #-}
import Rebase.Prelude hiding (toList, on, (<+>), Max)
import qualified Rebase.Prelude as Prelude
import qualified Data.Set as Set
import qualified Data.Sequence as Seq
import qualified Data.Sequence.NonEmpty as NESeq
import Data.Sequence (Seq(..))
import Data.Sequence.NonEmpty (NESeq(..))
import Control.Lens hiding ((<|))
import Data.Foldable (toList)
import Data.Ratio
import Text.Printf
import Graphics.Vty
import Data.Time.LocalTime
import Control.Monad.RWS
import Data.Time.Calendar.OrdinalDate
import qualified Data.Text as Text
import Data.Text.Format.Numbers
import Rebase.Data.Map.Strict (Map)
import qualified Rebase.Data.Map.Strict as Map

import Brick
import Brick.Types
import Data.Text (unpack)
import Control.Lens
import Control.Monad (void, forever)
import Control.Concurrent (threadDelay, forkIO)
import qualified Graphics.Vty as V
import Brick.Widgets.ProgressBar as P
import Brick.BChan
import Brick.Widgets.Center
import Brick.Widgets.Border
import Brick.Main
  ( App(..)
  , showFirstCursor
  , customMain
  , continue
  , halt
  )
import qualified Brick.Main as M
import Brick.AttrMap
  ( attrMap
  )
import Brick.Types
  ( Widget
  , Next
  , EventM
  , BrickEvent(..)
  )
import Brick.Widgets.Core
  ( (<=>)
  , str
  )
import Brick.AttrMap as A
import Brick.Util (fg, bg, on, clamp)
import Brick.Widgets.Core
import Brick.Widgets.Table
import Text.Pretty.Simple (pPrint)

import GHC.Generics (Generic)
import Generic.Data
import Generic.Data.Microsurgery

data Performance = Achieved {
  achievedReps :: Integer,
  achievedWeight :: Integer
} deriving Generic deriving (Read, Show) via Surgery Derecordify Performance

data LiftRecord = LiftRecord {
  liftRecordName :: Text,
  stats :: [Performance]
} deriving Generic deriving (Read, Show) via Surgery Derecordify LiftRecord

-- "I highly recommend having a goal in mind for these last sets. Sit down the
-- night before, or the week before, and think of the number of reps you’d like
-- to hit. See yourself doing it. Write it down and visualize the bar in your
-- hands or on your back. When it’s time, let yourself go and attack the
-- weight."

computeRepGoal :: Integer -> [Performance] -> (Integer)
computeRepGoal targetWeight stats = head $ filter isPR [2..]
  where
    isPR n = computeOneRepMax (Achieved n targetWeight) > computeOneRepMax (bestPerformance stats)

bestPerformance :: [Performance] -> Performance
bestPerformance = head . sortBy (flip $ comparing computeOneRepMax)

-- The formula from Jim Wendler's 5-3-1:
-- 1RM = Weight x Reps x 0.0333 + Weight.
computeOneRepMax :: Performance -> Rational
computeOneRepMax Achieved{..} = toRational achievedWeight * (realToFrac achievedReps * 0.0333 + 1)

showRealFrac' :: (RealFrac n, PrintfArg n ) => n -> String
showRealFrac' n = printf  "%.0f" n

showRealFrac :: (RealFrac n, PrintfArg n ) => n -> String
showRealFrac n = printf format n
  where
    format = if floor (n * 10) `mod` 10 == (0 :: Integer) then "%.0f" else "%.2f"

showRational :: Rational -> String
showRational = (showRealFrac :: Float -> String) . fromRational

data WeekSelection = Week1 | Week2 | Week3 | Week4 deriving (Enum, Bounded, Show, Eq)

data Programming = Wendler | Accessory Int Int | Unspecified

data Lift = Lift {
  liftName :: Text,
  liftProgramming :: Programming
}

data Session = Session {
  sessionName :: Text,
  sessionLifts :: NESeq Lift
}

-- TODO: State contains chosen repmax formula
-- TODO: State contains performances
data St = St {
  _lifts :: [LiftRecord],
  _week :: WeekSelection,
  _sessions :: NESeq (Session),
  _selectedSession :: Int
}
makeLenses ''St

-- TODO: Event for inotify on edited text file (as input interface)
data CustomEvent = CustomEvent

data SetScheme = SetScheme { ratioMax :: Rational, numberOfReps :: Integer, amrap :: Bool }

setScheme, setSchemeAMRAP :: Rational -> Integer -> SetScheme
setScheme percentage reps = SetScheme (percentage / 100) reps False
setSchemeAMRAP percentage reps = SetScheme (percentage / 100) reps True

data SetOption = Fresher | Heavier

warmup :: [SetScheme]
warmup = [setScheme 40 5, setScheme 50 5, setScheme 60 3]

downSets :: [SetScheme]
downSets = take 5 $ repeat $ setScheme 85 5

liftScheme :: SetOption -> WeekSelection -> [SetScheme]
liftScheme Fresher Week1 = warmup ++ [setScheme 65 5, setScheme 75 5, setSchemeAMRAP 85 5] ++ downSets
liftScheme Heavier Week1 = warmup ++ [setScheme 75 5, setScheme 80 5, setSchemeAMRAP 85 5] ++ downSets
liftScheme Fresher Week2 = warmup ++ [setScheme 70 3, setScheme 80 3, setSchemeAMRAP 90 3] ++ downSets
liftScheme Heavier Week2 = warmup ++ [setScheme 80 3, setScheme 85 3, setSchemeAMRAP 90 3] ++ downSets
liftScheme _______ Week3 = warmup ++ [setScheme 75 5, setScheme 85 3, setSchemeAMRAP 95 1] ++ downSets
liftScheme _______ Week4 = (flip setScheme 5) <$> [40, 50, 60, 70, 80] -- Wendler stops at 60

data ListPosition = FirstInList | NotFirstInList deriving (Eq)

annotatePosition :: [a] -> [(ListPosition, a)]
annotatePosition [] = undefined
annotatePosition (x:xs) = (FirstInList, x) : map (NotFirstInList,) xs

lookup' :: Int -> NESeq a -> a
lookup' i seq = fromJust $ NESeq.lookup (i `mod` NESeq.length seq) seq

liftRows :: WeekSelection -> (Lift, Maybe LiftRecord) -> [[Widget n]]
liftRows week (Lift _name Wendler, Just liftRecord) =
  (flip map) (annotatePosition $ liftScheme Fresher week) $ wendlerLiftRow liftRecord
liftRows _ (Lift name (Accessory sets reps), liftRecords) =
  (flip map) (annotatePosition $ take sets $ repeat reps) $ accessoryLiftRow liftRecords name reps
liftRows _ (Lift name Unspecified, _) = [unspecifiedLiftRow name]

accessoryLiftRow :: Maybe LiftRecord -> Text -> Int -> (ListPosition, Int) -> [Widget n]
accessoryLiftRow liftRecords name targetReps (position, _) =
  map (padLeftRight 2)
  [ txt $ case position of FirstInList -> name; NotFirstInList -> " "
  , str $ printf "%2d @ %3s" targetReps $ fromMaybe (" " :: String) (show <$> targetWeight)
  , txt " "
  , txt " "
  , txt " "
  , txt " "
  , txt " "
  ]
  where
    targetWeight :: Maybe Integer
    targetWeight = join $ fmap achievedWeight . find ((== targetReps) . fromIntegral . achievedReps) . stats <$> liftRecords

unspecifiedLiftRow :: Text -> [Widget n]
unspecifiedLiftRow name =
  map (padLeftRight 2)
  [ txt name
  , txt " "
  , txt " "
  , txt " "
  , txt " "
  , txt " "
  , txt " "
  ]

wendlerLiftRow :: LiftRecord -> (ListPosition, SetScheme) -> [Widget n]
wendlerLiftRow LiftRecord{..} (position, (SetScheme targetRatio targetReps amrap)) =
  let best              = bestPerformance stats
      computedTarget    = targetRatio * (90 % 100) * computeOneRepMax best
      targetWeight      = ceilingN 5 computedTarget
      repGoal           = computeRepGoal targetWeight stats
      goalTo1RM g       = computeOneRepMax $ Achieved g targetWeight
      showGoal g        = printf "%2d @ %d ≈ 1 @ %s" g targetWeight (showRational (goalTo1RM g))
      pad               = padBottom (Pad $ if amrap then 1 else 0)
      actualIntensity   = fromRational $ 100 * computedTarget / computeOneRepMax best :: Float
  in
  map (padLeftRight 2)
  [
    pad $ txt $ case position of FirstInList -> liftRecordName; NotFirstInList -> " ",
    pad $ str $ printf "%2d%s @ %d" targetReps (if amrap then "+" else "" :: Text) targetWeight,
    pad $ str $ showPlates targetWeight,
    pad $ str $ showRealFrac' actualIntensity ++ "%",
    str $ if amrap then showGoal repGoal ++ "\n" ++ showGoal (repGoal + 1) else " ",
    pad $ txt " ",
    pad $ txt " "
  ]


drawUI :: St -> [Widget ()]
drawUI st = [vCenter $ vBox [hCenter $ hBox [header, oneRepMaxTable], withVScrollBarHandles $ withVScrollBars OnRight $ viewport () Vertical $ hCenter lastSetTable]]
  where
    session :: Session
    session = view sessions st & lookup' (view selectedSession st)
    sameName a b = liftName a == liftRecordName b
    lifts' :: [(Lift, Maybe LiftRecord)]
    lifts' = session & sessionLifts & toList & map (\lift -> (lift, find (sameName lift) $ view lifts st))
    liftNames :: [Text]
    liftNames = session & sessionLifts & fmap liftName & toList
    liftRecords :: [LiftRecord]
    liftRecords = filter ((`elem` liftNames) . liftRecordName) (view lifts st)
    lastSetTable = renderTable $ table $ map (padLeftRight 1 . txt) ["Lift", "Set", "Plates", "%1RM", "Goal", "Done", "Rest"] : concatMap (liftRows (view week st)) lifts'
    header = renderTable $ table $ map (padLeftRight 1 . txt)
      ["Date", "Time", "Bodyweight", "Week", "Session"]
      : [ map (padLeftRight 2 . txt) ["        ", "     ", "     ", weekNumber, sessionName session] ]
    weekNumber = case (view week st) of Week1 -> "1"; Week2 -> "2"; Week3 -> "3"; Week4 -> "4"
    oneRepMaxTable = renderTable $ table $ map (padLeftRight 1 . txt) ["Lift", "Achieved Best", "Computed 1RM"] : map toRow liftRecords
    toRow LiftRecord{..} =
      let best@Achieved{..} = bestPerformance stats
      in
      map (padLeftRight 2)
      [
        txt $ liftRecordName,
        str $ printf "%d @ %d" achievedReps achievedWeight,
        str $ showRational $ computeOneRepMax best
      ]

showPlates :: Integer -> String
showPlates (fromIntegral -> wt) = fromMaybe " " $ fmap (ourShow . reverse) $ showPlates' (wt - 45) ourPlates []
  where
    ourPlates = [45,45,25,25,10,10,10,5,5,2.5]
    ourShow :: [Rational] -> String
    ourShow = concat . intersperse " " . map showRational
    showPlates' 0 _ [] = Nothing
    showPlates' 0 _ used = Just used
    showPlates' need (have:avail) used | have * 2 > need = showPlates' need avail used
    showPlates' need (have:avail) used = showPlates' (need - have * 2) avail (have:used)
    showPlates' _ _ _ = Nothing

ceilingN :: Integer -> Rational -> Integer
ceilingN n x = ceiling (x / toRational n) * n

succ' :: (Enum a, Bounded a, Eq a) => a -> a
succ' x | x == maxBound = minBound
succ' x = succ x

clipSelectedSession :: St -> St
clipSelectedSession st = st & selectedSession %~ maybeReset
  where
    len = view (sessions . to NESeq.length) st
    maybeReset n | n < 0 = len + n
    maybeReset n | n < len = n
    maybeReset _ | otherwise = 0

handleEvent :: St -> BrickEvent () CustomEvent -> EventM () (Next St)
handleEvent st e = case e of
  VtyEvent (V.EvKey V.KEsc [])         -> halt st
  VtyEvent (V.EvKey V.KDown _)         -> M.vScrollBy (M.viewportScroll ()) 5    >> continue st
  VtyEvent (V.EvKey V.KUp _)           -> M.vScrollBy (M.viewportScroll ()) (-5) >> continue st
  VtyEvent (V.EvKey (V.KChar 'w') _)   -> continue $ st & week %~ succ'
  VtyEvent (V.EvKey (V.KChar 'n') _)   -> continue $ st & selectedSession %~ (+1) & clipSelectedSession
  VtyEvent (V.EvKey (V.KChar 'p') _)   -> continue $ st & selectedSession %~ (subtract 1) & clipSelectedSession
  VtyEvent _                           -> continue st
  AppEvent _                           -> continue st
  _ -> continue st

theApp :: App St CustomEvent ()
theApp =
    App { appDraw = drawUI
        , appChooseCursor = showFirstCursor
        , appHandleEvent = handleEvent
        , appStartEvent = return
        , appAttrMap = const $ A.attrMap V.defAttr []
        }

-- Lift Ratios
-- from https://www.catalystathletics.com/article/1786/The-Relation-of-Snatch-Clean-Jerk-and-Squat-Weights/
squat, deadlift, press :: Rational
squat = 345
deadlift = 453
press = 140

cleanAndJerk = squat * 80 / 100
frontSquat = squat * 85 / 100
overheadSquat = squat * 65 / 100
snatch = squat * 60 / 100
powerSnatch = snatch * 80 / 100

clean = deadlift * 70 / 100
powerClean = clean * 80 / 100

pushPress = press * 100 / 75
jerk = pushPress / 100 * 85

program :: NESeq Session
program =

   ( Session "5-3-1 Deadlift/Press" $

     Lift "Deadlift" Wendler
     :<|| Seq.fromList
     [ Lift "Press" Wendler
     , Lift "Row" $ Accessory 5 10
     , Lift "Dip" $ Unspecified
     , Lift "Good Morning" $ Accessory 3 10
     ]
   )
  :<|| Seq.fromList

  [ Session "Front Squat" $

    Lift "Front Squat" Wendler :<||
    Seq.fromList
    [ Lift "Press front & back" $ Accessory 1 10
    , Lift "Squat" $ Accessory 1 5
    , Lift "Good Morning" $ Accessory 1 5
    , Lift "Front Squat" $ Accessory 1 5
    , Lift "SLDL" $ Accessory 1 5
    , Lift "Snatch high pull from thighs" $ Accessory 1 5
    , Lift "Snatch high pull from knees" $ Accessory 1 5
    , Lift "Snatch high pull from shins" $ Accessory 1 5
    , Lift "Muscle Snatch" $ Accessory 1 5
    , Lift "Overhead Squat" $ Accessory 1 5
    , Lift "Left-Arm Snatch" $ Accessory 5 3
    , Lift "Right-Arm Snatch" $ Accessory 5 3
    ]

    , Session "5-3-1 Squat/Bench" $

      Lift "Squat" Wendler
      :<|| Seq.fromList
      [ Lift "Bench" Wendler
      , Lift "Row" $ Accessory 5 10
      , Lift "Back Raise" $ Accessory 3 15
      , Lift "DB Press" $ Accessory 3 10
      , Lift "Rear Lateral" $ Accessory 3 20
      ]

    , Session "Power Clean" $

      Lift "Power Clean" Wendler
      :<|| Seq.fromList
      [ Lift "Push Press" Wendler
      ]

    , Session "Power Lifting" $

      Lift "Deadlift" Wendler
      :<|| Seq.fromList
      [ Lift "Squat" Wendler
      ,  Lift "Bench" Wendler
      ,  Lift "Press" Wendler
      ]

    , Session "Accessories" $

      Lift "Row" (Accessory 5 10)
      :<|| Seq.fromList
      [ Lift "DB Press" $ Accessory 3 10
      , Lift "Dip" $ Unspecified
      , Lift "Good Morning" $ Accessory 3 10
      , Lift "Rear Lateral" $ Accessory 3 20
      , Lift "Back Raise" $ Accessory 3 15
      ]

    , Session "Naked Bar Warmup" $

      Lift "Press front & back" (Accessory 1 10)
      :<|| Seq.fromList
      [ Lift "Back Squat" $ Accessory 1 5
      , Lift "Good Morning" $ Accessory 1 5
      , Lift "Front Squat" $ Accessory 1 5
      , Lift "SLDL" $ Accessory 1 5
      , Lift "Snatch high pull from thighs" $ Accessory 1 5
      , Lift "Snatch high pull from knees" $ Accessory 1 5
      , Lift "Snatch high pull from shins" $ Accessory 1 5
      , Lift "Muscle Snatch" $ Accessory 1 5
      , Lift "Overhead Squat" $ Accessory 1 5
      , Lift "Left-Arm Snatch" $ Accessory 5 3
      , Lift "Right-Arm Snatch" $ Accessory 5 3
      ]
    ]

main :: IO ()
main = do
    let buildVty = V.mkVty V.defaultConfig
    vty <- buildVty
    -- liftIO $ setMode (outputIface vty) Mouse True
    chan <- newBChan 10
    initial <- read <$> readFile "lifts.dat" :: IO [LiftRecord]
    void $ customMain vty buildVty (Just chan) (theApp) (St initial Week1 program 0)