"""
Calculate the solar position using the NREL SPA algorithm either using
numpy arrays or compiling the code to machine language with numba.
"""

# Contributors:
# Created by Tony Lorenzo (@alorenzo175), Univ. of Arizona, 2015

from __future__ import division
import os
import threading
import warnings

import numpy as np

# this block is a way to use an environment variable to switch between
# compiling the functions with numba or just use numpy
def nocompile(*args, **kwargs):
    return lambda func: func

if os.getenv('PVLIB_USE_NUMBA', '0') != '0':
    try:
        from numba import jit, __version__
    except ImportError:
        warnings.warn('Could not import numba, falling back to numpy ' +
                      'calculation')
        jcompile = nocompile
        USE_NUMBA = False
    else:
        major, minor = __version__.split('.')[:2]
        if int(major + minor) >= 17:
            # need at least numba >= 0.17.0
            jcompile = jit
            USE_NUMBA = True
        else:
            warnings.warn('Numba version must be >= 0.17.0, falling back to ' +
                          'numpy')
            jcompile = nocompile
            USE_NUMBA = False
else:
    jcompile = nocompile
    USE_NUMBA = False

TABLE_1_DICT = {
    'L0': np.array(
        [[175347046.0, 0.0, 0.0],
         [3341656.0, 4.6692568, 6283.07585],
         [34894.0, 4.6261, 12566.1517],
         [3497.0, 2.7441, 5753.3849],
         [3418.0, 2.8289, 3.5231],
         [3136.0, 3.6277, 77713.7715],
         [2676.0, 4.4181, 7860.4194],
         [2343.0, 6.1352, 3930.2097],
         [1324.0, 0.7425, 11506.7698],
         [1273.0, 2.0371, 529.691],
         [1199.0, 1.1096, 1577.3435],
         [990.0, 5.233, 5884.927],
         [902.0, 2.045, 26.298],
         [857.0, 3.508, 398.149],
         [780.0, 1.179, 5223.694],
         [753.0, 2.533, 5507.553],
         [505.0, 4.583, 18849.228],
         [492.0, 4.205, 775.523],
         [357.0, 2.92, 0.067],
         [317.0, 5.849, 11790.629],
         [284.0, 1.899, 796.298],
         [271.0, 0.315, 10977.079],
         [243.0, 0.345, 5486.778],
         [206.0, 4.806, 2544.314],
         [205.0, 1.869, 5573.143],
         [202.0, 2.458, 6069.777],
         [156.0, 0.833, 213.299],
         [132.0, 3.411, 2942.463],
         [126.0, 1.083, 20.775],
         [115.0, 0.645, 0.98],
         [103.0, 0.636, 4694.003],
         [102.0, 0.976, 15720.839],
         [102.0, 4.267, 7.114],
         [99.0, 6.21, 2146.17],
         [98.0, 0.68, 155.42],
         [86.0, 5.98, 161000.69],
         [85.0, 1.3, 6275.96],
         [85.0, 3.67, 71430.7],
         [80.0, 1.81, 17260.15],
         [79.0, 3.04, 12036.46],
         [75.0, 1.76, 5088.63],
         [74.0, 3.5, 3154.69],
         [74.0, 4.68, 801.82],
         [70.0, 0.83, 9437.76],
         [62.0, 3.98, 8827.39],
         [61.0, 1.82, 7084.9],
         [57.0, 2.78, 6286.6],
         [56.0, 4.39, 14143.5],
         [56.0, 3.47, 6279.55],
         [52.0, 0.19, 12139.55],
         [52.0, 1.33, 1748.02],
         [51.0, 0.28, 5856.48],
         [49.0, 0.49, 1194.45],
         [41.0, 5.37, 8429.24],
         [41.0, 2.4, 19651.05],
         [39.0, 6.17, 10447.39],
         [37.0, 6.04, 10213.29],
         [37.0, 2.57, 1059.38],
         [36.0, 1.71, 2352.87],
         [36.0, 1.78, 6812.77],
         [33.0, 0.59, 17789.85],
         [30.0, 0.44, 83996.85],
         [30.0, 2.74, 1349.87],
         [25.0, 3.16, 4690.48]]),
    'L1': np.array(
        [[628331966747.0, 0.0, 0.0],
         [206059.0, 2.678235, 6283.07585],
         [4303.0, 2.6351, 12566.1517],
         [425.0, 1.59, 3.523],
         [119.0, 5.796, 26.298],
         [109.0, 2.966, 1577.344],
         [93.0, 2.59, 18849.23],
         [72.0, 1.14, 529.69],
         [68.0, 1.87, 398.15],
         [67.0, 4.41, 5507.55],
         [59.0, 2.89, 5223.69],
         [56.0, 2.17, 155.42],
         [45.0, 0.4, 796.3],
         [36.0, 0.47, 775.52],
         [29.0, 2.65, 7.11],
         [21.0, 5.34, 0.98],
         [19.0, 1.85, 5486.78],
         [19.0, 4.97, 213.3],
         [17.0, 2.99, 6275.96],
         [16.0, 0.03, 2544.31],
         [16.0, 1.43, 2146.17],
         [15.0, 1.21, 10977.08],
         [12.0, 2.83, 1748.02],
         [12.0, 3.26, 5088.63],
         [12.0, 5.27, 1194.45],
         [12.0, 2.08, 4694.0],
         [11.0, 0.77, 553.57],
         [10.0, 1.3, 6286.6],
         [10.0, 4.24, 1349.87],
         [9.0, 2.7, 242.73],
         [9.0, 5.64, 951.72],
         [8.0, 5.3, 2352.87],
         [6.0, 2.65, 9437.76],
         [6.0, 4.67, 4690.48]]),
    'L2': np.array(
        [[52919.0, 0.0, 0.0],
         [8720.0, 1.0721, 6283.0758],
         [309.0, 0.867, 12566.152],
         [27.0, 0.05, 3.52],
         [16.0, 5.19, 26.3],
         [16.0, 3.68, 155.42],
         [10.0, 0.76, 18849.23],
         [9.0, 2.06, 77713.77],
         [7.0, 0.83, 775.52],
         [5.0, 4.66, 1577.34],
         [4.0, 1.03, 7.11],
         [4.0, 3.44, 5573.14],
         [3.0, 5.14, 796.3],
         [3.0, 6.05, 5507.55],
         [3.0, 1.19, 242.73],
         [3.0, 6.12, 529.69],
         [3.0, 0.31, 398.15],
         [3.0, 2.28, 553.57],
         [2.0, 4.38, 5223.69],
         [2.0, 3.75, 0.98]]),
    'L3': np.array(
        [[289.0, 5.844, 6283.076],
         [35.0, 0.0, 0.0],
         [17.0, 5.49, 12566.15],
         [3.0, 5.2, 155.42],
         [1.0, 4.72, 3.52],
         [1.0, 5.3, 18849.23],
         [1.0, 5.97, 242.73]]),
    'L4': np.array(
        [[114.0, 3.142, 0.0],
         [8.0, 4.13, 6283.08],
         [1.0, 3.84, 12566.15]]),
    'L5': np.array(
        [[1.0, 3.14, 0.0]]),
    'B0': np.array(
        [[280.0, 3.199, 84334.662],
         [102.0, 5.422, 5507.553],
         [80.0, 3.88, 5223.69],
         [44.0, 3.7, 2352.87],
         [32.0, 4.0, 1577.34]]),
    'B1': np.array(
        [[9.0, 3.9, 5507.55],
         [6.0, 1.73, 5223.69]]),
    'R0': np.array(
        [[100013989.0, 0.0, 0.0],
         [1670700.0, 3.0984635, 6283.07585],
         [13956.0, 3.05525, 12566.1517],
         [3084.0, 5.1985, 77713.7715],
         [1628.0, 1.1739, 5753.3849],
         [1576.0, 2.8469, 7860.4194],
         [925.0, 5.453, 11506.77],
         [542.0, 4.564, 3930.21],
         [472.0, 3.661, 5884.927],
         [346.0, 0.964, 5507.553],
         [329.0, 5.9, 5223.694],
         [307.0, 0.299, 5573.143],
         [243.0, 4.273, 11790.629],
         [212.0, 5.847, 1577.344],
         [186.0, 5.022, 10977.079],
         [175.0, 3.012, 18849.228],
         [110.0, 5.055, 5486.778],
         [98.0, 0.89, 6069.78],
         [86.0, 5.69, 15720.84],
         [86.0, 1.27, 161000.69],
         [65.0, 0.27, 17260.15],
         [63.0, 0.92, 529.69],
         [57.0, 2.01, 83996.85],
         [56.0, 5.24, 71430.7],
         [49.0, 3.25, 2544.31],
         [47.0, 2.58, 775.52],
         [45.0, 5.54, 9437.76],
         [43.0, 6.01, 6275.96],
         [39.0, 5.36, 4694.0],
         [38.0, 2.39, 8827.39],
         [37.0, 0.83, 19651.05],
         [37.0, 4.9, 12139.55],
         [36.0, 1.67, 12036.46],
         [35.0, 1.84, 2942.46],
         [33.0, 0.24, 7084.9],
         [32.0, 0.18, 5088.63],
         [32.0, 1.78, 398.15],
         [28.0, 1.21, 6286.6],
         [28.0, 1.9, 6279.55],
         [26.0, 4.59, 10447.39]]),
    'R1': np.array(
        [[103019.0, 1.10749, 6283.07585],
         [1721.0, 1.0644, 12566.1517],
         [702.0, 3.142, 0.0],
         [32.0, 1.02, 18849.23],
         [31.0, 2.84, 5507.55],
         [25.0, 1.32, 5223.69],
         [18.0, 1.42, 1577.34],
         [10.0, 5.91, 10977.08],
         [9.0, 1.42, 6275.96],
         [9.0, 0.27, 5486.78]]),
    'R2': np.array(
        [[4359.0, 5.7846, 6283.0758],
         [124.0, 5.579, 12566.152],
         [12.0, 3.14, 0.0],
         [9.0, 3.63, 77713.77],
         [6.0, 1.87, 5573.14],
         [3.0, 5.47, 18849.23]]),
    'R3': np.array(
        [[145.0, 4.273, 6283.076],
         [7.0, 3.92, 12566.15]]),
    'R4': np.array(
        [[4.0, 2.56, 6283.08]])
}

TABLE_1_DICT['L1'].resize((64, 3))
TABLE_1_DICT['L2'].resize((64, 3))
TABLE_1_DICT['L3'].resize((64, 3))
TABLE_1_DICT['L4'].resize((64, 3))
TABLE_1_DICT['L5'].resize((64, 3))

TABLE_1_DICT['B1'].resize((5, 3))

TABLE_1_DICT['R1'].resize((40, 3))
TABLE_1_DICT['R2'].resize((40, 3))
TABLE_1_DICT['R3'].resize((40, 3))
TABLE_1_DICT['R4'].resize((40, 3))

HELIO_LONG_TABLE = np.array([TABLE_1_DICT['L0'],
                             TABLE_1_DICT['L1'],
                             TABLE_1_DICT['L2'],
                             TABLE_1_DICT['L3'],
                             TABLE_1_DICT['L4'],
                             TABLE_1_DICT['L5']])

HELIO_LAT_TABLE = np.array([TABLE_1_DICT['B0'],
                            TABLE_1_DICT['B1']])

HELIO_RADIUS_TABLE = np.array([TABLE_1_DICT['R0'],
                               TABLE_1_DICT['R1'],
                               TABLE_1_DICT['R2'],
                               TABLE_1_DICT['R3'],
                               TABLE_1_DICT['R4']])

NUTATION_ABCD_ARRAY = np.array([
    [-171996, -174.2, 92025, 8.9],
    [-13187, -1.6, 5736, -3.1],
    [-2274, -0.2, 977, -0.5],
    [2062, 0.2, -895, 0.5],
    [1426, -3.4, 54, -0.1],
    [712, 0.1, -7, 0],
    [-517, 1.2, 224, -0.6],
    [-386, -0.4, 200, 0],
    [-301, 0, 129, -0.1],
    [217, -0.5, -95, 0.3],
    [-158, 0, 0, 0],
    [129, 0.1, -70, 0],
    [123, 0, -53, 0],
    [63, 0, 0, 0],
    [63, 0.1, -33, 0],
    [-59, 0, 26, 0],
    [-58, -0.1, 32, 0],
    [-51, 0, 27, 0],
    [48, 0, 0, 0],
    [46, 0, -24, 0],
    [-38, 0, 16, 0],
    [-31, 0, 13, 0],
    [29, 0, 0, 0],
    [29, 0, -12, 0],
    [26, 0, 0, 0],
    [-22, 0, 0, 0],
    [21, 0, -10, 0],
    [17, -0.1, 0, 0],
    [16, 0, -8, 0],
    [-16, 0.1, 7, 0],
    [-15, 0, 9, 0],
    [-13, 0, 7, 0],
    [-12, 0, 6, 0],
    [11, 0, 0, 0],
    [-10, 0, 5, 0],
    [-8, 0, 3, 0],
    [7, 0, -3, 0],
    [-7, 0, 0, 0],
    [-7, 0, 3, 0],
    [-7, 0, 3, 0],
    [6, 0, 0, 0],
    [6, 0, -3, 0],
    [6, 0, -3, 0],
    [-6, 0, 3, 0],
    [-6, 0, 3, 0],
    [5, 0, 0, 0],
    [-5, 0, 3, 0],
    [-5, 0, 3, 0],
    [-5, 0, 3, 0],
    [4, 0, 0, 0],
    [4, 0, 0, 0],
    [4, 0, 0, 0],
    [-4, 0, 0, 0],
    [-4, 0, 0, 0],
    [-4, 0, 0, 0],
    [3, 0, 0, 0],
    [-3, 0, 0, 0],
    [-3, 0, 0, 0],
    [-3, 0, 0, 0],
    [-3, 0, 0, 0],
    [-3, 0, 0, 0],
    [-3, 0, 0, 0],
    [-3, 0, 0, 0],
])

NUTATION_YTERM_ARRAY = np.array([
    [0, 0, 0, 0, 1],
    [-2, 0, 0, 2, 2],
    [0, 0, 0, 2, 2],
    [0, 0, 0, 0, 2],
    [0, 1, 0, 0, 0],
    [0, 0, 1, 0, 0],
    [-2, 1, 0, 2, 2],
    [0, 0, 0, 2, 1],
    [0, 0, 1, 2, 2],
    [-2, -1, 0, 2, 2],
    [-2, 0, 1, 0, 0],
    [-2, 0, 0, 2, 1],
    [0, 0, -1, 2, 2],
    [2, 0, 0, 0, 0],
    [0, 0, 1, 0, 1],
    [2, 0, -1, 2, 2],
    [0, 0, -1, 0, 1],
    [0, 0, 1, 2, 1],
    [-2, 0, 2, 0, 0],
    [0, 0, -2, 2, 1],
    [2, 0, 0, 2, 2],
    [0, 0, 2, 2, 2],
    [0, 0, 2, 0, 0],
    [-2, 0, 1, 2, 2],
    [0, 0, 0, 2, 0],
    [-2, 0, 0, 2, 0],
    [0, 0, -1, 2, 1],
    [0, 2, 0, 0, 0],
    [2, 0, -1, 0, 1],
    [-2, 2, 0, 2, 2],
    [0, 1, 0, 0, 1],
    [-2, 0, 1, 0, 1],
    [0, -1, 0, 0, 1],
    [0, 0, 2, -2, 0],
    [2, 0, -1, 2, 1],
    [2, 0, 1, 2, 2],
    [0, 1, 0, 2, 2],
    [-2, 1, 1, 0, 0],
    [0, -1, 0, 2, 2],
    [2, 0, 0, 2, 1],
    [2, 0, 1, 0, 0],
    [-2, 0, 2, 2, 2],
    [-2, 0, 1, 2, 1],
    [2, 0, -2, 0, 1],
    [2, 0, 0, 0, 1],
    [0, -1, 1, 0, 0],
    [-2, -1, 0, 2, 1],
    [-2, 0, 0, 0, 1],
    [0, 0, 2, 2, 1],
    [-2, 0, 2, 0, 1],
    [-2, 1, 0, 2, 1],
    [0, 0, 1, -2, 0],
    [-1, 0, 1, 0, 0],
    [-2, 1, 0, 0, 0],
    [1, 0, 0, 0, 0],
    [0, 0, 1, 2, 0],
    [0, 0, -2, 2, 2],
    [-1, -1, 1, 0, 0],
    [0, 1, 1, 0, 0],
    [0, -1, 1, 2, 2],
    [2, -1, -1, 2, 2],
    [0, 0, 3, 2, 2],
    [2, -1, 0, 2, 2],
])

@jcompile('float64(int64, int64, int64, int64, int64, int64, int64)',
          nopython=True)
def julian_day_dt(year, month, day, hour, minute, second, microsecond):
    """This is the original way to calculate the julian day from the NREL paper.
    However, it is much faster to convert to unix/epoch time and then convert
    to julian day. Note that the date must be UTC."""
    if month <= 2:
        year = year-1
        month = month+12
    a = int(year/100)
    b = 2 - a + int(a * 0.25)
    frac_of_day = (microsecond + (second + minute * 60 + hour * 3600)
                   ) * 1.0 / (3600*24)
    d = day + frac_of_day
    jd = (int(365.25 * (year + 4716)) + int(30.6001 * (month + 1)) + d +
          b - 1524.5)
    return jd

@jcompile('float64(float64)', nopython=True)
def julian_day(unixtime):
    jd = unixtime * 1.0 / 86400 + 2440587.5
    return jd

@jcompile('float64(float64, float64)', nopython=True)
def julian_ephemeris_day(julian_day, delta_t):
    jde = julian_day + delta_t * 1.0 / 86400
    return jde

@jcompile('float64(float64)', nopython=True)
def julian_century(julian_day):
    jc = (julian_day - 2451545) * 1.0 / 36525
    return jc

@jcompile('float64(float64)', nopython=True)
def julian_ephemeris_century(julian_ephemeris_day):
    jce = (julian_ephemeris_day - 2451545) * 1.0 / 36525
    return jce

@jcompile('float64(float64)', nopython=True)
def julian_ephemeris_millennium(julian_ephemeris_century):
    jme = julian_ephemeris_century * 1.0 / 10
    return jme

@jcompile('float64(float64)', nopython=True)
def heliocentric_longitude(jme):
    l0 = 0.0
    l1 = 0.0
    l2 = 0.0
    l3 = 0.0
    l4 = 0.0
    l5 = 0.0

    for row in range(HELIO_LONG_TABLE.shape[1]):
        l0 += (HELIO_LONG_TABLE[0, row, 0]
               * np.cos(HELIO_LONG_TABLE[0, row, 1]
                        + HELIO_LONG_TABLE[0, row, 2] * jme)
               )
        l1 += (HELIO_LONG_TABLE[1, row, 0]
               * np.cos(HELIO_LONG_TABLE[1, row, 1]
                        + HELIO_LONG_TABLE[1, row, 2] * jme)
               )
        l2 += (HELIO_LONG_TABLE[2, row, 0]
               * np.cos(HELIO_LONG_TABLE[2, row, 1]
                        + HELIO_LONG_TABLE[2, row, 2] * jme)
               )
        l3 += (HELIO_LONG_TABLE[3, row, 0]
               * np.cos(HELIO_LONG_TABLE[3, row, 1]
                        + HELIO_LONG_TABLE[3, row, 2] * jme)
               )
        l4 += (HELIO_LONG_TABLE[4, row, 0]
               * np.cos(HELIO_LONG_TABLE[4, row, 1]
                        + HELIO_LONG_TABLE[4, row, 2] * jme)
               )
        l5 += (HELIO_LONG_TABLE[5, row, 0]
               * np.cos(HELIO_LONG_TABLE[5, row, 1]
                        + HELIO_LONG_TABLE[5, row, 2] * jme)
               )

    l_rad = (l0 + l1 * jme + l2 * jme**2 + l3 * jme**3 + l4 * jme**4 +
             l5 * jme**5)/10**8
    l = np.rad2deg(l_rad)
    return l % 360

@jcompile('float64(float64)', nopython=True)
def heliocentric_latitude(jme):
    b0 = 0.0
    b1 = 0.0
    for row in range(HELIO_LAT_TABLE.shape[1]):
        b0 += (HELIO_LAT_TABLE[0, row, 0]
               * np.cos(HELIO_LAT_TABLE[0, row, 1]
                        + HELIO_LAT_TABLE[0, row, 2] * jme)
               )
        b1 += (HELIO_LAT_TABLE[1, row, 0]
               * np.cos(HELIO_LAT_TABLE[1, row, 1]
                        + HELIO_LAT_TABLE[1, row, 2] * jme)
               )

    b_rad = (b0 + b1 * jme)/10**8
    b = np.rad2deg(b_rad)
    return b

@jcompile('float64(float64)', nopython=True)
def heliocentric_radius_vector(jme):
    r0 = 0.0
    r1 = 0.0
    r2 = 0.0
    r3 = 0.0
    r4 = 0.0
    for row in range(HELIO_RADIUS_TABLE.shape[1]):
        r0 += (HELIO_RADIUS_TABLE[0, row, 0]
               * np.cos(HELIO_RADIUS_TABLE[0, row, 1]
                        + HELIO_RADIUS_TABLE[0, row, 2] * jme)
               )
        r1 += (HELIO_RADIUS_TABLE[1, row, 0]
               * np.cos(HELIO_RADIUS_TABLE[1, row, 1]
                        + HELIO_RADIUS_TABLE[1, row, 2] * jme)
               )
        r2 += (HELIO_RADIUS_TABLE[2, row, 0]
               * np.cos(HELIO_RADIUS_TABLE[2, row, 1]
                        + HELIO_RADIUS_TABLE[2, row, 2] * jme)
               )
        r3 += (HELIO_RADIUS_TABLE[3, row, 0]
               * np.cos(HELIO_RADIUS_TABLE[3, row, 1]
                        + HELIO_RADIUS_TABLE[3, row, 2] * jme)
               )
        r4 += (HELIO_RADIUS_TABLE[4, row, 0]
               * np.cos(HELIO_RADIUS_TABLE[4, row, 1]
                        + HELIO_RADIUS_TABLE[4, row, 2] * jme)
               )

    r = (r0 + r1 * jme + r2 * jme**2 + r3 * jme**3 + r4 * jme**4)/10**8
    return r

@jcompile('float64(float64)', nopython=True)
def geocentric_longitude(heliocentric_longitude):
    theta = heliocentric_longitude + 180.0
    return theta % 360

@jcompile('float64(float64)', nopython=True)
def geocentric_latitude(heliocentric_latitude):
    beta = -1.0*heliocentric_latitude
    return beta

@jcompile('float64(float64)', nopython=True)
def mean_elongation(julian_ephemeris_century):
    x0 = (297.85036
          + 445267.111480 * julian_ephemeris_century
          - 0.0019142 * julian_ephemeris_century**2
          + julian_ephemeris_century**3 / 189474)
    return x0

@jcompile('float64(float64)', nopython=True)
def mean_anomaly_sun(julian_ephemeris_century):
    x1 = (357.52772
          + 35999.050340 * julian_ephemeris_century
          - 0.0001603 * julian_ephemeris_century**2
          - julian_ephemeris_century**3 / 300000)
    return x1

@jcompile('float64(float64)', nopython=True)
def mean_anomaly_moon(julian_ephemeris_century):
    x2 = (134.96298
          + 477198.867398 * julian_ephemeris_century
          + 0.0086972 * julian_ephemeris_century**2
          + julian_ephemeris_century**3 / 56250)
    return x2

@jcompile('float64(float64)', nopython=True)
def moon_argument_latitude(julian_ephemeris_century):
    x3 = (93.27191
          + 483202.017538 * julian_ephemeris_century
          - 0.0036825 * julian_ephemeris_century**2
          + julian_ephemeris_century**3 / 327270)
    return x3

@jcompile('float64(float64)', nopython=True)
def moon_ascending_longitude(julian_ephemeris_century):
    x4 = (125.04452
          - 1934.136261 * julian_ephemeris_century
          + 0.0020708 * julian_ephemeris_century**2
          + julian_ephemeris_century**3 / 450000)
    return x4

@jcompile('float64(float64, float64, float64, float64, float64, float64)',
          nopython=True)
def longitude_nutation(julian_ephemeris_century, x0, x1, x2, x3, x4):
    delta_psi_sum = 0
    for row in range(NUTATION_YTERM_ARRAY.shape[0]):
        a = NUTATION_ABCD_ARRAY[row, 0]
        b = NUTATION_ABCD_ARRAY[row, 1]
        argsin = (NUTATION_YTERM_ARRAY[row, 0]*x0 +
                  NUTATION_YTERM_ARRAY[row, 1]*x1 +
                  NUTATION_YTERM_ARRAY[row, 2]*x2 +
                  NUTATION_YTERM_ARRAY[row, 3]*x3 +
                  NUTATION_YTERM_ARRAY[row, 4]*x4)
        term = (a + b * julian_ephemeris_century) * np.sin(np.radians(argsin))
        delta_psi_sum += term
    delta_psi = delta_psi_sum*1.0/36000000
    return delta_psi

@jcompile('float64(float64, float64, float64, float64, float64, float64)',
          nopython=True)
def obliquity_nutation(julian_ephemeris_century, x0, x1, x2, x3, x4):
    delta_eps_sum = 0.0
    for row in range(NUTATION_YTERM_ARRAY.shape[0]):
        c = NUTATION_ABCD_ARRAY[row, 2]
        d = NUTATION_ABCD_ARRAY[row, 3]
        argcos = (NUTATION_YTERM_ARRAY[row, 0]*x0 +
                  NUTATION_YTERM_ARRAY[row, 1]*x1 +
                  NUTATION_YTERM_ARRAY[row, 2]*x2 +
                  NUTATION_YTERM_ARRAY[row, 3]*x3 +
                  NUTATION_YTERM_ARRAY[row, 4]*x4)
        term = (c + d * julian_ephemeris_century) * np.cos(np.radians(argcos))
        delta_eps_sum += term
    delta_eps = delta_eps_sum*1.0/36000000
    return delta_eps

@jcompile('float64(float64)', nopython=True)
def mean_ecliptic_obliquity(julian_ephemeris_millennium):
    U = 1.0*julian_ephemeris_millennium/10
    e0 = (84381.448 - 4680.93 * U - 1.55 * U**2
          + 1999.25 * U**3 - 51.38 * U**4 - 249.67 * U**5
          - 39.05 * U**6 + 7.12 * U**7 + 27.87 * U**8
          + 5.79 * U**9 + 2.45 * U**10)
    return e0

@jcompile('float64(float64, float64)', nopython=True)
def true_ecliptic_obliquity(mean_ecliptic_obliquity, obliquity_nutation):
    e0 = mean_ecliptic_obliquity
    deleps = obliquity_nutation
    e = e0*1.0/3600 + deleps
    return e

@jcompile('float64(float64)', nopython=True)
def aberration_correction(earth_radius_vector):
    deltau = -20.4898 / (3600 * earth_radius_vector)
    return deltau

@jcompile('float64(float64, float64, float64)', nopython=True)
def apparent_sun_longitude(geocentric_longitude, longitude_nutation,
                           aberration_correction):
    lamd = geocentric_longitude + longitude_nutation + aberration_correction
    return lamd

@jcompile('float64(float64, float64)', nopython=True)
def mean_sidereal_time(julian_day, julian_century):
    v0 = (280.46061837 + 360.98564736629 * (julian_day - 2451545)
          + 0.000387933 * julian_century**2 - julian_century**3 / 38710000)
    return v0 % 360.0

@jcompile('float64(float64, float64, float64)', nopython=True)
def apparent_sidereal_time(mean_sidereal_time, longitude_nutation,
                           true_ecliptic_obliquity):
    v = mean_sidereal_time + longitude_nutation * np.cos(
        np.radians(true_ecliptic_obliquity))
    return v

@jcompile('float64(float64, float64, float64)', nopython=True)
def geocentric_sun_right_ascension(apparent_sun_longitude,
                                   true_ecliptic_obliquity,
                                   geocentric_latitude):
    num = (np.sin(np.radians(apparent_sun_longitude))
           * np.cos(np.radians(true_ecliptic_obliquity))
           - np.tan(np.radians(geocentric_latitude))
           * np.sin(np.radians(true_ecliptic_obliquity)))
    alpha = np.degrees(np.arctan2(num, np.cos(
        np.radians(apparent_sun_longitude))))
    return alpha % 360

@jcompile('float64(float64, float64, float64)', nopython=True)
def geocentric_sun_declination(apparent_sun_longitude, true_ecliptic_obliquity,
                               geocentric_latitude):
    delta = np.degrees(np.arcsin(np.sin(np.radians(geocentric_latitude)) *
                                 np.cos(np.radians(true_ecliptic_obliquity)) +
                                 np.cos(np.radians(geocentric_latitude)) *
                                 np.sin(np.radians(true_ecliptic_obliquity)) *
                                 np.sin(np.radians(apparent_sun_longitude))))
    return delta

@jcompile('float64(float64, float64, float64)', nopython=True)
def local_hour_angle(apparent_sidereal_time, observer_longitude,
                     sun_right_ascension):
    """Measured westward from south"""
    H = apparent_sidereal_time + observer_longitude - sun_right_ascension
    return H % 360

@jcompile('float64(float64)', nopython=True)
def equatorial_horizontal_parallax(earth_radius_vector):
    xi = 8.794 / (3600 * earth_radius_vector)
    return xi

@jcompile('float64(float64)', nopython=True)
def uterm(observer_latitude):
    u = np.arctan(0.99664719 * np.tan(np.radians(observer_latitude)))
    return u

@jcompile('float64(float64, float64, float64)', nopython=True)
def xterm(u, observer_latitude, observer_elevation):
    x = (np.cos(u) + observer_elevation / 6378140
         * np.cos(np.radians(observer_latitude)))
    return x

@jcompile('float64(float64, float64, float64)', nopython=True)
def yterm(u, observer_latitude, observer_elevation):
    y = (0.99664719 * np.sin(u) + observer_elevation / 6378140
         * np.sin(np.radians(observer_latitude)))
    return y

@jcompile('float64(float64, float64,float64, float64)', nopython=True)
def parallax_sun_right_ascension(xterm, equatorial_horizontal_parallax,
                                 local_hour_angle, geocentric_sun_declination):
    num = (-xterm * np.sin(np.radians(equatorial_horizontal_parallax))
           * np.sin(np.radians(local_hour_angle)))
    denom = (np.cos(np.radians(geocentric_sun_declination))
             - xterm * np.sin(np.radians(equatorial_horizontal_parallax))
             * np.cos(np.radians(local_hour_angle)))
    delta_alpha = np.degrees(np.arctan2(num, denom))
    return delta_alpha

@jcompile('float64(float64, float64)', nopython=True)
def topocentric_sun_right_ascension(geocentric_sun_right_ascension,
                                    parallax_sun_right_ascension):
    alpha_prime = geocentric_sun_right_ascension + parallax_sun_right_ascension
    return alpha_prime

@jcompile('float64(float64, float64, float64, float64, float64, float64)',
          nopython=True)
def topocentric_sun_declination(geocentric_sun_declination, xterm, yterm,
                                equatorial_horizontal_parallax,
                                parallax_sun_right_ascension,
                                local_hour_angle):
    num = ((np.sin(np.radians(geocentric_sun_declination)) - yterm
            * np.sin(np.radians(equatorial_horizontal_parallax)))
           * np.cos(np.radians(parallax_sun_right_ascension)))
    denom = (np.cos(np.radians(geocentric_sun_declination)) - xterm
             * np.sin(np.radians(equatorial_horizontal_parallax))
             * np.cos(np.radians(local_hour_angle)))
    delta = np.degrees(np.arctan2(num, denom))
    return delta

@jcompile('float64(float64, float64)', nopython=True)
def topocentric_local_hour_angle(local_hour_angle,
                                 parallax_sun_right_ascension):
    H_prime = local_hour_angle - parallax_sun_right_ascension
    return H_prime

@jcompile('float64(float64, float64, float64)', nopython=True)
def topocentric_elevation_angle_without_atmosphere(observer_latitude,
                                                   topocentric_sun_declination,
                                                   topocentric_local_hour_angle
                                                   ):
    e0 = np.degrees(np.arcsin(
        np.sin(np.radians(observer_latitude))
        * np.sin(np.radians(topocentric_sun_declination))
        + np.cos(np.radians(observer_latitude))
        * np.cos(np.radians(topocentric_sun_declination))
        * np.cos(np.radians(topocentric_local_hour_angle))))
    return e0

@jcompile('float64(float64, float64, float64, float64)', nopython=True)
def atmospheric_refraction_correction(local_pressure, local_temp,
                                      topocentric_elevation_angle_wo_atmosphere,
                                      atmos_refract):
    # switch sets delta_e when the sun is below the horizon
    switch = topocentric_elevation_angle_wo_atmosphere >= -1.0 * (
        0.26667 + atmos_refract)
    delta_e = ((local_pressure / 1010.0) * (283.0 / (273 + local_temp))
               * 1.02 / (60 * np.tan(np.radians(
                   topocentric_elevation_angle_wo_atmosphere
                   + 10.3 / (topocentric_elevation_angle_wo_atmosphere
                             + 5.11))))) * switch
    return delta_e

@jcompile('float64(float64, float64)', nopython=True)
def topocentric_elevation_angle(topocentric_elevation_angle_without_atmosphere,
                                atmospheric_refraction_correction):
    e = (topocentric_elevation_angle_without_atmosphere
         + atmospheric_refraction_correction)
    return e

@jcompile('float64(float64)', nopython=True)
def topocentric_zenith_angle(topocentric_elevation_angle):
    theta = 90 - topocentric_elevation_angle
    return theta

@jcompile('float64(float64, float64, float64)', nopython=True)
def topocentric_astronomers_azimuth(topocentric_local_hour_angle,
                                    topocentric_sun_declination,
                                    observer_latitude):
    num = np.sin(np.radians(topocentric_local_hour_angle))
    denom = (np.cos(np.radians(topocentric_local_hour_angle))
             * np.sin(np.radians(observer_latitude))
             - np.tan(np.radians(topocentric_sun_declination))
             * np.cos(np.radians(observer_latitude)))
    gamma = np.degrees(np.arctan2(num, denom))
    return gamma % 360

@jcompile('float64(float64)', nopython=True)
def topocentric_azimuth_angle(topocentric_astronomers_azimuth):
    phi = topocentric_astronomers_azimuth + 180
    return phi % 360

@jcompile('float64(float64)', nopython=True)
def sun_mean_longitude(julian_ephemeris_millennium):
    M = (280.4664567 + 360007.6982779 * julian_ephemeris_millennium
         + 0.03032028 * julian_ephemeris_millennium**2
         + julian_ephemeris_millennium**3 / 49931
         - julian_ephemeris_millennium**4 / 15300
         - julian_ephemeris_millennium**5 / 2000000)
    return M

@jcompile('float64(float64, float64, float64, float64)', nopython=True)
def equation_of_time(sun_mean_longitude, geocentric_sun_right_ascension,
                     longitude_nutation, true_ecliptic_obliquity):
    E = (sun_mean_longitude - 0.0057183 - geocentric_sun_right_ascension +
         longitude_nutation * np.cos(np.radians(true_ecliptic_obliquity)))
    # limit between 0 and 360
    E = E % 360
    # convert to minutes
    E *= 4
    greater = E > 20
    less = E < -20
    other = (E <= 20) & (E >= -20)
    E = greater * (E - 1440) + less * (E + 1440) + other * E
    return E

@jcompile('void(float64[:], float64[:], float64[:,:])', nopython=True,
          nogil=True)
def solar_position_loop(unixtime, loc_args, out):
    """Loop through the time array and calculate the solar position"""
    lat = loc_args[0]
    lon = loc_args[1]
    elev = loc_args[2]
    pressure = loc_args[3]
    temp = loc_args[4]
    delta_t = loc_args[5]
    atmos_refract = loc_args[6]
    sst = loc_args[7]
    esd = loc_args[8]

    for i in range(unixtime.shape[0]):
        utime = unixtime[i]
        jd = julian_day(utime)
        jde = julian_ephemeris_day(jd, delta_t)
        jc = julian_century(jd)
        jce = julian_ephemeris_century(jde)
        jme = julian_ephemeris_millennium(jce)
        R = heliocentric_radius_vector(jme)
        if esd:
            out[0, i] = R
            continue
        L = heliocentric_longitude(jme)
        B = heliocentric_latitude(jme)
        Theta = geocentric_longitude(L)
        beta = geocentric_latitude(B)
        x0 = mean_elongation(jce)
        x1 = mean_anomaly_sun(jce)
        x2 = mean_anomaly_moon(jce)
        x3 = moon_argument_latitude(jce)
        x4 = moon_ascending_longitude(jce)
        delta_psi = longitude_nutation(jce, x0, x1, x2, x3, x4)
        delta_epsilon = obliquity_nutation(jce, x0, x1, x2, x3, x4)
        epsilon0 = mean_ecliptic_obliquity(jme)
        epsilon = true_ecliptic_obliquity(epsilon0, delta_epsilon)
        delta_tau = aberration_correction(R)
        lamd = apparent_sun_longitude(Theta, delta_psi, delta_tau)
        v0 = mean_sidereal_time(jd, jc)
        v = apparent_sidereal_time(v0, delta_psi, epsilon)
        alpha = geocentric_sun_right_ascension(lamd, epsilon, beta)
        delta = geocentric_sun_declination(lamd, epsilon, beta)
        if sst:
            out[0, i] = v
            out[1, i] = alpha
            out[2, i] = delta
            continue
        m = sun_mean_longitude(jme)
        eot = equation_of_time(m, alpha, delta_psi, epsilon)
        H = local_hour_angle(v, lon, alpha)
        xi = equatorial_horizontal_parallax(R)
        u = uterm(lat)
        x = xterm(u, lat, elev)
        y = yterm(u, lat, elev)
        delta_alpha = parallax_sun_right_ascension(x, xi, H, delta)
        delta_prime = topocentric_sun_declination(delta, x, y, xi, delta_alpha,
                                                  H)
        H_prime = topocentric_local_hour_angle(H, delta_alpha)
        e0 = topocentric_elevation_angle_without_atmosphere(lat, delta_prime,
                                                            H_prime)
        delta_e = atmospheric_refraction_correction(pressure, temp, e0,
                                                    atmos_refract)
        e = topocentric_elevation_angle(e0, delta_e)
        theta = topocentric_zenith_angle(e)
        theta0 = topocentric_zenith_angle(e0)
        gamma = topocentric_astronomers_azimuth(H_prime, delta_prime, lat)
        phi = topocentric_azimuth_angle(gamma)
        out[0, i] = theta
        out[1, i] = theta0
        out[2, i] = e
        out[3, i] = e0
        out[4, i] = phi
        out[5, i] = eot

def solar_position_numba(unixtime, lat, lon, elev, pressure, temp, delta_t,
                         atmos_refract, numthreads, sst=False, esd=False):
    """Calculate the solar position using the numba compiled functions
    and multiple threads. Very slow if functions are not numba compiled.
    """
    # these args are the same for each thread
    loc_args = np.array([lat, lon, elev, pressure, temp, delta_t,
                         atmos_refract, sst, esd])

    # construct dims x ulength array to put the results in
    ulength = unixtime.shape[0]
    if sst:
        dims = 3
    elif esd:
        dims = 1
    else:
        dims = 6
    result = np.empty((dims, ulength), dtype=np.float64)

    if unixtime.dtype != np.float64:
        unixtime = unixtime.astype(np.float64)

    if ulength < numthreads:
        warnings.warn('The number of threads is more than the length of '
                      'the time array. Only using %s threads.'.format(ulength))
        numthreads = ulength

    if numthreads <= 1:
        solar_position_loop(unixtime, loc_args, result)
        return result

    # split the input and output arrays into numthreads chunks
    split0 = np.array_split(unixtime, numthreads)
    split2 = np.array_split(result, numthreads, axis=1)
    chunks = [[a0, loc_args, split2[i]] for i, a0 in enumerate(split0)]
    # Spawn one thread per chunk
    threads = [threading.Thread(target=solar_position_loop, args=chunk)
               for chunk in chunks]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join()
    return result

def solar_position_numpy(unixtime, lat, lon, elev, pressure, temp, delta_t,
                         atmos_refract, numthreads, sst=False, esd=False):
    """Calculate the solar position assuming unixtime is a numpy array. Note
    this function will not work if the solar position functions were
    compiled with numba.
    """

    jd = julian_day(unixtime)
    jde = julian_ephemeris_day(jd, delta_t)
    jc = julian_century(jd)
    jce = julian_ephemeris_century(jde)
    jme = julian_ephemeris_millennium(jce)
    R = heliocentric_radius_vector(jme)
    if esd:
        return (R, )
    L = heliocentric_longitude(jme)
    B = heliocentric_latitude(jme)
    Theta = geocentric_longitude(L)
    beta = geocentric_latitude(B)
    x0 = mean_elongation(jce)
    x1 = mean_anomaly_sun(jce)
    x2 = mean_anomaly_moon(jce)
    x3 = moon_argument_latitude(jce)
    x4 = moon_ascending_longitude(jce)
    delta_psi = longitude_nutation(jce, x0, x1, x2, x3, x4)
    delta_epsilon = obliquity_nutation(jce, x0, x1, x2, x3, x4)
    epsilon0 = mean_ecliptic_obliquity(jme)
    epsilon = true_ecliptic_obliquity(epsilon0, delta_epsilon)
    delta_tau = aberration_correction(R)
    lamd = apparent_sun_longitude(Theta, delta_psi, delta_tau)
    v0 = mean_sidereal_time(jd, jc)
    v = apparent_sidereal_time(v0, delta_psi, epsilon)
    alpha = geocentric_sun_right_ascension(lamd, epsilon, beta)
    delta = geocentric_sun_declination(lamd, epsilon, beta)
    if sst:
        return v, alpha, delta
    m = sun_mean_longitude(jme)
    eot = equation_of_time(m, alpha, delta_psi, epsilon)
    H = local_hour_angle(v, lon, alpha)
    xi = equatorial_horizontal_parallax(R)
    u = uterm(lat)
    x = xterm(u, lat, elev)
    y = yterm(u, lat, elev)
    delta_alpha = parallax_sun_right_ascension(x, xi, H, delta)
    delta_prime = topocentric_sun_declination(delta, x, y, xi, delta_alpha, H)
    H_prime = topocentric_local_hour_angle(H, delta_alpha)
    e0 = topocentric_elevation_angle_without_atmosphere(lat, delta_prime,
                                                        H_prime)
    delta_e = atmospheric_refraction_correction(pressure, temp, e0,
                                                atmos_refract)
    e = topocentric_elevation_angle(e0, delta_e)
    theta = topocentric_zenith_angle(e)
    theta0 = topocentric_zenith_angle(e0)
    gamma = topocentric_astronomers_azimuth(H_prime, delta_prime, lat)
    phi = topocentric_azimuth_angle(gamma)
    return theta, theta0, e, e0, phi, eot

def solar_position(unixtime, lat, lon, elev, pressure, temp, delta_t,
                   atmos_refract, numthreads=8, sst=False, esd=False):

    """
    Calculate the solar position using the
    NREL SPA algorithm described in [1].

    If numba is installed, the functions can be compiled
    and the code runs quickly. If not, the functions
    still evaluate but use numpy instead.

    Parameters
    ----------
    unixtime : numpy array
        Array of unix/epoch timestamps to calculate solar position for.
        Unixtime is the number of seconds since Jan. 1, 1970 00:00:00 UTC.
        A pandas.DatetimeIndex is easily converted using .astype(np.int64)/10**9
    lat : float
        Latitude to calculate solar position for
    lon : float
        Longitude to calculate solar position for
    elev : float
        Elevation of location in meters
    pressure : int or float
        avg. yearly pressure at location in millibars;
        used for atmospheric correction
    temp : int or float
        avg. yearly temperature at location in
        degrees C; used for atmospheric correction
    delta_t : float, optional
        If delta_t is None, uses spa.calculate_deltat
        using time.year and time.month from pandas.DatetimeIndex.
        For most simulations specifing delta_t is sufficient.
        Difference between terrestrial time and UT1.
        *Note: delta_t = None will break code using nrel_numba,
        this will be fixed in a future version.
        By default, use USNO historical data and predictions
    atmos_refrac : float, optional
        The approximate atmospheric refraction (in degrees)
        at sunrise and sunset.
    numthreads: int, optional, default 8
        Number of threads to use for computation if numba>=0.17
        is installed.
    sst : bool, default False
        If True, return only data needed for sunrise, sunset, and transit
        calculations.
    esd : bool, default False
        If True, return only Earth-Sun distance in AU

    Returns
    -------
    Numpy Array with elements:
        apparent zenith,
        zenith,
        elevation,
        apparent_elevation,
        azimuth,
        equation_of_time

    References
    ----------
    [1] I. Reda and A. Andreas, Solar position algorithm for solar radiation
    applications. Solar Energy, vol. 76, no. 5, pp. 577-589, 2004.

    [2] I. Reda and A. Andreas, Corrigendum to Solar position algorithm for
    solar radiation applications. Solar Energy, vol. 81, no. 6, p. 838, 2007.
    """
    if USE_NUMBA:
        do_calc = solar_position_numba
    else:
        do_calc = solar_position_numpy

    result = do_calc(unixtime, lat, lon, elev, pressure,
                     temp, delta_t, atmos_refract, numthreads,
                     sst, esd)

    if not isinstance(result, np.ndarray):
        try:
            result = np.array(result)
        except Exception:
            pass

    return result

def transit_sunrise_sunset(dates, lat, lon, delta_t, numthreads):
    """
    Calculate the sun transit, sunrise, and sunset
    for a set of dates at a given location.

    Parameters
    ----------
    dates : array
        Numpy array of ints/floats corresponding to the Unix time
        for the dates of interest, must be midnight UTC (00:00+00:00)
        on the day of interest.
    lat : float
        Latitude of location to perform calculation for
    lon : float
        Longitude of location
    delta_t : float
        Difference between terrestrial time and UT. USNO has tables.
    numthreads : int
        Number to threads to use for calculation (if using numba)

    Returns
    -------
    tuple : (transit, sunrise, sunset) localized to UTC

    """

    if ((dates % 86400) != 0.0).any():
        raise ValueError('Input dates must be at 00:00 UTC')

    utday = (dates // 86400) * 86400
    ttday0 = utday - delta_t
    ttdayn1 = ttday0 - 86400
    ttdayp1 = ttday0 + 86400

    # index 0 is v, 1 is alpha, 2 is delta
    utday_res = solar_position(utday, 0, 0, 0, 0, 0, delta_t,
                               0, numthreads, sst=True)
    v = utday_res[0]

    ttday0_res = solar_position(ttday0, 0, 0, 0, 0, 0, delta_t,
                                0, numthreads, sst=True)
    ttdayn1_res = solar_position(ttdayn1, 0, 0, 0, 0, 0, delta_t,
                                 0, numthreads, sst=True)
    ttdayp1_res = solar_position(ttdayp1, 0, 0, 0, 0, 0, delta_t,
                                 0, numthreads, sst=True)
    m0 = (ttday0_res[1] - lon - v) / 360
    cos_arg = ((np.sin(np.radians(-0.8333)) - np.sin(np.radians(lat))
               * np.sin(np.radians(ttday0_res[2]))) /
               (np.cos(np.radians(lat)) * np.cos(np.radians(ttday0_res[2]))))
    cos_arg[abs(cos_arg) > 1] = np.nan
    H0 = np.degrees(np.arccos(cos_arg)) % 180

    m = np.empty((3, len(utday)))
    m[0] = m0 % 1
    m[1] = (m[0] - H0 / 360)
    m[2] = (m[0] + H0 / 360)

    # need to account for fractions of day that may be the next or previous
    # day in UTC
    add_a_day = m[2] >= 1
    sub_a_day = m[1] < 0
    m[1] = m[1] % 1
    m[2] = m[2] % 1
    vs = v + 360.985647 * m
    n = m + delta_t / 86400

    a = ttday0_res[1] - ttdayn1_res[1]
    a[abs(a) > 2] = a[abs(a) > 2] % 1
    ap = ttday0_res[2] - ttdayn1_res[2]
    ap[abs(ap) > 2] = ap[abs(ap) > 2] % 1
    b = ttdayp1_res[1] - ttday0_res[1]
    b[abs(b) > 2] = b[abs(b) > 2] % 1
    bp = ttdayp1_res[2] - ttday0_res[2]
    bp[abs(bp) > 2] = bp[abs(bp) > 2] % 1
    c = b - a
    cp = bp - ap

    alpha_prime = ttday0_res[1] + (n * (a + b + c * n)) / 2
    delta_prime = ttday0_res[2] + (n * (ap + bp + cp * n)) / 2
    Hp = (vs + lon - alpha_prime) % 360
    Hp[Hp >= 180] = Hp[Hp >= 180] - 360

    h = np.degrees(np.arcsin(np.sin(np.radians(lat)) *
                             np.sin(np.radians(delta_prime)) +
                             np.cos(np.radians(lat)) *
                             np.cos(np.radians(delta_prime))
                             * np.cos(np.radians(Hp))))

    T = (m[0] - Hp[0] / 360) * 86400
    R = (m[1] + (h[1] + 0.8333) / (360 * np.cos(np.radians(delta_prime[1])) *
                                   np.cos(np.radians(lat)) *
                                   np.sin(np.radians(Hp[1])))) * 86400
    S = (m[2] + (h[2] + 0.8333) / (360 * np.cos(np.radians(delta_prime[2])) *
                                   np.cos(np.radians(lat)) *
                                   np.sin(np.radians(Hp[2])))) * 86400

    S[add_a_day] += 86400
    R[sub_a_day] -= 86400

    transit = T + utday
    sunrise = R + utday
    sunset = S + utday

    return transit, sunrise, sunset

def earthsun_distance(unixtime, delta_t, numthreads):
    """
    Calculates the distance from the earth to the sun using the
    NREL SPA algorithm described in [1].

    Parameters
    ----------
    unixtime : numpy array
        Array of unix/epoch timestamps to calculate solar position for.
        Unixtime is the number of seconds since Jan. 1, 1970 00:00:00 UTC.
        A pandas.DatetimeIndex is easily converted using .astype(np.int64)/10**9
    delta_t : float
        Difference between terrestrial time and UT. USNO has tables.
    numthreads : int
        Number to threads to use for calculation (if using numba)

    Returns
    -------
    R : array
        Earth-Sun distance in AU.

    References
    ----------
    [1] Reda, I., Andreas, A., 2003. Solar position algorithm for solar
    radiation applications. Technical report: NREL/TP-560- 34302. Golden,
    USA, http://www.nrel.gov.
    """

    R = solar_position(unixtime, 0, 0, 0, 0, 0, delta_t,
                       0, numthreads, esd=True)[0]

    return R

def calculate_deltat(year, month):
    """Calculate the difference between Terrestrial Dynamical Time (TD)
    and Universal Time (UT).

    Note: This function is not yet compatible for calculations using
    Numba.

    Equations taken from http://eclipse.gsfc.nasa.gov/SEcat5/deltatpoly.html
    """

    plw = 'Deltat is unknown for years before -1999 and after 3000. ' \
          'Delta values will be calculated, but the calculations ' \
          'are not intended to be used for these years.'

    try:
        if np.any((year > 3000) | (year < -1999)):
            warnings.warn(plw)
    except ValueError:
        if (year > 3000) | (year < -1999):
            warnings.warn(plw)
    except TypeError:
        return 0

    y = year + (month - 0.5)/12

    deltat = np.where(year < -500,

                      -20+32*((y-1820)/100)**2, 0)

    deltat = np.where((-500 <= year) & (year < 500),

                      10583.6-1014.41*(y/100)
                      + 33.78311*(y/100)**2
                      - 5.952053*(y/100)**3
                      - 0.1798452*(y/100)**4
                      + 0.022174192*(y/100)**5
                      + 0.0090316521*(y/100)**6, deltat)

    deltat = np.where((500 <= year) & (year < 1600),

                      1574.2-556.01*((y-1000)/100)
                      + 71.23472*((y-1000)/100)**2
                      + 0.319781*((y-1000)/100)**3
                      - 0.8503463*((y-1000)/100)**4
                      - 0.005050998*((y-1000)/100)**5
                      + 0.0083572073*((y-1000)/100)**6, deltat)

    deltat = np.where((1600 <= year) & (year < 1700),

                      120-0.9808*(y-1600)
                      - 0.01532*(y-1600)**2
                      + (y-1600)**3/7129, deltat)

    deltat = np.where((1700 <= year) & (year < 1800),

                      8.83+0.1603*(y-1700)
                      - 0.0059285*(y-1700)**2
                      + 0.00013336*(y-1700)**3
                      - (y-1700)**4/1174000, deltat)

    deltat = np.where((1800 <= year) & (year < 1860),

                      13.72-0.332447*(y-1800)
                      + 0.0068612*(y-1800)**2
                      + 0.0041116*(y-1800)**3
                      - 0.00037436*(y-1800)**4
                      + 0.0000121272*(y-1800)**5
                      - 0.0000001699*(y-1800)**6
                      + 0.000000000875*(y-1800)**7, deltat)

    deltat = np.where((1860 <= year) & (year < 1900),

                      7.6+0.5737*(y-1860)
                      - 0.251754*(y-1860)**2
                      + 0.01680668*(y-1860)**3
                      - 0.0004473624*(y-1860)**4
                      + (y-1860)**5/233174, deltat)

    deltat = np.where((1900 <= year) & (year < 1920),

                      -2.79+1.494119*(y-1900)
                      - 0.0598939*(y-1900)**2
                      + 0.0061966*(y-1900)**3
                      - 0.000197*(y-1900)**4, deltat)

    deltat = np.where((1920 <= year) & (year < 1941),

                      21.20+0.84493*(y-1920)
                      - 0.076100*(y-1920)**2
                      + 0.0020936*(y-1920)**3, deltat)

    deltat = np.where((1941 <= year) & (year < 1961),

                      29.07+0.407*(y-1950)
                      - (y-1950)**2/233
                      + (y-1950)**3/2547, deltat)

    deltat = np.where((1961 <= year) & (year < 1986),

                      45.45+1.067*(y-1975)
                      - (y-1975)**2/260
                      - (y-1975)**3/718, deltat)

    deltat = np.where((1986 <= year) & (year < 2005),

                      63.86+0.3345*(y-2000)
                      - 0.060374*(y-2000)**2
                      + 0.0017275*(y-2000)**3
                      + 0.000651814*(y-2000)**4
                      + 0.00002373599*(y-2000)**5, deltat)

    deltat = np.where((2005 <= year) & (year < 2050),

                      62.92+0.32217*(y-2000)
                      + 0.005589*(y-2000)**2, deltat)

    deltat = np.where((2050 <= year) & (year < 2150),

                      -20+32*((y-1820)/100)**2
                      - 0.5628*(2150-y), deltat)

    deltat = np.where(year > 2150,

                      -20+32*((y-1820)/100)**2, deltat)

    deltat = np.asscalar(deltat) if np.isscalar(year) & np.isscalar(month)\
        else deltat

    return deltat


from __future__ import division
from build.el_propagator_xf import el_run
from mqc.mqc import MQC
from misc import eps, au_to_K, au_to_A, call_name, typewriter
import random, os, shutil, textwrap
import numpy as np
import pickle

class Auxiliary_Molecule(object):
    """ Class for auxiliary molecule that is used for the calculation of decoherence term

        :param object molecule: Molecule object
    """
    def __init__(self, molecule, l_xf1d):
        # Initialize auxiliary molecule
        if (l_xf1d):

            self.nat = 1
            self.ndim = 1
            self.symbols = ['XX']

            self.mass = np.zeros((self.nat))
            self.mass[0] = 1. / np.sum(1. / molecule.mass[0:molecule.nat_qm])

        else:

            self.nat = molecule.nat_qm
            self.ndim = molecule.ndim
            self.symbols = np.copy(molecule.symbols[0:molecule.nat_qm])

            self.mass = np.copy(molecule.mass[0:molecule.nat_qm])

        self.pos = np.zeros((molecule.nst, self.nat, self.ndim))
        self.vel = np.zeros((molecule.nst, self.nat, self.ndim))
        self.vel_old = np.copy(self.vel)

class SHXF(MQC):
    """ Class for DISH-XF dynamics

        :param object molecule: Molecule object
        :param object thermostat: Thermostat object
        :param integer istate: Initial state
        :param double dt: Time interval
        :param integer nsteps: Total step of nuclear propagation
        :param integer nesteps: Total step of electronic propagation
        :param string elec_object: Electronic equation of motions
        :param string propagator: Electronic propagator
        :param boolean l_print_dm: Logical to print BO population and coherence
        :param boolean l_adj_nac: Adjust nonadiabatic coupling to align the phases
        :param string hop_rescale: Velocity rescaling method after successful hop
        :param string hop_reject: Velocity rescaling method after frustrated hop
        :param double rho_threshold: Electronic density threshold for decoherence term calculation
        :param sigma: Width of nuclear wave packet of auxiliary trajectory
        :type sigma: double or double,list
        :param init_coef: Initial BO coefficient
        :type init_coef: double, list or complex, list
        :param boolean l_econs_state: Logical to use identical total energies for all auxiliary trajectories
        :param string aux_econs_viol: How to treat trajectories violating the total energy conservation
        :param string unit_dt: Unit of time interval
        :param integer out_freq: Frequency of printing output
        :param integer verbosity: Verbosity of output
    """
    def __init__(self, molecule, thermostat=None, istate=0, dt=0.5, nsteps=1000, nesteps=20, \
        elec_object="density", propagator="rk4", l_print_dm=True, l_adj_nac=True, hop_rescale="augment", \
        hop_reject="reverse", rho_threshold=0.01, sigma=None, l_xf1d=False, init_coef=None, \
        l_econs_state=True, aux_econs_viol="fix", unit_dt="fs", out_freq=1, verbosity=0):
        # Initialize input values
        super().__init__(molecule, thermostat, istate, dt, nsteps, nesteps, \
            elec_object, propagator, l_print_dm, l_adj_nac, init_coef, unit_dt, out_freq, verbosity)

        # Initialize SH variables
        self.rstate = istate
        self.rstate_old = self.rstate

        self.rand = 0.
        self.prob = np.zeros(self.mol.nst)
        self.acc_prob = np.zeros(self.mol.nst + 1)

        self.l_hop = False
        self.l_reject = False

        self.hop_rescale = hop_rescale.lower()
        if not (self.hop_rescale in ["energy", "velocity", "momentum", "augment"]):
            error_message = "Invalid rescaling method for accepted hop!"
            error_vars = f"hop_rescale = {self.hop_rescale}"
            raise ValueError (f"( {self.md_type}.{call_name()} ) {error_message} ( {error_vars} )")

        self.hop_reject = hop_reject.lower()
        if not (self.hop_reject in ["keep", "reverse"]):
            error_message = "Invalid rescaling method for frustrated hop!"
            error_vars = f"hop_reject = {self.hop_reject}"
            raise ValueError (f"( {self.md_type}.{call_name()} ) {error_message} ( {error_vars} )")

        # Check error for incompatible cases
        if (self.mol.l_nacme):
            # No analytical nonadiabatic couplings exist
            if (self.hop_rescale in ["velocity", "momentum", "augment"]):
                error_message = "NACVs are not available with current QM object, only isotropic rescaling is possible!"
                error_vars = f"hop_rescale = {self.hop_rescale}"
                raise ValueError (f"( {self.md_type}.{call_name()} ) {error_message} ( {error_vars} )")
            # TODO : This error will be used after adding the 'flip' option for hop_reject
#            if (self.hop_reject == "reverse"):
#                error_message = "NACVs are not available with current QM object, only keep rescaling is possible!"
#                error_vars = f"hop_reject = {self.hop_reject}"
#                raise ValueError (f"( {self.md_type}.{call_name()} ) {error_message} ( {error_vars} )")

        # Initialize XF related variables
        self.force_hop = False
        self.l_econs_state = l_econs_state
        self.l_xf1d = l_xf1d
        self.l_coh = [False] * self.mol.nst
        self.l_first = [False] * self.mol.nst
        self.l_fix = [False] * self.mol.nst
        self.l_collapse = False
        self.rho_threshold = rho_threshold
        self.aux_econs_viol = aux_econs_viol

        if not (self.aux_econs_viol in ["fix", "collapse"]):
            error_message = "Invalid method to treat auxiliary trajectories that violate the total energy conservation!"
            error_vars = f"aux_econs_viol = {self.aux_econs_viol}"
            raise ValueError (f"( {self.md_type}.{call_name()} ) {error_message} ( {error_vars} )")

        self.sigma = sigma
        if (self.sigma == None):
            error_message = "Sigma for auxiliary trajectories must be set in running script!"
            error_vars = f"sigma = {self.sigma}"
            raise ValueError (f"( {self.md_type}.{call_name()} ) {error_message} ( {error_vars} )")

        if (isinstance(self.sigma, float)):
            # uniform value for sigma
            pass
        elif (isinstance(self.sigma, list)):
            # atom-resolved values for sigma
            if (len(self.sigma) != self.mol.nat_qm):
                error_message = "Number of elements for sigma must be equal to number of atoms!"
                error_vars = f"len(sigma) = {len(self.sigma)}"
                raise ValueError (f"( {self.md_type}.{call_name()} ) {error_message} ( {error_vars} )")
            if (self.l_xf1d):
                error_message = "Sigma must be float, not list in XF-1D scheme!"
                error_vars = f"sigma = {self.sigma}"
                raise TypeError (f"( {self.md_type}.{call_name()} ) {error_message} ( {error_vars} )")
        else:
            error_message = "Type of sigma must be float or list consisting of float!"
            error_vars = f"sigma = {self.sigma}"
            raise TypeError (f"( {self.md_type}.{call_name()} ) {error_message} ( {error_vars} )")

        self.upper_th = 1. - self.rho_threshold
        self.lower_th = self.rho_threshold

        # Initialize auxiliary molecule object
        self.aux = Auxiliary_Molecule(self.mol, self.l_xf1d)
        self.pos_0 = np.zeros((self.aux.nat, self.aux.ndim))
        self.phase = np.zeros((self.mol.nst, self.aux.nat, self.aux.ndim))

        # Debug variables
        self.dotpopdec = np.zeros(self.mol.nst)
        self.dotpopnac = np.zeros(self.mol.nst)
        self.qmom = np.zeros((self.aux.nat, self.aux.ndim))

        # Initialize event to print
        self.event = {"HOP": [], "DECO": []}

    def run(self, qm, mm=None, output_dir="./", l_save_qm_log=False, l_save_mm_log=False, l_save_scr=True, restart=None):
        """ Run MQC dynamics according to decoherence-induced surface hopping dynamics

            :param object qm: QM object containing on-the-fly calculation infomation
            :param object mm: MM object containing MM calculation infomation
            :param string output_dir: Name of directory where outputs to be saved.
            :param boolean l_save_qm_log: Logical for saving QM calculation log
            :param boolean l_save_mm_log: Logical for saving MM calculation log
            :param boolean l_save_scr: Logical for saving scratch directory
            :param string restart: Option for controlling dynamics restarting
        """
        # Initialize PyUNIxMD
        base_dir, unixmd_dir, qm_log_dir, mm_log_dir =\
             self.run_init(qm, mm, output_dir, l_save_qm_log, l_save_mm_log, l_save_scr, restart)
        bo_list = [self.rstate]
        qm.calc_coupling = True
        self.print_init(qm, mm, restart)

        if (restart == None):
            # Initialize decoherence variables
            self.append_sigma()

            # Calculate initial input geometry at t = 0.0 s
            self.istep = -1
            self.mol.reset_bo(qm.calc_coupling)
            qm.get_data(self.mol, base_dir, bo_list, self.dt, self.istep, calc_force_only=False)
            if (self.mol.l_qmmm and mm != None):
                mm.get_data(self.mol, base_dir, bo_list, self.istep, calc_force_only=False)
            if (not self.mol.l_nacme):
                self.mol.get_nacme()

            self.hop_prob()
            self.hop_check(bo_list)
            self.evaluate_hop(bo_list)
            if (qm.re_calc and self.l_hop):
                qm.get_data(self.mol, base_dir, bo_list, self.dt, self.istep, calc_force_only=True)
                if (self.mol.l_qmmm and mm != None):
                    mm.get_data(self.mol, base_dir, bo_list, self.istep, calc_force_only=True)

            self.update_energy()

            self.check_decoherence()
            self.check_coherence()
            self.aux_propagator()
            self.get_phase()
            if (self.l_collapse):
                self.check_decoherence()
                self.check_coherence()

            self.write_md_output(unixmd_dir, self.istep)
            self.print_step(self.istep)

        elif (restart == "write"):
            # Reset initial time step to t = 0.0 s
            self.istep = -1
            self.write_md_output(unixmd_dir, self.istep)
            self.print_step(self.istep)

        elif (restart == "append"):
            # Set initial time step to last successful step of previous dynamics
            self.istep = self.fstep

        self.istep += 1

        # Main MD loop
        for istep in range(self.istep, self.nsteps):

            self.calculate_force()
            self.cl_update_position()

            self.mol.backup_bo()
            self.mol.reset_bo(qm.calc_coupling)
            qm.get_data(self.mol, base_dir, bo_list, self.dt, istep, calc_force_only=False)
            if (self.mol.l_qmmm and mm != None):
                mm.get_data(self.mol, base_dir, bo_list, istep, calc_force_only=False)

            if (not self.mol.l_nacme and self.l_adj_nac):
                self.mol.adjust_nac()

            self.calculate_force()
            self.cl_update_velocity()

            if (not self.mol.l_nacme):
                self.mol.get_nacme()

            el_run(self)

            self.hop_prob()
            self.hop_check(bo_list)
            self.evaluate_hop(bo_list)
            if (self.l_hop):
                if (qm.re_calc):
                    qm.get_data(self.mol, base_dir, bo_list, self.dt, istep, calc_force_only=True)
                if (self.mol.l_qmmm and mm != None):
                    mm.get_data(self.mol, base_dir, bo_list, istep, calc_force_only=True)

            if (self.thermo != None):
                self.thermo.run(self)

            self.update_energy()

            self.check_decoherence()
            self.check_coherence()
            self.aux_propagator()
            self.get_phase()
            if (self.l_collapse):
                self.check_decoherence()
                self.check_coherence()

            if ((istep + 1) % self.out_freq == 0):
                self.write_md_output(unixmd_dir, istep)
            if ((istep + 1) % self.out_freq == 0 or len(self.event["HOP"]) > 0 or len(self.event["DECO"]) > 0):
                self.print_step(istep)
            if (istep == self.nsteps - 1):
                self.write_final_xyz(unixmd_dir, istep)

            self.fstep = istep
            restart_file = os.path.join(base_dir, "RESTART.bin")
            with open(restart_file, 'wb') as f:
                pickle.dump({'qm':qm, 'md':self}, f)

        # Delete scratch directory
        if (not l_save_scr):
            tmp_dir = os.path.join(unixmd_dir, "scr_qm")
            if (os.path.exists(tmp_dir)):
                shutil.rmtree(tmp_dir)

            if (self.mol.l_qmmm and mm != None):
                tmp_dir = os.path.join(unixmd_dir, "scr_mm")
                if (os.path.exists(tmp_dir)):
                    shutil.rmtree(tmp_dir)

    def hop_prob(self):
        """ Routine to calculate hopping probabilities

            :param integer istep: Current MD step
        """
        # Reset surface hopping variables
        self.rstate_old = self.rstate

        self.prob = np.zeros(self.mol.nst)
        self.acc_prob = np.zeros(self.mol.nst + 1)

        self.l_hop = False
        self.force_hop = False

        accum = 0.

        if (self.mol.rho.real[self.rstate, self.rstate] < self.lower_th):
            self.force_hop = True

        for ist in range(self.mol.nst):
            if (ist != self.rstate):
                if (self.force_hop):
                    self.prob[ist] = self.mol.rho.real[ist, ist] / self.upper_th
                else:
                    self.prob[ist] = - 2. * self.mol.rho.real[ist, self.rstate] * \
                        self.mol.nacme[ist, self.rstate] * self.dt / self.mol.rho.real[self.rstate, self.rstate]

                if (self.prob[ist] < 0.):
                    self.prob[ist] = 0.
                accum += self.prob[ist]
            self.acc_prob[ist + 1] = accum
        psum = self.acc_prob[self.mol.nst]

        if (psum > 1.):
            self.prob /= psum
            self.acc_prob /= psum

    def hop_check(self, bo_list):
        """ Routine to check hopping occurs with random number

            :param integer,list bo_list: List of BO states for BO calculation
        """
        self.rand = random.random()
        for ist in range(self.mol.nst):
            if (ist == self.rstate):
                continue
            if (self.rand > self.acc_prob[ist] and self.rand <= self.acc_prob[ist + 1]):
                self.l_hop = True
                self.rstate = ist
                bo_list[0] = self.rstate

    def evaluate_hop(self, bo_list):
        """ Routine to evaluate hopping and velocity rescaling

            :param integer,list bo_list: List of BO states for BO calculation
            :param integer istep: Current MD step
        """
        if (self.l_hop):
            # Calculate potential difference between hopping states
            pot_diff = self.mol.states[self.rstate].energy - self.mol.states[self.rstate_old].energy

            # Solve quadratic equation for scaling factor of velocities
            a = 1.
            b = 1.
            det = 1.
            if (self.hop_rescale == "velocity"):
                a = np.sum(self.mol.mass[0:self.mol.nat_qm] * np.sum(self.mol.nac[self.rstate_old, self.rstate] ** 2., axis=1))
                b = 2. * np.sum(self.mol.mass[0:self.mol.nat_qm] * np.sum(self.mol.nac[self.rstate_old, self.rstate] \
                    * self.mol.vel[0:self.mol.nat_qm], axis=1))
                c = 2. * pot_diff
                det = b ** 2. - 4. * a * c
            elif (self.hop_rescale == "momentum"):
                a = np.sum(1. / self.mol.mass[0:self.mol.nat_qm] * np.sum(self.mol.nac[self.rstate_old, self.rstate] ** 2., axis=1))
                b = 2. * np.sum(np.sum(self.mol.nac[self.rstate_old, self.rstate] * self.mol.vel[0:self.mol.nat_qm], axis=1))
                c = 2. * pot_diff
                det = b ** 2. - 4. * a * c
            elif (self.hop_rescale == "augment"):
                a = np.sum(1. / self.mol.mass[0:self.mol.nat_qm] * np.sum(self.mol.nac[self.rstate_old, self.rstate] ** 2., axis=1))
                b = 2. * np.sum(np.sum(self.mol.nac[self.rstate_old, self.rstate] * self.mol.vel[0:self.mol.nat_qm], axis=1))
                c = 2. * pot_diff
                det = b ** 2. - 4. * a * c

            # Default: hopping is allowed
            self.l_reject = False

            # Velocities cannot be adjusted when zero kinetic energy is given
            if (self.hop_rescale == "energy" and self.mol.ekin_qm < eps):
                self.l_reject = True
            # Clasically forbidden hop due to lack of kinetic energy
            if (self.mol.ekin_qm < pot_diff):
                self.l_reject = True
            # Kinetic energy is enough, but there is no solution for scaling factor
            if (det < 0.):
                self.l_reject = True
            # When kinetic energy is enough, velocities are always rescaled in 'augment' case
            if (self.hop_rescale == "augment" and self.mol.ekin_qm > pot_diff):
                self.l_reject = False

            if (self.l_reject):
                # Record event for frustrated hop
                if (self.mol.ekin_qm < pot_diff):
                    self.event["HOP"].append(f"Reject hopping: smaller kinetic energy than potential energy difference between {self.rstate} and {self.rstate_old}")
                # Set scaling constant with respect to 'hop_reject'
                if (self.hop_reject == "keep"):
                    self.event["HOP"].append("Reject hopping: no solution to find rescale factor, velocity is not changed")
                elif (self.hop_reject == "reverse"):
                    # x = - 1 when 'hop_rescale' is 'energy', otherwise x = - b / a
                    self.event["HOP"].append("Reject hopping: no solution to find rescale factor, velocity is reversed along coupling direction")
                    x = - b / a
                # Recover old running state
                self.l_hop = False

                if (self.force_hop):
                    self.event["HOP"].append(f"Collapse density: reset the density according to the current state {self.rstate_old}")
                    self.set_decoherence(self.rstate_old)

                self.force_hop = False

                self.rstate = self.rstate_old
                bo_list[0] = self.rstate
            else:
                if (self.hop_rescale == "energy" or (det < 0. and self.hop_rescale == "augment")):
                    if (det < 0.):
                        self.event["HOP"].append("Accept hopping: no solution to find rescale factor, but velocity is simply rescaled")
                    x = np.sqrt(1. - pot_diff / self.mol.ekin_qm)
                else:
                    if (b < 0.):
                        x = 0.5 * (- b - np.sqrt(det)) / a
                    else:
                        x = 0.5 * (- b + np.sqrt(det)) / a

            # Rescale velocities for QM atoms
            if (not (self.hop_reject == "keep" and self.l_reject)):
                if (self.hop_rescale == "energy"):
                    self.mol.vel[0:self.mol.nat_qm] *= x

                elif (self.hop_rescale == "velocity"):
                    self.mol.vel[0:self.mol.nat_qm] += x * self.mol.nac[self.rstate_old, self.rstate]

                elif (self.hop_rescale == "momentum"):
                    self.mol.vel[0:self.mol.nat_qm] += x * self.mol.nac[self.rstate_old, self.rstate] / \
                        self.mol.mass[0:self.mol.nat_qm].reshape((-1, 1))

                elif (self.hop_rescale == "augment"):
                    if (det > 0. or self.mol.ekin_qm < pot_diff):
                        self.mol.vel[0:self.mol.nat_qm] += x * self.mol.nac[self.rstate_old, self.rstate] / \
                            self.mol.mass[0:self.mol.nat_qm].reshape((-1, 1))
                    else:
                        self.mol.vel[0:self.mol.nat_qm] *= x

            # Update kinetic energy
            self.mol.update_kinetic()

        # Record hopping event
        if (self.rstate != self.rstate_old):
            if (self.force_hop):
                self.event["HOP"].append(f"Accept hopping: force hop {self.rstate_old} -> {self.rstate}")
            else:
                self.event["HOP"].append(f"Accept hopping: hop {self.rstate_old} -> {self.rstate}")

    def calculate_force(self):
        """ Routine to calculate the forces
        """
        self.rforce = np.copy(self.mol.states[self.rstate].force)

    def update_energy(self):
        """ Routine to update the energy of molecules in surface hopping dynamics
        """
        # Update kinetic energy
        self.mol.update_kinetic()
        self.mol.epot = self.mol.states[self.rstate].energy
        self.mol.etot = self.mol.epot + self.mol.ekin

    def check_decoherence(self):
        """ Routine to check if the electronic state is decohered
        """
        if (self.l_hop):
            if (True in self.l_coh):
                self.event["DECO"].append(f"Destroy auxiliary trajectories: hopping occurs")
            self.l_coh = [False] * self.mol.nst
            self.l_first = [False] * self.mol.nst
            self.l_fix = [False] * self.mol.nst
        else:
            for ist in range(self.mol.nst):
                if (self.l_coh[ist]):
                    rho = self.mol.rho.real[ist, ist]
                    if (rho > self.upper_th):
                        self.set_decoherence(ist)
                        return

    def check_coherence(self):
        """ Routine to check coherence among BO states
        """
        count = 0
        tmp_st = ""
        for ist in range(self.mol.nst):
            rho = self.mol.rho.real[ist, ist]
            if (rho > self.upper_th or rho < self.lower_th):
                self.l_coh[ist] = False
            else:
                if (self.l_coh[ist]):
                    self.l_first[ist] = False
                else:
                    self.l_first[ist] = True
                    tmp_st += f"{ist}, "
                self.l_coh[ist] = True
                count += 1

        if (count < 2):
            self.l_coh = [False] * self.mol.nst
            self.l_first = [False] * self.mol.nst
            tmp_st = ""

        if (len(tmp_st) >= 1):
            tmp_st = tmp_st.rstrip(', ')
            self.event["DECO"].append(f"Generate auxiliary trajectory on {tmp_st} state")

    def set_decoherence(self, one_st):
        """ Routine to reset coefficient/density if the state is decohered

            :param integer one_st: State index that its population is one
        """
        self.phase = np.zeros((self.mol.nst, self.aux.nat, self.aux.ndim))
        self.mol.rho = np.zeros((self.mol.nst, self.mol.nst), dtype=np.complex128)
        self.mol.rho[one_st, one_st] = 1. + 0.j

        self.l_coh = [False] * self.mol.nst
        self.l_first = [False] * self.mol.nst
        self.l_fix = [False] * self.mol.nst

        self.event["DECO"].append(f"Destroy auxiliary trajectories: decohered to {one_st} state")

        if (self.elec_object == "coefficient"):
            for ist in range(self.mol.nst):
                if (ist == one_st):
                    self.mol.states[ist].coef /= np.absolute(self.mol.states[ist].coef).real
                else:
                    self.mol.states[ist].coef = 0. + 0.j

    def aux_propagator(self):
        """ Routine to propagate auxiliary molecule
        """
        # Get auxiliary position
        for ist in range(self.mol.nst):
            if (self.l_coh[ist]):
                if (self.l_first[ist]):
                    if (self.l_xf1d):
                        self.aux.pos[ist] = np.zeros((self.aux.nat, self.aux.ndim))
                    else:
                        self.aux.pos[ist] = self.mol.pos[0:self.aux.nat]
                else:
                    if (self.l_xf1d):
                        self.aux.pos[ist] += self.aux.vel[ist] * self.dt
                    else:
                        if (ist == self.rstate):
                            self.aux.pos[ist] = self.mol.pos[0:self.aux.nat]
                        else:
                            self.aux.pos[ist] += self.aux.vel[ist] * self.dt

        self.pos_0 = np.copy(self.aux.pos[self.rstate])

        # Get auxiliary velocity
        self.l_collapse = False
        self.aux.vel_old = np.copy(self.aux.vel)
        for ist in range(self.mol.nst):
            # Calculate propagation factor alpha
            if (self.l_coh[ist]):
                if (self.l_fix[ist]):
                    alpha = 0.
                else:
                    if (ist == self.rstate):
                        alpha = self.mol.ekin_qm
                    else:
                        if (self.l_first[ist]):
                            alpha = self.mol.ekin_qm
                            if (self.l_econs_state):
                                alpha += self.mol.states[self.rstate].energy - self.mol.states[ist].energy
                        else:
                            ekin_old = np.sum(0.5 * self.aux.mass * np.sum(self.aux.vel_old[ist] ** 2, axis=1))
                            alpha = ekin_old + self.mol.states[ist].energy_old - self.mol.states[ist].energy
                    if (alpha < 0.):
                        alpha = 0.
                        if (self.aux_econs_viol == "fix"):
                            self.l_fix[ist] = True
                            self.event["DECO"].append(f"Energy conservation violated, the auxiliary trajectory on state {ist} is fixed.")
                        elif (self.aux_econs_viol == "collapse"):
                            self.l_collapse = True
                            self.collapse(ist)
                            self.event["DECO"].append(f"Energy conservation violated, collapse the {ist} state coefficient/density to zero.")

                # Calculate auxiliary velocity from alpha
                if (self.l_xf1d):
                    alpha /= 0.5 * self.aux.mass[0]
                    self.aux.vel[ist] = np.sqrt(alpha)
                else:
                    alpha /= self.mol.ekin_qm
                    self.aux.vel[ist] = self.mol.vel[0:self.aux.nat] * np.sqrt(alpha)

    def collapse(self, cstate):
        """ Routine to collapse coefficient/density of a state to zero
        """
        fac = 1. - self.mol.rho.real[cstate, cstate]

        if (self.elec_object == "coefficient"):
            for ist in range(self.mol.nst):
                if (ist == cstate):
                    self.mol.states[ist].coef = 0. + 0.j
                else:
                    self.mol.states[ist].coef /= np.sqrt(fac)

        self.mol.rho[cstate,:] = 0. + 0.j
        self.mol.rho[:,cstate] = 0. + 0.j
        self.mol.rho /= fac

    def get_phase(self):
        """ Routine to calculate phase term
        """
        for ist in range(self.mol.nst):
            if (self.l_coh[ist]):
                if (self.l_first[ist]):
                    self.phase[ist] = 0.
                else:
                    for iat in range(self.aux.nat):
                        self.phase[ist, iat] += self.aux.mass[iat] * \
                            (self.aux.vel[ist, iat] - self.aux.vel_old[ist, iat])

    def append_sigma(self):
        """ Routine to append sigma values when single float number is provided
        """
        # Create a list from single float number
        if (isinstance(self.sigma, float)):
            sigma = self.sigma
            self.sigma = self.aux.nat * [sigma]

    def write_md_output(self, unixmd_dir, istep):
        """ Write output files

            :param string unixmd_dir: PyUNIxMD directory
            :param integer istep: Current MD step
        """
        # Write the common part
        super().write_md_output(unixmd_dir, istep)

        # Write hopping-related quantities
        self.write_sh(unixmd_dir, istep)

        # Write time-derivative BO population
        self.write_dotpop(unixmd_dir, istep)

        # Write decoherence information
        self.write_dec(unixmd_dir, istep)

    def write_sh(self, unixmd_dir, istep):
        """ Write hopping-related quantities into files

            :param string unixmd_dir: PyUNIxMD directory
            :param integer istep: Current MD step
        """
        # Write SHSTATE file
        tmp = f'{istep + 1:9d}{"":14s}{self.rstate}'
        typewriter(tmp, unixmd_dir, "SHSTATE", "a")

        # Write SHPROB file
        tmp = f'{istep + 1:9d}' + "".join([f'{self.prob[ist]:15.8f}' for ist in range(self.mol.nst)])
        typewriter(tmp, unixmd_dir, "SHPROB", "a")

    def write_dotpop(self, unixmd_dir, istep):
        """ Write time-derivative BO population

            :param string unixmd_dir: PyUNIxMD directory
            :param integer istep: Current MD step
        """
        if (self.verbosity >= 1):
            # Write NAC term in DOTPOPNAC
            tmp = f'{istep + 1:9d}' + "".join([f'{pop:15.8f}' for pop in self.dotpopnac])
            typewriter(tmp, unixmd_dir, "DOTPOPNAC", "a")

            # Write decoherence term in DOTPOPDEC
            tmp = f'{istep + 1:9d}' + "".join([f'{pop:15.8f}' for pop in self.dotpopdec])
            typewriter(tmp, unixmd_dir, "DOTPOPDEC", "a")

    def write_dec(self, unixmd_dir, istep):
        """ Write XF-based decoherence information

            :param string unixmd_dir: PyUNIxMD directory
            :param integer istep: Current MD step
        """
        # Write auxiliary trajectories
        if (self.verbosity >= 2 and True in self.l_coh):
            # Write quantum momenta
            tmp = f'{self.aux.nat:6d}\n{"":2s}Step:{istep + 1:6d}{"":12s}Momentum (au)' + \
                "".join(["\n" + f'{self.aux.symbols[iat]:5s}' + \
                "".join([f'{self.qmom[iat, isp]:15.8f}' for isp in range(self.aux.ndim)]) for iat in range(self.aux.nat)])
            typewriter(tmp, unixmd_dir, f"QMOM", "a")

            # Write auxiliary variables
            for ist in range(self.mol.nst):
                if (self.l_coh[ist]):
                    # Write auxiliary phase
                    tmp = f'{self.aux.nat:6d}\n{"":2s}Step:{istep + 1:6d}{"":12s}Phase (au)' + \
                        "".join(["\n" + f'{self.aux.symbols[iat]:5s}' + \
                        "".join([f'{self.phase[ist, iat, isp]:15.8f}' for isp in range(self.aux.ndim)]) for iat in range(self.aux.nat)])
                    typewriter(tmp, unixmd_dir, f"AUX_PHASE_{ist}", "a")

                    # Write auxiliary trajectory movie files
                    tmp = f'{self.aux.nat:6d}\n{"":2s}Step:{istep + 1:6d}{"":12s}Position(A){"":34s}Velocity(au)' + \
                        "".join(["\n" + f'{self.aux.symbols[iat]:5s}' + \
                        "".join([f'{self.aux.pos[ist, iat, isp] * au_to_A:15.8f}' for isp in range(self.aux.ndim)]) + \
                        "".join([f"{self.aux.vel[ist, iat, isp]:15.8f}" for isp in range(self.aux.ndim)]) for iat in range(self.aux.nat)])
                    typewriter(tmp, unixmd_dir, f"AUX_MOVIE_{ist}.xyz", "a")

    def print_init(self, qm, mm, restart):
        """ Routine to print the initial information of dynamics

            :param object qm: QM object containing on-the-fly calculation infomation
            :param object mm: MM object containing MM calculation infomation
            :param string restart: Option for controlling dynamics restarting
        """
        # Print initial information about molecule, qm, mm and thermostat
        super().print_init(qm, mm, restart)

        # Print dynamics information for start line
        dynamics_step_info = textwrap.dedent(f"""\

        {"-" * 118}
        {"Start Dynamics":>65s}
        {"-" * 118}
        """)

        # Print INIT for each step
        INIT = f" #INFO{'STEP':>8s}{'State':>7s}{'Kinetic(H)':>14s}{'Potential(H)':>15s}{'Total(H)':>13s}{'Temperature(K)':>17s}{'Norm.':>8s}"
        dynamics_step_info += INIT

        # Print DEBUG1 for each step
        if (self.verbosity >= 1):
            DEBUG1 = f" #DEBUG1{'STEP':>6s}{'Rand.':>11s}{'Acc. Hopping Prob.':>28s}"
            dynamics_step_info += "\n" + DEBUG1

        print (dynamics_step_info, flush=True)

    def print_step(self, istep):
        """ Routine to print each steps infomation about dynamics

            :param integer istep: Current MD step
        """
        ctemp = self.mol.ekin * 2. / float(self.mol.ndof) * au_to_K
        norm = 0.
        for ist in range(self.mol.nst):
            norm += self.mol.rho.real[ist, ist]

        # Print INFO for each step
        INFO = f" INFO{istep + 1:>9d}{self.rstate:>5d}"
        INFO += f"{self.mol.ekin:16.8f}{self.mol.epot:15.8f}{self.mol.etot:15.8f}"
        INFO += f"{ctemp:13.6f}"
        INFO += f"{norm:11.5f}"
        print (INFO, flush=True)

        # Print DEBUG1 for each step
        if (self.verbosity >= 1):
            DEBUG1 = f" DEBUG1{istep + 1:>7d}"
            DEBUG1 += f"{self.rand:11.5f}"
            for ist in range(self.mol.nst):
                DEBUG1 += f"{self.acc_prob[ist]:12.5f} ({self.rstate}->{ist})"
            print (DEBUG1, flush=True)

        # Print event in SHXF
        for category, events in self.event.items():
            if (len(events) != 0):
                for ievent in events:
                    print (f" {category}{istep + 1:>9d}  {ievent}", flush=True)
        self.event["HOP"] = []
        self.event["DECO"] = []