"""Models for ground imaging dynamics."""

from typing import TYPE_CHECKING, Iterable, Optional

import numpy as np
from Basilisk.simulation import (
    groundLocation,
    partitionedStorageUnit,
    simpleInstrument,
    simplePowerSink,
    simpleStorageUnit,
    spaceToGroundTransmitter,
)
from Basilisk.utilities import orbitalMotion

from bsk_rl.sim import world
from bsk_rl.sim.dyn import BasicDynamicsModel
from bsk_rl.utils import actuator_primitives as aP
from bsk_rl.utils.functional import aliveness_checker, default_args

if TYPE_CHECKING:  # pragma: no cover
    from bsk_rl.sim.world import WorldModel


class ImagingDynModel(BasicDynamicsModel):
    """Equips the satellite with an instrument, storage unit, and transmitter."""

    def __init__(self, *args, **kwargs) -> None:
        """Equips the satellite with an instrument, storage unit, and transmitter.

        This dynamics model is used with :class:`~bsk_rl.sats.ImagingSatellite`. It
        provides the satellite with the ability to take images of a point target. To
        enable downlink, use :class:`GroundStationDynModel` and :class:`~bsk_rl.sim.world.GroundStationWorldModel`.
        """
        super().__init__(*args, **kwargs)

    @property
    def storage_level(self):
        """Storage level [bits]."""
        return self.storageUnit.storageUnitDataOutMsg.read().storageLevel

    @property
    def storage_level_fraction(self):
        """Storage level as a fraction of capacity."""
        return self.storage_level / self.storageUnit.storageCapacity

    def _setup_dynamics_objects(self, **kwargs) -> None:
        super()._setup_dynamics_objects(**kwargs)
        self.setup_instrument_power_sink(**kwargs)
        self.setup_transmitter_power_sink(**kwargs)
        self.setup_instrument(**kwargs)
        self.setup_transmitter(**kwargs)
        self.setup_storage_unit(**kwargs)
        self.setup_imaging_target(**kwargs)

    @default_args(instrumentBaudRate=8e6)
    def setup_instrument(
        self, instrumentBaudRate: float, priority: int = 895, **kwargs
    ) -> None:
        """Set up the instrument data collection model.

        Args:
            instrumentBaudRate: [bits] Data generated by an image.
            priority: Model priority.
            kwargs: Passed to other setup functions.
        """
        self.instrument = simpleInstrument.SimpleInstrument()
        self.instrument.ModelTag = "instrument" + self.satellite.name
        self.instrument.nodeBaudRate = (
            instrumentBaudRate / self.dyn_rate
        )  # makes imaging instantaneous
        self.instrument.nodeDataName = "Instrument" + self.satellite.name
        self.simulator.AddModelToTask(
            self.task_name, self.instrument, ModelPriority=priority
        )

    @default_args(
        transmitterBaudRate=-8e6, transmitterNumBuffers=100, transmitterPacketSize=None
    )
    def setup_transmitter(
        self,
        transmitterBaudRate: float,
        instrumentBaudRate: float,
        transmitterNumBuffers: int,
        transmitterPacketSize: Optional[float] = None,
        priority: int = 798,
        **kwargs,
    ) -> None:
        """Set up the transmitter model for downlinking data.

        Args:
            transmitterBaudRate: [baud] Rate of data downlink. Should be negative.
            instrumentBaudRate: [bits] Image size, used to set packet size if ``transmitterPacketSize`` is not specified.
            transmitterNumBuffers: Number of transmitter buffers
            transmitterPacketSize: [bits] Minimum data amount to be downlinked for a partition. If ``None``, set as ``instrumentBaudRate``.
            priority: Model priority.
            kwargs: Passed to other setup functions.
        """
        if transmitterBaudRate > 0:
            self.logger.warning("transmitterBaudRate should probably be negative.")
        self.transmitter = spaceToGroundTransmitter.SpaceToGroundTransmitter()
        self.transmitter.ModelTag = "transmitter" + self.satellite.name
        self.transmitter.nodeBaudRate = transmitterBaudRate  # baud
        if transmitterPacketSize is None:
            # set packet size equal to the size of a single image
            self.transmitter.packetSize = -instrumentBaudRate  # bits
        else:
            self.transmitter.packetSize = transmitterPacketSize  # bits
        self.transmitter.numBuffers = transmitterNumBuffers
        self.simulator.AddModelToTask(
            self.task_name, self.transmitter, ModelPriority=priority
        )

    @default_args(instrumentPowerDraw=-30.0)
    def setup_instrument_power_sink(
        self, instrumentPowerDraw: float, priority: int = 897, **kwargs
    ) -> None:
        """Set the instrument power sink parameters.

        The instrument draws power when in an imaging task, representing the power cost
        of operating the instrument.

        Args:
            instrumentPowerDraw: [W] Power draw when instrument is enabled.
            priority: Model priority.
            kwargs: Passed to other setup functions.
        """
        if instrumentPowerDraw > 0:
            self.logger.warning(
                "instrumentPowerDraw should probably be zero or negative."
            )
        self.instrumentPowerSink = simplePowerSink.SimplePowerSink()
        self.instrumentPowerSink.ModelTag = "insPowerSink" + self.satellite.name
        self.instrumentPowerSink.nodePowerOut = instrumentPowerDraw
        self.simulator.AddModelToTask(
            self.task_name, self.instrumentPowerSink, ModelPriority=priority
        )
        self.powerMonitor.addPowerNodeToModel(self.instrumentPowerSink.nodePowerOutMsg)

    @default_args(transmitterPowerDraw=-15.0)
    def setup_transmitter_power_sink(
        self, transmitterPowerDraw: float, priority: int = 896, **kwargs
    ) -> None:
        """Set the transmitter power sink parameters.

        The transmitter draws power when in a downlink task, representing the power cost
        of downlinking data.

        Args:
            transmitterPowerDraw: [W] Power draw when transmitter is enabled.
            priority: Model priority.
            kwargs: Passed to other setup functions.
        """
        if transmitterPowerDraw > 0:
            self.logger.warning(
                "transmitterPowerDraw should probably be zero or negative."
            )
        self.transmitterPowerSink = simplePowerSink.SimplePowerSink()
        self.transmitterPowerSink.ModelTag = "transPowerSink" + self.satellite.name
        self.transmitterPowerSink.nodePowerOut = transmitterPowerDraw
        self.simulator.AddModelToTask(
            self.task_name, self.transmitterPowerSink, ModelPriority=priority
        )
        self.powerMonitor.addPowerNodeToModel(self.transmitterPowerSink.nodePowerOutMsg)

    @default_args(
        dataStorageCapacity=20 * 8e6,
        bufferNames=None,
        storageUnitValidCheck=False,
        storageInit=0,
    )
    def setup_storage_unit(
        self,
        dataStorageCapacity: int,
        storageUnitValidCheck: bool,
        storageInit: int,
        transmitterNumBuffers: Optional[int] = None,
        bufferNames: Optional[Iterable[str]] = None,
        priority: int = 699,
        **kwargs,
    ) -> None:
        """Configure the storage unit and its buffers.

        Separate buffers can be used to track imaging of different targets. Often, the
        buffer names will be set up by satellite based on the scenario configuration.

        Args:
            dataStorageCapacity: [bits] Maximum data that can be stored.
            transmitterNumBuffers: Number of unit buffers. Not necessary if ``bufferNames``
                are given.
            bufferNames: List of buffer names to use. Named by number if ``None``.
            storageUnitValidCheck: If ``True``, enforce that the storage level is below
                the storage capacity when checking aliveness.
            storageInit: [bits] Initial storage level.
            priority: Model priority.
            kwargs: Passed to other setup functions.
        """
        self.storageUnit = partitionedStorageUnit.PartitionedStorageUnit()
        self.storageUnit.ModelTag = "storageUnit" + self.satellite.name
        self.storageUnit.storageCapacity = dataStorageCapacity  # bits
        self.storageUnit.addDataNodeToModel(self.instrument.nodeDataOutMsg)
        self.storageUnit.addDataNodeToModel(self.transmitter.nodeDataOutMsg)
        self.storageUnitValidCheck = storageUnitValidCheck
        # Add all of the targets to the data buffer
        if bufferNames is None:
            for buffer_idx in range(transmitterNumBuffers):
                self.storageUnit.addPartition(str(buffer_idx))
        else:
            if transmitterNumBuffers is not None and transmitterNumBuffers != len(
                bufferNames
            ):
                raise ValueError(
                    "transmitterNumBuffers cannot be different than len(bufferNames)."
                )
            for buffer_name in bufferNames:
                self.storageUnit.addPartition(buffer_name)

        if storageInit != 0:
            if storageInit > dataStorageCapacity or storageInit < 0:
                self.logger.warning(
                    f"Initial storage level {storageInit} incompatible with its capacity {dataStorageCapacity}."
                )
            self.storageUnit.setDataBuffer(["STORED DATA"], [int(storageInit)])

        # Add the storage unit to the transmitter
        self.transmitter.addStorageUnitToTransmitter(
            self.storageUnit.storageUnitDataOutMsg
        )

        self.simulator.AddModelToTask(
            self.task_name, self.storageUnit, ModelPriority=priority
        )

    @aliveness_checker
    def data_storage_valid(self) -> bool:
        """Check that the buffer has not run out of space.

        Only is checked if ``storageUnitValidCheck`` is ``True``; otherwise, a full storage
        unit will prevent additional data from being stored but will not cause the satellite
        to be considered dead.
        """
        storage_check = self.storageUnitValidCheck
        if storage_check:
            return self.storage_level < self.storageUnit.storageCapacity or np.isclose(
                self.storage_level, self.storageUnit.storageCapacity
            )
        else:
            return True

    @default_args(
        groundLocationPlanetRadius=orbitalMotion.REQ_EARTH * 1e3,
        imageTargetMinimumElevation=np.radians(45.0),
        imageTargetMaximumRange=-1,
    )
    def setup_imaging_target(
        self,
        groundLocationPlanetRadius: float,
        imageTargetMinimumElevation: float,
        imageTargetMaximumRange: float,
        priority: int = 2000,
        **kwargs,
    ) -> None:
        """Add a generic imaging target to dynamics.

        The target must be updated with a particular location when used.

        Args:
            groundLocationPlanetRadius: [m] Radius of ground locations from center of planet.
            imageTargetMinimumElevation: [rad] Minimum elevation angle from target to
                satellite when imaging.
            imageTargetMaximumRange: [m] Maximum range from target to satellite when
                imaging. -1 to disable.
            priority: Model priority.
            kwargs: Passed to other setup functions.
        """
        self.imagingTarget = groundLocation.GroundLocation()
        self.imagingTarget.ModelTag = "ImagingTarget"
        self.imagingTarget.planetRadius = groundLocationPlanetRadius
        self.imagingTarget.specifyLocation(0.0, 0.0, 1000.0)
        self.imagingTarget.planetInMsg.subscribeTo(
            self.world.gravFactory.spiceObject.planetStateOutMsgs[self.world.body_index]
        )
        self.imagingTarget.minimumElevation = imageTargetMinimumElevation
        self.imagingTarget.maximumRange = imageTargetMaximumRange

        self.simulator.AddModelToTask(
            self.world.world_task_name,
            self.imagingTarget,
            ModelPriority=priority,
        )
        self.imagingTarget.addSpacecraftToModel(self.scObject.scStateOutMsg)

    def reset_for_action(self) -> None:
        """Shut off power sinks unless the transmitter or instrument is being used."""
        super().reset_for_action()
        self.transmitter.dataStatus = 0
        self.transmitterPowerSink.powerStatus = 0
        self.instrumentPowerSink.powerStatus = 0


class ContinuousImagingDynModel(ImagingDynModel):
    """Equips the satellite for continuous nadir imaging."""

    def __init__(self, *args, **kwargs) -> None:
        """Equips the satellite for continuous nadir imaging.

        Equips satellite with an instrument, storage unit, and transmitter
        for continuous nadir imaging. A single data buffer is used for storage, and data
        is accumulated continuously while imaging. The imaging target is fixed at the
        center of the Earth for nadir imaging.
        """
        super().__init__(*args, **kwargs)

    @default_args(instrumentBaudRate=8e6)
    def setup_instrument(
        self, instrumentBaudRate: float, priority: int = 895, **kwargs
    ) -> None:
        """Set up the continuous instrument model.

        Args:
            instrumentBaudRate: [baud] Data generation rate step when continuously imaging.
            priority: Model priority.
            kwargs: Passed to other setup functions.
        """
        self.instrument = simpleInstrument.SimpleInstrument()
        self.instrument.ModelTag = "instrument" + self.satellite.name
        self.instrument.nodeBaudRate = instrumentBaudRate  # make imaging instantaneous
        self.instrument.nodeDataName = "Instrument" + self.satellite.name
        self.simulator.AddModelToTask(
            self.task_name, self.instrument, ModelPriority=priority
        )

    @default_args(
        dataStorageCapacity=20 * 8e6,
        storageUnitValidCheck=False,
        storageInit=0,
    )
    def setup_storage_unit(
        self,
        dataStorageCapacity: int,
        storageUnitValidCheck: bool,
        storageInit: int,
        priority: int = 699,
        **kwargs,
    ) -> None:
        """Configure the storage unit and its buffers.

        Args:
            dataStorageCapacity: [bits] Maximum data that can be stored.
            storageUnitValidCheck: If True, check that the storage level is below the
                storage capacity.
            storageInit: [bits] Initial storage level.
            priority: Model priority.
            kwargs: Passed to other setup functions.
        """
        self.storageUnit = simpleStorageUnit.SimpleStorageUnit()
        self.storageUnit.ModelTag = "storageUnit" + self.satellite.name
        self.storageUnit.storageCapacity = dataStorageCapacity  # bits
        self.storageUnit.addDataNodeToModel(self.instrument.nodeDataOutMsg)
        self.storageUnit.addDataNodeToModel(self.transmitter.nodeDataOutMsg)
        self.storageUnitValidCheck = storageUnitValidCheck
        if storageInit > dataStorageCapacity or storageInit < 0:
            self.logger.warning(
                f"Initial storage level {storageInit} incompatible with its capacity {dataStorageCapacity}."
            )
        self.storageUnit.setDataBuffer(storageInit)

        # Add the storage unit to the transmitter
        self.transmitter.addStorageUnitToTransmitter(
            self.storageUnit.storageUnitDataOutMsg
        )

        self.simulator.AddModelToTask(
            self.task_name, self.storageUnit, ModelPriority=priority
        )

    @default_args(imageTargetMaximumRange=-1)
    def setup_imaging_target(
        self,
        imageTargetMaximumRange: float = -1,
        priority: int = 2000,
        **kwargs,
    ) -> None:
        """Add a imaging target at the center of the Earth.

        Args:
            imageTargetMaximumRange: [m] Maximum range from target to satellite when
                imaging. -1 to disable.
            priority: Model priority.
            kwargs: Passed to other setup functions.
        """
        self.imagingTarget = groundLocation.GroundLocation()
        self.imagingTarget.ModelTag = "scanningTarget"
        self.imagingTarget.planetRadius = 1e-6
        self.imagingTarget.specifyLocation(0, 0, 0)
        self.imagingTarget.planetInMsg.subscribeTo(
            self.world.gravFactory.spiceObject.planetStateOutMsgs[self.world.body_index]
        )
        self.imagingTarget.minimumElevation = np.radians(-90)
        self.imagingTarget.maximumRange = imageTargetMaximumRange

        self.simulator.AddModelToTask(
            self.world.world_task_name,
            self.imagingTarget,
            ModelPriority=priority,
        )
        self.imagingTarget.addSpacecraftToModel(self.scObject.scStateOutMsg)


class GroundStationDynModel(ImagingDynModel):
    """Model that connects satellite to world ground stations."""

    def __init__(self, *args, **kwargs) -> None:
        """Model that connects satellite to world ground stations.

        This model enables the use of ground stations defined in :class:`~bsk_rl.sim.world.GroundStationWorldModel`
        for data downlink.
        """
        super().__init__(*args, **kwargs)

    @classmethod
    def _requires_world(cls) -> list[type["WorldModel"]]:
        return super()._requires_world() + [world.GroundStationWorldModel]

    def _setup_dynamics_objects(self, **kwargs) -> None:
        super()._setup_dynamics_objects(**kwargs)
        self.setup_ground_station_locations()

    def setup_ground_station_locations(self) -> None:
        """Connect the transmitter to ground stations."""
        for groundStation in self.world.groundStations:
            groundStation.addSpacecraftToModel(self.scObject.scStateOutMsg)
            self.transmitter.addAccessMsgToTransmitter(groundStation.accessOutMsgs[-1])

            if hasattr(self.satellite, "add_location_for_access_checking"):
                self.satellite.add_location_for_access_checking(
                    object=groundStation.ModelTag,
                    r_LP_P=np.array(groundStation.r_LP_P_Init).flatten(),
                    min_elev=groundStation.minimumElevation,
                    type="ground_station",
                )


__doc_title__ = "Ground Imaging"
__all__ = [
    "ImagingDynModel",
    "ContinuousImagingDynModel",
    "GroundStationDynModel",
]
