/* Copyright (C) 2012-2019 IBM Corp.
 * This program is Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. See accompanying LICENSE file.
 */

#include <fstream>
#include <unistd.h>

#include <NTL/ZZX.h>
#include <NTL/vector.h>

#include <helib/helib.h>
#include <helib/debugging.h>

#include "gtest/gtest.h"
#include "test_common.h"

namespace {
struct Parameters {
    Parameters(long r, long p, long c, long mm) :
        r(r),
        p(p),
        c(c),
        mm(mm)
    {};
    long r; // lifting
    long p; // plaintext base
    long c; // number of columns in the key-switching matrices
    long mm; // cyclotomic index

    friend std::ostream& operator<<(std::ostream& os, const Parameters& params)
    {
        return os << "{"
            << "r=" << params.r << ","
            << "p=" << params.p << ","
            << "c=" << params.c << ","
            << "mm=" << params.mm
            << "}";
    };
};

class GTestIO : public ::testing::TestWithParam<Parameters> {
    protected:
#define N_TESTS 3
        static constexpr long ms[N_TESTS][10] = {
            //nSlots  m   phi(m) ord(2)
            //  {   2,    7,    6,    3,   0,0,0,0,0,0},
            {   6,   31,   30,    5,   0,0,0,0,0,0},
            {   6,  127,  126,    7,  127,  1,  108,  24,   6,-3}, // gens=108(6), 24(!3)
            {  60, 1023,  600,   10,   11, 93,  838, 584,  10, 6}, // gens=838(10),584(6)
            //  {  378,  5461,  5292, 14}, // gens=3(126),509(3)
            //  {  630,  8191,  8190, 13}, // gens=39(630)
            //  {  600, 13981, 12000, 20}, // gens=10(30),23(10),3(!2)
            //  {  682, 15709, 15004, 22} // gens=5(682)
        };

        static std::string keyFilePath;
        static std::string getTestResourcePath()
        {
            std::string testResourcePath;
            std::string executablePath(helib_test::path_of_executable);
            std::size_t found;
            if((found = executablePath.find_last_of("/")) == std::string::npos) {
                // Then there's no / in the filepath
                testResourcePath = "../test_resources";
            } else {
                testResourcePath = executablePath.substr(0, found) + "/../test_resources";
            }
            return testResourcePath;
        };

        GTestIO () :
            r(GetParam().r),
            p(GetParam().p),
            c(GetParam().c),
            w(64),
            L(100),
            mm(GetParam().mm),
            useTable(mm == 0 && p == 2),
            ptxtSpace(NTL::power_long(p,r)),
            numTests(useTable ? N_TESTS : 1),
            contexts(numTests),
            sKeys(numTests),
            ctxts(numTests),
            eas(numTests),
            ptxts(numTests)
            {};

        const long r;
        const long p;
        const long c;
        const long w;
        const long L;
        const long mm;
        const bool useTable;
        const long ptxtSpace;
        const long numTests;

        std::vector<std::unique_ptr<helib::Context>> contexts;
        std::vector<std::unique_ptr<helib::SecKey>> sKeys;
        std::vector<std::unique_ptr<helib::Ctxt>> ctxts;
        std::vector<std::unique_ptr<helib::EncryptedArray>> eas;
        std::vector<std::vector<NTL::ZZX>> ptxts;

        virtual void TearDown() override
        {
            helib::cleanupGlobals();
        }

    public:
        static void SetUpTestCase() {
            keyFilePath = getTestResourcePath() + "/iotest.txt";
        };

        static void TearDownTestCase() {
            unlink(keyFilePath.c_str()); // clean up before exiting
        };
};

constexpr long GTestIO::ms[N_TESTS][10];
std::string GTestIO::keyFilePath;

// Testing the I/O of the important classes of the library
// (context, keys, ciphertexts).
TEST_P(GTestIO, importantClassesRemainConsistentUnderIo)
{
  // first loop: generate stuff and write it to file

  // open file for writing
  {std::fstream keyFile(keyFilePath, std::fstream::out|std::fstream::trunc);
   ASSERT_TRUE(keyFile.is_open());
  for (long i=0; i<numTests; i++) {
    long m = (mm==0)? ms[i][1] : mm;

    if(!helib_test::noPrint) {
        std::cout << "Testing IO: m="<<m<<", p^r="<<p<<"^"<<r<<std::endl;
    }

    NTL::Vec<long> mvec(NTL::INIT_SIZE,2);
    mvec[0] = ms[i][4];  mvec[1] = ms[i][5];
    std::vector<long> gens(2);
    gens[0] = ms[i][6];  gens[1] = ms[i][7];
    std::vector<long> ords(2);
    ords[0] = ms[i][8];  ords[1] = ms[i][9];

    if (useTable && gens[0]>0)
      contexts[i].reset(new helib::Context(m, p, r, gens, ords));
    else
      contexts[i].reset(new helib::Context(m, p, r));
    if(!helib_test::noPrint) {
        contexts[i]->zMStar.printout();
    }

    helib::buildModChain(*contexts[i], L, c);  // Set the modulus chain
    if (mm==0 && m==1023) contexts[i]->makeBootstrappable(mvec);

    // Output the Context to file
    helib::writeContextBase(keyFile, *contexts[i]);
    if(!helib_test::noPrint) {
        helib::writeContextBase(std::cout, *contexts[i]);
        std::cout << std::endl;
    }
    keyFile << *contexts[i] << std::endl;

    sKeys[i].reset(new helib::SecKey(*contexts[i]));
    sKeys[i]->GenSecKey(0,ptxtSpace); // A +-1/0 secret key
    addSome1DMatrices(*sKeys[i]);// compute key-switching matrices that we need
    const helib::PubKey publicKey = *sKeys[i];
    eas[i].reset(new helib::EncryptedArray(*contexts[i]));

    long nslots = eas[i]->size();

    // Output the secret key to file, twice. Below we will have two copies
    // of most things.
    keyFile << *sKeys[i] << std::endl;;
    keyFile << *sKeys[i] << std::endl;;

    std::vector<NTL::ZZX> b;
    long p2r = eas[i]->getContext().alMod.getPPowR();
    NTL::ZZX poly = helib::RandPoly(0,NTL::to_ZZ(p2r)); // choose a random constant polynomial
    eas[i]->decode(ptxts[i], poly);

    ctxts[i].reset(new helib::Ctxt(publicKey));
    eas[i]->encrypt(*ctxts[i], publicKey, ptxts[i]);
    eas[i]->decrypt(*ctxts[i], *sKeys[i], b);
    ASSERT_EQ(ptxts[i].size(), b.size());
    for (long j = 0; j < nslots; j++) {
        ASSERT_EQ(ptxts[i][j], b[j]);
    }

    // output the plaintext
    keyFile << "[ ";
    for (long j = 0; j < nslots; j++) keyFile << ptxts[i][j] << " ";
    keyFile << "]\n";

    eas[i]->encode(poly,ptxts[i]);
    keyFile << poly << std::endl;

    // Output the ciphertext to file
    keyFile << *ctxts[i] << std::endl;
    keyFile << *ctxts[i] << std::endl;
    // std::cerr << "okay " << i << std::endl<< std::endl;
  }
  keyFile.close();}
  // std::cerr << "so far, so good\n\n";

  // second loop: read from input and repeat the computation

  // open file for read
  {std::fstream keyFile(keyFilePath, std::fstream::in);
  for (long i=0; i<numTests; i++) {

    // Read context from file
    unsigned long m1, p1, r1;
    std::vector<long> gens, ords;
    helib::readContextBase(keyFile, m1, p1, r1, gens, ords);
    helib::Context tmpContext(m1, p1, r1, gens, ords);
    keyFile >> tmpContext;
    ASSERT_EQ (*contexts[i], tmpContext);
    // std::cerr << i << ": context matches input\n";

    // We define some things below wrt *contexts[i], not tmpContext.
    // This is because the various operator== methods check equality of
    // references, not equality of the referenced Context objects.
    helib::Context& context = *contexts[i];
    helib::SecKey secretKey(context);
    helib::SecKey secretKey2(tmpContext);
    const helib::PubKey& publicKey = secretKey;
    const helib::PubKey& publicKey2 = secretKey2;

    keyFile >> secretKey;
    keyFile >> secretKey2;
    ASSERT_EQ(secretKey, *sKeys[i]);
    // std::cerr << "   secret key matches input\n";

    helib::EncryptedArray ea(context);
    helib::EncryptedArray ea2(tmpContext);

    long nslots = ea.size();

    // Read the plaintext from file
    std::vector<NTL::ZZX> a;
    a.resize(nslots);
    ASSERT_EQ(nslots, (long)ptxts[i].size());
    helib::seekPastChar(keyFile, '['); // defined in NumbTh.cpp
    for (long j = 0; j < nslots; j++) {
      keyFile >> a[j];
      ASSERT_EQ(a[j], ptxts[i][j]);
    }
    helib::seekPastChar(keyFile, ']');
    // std::cerr << "   ptxt matches input\n";

    // Read the encoded plaintext from file
    NTL::ZZX poly1, poly2;
    keyFile >> poly1;
    eas[i]->encode(poly2,a);
    ASSERT_EQ(poly1, poly2);
    // std::cerr << "   eas[i].encode(a)==poly1 okay\n";

    ea.encode(poly2,a);
    ASSERT_EQ(poly1, poly2);
    // std::cerr << "   ea.encode(a)==poly1 okay\n";

    ea2.encode(poly2,a);
    ASSERT_EQ(poly1, poly2);
    // std::cerr << "   ea2.encode(a)==poly1 okay\n";

    eas[i]->decode(a,poly1);
    ASSERT_EQ(nslots, (long)a.size());
    for (long j = 0; j < nslots; j++) {
        ASSERT_EQ(a[j], ptxts[i][j]);
    }
    // std::cerr << "   eas[i].decode(poly1)==ptxts[i] okay\n";

    ea.decode(a,poly1);
    ASSERT_EQ(nslots, (long)a.size());
    for (long j = 0; j < nslots; j++) {
        ASSERT_EQ(a[j], ptxts[i][j]);
    }
    // std::cerr << "   ea.decode(poly1)==ptxts[i] okay\n";

    ea2.decode(a,poly1);
    ASSERT_EQ(nslots, (long)a.size());
    for (long j = 0; j < nslots; j++) {
        ASSERT_EQ(a[j], ptxts[i][j]);
    }
    // std::cerr << "   ea2.decode(poly1)==ptxts[i] okay\n";

    // Read ciperhtext from file
    helib::Ctxt ctxt(publicKey);
    helib::Ctxt ctxt2(publicKey2);
    keyFile >> ctxt;
    keyFile >> ctxt2;
    ASSERT_TRUE(ctxts[i]->equalsTo(ctxt,/*comparePkeys=*/false));
    // std::cerr << "   ctxt matches input\n";

    sKeys[i]->Decrypt(poly2,*ctxts[i]);
    ASSERT_EQ(poly1, poly2);
    // std::cerr << "   sKeys[i]->decrypt(*ctxts[i]) == poly1 okay\n";

    secretKey.Decrypt(poly2,*ctxts[i]);
    ASSERT_EQ(poly1, poly2);
    // std::cerr << "   secretKey.decrypt(*ctxts[i]) == poly1 okay\n";

    secretKey.Decrypt(poly2,ctxt);
    ASSERT_EQ(poly1, poly2);
    // std::cerr << "   secretKey.decrypt(ctxt) == poly1 okay\n";

    secretKey2.Decrypt(poly2,ctxt2);
    ASSERT_EQ(poly1, poly2);
    // std::cerr << "   secretKey2.decrypt(ctxt2) == poly1 okay\n";

    eas[i]->decrypt(ctxt, *sKeys[i], a);
    ASSERT_EQ(nslots, (long)a.size());
    for (long j = 0; j < nslots; j++) {
        ASSERT_EQ(a[j], ptxts[i][j]);
    }
    // std::cerr << "   eas[i].decrypt(ctxt, *sKeys[i])==ptxts[i] okay\n";

    ea.decrypt(ctxt, secretKey, a);
    ASSERT_EQ(nslots, (long)a.size());
    for (long j = 0; j < nslots; j++) {
        ASSERT_EQ(a[j], ptxts[i][j]);
    }
    // std::cerr << "   ea.decrypt(ctxt, secretKey)==ptxts[i] okay\n";

    ea2.decrypt(ctxt2, secretKey2, a);
    ASSERT_EQ(nslots, (long)a.size());
    for (long j = 0; j < nslots; j++) {
        ASSERT_EQ(a[j], ptxts[i][j]);
    }
    // std::cerr << "   ea2.decrypt(ctxt2, secretKey2)==ptxts[i] okay\n";

  }}
}

INSTANTIATE_TEST_SUITE_P(someSmallParameters, GTestIO, ::testing::Values(
            //FAST
            Parameters(1, 2, 2, 91)
            ));

} // namespace
