{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE OverloadedLists       #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE TemplateHaskell       #-}
module RCNN where

import           Control.Applicative               (ZipList (..))
import           Control.Lens                      (_1, _2, makePrisms)
import           Control.Monad.Trans.Resource
import qualified Data.Array.Repa                   as Repa
import           Formatting                        (sformat, string, (%))
import           Options.Applicative               (Parser, ReadM, auto,
                                                    eitherReader, help, long,
                                                    metavar, option,
                                                    showDefault, strOption,
                                                    value)
import           RIO
import           RIO.Directory                     (canonicalizePath,
                                                    doesFileExist)
import qualified RIO.HashSet                       as S
import           RIO.List                          (lastMaybe, unzip, unzip3)
import           RIO.List.Partial                  (maximum)
import qualified RIO.NonEmpty                      as RNE
import qualified RIO.NonEmpty.Partial              as RNE
import qualified RIO.Text                          as T
import qualified RIO.Vector.Boxed                  as V

import           MXNet.Base
import           MXNet.Base.ParserUtils            (decimal, endOfInput, list,
                                                    parseOnly, rational)
import           MXNet.NN
import qualified MXNet.NN.DataIter.Anchor          as Anchor
import qualified MXNet.NN.DataIter.Coco            as Coco
import qualified MXNet.NN.Initializer              as I
import           MXNet.NN.ModelZoo.RCNN.FasterRCNN

instance Coco.HasDatasetConfig (FeiApp t n Coco.CocoConfig) where
    type DatasetTag (FeiApp t n Coco.CocoConfig) = "coco"
    datasetConfig = fa_extra

data CommonArgs = CommonArgs
    { ds_base_path       :: String
    , ds_img_size        :: Int
    , ds_img_pixel_means :: [Float]
    , ds_img_pixel_stds  :: [Float]
    }
    deriving Show

data ExtraArgs = TrainArgs
    { pg_train_epochs         :: Int
    , pg_train_iter_per_epoch :: Int
    }
    | NoExtraArgs
    deriving Show

apRcnn :: (Parser RcnnConfiguration, Parser RcnnConfiguration)
apRcnn = (train, infr)
    where
        train = RcnnConfigurationTrain
                <$> backbone
                <*> batch_size
                <*> feature_strides
                <*> strOption        (long "pretrained"        <> metavar "PATH"
                                                               <> value ""
                                                               <> help "path to pretrained model")
                <*> option floatx4   (long "bbox-reg-stds"     <> metavar "BBOX_STDS"
                                                               <> value (0.1, 0.1, 0.2, 0.2))
                <*> rpn_anchor_scales
                <*> rpn_anchor_ratios
                <*> rpn_base_size
                <*> rpn_pre_topk
                <*> rpn_post_topk
                <*> rpn_nms_threshold
                <*> rpn_min_size
                <*> option auto      (long "rpn-batch-rois"    <> metavar "BATCH-ROIS"
                                                               <> showDefault
                                                               <> value 256
                                                               <> help "rpn number of rois per batch")
                <*> option auto      (long "rpn-fg-fraction"   <> metavar "FG-FRACTION"
                                                               <> showDefault
                                                               <> value 0.5
                                                               <> help "rpn foreground fraction")
                <*> option auto      (long "rpn-fg-overlap"    <> metavar "FG-OVERLAP"
                                                               <> showDefault
                                                               <> value 0.7
                                                               <> help "rpn foreground iou threshold")
                <*> option auto      (long "rpn-bg-overlap"    <> metavar "BG-OVERLAP"
                                                               <> showDefault
                                                               <> value 0.3
                                                               <> help "rpn background iou threshold")
                <*> option auto      (long "rpn-allowed-border"<> metavar "ALLOWED-BORDER"
                                                               <> showDefault
                                                               <> value 0
                                                               <> help "rpn allowed border")
                <*> rcnn_num_classes
                <*> rcnn_pool_sized
                <*> rcnn_batch_rois
                <*> option auto      (long "rcnn-fg-fraction"  <> metavar "FG-FRACTION"
                                                               <> showDefault
                                                               <> value 0.25
                                                               <> help "rcnn foreground fraction")
                <*> option auto      (long "rcnn-fg-overlap"   <> metavar "FG-OVERLAP"
                                                               <> showDefault
                                                               <> value 0.5
                                                               <> help "rcnn foreground iou threshold")
                <*> option auto      (long "rcnn-max-num-gt"   <> metavar "NUM-GT"
                                                               <> showDefault
                                                               <> value 100
                                                               <> help "rcnn max number of gt")

        infr = RcnnConfigurationInference
                <$> backbone
                <*> batch_size
                <*> feature_strides
                <*> strOption       (long "checkpoint" <> metavar "PATH"
                                                       <> value ""
                                                       <> help "path to a saved model")
                <*> option floatx4  (long "bbox-reg-stds" <> metavar "BBOX_STDS" <> value (0.1, 0.1, 0.2, 0.2))
                <*> rpn_anchor_scales
                <*> rpn_anchor_ratios
                <*> rpn_base_size
                <*> rpn_pre_topk
                <*> rpn_post_topk
                <*> rpn_nms_threshold
                <*> rpn_min_size
                <*> rcnn_num_classes
                <*> rcnn_pool_sized
                <*> rcnn_batch_rois
                <*> option auto     (long "rcnn-force-nms" <> metavar "FORCE_NMS" <> value False)
                <*> option auto     (long "rcnn-nms-threshold" <> metavar "NMS_THRESH" <> value 0.5)
                <*> option auto     (long "rcnn-nms-topk" <> metavar "NMS_TOPK" <> value (-1))


        rpn_anchor_scales = option intList   (long "rpn-anchor-scales" <> metavar "SCALES"
                                                              <> showDefault
                                                              <> value [8,16,32]
                                                              <> help "rpn anchor scales")
        rpn_anchor_ratios = option floatList (long "rpn-anchor-ratios" <> metavar "RATIOS"
                                                              <> showDefault
                                                              <> value [0.5,1,2]
                                                              <> help "rpn anchor ratios")
        rpn_base_size     = option auto      (long "rpn-anchor-bsize"  <> metavar "BSIZE"
                                                              <> showDefault
                                                              <> value 16
                                                              <> help "rpn anchor base size")
        rpn_pre_topk      = option auto      (long "rpn-pre-nms-topk"  <> metavar "PRE-NMS-TOPK"
                                                              <> showDefault
                                                              <> value 12000
                                                              <> help "rpn nms pre-top-k")
        rpn_post_topk     = option auto      (long "rpn-post-nms-topk" <> metavar "POST-NMS-TOPK"
                                                              <> showDefault
                                                              <> value 2000
                                                              <> help "rpn nms post-top-k")
        rpn_nms_threshold = option auto      (long "rpn-nms-thresh" <> metavar "NMS-THRESH"
                                                              <> showDefault
                                                              <> value 0.7
                                                              <> help "rpn nms threshold")
        rpn_min_size      = option auto      (long "rpn-min-size" <> metavar "MIN-SIZE"
                                                               <> showDefault
                                                               <> value 16
                                                               <> help "rpn min size")
        rcnn_num_classes  = option auto      (long "rcnn-num-classes"  <> metavar "NUM-CLASSES"
                                                              <> showDefault
                                                              <> value 81
                                                              <> help "rcnn number of classes")
        rcnn_batch_rois   = option auto      (long "rcnn-batch-rois"   <> metavar "BATCH_ROIS"
                                                              <> showDefault
                                                              <> value 128
                                                              <> help "rcnn batch rois")
        rcnn_pool_sized   = option auto      (long "rcnn-pooled-size"  <> metavar "POOLED-SIZE"
                                                              <> showDefault
                                                              <> value 14
                                                              <> help "rcnn pooled size")
        feature_strides   = option intList   (long "strides"  <> metavar "STRIDE"
                                                              <> showDefault
                                                              <> value [16]
                                                              <> help "feature stride")
        batch_size        = option auto      (long "batch-size" <> metavar "BATCH-SIZE"
                                                              <> showDefault
                                                              <> value 1
                                                              <> help "batch size")
        backbone          = option auto      (long "backbone" <> metavar "BACKBONE"
                                                              <> value VGG16
                                                              <> help "vgg-16 or resnet-50")

apCommon :: Parser CommonArgs
apCommon = CommonArgs
       <$> strOption        (long "base"              <> metavar "PATH"
                                                      <> help "path to the dataset")
       <*> option auto      (long "img-size"          <> metavar "SIZE"
                                                      <> showDefault
                                                      <> value 1024
                                                      <> help "long side of image")
       <*> option floatList (long "img-pixel-means"   <> metavar "RGB-MEAN"
                                                      <> showDefault
                                                      <> value [0,0,0]
                                                      <> help "RGB mean of images")
       <*> option floatList (long "img-pixel-stds"    <> metavar "RGB-STDS"
                                                      <> showDefault
                                                      <> value [1,1,1]
                                                      <> help "RGB std-dev of images")

apTrain :: Parser ExtraArgs
apTrain = TrainArgs
       <$> option auto      (long "train-epochs"      <> metavar "EPOCHS"
                                                      <> value 500
                                                      <> help "number of epochs to train")
       <*> option auto      (long "train-iter-per-epoch" <> metavar "ITER-PER-EPOCH"
                                                      <> value 100
                                                      <> help "number of iter per epoch")

floatList :: ReadM [Float]
floatList = eitherReader $ parseOnly (list rational<* endOfInput) . T.pack

floatx4 :: ReadM (Float, Float, Float, Float)
floatx4 = let p = do fs <- list rational
                     case fs of
                       [a, b, c, d] -> return (a, b, c, d)
                       _            -> fail "should be exactly 4 floats"
           in eitherReader $ parseOnly (p <* endOfInput) . T.pack

intList :: ReadM [Int]
intList   = eitherReader $ parseOnly (list decimal <* endOfInput) . T.pack

toTriple [a, b, c] = (a, b, c)
toTriple x         = error (show x)


default_initializer :: Initializer Float
default_initializer name = case name of
    "features.rpn.rpn_conv_3x3.weight"  -> I.normal 0.01 name
    "features.rpn.rpn_conv_3x3.bias"    -> I.zeros name
    "features.rpn.rpn_cls_score.weight" -> I.normal 0.01 name
    "features.rpn.rpn_cls_score.bias"   -> I.zeros name
    "features.rpn.rpn_bbox_pred.weight" -> I.normal 0.01 name
    "features.rpn.rpn_bbox_pred.bias"   -> I.zeros name
    "features.rcnn.rcnn_cls_score.weight"     -> I.normal 0.01 name
    "features.rcnn.rcnn_cls_score.bias"       -> I.zeros name
    "features.rcnn.rcnn_bbox_feature.weight"  -> I.normal 0.01 name
    "features.rcnn.rcnn_bbox_feature.bias"    -> I.zeros name
    "features.rcnn.rcnn_bbox_pred.weight"     -> I.normal 0.001 name
    "features.rcnn.rcnn_bbox_pred.bias"       -> I.zeros name
    "features.fpn.0.conv1.weight"             -> I.xavier 1 I.XavierUniform I.XavierIn name
    "features.fpn.0.conv2.weight"             -> I.xavier 1 I.XavierUniform I.XavierIn name
    "features.fpn.1.conv1.weight"             -> I.xavier 1 I.XavierUniform I.XavierIn name
    "features.fpn.1.conv2.weight"             -> I.xavier 1 I.XavierUniform I.XavierIn name
    "features.fpn.2.conv1.weight"             -> I.xavier 1 I.XavierUniform I.XavierIn name
    "features.fpn.2.conv2.weight"             -> I.xavier 1 I.XavierUniform I.XavierIn name
    "features.fpn.3.conv1.weight"             -> I.xavier 1 I.XavierUniform I.XavierIn name
    "features.fpn.3.conv2.weight"             -> I.xavier 1 I.XavierUniform I.XavierIn name
    _ | T.isSuffixOf ".running_mean" name -> I.zeros name
      | T.isSuffixOf ".running_var"  name -> I.ones name
      | T.isSuffixOf ".beta"         name -> I.zeros name
      | T.isSuffixOf ".gamma"        name -> I.ones  name
      | otherwise -> I.zeros name

loadWeights :: (DType a, MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack)
            => String -> Module t a m ()
loadWeights weights_path = do
    weights_path <- liftIO $ canonicalizePath weights_path
    e <- liftIO $ doesFileExist (weights_path ++ ".params")
    if not e
        then lift . logInfo . display $ sformat ("'" % string % ".params' doesn't exist.") weights_path
        else loadState weights_path ["features.9.gamma",
                                     "features.9.beta",
                                     "features.9.running_var",
                                     "features.9.running_mean",
                                     "output.weight",
                                     "output.bias"
                                    ]

data Stage = TRAIN
    | INFERENCE

fixedParams :: Backbone -> Stage -> SymbolHandle -> IO (HashSet Text)
fixedParams backbone stage symbol = do
    argnames <- listArguments symbol
    return $ case (stage, backbone) of
        (INFERENCE, _)
            -> S.fromList argnames
        (TRAIN, VGG16)
            -> S.fromList [n | n <- argnames
                          -- fix conv_1_1, conv_1_2, conv_2_1, conv_2_2
                          ,  elemM [0, 2, 5, 7] (layer n)]
        (TRAIN, r) | r `elem` ([RESNET50, RESNET50FPN, RESNET101] :: [Backbone])
            -> S.fromList [n | n <- argnames
                          -- fix conv_0, stage_1_*, *_gamma, *_beta
                          , let layer_idx = layer n
                          , elemM [0, 1, 5] layer_idx ||
                            (leqM 9 layer_idx && elemM ["gamma", "beta"] (lastName n))]

  where
    toMaybe = either (const Nothing) Just
    layer param = case T.split (=='.') param of
                    "features":n:_ -> toMaybe $ parseOnly decimal n
                    _              -> Nothing
    lastName = lastMaybe . T.split (=='.')
    elemM :: Eq a => [a] -> Maybe a -> Bool
    elemM b = isJust . (>>= guard) . liftM (`elem` b)
    leqM  n = isJust . (>>= guard) . liftM (<= n)

data App c = App LogFunc c
makePrisms ''App

instance HasLogFunc (App c) where
    logFuncL = _App . _1

instance Coco.HasDatasetConfig (App Coco.CocoConfig) where
    type DatasetTag (App Coco.CocoConfig) = "coco"
    datasetConfig = _App . _2

runApp :: c -> ReaderT (App c) (ResourceT IO) a -> IO a
runApp conf body = do
    logopt <- logOptionsHandle stdout False
    runResourceT $ withLogFunc logopt $ \logfunc ->
        flip runReaderT (App logfunc conf) body

generateTargets :: (SymbolHandle -> Layer (NonEmpty SymbolHandle))
                -> Coco.ImageInfo
                -> NonEmpty Int
                -> Anchor.Configuration
                -> [Anchor.GTBox Repa.U]
                -> IO (NDArray Float, NDArray Float, NDArray Float)
generateTargets feature_net im_info strides anchor_conf gt_boxes = do
    feats  <- runLayerBuilder $ variable "data" >>= feature_net

    -- there should equally number of features and strides, and pair them.
    let feat_stride = RNE.zip feats strides

    layers <- mapM (uncurry make) feat_stride
    let (cls_targets, box_targets, box_masks) = unzip3 $ RNE.toList layers
    cls_targets <- mapM fromRepa cls_targets
    box_targets <- mapM fromRepa box_targets
    box_masks   <- mapM fromRepa box_masks
    cls_targets <- concat_ 0 cls_targets
    box_targets <- concat_ 0 box_targets
    box_masks   <- concat_ 0 box_masks
    return (cls_targets, box_targets, box_masks)

  where
    [img_height, img_width, _] = Repa.toList im_info
    -- we have padded the image to a square
    img_size = floor (max img_height img_width)
    base_size = anchor_conf ^. Anchor.conf_anchor_base_size
    scales    = anchor_conf ^. Anchor.conf_anchor_scales
    ratios    = anchor_conf ^. Anchor.conf_anchor_ratios
    make :: SymbolHandle -> Int -> IO (Anchor.Labels, Anchor.Targets, Anchor.Weights)
    make feat stride = do
        (_, outputs, _, _) <- inferShape feat [("data", STensor [1,3,img_size,img_size])]
        let [(_, STensor [_, _, h, w])] = outputs
            anchors = Anchor.anchors (h, w) stride base_size scales ratios
        runReaderT (Anchor.assign (V.fromList gt_boxes) img_size img_size anchors) anchor_conf

padLength :: DType a => [NDArray a] -> a -> IO [NDArray a]
padLength arrays value = do
    shps <- mapM ndshape arrays
    let max_num = maximum $ map RNE.head shps
    forM (zip arrays shps) $ \(a, n :| shp) ->
        if n == max_num
        then return a
        else do
            padding <- full value ((max_num - n) :| shp)
            concat_ 0 [a, padding]

withRpnTargets :: MonadIO m
               => RcnnConfiguration
               -> (String, Coco.ImageTensor, Coco.ImageInfo, Coco.GTBoxes)
               -> m (String, [NDArray Float])
withRpnTargets RcnnConfigurationTrain{..} dat = liftIO $ do
    (cls_targets, box_targets, box_weights) <-
        generateTargets extract info (RNE.fromList feature_strides) conf (V.toList gt)
    imgA  <- fromRepa img
    infoA <- fromRepa info
    gtA   <- stack 0 . V.toList =<< mapM fromRepa gt
    return (filename, [gtA, imgA, infoA, cls_targets, box_targets, box_weights])
    where
        (filename, img, info, gt) = dat
        conf = Anchor.Configuration
               { Anchor._conf_anchor_scales    = rpn_anchor_scales
               , Anchor._conf_anchor_ratios    = rpn_anchor_ratios
               , Anchor._conf_anchor_base_size = rpn_anchor_base_size
               , Anchor._conf_allowed_border   = rpn_allowd_border
               , Anchor._conf_fg_num           = floor $ (rpn_fg_fraction * fromIntegral rpn_batch_rois)
               , Anchor._conf_batch_num        = rpn_batch_rois
               , Anchor._conf_bg_overlap       = rpn_bg_overlap
               , Anchor._conf_fg_overlap       = rpn_fg_overlap
               }
        extract = sequential "features" . features1 backbone


withRpnTargets'Mask :: MonadIO m
                    => RcnnConfiguration
                    -> (String, Coco.ImageTensor, Coco.ImageInfo, Coco.GTBoxes, Coco.Masks)
                    -> m (String, [NDArray Float])
withRpnTargets'Mask conf dat = do
    let (filename, img, info, gt, msks) = dat
    (_, ret) <- withRpnTargets conf (filename, img, info, gt)
    liftIO $ do
        msksA <- stack 0 . V.toList =<< mapM fromRepa msks
        msksA <- divScalar 255 =<< cast #float32 msksA :: IO (NDArray Float)
        return (filename, msksA : ret)

toListNDArray :: MonadIO m
               => (String, Coco.ImageTensor, Coco.ImageInfo, Coco.GTBoxes)
               -> m (String, [NDArray Float])
toListNDArray (filename, img, info, gt) = liftIO $ do
    imgA  <- fromRepa img
    infoA <- fromRepa info
    gtA   <- stack 0 . V.toList =<< mapM fromRepa gt
    return (filename, [gtA, imgA, infoA])

concatBatch :: MonadIO m => [(String, [NDArray Float])] -> m ([String], [NDArray Float])
concatBatch batch = liftIO $ do
    let (filenames, tensors) = unzip batch
        gt : others = unzipList tensors
    -- gt in the batch may not have the same number
    -- must be padded with -1 before stacking
    gt     <- stack 0 =<< padLength gt (-1)
    -- other tensors can be simply stacked
    others <- mapM (stack 0) others
    return (filenames, gt : others)


concatBatch'Mask :: MonadIO m => [(String, [NDArray Float])] -> m ([String], [NDArray Float])
concatBatch'Mask batch = liftIO $ do
    let (filenames, tensors) = unzip batch
        mask_gt : box_gt : others = unzipList tensors
    mask_gt <- stack 0 =<< padLength mask_gt 0
    box_gt  <- stack 0 =<< padLength box_gt (-1)
    others <- mapM (stack 0) others
    return (filenames, mask_gt : box_gt : others)

unzipList :: [[a]] -> [[a]]
unzipList = getZipList . traverse ZipList