import numpy as np
from input import parameters as para
if para.Target=="CPU":  import numpy as ncp
from lib import boundary as BC
import h5py as hp
from lib.PoissonSolver import multigrid3D
from lib import grid as gd
from lib.glob import * 
if para.Target=="GPU": 
    import cupy as ncp
    ncp.cuda.Device(para.device).use()



def initFields():
    
    pData, rData, iTemp = multigrid3D.initMG()

    if para.restart:
        if para.ProbType=="RBC": U, V, W, P, T = deadstate()
        if para.ProbType=="Hydro": U, V, W, P = deadstate()

        if para.ProbType=="RBC": U[x0, y0, z0], V[x0, y0, z0], W[x0, y0, z0], P[x0, y0, z0], T[x0, y0, z0], time = restartrun()
        if para.ProbType=="Hydro": U[x0, y0, z0], V[x0, y0, z0], W[x0, y0, z0], P[x0, y0, z0], time = restartrun()
        
    else:
        time = 0
        if para.ICtype==0 and para.ProbType=="RBC": U, V, W, P, T = deadstate()
        if para.ICtype==1 and para.ProbType=="RBC": U, V, W, P, T = randomIC()
        if para.ICtype==0 and para.ProbType=="Hydro": U, V, W, P = deadstate()
        if para.ICtype==1 and para.ProbType=="Hydro": U, V, W, P = randomIC()

    fwTime = para.fwInt
    rwTime = para.rwInt

    if para.restart:
        fwTime = time + para.fwInt
        rwTime = time + para.rwInt

    if para.ProbType=="RBC": 
        BC.imposeBCs(U), BC.imposeBCs(V), BC.imposeBCs(W), BC.imposePBCs(P), BC.imposeTBCs(T)
        return U, V, W, P, T, time, fwTime, rwTime

    if para.ProbType=="Hydro": 
        BC.imposeBCs(U), BC.imposeBCs(V), BC.imposeBCs(W), BC.imposePBCs(P)
        return U, V, W, P, time, fwTime, rwTime




def restartrun():
    filename = "output/Restart.h5"
    def hdf5_reader(filename,dataset):
        file_V1_read = hp.File(filename, "r")
        dataset_V1_read = file_V1_read["/"+dataset]
        V1=dataset_V1_read[:,:,:]
        return V1
    
    Vx = hdf5_reader(filename, "Vx")
    Vy = hdf5_reader(filename, "Vy")
    Vz = hdf5_reader(filename, "Vz")
    Press = hdf5_reader(filename, "P")
    if para.ProbType == "RBC": Temp = hdf5_reader(filename, "T")

    if para.ProbType == "RBC":
        if para.Target== "GPU": U, V, W, P, T = ncp.asarray(Vx), ncp.asarray(Vy), ncp.asarray(Vz), ncp.asarray(Press), ncp.asarray(Temp)
        if para.Target== "CPU": U, V, W, P, T = Vx, Vy, Vz, Press, Temp

    if para.ProbType == "Hydro":
        if para.Target== "GPU": U, V, W, P = ncp.asarray(Vx), ncp.asarray(Vy), ncp.asarray(Vz), ncp.asarray(Press)
        if para.Target== "CPU": U, V, W, P = Vx, Vy, Vz, Press
    
    f = hp.File(filename, 'r')
    time = f.get('Time')
    time = np.array(time)
    f.close()
    
    if para.Target=="GPU": time = ncp.array(time)

    if para.ProbType=="RBC": return U, V, W, P, T, time 
    if para.ProbType=="Hydro": return U, V, W, P, time 


def deadstate():
    P = ncp.zeros(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)
    U = ncp.zeros(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)   #np.random.rand(gd.Nx+2, gd.Ny+2, gd.Nz+2) 
    V = ncp.zeros(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype) 
    W = ncp.zeros(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)
    if para.ProbType=="RBC": T = ncp.zeros(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)   #T[1:-1, 1:-1, 1:-1] = (1.0 - 0.5/(gd.Nz+1)) - z[1:-1]

    if para.ProbType=="RBC": return U, V, W, P, T
    if para.ProbType=="Hydro": return U, V, W, P


def randomIC():
    P = ncp.zeros(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)
    U = 0.987*ncp.ones(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)   
    V = 0.879*ncp.ones(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)
    W = 0.789*ncp.ones(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)
    if para.ProbType=="RBC": T = ncp.zeros(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)   #T[1:-1, 1:-1, 1:-1] = (1.0 - 0.5/(gd.Nz+1)) - z[1:-1]


    if para.ProbType=="RBC": return U, V, W, P, T
    if para.ProbType=="Hydro": return U, V, W, P

import numpy as np
from input import parameters as para
if para.Target=="CPU":  import numpy as ncp
from lib import boundary as BC
import h5py as hp
from lib.PoissonSolver import multigrid3D
from lib import grid as gd
from lib.glob import * 
if para.Target=="GPU": 
    import cupy as ncp
    ncp.cuda.Device(para.device).use()



def initFields():
    
    pData, rData, iTemp = multigrid3D.initMG()

    if para.restart:
        if para.ProbType=="RBC": U, V, W, P, T = deadstate()
        if para.ProbType=="Hydro": U, V, W, P = deadstate()

        if para.ProbType=="RBC": U[x0, y0, z0], V[x0, y0, z0], W[x0, y0, z0], P[x0, y0, z0], T[x0, y0, z0], time = restartrun()
        if para.ProbType=="Hydro": U[x0, y0, z0], V[x0, y0, z0], W[x0, y0, z0], P[x0, y0, z0], time = restartrun()
        
    else:
        time = 0
        if para.ICtype==0 and para.ProbType=="RBC": U, V, W, P, T = deadstate()
        if para.ICtype==1 and para.ProbType=="RBC": U, V, W, P, T = randomIC()
        if para.ICtype==0 and para.ProbType=="Hydro": U, V, W, P = deadstate()
        if para.ICtype==1 and para.ProbType=="Hydro": U, V, W, P = randomIC()

    fwTime = para.fwInt
    rwTime = para.rwInt

    if para.restart:
        fwTime = time + para.fwInt
        rwTime = time + para.rwInt

    if para.ProbType=="RBC": 
        BC.imposeBCs(U), BC.imposeBCs(V), BC.imposeBCs(W), BC.imposePBCs(P), BC.imposeTBCs(T)
        return U, V, W, P, T, time, fwTime, rwTime

    if para.ProbType=="Hydro": 
        BC.imposeBCs(U), BC.imposeBCs(V), BC.imposeBCs(W), BC.imposePBCs(P)
        return U, V, W, P, time, fwTime, rwTime




def restartrun():
    filename = "output/Restart.h5"
    def hdf5_reader(filename,dataset):
        file_V1_read = hp.File(filename, "r")
        dataset_V1_read = file_V1_read["/"+dataset]
        V1=dataset_V1_read[:,:,:]
        return V1
    
    Vx = hdf5_reader(filename, "Vx")
    Vy = hdf5_reader(filename, "Vy")
    Vz = hdf5_reader(filename, "Vz")
    Press = hdf5_reader(filename, "P")
    if para.ProbType == "RBC": Temp = hdf5_reader(filename, "T")

    if para.ProbType == "RBC":
        if para.Target== "GPU": U, V, W, P, T = ncp.asarray(Vx), ncp.asarray(Vy), ncp.asarray(Vz), ncp.asarray(Press), ncp.asarray(Temp)
        if para.Target== "CPU": U, V, W, P, T = Vx, Vy, Vz, Press, Temp

    if para.ProbType == "Hydro":
        if para.Target== "GPU": U, V, W, P = ncp.asarray(Vx), ncp.asarray(Vy), ncp.asarray(Vz), ncp.asarray(Press)
        if para.Target== "CPU": U, V, W, P = Vx, Vy, Vz, Press
    
    f = hp.File(filename, 'r')
    time = f.get('Time')
    time = np.array(time)
    f.close()
    
    if para.Target=="GPU": time = ncp.array(time)

    if para.ProbType=="RBC": return U, V, W, P, T, time 
    if para.ProbType=="Hydro": return U, V, W, P, time 


def deadstate():
    P = ncp.zeros(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)
    U = ncp.zeros(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)   #np.random.rand(gd.Nx+2, gd.Ny+2, gd.Nz+2) 
    V = ncp.zeros(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype) 
    W = ncp.zeros(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)
    if para.ProbType=="RBC": T = ncp.zeros(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)   #T[1:-1, 1:-1, 1:-1] = (1.0 - 0.5/(gd.Nz+1)) - z[1:-1]

    if para.ProbType=="RBC": return U, V, W, P, T
    if para.ProbType=="Hydro": return U, V, W, P


def randomIC():
    P = ncp.zeros(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)
    U = 0.987*ncp.ones(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)   
    V = 0.879*ncp.ones(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)
    W = 0.789*ncp.ones(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)
    if para.ProbType=="RBC": T = ncp.zeros(([gd.Nx+2, gd.Ny+2, gd.Nz+2]), dtype=para.dtype)   #T[1:-1, 1:-1, 1:-1] = (1.0 - 0.5/(gd.Nz+1)) - z[1:-1]


    if para.ProbType=="RBC": return U, V, W, P, T
    if para.ProbType=="Hydro": return U, V, W, P

#np.random.rand(gd.Nx+2, gd.Ny+2, gd.Nz+2) #T[1:-1, 1:-1, 1:-1] = (1.0 - 0.5/(gd.Nz+1)) - z[1:-1] #T[1:-1, 1:-1, 1:-1] = (1.0 - 0.5/(gd.Nz+1)) - z[1:-1]