from typing import IO, Union

import click


@click.group("testing")
def cli():
    """Manipulate test data."""
    pass


FILENAMES = [
    "advance/advance_compare_20171018-134445.csv.zip",
    "iea/372f7e29-en.zip",
    "iea/8624f431-en.zip",
    "iea/cac5fa90-en.zip",
    "ssp/SSP-Review-Phase-1.csv.gz",
    "ssp/SspDb_country_data_2013-06-12.csv.zip",
]


@cli.command("fuzz-private-data")
@click.option("--frac", type=float, default=1.0, help="Fraction of rows (default: all)")
@click.argument("filename", metavar="FILENAME", type=click.Choice(FILENAMES))
def fuzz_private_data(filename, frac: float):  # pragma: no cover
    """Create random data for testing.

    This command creates data files in message_ix_models/data/test/… based on
    corresponding private files in message_data/data/…. This supports testing of code in
    message_ix_models that handles these files.

    The files are identical in structure and layout, except the values are "fuzzed", or
    replaced with random values.

    To see valid FILENAMES, run the command with no arguments.
    """
    import zipfile
    from pathlib import Path
    from tempfile import TemporaryDirectory

    import dask.dataframe as dd
    import pandas as pd
    from numpy import char, random

    from message_ix_models.project.advance.data import NAME
    from message_ix_models.util import package_data_path, private_data_path

    # Paths
    p = Path(filename)
    path_in = private_data_path(p)
    path_out = package_data_path("test", p)

    # Read the data
    with TemporaryDirectory() as td:
        td_path = Path(td)
        if "advance" in filename:
            # Manually unpack one member of the multi-member archive `path_in`
            target: Union[IO, Path, str] = zipfile.ZipFile(path_in).extract(
                NAME, path=td_path
            )
        elif "iea" in filename:
            # Manually unpack so that dask.dataframe.read_csv() can be used
            from message_ix_models.tools.iea.web import unpack_zip

            target = unpack_zip(path_in)
        else:
            target = path_in

        # - Read the data using dask & pyarrow.
        # - Subset the data if `frac` < 1.0.
        # - Compute the resulting pandas.DataFrame
        df = dd.read_csv(target, engine="pyarrow").sample(frac=frac).compute()

    # Determine columns in which to replace numerical data
    if "iea" in filename:
        # Specific column
        cols = ["Value"]
    else:
        # All columns with numeric names, for instance 2000, 2001, etc.
        cols = list(filter(char.isnumeric, df.columns))

    # Shape of random data
    size = (df.shape[0], len(cols))
    # - Generate random data of this shape.
    # - Keep only the elements corresponding to non-NA elements of `df`.
    # - Update `df` with these values.*
    generator = random.default_rng()
    df.update(df.where(df.isna(), pd.DataFrame(generator.random(size), columns=cols)))

    # Write to file, keeping only a few decimal points
    path_out.parent.mkdir(parents=True, exist_ok=True)

    if "advance" in filename:
        zf = zipfile.ZipFile(path_out, "w", compression=zipfile.ZIP_BZIP2)
        target = zf.open(NAME)
    else:
        target = path_out

    df.to_csv(target, index=False, float_format="%.2f")
