#!/usr/bin/env python
# coding: utf-8

# In[ ]:


# https://gml.noaa.gov/grad/solcalc/solareqns.PDF

# https://github.com/SatAgro/suntime

# https://stackoverflow.com/questions/19615350/calculate-sunrise-and-sunset-times-for-a-given-gps-coordinate-within-postgresql


# In[ ]:


import math
import numpy as np
import datetime
import pytz
import xarray as xr
from numba import vectorize


# In[ ]:

@vectorize
def rise_set_func(lat, lon, day, mo, yr):
        
# INPUTS
#
# lat: float32, latitude [degrees] of point of interest
# lon: float32, longitude [degrees] of point of interest
# day: integer, day of month
# mo: integer, month of year
# yr: integer, year of interest

# OUTPUTS
#
# sunrise: float32, UTC hour at which sunrise occurs at point of interest
# sunset: float32, UTC hour at which sunset occurs at point of interest
    
    
    # set UTC second, minute, hours to 0 for data point
    sc, mn, hr = 0, 0, 0
            
            
    # (1) CALCULATE FRACTIONAL YEAR, g [radians]    
            
    # number of days in a year is 365, unless it is a leap year -> 366
    nd = 365
    if yr%4 == 0:
        nd+=1
    #nd = 366 - yr%4%2
        
    # calculate day of year from datetime
    d1 = math.floor( 275*mo / 9)
    d2 = math.floor( (mo+9) / 12)
    d3 = ( 1 + math.floor( (yr - 4*math.floor(yr/4) + 2) / 3) )
        
    d = d1 - (d2*d3) + day - 30
        
    # get g
    #print(f'{nd = }, {hr = }')
    g = (2*math.pi/nd) * (d - 1 + (hr - 12)/24 )
        
        
    # (2) ESTIMATE EQUATION OF TIME, eqt [minutes]
    eqt = 229.18*(0.000075 + 0.001868*math.cos(g) - 0.032077*math.sin(g)                     - 0.014615*math.cos(2*g) - 0.040849*math.sin(2*g))
        
        
    # (3) ESTIMATE SOLAR DECLINATION ANGLE, sd [radians]
    sd = 0.006918 - 0.399912*math.cos(g) + 0.070257*math.sin(g)             - 0.006758*math.cos(2*g) + 0.000907*math.sin(2*g)                 - 0.002697*math.cos(3*g) + 0.00148*math.sin(3*g)
        
        
    # (4) GET SOLAR HOUR ANGLE, ha, AT SUNRISE/SUNSET [degrees]
        
    # zenith at sunrise/sunset [degrees]
    phi = 90.833
    # conversions from degrees to radians
    rphi = (math.pi/180)*phi
    rlat = (math.pi/180)*lat
        
    ha = np.arccos(math.cos(rphi)/(math.cos(rlat)*math.cos(sd))                     - math.tan(rlat)*math.tan(sd))
    ha = (180/math.pi)*ha

    # (6) GET UTC TIME OF SUNRISE AND SUNSET AT POINT OF INTEREST [minutes]
    sunrise = (720 - 4*(lon + ha) - eqt)/60
    sunset = (720 - 4*(lon - ha) - eqt)/60
        
    return sunrise#, sunset
    


# In[ ]:


def rise_set(lat, lon, time):
        
# INPUTS
#
# lat: float32 xr.DataArray, latitudes of interest
# lon: float32 xr.DataArray, longitudes of interest
# time: datetime64[ns] xr.DataArray, times of interest 
        
    day = time.dt.day
    mo = time.dt.month
    yr = time.dt.year
        
    return xr.apply_ufunc(
        rise_set_func,
        lat, lon, day, mo, yr,
        dask="parallelized"
    )