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


#č nazvy proměnných jsou v angličtině
#č Ale komenty teda ne)

import numpy as np
import matplotlib.tri as mtri
import matplotlib.colors as mcolor
import matplotlib.patches as mpatches

from . import misc

def get_simplices_colors(sb, simplices): #simplices = bx.tri.simplices
    s = '#a7ffb5' #'xkcd:light seafoam green' #a7ffb5
    f = '#fdc1c5' #'xkcd: pale rose' # (#fdc1c5)
    m = '#FFF39A' #'xkcd: dark cream' # (255, 243, 154, 255)
    
    in_failure = np.isin(simplices, sb.failure_points)
    has_failure = in_failure.any(axis=1)
    all_failure = in_failure.all(axis=1)
    return np.where(has_failure, np.where(all_failure, f, m), s)

def get_events(sb, simplices): #simplices = bx.tri.simplices
    """
    Metoda musí simplexům přiřazovat jev 
    0=success, 1=failure, 2=mix
    """
    in_failure = np.isin(simplices, sb.failure_points)
    has_failure = in_failure.any(axis=1)
    all_failure = in_failure.all(axis=1)
    return np.int8(np.where(has_failure, np.where(all_failure, 1, 2), 0))

#č napadlo mě, že bych mohl matplotlibovskému Axes
#č přiřazovat (připsavat, zadávat) atribut 'space'
#č Daválo by to smysl, ne? U všeho ostatního, u sample boksů
#č ne vždy na jedném sabplotu někdo potřebuje
#č kreslit z různejch prostoru.
#č Zkrátka, funkce v tomto modulu požadujou aby 
#č ax.space a ax.sample_box byl nastaven!

# ax.space and ax.sample_box attributes should (expected to) be set up!


def scatter_sample(ax, sample, **kwargs):
    xy = getattr(sample, ax.space)
    x, y = xy[:,:2].T
    return ax.scatter(x, y, **kwargs)

def plot_sample(ax, sample, *args, **kwargs):
    xy = getattr(sample, ax.space)
    x, y = xy[:,:2].T
    return ax.plot(x, y, *args, **kwargs)



def scatter_points(ax, **kwargs):
    xy = getattr(ax.sample_box, ax.space)
    nsim = len(xy)
    x, y = xy[:,:2].T
    
    failsi = ax.sample_box.failsi
    
    try: # proxy denotes to implicitly-known values
        proxy = ax.sample_box.proxy
    except AttributeError:
        proxy = np.full(nsim, False, dtype=np.bool)
    
    
    mask = np.all((~failsi, ~proxy), axis=0)
    success = ax.scatter(x[mask], y[mask], c='g', marker='P', **kwargs)
    
    mask = np.all((failsi, ~proxy), axis=0)
    failures = ax.scatter(x[mask], y[mask], c='r', marker='X', **kwargs)
    
    mask = np.all((~failsi, proxy), axis=0)
    proxy_successes = ax.scatter(x[mask], y[mask], c='#77AC30', marker='h', **kwargs)
    
    mask = np.all((failsi, proxy), axis=0)
    proxy_failures = ax.scatter(x[mask], y[mask], c='#D95319', marker='H', **kwargs)
    
    return success, failures, proxy_successes, proxy_failures
    


def triplot(ax, **kwargs):
    xy = getattr(ax.sample_box, ax.space)
    x, y = xy[:,:2].T
    
    return ax.triplot(x, y, **kwargs)


def tripcolor(ax, sfm_colors=None, **kwargs):
    xy = getattr(ax.sample_box, ax.space)
    x, y = xy[:,:2].T
    tri = mtri.Triangulation(x, y)
    
    if sfm_colors is None:
        # make a color map of fixed colors
        s = '#a7ffb5' #'xkcd:light seafoam green' #a7ffb5
        f = '#fdc1c5' #'xkcd: pale rose' # (#fdc1c5)
        m = '#FFF39A' #'xkcd: dark cream' # (255, 243, 154, 255)
        sfm_colors = [s, f, m]
    
    if 'cmap' not in kwargs:
        # 0=success, 1=failure, 2=mix
        kwargs['cmap'] = mcolor.ListedColormap(sfm_colors)
    if 'norm' not in kwargs:
        kwargs['norm'] = mcolor.NoNorm()
    
    # same as facecolors
    C = get_events(ax.sample_box, tri.get_masked_triangles())
    
    #č tak to má bejt, aby MPL jednoznačně bral barvy jako barvy obličejů
    #č jenomže to může zlobit
    return ax.tripcolor(tri, facecolors=C, **kwargs)
    

def plot_boundaries(ax, fmt='-b', nrod=200, **kwargs):
    xmin, xmax = ax.get_xlim()
    ymin, ymax = ax.get_ylim()
    limits = np.array([[xmin, ymin], [xmax, ymax]])
    bounds = ax.sample_box.get_2D_boundary(nrod=nrod, viewport_sample=limits,\
                                             viewport_space=ax.space)
    lines = []
    for bound in bounds:
        xy = getattr(bound, ax.space)
        x, y = xy[:,:2].T
        lines.append(ax.plot(x, y, fmt, **kwargs))
    
    return lines


def qhull_polygon(ax, qhull, **kwargs):
    x, y = qhull.points[qhull.vertices].T
    return ax.fill(x, y, **kwargs)


def qhull_plot(ax, qhull, fmt='-', ns=100, **kwargs):
    if (qhull.sample.nvar == 2) and (ax.space == qhull.space):
        points = qhull.points
        lines = []
        for simplex in qhull.simplices:
            xy = points[simplex]
            x, y = xy.T
            lines.append(ax.plot(x, y, fmt, **kwargs))
        return lines
                    
#    else:
#        #оӵ кулэ ӧвӧл обновлять экран карыны
#        for simplex in qhull.simplices:
#            start_id, end_id = simplex
#            
#            x_bound = np.linspace(sampled_plan_tri[start_id,0], sampled_plan_tri[end_id,0], ns, endpoint=True)
#            y_bound = np.linspace(sampled_plan_tri[start_id,1], sampled_plan_tri[end_id,1], ns, endpoint=True)
#            
#            # sample compatible
#            #оӵ малы транспонировать кароно? Озьы кулэ!
#            bound_tri = np.vstack((x_bound, y_bound)).T
#            #č vytvořme sample
#            bound = ax.sample_box.f_model.new_sample(bound_tri, space=ax.sample_box.tri_space)
#            
#            xy = getattr(bound, ax.space)
#            x, y = xy.T
#            lines.append(ax.plot(x, y, fmt, **kwargs))
                

def dhull_plot(ax, hull, **kwargs):
    #č zatím uděláme jen pro 2D infinite lajny
    design_points = hull.get_design_points()
    lines = []
    if (hull.sample.nvar == 2) and (ax.space == hull.space):
        for equation in hull.equations:
            #č ve 2D bych očekával v rovnici pouze 3 hodnoty (já potřebuji směry)
            x, y, offset = equation
            design_point = [-x*offset, -y*offset] #č offset je prej zápornej
            slope = np.divide(-x, y)
            lines.append(ax.axline(design_point, slope=slope, **kwargs))
    return lines
            
def bhull_plot(ax, bhull, **kwargs):
    if ax.space == bhull.space:
        point = bhull.mins[:2]
        x1, y1 = point
        x2, y2 = bhull.maxs[:2]
        frame = mpatches.Rectangle(point, x2-x1, y2-y1, fill=False, **kwargs)
        return ax.add_patch(frame) 

def shull_plot(ax, shull, **kwargs):
    r = shull.get_r()
    return gcircle(ax, r=r, fill=False, **kwargs)


def convex_plot(ax, fmt='-m', ns=100, **kwargs):
    if ax.sample_box.nvar == 2:
        simplices = ax.sample_box.convex_hull.simplices
        # convex hull should be made in the same space as triangulation, 
        # Will we take coordinates in the triangulation space, I guess?
        sampled_plan_tri = getattr(ax.sample_box, ax.sample_box.tri_space)
        
        # hmm...
        lines = []
        if ax.space == ax.sample_box.tri_space:
            for simplex in simplices:
                xy = sampled_plan_tri[simplex]
                x, y = xy.T
                lines.append(ax.plot(x, y, fmt, **kwargs))
                    
    else:
        #оӵ кулэ ӧвӧл обновлять экран карыны
        for simplex in simplices:
            start_id, end_id = simplex
            
            x_bound = np.linspace(sampled_plan_tri[start_id,0], sampled_plan_tri[end_id,0], ns, endpoint=True)
            y_bound = np.linspace(sampled_plan_tri[start_id,1], sampled_plan_tri[end_id,1], ns, endpoint=True)
            
            # sample compatible
            #оӵ малы транспонировать кароно? Озьы кулэ!
            bound_tri = np.vstack((x_bound, y_bound)).T
            #č vytvořme sample
            bound = ax.sample_box.f_model.new_sample(bound_tri, space=ax.sample_box.tri_space)
            
            xy = getattr(bound, ax.space)
            x, y = xy.T
            lines.append(ax.plot(x, y, fmt, **kwargs))
                
                
    return lines






def gcircle(ax, r=1, nrod=200, **kwargs):
    if ax.space == 'G':
        circle = mpatches.Circle((0,0), r, **kwargs)
    else:
        phi = np.linspace(0, 6.283185307, nrod, endpoint=True)
        cos_phi = np.cos(phi)
        sin_phi = np.sin(phi)
        
        sample_G = np.array((cos_phi, sin_phi)).T * r

        f_model = ax.sample_box.f_model
        sample = f_model.new_sample(sample_G, space='G', extend=True)
        xy = getattr(sample, ax.space)[:,:2]
        circle = mpatches.Polygon(xy, **kwargs)
    
    #č vrací add_patch něco?
    return ax.add_patch(circle)


def uframe(ax, **kwargs):
    if ax.space in ('P', 'aP', 'U', 'aU'):
        alpha = ax.sample_box.alpha
        frame = mpatches.Rectangle((0,0), alpha[0], alpha[1], fill=False, **kwargs)
        return ax.add_patch(frame) 



def isocurves(ax, ngrid=200, limits=None, ncurves=5, **kwargs):
    if limits is None:
        #xmin, xmax = ax.get_xlim()
        #ymin, ymax = ax.get_ylim()
        
        sample_G = np.array([[-ncurves, -ncurves], [ncurves, ncurves]])
        sample = ax.sample_box.f_model.new_sample(sample_G, space='G', extend=True)
        xy = getattr(sample, ax.space)[:,:2]
        xmin, ymin = np.min(xy, axis=0)
        xmax, ymax = np.max(xy, axis=0)
        
    else: # G-čko zlobí
        xmin, xmax, ymin, ymax = limits
    
    
    
    x = np.linspace(xmin, xmax, ngrid)
    y = np.linspace(ymin, ymax, ngrid)
    X, Y = np.meshgrid(x, y)
    XY = np.vstack((X.flatten(), Y.flatten())).T
    Z = ax.sample_box.f_model.sample_pdf(XY, ax.space)
    
    
    const = 1 / (xmax - xmin) / (ymax - ymin)
    r_levels = np.arange(ncurves) + 1
    levels = misc.isolevels_2d(Z, const, np.flip(r_levels), from_top=False)
    return ax.contour(X, Y, Z.reshape(ngrid, ngrid), levels, **kwargs)



def center_spines(ax):
    # Move the left and bottom spines to x = 0 and y = 0, respectively.
    ax.spines["left"].set_position(("data", 0))
    ax.spines["bottom"].set_position(("data", 0))
    
    # Show the left and bottom spines
    ax.spines["left"].set_visible(True)
    ax.spines["bottom"].set_visible(True)
    
    # Hide the top and right spines.
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    
    # Draw arrows (as black triangles: ">k"/"^k") at the end of the axes.  In each
    # case, one of the coordinates (0) is a data coordinate (i.e., y = 0 or x = 0,
    # respectively) and the other one (1) is an axes coordinate (i.e., at the very
    # right/top of the axes).  Also, disable clipping (clip_on=False) as the marker
    # actually spills out of the axes.
    ax.plot(1, 0, ">k", transform=ax.get_yaxis_transform(), clip_on=False)
    ax.plot(0, 1, "^k", transform=ax.get_xaxis_transform(), clip_on=False)


def fix_locator(ax, loc):
    loc_x = mpl.ticker.FixedLocator(loc)
    ax.xaxis.set_major_locator(loc_x)
    loc_y = mpl.ticker.FixedLocator(loc)
    ax.yaxis.set_major_locator(loc_y)

#č něco_kulatého
def curly(ax, linewidths=[0.7, 0.5, 0.4, 0.3, 0.2, 0.1], nrid=200, color='k', **kwargs):
    if ax.space in ('U', 'aU'):
        return None
        
    elif (ax.sample_box.nvar==2) and (ax.space in ('Rn', 'aRn', 'R', 'aR', 'P', 'aP')):
        isocurves(ax, ngrid=nrid, limits=None, ncurves=len(linewidths),\
                     linewidths=np.flip(linewidths), colors=[color], **kwargs)
        
    else:
        for i, lw in zip(range(len(linewidths)), linewidths):
            gcircle(ax, r=i+1, nrod=nrid, color=color, linewidth=lw, fill=False)


def setup(ax, lw=1): #č šup
    #ax.set_xlabel('$x_{1}$')
    #ax.set_ylabel('$x_{2}$')
    
    #ax.set_frame_on(False) #č pak se mi nezobrazí osy
    ax.set_aspect(1)
    #ax.set_box_aspect(1)
    if ax.space in ('P', 'aP', 'U', 'aU'):
        ax.margins(0)
        ax.set_frame_on(False)
        uframe(ax, linewidth=lw) 
    else:
        center_spines(ax)