import PospacAnalysis as pa 
import pandas as pd 
import shutil
import numpy as np

"""
This script takes the saved csv with the adjustment parameters, makes a copy of the truth file, and adjusts the LLH of the copy 
"""

#Get adjustment parameters by asking user which directory to work with
workingDir = pa.getDirFromUser()
plotDir = pa.getDirFromUser(workingDir + '/Plots')
adjustFile = workingDir + '/Plots/' + plotDir + '/adjustment_params.csv'
df = pd.read_csv(adjustFile) #Convert the csv to a dataframe for easier use

e = df['Median East'].values
n = df['Median North'].values
u = df['Median Up'].values

xv = df['AlongCross X'].values
yv = df['AlongCross Y'].values
zv = df['AlongCross Z'].values

 #Get just the name of the truth file without the path and .txt extension. Then take all words after the truth name. These are the adjustments.
 #e.g. for pospac_singlebase_alongCross_medianENU the adjustments array would be ['alongCross','medianENU']
adjustments = plotDir.split('_')[::-1][:2][::-1]


#Ask user for which truth file to use and create a dataframe from the truth file
truthDir = workingDir + '/Truth_Files'
truthFile = pa.getDirFromUser(truthDir)
filename = 'Adjusted_Truth_Files/' + 'copiedTruth.txt'
shutil.copy(truthDir + '/' + truthFile,filename) #Copy the truth file to a new directory that we want to save the adjusted file in
truth_df = pd.read_csv(filename,delimiter=' ',header=None)


#Get LLH and convert them from a dataframe to an array
lat = truth_df.iloc[:,1].values
lon = truth_df.iloc[:,2].values
hgt = truth_df.iloc[:,3].values


#For each adjustment, adjust the LLH according to the adjustment parameters
for adj in adjustments:
    if 'median' in adj:
        print("Doing Median Adjustment") #Print the adjustment so the user can verify the correct adjustments are being done in the correct order
        enuCorr = np.zeros((len(truth_df),3))
        enuCorr[:,0] = enuCorr[:,0] + e
        enuCorr[:,1] = enuCorr[:,1] + n
        enuCorr[:,2] = enuCorr[:,2] + u

        llh = pa.llh_enu_adjust(lat,lon,hgt,enuCorr)
        lat = llh[:,0]
        lon = llh[:,1]
        hgt = llh[:,2]

    if adj == 'alongCross':
        print("Doing Along Cross Adjustment") #Print the adjustment so the user can verify the correct adjustments are being done in the correct order
        enuCorr = np.zeros((len(truth_df),3))
        llh = pa.alongCrossAdjust(xv,yv,zv,truth_df.to_numpy())
        lat = llh[:,0]
        lon = llh[:,1]
        hgt = llh[:,2]

#Replace the truth dataframe LLH with the adjusted LLH
truth_df.iloc[:,1] = lat
truth_df.iloc[:,2] = lon
truth_df.iloc[:,3] = hgt

#Write the new adjusted LLH to the copied truth file and save it with the name of the directory (i.e. 2022-02-22-SantaNella-van-test) 
#Save it in Adjusted_Truth_Files main directory
writefile = 'Adjusted_Truth_Files/' + workingDir + '.txt'
truth_df.to_csv(writefile, sep=' ',header=None,index=False)