import math
from astropy.io import fits
import numpy as np
import math
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
from matplotlib.colors import LinearSegmentedColormap as LSC
from scipy.ndimage.interpolation import shift

c_kms = 2.99792458e5 # [km s^-1]

def read_SMA(fname):
    hdu_list = fits.open(fname)
    hdu = hdu_list[0]
    header = hdu.header

    vs = (np.arange(header["NAXIS3"])*header["CDELT3"] + header["CRVAL3"]) * 1e-3 # [km/s]

    data = hdu.data[0] # Get rid of the zombie first dimension.

    return (vs, data, header)

def read_ALMA(fname):
    "Return the beam parameters as well. Check the CASAMBM too."
    hdu_list = fits.open(fname)

    hdu = hdu_list[0]
    header = hdu.header

    # Check to see if there is CASAMBM keyword, and if it's true
    # If so, load the second header, which contains the beam info
    # Then, use the median beam, as determined by the major axis
    CASAMBM = header.get("CASAMBM", False)


    # Determine the frequencies.
    freq = header["CRVAL3"] + np.arange(header["NAXIS3"]) * header["CDELT3"]

    f0 = header["RESTFRQ"] # [Hz]

    # Convert all frequencies to velocity [km/s]
    vs = c_kms * (f0 - freq)/f0

    data = hdu.data[0] # Get rid of the zombie first dimension.


    if CASAMBM:
        # Get the beam info from the record array
        data2 = hdu_list[1].data
        BMAJ = np.median(data2["BMAJ"])
        BMIN = np.median(data2["BMIN"])
        BPA = np.median(data2["BPA"])
    else:
        # Get the beam info from the header, like normal
        BMAJ = 3600 * header["BMAJ"]
        BMIN = 3600 * header["BMIN"]
        BPA = header["BPA"]

    beam = (BMAJ, BMIN, BPA)

    return (vs, data, header, beam)


def get_coords(data, header, radius, mu_RA=0.0, mu_DEC=0.0):
    '''
    Given a FITS header object, return the RA and DEC vectors.

    Alternatively, also specify a truncation radius and this routine will return the slice indices.

    If mu_RA, mu_DEC specified, shift the image by this amount.
    '''
    nx = header["NAXIS1"]
    ny = header["NAXIS2"]
    nchan = header["NAXIS3"]

    assert nx % 2 == 0 and ny % 2 == 0 , "We don't have an even number of pixels, assumptions in the routine are violated."

    # Midpoints
    mx = int(nx/2)
    my = int(ny/2)

    # Phase center in decimal degrees
    RA = header["CRVAL1"]
    DEC = header["CRVAL2"]

    # from: http://tdc-www.harvard.edu/wcstools/wcstools.wcs.html
    # CRPIX1 and CRPIX2 are the pixel coordinates of the reference point to which
    # the projection and rotation refer, i.e, integers
    # Indexed from 1, I'm assuming
    # If the image is 256 x 256, CRPIX = 129, then it must be indexed from 1.

    # CRVAL1 and CRVAL2 give the center coordinate as right ascension or longitude
    # in decimal degrees.

    # CDELT1 and CDELT2 indicate the plate scale in degrees per pixel
    # This means that to get proper RA, DEC coordinates you would need to add a
    # cos(DEC) in there or use some WCS tools. But since we just care about offsets,
    # this should be fine

    xx = np.arange(header["NAXIS1"])
    yy = np.arange(header["NAXIS2"])

    # RA coordinates
    CDELT1 = header["CDELT1"]
    CRPIX1 = header["CRPIX1"] - 1. # Now indexed from 0

    # DEC coordinates
    CDELT2 = header["CDELT2"]
    CRPIX2 = header["CRPIX2"] - 1. # Now indexed from 0

    # Alternatively
    slice_RA = math.floor(math.fabs(radius/CDELT1))
    slice_DEC = math.floor(math.fabs(radius/CDELT2))

    assert slice_RA < mx and slice_DEC < my, "Selected radius must be smaller than half size of image."

    decl = mx-slice_DEC
    decr = mx+slice_DEC

    ral = mx-slice_RA