"""Utility functions for handling parquet metadata files"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd
import pyarrow.dataset as pds
from upath import UPath

from hats.io import file_io, paths
from hats.io.file_io.file_pointer import get_upath
from hats.pixel_math.healpix_pixel import HealpixPixel
from hats.pixel_math.healpix_pixel_function import get_pixel_argsort


def write_parquet_metadata(
    catalog_path: str | Path | UPath,
    order_by_healpix=True,
    output_path: str | Path | UPath | None = None,
):
    """Generate parquet metadata, using the already-partitioned parquet files
    for this catalog.

    For more information on the general parquet metadata files, and why we write them, see
    https://arrow.apache.org/docs/python/parquet.html#writing-metadata-and-common-metadata-files

    Args:
        catalog_path (str): base path for the catalog
        order_by_healpix (bool): use False if the dataset is not to be reordered by
            breadth-first healpix pixel (e.g. secondary indexes)
        output_path (str): base path for writing out metadata files
            defaults to `catalog_path` if unspecified

    Returns:
        sum of the number of rows in the dataset.
    """
    ignore_prefixes = [
        "intermediate",
        "_common_metadata",
        "_metadata",
    ]

    catalog_path = get_upath(catalog_path)
    dataset_subdir = catalog_path / "dataset"
    (dataset_path, dataset) = file_io.read_parquet_dataset(
        dataset_subdir,
        ignore_prefixes=ignore_prefixes,
        exclude_invalid_files=True,
    )
    metadata_collector = []
    # Collect the healpix pixels so we can sort before writing.
    healpix_pixels = []
    total_rows = 0

    for single_file in dataset.files:
        relative_path = single_file[len(dataset_path) + 1 :]
        single_metadata = file_io.read_parquet_metadata(dataset_subdir / relative_path)

        # Users must set the file path of each chunk before combining the metadata.
        single_metadata.set_file_path(relative_path)

        if order_by_healpix:
            healpix_pixel = paths.get_healpix_from_path(relative_path)

            healpix_pixels.append(healpix_pixel)
        metadata_collector.append(single_metadata)
        total_rows += single_metadata.num_rows

    ## Write out the two metadata files
    if output_path is None:
        output_path = catalog_path
    if order_by_healpix:
        argsort = get_pixel_argsort(healpix_pixels)
        metadata_collector = np.array(metadata_collector)[argsort]
    catalog_base_dir = get_upath(output_path)
    file_io.make_directory(catalog_base_dir / "dataset", exist_ok=True)
    metadata_file_pointer = paths.get_parquet_metadata_pointer(catalog_base_dir)
    common_metadata_file_pointer = paths.get_common_metadata_pointer(catalog_base_dir)

    file_io.write_parquet_metadata(
        dataset.schema,
        metadata_file_pointer,
        metadata_collector=metadata_collector,
        write_statistics=True,
    )
    file_io.write_parquet_metadata(dataset.schema, common_metadata_file_pointer)
    return total_rows


def read_row_group_fragments(metadata_file: str):
    """Generator for metadata fragment row groups in a parquet metadata file.

    Args:
        metadata_file (str): path to `_metadata` file.
    """
    metadata_file = get_upath(metadata_file)
    if not file_io.is_regular_file(metadata_file):
        metadata_file = paths.get_parquet_metadata_pointer(metadata_file)

    dataset = pds.parquet_dataset(metadata_file, filesystem=metadata_file.fs)

    for frag in dataset.get_fragments():
        yield from frag.row_groups


def _nonemin(value1, value2):
    """Similar to numpy's nanmin, but excludes `None` values.

    NB: If both values are `None`, this will still return `None`"""
    if value1 is None:
        return value2
    if value2 is None:
        return value1
    return min(value1, value2)


def _nonemax(value1, value2):
    """Similar to numpy's nanmax, but excludes `None` values.

    NB: If both values are `None`, this will still return `None`"""
    if value1 is None:
        return value2
    if value2 is None:
        return value1
    return max(value1, value2)


def _pick_columns(
    first_row_group,
    exclude_hats_columns: bool = True,
    exclude_columns: list[str] = None,
    include_columns: list[str] = None,
):
    """Convenience method to find the desired columns and their indexes, given
    some conventional user preferences."""

    if include_columns is None:
        include_columns = []

    if exclude_columns is None:
        exclude_columns = []
    if exclude_hats_columns:
        exclude_columns.extend(["Norder", "Dir", "Npix", "_healpix_29"])

    column_names = [
        first_row_group.column(col).path_in_schema for col in range(0, first_row_group.num_columns)
    ]
    column_names = [name.removesuffix(".list.element") for name in column_names]
    good_column_indexes = [
        index
        for index, name in enumerate(column_names)
        if (len(include_columns) == 0 or name in include_columns)
        and not (len(exclude_columns) > 0 and name in exclude_columns)
    ]
    column_names = [column_names[i] for i in good_column_indexes]

    return good_column_indexes, column_names


def aggregate_column_statistics(
    metadata_file: str | Path | UPath,
    exclude_hats_columns: bool = True,
    exclude_columns: list[str] = None,
    include_columns: list[str] = None,
    include_pixels: list[HealpixPixel] = None,
):
    """Read footer statistics in parquet metadata, and report on global min/max values.

    Args:
        metadata_file (str | Path | UPath): path to `_metadata` file
        exclude_hats_columns (bool): exclude HATS spatial and partitioning fields
            from the statistics. Defaults to True.
        exclude_columns (List[str]): additional columns to exclude from the statistics.
        include_columns (List[str]): if specified, only return statistics for the column
            names provided. Defaults to None, and returns all non-hats columns.
        include_pixels (list[HealpixPixel]): if specified, only return statistics
            for the pixels indicated. Defaults to none, and returns all pixels.
    Returns:
        dataframe with global summary statistics
    """
    total_metadata = file_io.read_parquet_metadata(metadata_file)
    num_row_groups = total_metadata.num_row_groups
    first_row_group = total_metadata.row_group(0)

    good_column_indexes, column_names = _pick_columns(
        first_row_group=first_row_group,
        exclude_hats_columns=exclude_hats_columns,
        exclude_columns=exclude_columns,
        include_columns=include_columns,
    )
    if not good_column_indexes:
        return pd.DataFrame()

    extrema = None

    for row_group_index in range(0, num_row_groups):
        row_group = total_metadata.row_group(row_group_index)
        if include_pixels is not None:
            pixel = paths.get_healpix_from_path(row_group.column(0).file_path)
            if pixel not in include_pixels:
                continue
        row_stats = [
            (
                (None, None, 0, 0)
                if row_group.column(col).statistics is None
                else (
                    row_group.column(col).statistics.min,
                    row_group.column(col).statistics.max,
                    row_group.column(col).statistics.null_count,
                    row_group.column(col).num_values,
                )
            )
            for col in good_column_indexes
        ]
        if extrema is None:
            extrema = row_stats
        ## This is annoying, but avoids extra copies, or none comparison.
        else:
            extrema = [
                (
                    (_nonemin(extrema[col][0], row_stats[col][0])),
                    (_nonemax(extrema[col][1], row_stats[col][1])),
                    extrema[col][2] + row_stats[col][2],
                    extrema[col][3] + row_stats[col][3],
                )
                for col in range(0, len(good_column_indexes))
            ]

    if extrema is None:
        return pd.DataFrame()

    stats_lists = np.array(extrema).T

    frame = pd.DataFrame(
        {
            "column_names": column_names,
            "min_value": stats_lists[0],
            "max_value": stats_lists[1],
            "null_count": stats_lists[2],
            "row_count": stats_lists[3],
        }
    ).set_index("column_names")
    return frame


def per_pixel_statistics(
    metadata_file: str | Path | UPath,
    exclude_hats_columns: bool = True,
    exclude_columns: list[str] = None,
    include_columns: list[str] = None,
    include_stats: list[str] = None,
    multi_index=False,
    include_pixels: list[HealpixPixel] = None,
):
    """Read footer statistics in parquet metadata, and report on statistics about
    each pixel partition.

    Args:
        metadata_file (str | Path | UPath): path to `_metadata` file
        exclude_hats_columns (bool): exclude HATS spatial and partitioning fields
            from the statistics. Defaults to True.
        exclude_columns (List[str]): additional columns to exclude from the statistics.
        include_columns (List[str]): if specified, only return statistics for the column
            names provided. Defaults to None, and returns all non-hats columns.
        include_pixels (list[HealpixPixel]): if specified, only return statistics
            for the pixels indicated. Defaults to none, and returns all pixels.
        include_stats (List[str]): if specified, only return the kinds of values from list
            (min_value, max_value, null_count, row_count). Defaults to None, and returns all values.
        multi_index (bool): should the returned frame be created with a multi-index, first on
            pixel, then on column name? Default is False, and instead indexes on pixel, with
            separate columns per-data-column and stat value combination.
    Returns:
        dataframe with granular per-pixel statistics
    """
    total_metadata = file_io.read_parquet_metadata(metadata_file)
    num_row_groups = total_metadata.num_row_groups
    first_row_group = total_metadata.row_group(0)

    good_column_indexes, column_names = _pick_columns(
        first_row_group=first_row_group,
        exclude_hats_columns=exclude_hats_columns,
        exclude_columns=exclude_columns,
        include_columns=include_columns,
    )
    if not good_column_indexes:
        return pd.DataFrame()

    all_stats = ["min_value", "max_value", "null_count", "row_count"]

    if include_stats is None or len(include_stats) == 0:
        include_stats = all_stats
    else:
        for stat in include_stats:
            if stat not in all_stats:
                raise ValueError(f"include_stats must be from list {all_stats} (found {stat})")

    stat_mask = np.array([ind for ind, stat in enumerate(all_stats) if stat in include_stats])
    pixels = []
    leaf_stats = []

    for row_group_index in range(0, num_row_groups):
        row_group = total_metadata.row_group(row_group_index)
        pixel = paths.get_healpix_from_path(row_group.column(0).file_path)
        if include_pixels is not None and pixel not in include_pixels:
            continue
        row_stats = [
            (
                [None, None, 0, 0]
                if row_group.column(col).statistics is None
                else [
                    row_group.column(col).statistics.min,
                    row_group.column(col).statistics.max,
                    row_group.column(col).statistics.null_count,
                    row_group.column(col).num_values,
                ]
            )
            for col in good_column_indexes
        ]
        row_stats = np.take(row_stats, stat_mask, axis=1)
        pixels.append(pixel)
        leaf_stats.append(row_stats)

    stats_lists = np.array(leaf_stats)
    original_shape = stats_lists.shape

    if multi_index:
        stats_lists = stats_lists.reshape((original_shape[0] * original_shape[1], original_shape[2]))
        frame = pd.DataFrame(
            stats_lists,
            index=pd.MultiIndex.from_product([pixels, column_names], names=["pixel", "column"]),
            columns=include_stats,
        )
    else:
        stats_lists = stats_lists.reshape((original_shape[0], original_shape[1] * original_shape[2]))
        mod_col_names = [[f"{col_name}: {stat}" for stat in include_stats] for col_name in column_names]
        mod_col_names = np.array(mod_col_names).flatten()
        frame = pd.DataFrame(stats_lists, index=pixels, columns=mod_col_names)
    return frame
