from samwich.waketrackers import track
from samwich.dataloaders import PlanarData
import os
import numpy as np
import pyFAST.input_output.vtk_file as vtk_file

# Temporary inputs - Used only to make the specification of file names easier
# Main inputs
numPlans = 1 # Number of planes to be processed. Assumed to be equally spaced
dt = 3.6       # Time step between vtk files
folderPath = r'G:\FAST_Farm_Template\farm_9T_Mann_baseline_avg_V_curl_v2\Cond00_v10.0_PL0.2_TI6\Case00_wdir0.0_fixed_kcurl6\Seed_0' # Path to the folder with the FAST.Farm simulation
rootNameVTK = 'FFarm_mod.Low.DisYZ1.'     # Root name of the test case
outFolder = os.path.join(folderPath, 'wakescenter') # Output folder. Output files will have the same name as the vtks, but with .csv extensions.
rotorDiam = 240 # Rotor diameter
rootPath = os.path.join(folderPath, 'vtk_ff') # Subfolder with vtk's generated by FAST.Farm



# Create the output folder if it does not exist
if not os.path.exists(outFolder):
    os.makedirs(outFolder)





# List all vtk files

planes_xPos = np.zeros(1)+2519.22

listNames = []
listXPos  = []
for i in range(1):
    # Some versions of FAST.Farm write VTk names as 01 and others 1
    nameVTK = f'{rootNameVTK}{i+1000:05d}'
    nameVTK_wo0 = f'{rootNameVTK}{i+1000}'

    # List all files with nameVTK in the name
    thisListOfNames = [f for f in os.listdir(rootPath) if os.path.isfile(os.path.join(rootPath, f)) and nameVTK in f or nameVTK_wo0 in f]
    listNames.extend(thisListOfNames)
    listXPos.extend([planes_xPos[i]]*len(thisListOfNames))    

list_vtk = []
for vtkName in listNames[0:1]:
    filevtk=vtk_file.VTKFile(rootPath + '/' + vtkName)
    filevtk.yp_grid = filevtk.yp_grid[3:14]
    filevtk.zp_grid = filevtk.zp_grid[0:9]
    filevtk.point_data_grid['Velocity']=filevtk.point_data_grid['Velocity'][:,3:14,0:9,:]
    list_vtk.append(filevtk)


# Check how many files
print(len(list_vtk))




# Sort listNames alphabetically. That is enough to make sure that things are sorted by time
listNames = sorted(listNames)

for i in range(numPlans):
    # Get indices of files for this plan
    ind4plan = [j for j, nm in enumerate(listNames) if f'{rootNameVTK}{i+1000:05d}' in nm or f'{rootNameVTK}{i+1000+1}' in nm]
    # Write to a csv file
    outFlPath = os.path.join(outFolder, f'{rootNameVTK}{i+1000:05d}.csv')    

    with open(outFlPath, 'w') as f:
        f.write('t,x,y,z')        

        dict4wake = {'y':[], 'z':[], 'u':[], 'v':[], 'w':[]}   # Used by SAMWICH
        for iTime, iP in enumerate(ind4plan):            
            time = dt*int(listNames[iP].split('.')[-2]) # TIme
            vtk = list_vtk[iP]            

            # 2D arrays with shape len(vtk.yp_grid): rows, len(vtk.zp_grid): columns
            dict4wake['u'] = vtk.point_data_grid['Velocity'][0, :, :, 0]
            dict4wake['v'] = vtk.point_data_grid['Velocity'][0, :, :, 1]
            dict4wake['w'] = vtk.point_data_grid['Velocity'][0, :, :, 2]
            xc = listXPos[i] # X position of the wake center (constant and equal to plane position)

            # Need 2D arrays with same shape as u, v, w
            dict4wake['z'], dict4wake['y'] = np.meshgrid(vtk.zp_grid, vtk.yp_grid)
            inp=dict4wake.copy()
            wakedata = PlanarData(inp)
            wake = track(wakedata.sliceI(), method='ConstantArea', verbose=True)
            
            wake.clear_plot() 
            wake.plot_contour(vmin=1,vmax=13)
            
            # Compute wake centers
            refArea = np.pi*240**2/4
            
            wake.remove_shear(alpha=0.2, zref=135, Uref=10)
            wake.clear_plot()  # force replot for python notebook
            wake.plot_contour(vmin=-5,vmax=5)             
            yc, zc = wake.find_centers(refArea,min_contour_points=5,verbosity=2) # y and z coordinates of the wake center

            
            # Write to file
            f.write(f'{time:.3f},{xc:.3f},{yc[0]:.3f},{zc[0]:.3f}')
            f.flush()
            wake.clear_plot()  
            wake.plot_contour(vmin=-5,vmax=5,outline=True)