shorten-schools

#!/usr/bin/env python3
"""
Zkrátí v databázi oficiální dlouhá jména škol na něco čitelnějšího, uloží
do sloupce places.name.

Algoritmus se jména snaží dostat do podoby ZKRÁCENÉ_JMÉNO, kde
ZKRÁCENÉ_JMÉNO = NÁZEV MÍSTO
NÁZEV = např. "SŠ", "ZŠ T. G. Masaryka", "SPŠ strojnická a SOŠ profesora Švejcara"
MÍSTO = MĚSTO [ULICE [Č.P.]]
  např. "Slatinice", "Praha 7", "Olomouc, Svatoplukova"

Může existovat víc možností zkrácení, např.
    ZŠ a MŠ Olomouc, Svatoplukova 11
    ZŠ a MŠ Olomouc, Svatoplukova
    ZŠ a MŠ Olomouc
Algoritmus vytvoří všechny varianty jmen a pak kontroluje, jestli při použití
nejkratší varianty ("ZŠ a MŠ Olomouc") nenastane konflikt jmen s jinou školou.
Pokud ano, zkusí použít pro obě školy delší variantu názvu. Toto se opakuje,
dokud se konflikty nevyřeší.
"""

import copy
import random
import re
import sys
import argparse

from sqlalchemy.orm import aliased

import mo.db as db


def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


def sorted_by_length(schools):
    schools2 = copy.copy(schools)
    schools2.sort(key=lambda sc: len(sc["names"][-1]))
    return schools2


def summarize(schools, k=5):
    lens = [len(sc["names"][-1]) for sc in schools]
    avg_len = sum(lens) / len(schools)
    eprint("Average length:", avg_len)
    eprint("Maximum length:", max(lens))

    names_by_lens = sorted_by_length(schools)

    eprint()
    eprint(f"{k} longest:")
    for sc in names_by_lens[::-1][:k]:
        eprint(f'{sc["names"][-1]} (@{sc["city"]})')

    random.shuffle(names_by_lens)
    eprint()
    eprint(f"{k} random:")
    for sc in names_by_lens[:k]:
        eprint(f'Old: {sc["names"][0]}')
        eprint(f'{sc["names"][-1]}')
        eprint()


def remove_formalities(name):
    for formality in formalities:
        name = re.sub(formality, "", name, flags=re.IGNORECASE)

    return name


def shorten_name(name):
    for re_from, re_to in school_kinds:
        name = re.sub(re_from, re_to, name, flags=re.IGNORECASE)

    return name


def partition(name, city):
    """Rozdělí název školy na část před názvem města a část po názvu města"""

    # Zkouší drobné úpravy názvu města
    for rule in city_rules:
        # Pro slova jako "Táborské" chceme odstranit i zbytek slova, nejen "Tábor"
        pat = r"\b{}\w*\b".format(city)

        if re.search(pat, name) is not None:
            parts = re.split(pat, name)

            if len(parts) != 2:
                # Název města se vyskytuje víckrát, není jasné, co dělat
                return None
            else:
                ok = True
                for kind, _ in school_kinds:
                    if kind.lower() in parts[1].lower():
                        ok = False

                if not ok:
                    # Názvová část školy pokračuje i po názvu města (např. "Táborské gymnázium"),
                    # nelze automaticky vyřešit
                    return None
                else:
                    return parts

        if rule is not None:
            city = re.sub(rule[0], rule[1], city)

    # Nenašli jsme název města
    return [name]


def remove_house_number(name):
    name, n = re.subn(r"(, ([^\W\d_]| |\.)+) [0-9/]+[a-z]?$", r"\1", name)
    # True, pokud se název změnil
    return name, n > 0


def should_have_comma_after_name(p_name):
    # Čárku chceme v případech jako
    # "Základní škola generála Zdeňka Škarvady, Ostrava-Poruba"
    # ale ne pro
    # "Základní škola Dolní Ředice, okres Pardubice"

    for sk in school_kinds:
        if p_name.endswith(sk):
            return False

    return True


def postprocess_name_part(p_name):
    # Vyřeší okrajové případy části názvu před městem

    p_name = p_name.strip(" ,-")
    p_name = re.sub(
        " v$", "", p_name
    )  # Pro případy jako "G v Kroměříži" -> "G v, Kroměříž"

    if should_have_comma_after_name(p_name):
        p_name += ","

    return p_name


def shorten_all(schools):
    for sc in schools:
        sc["names"].append(remove_formalities(sc["names"][-1]))
        sc["parts"] = partition(sc["names"][-1], sc["city"])

    eprint("Total schools: {}".format(len(schools)))

    n_split = 0

    for sc in schools:
        sc["names"].append(shorten_name(sc["names"][-1]))
        if sc["parts"] is not None:
            if len(sc["parts"]) == 1:
                # Název města nenalezen v názvu školy
                p_name = postprocess_name_part(sc["names"][-1])
                sc["names"].append(f"{p_name} {sc['city']}")
            else:
                # Když máme rozdělení, můžeme zkusit odstanit číslo popisné
                # a případně i celý název ulice
                n_split += 1
                assert len(sc["parts"]) == 2

                p_name, p_place = sc["parts"]

                p_name = shorten_name(p_name)
                p_name = postprocess_name_part(p_name)

                p_place2, changed = remove_house_number(p_place)

                if changed:
                    sc["names"].append(
                        f"{p_name} {sc['city']}, {p_place2.strip(' ,-')}"
                    )

                if "Praha" not in sc["city"]:  # např. "G Praha 2" nechceme
                    sc["names"].append(f"{p_name} {sc['city']}")

    eprint(f"Successfully split up {n_split} schools")

    return schools


def is_conflict(names1, names2):
    return any([(name in names1) for name in names2])


def remove_conflicts(shortened):
    """Vrátí se k delším variantám jmen, pokud se vyskytly konflikty"""
    again = True

    while again:
        shortened.sort(key=lambda sc: sc["names"][-1])
        eprint("----------------------------")
        n_conflicts = 0
        again = False

        bad_names = set()

        for sc1, sc2 in zip(shortened, shortened[1:]):
            if is_conflict(sc1["names"], sc2["names"]):
                n_conflicts += 1

                if sc1["names"][0] != sc2["names"][0]:
                    bad_names.add(sc1["names"][-1])
                    again = True

        for sc in shortened:
            if sc["names"][-1] in bad_names:
                assert len(sc["names"]) > 1
                sc["names"].pop()

        eprint(f"Found {n_conflicts} conflicts")

    # Hack - tato zkrácení vždy chceme aplikovat, předpokládáme, že nevzniknou konflikty
    for sc in shortened:
        sc["names"].append(remove_formalities(shorten_name(sc["names"][-1])))

    eprint("Done (possible unremovable conflicts)")


city_rules = [
    (r"(\w)-(\w)", r"\1 - \2"),  # Mezery kolem pomlček jsou někdy nekonzistentní
    ("Praha", "v Praze"),
    ("v Praze 4", "v Praze 12"),
    (r"v Praze [0-9]+", "v Praze"),
    ("v Praze", "Praha"),
    None,  # Dummy
]

school_kinds = [
    ("Gymnázium", "G"),
    ("Vyšší odborná škola", "VOŠ"),
    ("Střední odborná škola", "SOŠ"),
    ("Střední zdravotnická škola", "SZŠ"),
    ("Střední průmyslová škola", "SPŠ"),
    ("Střední pedagogická škola", "SPŠ"),
    ("Střední odborné učiliště", "SOU"),
    ("Střední škola", "SŠ"),
    ("Základní škola", "ZŠ"),
    ("Základní umělecká škola", "ZUŠ"),
    ("Mateřská škola", "MŠ"),
    # Nechceme mít zvlášť "ZŠ Nový Rychnov" a "ZŠ a MŠ Nový Rychnov" odlišené jen "MŠ"
    ("ZŠ a MŠ", "ZŠ"),
    ("MŠ a ZŠ", "ZŠ"),
]

formalities = [
    r",?-? ?příspěvková organizace",
    r",? s\.r\.o\.",
    r",? o\.p\.s\.",
    r" s právem státní jazykové zkoušky",
    r",? ?okres .+$",
]


def main():
    parser = argparse.ArgumentParser(
        description="Automaticky zkrátí jména škol v databázi"
    )
    parser.add_argument(
        "-n",
        "--dry-run",
        action="store_true",
        help="Jen uloží vygenerovaná zkrácení do 'prejmenovani.tsv', nemění databázi",
    )
    parser.add_argument(
        "--restore", action="store_true", help="Vrátí se k oficiálním názvům"
    )
    args = parser.parse_args()

    session = db.get_session()

    school_place_t = aliased(db.Place)
    parent_place_t = aliased(db.Place)

    schools_q = (
        session.query(db.School, school_place_t, parent_place_t)
        .filter(db.School.place_id == school_place_t.place_id)
        .filter(parent_place_t.place_id == school_place_t.parent)
        .all()
    )

    if args.restore:
        eprint("Vracím se k původním názvům.")
        for school, place, parent_place in schools_q:
            place.name = school.official_name

        session.commit()
        return

    schools = []

    for school, place, parent_place in schools_q:

        # Parent má být škola
        assert parent_place.level == 3

        # Toto platí před prvním spuštením skriptu, pak už ne (změníme place.name)
        # assert place.name == school.official_name

        schools.append(
            {
                "place_id": school.place_id,
                "names": [school.official_name],
                "city": parent_place.name,
                "db_place": place,
            }
        )

    shortened = shorten_all(schools)
    remove_conflicts(shortened)

    summarize(shortened, k=10)

    if args.dry_run:
        filename = "prejmenovani.tsv"
        with open(filename, "w") as f:
            shortened.sort(key=lambda sc: sc["names"][0])

            for sc in shortened:
                f.write(f"{sc['names'][-1]}\t{sc['names'][0]}\n")

        print(f"Seznam všech přejmenování uložen do {filename}.")
        return

    # Zapsat do DB
    for sc in shortened:
        sc["db_place"].name = sc["names"][-1]
    session.commit()


if __name__ == "__main__":
    main()