#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Jan 10 12:02:37 EST 2023
#from: https://stackoverflow.com/questions/34638040/plot-arbitrary-2-d-function-in-python-pyplot-like-matlabs-ezplot
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    tt = Timer('start ' + ' '.join(sys.argv))
import sys, os.path, os, glob, datetime
#import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
from sympy.plotting import plot_implicit
from sympy.parsing.sympy_parser import parse_expr
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here

def ezplot(s):
    #Parse doesn't parse = sign so split
    lhs, rhs = s.replace("^","**").split("=")
    eqn_lhs = parse_expr(lhs)
    eqn_rhs = parse_expr(rhs)

    plt.ion() #interactive on 
    p = plot_implicit(eqn_lhs-eqn_rhs, line_color='C0', linewidth=1.5)
    fig, axes = p._backend.fig, p._backend.ax  # get matplotib's figure and ax
    axes[0].set_aspect('equal')
    fig.set_size_inches(5, 4)
    plt.ioff() #interactive off

 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    ezplot('x^2 + y^2 = y + 5')
    ezplot('y = x^2')
    #plt.gca().set_aspect('equal')
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv:
        pass
    else:
        plt.show()