diff --git a/wafo/spectrum/models.py b/wafo/spectrum/models.py
index 45619aa..59a85bd 100644
--- a/wafo/spectrum/models.py
+++ b/wafo/spectrum/models.py
@@ -540,7 +540,7 @@ class Jonswap(ModelSpectrum):
         if self.gamma is None or not isfinite(self.gamma) or self.gamma < 1:
             self.gamma = jonswap_peakfact(Hm0, Tp)
 
-        self._preCalculateAg()
+        self._pre_calculate_ag()
 
         if chk_seastate:
             self.chk_seastate()
@@ -568,57 +568,62 @@ class Jonswap(ModelSpectrum):
         Gf = self.peak_e_factor(wn)
         return Gf * _gengamspec(wn, self.N, self.M)
 
-    def _preCalculateAg(self):
+    def _parametric_ag(self):
+        self.method = 'parametric'  # Original normalization
+        # NOTE: that  Hm0**2/16 generally is not equal to intS(w)dw
+        # with this definition of Ag if sa or sb are changed from the
+        # default values
+        N = self.N
+        M = self.M
+        gammai = self.gamma
+        parametersOK = (3 <= N and N <= 50 or 2 <= M and M <= 9.5 and
+                        1 <= gammai and gammai <= 20)
+        f1NM = 4.1 * (N - 2 * M ** 0.28 + 5.3) ** (-1.45 * M ** 0.1 + 0.96)
+        f2NM = ((2.2 * M ** (-3.3) + 0.57) * N ** (-0.58 * M ** 0.37 + 0.53) -
+                1.04 * M ** (-1.9) + 0.94)
+        self.Ag = (1 + f1NM * log(gammai) ** f2NM) / gammai
+        if not parametersOK:
+            raise ValueError('Not knowing the normalization because N, ' +
+                             'M or peakedness parameter is out of bounds!')
+        # elseif N == 5 && M == 4,
+        #     options.Ag = (1+1.0*log(gammai).**1.16)/gammai
+        #     options.Ag = (1-0.287*log(gammai))
+        #     options.normalizeMethod = 'Three'
+        # elseif  N == 4 && M == 4,
+        #     options.Ag = (1+1.1*log(gammai).**1.19)/gammai
+        if self.sigmaA != 0.07 or self.sigmaB != 0.09:
+            warnings.warn('Use integration to calculate Ag when ' +
+                          'sigmaA!=0.07 or sigmaB!=0.09')
+
+    def _custom_ag(self):
+        self.method = 'custom'
+        if self.Ag <= 0:
+            raise ValueError('Ag must be larger than 0!')
+
+    def _integrate_ag(self):
+        # normalizing by integration
+        self.method = 'integration'
+        if self.wnc < 1.0:
+            raise ValueError('Normalized cutoff frequency, wnc, ' +
+                             'must be larger than one!')
+        area1, unused_err1 = integrate.quad(self._localspec, 0, 1)
+        area2, unused_err2 = integrate.quad(self._localspec, 1, self.wnc)
+        area = area1 + area2
+        self.Ag = 1.0 / area
+
+    def _pre_calculate_ag(self):
         ''' PRECALCULATEAG Precalculate normalization.
         '''
         if self.gamma == 1:
             self.Ag = 1.0
             self.method = 'parametric'
         elif self.Ag is not None:
-            self.method = 'custom'
-            if self.Ag <= 0:
-                raise ValueError('Ag must be larger than 0!')
-        elif self.method[0] == 'i':
-            # normalizing by integration
-            self.method = 'integration'
-            if self.wnc < 1.0:
-                raise ValueError('Normalized cutoff frequency, wnc, ' +
-                                 'must be larger than one!')
-            area1, unused_err1 = integrate.quad(self._localspec, 0, 1)
-            area2, unused_err2 = integrate.quad(self._localspec, 1, self.wnc)
-            area = area1 + area2
-            self.Ag = 1.0 / area
-        elif self.method[1] == 'p':
-            self.method = 'parametric'
-            # Original normalization
-            # NOTE: that  Hm0**2/16 generally is not equal to intS(w)dw
-            # with this definition of Ag if sa or sb are changed from the
-            # default values
-            N = self.N
-            M = self.M
-            gammai = self.gamma
-            parametersOK = (3 <= N and N <= 50) or (
-                2 <= M and M <= 9.5) and (1 <= gammai and gammai <= 20)
-            if parametersOK:
-                f1NM = 4.1 * \
-                    (N - 2 * M ** 0.28 + 5.3) ** (-1.45 * M ** 0.1 + 0.96)
-                f2NM = (2.2 * M ** (-3.3) + 0.57) * \
-                    N ** (-0.58 * M ** 0.37 + 0.53) - 1.04 * M ** (-1.9) + 0.94
-                self.Ag = (1 + f1NM * log(gammai) ** f2NM) / gammai
-
-            # elseif N == 5 && M == 4,
-            # options.Ag = (1+1.0*log(gammai).**1.16)/gammai
-            # options.Ag = (1-0.287*log(gammai))
-            # options.normalizeMethod = 'Three'
-            # elseif  N == 4 && M == 4,
-            # options.Ag = (1+1.1*log(gammai).**1.19)/gammai
-            else:
-                raise ValueError('Not knowing the normalization because N, ' +
-                                 'M or peakedness parameter is out of bounds!')
-
-            if self.sigmaA != 0.07 or self.sigmaB != 0.09:
-                warnings.warn('Use integration to calculate Ag when ' +
-                              'sigmaA!=0.07 or sigmaB!=0.09')
+            self._custom_ag()
+        else:
+            norm_ag = dict(i=self._integrate_ag,
+                           p=self._parametric_ag,
+                           c=self._custom_ag)[self.method[0]]
+            norm_ag()
 
     def peak_e_factor(self, wn):
         ''' PEAKENHANCEMENTFACTOR
@@ -1523,25 +1528,7 @@ class Spreading(object):
         spreadfun = self._spreadfun[self.type[0]]
         return spreadfun(theta, w, wc)
 
-    def chk_input(self, theta, w=1, wc=1):  # [s_par,TH,phi0,Nt] =
-        ''' CHK_INPUT
-
-        CALL [s_par,TH,phi0,Nt] = inputchk(theta,w,wc)
-        '''
-
-        wn = atleast_1d(w / wc)
-        theta = theta.ravel()
-        Nt = len(theta)
-
-        # Make sure theta is from -pi to pi
-        phi0 = 0.0
-        theta = mod(theta + pi, 2 * pi) - pi
-
-        if hasattr(self.theta0, '__call__'):
-            th0 = self.theta0(wn.flatten())
-        else:
-            th0 = atleast_1d(self.theta0).flatten()
-
+    def _normalize_angle(self, wn, theta, th0):
         Nt0 = th0.size
         Nw = wn.size
         isFreqDepDir = (Nt0 == Nw)
@@ -1557,20 +1544,31 @@ class Spreading(object):
             TH = mod(theta - th0 + pi, 2 * pi) - pi  # make sure -pi<=TH<pi
             if self.method is not None:  # frequency dependent spreading
                 TH = TH[:, newaxis]
+        return TH
 
-        # Get spreading parameter
-        s = self.spread_par_s(wn)
+    def _get_main_direction(self, wn):
+        if hasattr(self.theta0, '__call__'):
+            return self.theta0(wn.flatten())
+        return atleast_1d(self.theta0).flatten()
 
-        if self.type[0] == 'c':  # cos2s
-            s_par = s
-        else:
-            # First Fourier coefficient of the directional spreading function.
-            r1 = abs(s / (s + 1))
-            # Find distribution parameter from first Fourier coefficient.
-            s_par = self.fourier2distpar(r1)
-        if self.method is not None:
-            s_par = s_par[newaxis, :]
-        return s_par, TH, phi0, Nt
+    def chk_input(self, theta, w=1, wc=1):
+        ''' CHK_INPUT
+
+        CALL [s_par,TH,phi0,Nt] = inputchk(theta,w,wc)
+        '''
+
+        wn = atleast_1d(w / wc)
+        theta = theta.ravel()
+        Nt = len(theta)
+
+        # Make sure theta is from -pi to pi
+        phi0 = 0.0
+        theta = mod(theta + pi, 2 * pi) - pi
+        theta0 = self._get_main_direction(wn)
+
+        TH = self._normalize_angle(wn, theta, theta0)
+        s = self.spread_parameter_s(wn)
+        return s, TH, phi0, Nt
 
     def cos2s(self, theta, w=1, wc=1):  # [D, phi0] =
         ''' COS2S spreading function
@@ -1875,65 +1873,91 @@ class Spreading(object):
         r = clip(r1, 0., 1.0)
         return where(r <= 0, inf, sqrt(-2.0 * log(r)))
 
-    def spread_par_s(self, wn):
-        ''' Return spread parameter, S, of COS2S function
+    def _init_frequency_dependent_spreading(self, wn):
+        wn_lo, wn_up = self.wn_lo, self.wn_up
+        wn_c = self.wn_c
+        spa, spb = self.s_a, self.s_b
+        ma, mb = self.m_a, self.m_b
+
+        # Mitsuyasu et. al and Hasselman et. al parametrization   of
+        # frequency dependent spreading
+        s = where(wn <= wn_c, spa * wn ** ma, spb * wn ** mb)
+        s[wn <= wn_lo] = 0.0
+        return s, spb, wn_up, mb
+
+    def _donelan_spread(self, wn):
+        # Donelan et. al. parametrization for B in SECH-2
+        s, spb, wn_up, mb = self._init_frequency_dependent_spreading(wn)
+        k = flatnonzero(wn_up < wn)
+        s[k] = spb * (wn_up) ** mb
+        # Convert to S-paramater in COS-2S distribution
+        r1 = self.r1ofsech2(s)
+        s = r1 / (1. - r1)
+        return s
+
+    def _banner_spread(self, wn):
+        # Donelan et. al. parametrization for B in SECH-2
+        s, spb, wn_up, mb = self._init_frequency_dependent_spreading(wn)
+        k = flatnonzero(wn_up < wn)
+        # Banner parametrization  for B in SECH-2
+        s3m = spb * (wn_up) ** mb
+        s3p = self._donelan(wn_up)
+        #  Scale so that parametrization will be continous
+        scale = s3m / s3p
+        s[k] = scale * self.donelan(wn[k])
+        r1 = self.r1ofsech2(s)
+        # Convert to S-paramater in COS-2S distribution
+        s = r1 / (1. - r1)
+
+        return s
+
+    def _mitsuyasu_spread(self, wn):
+        s, _spb, wn_up, _mb = self._init_frequency_dependent_spreading(wn)
+        k = flatnonzero(wn_up < wn)
+        s[k] = 0
+        return s
+
+    def _frequency_independent_spread(self, _wn):
+        """
+        no frequency dependent spreading,
+        but possible frequency dependent direction
+        """
+        return atleast_1d(self.s_a)
+
+    def spread_parameter_s(self, wn):
+        ''' Return spread parameter, S, equivalent for the COS2S function
 
         Parameters
         ----------
         wn  : array_like
             normalized frequencies.
+
         Returns
         -------
         S   : ndarray
             spread parameter of COS2S functions
         '''
-        if self.method is None:
-            # no frequency dependent spreading,
-            # but possible frequency dependent direction
-            s = atleast_1d(self.s_a)
-        else:
-            wn_lo = self.wn_lo
-            wn_up = self.wn_up
-            wn_c = self.wn_c
-
-            spa = self.s_a
-            spb = self.s_b
-            ma = self.m_a
-            mb = self.m_b
-
-            # Mitsuyasu et. al and Hasselman et. al parametrization   of
-            # frequency dependent spreading
-            s = where(wn <= wn_c, spa * wn ** ma, spb * wn ** mb)
-            s[wn <= wn_lo] = 0.0
-
-            k = flatnonzero(wn_up < wn)
-            if k.size > 0:
-                if self.method[0] == 'd':
-                    # Donelan et. al. parametrization for B in SECH-2
-                    s[k] = spb * (wn_up) ** mb
-
-                    # Convert to S-paramater in COS-2S distribution
-                    r1 = self.r1ofsech2(s)
-                    s = r1 / (1. - r1)
-
-                elif self.method[0] == 'b':
-                    # Banner parametrization  for B in SECH-2
-                    s3m = spb * (wn_up) ** mb
-                    s3p = self._donelan(wn_up)
-                    #  Scale so that parametrization will be continous
-                    scale = s3m / s3p
-                    s[k] = scale * self.donelan(wn[k])
-                    r1 = self.r1ofsech2(s)
-
-                    # Convert to S-paramater in COS-2S distribution
-                    s = r1 / (1. - r1)
-                else:
-                    s[k] = 0.0
+
+        spread = dict(b=self._banner_spread,
+                      d=self._donelan_spread,
+                      m=self._mitsuyasu_spread
+                      ).get(self.method[0],
+                            self._frequency_independent_spread)
+        s = spread(wn)
 
         if any(s < 0):
             raise ValueError('The COS2S spread parameter, S(w), ' +
                              'value must be larger than 0')
-        return s
+        if self.type[0] == 'c':  # cos2s
+            s_par = s
+        else:
+            # First Fourier coefficient of the directional spreading function.
+            r1 = abs(s / (s + 1))
+            # Find distribution parameter from first Fourier coefficient.
+            s_par = self.fourier2distpar(r1)
+        if self.method is not None:
+            s_par = s_par[newaxis, :]
+        return s_par
 
     def _donelan(self, wn):
         ''' High frequency decay of B of sech2 paramater