"""
Assign all 52 comunas of the Region Metropolitana to 20 groups,
balancing total population per group.

Uses LPT (Longest Processing Time) greedy algorithm:
sort comunas by population descending, then assign each to the
group with the smallest current total.
"""

import heapq
import random

# All 52 comunas of the Region Metropolitana (approx. 2017 Census / INE)
RM_COMUNAS = {
    # Provincia de Santiago (32)
    "Santiago": 404495,
    "Cerrillos": 80265,
    "Cerro Navia": 148312,
    "Conchali": 133256,
    "El Bosque": 162505,
    "Estacion Central": 147041,
    "Huechuraba": 98671,
    "Independencia": 100281,
    "La Cisterna": 90119,
    "La Florida": 366916,
    "La Granja": 116571,
    "La Pintana": 177335,
    "La Reina": 96762,
    "Las Condes": 294838,
    "Lo Barnechea": 105833,
    "Lo Espejo": 98804,
    "Lo Prado": 96498,
    "Macul": 116534,
    "Maipu": 521627,
    "Nunoa": 208237,
    "Pedro Aguirre Cerda": 101297,
    "Penalolen": 241599,
    "Providencia": 142079,
    "Pudahuel": 230293,
    "Quilicura": 210016,
    "Quinta Normal": 110026,
    "Recoleta": 157851,
    "Renca": 147151,
    "San Joaquin": 97625,
    "San Miguel": 107954,
    "San Ramon": 94906,
    "Vitacura": 85384,
    # Provincia de Cordillera (3)
    "Puente Alto": 568106,
    "San Jose de Maipo": 18225,
    "Pirque": 26820,
    # Provincia de Chacabuco (3)
    "Colina": 146207,
    "Lampa": 106559,
    "Tiltil": 17647,
    # Provincia de Maipo (4)
    "San Bernardo": 301313,
    "Buin": 81857,
    "Calera de Tango": 26476,
    "Paine": 82539,
    # Provincia de Melipilla (5)
    "Melipilla": 107882,
    "Alhue": 5765,
    "Curacavi": 30885,
    "Maria Pinto": 12572,
    "San Pedro": 9522,
    # Provincia de Talagante (5)
    "Talagante": 72704,
    "El Monte": 34842,
    "Isla de Maipo": 36174,
    "Padre Hurtado": 62025,
    "Penaflor": 90382,
}

NUM_GROUPS = 1


def main(seed=None):
    rng = random.Random(seed)

    # Sort comunas by population descending (LPT)
    sorted_comunas = sorted(RM_COMUNAS.items(), key=lambda x: x[1], reverse=True)

    # Min-heap: (current_total, group_index, [comunas])
    groups = [(0, i, []) for i in range(NUM_GROUPS)]
    heapq.heapify(groups)

    for name, pop in sorted_comunas:
        total, idx, members = heapq.heappop(groups)
        members.append((name, pop))
        heapq.heappush(groups, (total + pop, idx, members))

    # Collect and shuffle group assignment
    results = [(idx, members, total) for total, idx, members in groups]
    rng.shuffle(results)

    # Print
    print(f"{'Grupo':<7} {'Comunas':<65} {'Total':<10}")
    print("-" * 85)

    totals = []
    for g, (_, members, total) in enumerate(results, 1):
        totals.append(total)
        comunas_str = " + ".join(f"{n} ({p:,})" for n, p in members)
        print(f"{g:<7} {comunas_str:<65} {total:,}")

    print("-" * 85)
    print(f"Min: {min(totals):,}  |  Max: {max(totals):,}  |  "
          f"Spread: {max(totals) - min(totals):,}  |  Avg: {sum(totals) // len(totals):,}")


if __name__ == "__main__":
    main(756)