import numpy as np

from pydrodynamics.utils import state_object_to_array, world_to_body

class Environment:
	"""
		Calculates the external forces and moments generated by the surrounding water, for example:
		- Buoyancy
		- Gravity
		- Drag
		- Lift
		- Added mass
	"""
	def __init__(self, params):
		self.params = params

		# Construct drag matrix
		projected_area_vector = np.array([
			params.get('apx'),
			params.get('apy'),
			params.get('apz'),
			params.get('apz'),
			params.get('apz'),
			params.get('apy'),
		]).T
		drag_xu = params.get('drag_xu')
		drag_xv = params.get('drag_xv')
		drag_xw = params.get('drag_xw')
		drag_xp = params.get('drag_xp')
		drag_xq = params.get('drag_xq')
		drag_xr = params.get('drag_xr')
		drag_yu = params.get('drag_yu')
		drag_yv = params.get('drag_yv')
		drag_yw = params.get('drag_yw')
		drag_yp = params.get('drag_yp')
		drag_yq = params.get('drag_yq')
		drag_yr = params.get('drag_yr')
		drag_zu = params.get('drag_zu')
		drag_zv = params.get('drag_zv')
		drag_zw = params.get('drag_zw')
		drag_zp = params.get('drag_zp')
		drag_zq = params.get('drag_zq')
		drag_zr = params.get('drag_zr')
		drag_ku = params.get('drag_ku')
		drag_kv = params.get('drag_kv')
		drag_kw = params.get('drag_kw')
		drag_kp = params.get('drag_kp')
		drag_kq = params.get('drag_kq')
		drag_kr = params.get('drag_kr')
		drag_mu = params.get('drag_mu')
		drag_mv = params.get('drag_mv')
		drag_mw = params.get('drag_mw')
		drag_mp = params.get('drag_mp')
		drag_mq = params.get('drag_mq')
		drag_mr = params.get('drag_mr')
		drag_nu = params.get('drag_nu')
		drag_nv = params.get('drag_nv')
		drag_nw = params.get('drag_nw')
		drag_np = params.get('drag_np')
		drag_nq = params.get('drag_nq')
		drag_nr = params.get('drag_nr')
		self.drag_matrix = np.array([
			[drag_xu, drag_xv, drag_xw, drag_xp, drag_xq, drag_xr],
			[drag_yu, drag_yv, drag_yw, drag_yp, drag_yq, drag_yr],
			[drag_zu, drag_zv, drag_zw, drag_zp, drag_zq, drag_zr],
			[drag_ku, drag_kv, drag_kw, drag_kp, drag_kq, drag_kr],
			[drag_mu, drag_mv, drag_mw, drag_mp, drag_mq, drag_mr],
			[drag_nu, drag_nv, drag_nw, drag_np, drag_nq, drag_nr]
		]) * -0.5 * params.get('density') * projected_area_vector

	def calculate(self, state):
		"""
			Calculate the external forces and moments acting on the vehicle, all in body-fixed frame:

			Input:
                state: The current state of the vehicle, including position, orientation, linear velocity, angular velocity, and voltage.
            Output:
				tau_env: External forces and moments from the environment in (6,) format, including:
					- tau_hydrostatic: Hydrostatic forces and moments from gravity and buoyancy
					- tau_drag: Drag stemming from hydrodynamic damping
					- tau_lift: Lift forces and moments
					- tau_added_mass: Added mass forces and moments
		"""
		# Convert state to numpy array
		state = state_object_to_array(state)

		tau_hydrostatic = self.calculate_hydrostatic(state)
		tau_drag = self.calculate_drag(state)
		tau_lift = 0 # self.calculate_lift(state)
		tau_added_mass = 0 # self.calculate_added_mass(state)

		tau_env = tau_hydrostatic + tau_drag + tau_lift + tau_added_mass
		return tau_env

	def calculate_hydrostatic(self, state):
		"""
			Calculate the hydrostatic forces and moments acting on the vehicle.
			Includes buoyancy and gravity.

			Input:
				state: The current state of the vehicle.
			Output:
				tau_hydrostatic: Hydrostatic forces and moments in body-fixed frame.
		"""
		# Calculate rotation matrix from world frame to body-fixed frame
		r_trans = world_to_body(state)[0]

		# Calculate gravity and buoyancy forces in body-fixed frame, up is negative z
		gravity_force = np.matmul(r_trans, np.array([0, 0, self.params.get('mass') * self.params.get('gravity')]))
		buoyancy_force = np.matmul(r_trans, np.array([0, 0, -1 * self.params.get('volume') * self.params.get('density') * self.params.get('gravity')]))

		# Make vector for CoM and CoB
		com = np.array([self.params.get('xg'), self.params.get('yg'), self.params.get('zg')])
		cob = np.array([self.params.get('xb'), self.params.get('yb'), self.params.get('zb')])

		# Calculate the total hydrostatic force
		hydrostatic_forces = gravity_force + buoyancy_force
		hydrostatic_moments = np.cross(com, gravity_force) + np.cross(cob, buoyancy_force)

		# Merge forces and moments into a single vector
		return np.concatenate((hydrostatic_forces, hydrostatic_moments), axis=0).reshape((6,))
	
	def calculate_drag(self, state):
		"""
			Calculate the drag forces and moments acting on the vehicle.

			Input:
				state: The current state of the vehicle.
			Output:
				tau_drag: Drag forces and moments in body-fixed frame.
		"""
		# Extract velocities and convert angular velocities to radians
		vels = state[6:12].reshape((6,))
		vels[3:6] = np.radians(vels[3:6])
		vels *= np.abs(vels)

		# Calculate drag forces and moments
		return np.matmul(self.drag_matrix, vels).reshape((6,))
