import numpy as np

import PIL.Image
import pythreejs

import ipyvolume as ipv
from .datasets import UrlCached

def _randomSO3():
    """return random rotatation matrix, algo by James Arvo"""
    u1 = np.random.random()
    u2 = np.random.random()
    u3 = np.random.random()
    R = np.array([[np.cos(2*np.pi*u1), np.sin(2*np.pi*u1), 0], [-np.sin(2*np.pi*u1), np.cos(2*np.pi*u1), 0], [0, 0, 1]])
    v = np.array([np.cos(2*np.pi*u2)*np.sqrt(u3), np.sin(2*np.pi*u2)*np.sqrt(u3), np.sqrt(1-u3)])
    H = np.identity(3)-2*v*np.transpose([v])
    return - np.dot(H,  R)

def spherical_galaxy_orbit(orbit_x, orbit_y, orbit_z, N_stars=100, sigma_r=1, orbit_visible=False, orbit_line_interpolate=5, N_star_orbits=10, color=[255, 220, 200], size_star=1, scatter_kwargs={}):
    """Create a fake galaxy around the points orbit_x/y/z with N_stars around it"""
    if orbit_line_interpolate > 1:
        import scipy.interpolate
        x = np.linspace(0, 1, len(orbit_x))
        x_smooth = np.linspace(0, 1, len(orbit_x)*orbit_line_interpolate)
        kind = 'quadratic'
        orbit_x_line = scipy.interpolate.interp1d(x, orbit_x, kind)(x_smooth)
        orbit_y_line = scipy.interpolate.interp1d(x, orbit_y, kind)(x_smooth)
        orbit_z_line = scipy.interpolate.interp1d(x, orbit_z, kind)(x_smooth)
    else:
        orbit_x_line = orbit_x
        orbit_y_line = orbit_y
        orbit_z_line = orbit_z
    line = ipv.plot(orbit_x_line, orbit_y_line, orbit_z_line, visible=orbit_visible)
    x = np.repeat(orbit_x, N_stars).reshape((-1, N_stars))
    y = np.repeat(orbit_y, N_stars).reshape((-1, N_stars))
    z = np.repeat(orbit_z, N_stars).reshape((-1, N_stars))
    xr, yr, zr = np.random.normal(0, scale=sigma_r, size=(3, N_stars))# + 
    r = np.sqrt(xr**2 + yr**2 + zr**2)
    
    for i in range(N_stars):
        a = np.linspace(0, 1, x.shape[0]) * 2 * np.pi * N_star_orbits
        xo = r[i] * np.sin(a)
        yo = r[i] * np.cos(a)
        zo = a * 0
        xo, yo, zo = np.dot(_randomSO3(), [xo, yo, zo])
        #print(x.shape, xo.shape)
        x[:, i] += xo
        y[:, i] += yo
        z[:, i] += zo
    
    
    sprite = ipv.scatter(x, y, z, texture=radial_sprite((64, 64), color), marker='square_2d', size=size_star, **scatter_kwargs)
    with sprite.material.hold_sync():
        sprite.material.blending = pythreejs.BlendingMode.CustomBlending
        sprite.material.blendSrc = pythreejs.BlendFactors.SrcColorFactor
        sprite.material.blendDst = pythreejs.BlendFactors.OneFactor
        sprite.material.blendEquation = 'AddEquation'
        sprite.material.transparent = True
        sprite.material.depthWrite = False
        sprite.material.alphaTest = 0.1
    return sprite, line

def radial_sprite(shape, color):
    color = np.array(color)
    ara = np.zeros(shape[:2] + (4,), dtype=np.uint8)
    x = np.linspace(-1, 1, shape[0])
    y = np.linspace(-1, 1, shape[1])
    x, y = np.meshgrid(x, y)
    s = 0.5
    radius = np.sqrt(x**2+y**2)
    amplitude = np.maximum(0, np.exp(-radius**2/s**2)).T
    ara[...,3] = (amplitude * 255)
    ara[...,:3] = color * amplitude.reshape(shape + (1,))
    im = PIL.Image.fromarray(ara, 'RGBA')
    return im

def stars(N=1000, radius=100000, thickness=3, seed=42, color=[255, 240, 240]):
    import ipyvolume as ipv
    rng = np.random.RandomState(seed)
    x, y, z = rng.normal(size=(3, N))
    r = np.sqrt(x**2 + y**2 + z**2)/(radius + thickness * radius * np.random.random(N))
    x /= r
    y /= r
    z /= r
    return ipv.scatter(x, y, z, texture=radial_sprite((64, 64), color), marker='square_2d', grow_limits=False, size=radius*0.7/100)


milkyway_url = 'https://www.nasa.gov/sites/default/files/images/620057main_milkyway_full.jpg'
milkyway_image = UrlCached(milkyway_url)



def plot_milkyway(R_sun=8, size=100):
    mw_image = PIL.Image.open(milkyway_image.fetch())
    rescale = 40
    t = np.linspace(0, 1, 100)
    xmw = np.linspace(0, 1, 10)
    ymw = np.linspace(0, 1, 10)
    xmw, ymw = np.meshgrid(xmw, ymw)
    zmw = xmw * 0 + 0.01
    mw = mesh = ipv.plot_mesh((xmw-0.5)*rescale, (ymw-0.5)*rescale+R_sun, zmw, u=xmw, v=ymw, texture=mw_image, wireframe=False)
    mw.material.blending = pythreejs.BlendingMode.CustomBlending
    mw.material.blendSrc = pythreejs.BlendFactors.SrcColorFactor
    mw.material.blendDst = pythreejs.BlendFactors.OneFactor
    mw.material.blendEquation = 'AddEquation'
    mw.material.transparent = True
    mw.material.depthWrite = False
    mw.material.alphaTest = 0.1
    ipv.xyzlim(size)
    return mesh
[((11, 9, 11, 27), 'numpy.random.random', 'np.random.random', ({}, {}), '()', True, 'import numpy as np\n'), ((12, 9, 12, 27), 'numpy.random.random', 'np.random.random', ({}, {}), '()', True, 'import numpy as np\n'), ((13, 9, 13, 27), 'numpy.random.random', 'np.random.random', ({}, {}), '()', True, 'import numpy as np\n'), ((33, 11, 33, 84), 'ipyvolume.plot', 'ipv.plot', (), '', True, 'import ipyvolume as ipv\n'), ((37, 17, 37, 70), 'numpy.random.normal', 'np.random.normal', (), '', True, 'import numpy as np\n'), ((38, 8, 38, 38), 'numpy.sqrt', 'np.sqrt', ({(38, 16, 38, 37): 'xr ** 2 + yr ** 2 + zr ** 2'}, {}), '(xr ** 2 + yr ** 2 + zr ** 2)', True, 'import numpy as np\n'), ((64, 12, 64, 27), 'numpy.array', 'np.array', ({(64, 21, 64, 26): 'color'}, {}), '(color)', True, 'import numpy as np\n'), ((65, 10, 65, 52), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((66, 8, 66, 36), 'numpy.linspace', 'np.linspace', ({(66, 20, 66, 22): '-1', (66, 24, 66, 25): '1', (66, 27, 66, 35): 'shape[0]'}, {}), '(-1, 1, shape[0])', True, 'import numpy as np\n'), ((67, 8, 67, 36), 'numpy.linspace', 'np.linspace', ({(67, 20, 67, 22): '-1', (67, 24, 67, 25): '1', (67, 27, 67, 35): 'shape[1]'}, {}), '(-1, 1, shape[1])', True, 'import numpy as np\n'), ((68, 11, 68, 28), 'numpy.meshgrid', 'np.meshgrid', ({(68, 23, 68, 24): 'x', (68, 26, 68, 27): 'y'}, {}), '(x, y)', True, 'import numpy as np\n'), ((70, 13, 70, 31), 'numpy.sqrt', 'np.sqrt', ({(70, 21, 70, 30): 'x ** 2 + y ** 2'}, {}), '(x ** 2 + y ** 2)', True, 'import numpy as np\n'), ((79, 10, 79, 37), 'numpy.random.RandomState', 'np.random.RandomState', ({(79, 32, 79, 36): 'seed'}, {}), '(seed)', True, 'import numpy as np\n'), ((96, 8, 96, 30), 'numpy.linspace', 'np.linspace', ({(96, 20, 96, 21): '0', (96, 23, 96, 24): '1', (96, 26, 96, 29): '100'}, {}), '(0, 1, 100)', True, 'import numpy as np\n'), ((97, 10, 97, 31), 'numpy.linspace', 'np.linspace', ({(97, 22, 97, 23): '0', (97, 25, 97, 26): '1', (97, 28, 97, 30): '10'}, {}), '(0, 1, 10)', True, 'import numpy as np\n'), ((98, 10, 98, 31), 'numpy.linspace', 'np.linspace', ({(98, 22, 98, 23): '0', (98, 25, 98, 26): '1', (98, 28, 98, 30): '10'}, {}), '(0, 1, 10)', True, 'import numpy as np\n'), ((99, 15, 99, 36), 'numpy.meshgrid', 'np.meshgrid', ({(99, 27, 99, 30): 'xmw', (99, 32, 99, 35): 'ymw'}, {}), '(xmw, ymw)', True, 'import numpy as np\n'), ((101, 16, 101, 127), 'ipyvolume.plot_mesh', 'ipv.plot_mesh', (), '', True, 'import ipyvolume as ipv\n'), ((109, 4, 109, 20), 'ipyvolume.xyzlim', 'ipv.xyzlim', ({(109, 15, 109, 19): 'size'}, {}), '(size)', True, 'import ipyvolume as ipv\n'), ((16, 8, 16, 22), 'numpy.identity', 'np.identity', ({(16, 20, 16, 21): '(3)'}, {}), '(3)', True, 'import numpy as np\n'), ((17, 13, 17, 26), 'numpy.dot', 'np.dot', ({(17, 20, 17, 21): 'H', (17, 24, 17, 25): 'R'}, {}), '(H, R)', True, 'import numpy as np\n'), ((81, 8, 81, 35), 'numpy.sqrt', 'np.sqrt', ({(81, 16, 81, 34): '(x ** 2 + y ** 2 + z ** 2)'}, {}), '(x ** 2 + y ** 2 + z ** 2)', True, 'import numpy as np\n'), ((15, 82, 15, 95), 'numpy.sqrt', 'np.sqrt', ({(15, 90, 15, 94): '1 - u3'}, {}), '(1 - u3)', True, 'import numpy as np\n'), ((16, 27, 16, 44), 'numpy.transpose', 'np.transpose', ({(16, 40, 16, 43): '[v]'}, {}), '([v])', True, 'import numpy as np\n'), ((34, 8, 34, 35), 'numpy.repeat', 'np.repeat', ({(34, 18, 34, 25): 'orbit_x', (34, 27, 34, 34): 'N_stars'}, {}), '(orbit_x, N_stars)', True, 'import numpy as np\n'), ((35, 8, 35, 35), 'numpy.repeat', 'np.repeat', ({(35, 18, 35, 25): 'orbit_y', (35, 27, 35, 34): 'N_stars'}, {}), '(orbit_y, N_stars)', True, 'import numpy as np\n'), ((36, 8, 36, 35), 'numpy.repeat', 'np.repeat', ({(36, 18, 36, 25): 'orbit_z', (36, 27, 36, 34): 'N_stars'}, {}), '(orbit_z, N_stars)', True, 'import numpy as np\n'), ((42, 20, 42, 29), 'numpy.sin', 'np.sin', ({(42, 27, 42, 28): 'a'}, {}), '(a)', True, 'import numpy as np\n'), ((43, 20, 43, 29), 'numpy.cos', 'np.cos', ({(43, 27, 43, 28): 'a'}, {}), '(a)', True, 'import numpy as np\n'), ((71, 30, 71, 53), 'numpy.exp', 'np.exp', ({(71, 37, 71, 52): '(-radius ** 2 / s ** 2)'}, {}), '(-radius ** 2 / s ** 2)', True, 'import numpy as np\n'), ((14, 19, 14, 37), 'numpy.cos', 'np.cos', ({(14, 26, 14, 36): '2 * np.pi * u1'}, {}), '(2 * np.pi * u1)', True, 'import numpy as np\n'), ((14, 39, 14, 57), 'numpy.sin', 'np.sin', ({(14, 46, 14, 56): '2 * np.pi * u1'}, {}), '(2 * np.pi * u1)', True, 'import numpy as np\n'), ((14, 85, 14, 103), 'numpy.cos', 'np.cos', ({(14, 92, 14, 102): '2 * np.pi * u1'}, {}), '(2 * np.pi * u1)', True, 'import numpy as np\n'), ((15, 18, 15, 36), 'numpy.cos', 'np.cos', ({(15, 25, 15, 35): '2 * np.pi * u2'}, {}), '(2 * np.pi * u2)', True, 'import numpy as np\n'), ((15, 37, 15, 48), 'numpy.sqrt', 'np.sqrt', ({(15, 45, 15, 47): 'u3'}, {}), '(u3)', True, 'import numpy as np\n'), ((15, 50, 15, 68), 'numpy.sin', 'np.sin', ({(15, 57, 15, 67): '2 * np.pi * u2'}, {}), '(2 * np.pi * u2)', True, 'import numpy as np\n'), ((15, 69, 15, 80), 'numpy.sqrt', 'np.sqrt', ({(15, 77, 15, 79): 'u3'}, {}), '(u3)', True, 'import numpy as np\n'), ((81, 67, 81, 86), 'numpy.random.random', 'np.random.random', ({(81, 84, 81, 85): 'N'}, {}), '(N)', True, 'import numpy as np\n'), ((14, 65, 14, 83), 'numpy.sin', 'np.sin', ({(14, 72, 14, 82): '2 * np.pi * u1'}, {}), '(2 * np.pi * u1)', True, 'import numpy as np\n'), ((41, 12, 41, 41), 'numpy.linspace', 'np.linspace', ({(41, 24, 41, 25): '(0)', (41, 27, 41, 28): '(1)', (41, 30, 41, 40): 'x.shape[0]'}, {}), '(0, 1, x.shape[0])', True, 'import numpy as np\n')]