from __future__ import print_function
import unittest
import operator
import numpy as np
import pyqg

class ZonalJetTester(unittest.TestCase):
    def setUp(self):
        a = 0.5
        c = 50.
        nx = 256
        b=nx/2.
        U1 = 0.1 + a * np.exp(-((np.arange(nx,dtype=float)-b)**2)/(2*c**2))
        U2 = 0.05 + a * np.exp(-((np.arange(nx,dtype=float)-b)**2)/(2*c**2))
        U=[U1, U2, 0.0*np.zeros((nx))]
        self.m = pyqg.LayeredModel(
            nx=nx,
            nz=3,
            U=U,
            V=[.1, .05, .0],
            rho=[.1, .3, .5],
            H=[.1, .1, .3],
            f=1.,
            beta=0.)

        self.atol = 1.e-16

        # creates stretching matrix from scratch
        self.S = np.zeros((self.m.nz, self.m.nz))

        F11 = self.m.f2 / self.m.gpi[0] / self.m.Hi[0]
        F12 = self.m.f2 / self.m.gpi[0] / self.m.Hi[1]
        F22 = self.m.f2 / self.m.gpi[1] / self.m.Hi[1]
        F23 = self.m.f2 / self.m.gpi[1] / self.m.Hi[2]

        self.S[0, 0], self.S[0, 1] = -F11, F11
        self.S[1, 0], self.S[1, 1], self.S[1, 2] = F12, -(F12 + F22), F22
        self.S[2, 1], self.S[2, 2] = F23, -F23

    def test_stretching(self):
        """ Check if stretching matrix is consistent
                and satisfies basic properties """

        # the columns of the S must add to zero (i.e, S is singular)
        err_msg = ' Zero is not an eigenvalue of S'
        assert np.all(self.m.S.sum(axis=1) == 0.), err_msg

        # the matrix Hi * S must by a symmetric matrix
        HS = np.dot(np.diag(self.m.Hi), self.m.S)
        np.testing.assert_allclose(
            HS, HS.T, atol=self.atol, err_msg=' Hi*S is not symmetric')

        np.testing.assert_allclose(
            self.m.S,
            self.S,
            atol=self.atol,
            err_msg=' Unmatched stretching matrix')

    def test_init_background(self):
        """ Check the initialization of the mean PV gradients  """
        Qy = -np.einsum('ij...,j...->i...', self.S, self.m.Ubg)
        Qx = np.einsum('ij...,j...->i...', self.S, self.m.Vbg)

        np.testing.assert_allclose(
            Qy, self.m.Qy, atol=self.atol, err_msg=' Unmatched Qy')
        np.testing.assert_allclose(
            Qx, self.m.Qx, atol=self.atol, err_msg=' Unmatched Qx ')
    
    def test_jet_pv(self):
        """ Check that the meridional mean PV gradient includes impact of jet """
        m = self.m
        Qy_jet = np.gradient(np.gradient(m.Ubg, m.dy, axis=1), m.dy, axis=1)
        Qy = m.beta - np.dot(m.S, m.Ubg) + Qy_jet
        np.testing.assert_allclose(
            Qy, self.m.Qy, atol=self.atol, err_msg=' Unmatched Qy')
        np.testing.assert_equal(
            np.any(np.not_equal(0*Qy_jet, Qy_jet)), True, err_msg=' Unmatched Jet PV gradient')

    def test_inversion_matrix(self):
        """ Check the inversion matrix """

        # it suffices to test for a single wavenumber
        M = self.S - np.eye(self.m.nz) * self.m.wv2[5, 5]
        Minv = np.zeros_like(M)

        detM = np.linalg.det(M)

        Minv[0, 0] = M[1, 1] * M[2, 2] - M[1, 2] * M[2, 1]
        Minv[0, 1] = M[0, 2] * M[2, 1] - M[0, 1] * M[2, 2]
        Minv[0, 2] = M[0, 1] * M[1, 2] - M[0, 2] * M[1, 1]

        Minv[1, 0] = M[1, 2] * M[2, 0] - M[1, 0] * M[2, 2]
        Minv[1, 1] = M[0, 0] * M[2, 2] - M[0, 2] * M[2, 0]
        Minv[1, 2] = M[0, 2] * M[1, 0] - M[0, 0] * M[1, 2]

        Minv[2, 0] = M[1, 0] * M[2, 1] - M[1, 1] * M[2, 0]
        Minv[2, 1] = M[0, 1] * M[2, 0] - M[0, 0] * M[2, 1]
        Minv[2, 2] = M[0, 0] * M[1, 1] - M[0, 1] * M[1, 0]

        Minv = Minv / detM

        np.testing.assert_allclose(
            self.m.a[:, :, 5, 5],
            Minv,
            atol=self.atol,
            err_msg=' Unmatched inversion matrix ')


if __name__ == "__main__":
    unittest.main()
