import numpy as np

kmin = 1e-1
kmax = 2e0
Nk = 50
AliasingCorr=False

def get_power_spectrum(field, kmin=kmin, kmax=kmax, Nk=Nk):
    from pysbmy.power import PowerSpectrum
    from pysbmy.fft import FourierGrid
    from pysbmy.correlations import get_autocorrelation


    G = FourierGrid(
        field.L0,
        field.L1,
        field.L2,
        field.N0,
        field.N1,
        field.N2,
        k_modes=np.concat([PowerSpectrum(field.L0,field.L1,field.L2,field.N0,field.N1,field.N2,).FourierGrid.k_modes[:10],np.logspace(
            np.log10(kmin),
            np.log10(kmax),
            Nk,
        )]),
        kmax=kmax,
    )
    k = G.k_modes[1:]
    Pk, _ = get_autocorrelation(field, G, AliasingCorr)
    Pk = Pk[1:]

    return k, Pk


def get_cross_correlations(field_A, field_B, kmin=kmin, kmax=kmax, Nk=Nk):
    from pysbmy.power import PowerSpectrum
    from pysbmy.fft import FourierGrid
    from pysbmy.correlations import get_crosscorrelation


    G = FourierGrid(
        field_A.L0,
        field_A.L1,
        field_A.L2,
        field_A.N0,
        field_A.N1,
        field_A.N2,
        k_modes=np.concat([PowerSpectrum(field_A.L0,field_A.L1,field_A.L2,field_A.N0,field_A.N1,field_A.N2,).FourierGrid.k_modes[:10],np.logspace(
            np.log10(kmin),
            np.log10(kmax),
            Nk,
        )]),
        kmax=kmax,
    )
    k = G.k_modes[1:]
    _, _, Rks, _ = get_crosscorrelation(field_A, field_B, G, AliasingCorr)
    Rks = Rks[1:] 

    return k, Rks


def add_power_spectrum_to_plot(ax, field, Pk_ref=None, plot_args={}, power_args={}):
    k, Pk = get_power_spectrum(field, **power_args)
    if Pk_ref is not None:
        ax.plot(k, Pk/Pk_ref-1, **plot_args)
    else:
        ax.plot(k, Pk, **plot_args)
    return ax

def add_cross_correlations_to_plot(ax, field_A, field_B, plot_args={}, power_args={}):
    k, Rks = get_cross_correlations(field_A, field_B, **power_args)
    ax.plot(k, Rks, **plot_args)
    return ax


def plot_power_spectra(filenames,
                       labels=None,
                       colors=None,
                       linestyles=None,
                       markers=None,
                       Pk_ref=None,
                       ylims=[0.9,1.1],
                       yticks = np.linspace(0.9,1.1,11),
                       bound1=0.01,
                       bound2=0.02,
                       kmin=kmin,
                       kmax=kmax,
                       Nk=Nk,
                       figsize=(8,4),
                       dpi=300,
                       ax=None,
                       fig=None,):
    
    import matplotlib.pyplot as plt
    from pysbmy.field import read_field 

    if fig is None or ax is None:
        fig, ax = plt.subplots(figsize=figsize, dpi=dpi)

    if labels is None:
        labels = [None for f in filenames]
    if colors is None:
        colors = [None for f in filenames]
    if linestyles is None:
        linestyles = [None for f in filenames]
    if markers is None:
        markers = [None for f in filenames]
    
    for i, filename in enumerate(filenames):
        field = read_field(filename)
        add_power_spectrum_to_plot(ax=ax,
                                   field=field,
                                   Pk_ref=Pk_ref,
                                   plot_args=dict(label=labels[i],
                                                  color=colors[i],
                                                  linestyle=linestyles[i],
                                                  marker=markers[i],),
                                   power_args=dict(kmin=kmin,
                                                   kmax=kmax,
                                                   Nk=Nk),
                                    )
    ax.set_xscale('log')
    ax.set_xlim(kmin, kmax)
    if ylims is not None:
        ax.set_ylim(ylims)
    if yticks is not None:
        ax.set_yticks(yticks)
    ax.set_xlabel(r'$k$ [$h/\mathrm{Mpc}$]')

    if Pk_ref is not None:
        ax.set_ylabel(r'$P(k)/P_\mathrm{ref}(k)-1$')
    else:
        ax.set_ylabel('$P(k)$')
    
    if bound1 is not None:
        ax.axhspan(1 - bound1, 1 + bound1, color="grey", alpha=0.2)
    if bound2 is not None:
        ax.axhspan(1 - bound2, 1 + bound2, color="grey", alpha=0.1)

    ax.grid(which="major",alpha=0.5)
    ax.grid(which="minor",alpha=0.2)

    return fig, ax


def plot_cross_correlations(filenames_A,
                            filename_B,
                            labels=None,
                            colors=None,
                            linestyles=None,
                            markers=None,
                            ylims=[0.99, 1.001],
                            yticks = np.linspace(0.99,1.001,12),
                            bound1=0.001,
                            bound2=0.002,
                            kmin=kmin,
                            kmax=kmax,
                            Nk=Nk,
                            figsize=(8,4),
                            dpi=300,
                            ax=None,
                            fig=None,):
    
    import matplotlib.pyplot as plt
    from pysbmy.field import read_field 

    if fig is None or ax is None:
        fig, ax = plt.subplots(figsize=figsize, dpi=dpi)

    if labels is None:
        labels = [None for f in filenames_A]
    if colors is None:
        colors = [None for f in filenames_A]
    if linestyles is None:
        linestyles = [None for f in filenames_A]
    if markers is None:
        markers = [None for f in filenames_A]

    field_B = read_field(filename_B)
    
    for i, filename_A in enumerate(filenames_A):
        field_A = read_field(filename_A)
        add_cross_correlations_to_plot(ax=ax,
                                       field_A=field_A,
                                       field_B=field_B,
                                       plot_args=dict(label=labels[i],
                                                       color=colors[i],
                                                       linestyle=linestyles[i],
                                                       marker=markers[i],),
                                       power_args=dict(kmin=kmin,
                                                        kmax=kmax,
                                                        Nk=Nk),
                                        )   
    ax.set_xscale('log')
    ax.set_xlim(kmin, kmax)
    if ylims is not None:
        ax.set_ylim(ylims)
    if yticks is not None:
        ax.set_yticks(yticks)
    ax.set_xlabel(r'$k$ [$h/\mathrm{Mpc}$]')
    ax.set_ylabel('$R(k)$')
    
    if bound1 is not None:
        ax.axhspan(1 - bound1, 1 + bound1, color="grey", alpha=0.2)
    if bound2 is not None:
        ax.axhspan(1 - bound2, 1 + bound2, color="grey", alpha=0.1)

    ax.grid(which="major",alpha=0.5)
    ax.grid(which="minor",alpha=0.2)

    return fig, ax


if __name__ == "__main__":
    from argparse import ArgumentParser
    parser = ArgumentParser(description='Plot power spectra of fields')

    parser.add_argument('-ps', '--power_spectrum', action='store_true', help='Plot power spectra.')
    parser.add_argument('-cc', '--cross_correlation', action='store_true', help='Plot cross correlations.')
    parser.add_argument('-d', '--directory', type=str, required=True, help='Directory containing the fields files.')
    parser.add_argument('-ref', '--reference', type=str, default=None, help='Reference field file.')
    parser.add_argument('-f', '--filenames', type=str, nargs='+', required=True, help='Field files to be plotted.')
    parser.add_argument('-o', '--output', type=str, default=None, help='Output plot file name.')
    parser.add_argument('-l', '--labels', type=str, nargs='+', default=None, help='Labels for each field.')
    parser.add_argument('-c', '--colors', type=str, nargs='+', default=None, help='Colors for each field.')
    parser.add_argument('-ls', '--linestyles', type=str, nargs='+', default=None, help='Linestyles for each field.')
    parser.add_argument('-m', '--markers', type=str, nargs='+', default=None, help='Markers for each field.')
    parser.add_argument('-t','--title', type=str, default=None, help='Title of the plot.')

    args = parser.parse_args()

    if not args.power_spectrum and not args.cross_correlation:
        print('You must choose between power_spectrum and cross_correlation.')
        exit()

    if args.reference is not None:
        from pysbmy.field import read_field
        Pk_ref = get_power_spectrum(read_field(args.directory+args.reference), kmin=kmin, kmax=kmax, Nk=Nk)[1]
    else:
        Pk_ref = None
    
    filenames = [args.directory+f for f in args.filenames]

    if args.power_spectrum and args.cross_correlation:
        import matplotlib.pyplot as plt
        fig, axes = plt.subplots(2, 1, figsize=(8,8))
        plot_power_spectra(filenames=filenames,
                           labels=args.labels,
                           colors=args.colors,
                           linestyles=args.linestyles,
                           markers=args.markers,
                           Pk_ref=Pk_ref,
                        #    ylims=[0.9,1.1],
                        #    yticks = np.linspace(0.9,1.1,11),
                        #    bound1=0.01,
                        #    bound2=0.02,
                           ylims=None,
                           yticks = None,
                           bound1=None,
                           bound2=None,
                           kmin=kmin,
                           kmax=kmax,
                           Nk=Nk,
                           ax=axes[0],
                           fig=fig)
        
        plot_cross_correlations(filenames_A=filenames,
                                filename_B=args.directory+args.reference,
                                labels=args.labels,
                                colors=args.colors,
                                linestyles=args.linestyles,
                                markers=args.markers,
                                # ylims=[0.99, 1.001],
                                # yticks = np.linspace(0.99,1.001,12),
                                # bound1=0.001,
                                # bound2=0.002,
                                ylims=None,
                                yticks = None,
                                bound1=None,
                                bound2=None,
                                kmin=kmin,
                                kmax=kmax,
                                Nk=Nk,
                                ax=axes[1],
                                fig=fig)
        
        axes[1].legend(loc='lower left')
        axes[0].set_title("Power Spectrum")
        axes[1].set_title("Cross Correlations")
        
        if args.title is not None:
            fig.suptitle(args.title)
    
    elif args.power_spectrum:
        fig, ax = plot_power_spectra(filenames=filenames,
                                     labels=args.labels,
                                     colors=args.colors,
                                     linestyles=args.linestyles,
                                     markers=args.markers,
                                     Pk_ref=Pk_ref,
                                     ylims=[0.9,1.1],
                                     yticks = np.linspace(0.9,1.1,11),
                                     bound1=0.01,
                                     bound2=0.02,
                                     kmin=kmin,
                                     kmax=kmax,
                                     Nk=Nk)
        ax.legend()
        if args.title is not None:
            ax.set_title(args.title)
    
    elif args.cross_correlation:
        fig, ax = plot_cross_correlations(filenames_A=filenames,
                                          filename_B=args.reference,
                                          labels=args.labels,
                                          colors=args.colors,
                                          linestyles=args.linestyles,
                                          markers=args.markers,
                                          ylims=[0.99, 1.001],
                                          yticks = np.linspace(0.99,1.001,12),
                                          bound1=0.001,
                                          bound2=0.002,
                                          kmin=kmin,
                                          kmax=kmax,
                                          Nk=Nk)
        ax.legend(loc='lower left')
        if args.title is not None:
            ax.set_title(args.title)
    
    if args.output is not None:
        fig.savefig(args.output)
    else:
        fig.savefig(args.directory+'power_spectrum.png')