import numpy as np
from pyat.AcousticPara import GetEnv, Coord
from pyat.DataBase import DataBase
from pyat.bellhop import call_Bellhop
from pyat.krakenc import call_krakenc
from pyat.ram import call_ram
import matplotlib.pyplot as plt

def call_tl(database, freq, sourceDepth, sourceLat, sourceLon,
 receiveLat=None, receiveLon=None, range=None, azi=None, modelType='kraken', isPlot=False):
    """ Input:
    freq - float, Frequency
    sourceDepth - float, Source depth
    sourceLat - float, Source lat
    sourceLon - float, Source lon
    receiveDepth - float, Receive depth
    receiveLat - float, Receive lat
    receiveLon - float, Receive lon
    range - float, Range from source to receive
    azi - float, azimuth angle from receive to source
    timeType - char, #全年平均 'Annual', #1-12月'January''February''March''April''May''June''July''August''September'
        October': '10''November''December', #北半球四季'Winter''Spring''Summer''Autumn'
    modelType - char, 'bellhop'、'kraken'、'ram'
    isPlot - plot or not
    """
    """ Output:
    ranges 距离
    depths 深度
    TL 传播损失
    bathym类 地形数据 bathym.r距离 bathym.d深度
    """

    bathym, ssp, sed, base = [], [], [], []
    if (receiveLat != None) and (receiveLon != None):
        [bathym, ssp, sed, base] = GetEnv(database, 25, Coord(sourceLon, sourceLat), coordR=Coord(receiveLon, receiveLat))
    elif (range != None) and (azi != None):
        [bathym, ssp, sed, base] = GetEnv(database, 25, Coord(sourceLon, sourceLat), R=range, azi=azi)
    else:
        print('no receive info!')
        return [], [], [], bathym
    if modelType == 'bellhop':
        sd = np.array([sourceDepth])           # Source depth
        [Pos, TL] = call_Bellhop(freq, sd, bathym, ssp, sed, base)
        # Plot
        if isPlot:
            plt.pcolor(Pos.r.range, Pos.r.depth, TL[0,0,:,:], vmin=40, vmax=130, cmap='jet_r', shading='auto')
            plt.xlabel("Range/m")
            plt.ylabel("Depth/m")
            plt.gca().invert_yaxis()
            plt.colorbar()
            plt.show()
        return Pos.r.range, Pos.r.depth, TL[0,0,:,:], bathym
    elif modelType == 'kraken':
        sd = np.array([sourceDepth])           # Source depth
        [Pos, TL] = call_krakenc(freq, sd, bathym, ssp, sed, base)
        # Plot
        if isPlot:
            plt.pcolor(Pos.r.range, Pos.r.depth, TL[0,0,:,:], vmin=40, vmax=130, cmap='jet_r', shading='auto')
            plt.xlabel("Range/m")
            plt.ylabel("Depth/m")
            plt.gca().invert_yaxis()
            plt.colorbar()
            plt.show()
        return Pos.r.range, Pos.r.depth, TL[0,0,:,:], bathym
    elif modelType == 'ram':
        # 调用ram模块
        [rd, rr, TL] = call_ram(freq, sourceDepth, bathym, ssp, sed)
        # Plot
        if isPlot:
            plt.pcolormesh(rr, rd, TL, vmin=60, vmax=130, cmap='jet_r', shading='auto')
            plt.xlabel("Range/m")
            plt.ylabel("Depth/m")
            plt.gca().invert_yaxis()
            plt.colorbar()
            plt.show()
        return rr, rd, TL, bathym
    else:
        return [], [], [], bathym



# # 创建类
# database = DataBase()
# # 设置时间段
# database.set_timePeriod('06')
# # 设置时间段
# #database.set_timePeriod(timeIndex)
# range, depth, tl, bathym = call_tl(database=database,freq=200, sourceDepth=60, sourceLat=19, sourceLon=127,
#                                    receiveLat=17.5, receiveLon=127.1,  modelType='ram', isPlot=True)