from readwrite import write_env, read_shd, get_rays
from env import Pos, Source, Dom, cInt, SSPraw, SSP, HS, BotBndry, TopBndry, Bndry, Box, Beam
from os import system
import numpy as np
from config import AtBinPath


def write_ssp(sspfile, ssp, bathm, NZmax):
    if sspfile[-3:] != 'ssp':
        sspfile += '.ssp'
    with open(sspfile, 'w') as f:
        f.write(str(len(ssp)) + '\r\n')
        for val in bathm.r:
            f.write(str(val) + '\t')
        f.write('\r\n')
        for row in range(NZmax):
            for col in range(len(ssp)):
                if len(ssp[col].z) <= row:
                    f.write('{:4.6f}'.format(1800) + '\t')
                else:
                    f.write('{:6f}'.format(ssp[col].c[row]) + '\t')
            f.write('\r\n')
        return


def write_bathy(btyfile, bathm):
    if btyfile[-3:] != 'bty':
        btyfile += '.bty'
    with open(btyfile, 'w') as f:
        f.write('CS\r\n')
        f.write(str(len(bathm.d)) + '\r\n')
        for row in range(len(bathm.r)):
            f.write('  ' + str(bathm.r[row]) + ' ' + str(bathm.d[row]) + '\r\n')
    return


def calZmax(ssp):
    d = np.array([], dtype=int)
    for it in range(len(ssp)):
        d = np.append(d, len(ssp[it].z))
    NZmax =max(d)
    ssp_idx = np.argmax(d)
    Zmax = ssp[ssp_idx].z[-1]
    return NZmax, Zmax, ssp_idx

def call_Bellhop(freq, sd, rd, bathm, ssp, sed, base):
    # Source and receiving position
    filename = 'envB'
    Rmax = bathm.r[-1]  # unit: km
    dr = (1500 / freq) / 1000  # wavelength / 10, unit: km
    ran = np.linspace(0, Rmax, round(Rmax / dr))  # Receiving range (unit: km)
    NZmax, Zmax, ssp_idx = calZmax(ssp)
    RD = np.linspace(0, Zmax, round(Zmax / 10) + 1)
    pos = Pos(Source(sd), Dom(ran, rd))
    # Range of phase velocity
    cInt.Low = 1400
    cInt.High = 15000

    # The number of media
    NMedia = 1
    ssp_raw = []
    depth = [0]
    # Sound speed profile
    Z = ssp[ssp_idx].z  # Depth
    Cp = ssp[ssp_idx].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 is ignored
    '''
    if sed != None:
        Z =  sed[0].z + ssp[0].z[-1]  # Depth
        Cp = sed[0].cp     # Speed of P-wave
        Cs = sed[0].cs     # Speed of S-wave
        Rho = sed[0].rho   # Density of the media
        Ap = sed[0].a_p    # Attenuation of P-wave
        As = sed[0].a_s    # Attenuation of S-wave
        ssp_raw.append(SSPraw(Z, Cp, Cs, Rho, Ap, As))
        depth.append(Z[-1])
    '''

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

    #  Bottom option
    hs = HS(base[0].cp, base[0].cs, base[0].rho, base[0].a_p, base[0].a_s)
    Opt_bot = 'A~'
    bottom = BotBndry(Opt_bot, hs)
    top = TopBndry(Opt_top)
    bdy = Bndry(top, bottom)
    # Beam params
    run_type = 'C'  # incoherently sum beams, see AT docs for more info
    nbeams = 3000
    alpha = np.linspace(-45, 45, 2)  # min and max launch angle -20 degrees to 20 degrees
    box = Box(Zmax, max(bathm.r))  # bound the region you let the beams go, depth in meters and range in km
    deltas = 0  # length step of ray trace, 0 means automatically choose
    beam = Beam(RunType=run_type, Nbeams=nbeams, alpha=alpha, box=box, deltas=deltas)  # package

    # Write *.env file
    write_env(filename + '.env', 'BELLHOP', 'Pekeris profile', freq, sspB, bdy, pos, beam, cInt, Rmax)
    write_ssp(filename, ssp, bathm, NZmax)
    write_bathy(filename, bathm)

    system(AtBinPath + "/bellhop " + 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

def call_Bellhop_Rays(freq, sd, rd, bathm, ssp, sed, base):
    # Source and receiving position
    filename = 'cz'
    Rmax = bathm.r[-1]  # unit: km
    dr = (1500 / freq) / 1000  # wavelength / 10, unit: km
    ran = np.linspace(dr, Rmax, round(Rmax / dr))  # Receiving range (unit: km)
    NZmax, Zmax, ssp_idx = calZmax(ssp)
    RD = np.linspace(0, Zmax, round(Zmax / 10) + 1)
    pos = Pos(Source(sd), Dom(ran, rd))
    # Range of phase velocity
    cInt.Low = 1400
    cInt.High = 15000

    # The number of media
    NMedia = 1
    ssp_raw = []
    depth = [0]
    # Sound speed profile
    Z = ssp[ssp_idx].z  # Depth
    Cp = ssp[ssp_idx].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 is ignored
    '''
    if sed != None:
        Z =  sed[0].z + ssp[0].z[-1]  # Depth
        Cp = sed[0].cp     # Speed of P-wave
        Cs = sed[0].cs     # Speed of S-wave
        Rho = sed[0].rho   # Density of the media
        Ap = sed[0].a_p    # Attenuation of P-wave
        As = sed[0].a_s    # Attenuation of S-wave
        ssp_raw.append(SSPraw(Z, Cp, Cs, Rho, Ap, As))
        depth.append(Z[-1])
    '''

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

    #  Bottom option
    hs = HS(base[0].cp, base[0].cs, base[0].rho, base[0].a_p, base[0].a_s)
    Opt_bot = 'A~'
    bottom = BotBndry(Opt_bot, hs)
    top = TopBndry(Opt_top)
    bdy = Bndry(top, bottom)
    # Beam params
    nbeams = 301
    alpha = np.linspace(-10, 10, 2)  # min and max launch angle -20 degrees to 20 degrees
    box = Box(Zmax, max(bathm.r))  # bound the region you let the beams go, depth in meters and range in km
    deltas = 0  # length step of ray trace, 0 means automatically choose
    run_type = 'R'
    beam = Beam(RunType=run_type, Nbeams=nbeams, alpha=alpha, box=box, deltas=deltas)  # package

    # Write *.env file
    write_env(filename + '.env', 'BELLHOP', 'Pekeris profile', freq, sspB, bdy, pos, beam, cInt, Rmax)
    write_bathy(filename, bathm)

    system(AtBinPath + "/bellhop " + filename)
    # Read sound field
    return get_rays(filename + ".ray")

def find_cvgcRays(rays_total):
    Rays = []
    num_bnc_min = np.Inf
    for ray in rays_total[0]:
        if ray.num_top_bnc + ray.num_bot_bnc < num_bnc_min and max(ray.xy[1,:]) > 100:
            num_bnc_min = ray.num_top_bnc + ray.num_bot_bnc

    for ray in rays_total[0]:
        if ray.num_top_bnc + ray.num_bot_bnc > num_bnc_min + 1:
            continue
        elif max(ray.xy[1,:]) < 100:
            continue
        else:
            Rays.append(ray)
    return Rays

# if __name__=="__main__":
    ''' ********************************************** '''
    ''' ******************** Test ******************** '''
    # freq = 25                        # Frequency
    # sd = np.array([10.0])            # Source depth
    # rd = np.linspace(0, 4000, 401)   # Receiving depth
    # # 定义坐标
    # lat = 19.00
    # lon = 116.00
    # # 定义时间段
    # timeIndex = db.timeDict['Annual']
    # # 创建类
    # data = db.DataBase()
    # # 设置时间段
    # data.set_timePeriod(timeIndex)

    # [bathm, ssp, sed, base] = GetEnv(data, freq, Coord(lon, lat), R=100, azi=90)
    # [Pos, TL] = call_Bellhop(freq, sd, rd, bathm, ssp, sed, base)

    # #[x, x, x, x, Pos, pressure] = read_shd('envB.shd')
    # #TL = -10 * np.log10(np.square(np.abs(pressure)))  # Transmission Loss
    # # 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.plot(bathm.r*1000, bathm.d, color='k')
    # plt.show()