// PoleZeroPlot.java

import java.awt.*;

public class PoleZeroPlot extends Canvas {

  float[] pReal, pImag, z;
  int order = 0;
  int gridIntervals = 10;
  int zSize = 4;  // zero symbol size
  int pSize = 3;  // pole symbol size
  float scale;
  Color plotColor = Color.blue;
  Color axisColor = Color.darkGray;
  Color circColor = Color.red;
  Color gridColor = Color.darkGray;
  Color bgColor   = Color.lightGray;
  int vertSpace = 20;
  int horzSpace = 20;

  public PoleZeroPlot() {
  }

  public void setPlotColor(Color c) {
    if (c != null) plotColor = c;
  }

  public Color getPlotColor() {
    return plotColor;
  }

  public void setAxisColor(Color c) {
    if (c != null) axisColor = c;
  }

  public Color getAxisColor() {
    return axisColor;
  }

  public void setGridColor(Color c) {
    if (c != null) gridColor = c;
  }

  public Color getGridColor() {
    return gridColor;
  }

  public void setCircColor(Color c) {
    if (c != null) circColor = c;
  }

  public Color getCircColor() {
    return circColor;
  }

  public void setBgColor(Color c) {
    if (c != null) bgColor = c;
  }

  public Color getBgColor() {
    return bgColor;
  }

  public void setPolesAndZeros(float[] pr, float[] pi, float[] zr) {
    order = pr.length - 1; // number of poles/zeros = filter order
    pReal = new float[order + 1];
    pImag = new float[order + 1];
    z     = new float[order + 1];
    for (int i = 1; i <= order; i++) {
      pReal[i] = pr[i];
      pImag[i] = pi[i];
      z[i]     = zr[i];
    }
    repaint();
  }

  public void paint(Graphics g) {
    int x, y;
    int xc = getSize().width/2;
    int yc = getSize().height/2;
    int width  = getSize().width - 2*horzSpace;
    int height = getSize().height - 2*vertSpace;
    int radius = Math.min(width/2, height/2);
    int top    = yc - radius;
    int bottom = yc + radius;
    int left   = xc - radius;
    int right  = xc + radius;
    scale = 2*radius/(float)gridIntervals;
    setBackground(bgColor);
    g.setColor(gridColor);
    // grid lines
    for (int i = 0; i <= gridIntervals; i++) {
      x = left + Math.round(i*scale);
      y = top  + Math.round(i*scale);
      g.drawLine(x, top, x, bottom); // vertical grid line
      g.drawLine(left, y, right, y); // horizontal grid line
    }
    g.setColor(axisColor);
    g.drawLine(xc, top - vertSpace, xc, bottom + vertSpace); // vertical axis
    g.setFont(new Font("Sans serif", Font.BOLD, 10));
    FontMetrics fm = g.getFontMetrics();
    int h = fm.getMaxAscent();
    g.drawString("Im", xc + 4, top - vertSpace + h);
    g.drawLine(left - horzSpace, yc, right + horzSpace, yc); // horizontal axis
    int w = fm.stringWidth("Re");
    g.drawString("Re", right + horzSpace - w, yc + h + 4);
    g.setColor(circColor);
    g.drawOval(left, top, 2*radius, 2*radius); // unit circle
    if (order > 0) {
      g.setColor(plotColor);
      // plot zeros
      for (int i = 1; i <= order; i++) {
        x = xc + Math.round(radius*z[i]);
        g.drawOval(x - zSize, yc - zSize, 2*zSize, 2*zSize);
      }
      // plot poles
      for (int i = 1; i <= order; i++) {
        x = xc + Math.round(radius*pReal[i]);
        y = yc - Math.round(radius*pImag[i]);
        g.drawLine(x - pSize, y - pSize, x + pSize, y + pSize);
        g.drawLine(x - pSize, y + pSize, x + pSize, y - pSize);
      }
    }
  }
}