#
#     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 unittest
import os
import numpy as np
from gbm.spectra.fitting import SpectralFitterPgstat, SpectralFitterChisq
from gbm.spectra.fitting import SpectralFitterPstat, SpectralFitterCstat
from gbm.spectra.functions import Comptonized
from gbm.data import PHA, RSP

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


class TestSpectralFitter(unittest.TestCase):
    
    pha = PHA.open(os.path.join(data_dir, 'sim_pha.pha'))
    bkgd = np.load(os.path.join(data_dir, 'sim_bkgd.npy'), allow_pickle=True)[0]
    rsp = RSP.open(os.path.join(data_dir, 'glg_cspec_n9_bn090131090_v00.rsp2'))
    
    def test_pgstat(self):
        fitter = SpectralFitterPgstat([self.pha], [self.bkgd], [self.rsp],
                                      method='TNC')
        fitter.fit(Comptonized(), options={'maxiter': 10000})
        self.assertEqual(fitter.dof, 117)
        self.assertEqual(fitter.function_name, 'Comptonized')
        self.assertEqual(fitter.num_components, 1)
        self.assertEqual(fitter.detectors, ['n9'])
        self.assertAlmostEqual(fitter.energy_range[0], 8.0, places=0)
        self.assertAlmostEqual(fitter.energy_range[1], 908.0, places=0)
        self.assertEqual(fitter.num_sets, 1)
        self.assertEqual(fitter.success, True)
        self.assertAlmostEqual(fitter.parameters[0], 1.0, delta=0.2)
        self.assertAlmostEqual(fitter.parameters[1], 220., delta=100.)
        self.assertAlmostEqual(fitter.parameters[2], -1.2, delta=0.2)
        self.assertEqual(fitter.jacobian.size, 3)
        self.assertEqual(fitter.hessian.size, 9)
        self.assertEqual(fitter.covariance.size, 9)
    
    def test_chisq(self):
        fitter = SpectralFitterChisq([self.pha], [self.bkgd], [self.rsp], 
                                      method='TNC')
        fitter.fit(Comptonized(), options={'maxiter': 10000})
        self.assertEqual(fitter.dof, 117)
        self.assertEqual(fitter.function_name, 'Comptonized')
        self.assertEqual(fitter.num_components, 1)
        self.assertEqual(fitter.detectors, ['n9'])
        self.assertAlmostEqual(fitter.energy_range[0], 8.0, places=0)
        self.assertAlmostEqual(fitter.energy_range[1], 908.0, places=0)
        self.assertEqual(fitter.num_sets, 1)
        self.assertEqual(fitter.success, True)
        self.assertAlmostEqual(fitter.parameters[0], 1.0, delta=0.1)
        self.assertAlmostEqual(fitter.parameters[1], 220.0, delta=10.)
        self.assertAlmostEqual(fitter.parameters[2], -1.2, delta=0.1)
        self.assertEqual(fitter.jacobian.size, 3)
        self.assertEqual(fitter.hessian.size, 9)
        self.assertEqual(fitter.covariance.size, 9)

    def test_cstat(self):
        fitter = SpectralFitterCstat([self.pha], [self.bkgd], [self.rsp], 
                                      method='TNC')
        fitter.fit(Comptonized(), options={'maxiter': 10000})
        self.assertEqual(fitter.dof, 117)
        self.assertEqual(fitter.function_name, 'Comptonized')
        self.assertEqual(fitter.num_components, 1)
        self.assertEqual(fitter.detectors, ['n9'])
        self.assertAlmostEqual(fitter.energy_range[0], 8.0, places=0)
        self.assertAlmostEqual(fitter.energy_range[1], 908.0, places=0)
        self.assertEqual(fitter.num_sets, 1)
        self.assertEqual(fitter.success, True)
        self.assertAlmostEqual(fitter.parameters[0], 1.0, delta=0.1)
        self.assertAlmostEqual(fitter.parameters[1], 220.0, delta=30.)
        self.assertAlmostEqual(fitter.parameters[2], -1.2, delta=0.1)
        self.assertEqual(fitter.jacobian.size, 3)
        self.assertEqual(fitter.hessian.size, 9)
        self.assertEqual(fitter.covariance.size, 9)

    def test_pstat(self):
        fitter = SpectralFitterPstat([self.pha], [self.bkgd], [self.rsp], 
                                      method='TNC')
        fitter.fit(Comptonized(), options={'maxiter': 10000})
        self.assertEqual(fitter.dof, 117)
        self.assertEqual(fitter.function_name, 'Comptonized')
        self.assertEqual(fitter.num_components, 1)
        self.assertEqual(fitter.detectors, ['n9'])
        self.assertAlmostEqual(fitter.energy_range[0], 8.0, places=0)
        self.assertAlmostEqual(fitter.energy_range[1], 908.0, places=0)
        self.assertEqual(fitter.num_sets, 1)
        self.assertEqual(fitter.success, True)
        self.assertAlmostEqual(fitter.parameters[0], 1.0, delta=0.1)
        self.assertAlmostEqual(fitter.parameters[1], 220.0, delta=30.)
        self.assertAlmostEqual(fitter.parameters[2], -1.2, delta=0.1)
        self.assertEqual(fitter.jacobian.size, 3)
        self.assertEqual(fitter.hessian.size, 9)
        self.assertEqual(fitter.covariance.size, 9)