from pywavan import powspec, nb_scale
from astropy.io import fits
import numpy as np
from matplotlib import pyplot as plt
import aplpy


def make_images(im, wt, M, meanim, label):
    """ show the Gaussian and coherent part of the image """

    total = np.sum(wt[:M, :, :], axis=0).real + meanim
    coherent = np.sum(wt[M : 2 * M, :, :], axis=0).real + meanim
    Gaussian = np.sum(wt[2 * M : 3 * M, :, :], axis=0).real + meanim

    fig_all = plt.figure(1, figsize=(16, 4))

    # original image
    fig = aplpy.FITSFigure(fits.PrimaryHDU(im), figure=fig_all, subplot=(1, 4, 1))
    fig.show_colorscale(cmap="cubehelix")
    fig.add_colorbar()
    fig.axis_labels.hide()
    fig.tick_labels.hide()
    fig.set_title("Original")

    # wavelet image total (should be same as original image)
    fig = aplpy.FITSFigure(fits.PrimaryHDU(total), figure=fig_all, subplot=(1, 4, 2))
    fig.show_colorscale(cmap="cubehelix")
    fig.add_colorbar()
    fig.axis_labels.hide()
    fig.tick_labels.hide()
    fig.set_title("wavelet")

    # gaussian component
    fig = aplpy.FITSFigure(fits.PrimaryHDU(Gaussian), figure=fig_all, subplot=(1, 4, 3))
    fig.show_colorscale(cmap="cubehelix")
    fig.add_colorbar()
    fig.axis_labels.hide()
    fig.tick_labels.hide()
    fig.set_title("Gaussian")

    # coherent component
    fig = aplpy.FITSFigure(fits.PrimaryHDU(coherent), figure=fig_all, subplot=(1, 4, 4))
    fig.show_colorscale(cmap="cubehelix")
    fig.add_colorbar()
    fig.axis_labels.hide()
    fig.tick_labels.hide()
    fig.set_title("Coherent")

    plt.tight_layout()
    plt.savefig("reconstructed_image_{}.png".format(label))
    plt.close()


def scale_images(thingy, M, label, scale=14, mode="wt"):
    """ visualize wt or S11a for a specific scale. Remark S11a = wt^2"""
    total = thingy[scale, :, :].real
    coherent = thingy[M + scale, :, :].real
    Gaussian = thingy[2 * M + scale, :, :].real

    # make images
    fig_all = plt.figure(1, figsize=(12, 4))

    # wavelet image on scale
    fig = aplpy.FITSFigure(fits.PrimaryHDU(total), figure=fig_all, subplot=(1, 3, 1))
    limit = max(np.max(total), abs(np.min(total)))
    fig.show_colorscale(cmap="PiYG", vmin=-limit, vmax=limit)
    fig.add_colorbar()
    fig.axis_labels.hide()
    fig.tick_labels.hide()
    fig.set_title("wavelet")

    # gaussian component
    fig = aplpy.FITSFigure(fits.PrimaryHDU(Gaussian), figure=fig_all, subplot=(1, 3, 2))
    limit = max(np.max(Gaussian), abs(np.min(Gaussian)))
    fig.show_colorscale(cmap="PiYG", vmin=-limit, vmax=limit)
    fig.add_colorbar()
    fig.axis_labels.hide()
    fig.tick_labels.hide()
    fig.set_title("Gaussian")

    # coherent component
    fig = aplpy.FITSFigure(fits.PrimaryHDU(coherent), figure=fig_all, subplot=(1, 3, 3))
    limit = max(np.max(coherent), abs(np.min(coherent)))
    fig.show_colorscale(cmap="PiYG", vmin=-limit, vmax=limit)
    fig.add_colorbar()
    fig.axis_labels.hide()
    fig.tick_labels.hide()
    fig.set_title("Coherent")

    plt.tight_layout()
    plt.savefig("imag_{}_scale{}_{}.png".format(mode, scale, label))
    plt.close()


def plot_each_scale(S11a, wav_k, q, label, coherent=False, reso=1):
    """ plot histogram at a certain scale """
    nsize = len(S11a[0, 0, :])
    M = len(q)
    for scl in range(0, M):
        plt.figure(figsize=(6, 6))
        # determine bins (large scales should have less bins)
        nbins = np.int(nsize ** 2.0 * (wav_k[scl] * reso) ** 2.0)
        nbins = max(9, nbins)
        nbins = min(500, nbins)
        # calc histogram gaussian component w.r.t. its mean value (easier to compare)
        intermit = (S11a[2 * M + scl, :, :]) / np.mean(S11a[2 * M + scl, :, :])
        histo, edges = np.histogram(intermit, density=True, bins=nbins)
        plt.hist(
            (edges[:-1] + edges[1:]) / 2.0,
            bins=edges,
            weights=histo,
            alpha=0.5,
            label="Gaussian",
            color="red",
        )
        # calc histogram coherent component
        if coherent and (np.mean(S11a[M + scl, :, :]) > 0):
            intermit_C = (S11a[M + scl, :, :]) / np.mean(S11a[M + scl, :, :])
            histo_C, edges_C = np.histogram(intermit_C, density=True, bins=nbins)
            plt.hist(
                (edges_C[:-1] + edges_C[1:]) / 2.0,
                bins=edges_C,
                weights=histo_C,
                alpha=0.5,
                label="coherent",
                color="blue",
            )
        else:
            histo_C = []
        # plot mean
        plt.plot([1, 1], [0, 5])
        plt.xlabel(r"$I/\langle I \rangle$")
        plt.ylabel("PDF")
        plt.title(
            "scale {}, k= {}, l={} pc".format(
                scl, np.str(wav_k[scl])[:7], 1000.0 / nsize / wav_k[scl]
            )
        )
        plt.legend()
        # avoid first bin dominating the range
        gauss_max = 0
        if len(histo) > 1:
            gauss_max = max(histo[1:])
        coh_max = 0
        if coherent & len(histo_C) > 1:
            coh_max = max(histo_C[1:])
        if max(gauss_max, coh_max) > 0:
            plt.ylim(0, max(gauss_max, coh_max))
        plt.xlim(0, 4)
        plt.ylim(0, 1.0)
        plt.tight_layout()
        plt.savefig("S11a_scale{}_q_{}_{}.png".format(scl, q[scl], label))
        plt.close()


def plot_components_power_spectrum(
    tab_k, wav_k, spec_k, S1a, label="", fit_min=7, fit_max=15, nsize=1024, lvlmin=8
):

    # plot power spectra
    plt.figure(figsize=(5, 5))
    plt.plot(tab_k, spec_k, color="black", label="Fourier PS", linewidth=1.5)
    plt.plot(wav_k, S1a[0, :], "s", color="black", label="Wavelet PS", markersize=6)
    plt.plot(wav_k, S1a[1, :], "D", color="blue", label="Coherent PS", markersize=4)
    plt.plot(wav_k, S1a[2, :], "^", color="red", label="Gaussian PS", markersize=5)

    scl_min = fit_min
    scl_max = fit_max + 1

    # power law fit to gaussian component
    coefw, cov = np.polyfit(
        np.log(wav_k[scl_min:scl_max]), np.log(S1a[2, scl_min:scl_max]), deg=1, cov=True
    )
    yfitw = np.exp(coefw[1]) * wav_k[scl_min:scl_max] ** coefw[0]
    plt.plot(wav_k[scl_min:scl_max], yfitw, "-", color=("red"), linewidth=4, alpha=0.5)
    print("Gaussian Power law", coefw[0])

    # Coherent power law fit
    coefcw, cov = np.polyfit(
        np.log(wav_k[scl_min:scl_max]), np.log(S1a[1, scl_min:scl_max]), deg=1, cov=True
    )
    yfitcw = np.exp(coefcw[1]) * wav_k[scl_min:scl_max] ** coefcw[0]
    plt.plot(wav_k[scl_min:scl_max], yfitcw, "-", color="blue", linewidth=4, alpha=0.5)
    print("Coherent Power law", coefcw[0])

    # Fourier power law fit
    limit = np.where((tab_k >= wav_k[fit_min]) & (tab_k <= wav_k[fit_max]))
    coef, cov = np.polyfit(np.log(tab_k[limit]), np.log(spec_k[limit]), deg=1, cov=True)
    yfit = np.exp(coef[1]) * tab_k[limit] ** coef[0]
    plt.plot(tab_k[limit], yfit, "-", color="black", linewidth=2, alpha=0.5)
    print("Fourier Power law", coef[0])

    # show resolution limits
    # Gaussian part not accurate below levelmin due to the way AMR works
    sim_res_eff = 2 ** lvlmin / 10
    sim_res_lvl_min = 2 ** lvlmin
    plt.plot(
        [sim_res_lvl_min / nsize, sim_res_lvl_min / nsize],
        [1e-6, 1e8],
        color="green",
        ls="--",
        label="dx(levelmin)",
    )
    plt.plot(
        [sim_res_eff / nsize, sim_res_eff / nsize],
        [1e-6, 1e8],
        color="green",
        ls="-",
        label="dx(levelmin) x 10",
    )

    plt.xscale("log")
    plt.yscale("log")
    plt.xlabel(r"$k$ (pixel$^{-1}$)")
    plt.ylabel(r"$P(k)$")
    plt.ylim(1e-6, 1e8)
    plt.legend(loc="lower left")
    plt.savefig("powerspectrum_{}.png".format(label), bbox_inches="tight")
    plt.close()


def save_results(wt, S11a, wav_k, S1a, q, label):
    np.save("wav_k_{}.npy".format(label), wav_k)
    np.save("S1a_{}.npy".format(label), S1a)
    np.save("wt_{}.npy".format(label), wt)
    np.save("S11a_{}.npy".format(label), S11a)
    np.save("q_{}.npy".format(label), q)


def load_results(label):
    wav_k = np.load("wav_k_{}.npy".format(label))
    S1a = np.load("S1a_{}.npy".format(label))
    wt = np.load("wt_{}.npy".format(label))
    S11a = np.load("S11a_{}.npy".format(label))
    q = np.load("q_{}.npy".format(label))
    return wav_k, S1a, wt, S11a, q


def analyse_sim(im):
    """ Do the MnGseg analysis """
    meanim = np.mean(im)
    imzm = im - meanim
    M = nb_scale(im.shape)

    fit_min = 5  # minimal scale for fitting the power spectrum slope
    fit_max = 11  # max scale
    label = "final"  # label to identify parameter setup

    # after a lot of trials I found a fixed q=2 is a good value
    q = [2.0] * nb_scale(imzm.shape)
    # wt, S11a, wav_k, S1a, q = fan_trans(imzm, reso=1, q=q, qdyn=False)
    # alternatively you can let pywavan determine it automatically by setting skewl
    # q=[3.0]*nb_scale(imzm.shape)
    # wt, S11a, wav_k, S1a, q = fan_trans(imzm, reso=1, q=q, qdyn=True, skewl=0.4)
    # print(q)

    # save results because it can be long to calculate (especially if qdyn=True). Remark that wt and S11a are quite big
    # save_results(wt, S11a, wav_k, S1a, q, label)
    wav_k, S1a, wt, S11a, q = load_results(label)

    # make images of the Gaussian and coherent part
    make_images(im, wt, M, meanim, label)

    # (optional) make the image for each scale
    for s in range(fit_min, fit_max + 1):
        scale_images(wt, M, label, scale=s)

    # (optional) plot the histogram for each scale
    plot_each_scale(S11a, wav_k, q, label, coherent=True)

    # calc Fourier power spectrum for comparison
    tab_k, spec_k = powspec(imzm, reso=1)

    # plot the powerspectrum for each component and fit the slope
    plot_components_power_spectrum(tab_k, wav_k, spec_k, S1a, label, fit_min, fit_max)