SPATZ/spatz/sensors/antenna/tx_gain.py

from numpy.typing import ArrayLike
from typing import List, AnyStr
from numpy import matrix
from typing import List
import re
from io import StringIO
import numpy as np

import pandas as pd
import math

from spatz.sensors import Sensor
from spatz.simulation import Simulation
from spatz.transforms import Transform
from spatz.dataset import Dataset
from spatz.logger import Logger
import time

GAIN_NAME = "Abs(Dir.)"
#GAIN_NAME = "Abs(Gain)"

'''
Class representing a CST gain pattern
This (and the sensor below) follow the convetions laid out by https://www.antenna-theory.com/basics/radpattern.php.
I.e, theta represents the elevation angle and goes from 0 to 180 deg, Phi represents the azimuth angle.

The data is interpolated, you will have to specify the step size for this to work correctly.
'''

class GainPattern():
    def __init__(self, filepath: str, step_size: int): 
        self._stepsize = step_size
        # This is a cursed parser. If it breaks, though luck.
        with open(filepath,"r") as file:
            # Read Header
            header = file.readline()
            header = re.sub(r'\[(.*?)\]',",",header).replace(" ","").replace(",\n",'\n')

            # Discard ---- line
            file.readline()

            # Parse to DF
            lines = file.readlines()
            clean_csv = [header]
            start_time = time.time()
            num_lines = len(lines)
            for i,line in enumerate(lines):       
                if(i % step_size == 0 or i == num_lines-1):
                    cleaned = re.sub(r'\s+',',',line).removeprefix(',').removesuffix(',').strip()
                    clean_csv.append(cleaned + '\n')

            clean_csv = ''.join(clean_csv)
            filelike = StringIO(clean_csv)
            self._df = pd.read_csv(filelike)
            
            print(f"Processed {num_lines} lines in {(time.time()-start_time):.1f}s.")
            print(f"Used {num_lines // step_size} lines due to step size")

            self._df.to_csv("gainpattern.csv")




    def get_phi_cut(self, phi:float) -> ArrayLike: #Return farfield cut with phi = const (Looking from the side)
            assert 0 <= phi < 180
            sub_df  = self._df.loc[self._df["Phi"] == phi]
            angles = sub_df["Theta"]              
            gain = sub_df[GAIN_NAME]
            return angles,gain

    def get_theta_cut(self, theta:float) -> ArrayLike: #Return farfield cut with theta = const (looking from the top)
            assert 0<= theta < 180
            sub_df_left = self._df.loc[self._df["Theta"] == theta]
            angles_l = sub_df_left["Phi"]
            gain_l = sub_df_left[GAIN_NAME]

            sub_df_right = self._df.loc[self._df["Theta"] == (360-theta)]
            angles_r = sub_df_right["Phi"]+180
            gain_r = sub_df_right[GAIN_NAME]

            angles = pd.concat([angles_l,angles_r])
            gain = pd.concat([gain_l,gain_r])

            return angles,gain


        
    def __get_gain_internal(self,phi_step:float,theta_step:float):
        assert phi_step%self._stepsize ==0
        assert theta_step%self._stepsize==0

        row = self._df.loc[(self._df["Theta"] == theta_step) & (self._df["Phi"] == phi_step)].iloc[0]
        return row[GAIN_NAME]

    def get_gain(self, phi, theta) -> float:
        assert 0 <= phi < 360
        assert 0 <= theta < 180
        
        #Interpolate using binlinear interpolation https://en.wikipedia.org/wiki/Bilinear_interpolation
        phi_lower = math.floor(phi/self._stepsize)*self._stepsize
        phi_upper = phi_lower + self._stepsize
        theta_lower = math.floor(theta/self._stepsize)*self._stepsize
        theta_upper = theta_lower + self._stepsize

        G11 = self.__get_gain_internal(phi_lower,theta_lower)
        G12 = self.__get_gain_internal(phi_lower,theta_upper)
        G21 = self.__get_gain_internal(phi_upper,theta_lower)
        G22 = self.__get_gain_internal(phi_upper,theta_upper)

        v1 = np.array([phi_upper-phi,phi-phi_lower])
        v2 = np.array([[theta_upper-theta],[theta-theta_lower]])
        A = np.array([[G11,G12],[G21,G22]])

        interpolated = 1/(self._stepsize*self._stepsize) * v1 @ A @ v2

        return interpolated[0]
        

'''
Sensor to simulate TX antenna gain in direction of ground station

Returns the gain in dBi per timestep.

'''

class AntennaTxGain(Sensor):

    def __init__(self, dataset: Dataset, logger: Logger, transforms: List[Transform] = [], gain_pattern_path = "data/gain_pattern/farfield.txt", step_size=1):
        super().__init__(dataset, logger, transforms)
        self._pattern = GainPattern(gain_pattern_path,step_size)

    def _get_data(self) -> ArrayLike | float:
        magic_matrix = np.array([
             [0,1,0],
             [1,0,0],
             [0,0,-1]
        ]) 
        
        # Get current position of rocket in FL Frame (Launcher Frame).
        pos_fl = self._dataset.fetch_values(['x', 'y', 'z']) #X,Y,Z is in FL (Launcher frame) -> Z is up, X is east
        
        gs_offset_fl = np.array([-1810,-1500,100]) #Radar hill is approx 1.81km west. 1.5km south, 100higher
        rocket_to_gs_fl = pos_fl-gs_offset_fl
        rocket_to_gs_fl_n = rocket_to_gs_fl/np.linalg.norm(rocket_to_gs_fl)

        # Rocket in body frame is simply [1,0,0]^T by definition
        rocket_b = np.array([1,0,0])
        rocket_fl = magic_matrix @ np.linalg.inv(self._dataset.launch_rail_to_body()) @ rocket_b
        rocket_fl_n = rocket_fl / np.linalg.norm(rocket_fl)

        

        # Angle between rocket and pos returns elevation angle (Phi). Assume a rotation of 0° for now to get theta
        theta = 180-np.rad2deg(np.arccos(np.clip(np.dot(rocket_to_gs_fl_n,rocket_fl_n),-1.0,1.0))) #Clip trick from: https://stackoverflow.com/questions/2827393/angles-between-two-n-dimensional-vectors-in-python
        
        self._log("rocket_x",rocket_fl_n[0])
        self._log("rocket_y",rocket_fl_n[1])
        self._log("rocket_z",rocket_fl_n[2])
        self._log("pos_x",rocket_to_gs_fl_n[0])
        self._log("pos_y",rocket_to_gs_fl_n[1])
        self._log("pos_z",rocket_to_gs_fl_n[2])
        self._log("theta",theta)
        #return phi
        
        #Get Theta cut for this angle
        angles, gains = self._pattern.get_theta_cut(np.round(theta))

        min_gain = np.min(gains)
        #min_ix = np.argmin(gains)
        #min_angle = angles[min_ix]
        #self._log("works_case_angle",min_angle)  


        #min_gain = self._pattern.get_gain(45,theta)
        
        # Fetch gain in this direction
        return min_gain

    def _sensor_specific_effects(self, x: ArrayLike) -> ArrayLike:
        return x
    
    def _get_name(self) -> AnyStr:
        return 'antenna/tx_gain'
        

if  __name__ == '__main__':
    pattern = GainPattern("data/gain_pattern/farfield_all.txt")
    print(pattern.get_gain(0,12))
    print(pattern.get_gain(0,16))
    print(pattern.get_gain(6,12))
    print(pattern.get_gain(0,10))
    print(pattern.get_theta_cut(90))