# This file is part of pipe_base.
#
# Developed for the LSST Data Management System.
# This product includes software developed by the LSST Project
# (https://www.lsst.org).
# See the COPYRIGHT file at the top-level directory of this distribution
# for details of code ownership.
#
# 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 <http://www.gnu.org/licenses/>.

"""Tests of the Instrument class.
"""

import datetime
import math
import unittest

from lsst.daf.butler import DataCoordinate, DimensionPacker, DimensionUniverse, Registry, RegistryConfig
from lsst.daf.butler.formatters.json import JsonFormatter
from lsst.pex.config import Config
from lsst.pipe.base import Instrument
from lsst.utils.introspection import get_full_type_name


class DummyInstrument(Instrument):
    @classmethod
    def getName(cls):
        return "DummyInstrument"

    def register(self, registry, update=False):
        detector_max = 2
        visit_max = 10
        exposure_max = 8
        record = {
            "instrument": self.getName(),
            "class_name": get_full_type_name(DummyInstrument),
            "detector_max": detector_max,
            "visit_max": visit_max,
            "exposure_max": exposure_max,
        }
        with registry.transaction():
            registry.syncDimensionData("instrument", record, update=update)

    def getRawFormatter(self, dataId):
        return JsonFormatter


class BadInstrument(DummyInstrument):
    """Instrument with wrong class name."""

    raw_definition = ("raw2", ("instrument", "detector", "exposure"), "StructuredDataDict")

    @classmethod
    def getName(cls):
        return "BadInstrument"

    def register(self, registry, update=False):
        # Register a bad class name
        record = {
            "instrument": self.getName(),
            "class_name": "builtins.str",
            "detector_max": 1,
        }
        registry.syncDimensionData("instrument", record, update=update)


class UnimportableInstrument(DummyInstrument):
    """Instrument with class name that does not exist."""

    @classmethod
    def getName(cls):
        return "NoImportInstr"

    def register(self, registry, update=False):
        # Register a bad class name
        record = {
            "instrument": self.getName(),
            "class_name": "not.importable",
            "detector_max": 1,
        }
        registry.syncDimensionData("instrument", record, update=update)


class DimensionPackerTestConfig(Config):
    packer = Instrument.make_dimension_packer_config_field()


class InstrumentTestCase(unittest.TestCase):
    """Test for Instrument."""

    def setUp(self):
        self.instrument = DummyInstrument()
        self.name = "DummyInstrument"

    def test_basics(self):
        self.assertEqual(self.instrument.getName(), self.name)
        self.assertEqual(self.instrument.getRawFormatter({}), JsonFormatter)
        self.assertIsNone(DummyInstrument.raw_definition)
        raw = BadInstrument.raw_definition
        self.assertEqual(raw[2], "StructuredDataDict")

    def test_register(self):
        """Test that register() sets appropriate Dimensions."""
        registryConfig = RegistryConfig()
        registryConfig["db"] = "sqlite://"
        registry = Registry.createFromConfig(registryConfig)
        # Check that the registry starts out empty.
        self.instrument.importAll(registry)
        self.assertFalse(list(registry.queryDimensionRecords("instrument")))

        # Register and check again.
        self.instrument.register(registry)
        instruments = list(registry.queryDimensionRecords("instrument"))
        self.assertEqual(len(instruments), 1)
        self.assertEqual(instruments[0].name, self.name)
        self.assertEqual(instruments[0].detector_max, 2)
        self.assertIn("DummyInstrument", instruments[0].class_name)

        self.instrument.importAll(registry)
        from_registry = DummyInstrument.fromName("DummyInstrument", registry)
        self.assertIsInstance(from_registry, Instrument)
        with self.assertRaises(LookupError):
            Instrument.fromName("NotThrere", registry)

        # Register a bad instrument.
        BadInstrument().register(registry)
        with self.assertRaises(TypeError):
            Instrument.fromName("BadInstrument", registry)

        UnimportableInstrument().register(registry)
        with self.assertRaises(ImportError):
            Instrument.fromName("NoImportInstr", registry)

        # This should work even with the bad class name.
        self.instrument.importAll(registry)

    def test_defaults(self):
        self.assertEqual(self.instrument.makeDefaultRawIngestRunName(), "DummyInstrument/raw/all")
        self.assertEqual(
            self.instrument.makeUnboundedCalibrationRunName("a", "b"), "DummyInstrument/calib/a/b/unbounded"
        )
        self.assertEqual(
            self.instrument.makeCuratedCalibrationRunName("2018-05-04", "a"),
            "DummyInstrument/calib/a/curated/20180504T000000Z",
        )
        self.assertEqual(self.instrument.makeCalibrationCollectionName("c"), "DummyInstrument/calib/c")
        self.assertEqual(self.instrument.makeRefCatCollectionName(), "refcats")
        self.assertEqual(self.instrument.makeRefCatCollectionName("a"), "refcats/a")
        self.assertEqual(self.instrument.makeUmbrellaCollectionName(), "DummyInstrument/defaults")

        instrument = DummyInstrument(collection_prefix="Different")
        self.assertEqual(instrument.makeCollectionName("a"), "Different/a")
        self.assertEqual(self.instrument.makeCollectionName("a"), "DummyInstrument/a")

    def test_collection_timestamps(self):
        self.assertEqual(
            Instrument.formatCollectionTimestamp("2018-05-03"),
            "20180503T000000Z",
        )
        self.assertEqual(
            Instrument.formatCollectionTimestamp("2018-05-03T14:32:16"),
            "20180503T143216Z",
        )
        self.assertEqual(
            Instrument.formatCollectionTimestamp("20180503T143216Z"),
            "20180503T143216Z",
        )
        self.assertEqual(
            Instrument.formatCollectionTimestamp(datetime.datetime(2018, 5, 3, 14, 32, 16)),
            "20180503T143216Z",
        )
        formattedNow = Instrument.makeCollectionTimestamp()
        self.assertIsInstance(formattedNow, str)
        datetimeThen1 = datetime.datetime.strptime(formattedNow, "%Y%m%dT%H%M%S%z")
        self.assertEqual(datetimeThen1.tzinfo, datetime.timezone.utc)

        with self.assertRaises(TypeError):
            Instrument.formatCollectionTimestamp(0)

    def test_dimension_packer_config_defaults(self):
        """Test the ObservationDimensionPacker class and the Instrument-based
        systems for constructing it, in the case where the Instrument-defined
        default is used.
        """
        registry_config = RegistryConfig()
        registry_config["db"] = "sqlite://"
        registry = Registry.createFromConfig(registry_config)
        self.instrument.register(registry)
        config = DimensionPackerTestConfig()
        instrument_data_id = registry.expandDataId(instrument=self.name)
        record = instrument_data_id.records["instrument"]
        self.check_dimension_packers(
            registry.dimensions,
            # Test just one packer-construction signature here as that saves us
            # from having to insert dimension records for visits, exposures,
            # and detectors for this test.
            # Note that we don't need to pass any more than the instrument in
            # the data ID yet, because we're just constructing packers, not
            # calling their pack method.
            visit_packers=[config.packer.apply(instrument_data_id, is_exposure=False)],
            exposure_packers=[config.packer.apply(instrument_data_id, is_exposure=True)],
            n_detectors=record.detector_max,
            n_visits=record.visit_max,
            n_exposures=record.exposure_max,
        )

    def test_dimension_packer_config_override(self):
        """Test the ObservationDimensionPacker class and the Instrument-based
        systems for constructing it, in the case where configuration overrides
        the Instrument's default.
        """
        registry_config = RegistryConfig()
        registry_config["db"] = "sqlite://"
        registry = Registry.createFromConfig(registry_config)
        # Intentionally do not register instrument or insert any other
        # dimension records to ensure we don't need them in this mode.
        config = DimensionPackerTestConfig()
        config.packer["observation"].n_observations = 16
        config.packer["observation"].n_detectors = 4
        config.packer.name = "observation"
        instrument_data_id = DataCoordinate.standardize(instrument=self.name, universe=registry.dimensions)
        full_data_id = DataCoordinate.standardize(instrument_data_id, detector=1, visit=3, exposure=7)
        visit_data_id = full_data_id.subset(full_data_id.universe.extract(["visit", "detector"]))
        exposure_data_id = full_data_id.subset(full_data_id.universe.extract(["exposure", "detector"]))
        self.check_dimension_packers(
            registry.dimensions,
            # Note that we don't need to pass any more than the instrument in
            # the data ID yet, because we're just constructing packers, not
            # calling their pack method.
            visit_packers=[
                config.packer.apply(visit_data_id),
                config.packer.apply(full_data_id),
                config.packer.apply(instrument_data_id, is_exposure=False),
            ],
            exposure_packers=[
                config.packer.apply(instrument_data_id, is_exposure=True),
                config.packer.apply(exposure_data_id),
            ],
            n_detectors=config.packer["observation"].n_detectors,
            n_visits=config.packer["observation"].n_observations,
            n_exposures=config.packer["observation"].n_observations,
        )

    def check_dimension_packers(
        self,
        universe: DimensionUniverse,
        visit_packers: list[DimensionPacker],
        exposure_packers: list[DimensionPacker],
        n_detectors: int,
        n_visits: int,
        n_exposures: int,
    ) -> None:
        """Test the behavior of one or more dimension packers constructed by
        an instrument.

        Parameters
        ----------
        universe : `lsst.daf.butler.DimensionUniverse`
            Data model for butler data IDs.
        visit_packers : `list` [ `DimensionPacker` ]
            Packers with ``{visit, detector}`` dimensions to test.
        exposure_packers : `list` [ `DimensionPacker` ]
            Packers with ``{exposure, detector}`` dimensions to test.
        n_detectors : `int`
            Number of detectors all packers have been configured to reserve
            space for.
        n_visits : `int`
            Number of visits all packers in ``visit_packers`` have been
            configured to reserve space for.
        n_exposures : `int`
            Number of exposures all packers in ``exposure_packers`` have been
            configured to reserve space for.
        """
        full_data_id = DataCoordinate.standardize(
            instrument=self.name, detector=1, visit=3, exposure=7, universe=universe
        )
        visit_data_id = full_data_id.subset(full_data_id.universe.extract(["visit", "detector"]))
        exposure_data_id = full_data_id.subset(full_data_id.universe.extract(["exposure", "detector"]))
        for n, packer in enumerate(visit_packers):
            with self.subTest(n=n):
                packed_id, max_bits = packer.pack(full_data_id, returnMaxBits=True)
                self.assertEqual(packed_id, full_data_id["detector"] + full_data_id["visit"] * n_detectors)
                self.assertEqual(max_bits, math.ceil(math.log2(n_detectors * n_visits)))
                self.assertEqual(visit_data_id, packer.unpack(packed_id))
                with self.assertRaisesRegex(ValueError, f"Detector ID {n_detectors} is out of bounds"):
                    packer.pack(instrument=self.name, detector=n_detectors, visit=3)
                with self.assertRaisesRegex(ValueError, f"Visit ID {n_visits} is out of bounds"):
                    packer.pack(instrument=self.name, detector=1, visit=n_visits)
        for n, packer in enumerate(exposure_packers):
            with self.subTest(n=n):
                packed_id, max_bits = packer.pack(full_data_id, returnMaxBits=True)
                self.assertEqual(packed_id, full_data_id["detector"] + full_data_id["exposure"] * n_detectors)
                self.assertEqual(max_bits, math.ceil(math.log2(n_detectors * n_exposures)))
                self.assertEqual(exposure_data_id, packer.unpack(packed_id))
                with self.assertRaisesRegex(ValueError, f"Detector ID {n_detectors} is out of bounds"):
                    packer.pack(instrument=self.name, detector=n_detectors, exposure=3)
                with self.assertRaisesRegex(ValueError, f"Exposure ID {n_exposures} is out of bounds"):
                    packer.pack(instrument=self.name, detector=1, exposure=n_exposures)


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