import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
import scipy.special

"""
Explanation: Getting an Overview of Regular 3D Data
In this notebook, we're going to talk a little bit about how you might get an overview of regularized 3D data, specifically using matplotlib.
In a subsequent notebook we'll address the next few steps, specifically how you might use tools like ipyvolume and yt.
To start with, let's generate some fake data!  (Now, I say 'fake,' but that's a bit pejorative, isn't it?  Data is data!  Ours is just synthetic.)
End of explanation
"""


N = 64
x = np.mgrid[-1.0:1.0:N*1j][:,None,None]
y = np.mgrid[-1.0:1.0:N*1j][None,:,None]
z = np.mgrid[-1.0:1.0:N*1j][None,None,:]

r = np.sqrt(x*x + y*y + z*z)
theta = np.arctan2(np.sqrt(x*x + y*y), z)
phi = np.arctan2(y, x)

np.abs(x - r * np.sin(theta)*np.cos(phi)).max()

np.abs(y - r * np.sin(theta)*np.sin(phi)).max()

np.abs(z - r * np.cos(theta)).max()

data = {}
for n in [1, 4]:
    for m in range(n + 1):
        data[f"sph_n{n}_m{m}"] = np.absolute(scipy.special.sph_harm(m, n, phi, theta))

"""
Explanation: We'll use the scipy spherical harmonics function to make some data, but first we need a reference coordinate system.  We'll start with $x, y, z$ and then transform them into spherical coordinates.
Note: we'll be using the convention that $\theta \in [0, \pi]$ and $\phi \in[0,2\pi)$, which is reverse from what SciPy expects.  So if you compare to the docstring for sph_harm, keep that in mind.  Feel free to switch the definitions if you like!
End of explanation
"""


plt.imshow(data["sph_n4_m4"][:,:,N//4], norm=LogNorm())
plt.colorbar()

phi.min(), phi.max()

plt.imshow(data["sph_n1_m0"].max(axis=0), norm=LogNorm())
plt.colorbar()

"""
Explanation: Now we have some data!  And, we can use matplotlib to visualize it in reduced form.  Let's try this out:
End of explanation
"""


import ipywidgets

@ipywidgets.interact(dataset = list(sorted(data.keys())), slice_position = (0, N, 1))
def make_plots(dataset, slice_position):
    plt.imshow(data[dataset][slice_position,:,:], norm=LogNorm())
    plt.colorbar()

"""
Explanation: This is getting a bit cumbersome, though!  Let's try using the ipywidgets library to speed this up just a bit.
We're going to use the ipywidgets.interact decorator around our function to add some inputs.  This is a pretty powerful decorator, as it sets up new widgets based on the info that you feed it, and then re-executes the function every time those inputs change.
End of explanation
"""


print(min(_.min() for _ in data.values()), max(_.max() for _ in data.values()))

"""
Explanation: We still have some artifacts here we want to get rid of; let's see if we can restrict our colorbar a bit.
End of explanation
"""


@ipywidgets.interact(dataset = list(sorted(data.keys())), slice_position = (0, N, 1))
def make_plots(dataset, slice_position):
    plt.imshow(data[dataset][slice_position,:,:], norm=LogNorm(vmin=1e-5, vmax=1.0))
    plt.colorbar(extend = 'min')

"""
Explanation: Typically in these cases, the more interesting values are the ones at the top -- the bottom are usually falling off rather quickly to zero.  So let's set our maximum, and then drop 5 orders of magnitude for the minimum.  I'm changing the colorbar's "extend" value to reflect this.
End of explanation
"""


@ipywidgets.interact(dataset = list(sorted(data.keys())), x = (-1.0, 1.0, 2.0/N), y = (-1.0, 1.0, 2.0/N), z = (-1.0, 1.0, 2.0/N))
def make_plots(dataset, x, y, z):
    xi, yi, zi = (int(_*N + 1.0) for _ in (x, y, z))
    fig, axes = plt.subplots(nrows=2, ncols=2, dpi = 200)
    datax = data[dataset][xi,:,:]
    datay = data[dataset][:,yi,:]
    dataz = data[dataset][:,:,zi]
    vmax = max(_.max() for _ in (datax, datay, dataz))
    vmin = max( min(_.min() for _ in (datax, datay, dataz)), vmax / 1e5)
    imx = axes[0][0].imshow(datax, norm=LogNorm(vmin=vmin, vmax=vmax), extent = [-1.0, 1.0, -1.0, 1.0])
    imy = axes[0][1].imshow(datay, norm=LogNorm(vmin=vmin, vmax=vmax), extent = [-1.0, 1.0, -1.0, 1.0])
    imz = axes[1][0].imshow(dataz, norm=LogNorm(vmin=vmin, vmax=vmax), extent = [-1.0, 1.0, -1.0, 1.0])
    fig.delaxes(axes[1][1])
    fig.colorbar(imx, ax=axes, extend = 'min', fraction = 0.1)

import plotly.graph_objects as go

plt.hist(data["sph_n4_m3"].flatten())

iso_data=go.Isosurface(
    x=(x * np.ones((N,N,N))).flatten(),
    y=(y * np.ones((N,N,N))).flatten(),
    z=(z * np.ones((N,N,N))).flatten(),
    value=data["sph_n4_m3"].flatten(),
    isomin=0,
    isomax=data["sph_n4_m3"].max(),
    surface_count=5, # number of isosurfaces, 2 by default: only min and max
    colorbar_nticks=5, # colorbar ticks correspond to isosurface values
    caps=dict(x_show=False, y_show=False))
fig = go.Figure(data = iso_data)
fig

"""
Explanation: We're going to do one more thing for getting an overview, and then we'll see if we can do some other, cooler things with it using plotly.
We're going to change our slice_position to be in units of actual coordinates, instead of integers, and we'll add on a multiplot so we can see all three at once.
End of explanation
"""


plt.imshow(data["sph_n4_m3"].sum(axis=0), extent=[-1.0, 1.0, -1.0, 1.0])

"""
Explanation: One thing I've run into with plotly while making this notebook has been that in many cases, the 3D plots strain a bit under large data sizes.  This is to be expected, and is completely understandable!  One of the really nice things about regular mesh data like this is that you can usually cut it down quite effectively with slices.  Unfortunately, what I have found -- and I may have done something completely wrong! -- is that plotly some times appears to almost work, and then doesn't quite make it when I throw too much data at it.  I've found that it seems to work best in the neighborhood of $64^3$ zones, maybe a bit more.
Other Summary Techniques
There are, of course, other ways you can take a look at a set of values!  Given a regular mesh, it's straightforward with numpy to apply any of the reduction operations along one of the axes.  For instance, you might take the min, the max, the sum, the mean and so forth.  If we do this with our spherical harmonics data:
End of explanation
"""


