from os import system
import numpy as np
from pyat.readwrite import write_env, write_fieldflp, read_shd
from pyat.env import Pos, Source, Dom, cInt, SSPraw, SSP, HS, BotBndry, TopBndry, Bndry
from pyat.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_krakenc(freq, sd, rd, bathm, ssp, sed, base):
    filename = "envK"
    envFilename = "envK.env"
    model = "KRAKEN"
    title = "Krakenc Waveguide"

    # Source and receiving position
    Rmax = bathm.r[-1]       # unit: km
    dr = (1500 / freq / 3) / 1000  # wavelength / 10, unit: km
    ran = np.linspace(0, Rmax, round(Rmax / dr))  # Receiving range (unit: km)
    Zmax = calZmax(ssp, sed)
    RD = np.linspace(0, Zmax, round(Zmax) + 1)
    pos = Pos(Source(sd), Dom(ran, RD))

    # Range of phase velocity
    cInt.Low = 1400
    cInt.High = 15000

    # Write environmental file
    for it in range(len(ssp)):
        # The number of media
        if sed != None:
            NMedia = 2
        else:
            NMedia = 1

        ssp_raw = []
        depth = [0]
        # Sound speed profile
        Z = ssp[it].z                   # Depth
        Cp = ssp[it].c                  # Speed of P-wave
        Cs = np.zeros(np.shape(Z))      # Speed of S-wave
        Rho = np.ones(np.shape(Z))      # Density of the media
        Ap = np.zeros(np.shape(Z))      # Attenuation of P-wave
        As = np.zeros(np.shape(Z))      # Attenuation of S-wave
        ssp_raw.append(SSPraw(Z, Cp, Cs, Rho, Ap, As))
        depth.append(Z[-1])

        # Sediment
        if sed != None:
            Z =  sed[it].z + ssp[it].z[-1]  # Depth
            Cp = sed[it].cp     # Speed of P-wave
            Cs = sed[it].cs     # Speed of S-wave
            Rho = sed[it].rho   # Density of the media
            Ap = sed[it].a_p    # Attenuation of P-wave
            As = sed[it].a_s    # Attenuation of S-wave
            ssp_raw.append(SSPraw(Z, Cp, Cs, Rho, Ap, As))
            depth.append(Z[-1])

        # Make sure that the base depth is consistent in all profiles
        if depth[-1] != Zmax:
            Z = [depth[-1], Zmax]           # Depth
            Cp = [Cp[-1], base[it].cp]      # Speed of P-wave
            Cs = [Cs[-1], base[it].cs]      # Speed of S-wave
            Rho = [Rho[-1], base[it].rho]   # Density of the media
            Ap = [Ap[-1], base[it].a_p]     # Attenuation of P-wave
            As = [As[-1], base[it].a_s]     # Attenuation of S-wave
            ssp_raw.append(SSPraw(Z, Cp, Cs, Rho, Ap, As))
            depth.append(Z[-1])
            NMedia += 1

        Opt_top = 'SVW .'
        N = np.zeros(NMedia, np.int8)
        Sigma = np.zeros(NMedia+1)
        sspK = SSP(ssp_raw, depth, NMedia, Opt_top, N, Sigma)

        # Bottom option
        hs = HS(base[it].cp, base[it].cs, base[it].rho, base[it].a_p, base[it].a_s)
        Opt_bot = 'A'
        bottom = BotBndry(Opt_bot, hs)
        top = TopBndry(Opt_top)
        bdy = Bndry(top, bottom)

        # Write *.env file
        if it == 0:
            write_env(envFilename, model, title, freq, sspK, bdy, pos, [], cInt, Rmax)
        else:
            write_env(envFilename, model, title, freq, sspK, bdy, pos, [], cInt, Rmax, None, 'a')

    # Call krakenc.exe
    system(AtBinPath + "/krakenc " + filename)
    # Write *.flp file
    Opt_flp = 'RC'
    pos.r.depth = rd
    write_fieldflp(filename, Opt_flp, pos, NProf=len(ssp), rProf=bathm.r)
    # Call field.exe
    system(AtBinPath + "/field " + filename)
    # Read sound field
    [x, x, x, x, Pos1, pressure] = read_shd(filename+'.shd')
    TL = -10 * np.log10(np.square(np.abs(pressure)))  # Transmission Loss

    return Pos1, TL

# if __name__=="__main__":
    ''' ********************************************************************************* '''
    ''' ***************************** Input Parameters ********************************** '''
    ''' ********************************************************************************* '''
    # freq = 25                       # Frequency
    # sd = np.array([40.0])           # Source depth
    # rd = np.linspace(0, 500, 501)   # 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=50, azi=0)

    # [Pos, TL] = call_krakenc(freq, sd, rd, bathm, ssp, sed, base)

    # # Plot
    # 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()