import numpy as np
import os
import sys

from pathlib import Path

def write_header(file):
    header = [
'/*--------------------------------\*- C++ -\*-------------------------------*\\',
'| =========                 |                                                 |',
'| \\      /  F ield         | OpenFOAM: The Open Source CFD Toolbox           |',
'|  \\    /   O peration     | Version:  4.x                                   |',
'|   \\  /    A nd           | Web:      www.OpenFOAM.org                      |',
'|    \\/     M anipulation  |                                                 |',
'\*---------------------------------------------------------------------------*/',
'FoamFile',
'{',
'    version     2.0;',
'    format      ascii;',
'    class       dictionary;',
'    object      blockMeshDict;',
'}',
'// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * //\n',
'scale 1;\n']

    with open(file, 'w') as f:
        for line in header:
            f.write(line + '\n')

        f.close()

def write_vertices_variables(file, Lx, Lz, Hmin_x):
# write the initial variables for the calculating the vertices
    with open(file, 'a') as f:
        f.write('xmin 0;\n')
        f.write('xmax ' + str(Lx) + ';\n')
        f.write('zmin 0;\n')
        f.write('zmax ' + str(Lz) + ';\n\n')
        f.write('ymin ' +str(-Hmin_x) + ';\n')
        f.write('ymax ' +str(Hmin_x) + ';\n\n')
    
        f.close()

def write_vertices(file):
# writes the vertices based on the variables written by the previous method
    with open(file, 'a') as f:
        f.write('vertices\n(\n')
        f.write('\t($xmin  $ymin  $zmin) //0\n')
        f.write('\t($xmax  $ymin  $zmin) //1\n')
        f.write('\t($xmax  $ymax  $zmin) //2\n')
        f.write('\t($xmin  $ymax  $zmin) //3\n')
        f.write('\t($xmin  $ymin  $zmax) //4\n')
        f.write('\t($xmax  $ymin  $zmax) //5\n')
        f.write('\t($xmax  $ymax  $zmax) //6\n')
        f.write('\t($xmin  $ymax  $zmax) //7\n')
        f.write(');\n\n')

        f.close()

def write_block_variables(file, dim, Lx, Lz, Hmin_x, Hmin_z, xcells, ycells, zcells, x_stretch, y_stretch, z_stretch):
# writes the variables needed to define the blocks
    with open(file, 'a') as f:
        f.write('xcells ' + str(int(Lx*xcells)) + ';\n')
        f.write('ycells '  + str(int(2*Hmin_x*ycells)) + ';\n')
        
        if dim == 2:
            f.write('zcells 1;\n\n')
        else:
            f.write('zcells '  + str(int(Lz*zcells)) + ';\n')
        
        f.write('stretchX ' + str(x_stretch) + ';\n')
        f.write('stretchY ' + str(y_stretch) + ';\n')
        f.write('stretchZ ' + str(z_stretch) + ';\n\n')
        f.write('stretchXI ' + str(1/x_stretch) + ';\n')
        f.write('stretchYI ' + str(1/y_stretch) + ';\n')
        f.write('stretchZI ' + str(1/z_stretch) + ';\n\n')

        f.close()

def write_blocks(file, dim):
# writes the blocks
    with open(file, 'a') as f:
        f.write('blocks\n(\n')
        f.write('//block0 \n')
        f.write('\thex (0 1 2 3 4 5 6 7) ($xcells $ycells $zcells) simpleGrading \n\t(\n')
        
        if dim == 2:
            f.write('\t// x-direction grading\n')
            f.write('\t1\n')
        else:
            # x-direction grading
            f.write('\t// x-direction grading\n')
            f.write('\t1\n')
            # f.write('\t// z-direction grading\n')
            # f.write('\t// x-direction grading\n\t(\n')
            # f.write('\t\t(0.2 0.25 $stretchX)\n')
            # f.write('\t\t(0.6 0.50 1)\n')
            # f.write('\t\t(0.2 0.25 $stretchXI)\n')
            # f.write('\t)\n\n')

        # y-direction grading
        f.write('\t// y-direction grading\n\t(\n')
        f.write('\t\t(0.2 0.25 $stretchY)\n')
        f.write('\t\t(0.6 0.50 1)\n')
        f.write('\t\t(0.2 0.25 $stretchYI)\n')
        f.write('\t)\n\n')
    
        if dim == 2:
            f.write('\t// z-direction grading\n')
            f.write('\t1\n')
        else:
            # z-direction grading          
            f.write('\t// z-direction grading\n\t(\n')
            f.write('\t\t(0.2 0.25 $stretchZ)\n')
            f.write('\t\t(0.6 0.50 1)\n')
            f.write('\t\t(0.2 0.25 $stretchZI)\n')
            f.write('\t)\n')

        f.write('\t)\n);\n\n')

        f.close()

def write_edges(file, dim, sine, curved_inlet, Lx, Lz, Hmin_x, a_x, a_z, discretization):
# writes the edges

    x_points = np.linspace(0, Lx, int(discretization))
    z_points = np.linspace(0, Lz, int(discretization))

    if sine:
        y_x_top = Hmin_x + 2*a_x*np.sin(np.pi*x_points/Lx)**2
        y_x_bottom = -y_x_top
        y_z_top = Hmin_x + 2*a_z*np.sin(np.pi*z_points/Lx)**2
        y_z_bottom = -y_z_top
    else:
        y_x_top = Hmin_x + 2*a_x*np.sin(np.pi*x_points/Lx)
        y_x_bottom = -y_x_top
        y_z_top = Hmin_x + 2*a_z*np.sin(np.pi*z_points/Lx)
        y_z_bottom = -y_z_top

    if dim == 3:
        with open(file, 'a') as f:
            f.write('edges\n(\n')

            # Polylines in the x-direction
            # polyline between 0 1
            f.write('\tpolyLine 0 1\n\t(\n')
            for i, y in enumerate(y_x_bottom):
                f.write('\t\t(' + str(np.round(x_points[i],3)) + ' ' + str(np.round(y,4)) + ' ' + ' $zmin)\n')
            f.write('\t)\n\n')

            # polyline between 3 2
            f.write('\tpolyLine 3 2\n\t(\n')
            for i, y in enumerate(y_x_top):
                f.write('\t\t(' + str(np.round(x_points[i],3)) + ' ' + str(np.round(y,4)) + ' ' + ' $zmin)\n')
            f.write('\t)\n\n')

            # polyline between 4 5
            f.write('\tpolyLine 4 5\n\t(\n')
            for i, y in enumerate(y_x_bottom):
                f.write('\t\t(' + str(np.round(x_points[i],3)) + ' ' + str(np.round(y,4)) + ' ' + ' $zmax)\n')
            f.write('\t)\n\n')

            # polyline between 7 6
            f.write('\tpolyLine 7 6\n\t(\n')
            for i, y in enumerate(y_x_top):
                f.write('\t\t(' + str(np.round(x_points[i],3)) + ' ' + str(np.round(y,4)) + ' ' + ' $zmax)\n')
            f.write('\t)\n\n')
            
            if curved_inlet:
                # Polyline in the z-direction
                # polyline between 0 4
                f.write('\tpolyLine 0 4\n\t(\n')
                for i, y in enumerate(y_z_bottom):
                    f.write('\t\t( $xmin ' + str(np.round(y,4)) + ' ' + str(np.round(z_points[i],3)) + ')\n')
                f.write('\t)\n\n')

                # polyline between 3 4
                f.write('\tpolyLine 3 7\n\t(\n')
                for i, y in enumerate(y_z_top):
                    f.write('\t\t( $xmin ' + str(np.round(y,4)) + ' ' + str(np.round(z_points[i],3)) + ')\n')
                f.write('\t)\n\n')

                # polyline between 1 5
                f.write('\tpolyLine 1 5\n\t(\n')
                for i, y in enumerate(y_z_bottom):
                    f.write('\t\t( $xmax ' + str(np.round(y,4)) + ' ' + str(np.round(z_points[i],3)) + ')\n')
                f.write('\t)\n\n')

                # polyline between 2 6
                f.write('\tpolyLine 2 6\n\t(\n')
                for i, y in enumerate(y_z_top):
                    f.write('\t\t( $xmax ' + str(np.round(y,4)) + ' ' + str(np.round(z_points[i],3)) + ')\n')
                f.write('\t)\n')

            f.write(');\n\n')
    else:
        with open(file, 'a') as f:
            f.write('edges\n(\n')

            # Polylines in the x-direction
            # polyline between 0 1
            f.write('\tpolyLine 0 1\n\t(\n')
            for i, y in enumerate(y_x_bottom):
                f.write('\t\t(' + str(np.round(x_points[i],3)) + ' ' + str(np.round(y,4)) + ' ' + ' $zmin)\n')
            f.write('\t)\n\n')

            # polyline between 3 2
            f.write('\tpolyLine 3 2\n\t(\n')
            for i, y in enumerate(y_x_top):
                f.write('\t\t(' + str(np.round(x_points[i],3)) + ' ' + str(np.round(y,4)) + ' ' + ' $zmin)\n')
            f.write('\t)\n\n')

            # polyline between 4 5
            f.write('\tpolyLine 4 5\n\t(\n')
            for i, y in enumerate(y_x_bottom):
                f.write('\t\t(' + str(np.round(x_points[i],3)) + ' ' + str(np.round(y,4)) + ' ' + ' $zmax)\n')
            f.write('\t)\n\n')

            # polyline between 7 6
            f.write('\tpolyLine 7 6\n\t(\n')
            for i, y in enumerate(y_x_top):
                f.write('\t\t(' + str(np.round(x_points[i],3)) + ' ' + str(np.round(y,4)) + ' ' + ' $zmax)\n')
            f.write('\t)\n\n')
            
            f.write(');\n\n')

def write_boundary_variables(file):
# writes variables for cyclic boundary condition
    with open(file, 'a') as f:
        f.write('// needed for the seperation vector \n')
        f.write('minusXmax   #eval{ -1*$xmax };\n\n')

        f.close()

def write_inlet_boundary(file):
# writes the inlet boundary
    with open(file, 'a') as f:
        f.write('\tinlet\n\t{\n')
        f.write('\t\t//type patch;\n')
        f.write('\t\ttype cyclic;\n')
        f.write('\t\ttransform translational;\n')
        f.write('\t\tneighbourPatch outlet;\n')
        f.write('\t\tseparationVector ($xmax 0 0 );\n')
        f.write('\t\tfaces\n\t\t(\n')
        f.write('\t\t\t(0 4 7 3)\n')
        f.write('\t\t);\n')
        f.write('\t}\n')

        f.close()

def write_outlet_boundary(file):
# writes the outlet boundary
    with open(file, 'a') as f:
        f.write('\toutlet\n\t{\n')
        f.write('\t\t//type patch;\n')
        f.write('\t\ttype cyclic;\n')
        f.write('\t\ttransform translational;\n')
        f.write('\t\tneighbourPatch inlet;\n')
        f.write('\t\tseparationVector ( $minusXmax 0 0 );\n')
        f.write('\t\tfaces\n\t\t(\n')
        f.write('\t\t\t(5 1 2 6)\n')
        f.write('\t\t);\n')
        f.write('\t}\n')

        f.close()

def write_wall_boundary(file,dim):
# writes the wall boundary
    with open(file, 'a') as f:
        f.write('\twalls\n\t{\n')
        f.write('\t\ttype wall;\n')
        f.write('\t\tfaces\n\t\t(\n')
        if dim == 2:
            f.write('\t\t\t(2 3 7 6)\n')
            f.write('\t\t\t(0 1 5 4)\n')
        else:
            f.write('\t\t\t(2 3 7 6)\n')
            f.write('\t\t\t(0 1 5 4)\n')
            f.write('\t\t\t(4 5 6 7)\n')
            f.write('\t\t\t(1 0 3 2)\n')

        f.write('\t\t);\n')
        f.write('\t}\n')

def write_empty_boundary(file):
# writes the wall boundary
    with open(file, 'a') as f:
        f.write('\tfrontAndBack\n\t{\n')
        f.write('\t\ttype empty;\n')
        f.write('\t\tfaces\n\t\t(\n')
        f.write('\t\t\t(4 5 6 7)\n')
        f.write('\t\t\t(1 0 3 2)\n')
        f.write('\t\t);\n')
        f.write('\t}\n')

def write_boundaries(file, dim):
# write the boundaries
    with open(file, 'a') as f:
        f.write('boundary\n(\n')

    write_inlet_boundary(file)
    write_outlet_boundary(file)
    write_wall_boundary(file,dim)
    if dim == 2:
        write_empty_boundary(file)
    
    with open(file, 'a') as f:
        f.write(');\n\n')
        f.write('mergePatchPairs\n(\n);\n')

def main(argv):

    file_path = Path(r'./system/blockMeshDict')

    write_header(file_path)

#-----------------------------------------------------------
    #change these variables for a different geometry channel
    
    # 2D or 3D geometry
    dimensions = 2

    # whether sine or arc
    sinusoidal = False

    # curved or straight inlet and outlet
    curved_inlet = False

    # channel parameters
    Lx = 1.0
    Lz = 0.02

    # factors to calculate other parameters
    Hmin_over_Hmax_x = 0.3
    Hmin_over_Hmax_z = 0.6
    Lx_over_a = 8
    Lz_over_a = 26

    # calculate parameters needed for channel construction
    a_x = Lx/Lx_over_a # factor of max height for x
    a_z = Lz/Lz_over_a # factor of min height for z

    # amplitude of the corrugation
    Hmax_x = 4*Lx/(Lx_over_a*(1-Hmin_over_Hmax_x))
    Hmax_z = 4*Lz/(Lz_over_a*(1-Hmin_over_Hmax_z))

    Hmin_x = np.round(Hmin_over_Hmax_x*Hmax_x,2)
    Hmin_z = np.round(Hmin_over_Hmax_z*Hmax_z,2)

    Havg_x = (Hmax_x + Hmin_x)/2
    Havg_z = (Hmax_z + Hmin_z)/2

    print('Havg_x: %f, Havg_z: %f' % (Havg_x, Havg_z))

    # discretization of the polyline
    discretization = 40

    # number of cells in each direction
    xcells = 80
    ycells = 150
    zcells = 80

    # stretch of the cells in a particular direction
    z_stretch = 10
    y_stretch = 10
    x_stretch = 10

    write_vertices_variables(file_path, Lx, Lz, Hmin_x)
    write_vertices(file_path)
    write_block_variables(file_path, dimensions, Lx, Lz, Hmin_x, Hmin_z, zcells, ycells, xcells, z_stretch, y_stretch, x_stretch)
    write_blocks(file_path, dimensions)
    write_edges(file_path, dimensions, sinusoidal, curved_inlet, Lx, Lz, Hmin_x, a_x, a_z, discretization)
    write_boundary_variables(file_path)
    write_boundaries(file_path, dimensions)


if __name__ == "__main__":
   main(sys.argv[1:])