geinf/swp2.py

import matplotlib.pyplot as plt
import os
import numpy as np
import math
import json

def log_facto(k):
    """
    Using the Stirling's approximation
    """
    k = int(k)
    if k > 1e6:
        return k * np.log(k) - k + np.log(2*math.pi*k)/2
    val = 0
    for i in range(2, k+1):
        val += np.log(i)
    return val

def parse_stwp_theta_file(stwp_theta_file, breaks, mu, tgen, relative_theta_scale = False):
    with open(stwp_theta_file, "r") as swp_file:
        # Read the first line
        line = swp_file.readline()
        L = float(line.split()[2])
        rands = swp_file.readline()
        line = swp_file.readline()
        # skip empty lines before SFS
        while line == "\n":
            line = swp_file.readline()
        sfs = np.array(line.split()).astype(float)
        # Process lines until the end of the file
        while line:
            # check at each line
            if line.startswith("dim") :
                dim = int(line.split()[1])
                if dim == breaks+1:
                    likelihood = line.split()[5]
                    groups = line.split()[6:6+dim]
                    theta_site = line.split()[6+dim:6+dim+1+dim]
                elif dim < breaks+1:
                    line = swp_file.readline()
                    continue
                elif dim > breaks+1:
                    break
                    #return 0,0,0
            # Read the next line
            line = swp_file.readline()
    #### END of parsing
    # quit this file if the number of dimensions is incorrect
    if dim < breaks+1:
        return 0,0,0,0,0,0
    # get n, the last bin of the last group
    # revert the list of groups as the most recent times correspond
    # to the closest and last leafs of the coal. tree.
    groups = groups[::-1]
    theta_site = theta_site[::-1]
    # store thetas for later use
    grps = groups.copy()
    thetas = {}
    for i in range(len(groups)):
        grps[i] = grps[i].split(',')
        thetas[i] = [float(theta_site[i]), grps[i], likelihood]
    # initiate the dict of times
    t = {}
    # list of thetas
    theta_L = []
    sum_t = 0
    for group_nb, group in enumerate(groups):
        ###print(group_nb, group, theta_site[group_nb], len(theta_site))
        # store all the thetas one by one, with one theta per group
        theta_L.append(float(theta_site[group_nb]))
        # if the group is of size 1
        if len(group.split(',')) == 1:
            i = int(group)
        # if the group size is >1, take the first elem of the group
        # i is the first bin of each group, straight after a breakpoint
        else:
            i = int(group.split(",")[0])
            j = int(group.split(",")[-1])
        t[i] = 0
        #t =
        if len(group.split(',')) == 1:
            k = i
            if relative_theta_scale:
                t[i] += ((theta_L[group_nb] ) / (k*(k-1)))
            else:
                t[i] += ((theta_L[group_nb] ) / (k*(k-1)) * tgen) / mu
        else:
            for k in range(j, i-1, -1 ):
                if relative_theta_scale:
                    t[i] += ((theta_L[group_nb] ) / (k*(k-1)))
                else:
                    t[i] += ((theta_L[group_nb] ) / (k*(k-1)) * tgen) / mu
        # we add the cumulative times at the end
        t[i] += sum_t
        sum_t = t[i]
    # build the y axis (sizes)
    y = []
    for theta in theta_L:
        if relative_theta_scale:
            size = theta
        else:
            # with size N = theta/4mu
            size = theta / (4*mu)
        y.append(size)
        y.append(size)
    # build the time x axis
    x = [0]
    for time in range(0, len(t.values())-1):
        x.append(list(t.values())[time])
        x.append(list(t.values())[time])
    x.append(list(t.values())[len(t.values())-1])

    return x,y,likelihood,thetas,sfs,L

def plot_straight_x_y(x,y):
    x_1 = [x[0]]
    y_1 = []
    for i in range(0, len(y)-1):
        x_1.append(x[i])
        x_1.append(x[i])
        y_1.append(y[i])
        y_1.append(y[i])
    y_1 = y_1+[y[-1],y[-1]]
    x_1.append(x[-1])
    return x_1, y_1

def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title",
    theta_scale = True, ax = None, input = None, output = None):
    my_dpi = 500
    L = full_dict["L"]
    if ax is None:
        # intialize figure
        #my_dpi = 300
        fnt_size = 18
        # plt.rcParams['font.size'] = fnt_size
        fig, ax1 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
    else:
        fnt_size = 12
        # plt.rcParams['font.size'] = fnt_size
        ax1 = ax[1][0,0]
    ax1.set_yscale('log')
    ax1.set_xscale('log')
    plot_handles = []
    best_plot = full_dict['all_epochs']['best']
    p0, = ax1.plot(best_plot[0], best_plot[1], linestyle = "-",
    alpha=1, lw=2, label = str(best_plot[2])+' brks | Lik='+best_plot[3])
    plot_handles.append(p0)
    #ax1.grid(True,which="both", linestyle='--', alpha = 0.3)
    for k, plot_Lk in enumerate(full_dict['all_epochs']['plots']):
        plot_Lk = str(full_dict['all_epochs']['plots'][k][3])
        # plt.rcParams['font.size'] = fnt_size
        p, = ax1.plot(full_dict['all_epochs']['plots'][k][0], full_dict['all_epochs']['plots'][k][1], linestyle = "-",
        alpha=1/(k+1), lw=1.5, label = str(full_dict['all_epochs']['plots'][k][2])+' brks | Lik='+plot_Lk)
        plot_handles.append(p)
    if theta_scale:
        ax1.set_xlabel("Coal. time", fontsize=fnt_size)
        ax1.set_ylabel("Pop. size scaled by N0", fontsize=fnt_size)
        # recent_scale_lower_bound = 0.01
        # recent_scale_upper_bound = 0.1
        # ax1.axvline(x=recent_scale_lower_bound)
        # ax1.axvline(x=recent_scale_upper_bound)
    else:
        # years
        if ax is not None:
            plt.set_xlabel("Time (years)", fontsize=fnt_size)
            plt.set_ylabel("Effective pop. size (Ne)", fontsize=fnt_size)
        else:
            plt.xlabel("Time (years)", fontsize=fnt_size)
            plt.ylabel("Effective pop. size (Ne)", fontsize=fnt_size)
    # x_ticks = ax1.get_xticks()
    # ax1.set_xticklabels([f'{k:.0e}\n{k/(mu):.0e}\n{k/(mu)*tgen:.0e}' for k in x_ticks], fontsize = fnt_size*0.5)
    # ax1.set_xticklabels([f'{k}\n{k/(mu)}\n{k/(mu)*tgen}' for k in x_ticks], fontsize = fnt_size*0.8)
    # plt.rcParams['font.size'] = fnt_size
    # print(fnt_size, "rcParam font.size=", plt.rcParams['font.size'])
    ax1.legend(handles = plot_handles, loc='best', fontsize = fnt_size*0.5)
    ax1.set_title(title)
    breaks = len(full_dict['all_epochs']['plots'])
    if ax is None:
        plt.savefig(title+'_'+str(breaks+1)+'_epochs.pdf')
    # plot likelihood against nb of breakpoints
    if ax is None:
        fig, ax2 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
        # plt.rcParams['font.size'] = fnt_size
    else:
        #plt.rcParams['font.size'] = fnt_size
        ax2 = ax[0][0,1]
    # Retrieve the default color cycle from rcParams
    default_colors = plt.rcParams['axes.prop_cycle'].by_key()['color']
    # Create an array of colors from the default color cycle
    colors = [default_colors[i % len(default_colors)] for i in range(len(full_dict['Ln_Brks'][0]))]
    ax2.plot(full_dict['Ln_Brks'][0], full_dict['Ln_Brks'][1], "--", lw=1, color="black", zorder=1)
    ax2.scatter(full_dict['Ln_Brks'][0], full_dict['Ln_Brks'][1], s=50, c=colors, marker='o', zorder=2)
    ax2.axhline(y=full_dict['best_Ln'], linestyle = "-.", color = "red", label = "$-\log\mathcal{L}$ = "+str(round(full_dict['best_Ln'], 2)))
    ax2.set_yscale('log')
    ax2.set_xlabel("# breakpoints", fontsize=fnt_size)
    ax2.set_ylabel("$-\log\mathcal{L}$", fontsize=fnt_size)
    ax2.legend(loc='best', fontsize = fnt_size*0.5)
    ax2.set_title(title+" Likelihood gain from # breakpoints")
    if ax is None:
        plt.savefig(title+'_Breakpts_Likelihood.pdf')
    # AIC
    if ax is None:
        fig, ax3 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
        # plt.rcParams['font.size'] = '18'
    else:
        #plt.rcParams['font.size'] = fnt_size
        ax3 = ax[1][0,1]
    AIC = full_dict['AIC_Brks']
    # ax3.plot(AIC[0], AIC[1], 'o', linestyle = "dotted", lw=2)
    ax3.plot(AIC[0], AIC[1], "--", lw=1, color="black", zorder=1)
    ax3.scatter(AIC[0], AIC[1], s=50, c=colors, marker='o', zorder=2)
    ax3.axhline(y=full_dict['best_AIC'], linestyle = "-.", color = "red",
    label = "Min. AIC = "+str(round(full_dict['best_AIC'], 2)))
    ax3.set_yscale('log')
    ax3.set_xlabel("# breakpoints", fontsize=fnt_size)
    ax3.set_ylabel("AIC")
    ax3.legend(loc='best', fontsize = fnt_size*0.5)
    ax3.set_title(title+" AIC")
    if ax is None:
        plt.savefig(title+'_Breakpts_Likelihood_AIC.pdf')
    else:
        # return plots
        return ax[0], ax[1]

def save_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_scale = True, input = None, output = None):
    #scenari = {}
    cpt = 0
    epochs = {}
    plots = {}
    # store ['best'], and [0] for epoch 0 etc...
    for file_name in os.listdir(folder_path):
        breaks = 0
        cpt +=1
        if os.path.isfile(os.path.join(folder_path, file_name)):
            x, y, likelihood, theta, sfs, L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks,
                                                             tgen = tgen,
                                                             mu = mu, relative_theta_scale = theta_scale)
            SFS_stored = sfs
            L_stored = L
            while not (x == 0 and y == 0):
                if breaks not in epochs.keys():
                    epochs[breaks] = {}
                epochs[breaks][likelihood] = x,y
                breaks += 1
                x,y,likelihood,theta,sfs,L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks,
                                                                 tgen = tgen,
                                                                  mu = mu, relative_theta_scale = theta_scale)
            if x == 0:
                # last break did not work, then breaks = breaks-1
                breaks -= 1
    print("\n*******\n"+title+"\n--------\n"+"mu="+str(mu)+"\ntgen="+str(tgen)+"\nbreaks="+str(breaks)+"\n*******\n")
    print(cpt, "theta file(s) have been scanned.")
    brkpt_lik = []
    top_plots = {}
    for epoch, scenari in epochs.items():
        # sort starting by the smallest -log(Likelihood)
        best10_scenari = (sorted(list(scenari.keys())))[:10]
        greatest_likelihood = best10_scenari[0]
        # store the tuple breakpoints and likelihood for later plot
        brkpt_lik.append((epoch, greatest_likelihood))
        x, y = scenari[greatest_likelihood]
        #without breakpoint
        if epoch == 0:
            # do something with the theta without bp and skip the plotting
            N0 = y[0]
            #continue
        if theta_scale:
            for i in range(len(y)):
                # divide by N0
                y[i] = y[i]/N0
                x[i] = x[i]/N0
        top_plots[greatest_likelihood] = x,y,epoch
    plots_likelihoods = list(top_plots.keys())
    for i in range(len(plots_likelihoods)):
        plots_likelihoods[i] = float(plots_likelihoods[i])
    best10_plots = sorted(plots_likelihoods)[:10]
    top_plot_lik = str(best10_plots[0])
    # store x,y,brks,likelihood
    plots['best'] = (top_plots[top_plot_lik][0], top_plots[top_plot_lik][1], str(top_plots[top_plot_lik][2]), top_plot_lik)
    plots['plots'] = []
    for k, plot_Lk in enumerate(best10_plots[1:]):
        plot_Lk = str(plot_Lk)
        plots['plots'].append([top_plots[plot_Lk][0], top_plots[plot_Lk][1], str(top_plots[plot_Lk][2]), plot_Lk])
    # plot likelihood against nb of breakpoints
    # best possible likelihood from SFS
    # Segregating sites
    S = sum(SFS_stored)
    # Number of kept sites from which the SFS is computed
    L = L_stored
    # number of monomorphic sites
    S0 = L-S
    # print("SFS", SFS_stored)
    # print("S", S, "L", L, "S0=", S0)
    # compute Ln
    Ln = log_facto(S+S0) - log_facto(S0) + np.log(float(S0)/(S+S0)) * S0
    for xi in range(0, len(SFS_stored)):
        p_i = SFS_stored[xi] / float(S+S0)
        Ln += np.log(p_i) * SFS_stored[xi] - log_facto(SFS_stored[xi])
    # basic plot likelihood
    Ln_Brks = [list(np.array(brkpt_lik)[:, 0]), list(np.array(brkpt_lik)[:, 1].astype(float))]
    best_Ln = -Ln
    AIC = []
    for brk in np.array(brkpt_lik)[:, 0]:
        brk = int(brk)
        AIC.append((2*brk+1)+2*np.array(brkpt_lik)[brk, 1].astype(float))
    AIC_Brks = [list(np.array(brkpt_lik)[:, 0]), AIC]
    # AIC = 2*k - 2ln(L) ; where k is the number of parameters, here brks+1
    AIC_ln = 2*(len(brkpt_lik)+1) - 2*Ln
    best_AIC = AIC_ln
    selected_brks_nb = AIC.index(min(AIC))
    # to return : plots ; Ln_Brks ; AIC_Brks ; best_Ln ; best_AIC
    # 'plots' dict keys: 'best', {epochs}('0', '1',...)
    if input == None:
        saved_plots = {"S":S, "S0":S0, "L":L, "mu":mu, "tgen":tgen,
                       "all_epochs":plots, "Ln_Brks":Ln_Brks,
                       "AIC_Brks":AIC_Brks, "best_Ln":best_Ln,
                       "best_AIC":best_AIC, "best_epoch_by_AIC":selected_brks_nb}
    else:
        # if the dict has to be loaded from input
        with open(input, 'r') as json_file:
            saved_plots = json.load(json_file)
            saved_plots["S"] = S
            saved_plots["S0"] = S0
            saved_plots["L"] = L
            saved_plots["mu"] = mu
            saved_plots["tgen"] = tgen
            saved_plots["all_epochs"] = plots
            saved_plots["Ln_Brks"] = Ln_Brks
            saved_plots["AIC_Brks"] = AIC_Brks
            saved_plots["best_Ln"] = best_Ln
            saved_plots["best_AIC"] = best_AIC
            saved_plots["best_epoch_by_AIC"] = selected_brks_nb
    if output == None:
        output = title+"_plotdata.json"
    with open(output, 'w') as json_file:
        json.dump(saved_plots, json_file)
    return saved_plots

def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
    breaks_max = 10, input = None, output = None):
    """
    Save theta values as is to do basic plots.
    """
    cpt = 0
    epochs = {}
    len_sfs = 0
    for file_name in os.listdir(folder_path):
        cpt +=1
        if os.path.isfile(os.path.join(folder_path, file_name)):
            for k in range(breaks_max+1):
                x,y,likelihood,thetas,sfs,L = parse_stwp_theta_file(folder_path+file_name, breaks = k,
                                                                 tgen = tgen,
                                                                 mu = mu, relative_theta_scale = theta_scale)
                if thetas == 0:
                    continue
                if len(thetas)-1 != k:
                    continue
                if k not in epochs.keys():
                    epochs[k] = {}
                likelihood = str(eval(thetas[k][2]))
                epochs[k][likelihood] = thetas
                #epochs[k] = thetas
    print("\n*******\n"+title+"\n--------\n"+"mu="+str(mu)+"\ntgen="+str(tgen)+"\nbreaks="+str(k)+"\n*******\n")
    print(cpt, "theta file(s) have been scanned.")
    plots = []
    best_epochs = {}
    for epoch in epochs:
        likelihoods = []
        for key in epochs[epoch].keys():
            likelihoods.append(key)
        likelihoods.sort()
        minLogLn = str(likelihoods[0])
        best_epochs[epoch] = epochs[epoch][minLogLn]
    for epoch, theta in best_epochs.items():
        groups = np.array(list(theta.values()), dtype=object)[:, 1].tolist()
        x = []
        y = []
        thetas = np.array(list(theta.values()), dtype=object)[:, 0]
        for i,group in enumerate(groups):
            x += group[::-1]
            y += list(np.repeat(thetas[i], len(group)))
            if epoch == 0:
                N0 = y[0]
                # compute the proportion of information used at each bin of the SFS
                sum_theta_i = 0
                for i in range(2, len(y)+2):
                    sum_theta_i+=y[i-2] / (i-1)
                prop = []
                for k in range(2, len(y)+2):
                    prop.append(y[k-2] / (k - 1) / sum_theta_i)
                prop = prop[::-1]
        if theta_scale :
            # normalise to N0 (N0 of epoch1)
            for i in range(len(y)):
                y[i] = y[i]/N0
        # x_plot, y_plot = plot_straight_x_y(x, y)
        p = x, y
        # add plot to the list of all plots to superimpose
        plots.append(p)
    cumul = 0
    prop_cumul = []
    for val in prop:
        prop_cumul.append(val+cumul)
        cumul = val+cumul
    prop = prop_cumul

    lines_fig2 = []
    for epoch, theta in best_epochs.items():
        groups = np.array(list(theta.values()), dtype=object)[:, 1].tolist()
        x = []
        y = []
        thetas = np.array(list(theta.values()), dtype=object)[:, 0]
        for i,group in enumerate(groups):
            x += group[::-1]
            y += list(np.repeat(thetas[i], len(group)))
            if epoch == 0:
                N0 = y[0]
        if theta_scale :
            for i in range(len(y)):
                y[i] = y[i]/N0
        x_2 = []
        T = 0
        for i in range(len(x)):
            x[i] = int(x[i])
        # compute the times as: theta_k / (k*(k-1))
        for i in range(0, len(x)):
            T += y[i] / (x[i]*(x[i]-1))
            x_2.append(T)
        # Save plotting (fig 2)
        x_2 = [0]+x_2
        y = [y[0]]+y
        # x2_plot, y2_plot = plot_straight_x_y(x_2, y)
        p2 = x_2, y
        lines_fig2.append(p2)
    if input == None:
        saved_plots = {"raw_stairs":plots, "scaled_stairs":lines_fig2,
                        "prop":prop}
    else:
        # if the dict has to be loaded from input
        with open(input, 'r') as json_file:
            saved_plots = json.load(json_file)
        saved_plots["raw_stairs"] = plots
        saved_plots["scaled_stairs"] = lines_fig2
        saved_plots["prop"] = prop
    if output == None:
        output = title+"_plotdata.json"
    with open(output, 'w') as json_file:
        json.dump(saved_plots, json_file)
    return saved_plots

def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax = None, n_ticks = 10, subset = None, theta_scale = False):
    # nb of plot_lines represent the number of epochs stored (len(plot_lines) = #breaks+1)
    nb_epochs = len(plot_lines)
    # fig 2 & 3
    if ax is None:
        my_dpi = 500
        fnt_size = 18
        fig2, ax2 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
        fig3, ax3 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
    else:
        # plt.rcParams['font.size'] = fnt_size
        fnt_size = 12
        # place of plots on the grid
        ax2 = ax[1,0]
        ax3 = ax[1,1]
    lines_fig2 = []
    lines_fig3 = []
    #plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
    if swp2_lines:
        for k in range(len(swp2_lines[0])):
            swp2_lines[0][k] = swp2_lines[0][k]/tgen*mu
        for k in range(len(swp2_lines[1])):
            swp2_lines[1][k] = swp2_lines[1][k]*4*mu
        x2_plot, y2_plot = plot_straight_x_y(swp2_lines[0],swp2_lines[1])
        p2, = ax2.plot(x2_plot, y2_plot, linestyle="-", alpha=0.75, lw=2, label = 'swp2', color="black")
        lines_fig2.append(p2)
        # Plotting (fig 3) which is the same but log scale for x
        p3, = ax3.plot(x2_plot, y2_plot, linestyle="-", alpha=0.75, lw=2, label = 'swp2', color="black")
        lines_fig3.append(p3)
    for breaks, plot in enumerate(plot_lines):
        x,y=plot
        x2_plot, y2_plot = plot_straight_x_y(x,y)
        if subset is not None:
            if breaks in subset:
                masking_alpha = 0.75
            else:
                masking_alpha = 0
        p2, = ax2.plot(x2_plot, y2_plot, 'o', linestyle="-", alpha=masking_alpha, lw=2, label = str(breaks)+' brks')
        # Plotting (fig 3) which is the same but log scale for x
        p3, = ax3.plot(x2_plot, y2_plot, 'o', linestyle="-", alpha=masking_alpha, lw=2, label = str(breaks)+' brks')
        if subset is not None and breaks in subset:
            # store for legend
            lines_fig2.append(p2)
            lines_fig3.append(p3)
    ax3.axvline(x=500/tgen*mu, linestyle="--")
    if theta_scale:
        xlabel = "Theta scaled by N0"
        ylabel = "Theta scaled by N0"
    else:
        xlabel = "t"
        ylabel = r"$\theta$"
    if ax is None:
        # if not ax, then use the plt syntax, not ax...
        plt.xlabel(xlabel, fontsize=fnt_size)
        plt.ylabel(ylabel, fontsize=fnt_size)
        plt.xlim(left=0)
        x_ticks = list(plt.xticks())[0]
        plt.gca().set_xticks(x_ticks)
        plt.gca().set_xticklabels([f'{k:.0e}\n{k/(mu):.0e}\n{k/(mu)*tgen:.0e}' for k in x_ticks], fontsize = fnt_size*0.5)
        plt.title(title, fontsize=fnt_size)
        plt.legend(handles=lines_fig2, loc='best', fontsize = fnt_size*0.5)
        plt.text(-0.13, -0.135, 'Coal. time\nGen. time\nYears', ha='left', va='bottom', transform=ax3.transAxes)
        plt.subplots_adjust(bottom=0.2)  # Adjust the value as needed
        plt.savefig(title+'_plotB_'+str(nb_epochs)+'_epochs.pdf')
        # close fig2 to save memory
        plt.close(fig2)
    else:
        # when ax subplotting is used
        ax2.set_xlabel(xlabel, fontsize=fnt_size)
        ax2.set_ylabel(ylabel, fontsize=fnt_size)
        ax2.set_title(title, fontsize=fnt_size)
        ax2.legend(handles=lines_fig2, loc='best', fontsize = fnt_size*0.5)
    ax3.set_xscale('log')
    ax3.set_yscale('log')
    ax3.set_xlabel("time log scale", fontsize=fnt_size)
    ax3.set_ylabel("theta", fontsize=fnt_size)
    ax3.set_title(title, fontsize=fnt_size)
    ax3.legend(handles=lines_fig3, loc='best', fontsize = fnt_size*0.5)
    x_ticks = list(ax3.get_xticks())
    ax3.set_xlim(left=min(x_ticks))
    ax3.set_xticks(x_ticks)
    ax3.set_xticklabels([f'{k:.0e}\n{k/(mu):.0e}\n{k/(mu)*tgen:.0e}' for k in x_ticks], fontsize = fnt_size*0.5)
    plt.text(-0.13, -0.135, 'Coal. time\nGen. time\nYears', ha='left', va='bottom', transform=ax3.transAxes)
    plt.subplots_adjust(bottom=0.2)  # Adjust the value as needed
    if ax is None:
        # nb of plot_lines represent the number of epochs stored (len(plot_lines) = #breaks+1)
        plt.savefig(title+'_plotC_'+str(nb_epochs)+'_epochs_log.pdf')
        # close fig3 to save memory
        plt.close(fig3)
    return ax

def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10, rescale = False, subset = None, max_breaks = None):
    if max_breaks:
        nb_breaks = max_breaks
    else:
        nb_breaks = len(plot_lines)+1
    # multiple fig
    if ax is None:
        # intialize figure 1
        my_dpi = 500
        fnt_size = 18
        # plt.rcParams['font.size'] = fnt_size
        fig, ax1 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
        plt.subplots_adjust(bottom=0.2)  # Adjust the value as needed
    else:
        fnt_size = 12
        # plt.rcParams['font.size'] = fnt_size
        ax1 = ax[0, 0]
        plt.subplots_adjust(wspace=0.3, hspace=0.3)
    plots = []
    for breaks, plot in enumerate(plot_lines):
        if max_breaks and breaks > max_breaks:
            # stop plotting if it exceeds the limit
            continue
        x,y = plot
        x_plot, y_plot = plot_straight_x_y(x,y)
        p, = ax1.plot(x_plot, y_plot, 'o', linestyle="-", alpha=0.75, lw=2, label = str(breaks)+' brks')

        # add plot to the list of all plots to superimpose
        plots.append(p)
    x_ticks = x
    # print(x_ticks)
    #print(prop, "\n", sum(prop))
    #ax.legend(handles=[p0]+plots)
    ax1.set_xlabel("# bin & cumul. prop. of sites", fontsize=fnt_size)
    # Set the x-axis locator to reduce the number of ticks to 10
    ax1.set_ylabel(r'$\theta_k$', fontsize=fnt_size, rotation = 90)
    ax1.set_title(title, fontsize=fnt_size)
    ax1.legend(handles=plots, loc='best', fontsize = fnt_size*0.5)
    ax1.set_xticks(x_ticks)
    step = len(x_ticks)//(n_ticks-1)
    values = x_ticks[::step]
    new_prop = []
    for val in values:
        new_prop.append(prop[int(val)-2])
    new_prop = new_prop[::-1]
    ax1.set_xticks(values)
    ax1.set_xticklabels([f'{values[k]}\n{val:.2f}' for k, val in enumerate(new_prop)], fontsize = fnt_size*0.8)
    if ax is None:
        # nb of plot_lines represent the number of epochs stored (len(plot_lines) = #breaks+1)
        plt.savefig(title+'_raw_'+str(nb_breaks)+'_breaks.pdf')
        plt.close(fig)
    # return plots
    return ax

def combined_plot(folder_path, mu, tgen, breaks, title = "Title", theta_scale = False, selected_breaks = []):
    my_dpi = 300
    saved_plots_dict = save_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, output = title+"_plotdata.json")
    nb_of_epochs = len(saved_plots_dict["all_epochs"]["plots"])
    best_epoch = saved_plots_dict["best_epoch_by_AIC"]
    print("Best epoch based on AIC =", best_epoch)
    save_k_theta(folder_path, mu, tgen, title, theta_scale, breaks_max = nb_of_epochs, input = title+"_plotdata.json", output = title+"_plotdata.json")

    with open(title+"_plotdata.json", 'r') as json_file:
        loaded_data = json.load(json_file)

    # START OF COMBINED PLOT CODE
        
    # # plot page 1 of summary
    # fig1, ax1 = plt.subplots(2, 2, figsize=(5000/my_dpi, 2970/my_dpi), dpi=my_dpi)
    # # fig1.tight_layout()
    # # Adjust absolute space between the top and bottom rows
    # fig1.subplots_adjust(hspace=0.35)  # Adjust this value based on your requirement
    # # plot page 2 of summary
    # fig2, ax2 = plt.subplots(2, 2, figsize=(5000/my_dpi, 2970/my_dpi), dpi=my_dpi)
    # # fig2.tight_layout()
    # ax1 = plot_raw_stairs(plot_lines = loaded_data['raw_stairs'],
    #                         prop = loaded_data['prop'], title = title, ax = ax1)
    # ax1 = plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'],
    #                             prop = loaded_data['prop'], title = title, ax = ax1, subset=[loaded_data['best_epoch_by_AIC']]+selected_breaks)
    # ax2 = plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'],
    #                         prop = loaded_data['prop'], title = title, ax = ax2)
    # ax1, ax2 = plot_all_epochs_thetafolder(loaded_data, mu, tgen, title, theta_scale, ax = [ax1, ax2])
    
    # fig1.savefig(title+'_combined_p1.pdf')
    # print("Wrote", title+'_combined_p1.pdf')
    # fig2.savefig(title+'_combined_p2.pdf')
    # print("Wrote", title+'_combined_p2.pdf')

    # END OF COMBINED PLOT CODE


    # Start of Parsing real swp2 output
    folder_splitted = folder_path.split("/")
    swp2_summary = "/".join(folder_splitted[:-2])+'/'+folder_splitted[-3]+".final.summary"
    swp2_vals = parse_stairwayplot_output_summary(stwplt_out = swp2_summary)
    swp2_x, swp2_y = swp2_vals[0], swp2_vals[1]
    # End of Parsing real swp2 output
    plot_raw_stairs(plot_lines = loaded_data['raw_stairs'],
                            prop = loaded_data['prop'], title = title, ax = None, max_breaks = breaks)
    plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'], mu = mu, tgen = tgen, subset=[loaded_data['best_epoch_by_AIC']]+selected_breaks,
    # plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'], subset=list(range(0,3))+[loaded_data['best_epoch_by_AIC']]+selected_breaks,
                            prop = loaded_data['prop'], title = title, swp2_lines = [swp2_x, swp2_y], ax = None)
    plot_all_epochs_thetafolder(loaded_data, mu, tgen, title, theta_scale, ax = None)

    # plt.close(fig1)
    # plt.close(fig2)
def parse_stairwayplot_output_summary(stwplt_out, xlim = None, ylim = None, title = "default title", plot = False):
    #col 5
    year = []
    # col 6
    ne_median = []
    ne_2_5 = []
    ne_97_5 = []
    ne_12_5 = []
    # col 10
    ne_87_5 = []
    with open(stwplt_out, "r") as stwplt_stream:
        for line in stwplt_stream:
            ## Line format
            # mutation_per_site	n_estimation	theta_per_site_median	theta_per_site_2.5%	theta_per_site_97.5%	year	Ne_median	Ne_2.5%	Ne_97.5%	Ne_12.5%	Ne_87.5%
            if not line.startswith("mutation_per_site"):
                #not header
                values = line.strip().split()
                year.append(float(values[5]))
                ne_median.append(float(values[6]))
                ne_2_5.append(float(values[7]))
                ne_97_5.append(float(values[8]))
                ne_12_5.append(float(values[9]))
                ne_87_5.append(float(values[10]))

    vals = [year, ne_median, ne_2_5, ne_97_5, ne_12_5, ne_87_5]
    if plot :
        # plot parsed data
        label = ["Ne median", "Ne 2.5%",	"Ne 97.5%",	"Ne 12.5%",	"Ne 87.5%"]
        for i in range(1, 5):
            fig, = plt.plot(year, vals[i], '--', alpha = 0.4)
            fig.set_label(label[i])
        #     # last plot is median
        fig, = plt.plot(year, ne_median, 'r-', lw=2)
        fig.set_label(label[0])
        plt.legend()
        plt.ylabel("Individuals (Ne)")
        plt.xlabel("Time (years)")
        if xlim:
            plt.xlim(xlim)
        if ylim:
            plt.ylim(ylim)
        plt.title(title)
        plt.show()
        plt.close()
    return vals

if __name__ == "__main__":

    if len(sys.argv) != 4:
        print("Need 3 args: ThetaFolder MutationRate GenerationTime")
        exit(0)

    folder_path = sys.argv[1]
    mu = sys.argv[2]
    tgen = sys.argv[3]

    plot_all_epochs_thetafolder(folder_path, mu, tgen)