#plot_prf.py

from datetime import datetime
import os, numpy as np, pandas as pd, sys, re, math
import matplotlib.pyplot as plt

"""adding path to look for private modules"""
modules_lib = os.path.join(os.path.dirname(__file__), "modules")
if not modules_lib in sys.path:
    sys.path.append(modules_lib)
"""importing private modules"""
from wrappers import TraceCalls
from find_tools import findfile, findfiles
from prf_plot_tools import extract_out_file, calculate_size, excluded_regions
from prf_plot_tools import prepare_refinement_result, extract_pcr_file

@TraceCalls()
def main(path = "", printsize = "two_in_docx"):
    """ docstring! """
    def help():
        available_sizes = ["two_in_docx two_in_ppt one_in_docx one_in_ppt".split()]
        print("the available commands are:")
        for s in available_sizes:
            print("\t" + s)

    printsize = "two_in_ppt" # graphs are 17 cm wide, perfect for two in .ppt
    printsize = "one_in_docx" # graphs are 17 cm wide, perfect for one in word
    printsize = "one_in_ppt" # single grpah in .ppt
    printsize = "two_in_docx" # graphs are 8 cm wide, perfect for two in word

    """ path bookkeeping """
    if path == "":
        path = os.getcwd()
    os.chdir(path)
    folders = [folder for folder in os.listdir() if os.path.isdir(folder) ]

    for folder in folders:
        print("found folder {}".format(folder))
    print("")

    specified_size = True
    if printsize == "two_in_docx":
        graph_width = 8 # cm
        graph_height = 6 # cm
        bigger_text = 10 # do not use
        medium_text = 8 # still quite big
        smaller_text = 6 # better
        micro_text = 4 # 4 good for refinement info.
        subplot_right_space = 6.5
        dotsize = 1 #
        linewidth = 1
        refinfo_pos = (1.02,-0.07)
        ylabelpos_pre = (0.01,1.05)
        ylabelpos_post = (-0.27,0.50)
        xlabelpos_pre = (1.05,0.05)
        xlabelpos_post = (0.50,-0.12)

    elif printsize == "two_in_ppt":
        graph_width = 15 # cm
        graph_height = graph_width*3/4 # cm
        bigger_text = 14 # do not use
        medium_text = 12 # still quite big
        smaller_text = 10 # better
        micro_text = 8 # 4 good for refinement info.
        dotsize = 2
        linewidth = 2
        subplot_right_space = 7.5
        refinfo_pos = (1.03,-0.06)
        ylabelpos_pre = (-0.008,1.05)
        ylabelpos_post = (-0.18,0.50)
        xlabelpos_pre = (1.05,-0.043)
        xlabelpos_post = (0.50,-0.07)

    elif printsize == "one_in_docx":
        graph_width = 17 # cm
        graph_height = graph_width*3/4 # cm
        bigger_text = 14 # do not use
        medium_text = 12 # still quite big
        smaller_text = 10 # better
        micro_text = 8 # 4 good for refinement info.
        dotsize = 2
        linewidth = 2
        subplot_right_space = 6.6
        refinfo_pos = (1.015,-0.055)
        ylabelpos_pre = (-0.007,1.05)
        ylabelpos_post = (-0.14,0.50)
        xlabelpos_pre = (1.05,-0.043)
        xlabelpos_post = (0.50,-0.07)

    elif printsize == "one_in_ppt":
        graph_width = 32 # cm
        graph_height = graph_width*2/5 # cm
        bigger_text = 20 # do not use
        medium_text = 18 # still quite big
        smaller_text = 16 # better
        micro_text = 10 # 4 good for refinement info.
        dotsize = 3
        linewidth = 2
        subplot_right_space = 4.5
        refinfo_pos = (1.01,-0.065)
        ylabelpos_pre = (-0.009,1.05)
        ylabelpos_post = (-0.09,0.50)
        xlabelpos_pre = (1.05,-0.043)
        xlabelpos_post = (0.50,-0.07)

    for folder in folders:
        print("\nworking on", folder); os.chdir(path); os.chdir(folder)

        prf = findfile(".prf"); pcr = findfile(".pcr"); out = findfile(".out");
        if not os.path.exists(prf): print("!-!-!-! skipping \"" + folder + "\" + --> no .prf was found"); continue
        if not os.path.exists(pcr): print("!-!-!-! skipping \"" + folder + "\" + --> no .pcr was found"); continue
        if not os.path.exists(out): print("!-!-!-! skipping \"" + folder + "\" + --> no .out was found"); continue


        try:
            """ try if we can extract DMC data to plot """
            with open(prf, "r") as reading:
                for num,line in enumerate(reading):
                    if num >= 5:
                        numbers = [float(number) for number in line.strip().split() ]
                        if len(yobs) != entries:
                            yobs += numbers
                        elif len(ycalc) != entries:
                            ycalc += numbers
                        else:
                            bragg_reflections += numbers
                    elif num == 2:
                        end,start,step = [float(number) for number in line.strip().split()[0:3]]
                        angles = np.arange(start,end+step,step)
                        entries = len(angles)
                        yobs = list()
                        ycalc = list()
                        bragg_reflections = list()
            """ UNDER DEVELOPMENT: plot indicating lines for bragg reflections """
            # last_value = 0
            # b_r = bragg_reflections
            # for i,value in enumerate(b_r):
            #     if value > 10000:
            #         continue
            #     elif value < last_value:
            #         b_r = b_r[:i]
            #         break
            #     last_value = value
            # br = pd.DataFrame()
            # midway = int(len(b_r)/2)
            # for col_name, col in zip(["ID","2theta"],[ b_r[:midway], b_r[midway:]]):
            #     br[col_name] = col
            """ UNDER DEVELOPMENT: plot indicating lines for bragg reflections """


            """ preparing dataframe """
            df = pd.DataFrame()
            for col_name, col in zip(["2Theta", "Yobs", "Ycal"], [angles, yobs, ycalc]):
                df[col_name] = col
            print("found a DMC .prf")
        except:
            """ if DMC extraction fails, extract x-ray data """
            nb_excluded_regions = excluded_regions()
            df = pd.read_csv(prf, skiprows = 3+nb_excluded_regions, delimiter="[\t]+", engine="python")
            df = df[["2Theta","Yobs","Ycal"]]
            for col in df.columns.values:
                df[col] = pd.to_numeric(df[col], errors="coerce")
            df = df.dropna(axis=0)
            print("found a R1D1 .prf")

        """ check if model data are actual values! otherwise skip datafolder """
        if df.empty: sys.stdout.write("!-!-!-! skipping folder \"" + folder + "\". DataFrame is empty (values in prf might be NaNs)\n"); continue

        df["Yobs-Ycal"] = df["Yobs"]-df["Ycal"]

        """ shortening the view of the diffractogram """
        df = df[df["2Theta"] < 92.900]
        df = df[df["2Theta"] > 20.00]

        """ style of pandas plotting """
        plt.style.use("seaborn")

        """ plot scatter and line plots of observed and calculated """
        fig, ax1 = plt.subplots(figsize = cm2inch(graph_width, graph_height) )
        plt.subplots_adjust(right=subplot_right_space)

        """ extraction of:
        phase names (dict)
        phase fractions (dict)
        wavelength (float)
        refined parameter value and error (dict) """
        wavelength, frac_info, phases, refined_par = extract_out_file()

        """ calculates:
        ab and c-sizes (dict of dicts) """
        size_info = calculate_size(
            wavelength = wavelength,
            phases = phases,
            refined_par = refined_par,
            )

        """ extracting spacegroup from PCR file """
        spacegroups = extract_pcr_file()

        """ preparing information such as refined parameters, phases, phase fractions
        crystallite sizes and so to be printed on canvas. """
        combined_text = prepare_refinement_result(
            phases = phases,
            refined_par = refined_par,
            size_info = size_info,
            frac_info = frac_info,
            spacegroups = spacegroups,
            )

        """ adjusting the position of refinement info. """
        # text_xpos, text_ypos = fractional_position


        y = "Yobs"
        ax1 = df.plot(
            "2Theta",
            y,
            kind="scatter",
            s = dotsize,
            ax = ax1,
            fontsize = smaller_text,  # xtick and ytick labels
            label = y, # for legend
            c = "r",
            alpha=0.5,
            )
        df.plot(
            "2Theta",
            "Ycal",
            linewidth=linewidth,
            ax = ax1,
            alpha=0.8,
            )
        df.plot(
            "2Theta",
            "Yobs-Ycal",
            linewidth = linewidth,
            ax=ax1,
            alpha=0.3,
            )

        """ plotting refinement info on canvas"""
        ttt = ax1.text(
            refinfo_pos[0],
            refinfo_pos[1],
            combined_text,
            size=micro_text,
            fontname = "monospace",
            picker=True,
            transform = ax1.transAxes,
            # zorder = 0,
            )

        """ making the refinement info have a semitransparent white background """
        # ttt.set_bbox(dict(facecolor='white', alpha=0.5, edgecolor='white'))

        """ changing the position of the legens to make Yobs appear at the top """
        handles,labels = ax1.get_legend_handles_labels()
        handles = [handles[2], handles[0], handles[1]]
        labels = [labels[2], labels[0], labels[1]]

        """ setting legend to top left corner """
        leg = plt.legend(handles,labels, loc="best", prop = {"size": micro_text})


        """ setting title to folder name """
        title = os.path.basename(folder)
        plt.title(title, size = medium_text)

        """ setting dummy y- and xlabel names. These are changed later
        this is done due to the effect of using plt.tight_layout()
        in combination with rotation and shifting of the text object"""
        ylabel = plt.ylabel("Counts", size = smaller_text)
        plt.xlabel(r"2$\theta$", size = smaller_text)

        """ repositioning and rotation the ylabel """
        ax1.yaxis.set_label_coords(*ylabelpos_pre)
        ylabel.set_rotation(0)

        """ Xlabel must be here otherwise plt.tight_layout is not given the
        additional white space for the ref_info text """
        ax1.xaxis.set_label_coords(*xlabelpos_pre)

        """ name of output file """
        figname = folder+".png"
        figname = os.path.join(path,figname)

        """ scientific notation of major ytick values """
        import matplotlib.ticker as mtick
        ax1.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.0e'))

        """ ensure nothing is left outside the plot """
        plt.tight_layout()

        """ repositioning and renaming xlabel and ylabel
        _after_ the use of plt.tight_layout()"""
        ylabel.set_rotation(90)
        ax1.yaxis.set_label_coords(*ylabelpos_post)
        ax1.xaxis.set_label_coords(*xlabelpos_post)
        plt.xlabel(r"scattering angle 2$\theta$", size = smaller_text)
        ylabel = plt.ylabel("signal count", size = smaller_text)

        # import pdb; pdb.set_trace(); print("None =",None)

        """ saving the figure in high resolution """
        plt.savefig(figname, dpi=500, format = "png")

        """ avoid memory leak """
        plt.close()



def cm2inch(*tupl):
    """ helper function to convert figsize-cm-tuple to inches """
    inch = 2.54
    if isinstance(tupl[0], tuple):
        return tuple(i/inch for i in tupl[0])
    else:
        return tuple(i/inch for i in tupl)



if __name__ == "__main__":
    # import pdb; pdb.set_trace(); import time; time.sleep(1); print("sys.argv =\n",sys.argv,"\n")
    if len(sys.argv) > 1:
        main(path = sys.argv[1])
    else:
        main()
else:
    if len(sys.argv) > 1:
        main(path = sys.argv[1])










# print("\n\nEND_OF_SCRIPT")
# print(sys.version_info)
# print(sys.executable)
