import argparse
import math

import numpy as np

from utils.numpy.save import save


def run(
        height: int,
        width: int,
        channels: int,
        num_samples: int,
        save_folder: str,
        batch_size: int = None,
        start_index: int = 0,
        n_processes: int = 8,
        seed_entropy: int = None
) -> None:
    assert isinstance(height, int), f'height must be an int: {height}'
    assert height >= 0, f'height must be greater than or equal to 0: {height}'
    assert isinstance(width, int), f'width must be an int: {width}'
    assert width >= 0, f'width must be greater than or equal to 0: {width}'
    assert isinstance(channels, int), f'channels must be an int: {channels}'
    assert channels >= 0, f'channels must be greater than or equal to 0: {channels}'
    assert isinstance(num_samples, int), f'num_samples must be an int: {num_samples}'
    assert num_samples >= 0, f'num_samples must be greater than or equal to 0: {num_samples}'
    assert isinstance(save_folder, str), f'save_folder must be a str: {save_folder}'
    assert batch_size is None or (isinstance(batch_size, int) and batch_size > 0), \
        f'batch_size must be None or an int greater than 0: {batch_size}'
    assert isinstance(start_index, int), f'start_index must be an int: {start_index}'
    assert start_index >= 0, f'start_index must be greater than or equal to 0: {start_index}'
    assert isinstance(n_processes, int), f'n_processes must be an int: {n_processes}'
    assert n_processes > 0, f'n_processes must be greater than 0: {n_processes}'
    assert seed_entropy is None or isinstance(seed_entropy, int), f'seed_entropy must be None or an int: {seed_entropy}'
    print(
        f'scripts.create_noises.run - '
        f'height: {height}, '
        f'width: {width}, '
        f'channels: {channels}, '
        f'num_samples: {num_samples}, '
        f'save_folder: {save_folder}, '
        f'batch_size: {batch_size}, '
        f'start_index: {start_index}, '
        f'n_processes: {n_processes}, '
        f'seed_entropy: {seed_entropy}'
    )

    seed_sequence: np.random.SeedSequence = np.random.SeedSequence(seed_entropy)
    generator: np.random.Generator = np.random.default_rng(seed=seed_sequence)

    batch_size: int = batch_size if batch_size is not None else num_samples
    num_batches: int = math.ceil(num_samples / batch_size)

    for index in range(num_batches):
        print(f'batch: {index + 1}/{num_batches}')
        start: int = index * batch_size
        end: int = min((index + 1) * batch_size, num_samples)
        noises: np.ndarray = generator.standard_normal((end - start, channels, height, width), float)
        save(
            folder=save_folder,
            data=noises,
            start_index=start_index + start,
            n_processes=n_processes
        )
    print('done')


def run_from_config(config: dict[str, any]) -> None:
    run(**config)


def parse_args() -> argparse.Namespace:
    parser: argparse.ArgumentParser = argparse.ArgumentParser()
    parser.add_argument('--height', type=int, required=True)
    parser.add_argument('--width', type=int, required=True)
    parser.add_argument('--channels', type=int, required=True)
    parser.add_argument('--num_samples', type=int, required=True)
    parser.add_argument('--save_folder', type=str, required=True)
    parser.add_argument('--batch_size', type=int, default=None)
    parser.add_argument('--start_index', type=int, default=0)
    parser.add_argument('--n_processes', type=int, default=8)
    parser.add_argument('--seed_entropy', type=int, default=None)
    return parser.parse_args()


def get_config_from_args(args: argparse.Namespace) -> dict[str, any]:
    return {
        'height': args.height,
        'width': args.width,
        'channels': args.channels,
        'num_samples': args.num_samples,
        'save_folder': args.save_folder,
        'batch_size': args.batch_size,
        'start_index': args.start_index,
        'n_processes': args.n_processes,
        'seed_entropy': args.seed_entropy
    }


def main() -> None:
    args: argparse.Namespace = parse_args()
    config: dict[str, any] = get_config_from_args(args)
    run_from_config(config)


if __name__ == '__main__':
    main()
