import matplotlib.pyplot as plt
import numpy as np
from scipy.io.netcdf import netcdf_file
#from pyatmlib.radiometry.planck import btemp_power_wavenum as rad2bt

class constants():
    pass
constants.k = 1.380650e-23
constants.k_units = 'J K^-1'
constants.c = 2.997925e8
constants.c_units = 'm s^-1'
constants.h = 6.626069e-34
constants.h_units = 'J s'
constants._A = constants.h * constants.c / constants.k
constants._B = 2 * constants.h * constants.c **2

def rad2bt(radiance, wavenum):
    """ BT = btemp_power_wavenum(radiance, wavenum):
    Given an input radiance in [W m^-2 sr^-1 (cm^-1)^-1] and a 
    wavenumber in [cm^-1], compute the equivalent black body 
    temperature (the "brightness temp").

    """

    # factor of 1e8 is from converting (cm^-1)^-1 to (m^-1)^-1:
    # 1e6 from converting wavenum * 1e2 from converting radiance
    tmp = np.log( constants._B*np.power(wavenum,3)*1e8/radiance + 1 );
    BT = constants._A * 100 * wavenum / tmp;
    return BT

def get_data(ncfilename, min_wnum=895.0, max_wnum=905.0):

    nc = netcdf_file(ncfilename,'r')

    lat = nc.variables['Latitude'].data.astype(np.float)
    lon = nc.variables['Longitude'].data.astype(np.float)

    rad = nc.variables['radiance'].data.T
    wnum = nc.variables['wavenumber'].data

    nc.close()

    bt = rad2bt(rad*1e-3, wnum[:,np.newaxis])
    a = wnum.searchsorted(min_wnum)
    b = wnum.searchsorted(max_wnum)
    bt1 = bt[a:b,:].mean(0)
    
    return lat, lon, bt1

def make_a_plot(lat, lon, fovs, hAx=None, **kw):
    if hAx:
        hAx.scatter(lon, lat, c=fovs, s=20, edgecolor='none', **kw)
        hAx.set_xlabel('Deg. Longitude')
        hAx.set_ylabel('Deg. Latitude')
        hAx.grid(1)        
    else:
        plt.scatter(lon, lat, c=fovs, s=20, edgecolor='none', **kw)
        plt.xlabel('Deg. Longitude')
        plt.ylabel('Deg. Latitude')
        plt.grid(1)

def make_the_plots():

    file_list = [
        'SHIS_rdr20140916T154110sdr20140916T161211_rad.nc', 
        'SHIS_rdr20140916T174114sdr20140916T181215_rad.nc', 
        'SHIS_rdr20140916T194118sdr20140916T201221_rad.nc', 
        'SHIS_rdr20140916T221123sdr20140916T224224_rad.nc', 
        'SHIS_rdr20140917T001127sdr20140917T004235_rad.nc', 
        'SHIS_rdr20140917T021131sdr20140917T024300_rad.nc', 
        'SHIS_rdr20140917T024132sdr20140917T031234_rad.nc', ]

    # lengths: [1694, 1918, 1610, 2114, 2016, 1806, 2072]
    slicers = [slice(1200, 1700, None), 
               slice( 700, 1100, None),
               slice( 950, 1350, None),
               slice(   0,  650, None), 
               slice( 850, 1330, None), 
               slice(1600, None, None),
               slice(None,  300, None) ]

    hFig1 = plt.figure(1)
    hFig1.clf()
    for n, (s,f) in enumerate(zip(slicers,file_list)):
        lat, lon, fovs = get_data(f)
        make_a_plot(lat[s], lon[s], fovs[s], vmin=200, vmax=300)
    plt.colorbar()
    plt.savefig('eye_bt_collage.png')


    slicers = [slice( 600, 1900, None), 
               slice( 400, 1400, None),
               slice( 700, 1500, None),
               slice(   0,  950, None), 
               slice( 700, 1600, None), 
               slice(1400, None, None),
               slice(None,  500, None) ]
    
    hFig2 = plt.figure(2)

    # the silliness with 'plotters' is because the 5th and 6th files are both 
    # eyepass 6, because it split over 2 files.
    plotters = [0,1,2,3,4,5,5]

    hFig2.clf()
    hAxs = [hFig2.add_subplot(3,3,n+1) for n in range(9)]
    for n, (s,f) in enumerate(zip(slicers,file_list)):
        lat, lon, fovs = get_data(f)
        make_a_plot(lat[s], lon[s], fovs[s], vmin=200, vmax=300, hAx=hAxs[plotters[n]])
    plt.savefig('eye_bt_grid_wholerange.png')

    hFig2.clf()
    hAxs = [hFig2.add_subplot(3,3,n+1) for n in range(9)]
    for n, (s,f) in enumerate(zip(slicers,file_list)):
        lat, lon, fovs = get_data(f)
        make_a_plot(lat[s], lon[s], fovs[s], vmin=190, vmax=230, hAx=hAxs[plotters[n]])
        xlim = hAxs[n].get_xlim()
        ylim = hAxs[n].get_ylim()
        # make square axes ranges
        if (xlim[1]-xlim[0]) > (ylim[1]-ylim[0]):
            width = xlim[1]-xlim[0]
            ymid = (ylim[1]+ylim[0])*0.5
            hAxs[n].set_ylim([ymid-0.5*width, ymid+0.5*width])
        else:
            width = ylim[1]-ylim[0]
            xmid = (xlim[1]+xlim[0])*0.5
            hAxs[n].set_xlim([xmid-0.5*width, xmid+0.5*width])
    plt.savefig('eye_bt_grid_coldrange.png')

    hFig2.clf()
    hAxs = [hFig2.add_subplot(3,3,n+1) for n in range(9)]
    for n, (s,f) in enumerate(zip(slicers,file_list)):
        lat, lon, fovs = get_data(f)
        make_a_plot(lat[s], lon[s], fovs[s], vmin=200, vmax=300, hAx=hAxs[plotters[n]])
        xlim = hAxs[n].get_xlim()
        ylim = hAxs[n].get_ylim()
        # make square axes ranges
        if np.cos(np.deg2rad(ylim[0]))*(xlim[1]-xlim[0]) > (ylim[1]-ylim[0]):
            width = xlim[1]-xlim[0]
            ymid = (ylim[1]+ylim[0])*0.5
            hAxs[n].set_ylim([ymid-0.5*width, ymid+0.5*width])
        else:
            width = ylim[1]-ylim[0]
            xmid = (xlim[1]+xlim[0])*0.5
            hAxs[n].set_xlim([xmid-0.5*width, xmid+0.5*width])
    plt.savefig('eye_bt_grid_wholerange.png')