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

# Axes level functions (for Matplotlib)

#č Založme nový modul pro funkce (procedury),
#č které sice stejně jako mart pracují na urovni axes,
#č ale na rozdil od atomických funkcí mart modulu
#č udělaj z prazdných os hotový obrazek.
#č mplot.show2D() může použivat tuhle nabídku obrázků.
#
#č funkce v tomto modulu dostávájí jako parameter jedině ax
#č ax ale má nastavené atributy .space a .sample_box

import numpy as np
from . import mart
from . import mgraph

from matplotlib import colors as mcolors


# it is mostly for qt_plot, it offers availiable options to user
__all__ =   [
            'candidates_plot', 'candidates_sampling_plot',
            'convex_hull_plot', 'tri_plot', 'tri_nodes_plot',
            'tri_R_plot', 'tri_GK_plot',
            'tri_R_nodes_plot', 'tri_GK_nodes_plot',
            'convergence_diagram', 'convergence_legend', 
            'just_points', 'just_points_really',
            'base_drawing', 'candidates_drawing',
            'just_qhull', 'qhull_plot', 'qhull_infinite',
            'dhull_scheme_plot', 'dhull_random_plot',
            'bhull_plot', 'bhull_infinite',
            'shull_plot',
            'completehull_plot'
            ]





def candidates_sampling_plot(ax, linewidths=[0.7, 0.5, 0.4, 0.3, 0.2, 0.1]):
    from .. import simplex as six
    from .. import estimation as stm
    
    def _stm_draw_nodes(*args, **kwargs):
        #č v tomto plot bude šilený bordel v tom, kdo kolbek spouští
        #č a kdo mu co tam posílá. Na event a na nodes spolehat ale můžeme
        #
        # callback's signature: sx, indices=, simplex=, nodes=, cell_stats=
        # positional "sx" is Tri object itself
        # "indices" are numbers of simplex vertices
        # "simplex" are vertices itself
        # "nodes" is what we really want to draw
        
        event = kwargs['cell_stats']['event']
        
        if event in ('mix', 'outside'):
            mart.plot_sample(ax, kwargs['nodes'], ls='', marker='.',\
                         mec="#00007E", mfc="#00007E", ms=1.5, alpha=0.5,\
                         rasterized=True)
        
    data = stm.fast_simplex_estimation(ax.sample_box, model_space=ax.space,\
                                  sampling_space=ax.space, \
                                   weighting_space=ax.space,\
                                    outside_budget=1000, \
                                     simplex_budget=100,\
                                    callback=_stm_draw_nodes, design=None)    
    
    mart.setup(ax)
    mart.curly(ax, linewidths=linewidths)
    mart.triplot(ax, color="#B2B2B2", lw=0.5, zorder=1000)
    mart.plot_points(ax, ms=5, zorder=100500)
    try:
        mart.plot_boundaries(ax, lw=0.7, zorder=10500)
    except:
        pass

def candidates_plot(ax):
    tri_plot(ax, tri_space=None, linewidths=[0.7, 0.5, 0.4, 0.3], data_offset=0.4)
    blue_colors = ["#BDBDFF", "#9999FF", "#00007E", "#000057"]
    blue_cmap = mcolors.LinearSegmentedColormap.from_list("bluecmap", blue_colors)
    mart.scatter_candidates(ax, s=1.5, marker='.', cmap=blue_cmap, rasterized=True)
    mart.plot_the_best_candidate(ax, "^", color='#000057')
    mart.plot_points(ax, ms=5, zorder=1005000)
    
#    if ax.space in ('G', 'GK'):
#        ax.xaxis.set_label_coords(1, 0.47)
#        ax.set_xlabel("$x_1$")
#        ax.text(0.47, 1, '$x_2$', ha='right',va='top', transform=ax.transAxes)
#    else:
#        ax.set_xlabel("$x_1$")
#        ax.set_ylabel("$x_2$")

def convergence_diagram(ax, sources=['box', 'user'], apply_proxy=False):
    #č pokorně jedeme použiť guessbox
    #č nic jiného nebylo pořádně implementováno
    from .. import stm_df 
    df = stm_df.get_tri_data_frame(ax.sample_box, sources, apply_proxy)
    try:
        pf_exact = ax.sample_box.pf_exact
        pf_exact_method = ax.sample_box.pf_exact_method
        mgraph.tri_estimation_plot(ax, df, pf_exact=pf_exact, \
                            pf_exact_method=pf_exact_method, plot_outside=True)
    except:
        mgraph.tri_estimation_plot(ax, df, plot_outside=True)
    ax.margins(0)
    ax.set_yscale('log')
    ax.set_xlabel("Number of points")
    ax.set_ylabel("Probability measure")
    
    
def convergence_legend(ax):
    convergence_diagram(ax)
    ax.legend(bbox_to_anchor=(0.5, -0.15), ncol=2, loc='upper center')

#č ten [plot] zásadně vytvaří své struktury, nepouzívá oné ze skříňky,
#č protože já vím, že v těch obrázcích, ve kterých chcu ho použit,
#č můde být třeba použit řez a skříňka tedy potřebné struktury může nemít
def tri_nodes_plot(ax, tri_space=None, tn_scheme=None,\
                     linewidths=[0.7, 0.5, 0.4, 0.3, 0.2, 0.1]):
    from .. import simplex as six
    if tri_space is None:
        tri_space = ax.space
    
    #č já tuhle funkciju potřebuju ne abych kreslil bodíky ve vrcholech
    if tn_scheme is None:
        try:
            tn_scheme = ax.sample_box.Tri.tn_scheme
        except:
            import quadpy
            tn_scheme = quadpy.tn.grundmann_moeller(sample_box.nvar, 5)
            
    
    def _draw_nodes(*args, **kwargs):
        # callback's signature: sx, indices=, simplex=, nodes=, cell_stats=
        # positional "sx" is Tri object itself
        # "indices" are numbers of simplex vertices
        # "simplex" are vertices itself
        # "nodes" is what we really want to draw
        
        event = kwargs['cell_stats']['event']
        
        if event == 'mix':
            color = '#FFF39A' #'xkcd: dark cream' # (255, 243, 154, 255)
        elif event == 'failure':
            color = '#fdc1c5' #'xkcd: pale rose' # (#fdc1c5)
        elif event == 'success':
            color = '#a7ffb5' #'xkcd:light seafoam green' #a7ffb5
        else:
            assert 100500 < 0 #оӵ мар лэсьтӥське татын?
            
        mart.plot_sample(ax, kwargs['nodes'], ls='', marker='.',\
                         mew=0, mfc=color, ms=2, alpha=0.5, rasterized=True)
        
        
        
    
    #č vytvařím vlastní tringulaciju zde vécemeně kvůli callbackům
    #č jde, samozřejmě všecko udělat jínak, 
    #č ale nechcu zrovna zde z toho robiť vědu
    Tri = six.JustCubatureTriangulation(ax.sample_box, tn_scheme=tn_scheme, \
                                    tri_space=tri_space, issi=None,\
                                    weighting_space=None, \
                                    incremental=False,\
                                    on_add_simplex=_draw_nodes,\
                                    on_delete_simplex=None)
    
    mart.setup(ax)
    mart.curly(ax, linewidths=linewidths)
    Tri.integrate()
    if tri_space == ax.space:
        mart.triplot(ax, color="#B2B2B2", lw=0.5, zorder=100)
    else:
        mart.tri_plot(ax, Tri=Tri, color="#B2B2B2", lw=0.5, zorder=100)
    
    mart.plot_points(ax, ms=2.5, zorder=100500)
    try:
        mart.plot_boundaries(ax, lw=0.7, zorder=1050)
    except:
        pass
    mart.setup_labels(ax)

def tri_R_nodes_plot(ax, **kwargs):
    tri_nodes_plot(ax, tri_space='R', **kwargs)
    
def tri_GK_nodes_plot(ax, **kwargs):
    tri_nodes_plot(ax, tri_space='GK', **kwargs)


#č ten [plot] zásadně vytvaří své struktury, nepouzívá oné ze skříňky,
#č protože já vím, že v těch obrázcích, ve kterých chcu ho použit,
#č můde být třeba použit řez a skříňka tedy potřebné struktury může nemít
def tri_plot(ax, tri_space=None, linewidths=[0.7, 0.5, 0.4, 0.3, 0.2, 0.1], data_offset=0.5):
    from .. import simplex as six
    if tri_space is None:
        tri_space = ax.space
    
    Tri = six.JustCubatureTriangulation(ax.sample_box, tn_scheme=None, \
                                    tri_space=tri_space, issi=None,\
                                    weighting_space=None, \
                                    incremental=False,\
                                    on_add_simplex=None,\
                                    on_delete_simplex=None)
    
    mart.setup(ax)
    mart.curly(ax, linewidths=linewidths)
    if tri_space == ax.space:
        mart.triplot(ax, color="#B2B2B2", lw=0.5, zorder=100)
    else:
        mart.tri_plot(ax, Tri=Tri, color="#B2B2B2", lw=0.5, zorder=100)
    
    mart.plot_points(ax, ms=2.5, zorder=100500)
    try:
        mart.plot_boundaries(ax, lw=0.7, zorder=1050)
    except:
        pass
    mart.setup_labels(ax, data_offset)
        
    

def tri_R_plot(ax, **kwargs):
    tri_plot(ax, tri_space='R', **kwargs)
    
def tri_GK_plot(ax, **kwargs):
    tri_plot(ax, tri_space='GK', **kwargs)

#č ten [plot] zásadně vytvaří svou obálku, nepouzívá onou ze skříňky,
#č protože já vím, že v těch obrázcích, ve kterých chcu ho použit,
#č můde být třeba použit řez a skříňka tedy potřebné struktury může nemít
def convex_hull_plot(ax, tri_space=None, linewidths=[0.7, 0.5, 0.4, 0.3, 0.2, 0.1]):
    from .. import convex_hull as khull
    if tri_space is None:
        tri_space = ax.space
    
    mart.setup(ax)
    mart.curly(ax, linewidths=linewidths)
    qhull = khull.QHull(ax.sample_box, space=tri_space, incremental=False)
    mart.qhull_plot(ax, qhull, color="#B2B2B2", lw=0.7, zorder=100)
    mart.plot_points(ax, ms=2.5, zorder=100500)
    try:
        mart.plot_boundaries(ax, lw=0.7, zorder=1050)
    except:
        pass
    mart.setup_labels(ax)

def just_points(ax):
    ax.set_xlabel('$x_{1}$')
    ax.set_ylabel('$x_{2}$')
    
    ax.set_aspect(1)
    #ax.set_box_aspect(1)
    mart.scatter_points(ax)
    
def just_points_really(ax):
    ax.set_aspect(1)
    ax.set_frame_on(False)
    mart.scatter_points(ax)

def base_drawing(ax):
    mart.setup(ax)
    mart.curly(ax)
    try:
        mart.triplot(ax, color="grey", linewidth=0.4)
    except:
        pass
    mart.plot_boundaries(ax, linewidth=0.7)
    mart.scatter_points(ax)
    
    
def candidates_drawing(ax):
    try:
        mart.triplot(ax, color="grey", linewidth=0.4)
    except:
        pass
    mart.plot_boundaries(ax, linewidth=0.7)
    
    #č ax.scatter posílá parameter cmap Collections třídě.
    #č Třída mimo jiného dědí cm.ScalarMappable,
    #č která inicializaci deleguje funkci cm.get_cmap() ve (svém) modulu cm.
    # cmap='viridis_r' #cmap='plasma',
    mart.scatter_candidates(ax, s=5, marker='.', cmap='plasma_r',\
     alpha=None, linewidths=None, edgecolors=None, plotnonfinite=False, rasterized=True)
    mart.plot_the_best_candidate(ax, "^", color='#3D0D5B')
    
        
    
    

# defaults
hezkymodře = (85/255, 70/255, 1, 1)
inside_color = [(185/255, 228/255, 14/255, 1)]
outside_color = [(128/255, 128/255, 128/255, 0.5)]
hull_colors = (hezkymodře, inside_color, outside_color) # border_inside_outside

def just_qhull(ax, lw=1.5, **kwargs):
    from .. import convex_hull as khull
    qhull = khull.QHull(ax.sample_box, space=ax.space, incremental=False)
    border_color, _, __ = hull_colors

    # setup
    ax.set_aspect(1)
    #ax.set_box_aspect(3/4)
    #ax.set_xlim(-lim, lim)
    #ax.set_ylim(-lim, lim)
    ax.set_xlabel('$x_1$')
    ax.set_ylabel('$x_2$')
    
    mart.qhull_plot(ax, qhull, color=border_color, lw=lw, zorder=100, **kwargs)
    
    # finally samples
    #
    #mart.scatter_sample(ax, f, c='g', marker='P', zorder=1000) # why?
    mart.scatter_points(ax, zorder=100500)



def _hull_model_plot(ax, hull, hull_plot, ns=50000, lim=3, lw=1.5, s=4, hull_colors=hull_colors, **kwargs):
    border_color, inside_color, outside_color = hull_colors

    # setup
    ax.set_aspect(1)
    #ax.set_box_aspect(3/4)
    ax.set_xlim(-lim, lim)
    ax.set_ylim(-lim, lim)
    ax.set_xlabel('$x_1$')
    ax.set_ylabel('$x_2$')
    
    nodes = ax.sample_box.f_model(ns)
    mask = hull.is_outside(nodes)
    mart.scatter_sample(ax, nodes[mask], c=outside_color, marker='.', s=s)
    mart.scatter_sample(ax, nodes[~mask], c=inside_color, marker='.', s=s)
    hull_plot(ax, hull, color=border_color, lw=lw, zorder=100, **kwargs)
    
    # finally samples
    #
    #mart.scatter_sample(ax, f, c='g', marker='P', zorder=1000) # why?
    mart.scatter_points(ax, zorder=100500)
    

    
    
    
def qhull_plot(ax, **kwargs):
    from .. import convex_hull as khull
    qhull = khull.QHull(ax.sample_box, space=ax.space, incremental=False)
    _hull_model_plot(ax, qhull, mart.qhull_plot, **kwargs)
    
    #mart.qhull_polygon(ax, qhull, fc="white", ec="black", lw=0.75)


def qhull_infinite(ax, **kwargs):
    from .. import convex_hull as khull
    qhull = khull.QHull(ax.sample_box, space=ax.space, incremental=False)
    _hull_model_plot(ax, qhull, mart.dhull_plot, **kwargs)




def dhull_scheme_plot(ax, **kwargs):
    from .. import convex_hull as khull
    import quadpy
    
    scheme = quadpy.un.stroud_un_7_1(2)
    dhull = khull.DirectHull(ax.sample_box, scheme.points, space=ax.space)
    _hull_model_plot(ax, dhull, mart.dhull_plot, **kwargs)
    

def dhull_random_plot(ax, ndir=8, rand_dir=None, **kwargs):
    from .. import convex_hull as khull
    from .. import sball
    
    if rand_dir is None:
        rand_dir = sball.get_random_directions(ndir, 2)
    dhull = khull.DirectHull(ax.sample_box, rand_dir, space=ax.space)
    _hull_model_plot(ax, dhull, mart.dhull_plot, **kwargs)  
    
    
def bhull_plot(ax, **kwargs):
    from .. import convex_hull as khull
    
    bhull = khull.BrickHull(ax.sample_box, space=ax.space)
    _hull_model_plot(ax, bhull, mart.bhull_plot, **kwargs)
    
def bhull_infinite(ax, **kwargs):
    from .. import convex_hull as khull
    
    bhull = khull.BrickHull(ax.sample_box, space=ax.space)
    _hull_model_plot(ax, bhull, mart.dhull_plot, **kwargs)    
    
    
def shull_plot(ax, **kwargs):
    from .. import convex_hull as khull
    
    shull = khull.GBall(ax.sample_box)
    _hull_model_plot(ax, shull, mart.shull_plot, **kwargs)
    
    

def completehull_plot(ax, ndir=3, rand_dir=None, **kwargs):
    from .. import convex_hull as khull
    from .. import sball
    
    if rand_dir is None:
        rand_dir = sball.get_random_directions(ndir, 2)
    
    copletehull = khull.CompleteHull(ax.sample_box, rand_dir, space=ax.space)
    shull = khull.GBall(ax.sample_box)
    #_hull_model_plot(ax, copletehull, mart.dhull_plot, **kwargs)
    #_hull_model_plot(ax, copletehull, mart.shull_plot, ns=0, **kwargs)
    
    ns=50000
    lim=3
    lw=1.5
    s=4
    #hull_colors=hull_colors, **kwargs):
    border_color, inside_color, outside_color = hull_colors

    # setup
    ax.set_aspect(1)
    #ax.set_box_aspect(3/4)
    ax.set_xlim(-lim, lim)
    ax.set_ylim(-lim, lim)
    ax.set_xlabel('$x_1$')
    ax.set_ylabel('$x_2$')
    
    nodes = ax.sample_box.f_model(ns)
    completemask = copletehull.is_outside(nodes)
    smask = shull.is_outside(nodes)
    mask = np.any([completemask, smask], axis=0)
    mart.scatter_sample(ax, nodes[mask], c=outside_color, marker='.', s=s)
    mart.scatter_sample(ax, nodes[~mask], c=inside_color, marker='.', s=s)
    mart.dhull_plot(ax, copletehull, color=border_color, lw=lw, zorder=100, **kwargs)
    mart.shull_plot(ax, shull, color=border_color, lw=lw/2, ls='--', zorder=100, **kwargs)
    
    # finally samples
    #
    #mart.scatter_sample(ax, f, c='g', marker='P', zorder=1000) # why?
    mart.scatter_points(ax, zorder=100500)