import matplotlib.pyplot as plt
import numpy as np
from scipy import optimize
import itertools
import sys
import math

def readData(filename):
    f = open(filename, 'r')
    titles = f.readline().rstrip('\n\r').split(',')
    dots = []
    lines = f.readlines()
    f.close()
    for line in lines:
        dots.append([float(s) for s in line.rstrip('\\').split(',')])
    #transposing dots for convinience
    dots = map(list, zip(*dots))
    return (titles, dots)

def fit(f, x, y, *args):
    popt = optimize.curve_fit(f, x, y, args)
    newx = np.linspace(x [0], x [-1], 100)
    newy = f(newx, *(popt [0]))
    return (newx, newy)

def fitPlot(title, columns, dots, f, *args):
    #building x and xerror
    x = np.array(dots [0])
    firstValue = 1
    xerror = np.array([0.0] * len(x))
    if '\delta' in columns [1]:
        firstValue += 1
        xerror = np.array(dots [1])
    #each plot has individual marker
    marker = itertools.cycle(('x', 'o', 'D', '*', 'v', '^', '+', 'h'))
    #plotting
    for i in xrange(firstValue, len(dots)):
        if '\delta' in columns [i]:
            plt.errorbar(x, dots [i - 1], yerr = dots [i], fmt = None, ecolor = 'black')
        else:
            plt.plot(x, dots [i], marker = marker.next(), linestyle = '', color = 'black', label = columns [i])
            #fitting
            fitx, fity = fit(f, x, np.array(dots [i]), *args)
            plt.plot(fitx, fity, color = 'black')
            plt.errorbar(x, dots [i], xerr = xerror, fmt = None, ecolor = 'black')
    #some eye-candy
    plt.title(title)
    plt.xlabel(columns [0])
    plt.legend(loc = 'best', fancybox = True, numpoints = 1)
    plt.grid(True)

if __name__ == "__main__":
    columns, dots = readData('plot_data.csv')
    #example of polynomial fitting
    #f = lambda x, a, b, c: (a * x + b) * x + c
    #example of elementary function fitting
    f = lambda x, a, b, c: [a * math.exp(arg * arg * b) + c for arg in x]
    a, b, c = 1.0, -0.25, 1.0
    fitPlot('$simple\,plot$', columns, dots, f, a, b, c)
    plt.ylabel('$F\,and\,U\,values$')
    plt.show()