from matplotlib import pyplot as plt
from numpy import arange, zeros, linspace, sign, logspace, log10

from idmlin import IDMLin


t = 1.5
amin, amax, acount = log10(0.3), log10(3), 20
bmin, bmax, bcount = log10(1), log10(4), 20

arng = logspace(amin, amax, acount)
brng = logspace(bmin, bmax, bcount)
s0 = 2
v0 = 30
re = 0.1

ns = (200,)

f, axarr = plt.subplots(2, len(ns)+1, sharex=True, sharey=True)
#f, axarr = plt.subplots(2, 2, sharex=True, sharey=True)

allstab = zeros((len(arng), len(brng)))
alloss = zeros((len(arng), len(brng)))

for i, n in enumerate(ns):
    stab = zeros((len(arng), len(brng)))
    oss = zeros((len(arng), len(brng)))

    def isstable(a, b):
            idm = IDMLin(t=t, a=a, b=b, s0=s0)
            idm.go(v0, re)
            sys = idm.sys()
            ils = idm.isloopstable(n)
            imr = idm.maxresponse()
            print a, b, ils, imr

            return 1 if ils else 0, imr
            #return 1 if ils else 0, 1 if imr<0 else 0

    for ai, a in enumerate(arng):
        for bi, b in enumerate(brng):
            stab[ai, bi], oss[ai, bi] = isstable(a,b)


    axarr[0, i*1].imshow(stab, cmap=plt.cm.Reds, 
                    interpolation='none', 
                    extent=[bmin, bmax, amax, amin], 
                    aspect='auto',
                    )
    axarr[1, i*1].imshow(oss, cmap=plt.cm.Reds, 
                    interpolation='none', 
                    extent=[bmin, bmax, amax, amin], 
                    aspect='auto',
                    )
    allstab += stab
    alloss += oss

axarr[0, -1].imshow(allstab, cmap=plt.cm.Reds, 
                interpolation='none', 
                extent=[bmin, bmax, amax, amin], 
                aspect='auto',
                )
axarr[1, -1].imshow(alloss, cmap=plt.cm.Reds, 
                interpolation='none', 
                extent=[bmin, bmax, amax, amin], 
                aspect='auto',
                )
allstab += stab
alloss += oss

plt.show()