import numpy
from matplotlib import pyplot

BFDcoeffs = { 1: {'alpha': [1, -1], 'beta': 1},
              2: { 'alpha': [3, -4, 1], 'beta': 2 },
              3: { 'alpha': [11, -18, 9, -2], 'beta': 6 },
              4: { 'alpha': [25, -48, 36, -16, 3], 'beta': 12 },
              5: { 'alpha': [137, -300, 300, -200, 75, -12], 'beta': 60 },
              6: { 'alpha': [147, -360, 450, -400, 225, -72], 'beta': 60 } }

plotWindow = { 1: { 'realPart': [-2, 3], 'imagPart': [-2, 2] },
               2: { 'realPart': [-2, 5], 'imagPart': [-3, 3] },
               3: { 'realPart': [-4, 8], 'imagPart': [-5, 5] },
               4: { 'realPart': [-4, 14], 'imagPart': [-8, 8] },
               5: { 'realPart': [-10, 25], 'imagPart': [-15, 15] },
               6: { 'realPart': [-20, 40], 'imagPart': [-30, 30] } }

# Returns > 1 if argument is not in region of absolute stability
def stabilityFunction(hTimesLambda, s):
    stabPolyCoeffs = list(BFDcoeffs[s]['alpha'])
    stabPolyCoeffs[0] -= hTimesLambda * BFDcoeffs[s]['beta']
    return max(abs(numpy.roots(stabPolyCoeffs)))

# Main program
for s in range(1,7):
    x = numpy.linspace(*plotWindow[s]['realPart'], num=400)
    y = numpy.linspace(*plotWindow[s]['imagPart'], num=400)
    [X,Y] = numpy.meshgrid(x,y)
    Z = numpy.zeros(X.shape)
    for m in range(X.shape[0]):
        for n in range(X.shape[1]):
            Z[m,n] = stabilityFunction(X[m,n] + 1j * Y[m,n], s)
    pyplot.contour(X, Y, Z, [1], colors='k')
    pyplot.contourf(X, Y, Z, [0,1], colors=[[1, 0.5, 0.8]])
    pyplot.plot(plotWindow[s]['realPart'], [0, 0], 'k--')
    pyplot.plot([0, 0], plotWindow[s]['imagPart'], 'k--')
    pyplot.gca().tick_params(labelsize = 20)
    pyplot.savefig('Stability_region_for_BDF%d.svg' % s)
    pyplot.clf()