import numpy as np
import numpy.random as ran
import pylab as plt
import sys
from response_plot_files import *

# Least-squares fit of data to a polynomial.
# The data is assumed to be approximately linear.

def GetData(data_file, targets=[0, 1]) :
    # Targets is the list of only two data arrays to be used as x and y.
    file_in = DataFromFile(data_file, ',')      # Delimiter is assumed to be a comma.
    return file_in.data[targets[0]], file_in.data[targets[1]]


def FitDataToPolynomial(x, y, order) :
    coeffs, residuals, rank, singular_values, rcond = np.polyfit(x, y, order, full=True)
    crev = coeffs[:: -1]    # Reverse the order of the coeffs array
    sigma = np.sqrt(residuals[0]/len(x))
    print('sigma = ' + str(sigma))
    z =np.zeros( (order+1, len(x)) )
    ztot = np.zeros(len(x), float)
    for n in range(0, order+1, 1) :
        z[n] += crev[n] * np.power(x, n)
        ztot += z[n]
    return ztot, z, crev, sigma

def PlotFit(x, y, z, ztot, coeffs, sigma, title, text_loc=[0.3, 0.7]) :
    # text_loc is the location of the text box where [0,0] is the lower left and [1,1] is the upper right.
    plt.plot(x, y, label='Data')
    coeff_string = 'Fit Coefficients\n'
    n = 0
    for c in coeffs :
        coeff_string += str(n) + ':   ' + str(coeffs[n].round(decimals=4)) + '\n'
        plt.plot(x, z[n], label='$x^{' + str(n) + '}$ piece')
        n += 1
    coeff_string += '$\sigma =$ ' + str(sigma.round(decimals=4)) + '\n'
    plt.plot(x, ztot, label = 'Total Fit')
    #plt.legend(loc='best')
    plt.legend(loc='upper left')
    plt.title( title )
    y_min = np.min(ztot)
    y_max = np.max(ztot)
    for zn in z :
        zn_max = np.max(zn)
        zn_min = np.min(zn)
        if zn_max > y_max :
            y_max = zn_max
        if zn_min < y_min :
            y_min = zn_min
    yloc = y_min + text_loc[1] * (y_max - y_min)
    x_min = np.min(x)
    x_max = np.max(x)
    xloc = x_min + text_loc[0] * (x_max - x_min)
    plt.text(xloc, yloc, coeff_string)
    plt.show()
    plt.close()
    return

def MakeResidualHistogram(data, fit, bins, sigma) :
    # Make a histogram of the abs(error)
    #e2 = np.abs(data - fit)
    e2 = data - fit
    hist = np.histogram(e2, bins=bins)
    # Histogram returns the array of values and an array of bin edges, which has one more value.
    # Create bin center.
    bincenters = 0.5*(hist[1][1:]+hist[1][:-1])
    histo_max = np.max(hist[0])
    # Simulate the set of error, but with more samples for less noise on the noise.
    #   The mean should be 0.0, but it will fluctuate about that from run to run.
    normal_dist = ran.normal(0.0, sigma, len(e2)*3000) * histo_max * np.sqrt(2.0*np.pi)*sigma
    normal_hist = np.histogram(normal_dist, bins=bins)
    normal_hist_max = np.max(normal_hist[0])
    #print bins, len(bincenters), len(hist[0])
    #plt.plot(bincenters, hist[0], label='Data')
    plt.plot(bincenters, hist[0][1:], label='Data')
    #plt.plot(bincenters, normal_hist[0]*histo_max/normal_hist_max, label='Simulation')
    plt.plot(bincenters, normal_hist[0][1:]*histo_max/normal_hist_max, label='Simulation')
    the_title = 'Error Histogram with $\sigma = $' + str(sigma)
    plt.title(the_title)
    plt.legend(loc='best')
    plt.show()
    plt.close()


def Test() :
    order = 3   # polynomial order
    #graph = 'plot'
    graph = 'error'

    # Choose data file and columns of data to use.
    choice = 1
    if choice ==1 :
        data_file = 'linear_data.csv'
        targets = [1, 2]
        # text_loc is the location of the text box where [0,0] is the lower left and [1,1] is the upper right.
        text_loc = [0.3, 0.7]
    elif choice == 2 :
        data_file = 'xy.csv'
        # text_loc is the location of the text box where [0,0] is the lower left and [1,1] is the upper right.
        text_loc = [0.3, 0.7]
        targets = [0, 1]
    x, y = GetData(data_file, targets=targets)
    graph_title = 'Polynomial of order ' + str(order) + ' fit to arrays ' + str(targets[0]) + ' and ' + str(targets[1])
    graph_title += ' of ' + data_file
    #x, y, graph_title = MakeData(data_coeffs, xmax, xstep, noise_amp, noise_sigma)
    fit, fit_n, coeffs, sigma = FitDataToPolynomial(x, y, order)
    bins = np.int(len(y)/10)

    # Only one of the following two operations can be performed without a plot error.
    if graph == 'plot' :
        PlotFit(x, y, fit_n, fit, coeffs, sigma, graph_title, text_loc=text_loc)
    elif graph == 'error' :
        MakeResidualHistogram(y, fit, bins, sigma)

#-------------------------------------------------------------------------------------------

if __name__ == '__main__' :
    Test()