import toolviper.utils.logger as logger
from typing import Any, Dict, List, Tuple, Union

import dask, dask.array
import numpy as np
import xarray as xr
import pandas as pd

try:
    from casacore import tables
except ImportError:
    import xradio._utils._casacore.casacore_from_casatools as tables

from .read import (
    read_flat_col_chunk,
    read_col_chunk,
    convert_casacore_time,
    extract_table_attributes,
    add_units_measures,
)

from .table_query import open_table_ro, open_query
from xradio._utils.list_and_array import (
    unique_1d,
    pairing_function,
    inverse_pairing_function,
)

rename_msv2_cols = {
    "ANTENNA1": "antenna1_id",
    "ANTENNA2": "antenna2_id",
    "FEED1": "feed1_id",
    "FEED2": "feed2_id",
    # optional cols:
    "WEIGHT_SPECTRUM": "WEIGHT",
    "CORRECTED_DATA": "VIS_CORRECTED",
    "DATA": "VIS",
    "MODEL_DATA": "VIS_MODEL",
    "FLOAT_DATA": "AUTOCORR",
}


def rename_vars(mvars: Dict[str, xr.DataArray]) -> Dict[str, xr.DataArray]:
    """
    Apply rename rules. Also preserve ordering of data_vars

    Note: not using xr.DataArray.rename because we have optional
    column renames and rename complains if some of the names passed
    are not present in the dataset

    Parameters
    ----------
    mvars : Dict[str, xr.DataArray]
        dictionary of data_vars to be used to create an xr.Dataset

    Returns
    -------
    Dict[str, xr.DataArray]
        similar dictionary after applying MSv2 => MSv3/ngCASA renaming rules
    """
    renamed = {
        rename_msv2_cols[name] if name in rename_msv2_cols else name: var
        for name, var in mvars.items()
    }

    return renamed


def redim_id_data_vars(mvars: Dict[str, xr.DataArray]) -> Dict[str, xr.DataArray]:
    """
    Changes:
    Several id data variables to drop its baseline dim
    The antenna id data vars:
     From MS (antenna1_id(time, baseline), antenna2_id(time,baseline)
     To cds (baseline_ant1_id(baseline), baseline_ant2_id(baseline)

    Parameters
    ----------
    mvars : Dict[str, xr.DataArray]
        data variables being prepared for a partition xds

    Returns
    -------
    Dict[str, xr.DataArray]
        data variables with the ant id ones modified to cds type
    """
    # Vars to drop baseline dim
    var_names = [
        "ARRAY_ID",
        "OBSERVATION_ID",
        "PROCESSOR_ID",
        "SCAN_NUMBER",
        "STATE_ID",
    ]
    for vname in var_names:
        if "baseline" in mvars[vname].coords:
            mvars[vname] = mvars[vname].sel(baseline=0, drop=True)

    for idx in ["1", "2"]:
        new_name = f"baseline_ant{idx}_id"
        mvars[new_name] = mvars.pop(f"antenna{idx}_id")
        if "time" in mvars[new_name].coords:
            mvars[new_name] = mvars[new_name].sel(time=0, drop=True)

    return mvars


def get_partition_ids(mtable: tables.table, taql_where: str) -> Dict:
    """
    Get some of the partition IDs that we have to retrieve from some
    of the top level ID/sorting cols of the main table of the MS.

    Parameters
    ----------
    mtable : tables.table
        MS main table
    taql_where : str
        where part that defines the partition in TaQL

    Returns
    -------
    Dict
        ids of array, observation, and processor
    """

    taql_ids = f"select DISTINCT ARRAY_ID, OBSERVATION_ID, PROCESSOR_ID from $mtable {taql_where}"
    with open_query(mtable, taql_ids) as query:
        # array_id, observation_id, processor_id
        array_id = unique_1d(query.getcol("ARRAY_ID"))
        obs_id = unique_1d(query.getcol("OBSERVATION_ID"))
        proc_id = unique_1d(query.getcol("PROCESSOR_ID"))
        check_vars = [
            (array_id, "array_id"),
            (obs_id, "observation_id"),
            (proc_id, "processor_id"),
        ]
        for var, var_name in check_vars:
            if len(var) != 1:
                logger.warning(
                    f"Did not get exactly one {var_name} (got {var} for this partition. TaQL: {taql_where}"
                )

    pids = {
        "array_id": list(array_id),
        "observation_id": list(obs_id),
        "processor_id": list(proc_id),
    }
    return pids


def read_expanded_main_table(
    infile: str,
    ddi: int = 0,
    scan_state: Union[Tuple[int, int], None] = None,
    ignore_msv2_cols: Union[list, None] = None,
    chunks: Tuple[int, ...] = (400, 200, 100, 2),
) -> Tuple[xr.Dataset, Dict[str, Any], Dict[str, Any]]:
    """
    Reads one partition from the main table, all columns.
    This is the expanded version (time, baseline) dims.

    Chunk tuple: (time, baseline, freq, pol)

    Parameters
    ----------
    infile : str

    ddi : int  (Default value = 0)

    scan_state : Union[Tuple[int, int], None] (Default value = None)

    ignore_msv2_cols: Union[list, None] (Default value = None)

    chunks: Tuple[int, ...] (Default value = (400, 200, 100, 2))


    Returns
    -------
    Tuple[xr.Dataset, Dict[str, Any], Dict[str, Any]]
    """
    if ignore_msv2_cols is None:
        ignore_msv2_cols = []

    taql_where = f"where DATA_DESC_ID = {ddi}"
    if scan_state:
        # get partitions by scan/state
        scan, state = scan_state
        if type(state) == np.ndarray:
            state_ids_or = " OR STATE_ID = ".join(np.char.mod("%d", state))
            taql_where += f" AND (STATE_ID = {state_ids_or})"
        elif state:
            taql_where += f" AND SCAN_NUMBER = {scan} AND STATE_ID = {state}"
        elif scan:
            # scan can also be None, when partition_scheme='intent'
            # but the STATE table is empty!
            taql_where += f" AND SCAN_NUMBER = {scan}"

    with open_table_ro(infile) as mtable:
        # one partition, select just the specified ddi (+ scan/subscan)
        taql_main = f"select * from $mtable {taql_where}"
        with open_query(mtable, taql_main) as tb_tool:
            if tb_tool.nrows() == 0:
                tb_tool.close()
                mtable.close()
                return xr.Dataset(), {}, {}

            xds, attrs = read_main_table_chunks(
                infile, tb_tool, taql_where, ignore_msv2_cols, chunks
            )
            part_ids = get_partition_ids(tb_tool, taql_where)

    return xds, part_ids, attrs


def read_main_table_chunks(
    infile: str,
    tb_tool: tables.table,
    taql_where: str,
    ignore_msv2_cols: Union[list, None] = None,
    chunks: Tuple[int, ...] = (400, 200, 100, 2),
) -> Tuple[xr.Dataset, Dict[str, Any]]:
    """
    Iterates through the time,baseline chunks and reads slices from
    all the data columns.

    Parameters
    ----------
    infile : str

    tb_tool : tables.table

    taql_where : str

    ignore_msv2_cols : Union[list, None] (Default value = None)

    chunks: Tuple[int, ...] (Default value = (400, 200, 100, 2))


    Returns
    -------
    Tuple[xr.Dataset, Dict[str, Any]]
    """
    baselines = get_baselines(tb_tool)

    col_names = tb_tool.colnames()
    cshapes = [
        np.array(tb_tool.getcell(col, 0)).shape
        for col in col_names
        if tb_tool.iscelldefined(col, 0)
    ]
    chan_cnt, pol_cnt = [(cc[0], cc[1]) for cc in cshapes if len(cc) == 2][0]

    unique_times, tol = get_utimes_tol(tb_tool, taql_where)

    tvars = {}
    n_baselines = len(baselines)
    n_unique_times = len(unique_times)
    n_time_chunks = chunks[0]
    n_baseline_chunks = chunks[1]
    # loop over time chunks
    for time_chunk in range(0, n_unique_times, n_time_chunks):
        time_start = unique_times[time_chunk] - tol
        time_end = (
            unique_times[min(n_unique_times, time_chunk + n_time_chunks) - 1] + tol
        )

        # chunk time length
        ctlen = min(n_unique_times, time_chunk + n_time_chunks) - time_chunk

        bvars = {}
        # loop over baseline chunks
        for baseline_chunk in range(0, n_baselines, n_baseline_chunks):
            cblen = min(n_baselines - baseline_chunk, n_baseline_chunks)

            # read the specified chunk of data
            # def read_chunk(infile, ddi, times, blines, chans, pols):
            ttql = f"TIME BETWEEN {time_start} and {time_end}"
            ant1_start = baselines[baseline_chunk][0]
            ant1_end = baselines[cblen + baseline_chunk - 1][0]
            atql = f"ANTENNA1 BETWEEN {ant1_start} and {ant1_end}"
            ts_taql = f"select * from $mtable {taql_where} AND {ttql} AND {atql}"
            with open_query(None, ts_taql) as query_times_ants:
                tidxs = (
                    np.searchsorted(
                        unique_times, query_times_ants.getcol("TIME", 0, -1)
                    )
                    - time_chunk
                )
                ts_ant1, ts_ant2 = (
                    query_times_ants.getcol("ANTENNA1", 0, -1),
                    query_times_ants.getcol("ANTENNA2", 0, -1),
                )

            ts_bases = np.column_stack((ts_ant1, ts_ant2))

            bidxs = get_baseline_indices(baselines, ts_bases) - baseline_chunk

            # some antenna 2's will be out of bounds for this chunk, store rows that are in bounds
            didxs = np.where(
                (bidxs >= 0) & (bidxs < min(chunks[1], n_baselines - baseline_chunk))
            )[0]

            delayed_params = (infile, ts_taql, (ctlen, cblen), tidxs, bidxs, didxs)

            read_all_cols_bvars(
                tb_tool, chunks, chan_cnt, ignore_msv2_cols, delayed_params, bvars
            )

        concat_bvars_update_tvars(bvars, tvars)

    dims = ["time", "baseline", "freq", "pol"]
    mvars = concat_tvars_to_mvars(dims, tvars, pol_cnt, chan_cnt)

    mcoords = {
        "time": xr.DataArray(convert_casacore_time(unique_times), dims=["time"]),
        "baseline": xr.DataArray(np.arange(n_baselines), dims=["baseline"]),
    }

    # add xds global attributes
    cc_attrs = extract_table_attributes(infile)
    attrs = {"other": {"msv2": {"ctds_attrs": cc_attrs, "bad_cols": ignore_msv2_cols}}}
    # add per data var attributes
    mvars = add_units_measures(mvars, cc_attrs)
    mcoords = add_units_measures(mcoords, cc_attrs)

    mvars = rename_vars(mvars)
    mvars = redim_id_data_vars(mvars)
    xds = xr.Dataset(mvars, coords=mcoords)

    return xds, attrs


def get_utimes_tol(mtable: tables.table, taql_where: str) -> Tuple[np.ndarray, float]:
    taql_utimes = f"select DISTINCT TIME from $mtable {taql_where}"
    with open_query(mtable, taql_utimes) as query_utimes:
        utimes = unique_1d(query_utimes.getcol("TIME", 0, -1))
        # add a tol around the time ranges returned by taql
        if len(utimes) < 2:
            tol = 1e-5
        else:
            tol = np.diff(utimes).min() / 4

    return utimes, tol


def get_baselines(tb_tool: tables.table) -> np.ndarray:
    """
    Gets the unique baselines from antenna 1 and antenna 2 ids.

    Uses a pairing function and inverse pairing function to decrease the
    computation time of finding unique values.

    Parameters
    ----------
    tb_tool : tables.table
        MeasurementSet table to get the antenna ids.

    Returns
    -------
    unique_baselines : np.ndarray
        a 2D array of unique antenna pairs
        (baselines) from the MeasurementSet table provided.
    """
    ant1, ant2 = tb_tool.getcol("ANTENNA1", 0, -1), tb_tool.getcol("ANTENNA2", 0, -1)

    baselines = np.column_stack((ant1, ant2))

    # Using pairing function to reduce the computation time of finding unique values.
    baselines_paired = pairing_function(baselines)
    unique_baselines_paired = pd.unique(baselines_paired)
    unique_baselines = inverse_pairing_function(unique_baselines_paired)

    # Sorting the unique baselines.
    unique_baselines = unique_baselines[unique_baselines[:, 1].argsort()]
    unique_baselines = unique_baselines[
        unique_baselines[:, 0].argsort(kind="mergesort")
    ]

    return unique_baselines


def get_baseline_indices(
    unique_baselines: np.ndarray, baseline_set: np.ndarray
) -> np.ndarray:
    """
    Finds the baseline indices of a set of baselines using the unique baselines.

    Uses a pairing function to reduce the number of values so it's more
    efficient to find the indices.

    Parameters
    ----------
    unique_baselines : np.ndarray
        a 2D array of unique antenna pairs (baselines).
    baseline_set : np.ndarray
        a 2D array of antenna pairs (baselines). This array may contain duplicates.

    Returns
    -------
    baseline_indices : np.ndarray
        the indices of the baseline set that
        correspond to the unique baselines.
    """
    unique_baselines_paired = pairing_function(unique_baselines)
    baseline_set_paired = pairing_function(baseline_set)

    # Pairing function doesn't preserve order so they need to be sorted.
    unique_baselines_sorted = np.argsort(unique_baselines_paired)
    sorted_indices = np.searchsorted(
        unique_baselines_paired[unique_baselines_sorted],
        baseline_set_paired,
    )
    baseline_indices = unique_baselines_sorted[sorted_indices]

    return baseline_indices


def read_all_cols_bvars(
    tb_tool: tables.table,
    chunks: Tuple[int, ...],
    chan_cnt: int,
    ignore_msv2_cols: bool,
    delayed_params: Tuple,
    bvars: Dict[str, xr.DataArray],
) -> None:
    """
    Loops over each column and create delayed dask arrays

    Parameters
    ----------
    tb_tool : tables.table

    chunks : Tuple[int, ...]

    chan_cnt : int

    ignore_msv2_cols : bool

    delayed_params : Tuple

    bvars : Dict[str, xr.DataArray]


    Returns
    -------

    """

    col_names = tb_tool.colnames()
    for col in col_names:
        if (col in ignore_msv2_cols + ["TIME"]) or (not tb_tool.iscelldefined(col, 0)):
            continue
        if col not in bvars:
            bvars[col] = []

        cdata = tb_tool.getcol(col, 0, 1)[0]

        if len(cdata.shape) == 0:
            infile, ts_taql, (ctlen, cblen), tidxs, bidxs, didxs = delayed_params
            cshape = (ctlen, cblen)
            delayed_col = infile, ts_taql, col, cshape, tidxs, bidxs, didxs
            delayed_array = dask.delayed(read_col_chunk)(*delayed_col, None, None)
            bvars[col] += [dask.array.from_delayed(delayed_array, cshape, cdata.dtype)]

        elif col == "UVW":
            infile, ts_taql, (ctlen, cblen), tidxs, bidxs, didxs = delayed_params
            cshape = (ctlen, cblen, 3)
            delayed_3 = infile, ts_taql, col, cshape, tidxs, bidxs, didxs
            delayed_array = dask.delayed(read_col_chunk)(*delayed_3, None, None)
            bvars[col] += [dask.array.from_delayed(delayed_array, cshape, cdata.dtype)]

        elif len(cdata.shape) == 1:
            pol_list = []
            dd = 2 if cdata.shape == chan_cnt else 3
            for pc in range(0, cdata.shape[0], chunks[dd]):
                pols = (pc, min(cdata.shape[0], pc + chunks[dd]) - 1)
                infile, ts_taql, (ctlen, cblen), tidxs, bidxs, didxs = delayed_params
                cshape = (
                    ctlen,
                    cblen,
                ) + (pols[1] - pols[0] + 1,)
                delayed_cs = infile, ts_taql, col, cshape, tidxs, bidxs, didxs
                delayed_array = dask.delayed(read_col_chunk)(*delayed_cs, pols, None)
                pol_list += [
                    dask.array.from_delayed(delayed_array, cshape, cdata.dtype)
                ]
            bvars[col] += [dask.array.concatenate(pol_list, axis=2)]

        elif len(cdata.shape) == 2:
            chan_list = []
            for cc in range(0, cdata.shape[0], chunks[2]):
                chans = (cc, min(cdata.shape[0], cc + chunks[2]) - 1)
                pol_list = []
                for pc in range(0, cdata.shape[1], chunks[3]):
                    pols = (pc, min(cdata.shape[1], pc + chunks[3]) - 1)
                    (
                        infile,
                        ts_taql,
                        (ctlen, cblen),
                        tidxs,
                        bidxs,
                        didxs,
                    ) = delayed_params
                    cshape = (
                        ctlen,
                        cblen,
                    ) + (chans[1] - chans[0] + 1, pols[1] - pols[0] + 1)
                    delayed_cs = infile, ts_taql, col, cshape, tidxs, bidxs, didxs
                    delayed_array = dask.delayed(read_col_chunk)(
                        *delayed_cs, chans, pols
                    )
                    pol_list += [
                        dask.array.from_delayed(delayed_array, cshape, cdata.dtype)
                    ]
                chan_list += [dask.array.concatenate(pol_list, axis=3)]
            bvars[col] += [dask.array.concatenate(chan_list, axis=2)]


def concat_bvars_update_tvars(
    bvars: Dict[str, xr.DataArray], tvars: Dict[str, xr.DataArray]
) -> None:
    """
    concats all the dask chunks from each baseline. This is intended to
    be called iteratively, for every time chunk iteration, once all the
    baseline chunks have been read.

    Parameters
    ----------
    bvars: Dict[str, xr.DataArray]

    tvars: Dict[str, xr.DataArray]


    Returns
    -------

    """
    for kk in bvars.keys():
        if len(bvars[kk]) == 0:
            continue
        if kk not in tvars:
            tvars[kk] = []
        tvars[kk] += [dask.array.concatenate(bvars[kk], axis=1)]


def concat_tvars_to_mvars(
    dims: List[str], tvars: Dict[str, xr.DataArray], pol_cnt: int, chan_cnt: int
) -> Dict[str, xr.DataArray]:
    """
    Concat into a single dask array all the dask arrays from each time
    chunk to make the final arrays of the xds.

    Parameters
    ----------
    dims : List[str]
        dimension names
    tvars : Dict[str, xr.DataArray]
        variables as lists of dask arrays per time chunk
    pol_cnt : int
        len of pol axis/dim
    chan_cnt : int
        len of freq axis/dim (chan indices)

    Returns
    -------
    Dict[str, xr.DataArray]
        variables as concated dask arrays
    """

    mvars = {}
    for tvr in tvars.keys():
        data_var = tvr
        if tvr == "UVW":
            mvars[data_var] = xr.DataArray(
                dask.array.concatenate(tvars[tvr], axis=0),
                dims=dims[:2] + ["uvw_coords"],
            )
        elif len(tvars[tvr][0].shape) == 3 and (tvars[tvr][0].shape[-1] == pol_cnt):
            mvars[data_var] = xr.DataArray(
                dask.array.concatenate(tvars[tvr], axis=0), dims=dims[:2] + ["pol"]
            )
        elif len(tvars[tvr][0].shape) == 3 and (tvars[tvr][0].shape[-1] == chan_cnt):
            mvars[data_var] = xr.DataArray(
                dask.array.concatenate(tvars[tvr], axis=0), dims=dims[:2] + ["freq"]
            )
        else:
            mvars[data_var] = xr.DataArray(
                dask.array.concatenate(tvars[tvr], axis=0),
                dims=dims[: len(tvars[tvr][0].shape)],
            )

    return mvars


def read_flat_main_table(
    infile: str,
    ddi: Union[int, None] = None,
    scan_state: Union[Tuple[int, int], None] = None,
    rowidxs=None,
    ignore_msv2_cols: Union[List, None] = None,
    chunks: Tuple[int, ...] = (22000, 512, 2),
) -> Tuple[xr.Dataset, Dict[str, Any], Dict[str, Any]]:
    """
    Read main table using flat structure: no baseline dimension. Vis
    dimensions are: row, freq, pol
    (experimental, perhaps to be deprecated/removed). Works but some
    features may be missing and/or flaky.

    Chunk tuple: (row, freq, pol)

    Parameters
    ----------
    infile : str

    ddi : Union[int, None] (Default value = None)

    scan_state : Union[Tuple[int, int], None] (Default value = None)

    rowidxs : np.ndarray (Default value = None)

    ignore_msv2_cols : Union[List, None] (Default value = None)

    chunks : Tuple[int, ...] (Default value = (22000, 512, 2))

    Returns
    -------
    Tuple[xr.Dataset, Dict[str, Any], Dict[str, Any]]
    """
    taql_where = f"where DATA_DESC_ID = {ddi}"
    if scan_state:
        # TODO: support additional intent/scan/subscan conditions if
        # we keep this read_flat functionality
        scans, states = scan_state
        # get row indices relative to full main table
        if type(states) == np.ndarray:
            state_ids_or = " OR STATE_ID = ".join(np.char.mod("%d", states))
            taql_where += f" AND (STATE_ID = {state_ids_or})"
        elif states is not None:
            taql_where += f" AND STATE_ID = {states}"
        elif scans is not None:
            taql_where += f" AND SCAN_NUMBER = {scans}"

    mtable = tables.table(
        infile, readonly=True, lockoptions={"option": "usernoread"}, ack=False
    )

    # get row indices relative to full main table
    if rowidxs is None:
        taql_rowid = f"select rowid() as ROWS from $mtable {taql_where}"
        with open_query(mtable, taql_rowid) as query_rows:
            rowidxs = query_rows.getcol("ROWS")

    nrows = len(rowidxs)
    if nrows == 0:
        mtable.close()
        return xr.Dataset(), {}, {}

    part_ids = get_partition_ids(mtable, taql_where)

    taql_cols = f"select * from $mtable {taql_where}"
    with open_query(mtable, taql_cols) as query_cols:
        cols = query_cols.colnames()
        ignore = [
            col
            for col in cols
            if (not query_cols.iscelldefined(col, 0))
            or (query_cols.coldatatype(col) == "record")
        ]
        cdata = dict(
            [
                (col, query_cols.getcol(col, 0, 1))
                for col in cols
                if (col not in ignore)
                and not (ignore_msv2_cols and col in ignore_msv2_cols)
            ]
        )
        chan_cnt, pol_cnt = [
            (cdata[cc].shape[1], cdata[cc].shape[2])
            for cc in cdata
            if len(cdata[cc].shape) == 3
        ][0]

    mtable.close()

    mvars, mcoords, bvars, xds = {}, {}, {}, xr.Dataset()
    # loop over row chunks
    for rc in range(0, nrows, chunks[0]):
        crlen = min(chunks[0], nrows - rc)  # chunk row length
        rcidxs = rowidxs[rc : rc + chunks[0]]

        # loop over each column and create delayed dask arrays
        for col in cdata.keys():
            if col not in bvars:
                bvars[col] = []

            if len(cdata[col].shape) == 1:
                delayed_array = dask.delayed(read_flat_col_chunk)(
                    infile, col, (crlen,), rcidxs, None, None
                )
                bvars[col] += [
                    dask.array.from_delayed(delayed_array, (crlen,), cdata[col].dtype)
                ]

            elif col == "UVW":
                delayed_array = dask.delayed(read_flat_col_chunk)(
                    infile, col, (crlen, 3), rcidxs, None, None
                )
                bvars[col] += [
                    dask.array.from_delayed(delayed_array, (crlen, 3), cdata[col].dtype)
                ]

            elif len(cdata[col].shape) == 2:
                pol_list = []
                dd = 1 if cdata[col].shape[1] == chan_cnt else 2
                for pc in range(0, cdata[col].shape[1], chunks[dd]):
                    plen = min(chunks[dd], cdata[col].shape[1] - pc)
                    delayed_array = dask.delayed(read_flat_col_chunk)(
                        infile, col, (crlen, plen), rcidxs, None, pc
                    )
                    pol_list += [
                        dask.array.from_delayed(
                            delayed_array, (crlen, plen), cdata[col].dtype
                        )
                    ]
                bvars[col] += [dask.array.concatenate(pol_list, axis=1)]

            elif len(cdata[col].shape) == 3:
                chan_list = []
                for cc in range(0, chan_cnt, chunks[1]):
                    clen = min(chunks[1], chan_cnt - cc)
                    pol_list = []
                    for pc in range(0, cdata[col].shape[2], chunks[2]):
                        plen = min(chunks[2], cdata[col].shape[2] - pc)
                        delayed_array = dask.delayed(read_flat_col_chunk)(
                            infile, col, (crlen, clen, plen), rcidxs, cc, pc
                        )
                        pol_list += [
                            dask.array.from_delayed(
                                delayed_array, (crlen, clen, plen), cdata[col].dtype
                            )
                        ]
                    chan_list += [dask.array.concatenate(pol_list, axis=2)]
                bvars[col] += [dask.array.concatenate(chan_list, axis=1)]

    # now concat all the dask chunks from each time to make the xds
    mvars = {}
    for kk in bvars.keys():
        data_var = kk
        if len(bvars[kk]) == 0:
            ignore += [kk]
            continue
        if kk == "UVW":
            mvars[data_var] = xr.DataArray(
                dask.array.concatenate(bvars[kk], axis=0), dims=["row", "uvw_coords"]
            )
        elif len(bvars[kk][0].shape) == 2 and (bvars[kk][0].shape[-1] == pol_cnt):
            mvars[data_var] = xr.DataArray(
                dask.array.concatenate(bvars[kk], axis=0), dims=["row", "pol"]
            )
        elif len(bvars[kk][0].shape) == 2 and (bvars[kk][0].shape[-1] == chan_cnt):
            mvars[data_var] = xr.DataArray(
                dask.array.concatenate(bvars[kk], axis=0), dims=["row", "chan"]
            )
        else:
            mvars[data_var] = xr.DataArray(
                dask.array.concatenate(bvars[kk], axis=0),
                dims=["row", "freq", "pol"][: len(bvars[kk][0].shape)],
            )

    mvars["time"] = xr.DataArray(
        convert_casacore_time(mvars["TIME"].values), dims=["row"]
    ).chunk({"row": chunks[0]})

    # add xds global attributes
    cc_attrs = extract_table_attributes(infile)
    attrs = {"other": {"msv2": {"ctds_attrs": cc_attrs, "bad_cols": ignore}}}
    # add per data var attributes
    mvars = add_units_measures(mvars, cc_attrs)
    mcoords = add_units_measures(mcoords, cc_attrs)

    mvars = rename_vars(mvars)
    mvars = redim_id_data_vars(mvars)
    xds = xr.Dataset(mvars, coords=mcoords)

    return xds, part_ids, attrs
