from pathlib import Path
import json


def calculate_dataset_size(pregenerated_data_path: Path) -> list:
    """
    Calculates size of each training epochs, and number of distinct epoch datasets.
    :param pregenerated_data_path: Path to pregenerated dataset.
    :return: list of numbers of samples per epoch
    """
    samples_per_epoch = []
    if (pregenerated_data_path / "metrics.json").is_file():
        metrics = json.loads((pregenerated_data_path / "metrics.json").read_text())
        return [metrics["num_training_examples"]]
    for i in range(1000):
        epoch_path = pregenerated_data_path / f"epoch_{i}"
        metrics_file = epoch_path / "metrics.json"
        if epoch_path.is_dir() and metrics_file.is_file():
            metrics = json.loads(metrics_file.read_text())
            samples_per_epoch.append(metrics["num_training_examples"])
        else:
            if i == 0:
                raise FileNotFoundError("No training data was found!")
            break
    return samples_per_epoch