import gviz_api
from django.db.models.query import ValuesQuerySet
 
# mapping of model field types to google vizualization column types (all not 
# in list are 'string' type) valid google viz data types are:  'string' 
# 'number' 'boolean' 'date' 'datetime' 'timeofday'
fieldmap = {'DateField':'date','DateTimeField':'datetime','BooleanField':'boolean',
    'IntegerField':'number','DecimalField':'number','BigIntegerField':'number',
    'FloatField':'number','TimeField':'timeofday','NullBooleanField':'boolean'}
 
# { 'queryname' : [query object, { field_name : (datatype, descriptive_name) }, 
# { name of field with choices : { choices } }] }
queries = {}
 
def add_query(qry, qryName, makeStr=set()):
    """Takes a Django QuerySet or ValuesQuerySet and an (arbitrary, unique) name for it,
    and adds it to the collection of queries to use in Google Visualization DataTables.  Also
    takes an optional third argument of a set of column names to be put in the DataTable as 
    strings, ignoring their data type in the model."""
    
    if type(qry) is ValuesQuerySet:
        add_values_qry(qry, qryName, makeStr)
    else:
        add_full_qry(qry, qryName, makeStr)
    
 
def get_viz_json(qryName):
    """Takes a previously added query name and returns JSon representation of a 
    Google Visualization DataTable."""
    tbl = gviz_api.DataTable(queries[qryName][1])
    
    if type(queries[qryName][0]) is ValuesQuerySet:
        tbl.LoadData(get_viz_data_values(qryName))
    else:
        tbl.LoadData(get_viz_data_full(qryName))
    
    json = tbl.ToJSon()
    return json
 
def add_values_qry(qry, qryName, makeStr):
    model = qry.model
    desc = {} # { field_name : (datatype, descriptive_name) }
    choices = {}
    
    flds = qry.field_names
    for f in flds:
        fld = model._meta.get_field(f)
        t = fld.get_internal_type() # returns Django model field type as string
        gt = 'string' # default to string type
        
        if (not f in makeStr) and fieldmap.has_key(t):
            gt = fieldmap[t]
        
        if fld.choices:
            choices[f] = dict(fld.choices)
            desc[f] = ('string', fld.verbose_name)
        else:
            desc[f] = (gt, fld.verbose_name)
        
    
    if qry.aggregate_names:
        flds_agg = qry.aggregate_names # if query has annotations, add them as 'number' type
        for fa in flds_agg:
            descriptive_name = fa.replace('_', ' ')
            desc[fa] = ('number', descriptive_name)
        
    
    queries[qryName] = [qry, desc, choices]
 
def add_full_qry(qry, qryName, makeStr):
    model = qry.model
    desc = {} # { field_name : (datatype, descriptive_name) }
    choices = {}
    model_name = model.__name__
    flds = model._meta.get_all_field_names()
    
    for f in flds:
        try:
            fld = model._meta.get_field(f)
        except:
            continue # in case of FK field
        
        t = fld.get_internal_type() # returns Django model field type as string
        if t == 'ForeignKey':
            continue
        
        gt = 'string' # default to string
        
        if (not f in makeStr) and fieldmap.has_key(t):
            gt = fieldmap[t]
            
        desc[f] = (gt, fld.verbose_name)
        
        if fld.choices:
            choices[f] = dict(fld.choices)
            desc[f] = ('string', fld.verbose_name)
        else:
            desc[f] = (gt, fld.verbose_name)
        
    
    try:
        flds_agg = qry.query.aggregates.keys()
        for fa in flds_agg:
            descriptive_name = fa.replace('_', ' ')
            desc[fa] = ('number', descriptive_name)
        
    except:
        pass
    
    queries[qryName] = [qry, desc, choices]
 
def get_viz_data_full(qryName):
    """Takes a QuerySet name and returns a dataset to load into a Google Visualization DataTable."""
    data = []
    qry = queries[qryName][0]
    qry.update() # force query to re-evaluate
    flds = queries[qryName][1].keys()
    choices = queries[qryName][2]
    
    for itm in qry:
        row = {}
        for fld in flds:
            if choices.has_key(fld):
                row[fld] = choices[fld][itm.__getattribute__(fld)]
            else:
                row[fld] = itm.__getattribute__(fld)
            
        
        data.append(row)
    
    return data
 
def get_viz_data_values(qryName):
    """Takes a ValuesQuerySet name and returns a dataset to load into a Google Visualization DataTable."""
    data = []
    qry = queries[qryName][0]
    qry.update() # force query to re-evaluate
    flds = queries[qryName][1].keys()
    choices = queries[qryName][2]
    
    for itm in qry:
        row = {}
        for fld in flds:
            if choices.has_key(fld):
                row[fld] = choices[fld][itm[fld]]
            else:
                row[fld] = itm[fld]
            
                  
        data.append(row)
    
    return data