#!/usr/bin/python3

import argparse
import sys

import numpy as np

import teukolsky_data

parser = argparse.ArgumentParser('Find A_max for Teukolsky-type data')

parser.add_argument('--coeffs-radial',    type = int,   default = 64)
parser.add_argument('--coeffs-angular',   type = int,   default = 16)
parser.add_argument('--bounds-max',       type = float, default = np.nan)
parser.add_argument('--bounds-min',       type = float, default = np.nan)
parser.add_argument('--bounds-tolerance', type = float, default = 1e-8)
parser.add_argument('--solve-tolerance',  type = float, default = 1e-12)
parser.add_argument('--scale-factor',     type = float, default = 3.0)
parser.add_argument('--max-iter',         type = int,   default = 32)
parser.add_argument('--family',           type = int,   default = 0)

args = parser.parse_args()

bounds    = [args.bounds_min, args.bounds_max]
nb_coeffs = [args.coeffs_radial, args.coeffs_angular]

if not (np.isnan(bounds[0]) or np.isnan(bounds[1])):
    A = 0.5 * (bounds[0] + bounds[1])
else:
    A = 1.0

def converged(bounds, tol):
    return (np.isfinite(bounds[0]) and np.isfinite(bounds[1]) and
            np.abs(bounds[0] - bounds[1]) < tol)

sys.stderr.write('Solving with %dx%d coeffs\n' % (nb_coeffs[0], nb_coeffs[1]))

prev_coeffs = None
while not converged(bounds, args.bounds_tolerance):
    sys.stderr.write('Trying A=%g (%s)\n' % (A, bounds))
    try:
        td_args = {
            'amplitude'             : A,
            'nb_coeffs'             : nb_coeffs,
            'basis_scale_factor'    : [args.scale_factor, args.scale_factor],
            'atol'                  : args.solve_tolerance,
            'max_iter'              : args.max_iter,
            'family'                : args.family,
        }

        #if prev_coeffs:
        #    args['coeffs_init'] = prev_coeffs

        td = teukolsky_data.TeukolskyData(**td_args)

        bounds[0] = A
        sys.stderr.write('A=%16.16g converge\n' % A)
        prev_coeffs = [c[:] for c in td.coeffs]
    except Exception as e:
        bounds[1] = A
        sys.stderr.write('A=%16.16g diverge\n' % A)

    if np.isnan(bounds[0]):
        A *= 0.5
    elif np.isnan(bounds[1]):
        A *= 2
    else:
        A = 0.5 * (bounds[1] + bounds[0])

sys.stdout.write('%16.16g %16.16g\n' % (bounds[0], bounds[1]))