import numpy as np
import math
from scipy.fftpack import fft
from matplotlib import pyplot as plt
from scipy.interpolate import interp1d
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 Base:
    def __init__(self, cp, cs, rho, a_p, a_s):
        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

def GetSeaDepth(data, N, coordS, R, azi):
    r = np.linspace(0, R, N)
    lon = coordS.lon + r * math.sin(azi) / (111 * math.cos(coordS.lat / 180 * math.pi))
    lat = coordS.lat + r * math.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 GetEnv(data, freq, coordS, coordR=None, R=None, azi=None):
    if coordR == None:
        azi = azi/180*math.pi
        lon = coordS.lon + R * math.sin(azi) / (111 * math.cos(coordS.lat/180*math.pi))
        lat = coordS.lat + R * math.cos(azi) / 111
        coordR = Coord(lon, lat)

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

    # Get Bathymetry
    N = math.ceil(1000 * R / (1500 / freq))
    #print("N = {}".format(N))
    if N == 1:
        N1 = 2
    else:
        r, z = GetSeaDepth(data, N, coordS, R, azi)
        Z = np.square(abs(fft(z-np.mean(z))))
        s1 = np.sum(Z[0:math.ceil(N / 2)-1])
        s2 = 0
        for it in range(math.ceil(N / 2)):
            s2 = s2 + Z[it]
            if s2 / s1 > 0.98:
                break
        N1 = max(2 * it, 1)
        # plt.plot(r, z)

    #print("N1 = {}".format(N1))
    r, z = GetSeaDepth(data, N1, coordS, R, azi)
    bathym = BathyM(r, z)
    #plt.plot(r, z)
    #plt.gca().invert_yaxis()
    #plt.show()

    # Get Sound Speed Profile

    # N2 = math.ceil(1000 * R / (30 * 1500 / freq))
    N2 = N1
#    N2=math.ceil(N2/2)
    r = np.linspace(0, R, N2)
    lon = coordS.lon + r * math.sin(azi) / (111 * math.cos(coordS.lat / 180 * math.pi))
    lat = coordS.lat + r * math.cos(azi) / 111
    ssp = []

    z_ssp = np.arange(0, 30, 1)
    z_ssp = np.append(z_ssp, np.arange(30, 100, 10))
    z_ssp = np.append(z_ssp, np.arange(100, 2000, 50))
    z_ssp = np.append(z_ssp, np.arange(2000, 10000, 100))
    z_ssp = np.append(z_ssp, np.ceil(bathym.d))
    z_ssp = np.sort(z_ssp)
    z_ssp = np.unique(z_ssp)


    for it in range(N2):
        data.set_coord(lat[it], lon[it])
        velFullPf = data.get_velFullProfile()
        itp1 = interp1d(velFullPf.x, velFullPf.y)
        z_idx = (z_ssp <= velFullPf.x[-1]) & (z_ssp >= 0)
        ssp_z = z_ssp[z_idx]
        ssp_c = itp1(ssp_z)
        ssp.append(SSP_(r[it], ssp_z, ssp_c))
        #ssp.append(SSP_(r[it], velFullPf.x, velFullPf.y))
        #plt.plot(velFullPf.y, velFullPf.x)

    # plt.gca().invert_yaxis()
    # plt.show()

    # Sediment
    sed = []
    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])
    for it in range(N2):
        sed.append(Sed(z, cp, cs, rho, a_p, a_s))

    # Base
    base = []
    Cb_p = 1550     # Sound speed of P-wave
    Cb_s = 0        # Sound speed of S-wave
    Rhob = 1.4      # Density of sea bottom
    Ab_p = 0.2      # Attenuation of P-wave
    Ab_s = 0        # Attenuation of S-wave
    for it in range(N2):
        base.append(Base(Cb_p, Cb_s, Rhob, Ab_p, Ab_s))

    return bathym, ssp, sed, base

# if __name__ == "__main__":
#     freq = 25
#     # 定义坐标
#     lat = 18.00
#     lon = 114.00
#     # 定义时间段
#     timeIndex = db.timeDict['Annual']
#     # 创建类
#     data = db.DataBase()
#     # 设置时间段
#     data.set_timePeriod(timeIndex)

#     #[bathym, ssp, sed, base] = GetEnv(data, freq, Coord(lon, lat), R=111, azi=-90)
#     [bathym, ssp, sed, base] = GetEnv(data, freq, Coord(lon, lat), Coord(113,19))