import numpy as np
import matplotlib.pyplot as plt
import sphviewer

n1 = 10000 #number of particles to make the disk
n2 = n1/3  # number of particles to make the background
r  = np.random.rand(n1)
phi = 2*np.pi*np.random.rand(n1)
pos = np.zeros([3,n1], dtype=np.float32)
pos[0,:] = r*np.cos(phi)
pos[1,:] = r*np.sin(phi)
pos[2,:] = 0.1*np.random.rand(n1)

background = -2+4*np.random.rand(3,n2)
pos = np.concatenate((pos,background),axis=1)

fig = plt.figure(1, figsize=(10,5))
ax1  = fig.add_subplot(121)
ax2  = fig.add_subplot(122)
ax1.plot(pos[0,:], pos[1,:], 'k.')
ax2.plot(pos[0,:], pos[2,:], 'k.')


mass = np.ones(n1+n2)

Particles = sphviewer.Particles(pos,mass)

Scene = sphviewer.Scene(Particles)

Render = sphviewer.Render(Scene)

img = Render.get_image()
extent = Render.get_extent()
fig = plt.figure(1,figsize=(5,5))
ax1 = fig.add_subplot(111)
ax1.imshow(img, extent=extent, origin='lower', cmap='hot')
ax1.set_xlabel('X', size=15)
ax1.set_ylabel('Y', size=15)

Render.set_logscale()
img = Render.get_image()
extent = Render.get_extent()
fig = plt.figure(1,figsize=(5,5))
ax1 = fig.add_subplot(111)
ax1.imshow(img, extent=extent, origin='lower', cmap='hot', vmax=1.5)
ax1.set_xlabel('X', size=15)
ax1.set_ylabel('Y', size=15)

camera = Scene.Camera.get_params()
print camera['r']

Scene.update_camera(r=2.00)
Render = sphviewer.Render(Scene)
Render.set_logscale()
img = Render.get_image()
extent = Render.get_extent()
fig = plt.figure(1,figsize=(5,5))
ax1 = fig.add_subplot(111)
ax1.imshow(img, extent=extent, origin='lower', cmap='hot', vmax=0.5)
ax1.set_xlabel('X', size=15)
ax1.set_ylabel('Y', size=15)

Scene.update_camera(r='infinity')
Render = sphviewer.Render(Scene)
Render.set_logscale()
img = Render.get_image()
extent = Render.get_extent()
fig = plt.figure(1,figsize=(5,5))
ax1 = fig.add_subplot(111)
ax1.imshow(img, extent=extent, origin='lower', cmap='hot', vmax=0.5)
ax1.set_xlabel('X', size=15)
ax1.set_ylabel('Y', size=15)

Scene.update_camera(r='infinity')
Render = sphviewer.Render(Scene)
Render.set_logscale()
img1 = Render.get_image()
extent1 = Render.get_extent()

Scene.update_camera(r='infinity', extent=[-1,1,-1,1])
Render = sphviewer.Render(Scene)
Render.set_logscale()
img2 = Render.get_image()
extent2 = Render.get_extent()

Scene.update_camera(r='infinity', extent=[-4,4,-4,4])
Render = sphviewer.Render(Scene)
Render.set_logscale()
img3 = Render.get_image()
extent3 = Render.get_extent()

fig = plt.figure(1,figsize=(15,5))
ax1 = fig.add_subplot(131)
ax2 = fig.add_subplot(132)
ax3 = fig.add_subplot(133)

ax1.imshow(img1, extent=extent1, origin='lower', cmap='hot', vmax=0.5)
ax1.set_title('Automatic extent')
ax1.set_xlabel('X', size=15)
ax1.set_ylabel('Y', size=15)

ax2.imshow(img2, extent=extent2, origin='lower', cmap='hot', vmax=0.5)
ax2.set_title('extent=[-1,1,-1,1]')
ax2.set_xlabel('X', size=15)
ax2.set_ylabel('Y', size=15)

ax3.imshow(img3, extent=extent3, origin='lower', cmap='hot', vmax=0.5)
ax3.set_title('extent=[-4,4,-4,4]')
ax3.set_xlabel('X', size=15)
ax3.set_ylabel('Y', size=15)

fig = plt.figure(1,figsize=(15,5))
t = [-90,-45, 0]
for i in xrange(1,4):
    Scene.update_camera(r=2., t=t[i-1], extent=[-2,2,-2,2])
    Render = sphviewer.Render(Scene)
    Render.set_logscale()
    img = Render.get_image()
    extent = Render.get_extent()
    ax = fig.add_subplot(1,3,i)
    ax.imshow(img, extent=extent, origin='lower', cmap='hot', vmax=0.5)
    ax.set_xlabel('X', size=15)
    ax.set_ylabel('Y', size=15)

fig = plt.figure(1,figsize=(15,5))
ax1 = fig.add_subplot(131)
ax2 = fig.add_subplot(132)
ax3 = fig.add_subplot(133)
Particles.plot('xy',axis=ax1, color='r', markersize=1.)
Particles.plot('xz',axis=ax2, color='b', markersize=1.)
Particles.plot('yz',axis=ax3, color='g', markersize=1.)

fig = plt.figure(1,figsize=(5,5))
ax1 = fig.add_subplot(111)
Render.histogram(bins=100, log=True)

fig = plt.figure(1,figsize=(15,5))
plane = ['xy','xz','yz']
for i in xrange(3):
    ax = fig.add_subplot(1,3,i)
    Particles .plot(plane[i], axis=ax, markersize=0.5)
    Scene .update_camera(r=1.5)
    Scene.Camera.plot(plane[i], axis=ax)
plt.show()