MapComplete/scripts/generatePmTilesExtracts.ts

import Script from "./Script"
import { existsSync, mkdirSync, writeFileSync } from "fs"
import { Utils } from "../src/Utils"
import { OfflineBasemapManager } from "../src/Logic/OfflineBasemapManager"
import { PmTilesExtractGenerator } from "./pmTilesExtractGenerator"

class GeneratePmTilesExtracts extends Script {
    private targetDir: string
    private skipped: number = 0

    private extractsGenerator: PmTilesExtractGenerator

    constructor() {
        super(
            "Generates many pmtiles-archive from planet-latest.pmtiles. Expects the `pmtiles`-executable to be at `/data/pmtiles`." +
                "\n\n" +
                "The first argument should be the 'targetDirectory', where many subdirectories will be made containing the subarchives. A second argument (defaulting to [targetDir]/planet-latest.pmtiles) is the source data"
        )
    }

    public *generateColumnIfNeeded(
        z: number,
        x: number,
        boundary: number,
        maxzoom?: number
    ): Generator<Promise<any>> {
        const lastFileForColumn = this.extractsGenerator.getFilename(z, x, boundary - 1)
        if (existsSync(lastFileForColumn)) {
            // Skip this column, already exists
            console.log("Skipping column ", x, "at zoom", z)
            this.skipped += boundary
            return
        }
        console.log("Starting column", x, "at zoom", z, "as", lastFileForColumn, "does not exist")

        for (let y = 0; y < boundary; y++) {
            yield this.extractsGenerator.generateArchive(z, x, y, maxzoom)
        }
    }

    private *generateField(z: number, maxzoom?: number): Generator<Promise<void>> {
        const boundary = Math.max(1, 2 << (z - 1))
        for (let x = 0; x < boundary; x++) {
            if (!existsSync(this.targetDir + "/" + z + "/" + x)) {
                mkdirSync(this.targetDir + "/" + z + "/" + x)
            }
            for (const promise of this.generateColumnIfNeeded(z, x, boundary, maxzoom)) {
                yield promise
            }
        }
    }

    private *generateAll(): Generator<Promise<void>> {
        const zoomlevels: Record<number, number> = OfflineBasemapManager.zoomelevels
        for (const key in zoomlevels) {
            const minzoom: number = Number(key)
            const maxzoom: number | undefined = zoomlevels[key]
            if (!existsSync(this.targetDir + "/" + key)) {
                mkdirSync(this.targetDir + "/" + key)
            }
            for (const promise of this.generateField(minzoom, maxzoom)) {
                yield promise
            }
        }
    }

    createBatch<T>(generator: Generator<T>, length: number): T[] {
        const batch = []
        do {
            const next = generator.next()
            if (next.done) {
                return batch
            }
            batch.push(next.value)
        } while (batch.length < length)
        return batch
    }

    async main(args: string[]): Promise<void> {
        this.targetDir = args[0]
        const sourceFile = this.targetDir + "/planet-latest.pmtiles"
        if (!this.targetDir) {
            console.log("Please specify a target directory. Did you forget '--' in vite-node?")
            return
        }
        this.extractsGenerator = new PmTilesExtractGenerator(sourceFile, this.targetDir)

        let estimate = 0
        for (const key in OfflineBasemapManager.zoomelevels) {
            const z: number = Number(key)
            const boundary = 2 << z
            estimate += boundary * boundary
        }
        console.log("Target dir is:", this.targetDir)
        const numberOfThreads = 24
        const generator = this.generateAll()
        let batch: Promise<void>[] = []
        let done = 0
        const startDate = new Date()
        do {
            batch = this.createBatch(generator, numberOfThreads)
            await Promise.all(batch)
            done += batch.length
            const now = new Date()
            const timeElapsed = (now.getTime() - startDate.getTime()) / 1000
            const speed = ("" + done / timeElapsed).substring(0, 5)
            const perc = ("" + (100 * (done + this.skipped)) / estimate).substring(0, 5)
            const etaSecond = Math.floor(((estimate - done - this.skipped) * timeElapsed) / done)
            console.log(
                "Completed",
                numberOfThreads,
                `processes;  ${
                    done + this.skipped
                } / ${estimate}, ${perc}%, ${speed} tile/second, ETA: ${Utils.toHumanTime(
                    etaSecond
                )}`
            )
        } while (batch.length > 0)
        writeFileSync(
            this.targetDir + "/Last_pm_tile_extracts.txt",
            [
                new Date().getTime() + "",
                new Date().toISOString(),
                "# The script converting the planet-latest.pmtiles into sub-archives has been successfully executed at the stated time",
            ].join("\n"),
            "utf-8"
        )
    }
}

new GeneratePmTilesExtracts().run()