# -------------------------------------------------------------------------------
# Generated by PixieDust code generator
# 
# Licensed under the Apache License, Version 2.0 (the 'License');
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
# http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an 'AS IS' BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Inherited from maven-artifact https://github.com/hamnis/maven-artifact
# -------------------------------------------------------------------------------

from pixiedust.display.display import Display
from pixiedust.utils import Logger

from wordcloud import WordCloud
import cStringIO
import base64

@Logger()
class WordCloudDisplayHandler(Display):
    """
        Main Render method
    """
    def doRender(self,handlerId):
        # convert from dataframe to dict
        dfdict = {}
        df = self.entity.toPandas()
        for x in range(len(df)):
            currentid = df.iloc[x,0] or 'NoStreet'
            currentvalue = df.iloc[x,1]
            dfdict.setdefault(currentid, 0)
            dfdict[currentid] = dfdict[currentid] + currentvalue

        # create word cloud from dict
        wc = WordCloud(background_color="white").fit_words(dfdict)

        # encode word cloud image to base64 string
        b = cStringIO.StringIO()
        wc.to_image().save(b, format="PNG")
        img_str = base64.b64encode(b.getvalue())

        self._addHTMLTemplate("wordcloud.html", base64image=img_str.decode("ascii"))
