import csv
import matplotlib.pyplot as plt

"""
Smoothing by different methods and parameters
"""


#=============================
# parameters
#=============================
csvfile = 'xrd.csv'
outfile_template = 'smoothing-{}-{}.csv'
xmin = 37.0
xmax = 46.0


def SmoothingBySimpleAverage(y, n):
    n2 = int(n / 2);
    ndata = len(y);
    ys = []
    for i in range(0, ndata):
        c = 0;
        ys.append(0.0);
        for k in range(i - n2, i + n2 + 1):
            if k < 0 or k >= ndata:
                continue
            ys[i] += y[k]
            c += 1
        if c > 0:
            ys[i] /= c;
        else:
            ys[i] = y[i]
    return ys;


def SmoothingByPolynomialFit(y, n):
    m = int(n / 2);
    W23 = (4.0 * m * m - 1.0) * (2.0 * m + 3.0) / 3.0;
    w23j = [0.0]*n
    for j in range(-m, m+1):
        w23j[j + m] = (3.0 * m * (m+1.0) - 1.0 - 5.0 * j * j) / W23

    ndata = len(y)
    ys = []
    for i in range(0, ndata):
        c = 0.0;
        ys.append(0.0);
        for j in range(-m, m+1):
            k = i + j
            if k < 0 or k >= ndata:
                continue
            ys[i] += w23j[j+m] * y[k]
            c += w23j[j+m]
        if c > 0:
            ys[i] /= c
        else:
            ys[i] = y[i]
    return ys;

def savecsv(outfile, x, y, ys):
    try: 
        print("Write to [{}]".format(outfile))
        f = open(outfile, 'w')
    except:
#    except IOError:
        print("Error: Can not write to [{}]".format(outfile))
    else:
        fout = csv.writer(f, lineterminator='\n')
        fout.writerow(('x', 'y(raw)', 'y(smooth)'))
#        fout.writerows(data)
        for i in range(0, len(x)):
            fout.writerow((x[i], y[i], ys[i]))
        f.close()


def main():
    global csvfile
    global outfile

    print("Read data from [{}]".format(csvfile))
    x = []
    y = []
    with open(csvfile) as f:
        fin = csv.reader(f)
        next(fin)
        c = 0
        for row in fin:
            if c == 0:
                label = row
                print("{}\t{}".format(label[0], label[1]))
            else:
                xx = float(row[0])
                yy = float(row[1])
                if xx < xmin or xmax < xx:
                    continue

                print("{}\t{}".format(row[0], row[1]))
                x.append(xx)
                y.append(yy)
            c += 1
    ndata = len(x)
    print("")

#=============================
# prepare graph
#=============================
    fig, ax = plt.subplots(3, 3, figsize = (8, 8))


    ys = []
    print("Smoothing by simple moving average")
    icase = 0
    nsmooth = 3;
    method = 'simple'
    print("nsmooth={}".format(nsmooth))
    ys.append(SmoothingBySimpleAverage(y, nsmooth))
    for i in range(0, ndata):
         print("{}\t{}\t{}".format(x[i], y[i], ys[icase][i]))
    savecsv(outfile_template.format(method, nsmooth), x, y, ys[icase])
    print("")
    ax[0, 0].plot(x, ys[0], label = "{}p {}".format(nsmooth, method), linewidth = 0.5)
    ax[0, 0].plot(x, y,     linestyle = 'none', marker = 'o', markersize = 0.3)
    ax[0, 0].set_title("{}p {}".format(nsmooth, method))

    icase = 1
    nsmooth = 11;
    method = 'simple'
    print("nsmooth={}".format(nsmooth))
    ys.append(SmoothingBySimpleAverage(y, nsmooth))
    for i in range(0, ndata):
         print("{}\t{}\t{}".format(x[i], y[i], ys[icase][i]))
    savecsv(outfile_template.format(method, nsmooth), x, y, ys[icase])
    print("")
#    ax[0, 1].plot(x, y,     linestyle = 'none', marker = 'o', markersize = 0.5)
    ax[0, 1].plot(x, ys[1], label = "{}p {}".format(nsmooth, method), linewidth = 0.5)
    ax[0, 1].set_title("{}p {}".format(nsmooth, method))

    icase = 2
    nsmooth = 31;
    method = 'simple'
    print("nsmooth={}".format(nsmooth))
    ys.append(SmoothingBySimpleAverage(y, nsmooth))
    for i in range(0, ndata):
         print("{}\t{}\t{}".format(x[i], y[i], ys[icase][i]))
    savecsv(outfile_template.format(method, nsmooth), x, y, ys[icase])
    print("")
#    ax[0, 2].plot(x, y,     linestyle = 'none', marker = 'o', markersize = 0.5)
    ax[0, 2].plot(x, ys[2], label = "{}p {}".format(nsmooth, method), linewidth = 0.5)
    ax[0, 2].set_title("{}p {}".format(nsmooth, method))

    icase = 3
    nsmooth = 51;
    method = 'simple'
    print("nsmooth={}".format(nsmooth))
    ys.append(SmoothingBySimpleAverage(y, nsmooth))
    for i in range(0, ndata):
         print("{}\t{}\t{}".format(x[i], y[i], ys[icase][i]))
    savecsv(outfile_template.format(method, nsmooth), x, y, ys[icase])
    print("")
#    ax[1, 0].plot(x, y,     linestyle = 'none', marker = 'o', markersize = 0.5)
    ax[1, 0].plot(x, ys[3], label = "{}p {}".format(nsmooth, method), linewidth = 0.5)
    ax[1, 0].set_title("{}p {}".format(nsmooth, method))

    print("Smoothing by polynomial fit")
    icase = 4
    nsmooth = 11;
    method = 'polynomialfit'
    print("nsmooth={}".format(nsmooth))
    ys.append(SmoothingByPolynomialFit(y, nsmooth))
    for i in range(0, ndata):
         print("{}\t{}\t{}".format(x[i], y[i], ys[icase][i]))
    savecsv(outfile_template.format(method, nsmooth), x, y, ys[icase])
    print("")
#    ax[1, 1].plot(x, y,     linestyle = 'none', marker = 'o', markersize = 0.5)
    ax[1, 1].plot(x, ys[4], label = "{}p {}".format(nsmooth, method), linewidth = 0.5)
    ax[1, 1].set_title("{}p {}".format(nsmooth, method))

    icase = 5
    nsmooth = 31;
    method = 'polynomialfit'
    print("nsmooth={}".format(nsmooth))
    ys.append(SmoothingByPolynomialFit(y, nsmooth))
    for i in range(0, ndata):
         print("{}\t{}\t{}".format(x[i], y[i], ys[icase][i]))
    savecsv(outfile_template.format(method, nsmooth), x, y, ys[icase])
    print("")
#    ax[1, 2].plot(x, y,     linestyle = 'none', marker = 'o', markersize = 0.5)
    ax[1, 2].plot(x, ys[5], label = "{}p {}".format(nsmooth, method), linewidth = 0.5)
    ax[1, 2].set_title("{}p {}".format(nsmooth, method))

    icase = 6
    nsmooth = 51;
    method = 'polynomialfit'
    print("nsmooth={}".format(nsmooth))
    ys.append(SmoothingByPolynomialFit(y, nsmooth))
    for i in range(0, ndata):
         print("{}\t{}\t{}".format(x[i], y[i], ys[icase][i]))
    savecsv(outfile_template.format(method, nsmooth), x, y, ys[icase])
    print("")
#    ax[2, 0].plot(x, y,     linestyle = 'none', marker = 'o', markersize = 0.5)
    ax[2, 0].plot(x, ys[6], label = "{}p {}".format(nsmooth, method), linewidth = 0.5)
    ax[2, 0].set_title("{}p {}".format(nsmooth, method))

    plt.tight_layout()

    plt.pause(0.1)
    print("Press ENTER to exit>>", end = '')
    input()

    exit()


if __name__ == '__main__':
    main()
 