import os
import torch
import json

# local file imports
import domain_setup
import PINN_setup_and_training
import post_processing


parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if os.name == "nt":
    results_path = parent_dir + "\PINN_results"
else:
    results_path = parent_dir + "/PINN_results"

version_name = input("What version are you working on? (this will create a folder in PINN_results) ")
if not os.path.exists(results_path+f"/{version_name}"):
    os.makedirs(results_path+f"/{version_name}")


### INPUT PARAMETERS ###
input_variables = {
    "u0"    : 1,   # inlet velocity
    "rho0"  : 1,
    "nu0"   : 0.01,
    "N_col" : 4400,
    "N_BC"  : 500,
    "N_circ": 400,
    "r"     : 0.5,
    "N_data": 1000,
    "adam_epoch" : 700,
    "lbfgs_epochs" : 10, 
    "epoch_split" : 200 
## ADD in number of epochs for adam and LFBGS training
}
# This is where the Cyl 100 data is
data_path = "/home/james/Downloads"

### DOMAIN SETUP ###
if (input("Create new data points? y/n: ")) == "y":
    X_col, BC_X, BC_Y= domain_setup.Uniform_with_Cylinder(input_variables["u0"], input_variables["N_col"], input_variables["N_BC"], input_variables["N_circ"], input_variables["r"])
    X_D, U_D, PXY_D, P_D = domain_setup.Data(input_variables["N_data"], data_path)
    Tensors = {"X_col" : X_col,
               "BC_X"  : BC_X,
               "BC_Y"  : BC_Y, 
               "X_D"   : X_D,
            "U_D"   : U_D,
            "PXY_D" : PXY_D,
            "P_D"   : P_D}

if (input("Plot domain? y/n: ")== "y"):
    domain_setup.Domain_Plot(X_col, BC_X, X_D, results_path, version_name)




### SETUP AND RUN MODEL ###

Model = PINN_setup_and_training.PINN(Tensors, input_variables)

Train_Load = input("train or load a model? T/L: ")
if (Train_Load)== "T":
    L_hist, L_data_hist, L_UV_data_hist, L_P_data_hist, L_PDE_hist, L_BC_hist = Model.train(input_variables["adam_epoch"],input_variables["lbfgs_epochs"],input_variables["epoch_split"])
elif Train_Load == "L":
    print("What version would you like to load? enter from the version list below.")
    print(os.listdir(results_path))
    To_Load = "/"+input()+"/Model_Parameters.pt"
    Model.model.load_state_dict(torch.load((results_path+To_Load)))
    Model.eval()



### POST PROCESSING ###

# pass in model + data + input params
post_processing.Plotting(Model, Tensors, results_path, version_name)

### SAVING THE MODEL ###
if input("would you like to save the models domain, input variables and weights? y/n: ") == "y":
    torch.save(Model.model.state_dict(),(results_path+f"/{version_name}/Model_Parameters.pt"))
    torch.save(Tensors, (results_path+f"/{version_name}/Domain_Tensors.pt"))
    with open((results_path+f"/{version_name}/Input_Variables.json"),'w') as f:
        json.dump(input_variables, f, indent = 4)