#!/usr/bin/env python3

# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

import argparse
import glob
import os

import numpy as np
import rasterio as rio
from tqdm import tqdm
from tqdm.contrib.concurrent import thread_map

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("roots", nargs="+", help="directories to search for files")
    parser.add_argument("--suffix", default=".tif", help="file suffix")
    parser.add_argument("--sort", action="store_true", help="sort by class frequency")
    parser.add_argument(
        "--weights", action="store_true", help="print weights instead of ratios"
    )
    parser.add_argument(
        "--total-classes", type=int, default=256, help="total number of classes"
    )
    parser.add_argument(
        "--keep-classes",
        type=float,
        default=1,
        help="keep classes with percentage higher than this",
    )
    parser.add_argument(
        "--ignore-index", type=int, default=0, help="fill value to ignore"
    )
    parser.add_argument("--num-workers", type=int, default=10, help="number of threads")
    args = parser.parse_args()

    def class_counts(path: str) -> "np.typing.NDArray[np.float64]":
        """Calculate the number of values in each class.

        Args:
            path: Path to an image file.

        Returns:
            Counts of each class.
        """
        global args

        counts = np.zeros(args.total_classes)
        with rio.open(path, "r") as src:
            x = src.read()
            unique, unique_counts = np.unique(x, return_counts=True)
            counts[unique] = unique_counts

        return counts

    paths = []
    for root in args.roots:
        paths.extend(
            glob.glob(os.path.join(root, "**", f"*{args.suffix}"), recursive=True)
        )

    if args.num_workers > 0:
        counts = thread_map(class_counts, paths, max_workers=args.num_workers)
    else:
        counts = []
        for path in tqdm(paths):
            counts.append(class_counts(path))

    counts = np.sum(counts, axis=0)

    if 0 <= args.ignore_index < args.total_classes:
        counts[args.ignore_index] = 0

    if args.sort:
        indices = np.argsort(counts)
        indices = indices[::-1]
        counts = counts[indices]
    else:
        indices = np.arange(args.total_classes)

    counts /= np.sum(counts)
    keep = counts > args.keep_classes / 100
    indices = indices[keep]
    counts = counts[keep]

    if args.weights:
        counts = 1 / counts
        counts /= np.sum(counts)

    print(indices)
    print(counts)
