import numpy as np

import matplotlib as mpl
from matplotlib import pyplot as plt

import sympy
from sympy import lambdify, latex
from sympy.parsing.sympy_parser import parse_expr

import argparse
from tokenize import TokenError
from sys import exit

parser = argparse.ArgumentParser()
parser.add_argument('f')
parser.add_argument('g')
parser.add_argument('a', type=float)
parser.add_argument('b', type=float)
parser.add_argument('-y_a', type=float)
parser.add_argument('-y_b', type=float)
parser.add_argument('n_segs', type=int)
parser.add_argument('n_iter', type=int)
parser.add_argument('-tol', type=float, default=1e-2)
parser.add_argument('-thresh', type=float, default=1e4)
parser.add_argument('-to_x', type=float)
parser.add_argument('-to_y', type=float)
args = parser.parse_args()

# Parse component maps
try:
    f_expr = parse_expr(args.f)
    g_expr = parse_expr(args.g)
except TokenError as e:
    print('Failed to parse function: "{}"'.format(e))
    exit()
symbols = sympy.symbols(['x', 'y'])
f = lambdify(symbols, f_expr, 'numpy')
g = lambdify(symbols, g_expr, 'numpy')

# Local store parameters
x_a, x_b = args.a, args.b
y_a = args.y_a if args.y_a else x_a
y_b = args.y_b if args.y_b else x_b
n_segs = args.n_segs
n_iter = args.n_iter
tol, thresh = args.tol, args.thresh
to_x, to_y = args.to_x, args.to_y

# Construct mesh of initial conditions
x = np.linspace(x_a, x_b, n_segs)
y = np.linspace(y_b, y_a, n_segs)
x, y = np.meshgrid(x, y)

# Increment initial conditions
for _ in range(n_iter):
    x, y = f(x, y), g(x, y)

if to_x is None and to_y is None:
    # Converged if finite
    data = np.isfinite(x + y)
else:
    data = np.abs(to_x - x)**2 + np.abs(to_y - y)**2 >= tol

# Plot
plt.imshow(data, extent=(x_a, x_b, y_a, y_b), cmap=mpl.cm.gray)
plt.show()