load(draw)$

load(dynamics)$

normalize(vec) := block([d:sqrt(vec.vec)], if d#0 then vec/d else [0,0])$

get_trajectory_time(expr, xrange, yrange, point) := block(
  [t%, xvar:first(xrange), yvar:first(yrange), trajectory],
  trajectory : rk(expr,
                  [xvar, yvar],
                  point,
                  [t%, 0, 10, 0.1]),
  points(map(rest, trajectory)))$

get_trajectory(expr, xrange, yrange, point) := block(
  [trajectory_pl, trajectory_mi, xstep, yvar],
  xstep : (xrange[3]-xrange[2])/50.0,
  yvar : yrange[1],
  trajectory_pl : rk(expr, [yvar], [point[2]], [xrange[1], point[1], xrange[3], xstep]),
  trajectory_mi : rk(expr, [yvar], [point[2]], [xrange[1], point[1], xrange[2], -xstep]),
  apply('points, [append(reverse(trajectory_mi), trajectory_pl)]))$

get_field_lines_time(expr, xrange, yrange) := block(
  [points_df : [], xvar, xstart, xend, xstep, yvar, ystart, yend, ystep],
  [xvar, xstart, xend] : xrange,
  [yvar, ystart, yend] : yrange,
  xstep : (xend-xstart)/20,
  ystep : (yend-ystart)/20,
  for x%:xstart thru xend step xstep do (
    for y%:ystart thru yend step ystep do (
      points_df : cons(
        points([[x%, y%], [x%, y%] + min(xstep, ystep)/2*normalize(subst([xvar=x%, yvar=y%], expr))]),
        points_df))),
  points_df)$

get_field_lines(expr, xrange, yrange) := block(
  [points_df : [], xvar, xstart, xend, xstep, yvar, ystart, yend, ystep],
  [xvar, xstart, xend] : xrange,
  [yvar, ystart, yend] : yrange,
  xstep : (xend-xstart)/20,
  ystep : (yend-ystart)/20,
  for x%:xstart thru xend step xstep do (
    for y%:ystart thru yend step ystep do (
      points_df : cons(
        points([[x%, y%], [x%, y%] + min(xstep, ystep)/2*normalize([1, subst([xvar=x%, yvar=y%], expr)])]),
        points_df))),
  points_df)$

draw_df(expr, xvar, xstart, xend, yvar, ystart, yend, [options]) := block(
  [points_df, trajectory:[], xrange, yrange, comm,
   colors:['blue, 'red, 'green, 'yellow]],

  xrange : [xvar, xstart, xend],
  yrange : [yvar, ystart, yend],

  if listp(expr) then
    points_df : get_field_lines_time(expr, xrange, yrange)
  else
    points_df : get_field_lines(expr, xrange, yrange),

  for i:1 thru length(options) do (
    trajectory : append(['point_size=1, 'point_type='multiply, 'color='black, points([options[i]])],
      trajectory),
    if listp(expr) then
      trajectory : cons(get_trajectory_time(expr, xrange, yrange, options[i]), trajectory)
    else
      trajectory : cons(get_trajectory(expr, xrange, yrange, options[i]), trajectory),
    trajectory : append(['point_size=0, 'color=colors[mod(i-1, length(colors))+1]], trajectory)),
    append(['points_joined=true, 'head_length=0.12, 'point_size=0],
      points_df,
      ['color = blue, 'line_width = 2],
      trajectory,
      ['xrange = [xstart, xend]],
      ['yrange = [ystart, yend]],
      ['xaxis = true, 'yaxis = true])
  )$

drawdf_tr(expr) := apply(draw_df, args(expr))$

put(df, drawdf_tr, draw2d_transform)$