""" Wind Generated Noise Prediction """
import numpy as np
from scipy.interpolate import interp1d
from pyat import DataBase as db
import matplotlib.pyplot as plt

class MdlPara():
    def __init__(self, title, lat, lon, seaDepth, month, freq, vc, a, b, vs):
        self.title = title
        self.lat = lat
        self.lon = lon
        self.seaDepth = seaDepth
        self.month = month
        self.freq = freq
        self.vc = vc
        self.a = a
        self.b = b
        self.vs = vs

def loadMdlPara(filename):
    f = open(filename, mode='r')

    mdlPara = []
    eof = False
    while (1):
        # Title
        while (1):
            line = f.readline()
            if line == '':
                eof = True
                break
            title = line.strip()
            if title != '':
                break
        if eof:
            break

        # Latitude and Longitude
        line = f.readline()
        line = line.strip().split()
        lat = np.double(line[0])
        lon = np.double(line[1])

        # Sea Depth
        line = f.readline()
        line = line.strip().split()
        seaDepth = np.double(line[0])

        # Month
        line = f.readline()
        line = line.strip().split()
        month = np.int8(line[0])

        # Model Parameters
        freq = []
        vc = []
        a = []
        b = []
        vs = []
        while (1):
            line = f.readline()
            line = line.strip().split()
            if line[0] == '0':
                break
            freq.append(np.double(line[0]))
            vc.append(np.double(line[1]))
            a.append(np.double(line[2]))
            b.append(np.double(line[3]))
            vs.append(np.double(line[4]))

        mdlPara.append(MdlPara(title, lat, lon, seaDepth, month, freq, vc, a, b, vs))

    return mdlPara


def WGNMdl(mdlPara, windspeed):
    f = np.array(mdlPara.freq)
    spl = np.zeros(len(f))
    for it in range(len(f)):
        if windspeed < mdlPara.vc[it]:
            spl[it] = mdlPara.a[it] + mdlPara.b[it] * 10 * np.log10(mdlPara.vc[it])
        elif (mdlPara.vs[it] is not np.NaN) & (windspeed > mdlPara.vs[it]):
            spl[it] = mdlPara.a[it] + mdlPara.b[it] * 10 * np.log10(mdlPara.vs[it])
        else:
            spl[it] = mdlPara.a[it] + mdlPara.b[it] * 10 * np.log10(windspeed)

    return f, spl


def WGNPd(dataFile, seaDepth, windspeed, month, freq):
    if min(freq) < 50:
        raise UserWarning("The minimum frequency must be larger than 50 Hz")

    mdlPara = loadMdlPara(dataFile)
    
    z = seaDepth

    Q = []
    for it in range(len(mdlPara)):
        q1 = abs(mdlPara[it].seaDepth - z) / (mdlPara[it].seaDepth + z)
        q2 = min(abs(month-mdlPara[it].month), 12-abs(month-mdlPara[it].month)) / 6
        alpha = 0.7
        Q.append(alpha * q1 + (1-alpha) * q2)
    idx = np.argmin(Q)
    # print([z, month, idx])

    f0, spl0 = WGNMdl(mdlPara[idx], windspeed)
    if max(freq) > max(f0):
        spl0 = np.append(spl0, spl0[-1] - 5 * np.log2(max(freq)/max(f0)))
        f0 = np.append(f0, max(freq))
    if min(f0) > 200:
        spl0 = np.insert(spl0, 0, spl0[0] + 5 * np.log2(min(f0)/200))
        f0 = np.insert(f0, 0, 200)
    if min(f0) > 50:
        spl0 = np.insert(spl0, 0, spl0[0] + 9 * np.log2(min(f0)/50))
        f0 = np.insert(f0, 0, 50)

    itp1 = interp1d(np.log10(f0), spl0)
    spl = itp1(np.log10(freq))

    return freq, spl


# if __name__=="__main__":
#     lat = 19
#     lon = 116
#     windspeed = 5      # unit: m/s
#     month = 3
#     freq = 50 * np.power(2, np.arange(25)/3)

#     # 定义时间段
#     timeIndex = db.timeDict['Annual']
#     # 创建类
#     data = db.DataBase()
#     # 设置时间段
#     data.set_timePeriod(timeIndex)

#     f, spl = WGNPd(data, lat, lon, windspeed, month, freq)

#     plt.semilogx(f, spl)
#     plt.xlabel('Frequency / Hz')
#     plt.ylabel('Noise Level / dB')
#     plt.show()