"""Class to process neut setup.
Neut is run as a program external to the main comet code.
This module does normalization calculations that need only be done once per
job type. The results should be fed back in via the config file for neutMC
"""



import cometModule
import cometGeometry
import cometRun
import geomNameCheck
import os


class neutMC(cometModule.cometModule):

    """A class to run neut Monte Carlo setup
    """

###########################################################

    def checkConfig(self):

        """Check and fix the Monte Carlo type."""

        # The Monte Carlo type must be set to neut
        if self.config.options.has_option("SimG4","mc_type"):
            if (self.config.options.get("SimG4","mc_type") == "Neut_RooTracker"):
                return True
            else:
                self.config.log.error("The config option SimG4/mc_type is not set to Neut, but running neutMC module. Reset this option to Neut_RooTracker")
                self.config.options.set("SimG4","mc_type","Neut_RooTracker")
        else:
            # The option doesn't exist, set it.
            self.config.log.error("The config option SimG4/mc_type is not set, but running neutMC module. Setting this option to Neut_RooTracker")
            # Create the section if it doesn't exist.
            if not self.config.options.has_section("SimG4"):
                self.config.options.add_section("SimG4")
            # Set the option
            self.config.options.set("SimG4","mc_type","Neut_RooTracker")

###########################################################    
    
    def getOutputFilename(self):

        """Get the output filename"""

        self.outname = self.createFileName("root")
        return self.outname

###########################################################    

    def  getModuleName(self):

        """Return module name"""
        return "neutMC"

###########################################################    

    def  getShortModuleName(self):

        """Return module name"""
        return "numc"

#######################################################################

    def checkModuleConfig(self):

        """Checks the configuration"""

        # Make sure that the enviromnet setup script exists.
        if self.config.options.has_option("software","neut_setup_script"):
            script = self.config.options.get("software","neut_setup_script")
            if not os.path.isfile(script):
                self.config.log.error("Cannot locate specified neut setup script "+script)
                return False
        else:
            self.config.log.error("Neut MC being run, but setup script not set in the config file. Please use the software/neut_setup_script option")
            return False
        
        if self.config.options.has_option("neutrino","flux_file_path"):
            # Place holder for compilation
            path = self.config.options.get("neutrino","flux_file_path")

            if self.config.options.has_option("neutrino","flux_file_start"):
                path = self.config.options.get("neutrino","flux_file_path")
                first = self.config.options.get("neutrino","flux_file_start")
                fname = path + "." + first + ".root"
                if not os.path.exists(fname):
                    self.config.log.error("Cannot find file "+fname+" please check neutrin/flux_file_path and neutrino/flux_file_start. Cannot continue")
                    return False
            else:
                self.config.log.error("Neut MC being run, but neutrino/flux_file_start not specified. Cannot continue")
                return False
    
            if self.config.options.has_option("neutrino","flux_file_stop"):
                path = self.config.options.get("neutrino","flux_file_path")
                lastf = self.config.options.get("neutrino","flux_file_stop")
                fname = path + "." + lastf + ".root"
                if not os.path.exists(fname):
                    self.config.log.error("Cannot find file "+fname+" please check neutrin/flux_file_path and neutrino/flux_file_stop. Cannot continue")
                    return False
            else:
                self.config.log.error("Neut MC being run, but neutrino/flux_file_stop not specified. Cannot continue")
                return False

        elif self.config.options.has_option("neutrino","flux_file"):
            flux_file_name = self.config.options.get("neutrino","flux_file")
        else:
            self.config.log.error("Neut MC being run, but neutrino/flux_file_path and neutrino/flux_file not specified. Cannot continue")
            return False

        if self.config.options.has_option("neutrino","neutrino_type"):
            nu_type = self.config.options.get("neutrino","neutrino_type").upper()
            if nu_type == "BEAM" or nu_type == "NUE" or nu_type == "NUEBAR"or nu_type == "NUMU" or nu_type == "NUMUBAR":
                path = "hello"
            else:
                self.config.log.error("Unrecognised neutrino type requested. Cannot continue")
                return False

        if self.config.options.has_option("neutrino","neut_norm_factor"):
            # placeholder
            factor = self.config.options.get("neutrino","neut_norm_factor")
        else:
            if not self.config.options.has_option("neutrino","maxint_file"):
                self.config.log.error("Neut MC being run, but neutrino/neut_norm_factor nor setup ROOT file not specified. Cannot continue")
                return False       


        if self.config.options.has_option("neutrino","pot"):
            if self.config.options.has_option("neutrino","neut_pot_exp"):
                # placeholder
                factor = self.config.options.get("neutrino","neut_pot_exp")
            else:
	        if not self.config.options.has_option("neutrino","maxint_file"):
                    self.config.log.error("Neut MC being run using pot to set event number, but neutrino/neut_pot_exp nor setup ROOT file not specified. Cannot continue")
                    return False      

        if self.config.options.has_option("neutrino","flux_region"):
            region = self.config.options.get("neutrino","flux_region").upper()
            if not(region == "MAGNET" or region == "BASKET" or region == "INGRID" or region == "SAND"):
                self.config.log.error("Neut MC being run with invalid fluxregion. Please selecet magnet, basket or sand. Cannot continue")
                return False

        # Unless the the force_geom_name has been specified check that the
        # master_volume name matches to one in a list of known volumes. 
        if geomNameCheck.CheckTopVolumeName(self.config) == False: 
            return False

        return True

###########################################################

    def checkOutput(self):

        """Method to check the output of Neut"""

        # Set the randiom number seed if used.
        if self.config.options.has_option("neutrino","random_seed"):
            seed = self.config.options.get("neutrino","random_seed")
            self.jobSeed = seed

        if (self.run):
            return self.run.rtc == 0
        return True

###########################################################    

    def process(self):

        """Run the jobs for Neut"""

        # Get the geometry; either use geometry from SimG4 or sand geometry (if SAND region specified)
        geoFile = ""
        region = "MAGNET"
        if self.config.options.has_option("neutrino","flux_region"):
            region = self.config.options.get("neutrino","flux_region").upper()
        if region == "SAND":
            geoFile = self._generateSandGeometry()
        else:
            geoFile = self._generateGeometry()
        if not geoFile:
            self.config.log.stop("neutSetup: Error generating geometry file")
        if not os.path.isfile(geoFile):
            self.config.log.stop("neutSetup: Error generating geometry file")

        # MC the job commands.

        # Start by executing the setup script
        script = self.config.options.get("software","neut_setup_script")
        job = "source "+script+"\n"

        # Run the setup script inside $NEUTGEOM
        job += "$NEUTGEOM/setup.sh\n"

        # Copy the neut.card file to the working directory
        if self.config.options.has_option("neutrino","neut_card"):
            card = self.config.options.get("neutrino","neut_card")
        else:
            card = "$NEUTGEOM/neut.card"

        job += "cp "+card+" neut.card\n "

        # Random seeds for NEUT
        if self.config.options.has_option("neutrino","neut_seed1"):
            neut_seed1 = self.config.options.get("neutrino","neut_seed1")
            ran1 = neut_seed1
        else:
            ran1 = str(-1290784521)
            
        if self.config.options.has_option("neutrino","neut_seed2"):
            neut_seed2 = self.config.options.get("neutrino","neut_seed2")
            ran2 = neut_seed2
        else:
            ran2 = str(387739155)
            
        if self.config.options.has_option("neutrino","neut_seed3"):
            neut_seed3 = self.config.options.get("neutrino","neut_seed3")
            ran3 = neut_seed3
        else:
            ran3 = str(39827)

        # Seeds for restarting sequence
        ran4 = str(0)
        ran5 = str(0)

        ranseed = "\t" + ran1 + "\t" + ran2 + "\t" +ran3 + "\t" + ran4 + "\t" + ran5 + " "
        self.config.log.output("neutMC using random seed "+ranseed)
        job += "echo "+ranseed+">random.tbl\n"

        # Run Neut
        job += "$NEUTGEOM/genev"

        # Specify Geometry
        job += " -g "+geoFile

        # Specify jnubeam path
        if self.config.options.has_option("neutrino","flux_file_path"):
            path = self.config.options.get("neutrino","flux_file_path")
            first = self.config.options.get("neutrino","flux_file_start")
            lastf = self.config.options.get("neutrino","flux_file_stop")
            job += " -s "+path+" "+first+" "+lastf
        elif self.config.options.has_option("neutrino","flux_file"):
            flux_file_name = self.config.options.get("neutrino","flux_file")
            job += " -j "+flux_file_name

        # Specify flux region in job command
        if (region =="BASKET"):
            job += " -d 5 "
        elif (region =="MAGNET"):
            job += " -d 6 "
        elif (region =="INGRID"):
            job += " -d 3 "
        elif (region =="SAND"):
            job += " -d 13 "

        #Specify the flux tree
        if self.config.options.has_option("neutrino","flux_tree"):
            tree = self.config.options.get("neutrino","flux_tree")
            job += " -t "+tree+" "

	# Define the master volume for generation.
        masterVolume = ""
        if  self.config.options.has_option("neutrino","master_volume"):
           masterVolume =  self.config.options.get("neutrino","master_volume")
        if masterVolume:
            job += " -v +"+masterVolume

        # Specify the neutrino type
        if self.config.options.has_option("neutrino","neutrino_type"):
            nu_type = self.config.options.get("neutrino","neutrino_type").upper()
            if nu_type == "BEAM":
                job += " "
            elif nu_type == "NUE":
                job += " -p 12"
            elif nu_type == "NUEBAR":
                job += " -p -12"
            elif nu_type == "NUMU":
                job += " -p 14"
            elif nu_type == "NUMUBAR":
                job += " -p -14"

        # Normalization factor
        if self.config.options.has_option("neutrino","neut_norm_factor"):
            factor = self.config.options.get("neutrino","neut_norm_factor")
            job += " -m "+factor

        # rooTracker format
        job += " -f rootracker"

        # Random seed for the root side
        if self.config.options.has_option("neutrino","random_seed"):
            seed = self.config.options.get("neutrino","random_seed")
            job += " -r "+seed

        # Number of events
        if self.config.options.has_option("neutrino","pot"):
            if self.config.options.has_option("neutrino","neut_pot_exp"):
                factor = self.config.options.get("neutrino","neut_pot_exp")
                job += " -N "+factor

            pot = self.config.options.get("neutrino","pot")
            # Neut pot in units of 1e21 pot
            potn = float(pot)/1e21
            job += " -e "+str(potn)+" -w 1 " 
        elif self.config.options.has_option("neutrino","num_events"):
            num_events = self.config.options.get("neutrino","num_events")
            job += " -n  "+num_events+" -w 1 "
        else:
            # We just run through the specified jnubeam files
            job += " -w 0 "
            
        # Output file 
	outfile = self.getOutputFilename()

        job += " -o "+outfile

        # Maximum Interaction Probability file 
        if self.config.options.has_option("neutrino","maxint_file"):
            int_file = self.config.options.get("neutrino","maxint_file")
            job += " -i "+int_file

        # Start at random position in flux file
        if self.config.options.has_option("neutrino","random_start"):
            rnd_start = self.config.options.get("neutrino","random_start")
            job += " -Q "+rnd_start

        # Done
        job += "\n"
        
	# Create an emtpy output file to finish things off.
        job += "touch "+outfile+"\n"
        
        self.run = cometRun.cometRun(job,self.config)
        self.run.setCallName(self.getModuleName())
        # We don't want to configure cmt programs this time.
        self.run.setup = False

        self.run.run()

        self.config.log.output(self.run.stdout)
        self.config.log.error(self.run.stderr)

        
############################################################

    def _generateGeometry(self):
        
        """Method to generate the geometry file
        We run a very quick SimG4 job to do this."""
        
        geoname = self.createFileName("geo")
        macname = self.createFileName("mac")
        
        macContents = ""
        
        # Geometry Switches.
        cometGeom = cometGeometry.cometGeometry(self.config)
        geometryControl = cometGeom.getMACCommands()
        if geometryControl:
            macContents += geometryControl
        else:
            # Warning already printed.
            return False
                                                        
        # Now generate a single muon to run
        macContents += """
/comet/update
"""
        try:
            macFile = open(macname,"w")
            macFile.write(macContents)
            
        except:
            self.config.log.stop("SimG4: Cannot write mac file "+macname)
            
        macFile.close()

        job = "SIMG4.exe " + " -o " + geoname + " " + macname 
        
        self.run = cometRun.cometRun(job,self.config)
        self.run.setCallName(self.getModuleName()+"-geo")
        self.run.run()
        self.config.log.output(self.run.stdout)
        self.config.log.error(self.run.stderr)

        # Remove the mac file if required.
        cleanup = True
        if self.config.options.has_option("neutrino","cleanup"):
            cleanup = self.config.options.getboolean("neutrino","cleanup")

        if cleanup:
            os.unlink(macname)
            
        return geoname+".root"

############################################################

    def _generateSandGeometry(self):
        
        """Method to generate the sand geometry file"""

        job = "root -q -b -l ${SIMSANDROOT}/scripts/pit_geo_mm.C"

        
        self.run = cometRun.cometRun(job,self.config)
        self.run.setCallName(self.getModuleName()+"-sandgeo")
        self.run.run()
        self.config.log.output(self.run.stdout)
        self.config.log.error(self.run.stderr)

        return "pit_geometry_v3.root"