import numpy as np
from numpy.polynomial.polynomial import Polynomial
from scipy.signal import find_peaks, argrelextrema
from scipy.optimize import curve_fit

def findMaxIndices(lc, num=3, distance=100, height=(None, None), sortByHighest=False):
    if(hasattr(lc, "power")):
        data = lc.power
    elif(hasattr(lc, "flux")):
        data = lc.flux
    else:
        return np.zeros(num)-1
    peak_indices, peak_dict = find_peaks(data, height=height,
                                         distance=distance)
    peak_heights = peak_dict["peak_heights"]
    if(sortByHighest):
        return peak_indices[np.argsort(peak_heights)[-num:][::-1]]
    else:
        return peak_indices[0:num]

def findNonNanTimeDeltas(lc, column):
    timedeltas = []
    nonNanIndices = np.where(~np.isnan(lc[column]))[0]

    if len(nonNanIndices) == 0:
        return timedeltas

    start = nonNanIndices[0]
    end = start

    for i in range(1, len(nonNanIndices)):
        if nonNanIndices[i] == nonNanIndices[i-1] + 1:
            end = nonNanIndices[i]
        else:
            timedeltas.append((lc.time[start], lc.time[end]))
            start = nonNanIndices[i]
            end = start

    timedeltas.append((lc.time[start], lc.time[end]))

    return [timedelta[1] - timedelta[0] for timedelta in timedeltas]

def getTimeDeltaInSeconds(timedelta):
    return int(timedelta.to_value("sec"))

def getTotalValidDataInSeconds(lc, column):
    tds = findNonNanTimeDeltas(lc, column)
    validDataTime = 0
    for td in tds:
        validDataTime += getTimeDeltaInSeconds(td)
    return validDataTime, tds

def plotFlarePeaks(ax, peaks, flux = None):
    if(flux is None):
        for peak in peaks:
            ax.plot(peak["FlarePeakTime"], peak["FlarePeak"], "x", color="red")
    else:
        for peak in peaks:
            ax.plot(peak["FlarePeakTime"], flux[peak["StandardIndex"]], "x", color="red")

def plotFlareFits(ax, fits):
    for fit in fits:
        ax.plot(fit["FlareFitTime"], fit["FlareFit"], "g--")

def markFlare(ax, fits, flux):
    for fit in fits:
        ax.plot(fit["FlareFitTime"], flux[fit["StandardIndex"]], color="red")

def convertStarndardIndexToFoldedIndex(foldedLc, standardIndex):
    foldedIndex = standardIndex; cycle = 0
    for i in range(max(foldedLc.cycle)+1):
        if(foldedIndex - len(foldedLc.phase[foldedLc.cycle == i]) >= 0):
            foldedIndex = foldedIndex - len(foldedLc.phase[foldedLc.cycle == i])
        else:
            cycle = i
            break

    return cycle, foldedIndex

def singleSine(t, A, w, p, c):
    return A * np.sin(w*t + p) + c

def fitSingleSine(phase, flux):
    initGuess = [np.ptp(flux)/2, 2 * np.pi / (np.max(phase) - np.min(phase)), 0, np.mean(flux)]
    popt, _ = curve_fit(singleSine, phase, flux, p0=initGuess, maxfev=300)
    return singleSine(phase, *popt)

def polynomial(phase, flux, degree):
    return Polynomial.fit(phase, flux, degree).convert().coef

def fitPolynomial(phase, flux, degree):
    return sum(p * phase**i for i, p in enumerate(polynomial(phase, flux, degree)))

def compureRSS(fit, flux):
    residuals = flux - fit
    return np.sum(residuals**2)

def computeTotalSoS(flux):
    return np.sum((flux - np.mean(flux))**2)

def computeAIC(rss, numParams, numDataPoints):
    return 2 * numParams + numDataPoints * np.log(rss/numDataPoints)

def computeBIC(rss, numParams, numDataPoints):
    return numParams * np.log(numDataPoints) + numDataPoints * np.log(rss/numDataPoints)

def getFoldedBestFit(foldedLc):
    flux = foldedLc.flux
    filt = ~np.isnan(flux)
    multi = 1
    if(np.mean(flux).value > 10):
        multi = float(np.mean(flux).value)
        flux = foldedLc.normalize().flux
    phase = foldedLc.phase[filt].value
    flux = flux[filt]
    polyDegree = 7
    fitThreshold = 0.0
    smoothed_flux = np.convolve(flux, np.ones(len(flux)//100)/(len(flux)//100), mode="valid")
    peaks, _ = find_peaks(smoothed_flux, height=np.mean(smoothed_flux))
    #print("Num peaks: ", peaks)

    try:
        retFit = fitSingleSine(phase, flux)
        r2 = 1 - (compureRSS(retFit, flux)/computeTotalSoS(flux))
        fitType = "sine"
        #print("R2: ", r2)
        if(r2 < fitThreshold):
            raise Exception("Sine fit is suboptimal")
        #print("Single sine preferred")
    except:
        retFit = fitPolynomial(phase, flux, polyDegree)
        fitType = "poly"
        #print("Polynomial preferred")

    retFit *= multi
    return phase, retFit, fitType

def getFoldedFitPeakValley(sineFit):
    minima = argrelextrema(sineFit, np.less)[0]
    maxima = argrelextrema(sineFit, np.greater)[0]
    #print(minima)
    #print(maxima)
    if(len(minima) > 0 and len(maxima) > 0):
        if(minima[0] < maxima[0] and minima[-1] > maxima[-1]):
            if(sineFit[minima[0]] > sineFit[minima[-1]]):
                minima = np.delete(minima, 0)
            else:
                minima = np.delete(minima, -1)
        elif(maxima[0] < minima[0] and maxima[-1] > minima[-1]):
            if(sineFit[maxima[0]] > sineFit[maxima[-1]]):
                maxima = np.delete(maxima, -1)
            else:
                maxima = np.delete(maxima, 0)

    return minima, maxima

def findNearestIndexOfValue(array, value):
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()
    return idx

def getPhaseRangesNearPeak(maxArgs, phase):
    minPhases = phase[maxArgs[0]]
    maxPhases = phase[maxArgs[1]]
    totalPhase = abs(phase[0]) + abs(phase[-1])

    phasePart = totalPhase * 0.3 / (len(minPhases) + len(maxPhases))
    minPhasesBounds = []
    for minPhase in minPhases:
        minPhaseBounds = []
        if(minPhase - phasePart < phase[0]):
            minPhaseBounds.append((phase[0], minPhase + phasePart))
            minPhaseBounds.append((phase[-1] + (minPhase - phasePart - phase[0]), phase[-1]))
        elif(minPhase + phasePart > phase[-1]):
            minPhaseBounds.append((minPhase - phasePart, phase[-1]))
            minPhaseBounds.append((phase[0], phase[0] + (minPhase + phasePart - phase[-1])))
        else:
            minPhaseBounds.append((minPhase - phasePart, minPhase + phasePart))
        minPhasesBounds.append(minPhaseBounds)

    maxPhasesBounds = []
    for maxPhase in maxPhases:
        maxPhaseBounds = []
        if(maxPhase - phasePart < phase[0]):
            maxPhaseBounds.append((phase[0], maxPhase + phasePart))
            maxPhaseBounds.append((phase[-1] + (maxPhase - phasePart - phase[0]), phase[-1]))
        elif(maxPhase + phasePart > phase[-1]):
            maxPhaseBounds.append((maxPhase - phasePart, phase[-1]))
            maxPhaseBounds.append((phase[0], phase[0] + (maxPhase + phasePart - phase[-1])))
        else:
            maxPhaseBounds.append((maxPhase - phasePart, maxPhase + phasePart))
        maxPhasesBounds.append(maxPhaseBounds)

    minPhasesBoundsIndices = []
    for minPhaseBounds in minPhasesBounds:
        minPhaseBoundsIndices = []
        for l in minPhaseBounds:
            minPhaseBoundsIndices.append([findNearestIndexOfValue(phase, lv) for lv in l])
        minPhasesBoundsIndices.append(minPhaseBoundsIndices)

    maxPhasesBoundsIndices = []
    for maxPhaseBounds in maxPhasesBounds:
        maxPhaseBoundsIndices = []
        for l in maxPhaseBounds:
            maxPhaseBoundsIndices.append([findNearestIndexOfValue(phase, lv) for lv in l])
        maxPhasesBoundsIndices.append(maxPhaseBoundsIndices)

    return minPhasesBoundsIndices, maxPhasesBoundsIndices

def plotPhaseRangesNearPeak(indices, phase, ax=None):
    minPhasesBoundsIndices = indices[0]
    maxPhasesBoundsIndices = indices[1]

    if(ax is None):
        import matplotlib.pyplot as plt
        for minPhaseBoundsIndices in minPhasesBoundsIndices:
            for pair in minPhaseBoundsIndices:
                for l in pair:
                    plt.axes.axvline(phase[l], color="violet")
        for maxPhaseBoundsIndices in maxPhasesBoundsIndices:
            for pair in maxPhaseBoundsIndices:
                for l in pair:
                    plt.axes.axvline(phase[l], color="orange")
    else:
        for minPhaseBoundsIndices in minPhasesBoundsIndices:
            for pair in minPhaseBoundsIndices:
                for l in pair:
                    ax.axvline(phase[l], color="violet")
        for maxPhaseBoundsIndices in maxPhasesBoundsIndices:
            for pair in maxPhaseBoundsIndices:
                for l in pair:
                    ax.axvline(phase[l], color="orange")

def getEpochTime(lc):
    return lc.time.value[argrelextrema(np.asarray(lc.flux.value), np.less, order=500)[0]][0]