import numpy as np
from astropy.io import fits
from astropy.wcs import WCS


def flatten(f, x, y, size, hduid=0, channel=0, freqaxis=3, verbose=True):
    """
    Flatten a fits file so that it becomes a 2D image. Return new header and
    data
    This version also makes a sub-image of specified size.
    """

    naxis = f[hduid].header['NAXIS']
    if naxis < 2:
        raise RuntimeError('Can\'t make map from this')

    if verbose:
        print(f[hduid].data.shape)
    data_shape = f[hduid].data.shape[-2:]
    by, bx = data_shape
    xmin = int(x - size / 2)
    if xmin < 0:
        xmin = 0
    xmax = int(x + size / 2)
    if xmax > bx:
        xmax = bx
    ymin = int(y - size / 2)
    if ymin < 0:
        ymin = 0
    ymax = int(y + size / 2)
    if ymax > by:
        ymax = by

    if ymax <= ymin or xmax <= xmin:
        # this can only happen if the required position is not on the map
        print(xmin, xmax, ymin, ymax)
        raise RuntimeError('Failed to make subimage! Required position not on the map.')

    w = WCS(f[hduid].header)
    wn = WCS(naxis=2)

    wn.wcs.crpix[0] = w.wcs.crpix[0] - xmin
    wn.wcs.crpix[1] = w.wcs.crpix[1] - ymin
    wn.wcs.cdelt = w.wcs.cdelt[0:2]
    try:
        wn.wcs.pc = w.wcs.pc[0:2, 0:2]
    except AttributeError:
        pass  # pc is not present
    wn.wcs.crval = w.wcs.crval[0:2]
    wn.wcs.ctype[0] = w.wcs.ctype[0]
    wn.wcs.ctype[1] = w.wcs.ctype[1]

    header = wn.to_header()
    header["NAXIS"] = 2

    slice = []
    for i in range(naxis, 0, -1):
        if i == 1:
            slice.append(np.s_[xmin:xmax])
        elif i == 2:
            slice.append(np.s_[ymin:ymax])
        elif i == freqaxis:
            slice.append(channel)
        else:
            slice.append(0)
    if verbose:
        print(slice)

    hdu = fits.PrimaryHDU(f[hduid].data[tuple(slice)], header)
    copy = ('EQUINOX', 'EPOCH', 'BMAJ', 'BMIN', 'BPA')
    for k in copy:
        r = f[hduid].header.get(k)
        if r:
            hdu.header[k] = r
    if 'TAN' in hdu.header['CTYPE1']:
        hdu.header['LATPOLE'] = f[hduid].header['CRVAL2']
    hdulist = fits.HDUList([hdu])
    return hdulist


def extract_subimage(filename, ra, dec, size, hduid=0, verbose=True):
    if verbose:
        print('Opening', filename)
    orighdu = fits.open(filename)
    psize = int(size / orighdu[hduid].header['CDELT2'])  # size in pixels
    if verbose:
        print(f"Size in Pixels: {psize}")
        print(f"Size in Arcseconds: {size}")

    ndims = orighdu[hduid].header['NAXIS']
    pvect = np.zeros((1, ndims))
    lwcs = WCS(orighdu[hduid].header)
    pvect[0][0] = ra
    pvect[0][1] = dec
    imc = lwcs.wcs_world2pix(pvect, 0)
    x = imc[0][0]
    y = imc[0][1]
    hdu = flatten(orighdu, x, y, psize, hduid=hduid, verbose=verbose)
    return hdu
