# -*- coding: utf-8 -*-
from multiprocessing import Process, Pool
from functools import partial
from os import system
import numpy as np
from AcousticPara import GetEnv, Coord
from scipy.signal import savgol_filter, medfilt
from DataBase import DataBase
from readwrite import write_env, write_bathy, read_shd, get_rays
from env import Pos, Source, Dom, cInt, SSPraw, SSP, HS, BotBndry, TopBndry, Bndry, Box, Beam
import matplotlib.pyplot as plt
from 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, bathm, ssp, sed, base,azi):
    # Source and receiving position
    filename = 'cz'+str(azi)
    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)
    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 = 1000
    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, bathm, ssp, sed, base,azi):
    # Source and receiving position
    filename = 'cz_ray'+str(azi)
    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)
    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, bathm, ssp, sed, base, azi, 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, bathm, ssp, sed, base, azi)
    Rays = find_cvgcRays(rays_total)
    [Pos, TL] = call_Bellhop(freq, sd, bathm, ssp, sed, base, azi)
    [x, x, x, x, Pos, pressure] = read_shd('cz'+str(azi)+'.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 child(database, freq, lat, lon, sd, isPlot, i):
    descripts = ["正北方向", "东北方向", "正东方向", "东南方向", "正南方向", "西南方向", "正西方向", "西北方向"]
    theta = np.arange(0, 359, 45)
    [bathm, ssp, sed, base] = GetEnv(database, 25, Coord(lon, lat), R=100, azi=theta[i])
    if np.mean(bathm.d) < 2000:
        # 不会出现会聚区
        ddd=[]
        return None
    else:
        ranges, depths, TL, bathym, Rays, czdetail = call_fangwei(freq, sd, bathm, ssp, sed, base, azi=theta[i],isPlot=isPlot)
        czdetail.descript = descripts[i]
        return CzData(theta[i], ranges, depths, TL, bathym, Rays, czdetail)

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

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

    # def child_0(interval):
    #     czdata=child(interval)
    #     return czdata
    ''' 计算单个方向的会聚区 '''
    if azi != None:
        [bathm, ssp, sed, base] = GetEnv(database, 25, Coord(lon, lat), R=100, azi=azi)
        if np.mean(bathm.d) < 2000:
            return None
        else:
            ranges, depths, TL, bathym, Rays, czdetail = call_fangwei(freq, sd, bathm, ssp, sed, base, azi,
                                                                      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:
        pool = Pool(processes=8)
        applyResults = []
        for pi in range(8):
            applyResults.append(pool.apply_async(partial(child, database, freq, lat, lon, sd, isPlot), (pi,)))
            # pi=Process(target=child_0,args=(pi,))
            # pi.start()
        pool.close()
        pool.join()
        for applyResult in applyResults:
            results.append(applyResult.get())
        return results
# timeIndex = '20190101'
# # 设置时间段
#
# STPath = '/home/gan/下载/code'
# ETOPO1Path = '/home/gan/下载/code/20220920/ETOPO1'
# database = DataBase(ETOPO1Path, STPath)
# database.set_timePeriod(timeIndex)
# res = call_cz(database, 300, 13, 115, 0, 60, True)
# a = 1