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()

[((243, 254), 'time.time', 'time.time', ([], {}), '()\n', (252, 254), False, 'import time\n'), ((410, 435), 'os.listdir', 'os.listdir', (['DOWNLOAD_PATH'], {}), '(DOWNLOAD_PATH)\n', (420, 435), False, 'import os\n'), ((839, 850), 'time.time', 'time.time', ([], {}), '()\n', (848, 850), False, 'import time\n'), ((1265, 1326), 'tqdm.contrib.concurrent.process_map', 'process_map', (['create_json_for_state', 'data_files'], {'max_workers': '(4)'}), '(create_json_for_state, data_files, max_workers=4)\n', (1276, 1326), False, 'from tqdm.contrib.concurrent import process_map\n'), ((1554, 1569), 'tqdm.tqdm', 'tqdm', (['state_ids'], {}), '(state_ids)\n', (1558, 1569), False, 'from tqdm import tqdm\n'), ((1978, 2012), 'aggregator.aggregate', 'aggregate', (['*args'], {'hide_output': '(True)'}), '(*args, hide_output=True)\n', (1987, 2012), False, 'from aggregator import aggregate\n'), ((457, 490), 'os.path.join', 'os.path.join', (['DOWNLOAD_PATH', 'file'], {}), '(DOWNLOAD_PATH, file)\n', (469, 490), False, 'import os\n'), ((503, 528), 'os.path.isfile', 'os.path.isfile', (['file_path'], {}), '(file_path)\n', (517, 528), False, 'import os\n'), ((1689, 1701), 'json.load', 'json.load', (['f'], {}), '(f)\n', (1698, 1701), False, 'import json\n'), ((1069, 1133), 'os.path.join', 'os.path.join', (['DOWNLOAD_PATH', 'f"""CenPop2020_Mean_BG{state_id}.txt"""'], {}), "(DOWNLOAD_PATH, f'CenPop2020_Mean_BG{state_id}.txt')\n", (1081, 1133), False, 'import os\n'), ((1151, 1238), 'os.path.join', 'os.path.join', (['DOWNLOAD_PATH', 'f"""tl_2020_{state_id}_bg"""', 'f"""tl_2020_{state_id}_bg.shp"""'], {}), "(DOWNLOAD_PATH, f'tl_2020_{state_id}_bg',\n    f'tl_2020_{state_id}_bg.shp')\n", (1163, 1238), False, 'import os\n'), ((1612, 1659), 'os.path.join', 'os.path.join', (['DOWNLOAD_PATH', 'f"""{state_id}.json"""'], {}), "(DOWNLOAD_PATH, f'{state_id}.json')\n", (1624, 1659), False, 'import os\n'), ((533, 556), 'pathlib.Path', 'pathlib.Path', (['file_path'], {}), '(file_path)\n', (545, 556), False, 'import pathlib\n'), ((1776, 1807), 'json.dumps', 'json.dumps', (["geojson['features']"], {}), "(geojson['features'])\n", (1786, 1807), False, 'import json\n')]