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
from multiprocessing import Pool
from scipy.stats import norm
import math


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_p(filename):
    system(AtBinPath + "/bellhop " + filename)

def call_Bellhop(freq, sd, bathm, ssp, sed, base):
    # Source and receiving position
    filename = 'envB'
    Rmax = bathm.r[-1]  # unit: km
    dr = 0.1  # 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

    NAlphaRange = 12
    Alpha = alphadiv(NAlphaRange,Rmax) # min and max launch angle -20 degrees to 20 degrees
    totalBeams = beamsnumber(freq,Rmax,max(bathm.r))
    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

    # Write *.env file
    ialpha = 0
    Filenames = []
    for iAlphaRange in range(len(Alpha)-1):
        alpha = np.array([Alpha[iAlphaRange], Alpha[iAlphaRange + 1]])
        nbeams = int(totalBeams * (alpha[1]-alpha[0])/180)

        beam = Beam(RunType=run_type, Nbeams=nbeams, alpha=alpha, box=box, deltas=deltas)  # package
        filenameI = filename+ str(iAlphaRange)
        write_env(filenameI + '.env', 'BELLHOP', 'Pekeris profile', freq, sspB, bdy, pos, beam, cInt, Rmax)
        write_ssp(filenameI, ssp, bathm, NZmax)
        write_bathy(filenameI, bathm)
        Filenames.append(filenameI)

    pool = Pool(NAlphaRange)
    pool.map(call_Bellhop_p, Filenames)

    # Read sound field
    [x, x, x, x, Pos1, pressure] = read_shd(filename + '0.shd')
    for iAlphaRange in range(NAlphaRange-1):
        [x, x, x, x, Pos1, pressure1] = read_shd(filename + str(iAlphaRange) + '.shd')
        pressure = pressure + pressure1
    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

def alphadiv(NalphaRange,Rmax):
    angle = [-90,90]
    if Rmax >100 and Rmax<=1000:
        sigma = 20
    elif Rmax >=1000:
        sigma = 15
    else :
        sigma = 25
    for i in range(1,NalphaRange):
        angle.append(norm.ppf(i/NalphaRange,loc=0, scale=sigma))

    angle.sort()
    return angle



def beamsnumber(freq,Rmax,depth):
    if Rmax >700 :
        nbeams = max(int(0.01 * Rmax*1000 * freq /1500),300)
    elif Rmax >300 and Rmax<=700:
        nbeams = max(int(0.05 * Rmax*1000 * freq / 1500), 300)
    elif Rmax >100 and Rmax<=300:
        nbeams = max(int(0.1 * Rmax*1000 * freq / 1500), 300)
    else:
        nbeams = max(int(0.3 * Rmax*1000 * freq / 1500), 300)

    d_theta_recommended =math.atan(depth/(10.0 * Rmax*1000))
    nbeams = min(int(math.pi/d_theta_recommended),nbeams,3000)

    return nbeams

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