"""Count the chemical-substance codes under ATC J01 (antibacterials for systemic use).

Crawls the public ATC/DDD Index of the WHO Collaborating Centre for Drug
Statistics Methodology (https://atcddd.fhi.no/atc_ddd_index/) from J01 down to
the fifth level. A fifth-level code is a combination when its name contains
"combination" or joins substances with " and ". Prints totals as JSON.

Usage: python count_atc_j01.py
"""

import json
import re
import time
import urllib.request

BASE = "https://atcddd.fhi.no/atc_ddd_index/?code={}&showdescription=no"


def children(code):
    req = urllib.request.Request(BASE.format(code), headers={"User-Agent": "Mozilla/5.0"})
    page = urllib.request.urlopen(req, timeout=60).read().decode("utf-8", "replace")
    time.sleep(0.5)
    found = re.findall(r'code=(J01[A-Z0-9]+)&amp;showdescription=(?:no|yes)">([^<]+)<', page)
    found += re.findall(r'code=(J01[A-Z0-9]+)&showdescription=(?:no|yes)">([^<]+)<', page)
    return {c: n.strip() for c, n in found if len(c) > len(code) and c.startswith(code)}


def main():
    substances = {}
    for group in children("J01"):
        for subgroup in children(group):
            if len(subgroup) == 5:
                for code, name in children(subgroup).items():
                    if len(code) == 7:
                        substances[code] = name
    combos = {c: n for c, n in substances.items() if "combination" in n.lower() or " and " in n}
    single = sorted({n.lower() for c, n in substances.items() if c not in combos})
    print(json.dumps({
        "retrieved": time.strftime("%Y-%m-%d"),
        "fifth_level_codes": len(substances),
        "combination_codes": len(combos),
        "single_substance_codes": len(substances) - len(combos),
        "distinct_single_substance_names": len(single),
    }, indent=2))


if __name__ == "__main__":
    main()