import numpy as np
import settings as stg


class AcousticInversionMethodHighConcentration():

    """ Thi class compute acoustic inversion method adapted for high suspended sediments concentration
     For instance, the case of dam flush downstream Rhone-Isere confluence (07/01/2018) evaluated at ~ 10g/L  """

    def __init__(self):

        pass

    # ==========================================
    #                 Functions
    # ==========================================

    # ---------- Computing sound speed ------------- #
    def water_velocity(self, T):
        """Computing sond speed from Bilaniuk and Wong 1993"""
        C = 1.40238744 * 1e3 + 5.03836171 * T - 5.81172916 * 1e-2 * T ** 2 + 3.34638117 * 1e-4 * T ** 3 - \
            1.48259672 * 1e-6 * T ** 4 + 3.16585020 * 1e-9 * T ** 5
        return C

    # -------- Computing water attenuation coefficient ----------- #
    def water_attenuation(self, freq1, freq2, T):
        """Computing attenuation from François and Garrison 1982"""
        if T > 20:
            alpha = (3.964 * 1e-4 - 1.146 * 1e-5 * T + 1.45 * 1e-7 * T ** 2 - 6.5 * 1e-10 * T ** 3) * 1e-3 * \
                    (np.log(10) / 20) * (np.array([freq1, freq2]) * 1e-3) ** 2
        else:
            alpha = (4.937 * 1e-4 - 2.59 * 1e-5 * T + 9.11 * 1e-7 * T ** 2 - 1.5 * 1e-8 * T ** 3) * 1e-3 * \
                    (np.log(10) / 20) * (np.array([freq1, freq2]) * 1e-3) ** 2
        return alpha

    # ---------- Conmpute FBC ----------
    # def compute_FCB(self):
    #     # print(self.BS_averaged_cross_section_corr.V.shape)
    #     # print(self.r_2D.shape)
    #     FCB = np.zeros((256, 4, 1912))
    #     for f in range(4):
    #         # print(self.alpha_w_function(self.Freq[f], self.temperature))
    #         FCB[:, f, :] = np.log(self.BS_averaged_cross_section_corr.V[:, f, :]) + np.log(self.r_3D[:, f, :]) + \
    #                        np.log(2 * self.alpha_w_function(self.Freq[f], self.temperature) * self.r_3D[:, f, :])
    #     return FCB

    # ------------- Computing ks ------------- #

    def ks(self, ind):
        ks0 = np.array([0.0446681, 0.11490388, 0.35937832, 2.5025668])
        ks = ks0[ind]
        return ks

    # ------------- Computing sv ------------- #

    def sv(self):
        pass

    # ------------- Computing X ------------- #
    def X_exponent(self, ind):
        X0 = [3.688664080521131, 3.451418735488537, 0, 3.276577078823426, 0, 0]
        X = X0[ind]
        return X

    # -------------- Computing Kt -------------- #
    def kt_corrected(self, r, water_velocity, RxGain, TxGain, kt_ref):
        """Computing the instrument constant Kt that depends on gain and temperature"""
        # Cell size
        delta_r = r[1] - r[0]
        # Pulse length
        tau = 2 * delta_r / 1500
        # Sound speed
        cel = water_velocity
        # Reference pulse length
        tau_ref = 13.33333e-6
        # Reference sound speed
        c_ref = 1475
        # Gain
        gain = 10 ** ((RxGain + TxGain) / 20)
        # Computing Kt
        kt0 = kt_ref * gain * np.sqrt(tau * cel / (tau_ref * c_ref))    # 1D numpy array
        kt = np.reshape(kt0, (1, 2))     # convert to 2d numpy array to compute J_cross_section
        return kt

    # ------------- Computing J_cross_section ------------- #
    def j_cross_section(self, BS, r2D, kt):
        J_cross_section = np.zeros((2, BS.shape[0], BS.shape[2]))       # 2 because it's a pair of frequencies
        for k in range(2):
            J_cross_section[k, :, :] = (3 / (16 * np.pi)) * ((BS[:, k, :]**2 * r2D**2) / kt[k, :, :]**2)
        # J_cross_section[J_cross_section == 0] = np.nan
        print("compute j_cross_section finished")
        return J_cross_section

    # ------------- Computing alpha_s ------------- #

    def alpha_s(self):
        pass

    # ------------- Computing interpolation of fine SSC data obtained from water sampling -------------
    # -------------           collected at various depth in the vertical sample           -------------
    def M_profile_SCC_fine_interpolated(self):
        pass

    # ------------- Computing zeta ------------- #
    def zeta(self, ind1, ind2):
         # zeta = alpha_s / ((1/r) * (M_profile_SSC_fine))
        zeta0 = np.array([0.04341525, 0.04832906, 0.0847188, np.nan])
        zeta = zeta0[[ind1, ind2]]
        return zeta

    # ------------- Computing VBI ------------- #
    def VBI_cross_section(self, freq1, freq2,
                          zeta_freq1, zeta_freq2,
                          j_cross_section_freq1, j_cross_section_freq2,
                          r2D,
                          water_attenuation_freq1, water_attenuation_freq2,
                          X):

        # print('self.zeta_exp[ind_j].shape', self.zeta_exp[ind_j])
        # print('np.log(self.j_cross_section[:, ind_i, :]).shape', np.log(self.j_cross_section[:, ind_i, :]).shape)
        # print('self.r_3D[:, ind_i, :]', self.r_3D[:, ind_i, :].shape)
        # print('self.water_attenuation[ind_i]', self.water_attenuation[ind_i])
        # print('self.x_exp[0.3-1 MHz]', self.x_exp['0.3-1 MHz'].values[0])
        # print("start computing VBI")

        logVBI = ((zeta_freq2 *
                  np.log(j_cross_section_freq1 * np.exp(4 * r2D * water_attenuation_freq1) /
                         (freq1 ** X)) -
                  zeta_freq1 *
                  np.log(j_cross_section_freq2 * np.exp(4 * r2D * water_attenuation_freq2) /
                         (freq2 ** X))) /
                  (zeta_freq2 - zeta_freq1))

        print("compute VBI finished")

        return np.exp(logVBI)

    # ------------- Computing SSC fine ------------- #
    def SSC_fine(self, zeta, r2D, VBI, freq, X, j_cross_section):
        SSC_fine = (1/zeta) * ( 1/(4 * r2D) * np.log((VBI * freq**X) / j_cross_section) )
        print("compute SSC fine finished")
        return SSC_fine

    # ------------- Computing SSC sand ------------- #
    def SSC_sand(self, VBI, freq, X, ks):
        SSC_sand = (16 * np.pi * VBI * freq ** X) / (3 * ks**2)
        print("compute SSC sand finished")
        return SSC_sand