/*
 * This code demonstrates a solution to the problem described at
 * https://stackoverflow.com/questions/61393463/
 * Copyright (c) 2020, John McClane. All rights reserved.
 *   
 */

import java.util.*;
import java.util.function.Supplier;
import java.util.function.BiConsumer;
import java.math.BigInteger;

abstract class PartitionGenerator implements Supplier<int[]>{
	public static final Random rand = new Random();
	protected final int numberCount;
	protected final int min;
	protected final int range;
	protected final int sum; // shifted sum
	protected final boolean sorted;
	
	protected PartitionGenerator(int numberCount, int min, int max, int sum, boolean sorted) {
		if (numberCount <= 0)
			throw new IllegalArgumentException("Number count should be positive");
		this.numberCount = numberCount;
		this.min = min;
		range = max - min;
		if (range < 0)
			throw new IllegalArgumentException("min > max");
		sum -= numberCount * min;
		if (sum < 0)
			throw new IllegalArgumentException("Sum is too small");
		if (numberCount * range < sum)
			throw new IllegalArgumentException("Sum is too large");
		this.sum = sum;
		this.sorted = sorted;
	}
	
	// Whether this generator returns sorted arrays (i.e. combinations)
	public final boolean isSorted() {
		return sorted;
	}
	
	public interface GeneratorFactory {
		PartitionGenerator create(int numberCount, int min, int max, int sum);
	}
}

// Permutations with repetition (i.e. unsorted vectors) with given sum
class PermutationPartitionGenerator extends PartitionGenerator {
	private final double[][] distributionTable;

	public PermutationPartitionGenerator(int numberCount, int min, int max, int sum) {
		super(numberCount, min, max, sum, false);
		distributionTable = calculateSolutionCountTable();
	}

	private double[][] calculateSolutionCountTable() {
		double[][] table = new double[numberCount + 1][sum + 1];
		BigInteger[] a = new BigInteger[sum + 1];
		BigInteger[] b = new BigInteger[sum + 1];
		for (int i = 1; i <= sum; i++)
			a[i] = BigInteger.ZERO;
		a[0] = BigInteger.ONE;
		table[0][0] = 1.0;
		for (int n = 1; n <= numberCount; n++) {
			double[] t = table[n];
			for (int s = 0; s <= sum; s++) {
				BigInteger z = BigInteger.ZERO;
				for (int i = Math.max(0, s - range); i <= s; i++)
					z = z.add(a[i]);
				b[s] = z;
				t[s] = z.doubleValue();
			}
			// swap a and b
			BigInteger[] c = b;
			b = a;
			a = c;
		}
		return table;
	}

	@Override
	public int[] get() {
		int[] p = new int[numberCount];
		int s = sum; // current sum
		for (int i = numberCount - 1; i >= 0; i--) {
			double t = rand.nextDouble() * distributionTable[i + 1][s];
			double[] tableRow = distributionTable[i];
			int oldSum = s;
			// lowerBound is introduced only for safety, it shouldn't be crossed 
			int lowerBound = s - range;
			if (lowerBound < 0)
				lowerBound = 0;
			s++;
			do
				t -= tableRow[--s];
			// s can be equal to lowerBound here with t > 0 only due to imprecise subtraction
			while (t > 0 && s > lowerBound);
			p[i] = min + (oldSum - s);
		}
		assert s == 0;
		return p;
	}

	public static final PartitionGenerator.GeneratorFactory factory = (numberCount, min, max,sum) ->
		new PermutationPartitionGenerator(numberCount, min, max, sum);
}

// Combinations with repetition (i.e. sorted vectors) with given sum 
class CombinationPartitionGenerator extends PartitionGenerator {
	private final double[][][] distributionTable;

	public CombinationPartitionGenerator(int numberCount, int min, int max, int sum) {
		super(numberCount, min, max, sum, true);
		distributionTable = calculateSolutionCountTable();
	}

	private double[][][] calculateSolutionCountTable() {
		double[][][] table = new double[numberCount + 1][range + 1][sum + 1];
		BigInteger[][] a = new BigInteger[range + 1][sum + 1];
		BigInteger[][] b = new BigInteger[range + 1][sum + 1];
		double[][] t = table[0];
		for (int m = 0; m <= range; m++) {
			a[m][0] = BigInteger.ONE;
			t[m][0] = 1.0;
			for (int s = 1; s <= sum; s++) {
				a[m][s] = BigInteger.ZERO;
				t[m][s] = 0.0;
			}
		}
		for (int n = 1; n <= numberCount; n++) {
			t = table[n];
			for (int m = 0; m <= range; m++)
				for (int s = 0; s <= sum; s++) {
					BigInteger z;
					if (m == 0)
						z = a[0][s];
					else {
						z = b[m - 1][s];
						if (m <= s)
							z = z.add(a[m][s - m]);
					}
					b[m][s] = z;
					t[m][s] = z.doubleValue();
				}
			// swap a and b
			BigInteger[][] c = b;
			b = a;
			a = c;
		}
		return table;
	}

	@Override
	public int[] get() {
		int[] p = new int[numberCount];
		int m = range; // current max
		int s = sum; // current sum
		for (int i = numberCount - 1; i >= 0; i--) {
			double t = rand.nextDouble() * distributionTable[i + 1][m][s];
			double[][] tableCut = distributionTable[i];
			if (s < m)
				m = s;
			s -= m;
			while (true) {
				t -= tableCut[m][s];
				// m can be 0 here with t > 0 only due to imprecise subtraction
				if (t <= 0 || m == 0)
					break;
				m--;
				s++;
			}
			p[i] = min + m;
		}
		assert s == 0;
		return p;
	}

	public static final PartitionGenerator.GeneratorFactory factory = (numberCount, min, max, sum) ->
		new CombinationPartitionGenerator(numberCount, min, max, sum);
}

class SmithTromblePartitionGenerator extends PartitionGenerator {
	public SmithTromblePartitionGenerator(int numberCount, int min, int max, int sum) {
		super(numberCount, min, max, sum, false);
	}

	@Override
	public int[] get() {
		List<Integer> ls = new ArrayList<>(numberCount + 1);
		int[] ret = new int[numberCount];
		int increasedSum = sum + numberCount;
		while (true) {
			ls.add(0);
			while (ls.size() < numberCount) {
				int c = 1 + rand.nextInt(increasedSum - 1);
				if (!ls.contains(c))
					ls.add(c);
			}
			Collections.sort(ls);
			ls.add(increasedSum);
			boolean good = true;
			for (int i = 0; i < numberCount; i++) {
				int x = ls.get(i + 1) - ls.get(i) - 1;
				if (x > range) {
					good = false;
					break;
				}
				ret[i] = x;
			}
			if (good) {
				for (int i = 0; i < numberCount; i++)
					ret[i] += min;
				return ret;
			}
			ls.clear();
		}
	}

	public static final PartitionGenerator.GeneratorFactory factory = (numberCount, min, max, sum) ->
		new SmithTromblePartitionGenerator(numberCount, min, max, sum);
}

// Enumerates all partitions with given parameters
class SequentialEnumerator extends PartitionGenerator {
	private final int max;
	private final int[] p;
	private boolean finished;

	public SequentialEnumerator(int numberCount, int min, int max, int sum, boolean sorted) {
		super(numberCount, min, max, sum, sorted);
		this.max = max;
		p = new int[numberCount];
		startOver();
	}

	private void startOver() {
		finished = false;
		int unshiftedSum = sum + numberCount * min;
		fillMinimal(0, Math.max(min, unshiftedSum - (numberCount - 1) * max), unshiftedSum);
	}

	private void fillMinimal(int beginIndex, int minValue, int fillSum) {
		int fillRange = max - minValue;
		if (fillRange == 0)
			Arrays.fill(p, beginIndex, numberCount, max);
		else {
			int fillCount = numberCount - beginIndex;
			fillSum -= fillCount * minValue;
			int maxCount = fillSum / fillRange;
			int maxStartIndex = numberCount - maxCount;
			Arrays.fill(p, maxStartIndex, numberCount, max);
			fillSum -= maxCount * fillRange;
			Arrays.fill(p, beginIndex, maxStartIndex, minValue);
			if (fillSum != 0)
				p[maxStartIndex - 1] = minValue + fillSum;
		}
	}

	@Override
	public int[] get() { // returns null when there is no more partition, then starts over
		if (finished) {
			startOver();
			return null;
		}
		int[] pCopy = p.clone();
		if (numberCount > 1) {
			int i = numberCount;
			int s = p[--i];
			while (i > 0) {
				int x = p[--i];
				if (x == max) {
					s += x;
					continue;
				}
				x++;
				s--;
				int minRest = sorted ? x : min;
				if (s < minRest * (numberCount - i - 1)) {
					s += x;
					continue;
				}
				p[i++]++;
				fillMinimal(i, minRest, s);
				return pCopy;
			}
		}
		finished = true;
		return pCopy;
	}

	public static final PartitionGenerator.GeneratorFactory permutationFactory = (numberCount, min, max, sum) ->
		new SequentialEnumerator(numberCount, min, max, sum, false);
	public static final PartitionGenerator.GeneratorFactory combinationFactory = (numberCount, min, max, sum) ->
		new SequentialEnumerator(numberCount, min, max, sum, true);
}

class Test {
	private final int numberCount;
	private final int min;
	private final int max;
	private final int sum;
	private final int repeatCount;
	private final BiConsumer<PartitionGenerator, Test> procedure;

	public Test(int numberCount, int min, int max, int sum, int repeatCount,
			BiConsumer<PartitionGenerator, Test> procedure) {
		this.numberCount = numberCount;
		this.min = min;
		this.max = max;
		this.sum = sum;
		this.repeatCount = repeatCount;
		this.procedure = procedure;
	}

	@Override
	public String toString() {
		return String.format("=== %d numbers from [%d, %d] with sum %d, %d iterations ===",
				numberCount, min, max, sum,	repeatCount);
	}

	private static class GeneratedVector {
		final int[] v;

		GeneratedVector(int[] vect) {
			v = vect;
		}

		@Override
		public int hashCode() {
			return Arrays.hashCode(v);
		}

		@Override
		public boolean equals(Object obj) {
			if (this == obj)
				return true;
			return Arrays.equals(v, ((GeneratedVector)obj).v);
		}

		@Override
		public String toString() {
			return Arrays.toString(v);
		}
	}

	private static final Comparator<Map.Entry<GeneratedVector, Integer>> lexicographical = (e1, e2) -> {
		int[] v1 = e1.getKey().v;
		int[] v2 = e2.getKey().v;
		int len = v1.length;
		int d = len - v2.length;
		if (d != 0)
			return d;
		for (int i = 0; i < len; i++) {
			d = v1[i] - v2[i];
			if (d != 0)
				return d;
		}
		return 0;
	};
	
	private static final Comparator<Map.Entry<GeneratedVector, Integer>> byCount =
			Comparator.<Map.Entry<GeneratedVector, Integer>>comparingInt(Map.Entry::getValue)
			.thenComparing(lexicographical);
	
	public static int SHOW_MISSING_LIMIT = 10;

	private static void checkMissingPartitions(Map<GeneratedVector, Integer> map, PartitionGenerator reference) {
		int missingCount = 0;
		while (true) {
			int[] v = reference.get();
			if (v == null)
				break;
			GeneratedVector gv = new GeneratedVector(v);
			if (!map.containsKey(gv)) {
				if (missingCount == 0)
					System.out.println(" Missing:");
				if (++missingCount > SHOW_MISSING_LIMIT) {
					System.out.println("  . . .");
					break;
				}
				System.out.println(gv);
			}
		}
	}

	public static final BiConsumer<PartitionGenerator, Test> distributionTest(boolean sortByCount) {
		return (PartitionGenerator gen, Test test) -> {
			System.out.print("\n" + getName(gen) + "\n\n");
			Map<GeneratedVector, Integer> combos = new HashMap<>();
			// There's no point of checking permus for sorted generators
			// because they are the same as combos for them
			Map<GeneratedVector, Integer> permus = gen.isSorted() ? null : new HashMap<>();
			for (int i = 0; i < test.repeatCount; i++) {
				int[] v = gen.get();
				if (v == null && gen instanceof SequentialEnumerator)
					break;
				if (permus != null) {
					permus.merge(new GeneratedVector(v), 1, Integer::sum);
					v = v.clone();
					Arrays.sort(v);
				}
				combos.merge(new GeneratedVector(v), 1, Integer::sum);
			}
			Set<Map.Entry<GeneratedVector, Integer>> sortedEntries = new TreeSet<>(
					sortByCount ? byCount : lexicographical);
			System.out.println("Combos" + (gen.isSorted() ? ":" : " (don't have to be uniform):"));
			sortedEntries.addAll(combos.entrySet());
			for (Map.Entry<GeneratedVector, Integer> e : sortedEntries)
				System.out.println(e);
			checkMissingPartitions(combos, test.getGenerator(SequentialEnumerator.combinationFactory));
			if (permus != null) {
				System.out.println("\nPermus:");
				sortedEntries.clear();
				sortedEntries.addAll(permus.entrySet());
				for (Map.Entry<GeneratedVector, Integer> e : sortedEntries)
					System.out.println(e);
				checkMissingPartitions(permus, test.getGenerator(SequentialEnumerator.permutationFactory));
			}
		};
	}

	public static final BiConsumer<PartitionGenerator, Test> correctnessTest =
		(PartitionGenerator gen, Test test) -> {
		String genName = getName(gen);
		for (int i = 0; i < test.repeatCount; i++) {
			int[] v = gen.get();
			if (v == null && gen instanceof SequentialEnumerator)
				v = gen.get();
			if (v.length != test.numberCount)
				throw new RuntimeException(genName + ": array of wrong length");
			int s = 0;
			if (gen.isSorted()) {
				if (v[0] < test.min || v[v.length - 1] > test.max)
					throw new RuntimeException(genName + ": generated number is out of range");
				int prev = test.min;
				for (int x : v) {
					if (x < prev)
						throw new RuntimeException(genName + ": unsorted array");
					s += x;
					prev = x;
				}
			} else
				for (int x : v) {
					if (x < test.min || x > test.max)
						throw new RuntimeException(genName + ": generated number is out of range");
					s += x;
				}
			if (s != test.sum)
				throw new RuntimeException(genName + ": wrong sum");
		}
		System.out.format("%30s :   correctness test passed%n", genName);
	};
	
	public static final BiConsumer<PartitionGenerator, Test> performanceTest =
		(PartitionGenerator gen, Test test) -> {
		long time = System.nanoTime();
		for (int i = 0; i < test.repeatCount; i++)
			gen.get();
		time = System.nanoTime() - time;
		System.out.format("%30s : %8.3f s %10.0f ns/test%n", getName(gen), time * 1e-9, time * 1.0 / test.repeatCount);
	};

	public PartitionGenerator getGenerator(PartitionGenerator.GeneratorFactory factory) {
		return factory.create(numberCount, min, max, sum);
	}

	public static String getName(PartitionGenerator gen) {
		String name = gen.getClass().getSimpleName();
		if (gen instanceof SequentialEnumerator)
			return (gen.isSorted() ? "Sorted " : "Unsorted ") + name;
		else
			return name;
	}

	public static PartitionGenerator.GeneratorFactory[] factories = { SmithTromblePartitionGenerator.factory,
			PermutationPartitionGenerator.factory, CombinationPartitionGenerator.factory,
			SequentialEnumerator.permutationFactory, SequentialEnumerator.combinationFactory };

	public static void main(String[] args) {
		Test[] tests = {
							new Test(3, 0, 3, 5, 3_000, distributionTest(false)),
							new Test(3, 0, 6, 12, 3_000, distributionTest(true)),
							new Test(50, -10, 20, 70, 1_000, correctnessTest),
							new Test(7, 3, 10, 42, 25_000, performanceTest),
							new Test(20, 3, 10, 120, 1_500, performanceTest)
					   };
		for (Test t : tests) {
			System.out.println(t);
			for (PartitionGenerator.GeneratorFactory factory : factories) {
				PartitionGenerator candidate = t.getGenerator(factory);
				t.procedure.accept(candidate, t);
			}
			System.out.println();
		}
	}
}
