/***************************************************************************
* Copyright (c) 2016, Johan Mabille, Sylvain Corlay and Wolf Vollprecht    *
*                                                                          *
* Distributed under the terms of the BSD 3-Clause License.                 *
*                                                                          *
* The full license is in the file LICENSE, distributed with this software. *
****************************************************************************/

#include "gtest/gtest.h"

#include <complex>
#include "xtensor/xarray.hpp"
#include "xtensor/xbuilder.hpp" 
#include "xtensor/xcomplex.hpp"

namespace xt
{
    using namespace std::complex_literals;

    TEST(xcomplex, expression)
    {
        xarray<std::complex<double>> e =
            {{1.0       , 1.0 + 1.0i},
             {1.0 - 1.0i, 1.0       }};

        // Test real expression
        auto r = real(e);
        auto i = imag(e);

        ASSERT_EQ(r.dimension(), 2);
        ASSERT_EQ(i.dimension(), 2);

        ASSERT_EQ(r.shape()[0], 2);
        ASSERT_EQ(r.shape()[1], 2);
        ASSERT_EQ(i.shape()[0], 2);
        ASSERT_EQ(i.shape()[1], 2);

        ASSERT_EQ(i(0, 0), 0);
        ASSERT_EQ(i(0, 1), 1);
        ASSERT_EQ(i(1, 0), -1);
        ASSERT_EQ(i(1, 1), 0);

        // Test assignment to an array
        xarray<double> ar = r;
        EXPECT_TRUE(all(equal(ar, ones<double>({2, 2}))));
    }

    TEST(xcomplex, lvalue)
    {
        xarray<std::complex<double>> e =
            {{1.0       , 1.0 + 1.0i},
             {1.0 - 1.0i, 1.0       }};

        // Test assigning an expression to the complex view
        real(e) = zeros<double>({2, 2});
        xarray<std::complex<double>> expect1 = 
            {{0.0       , 0.0 + 1.0i},
             {0.0 - 1.0i, 0.0       }};
        EXPECT_TRUE(all(equal(e, expect1)));

        imag(e) = zeros<double>({2, 2});
        EXPECT_TRUE(all(equal(e, zeros<std::complex<double>>({2, 2}))));
    }

    TEST(xcomplex, scalar_assignmnent)
    {
        xarray<std::complex<double>> e =
            {{1.0       , 1.0 + 1.0i},
             {1.0 - 1.0i, 1.0       }};

        // Test assigning an expression to the complex view
        real(e) = 0.0;
        xarray<std::complex<double>> expect1 = 
            {{0.0       , 0.0 + 1.0i},
             {0.0 - 1.0i, 0.0       }};
        EXPECT_TRUE(all(equal(e, expect1)));
    }

    TEST(xcomplex, noncomplex)
    {
        xarray<double> e = ones<double>({2, 2});
        auto r = real(e);
        auto i = imag(e);
        EXPECT_TRUE(all(equal(r, e)));
        EXPECT_TRUE(all(equal(i, zeros<double>({2, 2}))));
    }

    TEST(xcomplex, scalar)
    {
        double d = 1.0;
        ASSERT_EQ(1.0, real(d));
        ASSERT_EQ(0.0, imag(d));
        real(d) = 2.0;
        ASSERT_EQ(2.0, d);
    }

    TEST(xcomplex, pointer)
    {
        xarray<std::complex<double>> e =
            {{1.0       , 1.0 + 1.0i},
             {1.0 - 1.0i, 1.0       }};
        auto r = real(e);
        auto it = r.begin();
        EXPECT_EQ(*(it.operator->()), 1.0);
    }

    TEST(xcomplex, abs_angle_conj)
    {
        xarray<std::complex<double>> cmplarg_0 = {{0.40101756 + 0.71233018i, 0.62731701 + 0.42786349i, 0.32415089 + 0.2977805i},
                                                  {0.24475928 + 0.49208478i, 0.69475518 + 0.74029639i, 0.59390240 + 0.35772892i},
                                                  {0.63179202 + 0.41720995i, 0.44025718 + 0.65472131i, 0.08372648 + 0.37380143i}};
        auto cmplres = xt::abs(cmplarg_0);
        xarray<double> cmplexpected = {{0.81745298, 0.75933774, 0.44016704},
                                       {0.54959488, 1.01524554, 0.69331814},
                                       {0.75711643, 0.78897806, 0.38306348}};

        EXPECT_TRUE(allclose(cmplexpected, cmplres));

        auto cmplres_angle = xt::angle(cmplarg_0);
        xarray<double> cmplexpected_angle = {{1.05805307, 0.59857922, 0.74302273},
                                             {1.10923689, 0.81712241, 0.54213553},
                                             {0.58362348, 0.97881125, 1.35044673}};
        EXPECT_TRUE(allclose(cmplexpected_angle, cmplres_angle));

        auto cmplres_conj = xt::conj(cmplarg_0);
        xarray<std::complex<double>> cmplexpected_conj = {{0.40101756 - 0.71233018i, 0.62731701 - 0.42786349i, 0.32415089 - 0.2977805i},
                                                          {0.24475928 - 0.49208478i, 0.69475518 - 0.74029639i, 0.59390240 - 0.35772892i},
                                                          {0.63179202 - 0.41720995i, 0.44025718 - 0.65472131i, 0.08372648 - 0.37380143i}};
        EXPECT_TRUE(allclose(imag(cmplexpected_conj), imag(cmplres_conj)));
        EXPECT_TRUE(allclose(real(cmplexpected_conj), real(cmplres_conj)));

        auto cmplres_norm = xt::norm(cmplarg_0);
        xarray<double> fieldnorm = {{0.66822937, 0.5765938, 0.19374703},
                                    {0.30205453, 1.0307235, 0.48069004},
                                    {0.57322529, 0.62248637, 0.14673763}};

        EXPECT_TRUE(allclose(fieldnorm, cmplres_norm));
    }

    TEST(xcomplex, conj_real)
    {
        xarray<double> A = {{0.81745298, 0.75933774, 0.44016704},
                            {0.54959488, 1.01524554, 0.69331814},
                            {0.75711643, 0.78897806, 0.38306348}};
        xarray<double> B = xt::conj(A);
        EXPECT_EQ(A, B);
    }
}