'''
    Exosphere simulation code as described in Kegerreis et al. (2016) and
    the README.pdf documentation.

    Publically available at icc.dur.ac.uk/index.php?content=Research/Topics/O13

    Example plotting functions.
'''

from    __future__ import division
import  numpy as np
import  matplotlib.pyplot as plt
from    mpl_toolkits.axes_grid1 import make_axes_locatable
from    matplotlib.colors import LinearSegmentedColormap
import  matplotlib.gridspec as gridspec
import  matplotlib.ticker as mticker
import  sys
import  latexify
latexify.latexify()

# Workaround for the weird latex error
from    matplotlib import rc
rc('text', usetex=False)


# //////////////////////////////////////////////////////////////////////////// #
#                               0. Constants                                   #
# //////////////////////////////////////////////////////////////////////////// #


# Standard
pi     = 3.14159265358979323846
conv   = pi/180            # Degrees to radians (rad/deg)
k      = 1.3806485e-23     # Boltzmann constant (J K^-1)
G      = 6.67384e-11       # Gravitational const. (m^3 kg^-1 s^-2)
h      = 6.626176e-34      # Planck's constant (J s)
amu    = 1.66053892e-27    # Atomic mass unit (kg)
year   = 365.25*24*60*60   # Earth year (s) (= 12.37 lunar days)
# Moon
R      = 1738e3            # Radius of the Moon (m)
M      = 7.34767e22        # Mass of the Moon (kg)
day    = 2551443           # Lunar Synodic day (s)
g_0    = G*M/(R*R)         # Surface gravity = 1.6234 (m s^-2)
v_esc  = np.sqrt(2*G*M/R)  # Escape velocity = 2375.49 (m s^-1)


dir_cur = "/cosma/home/jkeger/exosphere/results/"
dir_bak = "/cosma/home/jkeger/exosphere/results_backup/"


# //////////////////////////////////////////////////////////////////////////// #
#                               1. Functions                                   #
# //////////////////////////////////////////////////////////////////////////// #


# =========================
# Useful Functions
# =========================
# Make new colour map, removing the default dark ends of jet:
c_start     = 0.15
c_end       = 0.95
cmap        = plt.get_cmap('jet')
new_colours = cmap(np.linspace(c_start, c_end, cmap.N))
new_cmap    = LinearSegmentedColormap.from_list('Upper Half', new_colours)
# ==========


def log_modulus_transform(x):
    '''Return the log-modulus transformation of x, preserves 0 and -ve values'''
    return  np.sign(x) * np.log10(abs(x) + 1)
# ==========


def local_time_labels(y_height):
    '''Add vertical lines and text labels to a plot for sunrise and sunset'''
    # Mark sunset/rise
    plt.axvline(90, ls='--', lw=1.3, c='k', alpha=0.8)
    plt.axvline(270, ls='--', lw=1.3, c='k', alpha=0.8)

    fsize   = 16
    colour  = '0.4'

    plt.text(-14, y_height, 'noon', fontsize=fsize, color=colour)
    plt.text(69, y_height, 'sunset', fontsize=fsize, color=colour)
    plt.text(149, y_height, 'midnight', fontsize=fsize, color=colour)
    plt.text(245, y_height, 'sunrise', fontsize=fsize, color=colour)
    plt.text(342, y_height, 'noon', fontsize=fsize, color=colour)
# ==========


# =========================
# Main Functions
# =========================
def density_2D():               # 1
    '''Plot output from the record_dist = "density" in 2D'''
#    global  dir_cur
    # Use a results_backup folder instead of results
    dir_cur     = dir_bak + "123456/"

    A2_density  = np.array(np.loadtxt(dir_cur + "density.txt"))

    A2_density  /= np.nanmax(A2_density)

    # Flip so theta=0 at the top
    A2_density  = np.flipud(A2_density)

    num_theta   = len(A2_density)
    num_phi     = len(A2_density[0])

    A1_theta    = np.arange(0, 180+1, 180/num_theta)
    A1_phi      = np.arange(0, 360+1, 360/num_phi)

    # Scale bin contents by bin area
    for i_theta in range(int(num_theta/2)):
        bin_area    = R**2 * (np.cos(A1_theta[i_theta]*conv) -
                              np.cos(A1_theta[i_theta+1]*conv)) * conv

        A2_density[i_theta,:]       /= bin_area
        A2_density[-(i_theta+1),:]  /= bin_area

    # Convert simulation to argon particles (by matching equator sunrise)
    A2_density  *= 2e4 / A2_density[int(num_theta/2), int(3/4*num_phi)]

    # Log scale
    A2_density[A2_density==0]   = np.nan
    A2_density                  = np.log10(A2_density)



    plt.figure(figsize=(10,5))
    ax  = plt.gca()

    plt.ylabel(r'Colatitude ($ ^\circ$)')
    plt.xlabel(r'Local Time of Day ($^\circ$)')
    plt.xlim(0, 360)
    plt.ylim(0, 180)

    # weird stuff for aligning color bar with image
    im  = ax.imshow(A2_density, interpolation='nearest', cmap=new_cmap,
                    extent=(A1_phi.min(), A1_phi.max(), A1_theta.min(),
                            A1_theta.max()))

    divider = make_axes_locatable(ax)
    cax     = divider.append_axes("right", size="5%", pad=0.1)
    cbar    = plt.colorbar(im, cax=cax)
    c_label = r'Density at the Surface (cm$^{-3}$)'
    c_label = r'log$_{10}$( '+c_label+' )'
    cbar.set_label(c_label)
    # Set colour bar ticks
    if not True:
        cbar.set_ticks(np.arange(3.7, 4.5+0.1, 0.2))

    ax.invert_yaxis()
    ax.set_yticks([0, 30, 60, 90, 120, 150, 179])
    ax.set_yticklabels(np.arange(0, 181, 30))
    ax.set_xticks(np.arange(0, 360+1, 45))

    ax.set_xticks([0, 45, 90, 135, 180, 225, 270, 315, 359])
    ax.set_xticklabels([0, 45, 90, 135, 180, 225, 270, 315, 360])

    # Mark sunset/rise
    plt.sca(ax)
    local_time_labels(202)

    plt.show()
# ========


def local_time_dist():          # 2
    '''Plot the density at LADEE's latitudes in 1D from multiple simulations.'''

    # Lists of simulation results directories and labels
    A1_dir_sim  = [
                   # Various Q (Fig. 6 Left Panel)
                   "801707_1",
                   "801707_2",
                   "801707_3",
                   "801707_4",
                   "801707_5",
                   "801707_6",
                   "801707_7",
                   "801707_8",
                   "801707_9",

                   # Various P_sq, Q=28 (Right Panel)
#                   "234567_1",
#                   "234567_2",
#                   "234567_3",
#                   "234567_4",
#                   "234567_5",
#                   "234567_6",
                   ]
    A1_label    = [
                   # Various Q (Fig. 6 Left Panel)
                   '24',
                   '25',
                   '26',
                   '27',
                   '28',
                   '29',
                   '30',
                   '31',
                   '32',

                   # Various P_sq, Q=28 (Right Panel)
#                   '0',
#                   r'1$\times$10$^{-3}$',
#                   r'3$\times$10$^{-3}$',
#                   r'5$\times$10$^{-3}$',
#                   r'7$\times$10$^{-3}$',
#                   r'9$\times$10$^{-3}$',
                   ]


    # Check list lengths
    if len(A1_dir_sim) != len(A1_label):
        print 'Different numbers of simulations and labels selected!'


    num_sim     = len(A1_dir_sim)
    num_phi_    = 180

    A2_density_sim    = np.zeros((num_sim, 360))
    A2_density_tmp      = 0


    # Load in all simulations, make into 1D arrays
    for i_sim in range(num_sim):
        dir_cur = dir_bak + A1_dir_sim[i_sim]
        A2_density_tmp  = np.loadtxt(dir_cur + "/density.txt")

        num_theta   = len(A2_density_tmp)
        A1_theta    = np.arange(0, 180+1, 180/num_theta)

        # Scale bin contents by bin area
        for i_theta in range(int(num_theta/2)):
            bin_area    = R**2 * (np.cos(A1_theta[i_theta]*conv) -
                                  np.cos(A1_theta[i_theta+1]*conv)) * conv

            A2_density_tmp[i_theta,:]       /= bin_area
            A2_density_tmp[-(i_theta+1),:]  /= bin_area


        # Sum values to make 1D array at the same latitudes as LADEE
        A1_density_tmp  = np.zeros(360)
        A1_count        = np.zeros(360)
        for phi_ in range(360):
            # Boundary LADEE theta at this phi_
            theta_1 = int(90 + 24 * np.sin((phi_ + 70)*conv))
            theta_2 = int(90 + 24 * np.sin((phi_ + 20)*conv))

            if theta_2 < theta_1:
                tmp     = theta_2
                theta_2 = theta_1
                theta_1 = tmp

            # Add the densities with the right theta, phi_
            # Near the crossovers
            if 20 < phi_ and phi_ < 70:
                for theta in range(min(theta_1, theta_2), 114):
                    A1_density_tmp[phi_]    += A2_density_tmp[theta, phi_]
                    # Record number of additions in each bin
                    A1_count[phi_]  += 1

            elif 200 < phi_ and phi_ < 250:
                for theta in range(66, max(theta_1, theta_2)):
                    A1_density_tmp[phi_]    += A2_density_tmp[theta, phi_]
                    # Record number of additions in each bin
                    A1_count[phi_]  += 1

            else:
                for theta in range(theta_1, theta_2):
                    A1_density_tmp[phi_]    += A2_density_tmp[theta, phi_]
                    # Record number of additions in each bin
                    A1_count[phi_]  += 1

        # Get the means
        A1_density_tmp  /= A1_count

        A2_density_sim[i_sim]   = A1_density_tmp

    # Convert simulation to argon particles (by matching equator sunrise)
    for i_sim in range(num_sim):
        A2_density_sim[i_sim]   *= 2e4 / np.nanmax(A2_density_sim[i_sim])



    fig, ax1    = plt.subplots(figsize=(8,8))

    # Plot simulation results
    colours = new_cmap(np.linspace(0, 1, num_sim))
    for i_sim in range(num_sim):
        plt.plot(np.arange(360)+0.5, A2_density_sim[i_sim], c=colours[i_sim],
                 label=A1_label[i_sim], lw=1.7)


    legend  = plt.legend(loc='lower left', fontsize=12, ncol=2,
                         title='Desorption Energy')
#    legend  = plt.legend(loc='lower left', fontsize=12, ncol=1,
#                         title='Squirrelling\nProbability')


    plt.setp(legend.get_title(), fontsize=12)

    ax1.set_xticks(np.arange(0, 361, 45))

    plt.xlim(0, 360)
    plt.xlabel(r'Local Time of Day ($^\circ$)')
    plt.ylabel(r'Density at the Surface (cm$^{-3}$)')
    plt.ylim(1.2e2, 3.9e4)

    plt.yscale('log')

    # Mark and label sunset/rise
    local_time_labels(4.1e4)

    # Inset zoom-in plot
    if  True:
        #                 left, bottom, width, height
        ax2 = fig.add_axes([0.42, 0.72, 0.23, 0.17])

        phi__1  = 263
        phi__2  = 277

        # Plot simulation results
        for i_sim in range(num_sim):
            ax2.plot(np.arange(360), A2_density_sim[i_sim], c=colours[i_sim],
                     lw=2.5)

        ax2.set_xlim(phi__1, phi__2)
        ax2.set_xticks([265, 270, 275])
        ax2.set_ylim(1.04e4, 2.3e4)
        ax2.set_yscale('log')

        plt.axvline(270, ls='--', lw=1.5, c='grey', alpha=0.5)

    # Bigger axis ticks
    ax1.tick_params('both', width=1.2, length=8, which='major')
    ax1.tick_params('both', width=1.2, length=4, which='minor')
    # Revert x ticks to normal (this weird order seems to work best)
    ax1.xaxis.set_tick_params(width=1, length=4, which='major')

    plt.show()
# ========


def bulge():                    # 3
    '''Plot the bulge at several local times of day.'''
    global  dir_cur
    # Use a results_backup folder instead of results
    dir_cur = dir_bak+"123456/"

    num_phi_    = 18
    num_phi_set = 30

    # Times of day to plot
    A1_i_phi_   = np.arange(num_phi_)
    A1_i_phi_   = [1, 3, 14, 16]


    plt.figure(figsize=(8,8))
    ax  = plt.gca()

    num_phi__in = 360
    num_theta   = 180
    num_phi     = 360

    phi__width  = num_phi__in/num_phi_

    A1_phi_     = np.arange(0, 360+1, phi__width, dtype=int)

    A3_density  = np.zeros((num_phi__in, num_theta, num_phi))

    # Make and save the arrays, sum multiple phi_ maps together
    # (only needed once for a given simulation!)
    if  True:
        for i_phi_ in range(num_phi_):
            phi_        = A1_phi_[i_phi_]
            A2_density  = np.array(np.loadtxt(dir_cur+"densities/density_%d.txt"
                                              % phi_))

            for phi_ in range(A1_phi_[i_phi_], A1_phi_[i_phi_+1]):
                A2_density  += np.array(np.loadtxt(dir_cur+
                                                   "densities/density_%d.txt"
                                                   % phi_))

            # Rotate map to put phi=0 in the centre
            A2_density          = np.roll(A2_density, int(num_phi/2), axis=1)
            A3_density[i_phi_]  = A2_density

            # Save the new arrays
            np.savetxt(dir_cur+"densities/A2_density_%d.txt" % A1_phi_[i_phi_],
                       A2_density)

            print 'Saved A2_density_%s.txt ' % A1_phi_[i_phi_]

        print
    # Load the (pre-made) arrays
    else:
        print 'Loading arrays from %s ...' % dir_cur

        for i_phi_ in range(num_phi_):
            A3_density[i_phi_]  = np.array(np.loadtxt(
                    dir_cur+"densities/A2_density_%d.txt" % A1_phi_[i_phi_]))

    colours = new_cmap(np.linspace(0, 1, num_phi_))

    # Plot the 1D densities at LADEE latitudes
    num_phi = num_phi_set
    A1_phi  = np.arange(0, 360, 360/num_phi) + 180/num_phi
    # Add extra points at start and end to make it wrap round
    A1_phi  = np.append(A1_phi[0] - (A1_phi[1]-A1_phi[0]), A1_phi)
    A1_phi  = np.append(A1_phi, A1_phi[-1] + (A1_phi[1]-A1_phi[0]))

    # Scale by bin area
    A1_theta    = np.arange(0, 180+1, 180/num_theta)
    for i_theta in range(int(num_theta/2)):
        bin_area    = R**2 * (np.cos(A1_theta[i_theta]*conv) -
                              np.cos(A1_theta[i_theta+1]*conv)) * conv

        A3_density[:, i_theta,: ]       /= bin_area
        A3_density[:, -(i_theta+1), :]  /= bin_area

    # Normalise density to match the data post-sunrise (interested in shape not
    # magnitude which is shown by the local time of day distribution)
    A3_density  *= 94 / np.amax(A3_density[:, 90, :])

    # Plot the bulge for each time of day
    for i_phi_ in A1_i_phi_:
        # Sum values to make 1D array at the same latitudes as LADEE
        A1_density_tmp  = np.zeros(360)
        A1_count        = np.zeros(360)

        # Midpoint phi_
        phi_    = (i_phi_+0.5) * 360/num_phi_

        # Boundary LADEE theta at this phi_
        theta_1 = int(90 + 24 * np.sin((phi_ + 70)*conv))
        theta_2 = int(90 + 24 * np.sin((phi_ + 20)*conv))

        if theta_2 < theta_1:
            tmp     = theta_2
            theta_2 = theta_1
            theta_1 = tmp

        # Add the densities with the right theta, phi_
        # Near the crossovers
        if 20 < phi_ and phi_ < 70:
            A1_density  = np.sum(A3_density[i_phi_, min(theta_1, theta_2):114],
                                 axis=0)
        elif 200 < phi_ and phi_ < 250:
            A1_density  = np.sum(A3_density[i_phi_, 66:max(theta_1, theta_2)],
                                 axis=0)
        else:
            A1_density  = np.sum(A3_density[i_phi_, theta_1:theta_2], axis=0)


        # Smooth out over multiple phi bins
        A1_density  = np.sum(A1_density.reshape(-1, int(360/num_phi)), axis=1)

        # Add extra points at start and end to make it wrap round
        A1_density  = np.append(A1_density[-1], A1_density)
        A1_density  = np.append(A1_density, A1_density[0])


        plt.plot(A1_phi, A1_density, lw=2, c=colours[i_phi_],
                 label=(r'%d$^\circ$ $-$ %d$^\circ$' %
                        (i_phi_*360/num_phi_, ((i_phi_+1)*360/num_phi_))))


    legend  = plt.legend(loc='upper right', fontsize=12, title='Time of Day')
    plt.setp(legend.get_title(), fontsize=12)

    plt.xlim(0, 360)
    plt.ylim(5e2, 8.9e3)

    plt.xlabel(r'Selenographic Longitude $(^\circ)$')
    plt.ylabel(r'Nomalised Density at the Surface')

    # Overwrite the phi axis labels to match the rolled-round density array
    ax.set_xticks(np.arange(0, 361, 45))
    ax.set_xticklabels(['180', '225', '270', '315', '0', '45', '90', '135',
                        '180'])

    # Force standard form axis labels:
    formatter = mticker.ScalarFormatter(useMathText=True)
    ax.yaxis.set_major_formatter(formatter)
    ax.get_yaxis().get_major_formatter().set_powerlimits((0, 0))

    plt.tight_layout()



    plt.show()
# ========


def seasonal_long_term():       # 4
    '''Plot the long term variation within +/-30 deg lat from the seasonal
       cold traps.'''
    global  dir_cur
    # Use a results_backup folder instead of results
#    dir_cur = dir_bak+"803043/"
#    dir_cur = dir_bak+"803052/"

    dir_cur = dir_bak+"803595/" # low rates
#    dir_cur = dir_bak+"804194/" # high rates

    A1_bin_source   = np.array(np.loadtxt(dir_cur + "bin_time_start.txt"))
    A1_bin_capt     = np.array(np.loadtxt(dir_cur + "bin_time_capt.txt"))
    A1_bin_dest     = np.array(np.loadtxt(dir_cur + "bin_time_dest.txt"))
    A1_bin_esc      = np.array(np.loadtxt(dir_cur + "bin_time_esc.txt"))

    A1_bin_seasonal_capt_N  = np.array(np.loadtxt(dir_cur +
            "bin_time_seasonal_capt_N.txt"))
    A1_bin_seasonal_rel_N   = np.array(np.loadtxt(dir_cur +
            "bin_time_seasonal_rel_N.txt"))
    A1_bin_seasonal_capt_S  = np.array(np.loadtxt(dir_cur +
            "bin_time_seasonal_capt_S.txt"))
    A1_bin_seasonal_rel_S   = np.array(np.loadtxt(dir_cur +
            "bin_time_seasonal_rel_S.txt"))

    t_end   = 200   # year

    # Rate units: N per bin_width
    bin_width   = t_end / 1e4

    # Loss and total rate from the +/-30 deg lat region
    A1_bin_loss     = A1_bin_dest + A1_bin_esc
    A1_bin_total    = A1_bin_source - A1_bin_loss

    # Get total number of particles from the total rate
    A1_bin_number   = np.cumsum(A1_bin_total)

    # x axis
    A1_t_bin_mid    = np.arange(0, t_end, bin_width) + bin_width/2
    A1_t_bin_end    = np.arange(0, t_end, bin_width) + bin_width


    fig = plt.figure(figsize=(8,8))
    gs  = gridspec.GridSpec(8, 8)
    gs.update(hspace=0)
    ax1 = plt.subplot(gs[:5, :])
    ax2 = plt.subplot(gs[5:, :])

    # Mean post-eqm seasonal cold traps (with unscaled number)
    eqm_year    = 100
    N_year      = 2

    # Sum every N_year section together to get mean N_year section
    A1_bin_number           = np.sum(
            A1_bin_number[eqm_year/bin_width:].reshape(
            -1, N_year/bin_width), axis=0)
    A1_bin_seasonal_capt_N  = np.sum(
            A1_bin_seasonal_capt_N[eqm_year/bin_width:].reshape(
            -1, N_year/bin_width), axis=0)
    A1_bin_seasonal_rel_N   = np.sum(
            A1_bin_seasonal_rel_N[eqm_year/bin_width:].reshape(
            -1, N_year/bin_width), axis=0)
    A1_bin_seasonal_capt_S  = np.sum(
            A1_bin_seasonal_capt_S[eqm_year/bin_width:].reshape(
            -1, N_year/bin_width), axis=0)
    A1_bin_seasonal_rel_S   = np.sum(
            A1_bin_seasonal_rel_S[eqm_year/bin_width:].reshape(
            -1, N_year/bin_width), axis=0)

    A1_t_bin_end    = np.arange(0, N_year, bin_width) + bin_width
    A1_t_bin_mid    = np.arange(0, N_year, bin_width) + bin_width/2

    # Normalise so the mean is 1
    scale   = 1 / np.mean(A1_bin_number)

    ax1.plot(A1_t_bin_end, A1_bin_number*scale, c='purple', lw=2, ls='-',
             label='small permanent traps')

    # Normalise captured/released numbers to be fractions of the total exosphere
    # (0.187 comes from eqm distribution, how much of the exosphere is in +/-30)
    scale   = 1 / np.mean(A1_bin_number) * 0.187

    ax2.plot(A1_t_bin_mid, A1_bin_seasonal_capt_N*scale, c='brown', lw=2,
             ls='-',  label='N capture')
    ax2.plot(A1_t_bin_mid, A1_bin_seasonal_rel_N*scale, c='r', lw=2, ls='-',
             label='N release')

    ax2.plot(A1_t_bin_mid, A1_bin_seasonal_capt_S*scale, c='b', lw=2, ls='-',
             label='S capture')
    ax2.plot(A1_t_bin_mid, A1_bin_seasonal_rel_S*scale, c='c', lw=2, ls='-',
             label='S release')

    ax2.legend(loc='upper center', fontsize=12)

    ax1.set_xticks(np.arange(0.75, 1.75+0.01, 0.25))
    ax1.set_xticklabels([])
    ax2.set_xticks(np.arange(0.75, 1.75+0.01, 0.25))
    ax2.set_xticklabels([r'$-$0.25', '0', '0.25' ,'0.5', '0.75'])

    ax1.set_xlim(0.75, N_year-0.25)
    ax2.set_xlim(0.75, N_year-0.25)

    ax1.set_ylim(0.71, 1.29)
    ax2.set_ylim(0, 0.15)

    ax2.set_xlabel('Time Since Lunar Vernal Equinox (year)')
    ax1.set_ylabel(r'Norm. Density at the Surface')
    ax2.set_ylabel(r'Norm. No. of Particles')

    plt.show()
# ========


# //////////////////////////////////////////////////////////////////////////// #
#                               2. Main Program                                #
# //////////////////////////////////////////////////////////////////////////// #


# System argument selects which function to run:
if len(sys.argv) == 2:
    func_num    = sys.argv[1]

else:
    print 'Enter a func_num!'

# ========
if   func_num   == '1':
    density_2D()
elif func_num   == '2':
    local_time_dist()
elif func_num   == '3':
    bulge()
elif func_num   == '4':
    seasonal_long_term()
# ========

else:
    print 'No valid system argument...\n'











































