diff --git a/src/encode_coding_tree.c b/src/encode_coding_tree.c
index 925d80da..5277768c 100644
--- a/src/encode_coding_tree.c
+++ b/src/encode_coding_tree.c
@@ -1242,14 +1242,14 @@ bool uvg_write_split_flag(
   cabac_data_t* cabac,
   const cu_info_t * left_cu,
   const cu_info_t * above_cu,
-  uint8_t split_flag,
   const cu_loc_t* const cu_loc,
+  const uint32_t split_tree,
   int depth,
   enum uvg_tree_type tree_type,
   double* bits_out)
 {
-  uint16_t abs_x = (cu_loc->x + state->tile->offset_x) >> (tree_type == UVG_CHROMA_T);
-  uint16_t abs_y = (cu_loc->y + state->tile->offset_y) >> (tree_type == UVG_CHROMA_T);
+  uint16_t abs_x = cu_loc->x + (state->tile->offset_x >> (tree_type == UVG_CHROMA_T));
+  uint16_t abs_y = cu_loc->y + (state->tile->offset_y >> (tree_type == UVG_CHROMA_T));
   double bits = 0;
   const encoder_control_t* const ctrl = state->encoder_control;
   // Implisit split flag when on border
@@ -1271,23 +1271,23 @@ bool uvg_write_split_flag(
   bool top_right_available = ((abs_x + cu_width - 1) < (ctrl->in.width >> (tree_type == UVG_CHROMA_T)));
 
   if (!bottom_left_available && !top_right_available && allow_qt) {
-    implicit_split_mode = UVG_QUAD_SPLIT;
+    implicit_split_mode = QT_SPLIT;
   }
   else if (!bottom_left_available && allow_btt) {
-    implicit_split_mode = UVG_HORZ_SPLIT;
+    implicit_split_mode = BT_HOR_SPLIT;
   }
   else if (!top_right_available && allow_btt) {
-    implicit_split_mode = UVG_VERT_SPLIT;
+    implicit_split_mode = BT_VER_SPLIT;
   }
   else if (!bottom_left_available || !top_right_available) {
-    implicit_split_mode = UVG_QUAD_SPLIT;
+    implicit_split_mode = QT_SPLIT;
   }
   
   // Check split conditions
   if (implicit_split_mode != UVG_NO_SPLIT) {
     no_split = th_split = tv_split = false;
-    bh_split = (implicit_split_mode == UVG_HORZ_SPLIT);
-    bv_split = (implicit_split_mode == UVG_VERT_SPLIT);
+    bh_split = (implicit_split_mode == BT_HOR_SPLIT);
+    bv_split = (implicit_split_mode == BT_VER_SPLIT);
   }
 
   if (!allow_btt) {
@@ -1296,17 +1296,18 @@ bool uvg_write_split_flag(
 
   bool allow_split = allow_qt | bh_split | bv_split | th_split | tv_split;
 
-  split_flag |= implicit_split_mode != UVG_NO_SPLIT;
+  int split_flag = (split_tree >> (depth * 3)) & 7;
+
+  split_flag = implicit_split_mode != UVG_NO_SPLIT ? implicit_split_mode : split_flag;
 
   int split_model = 0;
   if (no_split && allow_split) {
     // Get left and top block split_flags and if they are present and true, increase model number
-    // ToDo: should use height and width to increase model, PU_GET_W() ?
-    if (left_cu && left_cu->depth > depth) {
+    if (left_cu && (1 << left_cu->log2_height) < cu_height) {
       split_model++;
     }
 
-    if (above_cu && above_cu->depth > depth) {
+    if (above_cu && (1 << above_cu->log2_width) < cu_width) {
       split_model++;
     }
 
@@ -1322,15 +1323,11 @@ bool uvg_write_split_flag(
     split_model += 3 * (split_num >> 1);
 
     cabac->cur_ctx = &(cabac->ctx.split_flag_model[split_model]);
-    if(cabac->only_count && !split_flag) {
 
-      //printf("%hu %hu %d %d %d\n", state->search_cabac.ctx.split_flag_model[split_model].state[0], state->search_cabac.ctx.split_flag_model[split_model].state[1],
-      //  split_model, x, y);
-    }
-    CABAC_FBITS_UPDATE(cabac, &(cabac->ctx.split_flag_model[split_model]), split_flag, bits, "split_flag");
+    CABAC_FBITS_UPDATE(cabac, &(cabac->ctx.split_flag_model[split_model]), split_flag != 0, bits, "split_flag");
   }
 
-  bool qt_split = split_flag || implicit_split_mode == UVG_QUAD_SPLIT;
+  bool qt_split = split_flag == UVG_QUAD_SPLIT;
 
   if (!(implicit_split_mode == UVG_NO_SPLIT) && (allow_qt && allow_btt)) {
     split_model = (left_cu && GET_SPLITDATA(left_cu, depth)) + (above_cu && GET_SPLITDATA(above_cu, depth)) + (depth < 2 ? 0 : 3);
@@ -1342,17 +1339,17 @@ bool uvg_write_split_flag(
 
     split_model = 0;
 
-    // Get left and top block split_flags and if they are present and true, increase model number
-    if (left_cu && GET_SPLITDATA(left_cu, depth) == 1) {
+    // TODO: These are incorrect
+    if (left_cu && (1 << left_cu->log2_height) > cu_height) {
       split_model++;
     }
 
-    if (above_cu && GET_SPLITDATA(above_cu, depth) == 1) {
+    if (above_cu && (1 << above_cu->log2_width) > cu_width) {
       split_model++;
     }
 
     split_model += (depth > 2 ? 0 : 3);
-    CABAC_FBITS_UPDATE(cabac, &(cabac->ctx.qt_split_flag_model[split_model]), split_flag, bits, "split_cu_mode");
+    CABAC_FBITS_UPDATE(cabac, &(cabac->ctx.qt_split_flag_model[split_model]), qt_split, bits, "split_cu_mode");
   }
   if (bits_out) *bits_out += bits;
   return split_flag;
@@ -1414,9 +1411,9 @@ void uvg_encode_coding_tree(
       state,
       cabac,
       left_cu,
-      above_cu,
-      (cur_cu->split_tree >> (split_tree.current_depth * 3)) & 7,
+      above_cu, 
       cu_loc,
+      cur_cu->split_tree,
       depth,
       tree_type,
       NULL);
@@ -1725,7 +1722,7 @@ double uvg_mock_encode_coding_unit(
       left_cu = LCU_GET_CU_AT_PX(lcu, x_local - 1, y_local);
     }
     else {
-      left_cu = uvg_cu_array_at_const(state->tile->frame->chroma_cu_array, (x >> 1) - 1, y >> 1);
+      left_cu = uvg_cu_array_at_const(state->tile->frame->chroma_cu_array, x  - 1, y);
     }
   }
   if (y) {
@@ -1733,7 +1730,7 @@ double uvg_mock_encode_coding_unit(
       above_cu = LCU_GET_CU_AT_PX(lcu, x_local, y_local-1);
     }
     else {
-      above_cu = uvg_cu_array_at_const(state->tile->frame->chroma_cu_array, x >> 1, (y >> 1) - 1);
+      above_cu = uvg_cu_array_at_const(state->tile->frame->chroma_cu_array, x, y - 1);
     }
   }
   
@@ -1748,8 +1745,8 @@ double uvg_mock_encode_coding_unit(
       cabac,
       left_cu,
       above_cu,
-      0,
       cu_loc,
+      cur_cu->split_tree,
       depth,
       tree_type,
       &bits);
diff --git a/src/encode_coding_tree.h b/src/encode_coding_tree.h
index a7fe896b..86605e4d 100644
--- a/src/encode_coding_tree.h
+++ b/src/encode_coding_tree.h
@@ -101,8 +101,8 @@ bool uvg_write_split_flag(
   cabac_data_t* cabac,
   const cu_info_t* left_cu,
   const cu_info_t* above_cu,
-  uint8_t split_flag,
   const cu_loc_t* const cu_loc,
+  const uint32_t split_tree,
   int depth,
   enum uvg_tree_type tree_type,
   double* bits_out);
diff --git a/src/search.c b/src/search.c
index 3fd67cc2..fb28085d 100644
--- a/src/search.c
+++ b/src/search.c
@@ -193,6 +193,9 @@ static void lcu_fill_cu_info(lcu_t *lcu, int x_local, int y_local, int width, in
       to->violates_lfnst_constrained_luma = cu->violates_lfnst_constrained_luma;
       to->violates_lfnst_constrained_chroma = cu->violates_lfnst_constrained_chroma;
 
+      to->log2_height = cu->log2_height;
+      to->log2_width = cu->log2_width;
+
       if (cu->type == CU_INTRA) {
         to->intra.mode        = cu->intra.mode;
         to->intra.mode_chroma = cu->intra.mode_chroma;
@@ -1256,6 +1259,14 @@ static double search_cu(
   }
 
   if (cur_cu->type == CU_INTRA || cur_cu->type == CU_INTER || cur_cu->type == CU_IBC) {
+  // The cabac functions assume chroma locations whereas the search uses luma locations
+  // for the chroma tree, therefore we need to shift the chroma coordinates here for
+  // passing to the bit cost calculating functions.
+  cu_loc_t chroma_loc = *cu_loc;
+  chroma_loc.y >>= 1;
+  chroma_loc.x >>= 1;
+
+  if (cur_cu->type == CU_INTRA || cur_cu->type == CU_INTER) {
     double bits = 0;
     cabac_data_t* cabac  = &state->search_cabac;
     cabac->update = 1;
@@ -1263,7 +1274,9 @@ static double search_cu(
     bits += uvg_mock_encode_coding_unit(
       state,
       cabac,
-      cu_loc, lcu, cur_cu,
+      tree_type != UVG_CHROMA_T ? cu_loc : &chroma_loc, 
+      lcu,
+      cur_cu,
       tree_type);
 
     
@@ -1310,6 +1323,8 @@ static double search_cu(
 
   // Recursively split all the way to max search depth.
   if (can_split_cu) {
+    const split_tree_t new_split = { split_tree.split_tree | QT_SPLIT << (split_tree.current_depth * 3), split_tree.current_depth + 1 };
+
     int half_cu = cu_width >> (tree_type != UVG_CHROMA_T);
     double split_cost = 0.0;
     int cbf = cbf_is_set_any(cur_cu->cbf, split_tree.current_depth);
@@ -1345,9 +1360,9 @@ static double search_cu(
         state,
         &state->search_cabac,
         left_cu,
-        above_cu,
-        1,
-        cu_loc,
+        above_cu, 
+        tree_type != UVG_CHROMA_T ? cu_loc : &chroma_loc,
+        new_split.split_tree,
         depth,
         tree_type,
         &split_bits);
@@ -1362,7 +1377,6 @@ static double search_cu(
     // It is ok to interrupt the search as soon as it is known that
     // the split costs at least as much as not splitting.
     if (cur_cu->type == CU_NOTSET || cbf || state->encoder_control->cfg.cu_split_termination == UVG_CU_SPLIT_TERMINATION_OFF) {
-      const split_tree_t new_split = { split_tree.split_tree | QT_SPLIT << (split_tree.current_depth * 3), split_tree.current_depth + 1};
       cu_loc_t new_cu_loc;
       if (split_cost < cost) {
         uvg_cu_loc_ctor(&new_cu_loc, x, y, half_cu, half_cu);
@@ -1407,8 +1421,7 @@ static double search_cu(
         double bits = 0;
         uvg_write_split_flag(state, &state->search_cabac,
                              x > 0 ? LCU_GET_CU_AT_PX(lcu, SUB_SCU(x) - 1, SUB_SCU(y)) : NULL,
-                             y > 0 ? LCU_GET_CU_AT_PX(lcu, SUB_SCU(x), SUB_SCU(y) - 1) : NULL,
-                             0, cu_loc, depth, tree_type, &bits);
+                             y > 0 ? LCU_GET_CU_AT_PX(lcu, SUB_SCU(x), SUB_SCU(y) - 1) : NULL, cu_loc, cur_cu->split_tree, depth, tree_type, &bits);
 
         cur_cu->intra = cu_d1->intra;
         cur_cu->type = CU_INTRA;