#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Feb 23 21:55:05 2020

@author: Hanshen Su
"""

#import needed libs

import numpy.random as rnd
import numpy as np

import scipy
from scipy import constants

import matplotlib.pyplot as plt
import matplotlib.cm
from matplotlib.patches import Rectangle

from itertools import cycle
from mpl_toolkits.mplot3d import Axes3D

import os
from pprint import pprint
import shutil
import subprocess
import urllib.request

import h5py

import openmc.data

colors = cycle("bgrcmykbgrcmykbgrcmykbgrcmyk")



""" Define Constants"""
neuron_mass = constants.neutron_mass
u235_mass = openmc.data.atomic_mass('U235')
uranium_density = 19.1
avogadro =constants.Avogadro



initial_energy = 1000000
radius_reactor = 0.02866218
# = 0.08

# import data from library


u235 = openmc.data.IncidentNeutron.from_endf("/Users/a1234/Database/ENDF-B-VIII.0/neutrons/n-092_U_235.endf")

u235_total = u235[1] # I don't know why, this total cross-section is wrong
                     # use the sum of other reaction instead
u235_elastic = u235[2]
u235_level = u235[4]
u235_misc = u235[5]
u235_n2n = u235[16]
u235_n3n = u235[17]
u235_fission = u235[18]
u235_ngamma = u235[102]

u235_Q = u235_level.q_value


""" Define Physical Functions"""


def macroscopic_cross_section(energy):
    '''
    This function defines the total macroscopic cross section as function of energy
    note that the macroscopic corss section is defined as number of atoms per 
    volume multiply by the microscopic corss section
    '''
    microscopic_cross_section = u235_elastic.xs['0K'](energy) \
                                + u235_level.xs['0K'](energy) \
                                + u235_fission.xs['0K'](energy) \
                                + u235_ngamma.xs['0K'](energy)
    atom_numbers = (uranium_density * avogadro/u235_mass)*1E6#from per cm^3 to m^3
    return atom_numbers*microscopic_cross_section*10**(-28)#convert from barn to m^2

def mean_freepath(energy):
    return 1/macroscopic_cross_section(energy)
    
# Calculate speed from energy
def speed(energy,mass):
    '''
    calculate speed under non-relativisitc limit
    '''
    return (2 * energy / mass)**(1/2)

def length_distribution(energy):
    '''
    Apply inverser CDF method so that lenth of travel is exponentially distributed
    '''
    lamda = mean_freepath(macroscopic_cross_section(energy))
    return -lamda*np.log(rnd.uniform(0,1))

def watts_distribution():
    '''
    energy spctra of neutron in fission reaction
    '''
    a = 1
    b = 2
    K = 1 + (b/(8*a))
    L = 1/a*(K+(K**2-1)**(1/2))
    M = a*L - 1
    p1 = rnd.uniform(0,1)
    p2 = rnd.uniform(0,1)
    x = -np.log(p1)
    y = -np.log(p2)

    if (y-M*(x+1))**2 <= b*L*x:
        return 1000000*L*x
    else:
        return watts_distribution()

""" Define Main Functions """
#
def unit_vector():

    """
    gives unit vector in arbitary direction
    depend on the shape of reactor
    sphere
    """
    sample = rnd.uniform(-1,1,size=3)
    if np.linalg.norm(sample)<=1:
        return sample/np.linalg.norm(sample)
    else:
        return unit_vector()
    
# Spawn a neutron 
def spawn_neutron(R):
    """
    spawn a neutron within the reactor
    depend on the shape of reactor(R the radius of reactor)
    sphere
    """
    sample = rnd.uniform(-1,1,size=3)*R
    if np.linalg.norm(sample)<=R:
        return sample
    else:
        return spawn_neutron(R)

def spawn_initial_neutrons(N,R):
    '''
    this function generates a group of neutron with 
    random position in the reactor (spawn_neutron function)
    and with random unit vector (unit_vector function) 
    '''
    neutron_seq = np.empty((0,7))
    for i in range(N):
        energy = watts_distribution()
        neutron = np.hstack((spawn_neutron(R),np.array([energy]),unit_vector()))
        neutron_seq = np.vstack((neutron_seq,neutron))
    return neutron_seq

def u235_event(energy):
    '''
    gives chance of undergoing elastic scattering, excitation, etc, 
    by calculating microscopic cross-section of dividing an individual 
    event by the total cross-section 
    '''
    
    chance_elastic = u235_elastic.xs["0K"](energy)
    chance_level = u235_level.xs["0K"](energy)
    chance_misc = u235_misc.xs["0K"](energy)#misc?
    chance_n2n = u235_n2n.xs["0K"](energy)#n2n?
    chance_n3n = u235_n3n.xs["0K"](energy)#n3n?
    chance_fission = u235_fission.xs["0K"](energy)
    #chance_total = u235_total.xs["0K"](energy)
    chance_gamma = u235_ngamma.xs["0K"](energy)
    chance_real_total = chance_elastic+chance_level+chance_misc+chance_n2n+\
                        chance_n3n+chance_fission+chance_gamma
                            
    return np.array([chance_elastic/chance_real_total,
                     chance_level/chance_real_total,
                     chance_misc/chance_real_total,
                     chance_n2n/chance_real_total,
                     chance_n3n/chance_real_total,
                     chance_fission/chance_real_total,
                     chance_gamma/chance_real_total])
    
def u235_event_rnd(energy):
    '''
    generates event outcome with their probabilities given previously
    '''
    
    outcome = ["elastic","level","misc","n2n","n3n","fission","ngamma"]
    return rnd.choice(outcome,1,p=u235_event(energy))[0]

def u235_elastic_angle(vector):
    '''
    gives unit vector after elastic scattering in the lab frame
    '''

    x = vector[0]
    y = vector[1]
    z = vector[2]
    coscm = rnd.uniform(-1,1)
    A = 235
    coslab = (1 + A*coscm) * (1+A**2+2*A*coscm)**(-1/2)
    
    chi = rnd.uniform(0,1) * 2*np.pi
    cosx = np.cos(chi)
    sinx = np.sin(chi)
    a = (1-z**2)**(1/2)
    z1 = z*coslab + a*cosx
    b = coslab - z*z1
    y1 = 1/(1-z**2) * (y*b + x*a*sinx)
    x1 = 1/(1-z**2) * (x*b - y*a*sinx)
    return np.array([x1,y1,z1])/np.linalg.norm([x1,y1,z1])    
    
def u235_elastic_energy(energy):
    '''
    gived energy after collision in lab
    '''
    A = 235
    a = (A-1)/(A+1)
    
    return 1/2*energy*((1-a**2)*rnd.uniform(-1,1)+1+a**2)

def u235_inelastic_angle_plus(vector,energy):
    x = vector[0]
    y = vector[1]
    z = vector[2]
    coscm = rnd.uniform(-1,1)
    A = 235
    Q = u235_Q#U235 is excited from ground state
    g = (A**2 + A*(A+1)*Q/energy)**(-1/2)
    coslab = (g + coscm) * (1+g**2+2*g*coscm)**(-1/2)
    
    chi = rnd.uniform(0,1) * 2*np.pi
    cosx = np.cos(chi)
    sinx = np.sin(chi)
    a = (1-z**2)**(1/2)
    z1 = z*coslab + a*cosx
    b = coslab - z*z1
    y1 = 1/(1-z**2) * (y*b + x*a*sinx)
    x1 = 1/(1-z**2) * (x*b - y*a*sinx)
    return np.array([x1,y1,z1])/np.linalg.norm([x1,y1,z1]),(A+1)**(-2)*(coslab*(energy)**(1/2)+\
                                            (energy*(coslab**2+A**2-1) + A*(A+1)*Q)**(1/2))**2 

def u235_fission_new(position):
    release = rnd.choice([2,3],p=(0.6,0.4))#gives expectation of neutron release during fission
    neutron_seq = np.empty((0,7))
    for i in range(release):
        neutron = np.hstack((position,np.array([watts_distribution()]),unit_vector()))
        neutron_seq = np.vstack((neutron_seq,neutron))
    return neutron_seq#add to the initial neutron list?

def inside(position):
    '''
    decide whehter neutron is inside the reactor or not,
    if outside, delete immediately
    '''
    if np.linalg.norm(position)<radius_reactor:
        return True
    else:
        return False
 
# Cycle zero

#generate 1000 initial neutrons
initial_neutrions = spawn_initial_neutrons(1000,radius_reactor)

""" Define Main Functions """
def random_walk(n,neutrions_seq):

    next_seq = np.empty((0,7))
    counter = 0
    for neutron in neutrions_seq:
        
        alive = True
        direction = neutron[4:]
        energy = neutron[3]
        position = neutron[:3]
        
        while alive == True:
            #direction = neutron[4:]
            #energy = neutron[3]
            #position = neutron[:3]
            
            move = direction * length_distribution(energy)
            
            position = position + move
            
            dice = u235_event_rnd(energy)
            counter +=1
            
            if inside(position)==False:
                alive = False
            
            elif dice == "elastic":
                energy = u235_elastic_energy(energy)
                direction = u235_elastic_angle(direction)                
                
            elif dice == "level":
                if energy > -u235_Q:
                    direction, energy = u235_inelastic_angle_plus(direction,energy)
                else:
                    alive = False
                    #energy = u235_elastic_energy(energy)
                    #direction = u235_elastic_angle(direction) 
                
                """pass"""    
            elif dice == "misc":
                energy = u235_elastic_energy(energy)
                direction = u235_elastic_angle(direction)
            elif dice == "n2n":
                energy = u235_elastic_energy(energy)
                direction = u235_elastic_angle(direction)
            elif dice == "n3n":
                energy = u235_elastic_energy(energy)
                direction = u235_elastic_angle(direction)
                """pass"""
                
            elif dice == "fission":
                alive = False
                
                new_neutron = u235_fission_new(position)
                next_seq = np.vstack((next_seq,new_neutron))#use previous function to add new neutron
                
            elif dice == "ngamma":
                alive = False
        
    return next_seq
            

def cycle_0(n,initial_neutrions_0):#iterations
    x = []

    for i in range(n):
        if i == 0:
            sample = random_walk(0,initial_neutrions_0)[:1000]
            x.append(sample)
        else:
            
            sample = random_walk(0,x[i-1])[:1000]
            x.append(sample)
    return x

#visulaize results


test = cycle_0(1,initial_neutrions)
source_distribution = test[-1]
path1 = test[0][:,:3]
path2 = test[-1][:,:3]


# a=random_walk(0,initial_neutrions)
# a1=random_walk(0,a)
# a2=random_walk(0,a1)
# a3=random_walk(0,a2)
# a4=random_walk(0,a3)


# print()
#%%
path1_dis = []
for i in path1:
    path1_dis.append(np.linalg.norm(i))

path2_dis = []
for i in path2:
    path2_dis.append(np.linalg.norm(i))

fig1, ax1 = plt.subplots(1, 1)


ax1.hist(path1_dis, 50, density=True, alpha=0.75)

# x = 0.5*(bins[1:]+ bins[:-1])

# ax2.plot(x,y)

# ax2.legend(loc="best", shadow=True, fontsize='large')
plt.show()

fig1, ax1 = plt.subplots(1, 1)


ax1.hist(path2_dis, 50, density=True, alpha=0.75)

plt.show


#%%
"""
path = random_walk(0,100)
start = path[:1]
stop = path[-1:]
"""
fig = plt.figure(figsize=(20,20))
ax = Axes3D(fig)
ax.grid(False)
ax.xaxis.pane.fill = ax.yaxis.pane.fill = ax.zaxis.pane.fill = False
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")





ax.scatter3D(path1[:,0], path1[:,1], path1[:,2], 
            c="blue", alpha=0.25,s=1)
#ax.plot3D(path[:,0], path[:,1], path[:,2], 
#          c="blue", alpha=0.5, lw=0.5)
#ax.plot3D(start[:,0], start[:,1], start[:,2], 
#          c="red", marker="+")
#ax.plot3D(stop[:,0], stop[:,1], stop[:,2], 
#          c="black", marker="o")
#plt.title("3D Random Walk")



plt.show()
fig = plt.figure(figsize=(20,20))
ax = Axes3D(fig)
ax.grid(False)
ax.xaxis.pane.fill = ax.yaxis.pane.fill = ax.zaxis.pane.fill = False
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")





ax.scatter3D(path2[:,0], path2[:,1], path2[:,2], 
            c="blue", alpha=0.25,s=1)

plt.show()



# for i, col in zip(np.arange(10), colors):
#     # Simulate steps in 3D
#     path = random_walk(0,10)
#     start = path[:1]
#     stop = path[-1:]
#     # Plot the path
#     ax.scatter3D(path[:,0], path[:,1], path[:,2],
#                  c=col,alpha=0.15,s=1);
#     ax.plot3D(path[:,0], path[:,1], path[:,2], 
#               c=col, alpha=0.25,lw=0.25)
#     ax.plot3D(start[:,0], start[:,1], start[:,2],
#               c=col, marker="+")
#     ax.plot3D(stop[:,0], stop[:,1], stop[:,2],
#               c=col, marker="o");

# plt.title("3D Random Walk - Multiple runs")
# plt.savefig("random walk.svg")





































