#required packages
import numpy as np
import matplotlib.pyplot as plt
import requests
import json
import h5py

def get(path, params=None):
    # make HTTP GET request to path
    headers = {"api-key":"2566d8dd2bcf9aefbb3d8b01080a877b"}
    r = requests.get(path, params=params, headers=headers)

    # raise exception if response code is not HTTP SUCCESS (200)
    r.raise_for_status()

    if r.headers['content-type'] == 'application/json':
        return r.json() # parse json responses automatically

    if 'content-disposition' in r.headers:
        filename = r.headers['content-disposition'].split("filename=")[1]
        with open(filename, 'wb') as f:
            f.write(r.content)
        return filename # return the filename string

    return r

def getsnapshot(id,num):
    params = {'dm':'Coordinates'}
    url = "http://www.illustris-project.org/api/Illustris-1/snapshots/" + str(num) + "/subhalos/" + str(id)
    subhalo = get(url)                                     # read the parameters of the subhalo in the json format
    return subhalo,get(url+"/cutout.hdf5", params=params)  # read the cutout of the subhalo into memory

nbins=20
maxradius=50
xp   = 0*np.ones((3,nbins))
xpsq = 0*xp
e    = 0*xp
count= 0*xp

id=1030 # 1030   75      543681
snapshots=[85,103,135]

for i in range(3):                             # loop through three snapshots/redshifts
    sub,filename=getsnapshot(id,snapshots[i])  # get json header and data file  
    with h5py.File(filename) as f:
        dx = f['PartType1']['Coordinates'][:,0] - sub['pos_x']
        dy = f['PartType1']['Coordinates'][:,1] - sub['pos_y']
        dz = f['PartType1']['Coordinates'][:,2] - sub['pos_z']
        
    rr = np.sqrt(dx**2 + dy**2 + dz**2)        # radii of all dm particles
  
    """ create a histogram with 20 bins
        num = counts
    and   x = bin boundaries
    are used """

    num,x=np.histogram(rr,bins=nbins,range=(0,maxradius))
    
    xp[i,:]=(x[1:]+x[:-1])/2                   # the middle of each bin
    xpsq[i,:]=(4*np.pi/3)*x[1:]**3-x[:-1]**3   # spherical shell volume
    count[i,:]=num                             # counts
    e[i,:]=np.sqrt(num)                        # Poisson error bar

from scipy.optimize import curve_fit
def dark(x,a,rs):
    return a/((x/rs)*(1+(x/rs))**4)

# example: z=0
x=xp[2,:]
y=count[2,:]/xpsq[2,:]
err=e[2,:]/xpsq[i,:]
a=0.1
rs=12

keep = err>0 
popt, pcov = curve_fit(dark, x[keep], y[keep], sigma=err[keep])

class Particle(object):
    
    def __init__(self, mass, x, y, z, vx, vy, vz):
        self.mass = mass
        self.x = x
        self.y = y
        self.z = z
        self.vx = vx
        self.vy = vy
        self.vz = vz
    
    def kinetic_energy(self):
        return 0.5 * self.mass * np.sqrt(self.vx ** 2 +
                                         self.vy ** 2 +
                                         self.vz ** 2)
    
    def distance(self, other):
        return np.sqrt((self.x - other.x) ** 2 +
                       (self.y - other.y) ** 2 +
                       (self.z - other.z) ** 2)
    
    def move(self, dt):
        self.x = self.x + self.vx * dt
        self.y = self.y + self.vy * dt
        self.z = self.z + self.vz * dt
        
class ChargedParticle(Particle):
    
    def __init__(self, mass, x, y, z, vx, vy, vz, charge):
        Particle.__init__(self, mass, x, y, z, vx, vy, vz)
        self.charge = charge

p1 = Particle(2., 1., 2., 3., 1., -3, -1)

p1.move(1.)
print(p1.x, p1.y, p1.z)

p1.move(1.)
print(p1.x, p1.y, p1.z)