"""
Script for displaying continuum images, deprojecting, and azimuthally averaging
This script is written for Python 3.6 (but should work in Python 2.7) 
"""

import numpy as np
from astropy.io import fits
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
from matplotlib.patches import Ellipse
from scipy.ndimage.interpolation import shift
import copy
import matplotlib as mpl
import matplotlib.colors as col
import yaml
from skimage import transform, measure
from scipy import ndimage

mpl.rcParams.update({'font.size':9})


class Continuum:
    def __init__(self, filename, offsetx, offsety, PA, incl, src_distance):
        """
        filename: fits file name
        offsetx: offset from phase center in x-direction in arcsec (positive is eastward). Note that the phase center occurs at pixelimsize/2+1 
        offsety: offset from phase center in y-directino in arcsec (positive is northward)
        PA: position angle, defined as angle that semimajor axis makes with north (deg)
        incl: inclination angle in degrees (0 is face-on, 90 is edge on)
        src_distance: source distance in pc
        """
        hdulist = fits.open(filename)
        #we assume here that the image is square
        self.image = np.squeeze(hdulist[0].data).astype(np.float64) #stored as (y,x)
        assert self.image.shape[0]==self.image.shape[1]
        self.image = np.nan_to_num(np.squeeze(hdulist[0].data))
        self.image[np.where(np.abs(self.image)>1.e6)]=0
        self.header = hdulist[0].header
        self.bmaj = self.header['BMAJ']*3600 #major axis of synthesized beam in arcsec
        self.bmin = self.header['BMIN']*3600 #minor axis of synthesized beam in arcsec
        self.bpa = self.header['BPA'] #position angle of synthesized beam in degrees
        self.freq = self.header['CRVAL3'] #continuum frequency in Hz
        self.imsize = self.header['NAXIS1']#image size along one dimension in pixels
        self.delt_x = self.header['CDELT1']*3600 #size of pixel in x direction in arcsec (this is negative because RA is positive going left)
        self.delt_y =self.header['CDELT2']*3600 #size of pixel in y direction in arcsec
        hdulist.close()


        self.offsetx = float(offsetx)
        self.offsety = float(offsety)
        self.PA = float(PA)
        self.incl = float(incl)
        self.src_distance = float(src_distance)

        RAcosDEC = self.delt_x*(np.arange(self.imsize)-(self.header['CRPIX1']-1.0))-self.offsetx #For an even number of pixels, the phase center corresponds to pixel = imsize/2+1

        DEC = self.delt_y*(np.arange(self.imsize)-(self.header['CRPIX2']-1.0))-self.offsety

        # Define a coordinate system based on the geometry

        self.xx, self.yy = np.meshgrid(RAcosDEC, DEC) # indexing is y,x 
        inclr = np.radians(self.incl)
        PAr = np.radians(self.PA)

        self.x_prime = (self.xx*np.cos(PAr) - self.yy*np.sin(PAr)) / np.cos(inclr) #this does the rotation so that the major axis of the disk image is aligned with the y-axis
        self.y_prime = (self.xx*np.sin(PAr) + self.yy*np.cos(PAr))

        # convert into polar coordinates (in physical dimensions)
        self.r = self.src_distance * np.sqrt(self.x_prime**2 + self.y_prime**2) #in pc
        self.theta = np.degrees(np.arctan2(self.y_prime, self.x_prime)) #90 deg corresponds to the major axis, angles increase in the clockwise direction when viewing the image (because of the orientation of the horizontal axis)

    def get_TB(self, Ivals):
        #convert Ivals from Jy/beam to watts/ster
        # a jansky is 10^(-26) watts per square meter per hertz
        #this is all in mks units
        h = 6.62607004e-34 #m^2 kg/s
        c = 299792458 #m/s
        k = 1.38064852e-23 #m^2 kg s^-2 K^-1
        beamsize = np.pi*self.bmin*self.bmaj/(4*np.log(2)) #arcsec^2
        B_nu = Ivals*1.e-26*1/beamsize*(206264.806**2)
        return h*self.freq/(k*np.log(2*h*self.freq**3/c**2*1/B_nu+1.))

    def radialprofile(self,radialbins,yaxis = 'intensity', theta_exclusion = np.array([]), high_incl = False): 
        """
        Parameters
        ==========
        """
        rbins = radialbins
        rwidth = (rbins[1]-rbins[0])

        radial_profile = np.zeros( len(rbins) )
        rad_prof_scatter = np.zeros_like(radial_profile)

        for i in range(len(rbins)):
            #for azimuthal asymmetries
            if len(theta_exclusion)==2:
                annulus_intensities = self.image[(self.r>=(radialbins[i]-0.5*rwidth)) & (self.r<(radialbins[i]+0.5*rwidth))&((self.theta<=theta_exclusion[0]) | (self.theta>=theta_exclusion[1]))]
            #for high inclination disks 
            elif high_incl:
                annulus_intensities = self.image[(self.r>=(rbins[i]-0.5*rwidth)) & (self.r<(rbins[i]+0.5*rwidth))&(np.abs(self.theta)<=110)&(np.abs(self.theta)>=70)]
            else:
                annulus_intensities = self.image[(self.r>=(rbins[i]-0.5*rwidth)) & (self.r<(rbins[i]+0.5*rwidth))]

            radial_profile[i] = np.average(annulus_intensities)
            rad_prof_scatter[i] = np.std(annulus_intensities)

        if yaxis == 'temperature': #return brightness temperature (K)
            return self.get_TB(radial_profile)
        elif yaxis == 'intensity': #return intensity profile (mJy/beam)
            return radial_profile, rad_prof_scatter

    def cumefluxprofile(self, radialbins):

        rbins = radialbins
        rwidth = (rbins[1]-rbins[0])

        flux_bin = np.zeros( len(rbins) )
 
        for i in range(len(rbins)):
            annulus_intensities = self.image[(self.r>=(rbins[i]-0.5*rwidth)) & (self.r<(rbins[i]+0.5*rwidth))]

            flux_bin[i] = np.sum(annulus_intensities)

        beamsize = np.pi*self.bmin*self.bmaj/(4*np.log(2))
        
        return np.cumsum(flux_bin)*(self.delt_y)**2/beamsize

    def radialcontrast(self,radialbins): 
        """
        Parameters
        ==========
        """
        rbins = radialbins
        rwidth = (rbins[1]-rbins[0])

        radial_contrast = np.zeros( len(rbins) )

        for i in range(len(rbins)):
            annulus_intensities = self.image[(self.r>=(rbins[i]-0.5*rwidth)) & (self.r<(rbins[i]+0.5*rwidth))]

            radial_contrast[i] = np.max(annulus_intensities)/np.min(annulus_intensities)

        return radial_contrast

    def profilepercentile(self,radialbins,percentile = 50): 
        """
        Parameters
        ==========
        """
        rbins = radialbins
        rwidth = (rbins[1]-rbins[0])

        radial_profile = np.zeros( len(rbins) )


        for i in range(len(rbins)):
            annulus_intensities = self.image[(self.r>=(rbins[i]-0.5*rwidth)) & (self.r<(rbins[i]+0.5*rwidth))]

            radial_profile[i] = np.percentile(annulus_intensities, percentile)

        return radial_profile

    def profilemin(self,radialbins): 
        """
        Parameters
        ==========
        """
        rbins = radialbins
        rwidth = (rbins[1]-rbins[0])

        radial_profile = np.zeros( len(rbins) )


        for i in range(len(rbins)):
            annulus_intensities = self.image[(self.r>=(rbins[i]-0.5*rwidth)) & (self.r<(rbins[i]+0.5*rwidth))]

            radial_profile[i] = np.min(annulus_intensities)

        return radial_profile

    def azunwrap(self,radialbins, tbins = -179.+2*np.arange(180), yaxis = 'intensity', theta_exclusion = np.array([])): 
        """
        Parameters
        ==========
        """
        rbins = np.copy(radialbins)
        rwidth = (rbins[1]-rbins[0])
        twidth = np.abs((tbins[1]-tbins[0]))

        # create a binned (r, theta) "map" 
        rtmap = np.zeros( (len(tbins), len(rbins)) )
        radial_profile = np.zeros( len(rbins) )
        rad_prof_scatter = np.zeros_like(radial_profile)

        for i in range(len(rbins)):
            annulus_thetas = self.theta[(self.r>=(rbins[i]-0.5*rwidth)) & (self.r<(rbins[i]+0.5*rwidth))]
            if len(theta_exclusion)==2:
                annulus_intensities = self.image[(self.r>=(radialbins[i]-0.5*rwidth)) & (self.r<(radialbins[i]+0.5*rwidth))&((self.theta<=theta_exclusion[0]) | (self.theta>=theta_exclusion[1]))]
            else:
                annulus_intensities = self.image[(self.r>=(rbins[i]-0.5*rwidth)) & (self.r<(rbins[i]+0.5*rwidth))]

            radial_profile[i] = np.average(annulus_intensities)
            rad_prof_scatter[i] = np.std(annulus_intensities)
            if len(theta_exclusion)==2:
                annulus_intensities = self.image[(self.r>=(rbins[i]-0.5*rwidth)) & (self.r<(rbins[i]+0.5*rwidth))]
            for j in range(len(tbins)):
                wedge = annulus_intensities[(annulus_thetas>=(tbins[j]-0.5*twidth)) & (annulus_thetas<(tbins[j]+0.5*twidth))]
                if len(wedge) > 0:
                    rtmap[j, i] = np.average(wedge)
                else: 
                    rtmap[j, i] = -1e10
        
        # "fix" the inner disk where there's too few azimuthal bins (linear interpolation)
        for i in range(len(rbins)):
            tslice = rtmap[:, i]
            if np.any(tslice <= -1e5):
                xslc = tbins[tslice >= -1e5]
                yslc = tslice[tslice >= -1e5]
                #we need the interpolating values to make at least a full 360 so that we can get the whole circle
                xslc_extended = np.pad(xslc, 1, mode = 'wrap')
                xslc_extended[0]-=360.
                xslc_extended[-1]+=360
                yslc_extended = np.pad(yslc, 1, mode = 'wrap')

                rzfunc = interp1d(xslc_extended, yslc_extended, bounds_error=True)
                fix_slice = rzfunc(tbins)
                rtmap[:, i] = fix_slice

        if yaxis == 'temperature': #return brightness temperature (K)
            return self.get_TB(radial_profile), rtmap
        elif yaxis == 'intensity': #return intensity profile (mJy/beam)
            return radial_profile, rad_prof_scatter, rtmap
        else:
            return rtmap

    def beam_plot(self, x, y, ax, flip = False, color = 'white'):
        """Retrieves parameters for drawing the synthesized beam"""
        beam = Ellipse(xy=(x, y), width=self.bmin, height=self.bmaj, angle=self.bpa, facecolor=color, edgecolor = color)
        if flip:
            #this option is for if the axes are labeled with the RA and DEC, since the image gets flipped around
            beam = Ellipse(xy=(x, y), width=self.bmin, height=self.bmaj, angle=-self.bpa, facecolor=color, edgecolor = color)
        ax.add_artist(beam) 
        
    def circularize(self):
        shifted = ndimage.shift(self.image.T, (-self.offsetx/self.delt_x+0.5,-self.offsety/self.delt_y+0.5))
        rotated = transform.rotate(shifted, 180-self.PA) #ccw
        rescaled = transform.rescale(rotated, (1/np.cos(self.incl*np.pi/180.),1), preserve_range = False)
        return rescaled

    def extract_ring_old(self,rmin,rmax,tbins = -179.5+np.arange(360), twidth = 1, extract_type = 'max', return_rvals = False): 
        """
        Extracts pixels corresponding to local maxima 

        Parameters
        ==========
        rmin: Inner radius of annulus over which local maxima/minima are being identified
        rmax: Outer radius of annulus over which local maxima/minima are being identified
        tbins: Array denoting the centers of the azimuthal bins in which the local max/min is being identified(in degrees)
        twidth: Width (in degrees) of each azimuthal search bin (should be less than or equal to the separation between the bin centers)
        extract_type: 'max' yields local maxima, 'min' yields local minima

        Returns
        ======
        xcoord, ycoord: Arrays of coordinates (in arcsec) of the pixels correponding to the local maxima/minima
        """
        
        xcoord = []
        ycoord = []
        rvals = []
        assert twidth<=np.min(np.abs(np.diff(tbins)))
        annulus_intensities = self.image[(self.r>rmin) & (self.r<rmax)]
        annulus_thetas = self.theta[(self.r>rmin) & (self.r<rmax)]
        for j in range(len(tbins)):
            wedge = annulus_intensities[(annulus_thetas>=(tbins[j]-0.5*twidth)) & (annulus_thetas<(tbins[j]+0.5*twidth))]
            if len(wedge) > 0.:
                if extract_type=='max':
                    index = np.where(wedge==np.max(wedge))[0][0]
                elif extract_type=='min':
                    index = np.where(wedge==np.min(wedge))[0][0]
                xcoord.append(self.xx[(self.r>rmin) & (self.r<rmax)][(annulus_thetas>=(tbins[j]-0.5*twidth)) & (annulus_thetas<(tbins[j]+0.5*twidth))][index])
                ycoord.append(self.yy[(self.r>rmin) & (self.r<rmax)][(annulus_thetas>=(tbins[j]-0.5*twidth)) & (annulus_thetas<(tbins[j]+0.5*twidth))][index])
                if return_rvals:
                     rvals.append(self.r[(self.r>rmin) & (self.r<rmax)][(annulus_thetas>=(tbins[j]-0.5*twidth)) & (annulus_thetas<(tbins[j]+0.5*twidth))][index])
        if return_rvals:
            return np.array(xcoord)+self.offsetx, np.array(ycoord)+self.offsety, np.array(rvals)
             

        else:
            return np.array(xcoord)+self.offsetx, np.array(ycoord)+self.offsety

    def extract_ring(self,rmin,rmax,tbins = -179.5+np.arange(360), extract_type = 'max'): 
        """
        Extracts pixels corresponding to local maxima 

        Parameters
        ==========
        rmin: Inner radius of annulus over which local maxima/minima are being identified
        rmax: Outer radius of annulus over which local maxima/minima are being identified
        tbins: Array denoting the centers of the azimuthal bins in which the local max/min is being identified(in degrees)
        twidth: Width (in degrees) of each azimuthal search bin (should be less than or equal to the separation between the bin centers)
        extract_type: 'max' yields local maxima, 'min' yields local minima

        Returns
        ======
        xcoord, ycoord: Arrays of coordinates (in arcsec) of the pixels correponding to the local maxima/minima
        """
        xcoord = []
        ycoord = []


        xdiff = self.x_prime[0,0]-self.x_prime[1,0]
        ydiff = self.y_prime[0,0]-self.y_prime[0,1]
        tbins=tbins*np.pi/180.
        for j in range(len(tbins)):
            xpstart = rmin*np.cos(tbins[j])/self.src_distance*np.cos(self.incl*np.pi/180.)
            ypstart = rmin*np.sin(tbins[j])/self.src_distance

            PAr = self.PA*np.pi/180.
            xstart = np.int((xpstart*np.cos(-PAr)-ypstart*np.sin(-PAr)+self.offsetx)/self.delt_x+self.imsize/2)
            ystart = np.int((xpstart*np.sin(-PAr)+ypstart*np.cos(-PAr)+self.offsety)/self.delt_y+self.imsize/2)

        
            xpend = rmax*np.cos(tbins[j])/self.src_distance*np.cos(self.incl*np.pi/180.)
            ypend = rmax*np.sin(tbins[j])/self.src_distance
            xend = np.int((xpend*np.cos(-PAr)-ypend*np.sin(-PAr)+self.offsetx)/self.delt_x+self.imsize/2)
            yend = np.int((xpend*np.sin(-PAr)+ypend*np.cos(-PAr)+self.offsety)/self.delt_y+self.imsize/2)
            profile = measure.profile_line(self.image, (ystart,xstart), (yend, xend), order= 0)
            if extract_type=='max':
                index = np.where(profile==np.max(profile))[0][0]
            elif extract_type=='min':
                index = np.where(profile==np.min(profile))[0][0]
            if index!=0 and index!=len(profile)-1:
                loc = np.where(self.image[(self.r>rmin) & (self.r<rmax)]==profile[index])


                xcoord.append(self.xx[(self.r>rmin) & (self.r<rmax)][loc[0][0]])
                ycoord.append(self.yy[(self.r>rmin) & (self.r<rmax)][loc[0][0]])
            
        return np.array(xcoord)+self.offsetx, np.array(ycoord)+self.offsety

    def plot_extracted_ring(self,cont,xcoords, ycoords, gamma, size, cmap = 'magma'):
        f = plt.figure(figsize=(7,7))
        LBmap = plt.subplot(111, aspect = 'equal', adjustable = 'box')
        self.plot_cont_intensity(LBmap, size, cmap = cmap, vmin = 2.e-2, norm = col.PowerNorm(gamma = gamma))
        plt.scatter(xcoords-self.offsetx, ycoords-self.offsety, s = 1, color = 'white')
        return LBmap

    def subtractprofile(self, radialbins, percentile = 50):
        profile = self.profilepercentile(radialbins,percentile)
        extended_prof = np.insert(profile, 0, np.max(self.image))
        extended_bins = np.insert(radialbins, 0,0)
        radinterp = interp1d(extended_bins, extended_prof, bounds_error = False, fill_value = 0)
        return self.image - radinterp(self.r)

    def contours_only(self, ax, size, levels, colors='pink', flip=False):

        im = self.image
        imshifted = shift(im, np.array([-self.offsety/self.delt_y,-self.offsetx/self.delt_x]))

        imgdata = imshifted*1000 #convert to mJy/beam if this is an intensity image

        size_pix = size/np.abs(self.delt_x) #size of zoomed-in image in pixels

        padding = int(0.5*(self.imsize-size_pix))

       
        zoomimgdata = imgdata[padding:self.imsize - padding, padding:self.imsize - padding]
        if flip:
            zoomimgdata = np.fliplr(zoomimgdata) #since reversing the x-axis later will flip the image around, we have to flip it back
        cont = ax.contour(zoomimgdata,levels,linestyles='solid',colors=colors, extent=[-size/2., size/2., -size/2., size/2.], linewidths = .5)

        #plt.gca().invert_xaxis() #since RA increases right to left 


    
    def plot_cont_intensity(self, ax, size, im = None, cmap ='magma', vmin =None, vmax = None, alpha = False, gamma = 1.0, levels = np.array([]), labels_off = False, show_beam = True, norm = None, quantity='intensity'):
        """Plots continuum image"""        
        
        if im is None:
            im = self.image
        imshifted = shift(im, np.array([-self.offsety/self.delt_y,-self.offsetx/self.delt_x]))

        if quantity=='intensity':
            imgdata = imshifted*1000 #convert to mJy/beam if this is an intensity image
        elif quantity=='temperature':
            imgdata = np.nan_to_num(self.get_TB(imshifted))
        size_pix = size/np.abs(self.delt_x) #size of zoomed-in image in pixels

        padding = int(0.5*(self.imsize-size_pix))

        std =np.std(imgdata[int(0.2*padding):padding, int(0.2*padding):padding])
        n = np.arange(4,26)
        #levels = std*np.append([4,8],2**n)
        
        zoomimgdata = imgdata[padding:self.imsize - padding, padding:self.imsize - padding]
        
        zoomimgdata = np.fliplr(zoomimgdata) #since reversing the x-axis later will flip the image around, we have to flip it back
        if show_beam:
            self.beam_plot(0.4*size,-0.4*size,ax,flip =  True)

        if not vmin:
            vmin = 0
        if not vmax:
            vmax = np.max(zoomimgdata)

        if norm:
            finalimage = ax.imshow(zoomimgdata,origin='lower',cmap=cmap,norm = norm,  interpolation='None', extent=[-size/2., size/2., -size/2., size/2.])
        else:
            finalimage = ax.imshow(zoomimgdata,origin='lower',cmap=cmap,vmin = vmin, vmax=vmax, interpolation='None', extent=[-size/2., size/2., -size/2., size/2.])
        if len(levels)>0: 
            cont = ax.contour(zoomimgdata,levels,linestyles='solid',colors='gray', extent=[-size/2., size/2., -size/2., size/2.], linewidths = .75)

        
        if labels_off:
            ax.set_xticks([])
            ax.set_yticks([])
        else:
            ax.xaxis.set_ticks_position('both')
            ax.yaxis.set_ticks_position('both')
            if size >=5:
                ticks = np.arange(-np.floor(size/2.), np.floor(size/2.)+0.001, 2)
            elif 2<size<5:
                ticks = np.arange(-np.floor(size/2.), np.floor(size/2.)+0.001, 1)
            elif 1<=size<=2:
                ticks = np.arange(-np.floor(size)/2., np.floor(size)/2.+0.001, 0.5)
            else:
                ticks = np.arange(-np.floor(size*5), np.floor(size*5)+0.001, 2)/10.


            ax.tick_params(axis = "both", direction = 'in',which = 'major', colors = 'white', length = 3)
            ax.set_xticks(ticks)
            ax.set_yticks(ticks)

            for xtick in ax.get_xticklabels():
                xtick.set_color('black')
            for ytick in ax.get_yticklabels():
                ytick.set_color('black')



        plt.gca().invert_xaxis() #since RA increases right to left 
        return finalimage

    def plotbeamprofile(self, x,y, height, ax):
        bmin = self.bmin*self.src_distance
        sigma = bmin/(2*np.sqrt(2*np.log(2.)))
        xvals = np.linspace(x-15, x+15)
        ax.plot(xvals, y+height*np.exp(-(xvals-x)**2/(2*sigma**2)), color = 'gray')
                    
    def extract_spiral(self,residualimage, rmin,rmax,r_search_func,tbins = -177.5+5*np.arange(72), outfile = None): 
        ntheta = len(tbins)
        xcoord = []
        ycoord = []
        thetas = []
        rvals = []
        tbins=tbins*np.pi/180.
        for i in range(ntheta):
            r_lower, r_upper = r_search_func(tbins[i]*180/np.pi, rmin, rmax)

            xpstart = r_lower*np.cos(tbins[i])/self.src_distance*np.cos(self.incl*np.pi/180.)
            ypstart = r_lower*np.sin(tbins[i])/self.src_distance

            PAr = self.PA*np.pi/180.
            xstart = np.int((xpstart*np.cos(-PAr)-ypstart*np.sin(-PAr)+self.offsetx)/self.delt_x+self.imsize/2)
            ystart = np.int((xpstart*np.sin(-PAr)+ypstart*np.cos(-PAr)+self.offsety)/self.delt_y+self.imsize/2)

            xpend = r_upper*np.cos(tbins[i])/self.src_distance*np.cos(self.incl*np.pi/180.)
            ypend = r_upper*np.sin(tbins[i])/self.src_distance
            xend = np.int((xpend*np.cos(-PAr)-ypend*np.sin(-PAr)+self.offsetx)/self.delt_x+self.imsize/2)
            yend = np.int((xpend*np.sin(-PAr)+ypend*np.cos(-PAr)+self.offsety)/self.delt_y+self.imsize/2)
        

            profile = measure.profile_line(residualimage, (ystart,xstart), (yend, xend), order= 0)
            index = np.where(profile==np.max(profile))[0][0]
        
            if index!=0 and index!=len(profile)-1:
                loc = np.where(residualimage[(self.r>r_lower) & (self.r<r_upper)]==profile[index])
                xcoord.append(self.xx[(self.r>r_lower) & (self.r<r_upper)][loc[0][0]])
                ycoord.append(self.yy[(self.r>r_lower) & (self.r<r_upper)][loc[0][0]])
                rvals.append(self.r[(self.r>r_lower) & (self.r<r_upper)][loc[0][0]])
            
                thetas.append(tbins[i])
        xcoord = np.array(xcoord)
        ycoord = np.array(ycoord)
        thetas = np.array(thetas)
        rvals = np.array(rvals)
        if outfile:
            arr = np.row_stack((xcoord, ycoord, thetas, rvals)).T
            np.savetxt(outfile, arr)
        
        
        return xcoord, ycoord, thetas, rvals

