import numpy as np
from ...util.math import on_finite, linspc_, logspc_, linrng_, logrng_, flattengrid
from ...util.plot import bin_centres_to_edges
from ...util.distributions import sample_kde
import logging

logging.getLogger(__name__).addHandler(logging.NullHandler())
logger = logging.getLogger(__name__)


class DensityGrid(object):
    def __init__(
        self, x, y, extent=None, bins=50, logx=False, logy=False, coverage_scale=1.2
    ):
        # limits
        self.logx = logx
        self.logy = logy
        self.bins = bins
        self.coverage_scale = coverage_scale

        if extent is None:
            self.xmin, self.xmax, self.ymin, self.ymax = self.extent_from_xy(x, y)
        else:
            self.xmin, self.xmax, self.ymin, self.ymax = extent

        self.xstep = self.get_xstep()
        self.ystep = self.get_ystep()

        if self.logx:
            assert self.xmin > 0.0
            assert (self.xmin / self.xstep) > 0.0
        if self.logy:
            assert self.ymin > 0.0
            assert (self.ymin / self.ystep) > 0.0

        self.calculate_grid()

    def calculate_grid(self):
        self.get_centre_grid()
        self.get_edge_grid()

    def get_ystep(self):
        if self.logy:
            return (self.ymax / self.ymin) / self.bins
        else:
            return (self.ymax - self.ymin) / self.bins

    def get_xstep(self):
        if self.logx:
            return (self.xmax / self.xmin) / self.bins
        else:
            return (self.xmax - self.xmin) / self.bins

    def extent_from_xy(self, x, y, coverage_scale=None):
        cov = coverage_scale or self.coverage_scale
        expand_grid = (cov - 1.0) / 2
        return [
            *[linrng_, logrng_][self.logx](x, exp=expand_grid),
            *[linrng_, logrng_][self.logy](y, exp=expand_grid),
        ]

    def get_xrange(self):
        return self.xmin, self.xmax

    def get_yrange(self):
        return self.ymin, self.ymax

    def get_extent(self):
        return [*self.get_xrange(), *self.get_yrange()]

    def get_range(self):
        return [[*self.get_xrange()], [*self.get_yrange()]]

    def update_grid_centre_ticks(self):
        if self.logx:
            self.grid_xc = logspc_(self.xmin, self.xmax, 1.0, self.bins)
        else:
            self.grid_xc = linspc_(self.xmin, self.xmax, self.xstep, self.bins)

        if self.logy:
            self.grid_yc = logspc_(self.ymin, self.ymax, 1.0, self.bins)
        else:
            self.grid_yc = linspc_(self.ymin, self.ymax, self.ystep, self.bins)

    def update_grid_edge_ticks(self):
        self.update_grid_centre_ticks()
        if self.logx:
            self.grid_xe = np.exp(bin_centres_to_edges(np.log(np.sort(self.grid_xc))))
        else:
            self.grid_xe = bin_centres_to_edges(np.sort(self.grid_xc))

        if self.logy:
            self.grid_ye = np.exp(bin_centres_to_edges(np.log(np.sort(self.grid_yc))))
        else:
            self.grid_ye = bin_centres_to_edges(np.sort(self.grid_yc))

    def get_centre_grid(self):
        self.update_grid_centre_ticks()
        self.grid_xci, self.grid_yci = np.meshgrid(self.grid_xc, self.grid_yc)

    def get_edge_grid(self):
        self.update_grid_edge_ticks()
        self.grid_xei, self.grid_yei = np.meshgrid(self.grid_xe, self.grid_ye)

    def get_hex_extent(self):
        if self.logx:
            xex = [np.log(self.xmin / self.xstep), np.log(self.xmax * self.xstep)]
        else:
            xex = [self.xmin - self.xstep, self.xmax + self.xstep]

        if self.logy:
            yex = [np.log(self.ymin / self.ystep), np.log(self.ymax * self.ystep)]
        else:
            yex = [self.ymin - self.ystep, self.ymax + self.ystep]
        return xex + yex

    def kdefrom(self, xy, xtransform=lambda x: x, ytransform=lambda x: x):
        assert np.isfinite(self.grid_xc).all() and np.isfinite(self.grid_yc).all()
        arr = xy.copy()
        # generate x grid over range spanned by log(x)
        arr[:, 0] = xtransform(xy[:, 0])
        # generate y grid over range spanned by log(y)
        arr[:, 1] = ytransform(xy[:, 1])
        zi = sample_kde(
            arr,
            flattengrid(
                np.meshgrid(xtransform(self.grid_xc), ytransform(self.grid_yc))
            ),
        )
        zi = zi.reshape(self.grid_xci.shape)
        return zi
