#!/usr/bin/python3

import argparse
import numpy as np
import sys

import teukolsky_data

parser = argparse.ArgumentParser('Solve the constraints 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('--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)
parser.add_argument('--upper-branch',     action = 'store_true')
parser.add_argument('--print-coeffs',     action = 'store_true')

parser.add_argument('amplitude', type = float)

args = parser.parse_args()

nb_coeffs       = [args.coeffs_radial, args.coeffs_angular]
scale_factor    = [args.scale_factor, args.scale_factor]

td = teukolsky_data.TeukolskyData(family = args.family, amplitude = args.amplitude,
                                  nb_coeffs = nb_coeffs, basis_scale_factor = scale_factor,
                                  atol = args.solve_tolerance, max_iter = args.max_iter,
                                  solution_branch = int(args.upper_branch))

r     = np.linspace(1e-3, 20, 256)
theta = np.linspace(0, 2.0 * np.pi, 32)

res = np.abs(td.eval_residual(r, theta))
sys.stderr.write('residual: %g %g %g\n' % (np.max(res[0]), np.max(res[1]), np.max(res[2])))

if args.print_coeffs:
    for c in td.coeffs:
        for j in range(nb_coeffs[1]):
            for i in range(nb_coeffs[0]):
                sys.stdout.write('%g ' % c[j, i])
            sys.stdout.write('\n')

        sys.stdout.write('\n')

sys.exit(0)

import matplotlib
matplotlib.use('Agg')
matplotlib.rc('text', usetex = True)
matplotlib.rc('font', size = 8)
matplotlib.rc('text.latex', preamble = r'\usepackage{times}\usepackage{mathrsfs}')
matplotlib.rcParams['font.family'] = 'serif'

from matplotlib import pyplot

r     = np.linspace(0, 20, 1000)
theta = np.linspace(0, 2.0 * np.pi, 200)
R, Theta = np.meshgrid(r, theta)
X = R * np.sin(Theta)
Z = R * np.cos(Theta)

varnames = ['$u$', r'$K_{r}^r$', r'$K_{\varphi}^\varphi$']
vals = [td.eval_psi(r, theta) - 1.0, td.eval_krr(r, theta), td.eval_kpp(r, theta)]

ax_lims = [[0, 1.0], [-1, 2], [-1, 1]]

fig  = pyplot.figure(figsize = (6, 4))
axes = [[None] * 2] * 3
for i in range(3):
    data = vals[i]
    maxval = np.max(np.abs(data))

    for j in range(2):
        ax = fig.add_subplot(2, 3, j * 3 + i + 1)
        axes[i][j] = ax

        if j == 0:
            ax.pcolormesh(X, Z, data, vmin = -maxval, vmax = maxval)
            ax.set_xlim(-4, 4)
            ax.set_ylim(-4, 4)
            ax.set_aspect('equal')
            ax.set_title(varnames[i])
            ax.set_xlabel('$x$')
            ax.set_ylabel('$z$')
        else:
            ax.plot(r,  data[0], 'b--')
            ax.plot(-r, data[100], 'b--')

            ax.plot(r,  data[50], 'g-')
            ax.plot(-r, data[150], 'g-')

            ax.set_xlim(-8, 8)
            #ax.set_ylim(*ax_lims[i])

fig.canvas.draw()
fig.savefig(sys.argv[1], dpi = 600)