
from datetime import datetime, timedelta
import os

import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
import pandas
import xarray
import scipy.ndimage as ndi

from . import io, map


def plot_domain_example():
    path = '/nas/smartcarb/output/'
    mod_fname = os.path.join(path, 'cosmo2D_2015-01/cosmo_2d_2015042311.nc')
    l2_fname = os.path.join(
        path, 'level2/Sentinel_7_CO2/Sentinel_7_CO2_2015042311_o1670_l0401.nc'
    )

    d = xarray.open_dataset(mod_fname, decode_times=False)
    values = io.read_total_tracer(mod_fname, 'CO2', time=datetime(2015,7,3,11))

    d2 = io.read_level2(l2_fname, 0.5, 0.01)

    with plt.rc_context({
            'axes.autolimit_mode': 'round_numbers',
            'font.size': 9,
            'xtick.labelsize': 'small',
            'ytick.labelsize': 'small',
            #'text.usetex': True,
            #'text.latex.unicode': True,
        }):
        fig = map.make_cosmo_map(d.rlon, d.rlat, values, vmin=404, vmax=408,
                                 clct=d.CLCT[0], label='XCO$_2$ (ppm)',
                                 fig_size=5)

        map.make_level2_map(d2.lonc, d2.latc, d2.CO2, vmin=403, vmax=408,
                            fig=fig, bg_color=None, clct=d2.clouds, 
                            label='XCO$_2$ (ppm)', do_zoom=False, xlim=None,
                            ylim=None)

        ax, cax = fig.get_axes()

        ha= [
            'right',
            'left',
            'right',
            'right',
            'right',
            'right',
            'right',
        ]
        map.add_city_labels(ax, size='medium', va=None, ha=ha)





def plot_emission_heights(filename, name='JAE', do_diurnal=False):
    """
    Plot time series of emission heights (annual or diurnal).
    """
    d = xarray.open_dataset(filename)

    # find index for id/name
    for i,s in enumerate(np.array(d.id, str)):
        if s == name:
            break

    H = d.H[i]
    h = np.array(d.h_ts)[i]
    h30 = ndi.convolve(h, np.ones(720)/720.0)
    times = [datetime(2015,1,1) + timedelta(hours=hour) for hour in
             range(24*365)]

    ts = pandas.Series(h, times)

    fig = plt.figure()
    ax = fig.add_subplot(111)

    if do_diurnal:
        hours = np.arange(24)
        heights = [np.mean(ts[ts.index.hour == hour]) for hour in hours]

        ax.plot(hours, heights)
        ax.set_xlabel('hour of day (UTC)')
        ax.set_xlim(0,23)
        ax.set_xticks(np.arange(0,24,3))

    else:
        ax.plot(ts.index, ts, label='hourly')
        ax.plot(times, h30, label='30-day mean')
        ax.plot([times[0], times[-1]], [H,H], 'k-', label='stack height')

        ax.xaxis.set_major_formatter(mdates.DateFormatter('%b'))

        ax.set_ylim(0,2000)
        ax.set_xlabel('2015')
        plt.legend()

    ax.set_ylabel('height (m)')
    plt.tight_layout()


