import torch


def forward_rate_constants_func(self):
    """Update forward_rate_constants
    """

    ln10 = torch.log(torch.Tensor([10.0])).to(self.device)

    self.forward_rate_constants = (self.Arrhenius_A *
                                   torch.exp(self.Arrhenius_b * torch.log(self.T) -
                                             self.Arrhenius_Ea * 4184.0 / self.R / self.T)
                                   * self.C_M2)

    for i in self.list_reaction_type4:

        reaction = self.reaction[i]

        # high pressure
        self.kinf = (reaction['A'] *
                     torch.exp(reaction['b'] * torch.log(self.T)
                               - reaction['Ea'] * 4184.0 / self.R / self.T))

        # low pressure
        self.k0 = (self.reaction[i]['A_0'] *
                   torch.exp(reaction['b_0'] * torch.log(self.T)
                             - reaction['Ea_0'] * 4184.0 / self.R / self.T))

        Pr = self.k0 * self.C_M[:, i: i + 1] / self.kinf
        lPr = torch.log10(Pr)

        self.k = self.kinf * Pr / (1 + Pr)

        if 'Troe' in self.reaction[i]:
            A = reaction['Troe']['A']
            T1 = reaction['Troe']['T1']
            T2 = reaction['Troe']['T2']
            T3 = reaction['Troe']['T3']

            F_cent = ((1 - A) * torch.exp(-self.T / T3) +
                      A * torch.exp(-self.T / T1) + torch.exp(-T2 / self.T))

            lF_cent = torch.log10(F_cent)
            C = -0.4 - 0.67 * lF_cent
            N = 0.75 - 1.27 * lF_cent
            f1 = (lPr + C) / (N - 0.14 * (lPr + C))

            F = torch.exp(ln10 * lF_cent / (1 + f1 * f1))
            self.k = self.k * F

        self.forward_rate_constants[:, i: i + 1] = self.k

    for i in self.list_reaction_type5:

        reaction = self.reaction[i]

        self.kk = [[None]] * self.n_rate_constants[i]

        # calculate rate expressions at all given pressures
        for j in range(self.n_rate_constants[i]):
            self.kk[j] = (reaction['p_dep']['A'][j] *
                          torch.exp(reaction['b'][j] * torch.log(self.T)
                                    - reaction['Ea'][j] * 4184.0 / self.R / self.T))

        # jhigh1 corresponds to the first Arrhenius expression given at the minumum pressure
        # higher than actual pressure. Considering multiple rate expressions may be given
        # at the same pressure, we need jhigh2, which corresponds to the last Arrhenius
        # expression given at the minumum pressure higher than actual pressure.
        jhigh1 = self.n_rate_constants[i]
        for j in range(self.n_rate_constants[i]):
            if self.P[0] <= reaction['P'][j]:
                jhigh1 = j
                break

        if jhigh1 != self.n_rate_constants[i]:
            for j in range(self.n_rate_constants[i] - 1, -1, -1):
                if reaction['P'][j] == reaction['P'][jhigh1]:
                    jhigh2 = j
                    break

        # jlow1 corresponds to the last Arrhenius expression given at the maximum pressure
        # lower than actual pressure while jlow2 corresponds to the first.
        jlow1 = -1
        for j in range(self.n_rate_constants[i] - 1, -1, -1):
            if self.P[0] >= reaction['P'][j]:
                jlow1 = j
                break

        if jlow1 != -1:
            for j in range(self.n_rate_constants[i]):
                if reaction['P'][j] == reaction['P'][jlow1]:
                    jlow2 = j
                    break

        # This is the case where the actual pressure is higher than all given pressures.
        if jhigh1 == self.n_rate_constants[i]:
            for j in range(self.n_rate_constants[i]):
                if reaction['P'][j] == reaction['P'][jhigh1 - 1]:
                    jhigh2 = j
                    break
            self.k = self.kk[jhigh1 - 1]

            if jhigh2 != jhigh1 - 1:
                for j in range(jhigh2, jhigh1 - 1):
                    self.k = self.k + self.kk[j]

        # This is the case where the actual pressure is lower than all given pressures.
        if jlow1 == -1:
            for j in range(self.n_rate_constants[i] - 1, -1, -1):
                if reaction['P'][j] == reaction['P'][0]:
                    jlow2 = j
                    break
            self.k = self.kk[0]

            if jlow2 != 0:
                for j in range(1, jlow2 + 1):
                    self.k = self.k + self.kk[j]

        # This is the case where the actual pressure is higher than the minimum
        # given pressure and lower than the maximum given pressure.
        if jhigh1 != self.n_rate_constants[i] and jlow1 != -1:

            self.k1 = self.kk[jlow1]

            self.k2 = self.kk[jhigh1]

            if jhigh1 != jhigh2:
                for j in range(jhigh1 + 1, jhigh2 + 1):
                    self.k2 = self.k2 + self.kk[j]

            if jlow1 != jlow2:
                for j in range(jlow2, jlow1):
                    self.k1 = self.k1 + self.kk[j]

            logk = (torch.log(self.k1) + (torch.log(self.k2) - torch.log(self.k1))
                    * (torch.log(self.P[0]) - torch.log(reaction['P'][jlow1])) /
                    (torch.log(reaction['P'][jhigh1]) - torch.log(reaction['P'][jlow1])))

            self.k = torch.exp(logk)

        self.forward_rate_constants[:, i: i + 1] = self.k

    self.forward_rate_constants = self.forward_rate_constants * self.uq_A.abs()


def forward_rate_constants_func_vec(self):
    """Update forward_rate_constants
    """

    ln10 = torch.log(torch.Tensor([10.0])).to(self.device)

    self.forward_rate_constants = (self.Arrhenius_A *
                                   torch.exp(self.Arrhenius_b * torch.log(self.T) -
                                             self.Arrhenius_Ea * 4184.0 / self.R / self.T)
                                   * self.C_M2)

    # dealing with type 4
    self.k0 = (self.Arrhenius_A0 *
               torch.exp(self.Arrhenius_b0 * torch.log(self.T) -
                         self.Arrhenius_Ea0 * 4184.0 / self.R / self.T))

    self.kinf = (self.Arrhenius_Ainf *
                 torch.exp(self.Arrhenius_binf * torch.log(self.T) -
                           self.Arrhenius_Eainf * 4184.0 / self.R / self.T))

    Pr = self.k0 * self.C_M_type4 / self.kinf

    # transfer the size to match the overall rate matrix
    self.falloff_matrix = (torch.mm(Pr / (1 + Pr), self.mat_transfer_type4) +
                           (1 - self.is_falloff) * self.identity_mat)

    # dealing with Troe
    lPr = torch.mm(torch.log10(Pr), self.mat_transfer_type4_to_Troe)
    F_cent = ((1 - self.Troe_A) * torch.exp(-self.T / self.Troe_T3) +
              self.Troe_A * torch.exp(-self.T / self.Troe_T1) +
              torch.exp(-self.Troe_T2 / self.T))

    lF_cent = torch.log10(F_cent)
    C = -0.4 - 0.67 * lF_cent
    N = 0.75 - 1.27 * lF_cent
    f1 = (lPr + C) / (N - 0.14 * (lPr + C))

    F = torch.exp(ln10 * lF_cent / (1 + f1 * f1))

    # transfer the size to match the overall rate matrix
    self.Troe_matrix = (torch.mm(F, self.mat_transfer_type4_Troe) +
                        (1 - self.is_Troe_falloff) * self.identity_mat)

    self.forward_rate_constants = (
        self.forward_rate_constants * self.Troe_matrix * self.falloff_matrix)

    # for uncertainty quantification
    self.forward_rate_constants = self.forward_rate_constants * self.uq_A.abs()


def forward_rate_constants_func_matrix(self):
    self.kf = (self.Arrhenius_A *
               torch.exp(self.Arrhenius_b * torch.log(self.T) -
                         self.Arrhenius_Ea * 4184.0 / self.R / self.T))


def equilibrium_constants_func(self):
    vk = (-self.reactant_stoich_coeffs + self.product_stoich_coeffs)
    delta_S_over_R = torch.mm(self.S0, vk) / self.R
    delta_H_over_RT = torch.mm(self.H, vk) / self.R / self.T

    self.equilibrium_constants = (torch.exp(
        delta_S_over_R - delta_H_over_RT + torch.log(self.one_atm / self.R / self.T) * vk.sum(dim=0)))


def reverse_rate_constants_func(self):
    self.reverse_rate_constants = (
        self.forward_rate_constants / self.equilibrium_constants * self.is_reversible)


def wdot_func(self):
    # eps = 1e-300
    nt=self.forward_rate_constants.size()[0]
    temp_forward=torch.ones(nt,self.n_reactions).to(self.device)
    temp_reverse=torch.ones(nt,self.n_reactions).to(self.device)    
    for i in range(self.n_reactions):
        multiply_temp_forward=torch.ones(nt).to(self.device)
        multiply_temp_reverse=torch.ones(nt).to(self.device)
        list_forward=[]
        list_reverse=[]
        for j in range(self.n_species):
            if self.reactant_orders[j,i]!=0:
                list_forward.append(j)
            if self.product_stoich_coeffs[j,i]!=0:
                list_reverse.append(j)
            multiply_temp_forward*=self.C[:,j]**self.reactant_orders[j,i]
            multiply_temp_reverse*=self.C[:,j]**self.product_stoich_coeffs[j,i]
        temp_forward[:,i]=multiply_temp_forward
        temp_reverse[:,i]=multiply_temp_reverse
        len_for=len(list_forward)
        len_rev=len(list_reverse)

        for t in range (nt):
            counter_for=0
            counter_rev=0
            flag_negative_forward=0
            flag_negative_reverse=0
            for j in list_forward:
                if self.C[t,j]<0:
                    counter_for+=1
            if counter_for>1 or (len_for==1 and counter_for==1):
                flag_negative_forward=1
            for j in list_reverse:
                if self.C[t,j]<0:
                    counter_rev+=1
            if counter_rev>1 or (len_rev==1 and counter_rev==1):
                flag_negative_reverse=1
                
            if flag_negative_forward==1 and temp_forward[t,i]>0:
                temp_forward[t,i]=0
            if flag_negative_reverse==1 and temp_reverse[t,i]>0:
                temp_reverse[t,i]=0
    self.forward_rates_of_progress = self.forward_rate_constants * temp_forward
    self.reverse_rates_of_progress = self.reverse_rate_constants * temp_reverse 
    
    # for i in range(self.n_reactions):
    #     list_forward=[]
    #     list_reverse=[]
    #     for j in range(self.n_species):
    #         if self.reactant_orders[j,i]!=0:
    #             list_forward.append(j)
    #         if self.product_stoich_coeffs[j,i]!=0:
    #             list_reverse.append(j)
    #     for t in range (nt):
    #         flag_negative_forward=1
    #         flag_negative_reverse=1
    #         for j in list_forward:
    #             if self.C[t,j]>0:
    #                 flag_negative_forward=0
    #         for j in list_reverse:
    #             if self.C[t,j]>0:
    #                 flag_negative_reverse=0
    #         if flag_negative_forward==1 and self.forward_rates_of_progress[t,i]>0:
    #             self.forward_rates_of_progress[t,i]=0
    #         if flag_negative_reverse==1 and self.reverse_rates_of_progress[t,i]>0:
    #             self.reverse_rates_of_progress[t,i]=0
    
    # nt=self.forward_rate_constants.size()[0]
    # temp_forward=torch.ones_like(self.forward_rate_constants)
    # temp_reverse=torch.ones_like(self.reverse_rate_constants)   
    # for t in range(nt):
    #     for i in range(self.n_reactions):
    #         multiply_temp_forward=1
    #         multiply_temp_reverse=1
    #         for j in range(self.n_species):
         
    #             multiply_temp_forward*=self.C[t,j]**self.reactant_orders[j,i]
    #             multiply_temp_reverse*=self.C[t,j]**self.product_stoich_coeffs[j,i]
    #         temp_forward[t,i]=multiply_temp_forward
    #         temp_reverse[t,i]=multiply_temp_reverse        
    # self.forward_rates_of_progress = self.forward_rate_constants * temp_forward
    # self.reverse_rates_of_progress = self.reverse_rate_constants * temp_reverse     
    
    # C_eps=torch.zeros_like(self.C).fill_(eps)
    # C_temp=torch.where(self.C<0,C_eps,self.C+eps)
    # self.forward_rates_of_progress = self.forward_rate_constants * \
    #     torch.exp(torch.mm(torch.log(C_temp), self.reactant_stoich_coeffs))
    # self.reverse_rates_of_progress = self.reverse_rate_constants * \
    #     torch.exp(torch.mm(torch.log(C_temp), self.product_stoich_coeffs))              
    
    
    # self.forward_rates_of_progress = self.forward_rate_constants * \
    #     torch.exp(torch.mm(torch.log(self.C+eps), self.reactant_stoich_coeffs))

    # self.reverse_rates_of_progress = self.reverse_rate_constants * \
    #     torch.exp(torch.mm(torch.log(self.C+eps), self.product_stoich_coeffs))     
        
    self.qdot = self.forward_rates_of_progress - self.reverse_rates_of_progress

    self.wdot = torch.mm(self.qdot, self.net_stoich_coeffs.T)

    self.net_production_rates = self.wdot


def Ydot_func(self):

    self.Ydot = self.wdot / self.density_mass * self.molecular_weights.T

          
           

def Xdot_func(self):

    self.Xdot = self.Ydot * self.mean_molecular_weight / self.molecular_weights.T


def Tdot_func(self):

    self.Tdot = -((self.partial_molar_enthalpies * self.wdot).sum(dim=1) /
                  self.density_mass.T / self.cp_mass).T


def TXdot_func(self):

    self.Tdot_func()

    self.Xdot_func()

    self.TXdot = torch.cat((self.Tdot, self.Xdot), dim=1)

    return self.TXdot


def TYdot_func(self):

    self.Tdot_func()

    self.Ydot_func()

    self.TYdot = torch.cat((self.Tdot, self.Ydot), dim=1)

    return self.TYdot
