#############################################################################
##Mooring simulator, 2014                                                ####
##Edited by Arnaud Le Fur, IFREMER                                       ####
#############################################################################

from PyQt4 import QtGui
import copy
import numpy as np
from math import cos, sin, atan, sqrt
from mySimulateCanvas import MySimulateCanvas

class Simulate_wind(QtGui.QWidget):
    """Classe contenant la fenetre de simulation """

    def __init__(self, data, parent=None):
        """initialisation, calcul et affichage des resultats de la simulation
        Input : data : Tableau reunissant toutes les donnees necessaires a la simulation
                data[0] : Liste des objets constituants le mouillage
                data[1] : Les profondeurs et valeurs de courants
                data[2] : Liste des profondeur fixe des instruments
                data[3] : Trainee induite par les elements lors de la chute du lest
                data[4] : Largeur de l ecran
                data[5] : Chemin du programme
                data[6] : Liste des ratio de clampage des instruments
                """
        super(Simulate_wind, self).__init__(parent)

        l_max = 10.  # Longueur max d un element
        count = 0
        self.release_ind = 0
        self.data = data
        self.table_match = np.zeros(len(self.data[0]))

        # Corrige la saisie de lutilisateur en cas de saisie positive du poids du lest
        if float(eval(self.data[0][-1].mass)) > 0:
            self.data[0][-1].mass = str(-1*float(eval(self.data[0][-1].mass)))

        # Calcul l inventaire des elements utilise sur la ligne
        self.inventory = self.calculate_inventory()
        # Calcul des profondeurs en statique, sans courant, ni allongement
        self.static_depth = self.calculate_static()
        # Calcul du temps, de la vitesse de chute et des tensions max lors de cette chute
        self.T, self.V_chute_t, self.T_max = self.resolve_max_Tension()
        self.percent_max_T = []
        self.weight_kg = []
        self.old_depth = []
        # Calcul du poids minimal,ideal,maximal pour le lest
        self.max_value_anchor = self.Anchor_max_value()

        #Calcul_des_tensions_max/Breaking_strenght#
        for i in range(len(self.T_max)):
            if type(self.data[0][i]) != list:
                if hasattr(self.data[0][i], "breaking_strength"):
                    self.percent_max_T.append(
                        100*self.T_max[i]/(9.81*eval(self.data[0][i].breaking_strength)))
                else:
                    self.percent_max_T.append(0.0)
                # On recupere egalement le poids de chaque element en kg
                self.weight_kg.append(eval(self.data[0][i].mass))
            else:
                if hasattr(self.data[0][i][0], "breaking_strength"):
                    self.percent_max_T.append(
                        100*self.T_max[i]/(9.81*eval(self.data[0][i][0].breaking_strength)))
                else:
                    self.percent_max_T.append(0.0)
                self.weight_kg.append(
                    eval(self.data[0][i][0].mass)+eval(self.data[0][i][1].mass))

        # On calcul la flottabilite en cas de rupture de la partie superieure en kg
        self.get_buoyancy_kg()

        # On met en place une table de correspondance entre la ligne dorigine et la ligne decoupee en troncon
        for i in range(len(self.data[0])):
            self.table_match[i] = i+count
            if type(self.data[0][i]) != list:
                if eval(self.data[0][i].length) > l_max:
                    coeff = round(eval(self.data[0][i].length)/l_max)
                    count = count+coeff-1
            else:
                if eval(self.data[0][i][0].length) > l_max:
                    coeff = round(eval(self.data[0][i][0].length)/l_max)
                    count = count+coeff-1

        i = 0
        # On decoupe les cables en troncon de longueur minimale : cable de 100m = 10 x cable de 10 m
        while i < len(self.data[0]):
            if type(self.data[0][i]) != list:
                if eval(self.data[0][i].length) > l_max:
                    #on met a jour le poids, la longueur, et la surface#
                    coeff = round(eval(self.data[0][i].length)/l_max)
                    new_length = str(eval(self.data[0][i].length)/coeff)
                    new_mass = str(eval(self.data[0][i].mass)/coeff)
                    new_projected_area = str(
                        eval(self.data[0][i].projected_area)/coeff)
                    self.data[0][i].length = new_length
                    self.data[0][i].mass = new_mass
                    self.data[0][i].projected_area = new_projected_area

                    # on ajoute le nouveau troncon en position i+1
                    for j in range(int(coeff)-1):
                        self.data[0].insert(
                            i+1, copy.deepcopy(self.data[0][i]))
                        self.data[2].insert(i+1, 0)

            else:  # Element clampes
                if eval(self.data[0][i][0].length) > l_max:
                    coeff = round(eval(self.data[0][i][0].length)/l_max)
                    new_length = str(eval(self.data[0][i][0].length)/coeff)
                    new_mass = str(eval(self.data[0][i][0].mass)/coeff)
                    new_projected_area = str(
                        eval(self.data[0][i][0].projected_area)/coeff)
                    self.data[0][i][0].length = new_length
                    self.data[0][i][0].mass = new_mass
                    self.data[0][i][0].projected_area = new_projected_area
                    count = 0
                    for j in range(int(coeff)-1):
                        self.data[0].insert(
                            i+1, copy.deepcopy(self.data[0][i][0]))
                        self.data[2].insert(i+1, 0)
                        count = count+1
            i = i+1

        self.depth = self.make_vector()  # Initialisation du vecteur profondeur
        self.current_vector(self.depth)  # Initialisation du vecteur de courant
        alpha, F, P = self.resolve()  # Resolution du systeme
        self.calculate_stretch(F)  # Calcul des allongements
        self.calculate_ratio(alpha)  # Calcul des ratios de clampage
        self.plot_figure(alpha, P, F)  # Plot des 4 graphiques de la simulation
        # Genere le tableau recapitulatif de la simulation
        self.make_table(alpha, F)

    def calculate_inventory(self):
        """Cette fonction permet de retourner un tableau contenant la liste des elements,longueurs et le nombre dutilisation
        Output : Result : Tableau contenant l inventaire"""
        inventory_name = []
        inventory_length = []
        Result = [["Name", "Length (m) / Number"]]
        #On recupere d abord tout les noms des elements#
        for i in range(len(self.data[0])):
            if type(self.data[0][i]) != list:
                if not(self.data[0][i].name) in inventory_name:
                    inventory_name.append(self.data[0][i].name)
                    inventory_length.append(0)
            else:
                if not(self.data[0][i][0].name) in inventory_name:
                    inventory_name.append(self.data[0][i][0].name)
                    inventory_length.append(0)
                if not(self.data[0][i][1].name) in inventory_name:
                    inventory_name.append(self.data[0][i][1].name)
                    inventory_length.append(0)
        #On incremente ensuite soit son nombre dutilisation soit la longueur si c est un cable#
        for i in range(len(self.data[0])):
            if type(self.data[0][i]) != list:
                ind = inventory_name.index(self.data[0][i].name)
                if self.data[0][i].__class__.__name__ == "Ropes":
                    inventory_length[ind] = inventory_length[ind] + \
                        eval(self.data[0][i].length)
                else:
                    inventory_length[ind] = inventory_length[ind]+1
            else:  # Elements clampes
                ind = inventory_name.index(self.data[0][i][0].name)
                if self.data[0][i][0].__class__.__name__ == "Ropes":
                    inventory_length[ind] = inventory_length[ind] + \
                        eval(self.data[0][i][0].length)
                else:
                    inventory_length[ind] = inventory_length[ind]+1
                ind = inventory_name.index(self.data[0][i][1].name)
                if self.data[0][i][1].__class__.__name__ == "Ropes":
                    inventory_length[ind] = inventory_length[ind] + \
                        eval(self.data[0][i][1].length)
                else:
                    inventory_length[ind] = inventory_length[ind]+1
        for i in range(len(inventory_name)):
            Result.append([inventory_name[i], inventory_length[i]])
        return Result

    def calculate_static(self):
        """Cette fonction permet de sauvegarder les longueurs d origines et retourne un vecteur avec les profondeurs(en haut de l element)
        en statique : sans allongement ni courant
        Output : subduct : vecteur de profondeur statique"""
        subduct = [self.data[1][1][-1]+eval(self.data[0][-1].length)]
        self.original_length = []
        for b in self.data[0]:
            if type(b) != list:
                self.original_length.append(round(eval(b.length), 1))
            else:
                self.original_length.append(
                    [round(eval(b[0].length), 1), round(eval(b[1].length), 1)])

        for i in range(2, len(self.data[0])+1):

            if type(self.data[0][-i]) != list:
                subduct.insert(0, subduct[0]+eval(self.data[0][-i].length))
            else:
                subduct.insert(0, subduct[0]+eval(self.data[0][-i][0].length))
        return subduct

    def resolve_max_Tension(self):
        """Cette fonction permet de calculer la vitesse, le temps de chute et les tensions maximales lors de cette chute
        Output : T : Vecteur temps
                V_chute_t : Vecteur vitesse de chute en fonction du temps
                T_max : Tension max pour chaque element """

        Tension = []
        T_max = []
        g = 9.81
        h = 0
        T = [0]
        t = 0.
        mass_lest = eval(self.data[0][-1].mass)
        floor_depth = abs(self.data[1][1][len(self.data[1][0])-1])
        #Calcul de la flottabilite#
        for i in range(len(self.data[0])):
            Tension.append(0)
            for j in range(i+1):
                if j == len(self.data[0]):
                    break
                else:
                    if type(self.data[0][j]) != list:
                        Tension[i] = Tension[i]+g*eval(self.data[0][j].mass)
                    else:
                        Tension[i] = Tension[i]+g * \
                            (eval(self.data[0][j][0].mass) +
                            eval(self.data[0][j][1].mass))

        # Trainee induite par les elements superieures au lest
        a = (self.data[3][-1])/(-mass_lest)
        b = g-(Tension[-2]/(-mass_lest))  # Tension au niveau du lest
        if b > 0:
            i = 0
            cond = 0
            t = 0.1
            Time = 0.
            V_chute_t = [0]
            while cond == 0:
                # Evaluation de la variation de la vitesse de chute en fonction du temps cf L.MARIE
                Time = Time+t
                T.append(Time)
                V_chute_t.append((b-a*V_chute_t[-1]**2)*t+V_chute_t[-1])
                h = h+V_chute_t[i]*t
                if h >= floor_depth:  # On touche le fond
                    cond = 1
                i = i+1

            v_chute = V_chute_t[-1]  # Vitesse en regime permanent
            for i in range(len(Tension)):
                T_max.append(0)
                # Calcul des Tensions max
                T_max[i] = (self.data[3][i]*(v_chute**2)+Tension[i])
            T_max[-1] = 0.0
            return T, V_chute_t, T_max
        else:
            QtGui.QMessageBox.warning(
                self, 'Message', "Your Anchor is too light")

    def Anchor_max_value(self):
        """Cette fonction permet de calculer quel est le poids maximal du lest qui entraine une Tension max
        egale a la charge de rupture
        Output : masse maximale"""
        vec_v_chute = []
        Tension = []
        g = 9.81
        #Calcul de la flottabilite#
        for i in range(len(self.data[0])):
            Tension.append(0)
            for j in range(i+1):
                if j == len(self.data[0]):
                    break
                else:
                    if type(self.data[0][j]) != list:
                        Tension[i] = Tension[i]+g*eval(self.data[0][j].mass)
                    else:
                        Tension[i] = Tension[i]+g * \
                            (eval(self.data[0][j][0].mass) +
                            eval(self.data[0][j][1].mass))
        #On evalue la vitesse de chute maximale que supporte les elements du mouillage, depend de leur charge de rupture#
        for i in range(len(self.data[0])):
            if type(self.data[0][i]) != list:
                if hasattr(self.data[0][i], "breaking_strength"):
                    v_chute = sqrt(
                        (eval(self.data[0][i].breaking_strength)*g-Tension[i])/(self.data[3][i]))
                    vec_v_chute.append(v_chute)
            else:
                if hasattr(self.data[0][i][0], "breaking_strength"):
                    v_chute = sqrt(
                        (eval(self.data[0][i][0].breaking_strength)*g-Tension[i])/(self.data[3][i]))
                    vec_v_chute.append(v_chute)

        # On prend l element qui supporte la vitesse de chute la plus petite, cas le plus defavorable
        mass_max = (Tension[-2]+self.data[3][-1]*(min(vec_v_chute))**2)/g
        return mass_max

    def get_buoyancy_kg(self):
        """Cette fonction permet de calculer la flottabilite restante en cas de sectionnement de la partie superieure en kg """
        release_ind = 0
        for i in range(len(self.data[0])):
            # On recupere l indice du largueur car ce calcule nest pas pertinent en dessous de celui-ci
            if(self.data[0][i].__class__.__name__) == "Releases":
                release_ind = i
        self.Buoyancy_kg = [self.weight_kg[release_ind]]
        for i in reversed(range(release_ind)):
            self.Buoyancy_kg.insert(0, self.Buoyancy_kg[0]+self.weight_kg[i])
        while len(self.Buoyancy_kg) != len(self.weight_kg):
            self.Buoyancy_kg.append(0.0)

    def make_vector(self):
        """Cette fonction permet de calculer la profondeur de chaque element(en son centre)
        Output : depth : vecteur profondeur"""
        depth = [self.data[1][1][len(self.data[1][0])-1]]

        #Calcul de la profondeur en haut de l element#
        for i in range(2, len(self.data[0])+2):
            depth.insert(0, 0)
            if type(self.data[0][-i+1]) != list:
                depth[-i] = depth[-i+1]+eval(self.data[0][-i+1].length)
            else:
                depth[-i] = depth[-i+1]+eval(self.data[0][-i+1][0].length)
        depth.remove(self.data[1][1][len(self.data[1][0])-1])
        #On enleve la moitie pour obtenir la profondeur du milieu de lelement#
        for i in range(len(self.data[0])):
            if type(self.data[0][i]) != list:
                depth[i] = depth[i]-eval(self.data[0][i].length)/2
            else:
                depth[i] = depth[i]-eval(self.data[0][i][0].length)/2
        return depth

    def current_vector(self, depth):
        """Cette fonction permet dobtenir une valeur de courant par interpolation pour chaque profondeur du vecteur depth
        Input : Vecteur profondeur
        Output : Vecteur courant """
        current_vector = []
        # On interpole lineairement, on calcul donc les coefficients k et b de la droite y=k*x+b
        for i in range(len(depth)):
            a = 0

            while(depth[i] < self.data[1][1][a]):
                a = a+1
            k = (self.data[1][0][a-1]-self.data[1][0][a]) / \
                (self.data[1][1][a-1]-self.data[1][1][a])
            b = self.data[1][0][a-1]-k*self.data[1][1][a-1]
            current_vector.append(k*depth[i]+b)

        return current_vector

    def update_depth(self, depth, alpha):
        """Cette fonction permet de mettre a jour le vecteur profondeur en fonction de l inclinaison de la ligne
        Input : depth : vecteur profondeur d origine
                alpha : vecteur angle
        Output : depth : vecteur profondeur mis a jour"""
        depth = [self.data[1][1][len(self.data[1][0])-1]]

        #Calcul de la profondeur en haut de l element#
        for i in range(2, len(self.data[0])+2):
            depth.insert(0, 0)
            if type(self.data[0][-i+1]) != list:
                depth[-i] = depth[-i+1] + \
                    cos(alpha[-i+1])*eval(self.data[0][-i+1].length)
            else:
                depth[-i] = depth[-i+1] + \
                    cos(alpha[-i+1])*eval(self.data[0][-i+1][0].length)
        depth.remove(self.data[1][1][len(self.data[1][0])-1])

        #Calcul de la profondeur au milieu de l element#
        for i in range(len(self.data[0])):
            if type(self.data[0][i]) != list:
                depth[i] = depth[i]-cos(alpha[i]) * \
                    eval(self.data[0][i].length)/2
            else:
                depth[i] = depth[i]-cos(alpha[i]) * \
                    eval(self.data[0][i][0].length)/2

        return depth

    def resolve(self):
        """Resolution du comportement de la ligne de mouillage
        Output : F : Effort sur chaque element
                alpha : Inclinaison
                P : Poids en newton
        """
        alpha = []
        alpha_new = []
        Fx = []
        Fy = []
        F = []
        Tn = []
        Tt = []
        P = []
        new_depth = []
        rho = 1028
        cond = 1
        fact_convergence = 1
        it = 0
        g = 9.81
        for i in range(len(self.depth)):
            alpha.append(0)
            alpha_new.append(0)
            Fx.append(0)
            Fy.append(0)
            Tn.append(0)
            Tt.append(0)
            F.append(0)
            new_depth.append(self.depth[i])  # Vecteur initialise precedemment
            #Calculate Weight#
            if type(self.data[0][i]) != list:
                P.append(eval(self.data[0][i].mass)*g)
            else:
                inter = [eval(self.data[0][i][0].mass)*g,
                        eval(self.data[0][i][1].mass)*g]
                P.append(inter)

        while (cond == 1):  # Tant que le systeme nest pas stable
            it = it+1
            for i in range(len(alpha)):  # On met a jour la valeur de alpha
                alpha[i] = alpha_new[i]
            # On met a jour le vecteur profondeur
            new_depth = self.update_depth(new_depth, alpha)
            # On met a jour le vecteur de courant
            V = self.current_vector(new_depth)
            #On commence par la tete de mouillage#
            Tn[0] = 0.5*rho*eval(self.data[0][0].nl_drag_cf)*eval(self.data[0]
                                                                [0].projected_area)*(V[0]*cos(alpha[0]))**2  # Calcul trainee normal
            Tt[0] = 0.5*rho*eval(self.data[0][0].tl_drag_cf)*eval(self.data[0]
                                                                [0].projected_area)*(V[0]*sin(alpha[0]))**2  # Calcul trainee tangentielle
            Fx[0] = Tn[0]*cos(alpha[0])+Tt[0] * \
                sin(alpha[0])  # Resultante selon x
            Fy[0] = Tt[0]*cos(alpha[0])-Tn[0]*sin(alpha[0]) + \
                P[0]  # Resultante selon y
            alpha_new[0] = atan(Fx[0]/Fy[0])  # Angle
            F[0] = sqrt(Fx[0]**2+Fy[0]**2)  # Norme
            #On calcule element par element en partant du haut#
            for i in range(1, len(self.depth)-1):
                if type(self.data[0][i]) != list:
                    Tn[i] = 0.5*rho*eval(self.data[0][i].nl_drag_cf)*eval(
                        self.data[0][i].projected_area)*(V[i]*cos(alpha[i]))**2  # Calcul trainee normal
                    Tt[i] = 0.5*rho*eval(self.data[0][i].tl_drag_cf)*eval(self.data[0][i].projected_area)*(
                        V[i]*sin(alpha[i]))**2  # Calcul trainee tangentielle
                    Fx[i] = Tn[i]*cos(alpha[i])+Tt[i] * \
                        sin(alpha[i])+Fx[i-1]  # Resultante selon x
                    Fy[i] = Tt[i]*cos(alpha[i])-Tn[i]*sin(alpha[i]) + \
                        P[i]+Fy[i-1]  # Resultante selon y
                    alpha_new[i] = fact_convergence*atan(Fx[i]/Fy[i])  # Angle
                    F[i] = sqrt(Fx[i]**2+Fy[i]**2)  # Norme
                #ELement clampe#
                else:

                    Tn[i] = 0.5*rho*(V[i]*cos(alpha[i])**2)*(eval(self.data[0][i][0].nl_drag_cf)*eval(self.data[0][i][0].projected_area) +  # Calcul trainee normal
                                                            eval(self.data[0][i][1].nl_drag_cf)*eval(self.data[0][i][1].projected_area))

                    Tt[i] = 0.5*rho*(V[i]*sin(alpha[i])**2)*(eval(self.data[0][i][0].tl_drag_cf)*eval(self.data[0][i][0].projected_area) +  # Calcul trainee tangentielle
                                                            eval(self.data[0][i][1].tl_drag_cf)*eval(self.data[0][i][1].projected_area))

                    Fx[i] = Tn[i]*cos(alpha[i])+Tt[i] * \
                        sin(alpha[i])+Fx[i-1]  # Resultante selon x
                    Fy[i] = Tt[i]*cos(alpha[i])-Tn[i]*sin(alpha[i]) + \
                        P[i][0]+P[i][1]+Fy[i-1]  # Resultante selon y
                    alpha_new[i] = fact_convergence*atan(Fx[i]/Fy[i])  # Angle
                    F[i] = sqrt(Fx[i]**2+Fy[i]**2)  # Norme

            cond = 0

            for i in range(len(self.depth)):  # Condition d arret
                if abs(alpha[i]-alpha_new[i]) > 0.0001:
                    cond = 1
                    break

        return alpha, F, P

    def calculate_stretch(self, F):
        """Cette fonction permet de calculer l allongement dun cable en fonction de leffort qui lui est applique
        Input : F : Effort subit par le cable """
        new_stretch = 0
        self.percent_stretch = []
        for i in range(len(self.data[0])):
            new_stretch = 0
            if type(self.data[0][i]) != list:
                if self.data[0][i].__class__.__name__ == "Ropes":
                    # On calcule  le pourcentage de la charge de rupture qui est applique au cable
                    ratio = 100/9.81 * \
                        (F[i]/eval(self.data[0][i].breaking_strength))
                    new_stretch = (round(eval(self.data[0][i].poly_1), 2)*ratio**2+round(eval(
                        self.data[0][i].poly_2), 2)*ratio+round(eval(self.data[0][i].poly_3), 2))/100
                    # En fonction des 3 coefficients polynomiales entre pour le materiau on calcul lallongement : allongement (%) = coeff1*(%Charge de rupture)^2+coeff2*(%Charge de rupture)+coeff3
                    # On met a jour la nouvelle longueur
                    self.data[0][i].length = str(
                        eval(self.data[0][i].length)*(1+new_stretch))
            #Element clampe#
            else:
                if self.data[0][i][0].__class__.__name__ == "Ropes":
                    ratio = 100/9.81 * \
                        (F[i]/eval(self.data[0][i][0].breaking_strength))
                    new_stretch = (round(eval(self.data[0][i][0].poly_1), 2)*ratio**2+round(eval(
                        self.data[0][i][0].poly_2), 2)*ratio+round(eval(self.data[0][i][0].poly_3), 2))/100
                    self.data[0][i][0].length = str(
                        eval(self.data[0][i][0].length)*(1+new_stretch))
            self.percent_stretch.append(100*new_stretch)

    def calculate_ratio(self, alpha):
        """Cette fonction permet de calculer le ratio de clampage definie par lutilisateur comme automatique
        Input : alpha : vecteur angle d inclinaison"""

        #Calcul des coordonnees x et y en fonction de l inclinaison#
        y_float = [self.data[1][1][len(self.data[1][0])-1]]
        for i in range(len(self.data[2])):
            if type(self.data[2][i]) != list:
                if self.data[2][i] > 0:
                    self.data[2][i] = self.data[2][i]+y_float[-1]
            else:
                if self.data[2][i][1] > 0:
                    self.data[2][i][1] = self.data[2][i][1]+y_float[-1]

        x_float = [0]
        for i in range(2, len(self.data[0])+2):
            y_float.insert(0, 0)
            x_float.insert(0, 0)
            if type(self.data[0][-i+1]) != list:
                y_float[-i] = y_float[-i+1] + \
                    cos(alpha[-i+1])*eval(self.data[0][-i+1].length)
                x_float[-i] = x_float[-i+1] + \
                    sin(alpha[-i+1])*eval(self.data[0][-i+1].length)
            else:
                y_float[-i] = y_float[-i+1] + \
                    cos(alpha[-i+1])*eval(self.data[0][-i+1][0].length)
                x_float[-i] = x_float[-i+1] + \
                    sin(alpha[-i+1])*eval(self.data[0][-i+1][0].length)
        y_float.remove(self.data[1][1][len(self.data[1][0])-1])
        x_float.remove(x_float[-1])
        self.myxfloat = x_float
        self.myyfloat = y_float
        # Coordonnee x de la bouee de tete
        self.x_instru = [x_float[0] -
                        sin(alpha[0])*eval(self.data[0][0].length)/2]
        # Coordonnee y de la bouee de tete
        self.y_instru = [y_float[0] -
                        cos(alpha[0])*eval(self.data[0][0].length)/2]
        self.x_instru_top = []
        self.y_instru_top = []
        self.name_instru = [self.data[0][0].name]  # Nom de la bouee de tete
        for i in range(len(self.data[6])):
            if self.data[6][i] == 0.0:  # Ratio defini comme auto
                calc = 0
                # Si la profondeur choisi par lutilisateur est plus haut que le support de clampage
                if self.data[2][self.find_match(i)][1] > self.myyfloat[self.find_match(i)]:
                    # le ratio est fixe au minimum c est a dire 0
                    self.data[6][i] = 0
                    calc = 1
                # Si la profondeur choisi par lutilisateur est plus basse que le support de clampage
                if self.data[2][self.find_match(i)][1] < self.myyfloat[self.find_match(i+1)]:
                    # le ratio est fixe au minimum c est a dire 1
                    self.data[6][i] = 1
                    calc = 1
                if calc == 0:
                    malongueur = eval(
                        self.data[0][self.find_match(i)][0].length)
                    lon = 0
                    p = 1
                    for j in range(int(self.table_match[i])+1, int(self.table_match[i+1])):
                        malongueur = malongueur+eval(self.data[0][j].length)
                    while(self.myyfloat[self.find_match(i)+p] > self.data[2][self.find_match(i)][1]):
                        if type(self.data[0][self.find_match(i)+p-1]) != list:
                            lon = lon + \
                                eval(self.data[0]
                                    [self.find_match(i)+p-1].length)
                        else:
                            lon = lon + \
                                eval(
                                    self.data[0][self.find_match(i)+p-1][0].length)
                        p = p+1
                    ratio = ((self.myyfloat[self.find_match(i)+p-1]-self.data[2][self.find_match(i)][1])/(cos(alpha[self.find_match(i)+p-1]))
                            + 0.5*eval(self.data[0][self.find_match(i)][1].length)+lon)/malongueur
                    self.data[6][i] = ratio

                if self.data[6][i] >= 1 or self.data[6][i] <= 0:
                    QtGui.QMessageBox.warning(
                        self, 'Message', "An instrument is clamped beyond it support")

        #Calcul de la profondeur des instruments#
        for i in range(len(self.data[6])):
            if type(self.data[0][self.find_match(i)]) != list:
                if (self.data[0][self.find_match(i)]).__class__.__name__ == "Instruments":
                    self.name_instru.append(
                        self.data[0][self.find_match(i)].name)
                    self.x_instru.append(self.myxfloat[self.find_match(
                        i)]-0.5*eval(self.data[0][self.find_match(i)].length)*sin(alpha[self.find_match(i)]))
                    self.y_instru.append(self.myyfloat[self.find_match(
                        i)]-0.5*eval(self.data[0][self.find_match(i)].length)*cos(alpha[self.find_match(i)]))
                    # coordonnee x en haut de l instrument
                    self.x_instru_top.append(self.myxfloat[self.find_match(i)])
                    # coordonnee y en haut de l instrument
                    self.y_instru_top.append(self.myyfloat[self.find_match(i)])
        #Element clampe#
            else:
                if (self.data[0][self.find_match(i)][1]).__class__.__name__ == "Instruments":
                    self.name_instru.append(
                        self.data[0][self.find_match(i)][1].name)
                    malongueur = eval(
                        self.data[0][self.find_match(i)][0].length)
                    lon = eval(self.data[0][self.find_match(i)][0].length)
                    lon2 = lon
                    p = 0
                    ratio = self.data[6][i]
                    for j in range(int(self.table_match[i])+1, int(self.table_match[i+1])):
                        malongueur = malongueur+eval(self.data[0][j].length)
                    while(lon < malongueur*self.data[6][i]):
                        ratio = ratio-lon2/malongueur
                        p = p+1
                        lon = lon + \
                            eval(self.data[0]
                                [int(self.table_match[i])+p].length)
                        lon2 = eval(
                            self.data[0][int(self.table_match[i])+p].length)

                    self.x_instru.append(self.myxfloat[self.find_match(
                        i)+p]-ratio*malongueur*sin(alpha[self.find_match(i)+p]))
                    self.y_instru.append(self.myyfloat[self.find_match(
                        i)+p]-ratio*malongueur*cos(alpha[self.find_match(i)+p]))
                    self.x_instru_top.append(self.myxfloat[self.find_match(i)+p]-(ratio*malongueur-0.5*eval(
                        self.data[0][self.find_match(i)][1].length))*sin(alpha[self.find_match(i)+p]))  # coordonnee x en haut de l instrument
                    self.y_instru_top.append(self.myyfloat[self.find_match(i)+p]-(ratio*malongueur-0.5*eval(
                        self.data[0][self.find_match(i)][1].length))*cos(alpha[self.find_match(i)+p]))  # coordonnee y en haut de l instrument

    def calculate_subduction(self):
        """Cette fonction permet de calculer la subduction entrainee par le courant et les allongements
        Output : vecteur subduct """
        subduct = []
        for i in range(len(self.static_depth)):
            subduct.append(self.static_depth[i]-self.myyfloat[i])
        return subduct

    def find_match(self, ind):
        """Cette fonction utilise la table de correspondance
        Input : Indice sans troncon de cable
        Output : Indice avec troncon"""
        return int(self.table_match[ind])

    def resolve_backup(self, P):
        """Cette fonction permet de resoudre le backup jusqu au largueurs
        Input : poids en newton
        Output : backup,zero de flottabilite"""
        result2 = []
        zero = []
        for i in range(len(self.data[0])):
            if(self.data[0][i].__class__.__name__) == "Releases":
                self.release_ind = i+1

        for i in range(self.release_ind):
            result = 0
            for j in range(i, self.release_ind):
                if type(P[j]) != list:
                    result = result+P[j]
                else:
                    result = result+P[j][0]+P[j][1]
            result2.append(result)
        for i in range(len(result2)):
            zero.append(0)
        return result2, zero

    def plot_figure(self, alpha, P, F):
        """Cette fonction permet de tracer les quatres graphes
        Input : alpha : angle d inclinaison
                : P : Poids
                F : Effort
                """

        if hasattr(self, "plot_mooring"):
            self.plot_mooring.hide()

        backup, zero = self.resolve_backup(P)

        for i in range(len(self.T_max)):
            self.old_depth.append(self.myyfloat[self.find_match(i)])

        self.groupLayout = QtGui.QVBoxLayout()
        self.scrollArea = QtGui.QScrollArea()
        self.groupLayout.addWidget(self.scrollArea)
        self.scrolledWidget = QtGui.QWidget()  # Zone de defilement verticale
        self.scrolledWidget.setMinimumWidth(self.data[4])
        self.layout = QtGui.QVBoxLayout(self.scrolledWidget)

        self.widget_graph = QtGui.QWidget(self)

        sc1 = MySimulateCanvas(self.myxfloat, self.myyfloat, backup, self.T_max, self.T,
                            self.old_depth, self.V_chute_t, self.release_ind, self.name_instru,
                            self.x_instru, self.y_instru, self.data[5], self.widget_graph)
        self.layout.addWidget(sc1)
        min_value_anchor = F[-2]/9.81  # Valeur minimale lest
        # Valeur safe lest formule visbeck runmoor
        safe_value_anchor = (1.5/9.81) * \
            (F[-2]*(cos(alpha[-2])+0.6*sin(alpha[-2])))
        self.anchor_value = [round(safe_value_anchor, 1), -1*round(
            float(self.data[0][-1].mass), 1), round(max(self.percent_max_T), 1)]
        #Affichage valeurs lest#
        self.widget_anchors = QtGui.QWidget(self)
        self.anchors_layout = QtGui.QGridLayout()
        label1 = QtGui.QLabel("Anchors's WET weight (kg)")
        label1.setStyleSheet(
            "QLabel { background-color : white; color : black;border: 2px solid black }")
        label2 = QtGui.QLabel("Min")
        label2.setStyleSheet(
            "QLabel { background-color : white; color : black;border: 2px solid black }")
        label3 = QtGui.QLabel("Max")
        label3.setStyleSheet(
            "QLabel { background-color : white; color : black;border: 2px solid black }")
        label4 = QtGui.QLabel(str(round(min_value_anchor, 1)))
        label4.setStyleSheet(
            "QLabel { background-color : white; color : black;border: 1px solid black }")
        label5 = QtGui.QLabel(str(round(self.max_value_anchor, 1)))
        label5.setStyleSheet(
            "QLabel { background-color : white; color : black;border: 1px solid black }")
        label6 = QtGui.QLabel("Safe Anchor's weight")
        label6.setStyleSheet(
            "QLabel { background-color : white; color : black;border: 2px solid black }")
        label7 = QtGui.QLabel("Selected Anchor's weight")
        label7.setStyleSheet(
            "QLabel { background-color : white; color : black;border: 2px solid black }")
        label8 = QtGui.QLabel(str(self.anchor_value[0]))
        label8.setStyleSheet(
            "QLabel { background-color : white; color : black;border: 1px solid black }")
        label9 = QtGui.QLabel(str(self.anchor_value[1]))
        # Valeur choisie < valeur recommande
        if self.anchor_value[1] < self.anchor_value[0]:
            label9.setStyleSheet(
                "QLabel { background-color : red; color : white ;border: 2px solid black }")
        else:
            label9.setStyleSheet(
                "QLabel { background-color : green; color : white;border: 1px solid black }")

        label10 = QtGui.QLabel("Max Launch Tension / Breaking Strength (%)")
        label10.setStyleSheet(
            "QLabel { background-color : white; color : black;border: 2px solid black }")
        label11 = QtGui.QLabel(str(self.anchor_value[2]))
        if round(max(self.percent_max_T), 1) > 40:  # T_max/Ultimate_load >40%
            label11.setStyleSheet(
                "QLabel { background-color : red; color : white ;border: 2px solid black }")
        else:
            label11.setStyleSheet(
                "QLabel { background-color : green; color : white;border: 1px solid black }")

        self.anchors_layout.addWidget(label2, 0, 1)
        self.anchors_layout.addWidget(label3, 0, 3)
        self.anchors_layout.addWidget(label1, 1, 0)
        self.anchors_layout.addWidget(label4, 1, 1)
        self.anchors_layout.addWidget(label5, 1, 3)
        self.anchors_layout.addWidget(label6, 0, 2)
        self.anchors_layout.addWidget(label7, 0, 4)
        self.anchors_layout.addWidget(label8, 1, 2)
        self.anchors_layout.addWidget(label9, 1, 4)
        self.anchors_layout.addWidget(label10, 2, 0)
        self.anchors_layout.addWidget(label11, 2, 1)

        self.widget_anchors.setLayout(self.anchors_layout)
        self.layout.addWidget(self.widget_anchors)

    def make_table(self, alpha, F):
        """Cette fonction permet de creer le tableau qui contient toutes les informations de la simulation
        Input : alpha : Vecteur angle d inclinaison
                F : Vecteur effort """
        name = []
        length = []
        length_stretch = []
        Tab_final = []
        self.Tab_report = []
        ind_instrum = []
        my_prof = []
        F_Kp = []
        warn = []
        count = 0
        release_ind = 0
        #Reconstitue les troncons en un seul cable#
        for i in range(len(self.table_match)):
            for j in range(i, int(self.table_match[i])-count):
                if type(self.data[0][i-1]) != list:
                    self.data[0][i-1].length = str(
                        eval(self.data[0][i-1].length)+eval(self.data[0][i].length))
                    self.data[0][i-1].mass = str(eval(self.data[0]
                                                    [i-1].mass)+eval(self.data[0][i].mass))
                    self.data[0][i-1].projected_area = str(
                        eval(self.data[0][i-1].projected_area)+eval(self.data[0][i].projected_area))
                else:
                    self.data[0][i-1][0].length = str(
                        eval(self.data[0][i-1][0].length)+eval(self.data[0][i].length))
                    self.data[0][i-1][0].mass = str(
                        eval(self.data[0][i-1][0].mass)+eval(self.data[0][i].mass))
                    self.data[0][i-1][0].projected_area = str(
                        eval(self.data[0][i-1][0].projected_area)+eval(self.data[0][i].projected_area))
                del(self.data[0][i])
                del(alpha[i])
                del(F[i])
                del(self.myxfloat[i])
                del(self.myyfloat[i])
                del(self.percent_stretch[i])
                count = count+1

        size = len(self.data[0])-1
        for i in range(len(alpha)):
            alpha[i] = alpha[i]*(180/(3.14))  # Converti radian en degre

        subduct = self.calculate_subduction()  # Calcule de la subduction

        for p in self.data[0]:
            if type(p) != list:
                name.append(p.name)
            else:
                # Recuperation des noms
                name.append(p[0].name+' '+'+'+' '+p[1].name)

        count = 0
        for i in range(len(self.data[0])):

            F_Kp.append(F[i]/9.81)  # Convertion des efforts en Kp
            if type(self.data[0][i]) != list:
                length.append(self.original_length[i])
                my_prof.append(str(round(self.myyfloat[i], 1)))
                # Detection des instruments
                if self.data[0][i].__class__.__name__ == "Instruments":
                    ind_instrum.append(i)
                    count = count+1
                if self.data[0][i].__class__.__name__ == "Releases":  # Detection du release
                    release_ind = i

                if self.data[0][i].__class__.__name__ == "Ropes":
                    length_stretch.append(str(round(eval(self.data[0][i].length), 1))+' '+'('+str(
                        round(self.percent_stretch[i], 1))+'%'+')')  # Recuperation des longueurs allongees
                else:
                    length_stretch.append('')

            #Element clampe#
            else:
                length.append(self.original_length[i][0])
                if self.data[0][i][1].__class__.__name__ == "Instruments":
                    ind_instrum.append(i)
                    my_prof.append(str(round(self.myyfloat[i], 1))+' '+'+'+' '+str(round(self.y_instru_top[count], 1))+' '+'ratio = '+str(
                        round(self.data[6][i], 2)))  # Affichage profondeur support+profondeur element clampe+clamp_ratio
                    if self.data[6][i] >= 1 or self.data[6][i] <= 0:
                        warn.append(i)
                    count = count+1
                else:
                    my_prof.append(str(round(self.myyfloat[i], 1)))
                if self.data[0][i].__class__.__name__ == "Ropes":
                    length_stretch.append(
                        self.data[0][i][0].length+' '+'('+self.percent_stretch[i]+')')
                else:
                    length_stretch.append('')

        Tab = np.zeros((size+1, 8))
        Tab[:, 0] = length
        Tab[:, 1] = F_Kp
        Tab[:, 2] = alpha
        Tab[:, 3] = self.myxfloat
        Tab[:, 4] = subduct
        Tab[:, 5] = self.Buoyancy_kg
        Tab[:, 6] = self.weight_kg
        Tab[:, 7] = self.percent_max_T

        for i in range(Tab.shape[0]):
            inter = [name[i]]  # On commence par les noms
            for j in range(1):
                inter.append(Tab[i][j])  # Longueur
            inter.append(my_prof[i])  # Profondeur
            for j in range(1, Tab.shape[1]):  # Reste du tableau
                inter.append(Tab[i][j])
            inter.append(length_stretch[i])
            Tab_final.append(inter)
        Tab_final.insert(0, ["Name", "Static Length (m)", "Depth (m)", "Tension (Kp)", "Angle (deg)",
                            "Dx (m)", "Dz (m)", "Buoy (kg)", "Weight (kg)", "Launch Tension (%)", "Length Stretched (m)"])
        tab_grid = QtGui.QGridLayout()
        inter = []
        for j in range(len(Tab_final[0])):  # Affichage du nom des categories
            label = QtGui.QLabel(Tab_final[0][j])
            label.setStyleSheet(
                "QLabel { background-color : white; color : black;border: 2px solid black }")
            tab_grid.addWidget(label, 0, j)
            inter.append(Tab_final[0][j])

        self.Tab_report.append(inter)
        for k in range(1, len(Tab_final)):
            # Affichage de chaque element sauf terminaison
            if self.data[0][k-1].__class__.__name__ != "Terminals":
                inter = []
                for l in range(len(Tab_final[k])):
                    if type(Tab_final[k][l]) != str:
                        Tab_final[k][l] = str(round(Tab_final[k][l], 1))
                    inter.append(Tab_final[k][l])
                    label = QtGui.QLabel(Tab_final[k][l])
                    # On repere en bleu ciel : la bouee de tete, les instruments, les largueurs, et le lest
                    if k-1 in ind_instrum or k == 1 or k-1 == release_ind or k == len(Tab_final)-1:
                        label.setStyleSheet(
                            "QLabel { background-color : rgb(168, 211, 255); color : black;border: 1px solid black }")
                        if k-1 in warn:  # Un instrument clampe au dela de son support
                            label.setStyleSheet(
                                "QLabel { background-color : red; color : black;border: 1px solid black }")

                    else:
                        label.setStyleSheet(
                            "QLabel { background-color : white; color : black;border: 1px solid black }")
                    tab_grid.addWidget(label, k, l)
                self.Tab_report.append(inter)
        tab_grid.setSpacing(0)

        self.layout.addLayout(tab_grid)

        self.scrollArea.setWidget(self.scrolledWidget)
        self.setLayout(self.groupLayout)


import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from mmcv.cnn import kaiming_init
from mmdet.ops import DeformConv

class NovelKQRAttention(nn.Module):
    """Modified GeneralizedAttention module.

    See 'An Empirical Study of Spatial Attention Mechanisms in Deep Networks'
    (https://arxiv.org/abs/1711.07971) for details.

    Args:
        in_dim (int): Channels of the input feature map.
        spatial_range (int): The spatial range.
            -1 indicates no spatial range constraint.
        num_heads (int): The head number of empirical_attention module.
        position_embedding_dim (int): The position embedding dimension.
        position_magnitude (int): A multiplier acting on coord difference.
        kv_stride (int): The feature stride acting on key/value feature map.
        q_stride (int): The feature stride acting on query feature map.
        attention_type (str): A binary indicator string for indicating which
            items in generalized empirical_attention module are used.
            '1000' indicates 'query and key content' (appr - appr) item,
            '0100' indicates 'query content and relative position'
              (appr - position) item,
            '0010' indicates 'key content only' (bias - appr) item,
            '0001' indicates 'relative position only' (bias - position) item.
    """

    def __init__(self,
                 in_dim,
                 spatial_range=-1,
                 num_heads=9,
                 position_embedding_dim=-1,
                 position_magnitude=1,
                 kv_stride=2,
                 q_stride=1,
                 attention_type='0110',
                 deformable_group=1,
                 dconv_stride = 1,
                 dconv_groups =1,
                 dconv_learnable_vector = True):

        super(NovelKQRAttention, self).__init__()

        # hard range means local range for non-local operation
        self.position_embedding_dim = (
            position_embedding_dim if position_embedding_dim > 0 else in_dim)

        self.position_magnitude = position_magnitude
        self.num_heads = num_heads
        self.channel_in = in_dim
        self.spatial_range = spatial_range
        self.kv_stride = kv_stride
        self.q_stride = q_stride
        self.attention_type = [bool(int(_)) for _ in attention_type]
        self.qk_embed_dim = in_dim // num_heads
        self.deformable_group = deformable_group
        self.dconv_stride = dconv_stride
        self.dconv_learnable_vector = dconv_learnable_vector
        self.dconv_groups = dconv_groups
        out_c = self.qk_embed_dim * num_heads

        if self.attention_type[2]:
            self.key_conv = nn.Conv2d(
                in_channels=in_dim,
                out_channels=out_c,
                kernel_size=1,
                bias=False)
            self.key_conv.kaiming_init = True

        if self.attention_type[1]:
            self.query_conv_offset = nn.Conv2d(
                in_channels=in_dim,
                out_channels=self.deformable_group * 18,
                kernel_size=3,
                padding=1,
                stride = self.dconv_stride,
                dilation=1,
                bias=False)

            self.query_dconv=DeformConv(
                in_channels =in_dim,
                out_channels = out_c,
                kernel_size=3,
                stride = self.dconv_stride,
                padding = 1,
                dilation =1,
                deformable_groups = self.deformable_group,
                groups= self.dconv_groups,
                bias = False)

        self.v_dim = in_dim // num_heads
        self.value_conv = nn.Conv2d(
            in_channels=in_dim,
            out_channels=self.v_dim * num_heads,
            kernel_size=1,
            bias=False)
        self.value_conv.kaiming_init = True

        if self.attention_type[2]:
            stdv = 1.0 / math.sqrt(self.qk_embed_dim * 2)
            appr_bias_value = -2 * stdv * torch.rand(out_c) + stdv
            self.appr_bias = nn.Parameter(appr_bias_value)

        if self.attention_type[1] and self.dconv_learnable_vector == True :
            stdv = 1.0 / math.sqrt(self.qk_embed_dim * 2)
            appr_bias_qRelPos = -2 * stdv * torch.rand(out_c) + stdv
            self.appr_bias_qRelPos = nn.Parameter(appr_bias_qRelPos)

        self.proj_conv = nn.Conv2d(
            in_channels=self.v_dim * num_heads,
            out_channels=in_dim,
            kernel_size=1,
            bias=True)
        self.proj_conv.kaiming_init = True
        self.gamma = nn.Parameter(torch.zeros(1))

        if self.spatial_range >= 0:
            # only works when non local is after 3*3 conv
            if in_dim == 256:
                max_len = 84
            elif in_dim == 512:
                max_len = 42

            max_len_kv = int((max_len - 1.0) / self.kv_stride + 1)
            local_constraint_map = np.ones(
                (max_len, max_len, max_len_kv, max_len_kv), dtype=np.int)
            for iy in range(max_len):
                for ix in range(max_len):
                    local_constraint_map[iy, ix,
                                         max((iy - self.spatial_range) //
                                             self.kv_stride, 0):min(
                                                 (iy + self.spatial_range +
                                                  1) // self.kv_stride +
                                                 1, max_len),
                                         max((ix - self.spatial_range) //
                                             self.kv_stride, 0):min(
                                                 (ix + self.spatial_range +
                                                  1) // self.kv_stride +
                                                 1, max_len)] = 0

            self.local_constraint_map = nn.Parameter(
                torch.from_numpy(local_constraint_map).byte(),
                requires_grad=False)

        if self.q_stride > 1:
            self.q_downsample = nn.AvgPool2d(
                kernel_size=1, stride=self.q_stride)
        else:
            self.q_downsample = None

        if self.kv_stride > 1:
            self.kv_downsample = nn.AvgPool2d(
                kernel_size=1, stride=self.kv_stride)
        else:
            self.kv_downsample = None

        self.init_weights()

    def forward(self, x_input):
        num_heads = self.num_heads

        # use empirical_attention
        if self.q_downsample is not None:
            x_q = self.q_downsample(x_input)
        else:
            x_q = x_input
        n, _, h, w = x_q.shape

        if self.dconv_stride > 1:
            h,w = h//self.dconv_stride, w//self.dconv_stride

        if self.kv_downsample is not None:
            x_kv = self.kv_downsample(x_input)
        else:
            x_kv = x_input
        _, _, h_kv, w_kv = x_kv.shape

        if self.attention_type[2]:
            proj_key = self.key_conv(x_kv).view(
                (n, num_heads, self.qk_embed_dim, h_kv * w_kv))

        if self.attention_type[1]:
            offset = self.query_conv_offset(x_q)
            proj_query_relativePos =self.query_dconv(x_q,offset).view(n,num_heads,self.qk_embed_dim,h*w)

        # accelerate for saliency only
        if (np.sum(self.attention_type) == 1) and self.attention_type[2]:
            appr_bias = self.appr_bias.\
                view(1, num_heads, 1, self.qk_embed_dim).\
                repeat(n, 1, 1, 1)

            energy = torch.matmul(appr_bias, proj_key).\
                view(n, num_heads, 1, h_kv * w_kv)

            h = 1
            w = 1
        else:
            if not self.attention_type[0]:
                energy = torch.zeros(
                    n,
                    num_heads,
                    h,
                    w,
                    h_kv,
                    w_kv,
                    dtype=x_input.dtype,
                    device=x_input.device)

            # Key Content
            if self.attention_type[2]:
                appr_bias = self.appr_bias.\
                    view(1, num_heads, 1, self.qk_embed_dim).\
                    repeat(n, 1, 1, 1)

                energy += torch.matmul(appr_bias, proj_key).\
                        view(n, num_heads, 1, 1, h_kv, w_kv)

            # Query Content and Relative Position
            if self.attention_type[1]:

                if self.dconv_learnable_vector == False :
                    energy+=proj_query_relativePos.view(n,num_heads,h,w,1,1)

                else:
                    appr_bias_qRelPos = self.appr_bias_qRelPos.\
                        view(1,num_heads,1,self.qk_embed_dim).\
                        repeat(n,1,1,1)
                    energy+= torch.matmul(appr_bias_qRelPos,proj_query_relativePos).\
                            view(n,num_heads,h,w,1,1)

            energy = energy.view(n,num_heads,h*w,h_kv*w_kv)

        if self.spatial_range >= 0:
            cur_local_constraint_map = \
                self.local_constraint_map[:h, :w, :h_kv, :w_kv].\
                contiguous().\
                view(1, 1, h*w, h_kv*w_kv)

            energy = energy.masked_fill_(cur_local_constraint_map,
                                         float('-inf'))

        attention = F.softmax(energy, 3)

        proj_value = self.value_conv(x_kv)
        proj_value_reshape = proj_value.\
            view((n, num_heads, self.v_dim, h_kv * w_kv)).\
            permute(0, 1, 3, 2)

        out = torch.matmul(attention, proj_value_reshape).\
            permute(0, 1, 3, 2).\
            contiguous().\
            view(n, self.v_dim * self.num_heads, h, w)

        out = self.proj_conv(out)

        if self.dconv_stride > 1:
            out = F.interpolate(out,scale_factor=self.dconv_stride)

        if self.q_stride > 1:
            out = F.interpolate(out,size=x_input.size()[2:])

        out = self.gamma * out + x_input

        return out

    def init_weights(self):
        for m in self.modules():
            if hasattr(m, 'kaiming_init') and m.kaiming_init:
                kaiming_init(
                    m,
                    mode='fan_in',
                    nonlinearity='leaky_relu',
                    bias=0,
                    distribution='uniform',
                    a=1)


import logging

import numpy as np
import pandas as pd
from numba import njit
from scipy.special import expn
from scipy.interpolate import PchipInterpolator
from collections import Counter as counter
from tardis import constants as const

from tardis.plasma.properties.base import (
    ProcessingPlasmaProperty,
    HiddenPlasmaProperty,
    BaseAtomicDataProperty,
)
from tardis.plasma.exceptions import IncompleteAtomicData
from tardis.plasma.properties.continuum_processes import (
    get_ground_state_multi_index,
    K_B,
    BETA_COLL,
    H,
    A0,
    M_E,
    C,
)

logger = logging.getLogger(__name__)

__all__ = [
    "Levels",
    "Lines",
    "LinesLowerLevelIndex",
    "LinesUpperLevelIndex",
    "AtomicMass",
    "IonizationData",
    "ZetaData",
    "NLTEData",
    "MacroAtomData",
    "PhotoIonizationData",
    "YgData",
    "YgInterpolator",
    "LevelIdxs2LineIdx",
    "LevelIdxs2TransitionIdx",
    "TwoPhotonData",
    "ContinuumInteractionHandler",
]

class Levels(BaseAtomicDataProperty):
    """
    Attributes
    ----------
    levels : pandas.MultiIndex
        (atomic_number, ion_number, level_number)
        Index of filtered atomic data. Index used for all other attribute dataframes for this class
    excitation_energy : pandas.DataFrame, dtype float
        Excitation energies of atomic levels.
        Index is levels.
    metastability : pandas.DataFrame, dtype bool
        Records whether atomic levels are metastable.
        Index is levels.
    g : pandas.DataFrame (index=levels), dtype float
        Statistical weights of atomic levels.
    """

    outputs = ("levels", "excitation_energy", "metastability", "g")
    latex_name = (
        r"\textrm{levels}",
        r"\epsilon_{\textrm{k}}",
        r"\textrm{metastability}",
        "g",
    )

    def _filter_atomic_property(self, levels, selected_atoms):
        return levels
        # return levels[levels.atomic_number.isin(selected_atoms)]

    def _set_index(self, levels):
        # levels = levels.set_index(['atomic_number', 'ion_number',
        #                          'level_number'])
        return (
            levels.index,
            levels["energy"],
            levels["metastable"],
            levels["g"],
        )

class Lines(BaseAtomicDataProperty):
    """
    Attributes
    ----------
    lines : pandas.DataFrame
        Atomic lines data. Columns are wavelength, atomic_number,ion_number,
        f_ul, f_lu, level_number_lower, level_number_upper, nu, B_lu, B_ul, A_ul,
        wavelength. Index is line_id.
    nu : pandas.DataFrame, dtype float
        Line frequency data. Index is line_id.
    f_lu : pandas.DataFrame, dtype float
        Transition probability data. Index is line_id.
    wavelength_cm : pandas.DataFrame, dtype float
        Line wavelengths in cm. Index is line_id.
    """

    # Would like for lines to just be the line_id values
    outputs = ("lines", "nu", "f_lu", "wavelength_cm")

    latex_name = (
        r"\textrm{lines}",
        r"\nu",
        r"f_lu",
        r"\lambda_{cm}",
    )

    def _filter_atomic_property(self, lines, selected_atoms):
        # return lines[lines.atomic_number.isin(selected_atoms)]
        return lines

    def _set_index(self, lines):
        # lines.set_index('line_id', inplace=True)
        return lines, lines["nu"], lines["f_lu"], lines["wavelength_cm"]

class MacroAtomData(BaseAtomicDataProperty):
    outputs = ("macro_atom_data",)

    def _filter_atomic_property(self, macro_atom_data, selected_atoms):
        return macro_atom_data

    def _set_index(self, macro_atom_data):
        return macro_atom_data

class PhotoIonizationData(ProcessingPlasmaProperty):
    """
    Attributes
    ----------
    photo_ion_cross_sections : pandas.DataFrame, dtype float
        Photoionization cross sections as a function of frequency.
        Columns are nu, x_sect, index=('atomic_number','ion_number','level_number')
    photo_ion_block_references : numpy.ndarray, dtype int
        Indices where the photoionization data for
        a given level starts. Needed for calculation
        of recombination rates.
    nu_i : pandas.Series, dtype float
        Threshold frequencies for ionization
    energy_i : pandas.Series, dtype float
        Energies of levels with bound-free transitions. Needed to calculate
        for example internal transition probabilities in the macro atom scheme.
    photo_ion_index : pandas.MultiIndex, dtype int
        Atomic, ion and level numbers for which photoionization data exists.
    level2continuum_idx : pandas.Series, dtype int
        Maps a level MultiIndex (atomic_number, ion_number, level_number) to
        the continuum_idx of the corresponding bound-free continuum (which are
        sorted by decreasing frequency).
    level_idxs2continuum_idx : pandas.DataFrame, dtype int
        Maps a source_level_idx destination_level_idx pair to a continuum_idx.
    """

    outputs = (
        "photo_ion_cross_sections",
        "photo_ion_block_references",
        "photo_ion_index",
        "nu_i",
        "energy_i",
        "photo_ion_idx",
        "level2continuum_idx",
        "level_idxs2continuum_idx",
    )
    latex_name = (
        r"\xi_{\textrm{i}}(\nu)",
        "",
        "",
        r"\nu_i",
        r"\epsilon_i",
        "",
        "",
    )

    def calculate(self, atomic_data, continuum_interaction_species):
        #photoionization_data = atomic_data.photoionization_data.set_index(
        #    ["atomic_number", "ion_number", "level_number"]
        #)
        photoionization_data = atomic_data.photoionization_data
        mask_selected_species = photoionization_data.index.droplevel(
            "level_number"
        ).isin(continuum_interaction_species)
        photoionization_data = photoionization_data[mask_selected_species]
        phot_nus = photoionization_data["nu"]
        block_references = np.pad(
            phot_nus.groupby(level=[0, 1, 2]).count().values.cumsum(), [1, 0]
        )
        photo_ion_index = photoionization_data.index.unique()
        nu_i = photoionization_data.groupby(level=[0, 1, 2]).first().nu
        energy_i = atomic_data.levels.loc[photo_ion_index].energy

        source_idx = atomic_data.macro_atom_references.loc[
            photo_ion_index
        ].references_idx
        destination_idx = atomic_data.macro_atom_references.loc[
            get_ground_state_multi_index(photo_ion_index)
        ].references_idx
        photo_ion_idx = pd.DataFrame(
            {
                "source_level_idx": source_idx.values,
                "destination_level_idx": destination_idx.values,
            },
            index=photo_ion_index,
        )

        level2continuum_edge_idx = pd.Series(
            np.arange(len(nu_i)),
            nu_i.sort_values(ascending=False).index,
            name="continuum_idx",
        )

        level_idxs2continuum_idx = photo_ion_idx.copy()
        level_idxs2continuum_idx["continuum_idx"] = level2continuum_edge_idx
        level_idxs2continuum_idx = level_idxs2continuum_idx.set_index(
            ["source_level_idx", "destination_level_idx"]
        )
        return (
            photoionization_data,
            block_references,
            photo_ion_index,
            nu_i,
            energy_i,
            photo_ion_idx,
            level2continuum_edge_idx,
            level_idxs2continuum_idx,
        )

class ContinuumInteractionHandler(ProcessingPlasmaProperty):
    outputs = (
        "get_current_bound_free_continua",
        "determine_bf_macro_activation_idx",
        "determine_continuum_macro_activation_idx",
    )

    def calculate(
        self,
        photo_ion_cross_sections,
        level2continuum_idx,
        photo_ion_idx,
        k_packet_idx,
    ):
        nus = photo_ion_cross_sections.nu.loc[
            level2continuum_idx.index
        ]  # Sort by descending frequency
        nu_mins = nus.groupby(level=[0, 1, 2], sort=False).first().values
        nu_maxs = nus.groupby(level=[0, 1, 2], sort=False).last().values

        @njit(error_model="numpy", fastmath=True)
        def get_current_bound_free_continua(nu):
            """
            Determine bound-free continua for which absorption is possible.

            Parameters
            ----------
            nu : float
                Comoving frequency of the r-packet.

            Returns
            -------
            numpy.ndarray, dtype int
                Continuum ids for which absorption is possible for frequency `nu`.
            """
            # searchsorted would be faster but would need stricter format for photoionization data
            current_continua = np.where(
                np.logical_and(nu >= nu_mins, nu <= nu_maxs)
            )[0]
            return current_continua

        destination_level_idxs = photo_ion_idx.loc[
            level2continuum_idx.index, "destination_level_idx"
        ].values

        @njit(error_model="numpy", fastmath=True)
        def determine_bf_macro_activation_idx(
            nu, chi_bf_contributions, active_continua
        ):
            """
            Determine the macro atom activation level after bound-free absorption.

            Parameters
            ----------
            nu : float
                Comoving frequency of the r-packet.
            chi_bf_contributions : numpy.ndarray, dtype float
                Cumulative distribution of bound-free opacities at frequency
                `nu`.
            active_continua : numpy.ndarray, dtype int
                Continuum ids for which absorption is possible for frequency `nu`.

            Returns
            -------
            float
                Macro atom activation idx.
            """
            # Perform a MC experiment to determine the continuum for absorption
            index = np.searchsorted(chi_bf_contributions, np.random.random())
            continuum_id = active_continua[index]

            # Perform a MC experiment to determine whether thermal or
            # ionization energy is created
            nu_threshold = nu_mins[continuum_id]
            fraction_ionization = nu_threshold / nu
            if (
                np.random.random() < fraction_ionization
            ):  # Create ionization energy (i-packet)
                destination_level_idx = destination_level_idxs[continuum_id]
            else:  # Create thermal energy (k-packet)
                destination_level_idx = k_packet_idx
            return destination_level_idx

        @njit(error_model="numpy", fastmath=True)
        def determine_continuum_macro_activation_idx(
            nu, chi_bf, chi_ff, chi_bf_contributions, active_continua
        ):
            """
            Determine the macro atom activation level after a continuum absorption.

            Parameters
            ----------
            nu : float
                Comoving frequency of the r-packet.
            chi_bf : numpy.ndarray, dtype float
                Bound-free opacity.
            chi_bf : numpy.ndarray, dtype float
                Free-free opacity.
            chi_bf_contributions : numpy.ndarray, dtype float
                Cumulative distribution of bound-free opacities at frequency
                `nu`.
            active_continua : numpy.ndarray, dtype int
                Continuum ids for which absorption is possible for frequency `nu`.

            Returns
            -------
            float
                Macro atom activation idx.
            """
            fraction_bf = chi_bf / (chi_bf + chi_ff)
            # TODO: In principle, we can also decide here whether a Thomson
            # scattering event happens and need one less RNG call.
            if np.random.random() < fraction_bf:  # Bound-free absorption
                destination_level_idx = determine_bf_macro_activation_idx(
                    nu, chi_bf_contributions, active_continua
                )
            else:  # Free-free absorption (i.e. k-packet creation)
                destination_level_idx = k_packet_idx
            return destination_level_idx

        return (
            get_current_bound_free_continua,
            determine_bf_macro_activation_idx,
            determine_continuum_macro_activation_idx,
        )

class TwoPhotonData(ProcessingPlasmaProperty):
    outputs = ("two_photon_data", "two_photon_idx")
    """
    Attributes
    ----------
    two_photon_data : pandas.DataFrame, dtype float
    A DataFrame containing the *two photon decay data* with:
        index: atomic_number, ion_number, level_number_lower, level_number_upper
        columns: A_ul[1/s], nu0[Hz], alpha, beta, gamma
        alpha, beta, gamma are fit coefficients for the frequency dependent
        transition probability A(y) of the two photon decay. See Eq. 2 in
        Nussbaumer & Schmutz (1984).
    two_photon_idx : pandas.DataFrame, dtype int
    """

    def calculate(self, atomic_data, continuum_interaction_species):
        two_photon_data = atomic_data.two_photon_data
        mask_selected_species = two_photon_data.index.droplevel(
            ["level_number_lower", "level_number_upper"]
        ).isin(continuum_interaction_species)
        if not mask_selected_species.sum():
            raise IncompleteAtomicData(
                "two photon transition data for the requested "
                "continuum_interactions species: {}".format(
                    continuum_interaction_species.values.tolist()
                )
            )
        two_photon_data = two_photon_data[mask_selected_species]
        index_lower = two_photon_data.index.droplevel("level_number_upper")
        index_upper = two_photon_data.index.droplevel("level_number_lower")
        source_idx = atomic_data.macro_atom_references.loc[
            index_upper
        ].references_idx
        destination_idx = atomic_data.macro_atom_references.loc[
            index_lower
        ].references_idx
        two_photon_idx = pd.DataFrame(
            {
                "source_level_idx": source_idx.values,
                "destination_level_idx": destination_idx.values,
            },
            index=two_photon_data.index,
        )
        if len(two_photon_data) != 1:
            raise NotImplementedError(
                "Currently only one two-photon decay is supported but there "
                f"are {len(two_photon_data)} in the atomic data."
            )
        return two_photon_data, two_photon_idx

class LinesLowerLevelIndex(HiddenPlasmaProperty):
    """
    Attributes
    ----------
    lines_lower_level_index : numpy.ndrarray, dtype int
        Levels data for lower levels of particular lines
    """

    outputs = ("lines_lower_level_index",)

    def calculate(self, levels, lines):
        levels_index = pd.Series(
            np.arange(len(levels), dtype=np.int64), index=levels
        )
        lines_index = lines.index.droplevel("level_number_upper")
        return np.array(levels_index.loc[lines_index])

class LinesUpperLevelIndex(HiddenPlasmaProperty):
    """
    Attributes
    ----------
    lines_upper_level_index : numpy.ndarray, dtype int
        Levels data for upper levels of particular lines
    """

    outputs = ("lines_upper_level_index",)

    def calculate(self, levels, lines):
        levels_index = pd.Series(
            np.arange(len(levels), dtype=np.int64), index=levels
        )
        lines_index = lines.index.droplevel("level_number_lower")
        return np.array(levels_index.loc[lines_index])

class LevelIdxs2LineIdx(HiddenPlasmaProperty):
    """
    Attributes
    ----------
    level_idxs2line_idx : pandas.Series, dtype int
        Maps a source_level_idx destination_level_idx pair to a line_idx.
    """

    outputs = ("level_idxs2line_idx",)

    def calculate(self, atomic_data):
        index = pd.MultiIndex.from_arrays(
            [
                atomic_data.lines_upper2level_idx,
                atomic_data.lines_lower2level_idx,
            ],
            names=["source_level_idx", "destination_level_idx"],
        )
        level_idxs2line_idx = pd.Series(
            np.arange(len(index)), index=index, name="lines_idx"
        )
        return level_idxs2line_idx

class LevelIdxs2TransitionIdx(HiddenPlasmaProperty):
    """
    Attributes
    ----------
    level_idxs2transition_idx : pandas.DataFrame, dtype int
       Maps a source_level_idx destination_level_idx pair to a transition_idx
       and transition type.
    """

    outputs = ("level_idxs2transition_idx",)

    def calculate(self, level_idxs2line_idx, level_idxs2continuum_idx):
        level_idxs2line_idx = level_idxs2line_idx.to_frame()
        level_idxs2line_idx.insert(1, "transition_type", -1)

        level_idxs2continuum_idx = level_idxs2continuum_idx.copy()
        level_idxs2continuum_idx.insert(1, "transition_type", -2)
        level_idxs2continuum_idx = level_idxs2continuum_idx.rename(
            columns=({"continuum_idx": "lines_idx"})
        )

        names = level_idxs2continuum_idx.index.names
        level_idxs2continuum_idx = level_idxs2continuum_idx.swaplevel()
        level_idxs2continuum_idx.index.names = names

        # TODO: This should probably be defined somewhere else.
        # One possibility would be to attach it to the cooling properties as
        # a class attribute.
        index_cooling = pd.MultiIndex.from_product(
            [["k"], ["ff", "adiabatic", "bf"]], names=names
        )
        num_cool = len(index_cooling)
        level_idxs2cooling_idx = pd.DataFrame(
            {
                "lines_idx": np.ones(num_cool, dtype=int) * -1,
                "transition_type": np.arange(-3, -3 - num_cool, -1),
            },
            index=index_cooling,
        )
        level_idxs2transition_idx = pd.concat(
            [
                level_idxs2continuum_idx,
                level_idxs2line_idx,
                level_idxs2cooling_idx,
            ]
        )

        # TODO: Add two-photon processes

        return level_idxs2transition_idx

class AtomicMass(ProcessingPlasmaProperty):
    """
    Attributes
    ----------
    atomic_mass : pandas.Series
        Atomic masses of the elements used. Indexed by atomic number.
    """

    outputs = ("atomic_mass",)

    def calculate(self, atomic_data, selected_atoms):
        if getattr(self, self.outputs[0]) is not None:
            return (getattr(self, self.outputs[0]),)
        else:
            return atomic_data.atom_data.loc[selected_atoms].mass

class IonizationData(BaseAtomicDataProperty):
    """
    Attributes
    ----------
    ionization_data : pandas.Series
        Holding ionization energies
        Indexed by atomic number, ion number.
    """

    outputs = ("ionization_data",)

    def _filter_atomic_property(self, ionization_data, selected_atoms):
        mask = ionization_data.index.isin(selected_atoms, level="atomic_number")
        ionization_data = ionization_data[mask]
        counts = ionization_data.groupby(level="atomic_number").count()

        if np.alltrue(counts.index == counts):
            return ionization_data
        else:
            raise IncompleteAtomicData(
                f"ionization data for the ion ({str(counts.index[counts.index != counts])}, {str(counts[counts.index != counts])})"
            )

    def _set_index(self, ionization_data):
        return ionization_data

class ZetaData(BaseAtomicDataProperty):
    """
    Attributes
    ----------
    zeta_data : pandas.DataFrame, dtype float
        Zeta data for the elements used. Indexed by atomic number, ion number.
        Columns are temperature values up to 40,000 K in iterations of 2,000 K.
        The zeta value represents the fraction of recombination events
        from the ionized state that go directly to the ground state.
    """

    outputs = ("zeta_data",)

    def _filter_atomic_property(self, zeta_data, selected_atoms):
        zeta_data["atomic_number"] = zeta_data.index.codes[0] + 1
        zeta_data["ion_number"] = zeta_data.index.codes[1] + 1
        zeta_data = zeta_data[zeta_data.atomic_number.isin(selected_atoms)]
        zeta_data_check = counter(zeta_data.atomic_number.values)
        keys = np.array(list(zeta_data_check.keys()))
        values = np.array(zeta_data_check.values())
        if np.alltrue(keys + 1 == values):
            return zeta_data
        else:
            #            raise IncompleteAtomicData('zeta data')
            # This currently replaces missing zeta data with 1, which is necessary with
            # the present atomic data. Will replace with the error above when I have
            # complete atomic data.
            missing_ions = []
            updated_index = []
            for atom in selected_atoms:
                for ion in range(1, atom + 2):
                    if (atom, ion) not in zeta_data.index:
                        missing_ions.append((atom, ion))
                    updated_index.append([atom, ion])
            logger.warn(
                f"Zeta_data missing - replaced with 1s. Missing ions: {missing_ions}"
            )
            updated_index = np.array(updated_index)
            updated_dataframe = pd.DataFrame(
                index=pd.MultiIndex.from_arrays(
                    updated_index.transpose().astype(int)
                ),
                columns=zeta_data.columns,
            )
            for value in range(len(zeta_data)):
                updated_dataframe.loc[
                    zeta_data.atomic_number.values[value],
                    zeta_data.ion_number.values[value],
                ] = zeta_data.loc[
                    zeta_data.atomic_number.values[value],
                    zeta_data.ion_number.values[value],
                ]
            updated_dataframe = updated_dataframe.astype(float)
            updated_index = pd.DataFrame(updated_index)
            updated_dataframe["atomic_number"] = np.array(updated_index[0])
            updated_dataframe["ion_number"] = np.array(updated_index[1])
            updated_dataframe.fillna(1.0, inplace=True)
            return updated_dataframe

    def _set_index(self, zeta_data):
        return zeta_data.set_index(["atomic_number", "ion_number"])

class NLTEData(ProcessingPlasmaProperty):
    """
    Attributes
    ----------
    nlte_data :
        #Finish later (need atomic dataset with NLTE data).
    """

    outputs = ("nlte_data",)

    def calculate(self, atomic_data):
        if getattr(self, self.outputs[0]) is not None:
            return (getattr(self, self.outputs[0]),)
        else:
            return atomic_data.nlte_data

class YgData(ProcessingPlasmaProperty):
    """
    Attributes
    ----------
    yg_data : pandas.DataFrame
        Table of thermally averaged effective collision strengths
        (divided by the statistical weight of the lower level) Y_ij / g_i .
        Columns are temperatures.
    t_yg : numpy.ndarray
        Temperatures at which collision strengths are tabulated.
    yg_index : Pandas MultiIndex
    delta_E_yg : pandas.DataFrame
        Energy difference between upper and lower levels coupled by collisions.
    yg_idx : pandas.DataFrame
        Source_level_idx and destination_level_idx of collision transitions.
        Indexed by atomic_number, ion_number, level_number_lower,
        level_number_upper.
    """

    outputs = ("yg_data", "t_yg", "yg_index", "delta_E_yg", "yg_idx")
    latex_name = (
        r"\frac{Y_{ij}}{g_i}",
        r"T_\textrm{Yg}",
        r"\textrm{yg_index}",
        r"\delta E_{ij}",
        r"\textrm{yg_idx}",
    )

    def calculate(self, atomic_data, continuum_interaction_species):
        yg_data = atomic_data.yg_data

        mask_selected_species = yg_data.index.droplevel(
            ["level_number_lower", "level_number_upper"]
        ).isin(continuum_interaction_species)
        yg_data = yg_data[mask_selected_species]

        t_yg = yg_data.columns.values.astype(float)
        yg_data.columns = t_yg
        approximate_yg_data = self.calculate_yg_van_regemorter(
            atomic_data, t_yg, continuum_interaction_species
        )

        yg_data = yg_data.combine_first(approximate_yg_data)

        energies = atomic_data.levels.energy
        index = yg_data.index
        lu_index = index.droplevel("level_number_lower")
        ll_index = index.droplevel("level_number_upper")
        delta_E = energies.loc[lu_index].values - energies.loc[ll_index].values
        delta_E = pd.Series(delta_E, index=index)

        source_idx = atomic_data.macro_atom_references.loc[
            ll_index
        ].references_idx
        destination_idx = atomic_data.macro_atom_references.loc[
            lu_index
        ].references_idx
        yg_idx = pd.DataFrame(
            {
                "source_level_idx": source_idx.values,
                "destination_level_idx": destination_idx.values,
            },
            index=index,
        )
        return yg_data, t_yg, index, delta_E, yg_idx

    @staticmethod
    def calculate_yg_van_regemorter(
        atomic_data, t_electrons, continuum_interaction_species
    ):
        """
        Calculate collision strengths in the van Regemorter approximation.

        This function calculates thermally averaged effective collision
        strengths (divided by the statistical weight of the lower level)
        Y_ij / g_i using the van Regemorter approximation.

        Parameters
        ----------
        atomic_data : tardis.io.atom_data.AtomData
        t_electrons : numpy.ndarray
        continuum_interaction_species : pandas.MultiIndex

        Returns
        -------
        pandas.DataFrame
            Thermally averaged effective collision strengths
            (divided by the statistical weight of the lower level) Y_ij / g_i

        Notes
        -----
        See Eq. 9.58 in [2].

        References
        ----------
        .. [1] van Regemorter, H., “Rate of Collisional Excitation in Stellar
               Atmospheres.”, The Astrophysical Journal, vol. 136, p. 906, 1962.
               doi:10.1086/147445.
        .. [2] Hubeny, I. and Mihalas, D., "Theory of Stellar Atmospheres". 2014.
        """
        I_H = atomic_data.ionization_data.loc[(1, 1)]

        mask_selected_species = atomic_data.lines.index.droplevel(
            ["level_number_lower", "level_number_upper"]
        ).isin(continuum_interaction_species)
        lines_filtered = atomic_data.lines[mask_selected_species]
        f_lu = lines_filtered.f_lu.values
        nu_lines = lines_filtered.nu.values

        yg = f_lu * (I_H / (H * nu_lines)) ** 2
        coll_const = A0 ** 2 * np.pi * np.sqrt(8 * K_B / (np.pi * M_E))
        yg = 14.5 * coll_const * t_electrons * yg[:, np.newaxis]

        u0 = nu_lines[np.newaxis].T / t_electrons * (H / K_B)
        gamma = 0.276 * np.exp(u0) * expn(1, u0)
        gamma[gamma < 0.2] = 0.2
        yg *= u0 * gamma / BETA_COLL
        yg = pd.DataFrame(yg, index=lines_filtered.index, columns=t_electrons)

        return yg

class YgInterpolator(ProcessingPlasmaProperty):
    """
    Attributes
    ----------
    yg_interp : scipy.interpolate.PchipInterpolator
        Interpolates the thermally averaged effective collision strengths
        (divided by the statistical weight of the lower level) Y_ij / g_i as
        a function of electron temperature.
    """

    outputs = ("yg_interp",)
    latex_name = ("\\frac{Y_ij}{g_i}_{\\textrm{interp}}",)

    def calculate(self, yg_data, t_yg):
        yg_interp = PchipInterpolator(t_yg, yg_data, axis=1, extrapolate=True)
        return yg_interp


import sys
import base64

import numpy as np

from . import _marching_cubes_lewiner_luts as mcluts
from . import _marching_cubes_lewiner_cy

if sys.version_info >= (3, ):
    base64decode = base64.decodebytes
else:
    base64decode = base64.decodestring

def marching_cubes_lewiner(volume, level=None, spacing=(1., 1., 1.),
                           gradient_direction='descent', step_size=1,
                           allow_degenerate=True, use_classic=False):
    """
    Lewiner marching cubes algorithm to find surfaces in 3d volumetric data.

    In contrast to ``marching_cubes_classic()``, this algorithm is faster,
    resolves ambiguities, and guarantees topologically correct results.
    Therefore, this algorithm generally a better choice, unless there
    is a specific need for the classic algorithm.

    Parameters
    ----------
    volume : (M, N, P) array
        Input data volume to find isosurfaces. Will internally be
        converted to float32 if necessary.
    level : float
        Contour value to search for isosurfaces in `volume`. If not
        given or None, the average of the min and max of vol is used.
    spacing : length-3 tuple of floats
        Voxel spacing in spatial dimensions corresponding to numpy array
        indexing dimensions (M, N, P) as in `volume`.
    gradient_direction : string
        Controls if the mesh was generated from an isosurface with gradient
        descent toward objects of interest (the default), or the opposite,
        considering the *left-hand* rule.
        The two options are:
        * descent : Object was greater than exterior
        * ascent : Exterior was greater than object
    step_size : int
        Step size in voxels. Default 1. Larger steps yield faster but
        coarser results. The result will always be topologically correct
        though.
    allow_degenerate : bool
        Whether to allow degenerate (i.e. zero-area) triangles in the
        end-result. Default True. If False, degenerate triangles are
        removed, at the cost of making the algorithm slower.
    use_classic : bool
        If given and True, the classic marching cubes by Lorensen (1987)
        is used. This option is included for reference purposes. Note
        that this algorithm has ambiguities and is not guaranteed to
        produce a topologically correct result. The results with using
        this option are *not* generally the same as the
        ``marching_cubes_classic()`` function.

    Returns
    -------
    verts : (V, 3) array
        Spatial coordinates for V unique mesh vertices. Coordinate order
        matches input `volume` (M, N, P).
    faces : (F, 3) array
        Define triangular faces via referencing vertex indices from ``verts``.
        This algorithm specifically outputs triangles, so each face has
        exactly three indices.
    normals : (V, 3) array
        The normal direction at each vertex, as calculated from the
        data.
    values : (V, ) array
        Gives a measure for the maximum value of the data in the local region
        near each vertex. This can be used by visualization tools to apply
        a colormap to the mesh.

    Notes
    -----
    The algorithm [1] is an improved version of Chernyaev's Marching
    Cubes 33 algorithm. It is an efficient algorithm that relies on
    heavy use of lookup tables to handle the many different cases,
    keeping the algorithm relatively easy. This implementation is
    written in Cython, ported from Lewiner's C++ implementation.

    To quantify the area of an isosurface generated by this algorithm, pass
    verts and faces to `skimage.measure.mesh_surface_area`.

    Regarding visualization of algorithm output, to contour a volume
    named `myvolume` about the level 0.0, using the ``mayavi`` package::

      >>> from mayavi import mlab # doctest: +SKIP
      >>> verts, faces, normals, values = marching_cubes_lewiner(myvolume, 0.0) # doctest: +SKIP
      >>> mlab.triangular_mesh([vert[0] for vert in verts],
      ...                      [vert[1] for vert in verts],
      ...                      [vert[2] for vert in verts],
      ...                      faces) # doctest: +SKIP
      >>> mlab.show() # doctest: +SKIP

    Similarly using the ``visvis`` package::

      >>> import visvis as vv # doctest: +SKIP
      >>> verts, faces, normals, values = marching_cubes_lewiner(myvolume, 0.0) # doctest: +SKIP
      >>> vv.mesh(np.fliplr(verts), faces, normals, values) # doctest: +SKIP
      >>> vv.use().Run() # doctest: +SKIP

    References
    ----------
    .. [1] Thomas Lewiner, Helio Lopes, Antonio Wilson Vieira and Geovan
           Tavares. Efficient implementation of Marching Cubes' cases with
           topological guarantees. Journal of Graphics Tools 8(2)
           pp. 1-15 (december 2003).
           DOI: 10.1080/10867651.2003.10487582

    See Also
    --------
    skimage.measure.marching_cubes_classic
    skimage.measure.mesh_surface_area
    """

    # Check volume and ensure its in the format that the alg needs
    if not isinstance(volume, np.ndarray) or (volume.ndim != 3):
        raise ValueError('Input volume should be a 3D numpy array.')
    if volume.shape[0] < 2 or volume.shape[1] < 2 or volume.shape[2] < 2:
        raise ValueError("Input array must be at least 2x2x2.")
    volume = np.ascontiguousarray(volume, np.float32)  # no copy if not necessary

    # Check/convert other inputs:
    # level
    if level is None:
        level = 0.5 * (volume.min() + volume.max())
    else:
        level = float(level)
        if level < volume.min() or level > volume.max():
            raise ValueError("Surface level must be within volume data range.")
    # spacing
    if len(spacing) != 3:
        raise ValueError("`spacing` must consist of three floats.")
    # step_size
    step_size = int(step_size)
    if step_size < 1:
        raise ValueError('step_size must be at least one.')
    # use_classic
    use_classic = bool(use_classic)

    # Get LutProvider class (reuse if possible)
    L = _get_mc_luts()

    # Apply algorithm
    func = _marching_cubes_lewiner_cy.marching_cubes
    vertices, faces , normals, values = func(volume, level, L, step_size, use_classic)

    if not len(vertices):
        raise RuntimeError('No surface found at the given iso value.')

    # Output in z-y-x order, as is common in skimage
    vertices = np.fliplr(vertices)
    normals = np.fliplr(normals)

    # Finishing touches to output
    faces.shape = -1, 3
    if gradient_direction == 'descent':
        # MC implementation is right-handed, but gradient_direction is left-handed
        faces = np.fliplr(faces)
    elif not gradient_direction == 'ascent':
        raise ValueError("Incorrect input %s in `gradient_direction`, see "
                         "docstring." % (gradient_direction))
    if spacing != (1, 1, 1):
        vertices = vertices * np.r_[spacing]

    if allow_degenerate:
        return vertices, faces, normals, values
    else:
        fun = _marching_cubes_lewiner_cy.remove_degenerate_faces
        return fun(vertices.astype(np.float32), faces, normals, values)

def _to_array(args):
    shape, text = args
    byts = base64decode(text.encode('utf-8'))
    ar = np.frombuffer(byts, dtype='int8')
    ar.shape = shape
    return ar

# Map an edge-index to two relative pixel positions. The ege index
# represents a point that lies somewhere in between these pixels.
# Linear interpolation should be used to determine where it is exactly.
#   0
# 3   1   ->  0x
#   2         xx
EDGETORELATIVEPOSX = np.array([ [0,1],[1,1],[1,0],[0,0], [0,1],[1,1],[1,0],[0,0], [0,0],[1,1],[1,1],[0,0] ], 'int8')
EDGETORELATIVEPOSY = np.array([ [0,0],[0,1],[1,1],[1,0], [0,0],[0,1],[1,1],[1,0], [0,0],[0,0],[1,1],[1,1] ], 'int8')
EDGETORELATIVEPOSZ = np.array([ [0,0],[0,0],[0,0],[0,0], [1,1],[1,1],[1,1],[1,1], [0,1],[0,1],[0,1],[0,1] ], 'int8')

def _get_mc_luts():
    """ Kind of lazy obtaining of the luts.
    """
    if not hasattr(mcluts, 'THE_LUTS'):

        mcluts.THE_LUTS = _marching_cubes_lewiner_cy.LutProvider(
                EDGETORELATIVEPOSX, EDGETORELATIVEPOSY, EDGETORELATIVEPOSZ,

                _to_array(mcluts.CASESCLASSIC), _to_array(mcluts.CASES),

                _to_array(mcluts.TILING1), _to_array(mcluts.TILING2), _to_array(mcluts.TILING3_1), _to_array(mcluts.TILING3_2),
                _to_array(mcluts.TILING4_1), _to_array(mcluts.TILING4_2), _to_array(mcluts.TILING5), _to_array(mcluts.TILING6_1_1),
                _to_array(mcluts.TILING6_1_2), _to_array(mcluts.TILING6_2), _to_array(mcluts.TILING7_1),
                _to_array(mcluts.TILING7_2), _to_array(mcluts.TILING7_3), _to_array(mcluts.TILING7_4_1),
                _to_array(mcluts.TILING7_4_2), _to_array(mcluts.TILING8), _to_array(mcluts.TILING9),
                _to_array(mcluts.TILING10_1_1), _to_array(mcluts.TILING10_1_1_), _to_array(mcluts.TILING10_1_2),
                _to_array(mcluts.TILING10_2), _to_array(mcluts.TILING10_2_), _to_array(mcluts.TILING11),
                _to_array(mcluts.TILING12_1_1), _to_array(mcluts.TILING12_1_1_), _to_array(mcluts.TILING12_1_2),
                _to_array(mcluts.TILING12_2), _to_array(mcluts.TILING12_2_), _to_array(mcluts.TILING13_1),
                _to_array(mcluts.TILING13_1_), _to_array(mcluts.TILING13_2), _to_array(mcluts.TILING13_2_),
                _to_array(mcluts.TILING13_3), _to_array(mcluts.TILING13_3_), _to_array(mcluts.TILING13_4),
                _to_array(mcluts.TILING13_5_1), _to_array(mcluts.TILING13_5_2), _to_array(mcluts.TILING14),

                _to_array(mcluts.TEST3), _to_array(mcluts.TEST4), _to_array(mcluts.TEST6),
                _to_array(mcluts.TEST7), _to_array(mcluts.TEST10), _to_array(mcluts.TEST12),
                _to_array(mcluts.TEST13), _to_array(mcluts.SUBCONFIG13),
                )

    return mcluts.THE_LUTS