utils.py

import pandas as pd
import pygame as pg
import math

import os

import torch 

from collections import deque

from Const import*
from Class_Def import Celestial_Core, Planet, SpaceShip
device = DEVICE

def update_position(planets,rockets):
	
	total_fx = [0] * len(planets+rockets)
	total_fy = [0] * len(planets+rockets)

	for n,i in enumerate(planets+rockets):
		for j in planets+rockets:
			if i != j:
				fx, fy = i.attraction(j)
				total_fx[n-1] += fx
				total_fy[n-1] += fy
		if (n>=len(planets)):
			# print(i.vx, i.vy)
			if i.reactor1 :
				total_fx[n-1] += - i.accelerator()
			if i.reactor2 :
				total_fx[n-1] += - i.accelerator() * math.cos(2.*math.pi/3.)
				total_fy[n-1] += - i.accelerator() * math.sin(2.*math.pi/3.)
			if i.reactor3 :
				total_fx[n-1] += - i.accelerator() * math.cos(- 2.*math.pi/3.)
				total_fy[n-1] += - i.accelerator() * math.sin(- 2.*math.pi/3.)
				
	for n,i in enumerate(planets+rockets):
		i.vx += total_fx[n-1] / i.mass * TIMESTEP
		i.vy += total_fy[n-1] / i.mass * TIMESTEP

		i.x += i.vx * TIMESTEP
		i.y += i.vy * TIMESTEP
		i.orbit.append((i.x, i.y))
		while(len(i.orbit)>600):
			i.orbit.pop(0)
		
	for i in rockets:
		for j in planets:
			if (i.x - j.x)**2 + (i.y - j.y)**2 < (i.radius + j.radius)**2:
				del i
				# rockets.remove(i)
	return planets,rockets

def Solar_System():
	sun = Planet(0, 0, 1.98892 * 10**30, 695_700e3, 30, YELLOW)
	
	earth = Planet(-1 * AU , 0, 5.9742 * 10**24, 6_371e3, 16, BLUE)
	earth.vy = 29.783 * 1000 

	moon = Planet((-1 - 0.00256954)*AU , 0, 7.347 * 10**22, 1_737e3, 10, WHITE)
	moon.vy = (29.783+1.022) * 1000 
    
	mars = Planet(-1.524 * AU, 0, 6.39* 10**23 , 3_390e3, 12, RED)
	mars.vy = 24.077 * 1000

	mercury = Planet(0.387 * AU, 0, 3.30 * 10**23, 2_440e3, 8, DARK_GREY)
	mercury.vy = -47.4 * 1000

	venus = Planet(0.723 * AU, 0, 4.8685 * 10**24, 6_052e3, 14, WHITE)
	venus.vy = -35.02 * 1000
    
	jupiter = Planet(5.19 * AU, 0, 1.898* 10**27, 69_911e3, 25, JUP)
	jupiter.vy = -13.06 * 1000
    
	saturn = Planet(9.536 * AU, 0, 5.684 * 10**26, 58_232e3, 20, GREEN)
	saturn.vy = -9.64 * 1000

	neptune = Planet(30.07 * AU, 0, 1.024 * 10**26, 24_622e3, 18, BLUE)
	neptune.vy = -5.43 * 1000

	return  [sun,mercury,venus,earth,moon,mars,jupiter,saturn,neptune]

def Launch(planet):
	
	rocket = SpaceShip(planet.x , planet.y  + 0.00256954*AU, 30_000, 12, RED)
	rocket.vx = planet.vx -1.022*1000
	rocket.vy = planet.vy 
	
	return rocket



def get_distance(core1, core2):
	return math.sqrt( (core1.x-core2.x)**2 + (core1.y-core2.y)**2)

def draw_dist(win, shift_x , shift_y, core1, core2, Scale):

	x1 = (core1.x+shift_x) * Scale + WIDTH / 2 
	y1 = (core1.y+shift_y) * Scale + HEIGHT / 2

	x2 = (core2.x+shift_x) * Scale + WIDTH / 2 
	y2 = (core2.y+shift_y) * Scale + HEIGHT / 2
		
	pg.draw.line(win, WHITE, [x2, y2],[x1, y1], 3)

		

def is_orbiting(core1, core2, start=0):
	hist_dist = []
	#if mae between core1 and core 2 get_distance < 1% return True :
	
	for i in range(start,len(core1.orbit)):
		hist_dist.append(math.sqrt( (core1.orbit[i][0] - core2.orbit[i][0])**2 + (core1.orbit[i][1]-core2.orbit[i][1])**2))

	if (max(hist_dist) - min(hist_dist)) < 0.01 * min(hist_dist):
		return True
	else:
		return False

def magnitude(vx, vy):
	return math.sqrt(vx**2 + vy**2)



# --- Helper function to get the current state ---
def get_state(rocket, target_planet):
    """Calculates the state vector for the agent."""
    if rocket is None or target_planet is None:
        # Return a zero state if objects are missing
        return torch.zeros((1, STATE_SIZE), dtype=torch.float32, device=device)
    # Relative position (normalize by AU for consistency)
    rel_x = (rocket.x - target_planet.x) / AU
    rel_y = (rocket.y - target_planet.y) / AU
    # Relative velocity (needs normalization - find typical max velocity?)
    # Let's estimate a max velocity (e.g., around Earth orbital speed ~30km/s = 3e4 m/s)
    # Or maybe relative to ideal orbital speed? This needs tuning.
    norm_factor_vel = 5e4 # Tune this normalization factor
    rel_vx = (rocket.vx - target_planet.vx) / norm_factor_vel
    rel_vy = (rocket.vy - target_planet.vy) / norm_factor_vel
    
    # State vector: [dx, dy, dvx, dvy]
    state_list = [rel_x, rel_y, rel_vx, rel_vy] 
    
    # Add other components if STATE_SIZE in Const.py is > 4
    # e.g., state_list.append(rocket.fuel / rocket.initial_fuel) # If fuel becomes relevant

    # Convert to PyTorch tensor
    state_tensor = torch.tensor(state_list, dtype=torch.float32, device=device).unsqueeze(0) # Add batch dimension
    return state_tensor



def get_state_sequence(
    action_history: deque,
    rel_pos_x_history: deque,
    rel_pos_y_history: deque,
    rel_vel_x_history: deque,
    rel_vel_y_history: deque,
    seq_length: int = SEQ_LENGTH_DEFAULT,
    padding_value: float = PADDING_VALUE,
    action_padding_value: int = ACTION_PADDING_VALUE
):
    """
    Creates a state tensor representing the history of actions, relative positions,
    and relative velocities, suitable for sequence models like attention networks.

    Args:
        action_history (deque): Deque of past actions taken (integers).
        rel_pos_x_history (deque): Deque of past relative x positions (normalized).
        rel_pos_y_history (deque): Deque of past relative y positions (normalized).
        rel_vel_x_history (deque): Deque of past relative x velocities (normalized).
        rel_vel_y_history (deque): Deque of past relative y velocities (normalized).
        seq_length (int): The desired fixed sequence length for the output tensor.
        padding_value (float): Value used for padding position/velocity sequences.
        action_padding_value (int): Value used for padding the action sequence.

    Returns:
        torch.Tensor: A tensor of shape (1, 5, seq_length) containing the padded sequences.
                      Features are ordered [action, rel_pos_x, rel_pos_y, rel_vel_x, rel_vel_y].
        torch.Tensor: A boolean mask tensor of shape (1, seq_length) where True indicates
                      a real data point and False indicates padding. Useful for attention mechanisms.
    """
    sequences = []
    # Combine all history deques for easier iteration
    histories = [
        action_history, rel_pos_x_history, rel_pos_y_history, rel_vel_x_history, rel_vel_y_history
    ]
    # Define padding value for each sequence type
    paddings = [action_padding_value] + [padding_value] * 4

    actual_len = 0 # Store the actual length before padding

    for i, (history, pad_val) in enumerate(zip(histories, paddings)):
        # Convert deque to list and get the last seq_length items
        recent_history = list(history)[-seq_length:]

        # Store the actual length of the sequence before padding (use any history's length)
        if i == 0:
            actual_len = len(recent_history)

        # Calculate padding needed (prepends padding)
        pad_len = seq_length - len(recent_history)

        # Create the padded sequence, ensuring consistent data types
        if isinstance(pad_val, float):
            # Pad with float, convert history to float
            processed_history = [float(x) for x in recent_history]
            padded_sequence = ([pad_val] * pad_len) + processed_history
        else: # Assuming integer padding (like for actions)
             # Pad with int, convert history to int
            processed_history = [int(x) for x in recent_history]
            padded_sequence = ([pad_val] * pad_len) + processed_history

        sequences.append(padded_sequence)

    # Create the attention mask (True for real data, False for padding)
    mask_pad_len = seq_length - actual_len
    # Mask has shape (1, seq_length)
    mask = torch.tensor(([False] * mask_pad_len) + ([True] * actual_len), device=device, dtype=torch.bool).unsqueeze(0)

    # Stack sequences into a tensor
    # The resulting tensor shape will be (5, seq_length)
    # Ensure the tensor is float32 for neural network input
    state_tensor_stacked = torch.tensor(sequences, dtype=torch.float32, device=device)

    # Add the batch dimension -> (1, 5, seq_length)
    state_tensor_batch = state_tensor_stacked.unsqueeze(0)

    return state_tensor_batch, mask

# --- Placeholder function to calculate current relative state values ---
# You'll call this in your main loop before appending to the history deques.
def calculate_current_relative_state(rocket, target_planet):
    """Calculates the *current* relative state values."""
    if rocket is None or target_planet is None:
        # Return zeros or handle appropriately if objects are missing
        # Returning NaNs might be better to signal invalid state if not handled earlier
        return math.nan, math.nan, math.nan, math.nan

    # Relative position (normalize by AU) [cite: 122]
    rel_x = (rocket.x - target_planet.x) / AU
    rel_y = (rocket.y - target_planet.y) / AU
    # Relative velocity (normalized) [cite: 124]
    rel_vx = (rocket.vx - target_planet.vx) / NORM_FACTOR_VEL
    rel_vy = (rocket.vy - target_planet.vy) / NORM_FACTOR_VEL

    return rel_x, rel_y, rel_vx, rel_vy

# --- Function to reset a specific rocket (useful after crash/OOB) ---
def reset_rocket_state(rocket, planets):
	"""Resets a single rocket's position/velocity near a start planet."""
	start_planet_idx = 3 # Assuming Earth is index 3
	if start_planet_idx >= len(planets): start_planet_idx = 0 
	start_planet = planets[start_planet_idx]


	rocket.x = start_planet.x
	rocket.y = start_planet.y + 0.00256954 * AU
	rocket.vx = start_planet.vx -1.022*1000
	rocket.vy = start_planet.vy 
	# Reset reactors and potentially other state if needed
	rocket.reactor1 = False
	rocket.reactor2 = False
	rocket.reactor3 = False
	# Reset mass/fuel if tracked later
	# rocket.fuel = rocket.initial_fuel 
	# rocket.mass = rocket.structure_mass + rocket.fuel

	print("Rocket state reset.")
	return rocket

# --- Function to reset the environment for a new episode ---
def reset_environment(planets):
    """Resets rocket position/velocity for a new episode."""
    # Choose a starting planet (e.g., Earth)
    start_planet_idx = 3 # Assuming Earth is index 3
    if start_planet_idx >= len(planets):
        start_planet_idx = 0 # Default to Sun if Earth doesn't exist

    start_planet = planets[start_planet_idx]

    # Start in a somewhat random position near the starting planet
    angle = random.uniform(0, 2 * math.pi)
    dist_factor = random.uniform(1.5, 3.0) # Start further away than just radius
    start_dist = start_planet.radius * dist_factor if hasattr(start_planet, 'radius') else 1e7
    
    start_x = start_planet.x + start_dist * math.cos(angle)
    start_y = start_planet.y + start_dist * math.sin(angle)

    # Create the rocket (assuming SpaceShip class exists)
    # Give it some initial velocity, perhaps tangential to start planet
    tangential_speed = math.sqrt(G * start_planet.mass / start_dist) * 0.5 # Start slower than orbit
    start_vx = start_planet.vx - tangential_speed * math.sin(angle)
    start_vy = start_planet.vy + tangential_speed * math.cos(angle)
    
    rocket = SpaceShip(start_x, start_y, 500, 5, BLUE) # Use default parameters
    rocket.vx = start_vx
    rocket.vy = start_vy

    # Reset target (optional, could keep same target)
    # Target = random.randint(0, len(planets) - 1) 

    return [rocket] # Return list containing the single rocket




def save_distance_data(all_acquisitions, save_folder="metrics"):
    """Saves recorded distance data to a CSV file using pandas.
       Each column represents one acquisition (run between resets).
    """
    if not all_acquisitions or not any(all_acquisitions): # Check if empty or contains only empty lists
        print("No distance data recorded to save.")
        return

    try:
        # Create timestamped folder
        filename = os.path.join(save_folder, "distances_log.csv") # New filename maybe

        # Prepare data for DataFrame: dictionary where keys are column names
        data_dict = {}
        for i, acq in enumerate(all_acquisitions):
            # Only include non-empty acquisitions if desired, or handle them all
            if acq: # Ensure the acquisition list is not empty
                 data_dict[f"Acquisition_{i+1}"] = pd.Series(acq)
            # else: # Optionally handle empty acquisitions if needed
            #     data_dict[f"Acquisition_{i+1}"] = pd.Series([]) # Add empty series

        if not data_dict:
             print("All recorded acquisitions were empty. Nothing to save.")
             return

        # Create DataFrame directly from dictionary of Series
        # Pandas handles unequal lengths by padding with NaN automatically
        df = pd.DataFrame(data_dict)

        # Save DataFrame to CSV
        df.to_csv(filename, index=False, float_format='%.2f') # Save without index, format floats

        print(f"Distance data saved to {filename} using pandas.")

    except ImportError:
         print("Error: pandas library not found. Please install it (`pip install pandas`) to use this save function.")
    except Exception as e:
        print(f"Error saving distance data using pandas: {e}")
        import traceback
        traceback.print_exc()