#!/usr/bin/python3

import argparse
import random
import sys
from collections import defaultdict
from enum import Enum


class Variant(Enum):
    CIRCULANT = 'circulant'
    RANDOM = 'random'
    CYCLE_AND_RANDOM = 'cycle_and_random'
    PUSH_DISKS = 'push_disks'

    def __str__(self):
        return self.value


N_MIN = 1
N_MAX = 200
M_MIN = 1
M_MAX = 500
BEADS_MIN = 0
BEADS_MAX = 10000
# 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_m(string):
    return bounded_int(string, M_MIN, M_MAX, 'm')


def bounded_beads(string):
    return bounded_int(string, BEADS_MIN, BEADS_MAX, 'beads bound')


def bounded_target(string):
    return bounded_int(string, 1, BEADS_MAX // N_MAX, 'target answer')


def _parse_args():
    parser = argparse.ArgumentParser('')
    parser.add_argument(
        'n',
        type=bounded_n,
        default=200,
        help='The value of n.'
    )
    parser.add_argument(
        'm',
        type=bounded_m,
        default=500,
        help='The value of m.'
    )
    parser.add_argument(
        '--beads-min',
        type=bounded_beads,
        default=BEADS_MIN,
        help='Minimum number of beads on an edge',
    )
    parser.add_argument(
        '--beads-max',
        type=bounded_beads,
        default=BEADS_MAX,
        help='Maximum number of beads on an edge'
    )
    parser.add_argument(
        '--variant',
        type=Variant,
        choices=list(Variant),
        default=Variant.RANDOM,
        help='Type of test case to generate.',
    )
    parser.add_argument(
        '--push-disks-target',
        type=bounded_target,
        default=100,
        help='The number of disks to start with on each node when pushing '
             'disks to edges.',
    )
    parser.add_argument(
        "--force-no-isolated",
        action="store_true",
        default=False,
        help='Ensure that there are no isolated nodes.'
    )
    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.m > (args.n * (args.n - 1)) // 2:
        raise ValueError('Too many edges for the number of nodes')
    if args.beads_max < args.beads_min:
        raise ValueError('Beads max cannot be less than beads min')
    if args.variant is Variant.CYCLE_AND_RANDOM and args.n > args.m:
        raise ValueError('Need at least as many edges as nodes to make a cycle')
    if args.force_no_isolated and args.m < args.n:
        raise ValueError(
            'Need at least as many edges as nodes to use --force-no-isolated. '
            'We could support m=n-1, but I did not want to code making a tree.'
        )


def _generate_random(n, m, beads_min, beads_max):
    edge_space = []
    for i in range(1, n + 1):
        for j in range(i+1, n+1):
            edge_space.append((i, j))
    edges = RNG.sample(edge_space, k=m)
    lines = [f"{n} {m}"]
    for edge in edges:
        n1, n2 = edge
        beads = RNG.randint(beads_min, beads_max)
        lines.append(f"{n1} {n2} {beads}")
    return lines


def _generate_circulant(n, m, beads_min, beads_max):
    edges = []
    offset = 1
    while len(edges) < m:
        for i in range(n):
            u = i + 1
            v = ((i + offset) % n) + 1
            beads = RNG.randint(beads_min, beads_max)
            edges.append((u, v, beads))
            if len(edges) >= m:
                break
        offset += 1
    lines = [f"{n} {m}"]
    for edge in edges:
        u, v, b = edge
        lines.append(f"{u} {v} {b}")
    return lines


def _generate_cycle_and_random(n, m, beads_min, beads_max):
    edge_space = set()
    for i in range(1, n+1):
        for j in range(i+1, n+1):
            edge_space.add((i, j))
    edges = []
    for i in range(n):
        edges.append((i+1, ((i+1) % n)+1))
        edge_space.remove(tuple(sorted(edges[-1])))
    edges.extend(RNG.sample(sorted(edge_space), k=m-n))
    lines = [f"{n} {m}"]
    for edge in edges:
        u, v = edge
        beads = RNG.randint(beads_min, beads_max)
        lines.append(f"{u} {v} {beads}")
    return lines


def _generate_random_push_disks(n, m, push_disks_target, force_no_isolated=False):
    edge_space = []
    for i in range(1, n + 1):
        for j in range(i + 1, n + 1):
            edge_space.append((i, j))
    nodes_seen = set()
    edges_chosen = set()
    if force_no_isolated:
        while len(nodes_seen) < n:
            edge = RNG.choice(edge_space)
            u, v = edge
            while edge in edges_chosen or (u in nodes_seen and v in nodes_seen):
                edge = RNG.choice(edge_space)
                u, v = edge
            edges_chosen.add(edge)
            nodes_seen.add(u)
            nodes_seen.add(v)
    edge_space = [edge for edge in edge_space if edge not in edges_chosen]
    edges = list(edges_chosen)
    edges.extend(RNG.sample(edge_space, k=m - len(edges_chosen)))
    edge_disks = defaultdict(int)
    nodes = defaultdict(list)
    for edge in edges:
        u, v = edge
        nodes[u].append(edge)
        nodes[v].append(edge)
    for node_edges in nodes.values():
        for i in range(push_disks_target):
            edge = RNG.choice(node_edges)
            edge_disks[edge] += 1
    lines = [f"{n} {m}"]
    for (u, v), disks in edge_disks.items():
        lines.append(f"{u} {v} {disks}")
    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.m, args.beads_min, args.beads_max)
        case Variant.CIRCULANT:
            output_lines = _generate_circulant(args.n, args.m, args.beads_min, args.beads_max)
        case Variant.CYCLE_AND_RANDOM:
            output_lines = _generate_cycle_and_random(args.n, args.m, args.beads_min, args.beads_max)
        case Variant.PUSH_DISKS:
            output_lines = _generate_random_push_disks(args.n, args.m, args.push_disks_target, args.force_no_isolated)
        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()