Неверный indexPath при работе heightForRow

ios
tableview

#1

Сломал уже всю голову, но не могу понять, почему возвращается неверный индекс ячейки при работе с heightForRow.

//Записываем номер ячейки

func tableView(_ tableView: UITableView, didSelectRowAt indexPath: IndexPath) {
if indexPath.row == selectedIndex {
        selectedIndex = -1
    } else {
        selectedIndex = indexPath.row
    }
    
    tableView.beginUpdates()
    tableView.endUpdates()
}

//Разворачиваем ячейку

func tableView(_ tableView: UITableView, heightForRowAt indexPath: IndexPath) -> CGFloat {
    if indexPath.row == selectedIndex {
        return 204
    } else {
        return 44
    }
}

//Здесь ищу ячейку по indexPath и все равно изображение присваивается на какую попало после того как любая из ячеек развернулась.

func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
    let cell = tableView.dequeueReusableCell(withIdentifier: "Cell", for: indexPath) as! SupplementTableViewCell
    
    cell.labelFullName.text = shift[indexPath.row].fullName
    
    if let cell = tableView.cellForRow(at: indexPath) as? SupplementTableViewCell {
        localShift.filter{ $0.fullName == shift[indexPath.row].fullName }.forEach { _ in
            cell.statusImageView.image = #imageLiteral(resourceName: "on")
        }
    }
    
    return cell
}

Скорее всего, после того как мы развернули ячейку, индекс считается неправильно, но почему он считается неправильно даже после поиска через cellForRow и что с этим сделать!?