<gh_stars>0
from pathlib import Path
import numpy as np
import h5py
from datetime import datetime
from input import parameters as para
from lib import cfl as cfl
from lib import write_field as wf
from lib import avg_profile as avgp
from lib import grid as gd
from lib import timehistory as tmsr
from lib import initial as initC
from src import rbc, hydro
from lib import details as PrintSimulationDetails

PrintSimulationDetails

def main():

    if para.ProbType == "RBC": U, V, W, P, T, time, fwTime, rwTime = initC.initFields()
    if para.ProbType == "Hydro": U, V, W, P, time, fwTime, rwTime = initC.initFields()
        
    if para.ProbType == "RBC": Re, Nu, Sw = tmsr.ReNuSw(U, V, W, T)
    if para.ProbType == "Hydro": Energy = tmsr.energy(U, V, W)

    maxDiv = tmsr.getDiv(U, V, W)

    Path("output").mkdir(parents=True, exist_ok=True)
    f = open('output/TimeSeries.dat', 'a')
    if para.ProbType == "RBC":
        f.write("# Ra = %.3e \t Pr = %.3f \t Ta = %.3e \t Nx = %i \t Ny = %i \t Nz = %i \n" %(para.Ra, para.Pr, para.Ta, gd.Nx, gd.Ny, gd.Nz)) 
        f.write('# time, Re, Nu, Sw, Divergence, dt \n')
        f.write("%.5e\t%.5e\t%.5e\t%.5e\t%.5e\t%.5e \n" %(time, Re, Nu, Sw, maxDiv, para.dt))

    if para.ProbType == "Hydro":
        f.write("# %f \t %f \t %i \t %i \t %i \n" %(para.nu, para.Omega, gd.Nx, gd.Ny, gd.Nz)) 
        f.write('# time, Energy, Divergence, dt \n')
        f.write("%.5e\t%.5e\t%.5e\t%.5e \n" %(time, Energy, maxDiv, para.dt))

    f.close()
    
    if para.ProbType == "RBC":        
        print('# time, Reynolds, Nusselt, Divergence')
        print("%.5f    %.5f    %.5f    %.3e" %(time, Re, Nu, maxDiv))

    if para.ProbType == "Hydro":        
        print('# time \t\t Energy \t Divergence')
        print("%.5f\t\t%.5f\t\t%.3e" %(time, Energy, maxDiv))

    iCnt, Sw, jrCnt, dtnew, dtchange = 0, 0.0, 0, 1.0, 0

    if para.ProbType == "RBC": wf.writeSoln_RBC(U, V, W, P, T, time, dtchange)
    if para.ProbType == "Hydro": wf.writeSoln_Hydro(U, V, W, P, time, dtchange)

    if abs(time - para.tMax) < 0.5*para.dt or time >= para.tMax:
        print("# Error! Final time is less than the start time")
        quit()

    ts1 = datetime.now()
    while True:

        if para.restart or iCnt > 1: dtnew = cfl.cflcondition(U, V, W)
        
        dt = min(para.dt, dtnew)

        if jrCnt==0 and dtnew<para.dt: dtchange, jrCnt = 1, 1
            
        if para.ProbType == "RBC": U, V, W, P, T = rbc.solveRBC3D(U, V, W, P, T, dt)
        if para.ProbType == "Hydro": U, V, W, P = hydro.solveHydro3D(U, V, W, P, dt)

        time, iCnt = time + dt, iCnt + 1
        
        if abs(rwTime - time) < 0.5*dt:
            rwTime = rwTime + para.rwInt
            if iCnt > 1:
                if para.ProbType == "RBC": wf.writeRestart_RBC(U, V, W, P, T, time, dtchange)                   
                if para.ProbType == "Hydro": wf.writeRestart_Hydro(U, V, W, P, time, dtchange)
                    
        if abs(fwTime - time) < 0.5*dt:
            fwTime = fwTime + para.fwInt
            if para.ProbType == "RBC": wf.writeSoln_RBC(U, V, W, P, T, time, dtchange)
            if para.ProbType == "Hydro": wf.writeSoln_Hydro(U, V, W, P, time, dtchange)
                
        if iCnt % para.opInt == 0:
            if para.ProbType == "RBC": Re, Nu, Sw = tmsr.ReNuSw(U, V, W, T)
            if para.ProbType == "Hydro": Energy = tmsr.energy(U, V, W)
     
            maxDiv = tmsr.getDiv(U, V, W)

            f = open('output/TimeSeries.dat', 'a')
            if para.ProbType == "RBC": f.write("%.6e\t%.6e\t%.6e\t%.6e\t%.6e\t%.6e \n" %(time, Re, Nu, Sw, maxDiv, dt))
            if para.ProbType == "Hydro": f.write("%.6e\t%.6e\t%.6e\t%.6e \n" %(time, Energy, maxDiv, dt))
            f.close()

            if para.ProbType == "RBC": print("%.5f    %.5f    %.5f    %.3e" %(time, Re, Nu, maxDiv))
            if para.ProbType == "Hydro": print("%.5f\t\t%.5f\t\t%.3e" %(time, Energy, maxDiv))

            if abs(maxDiv)>1e4: 
                print("Error! Simulation blowing up. Aborting!")
                quit() 

        if para.AvgProf and para.ProbType == "RBC": avgp.avg_profile(U, V, W, T, iCnt)
        
        if abs(time - para.tMax) < 0.2*dt or time >= para.tMax:
            print("\nSimulation completed!")
            break   

    ts2 = datetime.now()

    print("Simulation time",ts2-ts1)

main()








