"""insolation.py

This module contains general-purpose routines for computing incoming
solar radiation at the top of the atmosphere.

Currently, only daily average insolation is computed.

Ported and modified from MATLAB code daily_insolation.m
Original authors:
    Ian Eisenman and Peter Huybers, Harvard University, August 2006
Available online at http://eisenman.ucsd.edu/code/daily_insolation.m

If using calendar days, solar longitude is found using an
approximate solution to the differential equation representing conservation
of angular momentum (Kepler's Second Law).  Given the orbital parameters
and solar longitude, daily average insolation is calculated exactly
following Berger 1978.

References:
Berger A. and Loutre M.F. (1991). Insolation values for the climate of
 the last 10 million years. Quaternary Science Reviews, 10(4), 297-317.
Berger A. (1978). Long-term variations of daily insolation and
 Quaternary climatic changes. Journal of Atmospheric Science, 35(12),
 2362-2367.
"""

import numpy as np
from climlab import constants as const


def daily_insolation(lat, day, orb=const.orb_present, S0=None, day_type=1):
    """Compute daily average insolation given latitude, time of year and orbital parameters.

    Orbital parameters can be computed for any time in the last 5 Myears with
    ecc,long_peri,obliquity = orbital.lookup_parameters(kyears)

    Inputs:
    lat:      Latitude in degrees (-90 to 90).
    day:      Indicator of time of year, by default day 1 is Jan 1.
    orb:    a dictionary with three members (as provided by orbital.py)
        ecc:      eccentricity (dimensionless)
        long_peri:    longitude of perihelion (precession angle) (degrees)
        obliquity:  obliquity angle (degrees)
    S0:       Solar constant in W/m^2, will try to read from constants.py
    day_type: Convention for specifying time of year (+/- 1,2) [optional].
        day_type=1 (default): day input is calendar day (1-365.24), where day 1
        is January first.  The calendar is referenced to the vernal equinox
        which always occurs at day 80.
        day_type=2: day input is solar longitude (0-360 degrees). Solar
        longitude is the angle of the Earth's orbit measured from spring
        equinox (21 March). Note that calendar days and solar longitude are
        not linearly related because, by Kepler's Second Law, Earth's
        angular velocity varies according to its distance from the sun.
    Default values for orbital parameters are present-day

    Output:
    Fsw = Daily average solar radiation in W/m^2.

    Dimensions of output are (lat.size, day.size, ecc.size)

    Code is fully vectorized to handle array input for all arguments.
    Orbital arguments should all have the same sizes.
    This is automatic if computed from orbital.OrbitalTable.lookup_parameters()

    e.g. to compute the timeseries of insolation at 65N at summer solstice over the past 5 Myears
        from climlab.orbital import OrbitalTable
        table = OrbitalTable()
        years = np.linspace(0, 5000, 5001)
        orb = table.lookup_parameters( years )
        S65 = orbital.daily_insolation( 65, 172, orb )
     """

    # If input argument S0 is not given, use the standard Earth value
    if S0 is None:
        S0 = const.S0
    
    # Inputs can be scalar or vector. If scalar, convert to 0d numpy arrays
    lat = np.array( lat )
    day = np.array( day )
    ecc = np.array( orb['ecc'] )
    long_peri = np.array( orb['long_peri'] )
    obliquity = np.array( orb['obliquity'] )
    
    # Convert precession angle and latitude to radians
    phi = np.deg2rad( lat )
    
    # lambda_long (solar longitude) is the angular distance along Earth's orbit measured from spring equinox (21 March)
    if day_type==1: # calendar days
        lambda_long = solar_longitude( day, orb )
    elif day_type==2: #solar longitude (1-360) is specified in input, no need to convert days to longitude
        lambda_long = np.tile( np.expand_dims( np.deg2rad( day ), axis=1 ), [1, ecc.size] )  
    else:
        raise ValueError('Invalid day_type.')
        
    # Compute declination angle of the sun
    delta = np.array( np.arcsin( np.sin( np.deg2rad(obliquity) ) * np.sin( lambda_long ) ) ) 
    
    # recast all the arrays to dimensions (lat.size, day.size, ecc.size)
    delta_big = np.tile( delta, [ lat.size, 1, 1 ] )
    lambda_long_big = np.tile( lambda_long, [ lat.size, 1, 1 ] )
    phi_day = np.tile( np.expand_dims( phi, axis=1 ), [1, day.size] )
    phi_big = np.tile( np.expand_dims( phi_day, axis=2 ), [1, 1, ecc.size] )
    
    #  suppress warning message generated by arccos here!
    oldsettings = np.seterr(invalid='ignore')
    # Compute Ho, the hour angle at sunrise / sunset
    #  Check for no sunrise or no sunset: Berger 1978 eqn (8),(9)
    Ho = np.where( abs( delta_big ) - np.math.pi / 2. + abs( phi_big ) < 0., 
                  np.arccos( -np.tan( phi_big ) * np.tan( delta_big ) ), 
            np.where( phi_big * delta_big > 0. , np.math.pi, 0. ) )
    # this is not really the daily average cosine of the zenith angle...
    #  it's the integral from sunrise to sunset of that quantity...
    coszen = (Ho*np.sin(phi_big)*np.sin(delta_big) + 
              np.cos(phi_big)*np.cos(delta_big)*np.sin(Ho))
    # Compute insolation: Berger 1978 eq (10)
    Fsw = S0/np.math.pi*( (1. + ecc*np.cos(lambda_long_big - 
          np.deg2rad(long_peri)))**2 / (1. - ecc**2)**2 * coszen)
            
    #  Remove singleton dimensions and return
    return np.squeeze( Fsw )


def solar_longitude( day, orb=const.orb_present, days_per_year = None ):
    '''Estimate solar longitude (lambda = 0 at spring equinox) from calendar day 
    using an approximation from Berger 1978 section 3.
    
    Works for both scalar and vector orbital parameters.
    
    Reads the length of the year from constants.py if available.
    '''
    
    if days_per_year is None:
        days_per_year = const.days_per_year
    
    day = np.array(day)
    ecc = np.array(orb['ecc'])
    long_peri_rad = np.deg2rad( np.array(orb['long_peri']) )
    delta_lambda_long_m = ( ( np.tile( np.expand_dims( day, axis=1 ), [1, ecc.size]) - 80. ) * 2. *
        np.math.pi / days_per_year )
    beta = ( 1 - ecc**2 )**(1./2.)
    lambda_long_m0 = ( -2. * ( (1./2. * ecc + 1./8. * ecc**3 ) * (1. + beta) * np.sin(-long_peri_rad) - 
        1./4.* (ecc**2) * ( 1./2. + beta ) * np.sin( -2. * long_peri_rad ) + 1./8. * (ecc**3) * 
        (1./3. + beta) * ( np.sin(-3. * long_peri_rad) ) ) )
    lambda_long_m = np.tile( lambda_long_m0, [day.size, 1] ) + delta_lambda_long_m
    lambda_long = ( lambda_long_m + ( 2*ecc - 1/4. * (ecc**3)) * np.sin(lambda_long_m - long_peri_rad) + 
        (5./4.) * (ecc**2) * np.sin(2 * ( lambda_long_m - long_peri_rad )) + (13./12.) * (ecc**3) 
        * np.sin(3*( lambda_long_m - long_peri_rad )) )
    
    # lambda_long should have dimensions lambda_long.size = ( day.size, ecc.size )
    return lambda_long

