import pandas as pd
import vaex as vx
from PIL import Image, ImageDraw, ImageFont
from skimage import exposure
from skimage.util import img_as_ubyte
import numpy as np
from matplotlib import cm
import sys

base_size = 1080, 1920


def scale_to(width=None, height=None):
    isnone = (width is None, height is None)
    ret = {
        (False, False): lambda w, h: (w, h),
        (True, True): lambda w, h: (width, height),
        (False, True): lambda w, h: (width, width * (h / w)),
        (True, False): lambda w, h: (height * (w / h), height),
    }
    return lambda *args: tuple(map(int, ret[isnone](*args)))


# xz -1 1
bining = {
    ("zx", -1, 1): scale_to(width=base_size[0]),  # main view, top down
    # ('yx',1,1): lambda size,w,h: (size,int(size*(w/h))), #
    # ('zy',-1,1): lambda size,w,h: (int(size*(h/w)),size), #
}

print("Loading stars.csv")
stars = pd.read_csv(
    "stars.csv",
    names=["id", "name", "num_bodies", "has_scoopable", "mult", "x", "y", "z"],
    usecols=["id", "num_bodies", "x", "y", "z", "mult"],
    index_col=0,
)
stars = vx.from_pandas(stars, copy_index=False)

filename = "heuristic.png"
fnt = ImageFont.truetype(r"FiraCode-Regular", 40)

for (binby_key, m1, m2), calcshape in bining.items():
    binby = [m1 * stars[binby_key[0]], m2 * stars[binby_key[1]]]

    mm = [binby[0].minmax(), binby[1].minmax()]

    w, h = [mm[0][1] - mm[0][0], mm[1][1] - mm[1][0]]
    shape = calcshape(w, h)
    hm_all = stars.sum("num_bodies", binby=binby, shape=shape, limits="minmax")
    hm_all_mask = hm_all != 0
    hm_all = exposure.equalize_hist(hm_all)
    hm_all -= hm_all.min()
    hm_all /= hm_all.max()

    hm_boost = stars.sum(
        "astype(mult>1.0,'int')", binby=binby, shape=shape, limits="minmax"
    )
    hm_boost_mask = hm_boost != 0

    hm_boost = exposure.equalize_hist(hm_boost)
    hm_boost -= hm_boost.min()
    hm_boost /= hm_boost.max()

    # R = cm.Reds_r()
    G = cm.Greens_r(hm_all)
    B = cm.Blues_r(hm_boost)

    img = np.zeros((base_size[0], base_size[1], 4))
    img[:, :, :] = 0.0
    img[:, :, 3] = 1.0
    canvas = img[: shape[0], : shape[1], :]
    canvas[hm_all_mask] = G[hm_all_mask]
    canvas[hm_boost_mask] = B[hm_boost_mask]
    pil_img = Image.fromarray(img_as_ubyte(img))
    draw = ImageDraw.Draw(pil_img)
    messages = ["Hello World"]
    draw.multiline_text((shape[0], 0), "\n".join(messages), font=fnt)
    pil_img.save(filename)