from importlib_metadata import version
import logging, multiprocessing, psutil
from typing import Any, Dict, List, Tuple, Union

import numpy as np
import xarray as xr

from .cds import CASAVisSet
from .stokes_types import stokes_types


def make_coords(
    xds: xr.Dataset, ddi: int, subtables: Tuple[xr.Dataset, ...]
) -> Dict[str, np.ndarray]:
    """Make the coords to be added to a partition or chunk (besides
    the time, baseline) basic structure

    Grabs:
    - channel (center) frequency values from the spw subtable
    - pol idxs from the pol+ddi subtables -> pol names via the stokes_types
    - antenna IDs from antenna subtable

    :param: sub-xds as (ant_xds, ddi_xds, spw_xds, pol_xds)
    """
    ant_xds, ddi_xds, spw_xds, pol_xds = subtables
    freq = spw_xds.chan_freq.values[
        ddi_xds.spectral_window_id.values[ddi], : xds.freq.shape[0]
    ]
    pol_ids = pol_xds.corr_type.values[
        ddi_xds.polarization_id.values[ddi], : xds.pol.shape[0]
    ]
    pol_names = np.vectorize(stokes_types.get)(pol_ids)
    ant_id = ant_xds.antenna_id.values
    coords = {
        "freq": freq,
        "pol": pol_names,
        "antenna_id": ant_id
        # These will be metainfo in partitions
        # "spw_id": [ddi_xds["spectral_window_id"].values[ddi]],
        # "pol_id": [ddi_xds["polarization_id"].values[ddi]],
    }
    return coords


def vis_xds_packager_cds(
    subtables: List[Tuple[str, xr.Dataset]],
    partitions: Dict[Any, xr.Dataset],
    descr_add: str = "",
):
    """Takes a a list of subtable xds datasets and a dictionary of data
    partition xds datasets and and packages them as a CASA vis dataset
    (cds)

    :param partitions: data partiions as xds datasets
    :param subtables: subtables as xds datasets
    :param descr_add: substring to add to the short descr string of the cds
    :return: A "cds" - container for the metainfo subtables and data partitions
    """
    vers = version("xradio")

    return CASAVisSet(
        subtables,
        partitions,
        f"CASA vis set produced by xradio {vers}/{descr_add}",
    )


def vis_xds_packager_mxds(
    partitions: Dict[Any, xr.Dataset],
    subtables: List[Tuple[str, xr.Dataset]],
    add_global_coords: bool = True,
) -> xr.Dataset:
    """Takes a dictionary of data partition xds datasets and a list of
    subtable xds datasets and packages them as a dataset of datasets
    (mxds)

    :param partitions: data partiions as xds datasets
    :param subtables: subtables as xds datasets
    :add_global_coords: whether to add coords to the output mxds
    :return: A "mxds" - xr.dataset of datasets
    """
    mxds = xr.Dataset(attrs={"metainfo": subtables, "partitions": partitions})

    if add_global_coords:
        mxds = mxds.assign_coords(make_global_coords(mxds))

    return mxds


def make_global_coords(mxds: xr.Dataset):
    coords = {}
    metainfo = mxds.attrs["metainfo"]
    if "antenna" in metainfo:
        coords["antenna_ids"] = metainfo["antenna"].antenna_id.values
        coords["antennas"] = xr.DataArray(
            metainfo["antenna"].name.values, dims=["antenna_ids"]
        )
    if "field" in metainfo:
        coords["field_ids"] = metainfo["field"].field_id.values
        coords["fields"] = xr.DataArray(
            metainfo["field"].name.values, dims=["field_ids"]
        )
    if "feed" in mxds.attrs:
        coords["feed_ids"] = metainfo["feed"].feed_id.values
    if "observation" in metainfo:
        coords["observation_ids"] = metainfo["observation"].observation_id.values
        coords["observations"] = xr.DataArray(
            metainfo["observation"].project.values, dims=["observation_ids"]
        )
    if "polarization" in metainfo:
        coords["polarization_ids"] = metainfo["polarization"].pol_setup_id.values
    if "source" in metainfo:
        coords["source_ids"] = metainfo["source"].source_id.values
        coords["sources"] = xr.DataArray(
            metainfo["source"].name.values, dims=["source_ids"]
        )
    if "spectral_window" in metainfo:
        coords["spw_ids"] = metainfo["spectral_window"].spw_id.values
    if "state" in metainfo:
        coords["state_ids"] = metainfo["state"].state_id.values

    return coords


def expand_xds(xds: xr.Dataset) -> xr.Dataset:
    """
    expand single (row) dimension of xds to (time, baseline)
    """
    assert "baseline" not in xds.coords

    txds = xds.copy()
    unique_baselines, baselines = np.unique(
        [txds.antenna1.values, txds.antenna2.values], axis=1, return_inverse=True
    )
    txds["baseline"] = xr.DataArray(baselines.astype("int32"), dims=["row"])

    try:
        txds = (
            txds.set_index(row=["time", "baseline"])
            .unstack("row")
            .transpose("time", "baseline", ...)
        )

        # unstack changes type to float when it needs to introduce NaNs, so
        # we need to reset to the proper type. Avoid if possible, as the
        # astype are costly
        for dv in txds.data_vars:
            if txds[dv].dtype != xds[dv].dtype:
                txds[dv] = txds[dv].astype(xds[dv].dtype)
    except Exception as exc:
        print(
            f"WARNING: Cannot expand rows to (time, baseline), "
            f"possibly duplicate values in (time, baseline). Exception: {exc}"
        )
        txds = xds.copy()

    return txds


def flatten_xds(xds: xr.Dataset) -> xr.Dataset:
    """
    flatten (time, baseline) dimensions of xds back to single dimension (row)
    """
    nan_int = np.array([np.nan]).astype("int32")[0]
    txds = xds.copy()

    # flatten the time x baseline dimensions of main table
    if ("time" in xds.dims) and ("baseline" in xds.dims):
        txds = xds.stack({"row": ("time", "baseline")}).transpose("row", ...)
        txds = txds.where(
            (txds.state_id != nan_int) & (txds.field_id != nan_int), drop=True
        )  # .unify_chunks()
        for dv in list(xds.data_vars):
            txds[dv] = txds[dv].astype(xds[dv].dtype)

    return txds


####################################
# xautomatically compute best data chunking
def optimal_chunking(
    ndim: Union[int, None] = None,
    didxs: Union[Tuple[int], List[int], None] = None,
    chunk_size: str = "auto",
    data_shape: Union[tuple, None] = None,
) -> tuple:
    """
    Determine the optimal chunk shape for reading an MS or Image based
    on machine resources and intended operations

    :param ndim: number of dimensions to chunk. An MS is 3, an
    expanded MS is 4. An image could be anywhere from 2 to 5. Not
    needed if data_shape is given.
    :param didxs: dimension indices over which subsequent operations
    will be performed. Values should be less than ndim. Tries to
    reduce inter-process communication of data contents. Needs to
    know the shape to do this well. Default None balances chunk size
    across all dimensions.
    :param chunk_size: target chunk size ('large', 'small', 'auto').
    Default 'auto' tries to guess by looking at CPU core count and
    available memory.
    :param data_shape: shape of the total MS DDI or Image data. Helps
    to know. Default None does not optimize based on shape

    :return: optimal chunking for reading the ms (row, chan, pol)
    """
    assert (ndim is not None) or (
        data_shape is not None
    ), "either ndim or data_shape must be given"
    assert chunk_size in ["large", "small", "auto"], "invalid chunk_size parameter"
    if ndim is None:
        ndim = len(data_shape)

    opt_dims = (
        didxs if (didxs is not None) and (len(didxs) > 0) else np.arange(ndim)
    )  # maximize these dim chunk sizes
    nonopt_dims = np.setdiff1d(np.arange(ndim), opt_dims)  # at the expense of these

    max_chunk_sizes = (
        data_shape
        if data_shape is not None
        else [dd for ii, dd in enumerate([10000, 10000, 10000, 4, 10]) if ii < ndim]
    )
    min_chunk_sizes = (
        np.ceil(np.array(data_shape) / 80).astype(int)
        if data_shape is not None
        else [1000, 1, 1]
        if ndim == 3
        else [dd for ii, dd in enumerate([10, 10, 1, 1, 1]) if ii < ndim]
    )
    target_size = (
        175 * 1024**2 / 8
    )  # ~175 MB chunk worst case with 8-byte DATA column
    bytes_per_core = int(
        round(
            ((psutil.virtual_memory().available * 0.10) / multiprocessing.cpu_count())
        )
    )
    if data_shape is not None:
        bytes_per_core = min(
            bytes_per_core, np.prod(data_shape) * 8 / 2
        )  # ensure at least two chunks
    if chunk_size == "large":
        target_size = target_size * 6  # ~1 GB
    if chunk_size == "auto":
        target_size = max(min(target_size * 6, bytes_per_core / 8), target_size)

    # start by setting the optimized dims to their max size and non-optimized dims to their min size
    chunks = np.zeros((ndim), dtype="int")
    chunks[opt_dims] = np.array(max_chunk_sizes)[opt_dims]
    chunks[nonopt_dims] = np.array(min_chunk_sizes)[nonopt_dims]

    # iteratively walk towards an optimal chunk size
    # iteration is needed because rounding to nearest integer index can make a big different (2x) in chunk size
    # for small dimensions like pol
    for ii in range(10):
        # if the resulting size is too big, reduce the sizes of the optimized dimensions
        if (np.prod(chunks) > target_size) and (len(opt_dims) > 0):
            chunks[opt_dims] = np.round(
                chunks[opt_dims]
                * (target_size / np.prod(chunks)) ** (1 / len(opt_dims))
            )
        # else if the resulting size is too small, increase the sizes of the non-optimized dimensions
        elif (np.prod(chunks) < target_size) and (len(nonopt_dims) > 0):
            chunks[nonopt_dims] = np.round(
                chunks[nonopt_dims]
                * (target_size / np.prod(chunks)) ** (1 / len(nonopt_dims))
            )
        chunks = np.min((chunks, max_chunk_sizes), axis=0)
        chunks = np.max((chunks, min_chunk_sizes), axis=0)

    return tuple(chunks)


def calc_optimal_ms_chunk_shape(
    memory_available_in_bytes, shape, element_size_in_bytes, column_name
) -> int:
    """
    Calculates the max number of rows (1st dim in shape) of a variable
    that can be fit in the memory for a thread.
    """
    factor = 0.8  # Account for memory used by other objects in thread.
    # total_mem = np.prod(shape)*element_size_in_bytes
    single_row_mem = np.prod(shape[1:]) * element_size_in_bytes

    try:
        assert single_row_mem < factor * memory_available_in_bytes
    except AssertionError as err:
        logging.exception(
            "Not engough memory in a thread to contain a row of "
            + column_name
            + ". Need at least "
            + str(single_row_mem / factor)
            + " bytes."
        )
        raise err

    rows_chunk_size = int((factor * memory_available_in_bytes) / single_row_mem)

    if rows_chunk_size > shape[0]:
        rows_chunk_size = shape[0]

    logging.debug(
        "Numbers of rows in chunk for " + column_name + ": " + str(rows_chunk_size)
    )

    return rows_chunk_size
