<reponame>Melanee-Melanee/X-Caps
'''
Encoding Visual Attributes in Capsules for Explainable Medical Diagnoses (X-Caps)
Original Paper by <NAME>, <NAME>, and <NAME> (https://arxiv.org/abs/1909.05926)
Code written by: <NAME>
If you use significant portions of this code or the ideas from our paper, please cite it :)
If you have any questions, please email me at <EMAIL>.

This file contains the functions needed to convert the LIDC-IDRI dataset to the expected format.
'''

import os
from fnmatch import filter as fnf
from glob import glob
import xml.etree.ElementTree as ET

from tqdm import tqdm
try:
    import pydicom as pydcm
except:
    import dicom as pydcm
import numpy as np
from scipy.ndimage.morphology import binary_fill_holes
from PIL import Image

from utils import safe_mkdir

def create_cropped_nodules(IMG_ROOT, OUT_ROOT):
    safe_mkdir(OUT_ROOT)

    DEFAULT_PIXEL_SPACING = 0.787109
    DEFAULT_SLICE_THICKNESS = 2.5
    CROP_EXTRA_AMT = (np.sqrt(2)-1)/2

    print('Finding all xml files in LIDC-IDRI')
    matches = []
    for d1 in tqdm(sorted(os.listdir(IMG_ROOT))):
        for d2 in sorted(os.listdir(os.path.join(IMG_ROOT, d1))):
            if d2 == 'AdamsMasks':
                continue
            for d3 in sorted(os.listdir(os.path.join(IMG_ROOT, d1, d2))):
                for f in fnf(os.listdir(os.path.join(IMG_ROOT, d1, d2, d3)), '*.xml'):
                    matches.append(os.path.join(IMG_ROOT, d1, d2, d3, f))

    print('\nCreating cropped images of all nodules in LIDC-IDRI')
    for xml_file in tqdm(matches):
        # Load dicom image
        img_path = os.path.dirname(xml_file)
        dcm_imgs = []
        for dir, _, files in os.walk(img_path):
            for file in fnf(files, '*.dcm'):
                dcm_imgs.append(os.path.join(dir, file))

        # Get ref file
        RefDs = pydcm.read_file(dcm_imgs[0])

        # Load dimensions based on the number of rows, columns, and slices (along the Z axis)
        ConstPixelDims = (int(RefDs.Rows), int(RefDs.Columns), len(dcm_imgs))
        if int(RefDs.Rows) > 512 or int(RefDs.Columns) > 512:
            continue

        # Load spacing values (in mm)
        try:
            pixel_space = [float(RefDs.PixelSpacing[0]), float(RefDs.PixelSpacing[1])]
        except AttributeError as e:
            if str(e) == "'FileDataset' object has no attribute 'PixelSpacing'":
                pixel_space = [DEFAULT_PIXEL_SPACING, DEFAULT_PIXEL_SPACING]
            else:
                raise NotImplementedError('Unhandled exception in pixel spacing.')

        # Load slice thickness (in mm)
        try:
            slice_thick = float(RefDs.SliceThickness)
        except AttributeError as e:
            if str(e) == "'FileDataset' object has no attribute 'SliceThickness'":
                if os.path.basename(xml_file)[:-4] == '243' or os.path.basename(xml_file)[:-4] == '244' or \
                                os.path.basename(xml_file)[:-4] == '070':
                    slice_thick = 2.5
                elif os.path.basename(xml_file)[:-4] == '135':
                    slice_thick = 2.0
                elif os.path.basename(xml_file)[:-4] == '043':
                    slice_thick = 1.8
                else:
                    slice_thick = DEFAULT_SLICE_THICKNESS
            else:
                raise NotImplementedError('Unhandled exception in slice thickness.')

        ConstPixelSpacing = (pixel_space[0], pixel_space[1], slice_thick)

        x = np.arange(0.0, (ConstPixelDims[0] + 1) * ConstPixelSpacing[0], ConstPixelSpacing[0])
        y = np.arange(0.0, (ConstPixelDims[1] + 1) * ConstPixelSpacing[1], ConstPixelSpacing[1])
        z = np.arange(0.0, (ConstPixelDims[2] + 1) * ConstPixelSpacing[2], ConstPixelSpacing[2])

        # The array is sized based on 'ConstPixelDims'
        ArrayDicom = np.zeros(ConstPixelDims, dtype=RefDs.pixel_array.dtype)

        # loop through all the DICOM files
        sop_ids = dict()
        for filenameDCM in dcm_imgs:
            # read the file
            ds = pydcm.read_file(filenameDCM)
            # store the raw image data
            ArrayDicom[:, :, dcm_imgs.index(filenameDCM)] = ds.pixel_array
            sop_ids[ds.SOPInstanceUID] = filenameDCM

        # Load attributes
        tree = ET.parse(xml_file)
        root = tree.getroot()

        unique_nodule_list = []
        curr_nodule = -1
        for s in root.findall('{http://www.nih.gov}ResponseHeader'):
            try:
                study_id = s.find('{http://www.nih.gov}StudyInstanceUID').text
            except:
                study_id = -1
        for r_num, rad in enumerate(root.findall('{http://www.nih.gov}readingSession')):
            try:
                rad_id = rad.find('{http://www.nih.gov}servicingRadiologistID').text
                if rad_id == 'anon':
                    rad_id = 'anon-{:02d}'.format(r_num)
            except:
                rad_id = -1
            for nodule in rad.findall('{http://www.nih.gov}unblindedReadNodule'):
                nodule_id = nodule.find('{http://www.nih.gov}noduleID').text
                sub = -1; ist = -1; cal = -1; sph = -1; mar = -1; lob = -1; spi = -1; tex = -1; mal = -1
                for charac in nodule.findall('{http://www.nih.gov}characteristics'):
                    try:
                        sub = int(charac.find('{http://www.nih.gov}subtlety').text)
                    except:
                        pass
                    try:
                        ist = int(charac.find('{http://www.nih.gov}internalStructure').text)
                    except:
                        pass
                    try:
                        cal = int(charac.find('{http://www.nih.gov}calcification').text)
                    except:
                        pass
                    try:
                        sph = int(charac.find('{http://www.nih.gov}sphericity').text)
                    except:
                        pass
                    try:
                        mar = int(charac.find('{http://www.nih.gov}margin').text)
                    except:
                        pass
                    try:
                        lob = int(charac.find('{http://www.nih.gov}lobulation').text)
                    except:
                        pass
                    try:
                        spi = int(charac.find('{http://www.nih.gov}spiculation').text)
                    except:
                        pass
                    try:
                        tex = int(charac.find('{http://www.nih.gov}texture').text)
                    except:
                        pass
                    try:
                        mal = int(charac.find('{http://www.nih.gov}malignancy').text)
                    except:
                        pass
                slices = []
                x_min = 999999; x_max = -9999999; y_min = 999999; y_max = -9999999
                slice_list = nodule.findall('{http://www.nih.gov}roi')
                GT = np.zeros((ConstPixelDims[0], ConstPixelDims[1], len(slice_list)), dtype=np.uint8)
                for i, roi in enumerate(slice_list):
                    z_pos = -1*float(roi.find('{http://www.nih.gov}imageZposition').text)
                    sop_id = roi.find('{http://www.nih.gov}imageSOP_UID').text
                    for edges in roi.findall('{http://www.nih.gov}edgeMap'):
                        x_pos = int(edges.find('{http://www.nih.gov}xCoord').text)
                        y_pos = int(edges.find('{http://www.nih.gov}yCoord').text)
                        GT[y_pos,x_pos, i] = 1
                        if x_pos < x_min:
                            x_min = x_pos
                        if x_pos > x_max:
                            x_max = x_pos
                        if y_pos < y_min:
                            y_min = y_pos
                        if y_pos > y_max:
                            y_max = y_pos
                    slices.append([sop_id, z_pos])
                    GT[:,:,i] = binary_fill_holes(GT[:,:,i])

                np_slices = np.asarray(slices)
                sorted_slices = np_slices[np_slices[:, 1].argsort()]
                sorted_GT = GT[:,:,np_slices[:, 1].argsort()]

                mean_x = np.mean((x_min, x_max))
                mean_y = np.mean((y_min, y_max))
                mean_z = np.mean((float(sorted_slices[0][1]), float(sorted_slices[-1][1])))
                width = abs(x_max - x_min)
                height = abs(y_max - y_min)
                depth = abs(float(sorted_slices[-1][1]) - float(sorted_slices[0][1]))
                this_nodule = -1
                matched_list = []
                for i, nod_coords in enumerate(unique_nodule_list):
                    if (abs(nod_coords[0] - mean_x) < (nod_coords[3]+width)/4 or abs(nod_coords[0] - mean_x) <= 3) and \
                       (abs(nod_coords[1] - mean_y) < (nod_coords[4]+height)/4 or abs(nod_coords[1] - mean_y) <= 3) and \
                       (abs(nod_coords[2] - mean_z) < (nod_coords[5]+depth)/4 or abs(nod_coords[2] - mean_z) <= 3*slice_thick):
                        # Check for multiple matches
                        matched_list.append([i, np.sqrt((nod_coords[0] - mean_x)**2 + (nod_coords[1] - mean_y)**2 +
                                                        (nod_coords[2] - mean_z)**2)])
                if matched_list:
                    matched_list = np.asarray(matched_list)
                    for match in matched_list[matched_list[:, 1].argsort()]:
                        if not glob(os.path.join(OUT_ROOT, '{}_{}'.format(os.path.basename(os.path.dirname(os.path.dirname(os.path.dirname(xml_file)))), study_id),
                                   'nodule_{:03d}'.format(int(match[0])), 'rad-{}_*'.format(rad_id))):
                            this_nodule = int(match[0])
                            break
                if this_nodule == -1:
                    unique_nodule_list.append([mean_x, mean_y, mean_z, width, height, depth])
                    curr_nodule += 1
                    this_nodule = curr_nodule

                out_dir = os.path.join(OUT_ROOT, '{}_{}'.format(os.path.basename(os.path.dirname(os.path.dirname(os.path.dirname(xml_file)))), study_id),
                               'nodule_{:03d}'.format(this_nodule), 'rad-{}_sub-{}_ist-{}_cal-{}_sph-{}_mar-{}_lob-{}_'
                               'spi-{}_tex-{}_mal-{}'.format(rad_id, sub, ist, cal, sph, mar, lob, spi, tex, mal))
                safe_mkdir(out_dir)

                for i, box in enumerate(sorted_slices):
                    try:
                        dcm_slice = int(dcm_imgs.index(sop_ids[box[0]]))
                    except Exception as e:
                        box[0] = correct_images(box[0])
                        if box[0] != '-1' and box[0] != '-2':
                            dcm_slice = int(dcm_imgs.index(sop_ids[box[0]]))
                        elif box[0] == '-2':
                            continue # This option is for images which cannot be corrected but are unimportant
                        else:
                            print('Unable to locate correct .dcm slice for {}: {}. Please correct by hand.'.format(out_dir, box[0]))
                            print(e)
                            continue

                    h_extra = int(height * CROP_EXTRA_AMT)
                    w_extra = int(width * CROP_EXTRA_AMT)
                    # These are to handle single pixel annotations
                    if h_extra < 2:
                        h_extra = 2
                    if w_extra < 2:
                        w_extra = 2
                    crop = ArrayDicom[y_min-h_extra:y_max+h_extra, x_min-w_extra:x_max+w_extra, dcm_slice]
                    crop_GT = sorted_GT[y_min-h_extra:y_max+h_extra, x_min-w_extra:x_max+w_extra, i]

                    try:
                        # NOTE: Make sure to change values back to int16 from uint16 when reading the images in!!!
                        im = Image.fromarray(crop.astype('<i2'))
                        im.save(os.path.join(out_dir, '{:03d}.tif').format(i))
                        gt_im = Image.fromarray(crop_GT)
                        gt_im.save(os.path.join(out_dir, 'gt_{:03d}.tif').format(i))
                    except Exception as e:
                        print('Unable to create image for: {}.'.format(out_dir))
                        print(e)

    print('\n\nCompleted creating all nodule cropped images for LIDC-IDRI dataset!')

def create_master_list(IMG_ROOT, OUT_ROOT):
    MIN_RADS = 3
    safe_mkdir(OUT_ROOT)

    num_chars_data = np.zeros((9,6))

    print('Finding all nodule files in LIDC-IDRI with >= {} radiologists\' characteristics data'.format(MIN_RADS))
    matches = []
    nodules_total = 0
    nodules_to_use = 0
    non_mal_count = 0
    mal_count = 0

    for study_dir in tqdm(sorted(os.listdir(IMG_ROOT))):
        nodule_list = sorted(os.listdir(os.path.join(IMG_ROOT, study_dir)))
        nodules_total += len(nodule_list)
        for nodule_dir in nodule_list:
            rad_dirs = sorted(os.listdir(os.path.join(IMG_ROOT, study_dir, nodule_dir)))
            rads = len(rad_dirs)
            if rads < MIN_RADS:
                continue
            temp_list = []
            temp_char_data = np.zeros((9,6))
            for rad_dir in rad_dirs:
                if not os.listdir(os.path.join(IMG_ROOT, study_dir, nodule_dir, rad_dir)):
                    rads -= 1 # Make sure there is actually image(s) for this rad
                else:
                    split_names = rad_dir.split('_')
                    all_chars = True
                    for i in range(1,10):
                        if split_names[i][-3:] == '--1':
                            all_chars = False
                        else:
                            temp_char_data[i-1][int(split_names[i][-1])-1] += 1
                    if not all_chars:
                        rads -= 1
                    else:
                        temp_list.append(os.path.join(IMG_ROOT, study_dir, nodule_dir, rad_dir))

            if rads >= MIN_RADS:
                # Compute mean mal score
                char_data_totals = np.zeros((9,rads))
                char_data_stats = np.zeros((9,2))

                for i in range(temp_char_data.shape[0]):
                    c = 0
                    for j in range(temp_char_data.shape[1]):
                        for k in range(int(temp_char_data[i,j])):
                            char_data_totals[i,c] = (j + 1)
                            c += 1
                char_data_stats[:,0] = np.mean(char_data_totals, axis=1)
                char_data_stats[:, 1] = np.std(char_data_totals, axis=1)
                mean_mal = char_data_stats[8,0]

                if mean_mal != 3.:
                    if rads > 4:
                        print('Encountered rads > 4: {}'.format(os.path.join(IMG_ROOT, study_dir, nodule_dir))) # Sanity check for nodule matching
                    if mean_mal > 3.:
                        mal_count += 1
                    else:
                        non_mal_count += 1
                    num_chars_data += temp_char_data
                    nodules_to_use += 1
                    for file_path in temp_list:
                        matches.append([file_path] + list(np.ndarray.flatten(char_data_stats)))

    print('Found {} total nodules.'.format(nodules_total))
    print('Found {} nodules with characteristics and determinable (not score 3) avg malignancy scores.'.format(nodules_to_use))
    print('{} nodules average score below 3.0, {} nodules above 3.0.'.format(non_mal_count, mal_count))

    np.savetxt(os.path.join(OUT_ROOT, 'nodule_characteristics_counts.csv'),
               np.concatenate((np.expand_dims(np.asarray(['sub', 'ist', 'cal', 'sph', 'mar', 'lob', 'spi',
               'tex', 'mal', 'Totals:']), axis=1), np.vstack((num_chars_data.astype(np.int64),
               np.asarray(['Nodules', nodules_to_use, 'Benign', non_mal_count, 'Malig', mal_count])))), axis=1),
               fmt='%s,%s,%s,%s,%s,%s,%s', delimiter=',', header="Characteristics,1,2,3,4,5,6")

    np.savetxt(os.path.join(OUT_ROOT, 'master_nodule_list.csv'), np.asarray(matches),
               fmt='%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s')

def correct_images(sname):
    # This function contains a list of images which previously required manual correction within the LIDC-IDRI Dataset.

    # Manually correct for image 0017
    if sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.305973183883758685859912046949':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.139636132253744151113715840194'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.975363198401611311891539311888':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.225900589792147134785051710110'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.153194632177600377201998652445':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.222098252047357192090439228841'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.329142402178255247031380957411':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.276070543618203204841799986172'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.109012962923033337571132618784':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.201369290021439277502674762620'
    # Manually correct for image 0365
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.112512486762301518180007539984':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.249086187399161659167414756279'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.448378396789516014605561762604':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.216758182207805904911618558070'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.303256875597167746646589593562':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.180833422259316497536094826188'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.261962165647171557143883123825':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.680655051010882131364380217685'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.428441304577336024295581627835':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.282568083479753958511921318301'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.329221218419947342986803210392':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.130147884776737463511106208477'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.162079731049618854270820976684':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.317603920309971419052997711476'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.785736194417664146622972784664':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.180833422259316497536094826188'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.652444697985639935050732394135':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.247436296095192529771061686046'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.126539353916279887972936951408':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.119827225411893011639439591713'
    # Manually correct for image 0566
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.514074599988412913277488312051':
        sname = '-2'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.292826962115570345472638642623':
        sname = '-2'
    # Manually correct for image 0659
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.294658615382614203741435957661':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.248517083496561594434577071132'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.883951945165118277793500546792':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.218208995904859324781331654067'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.227260896757583835259462034815':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.183101167412396355129144409796'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.739565975013005403715405771404':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.276779570196705787348278946110'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.244522908325555679363936146772':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.576955311322527292170312066972'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.125776849447531170933991444187':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.944721105102210115761068591710'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.156959022761131412720241221222':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.685786061228252640465903515314'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.157441085111648851876365968475':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.218208995904859324781331654067'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.137224073243631437732289379681':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.327828856516446064398338817575'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.196786590005502760794118627532':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.205919555392658132555723231924'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.271749899549008749493412118500':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.467531607505823612652093494995'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.132283326090716626749170288137':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.715346929996135559455398127585'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.292014563425807316410737237443':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.251646556878192917000905983161'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.281350813740489812658551562167':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.188830155395223944149966050821'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.316094521169588935447289217773':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.318926701435673382024116339995'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.400249926333575297612413406645':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.247894320876850135016381965868'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.381007938661788498734279329156':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.263902120137234774391883090194'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.158922411981395099005780254611':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.266102462639631998797024975317'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.153968349496782778041856013116':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.313130758239406881022967921981'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.199135326994407563129497784698':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.137821430143892810553323149499'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.328454607276840155088910752459':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.277464141419855638903368659937'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.105901452377957975094355467039':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.270022323575518362032565947858'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.230181645068532680519497368825':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.101175636735586811268012081787'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.173090595736867429956574661962':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.311147794796034131535570099457'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.426419361480558838333009902353':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.283215380710563114133061955920'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.149783315493297937843600113966':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.211773626788832944113459632641'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.333349896902589057387703875126':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.324942356299228484760469569592'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.259329619426001073359049716159':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.218894908479906137103265765511'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.204211053191501804382709873157':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.219041363289039597488091781264'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.365700870941618176907116849738':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.269625131313796127254468189745'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.174649660921460497526396207837':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.499186182774918820678569631767'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.628619440608225619886544814747':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.288908400826505634061200144991'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.165071066866482679435986323504':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.138555834428706707378735123427'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.932783428097248153076463331304':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.403741639352083297611557443868'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.232252382783080336041314614357':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.242386520761336203399531222995'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.304249587531812156369799852687':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.233668843426769210066014174740'
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.547584917033319141420515123587':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.502898181085822091725452394574'
    # Manually correct for image 0931
    elif sname == '1.3.6.1.4.1.14519.5.2.1.6279.6001.265313295605480688537936547605':
        sname = '1.3.6.1.4.1.14519.5.2.1.6279.6001.290994721708875046196354781651'
    else:
        sname = '-1'
    return sname


import math
import numpy as np
import pylab as py
from astropy.table import Table
from astropy.io import fits
import pickle, glob
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import os
from scipy import interpolate
import kai

def setup_phot(imageRoot, silent=False,
               apertures=[25,50,75,100,125,150,175,200],
               sky_annulus=200, sky_dannulus=50, zmag=0):

    from pyraf import iraf as ir

    # Load image header
    hdr = fits.getheader(imageRoot + '.fits')

    ir.digiphot()
    ir.daophot()
    ir.unlearn('phot')
    ir.unlearn('datapars')
    ir.unlearn('centerpars')
    ir.unlearn('fitskypars')
    ir.unlearn('photpars')

    ##########
    # Set up datapars
    ##########
    ir.datapars.fwhmpsf = 5.0 # shouldn't really matter
    ir.datapars.sigma = 'INDEF'
    ir.datapars.datamin = 'INDEF'

    if os.path.exists(imageRoot + '.max'):
        max_file = open(imageRoot + '.max', 'r')
        max_line = max_file.readline()
        max = float(max_line)
        ir.datapars.datamax = max

        if not silent:
            print( 'Set ir.datapars.datamax = %d' % max)

    # Pull gain from the header
    ir.datapars.gain = 'GAIN'
    ir.datapars.epadu = 'INDEF'

    # Assumes 43.1 electrons per read of noise
    nreads = 1.0
    if int(hdr['SAMPMODE']) == 3:
        nreads = int(hdr['MULTISAM'])

    ir.datapars.ccdread = ''
    ir.datapars.readnoise = 43.1 * math.sqrt(2.0) / math.sqrt(nreads)

    # Get exposure times from header
    ir.datapars.exposure = ''
    ir.datapars.itime = float(hdr['ITIME']) * int(hdr['COADDS'])

    # Other Header keywords
    ir.datapars.airmass = 'AIRMASS'
    ir.datapars.filter = 'FWINAME'
    ir.datapars.obstime = 'EXPSTART'

    ##########
    # Setup centerpars. We will use *.coo file for initial guess.
    ##########
    ir.centerpars.calgorithm = 'centroid'

    ##########
    # Setup fitskypars
    ##########
    ir.fitskypars.salgorithm = 'centroid'
    ir.fitskypars.annulus = sky_annulus
    ir.fitskypars.dannulus = sky_dannulus

    ##########
    # Setup photpars
    ##########
    # Setup a zeropoint... this assumes Strehl = 1, but good enough for now.
    ir.photpars.zmag = zmag
    ir.photpars.apertures = ','.join([str(aa) for aa in apertures])

    ##########
    # Setup phot
    ##########
    ir.phot.interactive = 'no'
    ir.phot.radplots = 'no'
    ir.phot.verify = 'No'

    if silent:
        ir.phot.verbose = 'no'
    else:
        ir.phot.verbose = 'yes'

def run_phot(imageRoot, silent=False,
             apertures=[25,50,75,100,125,150,175,200],
             sky_annulus=200, sky_dannulus=50, zmag=0):

    from pyraf import iraf as ir

    setup_phot(imageRoot, apertures=apertures, zmag=zmag, silent=silent,
               sky_annulus=sky_annulus, sky_dannulus=sky_dannulus)

    image = imageRoot + '.fits'
    coords = imageRoot + '.coo'

    # Output into current directory, not data directory
    rootSplit = imageRoot.split('/')
    output = rootSplit[-1] + '.phot.mag'

    ir.phot(image, coords, output)

    (radius, flux, mag, merr) = get_phot_output(output, silent=silent)

    return (radius, flux, mag, merr)

def get_phot_output(output, silent=False):
    from pyraf import iraf as ir

    # Now get the results using txdump
    radStr = ir.txdump(output, 'RAPERT', 'yes', Stdout=1)
    fluxStr = ir.txdump(output, 'FLUX', 'yes', Stdout=1)
    magStr = ir.txdump(output, 'MAG', 'yes', Stdout=1)
    merrStr = ir.txdump(output, 'MERR', 'yes', Stdout=1)
    pierStr = ir.txdump(output, 'PIER', 'yes', Stdout=1)

    radFields = radStr[0].split()
    fluxFields = fluxStr[0].split()
    magFields = magStr[0].split()
    merrFields = merrStr[0].split()
    pierFields = pierStr[0].split()

    count = len(radFields)

    radius = np.zeros(count, dtype=float)
    flux = np.zeros(count, dtype=float)
    mag = np.zeros(count, dtype=float)
    merr = np.zeros(count, dtype=float)

    for rr in range(count):
        radius[rr] = float(radFields[rr])

        if (int(pierFields[rr]) != 0 or magFields[rr] == 'INDEF' or
            merrFields[rr] == 'INDEF'):
            print( 'Problem in image: ' + output)

            # Error
            flux[rr] = 0
            mag[rr] = 0
            merr[rr] = 0
        else:
            flux[rr] = float(fluxFields[rr])
            mag[rr] = float(magFields[rr])
            merr[rr] = float(merrFields[rr])

    if not silent:
        print( '%6s  %10s  %6s  %6s' % ('Radius', 'Flux', 'Mag', 'MagErr'))
        for ii in range(count):
            print( '%8.1f  %10d  %6.3f  %6.3f' % \
                (radius[ii], flux[ii], mag[ii], merr[ii]))

    return (radius, flux, mag, merr)

def get_filter_profile(filter):
    """
    Returns the wavelength (in microns) and the transmission for
    the specified NIRC2 filter.

    Example:
    (wave, trans) = kai.photometry.get_filter_profile('Kp')
    py.clf()
    py.plot(wave, trans)
    py.xlabel('Wavelength (microns)')
    py.ylabel('Transmission')
    """
    base_path = os.path.dirname(kai.__file__)
    rootDir = base_path + '/filters/'

    filters = ['J', 'H', 'K', 'Kcont', 'Kp', 'Ks', 'Lp', 'Ms',
               'Hcont', 'Brgamma', 'FeII']

    if filter not in filters:
        print( 'Could not find profile for filter %s.' % filter)
        print( 'Choices are: ', filters)
        return

    table = Table.read(rootDir + filter + '.dat', format='ascii')

    wavelength = table[table.colnames[0]]
    transmission = table[table.colnames[1]]

    # Lets fix wavelength array for duplicate values
    diff = np.diff(wavelength)
    idx = np.where(diff <= 0)[0]
    wavelength[idx+1] += 1.0e-7

    # Get rid of all entries with negative transmission
    idx = np.where(transmission > 1)[0]
    wavelength = wavelength[idx]
    transmission = transmission[idx] / 100.0 # convert from % to ratio

    return (wavelength, transmission)

def test_filter_profile_interp():
    """
    Plot up the filter transmission curves and their interpolations
    for the three K-band filters (K, Kp, Ks).
    """
    # Get the transmission curve for NIRC2 filters and atmosphere.
    K_wave, K_trans = get_filter_profile('K')
    Kp_wave, Kp_trans = get_filter_profile('Kp')
    Ks_wave, Ks_trans = get_filter_profile('Ks')
    J_wave, J_trans = get_filter_profile('J')
    H_wave, H_trans = get_filter_profile('H')
    Lp_wave, Lp_trans = get_filter_profile('Lp')

    # We will need to resample these transmission curves.
    print( 'Creating interp object')
    K_interp = interpolate.splrep(K_wave, K_trans, k=1, s=0)
    Kp_interp = interpolate.splrep(Kp_wave, Kp_trans, k=1, s=0)
    Ks_interp = interpolate.splrep(Ks_wave, Ks_trans, k=1, s=0)
    J_interp = interpolate.splrep(J_wave, J_trans, k=1, s=0)
    H_interp = interpolate.splrep(H_wave, H_trans, k=1, s=0)
    Lp_interp = interpolate.splrep(Lp_wave, Lp_trans, k=1, s=0)

    K_wave_new = np.arange(K_wave.min(), K_wave.max(), 0.0005)
    Kp_wave_new = np.arange(Kp_wave.min(), Kp_wave.max(), 0.0005)
    Ks_wave_new = np.arange(Ks_wave.min(), Ks_wave.max(), 0.0005)
    J_wave_new = np.arange(J_wave.min(), J_wave.max(), 0.0005)
    H_wave_new = np.arange(H_wave.min(), H_wave.max(), 0.0005)
    Lp_wave_new = np.arange(Lp_wave.min(), Lp_wave.max(), 0.0005)

    print( 'Interpolating')
    K_trans_new = interpolate.splev(K_wave_new, K_interp)
    Kp_trans_new = interpolate.splev(Kp_wave_new, Kp_interp)
    Ks_trans_new = interpolate.splev(Ks_wave_new, Ks_interp)
    J_trans_new = interpolate.splev(J_wave_new, J_interp)
    H_trans_new = interpolate.splev(H_wave_new, H_interp)
    Lp_trans_new = interpolate.splev(Lp_wave_new, Lp_interp)

    print( 'Plotting')
#     py.figure(2, figsize=(4,4))
#     py.subplots_adjust(left=0.2, bottom=0.14, top=0.95, right=0.94)
    py.clf()
    py.plot(K_wave, K_trans, 'bo', ms=4, label='_nolegend_', mec='blue')
    py.plot(K_wave_new, K_trans_new, 'b-', label='K', linewidth=2)

    py.plot(Kp_wave, Kp_trans, 'ro', ms=4, label='_nolegend_', mec='red')
    py.plot(Kp_wave_new, Kp_trans_new, 'r-', label='Kp', linewidth=2)

    py.plot(Ks_wave, Ks_trans, 'go', ms=4, label='_nolegend_', mec='green')
    py.plot(Ks_wave_new, Ks_trans_new, 'g-', label='Ks', linewidth=2)

    py.plot(J_wave, J_trans, 'go', ms=4, label='_nolegend_', mec='green')
    py.plot(J_wave_new, J_trans_new, 'g-', label='J', linewidth=2)

    py.plot(H_wave, H_trans, 'go', ms=4, label='_nolegend_', mec='green')
    py.plot(H_wave_new, H_trans_new, 'g-', label='H', linewidth=2)

    py.plot(Lp_wave, Lp_trans, 'go', ms=4, label='_nolegend_', mec='green')
    py.plot(Lp_wave_new, Lp_trans_new, 'g-', label='Lp', linewidth=2)

    py.legend(loc='lower right', numpoints=1, markerscale=0.1)
    py.xlabel('Wavelength (microns)')
    py.ylabel('Transmission (%)')

#     py.axis([2.110, 2.120, 0.928, 0.945])

def test_atmosphere_profile_interp():
    atmDir = '/u/jlu/data/w51/09jun26/weather/atmosphere_transmission.dat'
    atmData = Table.read(atmDir, format='ascii')
    atm_wave = atmData[atmData.colnames[0]]
    atm_trans = atmData[atmData.colnames[1]]

    atm_interp = interpolate.splrep(atm_wave, atm_trans, k=1, s=1)

    atm_wave_new = np.arange(2.0, 2.4, 0.0005)
    atm_trans_new = interpolate.splev(atm_wave_new, atm_interp)

    py.clf()
    py.plot(atm_wave, atm_trans, 'r.', ms=2)
    py.plot(atm_wave_new, atm_trans_new, 'b-')
    py.xlabel('Wavelength (microns)')
    py.ylabel('Transmission (%)')
    py.xlim(2, 2.4)


from pathlib import Path
from random import choice
from statistics import mean
from typing import Any, Counter

import typer
from click import Choice as click_Choice
from tqdm import tqdm  # type: ignore
from typer import echo as typer_echo
from typer import prompt as typer_prompt
from typer import secho as typer_secho
from typer import style as typer_style

from wordle_solver.solver import (
    calculate_entropies,
    get_initial_hints_dict,
    map_all_hints,
)
from wordle_solver.utils import DEFAULT_WORDLIST_PATH, calculate_hints, load_word_list

app = typer.Typer(add_completion=False)  # pragma: no cover
state: dict[Any, Any] = {}  # pragma: no cover

def get_hint_feedback(chosen_word: str) -> tuple[int, ...]:  # pragma: no cover
    while True:
        typer_echo("b = black (grey) / y = yellow / g = green")

        hint_input = typer_prompt(
            "Please provide the result for '"
            + typer_style(chosen_word, fg="bright_magenta")
            + "'"
        )

        if not all(char in set("byg") for char in hint_input):
            typer_secho(
                "Invalid input. Please only use the characters 'b', 'y', or 'g'",
                fg="red",
            )
            typer_echo("Try again...")
            continue

        if len(hint_input) != 5:
            typer_secho("Invalid input. Input must be 5 characters long.", fg="red")
            typer_echo("Try again...")
            continue

        break

    input_to_output: dict[str, int] = {"g": 2, "y": 1, "b": 0}

    return tuple(input_to_output[input_char] for input_char in hint_input)

@app.command()  # pragma: no cover
def stats(
    word_list: Path = typer.Option(
        DEFAULT_WORDLIST_PATH,
        "--word-list",
        "-w",
        help="The directory location containing the word list files.",
    ),
) -> None:
    _, all_guesses, all_solutions = load_word_list(word_list)
    initial_hint_dict = get_initial_hints_dict(all_guesses, all_solutions)

    temp_entropies = calculate_entropies(all_guesses, all_solutions, initial_hint_dict)
    top_initial_guesses: list[str] = list(temp_entropies.keys())[0:10]

    for initial_guess in top_initial_guesses:
        guesses_to_solved: list[int] = []

        for secret in tqdm(all_solutions):
            solved = False
            iteration_count = 1

            # The first guess is the most computational, so try to keep it fast
            hint = calculate_hints(initial_guess, secret)
            local_solutions = list(initial_hint_dict[initial_guess][hint])
            local_hint_dict = map_all_hints(all_guesses, local_solutions)

            while not solved:
                iteration_count += 1
                total_remaining_solutions = len(local_solutions)

                if total_remaining_solutions == 1:
                    chosen_guess = local_solutions[0]
                elif total_remaining_solutions == 2:
                    chosen_guess = choice(local_solutions)
                else:
                    entropies = calculate_entropies(
                        all_guesses, local_solutions, local_hint_dict
                    )
                    chosen_guess = list(entropies.keys())[0]

                # typer_echo(f"{chosen_guess=}")
                hint = calculate_hints(chosen_guess, secret)
                if hint == (2, 2, 2, 2, 2):
                    solved = True
                    guesses_to_solved.append(iteration_count)
                    break

                local_solutions = list(local_hint_dict[chosen_guess][hint])
                local_hint_dict = map_all_hints(all_guesses, local_solutions)

        typer_echo(f"Using guess: {initial_guess}")
        typer_echo(f"Total words solved: {len(guesses_to_solved)}")
        typer_echo(f"Average guess: {round(mean(guesses_to_solved), 3)}")
        typer_echo(f"Total counts: {Counter(guesses_to_solved)}")

@app.command()  # pragma: no cover
def solver(
    word_list: Path = typer.Option(
        DEFAULT_WORDLIST_PATH,
        "--word-list",
        "-w",
        help="The directory location containing the word list files.",
    ),
) -> None:
    """Solve a word on Wordle."""
    _, guesses, solutions = load_word_list(word_list)
    hint_dict = get_initial_hints_dict(guesses, solutions)

    for _ in range(6):
        total_remaining_solutions = len(solutions)
        if total_remaining_solutions == 0:
            # TODO: Error out?
            pass
        elif total_remaining_solutions == 1:
            typer_echo("There is only 1 reamining possible answer:")
            typer_echo(typer_style(solutions, bg="magenta"))

            chosen_guess = solutions[0]
        elif total_remaining_solutions == 2:
            typer_echo("There are 2 remaining possible answers:")
            typer_echo(typer_style(solutions, bg="magenta"))

            chosen_guess = typer_prompt(
                "Please select a guess word:",
                type=click_Choice(solutions),
            )
        else:
            entropies = calculate_entropies(guesses, solutions, hint_dict)
            top_8_entropies = dict(list(entropies.items())[0:8])

            typer_echo(
                f"There are {total_remaining_solutions} remaining possible answers."
            )
            if total_remaining_solutions < 8:
                typer_echo(typer_style(solutions, bg="magenta"))

            typer_echo("The top 8 guesses are:")
            typer_echo(typer_style(top_8_entropies, bg="blue"))

            chosen_guess = typer_prompt(
                "Please select a guess word:",
                type=click_Choice(list(top_8_entropies.keys())),
            )

        hint_feedback = get_hint_feedback(chosen_guess)
        if hint_feedback == (2, 2, 2, 2, 2):
            typer_secho("Victory!", fg="bright_green")
            return

        solutions = list(hint_dict[chosen_guess][hint_feedback])
        hint_dict = map_all_hints(guesses, solutions)

@app.callback()  # pragma: no cover
def main() -> None:
    # """
    # Manage users in the awesome CLI app.
    # """
    pass

if __name__ == "__main__":
    app()


<gh_stars>0
# -*- coding: utf-8 -*-
"""Approval-based committee (ABC) voting rules"""

from __future__ import print_function
import sys
import functools
from itertools import combinations
try:
    from gmpy2 import mpq as Fraction
except ImportError:
    print("Warning: module gmpy2 not found, "
          + "resorting to Python's fractions.Fraction")
    from fractions import Fraction
from abcvoting import abcrules_gurobi
from abcvoting.misc import sort_committees
from abcvoting.misc import hamming
from abcvoting.misc import enough_approved_candidates
from abcvoting.misc import str_committees_header
from abcvoting.misc import str_candset, str_candsets
from abcvoting.misc import header
from abcvoting import scores

########################################################################

class UnknownRuleIDError(ValueError):
    """Exception raised if unknown rule id is used"""

    def __init__(self, expression, message):
        self.expression = expression
        self.message = "Rule ID \"" + str(message) + "\" is not known."

class ABCRule():
    """Class for ABC rules containing basic information and function call"""
    def __init__(self, rule_id, shortname, longname, fct,
                 algorithms=["standard"], resolute=[True, False]):
        self.rule_id = rule_id
        self.shortname = shortname
        self.longname = longname
        self.fct = fct
        self.algorithms = algorithms
        # algorithms should be sorted by speed (fastest first)
        self.resolute = resolute

        assert len(resolute) > 0
        assert len(algorithms) > 0

    def compute(self, profile, committeesize, **kwargs):
        return self.fct(profile, committeesize, **kwargs)

    def fastest_algo(self):
        for algo in self.algorithms:
            if algo == "gurobi" and not abcrules_gurobi.available:
                continue
            return algo

########################################################################

def compute(rule_id, profile, committeesize, **kwargs):
    try:
        return rules[rule_id].compute(profile, committeesize, **kwargs)
    except KeyError:
        raise NotImplementedError("ABC rule " + str(rule_id) + " not known.")

# computes arbitrary Thiele methods via branch-and-bound
def compute_thiele_method(scorefct_str, profile, committeesize,
                          algorithm="branch-and-bound",
                          resolute=False, verbose=0):
    """Thiele methods

    Compute winning committees of the Thiele method specified
    by the score function (scorefct_str)
    """
    enough_approved_candidates(profile, committeesize)
    scorefct = scores.get_scorefct(scorefct_str, committeesize)

    # optional output
    if verbose:
        print(header(rules[scorefct_str].longname))
        if resolute:
            print("Computing only one winning committee (resolute=True)\n")
    if verbose >= 3:
        if algorithm == "gurobi":
            print("Using the Gurobi ILP solver\n")
        if algorithm == "branch-and-bound":
            print("Using a branch-and-bound algorithm\n")
    # end of optional output

    if algorithm == "gurobi":
        committees = abcrules_gurobi.__gurobi_thiele_methods(
            profile, committeesize, scorefct, resolute)

        committees = sort_committees(committees)
    elif algorithm == "branch-and-bound":
        committees = __thiele_methods_branchandbound(
            profile, committeesize, scorefct_str, resolute)
    else:
        raise NotImplementedError(
            "Algorithm " + str(algorithm)
            + " not specified for compute_thiele_method")

    # optional output
    if verbose >= 2:
        print("Optimal " + scorefct_str.upper() + "-score: "
              + str(scores.thiele_score(scorefct_str, profile, committees[0])))
        print()
    if verbose:
        print(str_committees_header(committees, winning=True))
        print(str_candsets(committees, names=profile.names))
    # end of optional output

    return committees

# computes arbitrary Thiele methods via branch-and-bound
def __thiele_methods_branchandbound(profile, committeesize,
                                    scorefct_str, resolute):
    """Branch-and-bound algorithm to compute winning committees
    for Thiele methods"""
    enough_approved_candidates(profile, committeesize)
    scorefct = scores.get_scorefct(scorefct_str, committeesize)

    best_committees = []
    init_com = compute_seq_thiele_method(
        profile, committeesize, scorefct_str, resolute=True)[0]
    best_score = scores.thiele_score(scorefct_str, profile, init_com)
    part_coms = [[]]
    while part_coms:
        part_com = part_coms.pop(0)
        # potential committee, check if at least as good
        # as previous best committee
        if len(part_com) == committeesize:
            score = scores.thiele_score(scorefct_str, profile, part_com)
            if score == best_score:
                best_committees.append(part_com)
            elif score > best_score:
                best_committees = [part_com]
                best_score = score
        else:
            if len(part_com) > 0:
                largest_cand = part_com[-1]
            else:
                largest_cand = -1
            missing = committeesize - len(part_com)
            marg_util_cand = scores.marginal_thiele_scores_add(
                scorefct, profile, part_com)
            upper_bound = (
                sum(sorted(marg_util_cand[largest_cand + 1:])[-missing:])
                + scores.thiele_score(scorefct_str, profile, part_com))
            if upper_bound >= best_score:
                for c in range(largest_cand + 1,
                               profile.num_cand - missing + 1):
                    part_coms.insert(0, part_com + [c])

    committees = sort_committees(best_committees)
    if resolute:
        committees = [committees[0]]

    return committees

# Sequential PAV
def compute_seqpav(profile, committeesize, algorithm="standard",
                   resolute=True, verbose=0):
    """Sequential PAV (seq-PAV)"""
    return compute_seq_thiele_method(
        profile, committeesize, 'pav', algorithm=algorithm,
        resolute=resolute, verbose=verbose)

def compute_seqslav(profile, committeesize, algorithm="standard",
                    resolute=True, verbose=0):
    """Sequential Sainte-Lague Approval Voting (SLAV)"""
    return compute_seq_thiele_method(
        profile, committeesize, "slav", algorithm=algorithm,
        resolute=resolute, verbose=verbose)

# Reverse Sequential PAV
def compute_revseqpav(profile, committeesize, algorithm="standard",
                      resolute=True, verbose=0):
    """Reverse sequential PAV (revseq-PAV)"""
    return compute_revseq_thiele_method(
        profile, committeesize, 'pav', algorithm=algorithm,
        resolute=resolute, verbose=verbose)

def compute_seqcc(profile, committeesize, algorithm="standard",
                  resolute=True, verbose=0):
    """Sequential Chamberlin-Courant (seq-CC)"""
    return compute_seq_thiele_method(
        profile, committeesize, 'cc', algorithm=algorithm,
        resolute=resolute, verbose=verbose)

def compute_sav(profile, committeesize, algorithm="standard",
                resolute=False, verbose=0):
    """Satisfaction Approval Voting (SAV)"""
    if algorithm == "standard":
        return __separable("sav", profile, committeesize, resolute, verbose)
    else:
            raise NotImplementedError(
                "Algorithm " + str(algorithm)
                + " not specified for compute_sav")

# Approval Voting (AV)
def compute_av(profile, committeesize, algorithm="standard",
               resolute=False, verbose=0):
    """Approval Voting"""
    if algorithm == "standard":
        return __separable("av", profile, committeesize, resolute, verbose)
    else:
        raise NotImplementedError(
            "Algorithm " + str(algorithm)
            + " not specified for compute_av")

def __separable(rule_id, profile, committeesize, resolute, verbose):
    enough_approved_candidates(profile, committeesize)

    appr_scores = [0] * profile.num_cand
    for pref in profile:
        for cand in pref:
            if rule_id == "sav":
                # Satisfaction Approval Voting
                appr_scores[cand] += Fraction(pref.weight, len(pref))
            elif rule_id == "av":
                # (Classic) Approval Voting
                appr_scores[cand] += pref.weight
            else:
                raise UnknownRuleIDError(rule_id)

    # smallest score to be in the committee
    cutoff = sorted(appr_scores)[-committeesize]

    certain_cands = [c for c in range(profile.num_cand)
                     if appr_scores[c] > cutoff]
    possible_cands = [c for c in range(profile.num_cand)
                      if appr_scores[c] == cutoff]
    missing = committeesize - len(certain_cands)
    if len(possible_cands) == missing:
        # candidates with appr_scores[c] == cutoff
        # are also certain candidates because all these candidates
        # are required to fill the committee
        certain_cands = sorted(certain_cands + possible_cands)
        possible_cands = []
        missing = 0

    if resolute:
        committees = sort_committees(
            [(certain_cands + possible_cands[:missing])])
    else:
        committees = sort_committees(
            [(certain_cands + list(selection))
             for selection
             in combinations(possible_cands, missing)])

    # optional output
    if verbose:
        print(header(rules[rule_id].longname))
        if resolute:
            print("Computing only one winning committee (resolute=True)\n")
    if verbose >= 2:
        print("Scores of candidates:")
        for c in range(profile.num_cand):
            print(profile.names[c] + ": " + str(appr_scores[c]))

        print("\nCandidates are contained in winning committees")
        print("if their score is >= " + str(cutoff) + ".")

        if len(certain_cands) > 0:
            print("\nThe following candidates are contained in")
            print("every winning committee:")
            namedset = [profile.names[c] for c in certain_cands]
            print(" " + ", ".join(map(str, namedset)))
            print()

        if len(possible_cands) > 0:
            print("The following candidates are contained in")
            print("some of the winning committees:")
            namedset = [profile.names[c] for c in possible_cands]
            print(" " + ", ".join(map(str, namedset)))
            print("(" + str(missing) + " of those candidates is contained\n"
                  + " in every winning committee.)\n")
    if verbose:
        print(str_committees_header(committees, winning=True))
        print(str_candsets(committees, names=profile.names))
    # end of optional output

    return committees

def __seq_thiele_resolute(profile, committeesize, scorefct_str, verbose):
    """Compute a *resolute* reverse sequential Thiele method

    Tiebreaking between candidates in favor of candidate with smaller
    number/index (candidates with larger numbers get deleted first).
    """
    committee = []

    scorefct = scores.get_scorefct(scorefct_str, committeesize)

    # optional output
    if verbose >= 2:
        output = "starting with the empty committee (score = "
        output += str(scores.thiele_score(
            scorefct_str, profile, committee)) + ")"
        print(output + "\n")
    # end of optional output

    # build a committee starting with the empty set
    for _ in range(committeesize):
        additional_score_cand = scores.marginal_thiele_scores_add(
            scorefct, profile, committee)
        next_cand = additional_score_cand.index(max(additional_score_cand))
        committee.append(next_cand)
        # optional output
        if verbose >= 2:
            output = "adding candidate number "
            output += str(len(committee)) + ": "
            output += profile.names[next_cand] + "\n"
            output += " score increases by "
            output += str(max(additional_score_cand))
            output += " to a total of "
            output += str(scores.thiele_score(
                scorefct_str, profile, committee))
            tied_cands = [c for c in range(len(additional_score_cand))
                          if (c > next_cand and
                              (additional_score_cand[c]
                               == max(additional_score_cand)))]
            if len(tied_cands) > 0:
                output += " tie broken in favor of " + str(next_cand)
                output += " candidates " + str_candset(tied_cands)
                output += " would increase the score by the same amount ("
                output += str(max(additional_score_cand)) + ")"
            print(output + "\n")
        # end of optional output
    return [sorted(committee)]

def __seq_thiele_irresolute(profile, committeesize, scorefct_str):
    """Compute an *irresolute* reverse sequential Thiele method

    Consider all possible ways to break ties between candidates
    (aka parallel universe tiebreaking)
    """
    scorefct = scores.get_scorefct(scorefct_str, committeesize)

    comm_scores = {(): 0}
    # build committees starting with the empty set
    for _ in range(committeesize):
        comm_scores_next = {}
        for committee, score in comm_scores.items():
            # marginal utility gained by adding candidate to the committee
            additional_score_cand = scores.marginal_thiele_scores_add(
                scorefct, profile, committee)
            for c in range(profile.num_cand):
                if additional_score_cand[c] >= max(additional_score_cand):
                    next_comm = tuple(sorted(committee + (c,)))
                    comm_scores_next[next_comm] = (
                        score + additional_score_cand[c])
        comm_scores = comm_scores_next
    return sort_committees(list(comm_scores.keys()))

def compute_seq_thiele_method(profile, committeesize, scorefct_str,
                              algorithm="standard", resolute=True, verbose=0):
    """Sequential Thiele methods"""

    enough_approved_candidates(profile, committeesize)

    if algorithm != "standard":
        raise NotImplementedError(
            "Algorithm " + str(algorithm)
            + " not specified for compute_seq_thiele_method")

    # optional output
    if verbose:
        print(header(rules["seq" + scorefct_str].longname))
        if resolute:
            print("Computing only one winning committee (resolute=True)\n")
    # end of optional output

    if resolute:
        committees = __seq_thiele_resolute(
            profile, committeesize, scorefct_str, verbose=verbose)
    else:
        committees = __seq_thiele_irresolute(
            profile, committeesize, scorefct_str)

    # optional output
    if verbose:
        print(str_committees_header(committees, winning=True))
        print(str_candsets(committees, names=profile.names))
        if resolute or len(committees) == 1:
            print(scorefct_str.upper() + "-score of winning committee:",
                  end="")
        else:
            print(scorefct_str.upper() + "-score of winning committees:")
        for comm in committees:
            print(" " + str(scores.thiele_score(scorefct_str, profile, comm)))
        print()
    # end of optional output

    return committees

def __revseq_thiele_irresolute(profile, committeesize, scorefct_str):
    """Compute an *irresolute* sequential Thiele method

    Consider all possible ways to break ties between candidates
    (aka parallel universe tiebreaking)
    """
    scorefct = scores.get_scorefct(scorefct_str, committeesize)

    allcandcomm = tuple(range(profile.num_cand))
    comm_scores = {allcandcomm: scores.thiele_score(
        scorefct_str, profile, allcandcomm)}

    for _ in range(profile.num_cand - committeesize):
        comm_scores_next = {}
        for committee, score in comm_scores.items():
            marg_util_cand = scores.marginal_thiele_scores_remove(
                scorefct, profile, committee)
            score_reduction = min(marg_util_cand)
            # find smallest elements in marg_util_cand and return indices
            cands_to_remove = [cand for cand in range(profile.num_cand)
                               if marg_util_cand[cand] == min(marg_util_cand)]
            for c in cands_to_remove:
                next_comm = tuple(set(committee) - set([c]))
                comm_scores_next[next_comm] = score - score_reduction
            comm_scores = comm_scores_next
    return sort_committees(list(comm_scores.keys()))

def __revseq_thiele_resolute(profile, committeesize, scorefct_str, verbose):
    """Compute a *resolute* reverse sequential Thiele method

    Tiebreaking between candidates in favor of candidate with smaller
    number/index (candidates with smaller numbers are added first).
    """
    scorefct = scores.get_scorefct(scorefct_str, committeesize)

    committee = set(range(profile.num_cand))

    # optional output
    if verbose >= 2:
        output = "full committee (" + str(len(committee))
        output += " candidates) has a total score of "
        output += str(scores.thiele_score(
            scorefct_str, profile, committee))
        print(output + "\n")
    # end of optional output

    for _ in range(profile.num_cand - committeesize):
        marg_util_cand = scores.marginal_thiele_scores_remove(
            scorefct, profile, committee)
        score_reduction = min(marg_util_cand)
        # find smallest elements in marg_util_cand and return indices
        cands_to_remove = [cand for cand in range(profile.num_cand)
                           if marg_util_cand[cand] == min(marg_util_cand)]
        committee.remove(cands_to_remove[-1])

        # optional output
        if verbose >= 2:
            rem_cand = cands_to_remove[-1]
            output = "removing candidate number "
            output += str(profile.num_cand - len(committee)) + ": "
            output += profile.names[rem_cand] + "\n"
            output += " score decreases by "
            output += str(score_reduction)
            output += " to a total of "
            output += str(scores.thiele_score(
                scorefct_str, profile, committee))
            if len(cands_to_remove) > 1:
                output += " (tie between candidates "
                output += str_candset(cands_to_remove) + ")\n"
            print(output + "\n")
        # end of optional output

    return [sorted(list(committee))]

def compute_revseq_thiele_method(profile, committeesize,
                                 scorefct_str, algorithm="standard",
                                 resolute=True, verbose=0):
    """Reverse sequential Thiele methods"""
    enough_approved_candidates(profile, committeesize)

    if algorithm != "standard":
        raise NotImplementedError(
            "Algorithm " + str(algorithm)
            + " not specified for compute_revseq_thiele_method")

    # optional output
    if verbose:
        print(header(rules["revseq" + scorefct_str].longname))
        if resolute:
            print("Computing only one winning committee (resolute=True)\n")
    # end of optional output

    if resolute:
        committees = __revseq_thiele_resolute(
            profile, committeesize, scorefct_str, verbose=verbose)
    else:
        committees = __revseq_thiele_irresolute(
            profile, committeesize, scorefct_str)

    # optional output
    if verbose:
        print(str_committees_header(committees, winning=True))
        print(str_candsets(committees, names=profile.names))
    if verbose >= 2:
        if resolute or len(committees) == 1:
            print("PAV-score of winning committee:", end="")
        else:
            print("PAV-score of winning committees:")
        for comm in committees:
            print(" " + str(scores.thiele_score(scorefct_str, profile, comm)))
        print()
    # end of optional output

    return committees

def __minimaxav_bruteforce(profile, committeesize):
    """Brute-force algorithm for computing Minimax AV (MAV)"""
    opt_committees = []
    opt_mavscore = profile.num_cand + 1
    for comm in combinations(list(range(profile.num_cand)), committeesize):
        score = scores.mavscore(profile, comm)
        if score < opt_mavscore:
            opt_committees = [comm]
            opt_mavscore = score
        elif scores.mavscore(profile, comm) == opt_mavscore:
            opt_committees.append(comm)

    committees = sort_committees(opt_committees)

    return committees

# Minimax Approval Voting
def compute_mav(profile, committeesize, algorithm="brute-force",
                resolute=False, verbose=0):
    """Minimax AV (MAV)"""
    enough_approved_candidates(profile, committeesize)

    # optional output
    if verbose:
        print(header(rules["mav"].longname))
        if resolute:
            print("Computing only one winning committee (resolute=True)\n")
    if verbose >= 3:
        if algorithm == "gurobi":
            print("Using the Gurobi ILP solver\n")
        if algorithm == "brute-force":
            print("Using a brute-force algorithm\n")
    # end of optional output

    if algorithm == "gurobi":
        committees = abcrules_gurobi.__gurobi_minimaxav(
            profile, committeesize, resolute)
        committees = sort_committees(committees)
    elif algorithm == "brute-force":
        committees = __minimaxav_bruteforce(profile, committeesize)
        if resolute:
            committees = [committees[0]]
    else:
        raise NotImplementedError("Algorithm " + str(algorithm)
                                  + " not specified for compute_mav")

    opt_mavscore = scores.mavscore(profile, committees[0])

    # optional output
    if verbose:
        print("Minimum maximal distance: " + str(opt_mavscore))

        print(str_committees_header(committees, winning=True))
        print(str_candsets(committees, names=profile.names))

        print("Corresponding distances to voters:")
        for comm in committees:
            print([hamming(pref, comm) for pref in profile])
        print()
    # end of optional output

    return committees

# Lexicographic Minimax Approval Voting
def compute_lexmav(profile, committeesize, algorithm="brute-force",
                   resolute=False, verbose=0):
    """Lexicographic Minimax AV"""
    enough_approved_candidates(profile, committeesize)

    if not profile.has_unit_weights():
        raise ValueError(rules["lexmav"].shortname +
                         " is only defined for unit weights (weight=1)")

    if algorithm != "brute-force":
        raise NotImplementedError(
            "Algorithm " + str(algorithm)
            + " not specified for compute_lexmav")

    opt_committees = []
    opt_distances = [profile.num_cand + 1] * len(profile)
    for comm in combinations(list(range(profile.num_cand)), committeesize):
        distances = sorted([hamming(pref, comm)
                            for pref in profile],
                           reverse=True)
        for i in range(len(distances)):
            if opt_distances[i] < distances[i]:
                break
            if opt_distances[i] > distances[i]:
                opt_distances = distances
                opt_committees = [comm]
                break
        else:
            opt_committees.append(comm)

    committees = sort_committees(opt_committees)
    if resolute:
        committees = [committees[0]]

    # optional output
    if verbose:
        print(header(rules["lexmav"].longname))
        if resolute:
            print("Computing only one winning committee (resolute=True)\n")

        print("Minimum maximal distance: " + str(max(opt_distances)))

        print(str_committees_header(committees, winning=True))
        print(str_candsets(committees, names=profile.names))

        print("Corresponding distances to voters:")
        for comm in committees:
            print([hamming(pref, comm) for pref in profile])
        print()
    # end of optional output

    return committees

# Proportional Approval Voting
def compute_pav(profile, committeesize, algorithm="branch-and-bound",
                resolute=False, verbose=0):
    """Proportional Approval Voting (PAV)"""
    return compute_thiele_method(
        'pav', profile, committeesize, algorithm=algorithm,
        resolute=resolute, verbose=verbose)

# Sainte-Lague Approval Voting
def compute_slav(profile, committeesize, algorithm="branch-and-bound",
                 resolute=False, verbose=0):
    """Sainte-Lague Approval Voting (SLAV)"""
    return compute_thiele_method(
        'slav', profile, committeesize, algorithm=algorithm,
        resolute=resolute, verbose=verbose)

# Chamberlin-Courant
def compute_cc(profile, committeesize, algorithm="branch-and-bound",
               resolute=False, verbose=0):
    """Approval Chamberlin-Courant (CC)"""
    return compute_thiele_method(
        'cc', profile, committeesize, algorithm=algorithm,
        resolute=resolute, verbose=verbose)

def compute_monroe(profile, committeesize, algorithm="brute-force",
                   resolute=False, verbose=0):
    """Monroe's rule"""
    enough_approved_candidates(profile, committeesize)

    # optional output
    if verbose:
        print(header(rules["monroe"].longname))
        if resolute:
            print("Computing only one winning committee (resolute=True)\n")
    if verbose >= 3:
        if algorithm == "gurobi":
            print("Using the Gurobi ILP solver\n")
        if algorithm == "brute-force":
            print("Using a brute-force algorithm\n")
    # end of optional output

    if not profile.has_unit_weights():
        raise ValueError(rules["monroe"].shortname +
                         " is only defined for unit weights (weight=1)")

    if algorithm == "gurobi":
        committees = abcrules_gurobi.__gurobi_monroe(
            profile, committeesize, resolute)
        committees = sort_committees(committees)
    elif algorithm == "brute-force":
        committees = __monroe_bruteforce(
            profile, committeesize, resolute)
    else:
        raise NotImplementedError(
            "Algorithm " + str(algorithm)
            + " not specified for compute_monroe")

    # optional output
    if verbose:
        print("Optimal Monroe score: "
              + str(scores.monroescore(profile, committees[0])) + "\n")

        print(str_committees_header(committees, winning=True))
        print(str_candsets(committees, names=profile.names))
    # end of optional output

    return committees

# Monroe's rule, computed via (brute-force) matching
def __monroe_bruteforce(profile, committeesize, resolute):
    """Brute-force computation of Monroe's rule"""
    opt_committees = []
    opt_monroescore = -1
    for comm in combinations(list(range(profile.num_cand)), committeesize):
        score = scores.monroescore(profile, comm)
        if score > opt_monroescore:
            opt_committees = [comm]
            opt_monroescore = score
        elif scores.monroescore(profile, comm) == opt_monroescore:
            opt_committees.append(comm)

    committees = sort_committees(opt_committees)
    if resolute:
        committees = [committees[0]]

    return committees

def compute_greedy_monroe(profile, committeesize,
                          algorithm="standard", resolute=True, verbose=0):
    """"Greedy Monroe"""
    enough_approved_candidates(profile, committeesize)
    if not profile.has_unit_weights():
        raise ValueError(rules["greedy-monroe"].shortname +
                         " is only defined for unit weights (weight=1)")

    if not resolute:
        raise NotImplementedError(
            "compute_greedy_monroe does not support resolute=False.")

    if algorithm != "standard":
        raise NotImplementedError(
            "Algorithm " + str(algorithm)
            + " not specified for compute_greedy_monroe")

    num_voters = len(profile)
    committee = []

    # remaining voters
    remaining_voters = list(range(num_voters))
    remaining_cands = set(range(profile.num_cand))

    assignment = []
    for t in range(committeesize):
        maxapprovals = -1
        selected = None
        for c in remaining_cands:
            approvals = len([i for i in remaining_voters
                             if c in profile[i]])
            if approvals > maxapprovals:
                maxapprovals = approvals
                selected = c

        # determine how many voters are removed (at most)
        if t < num_voters - committeesize * (num_voters // committeesize):
            num_remove = num_voters // committeesize + 1
        else:
            num_remove = num_voters // committeesize

        # only voters that approve the chosen candidate
        # are removed
        to_remove = [i for i in remaining_voters
                     if selected in profile[i]]
        if len(to_remove) > num_remove:
            to_remove = to_remove[:num_remove]
        assignment.append((selected, to_remove))
        remaining_voters = [i for i in remaining_voters
                            if i not in to_remove]
        committee.append(selected)
        remaining_cands.remove(selected)

    committees = sort_committees([committee])

    # optional output
    if verbose:
        print(header(rules["greedy-monroe"].longname))

    if verbose >= 2:
        score1 = scores.monroescore(profile, committees[0])

        score2 = len(profile) - len(remaining_voters)
        print("The Monroe assignment computed by <NAME>")
        print("has a Monroe score of " + str(score2) + ".")

        if score1 > score2:
            print("Monroe assignment found by Greedy Monroe is not "
                  + "optimal for the winning committee,")
            print("i.e., by redistributing voters to candidates a higher "
                  + "satisfaction is possible "
                  + "(without changing the committee).")
            print("Optimal Monroe score of the winning committee is "
                  + str(score1) + ".")

        # build actual Monroe assignment for winning committee
        for t, district in enumerate(assignment):
            cand, voters = district
            if t < num_voters - committeesize * (num_voters // committeesize):
                missing = num_voters // committeesize + 1 - len(voters)
            else:
                missing = num_voters // committeesize - len(voters)
            for _ in range(missing):
                v = remaining_voters.pop()
                voters.append(v)

        print("Assignment (unsatisfatied voters marked with *):\n")
        for cand, voters in assignment:
            print(" candidate " + profile.names[cand]
                  + " assigned to: ", end="")
            output = ""
            for v in sorted(voters):
                output += str(v)
                if cand not in profile[v].approved:
                    output += "*"
                output += ", "
            print(output[:-2])
        print()

    if verbose:
        print(str_committees_header(committees, winning=True))
        print(str_candsets(committees, names=profile.names))
    # end of optional output

    return committees

def __seqphragmen_resolute(profile, committeesize, verbose,
                           start_load=None, partial_committee=None):
    """Algorithm for computing resolute seq-Phragmen  (1 winning committee)"""
    approvers_weight = {}
    for c in range(profile.num_cand):
        approvers_weight[c] = sum(pref.weight for pref in profile if c in pref)

    load = start_load
    if load is None:
        load = {v: 0 for v, _ in enumerate(profile)}

    committee = partial_committee
    if partial_committee is None:
        committee = []  # build committees starting with the empty set

    for _ in range(len(committee), committeesize):
        approvers_load = {}
        for c in range(profile.num_cand):
            approvers_load[c] = sum(pref.weight * load[v]
                                    for v, pref in enumerate(profile)
                                    if c in pref)
        new_maxload = [Fraction(approvers_load[c] + 1, approvers_weight[c])
                       if approvers_weight[c] > 0 else committeesize + 1
                       for c in range(profile.num_cand)]
        # exclude committees already in the committee
        large = max(new_maxload) + 1
        for c in range(profile.num_cand):
            if c in committee:
                new_maxload[c] = large
        # find smallest maxload
        opt = min(new_maxload)
        next_cand = new_maxload.index(opt)
        # compute new loads and add new candidate
        for v, pref in enumerate(profile):
            if next_cand in pref:
                load[v] = new_maxload[next_cand]
            else:
                load[v] = load[v]
        committee = sorted(committee + [next_cand])

        # optional output
        if verbose >= 2:
            output = "adding candidate number "
            output += str(len(committee)) + ": "
            output += profile.names[next_cand] + "\n"
            output += " maximum load increased to "
            output += str(opt)
            print(output)
            print(" load distribution:")
            output = "  ("
            for v, _ in enumerate(profile):
                output += str(load[v]) + ", "
            print(output[:-2] + ")")
            tied_cands = [c for c in range(profile.num_cand)
                          if (c > next_cand and
                              (new_maxload[c] == new_maxload))]
            if len(tied_cands) > 0:
                output = " tie broken in favor of " + profile.names[next_cand]
                output += ",\n candidates " + str_candset(tied_cands)
                output += " would increase the load to the same amount ("
                output += str(new_maxload) + ")"
                print(output)
            print()
        # end of optional output

    comm_loads = {tuple(committee): load}
    return [committee], comm_loads

def __seqphragmen_irresolute(profile, committeesize,
                             start_load=None, partial_committee=None):
    """Algorithm for computing irresolute seq-Phragmen (>=1 winning committees)
    """
    approvers_weight = {}
    for c in range(profile.num_cand):
        approvers_weight[c] = sum(pref.weight for pref in profile if c in pref)

    load = start_load
    if load is None:
        load = {v: 0 for v, _ in enumerate(profile)}

    if partial_committee is None:
        partial_committee = []  # build committees starting with the empty set
    comm_loads = {tuple(partial_committee): load}

    for _ in range(len(partial_committee), committeesize):
        comm_loads_next = {}
        for committee, load in comm_loads.items():
            approvers_load = {}
            for c in range(profile.num_cand):
                approvers_load[c] = sum(pref.weight * load[v]
                                        for v, pref in enumerate(profile)
                                        if c in pref)
            new_maxload = [
                Fraction(approvers_load[c] + 1, approvers_weight[c])
                if approvers_weight[c] > 0 else committeesize + 1
                for c in range(profile.num_cand)]
            # exclude committees already in the committee
            for c in range(profile.num_cand):
                if c in committee:
                    new_maxload[c] = sys.maxsize
            # compute new loads
            # and add new committees
            for c in range(profile.num_cand):
                if new_maxload[c] <= min(new_maxload):
                    new_load = {}
                    for v, pref in enumerate(profile):
                        if c in pref:
                            new_load[v] = new_maxload[c]
                        else:
                            new_load[v] = load[v]
                    new_comm = tuple(sorted(committee + (c,)))
                    comm_loads_next[new_comm] = new_load
        comm_loads = comm_loads_next

    committees = sort_committees(list(comm_loads.keys()))
    return committees, comm_loads

def compute_seqphragmen(profile, committeesize, algorithm="standard",
                        resolute=True, verbose=False):
    """Phragmen's sequential rule (seq-Phragmen)"""
    enough_approved_candidates(profile, committeesize)

    if algorithm != "standard":
        raise NotImplementedError(
            "Algorithm " + str(algorithm)
            + " not specified for compute_seqphragmen")

    # optional output
    if verbose:
        print(header(rules["seqphrag"].longname))
        if resolute:
            print("Computing only one winning committee (resolute=True)\n")
    # end of optional output

    if resolute:
        committees, comm_loads = __seqphragmen_resolute(
            profile, committeesize, verbose)
    else:
        committees, comm_loads = __seqphragmen_irresolute(
            profile, committeesize)

    # optional output
    if verbose:
        print(str_committees_header(committees, winning=True))
        print(str_candsets(committees, names=profile.names))
    if verbose >= 2:
        if resolute or len(committees) == 1:
            print("corresponding load distribution:")
        else:
            print("corresponding load distributions:")
        for comm in committees:
            output = "("
            for v, _ in enumerate(profile):
                output += str(comm_loads[tuple(comm)][v]) + ", "
            print(output[:-2] + ")")
    # end of optional output

    return committees

def __rule_x_get_min_q(profile, budget, cand):
    rich = set([v for v, pref in enumerate(profile)
                if cand in pref])
    poor = set()

    while len(rich) > 0:
        poor_budget = sum(budget[v] for v in poor)
        q = Fraction(1 - poor_budget, len(rich))
        new_poor = set([v for v in rich
                        if budget[v] < q])
        if len(new_poor) == 0:
            return q
        rich -= new_poor
        poor.update(new_poor)

    return None  # not sufficient budget available

def compute_rule_x(profile, committeesize, algorithm="standard",
                   resolute=True, verbose=0):
    """Rule X

    See https://arxiv.org/pdf/1911.11747.pdf, page 7
    """
    enough_approved_candidates(profile, committeesize)
    if not profile.has_unit_weights():
        raise ValueError(rules["rule-x"].shortname +
                         " is only defined for unit weights (weight=1)")

    if algorithm != "standard":
        raise NotImplementedError(
            "Algorithm " + str(algorithm)
            + " not specified for compute_rule_x")

    # optional output
    if verbose:
        print(header(rules["rule-x"].longname))
        if resolute:
            print("Computing only one winning committee (resolute=True)\n")
    # end of optional output

    start_budget = {v: Fraction(committeesize, len(profile))
                    for v, _ in enumerate(profile)}
    cands = range(profile.num_cand)
    commbugdets = [(set(), start_budget)]
    final_committees = set()

    # optional output
    if resolute and verbose >= 2:
        print("Phase 1:\n")
        print("starting budget:")
        output = "  ("
        for v, _ in enumerate(profile):
            output += str(start_budget[v]) + ", "
        print(output[:-2] + ")\n")
    # end of optional output

    for _ in range(committeesize):
        next_commbudgets = []
        for committee, budget in commbugdets:

            curr_cands = set(cands) - committee
            min_q = {}
            for c in curr_cands:
                q = __rule_x_get_min_q(profile, budget, c)
                if q is not None:
                    min_q[c] = q

            if len(min_q) > 0:  # one or more candidates are affordable
                next_cands = [c for c in min_q.keys()
                              if min_q[c] == min(min_q.values())]
                for next_cand in next_cands:
                    new_budget = dict(budget)
                    for v, pref in enumerate(profile):
                        if next_cand in pref:
                            new_budget[v] -= min(budget[v], min_q[next_cand])
                    new_comm = set(committee)
                    new_comm.add(next_cand)
                    next_commbudgets.append((new_comm, new_budget))

                    # optional output
                    if resolute and verbose >= 2:
                        output = "adding candidate number "
                        output += str(len(committee)) + ": "
                        output += profile.names[next_cand] + "\n"
                        output += " with maxmimum cost per voter q = "
                        output += str(min(min_q.values()))
                        print(output)
                        print(" remaining budget:")
                        output = "  ("
                        for v, _ in enumerate(profile):
                            output += str(new_budget[v]) + ", "
                        print(output[:-2] + ")")
                        if len(next_cands) > 1:
                            output = " tie broken in favor of "
                            output += profile.names[next_cand] + ","
                            output += "\n candidates "
                            output += str_candset(next_cands[1:])
                            output += " are tied"
                            print(output)
                        print()
                    # end of optional output

                    if resolute:
                        break

            else:  # no affordable candidates remain
                # fill committee via seq-Phragmen

                # optional output
                if resolute and verbose >= 2:
                    print("Phase 2 (seq-Phragmén):\n")
                # end of optional output

                start_load = {}
                # translate budget to loads
                for v in range(len(profile)):
                    start_load[v] = (Fraction(committeesize, len(profile))
                                     - budget[v])

                # optional output
                if resolute and verbose >= 2:
                    print("starting loads (= budget spent):")
                    output = "  ("
                    for v, _ in enumerate(profile):
                        output += str(start_load[v]) + ", "
                    print(output[:-2] + ")\n")
                # end of optional output

                if resolute:
                    committees, _ = __seqphragmen_resolute(
                        profile, committeesize, verbose=verbose,
                        partial_committee=list(committee),
                        start_load=start_load)
                else:
                    committees, _ = __seqphragmen_irresolute(
                        profile, committeesize,
                        partial_committee=list(committee),
                        start_load=start_load)
                final_committees.update([tuple(comm) for comm in committees])
                # after filling the remaining spots these committees
                # have size committeesize

            commbugdets = next_commbudgets

    final_committees.update([tuple(comm) for comm, _ in commbugdets])

    committees = sort_committees(final_committees)
    if resolute:
        committees = committees[:1]

    # optional output
    if verbose:
        print(str_committees_header(committees, winning=True))
        print(str_candsets(committees, names=profile.names))
    # end of optional output

    return committees

def compute_optphragmen(profile, committeesize,
                        algorithm="gurobi", resolute=False, verbose=0):
    enough_approved_candidates(profile, committeesize)

    # optional output
    if verbose:
        print(header(rules["optphrag"].longname))
        if resolute:
            print("Computing only one winning committee (resolute=True)\n")
    if verbose >= 3:
        if algorithm == "gurobi":
            print("Using the Gurobi ILP solver")
    # end of optional output

    if algorithm != "gurobi":
        raise NotImplementedError("Algorithm " + str(algorithm)
                                  + " not specified for compute_optphragmen")

    committees = abcrules_gurobi.__gurobi_optphragmen(
        profile, committeesize, resolute=resolute, verbose=verbose)
    committees = sort_committees(committees)

    # optional output
    if verbose:
        print(str_committees_header(committees, winning=True))
        print(str_candsets(committees, names=profile.names))
    # end of optional output

    return committees

def compute_phragmen_enestroem(profile, committeesize, algorithm="standard",
                               resolute=True, verbose=0):
    """"Phragmen-Enestroem (aka Phragmen's first method, Enestroem's method)

    In every round the candidate with the highest combined budget of
    their supporters is put in the committee.
    Method described in:
    https://arxiv.org/pdf/1611.08826.pdf (Section 18.5, Page 59)
    """
    enough_approved_candidates(profile, committeesize)
    if not profile.has_unit_weights():
        raise ValueError(rules["phrag-enestr"].shortname +
                         " is only defined for unit weights (weight=1)")

    if algorithm != "standard":
        raise NotImplementedError(
            "Algorithm " + str(algorithm)
            + " not specified for compute_phragmen_enestroem")

    num_voters = len(profile)

    start_budget = {i: Fraction(profile[i].weight)
                    for i in range(num_voters)}
    price = Fraction(sum(start_budget.values()), committeesize)

    cands = range(profile.num_cand)

    committees = [(start_budget, set())]
    for _ in range(committeesize):
        # here the committees with i+1 candidates are
        # stored (together with budget)
        next_committees = []
        # loop in case multiple possible committees
        # with i filled candidates
        for committee in committees:
            budget, comm = committee
            curr_cands = set(cands) - comm
            support = {c: 0 for c in curr_cands}
            for nr, pref in enumerate(profile):
                voting_power = budget[nr]
                if voting_power <= 0:
                    continue
                for cand in pref:
                    if cand in curr_cands:
                        support[cand] += voting_power
            max_support = max(support.values())
            winners = [c for c, s in support.items()
                       if s == max_support]
            for cand in winners:
                b = dict(budget)  # copy of budget
                if max_support > price:  # supporters can afford it
                    # (voting_power - price) / voting_power
                    multiplier = Fraction(max_support - price,
                                          max_support)
                else:  # set supporters to 0
                    multiplier = 0
                for nr, pref in enumerate(profile):
                    if cand in pref:
                        b[nr] *= multiplier
                c = comm.union([cand])  # new committee with candidate
                next_committees.append((b, c))

        if resolute:
            committees = [next_committees[0]]
        else:
            committees = next_committees
    committees = [comm for b, comm in committees]
    committees = sort_committees(committees)
    if resolute:
        committees = [committees[0]]

    # optional output
    if verbose:
        print(header(rules["phrag-enestr"].longname))

        print(str_committees_header(committees, winning=True))
        print(str_candsets(committees, names=profile.names))
    # end of optional output

    return committees

__RULESINFO = [
    ("av", "AV", "Approval Voting (AV)", compute_av,
     ["standard"], [True, False]),
    ("sav", "SAV", "Satisfaction Approval Voting (SAV)", compute_sav,
     ["standard"], [True, False]),
    ("pav", "PAV", "Proportional Approval Voting (PAV)", compute_pav,
     ["gurobi", "branch-and-bound"], [True, False]),
    ("slav", "SLAV", "Sainte-Laguë Approval Voting (SLAV)", compute_slav,
     ["gurobi", "branch-and-bound"], [True, False]),
    ("cc", "CC", "Approval Chamberlin-Courant (CC)", compute_cc,
     ["gurobi", "branch-and-bound"], [True, False]),
    ("geom2", "2-Geometric", "2-Geometric Rule",
     functools.partial(compute_thiele_method, "geom2"),
     ["gurobi", "branch-and-bound"], [True, False]),
    ("seqpav", "seq-PAV", "Sequential Proportional Approval Voting (seq-PAV)",
     compute_seqpav, ["standard"], [True, False]),
    ("revseqpav", "revseq-PAV",
     "Reverse Sequential Proportional Approval Voting (revseq-PAV)",
     compute_revseqpav, ["standard"], [True, False]),
    ("seqslav", "seq-SLAV",
     "Sequential Sainte-Laguë Approval Voting (seq-SLAV)",
     compute_seqslav, ["standard"], [True, False]),
    ("seqcc", "seq-CC", "Sequential Approval Chamberlin-Courant (seq-CC)",
     compute_seqcc, ["standard"], [True, False]),
    ("seqphrag", "seq-Phragmén", "Phragmén's Sequential Rule (seq-Phragmén)",
     compute_seqphragmen, ["standard"], [True, False]),
    ("optphrag", "opt-Phragmén", "Phragmén's Optimization Rule (opt-Phragmén)",
     compute_optphragmen, ["gurobi"], [True, False]),
    ("monroe", "Monroe", "Monroe's Approval Rule (Monroe)",
     compute_monroe, ["gurobi", "brute-force"], [True, False]),
    ("greedy-monroe", "Greedy Monroe", "Greedy Monroe",
     compute_greedy_monroe, ["standard"], [True]),
    ("mav", "MAV", "Minimax Approval Voting (MAV)",
     compute_mav, ["gurobi", "brute-force"], [True, False]),
    ("lexmav", "lex-MAV", "Lexicographic Minimax Approval Voting (lex-MAV)",
     compute_lexmav, ["brute-force"], [True, False]),
    ("rule-x", "Rule X", "Rule X",
     compute_rule_x, ["standard"], [True, False]),
    ("phrag-enestr", "Phragmén-Eneström", "Method of Phragmén-Eneström",
     compute_phragmen_enestroem, ["standard"], [True, False])]
rules = {}
for ruleinfo in __RULESINFO:
    rules[ruleinfo[0]] = ABCRule(*ruleinfo)
# TODO: add other thiele methods