from ess import fieldmap
import os

"""
Example

The parameters in the following six rows define the path to your fieldmap, no name is needed, this script
scans the folder and converts all files, assumes all are binary fieldmaps, to ASCII fieldmaps.
The dimension of the field is also needed.
If you need to cut the field, the n_enter and n_exit should be adjusted.
"""
plot = False  # Make comparison plots
mypath = "maps"  # where the fieldmaps are stored
n_enter = 0  # Number of lines to be removed from entrance of fieldmap
n_exit = 0  # Number of lines to be removed from exit of fieldmap


if plot:
    from matplotlib import pyplot as plt

for f in os.listdir(mypath):

    print(f)
    fullpath = os.path.join(mypath, f)

    ftype = ""
    if "ascii" in f:
        ftype = "a"
    elif "binary" in f:
        ftype = "b"
    fdim = 0
    if "1D" in f:
        fdim = 1
    elif "2D" in f:
        fdim = 2
    elif "3D" in f:
        fdim = 3

    if n_enter or n_exit:
        Header, Field = fieldmap.cut_fieldmap_length(
            fullpath, fdim, n_enter, n_exit, ftype
        )
    else:
        Header, Field = fieldmap.read_fieldmap(fullpath, fdim, ftype)

    fieldmap.save_to_ascii(Header, Field, "new_" + f)
    oldfield = fieldmap.field_on_axis(Header, Field, fdim)
    origfield = fieldmap.field_on_axis_from_file(fullpath, fdim, "a")
    newfield = fieldmap.field_on_axis_from_file("new_" + f, fdim, "a")
    if plot:
        # Plot for comparison
        plt.figure()
        plt.title(f)
        if fdim == 1:
            plt.plot(oldfield, label="old")
            plt.plot(newfield, label="new")
            plt.plot(origfield, label="orig")
        else:
            plt.plot(oldfield[0, :], oldfield[1, :], label="old")
            plt.plot(newfield[0, :], newfield[1, :], label="new")
            plt.plot(origfield[0, :], origfield[1, :], label="orig")
        plt.legend()
    else:
        if fdim == 1:
            avg_diff = sum(newfield - oldfield) / len(newfield)
        else:
            avg_diff = sum(newfield[1, :] - oldfield[1, :]) / len(newfield[1, :])
        print("Average diff", avg_diff)

if plot:
    plt.show()