import logging
import pdb

def plot_rti(sTime, rad, eTime=None, bmnum=7, fileType='fitacf',
                    params=['power','velocity','width'], scales=[], channel=None,
                    coords='gate', colors='lasse', yrng=-1, gsct=False,
                    low_gray=False, show=True, filtered=False,
                    fileName=None, txfreq_lims=None, myFile=None,
                    xtick_size=9, ytick_size=9,
                    xticks=None, axvlines=None,
                    plot_terminator=False, data_dict=None):

    """ Create an rti plot for a specified radar and time period.

    Parameters
    ----------
    sTime : datetime
        a datetime object indicating the start time which you would like
        to plot
    rad : str
        the 3 letter radar code, e.g. 'bks'
    eTime : Optional[datetime]
        a datetime object indicating th end time you would like plotted.
        If this is None, 24 hours will be plotted.  default = None.
    bmnum : Optional[int]
        The beam to plot.  default: 7
    fileType : Optional[str]
        The file type to be plotted, one of ['fitex', 'fitacf', 'lmfit'].
        default = 'fitex'.
    params : Optional[list]
        a list of the fit parameters to plot, allowable values are:
        ['velocity', 'power', 'width', 'elevation', 'phi0'].  default:
        ['velocity', 'power', 'width']
    scales : Optional[list]
        a list of the min/max values for the color scale for each param.
        If omitted, default scales will be used.  If present, the list
        should be n x 2 where n is the number of elements in the params
        list.  Use an empty list for default range, e.g. [[-250,300],[],[]].
        default: [[-200,200], [0,30],[0,150]]
    channel : Optional[char]
        the channel you wish to plot, e.g. 'a', 'b', 'c', ...  default: 'a'
    coords : Optional[str]
        the coordinates to use for the y axis.  The allowable values are
        'gate', 'rng', 'geo', 'mag' default: 'gate'
    colors : Optional[str]
        a string indicating what color bar to use, valid inputs are
        'lasse' or 'aj'. Also can be a list of matplotlib colormaps names,
        for example ['jet','jet','jet'] if len(param)==3. default: 'lasse'.
    yrng : Optional[list or -1]
        a list indicating the min and max values for the y axis in the
        chosen coordinate system, or a -1 indicating to plot everything.
        default: -1.
    gsct : Optional[boolean]
        a flag indicating whether to plot ground scatter as gray. default:
        False (ground scatter plotted normally)
    low_gray : Optional[boolean]
        a flag indicating whether to plot low velocity scatter as gray.
        default: False (low velocity scatter plotted normally)
    show : Optional[boolean]
        a flag indicating whether to display the figure on the screen.
        This can cause problems over ssh.  default = True
    retfig : Optional[boolean]
        a flag indicating that you want the figure to be returned from
        the function.  Only the last figure in the list of frequency bands
        will be returned.  default = False
    filtered : Optional[boolean]
        a flag indicating whether to boxcar filter the data.  default:
        False (no filter)
    fileName : Optional[string]
        If you want to plot for a specific file, indicate the name of the
        file as fileName.  Include the type of the file in custType.
    txfreq_lims : Optional[list]
        a list of the min/max values for the transmitter frequencies in
        kHz.  If omitted, the default band will be used.  If more than
        one band is specified, retfig will cause only the last one to be
        returned.  default: [[8000,20000]]
    myFile : Optional[pydarn.sdio.radDataTypes.radDataPtr]
        contains the pipeline to the data we want to plot. If specified,
        data will be plotted from the file pointed to by myFile. default: None
    xtick_size : Optional[int]
        fontsize of xtick labels
    ytick_size : Optional[int]
        fontsize of ytick labels
    xticks : Optional[list]
        datetime.datetime objects indicating the location of xticks
    axvlines : Optoinal[list]
        datetime.datetime objects indicating the location vertical lines
        marking the plot
    plot_terminator : Optional[boolean]
        Overlay the day/night terminator.
    data_dict : Optinal[dict, default value is None]
        This takes the output of read_data function in this module.

    Returns
    -------
    A figure object

    Example
    -------
        import datetime as dt
        pydarn.plotting.rti.plot_rti(dt.datetime(2013,3,16), 'bks',
                                     eTime=dt.datetime(2013,3,16,14,30),
                                     bmnum=12, fileType='fitacf',
                                     scales=[[-500,500],[],[]], coords='geo',
                                     colors='aj', filtered=True, show=True)

    Written by AJ 20121002
    Modified by Matt W. 20130715
    Modified by Nathaniel F. 20131031 (added plot_terminator)
    Modified by ASR 20150917 (refactored)
    Modified by Muhamamd 20160728

    """

    import os
    from davitpy import pydarn
    from davitpy import utils
    import numpy as np
    from datetime import datetime, timedelta
    from matplotlib import pyplot
    from matplotlib.dates import DateFormatter
    import matplotlib.cm as cm

    #Time how long this is going to take
    timing_start = datetime.now()

    #NOTE TO DEVS: List of available params. Can be simply expanded
    #as more parameters are added to SuperDARN data set (like index
    #of refraction)
    available_params = ['power','velocity','width','elevation','phi0','velocity_error']
    default_scales = [[0,30],[-200,200],[0,150],[0,50],[-np.pi, np.pi],[0,200]]

    available_text = 'Allowable parameters are '
    for p in available_params:
        available_text = available_text + p + ', '
    available_text = available_text[:-2]

    #Check the inputs
    assert(isinstance(sTime, datetime)), logging.error(
        'sTime must be a datetime object')
    assert(isinstance(rad, str) and len(rad) == 3), logging.error(
        'rad must be a string 3 chars long')
    assert(isinstance(eTime, datetime) or
           eTime is None), (
        logging.error('eTime must be a datetime object or None'))
    if eTime is None:
        eTime = sTime + timedelta(days=1)
    assert(sTime < eTime), logging.error("eTime must be greater than sTime!")
    assert(coords == 'gate' or coords == 'rng' or coords == 'geo' or
           coords == 'mag'), logging.error("coords must be one of 'gate', "
                                           "'rng', 'geo', 'mag'")
    assert(isinstance(bmnum, int)), logging.error('beam must be integer')
    assert(0 < len(params) < 6), (
        logging.error('must input between 1 and 5 params in LIST form'))
    for i in range(0, len(params)):
        assert(params[i] in available_params), (
            logging.error(available_text))
    for i in range(0, len(scales)):
        assert(isinstance(scales[i], list)), (
            logging.error('each item in scales must be a list of upper and '
                          'lower bounds on paramaters.'))
    assert(scales == [] or
           len(scales) == len(params)), (
        logging.error('if present, scales must have same number of elements '
                      'as params'))
    assert(yrng == -1 or
           (isinstance(yrng, list) and
            yrng[0] <= yrng[1])), (
        logging.error('yrng must equal -1 or be a list with the 2nd element '
                      'larger than the first'))
    assert((colors == 'lasse' or colors == 'aj')) or isinstance(colors,list), (
        logging.error("Valid inputs for color are 'lasse' and 'aj' or a list of matplotlib colormaps"))

    assert( (isinstance(txfreq_lims,list) and len(txfreq_lims)==2) or
            isinstance(txfreq_lims,type(None))), (
        logging.error("txfreq_lims must be a list with the start and end frequencies") )

    # Assign any default color scale parameter limits.
    tscales = []
    for i in range(0, len(params)):
        if(scales == [] or scales[i] == []):
            if(params[i] in available_params):
                ind = available_params.index(params[i])
                tscales.append(default_scales[ind])
        else: tscales.append(scales[i])
    scales = tscales

    # Assign default frequency band.
    if txfreq_lims is None:
        tband = [8000, 20000]
    else:
        assert(txfreq_lims[0] < txfreq_lims[1]), (
            logging.error("Starting frequency must be less "
                          "than ending frequency!"))
        tband = txfreq_lims

    # Open the file if a pointer was not given to us
    # if fileName is specified then it will be read.
    if not myFile:
        from davitpy.pydarn.sdio import radDataOpen
        myFile = radDataOpen(sTime, rad, eTime, channel=channel, bmnum=bmnum,
                             fileType=fileType, filtered=filtered,
                             fileName=fileName)

        # Check that we have data available now that we may have tried
        # to read it using radDataOpen.
        if myFile is None:
            logging.error('No files available for the requested '
                          'time/radar/filetype combination')
            return None

    # Make sure that we will only plot data for the time range specified
    # by sTime and eTime.
    if myFile.sTime <= sTime and myFile.eTime > sTime and myFile.eTime >= eTime:
        myFile.sTime = sTime
        myFile.eTime = eTime
    else:
        # If the times range is not covered by the file, warn the user.
        logging.warning('Data not available in myFile for the whole of sTime to eTime!')

    # Finally we can start reading the data file
    myBeam = myFile.readRec()
    if myBeam is None:
        logging.error('Problem reading the data.')
        return None

    # Now read the data that we need to make the plots
    if data_dict is None:
        data_dict = read_data(myFile, bmnum, params, tband)
    else:
        # Make sure that we only plot data for the time range specified
        selected_indices = [x for x in xrange(len(data_dict['datetime'])) if \
                (data_dict['datetime'][x] >= sTime and data_dict['datetime'][x] <= eTime)]
        sindx = selected_indices[0]
        eindx = selected_indices[-1]
        for keywrd in data_dict.keys():
            data_dict[keywrd] = data_dict[keywrd][sindx:eindx+1]

    # Check to ensure that data exists for the requested frequency
    # band else continue on to the next range of frequencies
    if len(data_dict['freq']) == 0:
        logging.error('No data found in frequency range ' +
                        str(tbands[0]) + ' kHz to ' +
                        str(tbands[1]) + ' kHz')
        return None

    # Create a figure.
    rti_fig = pyplot.figure(figsize=(11, 8.5))

    # Create the axes for noise, tx freq, and cpid.
    noise_pos = [.1, .88, .76, .06]
    freq_pos = [.1, .82, .76, .06]
    cpid_pos = [.1, .77, .76, .05]

    skynoise_ax = rti_fig.add_axes(noise_pos, label='sky')
    searchnoise_ax = rti_fig.add_axes(noise_pos, label='search',
                                      frameon=False)
    freq_ax = rti_fig.add_axes(freq_pos, label='freq')
    nave_ax = rti_fig.add_axes(freq_pos, label='nave', frameon=False)
    cpid_ax = rti_fig.add_axes(cpid_pos)

    # Give the plot a title.
    rti_title(rti_fig, sTime, rad, fileType, bmnum, eTime=eTime)

#    # Plot the sky noise.
#    plot_skynoise(skynoise_ax, data_dict['datetime'],
#                  data_dict['nsky'])
#    # Plot the search noise.
#    plot_searchnoise(searchnoise_ax, data_dict['datetime'],
#                        data_dict['nsch'])
    # plot the frequency bar.
    plot_freq(freq_ax, data_dict['datetime'],
                data_dict['freq'])
    # Plot the nave data.
    plot_nave(nave_ax, data_dict['datetime'],
                data_dict['nave'])
    # Plot the cpid bar
    plot_cpid(cpid_ax, data_dict['datetime'],
                data_dict['cpid'], data_dict['mode'])

    # Plot each of the parameter panels.
    figtop = .77
    if ((eTime - sTime) <= timedelta(days=1)) and \
            (eTime.day == sTime.day):
        figheight = .72 / len(params)
    elif ((eTime - sTime) > timedelta(days=1)) or \
            (eTime.day != sTime.day):
        figheight = .70 / len(params)

    for p in range(len(params)):
        # Use draw_axes to create and set formatting of the axes to
        # plot to.
        pos = [.1, figtop - figheight * (p + 1) + .02, .76,
                figheight - .02]
        ax = draw_axes(rti_fig, data_dict['datetime'], rad,
                       data_dict['cpid'], bmnum,
                       data_dict['nrang'],
                       data_dict['frang'], data_dict['rsep'],
                       p == len(params) - 1, yrng=yrng, coords=coords,
                       pos=pos, xtick_size=xtick_size,
                       ytick_size=ytick_size, xticks=xticks,
                       axvlines=axvlines)

        if(params[p] == 'velocity'): pArr = data_dict['vel']
        elif(params[p] == 'power'): pArr = data_dict['pow']
        elif(params[p] == 'width'): pArr = data_dict['wid']
        elif(params[p] == 'elevation'): pArr = data_dict['elev']
        elif(params[p] == 'phi0'): pArr = data_dict['phi0']
        elif(params[p] == 'velocity_error'):
            pArr = data_dict['velocity_error']
        if(pArr == []): continue

        # Generate the color map.

        if colors in ['aj','lasse']:
            cmap, norm, bounds = utils.plotUtils.genCmap(params[p], scales[p],
                                                         colors=colors,
                                                         lowGray=low_gray)
        else:
            from matplotlib import colors as mpl_colors
            norm = mpl_colors.Normalize(vmin=scales[p][0], vmax=scales[p][1])
            cmap = cm.get_cmap(colors[p])

        # Plot the data to the axis object.
        pcoll = rti_panel(ax,data_dict,pArr,gsct,rad,bmnum,coords,cmap,norm,
                          plot_terminator=plot_terminator)

        # Set xaxis formatting depending on amount of data plotted.
        if ((eTime - sTime) <= timedelta(days=1)) and \
                (eTime.day == sTime.day):
            ax.xaxis.set_major_formatter(DateFormatter('%H:%M'))
        elif ((eTime - sTime) > timedelta(days=1)) or \
                (eTime.day != sTime.day):
            ax.xaxis.set_major_formatter(DateFormatter('%d/%m/%y \n%H:%M'))
        ax.set_xlabel('UT')

        # Draw the colorbar.
        cb = utils.drawCB(rti_fig, pcoll, cmap, norm, map_plot=0,
                          pos=[pos[0] + pos[2] + .02, pos[1], 0.02,
                          pos[3]])

        if colors in ['aj','lasse']:
            # Label the colorbar.
            l = []
            # Define the colorbar labels.
            for i in range(0, len(bounds)):
                if(params[p] == 'phi0'):
                    ln = 4
                    if(bounds[i] == 0): ln = 3
                    elif(bounds[i] < 0): ln = 5
                    l.append(str(bounds[i])[:ln])
                    continue
                if((i == 0 and
                    (params[p] == 'velocity' or
                     params[p] == 'velocity_error')) or i == len(bounds) - 1):
                    l.append(' ')
                    continue
                l.append(str(int(bounds[i])))
            cb.ax.set_yticklabels(l)

            # Set colorbar ticklabel size.
            for t in cb.ax.get_yticklabels():
                t.set_fontsize(9)

        # Set colorbar label.
        if(params[p] == 'velocity'):
            cb.set_label('Velocity [m/s]', size=10)
        if(params[p] == 'grid'): cb.set_label('Velocity [m/s]', size=10)
        if(params[p] == 'power'): cb.set_label('SNR [dB]', size=10)
        if(params[p] == 'width'): cb.set_label('Spec Wid [m/s]', size=10)
        if(params[p] == 'elevation'): cb.set_label('Elev [deg]', size=10)
        if(params[p] == 'phi0'): cb.set_label('Phi0 [rad]', size=10)
        if(params[p] == 'velocity_error'):
            cb.set_label('Velocity Error [m/s]', size=10)

    if show:
        rti_fig.show()

    logging.info('plotting took:', datetime.now() - timing_start)

    return rti_fig

def draw_axes(myFig, times, rad, cpid, bmnum, nrang, frang, rsep, bottom,
              yrng=-1, coords='gate', pos=[.1, .05, .76, .72], xtick_size=9,
              ytick_size=9, xticks=None, axvlines=None):
    """ Draws empty axes for an rti plot.

    Parameters
    ----------
    myFig :
        the MPL figure we are plotting to
    times : list
        a list of datetime objects referencing the beam soundings
    rad : str
        3 letter radar code
    cpid : list
        list of the cpids or the beam soundings
    bmnum : int
        beam number being plotted
    nrang : list
        list of nrang for the beam soundings
    frang : list
        list of frang of the beam soundings
    rsep : list
        list of rsep of the beam soundings
    bottom : bool
        flag indicating if we are at the bottom of the figure
    yrng : Optional[list]
        range of y axis, -1=autoscale (default)
    coords : Optional[ ]
        y axis coordinate system, acceptable values are 'geo',
        'mag', 'gate', 'rng'
    pos : Optional[ ]
        position of the plot
    xtick_size : Optional[ ]
        fontsize of xtick labels
    ytick_size : Optional[ ]
        fontsize of ytick labels
    xticks : Optional[list]
        datetime.datetime objects indicating the location of xticks
    axvlines : Optional[list]
        datetime.datetime objects indicating the location vertical
        lines marking the plot

    Returns
    -------
    ax :
        an axes object

    Example
    -------
        ax = draw_axes(myFig,times,rad,cpid,beam,nrang,frang,rsep,0)

    Written by AJ 20121002
    Modified by ASR 20150917 (refactored)

    """

    from davitpy import pydarn
    from matplotlib.ticker import MultipleLocator,FormatStrFormatter
    from matplotlib.dates import SecondLocator,DateFormatter,date2num
    from matplotlib.lines import Line2D
    import numpy as np

    nrecs = len(times)
    # Add an axes object to the figure.
    ax = myFig.add_axes(pos)
    ax.yaxis.set_tick_params(direction='out')
    ax.xaxis.set_tick_params(direction='out')
    ax.yaxis.set_tick_params(direction='out', which='minor')
    ax.xaxis.set_tick_params(direction='out', which='minor')

    # Draw the axes.
    ax.plot_date(date2num(times), np.arange(len(times)),
                 fmt='w', tz=None, xdate=True, ydate=False, alpha=0.0)

    # Determine the yaxis min/max unless it's been specified
    if(yrng == -1):
        ymin, ymax = 99999999, -999999999
        if(coords != 'gate'):
            oldCpid = -99999999
            for i in range(len(cpid)):
                if(cpid[i] == oldCpid): continue
                oldCpid = cpid[i]
                if(coords == 'geo' or coords == 'mag'):
                    # HACK NOT SURE IF YOU CAN DO THIS(Formatting)!
                    site = pydarn.radar.network().getRadarByCode(rad) \
                        .getSiteByDate(times[i])
                    myFov = pydarn.radar.radFov.fov(site=site, ngates=nrang[i],
                                                    nbeams=site.maxbeam,
                                                    rsep=rsep[i],
                                                    coords=coords,
                                                    date_time=times[i])
                    if(myFov.latFull[bmnum].max() > ymax):
                        ymax = myFov.latFull[bmnum].max()
                    if(myFov.latFull[bmnum].min() < ymin):
                        ymin = myFov.latFull[bmnum].min()
                else:
                    ymin = 0
                    if(nrang[i] * rsep[i] + frang[i] > ymax):
                        ymax = nrang[i] * rsep[i] + frang[i]

        else:
            ymin, ymax = 0, max(nrang)
    else:
        ymin, ymax = yrng[0], yrng[1]

    # Format the xaxis.
    xmin = date2num(times[0])
    xmax = date2num(times[len(times) - 1])
    xrng = (xmax - xmin)
    inter = int(round(xrng / 6. * 86400.))
    inter2 = int(round(xrng / 24. * 86400.))
    ax.xaxis.set_minor_locator(SecondLocator(interval=inter2))
    ax.xaxis.set_major_locator(SecondLocator(interval=inter))

    # If axis is at bottom of figure, draw xticks.
    if(not bottom):
        for tick in ax.xaxis.get_major_ticks():
            tick.label.set_fontsize(0)
    else:
        if xticks is not None:
            ax.xaxis.set_ticks(xticks)

    if axvlines is not None:
        for line in axvlines:
            ax.axvline(line, color='0.25', ls='--')

        for tick in ax.xaxis.get_major_ticks():
            tick.label.set_fontsize(xtick_size)
        ax.xaxis.set_major_formatter(DateFormatter('%H:%M'))
        ax.xaxis.set_label_text('UT')

    # Set ytick size.
    for tick in ax.yaxis.get_major_ticks():
        tick.label.set_fontsize(ytick_size)
    # Format yaxis depending on coords.
    if(coords == 'gate'):
        ax.yaxis.set_label_text('Range gate', size=10)
        ax.yaxis.set_major_formatter(FormatStrFormatter('%d'))
        ax.yaxis.set_major_locator(MultipleLocator((ymax - ymin) / 5.))
        ax.yaxis.set_minor_locator(MultipleLocator((ymax - ymin) / 25.))
    elif(coords == 'geo' or coords == 'mag'):
        if(coords == 'mag'):
            ax.yaxis.set_label_text('Mag Lat [deg]', size=10)
        else:
            ax.yaxis.set_label_text('Geo Lat [deg]', size=10)
    elif(coords == 'rng'):
        ax.yaxis.set_label_text('Slant Range [km]', size=10)
        ax.yaxis.set_major_formatter(FormatStrFormatter('%d'))
        ax.yaxis.set_major_locator(MultipleLocator(1000))
        ax.yaxis.set_minor_locator(MultipleLocator(250))
    ax.set_ylim(bottom=ymin, top=ymax)

    return ax

def rti_title(fig, sTime, rad, fileType, beam, eTime=None, xmin=.1, xmax=.86):
    """Draws title for an rti plot.

    Parameters
    ----------
    fig : matplotlib.figure.Figure
        a matplotlib.figure.Figure object
    sTime : datetime
        the start time for the data being plotted as a datetime object
    rad : str
        the 3 letter radar code
    fileType : str
        the file type being plotted
    beam : int
        the beam number being plotted
    eTime : Optional[datetime]
        the end time for the data being plotted as a datetime object
    xmin : Optional[ ]
        minimum x value o the plot in page coords
    xmax : Optional[ ]
        maximum x value o the plot in page coords

    Returns
    -------
    Nothing.

    Example
    -------
        import datetime as dt
        from matplotlib import pyplot
        fig = pyplot.figure()
        rti_title(fig,dt.datetime(2011,1,1),'bks','fitex',7)

    Written by AJ 20121002
    Modified by ASR 20150916

    """
    from davitpy import pydarn
    from datetime import timedelta
    import calendar

    # Obtain the davitpy.pydarn.radar.radStruct.radar object for rad.
    r = pydarn.radar.network().getRadarByCode(rad)

    # Plot the main title
    fig.text(xmin, .95, r.name + '  (' + fileType + ')', ha='left', weight=550)

    # Determine what time information should be plotted in the secondary title
    if ((eTime is not None) and
        (((eTime - sTime) > timedelta(days=1)) or
         (eTime.day != sTime.day))):
        title_text = str(sTime.day) + '/' \
            + calendar.month_name[sTime.month][:3] + '/' \
            + str(sTime.year) + ' - ' + str(eTime.day) + '/' \
            + calendar.month_name[eTime.month][:3] + '/' \
            + str(eTime.year)

    else:
        title_text = str(sTime.day) + '/' \
            + calendar.month_name[sTime.month][:3] + '/' \
            + str(sTime.year)

    # Plot the secondary title.
    fig.text((xmin + xmax) / 2., .95, title_text, weight=550,
             size='large', ha='center')
    fig.text(xmax, .95, 'Beam ' + str(beam), weight=550, ha='right')

def plot_cpid(ax, times, cpid, mode):
    """Plots control program ID (cpid) panel at position pos.

    Parameters
    ----------
    ax :
        a MPL axis object to plot to
    times : list
        a list of the times of the beam soundings
    cpid : list
        a list of the cpids of th beam soundings
    mode : list
        a list of the ifmode param

    Returns
    -------
    Nothing.

    Example
    -------
        plot_cpid(ax,times,cpid,mode)

    Written by AJ 20121002
    Modified by ASR 20150916

    """
    from davitpy import pydarn
    from matplotlib.ticker import MultipleLocator
    from matplotlib.dates import SecondLocator
    from matplotlib.dates import date2num
    import numpy as np
    oldCpid = -9999999

    # Format the yaxis.
    ax.yaxis.tick_left()
    ax.yaxis.set_tick_params(direction='out')
    ax.set_ylim(bottom=0, top=1)
    ax.yaxis.set_minor_locator(MultipleLocator(1))
    ax.yaxis.set_tick_params(direction='out', which='minor')

    # Draw the axes.
    ax.plot_date(date2num(times), np.arange(len(times)),
                 fmt='w', tz=None, xdate=True, ydate=False, alpha=0.0)

    # Label the CPIDs.
    for i in range(0, len(times)):
        if(cpid[i] != oldCpid):
            ax.plot_date([date2num(times[i]),date2num(times[i])],
                         [0, 1], fmt='k-', tz=None, xdate=True, ydate=False)
            oldCpid = cpid[i]
            s = ' ' + pydarn.radar.radUtils.getCpName(oldCpid)
            istr = ' '
            if(mode[i] == 1): istr = ' IF'
            if(mode == 0): istr = ' RF'
            ax.text(times[i], .5, ' ' + str(oldCpid) + s + istr, ha='left',
                    va='center', size=10)

    # Format the xaxis.
    xmin = date2num(times[0])
    xmax = date2num(times[len(times) - 1])
    xrng = (xmax - xmin)
    inter = int(round(xrng / 6. * 86400.))
    inter2 = int(round(xrng / 24. * 86400.))
    ax.xaxis.set_minor_locator(SecondLocator(interval=inter2))
    ax.xaxis.set_major_locator(SecondLocator(interval=inter))

    for tick in ax.xaxis.get_major_ticks():
        tick.label.set_fontsize(0)

    # Identify the CPID axis with a label.
    fig = ax.get_figure()
    bb = ax.get_position()
    x0 = bb.x0
    y0 = bb.y0
    height = bb.height
    width = bb.width
    pos = [x0, y0, width, height]
    fig.text(pos[0] - .07, pos[1] + pos[3] / 2., 'CPID', ha='center',
             va='center', size=8.5, rotation='vertical')
    ax.set_yticks([])

def plot_skynoise(ax, times, sky, xlim=None, xticks=None):
    """Plots a noise panel at position pos.

    Parameters
    ----------
    ax :
        a MPL axis object to plot to
    times : list
        a list of the times of the beam soundings
    sky: list
        a lsit of the noise.sky of the beam soundings
    search : list
        a list of the noise.search param
    xlim : Optional[list]
        2-element limits of the x-axis.  None for default.
    xticks : Optional[list]
        List of xtick poisitions.  None for default.

    Returns
    -------
    Nothing

    Example
    -------
        plot_skynoise(ax,times,sky)

    Written by AJ 20121002
    Modified by NAF 20131101
    Modified by ASR 20150916

    """

    from matplotlib.ticker import MultipleLocator
    from matplotlib.dates import date2num
    from matplotlib.lines import Line2D
    import numpy as np

    # Format the yaxis.
    ax.yaxis.tick_left()
    ax.yaxis.set_tick_params(direction='out')
    ax.set_ylim(bottom=0, top=6)
    ax.yaxis.set_minor_locator(MultipleLocator())
    ax.yaxis.set_tick_params(direction='out', which='minor')

    # Plot the sky noise data.
#    if np.all(sky):
#        sky_log10 = np.log10(sky)
#    else:
#        sky_log10 = np.array([0]*len(sky))
#    #sky_log10 = [np.log10(x) if x>0 else np.nan for x in sky]
    sky_log10 = np.log10(sky)
    ax.plot_date(date2num(times), sky_log10, fmt='k-',
                 tz=None, xdate=True, ydate=False)

    # Format the xaxis.
    if xlim is not None: ax.set_xlim(xlim)
    if xticks is not None: ax.set_xticks(xticks)

    # Add labels to identify the noise axis.
    fig = ax.get_figure()
    bb = ax.get_position()
    x0 = bb.x0
    y0 = bb.y0
    height = bb.height
    width = bb.width
    pos = [x0, y0, width, height]
    fig.text(pos[0] - .01, pos[1] + .004, '10^0', ha='right', va='bottom',
             size=8)
    fig.text(pos[0] - .01, pos[1] + pos[3], '10^6', ha='right', va='top',
             size=8)
    fig.text(pos[0] - .07, pos[1] + pos[3] / 2., 'N.Sky', ha='center',
             va='center', size=8.5, rotation='vertical')
    l = Line2D([pos[0] - .06, pos[0] - .06], [pos[1] + .01,
               pos[1] + pos[3] - .01], transform=fig.transFigure,
               clip_on=False, ls='-', color='k', lw=1.5)
    ax.add_line(l)
    ax.set_xticklabels([' '])
    # Only use 2 major yticks.
    ax.set_yticks([0, 6])
    ax.set_yticklabels([' ', ' '])

def plot_searchnoise(ax, times, search, xlim=None, xticks=None,
                     ytickside='right'):
    """Plots a noise panel at position pos.

    Parameters
    ----------
    ax :
        a MPL axis object to plot to
    times : list
        a list of the times of the beam soundings
    sky : list
        a lsit of the noise.sky of the beam soundings
    search : list
        a list of the noise.search param
    xlim : Optional[list]
        2-element limits of the x-axis.  None for default.
    xticks : Optional[list]
        List of xtick poisitions.  None for default.
    ytickside : Optional[string]
        Default is right.

    Returns
    -------
    Nothing

    Example
    -------
        plot_searchnoise(ax,times,search)

    Written by AJ 20121002
    Modified by NAF 20131101
    Modified by ASR 20150916

    """

    from matplotlib.ticker import MultipleLocator
    from matplotlib.dates import date2num
    from matplotlib.lines import Line2D
    import numpy as np

    # Format the yaxis.
    ax.yaxis.tick_left()
    ax.yaxis.set_tick_params(direction='out')
    ax.set_ylim(bottom=0, top=6)
    ax.yaxis.set_minor_locator(MultipleLocator())
    ax.yaxis.set_tick_params(direction='out', which='minor')

    # Plot the search noise data.
#    if np.all(search):
#        search_log10 = np.log10(search)
#    else:
#        search_log10 = np.array([0]*len(search))
#    search_log10 = [np.log10(x) if x>0 else np.nan for x in search]
    search_log10 = np.log10(search)
    ax.plot_date(date2num(times), search_log10,
                 fmt='k:', tz=None, xdate=True, ydate=False, lw=1.5)

    # Format the xaxis.
    if xlim is not None: ax.set_xlim(xlim)
    if xticks is not None: ax.set_xticks(xticks)

    # Add labels to identify the noise axis.
    fig = ax.get_figure()
    bb = ax.get_position()
    x0 = bb.x0
    y0 = bb.y0
    height = bb.height
    width = bb.width
    pos = [x0, y0, width, height]

    fig.text(pos[0] + pos[2] + .01, pos[1] + .004, '10^0', ha='left',
             va='bottom', size=8)
    fig.text(pos[0] + pos[2] + .01, pos[1] + pos[3], '10^6', ha='left',
             va='top', size=8)
    fig.text(pos[0] + pos[2] + .06, pos[1] + pos[3] / 2., 'N.Sch', ha='center',
             va='center', size=8.5, rotation='vertical')

    l = Line2D([pos[0] + pos[2] + .07, pos[0] + pos[2] + .07],
                     [pos[1] + .01, pos[1] + pos[3] - .01],
                     transform=fig.transFigure, clip_on=False, ls=':',
                     color='k', lw=1.5)
    ax.add_line(l)
    ax.set_xticklabels([' '])
    # use only 2 major yticks
    ax.set_yticks([0, 6])
    ax.set_yticklabels([' ', ' '])
    if ytickside == 'right':
        ax.yaxis.tick_right()

def plot_freq(ax, times, freq, xlim=None, xticks=None):
    """Plots the tx frequency data to an axis object.

    Parameters
    ----------
    ax :
        a MPL axis object to plot to
    times : list
        a list of the times of the beam soundings
    freq : list
        a lsit of the tfreq of the beam soundings
    xlim : Optional[list]
        2-element limits of the x-axis.  None for default.
    xticks : Optional[list]
        List of xtick poisitions.  None for default.

    Returns
    -------
    Nothing.

    Example
    -------
        plot_freq(ax, times, tfreq)

    Written by AJ 20121002
    Modified by NAF 20131101
    Modified by ASR 20150916

    """

    from matplotlib.ticker import MultipleLocator
    from matplotlib.dates import date2num
    from matplotlib.lines import Line2D

    # Format the yaxis.
    ax.yaxis.tick_left()
    ax.yaxis.set_tick_params(direction='out')
    ax.set_ylim(bottom=8, top=20)
    ax.yaxis.set_minor_locator(MultipleLocator())
    ax.yaxis.set_tick_params(direction='out', which='minor')

    # Plot the TX frequency.
    ax.plot_date(date2num(times), freq, fmt='k-',
                 tz=None, xdate=True, ydate=False, markersize=2)

    # Format the xaxis.
    if xlim is not None: ax.set_xlim(xlim)
    if xticks is not None: ax.set_xticks(xticks)

    # Add labels to identify the frequency axis.
    fig = ax.get_figure()
    bb = ax.get_position()
    x0 = bb.x0
    y0 = bb.y0
    height = bb.height
    width = bb.width
    pos = [x0, y0, width, height]
    fig.text(pos[0] - .01, pos[1] + .005, '10', ha='right', va='bottom',
             size=8)
    fig.text(pos[0] - .01, pos[1] + pos[3] - .015, '16', ha='right', va='top',
             size=8)
    fig.text(pos[0] - .07, pos[1] + pos[3] / 2., 'Freq', ha='center',
             va='center', size=9, rotation='vertical')
    fig.text(pos[0] - .05, pos[1] + pos[3] / 2., '[MHz]', ha='center',
             va='center', size=7, rotation='vertical')
    l = Line2D([pos[0] - .04, pos[0] - .04], [pos[1] + .01,
               pos[1] + pos[3] - .01], transform=fig.transFigure,
               clip_on=False, ls='-', color='k', lw=1.5)
    ax.add_line(l)
    ax.set_xticklabels([' '])
    # use only 2 major yticks
    ax.set_yticks([10, 16])
    ax.set_yticklabels([' ', ' '])

def plot_nave(ax, times, nave, xlim=None, xticks=None, ytickside='right'):
    """Plots the number of averages (nave) data to an axis object.

    Parameters
    ----------
    ax :
        a MPL axis object to plot to
    times : list
        a list of the times of the beam soundings
    nave : list
        a lsit of the nave of the beam soundings
    xlim : Optional[list]
        2-element limits of the x-axis.  None for default.
    xticks : Optional[list]
        List of xtick poisitions.  None for default.
    ytickside : Optional[str]
        Default is right.

    Returns
    -------
    Nothing.

    Example
    -------
        plot_nave(ax, times, nave)

    Written by AJ 20121002
    Modified by NAF 20131101
    Modified by ASR 20150916

    """

    from matplotlib.ticker import MultipleLocator
    from matplotlib.dates import date2num
    from matplotlib.lines import Line2D

    # Format the yaxis
    ax.yaxis.tick_left()
    ax.yaxis.set_tick_params(direction='out')
    ax.set_ylim(bottom=0, top=80)
    ax.yaxis.set_minor_locator(MultipleLocator(base=5))
    ax.yaxis.set_tick_params(direction='out', which='minor')

    # Plot the number of averages.
    ax.plot_date(date2num(times), nave, fmt='k:',
                 tz=None, xdate=True, ydate=False, markersize=2)

    # Format the xaxis.
    if xlim is not None: ax.set_xlim(xlim)
    if xticks is not None: ax.set_xticks(xticks)

    # Add labels to identify the nave axis.
    fig = ax.get_figure()
    bb = ax.get_position()
    x0 = bb.x0
    y0 = bb.y0
    height = bb.height
    width = bb.width
    pos = [x0, y0, width, height]
    fig.text(pos[0] + pos[2] + .01, pos[1] - .004, '0', ha='left', va='bottom',
             size=8)
    fig.text(pos[0] + pos[2] + .01, pos[1] + pos[3], '80', ha='left', va='top',
             size=8)
    fig.text(pos[0] + pos[2] + .06, pos[1] + pos[3] / 2., 'Nave', ha='center',
             va='center', size=8.5, rotation='vertical')

    l = Line2D([pos[0] + pos[2] + .07, pos[0] + pos[2] + .07],
               [pos[1] + .01, pos[1] + pos[3] - .01],
               transform=fig.transFigure, clip_on=False, ls=':',
               color='k', lw=1.5)
    ax.add_line(l)
    ax.set_xticklabels([' '])
    # use only 2 major yticks
    ax.set_yticks([0, 80])
    ax.set_yticklabels([' ', ' '])
    if ytickside == 'right':
        ax.yaxis.tick_right()

def read_data(myPtr, bmnum, params, tbands):
    """Reads data from the file pointed to by myPtr

    Parameter
    ---------
    myPtr :
        a davitpy file pointer object
    bmnum : int
        beam number of data to read in
    params : list
        a list of the parameters to read
    tbands : list
        a list of the frequency bands to separate data into

    Returns
    -------
    A dictionary of the data. Data is stored in lists and separated in
    to tbands.

    Example
    -------
        from davitpy import pydarn
        from datetime import datetime
        myPtr = pydarn.sdio.radDataOpen(datetime(2012,11,24),'sas')
        myBeam = myPtr.readRec()
        data_dict = read_data(myPtr, myBeam, 7, ['velocity'], [8000,20000])

    Written by ASR 20150914

    """

    import numpy as np

    # Initialize some things.
    data = dict()
    data_keys = ['vel', 'pow', 'wid', 'elev', 'phi0', 'datetime', 'freq', 'cpid',
                 'nave', 'nsky', 'nsch', 'slist', 'mode', 'rsep', 'nrang',
                 'frang', 'gsflg', 'velocity_error']
    for d in data_keys:
        data[d] = []

    # Read the parameters of interest.
    myPtr.rewind()
    myBeam = myPtr.readRec()
    while(myBeam is not None):
        if(myBeam.time > myPtr.eTime): break
        if(myBeam.bmnum == bmnum and (myPtr.sTime <= myBeam.time)):
            if (myBeam.prm.tfreq >= tbands[0] and
                myBeam.prm.tfreq <= tbands[1]):
                data['datetime'].append(myBeam.time)
                data['cpid'].append(myBeam.cp)
                data['nave'].append(myBeam.prm.nave)
                data['nsky'].append(myBeam.prm.noisesky)
                data['rsep'].append(myBeam.prm.rsep)
                data['nrang'].append(myBeam.prm.nrang)
                data['frang'].append(myBeam.prm.frang)
                data['nsch'].append(myBeam.prm.noisesearch)
                data['freq'].append(myBeam.prm.tfreq / 1e3)
                data['slist'].append(myBeam.fit.slist)
                data['mode'].append(myBeam.prm.ifmode)
                data['gsflg'].append(myBeam.fit.gflg)
                # To save time and RAM, only keep the data specified
                # in params.
                if('velocity' in params):
                    data['vel'].append(myBeam.fit.v)
                if('power' in params):
                    data['pow'].append(myBeam.fit.p_l)
                if('width' in params):
                    data['wid'].append(myBeam.fit.w_l)
                if('elevation' in params):
                    data['elev'].append(myBeam.fit.elv)
                if('phi0' in params):
                    data['phi0'].append(myBeam.fit.phi0)
                if('velocity_error' in params):
                    data['velocity_error'].append(myBeam.fit.v_e)

        myBeam = myPtr.readRec()
    return data

def rti_panel(ax, data_dict, pArr, gsct, rad, bmnum, coords, cmap,
              norm, plot_terminator=True):
    """Plots the data given by pArr to an axis object.

    Parameters
    ----------
    ax :
        a MPL axis object to plot to
    data_dict :
        the data dictionary returned by pydarn.plotting.read_data
    pArr : list
        the list of data to be plotted (e.g. data_dict['vel'] for
        velocity)
    gsct : bool
        a boolean stating whether to flag ground scatter data or not
    rad : str
        the 3 letter radar code
    bmnum : int
        The beam number of the data to plot
    coords : str
        plotting coordinates ('gate', 'range', 'geo', 'mag')
    cmap :
        a matplotlib.colors.ListedColormap (such as that returned
        by utils.plotUtils.genCmap)
    norm :
        a matplotlib.colors.BoundaryNorm (such as that returned by
        utils.plotUtils.genCmap)
    plot_terminator : Optional[bool]
        A boolean stating whether or not to plot the terminator; default
        is true.

    Returns
    -------
    pcoll
        the polygon collection returned by matplotib.pyplot.pcolormesh.

    Written by ASR 20150916

    """
    from davitpy import pydarn
    import matplotlib
    from matplotlib.dates import date2num, num2date
    import numpy as np

    # Initialize things.
    rmax = max(data_dict['nrang'])
    tmax = (len(data_dict['datetime'])) * 2
    data = np.zeros((tmax, rmax)) * np.nan
    x = np.zeros(tmax)
    tcnt = 0

    # Build a list of datetimes to plot each data point at.
    dt_list = []
    for i in range(len(data_dict['datetime'])):
        x[tcnt] = date2num(data_dict['datetime'][i])
        dt_list.append(data_dict['datetime'][i])

        if(i < len(data_dict['datetime']) - 1):
            if(date2num(data_dict['datetime'][i + 1]) - x[tcnt] > 4. / 1440.):
                tcnt += 1
                # 1440 minutes in a day, hardcoded 1 minute step per data point
                # but only if time between data points is > 4 minutes
                x[tcnt] = x[tcnt - 1] + 1. / 1440.
                dt_list.append(num2date(x[tcnt]))
        tcnt += 1

        if(pArr[i] == [] or pArr[i] is None): continue

        if data_dict['slist'][i] is not None:
            for j in range(len(data_dict['slist'][i])):
                if(not gsct or data_dict['gsflg'][i][j] == 0):
                    data[tcnt][data_dict['slist'][i][j]] = pArr[i][j]
                elif gsct and data_dict['gsflg'][i][j] == 1:
                    data[tcnt][data_dict['slist'][i][j]] = -100000.

    # For geo or mag coords, get radar FOV lats/lons.
    if (coords != 'gate' and coords != 'rng') or plot_terminator is True:
        site = pydarn.radar.network().getRadarByCode(rad) \
            .getSiteByDate(data_dict['datetime'][0])
        myFov = pydarn.radar.radFov.fov(site=site, ngates=rmax,
                                        nbeams=site.maxbeam,
                                        rsep=data_dict['rsep'][0],
                                        coords=coords,
                                        date_time=data_dict['datetime'][0])
        myLat = myFov.latCenter[bmnum]
        myLon = myFov.lonCenter[bmnum]

    # Determine the yaxis range limits to plot data to.
    if(coords == 'gate'):
        y = np.linspace(0, rmax, rmax + 1)
    elif(coords == 'rng'):
        y = np.linspace(data_dict['frang'][0],
                           rmax * data_dict['rsep'][0],
                           rmax + 1)
    else:
        y = myFov.latFull[bmnum]

    # Generate a mesh of x and y coords to plot data to.
    X, Y = np.meshgrid(x[:tcnt], y)

    # Calculate terminator as required.
    if plot_terminator:
        daylight = np.ones([len(dt_list), len(myLat)], np.bool)
        for tm_inx in range(len(dt_list)):
            tm = dt_list[tm_inx]
            term_lats, tau, dec = daynight_terminator(tm, myLon)
            if dec > 0:
                # NH Summer
                day_inx = np.where(myLat < term_lats)[0]
            else:
                day_inx = np.where(myLat > term_lats)[0]

            if day_inx.size != 0:
                daylight[tm_inx, day_inx] = False
        daylight = np.ma.array(daylight, mask=daylight)
        ax.pcolormesh(X, Y, daylight.T, lw=0, alpha=0.10,
                      cmap=matplotlib.cm.binary_r, zorder=99)

    # Mask the nan's in the data array so they aren't plotted.
    Zm = np.ma.masked_where(np.isnan(data[:tcnt][:].T), data[:tcnt][:].T)
    # Set colormap so that masked data (bad) is transparent.
    cmap.set_bad('w', alpha=0.0)

    # Now let's plot all data.
    pcoll = ax.pcolormesh(X, Y, Zm, lw=0.01, edgecolors='None',
                          cmap=cmap, norm=norm)

    return pcoll

def daynight_terminator(date, lons):
    """ Return the coordinates of day/night terminator for RTI plotting.

    Parameters
    ----------
    date : datetime.datetime
        a datetime.datetime object (assumed UTC)
    lons : list
        a numpy array of lons

    Returns
    -------
    lat
        the latitude of the day night terminator
    tau
        grenwich hour angle
    dec
        solar declination

    """
    import mpl_toolkits.basemap.solar as solar
    import numpy as np

    dg2rad = np.pi / 180.
    # compute greenwich hour angle and solar declination
    # from datetime object (assumed UTC).
    tau, dec = solar.epem(date)
    # compute day/night terminator from hour angle, declination.
    longitude = lons + tau
    lats = np.arctan(-np.cos(longitude * dg2rad) /
                     np.tan(dec * dg2rad)) / dg2rad
    return lats, tau, dec


"""
This module contains custom filter backends for Django REST Framework.

The backends included in this module are:

- TodaySeanceFilterBackend: A filter backend that filters seances by the ones that occur today.
- DateRangeFilterBackend: A filter backend that filters seances by a date range.

"""

from datetime import datetime, date, time, timedelta
from django_filters import rest_framework as filters
from rest_framework.filters import BaseFilterBackend

class TodaySeanceFilterBackend(BaseFilterBackend):
    """
    A filter backend that filters seances by the ones that occur today.

    This filter backend retrieves the seances that occur on the current date. It does this by
    comparing the date of the seance with the current date. If a seance's date is equal to the current
    date, then it is included in the queryset.

    Attributes:
        None

    Methods:
        filter_queryset(request, queryset, view): Filters the queryset based on the query parameters
            provided in the request object. Returns the filtered queryset.
    """

    def filter_queryset(self, request, queryset, view):
        """
            Filters the queryset based on the query parameters provided in the request object.
            Returns the filtered queryset.

            Args:
                request (Request): The request object containing the query parameters.
                queryset (QuerySet): The queryset to filter.
                view (APIView): The view class that called the filter.

            Returns:
                QuerySet: The filtered queryset.
            """

        today = date.today()
        start_time__gte = datetime.combine(today, datetime.min.time())
        start_time__lte = datetime.combine(today, datetime.max.time())

        hall = request.query_params.get('hall')
        start_time = request.query_params.get('start_time')
        end_time = request.query_params.get('end_time')

        if hall:
            queryset = queryset.filter(hall=hall)

        if start_time and end_time:
            start_time = datetime.strptime(start_time, '%Y-%m-%dT%H:%M:%S')
            end_time = datetime.strptime(end_time, '%Y-%m-%dT%H:%M:%S')
            queryset = queryset.filter(start_time__range=(start_time, end_time))

        if start_time and not end_time:
            start_time_exact = datetime.strptime(start_time, '%Y-%m-%dT%H:%M:%S')
            queryset = queryset.filter(
                start_time__gte=start_time__gte,
                start_time__lte=start_time__lte,
                start_time__time__gte=start_time_exact.time()
            )

        return queryset

class DateRangeFilterBackend(BaseFilterBackend):
    """
    A filter backend that filters seances by a date range.

    This filter backend retrieves the seances that occur between two dates. It does this by
    filtering the queryset by the start_time that is greater than or equal to the start_time__gte
    parameter and less than or equal to the start_time__lte parameter.

    Attributes:
        None

    Methods:
        filter_queryset(request, queryset, view): Filters the queryset based on the query parameters
            provided in the request object. Returns the filtered queryset.
    """

    def filter_queryset(self, request, queryset, view):
        """
        Filters the queryset based on the query parameters provided in the request object.
        Returns the filtered queryset.

        Args:
            request (Request): The request object containing the query parameters.
            queryset (QuerySet): The queryset to filter.
            view (APIView): The view class that called the filter.

        Returns:
            QuerySet: The filtered queryset.
        """

        start_time__gte = request.query_params.get('start_time__gte')
        start_time__lte = request.query_params.get('start_time__lte')

        if start_time__gte and start_time__lte:
            start_time__gte = datetime.strptime(start_time__gte, '%Y-%m-%dT%H:%M:%S')
            start_time__lte = datetime.strptime(start_time__lte, '%Y-%m-%dT%H:%M:%S')
            queryset = queryset.filter(start_time__gte=start_time__gte, start_time__lte=start_time__lte)

        return queryset


import datetime

from uuid import uuid4

from django.utils.timezone import now

from faker import Faker

from rest_framework import status
from rest_framework.reverse import reverse
from rest_framework.test import APIClient

from datahub.company.test.factories import AdviserFactory
from datahub.core.test_utils import (
    APITestMixin,
    format_date_or_datetime,
)

from datahub.task.test.factories import TaskFactory
from datahub.task.test.utils import BaseEditTaskTests, BaseListTaskTests, BaseTaskTests

class TestListTask(BaseListTaskTests):
    """Test the LIST task endpoint"""

    reverse_url = 'api-v4:task:collection'

    def test_get_all_tasks_returns_ordered_results_when_sortby_param_is_used(self):
        earliest_task = TaskFactory(due_date=datetime.date.today() - datetime.timedelta(days=1))
        latest_task = TaskFactory(due_date=datetime.date.today() + datetime.timedelta(days=1))
        middle_task = TaskFactory(due_date=datetime.date.today())

        url = f'{reverse(self.reverse_url)}?sortby=due_date'

        response = self.api_client.get(url).json()

        ordered_ids_response = [result['id'] for result in response['results']]
        assert ordered_ids_response == [
            str(earliest_task.id),
            str(middle_task.id),
            str(latest_task.id),
        ]

class TestGetTask(APITestMixin):
    """Test the GET task endpoint"""

    def test_get_task_return_404_when_task_id_unknown(self):
        TaskFactory()

        url = reverse(
            'api-v4:task:item',
            kwargs={'pk': uuid4()},
        )
        response = self.api_client.get(url)
        assert response.status_code == status.HTTP_404_NOT_FOUND

    def test_get_task_return_task_when_task_id_valid(self):
        task = TaskFactory()

        url = reverse(
            'api-v4:task:item',
            kwargs={'pk': task.id},
        )
        response = self.api_client.get(url).json()
        expected_response = {
            'id': str(task.id),
            'title': task.title,
            'description': task.description,
            'due_date': task.due_date,
            'reminder_days': task.reminder_days,
            'email_reminders_enabled': task.email_reminders_enabled,
            'advisers': [
                {
                    'id': str(adviser.id),
                    'name': adviser.name,
                    'first_name': adviser.first_name,
                    'last_name': adviser.last_name,
                }
                for adviser in task.advisers.all()
            ],
            'archived': task.archived,
            'archived_by': task.archived_by,
            'archived_reason': task.archived_reason,
            'created_by': {
                'name': task.created_by.name,
                'first_name': task.created_by.first_name,
                'last_name': task.created_by.last_name,
                'id': str(task.created_by.id),
            },
            'modified_by': {
                'name': task.modified_by.name,
                'first_name': task.modified_by.first_name,
                'last_name': task.modified_by.last_name,
                'id': str(task.modified_by.id),
            },
            'created_on': format_date_or_datetime(task.created_on),
            'modified_on': format_date_or_datetime(task.modified_on),
            'investment_project_task': None,
        }
        assert response == expected_response

class TestAddTask(APITestMixin):
    """Test the POST task endpoint"""

    def test_create_task_with_missing_mandatory_fields_returns_bad_request(self):
        url = reverse('api-v4:task:collection')

        response = self.api_client.post(
            url,
            data={},
        )
        assert response.status_code == status.HTTP_400_BAD_REQUEST
        assert list(response.json().keys()) == ['title', 'advisers']

    def test_create_task_with_unknown_advisor_id_returns_bad_request(self):
        faker = Faker()

        url = reverse('api-v4:task:collection')

        response = self.api_client.post(
            url,
            data={'title': faker.word(), 'advisers': [uuid4()]},
        )

        assert response.status_code == status.HTTP_400_BAD_REQUEST
        assert list(response.json().keys()) == ['advisers']

    def test_create_task_with_empty_advisors_returns_bad_request(self):
        faker = Faker()

        url = reverse('api-v4:task:collection')

        response = self.api_client.post(
            url,
            data={'title': faker.word(), 'advisers': []},
        )

        assert response.status_code == status.HTTP_400_BAD_REQUEST
        assert list(response.json().keys()) == ['advisers']

    def test_create_task_with_valid_mandatory_fields_returns_success_created(self):
        faker = Faker()

        adviser = AdviserFactory()

        url = reverse('api-v4:task:collection')

        new_task = {
            'title': faker.word(),
            'description': faker.word(),
            'due_date': now().date(),
            'reminder_days': 3,
            'email_reminders_enabled': True,
            'advisers': [adviser.id],
        }

        post_response = self.api_client.post(
            url,
            data=new_task,
        )
        post_response_json = post_response.json()

        assert post_response.status_code == status.HTTP_201_CREATED

        get_url = reverse(
            'api-v4:task:item',
            kwargs={'pk': post_response_json['id']},
        )
        get_response = self.api_client.get(get_url)

        assert get_response.status_code == status.HTTP_200_OK
        expected_response = {
            'id': post_response_json['id'],
            'title': post_response_json['title'],
            'description': post_response_json['description'],
            'due_date': post_response_json['due_date'],
            'reminder_days': post_response_json['reminder_days'],
            'email_reminders_enabled': post_response_json['email_reminders_enabled'],
            'advisers': [
                {
                    'id': str(adviser.id),
                    'name': adviser.name,
                    'first_name': adviser.first_name,
                    'last_name': adviser.last_name,
                },
            ],
            'archived': post_response_json['archived'],
            'archived_by': post_response_json['archived_by'],
            'archived_reason': post_response_json['archived_reason'],
            'created_by': {
                'name': self.user.name,
                'first_name': self.user.first_name,
                'last_name': self.user.last_name,
                'id': str(self.user.id),
            },
            'modified_by': {
                'name': self.user.name,
                'first_name': self.user.first_name,
                'last_name': self.user.last_name,
                'id': str(self.user.id),
            },
            'created_on': post_response_json['created_on'],
            'modified_on': post_response_json['modified_on'],
            'investment_project_task': None,
        }
        assert get_response.json() == expected_response

class TestEditTask(BaseEditTaskTests):
    reverse_url = 'api-v4:task:item'

    """Test the PATCH task endpoint"""

    def test_edit_task_return_404_when_task_id_unknown(self):
        url = reverse('api-v4:task:item', kwargs={'pk': uuid4()})

        response = self.api_client.patch(
            url,
            data={'title': 'abc'},
        )
        assert response.status_code == status.HTTP_404_NOT_FOUND

    def test_edit_task_with_valid_fields_returns_success(self):
        adviser = AdviserFactory()
        task = TaskFactory(created_by=adviser)
        new_adviser = AdviserFactory()

        url = reverse('api-v4:task:item', kwargs={'pk': task.id})

        response = self.adviser_api_client(adviser).patch(
            url,
            data={'advisers': [new_adviser.id]},
        )
        assert response.status_code == status.HTTP_200_OK
        assert response.json()['advisers'][0]['id'] == str(new_adviser.id)

class TestArchiveTask(BaseTaskTests):
    """Test the archive POST endpoint for task"""

    def test_archive_task_without_reason_returns_bad_request(self):
        adviser = AdviserFactory()
        task = TaskFactory(created_by=adviser)

        url = reverse('api-v4:task:task_archive', kwargs={'pk': task.id})

        response = self.adviser_api_client(adviser).post(url, data={})

        assert response.status_code == status.HTTP_400_BAD_REQUEST
        assert response.data == {
            'reason': ['This field is required.'],
        }

    def test_archive_task_returns_unauthorized_when_user_not_authenticated(self):
        adviser = AdviserFactory()
        task = TaskFactory(advisers=[adviser])

        url = reverse('api-v4:task:task_archive', kwargs={'pk': task.id})

        response = APIClient().post(url, data={'reason': 'completed'})

        assert response.status_code == status.HTTP_401_UNAUTHORIZED

    def test_archive_task_returns_success_when_user_is_creator_but_not_assigned_to_task(self):
        adviser = AdviserFactory()
        task = TaskFactory(created_by=adviser)

        url = reverse('api-v4:task:task_archive', kwargs={'pk': task.id})

        response = self.adviser_api_client(adviser).post(url, data={'reason': 'completed'})

        assert response.status_code == status.HTTP_200_OK

    def test_archive_task_returns_success_when_user_is_not_creator_but_is_assigned_to_task(
        self,
    ):
        adviser = AdviserFactory()
        task = TaskFactory(advisers=[adviser])

        url = reverse('api-v4:task:task_archive', kwargs={'pk': task.id})

        response = self.adviser_api_client(adviser).post(url, data={'reason': 'completed'})

        assert response.status_code == status.HTTP_200_OK

    def test_archive_task_returns_success_when_user_is_creator_and_is_assigned_to_task(self):
        adviser = AdviserFactory()
        task = TaskFactory(advisers=[adviser], created_by=adviser)

        url = reverse('api-v4:task:task_archive', kwargs={'pk': task.id})

        response = self.adviser_api_client(adviser).post(url, data={'reason': 'completed'})

        assert response.status_code == status.HTTP_200_OK

    def test_archive_task_with_valid_reason_returns_success(self):
        adviser = AdviserFactory()
        task = TaskFactory(created_by=adviser)
        url = reverse('api-v4:task:task_archive', kwargs={'pk': task.id})
        response = self.adviser_api_client(adviser).post(url, data={'reason': 'completed'})

        assert response.status_code == status.HTTP_200_OK
        assert response.data['archived']
        assert response.data['archived_by'] == {
            'id': str(adviser.id),
            'first_name': adviser.first_name,
            'last_name': adviser.last_name,
            'name': adviser.name,
        }
        assert response.data['archived_reason'] == 'completed'
        assert response.data['id'] == str(task.id)


from __future__ import absolute_import
from __future__ import unicode_literals

import logging
import os
import signal
import time
import traceback
from multiprocessing import Event
from multiprocessing import Lock
from multiprocessing import Process

import six
from kafka import KafkaConsumer
from kafka.common import ConsumerTimeout
from retrying import retry

from yelp_kafka import metrics
from yelp_kafka.consumer import KafkaSimpleConsumer
from yelp_kafka.error import ConsumerGroupError
from yelp_kafka.error import PartitionerError
from yelp_kafka.error import PartitionerZookeeperError
from yelp_kafka.error import ProcessMessageError
from yelp_kafka.metrics_responder import MetricsResponder
from yelp_kafka.partitioner import Partitioner
from yelp_kafka.utils import get_default_responder_if_available
from yelp_kafka.utils import retry_if_kafka_unavailable_error

DEFAULT_REFRESH_TIMEOUT_IN_SEC = 0.5
CONSUMER_GROUP_INTERNAL_TIMEOUT = 100  # milliseconds

class ConsumerGroup(object):
    """Single process consumer group.
    Support partitions distribution for a single topic
    between many consumer instances.

    If the topic consists of only one partition, only one consumer belonging
    to the group will be able to consume messages. The other consumers will
    stay idle, ready to take over the active consumer in case of failures.
    If the topic consists of many partitions and only one consumer has
    joined the group, the consumer will consume messages
    from all the partitions. Finally if the topic consists of many
    partitions and there are many consumers, each consumer will pick up
    a subset of partitions and consume from them.

    .. warning::
        Messages will be named tuples with fields "partition",
        "offset", "key", "value". This message format is different from those
        yielded by :py:class:`yelp_kafka.consumer_group.KafkaConsumerGroup`.

    .. warning:: This class is considered deprecated in favor of :py:class:`yelp_kafka.consumer_group.KafkaConsumerGroup`.

    Example:

    .. code-block:: python

       from yelp_kafka import discovery
       from yelp_kafka.config import KafkaConsumerConfig
       from yelp_kafka.consumer_group import ConsumerGroup

       cluster = discovery.get_local_cluster('standard')
       config = KafkaConsumerConfig('my_group', cluster)

       def my_process_function(message):
           partition, offset, key, value = message
           print partition, offset, key, value

       consumer = ConsumerGroup('test_topic', config, my_process_function)
       consumer.run()

    :param topics: topics to consume from
    :type topics: list
    :param config: yelp_kakfa config. See :py:mod:`yelp_kafka.config`
    :type config: dict
    :param process_func: function used to process the message.
        This function should accept a :py:data:`yelp_kafka.consumer.Message`
        as parameter.
    """

    def __init__(self, topic, config, process_func):
        self.log = logging.getLogger(self.__class__.__name__)
        self.config = config
        self.topic = topic
        self.partitioner = Partitioner(
            config,
            topic if isinstance(topic, list) else [topic],
            self._acquire,
            self._release
        )
        self.consumer = None
        self.process = process_func

    def run(self, refresh_timeout=DEFAULT_REFRESH_TIMEOUT_IN_SEC):
        """Create the group, instantiate a consumer and consume message
        from kafka. The refresh timeout shouldn't be too high in order to
        allow the consumers to stay updated with the group changes.

        :param refresh_timeout: waiting timeout in secs to refresh
            the consumer group. It should never be greater than
            'zk_partitioner_cooldown' to avoid wasting time resetting
            the partitioner many times upon changes.
            See :py:mod:`yelp_kafka.config`
            Default: 5 seconds.
        """
        with self.partitioner:
            while True:
                self.consume(refresh_timeout)

    def consume(self, refresh_timeout):
        """Consume messages from kafka and refresh the group
        upon timeout expiration.

        :param refresh_timeout: refresh period for consumer group
        """
        timeout = time.time() + refresh_timeout
        if self.consumer:
            for message in self.consumer:
                try:
                    self.process(message)
                except:
                    trace = traceback.format_exc()
                    self.log.exception(
                        "Error processing message: %s",
                        message,
                    )
                    raise ProcessMessageError(
                        "Error processing message: {0}\n\nException "
                        "caught inside processing message function:\n{1}"
                        .format(
                            message,
                            trace,
                        )
                    )
                if time.time() > timeout:
                    break
        try:
            self.partitioner.refresh()
        except (PartitionerZookeeperError, PartitionerError):
            self.log.exception("Encountered a partitioner error")
            raise

    def _acquire(self, partitions):
        """Create a consumer ready to consume from kafka.

        .. note: This function should only by called by the partitioner.
           The partitioner executes _acquire when the partition group
           changes and the partitions have been acquired.
        """
        if partitions.get(self.topic):
            self.consumer = KafkaSimpleConsumer(
                self.topic,
                self.config,
                partitions[self.topic]
            )
            try:
                # We explicitly catch and log the exception.
                self.consumer.connect()
            except:
                self.log.exception(
                    "Consumer topic %s, partition %s, config %s:"
                    " failed connecting to kafka",
                    self.topic,
                    partitions,
                    self.config,
                )
                raise

    def _release(self, partitions):
        """Release the consumer.

        .. note: This function should only by called by the partitioner.
           The partitioner executes _release when it has to release the
           acquired partitions.
        """
        if self.consumer:
            self.consumer.close()
            self.consumer = None

class KafkaConsumerGroup(object):
    """KafkaConsumerGroup allows you to efficiently consume from all the
    partitions of a topic without having to manually use the
    :py:class:`yelp_kafka.partitioner.Partitioner` class. You can spin up
    multiple KafkaConsumerGroups, and they will co-ordinate via the Partitioner
    to divvy up the available partitions between each other.

    This class works by attempting to rebalance before each call to `next()`. In
    the event that rebalancing does occur and that you have enabled
    auto-committing, any messages marked as done using `task_done()` will be
    committed before repartitioning. To commit messages immediately, you can
    call `commit()`.

    If metrics_reporter is enabled in config, the metrics prefix in
    will be: "yelp_kafka.KafkaConsumerGroup."

    .. warning::
        Do not create multiple KafkaConsumerGroups in the same process;
        the Partitioner class does not work if there are multiple instances of it in
        the same process.

    .. warning::
        Messages will be instances of kafka-python KafkaMessage, which is
        a named tuple with the fields "topic", "partition", "offset", "key",
        "value". This message format is different from those yielded by
        :py:class:`yelp_kafka.consumer_group.ConsumerGroup`.

    :param topics: a list of topics to consume from.
    :type topics: list
    :param config: yelp_kakfa consumer config.
    :type config: :py:class:`yelp_kafka.config.KafkaConsumerConfig`
    :param metrics_responder: A metric responder to report metrics, defaults to
        use :py:class:`yelp_kafka.yelp_metrics_responder.MeteoriteMetricsResponder`, if
        the import of yelp_meteorite is successful.
    :type metrics_responder: class which implements metric_responder.MetricsResponder
    """

    METRIC_PREFIX = 'yelp_kafka.KafkaConsumerGroup'

    def __init__(self, topics, config, metrics_responder=None):
        assert isinstance(topics, list), "Topics must be a list"

        self.log = logging.getLogger(self.__class__.__name__)
        self.topics = topics
        self.partitioner = Partitioner(
            config,
            topics,
            self._acquire,
            self._release
        )
        self.consumer = None

        consumer_config = config.get_kafka_consumer_config()
        self.log.debug(
            "Using config: %s",
            consumer_config,
        )

        self.metrics_responder = metrics_responder or get_default_responder_if_available()
        assert not metrics_responder or isinstance(metrics_responder, MetricsResponder), \
            "Metric Reporter is not of type yelp_kafka.metrics_responder.MetricsResponder"

        if self.metrics_responder:
            self._setup_metrics_responder(config)
            consumer_config['metrics_responder'] = self._send_to_metrics_responder

        self.pre_rebalance_callback = config.pre_rebalance_callback
        self.post_rebalance_callback = config.post_rebalance_callback

        # Intercept the user's timeout and pass in our own instead. We do this
        # in order to periodically refresh the partitioner when calling next()
        self.iter_timeout = consumer_config['consumer_timeout_ms']
        consumer_config['consumer_timeout_ms'] = CONSUMER_GROUP_INTERNAL_TIMEOUT
        self.config = consumer_config

    def _setup_metrics_responder(self, config):
        extra_dimensions = config.metrics_dimensions
        self.timers = {}
        for name in metrics.TIME_METRIC_NAMES:
            topic_name = self.METRIC_PREFIX + '.' + name
            timer = self.metrics_responder.get_timer_emitter(topic_name, extra_dimensions)
            self.timers[name] = timer

        self.counters = {}
        for name in metrics.FAILURE_COUNT_METRIC_NAMES:
            topic_name = self.METRIC_PREFIX + '.' + name
            counter = self.metrics_responder.get_counter_emitter(topic_name, extra_dimensions)
            self.counters[name] = counter

    def _send_to_metrics_responder(self, key, value):
        if key in self.timers:
            # kafka-python emits time in seconds, but yelp_meteorite wants
            # milliseconds
            time_in_ms = value * 1000
            self.metrics_responder.record(self.timers[key], time_in_ms)
        elif key in self.counters:
            self.metrics_responder.record(self.counters[key], 1)
        else:
            self.log.warn("Unknown metric: {0}".format(key))

    def start(self):
        self.partitioner.start()

    def stop(self):
        self.partitioner.stop()
        self.consumer.close()

    def next(self):
        start_time = time.time()
        while self._should_keep_trying(start_time):
            self.partitioner.refresh()
            try:
                return self.consumer.next()
            except ConsumerTimeout:
                # This is due to the internal timeout, not the user's provided
                # one.
                pass
        error_msg = "KafkaConsumerGroup timed out after {0} ms"
        raise ConsumerTimeout(error_msg.format(self.iter_timeout))

    def _should_keep_trying(self, start_time):
        if self.iter_timeout < 0:
            return True
        elapsed_seconds = time.time() - start_time
        return elapsed_seconds * 1000 < self.iter_timeout

    def task_done(self, message):
        return self.consumer.task_done(message)

    def commit(self):
        return self.consumer.commit()

    def _acquire(self, partitions):
        if not self.consumer:
            self.consumer = KafkaConsumer(partitions, **self.config)
        else:
            self.consumer.set_topic_partitions(partitions)
        if self.post_rebalance_callback:
            self.post_rebalance_callback(partitions)

    # set_topic_partitions causes a metadata request, which may fail on the
    # first try.
    @retry(stop_max_attempt_number=2, retry_on_exception=retry_if_kafka_unavailable_error)
    def _release(self, partitions):
        if self.pre_rebalance_callback:
            self.pre_rebalance_callback(partitions)
        if self._auto_commit_enabled():
            self.consumer.commit()
        self.consumer.set_topic_partitions({})

    def _auto_commit_enabled(self):
        return self.config['auto_commit_enable']

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, type, value, traceback):
        self.stop()

    def __iter__(self):
        return self

    def __next__(self):
        return self.next()

class MultiprocessingConsumerGroup(object):
    """Multiprocessing consumer group allows to consume
    from multiple topics at once. It spawns a python process
    for each assigned partition.
    It also implements monitoring for the running consumers and
    is able to restart restart these upon failures.

    .. note: This class is thread safe.

    Example:

    .. code-block:: python

       from threading import Thread
       from yelp_kafka import discovery
       from yelp_kafka.config import KafkaConsumerConfig
       from yelp_kafka.consumer import KafkaConsumer
       from yelp_kafka.consumer_group import MultiprocessingConsumerGroup

       class MyConsumer(KafkaConsumerBase):

           def __init__(topic, config, partitions):
               super(MyConsumer, self).__init__(topic, config, partitions)

           def initialize(self):
               print "Initializing.. Usually you want to set the logger here"

           def dispose(self):
               print "Dying..."

           def process(self, message):
               partition, offset, key, value = message
               print partition, offset, key, value

       cluster = discovery.get_local_cluster('standard')
       config = KafkaConsumerConfig('my_group', cluster)

       group = MultiprocessingConsumerGroup(
           ['topic1', 'topic2'],
           config,
           MyConsumer,
       )
       group_thread = Thread(target=group.start_group)
       group_thread.start()
       # Do some other cool stuff here like sleep
       time.sleep(600)
       group.stop_group()

    :param topics: a list of topics to consume from.
    :type topics: list
    :param config: yelp_kakfa config. See :py:mod:`yelp_kafka.config`
    :type config: dict
    :param consumer_factory: the function used to instantiate the consumer.
        the consumer_factory must have the same argument list of
        :py:class:`yelp_kafka.consumer.KafkaConsumerBase`. It has to return
        an instance of a subclass of
        :py:class:`yelp_kafka.consumer.KafkaConsumerBase`.
    """

    def __init__(self, topics, config, consumer_factory):
        self.config = config
        self.termination_flag = None
        self.partitioner = Partitioner(
            config,
            topics,
            self.acquire,
            self.release,
        )
        self.consumers = None
        self.consumers_lock = Lock()
        self.consumer_procs = {}
        self.consumer_factory = consumer_factory
        self.log = logging.getLogger(self.__class__.__name__)
        self.pre_rebalance_callback = config.pre_rebalance_callback
        self.post_rebalance_callback = config.post_rebalance_callback

    def start_group(self, refresh_timeout=DEFAULT_REFRESH_TIMEOUT_IN_SEC):
        """Start the consumer group.

        :param refresh_timeout: waiting timeout in secs to refresh
            the consumer group. It should never be greater than
            'zk_partitioner_cooldown' to avoid wasting time resetting
            the partitioner many times upon changes.
            See :py:mod:`yelp_kafka.config`
            Default: 5 seconds.

        .. note: this function does not return. You may want to run it into
            a separate thread.

        """
        # Create the termination flag
        self.termination_flag = Event()

        with self.partitioner:
            while not self.termination_flag.is_set():
                self.termination_flag.wait(refresh_timeout)
                self.monitor()
                try:
                    self.partitioner.refresh()
                except (PartitionerZookeeperError, PartitionerError):
                    self.log.exception("Encountered a partitioner error")
                    raise

    def stop_group(self):
        """Set the termination flag to stop the group.

        :raises: ConsumerGroupError is the group has not been started, yet.
        """
        if not self.termination_flag:
            raise ConsumerGroupError("Group not running")
        self.termination_flag.set()

    def acquire(self, partitions):
        """Acquire kafka topics-[partitions] and start the
        consumers for them.
        """
        self.log.debug("Acquired partitions: %s", partitions)
        with self.consumers_lock:
            self.consumers = list(self.start(
                partitions
            ))
            self.log.debug("Allocated consumers %s", self.consumers)
        if self.post_rebalance_callback:
            self.post_rebalance_callback(partitions)

    def start(self, acquired_partitions):
        """Start a consumer process for each acquired partition.
        Use consumer_factory to create a consumer instance. Then
        start a new process on the method run.

        :param acquired_partitions: acquired topics partitions
        :type: dict {<topic>: <[partitions]>}
        """
        # TODO KAFKA-72: We create a consumer process for each partition.
        # There can be too many consumers in case we have a lot of
        # partitions and just a worker instance.
        # We should make this smarter and allow the user to decide
        # how many processes/partitions allocate.
        for topic, partitions in six.iteritems(acquired_partitions):
            for p in partitions:
                self.log.info(
                    "Creating consumer topic = %s, config = %s,"
                    " partition = %s", topic, self.config, p
                )
                consumer = self.consumer_factory(topic, self.config, [p])
                self.consumer_procs[self.start_consumer(consumer)] = consumer
        return self.consumer_procs.values()

    def start_consumer(self, consumer):
        """Create a new consumer process"""
        try:
            proc = Process(
                target=consumer.run,
                name='Consumer-{0}-{1}'.format(
                    consumer.topic,
                    consumer.partitions,
                ),
            )
            proc.daemon = True
            proc.start()
        except Exception:
            self.log.error(
                "Impossible to start a new consumer."
                "Topic: %s. Partition: %s.",
                consumer.topic,
                consumer.partitions,
            )
            raise ConsumerGroupError(
                "Error starting a new consumer."
                "Topic: {topic}. Partition: {partitions}.".format(
                    topic=consumer.topic,
                    partitions=consumer.partitions,
                ),
            )
        return proc

    def release(self, partitions):
        """Terminate all the consumer processes"""
        if self.pre_rebalance_callback:
            self.pre_rebalance_callback(partitions)
        self.log.info("Terminating consumer group")
        for consumer in six.itervalues(self.consumer_procs):
            consumer.terminate()

        timeout = time.time() + self.config.max_termination_timeout_secs
        while (
            time.time() <= timeout and any(
                [proc.is_alive() for proc in six.iterkeys(self.consumer_procs)]
            )
        ):
            continue

        for proc, consumer in six.iteritems(self.consumer_procs):
            if proc.is_alive():
                os.kill(proc.pid, signal.SIGKILL)
                self.log.error(
                    "Process %s, topic %s, partitions %s:"
                    "killed due to timeout",
                    proc.name,
                    consumer.topic,
                    consumer.partitions,
                )
        self.consumer_procs.clear()
        with self.consumers_lock:
            self.consumers = None

    def monitor(self):
        """Respawn consumer processes upon failures."""
        # We don't use the iterator because the dict may change during the loop
        for proc, consumer in self.consumer_procs.items():
            if not proc.is_alive():
                self.log.error(
                    "consumer process %s topic %s partitions %s: "
                    "died exit status %s",
                    proc.name,
                    consumer.topic,
                    consumer.partitions, proc.exitcode,
                )
                # Restart consumer process
                self.consumer_procs[self.start_consumer(consumer)] = consumer
                # Remove dead process
                del self.consumer_procs[proc]

    def get_consumers(self):
        """Get a copy of the allocated consumers.

        :returns: A copy of allocated consumers
        :rtype: list
        """
        with self.consumers_lock:
            if self.consumers is not None:
                return self.consumers[:]
            return None


from telegram.ext import Updater, InlineQueryHandler, CommandHandler
import requests
# import re
import sqlite3
from scraping import get_product_info_flipkart
import telegram.ext
from datetime import timedelta
import logging
from time import sleep

def gather_name_list(chat_id):
    logging.info("Gathering list of items for "+str(chat_id))
    conn = sqlite3.connect("telegram_database.db")
    c = conn.cursor()

    table_name = "_" + str(chat_id)

    c.execute(
        'create table if not exists ' + table_name + ' (link text,name text,price_with_currency text,price real,in_stock text)')
    c.execute('SELECT * FROM ' + table_name)
    lis = c.fetchall()
    product_names = []
    for j in lis:
        product_names.append(j[1])
    # print(product_names)
    product_names.sort()
    list_message = "You are tracking the following products:\n\n"
    for i in enumerate(product_names):
        list_message += str(i[0] + 1) + ". " + i[1] + "\n\n"
    list_message = list_message.strip() + ""
    conn.commit()
    conn.close()
    return (list_message, product_names)

def add(bot, update):
    chat_id = update.message.chat_id
    logging.info(str(chat_id)+" is trying add a product.")
    # print(chat_id)
    try:
        link = str(update.message.text).strip().split()[1]
        # print(link)
    except IndexError:
        logging.info(str(chat_id)+" didn't enter any product link.")
        bot.send_message(
            chat_id=chat_id, text="Please enter a flipkart product link.", parse_mode="html")
        return
    if "flipkart.com" not in link:
        logging.info(str(chat_id)+" didn't enter a valid product link")
        bot.send_message(chat_id=chat_id, text=""+link +
                         " Doesn't seem to be a valid flipkart link.", parse_mode="html")
        return
    conn = sqlite3.connect("telegram_database.db")
    c = conn.cursor()
    table_name = "_"+str(chat_id)
    c.execute(
        'create table if not exists ' + table_name + ' (link text,name text,price_with_currency text,price real,in_stock text)')
    conn.commit()
    conn.close()
    try:
        product_info = get_product_info_flipkart(link)
    except IndexError:
        logging.info(str(chat_id)+" didn't enter a valid product link")
        bot.send_message(chat_id=chat_id, text=link +
                         " is not a link to product on flipkart.\nEnter / help commands to view all commands and their formats")
    if product_info:
        conn = sqlite3.connect("telegram_database.db")
        c = conn.cursor()
        c.execute('SELECT * FROM ' + table_name +
              ' WHERE name=?', (product_info["name"],))
        match = c.fetchone()
        if match:
            logging.info(str(chat_id)+" entered an already present link.")
            bot.send_message(chat_id=chat_id, text='You are already tracking '+match[1]+'')
        else:
            logging.info("Product addition was a success")
            c.execute('INSERT INTO ' + table_name + ' values (?,?,?,?,?)', (link,product_info['name'], product_info["price_with_currency"], product_info["price_in_num"], product_info['availability']))
            bot.send_message(chat_id=chat_id, text="You just added " + product_info["name"] + " for tracking.\nCurrently:\n   Price:" +product_info["price_with_currency"].rjust(10) + "\nIn Stock:" + product_info["availability"].rjust(10), parse_mode='html')

        conn.commit()
        conn.close()
    else:
        bot.send_message(chat_id=chat_id,text="Something didn't work there!! Try again after a minute or so.")

def list_names(bot, update):
    chat_id = update.message.chat_id
    logging.info("Listing "+str(chat_id)+"'s Products")
    list_message, product_names = gather_name_list(chat_id)
    if len(product_names) == 0:
        bot.send_message(
            chat_id=chat_id, text="You are not tracking any items.", parse_mode='html')
        return
    bot.send_message(chat_id=chat_id, text=list_message, parse_mode='html')

def remove(bot, update):
    chat_id = update.message.chat_id
    logging.info(str(chat_id)+" is trying to remove a product.")
    # print(chat_id)
    list_message, product_names = gather_name_list(chat_id)
    try:
        num = int(str(update.message.text).strip().split()[1])
        # print(num)
    except IndexError:
        num = 0
    except ValueError:
        num = 0

    if len(product_names) == 0:
        bot.send_message(
            chat_id=chat_id, text="You are not tracking any items.", parse_mode='html')
        return
    if not num:
        bot.send_message(chat_id=chat_id, text=list_message, parse_mode='html')
        bot.send_message(
            chat_id=chat_id, text="\nSelect the product to remove. \nEnter /remove --no. corresponding to the product--  to remove the product", parse_mode='html')
        return
    if num > len(product_names) or num < 0:
        bot.send_message(
            chat_id=chat_id, text="No Product with that number.", parse_mode='html')
        return

    conn = sqlite3.connect("telegram_database.db")
    c = conn.cursor()
    table_name = "_"+str(chat_id)
    c.execute("DELETE FROM "+table_name +
              " WHERE name=?", (product_names[num-1],))
    conn.commit()
    conn.close()
    bot.send_message(chat_id=chat_id, text="You removed " +
                     product_names[num-1]+" from tracking.", parse_mode='html')
    logging.info(str(chat_id)+" removed a product.")

def daily_checker(bot, job):
    conn = sqlite3.connect("telegram_database.db")
    c = conn.cursor()
    c.execute('SELECT name from sqlite_master where type= "table"')
    all_users = c.fetchall()
    conn.close()
    logging.info("There are currently "+str(len(all_users)) +
                 " accounts using the bot.")
    total_products = 0
    for j in all_users:

        conn = sqlite3.connect("telegram_database.db")
        c = conn.cursor()

        c.execute("SELECT * FROM "+j[0])

        all_items = c.fetchall()
        conn.close()

        chat_id = int(j[0][1:])
        # print("Doing "+str(chat_id)+"'s Tracking.")
        # print(dic)
        logging.info("Tracking "+str(chat_id)+"'s products.")
        for i in all_items:
            sleep(2)
            total_products += 1
            dic = get_product_info_flipkart(i[0])
            # print(dic)
            if not dic:
                continue
            current_price = dic['price_in_num']
            current_in_stock = dic['availability']
            price_change = current_price - i[3]
            # price_change= -1000
            # print(dic['name'])
            if price_change > 0:
                bot.send_message(chat_id=chat_id, text="**"+dic['name']+"**:\n\n"+"*6 hours back*:\n```   Price:"+i[2].rjust(10)+"\nIn Stock:"+i[4].rjust(
                    10)+"```\n*Now*:\n```   Price:"+dic["price_with_currency"].rjust(10)+"\nIn Stock:"+dic["availability"].rjust(10)+"```", parse_mode='markdown')

                bot.send_message(chat_id=chat_id, text="```diff\n-Price has increased by " +
                                 dic["price_with_currency"][0]+str(price_change)+"```", parse_mode='markdown')
                logging.info("Price of "+dic['name']+" increased.")

            elif price_change < 0:
                bot.send_message(chat_id=chat_id, text="**"+dic['name']+"**:\n\n"+"*6 hours back*:\n```   Price:"+i[2].rjust(10)+"\nIn Stock:"+i[4].rjust(
                    10)+"```\n*Now*:\n```   Price:"+dic["price_with_currency"].rjust(10)+"\nIn Stock:"+dic["availability"].rjust(10)+"```", parse_mode='markdown')

                bot.send_message(chat_id=chat_id, text="```yaml\nPrice has decreased by " +
                                 dic["price_with_currency"][0]+str(abs(price_change))+"```", parse_mode='markdown')
                logging.info("Price of "+dic['name']+" decreased.")

            else:
                logging.info("No change in price for "+dic['name'])

            if dic['availability'] == 'Yes' and i[4] == 'No':
                if price_change != 0:
                    bot.send_message(
                        chat_id=chat_id, text="```yaml\nProduct is in stock.```", parse_mode='markdown')
                else:
                    bot.send_message(chat_id=chat_id, text="**"+dic['name']+"**:\n"+"*6 hours back*:\n```   Price:"+i[2].rjust(10)+"\nIn Stock:"+i[4].rjust(
                        10)+"```\n*Now*:\n```   Price:"+dic["price_with_currency"].rjust(10)+"\nIn Stock:"+dic["availability"].rjust(10)+"```", parse_mode='markdown')
                    bot.send_message(
                        chat_id=chat_id, text="```yaml\nProduct is in stock.```", parse_mode='markdown')
                    logging.info("Product just came in to stock.")

            if price_change < 0 and dic['availability'] == "No":
                bot.send_message(
                    chat_id=chat_id, text="```diff\n-But the product is out of stock.```", parse_mode='markdown')
                logging.info("Product went out of stock.")
            if price_change != 0:
                bot.send_message(chat_id=chat_id,
                                 text="Link:\n"+i[0])
            # Updating the database with the current values.
            conn = sqlite3.connect("telegram_database.db")
            c = conn.cursor()
            c.execute("UPDATE "+j[0]+" SET price_with_currency='" +
                      dic['price_with_currency']+"' WHERE link=?", (i[0],))
            c.execute("UPDATE "+j[0]+" SET price=" +
                      str(dic["price_in_num"])+" WHERE link=?", (i[0],))
            c.execute("UPDATE "+j[0]+" SET in_stock='" +
                      dic["availability"]+"' WHERE link=?", (i[0],))
            conn.commit()
            conn.close()

    logging.info("Total number of products tracked : "+str(total_products))

def check_now(bot, update):
    chat_id = update.message.chat_id
    # print(chat_id)
    list_message, product_names = gather_name_list(chat_id)
    try:
        num = int(str(update.message.text).strip().split()[1])
        print(num)
    except IndexError:
        num = 0
    except ValueError:
        num = 0
    conn = sqlite3.connect("telegram_database.db")
    c = conn.cursor()
    table_name = "_"+str(chat_id)
    c.execute('create table if not exists ' + table_name +
              ' (link text,name text,price_with_currency text,price real,in_stock text)')
    c.execute("SELECT * FROM "+table_name)
    all_items = c.fetchall()
    logging.info("Doing "+str(chat_id)+"'s Tracking.")
    if len(all_items) == 0:
        bot.send_message(
            chat_id=chat_id, text="You are not tracking any items.", parse_mode='html')
        return
    if num == 0:
        message = "__**You are tracking the following items.**__\n"
        for j in enumerate(all_items):
            message += str(j[0]+1)+". "+j[1][1]+"\n\n"
        message += "-1. All Items\n" + \
            "Enter `/check <no. corresponding to product>` to check that product."
        bot.send_message(chat_id=chat_id, text=message, parse_mode='markdown')
    else:
        if num > len(all_items) or num < -2:
            bot.send_message(
                chat_id=chat_id, text="No Product with that number.", parse_mode='html')
            return
        else:
            if num == -1:
                for i in all_items:
                    sleep(2)
                    dic = get_product_info_flipkart(i[0])
                    if not dic:
                        bot.send_message(chat_id=chat_id,text=i[1]+" wasn't tracked due to some error. Try again after a while.")
                        continue
                    current_price = dic['price_in_num']
                    current_in_stock = dic['availability']
                    bot.send_message(chat_id=chat_id, text=""+dic['name']+":\n"+"Last Time Checked:\n   Price:"+i[2].rjust(10)+"\nIn Stock:"+i[4].rjust(
                        10)+"\nNow:\n   Price:"+dic["price_with_currency"].rjust(10)+"\nIn Stock:"+dic["availability"].rjust(10)+"\nLink:\n"+i[0], parse_mode='html')
                    # await user.send()
            else:
                dic = get_product_info_flipkart(all_items[num-1][0])
                if not dic:
                    bot.send_message(chat_id=chat_id,text=all_items[num-1][0]+" wasn't tracked due to some error. Try again after a while.")
                    return
                current_price = dic['price_in_num']
                current_in_stock = dic['availability']
                bot.send_message(chat_id=chat_id, text=""+dic['name']+":\n"+"Last Time Checked:\n   Price:"+all_items[num-1][2].rjust(10)+"\nIn Stock:"+all_items[num-1][4].rjust(
                    10)+"\nNow:\n   Price:"+dic["price_with_currency"].rjust(10)+"\nIn Stock:"+dic["availability"].rjust(10)+"\nLink:\n"+all_items[num-1][0], parse_mode='html')

def start(bot, update):
    chat_id = update.message.chat_id
    bot.send_message(chat_id=chat_id, text="This is a bot to track price and stock status of a product on flipkart.\nCommands:\n\n  /add --link to product--\n  /remove\n  /check: (To check current price and stock.)\n  /list\nThe bot will check the product price and stock every 6 hours.\nIf some command doesn't work, wait a minute or two and try again. ", parse_mode='html')
    logging.info(str(chat_id)+" started using the bot.")

def main():
    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s",
                        datefmt='%m/%d/%Y %I:%M:%S %p', handlers=[logging.FileHandler("output.log"), logging.StreamHandler()])

    updater = Updater("***TOKEN TO YOUR BOT***")
    dp = updater.dispatcher
    dp.add_handler(CommandHandler('add', add))
    dp.add_handler(CommandHandler('list', list_names))
    dp.add_handler(CommandHandler('remove', remove))
    dp.add_handler(CommandHandler('check', check_now))
    dp.add_handler(CommandHandler('start', start))
    dp.add_handler(CommandHandler('help', start))
    j = updater.job_queue
    j.run_repeating(daily_checker, interval=timedelta(hours=6), first=0)
    updater.start_polling()
    updater.idle()

if __name__ == '__main__':
    main()


from flask_sqlalchemy import SQLAlchemy
from sqlalchemy.ext.declarative import declarative_base

db = SQLAlchemy()
Base = declarative_base()

def init_app(app):
    app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = True
    app.config['SQLALCHEMY_DATABASE_URI'] = app.config.DATABASE_URL
    db.init_app(app)