class StreamsPartitionAssignor implements ConsumerPartitionAssignor, Configurable {

    static Map<String, List<TaskId>> assignTasksToThreads(final Collection<TaskId> statefulTasksToAssign,
                                                          final Collection<TaskId> statelessTasksToAssign,
                                                          final SortedSet<String> consumers,
                                                          final ClientState state) {
        final Map<String, List<TaskId>> assignment = new HashMap<>();
        for (final String consumer : consumers) {
            assignment.put(consumer, new ArrayList<>());
        }

        final List<TaskId> unassignedStatelessTasks = new ArrayList<>(statelessTasksToAssign);
        Collections.sort(unassignedStatelessTasks);

        final Iterator<TaskId> unassignedStatelessTasksIter = unassignedStatelessTasks.iterator();

        final int minStatefulTasksPerThread = (int) Math.floor(((double) statefulTasksToAssign.size()) / consumers.size());
        final PriorityQueue<TaskId> unassignedStatefulTasks = new PriorityQueue<>(statefulTasksToAssign);

        final Queue<String> consumersToFill = new LinkedList<>();
        // keep track of tasks that we have to skip during the first pass in case we can reassign them later
        // using tree-map to make sure the iteration ordering over keys are preserved
        final Map<TaskId, String> unassignedTaskToPreviousOwner = new TreeMap<>();

        if (!unassignedStatefulTasks.isEmpty()) {
            // First assign stateful tasks to previous owner, up to the min expected tasks/thread
            for (final String consumer : consumers) {
                final List<TaskId> threadAssignment = assignment.get(consumer);

                for (final TaskId task : getPreviousTasksByLag(state, consumer)) {
                    if (unassignedStatefulTasks.contains(task)) {
                        if (threadAssignment.size() < minStatefulTasksPerThread) {
                            threadAssignment.add(task);
                            unassignedStatefulTasks.remove(task);
                        } else {
                            unassignedTaskToPreviousOwner.put(task, consumer);
                        }
                    }
                }

                if (threadAssignment.size() < minStatefulTasksPerThread) {
                    consumersToFill.offer(consumer);
                }
            }

            // Next interleave remaining unassigned tasks amongst unfilled consumers
            while (!consumersToFill.isEmpty()) {
                final TaskId task = unassignedStatefulTasks.poll();
                if (task != null) {
                    final String consumer = consumersToFill.poll();
                    final List<TaskId> threadAssignment = assignment.get(consumer);
                    threadAssignment.add(task);
                    if (threadAssignment.size() < minStatefulTasksPerThread) {
                        consumersToFill.offer(consumer);
                    }
                } else {
                    throw new IllegalStateException("Ran out of unassigned stateful tasks but some members were not at capacity");
                }
            }

            // At this point all consumers are at the min capacity, so there may be up to N - 1 unassigned
            // stateful tasks still remaining that should now be distributed over the consumers
            if (!unassignedStatefulTasks.isEmpty()) {
                consumersToFill.addAll(consumers);

                // Go over the tasks we skipped earlier and assign them to their previous owner when possible
                for (final Map.Entry<TaskId, String> taskEntry : unassignedTaskToPreviousOwner.entrySet()) {
                    final TaskId task = taskEntry.getKey();
                    final String consumer = taskEntry.getValue();
                    if (consumersToFill.contains(consumer) && unassignedStatefulTasks.contains(task)) {
                        assignment.get(consumer).add(task);
                        unassignedStatefulTasks.remove(task);
                        // Remove this consumer since we know it is now at minCapacity + 1
                        consumersToFill.remove(consumer);
                    }
                }

                // Now just distribute the remaining unassigned stateful tasks over the consumers still at min capacity
                for (final TaskId task : unassignedStatefulTasks) {
                    final String consumer = consumersToFill.poll();
                    final List<TaskId> threadAssignment = assignment.get(consumer);
                    threadAssignment.add(task);
                }


                // There must be at least one consumer still at min capacity while all the others are at min
                // capacity + 1, so start distributing stateless tasks to get all consumers back to the same count
                while (unassignedStatelessTasksIter.hasNext()) {
                    final String consumer = consumersToFill.poll();
                    if (consumer != null) {
                        final TaskId task = unassignedStatelessTasksIter.next();
                        unassignedStatelessTasksIter.remove();
                        assignment.get(consumer).add(task);
                    } else {
                        break;
                    }
                }
            }
        }

        // Now just distribute tasks while circling through all the consumers
        consumersToFill.addAll(consumers);

        while (unassignedStatelessTasksIter.hasNext()) {
            final TaskId task = unassignedStatelessTasksIter.next();
            final String consumer = consumersToFill.poll();
            assignment.get(consumer).add(task);
            consumersToFill.offer(consumer);
        }

        return assignment;
    }

}