import os
import cProfile
import numpy as np
from unittest import TestCase
from dataclasses import fields, MISSING
from eos import options, reduction_reflectivity, logconfig
from orsopy import fileio

logconfig.setup_logging()
logconfig.update_loglevel(1)

# TODO: add unit tests for individual parts of reduction

class FullAmorTest(TestCase):
    @classmethod
    def setUpClass(cls):
        # generate map for option defaults
        cls._field_defaults = {}
        for opt in [options.ExperimentConfig, options.ReflectivityReductionConfig, options.ReflectivityOutputConfig]:
            defaults = {}
            for field in fields(opt):
                if field.default not in [None, MISSING]:
                    defaults[field.name] = field.default
                elif field.default_factory not in [None, MISSING]:
                    defaults[field.name] = field.default_factory()
            cls._field_defaults[opt.__name__] = defaults
        cls.pr = cProfile.Profile()

    @classmethod
    def tearDownClass(cls):
        cls.pr.dump_stats("profile_test.prof")

    def setUp(self):
        self.pr.enable()
        self.reader_config = options.ReaderConfig(
                year=2025,
                rawPath=["test_data"],
                )

    def tearDown(self):
        self.pr.disable()
        for fi in ['test_results/test.Rqz.ort', 'test_results/5952.norm']:
           try:
               os.unlink(fi)
           except FileNotFoundError:
               pass

    def test_time_slicing(self):
        experiment_config = options.ExperimentConfig(
                chopperPhase=-13.5,
                chopperPhaseOffset=-5,
                yRange=(18, 48),
                lambdaRange=(3., 11.5),
                mu=0,
                nu=0,
                muOffset=0.0,
                sampleModel='air | 10 H2O | D2O'
                )
        reduction_config = options.ReflectivityReductionConfig(
                qResolution=0.01,
                qzRange=(0.01, 0.15),
                thetaRange=(-0.75, 0.75),
                fileIdentifier=["6003-6005"],
                scale=[1],
                normalisationFileIdentifier=[],
                timeSlize=[300.0]
                )
        output_config = options.ReflectivityOutputConfig(
                outputFormats=[options.OutputFomatOption.Rqz_ort],
                outputName='test',
                outputPath='test_results',
                )
        config=options.ReflectivityConfig(self.reader_config, experiment_config, reduction_config, output_config)
        # run three times to get similar timing to noslicing runs
        reducer = reduction_reflectivity.ReflectivityReduction(config)
        reducer.reduce()
        reducer = reduction_reflectivity.ReflectivityReduction(config)
        reducer.reduce()
        reducer = reduction_reflectivity.ReflectivityReduction(config)
        reducer.reduce()

    def test_noslicing(self):
        experiment_config = options.ExperimentConfig(
                chopperPhase=-13.5,
                chopperPhaseOffset=-5,
                yRange=(18, 48),
                lambdaRange=(3., 11.5),
                mu=0,
                nu=0,
                muOffset=0.0,
                )
        reduction_config = options.ReflectivityReductionConfig(
                qResolution=0.01,
                thetaRange=(-0.75, 0.75),
                fileIdentifier=["6003", "6004", "6005"],
                scale=[1],
                normalisationFileIdentifier=["5952"],
                autoscale=(0.0, 0.05),
                )
        output_config = options.ReflectivityOutputConfig(
                outputFormats=[options.OutputFomatOption.Rqz_ort],
                outputName='test',
                outputPath='test_results',
                )
        config=options.ReflectivityConfig(self.reader_config, experiment_config, reduction_config, output_config)
        reducer = reduction_reflectivity.ReflectivityReduction(config)
        reducer.reduce()
        # run second time to reuse norm file
        reducer = reduction_reflectivity.ReflectivityReduction(config)
        reducer.reduce()

    def test_eventfilter(self):
        self.reader_config.year = 2026
        experiment_config = options.ExperimentConfig()
        reduction_config = options.ReflectivityReductionConfig(fileIdentifier=["826"],
                                                               logfilter=['polarization_config_label==2'])
        output_config = options.ReflectivityOutputConfig(
                outputFormats=[options.OutputFomatOption.Rqz_ort],
                outputName='test',
                outputPath='test_results',
                )
        config=options.ReflectivityConfig(self.reader_config, experiment_config, reduction_config, output_config)
        reducer = reduction_reflectivity.ReflectivityReduction(config)
        reducer.reduce()
        espin_up = reducer.dataset.data.events.shape[0]

        reduction_config.logfilter = ['polarization_config_label==3']
        output_config.append = True
        reducer = reduction_reflectivity.ReflectivityReduction(config)
        reducer.reduce()
        espin_down = reducer.dataset.data.events.shape[0]
        # measurement should have about 2x as many counts in spin_down
        self.assertAlmostEqual(espin_down/espin_up, 2., 2)

        # perform the same filter but remove pulses during which the switch occured
        reduction_config.logfilter = ['!polarization_config_label==3']
        output_config.append = True
        reducer = reduction_reflectivity.ReflectivityReduction(config)
        reducer.reduce()
        espin_down2 = reducer.dataset.data.events.shape[0]
        # measurement should have about 2x as many counts in spin_down
        self.assertLess(espin_down2, espin_down)

    def test_polsplitting(self):
        self.reader_config.year = 2026
        experiment_config = options.ExperimentConfig()
        reduction_config = options.ReflectivityReductionConfig(fileIdentifier=["826"])
        output_config = options.ReflectivityOutputConfig(
                outputFormats=[options.OutputFomatOption.Rqz_ort],
                outputName='test',
                outputPath='test_results',
                )
        config=options.ReflectivityConfig(self.reader_config, experiment_config, reduction_config, output_config)
        reducer = reduction_reflectivity.ReflectivityReduction(config)
        reducer.reduce()

        results = fileio.load_orso(os.path.join(output_config.outputPath, output_config.outputName+'.Rqz.ort'))
        self.assertEqual(len(results), 2)
        self.assertEqual(results[0].info.data_source.measurement.instrument_settings.polarization, 'po')
        self.assertEqual(results[1].info.data_source.measurement.instrument_settings.polarization, 'mo')
        espin_up = np.nansum(results[0].data[:,1])
        espin_down = np.nansum(results[1].data[:,1])
        # the total intensity should be around equal as events are doubled and monitor counts are doubled
        self.assertAlmostEqual(espin_down/espin_up, 1., 2)