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
[((201, 219), 'numpy.random.random', 'np.random.random', ([], {}), '()\n', (217, 219), True, 'import numpy as np\n'), ((229, 247), 'numpy.random.random', 'np.random.random', ([], {}), '()\n', (245, 247), True, 'import numpy as np\n'), ((257, 275), 'numpy.random.random', 'np.random.random', ([], {}), '()\n', (273, 275), True, 'import numpy as np\n'), ((1412, 1485), 'ipyvolume.plot', 'ipv.plot', (['orbit_x_line', 'orbit_y_line', 'orbit_z_line'], {'visible': 'orbit_visible'}), '(orbit_x_line, orbit_y_line, orbit_z_line, visible=orbit_visible)\n', (1420, 1485), True, 'import ipyvolume as ipv\n'), ((1680, 1733), 'numpy.random.normal', 'np.random.normal', (['(0)'], {'scale': 'sigma_r', 'size': '(3, N_stars)'}), '(0, scale=sigma_r, size=(3, N_stars))\n', (1696, 1733), True, 'import numpy as np\n'), ((1746, 1782), 'numpy.sqrt', 'np.sqrt', (['(xr ** 2 + yr ** 2 + zr ** 2)'], {}), '(xr ** 2 + yr ** 2 + zr ** 2)\n', (1753, 1782), True, 'import numpy as np\n'), ((2756, 2771), 'numpy.array', 'np.array', (['color'], {}), '(color)\n', (2764, 2771), True, 'import numpy as np\n'), ((2782, 2824), 'numpy.zeros', 'np.zeros', (['(shape[:2] + (4,))'], {'dtype': 'np.uint8'}), '(shape[:2] + (4,), dtype=np.uint8)\n', (2790, 2824), True, 'import numpy as np\n'), ((2833, 2861), 'numpy.linspace', 'np.linspace', (['(-1)', '(1)', 'shape[0]'], {}), '(-1, 1, shape[0])\n', (2844, 2861), True, 'import numpy as np\n'), ((2870, 2898), 'numpy.linspace', 'np.linspace', (['(-1)', '(1)', 'shape[1]'], {}), '(-1, 1, shape[1])\n', (2881, 2898), True, 'import numpy as np\n'), ((2910, 2927), 'numpy.meshgrid', 'np.meshgrid', (['x', 'y'], {}), '(x, y)\n', (2921, 2927), True, 'import numpy as np\n'), ((2953, 2977), 'numpy.sqrt', 'np.sqrt', (['(x ** 2 + y ** 2)'], {}), '(x ** 2 + y ** 2)\n', (2960, 2977), True, 'import numpy as np\n'), ((3296, 3323), 'numpy.random.RandomState', 'np.random.RandomState', (['seed'], {}), '(seed)\n', (3317, 3323), True, 'import numpy as np\n'), ((3872, 3894), 'numpy.linspace', 'np.linspace', (['(0)', '(1)', '(100)'], {}), '(0, 1, 100)\n', (3883, 3894), True, 'import numpy as np\n'), ((3905, 3926), 'numpy.linspace', 'np.linspace', (['(0)', '(1)', '(10)'], {}), '(0, 1, 10)\n', (3916, 3926), True, 'import numpy as np\n'), ((3937, 3958), 'numpy.linspace', 'np.linspace', (['(0)', '(1)', '(10)'], {}), '(0, 1, 10)\n', (3948, 3958), True, 'import numpy as np\n'), ((3974, 3995), 'numpy.meshgrid', 'np.meshgrid', (['xmw', 'ymw'], {}), '(xmw, ymw)\n', (3985, 3995), True, 'import numpy as np\n'), ((4037, 4163), 'ipyvolume.plot_mesh', 'ipv.plot_mesh', (['((xmw - 0.5) * rescale)', '((ymw - 0.5) * rescale + R_sun)', 'zmw'], {'u': 'xmw', 'v': 'ymw', 'texture': 'mw_image', 'wireframe': '(False)'}), '((xmw - 0.5) * rescale, (ymw - 0.5) * rescale + R_sun, zmw, u=\n    xmw, v=ymw, texture=mw_image, wireframe=False)\n', (4050, 4163), True, 'import ipyvolume as ipv\n'), ((4491, 4507), 'ipyvolume.xyzlim', 'ipv.xyzlim', (['size'], {}), '(size)\n', (4501, 4507), True, 'import ipyvolume as ipv\n'), ((503, 517), 'numpy.identity', 'np.identity', (['(3)'], {}), '(3)\n', (514, 517), True, 'import numpy as np\n'), ((553, 565), 'numpy.dot', 'np.dot', (['H', 'R'], {}), '(H, R)\n', (559, 565), True, 'import numpy as np\n'), ((3370, 3403), 'numpy.sqrt', 'np.sqrt', (['(x ** 2 + y ** 2 + z ** 2)'], {}), '(x ** 2 + y ** 2 + z ** 2)\n', (3377, 3403), True, 'import numpy as np\n'), ((479, 494), 'numpy.sqrt', 'np.sqrt', (['(1 - u3)'], {}), '(1 - u3)\n', (486, 494), True, 'import numpy as np\n'), ((522, 539), 'numpy.transpose', 'np.transpose', (['[v]'], {}), '([v])\n', (534, 539), True, 'import numpy as np\n'), ((1494, 1521), 'numpy.repeat', 'np.repeat', (['orbit_x', 'N_stars'], {}), '(orbit_x, N_stars)\n', (1503, 1521), True, 'import numpy as np\n'), ((1553, 1580), 'numpy.repeat', 'np.repeat', (['orbit_y', 'N_stars'], {}), '(orbit_y, N_stars)\n', (1562, 1580), True, 'import numpy as np\n'), ((1612, 1639), 'numpy.repeat', 'np.repeat', (['orbit_z', 'N_stars'], {}), '(orbit_z, N_stars)\n', (1621, 1639), True, 'import numpy as np\n'), ((1901, 1910), 'numpy.sin', 'np.sin', (['a'], {}), '(a)\n', (1907, 1910), True, 'import numpy as np\n'), ((1931, 1940), 'numpy.cos', 'np.cos', (['a'], {}), '(a)\n', (1937, 1940), True, 'import numpy as np\n'), ((3002, 3031), 'numpy.exp', 'np.exp', (['(-radius ** 2 / s ** 2)'], {}), '(-radius ** 2 / s ** 2)\n', (3008, 3031), True, 'import numpy as np\n'), ((295, 317), 'numpy.cos', 'np.cos', (['(2 * np.pi * u1)'], {}), '(2 * np.pi * u1)\n', (301, 317), True, 'import numpy as np\n'), ((315, 337), 'numpy.sin', 'np.sin', (['(2 * np.pi * u1)'], {}), '(2 * np.pi * u1)\n', (321, 337), True, 'import numpy as np\n'), ((361, 383), 'numpy.cos', 'np.cos', (['(2 * np.pi * u1)'], {}), '(2 * np.pi * u1)\n', (367, 383), True, 'import numpy as np\n'), ((415, 437), 'numpy.cos', 'np.cos', (['(2 * np.pi * u2)'], {}), '(2 * np.pi * u2)\n', (421, 437), True, 'import numpy as np\n'), ((434, 445), 'numpy.sqrt', 'np.sqrt', (['u3'], {}), '(u3)\n', (441, 445), True, 'import numpy as np\n'), ((447, 469), 'numpy.sin', 'np.sin', (['(2 * np.pi * u2)'], {}), '(2 * np.pi * u2)\n', (453, 469), True, 'import numpy as np\n'), ((466, 477), 'numpy.sqrt', 'np.sqrt', (['u3'], {}), '(u3)\n', (473, 477), True, 'import numpy as np\n'), ((3429, 3448), 'numpy.random.random', 'np.random.random', (['N'], {}), '(N)\n', (3445, 3448), True, 'import numpy as np\n'), ((341, 363), 'numpy.sin', 'np.sin', (['(2 * np.pi * u1)'], {}), '(2 * np.pi * u1)\n', (347, 363), True, 'import numpy as np\n'), ((1823, 1852), 'numpy.linspace', 'np.linspace', (['(0)', '(1)', 'x.shape[0]'], {}), '(0, 1, x.shape[0])\n', (1834, 1852), True, 'import numpy as np\n')]