import json
import os
import pathlib
import time

from tqdm import tqdm
from aggregator import aggregate
from download import DOWNLOAD_PATH, download_files, unzip_files
from tqdm.contrib.concurrent import process_map


def main():
    start = time.time()
    # print("Downloading files...")
    # download_files()

    # print("Unzipping shapefiles...")
    # unzip_files()

    state_ids = []
    for file in os.listdir(DOWNLOAD_PATH):
        file_path = os.path.join(DOWNLOAD_PATH, file)

        if os.path.isfile(file_path) and pathlib.Path(file_path).suffix == ".txt":
            state_ids.append(file[file.index("BG") + 2 : file.index(".")])

    # print("Computing population JSON heatmaps...")
    # compute_json_heatmaps(state_ids)

    print("Aggregating JSON files into one...")
    aggegrate_json_files(state_ids)

    end = time.time()
    print(f"Done in {(end - start):0.2f}s")


def compute_json_heatmaps(state_ids):
    data_files = []
    for state_id in state_ids:
        data_files.append(
            (
                state_id,
                os.path.join(DOWNLOAD_PATH, f"CenPop2020_Mean_BG{state_id}.txt"),
                os.path.join(DOWNLOAD_PATH, f"tl_2020_{state_id}_bg", f"tl_2020_{state_id}_bg.shp"),
            )
        )

    process_map(create_json_for_state, data_files, max_workers=4)


def aggegrate_json_files(state_ids):
    with open("public/data/pop.json", "w") as f:
        f.write("""{"type": "FeatureCollection", "features": [""")

    # state_ids = state_ids[:2]

    features = []
    for state_id in tqdm(state_ids):
        geojson = None
        with open(os.path.join(DOWNLOAD_PATH, f"{state_id}.json")) as f:
            geojson = json.load(f)

        with open("public/data/pop.json", "a") as f:
            f.write(json.dumps(geojson["features"])[1:-1] + ("," if state_id != state_ids[-1] else ""))

    with open("public/data/pop.json", "a") as f:
        f.write("]}")


def create_json_for_state(args):
    return aggregate(*args, hide_output=True)


if __name__ == "__main__":
    main()

[((13, 12, 13, 23), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((21, 16, 21, 41), 'os.listdir', 'os.listdir', ({(21, 27, 21, 40): 'DOWNLOAD_PATH'}, {}), '(DOWNLOAD_PATH)', False, 'import os\n'), ((33, 10, 33, 21), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((48, 4, 48, 65), 'tqdm.contrib.concurrent.process_map', 'process_map', (), '', False, 'from tqdm.contrib.concurrent import process_map\n'), ((58, 20, 58, 35), 'tqdm.tqdm', 'tqdm', ({(58, 25, 58, 34): 'state_ids'}, {}), '(state_ids)', False, 'from tqdm import tqdm\n'), ((71, 11, 71, 45), 'aggregator.aggregate', 'aggregate', (), '', False, 'from aggregator import aggregate\n'), ((22, 20, 22, 53), 'os.path.join', 'os.path.join', ({(22, 33, 22, 46): 'DOWNLOAD_PATH', (22, 48, 22, 52): 'file'}, {}), '(DOWNLOAD_PATH, file)', False, 'import os\n'), ((24, 11, 24, 36), 'os.path.isfile', 'os.path.isfile', ({(24, 26, 24, 35): 'file_path'}, {}), '(file_path)', False, 'import os\n'), ((61, 22, 61, 34), 'json.load', 'json.load', ({(61, 32, 61, 33): 'f'}, {}), '(f)', False, 'import json\n'), ((43, 16, 43, 80), 'os.path.join', 'os.path.join', ({(43, 29, 43, 42): 'DOWNLOAD_PATH', (43, 44, 43, 79): 'f"""CenPop2020_Mean_BG{state_id}.txt"""'}, {}), "(DOWNLOAD_PATH, f'CenPop2020_Mean_BG{state_id}.txt')", False, 'import os\n'), ((44, 16, 44, 99), 'os.path.join', 'os.path.join', ({(44, 29, 44, 42): 'DOWNLOAD_PATH', (44, 44, 44, 68): 'f"""tl_2020_{state_id}_bg"""', (44, 70, 44, 98): 'f"""tl_2020_{state_id}_bg.shp"""'}, {}), "(DOWNLOAD_PATH, f'tl_2020_{state_id}_bg',\n    f'tl_2020_{state_id}_bg.shp')", False, 'import os\n'), ((60, 18, 60, 65), 'os.path.join', 'os.path.join', ({(60, 31, 60, 44): 'DOWNLOAD_PATH', (60, 46, 60, 64): 'f"""{state_id}.json"""'}, {}), "(DOWNLOAD_PATH, f'{state_id}.json')", False, 'import os\n'), ((24, 41, 24, 64), 'pathlib.Path', 'pathlib.Path', ({(24, 54, 24, 63): 'file_path'}, {}), '(file_path)', False, 'import pathlib\n'), ((64, 20, 64, 51), 'json.dumps', 'json.dumps', ({(64, 31, 64, 50): "geojson['features']"}, {}), "(geojson['features'])", False, 'import json\n')]