#!/usr/bin/env python

'''

* Ostroumov, G. A. (1958). Free convection under the conditions of the
  internal problem. Technical Report NACA-TM-1407, NACA, eq. 5.18,
  p. 34

'''

from os.path import basename, splitext

from matplotlib.pyplot import subplots
import numpy as np
from pandas import DataFrame

from firedrake import (Mesh, FunctionSpace, SpatialCoordinate,
                       Function, TestFunction, DirichletBC, solve,
                       File)

from ufl import dx, inner, grad


def main(meshfile: str) -> None:

    mesh = Mesh(meshfile)

    V = FunctionSpace(mesh, 'CG', 1)
    x = SpatialCoordinate(V.mesh())[0]
    w = Function(V)
    v = TestFunction(V)

    solve((inner(grad(w), grad(v)) - inner(x, v)) * dx == 0, w,
          bcs=DirichletBC(V, 0.0, [1, 2]),
          solver_parameters={'ksp_type': 'cg'})

    x_axis = DirichletBC(V, None, [3]).nodes
    transect = DataFrame({'x': Function(V).interpolate(x).dat.data_ro[x_axis],
                          'w': w.dat.data_ro[x_axis]})
    wmax = transect.iloc[transect.w.idxmax()]
    print('Maximum:', wmax)

    name = splitext(basename(meshfile))[0]
    fig, ax = subplots()
    fig.suptitle(r'Velocity along the $x$-axis')
    ax.plot(transect.x, transect.w, marker='o', linestyle='None',
            label='Firedrake')
    s = np.linspace(0, 1)
    ax.plot(s, s/8*(1-s**2), label='exact')
    ax.plot(wmax.x, wmax.w,
            marker='o', linestyle='None', label='max(Firedrake)')
    ax.plot(1/np.sqrt(3), 1/12/np.sqrt(3),
            marker='o', linestyle='None', label='max(exact)')
    ax.set_ylabel(r'$w$')
    ax.set_ylim((0, None))
    ax.set_xlabel(r'$x$')
    ax.set_xlim((0, 1))
    ax.legend()
    fig.savefig(name + '.png')

    File(name + '.pvd').write(w)


if __name__ == '__main__':

    from sys import argv

    main(argv[1])