#!/usr/bin/env python
# coding: utf-8

import numpy as np

from . import sball # for Isocurves

def isolevels_2d(pdf, weighting_pdf_const, r_levels, from_top=None):
    """
    weighting_pdf_const = 1 / (xmax - xmin) / (ymax - ymin)
    """
    s_ball = sball.Sball(2) # nvar=2
    p_levels = []
    for r in r_levels:
        p_levels.append(1 - s_ball.get_pf(r))
    
    return isolevels(pdf, weighting_pdf_const, p_levels, from_top)




def isolevels(pdf, weighting_pdf_const, p_levels, from_top=None):
    """
    weighting_pdf_const = 1 / (xmax - xmin) / (ymax - ymin)
    """
    #č třeba P prostor doopravdy zlobí, takže zkusím nějak tak
    if from_top is None:
        weights = pdf / weighting_pdf_const
        p_all = np.sum(weights) / len(pdf)
        #č prečo víme, že celková pravděpodobnost může bejt nekoněčně velká
        if p_all <= 1:
            from_top = True
        else:
            from_top = False
            
    max_pdf = np.max(pdf)
    pdf_levels = []
    if from_top:
        # descending
        sorted_pdf = np.flip(np.sort(pdf))
        p_cumsum = np.cumsum(sorted_pdf) / weighting_pdf_const / len(pdf)
        for p in p_levels:
            # little bit tricky, didn't find numpy method for this
            mask = p_cumsum <= p
            level_down_bound = np.max(sorted_pdf[~mask], initial=0)
            level_up_bound = np.min(sorted_pdf[mask], initial=max_pdf)
            pdf_levels.append((level_down_bound + level_up_bound) / 2)
            
    else: # from bottom
        sorted_pdf = np.sort(pdf)
        p_cumsum = np.cumsum(sorted_pdf) / weighting_pdf_const / len(pdf)
        for p in p_levels:
            # little bit tricky, didn't find numpy method for this
            mask = p_cumsum <= 1-p
            level_down_bound = np.max(sorted_pdf[mask], initial=0)
            level_up_bound = np.min(sorted_pdf[~mask], initial=max_pdf)
            pdf_levels.append((level_down_bound + level_up_bound) / 2)

    return pdf_levels