with open(PATH, 'rb') as f:
    debugOutput = pkl.load(f)
isoColor, isoMag, fFid, fiducialLine = debugOutput
corrLists = list()

nMin = 5
nMax = 25
initShift = np.mean(fiducialLine[1]) - np.mean(isoMag)
isoMagOrig = isoMag.copy()
isoMag = isoMag + initShift

f = interp1d(isoMag, isoColor, bounds_error=False, fill_value='extrapolate')
for n in tqdm(range(nMin, nMax), desc="Guessing mu"):
    correlations = list()
    magShifts = np.linspace(-10,10, n)
    fMagEval = np.linspace(fiducialLine[1].min(), fiducialLine[1].max(), n)
    iMagEval = np.linspace(isoMag.min(), isoMag.max(), n)
    dtwX = np.zeros(shape=(n,1))
    dtwY = np.zeros(shape=(n,1))
    fColor = fFid(fMagEval)
    for magShift in magShifts: 
        iColor = f(iMagEval+magShift)


        dtwX[:, 0] = fColor
        dtwY[:, 0] = iColor
        dtw = fastdtw(dtwX, dtwY, dist=euclidean)

        minFiducialColor, maxFiducialColor = min(fMagEval), max(fMagEval)
        dists = list()
        for i, j in dtw[1]:
            if minFiducialColor <= iMagEval[j] <= maxFiducialColor:
                dists.append(np.sqrt((fColor[i] - iColor[j])**2 + (fMagEval[i] - iMagEval[j])**2))
        chi2 = np.sum([x**2 for x in dists])
        chi2nu = chi2/len(dists)
        correlations.append(chi2nu)
    magShift = magShifts[np.argmin(correlations)]
    corrLists.append(magShift)
magShift = np.mean(corrLists)
print(magShift)

n = 100
with plt.style.context(pubStyle):
    fig, axs = plt.subplots(1,2,figsize=(20,7))
    ax = axs[0]
    fMagEval = np.linspace(fiducialLine[1].min(), fiducialLine[1].max(), n)
    iMagEval = np.linspace(isoMag.min(), isoMag.max(), n)

    fColor = fFid(fMagEval)
    isoMagShifted = isoMag - magShift
    isoColorShifted = isoColor
    isoFuncShifted = interp1d(isoMagShifted, isoColorShifted, bounds_error=False, fill_value='extrapolate')
    iColor = isoFuncShifted(fMagEval)
    
    ax.plot(fColor, fMagEval, 'o-', markersize=3)
    ax.plot(iColor, fMagEval, 'o-', markersize=3)

    dtwX = np.zeros(shape=(n,1))
    dtwY = np.zeros(shape=(n,1))

    dtwX[:, 0] = fColor
    dtwY[:, 0] = iColor
    dtw = fastdtw(dtwX, dtwY, dist=euclidean)
    minFiducialColor, maxFiducialColor = min(fMagEval), max(fMagEval)
    dists = list()
    for i, j in dtw[1]:
        # if minFiducialColor <= iMagEval[j] <= maxFiducialColor:
        if ((j < len(iColor)-1) and j != 0) and ((i < len(fColor)-1) and i != 0):
            dists.append(np.sqrt((fColor[i] - iColor[j])**2 + (fMagEval[i] - fMagEval[j])**2))
            ax.plot([fColor[i], iColor[j]], [fMagEval[i], fMagEval[j]], color='black', alpha=0.1)
    chi2 = np.sum([x**2 for x in dists])
    chi2nu = chi2/len(dists)
    ax.invert_yaxis()
    
    axs[1].plot(fiducialLine[0], fiducialLine[1], label="Fiducial Line")
    axs[1].plot(isoColor, isoMagOrig, label="Isochrone")
    axs[1].legend(frameon=False)
    axs[1].invert_yaxis()
    
    axs[1].set_facecolor("#F0F1EB")
    axs[0].set_facecolor("#F0F1EB")
    fig.set_facecolor("#F0F1EB")
    
    axs[0].set_xlabel("F275W - F814W", fontsize=25)
    axs[1].set_xlabel("F275W - F814W", fontsize=25)
    axs[0].set_ylabel("F814W", fontsize=25)
    axs[1].set_ylabel("F814W", fontsize=25)
    
    fig.savefig("DTWChi2.png", dpi=200)

    # axs[1].set_xlim(ax.get_xlim())
    # axs[1].set_ylim(ax.get_ylim())