'''
Created on 24/04/2014

@author: MMPE
'''
import types
import os
import numpy as np


stc = "r m x_cg y_cg ri_x ri_y x_sh y_sh E G I_x I_y I_p k_x k_y A pitch x_e y_e"
fpm = 'r m x_cg y_cg ri_x ri_y pitch x_e y_e K_11 K_12 K_13 K_14 K_15 K_16 K_22'
fpm += ' K_23 K_24 K_25 K_26 K_33 K_34 K_35 K_36 K_44 K_45 K_46 K_55 K_56 K_66'

class StFile(object):
    """Read HAWC2 St (beam element structural data) file

    Methods are autogenerated for:

    - r : curved length distance from main_body node 1 [m]
    - m : mass per unit length [kg/m]
    - x_cg : xc2-coordinate from C1/2 to mass center [m]
    - y_cg : yc2-coordinate from C1/2 to mass center [m]
    - ri_x : radius of gyration related to elastic center. Corresponds to rotation about principal bending xe axis [m]
    - ri_y : radius of gyration related to elastic center. Corresponds to rotation about principal bending ye axis [m]
    - xs : xc2-coordinate from C1/2 to shear center [m]. The shear center is the point where external forces only contributes to pure bending and no torsion.
    - ys : yc2-coordinate from C1/2 to shear center [m]. The shear center is the point where external forces only contributes to pure bending and no torsion.
    - E : modulus of elasticity [N/m2]
    - G : shear modulus of elasticity [N/m2]
    - Ix : area moment of inertia with respect to principal bending xe axis [m4]. This is the principal bending axis most parallel to the xc2 axis
    - Iy : area moment of inertia with respect to principal bending ye axis [m4]
    - K : torsional stiffness constant with respect to ze axis at the shear center [m4/rad]. For a circular section only this is identical to the polar moment of inertia.
    - k_x : shear factor for force in principal bending xe direction [-]
    - k_y : shear factor for force in principal bending ye direction [-]
    - A : cross sectional area [m2]
    - pitch : structural pitch about z_c2 axis. This is the angle between the xc2 -axis defined with the c2_def command and the main principal bending axis xe.
    - xe : xc2-coordinate from C1/2 to center of elasticity [m]. The elastic center is the point where radial force (in the z-direction) does not contribute to bending around the x or y directions.
    - ye : yc2-coordinate from C1/2 to center of elasticity [m]. The elastic center is the point where radial force (in the

    The autogenerated methods have the following structure

    def xxx(radius=None, mset=1, set=1):
        Parameters:
        -----------
        radius : int, float, array_like or None, optional
            Radius/radii of interest\n
            If int, float or array_like: values are interpolated to requested radius/radii
            If None (default): Values of all radii specified in st file returned
        mset : int, optional
            Main set number
        set : int, optional
            Sub set number


    Examples
    --------
    >>> stfile = StFile(r"tests/test_files/DTU_10MW_RWT_Blade_st.dat")
    >>> print (stfile.m()) # Mass at nodes
    [ 1189.51054664  1191.64291781  1202.76694262  ... 15.42438683]
    >>> print (st.E(radius=36, mset=1, set=1))  # Elasticity interpolated to radius 36m
    8722924514.652649
    >>> print (st.E(radius=36, mset=1, set=2))  # Same for stiff blade set
    8.722924514652648e+17
    """

    def __init__(self, filename=None):

        # in case the user wants to create a new non-existing st file
        if filename is None:
            self.main_data_sets = {}
            return

        with open(filename) as fid:
            txt = fid.read()
#         Some files starts with first set ("#1...") with out specifying number of sets
#         no_maindata_sets = int(txt.strip()[0])
#         assert no_maindata_sets == txt.count("#")
        self.main_data_sets = {}
        for mset in txt.split("#")[1:]:
            mset_nr = int(mset.strip().split()[0])
            set_data_dict = {}

            for set_txt in mset.split("$")[1:]:
                set_lines = set_txt.split("\n")
                set_nr, no_rows = map(int, set_lines[0].split()[:2])
                assert set_nr not in set_data_dict
                try:
                    # HAWC2 will ignore everything after the 19th element,
                    # some users have placed comments here after a ;
                    linelst = [set_lines[i].split(';')[0].split() for i in range(1, no_rows + 1)]
                except Exception as e:
                    print('it went wrong at (set/subset):', mset_nr, set_nr,
                          'with', no_rows, 'rows')
                    raise e
                set_data_dict[set_nr] = np.array(linelst, dtype=float)
            self.main_data_sets[mset_nr] = set_data_dict

        if len(linelst[0])==len(stc.split()):
            self.cols = stc.split()
        elif len(linelst[0])==len(fpm.split()):
            self.cols = fpm.split()
        else:
            raise TypeError('wrong number of columns in st file')

        for i, name in enumerate(self.cols):
            setattr(self, name, lambda radius=None, mset=1, set=1,
                    column=i: self._value(radius, column, mset, set))

    def _value(self, radius, column, mset_nr=1, set_nr=1):
        st_data = self.main_data_sets[mset_nr][set_nr]
        if radius is None:
            radius = self.radius_st(None, mset_nr, set_nr)
        return np.interp(radius, st_data[:, 0], st_data[:, column])

    def radius_st(self, radius=None, mset=1, set=1):
        r = self.main_data_sets[mset][set][:, 0]
        if radius is None:
            return r
        return r[np.argmin(np.abs(r - radius))]

    def to_str(self, mset=1, set=1, precision='%12.5e '):
        d = self.main_data_sets[mset][set]
        return '\n'.join([(precision * d.shape[1]) % tuple(row) for row in d])

    def set_value(self, mset_nr, set_nr, **kwargs):
        for k, v in kwargs.items():
            column = self.cols.index(k)
            self.main_data_sets[mset_nr][set_nr][:, column] = v

    def save(self, filename, precision='%15.07e', encoding='utf-8'):
        """Save all data defined in main_data_sets to st file.
        """

        # when creating empty st object, cols is not set yet
        if not hasattr(self, 'cols'):
            if self.main_data_sets[1][1].shape[1]==19:
                self.cols = stc.split()
            elif self.main_data_sets[1][1].shape[1]==30:
                self.cols = fpm.split()
            else:
                c = self.main_data_sets[1][1].shape[1]
                raise TypeError(f'St input needs 19 (iso) or 30 (aniso/fpm) cols, not {c}')

        colwidth = len(precision % 1)
        sep = '=' * colwidth * len(self.cols) + '\n'
        colhead = ''.join([k.center(colwidth) for k in self.cols]) + '\n'

        nsets = len(self.main_data_sets)
        # fails if dirname is empty string
        if len(os.path.dirname(filename)) > 0:
            os.makedirs(os.path.dirname(filename), exist_ok=True)
        with open(filename, 'w', encoding=encoding) as fid:
            fid.write('%i ; number of sets, Nset\n' % nsets)
            for mset, set_data_dict in self.main_data_sets.items():
                fid.write('#%i ; set number\n' % mset)
                for set, set_array in set_data_dict.items():
                    dstr = self.to_str(mset=mset, set=set, precision=precision)
                    npoints = self.main_data_sets[mset][set].shape[0]
                    fid.write(sep + colhead + sep)
                    fid.write('$%i %i\n' % (set, npoints))
                    fid.write(dstr + '\n')

    def element_stiffnessmatrix(self, radius, mset_nr, set_nr, length):
        """Compute the element stiffness matrix

        Parameters
        ----------
        radius : float
            radius of element (used of obtain element properties
        length : float
            eleement length
        """

        # not supported for FPM format
        if len(self.cols)==30:
            return

        K = np.zeros((13, 13))
        "r m x_cg y_cg ri_x ri_y x_sh y_sh E G I_x I_y I_p k_x k_y A pitch x_e y_e"
        ES1, ES2, EMOD, GMOD, IX, IY, IZ, KX, KY, A = [getattr(self, n)(radius, mset_nr, set_nr)
                                                       for n in "x_sh,y_sh,E,G,I_x,I_y,I_p,k_x,k_y,A".split(",")]
        ELLGTH = length

        ETAX = EMOD * IX / (KY * GMOD * A * ELLGTH**2)
        ETAY = EMOD * IY / (KX * GMOD * A * ELLGTH**2)
        ROX = 1 / (1 + 12 * ETAX)
        ROY = 1 / (1 + 12 * ETAY)
        ROY = 1 / (1 + 12 * ETAX)
        K[1, 1] = 12 * EMOD * IY * ROY / ELLGTH**3
        K[1, 5] = 6 * EMOD * IY * ROY / ELLGTH**2
        K[1, 6] = -K[1, 1] * ES2
        K[1, 7] = -K[1, 1]
        K[1, 11] = K[1, 5]
        K[1, 12] = -K[1, 6]
        K[2, 2] = 12 * EMOD * IX * ROX / ELLGTH**3
        K[2, 4] = -6 * EMOD * IX * ROX / ELLGTH**2
        K[2, 6] = K[2, 2] * ES1
        K[2, 8] = -K[2, 2]
        K[2, 10] = K[2, 4]
        K[2, 12] = -K[2, 6]
        K[3, 3] = A * EMOD / ELLGTH
        K[3, 9] = -K[3, 3]
        K[4, 4] = 4 * EMOD * IX * (1 + 3 * ETAX) * ROX / ELLGTH
        K[4, 6] = K[2, 4] * ES1
        K[4, 8] = -K[2, 4]
        K[4, 10] = 2 * EMOD * IX * (1 - 6 * ETAX) * ROX / ELLGTH
        K[4, 12] = -K[4, 6]
        K[5, 5] = 4 * EMOD * IY * (1 + 3 * ETAY) * ROY / ELLGTH
        K[5, 6] = -K[1, 5] * ES2
        K[5, 7] = -K[1, 5]
        K[5, 11] = 2 * EMOD * IY * (1 - 6 * ETAY) * ROY / ELLGTH
        K[5, 12] = -K[5, 6]
        K[6, 6] = GMOD * IZ / ELLGTH + 12 * EMOD * (IX * ES1**2 * ROX + IY * ES2**2 * ROY) / ELLGTH**3
        K[6, 7] = K[1, 12]
        K[6, 8] = K[2, 12]
        K[6, 10] = -K[4, 12]
        K[6, 11] = -K[5, 12]
        K[6, 12] = -K[6, 6]
        K[7, 7] = K[1, 1]
        K[7, 11] = -K[1, 5]
        K[7, 12] = K[1, 6]
        K[8, 8] = K[2, 2]
        K[8, 10] = -K[2, 4]
        K[8, 12] = K[2, 6]
        K[9, 9] = K[3, 3]
        K[10, 10] = K[4, 4]
        K[10, 12] = -K[4, 6]
        K[11, 11] = K[5, 5]
        K[11, 12] = -K[5, 6]
        K[12, 12] = K[6, 6]
        K = K[1:, 1:]
        K = K + K.T - np.eye(12) * K
        return K

    def shape_function_ori(self, radius, mset_nr, set_nr, length, z):

        # not supported for FPM format
        if len(self.cols)==30:
            return

        XSC, YSC, EMOD, GMOD, IX, IY, IZ, KX, KY, AREA = [getattr(self, n)(radius, mset_nr, set_nr)
                                                          for n in "x_sh,y_sh,E,G,I_x,I_y,I_p,k_x,k_y,A".split(",")]

        etax = EMOD * IX / KY / GMOD / AREA / (length**2)
        etay = EMOD * IY / KX / GMOD / AREA / (length**2)
        rhox = 1 / (1 + 12 * etax)
        rhoy = 1 / (1 + 12 * etay)

        f1 = z / length
        f2 = 1 - f1
        f3x = f1 * (3 * f1 - 2 * f1**2 + 12 * etax) * rhox
        f4x = f2 * (3 * f2 - 2 * f2**2 + 12 * etax) * rhox
        f5x = length * (f1**2 - (1 - 6 * etax) * f1 - 6 * etax) * f1 * rhox
        f6x = length * (f2**2 - (1 - 6 * etax) * f2 - 6 * etax) * f2 * rhox
        f7x = 6 / length * f1 * f2 * rhox
        f8x = f1 * (3 * f1 - 2 * (1 - 6 * etax)) * rhox
        f9x = f2 * (3 * f2 - 2 * (1 - 6 * etax)) * rhox

        f3y = f1 * (3 * f1 - 2 * f1**2 + 12 * etay) * rhoy
        f4y = f2 * (3 * f2 - 2 * f2**2 + 12 * etay) * rhoy
        f5y = length * (f1**2 - (1 - 6 * etay) * f1 - 6 * etay) * f1 * rhoy
        f6y = length * (f2**2 - (1 - 6 * etay) * f2 - 6 * etay) * f2 * rhoy
        f7y = 6 / length * f1 * f2 * rhoy
        f8y = f1 * (3 * f1 - 2 * (1 - 6 * etay)) * rhoy
        f9y = f2 * (3 * f2 - 2 * (1 - 6 * etay)) * rhoy

        return np.array([[f4y, 0, 0, 0, -f7y, 0],
                         [0,  f4x, 0,  f7x, 0, 0],
                         [0, 0,   f2, 0, 0, 0],
                         [0, f6x, 0, f9x, 0, 0],
                         [-f6y, 0, 0, 0, f9y, 0],
                         [(f2 - f4y) * YSC, -(f2 - f4x) * XSC, 0, f7x * XSC, f7y * YSC, f2],
                         [f3y, 0, 0, 0, f7y, 0],
                         [0, f3x, 0, -f7x, 0, 0],
                         [0, 0,    f1, 0, 0, 0],
                         [0, -f5x, 0, f8x, 0, 0],
                         [f5y, 0, 0, 0, f8y, 0],
                         [(f1 - f3y) * YSC, -(f1 - f3x) * XSC, 0, -f7x * XSC, -f7y * YSC, f1]]).T


if __name__ == "__main__":
    import os
    cwd = os.path.dirname(__file__)
    st = StFile(os.path.join(cwd, r'tests/test_files/DTU_10MW_RWT_Blade_st.dat'))
    print(st.m())
    print(st.E(radius=36, mset=1, set=1))  # Elastic blade
    print(st.E(radius=36, mset=1, set=2))  # stiff blade
    #print (st.radius())
    xyz = np.array([st.x_e(), st.y_e(), st.r()]).T[:40]
    n = 2
    xyz = np.array([st.x_e(None, 1, n), st.y_e(None, 1, n), st.r(None, 1, n)]).T[:40]
    #print (xyz)
    print(np.sqrt(np.sum((xyz[1:] - xyz[:-1]) ** 2, 1)).sum())
    print(xyz[-1, 2])
    print(np.sqrt(np.sum((xyz[1:] - xyz[:-1]) ** 2, 1)).sum() - xyz[-1, 2])
    print(st.x_e(67.8883), st.y_e(67.8883))
    #print (np.sqrt(np.sum(np.diff(xyz, 0) ** 2, 1)))
    print(st.pitch(67.8883 - 0.01687))
    print(st.pitch(23.2446))

    # print (st.)
    # print (st.)
