1import random
2from collections import defaultdict
3
4import freeorion as fo
5
6import universe_statistics
7import universe_tables
8
9
10# REPEAT_RATE along with calculate_number_of_specials_to_place determines if there are multiple
11# specials in a single location.  There can only be at most 4 specials in a single location.
12# The probabilites break down as follows:
13# Count  Probability
14# one    (1 - REPEAT_RATE[0])
15# two    REPEAT_RATE[0] * (1 - REPEAT_RATE[1])
16# three  REPEAT_RATE[0] * REPEAT_RATE[1] *(1 - REPEAT_RATE[2])
17# four   REPEAT_RATE[0] * REPEAT_RATE[1] * REPEAT_RATE[2]
18REPEAT_RATE = {1: 0.08, 2: 0.05, 3: 0.01, 4: 0.00}
19
20
21def calculate_number_of_specials_to_place(objs):
22    """Return a list of number of specials to be placed at each obj"""
23    return [1 if random.random() > REPEAT_RATE[1] else
24            2 if random.random() > REPEAT_RATE[2] else
25            3 if random.random() > REPEAT_RATE[3] else 4
26            for _ in objs]
27
28
29def place_special(specials, obj):
30    """
31    Place at most a single special.
32    Return the number of specials placed.
33    """
34    # Calculate the conditional probabilities that each special is
35    # placed here given that a special will be placed here.
36    probs = [fo.special_spawn_rate(sp) for sp in specials]
37    total_prob = float(sum(probs))
38    if total_prob == 0:
39        # This shouldn't happen since special_spawn_rate > 0.0 is checked in distribute_specials()
40        return 0
41    thresholds = [x / total_prob for x in probs]
42
43    chance = random.random()
44    for threshold, special in zip(thresholds, specials):
45        if chance > threshold:
46            chance -= threshold
47            continue
48
49        fo.add_special(obj, special)
50        print("Special", special, "added to", fo.get_name(obj))
51        universe_statistics.specials_summary[special] += 1
52
53        return 1
54    return 0
55
56
57# TODO Bug:  distribute_specials forward checks that a special can be
58# placed, but it doesn't recursively check all previously placed
59# specials against the new special.
60def distribute_specials(specials_freq, universe_objects):
61    """
62    Adds start-of-game specials to universe objects.
63    """
64    # get basic chance for occurrence of specials from the universe tables
65    base_chance = universe_tables.SPECIALS_FREQUENCY[specials_freq]
66    if base_chance <= 0:
67        return
68
69    # get a list with all specials that have a spawn rate and limit both > 0 and a location condition defined
70    # (no location condition means a special shouldn't get added at game start)
71    specials = [sp for sp in fo.get_all_specials() if fo.special_spawn_rate(sp) > 0.0 and
72                fo.special_spawn_limit(sp) > 0 and fo.special_has_location(sp)]
73    if not specials:
74        return
75
76    # dump a list of all specials meeting that conditions and their properties to the log
77    print("Specials available for distribution at game start:")
78    for special in specials:
79        print("... {:30}: spawn rate {:2.3f} / spawn limit {}".
80              format(special, fo.special_spawn_rate(special), fo.special_spawn_limit(special)))
81
82    objects_needing_specials = [obj for obj in universe_objects if random.random() < base_chance]
83
84    track_num_placed = {obj: 0 for obj in universe_objects}
85
86    print("Base chance for specials is {}. Placing specials on {} of {} ({:1.4f})objects"
87          .format(base_chance, len(objects_needing_specials), len(universe_objects),
88                  float(len(objects_needing_specials)) / len(universe_objects)))
89
90    obj_tuple_needing_specials = set(zip(objects_needing_specials,
91                                         fo.objs_get_systems(objects_needing_specials),
92                                         calculate_number_of_specials_to_place(objects_needing_specials)))
93
94    # Equal to the largest distance in WithinStarlaneJumps conditions
95    # GALAXY_DECOUPLING_DISTANCE is used as follows.  For any two or more objects
96    # at least GALAXY_DECOUPLING_DISTANCE appart you only need to check
97    # fo.special_locations once and then you can place as many specials as possible,
98    # subject to number restrictions.
99    #
100    # Organize the objects into sets where all objects are spaced GALAXY_DECOUPLING_DISTANCE
101    # appart.  Place a special on each one.  Repeat until you run out of specials or objects.
102    GALAXY_DECOUPLING_DISTANCE = 6
103
104    while obj_tuple_needing_specials:
105        systems_needing_specials = defaultdict(set)
106        for (obj, system, specials_count) in obj_tuple_needing_specials:
107            systems_needing_specials[system].add((obj, system, specials_count))
108
109        print(" Placing in {} locations remaining.".format(len(systems_needing_specials)))
110
111        # Find a list of candidates all spaced GALAXY_DECOUPLING_DISTANCE apart
112        candidates = []
113        while systems_needing_specials:
114            random_sys = random.choice(list(systems_needing_specials.values()))
115            member = random.choice(list(random_sys))
116            obj, system, specials_count = member
117            candidates.append(obj)
118            obj_tuple_needing_specials.remove(member)
119            if specials_count > 1:
120                obj_tuple_needing_specials.add((obj, system, specials_count - 1))
121
122            # remove all neighbors from the local pool
123            for neighbor in fo.systems_within_jumps_unordered(GALAXY_DECOUPLING_DISTANCE, [system]):
124                if neighbor in systems_needing_specials:
125                    systems_needing_specials.pop(neighbor)
126
127        print("Caching specials_locations() at {} of {} remaining locations.".
128              format(str(len(candidates)), str(len(obj_tuple_needing_specials) + len(candidates))))
129        # Get the locations at which each special can be placed
130        locations_cache = {}
131        for special in specials:
132            # The fo.special_locations in the following line consumes most of the time in this
133            # function.  Decreasing GALAXY_DECOUPLING_DISTANCE will speed up the whole
134            # function by reducing the number of times this needs to be called.
135            locations_cache[special] = set(fo.special_locations(special, candidates))
136
137        # Attempt to apply a special to each candidate
138        # by finding a special that can be applied to it and hasn't been added too many times
139        for obj in candidates:
140
141            # check if the spawn limit for this special has already been reached (that is, if this special
142            # has already been added the maximal allowed number of times)
143            specials = [s for s in specials if universe_statistics.specials_summary[s] < fo.special_spawn_limit(s)]
144            if not specials:
145                break
146
147            # Find which specials can be placed at this one location
148            local_specials = [sp for sp in specials if obj in locations_cache[sp]]
149            if not local_specials:
150                universe_statistics.specials_repeat_dist[0] += 1
151                continue
152
153            # All prerequisites and the test have been met, now add this special to this universe object.
154            track_num_placed[obj] += place_special(local_specials, obj)
155
156    for num_placed in track_num_placed.values():
157        universe_statistics.specials_repeat_dist[num_placed] += 1
158