# pylint: disable=missing-class-docstring,missing-function-docstring
# pylint: disable=line-too-long,invalid-name,attribute-defined-outside-init
"""pure python trajectory calculation backend"""

import math
import warnings

from typing_extensions import Union, List, override

from py_ballisticcalc.engines.base_engine import (
    BaseIntegrationEngine,
    _TrajectoryDataFilter,
    _WindSock,
    BaseEngineConfigDict, _ShotProps
)
from py_ballisticcalc.exceptions import RangeError
from py_ballisticcalc.logger import logger
from py_ballisticcalc.trajectory_data import TrajectoryData, TrajFlag
from py_ballisticcalc.vector import Vector

__all__ = ('VelocityVerletIntegrationEngine',)


class VelocityVerletIntegrationEngine(BaseIntegrationEngine[BaseEngineConfigDict]):
    """Velocity Verlet integration engine for ballistic calculations."""
    DEFAULT_TIME_STEP = 0.001

    def __init__(self, config: BaseEngineConfigDict):
        super().__init__(config)
        self.integration_step_count = 0

    @override
    def get_calc_step(self) -> float:
        return super().get_calc_step() * self.DEFAULT_TIME_STEP

    @override
    def _integrate(self, props: _ShotProps, maximum_range: float, record_step: float,
                   filter_flags: Union[TrajFlag, int], time_step: float = 0.0) -> List[TrajectoryData]:
        """
        Calculate trajectory for specified shot.

        Args:
            shot_info (_ShotProps): Information about the shot.
            maximum_range (float): Feet down range to stop calculation.
            record_step (float): Frequency (in feet down range) to record TrajectoryData.
            filter_flags (Union[TrajFlag, int]): Bitfield for requesting special trajectory points.
            time_step (float, optional): Maximum time (in seconds) between TrajectoryData records.

        Returns:
            List[TrajectoryData]: list of TrajectoryData
        """

        _cMinimumVelocity = self._config.cMinimumVelocity
        _cMaximumDrop = -abs(self._config.cMaximumDrop)  # Ensure it's negative
        _cMinimumAltitude = self._config.cMinimumAltitude

        ranges: List[TrajectoryData] = []  # Record of TrajectoryData points to return
        time: float = .0
        drag: float = .0
        mach: float = .0
        density_factor: float = .0

        # region Initialize wind-related variables to first wind reading (if any)
        wind_sock = _WindSock(props.winds)
        wind_vector = wind_sock.current_vector()
        # endregion

        # region Initialize position, velocity, and acceleration
        relative_speed = props.muzzle_velocity_fps
        # x: downrange distance, y: drop, z: windage
        range_vector = Vector(.0, -props.cant_cosine * props.sight_height_ft, -props.cant_sine * props.sight_height_ft)
        velocity_vector: Vector = Vector(
            math.cos(props.barrel_elevation_rad) * math.cos(props.barrel_azimuth_rad),
            math.sin(props.barrel_elevation_rad),
            math.cos(props.barrel_elevation_rad) * math.sin(props.barrel_azimuth_rad)
        ).mul_by_const(relative_speed)  # type: ignore
        # Acceleration:
        density_factor, mach = props.get_density_and_mach_for_altitude(range_vector.y)
        relative_velocity = velocity_vector - wind_vector
        relative_speed = relative_velocity.magnitude()
        drag = density_factor * relative_speed * props.drag_by_mach(relative_speed / mach)
        acceleration_vector = self.gravity_vector - drag * relative_velocity # type: ignore[operator]
        # endregion

        # Ensure one iteration when record step is smaller than calc_step
        min_step = min(props.calc_step, record_step)

        data_filter = _TrajectoryDataFilter(filter_flags=filter_flags, range_step=record_step,
                                            initial_position=range_vector, initial_velocity=velocity_vector,
                                            barrel_angle_rad=props.barrel_elevation_rad,
                                            look_angle_rad=props.look_angle_rad,
                                            time_step=time_step
        )

        # region Trajectory Loop
        warnings.simplefilter("once")  # used to avoid multiple warnings in a loop
        termination_reason = None
        last_recorded_range = 0.0
        start_integration_step_count = self.integration_step_count
        while (range_vector.x <= maximum_range + min_step) or (
                last_recorded_range <= maximum_range - 1e-6):
            self.integration_step_count += 1

            # Update wind reading at current point in trajectory
            if range_vector.x >= wind_sock.next_range:  # require check before call to improve performance
                wind_vector = wind_sock.vector_for_range(range_vector.x)

            # Update air density at current point in trajectory
            density_factor, mach = props.get_density_and_mach_for_altitude(range_vector.y)

            # region Check whether to record TrajectoryData row at current point
            if (data := data_filter.should_record(range_vector, velocity_vector, mach, time)) is not None:
                ranges.append(self._make_row(
                    props, data.time, data.position, data.velocity, data.mach, data_filter.current_flag)
                )
                last_recorded_range = data.position.x
            # endregion

            # region Ballistic calculation step (point-mass)
            # Air resistance seen by bullet is ground velocity minus wind velocity relative to ground
            relative_velocity = velocity_vector - wind_vector
            relative_speed = relative_velocity.magnitude()  # Velocity relative to air
            # Time step is normalized by velocity so that we take smaller steps when moving faster
            delta_time = props.calc_step
            # Drag is a function of air density and velocity relative to the air
            drag = density_factor * relative_speed * props.drag_by_mach(relative_speed / mach)

            # region Verlet integration
            # 1. Update position using acceleration from the current step
            range_vector += (velocity_vector * delta_time +                           # type: ignore[operator]
                             acceleration_vector * delta_time * delta_time * 0.5)     # type: ignore[operator]
            # 2. Calculate the new acceleration a(t+Δt) at the new position
            new_acceleration_vector = self.gravity_vector - drag * relative_velocity  # type: ignore[operator]
            # 3. Update velocity using the average of the old a(t) and new a(t+Δt) accelerations
            velocity_vector += (acceleration_vector + new_acceleration_vector) * 0.5 * delta_time  # type: ignore
            acceleration_vector = new_acceleration_vector
            velocity = velocity_vector.magnitude()  # Velocity relative to ground
            time += delta_time
            # endregion Verlet integration
            # endregion ballistic calculation step

            if (velocity < _cMinimumVelocity
                or range_vector.y < _cMaximumDrop
                or props.alt0_ft + range_vector.y < _cMinimumAltitude
            ):
                if velocity < _cMinimumVelocity:
                    termination_reason = RangeError.MinimumVelocityReached
                elif range_vector.y < _cMaximumDrop:
                    termination_reason = RangeError.MaximumDropReached
                else:
                    termination_reason = RangeError.MinimumAltitudeReached
                break
        # endregion
        if (data := data_filter.should_record(range_vector, velocity_vector, mach, time)) is not None:
            ranges.append(self._make_row(
                props, data.time, data.position, data.velocity, data.mach, data_filter.current_flag)
            )
        # Ensure that we have at least two data points in trajectory, or 1 if filter_flags==NONE
        # ... as well as last point if we had an incomplete trajectory
        if (filter_flags and ((len(ranges) < 2) or termination_reason)) or len(ranges) == 0:
            if len(ranges) > 0 and ranges[-1].time == time:  # But don't duplicate the last point.
                pass
            else:
                ranges.append(self._make_row(
                    props, time, range_vector, velocity_vector, mach, TrajFlag.NONE)
                )
        logger.debug(f"Velocity Verlet ran {self.integration_step_count - start_integration_step_count} iterations")
        if termination_reason is not None:
            raise RangeError(termination_reason, ranges)
        return ranges
