from ramRW import PosRAM, PROF, SSPRAM, write_env_RAM, readTL_RAM
from os import system
import numpy as np
from config import AtBinPath

def calZmax(ssp, sed):
    d = np.array([])
    for it in range(len(ssp)):
        if sed == None:
            d = np.append(d, ssp[it].z[-1])
        else:
            d = np.append(d, ssp[it].z[-1] + sed[it].z[-1])
    Zmax = max(d)
    return Zmax

def call_ram(freq, sd, bathm, ssp, sed, dr=0.1):
    ''' ********************************************************************************* '''
    ''' ************************** Sound Field Calculation ****************************** '''
    ''' ********************************************************************************* '''
    title = "Waveguide"
    # Source and receiving position
    bathm.r = bathm.r * 1000        # km -> m
    rmax = max(bathm.r)             # Maximum receiving range (unit: m)
    zmax = round(max(bathm.d) * 2)  # Maximum receiving depth (uint: m)
    zmplt = round(max(bathm.d)*1.1)     # Maximum depth of plot (uint: m)
    dz = 10          # Step in depth direction: wavelength / 15
    dr = 1000* dr                     # Step in range direction
    Zmax = calZmax(ssp, sed)
    rd = Zmax
    pos = PosRAM(sd, rd, rmax, zmax, dr, dz, zmplt)

    # Sound speed profile for RAM
    ssp_ram = []
    for it in range(len(ssp)):
        rp = bathm.r[it]                            # Range of this profile
        Z = sed[it].z
        cw_p = PROF(ssp[it].z, ssp[it].c)           # Water sound speed profile
        cb_p = PROF(Z, sed[it].cp)                  # Bottom sound speed profile
        rhob_p = PROF(Z, sed[it].rho)               # Sea bottom density profile
        Z = np.append(Z, Z[-1] + 100)
        Ap = np.append(sed[it].a_p, 10)
        # Ap = sed[it].a_p
        attn_p = PROF(Z, Ap)               # Sea bottom attenuation profile

        ssp_ram.append(SSPRAM(rp, cw_p, cb_p, rhob_p, attn_p))


    # Write ram.in file
    write_env_RAM(title, freq, pos, bathm, ssp_ram)
    # Call ram.exe
    system(AtBinPath + "/ram")
    # Read tl.grid
    return readTL_RAM()


# if __name__=="__main__":
    # ''' ********************************************************************************* '''
    # ''' ***************************** Input Parameters ********************************** '''
    # ''' ********************************************************************************* '''
    # freq = 25                   # Frequency
    # sd = 40.0                   # Source depth
    # rd = 30.0                   # Receiving depth

    # # 定义坐标
    # lat = 18.00
    # lon = 114.00
    # # 定义时间段
    # timeIndex = db.timeDict['Annual']
    # # 创建类
    # data = db.DataBase()
    # # 设置时间段
    # data.set_timePeriod(timeIndex)

    # [bathm, ssp, sed, base] = GetEnv(data, 25, Coord(lon, lat), R=100, azi=5)

    # # 调用ram模块
    # [rd, rr, TL] = call_ram(freq, sd, rd, bathm, ssp, sed)

    # # Plot
    # plt.pcolormesh(rr, rd, TL, vmin=40, vmax=90, cmap='jet_r', shading='auto')
    # plt.xlabel("Range/m")
    # plt.ylabel("Depth/m")
    # plt.gca().invert_yaxis()
    # plt.colorbar()
    # plt.show()