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


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

import numpy as np
from matplotlib import colors
import matplotlib.tri as mtri

# copied from simplex module
# just don't want every time import simplex and its dependencies
def get_events(sx, simplices=None):
    """
    Metoda musí simplexům přiřazovat jev 
    0=success, 1=failure, 2=mix
    """
    if simplices is None:
        simplices = sx.tri.simplices
    
    in_failure = np.isin(simplices, sx.sample_box.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))


def get_g_model_wireframe_data(shape_share, limits, space='R', ngrid=50):
    xmin, xmax, ymin, ymax = limits
    
    # endpoint=True by default
    # we'll get len(.) == ngrid
    # ngrid 51 is an default maximum for MPL
    # (otherwise MPL will downsample)
    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
    
    if space == 'R':
        g_values = shape_share.gm(XY).g_values
    else:
        XY_sample = shape_share.f_model.new_sample(XY, space)
        g_values = shape_share.gm(XY_sample).g_values
    
    Z = g_values.reshape(ngrid, ngrid)
    return X, Y, Z


#č 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ů
#č nejsem jist, ale mám pocit, že na jedném sablotu 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_points(ax3d, zs=None, **kwargs):
    sample = getattr(ax3d.sample_box, ax3d.space)
    nsim = len(sample)
    if zs is not None:
        x, y = sample[:,:2].T
    elif ax3d.sample_box.nvar == 2:
        x, y = sample[:,:2].T
        #zs = ax3d.sample_box.pdf(ax3d.space)
        zs = ax3d.sample_box.g_values
    else:
        x, y, zs = sample[:,:3].T
    
    failsi = ax3d.sample_box.failsi
    
    try: # proxy denotes to implicitly-known values
        proxy = ax3d.sample_box.proxy
    except AttributeError:
        proxy = np.full(nsim, False, dtype=np.bool)
    
    
    mask = np.all((~failsi, ~proxy), axis=0)
    success = ax3d.scatter(x[mask], y[mask], zs=zs[mask], c='g', marker='P', **kwargs)
    
    mask = np.all((failsi, ~proxy), axis=0)
    failures = ax3d.scatter(x[mask], y[mask], zs=zs[mask], c='r', marker='X', **kwargs)
    
    mask = np.all((~failsi, proxy), axis=0)
    proxy_successes = ax3d.scatter(x[mask], y[mask], zs=zs[mask],\
                             c='#77AC30', marker='h', **kwargs)
    
    mask = np.all((failsi, proxy), axis=0)
    proxy_failures = ax3d.scatter(x[mask], y[mask], zs=zs[mask],\
                             c='#D95319', marker='H', **kwargs)
    
    return success, failures, proxy_successes, proxy_failures


def limit_state_wireframe(ax3d, **kwargs):
    xmin, xmax = ax3d.get_xlim()
    ymin, ymax = ax3d.get_ylim()
    limits = (xmin, xmax, ymin, ymax)
    
    if 'rcount' in kwargs:
        ngrid = kwargs['rcount']
    else:
        ngrid = 50
    
    X, Y, Z = get_g_model_wireframe_data(ax3d.sample_box,\
                             limits, space=ax3d.space, ngrid=ngrid)
        
    return ax3d.plot_wireframe(X, Y, Z, **kwargs)


def limit_state_surface(ax3d, **kwargs):
    xmin, xmax = ax3d.get_xlim()
    ymin, ymax = ax3d.get_ylim()
    limits = (xmin, xmax, ymin, ymax)
    
    if 'rcount' in kwargs:
        ngrid = kwargs['rcount']
    else:
        ngrid = 50
    
    X, Y, Z = get_g_model_wireframe_data(ax3d.sample_box,\
                             limits, space=ax3d.space, ngrid=ngrid)
    
    # make a color map of fixed colors
    if 'cmap' not in kwargs:
        kwargs['cmap'] = colors.ListedColormap(['red', 'green'])
    if 'norm' not in kwargs:
        kwargs['norm'] = colors.BoundaryNorm([0], kwargs['cmap'].N, extend='both')
        
    return ax3d.plot_surface(X, Y, Z, **kwargs)


def density_surface(ax3d, **kwargs):
    xmin, xmax = ax3d.get_xlim()
    ymin, ymax = ax3d.get_ylim()
    
    if 'rcount' in kwargs:
        ngrid = kwargs['rcount']
    else:
        ngrid = 50
    
    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 = ax3d.sample_box.sample_pdf(XY, ax3d.space)
    Z = z.reshape(ngrid, ngrid)
    
    return ax3d.plot_surface(X, Y, Z, **kwargs)
    

def density_colored_surface(ax3d, hull, colors=('#7735C2', '#808080'), **kwargs):
    xmin, xmax = ax3d.get_xlim()
    ymin, ymax = ax3d.get_ylim()
    
    inside_color, outside_color = colors
    
    if 'rcount' in kwargs:
        ngrid = kwargs['rcount']
    else:
        ngrid = 50
    
    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 = ax3d.sample_box.sample_pdf(XY, ax3d.space)
    Z = z.reshape(ngrid, ngrid)
    
    # for facecolors
    _x = np.linspace(xmin, xmax, ngrid-1, endpoint=False) + (xmax-xmin)/(ngrid-1)/2
    _y = np.linspace(ymin, ymax, ngrid-1, endpoint=False) + (ymax-ymin)/(ngrid-1)/2
    _X, _Y = np.meshgrid(_x, _y)
    _XY = np.vstack((_X.flatten(), _Y.flatten())).T
    facecolors = np.full(len(_XY), inside_color)
    mask = hull.is_outside(hull.sample.f_model.new_sample(_XY, space=ax3d.space))
    facecolors[mask] = outside_color
    facecolors = facecolors.reshape(ngrid-1, ngrid-1)
    
    return ax3d.plot_surface(X, Y, Z, facecolors = facecolors, **kwargs)

    
def tri_surface(ax3d, **kwargs):
    xy = getattr(ax3d.sample_box, ax3d.space)
    x, y = xy.T
    z = ax3d.sample_box.g_values
    if 'zs' in kwargs:
        z = kwargs.pop('zs')
    else:
        z = ax3d.sample_box.pdf(ax3d.space)
    return ax3d.plot_trisurf(x, y, z, **kwargs)
    



# from Matplotlib sources:
# "TODO: Support custom face colours"
    
#def tri_colored_surface(ax3d, **kwargs):
#    xy = getattr(ax3d.sample_box, ax3d.space)
#    x, y = xy.T
#    z = ax3d.sample_box.g_values
#    tri = mtri.Triangulation(x, y)
#    colors = get_simplices_colors(ax3d.sample_box, tri.get_masked_triangles())
#    return ax3d.plot_trisurf(tri, z, color=colors, **kwargs)


def tri_colored_surfaces(ax3d, **kwargs):
    xy = getattr(ax3d.sample_box, ax3d.space)
    x, y = xy.T
    if 'zs' in kwargs:
        z = kwargs.pop('zs')
    else:
        z = ax3d.sample_box.g_values
    if 'tri_colors' in kwargs:
        tri_colors = kwargs.pop('tri_colors')
    else:
        tri_colors = ("xkcd:light seafoam green", "xkcd:pale rose", "xkcd:dark cream")
    tri = mtri.Triangulation(x, y)
    triangles = tri.get_masked_triangles()
    events = get_events(ax3d.sample_box, triangles)
    # 0=success, 1=failure, 2=mix
    surf = list()
    if np.any(events==0):
        surf.append(ax3d.plot_trisurf(x, y, triangles[events==0], z, color=tri_colors[0], **kwargs))
    if np.any(events==1):
        surf.append(ax3d.plot_trisurf(x, y, triangles[events==1], z, color=tri_colors[1], **kwargs))
    if np.any(events==2):
        surf.append(ax3d.plot_trisurf(x, y, triangles[events==2], z, color=tri_colors[2], **kwargs))
    return surf



def wall(ax3d, lines2d, zmin=0, zmax=1, **kwargs):
    surf = list()
    for line2d in lines2d:
        x = line2d.get_xdata()
        y = line2d.get_ydata()
        _x = np.concatenate((x, x))
        _y = np.concatenate((y, y))
        _z = np.concatenate([np.full(len(x), zmin), np.full(len(x), zmax)])
        tri = mtri.Triangulation(_x, _z)
        triangles = tri.get_masked_triangles()
        surf.append(ax3d.plot_trisurf(_x, _y, _z, triangles=triangles, **kwargs))