# Script for 1st order closure model for CBL with no shear.

# Converted from MATLAB script
# April 2016
# Emily Ramnarine
# University of Washington
# Department of Atmospheric Sciences

from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint

# asymptotic Blackadar lengthscale
lmbda = 100

# surface sensible heat flux (W m^{-2})
surf_sens_heat_flux = 300

rho_ref = 1.2
Cp = 1000.
theta_flux0 = surf_sens_heat_flux / (rho_ref * Cp)

# height (z) parameters
# evenly spaced flux levels from 0-z_top m with spacing dz
# with initial theta (theta0) given at half-levels z_half
z_top = 1000
dz = 100
z_half = np.arange(dz/2.0, z_top-dz/2.0 + dz, dz)
z_full = np.arange(0, z_top + dz, dz)

# time (t) parameters
t_f = 3*3600
dt = 600
t_span = np.arange(0, t_f + dt, dt)

theta0 = 290 + 0.01 * z_half

# Specify Blackadar master lengthscale for interior flux points
# for use in ODE solver.
k = 0.4
length_blackadar = lmbda / (1 + lmbda / (k*z_full[1:len(z_half)]))


def d_theta_dt(theta, t, z_full, z_half, length_blackadar, theta_flux0,
               extras_out=False):
    '''
    theta tendency ODE solver used to calculate deepening of surf-heated BL.

    Arguments       Size        Units    Description
    t               1           s        time
    theta           nz x 1      K        at half-levels z_half

    Parameters
    z_full          (nz+1) x 1  m        height of flux levels
    z_half          nz x 1      m        height of half (thermo) levels
    length_blackadar(nz-1) x 1  m        Blackadar lengthscale at heights z(2)-z(nz)
    theta_flux0     1           K m/s    constant surface theta flux
    extras_out      1           None     whether K and theta_flux are returned

    Output
    tendency        nz x 1      m        tendency d(theta)/dt at the half-levels
    K               (nz-1) x 1  m^2/s    eddy diffusivity at interior flux levels
    theta_flux      (nz+1) x 1  K m/s    theta flux at flux levels
    '''
    nz = len(z_half)
    g = 9.8
    theta_ref = np.mean(theta)
    d_theta_dz = np.diff(theta) / np.diff(z_half)

    d_bouyancy_dz = (g/theta_ref) * d_theta_dz
    d_bouyancy_dz[d_bouyancy_dz > 0] = 0

    # We use the convective limit (Ri -> inf) of eddy diffusivity formula
    # K = length_blackadar^2*|du/dz|*sqrt(1-16Ri)
    # K ~ length_blackadar^2*sqrt(-16*db/dz) (unstable)
    # We take K=0 for stable stratification.

    K = length_blackadar**2 * (-16*d_bouyancy_dz)**0.5

    theta_flux = np.zeros((nz + 1,))
    theta_flux[0] = theta_flux0
    theta_flux[1:nz] = -K * d_theta_dz
    theta_flux[nz] = 0

    tendency = -np.diff(theta_flux) / np.diff(z_full)

    if extras_out:
        return tendency, K, theta_flux
    else:
        return tendency

# Use ODE solver to solve the equation.
theta = odeint(d_theta_dt, theta0, t_span, args=(
    z_full, z_half, length_blackadar, theta_flux0))
t, z = np.broadcast_arrays(t_span[:, None], z_half[None, :])

# Generate hourly profiles of theta, heat flux, and K
times = []
theta_hourly = []
K_hourly = []
theta_flux_hourly = []
for i in range(0, len(t_span)+1, int(3600/dt)):
    tendency, K, theta_flux = d_theta_dt(
        theta[i, :], t_span[i], z_full, z_half, length_blackadar, theta_flux0,
        extras_out=True)

    times.append('{:.0f} hr'.format(t_span[i]/3600))
    theta_hourly.append(theta[i, :])

    new_K = np.zeros_like(z_full)
    new_K[1:-1] = K[:]
    K_hourly.append(new_K)

    theta_flux_hourly.append(theta_flux)

# Plot the Data
plt.figure()

plt.subplot(2, 2, 1)
im = plt.pcolormesh(t/3600., z, theta, cmap='gray', shading='gouraud')
plt.contour(t/3600., z, theta, 10, colors='k')
plt.colorbar(im)
plt.xlabel('time (hr)')
plt.ylabel('z (m)')
plt.ylim(0, 1000)
plt.tick_params(axis='both', labelsize=8)
plt.title('Time-height section of $\\theta$')

plt.subplot(2, 2, 2)
for i in range(len(theta_hourly)):
    plt.plot(theta_hourly[i], z_half, label=times[i])
plt.legend(fontsize=8)
plt.xlabel('$\\theta$ (K)')
plt.ylabel('z (m)')
plt.tick_params(axis='both', labelsize=8)
plt.title('Hourly profiles of $\\theta$')
plt.text(0.5 + min(theta0), 0.75*z_top, '$\lambda$ = {} m'.format(lmbda))
unit = 'W m$^{-2}$'
plt.text(0.5 + min(theta0), 0.9*z_top,
         'H$_s$(0) = {} {}'.format(surf_sens_heat_flux, unit))

plt.subplot(2, 2, 3)
for i in range(len(theta_flux_hourly)):
    plt.plot(theta_flux_hourly[i], z_full, label=times[i])
plt.legend(fontsize=8)
plt.xlabel('Sensible heat flux (W m$^{-2}$)')
plt.xlim([0, 0.3])
plt.ylabel('z (m)')
plt.ylim([0, 1000])
plt.tick_params(axis='both', labelsize=8)
plt.title('Hourly heat flux profiles')

plt.subplot(2, 2, 4)
for i in range(len(K_hourly)):
    plt.plot(K_hourly[i], z_full, label=times[i])
plt.legend(fontsize=8)
plt.xlabel('K (m s$^{-1}$)')
plt.xlim([0, 100])
plt.ylabel('z (m)')
plt.ylim([0, 1000])
plt.tick_params(axis='both', labelsize=8)
plt.title('Hourly K profiles')

plt.tight_layout()
# plt.show()
plt.savefig('fig.png')