import numpy as np
from numpy import arange, array
import matplotlib as mp
import matplotlib.pyplot as pyplot

minus = u"\u2212"
blue = [0.2,0.4,0.6,1.0]
green = [0,0.46,0,0.8]
magenta = [0.6,0,0.4,0.8]
axes_color = "#505050"

def label_to_str(x):
  if type(x) is str:
    return x
  else:
    return minus+str(abs(x)) if x<0 else str(x)

def arrow(ax,x,y,dx,dy,lw=2,w=1,h=1):
  ax.arrow(x,y,dx,dy, shape='full', lw=lw, length_includes_head=True,
    head_width=0.22*w, head_length=0.5*h, color=axes_color,
    joinstyle='bevel', clip_on=False)

def plot(f,X=[-5,5],Y=[-5,5],name="plot",inc=1,incy=None,aspect=1,xticks=[],yticks=[],
  xsize=None, xshift=[0,0],yshift=[0,0], out=False, ticks_off=False,
  con=0.5,dx=0.01,lw_line=2.6,color=blue,xlabels=False,ylabels=False):

  lw_grid = 1.8
  fontsize = 14

  if xsize is None: xsize = abs(X[1]-X[0])
  if incy is None: incy = inc
  style = {
    "axes.linewidth": lw_grid,
    "grid.linewidth": lw_grid,
    "grid.linestyle": "solid",
    "grid.color": "#e0e0da",
    "lines.linewidth": 2,
    # "lines.markersize": 10,
    "xtick.labelsize": 12,
    "ytick.labelsize": 12,
    "figure.figsize": (4*xsize/8,20)
  }
  mp.rcParams.update(style)

  fig = pyplot.figure()
  ax = fig.add_subplot(1,1,1)

  ax.spines['bottom'].set_color(axes_color)
  ax.spines['top'].set_color(axes_color) 
  ax.spines['right'].set_color(axes_color)
  ax.spines['left'].set_color(axes_color)
  ax.spines['bottom'].set_zorder(2)
  ax.spines['top'].set_zorder(2)
  ax.spines['right'].set_zorder(2)
  ax.spines['left'].set_zorder(2)
  ax.spines['top'].set_color('none')
  ax.spines['right'].set_color('none')
  ax.xaxis.set_ticks_position('bottom')
  ax.yaxis.set_ticks_position('left')

  if not out:
    ax.spines['bottom'].set_position('zero')
    ax.spines['left'].set_position('zero')

  ax.xaxis.grid()
  ax.yaxis.grid()
  ax.set_axisbelow(True)
  ax.xaxis.set_tick_params(width=lw_grid, length=10,
    pad=8,color=axes_color,direction='inout')
  ax.yaxis.set_tick_params(width=lw_grid, length=10,
    pad=8,color=axes_color,direction='inout')
  ax.set_aspect(aspect)

  ax.axis(X+Y)
  if out:
    ax.set_xticks(list(arange(X[0],X[1]+inc,inc)))
    ax.set_yticks(list(arange(Y[0],Y[1]+incy,incy)))
  else:
    ax.set_xticks(list(arange(X[0],0,inc))+xticks+list(arange(inc,X[1]+inc,inc)))
    ax.set_yticks(list(arange(Y[0],0,incy))+yticks+list(arange(incy,Y[1]+incy,incy)))

  if xlabels is not False:
    ax.set_xticklabels([label_to_str(x) for x in xlabels])
  if ylabels is not False:
    ax.set_yticklabels([label_to_str(y) for y in ylabels])

  a = f if type(f) is list else [f]
  for f in a:
    x = arange(X[0], X[1], dx)
    y = array(map(f,x))
  
    pos = np.where(np.abs(np.diff(y)) >= con)[0]+1
    x = np.insert(x, pos, np.nan)
    y = np.insert(y, pos, np.nan)

    ax.plot(x,y, color=color, zorder=3,linewidth=lw_line)

  xposx = X[1]+(xshift[0]-0.5)*inc; xposy = (xshift[1]-0.68)*incy
  yposx = (yshift[0]+0.28)*inc; yposy = Y[1]+(yshift[1]-0.55)*incy
  if out:
    xposy = Y[0]+(xshift[1]-0.5)*incy
    yposx = X[0]+(yshift[0]-0.5)*inc
  else:
    arrow(ax,X[1]-inc,0,1,0,lw=lw_grid,w=incy,h=incy*aspect)
    arrow(ax,0,Y[1]-incy,0,incy,lw=lw_grid,w=incy*aspect,h=incy)

  ax.text(xposx,xposy,"x",fontsize=fontsize,style="italic",color="#202020")
  ax.text(yposx,yposy,"y",fontsize=fontsize,style="italic",color="#202020")

  # pyplot.savefig(name+".png",bbox_inches='tight')
  pyplot.savefig(name+".svg",bbox_inches='tight')


plot(lambda x: x*x-4*x+2, [-2,5], [-4,4], "plot",
  xlabels=["",-1,1,2,3,4],
  ylabels=["",-3,-2,-1,1,2,3])