# -*- coding: utf-8 -*-
"""
Created on Sun Feb 27 19:41:04 2022

@author: Craig Thompson

This file contains the functions used to run DA within openfoam
"""

from pathlib import Path
import shutil
from PyFoam.Execution.BasicRunner import BasicRunner
import numpy as np
from matplotlib import pyplot as plt

def RunSim(Forcing, Pdir):
    '''
    Runs an openfoam simulation based on run files from Pdir, allowing the
    addition of momentum via the Forcing array. The function then returns
    the velocity output of the simulation.

    Parameters
    ----------
    Forcing : array
        Vector of size 2xn indicating the additional momentum in x and y.
    Pdir : String
        Location of foam run files.

    Returns
    -------
    list
        2xn vector of velocities.

    '''
    
    #Generate forcing matrix U
    force_vectU = str(list(Forcing[0]))
    force_vectU = force_vectU.replace("[", '{')
    force_vectU = force_vectU.replace("]", '}')
    force_vectU = force_vectU.replace(" ", '')
    force_vectU = "double sfU[] = " + force_vectU + ' ;'
    
    #Save forcing matrix U
    fvf = open('SFU.txt','w')
    fvf.write(str(force_vectU))
    fvf.close()
    
    #Generate forcing matrix V
    force_vectV = str(list(Forcing[1]))
    force_vectV = force_vectV.replace("[", '{')
    force_vectV = force_vectV.replace("]", '}')
    force_vectV = force_vectV.replace(" ", '')
    force_vectV = "double sfV[] = " + force_vectV + ' ;'
    
    #Save forcing matrix V
    fvf = open('SFV.txt','w')
    fvf.write(str(force_vectV))
    fvf.close()
    
    #Run the sim
    run = BasicRunner(argv=["simpleFoam","-case",Pdir])
    run.start()
    
    #Finding final itteration number
    OFit = FinalIt(Pdir)
    
    #Reading the result from the simulation
    Sim_U = []
    Sim_V = []
    Sim = open(Pdir+"/"+str(OFit)+"/U")
    Simres = Sim.readlines()
    Simres = Simres[22:22+21600]		# Replace the end point of the splice with the number of OF cells
    for i in list(range(0,len(Simres))):
        Simres[i]=Simres[i].split()
        Simres[i][0]=Simres[i][0].replace('(','')
        Simres[i][2]=Simres[i][2].replace(')','')
        for ii in [0,1,2]:
            Simres[i][ii] = float(Simres[i][ii])
            if ii==0:
                Sim_U.append(Simres[i][ii])
            elif ii==1:
                Sim_V.append(Simres[i][ii])
    Sim.close()

    #Return velocity result
    return [Sim_U, Sim_V, OFit]

def FinalIt(Pdir):
    '''
    Determines the final itteration number of an openfoam simulation run in 
    Pdir.

    Parameters
    ----------
    Pdir : String
        Location of foam run files.

    Returns
    -------
    OFit : Integer
        Final itteration number.

    '''
    #Use the post proccesing file to determinet the final itteration
    OFitf = open(Pdir+"/postProcessing/yPlus1/1/yPlus.dat")
    lines = OFitf.readlines()
    lines = lines[2]
    OFit = (int(lines.split()[0]))
    OFitf.close()
    print(str(OFit))
    return OFit

def Delete_Folder(pth) :
    '''
    Delete the contents of a folder and the folder itself.

    Parameters
    ----------
    pth : pathlib.Path()
        Directory to be deleted.

    Returns
    -------
    None.

    '''
    
    #Itterate through the folder removing each file
    for sub in pth.iterdir() :
        if sub.is_dir() :
            Delete_Folder(sub)
        else :
            sub.unlink()
    pth.rmdir()

def RemoveFiles(Pdir):
    '''
    Remove the contents of an openfoam run.

    Parameters
    ----------
    Pdir : String
        Location of foam run files.

    Returns
    -------
    None.

    '''
    #Find case file number
    OFit = FinalIt(Pdir)
    #Remove old case folder
    CPath = Path(Pdir+"/1")
    Delete_Folder(CPath)
    #Make previous case the new restart file
    oldfilepath = Path(Pdir+"/"+str(OFit))
    oldfilepath.rename(Pdir+"/1")
    #Remove Dynamic code
    DCPath = Path(Pdir+"/dynamicCode")
    Delete_Folder(DCPath)
    #Remove post processing
    PPPath = Path(Pdir+"/postProcessing")
    Delete_Folder(PPPath)
    #Remove openfoam forcing
    SFUPath = Path("SFU.txt")
    SFVPath = Path("SFV.txt")
    SFUPath.unlink()
    SFVPath.unlink()
    return

def SaveFiles(Pdir,Sdir,name):
    '''
    Save or remove the relevent files from an openfoam run.

    Parameters
    ----------
    Pdir : String
        Location of foam run files.
    Sdir : String
        Location of save directory.
    name : integer
        Run name, added as a suffex.

    Returns
    -------
    None.

    '''
    #Find case file number
    OFit = FinalIt(Pdir)
    #Save old case folder
    CPath = Path(Pdir+"/1")
    CPath.rename(Sdir+"/C"+str(name))
    #Remove openfoam forcing
    SFUPath = Path("SFU.txt")
    SFVPath = Path("SFV.txt")
    SFUPath.unlink()
    SFVPath.unlink()
    #Coppy current forcing
    fExpPath = Path("fExp.txt")
    fPath = Path("f.txt")
    shutil.copy2(fExpPath,Sdir+"/fExp"+str(name)+".txt")
    shutil.copy2(fPath,Sdir+"/f"+str(name)+".txt")
    #Make previous case the new restart file
    oldfilepath = Path(Pdir+"/"+str(OFit))
    oldfilepath.rename(Pdir+"/1")
    #Remove Dynamic code
    DCPath = Path(Pdir+"/dynamicCode")
    Delete_Folder(DCPath)
    #Remove post processing
    PPPath = Path(Pdir+"/postProcessing")
    Delete_Folder(PPPath)
    return

def TxtRead(Filename, dimensions):
    '''
    Read experimental text file.

    Parameters
    ----------
    Filename : String
        Filename to read.
        
    dimensions : integer
        1,2,3.

    Returns
    -------
    list
        Content of file in x,y tuple.

    '''
    E1 = []
    E2 = []
    E3 = []
    Expf = open(Filename, 'r')
    for i in Expf:
        current = i.split(",")
        if dimensions==1:
            E1.append(float(current[0]))
        elif dimensions==2:
            E1.append(float(current[0]))
            E2.append(float(current[1]))
        elif dimensions==3:
            E1.append(float(current[0]))
            E2.append(float(current[1]))
            E3.append(float(current[2]))
        else:
            print('wrong dimension input')
    Expf.close()
    if dimensions == 1:
        return E1
    elif dimensions == 2:
        return [E1,E2]
    else:
        return [E1,E2,E3]
    
def IntRead(Filename, dimensions):
    '''
    Read text file and produce integer resopnce.

    Parameters
    ----------
    Filename : String
        Filename to read.
        
    dimensions : integer
        1,2,3.

    Returns
    -------
    list
        Content of file in x,y,z tuple.

    '''
    E1 = []
    E2 = []
    E3 = []
    Expf = open(Filename, 'r')
    for i in Expf:
        current = i.split(",")
        if dimensions==1:
            E1.append(int(current[0]))
        elif dimensions==2:
            E1.append(int(current[0]))
            E2.append(int(current[1]))
        elif dimensions==3:
            E1.append(int(current[0]))
            E2.append(int(current[1]))
            E3.append(int(current[2]))
        else:
            print('wrong dimension input')
    Expf.close()
    if dimensions == 1:
        return E1
    elif dimensions == 2:
        return [E1,E2]
    else:
        return [E1,E2,E3]
    
def PlotData(Sdir,name,Exp_x,Exp_y,L1,fExp):
    
    #Plot L1
    figL1 = plt.figure()
    plt.tricontourf(Exp_x,Exp_y,L1,25,cmap='binary')
    plt.colorbar()
    filename = Sdir + '/L1_' + str(name) 
    plt.title('L1 ' + str(name))
    plt.savefig(filename + '.png')
    plt.savefig(filename + '.eps')
    plt.close(figL1)
    
    #Plot Fx
    #figFx = plt.figure()
    #plt.tricontourf(Exp_x,Exp_y,fExp[0,:],np.linspace(-max(abs(fExp[0,:])),max(abs(fExp[0,:])),100),cmap='bwr')
    #plt.colorbar()
    #filename = Sdir + '/Fx_' + str(name) 
    #plt.title('Fx ' + str(name))
    #plt.savefig(filename + '.png')
    #plt.savefig(filename + '.eps')
    #plt.close(figFx)
    
    #Plot Fy
    #figFy = plt.figure()
    #plt.tricontourf(Exp_x,Exp_y,fExp[1,:],np.linspace(-max(abs(fExp[1,:])),max(abs(fExp[1,:])),100),cmap='bwr')
    #plt.colorbar()
    #filename = Sdir + '/Fy_' + str(name) 
    #plt.title('Fy ' + str(name))
    #plt.savefig(filename + '.png')
    #plt.savefig(filename + '.eps')
    #plt.close(figFy)

def ReadForcing(Filename):
    '''
    Read the OLD forcing file for restarts.

    Parameters
    ----------
    Filename : String
        Filename to read.

    Returns
    -------
    list
        Content of forcing file.

    '''
    ft = open(Filename,'r')
    f = ft.read()
    f = f.replace('((','').replace('))', '')
    f = f.split('), (')
    f = np.genfromtxt(f, delimiter=',')
    ft.close()
    return f


def ReadSFUSFV():
    '''
    Read forcing files SFU and SFV for use in restarts.

    Parameters
    ----------
    Filename : String
        Directory of SFU and SFV.

    Returns
    -------
    list
        Forcing used for restart.

    '''
    #SFU
    SFUft = open('SFU.txt','r')
    SFUf = SFUft.read()
    SFUf = SFUf.replace('double sfU[] = {','').replace('} ;','')
    SFU = np.fromstring(SFUf,sep=',')
    #SFV
    SFVft = open('SFV.txt','r')
    SFVf = SFVft.read()
    SFVf = SFVf.replace('double sfV[] = {','').replace('} ;','')
    SFV = np.fromstring(SFVf,sep=',')
    #Join to form f
    f = np.array([SFU.T,SFV.T])
    return f




