# coding: utf-8
"""
    Galaxy kiloparsec plotter
"""
from astropy.table import QTable
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np

def plot_radial(table: QTable):

    fig, ax = plt.subplots(subplot_kw=dict(projection="polar"), figsize=(6, 5))

    rbins = np.unique(table["r"])
    delta_r = rbins[1] - rbins[0]

    for r in rbins:

        mask = table["r"] == r
        phibins = table["phi"][mask].value

        C = np.log10(table["mass"][mask].value)
        N = len(C)
        # C = np.arange(N) + r.value
        # np.random.shuffle(C)
        C = C.reshape(1, N)
        P = np.zeros(shape=(2, N + 1))

        R = np.ones(shape=(2, N + 1)) * (r + delta_r / 2.0)
        R[0, :] -= delta_r
        R = R.value

        deltaphi = phibins[1] - phibins[0]
        P[0, :-1] = phibins - deltaphi / 2
        P[1, :-1] = phibins - deltaphi / 2
        P[0, -1] = P[0, 0]
        P[1, -1] = P[1, 0]

        pc = ax.pcolormesh(P, R, C, cmap="inferno")  # , vmin=6, vmax=9)
        print(P, R, C)
    fig.colorbar(pc)
    
def plot_hist2d(table, x, y):
    plt.figure()
    plt.hist2d(np.log10(table[x].value) - 6, np.log10(table[y].value), 
               bins=100, cmap="magma", weights=table["mass"], norm=mpl.colors.LogNorm())