import numpy as np
from scipy.interpolate import Akima1DInterpolator, interp1d
from netCDF4 import Dataset
from pyat.ProfileDataType import Profile

'''
************************************************************
该模块用于从数据库中, 获取剖面信息
************************************************************
'''

# 时间段标识



# def MunkFormula(z, v1, epsilon, z1):
#     eta = 2*(z-z1)/z1
#     return v1*(1+epsilon*(np.exp(-eta)+eta-1))

def soundSpeed(tem, sal, dep):
    '''
    声速计算公式
    '''
    # C = 1449.2 + 4.6*T - 0.055*T^2 + 0.00029*T^3 + (1.34 - 0.01*T)(S-35) + 0.017*D
    # 式中：T为海⽔测点的温度，单位为℃；S为海⽔测点的盐度，单位为‰；D为海⽔测点的深度，单位为m。
    vel = 1449.2 + 4.6*tem - 0.055*tem*tem + 0.00029 * \
        tem*tem*tem + (1.34 - 0.01*tem) * \
        (sal-35.0) + 0.017*dep
    return vel


class DataBase():
    """
    获取数据库中的剖面信息
    """
    def __init__(self,ETOPO1path,STpath):
        # 打开ETOPO1数据
        fileETOPO1 = ETOPO1path+'/ETOPO1_Bed_g_gmt4.nc'
        dataETOPO1 = Dataset(fileETOPO1)
        # 获取ETOPO1网格
        self.latGrid1 = dataETOPO1.variables['y'][:]
        self.lonGrid1 = dataETOPO1.variables['x'][:]
        # 获取海深数据
        self.botGrid = -dataETOPO1.variables['z'][:]
        self.STpath=STpath
        return

    def clearVar(self):
        """
        清空变量
        """
        self.botDep = None
        self.salRawPf = None
        self.temRawPf = None
        self.velRawPf = None
        self.velInterpPf = None
        self.velExtPf = None
        self.velFullPf = None
        self.isValid = True
        return

    def set_timePeriod(self, timeIndex):
        """
        设置时间段
        """
        self.timeIndex = timeIndex
        # 打开WOA18数据
        #fileWOA18sal = '.\WOA18\woa18_A5B7_s' + timeIndex + '_04.nc'
        #fileWOA18tem = '.\WOA18\woa18_A5B7_t' + timeIndex + '_04.nc'
        fileST=self.STpath+'/CA'+timeIndex+'.nc'
        dataST = Dataset(fileST)
        # 获取WOA18网格
        self.latGrid2 = dataST.variables['lat'][:]
        self.lonGrid2 = dataST.variables['lon'][:]
        # 获取深度数据
        self.depthGrid = dataST.variables['zlev'][:]*(-1)
        # 获取盐度数据
        self.salGrid = np.squeeze(dataST.variables['SALT'][:])
        # 获取温度数据
        self.temGrid = np.squeeze(
            dataST.variables['THETA'][:])
        self.clearVar()
        return

    def set_coord(self, lat, lon):
        """
        设置坐标
        """
        self.lat = lat
        self.lon = lon
        self.clearVar()
        return

    def interpBilinear(self, latId, lonId, dataGrid, dataTpye):
        """
        双线性插值
        """
        latId0 = int(np.ceil(latId))
        lonId0 = int(np.ceil(lonId))
        a = latId - latId0
        b = lonId - lonId0
        if (dataTpye == 1):
            data0 = dataGrid[latId0, lonId0]
            data1 = dataGrid[latId0+1, lonId0]
            data2 = dataGrid[latId0+1, lonId0+1]
            data3 = dataGrid[latId0, lonId0+1]
        else:
            data0 = dataGrid[:, latId0, lonId0]
            data1 = dataGrid[:, latId0+1, lonId0]
            data2 = dataGrid[:, latId0+1, lonId0+1]
            data3 = dataGrid[:, latId0, lonId0+1]
        data01 = (1-a)*data0+a*data1
        data32 = (1-a)*data3+a*data2
        data = (1-b)*data01+b*data32
        return data

    def normCoord(self, latGrid, lonGrid):
        """
        坐标单位化
        """
        lat0 = latGrid[0]
        lon0 = lonGrid[0]
        dlat = latGrid[1]-latGrid[0]
        dlon = lonGrid[1]-lonGrid[0]
        latId = (self.lat - lat0)/dlat
        lonId = (self.lon - lon0)/dlon
        return (latId, lonId)

    def get_bottomDepth(self):
        """
        获得海底深度
        """
        if (self.botDep == None):
            # 单位化经纬度
            (latId, lonId) = self.normCoord(self.latGrid1, self.lonGrid1)
            # 获取网格点海深
            self.botDep = self.interpBilinear(latId, lonId, self.botGrid, 1)
        return self.botDep

    def get_salRawProfile(self):
        """
        获得原始盐度剖面
        """
        if (self.salRawPf == None):
            botDep = self.get_bottomDepth()
            # 单位化经纬度
            (latId, lonId) = self.normCoord(self.latGrid2, self.lonGrid2)
            # 获取盐度剖面信息
            sal = self.interpBilinear(latId, lonId, self.salGrid, 2)
            # 获取有效数据段
            dep = self.depthGrid[sal.mask == False]
            sal = sal[sal.mask == False]
            dep = dep.flatten()
            sal = sal.flatten()
            # 若为空，则为无效数据
            if sal.size == 0:
                self.isValid = False
            # 剖面赋值
            salRawPf = Profile()
            salRawPf.set_data(dep, sal)
            salRawPf.set_coord(self.lat, self.lon)
            salRawPf.set_bottomDepth(botDep)
            self.salRawPf = salRawPf
        return self.salRawPf

    def get_temRawProfile(self):
        """
        获得原始温度剖面
        """
        if (self.temRawPf == None):
            botDep = self.get_bottomDepth()
            # 单位化经纬度
            (latId, lonId) = self.normCoord(self.latGrid2, self.lonGrid2)
            # 获取温度剖面信息
            tem = self.interpBilinear(latId, lonId, self.temGrid, 2)
            # 获取有效数据段
            dep = self.depthGrid[tem.mask == False]
            tem = tem[tem.mask == False]
            dep = dep.flatten()
            tem = tem.flatten()
            # 若为空，则为无效数据
            if tem.size == 0:
                self.isValid = False
            # 剖面赋值
            temRawPf = Profile()
            temRawPf.set_data(dep, tem)
            temRawPf.set_coord(self.lat, self.lon)
            temRawPf.set_bottomDepth(botDep)
            self.temRawPf = temRawPf
        return self.temRawPf

    def get_velRawProfile(self):
        """
        获得原始声速剖面
        """
        if (self.velRawPf == None):
            botDep = self.get_bottomDepth()
            # 获取温盐深数据
            dep = self.get_salRawProfile().x
            sal = self.get_salRawProfile().y
            tem = self.get_temRawProfile().y
            # 计算声速
            if self.isValid == True:
                vel = soundSpeed(tem, sal, dep)
            else:
                dep = np.array([0])
                vel = np.array([0])
            # 剖面赋值
            velRawPf = Profile()
            velRawPf.set_data(dep, vel)
            velRawPf.set_coord(self.lat, self.lon)
            velRawPf.set_bottomDepth(botDep)
            self.velRawPf = velRawPf
        return self.velRawPf

    def get_velInterpProfile(self):
        """
        获得插值后的声速剖面
        将非等距离采样填充为深度间隔为1m的序列, 并应用Akima插值算法, 对数据进行插值
        """
        if (self.velInterpPf == None):
            botDep = self.get_bottomDepth()
            if self.isValid == True:
                # 获得原始声速剖面数据
                dep = self.get_velRawProfile().x
                vel = self.get_velRawProfile().y
                # 对剖面进行垂向插值处理，将非等距离采样填充为深度间隔为1m的序列
                depInterp = np.arange(0, dep[-1]+1, 1)
                # 应用Akima插值算法，对数据进行插值
                if dep.size == 1:
                    velInterp = vel
                else:
                    func = Akima1DInterpolator(dep, vel)
                    velInterp = func(depInterp)
            else:
                depInterp = np.array([0])
                velInterp = np.array([0])
            # 剖面赋值
            velInterpPf = Profile()
            velInterpPf.set_data(depInterp, velInterp)
            velInterpPf.set_coord(self.lat, self.lon)
            velInterpPf.set_bottomDepth(botDep)
            self.velInterpPf = velInterpPf
        return self.velInterpPf

    def get_velExtProfile(self):
        """
        获得延拓后的声速剖面
        """
        if (self.velExtPf == None):
            botDep = self.get_bottomDepth()
            if self.isValid == True:
                # 获得原始声速剖面数据
                dep = self.get_velRawProfile().x
                vel = self.get_velRawProfile().y
                # 获得剖面最低处温盐值
                tem = self.get_temRawProfile().y[-1]
                sal = self.get_salRawProfile().y[-1]
                # 延拓部分深度序列
                depExt = np.arange(np.ceil(dep[-1]), botDep+1)
                # 按照声速公式延拓，假设温盐不变，仅深度变化
                velExt = soundSpeed(tem, sal, depExt)
            else:
                depExt = np.array([0])
                velExt = np.array([0])
            # 剖面赋值
            velExtPf = Profile()
            velExtPf.set_data(depExt, velExt)
            velExtPf.set_coord(self.lat, self.lon)
            velExtPf.set_bottomDepth(botDep)
            self.velExtPf = velExtPf
        return self.velExtPf

    def get_velFullProfile(self):
        '''
        获得全海深声速剖面
        '''
        if (self.velFullPf == None):
            botDep = self.get_bottomDepth()
            if self.isValid == True:
                # 插值后的速度剖面
                depInterp = self.get_velInterpProfile().x
                velInterp = self.get_velInterpProfile().y
                # 延拓后的速度剖面
                depExt = self.get_velExtProfile().x
                velExt = self.get_velExtProfile().y
                # 全海深声速剖面
                depFull = np.concatenate((depInterp, depExt))
                velFull = np.concatenate((velInterp, velExt))
            else:
                depFull = np.array([0])
                velFull = np.array([0])
            # 剖面赋值
            velFullPf = Profile()
            velFullPf.set_data(depFull, velFull)
            velFullPf.set_coord(self.lat, self.lon)
            velFullPf.set_bottomDepth(botDep)
            self.velFullPf = velFullPf
        return self.velFullPf

    def export_velFullProfile_txt(self, filePath):
        '''
        速度剖面导出txt文件
        '''
        velPf = self.get_velFullProfile()
        lat = velPf.lat
        lon = velPf.lon
        botDep = velPf.botDep
        dep = velPf.x
        vel = velPf.y
        # 导出文本
        line1 = 'lat[degree north], {0:>7.4f}, lon[degree east], {1:>7.4f}, BottomDepth[m], {2:>5.0f}\n'.format(
            lat, lon, botDep)
        line2 = 'depth[m], velocity[m/s]'
        np.savetxt(filePath, np.stack((dep, vel), -1),
                   fmt='%.5e', delimiter=',', comments='', header=line1+line2)

    def export_velFullProfile_nc(self, filePath):
        '''
        速度剖面导出nc文件
        '''
        velPf = self.get_velFullProfile()
        lat = velPf.lat
        lon = velPf.lon
        botDep = velPf.botDep
        dep = velPf.x
        vel = velPf.y

        data = Dataset(filePath, 'w', format='NETCDF4')
        data.createDimension('lon', 1)
        data.createDimension('lat', 1)
        data.createDimension('botDep', 1)
        data.createDimension('dep', dep.size)
        data.createDimension('vel', vel.size)

        data.createVariable('lon', 'f', ('lon'))
        data.createVariable('lat', 'f', ('lat'))
        data.createVariable('botDep', 'f', ('botDep'))
        data.createVariable('dep', 'f', ('dep'))
        data.createVariable('vel', 'f', ('vel'))

        data.variables['lat'][:] = lat
        data.variables['lon'][:] = lon
        data.variables['botDep'][:] = botDep
        data.variables['dep'][:] = dep
        data.variables['vel'][:] = vel

        data.close()