'''
beam_volume_tools.py
'''

import numpy as np
from scipy import ndimage, signal
from os import path

from spectral_cube import SpectralCube
from spectral_cube.spectral_cube import BaseSpectralCube
from astropy import units as u
from radio_beam import Beam, Beams
from astropy.convolution import convolve_fft #, convolve
from astropy.io import fits
from astropy import log


def measure_epsilon_from_psf(psf, beam, pixels_per_beam, max_npix_peak=100):
    if psf.max() <= 0:
        raise ValueError("Invalid PSF")
    center = np.unravel_index(np.argmax(psf), psf.shape)
    cy, cx = center

    cutout = psf[cy-max_npix_peak:cy+max_npix_peak+1, cx-max_npix_peak:cx+max_npix_peak+1]
    shape = cutout.shape
    sy, sx = shape
    Y, X = np.mgrid[0:sy, 0:sx]

    center = np.unravel_index(np.argmax(cutout), cutout.shape)
    cy, cx = center

    dy = (Y - cy)
    dx = (X - cx)
    # I guess these definitions already take into account the definition of PA (east from north)?
    costh = np.cos(beam.pa.to('rad'))
    sinth = np.sin(beam.pa.to('rad'))
    # Changed variable name to rminmaj (it was rmajmin)
    rminmaj =  beam.minor / beam.major

    rr = ((dx * costh + dy * sinth)**2 / rminmaj**2 +
          (dx * sinth - dy * costh)**2 / 1**2)**0.5
    rbin = (rr).astype(int)

    #From plots taking the abs looks better centered by ~ 1 pix.
    #radial_mean = ndimage.mean(cutout**2, labels=rbin, index=np.arange(max_npix_peak))
    radial_mean = ndimage.mean(np.abs(cutout), labels=rbin, index=np.arange(max_npix_peak))
    first_min_ind = signal.find_peaks(-radial_mean)[0][0]

    #cutout_posit = np.where(cutout > 0, cutout, 0.)
    radial_sum = ndimage.sum(cutout, labels=rbin, index=np.arange(first_min_ind))
    psf_sum = np.sum(radial_sum)

    clean_psf_sum = pixels_per_beam
    epsilon = clean_psf_sum/psf_sum
    log.debug(f'clean_psf_sum={clean_psf_sum}, psf_sum={psf_sum}, epsilon={epsilon}')

    return epsilon, clean_psf_sum, psf_sum


def epsilon_from_psf(psf_image, max_npix_peak=100, export_clean_beam=True,
                     verbose=False, beam_threshold=0.1, pbar=False, **kwargs):
    """
    Determine epsilon, the ratio of the clean beam volume to the dirty beam volume within the first null, for a cube's PSFs.

    Parameters
    ----------
    psf_image : casa image
        A cube image of PSFs made by CASA
    max_npix_peak : int
        The maximum separation to integrate within to estimate the beam
    export_clean_beam : bool
        Return the synthesized beam in addition to the epsilon values?
    kwargs :
        passed to `common_beam`
    """

    #if not path.exists(psf_image):
    #    print('CASA image file of PSF not found.')

    max_npix_peak = int(max_npix_peak)


    if isinstance(psf_image, BaseSpectralCube):
        psf = psf_image
    elif not hasattr(psf_image, 'beams') and not hasattr(psf_image, 'beam'):
        # assume it was a string, try to read it
        psf = SpectralCube.read(psf_image, format='casa_image')
    else:
        raise ValueError

    good_beams = psf.identify_bad_beams(beam_threshold)

    if hasattr(psf, 'beam'):
        common_beam = psf.beam
    else:
        common_beam = psf.beams[good_beams].common_beam(**kwargs)

    epsilon_arr = np.zeros(len(psf))

    if not pbar:
        pbar = lambda x: x

    for chan in pbar(range(len(psf))):

        if psf[chan].max() == 0:
            print("INVALID PSF for channel {chan}")
            epsilon_arr[chan] = 0

        epsilon, clean_psf_sum, psf_sum = measure_epsilon_from_psf(psf[chan],
                                                                   psf.beams[chan],
                                                                   psf.pixels_per_beam[chan],
                                                                   max_npix_peak)
        epsilon_arr[chan] = epsilon

        if verbose:
            print('\n')
            print('Clean beam area of channel {0} is {1} pixels:'.format(chan, clean_psf_sum))
            print('Dirty beam area of channel {0} is {1} pixels:'.format(chan, psf_sum))
            print('epsilon = Omega_clean / Omega_dirty = {}'.format(epsilon))

    if export_clean_beam:
        output = {'epsilon': epsilon_arr, 'clean_beam': common_beam}
    else:
        output = {'epsilon': epsilon_arr}
    return output


def conv_model(model_image, clean_beam, save_to_tmp_dir=False):
    if isinstance(model_image, BaseSpectralCube):
        model = model_image
    else:
        model = SpectralCube.read(model_image, format='casa_image')
    beam = clean_beam

    pix_scale = model.header['CDELT2']*u.deg
    pix_scale = pix_scale.to(u.arcsec)
    clean_beam_kernel = beam.as_kernel(pix_scale)

    omega_beam = beam.sr
    omega_pix = pix_scale.to('rad')**2
    npix_beam = (omega_beam/omega_pix).value

    # should we just use a delta function rather than try to hack correct pixel area?
    # alternately, we could deconvolve a pixel size.
    # What is technically correct?
    # What does CASA do?  (scary question)
    fwhm_gauss_pix = (4*np.log(2)/np.pi)**0.5 * pix_scale
    pix_beam = Beam(fwhm_gauss_pix, fwhm_gauss_pix, 0*u.deg)
    model = model.with_beam(pix_beam)

    conv = model.convolve_to(beam, save_to_tmp_dir=save_to_tmp_dir) * npix_beam

    return conv


def rescale(conv_model, epsilon, residual_image, savename=None,
            export_fits=True):
    if isinstance(residual_image, BaseSpectralCube):
        residual = residual_image
        if savename is None and export_fits:
            raise ValueError("Must specify savename if exporting")
    else:
        residual = SpectralCube.read(residual_image, format='casa_image')
        if savename is None:
            savename = residual_image.replace(".residual",
                                              ".image.rescaled.fits")

    header = conv_model.header

    epsilon = epsilon*u.dimensionless_unscaled

    restor = conv_model.unitless + residual*epsilon[:,None,None]

    if export_fits:
        print("Writing")
        restor.write(savename, overwrite=True)
        print("Done writing")

    return restor
