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


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 = DataBase.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