import numpy as N
import itertools,csv
from matplotlib.dates import datestr2num
from itertools import imap,izip,chain

string_conversions=[
    # conversion function,  numpy dtype
    ( int,                  N.dtype(int)    ),
    ( float,                N.dtype(float)  ),
    ( datestr2num,          N.dtype(float)  ),
    ]
format_change={ # (from_fn,to_fn) : conversion_fn for old converted
    (int,float) : float,
    }
def string_to_dt_cvt(s):
    """
    Converting data to the appropriate type
    """
    for fn,dt in string_conversions:
        try:
            v=fn(s)
            return fn,dt
        except:
            pass
    return str,N.dtype(object)


def get_generator_fn(str_cvt,dt_row,cvt=None):
    var_nm=["x%d" % i for i,x in enumerate(str_cvt)]
    fn_nm=[fn.__name__ for fn,dt in str_cvt]

    indent=" "*4
    header=["def get_data_iterator(row_iter,delim,res):"]
    olditer=[]
    if cvt:
        olditer=[indent+"for "+",".join(var_nm)+" in res['data']:"]
        olditer.extend([indent*2+"x%d=%s(x%d)" % (i,fn,i) for i,fn in cvt])
        olditer.append(indent*2+"print ("+",".join(var_nm)+")")
        olditer.append(indent*2+"yield ("+",".join(var_nm)+")")
    rowiter=[
        indent+"yield ("+",".join(["%s('%s')" % (f,r) for f,r in zip(fn_nm,dt_row)])+")",
        indent+"for row in row_iter:",
        indent*2+",".join(var_nm)+"=row.split(delim)",
        indent*2+"try:",
        indent*3+"yield ("+",".join(["%s(%s)" % (f,v) for f,v in zip(fn_nm,var_nm)])+")",
        indent*2+"except:",
        indent*3+"res['row']=row",
        indent*3+"return",
        indent+"res['finished']=True",
        ]
    return "\n".join(header+olditer+rowiter)

def load(fname,delim = ',',has_varnm = True, prn_report = True):
    """
    Loading data from a file using fromiter. Returns a recarray.
    """

    row_iter=open(fname,'rb')
    row0=map(str.strip,row_iter.next().split(delim))
    if not has_varnm:
        varnm = ['col%s' % str(i+1) for i in xrange(len(row0))]
        dt_row=row0
    else:
        varnm = [i.strip() for i in row0]
        dt_row=map(str.strip,row_iter.next().split(delim))

    str_cvt=[string_to_dt_cvt(item) for item in dt_row]

    res={}
    res['finished']=False
    cvt=None
    while not res['finished']:
        descr=[(name,dt) for name,(fn,dt) in zip(varnm,str_cvt)]
        generator_fn=get_generator_fn(str_cvt,dt_row,cvt)
        exec(compile(generator_fn,'<string>','exec'))
        data=N.fromiter(get_data_iterator(row_iter,delim,res),dtype=descr).view(N.recarray)
        if not res['finished']:
            res['data']=data
            dt_row=map(str.strip,res['row'].split(delim))
            new_cvt=[string_to_dt_cvt(item) for item in dt_row]
            conv=[(i,(f1,f2)) for i,((f1,d1),(f2,d2)) in enumerate(zip(str_cvt,new_cvt)) if f1!=f2]
            if set([x for i,x in conv])-set(format_change):
                raise "Unsupported string representation change in columns"
            cvt=[(i,format_change[x].__name__) for i,x in conv]
            str_cvt=new_cvt
            for i,(f1,f2) in conv:
                print "Inconsistent string representation: converting %s from %s to %s" % (varnm[i],f1.__name__,f2.__name__)

    # load report
    if prn_report:
        print "##########################################\n"
        print "Loaded file: %s\n" % fname
        print "Nr obs: %s\n" % data.shape[0]
        print "Variables and datatypes:\n"
        for i in data.dtype.descr:
            print "Varname: %s, Type: %s, Sample: %s" % (i[0], i[1], str(data[i[0]][0:3]))
            print "\n##########################################\n"

    return data

def show_dates(dates):
	return N.array([i.strftime('%d %b %y') for i in pylab.num2date(dates)])

if __name__ == '__main__':

	# creating data
	data = [['col1', 'col2', 'col3', 'col4', 'col5', 'col6'],
			['1','3','1/97','1.12','2.11','1.2'],
			['1','2','3/97','1.21','3.12','1.43'],
			['2','1','2/97','1.12','2.11','1.28'],
			['2.0','2','4/97','1.33','2.26','1.23'],
			['2.1','2','5/97','1.73','2.42','1.26']]

	# saving data to csv file
	f = open('testdata.csv','wb')
	output = csv.writer(f)
	for i in data:
		output.writerow(i)
	f.close()

	# opening data file with variable names
	ra = load('testdata.csv')