import sys
import numpy as np
from pylab import *
import lightkurve as lk
import matplotlib as mpl
import astropy.units as u
from astropy.wcs import WCS
from bs4 import BeautifulSoup
import warnings, os, requests
import matplotlib.pyplot as plt
from astroquery.mast import Tesscut
import matplotlib.gridspec as gridspec
from astropy.timeseries import LombScargle
from astropy.coordinates import SkyCoord, Angle
from mpl_toolkits.axes_grid1 import make_axes_locatable

from .ffi import use_pointing_model, load_pointing_model
from .mast import *

__all__ = []


class Visualize(object):
    """
    The main class for creating figures, movies, and interactive plots.
    Allows the user to have a grand ole time playing with their data!

    Parameters
    ----------
    obj :
        Object must have minimum attributes of 2D array of flux.
        Will allow for plotting of both postcards & tpfs.
    obj_type :
        Object type can be set to "tpf" or "postcard". Default is "tpf".
    """

    def __init__(self, object, obj_type="tpf"):
        self.obj      = object
        self.obj_type = obj_type.lower()

        if self.obj_type == "tpf":
            self.flux   = self.obj.tpf
            self.center = (np.nanmedian(self.obj.centroid_xs),
                             np.nanmedian(self.obj.centroid_ys))
            self.dimensions = self.obj.tpf[0].shape
        else:
            self.flux   = self.obj.flux
            self.center = self.obj.center_xy
            self.dimensions = self.obj.dimensions


    def aperture_contour(self, aperture=None, cadence=0, ap_color='w', ap_linewidth=4, ax=None, **kwargs):
        """
        Overplots the countour of an aperture on a target pixel file.
        Contribution from Gijs Mulders.

        Parameters
        ----------
        aperture : np.2darray, optional
            A 2D mask the same size as the target pixel file. Default
            is the eleanor default aperture.
        cadence : int, optional
            Index of the cadence in the TPF to plot.
            Default is 0 (the first cadence).
        ap_color : str, optional
            The color of the aperture contour. Takes a matplotlib color.
            Default is red.
        ap_linewidth : int, optional
            The linewidth of the aperture contour. Default is 4.
        ax : matplotlib.axes._subplots.AxesSubplot, optional
            Axes to plot on.
        """
        ax_given = ax
        if ax is None:
            fig, ax = plt.subplots(nrows=1)

        if aperture is None:
            aperture = self.obj.aperture

        ax.imshow(self.obj.tpf[cadence], origin='lower', **kwargs)

        f = lambda x,y: aperture[int(y),int(x) ]
        g = np.vectorize(f)

        x = np.linspace(0,aperture.shape[1], aperture.shape[1]*100)
        y = np.linspace(0,aperture.shape[0], aperture.shape[0]*100)
        X, Y= np.meshgrid(x[:-1],y[:-1])
        Z = g(X[:-1],Y[:-1])

        ax.contour(Z, [0.05], colors=ap_color, linewidths=[ap_linewidth],
                    extent=[0-0.5, x[:-1].max()-0.5,0-0.5, y[:-1].max()-0.5])

        if ax_given is None:
            return fig



    def pixel_by_pixel(self, colrange=None, rowrange=None, cmap='viridis',
                       data_type="corrected", mask=None, xlim=None,
                       ylim=None, color_by_pixel=False, color_by_aperture=True,
                       freq_range=[1/20., 1/0.1], view='frequency', FAP=None, aperture=None, ap_color='r',
                       ap_linewidth=2):
        """
        Creates a pixel-by-pixel light curve using the corrected flux.
        Contribution from Oliver Hall.

        Parameters
        ----------
        colrange : np.array, optional
             A list of start column and end column you're interested in
             zooming in on.
        rowrange : np.array, optional
             A list of start row and end row you're interested in zooming
             in on.
        cmap : str, optional
             Name of a matplotlib colormap. Default is 'viridis'.
        data_type : str, optional
             The type of flux used. Either: 'raw', 'corrected', 'amplitude',
             or 'periodogram'. If not, default set to 'corrected'.
        mask : np.array, optional
             Specifies the cadences used in the light curve. If not, default
             set to good quality cadences.
        xlim : np.array, optional
             Specifies the xlim on the subplots. If not, default is set to
             the entire light curve.
        ylim : np.array, optional
             Specifies the ylim on the subplots, If not, default is set to
             the entire light curve flux range.
        color_by_pixel : bool, optional
             Colors the light curve given the color of the pixel. If not,
             default is set to False.
        freq_range : list, optional
             List of minimum and maximum frequency to search in Lomb Scargle
             periodogram. Only used if data_type = 'periodogram'. If None,
             default = [1/20., 1/0.1].
        view : str, optional
            The type of view, i.e., x-axis units for periodogram visualization.
            Either 'frequency' or 'period'. If ‘frequency’, x-axis units will be
            frequency. If ‘period’, the x-axis units will be period.
            Only used if data_type = 'amplitude'. If None, default is 'frequency'.
        FAP: np.array, optional.
             False Alarm Probability levels to include in periodogram.
             Ensure that the values are < 1.
             For example: FAP = np.array([0.1, 0.01]), will plot the 10% and 1% FAP levels.
        """
        if self.obj.lite:
            print('This is an eleanor-lite object. No pixel_by_pixel visualization can be created.')
            print('Please create a regular eleanor.TargetData object (lite=False) to use this tool.')
            return

        if colrange is None:
            colrange = [0, self.dimensions[1]]

        if rowrange is None:
            rowrange = [0, self.dimensions[0]]

        nrows = int(np.round(colrange[1]-colrange[0]))
        ncols = int(np.round(rowrange[1]-rowrange[0]))

        if (colrange[1] > self.dimensions[1]) or (rowrange[1] > self.dimensions[0]):
            raise ValueError("Asking for more pixels than available in the TPF.")


        figure = plt.figure(figsize=(20,8))
        outer = gridspec.GridSpec(1,2, width_ratios=[1,4])

        inner = gridspec.GridSpecFromSubplotSpec(ncols, nrows, hspace=0.1, wspace=0.1,
                                                 subplot_spec=outer[1])

        # start from the highest row, to be consistent with the small pixel plot
        # using `ax.imshow(plotflux, origin='lower', ...)` below
        i, j = rowrange[1] - 1, colrange[0]

        if mask is None:
            q = self.obj.quality == 0
        else:
            q = mask == 0


        ## PLOTS TARGET PIXEL FILE ##

        ax = plt.subplot(outer[0])

        if aperture is None:
            aperture = self.obj.aperture

        plotflux = np.nanmedian(self.flux[:, rowrange[0]:rowrange[1],
                                          colrange[0]:colrange[1]], axis=0)
        c = ax.imshow(plotflux, origin='lower',
                      vmax=np.percentile(plotflux, 95),
                      cmap=cmap)
        divider = make_axes_locatable(ax)
        cax = divider.append_axes('right', size='5%', pad=0.15)
        plt.colorbar(c, cax=cax, orientation='vertical')

        f = lambda x,y: aperture[int(y),int(x) ]
        g = np.vectorize(f)

        x = np.linspace(colrange[0],colrange[1], nrows*100)
        y = np.linspace(rowrange[0],rowrange[1], ncols*100)
        X, Y= np.meshgrid(x[:-1],y[:-1])
        Z = g(X[:-1],Y[:-1])

        ax.contour(Z, [0.05],
                   colors=ap_color, linewidths=[ap_linewidth],
                    extent=[0-0.5, nrows-0.5,0-0.5, ncols-0.5])

        ## PLOTS PIXEL LIGHT CURVES ##
        for ind in range( int(nrows * ncols) ):
            if ind == 0:
                ax = plt.Subplot(figure, inner[ind])
                origax = ax
            else:
                ax = plt.Subplot(figure, inner[ind], sharex=origax)

            flux = self.flux[:,i,j]
            time = self.obj.time
            corr_flux = self.obj.corrected_flux(flux=flux)

            if data_type.lower() == 'corrected':
                y = corr_flux[q]/np.nanmedian(corr_flux[q])
                x = time[q]

            elif data_type.lower() == 'amplitude':
                lc = lk.LightCurve(time=time, flux=corr_flux)
                pg = lc.normalize().to_periodogram(minimum_frequency=freq_range[0],
                                                                     maximum_frequency=freq_range[1])

            elif data_type.lower() == 'raw':
                y = flux[q]/np.nanmedian(flux[q])
                x = time[q]

            elif data_type.lower() == 'periodogram':
                LS = LombScargle(time, corr_flux)
                freq, power = LS.autopower(minimum_frequency=freq_range[0],
                                                                     maximum_frequency=freq_range[1],
                                                                     method='fast')
                y = power
                x = 1/freq

                if (FAP is not None):
                    if np.all(FAP<1): # Ensure that the probabilities are all < 1
                        if type(FAP) == list: FAP = np.array(FAP)
                        FAPlevel = LS.false_alarm_level(FAP, method='baluev')

            if color_by_pixel is False:
                color = 'k'
            else:
                rgb = c.cmap(c.norm(self.flux[100,i,j]))
                color = matplotlib.colors.rgb2hex(rgb)

            if data_type.lower() == 'amplitude':
                pg.plot(view=view, scale='log', ax=ax, c=color)
            else:
                ax.plot(x, y, c=color)

            if (data_type.lower() == 'periodogram') & (FAP is not None):
                if np.all(FAP<1):
                    _ = [ax.axhline(f, color='k', ls='--', alpha=0.3) for f in FAPlevel]


            if color_by_aperture and aperture[i,j] > 0:
                for iax in ['top', 'bottom', 'left', 'right']:
                    ax.spines[iax].set_color(ap_color)
                    ax.spines[iax].set_linewidth(ap_linewidth)

            j += 1
            if j == colrange[1]:
                i -= 1
                j  = colrange[0]

            if data_type.lower() != 'amplitude':
                if ylim is None:
                    ax.set_ylim(np.percentile(y, 1), np.percentile(y, 99))
                else:
                    ax.set_ylim(ylim[0], ylim[1])

                if xlim is None:
                    ax.set_xlim(np.min(x)-0.1, np.max(x)+0.1)
                else:
                    ax.set_xlim(xlim[0], xlim[1])

            ax.xaxis.label.set_visible(False)
            ax.yaxis.label.set_visible(False)
            ax.set_xticks([])
            ax.set_yticks([])

            figure.add_subplot(ax)

        return figure


    def tess_the_movie(self):
        """
        Opens the link to Ethan Kruse's TESS: The Movie YouTube videos for
        the sector your target is observed in.

        Parameters
        ----------

        Attributes
        ----------
        movie_url : str

        """
        def type_of_script():
            try:
                ipy_str = str(type(get_ipython()))
                if 'zmqshell' in ipy_str:
                    return 'jupyter'
                if 'terminal' in ipy_str:
                    return 'ipython'
            except:
                return 'terminal'

        sector = self.obj.source_info.sector

        base="https://www.youtube.com/results?search_query="
        query="TESS+the+movie+sector+{0}+ethankruse".format(sector)

        soup = BeautifulSoup(requests.get(base+query).text, "html.parser").find_all('script')[26]

        items = soup.text
        items = items.split('\n')[1].split('title')

        good_sector=0

        for subitem in items:
            j = subitem.find('Sector')
            if j > 0 and 'TESS: The Movie' in subitem:

                sect = subitem[j:j+100].split(',')[0].split(' ')[-1]

                if int(sect) == int(sector):
                    i = subitem.find('/watch?v')
                    ext = subitem[i:i+100].split('"')[0]
                    good_sector=1
                    break

        if good_sector == 1:
            self.movie_url = 'https://www.youtube.com{0}'.format(ext)

            call_location = type_of_script()

            if (call_location == 'terminal') or (call_location == 'ipython'):
                os.system('python -m webbrowser -t "{0}"'.format(self.movie_url))

            elif (call_location == 'jupyter'):
                from IPython.display import YouTubeVideo
                id = self.movie_url.split('=')[-1]
                return YouTubeVideo(id=id, width=900, height=500)

        else:
            print('No movie is available yet.')
            return

    def download_tpf(self, tic=None):
        """
        Downloads the Target Pixel File (TPF) if None is passed in to
        plot_gaia_overlay or add_gaia_figure_elements.

        Utilizes astroquery.mast.TesscutClass.download_cutouts() instead of lightkurve.
        Downloads to the same path as lightkurve.search_tesscut().downlaod() does.
        """
        tesscut_path = os.path.join(os.path.expanduser('~'), '.lightkurve-cache/tesscut')
        tpf_path_table = Tesscut.download_cutouts(objectname=f'TIC {tic}',
                                                  size=(self.obj.tpf.shape[1], self.obj.tpf.shape[2]),
                                                  sector=self.obj.source_info.sector,
                                                  path=tesscut_path)
        tpf = lk.read(tpf_path_table['Local Path'].data[0])
        return tpf


    def plot_gaia_overlay(self, tic=None, tpf=None, cadence=0, ax=None, magnitude_limit=18):
        """
        Check if the source is contaminated.

        Parameters
        ----------
        tic : int, optional
            The TIC ID of the source.
        tpf : lightkurve.TargetPixelFile, optional
            Target Pixel File object. If None, searches and downloads TPF from MAST.
            Default is None.
        cadence : int, optional
            Index of the cadence in the TPF to plot.
            Default is 0 (the first cadence).
        ax : matplotlib.axes._subplots.AxesSubplot, optional
            Axes to plot on.
        magnitude_limit : float, optional
            The Gaia G magnitude limit for the Gaia sources to be plotted.
            Default is 18.
        """
        if tic is None:
            tic = self.obj.source_info.tic

        if tpf is None:
            tpf = self.download_tpf(tic)

        if ax is None:
            ax = tpf[cadence].plot(show_colorbar=False, title=f'TIC {tic}')
        else:
            ax = tpf[cadence].plot(ax=ax, show_colorbar=False, title=f'TIC {tic}')

        ax = self.add_gaia_figure_elements(fig_or_ax=ax, individual=False, tpf=tpf, magnitude_limit=magnitude_limit)
        return ax

    def add_gaia_figure_elements(self, fig_or_ax, individual=True, tic=None, tpf=None, magnitude_limit=18):
        """
        Make the Gaia Figure Elements

        Parameters
        ----------
        fig_or_ax : matplotlib.figure.Figure or matplotlib.axes._subplots.AxesSubplot
            Figure or Axes to overlay Gaia sources on.
        individual: bool, optional
            Whether the add_gaia_figure_elements() method is applied individually
            (i.e., not being called as a helper function of the plot_gaia_overlay() method).
            Default is True.
        tic : int, optional
            The TIC ID of the source.
        tpf : lightkurve.TargetPixelFile, optional
            Target Pixel File object. If None, searches and downloads TPF from MAST.
            Default is None.
        magnitude_limit : float, optional
            The Gaia G magnitude limit for the Gaia sources to be plotted.
            Default is 18.
        """
        if tic is None:
            tic = self.obj.source_info.tic

        if tpf is None:
            tpf = self.download_tpf(tic)

        # Get the positions of the Gaia sources
        c1 = SkyCoord(tpf.ra, tpf.dec, frame='icrs', unit='deg')
        # Use pixel scale for query size
        pix_scale = 21.0
        # We are querying with a diameter as the radius, overfilling by 2x.
        from astroquery.vizier import Vizier
        Vizier.ROW_LIMIT = -1
        result = Vizier.query_region(c1, catalog=["I/345/gaia2"],
                                     radius=Angle(np.max(tpf.shape[1:]) * pix_scale, "arcsec"))
        no_targets_found_message = ValueError('Either no sources were found in the query region '
                                              'or Vizier is unavailable')
        too_few_found_message = ValueError('No sources found brighter than {:0.1f}'.format(magnitude_limit))
        if result is None:
            raise no_targets_found_message
        elif len(result) == 0:
            raise too_few_found_message
        result = result["I/345/gaia2"].to_pandas()
        result = result[result.Gmag < magnitude_limit]
        if len(result) == 0:
            raise no_targets_found_message
        radecs = np.vstack([result['RA_ICRS'], result['DE_ICRS']]).T

        coords = tpf.wcs.all_world2pix(radecs, 0)
        # Correct the pixel coordinates when TPF is truncated to the border of the postcard and
        # the star is not located in the center of the TPF
        if ((self.obj.source_info.tc == False and individual) and
                (self.obj.tpf_star_x != int(np.floor(self.obj.tpf.shape[1] / 2)) or
                 self.obj.tpf_star_y != int(np.floor(self.obj.tpf.shape[2] / 2)))):
            coords[:, 0] -= int(np.floor(self.obj.tpf.shape[1] / 2) + 1) - self.obj.tpf_star_x
            coords[:, 1] -= int(np.floor(self.obj.tpf.shape[2] / 2) + 1) - self.obj.tpf_star_y

        try:
            year = ((tpf.time[0].jd - 2457206.375) * u.day).to(u.year)
        except:
            year = ((tpf.astropy_time[0].jd - 2457206.375) * u.day).to(u.year)
        pmra = ((np.nan_to_num(np.asarray(result.pmRA)) * u.milliarcsecond/u.year) * year).to(u.arcsec).value
        pmdec = ((np.nan_to_num(np.asarray(result.pmDE)) * u.milliarcsecond/u.year) * year).to(u.arcsec).value
        result.RA_ICRS += pmra
        result.DE_ICRS += pmdec

        # Gently size the points by their Gaia magnitude
        sizes = 10000.0 / 2**(result['Gmag']/2)

        target = tpf.wcs.world_to_pixel(c1)
        target_x = target[0]
        target_y = target[1]
        # Correct the pixel coordinates when TPF is truncated to the border of the postcard and
        # the star is not located in the center of the TPF
        if ((self.obj.source_info.tc == False and individual) and
                (self.obj.tpf_star_x != int(np.floor(self.obj.tpf.shape[1] / 2)) or
                 self.obj.tpf_star_y != int(np.floor(self.obj.tpf.shape[2] / 2)))):
            target_x -= int(np.floor(self.obj.tpf.shape[1] / 2) + 1) - self.obj.tpf_star_x
            target_y -= int(np.floor(self.obj.tpf.shape[2] / 2) + 1) - self.obj.tpf_star_y

        if isinstance(fig_or_ax, plt.Figure):
            ax = fig_or_ax.gca()
        elif isinstance(fig_or_ax, plt.Axes):
            ax = fig_or_ax
        else:
            raise ValueError("'fig_or_ax' must be a matplotlib Figure or Axes object.")

        if individual:
            ax.scatter(target_x, target_y, s=50, zorder=1000, c='k', marker='x')

            ax.scatter(coords[:, 0], coords[:, 1], c='firebrick', alpha=0.5, edgecolors='r', s=sizes)
            ax.scatter(coords[:, 0], coords[:, 1], c='None', edgecolors='r', s=sizes)
            ax.set_xlim([-0.5, tpf.shape[1]-0.5])
            ax.set_ylim([-0.5, tpf.shape[2]-0.5])

        else:
            ax.scatter(target_x+tpf.column, target_y+tpf.row, s=50, zorder=1000, c='k', marker='x')

            ax.scatter(coords[:, 0]+tpf.column, coords[:, 1]+tpf.row, c='firebrick', alpha=0.5, edgecolors='r', s=sizes)
            ax.scatter(coords[:, 0]+tpf.column, coords[:, 1]+tpf.row, c='None', edgecolors='r', s=sizes)
            ax.set_xlim([tpf.column-0.5, tpf.column+tpf.shape[1]-0.5])
            ax.set_ylim([tpf.row-0.5, tpf.row+tpf.shape[2]-0.5])

        return fig_or_ax
