#!/usr/bin/python3

import argparse
import math
import random
import sys
from enum import Enum


class Variant(Enum):
    NESTING_MAX = 'nesting_max'
    SLOW_MO = 'slow_mo'
    RANDOM = 'random'
    FIFO = 'fifo'
    ANTI_SORT_RIGHT = 'anti_sort_right'
    ANTI_SORT_LEFT = 'anti_sort_left'

    def __str__(self):
        return self.value


N_MIN = 1
N_MAX = 200000
DEFAULT_N = 20
LOCATIONS_MIN = 1
LOCATIONS_MAX = 2 * N_MAX
DEFAULT_LOCATIONS_MAX = 40
# We will reseed this with the provided seed or a random seed.
RNG = random.Random(0)


def bounded_int(string, val_min, val_max, name='Value'):
    value = int(string)
    if value < val_min or value > val_max:
        raise argparse.ArgumentTypeError(f'{name} must be in range [{val_min}, {val_max}]')
    return value


def bounded_n(string):
    return bounded_int(string, N_MIN, N_MAX, 'n')

def bounded_locations(string):
    return bounded_int(string, LOCATIONS_MIN, LOCATIONS_MAX, 'locations')


def _parse_args():
    parser = argparse.ArgumentParser('')
    parser.add_argument(
        'n',
        type=bounded_n,
        default=DEFAULT_N,
        help='The value of n.'
    )
    parser.add_argument(
        '--locations-min',
        type=bounded_locations,
        default=0,
        help='The minimum number of distinct locations to include.'
    )
    parser.add_argument(
        '--locations-max',
        type=bounded_locations,
        default=DEFAULT_LOCATIONS_MAX,
        help='The maximum number of distinct locations to include.',
    )
    parser.add_argument(
        '--variant',
        type=Variant,
        choices=list(Variant),
        default=Variant.RANDOM,
        help='Type of test case to generate.',
    )
    parser.add_argument(
        '--seed', type=int, default=random.randint(0, 10000),
        help='The random number to seed the random number generator with.'
    )
    parser.add_argument(
        '--test-name', type=str,
        help='The name for the test case. E.g., "025-small-cases" will produce files '
             '025-small-cases.in and 025-small-cases.desc. If no name is specified, '
             'output will be printed to stdout.'
    )
    return parser.parse_args()


def _validate_with_context(args):
    """Function for any validation that involves more than one argument."""
    if args.locations_min > args.locations_max:
        raise ValueError('locations-min must be less than or equal to locations-max.')
    if args.locations_min > 2*args.n:
        raise ValueError('locations_min must be less than or equal to 2*N')


def normalized(sequence):
    value_map = {}
    curr = 1
    for value in sequence:
        if value not in value_map:
            value_map[value] = curr
            curr += 1
    normalized = []
    for value in sequence:
        normalized.append(value_map[value])
    return normalized


def _random_location_sequence(num_passengers, locations_min, locations_max):
    location_sequence = []
    for i in range(locations_min):
        location_sequence.append(i + 1)
    while len(location_sequence) < 2 * num_passengers:
        location_sequence.append(RNG.randint(1, locations_max))
    RNG.shuffle(location_sequence)
    location_map = {}
    curr_location = 1
    for location in location_sequence:
        if location not in location_map:
            location_map[location] = curr_location
            curr_location += 1
    normalized_location_sequence = [location_map[location] for location in location_sequence]
    return normalized_location_sequence


def _generate_random(num_passengers, locations_min, locations_max):
    location_sequence = _random_location_sequence(num_passengers, locations_min, locations_max)
    passengers = []
    for i in range(num_passengers):
        passengers.append(i+1)
        passengers.append(i+1)
    RNG.shuffle(passengers)
    passenger_map = {}
    curr_value = 1
    for passenger in passengers:
        if passenger not in passenger_map:
            passenger_map[passenger] = curr_value
            curr_value += 1
    lines = []
    lines.append(f"{num_passengers}")
    for passenger, location in zip(passengers, location_sequence):
        lines.append(f"{passenger_map[passenger]} {location}")
    return lines


def _generate_nesting_max(num_passengers, locations_min, locations_max):
    location_sequence = _random_location_sequence(num_passengers, locations_min, locations_max)
    passengers = []
    for i in range(num_passengers):
        passengers.append(i+1)
    for i in range(num_passengers):
        passengers.append(num_passengers - i)
    lines = []
    lines.append(f"{num_passengers}")
    for passenger, location in zip(passengers, location_sequence):
        lines.append(f"{passenger} {location}")
    return lines


def _generate_slow_mo(num_passengers, locations_min, locations_max):
    """Passenger pickup/dropoff indexes are adversarial to Mo's Algorithm."""
    location_sequence = _random_location_sequence(num_passengers, locations_min, locations_max)
    passengers = [-1 for _ in range(2*num_passengers)]
    bucket_size = int(math.sqrt(2 * num_passengers))
    bucket_index = 0
    passenger = 1
    while passenger <= num_passengers:
        for j in range(bucket_size):
            index = bucket_index * bucket_size
            if j % 2 == 0:
                index = index + bucket_size - (j // 2) - 1
            else:
                index = index + (j // 2)
            passengers[index] = passenger
            passenger += 1
            if passenger > num_passengers:
                break
        bucket_index += 1
    passenger = 1
    index = 0
    while passenger <= num_passengers:
        if passengers[index] == -1:
            passengers[index] = passenger
            passenger += 1
        index += 1
    passenger_map = {}
    curr_passenger = 1
    for passenger in passengers:
        if passenger not in passenger_map:
            passenger_map[passenger] = curr_passenger
            curr_passenger += 1
    lines = []
    lines.append(f"{num_passengers}")
    for passenger, location in zip(passengers, location_sequence):
        lines.append(f"{passenger_map[passenger]} {location}")
    return lines


def _generate_fifo(num_passengers, locations_min, locations_max):
    """
    Generate a random location sequence for passengers served
    First In First Out after pickup up all passengers.
    """
    location_sequence = _random_location_sequence(num_passengers, locations_min, locations_max)
    passengers = []
    for i in range(num_passengers):
        passengers.append(i+1)
    for i in range(num_passengers):
        passengers.append(i+1)
    lines = []
    lines.append(f"{num_passengers}")
    for passenger, location in zip(passengers, location_sequence):
        lines.append(f"{passenger} {location}")
    return lines


def _generate_anti_sort_right(num_passengers, locations_min, locations_max):
    """
    Generate a random location sequence for passengers served in the following pattern:
    1 3 5 7 9 2 4 6 8 10 1 2 3 4 5 6 7 8 9 10
    """
    location_sequence = _random_location_sequence(num_passengers, locations_min, locations_max)
    passengers = []
    for i in range(0, num_passengers, 2):
        passengers.append(i + 1)
    for i in range(1, num_passengers, 2):
        passengers.append(i + 1)
    for i in range(num_passengers):
        passengers.append(i + 1)
    passengers = normalized(passengers)
    lines = []
    lines.append(f"{num_passengers}")
    for passenger, location in zip(passengers, location_sequence):
        lines.append(f"{passenger} {location}")
    return lines


def _generate_anti_sort_left(num_passengers, locations_min, locations_max):
    """
    Generate a random location sequence for passengers served in the following pattern:
    1 2 3 4 5 6 7 8 9 10 1 3 5 7 9 2 4 6 8 10
    """
    location_sequence = _random_location_sequence(num_passengers, locations_min, locations_max)
    passengers = []
    for i in range(num_passengers):
        passengers.append(i + 1)
    for i in range(0, num_passengers, 2):
        passengers.append(i + 1)
    for i in range(1, num_passengers, 2):
        passengers.append(i + 1)
    lines = []
    lines.append(f"{num_passengers}")
    for passenger, location in zip(passengers, location_sequence):
        lines.append(f"{passenger} {location}")
    return lines


def generate_output_lines(args):
    """Function to generate a list of output lines."""
    match args.variant:
        case Variant.RANDOM:
            output_lines = _generate_random(args.n, args.locations_min, args.locations_max)
        case Variant.NESTING_MAX:
            output_lines = _generate_nesting_max(args.n, args.locations_min, args.locations_max)
        case Variant.SLOW_MO:
            output_lines = _generate_slow_mo(args.n, args.locations_min, args.locations_max)
        case Variant.FIFO:
            output_lines = _generate_fifo(args.n, args.locations_min, args.locations_max)
        case Variant.ANTI_SORT_RIGHT:
            output_lines = _generate_anti_sort_right(args.n, args.locations_min, args.locations_max)
        case Variant.ANTI_SORT_LEFT:
            output_lines = _generate_anti_sort_left(args.n, args.locations_min, args.locations_max)
        case _:
            raise ValueError("Unhandled variant supplied: {args.variant}.")
    return output_lines


def main():
    global RNG
    args = _parse_args()
    _validate_with_context(args)
    RNG = random.Random(args.seed)

    output_lines = generate_output_lines(args)

    output = '\n'.join(output_lines) + '\n'
    command = ' '.join(sys.argv)
    if '--seed' not in sys.argv:
        command += f' --seed {args.seed}'
    if args.test_name is not None:
        test_input_file_name = args.test_name + '.in'
        test_desc_file_name = args.test_name + '.desc'
        with open(test_input_file_name, 'w') as test_input_file:
            test_input_file.write(output)
        with open(test_desc_file_name, 'w') as test_desc_file:
            test_desc_file.write(f'Produced by:\n\t{command}\n')
    else:
        sys.stdout.write(output)


if __name__ == '__main__':
    main()