# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: 'Python 3.6.9 64-bit (''.venv'': virtualenv)'
#     language: python
#     name: python36964bitvenvvirtualenv7d76ec95724e4c80a84536475a51dccd
# ---

# # Elevrad, toporad and stoporad
#
# Ensure that the conversion of the 3 functions to Python match with the IPW versions

# +
import os
import glob
import netCDF4 as nc
import pandas as pd
import numpy as np
import subprocess as sp
from spatialnc import ipw
from spatialnc.topo import get_topo_stats
from topocalc.shade import shade
from topocalc.horizon import horizon
from topocalc import gradient
import matplotlib.pyplot as plt
from datetime import datetime
import seaborn as sns

from smrf.data import loadTopo
from smrf.envphys.sunang import sunang
from smrf.envphys.albedo import albedo
from smrf.envphys.solar import toporad, irradiance
from smrf.utils import utils

sns.set_context('poster')

# +
# Topo, create an IPW image from topo, sky view factor and terrain config factor

# Load the netcdf dem and get all the geo coords for IPW
topo_path = '../tests/Lakes/topo/topo.nc'

topo_config = {
    'filename': topo_path,
    'northern_hemisphere': True,
    'gradient_method': 'gradient_d8',
    'sky_view_factor_angles': 72
}
topo = loadTopo.Topo(topo_config, tempDir='.')
dem = topo.dem

ts = get_topo_stats(topo_path)

csys = 'UTM'
nbits = 16

# Create the IPW image for dem
ipw_image = ipw.IPW()
ipw_image.new_band(dem)
ipw_image.add_geo_hdr(
    coordinates=[ts['u'], ts['v']],
    d=[ts['dv'], ts['du']],
    units=ts['units'],
    csys=csys
)
ipw_image.write('lakes_dem.ipw', nbits=nbits)

# Create the IPW image for sky view factor
ipw_image = ipw.IPW()
ipw_image.new_band(topo.sky_view_factor)
ipw_image.add_geo_hdr(
    coordinates=[ts['u'], ts['v']],
    d=[ts['dv'], ts['du']],
    units=ts['units'],
    csys=csys
)
ipw_image.write('svf.ipw', nbits=nbits)

# Create the IPW image for terrain configuration factor
ipw_image = ipw.IPW()
ipw_image.new_band(topo.terrain_config_factor)
ipw_image.add_geo_hdr(
    coordinates=[ts['u'], ts['v']],
    d=[ts['dv'], ts['du']],
    units=ts['units'],
    csys=csys
)
ipw_image.write('tcf.ipw', nbits=nbits)



# +
# Parameters for elevrad, toporad and stoporad

wavelength = 'vis'

date_time = pd.to_datetime('2/15/1990 20:00')
date_time = date_time.tz_localize('UTC')

min_storm_day = 25.0
wy_day, wyear = utils.water_day(date_time)
tz_min_west = int(np.abs(date_time.utcoffset().total_seconds()/60))

tau_elevation = 100
tau = 0.2
omega = 0.85
scattering_factor = 0.3
surface_albedo = 0.5
grain_size = 100.0
max_grain = 2000.0
dirt = 2.0
wavelength_vis = [0.28, 0.7]
wavelength_ir = [0.7, 2.8]

if wavelength == 'vis':
    solar_irradiance = irradiance.direct_solar_irradiance(date_time, w=wavelength_vis)
else:
    solar_irradiance = irradiance.direct_solar_irradiance(date_time, w=wavelength_ir)

cosz, azimuth, rad_vec = sunang(
    date_time,
    topo.basin_lat,
    topo.basin_long)

illum_ang = shade(
    topo.sin_slope,
    topo.aspect,
    azimuth,
    cosz)

alb_vis, alb_ir = albedo(min_storm_day * np.ones_like(dem), illum_ang, grain_size, max_grain)

# Create the IPW image for illumination angle
ipw_image = ipw.IPW()
ipw_image.new_band(illum_ang)
ipw_image.add_geo_hdr(
    coordinates=[ts['u'], ts['v']],
    d=[ts['dv'], ts['du']],
    units=ts['units'],
    csys=csys
)

ipw_image.write('illum_angle.ipw', nbits=nbits)

# Create the IPW image for albedo
ipw_image = ipw.IPW()
if wavelength == 'vis':
    ipw_image.new_band(alb_vis)
else:
    ipw_image.new_band(alb_ir)
ipw_image.add_geo_hdr(
    coordinates=[ts['u'], ts['v']],
    d=[ts['dv'], ts['du']],
    units=ts['units'],
    csys=csys
)

ipw_image.write('albedo.ipw', nbits=nbits)


# +
# Elevrad
# The only difference is bit resolution noise

# Call IPW elevrad
cmd = f"elevrad -z {tau_elevation} -t {tau} -w {omega} -g {scattering_factor} -r {surface_albedo} -s {solar_irradiance} -u {cosz} -n 16 lakes_dem.ipw > elevrad.ipw"

print(cmd)
visp = sp.Popen(cmd, shell=True)
visp.wait()

elevrad_ipw = ipw.IPW('elevrad.ipw')

# Python elevrad
elevrad_py = toporad.Elevrad(
    dem,
    solar_irradiance,
    cosz)

# Plot
fig, ax = plt.subplots(2, 3, figsize = (30, 20))
fig.subplots_adjust(left=0.02, bottom=0.06, right=0.95, top=0.94, wspace=0.05)
fig.suptitle('ELEVRAD')

# IPW originals
im0 = ax[0, 0].imshow(elevrad_ipw.bands[0].data)
ax[0, 0].set_title('IPW beam')
fig.colorbar(im0, ax=ax[0, 0])

im0 = ax[1, 0].imshow(elevrad_ipw.bands[1].data)
ax[1, 0].set_title('IPW diffuse')
fig.colorbar(im0, ax=ax[1, 0])

# Python elevrad
im0 = ax[0, 1].imshow(elevrad_py.beam)
ax[0, 1].set_title('Python beam')
fig.colorbar(im0, ax=ax[0, 1])

im0 = ax[1, 1].imshow(elevrad_py.diffuse)
ax[1, 1].set_title('Python diffuse')
fig.colorbar(im0, ax=ax[1, 1])

# difference
d = elevrad_ipw.bands[0].data - elevrad_py.beam
var_lim = max([abs(d.min()), abs(d.max())])
# im2 = ax[0, 2].imshow(d, cmap='RdBu', vmin=-var_lim, vmax=var_lim)
# fig.colorbar(im2, ax=ax[0, 2])
ax[0, 2].hist(d.flatten(), bins=100)
ax[0, 2].yaxis.tick_right()
ax[0, 2].set_title('Beam difference, ipw - python')

d = elevrad_ipw.bands[1].data - elevrad_py.diffuse
var_lim = max([abs(d.min()), abs(d.max())])
# im2 = ax[1, 2].imshow(d, cmap='RdBu', vmin=-var_lim, vmax=var_lim)
# fig.colorbar(im2, ax=ax[1, 2])
ax[1, 2].hist(d.flatten(), bins=100)
ax[1, 2].yaxis.tick_right()
ax[1, 2].set_title('Diffuse difference, ipw - python')

# remove the x and y ticks
for a in ax[:,:2].flatten():
    a.set_xticks([])
    a.set_yticks([])

# +
# toporad
# The only difference is bit resolution noise

# IPW toporad
# toporad requires a sun header that has the cosz and azimuth. Since this example hasn't done that, toporad sets cosz to 0.5

cmd = f"mux elevrad.ipw illum_angle.ipw svf.ipw tcf.ipw albedo.ipw | toporad -r {surface_albedo} -d > toporad.ipw"

print(cmd)
visp = sp.Popen(cmd, shell=True)
visp.wait()

toporad_ipw = ipw.IPW('toporad.ipw')

# Python toporad
erad = toporad.Elevrad(dem, solar_irradiance, cosz)

trad_beam, trad_diffuse = toporad.toporad(
    erad.beam,
    erad.diffuse,
    illum_ang,
    topo.sky_view_factor,
    topo.terrain_config_factor,
    0.5,
    surface_albedo=0.5)

# Plot
fig, ax = plt.subplots(2, 3, figsize = (30, 20))
fig.subplots_adjust(left=0.02, bottom=0.06, right=0.95, top=0.94, wspace=0.05)
fig.suptitle('TOPORAD')

# IPW originals
im0 = ax[0, 0].imshow(toporad_ipw.bands[0].data)
ax[0, 0].set_title('IPW beam')
fig.colorbar(im0, ax=ax[0, 0])

im0 = ax[1, 0].imshow(toporad_ipw.bands[1].data)
ax[1, 0].set_title('IPW diffuse')
fig.colorbar(im0, ax=ax[1, 0])

# Python elevrad
im0 = ax[0, 1].imshow(trad_beam)
ax[0, 1].set_title('Python beam')
fig.colorbar(im0, ax=ax[0, 1])

im0 = ax[1, 1].imshow(trad_diffuse)
ax[1, 1].set_title('Python diffuse')
fig.colorbar(im0, ax=ax[1, 1])

# difference
d = toporad_ipw.bands[0].data - trad_beam
var_lim = max([abs(d.min()), abs(d.max())])
# im2 = ax[0, 2].imshow(d, cmap='RdBu', vmin=-var_lim, vmax=var_lim)
# fig.colorbar(im2, ax=ax[0, 2])
ax[0, 2].hist(d.flatten(), bins=100)
ax[0, 2].yaxis.tick_right()
ax[0, 2].set_title('Beam difference, ipw - python')

d = toporad_ipw.bands[1].data - trad_diffuse
var_lim = max([abs(d.min()), abs(d.max())])
# im2 = ax[1, 2].imshow(d, cmap='RdBu', vmin=-var_lim, vmax=var_lim)
# fig.colorbar(im2, ax=ax[1, 2])
ax[1, 2].hist(d.flatten(), bins=100)
ax[1, 2].yaxis.tick_right()
ax[1, 2].set_title('Diffuse difference, ipw - python')

# remove the x and y ticks
for a in ax[:,:2].flatten():
    a.set_xticks([])
    a.set_yticks([])

# +
# stoporad for visible

# IPW stoporad
cmd_vis = f"stoporad -z {tau_elevation} -t {tau} -w {omega} -g {scattering_factor} -x {wavelength_vis[0]},{wavelength_vis[1]} -s {wy_day - min_storm_day} " \
        f"-d {wy_day} -f {tz_min_west} -y {wyear} -A {cosz},{azimuth} -a {grain_size} -m {max_grain} -c {dirt} -D stoporad_in.ipw > stoporad_vis.ipw"

print(cmd_vis)
visp = sp.Popen(cmd_vis, shell=True)
visp.wait()

stoporad_ipw = ipw.IPW('stoporad_vis.ipw')

# Python stoporad
srad_beam, srad_diffuse = toporad.stoporad_ipw(
    date_time,
    tau_elevation,
    tau,
    omega,
    scattering_factor,
    wavelength_range=wavelength_vis,
    start=min_storm_day,
    current_day=wy_day,
    time_zone=tz_min_west,
    year=wyear,
    cosz=cosz,
    azimuth=azimuth,
    grain_size=grain_size,
    max_grain=max_grain,
    dirt=dirt,
    topo=topo)

# Plot
fig, ax = plt.subplots(2, 3, figsize = (30, 20))
fig.subplots_adjust(left=0.02, bottom=0.06, right=0.95, top=0.94, wspace=0.05)
fig.suptitle('STOPORAD VIS')

# IPW originals
im0 = ax[0, 0].imshow(stoporad_ipw.bands[0].data)
ax[0, 0].set_title('IPW beam')
fig.colorbar(im0, ax=ax[0, 0])

im0 = ax[1, 0].imshow(stoporad_ipw.bands[1].data)
ax[1, 0].set_title('IPW diffuse')
fig.colorbar(im0, ax=ax[1, 0])

# Python elevrad
im0 = ax[0, 1].imshow(srad_beam)
ax[0, 1].set_title('Python beam')
fig.colorbar(im0, ax=ax[0, 1])

im0 = ax[1, 1].imshow(srad_diffuse)
ax[1, 1].set_title('Python diffuse')
fig.colorbar(im0, ax=ax[1, 1])

# difference
d = stoporad_ipw.bands[0].data - srad_beam
var_lim = max([abs(d.min()), abs(d.max())])
# im2 = ax[0, 2].imshow(d, cmap='RdBu', vmin=-var_lim, vmax=var_lim)
# fig.colorbar(im2, ax=ax[0, 2])
ax[0, 2].hist(d.flatten(), bins=100)
ax[0, 2].yaxis.tick_right()
ax[0, 2].set_title('Beam difference, ipw - python')

d = stoporad_ipw.bands[1].data - srad_diffuse
var_lim = max([abs(d.min()), abs(d.max())])
# im2 = ax[1, 2].imshow(d, cmap='RdBu', vmin=-var_lim, vmax=var_lim)
# fig.colorbar(im2, ax=ax[1, 2])
ax[1, 2].hist(d.flatten(), bins=100)
ax[1, 2].yaxis.tick_right()
ax[1, 2].set_title('Diffuse difference, ipw - python')

# remove the x and y ticks
for a in ax[:,:2].flatten():
    a.set_xticks([])
    a.set_yticks([])

# +
# stoporad for IR

# IPW stoporad
cmd_vis = f"stoporad -z {tau_elevation} -t {tau} -w {omega} -g {scattering_factor} -x {wavelength_ir[0]},{wavelength_ir[1]} -s {wy_day - min_storm_day} " \
        f"-d {wy_day} -f {tz_min_west} -y {wyear} -A {cosz},{azimuth} -a {grain_size} -m {max_grain} -c {dirt} -D stoporad_in.ipw > stoporad_ir.ipw"

print(cmd_vis)
visp = sp.Popen(cmd_vis, shell=True)
visp.wait()

stoporad_ipw = ipw.IPW('stoporad_ir.ipw')

# Python stoporad
srad_beam, srad_diffuse = toporad.stoporad_ipw(
    date_time,
    tau_elevation,
    tau,
    omega,
    scattering_factor,
    wavelength_range=wavelength_ir,
    start=min_storm_day,
    current_day=wy_day,
    time_zone=tz_min_west,
    year=wyear,
    cosz=cosz,
    azimuth=azimuth,
    grain_size=grain_size,
    max_grain=max_grain,
    dirt=dirt,
    topo=topo)

# Plot
fig, ax = plt.subplots(2, 3, figsize = (30, 20))
fig.subplots_adjust(left=0.02, bottom=0.06, right=0.95, top=0.94, wspace=0.05)
fig.suptitle('STOPORAD IR')

# IPW originals
im0 = ax[0, 0].imshow(stoporad_ipw.bands[0].data)
ax[0, 0].set_title('IPW beam')
fig.colorbar(im0, ax=ax[0, 0])

im0 = ax[1, 0].imshow(stoporad_ipw.bands[1].data)
ax[1, 0].set_title('IPW diffuse')
fig.colorbar(im0, ax=ax[1, 0])

# Python elevrad
im0 = ax[0, 1].imshow(srad_beam)
ax[0, 1].set_title('Python beam')
fig.colorbar(im0, ax=ax[0, 1])

im0 = ax[1, 1].imshow(srad_diffuse)
ax[1, 1].set_title('Python diffuse')
fig.colorbar(im0, ax=ax[1, 1])

# difference
d = stoporad_ipw.bands[0].data - srad_beam
var_lim = max([abs(d.min()), abs(d.max())])
# im2 = ax[0, 2].imshow(d, cmap='RdBu', vmin=-var_lim, vmax=var_lim)
# fig.colorbar(im2, ax=ax[0, 2])
ax[0, 2].hist(d.flatten(), bins=100)
ax[0, 2].yaxis.tick_right()
ax[0, 2].set_title('Beam difference, ipw - python')

d = stoporad_ipw.bands[1].data - srad_diffuse
var_lim = max([abs(d.min()), abs(d.max())])
# im2 = ax[1, 2].imshow(d, cmap='RdBu', vmin=-var_lim, vmax=var_lim)
# fig.colorbar(im2, ax=ax[1, 2])
ax[1, 2].hist(d.flatten(), bins=100)
ax[1, 2].yaxis.tick_right()
ax[1, 2].set_title('Diffuse difference, ipw - python')

# remove the x and y ticks
for a in ax[:,:2].flatten():
    a.set_xticks([])
    a.set_yticks([])


# +
# New version of stoporad that uses values already computed in SMRF -- visible

srad_beam, srad_diffuse = toporad.stoporad(
    date_time,
    topo,
    cosz,
    azimuth,
    illum_ang,
    alb_vis,
    wavelength_range=wavelength_vis,
    tau_elevation=tau_elevation,
    tau=tau,
    omega=omega,
    scattering_factor=scattering_factor)

# Plot
stoporad_ipw = ipw.IPW('stoporad_vis.ipw')

fig, ax = plt.subplots(2, 3, figsize = (30, 20))
fig.subplots_adjust(left=0.02, bottom=0.06, right=0.95, top=0.94, wspace=0.05)
fig.suptitle('STOPORAD VIS')

# IPW originals
im0 = ax[0, 0].imshow(stoporad_ipw.bands[0].data)
ax[0, 0].set_title('IPW beam')
fig.colorbar(im0, ax=ax[0, 0])

im0 = ax[1, 0].imshow(stoporad_ipw.bands[1].data)
ax[1, 0].set_title('IPW diffuse')
fig.colorbar(im0, ax=ax[1, 0])

# Python elevrad
im0 = ax[0, 1].imshow(srad_beam)
ax[0, 1].set_title('Python beam')
fig.colorbar(im0, ax=ax[0, 1])

im0 = ax[1, 1].imshow(srad_diffuse)
ax[1, 1].set_title('Python diffuse')
fig.colorbar(im0, ax=ax[1, 1])

# difference
d = stoporad_ipw.bands[0].data - srad_beam
var_lim = max([abs(d.min()), abs(d.max())])
# im2 = ax[0, 2].imshow(d, cmap='RdBu', vmin=-var_lim, vmax=var_lim)
# fig.colorbar(im2, ax=ax[0, 2])
ax[0, 2].hist(d.flatten(), bins=100)
ax[0, 2].yaxis.tick_right()
ax[0, 2].set_title('Beam difference, ipw - python')

d = stoporad_ipw.bands[1].data - srad_diffuse
var_lim = max([abs(d.min()), abs(d.max())])
# im2 = ax[1, 2].imshow(d, cmap='RdBu', vmin=-var_lim, vmax=var_lim)
# fig.colorbar(im2, ax=ax[1, 2])
ax[1, 2].hist(d.flatten(), bins=100)
ax[1, 2].yaxis.tick_right()
ax[1, 2].set_title('Diffuse difference, ipw - python')

# remove the x and y ticks
for a in ax[:,:2].flatten():
    a.set_xticks([])
    a.set_yticks([])

# +
# New version of stoporad that uses values already computed in SMRF -- IR

srad_beam, srad_diffuse = toporad.stoporad(
    date_time,
    topo,
    cosz,
    azimuth,
    illum_ang,
    alb_ir,
    wavelength_range=wavelength_ir,
    tau_elevation=tau_elevation,
    tau=tau,
    omega=omega,
    scattering_factor=scattering_factor)

# Plot
stoporad_ipw = ipw.IPW('stoporad_ir.ipw')

fig, ax = plt.subplots(2, 3, figsize = (30, 20))
fig.subplots_adjust(left=0.02, bottom=0.06, right=0.95, top=0.94, wspace=0.05)
fig.suptitle('STOPORAD VIS')

# IPW originals
im0 = ax[0, 0].imshow(stoporad_ipw.bands[0].data)
ax[0, 0].set_title('IPW beam')
fig.colorbar(im0, ax=ax[0, 0])

im0 = ax[1, 0].imshow(stoporad_ipw.bands[1].data)
ax[1, 0].set_title('IPW diffuse')
fig.colorbar(im0, ax=ax[1, 0])

# Python elevrad
im0 = ax[0, 1].imshow(srad_beam)
ax[0, 1].set_title('Python beam')
fig.colorbar(im0, ax=ax[0, 1])

im0 = ax[1, 1].imshow(srad_diffuse)
ax[1, 1].set_title('Python diffuse')
fig.colorbar(im0, ax=ax[1, 1])

# difference
d = stoporad_ipw.bands[0].data - srad_beam
var_lim = max([abs(d.min()), abs(d.max())])
# im2 = ax[0, 2].imshow(d, cmap='RdBu', vmin=-var_lim, vmax=var_lim)
# fig.colorbar(im2, ax=ax[0, 2])
ax[0, 2].hist(d.flatten(), bins=100)
ax[0, 2].yaxis.tick_right()
ax[0, 2].set_title('Beam difference, ipw - python')

d = stoporad_ipw.bands[1].data - srad_diffuse
var_lim = max([abs(d.min()), abs(d.max())])
# im2 = ax[1, 2].imshow(d, cmap='RdBu', vmin=-var_lim, vmax=var_lim)
# fig.colorbar(im2, ax=ax[1, 2])
ax[1, 2].hist(d.flatten(), bins=100)
ax[1, 2].yaxis.tick_right()
ax[1, 2].set_title('Diffuse difference, ipw - python')

# remove the x and y ticks
for a in ax[:,:2].flatten():
    a.set_xticks([])
    a.set_yticks([])

# +
# Perform some timing, 100 simulations
# IPW average time per run 0.68270476 sec
# Python average time per run 0.03738955 sec

sims = 100

cmd_vis = f"stoporad -z {tau_elevation} -t {tau} -w {omega} -g {scattering_factor} -x {wavelength_ir[0]},{wavelength_ir[1]} -s {wy_day - min_storm_day} " \
        f"-d {wy_day} -f {tz_min_west} -y {wyear} -A {cosz},{azimuth} -a {grain_size} -m {max_grain} -c {dirt} -D stoporad_in.ipw > stoporad_ir.ipw"

tstart = datetime.now()
for i in range(sims):
    visp = sp.Popen(cmd_vis, shell=True)
    visp.wait()
telapsed_ipw = datetime.now() - tstart
print(f"IPW average time per run {telapsed_ipw.total_seconds()/sims} sec")

tstart = datetime.now()
for i in range(100):
    srad_beam, srad_diffuse = toporad.stoporad(
        date_time,
        topo,
        cosz,
        azimuth,
        illum_ang,
        alb_vis,
        wavelength_range=wavelength_vis,
        tau_elevation=tau_elevation,
        tau=tau,
        omega=omega,
        scattering_factor=scattering_factor)
telapsed_python = datetime.now() - tstart
print(f"Python average time per run {telapsed_python.total_seconds()/sims} sec")
# -

# clean up IPW images
for fileName in glob.glob("*.ipw"):
  os.remove(fileName)

StarcoderdataJupyterScriptsDedupFiltered