Fix bugs. Add plot function.