"""
Rayleigh-Benard instability simulated by
Navier-Stokes solver D2Q9 coupled to thermic solver D2Q5
"""
import numpy as np
import sympy as sp
import pylbm
import sys

def printProgress (iteration, total, prefix = '', suffix = '', decimals = 1, barLength = 100):
    """
    Call in a loop to create terminal progress bar
    @params:
        iteration   - Required  : current iteration (Int)
        total       - Required  : total iterations (Int)
        prefix      - Optional  : prefix string (Str)
        suffix      - Optional  : suffix string (Str)
        decimals    - Optional  : positive number of decimals in percent complete (Int)
        barLength   - Optional  : character length of bar (Int)
    """
    formatStr       = "{0:." + str(decimals) + "f}"
    percents        = formatStr.format(100 * (iteration / float(total)))
    filledLength    = int(round(barLength * iteration / float(total)))
    bar             = '*' * filledLength + '-' * (barLength - filledLength)
    sys.stdout.write('\r%s |%s| %s%s %s' % (prefix, bar, percents, '%', suffix)),
    sys.stdout.flush()
    if iteration == total:
        sys.stdout.write('\n')
        sys.stdout.flush()

h5_save = True

X, Y, LA = sp.symbols('X, Y, LA')
rho, qx, qy, T = sp.symbols('rho, qx, qy, T')

def init_T(x, y):
    #return Td + (Tu-Td)/(ymax-ymin)*(y-ymin)
    return Td + (Tu-Td)/(ymax-ymin)*(y-ymin) + (Td-Tu) * (0.1*np.random.random_sample((x.shape[0],y.shape[1]))-0.05)

def bc_up(f, m, x, y):
    m[qx] = 0.
    m[qy] = 0.
    m[T] = Tu

def bc_down(f, m, x, y):
    np.random.seed(1)
    m[qx] = 0.
    m[qy] = 0.
    m[T] = Td# + (Td-Tu) * 5 * (0.1*np.random.random_sample((x.shape[0],1))-0.05)

def save(sol, filename, path, num):
    h5 = pylbm.H5File(sol.domain.mpi_topo, filename, path, num)
    h5.set_grid(sol.domain.x, sol.domain.y)
    h5.add_scalar('T', sol.m[T])
    h5.save()

# parameters
Tu = -0.5
Td =  0.5
xmin, xmax, ymin, ymax = 0., 2., 0., 1.
Ra = 2000
Pr = 0.71
Ma = 0.01
alpha = .005
if h5_save:
    dx = 1./256 # spatial step
else:
    dx = 1./128
la = 1. # velocity of the scheme
rhoo = 1.
g = 9.81

nu = np.sqrt(Pr*alpha*9.81*(Td-Tu)*(ymax-ymin)/Ra)
kappa = nu/Pr
eta = nu
snu = 1./(.5+3*nu)
seta = 1./(.5+3*eta)
sq = 8*(2-snu)/(8-snu)
se = seta
sf = [0., 0., 0., seta, se, sq, sq, snu, snu]
a = .5
skappa = 1./(.5+10*kappa/(4+a))
se = 1./(.5+np.sqrt(3)/3)
snu = se
sT = [0., skappa, skappa, se, snu]

dico = {
    'box': {
        'x': [xmin, xmax],
        'y': [ymin, ymax],
        'label': [-1, -1, 0, 1]
    },
    'space_step': dx,
    'scheme_velocity': la,
    'schemes': [
        {
            'velocities': list(range(9)),
            'conserved_moments': [rho, qx, qy],
            'polynomials': [
                1, X, Y,
                3*(X**2+Y**2)-4,
                0.5*(9*(X**2+Y**2)**2-21*(X**2+Y**2)+8),
                3*X*(X**2+Y**2)-5*X, 3*Y*(X**2+Y**2)-5*Y,
                X**2-Y**2, X*Y
            ],
            'relaxation_parameters': sf,
            'equilibrium': [
                rho, qx, qy,
                -2*rho + 3*(qx**2+qy**2),
                rho - 3*(qx**2+qy**2),
                -qx, -qy,
                qx**2 - qy**2, qx*qy
            ],
            'source_terms': {qy: alpha*g*T},
        },
        {
            'velocities': list(range(5)),
            'conserved_moments': T,
            'polynomials': [1, X, Y, 5*(X**2+Y**2) - 4, (X**2-Y**2)],
            'equilibrium': [T, T*qx, T*qy, a*T, 0.],
            'relaxation_parameters': sT,
        },
    ],
    'init': {
        rho: 1.,
        qx: 0.,
        qy: 0.,
        T: init_T,
    },
    'boundary_conditions': {
        0: {'method': {0: pylbm.bc.BouzidiBounceBack,
                       1: pylbm.bc.BouzidiAntiBounceBack},
            'value': bc_down},
        1: {'method': {0: pylbm.bc.BouzidiBounceBack,
                       1: pylbm.bc.BouzidiAntiBounceBack},
            'value': bc_up},
    },
    'generator': 'cython',
}

sol = pylbm.Simulation(dico)

x, y = sol.domain.x, sol.domain.y

if h5_save:
    Tf = 500.
    im = 0
    l = Tf / sol.dt / 32
    printProgress(im, l, prefix='Progress:', suffix='Complete', barLength=50)
    filename = 'Rayleigh_Benard'
    path = './data_' + filename
    save(sol, filename, path, im)
    while sol.t < Tf:
        for k in range(32):
            sol.one_time_step()
        im += 1
        printProgress(im, l, prefix='Progress:', suffix='Complete', barLength=50)
        save(sol, filename, path, im)
else:
    viewer = pylbm.viewer.matplotlib_viewer
    fig = viewer.Fig()
    ax = fig[0]
    image = ax.image(sol.m[T].T, cmap='cubehelix', clim=[Tu, Td+.25])

    def update(iframe):
        nrep = 64
        for i in range(nrep):
            sol.one_time_step()
        image.set_data(sol.m[T].T)
        ax.title = "Solution t={0:f}".format(sol.t)

    # run the simulation
    fig.animate(update, interval=1)
    fig.show()
