from pylab import *



# broj tacaka
n = 49
lmin = 1 / (n + 1) # donja granica po J0
lmax = 1 - lmin # gornja granica po J0
nstart = 2000 # zaletanje ka ustaljenom stanju
nrecord = 200 # broj tačaka koji se zapisuje
nfunc = 1001 # broj tačaka za prikazivanje funkcije

Jm = 1 # maksimalna struja kalema, normalizujem na 1



# LaTeX lettering na slikama
rc('text', usetex = True)
rc('font', family = 'serif')
rc('font', size = 12)
rcParams['text.latex.preamble']=[r'\usepackage{amsmath}']



# odredimo m1 i m2 iz Jm, J0 i D0
def initialization(J0, D0, Jm):
    '''inicijalizacija: racunamo m1 i m2 na bazi J0, D0 i Jm'''
    m1 = (Jm - J0) / D0
    m2 = (J0 - Jm) / (1 - D0)
    return m1, m2
    
    
    
# preslikavanje: j1 u funkciji j0 i parametara
def step(j0):
    '''preslikavanje j0 u j1, parametri su m1, m2 i Jm'''
    d = (Jm - j0) / m1
    if d > 1:
        return j0 + m1 # S switching pattern
    else:
        j1 = Jm + (1 - d) * m2
        if j1 < 0:
            return 0 # SDX switching pattern
        else:
            return j1 # SD switching pattern



# generišemo vektore koji prikazuju preslikavanje, gruba sila, može i samo dve
# do četiri tačke da bi bila prikazana slika, zavisi od slučaja, optimizujte 
# ako znate i hoćete, to je extra credit problem
def mapping():
    ''' preslikavanje j_k u j_{k+1}'''
    jk1 = empty(nfunc)
    jk0 = linspace(0, Jm, nfunc)
    # ovo nije Pythonic: moglo je korišćenjem map funkcije, pedagoski razlozi
    for i in range(len(jk0)):
        jk1[i] = step(jk0[i])
    return (jk0, jk1)



# glavni program
for iJ0, J0 in enumerate(linspace(lmin, lmax, n)):
    print(iJ0 + 1, '/', n)
    close('all')
    figure(2, figsize = (6,6))
    for iD0, D0 in enumerate(linspace(lmin, lmax, n)):
        m1, m2 = initialization(J0, D0, 1)
        close(1)
        figure(1, figsize = (6,6))
        xlim(-0.1, 1.1)
        xticks(linspace(0, 1, 11))
        ylim(-0.1, 1.1)
        yticks(linspace(0, 1, 11))
        xlabel(r'$j_k$')
        ylabel(r'$j_{k + 1}$')
        title(r'preslikavanje za $D_0 = {:6.4f}$, $J_0 = {:6.4f}$'.format(D0, J0))
        plot([0, Jm], [0, Jm], 'k', linewidth = 0.5)
        figdata = mapping()
        plot(figdata[0], figdata[1], 'b')
        j0 = 1
        for i in range(nstart):
            j1 = step(j0)
            j0 = j1
        js = empty(nrecord)
        for i in range(nrecord):
            js[i] = step(j0)
            figure(1)
            plot(j0, js[i], 'r.')
            j0 = js[i]
        savefig('mappings/mapping_{:04d}_{:04d}.png'.format(iJ0, iD0), bbox_inches = 'tight')
        figure(2)
        plot(D0 * ones(nrecord), js, 'b.')
    figure(2)
    xlim(-0.1, 1.1)
    xticks(linspace(0, 1, 11))
    ylim(-0.1, 1.1)
    yticks(linspace(0, 1, 11))
    xlabel(r'$D_0$')
    ylabel(r'$j_k$')
    title(r'bifurkacioni dijagram po $D_0$, $J_0 = {:6.4f}$'.format(J0))
    savefig('bifurcations/bifurcation_diagram_{:04d}.png'.format(iJ0), bbox_inches = 'tight')