#!/usr/bin/python
# -*- coding: UTF-8 -*-
import os
import json
import simplejson
from datetime import datetime, timedelta
import numpy
from pymongo import MongoClient
# import mongoengine
from asgiref.sync import async_to_sync
from channels.generic.websocket import WebsocketConsumer

# from plotter.models import Experiment, ExpDetail, PlotMeta, PlotData

# Here we create the db named "dbplots"
host = os.environ.get('HOST_MONGO', 'localhost')
# mongoengine.connect('dbplots', host=host, port=27017)

CLIENT = MongoClient('{}:27017'.format(host))
DB = CLIENT['dbplots']

class MainConsumer(WebsocketConsumer):

    def connect(self):
        self.group_name = 'main'
        async_to_sync(self.channel_layer.group_add)(
            self.group_name,
            self.channel_name
        )
        self.accept()

    def disconnect(self, close_code):
        async_to_sync(self.channel_layer.group_discard)(
            self.group_name,
            self.channel_name
        )

    def receive(self, text_data):
        pass

    def zmq_message(self, event):
        # Send message to WebSocket
        self.send(text_data=event['message'])

class PlotConsumer(WebsocketConsumer):

    def connect(self):

        if 'realtime' in self.scope['path']:
            self.realtime = True
            self.group_name = '{}_{}'.format(
                self.scope['url_route']['kwargs']['code'],
                self.scope['url_route']['kwargs']['plot'],
            )

            async_to_sync(self.channel_layer.group_add)(
                self.group_name,
                self.channel_name
            )
        else:
            self.realtime = False
        self.accept()

    def disconnect(self, close_code):

        if self.realtime:
            async_to_sync(self.channel_layer.group_discard)(
                self.group_name,
                self.channel_name
            )

    def receive(self, text_data):
        ret = {}
        dt = datetime.strptime(text_data, '%d-%m-%Y')
        code = self.scope['url_route']['kwargs']['code']
        plot = self.scope['url_route']['kwargs']['plot']
        exp = DB.experiment.find_one({'code': int(code)})
        det0 = DB.exp_detail.find_one({
            'experiment': exp['_id'],
            'date': dt-timedelta(days=1)
            })
        det1 = DB.exp_detail.find_one({
            'experiment': exp['_id'],
            'date': dt
            })
        
        if det1:
            meta1 = DB.plot_meta.find_one({
                'exp_detail': det1['_id'],
                'plot': plot
                })
            if meta1:
                if meta1['metadata']['type'] in ('pcolor', 'image'):
                    datas = DB.plot_data.find(
                        {'plot': meta1['_id']}, 
                        ['time', 'data'],
                        sort=[('time', -1)], 
                        limit=1)[0]
                    ret['time'] = [datas['time']]
                    ret['data'] = datas['data']
                    ret['metadata'] = meta1['metadata']
                else:
                    last = det1['last_time']
                    metas = [meta1['_id']]
                    if det0:
                        meta0 = DB.plot_meta.find_one({
                            'exp_detail': det0['_id'],
                            'plot': plot
                            })
                        if meta0:
                            metas.append(meta0['_id'])
                    total = DB.plot_data.count_documents({
                        'plot': {'$in': metas}, 
                        'time': {'$gt': last-12*60*60}
                    })
                    skip = 0 if total-720<0 else total-720
                    datas = DB.plot_data.find({
                        'plot': {'$in': metas}, 
                        'time': {'$gt': last-12*60*60}
                        }, ['time', 'data'], sort=[('time', 1)], limit=720, skip=skip)

                    dum = [(d['time'], d['data']) for d in datas]
                    ret['time'] = [d[0] for d in dum]
                    dum = numpy.array([d[1] for d in dum])
                    
                    if len(dum[0][0])==1:
                        ret['data'] = dum.T[0].tolist()
                    else:
                        ret['data'] = [t for t in map(list, list(zip(*dum.tolist())))]
                    ret['metadata'] = meta1['metadata']
                    
            self.send(simplejson.dumps(ret, ignore_nan=True))
        else:
            self.send(json.dumps({'interval': 0}))

    def zmq_message(self, event):
        # Send message to WebSocket
        self.send(text_data=event['message'])