# -*- coding: utf-8 -*-

import numpy
from matplotlib import pyplot

import hark


def cartesian2polar(xyz):
    """
    cartesian2polar(x, y, z)
    converts coordinate from xyz to r\theta\phi
    """
    x, y, z = xyz
    r = numpy.sqrt(sum(map(lambda x: x**2, [x, y, z])))
    theta = numpy.arctan2(z, numpy.sqrt(sum(map(lambda x: x**2, [x, y]))))
    phi = numpy.arctan2(y, x)
    return (r, theta, phi)


def r2d(r):
    return 180.0 / numpy.pi * r


class plotSource(hark.node.PythonNode):
    def __init__(self):
        #super().__init__(['SOURCES', 'WINDOW_LENGTH', 'UPDATE_INTERVAL'], ['OUTPUT'])
        super().__init__('SOURCES', 'OUTPUT')
        
        pyplot.ion()
        self.plotx = {}
        self.ploty = {}
        self.winlen = 300
        self.count = 0
    
    def forward(self, **kwargs):
        if isinstance(kwargs['SOURCES'], list):
            for src in kwargs['SOURCES']:
                if isinstance(src, hark.harklib.Source):
                    self.plotx.setdefault(src.id, [])
                    self.ploty.setdefault(src.id, [])
                    self.plotx[src.id].append(self.count)
                    self.ploty[src.id].append(r2d(cartesian2polar(src.x)[2]))
        
        if self.count % 10 == 0:
            for i, srcid in enumerate(self.plotx.keys()):
                if i == 0:
                    pyplot.clf()
                pyplot.plot(self.plotx[srcid], self.ploty[srcid], "." + "rgb"[srcid % 3])
            pyplot.xlim([self.count-self.winlen, self.count])
            pyplot.ylim([-180, 180])
            pyplot.xlabel("Time [frame]")
            pyplot.ylabel("Azimuth [deg]")
            pyplot.pause(0.01)
        
        self.count += 1
        return {'OUTPUT': numpy.array([1.0, 1.0, 1.0], dtype=numpy.float32)}

#    def __del__(self):
#        pyplot.close()

