## \example domino/marina_party.py
#
# This is a not very serious example that shows how to use domino from scratch
# to solve a problem. It is illustrative of creating your own type of particles,
# a pair score, particle states for DOMINO, a decorator, and using subset filter
# tables with DOMINO.
#
# Solution of the Marina Party Problem with DOMINO:
#
# - Girls attend a party
# - Each girl has dresses of different price
# - There can not be repeated dresses in the party within a group of connected
# people
# - Girls that are friends can call each other to agree
# - The objective is to maximize the total price of all dresses in the party
#
#

from __future__ import print_function
import IMP
import IMP.domino
import IMP.core
import random
import sys

IMP.setup_from_argv(sys.argv, "marina party")

class SumPricePairScore(IMP.PairScore):

    def evaluate_index(self, m, pair, accum):
        price0 = Price(m.get_particle(pair[0])).get_price()
        price1 = Price(m.get_particle(pair[1])).get_price()
        return price0 + price1

    def do_get_inputs(self, m, pis):
        return IMP.get_particles(m, pis)


class PriceStates(IMP.domino.ParticleStates):

    def __init__(self, prices):
        self.prices = prices
        IMP.domino.ParticleStates.__init__(self)

    def load_particle_state(self, i, particle):
        pr = Price(particle)
        pr.set_price(self.prices[i])

    def get_number_of_particle_states(self):
        return len(self.prices)


class Price(IMP.Decorator):
    price_key = IMP.IntKey("price")

    def __init__(self, p):
        if not self.get_is_setup(p):
            self.setup_particle(p)
        self.particle = p

    def setup_particle(self, p):
        p.add_attribute(self.price_key, 0)
        self.particle = p

    def get_price(self):
        return self.particle.get_value(self.price_key)

    def set_price(self, x):
        self.particle.set_value(self.price_key, int(x))

    def get_is_setup(self, p):
        return p.has_attribute(self.price_key)


def get_total_price(states_table, subset, assignment):
    total_price = 0
    for i, p in enumerate(subset):
        price_states = states_table.get_particle_states(p)
        price_states.load_particle_state(assignment[i], p)
        price = Price(p).get_price()
        total_price += price
    return total_price


def print_assignment(states_table, subset, assignment):
    total_price = 0
    print("########## solution assignment", assignment)
    for i, p in enumerate(subset):
        price_states = states_table.get_particle_states(p)
        price_states.load_particle_state(assignment[i], p)
        price = Price(p).get_price()
        print(p.get_name(), "price", price)


n_girls = 10
n_edges = 20
girls = []
prices = [100, 200, 400, 600, 800]


model = IMP.Model()

# prepare filter tables for DOMINO
states_table = IMP.domino.ParticleStatesTable()
sampler = IMP.domino.DominoSampler(model, states_table)
all_possible_states = PriceStates(prices)


# set_states
for i in range(n_girls):
    p = IMP.Particle(model, "girl-%d" % i)
    pr = Price(p)
    girls.append(p)
    # add possible dresses
    states_table.set_particle_states(p, all_possible_states)

    # dresses
    if len(prices) == 1:
        n_dresses = 1
    else:
        n_dresses = random.randrange(1, len(prices) + 1)

    # Each girl has a selection of dresses
    selection = random.sample(prices, n_dresses)
    allowed_states_indices = [prices.index(price) for price in selection]
    print(p.get_name(), "prices selected", selection, "indices", allowed_states_indices)
    list_states_table = IMP.domino.ListSubsetFilterTable(states_table)
    list_states_table.set_allowed_states(p, allowed_states_indices)
    sampler.add_subset_filter_table(list_states_table)

# create restraints
rs = []
for z in range(n_edges):
    # pair of friends
    i = random.randrange(0, n_girls)
    j = random.randrange(0, n_girls)
    friends = IMP.ParticlePair(girls[i], girls[j])
    # restraint
    score = SumPricePairScore()
    rs.append(IMP.core.PairRestraint(model, score, friends))
    # Exclusion states. Two girls can't have same dress
    ft = IMP.domino.ExclusionSubsetFilterTable()
    ft.add_pair(friends)
    sampler.add_subset_filter_table(ft)
sampler.set_restraints(rs)

subset = states_table.get_subset()
solutions = sampler.get_sample_assignments(subset)

if len(solutions) == 0:
    print("There are no solutions to the problem")
else:
    most_expensive = 0
    best_solution = solutions[0]
    for assignment in solutions:
        total_price = get_total_price(states_table, subset, assignment)
        if(total_price > most_expensive):
            most_expensive = total_price
            best_solution = assignment
    print(" There are", len(solutions), "possible solutions")
    print("=================> BEST SOLUTION <==============")
    print_assignment(states_table, subset, best_solution)
