from __future__ import division
import numpy as np 
import matplotlib.pyplot as plt 
import shellmodelutilities as smutil



# Set bin width and range
bin_width = 0.20
Emax = 14
Nbins = int(np.ceil(Emax/bin_width))
Emax_adjusted = bin_width*Nbins # Trick to get an integer number of bins
bins = np.linspace(0,Emax_adjusted,Nbins+1)

# Define list of calculation input files and corresponding label names
inputfile = "summary_Zn70_jun45.txt"

# Instantiate figure which we will fill
f_rho, ax_rho = plt.subplots(1,1)

# Read energy levels from file
levels = smutil.read_energy_levels(inputfile)

# Choose which [2*J,pi] combinations to include in partial level density plot
Jpi_list = [[0,-1],[2,-1],[4,-1],[6,-1],[8,-1],[10,-1],[12,-1],[14,-1],[16,-1],[18,-1],[20,-1],[22,-1],[24,-1],[26,-1],[28,-1],
            [0,+1],[2,+1],[4,+1],[6,+1],[8,+1],[10,+1],[12,+1],[14,+1],[16,+1],[18,+1],[20,+1],[22,+1],[24,+1],[26,+1],[28,+1]]

# Allocate (Ex,Jpi) matrix to store partial level density
rho_ExJpi = np.zeros((Nbins,len(Jpi_list)))
# Count number of levels for each (Ex, J, pi) pixel.
Egs = levels[0,0] # Ground state energy
for i_l in range(len(levels[:,0])):
    E, J, pi = levels[i_l]
    # Skip if level is outside range:
    if E-Egs >= Emax:
        continue
    i_Ex = int(np.floor((E-Egs)/bin_width))
    try:
        i_Jpi = Jpi_list.index([J,pi])
    except:
        continue
    rho_ExJpi[i_Ex,i_Jpi] += 1

rho_ExJpi /= bin_width # Normalize to bin width, to get density in MeV^-1


# Plot it
from matplotlib.colors import LogNorm # To get log scaling on the z axis
colorbar_object = ax_rho.pcolormesh(np.linspace(0,len(Jpi_list)-1,len(Jpi_list)), bins, rho_ExJpi, norm=LogNorm())
f_rho.colorbar(colorbar_object) # Add colorbar to plot

# Make the plot nice
ax_rho.set_xlabel(r"$\pi\cdot J\,\mathrm{(\hbar)}$")
ax_rho.set_ylabel(r'$E_x \, \mathrm{(MeV)}$')

# A bit of Python voodoo to get the x labels right:
Jpi_array = np.append(np.linspace(0,-int((len(Jpi_list)-1)/2),int(len(Jpi_list)/2)),np.linspace(0,int((len(Jpi_list)-1)/2),int(len(Jpi_list)/2))) # Array of pi*J for plot
def format_func(value, tick_number):
    if value >= 0 and value <= 28:
        return int(Jpi_array[int(value)])
    else:
        return None
ax_rho.set_xlim([0,29])
ax_rho.xaxis.set_major_formatter(plt.FuncFormatter(format_func))
ax_rho.set_xticks([0,2,4,6,8,10,12,14,15,17,19,21,23,25,27])

# Show plot
plt.show()

[((13, 7, 13, 43), 'numpy.linspace', 'np.linspace', ({(13, 19, 13, 20): '0', (13, 21, 13, 34): 'Emax_adjusted', (13, 35, 13, 42): 'Nbins + 1'}, {}), '(0, Emax_adjusted, Nbins + 1)', True, 'import numpy as np\n'), ((19, 16, 19, 33), 'matplotlib.pyplot.subplots', 'plt.subplots', ({(19, 29, 19, 30): '1', (19, 31, 19, 32): '1'}, {}), '(1, 1)', True, 'import matplotlib.pyplot as plt\n'), ((22, 9, 22, 45), 'shellmodelutilities.read_energy_levels', 'smutil.read_energy_levels', ({(22, 35, 22, 44): 'inputfile'}, {}), '(inputfile)', True, 'import shellmodelutilities as smutil\n'), ((68, 0, 68, 10), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((11, 12, 11, 35), 'numpy.ceil', 'np.ceil', ({(11, 20, 11, 34): 'Emax / bin_width'}, {}), '(Emax / bin_width)', True, 'import numpy as np\n'), ((64, 33, 64, 63), 'matplotlib.pyplot.FuncFormatter', 'plt.FuncFormatter', ({(64, 51, 64, 62): 'format_func'}, {}), '(format_func)', True, 'import matplotlib.pyplot as plt\n'), ((37, 15, 37, 42), 'numpy.floor', 'np.floor', ({(37, 24, 37, 41): '(E - Egs) / bin_width'}, {}), '((E - Egs) / bin_width)', True, 'import numpy as np\n'), ((49, 104, 49, 113), 'matplotlib.colors.LogNorm', 'LogNorm', ({}, {}), '()', False, 'from matplotlib.colors import LogNorm\n')]