from pylab import *

# A function for Mackenzie's model of speed of sound in water in m/s
def mackenzie_sos(tempC,depthM,salinityPPT):
    sos = 1448.96 + 4.591*tempC - 5.304*(10**-2)*(tempC**2) + 2.374*(10**-4)*(tempC**3) + 1.34*(salinityPPT-35) + 1.63*(10**-2)*depthM + 1.675*(10**-7)*(depthM**2) - 1.025*(10**-2)*tempC*(salinityPPT-35) - 7.139*(10**-13)*tempC*(depthM**3)
    return (sos)
# A function for the salinity of seawater in parts per thousand
def sal_of_seawater(latitude):
    salinityPPT = (1/45)*latitude + 36
    return (salinityPPT)
# A function for the temperature of seawater in degrees Celsius
TempSurfaceZone = 24-(11/45)*latitude # Temperature in the surface zone
def temp_of_seawater(TempSurfaceZone,depthM): # Temperature in the thermocline (dependent on the surface zone temp)
    tempC = ((2-TempSurfaceZone)/800)*depthM + (TempSurfaceZone-((2-TempSurfaceZone)/800)*200)
    return (tempC)
# A function for the distance to an objected detected using echolocation
def distance_to_object(ReturnTimeSec,sos):
    DistanceM = (ReturnTimeSec/2)*sos
    return (DistanceM)

# A welcome and prompt for patron to enter their preferred complexity for the program's info
print("Welcome to the South Bank Science Museum interactive echolocation program!")
patronType = float(input("Select which level of complexity you would like for information in the program. Enter 0 (simple), 1 (moderate), or 2 (complex): "))

# Separate welcome messages and an explanation of latitude for the 3 patron types
if patronType == 1: # Seasoned patron
    print("Regular")
    print("Latitude is the angle based on the distance somewhere is from the Earth's equator.")
elif patronType == 2: # Grizzled patron
    print("Expert")
    print("Latitude is the angle between a point on the Earth's surface and point on the surface at the equator. Latitude lines form concentric circles parallel to one another heading from the equator. A latitude in the northern hemisphere may be denoted by either 'N' or '+', and 'S' or '-' in the southern hemisphere.")
else: # Rookie patron
    print("Rookie")
    print("Latitude is the angle for the distance someone is away from the ring around Earth's widest point called the 'equator'.")

# Use of the variable i so the patron has the option to loop back to input another latitude if they desire
i = 1
while i == 1:
    # Common info on latitude for the 3 patron types and a prompt to enter the latitude
    print("Latitude starts at 0° at the equator and increases going towards the north and south poles where it is 90°, with an 'N' for north and 'S' for south. For example, the latitude in Brisbane is 27.4698° S.")
    latitude = float(input("Enter a latitude to explore the sperm whale's echolocation (without the N or S): "))
    # Messages if the patron enters a latitude outside the bounds of [0, 90]
    if latitude < 0:
        print("Sorry, this latitude was below 0. Please enter a value between 0 and 90° to run the program.")
    elif latitude > 90:
        print("Sorry, this latitude was above 90. Please enter a value between 0 and 90° to run the program.")
    else:
        salinityPPT = sal_of_seawater(latitude)
        print("The salinity at",latitude,"degrees latitude is",salinityPPT,"parts per thousand.")
        # Graphing seawater temperature vs depth at the given latitude for the grizzled patron
        if patronType == 2:
            # Plotting surface zone temperature
            DepthSurfaceM = arange(0,201,200)
            ArraySurfaceZone = array([TempSurfaceZone, TempSurfaceZone])
            plot(DepthSurfaceM, ArraySurfaceZone, linewidth = 2, label="Surface Zone")
            # Plotting thermocline temperature
            DepthThermoclineM = array([200, 1000.])
            XD = (2-TempSurfaceZone)/800 # Expression repeated multiple times in the following equation, storing it in a variable makes the program more efficient
            TempThermocline = array([(XD*DepthThermoclineM[0] + TempSurfaceZone-XD*200), (XD*DepthThermoclineM[1] + TempSurfaceZone-XD*200)])
            plot(DepthThermoclineM, TempThermocline, linewidth = 2, label="Thermocline")
            # Plotting deep zone temperature
            DepthDeepM = arange(1000,2001,200)
            TempDeepZone = zeros(6)
            plot(DepthDeepM, TempDeepZone+2, linewidth = 2, label="Deep Zone")
            # Graph communication
            title("Temperature vs. Depth at "+str(latitude)+"° Latitude")
            xlabel("Depth (m)")
            ylabel("Temperature (degrees C)")
            grid(True)
            legend()
            show()

        # Prompt patron to enter a depth
        depthM = float(input("Enter a depth in metres at which you would like to model sperm whale echolocation: "))
        # If the depth is in the surface zone
        if depthM >= 0 and depthM <= 200:
            tempC = TempSurfaceZone
        # If the depth is in the thermocline
        elif depthM > 200 and depthM < 1000:
            tempC = temp_of_seawater(TempSurfaceZone,depthM)
        # If the depth is in the deep zone
        else:
            tempC = 2
        # Calculating the speed of sound at patron's input latitude and depth
        sos = mackenzie_sos(tempC,depthM,salinityPPT)

        # Speed of sound, salinity, depth, and temperature data
        print("The speed of sound underwater is",round(sos,2),"m/s at",latitude,"° latitude and",depthM,"m depth, where the water is",round(tempC,2),"°C and the salinity is",round(salinityPPT,2),"parts per thousand.")

        # Separate explanations of return time for the 3 patron types
        if patronType == 1: # Seasoned patron
            print("Return time is how long the sperm whale waits between it emits a sonar signal and the signal returning.")
        elif patronType == 2: # Grizzled patron
            print("Return time is the duration of time from when the sperm whale emits its sonar which reflects off an object and returns and is detected by the whale's senses.")
        else: # Rookie patron
            print("Return time is how long the sperm whale waits after it sends out its sonar signal to when it returns and hears it.")

        # A loop and conditional so if the patron enters a return time of an object over 50km away, they are told to try again
        DistanceM = 50001
        while DistanceM > 50000:
            ReturnTimeSec = float(input("Enter a return time in seconds (sonar travels very fast, so a small time such as under a 1 minute is ideal): "))
            DistanceM = distance_to_object(ReturnTimeSec,sos)
            if DistanceM > 50000:
                print("The underwater object is too far away!")
            else:
                print("The sperm whale sends out a sonar signal and hears it",ReturnTimeSec,"seconds later, which means it detects an underwater object",round(DistanceM,0),"m away.")

    # i == 1 means the program loops back to where the patron can enter another latitude
    i = i + float(input("Would you like to explore sperm whale echolocation at another latitude? If YES, enter 0; if NO, enter any other number: "))

# Farewell message
print("Thank you for using the South Bank Science Museum sperm whale echolocation program!")