import os
from multiprocessing import Pool
from functools import partial
from time import time
from shutil import rmtree
import numpy as np
from scipy.interpolate import interp1d
from copy import deepcopy
from pyat import ramRW_new as ramRW
from pyat.DataBase import DataBase,soundSpeed


class BathyM:
    def __init__(self, ran=None, depth=None):
        self.r = ran
        self.d = depth


class SSP_:
    def __init__(self, r, z, c):
        self.r = r
        self.z = z
        self.c = c


class Sed:
    def __init__(self, z, cp, cs, rho, a_p, a_s):
        self.z = z
        self.cp = cp
        self.cs = cs
        self.rho = rho
        self.a_p = a_p
        self.a_s = a_s


class Coord:
    def __init__(self, lon, lat):
        self.lon = lon
        self.lat = lat


class RANDI_Para:
    def __init__(self, baseDir, freq, lon, lat, Range, dRange, dTheta, v_ship, l_ship, theta=0):
        self.baseDir = baseDir
        self.freq = freq
        self.lon = lon
        self.lat = lat
        self.Range = Range
        self.dRange = dRange
        self.dTheta = dTheta
        self.v_ship = v_ship
        self.l_ship = l_ship
        self.theta = theta

    def refreshTheta(self, theta):
        self.theta = theta
        return self


def RANDI_GetSeaDepth(data, N, coordS, R, azi):
    r = np.linspace(0, R, N)
    lon = coordS.lon + r * np.sin(azi) / (111 * np.cos(coordS.lat / 180 * np.pi))
    lat = coordS.lat + r * np.cos(azi) / 111

    z = np.zeros(N)
    for it in range(N):
        data.set_coord(lat[it], lon[it])
        z[it] = data.get_bottomDepth()
    return r, z


def RANDI_GetMeanProfile(data, N, coordS, R, azi):
    r = np.linspace(0, R, N)
    lon = coordS.lon + r * np.sin(azi) / (111 * np.cos(coordS.lat / 180 * np.pi))
    lat = coordS.lat + r * np.cos(azi) / 111
    sal = np.full([len(data.depthGrid), N], np.nan)
    tem = np.full([len(data.depthGrid), N], np.nan)
    for i in range(N):
        data.set_coord(lat[i], lon[i])
        salPf = data.get_salRawProfile()
        temPf = data.get_temRawProfile()
        if len(salPf.x) > 1:
            sal[0:len(salPf.x), i] = salPf.y
            tem[0:len(temPf.x), i] = temPf.y
    sal = np.ma.masked_array(sal, np.isnan(sal))
    tem = np.ma.masked_array(tem, np.isnan(tem))
    sal_m = np.mean(sal, 1)
    tem_m = np.mean(tem, 1)
    z = data.depthGrid[~sal_m.mask]
    sal_m = sal_m[~sal_m.mask]
    tem_m = tem_m[~tem_m.mask]
    return z, sal_m, tem_m


def RANDI_GetEnv(data, coordS, coordR=None, R=None, azi=None):
    if coordR == None:
        azi = azi / 180 * np.pi
        lon = coordS.lon + R * np.sin(azi) / (111 * np.cos(coordS.lat / 180 * np.pi))
        lat = coordS.lat + R * np.cos(azi) / 111
        coordR = Coord(lon, lat)

    if R == None:
        x = (coordR.lon - coordS.lon) * (111 * np.cos(coordS.lat / 180 * np.pi))
        y = (coordR.lat - coordS.lat) * 111
        R = np.sqrt(x * x + y * y)
        if y > 0:
            azi = np.arctan(x / y)
        elif y == 0:
            azi = np.sign(x) * np.pi / 2
        else:
            azi = np.arctan(x / y) + np.pi

    # Get Bathymetry
    N = int(np.ceil(R / 10) + 1)
    r, z = RANDI_GetSeaDepth(data, N, coordS, R, azi)
    bathym = BathyM(r, z)
    # plt.plot(r, z)
    # plt.gca().invert_yaxis()
    # plt.show()

    # Get Sound Speed Profile
    z_ssp, sal_m, tem_m = RANDI_GetMeanProfile(data, N, coordS, R, azi)
    if max(bathym.d) > z_ssp[-1]:
        z_ssp = np.append(z_ssp, max(bathym.d))
        sal_m = np.append(sal_m, sal_m[-1])
        tem_m = np.append(tem_m, tem_m[-1])
    c_ssp = soundSpeed(tem_m, sal_m, z_ssp)
    ssp = SSP_(0, z_ssp, c_ssp)
    # plt.plot(c_ssp, z_ssp)
    # plt.gca().invert_yaxis()
    # plt.show()

    # Sediment
    z = np.array([0, 50])
    cp = np.array([1550, 1590])
    cs = np.array([0, 0])
    rho = np.array([1.3, 1.5])
    a_p = np.array([0.1, 0.1])
    a_s = np.array([0, 0])
    sed = Sed(z, cp, cs, rho, a_p, a_s)

    return bathym, ssp, sed


def RANDI_writeRAMin(baseDir, data, freq, dTheta, lat, lon, Range, rd, ship_sd, SeaDepthLim):
    t1 = time()
    btyFile = open("bty", 'w')
    for theta in np.arange(0, 360, dTheta):
        [bathm, ssp, sed] = RANDI_GetEnv(data, Coord(lon, lat), R=Range, azi=theta)
        for dd in bathm.d:
            btyFile.write(str(dd) + ' ')
        btyFile.write('\n')
        idx = np.where(bathm.d < SeaDepthLim)
        if len(idx[0]) > 0:
            bathm.r = bathm.r[0:idx[0][0]]
            bathm.d = bathm.d[0:idx[0][0]]

        title = "Waveguide"
        # Source and receiving position
        bathm.r = bathm.r * 1000  # km -> m
        rmax = max(bathm.r)  # Maximum receiving range (unit: m)
        zmax = round(max(bathm.d) + 5 * 1500 / freq)  # Maximum receiving depth (uint: m)
        zmplt = 0  # round(max(bathm.d)*1.05)    # Maximum depth of plot (uint: m)
        dz = 1500 / freq / 4  # Step in depth direction: wavelength / 15
        dr = 5 * dz  # Step in range direction
        pos = ramRW.PosRAM(rd, ship_sd, rmax, zmax, dr, dz, zmplt)

        # Sound speed profile for RAM
        Z = sed.z
        cw_p = ramRW.PROF(ssp.z, ssp.c)  # Water sound speed profile
        cb_p = ramRW.PROF(Z, sed.cp)  # Bottom sound speed profile
        rhob_p = ramRW.PROF(Z, sed.rho)  # Sea bottom density profile

        Ap = sed.a_p
        if Z[-1] < 5 * 1500 / freq:
            Z = np.append(Z, 3 * 1500 / freq)
            Ap = np.append(Ap, Ap[-1])
        Z = np.append(Z, Z[-1] + 2 * 1500 / freq)
        Ap = np.append(Ap, 10)
        # Ap = sed[it].a_p
        attn_p = ramRW.PROF(Z, Ap)  # Sea bottom attenuation profile

        ssp_ram = [ramRW.SSPRAM(0, cw_p, cb_p, rhob_p, attn_p)]

        # Write ram.in file
        if not os.path.exists(baseDir + "/TL/temp/theta_{:1d}".format(theta)):
            os.makedirs(baseDir + "/TL/temp/theta_{:1d}".format(theta))
        envfil = baseDir + "/TL/temp/theta_{:1d}/ram.in".format(theta)
        ramRW.write_env_RAM(title, envfil, freq, pos, bathm, ssp_ram)
    t2 = time()
    #print(t2 - t1)
    btyFile.close()


def RANDI_callRAM(baseDir, NThreads, freq, dTheta):
    '''
    t1 = time()
    for theta in np.arange(0,360,5):
        os.chdir(baseDir + "/TL/theta_{:1d}".format(theta))
        os.system(baseDir + "/ram.exe")

    os.chdir(baseDir)
    t2 = time()
    print(t2-t1)
    '''
    t1 = time()
    func = partial(callRAM, baseDir)
    pool = Pool(NThreads)
    pool.map(func, np.arange(0, 360, dTheta))
    t2 = time()
    #print(t2 - t1)

    if os.path.exists(baseDir + f"/TL/freq_{round(freq)}"):
        rmtree(baseDir + f"/TL/freq_{round(freq)}")
    os.rename(baseDir + "/TL/temp", baseDir + f"/TL/freq_{round(freq)}")


def callRAM(baseDir, theta):
    AtBinPath = os.path.dirname(os.path.abspath(__file__))+'/at/bin'
    print(AtBinPath)
    os.chdir(baseDir + "/TL/temp/theta_{:1d}".format(theta))
    os.system(AtBinPath + "/ram")
    os.chdir(baseDir)

def RANDI_getShipNum(lat, lon, Range, dRange, theta):
    NRange = int(Range / dRange)
    #shipNum = np.random.rand(NRange)
    shipNum = 0.5*np.ones(NRange)
    return shipNum

def calNoise(randiP):
    rr = np.arange(randiP.dRange / 2, randiP.Range, randiP.dRange)
    SL = RossShipSL(randiP.freq, randiP.v_ship, randiP.l_ship)
    S = np.power(10, SL / 10)
    shipNum = RANDI_getShipNum(randiP.lat, randiP.lon, randiP.Range, randiP.dRange, randiP.theta)
    filePath = randiP.baseDir + f"/TL/freq_{round(randiP.freq)}/theta_{round(randiP.theta)}"
    [r, tl] = ramRW.readTLline_RAM(filePath)
    r = r / 1000
    I2 = np.power(10, tl / 10)
    I_ = interp1d(r, I2)
    I = I_(rr[rr < r[-1]])
    noise = sum(shipNum[rr < r[-1]] * S / I)
    return noise

def RANDI_calNL(randiP, NThreads):
    '''
    t1 = time()
    theta = np.arange(0, 360, dTheta)
    noise1 = np.zeros(len(theta))
    for i in range(len(theta)):
        randiP.theta = theta[i]
        noise1[i] = calNoise(randiP)
    NL1 = 10 * np.log10(sum(noise1))
    t2 = time()
    print(t2-t1)
    '''

    t1 = time()
    RANDI_P = []
    for theta in np.arange(0, 360, randiP.dTheta):
        randiP.theta = theta
        RANDI_P.append(deepcopy(randiP))
    pool = Pool(NThreads)
    noise = pool.map(calNoise, RANDI_P)
    NL = 10 * np.log10(sum(noise))
    t2 = time()
    #print(t2-t1)
    return NL


def RossShipSL(freq, v, l):
    if freq < 500:
        Lso = -10 * np.log10(np.power(10, -1.06 * np.log10(freq) - 14.34) +
                                      np.power(10, 3.32 * np.log10(freq) - 21.425))
    else:
        Lso = 173.2 - 18 * np.log10(freq)

    if freq <= 28.4:
        df = 8.1
    elif freq < 191.6:
        df = 22.3-9.77*np.log10(freq)
    else:
        df = 0

    dl = np.power(l, 1.15) / 3643
    Ls = Lso + 60*np.log10(v/12) + 20*np.log10(l/300) + df*dl + 3
    return Ls


def call_RANDI(ETOPO1Path, STPath, baseDir, freq, rd, lon, lat, timeType, NThreads):
    # 定义时间段
    timeIndex = timeType
    # 创建类
    data = DataBase(ETOPO1Path, STPath)
    # 设置时间段
    data.set_timePeriod(timeIndex)

    Range = 500  # Range of sound field calculation (Unit: km)
    ship_sd = 6.0  # Ship source depth (Unit: m)
    dTheta = 5  # The angle of sector (Unit: °)
    dRange = 10  # The
    SeaDepthLim = 20  # (Unut: m)
    v_ship = 12  # Ship speed (Unit: kt)
    l_ship = 300  # Ship length (Unit: ft)

    randiP = RANDI_Para(baseDir, freq, lon, lat, Range,  dRange, dTheta, v_ship, l_ship)

    RANDI_writeRAMin(baseDir, data, freq, dTheta, lat, lon, Range, rd, ship_sd, SeaDepthLim)

    RANDI_callRAM(baseDir, NThreads, freq, dTheta)

    NL = RANDI_calNL(randiP, NThreads)

    return NL