#
#     Authors: William Cleveland (USRA),
#              Adam Goldstein (USRA) and
#              Daniel Kocevski (NASA)
#
#     Portions of the code are Copyright 2020 William Cleveland and
#     Adam Goldstein, Universities Space Research Association
#     All rights reserved.
#
#     Written for the Fermi Gamma-ray Burst Monitor (Fermi-GBM)
#
#     This program 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.
#
#     This program 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 this program.  If not, see <https://www.gnu.org/licenses/>.
#
import numpy as np
import os
from unittest import TestCase

from gbm.data.pha import PHA, BAK

data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data')

class TestBAK(TestCase):
    filename = os.path.join(data_dir, 'glg_tte_n0_bn160509374_xspec_v00.bak')
    trigtime = 484477130.219298
    header_names = ['PRIMARY', 'EBOUNDS', 'SPECTRUM', 'GTI']
    numchans = 128
    trange = (484477128.411298, 484477165.275298)
    tcent = np.sum(trange)/2.0
    erange = (4.525152, 2000)
    gti = [(-1.80800002813339, 35.055999994278)]
    exposure = 36.864
    valid_channels = np.arange(128).tolist()

    def test_attributes(self):
        bak = BAK.open(self.filename)
        trigtime = bak.trigtime
        self.assertAlmostEqual(trigtime, self.trigtime, places=6)
        self.assertEqual(bak.is_gbm_file, True)
        self.assertEqual(bak.id, '160509374')
        self.assertEqual(bak.filename, os.path.basename(self.filename))
        self.assertEqual(bak.is_trigger, True)
        self.assertEqual(bak.detector, 'n0')
        self.assertEqual(bak.datatype, 'TTE')
        self.assertAlmostEqual(bak.energy_range[0], self.erange[0], places=4)
        self.assertAlmostEqual(bak.energy_range[1], self.erange[1], places=4)
        self.assertEqual(len(bak.gti), 1)
        self.assertAlmostEqual(bak.gti[0][0], self.gti[0][0], places=6)
        self.assertAlmostEqual(bak.gti[0][1], self.gti[0][1], places=6)
        self.assertCountEqual(bak.headers.keys(), self.header_names)
        self.assertEqual(bak.numchans, self.numchans)
        self.assertAlmostEqual(bak.time_range[0], self.trange[0]-trigtime, places=6)
        self.assertAlmostEqual(bak.time_range[1], self.trange[1]-trigtime, places=6)
        self.assertAlmostEqual(bak.tcent, self.tcent-trigtime, places=6)
        self.assertCountEqual(bak.valid_channels[0], self.valid_channels)
    
    def test_write(self):
        bak = BAK.open(self.filename)
        new_file = os.path.join(data_dir, 'glg_tte_n0_bn160509374_xspec_v01.bak')
        bak.write(os.path.dirname(self.filename), 
                  filename='glg_tte_n0_bn160509374_xspec_v01.bak')    
        bak2 = BAK.open(new_file)
        os.remove(new_file)
        
        self.assertCountEqual(bak2.time_range, bak.time_range)
        self.assertCountEqual(bak2.data.lo_edges, bak.data.lo_edges)
        self.assertCountEqual(bak2.data.hi_edges, bak.data.hi_edges)
        self.assertCountEqual(bak2.data.exposure, bak.data.exposure)
        self.assertCountEqual(bak2.data.rates, bak.data.rates)
        self.assertCountEqual(bak2.data.rate_uncertainty, bak.data.rate_uncertainty)

class TestPHA(TestCase):
    pha = PHA.open(os.path.join(data_dir, 'sim_pha.pha'))

    def test_write(self):
        self.pha.write(data_dir, filename='test_pha.pha')
        test_pha = PHA.open(os.path.join(data_dir, 'test_pha.pha'))
        os.remove(os.path.join(data_dir, 'test_pha.pha'))
        
        self.assertCountEqual(self.pha.energy_range, test_pha.energy_range)
        self.assertEqual(self.pha.exposure, test_pha.exposure)
        self.assertCountEqual(self.pha.gti, test_pha.gti)
        self.assertEqual(self.pha.numchans, test_pha.numchans)
        self.assertEqual(self.pha.tcent, test_pha.tcent)
        self.assertCountEqual(self.pha.time_range, test_pha.time_range)
        self.assertEqual(self.pha.trigtime, test_pha.trigtime)
        self.assertCountEqual(self.pha.valid_channels, test_pha.valid_channels)
        
        self.assertCountEqual(self.pha.data.counts, test_pha.data.counts)
        self.assertCountEqual(self.pha.data.lo_edges, test_pha.data.lo_edges)
        self.assertCountEqual(self.pha.data.hi_edges, test_pha.data.hi_edges)
        self.assertCountEqual(self.pha.data.exposure, test_pha.data.exposure)

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