import argparse
import random


def parse_args():
    parser = argparse.ArgumentParser(
        description=(
            "Generate a random knapsack dataset with a prescribed number of optimal "
            "solutions within a required size range."
        )
    )
    parser.add_argument("output", help="Output text file with one 'size priority' pair per line.")
    parser.add_argument("--items", type=int, required=True, help="Number of items to generate.")
    parser.add_argument(
        "--min-size",
        type=int,
        required=True,
        help="Minimal allowed sum of selected item sizes.",
    )
    parser.add_argument(
        "--max-size",
        type=int,
        required=True,
        help="Maximal allowed sum of selected item sizes.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Seed for reproducible randomization.",
    )
    return parser.parse_args()


def validate_inputs(item_count, min_size, max_size, desired_solutions):
    if item_count < 10:
        raise ValueError("We promise to generate at least 10 items.")
    if item_count > 63:
        raise ValueError("The knapsack solver stores selections in 64 bits, so --items must be at most 63.")
    if min_size < 1000:
        raise ValueError("--min-size must be at least 1000.")
    if max_size < min_size:
        raise ValueError("--max-size must be greater than or equal to --min-size.")
    if max_size - min_size >= min_size // 5:
        raise ValueError("Delta between max and min size must be less than 1/5 of min-size.")
    if max_size - min_size // item_count <= 3:
        raise ValueError("Delta between max and min size divided by item count must be greater than 3.")


def bucket_sizes(item_count, rng):
    '''
    We divide the items into 5 buckets:
    - items of size min_size // 2 + epsilon
    - items of size min_size // 3 + epsilon
    - items of size min_size // 4 + epsilon
    - small items to be selected to pad the sack
    - items of size in range [delta, min // 4] - garbage that won't be selected
    '''
    buckets = [2, 3, 4, 0, min(item_count // 5, item_count - 9)]
    # sum the weights of the first 4 buckets, and put the rest in the last bucket
    allocated = sum(buckets)
    while allocated < item_count:
        rng_index = rng.randint(0, 10)
        # the small items bucket [3] is much more likely to be selected (indices > 3 are cropped to 3)
        buckets[min(rng_index, 3)] += 1
        allocated += 1

    return buckets


def random_items(items, count, min_size, max_size, min_priority, max_priority, rng):
    for _ in range(count):
        size = rng.randint(min_size, max_size)
        priority = rng.randint(min_priority, max_priority)
        items.append((size, priority))


def generate_items(item_count, min_size, max_size, rng):
    delta = max_size - min_size
    epsilon = delta // item_count - 1
    items = []

    p = rng.randint(50, 100)

    sizes = bucket_sizes(item_count, rng)
    # the min/2 bucket (first, we generate minimal amount of items with high priority for the solution)
    random_items(items, 2, min_size // 2 + 1, min_size // 2 + epsilon, p * 6, p * 6, rng)
    random_items(items, sizes[0] - 2, min_size // 2 + 1, min_size // 2 + epsilon, 1, p * 6, rng)

    # the min/3 bucket
    random_items(items, 3, min_size // 3 + 1, min_size // 3 + epsilon, p * 4, p * 4, rng)
    random_items(items, sizes[1] - 3, min_size // 3 + 1, min_size // 3 + epsilon, 1, p * 4, rng)

    # the min/4 bucket
    random_items(items, 4, min_size // 4 + 1, min_size // 4 + epsilon, p * 3, p * 3, rng)
    random_items(items, sizes[2] - 4, min_size // 4 + 1, min_size // 4 + epsilon, 1, p * 3, rng)

    # the small items bucket
    lower = delta - 2
    for i in range(sizes[3]):
        random_items(items, 1, lower, delta - 1, 1, p, rng)
        lower = max(lower - 2 * epsilon, 1)

    # garbage items that won't be selected
    random_items(items, sizes[4], delta + 1, min_size // 4 - 1, 1, 10, rng)

    rng.shuffle(items)
    return items


def write_items(path, items):
    with open(path, "w", encoding="utf-8") as handle:
        for size, priority in items:
            handle.write(f"{size} {priority}\n")


def main():
    args = parse_args()
    rng = random.Random(args.seed)
    items = generate_items(args.items, args.min_size, args.max_size, rng)
    write_items(args.output, items)
    print(
        f"Generated {len(items)} items in {args.output} for size range "
        f"[{args.min_size}, {args.max_size}]."
    )


if __name__ == "__main__":
    main()
