#!/usr/bin/env python3

################################################################################
# 
#  Copyright 2009 National Technology & Engineering Solutions of Sandia, LLC
#  (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the U.S.
#  Government retains certain rights in this software.
# 
#  BSD Open Source License
#  All rights reserved.
# 
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
# 
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
# 
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
# 
#    3. Neither the name of the copyright holder nor the names of its
#       contributors may be used to endorse or promote products derived from
#       this software without specific prior written permission.
# 
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
#  POSSIBILITY OF SUCH DAMAGE.
# 
################################################################################

# import statements
import sys  # for exiting w/codes

# first try and import rstt. if that fails, search the environment variables
# for the appropriate path and try and import from there.
try:
    import rstt

# it looks like importing failed. try searching paths from environment variables for the module
except ModuleNotFoundError:
    import os, sys                                                  # import some required modules
    for d in ('RSTT_ROOT', 'RSTT_HOME', 'SLBM_ROOT', 'SLBM_HOME'):  # list of env vars to check, in order
        if os.getenv(d):                                            # if the environmental var is set
            sys.path.append(os.getenv(d) + "/lib")                  # add to the start of the module search path
    
    # try and import rstt again
    try:
        import rstt

    # if it still fails, raise an exception with a helpful error message
    except ModuleNotFoundError as e:
        e.msg += ". Have you set $RSTT_ROOT, $SLBM_ROOT, or installed RSTT with pip?"
        raise e from None


# declare some variables relevant to testing
PASSED = True   # bool for whether all the tests passed (default true)
EPS    = 1.e-6  # small value to allow for machine error differences


# function to check if two floats are equal to some high precision
def assertIsEqual(x: float, y: float, varname: str):

    # we have to specify global, here, because we're assigning a value to it
    global PASSED

    # print status message
    print("{:21s}  {:12.8f} == {:12.8f}".format(varname, x, y), end='')

    # if the two values are equal
    if (abs(x-y) < EPS):
        print(" OK")
    else:
        print(" FAIL")
        PASSED = False


# main program
if __name__ == "__main__":

    # wrap everything in a try-except block
    try:

        # print a quick message
        print("======================================================")
        print("Performing Python tests...")
        print("------------------------------------------------------")
        print("Variable                 Calculated       Reference")
        print("------------------------------------------------------")

        # arbitrary source (Meteor Crater, AZ)
        srcLatDeg =   35.0274  # latitude (degrees)
        srcLonDeg = -111.0228  # longitude (degrees)
        srcDepKm  =    1.2345  # depth (km)

        # arbitrary receiver (Albuquerque, NM)
        rcvLatDeg =   35.1053  # latitude (degrees)
        rcvLonDeg = -106.6294  # longitude (degrees)
        rcvDepKm  =   -1.6000  # depth (km)

        # convert lat/lon from degrees to radians
        srcLatRad = rstt.deg2rad(srcLatDeg)
        srcLonRad = rstt.deg2rad(srcLonDeg)
        rcvLatRad = rstt.deg2rad(rcvLatDeg)
        rcvLonRad = rstt.deg2rad(rcvLonDeg)

        # instantiate an RSTT object
        slbm = rstt.SlbmInterface()

        # load the velocity model
        slbm.loadVelocityModel("models/unittest.geotess")

        # loop over each type of phase
        phases = ["Pn", "Sn", "Pg", "Lg"]
        for i in range(len(phases)):

            # print status message
            if (i > 0):
                print("")
            print("Computing values for {:s} phase...".format(phases[i]))

            # create a great circle from source to the receiver
            slbm.createGreatCircle(phases[i],
                srcLatRad, srcLonRad, srcDepKm,
                rcvLatRad, rcvLonRad, rcvDepKm)

            # get the distance, travel time, and uncertainties from source --> receiver
            distRad               = slbm.getDistance()
            travelTimeSec         = slbm.getTravelTime()
            travelTimeUncertSec   = slbm.getTravelTimeUncertainty()
            travelTime1DUncertSec = slbm.getTravelTimeUncertainty1D()

            # make sure all the values are correct
            if   i == 0:
                assertIsEqual(distRad,                0.06290927, "distRad")
                assertIsEqual(travelTimeSec,         56.88549071, "travelTimeSec")
                assertIsEqual(travelTimeUncertSec,    3.12378957, "travelTimeUncertSec")
                assertIsEqual(travelTime1DUncertSec,  1.15423530, "travelTime1DUncertSec")

            elif i == 1:
                assertIsEqual(distRad,                0.06290927, "distRad")
                assertIsEqual(travelTimeSec,         98.95213093, "travelTimeSec")
                assertIsEqual(travelTimeUncertSec,    6.29314026, "travelTimeUncertSec")
                assertIsEqual(travelTime1DUncertSec,  2.47666863, "travelTime1DUncertSec")

            elif i == 2:
                assertIsEqual(distRad,                0.06290927, "distRad")
                assertIsEqual(travelTimeSec,         70.40335768, "travelTimeSec")
                assertIsEqual(travelTimeUncertSec,   15.58790453, "travelTimeUncertSec")
                assertIsEqual(travelTime1DUncertSec,  1.61320256, "travelTime1DUncertSec")

            elif i == 3:
                assertIsEqual(distRad,                0.06290927, "distRad")
                assertIsEqual(travelTimeSec,        119.65236722, "travelTimeSec")
                assertIsEqual(travelTimeUncertSec,   20.08714726, "travelTimeUncertSec")
                assertIsEqual(travelTime1DUncertSec,  1.98573827, "travelTime1DUncertSec")

    # catch possible exceptions
    except Exception as e:

        # raise the error we just caught
        raise(e)

    # return
    if PASSED:
        print("------------------------------------------------------")
        print("Python tests PASSED.")
        print("======================================================")
        sys.exit(0)

    else:
        print("------------------------------------------------------")
        print("Python tests FAILED.")
        print("======================================================")
        sys.exit(1)
