import astropy
from astropy.coordinates import Angle, SkyCoord
from casacore import images, tables
import copy
import dask.array as da
import numpy as np
import os
from typing import Union
import xarray as xr
from .common import _active_mask, _create_new_image, _object_name, _pointing_center
from ..common import _compute_sky_reference_pixel, _doppler_types
from ...._utils._casacore.tables import open_table_rw


def _compute_direction_dict(xds: xr.Dataset) -> dict:
    """
    Given xds metadata, compute the direction dict that is valid
    for a CASA image coordinate system
    """
    direction = {}
    xds_dir = xds.attrs["direction"]
    direction["system"] = xds_dir["reference"]["equinox"]
    direction["projection"] = xds_dir["projection"]
    direction["projection_parameters"] = xds_dir["projection_parameters"]
    direction["units"] = np.array(xds_dir["reference"]["units"], dtype="<U16")
    direction["crval"] = np.array(xds_dir["reference"]["value"])
    direction["cdelt"] = np.array(xds_dir["reference"]["cdelt"])
    direction["crpix"] = _compute_sky_reference_pixel(xds)
    direction["pc"] = np.array(xds_dir["pc"])
    direction["axes"] = ["Right Ascension", "Declination"]
    direction["conversionSystem"] = direction["system"]
    for s in ["longpole", "latpole"]:
        # longpole, latpole are numerical values in degrees in casa images
        direction[s] = Angle(str(xds_dir[s]["value"]) + xds_dir[s]["units"]).deg
    return direction


def _compute_linear_dict(xds: xr.Dataset) -> dict:
    linear = {}
    u = xds["u"].attrs
    v = xds["v"].attrs
    linear["crval"] = np.array([u["crval"], v["crval"]])
    linear["cdelt"] = np.array([u["cdelt"], v["cdelt"]])
    linear["axes"] = ["UU", "VV"]
    linear["units"] = [u["units"], v["units"]]
    lu = len(xds.coords["u"])
    lv = len(xds.coords["v"])
    linear["crpix"] = np.array([0.0, 0.0], dtype=np.float64)
    # np.interp() appears to require the x coordinate be monotonically increasing
    for i, w, z in zip([0, 1], [u, v], ["u", "v"]):
        x = xds.coords[z].values
        y = np.array(range(len(x)), dtype=np.float64)
        if w["cdelt"] < 0:
            x = x[::-1]
            y = y[::-1]
        linear["crpix"][i] = np.interp(0, x, y)
    linear["pc"] = np.array([[1.0, 0.0], [0.0, 1.0]])
    return linear


def _compute_spectral_dict(xds: xr.Dataset) -> dict:
    """
    Given xds metadata, compute the spectral dict that is valid
    for a CASA image coordinate system
    """
    spec = {}
    spec["formatUnit"] = ""
    spec["name"] = "Frequency"
    # spec["nativeType"] = _native_types.index(xds.frequency.attrs["native_type"])
    # FREQ
    spec["nativeType"] = 0
    spec["restfreq"] = xds.frequency.attrs["rest_frequency"]["value"]
    # spec["restfreqs"] = copy.deepcopy(xds.frequency.attrs["restfreqs"]["value"])
    spec["restfreqs"] = [spec["restfreq"]]
    spec["system"] = xds.frequency.attrs["frame"]
    spec["unit"] = xds.frequency.attrs["units"]
    spec["velType"] = _doppler_types.index(xds.velocity.attrs["doppler_type"])
    spec["velUnit"] = xds.velocity.attrs["units"]
    spec["version"] = 2
    spec["waveUnit"] = xds.frequency.attrs["wave_unit"]
    wcs = {}
    wcs["ctype"] = "FREQ"
    wcs["pc"] = 1.0
    wcs["crval"] = xds.frequency.attrs["crval"]
    wcs["cdelt"] = xds.frequency.attrs["cdelt"]
    wcs["crpix"] = (wcs["crval"] - xds.frequency.values[0]) / wcs["cdelt"]
    spec["wcs"] = wcs
    return spec


def _coord_dict_from_xds(xds: xr.Dataset) -> dict:
    coord = {}
    coord["telescope"] = xds.attrs["telescope"]["name"]
    coord["observer"] = xds.attrs["observer"]
    obsdate = {}
    obsdate["refer"] = xds.coords["time"].attrs["scale"]
    obsdate["type"] = "epoch"
    obsdate["m0"] = {}
    obsdate["m0"]["unit"] = xds.coords["time"].attrs["units"]
    obsdate["m0"]["value"] = xds.coords["time"].values[0]
    coord["obsdate"] = obsdate
    coord["pointingcenter"] = xds.attrs[_pointing_center].copy()
    telpos = {}
    telpos["refer"] = xds.attrs["telescope"]["position"]["ellipsoid"]
    if xds.attrs["telescope"]["position"]["ellipsoid"] == "GRS80":
        telpos["refer"] = "ITRF"
    for i in range(3):
        telpos[f"m{i}"] = {
            "unit": xds.attrs["telescope"]["position"]["units"][i],
            "value": xds.attrs["telescope"]["position"]["value"][i],
        }
    telpos["type"] = "position"
    coord["telescopeposition"] = telpos
    if "l" in xds.coords:
        coord["direction0"] = _compute_direction_dict(xds)
    else:
        coord["linear0"] = _compute_linear_dict(xds)
    coord["stokes1"] = {
        "axes": np.array(["Stokes"], dtype="<U16"),
        "cdelt": np.array([1.0]),
        "crpix": np.array([0.0]),
        "crval": np.array([1.0]),
        "pc": np.array([[1.0]]),
        "stokes": np.array(xds.polarization.values, dtype="<U16"),
    }
    coord["spectral2"] = _compute_spectral_dict(xds)
    coord["pixelmap0"] = np.array([0, 1])
    coord["pixelmap1"] = np.array([2])
    coord["pixelmap2"] = np.array([3])
    coord["pixelreplace0"] = np.array([0.0, 0.0])
    coord["pixelreplace1"] = np.array([0.0])
    coord["pixelreplace2"] = np.array([0.0])
    coord["worldmap0"] = np.array([0, 1], dtype=np.int32)
    coord["worldmap1"] = np.array([2], dtype=np.int32)
    coord["worldmap2"] = np.array([3], dtype=np.int32)
    # this probbably needs some verification
    coord["worldreplace0"] = [0.0, 0.0]
    coord["worldreplace1"] = np.array(coord["stokes1"]["crval"])
    coord["worldreplace2"] = np.array([xds.frequency.attrs["crval"]])
    return coord


def _history_from_xds(xds: xr.Dataset, image: str) -> None:
    nrows = len(xds.history.row) if "row" in xds.data_vars else 0
    if nrows > 0:
        # TODO need to implement nrows == 0 case
        with open_table_rw(os.sep.join([image, "logtable"])) as tb:
            tb.addrows(nrows + 1)
            for c in ["TIME", "PRIORITY", "MESSAGE", "LOCATION", "OBJECT_ID"]:
                vals = xds.history[c].values
                if c == "TIME":
                    k = time.time() + 40587 * 86400
                elif c == "PRIORITY":
                    k = "INFO"
                elif c == "MESSAGE":
                    k = (
                        "Wrote xds to "
                        + os.path.basename(image)
                        + " using cngi_io.xds_to_casa_image_2()"
                    )
                elif c == "LOCATION":
                    k = "cngi_io.xds_to_casa_image_2"
                elif c == "OBJECT_ID":
                    k = ""
                vals = np.append(vals, k)
                tb.putcol(c, vals)


def _imageinfo_dict_from_xds(xds: xr.Dataset) -> dict:
    ii = {}
    ap_sky = "sky" if "l" in xds.coords else "apeture"
    ii["image_type"] = (
        xds[ap_sky].attrs["image_type"] if "image_type" in xds[ap_sky].attrs else ""
    )
    ii["objectname"] = xds.attrs[_object_name]
    if "beam" in xds.data_vars:
        # multi beam
        pp = {}
        pp["nChannels"] = len(xds.frequency)
        pp["nStokes"] = len(xds.polarization)
        bu = xds.beam.attrs["units"]
        chan = 0
        polarization = 0
        bv = xds.beam.values
        for i in range(pp["nChannels"] * pp["nStokes"]):
            bp = bv[0][polarization][chan][:]
            b = {
                "major": {"unit": bu, "value": bp[0]},
                "minor": {"unit": bu, "value": bp[1]},
                "positionangle": {"unit": bu, "value": bp[2]},
            }
            pp["*" + str(pp["nChannels"] * polarization + chan)] = b
            chan += 1
            if chan >= pp["nChannels"]:
                chan = 0
                polarization += 1
        ii["perplanebeams"] = pp
    elif "beam" in xds.attrs and xds.attrs["beam"]:
        # do nothing if xds.attrs['beam'] is None
        ii["restoringbeam"] = copy.deepcopy(xds.attrs["beam"])
        for k in ["major", "minor", "pa"]:
            del ii["restoringbeam"][k]["type"]
            ii["restoringbeam"][k]["unit"] = ii["restoringbeam"][k]["units"]
            del ii["restoringbeam"][k]["units"]
        ii["restoringbeam"]["positionangle"] = copy.deepcopy(ii["restoringbeam"]["pa"])
        del ii["restoringbeam"]["pa"]
    return ii


def _write_casa_data(xds: xr.Dataset, image_full_path: str) -> None:
    sky_ap = "sky" if "sky" in xds else "apeture"
    if xds[sky_ap].shape[0] != 1:
        raise RuntimeError("XDS can only be converted if it has exactly one time plane")
    trans_coords = (
        ("frequency", "polarization", "m", "l")
        if sky_ap == "sky"
        else ("frequency", "polarization", "v", "u")
    )
    casa_image_shape = xds[sky_ap].isel(time=0).transpose(*trans_coords).shape[::-1]
    active_mask = xds.attrs["active_mask"] if _active_mask in xds.attrs else ""
    masks = []
    masks_rec = {}
    mask_rec = {
        "box": {
            "blc": np.array([1.0, 1.0, 1.0, 1.0]),
            "comment": "",
            "isRegion": 1,
            "name": "LCBox",
            "oneRel": True,
            "shape": np.array(casa_image_shape),
            "trc": np.array(casa_image_shape),
        },
        "comment": "",
        "isRegion": 1,
        "name": "LCPagedMask",
    }
    for m in xds.data_vars:
        attrs = xds[m].attrs
        if "image_type" in attrs and attrs["image_type"] == "Mask":
            masks_rec[m] = mask_rec
            masks_rec[m]["mask"] = f"Table: {os.sep.join([image_full_path, m])}"
            masks.append(m)
    myvars = [sky_ap]
    myvars.extend(masks)
    # xr.apply_ufunc seems to like stripping attributes from xds and its coordinates,
    # so make a deep copy of the pertinent thing to be handled in xr.apply_ufunc()
    nan_mask = xr.apply_ufunc(da.isnan, xds[sky_ap].copy(deep=True), dask="allowed")
    # test if nan pixels are all already masked in active_mask.
    # if so, don't worry about masking them again, just use active_mask
    # which already masks all nans
    arr_masks = {}
    do_mask_nans = False
    there_are_nans = nan_mask.any()
    if there_are_nans:
        has_active_mask = bool(active_mask)
        if not has_active_mask:
            do_mask_nans = True
        else:
            notted_active_mask = xr.apply_ufunc(
                da.logical_not, xds[active_mask].copy(deep=True), dask="allowed"
            )
            some_nans_are_not_already_masked = xr.apply_ufunc(
                da.logical_and, nan_mask, notted_active_mask, dask="allowed"
            )
            do_mask_nans = some_nans_are_not_already_masked.any()
    if do_mask_nans:
        mask_name = "mask_xds_nans"
        i = 0
        while mask_name in masks:
            mask_name = f"mask_xds_nans{i}"
            i += 1
        masks_rec[mask_name] = mask_rec
        masks_rec[mask_name][
            "mask"
        ] = f"Table: {os.sep.join([image_full_path, mask_name])}"
        masks.append(mask_name)
        arr_masks[mask_name] = nan_mask
        if active_mask:
            mask_name = f"mask_xds_nans_or_{active_mask}"
            i = 0
            while mask_name in masks:
                mask_name = f"mask_xds_nans{i}"
                i += 1
            masks_rec[mask_name] = mask_rec
            masks_rec[mask_name][
                "mask"
            ] = f"Table: {os.sep.join([image_full_path, mask_name])}"
            masks.append(mask_name)
            # This command strips attributes at various places for no
            # apparent reason, so make a copy of the xds to run it on
            arr_masks[mask_name] = xr.apply_ufunc(
                da.logical_or,
                xds[active_mask].copy(deep=True),
                nan_mask,
                dask="allowed",
            )
        active_mask = mask_name
    data_type = "complex" if "u" in xds.coords else "float"
    _write_initial_image(xds, image_full_path, active_mask, casa_image_shape[::-1])
    for v in myvars:
        _write_pixels(v, active_mask, image_full_path, xds)
    for name, v in arr_masks.items():
        _write_pixels(name, active_mask, image_full_path, xds, v)
    if masks:
        with open_table_rw(image_full_path) as tb:
            tb.putkeyword("masks", masks_rec)
            tb.putkeyword("Image_defaultmask", active_mask)


def _write_initial_image(
    xds: xr.Dataset, imagename: str, maskname: str, image_shape: tuple
) -> None:
    if not maskname:
        maskname = ""
    for dv in ["sky", "apeture"]:
        if dv in xds.data_vars:
            value = xds[dv][0, 0, 0, 0, 0].values.item()
            break
    image_full_path = os.path.expanduser(imagename)
    with _create_new_image(
        image_full_path, mask=maskname, shape=image_shape, value=value
    ) as casa_image:
        # just create the image, don't do anythong with it
        pass


def _write_image_block(xda: xr.DataArray, outfile: str, blc: tuple) -> None:
    """
    Write image xda chunk to the corresponding image table slice
    """
    # trigger the DAG for this chunk and return values while the table is
    # unlocked
    values = xda.compute().values
    with open_table_rw(outfile) as tb_tool:
        tb_tool.putcellslice(
            tb_tool.colnames()[0],
            0,
            values,
            blc,
            tuple(np.array(blc) + np.array(values.shape) - 1),
        )


def _write_pixels(
    v: str,
    active_mask: str,
    image_full_path: str,
    xds: xr.Dataset,
    value: xr.DataArray = None,
) -> None:
    flip = False
    if v == "sky" or v == "apeture":
        filename = image_full_path
    else:
        # mask
        flip = True
        filename = os.sep.join([image_full_path, v])
        # the default mask has already been written in _xds_to_casa_image()
        if not os.path.exists(filename):
            tb = tables.table(os.sep.join([image_full_path, active_mask]))
            tb.copy(filename, deep=True, valuecopy=True)
            tb.close()
    if "l" in xds.coords:
        trans_coords = ("frequency", "polarization", "m", "l")
    elif "u" in xds.coords:
        trans_coords = ("frequency", "polarization", "v", "u")
    else:
        raise RuntimeError(f"Unhandled coords combination {xds.coords.keys()}")
    trans_coord = ("frequency", "polarization", "m", "l")
    arr = xds[v] if v in xds.data_vars else value
    arr = arr.isel(time=0).transpose(*trans_coords)
    chunk_bounds = arr.chunks
    b = [0, 0, 0, 0]
    loc0, loc1, loc2, loc3 = (0, 0, 0, 0)
    for i0 in chunk_bounds[0]:
        b[0] = loc0
        s0 = slice(b[0], b[0] + i0)
        loc1 = 0
        for i1 in chunk_bounds[1]:
            b[1] = loc1
            s1 = slice(b[1], b[1] + i1)
            loc2 = 0
            for i2 in chunk_bounds[2]:
                b[2] = loc2
                s2 = slice(b[2], b[2] + i2)
                loc3 = 0
                for i3 in chunk_bounds[3]:
                    b[3] = loc3
                    blc = tuple(b)
                    s3 = slice(b[3], b[3] + i3)
                    sub_arr = arr[s0, s1, s2, s3]
                    if flip:
                        sub_arr = np.logical_not(sub_arr)
                    _write_image_block(sub_arr, filename, blc)
                    loc3 += i3
                loc2 += i2
            loc1 += i1
        loc0 += i0
