# PopulationPlotter.py generated by R. Yokoyama on 10/05/2018
from collections import OrderedDict
import ROOT
import copy

class PopulationPlotter(object):
    # A class for plotting the nuclear population of a given time
    # from a given list of DecayChain objects
    def __init__(self):
        self.low_n = 0
        self.low_z = 0
        self.low_a = 0
        self.up_n = 250
        self.up_z = 150
        self.up_a = 350
        self.count_max = -1

    def configure(self,config):
        # configures the plot range from a given yaml node
        self.low_n = config['low_n']
        self.low_z = config['low_z']
        self.low_a = config['low_a']
        self.up_n = config['up_n']
        self.up_z = config['up_z']
        self.up_a = config['up_a']
        self.count_max = config['count_max']

    def plotPopulation(self, decay_chain_list, time):
        # plot nuclear population at "time"
        # returns ROOT::TH2D object
        hist = ROOT.TH2D("population","population",
                         self.up_n - self.low_n, self.low_n - 0.5, self.up_n - 0.5,
                         self.up_z - self.low_z, self.low_z - 0.5, self.up_z - 0.5)
        for chain in decay_chain_list:
            nucl_after = filter(lambda x: x.time > time, chain.event_list)
            if len(nucl_after) == len(chain.event_list):
                nucleus = copy.copy(nucl_after[0])
            elif len(nucl_after) == 0:
                nucleus = copy.copy(chain.event_list[-1])
                nucleus.n = nucleus.n - 1 - nucleus.n_neutron
                nucleus.z = nucleus.z + 1
            else:
                nucleus = copy.copy(nucl_after[0])
            hist.Fill(nucleus.n,nucleus.z)

        hist.SetMinimum(1)
        hist.SetMaximum(self.count_max)
        return hist

    def plotPopulationA(self, decay_chain_list, time):
        # plot nuclear population at "time"
        # returns ROOT::TH1D object which is a projection on to mass (A) axis
        hist = ROOT.TH1D("populationA","populationA",self.up_a-self.low_a,self.low_a-0.5,self.up_a-0.5)
        for chain in decay_chain_list:
           nucl_after = filter(lambda x: x.time > time, chain.event_list)
           if len(nucl_after) == len(chain.event_list):
               nucleus = copy.copy(nucl_after[0])
           elif len(nucl_after) == 0:
               nucleus = copy.copy(chain.event_list[-1])
               nucleus.n = nucleus.n - 1 - nucleus.n_neutron
               nucleus.z = nucleus.z + 1
           else:
               nucleus = copy.copy(nucl_after[0])
           hist.Fill(nucleus.n+nucleus.z)
        hist.SetMinimum(0.1)
        hist.SetMaximum(self.count_max)
        return hist

    def plotPopulation1D(self, decay_chain_list, time):
        hists = OrderedDict()
        for z in list(reversed(range(self.low_z, self.up_z))):
            name = "population_z_" + str(z)
            hists.update({z : ROOT.TH1I(name,name,self.up_n-self.low_n,self.low_n-0.5,self.up_n-0.5)})
            hists[z].SetMinimum(0)
            hists[z].SetMaximum(self.count_max)

        for chain in decay_chain_list:
            nucl_after = filter(lambda x: x.time > time, chain.event_list)
            if len(nucl_after) == len(chain.event_list):
                nucleus = copy.copy(chain.initial_nuclide)
            elif len(nucl_after) == 0:
                nucleus = copy.copy(chain.event_list[-1])
                nucleus.n = nucleus.n - 1 - nucleus.n_neutron
                nucleus.z = nucleus.z + 1
            else:
                nucleus = copy.copy(nucl_after[0])
                #nucleus.n = nucleus.n - 1 - nucleus.n_neutron
                #nucleus.z = nucleus.z + 1
            if nucleus.z > self.low_z-1 and nucleus.z < self.up_z:
                hists[nucleus.z].Fill(nucleus.n)
        return hists

