"""Optional integration test: run with the recovered TOP Python environment.

Usage: python tests/serialize/top_reader.py build/tests/serialize/top_export_fixture
The compiled eq_poly_ester_gamma2 module must be available to top.load().
"""
import json
import os
from pathlib import Path
import subprocess
import sys
import tempfile

import numpy as np
import top


def read_model(directory):
    dati = directory / 'dati'
    dati.write_text('nt = 1\nlres = 40\niparity = 0\niCor = 1\n'
                    'dertype = CHEB\norderFD = 2\nmattype = FULL\n'
                    'nsol = 1\nshift = 0.6\nrootdir = ./\nm = 0\ntagdir = test\n')
    os.chdir(directory)
    solver = top.load('eq_poly_ester_gamma2')
    solver.read_dati(str(dati))
    model = solver.init_model(str(directory) + '/')
    return solver, model


def verify(directory, deformed):
    solver, model = read_model(directory)
    radius, theta = solver.get_grid()
    z = solver.get_zeta()[:, None]
    mu = np.cos(theta)[None, :]
    if deformed:
        rs = 1.0 - 0.2 * mu**2
        a = 0.5 * (5.0*z**3 - 3.0*z**5)
        expected_radius = 0.8*z + a*(rs-0.8)
        p2 = 0.5*(3.0*mu**2-1.0)
        h = (1-z*z)*(1+0.1*z*z*p2)
        hz = -2*z + 0.1*p2*(2*z-4*z**3)
        hzz = -2 + 0.1*p2*(2-12*z*z)
    else:
        expected_radius = np.broadcast_to(z, radius.shape)
        h = np.broadcast_to(np.sinc(z), radius.shape)
        with np.errstate(divide='ignore', invalid='ignore'):
            hz = np.pi*np.cos(np.pi*z)/(np.pi*z) - np.sin(np.pi*z)/(np.pi*z*z)
            hzz = -np.pi*np.sin(np.pi*z)/z - 2*np.cos(np.pi*z)/(z*z) + 2*np.sin(np.pi*z)/(np.pi*z**3)
        hz[0] = 0.0
        hzz[0] = -np.pi**2/3
        hz = np.broadcast_to(hz, radius.shape)
        hzz = np.broadcast_to(hzz, radius.shape)
    np.testing.assert_allclose(radius, expected_radius, rtol=0, atol=2e-12)
    np.testing.assert_allclose(model['hh'], h, rtol=0, atol=3e-12)
    np.testing.assert_allclose(model['hhz'], hz, rtol=0, atol=3e-9)
    np.testing.assert_allclose(model['hhzz'], hzz, rtol=0, atol=2e-5)
    np.testing.assert_allclose(model['lambda'], np.pi**2, rtol=3e-13)
    np.testing.assert_allclose(model['omga'], 0.2/np.pi if deformed else 0.0, rtol=3e-13, atol=1e-15)
    if not deformed:
        result = solver.run_arncheb(0.6)
        assert result.nsol == 1
        np.testing.assert_allclose(result.valps[0], 0.6168010124, rtol=0, atol=2e-8)
        return float(result.valps[0])


if __name__ == '__main__':
    if len(sys.argv) != 2:
        raise SystemExit(__doc__)
    executable = Path(sys.argv[1]).resolve()
    original_directory = Path.cwd()
    try:
        with tempfile.TemporaryDirectory(prefix='serif-top-reader-') as tmp:
            directory = Path(tmp)
            subprocess.run([str(executable), str(directory)], check=True)
            frequency = verify(directory / 'spherical', False)
            verify(directory / 'deformed', True)
            print(json.dumps({'reader_checks': 'passed', 'radial_frequency_top_units': frequency}))
    finally:
        os.chdir(original_directory)