content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE TypeFamilies #-}
  
  module Futhark.AD.Rev.Hist
    ( diffMinMaxHist,
      diffMulHist,
      diffAddHist,
      diffVecHist,
      diffHist,
    )
  where
  
  import Control.Monad
  import Futhark.AD.Rev.Monad
  import Futhark.Analysis.PrimExp.Convert
  import Futhark.Builder
  import Futhark.IR.SOACS
  import Futhark.Tools
  import Futhark.Transform.Rename
  
  getBinOpPlus :: PrimType -> BinOp
  getBinOpPlus (IntType x) = Add x OverflowUndef
  getBinOpPlus (FloatType f) = FAdd f
  getBinOpPlus _ = error "In getBinOpMul, Hist.hs: input not supported"
  
  getBinOpDiv :: PrimType -> BinOp
  getBinOpDiv (IntType t) = SDiv t Unsafe
  getBinOpDiv (FloatType t) = FDiv t
  getBinOpDiv _ = error "In getBinOpDiv, Hist.hs: input not supported"
  
  withinBounds :: [(SubExp, VName)] -> TPrimExp Bool VName
  withinBounds [] = TPrimExp $ ValueExp (BoolValue True)
  withinBounds [(q, i)] = (le64 i .<. pe64 q) .&&. (pe64 (intConst Int64 (-1)) .<. le64 i)
  withinBounds (qi : qis) = withinBounds [qi] .&&. withinBounds qis
  
  elseIf ::
    (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
    PrimType ->
    [(m (Exp (Rep m)), m (Exp (Rep m)))] ->
    [m (Body (Rep m))] ->
    m (Exp (Rep m))
  elseIf t [(c1, c2)] [bt, bf] =
    eIf
      (eCmpOp (CmpEq t) c1 c2)
      bt
      bf
  elseIf t ((c1, c2) : cs) (bt : bs) =
    eIf
      (eCmpOp (CmpEq t) c1 c2)
      bt
      $ eBody
      $ pure
      $ elseIf t cs bs
  elseIf _ _ _ = error "In elseIf, Hist.hs: input not supported"
  
  bindSubExpRes :: (MonadBuilder m) => Name -> [SubExpRes] -> m [VName]
  bindSubExpRes s =
    traverse
      ( \(SubExpRes cs se) -> do
          bn <- newVName s
          certifying cs $ letBindNames [bn] $ BasicOp $ SubExp se
          pure bn
      )
  
  nestedmap :: [SubExp] -> [PrimType] -> Lambda SOACS -> ADM (Lambda SOACS)
  nestedmap [] _ lam = pure lam
  nestedmap s@(h : r) pt lam = do
    params <- traverse (\tp -> newParam "x" $ Array tp (Shape s) NoUniqueness) pt
    body <- nestedmap r pt lam
    mkLambda params $
      fmap varsRes . letTupExp "res" . Op $
        Screma h (map paramName params) (mapSOAC body)
  
  -- \ds hs -> map2 lam ds hs
  mkF' :: Lambda SOACS -> [Type] -> SubExp -> ADM ([VName], [VName], Lambda SOACS)
  mkF' lam tps n = do
    lam' <- renameLambda lam
  
    ds_params <- traverse (newParam "ds_param") tps
    hs_params <- traverse (newParam "hs_param") tps
    let ds_pars = fmap paramName ds_params
    let hs_pars = fmap paramName hs_params
    lam_map <-
      mkLambda (ds_params <> hs_params) $
        fmap varsRes . letTupExp "map_f'" . Op $
          Screma n (ds_pars <> hs_pars) (mapSOAC lam')
  
    pure (ds_pars, hs_pars, lam_map)
  
  -- \ls as rs -> map3 (\li ai ri -> li `lam` ai `lam` ri) ls as rs
  mkF :: Lambda SOACS -> [Type] -> SubExp -> ADM ([VName], Lambda SOACS)
  mkF lam tps n = do
    lam_l <- renameLambda lam
    lam_r <- renameLambda lam
    let q = length $ lambdaReturnType lam
        (lps, aps) = splitAt q $ lambdaParams lam_l
        (ips, rps) = splitAt q $ lambdaParams lam_r
    lam' <- mkLambda (lps <> aps <> rps) $ do
      lam_l_res <- bodyBind $ lambdaBody lam_l
      forM_ (zip ips lam_l_res) $ \(ip, SubExpRes cs se) ->
        certifying cs $ letBindNames [paramName ip] $ BasicOp $ SubExp se
      bodyBind $ lambdaBody lam_r
  
    ls_params <- traverse (newParam "ls_param") tps
    as_params <- traverse (newParam "as_param") tps
    rs_params <- traverse (newParam "rs_param") tps
    let map_params = ls_params <> as_params <> rs_params
    lam_map <-
      mkLambda map_params $
        fmap varsRes . letTupExp "map_f" $
          Op $
            Screma n (map paramName map_params) $
              mapSOAC lam'
  
    pure (map paramName as_params, lam_map)
  
  mapout :: VName -> SubExp -> SubExp -> ADM VName
  mapout is n w = do
    par_is <- newParam "is" $ Prim int64
    is'_lam <-
      mkLambda [par_is] $
        fmap varsRes . letTupExp "is'"
          =<< eIf
            (toExp $ withinBounds $ pure (w, paramName par_is))
            (eBody $ pure $ eParam par_is)
            (eBody $ pure $ eSubExp w)
  
    letExp "is'" $ Op $ Screma n (pure is) $ mapSOAC is'_lam
  
  multiScatter :: [VName] -> VName -> [VName] -> ADM [VName]
  multiScatter dst is vs =
    doScatter "scatter_res" 1 dst (is : vs) $ pure . map (Var . paramName)
  
  multiIndex :: [VName] -> [DimIndex SubExp] -> ADM [VName]
  multiIndex vs s = do
    traverse
      ( \x -> do
          t <- lookupType x
          letExp "sorted" $ BasicOp $ Index x (fullSlice t s)
      )
      vs
  
  --
  -- special case of histogram with min/max as operator.
  -- Original, assuming `is: [n]i64` and `dst: [w]btp`
  --     let x = reduce_by_index dst minmax ne is vs
  -- Forward sweep:
  --     need to copy dst: reverse sweep might use it 7
  --       (see ex. in reducebyindexminmax6.fut where the first map requires the original dst to be differentiated).
  --     let dst_cpy = copy dst
  --     let (x, x_inds) = zip vs (iota n)
  --                       |> reduce_by_index (dst_cpy,-1s) argminmax (ne,-1) is
  --
  -- Reverse sweep:
  --     dst_bar += map2 (\i b -> if i == -1
  --                              then b
  --                              else 0
  --                     ) x_inds x_bar
  
  --     vs_ctrbs = map2 (\i b -> if i == -1
  --                              then 0
  --                              else vs_bar[i] + b
  --                     ) x_inds x_bar
  --     vs_bar <- scatter vs_bar x_inds vs_ctrbs
  diffMinMaxHist ::
    VjpOps -> VName -> StmAux () -> SubExp -> BinOp -> SubExp -> VName -> VName -> SubExp -> SubExp -> VName -> ADM () -> ADM ()
  diffMinMaxHist _ops x aux n minmax ne is vs w rf dst m = do
    let t = binOpType minmax
    vs_type <- lookupType vs
    let vs_elm_type = elemType vs_type
    let vs_dims = arrayDims vs_type
    let inner_dims = tail vs_dims
    let nr_dims = length vs_dims
    dst_type <- lookupType dst
    let dst_dims = arrayDims dst_type
  
    dst_cpy <-
      letExp (baseName dst <> "_copy") . BasicOp $
        Replicate mempty (Var dst)
  
    acc_v_p <- newParam "acc_v" $ Prim t
    acc_i_p <- newParam "acc_i" $ Prim int64
    v_p <- newParam "v" $ Prim t
    i_p <- newParam "i" $ Prim int64
    hist_lam_inner <-
      mkLambda [acc_v_p, acc_i_p, v_p, i_p] $
        fmap varsRes . letTupExp "idx_res"
          =<< eIf
            (eCmpOp (CmpEq t) (eParam acc_v_p) (eParam v_p))
            ( eBody
                [ eParam acc_v_p,
                  eBinOp (SMin Int64) (eParam acc_i_p) (eParam i_p)
                ]
            )
            ( eBody
                [ eIf
                    ( eCmpOp
                        (CmpEq t)
                        (eParam acc_v_p)
                        (eBinOp minmax (eParam acc_v_p) (eParam v_p))
                    )
                    (eBody [eParam acc_v_p, eParam acc_i_p])
                    (eBody [eParam v_p, eParam i_p])
                ]
            )
    hist_lam <- nestedmap inner_dims [vs_elm_type, int64, vs_elm_type, int64] hist_lam_inner
  
    dst_minus_ones <-
      letExp "minus_ones" . BasicOp $
        Replicate (Shape dst_dims) (intConst Int64 (-1))
    ne_minus_ones <-
      letSubExp "minus_ones" . BasicOp $
        Replicate (Shape inner_dims) (intConst Int64 (-1))
    iota_n <-
      letExp "red_iota" . BasicOp $
        Iota n (intConst Int64 0) (intConst Int64 1) Int64
  
    inp_iota <- do
      if nr_dims == 1
        then pure iota_n
        else do
          i <- newParam "i" $ Prim int64
          lam <-
            mkLambda [i] $
              fmap varsRes . letTupExp "res" =<< do
                pure $ BasicOp $ Replicate (Shape inner_dims) $ Var $ paramName i
  
          letExp "res" $ Op $ Screma n [iota_n] $ mapSOAC lam
  
    let hist_op = HistOp (Shape [w]) rf [dst_cpy, dst_minus_ones] [ne, if nr_dims == 1 then intConst Int64 (-1) else ne_minus_ones] hist_lam
    f' <- mkIdentityLambda [Prim int64, rowType vs_type, rowType $ Array int64 (Shape vs_dims) NoUniqueness]
    x_inds <- newVName (baseName x <> "_inds")
    auxing aux $
      letBindNames [x, x_inds] $
        Op $
          Hist n [is, vs, inp_iota] [hist_op] f'
  
    m
  
    x_bar <- lookupAdjVal x
  
    x_ind_dst <- newParam (baseName x <> "_ind_param") $ Prim int64
    x_bar_dst <- newParam (baseName x <> "_bar_param") $ Prim t
    dst_lam_inner <-
      mkLambda [x_ind_dst, x_bar_dst] $
        fmap varsRes . letTupExp "dst_bar"
          =<< eIf
            (toExp $ le64 (paramName x_ind_dst) .==. -1)
            (eBody $ pure $ eParam x_bar_dst)
            (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)
    dst_lam <- nestedmap inner_dims [int64, vs_elm_type] dst_lam_inner
  
    dst_bar <-
      letExp (baseName dst <> "_bar") . Op $
        Screma w [x_inds, x_bar] (mapSOAC dst_lam)
  
    updateAdj dst dst_bar
  
    vs_bar <- lookupAdjVal vs
  
    inds' <- traverse (letExp "inds" . BasicOp . Replicate (Shape [w]) . Var) =<< mk_indices inner_dims []
    let inds = x_inds : inds'
  
    par_x_ind_vs <- replicateM nr_dims $ newParam (baseName x <> "_ind_param") $ Prim int64
    par_x_bar_vs <- newParam (baseName x <> "_bar_param") $ Prim t
    vs_lam_inner <-
      mkLambda (par_x_bar_vs : par_x_ind_vs) $
        fmap varsRes . letTupExp "res"
          =<< eIf
            (toExp $ le64 (paramName $ head par_x_ind_vs) .==. -1)
            (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)
            ( eBody $
                pure $ do
                  vs_bar_i <-
                    letSubExp (baseName vs_bar <> "_el") . BasicOp $
                      Index vs_bar . Slice $
                        fmap (DimFix . Var . paramName) par_x_ind_vs
                  eBinOp (getBinOpPlus t) (eParam par_x_bar_vs) (eSubExp vs_bar_i)
            )
    vs_lam <- nestedmap inner_dims (vs_elm_type : replicate nr_dims int64) vs_lam_inner
  
    vs_bar_p <-
      letExp (baseName vs <> "_partial") . Op $
        Screma w (x_bar : inds) (mapSOAC vs_lam)
  
    q <-
      letSubExp "q"
        =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) dst_dims
  
    scatter_inps <- do
      -- traverse (letExp "flat" . BasicOp . Reshape [DimNew q]) $ inds ++ [vs_bar_p]
      -- ToDo: Cosmin asks: is the below the correct translation of the line above?
      forM (inds ++ [vs_bar_p]) $ \v -> do
        v_t <- lookupType v
        letExp "flat" . BasicOp . Reshape v $
          reshapeAll (arrayShape v_t) (Shape [q])
  
    vs_bar' <-
      fmap head $
        doScatter (baseName vs <> "_bar") nr_dims [vs_bar] scatter_inps $
          pure . map (Var . paramName)
    insAdj vs vs_bar'
    where
      mk_indices :: [SubExp] -> [SubExp] -> ADM [VName]
      mk_indices [] _ = pure []
      mk_indices [d] iotas = do
        reps <- traverse (letExp "rep" . BasicOp . Replicate (Shape [d])) iotas
        iota_d <-
          letExp "red_iota" . BasicOp $
            Iota d (intConst Int64 0) (intConst Int64 1) Int64
        pure $ reps ++ [iota_d]
      mk_indices (d : dims) iotas = do
        iota_d <-
          letExp "red_iota" . BasicOp $
            Iota d (intConst Int64 0) (intConst Int64 1) Int64
  
        i_param <- newParam "i" $ Prim int64
        lam <-
          mkLambda [i_param] $
            fmap varsRes $
              mk_indices dims $
                iotas ++ [Var $ paramName i_param]
  
        letTupExp "res" $ Op $ Screma d [iota_d] $ mapSOAC lam
  
  --
  -- special case of histogram with multiplication as operator.
  -- Original, assuming `is: [n]i64` and `dst: [w]btp`
  --     let x = reduce_by_index dst (*) ne is vs
  -- Forward sweep:
  --     dst does not need to be copied: dst is not overwritten
  --     let (ps, zs) = map (\v -> if v == 0 then (1,1) else (v,0)) vs
  --     let non_zero_prod = reduce_by_index nes (*) ne is ps
  --     let zero_count = reduce_by_index 0s (+) 0 is zs
  --     let h_part = map2 (\p c -> if c == 0 then p else 0
  --                       ) non_zero_prod zero_count
  --     let x = map2 (*) dst h_part
  --
  -- Reverse sweep:
  --     dst_bar += map2 (*) h_part x_bar
  
  --     let part_bar = map2 (*) dst x_bar
  --     vs_bar += map2 (\i v -> let zr_cts = zero_count[i]
  --                             let pr_bar = part_bar[i]
  --                             let nz_prd = non_zero_prod[i]
  --                             in if zr_cts == 0
  --                             then pr_bar * (nz_prd / v)
  --                             else if zr_cts == 1 and v == 0
  --                             then nz_prd * pr_bar
  --                             else 0
  --                    ) is vs
  diffMulHist ::
    VjpOps -> VName -> StmAux () -> SubExp -> BinOp -> SubExp -> VName -> VName -> SubExp -> SubExp -> VName -> ADM () -> ADM ()
  diffMulHist _ops x aux n mul ne is vs w rf dst m = do
    let t = binOpType mul
    vs_type <- lookupType vs
    let vs_dims = arrayDims vs_type
    let vs_elm_type = elemType vs_type
    dst_type <- lookupType dst
    let dst_dims = arrayDims dst_type
    let inner_dims = tail vs_dims
  
    v_param <- newParam "v" $ Prim t
    lam_ps_zs_inner <-
      mkLambda [v_param] $
        fmap varsRes . letTupExp "map_res"
          =<< eIf
            (eCmpOp (CmpEq t) (eParam v_param) (eSubExp $ Constant $ blankPrimValue t))
            (eBody $ fmap eSubExp [Constant $ onePrimValue t, intConst Int64 1])
            (eBody [eParam v_param, eSubExp $ intConst Int64 0])
    lam_ps_zs <- nestedmap vs_dims [vs_elm_type] lam_ps_zs_inner
    ps_zs_res <- eLambda lam_ps_zs [eSubExp $ Var vs]
    ps_zs <- bindSubExpRes "ps_zs" ps_zs_res
    let [ps, zs] = ps_zs
  
    lam_mul_inner <- binOpLambda mul t
    lam_mul <- nestedmap inner_dims [vs_elm_type, vs_elm_type] lam_mul_inner
    nz_prods0 <- letExp "nz_prd" $ BasicOp $ Replicate (Shape [w]) ne
    let hist_nzp = HistOp (Shape [w]) rf [nz_prods0] [ne] lam_mul
  
    lam_add_inner <- binOpLambda (Add Int64 OverflowUndef) int64
    lam_add <- nestedmap inner_dims [int64, int64] lam_add_inner
    zr_counts0 <- letExp "zr_cts" $ BasicOp $ Replicate (Shape dst_dims) (intConst Int64 0)
    zrn_ne <- letSubExp "zr_ne" $ BasicOp $ Replicate (Shape inner_dims) (intConst Int64 0)
    let hist_zrn = HistOp (Shape [w]) rf [zr_counts0] [if length vs_dims == 1 then intConst Int64 0 else zrn_ne] lam_add
  
    f' <- mkIdentityLambda [Prim int64, Prim int64, rowType vs_type, rowType $ Array int64 (Shape vs_dims) NoUniqueness]
    nz_prods <- newVName "non_zero_prod"
    zr_counts <- newVName "zero_count"
    auxing aux $
      letBindNames [nz_prods, zr_counts] $
        Op $
          Hist n [is, is, ps, zs] [hist_nzp, hist_zrn] f'
  
    p_param <- newParam "prod" $ Prim t
    c_param <- newParam "count" $ Prim int64
    lam_h_part_inner <-
      mkLambda [p_param, c_param] $
        fmap varsRes . letTupExp "h_part"
          =<< eIf
            (toExp $ 0 .==. le64 (paramName c_param))
            (eBody $ pure $ eParam p_param)
            (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)
    lam_h_part <- nestedmap dst_dims [vs_elm_type, int64] lam_h_part_inner
    h_part_res <- eLambda lam_h_part $ map (eSubExp . Var) [nz_prods, zr_counts]
    h_part' <- bindSubExpRes "h_part" h_part_res
    let [h_part] = h_part'
  
    lam_mul_inner' <- binOpLambda mul t
    lam_mul' <- nestedmap dst_dims [vs_elm_type, vs_elm_type] lam_mul_inner'
    x_res <- eLambda lam_mul' $ map (eSubExp . Var) [dst, h_part]
    x' <- bindSubExpRes "x" x_res
    auxing aux $ letBindNames [x] $ BasicOp $ SubExp $ Var $ head x'
  
    m
  
    x_bar <- lookupAdjVal x
  
    lam_mul'' <- renameLambda lam_mul'
    dst_bar_res <- eLambda lam_mul'' $ map (eSubExp . Var) [h_part, x_bar]
    dst_bar <- bindSubExpRes (baseName dst <> "_bar") dst_bar_res
    updateAdj dst $ head dst_bar
  
    lam_mul''' <- renameLambda lam_mul'
    part_bar_res <- eLambda lam_mul''' $ map (eSubExp . Var) [dst, x_bar]
    part_bar' <- bindSubExpRes "part_bar" part_bar_res
    let [part_bar] = part_bar'
  
    inner_params <- zipWithM newParam ["zr_cts", "pr_bar", "nz_prd", "a"] $ map Prim [int64, t, t, t]
    let [zr_cts, pr_bar, nz_prd, a_param] = inner_params
    lam_vsbar_inner <-
      mkLambda inner_params $
        fmap varsRes . letTupExp "vs_bar" =<< do
          eIf
            (eCmpOp (CmpEq int64) (eSubExp $ intConst Int64 0) (eParam zr_cts))
            (eBody $ pure $ eBinOp mul (eParam pr_bar) $ eBinOp (getBinOpDiv t) (eParam nz_prd) $ eParam a_param)
            ( eBody $
                pure $
                  eIf
                    ( eBinOp
                        LogAnd
                        (eCmpOp (CmpEq int64) (eSubExp $ intConst Int64 1) (eParam zr_cts))
                        (eCmpOp (CmpEq t) (eSubExp $ Constant $ blankPrimValue t) $ eParam a_param)
                    )
                    (eBody $ pure $ eBinOp mul (eParam nz_prd) (eParam pr_bar))
                    (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)
            )
  
    lam_vsbar_middle <- nestedmap inner_dims [int64, t, t, t] lam_vsbar_inner
  
    i_param <- newParam "i" $ Prim int64
    a_param' <- newParam "a" $ rowType vs_type
    lam_vsbar <-
      mkLambda [i_param, a_param'] $
        fmap varsRes . letTupExp "vs_bar"
          =<< eIf
            (toExp $ withinBounds $ pure (w, paramName i_param))
            ( buildBody_ $ do
                let i = fullSlice vs_type [DimFix $ Var $ paramName i_param]
                names <- traverse newVName ["zr_cts", "pr_bar", "nz_prd"]
                zipWithM_ (\name -> letBindNames [name] . BasicOp . flip Index i) names [zr_counts, part_bar, nz_prods]
                eLambda lam_vsbar_middle $ map (eSubExp . Var) names <> [eParam a_param']
            )
            (eBody $ pure $ pure $ zeroExp $ rowType dst_type)
  
    vs_bar <-
      letExp (baseName vs <> "_bar") $ Op $ Screma n [is, vs] $ mapSOAC lam_vsbar
  
    updateAdj vs vs_bar
  
  --
  -- special case of histogram with add as operator.
  -- Original, assuming `is: [n]i64` and `dst: [w]btp`
  --     let x = reduce_by_index dst (+) ne is vs
  -- Forward sweep:
  --     need to copy dst: reverse sweep might use it 7
  --       (see ex. in reducebyindexminmax6.fut where the first map requires the original dst to be differentiated).
  --     let dst_cpy = copy dst
  --     let x = reduce_by_index dst_cpy (+) ne is vs
  --
  -- Reverse sweep:
  --     dst_bar += x_bar
  --
  --     vs_bar += map (\i -> x_bar[i]) is
  diffAddHist ::
    VjpOps -> VName -> StmAux () -> SubExp -> Lambda SOACS -> SubExp -> VName -> VName -> SubExp -> SubExp -> VName -> ADM () -> ADM ()
  diffAddHist _ops x aux n add ne is vs w rf dst m = do
    let t = paramDec $ head $ lambdaParams add
  
    dst_cpy <-
      letExp (baseName dst <> "_copy") . BasicOp $
        Replicate mempty (Var dst)
  
    f <- mkIdentityLambda [Prim int64, t]
    auxing aux . letBindNames [x] . Op $
      Hist n [is, vs] [HistOp (Shape [w]) rf [dst_cpy] [ne] add] f
  
    m
  
    x_bar <- lookupAdjVal x
  
    updateAdj dst x_bar
  
    x_type <- lookupType x
    i_param <- newParam (baseName vs <> "_i") $ Prim int64
    let i = paramName i_param
    lam_vsbar <-
      mkLambda [i_param] $
        fmap varsRes . letTupExp "vs_bar"
          =<< eIf
            (toExp $ withinBounds $ pure (w, i))
            (eBody $ pure $ pure $ BasicOp $ Index x_bar $ fullSlice x_type [DimFix $ Var i])
            (eBody $ pure $ eSubExp ne)
  
    vs_bar <- letExp (baseName vs <> "_bar") $ Op $ Screma n [is] $ mapSOAC lam_vsbar
    updateAdj vs vs_bar
  
  -- Special case for vectorised combining operator. Rewrite
  --   reduce_by_index dst (map2 op) nes is vss
  -- to
  --   map3 (\dst_col vss_col ne ->
  --           reduce_by_index dst_col op ne is vss_col
  --        ) (transpose dst) (transpose vss) nes |> transpose
  -- before differentiating.
  diffVecHist ::
    VjpOps ->
    VName ->
    StmAux () ->
    SubExp ->
    Lambda SOACS ->
    VName ->
    VName ->
    VName ->
    SubExp ->
    SubExp ->
    VName ->
    ADM () ->
    ADM ()
  diffVecHist ops x aux n op nes is vss w rf dst m = do
    stms <- collectStms_ $ do
      rank <- arrayRank <$> lookupType vss
      let dims = [1, 0] ++ drop 2 [0 .. rank - 1]
  
      dstT <- letExp "dstT" $ BasicOp $ Rearrange dst dims
      vssT <- letExp "vssT" $ BasicOp $ Rearrange vss dims
      t_dstT <- lookupType dstT
      t_vssT <- lookupType vssT
      t_nes <- lookupType nes
  
      dst_col <- newParam "dst_col" $ rowType t_dstT
      vss_col <- newParam "vss_col" $ rowType t_vssT
      ne <- newParam "ne" $ rowType t_nes
  
      f <- mkIdentityLambda (Prim int64 : lambdaReturnType op)
      map_lam <-
        mkLambda [dst_col, vss_col, ne] $ do
          -- TODO Have to copy dst_col, but isn't it already unique?
          dst_col_cpy <-
            letExp "dst_col_cpy" . BasicOp $
              Replicate mempty (Var $ paramName dst_col)
          fmap (varsRes . pure) . letExp "col_res" $
            Op $
              Hist
                n
                [is, paramName vss_col]
                [HistOp (Shape [w]) rf [dst_col_cpy] [Var $ paramName ne] op]
                f
      histT <-
        letExp "histT" . Op $
          Screma (arraySize 0 t_dstT) [dstT, vssT, nes] $
            mapSOAC map_lam
      auxing aux . letBindNames [x] . BasicOp $ Rearrange histT dims
    foldr (vjpStm ops) m stms
  
  --
  -- a step in the radix sort implementation
  -- it assumes the key we are sorting
  -- after is [n]i64 and it is the first VName
  --
  -- local def radix_sort_step [n] 't (xs: [n]t) (get_bit: i32 -> t -> i32)
  --                                  (digit_n: i32): [n]t =
  --   let num x = get_bit (digit_n+1) x * 2 + get_bit digit_n x
  --   let pairwise op (a1,b1,c1,d1) (a2,b2,c2,d2) =
  --     (a1 `op` a2, b1 `op` b2, c1 `op` c2, d1 `op` d2)
  --   let bins = xs |> map num
  --   let flags = bins |> map (\x -> if x == 0 then (1,0,0,0)
  --                                  else if x == 1 then (0,1,0,0)
  --                                  else if x == 2 then (0,0,1,0)
  --                                  else (0,0,0,1))
  --   let offsets = scan (pairwise (+)) (0,0,0,0) flags
  --   let (na,nb,nc,_nd) = last offsets
  --   let f bin (a,b,c,d) = match bin
  --                         case 0 -> a-1
  --                         case 1 -> na+b-1
  --                         case 2 -> na+nb+c-1
  --                         case _ -> na+nb+nc+d-1
  --   let is = map2 f bins offsets
  --   in scatter scratch is xs
  radixSortStep :: [VName] -> [Type] -> SubExp -> SubExp -> SubExp -> ADM [VName]
  radixSortStep xs tps bit n w = do
    -- let is = head xs
    is <- mapout (head xs) n w
  
    num_param <- newParam "num" $ Prim int64
    num_lam <-
      mkLambda [num_param] $
        fmap varsRes . letTupExp "num_res"
          =<< eBinOp
            (Add Int64 OverflowUndef)
            ( eBinOp
                (And Int64)
                (eBinOp (AShr Int64) (eParam num_param) (eSubExp bit))
                (iConst 1)
            )
            ( eBinOp
                (Mul Int64 OverflowUndef)
                (iConst 2)
                ( eBinOp
                    (And Int64)
                    (eBinOp (AShr Int64) (eParam num_param) (eBinOp (Add Int64 OverflowUndef) (eSubExp bit) (iConst 1)))
                    (iConst 1)
                )
            )
  
    bins <- letExp "bins" $ Op $ Screma n [is] $ mapSOAC num_lam
    flag_param <- newParam "flag" $ Prim int64
    flag_lam <-
      mkLambda [flag_param] $
        fmap varsRes . letTupExp "flag_res"
          =<< elseIf
            int64
            (map ((,) (eParam flag_param) . iConst) [0 .. 2])
            (map (eBody . fmap iConst . (\i -> map (\j -> if i == j then 1 else 0) [0 .. 3])) ([0 .. 3] :: [Integer]))
  
    flags <- letTupExp "flags" $ Op $ Screma n [bins] $ mapSOAC flag_lam
  
    scan_params <- traverse (flip newParam $ Prim int64) ["a1", "b1", "c1", "d1", "a2", "b2", "c2", "d2"]
    scan_lam <-
      mkLambda scan_params $
        fmap subExpsRes . mapM (letSubExp "scan_res") =<< do
          uncurry (zipWithM (eBinOp $ Add Int64 OverflowUndef)) $ splitAt 4 $ map eParam scan_params
  
    scan <- scanSOAC $ pure $ Scan scan_lam $ map (intConst Int64) [0, 0, 0, 0]
    offsets <- letTupExp "offsets" $ Op $ Screma n flags scan
  
    ind <- letSubExp "ind_last" =<< eBinOp (Sub Int64 OverflowUndef) (eSubExp n) (iConst 1)
    let i = Slice [DimFix ind]
    nabcd <- traverse newVName ["na", "nb", "nc", "nd"]
    zipWithM_ (\abcd -> letBindNames [abcd] . BasicOp . flip Index i) nabcd offsets
  
    let vars = map Var nabcd
    map_params <- traverse (flip newParam $ Prim int64) ["bin", "a", "b", "c", "d"]
    map_lam <-
      mkLambda map_params $
        fmap varsRes . letTupExp "map_res"
          =<< elseIf
            int64
            (map ((,) (eParam $ head map_params) . iConst) [0 .. 2])
            ( zipWith
                ( \j p ->
                    eBody $
                      pure $ do
                        t <- letSubExp "t" =<< eBinOp (Sub Int64 OverflowUndef) (eParam p) (iConst 1)
                        foldBinOp (Add Int64 OverflowUndef) (intConst Int64 0) (t : take j vars)
                )
                [0 .. 3]
                (tail map_params)
            )
  
    nis <- letExp "nis" $ Op $ Screma n (bins : offsets) $ mapSOAC map_lam
  
    scatter_dst <- traverse (\t -> letExp "scatter_dst" $ BasicOp $ Scratch (elemType t) (arrayDims t)) tps
    multiScatter scatter_dst nis xs
    where
      iConst c = eSubExp $ intConst Int64 c
  
  --
  -- the radix sort implementation
  -- def radix_sort [n] 't (xs: [n]i64) =
  --   let iters = if n == 0 then 0 else 32
  --   in loop xs for i < iters do radix_sort_step xs i64.get_bit (i*2)
  radixSort :: [VName] -> SubExp -> SubExp -> ADM [VName]
  radixSort xs n w = do
    logw <- log2 =<< letSubExp "w1" =<< toExp (pe64 w + 1)
    -- ceil logw by (logw + 1) / 2
    iters <- letSubExp "iters" =<< toExp (untyped (pe64 logw + 1) ~/~ untyped (pe64 (intConst Int64 2)))
  
    types <- traverse lookupType xs
    params <- zipWithM (\x -> newParam (baseName x) . flip toDecl Nonunique) xs types
    i <- newVName "i"
    loopbody <- buildBody_ . localScope (scopeOfFParams params) $
      fmap varsRes $ do
        bit <- letSubExp "bit" =<< toExp (le64 i * 2)
        radixSortStep (map paramName params) types bit n w
  
    letTupExp "sorted" $
      Loop
        (zip params $ map Var xs)
        (ForLoop i Int64 iters)
        loopbody
    where
      log2 :: SubExp -> ADM SubExp
      log2 m = do
        params <- zipWithM newParam ["cond", "r", "i"] $ map Prim [Bool, int64, int64]
        let [cond, r, i] = params
  
        body <- buildBody_ . localScope (scopeOfFParams params) $ do
          r' <- letSubExp "r'" =<< toExp (le64 (paramName r) .>>. 1)
          cond' <- letSubExp "cond'" =<< toExp (bNot $ pe64 r' .==. 0)
          i' <- letSubExp "i'" =<< toExp (le64 (paramName i) + 1)
          pure $ subExpsRes [cond', r', i']
  
        cond_init <- letSubExp "test" =<< toExp (bNot $ pe64 m .==. 0)
  
        l <-
          letTupExp' "log2res" $
            Loop
              (zip params [cond_init, m, Constant $ blankPrimValue int64])
              (WhileLoop $ paramName cond)
              body
  
        let [_, _, res] = l
        pure res
  
  radixSort' :: [VName] -> SubExp -> SubExp -> ADM [VName]
  radixSort' xs n w = do
    iota_n <-
      letExp "red_iota" . BasicOp $
        Iota n (intConst Int64 0) (intConst Int64 1) Int64
  
    radres <- radixSort [head xs, iota_n] n w
    let [is', iota'] = radres
  
    i_param <- newParam "i" $ Prim int64
    let slice = [DimFix $ Var $ paramName i_param]
    map_lam <- mkLambda [i_param] $ varsRes <$> multiIndex (tail xs) slice
  
    sorted <- letTupExp "sorted" $ Op $ Screma n [iota'] $ mapSOAC map_lam
    pure $ iota' : is' : sorted
  
  --
  -- generic case of histogram.
  -- Original, assuming `is: [n]i64` and `dst: [w]btp`
  --   let xs = reduce_by_index dst odot ne is as
  -- Forward sweep:
  -- let h_part = reduce_by_index (replicate w ne) odot ne is as
  -- let xs = map2 odot dst h_part
  -- Reverse sweep:
  -- h_part_bar += f'' dst h_part
  -- dst_bar += f' dst h_part
  
  -- let flag = map (\i -> i == 0 || sis[i] != sis[i-1]) (iota n)
  -- let flag_rev = map (\i -> i==0 || flag[n-i]) (iota n)
  -- let ls = seg_scan_exc odot ne flag sas
  -- let rs = reverse sas |>
  --          seg_scan_exc odot ne flag_rev |> reverse
  -- let f_bar = map (\i -> if i < w && -1 < w
  --                        then h_part_bar[i]
  --                        else 0s
  --                 ) sis
  -- let sas_bar = f f_dst ls sas rs
  -- as_bar += scatter (Scratch alpha n) siota sas_bar
  -- Where:
  --  siota: 'iota n' sorted wrt 'is'
  --  sis: 'is' sorted wrt 'is'
  --  sas: 'as' sorted wrt 'is'
  --  f'' = vjpLambda xs_bar h_part (map2 odot)
  --  f' = vjpLambda xs_bar dst (map2 odot)
  --  f  = vjpLambda f_bar sas (map4 (\di li ai ri -> di odot li odot ai odot ri))
  --  0s is an alpha-dimensional array with 0 (possibly 0-dim)
  diffHist :: VjpOps -> [VName] -> StmAux () -> SubExp -> Lambda SOACS -> [SubExp] -> [VName] -> [SubExp] -> SubExp -> [VName] -> ADM () -> ADM ()
  diffHist ops xs aux n lam0 ne as w rf dst m = do
    as_type <- traverse lookupType $ tail as
    dst_type <- traverse lookupType dst
  
    nes <- traverse (letExp "new_dst" . BasicOp . Replicate (Shape $ pure $ head w)) ne
  
    h_map <- mkIdentityLambda $ Prim int64 : map rowType as_type
    h_part <- traverse (newVName . flip (<>) "_h_part" . baseName) xs
    auxing aux . letBindNames h_part . Op $
      Hist n as [HistOp (Shape w) rf nes ne lam0] h_map
  
    lam0' <- renameLambda lam0
    auxing aux . letBindNames xs . Op $
      Screma (head w) (dst <> h_part) (mapSOAC lam0')
  
    m
  
    xs_bar <- traverse lookupAdjVal xs
  
    (dst_params, hp_params, f') <- mkF' lam0 dst_type $ head w
    f'_adj_dst <- vjpLambda ops (map adjFromVar xs_bar) dst_params f'
    f'_adj_hp <- vjpLambda ops (map adjFromVar xs_bar) hp_params f'
  
    dst_bar' <- eLambda f'_adj_dst $ map (eSubExp . Var) $ dst <> h_part
    dst_bar <- bindSubExpRes "dst_bar" dst_bar'
    zipWithM_ updateAdj dst dst_bar
  
    h_part_bar' <- eLambda f'_adj_hp $ map (eSubExp . Var) $ dst <> h_part
    h_part_bar <- bindSubExpRes "h_part_bar" h_part_bar'
  
    lam <- renameLambda lam0
    lam' <- renameLambda lam0
  
    -- is' <- mapout (head as) n (head w)
    -- sorted <- radixSort' (is' : tail as) n $ head w
    sorted <- radixSort' as n $ head w
    let siota = head sorted
    let sis = head $ tail sorted
    let sas = drop 2 sorted
  
    iota_n <-
      letExp "iota" $ BasicOp $ Iota n (intConst Int64 0) (intConst Int64 1) Int64
  
    par_i <- newParam "i" $ Prim int64
    flag_lam <- mkFlagLam par_i sis
    flag <- letExp "flag" $ Op $ Screma n [iota_n] $ mapSOAC flag_lam
  
    -- map (\i -> (if flag[i] then (true,ne) else (false,vs[i-1]), if i==0 || flag[n-i] then (true,ne) else (false,vs[n-i]))) (iota n)
    par_i' <- newParam "i" $ Prim int64
    let i' = paramName par_i'
    g_lam <-
      mkLambda [par_i'] $
        fmap subExpsRes . mapM (letSubExp "scan_inps") =<< do
          im1 <- letSubExp "i_1" =<< toExp (le64 i' - 1)
          nmi <- letSubExp "n_i" =<< toExp (pe64 n - le64 i')
          let s1 = [DimFix im1]
          let s2 = [DimFix nmi]
  
          -- flag array for left scan
          f1 <- letSubExp "f1" $ BasicOp $ Index flag $ Slice [DimFix $ Var i']
  
          -- array for left scan
          r1 <-
            letTupExp' "r1"
              =<< eIf
                (eSubExp f1)
                (eBody $ fmap eSubExp ne)
                (eBody . fmap (eSubExp . Var) =<< multiIndex sas s1)
  
          -- array for right scan inc flag
          r2 <-
            letTupExp' "r2"
              =<< eIf
                (toExp $ le64 i' .==. 0)
                (eBody $ fmap eSubExp $ Constant (onePrimValue Bool) : ne)
                ( eBody $
                    pure $ do
                      eIf
                        (pure $ BasicOp $ Index flag $ Slice s2)
                        (eBody $ fmap eSubExp $ Constant (onePrimValue Bool) : ne)
                        ( eBody . fmap eSubExp . (Constant (blankPrimValue Bool) :) . fmap Var
                            =<< multiIndex sas s2
                        )
                )
  
          traverse eSubExp $ f1 : r1 ++ r2
  
    -- scan (\(f1,v1) (f2,v2) ->
    --   let f = f1 || f2
    --   let v = if f2 then v2 else g v1 v2
    --   in (f,v) ) (false,ne) (zip flags vals)
    scan_lams <-
      traverse
        ( \l -> do
            f1 <- newParam "f1" $ Prim Bool
            f2 <- newParam "f2" $ Prim Bool
            ps <- lambdaParams <$> renameLambda lam0
            let (p1, p2) = splitAt (length ne) ps
  
            mkLambda (f1 : p1 ++ f2 : p2) $
              fmap varsRes . letTupExp "scan_res" =<< do
                let f = eBinOp LogOr (eParam f1) (eParam f2)
                eIf
                  (eParam f2)
                  (eBody $ f : fmap eParam p2)
                  ( eBody . (f :) . fmap (eSubExp . Var)
                      =<< bindSubExpRes "gres"
                      =<< eLambda l (fmap eParam ps)
                  )
        )
        [lam, lam']
  
    let ne' = Constant (BoolValue False) : ne
  
    scansres <-
      letTupExp "adj_ctrb_scan" . Op $
        Screma n [iota_n] (scanomapSOAC (map (`Scan` ne') scan_lams) g_lam)
  
    let (_ : ls_arr, _ : rs_arr_rev) = splitAt (length ne + 1) scansres
  
    -- map (\i -> if i < w && -1 < w then (xs_bar[i], dst[i]) else (0,ne)) sis
    par_i'' <- newParam "i" $ Prim int64
    let i'' = paramName par_i''
    map_lam <-
      mkLambda [par_i''] $
        fmap varsRes . letTupExp "scan_res"
          =<< eIf
            (toExp $ withinBounds $ pure (head w, i''))
            (eBody . fmap (eSubExp . Var) =<< multiIndex h_part_bar [DimFix $ Var i''])
            ( eBody $ do
                map (\t -> pure $ BasicOp $ Replicate (Shape $ tail $ arrayDims t) (Constant $ blankPrimValue $ elemType t)) as_type
            )
  
    f_bar <- letTupExp "f_bar" $ Op $ Screma n [sis] $ mapSOAC map_lam
  
    (as_params, f) <- mkF lam0 as_type n
    f_adj <- vjpLambda ops (map adjFromVar f_bar) as_params f
  
    -- map (\i -> rs_arr_rev[n-i-1]) (iota n)
    par_i''' <- newParam "i" $ Prim int64
    let i''' = paramName par_i'''
    rev_lam <- mkLambda [par_i'''] $ do
      nmim1 <- letSubExp "n_i_1" =<< toExp (pe64 n - le64 i''' - 1)
      varsRes <$> multiIndex rs_arr_rev [DimFix nmim1]
  
    rs_arr <- letTupExp "rs_arr" $ Op $ Screma n [iota_n] $ mapSOAC rev_lam
  
    sas_bar <-
      bindSubExpRes "sas_bar"
        =<< eLambda f_adj (map (eSubExp . Var) $ ls_arr <> sas <> rs_arr)
  
    scatter_dst <- traverse (\t -> letExp "scatter_dst" $ BasicOp $ Scratch (elemType t) (arrayDims t)) as_type
    as_bar <- multiScatter scatter_dst siota sas_bar
  
    zipWithM_ updateAdj (tail as) as_bar
    where
      -- map (\i -> if i == 0 then true else is[i] != is[i-1]) (iota n)
      mkFlagLam :: LParam SOACS -> VName -> ADM (Lambda SOACS)
      mkFlagLam par_i sis =
        mkLambda [par_i] $
          fmap varsRes . letTupExp "flag" =<< do
            let i = paramName par_i
            eIf
              (toExp (le64 i .==. 0))
              (eBody $ pure $ eSubExp $ Constant $ onePrimValue Bool)
              ( eBody $
                  pure $ do
                    i_p <- letExp "i_p" =<< toExp (le64 i - 1)
                    vs <- traverse (letExp "vs" . BasicOp . Index sis . Slice . pure . DimFix . Var) [i, i_p]
                    let [vs_i, vs_p] = vs
                    toExp $ bNot $ le64 vs_i .==. le64 vs_p
              )
  