#!/usr/bin/env python3
"""Weber boundary Wronskian plus an independent finite-difference spectrum.

Supported baseline: Python 3.10+, mpmath 1.3.0, NumPy and SciPy (see the
guide's numerical environment). The eigenvalue comparison solves a Dirichlet
problem on a truncated real interval; mesh and endpoint refinements are
reported separately. No exact eigenvalues are used in the matrix assembly.
These refinements are numerical evidence, not certified error bounds.
Exit 0 = PASS, 1 = numerical FAIL, 2 = invalid input.
"""

import argparse
import json
import sys

import mpmath as mp
import numpy as np
from scipy.linalg import eigh_tridiagonal


def grid_spectrum(hbar, half_width, intervals, levels):
    # half_width is measured in oscillator lengths sqrt(hbar).
    length = half_width*np.sqrt(hbar)
    dx = 2*length/intervals
    x = np.linspace(-length, length, intervals+1)[1:-1]
    diagonal = 2*hbar*hbar/(dx*dx)+x*x
    off_diagonal = np.full(len(x)-1, -hbar*hbar/(dx*dx))
    return eigh_tridiagonal(diagonal, off_diagonal, select='i',
                            select_range=(0, levels-1), eigvals_only=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--hbar', default='1')
    parser.add_argument('--intervals', type=int, default=1200)
    parser.add_argument('--half-width', type=float, default=8.0)
    parser.add_argument('--levels', type=int, default=4)
    parser.add_argument('--tolerance', type=float, default=2e-7,
                        help='absolute spectral error measured in units of hbar')
    parser.add_argument('--json', action='store_true')
    args = parser.parse_args()
    mp.mp.dps = 50
    try:
        h = mp.mpf(args.hbar)
    except (ValueError, TypeError):
        parser.error('--hbar must be a positive finite real number')
    if not mp.isfinite(h) or not (mp.mpf('1e-6') <= h <= mp.mpf('1e6')):
        parser.error('--hbar must lie between 1e-6 and 1e6 for this floating-point grid')
    if args.intervals < 32 or args.intervals % 4 or not (1 <= args.levels <= 12):
        parser.error('--intervals must be a multiple of 4, at least 32; --levels must lie in 1..12')
    if not np.isfinite(args.half_width) or not (1 <= args.half_width <= 20):
        parser.error('--half-width must lie in 1..20 oscillator lengths')
    if not np.isfinite(args.tolerance) or args.tolerance <= 0:
        parser.error('--tolerance must be finite and positive')

    def boundary_wronskian(energy, x):
        nu = energy/(2*h)-mp.mpf(1)/2
        k = mp.sqrt(2/h)
        right = lambda t: mp.pcfd(nu, k*t)
        left = lambda t: mp.pcfd(nu, -k*t)
        return left(x)*mp.diff(right, x)-mp.diff(left, x)*right(x)

    # Off-spectrum: check sign, scale, and constancy away from a gamma pole.
    samples = []
    for ratio in ['0', '1.4', '2.5', '6.2']:
        energy = mp.mpf(ratio)*h
        exact = -2*mp.sqrt(mp.pi/h)*mp.rgamma(mp.mpf(1)/2-energy/(2*h))
        error = max(abs(boundary_wronskian(energy, mp.mpf(t)*mp.sqrt(h))-exact)
                    /(1+abs(exact)) for t in ['-0.4', '0', '0.7'])
        samples.append({'E_over_hbar': ratio, 'scaled_wronskian_error': mp.nstr(error, 8)})
    wronskian_error = max(mp.mpf(row['scaled_wronskian_error']) for row in samples)
    hf = float(h)
    n = args.intervals
    grids = [grid_spectrum(hf, args.half_width, m, args.levels) for m in [n, 2*n, 4*n]]
    extrapolated_1 = (4*grids[1]-grids[0])/3
    extrapolated_2 = (4*grids[2]-grids[1])/3
    # Increase the interval length and grid count together to keep dx fixed.
    extended = [grid_spectrum(hf, args.half_width*1.25, 5*m//4, args.levels)
                for m in [2*n, 4*n]]
    extended_extrapolated = (4*extended[1]-extended[0])/3
    exact_levels = np.array([2*j+1 for j in range(args.levels)], dtype=float)
    error = float(np.max(np.abs(extrapolated_2/hf-exact_levels)))
    refinement = float(np.max(np.abs((extrapolated_2-extrapolated_1)/hf)))
    endpoint = float(np.max(np.abs((extended_extrapolated-extrapolated_2)/hf)))
    passed = (np.isfinite(error+refinement+endpoint) and wronskian_error < mp.mpf('1e-40')
              and max(error, refinement, endpoint) < args.tolerance)
    result = {'status': 'PASS' if passed else 'FAIL', 'hbar': str(h),
              'half_width_in_oscillator_lengths': args.half_width,
              'grid_intervals': [n, 2*n, 4*n],
              'tolerance_in_hbar_units': args.tolerance,
              'wronskian_checks': samples,
              'grid_spectra_in_hbar_units': [(values/hf).tolist() for values in grids],
              'richardson_spectrum_in_hbar_units': (extrapolated_2/hf).tolist(),
              'exact_spectrum_in_hbar_units': exact_levels.tolist(),
              'maximum_spectral_error': error, 'richardson_refinement_change': refinement,
              'endpoint_refinement_change': endpoint}
    print(json.dumps(result, indent=2) if args.json else
          f"{result['status']}: E/hbar = {result['richardson_spectrum_in_hbar_units']}\n"
          f"Spectral error {error:.3e}; mesh refinement {refinement:.3e}; endpoint change {endpoint:.3e}")
    return 0 if passed else 1


if __name__ == '__main__':
    sys.exit(main())