#!/usr/bin/env python2
#
#   In this file test glue layer between Fortran 95 code and module generated by f2py from numpy.
#   Here and bellow `native solution` means solution that was obtained with native fortran code
#   and in opposite `glue solution` means native solution transmitted through glue layer.
#
#   (c) Daniel Bershatsky <daniel.bershatsky@skolkovotech.ru>, 2016
#

from __future__ import print_function
from nls import nls
from numpy import array, arange, diag, abs, max
from scipy.sparse import diags
from subprocess import Popen, PIPE


VERBOSE = True
FLOAT_THRESHOLD = 1e-3

NX, ITERS, ORDER = 200, 10000, 5
DT, DX = 1.0e-6, 1.0e-3


def test_make_banded_matrix():
    print('Test make_banded_matrix():')

def test_banded_matvec():
    print('Test banded_matvec():')

def test_make_laplacian():
    print('Test make_laplacian():')
    n, m, h = 8, 8, 0.1
    r = arange(0, h * m - 1e-10, h)
    D1 = diags((1, -8, 0, 8, -1), (-2, -1, 0, 1, 2), shape=(m, m)).toarray()
    D2 = diags((-1, 16, -30, 16, -1), (-2, -1, 0, 1, 2), shape=(m, m)).toarray()
    r[0] = 1
    D1[0, :] = 0
    D1[1, 1] += 1
    D2[0, :3] = [-60, 64, -4]
    D2[1, 1] += -1
    D = D2 / (24 * h ** 2) + diag(1.0 / r).dot(D1)  / (12 * h)
    L = nls.make_laplacian(n, 5, h)
    tolerance = 5.0e-5
    for k in xrange(0, 3):
        err = max(abs(diag(D, 2 - k) - L[k, 2 - k:]))
        print('Diagonal # ', 2 - k, ':    ', err)
        if err >= tolerance:
            print(diag(D, 2 - k))
            print(L[k, 2 - k:])
    for k in xrange(1, 3):
        err = max(abs(diag(D, -k) - L[2 + k, :-k]))
        print('Diagonal #', -k, ':    ', err)
        if err >= tolerance:
            print(diag(D, -k))
            print(L[2 + k, :-k])

def test_runge_kutta():
    """Call runge_kutta() fortran subroutine directly through f2py wrapper layer.
    """
    print('Test runge_kutta():')
    u0 = array([1, 1, 1, 1, 1, 1, 1, 1])
    op = nls.make_laplacian(8, 5, 0.01)
    u = nls.runge_kutta(0.00001, 0.0, 0.01, u0, op, 10)
    print(u)

def get_native_solution():
    pid = Popen(['./solve'], stdin=PIPE, stdout=PIPE, stderr=PIPE)
    output, _ = pid.communicate()
    return array([float(x) for x in output.split()])


def get_glue_solution():
    u = nls.solve_nls(DT, DX, NX, ORDER, ITERS)
    return abs(u.conj() * u)

def compare_solutions():
    native = get_native_solution()
    glue = get_native_solution()
    difference = abs(native - glue) < FLOAT_THRESHOLD
    different = reduce(lambda acc, val: acc if val[1] else acc + [val[0]], enumerate(difference), [])
    return all(difference), different

def report_solution_comparison(is_equal, different):
    print('Equal:', is_equal)
    
    if VERBOSE and not is_equal:
        print('Next elements are different:')
        print(different)
        print('Native:\n', native[different])
        print('Glue:\n', glue[different])

def test_glue_layer():
    """Run native fortran code and its python bindings in order to compare these solutions and test
    their equality.
    """
    print('Test glue layer:')
    is_equal, different = compare_solutions()
    report_solution_comparison(is_equal, different)

def test():
    test_make_banded_matrix()
    test_banded_matvec()
    test_make_laplacian()
    test_runge_kutta()
    test_glue_layer()


if __name__ == '__main__':
    test()
