#!/usr/bin/env python

import numpy as np
from netCDF4 import Dataset
from LBPM_Preprocessing_utils import *
import matplotlib.pyplot as plt

# Load the *nc image
# NOTE: for the variable name, you might want to check the *.nc file with ncdump -h
image_red = read_NetCDF_file_py('Refined_subdomain_1_raw.nc','segmented')
# This *nc file is not a raw CT image data, which has been processed by me. 
# It already has the format as :
# Whereas LBPM-WIA has the convention as:
# 0 -> solid phase
# 1 -> fluid phase
(lz_red,ly_red,lx_red) = image_red.shape

# Preprocess the image
# ---------------- By default: ---------------------------------------------
# 1. The flow direction is along positive z-axis
#    z_min: inlet;  z_max: outlet
# 2. The boundary condition, e.g. pressure boundary condition is applied
#    on x-y planes
# --------------------------------------------------------------------------

# Steps of preprocessing:
# 1. wrap boundaries along z-axis with solid nodes
# 1.1 (optional) save the raw image for morphological drainge & opeing
# 2. change the original z_min and z_max layers to reservoir layers
# 3. set z_min reservoir layer to non-wetting phase reservoir
# 4. set z_max reservoir layer to wetting phase reservoir
# 5. add porous plate layer before the wetting phase reservoir
# 6. saturate the pore space with wetting phase

# 1. wrap the boundaries
# TODO: make a variable for wall_width
image_red[:,0,:]        = 0
image_red[:,ly_red-1,:] = 0
image_red[:,:,0]        = 0
image_red[:,:,lx_red-1] = 0

# 1.1 (optional) Save the image for morphological analysis
#print "**Info: Save data for morphological drainage and imbibition."
#image_red.tofile('subdomain_1_reduced_morph_2x.bin') # dont use this

# 2. Saturate the medium with wetting phase
image_red[image_red==1] = 2
# 3. Set up reservoir layers
# Input parameters for the reservoir layers
NWR_len = 1 # The thickness of the non-wetting phaer reservoir
WR_len  = 1 # The thickness of the wetting phase reservoir
image_red[0:NWR_len,:]  = 1
image_red[-WR_len:,:]   = 2
# temp test
#plt.pcolormesh(image_red[:,1,:],cmap='hot')
#plt.axis('equal')
#plt.colorbar();
#plt.show()
# 4. setup up a porous plate saturated with **wetting phase**
# Input parameters for the porous plate
pp_size = 3 #NOTE: Although it is made as an variable, pp_size of 3 is recommended
pp_len  = 3 # the thickness of the porous plate
pp = 2*np.ones((pp_size,pp_size),dtype=np.int8)
#pp = np.lib.pad(pp,((1,1),(1,1)),'constant',constant_values=0)
pp = np.lib.pad(pp,((1,1),(1,1)),'constant',constant_values=0)
(ly_pp,lx_pp) = pp.shape
pp_pad = np.tile(pp,(ly_red/ly_pp,lx_red/lx_pp))
l_residual=ly_red%ly_pp
pp_pad = np.lib.pad(pp_pad,((l_residual%2,0),(l_residual%2,0)),'constant',constant_values=0)
pp_pad = np.lib.pad(pp_pad,((l_residual/2,l_residual/2),(l_residual/2,l_residual/2)),'constant',constant_values=0)
# 5. Create the simulation domain
image_red[-(WR_len+pp_len):-WR_len,:,:]=pp_pad
# temp test
#plt.pcolormesh(image_red[-5,:],cmap='hot')
#plt.axis('equal')
#plt.grid(True)
#plt.colorbar()
#plt.show()
# write to *.nc file for futher check
#print "**Info: Save data for LBPM simulations."
#write_NetCDF_file_py('subdomain_1_reduced_lbpm_2x.nc','segmented',\
#                     image_red,image_red.dtype.char,lx_red,ly_red,lz_red)
#image_red.tofile('subdomain_1_reduced_lbpm_2x.bin')