# -*- coding: utf-8 -*-

from os import system
import numpy as np
from pyat.AcousticPara import GetEnv, Coord
from scipy.signal  import savgol_filter, medfilt
from pyat.DataBase import DataBase
from pyat.readwrite import write_env, write_bathy, read_shd, get_rays
from pyat.env import Pos, Source, Dom, cInt, SSPraw, SSP, HS, BotBndry, TopBndry, Bndry, Box, Beam
import matplotlib.pyplot as plt
from pyat.config import AtBinPath

file_cz_data = open('cz_data.txt', 'w')

class CzRangeWidth:
    def __init__(self, minRange, maxRange, width):
        self.minRange = minRange
        self.maxRange = maxRange
        self.width = width

class CzDetail:
    def __init__(self, descript, directCz, reflectCz):
        self.descript = descript
        self.directCz = directCz
        self.reflectCz = reflectCz

class CzData:
    def __init__(self, azi, ranges, depths, TL, bathym, Rays, CzDetail):
        self.azi = azi
        self.ranges = ranges
        self.depths = depths
        self.TL = TL
        self.bathym = bathym
        self.Rays = Rays
        self.CzDetail = CzDetail

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

def call_Bellhop(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(0, Rmax, round(Rmax / dr))  # Receiving range (unit: km)
    Zmax, ssp_idx = calZmax(ssp, sed)
    """RD 已修改 8.26"""
    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
    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 已修改 8.26"""
    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[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 = 'I'  # incoherently sum beams, see AT docs for more info
    nbeams = 0
    alpha = np.linspace(-90, 90, 2)  # min and max launch angle -20 degrees to 20 degrees
    box = Box(Zmax, 100)  # 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
    """删除ssp 已修改 8.26"""
    # Write *.env file
    write_env(filename + '.env', 'BELLHOP', 'Pekeris profile', freq, sspK, bdy, pos, beam, cInt, Rmax)
    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_ray'
    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)
    Zmax, ssp_idx = calZmax(ssp, sed)
    """RD 已修改 8.26"""
    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
    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 已修改 8.26"""
    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[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 = 'R'  # incoherently sum beams, see AT docs for more info
    nbeams = 301
    alpha = np.linspace(-15, 15, 2)  # min and max launch angle -20 degrees to 20 degrees
    box = Box(Zmax, 100)  # 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
    """删除ssp 已修改 8.26"""
    # Write *.env file
    write_env(filename + '.env', 'BELLHOP', 'Pekeris profile', freq, sspK, 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 call_hjq(r, TL):
    TL_R = 20 * np.log10(r)
    out=[]
    TL[np.isnan(TL)] = 0
    TL_R[np.isnan(TL_R)] = 0
    dTL = (TL_R - TL) * (TL <= TL_R)
    dTL[np.isnan(dTL)] = 0
    test_fliter_s = savgol_filter(dTL, 103, 3)
    test_fliter = medfilt(test_fliter_s,13)

    signal=1
    for r_3 in range(len(r)):
        if r[r_3]>30000:
            break
    
    czrangewidth_list = []
    for i in range(r_3,len(test_fliter)):
        czrangewidth = CzRangeWidth(0,0,0)
        if(test_fliter[i]>=3):
            out.append(i)
            if i+1<len(test_fliter):
                if test_fliter[i+1]<3:
                    m = r[out[0]]
                    n = r[out[len(out) - 1]]
                    l=n-m
                    # file_cz_data.write('\n'"会聚区距离：%.1fm"%m +"-%.1fm" %n)
                    # file_cz_data.write('\n'"会聚区宽度：%.1fm\n"%l )
                    czrangewidth_list.append(CzRangeWidth(m,n,l))
                    out.clear()
                    signal=0
            else:
                m = r[out[0]]
                n = r[out[len(out) - 1]]
                l = n - m
                # file_cz_data.write('\n'"会聚区距离：%.1fm" % m + "-%.1fm" % n)
                # file_cz_data.write('\n'"会聚区宽度：%.1fm\n" % l )
                czrangewidth_list.append(CzRangeWidth(m,n,l))
                out.clear()
                signal = 0
    # if signal==1:
        # file_cz_data.write('\n'" 没有出现会聚区\n")
    #滤波器效果图：
    #plt.plot(r, TL)
    #plt.plot(r, TL_R)
    # plt.plot(r,dTL)
    # plt.plot(r, test_fliter)
    # plt.gca().invert_yaxis()
    # plt.show()
    
    czlen=len(czrangewidth_list)-1
    if czlen>=1:
        for i in range(1,czlen,1):
            if czrangewidth_list[i].minRange-czrangewidth_list[i-1].maxRange<=2000:
                czrangewidth_list[i].minRange = czrangewidth_list[i-1].minRange
                czrangewidth_list[i].width=czrangewidth_list[i].maxRange-czrangewidth_list[i].minRange
                czrangewidth_list[i - 1]=0

    while czrangewidth_list.count(0) > 0:
        for j in czrangewidth_list:
            if j == 0:
                czrangewidth_list.remove(j)
    return czrangewidth_list


def call_fangwei(freq, sd, rd, bathm, ssp, sed, base, isPlot=False):
    czdetail = CzDetail("", [], [])
    last=len(ssp[0].c)-1
    # if ssp[0].c[0] > ssp[0].c[last]:
        # file_cz_data.write('\n'"直达会聚区：无")
        # file_cz_data.write('\n'"反射会聚区：")
    # else:
        # file_cz_data.write('\n'"直达会聚区：")
    sed = None
    rays_total = call_Bellhop_Rays(freq, sd, rd, bathm, ssp, sed, base)
    Rays = find_cvgcRays(rays_total)
    [Pos, TL] = call_Bellhop(freq, sd, rd, bathm, ssp, sed, base)
    [x, x, x, x, Pos, pressure] = read_shd('cz.shd')
    TL = -10 * np.log10(np.square(np.abs(pressure)))  # Transmission Loss
    # 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()
    czrangewidth_list = call_hjq(Pos.r.range,TL[0, 0, 1,:])
    if ssp[0].c[0] > ssp[0].c[last]:
        czdetail.directCz = []
        czdetail.reflectCz = czrangewidth_list
    else:
        czdetail.directCz = czrangewidth_list
        czdetail.reflectCz = []
    return Pos.r.range, Pos.r.depth, TL[0, 0, :, :], bathm, Rays, czdetail

def call_cz(database, freq, lat, lon, azi, sourceDepth=10, receiveDepth=10, isPlot=False):
    sd = np.array([sourceDepth])  # Source depth
    rd = np.linspace(0, receiveDepth, 351)  # Receiving depth

    results = []
    ranges, depths, TL, bathym, Rays = [], [], [], [], []

    ''' 计算单个方向的会聚区 '''
    if azi != None:
        [bathm, ssp, sed, base] = GetEnv(database, 25, Coord(lon, lat), R=100, azi=azi)
        if np.mean(bathm.d) < 2000:
            return results
        else:
            ranges, depths, TL, bathym, Rays, czdetail = call_fangwei(freq, sd, rd, bathm, ssp, sed, base, isPlot=isPlot)
            czdetail.descript = "{:.1f}°方向".format(azi)
            results.append(CzData(0, ranges, depths, TL, bathym, Rays, czdetail))
            return results

    ''' 计算8个方向的会聚区 '''
    [bathm, ssp, sed, base] = GetEnv(database, 25, Coord(lon, lat), R=100, azi=0)
    if bathm.d[0] < 2000:
        # file_cz_data.write("各个方位都不会出现会聚区")
        return results
    else:
        descripts = ["正北方向", "东北方向", "正东方向", "东南方向", "正南方向", "西南方向", "正西方向", "西北方向"]
        theta = np.arange(0, 359, 45)
        for i in range(8):
            [bathm, ssp, sed, base] = GetEnv(database, 25, Coord(lon, lat), R=100, azi=theta[i])
            if np.mean(bathm.d) < 2000:
                # 不会出现会聚区
                continue
            ranges, depths, TL, bathym, Rays, czdetail = call_fangwei(freq, sd, rd, bathm, ssp, sed, base, isPlot=isPlot)
            czdetail.descript = descripts[i]
            results.append(CzData(theta[i], ranges, depths, TL, bathym, Rays, czdetail))

        return results