import numpy as np
from scipy.signal import find_peaks
from pyat.ProfileDataType import Channel, SpringLayer

'''
************************************************************
该模块用于分析剖面, 根据标准化剖面，得到声道, 跃层分析结果
************************************************************
'''


class ProfileTool():
    '''
    剖面分析工具, 根据标准化剖面，得到声道, 跃层分析结果
    '''

    def __init__(self):
        self.clearVar()

    def clearVar(self):
        """
        清空函数调用记录
        """
        self.maxTurn = None
        self.minTurn = None
        self.surfCh = None
        self.deepCh = None
        self.surfSL = None
        self.mainSL = None
        return

    def set_profile(self, velPf):
        '''
        设置剖面
        '''
        self.velPf = velPf
        self.clearVar()

    def findMaxTurn(self):
        '''
        寻找最大值的拐点
        '''
        if (self.maxTurn == None):
            dep = self.velPf.x
            vel = self.velPf.y
            # 寻找极大值，设定窗口宽度为20m
            peaks, properties = find_peaks(vel, width=20)
            if (peaks.size != 0):
                peaks = peaks[vel[peaks] > vel[0]]
            if (peaks.size != 0):
                # 定位极大值中最大的那个
                peak = peaks[np.argmax(vel[peaks])]
                self.maxTurn = peak
            else:
                self.maxTurn = -1
        return self.maxTurn
    
    def findMinTrun(self):
        '''
        寻找最小值的拐点
        '''
        if (self.minTurn == None):
            dep = self.velPf.x
            vel = self.velPf.y
            # 寻找最小值拐点，设定窗口宽度为20m
            peaks, properties = find_peaks(-vel, width=20)
            if (peaks.size != 0):
                peaks = peaks[vel[peaks] < vel[0]]
            if (peaks.size != 0):
                # 定位极小值中最小的那个
                peak = peaks[np.argmin(vel[peaks])]
                self.minTurn = peak
            else:
                self.minTurn = -1
        return self.minTurn
        
    def get_surfaceChannel(self):
        '''
        计算表面声道
        '''
        if (self.surfCh == None):
            dep = self.velPf.x
            vel = self.velPf.y
            maxTurn = self.findMaxTurn()
            minTurn = self.findMinTrun()
            if maxTurn != -1:
                self.surfCh = Channel()
                self.surfCh.exist = True
                self.surfCh.depAxis = dep[0]
                self.surfCh.velAxis = vel[0]
            else:
                self.surfCh = Channel()
                self.surfCh.exist = False
        return self.surfCh

    def get_deepChannel(self):
        '''
        计算深海声道
        '''
        if (self.deepCh == None):
            dep = self.velPf.x
            vel = self.velPf.y
            maxTurn = self.findMaxTurn()
            minTurn = self.findMinTrun()
            if minTurn != -1:
                self.deepCh = Channel()
                self.deepCh.exist = True
                self.deepCh.depAxis = dep[minTurn]
                self.deepCh.velAxis = vel[minTurn]
            else:
                self.deepCh = Channel()
                self.deepCh.exist = False
        return self.deepCh

    def get_springLayer(self):
        '''
        计算跃层
        '''
        if (self.surfSL == None):
            dep = self.velPf.x
            vel = self.velPf.y
            # 寻找拐点
            maxTurn = self.findMaxTurn()
            minTurn = self.findMinTrun()
            # 初始化跃层
            surfSL = SpringLayer()
            surfSL.exist = False
            mainSL = SpringLayer()
            mainSL.exist = False
            if (maxTurn != -1 and minTurn != -1):
                # 对应这种形状
                #  \
                #   \
                #   /
                #  /
                #  \
                #   \
                surfSL.exist = True
                surfSL.limit = np.array([dep[0], dep[maxTurn]])
                surfSL.intensity = (vel[maxTurn]-vel[0])/(dep[maxTurn]-dep[0])
                mainSL.exist = True
                mainSL.limit = np.array([dep[maxTurn], dep[minTurn]])
                mainSL.intensity = (
                    vel[minTurn]-vel[maxTurn])/(dep[minTurn]-dep[maxTurn])
            elif (maxTurn != -1 and minTurn == -1):
                # 对应这种形状
                #  \
                #   \
                #   /
                #  /
                # /
                surfSL.exist = True
                surfSL.limit = np.array([dep[0], dep[maxTurn]])
                surfSL.intensity = (vel[maxTurn]-vel[0])/(dep[maxTurn]-dep[0])
                mainSL.exist = True
                mainSL.limit = np.array([dep[maxTurn], dep[-1]])
                mainSL.intensity = (vel[-1]-vel[maxTurn]) / \
                    (dep[-1]-dep[maxTurn])
            elif (maxTurn == -1 and minTurn != -1):
                # 对应这种形状
                #   /
                #  /
                # /
                # \
                #  \
                #   \
                mainSL.exist = True
                mainSL.limit = np.array([dep[0], dep[minTurn]])
                mainSL.intensity = (vel[minTurn]-vel[0])/(dep[minTurn]-dep[0])
            elif (maxTurn == -1 and minTurn == -1):
                if vel[-1] < vel[0]:
                    # 对应这种形状
                    #   /
                    #  /
                    # /
                    mainSL.exist = True
                    mainSL.limit = np.array([dep[0], dep[-1]])
                    mainSL.intensity = (vel[-1]-vel[0])/(dep[-1]-dep[0])
                else:
                    # 对应这种形状
                    # \
                    #  \
                    #   \
                    surfSL.exist = True
                    surfSL.limit = np.array([dep[0], dep[-1]])
                    surfSL.intensity = (vel[-1]-vel[0])/(dep[-1]-dep[0])
            self.surfSL = surfSL
            self.mainSL = mainSL
        return (self.surfSL, self.mainSL)
    
    def export_profileParm(self, fileName):
        '''
        计算结果输出文件
        '''
        lat = self.velPf.lat
        lon = self.velPf.lon
        botDep = self.velPf.botDep
        # 计算表面声道
        surfCh = self.get_surfaceChannel()
        # 计算深海声道
        deepCh = self.get_deepChannel()
        # 计算跃层信息
        surfSL, mainSL = self.get_springLayer()
        f = open(fileName, mode='w')
        # 输出坐标深度
        line = 'lat[degree north], {0:>7.4f}, lon[degree east], {1:>7.4f}, BottomDepth[m], {2:>5.0f}\n'.format(
            lat, lon, botDep)
        f.write(line)
        f.write('----------------------------\n')
        if surfCh.exist:
            f.write('表面声道轴深度: {0:>.2f}\n'.format(surfCh.depAxis))
            f.write('表面声道声速: {0:>.4f}\n'.format(surfCh.velAxis))
        else:
            f.write('表面声道轴深度: NONE\n')
            f.write('表面声道声速: NONE\n')
        f.write('----------------------------\n')
        if deepCh.exist:
            f.write('深海声道轴深度: {0:>.2f}\n'.format(deepCh.depAxis))
            f.write('深海声道声速: {0:>.4f}\n'.format(deepCh.velAxis))
        else:
            f.write('深海声道轴深度: NONE\n')
            f.write('深海声道声速: NONE\n')
        f.write('----------------------------\n')
        if surfSL.exist:
            f.write('表面跃层区间: {0:>.2f}, {1:>.2f}\n'.format(surfSL.limit[0], surfSL.limit[1]))
            f.write('表面跃层强度: {0:>.4f}\n'.format(surfSL.intensity))
        else:
            f.write('表面跃层区间: NONE\n')
            f.write('表面跃层强度: NONE\n')
        f.write('----------------------------\n')
        if mainSL.exist:
            f.write('主跃层区间: {0:>.2f}, {1:>.2f}\n'.format(mainSL.limit[0], mainSL.limit[1]))
            f.write('主跃层强度: {0:>.4f}\n'.format(mainSL.intensity))
        else:
            f.write('主跃层区间: NONE\n')
            f.write('主跃层强度: NONE\n')
        f.close()