# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/plotting.manifolds.ipynb (unless otherwise specified).

__all__ = ['plot_manifold_2d', 'plot_manifold_3d', 'plot_manifold_featurespace', 'plot_manifold_3d_set']

# Cell
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['text.usetex'] = True
mpl.rcParams['text.latex.preamble'] = [r'\usepackage{amsmath}'] #for \text command
import numpy as np

# Cell
import numpy as np

def plot_manifold_2d(data, s, alpha, c, title, title_color = 'black'):
    fig, axs = plt.subplots(1,1, figsize = (20,15))

    axs.set_axis_off()
    s = 1
    alpha = .5
    selected_points = list(range(data.shape[0]))
    x = data[:, 0][selected_points]
    y = data[:, 1][selected_points]
    ax = fig.add_subplot(1,1,1)

    ax.scatter(x, y,  s=s, alpha=alpha, marker='.',vmin=0,vmax=np.pi, c=c )
    ax.set_xlabel(r'$\phi_1$', fontsize = 140)
    ax.set_ylabel(r'$\phi_2$', fontsize = 140)
    ax.set_title(title, fontsize = 200, color= title_color)

    xmin = np.float(np.format_float_positional(data[:, 0][selected_points].min(), precision=2, fractional=False))
    xmax = np.float(np.format_float_positional(data[:, 0][selected_points].max(), precision=2, fractional=False))
    ymin = np.float(np.format_float_positional(data[:, 1][selected_points].min(), precision=2, fractional=False))
    ymax = np.float(np.format_float_positional(data[:, 1][selected_points].max(), precision=2, fractional=False))
#     ax.set_xticks([xmin, xmax])
#     ax.set_yticks([ymin, ymax])
    ax.set_xticks([])
    ax.set_yticks([])
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    #ax.tick_params(labelsize=30)

# Cell
def plot_manifold_3d(data, s, alpha, c, title, title_color = 'black'):

    fig, axs = plt.subplots(1,1, figsize = (15,10))
    axs.set_axis_off()

    selected_points = list(range(data.shape[0]))
    x = data[:, 0][selected_points]
    y = data[:, 1][selected_points]
    z = data[:, 2][selected_points]

    ax = fig.add_subplot(1,1,1, projection='3d')
    ax.scatter(x, y, z, s=s,c = c, alpha=alpha, marker='.',vmin=0,vmax=np.pi)
    ax.set_xlabel(r'$\phi_1$', fontsize = 90)
    ax.set_ylabel(r'$\phi_2$', fontsize = 90)
    ax.set_zlabel(r'$\phi_3$', fontsize = 90)
    ax.set_title(title, fontsize = 150, color = title_color)

    xmin = np.float(np.format_float_positional(data[:, 0][selected_points].min(), precision=2, fractional=False))
    xmax = np.float(np.format_float_positional(data[:, 0][selected_points].max(), precision=2, fractional=False))
    ymin = np.float(np.format_float_positional(data[:, 1][selected_points].min(), precision=2, fractional=False))
    ymax = np.float(np.format_float_positional(data[:, 1][selected_points].max(), precision=2, fractional=False))
    zmin = np.float(np.format_float_positional(data[:, 2][selected_points].min(), precision=2, fractional=False))
    zmax = np.float(np.format_float_positional(data[:, 2][selected_points].max(), precision=2, fractional=False))

    #ax.set_xticks([xmin, xmax])
    #ax.set_yticks([ymin, ymax])
    #ax.set_zticks([zmin, zmax])
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_zticks([])
    #ax.tick_params(labelsize=30)
    ax.xaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
    ax.yaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
    ax.zaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
    # make the grid lines transparent
    ax.xaxis._axinfo["grid"]['color'] =  (1,1,1,0)
    ax.yaxis._axinfo["grid"]['color'] =  (1,1,1,0)
    ax.zaxis._axinfo["grid"]['color'] =  (1,1,1,0)


# Cell
import math
import matplotlib.pyplot as plt
import numpy as np

def plot_manifold_featurespace(data,title,ncord = 6):

    s = 100
    alpha = .5
    fig, axes = plt.subplots(ncord,ncord, figsize = (40,40))
    xmins = np.zeros(ncord)
    xmaxs = np.zeros(ncord)
    xmeans = np.zeros(ncord)
    for d in range(ncord):
        xmins[d] = data[:,d].min()
        xmaxs[d] = data[:,d].max()
        xmeans[d] = (xmins[d] +  xmaxs[d]) / 2

    gap = (xmaxs - xmins).max()

    for d in range(ncord):
        xmins[d] = xmeans[d] - gap / 2
        xmaxs[d] = xmeans[d] + gap / 2

    for d in range(ncord):
        print(d)
        for e in range(ncord):
            xticks = [math.ceil(xmins[d]), 0 ,math.ceil(xmaxs[d])]
            yticks = [math.ceil(xmins[e]), 0 ,math.ceil(xmaxs[e])]
            if d != e:
                axes[e,d].scatter(data[:,d],data[:,e], s = .1, alpha = .1)
                axes[e,d].set_xlim(xmins[d], xmaxs[d])
                axes[e,d].set_ylim(xmins[e], xmaxs[e])

                if d == 0:
                    #axes[e,d].set_xticks(xticks)
                    axes[e,d].set_yticks(yticks)
                else:
                    axes[e,d].set_yticks([])
                if e == (ncord-1):
                    axes[e,d].set_xticks(xticks)
                else:
                    axes[e,d].set_xticks([])
                axes[e,d].tick_params(labelsize=60)
            if d == e:

                axes[d,e].hist(data[:,d])
                axes[e,d].tick_params(labelsize=50)
                axes[e,d].tick_params(labelsize=50)
                if e == (ncord-1):
                    axes[e,d].set_xticks(xticks)
                else:
                    axes[e,d].set_xticks([])

    for d in range(ncord):
        axes[ncord- 1,d].set_xlabel(r'$\xi_{{{}}}$'.format(d+1), fontsize= 110)
        axes[d,0].set_ylabel(r'$\xi_{{{}}}$'.format(d+1), fontsize= 110)

    plt.tight_layout()
    fig = plt.gcf()
    fig.subplots_adjust(top=0.9)
    fig.suptitle(title,y=.95,x = .53, fontsize=120)




# Cell

def plot_manifold_3d_set(data, s, alpha, gb, titles,sub, title_colors, title):

    nb = len(sub)


    fig, axs = plt.subplots(math.floor(nb / 4) + 1,4, figsize = (40, 10 *( math.floor(nb / 4) + 1),))

    for r in range((math.floor(nb / 4)+1 )* 4):
        axs[math.floor(r / 4), r %4 ].set_axis_off()

    selected_points = list(range(data.shape[0]))
    x = data[:, 0][selected_points]
    y = data[:, 1][selected_points]
    z = data[:, 2][selected_points]
    for r in range(nb):
        print(r)
        ax = fig.add_subplot(math.floor(nb / 4)+1, 4 ,r+1, projection='3d')
        ax.scatter(x, y, z, c = gb[:,r], s=s, alpha=alpha, marker='.',vmin=0,vmax=np.pi)
        ax.set_xlabel(r'$\phi_1$', fontsize = 60)
        ax.set_ylabel(r'$\phi_2$', fontsize = 60)
        ax.set_zlabel(r'$\phi_3$', fontsize = 60)
        ax.set_title(r'$g_{{{}}}$'.format(sub[r]) + ' ' + str(titles[sub[r]]+1), fontsize = 60, color = title_colors[r])
        #ax.set_axis_off()
        xmin = np.float(np.format_float_positional(data[:, 0][selected_points].min(), precision=2, fractional=False))
        xmax = np.float(np.format_float_positional(data[:, 0][selected_points].max(), precision=2, fractional=False))
        ymin = np.float(np.format_float_positional(data[:, 1][selected_points].min(), precision=2, fractional=False))
        ymax = np.float(np.format_float_positional(data[:, 1][selected_points].max(), precision=2, fractional=False))
        zmin = np.float(np.format_float_positional(data[:, 2][selected_points].min(), precision=2, fractional=False))
        zmax = np.float(np.format_float_positional(data[:, 2][selected_points].max(), precision=2, fractional=False))
        ax.set_xticks([])
        ax.set_yticks([])
        ax.set_zticks([])
        ax.set_xticklabels([])
        ax.set_yticklabels([])
        ax.set_zticklabels([])
        ax.tick_params(labelsize=40)

        ax.xaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
        ax.yaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
        ax.zaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
        # make the grid lines transparent
        ax.xaxis._axinfo["grid"]['color'] =  (1,1,1,0)
        ax.yaxis._axinfo["grid"]['color'] =  (1,1,1,0)
        ax.zaxis._axinfo["grid"]['color'] =  (1,1,1,0)

    #plt.tight_layout()
    fig = plt.gcf()
    fig.subplots_adjust(top=0.9)
    fig.suptitle(title,y=.95,x = .53, fontsize=80)