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

import pandas as pd
import numpy as np
import xarray as xr

from casacore import tables

from .load import load_col_chunk
from .read_main_table import get_partition_ids, redim_id_data_vars, rename_vars
from .read import add_units_measures, convert_casacore_time, extract_table_attributes
from .write import revert_time
from .table_query import open_query, open_table_ro
from xradio.measurement_set._utils._ms._tables.read_main_table import (
    get_baselines,
    get_baseline_indices,
)
from xradio._utils.list_and_array import unique_1d


def load_expanded_main_table_chunk(
    infile: str,
    ddi: int,
    chunk: Dict[str, slice],
    ignore_msv2_cols: Union[list, None] = None,
) -> xr.Dataset:
    """
    Load a chunk of data from main table into memory, with expanded
    dims: (time, baseline, freq, pols)

    Parameters
    ----------
    infile : str
        Input MS path
    ddi : int
        DDI to load chunk from
    chunk : Dict[str, slice]
        specification of chunk to load
    ignore_msv2_cols : Union[list, None] (Default value = None)
        cols that should not be loaded (deprecated MSv2 or similar)

    Returns
    -------
    xr.Dataset
        Xarray datasets with chunk of visibility data, one per DDI (spw_id, pol_setup_id) pair
    """

    taql_where = f"where DATA_DESC_ID = {ddi}"
    taql_ddi = f"select * from $mtable {taql_where}"

    with open_table_ro(infile) as mtable:
        with open_query(mtable, taql_ddi) as tb_tool:
            if tb_tool.nrows() == 0:
                return xr.Dataset()

            xds, part_ids, attrs = load_expanded_ddi_chunk(
                infile, tb_tool, taql_where, chunk, ignore_msv2_cols
            )

    return xds, part_ids, attrs


def load_expanded_ddi_chunk(
    infile: str,
    tb_tool: tables.table,
    taql_pre: str,
    chunk: Dict[str, slice],
    ignore_msv2_cols: Union[list, None] = None,
) -> xr.Dataset:
    """
    Helper function to effectively load the chunk and produce an
    xr.Dataset from a DII once the table and initial query(ies) have
    been opened successfully.

    Parameters
    ----------
    infile : str
        Input MS path
    tb_tool : tables.table
        table query contrained to one DDI and chunk time range
    taql_pre : str
        TaQL query used for tb_tool, with some pre-selection of rows and columns
    chunk : Dict[str, slice]
        specification of data chunk to load
    ignore_msv2_cols : Union[list, None] (Default value = None)
        propagated from calling funtions

    Returns
    -------
    xr.Dataset
        An Xarray dataset with data variables as plain numpy
        arrays loaded directly from the MS columns
    """

    # read the specified chunk of data, figure out indices and lens
    utimes, times = get_chunk_times(taql_pre, chunk)
    baselines, blines = get_chunk_baselines(tb_tool, chunk)
    tidxs, bidxs, didxs, taql_where_chunk = get_chunk_data_indices(
        taql_pre, chunk, utimes, times, baselines, blines
    )

    ctlen = min(len(utimes), times[1] - times[0] + 1)
    cblen = min(len(baselines), blines[1] - blines[0] + 1)
    mvars = load_ddi_cols_chunk(
        ctlen, cblen, tidxs, bidxs, didxs, tb_tool, chunk, ignore_msv2_cols
    )

    mcoords = {
        "time": xr.DataArray(convert_casacore_time(utimes[:ctlen]), dims=["time"]),
        "baseline": xr.DataArray(np.arange(cblen), 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)

    part_ids = get_partition_ids(tb_tool, taql_where_chunk)

    # needs an ~equivalent to add_partition_attrs?
    return xds, part_ids, attrs


def load_ddi_cols_chunk(
    ctlen: int,
    cblen: int,
    tidxs: np.ndarray,
    bidxs: np.ndarray,
    didxs: np.ndarray,
    tb_tool: tables.table,
    chunk: Dict[str, slice],
    ignore_msv2_cols: Union[list, None] = None,
) -> Dict[str, np.ndarray]:
    """
    For a given chunk and DDI, load all the MSv2 columns

    Parameters
    ----------
    ctlen : int
        length of the time axis/dim of the chunk
    cblen : int
        length of the baseline axis of the chunk
    tidxs : np.ndarray
        time axis indices
    bidxs : np.ndarray
        baseline axis indices
    didxs : np.ndarray
        (effective) data indices, excluding missing baselines
    tb_tool : tables.table
        a table/TaQL query open and being used to load columns
    chunk : Dict[str, slice]
        data chunk specification
    ignore_msv2_cols : Union[list, None] (Default value = None)
        propagated from calling funtions

    Returns
    -------
    Dict[str, np.ndarray]
        columns loaded into memory as np arrays

    """
    cols = tb_tool.colnames()

    cshapes = [
        np.array(tb_tool.getcell(col, 0)).shape
        for col in cols
        if tb_tool.iscelldefined(col, 0)
    ]
    # Assumes shapes are consistent across columns - MSv2
    chan_cnt, pol_cnt = [(csh[0], csh[1]) for csh in cshapes if len(csh) == 2][0]

    dims = ["time", "baseline", "freq", "pol"]
    mvars = {}
    # loop over each column and load data
    for col in cols:
        if (col in ignore_msv2_cols + ["TIME"]) or not tb_tool.iscelldefined(col, 0):
            continue

        cdata = tb_tool.getcol(col, 0, 1)[0]
        cell_shape = cdata.shape
        if len(cell_shape) == 0:
            col_dims = dims[:2]
            mvars[col] = xr.DataArray(
                load_col_chunk(
                    tb_tool, col, (ctlen, cblen), tidxs, bidxs, didxs, None, None
                ),
                dims=col_dims,
            )

        elif col == "UVW":
            col_dims = dims[:2] + ["uvw_coords"]
            mvars[col] = xr.DataArray(
                load_col_chunk(
                    tb_tool, col, (ctlen, cblen, 3), tidxs, bidxs, didxs, None, None
                ),
                dims=col_dims,
            )

        elif len(cell_shape) == 1:
            pols, col_dims = get_col_1d_pols(cell_shape, dims, chan_cnt, pol_cnt, chunk)
            cshape = (ctlen, cblen) + (pols[1] - pols[0] + 1,)
            mvars[col] = xr.DataArray(
                load_col_chunk(tb_tool, col, cshape, tidxs, bidxs, didxs, pols, None),
                dims=col_dims,
            )

        elif len(cell_shape) == 2:
            chans, pols = get_col_2d_chans_pols(cell_shape, chan_cnt, pol_cnt, chunk)
            cshape = (ctlen, cblen) + (chans[1] - chans[0] + 1, pols[1] - pols[0] + 1)
            col_dims = dims
            mvars[col] = xr.DataArray(
                load_col_chunk(tb_tool, col, cshape, tidxs, bidxs, didxs, chans, pols),
                dims=col_dims,
            )

    return mvars


def get_chunk_times(
    taql_pre: str, chunk: Dict[str, slice]
) -> Tuple[np.ndarray, Tuple[int, int]]:
    """
    Produces time col/axis related values for a chunk: unique times,
    start/stop times.

    Parameters
    ----------
    taql_pre : str
        TaQL query used for tb_tool, with some pre-selection
        of rows and columns.
    chunk : Dict[str, slice]
        specification of data chunk to load

    Returns
    -------
    Tuple[np.ndarray, Tuple[int, int]]
        array of unique times + (firsr, last) time in the chunk
    """

    taql_utimes = f"select DISTINCT TIME from $mtable {taql_pre}"
    with open_query(None, 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

    if "time" in chunk:
        time_slice = chunk["time"]
        if (
            type(time_slice.start) == pd.Timestamp
            and type(time_slice.stop) == pd.Timestamp
        ):
            times = (
                revert_time(time_slice.start) - tol,
                revert_time(time_slice.stop) + tol,
            )
        elif (
            int(time_slice.start) == time_slice.start
            and int(time_slice.stop) == time_slice.stop
        ):
            # could be operator.index(time_slice.start):
            nutimes = len(utimes)
            times = (
                min(nutimes, int(time_slice.start)),
                min(nutimes, int(time_slice.stop)) - 1,
            )
        else:
            raise ValueError(
                f"Invalid time type. Not a timestamp and Cannot use as"
                f" index: {time_slice.start} (type: {type(time_slice.start)})"
            )
    else:
        times = (utimes[0], utimes[-1])

    return utimes, times


def get_chunk_baselines(
    tb_tool: tables.table, chunk: Dict[str, slice]
) -> Tuple[np.ndarray, Tuple[int, int]]:
    """
    Produces the basline col/axis related values for a chunk: an array of
    baselines and the start/stop baseline indices.

    Parameters
    ----------
    tb_tool : tables.table
        table/query opened with prev selections (time)
    chunk : Dict[str, slice]
        specification of data chunk to load

    Returns
    -------
    Tuple[np.ndarray, Tuple[int, int]]
        array of baselines + (first, last) baseline in the chunk
    """
    baselines = get_baselines(tb_tool)

    if "baseline" in chunk:
        baseline_chunk = chunk["baseline"]
        baseline_boundaries = (int(baseline_chunk.start), int(baseline_chunk.stop))
    else:
        baseline_boundaries = (baselines[0][0], baselines[-1][0] - 1)

    return baselines, baseline_boundaries


def get_chunk_data_indices(
    taql_pre: str,
    chunk: Dict[str, slice],
    utimes: np.ndarray,
    times: Tuple[int, int],
    baselines: np.ndarray,
    blines: Tuple[int, int],
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, str]:
    """
    Produces indices to pass to the casacore getcol(slice) functions to load
    the chunk of data. tidxs (time), bidxs (baseline), didxs (effective data
    indices, considering present/absent baselines).

    Time selection is added on top of that.

    Parameters
    ----------
    taql_pre : str
        TaQL query constraints to prepend/inject
    chunk : Dict[str, slice]
        specification of data chunk
    utimes : np.ndarray
        array of times in the chunk
    times : Tuple[int, int]
        start, stop time indices
    baselines : np.ndarray
        array of baselines inthe chunk
    blines : Tuple[int, int]
        start, stop baseline indices

    Returns
    -------
    Tuple[np.ndarray, np.ndarray, np.ndarray]
        indices along the time, baseline and data (time/baseline)
        axes + the full where... string defined for this chunk

    """

    taql_time = f"TIME BETWEEN {utimes[times[0]]} AND {utimes[times[1]]}"
    taql_ant = f"ANTENNA1 BETWEEN {blines[0]} and {blines[1]}"
    taql_where_chunk = f"{taql_pre} AND {taql_time} AND {taql_ant}"
    taql_chunk = f"select * from $mtable {taql_where_chunk}"
    with open_query(None, taql_chunk) as query_times_ants:
        logger.debug(
            f"Opened chunk query, with {query_times_ants.nrows()} rows. Query: {taql_chunk}"
        )
        tidxs = (
            np.searchsorted(utimes, query_times_ants.getcol("TIME", 0, -1)) - times[0]
        )
        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) - blines[0]

    # 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(blines[1] - blines[0] + 1, len(baselines) - blines[0]))
    )[0]

    return tidxs, bidxs, didxs, taql_where_chunk


def get_col_1d_pols(
    cell_shape: Tuple[int],
    dims: List[str],
    chan_cnt: int,
    pol_cnt: int,
    chunk: Dict[str, slice],
) -> Tuple[Tuple[int, int], List[str]]:
    """
    For a column with 1d array values, calculate the start/stop
    indices for the last dimension (either pol or freq).
    It also produces the adequate dimension names.

    Parameters
    ----------
    cell_shape : Tuple[int]
        shape of the column
    dims : List[str]
        full list of dataset dimensions
    chan_cnt : int
        number of channels
    pol_cnt : int
        number of pols
    chunk : Dict[str, slice]
        data chunk specification

    Returns
    -------
    Tuple[Tuple[int, int], List[str]]
        first and last pol/freq index of the chunk, and its
        dimension names

    """
    if cell_shape == chan_cnt:
        # chan/freq
        col_dims = dims[:2] + ["freq"]
        if "freq" in chunk:
            pols = (
                min(chan_cnt, chunk["freq"].start),
                min(chan_cnt, chunk["freq"].stop) - 1,
            )
        else:
            pols = (0, cell_shape[0])
    else:
        # pol
        col_dims = dims[:2] + ["pol"]
        if "pol" in chunk:
            pols = (
                min(pol_cnt, chunk["pol"].start),
                min(pol_cnt, chunk["pol"].stop) - 1,
            )
        else:
            pols = (0, cell_shape[0])

    return pols, col_dims


def get_col_2d_chans_pols(
    cell_shape: Tuple[int],
    chan_cnt: int,
    pol_cnt: int,
    chunk: Dict[str, slice],
) -> Tuple[Tuple[int, int], Tuple[int, int]]:
    """
    For a column with 2d array values (FLAG, DATA, WEIGHT_SPECTRUM,
    etc., calculate the the start/stop indices for the last two
    dimensions of the chunk (freq and pol).
    The dimension names can be assumed to be the full list of dims in
    visibilities (time, baseline, freq, pol).

    Parameters
    ----------
    cell_shape : Tuple[int]
        shape of the column
    chan_cnt : int
        number of channels
    pol_cnt : int
        number of pols
    chunk : Dict[str, slice]
        data chunk specification

    Returns
    -------
    Tuple[Tuple[int, int], Tuple[int, int]]
        first and last index for freq (channel) and pol axes of
        the chunk

    """
    if "freq" in chunk:
        chans = (
            min(chan_cnt, chunk["freq"].start),
            min(chan_cnt, chunk["freq"].stop) - 1,
        )
    else:
        chans = (0, cell_shape[0])

    if "pol" in chunk:
        pols = (
            min(pol_cnt, chunk["pol"].start),
            min(pol_cnt, chunk["pol"].stop) - 1,
        )
    else:
        pols = (0, cell_shape[1])

    return chans, pols
