snippet_mars.exosphere_demo

import numpy as np
import matplotlib.pyplot as plt
from irfpy.mars.exosphere import _EF_const, ZetaChamberlain1963



def main():
    '''
    Plotting for verification, comaprison and other purposes.

    :return:
    '''
    import matplotlib.pyplot as plt
    R_m = _EF_const.R_m

    currlam_c_tab = np.array([3, 7.5])

    fig1 = plt.figure(1, figsize=[8, 6])
    shen_tab_alpha = np.linspace(0, 1, 11)
    shen_tab_rho_a = np.array([0, 0.057, 0.357, 1.102, 2.673, 5.654, 10.951, 19.731, 34.26, 57.77, 98.89]) / 100

    anaalt = np.logspace(np.log10(200), np.log10(50 * R_m / 1000), 100)

    for il, lam_c in enumerate(currlam_c_tab):
        curralt, Tc, zeta_bal, zeta_esc = ZetaChamberlain1963.zeta_chamberlain1963_tab(lam_c)
        zeta_shen_bal, rho_t, rho_s, alpha = ZetaChamberlain1963.zeta_bal_shen1963(curralt * 1000 + R_m, Tc)

        lam, lamc, zb_chan, zs_chan, ze_chan = ZetaChamberlain1963.zeta_chamberlain1963(anaalt * 1000 + R_m, Tc)

        ax1 = fig1.add_subplot(111)
        p = ax1.plot(zeta_bal, curralt, '--')
        ax1.plot(zeta_esc, curralt, ':', c=p[0].get_color())
        ax1.plot(zeta_esc + zeta_bal, curralt, '-', label='C63 tab, Tc={:5.4g}K'.format(Tc), linewidth=2,
                 c=p[0].get_color())

        p1 = ax1.plot(zb_chan, anaalt, '--')
        ax1.plot(ze_chan, anaalt, ':', c=p1[0].get_color())
        ax1.plot(ze_chan + zb_chan, anaalt, '-', label='C63 ana, Tc={:5.4g}K'.format(Tc), c=p1[0].get_color())
        # ax1.plot(zs_chan, anaalt, '.', label='C63 ana sat, Tc={:5.4g}K'.format(Tc), c=p1[0].get_color())

        # p2 = ax1.plot(zeta_shen_bal, curralt, label='S63 bal, Tc={:5.4} K'.format(Tc))

        # ax1.plot(rho_t, curralt, '--', label='S63 rho_t, Tc={:5.4} K'.format(Tc), c = p2[0].get_color())
        # ax1.scatter(shen_tab_rho_a, ((R_m + 200e3)/shen_tab_alpha - R_m)/1000, facecolor=p2[0].get_color())

        # fig2 = plt.figure(2, figsize = [8, 6])
        # ax = fig2.add_subplot(111)
        # ax.plot(alpha, rho_t)
        # ax.set_xlim(0,1)
        # ax.set_ylim(1e-4,1)
        # ax.set_yscale('log')
        # ax.grid()

    ax1.set_title(
        'Chamberlain 1963 (C63) tabulated values for $\zeta_{bal}$ and $\zeta_{esc}$ and \n Shen 1963 (S63) analytical forms for $\zeta_{bal}$\nLambda_c translated to exobase temperature')
    ax1.set_xlabel('$\zeta_{bal}:$---  $\zeta_{esc}:\cdots$  $\zeta_{tot}:-$', fontsize=16)
    ax1.set_xscale('log')
    ax1.set_yscale('log')
    ax1.set_ylabel('Altitude (km)')
    ax1.set_ylim(200, 7e5)
    ax1.set_xlim(1e-4, 1)
    ax1.grid()
    ax1.legend()
    plt.show()

if __name__ == "__main__":
    main()