#!/usr/bin/python env
#

import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

from matplotlib.patches import Polygon

from matplotlib.ticker import MultipleLocator

import numpy as np
import h5py as h5

from colors import *
from mesh import *

# Routines

def read_fbuffs (scalar_filename, velocity_filename, R_scale=1.) :

    # Read scalar data

    with h5.File(scalar_filename, 'r') as f :

        Teff = f['dataset'][:,:,0]/1E3
        geff = f['dataset'][:,:,1]
        star = f['dataset'][:,:,2]

        x_min = f.attrs['x_min'][0]*R_scale
        x_max = f.attrs['x_max'][0]*R_scale

        y_min = f.attrs['y_min'][0]*R_scale
        y_max = f.attrs['y_max'][0]*R_scale

    # Read velocity data

    with h5.File(velocity_filename, 'r') as f :

        v_x = f['dataset'][:,:,0]
        v_y = f['dataset'][:,:,1]
        v_z = f['dataset'][:,:,2]

    # Mask data

    mask = star == 0.

    Teff_masked = np.ma.masked_where(mask, Teff)
    geff_masked = np.ma.masked_where(mask, geff)

    v_x_masked = np.ma.masked_where(mask, v_x)
    v_y_masked = np.ma.masked_where(mask, v_y)
    v_z_masked = np.ma.masked_where(mask, v_z)
    
    return {'Teff': Teff_masked,
            'geff': geff_masked,
            'v_x': v_x_masked,
            'v_y': v_y_masked,
            'v_z': v_z_masked,
            'x_min': x_min,
            'x_max': x_max,
            'y_min': y_min,
            'y_max': y_max,
            'n_x': mask.shape[0],
            'n_y': mask.shape[1]}

#

def plot_mesh (ax, m, inc, R_star, **kwargs) :

    # Rotate around the x axis

    vr_rot = np.empty([m['n_vert'],3])

    i = (90 - inc)/180.*np.pi

    R = np.array([[1.,0.,0.],
                  [0.,np.cos(i),np.sin(i)],
                  [0.,-np.sin(i),np.cos(i)]])

    for i in range(0, m['n_vert']) :
        vr = m['verts'][i,:]/R_star
        vr_rot[i,:] = R.dot(vr)

    # Determine the center of each triangle

    vr_c = np.empty([m['n_face'],3])

    for i in range(0, m['n_face']) :
        j_vt = m['faces'][i,:] - 1
        vr_1 = vr_rot[j_vt[0],:]
        vr_2 = vr_rot[j_vt[1],:]
        vr_3 = vr_rot[j_vt[2],:]
        vr_c[i,:] = (vr_1 + vr_2 + vr_3)/3.

    # Sort into order

    i_srt = np.argsort(vr_c[:,2])

    # Draw the triangles

    for i in range(0, m['n_face']) : 
        j_vt = m['faces'][i,:] - 1
        vr_1 = vr_rot[j_vt[0],:]
        vr_2 = vr_rot[j_vt[1],:]
        vr_3 = vr_rot[j_vt[2],:]
        xy = np.array([vr_1[0:2],vr_2[0:2],vr_3[0:2]])
        p = Polygon(xy, closed=True, zorder=i, **kwargs)
        ax.add_patch(p)

# 

def plot_mesh_panel (ax, i, j, m, inc, R_star, title, text) :

    plot_mesh(ax[i,j], m, inc, R_star, lw=0.25, fc=WHITE, ec=BLACK)

    ax[i,j].set_xlim(0,3.5)
    ax[i,j].set_ylim(-1.5,1.5)

    if i == 0 :
        ax[i,j].set_title(title)

    if i == 1 :
        ax[i,j].set_xlabel(r'$x\ [R_{\rm p}]$')

    if j == 0 :
        ax[i,j].set_ylabel(r'$y\ [R_{\rm p}]$')
        
    ax[i,j].text(0.05, 0.95, text, ha='left', va='top', fontsize=8, transform=ax[i,j].transAxes)


# 

def plot_fbuff_panel (ax, i, j, fb, var, title, text, vmin, vmax, ims, k, mapname) :

    extent = [fb['x_min'],fb['x_max'],fb['y_min'],fb['y_max']]

    ims[k] = ax[i,j].imshow(fb[var], extent=extent, vmin=vmin, vmax=vmax,
                            cmap=plt.get_cmap(mapname), aspect='auto')

    ax[i,j].set_xlim(0,3.5)
    ax[i,j].set_ylim(-1.5,1.5)

    if i == 0 :
        ax[i,j].set_title(title)

    if i == 1 :
        ax[i,j].set_xlabel(r'$x\ [R_{\rm p}]$')

    if j == 0 :
        ax[i,j].set_ylabel(r'$y\ [R_{\rm p}]$')

    ax[i,j].text(0.05, 0.95, text, ha='left', va='top', fontsize=8, transform=ax[i,j].transAxes)

# Read data

R_star = 10.23
R_sun = 6.955E10

m_pri = read_mesh('mesh_pri_lo.h5')
m_sec = read_mesh('mesh_sec_lo.h5')

fb_00 = read_fbuffs('scalar_i90.h5', 'velocity_i90.h5', 1./R_star)
fb_90 = read_fbuffs('scalar_i00.h5', 'velocity_i00.h5', 1./R_star)

Teff_min = min(np.min(fb_00['Teff']), np.min(fb_90['Teff']))
Teff_max = max(np.max(fb_00['Teff']), np.max(fb_90['Teff']))

geff_min = min(np.min(fb_00['geff']), np.min(fb_90['geff']))
geff_max = max(np.max(fb_00['geff']), np.max(fb_90['geff']))

print "Teff range:", Teff_min, Teff_max
print "geff range:", geff_min, geff_max

# Set up the plots

params = {'backend': 'pdf',
          'figure.figsize': [7, 3.5],
          'font.family':'serif',
          'font.size':10,
          'font.serif': 'Times Roman',
          'axes.titlesize': 'medium',
          'axes.labelsize': 'medium',
          'legend.fontsize': 8,
          'legend.frameon' : False,
          'text.usetex': True,
          'figure.dpi': 600,
          'lines.markersize': 4,
          'lines.linewidth': 1,
          'lines.antialiased': False,
          'path.simplify': False }

mpl.rcParams.update(params)

fig = plt.figure()

gs_f = gridspec.GridSpec(2, 3)

ax_f = np.empty([2,3], dtype=plt.Axes)

for i in [0,1] :
    for j in [0,1,2] :
        if j >= 1 :
            sharex = ax_f[i,0]
        else :
            sharex = None
        if i == 1 :
            sharey = ax_f[0,j]
        else :
            sharey = None
        ax_f[i,j] = fig.add_subplot(gs_f[i,j], sharex=sharex, sharey=sharey)
        ax_f[i,j].label_outer()

gs_c = gridspec.GridSpec(2, 1)

ax_c = np.empty([2,1], dtype=plt.Axes)

for i in [0,1] :
    ax_c[i,0] = fig.add_subplot(gs_c[i,0])

# Do the plots

plot_mesh_panel(ax_f, 0, 0, m_pri, 0., R_star*R_sun, r'Mesh', r'$i=0^{\circ}$')
plot_mesh_panel(ax_f, 1, 0, m_pri, 90., R_star*R_sun, r'Mesh', r'$i=90^{\circ}$')

ims = np.empty([4], dtype=mpl.image.AxesImage)

T_min = 29.5
T_max = 31.5

g_min = 3.75
g_max = 3.95

plot_fbuff_panel(ax_f, 0, 1, fb_00, 'Teff', r'$T_{\rm eff}$', r'$i=0^{\circ}$', T_min, T_max, ims, 0, 'rainbow')
plot_fbuff_panel(ax_f, 1, 1, fb_90, 'Teff', r'$T_{\rm eff}$', r'$i=90^{\circ}$', T_min, T_max, ims, 1, 'rainbow')

plot_fbuff_panel(ax_f, 0, 2, fb_00, 'geff', r'$\log g_{\rm eff}$', r'$i=0^{\circ}$', g_min, g_max, ims, 2, 'spring')
plot_fbuff_panel(ax_f, 1, 2, fb_90, 'geff', r'$\log g_{\rm eff}$', r'$i=90^{\circ}$', g_min, g_max, ims, 3, 'spring')

# Colorbars

cb_ticks = [29.5,30,30.5,31,31.5]

cb_Teff = plt.colorbar(ims[0], cax=ax_c[0,0], ticks=cb_ticks, label=r'$T_{\rm eff}\ [{\rm 10^{3}\,K}]$')

cb_ticks = [3.75,3.8,3.85,3.9,3.95]

cb_geff = plt.colorbar(ims[2], cax=ax_c[1,0], ticks=cb_ticks, label=r'$\log (g_{\rm eff}/{\rm cm\,s^{-2}})$')

cb_Teff.ax.tick_params(labelsize=8)
cb_geff.ax.tick_params(labelsize=8)

# Adjust dimensions

gs_f.tight_layout(fig, pad=0.5, rect=[0,0,0.88,1])
gs_c.tight_layout(fig, pad=0.5, rect=[0.88,0,1,1])

top = min(gs_f.top, gs_c.top)
bottom = max(gs_f.bottom, gs_c.bottom)

gs_f.update(top=top, bottom=bottom)
gs_c.update(top=top, bottom=bottom)

# Reset data limits for the maps (using aspect='equal' with
# adjustable='datalim' doesn't seem to properly account for the shape
# of the overall figure)

x_min = 0
x_max = 3

for i in [0,1] :
    for j in [0,1,2] :
        bounds = ax_f[i,j].get_position().bounds
        width = bounds[2]*fig.get_figwidth()
        height = bounds[3]*fig.get_figheight()
        aspect = height/width
        ax_f[i,j].set_xlim(x_min, x_max)
        ax_f[i,j].set_ylim(-0.5*(x_max-x_min)*aspect, 0.5*(x_max-x_min)*aspect)
    
# Write out

fig.savefig('fig_props.pdf')
