from IPython.display import HTML
from IPython.display import Image
from PIL import Image as ImagePIL

%pylab
%matplotlib inline

"""
Explanation: Неустойчивости
Тут будут все функции, имеющие отношение к неустойчивостям, чтобы не копировать их каждый раз.
Тесты на эти функции в соседнем ноутбуке.
End of explanation
"""


G = 4.32 #гравитационная постоянная в нужных единицах

def Qs(epicycl=None, sigma=None, star_density=None):
    '''Вычисление безразмерного параметра Тумре для звездного диска. 
    Зависит от плотности звезд, дисперсии скоростей и эпициклической частоты.
    Вычисляется для pi, а не 3.36, поскольку так в Rafikov(2001).'''
    return epicycl * sigma / (math.pi * G * star_density)


def Qg(epicycl=None, sound_vel=None, gas_density=None):
    '''Вычисление безразмерного параметра Тумре для газового диска. 
    Зависит от плотности газа и эпициклической частоты, скорости звука в газе.'''
    return epicycl * sound_vel / (math.pi * G * gas_density)

"""
Explanation: Одножидкостный критерий
Устойчиво, когда > 1:
$$Q_g = \frac{\Sigma_g^{cr}}{\Sigma_g}=\frac{\kappa c_g}{\pi G \Sigma_g}$$
$$Q_s = \frac{\Sigma_s^{cr}}{\Sigma_s}=\frac{\sigma_R}{\sigma_R^{min}}=\frac{\kappa \sigma_R}{\pi G \Sigma_s}$$
End of explanation
"""


from scipy.special import i0e, i1e

def inverse_hydro_Qeff_from_k(dimlK, Qg=None, Qs=None, s=None):
    return 2.*dimlK / Qs / (1 + dimlK**2) + 2*s*dimlK / Qg / (1 + dimlK**2 * s**2)

def inverse_kinem_Qeff_from_k(dimlK, Qg=None, Qs=None, s=None):
    return 2. / dimlK / Qs * (1 - i0e(dimlK ** 2)) + 2*s*dimlK / Qg / (1 + dimlK**2 * s**2)

"""
Explanation: Двухжидкостный критерий
Кинетическое приближение:
$$\frac{1}{Q_{\mathrm{eff}}}=\frac{2}{Q_{\mathrm{s}}}\frac{1}{\bar{k}}\left[1-e^{-\bar{k}^{2}}I_{0}(\bar{k}^{2})\right]+\frac{2}{Q_{\mathrm{g}}}s\frac{\bar{k}}{1+\bar{k}^{2}s^{2}}>1\,$$
Гидродинамическое приближение:
$$\frac{2\,\pi\, G\, k\,\Sigma_{\mathrm{s}}}{\kappa+k^{2}\sigma_{\mathrm{s}}}+\frac{2\,\pi\, G\, k\,\Sigma_{\mathrm{g}}}{\kappa+k^{2}c_{\mathrm{g}}}>1$$ или $$\frac{1}{Q_{\mathrm{eff}}}=\frac{2}{Q_{\mathrm{s}}}\frac{\bar{k}}{1+\bar{k}^{2}}+\frac{2}{Q_{\mathrm{g}}}s\frac{\bar{k}}{1+\bar{k}^{2}s^{2}}>1$$ для безразмерного волнового числа ${\displaystyle \bar{k}\equiv\frac{k\,\sigma_{\mathrm{s}}}{\kappa}},\, s=c/\sigma$
End of explanation
"""


from sympy import Symbol, solve

def findInvHydroQeffSympy(Qs, Qg, s):
    '''Решаем уравнение deriv()=0 чтобы найти максимум функции в гидродинамическом приближении.'''
    k = Symbol('k') #solve for complex because it may returns roots as 1.03957287978471 + 0.e-20*I 
    foo = 2./Qs*k/(1+k**2) + 2/Qg*s*k/(1+k**2 * s**2)
    foo2 = 2./Qs * (1-k)*(1+k * s**2)**2 + 2/Qg*s*(1-k*s**2)*(1+k)**2
    roots = solve(foo2.simplify(), k)
    roots = [np.sqrt(float(abs(re(r)))) for r in roots]
    _tmp = [foo.evalf(subs={k:r}) for r in roots]
    max_val = max(_tmp)
    return (roots[_tmp.index(max_val)], max_val)

def findInvHydroQeffBrute(Qs, Qg, s, krange):
    '''Находим максимум функции в гидродинамическом приближении перебором по сетке.'''
    _tmp = [inverse_hydro_Qeff_from_k(l, Qg=Qg, Qs=Qs, s=s) for l in krange]
    max_val = max(_tmp)
    root_for_max = krange[_tmp.index(max_val)]
    if abs(root_for_max-krange[-1]) < 0.5:
        print 'WARNING! For Qs={} Qg={} s={} root of max near the max of k-range'.format(Qs, Qg, s)
    return (root_for_max, max_val)

from scipy.optimize import brentq

def findInvHydroQeffBrentq(Qs, Qg, s, krange):
    '''Решение уравнения deriv(9) = 0 для нахождения максимума исходной функции. Запускается brentq на исходной сетке,
    в случае если на концах сетки разные знаки функции (промежуток содержит корень),
    затем выбираются лучшие корни, после чего ищется, какой их них дает максимум. Возвращается только этот корень.'''
    grid = krange
    args = [Qs, Qg, s]
    signs = [derivTwoFluidHydroQeff(x, *args) for x in grid]
    signs = map(lambda x: x / abs(x), signs)
    roots = []
    for i in range(0, signs.__len__() - 1):
        if signs[i] * signs[i + 1] < 0:
            roots.append(brentq(lambda x: derivTwoFluidHydroQeff(x, *args), grid[i], grid[i + 1]))
    original = [inverse_hydro_Qeff_from_k(l, Qg=Qg, Qs=Qs, s=s) for l in roots]
    root_for_max = roots[original.index(max(original))]
    if abs(root_for_max-krange[-1]) < 0.5:
        print 'WARNING! For Qs={} Qg={} s={} root of max near the max of k-range'.format(Qs, Qg, s)
    return (root_for_max, max(original))

def derivTwoFluidHydroQeff(dimlK, Qs, Qg, s):
    '''Производная по \bar{k} от левой части (9) для того, чтобы найти максимум.'''
    part1 = (1 - dimlK ** 2) / (1 + dimlK ** 2) ** 2
    part3 = (1 - (dimlK * s) ** 2) / (1 + (dimlK * s) ** 2) ** 2
    return (2 * part1 / Qs) + (2 * s * part3 / Qg)

"""
Explanation: Нахождение максимума:
Найти максимум функции в гидродинамическом приближении вообще просто - это многочлен
$$\frac{2}{Q_{\mathrm{s}}}\frac{\bar{k}}{1+\bar{k}^{2}}+\frac{2}{Q_{\mathrm{g}}}s\frac{\bar{k}}{1+\bar{k}^{2}s^{2}}>1$$
и у него можно найти максимум методами Sympy, взяв производную:
End of explanation
"""


def findInvKinemQeffBrute(Qs, Qg, s, krange):
    '''Находим максимум функции в кинематическом приближении перебором по сетке.'''
    _tmp = [inverse_kinem_Qeff_from_k(l, Qg=Qg, Qs=Qs, s=s) for l in krange]
    max_val = max(_tmp)
    root_for_max = krange[_tmp.index(max_val)]
    if abs(root_for_max-krange[-1]) < 0.5:
        print 'WARNING! For Qs={} Qg={} s={} root of max near the max of k-range'.format(Qs, Qg, s)
    return (root_for_max, max_val)

def findInvKinemQeffBrentq(Qs, Qg, s, krange):
    '''Решение уравнения deriv(13) = 0 для нахождения максимума исходной функции. Запускается brentq на исходной сетке,
    в случае если на концах сетки разные знаки функции (промежуток содержит корень),
    затем выбираются лучшие корни, после чего ищется, какой их них дает максимум. Возвращается только этот корень.'''
    grid = krange
    args = [Qs, Qg, s]
    signs = [derivTwoFluidKinemQeff(x, *args) for x in grid]
    signs = map(lambda x: x / abs(x), signs)
    roots = []
    for i in range(0, signs.__len__() - 1):
        if signs[i] * signs[i + 1] < 0:
            roots.append(brentq(lambda x: derivTwoFluidKinemQeff(x, *args), grid[i], grid[i + 1]))
    original = [inverse_kinem_Qeff_from_k(l, Qg=Qg, Qs=Qs, s=s) for l in roots]
    root_for_max = roots[original.index(max(original))]
    if abs(root_for_max-krange[-1]) < 0.5:
        print 'WARNING! For Qs={} Qg={} s={} root of max near the max of k-range'.format(Qs, Qg, s)
    return (root_for_max, max(original))


def derivTwoFluidKinemQeff(dimlK, Qs, Qg, s):
    '''Производная по \bar{k} от левой части (13) для того, чтобы найти максимум. Коррекция за ассимптотику производится
    с помощью встроенных функций бесселя, нормированных на exp.'''
    part1 = (1 - i0e(dimlK ** 2)) / (-dimlK ** 2)
    part2 = (2 * dimlK * i0e(dimlK ** 2) - 2 * dimlK * i1e(dimlK ** 2)) / dimlK
    part3 = (1 - (dimlK * s) ** 2) / (1 + (dimlK * s) ** 2) ** 2
    return 2 * (part1 + part2) / Qs + 2 * s * part3 / Qg

def calc_Qeffs_(Qss=None, Qgs=None, s_params=None, verbose=False):
    '''считает сразу все Qeff в кинематическом'''
    invQeff = []
    for Qs, Qg, s in zip(Qss, Qgs, s_params):
        qeff = findInvKinemQeffBrentq(Qs, Qg, s, np.arange(0.01, 60000., 1.))
        if verbose:
            print 'Qs = {:2.2f}; Qg = {:2.2f}; s = {:2.2f}; Qeff = {:2.2f}'.format(Qs, Qg, s, 1./qeff[1])
        invQeff.append(qeff[1])
    return invQeff

def calc_Qeffs(r_g_dens=None, gas_dens=None, epicycl=None, 
               sound_vel=None, star_density=None, sigma=None, verbose=False):
    '''считаем модельное Qeff в кинематическом'''
    Qgs = []
    Qss = []
    s_params = []
    for r, gd, sd in zip(r_g_dens, gas_dens, star_density):
        Qgs.append(Qg(epicycl=epicycl(r), sound_vel=sound_vel, gas_density=gd))
        Qss.append(Qs(epicycl=epicycl(r), sigma=sigma(r), star_density=sd))
        s_params.append(sound_vel/sigma(r))
    return calc_Qeffs_(Qss=Qss, Qgs=Qgs, s_params=s_params, verbose=verbose)

def plot_k_dependency(Qs=None, Qg=None, s=None, krange=None, ax=None, label=None, color=None):
    '''рисуется зависимость между волновыми числами и двухжидкостной неустойчивостью, показан максимум'''
    TFcriteria = []
    _tmp = [inverse_kinem_Qeff_from_k(dimlK, Qg=Qg, Qs=Qs, s=s) for dimlK in krange]
    root_for_max, max_val = findInvKinemQeffBrentq(Qs, Qg, s, krange)
    ax.plot(krange, _tmp, '-', label=label, color=color)
    ax.plot(root_for_max, max_val, 'o', color=color)

def plot_k_dependencies(r_g_dens=None, gas_dens=None, epicycl=None, 
               sound_vel=None, star_density=None, sigma=None, krange=None, show=False):
    '''рисуем много зависимостей сразу'''
    Qgs, Qss, s_params = [], [], []
    maxk = 30.
    fig = plt.figure(figsize=[16,8])
    ax = plt.gca()
    colors = cm.rainbow(np.linspace(0, 1, len(r_g_dens)))
    for r, gd, sd, color in zip(r_g_dens, gas_dens, star_density, colors):
        Qgs.append(Qg(epicycl=epicycl(r), sound_vel=sound_vel, gas_density=gd))
        Qss.append(Qs(epicycl=epicycl(r), sigma=sigma(r), star_density=sd))
        s_params.append(sound_vel/sigma(r))
        if show:
            print 'r={:7.3f} Qg={:7.3f} Qs={:7.3f} Qg^-1={:7.3f} Qs^-1={:7.3f} s={:7.3f}'.format(r, Qgs[-1], Qss[-1], 1./Qgs[-1], 1./Qss[-1], s_params[-1])
        plot_k_dependency(Qs=Qss[-1], Qg=Qgs[-1], s=s_params[-1], krange=krange, ax=ax, label=str(r), color=color)
        maxk = max(maxk, findInvKinemQeffBrentq(Qss[-1], Qgs[-1], s_params[-1], krange)[0]) #not optimal
    plt.legend()
    plt.xlim(0, maxk+100.)

"""
Explanation: Теперь кинематическое приближение:
$$\frac{2}{Q_{\mathrm{s}}}\frac{1}{\bar{k}}\left[1-e^{-\bar{k}^{2}}I_{0}(\bar{k}^{2})\right]+\frac{2}{Q_{\mathrm{g}}}s\frac{\bar{k}}{1+\bar{k}^{2}s^{2}}>1\,$$
Тут сложнее, честно уже не решить. остается два способа - брутфорсом и brentq, производная известна.
End of explanation
"""


def plot_WKB_dependency(Qs=None, Qg=None, s=None, krange=None, ax=None, label=None, color=None, r=None, scale=None, epicycl=None, sound_vel=None):
    '''рисуется зависимость между (k x r) и двухжидкостной неустойчивостью, показан максимум (см. уравнение выше), чтобы проверить справедливость WKB'''
    TFcriteria = []
    sigma = sound_vel/s
    _tmp = [inverse_kinem_Qeff_from_k(dimlK, Qg=Qg, Qs=Qs, s=s) for dimlK in krange]
    root_for_max, max_val = findInvKinemQeffBrentq(Qs, Qg, s, krange)
    factor = epicycl*r*scale/sigma
    ax.plot(np.array(krange)*factor, _tmp, '-', label=label, color=color)
    ax.plot(root_for_max*factor, max_val, 'o', color=color)
    
def plot_WKB_dependencies(r_g_dens=None, gas_dens=None, epicycl=None, 
               sound_vel=None, star_density=None, sigma=None, krange=None, scale=None):
    '''рисуем много зависимостей сразу для проверки WKB'''
    Qgs, Qss, s_params = [], [], []
    maxk = 30.
    fig = plt.figure(figsize=[16,8])
    ax = plt.gca()
    colors = cm.rainbow(np.linspace(0, 1, len(r_g_dens)))
    for r, gd, sd, color in zip(r_g_dens, gas_dens, star_density, colors):
        Qgs.append(Qg(epicycl=epicycl(r), sound_vel=sound_vel, gas_density=gd))
        Qss.append(Qs(epicycl=epicycl(r), sigma=sigma(r), star_density=sd))
        s_params.append(sound_vel/sigma(r))
        plot_WKB_dependency(Qs=Qss[-1], Qg=Qgs[-1], s=s_params[-1], krange=krange, ax=ax, label=str(r), 
                          color=color, r=r, scale=scale, epicycl=epicycl(r), sound_vel=sound_vel)
        maxk = max(maxk, findInvKinemQeffBrentq(Qss[-1], Qgs[-1], s_params[-1], krange)[0]) #not optimal
    plt.legend()
    plt.xlim(0, maxk+100.)

def get_invQeff_from_data(gas_data=None, epicycl=None, gas_approx=None, sound_vel=None, scale=None, sigma=None, star_density=None, verbose=False):
    '''рассчитывает из наблюдательных данных сразу много значений Qeff^-1 и возвращает набор (Qg, Qs, Qeff)'''
    Qgs = []
    Qss = []
    invQeff = []
    for ind, (r, gd) in enumerate(gas_data):
        if type(sound_vel) == tuple or type(sound_vel) == list: #учет случая разных скоростей звука 
            s_vel = sound_vel[ind]
        else:
            s_vel = sound_vel
        Qgs.append(Qg(epicycl=epicycl(gas_approx, r, scale), sound_vel=s_vel, gas_density=gd))
        Qss.append(Qs(epicycl=epicycl(gas_approx, r, scale), sigma=sigma(r), star_density=star_density(r)))
        qeff = findInvKinemQeffBrentq(Qss[-1], Qgs[-1], s_vel/sigma(r), np.arange(0.01, 60000., 1.))
        if verbose:
            print 'r = {:2.2f}; gas_d = {:2.2f}; epicycl = {:2.2f}; sig = {:2.2f}; star_d = {:2.2f}'.format(r, gd, epicycl(gas_approx, r, scale), 
                                                                                                          sigma(r), star_density(r))
            print '\tQs = {:2.2f}; Qg = {:2.2f}; Qeff = {:2.2f}'.format(Qss[-1], Qgs[-1], 1./qeff[1])
        invQeff.append(qeff[1])
    return zip(map(lambda l: 1./l, Qgs), map(lambda l: 1./l, Qss), invQeff)

"""
Explanation: У Rafikov написано следующее
"For the local analysis we employ the WKB approximation (or tight-winding approximation) which requires that kr ≫ 1 and allows us to neglect terms proportional to 1/r compared to the terms proportional to k."
Т.е. применялось ВКБ и действительно придется смотреть, какие там длины волн в максимуме
$$kr \gg 1$$
$$k=\frac{2\pi}{\lambda}$$
$${\displaystyle \bar{k}\equiv\frac{k\,\sigma_{\mathrm{s}}}{\kappa}}$$
$$k\times r = \frac{\bar{k}\varkappa}{\sigma}\times(r\times scale)$$
End of explanation
"""


