import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
try:
    from .. import hitdetector as hd
    from ..misc import sort_data
except(ImportError):
    import hits.hitdetector as hd
    from misc import sort_data
from scipy.interpolate import UnivariateSpline, BSpline


@sort_data
def isolate_anomaly(df, method='magnitude', time_res=0.01):
    """
    Accepts:

        a Pandas dataframe of shape:

                obmt    rate    w1_rate
            1.  float   float   float

        or equivalent.

    Isolates detected hits.

    Kwargs:

        time_res (float, default=0.01):
            half the width of the neighbourhoods to be generated.

    Returns:

        a tuple of reduced Pandas dataframes of shape:

                obmt    rate    w1_rate anomaly hits
            1.  float   float   float   bool    bool

        or equivalent.
    """
    working_df = df.copy()

    identify = hd.method_dict[method]

    working_df = identify(working_df)[0]

    hit_df = working_df[working_df['anomaly']]  # Isolate the hits.

    # Generate neighbourhoods around the hits with width 2*time_res.
    hit_neighbourhoods = [working_df[abs(working_df['obmt'] - time) <
                          time_res] for time in hit_df['obmt']]

    return tuple(hit_neighbourhoods)


@sort_data
def spline_anomaly(df, smooth=0.5, plot=False,
                   B=False, turning=False, filtered=False, threshold=1):
    """
    Accepts:

        a Pandas dataframe of shape:

                obmt    rate    w1_rate
            1.  float   float   float

        or equivalent.

    This dataframe should be a hit neighbourhood as generated by
    isolate_anomaly(). The function will work on larger datasets but
    there is nothing to be gained by splining these - and the time taken
    to do so would be significantly larger.

    Runs scipy's splining algorithms on the hit neighbourhoods to
    generate a spline to fit the data.

    Kwargs:

        smooth (float, default=0.5):
            smoothing factor for the generated splines.

        plot (bool, default=False):
            if True, generates a plot of the data and the spline.

        B (bool, default=False):
            if True, generates B splines.

        turning (bool, default=False):
            if True, plots all the turning points alongside a normal
            plot.

        filtered (bool, default=False):
            if True, plots the filterd turning points alongside a normal
            plot.

        threshold (float, default=1.0):
            the threshold for filtering turning points

    Returns:

        tuple of the arrays of knots and coefficients (in that order) of
        the fitted spline.
    """

    # Create spline.
    spl = UnivariateSpline(df['obmt'], df['rate'] - df['w1_rate'])
    spl.set_smoothing_factor(smooth)
    knots, coeffs = (spl.get_knots(), spl.get_coeffs())

    xs = np.linspace(df['obmt'].tolist()[0], df['obmt'].tolist()[-1], 10000)

    if B:
        spl = BSpline(knots, coeffs, 3)

    # Plot original data and spline.
    if plot or turning or filtered:
        plt.scatter(df['obmt'], df['rate'] - df['w1_rate'])
        plt.plot(xs, spl(xs))

        if filtered or turning:
            # Calls get_turning_points() to isolate the turning points.
            if turning:
                turning_points = get_turning_points(df)

            elif filtered:
                turning_points = filter_turning_points(get_turning_points(df),
                                                       threshold=threshold)

            plt.scatter(turning_points['obmt'],
                        turning_points['rate'] - turning_points['w1_rate'],
                        color='red')

        plt.show()

    return (knots, coeffs)
