#!/usr/bin/env python

import math
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.spatial import ConvexHull
from colour import Color
from matplotlib.patches import Polygon
import statistics as st

from granatum_sdk import Granatum

COLORS = ["#3891ea", "#29ad19", "#ac2d58", "#db7580", "#ed2310", "#ca2dc2", "#5f7575", "#7cc1b5", "#c3bd78", "#4ffa24"]

def main():
    gn = Granatum()
    sample_coords = gn.get_import("viz_data")
    value = gn.get_import("value")
#    print(value)
    coloring_type = gn.get_arg("coloring_type")
    bounding_stdev = gn.get_arg("bounding_stdev")

    coords = sample_coords.get("coords")
    dim_names = sample_coords.get("dimNames")

    df = pd.DataFrame(
        {"x": [a[0] for a in coords.values()], "y": [a[1] for a in coords.values()], "value": pd.Series(value)},
        index=coords.keys()
    )

#    print(df)

    if coloring_type == "categorical":
        uniq = df["value"].unique();
        num = uniq.shape[0]
        COLORS2 = plt.get_cmap('gist_rainbow')
        carr = [0]*df.shape[0]
        listcats = list(df["value"])     
        miny = min(list(df["y"]))
        maxy = max(list(df["y"]))
        scaley = (maxy-miny)/650
        print("Scaley = {}".format(scaley))

        for i, cat in enumerate(df["value"].unique()):
            dff = df[df["value"] == cat]
            xs = list(dff["x"])
            ys = list(dff["y"])
            #avgx = sum(dff["x"]) / len(dff["x"]) 
            #avgy = sum(dff["y"]) / len(dff["y"]) 
            #plt.scatter(x=dff["x"], y=dff["y"], s=5000 / df.shape[0], c=COLORS[i].hex_l, label=cat)
            #plt.scatter(x=dff["x"], y=dff["y"], s=5000 / df.shape[0], c=[abs(hash(cat)) % 256]*len(dff["x"]), cmap=COLORS2, label=cat)
            #plt.scatter(x=dff["x"], y=dff["y"], s=5000 / df.shape[0], c=abs(hash(cat)) % 256, cmap=COLORS2, label=cat)
            colorindex = abs(hash(cat)) % 256
            craw = COLORS2(colorindex/255.0)
            color = (craw[0], craw[1], craw[2], 0.2)
            whitetransparent = (1, 1, 1, 0.5)
            coloropaque = (craw[0], craw[1], craw[2], 1.0)
            if len(xs)>3:
                pts = list(zip(xs, ys))
                cent = np.mean(pts, axis=0)
                lengs = list(map(lambda p: math.sqrt((p[0]-cent[0])*(p[0]-cent[0])+(p[1]-cent[1])*(p[1]-cent[1])), pts))
                avgleng = st.mean(lengs)
                stdleng = st.stdev(lengs)*bounding_stdev
                rpts = []
                if(stdleng > 0.0):
                    for j, ln in enumerate(lengs):
                        if(ln - avgleng < stdleng):
                            rpts.append(pts[j])
                pts = rpts
                cent = np.mean(pts, axis=0)
                hull = ConvexHull(pts)
                ptslist = []
                for pt in hull.simplices:
                    ptslist.append(pts[pt[0]])
                    ptslist.append(pts[pt[1]])
                ptslist.sort(key=lambda p: np.arctan2(p[1]-cent[1], p[0]-cent[0]))
                ptslist = ptslist[0::2]
                ptslist.insert(len(ptslist), ptslist[0])
                lowestpt = ptslist[0]
                for pt in ptslist:
                    if(pt[1] < lowestpt[1]):
                        lowestpt = pt
                poly = Polygon(1.1*(np.array(ptslist)-cent)+cent, facecolor=color)
                poly.set_capstyle('round')
                plt.gca().add_patch(poly)
                plt.text(lowestpt[0], lowestpt[1]-scaley*10, cat, fontsize=6, ha="center", va="center", color="black", bbox=dict(boxstyle="round",fc=whitetransparent,ec=coloropaque))
            for j,x in enumerate(listcats):
                if x == cat:
                    carr[j] = int(abs(hash(cat)) % 256)
        
        plt.scatter(x=df["x"], y=df["y"], s=5000 / df.shape[0], c=carr, cmap=COLORS2)
        lgd = plt.legend(markerscale=6, loc='upper center', bbox_to_anchor=(0.5, -0.05), ncol=5)
#60 / (5000 / df.shape[0])
    elif coloring_type == "continuous":
        plt.scatter(x=df["x"], y=df["y"], s=5000 / df.shape[0], c=df["value"], cmap="Reds")
        plt.colorbar()

    plt.xlabel(dim_names[0])
    plt.ylabel(dim_names[1])
    # plt.tight_layout()

    gn.add_current_figure_to_results(
        "Scatter-plot",
        dpi=75,
        width=750,
        height=650,
#        savefig_kwargs={'bbox_extra_artists': (lgd,), 'bbox_inches': 'tight'}
        savefig_kwargs={'bbox_inches': 'tight'}
    )

    gn.commit()


if __name__ == "__main__":
    main()

matplotlib.pyplot.colorbar matplotlib.pyplot.xlabel matplotlib.pyplot.get_cmap matplotlib.pyplot.legend matplotlib.pyplot.ylabel matplotlib.pyplot.scatter matplotlib.pyplot.gca