#!/usr/bin/env python

#  This file is part of MAUS: http://micewww.pp.rl.ac.uk:8080/projects/maus
# 
#  MAUS is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
# 
#  MAUS is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
# 
#  You should have received a copy of the GNU General Public License
#  along with MAUS.  If not, see <http://www.gnu.org/licenses/>.

"""Script to test the BeamlineSimulation Application."""

import unittest
import json
import MapPyBeamlineSimulation
from CallG4bl import CallG4bl

"""
Test Process function by taking all particles in spill generated by g4bl and 
adjusting number of protons from target. 
""" #pylint: disable = W0105
TEST_1 = {
    "g4bl":{"q_1":1.16, "q_2":-1.45, "q_3":1.006, "d_1":-1.41,\
    "proton_absorber_thickness":11,'d_s':200,\
    "proton_number":500000000,'particles_per_spill':0,\
    "particle_charge":"all", 'rotation_angle':0, 'translation_z':20000.0,\
    "random_seed":0,"run_number":0,"get_magnet_currents_pa_cdb":False,\
    "proton_weight":1

    }
}

"""
Test Process function by taking specified number of particles per spill 
generated by g4bl and require that the buffer be refilled twice
""" #pylint: disable = W0105
TEST_2 = {
    "g4bl":{"q_1":1.16, "q_2":-1.45, "q_3":1.006, "d_1":-1.41,\
    "proton_absorber_thickness":11,'d_s':200,\
    "proton_number":500000000,'particles_per_spill':200,\
    "particle_charge":"all", 'rotation_angle':0, 'translation_z':20000.0,\
    "random_seed":0,"run_number":0,"get_magnet_currents_pa_cdb":False,\
    "proton_weight":1

    }
}

"""
Test Process function by taking one spill with more particles than 
generated by g4bl and requiring that the buffer be refilled twice
""" #pylint: disable = W0105
TEST_3 = {
    "g4bl":{"q_1":1.16, "q_2":-1.45, "q_3":1.006, "d_1":-1.41,\
    "proton_absorber_thickness":11,'d_s':200,\
    "proton_number":500000000,'particles_per_spill':500,\
    "particle_charge":"all", 'rotation_angle':0, 'translation_z':20000.0,\
    "random_seed":0,"run_number":0,"get_magnet_currents_pa_cdb":False,\
    "proton_weight":1

    }
}

"""
Test error messages in Birth function

""" #pylint: disable = W0105
TEST_4 = {
    "g4bl":{"q_1":1.16, "q_2":-1.45, "q_3":1.006, "d_1":-1.41,\
    "proton_absorber_thickness":-11,'d_s':200,
    "proton_number":10000,'particles_per_spill':200,\
    "particle_charge":"all", 'rotation_angle':0, 'translation_z':20000.0,\
    "random_seed":0,"run_number":0,"get_magnet_currents_pa_cdb":False,\
    "proton_weight":1

    }
}

"""
Test error messages in Birth function

""" #pylint: disable = W0105
TEST_5 = {
    "g4bl":{"q_1":1.16, "q_2":-1.45, "q_3":1.006, "d_1":-1.41,\
    "proton_absorber_thickness":11,'d_s':200,\
    "proton_number":-10000,'particles_per_spill':200,\
    "particle_charge":"all", 'rotation_angle':0, 'translation_z':20000.0,\
    "random_seed":0,"run_number":0,"get_magnet_currents_pa_cdb":False,\
    "proton_weight":1

    }
}

"""
Test error messages in Birth function

""" #pylint: disable = W0105
TEST_6 = {
    "g4bl":{"q_1":1.16, "q_2":-1.45, "q_3":1.006, "d_1":-1.41,\
    "proton_absorber_thickness":11,'d_s':200,\
    "proton_number":10000,'particles_per_spill':200,\
    "particle_charge":"1.21 jiggawatts", 'rotation_angle':0,\
    'translation_z':20000.0,\
    "random_seed":0,"run_number":0,"get_magnet_currents_pa_cdb":False,\
    "proton_weight":1

    }
}

"""
Test that process function stops if g4bl file is empty. 
""" #pylint: disable = W0105
TEST_7 = {
    "g4bl":{"q_1":0, "q_2":0, "q_3":0, "d_1":0,\
    "proton_absorber_thickness":11,'d_s':0,\
    "proton_number":100,'particles_per_spill':25,\
    "particle_charge":"negative", 'rotation_angle':0,\
    'translation_z':20000.0,\
    "random_seed":0,"run_number":0,"get_magnet_currents_pa_cdb":False,\
    "proton_weight":1

    }
}

"""
Test random seed - seed from spill

""" #pylint: disable = W0105
TEST_8 = {
    "g4bl":{"q_1":1.16, "q_2":-1.45, "q_3":1.006, "d_1":-1.41, "d_2":1., 
    "proton_absorber_thickness":11,'d_s':200,
    "proton_number":100,'particles_per_spill':10,
    "particle_charge":"positive", 'rotation_angle':0,
    'translation_z':20000.0,
    "random_seed":-2,"run_number":0,"get_magnet_currents_pa_cdb":False,
    "proton_weight":1,
    "seed_algorithm":"random_seed_and_spill_number",
    }
}

#pylint: disable = R0903
class MockCallG4bl(object):
    """ Class which does a minimal mock of a G4BL call """
    #pylint: disable=C0103, R0903, R0912, R0913, R0914, R0915, W0613
    def __init__(self, newline='', file_path='', theta=0, deltaz=0, 
                 path_g4bl='', output_path='', run_number=0,
                 get_magnet_currents_pa_cdb='', random_seed=0):
        """ Initialise the class """
        self.particles = {}
        for i in range(self.number_of_entries):
            self.particles['entry'+str(i)] = {}

    number_of_entries = 10


class TestMapPyBeamlineSimulation(unittest.TestCase): #pylint: disable = R0904
    """ set of tests for MapPyBeamlineSimulation """
    @classmethod
    def setUp(cls): #pylint: disable = C0103
        """Initialise beam_maker object"""
        cls.g4bl = MapPyBeamlineSimulation.MapPyBeamlineSimulation()

    def teardown(cls): #pylint: disable = E0213
        """Close beam_maker object"""
        cls.assertTrue(cls.g4bl.death())

    def test_birth(self):
        """ Test error messages """
        self.assertTrue(self.g4bl.birth(json.dumps(TEST_1)))
        self.assertFalse(self.g4bl.birth(json.dumps(TEST_4)))
        self.assertFalse(self.g4bl.birth(json.dumps(TEST_5)))
        self.assertFalse(self.g4bl.birth(json.dumps(TEST_6)))
        self.assertFalse(self.g4bl.birth(""))
        self.assertFalse(self.g4bl.birth(json.dumps({})))

    def test_birth_values(self):
        """ Test birth assigns correct values """
        self.g4bl.birth(json.dumps(TEST_1))
        self.assertEqual(self.g4bl.q_1, 1.16)
        self.assertEqual(self.g4bl.q_2, -1.45)
        self.assertEqual(self.g4bl.q_3, 1.006)
        self.assertEqual(self.g4bl.d_1, -1.41)
        self.assertEqual(self.g4bl.d_s, 200)
        self.assertEqual(self.g4bl.proton_absorber_thickness, 11)
        self.assertEqual(self.g4bl.proton_weight, 1)
        self.assertEqual(self.g4bl.particles_per_spill, 0)
        self.assertEqual(self.g4bl.particle_charge, 'all')

    def test_process_test_1(self):
        """ Test that process returns 'good' spill """
        print json.dumps(TEST_1)
        self.g4bl.birth(json.dumps(TEST_1))
        print self.g4bl.birth(json.dumps(TEST_1))
        spill_one = self.g4bl.process(json.dumps({}))
        spill = json.loads(spill_one)        
        self.assertTrue(spill["mc_events"]!=[])
#        spill_two = self.g4bl.process(json.dumps({}))
#        spill = json.loads(spill_two)
#        self.assertTrue(spill["mc_events"]!=[])

    def test_process_test_2(self):
        """ Test that process returns 'good' spill """
        self.g4bl.birth(json.dumps(TEST_2))
        spill_one = self.g4bl.process(json.dumps({}))
        spill = json.loads(spill_one)
        self.assertEqual(len(spill["mc_events"]),  200)
        spill_two = self.g4bl.process(json.dumps({}))
        spill = json.loads(spill_two)
        self.assertEqual(len(spill["mc_events"]),  200)

    def test_process_test_3(self):
        """ Test that process returns 'good' spill """
        self.g4bl.birth(json.dumps(TEST_3))
        spill_one = self.g4bl.process(json.dumps({}))
        spill = json.loads(spill_one)
        self.assertEqual(len(spill["mc_events"]),  500)

    def test_process_test_7(self):
        """ Test that process returns 'bad' spill """
        #with self.assertRaises(SystemExit):
        self.g4bl.birth(json.dumps(TEST_7))
        self.g4bl.process(json.dumps({}))

    #pylint: disable = C0103
    def test_mappybeamline_new_random_seed(self):
        """ Test mappybeamline random seed algorithm"""
        random_seeds = []
        for i in range(10):
            for j in range(10):
                for k in range(10):
                    self.g4bl.new_random_seed([i, j, k])
                    random_seeds.append(self.g4bl.random_seed)
        self.assertEqual(len(set(random_seeds)), len(random_seeds))


    def test_mappybeamline_random_seed(self):
        """ Test mappybeamline random seed from spill"""
        # substitute in mock-up for CallG4bl
        MapPyBeamlineSimulation.CallG4bl = MockCallG4bl
        self.g4bl = MapPyBeamlineSimulation.MapPyBeamlineSimulation()
        # run a selection of spill numbers/etc and check that the random seed
        # is indeed unique 
        self.g4bl.birth(json.dumps(TEST_8))
        random_seed = self.g4bl.random_seed
        try: # spill = 0, ev = 0
            self.g4bl.process(json.dumps({"spill_number":0}))
        except IndexError:
            pass
        self.assertNotEqual(self.g4bl.random_seed, random_seed)
        random_seed = self.g4bl.random_seed
        try: # spill = 0, ev = 0
            self.g4bl.process(json.dumps({"spill_number":0}))
        except IndexError:
            pass
        self.assertEqual(self.g4bl.random_seed, random_seed)
        random_seed = self.g4bl.random_seed
        try: # spill = 1, ev = 0; old random number
            self.g4bl.process(json.dumps({"spill_number":1}))
        except IndexError:
            pass
        self.assertNotEqual(self.g4bl.random_seed, random_seed)
        random_seed = self.g4bl.random_seed
        MockCallG4bl.number_of_entries = 5
        try: # spill = 1, ev = 5; new random number
            self.g4bl.process(json.dumps({"spill_number":1}))
        except IndexError:
            pass
        self.assertNotEqual(self.g4bl.random_seed, random_seed)
        MapPyBeamlineSimulation.CallG4bl = CallG4bl


if __name__ == "__main__":
    unittest.main()