from pNbody import *
from numpy import *
from numpy import random as RandomArray
import iclib



try:
  from scipy.integrate import quadrature
  is_scipy = True
except ImportError:
  is_scipy = False 





'''
# isotropic velocities
random2 = RandomArray.random([n])
random3 = RandomArray.random([n]) 
p = 2*pi*random2
costh = 1.-2.*random3			    
sinth = sqrt(1.-costh**2)
vx = v*sinth*cos(p)	     
vy = v*sinth*sin(p)				      
vz = v*costh 

# radial velocities
x = nb.pos[:,0] 
y = nb.pos[:,1]  
z = nb.pos[:,2] 
vx = v*x/r
vy = v*y/r
vz = v*z/r	    



Somes notes on the teneration of initial conditions:

Spherical case
--------------

1) if M(r) is known analitically :
  
  1.1) if it is invertible analitically
       --> for a given x between 0 and 1, find r from M(r)
  
  1.2) if it si not invertible anatitically
  
       
       1.2.1) use a Monte-Carlo approach
              (warning : may take time if the model is cuspy)
            
       1.2.2) invert M(r) numerically and create a vector of r and Mr
              then, use generic_Mr
	      
       1.2.3) for a given x between 0 and 1, find r from M(r),
              by inverting M(r). May take time.	      







'''

def invert(x,rmin,rmax,fct,args,eps=1e-10):
  '''
  return vector r that corresponds to
  fct(r,args)=x
  This routine uses a simple bissector algorithm
  '''
  n = len(x)
  rrmin = rmin*ones(n)
  rrmax = rmax*ones(n)
  xxmin = fct(rrmin,args) - x
  xxmax = fct(rrmax,args) - x

  if sum((xxmin*xxmax>=0))!=0:
    print "No max between rmin and rmax ! for some points"
    sys.exit()
  
  k = 0  
  
  while max(abs(rrmax-rrmin)) > eps:
    
    print "it = %3d err = %8.1e"%(k,max(abs(rrmax-rrmin)))
    
    k = k +1 
    rr = (rrmax+rrmin)/2.
    
    xx = fct(rr,args)-x  

    rrmax = where(xxmin*xx <= 0,rr,rrmax)
    rrmin = where(xxmin*xx >  0,rr,rrmin)
    xxmin = where(xxmin*xx >  0,xx,xxmin)

  return rr    


def box(n,a,b,c,name='box.dat',ftype='binary'):
  '''
  Return an Nbody object that contains n particules distributed
  in an homogeneous box of dimension a,b,c, centred at the origin
  radius rmax.
  '''
  
  if type(n) == ndarray:
    random = n
    n = len(n)
  else:
    random = RandomArray.random([n,3])  
  
  pos = random-[0.5,0.5,0.5]
  pos = pos*array([2*a,2*b,2*c])
  
  vel =  ones([n,3])*0.0
  mass = ones([n])*1./n
  
  nb = Nbody(status='new',p_name=name,pos=pos,vel=vel,mass=mass,ftype=ftype)

  return nb


def homodisk(n,a,b,dz,name='homodisk.dat',ftype='binary'):
  '''
  Return an Nbody object that contains n particules distributed
  in an homogeneous oval of radius a and b, and of thickness dz.
  '''

  if type(n) == ndarray:
    random1 = n[:,0]
    random2 = n[:,1]
    random3 = n[:,2]
    n = len(n)
  else:
    random1 = RandomArray.random([n])
    random2 = RandomArray.random([n])
    random3 = RandomArray.random([n]) 
    
  xx = random1**(1./2.)  
  
  theta = 2.* random2*pi  
  
  x = a* xx * cos(theta)
  y = b* xx * sin(theta) 
  z = dz*random3 - dz/2.
  
  pos = transpose(array([x,y,z]))
  vel =  ones([n,3])*0.0
  mass = ones([n])*1./n
  
  nb = Nbody(status='new',p_name=name,pos=pos,vel=vel,mass=mass,ftype=ftype)

  return nb


def homosphere(n,a,b,c,name='homosphere.dat',ftype='binary'):
  '''
  Return an Nbody object that contains n particules distributed
  in an homogeneous triaxial sphere of axis a,b,c.
  '''

  if type(n) == ndarray:
    random1 = n[:,0]
    random2 = n[:,1]
    random3 = n[:,2]
    n = len(n)
  else:
    random1 = RandomArray.random([n])
    random2 = RandomArray.random([n])
    random3 = RandomArray.random([n]) 
      
  xm = (random1)**(1./3.)
  phi = random2*pi*2.		 
  costh = 1.-2.*random3   
      
  	   
  sinth = sqrt(1.-costh**2) 	   
  axm = a*xm*sinth
  bxm = b*xm*sinth	     	   
  x = axm*cos(phi)	     	   
  y = bxm*sin(phi)	     	   
  z = c*xm*costh	     	   

  pos = transpose(array([x,y,z]))
  vel =  ones([n,3])*0.0
  mass = ones([n])*1./n
  
  nb = Nbody(status='new',p_name=name,pos=pos,vel=vel,mass=mass,ftype=ftype)

  return nb



def plummer(n,a,b,c,eps,rmax,M=1.,vel='no',name='plummer.dat',ftype='binary'):
  '''
  Return an Nbody object that contains n particules distributed
  in a triaxial plummer model of axis a,b,c and core radius eps
  and max radius of rmax.  
  
  rho = (1.+(r/eps)**2)**(-5/2)
  
  '''

  if type(n) == ndarray:
    random1 = n[:,0]
    random2 = n[:,1]
    random3 = n[:,2]
    n = len(n)
  else:
    random1 = RandomArray.random([n])
    random2 = RandomArray.random([n])
    random3 = RandomArray.random([n]) 
  
  # positions
  rmax = float(rmax)
  eps  = float(eps)
  
  eps = eps/rmax
  xm = 1./(1.+(eps)**2)*random1**(2./3.)    
  xm = eps*sqrt(xm/(1.-xm))	
  phi = 2*pi*random2
  
  costh = 1.-2.*random3			      
  sinth = sqrt(1.-costh**2)
  axm = rmax*a*xm*sinth
  bxm = rmax*b*xm*sinth
  x = axm*cos(phi)	     	   
  y = bxm*sin(phi) 					    
  z = rmax*c*xm*costh        
  
  pos = transpose(array([x,y,z]))
  # velocities
  if vel == 'yes':
    
    R  = sqrt(x**2+y**2)
    rho = (3.*M/(4.*pi*eps**3))*(1+(R**2+z**2)/eps**2)**(-5./2.)
    C2 = z**2+eps**2
    C  = sqrt(C2)
    
    TD = M*C/(R**2+C2)**(3./2.)
    sz = sqrt(eps**2/(8.*pi*C2)/rho) *TD
    
    vx = sz*RandomArray.standard_normal([n])
    vy = sz*RandomArray.standard_normal([n])
    vz = sz*RandomArray.standard_normal([n])
    vel = transpose(array([vx,vy,vz]))
   
  else: 
    vel = ones([n,3])*0.0
  
  # masses
  mass = ones([n])*M/n
  
  nb = Nbody(status='new',p_name=name,pos=pos,vel=vel,mass=mass,ftype=ftype)

  return nb
  
  
def kuzmin(n,eps,dz,name='kuzmin.dat',ftype='binary'):
  '''
  Return an Nbody object that contains n particules distributed
  in a kuzmin (infinitely thin) disk 
  
  rho = eps*M/(2*pi*(R**2+eps**2)**(3/2))
  
  '''

  if type(n) == ndarray:
    random1 = n[:,0]
    random2 = n[:,1]
    random3 = n[:,2]
    n = len(n)
  else:
    random1 = RandomArray.random([n])
    random2 = RandomArray.random([n])
    random3 = RandomArray.random([n]) 
      
  rmax = 1
  xx = random1
  xx = sqrt((eps/(1-xx))**2-eps**2)
  theta = 2.* random2*pi  
  
  x = rmax* xx * cos(theta)
  y = rmax* xx * sin(theta) 
  z = dz*random3 - dz/2.
  
  pos = transpose(array([x,y,z]))
  vel =  ones([n,3])*0.0
  mass = ones([n])*1./n

  nb = Nbody(status='new',p_name=name,pos=pos,vel=vel,mass=mass,ftype=ftype)

  return nb

def nfw_mr(r,args):
  '''
  Mass in the radius r for the nfw distribution
  
  rho = rhos/[ (r/rs)(1+r/rs)^2 ]
  
  mr  = 4*pi*rhos*rs**3 *(log(1+r/rs)-r/(rs+r))
  
  '''  
  rs      = args[0]
  rmax    = args[1]
    
  Mr = (log(1+r/rs)-r/(rs+r))
  Mx = (log(1+rmax/rs)-rmax/(rs+rmax))
  
  return Mr/Mx  
  

def nfw(n,a,b,c,rs,rmax,name='nfw.dat',ftype='binary'):
  '''
  Return an Nbody object that contains n particules following
  an nfw profile. 
  
  rho = 1/[ (r/rs)(1+r/rs)^2 ]
  
  '''

  if type(n) == ndarray:
    random1 = n[:,0]
    random2 = n[:,1]
    random3 = n[:,2]
    n = len(n)
  else:
    random1 = RandomArray.random([n])
    random2 = RandomArray.random([n])
    random3 = RandomArray.random([n]) 

  x = random1
  xm = invert(x,0,rmax,nfw_mr,[rs,rmax])    
  phi = 2*pi*random2
  
  costh = 1.-2.*random3		      
  sinth = sqrt(1.-costh**2)
  axm = a*xm*sinth
  bxm = b*xm*sinth
  x = axm*cos(phi)	     	   
  y = bxm*sin(phi) 					    
  z = c*xm*costh        
       
  
  pos = transpose(array([x,y,z]))
  vel =  ones([n,3])*0.0
  mass = ones([n])*1./n
  
  nb = Nbody(status='new',p_name=name,pos=pos,vel=vel,mass=mass,ftype=ftype)

  return nb



  

def dl2_mr(r,args):
  '''
  Mass in the radius r for the distribution
  
  rho = (1.-eps*(r/rmax)**2)
  '''  
  eps    = args[0]
  rmax = args[1]
  return ((4./3.)*r**3 - (4./5.)*eps*r**5/rmax**2)/(((4./3.) - (4./5.)*eps)*rmax**3)


def dl2(n,a,b,c,eps,rmax,name='dl2.dat',ftype='binary'):
  '''
  Return an Nbody object that contains n particules distributed as
  
  rho = (1.-eps*(r/rmax)**2)
  '''

  if type(n) == ndarray:
    random1 = n[:,0]
    random2 = n[:,1]
    random3 = n[:,2]
    n = len(n)
  else:
    random1 = RandomArray.random([n])
    random2 = RandomArray.random([n])
    random3 = RandomArray.random([n]) 
      
  x = random1
  xm = invert(x,0,rmax,dl2_mr,[eps,rmax])    
  phi = 2*pi*random2
  
  costh = 1.-2.*random3		      
  sinth = sqrt(1.-costh**2)
  axm = a*xm*sinth
  bxm = b*xm*sinth
  x = axm*cos(phi)	     	   
  y = bxm*sin(phi) 					    
  z = c*xm*costh        
       
  
  pos = transpose(array([x,y,z]))
  vel =  ones([n,3])*0.0
  mass = ones([n])*1./n
  
  nb = Nbody(status='new',p_name=name,pos=pos,vel=vel,mass=mass,ftype=ftype)

  return nb

def isothm_mr(r,args):
  '''
  Mass in the radius r for the distribution
  
  rho = 1/(1+r/rc)**2
  '''  
  rc   = args[0]
  rm   = args[1]
  
  cte = 2*rc**3*log(rc) + rc**3
  
  Mr = r *rc**2 - 2*rc**3*log(rc+r ) - rc**4/(rc+r ) + cte
  Mx = rm*rc**2 - 2*rc**3*log(rc+rm) - rc**4/(rc+rm) + cte
  
  return Mr/Mx


def isothm(n,rc,rmax,name='isothm.dat',ftype='binary'):
  '''
  Return an Nbody object that contains n particules distributed as

  rho = 1/(1+r/rc)**2
  '''

  if type(n) == ndarray:
    random1 = n[:,0]
    random2 = n[:,1]
    random3 = n[:,2]
    n = len(n)
  else:
    random1 = RandomArray.random([n])
    random2 = RandomArray.random([n])
    random3 = RandomArray.random([n]) 
      
  x = random1
  xm = invert(x,0,rmax,isothm_mr,[rc,rmax])    
  phi = 2*pi*random2
  
  costh = 1.-2.*random3		      
  sinth = sqrt(1.-costh**2)
  axm = xm*sinth
  bxm = xm*sinth
  x = axm*cos(phi)	     	   
  y = bxm*sin(phi) 					    
  z = xm*costh        
       
  
  pos = transpose(array([x,y,z]))
  vel =  ones([n,3])*0.0
  mass = ones([n])*1./n
  
  nb = Nbody(status='new',p_name=name,pos=pos,vel=vel,mass=mass,ftype=ftype)

  return nb


def pisothm_mr(r,args):
  '''
  Mass in the radius r for the distribution
  
  rho = 1/(1+(r/rc)**2)
  '''  
  rc    = args[0]
  rmn   = args[1]
  rmx   = args[2] 
 
  Mr  = rc**3 *(   r/rc - arctan( r/rc) )
  Mmn = rc**3 *( rmn/rc - arctan(rmn/rc) )
  Mmx = rc**3 *( rmx/rc - arctan(rmx/rc) )

  return (Mr-Mmn)/(Mmx-Mmn)


def pisothm(n,rc,rmax,rmin=0,name='pisothm.dat',ftype='binary'):
  '''
  Pseudo-isothermal sphere
  Mass in the radius r for the distribution

  rho = 1/(1+(r/rc)**2)
  '''

  if type(n) == ndarray:
    random1 = n[:,0]
    random2 = n[:,1]
    random3 = n[:,2]
    n = len(n)
  else:
    random1 = RandomArray.random([n])
    random2 = RandomArray.random([n])
    random3 = RandomArray.random([n])

  x = random1
  xm = invert(x,rmin,rmax,pisothm_mr,[rc,rmin,rmax])
  phi = 2*pi*random2

  costh = 1.-2.*random3
  sinth = sqrt(1.-costh**2)
  axm = xm*sinth
  bxm = xm*sinth
  x = axm*cos(phi)
  y = bxm*sin(phi)
  z = xm*costh


  pos = transpose(array([x,y,z]))
  vel =  ones([n,3])*0.0
  mass = ones([n])*1./n

  nb = Nbody(status='new',p_name=name,pos=pos,vel=vel,mass=mass,ftype=ftype)

  return nb

 


def shell(n,r,name='cell.dat',ftype='binary'):
  '''
  Shell of radius r
  '''


  if type(n) == ndarray:
    random2 = n[:,1]
    random3 = n[:,2]
    n = len(n)
  else:
    random2 = RandomArray.random([n])
    random3 = RandomArray.random([n]) 
    

  phi = random2*pi*2.		 
  costh = 1.-2.*random3 	   
    
  sinth = sqrt(1.-costh**2)	   

  x = r*sinth*cos(phi) 	   
  y = r*sinth*sin(phi) 	   
  z = r*costh   


  pos = transpose(array([x,y,z]))
  vel =  ones([n,3])*0.0
  mass = ones([n])*1./n
  
  nb = Nbody(status='new',p_name=name,pos=pos,vel=vel,mass=mass,ftype=ftype)

  return nb


def expd_mr(r,args):
  '''
  Mass in the radius r for the distribution
  
  Sigma = epx(-r)
  '''  
  
  rm   = args[0]
  
  Mr = 1-(1+r)*exp(-r)
  Mx = 1-(1+rm)*exp(-rm)
  
  return Mr/Mx
  

def expd(n,Hr,Hz,Rmax,Zmax,irand=1,name='expd.dat',ftype='binary'):
  '''
  Exonential disk
  
  rho = 1/(1+(r/rc)**2)
  '''  
  
  '''
  if type(n) == ndarray:
    random1 = n[:,0]
    random2 = n[:,1]
    random3 = n[:,2]
    n = len(n)
  else:
    random1 = RandomArray.random([n])
    random2 = RandomArray.random([n])
    random3 = RandomArray.random([n]) 
      
  x = random1
  xm = invert(x,0,rmax,expd_mr,[rmax]) 
  
  
  phi = 2*pi*random2
  
  Hr = 4
  axm = xm * Hr
  bxm = xm * Hr
  
  x = axm*cos(phi)	     	   
  y = bxm*sin(phi) 	
  
  				    
  z = zeros(n)   
  
  pos = transpose(array([x,y,z]))

  '''
  
  pos = iclib.exponential_disk(n,Hr,Hz,Rmax,Zmax,irand)
       
  
  vel =  ones([n,3])*0.0
  mass = ones([n])*1./n
  
  nb = Nbody(status='new',p_name=name,pos=pos,vel=vel,mass=mass,ftype=ftype)

  return nb



def miyamoto_nagai(n,a,b,Rmax,Zmax,irand=1,fct=None,fRmax=0,name='miyamoto.dat',ftype='binary'):
  '''
  Miyamoto Nagai distribution
  '''
  
  if fct==None:
    pos = iclib.miyamoto_nagai(n,a,b,Rmax,Zmax,irand)
  else:
    pos = iclib.miyamoto_nagai_f(n,a,b,Rmax,Zmax,irand,fct,fRmax)  
       
  vel =  ones([n,3])*0.0
  mass = ones([n])*1./n
  
  nb = Nbody(status='new',p_name=name,pos=pos,vel=vel,mass=mass,ftype=ftype)

  return nb

def burkert(n,rs,Rmax,irand=1,name='burkert.dat',ftype='binary'):
  '''
  Burkert distribution
  '''
  pos = iclib.burkert(n,rs,Rmax,irand)  
       
  vel =  ones([n,3])*0.0
  mass = ones([n])*1./n
  
  nb = Nbody(status='new',p_name=name,pos=pos,vel=vel,mass=mass,ftype=ftype)

  return nb


def nfwg(n,rs,gamma,Rmax,dR,Rs=None,irand=1,name='nfwg.dat',ftype='binary'):
  '''
  Return an Nbody object that contains n particules following
  an nfw modifed profile. 
  
  rho = 1/[ ((r/rs))**(gamma)(1+r/rs)^2 ]**(0.5*(3-gamma))
  
  '''
  
  if not is_scipy:
    raise "module scipy needed for function nfwg !"
  
  def rho(r,rs,gamma):
    return 1/(  (r/rs)**(gamma)  * (1+(r/rs)**2)**(0.5*(3.-gamma))  )    

  def rhor2(r,rs,gamma):
    return 4*pi*r**2 /(  (r/rs)**(gamma)  * (1+(r/rs)**2)**(0.5*(3.-gamma))  ) 

  def Mr(r,rs,gamma):
    return quadrature(rhor2, 0, r, args=(rs,gamma), tol=1.5e-10, maxiter=50)
   
  
  if Rs==None: 
    Rs  = arange(0,Rmax+dR,dR)		  # should use a non linear grid
  
  Mrs = zeros(len(Rs))
  ntot = len(Rs)

  for i in xrange(len(Rs)):
    Mrs[i] =  Mr(Rs[i],rs,gamma)[0]
    print Rs[i],Mrs[i],i,'/',ntot

  # normalisation
  Mrs = Mrs/Mrs[-1]
  Mrs[0] = 0

  
  # now use Mr
  nb = generic_Mr(n,rmax=Rmax,irand=irand,R=Rs,Mr=Mrs,name=name,ftype=ftype)


  return nb






def generic_alpha(n,a,e,rmax,irand=1,fct=None,name='generic_alpha.dat',ftype='binary'):
  '''
  Generic alpha distribution : rho ~ (r+e)^a
  '''
  
  pos = iclib.generic_alpha(n,a,e,rmax,irand)
       
  vel =  ones([n,3])*0.0
  mass = ones([n])*1./n
  
  nb = Nbody(status='new',p_name=name,pos=pos,vel=vel,mass=mass,ftype=ftype)

  return nb




def generic_Mr(n,rmax,irand=1,R=None,Mr=None,name='sphere_Mr.dat',ftype='binary'):
  '''
  Distribute particles in order to reproduce M(R) given by Mr
  '''
  
  if type(n) == ndarray:
    random1 = n[:,0]
    random2 = n[:,1]
    random3 = n[:,2]
    n = len(n)
  else:
    random1 = RandomArray.random([n])
    random2 = RandomArray.random([n])
    random3 = RandomArray.random([n])
        
  pos = iclib.generic_Mr(n,rmax,R.astype(float32),Mr.astype(float32),random1.astype(float32),random2.astype(float32),random3.astype(float32))
       
  vel =  ones([n,3])*0.0
  mass = ones([n])*1./n
  
  nb = Nbody(status='new',p_name=name,pos=pos,vel=vel,mass=mass,ftype=ftype)

  return nb










