######################################################################
# Copyright (c) 2007, Petteri Aimonen <jpa@kapsi.fi>
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice and this list of conditions.
#     * Redistributions in binary form must reproduce the above copyright
#       notice and this list of conditions in the documentation and/or other
#       materials provided with the distribution.

'''High-level data access functions'''

import datetime
from database import db_conn
cursor = db_conn.cursor()

####################
# Groups & sections

def changegroup(animalid, newgroup):
    '''Updates group for animal'''
    cursor.execute("UPDATE animals SET groupid=? WHERE animalid=?",
                   (newgroup, animalid))

def groupinfo(groupid):
    '''Get group information. Returns section, pen.'''
    section = groupid // 10 # Last number is per-pen
    pen = groupid % 10
    return section, pen

def getgroup(animalid):
    '''Get groupid based on animalid or None'''
    cursor.execute("SELECT groupid FROM animals WHERE animalid=?", (animalid,))
    row = cursor.fetchone()
    if row is None:
        return None
    else:
        return row[0]

def groups_in_section(section):
    '''Get a sorted list of groups in a section'''
    cursor.execute("SELECT groupid FROM animals " \
                   "WHERE SUBSTR(groupid, 1, LENGTH(groupid) - 1)=? " \
                   "GROUP BY groupid", (str(section),))
    r = [row[0] for row in cursor.fetchall()]
    r.sort()
    return r

def sectionlist():
    '''Get a sorted list of all sections'''
    cursor.execute("SELECT SUBSTR(groupid, 1, LENGTH(groupid) - 1) AS section" \
                   " FROM animals GROUP BY section")
    r = [int(row[0]) for row in cursor.fetchall()]
    r.sort()
    return r

def group_aggfeeds(groupid, fromdate, todate):
    '''Get [(date, min, max, avg),..] for group, aggregating
    from per-animal data.'''
    
    cursor.execute("SELECT date,MIN(grams),MAX(grams),SUM(grams)/COUNT(grams)" \
        " FROM feeds WHERE animalid IN (SELECT animalid FROM animals WHERE " \
        "groupid=? AND active=1) AND date>=? AND date<=? GROUP BY date " \
        "ORDER BY date",
        (groupid, date_todb(fromdate), date_todb(todate)))
    
    return [(date_fromdb(d), mi, ma, av) for d, mi, ma, av in cursor.fetchall()]

def section_aggfeeds(section, fromdate, todate):
    '''Get [(date, min, max, avg),..] for section, aggregating
    from per-animal data.'''
    
    cursor.execute("SELECT date,MIN(grams),MAX(grams),SUM(grams)/COUNT(grams)" \
        " FROM feeds WHERE animalid IN (SELECT animalid FROM animals WHERE " \
        "SUBSTR(groupid, 1, LENGTH(groupid) - 1)=? AND active=1) " \
        "AND date>=? AND date<=? GROUP BY date ORDER BY date",
        (str(section), date_todb(fromdate), date_todb(todate)))
    
    return [(date_fromdb(d), mi, ma, av) for d, mi, ma, av in cursor.fetchall()]


##################
# Animals

def animalcount():
    '''Get count of all active animals'''
    cursor.execute("SELECT COUNT(*) FROM animals WHERE active=1")
    return cursor.fetchone()[0]

def animalinfo(animalid):
    '''Get animal information. Returns groupid, active or None'''
    cursor.execute("SELECT groupid, active FROM animals WHERE animalid=?",
                   (animalid, ))
    row = cursor.fetchone()
    if row is None:
        return None
    else:
        groupid, active = row
        return groupid, bool(active)

def animalnotes(animalid):
    '''Get notes for animal or None'''
    cursor.execute("SELECT notes FROM animals WHERE animalid=?", (animalid,))
    row = cursor.fetchone()
    if row is None:
        return None
    else:
        return row[0]

def animals_in_group(groupid, activeonly = True):
    '''Get a sorted list of animalids in a group'''
    if activeonly:
        activewhere = "AND active=1 "
    else:
        activewhere = ""
    
    cursor.execute("SELECT animalid FROM animals WHERE groupid=? "
                   + activewhere + "ORDER BY animalid", (groupid,))
    return [row[0] for row in cursor.fetchall()]

def animals_in_section(section, activeonly = True):
    '''Get a sorted list of animalids in a section'''
    if activeonly:
        activewhere = "active=1 AND "
    else:
        activewhere = ""
    
    cursor.execute("SELECT animalid FROM animals WHERE " + activewhere
        + "SUBSTR(groupid, 1, LENGTH(groupid) - 1)=? ORDER BY animalid",
        (str(section), ))
    return [row[0] for row in cursor.fetchall()]

def animal_getfeeds(animalid, fromdate, todate):
    '''Get animal feeds in date order in a time interval
    (defined as datetime.dates)
    '''
    fromdate = date_todb(fromdate)
    todate = date_todb(todate)
    cursor.execute("SELECT date, grams FROM feeds WHERE animalid=? " \
        "AND date>=? AND date<=? ORDER BY date", (animalid, fromdate, todate))
    
    return [(date_fromdb(date), grams) for date, grams in cursor.fetchall()]

def animals_getmaxfeeds(animals, fromdate, todate):
    '''Get the maximum feed amount in timespan for animals in list.
    If animals is None, return maximum for all animals. If no feed
    data exists, return 0.
    '''
    fromdate = date_todb(fromdate)
    todate = date_todb(todate)
    
    if animals is not None:
        values = ','.join([str(int(aid)) for aid in animals])
        cursor.execute("SELECT MAX(grams) FROM feeds WHERE animalid IN" \
                " (%s) AND date>=? AND date<=?" % values, (fromdate, todate))
    else:
        cursor.execute("SELECT MAX(grams) FROM feeds WHERE " \
                "date>=? AND date<=?", (fromdate, todate))
    
    row = cursor.fetchone()
    if not row or not row[0]:
        return 0
    else:
        return row[0]

def get_inactiveanimals():
    '''Get list of all inactive animals'''
    cursor.execute("SELECT animalid FROM animals WHERE active=0")
    return [row[0] for row in cursor.fetchall()]

def getlastdate():
    '''Get last date that exists in database feeds or
    if no feed data exists datetime.date.today()'''
    cursor.execute("SELECT MAX(date) FROM feeds")
    row = cursor.fetchone()
    if not row or not row[0]:
        return datetime.date.today()
    else:
        return date_fromdb(row[0])

def remove_animal(animalid):
    '''Remove animal and any associated feed data'''
    cursor.execute("DELETE FROM feeds WHERE animalid=?", (animalid, ))
    cursor.execute("DELETE FROM animals WHERE animalid=?", (animalid, ))

##################
# Data conversion

def date_fromdb(string):
    '''Convert date from database string presentation to datetime.date'''
    year = int(string[0:4])
    month = int(string[4:6])
    day = int(string[6:8])
    return datetime.date(year, month, day)

def date_todb(date):
    '''Convert datetime.date to string representation for database'''
    return date.strftime("%Y%m%d")


###############
# Files

def filterfiles(filenames):
    '''Process the list of filenames. Return list of filenames that are not in
    the database.
    '''
    
    result = []
    for filename in filenames:
        cursor.execute("SELECT filename FROM files WHERE filename=?",
                       (filename, ))
        if not cursor.fetchone():
            result.append(filename)
    return result

def addfiles(filenames):
    '''Add filenames to database. Should only be called after the files have
    succesfully been processed.
    '''
    for filename in filenames:
        cursor.execute("INSERT INTO files VALUES (?)", (filename, ))


##############
# Groupfeed data

def groupfeed_getlast():
    '''Get date of last imported row asYYYYMMDD or None'''
    cursor.execute("SELECT date FROM groupfeed ORDER BY date DESC LIMIT 1")
    row = cursor.fetchone()
    if row:
        return date_fromdb(row[0])
    else:
        return None

# Pen to valve mapping. Keys are valves, and values are list of pens
pens_per_valve = {1: [1, 3], 2: [2, 4], 5: [5], 6: [6], 7: [7], 8: [8]}

def getvalve(groupid):
    '''Return valve number for groupid.'''
    section, pen = groupinfo(groupid)
    
    for valve, pens in pens_per_valve.items():
        if pen in pens:
            return section * 10 + valve
    
    raise Exception, "No valve for group %d" % groupid

def getgroupspervalve(valve):
    '''Get groupids that feed from valve'''
    section, valvepen = groupinfo(valve)
    pens = pens_per_valve.get(valvepen, [])
    return [section * 10 + pen for pen in pens]

def getvalvelist():
    '''Get list of all valves with animals'''
    cursor.execute("SELECT valve FROM groupfeed WHERE count>0 GROUP BY valve")
    return [row[0] for row in cursor.fetchall()]

def groupfeed_getfeeds(valve, fromdate, todate):
    '''Get groupfeed amounts per valve'''
    cursor.execute("SELECT date, grams FROM groupfeed WHERE valve=? AND " \
                   "date>=? AND date<=? ORDER BY date",
                   (valve, date_todb(fromdate), date_todb(todate)))
    return [(date_fromdb(d), g) for d, g in cursor.fetchall()]

def groupfeed_average(groups, fromdate, todate):
    '''Get groupfeed per-animal average for groups. Average may include
    other groups that are fed from same valve. Returns [(date, grams), ]
    If there are no animals in groupfeed group, grams is None.
    '''
    valves = set([getvalve(gid) for gid in groups]) # Unique valvelist
    vlist = ','.join([str(int(valve)) for valve in valves])
    cursor.execute("SELECT date, SUM(grams)/SUM(count) FROM groupfeed " \
        "WHERE valve IN (%s) AND date>=? AND date<=? GROUP BY date" % vlist
        , (date_todb(fromdate), date_todb(todate)))
    
    return [(date_fromdb(d), g) for d, g in cursor.fetchall()]

def groupfeed_animalcount(groups = None):
    '''Get number of animals in groupfeed data for groups. Count may include
    other groups. In groups is None, return count of all groupfeed animals.'''
    if groups is None:
        whereclause = ""
    else:
        valves = set([getvalve(gid) for gid in groups]) # Unique valvelist
        vlist = ','.join([str(int(valve)) for valve in valves])
        whereclause = "WHERE valve IN (%s) " % vlist
    
    cursor.execute("SELECT SUM(count) FROM groupfeed " + whereclause
                   + "GROUP BY date ORDER BY date DESC LIMIT 1")
    row = cursor.fetchone()
    if row:
        return row[0]
    else:
        return 0

##############
# Import log

def save_importlog(log, time = None):
    '''Save importlog to database. If time is not given, assume now.'''
    if not time:
        time = datetime.datetime.today()
    
    if log.strip() == '':
        return
    
    timespec = time.strftime("%Y-%m-%d %H:%M")
    cursor.execute("REPLACE INTO importlogs VALUES (?,?)", (timespec, log))

def cleanup_importlogs(time):
    '''Clean up old import logs. Keep ones newer than time.'''
    timespec = time.strftime("%Y-%m-%d %H:%M")
    cursor.execute("DELETE FROM importlogs WHERE time<=?", (timespec, ))

def get_importlogs():
    '''Get all importlogs as [(timespec, log), ]'''
    cursor.execute("SELECT time, log FROM importlogs ORDER BY time ASC")
    return cursor.fetchall()


##############
# Animal tree

from pymisc import TreeNode

def getanimaltree():
    '''Construct tree of animals.
    Tree levels are root->sections->groups->animals.
    Each node has node.userdata, like ('animal', 1234).'''
    root = TreeNode()
    
    for sectionid in sectionlist():
        sectnode = TreeNode(root, ('section', sectionid))
        
        for groupid in groups_in_section(sectionid):
            groupnode = TreeNode(sectnode, ('group', groupid))
            
            for animalid in animals_in_group(groupid):
                TreeNode(groupnode, ('animal', animalid))

    # Inactive animals are under separate category
    inactivenode = TreeNode(root, ('inactive',))
    for animalid in get_inactiveanimals():
        TreeNode(inactivenode, ('animal', animalid))

    return root


if __name__ == '__main__':
    assert date_fromdb("20070304") == datetime.date(2007, 3, 4)
    assert date_todb(datetime.date(2008,12,30)) == '20081230'