from pprint import pprint
from typing import List, Dict
from math import log, ceil
import argparse

import pandas as pd
from dwave.system import LeapHybridSampler
import dimod


def knapsack_bqm(cities, values, weights, total_capacity, value_r=0, weight_r=0):
    """
    build the knapsack binary quadratic model

    See # Q-Alpha comments for original contributions
    added parameters:
        value_r: the proportion of value contributed from the objects outside of
                the knapsack. For the standard knapsack problem this is 0,
                but in the case of GDP we might consider that a closed city
                retains some % of GDP
        weight_r: the proportion of weight contributed from the objects outside
                of the knapsack. For the standard knapsack problem this is 0,
                but in the case of sick people we might consider that a closed city
                retains some % of its sick people over time

    From DWave Knapsack examples
    Originally from <NAME>, NP-hard combinatorial problems as Ising spin glasses
    Workshop on Classical and Quantum Optimization; ETH Zuerich - August 20, 2014
    based on Lucas, Frontiers in Physics _2, 5 (2014)    
    """

    # Initialize BQM - use large-capacity BQM so that the problem can be
    # scaled by the user.
    bqm = dimod.AdjVectorBQM(dimod.Vartype.BINARY)

    # Lagrangian multiplier
    # First guess as suggested in Lucas's paper
    lagrange = max(values)

    # Number of objects
    x_size = len(values)

    # Lucas's algorithm introduces additional slack variables to handle
    # the inequality. max_y_index indicates the maximum index in the y
    # sum; hence the number of slack variables.
    max_y_index = ceil(log(total_capacity))

    # Slack variable list for Lucas's algorithm. The last variable has
    # a special value because it terminates the sequence.
    y = [2**n for n in range(max_y_index - 1)]
    y.append(total_capacity + 1 - 2**(max_y_index - 1))

    # Q-Alpha - calculate the extra constant in second part of problem hamiltonian
    C = sum([weight * weight_r for weight in weights])
    # Q-Alpha - change weights to weight*(1-weight_r)
    weights = [weight*(1-weight_r) for weight in weights]

    # Q-Alpha - change values to value*(1-value_r)
    values = [value*(1-value_r) for value in values]

    # Hamiltonian xi-xi terms
    for k in range(x_size):
        # Q-Alpha add final term lagrange * C * weights[k]
        bqm.set_linear(
            cities[k],
            lagrange * (weights[k] ** 2) - values[k] + lagrange * C * weights[k])

    # Hamiltonian xi-xj terms
    for i in range(x_size):
        for j in range(i + 1, x_size):
            key = (cities[i], cities[j])
            bqm.quadratic[key] = 2 * lagrange * weights[i] * weights[j]

    # Hamiltonian y-y terms
    for k in range(max_y_index):
        # Q-Alpha add final term -lagrange * C * y[k]
        bqm.set_linear('y' + str(k), lagrange *
                       (y[k]**2) - lagrange * C * y[k])

    # Hamiltonian yi-yj terms
    for i in range(max_y_index):
        for j in range(i + 1, max_y_index):
            key = ('y' + str(i), 'y' + str(j))
            bqm.quadratic[key] = 2 * lagrange * y[i] * y[j]

    # Hamiltonian x-y terms
    for i in range(x_size):
        for j in range(max_y_index):
            key = (cities[i], 'y' + str(j))
            bqm.quadratic[key] = -2 * lagrange * weights[i] * y[j]

    return bqm


def solve_cities(cities: List, gdps: List, sick: List, total_capacity: int,
                 value_r=0, weight_r=0, num_reads=1, verbose=False) -> Dict:
    """
    Solves problem: "Which cities should I should I shut down in order to stay
    within healthcare resources constraints while maximizing overall GDP"
    parameters:
        cities - list of city names
        gdps - corresponding list of GDP per city
        sick - corresponding number of sick people per city
        total_capacity - max capacity for sick people summed over all cities
        num_reads - number of samples to take
        verbose - whether to print out best result
    returns:
        (dict) - list of dictionaries with individual results and selected attributes
                    sorted in order of least energy first
    """
    if sum(sick) < total_capacity:
        print("Warning in solve_cities: Total number of sick people is less "
              + "than total capacity. There's no knapsack problem to solve!")

    bqm = knapsack_bqm(cities, gdps, sick, total_capacity,
                       value_r=value_r, weight_r=weight_r)
    sampler = LeapHybridSampler()
    samplesets = [sampler.sample(bqm) for _ in range(num_reads)]

    df = pd.DataFrame({'city': cities, 'gdp': gdps, 'sick': sick})
    df = df.set_index('city')

    solution_set = []
    for sampleset in samplesets:
        open_cities = []
        closed_cities = []
        for k, v in sampleset.first.sample.items():
            if k in cities:
                if v == 1:
                    open_cities.append(k)
                else:
                    closed_cities.append(k)
        solution_set.append({
            'open_cities': open_cities,
            'closed_cities': closed_cities,
            'energy': sampleset.first.energy,
            'salvaged_gdp': sum(df.loc[open_cities]['gdp']) + sum(df.loc[closed_cities]['gdp']) * value_r,
            'used_capacity': int(round(sum(df.loc[open_cities]['sick'])))
        })

    # do sorting from lowest to highest energy
    if num_reads > 1:
        energies = [solution['energy'] for solution in solution_set]
        solution_set = [x for _, x in sorted(zip(energies, solution_set))]

    if verbose:
        print('BEST SOLUTION')
        print('Open cities')
        print(solution_set[0]['open_cities'])
        print('\n')
        print('Closed cities')
        print(solution_set[0]['closed_cities'])
        print('\n')
        total_gdp = sum(df['gdp'])
        salvaged_gdp = solution_set[0]['salvaged_gdp']
        print(
            f'Salvaged GDP: {salvaged_gdp} ({(100*salvaged_gdp/total_gdp):.1f}%)')
        used_capacity = solution_set[0]['used_capacity']
        print(
            f'Used up hospital capacity: {used_capacity:d} of {total_capacity} ({(100*used_capacity/total_capacity):.1f}%)')

    return solution_set


def solve_cities_from_csv(filepath: str, total_capacity: int, value_r=0, weight_r=0,
                          num_reads=1, verbose=False) -> Dict:
    """
    Bridge to solve_cities given a csv file for the cities, gdps, and sick
    csv file needs to have header: city,gdp,sick
    """
    df = pd.read_csv(filepath)
    assert ','.join(
        df.columns) == 'city,gdp,sick', "Ensure csv header is city,gdp,sick"
    solution_set = solve_cities(list(df['city']), list(
        df['gdp']), list(df['sick']), total_capacity, value_r=value_r,
        weight_r=weight_r, num_reads=num_reads, verbose=verbose)
    return solution_set


def main():
    """ CLI
    """
    description = "Solve the problem: \"Which cities should I should I shut "
    description += "down in order to stay within healthcare resources "
    description += "constraints while maximizing overall GDP\""
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument('--data', '-d',
                        help="path to csv file with columns: city,gdp,sick. Header expected")
    parser.add_argument('--total-capacity', '-t', type=int,
                        help="Maximum capacity for sick people for all cities combined.")
    parser.add_argument('--num-reads', '-n', type=int, default=1,
                        help="Number of samples to return")
    args = parser.parse_args()
    solution_set = solve_cities_from_csv(
        args.data, args.total_capacity, value_r=0.8, weight_r=0.2,
        num_reads=args.num_reads, verbose=True)


if __name__ == '__main__':
    main()

from pprint import pprint
from typing import List, Dict
from math import log, ceil
import argparse

import pandas as pd
from dwave.system import LeapHybridSampler
import dimod


def knapsack_bqm(cities, values, weights, total_capacity, value_r=0, weight_r=0):
    """
    build the knapsack binary quadratic model

    See # Q-Alpha comments for original contributions
    added parameters:
        value_r: the proportion of value contributed from the objects outside of
                the knapsack. For the standard knapsack problem this is 0,
                but in the case of GDP we might consider that a closed city
                retains some % of GDP
        weight_r: the proportion of weight contributed from the objects outside
                of the knapsack. For the standard knapsack problem this is 0,
                but in the case of sick people we might consider that a closed city
                retains some % of its sick people over time

    From DWave Knapsack examples
    Originally from <NAME>, NP-hard combinatorial problems as Ising spin glasses
    Workshop on Classical and Quantum Optimization; ETH Zuerich - August 20, 2014
    based on Lucas, Frontiers in Physics _2, 5 (2014)    
    """

    # Initialize BQM - use large-capacity BQM so that the problem can be
    # scaled by the user.
    bqm = dimod.AdjVectorBQM(dimod.Vartype.BINARY)

    # Lagrangian multiplier
    # First guess as suggested in Lucas's paper
    lagrange = max(values)

    # Number of objects
    x_size = len(values)

    # Lucas's algorithm introduces additional slack variables to handle
    # the inequality. max_y_index indicates the maximum index in the y
    # sum; hence the number of slack variables.
    max_y_index = ceil(log(total_capacity))

    # Slack variable list for Lucas's algorithm. The last variable has
    # a special value because it terminates the sequence.
    y = [2**n for n in range(max_y_index - 1)]
    y.append(total_capacity + 1 - 2**(max_y_index - 1))

    # Q-Alpha - calculate the extra constant in second part of problem hamiltonian
    C = sum([weight * weight_r for weight in weights])
    # Q-Alpha - change weights to weight*(1-weight_r)
    weights = [weight*(1-weight_r) for weight in weights]

    # Q-Alpha - change values to value*(1-value_r)
    values = [value*(1-value_r) for value in values]

    # Hamiltonian xi-xi terms
    for k in range(x_size):
        # Q-Alpha add final term lagrange * C * weights[k]
        bqm.set_linear(
            cities[k],
            lagrange * (weights[k] ** 2) - values[k] + lagrange * C * weights[k])

    # Hamiltonian xi-xj terms
    for i in range(x_size):
        for j in range(i + 1, x_size):
            key = (cities[i], cities[j])
            bqm.quadratic[key] = 2 * lagrange * weights[i] * weights[j]

    # Hamiltonian y-y terms
    for k in range(max_y_index):
        # Q-Alpha add final term -lagrange * C * y[k]
        bqm.set_linear('y' + str(k), lagrange *
                       (y[k]**2) - lagrange * C * y[k])

    # Hamiltonian yi-yj terms
    for i in range(max_y_index):
        for j in range(i + 1, max_y_index):
            key = ('y' + str(i), 'y' + str(j))
            bqm.quadratic[key] = 2 * lagrange * y[i] * y[j]

    # Hamiltonian x-y terms
    for i in range(x_size):
        for j in range(max_y_index):
            key = (cities[i], 'y' + str(j))
            bqm.quadratic[key] = -2 * lagrange * weights[i] * y[j]

    return bqm


def solve_cities(cities: List, gdps: List, sick: List, total_capacity: int,
                 value_r=0, weight_r=0, num_reads=1, verbose=False) -> Dict:
    """
    Solves problem: "Which cities should I should I shut down in order to stay
    within healthcare resources constraints while maximizing overall GDP"
    parameters:
        cities - list of city names
        gdps - corresponding list of GDP per city
        sick - corresponding number of sick people per city
        total_capacity - max capacity for sick people summed over all cities
        num_reads - number of samples to take
        verbose - whether to print out best result
    returns:
        (dict) - list of dictionaries with individual results and selected attributes
                    sorted in order of least energy first
    """
    if sum(sick) < total_capacity:
        print("Warning in solve_cities: Total number of sick people is less "
              + "than total capacity. There's no knapsack problem to solve!")

    bqm = knapsack_bqm(cities, gdps, sick, total_capacity,
                       value_r=value_r, weight_r=weight_r)
    sampler = LeapHybridSampler()
    samplesets = [sampler.sample(bqm) for _ in range(num_reads)]

    df = pd.DataFrame({'city': cities, 'gdp': gdps, 'sick': sick})
    df = df.set_index('city')

    solution_set = []
    for sampleset in samplesets:
        open_cities = []
        closed_cities = []
        for k, v in sampleset.first.sample.items():
            if k in cities:
                if v == 1:
                    open_cities.append(k)
                else:
                    closed_cities.append(k)
        solution_set.append({
            'open_cities': open_cities,
            'closed_cities': closed_cities,
            'energy': sampleset.first.energy,
            'salvaged_gdp': sum(df.loc[open_cities]['gdp']) + sum(df.loc[closed_cities]['gdp']) * value_r,
            'used_capacity': int(round(sum(df.loc[open_cities]['sick'])))
        })

    # do sorting from lowest to highest energy
    if num_reads > 1:
        energies = [solution['energy'] for solution in solution_set]
        solution_set = [x for _, x in sorted(zip(energies, solution_set))]

    if verbose:
        print('BEST SOLUTION')
        print('Open cities')
        print(solution_set[0]['open_cities'])
        print('\n')
        print('Closed cities')
        print(solution_set[0]['closed_cities'])
        print('\n')
        total_gdp = sum(df['gdp'])
        salvaged_gdp = solution_set[0]['salvaged_gdp']
        print(
            f'Salvaged GDP: {salvaged_gdp} ({(100*salvaged_gdp/total_gdp):.1f}%)')
        used_capacity = solution_set[0]['used_capacity']
        print(
            f'Used up hospital capacity: {used_capacity:d} of {total_capacity} ({(100*used_capacity/total_capacity):.1f}%)')

    return solution_set


def solve_cities_from_csv(filepath: str, total_capacity: int, value_r=0, weight_r=0,
                          num_reads=1, verbose=False) -> Dict:
    """
    Bridge to solve_cities given a csv file for the cities, gdps, and sick
    csv file needs to have header: city,gdp,sick
    """
    df = pd.read_csv(filepath)
    assert ','.join(
        df.columns) == 'city,gdp,sick', "Ensure csv header is city,gdp,sick"
    solution_set = solve_cities(list(df['city']), list(
        df['gdp']), list(df['sick']), total_capacity, value_r=value_r,
        weight_r=weight_r, num_reads=num_reads, verbose=verbose)
    return solution_set


def main():
    """ CLI
    """
    description = "Solve the problem: \"Which cities should I should I shut "
    description += "down in order to stay within healthcare resources "
    description += "constraints while maximizing overall GDP\""
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument('--data', '-d',
                        help="path to csv file with columns: city,gdp,sick. Header expected")
    parser.add_argument('--total-capacity', '-t', type=int,
                        help="Maximum capacity for sick people for all cities combined.")
    parser.add_argument('--num-reads', '-n', type=int, default=1,
                        help="Number of samples to return")
    args = parser.parse_args()
    solution_set = solve_cities_from_csv(
        args.data, args.total_capacity, value_r=0.8, weight_r=0.2,
        num_reads=args.num_reads, verbose=True)


if __name__ == '__main__':
    main()

build the knapsack binary quadratic model      See # Q-Alpha comments for original contributions     added parameters:         value_r: the proportion of value contributed from the objects outside of                 the knapsack. For the standard knapsack problem this is 0,                 but in the case of GDP we might consider that a closed city                 retains some % of GDP         weight_r: the proportion of weight contributed from the objects outside                 of the knapsack. For the standard knapsack problem this is 0,                 but in the case of sick people we might consider that a closed city                 retains some % of its sick people over time      From DWave Knapsack examples     Originally from <NAME>, NP-hard combinatorial problems as Ising spin glasses     Workshop on Classical and Quantum Optimization; ETH Zuerich - August 20, 2014     based on Lucas, Frontiers in Physics _2, 5 (2014) # Initialize BQM - use large-capacity BQM so that the problem can be # scaled by the user. # Lagrangian multiplier # First guess as suggested in Lucas's paper # Number of objects # Lucas's algorithm introduces additional slack variables to handle # the inequality. max_y_index indicates the maximum index in the y # sum; hence the number of slack variables. # Slack variable list for Lucas's algorithm. The last variable has # a special value because it terminates the sequence. # Q-Alpha - calculate the extra constant in second part of problem hamiltonian # Q-Alpha - change weights to weight*(1-weight_r) # Q-Alpha - change values to value*(1-value_r) # Hamiltonian xi-xi terms # Q-Alpha add final term lagrange * C * weights[k] # Hamiltonian xi-xj terms # Hamiltonian y-y terms # Q-Alpha add final term -lagrange * C * y[k] # Hamiltonian yi-yj terms # Hamiltonian x-y terms Solves problem: "Which cities should I should I shut down in order to stay     within healthcare resources constraints while maximizing overall GDP"     parameters:         cities - list of city names         gdps - corresponding list of GDP per city         sick - corresponding number of sick people per city         total_capacity - max capacity for sick people summed over all cities         num_reads - number of samples to take         verbose - whether to print out best result     returns:         (dict) - list of dictionaries with individual results and selected attributes                     sorted in order of least energy first # do sorting from lowest to highest energy Bridge to solve_cities given a csv file for the cities, gdps, and sick     csv file needs to have header: city,gdp,sick CLI