#!/usr/bin/python3

import numpy
import pandas
import matplotlib.pyplot
from cycler import cycler

matplotlib.pyplot.ion()

cmap=[(0.10588235294117647, 0.6196078431372549, 0.4666666666666667), (0.8509803921568627, 0.37254901960784315, 0.00784313725490196), (0.4588235294117647, 0.4392156862745098, 0.7019607843137254), (0.9058823529411765, 0.1607843137254902, 0.5411764705882353), (0.4, 0.6509803921568628, 0.11764705882352941), (0.9019607843137255, 0.6705882352941176, 0.00784313725490196), (0.6509803921568628, 0.4627450980392157, 0.11372549019607843), (0.4, 0.4, 0.4)]

amps_pos = pandas.read_csv("prm_amplitudes_pos.dat", sep="\s+")
amps_neg = pandas.read_csv("prm_amplitudes_neg.dat", sep="\s+")

fig,ax=matplotlib.pyplot.subplots(2,1,figsize=(10,8))
ax[0].set_prop_cycle(cycler(color=cmap))
ax[1].set_prop_cycle(cycler(color=cmap))

nstates = len(amps_pos.index.values)
for i,state in enumerate(amps_pos.index.values):
    amplitudes = numpy.array(amps_pos.loc[state,:][1:])
    if (abs(amplitudes[0])>0):
        sgn = abs(amplitudes[0])/amplitudes[0]
    else:
        sgn = 1
    amplitudes = [x*sgn for x in amplitudes]
    width = 1./(nstates+4.)
    xoff = i*width - width*nstates/2.
    x = numpy.array(amps_pos.columns[1:],dtype=float) + xoff
    ax[0].bar(x, amplitudes, width, label="$"+state+"^{+}$")


nstates = len(amps_neg.index.values)    
for i,state in enumerate(amps_neg.index.values):
    amplitudes = numpy.array(amps_neg.loc[state,:][1:])
    if (abs(amplitudes[0])>0):
        sgn = abs(amplitudes[0])/amplitudes[0]
    else:
        sgn = 1
    amplitudes = [x*sgn for x in amplitudes]
    width = 1./(nstates+4.)
    xoff = i*width - width*nstates/2.
    x = numpy.array(amps_neg.columns[1:],dtype=float) + xoff
    ax[1].bar(x, amplitudes, width, label="$"+state+"^{-}$")


ax[0].text(0.05,0.95,"Positive-parity",transform=ax[0].transAxes);
ax[0].legend()
ax[0].set_xlabel("SP basis state");
ax[0].set_ylabel("Wavefunction amplitude");

ax[1].text(0.05,0.95,"Negative-parity",transform=ax[1].transAxes);
ax[1].legend()
ax[1].set_xlabel("SP basis state");
ax[1].set_ylabel("Wavefunction amplitude");
    
matplotlib.pyplot.show(block=True);
