#!/usr/bin/env python3
"""
==========================
Waterfall diagnostic plots
==========================
Create dynamic spectra or "waterfall" plots from the calibrated measurements
generated by the VLA CASA calibration pipeline using the CNGI IO tools.

:authors: <NAME>, <NAME>
"""
# TODO
# * Plot per SPW
# * Plot all frequency together

import os
import uuid
from pathlib import Path

import xarray
import numpy as np
from matplotlib import pyplot as plt

from . import (PATHS, savefig)
from . import cngi_io


OFF_FIELD = "3C48_OFF"
X_BAND_SPWS = "18~33"
U_BAND_SPWS = "2~17"


PIPE_FMT = "pipe_{0}"
MS_FILES_BY_DATE = {
        "030422": "TRFI0004_sb41515547_1_1.59641.79001768518.ms",
        "031422": "TRFI0004_sb41267655_1_1.59652.91350653935.ms",
        "032822": "TRFI0004_sb41267655_1_1.59666.90680377315.ms",
        "032922": "TRFI0004.sb41721231.eb41745688.59667.694925393516.ms",
}
ALL_DATES = list(MS_FILES_BY_DATE.keys())


def get_listobs_text(vis_filen):
    from casatasks import listobs
    assert Path(vis_filen).exists()
    tmp_filen = f"/dev/shm/listobs-{uuid.uuid1()}.txt"
    listobs(vis=vis_filen, listfile=tmp_filen, verbose=True)
    with open(tmp_filen, "r") as f:
        text = f.read()
    os.remove(tmp_filen)
    return text


def read_ms_data(vis_filen, band="X"):
    """
    Read a measurement set into an xarray dataset.

    Parameters
    ----------
    vis_filen : str
    band : str
    """
    assert Path(vis_filen).exists()
    if band == "X":
        rowmap = cngi_io.ms_selection(
                vis_filen, outfile=None, field=OFF_FIELD, spw=X_BAND_SPWS,
        )
    elif band == "U":
        rowmap = cngi_io.ms_selection(
                vis_filen, outfile=None, field=OFF_FIELD, spw=U_BAND_SPWS,
        )
    else:
        raise ValueError(f"Invalid band: {band}")
    mxds = cngi_io.read_ms(
            vis_filen,
            rowmap=rowmap,
            subtables=False,
            expand=True,
            chunks=(1000, 120, 2),
    )
    return mxds


def read_ms_from_pipe(date, **kwargs):
    ms_filen = MS_FILES_BY_DATE[date]
    vis_path = PATHS.data / f"pipe_{date}" / ms_filen
    return read_ms_data(str(vis_path), **kwargs)


def concat_by_spw(mxds):
    # Clip a few metadata attributes at the end.
    spwlist = list(mxds.attrs.keys())[:-3]
    xds_by_spw = [mxds.attrs[k] for k in spwlist]
    # Concatenate all SPWs together
    return xarray.concat(xds_by_spw, dim="chan", data_vars="all")


def calc_max_cross(xds):
    data = xds["CORRECTED_DATA"]
    xds["DMAG"] = (
            ((data.real**2 + data.imag**2)**0.5)
            .mean(axis=3)  # average across polarizations
            .max(axis=1)   # max across baselines
    )**(0.01)              # apply power scaling
    chan_ghz = xds["chan"] / 1e9  # Hz to GHz
    return (xds
            .assign_coords({"chan_GHz": chan_ghz})
            .sortby("chan")
    )


def plot_waterfall(xds, outname=None):
    assert "DMAG" in xds
    if outname is None:
        outname = "waterfall"
    fig, ax = plt.subplots(ncols=1, figsize=(8,4))
    xds["DMAG"].plot.pcolormesh(ax=ax, x="chan_GHz", y="time")
    #ax.set_xlabel()
    #ax.set_ylabel()
    savefig(outname)


def plot_all_waterfall(mxds):
    mxds


