from peewee import *
from playhouse.sqlite_ext import SqliteExtDatabase
import datetime
import json
import numpy as np

def coloursToJson(colours):
    colours2 = [(list(colour[0]), colour[1]) for colour in colours]
    return json.dumps(colours2)

def jsonToColours(string):
    data = json.loads(string)
    return [(np.array(d[0]), d[1]) for d in data]

db = SqliteExtDatabase('images.db')

class ColoursField(TextField):
#     db_field = 'colour'

    def db_value(self, value):
        return coloursToJson(value)

    def python_value(self, value):
        return jsonToColours(value) # convert str to UUID


class BaseModel(Model):
    class Meta:
        database = db
        
class Emotion(BaseModel):
    name = CharField(unique=True)
        
class Group(BaseModel):
    name = CharField(unique=True)

class Artwork(BaseModel):
    author = CharField()
    age = SmallIntegerField(index=True)
    gender = FixedCharField(max_length=1) # we should not really use this one
    group = ForeignKeyField(Group, related_name='artworks', index=True)
    emotion = ForeignKeyField(Emotion, related_name='artworks', index=True)
    created_date = DateTimeField(default=datetime.datetime.now)
    filename = CharField()
    colours = ColoursField() # serialised colours + percentages: [([r,g,b], percentage), ...]

def getAges():
    r = Artwork.select(fn.Distinct(Artwork.age)).dicts()
    return [a['age'] for a in r]

def getEmotionCountsFromArtworks(artworks):
    emotions = {}
    for artwork in artworks:
        e = artwork.emotion.name
        if e in emotions:
            emotions[e] += 1
        else:
            emotions[e] = 1
    return emotions

def getEmotionFractionsFromArtworks(artworks):
    emotions = getEmotionCountsFromArtworks(artworks)
    total = sum(emotions.values())
    for e in emotions:
        emotions[e] = emotions[e] / total
    return emotions

def getAgeFractionsFromArtworks(artworks):
    ages = {}
    for artwork in artworks:
        age = artwork.age
        if age in ages:
            ages[age] += 1
        else:
            ages[age] = 1
    
    total = sum(ages.values())
    for age in ages:
        ages[age] = ages[age] / total
    return ages